multimodalart/EchoMimic-zero

1import importlib2import os3import os.path as osp4import shutil5import sys6from pathlib import Path7 8import av9import numpy as np10import torch11import torchvision12from einops import rearrange13from PIL import Image14 15 16def seed_everything(seed):17    import random18 19    import numpy as np20 21    torch.manual_seed(seed)22    torch.cuda.manual_seed_all(seed)23    np.random.seed(seed % (2**32))24    random.seed(seed)25 26 27def import_filename(filename):28    spec = importlib.util.spec_from_file_location("mymodule", filename)29    module = importlib.util.module_from_spec(spec)30    sys.modules[spec.name] = module31    spec.loader.exec_module(module)32    return module33 34 35def delete_additional_ckpt(base_path, num_keep):36    dirs = []37    for d in os.listdir(base_path):38        if d.startswith("checkpoint-"):39            dirs.append(d)40    num_tot = len(dirs)41    if num_tot <= num_keep:42        return43    # ensure ckpt is sorted and delete the ealier!44    del_dirs = sorted(dirs, key=lambda x: int(x.split("-")[-1]))[: num_tot - num_keep]45    for d in del_dirs:46        path_to_dir = osp.join(base_path, d)47        if osp.exists(path_to_dir):48            shutil.rmtree(path_to_dir)49 50 51def save_videos_from_pil(pil_images, path, fps=8, audio_path=None):52    import av53 54    save_fmt = Path(path).suffix55    os.makedirs(os.path.dirname(path), exist_ok=True)56    width, height = pil_images[0].size57 58    if save_fmt == ".mp4":59        codec = "libx264"60        container = av.open(path, "w")61        stream = container.add_stream(codec, rate=fps)62 63        stream.width = width64        stream.height = height65 66        for pil_image in pil_images:67            # pil_image = Image.fromarray(image_arr).convert("RGB")68            av_frame = av.VideoFrame.from_image(pil_image)69            container.mux(stream.encode(av_frame))70        container.mux(stream.encode())71        container.close()72 73    elif save_fmt == ".gif":74        pil_images[0].save(75            fp=path,76            format="GIF",77            append_images=pil_images[1:],78            save_all=True,79            duration=(1 / fps * 1000),80            loop=0,81        )82    else:83        raise ValueError("Unsupported file type. Use .mp4 or .gif.")84 85 86def save_videos_grid(videos: torch.Tensor, path: str, audio_path=None, rescale=False, n_rows=6, fps=8):87    videos = rearrange(videos, "b c t h w -> t b c h w")88    height, width = videos.shape[-2:]89    outputs = []90 91    for x in videos:92        x = torchvision.utils.make_grid(x, nrow=n_rows)  # (c h w)93        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)  # (h w c)94        if rescale:95            x = (x + 1.0) / 2.0  # -1,1 -> 0,196        x = (x * 255).numpy().astype(np.uint8)97        x = Image.fromarray(x)98 99        outputs.append(x)100 101    os.makedirs(os.path.dirname(path), exist_ok=True)102 103    save_videos_from_pil(outputs, path, fps, audio_path=audio_path)104 105 106def read_frames(video_path):107    container = av.open(video_path)108 109    video_stream = next(s for s in container.streams if s.type == "video")110    frames = []111    for packet in container.demux(video_stream):112        for frame in packet.decode():113            image = Image.frombytes(114                "RGB",115                (frame.width, frame.height),116                frame.to_rgb().to_ndarray(),117            )118            frames.append(image)119 120    return frames121 122 123def get_fps(video_path):124    container = av.open(video_path)125    video_stream = next(s for s in container.streams if s.type == "video")126    fps = video_stream.average_rate127    container.close()128    return fps129 130 131def crop_and_pad(image, rect):132    x0, y0, x1, y1 = rect133    h, w = image.shape[:2]134 135    # 确保坐标在图像范围内136    x0, y0 = max(0, x0), max(0, y0)137    x1, y1 = min(w, x1), min(h, y1)138 139    # 计算原始框的宽度和高度140    width = x1 - x0141    height = y1 - y0142 143    # 使用较小的边长作为裁剪正方形的边长144    side_length = min(width, height)145 146    # 计算正方形框中心点147    center_x = (x0 + x1) // 2148    center_y = (y0 + y1) // 2149 150    # 重新计算正方形框的坐标151    new_x0 = max(0, center_x - side_length // 2)152    new_y0 = max(0, center_y - side_length // 2)153    new_x1 = min(w, new_x0 + side_length)154    new_y1 = min(h, new_y0 + side_length)155 156    # 最终裁剪框的尺寸修正（确保是正方形）157    if (new_x1 - new_x0) != (new_y1 - new_y0):158        side_length = min(new_x1 - new_x0, new_y1 - new_y0)159        new_x1 = new_x0 + side_length160        new_y1 = new_y0 + side_length161 162    # 裁剪图像163    cropped_image = image[new_y0:new_y1, new_x0:new_x1]164 165    return cropped_image