mmpose_pytorch_runner.py

Go to the documentation of this file.

 
# import torch
# import numpy as np
 
# from mmpose.apis import init_model
# from smutils.utils_image import crop_image
 
# class MMPosePyTorchRunner(object):
#     def __init__(self, model_cfg, checkpoint, device='cuda', batch_size=1):
#         """
#         input:
#             deploy_cfg : 
#         """
#         self.model = init_model(model_cfg, checkpoint, device=device)
#         self.batch_size = batch_size
#         self.test_pipeline = self.build_test_pipeline(self.model.cfg)
 
#     def build_test_pipeline(self, cfg):
#         from mmcv.transforms import Compose
        
#         test_pipeline = Compose(cfg.test_dataloader.dataset.pipeline)
#         return test_pipeline
 
#     def make_src(self, image, image_id, bboxes=None):
#         """
#         return:
#             crop_images
#             bboxes
#             track_ids
#             image_ids
#         """
#         if bboxes is None:
#             h, w = image.shape[:2]
#             return [image], np.array([[0, 0, w, h]]), np.array([0]), [image_id]
        
#         num_person, dim = bboxes.shape
#         box_index = 1 if dim == 6 else 0
#         track_ids = bboxes[:, 0] if dim == 6 else np.arange(num_person)
#         image_ids = [image_id]*num_person
#         crop_images = crop_image(image, bboxes[:, box_index:box_index+4])
 
#         return crop_images, bboxes[:, box_index:box_index+4], track_ids, image_ids
 
#     def divide_into_batches(self, data, batch_size):
#         if batch_size == 1:
#             return data
 
#         batches = [data[i:i+batch_size] for i in range(0, len(data), batch_size)]
#         return batches
 
#     def run_detector(self, images, bboxes=None):
#         is_batch = isinstance(images, (list, tuple))
 
#         images = [images] if not is_batch else images
#         if not is_batch:
#             bboxes = [bboxes]
#         elif bboxes is None:
#             bboxes = [None]*len(images)
 
#         crop_images, bbox_list, track_ids, image_ids = [], [], [], []
#         for image_id, (image, bbox) in enumerate(zip(images, bboxes)):            
#             crops, bbs, tids, imids = self.make_src(image, image_id, bbox)
#             crop_images.extend(crops)
#             bbox_list.append(bbs)
#             track_ids.append(tids)
#             image_ids.extend(imids)
 
#         bbox_list = np.vstack(bbox_list)
#         track_ids = np.hstack(track_ids)
 
#         crop_images = self.divide_into_batches(crop_images, self.batch_size)
#         with torch.no_grad():
#             results = [item for image in crop_images for item in self.model.test_step(self.task_processor.create_input(image, self.input_shape, test_pipeline=self.test_pipeline)[0])]
 
#         pose_results = []
#         pre_img_id = -1
#         for res, bbox, track_id, img_id in zip(results, bbox_list, track_ids, image_ids):
#             keypoints = res.pred_instances.keypoints
#             keypoint_scores = np.expand_dims(res.pred_instances.keypoint_scores, axis=-1)
#             keypoints = np.dstack((keypoints, keypoint_scores))
            
#             if pre_img_id != img_id:
#                 pre_img_id = img_id
#                 pose_results.append([])
 
#             data = {'track_id': track_id, 'bbox': bbox, 'keypoints': keypoints}
#             pose_results[-1].append(data)
 
#         if is_batch:
#             return pose_results
#         else:
#             return pose_results[0]