File size: 2,578 Bytes
53ccd32 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | 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_13_vllm.py", "文字-排布位置"),
#("rule_14_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() |