File size: 3,774 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 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 |
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"""
# Color palette for different brand categories
self.colors = {
'luxury': (218, 165, 32), # Gold
'sportswear': (0, 191, 255), # Deep Sky Blue
'tech': (169, 169, 169), # Dark Gray
'automotive': (220, 20, 60), # Crimson
'fashion': (186, 85, 211), # Medium Orchid
'watches': (184, 134, 11), # Dark Goldenrod
'default': (0, 255, 0) # Green
}
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
# Convert PIL to OpenCV format
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'])
# Draw bounding box
cv2.rectangle(img_cv, (int(x1), int(y1)), (int(x2), int(y2)), color, 3)
# Prepare label text
if show_confidence:
label = f"{brand_name} {confidence:.2f}"
else:
label = brand_name
# Calculate text size
font = cv2.FONT_HERSHEY_SIMPLEX
font_scale = 0.7
thickness = 2
(text_width, text_height), baseline = cv2.getTextSize(label, font, font_scale, thickness)
# Draw label background
cv2.rectangle(img_cv,
(int(x1), int(y1) - text_height - 10),
(int(x1) + text_width + 10, int(y1)),
color, -1)
# Draw label text
cv2.putText(img_cv, label,
(int(x1) + 5, int(y1) - 5),
font, font_scale, (255, 255, 255), thickness, cv2.LINE_AA)
# Convert back to PIL
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)
# 移除 category 標籤顯示，保持簡潔
formatted.append(f"{brand_name} ({confidence:.2f})")
return ", ".join(formatted)
print("✓ BrandVisualizationManager defined")
|