benchmark/app_video.py

import argparse
import os
import uuid

import cv2
import gradio as gr
import kornia
import numpy as np
import torch
from loguru import logger
from torchaudio.io import StreamReader
from torchaudio.io import StreamWriter

from benchmark.face_pipeline import alignFace
from benchmark.face_pipeline import FaceDetector
from benchmark.face_pipeline import inverse_transform_batch
from benchmark.face_pipeline import SoftErosion
from configs.train_config import TrainConfig
from models.model import HifiFace


class VideoSwap:
    def __init__(self, cfg, model=None):
        self.facedetector = FaceDetector(cfg.face_detector_weights)
        self.alignface = alignFace()
        self.work_dir = "."
        opt = TrainConfig()
        opt.use_ddp = False
        self.device = "cuda"
        self.ffmpeg_device = cfg.ffmpeg_device
        self.num_frames = 10
        self.kps_window = []
        checkpoint = (cfg.model_path, cfg.model_idx)
        if model is None:
            self.model = HifiFace(
                opt.identity_extractor_config, is_training=False, device=self.device, load_checkpoint=checkpoint
            )
        else:
            self.model = model
        self.model.eval()
        os.makedirs(self.work_dir, exist_ok=True)
        uid = uuid.uuid4()
        self.swapped_video = os.path.join(self.work_dir, f"tmp_{uid}.mp4")

        # model-idx_image-name_target-video-name.mp4
        swapped_with_audio_name = f"result_{uid}.mp4"

        # 带有音频的换脸视频
        self.swapped_video_with_audio = os.path.join(self.work_dir, swapped_with_audio_name)

        self.smooth_mask = SoftErosion(kernel_size=7, threshold=0.9, iterations=7).to(self.device)

    def yuv_to_rgb(self, img):
        img = img.to(torch.float)
        y = img[..., 0, :, :]
        u = img[..., 1, :, :]
        v = img[..., 2, :, :]
        y /= 255

        u = u / 255 - 0.5
        v = v / 255 - 0.5

        r = y + 1.14 * v
        g = y + -0.396 * u - 0.581 * v
        b = y + 2.029 * u

        rgb = torch.stack([r, g, b], -1)
        return rgb

    def rgb_to_yuv(self, img):
        r = img[..., 0, :, :]
        g = img[..., 1, :, :]
        b = img[..., 2, :, :]
        y = (0.299 * r + 0.587 * g + 0.114 * b) * 255
        u = (-0.1471 * r - 0.2889 * g + 0.4360 * b) * 255 + 128
        v = (0.6149 * r - 0.5149 * g - 0.1 * b) * 255 + 128
        yuv = torch.stack([y, u, v], -1)
        return torch.clamp(yuv, 0.0, 255.0, out=None).type(dtype=torch.uint8).transpose(3, 2).transpose(2, 1)

    def _geometry_transfrom_warp_affine(self, swapped_image, inv_att_transforms, frame_size, square_mask):
        swapped_image = kornia.geometry.transform.warp_affine(
            swapped_image,
            inv_att_transforms,
            frame_size,
            mode="bilinear",
            padding_mode="border",
            align_corners=True,
            fill_value=torch.zeros(3),
        )

        square_mask = kornia.geometry.transform.warp_affine(
            square_mask,
            inv_att_transforms,
            frame_size,
            mode="bilinear",
            padding_mode="zeros",
            align_corners=True,
            fill_value=torch.zeros(3),
        )
        return swapped_image, square_mask

    def smooth_kps(self, kps):
        self.kps_window.append(kps.flatten())
        self.kps_window = self.kps_window[1:]
        X = np.stack(self.kps_window, axis=1)
        y = self.kps_window[-1]
        y_cor = X @ np.linalg.inv(X.transpose() @ X - 0.0007 * np.eye(self.num_frames)) @ X.transpose() @ y
        self.kps_window[-1] = y_cor
        return y_cor.reshape((5, 2))

    def detect_and_align(self, image, src_is=False):
        detection = self.facedetector(image)
        if detection.score is None:
            self.kps_window = []
            return None, None
        max_score_ind = np.argmax(detection.score, axis=0)
        kps = detection.key_points[max_score_ind]
        if len(self.kps_window) < self.num_frames:
            self.kps_window.append(kps.flatten())
        else:
            kps = self.smooth_kps(kps)
        align_img, warp_mat = self.alignface.align_face(image, kps, 256)
        align_img = cv2.resize(align_img, (256, 256))
        align_img = align_img.transpose(2, 0, 1)
        align_img = torch.from_numpy(align_img).unsqueeze(0).to(self.device).float()
        align_img = align_img / 255.0
        if src_is:
            self.kps_window = []
        return align_img, warp_mat

    def inference(self, source_face, target_video, shape_rate, id_rate, iterations=1):
        video = cv2.VideoCapture(target_video)
        # 获取视频宽度
        frame_width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
        # 获取视频高度
        frame_height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
        # 获取帧率
        frame_rate = int(video.get(cv2.CAP_PROP_FPS))
        video.release()
        self.frame_size = (frame_height, frame_width)
        if self.ffmpeg_device == "cuda":
            self.decode_config = {"frames_per_chunk": 1, "decoder": "h264", "format": "yuv444p"}
            # self.decode_config = {
            #     "frames_per_chunk": 1,
            #     "decoder": "h264_cuvid",
            #     "decoder_option": {"gpu": "0"},
            #     "hw_accel": "cuda:0",
            # }

            self.encode_config = {
                "encoder": "h264_nvenc",  # GPU Encoder
                "encoder_format": "yuv444p",
                "encoder_option": {"gpu": "0", "cq": "10"},  # Run encoding on the cuda:0 device
                "hw_accel": "cuda:0",  # Data comes from cuda:0 device
                "frame_rate": frame_rate,
                "height": frame_height,
                "width": frame_width,
                "format": "yuv444p",
            }
        else:
            self.decode_config = {"frames_per_chunk": 1, "decoder": "h264", "format": "yuv444p"}

            self.encode_config = {
                "encoder": "libx264",
                "encoder_format": "yuv444p",
                "frame_rate": frame_rate,
                "height": frame_height,
                "width": frame_width,
                "format": "yuv444p",
            }
        src = source_face
        src, _ = self.detect_and_align(src, src_is=True)
        logger.info("start swapping")
        sr = StreamReader(target_video)
        if self.ffmpeg_device == "cpu":
            sr.add_basic_video_stream(**self.decode_config)
        else:
            sr.add_basic_video_stream(**self.decode_config)
            # sr.add_video_stream(**self.decode_config)
        sw = StreamWriter(self.swapped_video)
        sw.add_video_stream(**self.encode_config)
        with sw.open():
            for (chunk,) in sr.stream():
                # StreamReader cuda decode颜色格式默认为yuv需要转为rgb
                chunk = self.yuv_to_rgb(chunk)
                image = (chunk * 255).clamp(0, 255).to(torch.uint8)[0].cpu().numpy()
                align_img, warp_mat = self.detect_and_align(image)
                chunk = chunk.transpose(3, 2).transpose(2, 1).to(self.device)
                if align_img is None:
                    result_face = chunk
                else:
                    with torch.no_grad():
                        for _ in range(iterations):
                            swapped_face, m_r = self.model.forward(src, align_img, shape_rate, id_rate)
                            swapped_face = torch.clamp(swapped_face, 0, 1)
                            align_img = swapped_face
                        smooth_face_mask, _ = self.smooth_mask(m_r)
                    warp_mat = torch.from_numpy(warp_mat).float().unsqueeze(0)
                    inverse_warp_mat = inverse_transform_batch(warp_mat)
                    swapped_face, smooth_face_mask = self._geometry_transfrom_warp_affine(
                        swapped_face, inverse_warp_mat, self.frame_size, smooth_face_mask
                    )
                    result_face = (1 - smooth_face_mask) * chunk + smooth_face_mask * swapped_face
                result_face = self.rgb_to_yuv(result_face).to(self.ffmpeg_device)
                sw.write_video_chunk(0, result_face)

        # 将target_video中的音频转移到换脸视频上
        command = f"ffmpeg -loglevel error -i {self.swapped_video} -i {target_video} -c copy \
            -map 0 -map 1:1? -y -shortest {self.swapped_video_with_audio}"
        os.system(command)

        # 删除没有音频的换脸视频
        os.system(f"rm {self.swapped_video}")
        return self.swapped_video_with_audio


