| import os |
| import json |
| import subprocess |
| import pandas as pd |
| from pathlib import Path |
|
|
| GLOBAL_MODEL_PATH = "../stage2_object_v8_1200" |
|
|
| |
| TASK_CONFIGS = [ |
| ("rule_11_vllm.py", "文字占比"), |
| ("rule_12_vllm.py", "文字-样式数量合理"), |
| |
| |
| ("rule_17_vllm.py", "信息量"), |
| ("rule_18_vllm.py", "排布间距"), |
| ("rule_19_vllm.py", "内容构图"), |
| ] |
|
|
| SUMMARY_OUTPUT = "vllm_audit_summary.csv" |
|
|
| def run_vllm_task(script, input_dir): |
|
|
| cmd = [ |
| "python", script, |
| "--input_dir", input_dir, |
| "--model_path", GLOBAL_MODEL_PATH |
| ] |
| print(f"\n[EXEC] 正在运行: {script}") |
| print(f"[PATH] 输入目录: {Path(input_dir).name}") |
| |
| try: |
| subprocess.run(cmd, check=True) |
| return True |
| except subprocess.CalledProcessError as e: |
| print(f"[ERROR] 脚本 {script} 运行失败: {e}") |
| return False |
|
|
| def extract_metrics(script, input_dir): |
| input_path = Path(input_dir) |
| json_file = input_path / f"audit_result_{input_path.name}.json" |
| |
| if not json_file.exists(): |
| return { |
| "规则名称": script, |
| "对应目录": input_path.name, |
| "结果": "未找到 JSON 文件" |
| } |
|
|
| with open(json_file, 'r', encoding='utf-8') as f: |
| data = json.load(f) |
| |
| total = len(data) |
| unsuitable = sum(1 for item in data if item["label"] == "Unsuitable") |
| suitable = sum(1 for item in data if item["label"] == "Suitable") |
| recall = (unsuitable / total) if total > 0 else 0 |
|
|
| return { |
| "规则名称": script.replace("_vllm.py", ""), |
| "对应目录": input_path.name, |
| "样本总数": total, |
| "不通过数(Unsuitable)": unsuitable, |
| "通过数(Suitable)": suitable, |
| "召回率(违规检出率)": f"{recall:.2%}", |
| "所用模型": Path(GLOBAL_MODEL_PATH).name |
| } |
|
|
| def main(): |
| final_results = [] |
| |
| for script, input_dir in TASK_CONFIGS: |
| success = run_vllm_task(script, input_dir) |
| |
| metrics = extract_metrics(script, input_dir) |
| final_results.append(metrics) |
| |
| df = pd.DataFrame(final_results) |
| df.to_csv(SUMMARY_OUTPUT, index=False, encoding='utf-8-sig') |
| |
| print("\n" + "="*70) |
| print(f"所有任务运行完毕！汇总报告已保存至: {SUMMARY_OUTPUT}") |
| print("-" * 70) |
| print(df.to_string(index=False)) |
| print("="*70) |
|
|
| if __name__ == "__main__": |
| main() |
