()
| 17 | import tza |
| 18 | |
| 19 | def main(): |
| 20 | # Parse the command line arguments |
| 21 | cfg = parse_args(description='Preprocesses training and validation datasets.') |
| 22 | main_feature = get_main_feature(cfg.features) |
| 23 | aux_features = get_aux_features(cfg.features) |
| 24 | all_channels = get_dataset_channels(cfg.features) |
| 25 | num_main_channels = len(get_dataset_channels(main_feature)) |
| 26 | |
| 27 | # Initialize the PyTorch device |
| 28 | device = init_device(cfg) |
| 29 | |
| 30 | # Initialize the transfer function |
| 31 | transfer = get_transfer_function(cfg) |
| 32 | |
| 33 | # Initialize auxiliary feature inference |
| 34 | aux_infers = {} |
| 35 | for aux_result in set(cfg.aux_results): |
| 36 | aux_infer = Infer(cfg, device, aux_result, is_aux=True) |
| 37 | if (aux_infer.main_feature not in aux_features) or aux_infer.aux_features: |
| 38 | error(f'result {aux_result} does not correspond to an auxiliary feature') |
| 39 | aux_infers[aux_infer.main_feature] = aux_infer |
| 40 | |
| 41 | # Determine the input and target features |
| 42 | if cfg.clean_aux: |
| 43 | input_features = [main_feature] |
| 44 | target_features = cfg.features |
| 45 | else: |
| 46 | input_features = cfg.features |
| 47 | target_features = [main_feature] |
| 48 | |
| 49 | # Returns a preprocessed image (also changes the original image!) |
| 50 | def preprocess_image(image, exposure, prefilter=False): |
| 51 | # Apply the transfer function to the main feature |
| 52 | color = image[..., 0:num_main_channels] |
| 53 | color = torch.from_numpy(color).to(device) |
| 54 | if main_feature == 'hdr': |
| 55 | color *= exposure |
| 56 | color = transfer.forward(color) |
| 57 | color = torch.clamp(color, max=1.) |
| 58 | color = color.cpu().numpy() |
| 59 | image[..., 0:num_main_channels] = color |
| 60 | |
| 61 | # Prefilter the auxiliary features |
| 62 | if prefilter: |
| 63 | for aux_feature, aux_infer in aux_infers.items(): |
| 64 | aux_channels = get_dataset_channels(aux_feature) |
| 65 | aux_channel_indices = get_channel_indices(aux_channels, all_channels) |
| 66 | aux = image[..., aux_channel_indices] |
| 67 | aux = image_to_tensor(aux, batch=True).to(device) |
| 68 | aux = aux_infer(aux) |
| 69 | aux = tensor_to_image(aux) |
| 70 | image[..., aux_channel_indices] = aux |
| 71 | |
| 72 | # Convert to FP16 |
| 73 | return np.nan_to_num(image.astype(np.float16)) |
| 74 | |
| 75 | # Preprocesses a group of input and target images at different SPPs |
| 76 | def preprocess_sample_group(input_dir, output_tza, input_names, target_name): |
no test coverage detected