init

.gitattributes

@@ -0,0 +1,2 @@
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,12 @@
+venv
+*.pyc
+*.jpg
+*.png
+*.jpeg
+*.gif
+*.bmp
+*.tiff
+*.ico
+
+
+*.pt

README.md

@@ -0,0 +1,107 @@
+# Foot Calib Pos Image Processor
+
+This project uses the TvCalib library to calculate the homography matrix for a football pitch image. This matrix allows mapping image points onto a standard 2D representation of the pitch (minimap).
+
+The project also includes a pose estimation step (ViTPose) to detect players and calculate the average color of their torso.
+
+**The main result is a minimap where each player is represented by their colored skeleton, drawn at a *dynamically* reduced scale around their projected position on the pitch.**
+The position is determined by projecting a reference point (feet/bottom of bbox) using the homography. The skeleton is then drawn using its relative coordinates (original image), scaled, and translated. The scale used depends on the **player's vertical position (Y) on the minimap** (higher on the minimap = smaller) and a base factor adjustable via the `--target_avg_scale` option.
+
+It is also possible to visualize a minimap with the projection of the original image for comparison.
+
+## Features
+
+*   Homography calculation from a single image via TvCalib.
+*   Person detection (RT-DETR) and pose estimation (ViTPose).
+*   Calculation of the filtered average torso color for each player.
+*   Projection of each player's reference point (feet/bbox) onto the minimap.
+*   Generation of a minimap with the **original skeletons (colored, dynamically scaled based on projected Y position, and offset) drawn around the projected point**.
+*   (Optional) Generation of a minimap with the projected original image.
+*   Possibility to save the calculated homography matrix.
+
+## Project Structure
+
+```
+.
+├── .git/               # Git metadata
+├── .venv/              # Python virtual environment (recommended)
+├── common/             # Common Python modules (potentially)
+├── data/               # Data (input images, etc.)
+├── models/
+│   └── segmentation/
+│       └── train_59.pt # Pre-trained segmentation model (TO DOWNLOAD)
+├── tvcalib/            # Source code of the TvCalib library (or a fork/adaptation)
+│   └── infer/
+│       └── module.py   # Main module for TvCalib inference
+├── .gitignore          # Files ignored by Git
+├── main.py             # Main script entry point
+├── requirements.txt    # Python dependencies file
+├── visualizer.py       # Module for generating visualization minimaps
+├── pose_estimator.py   # Module for pose estimation and player data extraction
+└── README.md           # This file
+```
+
+## Installation
+
+1.  **Clone the repository:**
+    ```powershell
+    git clone <repository-url>
+    cd Foot_calib_pos_image_processor
+    ```
+
+2.  **Create a virtual environment (recommended):**
+    ```powershell
+    python -m venv venv
+    .\venv\Scripts\Activate.ps1
+    ```
+
+3.  **Install dependencies:**
+    ```powershell
+    pip install -r requirements.txt
+    ```
+    *(Make sure to install PyTorch with appropriate CUDA support if needed.)*
+
+4.  **Download the segmentation model:**
+    Place `train_59.pt` in `models/segmentation/`.
+
+5.  **(Automatic) Download detection/pose models:**
+    The RT-DETR and ViTPose models will be downloaded automatically.
+
+## Usage
+
+Run the `main.py` script providing the path to the image:
+
+```powershell
+python main.py path/to/your/image.jpg [OPTIONS]
+```
+
+**Options:**
+
+*   `image_path`: Path to the input image (required).
+*   `--output_homography PATH.npy`: Saves the calculated homography matrix.
+*   `--optim_steps NUMBER`: Number of optimization steps for calibration (default: 500, was 1000 in original README example).
+*   `--target_avg_scale FLOAT`: **Target average** scale factor for drawing skeletons (default: 0.35). The script attempts to adjust the internal base scale so that the resulting average scale (after inverse dynamic modulation) is close to this value.
+
+**Example:**
+
+```powershell
+# Simple usage (target average size 0.35)
+python main.py data/img3.png
+
+# Aim for larger skeletons on average (target 0.5)
+python main.py data/img2.png --target_avg_scale 0.5
+```
+
+The script will display:
+*   Time taken and homography matrix.
+*   Estimated internal base scale.
+*   Requested TARGET average scale.
+*   ACTUALLY applied FINAL average scale.
+*   Window: **Minimap with Original Projection**.
+*   Window: **Minimap with Offset Skeletons** (dynamically scaled inversely, targeting the average scale).
+*   Press any key to close.
+
+## Key Dependencies
+
+*   PyTorch, OpenCV, NumPy, PyTorch Lightning
+*   SoccerNet, Kornia, Hugging Face Transformers, Pillow 

app.py

@@ -0,0 +1,234 @@
+import gradio as gr
+import cv2
+import numpy as np
+import torch
+from pathlib import Path
+import time
+import traceback
+
+# 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
+    from visualizer import (
+        create_minimap_view,
+        create_minimap_with_offset_skeletons,
+        DYNAMIC_SCALE_MIN_MODULATION,
+        DYNAMIC_SCALE_MAX_MODULATION
+    )
+    from pose_estimator import get_player_data
+except ImportError as e:
+    print(f"Erreur d'importation : {e}")
+    print("Assurez-vous que les modules tvcalib, main, visualizer, pose_estimator sont accessibles.")
+    # On pourrait mettre des stubs ou lever une exception ici pour Gradio
+    raise e
+
+# --- Configuration Globale (Modèle, etc.) ---
+# Essayer de charger le modèle une seule fois globalement peut améliorer les performances
+# mais attention à la gestion de l'état dans les environnements multi-utilisateurs/threads de Spaces
+# Pour l'instant, on le chargera dans la fonction de traitement.
+# MODEL = None # Optionnel: Charger ici
+# DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+print(f"Utilisation du device : {DEVICE}")
+
+if not SEGMENTATION_MODEL_PATH.exists():
+    print(f"AVERTISSEMENT : Modèle de segmentation introuvable : {SEGMENTATION_MODEL_PATH}")
+    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.
+
+# --- Fonction Principale de Traitement ---
+def process_image_and_generate_minimaps(input_image_bgr, optim_steps, target_avg_scale):
+    """
+    Prend une image BGR (NumPy), les étapes d'optimisation et l'échelle cible,
+    retourne les deux minimaps (NumPy BGR).
+    """
+    global DEVICE # Utiliser le device défini globalement
+
+    print("\n--- Nouvelle requête ---")
+    print(f"Paramètres: optim_steps={optim_steps}, target_avg_scale={target_avg_scale}")
+
+    # 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."
+         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)
+         return placeholder, placeholder.copy() # Retourner deux placeholders
+
+    try:
+        # 1. Initialisation du modèle TvCalib (peut être lent si fait à chaque fois)
+        #    Pourrait être optimisé en chargeant globalement (voir commentaire plus haut)
+        print("Initialisation de TvCalibInferModule...")
+        start_init = time.time()
+        model = TvCalibInferModule(
+            segmentation_checkpoint=SEGMENTATION_MODEL_PATH,
+            image_shape=IMAGE_SHAPE, # Utilise la constante importée
+            optim_steps=int(optim_steps), # Assurer que c'est un entier
+            lens_dist=False
+        )
+        # Déplacer le modèle sur le bon device ici explicitement si nécessaire
+        # model.to(DEVICE) # TvCalibInferModule devrait gérer ça en interne ? A vérifier.
+        print(f"✓ Modèle chargé sur {next(model.model_calib.parameters()).device} en {time.time() - start_init:.3f}s")
+        model_device = next(model.model_calib.parameters()).device # Vérifier le device réel
+
+        # 2. Prétraitement de l'image
+        print("Prétraitement de l'image...")
+        start_preprocess = time.time()
+        # preprocess_image_tvcalib attend BGR, Gradio fournit BGR par défaut avec type="numpy"
+        # Assurez-vous que preprocess_image_tvcalib déplace bien le tenseur sur le bon device
+        image_tensor, image_bgr_resized, image_rgb_resized = preprocess_image_tvcalib(input_image_bgr)
+        # Vérifier/forcer le device du tenseur
+        image_tensor = image_tensor.to(model_device)
+        print(f"Temps de prétraitement TvCalib : {time.time() - start_preprocess:.3f}s")
+
+
+        # 3. Exécuter la calibration (Segmentation + Optimisation)
+        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")
+
+        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 None:
+            print("Aucune homographie n'a pu être calculée.")
+            # Retourner des placeholders avec message
+            placeholder = np.zeros((300, 500, 3), dtype=np.uint8)
+            cv2.putText(placeholder, "Homographie non calculee", (10, 150), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1, cv2.LINE_AA)
+            return placeholder, placeholder.copy()
+
+        if isinstance(homography, torch.Tensor):
+            homography_np = homography.detach().cpu().numpy()
+        else:
+            homography_np = np.array(homography) # Assurer que c'est un NumPy array
+        print("✓ Homographie calculée.")
+
+
+        # 4. Extraction des données joueurs
+        print("Extraction des données joueurs (pose+couleur)...")
+        start_pose = time.time()
+        # get_player_data attend une image BGR
+        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)")
+
+        # 5. Calcul de l'échelle de base
+        print("Calcul de l'échelle de base...")
+        # Reprend la logique de main.py pour estimer l'échelle de base
+        avg_modulation_expected = DYNAMIC_SCALE_MIN_MODULATION + \
+                                  (DYNAMIC_SCALE_MAX_MODULATION - DYNAMIC_SCALE_MIN_MODULATION) * (1.0 - 0.5)
+        estimated_base_scale = target_avg_scale
+        if avg_modulation_expected != 0:
+             estimated_base_scale = target_avg_scale / avg_modulation_expected
+        print(f"  Échelle de base interne estimée pour cible {target_avg_scale:.3f} : {estimated_base_scale:.3f}")
+
+        # 6. Génération des minimaps
+        print("Génération des minimaps...")
+        start_viz = time.time()
+        # 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_offset_skeletons, actual_avg_scale = create_minimap_with_offset_skeletons(
+            player_list,
+            homography_np,
+            base_skeleton_scale=estimated_base_scale
+        )
+        print(f"Temps de génération des minimaps : {time.time() - start_viz:.3f}s")
+        if actual_avg_scale is not None:
+             print(f"Échelle moyenne CIBLE demandée : {target_avg_scale:.3f}")
+             print(f"Échelle moyenne FINALE RÉELLEMENT appliquée : {actual_avg_scale:.3f}")
+
+
+        # Vérifier si les minimaps ont été créées (peuvent être None en cas d'erreur interne)
+        if minimap_original is None:
+            print("Erreur: La minimap originale n'a pas pu être générée.")
+            minimap_original = np.zeros((300, 500, 3), dtype=np.uint8)
+            cv2.putText(minimap_original, "Erreur Minimap Originale", (10, 150), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1, cv2.LINE_AA)
+
+        if minimap_offset_skeletons is None:
+            print("Erreur: La minimap squelettes n'a pas pu être générée.")
+            minimap_offset_skeletons = np.zeros((300, 500, 3), dtype=np.uint8)
+            cv2.putText(minimap_offset_skeletons, "Erreur Minimap Squelettes", (10, 150), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1, cv2.LINE_AA)
+
+        # Gradio attend des images RGB pour l'affichage, nos fonctions retournent probablement BGR (via OpenCV)
+        # Conversion BGR -> RGB si nécessaire
+        if minimap_original.shape[2] == 3: # Assurer que c'est une image couleur
+             minimap_original = cv2.cvtColor(minimap_original, cv2.COLOR_BGR2RGB)
+        if minimap_offset_skeletons.shape[2] == 3:
+             minimap_offset_skeletons = cv2.cvtColor(minimap_offset_skeletons, cv2.COLOR_BGR2RGB)
+
+        print("✓ Traitement terminé.")
+        return minimap_original, minimap_offset_skeletons
+
+    except Exception as e:
+        print(f"Erreur majeure lors du traitement : {e}")
+        traceback.print_exc()
+        # Retourner des placeholders avec message d'erreur général
+        placeholder = np.zeros((300, 500, 3), dtype=np.uint8)
+        cv2.putText(placeholder, f"Erreur: {e}", (10, 150), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 255), 1, cv2.LINE_AA)
+        return placeholder, placeholder.copy()
+
+# --- Interface Gradio ---
+with gr.Blocks() as demo:
+    gr.Markdown("# Foot Calib Pos Image Processor - Minimap Generator")
+    gr.Markdown(
+        "Upload a football pitch image to compute homography (TvCalib), "
+        "detect players (RT-DETR/ViTPose), and generate two minimap visualizations."
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_image = gr.Image(type="numpy", label="Input Image (.jpg, .png)")
+            optim_steps_slider = gr.Slider(
+                minimum=100, maximum=2000, step=50, value=500,
+                label="TvCalib Optimization Steps",
+                info="Number of iterations to refine homography."
+            )
+            target_scale_slider = gr.Slider(
+                minimum=0.1, maximum=2.5, step=0.05, value=0.35,
+                label="Target Average Skeleton Scale",
+                info="Adjusts the desired average size of skeletons on the minimap."
+            )
+            submit_button = gr.Button("Generate Minimaps", variant="primary")
+
+        with gr.Column(scale=2):
+            output_minimap_orig = gr.Image(type="numpy", label="Minimap with Original Projection", interactive=False)
+            output_minimap_skel = gr.Image(type="numpy", label="Minimap with Offset Skeletons", interactive=False)
+
+    # Connecter le bouton à la fonction de traitement
+    submit_button.click(
+        fn=process_image_and_generate_minimaps,
+        inputs=[input_image, optim_steps_slider, target_scale_slider],
+        outputs=[output_minimap_orig, output_minimap_skel]
+    )
+
+    # Ajouter des exemples (optionnel mais utile pour Spaces)
+    gr.Examples(
+        examples=[
+            ["data/img1.png", 500, 1.35],
+            ["data/img2.png", 1000, 1.5],
+            ["data/img3.png", 500, 0.8],
+             ["data/7.jpg", 500, 1], # Add .jpg examples
+             ["data/15.jpg", 800, 1.35],
+        ],
+        inputs=[input_image, optim_steps_slider, target_scale_slider],
+        outputs=[output_minimap_orig, output_minimap_skel], # Outputs won't be pre-calculated here, just to populate inputs
+        fn=process_image_and_generate_minimaps, # Function will be called if example is clicked
+        cache_examples=False # Important if processing is long or depends on external models
+    )
+
+
+# --- Lancement de l'application ---
+if __name__ == "__main__":
+    # share=True creates a temporary public link (useful for testing outside localhost)
+    # debug=True shows more Gradio logs in the console
+    demo.launch(debug=True) 

common/data/augmentation.py

@@ -0,0 +1,55 @@
+
+import numpy as np
+from typing import Callable, Any
+import methods.common.data.utils as utils
+import random
+
+
+
+class WarpAugmentation(Callable):
+    def __init__(
+        self,
+        warp_function: Any,
+        mode: str="train",
+        noise_translate: float=0.0,
+        noise_rotate: float=0.0
+    ):
+        # Template
+        self.template_grid = utils.gen_template_grid()
+        self.warp_fn = warp_function
+        self.mode = mode
+        self.noise_translate = noise_translate
+        self.noise_rotate = noise_rotate
+
+    def __call__(
+        self,
+        image: np.ndarray,
+        homography: np.ndarray,
+        frame_idx: int      
+    ):
+        warp_image, warp_grid, warp_homography = self.warp_fn(
+            mode=self.mode,
+            frame=image,
+            f_idx=frame_idx,
+            gt_homo=homography,
+            template=self.template_grid,
+            noise_trans=self.noise_translate,
+            noise_rotate=self.noise_rotate,
+            index=-1        # not really used ...
+        )
+        return warp_image, warp_grid, warp_homography
+
+
+
+
+class LeftRightFlipAugmentation(Callable):
+    def __init__(self, enabled: bool=False):
+        self.enabled = enabled
+        
+    def __call__(self, image, grid):
+        if (self.enabled):            
+            if (random.random() < 0.5):
+                image, grid = utils.put_lrflip_augmentation(image, grid)
+                return image, grid, True
+            
+        return image, grid, False

common/data/calib.py

@@ -0,0 +1,116 @@
+
+import cv2
+import numpy as np
+import torch
+import methods.common.data.utils as utils
+
+
+
+class Calibrator:
+    def __init__(
+        self,
+        num_classes: int=92,
+        nms_threshold: float=0.995
+    ):
+        self.template_grid = utils.gen_template_grid()
+        self.num_classes = num_classes
+        self.nms_threshold = nms_threshold
+
+
+    def find_homography(
+        self, 
+        heatmap_logits: torch.Tensor
+    ):
+        """ Extract keypoints from heatmap, and find homography matrix 
+        
+            heatmap_logits: torch.Tensor for individual frame (not a mini-batch!)
+        """
+        
+        pred_rgb, pred_keypoints, scores_heatmap = self.decode_heatmap(heatmap_logits)
+        homography = None
+        
+        # We need at least 4 point correspondences
+        if (pred_rgb.shape[0] >= 4):
+            src_pts, dst_pts = self.get_class_mapping(pred_rgb)
+            
+            # Find homography from point correspondences
+            homography, _ = cv2.findHomography(
+                src_pts.reshape(-1, 1, 2), 
+                dst_pts.reshape(-1, 1, 2), 
+                cv2.RANSAC, 
+                10
+            )
+            
+        return homography, pred_keypoints, scores_heatmap
+
+
+    def decode_heatmap(self, heatmap_logits: torch.Tensor):
+        """ Decode heatmap info keypoint set using non-maximum suppression        
+            heatmap_logits: torc.Tensor with shape <NUM_CLASSES; H; W>
+        """
+        
+        pred_heatmap = torch.softmax(heatmap_logits, dim=0)
+        arg = torch.argmax(pred_heatmap, dim=0).detach().cpu().numpy()
+        scores, pred_heatmap = torch.max(pred_heatmap, dim=0)
+                
+        # Convert to Numpy & get keypoints locations
+        scores = scores.detach().cpu().numpy()
+        pred_heatmap = pred_heatmap.detach().cpu().numpy()
+        pred_class_dict = self.get_class_dict(scores, pred_heatmap)
+        
+        # Colorize
+        num_classes = heatmap_logits.shape[0]
+        np_scores = np.clip(arg * 255.0 / num_classes, 0, 255).astype(np.uint8)
+        scores_heatmap = cv2.applyColorMap(np_scores, cv2.COLORMAP_HOT)
+        scores_heatmap = cv2.cvtColor(scores_heatmap, cv2.COLOR_BGR2RGB)
+        
+        # Produce image with keypoints
+        pred_keypoints = np.zeros_like(pred_heatmap, dtype=np.uint8)
+        pred_rgb = []
+        for _, (pk, pv) in enumerate(pred_class_dict.items()):
+            if (pv):
+                pred_keypoints[pv[1][0], pv[1][1]] = pk     # (H,W)
+                # camera view point sets (x, y, label) in rgb domain not heatmap domain
+                pred_rgb.append([pv[1][1] * 4, pv[1][0] * 4, pk])
+        pred_rgb = np.asarray(pred_rgb, dtype=np.float32)  # (?, 3)
+        
+        # Return list of point locations, and image of keypoints
+        return pred_rgb, pred_keypoints, scores_heatmap
+
+
+
+    def get_class_mapping(self, rgb):
+        src_pts = rgb.copy()
+        cls_map_pts = []
+
+        for ind, elem in enumerate(src_pts):
+            coords = np.where(elem[2] == self.template_grid[:, 2])[0]  # find correspondence
+            cls_map_pts.append(self.template_grid[coords[0]])
+        dst_pts = np.array(cls_map_pts, dtype=np.float32)
+
+        return src_pts[:, :2], dst_pts[:, :2]
+
+
+    def get_class_dict(self, scores, pred):
+        # Decode                               
+        pred_cls_dict = {k: [] for k in range(1, self.num_classes)}
+        for cls in range(1, self.num_classes):
+            pred_inds = (pred == cls)
+
+            # implies the current class does not appear in this heatmaps
+            if not np.any(pred_inds):
+                continue
+
+            values = scores[pred_inds]
+            max_score = values.max()
+            max_index = values.argmax()
+
+            indices = np.where(pred_inds)
+            coords = list(zip(indices[0], indices[1]))
+
+            # the only keypoint with max confidence is greater than threshold or not
+            if max_score >= self.nms_threshold:
+                pred_cls_dict[cls].append(max_score)
+                pred_cls_dict[cls].append(coords[max_index])
+
+        return pred_cls_dict

common/data/transforms.py

@@ -0,0 +1,23 @@
+from typing import Callable
+import torch
+import numpy as np
+
+
+
+class UnNormalize(object):
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+
+    def __call__(self, tensor):
+        """
+        Args:
+            tensor (Tensor): Tensor image of size (C, H, W) to be normalized.
+        Returns:
+            Tensor: Normalized image.
+        """
+        for t, m, s in zip(tensor, self.mean, self.std):
+            t.mul_(s).add_(m)
+            # The normalize code -> t.sub_(m).div_(s)
+        return tensor
+    

common/data/utils.py

@@ -0,0 +1,72 @@
+import cv2
+import numpy as np
+import random
+import torch
+
+
+
+def yards(x):
+    return x * 1.0936132983
+
+
+
+def to_torch(ndarray):
+    if type(ndarray).__module__ == 'numpy':
+        return torch.from_numpy(ndarray.copy())
+    elif not torch.is_tensor(ndarray):
+        raise ValueError("Cannot convert {} to torch tensor"
+                         .format(type(ndarray)))
+    return ndarray
+
+
+
+def gen_template_grid():
+    # === set uniform grid ===
+    # field_dim_x, field_dim_y = 105.000552, 68.003928 # in meter
+    field_dim_x, field_dim_y = 114.83, 74.37  # in yard
+    # field_dim_x, field_dim_y = 115, 74 # in yard
+    nx, ny = (13, 7)
+    x = np.linspace(0, field_dim_x, nx)
+    y = np.linspace(0, field_dim_y, ny)
+    xv, yv = np.meshgrid(x, y, indexing='ij')
+    uniform_grid = np.stack((xv, yv), axis=2).reshape(-1, 2)
+    uniform_grid = np.concatenate((uniform_grid, np.ones(
+        (uniform_grid.shape[0], 1))), axis=1)  # top2bottom, left2right
+    # TODO: class label in template, each keypoints is (x, y, c), c is label that starts from 1
+    for idx, pts in enumerate(uniform_grid):
+        pts[2] = idx + 1  # keypoints label
+    return uniform_grid
+
+
+def put_lrflip_augmentation(frame, unigrid):
+
+    frame_h, frame_w = frame.shape[0], frame.shape[1]
+    flipped_img = np.fliplr(frame)
+
+    # TODO: grid flipping and re-assign pixels class label, 1-91
+    for ind, pts in enumerate(unigrid):
+        pts[0] = frame_w - pts[0]
+        col = (pts[2] - 1) // 7  # get each column of uniform grid
+        pts[2] = pts[2] - (col - 6) * 2 * 7  # keypoints label
+        
+    return flipped_img, unigrid
+
+
+
+
+def make_grid(images, nrow=4):
+    num_images = len(images)
+    ih, iw = images[0].shape[0], images[0].shape[1]
+    rows = min(num_images, nrow)
+    cols = (num_images + nrow-1) // nrow
+    
+    result = np.zeros(shape=(cols*ih, rows*iw, 3), dtype=np.uint8)
+    for i in range(num_images):
+        cell_x = i%nrow
+        cell_y = i//nrow
+        result[
+            (cell_y+0)*ih:(cell_y+1)*ih,
+            (cell_x+0)*iw:(cell_x+1)*iw
+        ] = images[i]
+        
+    return result

common/infer/base.py

@@ -0,0 +1,43 @@
+
+from typing import Any, Dict
+from torch.utils.data import Dataset
+
+import numpy as np
+
+
+
+
+class InferDataModule:
+    def __init__(self):
+        pass
+
+
+    def get_inference_dataset(self) -> Dataset:
+        """ Return the dataset to run inference on """
+        pass
+    
+    
+
+class InferModule:
+    def __init__(self):
+        pass
+    
+    
+    def setup(
+        self,
+        datamodule: InferDataModule
+    ):
+        """ Initialize inference proces with the given datamodule """
+        pass
+        
+   
+    def predict(
+        self, 
+        x: Any
+    ) -> Dict:
+        """ Predict the calibration information for the given dataset sample """
+        return None
+    
+    
+    
+    

common/infer/module.py

@@ -0,0 +1,24 @@
+from typing import Any, Dict
+from methods.common.infer.base import *
+
+
+
+class LabelInferModule(InferModule):
+    def __init__(self):
+        pass        
+    
+    def setup(self, datamodule: InferDataModule):
+        pass
+    
+    
+    def predict(self, x: Any) -> Dict:
+        """
+            x - sample from dataset (including label)
+        """
+                    
+        # Extract homography matrix
+        result = {
+            "homography": x["homography"]
+        }
+        
+        return result

common/infer/sink.py

@@ -0,0 +1,42 @@
+
+from pathlib import Path
+from collections import defaultdict
+import pandas as pd
+from typing import Dict
+import cv2
+
+
+class PredictionsSink:
+    def __init__(
+        self,
+        target_filepath: Path
+    ):        
+        self.target_filepath = target_filepath        
+        def new_item():
+            return []
+        self.data = defaultdict(new_item) 
+        self.count = 0      
+
+
+    def write(self, item: Dict):
+        self.data["item"].append(self.count)        
+        for k,v in item.items():            
+            if ("image" in k):
+                self.write_image(self.count, k, v)
+            else:
+                self.data[k].append(v)
+                
+        self.count += 1
+            
+            
+    def flush(self):
+        df = pd.DataFrame(self.data)
+        self.target_filepath.parent.mkdir(parents=True, exist_ok=True)
+        df.to_csv(self.target_filepath)        
+        
+    def write_image(self, idx, name, image):
+        folder = self.target_filepath.parent / "images" / self.target_filepath.stem / name
+        filepath = folder / f"{idx:06d}.png"        
+        filepath.parent.mkdir(parents=True, exist_ok=True)        
+        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+        cv2.imwrite(filepath.as_posix(), image)

common/loggers/homography_previewer.py

@@ -0,0 +1,103 @@
+
+import numpy as np
+import skimage.segmentation as ss
+import cv2
+
+from torchvision import transforms
+
+from project.data.transforms import TensorToNumpy
+
+
+from methods.common.loggers.image_preview import ImagePreviewLogger
+from methods.common.data.transforms import UnNormalize
+from methods.common.data.calib import Calibrator
+
+
+class HomographyPreviewerLogger(ImagePreviewLogger):
+    def __init__(
+        self,
+        experiment,
+        num_rows: int=3,
+    ):
+        super().__init__(experiment, num_rows)
+
+        self.calib = Calibrator()        
+        self.to_image = transforms.Compose([
+            UnNormalize(
+                mean=[0.485, 0.456, 0.406], 
+                std=[0.229, 0.224, 0.225]
+            ),
+            TensorToNumpy()
+        ])
+
+
+
+
+    def draw_keypoints(self, image, image_keypoints, color=(255,0,0)):
+        """ Upscales keypoints map into image resolution and 
+            overlays it over the image
+        """
+
+        # Get keypoints image in target image resolution
+        a = ss.expand_labels(image_keypoints, distance=1)
+        a = cv2.resize(
+            a, 
+            dsize=(image.shape[1], image.shape[0]),
+            interpolation=cv2.INTER_NEAREST
+        )
+        a = np.expand_dims(a, axis=2)
+
+        # Alpha of keypoints image
+        a = (a > 0)*1.0
+
+        # Color of keypoints image
+        c = np.concatenate([a*color[0], a*color[1], a*color[2]], axis=2)
+
+        # Superimpose the keypoints
+        result = (1.0-a)*image + a*c
+        result = np.clip(result, 0, 255).astype(np.uint8)
+        return result
+
+
+    def draw_playfield(
+        self, 
+        image, 
+        image_playfield, 
+        homography, 
+        color=(255,0,0),
+        alpha=1.0,
+        flip=False
+    ):
+        """ Draws the playfield image under the homography matrix
+            over the target image
+        """
+        if (homography is None):
+            return image
+        
+        # Warp the playfield image
+        warp_field = cv2.warpPerspective(
+            image_playfield,
+            homography,
+            (image.shape[1], image.shape[0]),
+            cv2.INTER_LINEAR,
+            borderMode=cv2.BORDER_CONSTANT,
+            borderValue=(0)
+        )
+        
+        if (flip):
+            warp_field = np.fliplr(warp_field)
+
+        # Get the alpha
+        a = np.expand_dims((warp_field / 255.0), axis=2)
+
+        # Color of playfield
+        c = np.concatenate([a*color[0], a*color[1], a*color[2]], axis=2)
+        
+        # Draw with specified alpha
+        a = a * alpha
+
+        # Superimpose the playing field image
+        result = (1.0-a)*image + a*c
+        result = np.clip(result, 0, 255).astype(np.uint8)
+        return result
+

common/loggers/image_preview.py

@@ -0,0 +1,93 @@
+
+from project.base.logging import Logger
+import torch
+import torch.nn as nn
+import numpy as np
+
+import methods.common.data.utils as utils
+
+
+class ImagePreviewLogger(Logger):
+    def __init__(
+        self,
+        experiment,
+        num_rows: int=3
+    ):
+        self.experiment = experiment
+        if (experiment is not None):
+            self.tracker = experiment.tracker
+        self.num_rows = num_rows
+
+        # Will be sampled later
+        self.samples = None
+        self.images = None
+        
+
+    def on_training_start(self):
+        
+        datamodule = self.experiment.datamodule
+        module = self.experiment.module
+
+        # Get the images
+        self.samples, images = self.sample_images(datamodule, module)
+        if (len(images) > 0):
+            self.images = torch.concatenate(
+                images, 
+                dim=0
+            ).to(module.device)
+
+
+    def on_epoch_end(self, epoch, stats):
+        self.make_preview()
+
+
+    def make_preview(self):
+
+        # Get the model
+        model = self.experiment.module.model
+        if (isinstance(model, nn.DataParallel)):
+            model = model.module
+
+        # Get the preview results
+        model.eval()
+        items = self.process_images(model, self.images)
+        model.train()
+
+        # Arange into grid, and send to tracker
+        log_images = {}
+        # Unpack the list of dicts
+        for item in items:
+            for key,image in item.items():
+                if (not key in log_images):
+                    log_images[key] = []
+                log_images[key].append(image)
+
+        # Arange images into grids
+        result = {}
+        for key, images in log_images.items():
+            result[key] = utils.make_grid(images, nrow=self.num_rows)
+
+        # Send to tracker
+        self.tracker.write_images(result)
+
+
+    def sample_dataset(self, dataset, num_images):
+        idx = np.random.choice(
+            len(dataset), 
+            size=(num_images,), 
+            replace=False
+        )
+        samples = [ dataset[i] for i in idx ]
+        return samples
+
+
+
+
+    def sample_images(self, datamodule, module):
+        """ Sample our datasets """
+        return [], []
+    
+
+    def process_images(self, model, images):
+        """ Returns list of dict[key,image] items """
+        return []

main.py

@@ -0,0 +1,199 @@
+import argparse
+import cv2
+import numpy as np
+import torch
+from pathlib import Path
+import time
+import traceback
+
+# 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")
+
+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
+
+    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:
+        # 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}")
+
+        # Prétraiter l'image
+        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")
+
+        # 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
+            # Utiliser l'échelle de base ESTIMÉE
+            minimap_offset_skeletons, actual_avg_scale = create_minimap_with_offset_skeletons(
+                player_list, 
+                homography_np, 
+                base_skeleton_scale=estimated_base_scale # Utiliser l'estimation
+            )
+
+            # 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() 

pose_estimator.py

@@ -0,0 +1,265 @@
+import torch
+import numpy as np
+import cv2
+from PIL import Image
+from transformers import AutoProcessor, RTDetrForObjectDetection, VitPoseForPoseEstimation
+from pathlib import Path
+
+# --- Global variables for models and processor (lazy loading) ---
+person_processor = None
+person_model = None
+pose_processor = None
+pose_model = None
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Pose Estimator: Using device: {device}")
+
+# --- Constantes pour la couleur et le dessin ---
+# Utilisation de tuples BGR pour les couleurs
+DEFAULT_MARKER_COLOR = (255, 255, 255) # Blanc
+MIN_PIXELS_FOR_COLOR = 20 # Nombre minimum de pixels valides dans la ROI pour tenter de calculer la couleur
+CONFIDENCE_THRESHOLD_KEYPOINTS = 0.3 # Seuil pour considérer un keypoint fiable pour la ROI et le dessin
+SKELETON_THICKNESS = 2
+
+# Définition des segments du squelette (indices COCO 0-16)
+# 0:Nose, 1:L_Eye, 2:R_Eye, 3:L_Ear, 4:R_Ear, 5:L_Shoulder, 6:R_Shoulder, 
+# 7:L_Elbow, 8:R_Elbow, 9:L_Wrist, 10:R_Wrist, 11:L_Hip, 12:R_Hip, 
+# 13:L_Knee, 14:R_Knee, 15:L_Ankle, 16:R_Ankle
+SKELETON_EDGES = [
+    # Tête
+    (0, 1), (0, 2), (1, 3), (2, 4),
+    # Torse
+    (5, 6), (5, 11), (6, 12), (11, 12),
+    # Bras Gauche
+    (5, 7), (7, 9),
+    # Bras Droit
+    (6, 8), (8, 10),
+    # Jambe Gauche
+    (11, 13), (13, 15),
+    # Jambe Droite
+    (12, 14), (14, 16)
+]
+
+# Indices des keypoints pour le torse et les chevilles
+TORSO_KP_INDICES = [5, 6, 11, 12] # Épaules, Hanches
+LEFT_ANKLE_KP_INDEX = 15
+RIGHT_ANKLE_KP_INDEX = 16
+
+def _load_models():
+    """Loads the models if they haven't been loaded yet."""
+    global person_processor, person_model, pose_processor, pose_model
+    
+    if person_processor is None:
+        print("Loading RTDetr person detector model...")
+        person_processor = AutoProcessor.from_pretrained("PekingU/rtdetr_r50vd_coco_o365")
+        person_model = RTDetrForObjectDetection.from_pretrained("PekingU/rtdetr_r50vd_coco_o365", device_map=device)
+        print("✓ RTDetr loaded.")
+        
+    if pose_processor is None:
+        print("Loading ViTPose pose estimation model...")
+        pose_processor = AutoProcessor.from_pretrained("usyd-community/vitpose-base-simple")
+        pose_model = VitPoseForPoseEstimation.from_pretrained("usyd-community/vitpose-base-simple", device_map=device)
+        print("✓ ViTPose loaded.")
+
+def _is_color_greenish(bgr_pixel, threshold=10):
+    b, g, r = bgr_pixel
+    return g > b + threshold and g > r + threshold
+
+def _is_color_grayscale(bgr_pixel, tolerance=30):
+     b, g, r = bgr_pixel
+     min_val, max_val = min(b, g, r), max(b, g, r)
+     is_dark = max_val < 50
+     is_light = min_val > 200
+     is_low_saturation = (max_val - min_val) < tolerance
+     return is_dark or is_light or is_low_saturation
+
+def _get_average_color(roi_bgr):
+    """Calcule la couleur moyenne d'une ROI après filtrage."""
+    if roi_bgr is None or roi_bgr.size == 0:
+        return DEFAULT_MARKER_COLOR
+
+    try:
+        pixels = roi_bgr.reshape(-1, 3)
+        valid_pixels = []
+        for pixel in pixels:
+            if not _is_color_greenish(pixel) and not _is_color_grayscale(pixel):
+                valid_pixels.append(pixel)
+        
+        if len(valid_pixels) < MIN_PIXELS_FOR_COLOR:
+            return DEFAULT_MARKER_COLOR
+
+        avg_color = np.mean(valid_pixels, axis=0)
+        return tuple(map(int, avg_color))
+
+    except Exception as e:
+        print(f"  Erreur calcul couleur moyenne: {e}. Utilisation couleur défaut.")
+        return DEFAULT_MARKER_COLOR
+
+def get_player_data(image_bgr: np.ndarray) -> list:
+    """
+    Detects persons, estimates pose, calculates average torso color, 
+    and returns a list of data for each player.
+    
+    Args:
+        image_bgr: The input image in BGR format (NumPy array).
+
+    Returns:
+        A list of dictionaries, each containing:
+        {
+            'keypoints': np.ndarray (17, 2), 
+            'scores': np.ndarray (17,),
+            'bbox': np.ndarray (4,) [x1, y1, x2, y2],
+            'avg_color': tuple (b, g, r)
+        }
+        Returns an empty list if no persons are detected or an error occurs.
+    """
+    _load_models() 
+    player_list = []
+    height, width = image_bgr.shape[:2]
+
+    try:
+        image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
+        image_pil = Image.fromarray(image_rgb)
+
+        # --- Stage 1: Detect humans ---
+        inputs_det = person_processor(images=image_pil, return_tensors="pt").to(device)
+        with torch.no_grad():
+            outputs_det = person_model(**inputs_det)
+        results_det = person_processor.post_process_object_detection(
+            outputs_det, target_sizes=torch.tensor([(height, width)]), threshold=0.5
+        )
+        result_det = results_det[0]
+        person_boxes = result_det["boxes"][result_det["labels"] == 0].cpu().numpy()
+
+        if len(person_boxes) == 0:
+            print("No persons detected.")
+            return player_list
+
+        person_boxes_coco = person_boxes.copy()
+        person_boxes_coco[:, 2] = person_boxes_coco[:, 2] - person_boxes_coco[:, 0]
+        person_boxes_coco[:, 3] = person_boxes_coco[:, 3] - person_boxes_coco[:, 1]
+
+        # --- Stage 2: Detect keypoints ---
+        inputs_pose = pose_processor(image_pil, boxes=[person_boxes_coco], return_tensors="pt").to(device)
+        with torch.no_grad():
+            outputs_pose = pose_model(**inputs_pose)
+        pose_results = pose_processor.post_process_pose_estimation(outputs_pose, boxes=[person_boxes_coco])
+        image_pose_result = pose_results[0] 
+
+        if not image_pose_result:
+             print("Pose estimation did not return results.")
+             return player_list
+        
+        # --- Stage 3: Process each person --- 
+        for i, person_box_xyxy in enumerate(person_boxes):
+            if i >= len(image_pose_result): continue 
+            
+            pose_result = image_pose_result[i]
+            xy = pose_result['keypoints'].cpu().numpy()
+            scores = pose_result['scores'].cpu().numpy()
+
+            # Ensure xy shape is correct before proceeding
+            if xy.shape != (17, 2):
+                print(f"Person {i}: Unexpected keypoints shape {xy.shape}, skipping.")
+                continue
+
+            # -- Define Torso ROI --
+            reliable_torso_keypoints = xy[TORSO_KP_INDICES][scores[TORSO_KP_INDICES] > CONFIDENCE_THRESHOLD_KEYPOINTS]
+            x1_box, y1_box, x2_box, y2_box = map(int, person_box_xyxy)
+            box_h = y2_box - y1_box
+            box_w = x2_box - x1_box
+            if len(reliable_torso_keypoints) >= 3:
+                min_x_kp = int(np.min(reliable_torso_keypoints[:, 0]))
+                max_x_kp = int(np.max(reliable_torso_keypoints[:, 0]))
+                min_y_kp = int(np.min(reliable_torso_keypoints[:, 1]))
+                max_y_kp = int(np.max(reliable_torso_keypoints[:, 1]))
+                roi_x1 = max(x1_box, min_x_kp - 5); roi_y1 = max(y1_box, min_y_kp - 5)
+                roi_x2 = min(x2_box, max_x_kp + 5); roi_y2 = min(y2_box, max_y_kp + 5)
+            else: 
+                roi_x1 = x1_box; roi_y1 = y1_box + int(0.1 * box_h) 
+                roi_x2 = x2_box; roi_y2 = y1_box + int(0.6 * box_h) 
+            roi_x1 = max(0, roi_x1); roi_y1 = max(0, roi_y1)
+            roi_x2 = min(width, roi_x2); roi_y2 = min(height, roi_y2)
+
+            # -- Extract Average Color --
+            avg_color = DEFAULT_MARKER_COLOR
+            if roi_y2 > roi_y1 and roi_x2 > roi_x1: 
+                torso_roi = image_bgr[roi_y1:roi_y2, roi_x1:roi_x2]
+                avg_color = _get_average_color(torso_roi)
+            # else: # Pas besoin de message si ROI invalide, couleur par défaut suffit
+                # print(f"Person {i}: Invalid ROI, using default color.")
+            
+            # -- Store player data --
+            player_data = {
+                'keypoints': xy,
+                'scores': scores,
+                'bbox': person_box_xyxy, # Utiliser la bbox originale xyxy
+                'avg_color': avg_color
+            }
+            player_list.append(player_data)
+
+    except Exception as e:
+        print(f"Error during player data extraction: {e}")
+        import traceback
+        traceback.print_exc()
+        # Retourner une liste vide en cas d'erreur majeure
+        return []
+
+    return player_list
+
+# Example usage (optional, for testing the module directly)
+if __name__ == '__main__':
+    test_image_path = 'img3.png' 
+    
+    if not Path(test_image_path).exists():
+         print(f"Test image not found: {test_image_path}")
+    else:
+        print(f"Testing player data extraction with image: {test_image_path}")
+        input_img = cv2.imread(test_image_path)
+        
+        if input_img is None:
+            print(f"Failed to load test image: {test_image_path}")
+        else:
+            print("Getting player data...")
+            players = get_player_data(input_img)
+            print(f"✓ Found data for {len(players)} players.")
+
+            # --- Draw markers and info on original image for testing --- 
+            output_img_test = input_img.copy()
+            for idx, p_data in enumerate(players):
+                kps = p_data['keypoints']
+                scores = p_data['scores']
+                bbox = p_data['bbox']
+                color = p_data['avg_color']
+
+                # Determine reference point (ankles or bbox bottom mid)
+                l_ankle_pt = kps[LEFT_ANKLE_KP_INDEX]
+                r_ankle_pt = kps[RIGHT_ANKLE_KP_INDEX]
+                l_ankle_score = scores[LEFT_ANKLE_KP_INDEX]
+                r_ankle_score = scores[RIGHT_ANKLE_KP_INDEX]
+                
+                ref_point = None
+                if l_ankle_score > CONFIDENCE_THRESHOLD_KEYPOINTS and r_ankle_score > CONFIDENCE_THRESHOLD_KEYPOINTS:
+                    ref_point = tuple(map(int, (l_ankle_pt + r_ankle_pt) / 2))
+                elif l_ankle_score > CONFIDENCE_THRESHOLD_KEYPOINTS:
+                    ref_point = tuple(map(int, l_ankle_pt))
+                elif r_ankle_score > CONFIDENCE_THRESHOLD_KEYPOINTS:
+                    ref_point = tuple(map(int, r_ankle_pt))
+                else:
+                    x1, y1, x2, y2 = map(int, bbox)
+                    ref_point = (int((x1 + x2) / 2), y2) 
+
+                # Draw marker at reference point
+                if ref_point:
+                    cv2.circle(output_img_test, ref_point, 8, color, -1, cv2.LINE_AA)
+                    cv2.circle(output_img_test, ref_point, 8, (0,0,0), 1, cv2.LINE_AA) # Black outline
+                    # Draw player index
+                    cv2.putText(output_img_test, str(idx), (ref_point[0]+5, ref_point[1]-5), 
+                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,0), 2, cv2.LINE_AA)
+                    cv2.putText(output_img_test, str(idx), (ref_point[0]+5, ref_point[1]-5), 
+                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 1, cv2.LINE_AA)
+            
+            cv2.imshow("Original Image", input_img)
+            cv2.imshow("Player Markers Test", output_img_test)
+            print("Displaying test results. Press any key to exit.")
+            cv2.waitKey(0)
+            cv2.destroyAllWindows() 

