drawers/utils.py

"""
A utility module providing functions for drawing shapes on video frames.

This module includes functions to draw triangles and ellipses on frames, which can be used
to represent various annotations such as player positions or ball locations in sports analysis.
"""

import cv2 
import numpy as np
import sys 
sys.path.append('../')
from utils import get_center_of_bbox, get_bbox_width, get_foot_position

def draw_traingle(frame, bbox, color):
    """
    Draws a sharp, premium triangle on the frame.
    """
    y = int(bbox[1])
    x, _ = get_center_of_bbox(bbox)

    triangle_points = np.array([
        [x, y],
        [x - 12, y - 24],
        [x + 12, y - 24],
    ])
    
    # Outer glow / shadow
    cv2.drawContours(frame, [triangle_points + [0, 2]], 0, (0, 0, 0), cv2.FILLED)
    
    # Main triangle
    cv2.drawContours(frame, [triangle_points], 0, color, cv2.FILLED)
    cv2.drawContours(frame, [triangle_points], 0, (255, 255, 255), 1, cv2.LINE_AA)

    return frame

def draw_ellipse(frame, bbox, color, track_id=None):
    """
    Draws a premium ellipse with a subtle glow and a professional track ID box.
    """
    y2 = int(bbox[3])
    x_center, _ = get_center_of_bbox(bbox)
    width = get_bbox_width(bbox)
    
    axes = (int(width), int(0.35 * width))
    
    # Draw outer glow (semi-transparent)
    overlay = frame.copy()
    cv2.ellipse(
        overlay,
        center=(x_center, y2),
        axes=(axes[0] + 4, axes[1] + 2),
        angle=0.0,
        startAngle=-45,
        endAngle=235,
        color=(0, 0, 0),
        thickness=4,
        lineType=cv2.LINE_AA
    )
    cv2.addWeighted(overlay, 0.4, frame, 0.6, 0, frame)

    # Main highlight ellipse
    cv2.ellipse(
        frame,
        center=(x_center, y2),
        axes=axes,
        angle=0.0,
        startAngle=-45,
        endAngle=235,
        color=color,
        thickness=2,
        lineType=cv2.LINE_AA
    )

    if track_id is not None:
        # Premium ID Box
        rectangle_width = 44
        rectangle_height = 22
        x1_rect = int(x_center - rectangle_width // 2)
        y1_rect = int(y2 + 10)
        
        # Shadow for box
        cv2.rectangle(frame, (x1_rect + 2, y1_rect + 2), (x1_rect + rectangle_width + 2, y1_rect + rectangle_height + 2), (0,0,0), -1)
        # Main box
        cv2.rectangle(frame, (x1_rect, y1_rect), (x1_rect + rectangle_width, y1_rect + rectangle_height), color, -1)
        # White thin border
        cv2.rectangle(frame, (x1_rect, y1_rect), (x1_rect + rectangle_width, y1_rect + rectangle_height), (255,255,255), 1, cv2.LINE_AA)
        
        text = str(track_id)
        (tw, th), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_DUPLEX, 0.5, 1)
        tx = x1_rect + (rectangle_width - tw) // 2
        ty = y1_rect + (rectangle_height + th) // 2
        cv2.putText(frame, text, (tx, ty), cv2.FONT_HERSHEY_DUPLEX, 0.5, (0,0,0), 1, cv2.LINE_AA)

    return frame

