SCAIL-Pose/NLFPoseExtract/v1_process_pose.py

import os
import sys

# 动态添加项目根目录到 sys.path，这样就不需要 export PYTHONPATH
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)  # SCAIL_Pose 目录
if project_root not in sys.path:
    sys.path.insert(0, project_root)

import cv2
import torch
import pickle
import torchvision
import shutil
import glob
import random
from tqdm import tqdm   
import decord
from decord import VideoReader, cpu, gpu
from torchvision.transforms import ToPILImage
from PIL import Image
import numpy as np
import argparse
from NLFPoseExtract.nlf_render import render_nlf_as_images, collect_smpl_poses, shift_dwpose_according_to_nlf, p3d_single_p2d
from NLFPoseExtract.nlf_draw import intrinsic_matrix_from_field_of_view, process_data_to_COCO_format, p3d_to_p2d
from DWPoseProcess.dwpose import DWposeDetector
from concurrent.futures import ProcessPoolExecutor, as_completed
import multiprocessing
import traceback
from NLFPoseExtract.extract_nlfpose_batch import process_video_nlf
from NLFPoseExtract.reshape_utils_3d import reshapePool3d
try:
    import moviepy.editor as mpy
except:
    import moviepy as mpy
from torchvision.transforms.functional import center_crop, resize
from torchvision.transforms import InterpolationMode
import torchvision.transforms as TT
import copy
from NLFPoseExtract.align3d import solve_new_camera_params_central, solve_new_camera_params_down


def recollect_nlf(data):
    new_data = []
    for item in data:
        new_item = item.copy()
        if len(item['bboxes']) > 0:
            new_item['bboxes'] = item['bboxes'][:1]
            new_item['nlfpose'] = item['nlfpose'][:1]
        new_data.append(new_item)
    return new_data

def recollect_dwposes(poses):
    new_poses = []
    for pose in poses:
        new_pose = pose.copy()
        for i in range(1):
            bodies = pose["bodies"]
            faces = pose["faces"][i:i+1]
            hands = pose["hands"][2*i:2*i+2]
            candidate = bodies["candidate"][i:i+1]  # candidate是所有点的坐标和置信度
            subset = bodies["subset"][i:i+1]   # subset是认为的有效点
            new_pose = {
                "bodies": {
                    "candidate": candidate,
                    "subset": subset
                },
                "faces": faces,
                "hands": hands
            }
        new_poses.append(new_pose)
    return new_poses



def resize_for_rectangle_crop(arr, image_size, reshape_mode='random'):
    if arr.shape[3] / arr.shape[2] > image_size[1] / image_size[0]:
        arr = resize(arr, size=[image_size[0], int(arr.shape[3] * image_size[0] / arr.shape[2])], interpolation=InterpolationMode.BICUBIC)
    else:
        arr = resize(arr, size=[int(arr.shape[2] * image_size[1] / arr.shape[3]), image_size[1]], interpolation=InterpolationMode.BICUBIC)

    h, w = arr.shape[2], arr.shape[3]

    delta_h = h - image_size[0]
    delta_w = w - image_size[1]

    if reshape_mode == 'random' or reshape_mode == 'none':
        top = np.random.randint(0, delta_h + 1)
        left = np.random.randint(0, delta_w + 1)
    elif reshape_mode == 'center':
        top, left = delta_h // 2, delta_w // 2
    else:
        raise NotImplementedError
    arr = TT.functional.crop(
        arr, top=top, left=left, height=image_size[0], width=image_size[1]
    )
    return arr

