pose2video.py

import os
import sys
from datetime import datetime

import cv2
import numpy as np
import torch
from diffusers import AutoencoderKL, DDIMScheduler
from einops import repeat
from omegaconf import OmegaConf
from PIL import Image
from torchvision import transforms
from transformers import CLIPVisionModelWithProjection

from src.models.pose_guider import PoseGuider
from src.models.unet_2d_condition import UNet2DConditionModel
from src.models.unet_3d import UNet3DConditionModel
from src.pipelines.pipeline_pose2vid_long import Pose2VideoPipeline
from src.utils.util import get_fps, read_frames, save_videos_grid, save_videos_from_pil


def apply_mask_and_replace_background(original_video_path, mask_video_path, background_image_path, output_video_path):
    # 读取视频源文件
    cap_original = cv2.VideoCapture(original_video_path)
    cap_mask = cv2.VideoCapture(mask_video_path)

    # 获取原视频的宽高和帧率
    frame_width = int(cap_original.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap_original.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap_original.get(cv2.CAP_PROP_FPS)

    # 读取背景图像，并调整为与视频相同的分辨率
    background_image = cv2.imread(background_image_path)
    background_image = cv2.resize(background_image, (frame_width, frame_height))

    # 定义视频编码器
    # out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))
    pil_array_list = []
    while True:
        ret_original, frame_original = cap_original.read()
        ret_mask, frame_mask = cap_mask.read()

        # 如果任一视频流结束，退出循环
        if not ret_original or not ret_mask:
            break

        """
        # 将mask图像转换为二值图像
        ret, frame_mask = cv2.threshold(frame_mask, 127, 255, cv2.THRESH_TOZERO)
        # 创建遮罩层，并使用mask覆盖原视频中的人物部分
        masked_frame = cv2.bitwise_and(frame_original, frame_mask)
        # 使用逆向遮罩与背景合成，保证背景中的非人物部分覆盖
        inverse_mask = cv2.bitwise_not(frame_mask)
        background_part = cv2.bitwise_and(background_image, inverse_mask)
        # 将人物和背景合并在一起
        combined_frame = cv2.add(masked_frame, background_part)
        """

        frame_mask = frame_mask.astype(np.float32) / 255
        combined_frame = background_image * (1 - frame_mask) + frame_original * frame_mask
        combined_frame = combined_frame.astype(np.uint8)

        # 将cv2格式数据改为通过现有的av库来写入视频流数据
        pil_combined_frame = Image.fromarray(cv2.cvtColor(combined_frame, cv2.COLOR_BGR2RGB))
        pil_array_list.append(pil_combined_frame)
        # 写入合成帧到输出视频
        # out.write(combined_frame)

    save_videos_from_pil(pil_array_list, output_video_path, fps)
    # 释放资源
    # out.release()
    cap_original.release()
    cap_mask.release()
    cv2.destroyAllWindows()


def save_video_with_audio(input_video, input_audio, output_video):
    import subprocess
    # 构建 ffmpeg 命令
    ffmpeg_command = [
        "ffmpeg",
        "-y",  # 强制覆盖输出文件而不询问
        "-i", input_video,
        "-i", input_audio,
        "-c:v", "copy",
        "-c:a", "aac",
        "-map", "0:v:0",
        "-map", "1:a:0",
        "-strict", "experimental",
        "-shortest",
        output_video
    ]
    # 运行 ffmpeg 命令
    subprocess.run(ffmpeg_command, check=True)


class AnimateController:
    def __init__(
            self,
            config_path="configs/prompts/animation.yaml",
            weight_dtype=torch.float16,
    ):
        # Read pretrained weights path from config
        self.config = OmegaConf.load(config_path)
        self.pipeline = None
        self.rvm = None
        self.lama = None
        self.weight_dtype = weight_dtype

    def animate(
            self,
            ref_image,
            pose_video_path,
            width=512,
            height=768,
            length=24,
            num_inference_steps=25,
            cfg=3.5,
            seed=123,
            save_dir=f"output/demo"
    ):
        generator = torch.manual_seed(seed)
        if isinstance(ref_image, np.ndarray):
            ref_image = Image.fromarray(ref_image)
        if self.pipeline is None:
            vae = AutoencoderKL.from_pretrained(
                self.config.pretrained_vae_path,
            ).to("cuda", dtype=self.weight_dtype)

            reference_unet = UNet2DConditionModel.from_pretrained(
                self.config.pretrained_base_model_path,
                subfolder="unet",
            ).to(dtype=self.weight_dtype, device="cuda")

            inference_config_path = self.config.inference_config
            infer_config = OmegaConf.load(inference_config_path)
            denoising_unet = UNet3DConditionModel.from_pretrained_2d(
                self.config.pretrained_base_model_path,
                self.config.motion_module_path,
                subfolder="unet",
                unet_additional_kwargs=infer_config.unet_additional_kwargs,
            ).to(dtype=self.weight_dtype, device="cuda")

            pose_guider = PoseGuider(320, block_out_channels=(16, 32, 96, 256)).to(
                dtype=self.weight_dtype, device="cuda"
            )

            image_enc = CLIPVisionModelWithProjection.from_pretrained(
                self.config.image_encoder_path
            ).to(dtype=self.weight_dtype, device="cuda")
            sched_kwargs = OmegaConf.to_container(infer_config.noise_scheduler_kwargs)
            scheduler = DDIMScheduler(**sched_kwargs)

            # load pretrained weights
            denoising_unet.load_state_dict(
                torch.load(self.config.denoising_unet_path, map_location="cpu"),
                strict=False,
            )
            reference_unet.load_state_dict(
                torch.load(self.config.reference_unet_path, map_location="cpu"),
            )
            pose_guider.load_state_dict(
                torch.load(self.config.pose_guider_path, map_location="cpu"),
            )

            pipe = Pose2VideoPipeline(
                vae=vae,
                image_encoder=image_enc,
                reference_unet=reference_unet,
                denoising_unet=denoising_unet,
                pose_guider=pose_guider,
                scheduler=scheduler,
            )
            pipe = pipe.to("cuda", dtype=self.weight_dtype)
            self.pipeline = pipe

        pose_images = read_frames(pose_video_path)
        src_fps = get_fps(pose_video_path)

        pose_list = []
        pose_tensor_list = []
        pose_transform = transforms.Compose(
            [transforms.Resize((height, width)), transforms.ToTensor()]
        )
        for pose_image_pil in pose_images[:length]:
            pose_list.append(pose_image_pil)
            pose_tensor_list.append(pose_transform(pose_image_pil))

        video = self.pipeline(
            ref_image,
            pose_list,
            width=width,
            height=height,
            video_length=length,
            num_inference_steps=num_inference_steps,
            guidance_scale=cfg,
            generator=generator,
        ).videos

        # ref_image_tensor = pose_transform(ref_image)  # (c, h, w)
        # ref_image_tensor = ref_image_tensor.unsqueeze(1).unsqueeze(0)  # (1, c, 1, h, w)
        # ref_image_tensor = repeat(
        #     ref_image_tensor, "b c f h w -> b c (repeat f) h w", repeat=length
        # )
        # pose_tensor = torch.stack(pose_tensor_list, dim=0)  # (f, c, h, w)
        # pose_tensor = pose_tensor.transpose(0, 1)
        # pose_tensor = pose_tensor.unsqueeze(0)
        # video = torch.cat([ref_image_tensor, pose_tensor, video], dim=0)

        # save_dir = f"./output/demo"
        if not os.path.exists(save_dir):
            os.makedirs(save_dir, exist_ok=True)
        date_str = datetime.now().strftime("%Y%m%d")
        time_str = datetime.now().strftime("%H%M%S")
        out_path = os.path.join(save_dir, f"{date_str}T{time_str}.mp4")

        # add
        im_ratio = float(ref_image.size[1]) / ref_image.size[0]

        save_videos_grid(
            video,
            out_path,
            n_rows=1,
            fps=src_fps,
            im_ratio=im_ratio
        )

        torch.cuda.empty_cache()

        # rvm video
        sys.path.append("RobustVideoMatting")
        from RobustVideoMatting.inference import Converter
        out_path_alpha = os.path.join(save_dir, f"{date_str}T{time_str}_alpha.mp4")
        if self.rvm is None:
            # self.rvm = Converter("mobilenetv3", "pretrained_weights/rvm_mobilenetv3.pth", "cuda")
            self.rvm = Converter("resnet50", "pretrained_weights/rvm_resnet50.pth", "cuda")
        self.rvm.convert(input_source=out_path, output_alpha=out_path_alpha, progress=True)

        # rvm image
        out_path_image = os.path.join(save_dir, f"{date_str}T{time_str}_image")
        if os.path.exists(out_path_image):
            os.removedirs(out_path_image)
        os.makedirs(out_path_image, exist_ok=True)
        out_path_image_src = os.path.join(out_path_image, "src.png")
        ref_image.save(out_path_image_src)
        self.rvm.convert(input_source=out_path_image, output_alpha=out_path_image, output_type="png_sequence",
                         progress=True)

        # lama image
        mask_image_path = os.path.join(out_path_image, "0000.png")
        from lama_demo import lama
        if self.lama is None:
            self.lama = lama("pretrained_weights/big-lama.pt")
        image = cv2.imread(out_path_image_src)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        mask = cv2.imread(mask_image_path, cv2.IMREAD_GRAYSCALE)
        out_lama = self.lama.process(image, mask)
        out_lama = cv2.cvtColor(out_lama, cv2.COLOR_RGB2BGR)
        out_path_image_dst = os.path.join(out_path_image, "dst.png")
        cv2.imwrite(out_path_image_dst, out_lama)

        # merge alpha video
        out_path_replace = os.path.join(save_dir, f"{date_str}T{time_str}_replace.mp4")
        apply_mask_and_replace_background(out_path, out_path_alpha, out_path_image_dst, out_path_replace)

        # add music
        out_path_with_audio = os.path.join(save_dir, f"{date_str}T{time_str}_with_audio.mp4")
        pose_video_name = os.path.basename(pose_video_path)
        audio_name = pose_video_name.replace("_kps.mp4", ".mp3")
        audio_path = "./test-video_audio"
        audio_path = os.path.join(audio_path, audio_name)
        save_video_with_audio(out_path_replace, audio_path, out_path_with_audio)
        return out_path_with_audio

    def animate_v2(
            self,
            ref_image,
            video_path,
            width=512,
            height=768,
            length=24,
            num_inference_steps=25,
            cfg=3.5,
            seed=123,
            save_dir=f"output/demo"
    ):
        # pose video
        from tools.vid2pose import run_by_refImg_demo
        # run pose
        pose_video_path = video_path.replace("_15fps", "_kps")
        run_by_refImg_demo(video_path, ref_image, pose_video_path)

        generator = torch.manual_seed(seed)
        if isinstance(ref_image, np.ndarray):
            ref_image = Image.fromarray(ref_image)
        if self.pipeline is None:
            vae = AutoencoderKL.from_pretrained(
                self.config.pretrained_vae_path,
            ).to("cuda", dtype=self.weight_dtype)

            reference_unet = UNet2DConditionModel.from_pretrained(
                self.config.pretrained_base_model_path,
                subfolder="unet",
            ).to(dtype=self.weight_dtype, device="cuda")

            inference_config_path = self.config.inference_config
            infer_config = OmegaConf.load(inference_config_path)
            denoising_unet = UNet3DConditionModel.from_pretrained_2d(
                self.config.pretrained_base_model_path,
                self.config.motion_module_path,
                subfolder="unet",
                unet_additional_kwargs=infer_config.unet_additional_kwargs,
            ).to(dtype=self.weight_dtype, device="cuda")

            pose_guider = PoseGuider(320, block_out_channels=(16, 32, 96, 256)).to(
                dtype=self.weight_dtype, device="cuda"
            )

            image_enc = CLIPVisionModelWithProjection.from_pretrained(
                self.config.image_encoder_path
            ).to(dtype=self.weight_dtype, device="cuda")
            sched_kwargs = OmegaConf.to_container(infer_config.noise_scheduler_kwargs)
            scheduler = DDIMScheduler(**sched_kwargs)

            # load pretrained weights
            denoising_unet.load_state_dict(
                torch.load(self.config.denoising_unet_path, map_location="cpu"),
                strict=False,
            )
            reference_unet.load_state_dict(
                torch.load(self.config.reference_unet_path, map_location="cpu"),
            )
            pose_guider.load_state_dict(
                torch.load(self.config.pose_guider_path, map_location="cpu"),
            )

            pipe = Pose2VideoPipeline(
                vae=vae,
                image_encoder=image_enc,
                reference_unet=reference_unet,
                denoising_unet=denoising_unet,
                pose_guider=pose_guider,
                scheduler=scheduler,
            )
            pipe = pipe.to("cuda", dtype=self.weight_dtype)
            self.pipeline = pipe

        pose_images = read_frames(pose_video_path)
        src_fps = get_fps(pose_video_path)

        pose_list = []
        pose_tensor_list = []
        pose_transform = transforms.Compose(
            [transforms.Resize((height, width)), transforms.ToTensor()]
        )
        for pose_image_pil in pose_images[:length]:
            pose_list.append(pose_image_pil)
            pose_tensor_list.append(pose_transform(pose_image_pil))

        video = self.pipeline(
            ref_image,
            pose_list,
            width=width,
            height=height,
            video_length=length,
            num_inference_steps=num_inference_steps,
            guidance_scale=cfg,
            generator=generator,
        ).videos

        # ref_image_tensor = pose_transform(ref_image)  # (c, h, w)
        # ref_image_tensor = ref_image_tensor.unsqueeze(1).unsqueeze(0)  # (1, c, 1, h, w)
        # ref_image_tensor = repeat(
        #     ref_image_tensor, "b c f h w -> b c (repeat f) h w", repeat=length
        # )
        # pose_tensor = torch.stack(pose_tensor_list, dim=0)  # (f, c, h, w)
        # pose_tensor = pose_tensor.transpose(0, 1)
        # pose_tensor = pose_tensor.unsqueeze(0)
        # video = torch.cat([ref_image_tensor, pose_tensor, video], dim=0)

        # save_dir = f"./output/demo"
        if not os.path.exists(save_dir):
            os.makedirs(save_dir, exist_ok=True)
        date_str = datetime.now().strftime("%Y%m%d")
        time_str = datetime.now().strftime("%H%M%S")
        out_path = os.path.join(save_dir, f"{date_str}T{time_str}.mp4")

        # add
        im_ratio = float(ref_image.size[1]) / ref_image.size[0]

        save_videos_grid(
            video,
            out_path,
            n_rows=1,
            fps=src_fps,
            im_ratio=im_ratio
        )

        torch.cuda.empty_cache()

        # rvm video
        sys.path.append("RobustVideoMatting")
        from RobustVideoMatting.inference import Converter
        out_path_alpha = os.path.join(save_dir, f"{date_str}T{time_str}_alpha.mp4")
        if self.rvm is None:
            # self.rvm = Converter("mobilenetv3", "pretrained_weights/rvm_mobilenetv3.pth", "cuda")
            self.rvm = Converter("resnet50", "pretrained_weights/rvm_resnet50.pth", "cuda")
        self.rvm.convert(input_source=out_path, output_alpha=out_path_alpha, progress=True)

        # rvm image
        out_path_image = os.path.join(save_dir, f"{date_str}T{time_str}_image")
        if os.path.exists(out_path_image):
            os.removedirs(out_path_image)
        os.makedirs(out_path_image, exist_ok=True)
        out_path_image_src = os.path.join(out_path_image, "src.png")
        ref_image.save(out_path_image_src)
        self.rvm.convert(input_source=out_path_image, output_alpha=out_path_image, output_type="png_sequence",
                         progress=True)

        # lama image
        mask_image_path = os.path.join(out_path_image, "0000.png")
        from lama_demo import lama
        if self.lama is None:
            self.lama = lama("pretrained_weights/big-lama.pt")
        image = cv2.imread(out_path_image_src)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        mask = cv2.imread(mask_image_path, cv2.IMREAD_GRAYSCALE)
        out_lama = self.lama.process(image, mask)
        out_lama = cv2.cvtColor(out_lama, cv2.COLOR_RGB2BGR)
        out_path_image_dst = os.path.join(out_path_image, "dst.png")
        cv2.imwrite(out_path_image_dst, out_lama)

        # merge alpha video
        out_path_replace = os.path.join(save_dir, f"{date_str}T{time_str}_replace.mp4")
        apply_mask_and_replace_background(out_path, out_path_alpha, out_path_image_dst, out_path_replace)

        # add music
        out_path_with_audio = os.path.join(save_dir, f"{date_str}T{time_str}_with_audio.mp4")
        pose_video_name = os.path.basename(pose_video_path)
        audio_name = pose_video_name.replace("_kps.mp4", ".mp3")
        audio_path = "./test-video_audio"
        audio_path = os.path.join(audio_path, audio_name)
        save_video_with_audio(out_path_replace, audio_path, out_path_with_audio)
        return out_path_with_audio

def crop_image_center(image, desired_ratio):
    # Get original dimensions
    original_width, original_height = image.size

    # Determine the dimensions of the new crop area based on desired ratio
    if original_width / original_height > desired_ratio:
        # Width is larger than what we want, reduce it
        new_height = original_height
        new_width = int(desired_ratio * new_height)
    else:
        # Height is larger than what we want, reduce it
        new_width = original_width
        new_height = int(new_width / desired_ratio)

    # Calculate the top left corner position to start cropping
    left = (original_width - new_width) / 2
    top = (original_height - new_height) / 2
    right = (original_width + new_width) / 2
    bottom = (original_height + new_height) / 2

    # Perform the crop
    cropped_image = image.crop((left, top, right, bottom))

    return cropped_image


def batch_test(img_dir, video_dir, output_dir):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir, exist_ok=True)

    # hyperparameters
    width, height, length = 640, 960, 15 * 8
    ratio = width * 1.0 / height
    steps = 1
    controller = AnimateController()

    ref_img_path_list = os.listdir(img_dir)
    ref_img_path_list.sort()

    video_path_list = os.listdir(video_dir)
    video_path_list.sort()

    save_dir = f"./output/demo"
    for video_name in video_path_list:
        print(video_name)
        video_path = os.path.join(video_dir, video_name)
        output_path = os.path.join(output_dir, video_name.split(".mp4")[0])
        if not os.path.exists(output_path):
            os.makedirs(output_path, exist_ok=True)
        for i, ref_img_name in enumerate(ref_img_path_list):
            ref_img_path = os.path.join(img_dir, ref_img_name)
            ref_image = Image.open(ref_img_path)
            ref_image = ref_image.convert('RGB')
            cropped_image = crop_image_center(ref_image, ratio)
            print(i, ref_img_name)
            # run model
            result_video_path = controller.animate_v2(cropped_image, video_path, width, height, length, steps, 3.5, -1,
                                                   save_dir)
            import shutil
            # 拷贝文件
            shutil.copy(result_video_path, output_path)


