132 lines
5.0 KiB
Python
132 lines
5.0 KiB
Python
import os
|
|
import cv2
|
|
import numpy as np
|
|
import onnxruntime as ort
|
|
|
|
from data_generator import CLASSES
|
|
|
|
class FruitVegetableDetector:
|
|
def __init__(self, model_path=os.path.join("models", "model.onnx")):
|
|
self.model_path = model_path
|
|
self.session = None
|
|
self.input_name = None
|
|
self.load_model()
|
|
|
|
def load_model(self):
|
|
if not os.path.exists(self.model_path):
|
|
print(f"Model file {self.model_path} not found. Running in SIMULATED mode.")
|
|
self.session = None
|
|
return
|
|
|
|
try:
|
|
# Set providers. SSDLite ONNX Runtime runs extremely fast on CPU
|
|
providers = ['CPUExecutionProvider']
|
|
if 'DirectMLExecutionProvider' in ort.get_available_providers():
|
|
providers = ['DirectMLExecutionProvider', 'CPUExecutionProvider']
|
|
|
|
self.session = ort.InferenceSession(self.model_path, providers=providers)
|
|
self.input_name = self.session.get_inputs()[0].name
|
|
print(f"ONNX model loaded successfully from {self.model_path}")
|
|
except Exception as e:
|
|
print(f"Failed to load ONNX model: {e}. Falling back to SIMULATED mode.")
|
|
self.session = None
|
|
|
|
def detect(self, image_rgb, confidence_threshold=0.10):
|
|
"""
|
|
Runs inference on an RGB image.
|
|
Args:
|
|
image_rgb (np.ndarray): HWC RGB image.
|
|
confidence_threshold (float): Minimum confidence to report.
|
|
Returns:
|
|
list: List of predictions with class, confidence, and box [xmin, ymin, xmax, ymax]
|
|
"""
|
|
orig_h, orig_w, _ = image_rgb.shape
|
|
|
|
# If no model is loaded, run simulation fallback
|
|
if self.session is None:
|
|
return self._simulate_detection(orig_w, orig_h, confidence_threshold)
|
|
|
|
# Preprocessing: Resize to 320x320
|
|
resized = cv2.resize(image_rgb, (320, 320))
|
|
# Scale to [0.0, 1.0] and transpose to [C, H, W]
|
|
input_tensor = resized.astype(np.float32) / 255.0
|
|
input_tensor = np.transpose(input_tensor, (2, 0, 1)) # shape: (3, 320, 320)
|
|
|
|
# Run inference
|
|
outputs = self.session.run(None, {self.input_name: input_tensor})
|
|
boxes, scores, labels = outputs
|
|
|
|
predictions = []
|
|
for i in range(len(scores)):
|
|
score = float(scores[i])
|
|
if score < confidence_threshold:
|
|
continue
|
|
|
|
# Class ID is 1-indexed, convert to 0-indexed
|
|
class_id = int(labels[i]) - 1
|
|
if class_id < 0 or class_id >= len(CLASSES):
|
|
continue
|
|
|
|
class_name = CLASSES[class_id]
|
|
box_320 = boxes[i]
|
|
|
|
# Scale box back to original coordinates
|
|
xmin = int(round((box_320[0] / 320.0) * orig_w))
|
|
ymin = int(round((box_320[1] / 320.0) * orig_h))
|
|
xmax = int(round((box_320[2] / 320.0) * orig_w))
|
|
ymax = int(round((box_320[3] / 320.0) * orig_h))
|
|
|
|
# Clip to image boundaries
|
|
xmin = max(0, min(xmin, orig_w - 1))
|
|
ymin = max(0, min(ymin, orig_h - 1))
|
|
xmax = max(xmin + 1, min(xmax, orig_w))
|
|
ymax = max(ymin + 1, min(ymax, orig_h))
|
|
|
|
predictions.append({
|
|
"class": class_name,
|
|
"confidence": round(score, 3),
|
|
"box": [xmin, ymin, xmax, ymax]
|
|
})
|
|
|
|
# Sort predictions by confidence descending
|
|
predictions = sorted(predictions, key=lambda x: x["confidence"], reverse=True)
|
|
return predictions
|
|
|
|
def _simulate_detection(self, width, height, confidence_threshold):
|
|
"""
|
|
Generates simulated detections based on random chances.
|
|
Used as a fallback when the model hasn't been trained yet.
|
|
"""
|
|
# Seed a pseudo-random detection based on image size to look semi-stable
|
|
seed_factor = (width + height) % 6
|
|
|
|
# Simulate 1 to 2 detections
|
|
predictions = []
|
|
|
|
# 1. First main prediction (high confidence)
|
|
class_idx = seed_factor
|
|
conf = 0.85 + (seed_factor * 0.02)
|
|
|
|
# Simulate reasonable bounding box centered on scale
|
|
cx, cy = int(width * 0.5), int(height * 0.5)
|
|
w, h = int(width * 0.35), int(height * 0.45)
|
|
|
|
predictions.append({
|
|
"class": CLASSES[class_idx],
|
|
"confidence": round(conf, 3),
|
|
"box": [cx - w//2, cy - h//2, cx + w//2, cy + h//2]
|
|
})
|
|
|
|
# 2. Second prediction (low confidence, e.g. alternative suggestion)
|
|
alt_class_idx = (class_idx + 1) % len(CLASSES)
|
|
alt_conf = 0.12 + (seed_factor * 0.01)
|
|
|
|
if alt_conf >= confidence_threshold:
|
|
predictions.append({
|
|
"class": CLASSES[alt_class_idx],
|
|
"confidence": round(alt_conf, 3),
|
|
"box": [cx - w//2 + 20, cy - h//2 + 10, cx + w//2 + 20, cy + h//2 + 10]
|
|
})
|
|
|
|
return predictions
|