(self, index)
| 32 | self.extra_info = [] |
| 33 | |
| 34 | def __getitem__(self, index): |
| 35 | |
| 36 | if self.is_test: |
| 37 | img1 = frame_utils.read_gen(self.image_list[index][0]) |
| 38 | img2 = frame_utils.read_gen(self.image_list[index][1]) |
| 39 | img1 = np.array(img1).astype(np.uint8)[..., :3] |
| 40 | img2 = np.array(img2).astype(np.uint8)[..., :3] |
| 41 | img1 = torch.from_numpy(img1).permute(2, 0, 1).float() |
| 42 | img2 = torch.from_numpy(img2).permute(2, 0, 1).float() |
| 43 | return img1, img2, self.extra_info[index] |
| 44 | |
| 45 | if not self.init_seed: |
| 46 | worker_info = torch.utils.data.get_worker_info() |
| 47 | if worker_info is not None: |
| 48 | torch.manual_seed(worker_info.id) |
| 49 | np.random.seed(worker_info.id) |
| 50 | random.seed(worker_info.id) |
| 51 | self.init_seed = True |
| 52 | |
| 53 | index = index % len(self.image_list) |
| 54 | valid = None |
| 55 | if self.sparse: |
| 56 | flow, valid = frame_utils.readFlowKITTI(self.flow_list[index]) |
| 57 | else: |
| 58 | flow = frame_utils.read_gen(self.flow_list[index]) |
| 59 | |
| 60 | img1 = frame_utils.read_gen(self.image_list[index][0]) |
| 61 | img2 = frame_utils.read_gen(self.image_list[index][1]) |
| 62 | |
| 63 | flow = np.array(flow).astype(np.float32) |
| 64 | img1 = np.array(img1).astype(np.uint8) |
| 65 | img2 = np.array(img2).astype(np.uint8) |
| 66 | |
| 67 | # grayscale images |
| 68 | if len(img1.shape) == 2: |
| 69 | img1 = np.tile(img1[...,None], (1, 1, 3)) |
| 70 | img2 = np.tile(img2[...,None], (1, 1, 3)) |
| 71 | else: |
| 72 | img1 = img1[..., :3] |
| 73 | img2 = img2[..., :3] |
| 74 | |
| 75 | if self.augmentor is not None: |
| 76 | if self.sparse: |
| 77 | img1, img2, flow, valid = self.augmentor(img1, img2, flow, valid) |
| 78 | else: |
| 79 | img1, img2, flow = self.augmentor(img1, img2, flow) |
| 80 | |
| 81 | img1 = torch.from_numpy(img1).permute(2, 0, 1).float() |
| 82 | img2 = torch.from_numpy(img2).permute(2, 0, 1).float() |
| 83 | flow = torch.from_numpy(flow).permute(2, 0, 1).float() |
| 84 | |
| 85 | if valid is not None: |
| 86 | valid = torch.from_numpy(valid) |
| 87 | else: |
| 88 | valid = (flow[0].abs() < 1000) & (flow[1].abs() < 1000) |
| 89 | |
| 90 | return img1, img2, flow, valid.float() |
| 91 |
nothing calls this directly
no outgoing calls
no test coverage detected