def scale_faces(poses, pose_2d_ref):
    # 输入：两个list of dict，poses[0]['faces'].shape: 1, 68, 2  , poses_ref[0]['faces'].shape: 1, 68, 2
    # 根据脸部的中心点，对poses中的脸部关键点进行缩放
    # 也即：计算ref里面脸部中心点(idx: 30)到其他脸部关键点的中心距离， 计算poses里面脸部中心点到其他脸部关键点的中心距离，得到scale_n
    # 对scale_n 取一下0.8-1.5的上下界，然后应用在poses上
    # 注意：需要inplace改变poses

    ref = pose_2d_ref[0]
    pose_0 = poses[0]
        

    face_0 = pose_0['faces']  # shape: (1, 68, 2)
    face_ref = ref['faces']

    # 提取 numpy 数组
    face_0 = np.array(face_0[0])      # (68, 2)
    face_ref = np.array(face_ref[0])

    # 中心点（鼻尖或面部中心）
    center_idx = 30
    center_0 = face_0[center_idx]
    center_ref = face_ref[center_idx]

    # 计算到中心点的距离
    dist = np.linalg.norm(face_0 - center_0, axis=1)
    dist_ref = np.linalg.norm(face_ref - center_ref, axis=1)

    # 避免中心点自身的 0 距离影响
    dist = np.delete(dist, center_idx)
    dist_ref = np.delete(dist_ref, center_idx)

    mean_dist = np.mean(dist)
    mean_dist_ref = np.mean(dist_ref)

    if mean_dist < 1e-6:
        scale_n = 1.0
    else:
        scale_n = mean_dist_ref / mean_dist

    # 限制在 [0.8, 1.5]
    scale_n = np.clip(scale_n, 0.8, 1.5)

    for i, pose in enumerate(poses):
        face = pose['faces']
        # 提取 numpy 数组
        face = np.array(face[0])      # (68, 2)
        center = face[center_idx]
        scaled_face = (face - center) * scale_n + center
        poses[i]['faces'][0] = scaled_face

        body = pose['bodies']
        candidate = body['candidate']
        candidate_np = np.array(candidate[0])   # (14, 2)
        body_center = candidate_np[0]
        scaled_candidate = (candidate_np - body_center) * scale_n + body_center
        poses[i]['bodies']['candidate'][0] = scaled_candidate

    # inplace 修改
    pose['faces'][0] = scaled_face
    
    return scale_n


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Process video with NLF pose estimation')
    parser.add_argument('--subdir', type=str, default="../examples/001", help='Path to the subdirectory to process')
    parser.add_argument('--model_path', type=str, default='pretrained_weights/nlf_l_multi_0.3.2.torchscript', 
                        help='Path to NLF model')
    parser.add_argument('--use_align', action='store_true', help='Whether to use 2D keypoints from reference image for alignment')
    parser.add_argument('--resolution', type=int, nargs=2, default=[512, 896], 
                        metavar=('HEIGHT', 'WIDTH'),
                        help='Target resolution as [height, width], currently only [512, 896] are supported')
    args = parser.parse_args()
    
    subdir = args.subdir
    model_nlf = torch.jit.load(args.model_path).cuda().eval()
    decord.bridge.set_bridge("torch")

    # 设置路径
    mp4_path = os.path.join(subdir, 'driving.mp4')
    if not os.path.exists(mp4_path):
        raise FileNotFoundError(f"No video file found in {subdir}")
    
    if args.use_align:
        out_path_aligned = os.path.join(subdir, 'rendered_aligned.mp4')
    else:
        out_path_aligned = os.path.join(subdir, 'rendered.mp4')
    
    ref_image_path = os.path.join(subdir, 'ref_image.jpg')
    if not os.path.exists(ref_image_path):
        ref_image_path = os.path.join(subdir, 'ref_image.png')
    if not os.path.exists(ref_image_path):
        ref_image_path = os.path.join(subdir, 'ref.jpg')
    if not os.path.exists(ref_image_path):
        raise FileNotFoundError(f"No reference image found in {subdir}")

    print(f"Processing: {subdir}")
    print(f"Video: {mp4_path}")
    print(f"Reference: {ref_image_path}")
    print(f"Resolution: {args.resolution}")

    # 读取视频
    vr = VideoReader(mp4_path)
    vr_frames = vr.get_batch(list(range(len(vr))))   # T H W C
    sampling_image_size = args.resolution
    if vr_frames.shape[1] < vr_frames.shape[2]:
        target_H, target_W = sampling_image_size
    else:
        target_W, target_H = sampling_image_size
    vr_frames = resize_for_rectangle_crop(vr_frames.permute(0, 3, 1, 2), [target_H, target_W], reshape_mode='center').permute(0, 2, 3, 1)  # T H W C ->T C H W -> T H W C

    # 读取参考图片
    img_ref = cv2.imread(ref_image_path)
    img_ref = cv2.cvtColor(img_ref, cv2.COLOR_BGR2RGB)
    vr_frames_ref = torch.from_numpy(img_ref).unsqueeze(0)
    vr_frames_ref = resize_for_rectangle_crop(vr_frames_ref.permute(0, 3, 1, 2), [target_H, target_W], reshape_mode='center').permute(0, 2, 3, 1)  # 1 H W C ->1 C H W -> 1 H W C

    # 初始化检测器
    detector = DWposeDetector(use_batch=False).to(0)

    # 处理Driving视频
    print("Processing driving video...")
    detector_return_list = []
    pil_frames = []
    for i in tqdm(range(len(vr_frames)), desc="Detecting poses in video"):
        pil_frame = Image.fromarray(vr_frames[i].numpy())
        pil_frames.append(pil_frame)
        detector_result = detector(pil_frame)
        detector_return_list.append(detector_result)

    W, H = pil_frames[0].size
    poses, scores, det_results = zip(*detector_return_list)
    
    print("Running NLF on driving video...")
    nlf_results = process_video_nlf(model_nlf, vr_frames, det_results)

    # 处理ref图片
    print("Processing reference image...")
    detector_return_list_ref = []
    pil_frames_ref = []
    for i in range(len(vr_frames_ref)):
        pil_frame = Image.fromarray(vr_frames_ref[i].numpy())
        pil_frames_ref.append(pil_frame)
        detector_result = detector(pil_frame)
        detector_return_list_ref.append(detector_result)

    poses_ref, scores_ref, det_results_ref = zip(*detector_return_list_ref)
    
    print("Running NLF on reference image...")
    nlf_results_ref = process_video_nlf(model_nlf, vr_frames_ref, det_results_ref)

    # 进行对齐和渲染
    print("Aligning and rendering...")
    ori_camera_pose = intrinsic_matrix_from_field_of_view([target_H, target_W])
    ori_focal = ori_camera_pose[0, 0]
    pose_3d_first_driving_frame = nlf_results[0]['nlfpose'][0][0].cpu().numpy()  # 3D点 frame-idx bbox-idx detect-idx
    pose_3d_coco_first_driving_frame = process_data_to_COCO_format(pose_3d_first_driving_frame)

    if args.use_align:
        poses_2d_ref = poses_ref[0]['bodies']['candidate'][0][:14]
        poses_2d_ref[:, 0] = poses_2d_ref[:, 0] * target_W
        poses_2d_ref[:, 1] = poses_2d_ref[:, 1] * target_H

        poses_2d_subset = poses_ref[0]['bodies']['subset'][0][:14]
        pose_3d_coco_first_driving_frame = pose_3d_coco_first_driving_frame[:14]

        valid_indices = []
        valid_upper_indices = []
        valid_lower_indices = []
        upper_body_indices = [0, 2, 3, 5, 6]
        lower_body_indices = [9, 10, 12, 13]
        excluded_indices = [3, 4, 6, 7]  # 去除手
        for i in range(len(poses_2d_subset)):
            if poses_2d_subset[i] != -1.0 and np.sum(pose_3d_coco_first_driving_frame[i]) != 0:
                if i in upper_body_indices:
                    valid_upper_indices.append(i)
                if i in lower_body_indices:
                    valid_lower_indices.append(i)
        
        if len(valid_lower_indices) >= 4:
            print("Align feet")
            valid_indices = [1] + valid_lower_indices
        else:
            print("Align body")
            valid_indices = [1] + valid_upper_indices

        pose_2d_ref = poses_2d_ref[valid_indices]
        pose_3d_coco_first_driving_frame = pose_3d_coco_first_driving_frame[valid_indices]
        
        if len(valid_lower_indices) >= 4:
            new_camera_intrinsics, scale_m, scale_s = solve_new_camera_params_down(pose_3d_coco_first_driving_frame, ori_focal, [target_H, target_W], pose_2d_ref)
        else:
            new_camera_intrinsics, scale_m, scale_s = solve_new_camera_params_central(pose_3d_coco_first_driving_frame, ori_focal, [target_H, target_W], pose_2d_ref)
        
        # m 代表缩放了多少
        scale_face = scale_faces(list(poses), list(poses_ref))   # poses[0]['faces'].shape: 1, 68, 2  , poses_ref[0]['faces'].shape: 1, 68, 2

        print(f"Scale - m: {scale_m}, face: {scale_face}")

        nlf_results = recollect_nlf(nlf_results)
        poses = recollect_dwposes(list(poses))
        shift_dwpose_according_to_nlf(collect_smpl_poses(nlf_results), poses, ori_camera_pose, new_camera_intrinsics, target_H, target_W, scale_x=scale_m, scale_y=scale_m*scale_s)
        
        print("Rendering final video...")
        frames_np = render_nlf_as_images(nlf_results, poses, reshape_pool=None, intrinsic_matrix=new_camera_intrinsics)

    else:
        nlf_results = recollect_nlf(nlf_results)
        print("Rendering final video...")
        frames_np = render_nlf_as_images(nlf_results, poses, reshape_pool=None, intrinsic_matrix=ori_camera_pose)

    mpy.ImageSequenceClip(frames_np, fps=16).write_videofile(out_path_aligned)
    print(f"Done! Output saved to: {out_path_aligned}")