requirements.txt

@@ -0,0 +1,22 @@
+# Installez PyTorch en utilisant la commande spécifique si nécessaire :
+# 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
+torchvision
+torchaudio
+numpy
+opencv-python
+pytorch-lightning
+soccernet
+kornia
+
+# Ajouté car requis par sn_segmentation 
+
+# Dépendances pour l'estimation de pose (ViTPose)
+transformers
+supervision
+Pillow # Souvent une dépendance de transformers/supervision, mais explicite ici 
+accelerate
+# scikit-learn # Retiré car K-Means n'est plus utilisé
+gradio

tvcalib/cam_distr/tv_main_behind.py

@@ -0,0 +1,77 @@
+from math import pi
+from tvcalib.utils.data_distr import mean_std_with_confidence_interval
+
+
+def get_cam_distr(sigma_scale: float, batch_dim: int, temporal_dim: int):
+    cam_distr = {
+        "pan": {
+            "minmax": (pi / 4, 3 * pi / 4),  # in deg 45°, 135°
+            "dimension": (
+                batch_dim,
+                temporal_dim,
+            ),
+        },
+        "tilt": {
+            "minmax": (pi / 16, pi / 2),  # in deg 11.25°, 90°
+            "dimension": (
+                batch_dim,
+                temporal_dim,
+            ),
+        },
+        "roll": {
+            "minmax": (-pi / 32, pi / 32),  # in deg -5.6°, 5.6°
+            "dimension": (
+                batch_dim,
+                temporal_dim,
+            ),
+        },
+        "aov": {
+            "minmax": (pi / 22, pi / 2),  # (8.2°, 90°)
+            "dimension": (
+                batch_dim,
+                temporal_dim,
+            ),
+        },
+        "c_x": {
+            "minmax": (-32.5, -52.5),
+            "dimension": (
+                batch_dim,
+                1,
+            ),
+        },
+        "c_y": {
+            "minmax": (-5.0, 5.0),
+            "dimension": (
+                batch_dim,
+                1,
+            ),
+        },
+        "c_z": {
+            "minmax": (-35.0, -1.0),
+            "dimension": (
+                batch_dim,
+                1,
+            ),
+        },
+    }
+
+    for k, params in cam_distr.items():
+        cam_distr[k]["mean_std"] = mean_std_with_confidence_interval(
+            *params["minmax"], sigma_scale=sigma_scale
+        )
+    return cam_distr
+
+
+def get_dist_distr(batch_dim: int, temporal_dim: int, _sigma_scale: float = 2.57):
+    return {
+        "k1": {
+            "minmax": [0.0, 0.5],  # we clip min(0.0, x)
+            "mean_std": (0.0, _sigma_scale * 0.5),
+            "dimension": (batch_dim, temporal_dim),
+        },
+        "k2": {
+            "minmax": [-0.1, 0.1],
+            "mean_std": (0.0, _sigma_scale * 0.1),
+            "dimension": (batch_dim, temporal_dim),
+        },
+    }

tvcalib/cam_distr/tv_main_center.py

@@ -0,0 +1,78 @@
+import numpy as np
+from math import pi
+from ..utils.data_distr import mean_std_with_confidence_interval
+
+
+def get_cam_distr(sigma_scale: float, batch_dim: int, temporal_dim: int):
+    cam_distr = {
+        "pan": {
+            "minmax": (-pi / 4, pi / 4),  # in deg -45°, 45°
+            "dimension": (
+                batch_dim,
+                temporal_dim,
+            ),
+        },
+        "tilt": {
+            "minmax": (pi / 4, pi / 2),  # in deg 45°, 90°
+            "dimension": (
+                batch_dim,
+                temporal_dim,
+            ),
+        },
+        "roll": {
+            "minmax": (-pi / 18, pi / 18),  # in deg -10°, 10°
+            "dimension": (
+                batch_dim,
+                temporal_dim,
+            ),
+        },
+        "aov": {
+            "minmax": (pi / 22, pi / 2),  # (8.2°, 90°)
+            "dimension": (
+                batch_dim,
+                temporal_dim,
+            ),
+        },
+        "c_x": {
+            "minmax": (-12.0, 12.0),  # (-36, 36) for entire main tribune
+            "dimension": (
+                batch_dim,
+                1,
+            ),
+        },
+        "c_y": {
+            "minmax": (40.0, 110.0),
+            "dimension": (
+                batch_dim,
+                1,
+            ),
+        },
+        "c_z": {
+            "minmax": (-40.0, -5.0),
+            "dimension": (
+                batch_dim,
+                1,
+            ),
+        },
+    }
+
+    for k, params in cam_distr.items():
+        cam_distr[k]["mean_std"] = mean_std_with_confidence_interval(
+            *params["minmax"], sigma_scale=sigma_scale
+        )
+    return cam_distr
+
+
+def get_dist_distr(batch_dim: int, temporal_dim: int, _sigma_scale: float = 2.57):
+    return {
+        "k1": {
+            "minmax": [0.0, 0.5],  # we clip min(0.0, x)
+            "mean_std": (0.0, _sigma_scale * 0.5),
+            "dimension": (batch_dim, temporal_dim),
+        },
+        "k2": {
+            "minmax": [-0.1, 0.1],
+            "mean_std": (0.0, _sigma_scale * 0.1),
+            "dimension": (batch_dim, temporal_dim),
+        },
+    }

tvcalib/cam_distr/tv_main_left.py

@@ -0,0 +1,77 @@
+from math import pi
+from tvcalib.utils.data_distr import mean_std_with_confidence_interval
+
+
+def get_cam_distr(sigma_scale: float, batch_dim: int, temporal_dim: int):
+    cam_distr = {
+        "pan": {
+            "minmax": (-pi / 4, pi / 4),  # in deg -45°, 45°
+            "dimension": (
+                batch_dim,
+                temporal_dim,
+            ),
+        },
+        "tilt": {
+            "minmax": (pi / 4, pi / 2),  # in deg 45°, 90°
+            "dimension": (
+                batch_dim,
+                temporal_dim,
+            ),
+        },
+        "roll": {
+            "minmax": (-pi / 18, pi / 18),  # in deg -10°, 10°
+            "dimension": (
+                batch_dim,
+                temporal_dim,
+            ),
+        },
+        "aov": {
+            "minmax": (pi / 22, pi / 2),  # (8.2°, 90°)
+            "dimension": (
+                batch_dim,
+                temporal_dim,
+            ),
+        },
+        "c_x": {
+            "minmax": (-36 - 16.5, -36 + 16.5),
+            "dimension": (
+                batch_dim,
+                1,
+            ),
+        },
+        "c_y": {
+            "minmax": (40.0, 110.0),
+            "dimension": (
+                batch_dim,
+                1,
+            ),
+        },
+        "c_z": {
+            "minmax": (-40.0, -5.0),
+            "dimension": (
+                batch_dim,
+                1,
+            ),
+        },
+    }
+
+    for k, params in cam_distr.items():
+        cam_distr[k]["mean_std"] = mean_std_with_confidence_interval(
+            *params["minmax"], sigma_scale=sigma_scale
+        )
+    return cam_distr
+
+
+def get_dist_distr(batch_dim: int, temporal_dim: int, _sigma_scale: float = 2.57):
+    return {
+        "k1": {
+            "minmax": [0.0, 0.5],  # we clip min(0.0, x)
+            "mean_std": (0.0, _sigma_scale * 0.5),
+            "dimension": (batch_dim, temporal_dim),
+        },
+        "k2": {
+            "minmax": [-0.1, 0.1],
+            "mean_std": (0.0, _sigma_scale * 0.1),
+            "dimension": (batch_dim, temporal_dim),
+        },
+    }

tvcalib/cam_distr/tv_main_right.py

@@ -0,0 +1,77 @@
+from math import pi
+from tvcalib.utils.data_distr import mean_std_with_confidence_interval
+
+
+def get_cam_distr(sigma_scale: float, batch_dim: int, temporal_dim: int):
+    cam_distr = {
+        "pan": {
+            "minmax": (-pi / 4, pi / 4),  # in deg -45°, 45°
+            "dimension": (
+                batch_dim,
+                temporal_dim,
+            ),
+        },
+        "tilt": {
+            "minmax": (pi / 4, pi / 2),  # in deg 45°, 90°
+            "dimension": (
+                batch_dim,
+                temporal_dim,
+            ),
+        },
+        "roll": {
+            "minmax": (-pi / 18, pi / 18),  # in deg -10°, 10°
+            "dimension": (
+                batch_dim,
+                temporal_dim,
+            ),
+        },
+        "aov": {
+            "minmax": (pi / 22, pi / 2),  # (8.2°, 90°)
+            "dimension": (
+                batch_dim,
+                temporal_dim,
+            ),
+        },
+        "c_x": {
+            "minmax": (36 - 16.5, 36 + 16.5),
+            "dimension": (
+                batch_dim,
+                1,
+            ),
+        },
+        "c_y": {
+            "minmax": (40.0, 110.0),
+            "dimension": (
+                batch_dim,
+                1,
+            ),
+        },
+        "c_z": {
+            "minmax": (-40.0, -5.0),
+            "dimension": (
+                batch_dim,
+                1,
+            ),
+        },
+    }
+
+    for k, params in cam_distr.items():
+        cam_distr[k]["mean_std"] = mean_std_with_confidence_interval(
+            *params["minmax"], sigma_scale=sigma_scale
+        )
+    return cam_distr
+
+
+def get_dist_distr(batch_dim: int, temporal_dim: int, _sigma_scale: float = 2.57):
+    return {
+        "k1": {
+            "minmax": [0.0, 0.5],  # we clip min(0.0, x)
+            "mean_std": (0.0, _sigma_scale * 0.5),
+            "dimension": (batch_dim, temporal_dim),
+        },
+        "k2": {
+            "minmax": [-0.1, 0.1],
+            "mean_std": (0.0, _sigma_scale * 0.1),
+            "dimension": (batch_dim, temporal_dim),
+        },
+    }

tvcalib/cam_distr/tv_main_tribune.py

@@ -0,0 +1,77 @@
+from math import pi
+from tvcalib.utils.data_distr import mean_std_with_confidence_interval
+
+
+def get_cam_distr(sigma_scale: float, batch_dim: int, temporal_dim: int):
+    cam_distr = {
+        "pan": {
+            "minmax": (-pi / 4, pi / 4),  # in deg -45°, 45°
+            "dimension": (
+                batch_dim,
+                temporal_dim,
+            ),
+        },
+        "tilt": {
+            "minmax": (pi / 4, pi / 2),  # in deg 45°, 90°
+            "dimension": (
+                batch_dim,
+                temporal_dim,
+            ),
+        },
+        "roll": {
+            "minmax": (-pi / 18, pi / 18),  # in deg -10°, 10°
+            "dimension": (
+                batch_dim,
+                temporal_dim,
+            ),
+        },
+        "aov": {
+            "minmax": (pi / 22, pi / 2),  # (8.2°, 90°)
+            "dimension": (
+                batch_dim,
+                temporal_dim,
+            ),
+        },
+        "c_x": {
+            "minmax": (-40.0, 40.0),  # entire main tribune
+            "dimension": (
+                batch_dim,
+                1,
+            ),
+        },
+        "c_y": {
+            "minmax": (40.0, 110.0),
+            "dimension": (
+                batch_dim,
+                1,
+            ),
+        },
+        "c_z": {
+            "minmax": (-40.0, -5.0),
+            "dimension": (
+                batch_dim,
+                1,
+            ),
+        },
+    }
+
+    for k, params in cam_distr.items():
+        cam_distr[k]["mean_std"] = mean_std_with_confidence_interval(
+            *params["minmax"], sigma_scale=sigma_scale
+        )
+    return cam_distr
+
+
+def get_dist_distr(batch_dim: int, temporal_dim: int, _sigma_scale: float = 2.57):
+    return {
+        "k1": {
+            "minmax": [0.0, 0.5],  # we clip min(0.0, x)
+            "mean_std": (0.0, _sigma_scale * 0.5),
+            "dimension": (batch_dim, temporal_dim),
+        },
+        "k2": {
+            "minmax": [-0.1, 0.1],
+            "mean_std": (0.0, _sigma_scale * 0.1),
+            "dimension": (batch_dim, temporal_dim),
+        },
+    }

tvcalib/cam_modules.py

@@ -0,0 +1,583 @@
+from typing import Tuple, Dict, Union
+
+from pytorch_lightning import LightningModule
+import torch
+import kornia
+import torch.nn as nn
+from .utils.data_distr import FeatureScalerZScore
+
+
+class CameraParameterWLensDistDictZScore(LightningModule):
+    """Holds individual camera parameters including lens distortion parameters as nn.Modul"""
+
+    def __init__(self, cam_distr, dist_distr, device="cpu"):
+        super(CameraParameterWLensDistDictZScore, self).__init__()
+
+        self.cam_distr = cam_distr
+        self._device = device
+
+        # phi raw
+        self.param_dict = torch.nn.ParameterDict(
+            {
+                k: torch.nn.parameter.Parameter(
+                    torch.zeros(
+                        *cam_distr[k]["dimension"],
+                        device=device,
+                    ),
+                    requires_grad=False
+                    if ("no_grad" in cam_distr[k]) and (cam_distr[k]["no_grad"] == True)
+                    else True,
+                )
+                for k in cam_distr.keys()
+            }
+        )
+
+        # denormalization module to get phi_target
+        self.feature_scaler = torch.nn.ModuleDict(
+            {k: FeatureScalerZScore(*cam_distr[k]["mean_std"]) for k in cam_distr.keys()}
+        )
+
+        self.dist_distr = dist_distr
+        if self.dist_distr is not None:
+            self.param_dict_dist = torch.nn.ParameterDict(
+                {
+                    k: torch.nn.Parameter(torch.zeros(*dist_distr[k]["dimension"], device=device))
+                    for k in dist_distr.keys()
+                }
+            )
+            # TODO: modify later to dynamically cunstruct a tensor of shape (k_1,k_2,p_1,p_2[,k_3[,k_4,k_5,k_6[,s_1,s_2,s_3,s_4[,\tau_x,\tau_y]]]])
+            #
+
+            self.feature_scaler_dist_coeff = torch.nn.ModuleDict(
+                {k: FeatureScalerZScore(*dist_distr[k]["mean_std"]) for k in dist_distr.keys()}
+            )
+
+    def initialize(
+        self,
+        update_dict_cam: Union[Dict[str, Union[float, torch.tensor]], None],
+        update_dict_dist=None,
+    ):
+        """Initializes all camera parameters with zeros and replace specific values with provided values
+
+        Args:
+            update_dict_cam (Dict[str, Union[float, torch.tensor]]): Parameters to be updated
+        """
+
+        for k in self.param_dict.keys():
+            self.param_dict[k].data = torch.zeros(
+                *self.cam_distr[k]["dimension"], device=self._device
+            )
+        if self.dist_distr is not None:
+            for k in self.dist_distr.keys():
+                self.param_dict_dist[k].data = torch.zeros(
+                    *self.dist_distr[k]["dimension"], device=self._device
+                )
+
+        if update_dict_cam is not None and len(update_dict_cam) > 0:
+            for k, v in update_dict_cam.items():
+                self.param_dict[k].data = (
+                    torch.zeros(*self.cam_distr[k]["dimension"], device=self._device) + v
+                )
+        if update_dict_dist is not None:
+            raise NotImplementedError
+
+    def forward(self):
+        phi_dict = {}
+        for k, param in self.param_dict.items():
+            phi_dict[k] = self.feature_scaler[k](param)
+
+        if self.dist_distr is None:
+            return phi_dict, None
+
+        # This is a vector with 4, 5, 8, 12 or 14 elements with shape :math:`(*, n)` depending on the provided dict of coefficients
+        # assumes dict is ordered according (k_1,k_2,p_1,p_2[,k_3[,k_4,k_5,k_6[,s_1,s_2,s_3,s_4[,\tau_x,\tau_y]]]])
+        psi = torch.stack(
+            [
+                torch.clamp(
+                    self.feature_scaler_dist_coeff[k](param),
+                    min=self.dist_distr[k]["minmax"][0],
+                    max=self.dist_distr[k]["minmax"][1],
+                )
+                for k, param in self.param_dict_dist.items()
+            ],
+            dim=-1,  # stack individual features and not arbirary leading dimensions
+        )
+
+        return phi_dict, psi
+
+
+class SNProjectiveCamera:
+    def __init__(
+        self,
+        phi_dict: Dict[str, torch.tensor],
+        psi: torch.tensor,
+        principal_point: Tuple[float, float],
+        image_width: int,
+        image_height: int,
+        device: str = "cpu",
+        nan_check=True,
+    ) -> None:
+        """Projective camera defined as K @ R [I|-t] with lens distortion module and batch dimensions B,T.
+
+        Following Euler angles convention, we use a ZXZ succession of intrinsic rotations in order to describe
+        the orientation of the camera. Starting from the world reference axis system, we first apply a rotation
+        around the Z axis to pan the camera. Then the obtained axis system is rotated around its x axis in order to tilt the camera.
+        Then the last rotation around the z axis of the new axis system alows to roll the camera. Note that this z axis is the principal axis of the camera.
+
+        As T is not provided for camra location and lens distortion, these parameters are assumed to be fixed accross T.
+        phi_dict is a dict of parameters containing:
+        {
+            'aov_x, torch.Size([B, T])',
+            'pan, torch.Size([B, T])',
+            'tilt, torch.Size([B, T])',
+            'roll, torch.Size([B, T])',
+            'c_x, torch.Size([B, 1])',
+            'c_y, torch.Size([B, 1])',
+            'c_z, torch.Size([B, 1])',
+        }
+
+        Internally fuses B and T dimension to pseudo batch dimension.
+        {
+            'aov_x, torch.Size([B*T])',
+            'pan, torch.Size([B*T])',
+            'tilt, torch.Size([B*T])'
+            'roll, torch.Size([B*T])',
+            'c_x, torch.Size([B])',
+            'c_y, torch.Size([B])',
+            'c_z, torch.Size([B])',
+            }
+
+        aov_x, pan, tilt, roll are assumed in radian.
+
+        Note on lens distortion:
+            Lens distortion coefficients are independent from image resolution!
+            We I(dist_points(K_ndc, dist_coeff, points2d_ndc)) == I(dist_points(K_raster, dist_coeff, points2d_raster))
+
+        Args:
+            phi_dict (Dict[str, torch.tensor]): See example above
+            psi (Union[None, torch.Tensor]): distortion coefficients as concatinated vector according to https://kornia.readthedocs.io/en/latest/geometry.calibration.html of shape (B, T, {2, 4, 5,8,12, 14})
+            principal_point (Tuple[float, float]): Principal point assumed to be fixed across all samples (B,T,)
+            image_width (int): assumed to be fixed across all samples (B,T,)
+            image_height (int): assumed to be fixed across all samples (B,T,)
+        """
+
+        # fuse B and T dimension
+        phi_dict_flat = {}
+        for k, v in phi_dict.items():
+            if len(v.shape) == 2:
+                phi_dict_flat[k] = v.view(v.shape[0] * v.shape[1])
+            elif len(v.shape) == 3:
+                phi_dict_flat[k] = v.view(v.shape[0] * v.shape[1], v.shape[-1])
+
+        self.batch_dim, self.temporal_dim = phi_dict["pan"].shape
+        self.pseudo_batch_size = phi_dict_flat["pan"].shape[0]
+        self.phi_dict_flat = phi_dict_flat
+
+        self.principal_point = principal_point
+        self.image_width = image_width
+        self.image_height = image_height
+        self.device = device
+
+        self.psi = psi
+        if self.psi is not None:
+            if self.psi.shape[-1] != 2:
+                raise NotImplementedError
+
+            # :math:`(k_1,k_2,p_1,p_2[,k_3[,k_4,k_5,k_6[,s_1,s_2,s_3,s_4[,\tau_x,\tau_y]]]])`.
+            # psi is a vector with 2, 4, 5, 8, 12 or 14 elements with shape :math:`(*, n)`.
+            if self.psi.shape[-1] == 2:
+                # assume zero tangential coefficients
+                psi_ext = torch.zeros(*list(self.psi.shape[:-1]), 4)
+                psi_ext[..., :2] = self.psi
+                self.psi = psi_ext
+            self.lens_dist_coeff = self.psi.view(self.pseudo_batch_size, self.psi.shape[-1]).to(
+                self.device
+            )
+
+        self.intrinsics_ndc = self.construct_intrinsics_ndc()
+        self.intrinsics_raster = self.construct_intrinsics_raster()
+
+        self.rotation = self.rotation_from_euler_angles(
+            *[phi_dict_flat[k] for k in ["pan", "tilt", "roll"]]
+        )
+        self.position = torch.stack([phi_dict_flat[k] for k in ["c_x", "c_y", "c_z"]], dim=-1)
+        self.position = self.position.repeat_interleave(
+            int(self.pseudo_batch_size / self.batch_dim), dim=0
+        )  # (B, 3) # TODO: probably needs modification if B > 0?
+        self.P_ndc = self.construct_projection_matrix(self.intrinsics_ndc)
+        self.P_raster = self.construct_projection_matrix(self.intrinsics_raster)
+        self.phi_dict = phi_dict
+
+        self.nan_check = nan_check
+        super().__init__()
+
+    def construct_projection_matrix(self, intrinsics):
+        It = torch.eye(4, device=self.device)[:-1].repeat(self.pseudo_batch_size, 1, 1)
+        It[:, :, -1] = -self.position  # (B, 3, 4)
+        self.It = It
+        return intrinsics @ self.rotation @ It  #  # (B, 3, 4)
+
+    def construct_intrinsics_ndc(self):
+        # assume that the principal point is (0,0)
+        K = torch.eye(3, requires_grad=False, device=self.device)
+        K = K.reshape((1, 3, 3)).repeat(self.pseudo_batch_size, 1, 1)
+        K[:, 0, 0] = self.get_fl_from_aov_rad(self.phi_dict_flat["aov"], d=2)
+        K[:, 1, 1] = self.get_fl_from_aov_rad(
+            self.phi_dict_flat["aov"], d=2 * self.image_width / self.image_height
+        )
+        return K
+
+    def construct_intrinsics_raster(self):
+        # assume that the principal point is (W/2,H/2)
+        K = torch.eye(3, requires_grad=False, device=self.device)
+        K = K.reshape((1, 3, 3)).repeat(self.pseudo_batch_size, 1, 1)
+        K[:, 0, 0] = self.get_fl_from_aov_rad(self.phi_dict_flat["aov"], d=self.image_width)
+        K[:, 1, 1] = self.get_fl_from_aov_rad(self.phi_dict_flat["aov"], d=self.image_width)
+        K[:, 0, 2] = self.principal_point[0]
+        K[:, 1, 2] = self.principal_point[1]
+        return K
+
+    def __str__(self) -> str:
+        return f"aov_deg={torch.rad2deg(self.phi_dict['aov'])}, t={torch.stack([self.phi_dict[k] for k in ['c_x', 'c_y', 'c_z']], dim=-1)}, pan_deg={torch.rad2deg(self.phi_dict['pan'])} tilt_deg={torch.rad2deg(self.phi_dict['tilt'])} roll_deg={torch.rad2deg(self.phi_dict['roll'])}"
+
+    def str_pan_tilt_roll_fl(self, b, t):
+        r = f"FOV={torch.rad2deg(self.phi_dict['aov'][b, t]):.1f}°, pan={torch.rad2deg(self.phi_dict['pan'][b, t]):.1f}° tilt={torch.rad2deg(self.phi_dict['tilt'][b, t]):.1f}° roll={torch.rad2deg(self.phi_dict['roll'][b, t]):.1f}°"
+        return r
+
+    def str_lens_distortion_coeff(self, b):
+        # TODO: T! also need indivudual lens_dist_coeff for each t in T
+        # print(self.lens_dist_coeff.shape)
+        return f"lens dist coeff=" + " ".join(
+            [f"{x:.2f}" for x in self.lens_dist_coeff[b, :2]]
+        )  # print only radial lens dist. coeff
+
+    def __repr__(self) -> str:
+        return f"{self.__class__}:" + self.__str__()
+
+    def __len__(self):
+        return self.pseudo_batch_size  # e.g. self.intrinsics.shape[0]
+
+    def project_point2pixel(self, points3d: torch.tensor, lens_distortion: bool) -> torch.tensor:
+        """Project world coordinates to pixel coordinates.
+
+        Args:
+            points3d (torch.tensor): of shape (N, 3) or (1, N, 3)
+
+        Returns:
+            torch.tensor: projected points of shape (B, T, N, 2)
+        """
+        position = self.position.view(self.pseudo_batch_size, 1, 3)
+        point = points3d - position
+        rotated_point = self.rotation @ point.transpose(1, 2)  # (pseudo_batch_size, 3, N)
+        dist_point2cam = rotated_point[:, 2]  # (B, N) distance pixel to world point
+        dist_point2cam = dist_point2cam.view(self.pseudo_batch_size, 1, rotated_point.shape[-1])
+        rotated_point = rotated_point / dist_point2cam  # (B, 3, N) / (B, 1, N) -> (B, 3, N)
+
+        projected_points = self.intrinsics_raster @ rotated_point  # (B, 3, N)
+        # transpose vs view? here
+        projected_points = projected_points.transpose(-1, -2)  # cannot use view()
+        projected_points = kornia.geometry.convert_points_from_homogeneous(projected_points)
+        if lens_distortion:
+            if self.psi is None:
+                raise RuntimeError("Lens distortion requested, but deactivated in module")
+            projected_points = self.distort_points(projected_points, self.intrinsics_raster)
+
+        # reshape back from (pseudo_batch_size, N, 2) to (B, T, N, 2)
+        projected_points = projected_points.view(
+            self.batch_dim, self.temporal_dim, projected_points.shape[-2], 2
+        )
+        if self.nan_check:
+            if torch.isnan(projected_points).any().item():
+                print(self.phi_dict_flat)
+                print(projected_points)
+                raise RuntimeWarning("NaN in project_point2pixel")
+        return projected_points
+
+    def project_point2ndc(self, points3d: torch.tensor, lens_distortion: bool) -> torch.tensor:
+        """Project world coordinates to pixel coordinates.
+
+        Args:
+            points3d (torch.tensor): of shape (N, 3) or (1, N, 3)
+
+        Returns:
+            torch.tensor: projected points of shape (B, T, N, 2)
+        """
+        position = self.position.view(self.pseudo_batch_size, 1, 3)
+        point = points3d - position
+        rotated_point = self.rotation @ point.transpose(1, 2)  # (pseudo_batch_size, 3, N)
+        dist_point2cam = rotated_point[:, 2]  # (B, N) distance pixel to world point
+        dist_point2cam = dist_point2cam.view(self.pseudo_batch_size, 1, rotated_point.shape[-1])
+        rotated_point = rotated_point / dist_point2cam  # (B, 3, N) / (B, 1, N) -> (B, 3, N)
+
+        projected_points = self.intrinsics_ndc @ rotated_point  # (B, 3, N)
+        # transpose vs view? here
+        projected_points = projected_points.transpose(-1, -2)  # cannot use view()
+        projected_points = kornia.geometry.convert_points_from_homogeneous(projected_points)
+        if self.nan_check:
+            if torch.isnan(projected_points).any().item():
+                print(projected_points)
+                print(self.phi_dict_flat)
+                print("lens distortion", self.lens_dist_coeff)
+
+                raise RuntimeWarning("NaN in project_point2ndc before distort")
+        if lens_distortion:
+            if self.psi is None:
+                raise RuntimeError("Lens distortion requested, but deactivated in module")
+            projected_points = self.distort_points(projected_points, self.intrinsics_ndc)
+
+        # reshape back from (pseudo_batch_size, N, 2) to (B, T, N, 2)
+        projected_points = projected_points.view(
+            self.batch_dim, self.temporal_dim, projected_points.shape[-2], 2
+        )
+        if self.nan_check:
+            if torch.isnan(projected_points).any().item():
+                print(self.phi_dict_flat)
+                print(projected_points)
+                raise RuntimeWarning("NaN in project_point2ndc after distort")
+        return projected_points
+
+    def project_point2pixel_from_P(
+        self, points3d: torch.tensor, lens_distortion: bool
+    ) -> torch.tensor:
+        """Project world coordinates to pixel coordinates from the projection matrix.
+
+        Args:
+            points3d (torch.tensor): of shape (1, N, 3)
+
+        Returns:
+            torch.tensor: projected points of shape (B, T, N, 2)
+        """
+
+        points3d = kornia.geometry.conversions.convert_points_to_homogeneous(points3d).transpose(
+            1, 2
+        )  # (B, 4, N)
+        projected_points = torch.bmm(self.P_raster, points3d.repeat(self.pseudo_batch_size, 1, 1))
+        normalize_by = projected_points[:, -1].view(
+            self.pseudo_batch_size, 1, projected_points.shape[-1]
+        )
+        projected_points /= normalize_by
+        projected_points = projected_points.transpose(-1, -2)  # cannot use view()
+        projected_points = kornia.geometry.convert_points_from_homogeneous(projected_points)
+        if lens_distortion:
+            if self.psi is None:
+                raise RuntimeError("Lens distortion requested, but deactivated in module")
+            projected_points = self.distort_points(projected_points, self.intrinsics_raster)
+        # reshape back from (pseudo_batch_size, N, 2) to (B, T, N, 2)
+        projected_points = projected_points.view(
+            self.batch_dim, self.temporal_dim, projected_points.shape[-2], 2
+        )
+        return projected_points  # (B, T,  N, 2)
+
+    def project_point2ndc_from_P(
+        self, points3d: torch.tensor, lens_distortion: bool
+    ) -> torch.tensor:
+        """Project world coordinates to pixel coordinates from the projection matrix.
+
+        Args:
+            points3d (torch.tensor): of shape (1, N, 3)
+
+        Returns:
+            torch.tensor: projected points of shape (B, T, N, 2)
+        """
+
+        points3d = kornia.geometry.conversions.convert_points_to_homogeneous(points3d).transpose(
+            1, 2
+        )  # (B, 4, N)
+        projected_points = torch.bmm(self.P_ndc, points3d.repeat(self.pseudo_batch_size, 1, 1))
+        normalize_by = projected_points[:, -1].view(
+            self.pseudo_batch_size, 1, projected_points.shape[-1]
+        )
+        projected_points /= normalize_by
+        projected_points = projected_points.transpose(-1, -2)  # cannot use view()
+        projected_points = kornia.geometry.convert_points_from_homogeneous(projected_points)
+        if lens_distortion:
+            if self.psi is None:
+                raise RuntimeError("Lens distortion requested, but deactivated in module")
+            projected_points = self.distort_points(projected_points, self.intrinsics_ndc)
+        # reshape back from (pseudo_batch_size, N, 2) to (B, T, N, 2)
+        projected_points = projected_points.view(
+            self.batch_dim, self.temporal_dim, projected_points.shape[-2], 2
+        )
+        return projected_points  # (B, T,  N, 2)
+
+    def rotation_from_euler_angles(self, pan, tilt, roll):
+        # rotation matrices from a batch of pan tilt roll [rad] vectors of shape (?, )
+
+        mask = (
+            torch.eye(3, requires_grad=False, device=self.device)
+            .reshape((1, 3, 3))
+            .repeat(pan.shape[0], 1, 1)
+        )
+        mask[:, 0, 0] = -torch.sin(pan) * torch.sin(roll) * torch.cos(tilt) + torch.cos(
+            pan
+        ) * torch.cos(roll)
+        mask[:, 0, 1] = torch.sin(pan) * torch.cos(roll) + torch.sin(roll) * torch.cos(
+            pan
+        ) * torch.cos(tilt)
+        mask[:, 0, 2] = torch.sin(roll) * torch.sin(tilt)
+
+        mask[:, 1, 0] = -torch.sin(pan) * torch.cos(roll) * torch.cos(tilt) - torch.sin(
+            roll
+        ) * torch.cos(pan)
+        mask[:, 1, 1] = -torch.sin(pan) * torch.sin(roll) + torch.cos(pan) * torch.cos(
+            roll
+        ) * torch.cos(tilt)
+        mask[:, 1, 2] = torch.sin(tilt) * torch.cos(roll)
+
+        mask[:, 2, 0] = torch.sin(pan) * torch.sin(tilt)
+        mask[:, 2, 1] = -torch.sin(tilt) * torch.cos(pan)
+        mask[:, 2, 2] = torch.cos(tilt)
+
+        return mask
+
+    def get_homography_raster(self):
+        return self.P_raster[:, :, [0, 1, 3]].inverse()
+
+    def get_rays_world(self, x):
+        """_summary_
+
+        Args:
+            x (_type_): x of shape (B, 3, N)
+
+        Returns:
+            LineCollection: _description_
+        """
+        raise NotImplementedError
+        # TODO: verify
+        # ray_cam_trans = torch.bmm(self.rotation.inverse(), torch.bmm(self.intrinsics.inverse(), x))
+        # # unnormalized direction vector in euclidean points (x,y,z) based on camera origin (0,0,0)
+        # ray_cam_trans = torch.nn.functional.normalize(ray_cam_trans, p=2, dim=1)  # (B, 3, N)
+
+        # # shift support vector to origin in world space, i.e. the translation vector
+        # support = self.position.unsqueeze(-1).repeat(
+        #     ray_cam_trans.shape[0], 1, ray_cam_trans.shape[2]
+        # )  # (B, 3, N)
+        # return LineCollection(support=support, direction_norm=ray_cam_trans)
+
+    @staticmethod
+    def get_aov_rad(d: float, fl: torch.tensor):
+        # https://en.wikipedia.org/wiki/Angle_of_view#Calculating_a_camera's_angle_of_view
+        return 2 * torch.arctan(d / (2 * fl))  # in range [0.0, PI]
+
+    @staticmethod
+    def get_fl_from_aov_rad(aov_rad: torch.tensor, d: float):
+        return 0.5 * d * (1 / torch.tan(0.5 * aov_rad))
+
+    def undistort_points(self, points_pixel: torch.tensor, intrinsics, num_iters=5) -> torch.tensor:
+        """Compensate for lens distortion a set of 2D image points.
+
+        Wrapper for kornia.geometry.undistort_points()
+
+        Args:
+            points_pixel (torch.tensor): tensor of shape (B, N, 2)
+
+        Returns:
+            torch.tensor: undistorted points of shape (B, N, 2)
+        """
+        # print(points_pixel.shape, intrinsics.shape, self.lens_dist_coeff.shape)
+        batch_dim, temporal_dim, N, _ = points_pixel.shape
+        points_pixel = points_pixel.view(batch_dim * temporal_dim, N, 2)
+        true_batch_size = batch_dim
+
+        lens_dist_coeff = self.lens_dist_coeff
+        if true_batch_size < self.batch_dim:
+            intrinsics = intrinsics[:true_batch_size]
+            lens_dist_coeff = lens_dist_coeff[:true_batch_size]
+
+        return kornia.geometry.undistort_points(
+            points_pixel, intrinsics, dist=lens_dist_coeff, num_iters=num_iters
+        ).view(batch_dim, temporal_dim, N, 2)
+
+    def distort_points(self, points_pixel: torch.tensor, intrinsics) -> torch.tensor:
+        """Distortion of a set of 2D points based on the lens distortion model.
+
+        Wrapper for kornia.geometry.distort_points()
+
+        Args:
+            points_pixel (torch.tensor): tensor of shape (B, N, 2)
+
+        Returns:
+            torch.tensor: distorted points of shape (B, N, 2)
+        """
+        return kornia.geometry.distort_points(points_pixel, intrinsics, dist=self.lens_dist_coeff)
+
+    def undistort_images(self, images):
+        # images of shape (B, T, C, H, W)
+        true_batch_size, T = images.shape[:2]
+        images = images.view(true_batch_size * T, 3, self.image_height, self.image_width).to(
+            self.device
+        )
+        intrinsics = self.intrinsics_raster
+        lens_dist_coeff = self.lens_dist_coeff
+        if true_batch_size < self.batch_dim:
+            intrinsics = intrinsics[:true_batch_size]
+            lens_dist_coeff = lens_dist_coeff[:true_batch_size]
+
+        return kornia.geometry.calibration.undistort_image(
+            images, intrinsics, lens_dist_coeff
+        ).view(true_batch_size, self.temporal_dim, 3, self.image_height, self.image_width)
+
+    def get_parameters(self, true_batch_size=None):
+        """
+        Get dict of relevant camera parameters and homography matrix
+        :return: The dictionary
+        """
+        out_dict = {
+            "pan_degrees": torch.rad2deg(self.phi_dict["pan"]),
+            "tilt_degrees": torch.rad2deg(self.phi_dict["tilt"]),
+            "roll_degrees": torch.rad2deg(self.phi_dict["roll"]),
+            "position_meters": torch.stack([self.phi_dict[k] for k in ["c_x", "c_y", "c_z"]], dim=1)
+            .squeeze(-1)
+            .unsqueeze(-2)
+            .repeat(1, self.temporal_dim, 1),
+            "aov_radian": self.phi_dict["aov"],
+            "aov_degrees": torch.rad2deg(self.phi_dict["aov"]),
+            "x_focal_length": self.get_fl_from_aov_rad(self.phi_dict["aov"], d=self.image_width),
+            "y_focal_length": self.get_fl_from_aov_rad(self.phi_dict["aov"], d=self.image_width),
+            "principal_point": torch.tensor(
+                [[self.principal_point] * self.temporal_dim] * self.batch_dim
+            ),
+        }
+        out_dict["homography"] = self.get_homography_raster().unsqueeze(1) # (B, 1, 3, 3)
+
+        # expected for SN evaluation
+        out_dict["radial_distortion"] = torch.zeros(self.batch_dim, self.temporal_dim, 6)
+        out_dict["tangential_distortion"] = torch.zeros(self.batch_dim, self.temporal_dim, 2)
+        out_dict["thin_prism_distortion"] = torch.zeros(self.batch_dim, self.temporal_dim, 4)
+
+        if self.psi is not None:
+            # in case only k1 and k2 are provided
+            out_dict["radial_distortion"][..., :2] = self.psi[..., :2]
+
+        if true_batch_size is None or true_batch_size == self.batch_dim:
+            return out_dict
+
+        for k in out_dict.keys():
+            out_dict[k] = out_dict[k][:true_batch_size]
+
+        return out_dict
+
+    @staticmethod
+    def static_undistort_points(points, cam):
+
+        intrinsics = cam.intrinsics_raster
+        lens_dist_coeff = cam.lens_dist_coeff
+
+        true_batch_size = points.shape[0]
+        if true_batch_size < cam.batch_dim:
+            intrinsics = intrinsics[:true_batch_size]
+            lens_dist_coeff = lens_dist_coeff[:true_batch_size]
+        # points in homogenous coordinates
+        # (B, T, 3, S, N) -> (T, 3, S*N) -> (T, S*N, 3)
+        batch_size, T, _, S, N = points.shape
+        points = points.view(batch_size, T, 3, S * N).transpose(2, 3)
+        points[..., :2] = kornia.geometry.undistort_points(
+            points[..., :2].view(batch_size * T, S * N, 2),
+            intrinsics,
+            dist=lens_dist_coeff,
+            num_iters=1,
+        ).view(batch_size, T, S * N, 2)
+
+        # (T, S*N, 3) -> (T, 3, S*N) -> (B, T, 3, S, N)
+        points = points.transpose(2, 3).view(batch_size, T, 3, S, N)
+        return points

