feat: Intégration de la détection de ballon avec YOLO et amélioration des vérifications de modèles dans app.py et main.py.

app.py

@@ -5,13 +5,16 @@ import torch
 from pathlib import Path
 import time
 import traceback
+# Import YOLO
+from ultralytics import YOLO
 
 # Importer les éléments nécessaires depuis les autres modules du projet
 try:
     from tvcalib.infer.module import TvCalibInferModule
     # On essaie d'importer la fonction de pré-traitement depuis main.py
     # Si main.py n'est pas conçu pour être importé, il faudra peut-être copier/coller cette fonction ici
-    from main import preprocess_image_tvcalib, IMAGE_SHAPE, SEGMENTATION_MODEL_PATH
+    # Importer aussi les constantes YOLO depuis main.py (ou les redéfinir ici)
+    from main import preprocess_image_tvcalib, IMAGE_SHAPE, SEGMENTATION_MODEL_PATH, YOLO_MODEL_PATH, BALL_CLASS_INDEX
     from visualizer import (
         create_minimap_view,
         create_minimap_with_offset_skeletons,
@@ -40,6 +43,11 @@ if not SEGMENTATION_MODEL_PATH.exists():
     print("L'application risque de ne pas fonctionner. Assurez-vous que le fichier est présent.")
     # Gradio peut quand même démarrer, mais le traitement échouera.
 
+# Vérifier si le modèle YOLO existe aussi
+if not YOLO_MODEL_PATH.exists():
+    print(f"AVERTISSEMENT : Modèle YOLO introuvable : {YOLO_MODEL_PATH}")
+    print("L'application risque de ne pas fonctionner. Assurez-vous que le fichier est présent.")
+
 # --- Fonction Principale de Traitement ---
 def process_image_and_generate_minimaps(input_image_bgr, optim_steps, target_avg_scale):
     """
@@ -54,12 +62,23 @@ def process_image_and_generate_minimaps(input_image_bgr, optim_steps, target_avg
     # Vérifier si le modèle de segmentation existe (important car on ne peut pas l'afficher dans l'UI facilement)
     if not SEGMENTATION_MODEL_PATH.exists():
          # Retourner des images noires ou des messages d'erreur
-         error_msg = f"Erreur: Modèle {SEGMENTATION_MODEL_PATH} introuvable."
+         error_msg = f"Erreur: Modèle {SEGMENTATION_MODEL_PATH.name} introuvable."
          print(error_msg)
-         placeholder = np.zeros((300, 500, 3), dtype=np.uint8) # Placeholder noir
-         cv2.putText(placeholder, error_msg, (10, 150), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1, cv2.LINE_AA)
+         # Créer un placeholder plus informatif
+         placeholder = np.zeros((300, 500, 3), dtype=np.uint8)
+         cv2.putText(placeholder, "Model Error:", (10, 130), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 1, cv2.LINE_AA)
+         cv2.putText(placeholder, error_msg, (10, 155), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1, cv2.LINE_AA)
          return placeholder, placeholder.copy() # Retourner deux placeholders
 
+    # Vérifier aussi le modèle YOLO
+    if not YOLO_MODEL_PATH.exists():
+         error_msg = f"Erreur: Modèle {YOLO_MODEL_PATH.name} introuvable."
+         print(error_msg)
+         placeholder = np.zeros((300, 500, 3), dtype=np.uint8)
+         cv2.putText(placeholder, "Model Error:", (10, 130), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 1, cv2.LINE_AA)
+         cv2.putText(placeholder, error_msg, (10, 155), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1, cv2.LINE_AA)
+         return placeholder, placeholder.copy()
+
     try:
         # 1. Initialisation du modèle TvCalib (peut être lent si fait à chaque fois)
         #    Pourrait être optimisé en chargeant globalement (voir commentaire plus haut)
@@ -86,6 +105,31 @@ def process_image_and_generate_minimaps(input_image_bgr, optim_steps, target_avg
         image_tensor = image_tensor.to(model_device)
         print(f"Temps de prétraitement TvCalib : {time.time() - start_preprocess:.3f}s")
 
+        # --- Détection du ballon avec YOLO --- 
+        print("Chargement du modèle YOLO et détection du ballon...")
+        start_yolo = time.time()
+        ball_ref_point_img = None # Point de référence du ballon sur l'image originale redimensionnée
+        try:
+            # Charger le modèle YOLO (pourrait être chargé globalement pour la perf, mais attention)
+            yolo_model = YOLO(YOLO_MODEL_PATH) 
+            # Utiliser l'image BGR redimensionnée pour YOLO
+            results = yolo_model.predict(image_bgr_resized, classes=[BALL_CLASS_INDEX], verbose=False) 
+            
+            if results and len(results[0].boxes) > 0:
+                # Prendre la détection avec la plus haute confiance
+                best_ball_box = results[0].boxes[results[0].boxes.conf.argmax()]
+                x1, y1, x2, y2 = map(int, best_ball_box.xyxy[0].tolist())
+                conf = best_ball_box.conf[0].item()
+                
+                # Calculer le point de référence (centre bas de la bbox)
+                ball_ref_point_img = np.array([(x1 + x2) / 2, y2], dtype=np.float32)
+                print(f"  ✓ Ballon trouvé (conf: {conf:.2f}) à la bbox [{x1},{y1},{x2},{y2}]. Point réf: {ball_ref_point_img}")
+            else:
+                print("  Aucun ballon détecté.")
+                
+        except Exception as e_yolo:
+             print(f"  Erreur pendant la détection YOLO : {e_yolo}")
+        print(f"Temps de détection YOLO : {time.time() - start_yolo:.3f}s")
 
         # 3. Exécuter la calibration (Segmentation + Optimisation)
         print("Exécution de la segmentation...")
@@ -136,11 +180,12 @@ def process_image_and_generate_minimaps(input_image_bgr, optim_steps, target_avg
         # Minimap avec projection (image RGB attendue par la fonction)
         minimap_original = create_minimap_view(image_rgb_resized, homography_np)
 
-        # Minimap avec squelettes (utilise l'échelle estimée)
+        # Minimap avec squelettes ET LE BALLON (utilise l'échelle estimée)
         minimap_offset_skeletons, actual_avg_scale = create_minimap_with_offset_skeletons(
             player_list,
             homography_np,
-            base_skeleton_scale=estimated_base_scale
+            base_skeleton_scale=estimated_base_scale,
+            ball_ref_point_img=ball_ref_point_img # Passer le point de référence du ballon
         )
         print(f"Temps de génération des minimaps : {time.time() - start_viz:.3f}s")
         if actual_avg_scale is not None:
@@ -194,7 +239,7 @@ with gr.Blocks() as demo:
                 info="Number of iterations to refine homography."
             )
             target_scale_slider = gr.Slider(
-                minimum=0.1, maximum=2.5, step=0.05, value=0.35,
+                minimum=0.1, maximum=2.5, step=0.05, value=1,
                 label="Target Average Skeleton Scale",
                 info="Adjusts the desired average size of skeletons on the minimap."
             )

main.py

@@ -5,6 +5,8 @@ import torch
 from pathlib import Path
 import time
 import traceback
+# Import YOLO
+from ultralytics import YOLO
 
 # Assurez-vous que le répertoire tvcalib est dans le PYTHONPATH
 # ou exécutez depuis le répertoire tvcalib_image_processor
@@ -22,6 +24,10 @@ from pose_estimator import get_player_data
 # Constantes
 IMAGE_SHAPE = (720, 1280)  # Hauteur, Largeur
 SEGMENTATION_MODEL_PATH = Path("models/segmentation/train_59.pt")
+# Chemin vers le modèle YOLO pour la détection du ballon
+YOLO_MODEL_PATH = Path("models/detection/yolo_football.pt") 
+# Index de classe pour le ballon (basé sur votre exemple)
+BALL_CLASS_INDEX = 2 
 
 def preprocess_image_tvcalib(image_bgr):
     """Prétraite l'image BGR pour TvCalib et retourne le tenseur et l'image RGB redimensionnée."""
@@ -72,6 +78,12 @@ def main():
         print("Assurez-vous d'avoir copié train_59.pt dans le dossier models/segmentation/")
         return
 
+    # Vérifier l'existence du modèle YOLO
+    if not YOLO_MODEL_PATH.exists():
+        print(f"Erreur : Modèle YOLO introuvable : {YOLO_MODEL_PATH}")
+        print(f"Assurez-vous d'avoir téléchargé {YOLO_MODEL_PATH.name} et de l'avoir placé dans {YOLO_MODEL_PATH.parent}/")
+        return
+
     print("Initialisation de TvCalibInferModule...")
     try:
         model = TvCalibInferModule(
@@ -87,16 +99,49 @@ def main():
 
     print(f"Traitement de l'image : {args.image_path}")
     try:
+        # Vérification supplémentaire avant imread
+        image_path_obj = Path(args.image_path)
+        absolute_path = image_path_obj.resolve()
+        print(f"Tentative de lecture de l'image via cv2.imread depuis : {absolute_path}")
+        if not absolute_path.is_file():
+             print(f"ERREUR : Le chemin absolu {absolute_path} ne pointe pas vers un fichier existant juste avant imread !")
+             return # Arrêter ici si le fichier n'est pas trouvé à ce stade
+
         # Charger l'image (en BGR par défaut avec OpenCV)
         image_bgr_orig = cv2.imread(args.image_path)
         if image_bgr_orig is None:
-            raise FileNotFoundError(f"Impossible de lire le fichier image: {args.image_path}")
+            raise FileNotFoundError(f"Impossible de lire le fichier image: {args.image_path} (vérifié comme existant juste avant, problème avec imread)")
 
-        # Prétraiter l'image
+        # Prétraiter l'image pour TvCalib (redimensionne aussi)
         start_preprocess = time.time()
         image_tensor, image_bgr_resized, image_rgb_resized = preprocess_image_tvcalib(image_bgr_orig)
         print(f"Temps de prétraitement TvCalib : {time.time() - start_preprocess:.3f}s")
 
+        # --- Détection du ballon avec YOLO ---
+        print("\nChargement du modèle YOLO et détection du ballon...")
+        start_yolo = time.time()
+        ball_ref_point_img = None # Point de référence du ballon sur l'image originale redimensionnée
+        try:
+            yolo_model = YOLO(YOLO_MODEL_PATH)
+            # Utiliser l'image BGR redimensionnée pour YOLO
+            results = yolo_model.predict(image_bgr_resized, classes=[BALL_CLASS_INDEX], verbose=False) 
+            
+            if results and len(results[0].boxes) > 0:
+                # Prendre la détection avec la plus haute confiance
+                best_ball_box = results[0].boxes[results[0].boxes.conf.argmax()]
+                x1, y1, x2, y2 = map(int, best_ball_box.xyxy[0].tolist())
+                conf = best_ball_box.conf[0].item()
+                
+                # Calculer le point de référence (centre bas de la bbox)
+                ball_ref_point_img = np.array([(x1 + x2) / 2, y2], dtype=np.float32)
+                print(f"  ✓ Ballon trouvé (conf: {conf:.2f}) à la bbox [{x1},{y1},{x2},{y2}]. Point réf: {ball_ref_point_img}")
+            else:
+                print("  Aucun ballon détecté.")
+                
+        except Exception as e_yolo:
+             print(f"  Erreur pendant la détection YOLO : {e_yolo}")
+        print(f"Temps de détection YOLO : {time.time() - start_yolo:.3f}s")
+
         # Exécuter la segmentation
         print("Exécution de la segmentation...")
         start_segment = time.time()
@@ -155,12 +200,13 @@ def main():
             # 1. Minimap avec l'image originale (RGB)
             minimap_original = create_minimap_view(image_rgb_resized, homography_np)
             
-            # 2. Minimap avec les squelettes
-            # Utiliser l'échelle de base ESTIMÉE
+            # 2. Minimap avec les squelettes ET LE BALLON
+            # Utiliser l'échelle de base ESTIMÉE et passer les coordonnées du ballon
             minimap_offset_skeletons, actual_avg_scale = create_minimap_with_offset_skeletons(
                 player_list, 
                 homography_np, 
-                base_skeleton_scale=estimated_base_scale # Utiliser l'estimation
+                base_skeleton_scale=estimated_base_scale, # Utiliser l'estimation
+                ball_ref_point_img=ball_ref_point_img     # Passer le point de référence du ballon
             )
 
             # Afficher la cible et le résultat réel

requirements.txt

@@ -2,21 +2,34 @@
 # pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu121
 
 # Dépendances principales
-torch
+torch>=1.8.0
 torchvision
 torchaudio
 numpy
 opencv-python
-pytorch-lightning
+pytorch-lightning>=1.4.0
 soccernet
-kornia
+kornia>=0.6.0
 
 # Ajouté car requis par sn_segmentation 
 
 # Dépendances pour l'estimation de pose (ViTPose)
-transformers
+transformers>=4.0.0
 supervision
-Pillow # Souvent une dépendance de transformers/supervision, mais explicite ici 
+Pillow
 accelerate
 # scikit-learn # Retiré car K-Means n'est plus utilisé
-gradio
+gradio
+
+# Ball detection
+ultralytics
+
+# Optionnel mais recommandé pour la gestion des chemins
+pathlib
+
+# Dépendances possibles de TvCalib (à vérifier)
+# tensorboard
+# matplotlib
+# scipy
+# tqdm
+# scikit-image

visualizer.py

@@ -16,6 +16,11 @@ from pose_estimator import (LEFT_ANKLE_KP_INDEX, RIGHT_ANKLE_KP_INDEX,
 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)
@@ -163,14 +168,16 @@ def create_minimap_view(image_rgb, homography, minimap_size=(EXPECTED_W, EXPECTE
 
 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.
+       à 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:
@@ -280,12 +287,32 @@ def create_minimap_with_offset_skeletons(player_data_list, homography,
                     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)
 
-    # Calculer l'échelle moyenne finale
-    average_draw_scale = base_skeleton_scale # Valeur par défaut si aucun joueur n'est dessiné
-    if drawn_players_count > 0:
-        average_draw_scale = total_applied_scale / drawn_players_count
+    # --- É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, average_draw_scale # Retourner aussi l'échelle moyenne
+    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)