|
|
| from ultralytics import YOLO |
| import numpy as np |
| from typing import List, Dict |
| from PIL import Image |
|
|
| class YOLODetectionManager: |
| """Object detection using YOLOv11""" |
|
|
| def __init__(self, variant='m'): |
| print(f"Loading YOLOv11{variant} model...") |
| self.model = YOLO(f'yolo11{variant}.pt') |
| self.variant = variant |
| self.conf_threshold = 0.25 |
| self.iou_threshold = 0.45 |
| self.max_detections = 100 |
|
|
| |
| self.brand_relevant_classes = [ |
| 'handbag', 'bottle', 'cell phone', 'laptop', |
| 'backpack', 'tie', 'suitcase', 'cup', 'watch', |
| 'shoe', 'sneaker', 'boot' |
| ] |
|
|
| print(f"✓ YOLOv11{variant} loaded") |
|
|
| def detect(self, image: np.ndarray) -> List[Dict]: |
| """Detect objects in image""" |
| results = self.model.predict( |
| image, |
| conf=self.conf_threshold, |
| iou=self.iou_threshold, |
| max_det=self.max_detections, |
| verbose=False |
| ) |
|
|
| detections = [] |
|
|
| for result in results: |
| boxes = result.boxes |
| for box in boxes: |
| class_id = int(box.cls[0]) |
| class_name = result.names[class_id] |
| bbox = box.xyxy[0].cpu().numpy().tolist() |
| confidence = float(box.conf[0]) |
|
|
| detection = { |
| 'class_id': class_id, |
| 'class_name': class_name, |
| 'bbox': bbox, |
| 'confidence': confidence, |
| 'is_brand_relevant': class_name.lower() in self.brand_relevant_classes, |
| 'source': 'yolo' |
| } |
| detections.append(detection) |
|
|
| return detections |
|
|
| def filter_brand_relevant_objects(self, detections: List[Dict]) -> List[Dict]: |
| """Filter brand-relevant objects""" |
| return [det for det in detections if det['is_brand_relevant']] |
|
|
| print("✓ YOLODetectionManager defined") |
|
|
