#!/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()