| from numpy import sqrt |
| import torch |
| import pytorch_lightning as pl |
| import wandb |
| from argparse import ArgumentParser |
| from kornia.contrib import extract_tensor_patches |
| from torch.optim import Adam, SGD |
| from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts |
| from torchvision.utils import make_grid |
| from src.datasources import S2_ALL_12BANDS |
| from src.losses import mse_loss, mae_loss, ssim_loss, tv_loss, Shift |
| from src.modules import BicubicUpscaledBaseline |
| import pandas as pd |
|
|
|
|
| def detach_dict(dict): |
| """Detaches the tensors in a dictionary. |
| |
| Parameters |
| ---------- |
| dict : dict of Tensor |
| Dictionary containing the tensors to detach. |
| |
| Returns |
| ------- |
| dict of Tensor |
| Dictionary containing the detached tensors. |
| """ |
| for tensor_key in dict: |
| if dict[tensor_key].requires_grad: |
| dict[tensor_key] = dict[tensor_key].detach() |
| return dict |
|
|
|
|
| class LitModel(pl.LightningModule): |
|
|
| """Lightning Module wrapper. |
| Implements the forward, loss, train/val/test step functions and logging. |
| Parses and adds default model arguments. |
| """ |
|
|
| def __init__(self, backbone, **kws) -> None: |
| """Initialises the LightningModule with the backbone and the default arguments. |
| The hyperparameters are fetchd from the hparams attribute. |
| |
| Parameters |
| ---------- |
| backbone : nn.Module |
| Backbone model, as generated by train.py. |
| """ |
| super().__init__() |
| self.save_hyperparameters() |
| self.backbone = backbone |
| self.transform = self.hparams.transform |
|
|
| output_channels = [channel for channel in range(self.backbone.out_channels)] |
|
|
| self.baseline = BicubicUpscaledBaseline( |
| input_size=self.hparams.input_size, |
| output_size=self.backbone.output_size, |
| output_channels=output_channels, |
| interpolation="bicubic", |
| chip_size=self.hparams.chip_size, |
| ) |
|
|
| self.check_loss_functions_weights_are_valid() |
| self.create_shifters() |
| self.create_metrics() |
| self.enable_benchmark_logging() |
|
|
| def enable_benchmark_logging(self): |
| if self.hparams.benchmark: |
| self.logging = {} |
| self.logging['current_idx'] = 1 |
| self.logging['log'] = pd.DataFrame() |
| self.log_path = '../../log' |
|
|
| def check_loss_functions_weights_are_valid(self): |
| """Checks that the loss functions weights all sum to 1.""" |
| assert self.hparams.w_mse + self.hparams.w_mae + self.hparams.w_ssim == 1 |
|
|
| def create_shifters(self): |
| """Creates a shifter for the training and validation steps, and a shifter for the baseline.""" |
|
|
| self.shifter = Shift( |
| shift_by_px=self.hparams.shift_px, |
| mode=self.hparams.shift_mode, |
| step=self.hparams.shift_step, |
| use_cache=False, |
| ) |
|
|
| self.baseline_shifter = Shift( |
| shift_by_px=self.hparams.shift_px, |
| mode=self.hparams.shift_mode, |
| step=self.hparams.shift_step, |
| use_cache=False, |
| ) |
|
|
| def create_metrics(self): |
| """Creates metrics that use shifters when computing the losses (MSE, MAE, SSIM. TV) for the train, validation and test sets.""" |
| self.metrics_baseline = {} |
| self.metrics_hr = {} |
| |
| self.train_metrics = self.metrics_factory("train", self.shifter) |
| self.val_metrics = self.metrics_factory("val", self.shifter) |
| self.test_metrics = self.metrics_factory("test", self.shifter) |
| self.baseline_train_metrics = self.metrics_factory( |
| "train", self.baseline_shifter |
| ) |
| self.baseline_val_metrics = self.metrics_factory("val", self.baseline_shifter) |
| self.baseline_test_metrics = self.metrics_factory("test", self.baseline_shifter) |
|
|
| |
|
|
| def metrics_factory(self, prefix, shifter): |
| """Creates a dictionary of metrics that use the shifter when computing the losses (MSE, MAE, SSIM. TV) for the train, validation and test sets. |
| |
| Parameters |
| ---------- |
| prefix : str |
| Prefix for the metrics (train/val/test). |
| shifter : losses.Shift |
| Shifter to use when computing the losses. |
| |
| Returns |
| ------- |
| dict |
| Dictionary of metrics. |
| """ |
| return { |
| f"{prefix}/MSE": lambda y_hat, y: shifter.registered_loss(mse_loss)( |
| y_hat, y |
| ), |
| f"{prefix}/MAE": lambda y_hat, y: shifter.registered_loss(mae_loss)( |
| y_hat, y |
| ), |
| f"{prefix}/SSIM_loss": lambda y_hat, y,: shifter.registered_loss(ssim_loss)( |
| y_hat, y, window_size=5 |
| ), |
| f"{prefix}/TV": lambda y_hat, y=None: tv_loss(y_hat), |
| } |
|
|
| def forward(self, x, y=None): |
| """Forward pass of the model. |
| |
| Parameters |
| ---------- |
| x : Tensor |
| Input tensor (LR) of dimensions (batch_size, channels, height, width). |
| y : Tensor, optional |
| Target tensor (HR) of dimensions (batch_size, channels, height, width). |
| |
| Returns |
| ------- |
| Tensor |
| Output tensor (y_hat) of dimensions (batch_size, channels, height, width). |
| """ |
| return self.backbone(x, y) |
|
|
| def validation_log(self, y, m, baseline_m, prefix): |
|
|
| log = pd.DataFrame( |
| columns=[ |
| "MSE", |
| "MAE", |
| "SSIM_loss", |
| "TV", |
| "PSNR", |
| "Baseline MSE", |
| "Baseline PSNR", |
| ] |
| ) |
| mse, mae, ssim, tv, psnr = ( |
| m[f"{prefix}/MSE"], |
| m[f"{prefix}/MAE"], |
| m[f"{prefix}/SSIM_loss"], |
| m[f"{prefix}/TV"], |
| m[f"{prefix}/PSNR"], |
| ) |
| baseline_mse, baseline_psnr = ( |
| baseline_m[f"{prefix}/MSE"], |
| baseline_m[f"{prefix}/PSNR_baseline"], |
| ) |
| baseline_m = list(baseline_m.values()) |
| for batch_item in range(y.shape[0]): |
| log.loc[self.logging["current_idx"]] = ( |
| mse[batch_item].detach().item(), |
| mae[batch_item].detach().item(), |
| ssim[batch_item].detach().item(), |
| tv[batch_item].detach().item(), |
| psnr[batch_item].detach().item(), |
| baseline_mse[batch_item].detach().item(), |
| baseline_psnr[batch_item].detach().item(), |
| ) |
| self.logging["current_idx"] += 1 |
| self.logging["log"] = pd.concat([self.logging["log"], log]) |
|
|
| def loss(self, batch, metrics, baseline_metrics, prefix): |
| """Computes the loss for a batch using the provided metrics. |
| |
| Parameters |
| ---------- |
| batch : dict |
| Dictionary with the batches for LR/HR/etc. |
| metrics : dict of Callable |
| Dictionary with the metrics to use. |
| baseline_metrics : dict of Callable |
| Dictionary with the metrics to use for the baseline. |
| prefix : str |
| Prefix for the metrics (train/val/test). |
| |
| Returns |
| ------- |
| dict |
| Dictionary with the losses. |
| """ |
|
|
| batch = self.transform(batch) |
| x, y = batch["lr"], batch["hr"] |
| |
| x = self.apply_masks_to_x(x, batch) |
|
|
| y_hat = self(x, y) |
| y, y_hat = y[:, 0], y_hat[:, 0] |
|
|
| y_hat = self.bias_adjust(y, y_hat) |
|
|
| self.initialise_shifters() |
|
|
| m, mse, mae, ssim, tv = self.compute_metrics(y, y_hat, metrics, prefix) |
| y_hat_base = self.baseline(x)[:, 0] |
| y_hat_base = self.bias_adjust(y, y_hat_base) |
| baseline_m = self.compute_baseline_metrics( |
| y, y_hat_base, m, baseline_metrics, prefix |
| ) |
|
|
|
|
| loss = ( |
| (self.hparams.w_mse * mse) |
| + (self.hparams.w_mae * mae) |
| + (self.hparams.w_tv * tv) |
| + (self.hparams.w_ssim * ssim) |
| ) |
| |
| if self.hparams.benchmark: |
| self.validation_log(y, detach_dict(m), baseline_m, prefix) |
|
|
| return { |
| "y_hat": y_hat, |
| "y_hat_base": y_hat_base, |
| "y": y, |
| "metrics": m, |
| "baseline_metrics": baseline_m, |
| "loss": loss, |
| } |
|
|
| def apply_masks_to_x(self, x, batch): |
| """Applies the input masks to the input tensor. |
| |
| Parameters |
| ---------- |
| x : torch.Tensor |
| Input tensor of dimensions (batch_size, channels, height, width). |
| batch : dict |
| Dictionary with the batches for LR/Masks/HR/etc. |
| |
| Returns |
| ------- |
| torch.Tensor |
| Input tensor with the mask channels concatenated. |
| """ |
| if self.hparams.use_masks: |
| masks = batch["lrc"] |
| x = torch.cat([x, masks], dim=-3) |
| return x |
|
|
| @staticmethod |
| def bias_adjust(y, y_hat): |
| """Applies the brightness correction to the output. |
| |
| Parameters |
| ---------- |
| y : torch.Tensor |
| Target tensor (HR) of dimensions (batch_size, channels, height, width). |
| y_hat : torch.Tensor |
| Output tensor (SR) of dimensions (batch_size, channels, height, width). |
| |
| Returns |
| ------- |
| torch.Tensor |
| Output tensor with the brightness correction applied. |
| """ |
| bias = (y - y_hat).mean(dim=(-1, -2), keepdim=True) |
| return y_hat + bias |
|
|
| def initialise_shifters(self): |
| """Initialises the shifter for the HR/SR pairs and the HR/LR Baseline pairs.""" |
| self.shifter.y, self.shifter.y_hat = None, None |
| self.baseline_shifter.y, self.baseline_shifter.y_hat = None, None |
|
|
| |
| |
| |
| if self.current_epoch % 50 == 0: |
| self.clear_shifter_cache() |
|
|
| def compute_metrics(self, y, y_hat, metrics, prefix): |
| """Computes the metrics for a batch (MSE, MAE, SSIM, TV, PSNR). |
| |
| Parameters |
| ---------- |
| y : torch.Tensor |
| Target tensor (HR) of dimensions (batch_size, channels, height, width). |
| y_hat : torch.Tensor |
| Output tensor (SR) of dimensions (batch_size, channels, height, width). |
| metrics : dict of Callable |
| Dictionary with the metric functions to use. |
| prefix : str |
| Prefix for the metrics (train/val/test). |
| |
| Returns |
| ------- |
| dict, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor |
| Dictionary with the metrics (MSE, MAE, SSIM, TV, PSNR). |
| The metrics themselves, the MSE, MAE, SSIM, TV. |
| """ |
| metrics = { |
| metrics_function: metrics[metrics_function](y_hat, y) |
| for metrics_function in metrics |
| } |
| mse = metrics[f"{prefix}/MSE"].clone() |
| mae = metrics[f"{prefix}/MAE"].clone() |
| ssim = metrics[f"{prefix}/SSIM_loss"].clone() |
| tv = metrics[f"{prefix}/TV"].clone() |
| metrics[f"{prefix}/PSNR"] = -10.0 * torch.log10(mse) |
| return metrics, mse, mae, ssim, tv |
|
|
| def compute_baseline_metrics( |
| self, y, y_hat_base, metrics, baseline_metrics, prefix |
| ): |
| """Computes the baseline metrics for a batch (MSE, MAE, SSIM, TV, PSNR_baseline, PSNR_ratio). |
| |
| Parameters |
| ---------- |
| y : torch.Tensor |
| Target tensor (HR) of dimensions (batch_size, channels, height, width). |
| y_hat_base : torch.Tensor |
| Baseline tensor (single upscaled LR) of dimensions (batch_size, channels, height, width). |
| metrics : dict of Callable |
| Dictionary with the metric functions to use. |
| baseline_metrics : dict of Callable |
| Dictionary with the baseline metric functions to use. |
| prefix : str |
| Prefix for the metrics (train/val/test). |
| |
| Returns |
| ------- |
| dict |
| Dictionary with the baseline metrics (MSE, MAE, SSIM, TV, PSNR_baseline, PSNR_ratio). |
| """ |
| baseline_m = {k: baseline_metrics[k](y_hat_base, y) for k in baseline_metrics} |
|
|
| baseline_psnr = -10.0 * torch.log10(mse_loss(y_hat_base, y)) |
| baseline_m[f"{prefix}/PSNR_baseline"] = baseline_psnr |
|
|
| baseline_m[f"{prefix}/PSNR_ratio"] = ( |
| baseline_m[f"{prefix}/PSNR_baseline"] / metrics[f"{prefix}/PSNR"] |
| ) |
| baseline_m = detach_dict(baseline_m) |
| return baseline_m |
|
|
| def clear_shifter_cache(self): |
| """Clears the cache of the output and baseline shifters.""" |
|
|
| self.shifter.shift_cache = {} |
| self.baseline_shifter.shift_cache = {} |
|
|
| def training_step(self, batch): |
| """Training step. |
| Calls the forward pass, computes the loss and metrics for the training batch. |
| Logs the reduced metrics. |
| |
| Parameters |
| ---------- |
| batch : dict |
| Dictionary with the batches for LR/HR/etc. |
| |
| Returns |
| ------- |
| dict |
| Dictionary with the training loss and metrics. |
| """ |
| loss_output = self.loss( |
| batch, self.train_metrics, self.baseline_train_metrics, prefix="train" |
| ) |
|
|
| metrics, metrics_reduced = self.unpack_and_reduce_metrics( |
| loss_output, prefix="train" |
| ) |
|
|
| self.log_dict(metrics_reduced, on_step=True, on_epoch=True) |
| return {"loss": metrics_reduced["train/loss"], "metrics": metrics} |
|
|
| def unpack_and_reduce_metrics(self, loss_output, prefix): |
| """Unpacks the metrics for a split (prefix) and reduces them (mean). |
| |
| Parameters |
| ---------- |
| loss_output : dict |
| Dictionary with the loss and metrics. |
| prefix : str |
| Prefix for the metrics (train/val/test). |
| |
| Returns |
| ------- |
| dict |
| Dictionary with the reduced (mean) metrics. |
| """ |
| metrics = loss_output["metrics"] |
| metrics_reduced = { |
| metric_name: metric.mean() for metric_name, metric in metrics.items() |
| } |
| metrics_reduced[f"{prefix}/loss"] = loss_output["loss"].mean() |
| return metrics, metrics_reduced |
|
|
| def validation_step(self, batch, batch_idx): |
| """Validation step. |
| Calls the forward pass, computes the loss and metrics for the validation batch. |
| Logs the reduced metrics. |
| |
| Parameters |
| ---------- |
| batch : dict |
| Dictionary with the batches for LR/HR/etc. |
| |
| Returns |
| ------- |
| torch.Tensor |
| Validation loss. |
| """ |
| loss_output = self.loss( |
| batch, self.val_metrics, self.baseline_val_metrics, prefix="val" |
| ) |
| _, metrics_reduced = self.unpack_and_reduce_metrics(loss_output, prefix="val") |
| self.log_dict(metrics_reduced) |
| return loss_output["loss"] |
|
|
| def validation_epoch_end(self, validation_step_output): |
| """Validation epoch end. |
| Concatenates the validation losses and logs them as a histogram to WandB. |
| |
| Parameters |
| ---------- |
| validation_step_output : dict |
| Dictionary with the validation losses. |
| """ |
| instance_loss = torch.cat(validation_step_output).to("cpu") |
| try: |
| self.logger.experiment.log( |
| { |
| "val/loss_hist": wandb.Histogram(instance_loss), |
| "global_step": self.global_step, |
| } |
| ) |
| except ValueError as e: |
| print(e) |
|
|
| def test_step(self, batch, batch_idx): |
| """Test step.+ |
| Calls the forward pass, computes the loss and metrics for the test batch. |
| Logs the reduced metrics. |
| |
| Parameters |
| ---------- |
| batch : dict |
| Dictionary with the batches for LR/HR/etc. |
| |
| Returns |
| ------- |
| torch.Tensor |
| Test loss. |
| """ |
| loss_output = self.loss( |
| batch, self.test_metrics, self.baseline_test_metrics, prefix="test" |
| ) |
|
|
| _, metrics_reduced = self.unpack_and_reduce_metrics(loss_output, prefix="test") |
|
|
| self.log_dict(metrics_reduced) |
| return loss_output["loss"] |
|
|
| def configure_optimizers(self): |
| """Configures the optimizers for the model based on the hyperparameters. |
| Uses the learning rate and weight decay parameters for the Adam optimizer. |
| Uses the learning rate, weight decay and momentum parameters for the SGD optimizer. |
| |
| Uses the cos_anneal_T_0 hyperparameter for the cosine annealing scheduler. |
| See: https://arxiv.org/abs/1608.03983 |
| |
| Returns |
| ------- |
| dict |
| Dictionary containing the optimizer, scheduler and monitor key (val/loss). |
| """ |
| if self.hparams.optimizer == "adam": |
| optimizer = Adam( |
| self.parameters(), |
| lr=self.hparams.learning_rate, |
| weight_decay=self.hparams.weight_decay, |
| ) |
|
|
| elif self.hparams.optimizer == "sgd": |
| optimizer = SGD( |
| self.parameters(), |
| lr=self.hparams.learning_rate, |
| weight_decay=self.hparams.weight_decay, |
| momentum=self.hparams.momentum, |
| ) |
| scheduler = CosineAnnealingWarmRestarts( |
| optimizer, T_0=self.hparams.cos_anneal_T_0 |
| ) |
|
|
| return { |
| "optimizer": optimizer, |
| "lr_scheduler": scheduler, |
| "monitor": "val/loss", |
| } |
|
|
| @staticmethod |
| def add_model_specific_args(parent_parser): |
| """Add model-specific arguments (and their default values) to the parser. |
| |
| Parameters |
| ---------- |
| parent_parser : argparse.ArgumentParser |
| argparse.ArgumentParser to add to. |
| |
| Returns |
| ------- |
| ArgumentParser |
| ArgumentParser with the model-specific arguments added. |
| """ |
| parser = ArgumentParser(parents=[parent_parser], add_help=False) |
| |
| parser.add_argument("--model", default="srcnn", type=str) |
| parser.add_argument("--hidden_channels", default=64, type=int) |
| parser.add_argument("--residual_layers", default=4, type=int) |
| parser.add_argument("--homography_fc_size", default=128, type=int) |
| parser.add_argument("--kernel_size", default=3, type=int) |
| parser.add_argument("--zoom_factor", default=2, type=int) |
| parser.add_argument("--sr_kernel_size", default=1, type=int) |
|
|
| |
| parser.add_argument("--optimizer", default="adam", type=str) |
| parser.add_argument("--learning_rate", default=5e-4, type=float) |
| parser.add_argument("--learning_rate_decay", default=0.97, type=float) |
| parser.add_argument("--learning_rate_patience", default=3, type=int) |
| parser.add_argument("--momentum", type=float, default=0.9) |
| parser.add_argument("--weight_decay", type=float, default=1e-4) |
| parser.add_argument("--cos_anneal_T_0", type=int, default=300) |
|
|
| |
| parser.add_argument("--input_size", nargs="+", default=(400, 400), type=int) |
| parser.add_argument("--output_size", nargs="+", default=(1000, 1000), type=int) |
| parser.add_argument("--chip_size", nargs="+", type=int, default=(50, 50)) |
| parser.add_argument("--chip_stride", nargs="+", type=int, default=None) |
| parser.add_argument("--revisits", default=8, type=int) |
| parser.add_argument( |
| "--use_masks", |
| action="store_true", |
| help="Include SCL cloud masks as extra input bands.", |
| ) |
|
|
| |
| parser.add_argument("--registration_kind", default=None, type=str) |
| parser.add_argument("--shift_px", default=2, type=int) |
| parser.add_argument("--shift_mode", default="lanczos", type=str) |
| parser.add_argument("--shift_step", default=0.5, type=float) |
| parser.add_argument( |
| "--use_reference_frame", |
| action="store_true", |
| help="Pair each revisit with a reference (median) frame.", |
| ) |
|
|
| |
| parser.add_argument("--w_mse", default=1.0, type=float) |
| parser.add_argument("--w_mae", default=0.0, type=float) |
| parser.add_argument("--w_ssim", default=0.0, type=float) |
| parser.add_argument("--w_tv", default=0.0, type=float) |
|
|
| return parser |
|
|
|
|
| class ImagePredictionLogger(pl.Callback): |
| """Logs the model losses, inputs and outputs to WandB. |
| |
| Parameters |
| ---------- |
| pl : pl.LightningModule |
| Lightning module containing the model. |
| """ |
|
|
| def __init__( |
| self, |
| train_dataloader, |
| val_dataloader, |
| test_dataloader, |
| log_every_n_epochs=1, |
| window_size=None, |
| ): |
| """Initialize the logger. |
| |
| Parameters |
| ---------- |
| train_dataloader : torch.utils.data.DataLoader |
| Training dataloader. |
| val_dataloader : torch.utils.data.DataLoader |
| Validation dataloader. |
| test_dataloader : torch.utils.data.DataLoader |
| Test dataloader. |
| log_every_n_epochs : int, optional |
| Log every n epochs. |
| window_size : tuple of int, optional |
| Size of the window to use for the patches. |
| """ |
| super().__init__() |
| self.train_dataloader = train_dataloader |
| self.val_dataloader = val_dataloader |
| self.test_dataloader = test_dataloader |
| self.log_every_n_epochs = log_every_n_epochs |
| self.window_size = window_size |
| self.baseline_psnr = None |
|
|
| def _on_epoch_end(self, prefix, batch, trainer, pl_module): |
| """Log the current epoch's metrics and images. |
| |
| Parameters |
| ---------- |
| prefix : str |
| The split prefix (train, val, test). |
| batch : dict |
| A batch of data. |
| trainer : pytorch_lightning.trainer.Trainer |
| The trainer. |
| pl_module : pytorch_lightning.LightningModule |
| The model. |
| """ |
| lr, hr = batch["lr"], batch["hr"] |
| batch = self.move_batch_to_device(batch, pl_module.device) |
|
|
| metrics, baseline_metrics = self.fetch_metrics_for_split(prefix, pl_module) |
|
|
| loss_output, instance_psnr, baseline_psnr = self.get_losses_from_metrics( |
| batch, metrics, baseline_metrics, pl_module, prefix |
| ) |
| sr, sr_base = self.get_output_and_baseline(loss_output) |
|
|
| window_size, window_size_hr = self.calculate_window_size(lr, hr, pl_module) |
|
|
| lr, hr, sr, sr_base = self.extract_patches_from_images( |
| sr, sr_base, lr, hr, window_size, window_size_hr |
| ) |
|
|
| instance_psnr, baseline_psnr = self.extract_psnr_values_for_patches( |
| sr, sr_base, lr, instance_psnr, baseline_psnr |
| ) |
|
|
| lr = self.convert_lr_to_RGB_only_for_logging(lr) |
| self.log_image_previews_to_wandb( |
| sr, |
| sr_base, |
| lr, |
| hr, |
| instance_psnr, |
| baseline_psnr, |
| pl_module, |
| prefix, |
| trainer, |
| ) |
|
|
| def move_batch_to_device(self, batch, device): |
| """Move all tensors from batch to the specified device. |
| |
| Parameters |
| ---------- |
| batch : dict |
| A batch of data. |
| device : torch.device |
| The device to move the tensors to. |
| |
| Returns |
| ------- |
| dict |
| The batch with all tensors moved to the specified device. |
| """ |
| return {key: value.to(device) for key, value in batch.items()} |
|
|
| def fetch_metrics_for_split(self, prefix, pl_module): |
| """Fetch the metrics for the split. |
| |
| Parameters |
| ---------- |
| prefix : str |
| The prefix for the metrics. |
| pl_module : pytorch_lightning.LightningModule |
| The model. |
| |
| Returns |
| ------- |
| dict, dict |
| The metrics and baseline metrics for the current batch. |
| """ |
| if prefix == "val": |
| metrics = pl_module.val_metrics |
| baseline_metrics = pl_module.baseline_val_metrics |
| elif prefix == "train": |
| metrics = pl_module.train_metrics |
| baseline_metrics = pl_module.baseline_train_metrics |
| elif prefix == "test": |
| metrics = pl_module.test_metrics |
| baseline_metrics = pl_module.baseline_test_metrics |
| return metrics, baseline_metrics |
|
|
| def get_losses_from_metrics( |
| self, batch, metrics, baseline_metrics, pl_module, prefix |
| ): |
| """Get the losses that will be logged from the metrics. |
| |
| Parameters |
| ---------- |
| |
| batch : dict |
| The batch. |
| metrics : dict |
| The metrics for the current batch. |
| baseline_metrics : dict |
| The baseline metrics for the current batch. |
| pl_module : pytorch_lightning.LightningModule |
| The model. |
| prefix : str |
| The prefix for the metrics. |
| |
| Returns |
| ------- |
| torch.Tensor, torch.Tensor, torch.Tensor |
| The loss output, the PSNR values for the SR and Baseline images in the batch. |
| """ |
| loss_output = pl_module.loss(batch, metrics, baseline_metrics, prefix=prefix) |
| instance_psnr = loss_output["metrics"][f"{prefix}/PSNR"].to("cpu") |
| baseline_psnr = loss_output["baseline_metrics"][f"{prefix}/PSNR_baseline"].to( |
| "cpu" |
| ) |
|
|
| return loss_output, instance_psnr, baseline_psnr |
|
|
| def get_output_and_baseline(self, loss_output): |
| """Get the output and baseline images from the loss output. |
| |
| Parameters |
| ---------- |
| loss_output : dict |
| The loss output. |
| |
| Returns |
| ------- |
| torch.Tensor, torch.Tensor |
| The output and baseline images. |
| """ |
| sr = loss_output["y_hat"].to("cpu")[:, None] |
| sr_base = loss_output["y_hat_base"].to("cpu")[:, None] |
| return sr, sr_base |
|
|
| def calculate_window_size(self, lr, hr, pl_module): |
| """Calculate the window size for the patches. |
| |
| Parameters |
| ---------- |
| lr : torch.Tensor |
| The LR image. |
| hr : torch.Tensor |
| The HR image. |
| pl_module : pytorch_lightning.LightningModule |
| The model. |
| |
| Returns |
| ------- |
| tuple of int, tuple of int |
| The window size for lr and hr. |
| """ |
| _, _, _, lr_height, lr_width = lr.shape |
| hr_height, hr_width = hr.shape[-2:] |
| ratio_hr_lr_height, ratio_hr_lr_width = ( |
| hr_height / lr_height, |
| hr_width / lr_width, |
| ) |
| window_size = self.window_size or pl_module.hparams.input_size |
| assert window_size <= tuple(pl_module.hparams.chip_size), "window_size too big" |
| window_size_hr = ( |
| round(ratio_hr_lr_height * window_size[0]), |
| round(ratio_hr_lr_width * window_size[0]), |
| ) |
|
|
| return window_size, window_size_hr |
|
|
| def extract_patches_from_images( |
| self, sr, sr_base, lr, hr, window_size, window_size_hr |
| ): |
| """Extract the patches from the images. |
| |
| Parameters |
| ---------- |
| sr : torch.Tensor |
| The SR image. |
| sr_base : torch.Tensor |
| The baseline SR image. |
| hr : torch.Tensor |
| The HR image. |
| window_size : tuple of int |
| The window size. |
| window_size_hr : tuple of int |
| The window size for the HR image. |
| |
| Returns |
| ------- |
| torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor |
| The extracted SR, SR_base, LR, HR images. |
| """ |
| lr = self.batch_patch_extractor( |
| lr, window_size |
| ) |
| hr = self.batch_patch_extractor(hr, window_size_hr) |
| sr = self.batch_patch_extractor(sr, window_size_hr) |
| sr_base = self.batch_patch_extractor(sr_base, window_size_hr) |
| return lr, hr, sr, sr_base |
|
|
| @staticmethod |
| def batch_patch_extractor(x, window_size): |
| """Extract patches from a batch of images. |
| |
| Parameters |
| ---------- |
| x : Tensor |
| A batch of images (batch, tensors/revisits, channels, height, width). |
| window_size : tuple of int |
| The size of the patch to extract. |
| |
| Returns |
| ------- |
| Tensor |
| A batch of patches extracted from the images (batch*patches, tensors, channels, patch_height, patch_width), with patches extracted patches. |
| """ |
| batch, tensors, channels, height, width = x.shape |
| x = extract_tensor_patches( |
| x.view(batch, tensors * channels, height, width), |
| window_size, |
| stride=window_size, |
| ) |
| _, patches, _, patch_height, patch_width = x.shape |
| return x.view( |
| batch, patches, tensors, channels, patch_height, patch_width |
| ).view(batch * patches, tensors, channels, patch_height, patch_width) |
|
|
| def extract_psnr_values_for_patches( |
| self, sr, sr_base, lr, instance_psnr, baseline_psnr |
| ): |
| """Extract the PSNR values for the patches. |
| |
| Parameters |
| ---------- |
| sr : torch.Tensor |
| The SR image. |
| sr_base : torch.Tensor |
| The baseline SR image. |
| lr : torch.Tensor |
| The LR image. |
| instance_psnr : torch.Tensor |
| The instance PSNR values. |
| baseline_psnr : torch.Tensor |
| The baseline PSNR values. |
| |
| Returns |
| ------- |
| torch.Tensor, torch.Tensor |
| The extracted instance and baseline PSNR values. |
| """ |
| instance_psnr = ( |
| instance_psnr[None].expand(sr.shape[0] // lr.shape[0], -1).flatten() |
| ) |
| baseline_psnr = ( |
| baseline_psnr[None].expand(sr_base.shape[0] // lr.shape[0], -1).flatten() |
| ) |
|
|
| return instance_psnr, baseline_psnr |
|
|
| def convert_lr_to_RGB_only_for_logging(self, lr): |
| """Convert the LR images to RGB only for logging. |
| |
| Parameters |
| ---------- |
| lr : torch.Tensor |
| The LR image. |
| |
| Returns |
| ------- |
| torch.Tensor |
| The LR image in RGB. |
| """ |
| if ( |
| lr.shape[2] > 3 |
| ): |
| lr = lr[:, :, S2_ALL_12BANDS["true_color_zero_index"]] |
| return lr |
|
|
| def log_image_previews_to_wandb( |
| self, |
| sr, |
| sr_base, |
| lr, |
| hr, |
| instance_psnr, |
| baseline_psnr, |
| pl_module, |
| prefix, |
| trainer, |
| ): |
| """Log the current epoch's images to WandB. |
| |
| Parameters |
| ---------- |
| sr : torch.Tensor |
| The SR image. |
| sr_base : torch.Tensor |
| The baseline SR image. |
| lr : torch.Tensor |
| The LR image. |
| hr : torch.Tensor |
| The HR image. |
| instance_psnr : torch.Tensor |
| The instance PSNR values. |
| baseline_psnr : torch.Tensor |
| The baseline PSNR values. |
| pl_module : pl.LightningModule |
| The model. |
| prefix : str |
| The split prefix (train, val, test). |
| trainer : pytorch_lightning.trainer.Trainer |
| The trainer. |
| """ |
| number_of_rows_for_lr_revisits = int(sqrt(pl_module.hparams.revisits).round()) |
|
|
| trainer.logger.experiment.log( |
| { |
| f"{prefix}-lr": [ |
| wandb.Image( |
| make_grid( |
| lr_image, |
| nrow=number_of_rows_for_lr_revisits, |
| normalize=True, |
| scale_each=True, |
| ), |
| caption=f"lr", |
| ) |
| for lr_image in lr |
| ], |
| f"{prefix}-hr": [ |
| wandb.Image( |
| make_grid(hr_image, nrow=1, normalize=True, scale_each=True), |
| caption=f"hr", |
| ) |
| for hr_image in hr |
| ], |
| f"{prefix}-baseline_upscaled_bicubic": [ |
| wandb.Image( |
| make_grid( |
| baseline_image, nrow=1, normalize=True, scale_each=True |
| ), |
| caption=f"baseline:{str(baseline_psnr)}", |
| ) |
| for baseline_image, baseline_psnr in zip(sr_base, baseline_psnr) |
| ], |
| f"{prefix}-sr": [ |
| wandb.Image( |
| make_grid(image, nrow=1, normalize=True, scale_each=True), |
| caption=f"sr:{str(sr_psnr)}", |
| ) |
| for image, sr_psnr in zip(sr, instance_psnr) |
| ], |
| } |
| ) |
|
|
| def on_validation_epoch_end(self, trainer, pl_module): |
| """Called at the end of the validation epoch. |
| If the current epoch should be logged (i.e. the current epoch is a multiple of the log interval), |
| then _on_epoch_end is called for the next batch of the validation set. |
| |
| Parameters |
| ---------- |
| trainer : pytorch_lightning.trainer.Trainer |
| The trainer. |
| pl_module : pytorch_lightning.LightningModule |
| The model. |
| """ |
| if trainer.current_epoch % self.log_every_n_epochs == 0: |
| batch = next(iter(self.val_dataloader)) |
| self._on_epoch_end("val", batch, trainer, pl_module) |
|
|
| def on_test_epoch_end(self, trainer, pl_module): |
| """Called at the end of the test epoch. |
| |
| Parameters |
| ---------- |
| trainer : pytorch_lightning.trainer.Trainer |
| The trainer. |
| pl_module : pytorch_lightning.LightningModule |
| The model. |
| """ |
| for batch in iter(self.test_dataloader): |
| self._on_epoch_end("test", batch, trainer, pl_module) |
|
|
