detector.py

"""
detector.py - YOLO detection wrapper
"""

from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image
from config import MODEL_PATH


class RoadsideDetector:
    def __init__(self):
        print("🤖 Loading YOLO model...")
        self.model = YOLO(MODEL_PATH)
        print("✅ Model loaded successfully!")

    def detect_items(self, image):
        results = self.model(image)
        detections = []
        for r in results:
            for box in r.boxes:
                cls = int(box.cls[0])
                conf = float(box.conf[0])
                label = self.model.names[cls]
                x1, y1, x2, y2 = box.xyxy[0].tolist()
                detections.append({
                    "type": label,
                    "confidence": conf,
                    "bbox": [x1, y1, x2, y2]
                })
        return detections

    def draw_boxes(self, image, detections):
        img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)

        for det in detections:
            x1, y1, x2, y2 = [int(x) for x in det["bbox"]]
            label = f"{det['type']} {det['confidence']:.0%}"

            # Color by type
            if det["type"].lower() == "çöp":
                color = (0, 0, 255)  # red
            elif "kutusu" in det["type"].lower():
                color = (0, 165, 255)  # orange
            else:
                color = (0, 255, 255)  # yellow

            cv2.rectangle(img_cv, (x1, y1), (x2, y2), color, 2)
            cv2.putText(img_cv, label, (x1, max(0, y1 - 10)),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)

        return Image.fromarray(cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB))

    def calculate_severity(self, detections):
        count = len(detections)
        if count >= 5:
            return "HIGH"
        elif count >= 2:
            return "MEDIUM"
        elif count == 1:
            return "LOW"
        return "NONE"