Spaces:
Runtime error
Runtime error
File size: 1,361 Bytes
a7eca0b | 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 | #visualization/utils.py
import torch
import numpy as np
from imageio import mimsave
from PIL import Image, ImageDraw, ImageFont
def unnormalize_img(image: np.ndarray, std: tuple, mean: tuple) -> np.ndarray:
image = (image * std) + mean
image = (image * 255).astype('uint8')
return image.clip(0, 255)
def save_as_gif(
video_tensor: torch.Tensor,
save_path: str = 'sample.gif',
std: tuple = None,
mean: tuple = None,
):
frames = []
for video_frame in video_tensor:
frame_unnormalized = unnormalize_img(
image=video_frame.permute(1, 2, 0).numpy(),
std=std,
mean=mean,
)
frames.append(frame_unnormalized)
kargs = {'duration': 0.25}
mimsave(save_path, frames, 'GIF', **kargs)
return save_path
def display_gif(gif_path: str) -> Image:
return Image(filename=gif_path)
def draw_text_on_image(
image: np.ndarray,
text: str,
position: tuple = (20, 20),
color: tuple = (0, 0, 255),
font_size: int = 20,
) -> np.ndarray:
font = ImageFont.truetype(
font="fonts/OpenSans-Regular.ttf",
size=font_size,
)
pil_image = Image.fromarray(image)
draw = ImageDraw.Draw(pil_image)
draw.text(
xy=position,
text=text,
fill=color,
font=font,
)
return np.array(pil_image)
|