(self, index)
| 64 | return len(self.image_total_files) |
| 65 | |
| 66 | def __getitem__(self, index): |
| 67 | # Train set |
| 68 | if self.set == 'train': |
| 69 | loss = nn.L1Loss() |
| 70 | while True: |
| 71 | if index in self.blacklist: |
| 72 | index=random.randint(0,self.__len__()-1) |
| 73 | continue |
| 74 | |
| 75 | test_rgb_path = self.image_total_files[index] |
| 76 | file_idx = test_rgb_path.split('/')[-1].split('.')[0] # ~~ of ~~.png |
| 77 | |
| 78 | ref_pose_find_path = test_rgb_path.replace(f'rgb/{file_idx}.png',f't0/idx/{file_idx}.txt') |
| 79 | f = open(ref_pose_find_path,'r',encoding='utf8') |
| 80 | ref_pose_idx = int(f.readlines()[0]) |
| 81 | g2o_path = test_rgb_path.replace('/Query/Query_Seq_Train','/Reference/Ref_Seq_Train').replace(f'rgb/{file_idx}.png',f'raw/poses.g2o') |
| 82 | with open(g2o_path,'r',encoding = 'utf8') as f2: |
| 83 | while True: |
| 84 | line = f2.readline() |
| 85 | try: |
| 86 | if line.split()[0] == 'VERTEX_SE3:QUAT' and int(line.split()[1]) == ref_pose_idx: |
| 87 | ref_pose = line.split()[2:] |
| 88 | except: |
| 89 | break |
| 90 | ref_pose = torch.from_numpy(np.array(ref_pose).astype(float)) |
| 91 | change_pose_path = test_rgb_path.replace(f'rgb/{file_idx}.png',f'pose/{file_idx}.txt') |
| 92 | with open(change_pose_path,'r',encoding='utf8') as f3: |
| 93 | change_pose = f3.readline().split() |
| 94 | change_pose = torch.from_numpy(np.array(change_pose).astype(float)) |
| 95 | |
| 96 | distance = loss(ref_pose.cuda(),change_pose.cuda()) |
| 97 | if distance.item()<0.5: |
| 98 | break |
| 99 | else: |
| 100 | self.blacklist.append(index) |
| 101 | index=random.randint(0,self.__len__()-1) |
| 102 | # Test set |
| 103 | else: |
| 104 | test_rgb_path = self.image_total_files[index] |
| 105 | |
| 106 | # Get File Paths |
| 107 | test_depth_path = test_rgb_path.replace('rgb', 'depth') |
| 108 | ref_rgb_path = test_rgb_path.replace('rgb', 't0/rgb') |
| 109 | ref_depth_path = test_rgb_path.replace('rgb', 't0/depth') |
| 110 | change_segmentation_path = test_rgb_path.replace('rgb', 'change_segmentation') |
| 111 | |
| 112 | name = '_'.join(test_rgb_path.split('/')[-5:]) |
| 113 | |
| 114 | #### Color Transform #### |
| 115 | test_rgb = Image.open(test_rgb_path) |
| 116 | ref_rgb = Image.open(ref_rgb_path) |
| 117 | test_rgb = test_rgb.resize(self.crop_size, Image.BICUBIC) |
| 118 | ref_rgb = ref_rgb.resize(self.crop_size, Image.BICUBIC) |
| 119 | |
| 120 | # RGB, Color Transform for train set |
| 121 | if self.set == 'train': |
| 122 | test_rgb = self.color_transform(test_rgb) |
| 123 | ref_rgb = self.color_transform(ref_rgb) |
nothing calls this directly
no test coverage detected