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 PIL import Image11import numpy as np12import json13 14def ensure_dir(directory):15    if os.path.exists(directory):16        print(f"Directory already exists: {directory}")17    else:18        os.makedirs(directory)19        print(f"Created directory: {directory}")20    return directory21 22# [previous helper functions remain the same]23def get_video_dimensions(video_path):24    cmd = [25        'ffprobe',26        '-v', 'error',27        '-select_streams', 'v:0',28        '-show_entries', 'stream=width,height',29        '-of', 'csv=p=0',30        video_path31    ]32    result = subprocess.run(cmd, capture_output=True, text=True)33    width, height = map(int, result.stdout.strip().split(','))34    return width, height35 36def compile_frames_to_video(frame_dir, output_path, fps=30):37    """Compile frames into a video using H.264 codec."""38    cmd = [39        'ffmpeg', '-y',40        '-f', 'image2',41        '-r', str(fps),42        '-i', f'{frame_dir}/%08d.jpg',43        '-c:v', 'libx264',44        '-preset', 'medium',45        '-crf', '18',46        '-pix_fmt', 'yuv420p',47        output_path48    ]49    subprocess.run(cmd, check=True)50    print(f"Successfully compiled video: {output_path}")51 52def preprocess_videos(video_dir, dataset_name, square_crop=False, fps=24, quality_preset="medium", target_resolution=None):53    """54    Preprocess all videos with optional square cropping, customizable FPS, quality, and resolution.55    56    Args:57        video_dir (str): Directory containing input videos58        dataset_name (str): Name of the dataset59        square_crop (bool): Whether to crop videos to 1:1 aspect ratio (default: True)60        fps (int): Target frames per second (default: 30)61        quality_preset (str): Quality preset - "high", "medium", "low", "ultra_low" (default: "medium")62        target_resolution (int): Target resolution for the shorter side (e.g., 512, 256). None for original63    """64    result_dir = ensure_dir(f"../output/{dataset_name}_results")65 66    # 质量设置67    quality_settings = {68        "high": {"qscale": "2", "crf": "18"},      # 高质量69        "medium": {"qscale": "5", "crf": "23"},    # 中等质量70        "low": {"qscale": "10", "crf": "28"},      # 低质量71        "ultra_low": {"qscale": "15", "crf": "35"} # 超低质量72    }73    74    current_quality = quality_settings.get(quality_preset, quality_settings["medium"])75 76    for video_file in os.listdir(video_dir):77        if not video_file.endswith(".mp4"):78            continue79 80        video_name = os.path.splitext(video_file)[0]81        video_full_path = os.path.join(video_dir, video_file)82        folder_path = f"{result_dir}/{video_name}"83        84        frame_path = f"{folder_path}/crop_frame"85        output_video_path = f"{folder_path}/crop_original_video.mp4"86 87        # Skip if already processed88        if os.path.exists(frame_path) and os.listdir(frame_path) and os.path.exists(output_video_path):89            crop_status = "cropped" if square_crop else "original"90            print(f"{crop_status.capitalize()} frames and video already exist for {video_name}. Skipping preprocessing.")91            continue92 93        try:94            # Create output directory95            os.makedirs(frame_path, exist_ok=True)96 97            # Get video dimensions98            width, height = get_video_dimensions(video_full_path)99            100            # 构建视频滤镜101            filters = []102            103            if square_crop:104                # Calculate crop dimensions105                if width < height:106                    crop_size = width107                    x_offset = 0108                    y_offset = (height - width) // 2109                else:110                    crop_size = height111                    x_offset = (width - height) // 2112                    y_offset = 0113                filters.append(f'crop={crop_size}:{crop_size}:{x_offset}:{y_offset}')114            115            # 添加分辨率缩放116            if target_resolution:117                if square_crop:118                    # 方形裁剪后直接缩放到目标分辨率119                    filters.append(f'scale={target_resolution}:{target_resolution}')120                else:121                    # 保持宽高比缩放122                    filters.append(f'scale=-2:{target_resolution}:force_original_aspect_ratio=decrease')123            124            # 添加帧率125            filters.append(f'fps={fps}/1')126            127            # 组合所有滤镜128            filter_complex = ','.join(filters)129 130            # 提取帧的命令131            cmd = [132                'ffmpeg', '-i', video_full_path,133                '-vf', filter_complex,134                '-f', 'image2',135                '-qscale', current_quality["qscale"],  # 使用可调节的质量136                f'{frame_path}/%08d.jpg'137            ]138            139            resolution_info = f" (Resolution: {target_resolution})" if target_resolution else ""140            crop_info = "with square cropping" if square_crop else "without cropping"141            print(f"Processing {video_file} {crop_info} (FPS: {fps}, Quality: {quality_preset}{resolution_info})")142 143            subprocess.run(cmd, check=True)144            print(f"Successfully extracted frames for {video_file}")145 146            # Compile frames back into a video with optimized settings147            compile_frames_to_video_optimized(frame_path, output_video_path, fps, quality_preset)148 149        except Exception as e:150            print(f"Error preprocessing {video_file}: {str(e)}")151            continue152 153def compile_frames_to_video_optimized(frame_dir, output_path, fps=30, quality_preset="medium"):154    """Compile frames into a video with optimized quality settings."""155    156    # 质量设置 - CRF值（越高质量越低，文件越小）157    quality_crf = {158        "high": "18",159        "medium": "23", 160        "low": "28",161        "ultra_low": "35"162    }163    164    crf_value = quality_crf.get(quality_preset, "23")165    166    cmd = [167        'ffmpeg', '-y',168        '-f', 'image2',169        '-r', str(fps),170        '-i', f'{frame_dir}/%08d.jpg',171        '-c:v', 'libx264',172        '-preset', 'medium',  # 可以改为 'fast' 加速编码173        '-crf', crf_value,174        '-pix_fmt', 'yuv420p',175        output_path176    ]177    subprocess.run(cmd, check=True)178    print(f"Successfully compiled optimized video: {output_path} (Quality: {quality_preset})")179 180# 使用示例：181 182# 1. 保持原分辨率，降低质量183# preprocess_videos(video_dir, dataset_name, square_crop=True, fps=30, quality_preset="low")184 185# 2. 降低分辨率到512x512（方形裁剪）186# preprocess_videos(video_dir, dataset_name, square_crop=True, fps=30, quality_preset="medium", target_resolution=512)187 188# 3. 极度压缩：低分辨率 + 超低质量189# preprocess_videos(video_dir, dataset_name, square_crop=True, fps=30, quality_preset="ultra_low", target_resolution=256)190 191# 4. 不裁剪，但缩放到较小尺寸192# preprocess_videos(video_dir, dataset_name, square_crop=False, fps=30, quality_preset="low", target_resolution=480)193 194def process_npz_files(input_folder_path, output_folder_path):195    """196    Process all NPZ files in the specified folder and generate the required output format.197    198    Args:199        input_folder_path (str): Path to the folder containing NPZ files200        output_folder_path (str): Path where output files will be saved201    """202    # Get all NPZ files in the folder203    npz_files = sorted([f for f in os.listdir(input_folder_path) if f.endswith('.npz')])204    total_frames = len(npz_files)205    206    output = []207    208    for idx, npz_file in enumerate(npz_files):209        file_path = os.path.join(input_folder_path, npz_file)210        data = np.load(file_path, allow_pickle=True)211        212        # Process bodies data213        bodies = data['bodies']214        body_scores = data['body_scores'][0]215        216        # Process hands data217        hands = data['hands']218        hands_scores = data['hands_scores']219        220        # Process faces data221        faces = data['faces'][0]222        faces_scores = data['faces_scores'][0]223        224        # Convert coordinates to strings with space separation225        frame_data = []226        227        # Add body coordinates and scores228        for i in range(bodies.shape[0]):229            frame_data.extend([f"{bodies[i][0]:.8f}", f"{bodies[i][1]:.8f}"])230        for score in body_scores:231            frame_data.append(f"{score:.8f}")232            233        # Add hand coordinates and scores234        for hand in hands:235            for point in hand:236                frame_data.extend([f"{point[0]:.8f}", f"{point[1]:.8f}"])237        for hand_score in hands_scores:238            frame_data.extend([f"{score:.8f}" for score in hand_score])239            240        # Add face coordinates and scores241        for point in faces:242            frame_data.extend([f"{point[0]:.8f}", f"{point[1]:.8f}"])243        for score in faces_scores:244            frame_data.append(f"{score:.8f}")245            246        # Add frame count247        frame_count = idx / (total_frames - 1) if total_frames > 1 else 0248        frame_data.append(f"{frame_count:.8f}")249        250        # 验证这一帧的数据点数是否为385251        if len(frame_data) != 385:252            print(f"Warning: Frame {idx} in {input_folder_path} has {len(frame_data)} values instead of 385")253            continue  # 跳过这一帧254 255        # Join all data with spaces256        output.append(" ".join(frame_data))257    258    return " ".join(output) + "\n"