SCAIL-Pose/DWPoseProcess/AAUtils.py

# MooreAA 同样API
import importlib
import os
import os.path as osp
import shutil
import sys
from pathlib import Path
import decord
from decord import VideoReader, cpu, gpu
import av
import numpy as np
import torch
import torchvision
from einops import rearrange
from PIL import Image
import time
from fractions import Fraction
import cv2
import jsonlines
import random
import io



def seed_everything(seed):
    import random

    import numpy as np

    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed % (2**32))
    random.seed(seed)


def import_filename(filename):
    spec = importlib.util.spec_from_file_location("mymodule", filename)
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def delete_additional_ckpt(base_path, num_keep):
    dirs = []
    for d in os.listdir(base_path):
        if d.startswith("checkpoint-"):
            dirs.append(d)
    num_tot = len(dirs)
    if num_tot <= num_keep:
        return
    # ensure ckpt is sorted and delete the ealier!
    del_dirs = sorted(dirs, key=lambda x: int(x.split("-")[-1]))[: num_tot - num_keep]
    for d in del_dirs:
        path_to_dir = osp.join(base_path, d)
        if osp.exists(path_to_dir):
            shutil.rmtree(path_to_dir)
        

# def save_videos_from_pil(pil_images, path, fps=8):
#     if fps is None or fps <= 0 or fps > 240:
#         print(f"Warning: Invalid FPS {fps}")
#         return

#     save_fmt = Path(path).suffix
#     os.makedirs(os.path.dirname(path), exist_ok=True)
#     width, height = pil_images[0].size

#     if save_fmt == ".mp4":
#         try:
#             codec = "libx264"
#             container = av.open(path, "w")
#             stream = container.add_stream(codec, rate=fps)

#             stream.width = width
#             stream.height = height

#             for pil_image in pil_images:
#                 # pil_image = Image.fromarray(image_arr).convert("RGB")
#                 av_frame = av.VideoFrame.from_image(pil_image)
#                 container.mux(stream.encode(av_frame))
#             container.mux(stream.encode())
#             container.close()
#         except Exception as e:
#             print(f"Unexpected error while saving video {path}: {e}")
#             if os.path.exists(path):
#                 try:
#                     os.remove(path)
#                     print(f"Corrupted file {path} removed successfully.")
#                 except Exception as rm_e:
#                     print(f"Failed to remove corrupted file {path}: {rm_e}")

#     elif save_fmt == ".gif":
#         pil_images[0].save(
#             fp=path,
#             format="GIF",
#             append_images=pil_images[1:],
#             save_all=True,
#             duration=(1 / fps * 1000),
#             loop=0,
#         )
#     else:
#         raise ValueError("Unsupported file type. Use .mp4 or .gif.")


def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=6, fps=8):
    videos = rearrange(videos, "b c t h w -> t b c h w")
    height, width = videos.shape[-2:]
    outputs = []

    for x in videos:    # x: b c h w
        if x.shape[0] != 1:
            x = torchvision.utils.make_grid(x, nrow=n_rows)  # (c h w)
        else:
            x = x.squeeze(0)
        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)  # (h w c)
        if rescale:
            x = (x + 1.0) / 2.0  # -1,1 -> 0,1
        x = (x * 255).numpy().astype(np.uint8)
        x = Image.fromarray(x)

        outputs.append(x)

    os.makedirs(os.path.dirname(path), exist_ok=True)

    save_videos_from_pil(outputs, path, fps)


def read_frames(video_path):
    try:
        # 使用 decord 打开视频
        vr = VideoReader(video_path)
        frames = []

        # 逐帧解码
        for i in range(len(vr)):
            frame = vr[i]  # 获取帧，返回的是 mx.ndarray
            image = Image.fromarray(frame.asnumpy())  # 转换为 PIL 格式
            frames.append(image)

        return frames

    except Exception as e:
        print(f"Error reading frames from {video_path}: {e}")
        return None  # 返回 None 避免代码崩溃

def get_fps(video_path):

    # container = av.open(video_path)
    # video_stream = next(s for s in container.streams if s.type == "video")
    # fps = video_stream.average_rate
    # container.close()
    # print("pyav_fps")
    # print(fps)
    try:
        vr = decord.VideoReader(video_path)
        fps = vr.get_avg_fps()
        return Fraction(fps).limit_denominator(1001)
    except Exception as e:
        print(f"Error reading FPS from {video_path}: {e}")
        return None  # 返回 None 避免代码崩溃


def read_frames_and_fps(video_path):
    try:
        # 使用 decord 打开视频
        vr = VideoReader(video_path)
        fps = vr.get_avg_fps()
        frames = []

        # 逐帧解码
        for i in range(len(vr)):
            frame = vr[i]  # 获取帧，返回的是 mx.ndarray
            image = Image.fromarray(frame.asnumpy())  # 转换为 PIL 格式
            frames.append(image)

        processed_fps = Fraction(fps).limit_denominator(1001)
        return frames, processed_fps

    except Exception as e:
        print(f"Error reading frames from {video_path}: {e}")
        return None, None  # 返回 None 避免代码崩溃

