Pixcribe / brand_visualization_manager.py
DawnC's picture
Upload 22 files
6a3bd1f verified
raw
history blame
3.77 kB
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")