def draw_rounded_rect(frame, rect, color, thickness=1, radius=10):
    """
    Draws a rounded rectangle on the given frame using anti-aliased lines.
    Supports filled rounded rectangle if thickness < 0.
    """
    x, y, w, h = map(int, rect)
    
    if thickness < 0:
        cv2.rectangle(frame, (x + radius, y), (x + w - radius, y + h), color, -1)
        cv2.rectangle(frame, (x, y + radius), (x + w, y + h - radius), color, -1)
        cv2.circle(frame, (x + radius, y + radius), radius, color, -1)
        cv2.circle(frame, (x + w - radius, y + radius), radius, color, -1)
        cv2.circle(frame, (x + radius, y + h - radius), radius, color, -1)
        cv2.circle(frame, (x + w - radius, y + h - radius), radius, color, -1)
        return frame

    # Top-left corner
    cv2.ellipse(frame, (x + radius, y + radius), (radius, radius), 180, 0, 90, color, thickness, cv2.LINE_AA)
    # Top-right corner
    cv2.ellipse(frame, (x + w - radius, y + radius), (radius, radius), 270, 0, 90, color, thickness, cv2.LINE_AA)
    # Bottom-right corner
    cv2.ellipse(frame, (x + w - radius, y + h - radius), (radius, radius), 0, 0, 90, color, thickness, cv2.LINE_AA)
    # Bottom-left corner
    cv2.ellipse(frame, (x + radius, y + h - radius), (radius, radius), 90, 0, 90, color, thickness, cv2.LINE_AA)
    
    # Lines
    cv2.line(frame, (x + radius, y), (x + w - radius, y), color, thickness, cv2.LINE_AA)
    cv2.line(frame, (x, y + radius), (x, y + h - radius), color, thickness, cv2.LINE_AA)
    cv2.line(frame, (x + w, y + radius), (x + w, y + h - radius), color, thickness, cv2.LINE_AA)
    cv2.line(frame, (x + radius, y + h), (x + w - radius, y + h), color, thickness, cv2.LINE_AA)
    
    return frame

def draw_glass_panel(frame, rect, alpha=0.7, color=(15, 12, 10), radius=15):
    """
    Draws a premium glassmorphic panel.
    """
    x, y, w, h = map(int, rect)
    overlay = frame.copy()
    
    # Fill background
    cv2.rectangle(overlay, (x + radius, y), (x + w - radius, y + h), color, -1)
    cv2.rectangle(overlay, (x, y + radius), (x + w, y + h - radius), color, -1)
    cv2.circle(overlay, (x + radius, y + radius), radius, color, -1)
    cv2.circle(overlay, (x + w - radius, y + radius), radius, color, -1)
    cv2.circle(overlay, (x + radius, y + h - radius), radius, color, -1)
    cv2.circle(overlay, (x + w - radius, y + h - radius), radius, color, -1)
    
    # Apply alpha blending
    frame_new = cv2.addWeighted(overlay, alpha, frame, 1 - alpha, 0)
    
    # Mask area outside panel to preserve original
    mask = np.zeros(frame.shape[:2], dtype=np.uint8)
    cv2.rectangle(mask, (x + radius, y), (x + w - radius, y + h), 255, -1)
    cv2.rectangle(mask, (x, y + radius), (x + w, y + h - radius), 255, -1)
    cv2.circle(mask, (x + radius, y + radius), radius, 255, -1)
    cv2.circle(mask, (x + w - radius, y + radius), radius, 255, -1)
    cv2.circle(mask, (x + radius, y + h - radius), radius, 255, -1)
    cv2.circle(mask, (x + w - radius, y + h - radius), radius, 255, -1)
    
    mask_bool = mask > 0
    frame[mask_bool] = frame_new[mask_bool]
    
    # Add thin silver border
    draw_rounded_rect(frame, (x, y, w, h), (200, 200, 200), thickness=1, radius=radius)
    
    return frame

def draw_text_with_shadow(frame, text, pos, font_scale=0.6, color=(255, 255, 255), thickness=1):
    """
    Draws text with a small black shadow for readability.
    """
    # Shadow
    cv2.putText(frame, text, (pos[0]+1, pos[1]+1), cv2.FONT_HERSHEY_DUPLEX, font_scale, (0, 0, 0), thickness+1, cv2.LINE_AA)
    # Main text
    cv2.putText(frame, text, pos, cv2.FONT_HERSHEY_DUPLEX, font_scale, color, thickness, cv2.LINE_AA)
    return frame