| 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) |