tvcalib/data/dataset.py

@@ -0,0 +1,142 @@
+import kornia
+import torch
+import random
+import numpy as np
+from kornia.geometry.transform import resize
+from .utils import split_circle_central
+
+
+
+class InferenceDatasetCalibration(torch.utils.data.Dataset):
+    def __init__(self, keypoints_raw, image_width_source, image_height_source, object3d) -> None:
+        super().__init__()
+        self.keypoints_raw = keypoints_raw
+        self.w = image_width_source
+        self.h = image_height_source
+        self.object3d = object3d
+        self.split_circle_central = True
+
+    def __getitem__(self, idx):
+
+        keypoints_dict = self.keypoints_raw[idx]
+        if self.split_circle_central:
+            keypoints_dict = split_circle_central(keypoints_dict)
+        # add empty entries for non-visible segments
+        for l in self.object3d.segment_names:
+            if l not in keypoints_dict:
+                keypoints_dict[l] = []
+
+        per_sample_output = self.prepare_per_sample(
+            keypoints_dict, self.object3d, 4, 8, self.w, self.h, pad_pixel_position_xy=0.0
+        )
+        for k in per_sample_output.keys():
+            per_sample_output[k] = per_sample_output[k].unsqueeze(0)
+
+        return per_sample_output
+
+    def __len__(self):
+        return len(self.keypoints_raw)
+
+    @staticmethod
+    def prepare_per_sample(
+        keypoints_raw: dict,
+        model3d,
+        num_points_on_line_segments: int,
+        num_points_on_circle_segments: int,
+        image_width_source: int,
+        image_height_source: int,
+        pad_pixel_position_xy=0.0,
+    ):
+        r = {}
+        pixel_stacked = {}
+        for label, points in keypoints_raw.items():
+
+            num_points_selection = num_points_on_line_segments
+            if "Circle" in label:
+                num_points_selection = num_points_on_circle_segments
+
+            # rand select num_points_selection
+            if num_points_selection > len(points):
+                points_sel = points
+            else:
+                # random sample without replacement
+                points_sel = random.sample(points, k=num_points_selection)
+
+            if len(points_sel) > 0:
+                xx = torch.tensor([a["x"] for a in points_sel])
+                yy = torch.tensor([a["y"] for a in points_sel])
+                pixel_stacked[label] = torch.stack([xx, yy], dim=-1)  # (?, 2)
+                # scale pixel annotations from [0, 1] range to source image resolution
+                # as this ranges from [1, {image_height, image_width}] shift pixel one left
+                pixel_stacked[label][:, 0] = pixel_stacked[label][:, 0] * (image_width_source - 1)
+                pixel_stacked[label][:, 1] = pixel_stacked[label][:, 1] * (image_height_source - 1)
+
+        for segment_type, num_segments, segment_names in [
+            ("lines", model3d.line_segments.shape[1], model3d.line_segments_names),
+            ("circles", model3d.circle_segments.shape[1], model3d.circle_segments_names),
+        ]:
+
+            num_points_selection = num_points_on_line_segments
+            if segment_type == "circles":
+                num_points_selection = num_points_on_circle_segments
+            px_projected_selection = (
+                torch.zeros((num_segments, num_points_selection, 2)) + pad_pixel_position_xy
+            )
+            for segment_index, label in enumerate(segment_names):
+                if label in pixel_stacked:
+                    # set annotations to first positions
+                    px_projected_selection[
+                        segment_index, : pixel_stacked[label].shape[0], :
+                    ] = pixel_stacked[label]
+
+            randperm = torch.randperm(num_points_selection)
+            px_projected_selection_shuffled = px_projected_selection.clone()
+            px_projected_selection_shuffled[:, :, 0] = px_projected_selection_shuffled[
+                :, randperm, 0
+            ]
+            px_projected_selection_shuffled[:, :, 1] = px_projected_selection_shuffled[
+                :, randperm, 1
+            ]
+
+            is_keypoint_mask = (
+                (0.0 <= px_projected_selection_shuffled[:, :, 0])
+                & (px_projected_selection_shuffled[:, :, 0] < image_width_source)
+            ) & (
+                (0 < px_projected_selection_shuffled[:, :, 1])
+                & (px_projected_selection_shuffled[:, :, 1] < image_height_source)
+            )
+
+            r[f"{segment_type}__is_keypoint_mask"] = is_keypoint_mask.unsqueeze(0)
+
+            # reshape from (num_segments, num_points_selection, 2) to (3, num_segments, num_points_selection)
+            px_projected_selection_shuffled = (
+                kornia.geometry.conversions.convert_points_to_homogeneous(
+                    px_projected_selection_shuffled
+                )
+            )
+            px_projected_selection_shuffled = px_projected_selection_shuffled.view(
+                num_segments * num_points_selection, 3
+            )
+            px_projected_selection_shuffled = px_projected_selection_shuffled.transpose(0, 1)
+            px_projected_selection_shuffled = px_projected_selection_shuffled.view(
+                3, num_segments, num_points_selection
+            )
+            # (3, num_segments, num_points_selection)
+            r[f"{segment_type}__px_projected_selection_shuffled"] = px_projected_selection_shuffled
+
+            ndc_projected_selection_shuffled = px_projected_selection_shuffled.clone()
+            ndc_projected_selection_shuffled[0] = (
+                ndc_projected_selection_shuffled[0] / image_width_source
+            )
+            ndc_projected_selection_shuffled[1] = (
+                ndc_projected_selection_shuffled[1] / image_height_source
+            )
+            ndc_projected_selection_shuffled[1] = ndc_projected_selection_shuffled[1] * 2.0 - 1
+            ndc_projected_selection_shuffled[0] = ndc_projected_selection_shuffled[0] * 2.0 - 1
+            r[
+                f"{segment_type}__ndc_projected_selection_shuffled"
+            ] = ndc_projected_selection_shuffled
+
+        return r
+
+

tvcalib/data/utils.py

@@ -0,0 +1,166 @@
+
+from operator import itemgetter
+import torch
+import re
+import collections
+
+
+string_classes=str
+
+
+def split_circle_central(keypoints_dict):
+    # split "circle central" in  "circle central left" and "circle central right"
+
+    # assume main camera --> TODO behind the goal camera
+    if "Circle central" in keypoints_dict:
+        points_circle_central_left = []
+        points_circle_central_right = []
+
+        if "Middle line" in keypoints_dict:
+            p_index_ymin, _ = min(
+                enumerate([p["y"] for p in keypoints_dict["Middle line"]]),
+                key=itemgetter(1),
+            )
+            p_index_ymax, _ = max(
+                enumerate([p["y"] for p in keypoints_dict["Middle line"]]),
+                key=itemgetter(1),
+            )
+            p_ymin = keypoints_dict["Middle line"][p_index_ymin]
+            p_ymax = keypoints_dict["Middle line"][p_index_ymax]
+            p_xmean = (p_ymin["x"] + p_ymax["x"]) / 2
+
+            points_circle_central = keypoints_dict["Circle central"]
+            for p in points_circle_central:
+                if p["x"] < p_xmean:
+                    points_circle_central_left.append(p)
+                else:
+                    points_circle_central_right.append(p)
+        else:
+            # circle is partly shown on the left or right side of the image
+            # mean position is shown on the left part of the image --> label right
+            circle_x = [p["x"] for p in keypoints_dict["Circle central"]]
+            mean_x_circle = sum(circle_x) / len(circle_x)
+            if mean_x_circle < 0.5:
+                points_circle_central_right = keypoints_dict["Circle central"]
+            else:
+                points_circle_central_left = keypoints_dict["Circle central"]
+
+        if len(points_circle_central_left) > 0:
+            keypoints_dict["Circle central left"] = points_circle_central_left
+        if len(points_circle_central_right) > 0:
+            keypoints_dict["Circle central right"] = points_circle_central_right
+        if len(points_circle_central_left) == 0 and len(points_circle_central_right) == 0:
+            raise RuntimeError
+        del keypoints_dict["Circle central"]
+    return keypoints_dict
+
+
+def custom_list_collate(batch):
+    r"""
+    Function that takes in a batch of data and puts the elements within the batch
+    into a tensor with an additional outer dimension - batch size. The exact output type can be
+    a :class:`torch.Tensor`, a `Sequence` of :class:`torch.Tensor`, a
+    Collection of :class:`torch.Tensor`, or left unchanged, depending on the input type.
+    This is used as the default function for collation when
+    `batch_size` or `batch_sampler` is defined in :class:`~torch.utils.data.DataLoader`.
+    Here is the general input type (based on the type of the element within the batch) to output type mapping:
+    * :class:`torch.Tensor` -> :class:`torch.Tensor` (with an added outer dimension batch size)
+    * NumPy Arrays -> :class:`torch.Tensor`
+    * `float` -> :class:`torch.Tensor`
+    * `int` -> :class:`torch.Tensor`
+    * `str` -> `str` (unchanged)
+    * `bytes` -> `bytes` (unchanged)
+    * `Mapping[K, V_i]` -> `Mapping[K, default_collate([V_1, V_2, ...])]`
+    * `NamedTuple[V1_i, V2_i, ...]` -> `NamedTuple[default_collate([V1_1, V1_2, ...]), default_collate([V2_1, V2_2, ...]), ...]`
+    * `Sequence[V1_i, V2_i, ...]` -> `Sequence[default_collate([V1_1, V1_2, ...]), default_collate([V2_1, V2_2, ...]), ...]`
+    Args:
+        batch: a single batch to be collated
+    Examples:
+        >>> # Example with a batch of `int`s:
+        >>> default_collate([0, 1, 2, 3])
+        tensor([0, 1, 2, 3])
+        >>> # Example with a batch of `str`s:
+        >>> default_collate(['a', 'b', 'c'])
+        ['a', 'b', 'c']
+        >>> # Example with `Map` inside the batch:
+        >>> default_collate([{'A': 0, 'B': 1}, {'A': 100, 'B': 100}])
+        {'A': tensor([  0, 100]), 'B': tensor([  1, 100])}
+        >>> # Example with `NamedTuple` inside the batch:
+        >>> Point = namedtuple('Point', ['x', 'y'])
+        >>> default_collate([Point(0, 0), Point(1, 1)])
+        Point(x=tensor([0, 1]), y=tensor([0, 1]))
+        >>> # Example with `Tuple` inside the batch:
+        >>> default_collate([(0, 1), (2, 3)])
+        [tensor([0, 2]), tensor([1, 3])]
+
+        >>> # modification
+        >>> # Example with `List` inside the batch:
+        >>> default_collate([[0, 1, 2], [2, 3, 4]])
+        >>> [[0, 1, 2], [2, 3, 4]]
+        >>> # original behavior
+        >>> [[0, 2], [1, 3], [2, 4]]
+    """
+
+    np_str_obj_array_pattern = re.compile(r"[SaUO]")
+    default_collate_err_msg_format = "default_collate: batch must contain tensors, numpy arrays, numbers, dicts or lists; found {}"
+
+    elem = batch[0]
+    elem_type = type(elem)
+    if isinstance(elem, torch.Tensor):
+        out = None
+        if torch.utils.data.get_worker_info() is not None:
+            # If we're in a background process, concatenate directly into a
+            # shared memory tensor to avoid an extra copy
+            numel = sum(x.numel() for x in batch)
+            storage = elem.storage()._new_shared(numel)
+            out = elem.new(storage).resize_(len(batch), *list(elem.size()))
+        return torch.stack(batch, 0, out=out)
+    elif (
+        elem_type.__module__ == "numpy"
+        and elem_type.__name__ != "str_"
+        and elem_type.__name__ != "string_"
+    ):
+        if elem_type.__name__ == "ndarray" or elem_type.__name__ == "memmap":
+            # array of string classes and object
+            if np_str_obj_array_pattern.search(elem.dtype.str) is not None:
+                raise TypeError(default_collate_err_msg_format.format(elem.dtype))
+
+            return [torch.as_tensor(b) for b in batch]
+        elif elem.shape == ():  # scalars
+            return torch.as_tensor(batch)
+    elif isinstance(elem, float):
+        return torch.tensor(batch, dtype=torch.float64)
+    elif isinstance(elem, int):
+        return torch.tensor(batch)
+    elif isinstance(elem, string_classes):
+        return batch
+    elif isinstance(elem, collections.abc.Mapping):
+        try:
+            return elem_type({key: custom_list_collate([d[key] for d in batch]) for key in elem})
+        except TypeError:
+            # The mapping type may not support `__init__(iterable)`.
+            return {key: custom_list_collate([d[key] for d in batch]) for key in elem}
+    elif isinstance(elem, tuple) and hasattr(elem, "_fields"):  # namedtuple
+        return elem_type(*(custom_list_collate(samples) for samples in zip(*batch)))
+    elif isinstance(elem, collections.abc.Sequence):
+        # check to make sure that the elements in batch have consistent size
+        it = iter(batch)
+        elem_size = len(next(it))
+        if not all(len(elem) == elem_size for elem in it):
+            raise RuntimeError("each element in list of batch should be of equal size")
+        # transposed = list(zip(*batch))  # It may be accessed twice, so we use a list.
+
+        return batch
+
+        # if isinstance(elem, tuple):
+        #     return [
+        #         custom_list_collate(samples) for samples in transposed
+        #     ]  # Backwards compatibility.
+        # else:
+        #     try:
+        #         return elem_type([custom_list_collate(samples) for samples in transposed])
+        #     except TypeError:
+        #         # The sequence type may not support `__init__(iterable)` (e.g., `range`).
+        #         return [custom_list_collate(samples) for samples in transposed]
+
+    raise TypeError(default_collate_err_msg_format.format(elem_type))

tvcalib/infer/module.py

