AniTalker项目提交

AniTalker-kit/AniTalker-judge/Dockerfile

@@ -0,0 +1,46 @@
+# 使用nvidia的CUDA基础镜像
+FROM nvidia/cuda:11.7.1-runtime-ubuntu22.04
+
+# 设置工作目录
+WORKDIR /app
+
+# 安装依赖，同时设置非交互式前端和时区
+RUN apt-get update && apt-get install -y \
+    software-properties-common \
+    build-essential \
+    libgl1-mesa-glx \
+    && DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC add-apt-repository ppa:deadsnakes/ppa \
+    && apt-get update \
+    && DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get install -y \
+    python3.9 \
+    python3.9-distutils \
+    python3-pip \
+    && rm -rf /var/lib/apt/lists/*
+
+# 升级pip到最新版本，确保使用的是python3.9的pip
+RUN python3.9 -m pip install --upgrade pip
+
+# 设置Python版本为3.9
+RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.9 1
+
+
+# 安装Python依赖
+COPY requirements.txt /app/requirements.txt
+RUN pip install --no-cache-dir -r requirements.txt
+
+
+# 将Python脚本复制到容器中
+COPY PSNR_SSIM_FID_NIQE.py /app/PSNR_SSIM_FID_NIQE.py
+
+COPY inception_v3_google-0cc3c7bd.pth /app/inception_v3_google-0cc3c7bd.pth
+COPY niqe_modelparameters.mat /root/.cache/torch/hub/pyiqa/niqe_modelparameters.mat
+
+# 使脚本可执行
+RUN chmod +x /app/PSNR_SSIM_FID_NIQE.py
+
+# 设置CUDA_HOME环境变量
+ENV CUDA_HOME=/usr/local/cuda
+
+# 运行脚本
+ENTRYPOINT ["python3.9", "/app/PSNR_SSIM_FID_NIQE.py"]
+

AniTalker-kit/AniTalker-judge/PSNR_SSIM_FID_NIQE.py

@@ -0,0 +1,199 @@
+import imageio
+import torch
+from pytorch_msssim import ssim
+import numpy as np
+from PIL import Image
+import cv2
+import pyiqa
+from torchvision.models import inception_v3
+from torchvision import transforms
+import scipy.linalg as linalg
+import warnings
+import argparse
+import os
+
+# 忽略警告
+warnings.filterwarnings('ignore')
+
+# 加载视频帧
+def load_video_frames(video_path):
+    reader = imageio.get_reader(video_path)
+    frames = [frame for frame in reader]
+    reader.close()  # 关闭reader释放资源
+    return frames
+
+def resize_video(generated_frames, standard_size):
+    resized_generated_frames = [cv2.resize(frame, (standard_size[1], standard_size[0])) for frame in generated_frames]
+    return resized_generated_frames
+
+# 计算PSNR
+def calculate_psnr(standard_frames, generated_frames):
+    psnr_values = []
+    for i in range(len(standard_frames)):
+        mse = np.mean((standard_frames[i] - generated_frames[i]) ** 2)
+        if mse == 0:
+            psnr = float('inf')  # 对于完全相同的图像，PSNR为无穷大
+        else:
+            max_pixel = 255.0
+            psnr = 20 * np.log10(max_pixel / np.sqrt(mse))
+        psnr_values.append(psnr)
+    return np.mean(psnr_values)
+
+# 计算SSIM
+def calculate_ssim(standard_frames, generated_frames):
+    ssim_values = []
+    for i in range(len(standard_frames)):
+        # 规范化帧数据到[0, 1]并转换为float32
+        standard_frame = torch.tensor(standard_frames[i]).permute(2, 0, 1).unsqueeze(0).float() / 255.0
+        generated_frame = torch.tensor(generated_frames[i]).permute(2, 0, 1).unsqueeze(0).float() / 255.0
+        ssim_val = ssim(standard_frame, generated_frame, data_range=1.0, size_average=True)  # 正确的data_range
+        ssim_values.append(ssim_val.item())
+    return np.mean(ssim_values)
+
+# 定义一个函数来提取视频帧的特征，使用批量处理
+def extract_features(video_path, model, device, batch_size=16):
+    frames = []
+    cap = cv2.VideoCapture(video_path)
+    batch = []
+
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+        img = transforms.ToTensor()(img).unsqueeze(0).to(device)
+        batch.append(img)
+
+        if len(batch) == batch_size:
+            with torch.no_grad():
+                batch_tensor = torch.cat(batch, dim=0)
+                features = model(batch_tensor)
+                frames.extend(features.cpu().numpy())
+            batch = []  # 重置批量
+
+    # 处理剩余的帧
+    if batch:
+        with torch.no_grad():
+            batch_tensor = torch.cat(batch, dim=0)
+            features = model(batch_tensor)
+            frames.extend(features.cpu().numpy())
+
+    cap.release()  # 释放视频捕获资源
+    return frames
+
+# 计算每个视频特征的均值和协方差
+def calculate_statistics(features):
+    mean = np.mean(features, axis=0)
+    cov = np.cov(features, rowvar=False)
+    return mean, cov
+
+# 使用FID公式计算两个视频特征分布之间的FID分数
+def calculate_fid(standard_mean, standard_cov, generated_mean, generated_cov):
+    # 计算均值差
+    mean_diff = standard_mean - generated_mean
+
+    # 为了确保协方差矩阵正定，添加一个小的正则化项
+    epsilon = 1e-6
+    standard_cov += epsilon * np.eye(standard_cov.shape[0])
+    generated_cov += epsilon * np.eye(generated_cov.shape[0])
+
+    # 计算协方差矩阵的平方根
+    cov_sqrt = linalg.sqrtm(standard_cov @ generated_cov)
+
+    # 如果平方根返回复数，取其实部
+    if np.iscomplexobj(cov_sqrt):
+        cov_sqrt = cov_sqrt.real
+
+    # 计算FID
+    fid = np.sum(mean_diff ** 2) + np.trace(standard_cov + generated_cov - 2 * cov_sqrt)
+
+    # 检查是否出现NaN或负值
+    if np.isnan(fid) or fid < 0:
+        fid = 0  # 或者返回一个合理的默认值
+    return fid
+
+# 计算NIQE得分
+def calculate_niqe_score(frames):
+    niqe_metric = pyiqa.create_metric('niqe')
+    scores = []
+    for frame in frames:
+        frame_tensor = torch.from_numpy(frame).permute(2, 0, 1).unsqueeze(0).float() / 255.0  # 规范化到[0, 1]
+        score = niqe_metric(frame_tensor)
+        scores.append(score)
+    del niqe_metric  # 删除NIQE指标以释放资源
+    return torch.tensor(scores).mean().item()
+
+# 定义评价函数
+def evaluate_video(video_name, standard_path, generated_path, device):
+    path_standard_video = os.path.join(standard_path, video_name)
+    path_generated_video = os.path.join(generated_path, video_name)
+
+    standard_frames = load_video_frames(path_standard_video)
+    generated_frames = load_video_frames(path_generated_video)
+
+    if generated_frames[0].shape != standard_frames[0].shape:
+        generated_frames = resize_video(generated_frames, standard_frames[0].shape)
+
+    psnr_score = round(calculate_psnr(standard_frames, generated_frames), 6)
+    ssim_score = round(calculate_ssim(standard_frames, generated_frames), 6)
+
+    standard_features = extract_features(path_standard_video, inception_model, device)
+    generated_features = extract_features(path_generated_video, inception_model, device)
+
+    standard_mean, standard_cov = calculate_statistics(standard_features)
+    generated_mean, generated_cov = calculate_statistics(generated_features)
+
+    fid_score = round(calculate_fid(standard_mean, standard_cov, generated_mean, generated_cov), 6)
+    generated_niqe = round(calculate_niqe_score(generated_frames), 6)
+
+    return psnr_score, ssim_score, fid_score, generated_niqe
+
+# 命令行参数解析
+parser = argparse.ArgumentParser(description='Evaluate videos using various metrics.')
+parser.add_argument('stand_path', type=str, help='Path to the standard videos folder')
+parser.add_argument('generate_path', type=str, help='Path to the generated videos folder')
+parser.add_argument('--device', type=str, default='cuda:0', help='Device to use for computations (default: cuda:0)')
+args = parser.parse_args()
+
+# 视频路径
+standard_videos_path = args.stand_path
+generated_videos_path = args.generate_path
+device = torch.device(args.device)
+
+
+# 检查是否有可用的 GPU
+if device.type == 'cuda' and not torch.cuda.is_available():
+    raise ValueError("CUDA is not available, but device set to cuda. Please check your device settings.")
+
+
+# 加载预训练的Inception模型
+
+inception_model = inception_v3(pretrained=False)
+inception_model.load_state_dict(torch.load('/app/inception_v3_google-0cc3c7bd.pth'))
+inception_model.to(device).eval()
+
+
+# 视频列表
+standard_video_names = os.listdir(standard_videos_path)
+generated_video_names = os.listdir(generated_videos_path)
+
+
+# 获取两个列表的交集
+video_names = list(set(standard_video_names) & set(generated_video_names))
+
+
+# 输出指标名作为表头，设置列宽
+header = "Video Name".ljust(20) + "PSNR Score".ljust(12) + "SSIM Score".ljust(12) + "FID Score".ljust(12) + "NIQE Score".ljust(12)
+print(header)
+
+# 循环处理每个视频
+for video_name in video_names:
+    psnr_score, ssim_score, fid_score, niqe_score = evaluate_video(video_name, standard_videos_path, generated_videos_path, device)
+
+    # 格式化输出，确保列对齐
+    row = f"{video_name.ljust(20)}{psnr_score:<12}{ssim_score:<12}{fid_score:<12}{niqe_score:<12}"
+    print(row)
+
+    # 释放不再需要的资源
+    torch.cuda.empty_cache()  # 释放GPU缓存

AniTalker-kit/AniTalker-judge/README.md

@@ -0,0 +1,86 @@
+# Judge
+
+这一部分用于评价指标PSNR、SSIM、FID、NIQE的计算。
+
+仓库中提供了代码、样例和Dockerfile文件，其中缺少了模型参数文件inception_v3_google-0cc3c7bd.pth，可以在[Google Drive](https://drive.google.com/file/d/1urWE3mUroo2rOn-A2nBXWhuKQ2I1M2zd/view?usp=drive_link)或[下载链接](https://download.pytorch.org/models/inception_v3_google-0cc3c7bd.pth)下载。
+
+## Quick Start
+
+1. 安装docker，宿主机CUDA版本为11.7及以上
+2. 拉取镜像：
+
+```
+docker pull gre123/anitalkerjudge:v1
+```
+
+3. 参考以下docker run命令运行：
+
+```
+docker run --rm --gpus all \
+-v your_stand_videos_dir_path:/app/stand_videos_dir_path \
+-v your_generate_videos_dir_path:/app/generate_videos_dir_path \
+gre123/anitalkerjudge:v1 \
+/app/stand_videos_dir_path \
+/app/generate_videos_dir_path
+[--device your_device]（默认值为cuda:0）
+```
+
+其中，your_stand_videos_dir_path、your_generate_videos_dir_path为文件夹的绝对路径。
+注意：评测程序会对两个文件夹中名称相同的视频进行评测指标的计算，运行该程序需要确保对应的参照视频和生成视频名称相同且时长相同。
+
+## Install
+
+如果想要自己手动构建镜像，请按照下列步骤进行：
+
+1. 将项目代码拉取到本地
+
+```
+git clone https://github.com/RubyZh/talkingface-kit.git
+cd talkingface-kit/AniTalker-kit/AniTalker-judge
+```
+
+2. 下载缺少的模型参数文件inception_v3_google-0cc3c7bd.pth，放在根目录下。（可以在[Google Drive](https://drive.google.com/file/d/1urWE3mUroo2rOn-A2nBXWhuKQ2I1M2zd/view?usp=drive_link)或[下载链接](https://download.pytorch.org/models/inception_v3_google-0cc3c7bd.pth)下载）
+3. 在项目根目录下打开终端，运行以下命令：
+
+```
+docker build --rm -f "Dockerfile" -t <image-name>:<tag> "."
+```
+
+如果ubuntu镜像无法拉取，可尝试先单独拉取对应镜像：
+
+```
+docker pull nvidia/cuda:11.7.1-runtime-ubuntu22.04
+```
+
+3. 构建成功后，参考上面的docker run命令运行。
+
+## Run
+
+如果想要直接在本地运行，请按照下列步骤进行：
+
+1. 将项目代码拉取到本地
+
+```
+git clone https://github.com/RubyZh/talkingface-kit.git
+cd talkingface-kit/AniTalker-kit/AniTalker-judge
+```
+
+2. 下载缺少的模型参数文件inception_v3_google-0cc3c7bd.pth，放在根目录下。（可以在[Google Drive](https://drive.google.com/file/d/1urWE3mUroo2rOn-A2nBXWhuKQ2I1M2zd/view?usp=drive_link)或[下载链接](https://download.pytorch.org/models/inception_v3_google-0cc3c7bd.pth)下载）
+3. 打开Anaconda Prompt Shell，创建conda环境：
+
+```
+conda create -n AniTalker_judge python=3.9.21
+conda activate AniTalker_judge
+```
+
+3. 安装依赖：
+
+```
+pip install -r requirements.txt
+```
+
+4. 将视频放到项目文件夹下，在项目根目录下运行：
+
+```
+python PSNR_SSIM_FID_NIQE.py /path/to/raw/video/dir /path/to/generate/video/dir
+```

AniTalker-kit/AniTalker-judge/requirements.txt

@@ -0,0 +1,13 @@
+# python==3.9.21
+imageio[ffmpeg]==2.36.1
+torch==2.0.1 
+torchvision==0.15.2 
+torchaudio==2.0.2
+pytorch_msssim==1.0.0
+pillow==11.0.0
+opencv-python-headless==4.10.0.84
+pyiqa==0.1.13
+scipy==1.13.1
+ipykernel==6.29.5
+ipywidgets==8.1.5
+numpy==1.23.0

AniTalker-kit/AniTalker-judge/run_judge.txt

@@ -0,0 +1,19 @@
+首先拉取镜像：
+docker pull gre123/anitalkerjudge:v1
+
+接着运行docker run命令：
+docker run --rm --gpus all \
+-v your_stand_videos_dir_path:/app/stand_videos_dir_path \
+-v your_generate_videos_dir_path:/app/generate_videos_dir_path \
+gre123/anitalkerjudge:v1 \
+/app/stand_videos_dir_path \
+/app/generate_videos_dir_path
+[--device your_device]（默认值为cuda:0）
+
+其中，your_stand_videos_dir_path、your_generate_videos_dir_path为文件夹的绝对路径。
+
+注意：评测程序会对两个文件夹中名称相同的视频进行评测指标的计算，运行该程序需要确保对应的参照视频和生成视频名称相同且时长相同。
+
+例：
+
+docker run --rm --gpus all --memory="64g" --cpus="4" -v F:/AniTalker/AniTalker-judge/demo/raw:/app/demo/raw -v F:/AniTalker/AniTalker-judge/demo/synthetic:/app/demo/synthetic gre123/anitalkerjudge:v1 /app/demo/raw /app/demo/synthetic

AniTalker-kit/AniTalker/.gitignore

@@ -0,0 +1,16 @@
+.DS_Store
+README.md.bak
+__pycache__
+results
+*.ckpt
+ckpts
+audios_hubert
+scripts
+*TEMP*
+gfpgan
+enhance_face_test.py
+checkpoints
+tmp
+train_stage2.py
+espnet
+*.dat