video_benchmark.py

#!/usr/bin/env python3
import argparse
import subprocess
import ffmpeg
import os
import re
import csv
import zipfile
import sys
from collections import defaultdict

from benchmark_multi_inst import run_multi_instance_benchmark

# --- 配置数据 ---
DECODE_FILE_ARR = [
    ["sample_1920x1080_30fps_8bit.h264", "1", "812", "H264", "8", "1920x1080"],
]

ENCODE_FILE_ARR = [
    ["1280x720_nv12.yuv", "677", "H264", "8", "1280x720"],
    ["1920x1080_nv12.yuv", "812", "H264", "8", "1920x1080"],
    ["3840x2160_nv12.yuv", "812", "H264", "8", "3840x2160"],
    ["7680x4320_nv12.yuv", "77", "H264", "8", "7680x4320"],
    ["1280x720_nv12.yuv", "677", "HEVC", "8", "1280x720"],
    ["1920x1080_nv12.yuv", "812", "HEVC", "8", "1920x1080"],
    ["3840x2160_nv12.yuv", "812", "HEVC", "8", "3840x2160"],
    ["7680x4320_nv12.yuv", "77", "HEVC", "8", "7680x4320"],
    ["1280x720_p010le.yuv", "677", "HEVC", "10", "1280x720"],
    ["1920x1080_p010le.yuv", "812", "HEVC", "10", "1920x1080"],
    ["3840x2160_p010le.yuv", "812", "HEVC", "10", "3840x2160"],
    ["1920x1080_nv12.yuv", "1", "JPEG_1", "8", "1920x1080"],
    ["3840x2160_nv12.yuv", "1", "JPEG_1", "8", "3840x2160"],
    ["4096x2160_nv12.yuv", "1", "JPEG_1", "8", "4096x2160"],
    ["8192x4320_nv12.yuv", "1", "JPEG_1", "8", "8192x4320"],
    ["1920x1080_nv12.yuv", "2174", "JPEG", "8", "1920x1080"],
    ["1920x1080_nv12.yuv", "1701", "JPEG", "8", "1920x1080"],
    ["1920x1080_nv12.yuv", "2076", "JPEG", "8", "1920x1080"],
    ["1920x1080_nv12.yuv", "1581", "JPEG", "8", "1920x1080"],
    ["1920x1080_nv12.yuv", "2154", "JPEG", "8", "1920x1080"],
]

# 编解码器映射（键为小写）
NVENC_ENCODER_NAMES = {
    "h264": "h264_nvenc",
    "hevc": "hevc_nvenc",
    "vp9": "vp9_nvenc",
    "mjpeg": "mjpeg_nvenc",
}

VAAPI_ENCODER_NAMES = {
    "h264": "h264_vaapi",
    "hevc": "hevc_vaapi",
    "vp9": "vp9_vaapi",
    "mjpeg": "mjpeg_vaapi",
}

V4L2_CODEC_NAMES = {
    "h264": "h264_v4l2m2m",
    "hevc": "hevc_v4l2m2m",
    "vp9": "vp9_v4l2m2m",
    "avs2": "avs2_v4l2m2m",
    "mjpeg": "mjpeg_v4l2m2m",
}

CUVID_CODEC_NAMES = {
    "h264": "h264_cuvid",
    "hevc": "hevc_cuvid",
    "vp9": "vp9_cuvid",
    "avs2": "avs2_cuvid",
    "mjpeg": "mjpeg_cuvid",
}

# --- 辅助函数 ---

