color_mapper.py

import numpy as np
from typing import Dict, List, Tuple, Union, Any

class ColorMapper:
    """
    A class for consistent color mapping of object detection classes
    Provides color schemes for visualization in both RGB and hex formats
    """
    
    # Class categories for better organization
    CATEGORIES = {
        "person": [0],
        "vehicles": [1, 2, 3, 4, 5, 6, 7, 8],
        "traffic": [9, 10, 11, 12],
        "animals": [14, 15, 16, 17, 18, 19, 20, 21, 22, 23],
        "outdoor": [13, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33],
        "sports": [34, 35, 36, 37, 38],
        "kitchen": [39, 40, 41, 42, 43, 44, 45],
        "food": [46, 47, 48, 49, 50, 51, 52, 53, 54, 55],
        "furniture": [56, 57, 58, 59, 60, 61],
        "electronics": [62, 63, 64, 65, 66, 67, 68, 69, 70],
        "household": [71, 72, 73, 74, 75, 76, 77, 78, 79]
    }
    
    # Base colors for each category (in HSV for easier variation)
    CATEGORY_COLORS = {
        "person": (0, 0.8, 0.9),       # Red
        "vehicles": (210, 0.8, 0.9),   # Blue
        "traffic": (45, 0.8, 0.9),     # Orange
        "animals": (120, 0.7, 0.8),    # Green
        "outdoor": (180, 0.7, 0.9),    # Cyan
        "sports": (270, 0.7, 0.8),     # Purple
        "kitchen": (30, 0.7, 0.9),     # Light Orange
        "food": (330, 0.7, 0.85),      # Pink
        "furniture": (150, 0.5, 0.85), # Light Green
        "electronics": (240, 0.6, 0.9), # Light Blue
        "household": (60, 0.6, 0.9)    # Yellow
    }
    
    def __init__(self):
        """Initialize the ColorMapper with COCO class mappings"""
        self.class_names = self._get_coco_classes()
        self.color_map = self._generate_color_map()
    
    def _get_coco_classes(self) -> Dict[int, str]:
        """Get the standard COCO class names with their IDs"""
        return {
            0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 4: 'airplane',
            5: 'bus', 6: 'train', 7: 'truck', 8: 'boat', 9: 'traffic light',
            10: 'fire hydrant', 11: 'stop sign', 12: 'parking meter', 13: 'bench', 
            14: 'bird', 15: 'cat', 16: 'dog', 17: 'horse', 18: 'sheep', 19: 'cow',
            20: 'elephant', 21: 'bear', 22: 'zebra', 23: 'giraffe', 24: 'backpack',
            25: 'umbrella', 26: 'handbag', 27: 'tie', 28: 'suitcase', 29: 'frisbee',
            30: 'skis', 31: 'snowboard', 32: 'sports ball', 33: 'kite', 34: 'baseball bat',
            35: 'baseball glove', 36: 'skateboard', 37: 'surfboard', 38: 'tennis racket', 
            39: 'bottle', 40: 'wine glass', 41: 'cup', 42: 'fork', 43: 'knife',
            44: 'spoon', 45: 'bowl', 46: 'banana', 47: 'apple', 48: 'sandwich', 
            49: 'orange', 50: 'broccoli', 51: 'carrot', 52: 'hot dog', 53: 'pizza',
            54: 'donut', 55: 'cake', 56: 'chair', 57: 'couch', 58: 'potted plant', 
            59: 'bed', 60: 'dining table', 61: 'toilet', 62: 'tv', 63: 'laptop',
            64: 'mouse', 65: 'remote', 66: 'keyboard', 67: 'cell phone', 68: 'microwave', 
            69: 'oven', 70: 'toaster', 71: 'sink', 72: 'refrigerator', 73: 'book',
            74: 'clock', 75: 'vase', 76: 'scissors', 77: 'teddy bear', 78: 'hair drier', 
            79: 'toothbrush'
        }
    
    def _hsv_to_rgb(self, h: float, s: float, v: float) -> Tuple[int, int, int]:
        """
        Convert HSV color to RGB
        
        Args:
            h: Hue (0-360)
            s: Saturation (0-1)
            v: Value (0-1)
            
        Returns:
            Tuple of (R, G, B) values (0-255)
        """
        h = h / 60
        i = int(h)
        f = h - i
        p = v * (1 - s)
        q = v * (1 - s * f)
        t = v * (1 - s * (1 - f))
        
        if i == 0:
            r, g, b = v, t, p
        elif i == 1:
            r, g, b = q, v, p
        elif i == 2:
            r, g, b = p, v, t
        elif i == 3:
            r, g, b = p, q, v
        elif i == 4:
            r, g, b = t, p, v
        else:
            r, g, b = v, p, q
        
        return (int(r * 255), int(g * 255), int(b * 255))
    
    def _rgb_to_hex(self, rgb: Tuple[int, int, int]) -> str:
        """
        Convert RGB color to hex color code
        
        Args:
            rgb: Tuple of (R, G, B) values (0-255)
            
        Returns:
            Hex color code (e.g. '#FF0000')
        """
        return f'#{rgb[0]:02x}{rgb[1]:02x}{rgb[2]:02x}'
    
    def _find_category(self, class_id: int) -> str:
        """
        Find the category for a given class ID
        
        Args:
            class_id: Class ID (0-79)
            
        Returns:
            Category name
        """
        for category, ids in self.CATEGORIES.items():
            if class_id in ids:
                return category
        return "other"  # Fallback
    
    def _generate_color_map(self) -> Dict:
        """
        Generate a color map for all 80 COCO classes
        
        Returns:
            Dictionary mapping class IDs and names to color values
        """
        color_map = {
            'by_id': {},      # Map class ID to RGB and hex
            'by_name': {},    # Map class name to RGB and hex
            'categories': {}  # Map category to base color
        }
        
        # Generate colors for categories
        for category, hsv in self.CATEGORY_COLORS.items():
            rgb = self._hsv_to_rgb(hsv[0], hsv[1], hsv[2])
            hex_color = self._rgb_to_hex(rgb)
            color_map['categories'][category] = {
                'rgb': rgb,
                'hex': hex_color
            }
        
        # Generate variations for each class within a category
        for class_id, class_name in self.class_names.items():
            category = self._find_category(class_id)
            base_hsv = self.CATEGORY_COLORS.get(category, (0, 0, 0.8))  # Default gray
            
            # Slightly vary the hue and saturation within the category
            ids_in_category = self.CATEGORIES.get(category, [])
            if ids_in_category:
                position = ids_in_category.index(class_id) if class_id in ids_in_category else 0
                variation = position / max(1, len(ids_in_category) - 1)  # 0 to 1
                
                # Vary hue slightly (±15°) and saturation
                h_offset = 30 * variation - 15  # -15 to +15
                s_offset = 0.2 * variation  # 0 to 0.2
                
                h = (base_hsv[0] + h_offset) % 360
                s = min(1.0, base_hsv[1] + s_offset)
                v = base_hsv[2]
            else:
                h, s, v = base_hsv
            
            rgb = self._hsv_to_rgb(h, s, v)
            hex_color = self._rgb_to_hex(rgb)
            
            # Store in both mappings
            color_map['by_id'][class_id] = {
                'rgb': rgb,
                'hex': hex_color,
                'category': category
            }
            
            color_map['by_name'][class_name] = {
                'rgb': rgb,
                'hex': hex_color,
                'category': category
            }
        
        return color_map
    
    def get_color(self, class_identifier: Union[int, str], format: str = 'hex') -> Any:
        """
        Get color for a specific class
        
        Args:
            class_identifier: Class ID (int) or name (str)
            format: Color format ('hex', 'rgb', or 'bgr')
            
        Returns:
            Color in requested format
        """
        # Determine if identifier is an ID or name
        if isinstance(class_identifier, int):
            color_info = self.color_map['by_id'].get(class_identifier)
        else:
            color_info = self.color_map['by_name'].get(class_identifier)
        
        if not color_info:
            # Fallback color if not found
            return '#CCCCCC' if format == 'hex' else (204, 204, 204)
        
        if format == 'hex':
            return color_info['hex']
        elif format == 'rgb':
            return color_info['rgb']
        elif format == 'bgr':
            # Convert RGB to BGR for OpenCV
            r, g, b = color_info['rgb']
            return (b, g, r)
        else:
            return color_info['rgb']
    
    def get_all_colors(self, format: str = 'hex') -> Dict:
        """
        Get all colors in the specified format
        
        Args:
            format: Color format ('hex', 'rgb', or 'bgr')
            
        Returns:
            Dictionary mapping class names to colors
        """
        result = {}
        for class_id, class_name in self.class_names.items():
            result[class_name] = self.get_color(class_id, format)
        return result
    
    def get_category_colors(self, format: str = 'hex') -> Dict:
        """
        Get base colors for each category
        
        Args:
            format: Color format ('hex', 'rgb', or 'bgr')
            
        Returns:
            Dictionary mapping categories to colors
        """
        result = {}
        for category, color_info in self.color_map['categories'].items():
            if format == 'hex':
                result[category] = color_info['hex']
            elif format == 'bgr':
                r, g, b = color_info['rgb']
                result[category] = (b, g, r)
            else:
                result[category] = color_info['rgb']
        return result
    
    def get_category_for_class(self, class_identifier: Union[int, str]) -> str:
        """
        Get the category for a specific class
        
        Args:
            class_identifier: Class ID (int) or name (str)
            
        Returns:
            Category name
        """
        if isinstance(class_identifier, int):
            return self.color_map['by_id'].get(class_identifier, {}).get('category', 'other')
        else:
            return self.color_map['by_name'].get(class_identifier, {}).get('category', 'other')