330 lines
12 KiB
Python
330 lines
12 KiB
Python
import os
|
|
import cv2
|
|
import numpy as np
|
|
import time
|
|
from typing import List, Dict, Tuple, Optional
|
|
import importlib.util
|
|
import sys
|
|
|
|
from app.models import Detection
|
|
|
|
|
|
class PythonModelDetector:
|
|
"""Object detector using native Python models"""
|
|
|
|
def __init__(self, model_name: str, model_path: str, input_width: int, input_height: int, color: int = 0x00FF00, model_config: dict = None):
|
|
"""
|
|
Initialize detector with Python model
|
|
|
|
Args:
|
|
model_name: Name of the model
|
|
model_path: Path to the Python model file (.py)
|
|
input_width: Input width for the model
|
|
input_height: Input height for the model
|
|
color: RGB color for detection boxes (default: green)
|
|
model_config: Additional configuration to pass to the model
|
|
"""
|
|
self.model_name = model_name
|
|
self.input_width = input_width
|
|
self.input_height = input_height
|
|
self.color = color
|
|
self.model_config = model_config or {}
|
|
|
|
# Convert color from RGB to BGR (OpenCV uses BGR)
|
|
self.color_bgr = ((color & 0xFF) << 16) | (color & 0xFF00) | ((color >> 16) & 0xFF)
|
|
|
|
# Default confidence thresholds
|
|
self.conf_threshold = 0.25
|
|
self.nms_threshold = 0.45
|
|
|
|
# Load the Python model dynamically
|
|
self._load_python_model(model_path)
|
|
|
|
# Load class names if available
|
|
self.classes = []
|
|
model_dir = os.path.dirname(model_path)
|
|
classes_path = os.path.join(model_dir, "classes.txt")
|
|
if os.path.exists(classes_path):
|
|
with open(classes_path, 'r') as f:
|
|
self.classes = [line.strip() for line in f.readlines() if line.strip()]
|
|
|
|
def _load_python_model(self, model_path: str):
|
|
"""Load Python model dynamically"""
|
|
if not os.path.exists(model_path):
|
|
raise FileNotFoundError(f"Model file not found: {model_path}")
|
|
|
|
# Get model directory and file name
|
|
model_dir = os.path.dirname(model_path)
|
|
model_file = os.path.basename(model_path)
|
|
model_name = os.path.splitext(model_file)[0]
|
|
|
|
# Add model directory to system path
|
|
if model_dir not in sys.path:
|
|
sys.path.append(model_dir)
|
|
|
|
# Import the model module
|
|
spec = importlib.util.spec_from_file_location(model_name, model_path)
|
|
if spec is None:
|
|
raise ImportError(f"Failed to load model specification: {model_path}")
|
|
|
|
model_module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(model_module)
|
|
|
|
# Check if the module has the required interface
|
|
if not hasattr(model_module, "Model"):
|
|
raise AttributeError(f"Model module must define a 'Model' class: {model_path}")
|
|
|
|
# Create model instance with config
|
|
# Try to pass config to model constructor if it accepts parameters
|
|
import inspect
|
|
model_class = model_module.Model
|
|
sig = inspect.signature(model_class.__init__)
|
|
|
|
if len(sig.parameters) > 1: # Has parameters beyond 'self'
|
|
# Pass all config as keyword arguments
|
|
self.model = model_class(**self.model_config)
|
|
else:
|
|
# No parameters, create without arguments
|
|
self.model = model_class()
|
|
|
|
# Check if model has the required methods
|
|
if not hasattr(self.model, "predict"):
|
|
raise AttributeError(f"Model must implement 'predict' method: {model_path}")
|
|
|
|
def preprocess(self, img: np.ndarray) -> np.ndarray:
|
|
"""Preprocess image for model input"""
|
|
# Ensure BGR image
|
|
if len(img.shape) == 2: # Grayscale
|
|
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
|
|
elif img.shape[2] == 4: # BGRA
|
|
img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
|
|
|
|
# Resize to model input size
|
|
resized = cv2.resize(img, (self.input_width, self.input_height))
|
|
|
|
# Use model's preprocess method if available
|
|
if hasattr(self.model, "preprocess"):
|
|
return self.model.preprocess(resized)
|
|
|
|
# Default preprocessing: normalize to [0, 1]
|
|
return resized / 255.0
|
|
|
|
def detect(self, img: np.ndarray) -> Tuple[List[Detection], float]:
|
|
"""
|
|
Detect objects in an image
|
|
|
|
Args:
|
|
img: Input image in BGR format (OpenCV)
|
|
|
|
Returns:
|
|
List of Detection objects and inference time in milliseconds
|
|
"""
|
|
if img is None or img.size == 0:
|
|
return [], 0.0
|
|
|
|
# Original image dimensions
|
|
img_height, img_width = img.shape[:2]
|
|
|
|
# Measure inference time
|
|
start_time = time.time()
|
|
|
|
try:
|
|
# Run inference using model's predict method
|
|
# Note: Pass original image to model, let it handle preprocessing
|
|
# Expected return format from model's predict:
|
|
# List of dicts with keys: 'bbox', 'class_id', 'confidence'
|
|
# bbox: (x, y, w, h) normalized [0-1]
|
|
model_results = self.model.predict(img)
|
|
|
|
# Calculate inference time in milliseconds
|
|
inference_time = (time.time() - start_time) * 1000
|
|
|
|
# Convert model results to Detection objects
|
|
detections = []
|
|
|
|
for result in model_results:
|
|
# Skip low confidence detections
|
|
confidence = result.get('confidence', 0)
|
|
if confidence < self.conf_threshold:
|
|
continue
|
|
|
|
# Get bounding box (normalized coordinates)
|
|
bbox = result.get('bbox', [0, 0, 0, 0])
|
|
|
|
# Denormalize bbox to image coordinates
|
|
x = int(bbox[0] * img_width)
|
|
y = int(bbox[1] * img_height)
|
|
w = int(bbox[2] * img_width)
|
|
h = int(bbox[3] * img_height)
|
|
|
|
# Skip invalid boxes
|
|
if w <= 0 or h <= 0:
|
|
continue
|
|
|
|
# Get class ID and name
|
|
class_id = result.get('class_id', 0)
|
|
class_name = f"cls{class_id}"
|
|
if 0 <= class_id < len(self.classes):
|
|
class_name = self.classes[class_id]
|
|
|
|
# Create Detection object
|
|
label = f"[{self.model_name}] {class_name}"
|
|
detection = Detection(
|
|
label=label,
|
|
confidence=confidence,
|
|
x=x,
|
|
y=y,
|
|
width=w,
|
|
height=h,
|
|
color=self.color
|
|
)
|
|
detections.append(detection)
|
|
|
|
# Apply NMS if model doesn't do it internally
|
|
if hasattr(self.model, "applies_nms") and self.model.applies_nms:
|
|
return detections, inference_time
|
|
else:
|
|
# Convert detections to boxes and scores
|
|
boxes = [(d.x, d.y, d.width, d.height) for d in detections]
|
|
scores = [d.confidence for d in detections]
|
|
|
|
if boxes:
|
|
# Apply NMS
|
|
indices = self._non_max_suppression(boxes, scores, self.nms_threshold)
|
|
detections = [detections[i] for i in indices]
|
|
|
|
return detections, inference_time
|
|
|
|
except Exception as e:
|
|
print(f"Error during detection: {str(e)}")
|
|
return [], (time.time() - start_time) * 1000
|
|
|
|
def _non_max_suppression(self, boxes: List[Tuple], scores: List[float], threshold: float) -> List[int]:
|
|
"""Apply Non-Maximum Suppression to remove overlapping boxes"""
|
|
# Sort by score in descending order
|
|
indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
|
|
|
|
keep = []
|
|
while indices:
|
|
# Get index with highest score
|
|
current = indices.pop(0)
|
|
keep.append(current)
|
|
|
|
# No more indices to process
|
|
if not indices:
|
|
break
|
|
|
|
# Get current box
|
|
x1, y1, w1, h1 = boxes[current]
|
|
x2_1 = x1 + w1
|
|
y2_1 = y1 + h1
|
|
area1 = w1 * h1
|
|
|
|
# Check remaining boxes
|
|
i = 0
|
|
while i < len(indices):
|
|
# Get box to compare
|
|
idx = indices[i]
|
|
x2, y2, w2, h2 = boxes[idx]
|
|
x2_2 = x2 + w2
|
|
y2_2 = y2 + h2
|
|
area2 = w2 * h2
|
|
|
|
# Calculate intersection
|
|
xx1 = max(x1, x2)
|
|
yy1 = max(y1, y2)
|
|
xx2 = min(x2_1, x2_2)
|
|
yy2 = min(y2_1, y2_2)
|
|
|
|
# Calculate intersection area
|
|
w = max(0, xx2 - xx1)
|
|
h = max(0, yy2 - yy1)
|
|
intersection = w * h
|
|
|
|
# Calculate IoU
|
|
union = area1 + area2 - intersection + 1e-9 # Avoid division by zero
|
|
iou = intersection / union
|
|
|
|
# Remove box if IoU is above threshold
|
|
if iou > threshold:
|
|
indices.pop(i)
|
|
else:
|
|
i += 1
|
|
|
|
return keep
|
|
|
|
def close(self):
|
|
"""Close the model resources"""
|
|
if hasattr(self.model, "close"):
|
|
self.model.close()
|
|
self.model = None
|
|
|
|
|
|
class ModelManager:
|
|
"""Model manager for detectors"""
|
|
|
|
def __init__(self):
|
|
self.models = {}
|
|
|
|
def load(self, models_config: List[Dict]):
|
|
"""
|
|
Load models from configuration
|
|
|
|
Args:
|
|
models_config: List of model configurations
|
|
"""
|
|
# Basic color palette for different models
|
|
palette = [0x00FF00, 0xFF8000, 0x00A0FF, 0xFF00FF, 0x00FFFF, 0xFF0000, 0x80FF00]
|
|
|
|
for i, model_config in enumerate(models_config):
|
|
name = model_config.get("name")
|
|
path = model_config.get("path")
|
|
size = model_config.get("size", [640, 640])
|
|
|
|
if not name or not path or not os.path.exists(path):
|
|
print(f"Skipping model: {name} - Invalid configuration")
|
|
continue
|
|
|
|
try:
|
|
# Use color from palette
|
|
color = palette[i % len(palette)]
|
|
|
|
# Extract model-specific config (model_file, model_name, etc.)
|
|
# These will be passed to the Model class __init__
|
|
model_init_config = {}
|
|
if "model_file" in model_config:
|
|
model_init_config["model_file"] = model_config["model_file"]
|
|
if "display_name" in model_config:
|
|
model_init_config["model_name"] = model_config["display_name"]
|
|
|
|
# Create detector for Python model
|
|
detector = PythonModelDetector(
|
|
model_name=name,
|
|
model_path=path,
|
|
input_width=size[0],
|
|
input_height=size[1],
|
|
color=color,
|
|
model_config=model_init_config
|
|
)
|
|
|
|
self.models[name] = detector
|
|
print(f"Model loaded: {name} ({path})")
|
|
except Exception as e:
|
|
print(f"Failed to load model {name}: {str(e)}")
|
|
|
|
def get(self, name: str) -> Optional[PythonModelDetector]:
|
|
"""Get detector by name"""
|
|
return self.models.get(name)
|
|
|
|
def all(self) -> List[PythonModelDetector]:
|
|
"""Get all detectors"""
|
|
return list(self.models.values())
|
|
|
|
def close(self):
|
|
"""Close all detectors"""
|
|
for detector in self.models.values():
|
|
try:
|
|
detector.close()
|
|
except:
|
|
pass
|
|
self.models.clear() |