SignerX/SignVerse-2M

1import os2import shutil3import subprocess4import av5import torch6import torch.distributed as dist7import torch.multiprocessing as mp8from torch.nn.parallel import DistributedDataParallel9from utils.util import get_fps, read_frames, save_videos_from_pil10from utils.preprocess_video import *11from PIL import Image12import numpy as np13import json14 15def ensure_dir(directory):16    if os.path.exists(directory):17        print(f"Directory already exists: {directory}")18    else:19        os.makedirs(directory)20        print(f"Created directory: {directory}")21    return directory22 23# [previous helper functions remain the same]24def get_video_dimensions(video_path):25    cmd = [26        'ffprobe',27        '-v', 'error',28        '-select_streams', 'v:0',29        '-show_entries', 'stream=width,height',30        '-of', 'csv=p=0',31        video_path32    ]33    result = subprocess.run(cmd, capture_output=True, text=True)34    width, height = map(int, result.stdout.strip().split(','))35    return width, height36 37def compile_frames_to_video(frame_dir, output_path, fps=30):38    """Compile frames into a video using H.264 codec."""39    cmd = [40        'ffmpeg', '-y',41        '-f', 'image2',42        '-r', str(fps),43        '-i', f'{frame_dir}/%08d.jpg',44        '-c:v', 'libx264',45        '-preset', 'medium',46        '-crf', '18',47        '-pix_fmt', 'yuv420p',48        output_path49    ]50    subprocess.run(cmd, check=True)51    print(f"Successfully compiled video: {output_path}")52 53def preprocess_videos(video_dir, dataset_name, square_crop=False, fps=24, quality_preset="medium", target_resolution=None):54    """55    Preprocess all videos with optional square cropping, customizable FPS, quality, and resolution.56    57    Args:58        video_dir (str): Directory containing input videos59        dataset_name (str): Name of the dataset60        square_crop (bool): Whether to crop videos to 1:1 aspect ratio (default: True)61        fps (int): Target frames per second (default: 30)62        quality_preset (str): Quality preset - "high", "medium", "low", "ultra_low" (default: "medium")63        target_resolution (int): Target resolution for the shorter side (e.g., 512, 256). None for original64    """65    result_dir = ensure_dir(f"../output/{dataset_name}_results")66 67    # 质量设置68    quality_settings = {69        "high": {"qscale": "2", "crf": "18"},      # 高质量70        "medium": {"qscale": "5", "crf": "23"},    # 中等质量71        "low": {"qscale": "10", "crf": "28"},      # 低质量72        "ultra_low": {"qscale": "15", "crf": "35"} # 超低质量73    }74    75    current_quality = quality_settings.get(quality_preset, quality_settings["medium"])76 77    for video_file in os.listdir(video_dir):78        if not video_file.endswith(".mp4"):79            continue80 81        video_name = os.path.splitext(video_file)[0]82        video_full_path = os.path.join(video_dir, video_file)83        folder_path = f"{result_dir}/{video_name}"84        85        frame_path = f"{folder_path}/crop_frame"86        output_video_path = f"{folder_path}/crop_original_video.mp4"87 88        # Skip if already processed89        if os.path.exists(frame_path) and os.listdir(frame_path) and os.path.exists(output_video_path):90            crop_status = "cropped" if square_crop else "original"91            print(f"{crop_status.capitalize()} frames and video already exist for {video_name}. Skipping preprocessing.")92            continue93 94        try:95            # Create output directory96            os.makedirs(frame_path, exist_ok=True)97 98            # Get video dimensions99            width, height = get_video_dimensions(video_full_path)100            101            # 构建视频滤镜102            filters = []103            104            if square_crop:105                # Calculate crop dimensions106                if width < height:107                    crop_size = width108                    x_offset = 0109                    y_offset = (height - width) // 2110                else:111                    crop_size = height112                    x_offset = (width - height) // 2113                    y_offset = 0114                filters.append(f'crop={crop_size}:{crop_size}:{x_offset}:{y_offset}')115            116            # 添加分辨率缩放117            if target_resolution:118                if square_crop:119                    # 方形裁剪后直接缩放到目标分辨率120                    filters.append(f'scale={target_resolution}:{target_resolution}')121                else:122                    # 保持宽高比缩放123                    filters.append(f'scale=-2:{target_resolution}:force_original_aspect_ratio=decrease')124            125            # 添加帧率126            filters.append(f'fps={fps}/1')127            128            # 组合所有滤镜129            filter_complex = ','.join(filters)130 131            # 提取帧的命令132            cmd = [133                'ffmpeg', '-i', video_full_path,134                '-vf', filter_complex,135                '-f', 'image2',136                '-qscale', current_quality["qscale"],  # 使用可调节的质量137                f'{frame_path}/%08d.jpg'138            ]139            140            resolution_info = f" (Resolution: {target_resolution})" if target_resolution else ""141            crop_info = "with square cropping" if square_crop else "without cropping"142            print(f"Processing {video_file} {crop_info} (FPS: {fps}, Quality: {quality_preset}{resolution_info})")143 144            subprocess.run(cmd, check=True)145            print(f"Successfully extracted frames for {video_file}")146 147            # Compile frames back into a video with optimized settings148            compile_frames_to_video_optimized(frame_path, output_video_path, fps, quality_preset)149 150        except Exception as e:151            print(f"Error preprocessing {video_file}: {str(e)}")152            continue153 154def compile_frames_to_video_optimized(frame_dir, output_path, fps=30, quality_preset="medium"):155    """Compile frames into a video with optimized quality settings."""156    157    # 质量设置 - CRF值（越高质量越低，文件越小）158    quality_crf = {159        "high": "18",160        "medium": "23", 161        "low": "28",162        "ultra_low": "35"163    }164    165    crf_value = quality_crf.get(quality_preset, "23")166    167    cmd = [168        'ffmpeg', '-y',169        '-f', 'image2',170        '-r', str(fps),171        '-i', f'{frame_dir}/%08d.jpg',172        '-c:v', 'libx264',173        '-preset', 'medium',  # 可以改为 'fast' 加速编码174        '-crf', crf_value,175        '-pix_fmt', 'yuv420p',176        output_path177    ]178    subprocess.run(cmd, check=True)179    print(f"Successfully compiled optimized video: {output_path} (Quality: {quality_preset})")180 181# 使用示例：182 183# 1. 保持原分辨率，降低质量184# preprocess_videos(video_dir, dataset_name, square_crop=True, fps=30, quality_preset="low")185 186# 2. 降低分辨率到512x512（方形裁剪）187# preprocess_videos(video_dir, dataset_name, square_crop=True, fps=30, quality_preset="medium", target_resolution=512)188 189# 3. 极度压缩：低分辨率 + 超低质量190# preprocess_videos(video_dir, dataset_name, square_crop=True, fps=30, quality_preset="ultra_low", target_resolution=256)191 192# 4. 不裁剪，但缩放到较小尺寸193# preprocess_videos(video_dir, dataset_name, square_crop=False, fps=30, quality_preset="low", target_resolution=480)194 195def process_npz_files(input_folder_path, output_folder_path):196    """197    Process all NPZ files in the specified folder and generate the required output format.198    199    Args:200        input_folder_path (str): Path to the folder containing NPZ files201        output_folder_path (str): Path where output files will be saved202    """203    # Get all NPZ files in the folder204    npz_files = sorted([f for f in os.listdir(input_folder_path) if f.endswith('.npz')])205    total_frames = len(npz_files)206    207    output = []208    209    for idx, npz_file in enumerate(npz_files):210        file_path = os.path.join(input_folder_path, npz_file)211        data = np.load(file_path, allow_pickle=True)212        213        # Process bodies data214        bodies = data['bodies']215        body_scores = data['body_scores'][0]216        217        # Process hands data218        hands = data['hands']219        hands_scores = data['hands_scores']220        221        # Process faces data222        faces = data['faces'][0]223        faces_scores = data['faces_scores'][0]224        225        # Convert coordinates to strings with space separation226        frame_data = []227        228        # Add body coordinates and scores229        for i in range(bodies.shape[0]):230            frame_data.extend([f"{bodies[i][0]:.8f}", f"{bodies[i][1]:.8f}"])231        for score in body_scores:232            frame_data.append(f"{score:.8f}")233            234        # Add hand coordinates and scores235        for hand in hands:236            for point in hand:237                frame_data.extend([f"{point[0]:.8f}", f"{point[1]:.8f}"])238        for hand_score in hands_scores:239            frame_data.extend([f"{score:.8f}" for score in hand_score])240            241        # Add face coordinates and scores242        for point in faces:243            frame_data.extend([f"{point[0]:.8f}", f"{point[1]:.8f}"])244        for score in faces_scores:245            frame_data.append(f"{score:.8f}")246            247        # Add frame count248        frame_count = idx / (total_frames - 1) if total_frames > 1 else 0249        frame_data.append(f"{frame_count:.8f}")250        251        # 验证这一帧的数据点数是否为385252        if len(frame_data) != 385:253            print(f"Warning: Frame {idx} in {input_folder_path} has {len(frame_data)} values instead of 385")254            continue  # 跳过这一帧255 256        # Join all data with spaces257        output.append(" ".join(frame_data))258    259    return " ".join(output) + "\n"