def read_frames_and_fps_as_np(video_path):
    try:
        # 使用 decord 打开视频
        vr = VideoReader(video_path)
        fps = vr.get_avg_fps()
        frames = []

        # 逐帧解码
        for i in range(len(vr)):
            frame = vr[i]  # 获取帧，返回的是 mx.ndarray
            image = frame.asnumpy()  # 转换为 PIL 格式
            frames.append(image)
        processed_fps = Fraction(fps).limit_denominator(1001)
        return frames, processed_fps

    except Exception as e:
        print(f"Error reading frames from {video_path}: {e}")
        return None, None  # 返回 None 避免代码崩溃



def save_videos_from_pil(pil_images, path, fps=8):
    if fps is None or fps <= 0 or fps > 240:
        print(f"Warning: Invalid FPS {fps}")
        return

    save_fmt = Path(path).suffix
    os.makedirs(os.path.dirname(path), exist_ok=True)
    width, height = pil_images[0].size

    if save_fmt == ".mp4":
        try:
            codec = "libx264"
            container = av.open(path, "w")
            stream = container.add_stream(codec, rate=fps)

            stream.width = width
            stream.height = height

            for pil_image in pil_images:
                # pil_image = Image.fromarray(image_arr).convert("RGB")
                av_frame = av.VideoFrame.from_image(pil_image)
                container.mux(stream.encode(av_frame))
            container.mux(stream.encode())
            container.close()
        except Exception as e:
            print(f"Unexpected error while saving video {path}: {e}")
            if os.path.exists(path):
                try:
                    os.remove(path)
                    print(f"Corrupted file {path} removed successfully.")
                except Exception as rm_e:
                    print(f"Failed to remove corrupted file {path}: {rm_e}")

    elif save_fmt == ".gif":
        pil_images[0].save(
            fp=path,
            format="GIF",
            append_images=pil_images[1:],
            save_all=True,
            duration=(1 / fps * 1000),
            loop=0,
        )
    else:
        raise ValueError("Unsupported file type. Use .mp4 or .gif.")
    

def load_video_with_pose_from_first_frame(video_data, pose_data, sampling="uniform", duration=None, num_frames=99, wanted_fps=None, actual_fps=None,
               skip_frms_num=4., nb_read_frames=None):
    decord.bridge.set_bridge("torch")
    vr = VideoReader(uri=video_data, height=-1, width=-1)
    vr_pose = VideoReader(uri=pose_data, height=-1, width=-1)

    start = 0
    end = int(start + num_frames / wanted_fps * actual_fps)
    n_frms = num_frames + 1     # 要取到num_frames帧, +1把第一帧需要额外拿出来处理，让第一帧和后续帧不同

    if sampling == "uniform":
        indices = np.arange(start, end, (end - start) / n_frms).astype(int)
    else:
        raise NotImplementedError

    # get_batch -> T, H, W, C

    temp_frms = vr.get_batch(np.arange(start, end))
    temp_frms_pose = vr_pose.get_batch(np.arange(start, end))

    assert temp_frms is not None
    assert temp_frms_pose is not None

    tensor_frms = torch.from_numpy(temp_frms) if type(temp_frms) is not torch.Tensor else temp_frms
    tensor_frms = tensor_frms[torch.tensor((indices - start).tolist())]

    tensor_frms_pose = torch.from_numpy(temp_frms_pose) if type(temp_frms_pose) is not torch.Tensor else temp_frms_pose
    tensor_frms_pose = temp_frms_pose[torch.tensor((indices - start).tolist())]

    # print(f"n_frms: {n_frms}; tensor_frms.shape: {tensor_frms.shape} tensor_frms_pose.shape: {tensor_frms_pose.shape}")
    return pad_last_frame(tensor_frms, n_frms), pad_last_frame(tensor_frms_pose, n_frms)


def pad_last_frame(tensor, sampling_frms_num):
    # T, H, W, C
    if tensor.shape[0] < sampling_frms_num:
        # 复制最后一帧
        last_frame = tensor[-int(sampling_frms_num-tensor.shape[0]):]
        # 将最后一帧添加到第二个维度
        padded_tensor = torch.cat([tensor, last_frame], dim=0)
        return padded_tensor
    else:
        return tensor[:sampling_frms_num]
    

def load_video_sampling(video_data, pose_data, num_frames, wanted_fps):
    decord.bridge.set_bridge("torch")
    # 以video_data的为准
    vr = VideoReader(uri=video_data, height=-1, width=-1)
    actual_fps = vr.get_avg_fps()
    if video_data:
        video, pose = load_video_with_pose_from_first_frame(video_data, pose_data, sampling="uniform", duration=100000, num_frames=num_frames, wanted_fps=wanted_fps, actual_fps=actual_fps, skip_frms_num=0, nb_read_frames=None)
        return video, pose
    else:
        raise ValueError("mooreAA should have video data")