pvcracks.vae subpackage

This subpackage contains the Variational Autoencoder tools for PVCracks.

VAE helper modules

Created on Wed Apr 5 12:11:11 2023

@author: jlbraid

class pvcracks.vae.custom_dataset.CustomDataset(data, transform=None)[source]

class pvcracks.vae.custom_dataset.FixedHorizontalFlip[source]

class pvcracks.vae.custom_dataset.FixedRotation(angle)[source]

class pvcracks.vae.custom_dataset.FixedVerticalFlip[source]

Created on Mon Aug 7 12:04:49 2023

@authors: jlbraid, nrjost

pvcracks.vae.VAE_functions.VAE_output_for_images(model, images)[source]

Get the VAE outputs for the input images.

Parameters

model (torch.nn.Module) – The VAE model.
images (torch.Tensor) – The input images.

Returns

The VAE outputs.

Return type

torch.Tensor

Notes

pvcracks.vae.VAE_functions.decode_latent_vector(model, latent_vector)[source]

Decode a latent vector using the VAE decoder.

Parameters

model (torch.nn.Module) – The VAE model.
latent_vector (torch.Tensor) – The latent vector.

Returns

The reconstructed image.

Return type

torch.Tensor

Notes

pvcracks.vae.VAE_functions.encode_image(model, image)[source]

Encode an image using the VAE encoder.

Parameters

model (torch.nn.Module) – The VAE model.
image (torch.Tensor) – The input image tensor.

Returns

The latent vector.

Return type

torch.Tensor

Notes

pvcracks.vae.VAE_functions.generate_random_images(model, num_images, latent_dim, device='cuda')[source]

Generate random images using the VAE model.

Parameters

model (torch.nn.Module) – The VAE model.
num_images (int) – The number of images to generate.
latent_dim (int) – The dimension of the latent space.
device (str = "cpu" or "cuda") – The device to perform computations on.

Returns

The generated images and the random latent vectors.

Return type

tuple

Notes

pvcracks.vae.VAE_functions.initialize_model_optimizer(model, latent_dim, learning_rate, device)[source]

Initialize the VAE model and optimizer.

Parameters

model (torch.nn.Module) – The VAE model class.
latent_dim (int) – The dimension of the latent space.
learning_rate (float) – The learning rate for the optimizer.
device (str = "cpu" or "cuda") – The device to perform computations on.

Returns

The initialized model and optimizer.

Return type

tuple

Notes

The function initializes a VAE model and an Adam optimizer.

pvcracks.vae.VAE_functions.load_from_testloader(test, num_images=100)[source]

Load a batch of images from the test dataset.

Parameters

test (torch.utils.data.Dataset) – The test dataset.
num_images (int) – The number of images to load in a batch.

Returns

A batch of images.

Return type

torch.Tensor

Notes

pvcracks.vae.VAE_functions.plot_training_losses(num_epochs, train_losses, path)[source]

Plot the training losses per epoch.

Parameters

num_epochs (int) – The number of training epochs.
train_losses (list) – A list of training losses.
path (str or None) – Path to save the plot.

Return type

None

Notes

pvcracks.vae.VAE_functions.preprocess(impath)[source]

Preprocess an image from the specified path.

Parameters: impath (str) – Path to the image file.
Returns: The preprocessed image as a numpy array.
Return type: numpy.ndarray

Notes

This function reads an image, normalizes its pixel values to the range [0, 1], and converts it to a float32 numpy array.

pvcracks.vae.VAE_functions.set_seeds(seed=50, multiGPU=False)[source]

Set random seeds for reproducibility.

Parameters

seed (int) – The seed value to use for random number generators.
multiGPU (bool) – If True, sets seeds for multiple GPUs.

Return type

None

Notes

This function sets seeds for Python’s random module, numpy, and PyTorch (CPU and CUDA).

pvcracks.vae.VAE_functions.show_generated_images(generated_images, num_images, path)[source]

Display the generated images.

Parameters

generated_images (torch.Tensor) – The generated images.
num_images (int) – The number of images to display.
path (str or None) – Path to save the plot.

Return type

None

Notes

pvcracks.vae.VAE_functions.show_input_output_images(inputs, outputs, num_images, path)[source]

Display the input images and their VAE outputs.

Parameters

inputs (torch.Tensor) – The input images.
outputs (torch.Tensor) – The VAE output images.
num_images (int) – The number of images to display.
path (str) – Path to save the plot.

Return type

None

Notes

The function displays the input images and their VAE outputs side by side.

pvcracks.vae.VAE_functions.ssim_input_output(inputs, outputs, num_images, path)[source]

Compute the SSIM between input and output images and save the results.

Parameters

inputs (torch.Tensor) – The input images.
outputs (torch.Tensor) – The VAE output images.
num_images (int) – The number of images to compare.
path (str or None) – Path to save the results.

Returns

The SSIM comparison results.

Return type

pandas.DataFrame

Notes

The function computes the SSIM between input and output images and saves the results to a CSV file.

pvcracks.vae.VAE_functions.train_model(model)[source]

Train the VAE model.

Parameters: model (torch.nn.Module) – The VAE model to be trained.
Returns: The mean, log variance, training losses, and number of epochs.
Return type: tuple

Notes

The function trains the VAE model for a specified number of epochs and returns the training losses.

pvcracks.vae.VAE_functions.vae_loss(recon_x, x, mu, logvar, bce_weight, kld_weight, ssim_weight, device='cuda')[source]

Compute the loss for the Variational Autoencoder (VAE).

Parameters

recon_x (torch.Tensor) – The reconstructed images.
x (torch.Tensor) – The original input images.
mu (torch.Tensor) – The mean of the latent distribution.
logvar (torch.Tensor) – The log variance of the latent distribution.
bce_weight (float) – Weight for the Binary Cross Entropy (BCE) loss.
kld_weight (float) – Weight for the Kullback-Leibler Divergence (KLD) loss.
ssim_weight (float) – Weight for the Structural Similarity Index (SSIM) loss.
device (str) – The device to perform computations on.

Returns

The total loss.

Return type

torch.Tensor

Notes

The function combines BCE loss, KLD loss, and SSIM loss to compute the total VAE loss. Uses: from pytorch_ssim import SSIM don’t use pip installed version (not maintained). Use: https://github.com/Po-Hsun-Su/pytorch-ssim

class pvcracks.vae.VAE_model_1CH.Decoder(latent_dim)[source]

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class pvcracks.vae.VAE_model_1CH.Encoder(latent_dim)[source]

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class pvcracks.vae.VAE_model_1CH.VAE(latent_dim)[source]

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reparameterize(mu, logvar)[source]

class pvcracks.vae.VAE_model_3CH.Decoder(latent_dim)[source]

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class pvcracks.vae.VAE_model_3CH.Encoder(latent_dim)[source]

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class pvcracks.vae.VAE_model_3CH.VAE(latent_dim)[source]

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reparameterize(mu, logvar)[source]