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

Method prepare_targets

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

Source from the content-addressed store, hash-verified

1840 outputs_smpl_kp3d[:-1])]
1841
1842 def prepare_targets(self, data_batch):
1843
1844 data_batch_coco = []
1845 instance_dict = {}
1846 img_list = data_batch['img'].float()
1847 batch_size, _, input_img_h, input_img_w = img_list.shape
1848 device = img_list.device
1849 masks = torch.ones((batch_size, input_img_h, input_img_w),
1850 dtype=torch.bool,
1851 device=device)
1852
1853
1854 # cv2.imread(data_batch['img_metas'][img_id]['image_path']).shape
1855 for img_id in range(batch_size):
1856 img_h, img_w = data_batch['img_shape'][img_id]
1857 masks[img_id, :img_h, :img_w] = 0
1858
1859 if not self.inference:
1860 instance_body_bbox = torch.cat([data_batch['body_bbox_center'][img_id],\
1861 data_batch['body_bbox_size'][img_id]],dim=-1)
1862 instance_face_bbox = torch.cat([data_batch['face_bbox_center'][img_id],\
1863 data_batch['face_bbox_size'][img_id]],dim=-1)
1864 instance_lhand_bbox = torch.cat([data_batch['lhand_bbox_center'][img_id],\
1865 data_batch['lhand_bbox_size'][img_id]],dim=-1)
1866 instance_rhand_bbox = torch.cat([data_batch['rhand_bbox_center'][img_id],\
1867 data_batch['rhand_bbox_size'][img_id]],dim=-1)
1868
1869 instance_kp2d = data_batch['joint_img'][img_id].clone().float()
1870 instance_kp2d_mask = data_batch['joint_trunc'][img_id].clone().float()
1871 instance_kp2d[:,:,2:] = instance_kp2d_mask
1872 body_kp2d, _ = convert_kps(instance_kp2d, 'smplx_137', 'coco', approximate=True)
1873 lhand_kp2d, _ = convert_kps(instance_kp2d, 'smplx_137', 'smplx_lhand', approximate=True)
1874 rhand_kp2d, _ = convert_kps(instance_kp2d, 'smplx_137', 'smplx_rhand', approximate=True)
1875 face_kp2d, _ = convert_kps(instance_kp2d, 'smplx_137', 'smplx_face', approximate=True)
1876 body_kp2d[:,:,0] = body_kp2d[:,:,0]/cfg.output_hm_shape[2]
1877 body_kp2d[:,:,1] = body_kp2d[:,:,1]/cfg.output_hm_shape[1]
1878 body_kp2d = torch.cat([body_kp2d[:,:,:2].flatten(1),body_kp2d[:,:,2]],dim=-1)
1879
1880 lhand_kp2d[:,:,0] = lhand_kp2d[:,:,0]/cfg.output_hm_shape[2]
1881 lhand_kp2d[:,:,1] = lhand_kp2d[:,:,1]/cfg.output_hm_shape[1]
1882 lhand_kp2d = torch.cat([lhand_kp2d[:,:,:2].flatten(1),lhand_kp2d[:,:,2]],dim=-1)
1883
1884 rhand_kp2d[:,:,0] = rhand_kp2d[:,:,0]/cfg.output_hm_shape[2]
1885 rhand_kp2d[:,:,1] = rhand_kp2d[:,:,1]/cfg.output_hm_shape[1]
1886 rhand_kp2d = torch.cat([rhand_kp2d[:,:,:2].flatten(1),rhand_kp2d[:,:,2]],dim=-1)
1887
1888 face_kp2d[:,:,0] = face_kp2d[:,:,0]/cfg.output_hm_shape[2]
1889 face_kp2d[:,:,1] = face_kp2d[:,:,1]/cfg.output_hm_shape[1]
1890 face_kp2d = torch.cat([face_kp2d[:,:,:2].flatten(1),face_kp2d[:,:,2]],dim=-1)
1891
1892 instance_dict = {}
1893 instance_dict['boxes'] = instance_body_bbox.float()
1894 instance_dict['face_boxes'] = instance_face_bbox.float()
1895 instance_dict['lhand_boxes'] = instance_lhand_bbox.float()
1896 instance_dict['rhand_boxes'] = instance_rhand_bbox.float()
1897 instance_dict['keypoints'] = body_kp2d.float()
1898 instance_dict['lhand_keypoints'] = lhand_kp2d.float()
1899 instance_dict['rhand_keypoints'] = rhand_kp2d.float()

Callers 1

forwardMethod · 0.95

Calls 3

convert_kpsFunction · 0.90
NestedTensorClass · 0.90
cloneMethod · 0.80

Tested by

no test coverage detected