@@ -0,0 +1,518 @@
+import torch
+import torch.nn as nn
+from functools import partial
+
+from pathlib import Path
+from typing import Any, Dict, Tuple
+# Imports depuis le package common (supposé être au même niveau que tvcalib)
+from common.infer.base import *
+# from common.registry import Registry  # Toujours commenté car source inconnue
+# from common.utils import to_cuda      # Toujours commenté car source inconnue
+# import project as p                  # Supprimé car probablement lié au projet complet
+
+import torchvision.transforms as T
+# Imports relatifs à l'intérieur de tvcalib (restent relatifs)
+from ..sn_segmentation.src.custom_extremities import (
+    generate_class_synthesis, get_line_extremities
+)
+from ..models.segmentation import InferenceSegmentationModel
+from ..data.dataset import InferenceDatasetCalibration
+from ..data.utils import custom_list_collate
+from ..cam_modules import CameraParameterWLensDistDictZScore, SNProjectiveCamera
+from ..utils.linalg import distance_line_pointcloud_3d, distance_point_pointcloud
+from ..utils.objects_3d import SoccerPitchLineCircleSegments, SoccerPitchSNCircleCentralSplit
+from ..cam_distr.tv_main_center import get_cam_distr, get_dist_distr
+from ..utils.io import detach_dict, tensor2list
+# Import depuis le package common
+from common.data.utils import yards
+
+from kornia.geometry.conversions import convert_points_to_homogeneous
+from tqdm.auto import tqdm
+
+# Commenté car lié à la méthode 'robust' et peut introduire des dépendances
+# from methods.robust.loggers.preview import RobustPreviewLogger 
+
+import numpy as np
+
+
+
+class TvCalibInferModule(InferModule):
+    def __init__(
+        self,
+        segmentation_checkpoint: Path,
+        image_shape=(720,1280),
+        optim_steps=2000,
+        lens_dist: bool=False,
+        playfield_size=(105, 68),
+        make_images: bool=False
+
+    ):
+        self.image_shape = image_shape
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 
+        self.make_images = make_images
+        
+        # We use the logger to draw visualizations
+        # Commenté car la classe RobustPreviewLogger est commentée
+        # self.previewer = RobustPreviewLogger(
+        #     None, num_images=1
+        # )
+        
+        self.fn_generate_class_synthesis = partial(
+            generate_class_synthesis, 
+            radius=4
+        )
+        self.fn_get_line_extremities = partial(
+            get_line_extremities, 
+            maxdist=30, 
+            width=455, 
+            height=256, 
+            num_points_lines=4, 
+            num_points_circles=8
+        )
+        
+        # Segmentation model
+        self.model_seg = InferenceSegmentationModel(
+            segmentation_checkpoint,
+            self.device
+        )
+
+        self.object3d = SoccerPitchLineCircleSegments(
+            device=self.device, 
+            base_field=SoccerPitchSNCircleCentralSplit()
+        )
+        self.object3dcpu = SoccerPitchLineCircleSegments(
+            device="cpu", 
+            base_field=SoccerPitchSNCircleCentralSplit()
+        )
+
+        # Calibration module 
+        batch_size_calib = 1       
+        self.model_calib = TVCalibModule(
+            self.object3d,
+            get_cam_distr(1.96, batch_size_calib, 1),
+            get_dist_distr(batch_size_calib, 1) if lens_dist else None,
+            (image_shape[0], image_shape[1]),
+            optim_steps,
+            self.device,
+            log_per_step=False,
+            tqdm_kwqargs=None,
+        )
+        self.resize = T.Compose([
+            T.Resize(size=(256,455))
+        ])
+        self.offset = np.array([
+            [1, 0, playfield_size[0]/2.0 ],
+            [0, 1, playfield_size[1]/2.0 ],
+            [0, 0, 1]
+        ])
+
+    
+    
+    def setup(self, datamodule: InferDataModule):
+        pass
+    
+    
+    def predict(self, x: Any) -> Dict:
+        
+        """
+            1. Run segmentation & Pick keypoints
+            2. Calibrate based on selected points
+        """
+        
+        # Segment
+        image = x["image"]        
+        keypoints = self._segment(x["image"])
+        
+        # Calibrate
+        homo = self._calibrate(keypoints)
+        
+        # Rescale to 720p
+        image_720p = self.previewer.to_image(image.clone().detach().cpu())
+        
+         # Draw predicted playing field
+        if (homo is not None):            
+            # to yards
+            # Commenté car previewer est commenté
+            # to_yards = np.array([
+            #     [ yards(1.0), 0, 0 ],
+            #     [ 0, yards(1.0), 0 ],
+            #     [ 0, 0, 1]
+            # ])
+            #homo = to_yards @ homo
+                        
+            # Commenté car previewer est commenté
+            # try:
+            #     inv_homo = np.linalg.inv(homo) @ self.previewer.scale
+            #     image_720p = self.previewer.draw_playfield(
+            #         image_720p, 
+            #         self.previewer.image_playfield, 
+            #         inv_homo,
+            #         color=(255,0,0), alpha=1.0,
+            #         flip=False
+            #     )
+            # except:
+            #     # Homography might
+            #     pass
+            pass # Placeholder si l'homographie existe mais previewer est commenté
+        
+        result = {
+            "homography": homo
+        }
+        
+        if (self.make_images):
+            # result["image_720p"] = image_720p # Commenté car image_720p n'est pas modifié sans previewer
+            pass # Placeholder si make_images est True
+        
+        return result
+    
+    
+    def _segment(self, image):
+        
+        # Image -> <1;3;256;455>
+        image = self.resize(image)                
+        with torch.no_grad():
+            sem_lines = self.model_seg.inference(
+                image.unsqueeze(0).to(self.device)
+            )
+        # <B;256;455>
+        sem_lines = sem_lines.detach().cpu().numpy().astype(np.uint8)
+            
+        # Point selection
+        skeletons_batch = self.fn_generate_class_synthesis(sem_lines[0])
+        keypoints_raw_batch = self.fn_get_line_extremities(skeletons_batch)
+        
+        # Return the keypoints
+        return keypoints_raw_batch
+    
+    
+    def _calibrate(self, keypoints):
+        
+        # Just wrap around the keypoints
+        ds = InferenceDatasetCalibration(
+            [keypoints],
+            self.image_shape[1], self.image_shape[0],
+            self.object3d
+        )
+        
+        # Get the first item and optimize it
+        _batch_size = 1
+        x_dict = custom_list_collate([ds[0]])
+        try:
+            # La gestion de previous_params est faite dans self_optim_batch
+            per_sample_loss, cam, _ = self.model_calib.self_optim_batch(x_dict)
+            output_dict = tensor2list(
+                detach_dict({**cam.get_parameters(_batch_size), **per_sample_loss})
+            )
+            
+            homo = output_dict["homography"][0]
+            if (len(homo) > 0):
+                homo = np.array(homo[0])
+
+                to_yards = np.array([
+                    [ yards(1), 0, 0 ],
+                    [ 0, yards(1), 0 ],
+                    [ 0, 0, 1]
+                ])
+
+                # Shift the homography by half the playing field
+                homo = to_yards @ self.offset @ homo
+                            
+            else:
+                homo = None
+        except Exception as e:
+            print(f"Erreur lors de la calibration: {str(e)}")
+            homo = None
+            
+        return homo
+    
+    
+    
+    
+    
+    
+    
+class TVCalibModule(torch.nn.Module):
+    def __init__(
+        self,
+        model3d,
+        cam_distr,
+        dist_distr,
+        image_dim: Tuple[int, int],
+        optim_steps: int,
+        device="cpu",
+        tqdm_kwqargs=None,
+        log_per_step=False,
+        *args,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.image_height, self.image_width = image_dim
+        self.principal_point = (self.image_width / 2, self.image_height / 2)
+        self.model3d = model3d
+        self.cam_param_dict = CameraParameterWLensDistDictZScore(
+            cam_distr, dist_distr, device=device
+        )
+
+        self.lens_distortion_active = False if dist_distr is None else True
+        self.optim_steps = optim_steps
+        self._device = device
+
+        # Ajouter l'attribut pour stocker les paramètres précédents
+        self.previous_params = None
+
+        self.optim = torch.optim.AdamW(
+            self.cam_param_dict.param_dict.parameters(), lr=0.1, weight_decay=0.01
+        )
+        self.Scheduler = partial(
+            torch.optim.lr_scheduler.OneCycleLR,
+            max_lr=0.05,
+            total_steps=self.optim_steps,
+            pct_start=0.5,
+        )
+
+        if self.lens_distortion_active:
+            self.optim_lens_distortion = torch.optim.AdamW(
+                self.cam_param_dict.param_dict_dist.parameters(), lr=1e-3, weight_decay=0.01
+            )
+            self.Scheduler_lens_distortion = partial(
+                torch.optim.lr_scheduler.OneCycleLR,
+                max_lr=1e-3,
+                total_steps=self.optim_steps,
+                pct_start=0.33,
+                optimizer=self.optim_lens_distortion,
+            )
+
+        self.tqdm_kwqargs = tqdm_kwqargs
+        if tqdm_kwqargs is None:
+            self.tqdm_kwqargs = {}
+
+        self.hparams = {"optim": str(self.optim), "scheduler": str(self.Scheduler)}
+        self.log_per_step = log_per_step
+
+    def forward(self, x):
+
+        # individual camera parameters & distortion parameters
+        phi_hat, psi_hat = self.cam_param_dict()
+
+        cam = SNProjectiveCamera(
+            phi_hat,
+            psi_hat,
+            self.principal_point,
+            self.image_width,
+            self.image_height,
+            device=self._device,
+            nan_check=False,
+        )
+
+        # (batch_size, num_views_per_cam, 3, num_segments, num_points)
+        points_px_lines_true = x["lines__ndc_projected_selection_shuffled"].to(self._device)
+        batch_size, T_l, _, S_l, N_l = points_px_lines_true.shape
+
+        # project circle points
+        points_px_circles_true = x["circles__ndc_projected_selection_shuffled"].to(self._device)
+        _, T_c, _, S_c, N_c = points_px_circles_true.shape
+        assert T_c == T_l
+
+        ####################  line-to-point distance at pixel space ####################
+        # start and end point (in world coordinates) for each line segment
+        points3d_lines_keypoints = self.model3d.line_segments  # (3, S_l, 2) to (S_l * 2, 3)
+        points3d_lines_keypoints = points3d_lines_keypoints.reshape(3, S_l * 2).transpose(0, 1)
+        points_px_lines_keypoints = convert_points_to_homogeneous(
+            cam.project_point2ndc(points3d_lines_keypoints, lens_distortion=False)
+        )  # (batch_size, t_l, S_l*2, 3)
+
+        if batch_size < cam.batch_dim:  # actual batch_size smaller than expected, i.e. last batch
+            points_px_lines_keypoints = points_px_lines_keypoints[:batch_size]
+
+        points_px_lines_keypoints = points_px_lines_keypoints.view(batch_size, T_l, S_l, 2, 3)
+
+        lp1 = points_px_lines_keypoints[..., 0, :].unsqueeze(-2)  # -> (batch_size, T_l, 1, S_l, 3)
+        lp2 = points_px_lines_keypoints[..., 1, :].unsqueeze(-2)  # -> (batch_size, T_l, 1, S_l, 3)
+        # (batch_size, T, 3, S, N) -> (batch_size, T, 3, S*N) -> (batch_size, T, S*N, 3) -> (batch_size, T, S, N, 3)
+        pc = (
+            points_px_lines_true.view(batch_size, T_l, 3, S_l * N_l)
+            .transpose(2, 3)
+            .view(batch_size, T_l, S_l, N_l, 3)
+        )
+
+        if self.lens_distortion_active:
+            # undistort given points
+            pc = pc.view(batch_size, T_l, S_l * N_l, 3)
+            pc = pc.detach().clone()
+            pc[..., :2] = cam.undistort_points(
+                pc[..., :2], cam.intrinsics_ndc, num_iters=1
+            )  # num_iters=1 might be enough for a good approximation
+            pc = pc.view(batch_size, T_l, S_l, N_l, 3)
+
+        distances_px_lines_raw = distance_line_pointcloud_3d(
+            e1=lp2 - lp1, r1=lp1, pc=pc, reduce=None
+        )  # (batch_size, T_l, S_l, N_l)
+        distances_px_lines_raw = distances_px_lines_raw.unsqueeze(-3)
+        # (..., 1, S_l, N_l,), i.e. (batch_size, T, 1, S_l, N_l)
+        ####################  circle-to-point distance at pixel space ####################
+
+        # circle segments are approximated as point clouds of size N_c_star
+        points3d_circles_pc = self.model3d.circle_segments
+        _, S_c, N_c_star = points3d_circles_pc.shape
+        points3d_circles_pc = points3d_circles_pc.reshape(3, S_c * N_c_star).transpose(0, 1)
+        points_px_circles_pc = cam.project_point2ndc(points3d_circles_pc, lens_distortion=False)
+
+        if batch_size < cam.batch_dim:  # actual batch_size smaller than expected, i.e. last batch
+            points_px_circles_pc = points_px_circles_pc[:batch_size]
+
+        if self.lens_distortion_active:
+            # (batch_size, T_c, _, S_c, N_c)
+            points_px_circles_true = points_px_circles_true.view(
+                batch_size, T_c, 3, S_c * N_c
+            ).transpose(2, 3)
+            points_px_circles_true = points_px_circles_true.detach().clone()
+            points_px_circles_true[..., :2] = cam.undistort_points(
+                points_px_circles_true[..., :2], cam.intrinsics_ndc, num_iters=1
+            )
+            points_px_circles_true = points_px_circles_true.transpose(2, 3).view(
+                batch_size, T_c, 3, S_c, N_c
+            )
+
+        distances_px_circles_raw = distance_point_pointcloud(
+            points_px_circles_true, points_px_circles_pc.view(batch_size, T_c, S_c, N_c_star, 2)
+        )
+
+        distances_dict = {
+            "loss_ndc_lines": distances_px_lines_raw,  # (batch_size, T_l, 1, S_l, N_l)
+            "loss_ndc_circles": distances_px_circles_raw,  # (batch_size, T_c, 1, S_c, N_c)
+        }
+        return distances_dict, cam
+
+    def self_optim_batch(self, x, *args, **kwargs):
+
+        scheduler = self.Scheduler(self.optim)  # re-initialize lr scheduler for every batch
+        if self.lens_distortion_active:
+            scheduler_lens_distortion = self.Scheduler_lens_distortion()
+
+        # Initialiser avec les paramètres précédents si disponibles
+        if self.previous_params is not None:
+            print("Utilisation des paramètres précédents pour l'initialisation")
+            update_dict = {}
+            for k, v in self.previous_params.items():
+                update_dict[k] = v.detach().clone()
+            self.cam_param_dict.initialize(update_dict)
+        else:
+            print("Première frame : initialisation à zéro")
+            self.cam_param_dict.initialize(None)
+
+        self.optim.zero_grad()
+        if self.lens_distortion_active:
+            self.optim_lens_distortion.zero_grad()
+
+        keypoint_masks = {
+            "loss_ndc_lines": x["lines__is_keypoint_mask"].to(self._device),
+            "loss_ndc_circles": x["circles__is_keypoint_mask"].to(self._device),
+        }
+        num_actual_points = {
+            "loss_ndc_circles": keypoint_masks["loss_ndc_circles"].sum(dim=(-1, -2)),
+            "loss_ndc_lines": keypoint_masks["loss_ndc_lines"].sum(dim=(-1, -2)),
+        }
+
+        per_sample_loss = {}
+        per_sample_loss["mask_lines"] = keypoint_masks["loss_ndc_lines"]
+        per_sample_loss["mask_circles"] = keypoint_masks["loss_ndc_circles"]
+
+        per_step_info = {"loss": [], "lr": []}
+
+        # Paramètres pour les critères d'arrêt
+        loss_target = 0.001  # Réduit pour une meilleure précision potentielle
+        loss_patience = 10  # Nombre d'itérations pour vérifier la stagnation
+        loss_tolerance = 1e-4  # Tolérance pour la variation relative de loss
+        loss_history = []  # Historique des valeurs de loss
+        best_loss = float('inf')  # Meilleure loss obtenue
+        steps_without_improvement = 0  # Compteur d'itérations sans amélioration
+
+        # with torch.autograd.detect_anomaly():
+        with tqdm(range(self.optim_steps), **self.tqdm_kwqargs) as pbar:
+            for step in pbar:
+                self.optim.zero_grad()
+                if self.lens_distortion_active:
+                    self.optim_lens_distortion.zero_grad()
+
+                # forward pass
+                distances_dict, cam = self(x)
+
+                # distance calculate with masked input and output
+                losses = {}
+                for key_dist, distances in distances_dict.items():
+                    distances[~keypoint_masks[key_dist]] = 0.0
+                    per_sample_loss[f"{key_dist}_distances_raw"] = distances
+                    distances_reduced = distances.sum(dim=(-1, -2))
+                    distances_reduced = distances_reduced / num_actual_points[key_dist]
+                    distances_reduced[num_actual_points[key_dist] == 0] = 0.0
+                    distances_reduced = distances_reduced.squeeze(-1)
+                    per_sample_loss[key_dist] = distances_reduced
+                    loss = distances_reduced.mean(dim=-1)
+                    loss = loss.sum()
+                    losses[key_dist] = loss
+
+                loss_total_dist = losses["loss_ndc_lines"] + losses["loss_ndc_circles"]
+                loss_total = loss_total_dist
+                current_loss = loss_total.item()
+
+                # Mettre à jour l'historique des loss
+                loss_history.append(current_loss)
+
+                # Vérifier si on a une meilleure loss
+                if current_loss < best_loss:
+                    best_loss = current_loss
+                    steps_without_improvement = 0
+                else:
+                    steps_without_improvement += 1
+
+                # Critères d'arrêt (commentés pour forcer le nombre total d'étapes)
+                # if len(loss_history) >= loss_patience:
+                #     # Calculer la variation relative moyenne sur les dernières itérations
+                #     recent_losses = loss_history[-loss_patience:]
+                #     # Gérer le cas où toutes les pertes récentes sont nulles ou proches de zéro
+                #     max_recent_loss = max(max(recent_losses), 1e-9) # Evite division par zéro
+                #     loss_variation = abs(max(recent_losses) - min(recent_losses)) / max_recent_loss
+                #     
+                #     # Conditions d'arrêt
+                #     if (current_loss <= loss_target or  # On a atteint la valeur cible
+                #         loss_variation < loss_tolerance or  # La loss ne varie plus significativement
+                #         steps_without_improvement >= loss_patience):  # Pas d'amélioration depuis un moment
+                #         print(f"\nArrêt anticipé à l'itération {step+1}:")
+                #         print(f"Loss finale: {current_loss:.5f}")
+                #         print(f"Meilleure loss: {best_loss:.5f}")
+                #         print(f"Variation relative: {loss_variation:.6f}")
+                #         break
+
+                if self.log_per_step:
+                    per_step_info["lr"].append(scheduler.get_last_lr())
+                    per_step_info["loss"].append(distances_reduced)
+                if step % 50 == 0:
+                    pbar.set_postfix(
+                        loss=f"{loss_total_dist.detach().cpu().tolist():.5f}",
+                        loss_lines=f'{losses["loss_ndc_lines"].detach().cpu().tolist():.3f}',
+                        loss_circles=f'{losses["loss_ndc_circles"].detach().cpu().tolist():.3f}',
+                    )
+
+                loss_total.backward()
+                self.optim.step()
+                scheduler.step()
+                if self.lens_distortion_active:
+                    self.optim_lens_distortion.step()
+                    scheduler_lens_distortion.step()
+
+        # Sauvegarder les paramètres optimisés pour la prochaine frame
+        self.previous_params = {}
+        for k, v in self.cam_param_dict.param_dict.items():
+            self.previous_params[k] = v.detach().clone()
+
+        per_sample_loss["loss_ndc_total"] = torch.sum(
+            torch.stack([per_sample_loss[key_dist] for key_dist in distances_dict.keys()], dim=0),
+            dim=0,
+        )
+
+        if self.log_per_step:
+            per_step_info["loss"] = torch.stack(
+                per_step_info["loss"], dim=-1
+            )
+            per_step_info["lr"] = torch.tensor(per_step_info["lr"])
+        return per_sample_loss, cam, per_step_info

tvcalib/models/segmentation.py

@@ -0,0 +1,22 @@
+from typing import Union
+from pathlib import Path
+import torch
+
+from torchvision.models.segmentation import deeplabv3_resnet101
+from SoccerNet.Evaluation.utils_calibration import SoccerPitch
+
+
+
+class InferenceSegmentationModel:
+    def __init__(self, checkpoint: Union[str, Path], device) -> None:
+        self.device = device
+        self.model = deeplabv3_resnet101(
+            num_classes=len(SoccerPitch.lines_classes) + 1, aux_loss=True
+        )
+        checkpoint_data = torch.load(checkpoint, map_location=self.device, weights_only=False)
+        self.model.load_state_dict(checkpoint_data["model"], strict=False)
+        self.model.to(self.device)
+        self.model.eval()
+
+    def inference(self, img_batch):
+        return self.model(img_batch)["out"].argmax(1)

tvcalib/sn_segmentation/src/baseline_extremities.py

@@ -0,0 +1,311 @@
+import argparse
+import copy
+import json
+import os.path
+import random
+from collections import deque
+from pathlib import Path
+
+import cv2 as cv
+import numpy as np
+import torch
+import torch.backends.cudnn
+import torch.nn as nn
+from PIL import Image
+from torchvision.models.segmentation import deeplabv3_resnet50
+from tqdm import tqdm
+
+from SoccerNet.Evaluation.utils_calibration import SoccerPitch
+
+
+def generate_class_synthesis(semantic_mask, radius):
+    """
+    This function selects for each class present in the semantic mask, a set of circles that cover most of the semantic
+    class blobs.
+    :param semantic_mask: a image containing the segmentation predictions
+    :param radius: circle radius
+    :return: a dictionary which associates with each class detected a list of points ( the circles centers)
+    """
+    buckets = dict()
+    kernel = np.ones((5, 5), np.uint8)
+    semantic_mask = cv.erode(semantic_mask, kernel, iterations=1)
+    for k, class_name in enumerate(SoccerPitch.lines_classes):
+        mask = semantic_mask == k + 1
+        if mask.sum() > 0:
+            disk_list = synthesize_mask(mask, radius)
+            if len(disk_list):
+                buckets[class_name] = disk_list
+
+    return buckets
+
+
+def join_points(point_list, maxdist):
+    """
+    Given a list of points that were extracted from the blobs belonging to a same semantic class, this function creates
+    polylines by linking close points together if their distance is below the maxdist threshold.
+    :param point_list: List of points of the same line class
+    :param maxdist: minimal distance between two polylines.
+    :return: a list of polylines
+    """
+    polylines = []
+
+    if not len(point_list):
+        return polylines
+    head = point_list[0]
+    tail = point_list[0]
+    polyline = deque()
+    polyline.append(point_list[0])
+    remaining_points = copy.deepcopy(point_list[1:])
+
+    while len(remaining_points) > 0:
+        min_dist_tail = 1000
+        min_dist_head = 1000
+        best_head = -1
+        best_tail = -1
+        for j, point in enumerate(remaining_points):
+            dist_tail = np.sqrt(np.sum(np.square(point - tail)))
+            dist_head = np.sqrt(np.sum(np.square(point - head)))
+            if dist_tail < min_dist_tail:
+                min_dist_tail = dist_tail
+                best_tail = j
+            if dist_head < min_dist_head:
+                min_dist_head = dist_head
+                best_head = j
+
+        if min_dist_head <= min_dist_tail and min_dist_head < maxdist:
+            polyline.appendleft(remaining_points[best_head])
+            head = polyline[0]
+            remaining_points.pop(best_head)
+        elif min_dist_tail < min_dist_head and min_dist_tail < maxdist:
+            polyline.append(remaining_points[best_tail])
+            tail = polyline[-1]
+            remaining_points.pop(best_tail)
+        else:
+            polylines.append(list(polyline.copy()))
+            head = remaining_points[0]
+            tail = remaining_points[0]
+            polyline = deque()
+            polyline.append(head)
+            remaining_points.pop(0)
+    polylines.append(list(polyline))
+    return polylines
+
+
+def get_line_extremities(buckets, maxdist, width, height):
+    """
+    Given the dictionary {lines_class: points}, finds plausible extremities of each line, i.e the extremities
+    of the longest polyline that can be built on the class blobs,  and normalize its coordinates
+    by the image size.
+    :param buckets: The dictionary associating line classes to the set of circle centers that covers best the class
+    prediction blobs in the segmentation mask
+    :param maxdist: the maximal distance between two circle centers belonging to the same blob (heuristic)
+    :param width: image width
+    :param height: image height
+    :return: a dictionary associating to each class its extremities
+    """
+    extremities = dict()
+    for class_name, disks_list in buckets.items():
+        polyline_list = join_points(disks_list, maxdist)
+        max_len = 0
+        longest_polyline = []
+        for polyline in polyline_list:
+            if len(polyline) > max_len:
+                max_len = len(polyline)
+                longest_polyline = polyline
+        extremities[class_name] = [
+            {'x': longest_polyline[0][1] / width, 'y': longest_polyline[0][0] / height},
+            {'x': longest_polyline[-1][1] / width, 'y': longest_polyline[-1][0] / height}
+        ]
+    return extremities
+
+
+def get_support_center(mask, start, disk_radius, min_support=0.1):
+    """
+    Returns the barycenter of the True pixels under the area of the mask delimited by the circle of center start and
+    radius of disk_radius pixels.
+    :param mask: Boolean mask
+    :param start: A point located on a true pixel of the mask
+    :param disk_radius: the radius of the circles
+    :param min_support: proportion of the area under the circle area that should be True in order to get enough support
+    :return: A boolean indicating if there is enough support in the circle area, the barycenter of the True pixels under
+     the circle
+    """
+    x = int(start[0])
+    y = int(start[1])
+    support_pixels = 1
+    result = [x, y]
+    xstart = x - disk_radius
+    if xstart < 0:
+        xstart = 0
+    xend = x + disk_radius
+    if xend > mask.shape[0]:
+        xend = mask.shape[0] - 1
+
+    ystart = y - disk_radius
+    if ystart < 0:
+        ystart = 0
+    yend = y + disk_radius
+    if yend > mask.shape[1]:
+        yend = mask.shape[1] - 1
+
+    for i in range(xstart, xend + 1):
+        for j in range(ystart, yend + 1):
+            dist = np.sqrt(np.square(x - i) + np.square(y - j))
+            if dist < disk_radius and mask[i, j] > 0:
+                support_pixels += 1
+                result[0] += i
+                result[1] += j
+    support = True
+    if support_pixels < min_support * np.square(disk_radius) * np.pi:
+        support = False
+
+    result = np.array(result)
+    result = np.true_divide(result, support_pixels)
+
+    return support, result
+
+
+def synthesize_mask(semantic_mask, disk_radius):
+    """
+    Fits circles on the True pixels of the mask and returns those which have enough support : meaning that the
+    proportion of the area of the circle covering True pixels is higher that a certain threshold in order to avoid
+    fitting circles on alone pixels.
+    :param semantic_mask: boolean mask
+    :param disk_radius: radius of the circles
+    :return: a list of disk centers, that have enough support
+    """
+    mask = semantic_mask.copy().astype(np.uint8)
+    points = np.transpose(np.nonzero(mask))
+    disks = []
+    while len(points):
+
+        start = random.choice(points)
+        dist = 10.
+        success = True
+        while dist > 1.:
+            enough_support, center = get_support_center(mask, start, disk_radius)
+            if not enough_support:
+                bad_point = np.round(center).astype(np.int32)
+                cv.circle(mask, (bad_point[1], bad_point[0]), disk_radius, (0), -1)
+                success = False
+            dist = np.sqrt(np.sum(np.square(center - start)))
+            start = center
+        if success:
+            disks.append(np.round(start).astype(np.int32))
+            cv.circle(mask, (disks[-1][1], disks[-1][0]), disk_radius, 0, -1)
+        points = np.transpose(np.nonzero(mask))
+
+    return disks
+
+
+class SegmentationNetwork:
+    def __init__(self, model_file, mean_file, std_file, num_classes=29, width=640, height=360):
+        file_path = Path(model_file).resolve()
+        model = nn.DataParallel(deeplabv3_resnet50(pretrained=False, num_classes=num_classes))
+        self.init_weight(model, nn.init.kaiming_normal_,
+                         nn.BatchNorm2d, 1e-3, 0.1,
+                         mode='fan_in')
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        checkpoint = torch.load(str(file_path), map_location=self.device)
+        model.load_state_dict(checkpoint["model"])
+        model.eval()
+        self.model = model.to(self.device)
+        file_path = Path(mean_file).resolve()
+        self.mean = np.load(str(file_path))
+        file_path = Path(std_file).resolve()
+        self.std = np.load(str(file_path))
+        self.width = width
+        self.height = height
+
+    def init_weight(self, feature, conv_init, norm_layer, bn_eps, bn_momentum,
+                    **kwargs):
+        for name, m in feature.named_modules():
+            if isinstance(m, (nn.Conv2d, nn.Conv3d)):
+                conv_init(m.weight, **kwargs)
+            elif isinstance(m, norm_layer):
+                m.eps = bn_eps
+                m.momentum = bn_momentum
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def analyse_image(self, image):
+        """
+        Process image and perform inference, returns mask of detected classes
+        :param image: BGR image
+        :return: predicted classes mask
+        """
+        img = cv.resize(image, (self.width, self.height), interpolation=cv.INTER_LINEAR)
+        img = np.asarray(img, np.float32) / 255.
+        img = (img - self.mean) / self.std
+        img = img.transpose((2, 0, 1))
+        img = torch.from_numpy(img).to(self.device).unsqueeze(0)
+
+        cuda_result = self.model.forward(img.float())
+        output = cuda_result['out'].data[0].cpu().numpy()
+        output = output.transpose(1, 2, 0)
+        output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
+
+        return output
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Test')
+
+    parser.add_argument('-s', '--soccernet', default="./annotations/", type=str,
+                        help='Path to the SoccerNet-V3 dataset folder')
+    parser.add_argument('-p', '--prediction', default="./results_bis", required=False, type=str,
+                        help="Path to the prediction folder")
+    parser.add_argument('--split', required=False, type=str, default="test", help='Select the split of data')
+    parser.add_argument('--masks', required=False, type=bool, default=False, help='Save masks in prediction directory')
+    parser.add_argument('--resolution_width', required=False, type=int, default=455,
+                        help='width resolution of the images')
+    parser.add_argument('--resolution_height', required=False, type=int, default=256,
+                        help='height resolution of the images')
+    parser.add_argument('--checkpoint_dir', default="resources")
+    args = parser.parse_args()
+
+    lines_palette = [0, 0, 0]
+    for line_class in SoccerPitch.lines_classes:
+        lines_palette.extend(SoccerPitch.palette[line_class])
+
+    calib_net = SegmentationNetwork(
+        os.path.join(args.checkpoint_dir, "soccer_pitch_segmentation.pth"),
+        os.path.join(args.checkpoint_dir, "mean.npy"),
+        os.path.join(args.checkpoint_dir, "std.npy")
+    )
+
+    dataset_dir = os.path.join(args.soccernet, args.split)
+    if not os.path.exists(dataset_dir):
+        print("Invalid dataset path !")
+        exit(-1)
+
+    frames = [f for f in os.listdir(dataset_dir) if ".jpg" in f]
+    with tqdm(enumerate(frames), total=len(frames), ncols=160) as t:
+        for i, frame in t:
+
+            output_prediction_folder = os.path.join(args.prediction, args.split)
+            if not os.path.exists(output_prediction_folder):
+                os.makedirs(output_prediction_folder)
+            prediction = dict()
+            count = 0
+
+            frame_path = os.path.join(dataset_dir, frame)
+
+            frame_index = frame.split(".")[0]
+
+            image = cv.imread(frame_path)
+            semlines = calib_net.analyse_image(image)
+            if args.masks:
+                mask = Image.fromarray(semlines.astype(np.uint8)).convert('P')
+                mask.putpalette(lines_palette)
+                mask_file = os.path.join(output_prediction_folder, frame)
+                mask.convert("RGB").save(mask_file)
+            skeletons = generate_class_synthesis(semlines, 6)
+            extremities = get_line_extremities(skeletons, 40, args.resolution_width, args.resolution_height)
+
+            prediction = extremities
+            count += 1
+
+            prediction_file = os.path.join(output_prediction_folder, f"extremities_{frame_index}.json")
+            with open(prediction_file, "w") as f:
+                json.dump(prediction, f, indent=4)

tvcalib/sn_segmentation/src/custom_extremities.py

@@ -0,0 +1,322 @@
+import argparse
+import copy
+import itertools
+import json
+import os.path
+import random
+from collections import deque
+from pathlib import Path
+
+from pytorch_lightning import seed_everything
+
+seed_everything(seed=10, workers=True)
+
+import cv2 as cv
+import numpy as np
+import torch
+import torch.backends.cudnn
+import torch.nn as nn
+import torchvision.transforms as T
+
+from PIL import Image
+from torchvision.models.segmentation import deeplabv3_resnet101
+from tqdm import tqdm
+
+from SoccerNet.Evaluation.utils_calibration import SoccerPitch
+
+
+def generate_class_synthesis(semantic_mask, radius):
+    """
+    This function selects for each class present in the semantic mask, a set of circles that cover most of the semantic
+    class blobs.
+    :param semantic_mask: a image containing the segmentation predictions
+    :param radius: circle radius
+    :return: a dictionary which associates with each class detected a list of points ( the circles centers)
+    """
+    buckets = dict()
+    kernel = np.ones((5, 5), np.uint8)
+    semantic_mask = cv.erode(semantic_mask, kernel, iterations=1)
+    for k, class_name in enumerate(SoccerPitch.lines_classes):
+        mask = semantic_mask == k + 1
+        if mask.sum() > 0:
+            disk_list = synthesize_mask(mask, radius)
+            if len(disk_list):
+                buckets[class_name] = disk_list
+
+    return buckets
+
+
+def join_points(point_list, maxdist):
+    """
+    Given a list of points that were extracted from the blobs belonging to a same semantic class, this function creates
+    polylines by linking close points together if their distance is below the maxdist threshold.
+    :param point_list: List of points of the same line class
+    :param maxdist: minimal distance between two polylines.
+    :return: a list of polylines
+    """
+    polylines = []
+
+    if not len(point_list):
+        return polylines
+    head = point_list[0]
+    tail = point_list[0]
+    polyline = deque()
+    polyline.append(point_list[0])
+    remaining_points = copy.deepcopy(point_list[1:])
+
+    while len(remaining_points) > 0:
+        min_dist_tail = 1000
+        min_dist_head = 1000
+        best_head = -1
+        best_tail = -1
+        for j, point in enumerate(remaining_points):
+            dist_tail = np.sqrt(np.sum(np.square(point - tail)))
+            dist_head = np.sqrt(np.sum(np.square(point - head)))
+            if dist_tail < min_dist_tail:
+                min_dist_tail = dist_tail
+                best_tail = j
+            if dist_head < min_dist_head:
+                min_dist_head = dist_head
+                best_head = j
+
+        if min_dist_head <= min_dist_tail and min_dist_head < maxdist:
+            polyline.appendleft(remaining_points[best_head])
+            head = polyline[0]
+            remaining_points.pop(best_head)
+        elif min_dist_tail < min_dist_head and min_dist_tail < maxdist:
+            polyline.append(remaining_points[best_tail])
+            tail = polyline[-1]
+            remaining_points.pop(best_tail)
+        else:
+            polylines.append(list(polyline.copy()))
+            head = remaining_points[0]
+            tail = remaining_points[0]
+            polyline = deque()
+            polyline.append(head)
+            remaining_points.pop(0)
+    polylines.append(list(polyline))
+    return polylines
+
+
+def get_line_extremities(buckets, maxdist, width, height, num_points_lines, num_points_circles):
+    """
+    Given the dictionary {lines_class: points}, finds plausible extremities of each line, i.e the extremities
+    of the longest polyline that can be built on the class blobs,  and normalize its coordinates
+    by the image size.
+    :param buckets: The dictionary associating line classes to the set of circle centers that covers best the class
+    prediction blobs in the segmentation mask
+    :param maxdist: the maximal distance between two circle centers belonging to the same blob (heuristic)
+    :param width: image width
+    :param height: image height
+    :return: a dictionary associating to each class its extremities
+    """
+    extremities = dict()
+    for class_name, disks_list in buckets.items():
+        polyline_list = join_points(disks_list, maxdist)
+        max_len = 0
+        longest_polyline = []
+        for polyline in polyline_list:
+            if len(polyline) > max_len:
+                max_len = len(polyline)
+                longest_polyline = polyline
+        extremities[class_name] = [
+            {'x': longest_polyline[0][1] / width, 'y': longest_polyline[0][0] / height},
+            {'x': longest_polyline[-1][1] / width, 'y': longest_polyline[-1][0] / height}, 
+            
+        ]
+        num_points = num_points_lines
+        if "Circle" in class_name:
+            num_points = num_points_circles
+        if num_points > 2:
+            # equally spaced points along the longest polyline
+            # skip first and last as they already exist
+            for i in range(1, num_points - 1):
+                extremities[class_name].insert(
+                    len(extremities[class_name]) - 1,
+                    {'x': longest_polyline[i * int(len(longest_polyline) / num_points)][1] / width, 'y': longest_polyline[i * int(len(longest_polyline) / num_points)][0] / height}
+                )
+
+    return extremities
+
+
+def get_support_center(mask, start, disk_radius, min_support=0.1):
+    """
+    Returns the barycenter of the True pixels under the area of the mask delimited by the circle of center start and
+    radius of disk_radius pixels.
+    :param mask: Boolean mask
+    :param start: A point located on a true pixel of the mask
+    :param disk_radius: the radius of the circles
+    :param min_support: proportion of the area under the circle area that should be True in order to get enough support
+    :return: A boolean indicating if there is enough support in the circle area, the barycenter of the True pixels under
+     the circle
+    """
+    x = int(start[0])
+    y = int(start[1])
+    support_pixels = 1
+    result = [x, y]
+    xstart = x - disk_radius
+    if xstart < 0:
+        xstart = 0
+    xend = x + disk_radius
+    if xend > mask.shape[0]:
+        xend = mask.shape[0] - 1
+
+    ystart = y - disk_radius
+    if ystart < 0:
+        ystart = 0
+    yend = y + disk_radius
+    if yend > mask.shape[1]:
+        yend = mask.shape[1] - 1
+
+    for i in range(xstart, xend + 1):
+        for j in range(ystart, yend + 1):
+            dist = np.sqrt(np.square(x - i) + np.square(y - j))
+            if dist < disk_radius and mask[i, j] > 0:
+                support_pixels += 1
+                result[0] += i
+                result[1] += j
+    support = True
+    if support_pixels < min_support * np.square(disk_radius) * np.pi:
+        support = False
+
+    result = np.array(result)
+    result = np.true_divide(result, support_pixels)
+
+    return support, result
+
+
+def synthesize_mask(semantic_mask, disk_radius):
+    """
+    Fits circles on the True pixels of the mask and returns those which have enough support : meaning that the
+    proportion of the area of the circle covering True pixels is higher that a certain threshold in order to avoid
+    fitting circles on alone pixels.
+    :param semantic_mask: boolean mask
+    :param disk_radius: radius of the circles
+    :return: a list of disk centers, that have enough support
+    """
+    mask = semantic_mask.copy().astype(np.uint8)
+    points = np.transpose(np.nonzero(mask))
+    disks = []
+    while len(points):
+
+        start = random.choice(points)
+        dist = 10.
+        success = True
+        while dist > 1.:
+            enough_support, center = get_support_center(mask, start, disk_radius)
+            if not enough_support:
+                bad_point = np.round(center).astype(np.int32)
+                cv.circle(mask, (bad_point[1], bad_point[0]), disk_radius, (0), -1)
+                success = False
+            dist = np.sqrt(np.sum(np.square(center - start)))
+            start = center
+        if success:
+            disks.append(np.round(start).astype(np.int32))
+            cv.circle(mask, (disks[-1][1], disks[-1][0]), disk_radius, 0, -1)
+        points = np.transpose(np.nonzero(mask))
+
+    return disks
+
+class CustomNetwork:
+
+    def __init__(self, checkpoint):
+        print("Loading model" + checkpoint)
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model = deeplabv3_resnet101(num_classes=len(SoccerPitch.lines_classes) + 1, aux_loss=True)
+        self.model.load_state_dict(torch.load(checkpoint)["model"], strict=False)
+        self.model.to(self.device)
+        self.model.eval()
+        print("using", self.device)
+
+    def forward(self, img):
+        trf = T.Compose(
+            [
+                T.Resize(256),
+                #T.CenterCrop(224),
+                T.ToTensor(),
+                T.Normalize(
+                    mean = [0.485, 0.456, 0.406],
+                    std = [0.229, 0.224, 0.225]
+                    )
+            ]
+        )
+        img = trf(img).unsqueeze(0).to(self.device) 
+        result = self.model(img)["out"].detach().squeeze(0).argmax(0)
+        result = result.cpu().numpy().astype(np.uint8)
+        #print(result)
+        return result
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Test')
+
+    parser.add_argument('-s', '--soccernet', default="/nfs/data/soccernet/calibration/", type=str,
+                        help='Path to the SoccerNet-V3 dataset folder')
+    parser.add_argument('-p', '--prediction', default="sn-calib-test_endpoints", required=False, type=str,
+                        help="Path to the prediction folder")
+    parser.add_argument('--split', required=False, type=str, default="challenge", help='Select the split of data')
+    parser.add_argument('--masks', required=False, type=bool, default=False, help='Save masks in prediction directory')
+    parser.add_argument('--resolution_width', required=False, type=int, default=455,
+                        help='width resolution of the images')
+    parser.add_argument('--resolution_height', required=False, type=int, default=256,
+                        help='height resolution of the images')
+    parser.add_argument('--checkpoint', required=False, type=str, help="Path to the custom model checkpoint.")
+    parser.add_argument('--pp_radius', required=False, type=int, default=4,
+                        help='Post processing: Radius of circles that cover each segment.')
+    parser.add_argument('--pp_maxdists', required=False, type=int, default=30,
+                        help='Post processing: Maximum distance of circles that are allowed within one segment.')
+    parser.add_argument('--num_points_lines', required=False, type=int, default=2, choices=range(2,10),
+                        help='Post processing: Number of keypoints that represent a line segment')
+    parser.add_argument('--num_points_circles', required=False, type=int, default=2, choices=range(2,10),
+                        help='Post processing: Number of keypoints that represent a circle segment')
+    args = parser.parse_args()
+
+    lines_palette = [0, 0, 0]
+    for line_class in SoccerPitch.lines_classes:
+        lines_palette.extend(SoccerPitch.palette[line_class])
+
+    model = CustomNetwork(args.checkpoint)
+
+    dataset_dir = os.path.join(args.soccernet, args.split)
+    if not os.path.exists(dataset_dir):
+        print("Invalid dataset path !")
+        exit(-1)
+
+    radius = args.pp_radius
+    maxdists = args.pp_maxdists
+    
+    frames = [f for f in os.listdir(dataset_dir) if ".jpg" in f]
+    with tqdm(enumerate(frames), total=len(frames), ncols=160) as t:
+        for i, frame in t:
+
+            output_prediction_folder = os.path.join(str(args.prediction), f"np{args.num_points_lines}_nc{args.num_points_circles}_r{radius}_md{maxdists}", args.split)
+            if not os.path.exists(output_prediction_folder):
+                os.makedirs(output_prediction_folder)
+            prediction = dict()
+            count = 0
+
+            frame_path = os.path.join(dataset_dir, frame)
+
+            frame_index = frame.split(".")[0]
+
+            image = Image.open(frame_path)
+
+            semlines = model.forward(image)
+            #print(semlines.shape)
+            # print("\nsemlines", type(semlines), semlines.shape)
+            if args.masks:
+                mask = Image.fromarray(semlines.astype(np.uint8)).convert('P')
+                mask.putpalette(lines_palette)
+                mask_file = os.path.join(output_prediction_folder, frame)
+                mask.convert("RGB").save(mask_file)
+            skeletons = generate_class_synthesis(semlines, radius)
+
+            extremities = get_line_extremities(skeletons, maxdists, args.resolution_width, args.resolution_height, args.num_points_lines, args.num_points_circles)
+
+
+            prediction = extremities
+            count += 1
+
+            prediction_file = os.path.join(output_prediction_folder, f"extremities_{frame_index}.json")
+            with open(prediction_file, "w") as f:
+                json.dump(prediction, f, indent=4)
+

tvcalib/sn_segmentation/src/dataloader.py

@@ -0,0 +1,122 @@
+"""
+DataLoader used to train the segmentation network used for the prediction of extremities.
+"""
+
+import json
+import os
+import time
+from argparse import ArgumentParser
+
+import cv2 as cv
+import numpy as np
+from torch.utils.data import Dataset
+from tqdm import tqdm
+
+from SoccerNet.Evaluation.utils_calibration import SoccerPitch
+
+
+class SoccerNetDataset(Dataset):
+    def __init__(self,
+                 datasetpath,
+                 split="test",
+                 width=640,
+                 height=360,
+                 mean="../resources/mean.npy",
+                 std="../resources/std.npy"):
+        self.mean = np.load(mean)
+        self.std = np.load(std)
+        self.width = width
+        self.height = height
+
+        dataset_dir = os.path.join(datasetpath, split)
+        if not os.path.exists(dataset_dir):
+            print("Invalid dataset path !")
+            exit(-1)
+
+        frames = [f for f in os.listdir(dataset_dir) if ".jpg" in f]
+
+        self.data = []
+        self.n_samples = 0
+        for frame in frames:
+
+            frame_index = frame.split(".")[0]
+            annotation_file = os.path.join(dataset_dir, f"{frame_index}.json")
+            if not os.path.exists(annotation_file):
+                continue
+            with open(annotation_file, "r") as f:
+                groundtruth_lines = json.load(f)
+            img_path = os.path.join(dataset_dir, frame)
+            if groundtruth_lines:
+                self.data.append({
+                    "image_path": img_path,
+                    "annotations": groundtruth_lines,
+                })
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        item = self.data[index]
+
+        img = cv.imread(item["image_path"])
+        img = cv.resize(img, (self.width, self.height), interpolation=cv.INTER_LINEAR)
+
+        mask = np.zeros(img.shape[:-1], dtype=np.uint8)
+        img = np.asarray(img, np.float32) / 255.
+        img -= self.mean
+        img /= self.std
+        img = img.transpose((2, 0, 1))
+        for class_number, class_ in enumerate(SoccerPitch.lines_classes):
+            if class_ in item["annotations"].keys():
+                key = class_
+                line = item["annotations"][key]
+                prev_point = line[0]
+                for i in range(1, len(line)):
+                    next_point = line[i]
+                    cv.line(mask,
+                            (int(prev_point["x"] * mask.shape[1]), int(prev_point["y"] * mask.shape[0])),
+                            (int(next_point["x"] * mask.shape[1]), int(next_point["y"] * mask.shape[0])),
+                            class_number + 1,
+                            2)
+                    prev_point = next_point
+        return img, mask
+
+
+if __name__ == "__main__":
+
+    # Load the arguments
+    parser = ArgumentParser(description='dataloader')
+
+    parser.add_argument('--SoccerNet_path', default="./annotations/", type=str,
+                        help='Path to the SoccerNet-V3 dataset folder')
+    parser.add_argument('--tiny', required=False, type=int, default=None, help='Select a subset of x games')
+    parser.add_argument('--split', required=False, type=str, default="test", help='Select the split of data')
+    parser.add_argument('--num_workers', required=False, type=int, default=4,
+                        help='number of workers for the dataloader')
+    parser.add_argument('--resolution_width', required=False, type=int, default=1920,
+                        help='width resolution of the images')
+    parser.add_argument('--resolution_height', required=False, type=int, default=1080,
+                        help='height resolution of the images')
+    parser.add_argument('--preload_images', action='store_true',
+                        help="Preload the images when constructing the dataset")
+    parser.add_argument('--zipped_images', action='store_true', help="Read images from zipped folder")
+
+    args = parser.parse_args()
+
+    start_time = time.time()
+    soccernet = SoccerNetDataset(args.SoccerNet_path, split=args.split)
+    with tqdm(enumerate(soccernet), total=len(soccernet), ncols=160) as t:
+        for i, data in t:
+            img = soccernet[i][0].astype(np.uint8).transpose((1, 2, 0))
+            print(img.shape)
+            print(img.dtype)
+            cv.imshow("Normalized image", img)
+            cv.waitKey(0)
+            cv.destroyAllWindows()
+            print(data[1].shape)
+            cv.imshow("Mask", soccernet[i][1].astype(np.uint8))
+            cv.waitKey(0)
+            cv.destroyAllWindows()
+            continue
+    end_time = time.time()
+    print(end_time - start_time)

tvcalib/sn_segmentation/src/evaluate_extremities.py

@@ -0,0 +1,270 @@
+import argparse
+import json
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+from tqdm import tqdm
+
+from SoccerNet.Evaluation.utils_calibration import SoccerPitch
+
+
+def distance(point1, point2):
+    """
+    Computes euclidian distance between 2D points
+    :param point1
+    :param point2
+    :return: euclidian distance between point1 and point2
+    """
+    diff = np.array([point1['x'], point1['y']]) - np.array([point2['x'], point2['y']])
+    sq_dist = np.square(diff)
+    return np.sqrt(sq_dist.sum())
+
+
+def mirror_labels(lines_dict):
+    """
+    Replace each line class key of the dictionary with its opposite element according to a central projection by the
+    soccer pitch center
+    :param lines_dict: dictionary whose keys will be mirrored
+    :return: Dictionary with mirrored keys and same values
+    """
+    mirrored_dict = dict()
+    for line_class, value in lines_dict.items():
+        mirrored_dict[SoccerPitch.symetric_classes[line_class]] = value
+    return mirrored_dict
+
+
+def evaluate_detection_prediction(detected_lines, groundtruth_lines, threshold=2.):
+    """
+    Evaluates the prediction of extremities. The extremities associated to a class are unordered. The extremities of the
+    "Circle central" element is not well-defined for this task, thus this class is ignored.
+    Computes confusion matrices for a level of precision specified by the threshold.
+    A groundtruth extremity point is correctly classified if it lies at less than threshold pixels from the
+    corresponding extremity point of the prediction of the same class.
+    Computes also the euclidian distance between each predicted extremity and its closest groundtruth extremity, when
+    both the groundtruth and the prediction contain the element class.
+
+    :param detected_lines: dictionary of detected lines classes as keys and associated predicted extremities as values
+    :param groundtruth_lines: dictionary of annotated lines classes as keys and associated annotated points as values
+    :param threshold: distance in pixels that distinguishes good matches from bad ones
+    :return: confusion matrix, per class confusion matrix & per class localization errors
+    """
+    confusion_mat = np.zeros((2, 2), dtype=np.float32)
+    per_class_confusion = {}
+    errors_dict = {}
+    detected_classes = set(detected_lines.keys())
+    groundtruth_classes = set(groundtruth_lines.keys())
+
+    if "Circle central" in groundtruth_classes:
+        groundtruth_classes.remove("Circle central")
+    if "Circle central" in detected_classes:
+        detected_classes.remove("Circle central")
+
+    false_positives_classes = detected_classes - groundtruth_classes
+    for false_positive_class in false_positives_classes:
+        false_positives = len(detected_lines[false_positive_class])
+        confusion_mat[0, 1] += false_positives
+        per_class_confusion[false_positive_class] = np.array([[0., false_positives], [0., 0.]])
+
+    false_negatives_classes = groundtruth_classes - detected_classes
+    for false_negatives_class in false_negatives_classes:
+        false_negatives = len(groundtruth_lines[false_negatives_class])
+        confusion_mat[1, 0] += false_negatives
+        per_class_confusion[false_negatives_class] = np.array([[0., 0.], [false_negatives, 0.]])
+
+    common_classes = detected_classes - false_positives_classes
+
+    for detected_class in common_classes:
+
+        detected_points = detected_lines[detected_class]
+
+        groundtruth_points = groundtruth_lines[detected_class]
+
+        groundtruth_extremities = [groundtruth_points[0], groundtruth_points[-1]]
+        predicted_extremities = [detected_points[0], detected_points[-1]]
+        per_class_confusion[detected_class] = np.zeros((2, 2))
+
+        dist1 = distance(groundtruth_extremities[0], predicted_extremities[0])
+        dist1rev = distance(groundtruth_extremities[1], predicted_extremities[0])
+
+        dist2 = distance(groundtruth_extremities[1], predicted_extremities[1])
+        dist2rev = distance(groundtruth_extremities[0], predicted_extremities[1])
+        if dist1rev <= dist1 and dist2rev <= dist2:
+            # reverse order
+            dist1 = dist1rev
+            dist2 = dist2rev
+
+        errors_dict[detected_class] = [dist1, dist2]
+
+        if dist1 < threshold:
+            confusion_mat[0, 0] += 1
+            per_class_confusion[detected_class][0, 0] += 1
+        else:
+            # treat too far detections as false positives
+            confusion_mat[0, 1] += 1
+            per_class_confusion[detected_class][0, 1] += 1
+
+        if dist2 < threshold:
+            confusion_mat[0, 0] += 1
+            per_class_confusion[detected_class][0, 0] += 1
+
+        else:
+            # treat too far detections as false positives
+            confusion_mat[0, 1] += 1
+            per_class_confusion[detected_class][0, 1] += 1
+
+    return confusion_mat, per_class_confusion, errors_dict
+
+
+def scale_points(points_dict, s_width, s_height):
+    """
+    Scale points by s_width and s_height factors
+    :param points_dict: dictionary of annotations/predictions with normalized point values
+    :param s_width: width scaling factor
+    :param s_height: height scaling factor
+    :return: dictionary with scaled points
+    """
+    line_dict = {}
+    for line_class, points in points_dict.items():
+        scaled_points = []
+        for point in points:
+            new_point = {'x': point['x'] * (s_width-1), 'y': point['y'] * (s_height-1)}
+            scaled_points.append(new_point)
+        if len(scaled_points):
+            line_dict[line_class] = scaled_points
+    return line_dict
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(description='Test')
+
+    parser.add_argument('-s', '--soccernet', default="./annotations", type=str,
+                        help='Path to the SoccerNet-V3 dataset folder')
+    parser.add_argument('-p', '--prediction', default="./results_bis",
+                        required=False, type=str,
+                        help="Path to the prediction folder")
+    parser.add_argument('-t', '--threshold', default=10, required=False, type=int,
+                        help="Accuracy threshold in pixels")
+    parser.add_argument('--split', required=False, type=str, default="test", help='Select the split of data')
+    parser.add_argument('--resolution_width', required=False, type=int, default=960,
+                        help='width resolution of the images')
+    parser.add_argument('--resolution_height', required=False, type=int, default=540,
+                        help='height resolution of the images')
+    args = parser.parse_args()
+
+    accuracies = []
+    precisions = []
+    recalls = []
+    dict_errors = {}
+    per_class_confusion_dict = {}
+
+    dataset_dir = os.path.join(args.soccernet, args.split)
+    if not os.path.exists(dataset_dir):
+        print("Invalid dataset path !")
+        exit(-1)
+
+    annotation_files = [f for f in os.listdir(dataset_dir) if ".json" in f]
+
+    with tqdm(enumerate(annotation_files), total=len(annotation_files), ncols=160) as t:
+        for i, annotation_file in t:
+            frame_index = annotation_file.split(".")[0]
+            annotation_file = os.path.join(args.soccernet, args.split, annotation_file)
+            prediction_file = os.path.join(args.prediction, args.split, f"extremities_{frame_index}.json")
+
+            if not os.path.exists(prediction_file):
+                accuracies.append(0.)
+                precisions.append(0.)
+                recalls.append(0.)
+                continue
+
+            with open(annotation_file, 'r') as f:
+                line_annotations = json.load(f)
+
+            with open(prediction_file, 'r') as f:
+                predictions = json.load(f)
+
+            predictions = scale_points(predictions, args.resolution_width, args.resolution_height)
+            line_annotations = scale_points(line_annotations, args.resolution_width, args.resolution_height)
+
+            img_prediction = predictions
+            img_groundtruth = line_annotations
+            confusion1, per_class_conf1, reproj_errors1 = evaluate_detection_prediction(img_prediction,
+                                                                                        img_groundtruth,
+                                                                                        args.threshold)
+            confusion2, per_class_conf2, reproj_errors2 = evaluate_detection_prediction(img_prediction,
+                                                                                        mirror_labels(
+                                                                                            img_groundtruth),
+                                                                                        args.threshold)
+
+            accuracy1, accuracy2 = 0., 0.
+            if confusion1.sum() > 0:
+                accuracy1 = confusion1[0, 0] / confusion1.sum()
+
+            if confusion2.sum() > 0:
+                accuracy2 = confusion2[0, 0] / confusion2.sum()
+
+            if accuracy1 > accuracy2:
+                accuracy = accuracy1
+                confusion = confusion1
+                per_class_conf = per_class_conf1
+                reproj_errors = reproj_errors1
+            else:
+                accuracy = accuracy2
+                confusion = confusion2
+                per_class_conf = per_class_conf2
+                reproj_errors = reproj_errors2
+
+            accuracies.append(accuracy)
+            if confusion[0, :].sum() > 0:
+                precision = confusion[0, 0] / (confusion[0, :].sum())
+                precisions.append(precision)
+            if (confusion[0, 0] + confusion[1, 0]) > 0:
+                recall = confusion[0, 0] / (confusion[0, 0] + confusion[1, 0])
+                recalls.append(recall)
+
+            for line_class, errors in reproj_errors.items():
+                if line_class in dict_errors.keys():
+                    dict_errors[line_class].extend(errors)
+                else:
+                    dict_errors[line_class] = errors
+
+            for line_class, confusion_mat in per_class_conf.items():
+                if line_class in per_class_confusion_dict.keys():
+                    per_class_confusion_dict[line_class] += confusion_mat
+                else:
+                    per_class_confusion_dict[line_class] = confusion_mat
+
+    mRecall = np.mean(recalls)
+    sRecall = np.std(recalls)
+    medianRecall = np.median(recalls)
+    print(
+        f" On SoccerNet {args.split} set, recall mean value : {mRecall * 100:2.2f}% with standard deviation of {sRecall * 100:2.2f}% and median of {medianRecall * 100:2.2f}%")
+
+    mPrecision = np.mean(precisions)
+    sPrecision = np.std(precisions)
+    medianPrecision = np.median(precisions)
+    print(
+        f" On SoccerNet {args.split} set, precision mean value : {mPrecision * 100:2.2f}% with standard deviation of {sPrecision * 100:2.2f}% and median of {medianPrecision * 100:2.2f}%")
+
+    mAccuracy = np.mean(accuracies)
+    sAccuracy = np.std(accuracies)
+    medianAccuracy = np.median(accuracies)
+    print(
+        f" On SoccerNet {args.split} set, accuracy mean value : {mAccuracy * 100:2.2f}% with standard deviation of {sAccuracy * 100:2.2f}% and median of {medianAccuracy * 100:2.2f}%")
+
+    for line_class, confusion_mat in per_class_confusion_dict.items():
+        class_accuracy = confusion_mat[0, 0] / confusion_mat.sum()
+        class_recall = confusion_mat[0, 0] / (confusion_mat[0, 0] + confusion_mat[1, 0])
+        class_precision = confusion_mat[0, 0] / (confusion_mat[0, 0] + confusion_mat[0, 1])
+        print(
+            f"For class {line_class}, accuracy of {class_accuracy * 100:2.2f}%, precision of {class_precision * 100:2.2f}%  and recall of {class_recall * 100:2.2f}%")
+
+    for k, v in dict_errors.items():
+        fig, ax1 = plt.subplots(figsize=(11, 8))
+        ax1.hist(v, bins=30, range=(0, 60))
+        ax1.set_title(k)
+        ax1.set_xlabel("Errors in pixel")
+        os.makedirs(f"./results/", exist_ok=True)
+        plt.savefig(f"./results/{k}_detection_error.png")
+        plt.close(fig)

tvcalib/sn_segmentation/src/masks_gt2chen.py

@@ -0,0 +1,217 @@
+import pickle
+import torch
+from torch.utils.data import Dataset
+import torchvision
+import torchvision.transforms as T
+from PIL import Image
+import cv2
+import numpy as np
+import matplotlib.pyplot as plt
+import pandas as pd
+import h5py
+from tqdm.auto import tqdm
+from collections import defaultdict
+from pathlib import Path
+import json
+import os
+from argparse import ArgumentParser
+
+lines_classes = [
+    "Big rect. left bottom",
+    "Big rect. left main",
+    "Big rect. left top",
+    "Big rect. right bottom",
+    "Big rect. right main",
+    "Big rect. right top",
+    "Circle central",
+    "Circle left",
+    "Circle right",
+    # "Goal left crossbar",
+    # "Goal left post left ",
+    # "Goal left post right",
+    # "Goal right crossbar",
+    # "Goal right post left",
+    # "Goal right post right",
+    "Goal unknown",
+    "Line unknown",
+    "Middle line",
+    "Side line bottom",
+    "Side line left",
+    "Side line right",
+    "Side line top",
+    "Small rect. left bottom",
+    "Small rect. left main",
+    "Small rect. left top",
+    "Small rect. right bottom",
+    "Small rect. right main",
+    "Small rect. right top",
+]
+
+# RGB values
+palette = {
+    "Big rect. left bottom": (127, 0, 0),
+    "Big rect. left main": (102, 102, 102),
+    "Big rect. left top": (0, 0, 127),
+    "Big rect. right bottom": (86, 32, 39),
+    "Big rect. right main": (48, 77, 0),
+    "Big rect. right top": (14, 97, 100),
+    "Circle central": (0, 0, 255),
+    "Circle left": (255, 127, 0),
+    "Circle right": (0, 255, 255),
+    # "Goal left crossbar": (255, 255, 200),
+    # "Goal left post left ": (165, 255, 0),
+    # "Goal left post right": (155, 119, 45),
+    # "Goal right crossbar": (86, 32, 139),
+    # "Goal right post left": (196, 120, 153),
+    # "Goal right post right": (166, 36, 52),
+    "Goal unknown": (0, 0, 0),
+    "Line unknown": (0, 0, 0),
+    "Middle line": (255, 255, 0),
+    "Side line bottom": (255, 0, 255),
+    "Side line left": (0, 255, 150),
+    "Side line right": (0, 230, 0),
+    "Side line top": (230, 0, 0),
+    "Small rect. left bottom": (0, 150, 255),
+    "Small rect. left main": (254, 173, 225),
+    "Small rect. left top": (87, 72, 39),
+    "Small rect. right bottom": (122, 0, 255),
+    "Small rect. right main": (255, 255, 255),
+    "Small rect. right top": (153, 23, 153),
+}
+
+
+def create_target_from_annotation(width, height, annotation, classes, linewidth=4):
+    """Draw one-hot encoded segments according to the annotation.
+    Creates target that matches image size ([C+1]xHxW).
+    """
+    annotation_abs = defaultdict(list)
+    # unnormalize every point in every class k
+    for k in annotation.keys():
+        if k not in lines_classes:
+            continue
+        start = annotation[k][0].copy()
+        end = annotation[k][-1].copy()
+        for annotation_point in annotation[k]:
+            tup = annotation_point.copy()
+            tup["x"] *= width
+            tup["x"] = int(tup["x"])
+            tup["y"] *= height
+            tup["y"] = int(tup["y"])
+            annotation_abs[k].append(tup)
+
+    # draw lines between annotated points for each segment
+    # offset class +1 such that no classes detected will end in argmax 0
+    # otherwise argmax 0 will be another class
+    classes_segments = np.zeros(shape=(len(classes) + 1, height, width))
+    for cls, points in annotation_abs.items():
+        class_segments = np.zeros(shape=(height, width, 3))
+        for start, end in zip(points, points[1:]):
+            startxy = (start["x"], start["y"])
+            endxy = [end["x"], end["y"]]
+            class_segments = cv2.line(
+                class_segments, startxy, endxy, (1, 1, 1), linewidth
+            )
+        classes_segments[classes.index(cls) + 1] = class_segments[:, :, 1]
+
+    classes_segments = torch.Tensor(classes_segments)
+    return classes_segments
+
+
+class ExtremitiesDataset(Dataset):
+    def __init__(
+        self, root, split, annotations, filter_cam=None, extremities_prefix="", classes=lines_classes, palette=palette
+    ):
+        self.data_root = Path(root)
+        self.split = split
+
+        self.annotations_path = annotations
+
+        if filter_cam is None:
+            files = os.listdir(self.data_root / self.split)
+            self.annotations = sorted([fn for fn in files if fn.endswith("json")])
+            self.images = sorted([fn for fn in files if fn.endswith("jpg")])
+        else:
+            df = pd.read_json(self.data_root / self.split / "match_info_cam_gt.json").T
+            df = df.loc[df.camera == filter_cam]
+            assert len(df.index) > 0
+            df["image_file"] = df.index
+            df = df.sort_values(by=["image_file"])
+            df["annotation_file"] = df["image_file"].apply(
+                lambda s: extremities_prefix + s.split(".jpg")[0] + ".json"
+            )
+            self.annotations = df["annotation_file"].tolist()
+            self.images = df["image_file"].tolist()
+
+        self.classes = classes
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+        # see https://learnopencv.com/pytorch-for-beginners-semantic-segmentation-using-torchvision/
+
+        impath = self.data_root / self.split / self.images[idx]
+        annotation_path = self.annotations_path /  self.annotations[idx]
+        with open(annotation_path, "r") as f:
+            annotation = json.load(f)
+
+        img = Image.open(impath)  # .resize((1280, 720))
+        trf = T.Compose(
+            [
+                T.ToTensor(),
+                T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+            ]
+        )
+        # prepare batches
+        img = trf(img)
+
+        # see https://git.tib.eu/vid2pos/sccvsd/-/blob/master/utils/synthetic_util.py
+        # draw lines (linewidth=4 for 720p) -> hence we rescale first
+        target = create_target_from_annotation(1280, 720, annotation, self.classes)
+        target = target.long().argmax(dim=0).unsqueeze(0)
+        # to binary mask
+        target = target.bool().float()
+        # rescale target equivalent to cv2.resize() with default args (interpolation bilinear) -> same as in torchvision
+        # bilinear -> [0, 0.25, 0.5, 0.7, 1.0] are entries
+        target = torchvision.transforms.Resize((180, 320))(target)
+        # print(torch.unique(target))
+        # to uint8
+        target = (target * 255.0).to(torch.uint8)
+
+        return img, target, impath.name
+
+
+if __name__ == "__main__":
+
+    args = ArgumentParser()
+    args.add_argument("--data_dir", type=Path)
+    args.add_argument("--annotations", type=Path)
+    args.add_argument("--output_dir", type=Path)
+    args.add_argument("--extremities_prefix", type=str, default="")
+    args = args.parse_args()
+
+    data_dir = args.data_dir.parent
+    split = args.data_dir.name
+    output_dir = args.output_dir
+    if not output_dir.exists():
+        raise FileNotFoundError
+
+    dataset = ExtremitiesDataset(data_dir, split, args.annotations, filter_cam="Main camera center", extremities_prefix=args.extremities_prefix)
+
+    # img, edge_map, img_id = dataset[0]
+    # edge_map = edge_map.squeeze(0).numpy()
+    # Image.fromarray(edge_map).show()
+
+    image_src = []
+    edge_maps = np.zeros((len(dataset), 1, 180, 320), dtype=np.uint8)
+    for i, (_, edge_map, img_id) in enumerate(tqdm(dataset)):
+        edge_map = edge_map.numpy()
+        # Image.fromarray(edge_map).show()
+        edge_maps[i] = edge_map
+        image_src.append(img_id)
+
+    with h5py.File(output_dir / "seg_edge_maps.h5", "w") as f:
+        f.create_dataset("edge_map", data=edge_maps)
+
+    with open(output_dir / "seg_image_paths.pkl", "wb") as f:
+        pickle.dump(image_src, f)

tvcalib/sn_segmentation/src/masks_pred2chen.py

@@ -0,0 +1,150 @@
+import argparse
+import os.path
+import pickle
+
+import h5py
+import numpy as np
+import pandas as pd
+from PIL import Image
+from tqdm import tqdm
+import cv2
+
+from SoccerNet.Evaluation.utils_calibration import SoccerPitch
+
+from custom_extremities import CustomNetwork
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Test")
+
+    parser.add_argument(
+        "-s",
+        "--soccernet",
+        default="/nfs/data/soccernet/calibration/",
+        type=str,
+        help="Path to the SoccerNet-V3 dataset folder",
+    )
+    parser.add_argument(
+        "-p",
+        "--prediction",
+        default="/nfs/home/rhotertj/datasets/sn-calib-test_endpoints",
+        required=False,
+        type=str,
+        help="Path to the prediction folder",
+    )
+    parser.add_argument(
+        "--split",
+        required=False,
+        type=str,
+        default="challenge",
+        help="Select the split of data",
+    )
+    parser.add_argument(
+        "--resolution_width",
+        required=False,
+        type=int,
+        default=455,
+        help="width resolution of the images",
+    )
+    parser.add_argument(
+        "--resolution_height",
+        required=False,
+        type=int,
+        default=256,
+        help="height resolution of the images",
+    )
+    parser.add_argument(
+        "--checkpoint",
+        required=False,
+        type=str,
+        help="Path to the custom model checkpoint.",
+    )
+    parser.add_argument("--filter_cam", type=str, required=False)
+    args = parser.parse_args()
+
+    lines_palette = [0, 0, 0]
+    for line_class in SoccerPitch.lines_classes:
+        print(line_class, SoccerPitch.palette[line_class])
+        lines_palette.extend(SoccerPitch.palette[line_class])
+
+    print(lines_palette)
+
+    # exit(0)
+
+    dataset_dir = os.path.join(args.soccernet, args.split)
+    if not os.path.exists(dataset_dir):
+        print("Invalid dataset path !")
+        exit(-1)
+
+    match_info_file = os.path.join(args.soccernet, args.split, "match_info_cam_gt.json")
+    print(match_info_file)
+    if not os.path.exists(match_info_file):
+        exit(-1)
+    df = pd.read_json(match_info_file).T
+    if args.filter_cam:
+        df = df.loc[df.camera == args.filter_cam]
+    df["image_file"] = df.index
+    df = df.sort_values(by=["image_file"])
+    print(df)
+
+    frames = df["image_file"].tolist()
+
+    model = CustomNetwork(args.checkpoint)
+
+    image_src = []
+    edge_maps = np.zeros((len(frames), 1, 180, 320), dtype=np.uint8)
+
+    kernel = np.ones((4, 4), np.uint8)
+
+    with tqdm(enumerate(frames), total=len(frames), ncols=100) as t:
+        for i, frame in t:
+
+            output_prediction_folder = args.prediction
+            if not os.path.exists(output_prediction_folder):
+                os.makedirs(output_prediction_folder)
+
+            frame_path = os.path.join(dataset_dir, frame)
+
+            frame_index = frame.split(".")[0]
+
+            image = Image.open(frame_path)
+
+            semlines = model.forward(image)
+
+            # print(semlines.shape, np.unique(semlines))
+            # set class 9-15 (goal parts) to background
+            mask_goal = (semlines >= 9) & (semlines <= 15)
+            semlines[mask_goal] = 0
+
+            mask = Image.fromarray(semlines.astype(np.uint8)).convert("P")
+            mask.putpalette(lines_palette)
+
+            # to binary edge map
+            mask = np.asarray(mask.convert("L"))
+            mask[mask > 0] = 255
+
+            mask = Image.fromarray(mask)
+            mask = mask.resize((320, 180), resample=Image.NEAREST)
+            # expected linewith @ 720p resulution -> 4px
+
+            mask = np.asarray(mask)
+            # print(mask.shape)
+
+            mask = cv2.erode(mask, kernel, iterations=1)
+
+            # assert len(np.unique(mask)) == 2  # [0, 255]
+
+            # mask_file = os.path.join(output_prediction_folder, frame)
+            # mask.save(mask_file)
+            # print(mask)
+            # exit(0)
+
+            edge_maps[i] = mask
+            image_src.append(frame)
+
+    with h5py.File(
+        os.path.join(output_prediction_folder, "seg_edge_maps.h5"), "w"
+    ) as f:
+        f.create_dataset("edge_map", data=edge_maps)
+
+    with open(os.path.join(output_prediction_folder, "seg_image_paths.pkl"), "wb") as f:
+        pickle.dump(image_src, f)

tvcalib/sn_segmentation/src/segmentation/README.md

@@ -0,0 +1,23 @@
+# Semantic segmentation reference training scripts
+
+This directory is an edited copy of the [torchvision](https://github.com/pytorch/vision/tree/main/references/segmentation) reference scripts.
+
+# Usage
+
+Train from scratch:
+
+```
+python train.py -b 8 --model deeplabv3_resnet101 --aux-loss --weights-backbone ResNet101_Weights.IMAGENET1K_V1 --epochs 30 --output-dir "./checkpoints" --split train
+```
+
+Resume training:
+
+```
+python train.py -b 8 --model deeplabv3_resnet101 --aux-loss --weights-backbone ResNet101_Weights.IMAGENET1K_V1 --epochs 60 --start-epoch 30 --weights /path/to/checkpoints.pt
+```
+
+Evaluate checkpoint:
+
+```
+python train.py -b 4 --model deeplabv3_resnet101 --aux-loss --weights-backbone ResNet101_Weights.IMAGENET1K_V1 --test-only --weights /path/to/checkpoints.pt
+```

tvcalib/sn_segmentation/src/segmentation/coco_utils.py

@@ -0,0 +1,108 @@
+import copy
+import os
+
+import torch
+import torch.utils.data
+import torchvision
+from PIL import Image
+from pycocotools import mask as coco_mask
+from transforms import Compose
+
+
+class FilterAndRemapCocoCategories:
+    def __init__(self, categories, remap=True):
+        self.categories = categories
+        self.remap = remap
+
+    def __call__(self, image, anno):
+        anno = [obj for obj in anno if obj["category_id"] in self.categories]
+        if not self.remap:
+            return image, anno
+        anno = copy.deepcopy(anno)
+        for obj in anno:
+            obj["category_id"] = self.categories.index(obj["category_id"])
+        return image, anno
+
+
+def convert_coco_poly_to_mask(segmentations, height, width):
+    masks = []
+    for polygons in segmentations:
+        rles = coco_mask.frPyObjects(polygons, height, width)
+        mask = coco_mask.decode(rles)
+        if len(mask.shape) < 3:
+            mask = mask[..., None]
+        mask = torch.as_tensor(mask, dtype=torch.uint8)
+        mask = mask.any(dim=2)
+        masks.append(mask)
+    if masks:
+        masks = torch.stack(masks, dim=0)
+    else:
+        masks = torch.zeros((0, height, width), dtype=torch.uint8)
+    return masks
+
+
+class ConvertCocoPolysToMask:
+    def __call__(self, image, anno):
+        w, h = image.size
+        segmentations = [obj["segmentation"] for obj in anno]
+        cats = [obj["category_id"] for obj in anno]
+        if segmentations:
+            masks = convert_coco_poly_to_mask(segmentations, h, w)
+            cats = torch.as_tensor(cats, dtype=masks.dtype)
+            # merge all instance masks into a single segmentation map
+            # with its corresponding categories
+            target, _ = (masks * cats[:, None, None]).max(dim=0)
+            # discard overlapping instances
+            target[masks.sum(0) > 1] = 255
+        else:
+            target = torch.zeros((h, w), dtype=torch.uint8)
+        target = Image.fromarray(target.numpy())
+        return image, target
+
+
+def _coco_remove_images_without_annotations(dataset, cat_list=None):
+    def _has_valid_annotation(anno):
+        # if it's empty, there is no annotation
+        if len(anno) == 0:
+            return False
+        # if more than 1k pixels occupied in the image
+        return sum(obj["area"] for obj in anno) > 1000
+
+    if not isinstance(dataset, torchvision.datasets.CocoDetection):
+        raise TypeError(
+            f"This function expects dataset of type torchvision.datasets.CocoDetection, instead  got {type(dataset)}"
+        )
+
+    ids = []
+    for ds_idx, img_id in enumerate(dataset.ids):
+        ann_ids = dataset.coco.getAnnIds(imgIds=img_id, iscrowd=None)
+        anno = dataset.coco.loadAnns(ann_ids)
+        if cat_list:
+            anno = [obj for obj in anno if obj["category_id"] in cat_list]
+        if _has_valid_annotation(anno):
+            ids.append(ds_idx)
+
+    dataset = torch.utils.data.Subset(dataset, ids)
+    return dataset
+
+
+def get_coco(root, image_set, transforms):
+    PATHS = {
+        "train": ("train2017", os.path.join("annotations", "instances_train2017.json")),
+        "val": ("val2017", os.path.join("annotations", "instances_val2017.json")),
+        # "train": ("val2017", os.path.join("annotations", "instances_val2017.json"))
+    }
+    CAT_LIST = [0, 5, 2, 16, 9, 44, 6, 3, 17, 62, 21, 67, 18, 19, 4, 1, 64, 20, 63, 7, 72]
+
+    transforms = Compose([FilterAndRemapCocoCategories(CAT_LIST, remap=True), ConvertCocoPolysToMask(), transforms])
+
+    img_folder, ann_file = PATHS[image_set]
+    img_folder = os.path.join(root, img_folder)
+    ann_file = os.path.join(root, ann_file)
+
+    dataset = torchvision.datasets.CocoDetection(img_folder, ann_file, transforms=transforms)
+
+    if image_set == "train":
+        dataset = _coco_remove_images_without_annotations(dataset, CAT_LIST)
+
+    return dataset

tvcalib/sn_segmentation/src/segmentation/presets.py

@@ -0,0 +1,39 @@
+import torch
+import transforms as T
+
+
+class SegmentationPresetTrain:
+    def __init__(self, *, base_size, crop_size, hflip_prob=0.5, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
+        min_size = int(0.5 * base_size)
+        max_size = int(2.0 * base_size)
+
+        trans = [T.RandomResize(min_size, max_size)]
+        if hflip_prob > 0:
+            trans.append(T.RandomHorizontalFlip(hflip_prob))
+        trans.extend(
+            [
+                T.RandomCrop(crop_size),
+                T.PILToTensor(),
+                T.ConvertImageDtype(torch.float),
+                T.Normalize(mean=mean, std=std),
+            ]
+        )
+        self.transforms = T.Compose(trans)
+
+    def __call__(self, img, target):
+        return self.transforms(img, target)
+
+
+class SegmentationPresetEval:
+    def __init__(self, *, base_size, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
+        self.transforms = T.Compose(
+            [
+                T.RandomResize(base_size, base_size),
+                T.PILToTensor(),
+                T.ConvertImageDtype(torch.float),
+                T.Normalize(mean=mean, std=std),
+            ]
+        )
+
+    def __call__(self, img, target):
+        return self.transforms(img, target)

tvcalib/sn_segmentation/src/segmentation/soccerdata.py

@@ -0,0 +1,164 @@
+import torch
+from torch.utils.data import Dataset
+import torchvision
+import torchvision.transforms as T
+from PIL import Image
+import cv2
+import numpy as np
+import matplotlib.pyplot as plt
+
+from collections import defaultdict
+from pathlib import Path
+import json
+import os
+
+lines_classes = [
+    'Big rect. left bottom',
+    'Big rect. left main',
+    'Big rect. left top',
+    'Big rect. right bottom',
+    'Big rect. right main',
+    'Big rect. right top',
+    'Circle central',
+    'Circle left',
+    'Circle right',
+    'Goal left crossbar',
+    'Goal left post left ',
+    'Goal left post right',
+    'Goal right crossbar',
+    'Goal right post left',
+    'Goal right post right',
+    'Goal unknown',
+    'Line unknown',
+    'Middle line',
+    'Side line bottom',
+    'Side line left',
+    'Side line right',
+    'Side line top',
+    'Small rect. left bottom',
+    'Small rect. left main',
+    'Small rect. left top',
+    'Small rect. right bottom',
+    'Small rect. right main',
+    'Small rect. right top'
+]
+
+# RGB values
+palette = {
+    'Big rect. left bottom': (127, 0, 0),
+    'Big rect. left main': (102, 102, 102),
+    'Big rect. left top': (0, 0, 127),
+    'Big rect. right bottom': (86, 32, 39),
+    'Big rect. right main': (48, 77, 0),
+    'Big rect. right top': (14, 97, 100),
+    'Circle central': (0, 0, 255),
+    'Circle left': (255, 127, 0),
+    'Circle right': (0, 255, 255),
+    'Goal left crossbar': (255, 255, 200),
+    'Goal left post left ': (165, 255, 0),
+    'Goal left post right': (155, 119, 45),
+    'Goal right crossbar': (86, 32, 139),
+    'Goal right post left': (196, 120, 153),
+    'Goal right post right': (166, 36, 52),
+    'Goal unknown': (0, 0, 0),
+    'Line unknown': (0, 0, 0),
+    'Middle line': (255, 255, 0),
+    'Side line bottom': (255, 0, 255),
+    'Side line left': (0, 255, 150),
+    'Side line right': (0, 230, 0),
+    'Side line top': (230, 0, 0),
+    'Small rect. left bottom': (0, 150, 255),
+    'Small rect. left main': (254, 173, 225),
+    'Small rect. left top': (87, 72, 39),
+    'Small rect. right bottom': (122, 0, 255),
+    'Small rect. right main': (255, 255, 255),
+    'Small rect. right top': (153, 23, 153)
+}
+
+data_dir = Path("data/datasets")
+
+def create_target_from_annotation(width, height, annotation, classes):
+    """Draw one-hot encoded segments according to the annotation.
+    Creates target that matches image size ([C+1]xHxW).
+    """
+    annotation_abs = defaultdict(list)
+    # unnormalize every point in every class k
+    for k in annotation.keys():
+        start = annotation[k][0].copy()
+        end = annotation[k][-1].copy()
+        for annotation_point in annotation[k]:
+            tup = annotation_point.copy()
+            tup["x"] *= width
+            tup["x"] = int(tup["x"])
+            tup["y"] *= height
+            tup["y"] = int(tup["y"])
+            annotation_abs[k].append(tup)
+
+    # draw lines between annotated points for each segment
+    # offset class +1 such that no classes detected will end in argmax 0
+    # otherwise argmax 0 will be another class
+    classes_segments = np.zeros(shape=(len(classes) + 1, height, width))
+    for cls, points in annotation_abs.items():
+        class_segments = np.zeros(shape=(height, width, 3))
+        for start, end in zip(points, points[1:]):
+            startxy = (start["x"], start["y"])
+            endxy = [end["x"], end["y"]]
+            class_segments = cv2.line(class_segments, startxy, endxy, (1,1,1), 5)
+        classes_segments[classes.index(cls) + 1] = class_segments[:,:,1]
+
+    classes_segments = torch.Tensor(classes_segments)
+    return classes_segments
+
+class ExtremitiesDataset(Dataset):
+
+    def __init__(self, root, split, classes=lines_classes, palette=palette):
+        self.data_root = Path(root)
+        self.split = split
+        files = os.listdir(self.data_root / self.split)
+        self.annotations = sorted([fn for fn in files if fn.endswith("json")])
+        self.images = sorted([fn for fn in files if fn.endswith("jpg")])
+        #self.height, self.width = 224, 224
+        self.classes = classes
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+        # see https://learnopencv.com/pytorch-for-beginners-semantic-segmentation-using-torchvision/
+
+        impath = self.data_root / self.split / self.images[idx]
+        annotation_path = self.data_root / self.split / self.annotations[idx]
+        #print(impath)
+        #print(annotation_path)
+        with open(annotation_path, "r") as f:
+            annotation = json.load(f)
+        
+        # setup image, cast to device later in training
+        img = Image.open(impath)
+        trf = T.Compose(
+            [
+                T.Resize(256),
+                #T.CenterCrop(224),
+                T.ToTensor(),
+                T.Normalize(
+                    mean = [0.485, 0.456, 0.406],
+                    std = [0.229, 0.224, 0.225]
+                    )
+            ]
+        )
+        # prepare batches
+        img = trf(img)#.unsqueeze(0)
+        new_height, new_width = img.shape[-2], img.shape[-1]
+    
+        
+        target = create_target_from_annotation(new_width, new_height, annotation, self.classes)
+        #target = torchvision.transforms.functional.center_crop(target, 224)
+        target = target.long().argmax(dim=0)
+
+        return img, target
+
+if __name__ == "__main__":
+    data = ExtremitiesDataset(root=data_dir, split="test")
+    print(data[0][1])
+    target = data[0][1].unsqueeze(0).permute(1,2,0)
+    plt.imshow(target)

tvcalib/sn_segmentation/src/segmentation/train.py

@@ -0,0 +1,341 @@
+import datetime
+import os
+import time
+import warnings
+
+import numpy as np
+
+import presets
+import torch
+import torch.utils.data
+import torchvision
+import utils
+from coco_utils import get_coco
+from torch import nn
+from torchvision.transforms import functional as F, InterpolationMode
+import wandb
+
+from soccerdata import ExtremitiesDataset
+
+
+def get_dataset(dir_path, name, image_set, transform):
+    def sbd(*args, **kwargs):
+        kwargs["root"] = "."
+        print(kwargs)
+        return torchvision.datasets.SBDataset(*args, mode="segmentation", download=True, **kwargs)
+
+    paths = {
+        "voc": (dir_path, torchvision.datasets.VOCSegmentation, 21),
+        "voc_aug": (dir_path, sbd, 21),
+        "coco": (dir_path, get_coco, 21),
+    }
+    p, ds_fn, num_classes = paths[name]
+
+    ds = ds_fn(p, image_set=image_set, transforms=transform)
+    return ds, num_classes
+
+
+def get_transform(train, args):
+    if train:
+        return presets.SegmentationPresetTrain(base_size=520, crop_size=480)
+    elif args.weights and args.test_only:
+        weights = torchvision.models.get_weight(args.weights)
+        trans = weights.transforms()
+
+        def preprocessing(img, target):
+            img = trans(img)
+            size = F.get_dimensions(img)[1:]
+            target = F.resize(target, size, interpolation=InterpolationMode.NEAREST)
+            return img, F.pil_to_tensor(target)
+
+        return preprocessing
+    else:
+        return presets.SegmentationPresetEval(base_size=520)
+
+
+def criterion(inputs, target):
+    losses = {}
+    for name, x in inputs.items():
+        losses[name] = nn.functional.cross_entropy(x, target, ignore_index=255)
+
+    if len(losses) == 1:
+        return losses["out"]
+
+    return losses["out"] + 0.5 * losses["aux"]
+
+
+def evaluate(model, data_loader, device, num_classes):
+    model.eval()
+    confmat = utils.ConfusionMatrix(num_classes)
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    header = "Test:"
+    num_processed_samples = 0
+    losses = []
+    with torch.inference_mode():
+        for image, target in metric_logger.log_every(data_loader, 100, header):
+            image, target = image.to(device), target.to(device)
+            output = model(image)
+            loss = criterion(output, target).unsqueeze(0).detach().cpu()
+            losses.append(loss)
+            output = output["out"]
+
+            # 1xCx224x224
+
+            confmat.update(target.flatten(), output.argmax(1).flatten())
+            # FIXME need to take into account that the datasets
+            # could have been padded in distributed setup
+            num_processed_samples += image.shape[0]
+
+        confmat.reduce_from_all_processes()
+    
+    print(losses[0])
+    #print(losses)
+    losses = torch.cat(losses)
+    loss_argsort = torch.argsort(losses, descending=True)
+    loss_argsort.numpy()
+    np.save("losses_argsort.npy", loss_argsort)
+
+    num_processed_samples = utils.reduce_across_processes(num_processed_samples)
+    if (
+        hasattr(data_loader.dataset, "__len__")
+        and len(data_loader.dataset) != num_processed_samples
+        and torch.distributed.get_rank() == 0
+    ):
+        # See FIXME above
+        warnings.warn(
+            f"It looks like the dataset has {len(data_loader.dataset)} samples, but {num_processed_samples} "
+            "samples were used for the validation, which might bias the results. "
+            "Try adjusting the batch size and / or the world size. "
+            "Setting the world size to 1 is always a safe bet."
+        )
+
+    return confmat
+
+
+def train_one_epoch(model, criterion, optimizer, data_loader, lr_scheduler, device, epoch, print_freq, scaler=None):
+    model.train()
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    metric_logger.add_meter("lr", utils.SmoothedValue(window_size=1, fmt="{value}"))
+    header = f"Epoch: [{epoch}]"
+    for image, target in metric_logger.log_every(data_loader, print_freq, header):
+        image, target = image.to(device), target.to(device)
+        with torch.cuda.amp.autocast(enabled=scaler is not None):
+            output = model(image)
+            loss = criterion(output, target)
+            wandb.log({"loss":  loss})
+
+        optimizer.zero_grad()
+        if scaler is not None:
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            loss.backward()
+            optimizer.step()
+
+        lr_scheduler.step()
+
+        metric_logger.update(loss=loss.item(), lr=optimizer.param_groups[0]["lr"])
+
+
+def main(args):
+    wandb.init(project="sn-calibration", entity="rhotertj")
+    if args.output_dir:
+        utils.mkdir(args.output_dir)
+
+    utils.init_distributed_mode(args)
+    print(args)
+
+    device = torch.device(args.device)
+
+    if args.use_deterministic_algorithms:
+        torch.backends.cudnn.benchmark = False
+        torch.use_deterministic_algorithms(True)
+    else:
+        torch.backends.cudnn.benchmark = True
+
+    #dataset, num_classes = get_dataset(args.data_path, args.dataset, "train", get_transform(True, args))
+    dataset = ExtremitiesDataset(root="/nfs/data/soccernet/calibration", split="train")
+    num_classes = len(dataset.classes) + 1
+
+    dataset_test = ExtremitiesDataset(root="/nfs/data/soccernet/calibration", split="test")
+    #dataset_test, _ = get_dataset(args.data_path, args.dataset, "val", get_transform(False, args))
+
+    if args.distributed:
+        train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+        test_sampler = torch.utils.data.distributed.DistributedSampler(dataset_test, shuffle=False)
+    else:
+        train_sampler = torch.utils.data.RandomSampler(dataset)
+        test_sampler = torch.utils.data.SequentialSampler(dataset_test)
+
+    data_loader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        sampler=train_sampler,
+        num_workers=args.workers,
+        collate_fn=utils.collate_fn,
+        drop_last=True,
+    )
+
+    data_loader_test = torch.utils.data.DataLoader(
+        dataset_test, batch_size=1, sampler=test_sampler, num_workers=args.workers, collate_fn=utils.collate_fn
+    )
+
+
+    model = torchvision.models.segmentation.deeplabv3_resnet101(num_classes=num_classes, aux_loss=args.aux_loss)
+    if args.test_only or args.resume:
+        model.load_state_dict(torch.load(args.weights)["model"], strict=False)
+        
+    #model = torchvision.models.segmentation.__dict__[args.model](
+    #    weights=args.weights, num_classes=num_classes, aux_loss=args.aux_loss
+    #)
+    model.to(device)
+    if args.distributed:
+        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
+
+    model_without_ddp = model
+    if args.distributed:
+        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
+        model_without_ddp = model.module
+
+    params_to_optimize = [
+        {"params": [p for p in model_without_ddp.backbone.parameters() if p.requires_grad]},
+        {"params": [p for p in model_without_ddp.classifier.parameters() if p.requires_grad]},
+    ]
+    if args.aux_loss:
+        params = [p for p in model_without_ddp.aux_classifier.parameters() if p.requires_grad]
+        params_to_optimize.append({"params": params, "lr": args.lr * 10})
+    optimizer = torch.optim.SGD(params_to_optimize, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
+
+    scaler = torch.cuda.amp.GradScaler() if args.amp else None
+
+    iters_per_epoch = len(data_loader)
+    main_lr_scheduler = torch.optim.lr_scheduler.LambdaLR(
+        optimizer, lambda x: (1 - x / (iters_per_epoch * (args.epochs - args.lr_warmup_epochs))) ** 0.9
+    )
+
+    if args.lr_warmup_epochs > 0:
+        warmup_iters = iters_per_epoch * args.lr_warmup_epochs
+        args.lr_warmup_method = args.lr_warmup_method.lower()
+        if args.lr_warmup_method == "linear":
+            warmup_lr_scheduler = torch.optim.lr_scheduler.LinearLR(
+                optimizer, start_factor=args.lr_warmup_decay, total_iters=warmup_iters
+            )
+        elif args.lr_warmup_method == "constant":
+            warmup_lr_scheduler = torch.optim.lr_scheduler.ConstantLR(
+                optimizer, factor=args.lr_warmup_decay, total_iters=warmup_iters
+            )
+        else:
+            raise RuntimeError(
+                f"Invalid warmup lr method '{args.lr_warmup_method}'. Only linear and constant are supported."
+            )
+        lr_scheduler = torch.optim.lr_scheduler.SequentialLR(
+            optimizer, schedulers=[warmup_lr_scheduler, main_lr_scheduler], milestones=[warmup_iters]
+        )
+    else:
+        lr_scheduler = main_lr_scheduler
+
+    if args.resume:
+        #checkpoint = torch.load(args.resume, map_location="cpu")
+        #model_without_ddp.load_state_dict(checkpoint["model"], strict=not args.test_only)
+        if not args.test_only:
+            optimizer.load_state_dict(checkpoint["optimizer"])
+            lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
+            args.start_epoch = checkpoint["epoch"] + 1
+            if args.amp:
+                scaler.load_state_dict(checkpoint["scaler"])
+
+    if args.test_only:
+        # We disable the cudnn benchmarking because it can noticeably affect the accuracy
+        torch.backends.cudnn.benchmark = False
+        torch.backends.cudnn.deterministic = True
+        confmat = evaluate(model, data_loader_test, device=device, num_classes=num_classes)
+        print(confmat)
+        return
+
+    start_time = time.time()
+    for epoch in range(args.start_epoch, args.epochs):
+        if args.distributed:
+            train_sampler.set_epoch(epoch)
+        train_one_epoch(model, criterion, optimizer, data_loader, lr_scheduler, device, epoch, args.print_freq, scaler)
+        confmat = evaluate(model, data_loader_test, device=device, num_classes=num_classes)
+        print(confmat)
+        checkpoint = {
+            "model": model_without_ddp.state_dict(),
+            "optimizer": optimizer.state_dict(),
+            "lr_scheduler": lr_scheduler.state_dict(),
+            "epoch": epoch,
+            "args": args,
+        }
+        if args.amp:
+            checkpoint["scaler"] = scaler.state_dict()
+        utils.save_on_master(checkpoint, os.path.join(args.output_dir, f"model_{epoch}.pth"))
+        utils.save_on_master(checkpoint, os.path.join(args.output_dir, "checkpoint.pth"))
+
+    total_time = time.time() - start_time
+    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+    print(f"Training time {total_time_str}")
+
+
+def get_args_parser(add_help=True):
+    import argparse
+
+    parser = argparse.ArgumentParser(description="PyTorch Segmentation Training", add_help=add_help)
+
+    parser.add_argument("--data-path", default="/datasets01/COCO/022719/", type=str, help="dataset path")
+    parser.add_argument("--dataset", default="coco", type=str, help="dataset name")
+    parser.add_argument("--model", default="fcn_resnet101", type=str, help="model name")
+    parser.add_argument("--aux-loss", action="store_true", help="auxiliar loss")
+    parser.add_argument("--device", default="cuda", type=str, help="device (Use cuda or cpu Default: cuda)")
+    parser.add_argument(
+        "-b", "--batch-size", default=8, type=int, help="images per gpu, the total batch size is $NGPU x batch_size"
+    )
+    parser.add_argument("--epochs", default=30, type=int, metavar="N", help="number of total epochs to run")
+
+    parser.add_argument(
+        "-j", "--workers", default=16, type=int, metavar="N", help="number of data loading workers (default: 16)"
+    )
+    parser.add_argument("--lr", default=0.01, type=float, help="initial learning rate")
+    parser.add_argument("--momentum", default=0.9, type=float, metavar="M", help="momentum")
+    parser.add_argument(
+        "--wd",
+        "--weight-decay",
+        default=1e-4,
+        type=float,
+        metavar="W",
+        help="weight decay (default: 1e-4)",
+        dest="weight_decay",
+    )
+    parser.add_argument("--lr-warmup-epochs", default=0, type=int, help="the number of epochs to warmup (default: 0)")
+    parser.add_argument("--lr-warmup-method", default="linear", type=str, help="the warmup method (default: linear)")
+    parser.add_argument("--lr-warmup-decay", default=0.01, type=float, help="the decay for lr")
+    parser.add_argument("--print-freq", default=10, type=int, help="print frequency")
+    parser.add_argument("--output-dir", default=".", type=str, help="path to save outputs")
+    parser.add_argument("--resume", default="", type=str, help="path of checkpoint")
+    parser.add_argument("--start-epoch", default=0, type=int, metavar="N", help="start epoch")
+    parser.add_argument(
+        "--test-only",
+        dest="test_only",
+        help="Only test the model",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--use-deterministic-algorithms", action="store_true", help="Forces the use of deterministic algorithms only."
+    )
+    # distributed training parameters
+    parser.add_argument("--world-size", default=1, type=int, help="number of distributed processes")
+    parser.add_argument("--dist-url", default="env://", type=str, help="url used to set up distributed training")
+
+    parser.add_argument("--weights", default=None, type=str, help="the weights enum name to load")
+    parser.add_argument("--weights-backbone", default=None, type=str, help="the backbone weights enum name to load")
+
+    # Mixed precision training parameters
+    parser.add_argument("--amp", action="store_true", help="Use torch.cuda.amp for mixed precision training")
+    parser.add_argument("--split", default="train", type=str, help="Dataset split to be used for training")
+
+    return parser
+
+
+if __name__ == "__main__":
+    args = get_args_parser().parse_args()
+    main(args)

tvcalib/sn_segmentation/src/segmentation/transforms.py

@@ -0,0 +1,100 @@
+import random
+
+import numpy as np
+import torch
+from torchvision import transforms as T
+from torchvision.transforms import functional as F
+
+
+def pad_if_smaller(img, size, fill=0):
+    min_size = min(img.size)
+    if min_size < size:
+        ow, oh = img.size
+        padh = size - oh if oh < size else 0
+        padw = size - ow if ow < size else 0
+        img = F.pad(img, (0, 0, padw, padh), fill=fill)
+    return img
+
+
+class Compose:
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, image, target):
+        for t in self.transforms:
+            image, target = t(image, target)
+        return image, target
+
+
+class RandomResize:
+    def __init__(self, min_size, max_size=None):
+        self.min_size = min_size
+        if max_size is None:
+            max_size = min_size
+        self.max_size = max_size
+
+    def __call__(self, image, target):
+        size = random.randint(self.min_size, self.max_size)
+        image = F.resize(image, size)
+        target = F.resize(target, size, interpolation=T.InterpolationMode.NEAREST)
+        return image, target
+
+
+class RandomHorizontalFlip:
+    def __init__(self, flip_prob):
+        self.flip_prob = flip_prob
+
+    def __call__(self, image, target):
+        if random.random() < self.flip_prob:
+            image = F.hflip(image)
+            target = F.hflip(target)
+        return image, target
+
+
+class RandomCrop:
+    def __init__(self, size):
+        self.size = size
+
+    def __call__(self, image, target):
+        image = pad_if_smaller(image, self.size)
+        target = pad_if_smaller(target, self.size, fill=255)
+        crop_params = T.RandomCrop.get_params(image, (self.size, self.size))
+        image = F.crop(image, *crop_params)
+        target = F.crop(target, *crop_params)
+        return image, target
+
+
+class CenterCrop:
+    def __init__(self, size):
+        self.size = size
+
+    def __call__(self, image, target):
+        image = F.center_crop(image, self.size)
+        target = F.center_crop(target, self.size)
+        return image, target
+
+
+class PILToTensor:
+    def __call__(self, image, target):
+        image = F.pil_to_tensor(image)
+        target = torch.as_tensor(np.array(target), dtype=torch.int64)
+        return image, target
+
+
+class ConvertImageDtype:
+    def __init__(self, dtype):
+        self.dtype = dtype
+
+    def __call__(self, image, target):
+        image = F.convert_image_dtype(image, self.dtype)
+        return image, target
+
+
+class Normalize:
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+
+    def __call__(self, image, target):
+        image = F.normalize(image, mean=self.mean, std=self.std)
+        return image, target

tvcalib/sn_segmentation/src/segmentation/utils.py

@@ -0,0 +1,304 @@
+import datetime
+import errno
+import os
+import time
+from collections import defaultdict, deque
+
+import torch
+import torch.distributed as dist
+
+
+class SmoothedValue:
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.4f} ({global_avg:.4f})"
+        self.deque = deque(maxlen=window_size)
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """
+        Warning: does not synchronize the deque!
+        """
+        t = reduce_across_processes([self.count, self.total])
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = torch.tensor(list(self.deque))
+        return d.median().item()
+
+    @property
+    def avg(self):
+        d = torch.tensor(list(self.deque), dtype=torch.float32)
+        return d.mean().item()
+
+    @property
+    def global_avg(self):
+        return self.total / self.count
+
+    @property
+    def max(self):
+        return max(self.deque)
+
+    @property
+    def value(self):
+        return self.deque[-1]
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median, avg=self.avg, global_avg=self.global_avg, max=self.max, value=self.value
+        )
+
+
+class ConfusionMatrix:
+    def __init__(self, num_classes):
+        self.num_classes = num_classes
+        self.mat = None
+
+    def update(self, a, b):
+        n = self.num_classes
+        if self.mat is None:
+            self.mat = torch.zeros((n, n), dtype=torch.int64, device=a.device)
+        with torch.inference_mode():
+            k = (a >= 0) & (a < n)
+            inds = n * a[k].to(torch.int64) + b[k]
+            self.mat += torch.bincount(inds, minlength=n ** 2).reshape(n, n)
+
+    def reset(self):
+        self.mat.zero_()
+
+    def compute(self):
+        h = self.mat.float()
+        acc_global = torch.diag(h).sum() / h.sum()
+        acc = torch.diag(h) / h.sum(1)
+        iu = torch.diag(h) / (h.sum(1) + h.sum(0) - torch.diag(h))
+        return acc_global, acc, iu
+
+    def reduce_from_all_processes(self):
+        reduce_across_processes(self.mat)
+
+    def __str__(self):
+        acc_global, acc, iu = self.compute()
+        return ("global correct: {:.1f}\naverage row correct: {}\nIoU: {}\nmean IoU: {:.1f}").format(
+            acc_global.item() * 100,
+            [f"{i:.1f}" for i in (acc * 100).tolist()],
+            [f"{i:.1f}" for i in (iu * 100).tolist()],
+            iu.mean().item() * 100,
+        )
+
+
+class MetricLogger:
+    def __init__(self, delimiter="\t"):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            if not isinstance(v, (float, int)):
+                raise TypeError(
+                    f"This method expects the value of the input arguments to be of type float or int, instead  got {type(v)}"
+                )
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError(f"'{type(self).__name__}' object has no attribute '{attr}'")
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append(f"{name}: {str(meter)}")
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def add_meter(self, name, meter):
+        self.meters[name] = meter
+
+    def log_every(self, iterable, print_freq, header=None):
+        i = 0
+        if not header:
+            header = ""
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt="{avg:.4f}")
+        data_time = SmoothedValue(fmt="{avg:.4f}")
+        space_fmt = ":" + str(len(str(len(iterable)))) + "d"
+        if torch.cuda.is_available():
+            log_msg = self.delimiter.join(
+                [
+                    header,
+                    "[{0" + space_fmt + "}/{1}]",
+                    "eta: {eta}",
+                    "{meters}",
+                    "time: {time}",
+                    "data: {data}",
+                    "max mem: {memory:.0f}",
+                ]
+            )
+        else:
+            log_msg = self.delimiter.join(
+                [header, "[{0" + space_fmt + "}/{1}]", "eta: {eta}", "{meters}", "time: {time}", "data: {data}"]
+            )
+        MB = 1024.0 * 1024.0
+        for obj in iterable:
+            data_time.update(time.time() - end)
+            yield obj
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0:
+                eta_seconds = iter_time.global_avg * (len(iterable) - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available():
+                    print(
+                        log_msg.format(
+                            i,
+                            len(iterable),
+                            eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time),
+                            data=str(data_time),
+                            memory=torch.cuda.max_memory_allocated() / MB,
+                        )
+                    )
+                else:
+                    print(
+                        log_msg.format(
+                            i, len(iterable), eta=eta_string, meters=str(self), time=str(iter_time), data=str(data_time)
+                        )
+                    )
+            i += 1
+            end = time.time()
+        total_time = time.time() - start_time
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        print(f"{header} Total time: {total_time_str}")
+
+
+def cat_list(images, fill_value=0):
+    max_size = tuple(max(s) for s in zip(*[img.shape for img in images]))
+    batch_shape = (len(images),) + max_size
+    batched_imgs = images[0].new(*batch_shape).fill_(fill_value)
+    for img, pad_img in zip(images, batched_imgs):
+        pad_img[..., : img.shape[-2], : img.shape[-1]].copy_(img)
+    return batched_imgs
+
+
+def collate_fn(batch):
+    images, targets = list(zip(*batch))
+    batched_imgs = cat_list(images, fill_value=0)
+    batched_targets = cat_list(targets, fill_value=255)
+    return batched_imgs, batched_targets
+
+
+def mkdir(path):
+    try:
+        os.makedirs(path)
+    except OSError as e:
+        if e.errno != errno.EEXIST:
+            raise
+
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    import builtins as __builtin__
+
+    builtin_print = __builtin__.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop("force", False)
+        if is_master or force:
+            builtin_print(*args, **kwargs)
+
+    __builtin__.print = print
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
+
+
+def save_on_master(*args, **kwargs):
+    if is_main_process():
+        torch.save(*args, **kwargs)
+
+
+def init_distributed_mode(args):
+    print("Not using distributed mode")
+    args.distributed = False
+    return
+
+    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ["WORLD_SIZE"])
+        args.gpu = int(os.environ["LOCAL_RANK"])
+    elif "SLURM_PROCID" in os.environ:
+        args.rank = int(os.environ["SLURM_PROCID"])
+        args.gpu = args.rank % torch.cuda.device_count()
+    elif hasattr(args, "rank"):
+        pass
+    else:
+        print("Not using distributed mode")
+        args.distributed = False
+        return
+
+    args.distributed = True
+
+    torch.cuda.set_device(args.gpu)
+    args.dist_backend = "nccl"
+    print(f"| distributed init (rank {args.rank}): {args.dist_url}", flush=True)
+    torch.distributed.init_process_group(
+        backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank
+    )
+    torch.distributed.barrier()
+    setup_for_distributed(args.rank == 0)
+
+
+def reduce_across_processes(val):
+    if not is_dist_avail_and_initialized():
+        # nothing to sync, but we still convert to tensor for consistency with the distributed case.
+        return torch.tensor(val)
+
+    t = torch.tensor(val, device="cuda")
+    dist.barrier()
+    dist.all_reduce(t)
+    return t

