Encode images to smplx parameters Args: image: [bz, 3, 224, 224], range [0,1] image_hd(needed if key==body): a high res version of image, only for cropping parts from body image Returns: param_dict: dict key: image_type (body/head/
(self, image, image_hd, threthold=True, keep_local=True, copy_and_paste=False, body_only=False)
| 192 | |
| 193 | @torch.no_grad() |
| 194 | def forward(self, image, image_hd, threthold=True, keep_local=True, copy_and_paste=False, body_only=False): |
| 195 | ''' Encode images to smplx parameters |
| 196 | Args: |
| 197 | image: [bz, 3, 224, 224], range [0,1] |
| 198 | image_hd(needed if key==body): a high res version of image, only for cropping parts from body image |
| 199 | Returns: |
| 200 | param_dict: dict |
| 201 | key: image_type (body/head/hand) |
| 202 | value: param_dict |
| 203 | ''' |
| 204 | feature = {} |
| 205 | param_dict = {} |
| 206 | |
| 207 | # Encode features |
| 208 | part = key = 'body' |
| 209 | # encode feature |
| 210 | feature[key] = {} |
| 211 | feature[key][part] = self.Encoder[part](image) |
| 212 | |
| 213 | # for body image |
| 214 | if key == 'body': |
| 215 | fusion_weight = {} |
| 216 | f_body = feature['body']['body'] |
| 217 | # extract part feature |
| 218 | for part_name in ['head', 'left_hand', 'right_hand']: |
| 219 | feature['body'][f'{part_name}_share'] = self.Extractor[f'{part_name}_share'](f_body) |
| 220 | |
| 221 | # -- check if part crops are given, if not, crop parts by coarse body estimation |
| 222 | # - run without fusion to get coarse estimation, for cropping parts |
| 223 | # body only |
| 224 | body_dict = self.decompose_code(self.Regressor[part]( |
| 225 | feature[key][part]), self.param_list_dict[part+'_list']) |
| 226 | # head share |
| 227 | head_share_dict = self.decompose_code(self.Regressor['head'+'_share']( |
| 228 | feature[key]['head'+'_share']), self.param_list_dict['head'+'_share_list']) |
| 229 | # right hand share |
| 230 | right_hand_share_dict = self.decompose_code(self.Regressor['hand'+'_share']( |
| 231 | feature[key]['right_hand'+'_share']), self.param_list_dict['hand'+'_share_list']) |
| 232 | # left hand share |
| 233 | left_hand_share_dict = self.decompose_code(self.Regressor['hand'+'_share']( |
| 234 | feature[key]['left_hand'+'_share']), self.param_list_dict['hand'+'_share_list']) |
| 235 | # change the dict name from right to left |
| 236 | left_hand_share_dict['left_hand_pose'] = left_hand_share_dict.pop( |
| 237 | 'right_hand_pose') |
| 238 | left_hand_share_dict['left_wrist_pose'] = left_hand_share_dict.pop( |
| 239 | 'right_wrist_pose') |
| 240 | param_dict[key] = {**body_dict, **head_share_dict, |
| 241 | **left_hand_share_dict, **right_hand_share_dict} |
| 242 | if body_only: |
| 243 | param_dict['moderator_weight'] = None |
| 244 | return param_dict |
| 245 | |
| 246 | prediction_body_only = self.decode( |
| 247 | param_dict[key], param_type='body') |
| 248 | # crop |
| 249 | data = {key: {}} |
| 250 | for part_name in ['head', 'left_hand', 'right_hand']: |
| 251 | part = part_name.split('_')[-1] |
nothing calls this directly
no test coverage detected