import os import cv2 import numpy as np import pandas as pd import matplotlib matplotlib.use('Agg') # Headless backend for matplotlib import matplotlib.pyplot as plt from mplsoccer import Pitch import football_analytics.config as config class FootballVisualizer: def __init__(self): # Color palettes (BGR for OpenCV) self.colors = { 0: (230, 50, 50), # Team A: Bright Red/Coral 1: (50, 100, 240), # Team B: Bright Blue 2: (50, 240, 240), # Referee/Other: Yellow -1: (200, 200, 200), # Unknown: Light Gray "ball": (0, 255, 255), # Ball: Neon Yellow/Green "radar_bg": (34, 49, 43), # Dark Green Pitch Background "radar_line": (200, 215, 205), # Pitch Lines "possession_glow": (0, 255, 0) # Glowing green for player in possession } # Dimensions of the radar overlay self.radar_w = 320 self.radar_h = 200 self.radar_margin = 20 def draw_pitch_radar(self, players, ball_pos): """ Draws a 2D mini-pitch radar representation. Args: players (list of dict): [{id, team, x_pitch, y_pitch}] ball_pos (tuple): (x_pitch, y_pitch) or None Returns: np.ndarray: Mini-pitch image of size (radar_h, radar_w, 3) """ # Create dark green pitch background radar = np.zeros((self.radar_h, self.radar_w, 3), dtype=np.uint8) radar[:] = self.colors["radar_bg"] # Scale factors scale_x = self.radar_w / config.PITCH_WIDTH scale_y = self.radar_h / config.PITCH_HEIGHT def to_radar_coords(x, y): rx = int(x * scale_x) ry = int(y * scale_y) return rx, ry # Draw pitch boundaries cv2.rectangle(radar, (0, 0), (self.radar_w - 1, self.radar_h - 1), self.colors["radar_line"], 2) # Draw half-way line mid_x = int(self.radar_w / 2) cv2.line(radar, (mid_x, 0), (mid_x, self.radar_h), self.colors["radar_line"], 1) # Draw center circle center_circle_r = int(9.15 * scale_x) cv2.circle(radar, (mid_x, int(self.radar_h / 2)), center_circle_r, self.colors["radar_line"], 1) cv2.circle(radar, (mid_x, int(self.radar_h / 2)), 2, self.colors["radar_line"], -1) # Draw penalty areas (UEFA rules: 16.5m deep, 40.32m wide) # Left penalty area pa_w = int(16.5 * scale_x) pa_h = int(40.32 * scale_y) pa_y1 = int((config.PITCH_HEIGHT - 40.32) / 2 * scale_y) pa_y2 = pa_y1 + pa_h cv2.rectangle(radar, (0, pa_y1), (pa_w, pa_y2), self.colors["radar_line"], 1) # Goal area left ga_w = int(5.5 * scale_x) ga_h = int(18.32 * scale_y) ga_y1 = int((config.PITCH_HEIGHT - 18.32) / 2 * scale_y) ga_y2 = ga_y1 + ga_h cv2.rectangle(radar, (0, ga_y1), (ga_w, ga_y2), self.colors["radar_line"], 1) # Right penalty area cv2.rectangle(radar, (self.radar_w - pa_w, pa_y1), (self.radar_w, pa_y2), self.colors["radar_line"], 1) # Goal area right cv2.rectangle(radar, (self.radar_w - ga_w, ga_y1), (self.radar_w, ga_y2), self.colors["radar_line"], 1) # Draw players on radar for p in players: team = p["team"] px, py = p["x_pitch"], p["y_pitch"] rx, ry = to_radar_coords(px, py) # Skip if invalid coords if not (0 <= rx < self.radar_w and 0 <= ry < self.radar_h): continue color = self.colors.get(team, self.colors[-1]) # Draw player dot cv2.circle(radar, (rx, ry), 5, color, -1) cv2.circle(radar, (rx, ry), 6, (255, 255, 255), 1) # Draw ball on radar if ball_pos is not None: bx, by = ball_pos rx, ry = to_radar_coords(bx, by) if 0 <= rx < self.radar_w and 0 <= ry < self.radar_h: cv2.circle(radar, (rx, ry), 4, self.colors["ball"], -1) cv2.circle(radar, (rx, ry), 5, (0, 0, 0), 1) return radar def annotate_frame(self, frame, detections, team_assignments, ball_xyxy, current_possession_player, draw_possession=True): """ Annotates the frame with bounding boxes, team color coding, ball info, and overlays the 2D radar. Args: frame (np.ndarray): Input image frame. detections (sv.Detections): Player tracked detections. team_assignments (dict): track_id -> team_id. ball_xyxy (np.ndarray): Ball bounding box [x1, y1, x2, y2] or None. current_possession_player (int): ID of player currently in possession. Returns: np.ndarray: Annotated frame. """ annotated_frame = frame.copy() h_frame, w_frame, _ = frame.shape # 1. Annotate players for i, bbox in enumerate(detections.xyxy): track_id = int(detections.tracker_id[i]) if detections.tracker_id is not None else -1 class_id = int(detections.class_id[i]) if class_id != 0: # Skip ball (handle separately) continue team = team_assignments.get(track_id, -1) if team == 3: # Skip drawing staff/coaches completely continue color = self.colors.get(team, self.colors[-1]) x1, y1, x2, y2 = map(int, bbox) # Draw semi-transparent ellipse at the feet of the player for a premium tactical look feet_y = y2 center_x = int((x1 + x2) / 2) axis_x = int((x2 - x1) / 2) axis_y = int(axis_x / 3) # Perspective flattening # Make a glowing ellipse for possession is_possessor = (track_id == current_possession_player and track_id != -1) ellipse_color = self.colors["possession_glow"] if is_possessor else color thickness = 3 if is_possessor else 2 # Draw ellipse feet base cv2.ellipse(annotated_frame, (center_x, feet_y), (axis_x, axis_y), 0, 0, 360, ellipse_color, thickness) # Draw bounding box (subtle corner-only or thin lines) cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), color, 1) # Draw label background label_text = f"ID:{track_id}" if team == 0: label_text = f"A-{track_id}" elif team == 1: label_text = f"B-{track_id}" elif team == 2: label_text = f"REF" (w, h_text), _ = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.4, 1) cv2.rectangle(annotated_frame, (x1, y1 - h_text - 4), (x1 + w + 4, y1), color, -1) cv2.putText(annotated_frame, label_text, (x1 + 2, y1 - 2), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1, cv2.LINE_AA) # 2. Annotate Ball if ball_xyxy is not None and len(ball_xyxy) > 0: if ball_xyxy.ndim == 2: ball_box = ball_xyxy[0] else: ball_box = ball_xyxy bx1, by1, bx2, by2 = map(int, ball_box) # Draw glowing circle around the ball bc_x = int((bx1 + bx2) / 2) bc_y = int((by1 + by2) / 2) r = max(6, int((bx2 - bx1) / 2)) cv2.circle(annotated_frame, (bc_x, bc_y), r + 2, (0, 0, 0), 1) cv2.circle(annotated_frame, (bc_x, bc_y), r, self.colors["ball"], 2) # Draw a pointer above the ball pointer_y = by1 - 8 pointer_pts = np.array([[bc_x, by1 - 2], [bc_x - 4, pointer_y], [bc_x + 4, pointer_y]], dtype=np.int32) cv2.fillPoly(annotated_frame, [pointer_pts], self.colors["ball"]) # 3. Add possession banner overlay at the top middle if draw_possession: banner_w, banner_h = 360, 40 banner_x = int((w_frame - banner_w) / 2) banner_y = 15 # Transparent background for banner overlay = annotated_frame.copy() cv2.rectangle(overlay, (banner_x, banner_y), (banner_x + banner_w, banner_y + banner_h), (25, 25, 25), -1) cv2.addWeighted(overlay, 0.6, annotated_frame, 0.4, 0, annotated_frame) # Add possession text poss_text = "POSSESSION: CONTESTED" text_color = (200, 200, 200) if current_possession_player is not None: team_poss = team_assignments.get(current_possession_player, -1) if team_poss == 0: poss_text = f"POSSESSION: TEAM A (ID {current_possession_player})" text_color = (self.colors[0][2], self.colors[0][1], self.colors[0][0]) elif team_poss == 1: poss_text = f"POSSESSION: TEAM B (ID {current_possession_player})" text_color = (self.colors[1][2], self.colors[1][1], self.colors[1][0]) (w_txt, h_txt), _ = cv2.getTextSize(poss_text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2) tx = banner_x + int((banner_w - w_txt) / 2) ty = banner_y + int((banner_h + h_txt) / 2) cv2.putText(annotated_frame, poss_text, (tx, ty), cv2.FONT_HERSHEY_SIMPLEX, 0.5, text_color, 2, cv2.LINE_AA) return annotated_frame def overlay_radar_on_frame(self, frame, radar_img): """ Overlays the radar image in the top right corner with semi-transparency. """ h_frame, w_frame, _ = frame.shape # Position: Top-right corner with 20px padding x_offset = w_frame - self.radar_w - self.radar_margin y_offset = self.radar_margin # Extract Region of Interest (ROI) from frame roi = frame[y_offset:y_offset + self.radar_h, x_offset:x_offset + self.radar_w] # Blend ROI and radar image (0.7 radar + 0.3 frame background for transparency) blended = cv2.addWeighted(radar_img, 0.75, roi, 0.25, 0) # Put blended radar back onto frame frame[y_offset:y_offset + self.radar_h, x_offset:x_offset + self.radar_w] = blended # Draw simple border around radar cv2.rectangle(frame, (x_offset, y_offset), (x_offset + self.radar_w, y_offset + self.radar_h), (120, 130, 125), 2) return frame def save_team_heatmaps(self, positions_log, output_dir=None): """ Creates and saves kernel density heatmaps for each team. """ out_dir = output_dir or os.path.join(config.OUTPUT_DIR, "heatmaps") os.makedirs(out_dir, exist_ok=True) if not positions_log: print("[Visualizer] Position log empty. Heatmaps not created.") return df = pd.DataFrame(positions_log) # Pitch instance pitch = Pitch(pitch_type='uefa', pitch_color='#22312b', line_color='#c7d5cc', corner_arcs=True) for team_id, name in [(0, "team_A"), (1, "team_B")]: team_df = df[(df["team_id"] == team_id) & (df["object_id"] != -99)] if len(team_df) < 5: print(f"[Visualizer] Not enough positions logged for {name} to draw heatmap.") continue fig, ax = pitch.draw(figsize=(12, 8)) try: # Kernel Density Estimate Plot # sns.kdeplot requires x and y. Note: in mplsoccer Pitch, # x is along the width, y is along the height. # We can use pitch.kdeplot: pitch.kdeplot( team_df["x_pitch"].values, team_df["y_pitch"].values, ax=ax, cmap='hot', fill=True, levels=80, alpha=0.6, zorder=1 ) ax.set_title(f"Tactical Occupation Density - {name.replace('_', ' ').title()}", color='white', fontsize=18, pad=15) fig.patch.set_facecolor('#22312b') # Save plot filepath = os.path.join(out_dir, f"{name}_heatmap.png") fig.savefig(filepath, facecolor='#22312b', edgecolor='none', bbox_inches='tight') print(f"[Visualizer] Saved {name} heatmap to {filepath}") except Exception as e: print(f"[Visualizer] Error creating heatmap for {name}: {e}") finally: plt.close(fig)