Spaces:
Runtime error
Runtime error
File size: 1,653 Bytes
f54e27b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | 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
|