class ConfigPath:
    face_detector_weights = "/mnt/c/yangguo/useful_ckpt/face_detector/face_detector_scrfd_10g_bnkps.onnx"
    model_path = ""
    model_idx = 80000
    ffmpeg_device = "cuda"


def main():
    cfg = ConfigPath()
    parser = argparse.ArgumentParser(
        prog="benchmark", description="What the program does", epilog="Text at the bottom of help"
    )
    parser.add_argument("-m", "--model_path")
    parser.add_argument("-i", "--model_idx")
    parser.add_argument("-f", "--ffmpeg_device")

    args = parser.parse_args()

    cfg.model_path = args.model_path
    cfg.model_idx = int(args.model_idx)
    cfg.ffmpeg_device = args.ffmpeg_device

    infer = VideoSwap(cfg)

    def inference(source_face, target_video, shape_rate, id_rate):
        return infer.inference(source_face, target_video, shape_rate, id_rate)

    output = gr.Video(value=None, label="换脸结果")
    demo = gr.Interface(
        fn=inference,
        inputs=[
            gr.Image(shape=None, label="选脸图"),
            gr.Video(value=None, label="目标视频"),
            gr.Slider(
                minimum=0.0,
                maximum=1.0,
                value=1.0,
                step=0.1,
                label="3d结构相似度（1.0表示完全替换）",
            ),
            gr.Slider(
                minimum=0.0,
                maximum=1.0,
                value=1.0,
                step=0.1,
                label="人脸特征相似度（1.0表示完全替换）",
            ),
        ],
        outputs=output,
        title="HiConFace视频人脸融合系统",
        description="v1.0: developed by yiwise CV group",
    )
    demo.launch(server_name="0.0.0.0", server_port=7860)


if __name__ == "__main__":
    main()