60 lines
2.0 KiB
Python
60 lines
2.0 KiB
Python
# AI 추론 — YOLOv8 기반 불량(스크래치/이물/흑점/변형) 검출
|
|
import os
|
|
import numpy as np
|
|
|
|
from utils.path_helper import BASE_PATH
|
|
|
|
|
|
class Detector:
|
|
class_names = ["스크래치", "이물", "흑점", "변형"]
|
|
|
|
def __init__(self):
|
|
self._model = None
|
|
self.model_path = None
|
|
|
|
def load_model(self, model_path: str) -> bool:
|
|
if model_path and not os.path.isabs(model_path):
|
|
model_path = os.path.join(BASE_PATH, model_path)
|
|
try:
|
|
from ultralytics import YOLO # 지연 로딩 — 앱 시작 시 torch DLL 오류 방지
|
|
self._model = YOLO(model_path)
|
|
self.model_path = model_path
|
|
print(f"[AI] 모델 로드 성공: {model_path}")
|
|
return True
|
|
except Exception as e:
|
|
print(f"[AI] 모델 로드 실패: {e}")
|
|
self._model = None
|
|
return False
|
|
|
|
def is_loaded(self) -> bool:
|
|
return self._model is not None
|
|
|
|
def detect(self, image: np.ndarray) -> list:
|
|
"""
|
|
image: numpy array (BGR)
|
|
반환: [{"class_id": int, "class_name": str,
|
|
"confidence": float, "bbox": [x1,y1,x2,y2]}, ...]
|
|
"""
|
|
if self._model is None:
|
|
return []
|
|
try:
|
|
results = self._model(image, verbose=False)
|
|
detections = []
|
|
for result in results:
|
|
for box in result.boxes:
|
|
class_id = int(box.cls[0])
|
|
detections.append({
|
|
"class_id": class_id,
|
|
"class_name": (
|
|
self.class_names[class_id]
|
|
if class_id < len(self.class_names)
|
|
else str(class_id)
|
|
),
|
|
"confidence": float(box.conf[0]),
|
|
"bbox": box.xyxy[0].tolist(),
|
|
})
|
|
return detections
|
|
except Exception as e:
|
|
print(f"[AI] 추론 오류: {e}")
|
|
return []
|