detect_helmet.py

from ultralytics import YOLO
from PIL import Image
import os

MODEL_PATH = "models/yolo_helmet.pt"

def download_pretrained_yolo():
    print("Downloading YOLOv8n for helmet detection...")
    model = YOLO("yolov8n.pt")
    model.save(MODEL_PATH)
    return model

def load_helmet_model():
    if not os.path.exists(MODEL_PATH):
        print("Model not found, downloading base YOLOv8...")
        download_pretrained_yolo()
    model = YOLO(MODEL_PATH)
    return model

def detect_helmet_in_image(img_pil, conf_threshold=0.5):
    model = load_helmet_model()
    results = model(img_pil, conf=conf_threshold)
    result = results[0]
    annotated_img = result.plot()
    annotated_pil = Image.fromarray(annotated_img)

    missing_helmets = 0
    detections = []

    for box in result.boxes:
        cls_id = int(box.cls)
        conf = float(box.conf)
        label = result.names[cls_id]
        detections.append(f"{label}: {conf:.2f}")
        if "head" in label.lower() or "person" in label.lower():
            missing_helmets += 1

    return annotated_pil, missing_helmets, ", ".join(detections)