(self, idx)
| 467 | return len(self.datalist) |
| 468 | # 19493 |
| 469 | def __getitem__(self, idx): |
| 470 | # rank = self.rank |
| 471 | # local_rank = rank % torch.cuda.device_count() |
| 472 | # with open(f'index_log_{rank}.txt', 'a') as f: |
| 473 | # f.write(f'{rank}-{local_rank}-{idx}\n') |
| 474 | try: |
| 475 | data = copy.deepcopy(self.datalist[idx]) |
| 476 | except Exception as e: |
| 477 | print(f'[{self.__class__.__name__}] Error loading data {idx}') |
| 478 | print(e) |
| 479 | exit(0) |
| 480 | # data/datasets/coco_2017/train2017/000000029582.jpg' 45680 |
| 481 | img_path, img_shape, bbox = \ |
| 482 | data['img_path'], data['img_shape'], data['bbox'] |
| 483 | as_smplx = data['as_smplx'] |
| 484 | gender = data['gender'].copy() |
| 485 | for gender_str, gender_num in { |
| 486 | 'neutral': -1, 'male': 0, 'female': 1}.items(): |
| 487 | gender[gender==gender_str]=gender_num |
| 488 | gender = gender.astype(int) |
| 489 | |
| 490 | img_whole_bbox = np.array([0, 0, img_shape[1], img_shape[0]]) |
| 491 | img = load_img(img_path, order='BGR') |
| 492 | |
| 493 | num_person = len(data['bbox']) |
| 494 | data_name = self.__class__.__name__ |
| 495 | try: |
| 496 | # dist.barrier() |
| 497 | img, img2bb_trans, bb2img_trans, rot, do_flip = \ |
| 498 | augmentation_instance_sample(img, img_whole_bbox, self.data_split, data, data_name) |
| 499 | except Exception as e: |
| 500 | rank = self.rank |
| 501 | local_rank = rank % torch.cuda.device_count() |
| 502 | with open(f'index_log_{rank}.txt', 'a') as f: |
| 503 | f.write(f'{rank}-{local_rank}-{idx}\n') |
| 504 | f.write(f'[{self.__class__.__name__}] Error loading data {idx}\n') |
| 505 | f.write(f'Error in augmentation_instance_sample for {img_path}\n') |
| 506 | # print(f'[{self.__class__.__name__}] Error loading data {idx}') |
| 507 | # print(f'Error in augmentation_instance_sample for {img_path}') |
| 508 | raise e |
| 509 | cropped_img_shape = img.shape[:2] |
| 510 | |
| 511 | if self.data_split == 'train': |
| 512 | joint_cam = data['joint_cam'] # num, 137,4 |
| 513 | if joint_cam is not None: |
| 514 | dummy_cord = False |
| 515 | joint_cam[:,:,:3] = \ |
| 516 | joint_cam[:,:,:3] - joint_cam[:, self.joint_set['root_joint_idx'], None, :3] # root-relative |
| 517 | else: |
| 518 | # dummy cord as joint_cam |
| 519 | dummy_cord = True |
| 520 | joint_cam = np.zeros( |
| 521 | (num_person, self.joint_set['joint_num'], 4), |
| 522 | dtype=np.float32) |
| 523 | |
| 524 | joint_img = data['joint_img'] |
| 525 | # do rotation on keypoints |
| 526 | joint_img_aug, joint_cam_wo_ra, joint_cam_ra, joint_trunc = \ |
nothing calls this directly
no test coverage detected