| 16 | |
| 17 | |
| 18 | class FlowDataset(data.Dataset): |
| 19 | def __init__(self, aug_params=None, sparse=False): |
| 20 | self.augmentor = None |
| 21 | self.sparse = sparse |
| 22 | if aug_params is not None: |
| 23 | if sparse: |
| 24 | self.augmentor = SparseFlowAugmentor(**aug_params) |
| 25 | else: |
| 26 | self.augmentor = FlowAugmentor(**aug_params) |
| 27 | |
| 28 | self.is_test = False |
| 29 | self.init_seed = False |
| 30 | self.flow_list = [] |
| 31 | self.image_list = [] |
| 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: |
nothing calls this directly
no outgoing calls
no test coverage detected