update dashboard 6

.gitignore

@@ -11,6 +11,8 @@ __pycache__/
 *.pptx
 .dockerignore
 app_v0.py
+app_v1.py
 inference_v0.py
+inference_v1.py
 Dockerfile
 Notes.txt

app.py

@@ -14,6 +14,7 @@ ALL_MODEL_FILES = {
     "Modèle EfficientnetB0 Baseline": "efficientnet_fire.pt",
     "Modèle EfficientnetB0 FE": "efficientnet_fire_2.pt",
     "Modèle EfficientnetB0 FT": "efficientnet_fire_3.pt",
+    "Modèle YOLOv8": "yolov8_fire.pt",
     "Modèle Inception3": "inception3_fire.pt",
 }
 
@@ -55,7 +56,22 @@ MODEL_METRICS = {
         "fp": 440,
         "fn": 680,
     },
-    # "inception3_fire.pt": {...},
+    # "yolov8_fire.pt": {
+    #    "accuracy": 0.7717,
+    #    "precision": 0.8713,
+    #    "recall": 0.8142,
+    #    "f1": 0.8418,
+    #    "fp": 440,
+    #    "fn": 680,
+    # },
+    # "inception3_fire.pt": {
+    #    "accuracy": 0.7717,
+    #    "precision": 0.8713,
+    #    "recall": 0.8142,
+    #    "f1": 0.8418,
+    #    "fp": 440,
+    #    "fn": 680,
+    # },
 }
 
 # -----------------------------------------------------------
@@ -164,13 +180,12 @@ else:
     # Afficher l'image uploadée
     # -------------------------------------------------------
     image = Image.open(uploaded_file)
-    st.image(image, caption="Image chargée", use_container_width=True)
 
     # -------------------------------------------------------
     # Prédiction
     # -------------------------------------------------------
     with st.spinner("Analyse de l'image en cours..."):