tvcalib/utils/data_distr.py

@@ -0,0 +1,44 @@
+import torch
+
+
+def mean_std_with_confidence_interval(
+    vmin, vmax, sigma_scale: float, _steps=1000, round_decimals=4
+):
+    """Computes mean and std given min,max values with respect a confidence interval (sigma_scale).
+
+    sigma_scale = 1.65 -> 90% of samples are in range [vmin, vmax]
+    sigma_scale = 1.96 -> 95% of samples are in range [vmin, vmax]
+    sigma_scale = 2.58 -> 99% of samples are in range [vmin, vmax]
+    """
+
+    # sample from uniform distribution
+    x = torch.linspace(vmin, vmax, _steps)
+    mu = x.mean(dim=-1)
+    sigma = x.std(dim=-1)
+    return (round(mu.item(), round_decimals), round((sigma * sigma_scale).item(), round_decimals))
+
+
+class FeatureScalerZScore(torch.nn.Module):
+    def __init__(self, loc: float, scale: float) -> None:
+        # Transforms data from distribution parameterized by loc (mean) and scale (=sigma*scaling factor).
+        super(FeatureScalerZScore, self).__init__()
+
+        self.loc = loc
+        self.scale = scale
+
+    def forward(self, z):
+        """
+        Args:
+            z (Tensor): tensor of size (B, *) to be denormalized.
+        Returns:
+            x: tensor.
+        """
+        return self.denormalize(z)
+
+    def denormalize(self, z):
+        x = z * self.scale + self.loc
+        return x
+
+    def normalize(self, x):
+        z = (x - self.loc) / self.scale
+        return z