def parse_ffmpeg_benchmark_times(log_content):
    """解析 ffmpeg 日志获取所有 rtime 值（秒）"""
    times = []
    # 匹配 rtime=0.123456 (benchmark 输出的浮点秒数)
    time_matches = re.findall(r'rtime=(\d+\.?\d*)', log_content)
    
    for time_str in time_matches:
        try:
            total_sec = float(time_str)
            times.append(total_sec)
        except ValueError:
            continue
    
    # 备选：匹配 time=00:00:00.00 (时间戳格式)
    if not times:
        timestamp_matches = re.findall(r'time=(\d+:\d+:\d+\.\d+)', log_content)
        for time_str in timestamp_matches:
            parts = time_str.split(':')
            if len(parts) == 3:
                h, m, s = parts
                s_clean = s.rstrip('s')
                try:
                    total_sec = float(h) * 3600 + float(m) * 60 + float(s_clean)
                    times.append(total_sec)
                except ValueError:
                    continue
    return times

def parse_ffmpeg_log_frames(log_content):
    """获取总帧数"""
    frame_matches = re.findall(r'frame=\s*(\d+)', log_content)
    if frame_matches:
        return int(frame_matches[-1])
    return 0

def calculate_fps(frames, rtime):
    if rtime > 0:
        return round(frames / rtime, 2)
    return 0.0

def run_ffmpeg_command(cmd, log_file_path):
    """运行 ffmpeg 命令并将 stderr 保存到日志文件（即使失败也创建文件）"""
    try:
        with open(log_file_path, 'w') as log_f:
            process = subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=log_f, text=True)
        return process.returncode
    except Exception as e:
        # 即使出错也创建日志文件
        with open(log_file_path, 'w') as log_f:
            log_f.write(f"Error: {e}\n")
        return -1

def extract_zip():
    video_benchmark_dir = "./video_benchmark"
    zip_file = "video_benchmark.zip"
    if not os.path.exists(video_benchmark_dir):
        if os.path.exists(zip_file):
            print(f"Extracting {zip_file}...")
            with zipfile.ZipFile(zip_file, 'r') as zip_ref:
                zip_ref.extractall(video_benchmark_dir)
        else:
            print(f"Warning: {zip_file} not found. Ensure test files exist in {video_benchmark_dir}")
    else:
        print("video_benchmark dir exist")

def get_decoder_args(codec_type, file_ext):
    """获取解码器参数"""
    file_ext_lower = file_ext.lower()
    
    if codec_type == "cuda":
        return {"hwaccel": "cuda"}
    elif codec_type == "vaapi":
        return {"hwaccel": "vaapi"}
    elif codec_type == "v4l2":
        decoder = V4L2_CODEC_NAMES.get(file_ext_lower, file_ext_lower)
        return {"vcodec": decoder}
    elif codec_type == "cuvid":
        decoder = CUVID_CODEC_NAMES.get(file_ext_lower, file_ext_lower)
        return {"vcodec": decoder}
    return {}

def get_encoder_name(codec_type, file_ext):
    """获取编码器名称（修复大小写问题）"""
    # 转换为小写
    file_ext_lower = file_ext.lower()
    
    # 特殊处理：H264 -> h264, HEVC -> hevc, JPEG -> mjpeg
    if file_ext_lower in ["h264", "hevc", "vp9", "avs2"]:
        codec_key = file_ext_lower
    elif file_ext_lower in ["jpg", "jpeg", "avi", "mjpeg"]:
        codec_key = "mjpeg"
    else:
        codec_key = file_ext_lower
    
    if codec_type == "cuda":
        return NVENC_ENCODER_NAMES.get(codec_key, "")
    elif codec_type == "vaapi":
        return VAAPI_ENCODER_NAMES.get(codec_key, "")
    elif codec_type == "v4l2":
        return V4L2_CODEC_NAMES.get(codec_key, "")
    elif codec_type == "cuvid":
        return CUVID_CODEC_NAMES.get(codec_key, "")
    return ""

def run_single_decode_benchmark(input_file, codec_type, log_file):
    """运行单次解码测试"""
    file_ext = os.path.splitext(input_file)[1].replace('.', '')
    input_kwargs = get_decoder_args(codec_type, file_ext)
    
    try:
        stream = ffmpeg.input(input_file, **input_kwargs)
        cmd = ffmpeg.output(stream, 'pipe:', format='null').global_args('-benchmark', '-y').compile()
        run_ffmpeg_command(cmd, log_file)
        return True
    except Exception as e:
        print(f"FFmpeg command construction failed: {e}")
        return False

