File size: 2,012 Bytes
6a3bd1f |
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 |
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
# Brand-relevant classes
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")
|