MCPcopy Create free account
hub / github.com/MotrixLab/AiOS / prepare_targets

Method prepare_targets

models/aios/aios_smplx.py:3620–3710  ·  view source on GitHub ↗
(self, data_batch)

Source from the content-addressed store, hash-verified

3618 outputs_smpl_kp3d[:-1])]
3619
3620 def prepare_targets(self, data_batch):
3621
3622 data_batch_coco = []
3623 instance_dict = {}
3624 img_list = data_batch['img'].float()
3625 # input_img_h, input_img_w = data_batch['image_metas'][0]['batch_input_shape']
3626 batch_size, _, input_img_h, input_img_w = img_list.shape
3627 device = img_list.device
3628 masks = torch.ones((batch_size, input_img_h, input_img_w),
3629 dtype=torch.bool,
3630 device=device)
3631
3632 if self.num_body_points == 17:
3633 ed_convention = 'coco'
3634 elif self.num_body_points == 14:
3635 ed_convention = 'crowdpose'
3636
3637 # cv2.imread(data_batch['img_metas'][img_id]['image_path']).shape
3638 for img_id in range(batch_size):
3639 img_h, img_w = data_batch['img_shape'][img_id]
3640 masks[img_id, :img_h, :img_w] = 0
3641
3642 if not self.inference:
3643 instance_body_bbox = torch.cat([data_batch['body_bbox_center'][img_id],\
3644 data_batch['body_bbox_size'][img_id]],dim=-1)
3645 instance_face_bbox = torch.cat([data_batch['face_bbox_center'][img_id],\
3646 data_batch['face_bbox_size'][img_id]],dim=-1)
3647 instance_lhand_bbox = torch.cat([data_batch['lhand_bbox_center'][img_id],\
3648 data_batch['lhand_bbox_size'][img_id]],dim=-1)
3649 instance_rhand_bbox = torch.cat([data_batch['rhand_bbox_center'][img_id],\
3650 data_batch['rhand_bbox_size'][img_id]],dim=-1)
3651
3652 instance_kp2d = data_batch['joint_img'][img_id].clone().float()
3653 instance_kp2d_mask = data_batch['joint_trunc'][img_id].clone().float()
3654 instance_kp2d[:,:,2:] = instance_kp2d_mask
3655 body_kp2d, _ = convert_kps(instance_kp2d, 'smplx_137', 'coco', approximate=True)
3656 lhand_kp2d, _ = convert_kps(instance_kp2d, 'smplx_137', 'smplx_lhand', approximate=True)
3657 rhand_kp2d, _ = convert_kps(instance_kp2d, 'smplx_137', 'smplx_rhand', approximate=True)
3658 face_kp2d, _ = convert_kps(instance_kp2d, 'smplx_137', 'smplx_face', approximate=True)
3659 # from util.vis_utils import show_bbox
3660 # show_bbox(img_list[img_id],instance_kp2d.cpu().numpy(),data_batch['bbox_xywh'][img_id].cpu().numpy)
3661 body_kp2d[:,:,0] = body_kp2d[:,:,0]/cfg.output_hm_shape[2]
3662 body_kp2d[:,:,1] = body_kp2d[:,:,1]/cfg.output_hm_shape[1]
3663 body_kp2d = torch.cat([body_kp2d[:,:,:2].flatten(1),body_kp2d[:,:,2]],dim=-1)
3664
3665 lhand_kp2d[:,:,0] = lhand_kp2d[:,:,0]/cfg.output_hm_shape[2]
3666 lhand_kp2d[:,:,1] = lhand_kp2d[:,:,1]/cfg.output_hm_shape[1]
3667 lhand_kp2d = torch.cat([lhand_kp2d[:,:,:2].flatten(1),lhand_kp2d[:,:,2]],dim=-1)
3668
3669 rhand_kp2d[:,:,0] = rhand_kp2d[:,:,0]/cfg.output_hm_shape[2]
3670 rhand_kp2d[:,:,1] = rhand_kp2d[:,:,1]/cfg.output_hm_shape[1]
3671 rhand_kp2d = torch.cat([rhand_kp2d[:,:,:2].flatten(1),rhand_kp2d[:,:,2]],dim=-1)
3672
3673 face_kp2d[:,:,0] = face_kp2d[:,:,0]/cfg.output_hm_shape[2]
3674 face_kp2d[:,:,1] = face_kp2d[:,:,1]/cfg.output_hm_shape[1]
3675 face_kp2d = torch.cat([face_kp2d[:,:,:2].flatten(1),face_kp2d[:,:,2]],dim=-1)
3676
3677 instance_dict = {}

Callers 1

forwardMethod · 0.95

Calls 3

convert_kpsFunction · 0.90
NestedTensorClass · 0.90
cloneMethod · 0.80

Tested by

no test coverage detected