|
|
|
|
|
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") |
|
|
|