def run_encode_benchmark(yuv_file, output_file, resolution, codec_type, file_ext, depth, log_file):
    """运行编码测试"""
    w, h = resolution.split('x')
    pix_fmt = "p010le" if "10" in str(depth) else "nv12"
    profile = "main10" if "10" in str(depth) else None
    
    # 修复：获取编码器名称（处理大小写）
    encoder_name = get_encoder_name(codec_type, file_ext)
    
    if not encoder_name:
        print(f"Unknown encoder for {file_ext} on {codec_type}")
        # 即使失败也创建日志文件
        with open(log_file, 'w') as f:
            f.write(f"Error: Unknown encoder for {file_ext} on {codec_type}\n")
        return False

    try:
        input_kwargs = {
            'format': 'rawvideo',
            's': f"{w}x{h}",
            'pix_fmt': pix_fmt
        }
        
        global_args = ['-benchmark', '-y']
        output_kwargs = {'vcodec': encoder_name}
        
        if codec_type == "vaapi":
            input_kwargs['hwaccel'] = 'vaapi'
            global_args.extend(['-vaapi_device', '/dev/dri/renderD128'])
            output_kwargs['vf'] = 'format=nv12,hwupload' if pix_fmt == 'nv12' else 'format=p010le,hwupload'
        
        if profile:
            output_kwargs['profile'] = profile
            
        stream = ffmpeg.input(yuv_file, **input_kwargs)
        cmd = ffmpeg.output(stream, output_file, **output_kwargs).global_args(*global_args).compile()
        
        run_ffmpeg_command(cmd, log_file)
        return True
    except Exception as e:
        print(f"Encode command failed: {e}")
        # 即使失败也创建日志文件
        with open(log_file, 'w') as f:
            f.write(f"Error: {e}\n")
        return False

def aggregate_results(temp_csv, final_csv, is_encode=False, threads=6):
    """聚合 CSV 结果"""
    data = defaultdict(lambda: {"sum_1core": 0.0, "sum_multi": 0.0, "count": 0})
    
    try:
        with open(temp_csv, 'r', newline='') as f:
            reader = csv.reader(f)
            header = next(reader)
            for row in reader:
                if len(row) < 5:
                    continue
                
                fmt = row[2].strip()
                depth = row[3].strip()
                res = row[4].strip()
                
                try:
                    fps_1 = float(row[5]) if len(row) > 5 else 0.0
                    fps_multi = float(row[6]) if (not is_encode and len(row) > 6) else 0.0
                except ValueError:
                    fps_1 = 0.0
                    fps_multi = 0.0
                
                key = f"{fmt},{depth},{res}"
                data[key]["sum_1core"] += fps_1
                data[key]["sum_multi"] += fps_multi
                data[key]["count"] += 1
    except FileNotFoundError:
        print(f"Warning: {temp_csv} not found.")
        return

    with open(final_csv, 'w', newline='') as f:
        writer = csv.writer(f)
        if is_encode:
            writer.writerow(["Encode", "Bit Depth", "Resolution", "Fps/(BMZ 1 Core)"])
            for key, val in data.items():
                if val["count"] > 0:
                    avg_1 = val["sum_1core"] / val["count"]
                    parts = key.split(',')
                    writer.writerow([parts[0], parts[1], parts[2], f"{avg_1:.2f}"])
        else:
            writer.writerow(["Decode", "Bit Depth", "Resolution", "Fps/(BMZ 1 Core)", f"Fps/(BMZ {threads} Processes)"])
            for key, val in data.items():
                if val["count"] > 0:
                    avg_1 = val["sum_1core"] / val["count"]
                    avg_multi = val["sum_multi"] / val["count"]
                    parts = key.split(',')
                    writer.writerow([parts[0], parts[1], parts[2], f"{avg_1:.2f}", f"{avg_multi:.2f}"])

