biapy.engine.super_resolution

Super-resolution workflow for BiaPy.

This module defines the Super_resolution_Workflow class, which implements the training, validation, and inference pipeline for single-image super-resolution and image restoration tasks in BiaPy. It supports 2D and 3D data, multiple metrics (PSNR, SSIM, FID, IS, LPIPS), and handles data loading, model setup, predictions, and result saving for reconstructing high-resolution images from low-resolution inputs.

class biapy.engine.super_resolution.Super_resolution_Workflow(cfg, job_identifier, device, system_dict, args, **kwargs)[source]

Bases: Base_Workflow

Single-image super-resolution workflow where the goal is to reconstruct high-resolution (HR) images from low-resolution (LR) ones. If there is a difference in the size of the LR and HR images, typically determined by a scale factor (x2, x4), this task is known as single-image super-resolution. If the size of the LR and HR images is the same, this task is usually referred to as image restoration.

More details in our documentation.

Parameters:

cfg (YACS configuration) – Running configuration.
Job_identifier (str) – Complete name of the running job.
device (Torch device) – Device used.
args (argpase class) – Arguments used in BiaPy’s call.

define_activations_and_channels()[source]

Define the activations to be applied to the model output and the channels that the model will output.

This function must define the following variables:

self.model_output_channelsList of int: Number of channels for each output head of the model. E.g. [3] for a model with one head outputting 3 channels, [1, 5] for a model with two heads outputting 1 and 5 channels respectively, etc.
self.model_output_channel_infoList of str: Information about the output channels. A value per output head of the model must be defined.
self.separated_class_channelbool: Whether if we should expect a separated output channel for classification.
self.head_activationsList of str: Activations to be applied to the model output. A value per output channel (not output head) of the model must be defined. “linear” and “ce_sigmoid” will not be applied. E.g. [“linear”] for a model with one channel, [“linear”, “sigmoid”] for a model with two channels, etc.

Example of a correct definition of the function for a model with two output heads: 1) the first one will be predicting foreground and contours; 2) the second one will classify into 3 classes the predicted objects. In this case the following definition would be correct:

self.model_output_channels = [1, 3]
self.model_output_channel_info = ["mask", "class"]
self.separated_class_channel = True
self.head_activations = ["ce_sigmoid", "ce_sigmoid", "ce_softmax", "ce_softmax", "ce_softmax"]

define_metrics()[source]

Define the metrics to be used during training and test.

This function must define the following variables:

self.train_metricsList of functions: Metrics to be calculated during model’s training.
self.train_metric_namesList of str: Names of the metrics calculated during training.
self.train_metric_bestList of str: To know which value should be considered as the best one. Options must be: “max” or “min”.
self.test_metricsList of functions: Metrics to be calculated during model’s test/inference.
self.test_metric_namesList of str: Names of the metrics calculated during test/inference.
self.lossFunction: Loss function used during training and test.

metric_calculation(output: ndarray[tuple[int, ...], dtype[_ScalarType_co]] | Tensor, targets: ndarray[tuple[int, ...], dtype[_ScalarType_co]] | Tensor, train: bool = True, metric_logger: MetricLogger | None = None) → Dict[source]

Calculate the metrics defined in define_metrics() function.

Parameters:

output (Torch Tensor) – Prediction of the model.
targets (Torch Tensor) – Ground truth to compare the prediction with.
train (bool, optional) – Whether to calculate train or test metrics.
metric_logger (MetricLogger, optional) – Class to be updated with the new metric(s) value(s) calculated.

Returns:

out_metrics – Value of the metrics for the given prediction.

Return type:

dict

process_test_sample()[source]: Process a sample in the test/inference phase.

torchvision_model_call(in_img: Tensor, is_train: bool = False) → Tensor | None[source]

Call a regular Pytorch model.

Parameters:

in_img (torch.Tensor) – Input image to pass through the model.
is_train (bool, optional) – Whether if the call is during training or inference.

Returns:

prediction – Image prediction.

Return type:

torch.Tensor

after_merge_patches(pred)[source]

Execute steps needed after merging all predicted patches into the original image.

Parameters:: pred (Torch Tensor) – Model prediction.

after_full_image(pred: ndarray[tuple[int, ...], dtype[_ScalarType_co]])[source]

Execute steps needed after generating the prediction by supplying the entire image to the model.

Parameters:: pred (NDArray) – Model prediction.

after_all_images()[source]: Execute steps needed after predicting all images.

after_all_chunk_prediction_workflow_process()[source]: Place any code that needs to be done after predicting all patches in “by chunks” setting. This function is called on all ranks.

after_all_chunk_prediction_workflow_process_master_rank()[source]: Place any code that needs to be done after predicting all patches in “by chunks” setting, but only on the master rank. This function is called only on the master rank.