LTT/PRM

1import os2import imageio3import rembg4import torch5import numpy as np6import PIL.Image7from PIL import Image8from typing import Any9 10 11def remove_background(image: PIL.Image.Image,12    rembg_session: Any = None,13    force: bool = False,14    **rembg_kwargs,15) -> PIL.Image.Image:16    do_remove = True17    if image.mode == "RGBA" and image.getextrema()[3][0] < 255:18        do_remove = False19    do_remove = do_remove or force20    if do_remove:21        image = rembg.remove(image, session=rembg_session, **rembg_kwargs)22    return image23 24 25def resize_foreground(26    image: PIL.Image.Image,27    ratio: float,28) -> PIL.Image.Image:29    image = np.array(image)30    assert image.shape[-1] == 431    alpha = np.where(image[..., 3] > 0)32    y1, y2, x1, x2 = (33        alpha[0].min(),34        alpha[0].max(),35        alpha[1].min(),36        alpha[1].max(),37    )38    # crop the foreground39    fg = image[y1:y2, x1:x2]40    # pad to square41    size = max(fg.shape[0], fg.shape[1])42    ph0, pw0 = (size - fg.shape[0]) // 2, (size - fg.shape[1]) // 243    ph1, pw1 = size - fg.shape[0] - ph0, size - fg.shape[1] - pw044    new_image = np.pad(45        fg,46        ((ph0, ph1), (pw0, pw1), (0, 0)),47        mode="constant",48        constant_values=((0, 0), (0, 0), (0, 0)),49    )50 51    # compute padding according to the ratio52    new_size = int(new_image.shape[0] / ratio)53    # pad to size, double side54    ph0, pw0 = (new_size - size) // 2, (new_size - size) // 255    ph1, pw1 = new_size - size - ph0, new_size - size - pw056    new_image = np.pad(57        new_image,58        ((ph0, ph1), (pw0, pw1), (0, 0)),59        mode="constant",60        constant_values=((0, 0), (0, 0), (0, 0)),61    )62    new_image = PIL.Image.fromarray(new_image)63    return new_image64 65 66def images_to_video(67    images: torch.Tensor, 68    output_path: str, 69    fps: int = 30,70) -> None:71    # images: (N, C, H, W)72    video_dir = os.path.dirname(output_path)73    video_name = os.path.basename(output_path)74    os.makedirs(video_dir, exist_ok=True)75 76    frames = []77    for i in range(len(images)):78        frame = (images[i].permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)79        assert frame.shape[0] == images.shape[2] and frame.shape[1] == images.shape[3], \80            f"Frame shape mismatch: {frame.shape} vs {images.shape}"81        assert frame.min() >= 0 and frame.max() <= 255, \82            f"Frame value out of range: {frame.min()} ~ {frame.max()}"83        frames.append(frame)84    imageio.mimwrite(output_path, np.stack(frames), fps=fps, quality=10)85 86 87def save_video(88    frames: torch.Tensor,89    output_path: str,90    fps: int = 30,91) -> None:92    # images: (N, C, H, W)93    frames = [(frame.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8) for frame in frames]94    writer = imageio.get_writer(output_path, fps=fps)95    for frame in frames:96        writer.append_data(frame)97    writer.close()