Spaces:
Sleeping
Sleeping
File size: 7,380 Bytes
c8a046c |
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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 |
import numpy as np
from typing import List, Dict, Tuple
import cv2
from pathlib import Path
import yaml
class YOLOv11Detector:
"""YOLOv11 detector for car damage detection"""
def __init__(self, config_path: str = "config.yaml"):
"""Initialize YOLOv11 detector with configuration"""
with open(config_path, 'r') as f:
self.config = yaml.safe_load(f)
model_path = self.config['model']['path']
# Check which model file exists
if not Path(model_path).exists():
# Try to find available model files
model_dir = Path("models")
if (model_dir / "best.pt").exists():
model_path = str(model_dir / "best.pt")
print(f"Using best.pt model from training")
elif (model_dir / "last.pt").exists():
model_path = str(model_dir / "last.pt")
print(f"Using last.pt checkpoint model")
elif (model_dir / "best.onnx").exists():
model_path = str(model_dir / "best.onnx")
print(f"Using best.onnx model")
else:
raise FileNotFoundError(f"No model files found in models/ directory!")
self.model_path = model_path
self.device = self.config['model']['device']
self.confidence = self.config['model']['confidence']
self.iou_threshold = self.config['model']['iou_threshold']
self.classes = self.config['detection']['classes']
# Load model based on format
if model_path.endswith('.onnx'):
self._load_onnx_model()
else: # .pt format
self._load_pytorch_model()
def _load_pytorch_model(self):
"""Load PyTorch model using Ultralytics"""
from ultralytics import YOLO
self.model = YOLO(self.model_path)
# Set model to appropriate device
if self.device == 'cuda:0':
self.model.to('cuda')
print(f"Loaded PyTorch model: {self.model_path}")
def _load_onnx_model(self):
"""Load ONNX model using OpenCV DNN"""
self.net = cv2.dnn.readNet(self.model_path)
# Set backend based on device
if self.device == 'cuda:0':
self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
self.net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
else:
self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
self.net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)
print(f"Loaded ONNX model: {self.model_path}")
def detect(self, image: np.ndarray) -> Dict:
"""
Perform detection on image
Args:
image: Input image as numpy array (BGR format)
Returns:
Dictionary containing detection results
"""
if self.model_path.endswith('.onnx'):
return self._detect_onnx(image)
else:
return self._detect_pytorch(image)
def _detect_pytorch(self, image: np.ndarray) -> Dict:
"""Detection using PyTorch model"""
# Run YOLO inference
results = self.model(
image,
conf=self.confidence,
iou=self.iou_threshold,
device=self.device,
verbose=False
)
# Parse results
detections = {
'boxes': [],
'confidences': [],
'classes': [],
'class_ids': []
}
if len(results) > 0 and results[0].boxes is not None:
boxes = results[0].boxes
for box in boxes:
# Get box coordinates (xyxy format)
x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
# Get confidence and class
conf = float(box.conf[0].cpu().numpy())
cls_id = int(box.cls[0].cpu().numpy())
# Map class ID to class name
if cls_id < len(self.classes):
cls_name = self.classes[cls_id]
else:
cls_name = f"class_{cls_id}"
detections['boxes'].append([int(x1), int(y1), int(x2), int(y2)])
detections['confidences'].append(conf)
detections['classes'].append(cls_name)
detections['class_ids'].append(cls_id)
return detections
def _detect_onnx(self, image: np.ndarray) -> Dict:
"""Detection using ONNX model (compatible with original code)"""
height, width = image.shape[:2]
# Preprocess image for ONNX
blob = cv2.dnn.blobFromImage(
image, 1/255.0, (640, 640),
swapRB=True, crop=False
)
self.net.setInput(blob)
preds = self.net.forward()
preds = preds.transpose((0, 2, 1))
# Extract outputs
detections = self._extract_onnx_output(
preds=preds,
image_shape=(height, width),
input_shape=(640, 640)
)
return detections
def _extract_onnx_output(self, preds: np.ndarray, image_shape: Tuple[int, int],
input_shape: Tuple[int, int]) -> Dict:
"""Extract detection results from ONNX model output"""
class_ids, confs, boxes = [], [], []
image_height, image_width = image_shape
input_height, input_width = input_shape
x_factor = image_width / input_width
y_factor = image_height / input_height
rows = preds[0].shape[0]
for i in range(rows):
row = preds[0][i]
conf = row[4]
classes_score = row[4:]
_, _, _, max_idx = cv2.minMaxLoc(classes_score)
class_id = max_idx[1]
if classes_score[class_id] > self.confidence:
confs.append(float(conf))
label = self.classes[int(class_id)] if int(class_id) < len(self.classes) else f"class_{class_id}"
class_ids.append(label)
# Extract boxes
x, y, w, h = row[0].item(), row[1].item(), row[2].item(), row[3].item()
left = int((x - 0.5 * w) * x_factor)
top = int((y - 0.5 * h) * y_factor)
width = int(w * x_factor)
height = int(h * y_factor)
box = [left, top, left + width, top + height]
boxes.append(box)
# Apply NMS
if len(boxes) > 0:
indices = cv2.dnn.NMSBoxes(
[[b[0], b[1], b[2]-b[0], b[3]-b[1]] for b in boxes],
confs, self.confidence, self.iou_threshold
)
if len(indices) > 0:
indices = indices.flatten()
return {
'boxes': [boxes[i] for i in indices],
'confidences': [confs[i] for i in indices],
'classes': [class_ids[i] for i in indices],
'class_ids': list(range(len(indices)))
}
return {'boxes': [], 'confidences': [], 'classes': [], 'class_ids': []}
def detect_batch(self, images: List[np.ndarray]) -> List[Dict]:
"""Detect on multiple images"""
return [self.detect(img) for img in images] |