| 91 | self.transform = ToTensor() |
| 92 | |
| 93 | def __getitem__(self, idx): |
| 94 | image_path = self.image_files[idx] |
| 95 | depth_path = self.depth_files[idx] |
| 96 | |
| 97 | image = Image.open(image_path) |
| 98 | depth = Image.open(depth_path) |
| 99 | depth = cv2.imread(depth_path, cv2.IMREAD_ANYCOLOR | |
| 100 | cv2.IMREAD_ANYDEPTH) |
| 101 | print("dpeth min max", depth.min(), depth.max()) |
| 102 | |
| 103 | # print(np.shape(image)) |
| 104 | # print(np.shape(depth)) |
| 105 | |
| 106 | # depth[depth > 8] = -1 |
| 107 | |
| 108 | if self.do_kb_crop and False: |
| 109 | height = image.height |
| 110 | width = image.width |
| 111 | top_margin = int(height - 352) |
| 112 | left_margin = int((width - 1216) / 2) |
| 113 | depth = depth.crop( |
| 114 | (left_margin, top_margin, left_margin + 1216, top_margin + 352)) |
| 115 | image = image.crop( |
| 116 | (left_margin, top_margin, left_margin + 1216, top_margin + 352)) |
| 117 | # uv = uv[:, top_margin:top_margin + 352, left_margin:left_margin + 1216] |
| 118 | |
| 119 | image = np.asarray(image, dtype=np.float32) / 255.0 |
| 120 | # depth = np.asarray(depth, dtype=np.uint16) /1. |
| 121 | depth = depth[..., None] |
| 122 | sample = dict(image=image, depth=depth) |
| 123 | |
| 124 | # return sample |
| 125 | sample = self.transform(sample) |
| 126 | |
| 127 | if idx == 0: |
| 128 | print(sample["image"].shape) |
| 129 | |
| 130 | return sample |
| 131 | |
| 132 | def __len__(self): |
| 133 | return len(self.image_files) |