tvcalib/utils/io.py

@@ -0,0 +1,44 @@
+import yaml
+import json
+from typing import List
+import torch
+
+
+def tensor2list(d: dict):
+    tensor2list_lambda = lambda x: x.detach().cpu().numpy().tolist()
+    for k in d.keys():
+        if isinstance(d[k], torch.Tensor):
+            d[k] = tensor2list_lambda(d[k])
+        if isinstance(d[k], List):
+            if isinstance(d[k][0], torch.Tensor):
+                d[k] = [tensor2list_lambda(x) for x in d[k]]
+    return d
+
+
+def write_json(json_serializable_dict, fout, indent=2):
+    with open(fout, "w") as fw:
+        json.dump(json_serializable_dict, fw, indent=indent)
+
+
+def write_yaml(json_serializable_dict, fout):
+    with open(fout, "w") as fw:
+        yaml.dump(json_serializable_dict, fw, default_flow_style=False)
+
+
+def detach_dict(x_dict):
+    with torch.no_grad():
+        for k in x_dict.keys():
+            if isinstance(x_dict[k], torch.Tensor):
+                x_dict[k] = x_dict[k].detach().cpu()
+            elif isinstance(x_dict[k], dict):
+                x_dict[k] = detach_dict(x_dict[k])
+    return x_dict
+
+
+def tensor2list(xdict):
+    for k in xdict.keys():
+        if isinstance(xdict[k], torch.Tensor):
+            xdict[k] = xdict[k].numpy().tolist()
+        elif isinstance(xdict[k], dict):
+            xdict[k] = tensor2list(xdict[k])
+    return xdict

tvcalib/utils/linalg.py

@@ -0,0 +1,106 @@
+from typing import Optional, Union
+import torch
+from kornia.geometry.conversions import convert_points_from_homogeneous
+
+
+class LineCollection:
+    def __init__(
+        self,
+        support: torch.tensor,
+        direction_norm: torch.tensor,
+        direction: Optional[torch.tensor] = None,
+    ):
+        """Wrapper class to represent lines by support and direction vectors.
+
+        Args:
+            support (torch.tensor): with shape (*, {2,3})
+            direction_norm (torch.tensor): with shape (*, {2,3})
+            direction (Optional[torch.tensor], optional): Unnormalized direction vector. Defaults to None.
+        """
+        self.support = support
+        self.direction_norm = direction_norm
+        self.direction = direction
+
+    def __copy__(self):
+        return LineCollection(
+            self.support.clone(),
+            self.direction_norm.clone(),
+            self.direction.clone() if self.direction is not None else None,
+        )
+
+    def copy(self):
+        return self.__copy__()
+
+    def shape(self):
+        return f"support={self.support.shape} direction_norm={self.direction_norm.shape} direction={self.direction.shape if self.direction else None}"
+
+    def __repr__(self) -> str:
+        return f"{self.__class__} " + self.shape()
+
+
+def distance_line_pointcloud_3d(
+    e1: torch.Tensor,
+    r1: torch.Tensor,
+    pc: torch.Tensor,
+    reduce: Union[None, str] = None,
+) -> torch.Tensor:
+    """
+    Line to point cloud distance with arbitrary leading dimensions.
+
+    TODO. if cross = (0.0.0) -> distance=0 otherwise NaNs are returned
+
+    https://mathworld.wolfram.com/Point-LineDistance2-Dimensional.html
+    Args:
+        e1 (torch.Tensor): direction vector of shape (*, B, 1, 3)
+        r1 (torch.Tensor): support vector of shape (*, B, 1, 3)
+        pc (torch.Tensor): point cloud of shape (*, B, A, 3)
+        reduce (Union[None, str]): reduce distance for all points to one using 'mean' or 'min'
+    Returns:
+        distance of an infinite line to given points, (*, B, ) using reduce='mean' or reduce='min' or (*, B, A) if reduce=False
+    """
+
+    num_points = pc.shape[-2]
+    _sub = r1 - pc  # (*, B, A, 3)
+
+    cross = torch.cross(e1.repeat_interleave(num_points, dim=-2), _sub, dim=-1)  # (*, B, A, 3)
+
+    e1_norm = torch.linalg.norm(e1, dim=-1)
+    cross_norm = torch.linalg.norm(cross, dim=-1)
+
+    d = cross_norm / e1_norm
+    if reduce == "mean":
+        return d.mean(dim=-1)  # (*, B, )
+    elif reduce == "min":
+        return d.min(dim=-1)[0]  # (*, B, )
+
+    return d  # (B, A)
+
+
+def distance_point_pointcloud(points: torch.Tensor, pointcloud: torch.Tensor) -> torch.Tensor:
+    """Batched version for point-pointcloud distance calculation
+    Args:
+        points (torch.Tensor): N points in homogenous coordinates; shape (B, T, 3, S, N)
+        pointcloud (torch.Tensor): N_star points for each pointcloud; shape (B, T, S, N_star, 2)
+
+    Returns:
+        torch.Tensor: Minimum distance for each point N to pointcloud; shape (B, T, 1, S, N)
+    """
+
+    batch_size, T, _, S, N = points.shape
+    batch_size, T, S, N_star, _ = pointcloud.shape
+
+    pointcloud = pointcloud.reshape(batch_size * T * S, N_star, 2)
+
+    points = convert_points_from_homogeneous(
+        points.permute(0, 1, 3, 4, 2).reshape(batch_size * T * S, N, 3)
+    )
+
+    # cdist signature: (B, P, M), (B, R, M) -> (B, P, R)
+    distances = torch.cdist(points, pointcloud, p=2)  # (B*T*S, N, N_star)
+
+    distances = distances.view(batch_size, T, S, N, N_star)
+    distances = distances.unsqueeze(-4)
+
+    # distance to nearest point from point cloud (batch_size, T, 1, S, N, N_star)
+    distances = distances.min(dim=-1)[0]
+    return distances

tvcalib/utils/objects_3d.py

