|
|
|
|
|
import cv2 |
|
|
import numpy as np |
|
|
from PIL import Image, ImageDraw, ImageFont |
|
|
from typing import List, Tuple, Dict |
|
|
|
|
|
class BrandVisualizationManager: |
|
|
"""Visualize detected brands with bounding boxes and labels (like YOLO)""" |
|
|
|
|
|
def __init__(self): |
|
|
"""Initialize visualization manager""" |
|
|
|
|
|
self.colors = { |
|
|
'luxury': (218, 165, 32), |
|
|
'sportswear': (0, 191, 255), |
|
|
'tech': (169, 169, 169), |
|
|
'automotive': (220, 20, 60), |
|
|
'fashion': (186, 85, 211), |
|
|
'watches': (184, 134, 11), |
|
|
'default': (0, 255, 0) |
|
|
} |
|
|
|
|
|
print("✓ Brand Visualization Manager initialized") |
|
|
|
|
|
def draw_brand_detections(self, image: Image.Image, brand_detections: List[Dict], |
|
|
show_confidence: bool = True) -> Image.Image: |
|
|
"""Draw bounding boxes and labels for detected brands |
|
|
|
|
|
Args: |
|
|
image: PIL Image |
|
|
brand_detections: List of dicts with keys: 'name', 'confidence', 'bbox', 'category' |
|
|
show_confidence: Whether to show confidence scores |
|
|
|
|
|
Returns: |
|
|
Image with drawn bounding boxes |
|
|
""" |
|
|
if not brand_detections: |
|
|
return image |
|
|
|
|
|
|
|
|
img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) |
|
|
|
|
|
for detection in brand_detections: |
|
|
brand_name = detection.get('name', 'Unknown') |
|
|
confidence = detection.get('confidence', 0.0) |
|
|
bbox = detection.get('bbox') |
|
|
category = detection.get('category', 'default') |
|
|
|
|
|
if bbox is None: |
|
|
continue |
|
|
|
|
|
x1, y1, x2, y2 = bbox |
|
|
color = self.colors.get(category, self.colors['default']) |
|
|
|
|
|
|
|
|
cv2.rectangle(img_cv, (int(x1), int(y1)), (int(x2), int(y2)), color, 3) |
|
|
|
|
|
|
|
|
if show_confidence: |
|
|
label = f"{brand_name} {confidence:.2f}" |
|
|
else: |
|
|
label = brand_name |
|
|
|
|
|
|
|
|
font = cv2.FONT_HERSHEY_SIMPLEX |
|
|
font_scale = 0.7 |
|
|
thickness = 2 |
|
|
(text_width, text_height), baseline = cv2.getTextSize(label, font, font_scale, thickness) |
|
|
|
|
|
|
|
|
cv2.rectangle(img_cv, |
|
|
(int(x1), int(y1) - text_height - 10), |
|
|
(int(x1) + text_width + 10, int(y1)), |
|
|
color, -1) |
|
|
|
|
|
|
|
|
cv2.putText(img_cv, label, |
|
|
(int(x1) + 5, int(y1) - 5), |
|
|
font, font_scale, (255, 255, 255), thickness, cv2.LINE_AA) |
|
|
|
|
|
|
|
|
img_pil = Image.fromarray(cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB)) |
|
|
return img_pil |
|
|
|
|
|
def format_brand_list(self, brand_detections: List[Dict]) -> str: |
|
|
"""Format brand detections as readable text |
|
|
|
|
|
Args: |
|
|
brand_detections: List of brand detection dicts |
|
|
|
|
|
Returns: |
|
|
Formatted string with brand names and confidences |
|
|
""" |
|
|
if not brand_detections: |
|
|
return "No brands identified" |
|
|
|
|
|
formatted = [] |
|
|
for detection in brand_detections: |
|
|
brand_name = detection.get('name', 'Unknown') |
|
|
confidence = detection.get('confidence', 0.0) |
|
|
|
|
|
|
|
|
formatted.append(f"{brand_name} ({confidence:.2f})") |
|
|
|
|
|
return ", ".join(formatted) |
|
|
|
|
|
print("✓ BrandVisualizationManager defined") |
|
|
|