feat: QA self-check agent — 15 prompts, 5 judges, weighted scoring

2026-05-14 10:53:30 +08:00
parent 2b633174b9
commit f60a59b280
10 changed files with 875 additions and 0 deletions
@@ -0,0 +1,62 @@
+"""YOLO judge: object detection matching against expected objects"""
+import cv2, numpy as np
+from ultralytics import YOLO
+
+MODEL_PATH = "/Users/accusys/momentry_core_0.1/yolov8s.mlpackage"
+COCO = [
+    "person","bicycle","car","motorbike","aeroplane","bus","train","truck","boat",
+    "traffic light","fire hydrant","stop sign","parking meter","bench","bird","cat","dog",
+    "horse","sheep","cow","elephant","bear","zebra","giraffe","backpack","umbrella",
+    "handbag","tie","suitcase","frisbee","skis","snowboard","sports ball","kite",
+    "baseball bat","baseball glove","skateboard","surfboard","tennis racket","bottle",
+    "wine glass","cup","fork","knife","spoon","bowl","banana","apple","sandwich","orange",
+    "broccoli","carrot","hot dog","pizza","donut","cake","chair","sofa","pottedplant",
+    "bed","diningtable","toilet","tvmonitor","laptop","mouse","remote","keyboard",
+    "cell phone","microwave","oven","toaster","sink","refrigerator","book","clock",
+    "vase","scissors","teddy bear","hair drier","toothbrush",
+]
+
+_model = None
+
+def load():
+    global _model
+    if _model is None:
+        try:
+            _model = YOLO(MODEL_PATH, task="detect", verbose=False)
+        except:
+            _model = YOLO("yolov8s.pt")
+
+def score(frames, prompt):
+    load()
+    prompt_lower = prompt.lower()
+    
+    # Extract expected objects from prompt: check each COCO class
+    expected = [c for c in COCO if c in prompt_lower]
+    if not expected:
+        expected = ["person"]  # default fallback
+    
+    results = []
+    for img in frames:
+        # Convert PIL to numpy
+        arr = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+        dets = _model(arr, verbose=False, imgsz=640)
+        
+        found = []
+        if dets and len(dets) > 0 and dets[0].boxes is not None:
+            for cls_id in dets[0].boxes.cls.int().tolist():
+                cls_name = COCO[cls_id] if cls_id < len(COCO) else f"cls_{cls_id}"
+                found.append(cls_name)
+        
+        match = sum(1 for e in expected if e in found)
+        results.append({"expected": expected, "found": found, "match_count": match, "total": len(expected)})
+    
+    total_match = sum(r["match_count"] for r in results)
+    total_expected = sum(r["total"] for r in results) or 1
+    score_val = int(100 * total_match / total_expected)
+    
+    return {
+        "agent": "YOLO",
+        "score": score_val,
+        "reasoning": f"Found {total_match}/{total_expected} expected objects: {expected}",
+        "details": {"frames": results}
+    }