File size: 2,660 Bytes
7bef20f | 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 | """Utils functions for visualization."""
import torch
import torchvision.transforms.functional as F
from einops import rearrange
from PIL import Image, ImageDraw, ImageFont
def make_viz_from_samples(
original_images,
reconstructed_images
):
"""Generates visualization images from original images and reconstructed images.
Args:
original_images: A torch.Tensor, original images.
reconstructed_images: A torch.Tensor, reconstructed images.
Returns:
A tuple containing two lists - images_for_saving and images_for_logging.
"""
reconstructed_images = torch.clamp(reconstructed_images, 0.0, 1.0)
reconstructed_images = reconstructed_images * 255.0
reconstructed_images = reconstructed_images.cpu()
original_images = torch.clamp(original_images, 0.0, 1.0)
original_images *= 255.0
original_images = original_images.cpu()
diff_img = torch.abs(original_images - reconstructed_images)
to_stack = [original_images, reconstructed_images, diff_img]
images_for_logging = rearrange(
torch.stack(to_stack),
"(l1 l2) b c h w -> b c (l1 h) (l2 w)",
l1=1).byte()
images_for_saving = [F.to_pil_image(image) for image in images_for_logging]
return images_for_saving, images_for_logging
def make_viz_from_samples_generation(
generated_images,
):
generated = torch.clamp(generated_images, 0.0, 1.0) * 255.0
images_for_logging = rearrange(
generated,
"(l1 l2) c h w -> c (l1 h) (l2 w)",
l1=2)
images_for_logging = images_for_logging.cpu().byte()
images_for_saving = F.to_pil_image(images_for_logging)
return images_for_saving, images_for_logging
def make_viz_from_samples_t2i_generation(
generated_images,
captions,
):
generated = torch.clamp(generated_images, 0.0, 1.0) * 255.0
images_for_logging = rearrange(
generated,
"(l1 l2) c h w -> c (l1 h) (l2 w)",
l1=2)
images_for_logging = images_for_logging.cpu().byte()
images_for_saving = F.to_pil_image(images_for_logging)
# Create a new image with space for captions
width, height = images_for_saving.size
caption_height = 20 * len(captions) + 10
new_height = height + caption_height
new_image = Image.new("RGB", (width, new_height), "black")
new_image.paste(images_for_saving, (0, 0))
# Adding captions below the image
draw = ImageDraw.Draw(new_image)
font = ImageFont.load_default()
for i, caption in enumerate(captions):
draw.text((10, height + 10 + i * 20), caption, fill="white", font=font)
return new_image, images_for_logging |