@@ -0,0 +1,1674 @@
+from abc import ABCMeta
+from typing import List
+import kornia
+import torch
+import numpy as np
+import random
+from .linalg import LineCollection
+from torch.utils.data import Dataset, DataLoader
+from pytorch_lightning import LightningDataModule
+
+
+class SoccerPitchSN:
+    """Static class variables that are specified by the rules of the game"""
+
+    GOAL_LINE_TO_PENALTY_MARK = 11.0
+    PENALTY_AREA_WIDTH = 40.32
+    PENALTY_AREA_LENGTH = 16.5
+    GOAL_AREA_WIDTH = 18.32
+    GOAL_AREA_LENGTH = 5.5
+    CENTER_CIRCLE_RADIUS = 9.15
+    GOAL_HEIGHT = 2.44
+    GOAL_LENGTH = 7.32
+
+    lines_classes = [
+        "Big rect. left bottom",
+        "Big rect. left main",
+        "Big rect. left top",
+        "Big rect. right bottom",
+        "Big rect. right main",
+        "Big rect. right top",
+        "Circle central",
+        "Circle left",
+        "Circle right",
+        "Goal left crossbar",
+        "Goal left post left ",
+        "Goal left post right",
+        "Goal right crossbar",
+        "Goal right post left",
+        "Goal right post right",
+        "Goal unknown",
+        "Line unknown",
+        "Middle line",
+        "Side line bottom",
+        "Side line left",
+        "Side line right",
+        "Side line top",
+        "Small rect. left bottom",
+        "Small rect. left main",
+        "Small rect. left top",
+        "Small rect. right bottom",
+        "Small rect. right main",
+        "Small rect. right top",
+    ]
+
+    symetric_classes = {
+        "Side line top": "Side line bottom",
+        "Side line bottom": "Side line top",
+        "Side line left": "Side line right",
+        "Middle line": "Middle line",
+        "Side line right": "Side line left",
+        "Big rect. left top": "Big rect. right bottom",
+        "Big rect. left bottom": "Big rect. right top",
+        "Big rect. left main": "Big rect. right main",
+        "Big rect. right top": "Big rect. left bottom",
+        "Big rect. right bottom": "Big rect. left top",
+        "Big rect. right main": "Big rect. left main",
+        "Small rect. left top": "Small rect. right bottom",
+        "Small rect. left bottom": "Small rect. right top",
+        "Small rect. left main": "Small rect. right main",
+        "Small rect. right top": "Small rect. left bottom",
+        "Small rect. right bottom": "Small rect. left top",
+        "Small rect. right main": "Small rect. left main",
+        "Circle left": "Circle right",
+        "Circle central": "Circle central",
+        "Circle right": "Circle left",
+        "Goal left crossbar": "Goal right crossbar",
+        "Goal left post left ": "Goal right post right",
+        "Goal left post right": "Goal right post left",
+        "Goal right crossbar": "Goal left crossbar",
+        "Goal right post left": "Goal left post right",
+        "Goal right post right": "Goal left post left ",
+        "Goal unknown": "Goal unknown",
+        "Line unknown": "Line unknown",
+    }
+
+    # RGB values
+    palette = {
+        "Big rect. left bottom": (127, 0, 0),
+        "Big rect. left main": (102, 102, 102),
+        "Big rect. left top": (0, 0, 127),
+        "Big rect. right bottom": (86, 32, 39),
+        "Big rect. right main": (48, 77, 0),
+        "Big rect. right top": (14, 97, 100),
+        "Circle central": (0, 0, 255),
+        "Circle left": (255, 127, 0),
+        "Circle right": (0, 255, 255),
+        "Goal left crossbar": (255, 255, 200),
+        "Goal left post left ": (165, 255, 0),
+        "Goal left post right": (155, 119, 45),
+        "Goal right crossbar": (86, 32, 139),
+        "Goal right post left": (196, 120, 153),
+        "Goal right post right": (166, 36, 52),
+        "Goal unknown": (0, 0, 0),
+        "Line unknown": (0, 0, 0),
+        "Middle line": (255, 255, 0),
+        "Side line bottom": (255, 0, 255),
+        "Side line left": (0, 255, 150),
+        "Side line right": (0, 230, 0),
+        "Side line top": (230, 0, 0),
+        "Small rect. left bottom": (0, 150, 255),
+        "Small rect. left main": (254, 173, 225),
+        "Small rect. left top": (87, 72, 39),
+        "Small rect. right bottom": (122, 0, 255),
+        "Small rect. right main": (128, 128, 128), # (255, 255, 255)
+        "Small rect. right top": (153, 23, 153),
+    }
+
+    def __init__(self, pitch_length=105.0, pitch_width=68.0):
+        """
+        Initialize 3D coordinates of all elements of the soccer pitch.
+        :param pitch_length: According to FIFA rules, length belong to [90,120] meters
+        :param pitch_width: According to FIFA rules, length belong to [45,90] meters
+        """
+        self.PITCH_LENGTH = pitch_length
+        self.PITCH_WIDTH = pitch_width
+
+        self.center_mark = np.array([0, 0, 0], dtype="float")
+        self.halfway_and_bottom_touch_line_mark = np.array([0, pitch_width / 2.0, 0], dtype="float")
+        self.halfway_and_top_touch_line_mark = np.array([0, -pitch_width / 2.0, 0], dtype="float")
+        self.halfway_line_and_center_circle_top_mark = np.array(
+            [0, -SoccerPitchSN.CENTER_CIRCLE_RADIUS, 0], dtype="float"
+        )
+        self.halfway_line_and_center_circle_bottom_mark = np.array(
+            [0, SoccerPitchSN.CENTER_CIRCLE_RADIUS, 0], dtype="float"
+        )
+        self.bottom_right_corner = np.array(
+            [pitch_length / 2.0, pitch_width / 2.0, 0], dtype="float"
+        )
+        self.bottom_left_corner = np.array(
+            [-pitch_length / 2.0, pitch_width / 2.0, 0], dtype="float"
+        )
+        self.top_right_corner = np.array([pitch_length / 2.0, -pitch_width / 2.0, 0], dtype="float")
+        self.top_left_corner = np.array([-pitch_length / 2.0, -34, 0], dtype="float")
+
+        self.left_goal_bottom_left_post = np.array(
+            [-pitch_length / 2.0, SoccerPitchSN.GOAL_LENGTH / 2.0, 0.0], dtype="float"
+        )
+        self.left_goal_top_left_post = np.array(
+            [-pitch_length / 2.0, SoccerPitchSN.GOAL_LENGTH / 2.0, -SoccerPitchSN.GOAL_HEIGHT],
+            dtype="float",
+        )
+        self.left_goal_bottom_right_post = np.array(
+            [-pitch_length / 2.0, -SoccerPitchSN.GOAL_LENGTH / 2.0, 0.0], dtype="float"
+        )
+        self.left_goal_top_right_post = np.array(
+            [-pitch_length / 2.0, -SoccerPitchSN.GOAL_LENGTH / 2.0, -SoccerPitchSN.GOAL_HEIGHT],
+            dtype="float",
+        )
+
+        self.right_goal_bottom_left_post = np.array(
+            [pitch_length / 2.0, -SoccerPitchSN.GOAL_LENGTH / 2.0, 0.0], dtype="float"
+        )
+        self.right_goal_top_left_post = np.array(
+            [pitch_length / 2.0, -SoccerPitchSN.GOAL_LENGTH / 2.0, -SoccerPitchSN.GOAL_HEIGHT],
+            dtype="float",
+        )
+        self.right_goal_bottom_right_post = np.array(
+            [pitch_length / 2.0, SoccerPitchSN.GOAL_LENGTH / 2.0, 0.0], dtype="float"
+        )
+        self.right_goal_top_right_post = np.array(
+            [pitch_length / 2.0, SoccerPitchSN.GOAL_LENGTH / 2.0, -SoccerPitchSN.GOAL_HEIGHT],
+            dtype="float",
+        )
+
+        self.left_penalty_mark = np.array(
+            [-pitch_length / 2.0 + SoccerPitchSN.GOAL_LINE_TO_PENALTY_MARK, 0, 0], dtype="float"
+        )
+        self.right_penalty_mark = np.array(
+            [pitch_length / 2.0 - SoccerPitchSN.GOAL_LINE_TO_PENALTY_MARK, 0, 0], dtype="float"
+        )
+
+        self.left_penalty_area_top_right_corner = np.array(
+            [
+                -pitch_length / 2.0 + SoccerPitchSN.PENALTY_AREA_LENGTH,
+                -SoccerPitchSN.PENALTY_AREA_WIDTH / 2.0,
+                0,
+            ],
+            dtype="float",
+        )
+        self.left_penalty_area_top_left_corner = np.array(
+            [-pitch_length / 2.0, -SoccerPitchSN.PENALTY_AREA_WIDTH / 2.0, 0], dtype="float"
+        )
+        self.left_penalty_area_bottom_right_corner = np.array(
+            [
+                -pitch_length / 2.0 + SoccerPitchSN.PENALTY_AREA_LENGTH,
+                SoccerPitchSN.PENALTY_AREA_WIDTH / 2.0,
+                0,
+            ],
+            dtype="float",
+        )
+        self.left_penalty_area_bottom_left_corner = np.array(
+            [-pitch_length / 2.0, SoccerPitchSN.PENALTY_AREA_WIDTH / 2.0, 0], dtype="float"
+        )
+        self.right_penalty_area_top_right_corner = np.array(
+            [pitch_length / 2.0, -SoccerPitchSN.PENALTY_AREA_WIDTH / 2.0, 0], dtype="float"
+        )
+        self.right_penalty_area_top_left_corner = np.array(
+            [
+                pitch_length / 2.0 - SoccerPitchSN.PENALTY_AREA_LENGTH,
+                -SoccerPitchSN.PENALTY_AREA_WIDTH / 2.0,
+                0,
+            ],
+            dtype="float",
+        )
+        self.right_penalty_area_bottom_right_corner = np.array(
+            [pitch_length / 2.0, SoccerPitchSN.PENALTY_AREA_WIDTH / 2.0, 0], dtype="float"
+        )
+        self.right_penalty_area_bottom_left_corner = np.array(
+            [
+                pitch_length / 2.0 - SoccerPitchSN.PENALTY_AREA_LENGTH,
+                SoccerPitchSN.PENALTY_AREA_WIDTH / 2.0,
+                0,
+            ],
+            dtype="float",
+        )
+
+        self.left_goal_area_top_right_corner = np.array(
+            [
+                -pitch_length / 2.0 + SoccerPitchSN.GOAL_AREA_LENGTH,
+                -SoccerPitchSN.GOAL_AREA_WIDTH / 2.0,
+                0,
+            ],
+            dtype="float",
+        )
+        self.left_goal_area_top_left_corner = np.array(
+            [-pitch_length / 2.0, -SoccerPitchSN.GOAL_AREA_WIDTH / 2.0, 0], dtype="float"
+        )
+        self.left_goal_area_bottom_right_corner = np.array(
+            [
+                -pitch_length / 2.0 + SoccerPitchSN.GOAL_AREA_LENGTH,
+                SoccerPitchSN.GOAL_AREA_WIDTH / 2.0,
+                0,
+            ],
+            dtype="float",
+        )
+        self.left_goal_area_bottom_left_corner = np.array(
+            [-pitch_length / 2.0, SoccerPitchSN.GOAL_AREA_WIDTH / 2.0, 0], dtype="float"
+        )
+        self.right_goal_area_top_right_corner = np.array(
+            [pitch_length / 2.0, -SoccerPitchSN.GOAL_AREA_WIDTH / 2.0, 0], dtype="float"
+        )
+        self.right_goal_area_top_left_corner = np.array(
+            [
+                pitch_length / 2.0 - SoccerPitchSN.GOAL_AREA_LENGTH,
+                -SoccerPitchSN.GOAL_AREA_WIDTH / 2.0,
+                0,
+            ],
+            dtype="float",
+        )
+        self.right_goal_area_bottom_right_corner = np.array(
+            [pitch_length / 2.0, SoccerPitchSN.GOAL_AREA_WIDTH / 2.0, 0], dtype="float"
+        )
+        self.right_goal_area_bottom_left_corner = np.array(
+            [
+                pitch_length / 2.0 - SoccerPitchSN.GOAL_AREA_LENGTH,
+                SoccerPitchSN.GOAL_AREA_WIDTH / 2.0,
+                0,
+            ],
+            dtype="float",
+        )
+
+        x = -pitch_length / 2.0 + SoccerPitchSN.PENALTY_AREA_LENGTH
+        dx = SoccerPitchSN.PENALTY_AREA_LENGTH - SoccerPitchSN.GOAL_LINE_TO_PENALTY_MARK
+        y = -np.sqrt(
+            SoccerPitchSN.CENTER_CIRCLE_RADIUS * SoccerPitchSN.CENTER_CIRCLE_RADIUS - dx * dx
+        )
+        self.top_left_16M_penalty_arc_mark = np.array([x, y, 0], dtype="float")
+
+        x = pitch_length / 2.0 - SoccerPitchSN.PENALTY_AREA_LENGTH
+        dx = SoccerPitchSN.PENALTY_AREA_LENGTH - SoccerPitchSN.GOAL_LINE_TO_PENALTY_MARK
+        y = -np.sqrt(
+            SoccerPitchSN.CENTER_CIRCLE_RADIUS * SoccerPitchSN.CENTER_CIRCLE_RADIUS - dx * dx
+        )
+        self.top_right_16M_penalty_arc_mark = np.array([x, y, 0], dtype="float")
+
+        x = -pitch_length / 2.0 + SoccerPitchSN.PENALTY_AREA_LENGTH
+        dx = SoccerPitchSN.PENALTY_AREA_LENGTH - SoccerPitchSN.GOAL_LINE_TO_PENALTY_MARK
+        y = np.sqrt(
+            SoccerPitchSN.CENTER_CIRCLE_RADIUS * SoccerPitchSN.CENTER_CIRCLE_RADIUS - dx * dx
+        )
+        self.bottom_left_16M_penalty_arc_mark = np.array([x, y, 0], dtype="float")
+
+        x = pitch_length / 2.0 - SoccerPitchSN.PENALTY_AREA_LENGTH
+        dx = SoccerPitchSN.PENALTY_AREA_LENGTH - SoccerPitchSN.GOAL_LINE_TO_PENALTY_MARK
+        y = np.sqrt(
+            SoccerPitchSN.CENTER_CIRCLE_RADIUS * SoccerPitchSN.CENTER_CIRCLE_RADIUS - dx * dx
+        )
+        self.bottom_right_16M_penalty_arc_mark = np.array([x, y, 0], dtype="float")
+
+        # self.set_elevations(elevation)
+
+        self.point_dict = {}
+        self.point_dict["CENTER_MARK"] = self.center_mark
+        self.point_dict["L_PENALTY_MARK"] = self.left_penalty_mark
+        self.point_dict["R_PENALTY_MARK"] = self.right_penalty_mark
+        self.point_dict["TL_PITCH_CORNER"] = self.top_left_corner
+        self.point_dict["BL_PITCH_CORNER"] = self.bottom_left_corner
+        self.point_dict["TR_PITCH_CORNER"] = self.top_right_corner
+        self.point_dict["BR_PITCH_CORNER"] = self.bottom_right_corner
+        self.point_dict["L_PENALTY_AREA_TL_CORNER"] = self.left_penalty_area_top_left_corner
+        self.point_dict["L_PENALTY_AREA_TR_CORNER"] = self.left_penalty_area_top_right_corner
+        self.point_dict["L_PENALTY_AREA_BL_CORNER"] = self.left_penalty_area_bottom_left_corner
+        self.point_dict["L_PENALTY_AREA_BR_CORNER"] = self.left_penalty_area_bottom_right_corner
+
+        self.point_dict["R_PENALTY_AREA_TL_CORNER"] = self.right_penalty_area_top_left_corner
+        self.point_dict["R_PENALTY_AREA_TR_CORNER"] = self.right_penalty_area_top_right_corner
+        self.point_dict["R_PENALTY_AREA_BL_CORNER"] = self.right_penalty_area_bottom_left_corner
+        self.point_dict["R_PENALTY_AREA_BR_CORNER"] = self.right_penalty_area_bottom_right_corner
+
+        self.point_dict["L_GOAL_AREA_TL_CORNER"] = self.left_goal_area_top_left_corner
+        self.point_dict["L_GOAL_AREA_TR_CORNER"] = self.left_goal_area_top_right_corner
+        self.point_dict["L_GOAL_AREA_BL_CORNER"] = self.left_goal_area_bottom_left_corner
+        self.point_dict["L_GOAL_AREA_BR_CORNER"] = self.left_goal_area_bottom_right_corner
+
+        self.point_dict["R_GOAL_AREA_TL_CORNER"] = self.right_goal_area_top_left_corner
+        self.point_dict["R_GOAL_AREA_TR_CORNER"] = self.right_goal_area_top_right_corner
+        self.point_dict["R_GOAL_AREA_BL_CORNER"] = self.right_goal_area_bottom_left_corner
+        self.point_dict["R_GOAL_AREA_BR_CORNER"] = self.right_goal_area_bottom_right_corner
+
+        self.point_dict["L_GOAL_TL_POST"] = self.left_goal_top_left_post
+        self.point_dict["L_GOAL_TR_POST"] = self.left_goal_top_right_post
+        self.point_dict["L_GOAL_BL_POST"] = self.left_goal_bottom_left_post
+        self.point_dict["L_GOAL_BR_POST"] = self.left_goal_bottom_right_post
+
+        self.point_dict["R_GOAL_TL_POST"] = self.right_goal_top_left_post
+        self.point_dict["R_GOAL_TR_POST"] = self.right_goal_top_right_post
+        self.point_dict["R_GOAL_BL_POST"] = self.right_goal_bottom_left_post
+        self.point_dict["R_GOAL_BR_POST"] = self.right_goal_bottom_right_post
+
+        self.point_dict[
+            "T_TOUCH_AND_HALFWAY_LINES_INTERSECTION"
+        ] = self.halfway_and_top_touch_line_mark
+        self.point_dict[
+            "B_TOUCH_AND_HALFWAY_LINES_INTERSECTION"
+        ] = self.halfway_and_bottom_touch_line_mark
+        self.point_dict[
+            "T_HALFWAY_LINE_AND_CENTER_CIRCLE_INTERSECTION"
+        ] = self.halfway_line_and_center_circle_top_mark
+        self.point_dict[
+            "B_HALFWAY_LINE_AND_CENTER_CIRCLE_INTERSECTION"
+        ] = self.halfway_line_and_center_circle_bottom_mark
+        self.point_dict[
+            "TL_16M_LINE_AND_PENALTY_ARC_INTERSECTION"
+        ] = self.top_left_16M_penalty_arc_mark
+        self.point_dict[
+            "BL_16M_LINE_AND_PENALTY_ARC_INTERSECTION"
+        ] = self.bottom_left_16M_penalty_arc_mark
+        self.point_dict[
+            "TR_16M_LINE_AND_PENALTY_ARC_INTERSECTION"
+        ] = self.top_right_16M_penalty_arc_mark
+        self.point_dict[
+            "BR_16M_LINE_AND_PENALTY_ARC_INTERSECTION"
+        ] = self.bottom_right_16M_penalty_arc_mark
+
+        self.line_extremities = dict()
+        self.line_extremities["Big rect. left bottom"] = (
+            self.point_dict["L_PENALTY_AREA_BL_CORNER"],
+            self.point_dict["L_PENALTY_AREA_BR_CORNER"],
+        )
+        self.line_extremities["Big rect. left top"] = (
+            self.point_dict["L_PENALTY_AREA_TL_CORNER"],
+            self.point_dict["L_PENALTY_AREA_TR_CORNER"],
+        )
+        self.line_extremities["Big rect. left main"] = (
+            self.point_dict["L_PENALTY_AREA_TR_CORNER"],
+            self.point_dict["L_PENALTY_AREA_BR_CORNER"],
+        )
+        self.line_extremities["Big rect. right bottom"] = (
+            self.point_dict["R_PENALTY_AREA_BL_CORNER"],
+            self.point_dict["R_PENALTY_AREA_BR_CORNER"],
+        )
+        self.line_extremities["Big rect. right top"] = (
+            self.point_dict["R_PENALTY_AREA_TL_CORNER"],
+            self.point_dict["R_PENALTY_AREA_TR_CORNER"],
+        )
+        self.line_extremities["Big rect. right main"] = (
+            self.point_dict["R_PENALTY_AREA_TL_CORNER"],
+            self.point_dict["R_PENALTY_AREA_BL_CORNER"],
+        )
+
+        self.line_extremities["Small rect. left bottom"] = (
+            self.point_dict["L_GOAL_AREA_BL_CORNER"],
+            self.point_dict["L_GOAL_AREA_BR_CORNER"],
+        )
+        self.line_extremities["Small rect. left top"] = (
+            self.point_dict["L_GOAL_AREA_TL_CORNER"],
+            self.point_dict["L_GOAL_AREA_TR_CORNER"],
+        )
+        self.line_extremities["Small rect. left main"] = (
+            self.point_dict["L_GOAL_AREA_TR_CORNER"],
+            self.point_dict["L_GOAL_AREA_BR_CORNER"],
+        )
+        self.line_extremities["Small rect. right bottom"] = (
+            self.point_dict["R_GOAL_AREA_BL_CORNER"],
+            self.point_dict["R_GOAL_AREA_BR_CORNER"],
+        )
+        self.line_extremities["Small rect. right top"] = (
+            self.point_dict["R_GOAL_AREA_TL_CORNER"],
+            self.point_dict["R_GOAL_AREA_TR_CORNER"],
+        )
+        self.line_extremities["Small rect. right main"] = (
+            self.point_dict["R_GOAL_AREA_TL_CORNER"],
+            self.point_dict["R_GOAL_AREA_BL_CORNER"],
+        )
+
+        self.line_extremities["Side line top"] = (
+            self.point_dict["TL_PITCH_CORNER"],
+            self.point_dict["TR_PITCH_CORNER"],
+        )
+        self.line_extremities["Side line bottom"] = (
+            self.point_dict["BL_PITCH_CORNER"],
+            self.point_dict["BR_PITCH_CORNER"],
+        )
+        self.line_extremities["Side line left"] = (
+            self.point_dict["TL_PITCH_CORNER"],
+            self.point_dict["BL_PITCH_CORNER"],
+        )
+        self.line_extremities["Side line right"] = (
+            self.point_dict["TR_PITCH_CORNER"],
+            self.point_dict["BR_PITCH_CORNER"],
+        )
+        self.line_extremities["Middle line"] = (
+            self.point_dict["T_TOUCH_AND_HALFWAY_LINES_INTERSECTION"],
+            self.point_dict["B_TOUCH_AND_HALFWAY_LINES_INTERSECTION"],
+        )
+
+        self.line_extremities["Goal left crossbar"] = (
+            self.point_dict["L_GOAL_TR_POST"],
+            self.point_dict["L_GOAL_TL_POST"],
+        )
+        self.line_extremities["Goal left post left "] = (
+            self.point_dict["L_GOAL_TL_POST"],
+            self.point_dict["L_GOAL_BL_POST"],
+        )
+        self.line_extremities["Goal left post right"] = (
+            self.point_dict["L_GOAL_TR_POST"],
+            self.point_dict["L_GOAL_BR_POST"],
+        )
+
+        self.line_extremities["Goal right crossbar"] = (
+            self.point_dict["R_GOAL_TL_POST"],
+            self.point_dict["R_GOAL_TR_POST"],
+        )
+        self.line_extremities["Goal right post left"] = (
+            self.point_dict["R_GOAL_TL_POST"],
+            self.point_dict["R_GOAL_BL_POST"],
+        )
+        self.line_extremities["Goal right post right"] = (
+            self.point_dict["R_GOAL_TR_POST"],
+            self.point_dict["R_GOAL_BR_POST"],
+        )
+        self.line_extremities["Circle right"] = (
+            self.point_dict["TR_16M_LINE_AND_PENALTY_ARC_INTERSECTION"],
+            self.point_dict["BR_16M_LINE_AND_PENALTY_ARC_INTERSECTION"],
+        )
+        self.line_extremities["Circle left"] = (
+            self.point_dict["TL_16M_LINE_AND_PENALTY_ARC_INTERSECTION"],
+            self.point_dict["BL_16M_LINE_AND_PENALTY_ARC_INTERSECTION"],
+        )
+
+        self.line_extremities_keys = dict()
+        self.line_extremities_keys["Big rect. left bottom"] = (
+            "L_PENALTY_AREA_BL_CORNER",
+            "L_PENALTY_AREA_BR_CORNER",
+        )
+        self.line_extremities_keys["Big rect. left top"] = (
+            "L_PENALTY_AREA_TL_CORNER",
+            "L_PENALTY_AREA_TR_CORNER",
+        )
+        self.line_extremities_keys["Big rect. left main"] = (
+            "L_PENALTY_AREA_TR_CORNER",
+            "L_PENALTY_AREA_BR_CORNER",
+        )
+        self.line_extremities_keys["Big rect. right bottom"] = (
+            "R_PENALTY_AREA_BL_CORNER",
+            "R_PENALTY_AREA_BR_CORNER",
+        )
+        self.line_extremities_keys["Big rect. right top"] = (
+            "R_PENALTY_AREA_TL_CORNER",
+            "R_PENALTY_AREA_TR_CORNER",
+        )
+        self.line_extremities_keys["Big rect. right main"] = (
+            "R_PENALTY_AREA_TL_CORNER",
+            "R_PENALTY_AREA_BL_CORNER",
+        )
+
+        self.line_extremities_keys["Small rect. left bottom"] = (
+            "L_GOAL_AREA_BL_CORNER",
+            "L_GOAL_AREA_BR_CORNER",
+        )
+        self.line_extremities_keys["Small rect. left top"] = (
+            "L_GOAL_AREA_TL_CORNER",
+            "L_GOAL_AREA_TR_CORNER",
+        )
+        self.line_extremities_keys["Small rect. left main"] = (
+            "L_GOAL_AREA_TR_CORNER",
+            "L_GOAL_AREA_BR_CORNER",
+        )
+        self.line_extremities_keys["Small rect. right bottom"] = (
+            "R_GOAL_AREA_BL_CORNER",
+            "R_GOAL_AREA_BR_CORNER",
+        )
+        self.line_extremities_keys["Small rect. right top"] = (
+            "R_GOAL_AREA_TL_CORNER",
+            "R_GOAL_AREA_TR_CORNER",
+        )
+        self.line_extremities_keys["Small rect. right main"] = (
+            "R_GOAL_AREA_TL_CORNER",
+            "R_GOAL_AREA_BL_CORNER",
+        )
+
+        self.line_extremities_keys["Side line top"] = ("TL_PITCH_CORNER", "TR_PITCH_CORNER")
+        self.line_extremities_keys["Side line bottom"] = ("BL_PITCH_CORNER", "BR_PITCH_CORNER")
+        self.line_extremities_keys["Side line left"] = ("TL_PITCH_CORNER", "BL_PITCH_CORNER")
+        self.line_extremities_keys["Side line right"] = ("TR_PITCH_CORNER", "BR_PITCH_CORNER")
+        self.line_extremities_keys["Middle line"] = (
+            "T_TOUCH_AND_HALFWAY_LINES_INTERSECTION",
+            "B_TOUCH_AND_HALFWAY_LINES_INTERSECTION",
+        )
+
+        self.line_extremities_keys["Goal left crossbar"] = ("L_GOAL_TR_POST", "L_GOAL_TL_POST")
+        self.line_extremities_keys["Goal left post left "] = ("L_GOAL_TL_POST", "L_GOAL_BL_POST")
+        self.line_extremities_keys["Goal left post right"] = ("L_GOAL_TR_POST", "L_GOAL_BR_POST")
+
+        self.line_extremities_keys["Goal right crossbar"] = ("R_GOAL_TL_POST", "R_GOAL_TR_POST")
+        self.line_extremities_keys["Goal right post left"] = ("R_GOAL_TL_POST", "R_GOAL_BL_POST")
+        self.line_extremities_keys["Goal right post right"] = ("R_GOAL_TR_POST", "R_GOAL_BR_POST")
+        self.line_extremities_keys["Circle right"] = (
+            "TR_16M_LINE_AND_PENALTY_ARC_INTERSECTION",
+            "BR_16M_LINE_AND_PENALTY_ARC_INTERSECTION",
+        )
+        self.line_extremities_keys["Circle left"] = (
+            "TL_16M_LINE_AND_PENALTY_ARC_INTERSECTION",
+            "BL_16M_LINE_AND_PENALTY_ARC_INTERSECTION",
+        )
+
+    def points(self):
+        return [
+            self.center_mark,
+            self.halfway_and_bottom_touch_line_mark,
+            self.halfway_and_top_touch_line_mark,
+            self.halfway_line_and_center_circle_top_mark,
+            self.halfway_line_and_center_circle_bottom_mark,
+            self.bottom_right_corner,
+            self.bottom_left_corner,
+            self.top_right_corner,
+            self.top_left_corner,
+            self.left_penalty_mark,
+            self.right_penalty_mark,
+            self.left_penalty_area_top_right_corner,
+            self.left_penalty_area_top_left_corner,
+            self.left_penalty_area_bottom_right_corner,
+            self.left_penalty_area_bottom_left_corner,
+            self.right_penalty_area_top_right_corner,
+            self.right_penalty_area_top_left_corner,
+            self.right_penalty_area_bottom_right_corner,
+            self.right_penalty_area_bottom_left_corner,
+            self.left_goal_area_top_right_corner,
+            self.left_goal_area_top_left_corner,
+            self.left_goal_area_bottom_right_corner,
+            self.left_goal_area_bottom_left_corner,
+            self.right_goal_area_top_right_corner,
+            self.right_goal_area_top_left_corner,
+            self.right_goal_area_bottom_right_corner,
+            self.right_goal_area_bottom_left_corner,
+            self.top_left_16M_penalty_arc_mark,
+            self.top_right_16M_penalty_arc_mark,
+            self.bottom_left_16M_penalty_arc_mark,
+            self.bottom_right_16M_penalty_arc_mark,
+            self.left_goal_top_left_post,
+            self.left_goal_top_right_post,
+            self.left_goal_bottom_left_post,
+            self.left_goal_bottom_right_post,
+            self.right_goal_top_left_post,
+            self.right_goal_top_right_post,
+            self.right_goal_bottom_left_post,
+            self.right_goal_bottom_right_post,
+        ]
+
+    def sample_field_points(self, dist=0.1, dist_circles=0.2):
+        """
+        Samples each pitch element every dist meters, returns a dictionary associating the class of the element with a list of points sampled along this element.
+        :param dist: the distance in meters between each point sampled
+        :param dist_circles: the distance in meters between each point sampled on circles
+        :return:  a dictionary associating the class of the element with a list of points sampled along this element.
+        """
+        polylines = dict()
+        center = self.point_dict["CENTER_MARK"]
+        fromAngle = 0.0
+        toAngle = 2 * np.pi
+
+        if toAngle < fromAngle:
+            toAngle += 2 * np.pi
+        x1 = center[0] + np.cos(fromAngle) * SoccerPitchSN.CENTER_CIRCLE_RADIUS
+        y1 = center[1] + np.sin(fromAngle) * SoccerPitchSN.CENTER_CIRCLE_RADIUS
+        z1 = 0.0
+        point = np.array((x1, y1, z1))
+        polyline = [point]
+        length = SoccerPitchSN.CENTER_CIRCLE_RADIUS * (toAngle - fromAngle)
+        nb_pts = int(length / dist_circles)
+        dangle = dist_circles / SoccerPitchSN.CENTER_CIRCLE_RADIUS
+        for i in range(1, nb_pts):
+            angle = fromAngle + i * dangle
+            x = center[0] + np.cos(angle) * SoccerPitchSN.CENTER_CIRCLE_RADIUS
+            y = center[1] + np.sin(angle) * SoccerPitchSN.CENTER_CIRCLE_RADIUS
+            z = 0
+            point = np.array((x, y, z))
+            polyline.append(point)
+        polylines["Circle central"] = polyline
+        for key, line in self.line_extremities.items():
+
+            if "Circle" in key:
+                if key == "Circle right":
+                    top = self.point_dict["TR_16M_LINE_AND_PENALTY_ARC_INTERSECTION"]
+                    bottom = self.point_dict["BR_16M_LINE_AND_PENALTY_ARC_INTERSECTION"]
+                    center = self.point_dict["R_PENALTY_MARK"]
+                    toAngle = np.arctan2(top[1] - center[1], top[0] - center[0]) + 2 * np.pi
+                    fromAngle = np.arctan2(bottom[1] - center[1], bottom[0] - center[0]) + 2 * np.pi
+                elif key == "Circle left":
+                    top = self.point_dict["TL_16M_LINE_AND_PENALTY_ARC_INTERSECTION"]
+                    bottom = self.point_dict["BL_16M_LINE_AND_PENALTY_ARC_INTERSECTION"]
+                    center = self.point_dict["L_PENALTY_MARK"]
+                    fromAngle = np.arctan2(top[1] - center[1], top[0] - center[0]) + 2 * np.pi
+                    toAngle = np.arctan2(bottom[1] - center[1], bottom[0] - center[0]) + 2 * np.pi
+                if toAngle < fromAngle:
+                    toAngle += 2 * np.pi
+                x1 = center[0] + np.cos(fromAngle) * SoccerPitchSN.CENTER_CIRCLE_RADIUS
+                y1 = center[1] + np.sin(fromAngle) * SoccerPitchSN.CENTER_CIRCLE_RADIUS
+                z1 = 0.0
+                xn = center[0] + np.cos(toAngle) * SoccerPitchSN.CENTER_CIRCLE_RADIUS
+                yn = center[1] + np.sin(toAngle) * SoccerPitchSN.CENTER_CIRCLE_RADIUS
+                zn = 0.0
+                start = np.array((x1, y1, z1))
+                end = np.array((xn, yn, zn))
+                polyline = [start]
+                length = SoccerPitchSN.CENTER_CIRCLE_RADIUS * (toAngle - fromAngle)
+                nb_pts = int(length / dist_circles)
+                dangle = dist_circles / SoccerPitchSN.CENTER_CIRCLE_RADIUS
+                for i in range(1, nb_pts + 1):
+                    angle = fromAngle + i * dangle
+                    x = center[0] + np.cos(angle) * SoccerPitchSN.CENTER_CIRCLE_RADIUS
+                    y = center[1] + np.sin(angle) * SoccerPitchSN.CENTER_CIRCLE_RADIUS
+                    z = 0
+                    point = np.array((x, y, z))
+                    polyline.append(point)
+                polyline.append(end)
+                polylines[key] = polyline
+            else:
+                start = line[0]
+                end = line[1]
+
+                polyline = [start]
+
+                total_dist = np.sqrt(np.sum(np.square(start - end)))
+                nb_pts = int(total_dist / dist - 1)
+
+                v = end - start
+                v /= np.linalg.norm(v)
+                prev_pt = start
+                for i in range(nb_pts):
+                    pt = prev_pt + dist * v
+                    prev_pt = pt
+                    polyline.append(pt)
+                polyline.append(end)
+                polylines[key] = polyline
+        return polylines
+
+    def get_2d_homogeneous_line(self, line_name):
+        """
+        For lines belonging to the pitch lawn plane returns its 2D homogenous equation coefficients
+        :param line_name
+        :return: an array containing the three coefficients of the line
+        """
+        # ensure line in football pitch plane
+        if (
+            line_name in self.line_extremities.keys()
+            and "post" not in line_name
+            and "crossbar" not in line_name
+            and "Circle" not in line_name
+        ):
+            extremities = self.line_extremities[line_name]
+            p1 = np.array([extremities[0][0], extremities[0][1], 1], dtype="float")
+            p2 = np.array([extremities[1][0], extremities[1][1], 1], dtype="float")
+            line = np.cross(p1, p2)
+
+            return line
+        return None
+
+
+class SoccerPitchSNCircleCentralSplit:
+    """Static class variables that are specified by the rules of the game"""
+
+    GOAL_LINE_TO_PENALTY_MARK = 11.0
+    PENALTY_AREA_WIDTH = 40.32
+    PENALTY_AREA_LENGTH = 16.5
+    GOAL_AREA_WIDTH = 18.32
+    GOAL_AREA_LENGTH = 5.5
+    CENTER_CIRCLE_RADIUS = 9.15
+    GOAL_HEIGHT = 2.44
+    GOAL_LENGTH = 7.32
+
+    lines_classes = [
+        "Big rect. left bottom",
+        "Big rect. left main",
+        "Big rect. left top",
+        "Big rect. right bottom",
+        "Big rect. right main",
+        "Big rect. right top",
+        "Circle central left",
+        "Circle central right",
+        "Circle left",
+        "Circle right",
+        "Goal left crossbar",
+        "Goal left post left ",
+        "Goal left post right",
+        "Goal right crossbar",
+        "Goal right post left",
+        "Goal right post right",
+        "Goal unknown",
+        "Line unknown",
+        "Middle line",
+        "Side line bottom",
+        "Side line left",
+        "Side line right",
+        "Side line top",
+        "Small rect. left bottom",
+        "Small rect. left main",
+        "Small rect. left top",
+        "Small rect. right bottom",
+        "Small rect. right main",
+        "Small rect. right top",
+    ]
+
+    symetric_classes = {
+        "Side line top": "Side line bottom",
+        "Side line bottom": "Side line top",
+        "Side line left": "Side line right",
+        "Middle line": "Middle line",
+        "Side line right": "Side line left",
+        "Big rect. left top": "Big rect. right bottom",
+        "Big rect. left bottom": "Big rect. right top",
+        "Big rect. left main": "Big rect. right main",
+        "Big rect. right top": "Big rect. left bottom",
+        "Big rect. right bottom": "Big rect. left top",
+        "Big rect. right main": "Big rect. left main",
+        "Small rect. left top": "Small rect. right bottom",
+        "Small rect. left bottom": "Small rect. right top",
+        "Small rect. left main": "Small rect. right main",
+        "Small rect. right top": "Small rect. left bottom",
+        "Small rect. right bottom": "Small rect. left top",
+        "Small rect. right main": "Small rect. left main",
+        "Circle left": "Circle right",
+        "Circle central left": "Circle central right",
+        "Circle central right": "Circle central left",
+        "Circle right": "Circle left",
+        "Goal left crossbar": "Goal right crossbar",
+        "Goal left post left ": "Goal right post right",
+        "Goal left post right": "Goal right post left",
+        "Goal right crossbar": "Goal left crossbar",
+        "Goal right post left": "Goal left post right",
+        "Goal right post right": "Goal left post left ",
+        "Goal unknown": "Goal unknown",
+        "Line unknown": "Line unknown",
+    }
+
+    # RGB values
+    palette = {
+        "Big rect. left bottom": (127, 0, 0),
+        "Big rect. left main": (102, 102, 102),
+        "Big rect. left top": (0, 0, 127),
+        "Big rect. right bottom": (86, 32, 39),
+        "Big rect. right main": (48, 77, 0),
+        "Big rect. right top": (14, 97, 100),
+        "Circle central left": (0, 0, 255),
+        "Circle central right": (0, 255, 0),
+        "Circle left": (255, 127, 0),
+        "Circle right": (0, 255, 255),
+        "Goal left crossbar": (255, 255, 200),
+        "Goal left post left ": (165, 255, 0),
+        "Goal left post right": (155, 119, 45),
+        "Goal right crossbar": (86, 32, 139),
+        "Goal right post left": (196, 120, 153),
+        "Goal right post right": (166, 36, 52),
+        "Goal unknown": (0, 0, 0),
+        "Line unknown": (0, 0, 0),
+        "Middle line": (255, 255, 0),
+        "Side line bottom": (255, 0, 255),
+        "Side line left": (0, 255, 150),
+        "Side line right": (0, 230, 0),
+        "Side line top": (230, 0, 0),
+        "Small rect. left bottom": (0, 150, 255),
+        "Small rect. left main": (254, 173, 225),
+        "Small rect. left top": (87, 72, 39),
+        "Small rect. right bottom": (122, 0, 255),
+        "Small rect. right main": (255, 255, 255),
+        "Small rect. right top": (153, 23, 153),
+    }
+
+    def __init__(self, pitch_length=105.0, pitch_width=68.0):
+        """
+        Initialize 3D coordinates of all elements of the soccer pitch.
+        :param pitch_length: According to FIFA rules, length belong to [90,120] meters
+        :param pitch_width: According to FIFA rules, length belong to [45,90] meters
+        """
+        self.PITCH_LENGTH = pitch_length
+        self.PITCH_WIDTH = pitch_width
+
+        self.center_mark = np.array([0, 0, 0], dtype="float")
+        self.halfway_and_bottom_touch_line_mark = np.array([0, pitch_width / 2.0, 0], dtype="float")
+        self.halfway_and_top_touch_line_mark = np.array([0, -pitch_width / 2.0, 0], dtype="float")
+        self.halfway_line_and_center_circle_top_mark = np.array(
+            [0, -SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS, 0], dtype="float"
+        )
+        self.halfway_line_and_center_circle_bottom_mark = np.array(
+            [0, SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS, 0], dtype="float"
+        )
+        self.bottom_right_corner = np.array(
+            [pitch_length / 2.0, pitch_width / 2.0, 0], dtype="float"
+        )
+        self.bottom_left_corner = np.array(
+            [-pitch_length / 2.0, pitch_width / 2.0, 0], dtype="float"
+        )
+        self.top_right_corner = np.array([pitch_length / 2.0, -pitch_width / 2.0, 0], dtype="float")
+        self.top_left_corner = np.array([-pitch_length / 2.0, -34, 0], dtype="float")
+
+        self.left_goal_bottom_left_post = np.array(
+            [-pitch_length / 2.0, SoccerPitchSNCircleCentralSplit.GOAL_LENGTH / 2.0, 0.0],
+            dtype="float",
+        )
+        self.left_goal_top_left_post = np.array(
+            [
+                -pitch_length / 2.0,
+                SoccerPitchSNCircleCentralSplit.GOAL_LENGTH / 2.0,
+                -SoccerPitchSNCircleCentralSplit.GOAL_HEIGHT,
+            ],
+            dtype="float",
+        )
+        self.left_goal_bottom_right_post = np.array(
+            [-pitch_length / 2.0, -SoccerPitchSNCircleCentralSplit.GOAL_LENGTH / 2.0, 0.0],
+            dtype="float",
+        )
+        self.left_goal_top_right_post = np.array(
+            [
+                -pitch_length / 2.0,
+                -SoccerPitchSNCircleCentralSplit.GOAL_LENGTH / 2.0,
+                -SoccerPitchSNCircleCentralSplit.GOAL_HEIGHT,
+            ],
+            dtype="float",
+        )
+
+        self.right_goal_bottom_left_post = np.array(
+            [pitch_length / 2.0, -SoccerPitchSNCircleCentralSplit.GOAL_LENGTH / 2.0, 0.0],
+            dtype="float",
+        )
+        self.right_goal_top_left_post = np.array(
+            [
+                pitch_length / 2.0,
+                -SoccerPitchSNCircleCentralSplit.GOAL_LENGTH / 2.0,
+                -SoccerPitchSNCircleCentralSplit.GOAL_HEIGHT,
+            ],
+            dtype="float",
+        )
+        self.right_goal_bottom_right_post = np.array(
+            [pitch_length / 2.0, SoccerPitchSNCircleCentralSplit.GOAL_LENGTH / 2.0, 0.0],
+            dtype="float",
+        )
+        self.right_goal_top_right_post = np.array(
+            [
+                pitch_length / 2.0,
+                SoccerPitchSNCircleCentralSplit.GOAL_LENGTH / 2.0,
+                -SoccerPitchSNCircleCentralSplit.GOAL_HEIGHT,
+            ],
+            dtype="float",
+        )
+
+        self.left_penalty_mark = np.array(
+            [-pitch_length / 2.0 + SoccerPitchSNCircleCentralSplit.GOAL_LINE_TO_PENALTY_MARK, 0, 0],
+            dtype="float",
+        )
+        self.right_penalty_mark = np.array(
+            [pitch_length / 2.0 - SoccerPitchSNCircleCentralSplit.GOAL_LINE_TO_PENALTY_MARK, 0, 0],
+            dtype="float",
+        )
+
+        self.left_penalty_area_top_right_corner = np.array(
+            [
+                -pitch_length / 2.0 + SoccerPitchSNCircleCentralSplit.PENALTY_AREA_LENGTH,
+                -SoccerPitchSNCircleCentralSplit.PENALTY_AREA_WIDTH / 2.0,
+                0,
+            ],
+            dtype="float",
+        )
+        self.left_penalty_area_top_left_corner = np.array(
+            [-pitch_length / 2.0, -SoccerPitchSNCircleCentralSplit.PENALTY_AREA_WIDTH / 2.0, 0],
+            dtype="float",
+        )
+        self.left_penalty_area_bottom_right_corner = np.array(
+            [
+                -pitch_length / 2.0 + SoccerPitchSNCircleCentralSplit.PENALTY_AREA_LENGTH,
+                SoccerPitchSNCircleCentralSplit.PENALTY_AREA_WIDTH / 2.0,
+                0,
+            ],
+            dtype="float",
+        )
+        self.left_penalty_area_bottom_left_corner = np.array(
+            [-pitch_length / 2.0, SoccerPitchSNCircleCentralSplit.PENALTY_AREA_WIDTH / 2.0, 0],
+            dtype="float",
+        )
+        self.right_penalty_area_top_right_corner = np.array(
+            [pitch_length / 2.0, -SoccerPitchSNCircleCentralSplit.PENALTY_AREA_WIDTH / 2.0, 0],
+            dtype="float",
+        )
+        self.right_penalty_area_top_left_corner = np.array(
+            [
+                pitch_length / 2.0 - SoccerPitchSNCircleCentralSplit.PENALTY_AREA_LENGTH,
+                -SoccerPitchSNCircleCentralSplit.PENALTY_AREA_WIDTH / 2.0,
+                0,
+            ],
+            dtype="float",
+        )
+        self.right_penalty_area_bottom_right_corner = np.array(
+            [pitch_length / 2.0, SoccerPitchSNCircleCentralSplit.PENALTY_AREA_WIDTH / 2.0, 0],
+            dtype="float",
+        )
+        self.right_penalty_area_bottom_left_corner = np.array(
+            [
+                pitch_length / 2.0 - SoccerPitchSNCircleCentralSplit.PENALTY_AREA_LENGTH,
+                SoccerPitchSNCircleCentralSplit.PENALTY_AREA_WIDTH / 2.0,
+                0,
+            ],
+            dtype="float",
+        )
+
+        self.left_goal_area_top_right_corner = np.array(
+            [
+                -pitch_length / 2.0 + SoccerPitchSNCircleCentralSplit.GOAL_AREA_LENGTH,
+                -SoccerPitchSNCircleCentralSplit.GOAL_AREA_WIDTH / 2.0,
+                0,
+            ],
+            dtype="float",
+        )
+        self.left_goal_area_top_left_corner = np.array(
+            [-pitch_length / 2.0, -SoccerPitchSNCircleCentralSplit.GOAL_AREA_WIDTH / 2.0, 0],
+            dtype="float",
+        )
+        self.left_goal_area_bottom_right_corner = np.array(
+            [
+                -pitch_length / 2.0 + SoccerPitchSNCircleCentralSplit.GOAL_AREA_LENGTH,
+                SoccerPitchSNCircleCentralSplit.GOAL_AREA_WIDTH / 2.0,
+                0,
+            ],
+            dtype="float",
+        )
+        self.left_goal_area_bottom_left_corner = np.array(
+            [-pitch_length / 2.0, SoccerPitchSNCircleCentralSplit.GOAL_AREA_WIDTH / 2.0, 0],
+            dtype="float",
+        )
+        self.right_goal_area_top_right_corner = np.array(
+            [pitch_length / 2.0, -SoccerPitchSNCircleCentralSplit.GOAL_AREA_WIDTH / 2.0, 0],
+            dtype="float",
+        )
+        self.right_goal_area_top_left_corner = np.array(
+            [
+                pitch_length / 2.0 - SoccerPitchSNCircleCentralSplit.GOAL_AREA_LENGTH,
+                -SoccerPitchSNCircleCentralSplit.GOAL_AREA_WIDTH / 2.0,
+                0,
+            ],
+            dtype="float",
+        )
+        self.right_goal_area_bottom_right_corner = np.array(
+            [pitch_length / 2.0, SoccerPitchSNCircleCentralSplit.GOAL_AREA_WIDTH / 2.0, 0],
+            dtype="float",
+        )
+        self.right_goal_area_bottom_left_corner = np.array(
+            [
+                pitch_length / 2.0 - SoccerPitchSNCircleCentralSplit.GOAL_AREA_LENGTH,
+                SoccerPitchSNCircleCentralSplit.GOAL_AREA_WIDTH / 2.0,
+                0,
+            ],
+            dtype="float",
+        )
+
+        x = -pitch_length / 2.0 + SoccerPitchSNCircleCentralSplit.PENALTY_AREA_LENGTH
+        dx = (
+            SoccerPitchSNCircleCentralSplit.PENALTY_AREA_LENGTH
+            - SoccerPitchSNCircleCentralSplit.GOAL_LINE_TO_PENALTY_MARK
+        )
+        y = -np.sqrt(
+            SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS
+            * SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS
+            - dx * dx
+        )
+        self.top_left_16M_penalty_arc_mark = np.array([x, y, 0], dtype="float")
+
+        x = pitch_length / 2.0 - SoccerPitchSNCircleCentralSplit.PENALTY_AREA_LENGTH
+        dx = (
+            SoccerPitchSNCircleCentralSplit.PENALTY_AREA_LENGTH
+            - SoccerPitchSNCircleCentralSplit.GOAL_LINE_TO_PENALTY_MARK
+        )
+        y = -np.sqrt(
+            SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS
+            * SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS
+            - dx * dx
+        )
+        self.top_right_16M_penalty_arc_mark = np.array([x, y, 0], dtype="float")
+
+        x = -pitch_length / 2.0 + SoccerPitchSNCircleCentralSplit.PENALTY_AREA_LENGTH
+        dx = (
+            SoccerPitchSNCircleCentralSplit.PENALTY_AREA_LENGTH
+            - SoccerPitchSNCircleCentralSplit.GOAL_LINE_TO_PENALTY_MARK
+        )
+        y = np.sqrt(
+            SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS
+            * SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS
+            - dx * dx
+        )
+        self.bottom_left_16M_penalty_arc_mark = np.array([x, y, 0], dtype="float")
+
+        x = pitch_length / 2.0 - SoccerPitchSNCircleCentralSplit.PENALTY_AREA_LENGTH
+        dx = (
+            SoccerPitchSNCircleCentralSplit.PENALTY_AREA_LENGTH
+            - SoccerPitchSNCircleCentralSplit.GOAL_LINE_TO_PENALTY_MARK
+        )
+        y = np.sqrt(
+            SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS
+            * SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS
+            - dx * dx
+        )
+        self.bottom_right_16M_penalty_arc_mark = np.array([x, y, 0], dtype="float")
+
+        # self.set_elevations(elevation)
+
+        self.point_dict = {}
+        self.point_dict["CENTER_MARK"] = self.center_mark
+        self.point_dict["L_PENALTY_MARK"] = self.left_penalty_mark
+        self.point_dict["R_PENALTY_MARK"] = self.right_penalty_mark
+        self.point_dict["TL_PITCH_CORNER"] = self.top_left_corner
+        self.point_dict["BL_PITCH_CORNER"] = self.bottom_left_corner
+        self.point_dict["TR_PITCH_CORNER"] = self.top_right_corner
+        self.point_dict["BR_PITCH_CORNER"] = self.bottom_right_corner
+        self.point_dict["L_PENALTY_AREA_TL_CORNER"] = self.left_penalty_area_top_left_corner
+        self.point_dict["L_PENALTY_AREA_TR_CORNER"] = self.left_penalty_area_top_right_corner
+        self.point_dict["L_PENALTY_AREA_BL_CORNER"] = self.left_penalty_area_bottom_left_corner
+        self.point_dict["L_PENALTY_AREA_BR_CORNER"] = self.left_penalty_area_bottom_right_corner
+
+        self.point_dict["R_PENALTY_AREA_TL_CORNER"] = self.right_penalty_area_top_left_corner
+        self.point_dict["R_PENALTY_AREA_TR_CORNER"] = self.right_penalty_area_top_right_corner
+        self.point_dict["R_PENALTY_AREA_BL_CORNER"] = self.right_penalty_area_bottom_left_corner
+        self.point_dict["R_PENALTY_AREA_BR_CORNER"] = self.right_penalty_area_bottom_right_corner
+
+        self.point_dict["L_GOAL_AREA_TL_CORNER"] = self.left_goal_area_top_left_corner
+        self.point_dict["L_GOAL_AREA_TR_CORNER"] = self.left_goal_area_top_right_corner
+        self.point_dict["L_GOAL_AREA_BL_CORNER"] = self.left_goal_area_bottom_left_corner
+        self.point_dict["L_GOAL_AREA_BR_CORNER"] = self.left_goal_area_bottom_right_corner
+
+        self.point_dict["R_GOAL_AREA_TL_CORNER"] = self.right_goal_area_top_left_corner
+        self.point_dict["R_GOAL_AREA_TR_CORNER"] = self.right_goal_area_top_right_corner
+        self.point_dict["R_GOAL_AREA_BL_CORNER"] = self.right_goal_area_bottom_left_corner
+        self.point_dict["R_GOAL_AREA_BR_CORNER"] = self.right_goal_area_bottom_right_corner
+
+        self.point_dict["L_GOAL_TL_POST"] = self.left_goal_top_left_post
+        self.point_dict["L_GOAL_TR_POST"] = self.left_goal_top_right_post
+        self.point_dict["L_GOAL_BL_POST"] = self.left_goal_bottom_left_post
+        self.point_dict["L_GOAL_BR_POST"] = self.left_goal_bottom_right_post
+
+        self.point_dict["R_GOAL_TL_POST"] = self.right_goal_top_left_post
+        self.point_dict["R_GOAL_TR_POST"] = self.right_goal_top_right_post
+        self.point_dict["R_GOAL_BL_POST"] = self.right_goal_bottom_left_post
+        self.point_dict["R_GOAL_BR_POST"] = self.right_goal_bottom_right_post
+
+        self.point_dict[
+            "T_TOUCH_AND_HALFWAY_LINES_INTERSECTION"
+        ] = self.halfway_and_top_touch_line_mark
+        self.point_dict[
+            "B_TOUCH_AND_HALFWAY_LINES_INTERSECTION"
+        ] = self.halfway_and_bottom_touch_line_mark
+        self.point_dict[
+            "T_HALFWAY_LINE_AND_CENTER_CIRCLE_INTERSECTION"
+        ] = self.halfway_line_and_center_circle_top_mark
+        self.point_dict[
+            "B_HALFWAY_LINE_AND_CENTER_CIRCLE_INTERSECTION"
+        ] = self.halfway_line_and_center_circle_bottom_mark
+        self.point_dict[
+            "TL_16M_LINE_AND_PENALTY_ARC_INTERSECTION"
+        ] = self.top_left_16M_penalty_arc_mark
+        self.point_dict[
+            "BL_16M_LINE_AND_PENALTY_ARC_INTERSECTION"
+        ] = self.bottom_left_16M_penalty_arc_mark
+        self.point_dict[
+            "TR_16M_LINE_AND_PENALTY_ARC_INTERSECTION"
+        ] = self.top_right_16M_penalty_arc_mark
+        self.point_dict[
+            "BR_16M_LINE_AND_PENALTY_ARC_INTERSECTION"
+        ] = self.bottom_right_16M_penalty_arc_mark
+
+        self.line_extremities = dict()
+        self.line_extremities["Big rect. left bottom"] = (
+            self.point_dict["L_PENALTY_AREA_BL_CORNER"],
+            self.point_dict["L_PENALTY_AREA_BR_CORNER"],
+        )
+        self.line_extremities["Big rect. left top"] = (
+            self.point_dict["L_PENALTY_AREA_TL_CORNER"],
+            self.point_dict["L_PENALTY_AREA_TR_CORNER"],
+        )
+        self.line_extremities["Big rect. left main"] = (
+            self.point_dict["L_PENALTY_AREA_TR_CORNER"],
+            self.point_dict["L_PENALTY_AREA_BR_CORNER"],
+        )
+        self.line_extremities["Big rect. right bottom"] = (
+            self.point_dict["R_PENALTY_AREA_BL_CORNER"],
+            self.point_dict["R_PENALTY_AREA_BR_CORNER"],
+        )
+        self.line_extremities["Big rect. right top"] = (
+            self.point_dict["R_PENALTY_AREA_TL_CORNER"],
+            self.point_dict["R_PENALTY_AREA_TR_CORNER"],
+        )
+        self.line_extremities["Big rect. right main"] = (
+            self.point_dict["R_PENALTY_AREA_TL_CORNER"],
+            self.point_dict["R_PENALTY_AREA_BL_CORNER"],
+        )
+
+        self.line_extremities["Small rect. left bottom"] = (
+            self.point_dict["L_GOAL_AREA_BL_CORNER"],
+            self.point_dict["L_GOAL_AREA_BR_CORNER"],
+        )
+        self.line_extremities["Small rect. left top"] = (
+            self.point_dict["L_GOAL_AREA_TL_CORNER"],
+            self.point_dict["L_GOAL_AREA_TR_CORNER"],
+        )
+        self.line_extremities["Small rect. left main"] = (
+            self.point_dict["L_GOAL_AREA_TR_CORNER"],
+            self.point_dict["L_GOAL_AREA_BR_CORNER"],
+        )
+        self.line_extremities["Small rect. right bottom"] = (
+            self.point_dict["R_GOAL_AREA_BL_CORNER"],
+            self.point_dict["R_GOAL_AREA_BR_CORNER"],
+        )
+        self.line_extremities["Small rect. right top"] = (
+            self.point_dict["R_GOAL_AREA_TL_CORNER"],
+            self.point_dict["R_GOAL_AREA_TR_CORNER"],
+        )
+        self.line_extremities["Small rect. right main"] = (
+            self.point_dict["R_GOAL_AREA_TL_CORNER"],
+            self.point_dict["R_GOAL_AREA_BL_CORNER"],
+        )
+
+        self.line_extremities["Side line top"] = (
+            self.point_dict["TL_PITCH_CORNER"],
+            self.point_dict["TR_PITCH_CORNER"],
+        )
+        self.line_extremities["Side line bottom"] = (
+            self.point_dict["BL_PITCH_CORNER"],
+            self.point_dict["BR_PITCH_CORNER"],
+        )
+        self.line_extremities["Side line left"] = (
+            self.point_dict["TL_PITCH_CORNER"],
+            self.point_dict["BL_PITCH_CORNER"],
+        )
+        self.line_extremities["Side line right"] = (
+            self.point_dict["TR_PITCH_CORNER"],
+            self.point_dict["BR_PITCH_CORNER"],
+        )
+        self.line_extremities["Middle line"] = (
+            self.point_dict["T_TOUCH_AND_HALFWAY_LINES_INTERSECTION"],
+            self.point_dict["B_TOUCH_AND_HALFWAY_LINES_INTERSECTION"],
+        )
+
+        self.line_extremities["Goal left crossbar"] = (
+            self.point_dict["L_GOAL_TR_POST"],
+            self.point_dict["L_GOAL_TL_POST"],
+        )
+        self.line_extremities["Goal left post left "] = (
+            self.point_dict["L_GOAL_TL_POST"],
+            self.point_dict["L_GOAL_BL_POST"],
+        )
+        self.line_extremities["Goal left post right"] = (
+            self.point_dict["L_GOAL_TR_POST"],
+            self.point_dict["L_GOAL_BR_POST"],
+        )
+
+        self.line_extremities["Goal right crossbar"] = (
+            self.point_dict["R_GOAL_TL_POST"],
+            self.point_dict["R_GOAL_TR_POST"],
+        )
+        self.line_extremities["Goal right post left"] = (
+            self.point_dict["R_GOAL_TL_POST"],
+            self.point_dict["R_GOAL_BL_POST"],
+        )
+        self.line_extremities["Goal right post right"] = (
+            self.point_dict["R_GOAL_TR_POST"],
+            self.point_dict["R_GOAL_BR_POST"],
+        )
+        self.line_extremities["Circle right"] = (
+            self.point_dict["TR_16M_LINE_AND_PENALTY_ARC_INTERSECTION"],
+            self.point_dict["BR_16M_LINE_AND_PENALTY_ARC_INTERSECTION"],
+        )
+        self.line_extremities["Circle left"] = (
+            self.point_dict["TL_16M_LINE_AND_PENALTY_ARC_INTERSECTION"],
+            self.point_dict["BL_16M_LINE_AND_PENALTY_ARC_INTERSECTION"],
+        )
+
+        self.line_extremities_keys = dict()
+        self.line_extremities_keys["Big rect. left bottom"] = (
+            "L_PENALTY_AREA_BL_CORNER",
+            "L_PENALTY_AREA_BR_CORNER",
+        )
+        self.line_extremities_keys["Big rect. left top"] = (
+            "L_PENALTY_AREA_TL_CORNER",
+            "L_PENALTY_AREA_TR_CORNER",
+        )
+        self.line_extremities_keys["Big rect. left main"] = (
+            "L_PENALTY_AREA_TR_CORNER",
+            "L_PENALTY_AREA_BR_CORNER",
+        )
+        self.line_extremities_keys["Big rect. right bottom"] = (
+            "R_PENALTY_AREA_BL_CORNER",
+            "R_PENALTY_AREA_BR_CORNER",
+        )
+        self.line_extremities_keys["Big rect. right top"] = (
+            "R_PENALTY_AREA_TL_CORNER",
+            "R_PENALTY_AREA_TR_CORNER",
+        )
+        self.line_extremities_keys["Big rect. right main"] = (
+            "R_PENALTY_AREA_TL_CORNER",
+            "R_PENALTY_AREA_BL_CORNER",
+        )
+
+        self.line_extremities_keys["Small rect. left bottom"] = (
+            "L_GOAL_AREA_BL_CORNER",
+            "L_GOAL_AREA_BR_CORNER",
+        )
+        self.line_extremities_keys["Small rect. left top"] = (
+            "L_GOAL_AREA_TL_CORNER",
+            "L_GOAL_AREA_TR_CORNER",
+        )
+        self.line_extremities_keys["Small rect. left main"] = (
+            "L_GOAL_AREA_TR_CORNER",
+            "L_GOAL_AREA_BR_CORNER",
+        )
+        self.line_extremities_keys["Small rect. right bottom"] = (
+            "R_GOAL_AREA_BL_CORNER",
+            "R_GOAL_AREA_BR_CORNER",
+        )
+        self.line_extremities_keys["Small rect. right top"] = (
+            "R_GOAL_AREA_TL_CORNER",
+            "R_GOAL_AREA_TR_CORNER",
+        )
+        self.line_extremities_keys["Small rect. right main"] = (
+            "R_GOAL_AREA_TL_CORNER",
+            "R_GOAL_AREA_BL_CORNER",
+        )
+
+        self.line_extremities_keys["Side line top"] = ("TL_PITCH_CORNER", "TR_PITCH_CORNER")
+        self.line_extremities_keys["Side line bottom"] = ("BL_PITCH_CORNER", "BR_PITCH_CORNER")
+        self.line_extremities_keys["Side line left"] = ("TL_PITCH_CORNER", "BL_PITCH_CORNER")
+        self.line_extremities_keys["Side line right"] = ("TR_PITCH_CORNER", "BR_PITCH_CORNER")
+        self.line_extremities_keys["Middle line"] = (
+            "T_TOUCH_AND_HALFWAY_LINES_INTERSECTION",
+            "B_TOUCH_AND_HALFWAY_LINES_INTERSECTION",
+        )
+
+        self.line_extremities_keys["Goal left crossbar"] = ("L_GOAL_TR_POST", "L_GOAL_TL_POST")
+        self.line_extremities_keys["Goal left post left "] = ("L_GOAL_TL_POST", "L_GOAL_BL_POST")
+        self.line_extremities_keys["Goal left post right"] = ("L_GOAL_TR_POST", "L_GOAL_BR_POST")
+
+        self.line_extremities_keys["Goal right crossbar"] = ("R_GOAL_TL_POST", "R_GOAL_TR_POST")
+        self.line_extremities_keys["Goal right post left"] = ("R_GOAL_TL_POST", "R_GOAL_BL_POST")
+        self.line_extremities_keys["Goal right post right"] = ("R_GOAL_TR_POST", "R_GOAL_BR_POST")
+        self.line_extremities_keys["Circle right"] = (
+            "TR_16M_LINE_AND_PENALTY_ARC_INTERSECTION",
+            "BR_16M_LINE_AND_PENALTY_ARC_INTERSECTION",
+        )
+        self.line_extremities_keys["Circle left"] = (
+            "TL_16M_LINE_AND_PENALTY_ARC_INTERSECTION",
+            "BL_16M_LINE_AND_PENALTY_ARC_INTERSECTION",
+        )
+
+    def points(self):
+        return [
+            self.center_mark,
+            self.halfway_and_bottom_touch_line_mark,
+            self.halfway_and_top_touch_line_mark,
+            self.halfway_line_and_center_circle_top_mark,
+            self.halfway_line_and_center_circle_bottom_mark,
+            self.bottom_right_corner,
+            self.bottom_left_corner,
+            self.top_right_corner,
+            self.top_left_corner,
+            self.left_penalty_mark,
+            self.right_penalty_mark,
+            self.left_penalty_area_top_right_corner,
+            self.left_penalty_area_top_left_corner,
+            self.left_penalty_area_bottom_right_corner,
+            self.left_penalty_area_bottom_left_corner,
+            self.right_penalty_area_top_right_corner,
+            self.right_penalty_area_top_left_corner,
+            self.right_penalty_area_bottom_right_corner,
+            self.right_penalty_area_bottom_left_corner,
+            self.left_goal_area_top_right_corner,
+            self.left_goal_area_top_left_corner,
+            self.left_goal_area_bottom_right_corner,
+            self.left_goal_area_bottom_left_corner,
+            self.right_goal_area_top_right_corner,
+            self.right_goal_area_top_left_corner,
+            self.right_goal_area_bottom_right_corner,
+            self.right_goal_area_bottom_left_corner,
+            self.top_left_16M_penalty_arc_mark,
+            self.top_right_16M_penalty_arc_mark,
+            self.bottom_left_16M_penalty_arc_mark,
+            self.bottom_right_16M_penalty_arc_mark,
+            self.left_goal_top_left_post,
+            self.left_goal_top_right_post,
+            self.left_goal_bottom_left_post,
+            self.left_goal_bottom_right_post,
+            self.right_goal_top_left_post,
+            self.right_goal_top_right_post,
+            self.right_goal_bottom_left_post,
+            self.right_goal_bottom_right_post,
+        ]
+
+    def sample_field_points(self, dist=0.1, dist_circles=0.2):
+        """
+        Samples each pitch element every dist meters, returns a dictionary associating the class of the element with a list of points sampled along this element.
+        :param dist: the distance in meters between each point sampled
+        :param dist_circles: the distance in meters between each point sampled on circles
+        :return:  a dictionary associating the class of the element with a list of points sampled along this element.
+        """
+        polylines = dict()
+        center = self.point_dict["CENTER_MARK"]
+        fromAngle = 0.0
+        toAngle = 2 * np.pi
+
+        if toAngle < fromAngle:
+            toAngle += 2 * np.pi
+        x1 = center[0] + np.cos(fromAngle) * SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS
+        y1 = center[1] + np.sin(fromAngle) * SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS
+        z1 = 0.0
+        point = np.array((x1, y1, z1))
+        polyline = [point]
+        length = SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS * (toAngle - fromAngle)
+        nb_pts = int(length / dist_circles)
+        dangle = dist_circles / SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS
+        for i in range(1, nb_pts):
+            angle = fromAngle + i * dangle
+            x = center[0] + np.cos(angle) * SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS
+            y = center[1] + np.sin(angle) * SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS
+            z = 0
+            point = np.array((x, y, z))
+            polyline.append(point)
+
+        # split central circle in left and right
+        polylines["Circle central left"] = [p for p in polyline if p[0] < 0.0]
+        polylines["Circle central right"] = [p for p in polyline if p[0] >= 0.0]
+        for key, line in self.line_extremities.items():
+
+            if "Circle" in key:
+                if key == "Circle right":
+                    top = self.point_dict["TR_16M_LINE_AND_PENALTY_ARC_INTERSECTION"]
+                    bottom = self.point_dict["BR_16M_LINE_AND_PENALTY_ARC_INTERSECTION"]
+                    center = self.point_dict["R_PENALTY_MARK"]
+                    toAngle = np.arctan2(top[1] - center[1], top[0] - center[0]) + 2 * np.pi
+                    fromAngle = np.arctan2(bottom[1] - center[1], bottom[0] - center[0]) + 2 * np.pi
+                elif key == "Circle left":
+                    top = self.point_dict["TL_16M_LINE_AND_PENALTY_ARC_INTERSECTION"]
+                    bottom = self.point_dict["BL_16M_LINE_AND_PENALTY_ARC_INTERSECTION"]
+                    center = self.point_dict["L_PENALTY_MARK"]
+                    fromAngle = np.arctan2(top[1] - center[1], top[0] - center[0]) + 2 * np.pi
+                    toAngle = np.arctan2(bottom[1] - center[1], bottom[0] - center[0]) + 2 * np.pi
+                if toAngle < fromAngle:
+                    toAngle += 2 * np.pi
+                x1 = (
+                    center[0]
+                    + np.cos(fromAngle) * SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS
+                )
+                y1 = (
+                    center[1]
+                    + np.sin(fromAngle) * SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS
+                )
+                z1 = 0.0
+                xn = (
+                    center[0]
+                    + np.cos(toAngle) * SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS
+                )
+                yn = (
+                    center[1]
+                    + np.sin(toAngle) * SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS
+                )
+                zn = 0.0
+                start = np.array((x1, y1, z1))
+                end = np.array((xn, yn, zn))
+                polyline = [start]
+                length = SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS * (
+                    toAngle - fromAngle
+                )
+                nb_pts = int(length / dist_circles)
+                dangle = dist_circles / SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS
+                for i in range(1, nb_pts + 1):
+                    angle = fromAngle + i * dangle
+                    x = (
+                        center[0]
+                        + np.cos(angle) * SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS
+                    )
+                    y = (
+                        center[1]
+                        + np.sin(angle) * SoccerPitchSNCircleCentralSplit.CENTER_CIRCLE_RADIUS
+                    )
+                    z = 0
+                    point = np.array((x, y, z))
+                    polyline.append(point)
+                polyline.append(end)
+                polylines[key] = polyline
+            else:
+                start = line[0]
+                end = line[1]
+
+                polyline = [start]
+
+                total_dist = np.sqrt(np.sum(np.square(start - end)))
+                nb_pts = int(total_dist / dist - 1)
+
+                v = end - start
+                v /= np.linalg.norm(v)
+                prev_pt = start
+                for i in range(nb_pts):
+                    pt = prev_pt + dist * v
+                    prev_pt = pt
+                    polyline.append(pt)
+                polyline.append(end)
+                polylines[key] = polyline
+        return polylines
+
+    def get_2d_homogeneous_line(self, line_name):
+        """
+        For lines belonging to the pitch lawn plane returns its 2D homogenous equation coefficients
+        :param line_name
+        :return: an array containing the three coefficients of the line
+        """
+        # ensure line in football pitch plane
+        if (
+            line_name in self.line_extremities.keys()
+            and "post" not in line_name
+            and "crossbar" not in line_name
+            and "Circle" not in line_name
+        ):
+            extremities = self.line_extremities[line_name]
+            p1 = np.array([extremities[0][0], extremities[0][1], 1], dtype="float")
+            p2 = np.array([extremities[1][0], extremities[1][1], 1], dtype="float")
+            line = np.cross(p1, p2)
+
+            return line
+        return None
+
+
+class Abstract3dModel(metaclass=ABCMeta):
+    def __init__(self) -> None:
+
+        self.points = None  # keypoints: tensor of shape (N, 3)
+        self.points_sampled = (
+            None  # sampled points for each segment Dict[str: torch.tensor of shape (*, 3)]
+        )
+        self.points_sampled_palette = {}
+        self.segment_names = set(self.points_sampled_palette.keys())
+
+        self.line_segments = []  # tensor of shape (3, S_l, 2) containing 2 points
+        self.line_segments_names = []  # list of of respective names for each s in S_l
+        self.line_palette = []  # list of RGB tuples
+
+        self.circle_segments = None  # tensor of shape (3, S_c, num_points_per_circle)
+        self.circle_segments_names = []  # list of of respective names for each s in S_c
+        self.circle_palette = []  # list of RGB tuples
+
+
+class Meshgrid(Abstract3dModel):
+    def __init__(self, height=68, width=105):
+        self.points = kornia.utils.create_meshgrid(
+            height=height + 1, width=width + 1, normalized_coordinates=True
+        )
+        self.points = self.points.flatten(start_dim=-3, end_dim=-2)
+        self.points[:, :, 0] = self.points[:, :, 0] * width / 2
+        self.points[:, :, 1] = self.points[:, :, 1] * height / 2
+        self.points = kornia.geometry.conversions.convert_points_to_homogeneous(self.points)
+        self.points[:, :, -1] = 0.0  # set z=0)
+        self.points = self.points.squeeze(0)
+        self.points_sampled = {"meshgrid": self.points}
+
+
+class SoccerPitchLineCircleSegments(Abstract3dModel):
+    def __init__(
+        self,
+        base_field,
+        device="cpu",
+        N_cstar=128,
+        sampling_factor_lines=0.2,
+        sampling_factor_circles=0.8,
+    ) -> None:
+
+        if not (
+            isinstance(base_field, SoccerPitchSNCircleCentralSplit)
+            or isinstance(base_field, SoccerPitchSN)
+        ):
+            raise NotImplementedError
+
+        self.sampling_factor_lines = sampling_factor_lines
+        self.sampling_factor_circles = sampling_factor_circles
+
+        self._field_sncalib = base_field
+
+        self.device = device
+
+        # classical keypoints as single tensor
+        self.points = torch.from_numpy(np.stack(self._field_sncalib.points())).float().to(device)
+
+        # sampled points for each segment Dict[str: torch.tensor of shape (*, 3)]
+        self.points_sampled = self._field_sncalib.sample_field_points(
+            self.sampling_factor_lines, self.sampling_factor_circles
+        )
+        self.points_sampled = {
+            k: torch.from_numpy(np.stack(v)).float().to(device)
+            for k, v in self.points_sampled.items()
+        }
+        self.points_sampled_palette = self._field_sncalib.palette
+        self.segment_names = set(self.points_sampled_palette.keys())
+        self.cmap_01 = {k: [c / 255.0 for c in v] for k, v in self.points_sampled_palette.items()}
+
+        self.line_collection: List[LineCollection] = []
+        self.line_segments = []  # (3, S, 2)
+        self.line_segments_names = []
+
+        for line_name, (p0, p1) in self._field_sncalib.line_extremities.items():
+            if "Circle" not in line_name:
+                p0 = torch.from_numpy(p0).float().to(device)
+                p1 = torch.from_numpy(p1).float().to(device)
+                direction = p1 - p0
+                direction_norm = direction / torch.linalg.norm(direction)
+                self.line_collection.append(
+                    LineCollection(
+                        support=p0,
+                        direction=direction,
+                        direction_norm=direction_norm,
+                    )
+                )
+                self.line_segments_names.append(line_name)
+                self.line_segments.append(torch.stack([p0, p1], dim=1))
+
+        self.line_segments = torch.stack(self.line_segments, dim=-1).transpose(1, 2).to(device)
+        self.line_palette = [
+            self._field_sncalib.palette[self.line_segments_names[i]]
+            for i in range(len(self.line_segments_names))
+        ]
+
+        if isinstance(base_field, SoccerPitchSNCircleCentralSplit):
+            self.circle_segments_names = [
+                "Circle central left",
+                "Circle central right",
+                "Circle left",
+                "Circle right",
+            ]
+        elif isinstance(base_field, SoccerPitchSN):
+            self.circle_segments_names = [
+                "Circle central",
+                "Circle left",
+                "Circle right",
+            ]
+        else:
+            raise NotImplementedError
+
+        self.circle_segments = self._sample_points_from_circle_segments(
+            m=N_cstar
+        )  # (3, num_circles, num_points_per_circle)
+
+        self.circle_palette = [
+            self._field_sncalib.palette[self.circle_segments_names[i]]
+            for i in range(len(self.circle_segments_names))
+        ]
+
+    def _sample_points_from_circle_segments(self, m: int):
+
+        sampled_points = self._field_sncalib.sample_field_points(dist=1.0, dist_circles=0.05)
+        for key in self.circle_segments_names:
+            assert len(sampled_points[key]) >= m
+        return (
+            torch.stack(
+                [
+                    torch.from_numpy(np.stack(random.sample(sampled_points[key], k=m), axis=-1))
+                    for key in self.circle_segments_names
+                ],
+                dim=1,
+            )
+            .float()
+            .to(self.device)
+        )  # (3, S, m)
+
+
+if __name__ == "__main__":
+    from matplotlib import pyplot as plt
+    from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+
+    model3d = SoccerPitchLineCircleSegments()
+
+    fig = plt.figure(figsize=(20, 20))
+    ax = fig.add_subplot(projection="3d")
+    for s in range(len(model3d.line_collection)):
+        ax.quiver(
+            model3d.line_collection[s].support[0],
+            model3d.line_collection[s].support[1],
+            model3d.line_collection[s].support[2],
+            model3d.line_collection[s].direction[0],
+            model3d.line_collection[s].direction[1],
+            model3d.line_collection[s].direction[2],
+            arrow_length_ratio=0.05,
+            color=[x / 255.0 for x in model3d.line_palette[s]],
+            zorder=2000,
+            # length=68.0,
+            linewidths=3,
+            label=model3d.line_segments_names[s],
+            alpha=0.5,
+        )
+
+    plt.legend()
+    ax.set_xlim([-105 / 2, 105 / 2])
+    ax.set_ylim([-105 / 2, 105 / 2])
+    ax.set_zlim([-105 / 2, 105 / 2])
+    plt.show()
+
+    fig = plt.figure(figsize=(20, 20))
+    ax = fig.add_subplot(projection="3d")
+    for segment_name, sampled_points in model3d.points_sampled.items():
+        ax.scatter(
+            sampled_points[:, 0],
+            sampled_points[:, 1],
+            -sampled_points[:, 2],
+            zorder=3000,
+            color=[x / 255.0 for x in model3d.points_sampled_palette[segment_name]],
+            marker="x",
+            label=segment_name,
+        )
+
+    plt.legend()
+    ax.set_xlim([-105 / 2, 105 / 2])
+    ax.set_ylim([-105 / 2, 105 / 2])
+    ax.set_zlim([-105 / 2, 105 / 2])
+    plt.show()
+
+    fig = plt.figure(figsize=(20, 20))
+    ax = fig.add_subplot(projection="3d")
+    for s in range(model3d.line_segments.shape[1]):
+
+        if "crossbar" in model3d.line_segments_names[s]:
+            print(s, model3d.line_segments[:, s])
+            ax.scatter(
+                model3d.line_segments[0, s],
+                model3d.line_segments[1, s],
+                -model3d.line_segments[2, s],
+                zorder=3000,
+                color=[x / 255.0 for x in model3d.line_palette[s]],
+                marker="x",
+                label=model3d.line_segments_names[s],
+            )
+
+    plt.legend()
+    ax.set_xlim([-105 / 2, 105 / 2])
+    ax.set_ylim([-105 / 2, 105 / 2])
+    ax.set_zlim([-105 / 2, 105 / 2])
+    plt.savefig("soccer_field_line_segments.pdf")

visualizer.py

@@ -0,0 +1,298 @@
+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
+
+# 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, 
+                                         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.
+    
+    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.
+        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)
+
+    # 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
+
+    return minimap_bgr, average_draw_scale # Retourner aussi l'échelle moyenne
+
+# 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