Add DocString

app.py

@@ -37,9 +37,9 @@ id2label = {
 }
 label2id = {v: k for k, v in id2label.items()}
 num_labels = len(id2label)
-checkpoint = "nvidia/segformer-b4-finetuned-cityscapes-1024-1024"
-image_processor = SegformerImageProcessor()
-state_dict_path = f"runs/{checkpoint}_v1/best_model.pt"
+checkpoint = "nvidia/segformer-b3-finetuned-cityscapes-1024-1024"
+image_processor = SegformerImageProcessor(do_resize=False)
+state_dict_path = f"runs/{checkpoint}/best_model.pt"
 model = SegformerForSemanticSegmentation.from_pretrained(
     checkpoint,
     num_labels=num_labels,
@@ -58,6 +58,17 @@ model.eval()
 
 
 def load_and_prepare_images(image_name, segformer=False):
+    """
+    Charge et prépare les images, les masques et les prédictions associées pour une image donnée.
+
+    Args:
+        image_name (str): Le nom du fichier de l'image à charger.
+        segformer (bool, optional): Si True, prédit également le masque avec SegFormer. Par défaut False.
+
+    Returns:
+        tuple: Contient l'image originale redimensionnée, le masque réel, la prédiction FPN,
+               et la prédiction SegFormer si `segformer` est True.
+    """
     image_path = os.path.join(data_folder, "images", image_name)
     mask_name = image_name.replace("_leftImg8bit.png", "_gtFine_labelIds.png")
     mask_path = os.path.join(data_folder, "masks", mask_name)
@@ -82,35 +93,47 @@ def load_and_prepare_images(image_name, segformer=False):
 
 
 def predict_segmentation(image):
-    # Charger et préparer l'image
-    inputs = image_processor(images=image, return_tensors="pt")
+    """
+    Prédit la segmentation d'une image donnée à l'aide d'un modèle pré-entraîné.
 
-    # Utiliser GPU si disponible
+    Args:
+        image (PIL.Image.Image): L'image à segmenter.
+
+    Returns:
+        numpy.ndarray: La carte de segmentation prédite.
+    """
+
+    inputs = image_processor(images=image, return_tensors="pt")
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
-
-    # Déplacer les inputs sur le bon device et faire la prédiction
     pixel_values = inputs.pixel_values.to(device)
 
-    with torch.no_grad():  # Désactiver le calcul des gradients pour l'inférence
+    with torch.no_grad():
         outputs = model(pixel_values=pixel_values)
         logits = outputs.logits
 
-    # Redimensionner les logits à la taille de l'image d'origine
     upsampled_logits = nn.functional.interpolate(
         logits,
         size=image.size[::-1],  # (height, width)
         mode="bilinear",
         align_corners=False,
     )
-
-    # Obtenir la prédiction finale
     pred_seg = upsampled_logits.argmax(dim=1)[0].cpu().numpy()
 
     return pred_seg
 
 
 def process_image(image_name):
+    """
+    Traite une image en chargeant l'image originale, le masque réel, et les prédictions de masques.
+    Envoie la liste de tuple à l'interface "Predictions" de Gradio
+
+    Args:
+        image_name (str): Le nom de l'image à traiter.
+
+    Returns:
+        list: Une liste de tuples contenant l'image et son titre associé.
+    """
     original, true_mask, fpn_pred, segformer_pred = load_and_prepare_images(
         image_name, segformer=True
     )
@@ -131,6 +154,12 @@ def process_image(image_name):
 
 
 def create_cityscapes_label_colormap():
+    """
+    Crée une colormap pour les labels Cityscapes.
+
+    Returns:
+        numpy.ndarray: Un tableau 2D où chaque ligne représente la couleur RGB d'un label.
+    """
     colormap = np.zeros((256, 3), dtype=np.uint8)
     colormap[0] = [78, 82, 110]
     colormap[1] = [128, 64, 128]
@@ -147,68 +176,43 @@ def create_cityscapes_label_colormap():
 cityscapes_colormap = create_cityscapes_label_colormap()
 
 
-def blend_images(original_image, colored_segmentation, alpha=0.6):
-    blended_image = Image.blend(original_image, colored_segmentation, alpha)
-    return blended_image
-
-
 def colorize_mask(mask):
     return cityscapes_colormap[mask]
 
 
 # ---- Fin Partie Segmentation
 
-# def compare_masks(real_mask, fpn_mask, segformer_mask):
-#     """
-#     Compare les masques prédits par FPN et SegFormer avec le masque réel.
-#     Retourne un score IoU et une précision pixel par pixel pour chaque modèle.
-
-#     Args:
-#     real_mask (np.array): Le masque réel de référence
-#     fpn_mask (np.array): Le masque prédit par le modèle FPN
-#     segformer_mask (np.array): Le masque prédit par le modèle SegFormer
-
-#     Returns:
-#     dict: Dictionnaire contenant les scores IoU et les précisions pour chaque modèle
-#     """
-
-#     assert real_mask.shape == fpn_mask.shape == segformer_mask.shape, "Les masques doivent avoir la même forme"
-
-#     real_flat = real_mask.flatten()
-#     fpn_flat = fpn_mask.flatten()
-#     segformer_flat = segformer_mask.flatten()
-
-#     # Calcul du score de Jaccard (IoU)
-#     iou_fpn = jaccard_score(real_flat, fpn_flat, average='weighted')
-#     iou_segformer = jaccard_score(real_flat, segformer_flat, average='weighted')
-
-#     # Calcul de la précision pixel par pixel
-#     accuracy_fpn = accuracy_score(real_flat, fpn_flat)
-#     accuracy_segformer = accuracy_score(real_flat, segformer_flat)
-
-#     return {
-#         'FPN': {'IoU': iou_fpn, 'Precision': accuracy_fpn},
-#         'SegFormer': {'IoU': iou_segformer, 'Precision': accuracy_segformer}
-#     }
-
 # ---- Partie EDA
 
 
 def analyse_mask(real_mask, num_labels):
-    # Compter les occurrences de chaque classe
-    counts = np.bincount(real_mask.ravel(), minlength=num_labels)
+    """
+    Analyse la distribution des classes dans un masque réel.
 
-    # Calculer le nombre total de pixels
-    total_pixels = real_mask.size
+    Args:
+        real_mask (numpy.ndarray): Le masque de labels réels.
+        num_labels (int): Le nombre total de classes.
 
-    # Calculer les proportions
+    Returns:
+        dict: Un dictionnaire contenant les proportions des classes dans le masque.
+    """
+    counts = np.bincount(real_mask.ravel(), minlength=num_labels)
+    total_pixels = real_mask.size
     class_proportions = counts / total_pixels
-
-    # Créer un dictionnaire avec les proportions
     return dict(enumerate(class_proportions))
 
 
 def show_eda(image_name):
+    """
+    Affiche une analyse exploratoire de la distribution des classes pour une image et son masque associé.
+
+    Args:
+        image_name (str): Le nom de l'image à analyser.
+
+    Returns:
+        tuple: Contient l'image originale, le masque réel coloré et une figure Plotly représentant
+               la distribution des classes.
+    """
     original_image, true_mask, _ = load_and_prepare_images(image_name)
     class_proportions = analyse_mask(true_mask, num_labels)
     cityscapes_colormap = create_cityscapes_label_colormap()
@@ -266,17 +270,54 @@ def show_eda(image_name):
 
 
 class SegformerWrapper(nn.Module):
+    """
+    Un wrapper pour le modèle SegFormer qui renvoie uniquement les logits en sortie.
+
+    Args:
+        model (torch.nn.Module): Le modèle SegFormer pré-entraîné.
+    """
+
     def __init__(self, model):
+        """
+        Initialise le SegformerWrapper.
+
+        Args:
+            model (torch.nn.Module): Le modèle SegFormer pré-entraîné.
+        """
         super().__init__()
         self.model = model
 
     def forward(self, x):
+        """
+        Renvoie les logits du modèle au lieu de renvoyer un dictionnaire.
+
+        Args:
+            x (torch.Tensor): Les entrées du modèle.
+
+        Returns:
+            torch.Tensor: Les logits du modèle.
+        """
         output = self.model(x)
         return output.logits
 
 
 class SemanticSegmentationTarget:
+    """
+    Représente une classe cible pour la segmentation sémantique utilisée dans GradCAM.
+
+    Args:
+        category (int): L'index de la catégorie cible.
+        mask (numpy.ndarray): Le masque binaire indiquant les pixels d'intérêt.
+    """
+
     def __init__(self, category, mask):
+        """
+        Initialise la cible de segmentation sémantique.
+
+        Args:
+            category (int): L'index de la catégorie cible.
+            mask (numpy.ndarray): Le masque binaire indiquant les pixels d'intérêt.
+        """
         self.category = category
         self.mask = torch.from_numpy(mask)
         if torch.cuda.is_available():
@@ -305,12 +346,33 @@ class SemanticSegmentationTarget:
 
 
 def segformer_reshape_transform_huggingface(tensor, width, height):
+    """
+    Réorganise les dimensions du tenseur pour qu'elles correspondent au format attendu par GradCAM.
+
+    Args:
+        tensor (torch.Tensor): Le tenseur à réorganiser.
+        width (int): La nouvelle largeur.
+        height (int): La nouvelle hauteur.
+
+    Returns:
+        torch.Tensor: Le tenseur réorganisé.
+    """
     result = tensor.reshape(tensor.size(0), height, width, tensor.size(2))
     result = result.transpose(2, 3).transpose(1, 2)
     return result
 
 
 def explain_model(image_name, category_name):
+    """
+    Explique les prédictions du modèle SegFormer en utilisant GradCAM pour une image et une catégorie données.
+
+    Args:
+        image_name (str): Le nom de l'image à expliquer.
+        category_name (str): Le nom de la catégorie cible.
+
+    Returns:
+        matplotlib.figure.Figure: Une figure matplotlib contenant la carte de chaleur GradCAM superposée sur l'image originale.
+    """
     original_image, _, _ = load_and_prepare_images(image_name)
     rgb_img = np.float32(original_image) / 255
     img_tensor = transforms.ToTensor()(rgb_img)
@@ -379,6 +441,12 @@ import random
 
 
 def change_image():
+    """
+    Sélectionne et charge aléatoirement une image depuis un dossier spécifié.
+
+    Returns:
+        PIL.Image.Image: L'image sélectionnée.
+    """
     image_dir = (
         "data_sample/images"  # Remplacez par le chemin de votre dossier d'images
     )
@@ -388,6 +456,16 @@ def change_image():
 
 
 def apply_augmentation(image, augmentation_names):
+    """
+    Applique une ou plusieurs augmentations à une image.
+
+    Args:
+        image (PIL.Image.Image): L'image à augmenter.
+        augmentation_names (list of str): Les noms des augmentations à appliquer.
+
+    Returns:
+        PIL.Image.Image: L'image augmentée.
+    """
     augmentations = {
         "Horizontal Flip": A.HorizontalFlip(p=1),
         "Shift Scale Rotate": A.ShiftScaleRotate(p=1),
@@ -541,4 +619,4 @@ with gr.Blocks(title="Preuve de concept", theme=my_theme) as demo:
 
 
 # Lancer l'application
-demo.launch(favicon_path="favicon.ico", share=True)
+demo.launch(favicon_path="favicon.ico")