visualizer.py

import cv2
import numpy as np

# Dimensions du terrain en yards 
FIELD_LENGTH_YARDS = 114.83
FIELD_WIDTH_YARDS = 74.37

# Constantes de taille d'image attendue
EXPECTED_H, EXPECTED_W = 720, 1280

# Import des constantes d'indices depuis pose_estimator si nécessaire (ou les redéfinir ici)
from pose_estimator import (LEFT_ANKLE_KP_INDEX, RIGHT_ANKLE_KP_INDEX, 
                            CONFIDENCE_THRESHOLD_KEYPOINTS, DEFAULT_MARKER_COLOR, SKELETON_EDGES, SKELETON_THICKNESS)

# Constantes pour les marqueurs
MARKER_RADIUS = 6
MARKER_BORDER_THICKNESS = 1
MARKER_BORDER_COLOR = (0, 0, 0) # Noir
# Constantes pour le ballon
BALL_MARKER_RADIUS = 5
BALL_MARKER_COLOR = (255, 255, 255) # Blanc
BALL_BORDER_COLOR = (0, 0, 0) # Noir
BALL_BORDER_THICKNESS = 1

# Plage de modulation pour l'échelle dynamique inverseé
DYNAMIC_SCALE_MIN_MODULATION = 0.4 # Pour les joueurs les plus loin (haut de la minimap)
DYNAMIC_SCALE_MAX_MODULATION = 1.6 # Pour les joueurs les plus près (bas de la minimap)

def calculate_dynamic_scale(y_position, frame_height, min_scale=1.0, max_scale=2):
    """Calcule le facteur d'échelle en fonction de la position verticale (non utilisé dans cette version simplifiée)."""
    normalized_position = y_position / frame_height
    return min_scale + (max_scale - min_scale) * normalized_position

def _prepare_minimap_base(minimap_size=(EXPECTED_W, EXPECTED_H)):
    """Prépare le fond de la minimap (vert texturé verticalement dans la zone terrain) et calcule les métriques du terrain."""
    minimap_h, minimap_w = minimap_size[1], minimap_size[0]
    
    # Définir les couleurs et la largeur des bandes verticales
    base_green = (0, 60, 0)    # Vert foncé (fond)
    stripe_green = (0, 70, 0)   # 
    stripe_width = 5            # Largeur de chaque bande verticale (pixels)
    
    # Initialiser TOUTE la minimap avec la couleur de base
    minimap_bgr = np.full((minimap_h, minimap_w, 3), base_green, dtype=np.uint8)
    
    # --- Calculer les métriques et les limites du terrain D'ABORD --- 
    scale_x = minimap_w / FIELD_LENGTH_YARDS
    scale_y = minimap_h / FIELD_WIDTH_YARDS
    scale = min(scale_x, scale_y) * 0.9  # Marge
    
    field_width_px = int(FIELD_WIDTH_YARDS * scale)
    field_length_px = int(FIELD_LENGTH_YARDS * scale)
    offset_x = (minimap_w - field_length_px) // 2
    offset_y = (minimap_h - field_width_px) // 2

    # --- Dessiner les bandes VERTICALES alternées UNIQUEMENT dans la zone du terrain --- 
    for x in range(offset_x, offset_x + field_length_px, stripe_width * 2):
        # Coordonnées du rectangle vertical pour la bande claire
        start_x = x
        end_x = min(x + stripe_width, offset_x + field_length_px) # Ne pas dépasser la limite droite
        start_y = offset_y
        end_y = offset_y + field_width_px
        
        cv2.rectangle(minimap_bgr, (start_x, start_y), (end_x, end_y), stripe_green, thickness=-1)
        # La bande foncée suivante est déjà là (couleur de base)

    # --- Préparer la matrice S et les métriques à retourner --- 
    S = np.array([
        [scale, 0, offset_x],
        [0, scale, offset_y],
        [0, 0, 1]
    ], dtype=np.float32)
    
    metrics = {
        "scale": scale,
        "offset_x": offset_x,
        "offset_y": offset_y,
        "field_width_px": field_width_px,
        "field_length_px": field_length_px,
        "S": S
    }
    return minimap_bgr, metrics

