File size: 3,371 Bytes
fe43247 d1f1433 fe43247 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | """Utility helpers for data processing and visualization."""
import torch
import numpy as np
import os
import matplotlib.pyplot as plt
from config import vae_plots_path, unet_plots_path, cfg, device
from diffusers import DDPMScheduler
from tqdm import tqdm
def create_path_if_not_exists(path: str) -> None:
"""Create a directory if it does not already exist."""
if not os.path.exists(path):
os.makedirs(path)
def revert_images(imgs: torch.Tensor) -> np.ndarray:
"""Convert normalized tensors into displayable images."""
h = imgs.shape[-1]
imgs = imgs.cpu().detach().numpy()
min_vals = imgs.min(axis=(1, 2, 3))[:, np.newaxis, np.newaxis, np.newaxis]
max_vals = imgs.max(axis=(1, 2, 3))[:, np.newaxis, np.newaxis, np.newaxis]
imgs = ((max_vals - imgs) / (max_vals - min_vals)) * 255
if imgs.shape[1] == 1:
imgs = imgs.astype(int).reshape(-1, h, h)
return imgs
def plot_side_by_side(
images_y: torch.Tensor, images_pred: torch.Tensor, latents: torch.Tensor, epoch: int
) -> None:
"""Visualize input images, outputs and latent channels."""
images_y, images_pred = revert_images(images_y), revert_images(images_pred)
latents = revert_images(latents)
idx = np.random.randint(0, images_y.shape[0])
fig, axs = plt.subplots(1, 2)
# Plot input image and Output image
axs[0].imshow(images_y[idx], cmap="gray")
axs[0].axis("off")
axs[0].set_title("Input")
axs[1].imshow(images_pred[idx], cmap="gray")
axs[1].axis("off")
axs[1].set_title("Output")
plt.savefig(os.path.join(vae_plots_path, f"epoch_{epoch}_input_output.png"))
plt.clf()
latent_channels = latents.shape[1]
fig, axs = plt.subplots(1, 4)
# Plot the different latent channels
for i in range(latent_channels):
axs[i].imshow(latents[idx, i, :, :], cmap="gray")
axs[i].axis("off")
axs[i].set_title(f"Latent channel: {i}", fontsize=8)
plt.savefig(os.path.join(vae_plots_path, f"epoch_{epoch}_latent_channels.png"))
plt.clf()
def generate(
vae: torch.nn.Module,
unet: torch.nn.Module,
noise_scheduler: DDPMScheduler,
epoch: int,
) -> None:
"""Generate samples from the UNet model."""
def plot(recon_imgs: torch.Tensor, timesteps: int, epoch: int) -> None:
create_path_if_not_exists(os.path.join(unet_plots_path, f"epoch_{epoch}"))
recon_imgs = revert_images(recon_imgs.sample)
fig, axs = plt.subplots(2, 5)
for i in range(10):
axs[i // 5][i % 5].imshow(recon_imgs[i], cmap="gray")
axs[i // 5][i % 5].axis("off")
axs[i // 5][i % 5].set_title(str(i))
plt.suptitle(f"Timesteps: {timesteps}")
plt.savefig(
os.path.join(unet_plots_path, f"epoch_{epoch}", f"plot {timesteps}.png")
)
plt.clf()
latents = torch.randn((10, cfg.latent_channels, 8, 8)).to(device)
labels = torch.arange(10).to(device)
for t in tqdm(noise_scheduler.timesteps):
with torch.no_grad():
noise_pred = unet(
latents, t, class_labels=labels, encoder_hidden_states=None
).sample
latents = noise_scheduler.step(noise_pred, t, latents).prev_sample
recon_imgs = vae.decode(latents)
if t == 999 or t % 100 == 0:
plot(recon_imgs, t, epoch)
|