inference/detect.py

import os
import time
from ultralytics import YOLO
import streamlit as st
from inference.logger import get_logger

log = get_logger("detect")

DISPLAY = {
    "steel_coil": "Steel Coil",
    "steel_plate": "Steel Plate",
}

@st.cache_resource
def _model():
    base_dir = os.path.dirname(os.path.dirname(__file__))
    model_path = os.path.join(base_dir, "models", "poc1", "best.pt")

    if not os.path.exists(model_path):
        log.warning("POC1 model not found at %s — trying fallback path", model_path)
        model_path = "d:/vision-ai-manufacturing/models/poc1/best.pt"

    log.info("POC1 YOLO — loading from %s  exists=%s", model_path, os.path.exists(model_path))
    t0 = time.perf_counter()
    model = YOLO(model_path)
    log.info("POC1 YOLO — ready  total=%.2fs", time.perf_counter() - t0)
    return model

def detect_assets(image, conf=0.25):
    """image: file path, PIL.Image, or np.ndarray (BGR/RGB ok).
    Returns (annotated_rgb_ndarray, [{'name','confidence','box'}...])."""
    t0 = time.perf_counter()
    result = _model().predict(image, conf=conf, verbose=False)[0]
    log.debug("detect_assets — inference=%.3fs  detections=%d", time.perf_counter() - t0, len(result.boxes))

    annotated = result.plot()[:, :, ::-1]   # plot() returns BGR -> flip to RGB
    dets = []
    for b in result.boxes:
        slug = result.names[int(b.cls)]
        dets.append({
            "name": DISPLAY.get(slug, slug),
            "confidence": round(float(b.conf) * 100, 1),
            "box": [round(v) for v in b.xyxy[0].tolist()],
        })
    dets.sort(key=lambda d: d["confidence"], reverse=True)
    return annotated, dets