-        label, prob = predict_from_pil(
+        label, prob, annotated_image = predict_from_pil(
             image=image,
             model=model,
             device=device,
@@ -178,6 +193,17 @@ else:
             threshold=threshold
         )
 
+    # Image originale
+    st.image(image, caption="Image chargée", use_container_width=True)
+
+    # Si YOLO a fourni une image annotée, on l'affiche
+    if annotated_image is not None:
+        st.image(
+            annotated_image,
+            caption="Zones détectées (modèle de détection)",
+            use_container_width=True,
+        )
+
     # -------------------------------------------------------
     # Affichage du résultat avec couleur
     # -------------------------------------------------------

inference.py

@@ -3,59 +3,51 @@ inference.py
 ------------
 Module d'inférence pour des modèles de classification binaire : FIRE (1) / NO_FIRE (0).
 
-Compatible avec plusieurs architectures (ex. EfficientNet-B0, Inception v3),
-sans changer le code dès qu'on ajoute un nouveau fichier .pt, à condition :
+Supporte :
+- EfficientNet-B0 (classification)
+- Inception v3 (classification)
+- YOLO (Ultralytics, détection) pour localiser le feu
 
-- que le nom du fichier permette de deviner le type de modèle
-  (ex : "inception3_fire.pt" → Inception v3,
-        sinon → EfficientNet-B0 par défaut),
-- ou que le fichier .pt contienne un dictionnaire avec "state_dict" (et
-  éventuellement "model_name" / "model_key").
-
-Usage typique :
----------------
-from inference import load_model, get_val_transform, predict_from_path
-
-model, device = load_model("efficientnet_fire.pt")
-transform = get_val_transform()  # ou get_val_transform(image_size_personnalisée)
-
-label, prob = predict_from_path("mon_image.jpg", model, device, transform)
-print(label, prob)
+Retour principal des fonctions de prédiction :
+    predicted_label (str), fire_prob (float), annotated_image (PIL.Image ou None)
 """
 
 # ----------------------------
 # 1) Imports
 # ----------------------------
-import os                               # pour gérer les chemins de fichiers
-import torch                            # bibliothèque principale pour le deep learning
-import torch.nn as nn                   # pour définir les têtes de classification
-from torchvision import transforms      # pour les pré-traitements d'images
-from PIL import Image                   # pour charger les images depuis un fichier
-import timm                             # pour charger des architectures (EfficientNet, Inception, etc.)
+import os
+
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from PIL import Image
+import timm
+
+# YOLO (Ultralytics) – optionnel
+try:
+    from ultralytics import YOLO
+except ImportError:
+    YOLO = None  # si la lib n'est pas installée, on gère ça proprement
 
 
 # ----------------------------
 # 2) Constantes globales
 # ----------------------------
 
-# Taille d'entrée par défaut (utilisée si on ne précise rien d'autre).
-DEFAULT_IMAGE_SIZE = 224  # (224 x 224 pixels)
+# Taille d'entrée par défaut si rien n'est précisé
+DEFAULT_IMAGE_SIZE = 224
 
-# Moyennes et écarts-types d'ImageNet (pour normaliser les images)
-IMAGENET_MEAN = [0.485, 0.456, 0.406]   # moyenne des canaux R, G, B
-IMAGENET_STD = [0.229, 0.224, 0.225]    # écart-type des canaux R, G, B
+# Stats ImageNet
+IMAGENET_MEAN = [0.485, 0.456, 0.406]
+IMAGENET_STD = [0.229, 0.224, 0.225]
 
-# Mapping des classes numériques vers des labels lisibles
+# Mapping idx -> label lisible
 IDX_TO_LABEL = {
-    0: "no_fire",   # classe 0 → pas de feu
-    1: "fire"       # classe 1 → feu
+    0: "no_fire",
+    1: "fire",
 }
 
-# Registre des modèles supportés.
-# - model_key : identifiant interne (nos clés)
-# - timm_name : nom utilisé dans timm.create_model(...)
-# - image_size : taille d'entrée "recommandée" pour ce modèle (optionnelle)
-# - classifier_attr : nom de l'attribut contenant la dernière couche de classification
+# Registre des modèles de classification (timm)
 MODEL_REGISTRY = {
     "efficientnet_b0": {
         "timm_name": "efficientnet_b0",
@@ -69,12 +61,17 @@ MODEL_REGISTRY = {
     },
 }
 
-# Si on ne reconnaît pas de modèle particulier, on utilisera EfficientNet-B0 par défaut.
+# Modèle par défaut si on ne sait pas quoi choisir
 DEFAULT_MODEL_KEY = "efficientnet_b0"
 
+# IDs des classes "feu" pour YOLO (à adapter si besoin après entraînement)
+# Exemple : si model.names == {0: 'fire'} → [0]
+#          si model.names == {0: 'no_fire', 1: 'fire'} → [1]
+FIRE_CLASS_IDS = [0]
+
 
 # ----------------------------
-# 3) Utilitaires device
+# 3) Device
 # ----------------------------
 
 def get_device():
@@ -85,8 +82,7 @@ def get_device():
     """
     if torch.cuda.is_available():
         return torch.device("cuda")
-    else:
-        return torch.device("cpu")
+    return torch.device("cpu")
 
 
 # ----------------------------
@@ -96,62 +92,58 @@ def get_device():
 def infer_model_key_from_path(weights_path: str) -> str:
     """
     Devine une clé de modèle (model_key) à partir du nom de fichier.
+
     Exemple :
-    - "inception3_fire.pt"  → "inception_v3"
-    - "efficientnet_fire.pt" → "efficientnet_b0" (par défaut)
+    - "yolov8_fire.pt"        → "yolo"
+    - "inception3_fire.pt"    → "inception_v3"
+    - "efficientnet_fire.pt"  → "efficientnet_b0" (par défaut)
     """
     filename = os.path.basename(weights_path).lower()
 
-    # Règle simple : si "inception" est dans le nom → Inception v3
+    if "yolo" in filename:
+        return "yolo"
+
     if "inception" in filename:
         return "inception_v3"
 
-    # Sinon, par défaut → EfficientNet-B0
     return DEFAULT_MODEL_KEY
 
 
 def get_model_config(model_key: str) -> dict:
     """
     Renvoie la config du modèle à partir d'une model_key.
-    Si model_key n'est pas connue, on retourne la config du modèle par défaut.
+    Fallback : EfficientNet-B0 si model_key inconnue.
     """
     if model_key in MODEL_REGISTRY:
         return MODEL_REGISTRY[model_key]
-    # Fallback de sécurité : ne jamais casser si clé inconnue
     return MODEL_REGISTRY[DEFAULT_MODEL_KEY]
 
 
 # ----------------------------
-# 5) Construction de l'architecture
+# 5) Construction modèle (classification)
 # ----------------------------
 
 def build_model(model_key: str, num_classes: int = 2) -> torch.nn.Module:
     """
-    Construit l'architecture correspondant à model_key
-    et adapte la tête de classification à num_classes sorties.
+    Construit un modèle de classification (EfficientNet, Inception...)
+    et adapte la dernière couche à num_classes sorties.
     """
     config = get_model_config(model_key)
     timm_name = config["timm_name"]
     classifier_attr = config["classifier_attr"]
 
-    # On crée le backbone via timm (sans pré-entraînement, les poids viendront du .pt)
     model = timm.create_model(timm_name, pretrained=False)
 
-    # On adapte la tête de classification à notre problème binaire.
-    # On récupère l'ancienne couche de classification (fc, classifier, etc.)
     classifier = getattr(model, classifier_attr)
 
-    # Certains modèles ont déjà un nn.Linear, on récupère in_features
     if isinstance(classifier, nn.Linear):
         in_features = classifier.in_features
     else:
-        # Cas plus exotique : on essaie de deviner proprement, sinon on lève une erreur claire.
         raise ValueError(
             f"Impossible de déterminer in_features pour la tête du modèle '{timm_name}'. "
             f"Attribut '{classifier_attr}' de type {type(classifier)} non supporté."
         )
 
-    # On remplace par une couche linéaire adaptée à notre nombre de classes
     new_classifier = nn.Linear(in_features, num_classes)
     setattr(model, classifier_attr, new_classifier)
 
@@ -159,89 +151,78 @@ def build_model(model_key: str, num_classes: int = 2) -> torch.nn.Module:
 
 
 # ----------------------------
-# 6) Chargement du modèle
+# 6) Nettoyage de state_dict
 # ----------------------------
 
 def _clean_state_dict_keys(state_dict: dict) -> dict:
     """
-    Nettoie les clés d'un state_dict pour gérer plusieurs cas courants :
-    - clés préfixées par 'model.' (Lightning)
-    - clés préfixées par 'module.' (DataParallel)
+    Nettoie les clés pour gérer les prefixes 'model.' (Lightning) et 'module.' (DataParallel).
     """
-    new_state_dict = {}
+    new_state = {}
     for k, v in state_dict.items():
         new_key = k
         if new_key.startswith("model."):
             new_key = new_key[len("model."):]
         if new_key.startswith("module."):
             new_key = new_key[len("module."):]
-        new_state_dict[new_key] = v
-    return new_state_dict
+        new_state[new_key] = v
+    return new_state
 
 
+# ----------------------------
+# 7) Chargement du modèle
+# ----------------------------
+
 def load_model(weights_path: str, map_location=None, model_key: str | None = None):
     """
     Charge un modèle avec les poids entraînés.
 
-    Paramètres
-    ----------
-    weights_path : str
-        Chemin vers le fichier .pt (state_dict ou dict avec 'state_dict').
-    map_location : torch.device ou None
-        Device sur lequel charger les poids. Si None, on détecte automatiquement.
-    model_key : str ou None
-        Clé de modèle à utiliser (ex: 'efficientnet_b0', 'inception_v3').
-        Si None, on essaie de la déduire du nom de fichier.
+    - Pour YOLO (Ultralytics) : charge un modèle de détection.
+    - Pour EfficientNet / Inception : charge un modèle de classification binaire.
 
-    Retour
-    ------
-    model : torch.nn.Module
-        Le modèle prêt pour l'inférence.
+    Retour :
+    --------
+    model : torch.nn.Module ou YOLO
     device : torch.device
-        Le device utilisé (cuda ou cpu).
     """
-    # 1) Sélection du device
     device = map_location if map_location is not None else get_device()
 
-    # 2) Si la model_key n'est pas donnée, on essaye de l'inférer depuis le nom de fichier
+    # Détecter le type de modèle si non fourni
     if model_key is None:
         model_key = infer_model_key_from_path(weights_path)
 
-    # 3) Chargement brut du .pt
+    # 🔹 Cas YOLO : modèle de détection (Ultralytics)
+    if model_key == "yolo":
+        if YOLO is None:
+            raise ImportError(
+                "Le modèle YOLO est demandé mais la librairie 'ultralytics' "
+                "n'est pas installée. Ajoutez 'ultralytics' dans requirements.txt."
+            )
+        yolo_model = YOLO(weights_path)
+        # YOLO gère déjà souvent le device en interne, on essaye juste par sécurité
+        try:
+            yolo_model.to(device)
+        except Exception:
+            pass
+        return yolo_model, device
+
+    # 🔹 Cas classification (EfficientNet, Inception...)
     checkpoint = torch.load(weights_path, map_location=device)
 
-    # 4) Si le checkpoint est un dict complet (ex: {'state_dict': ..., 'model_name': ...})
     if isinstance(checkpoint, dict) and "state_dict" in checkpoint:
         state_dict = checkpoint["state_dict"]
-
-        # Optionnel : si un model_name ou model_key est stocké dedans, on peut le préférer.
-        if "model_key" in checkpoint and checkpoint["model_key"] in MODEL_REGISTRY:
-            model_key = checkpoint["model_key"]
-        elif "model_name" in checkpoint:
-            # Tentative : si model_name correspond au timm_name d'une entrée du registre
-            model_name_lower = str(checkpoint["model_name"]).lower()
-            for k, cfg in MODEL_REGISTRY.items():
-                if cfg["timm_name"].lower() == model_name_lower:
-                    model_key = k
-                    break
     else:
-        # Cas simple : le fichier .pt est directement un state_dict
         state_dict = checkpoint
 
-    # 5) Nettoyage des clés du state_dict (Lightning, DataParallel...)
     state_dict = _clean_state_dict_keys(state_dict)
 
-    # 6) Construction de l'architecture adaptée
     model = build_model(model_key=model_key, num_classes=2)
-
-    # 7) Application des poids (strict=False pour tolérer quelques petites diff de clés)
     missing, unexpected = model.load_state_dict(state_dict, strict=False)
 
-    # (Optionnel) DEBUG : on pourrait afficher missing / unexpected si besoin
+    # (optionnel) debug :
     # print("Missing keys:", missing)
     # print("Unexpected keys:", unexpected)
 
-    # 8) Envoi sur le bon device + mode eval
     model = model.to(device)
     model.eval()
 
@@ -249,21 +230,14 @@ def load_model(weights_path: str, map_location=None, model_key: str | None = Non
 
 
 # ----------------------------
-# 7) Transforms pour l'inférence
+# 8) Transforms pour l'inférence
 # ----------------------------
 
 def get_val_transform(image_size: int | None = None):
     """
     Renvoie les transformations à appliquer aux images pour l'inférence.
-    Paramètres
-    ----------
-    image_size : int ou None
-        Si None, on utilise DEFAULT_IMAGE_SIZE (224).
-        Sinon, on redimensionne en (image_size, image_size).
 
-    Retour
-    ------
-    transform : torchvision.transforms.Compose
+    Si image_size est None → DEFAULT_IMAGE_SIZE (224).
     """
     if image_size is None:
         image_size = DEFAULT_IMAGE_SIZE
@@ -277,26 +251,12 @@ def get_val_transform(image_size: int | None = None):
 
 
 # ----------------------------
-# 8) Prétraitement d'une image
+# 9) Prétraitement d'une image
 # ----------------------------
 
 def preprocess_image(image: Image.Image, transform=None, image_size: int | None = None):
     """
     Applique les transforms à une image PIL et ajoute une dimension batch.
-
-    Paramètres
-    ----------
-    image : PIL.Image.Image
-        Image brute chargée (par exemple via Image.open(...)).
-    transform : callable ou None
-        Transformations à appliquer (si None, on utilise get_val_transform(image_size)).
-    image_size : int ou None
-        Taille de redimensionnement si transform est None.
-
-    Retour
-    ------
-    image_tensor : torch.Tensor
-        Tenseur prêt pour l'inférence, de taille [1, 3, H, W].
     """
     if transform is None:
         transform = get_val_transform(image_size=image_size)
@@ -308,57 +268,39 @@ def preprocess_image(image: Image.Image, transform=None, image_size: int | None
 
 
 # ----------------------------
-# 9) Fonction de prédiction principale
+# 10) Prédiction depuis un tenseur (classification)
 # ----------------------------
 
 def predict_from_tensor(
     image_tensor: torch.Tensor,
     model: torch.nn.Module,
     device: torch.device,
-    threshold: float = 0.5
+    threshold: float = 0.5,
 ):
     """
-    Prédit la classe (fire/no_fire) à partir d'un tenseur déjà prétraité.
-
-    Paramètres
-    ----------
-    image_tensor : torch.Tensor
-        Tenseur d'images de taille [1, 3, H, W] (batch de 1 image).
-    model : torch.nn.Module
-        Modèle chargé (EfficientNet, Inception, etc.).
-    device : torch.device
-        Device sur lequel le modèle est (cuda ou cpu).
-    threshold : float
-        Seuil sur la probabilité de FEU pour décider entre no_fire / fire.
-
-    Retour
-    ------
-    predicted_label : str
-        "fire" ou "no_fire".
-    fire_prob : float
-        Probabilité prédite pour la classe "fire" (entre 0 et 1).
+    Prédit la classe (fire/no_fire) à partir d'un tenseur déjà prétraité
+    pour les modèles de classification (EfficientNet, Inception...).
     """
     image_tensor = image_tensor.to(device)
 
     with torch.no_grad():
         outputs = model(image_tensor)           # logits [1, 2]
-        probs = torch.softmax(outputs, dim=1)   # probabilités
+        probs = torch.softmax(outputs, dim=1)   # probas
 
         fire_prob = probs[0, 1].item()
-
-        if fire_prob >= threshold:
-            predicted_idx = 1  # feu
-        else:
-            predicted_idx = 0  # pas de feu
+        predicted_idx = 1 if fire_prob >= threshold else 0
 
     predicted_label = IDX_TO_LABEL[predicted_idx]
-
     return predicted_label, fire_prob
 
 
+# ----------------------------
+# 11) Prédiction depuis une image PIL
+# ----------------------------
+
 def predict_from_pil(
     image: Image.Image,
-    model: torch.nn.Module,
+    model,
     device: torch.device,
     transform=None,
     threshold: float = 0.5,
@@ -367,39 +309,66 @@ def predict_from_pil(
     """
     Prédit la classe à partir d'une image PIL.
 
-    Paramètres
-    ----------
-    image : PIL.Image.Image
-        Image chargée (par exemple via Image.open).
-    model : torch.nn.Module
-        Modèle chargé.
-    device : torch.device
-        Device (cuda ou cpu).
-    transform : callable ou None
-        Transformations à appliquer à l'image.
-    threshold : float
-        Seuil sur la probabilité de FEU.
-    image_size : int ou None
-        Taille utilisée si transform est None.
-
     Retour
     ------
     predicted_label : str
         "fire" ou "no_fire".
     fire_prob : float
         Probabilité de "fire".
+    annotated_image : PIL.Image.Image ou None
+        Image annotée (bounding boxes) pour les modèles de détection (YOLO),
+        None pour les modèles de classification.
     """
     if image.mode != "RGB":
         image = image.convert("RGB")
 
+    # 🔹 Cas YOLO : modèle de détection (Ultralytics)
+    if hasattr(model, "task") and getattr(model, "task", None) == "detect":
+        results = model(image)
+        result = results[0]
+
+        boxes = getattr(result, "boxes", None)
+        fire_prob = 0.0
+
+        if boxes is not None and len(boxes) > 0:
+            classes = boxes.cls  # ids des classes (tensor)
+            confs = boxes.conf   # scores de confiance (tensor)
+
+            # masque des boxes "feu"
+            mask_fire = torch.zeros_like(classes, dtype=torch.bool)
+            for cid in FIRE_CLASS_IDS:
+                mask_fire |= (classes == cid)
+
+            if mask_fire.any():
+                fire_prob = float(confs[mask_fire].max().item())
+
+        predicted_label = "fire" if fire_prob >= threshold else "no_fire"
+
+        # Image annotée avec les bounding boxes
+        annotated_image = None
+        try:
+            annotated_array = result.plot()          # numpy array BGR
+            annotated_image = Image.fromarray(annotated_array[..., ::-1])  # BGR -> RGB
+        except Exception:
+            annotated_image = None
+
+        return predicted_label, fire_prob, annotated_image
+
+    # 🔹 Cas classification classique
     image_tensor = preprocess_image(image, transform=transform, image_size=image_size)
+    predicted_label, fire_prob = predict_from_tensor(image_tensor, model, device, threshold=threshold)
+    annotated_image = None
+
+    return predicted_label, fire_prob, annotated_image
 
-    return predict_from_tensor(image_tensor, model, device, threshold=threshold)
 
+# ----------------------------
+# 12) Prédiction depuis un chemin de fichier
+# ----------------------------
 
 def predict_from_path(
     image_path: str,
-    model: torch.nn.Module,
+    model,
     device: torch.device,
     transform=None,
     threshold: float = 0.5,
@@ -407,32 +376,10 @@ def predict_from_path(
 ):
     """
     Prédit la classe à partir d'un chemin vers une image.
-
-    Paramètres
-    ----------
-    image_path : str
-        Chemin vers le fichier image (jpg, png, etc.).
-    model : torch.nn.Module
-        Modèle chargé.
-    device : torch.device
-        Device (cuda ou cpu).
-    transform : callable ou None
-        Transformations à appliquer.
-    threshold : float
-        Seuil sur la probabilité de FEU.
-    image_size : int ou None
-        Taille utilisée si transform est None.
-
-    Retour
-    ------
-    predicted_label : str
-        "fire" ou "no_fire".
-    fire_prob : float
-        Probabilité de "fire".
     """
     image = Image.open(image_path)
     return predict_from_pil(
-        image,
+        image=image,
         model=model,
         device=device,
         transform=transform,
@@ -442,31 +389,25 @@ def predict_from_path(
 
 
 # ----------------------------
-# 10) Exemple d'utilisation en script direct
+# 13) Exemple d'utilisation en script direct
 # ----------------------------
 
 if __name__ == "__main__":
-    """
-    Exemple simple pour tester le module en local :
-    python inference.py
-    """
-    weights_path = "efficientnet_fire.pt"  # à adapter si besoin
-
+    # Petit test local (à adapter)
+    weights_path = "efficientnet_fire.pt"
     if not os.path.exists(weights_path):
         print(f"[ERREUR] Fichier de poids introuvable : {weights_path}")
     else:
         model, device = load_model(weights_path)
         print(f"Modèle chargé sur le device : {device}")
 
-        # On utilise la taille par défaut (224) pour ce test
         transform = get_val_transform()
-
-        test_image_path = "example.jpg"  # à adapter si besoin
+        test_image_path = "example.jpg"
 
         if not os.path.exists(test_image_path):
             print(f"[INFO] Aucune image test trouvée à : {test_image_path}")
         else:
-            label, prob = predict_from_path(
+            label, prob, _ = predict_from_path(
                 test_image_path,
                 model=model,
                 device=device,