| 32 | |
| 33 | |
| 34 | class iBims(Dataset): |
| 35 | def __init__(self, config): |
| 36 | root_folder = config.ibims_root |
| 37 | with open(os.path.join(root_folder, "imagelist.txt"), 'r') as f: |
| 38 | imglist = f.read().split() |
| 39 | |
| 40 | samples = [] |
| 41 | for basename in imglist: |
| 42 | img_path = os.path.join(root_folder, 'rgb', basename + ".png") |
| 43 | depth_path = os.path.join(root_folder, 'depth', basename + ".png") |
| 44 | valid_mask_path = os.path.join( |
| 45 | root_folder, 'mask_invalid', basename+".png") |
| 46 | transp_mask_path = os.path.join( |
| 47 | root_folder, 'mask_transp', basename+".png") |
| 48 | |
| 49 | samples.append( |
| 50 | (img_path, depth_path, valid_mask_path, transp_mask_path)) |
| 51 | |
| 52 | self.samples = samples |
| 53 | # self.normalize = T.Normalize( |
| 54 | # mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) |
| 55 | self.normalize = lambda x : x |
| 56 | |
| 57 | def __getitem__(self, idx): |
| 58 | img_path, depth_path, valid_mask_path, transp_mask_path = self.samples[idx] |
| 59 | |
| 60 | img = np.asarray(Image.open(img_path), dtype=np.float32) / 255.0 |
| 61 | depth = np.asarray(Image.open(depth_path), |
| 62 | dtype=np.uint16).astype('float')*50.0/65535 |
| 63 | |
| 64 | mask_valid = np.asarray(Image.open(valid_mask_path)) |
| 65 | mask_transp = np.asarray(Image.open(transp_mask_path)) |
| 66 | |
| 67 | # depth = depth * mask_valid * mask_transp |
| 68 | depth = np.where(mask_valid * mask_transp, depth, -1) |
| 69 | |
| 70 | img = torch.from_numpy(img).permute(2, 0, 1) |
| 71 | img = self.normalize(img) |
| 72 | depth = torch.from_numpy(depth).unsqueeze(0) |
| 73 | return dict(image=img, depth=depth, image_path=img_path, depth_path=depth_path, dataset='ibims') |
| 74 | |
| 75 | def __len__(self): |
| 76 | return len(self.samples) |
| 77 | |
| 78 | |
| 79 | def get_ibims_loader(config, batch_size=1, **kwargs): |