63 lines
2.5 KiB
Python
63 lines
2.5 KiB
Python
"""YOLO judge: object detection matching against expected objects"""
|
|
import cv2, numpy as np, re
|
|
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 (word boundary)
|
|
expected = [c for c in COCO if re.search(r'\b' + re.escape(c) + r'\b', 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}
|
|
}
|