LabelOcr/test_icon.py

"""
象形图识别测试脚本
用法：python test_icon.py
"""
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from PIL import Image
from agent.agent import OcrAgent, image_to_base64, resize_image
import json

# 预期结果
CASES = [
    {"file": "pic/A.png", "name": "镀铜药水", "expected": {"GHS05", "GHS09"}},
    {"file": "pic/B.png", "name": "液氨",     "expected": {"GHS04", "GHS05", "GHS06", "GHS09"}},
    {"file": "pic/C.png", "name": "三氯苯",   "expected": {"GHS07", "GHS09"}},
]

agent = OcrAgent()

print("=" * 60)
for case in CASES:
    img = Image.open(case["file"]).convert("RGB")
    img = resize_image(img, max_size=1600)
    b64 = image_to_base64(img, quality=95)

    _, result_json = agent.extract_icon(b64)
    got = set(json.loads(result_json)["tag_images"])
    expected = case["expected"]

    hit     = got & expected        # 正确识别
    missed  = expected - got        # 漏识别
    extra   = got - expected        # 多识别

    status = "✓ PASS" if not missed and not extra else "✗ FAIL"
    print(f"\n[{status}] {case['name']}")
    print(f"  期望: {sorted(expected)}")
    print(f"  识别: {sorted(got)}")
    if missed: print(f"  漏识别: {sorted(missed)}")
    if extra:  print(f"  误报:   {sorted(extra)}")

print("\n" + "=" * 60)