def main():
    parser = argparse.ArgumentParser(description="Video Benchmark Script")
    parser.add_argument("-c", "--codec", default="cuda", choices=["cuda", "cuvid", "v4l2", "vaapi"], help="Codec backend")
    parser.add_argument("-n", "--threads", type=int, default=6, help="Number of concurrent processes")
    args = parser.parse_args()
    
    CODEC = args.codec
    THREADS = args.threads
    
    video_decode_benchmark_temp = f"video_{CODEC}_decode_benchmark_temp.csv"
    video_decode_benchmark = f"video_{CODEC}_decode_benchmark.csv"
    video_encode_benchmark_temp = f"video_{CODEC}_encode_benchmark_temp.csv"
    video_encode_benchmark = f"video_{CODEC}_encode_benchmark.csv"
    
    for f in [video_decode_benchmark_temp, video_encode_benchmark_temp]:
        if os.path.exists(f):
            os.remove(f)
            
    with open(video_decode_benchmark_temp, 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(["Test File", "Total Frames", "Decode", "Bit Depth", "Resolution", "Fps/(BMZ 1 Core)", f"Fps/(BMZ {THREADS} Processes)"])
        
    with open(video_encode_benchmark_temp, 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(["Test File", "Total Frames", "Encode", "Bit Depth", "Resolution", "Fps/(BMZ 1 Core)"])
        
    extract_zip()
    base_dir = "./video_benchmark"
    
    if not DECODE_FILE_ARR:
        print("Error: DECODE_FILE_ARR is empty. Please uncomment test files in script.")
        return

    for i, file_info in enumerate(DECODE_FILE_ARR):
        filename, encode_idx, frames, format_name, depth, resolution = file_info
        filepath = os.path.join(base_dir, filename)
        
        if not os.path.exists(filepath):
            print(f"Skip missing file: {filepath}")
            continue
            
        file_ext = os.path.splitext(filename)[1].replace('.', '')
        base_name = os.path.splitext(filename)[0]
        
        # 中间生成的 YUV 文件路径
        yuv_filename = os.path.join(base_dir, base_name + ".yuv")
        decode_log = os.path.join(base_dir, base_name + ".log")
        multi_log = os.path.join(base_dir, base_name + "_multi.log")
        encode_log = os.path.join(base_dir, base_name + "_encode.log")
        
        # --- 修改点 1: 编码输出文件改为临时文件，避免覆盖原文件 ---
        # 原代码：encode_result = os.path.join(base_dir, filename)
        # 新代码：添加 _encoded_tmp 后缀，确保不与原文件冲突
        encode_result = os.path.join(base_dir, f"{base_name}_encoded_tmp.{file_ext}")
        
        print(f"Processing: {filename} (Idx: {encode_idx})")
        
        # --- 解码测试 ---
        success = run_single_decode_benchmark(filepath, CODEC, decode_log)
        
        fps_single = 0.0
        fps_multi = 0.0
        error_occurred = False
        
        if not success or not os.path.exists(decode_log) or os.path.getsize(decode_log) == 0:
            error_occurred = True
            print(f"  -> Single instance failed or empty log")
        else:
            with open(decode_log, 'r') as f:
                log_content = f.read()
            
            d_times = parse_ffmpeg_benchmark_times(log_content)
            d_frames = parse_ffmpeg_log_frames(log_content)
            
            d_rtime = d_times[0] if d_times else 0.0
            
            if d_rtime > 0:
                fps_single = calculate_fps(d_frames, d_rtime)
            else:
                error_occurred = True
                print(f"  -> Failed to parse rtime from single log")
            
            # 多实例
            if not error_occurred:
                run_multi_instance_benchmark(filepath, CODEC, THREADS, multi_log)
                
                with open(multi_log, 'r') as f:
                    multi_log_content = f.read()
                
                m_times = parse_ffmpeg_benchmark_times(multi_log_content)
                
                if m_times:
                    avg_rtime = sum(m_times) / len(m_times)
                    fps_multi = calculate_fps(d_frames, avg_rtime) * THREADS
                else:
                    fps_multi = 0.0
                    print(f"  -> Failed to parse rtime from multi log")
                
        if error_occurred:
            fps_single = 0.0
            fps_multi = 0.0
            
        with open(video_decode_benchmark_temp, 'a', newline='') as f:
            writer = csv.writer(f)
            writer.writerow([filename, frames, format_name.strip(), depth.strip(), resolution.strip(), fps_single, fps_multi])
            
        # --- 编码测试 ---
        if encode_idx != "n/a" and not error_occurred:
            try:
                idx = int(encode_idx)
                if 0 <= idx < len(ENCODE_FILE_ARR):
                    enc_info = ENCODE_FILE_ARR[idx]
                    enc_yuv_name = enc_info[0]
                    enc_frames = enc_info[1]
                    enc_format = enc_info[2]
                    enc_depth = enc_info[3]
                    enc_resolution = enc_info[4]
                    
                    print(f"Encoding: {filename} -> {enc_format}")
                    
                    # 生成 YUV（使用与解码测试相同的解码器参数）
                    if not os.path.exists(yuv_filename):
                        try:
                            input_kwargs = get_decoder_args(CODEC, file_ext)
                            stream = ffmpeg.input(filepath, **input_kwargs)
                            cmd = ffmpeg.output(stream, yuv_filename, pix_fmt='nv12').global_args('-y').compile()
                            subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
                        except Exception as e:
                            print(f"YUV conversion failed: {e}")
                            enc_fps = 0.0
                            with open(video_encode_benchmark_temp, 'a', newline='') as f:
                                csv.writer(f).writerow([enc_yuv_name, enc_frames, enc_format.strip(), enc_depth.strip(), enc_resolution.strip(), enc_fps])
                            continue

                    # 编码 - 使用 Encode_File_Arr 中的格式
                    # 注意：输出文件现在是 encode_result (临时文件)，不会覆盖 filepath (原文件)
                    run_encode_benchmark(yuv_filename, encode_result, enc_resolution, CODEC, enc_format.strip(), enc_depth, encode_log)
                    
                    # 检查日志文件是否存在
                    if os.path.exists(encode_log):
                        with open(encode_log, 'r') as f:
                            enc_log_content = f.read()
                        e_frames = parse_ffmpeg_log_frames(enc_log_content)
                        e_times = parse_ffmpeg_benchmark_times(enc_log_content)
                        e_rtime = e_times[0] if e_times else 0.0
                        enc_fps = calculate_fps(e_frames, e_rtime)
                    else:
                        enc_fps = 0.0
                        print(f"  -> Encode log file not created")
                    
                    with open(video_encode_benchmark_temp, 'a', newline='') as f:
                        writer = csv.writer(f)
                        writer.writerow([enc_yuv_name, enc_frames, enc_format.strip(), enc_depth.strip(), enc_resolution.strip(), enc_fps])
                        
                    # --- 修改点 2: 清理临时文件 ---
                    # 现在 encode_result 是临时文件，可以安全删除，不会误删原文件
                    for f_to_rm in [yuv_filename, encode_result, encode_log]:
                        if os.path.exists(f_to_rm):
                            os.remove(f_to_rm)
            except ValueError:
                pass
                
        # 清理日志
        for f_to_rm in [decode_log, multi_log]:
            if os.path.exists(f_to_rm):
                os.remove(f_to_rm)
            
    print("Aggregating results...")
    aggregate_results(video_decode_benchmark_temp, video_decode_benchmark, is_encode=False, threads=THREADS)
    aggregate_results(video_encode_benchmark_temp, video_encode_benchmark, is_encode=True, threads=THREADS)
    
    for f in [video_decode_benchmark_temp, video_encode_benchmark_temp]:
        if os.path.exists(f):
            os.remove(f)
    
    print("Benchmark finished.")

if __name__ == "__main__":
    main()