| 20 | |
| 21 | |
| 22 | class BaseDataset(torch.utils.data.Dataset): |
| 23 | def __init__(self, odgt, opt, **kwargs): |
| 24 | # parse options |
| 25 | self.imgSizes = opt.imgSizes |
| 26 | self.imgMaxSize = opt.imgMaxSize |
| 27 | # max down sampling rate of network to avoid rounding during conv or pooling |
| 28 | self.padding_constant = opt.padding_constant |
| 29 | |
| 30 | # parse the input list |
| 31 | self.parse_input_list(odgt, **kwargs) |
| 32 | |
| 33 | # mean and std |
| 34 | self.normalize = transforms.Normalize( |
| 35 | mean=[0.485, 0.456, 0.406], |
| 36 | std=[0.229, 0.224, 0.225]) |
| 37 | |
| 38 | def parse_input_list(self, odgt, max_sample=-1, start_idx=-1, end_idx=-1): |
| 39 | if isinstance(odgt, list): |
| 40 | self.list_sample = odgt |
| 41 | elif isinstance(odgt, str): |
| 42 | self.list_sample = [json.loads(x.rstrip()) for x in open(odgt, 'r')] |
| 43 | |
| 44 | if max_sample > 0: |
| 45 | self.list_sample = self.list_sample[0:max_sample] |
| 46 | if start_idx >= 0 and end_idx >= 0: # divide file list |
| 47 | self.list_sample = self.list_sample[start_idx:end_idx] |
| 48 | |
| 49 | self.num_sample = len(self.list_sample) |
| 50 | assert self.num_sample > 0 |
| 51 | print('# samples: {}'.format(self.num_sample)) |
| 52 | |
| 53 | def img_transform(self, img): |
| 54 | # 0-255 to 0-1 |
| 55 | img = np.float32(np.array(img)) / 255. |
| 56 | img = img.transpose((2, 0, 1)) |
| 57 | img = self.normalize(torch.from_numpy(img.copy())) |
| 58 | return img |
| 59 | |
| 60 | def segm_transform(self, segm): |
| 61 | # to tensor, -1 to 149 |
| 62 | segm = torch.from_numpy(np.array(segm)).long() - 1 |
| 63 | return segm |
| 64 | |
| 65 | # Round x to the nearest multiple of p and x' >= x |
| 66 | def round2nearest_multiple(self, x, p): |
| 67 | return ((x - 1) // p + 1) * p |
| 68 | |
| 69 | |
| 70 | class TrainDataset(BaseDataset): |
nothing calls this directly
no outgoing calls
no test coverage detected