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