| """ |
| Defines a Loader class to load data from a file or file wildcard |
| """ |
|
|
| import argparse |
|
|
| import h5py |
| import torch |
| import numpy as np |
| import glob |
| from typing import Tuple |
|
|
|
|
| class Loader: |
| """ |
| Data loader class |
| """ |
|
|
| def __init__(self, **kwargs): |
| parser = Loader.add_argparse_args() |
| for action in parser._actions: |
| if action.dest in kwargs: |
| action.default = kwargs[action.dest] |
| args = parser.parse_args([]) |
| self.__dict__.update(vars(args)) |
|
|
| if type(self.label_vars) is str: |
| self.label_vars = [self.label_vars] |
|
|
| @staticmethod |
| def add_argparse_args(parent_parser=None): |
| """ |
| Add argeparse argument for the data loader |
| """ |
| parser = argparse.ArgumentParser( |
| prog='Loader', |
| usage=Loader.__doc__, |
| parents=[parent_parser] if parent_parser is not None else [], |
| add_help=False) |
|
|
| parser.add_argument('--input_var', default='p_f5.0_o0', help='Variable name for the label data') |
| parser.add_argument('--label_vars', nargs='*', default='c0', help='Variable name(s) for the label data') |
|
|
| parser.add_argument('--inputs_crop', type=int, default=[0, 1, 32, 96, 42, 2090], nargs='*', |
| help='Crop input data on load [layer_min layer_max x_min x_max y_min y_max]') |
| parser.add_argument('--labels_crop', type=int, default=[322, 830, 60, 1076], nargs='*', help='Crop label data on load [x_min x_max y_min y_max]') |
| parser.add_argument('--labels_resize', type=float, default=256.0 / 1016.0, help='scaling factor for labels image') |
|
|
| parser.add_argument('--data_scale', type=float, default=1.0, help='Data scaling factor') |
| parser.add_argument('--data_gain', type=float, default=1.8, help='Data gain factor in dB/20 at farthest point in data.') |
|
|
| return parser |
|
|
| def load_data(self, test_file_pattern: str, train_file_pattern: str = None) -> Tuple[torch.FloatTensor, torch.FloatTensor]: |
| """Loads training/testing data from file(s) |
| |
| Arguments: |
| test_file_pattern {str} -- testing dataset(s) pattern |
| train_file_pattern {str} -- training dataset(s) pattern |
| |
| Returns: |
| (test_inputs, test_labels, train_inputs, train_labels) -- None for values that are not loaded |
| """ |
|
|
| test_inputs, test_labels = self._load_data_files(test_file_pattern) |
| train_inputs, train_labels = self._load_data_files(train_file_pattern) |
|
|
| if train_file_pattern is not None and train_inputs is None: |
| raise ValueError('Failed to load train set') |
|
|
| if test_file_pattern is not None and test_inputs is None: |
| raise ValueError('Failed to load train set') |
|
|
| return test_inputs, test_labels, train_inputs, train_labels |
|
|
| def _load_data_files(self, file_pattern: str) -> Tuple[torch.FloatTensor, torch.FloatTensor]: |
| """ Perform actual data loading |
| |
| Args: |
| file_pattern: file name pattern |
| |
| Returns: |
| inputs and labels tensors |
| """ |
|
|
| inputs, labels = None, None |
|
|
| if file_pattern is None: |
| return inputs, labels |
|
|
| files = glob.glob(file_pattern) |
|
|
| if len(files) == 0: |
| raise ValueError(f'{file_pattern=} comes up empty') |
|
|
| |
| with h5py.File(files[0], 'r') as f: |
| if self.input_var not in f: |
| raise ValueError(f'input data key not in file: {self.input_var=}') |
|
|
| shape = list(f[self.input_var].shape) |
| if self.inputs_crop is not None: |
| for i in range(len(self.inputs_crop) // 2): |
| shape[-i - 1] = self.inputs_crop[-i * 2 - 1] - self.inputs_crop[-i * 2 - 2] |
|
|
| shape[0] *= len(files) |
|
|
| inputs = np.empty(shape, np.single) |
|
|
| if len(self.label_vars): |
| if not all([v in f for v in self.label_vars]): |
| raise ValueError(f'labels data key(s) not in file: {self.label_vars=}') |
|
|
| shape = list(f[self.label_vars[0]].shape) |
| shape[1] *= len(self.label_vars) |
| if self.labels_crop is not None: |
| for i in range(len(self.labels_crop) // 2): |
| shape[-i - 1] = self.labels_crop[-i * 2 - 1] - self.labels_crop[-i * 2 - 2] |
|
|
| shape[-1] = int(shape[-1] * self.labels_resize) |
| shape[-2] = int(shape[-2] * self.labels_resize) |
| shape[0] *= len(files) |
|
|
| labels = np.empty(shape, np.single) |
|
|
| |
| pos = 0 |
| for file in files: |
| with h5py.File(files[0], 'r') as f: |
| tmp_inputs = np.array(f[self.input_var]) |
|
|
| if self.inputs_crop is not None: |
| slc = [slice(None)] * 4 |
| for i in range(len(self.inputs_crop) // 2): |
| slc[-i - 1] = slice(self.inputs_crop[-i * 2 - 2], self.inputs_crop[-i * 2 - 1]) |
| tmp_inputs = tmp_inputs[tuple(slc)] |
|
|
| inputs[pos:pos + tmp_inputs.shape[0], ...] = tmp_inputs |
|
|
| if len(self.label_vars): |
| tmp_labels = [] |
| for v in self.label_vars: |
| tmp_labels.append(np.array(f[v])) |
| tmp_labels = np.concatenate(tmp_labels, axis=1) |
|
|
| if self.labels_crop is not None and self.labels_crop: |
| slc = [slice(None)] * 4 |
| for i in range(len(self.labels_crop) // 2): |
| slc[-i - 1] = slice(self.labels_crop[-i * 2 - 2], self.labels_crop[-i * 2 - 1]) |
| tmp_labels = tmp_labels[tuple(slc)] |
|
|
| if self.labels_resize != 1.0: |
| tmp_labels = torch.nn.Upsample(scale_factor=self.labels_resize, mode='nearest')(torch.from_numpy(tmp_labels)).numpy() |
|
|
| labels[pos:pos + tmp_labels.shape[0], ...] = tmp_labels |
|
|
| pos += tmp_inputs.shape[0] |
|
|
| inputs = inputs[:pos, ...] |
| if len(self.label_vars): |
| labels = labels[:pos, ...] |
|
|
| if self.data_scale != 1.0: |
| inputs *= self.data_scale |
|
|
| if self.data_gain != 0.0: |
| gain = 10.0 ** np.linspace(0, self.data_gain, inputs.shape[-1], np.single).reshape((1, 1, 1, -1)) |
| inputs *= gain |
|
|
| |
| |
| inputs = torch.from_numpy(inputs.copy()) |
| if len(self.label_vars): |
| labels = torch.from_numpy(labels) |
|
|
| return inputs, labels |
|
|