def _draw_field_lines(minimap_bgr, metrics):
    """Dessine les lignes du terrain et les buts sur la minimap."""
    scale = metrics["scale"]
    offset_x = metrics["offset_x"]
    offset_y = metrics["offset_y"]
    field_width_px = metrics["field_width_px"]
    field_length_px = metrics["field_length_px"]
    
    line_color = (255, 255, 255) # Blanc
    line_thickness = 1 
    goal_thickness = 1 # Épaisseur pour les poteaux de but
    goal_width_yards = 8 # Largeur standard du but

    center_x = offset_x + field_length_px // 2
    center_y = offset_y + field_width_px // 2
    penalty_area_width_px = int(SoccerPitchSN.PENALTY_AREA_WIDTH * scale)
    penalty_area_length_px = int(SoccerPitchSN.PENALTY_AREA_LENGTH * scale)
    goal_area_width_px = int(SoccerPitchSN.GOAL_AREA_WIDTH * scale)
    goal_area_length_px = int(SoccerPitchSN.GOAL_AREA_LENGTH * scale)
    center_circle_radius_px = int(SoccerPitchSN.CENTER_CIRCLE_RADIUS * scale)
    goal_width_px = int(goal_width_yards * scale)

    # Dessiner les lignes principales
    cv2.rectangle(minimap_bgr, (offset_x, offset_y), (offset_x + field_length_px, offset_y + field_width_px), line_color, line_thickness)
    cv2.line(minimap_bgr, (center_x, offset_y), (center_x, offset_y + field_width_px), line_color, line_thickness)
    cv2.circle(minimap_bgr, (center_x, center_y), center_circle_radius_px, line_color, line_thickness)
    cv2.circle(minimap_bgr, (center_x, center_y), 3, line_color, -1) # Point central
    cv2.rectangle(minimap_bgr, (offset_x, center_y - penalty_area_width_px//2), (offset_x + penalty_area_length_px, center_y + penalty_area_width_px//2), line_color, line_thickness)
    cv2.rectangle(minimap_bgr, (offset_x + field_length_px - penalty_area_length_px, center_y - penalty_area_width_px//2), (offset_x + field_length_px, center_y + penalty_area_width_px//2), line_color, line_thickness)
    cv2.rectangle(minimap_bgr, (offset_x, center_y - goal_area_width_px//2), (offset_x + goal_area_length_px, center_y + goal_area_width_px//2), line_color, line_thickness)
    cv2.rectangle(minimap_bgr, (offset_x + field_length_px - goal_area_length_px, center_y - goal_area_width_px//2), (offset_x + field_length_px, center_y + goal_area_width_px//2), line_color, line_thickness)

    # Dessiner les buts (rectangles épais sur les lignes de but)
    goal_y_top = center_y - goal_width_px // 2
    goal_y_bottom = center_y + goal_width_px // 2
    # But gauche
    cv2.rectangle(minimap_bgr, (offset_x-6 - goal_thickness // 2, goal_y_top), (offset_x + goal_thickness // 2, goal_y_bottom), line_color, thickness=goal_thickness)
    # But droit
    cv2.rectangle(minimap_bgr, (offset_x + field_length_px - goal_thickness // 2, goal_y_top), (offset_x +6 + field_length_px + goal_thickness // 2, goal_y_bottom), line_color, thickness=goal_thickness)

def create_minimap_view(image_rgb, homography, minimap_size=(EXPECTED_W, EXPECTED_H)):
    """Crée une vue minimap avec l'image RGB originale projetée et les lignes du terrain.
    
    Args:
        image_rgb: Image source en format RGB (720p attendu).
        homography: Matrice d'homographie (numpy array) pour projeter l'image.
        minimap_size: Taille de la minimap de sortie (largeur, hauteur).
        
    Returns:
        L'image de la minimap (numpy array BGR) ou None si l'homographie est invalide.
    """
    if homography is None:
        print("Avertissement : Homographie invalide, impossible de créer la minimap (vue originale).")
        return None

    h, w = image_rgb.shape[:2]
    if h != EXPECTED_H or w != EXPECTED_W:
        print(f"Avertissement : L'image RGB d'entrée n'est pas en {EXPECTED_W}x{EXPECTED_H}, redimensionnement...")
        image_rgb = cv2.resize(image_rgb, (EXPECTED_W, EXPECTED_H), interpolation=cv2.INTER_LINEAR)
    
    minimap_bgr, metrics = _prepare_minimap_base(minimap_size)
    S = metrics["S"]

    try:
        overlay = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
        overlay = cv2.convertScaleAbs(overlay, alpha=1.2, beta=10) 
        overlay = cv2.addWeighted(overlay, 0.5, np.zeros_like(overlay), 0.5, 0)
        
        H_minimap = S @ homography 
        warped = cv2.warpPerspective(overlay, H_minimap, minimap_size, flags=cv2.INTER_LINEAR)
        
        mask = cv2.cvtColor(warped, cv2.COLOR_BGR2GRAY)
        _, mask = cv2.threshold(mask, 1, 255, cv2.THRESH_BINARY)
        
        minimap_bgr = np.where(mask[..., None] > 0, warped, minimap_bgr)
        
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        cv2.drawContours(minimap_bgr, contours, -1, (255, 255, 255), 2) 
                
    except Exception as e:
        print(f"Erreur lors de la projection sur la mini-carte (vue originale) : {str(e)}")

    _draw_field_lines(minimap_bgr, metrics)
    return minimap_bgr

def create_minimap_with_offset_skeletons(player_data_list, homography, 
                                         base_skeleton_scale: float, 
                                         ball_ref_point_img: np.ndarray | None = None, # Ajout du paramètre optionnel pour le ballon
                                         minimap_size=(EXPECTED_W, EXPECTED_H)) -> tuple[np.ndarray | None, float | None]:
    """Crée une vue minimap en dessinant le squelette original (réduit/agrandi dynamiquement et inversé) 
       à la position projetée du joueur, trié par position Y. Dessine aussi le ballon s'il est fourni.
    
    Args:
        player_data_list: Liste de dictionnaires retournée par get_player_data.
        homography: Matrice d'homographie (numpy array).
        base_skeleton_scale: Facteur d'échelle de base pour dessiner les squelettes.
        ball_ref_point_img: Point de référence (x, y) du ballon dans l'image originale (optionnel).
        minimap_size: Taille de la minimap de sortie (largeur, hauteur).
        
    Returns:
        Tuple: (L'image de la minimap (numpy array BGR) ou None, Échelle moyenne appliquée ou None)
    """
    if homography is None:
        print("Avertissement : Homographie invalide, impossible de créer la minimap (squelettes décalés).")
        return None, None # Retourner None pour l'image et l'échelle

    minimap_bgr, metrics = _prepare_minimap_base(minimap_size)
    
    # --- Dessiner les lignes du terrain D'ABORD --- 
    _draw_field_lines(minimap_bgr, metrics)
    
    S = metrics["S"]
    H_minimap = S @ homography 

    players_to_draw = [] # Liste pour stocker les joueurs valides avec leur position Y

    # --- Étape 1 & 2: Calculer la position projetée pour tous les joueurs valides --- 
    for p_data in player_data_list:
        kps_img = p_data['keypoints'] 
        scores = p_data['scores']
        bbox = p_data['bbox'] 
        color = p_data['avg_color']

        # -- Calculer le point de référence sur l'image --
        l_ankle_pt = kps_img[LEFT_ANKLE_KP_INDEX]
        r_ankle_pt = kps_img[RIGHT_ANKLE_KP_INDEX]
        l_ankle_score = scores[LEFT_ANKLE_KP_INDEX]
        r_ankle_score = scores[RIGHT_ANKLE_KP_INDEX]
        ref_point_img = None
        if l_ankle_score > CONFIDENCE_THRESHOLD_KEYPOINTS and r_ankle_score > CONFIDENCE_THRESHOLD_KEYPOINTS:
            ref_point_img = (l_ankle_pt + r_ankle_pt) / 2
        elif l_ankle_score > CONFIDENCE_THRESHOLD_KEYPOINTS:
            ref_point_img = l_ankle_pt
        elif r_ankle_score > CONFIDENCE_THRESHOLD_KEYPOINTS:
            ref_point_img = r_ankle_pt
        else:
            x1, _, x2, y2 = bbox
            ref_point_img = np.array([(x1 + x2) / 2, y2], dtype=np.float32) 
        if ref_point_img is None: continue 

        # -- Projeter ce point de référence sur la minimap --
        try:
            point_to_transform = np.array([[ref_point_img]], dtype=np.float32)
            projected_point = cv2.perspectiveTransform(point_to_transform, H_minimap)
            mx, my = map(int, projected_point[0, 0])
            h_map, w_map = minimap_bgr.shape[:2]
            if not (0 <= mx < w_map and 0 <= my < h_map):
                 continue # Ignorer si hors des limites de la minimap
        except Exception as e:
            # print(f"Erreur lors de la projection du point de référence {ref_point_img}: {e}") # Optionnel: décommenter pour debug
            continue 
        
        # Stocker les données nécessaires pour le tri et le dessin
        players_to_draw.append({
            'data': p_data,
            'mx': mx,
            'my': my,
            'ref_point': ref_point_img
        })

    # --- Étape 3: Trier les joueurs par position Y (ordre croissant) --- 
    # Ceux avec Y plus petit (plus haut) seront dessinés en premier
    players_to_draw.sort(key=lambda p: p['my'])

    # Variables pour calculer l'échelle moyenne appliquée
    total_applied_scale = 0.0
    drawn_players_count = 0

    # --- Étape 4: Dessiner les joueurs dans l'ordre trié (MAINTENANT AU-DESSUS DES LIGNES) --- 
    for player_info in players_to_draw:
        p_data = player_info['data']
        mx = player_info['mx']
        my = player_info['my']
        ref_point_img = player_info['ref_point']
        kps_img = p_data['keypoints']
        scores = p_data['scores']
        # color = p_data['avg_color'] # Ignorer la couleur calculée
        drawing_color = (0, 0, 0) # Utiliser le noir pour tous les joueurs

        # -- Calculer l'échelle dynamique INVERSÉE pour CE joueur --
        minimap_height = minimap_bgr.shape[0]
        if minimap_height == 0: continue 
        ref_y_normalized = my / minimap_height 
        dynamic_modulation = DYNAMIC_SCALE_MIN_MODULATION + \
                             (DYNAMIC_SCALE_MAX_MODULATION - DYNAMIC_SCALE_MIN_MODULATION) * (1.0 - ref_y_normalized)
        dynamic_modulation = np.clip(dynamic_modulation, DYNAMIC_SCALE_MIN_MODULATION * 0.8, DYNAMIC_SCALE_MAX_MODULATION * 1.2)
        final_draw_scale = base_skeleton_scale * dynamic_modulation

        # Ajouter à la somme pour la moyenne
        total_applied_scale += final_draw_scale
        drawn_players_count += 1

        # -- Dessiner le squelette --
        kps_relative_to_ref = kps_img - ref_point_img
        for kp_idx1, kp_idx2 in SKELETON_EDGES:
            if scores[kp_idx1] > CONFIDENCE_THRESHOLD_KEYPOINTS and scores[kp_idx2] > CONFIDENCE_THRESHOLD_KEYPOINTS:
                pt1_map_offset = (mx, my) + kps_relative_to_ref[kp_idx1] * final_draw_scale 
                pt2_map_offset = (mx, my) + kps_relative_to_ref[kp_idx2] * final_draw_scale
                pt1_draw = tuple(map(int, pt1_map_offset))
                pt2_draw = tuple(map(int, pt2_map_offset))
                # Vérifier si les points sont dans les limites avant de dessiner (sécurité)
                h_map, w_map = minimap_bgr.shape[:2]
                if (0 <= pt1_draw[0] < w_map and 0 <= pt1_draw[1] < h_map and
                    0 <= pt2_draw[0] < w_map and 0 <= pt2_draw[1] < h_map):
                    cv2.line(minimap_bgr, pt1_draw, pt2_draw, drawing_color, SKELETON_THICKNESS, cv2.LINE_AA) # Utiliser drawing_color (noir)

    # --- Étape 5: Dessiner le ballon (si trouvé et projeté) --- 
    ball_mx, ball_my = None, None
    if ball_ref_point_img is not None:
        try:
            point_to_transform = np.array([[ball_ref_point_img]], dtype=np.float32)
            projected_ball_point = cv2.perspectiveTransform(point_to_transform, H_minimap)
            ball_mx, ball_my = map(int, projected_ball_point[0, 0])
            h_map, w_map = minimap_bgr.shape[:2]
            if not (0 <= ball_mx < w_map and 0 <= ball_my < h_map):
                 print("Avertissement : Ballon projeté hors des limites de la minimap.")
                 ball_mx, ball_my = None, None # Ne pas dessiner
            else:
                 # Dessiner la bordure noire
                 cv2.circle(minimap_bgr, (ball_mx, ball_my), BALL_MARKER_RADIUS + BALL_BORDER_THICKNESS, 
                            BALL_BORDER_COLOR, -1, cv2.LINE_AA)
                 # Dessiner le cercle blanc intérieur
                 cv2.circle(minimap_bgr, (ball_mx, ball_my), BALL_MARKER_RADIUS, 
                            BALL_MARKER_COLOR, -1, cv2.LINE_AA)
        except Exception as e:
            print(f"Erreur lors de la projection ou du dessin du ballon : {e}")
            ball_mx, ball_my = None, None

    # Calculer l'échelle moyenne si des joueurs ont été dessinés
    final_avg_scale = total_applied_scale / drawn_players_count if drawn_players_count > 0 else None

    return minimap_bgr, final_avg_scale

# Définition simplifiée de SoccerPitchSN juste pour les constantes de dimension
# (pour éviter d'importer toute la classe complexe)
class SoccerPitchSN:
    GOAL_LINE_TO_PENALTY_MARK = 11.0
    PENALTY_AREA_WIDTH = 42
    PENALTY_AREA_LENGTH = 19
    GOAL_AREA_WIDTH = 18.32
    GOAL_AREA_LENGTH = 5.5
    CENTER_CIRCLE_RADIUS = 10