import argparse import cv2 import numpy as np 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 from tvcalib.infer.module import TvCalibInferModule # Importer les fonctions de visualisation et les constantes de modulation from visualizer import ( create_minimap_view, create_minimap_with_offset_skeletons, DYNAMIC_SCALE_MIN_MODULATION, # Importer les constantes DYNAMIC_SCALE_MAX_MODULATION ) # Importer la fonction d'extraction des données joueurs 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.""" if image_bgr is None: raise ValueError("Impossible de charger l'image") # 1. Redimensionner en 720p si nécessaire h, w = image_bgr.shape[:2] if h != IMAGE_SHAPE[0] or w != IMAGE_SHAPE[1]: print(f"Redimensionnement de l'image vers {IMAGE_SHAPE[1]}x{IMAGE_SHAPE[0]}") image_bgr_resized = cv2.resize(image_bgr, (IMAGE_SHAPE[1], IMAGE_SHAPE[0]), interpolation=cv2.INTER_LINEAR) else: image_bgr_resized = image_bgr # 2. Convertir en RGB (pour TvCalib ET pour la visualisation originale) image_rgb_resized = cv2.cvtColor(image_bgr_resized, cv2.COLOR_BGR2RGB) # 3. Normalisation spécifique pour le modèle pré-entraîné (pour TvCalib) image_tensor = torch.from_numpy(image_rgb_resized).float() image_tensor = image_tensor.permute(2, 0, 1) # HWC -> CHW mean = torch.tensor([0.485, 0.456, 0.406]).view(-1, 1, 1) std = torch.tensor([0.229, 0.224, 0.225]).view(-1, 1, 1) image_tensor = (image_tensor / 255.0 - mean) / std # 4. Déplacer sur le bon device device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') image_tensor = image_tensor.to(device) # Retourner le tenseur pour TvCalib, l'image BGR et RGB redimensionnée return image_tensor, image_bgr_resized, image_rgb_resized def main(): parser = argparse.ArgumentParser(description="Exécute la méthode TvCalib sur une seule image.") parser.add_argument("image_path", type=str, help="Chemin vers l'image à traiter.") parser.add_argument("--output_homography", type=str, default=None, help="Chemin optionnel pour sauvegarder la matrice d'homographie (.npy).") parser.add_argument("--optim_steps", type=int, default=500, help="Nombre d'étapes d'optimisation pour la calibration (l'arrêt anticipé est désactivé).") parser.add_argument("--target_avg_scale", type=float, default=1, help="Facteur d'échelle MOYEN CIBLE pour dessiner les squelettes sur la minimap (défaut: 0.35). Le script ajuste l'échelle de base pour tenter d'atteindre cette moyenne.") args = parser.parse_args() if not Path(args.image_path).exists(): print(f"Erreur : Fichier image introuvable : {args.image_path}") return if not SEGMENTATION_MODEL_PATH.exists(): print(f"Erreur : Modèle de segmentation introuvable : {SEGMENTATION_MODEL_PATH}") 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( segmentation_checkpoint=SEGMENTATION_MODEL_PATH, image_shape=IMAGE_SHAPE, optim_steps=args.optim_steps, lens_dist=False # Gardons cela simple pour l'instant ) print(f"✓ Modèle chargé sur {next(model.model_calib.parameters()).device}") except Exception as e: print(f"Erreur lors de l'initialisation du modèle : {e}") return 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} (vérifié comme existant juste avant, problème avec imread)") # 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() with torch.no_grad(): keypoints = model._segment(image_tensor) print(f"Temps de segmentation : {time.time() - start_segment:.3f}s") # Exécuter la calibration (optimisation) print("Exécution de la calibration (optimisation)...") start_calibrate = time.time() homography = model._calibrate(keypoints) print(f"Temps de calibration : {time.time() - start_calibrate:.3f}s") if homography is not None: print("\n--- Homographie Calculée ---") if isinstance(homography, torch.Tensor): homography_np = homography.detach().cpu().numpy() else: homography_np = homography print(homography_np) if args.output_homography: try: np.save(args.output_homography, homography_np) print(f"\nHomographie sauvegardée dans : {args.output_homography}") except Exception as e: print(f"Erreur lors de la sauvegarde de l'homographie : {e}") # --- Extraction des données joueurs --- print("\nExtraction des données joueurs (pose+couleur)...") start_pose = time.time() player_list = get_player_data(image_bgr_resized) print(f"Temps d'extraction données joueurs : {time.time() - start_pose:.3f}s ({len(player_list)} joueurs trouvés)") # --- Calcul de l'échelle de base estimée --- print("\nCalcul de l'échelle de base pour atteindre la cible...") target_average_scale = args.target_avg_scale # Calculer la modulation moyenne attendue (hypothèse: joueur moyen au centre Y=0.5) # Logique inversée actuelle : MIN + (MAX - MIN) * (1.0 - norm_y) avg_modulation_expected = DYNAMIC_SCALE_MIN_MODULATION + \ (DYNAMIC_SCALE_MAX_MODULATION - DYNAMIC_SCALE_MIN_MODULATION) * (1.0 - 0.5) estimated_base_scale = target_average_scale # Valeur par défaut si modulation = 0 if avg_modulation_expected != 0: estimated_base_scale = target_average_scale / avg_modulation_expected else: print("Avertissement : Modulation moyenne attendue nulle, impossible d'estimer l'échelle de base.") print(f" Modulation dynamique moyenne attendue (pour Y=0.5) : {avg_modulation_expected:.3f}") print(f" Échelle de base interne estimée pour cible {target_average_scale:.3f} : {estimated_base_scale:.3f}") # --- Génération des DEUX minimaps --- print("\nGénération des minimaps (Originale et Squelettes Décalés)...") # 1. Minimap avec l'image originale (RGB) minimap_original = create_minimap_view(image_rgb_resized, homography_np) # 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 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 if actual_avg_scale is not None: print(f"\nÉchelle moyenne CIBLE demandée (--target_avg_scale) : {target_average_scale:.3f}") print(f"Échelle moyenne FINALE RÉELLEMENT appliquée (basée sur joueurs réels) : {actual_avg_scale:.3f}") # --- Affichage des résultats --- print("\nAffichage des résultats. Appuyez sur une touche pour quitter.") # Afficher la minimap originale if minimap_original is not None: cv2.imshow("Minimap avec Projection Originale", minimap_original) else: print("N'a pas pu générer la minimap originale.") # Afficher la minimap avec les squelettes décalés if minimap_offset_skeletons is not None: cv2.imshow("Minimap avec Squelettes Decales", minimap_offset_skeletons) else: print("N'a pas pu générer la minimap squelettes décalés.") cv2.waitKey(0) # Attend qu'une touche soit pressée else: print("\nAucune homographie n'a pu être calculée.") except Exception as e: print(f"Erreur lors du traitement de l'image : {e}") traceback.print_exc() finally: print("Fermeture des fenêtres OpenCV.") cv2.destroyAllWindows() if __name__ == "__main__": main()