hubconfCustom.py

from ultralytics import YOLO
import cv2
import time
import numpy as np

# Configuration
opt = {
    "weights": r"C:\Users\mihar\OneDrive\Documents\Deep Leaning Project\ultralytics\yolo11s.pt",
    "img-size": 640,
    "conf-thres": 0.25,
    "device": 'cpu'
}

# Load YOLO model only once
model = YOLO(opt['weights'])

def video_detection(frame, conf=0.25):
    start_time = time.time()
    
    # Run inference
    results = model(frame, conf=conf, imgsz=opt["img-size"], device=opt["device"])

    total_detections = 0
    output_frame = frame.copy()

    # Process detection results
    for result in results:
        if result.boxes is not None:
            boxes = result.boxes.xyxy.cpu().numpy()
            scores = result.boxes.conf.cpu().numpy()
            class_ids = result.boxes.cls.cpu().numpy().astype(int)

            for i, box in enumerate(boxes):
                x1, y1, x2, y2 = map(int, box)
                conf_score = scores[i]
                cls_id = class_ids[i]
                class_name = model.names[cls_id]

                label = f"{class_name} {conf_score:.2f}"
                color = (0, 255, 0)

                cv2.rectangle(output_frame, (x1, y1), (x2, y2), color, 2)
                cv2.putText(output_frame, label, (x1, y1 - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
                total_detections += 1

    # Calculate FPS
    fps_x = int(1 / (time.time() - start_time + 1e-5))  # Avoid division by zero
    cv2.putText(output_frame, f'Total Detected: {total_detections}', (10, 30),
                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
    cv2.putText(output_frame, f'FPS: {fps_x}', (10, 60),
                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)

    return [(output_frame, fps_x, output_frame.shape, total_detections)]