Encode images to smplx parameters Args: data: dict key: image_type (body/head/hand) value: image: [bz, 3, 224, 224], range [0,1] image_hd(needed if key==body): a high res version of image, only for cropping
(self, data, threthold=True, keep_local=True, copy_and_paste=False, body_only=False)
| 238 | |
| 239 | @torch.no_grad() |
| 240 | def encode(self, data, threthold=True, keep_local=True, copy_and_paste=False, body_only=False): |
| 241 | ''' Encode images to smplx parameters |
| 242 | Args: |
| 243 | data: dict |
| 244 | key: image_type (body/head/hand) |
| 245 | value: |
| 246 | image: [bz, 3, 224, 224], range [0,1] |
| 247 | image_hd(needed if key==body): a high res version of image, only for cropping parts from body image |
| 248 | head_image: optinal, well-cropped head from body image |
| 249 | left_hand_image: optinal, well-cropped left hand from body image |
| 250 | right_hand_image: optinal, well-cropped right hand from body image |
| 251 | Returns: |
| 252 | param_dict: dict |
| 253 | key: image_type (body/head/hand) |
| 254 | value: param_dict |
| 255 | ''' |
| 256 | for key in data.keys(): |
| 257 | assert key in ['body', 'head', 'hand'] |
| 258 | |
| 259 | feature = {} |
| 260 | param_dict = {} |
| 261 | |
| 262 | # Encode features |
| 263 | for key in data.keys(): |
| 264 | part = key |
| 265 | # encode feature |
| 266 | feature[key] = {} |
| 267 | feature[key][part] = self.Encoder[part](data[key]['image']) |
| 268 | |
| 269 | # for head/hand image |
| 270 | if key == 'head' or key == 'hand': |
| 271 | # predict head/hand-only parameters from part feature |
| 272 | part_dict = self.decompose_code(self.Regressor[part]( |
| 273 | feature[key][part]), self.param_list_dict[f'{part}_list']) |
| 274 | # if input is part data, skip feature fusion: share feature is the same as part feature |
| 275 | # then predict share parameters |
| 276 | feature[key][f'{key}_share'] = feature[key][key] |
| 277 | share_dict = self.decompose_code( |
| 278 | self.Regressor[f'{part}_share']( |
| 279 | feature[key][f'{part}_share']), |
| 280 | self.param_list_dict[f'{part}_share_list']) |
| 281 | # compose parameters |
| 282 | param_dict[key] = {**share_dict, **part_dict} |
| 283 | |
| 284 | # for body image |
| 285 | if key == 'body': |
| 286 | fusion_weight = {} |
| 287 | f_body = feature['body']['body'] |
| 288 | # extract part feature |
| 289 | for part_name in ['head', 'left_hand', 'right_hand']: |
| 290 | feature['body'][f'{part_name}_share'] = self.Extractor[f'{part_name}_share']( |
| 291 | f_body) |
| 292 | |
| 293 | # -- check if part crops are given, if not, crop parts by coarse body estimation |
| 294 | if 'head_image' not in data[key].keys() \ |
| 295 | or 'left_hand_image' not in data[key].keys() \ |
| 296 | or 'right_hand_image' not in data[key].keys(): |
| 297 | # - run without fusion to get coarse estimation, for cropping parts |