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