Method inference

datasets/INFERENCE_demo.py:107–168

(self, outs)

        return result
        
    def inference(self, outs):
        img_paths = self.img_paths
        for out in outs:
            ann_idx = out['image_idx']
            scores = out['scores'].clone().cpu().numpy()
            img_shape = out['img_shape'].cpu().numpy()[::-1] # w, h
            img = cv2.imread(img_paths[ann_idx]) # h, w
            scale = img.shape[1]/img_shape[0]
            body_bbox = out['body_bbox'].clone().cpu().numpy()
            body_bbox = body_bbox * scale
            joint_3d, _ =  convert_kps(out['smpl_kp3d'].clone().cpu().numpy(),src='smplx',dst='smplx', approximate=True)

            for i, score in enumerate(scores):
                if score < self.score_threshold:
                    break
                if i>self.num_person:
                    break
                save_name = img_paths[ann_idx].split('/')[-1]
                save_name = save_name.split('.')[0] 
                vert = out['smpl_verts'][i] + out['cam_trans'][i][None]
                # save mesh
                exist_result_path = glob(osp.join(self.mesh_path, save_name + '*'))
                if len(exist_result_path) == 0:
                    person_idx = 0
                else:
                    last_person_idx = max([
                        int(name.split('personId_')[1].split('.obj')[0])
                        for name in exist_result_path
                    ])
                    person_idx = last_person_idx + 1

                save_name += '_personId_' + str(person_idx) + '.obj'
                save_obj(osp.join(self.mesh_path, save_name), vert, faces=torch.tensor(self.body_model.faces.astype(np.int32)))
            
            if i == 0:
                save_name = img_paths[ann_idx].split('/')[-1][:-4]
                cv2.imwrite(os.path.join(self.result_img_dir,img_paths[ann_idx].split('/')[-1]), img)
            else:
                verts = out['smpl_verts'][:i] + out['cam_trans'][:i][:, None] 
                img = mmcv.imshow_bboxes(img, body_bbox[:i], show=False, colors='green') 
                render_smpl(
                    verts=verts[None],
                    body_model=self.body_model,
                    K= np.array(
                        [[5000, 0, img_shape[0]/2],
                         [0, 5000, img_shape[1]/2],
                         [0, 0, 1]]),
                    R=None,
                    T=None,
                    output_path=os.path.join(self.result_img_dir,img_paths[ann_idx].split('/')[-1]),
                    image_array=cv2.resize(img, (img_shape[0],img_shape[1]), cv2.INTER_CUBIC),
                    in_ndc=False,
                    alpha=0.9,
                    convention='opencv',
                    projection='perspective',
                    overwrite=True,
                    no_grad=True,
                    device='cuda',