shape_overlay.py

import cv2
import numpy as np
from enum import Enum

class ShapeType(Enum):
    CIRCLE = "circle"
    RECTANGLE = "rectangle"
    ARROW = "arrow"
    TRIANGLE = "triangle"
    CROSS = "cross"
    DIAMOND = "diamond"

class ShapeOverlay:
    def __init__(self):
        self.shapes = {}  # Store shape configurations for each track
        
    def set_shape_for_track(self, track_id, shape_type, size=20, color=(0, 255, 0), thickness=2):
        """Set the overlay shape for a specific track"""
        self.shapes[track_id] = {
            'type': shape_type,
            'size': size,
            'color': color,
            'thickness': thickness
        }
    
    def draw_shape(self, frame, center, track_id, default_shape=ShapeType.CIRCLE):
        """Draw the overlay shape for a track"""
        if track_id not in self.shapes:
            # Use default shape if not configured
            shape_config = {
                'type': default_shape,
                'size': 20,
                'color': (0, 255, 0),
                'thickness': 2
            }
        else:
            shape_config = self.shapes[track_id]
        
        shape_type = shape_config['type']
        size = shape_config['size']
        color = shape_config['color']
        thickness = shape_config['thickness']
        
        x, y = center
        
        if shape_type == ShapeType.CIRCLE:
            self._draw_circle(frame, (x, y), size, color, thickness)
        elif shape_type == ShapeType.RECTANGLE:
            self._draw_rectangle(frame, (x, y), size, color, thickness)
        elif shape_type == ShapeType.ARROW:
            self._draw_arrow(frame, (x, y), size, color, thickness)
        elif shape_type == ShapeType.TRIANGLE:
            self._draw_triangle(frame, (x, y), size, color, thickness)
        elif shape_type == ShapeType.CROSS:
            self._draw_cross(frame, (x, y), size, color, thickness)
        elif shape_type == ShapeType.DIAMOND:
            self._draw_diamond(frame, (x, y), size, color, thickness)
    
    def _draw_circle(self, frame, center, size, color, thickness):
        """Draw a circle overlay"""
        cv2.circle(frame, center, size, color, thickness)
    
    def _draw_rectangle(self, frame, center, size, color, thickness):
        """Draw a rectangle overlay"""
        x, y = center
        half_size = size // 2
        pt1 = (x - half_size, y - half_size)
        pt2 = (x + half_size, y + half_size)
        cv2.rectangle(frame, pt1, pt2, color, thickness)
    
    def _draw_arrow(self, frame, center, size, color, thickness):
        """Draw an arrow overlay pointing upward"""
        x, y = center
        # Arrow body (vertical line)
        cv2.line(frame, (x, y + size//2), (x, y - size//2), color, thickness)
        # Arrow head (triangle)
        points = np.array([
            [x, y - size//2],
            [x - size//4, y - size//4],
            [x + size//4, y - size//4]
        ], np.int32)
        cv2.fillPoly(frame, [points], color)
    
    def _draw_triangle(self, frame, center, size, color, thickness):
        """Draw a triangle overlay"""
        x, y = center
        points = np.array([
            [x, y - size//2],
            [x - size//2, y + size//2],
            [x + size//2, y + size//2]
        ], np.int32)
        cv2.polylines(frame, [points], True, color, thickness)
    
    def _draw_cross(self, frame, center, size, color, thickness):
        """Draw a cross overlay"""
        x, y = center
        half_size = size // 2
        # Horizontal line
        cv2.line(frame, (x - half_size, y), (x + half_size, y), color, thickness)
        # Vertical line
        cv2.line(frame, (x, y - half_size), (x, y + half_size), color, thickness)
    
    def _draw_diamond(self, frame, center, size, color, thickness):
        """Draw a diamond overlay"""
        x, y = center
        half_size = size // 2
        points = np.array([
            [x, y - half_size],
            [x + half_size, y],
            [x, y + half_size],
            [x - half_size, y]
        ], np.int32)
        cv2.polylines(frame, [points], True, color, thickness)

class AdvancedOverlay:
    def __init__(self):
        self.shape_overlay = ShapeOverlay()
        self.team_colors = {
            'team_a': (255, 0, 0),    # Red
            'team_b': (0, 0, 255),    # Blue
            'referee': (255, 255, 0), # Yellow
            'unknown': (0, 255, 0)    # Green
        }
        self.player_teams = {}  # Map track_id to team
        
    def assign_team(self, track_id, team):
        """Assign a player to a team"""
        self.player_teams[track_id] = team
        
        # Set shape color based on team
        if team in self.team_colors:
            color = self.team_colors[team]
            self.shape_overlay.set_shape_for_track(
                track_id, ShapeType.CIRCLE, color=color
            )
    
    def draw_player_overlay(self, frame, track_id, center, bbox=None, 
                           show_bbox=True, show_shape=True, show_id=True, 
                           show_team=False):
        """Draw comprehensive overlay for a player"""
        team = self.player_teams.get(track_id, 'unknown')
        color = self.team_colors[team]
        
        # Draw bounding box
        if show_bbox and bbox is not None:
            x, y, w, h = bbox
            cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
        
        # Draw shape overlay
        if show_shape:
            self.shape_overlay.draw_shape(frame, center, track_id)
        
        # Draw player ID
        if show_id:
            cv2.putText(frame, f"P{track_id}", 
                       (center[0] - 15, center[1] - 25),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
        
        # Draw team info
        if show_team:
            cv2.putText(frame, team.upper(), 
                       (center[0] - 20, center[1] + 30),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
    
    def draw_field_overlay(self, frame, field_lines=None):
        """Draw field lines and markings"""
        if field_lines is None:
            # Default field overlay (simple center line)
            height, width = frame.shape[:2]
            cv2.line(frame, (width//2, 0), (width//2, height), (255, 255, 255), 2)
        else:
            # Draw custom field lines
            for line in field_lines:
                pt1, pt2, color, thickness = line
                cv2.line(frame, pt1, pt2, color, thickness)
    
    def draw_statistics_overlay(self, frame, stats):
        """Draw game statistics overlay"""
        y_offset = 30
        for key, value in stats.items():
            text = f"{key}: {value}"
            cv2.putText(frame, text, (10, y_offset), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
            y_offset += 30