def batch_test_old(img_dir, video_dir, output_dir):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir, exist_ok=True)

    # hyperparameters
    width, height, length = 640, 960, 15 * 8
    steps = 10
    controller = AnimateController()

    ref_img_path_list = os.listdir(img_dir)
    ref_img_path_list.sort()

    pose_video_path_list = os.listdir(video_dir)
    pose_video_path_list.sort()
    # pose_video_path_list = [pose_video_path_list[5]]

    save_dir = f"./output/demo"
    from tqdm import tqdm
    for pose_video_name in pose_video_path_list:
        print(pose_video_name)
        pose_video_path = os.path.join(video_dir, pose_video_name)
        output_path = os.path.join(output_dir, pose_video_name.split(".mp4")[0])
        if not os.path.exists(output_path):
            os.makedirs(output_path, exist_ok=True)
        for i, ref_img_name in enumerate(ref_img_path_list):
            ref_img_path = os.path.join(img_dir, ref_img_name)
            ref_image = Image.open(ref_img_path)
            ref_image = ref_image.convert('RGB')
            print(i, ref_img_name)
            # run model
            result_video_path = controller.animate(ref_image, pose_video_path, width, height, length, steps, 3.5, -1,
                                                   save_dir)
            import shutil
            # 拷贝文件
            shutil.copy(result_video_path, output_path)


if __name__ == "__main__":
    batch_test("./test-img-0326", "./test-video_15fps", "output/demo-0423")
    # batch_test_old("./test-img-0415", "./test-video-0326_kps", "output/demo-0328-3")

    # # img
    # ref_img_path = "test-img-0326/e原图.png"
    # # video
    # pose_video_path = "old_tongyi_video_10_kps.mp4"
    # # hyperparameters
    # width, height, length = 640, 960, 15 * 8
    # steps = 25
    # controller = AnimateController()
    # save_dir = f"./output/demo"
    # #
    # ref_image = cv2.imread(ref_img_path)
    # ref_image = cv2.cvtColor(ref_image, cv2.COLOR_BGR2RGB)
    # #
    # # run model
    # controller.animate(ref_image, pose_video_path, width, height, length, steps, 3.5, -1, save_dir)