core/model.py

from ultralytics import YOLO
import streamlit as st
import cv2
from huggingface_hub import hf_hub_download

@st.cache_resource
def load_model():
    model_path = hf_hub_download(
        repo_id="thanhhiepvos/Fine_tuned_Vietnam_Traffic_Sign_Detection",
        filename="YOLO12X.pt",
        repo_type="model"
    )
    return YOLO(model_path)

def run_yolo_inference(model, image_path, conf_threshold=0.25, iou_threshold=0.7, conf_show=False):
    # Read image with OpenCV (BGR)
    img = cv2.imread(image_path)
    results = model(img,
                    conf=conf_threshold,
                    iou=iou_threshold
                    )[0]
    
    classes_with_conf = []

    for i, box in enumerate(results.boxes):
        x1, y1, x2, y2 = box.xyxy[0].int().tolist()
        conf = float(box.conf[0])
        cls = int(box.cls[0])
        if conf_show:
            label = f"{model.names[cls]} {conf:.2f}"
        else:
            label = f"{model.names[cls]}"
        classes_with_conf.append([cls, conf])

        # Draw bounding box
        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), 2)

        # Get text size
        (text_w, text_h), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1)

        # New Y position: just below the bounding box, plus spacing
        label_y = y2 + (i * (text_h + 6)) + 10

        # Background rectangle for the label
        cv2.rectangle(img, (x1, label_y), (x1 + text_w, label_y + text_h + 4), (255, 255, 255), -1)

        # Red text over white background
        cv2.putText(img, label, (x1, label_y + text_h + 2), cv2.FONT_HERSHEY_SIMPLEX,
                    0.6, (0, 0, 255), 1, cv2.LINE_AA)

    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB), classes_with_conf