Spaces:
Running on Zero
Running on Zero
File size: 2,606 Bytes
7405e85 | 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 | import io
from pathlib import Path
from typing import Union
import imageio.v2 as imageio
import numpy as np
import torch
import torchvision.transforms as tf
from einops import rearrange, repeat
from jaxtyping import Float, UInt8
from matplotlib.figure import Figure
from PIL import Image
from torch import Tensor
FloatImage = Union[
Float[Tensor, "height width"],
Float[Tensor, "channel height width"],
Float[Tensor, "batch channel height width"],
]
def fig_to_image(
fig: Figure,
dpi: int = 100,
device: torch.device = torch.device("cpu"),
) -> Float[Tensor, "3 height width"]:
buffer = io.BytesIO()
fig.savefig(buffer, format="raw", dpi=dpi)
buffer.seek(0)
data = np.frombuffer(buffer.getvalue(), dtype=np.uint8)
h = int(fig.bbox.bounds[3])
w = int(fig.bbox.bounds[2])
data = rearrange(data, "(h w c) -> c h w", h=h, w=w, c=4)
buffer.close()
return (torch.tensor(data, device=device, dtype=torch.float32) / 255)[:3]
def prep_image(image: FloatImage) -> UInt8[np.ndarray, "height width channel"]:
# Handle batched images.
if image.ndim == 4:
image = rearrange(image, "b c h w -> c h (b w)")
# Handle single-channel images.
if image.ndim == 2:
image = rearrange(image, "h w -> () h w")
# Ensure that there are 3 or 4 channels.
channel, _, _ = image.shape
if channel == 1:
image = repeat(image, "() h w -> c h w", c=3)
assert image.shape[0] in (3, 4)
image = (image.detach().clip(min=0, max=1) * 255).type(torch.uint8)
return rearrange(image, "c h w -> h w c").cpu().numpy()
def save_image(
image: FloatImage,
path: Union[Path, str],
) -> None:
"""Save an image. Assumed to be in range 0-1."""
# Create the parent directory if it doesn't already exist.
path = Path(path)
path.parent.mkdir(exist_ok=True, parents=True)
# Save the image.
Image.fromarray(prep_image(image)).save(path)
def load_image(
path: Union[Path, str],
) -> Float[Tensor, "3 height width"]:
return tf.ToTensor()(Image.open(path))[:3]
def save_video(
images: list[FloatImage],
path: Union[Path, str],
fps: None | int = None,
) -> None:
"""Save a video whose input frames are assumed to be in range 0-1."""
path = Path(path)
path.parent.mkdir(exist_ok=True, parents=True)
with imageio.get_writer(
str(path),
fps=fps or 30,
codec="libx264",
pixelformat="yuv420p",
output_params=["-crf", "0"],
) as writer:
for image in images:
writer.append_data(prep_image(image))
|