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,65 @@
+"""MaskFormer judge: scene classification via COCO-Stuff 171 class"""
+import torch, numpy as np
+from PIL import Image
+from transformers import MaskFormerImageProcessor, MaskFormerForInstanceSegmentation
+
+MODEL_ID = "facebook/maskformer-resnet50-coco-stuff"
+
+_model = None
+_processor = None
+_id2label = None
+
+INDOOR_STUFF = {"wall", "floor", "ceiling", "door", "window", "curtain", "desk", "table",
+                "furniture", "bed", "chair", "cabinet", "shelf", "carpet", "pillow"}
+OUTDOOR_STUFF = {"sky", "road", "river", "sea", "grass", "tree", "mountain", "pavement",
+                 "sand", "gravel", "snow", "cloud"}
+
+def load():
+    global _model, _processor, _id2label
+    if _model is None:
+        _processor = MaskFormerImageProcessor.from_pretrained(MODEL_ID)
+        _model = MaskFormerForInstanceSegmentation.from_pretrained(MODEL_ID).eval()
+        if torch.backends.mps.is_available():
+            _model = _model.to("mps")
+        _id2label = {int(k): v for k, v in _model.config.id2label.items()}
+
+def score(frames, prompt):
+    load()
+    results = []
+    for img in frames:
+        w, h = img.size
+        inputs = _processor(images=img, return_tensors="pt")
+        if torch.backends.mps.is_available():
+            inputs = {k: v.to("mps") for k, v in inputs.items()}
+        with torch.no_grad():
+            outputs = _model(**inputs)
+        seg = _processor.post_process_semantic_segmentation(outputs, target_sizes=[(h, w)])[0].cpu().numpy()
+        
+        classes, counts = np.unique(seg, return_counts=True)
+        total_px = h * w
+        stuff_found = []
+        indoor_px = outdoor_px = 0
+        for cid, cnt in zip(classes, counts):
+            lbl = _id2label.get(int(cid), f"_{cid}")
+            pct = 100 * cnt / total_px
+            if pct > 1.0:
+                stuff_found.append((lbl, pct))
+                if lbl in INDOOR_STUFF: indoor_px += cnt
+                if lbl in OUTDOOR_STUFF: outdoor_px += cnt
+        
+        is_indoor = bool(indoor_px > outdoor_px)
+        dominant = max(stuff_found, key=lambda x: x[1]) if stuff_found else ("unknown", 0)
+        results.append({
+            "is_indoor": is_indoor,
+            "dominant_stuff": dominant[0],
+            "dom_pct": round(dominant[1], 1),
+            "top_stuff": stuff_found[:5]
+        })
+    
+    is_indoor = results[0]["is_indoor"] if results else None
+    return {
+        "agent": "MaskFormer",
+        "score": 100 if is_indoor else 0,
+        "reasoning": f"Scene: {'indoor' if is_indoor else 'outdoor'} (dom={results[0]['dominant_stuff']})",
+        "details": {"frames": results}
+    }