From cdd6f9645bb7272b23b92cce02a84d768be48041 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gr=C3=A9goire=20Mureau?= <gregoire.mureau+pro@etik.com>
Date: Tue, 2 Sep 2025 10:46:13 +0200
Subject: [PATCH] Docker working. Improvements incoming.

---
 EV3                        |   0
 app/main.py                | 419 ++++++++++-----------
 app/requirements.txt       | 117 ++----
 app/streamlit_app.py       | 749 +++++++++++++++----------------------
 notebooks/requirements.txt |   6 +-
 utils/utils.py             |  20 +-
 6 files changed, 530 insertions(+), 781 deletions(-)
 create mode 100644 EV3

diff --git a/EV3 b/EV3
new file mode 100644
index 0000000..e69de29
diff --git a/app/main.py b/app/main.py
index 3487456..136db01 100644
--- a/app/main.py
+++ b/app/main.py
@@ -1,24 +1,27 @@
 import os
 import sys
+import traceback
 import tempfile
 import shutil
 from datetime import datetime
 from typing import Optional, List
 from pathlib import Path
 import mlflow
-from fastapi import FastAPI, HTTPException, File, UploadFile, Form
+from fastapi import FastAPI, HTTPException, File, UploadFile
 from fastapi.responses import FileResponse
 from pydantic import BaseModel
 import numpy as np
 import pickle
 import base64
+import matplotlib.pyplot as plt
+import matplotlib.image as mpimg
+
 
 # Configuration MLflow avec variables d'environnement
 def configure_mlflow():
     """Configure MLflow avec les variables d'environnement"""
     print("=== DEBUG VARIABLES D'ENVIRONNEMENT ===")
     
-    # ✅ NETTOYAGE automatique des variables d'environnement
     mlflow_tracking_uri = os.getenv("MLFLOW_TRACKING_URI", "").strip()
     mlflow_s3_endpoint_url = os.getenv("MLFLOW_S3_ENDPOINT_URL", "").strip()
     aws_access_key_id = os.getenv("AWS_ACCESS_KEY_ID", "").strip()
@@ -33,11 +36,9 @@ def configure_mlflow():
     if not mlflow_tracking_uri:
         raise ValueError("MLFLOW_TRACKING_URI not found in environment variables")
     
-    # Configuration MLflow avec URLs nettoyées
     mlflow.set_tracking_uri(mlflow_tracking_uri)
     print(f"MLflow Tracking URI: {mlflow_tracking_uri}")
     
-    # ✅ CORRECTION: Configuration EXPLICITE des identifiants avec nettoyage
     if not mlflow_s3_endpoint_url or not aws_access_key_id or not aws_secret_access_key:
         raise ValueError("Variables d'environnement MLflow manquantes (S3_ENDPOINT, ACCESS_KEY, SECRET_KEY)")
     
@@ -46,66 +47,68 @@ def configure_mlflow():
     os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key
     print(f"MLflow S3 Endpoint: {mlflow_s3_endpoint_url}")
     print("✅ Identifiants AWS configurés")
-    
-    # ✅ SKIP set_experiment pour éviter le problème 404
-    # mlflow.set_experiment("OC Projet 8")  # Commenté temporairement
-    print("✅ Configuration MLflow terminée (sans set_experiment)")
+    print("✅ Configuration MLflow terminée")
 
 # Import des fonctions utils
 sys.path.append('.')
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 from utils.utils import (
     load_cityscapes_config, 
-    load_best_model_from_registry,
+    load_model,
     run_inference_and_visualize,
     colorize_mask
 )
 
-# ✅ AJOUT: Import segmentation_models pour le chargement des modèles
-try:
-    import segmentation_models as sm
-    print("✅ segmentation_models importé - modèles compatibles")
-except ImportError as e:
-    print(f"⚠️ segmentation_models non disponible: {e}")
-    print("💡 Installez avec: pip install git+https://github.com/qubvel/segmentation_models.git")
-
 # Configuration
 MODEL_NAME = "StreetsSegmentation"
 SAMPLE_IMAGES_DIR = "./notebooks/content/data/test_images_sample"
 CITYSCAPES_CONFIG_PATH = "./cityscapes_config.json"
 TEMP_UPLOAD_DIR = "/tmp/uploads"
 
-# Modèles de données
+# Modèles de données mis à jour
+class ModelInfo(BaseModel):
+    name: str
+    run_id: str
+    encoder_name: str
+    input_size: tuple
+    rank: int
+
+class ComparisonResult(BaseModel):
+    success: bool
+    message: str
+    model1_info: ModelInfo
+    model2_info: ModelInfo
+    image_path: Optional[str] = None
+    mask_path: Optional[str] = None
+    model1_stats: Optional[dict] = None
+    model2_stats: Optional[dict] = None
+    model1_inference_time: float
+    model2_inference_time: float
+    speed_comparison: dict
+    ground_truth_available: bool = False
+    figure_data: Optional[str] = None
+
 class ImageInfo(BaseModel):
     filename: str
     display_name: str
     has_ground_truth: bool
 
-class PredictionResult(BaseModel):
-    success: bool
-    message: str
-    image_path: Optional[str] = None
-    mask_path: Optional[str] = None
-    prediction_stats: Optional[dict] = None
-    ground_truth_available: bool = False
-    figure_data: Optional[str] = None
-
 class AvailableImagesResponse(BaseModel):
     images: List[ImageInfo]
     total_count: int
 
-# Variables globales pour le modèle
-model = None
-encoder_name = None
-script_img_size = None
+# Variables globales pour les deux modèles
+best_model = None
+second_best_model = None
+best_model_info = {}
+second_best_model_info = {}
 mapping_config = None
-run_id = None
 
-def initialize_model():
-    """Initialise le modèle et la configuration au démarrage"""
-    global model, encoder_name, script_img_size, mapping_config, run_id
+def initialize_models():
+    """Initialise les deux meilleurs modèles au démarrage"""
+    global best_model, second_best_model, best_model_info, second_best_model_info, mapping_config
     
     try:
-        # Configuration MLflow
         print("🔄 Configuration de MLflow...")
         configure_mlflow()
         
@@ -113,76 +116,83 @@ def initialize_model():
         mapping_config = load_cityscapes_config(CITYSCAPES_CONFIG_PATH, verbose=False)
         print("✅ Configuration Cityscapes chargée")
         
-        print("🔄 Chargement du modèle depuis MLflow...")
-        print(f"🔍 Recherche du modèle: {MODEL_NAME}")
-        model, run_id, encoder_name, script_img_size = load_best_model_from_registry(MODEL_NAME)
-        print("✅ Modèle téléchargé depuis MLflow")
+        print("🔄 Chargement du meilleur modèle...")
+        best_model, best_run_id, best_encoder_name, best_img_size = load_model(
+            experiment_name="OC Projet 9", 
+            metric="test_mean_iou", 
+            top_n=1
+        )
+        best_model_info = {
+            "name": "Meilleur modèle",
+            "run_id": best_run_id,
+            "encoder_name": best_encoder_name,
+            "input_size": best_img_size,
+            "rank": 1
+        }
         
-        print(f"✅ Modèle chargé du run {run_id}")
-        print(f"📐 Encoder: {encoder_name}")
-        print(f"📐 Input shape: {script_img_size}")
-        print(f"🎨 Nombre de classes: {mapping_config['num_classes']}")
+        print("🔄 Chargement du deuxième meilleur modèle...")
+        second_best_model, second_run_id, second_encoder_name, second_img_size = load_model(
+            experiment_name="OC Projet 9", 
+            metric="test_mean_iou", 
+            top_n=2
+        )
+        second_best_model_info = {
+            "name": "Deuxième modèle",
+            "run_id": second_run_id,
+            "encoder_name": second_encoder_name,
+            "input_size": second_img_size,
+            "rank": 2
+        }
         
-        # Créer le dossier temporaire
         os.makedirs(TEMP_UPLOAD_DIR, exist_ok=True)
         
     except Exception as e:
         print(f"❌ Erreur lors de l'initialisation: {e}")
-        import traceback
         traceback.print_exc()
         raise e
 
 # Initialisation de l'API
 app = FastAPI(
-    title="Multi-Class Segmentation API",
-    description="API de segmentation sémantique pour véhicules autonomes",
-    version="1.0.0"
+    title="Multi-Class Segmentation Comparison API",
+    description="API de comparaison de segmentation sémantique avec deux modèles",
+    version="2.0.0"
 )
 
-# ✅ AJOUT: Initialiser le modèle de manière plus robuste
 @app.on_event("startup")
 async def startup_event():
-    """Initialise le modèle au démarrage de l'API"""
     try:
-        initialize_model()
-        print("🚀 API démarrée avec succès - Modèle chargé!")
+        initialize_models()
+        print("🚀 API démarrée avec succès - Modèles chargés!")
     except Exception as e:
-        print(f"⚠️ ATTENTION: Impossible de charger le modèle au démarrage: {e}")
-        print("L'API démarre quand même mais les prédictions ne fonctionneront pas.")
+        print(f"⚠️ ATTENTION: Impossible de charger les modèles au démarrage: {e}")
 
 @app.get("/")
 def root():
     return {
-        "message": "Multi-Class Segmentation API - Online",
-        "model_loaded": model is not None,
-        "model_run_id": run_id,
-        "encoder": encoder_name,
+        "message": "Multi-Class Segmentation Comparison API - Online",
+        "models_loaded": best_model is not None and second_best_model is not None,
+        "best_model": best_model_info if best_model else None,
+        "second_best_model": second_best_model_info if second_best_model else None,
         "num_classes": mapping_config['num_classes'] if mapping_config else None
     }
 
 @app.get("/health")
 def health_check():
-    """Health check endpoint pour nginx"""
-    if model is None or mapping_config is None:
-        raise HTTPException(status_code=503, detail="Model not loaded")
-    
+    if best_model is None or second_best_model is None or mapping_config is None:
+        raise HTTPException(status_code=503, detail="Models not loaded")
     return {
         "status": "healthy",
-        "model_loaded": True,
+        "models_loaded": True,
         "timestamp": datetime.utcnow().isoformat()
     }
 
 @app.get("/models", response_model=dict)
-def get_model_info():
-    """Retourne les informations du modèle chargé"""
-    if model is None:
-        raise HTTPException(status_code=503, detail="Model not loaded")
-    
+def get_models_info():
+    if best_model is None or second_best_model is None:
+        raise HTTPException(status_code=503, detail="Models not loaded")
     return {
-        "model_name": MODEL_NAME,
-        "run_id": run_id,
-        "encoder": encoder_name,
-        "input_size": script_img_size,
+        "best_model": best_model_info,
+        "second_best_model": second_best_model_info,
         "num_classes": mapping_config['num_classes'],
         "class_names": mapping_config['group_names'],
         "class_colors": mapping_config['group_colors'].tolist()
@@ -190,238 +200,199 @@ def get_model_info():
 
 @app.get("/sample-images", response_model=AvailableImagesResponse)
 def get_sample_images():
-    """Retourne la liste des images d'exemple disponibles"""
     try:
         original_dir = Path(SAMPLE_IMAGES_DIR) / "original"
         mask_dir = Path(SAMPLE_IMAGES_DIR) / "mask"
-        
         if not original_dir.exists():
             raise HTTPException(status_code=404, detail="Sample images directory not found")
-        
         images = []
         for img_path in original_dir.glob("*leftImg8bit.png"):
-            # Construire le nom du masque correspondant
             base_name = img_path.stem.replace("_leftImg8bit", "")
             mask_name = f"{base_name}_gtFine_labelIds.png"
             mask_path = mask_dir / mask_name
-            
-            # Créer un nom d'affichage plus convivial
             display_name = base_name.replace("_", " ").title()
-            
             images.append(ImageInfo(
                 filename=img_path.name,
                 display_name=display_name,
                 has_ground_truth=mask_path.exists()
             ))
-        
-        # Trier par nom pour un affichage cohérent
         images.sort(key=lambda x: x.filename)
-        
-        return AvailableImagesResponse(
-            images=images,
-            total_count=len(images)
-        )
-        
+        return AvailableImagesResponse(images=images, total_count=len(images))
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error listing sample images: {str(e)}")
 
 @app.get("/sample-image/{filename}")
 def get_sample_image(filename: str):
-    """Sert une image d'exemple pour prévisualisation"""
     try:
-        print(f"🔍 DEBUG: Demande d'image: {filename}")
-        
         original_dir = Path(SAMPLE_IMAGES_DIR) / "original"
         image_path = original_dir / filename
-        
-        print(f"🔍 DEBUG: Chemin construit: {image_path}")
-        print(f"🔍 DEBUG: Fichier existe: {image_path.exists()}")
-        
         if not image_path.exists():
-            print(f"❌ DEBUG: Image non trouvée: {image_path}")
             raise HTTPException(status_code=404, detail=f"Sample image {filename} not found")
-        
-        # Vérifier que c'est bien un fichier image attendu (sécurité)
         if not filename.endswith('_leftImg8bit.png'):
-            print(f"❌ DEBUG: Format invalide: {filename}")
             raise HTTPException(status_code=400, detail="Invalid image filename format")
-        
-        print(f"✅ DEBUG: Envoi de l'image: {filename}")
         return FileResponse(
             image_path,
             media_type="image/png",
             headers={"Content-Disposition": f"inline; filename={filename}"}
         )
-        
     except HTTPException:
-        # Re-raise les HTTPException pour qu'elles soient gérées correctement
         raise
     except Exception as e:
-        print(f"❌ Erreur lors du service de l'image {filename}: {e}")
-        import traceback
-        traceback.print_exc()
         raise HTTPException(status_code=500, detail=f"Error serving image: {str(e)}")
 
 class SampleImageRequest(BaseModel):
     filename: str
 
-@app.post("/predict-sample", response_model=PredictionResult)
-def predict_sample_image(request: SampleImageRequest):
-    """Lance une prédiction sur une image d'exemple"""
-    if model is None or mapping_config is None:
-        raise HTTPException(status_code=503, detail="Model not loaded")
-    
+@app.post("/compare-sample", response_model=ComparisonResult)
+def compare_sample_image(request: SampleImageRequest):
+    if best_model is None or second_best_model is None or mapping_config is None:
+        raise HTTPException(status_code=503, detail="Models not loaded")
     try:
-        # Vérifier que l'image existe
         original_dir = Path(SAMPLE_IMAGES_DIR) / "original"
         image_path = original_dir / request.filename
-        
         if not image_path.exists():
             raise HTTPException(status_code=404, detail=f"Sample image {request.filename} not found")
-        
-        # Construire le chemin du masque de vérité
         base_name = image_path.stem.replace("_leftImg8bit", "")
         mask_name = f"{base_name}_gtFine_labelIds.png"
         mask_path = Path(SAMPLE_IMAGES_DIR) / "mask" / mask_name
-        
         mask_path_str = str(mask_path) if mask_path.exists() else None
         
-        print(f"🔮 Lancement de l'inférence sur {request.filename}")
-        
-        # Lancer l'inférence (sans affichage pour l'API)
-        predicted_mask, fig = run_inference_and_visualize(
+        # Inférence modèle 1
+        pred1, t1, _ = run_inference_and_visualize(
             image_path=str(image_path),
             mask_path=mask_path_str,
-            model=model,
-            encoder_name=encoder_name,
-            img_size=script_img_size,
+            model=best_model,
+            encoder_name=None,
+            img_size=best_model_info["input_size"],
             mapping_config=mapping_config,
-            save_dir=None,
-            show_stats=False,
-            show_plot=False  # ← AJOUTÉ pour récupérer la figure
+            save_dir=None, show_stats=False, show_plot=False
         )
-        
-        # Sérialiser la figure
-        fig_bytes = pickle.dumps(fig)
-        fig_b64 = base64.b64encode(fig_bytes).decode()
-        
-        # Calculer les statistiques
-        unique_classes, counts = np.unique(predicted_mask, return_counts=True)
-        total_pixels = predicted_mask.size
-        
-        stats = {}
-        for class_id, count in zip(unique_classes, counts):
-            if class_id < len(mapping_config['group_names']):
-                percentage = (count / total_pixels) * 100
-                stats[mapping_config['group_names'][class_id]] = {
-                    "pixels": int(count),
-                    "percentage": round(percentage, 1)
-                }
-        
-        return PredictionResult(
-            success=True,
-            message="Prediction completed successfully",
+        # Inférence modèle 2
+        pred2, t2, _ = run_inference_and_visualize(
             image_path=str(image_path),
             mask_path=mask_path_str,
-            prediction_stats=stats,
-            ground_truth_available=mask_path.exists(),
+            model=second_best_model,
+            encoder_name=None,
+            img_size=second_best_model_info["input_size"],
+            mapping_config=mapping_config,
+            save_dir=None, show_stats=False, show_plot=False
+        )
+        # Figure combinée
+        fig, axes = plt.subplots(2, 3, figsize=(15, 8))
+        img_arr = mpimg.imread(str(image_path))
+        axes[0, 0].imshow(img_arr); axes[0, 0].set_title("Image originale"); axes[0, 0].axis('off')
+        axes[1, 0].imshow(img_arr); axes[1, 0].set_title("Image originale"); axes[1, 0].axis('off')
+        axes[0, 1].imshow(colorize_mask(pred1, mapping_config['group_colors']))
+        axes[0, 1].set_title(f"Prédiction {best_model_info['encoder_name']}"); axes[0, 1].axis('off')
+        axes[1, 1].imshow(colorize_mask(pred2, mapping_config['group_colors']))
+        axes[1, 1].set_title(f"Prédiction {second_best_model_info['encoder_name']}"); axes[1, 1].axis('off')
+        if mask_path_str and os.path.exists(mask_path_str):
+            gt = mpimg.imread(mask_path_str)
+            axes[0, 2].imshow(gt); axes[0, 2].set_title("Ground Truth"); axes[0, 2].axis('off')
+            axes[1, 2].imshow(gt); axes[1, 2].set_title("Ground Truth"); axes[1, 2].axis('off')
+        plt.tight_layout()
+        fig_b64 = base64.b64encode(pickle.dumps(fig)).decode()
+        # Stats
+        def calc(mask):
+            uc, cnts = np.unique(mask, return_counts=True)
+            stats = {}
+            for cid, c in zip(uc, cnts):
+                if cid < len(mapping_config['group_names']):
+                    stats[mapping_config['group_names'][cid]] = {
+                        "pixels": int(c),
+                        "percentage": round((c/mask.size)*100, 1)
+                    }
+            return stats
+        stats1, stats2 = calc(pred1), calc(pred2)
+        faster = best_model_info['encoder_name'] if t1 < t2 else second_best_model_info['encoder_name']
+        speedup = round(max(t1, t2) / min(t1, t2), 2)
+        speed_comp = {"faster_model": faster, "speedup_ratio": speedup, "time_difference_ms": abs(t1-t2)*1000}
+        return ComparisonResult(
+            success=True, message="Comparison completed successfully",
+            model1_info=ModelInfo(**best_model_info), model2_info=ModelInfo(**second_best_model_info),
+            image_path=str(image_path), mask_path=mask_path_str,
+            model1_stats=stats1, model2_stats=stats2,
+            model1_inference_time=t1, model2_inference_time=t2,
+            speed_comparison=speed_comp, ground_truth_available=mask_path.exists(),
             figure_data=fig_b64
         )
-        
     except Exception as e:
-        print(f"❌ Erreur lors de la prédiction: {e}")
-        import traceback
-        traceback.print_exc()
-        raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Comparison error: {str(e)}")
 
-@app.post("/predict-upload", response_model=PredictionResult)
-def predict_uploaded_image(file: UploadFile = File(...)):
-    """Lance une prédiction sur une image uploadée"""
-    if model is None or mapping_config is None:
-        raise HTTPException(status_code=503, detail="Model not loaded")
-    
-    # Vérifier le type de fichier
+@app.post("/compare-upload", response_model=ComparisonResult)
+def compare_uploaded_image(file: UploadFile = File(...)):
+    if best_model is None or second_best_model is None or mapping_config is None:
+        raise HTTPException(status_code=503, detail="Models not loaded")
     allowed_types = ["image/jpeg", "image/jpg", "image/png"]
     if file.content_type not in allowed_types:
-        raise HTTPException(
-            status_code=400, 
-            detail=f"Invalid file type. Allowed: {', '.join(allowed_types)}"
-        )
-    
+        raise HTTPException(status_code=400, detail=f"Invalid file type. Allowed: {', '.join(allowed_types)}")
     temp_file_path = None
     try:
-        # Créer un fichier temporaire
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
         temp_filename = f"upload_{timestamp}_{file.filename}"
         temp_file_path = os.path.join(TEMP_UPLOAD_DIR, temp_filename)
-        
-        # Sauvegarder le fichier uploadé
         with open(temp_file_path, "wb") as buffer:
             shutil.copyfileobj(file.file, buffer)
-        
-        print(f"🔮 Lancement de l'inférence sur {file.filename}")
-        
-        # Lancer l'inférence (sans masque de vérité)
-        predicted_mask, fig = run_inference_and_visualize(
-            image_path=temp_file_path,
-            mask_path=None,  # Pas de masque de vérité pour les uploads
-            model=model,
-            encoder_name=encoder_name,
-            img_size=script_img_size,
-            mapping_config=mapping_config,
-            save_dir=None,
-            show_stats=False,
-            show_plot=False  # ← AJOUTÉ pour récupérer la figure
+        # Inférences
+        pred1, t1, _ = run_inference_and_visualize(
+            image_path=temp_file_path, mask_path=None,
+            model=best_model, encoder_name=None,
+            img_size=best_model_info["input_size"], mapping_config=mapping_config,
+            save_dir=None, show_stats=False, show_plot=False,
         )
-        
-        # Sérialiser la figure
-        fig_bytes = pickle.dumps(fig)
-        fig_b64 = base64.b64encode(fig_bytes).decode()
-        
-        # Calculer les statistiques
-        unique_classes, counts = np.unique(predicted_mask, return_counts=True)
-        total_pixels = predicted_mask.size
-        
-        stats = {}
-        for class_id, count in zip(unique_classes, counts):
-            if class_id < len(mapping_config['group_names']):
-                percentage = (count / total_pixels) * 100
-                stats[mapping_config['group_names'][class_id]] = {
-                    "pixels": int(count),
-                    "percentage": round(percentage, 1)
-                }
-        
-        return PredictionResult(
-            success=True,
-            message="Prediction completed successfully",
+        pred2, t2, _ = run_inference_and_visualize(
+            image_path=temp_file_path, mask_path=None,
+            model=second_best_model, encoder_name=None,
+            img_size=second_best_model_info["input_size"], mapping_config=mapping_config,
+            save_dir=None, show_stats=False, show_plot=False,
+        )
+        # Figure
+        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
+        img_arr = mpimg.imread(temp_file_path)
+        axes[0].imshow(img_arr); axes[0].set_title("Image originale"); axes[0].axis('off')
+        axes[1].imshow(colorize_mask(pred1, mapping_config['group_colors']))
+        axes[1].set_title(f"Prédiction {best_model_info['encoder_name']}"); axes[1].axis('off')
+        axes[2].imshow(colorize_mask(pred2, mapping_config['group_colors']))
+        axes[2].set_title(f"Prédiction {second_best_model_info['encoder_name']}"); axes[2].axis('off')
+        plt.tight_layout()
+        fig_b64 = base64.b64encode(pickle.dumps(fig)).decode()
+        # Stats
+        def calc(mask):
+            uc, cnts = np.unique(mask, return_counts=True)
+            stats = {}
+            for cid, c in zip(uc, cnts):
+                if cid < len(mapping_config['group_names']):
+                    stats[mapping_config['group_names'][cid]] = {
+                        "pixels": int(c),
+                        "percentage": round((c/mask.size)*100, 1)
+                    }
+            return stats
+        stats1, stats2 = calc(pred1), calc(pred2)
+        faster = best_model_info['encoder_name'] if t1 < t2 else second_best_model_info['encoder_name']
+        speedup = round(max(t1, t2) / min(t1, t2), 2)
+        speed_comp = {"faster_model": faster, "speedup_ratio": speedup, "time_difference_ms": abs(t1-t2)*1000}
+        return ComparisonResult(
+            success=True, message="Comparison completed successfully",
+            model1_info=ModelInfo(**best_model_info), model2_info=ModelInfo(**second_best_model_info),
             image_path=temp_file_path,
-            prediction_stats=stats,
-            ground_truth_available=False,
+            model1_stats=stats1, model2_stats=stats2,
+            model1_inference_time=t1, model2_inference_time=t2,
+            speed_comparison=speed_comp, ground_truth_available=False,
             figure_data=fig_b64
         )
-        
     except Exception as e:
-        print(f"❌ Erreur lors de la prédiction: {e}")
-        import traceback
-        traceback.print_exc()
-        raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")
-    
+        raise HTTPException(status_code=500, detail=f"Comparison error: {str(e)}")
     finally:
-        # Nettoyer le fichier temporaire
         if temp_file_path and os.path.exists(temp_file_path):
-            try:
-                os.remove(temp_file_path)
-                print(f"🗑️ Fichier temporaire supprimé: {temp_filename}")
-            except Exception as e:
-                print(f"⚠️ Impossible de supprimer le fichier temporaire: {e}")
+            try: os.remove(temp_file_path)
+            except Exception: pass
 
 @app.get("/metrics")
 def get_metrics():
-    """Endpoint de métriques pour monitoring"""
     return {
-        "model_loaded": model is not None,
-        "model_name": MODEL_NAME,
-        "run_id": run_id,
+        "models_loaded": best_model is not None and second_best_model is not None,
+        "best_model_info": best_model_info,
+        "second_best_model_info": second_best_model_info,
         "num_classes": mapping_config['num_classes'] if mapping_config else None,
         "timestamp": datetime.utcnow().isoformat()
-    }
\ No newline at end of file
+    }
diff --git a/app/requirements.txt b/app/requirements.txt
index c78ee04..ef6d4de 100644
--- a/app/requirements.txt
+++ b/app/requirements.txt
@@ -8,37 +8,44 @@ uvicorn==0.35.0
 starlette==0.46.2
 pydantic==2.11.7
 pydantic_core==2.33.2
+python-multipart==0.0.20  # uploads FastAPI
 
 # ===== STREAMLIT =====
 streamlit==1.46.1
-altair==5.5.0  # Utilisé par Streamlit pour les graphiques
-blinker==1.9.0  # Streamlit dependency
+altair==5.5.0  # utilisé par Streamlit
+blinker==1.9.0 # dépendance Streamlit
 
-# ===== HTTP REQUESTS =====
+# ===== HTTP CLIENTS =====
 requests==2.32.4
-httpx==0.28.1  # Modern HTTP client (utilisé par FastAPI)
+httpx==0.28.1   # utilisé en interne par FastAPI
 httpcore==1.0.9
 h11==0.16.0
 
-# ===== MACHINE LEARNING CORE =====
-tensorflow==2.19.0
-keras==3.10.0
-numpy>=1.24.0,<2.0.0
+# ===== MACHINE LEARNING CORE (PyTorch) =====
+torch==2.8.0        # version stable PyTorch
+torchvision==0.23.0 # utile pour les modèles pré-entraînés
+torchaudio==2.8.0   # optionnel (mais recommandé pour compatibilité PyTorch)
+segmentation-models-pytorch==0.5.0
+albumentations==2.0.8
+# version PyTorch de ToTensor
+albumentations[imgaug,opencv]
 
-# Note: tensorflow-addons n'est plus nécessaire avec TensorFlow 2.19+
-# Les fonctionnalités sont maintenant intégrées dans TensorFlow core
+numpy>=1.24.0,<2.0.0
+scipy==1.16.0
+scikit-learn==1.7.0
 
 # ===== COMPUTER VISION =====
 pillow==11.3.0
-opencv-python-headless==4.9.0.80  # Pas dans votre freeze, mais nécessaire pour Docker
+opencv-python-headless==4.9.0.80
 scikit-image==0.25.2
 imageio==2.37.0
 
 # ===== DATA SCIENCE =====
 pandas==2.3.0
 matplotlib==3.10.3
-scipy==1.16.0
-scikit-learn==1.7.0
+networkx==3.5
+tifffile==2025.6.11
+lazy_loader==0.4
 
 # ===== MLFLOW & MODEL REGISTRY =====
 mlflow==3.2.0
@@ -50,34 +57,7 @@ boto3==1.39.2
 botocore==1.39.2
 s3transfer==0.13.0
 
-# ===== SEGMENTATION MODELS =====
-# Option 1: Version figée par commit (RECOMMANDÉE pour production)
-segmentation_models @ git+https://github.com/qubvel/segmentation_models.git@5d24bbfb28af6134e25e2c0b79e7727f6c0491d0
-
-# Option 2: Si vous préférez une version PyPI (décommentez et commentez la ligne du dessus)
-# segmentation-models==1.0.1
-
-# ===== SUPPORTING LIBRARIES =====
-# TensorFlow ecosystem
-tensorboard==2.19.0
-tensorboard-data-server==0.7.2
-protobuf==3.20.3
-flatbuffers==25.2.10
-h5py==3.14.0
-astunparse==1.6.3
-gast==0.6.0
-google-pasta==0.2.0
-keras-applications==1.0.8
-opt_einsum==3.4.0
-termcolor==2.5.0
-wrapt==1.17.2
-
-# Data processing
-networkx==3.5
-tifffile==2025.6.11
-lazy_loader==0.4
-
-# Utilities
+# ===== UTILITIES =====
 click==8.2.1
 tqdm==4.67.1
 joblib==1.5.1
@@ -85,57 +65,6 @@ python-dateutil
 packaging
 typing_extensions==4.14.0
 
-# File handling
-python-multipart==0.0.20  # Pour FastAPI file uploads
-
-# ===== AJOUTS SPÉCIFIQUES POUR UTILS/UTILS.PY =====
-# Pour la gestion des couleurs et légendes matplotlib
-cycler==0.12.1  # Pour matplotlib customization
-
-# Font handling pour matplotlib (nécessaire dans Docker)
+# ===== MATPLOTLIB ADDONS =====
+cycler==0.12.1
 fonttools==4.58.5
-
-# ===== REMOVED FROM ORIGINAL =====
-# Les packages suivants ont été supprimés car non utilisés dans votre app:
-# - awsebcli (AWS CLI - pas nécessaire dans le container)
-# - blessed, fabric, invoke (outils dev)
-# - jupyter, ipython, spyder (dev environment)
-# - nltk (pas d'NLP dans votre app)
-# - pytest (tests - à garder seulement pour dev)
-# - flask, gunicorn (vous utilisez FastAPI+uvicorn)
-# - docker (pas besoin dans le container)
-# - git* packages (dev tools)
-
-# =============================================================================
-# VERSION ALTERNATIVE - PLUS MINIMALISTE
-# =============================================================================
-
-# Si vous voulez une version encore plus légère, voici les STRICTEMENT nécessaires:
-# fastapi==0.115.14
-# uvicorn==0.35.0
-# streamlit==1.46.1
-# tensorflow==2.19.0
-# numpy==2.1.3
-# pillow==11.3.0
-# matplotlib==3.10.3
-# pandas==2.3.0
-# requests==2.32.4
-# mlflow==3.2.0
-# boto3==1.39.2
-# segmentation_models @ git+https://github.com/qubvel/segmentation_models.git@5d24bbfb28af6134e25e2c0b79e7727f6c0491d0
-
-# =============================================================================
-# NOTES SUR SEGMENTATION_MODELS
-# =============================================================================
-
-# OPTION A: Git avec commit hash (ACTUEL - RECOMMANDÉ)
-# Avantages: Version exacte, pas de surprises
-# Inconvénients: Nécessite git dans le container de build
-
-# OPTION B: Version PyPI officielle
-# segmentation-models==1.0.1
-# Avantages: Plus simple, pas besoin de git
-# Inconvénients: Peut être différent de votre version dev
-
-# OPTION C: Version locale (si vous avez des modifications)
-# Copier le package dans le container et faire pip install -e .
diff --git a/app/streamlit_app.py b/app/streamlit_app.py
index 6f9a090..3351498 100644
--- a/app/streamlit_app.py
+++ b/app/streamlit_app.py
@@ -10,25 +10,21 @@ import base64
 import pickle
 
 # Configuration de l'API
-# prod docker
 API_BASE_URL = os.getenv("MULTISEG_API_BASE_URL", "http://127.0.0.1:8080/api")
-# test local
-# API_BASE_URL = os.getenv("MULTISEG_API_BASE_URL", "http://127.0.0.1:8000")  # ← Port 8000 au lieu de 8080
 
 # Import des fonctions utils pour l'affichage local
 sys.path.append('/app')
-# APRÈS (pour test local)
 sys.path.append('.')
 from utils.utils import colorize_mask
 
 # Configuration de la page
 st.set_page_config(
-    page_title="Multi-Class Segmentation",
-    page_icon="🚗",
+    page_title="Multi-Class Segmentation Comparison",
+    page_icon="🔄",
     layout="wide"
 )
 
-# CSS personnalisé
+# CSS personnalisé mis à jour
 st.markdown("""
 <style>
 .main-header {
@@ -39,6 +35,23 @@ st.markdown("""
     color: white;
     text-align: center;
 }
+.comparison-card {
+    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
+    padding: 1.5rem;
+    border-radius: 10px;
+    border-left: 4px solid #667eea;
+    margin: 1rem 0;
+    color: #2c3e50;
+    font-weight: 500;
+}
+.speed-winner {
+    background: linear-gradient(135deg, #d4edda 0%, #c3e6cb 100%);
+    border-left-color: #28a745;
+}
+.speed-slower {
+    background: linear-gradient(135deg, #f8d7da 0%, #f1b0b7 100%);
+    border-left-color: #dc3545;
+}
 .metric-card {
     background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
     padding: 1rem;
@@ -60,33 +73,19 @@ st.markdown("""
 # En-tête principal
 st.markdown("""
 <div class="main-header">
-    <h1>🚗 Segmentation Multi-Classes</h1>
-    <p>Système de segmentation sémantique pour véhicules autonomes</p>
+    <h1>🔄 Comparaison de Segmentation Multi-Classes</h1>
+    <p>Système de comparaison entre deux modèles de segmentation sémantique</p>
 </div>
 """, unsafe_allow_html=True)
 
-# Initialisation des états
-if "prediction_result" not in st.session_state:
-    st.session_state.prediction_result = None
-if "selected_image" not in st.session_state:
-    st.session_state.selected_image = None
-if "uploaded_file" not in st.session_state:
-    st.session_state.uploaded_file = None
-
-def reset_state():
-    """Reset des états de session"""
-    st.session_state.prediction_result = None
-    st.session_state.selected_image = None
-    st.session_state.uploaded_file = None
-
-def get_model_info():
-    """Récupère les informations du modèle"""
+def get_models_info():
+    """Récupère les informations des deux modèles"""
     try:
         response = requests.get(f"{API_BASE_URL}/models")
         if response.status_code == 200:
             return response.json()
         else:
-            st.error(f"Erreur lors de la récupération des infos modèle: {response.status_code}")
+            st.error(f"Erreur lors de la récupération des infos modèles: {response.status_code}")
             return None
     except Exception as e:
         st.error(f"Erreur de connexion à l'API: {e}")
@@ -105,15 +104,14 @@ def get_sample_images():
         st.error(f"Erreur de connexion à l'API: {e}")
         return None
 
-def predict_sample_image(filename):
-    """Lance une prédiction sur une image d'exemple"""
+def compare_sample_image(filename):
+    """Lance une comparaison sur une image d'exemple"""
     try:
-        # CORRECTION: Envoyer en JSON au lieu de form data
         data = {"filename": filename}
         response = requests.post(
-            f"{API_BASE_URL}/predict-sample", 
-            json=data,  # ← Changé de 'data=' à 'json='
-            headers={"Content-Type": "application/json"}  # ← Ajout du header explicite
+            f"{API_BASE_URL}/compare-sample", 
+            json=data,
+            headers={"Content-Type": "application/json"}
         )
         
         if response.status_code == 200:
@@ -122,14 +120,14 @@ def predict_sample_image(filename):
             st.error(f"Erreur API: {response.status_code} - {response.text}")
             return None
     except Exception as e:
-        st.error(f"Erreur lors de la prédiction: {e}")
+        st.error(f"Erreur lors de la comparaison: {e}")
         return None
     
-def predict_uploaded_image(uploaded_file):
-    """Lance une prédiction sur une image uploadée"""
+def compare_uploaded_image(uploaded_file):
+    """Lance une comparaison sur une image uploadée"""
     try:
         files = {"file": (uploaded_file.name, uploaded_file.getvalue(), uploaded_file.type)}
-        response = requests.post(f"{API_BASE_URL}/predict-upload", files=files)
+        response = requests.post(f"{API_BASE_URL}/compare-upload", files=files)
         
         if response.status_code == 200:
             return response.json()
@@ -137,303 +135,253 @@ def predict_uploaded_image(uploaded_file):
             st.error(f"Erreur API: {response.status_code} - {response.text}")
             return None
     except Exception as e:
-        st.error(f"Erreur lors de la prédiction: {e}")
+        st.error(f"Erreur lors de la comparaison: {e}")
         return None
 
-def display_segmentation_plots(fig):
-    """Affiche chaque subplot séparément dans des colonnes"""
-    if fig is None:
-        st.warning("Aucune figure de segmentation disponible")
+def display_comparison_results(result, models_info):
+    """Affiche les résultats de comparaison entre les deux modèles"""
+    if not result or not result['success']:
+        st.error("Aucun résultat de comparaison disponible")
         return
     
-    axes = fig.get_axes()
-    num_plots = len(axes)
+    # Informations des modèles
+    st.markdown("### 🏆 Modèles Comparés")
     
-    st.markdown("### 🎨 Résultat de la Segmentation")
+    col1, col2 = st.columns(2)
     
-    # Créer les colonnes selon le nombre de subplots
-    if num_plots <= 4:
-        cols = st.columns(num_plots)
-    else:
-        cols = st.columns(4)  # Maximum 4 colonnes
+    with col1:
+        model1 = result['model1_info']
+        st.markdown(f"""
+        <div class="comparison-card">
+            <h4>🥇 Modèle #1</h4>
+            <strong>{model1['encoder_name']}</strong><br>
+            📐 Taille: {model1['input_size']}<br>
+            🆔 Run: {model1['run_id'][:8]}...
+        </div>
+        """, unsafe_allow_html=True)
     
-    for i, ax in enumerate(axes):
-        col_idx = i % len(cols)
-        
-        with cols[col_idx]:
-            # Récupérer le titre existant
-            title = ax.get_title()
-            # st.markdown(f"**{title}:**")
-            
-            # Créer une figure individuelle avec ce subplot
-            individual_fig, individual_ax = plt.subplots(figsize=(6, 6))
-            
-            # Copier toutes les images de l'axe original
-            for image in ax.get_images():
-                individual_ax.imshow(
-                    image.get_array(),
-                    cmap=image.get_cmap() if image.get_cmap() else None,
-                    alpha=image.get_alpha() if image.get_alpha() else None,
-                    vmin=image.get_clim()[0] if image.get_clim() else None,
-                    vmax=image.get_clim()[1] if image.get_clim() else None
-                )
-            
-            # Copier les propriétés de l'axe
-            individual_ax.set_title(title)
-            individual_ax.axis('off')
-            
-            # Ajuster la mise en page
-            individual_fig.tight_layout()
-            
-            # Afficher dans Streamlit
-            st.pyplot(individual_fig)
-            
-            # Libérer la mémoire
-            plt.close(individual_fig)
+    with col2:
+        model2 = result['model2_info']
+        st.markdown(f"""
+        <div class="comparison-card">
+            <h4>🥈 Modèle #2</h4>
+            <strong>{model2['encoder_name']}</strong><br>
+            📐 Taille: {model2['input_size']}<br>
+            🆔 Run: {model2['run_id'][:8]}...
+        </div>
+        """, unsafe_allow_html=True)
     
-    # Fermer la figure originale pour libérer la mémoire
-    plt.close(fig)
+    # Comparaison des temps d'inférence
+    st.markdown("### ⏱️ Comparaison des Performances")
+    
+    speed_comp = result['speed_comparison']
+    model1_time = result['model1_inference_time'] * 1000
+    model2_time = result['model2_inference_time'] * 1000
+    
+    col1, col2, col3 = st.columns(3)
+    
+    with col1:
+        winner_class = "speed-winner" if speed_comp['faster_model'] == model1['encoder_name'] else "speed-slower"
+        st.markdown(f"""
+        <div class="comparison-card {winner_class}">
+            <h5>🥇 {model1['encoder_name']}</h5>
+            <strong>{model1_time:.1f} ms</strong>
+        </div>
+        """, unsafe_allow_html=True)
+    
+    with col2:
+        st.markdown(f"""
+        <div class="comparison-card">
+            <h5>⚡ Comparaison</h5>
+            <strong>{speed_comp['faster_model']}</strong><br>
+            est {speed_comp['speedup_ratio']}x plus rapide<br>
+            <small>Différence: {speed_comp['time_difference_ms']:.1f} ms</small>
+        </div>
+        """, unsafe_allow_html=True)
+    
+    with col3:
+        winner_class = "speed-winner" if speed_comp['faster_model'] == model2['encoder_name'] else "speed-slower"
+        st.markdown(f"""
+        <div class="comparison-card {winner_class}">
+            <h5>🥈 {model2['encoder_name']}</h5>
+            <strong>{model2_time:.1f} ms</strong>
+        </div>
+        """, unsafe_allow_html=True)
+    
+    # Affichage des segmentations
+    if result.get('figure_data'):
+        st.markdown("### 🎨 Résultats de Segmentation")
+        try:
+            fig_bytes = base64.b64decode(result['figure_data'])
+            fig = pickle.loads(fig_bytes)
+            st.pyplot(fig)
+            plt.close(fig)
+        except Exception as e:
+            st.error(f"Erreur lors de l'affichage: {e}")
+    
+    # Comparaison des statistiques
+    st.markdown("### 📊 Comparaison des Statistiques")
+    
+    col1, col2 = st.columns(2)
+    
+    with col1:
+        st.markdown(f"**📈 {model1['encoder_name']} - Répartition des classes:**")
+        if result.get('model1_stats'):
+            display_model_stats_chart(result['model1_stats'], models_info, f"Modèle 1 - {model1['encoder_name']}")
+    
+    with col2:
+        st.markdown(f"**📈 {model2['encoder_name']} - Répartition des classes:**")
+        if result.get('model2_stats'):
+            display_model_stats_chart(result['model2_stats'], models_info, f"Modèle 2 - {model2['encoder_name']}")
+    
+    # Analyse comparative des classes
+    if result.get('model1_stats') and result.get('model2_stats'):
+        display_class_comparison(result['model1_stats'], result['model2_stats'], model1, model2)
 
-        
-def display_prediction_stats_with_chart(stats, model_info):
-    """Affiche les statistiques avec un graphique en barres"""
-    st.markdown("### 📊 Statistiques de Segmentation")
-    
+def display_model_stats_chart(stats, models_info, title):
+    """Affiche les statistiques d'un modèle sous forme de graphique"""
     if not stats:
         st.warning("Aucune statistique disponible")
         return
     
-    # Ratio 1:2 pour donner plus de place au graphique
-    col1, col2 = st.columns([1, 2])
+    # Préparer les données pour le graphique
+    class_names = list(stats.keys())
+    percentages = [data['percentage'] for data in stats.values()]
     
-    # Colonne 1: Cartes de statistiques compactes avec carrés colorés
-    with col1:
-        st.markdown("**Détails par classe:**")
+    # Récupérer les couleurs des classes
+    colors = []
+    if models_info and 'class_colors' in models_info and 'class_names' in models_info:
+        model_class_names = models_info['class_names']
+        model_class_colors = models_info['class_colors']
         
-        # Récupérer les couleurs des classes pour les carrés
-        class_colors_dict = {}
-        if model_info and 'class_colors' in model_info and 'class_names' in model_info:
-            model_class_names = model_info['class_names']
-            model_class_colors = model_info['class_colors']
-            
-            for class_name in stats.keys():
-                try:
-                    idx = model_class_names.index(class_name)
-                    color = model_class_colors[idx]
-                    # Convertir en hex pour CSS
-                    color_hex = f"#{color[0]:02x}{color[1]:02x}{color[2]:02x}"
-                    class_colors_dict[class_name] = color_hex
-                except (ValueError, IndexError):
-                    # Couleur par défaut si pas trouvée
-                    class_colors_dict[class_name] = "#808080"
-        
-        for class_name, data in stats.items():
-            # Récupérer la couleur pour cette classe
-            color_hex = class_colors_dict.get(class_name, "#808080")
-            
-            # ✨ NOUVEAU: Carré coloré + nom de classe
-            st.markdown(f"""
-            <div style="
-                background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
-                padding: 0.5rem 0.8rem;
-                border-radius: 6px;
-                border-left: 3px solid #667eea;
-                margin: 0.3rem 0;
-                color: #2c3e50;
-                font-weight: 500;
-                font-size: 0.85rem;
-                line-height: 1.3;
-            ">
-                <div style="
-                    font-weight: 600; 
-                    margin-bottom: 0.2rem;
-                    display: flex;
-                    align-items: center;
-                    gap: 0.5rem;
-                ">
-                    <div style="
-                        width: 12px;
-                        height: 12px;
-                        background-color: {color_hex};
-                        border-radius: 2px;
-                        border: 1px solid rgba(0,0,0,0.2);
-                        flex-shrink: 0;
-                    "></div>
-                    <span>{class_name}</span>
-                </div>
-                <div style="font-size: 0.8rem; color: #555;">
-                    🔢 {data['pixels']:,} pixels | 📊 {data['percentage']}%
-                </div>
-            </div>
-            """, unsafe_allow_html=True)
-    
-    # Colonne 2: Graphique en barres (reste identique)
-    with col2:
-        st.markdown("**Répartition visuelle:**")
-        
-        # Préparer les données pour le graphique
-        class_names = list(stats.keys())
-        percentages = [data['percentage'] for data in stats.values()]
-        
-        # Récupérer les couleurs des classes depuis model_info
-        colors = []
-        if model_info and 'class_colors' in model_info and 'class_names' in model_info:
-            model_class_names = model_info['class_names']
-            model_class_colors = model_info['class_colors']
-            
-            for class_name in class_names:
-                # Trouver l'index de cette classe dans le modèle
-                try:
-                    idx = model_class_names.index(class_name)
-                    color = model_class_colors[idx]
-                    # Convertir RGB en format matplotlib
-                    colors.append([c/255.0 for c in color])
-                except (ValueError, IndexError):
-                    # Couleur par défaut si pas trouvée
-                    colors.append([0.5, 0.5, 0.5])
-        else:
-            # Couleurs par défaut si pas d'info du modèle
-            colors = plt.cm.Set3(np.linspace(0, 1, len(class_names)))
-        
-        # Créer un graphique plus grand et mieux adapté
-        fig, ax = plt.subplots(figsize=(10, 6))  # Plus large
-        
-        bars = ax.bar(class_names, percentages, color=colors, alpha=0.8, edgecolor='white', linewidth=1)
-        
-        # Personnalisation du graphique améliorée
-        ax.set_xlabel('Classes de Segmentation', fontsize=11, fontweight='bold')
-        ax.set_ylabel('Pourcentage (%)', fontsize=11, fontweight='bold')
-        ax.set_title('Répartition des Classes Détectées', fontsize=13, fontweight='bold', pad=20)
-        
-        # Meilleure gestion des labels selon le nombre de classes
-        if len(class_names) <= 4:
-            plt.xticks(rotation=0, ha='center', fontsize=10)
-        elif len(class_names) <= 6:
-            plt.xticks(rotation=30, ha='right', fontsize=9)
-        else:
-            plt.xticks(rotation=45, ha='right', fontsize=8)
-        
-        plt.yticks(fontsize=10)
-        
-        # Valeurs sur les barres avec meilleur positionnement
-        for bar, percentage in zip(bars, percentages):
-            height = bar.get_height()
-            # Adapter la position du texte selon la hauteur de la barre
-            if height < max(percentages) * 0.1:
-                # Si la barre est très petite, mettre le texte au-dessus
-                va = 'bottom'
-                y_pos = height + max(percentages) * 0.02
-            else:
-                # Sinon, centrer dans la barre
-                va = 'center'
-                y_pos = height / 2
-            
-            ax.text(bar.get_x() + bar.get_width()/2., y_pos,
-                   f'{percentage:.1f}%',
-                   ha='center', va=va, fontsize=9, fontweight='bold',
-                   color='white' if va == 'center' else 'black',
-                   bbox=dict(boxstyle='round,pad=0.2', facecolor='black', alpha=0.7) if va == 'center' else None)
-        
-        # Améliorer la mise en page et le style
-        plt.tight_layout()
-        
-        # Grille plus subtile
-        ax.grid(axis='y', alpha=0.2, linestyle='--')
-        ax.set_axisbelow(True)
-        
-        # Limites et style des axes
-        ax.set_ylim(0, max(percentages) * 1.15)  # Plus d'espace au-dessus
-        ax.spines['top'].set_visible(False)
-        ax.spines['right'].set_visible(False)
-        ax.spines['left'].set_color('#666')
-        ax.spines['bottom'].set_color('#666')
-        
-        # Couleur de fond
-        fig.patch.set_facecolor('white')
-        ax.set_facecolor('#fafafa')
-        
-        # Afficher dans Streamlit
-        st.pyplot(fig)
-        plt.close(fig)       
-
-
-def display_class_legend(model_info):
-    """Affiche la légende des classes de segmentation"""
-    if model_info and 'class_colors' in model_info:
-        st.markdown("**Légende des classes:**")
-        
-        class_names = model_info.get('class_names', [])
-        class_colors = model_info.get('class_colors', [])
-        
-        # Créer une légende colorée
-        legend_cols = st.columns(len(class_names))
-        for i, (name, color) in enumerate(zip(class_names, class_colors)):
-            with legend_cols[i]:
-                color_hex = f"#{color[0]:02x}{color[1]:02x}{color[2]:02x}"
-                st.markdown(f"""
-                <div style="
-                    background-color: {color_hex}; 
-                    padding: 10px; 
-                    border-radius: 5px; 
-                    text-align: center;
-                    color: white;
-                    font-weight: bold;
-                    text-shadow: 1px 1px 2px rgba(0,0,0,0.7);
-                    margin: 2px;
-                ">
-                    {name}
-                </div>
-                """, unsafe_allow_html=True)
-
-# # Sidebar avec informations du modèle
-# with st.sidebar:
-#     st.markdown("## 🔧 Informations du Modèle")
-    
-#     model_info = get_model_info()
-#     if model_info:
-#         st.success("✅ Modèle chargé")
-#         st.write(f"**Nom:** {model_info.get('model_name', 'N/A')}")
-#         st.write(f"**Encoder:** {model_info.get('encoder', 'N/A')}")
-#         st.write(f"**Classes:** {model_info.get('num_classes', 'N/A')}")
-#         st.write(f"**Taille d'entrée:** {model_info.get('input_size', 'N/A')}")
-#         st.write(f"**Run ID:** {model_info.get('run_id', 'N/A')[:8]}...")
-#     else:
-#         st.error("❌ Modèle non disponible")
-
-#     st.markdown("---")
-#     st.markdown("## 🎯 Classes Segmentées")
-#     if model_info and 'class_names' in model_info:
-#         for i, class_name in enumerate(model_info['class_names']):
-#             st.write(f"{i}. {class_name}")
-            
-# Sidebar avec informations du modèle
-with st.sidebar:
-    st.markdown("## 🔧 Informations du Modèle")
-    
-    model_info = get_model_info()
-    if model_info:
-        st.success("✅ Modèle chargé")
-        st.write(f"**Nom:** {model_info.get('model_name', 'N/A')}")
-        st.write(f"**Encoder:** {model_info.get('encoder', 'N/A')}")
-        st.write(f"**Classes:** {model_info.get('num_classes', 'N/A')}")
-        st.write(f"**Taille d'entrée:** {model_info.get('input_size', 'N/A')}")
-        st.write(f"**Run ID:** {model_info.get('run_id', 'N/A')[:8]}...")
+        for class_name in class_names:
+            try:
+                idx = model_class_names.index(class_name)
+                color = model_class_colors[idx]
+                colors.append([c/255.0 for c in color])
+            except (ValueError, IndexError):
+                colors.append([0.5, 0.5, 0.5])
     else:
-        st.error("❌ Modèle non disponible")
+        colors = plt.cm.Set3(np.linspace(0, 1, len(class_names)))
+    
+    # Créer le graphique
+    fig, ax = plt.subplots(figsize=(8, 5))
+    
+    bars = ax.bar(class_names, percentages, color=colors, alpha=0.8, edgecolor='white', linewidth=1)
+    
+    ax.set_xlabel('Classes de Segmentation', fontsize=10, fontweight='bold')
+    ax.set_ylabel('Pourcentage (%)', fontsize=10, fontweight='bold')
+    ax.set_title(title, fontsize=11, fontweight='bold', pad=15)
+    
+    # Rotation des labels selon le nombre de classes
+    if len(class_names) <= 4:
+        plt.xticks(rotation=0, ha='center', fontsize=9)
+    elif len(class_names) <= 6:
+        plt.xticks(rotation=30, ha='right', fontsize=8)
+    else:
+        plt.xticks(rotation=45, ha='right', fontsize=8)
+    
+    plt.yticks(fontsize=9)
+    
+    # Valeurs sur les barres
+    for bar, percentage in zip(bars, percentages):
+        height = bar.get_height()
+        ax.text(bar.get_x() + bar.get_width()/2., height + max(percentages) * 0.01,
+               f'{percentage:.1f}%',
+               ha='center', va='bottom', fontsize=8, fontweight='bold')
+    
+    plt.tight_layout()
+    ax.grid(axis='y', alpha=0.2, linestyle='--')
+    ax.set_axisbelow(True)
+    
+    # Style des axes
+    ax.spines['top'].set_visible(False)
+    ax.spines['right'].set_visible(False)
+    
+    # Couleur de fond
+    fig.patch.set_facecolor('white')
+    ax.set_facecolor('#fafafa')
+    
+    st.pyplot(fig)
+    plt.close(fig)
+
+def display_class_comparison(stats1, stats2, model1, model2):
+    """Affiche une comparaison détaillée des classes entre les deux modèles"""
+    st.markdown("### 🔍 Analyse Comparative par Classe")
+    
+    # Combiner toutes les classes présentes
+    all_classes = set(stats1.keys()) | set(stats2.keys())
+    
+    comparison_data = []
+    for class_name in sorted(all_classes):
+        pct1 = stats1.get(class_name, {}).get('percentage', 0)
+        pct2 = stats2.get(class_name, {}).get('percentage', 0)
+        diff = pct1 - pct2
+        
+        comparison_data.append({
+            'class': class_name,
+            'model1_pct': pct1,
+            'model2_pct': pct2,
+            'difference': diff
+        })
+    
+    # Afficher sous forme de tableau interactif
+    col1, col2 = st.columns([2, 1])
+    
+    with col1:
+        st.markdown("**📋 Tableau Comparatif:**")
+        for data in comparison_data:
+            diff_color = "🟢" if abs(data['difference']) < 1 else ("🔴" if abs(data['difference']) > 5 else "🟡")
+            diff_text = f"{data['difference']:+.1f}%"
+            
+            st.markdown(f"""
+            **{data['class']}** {diff_color}
+            - {model1['encoder_name']}: {data['model1_pct']:.1f}%
+            - {model2['encoder_name']}: {data['model2_pct']:.1f}%
+            - Différence: {diff_text}
+            """)
+    
+    with col2:
+        st.markdown("**🎯 Légende:**")
+        st.markdown("🟢 Différence < 1%")
+        st.markdown("🟡 Différence 1-5%")
+        st.markdown("🔴 Différence > 5%")
+
+# Sidebar avec informations des modèles
+with st.sidebar:
+    st.markdown("## 🔧 Informations des Modèles")
+    
+    models_info = get_models_info()
+    if models_info:
+        st.success("✅ Modèles chargés")
+        
+        # Modèle 1
+        best_model = models_info.get('best_model', {})
+        st.markdown("**🥇 Meilleur Modèle:**")
+        st.write(f"**Encoder:** {best_model.get('encoder_name', 'N/A')}")
+        st.write(f"**Run ID:** {best_model.get('run_id', 'N/A')[:8]}...")
+        
+        # Modèle 2
+        second_model = models_info.get('second_best_model', {})
+        st.markdown("**🥈 Deuxième Modèle:**")
+        st.write(f"**Encoder:** {second_model.get('encoder_name', 'N/A')}")
+        st.write(f"**Run ID:** {second_model.get('run_id', 'N/A')[:8]}...")
+        
+        st.write(f"**Classes:** {models_info.get('num_classes', 'N/A')}")
+    else:
+        st.error("❌ Modèles non disponibles")
 
     st.markdown("---")
     st.markdown("## 🎯 Classes Segmentées")
     
-    # ✨ NOUVEAU: Affichage avec carrés colorés
-    if model_info and 'class_names' in model_info and 'class_colors' in model_info:
-        class_names = model_info['class_names']
-        class_colors = model_info['class_colors']
+    # Affichage avec carrés colorés
+    if models_info and 'class_names' in models_info and 'class_colors' in models_info:
+        class_names = models_info['class_names']
+        class_colors = models_info['class_colors']
         
         for i, (class_name, color) in enumerate(zip(class_names, class_colors)):
-            # Convertir RGB en hex
             color_hex = f"#{color[0]:02x}{color[1]:02x}{color[2]:02x}"
             
-            # Affichage avec carré coloré
             st.markdown(f"""
             <div style="
                 display: flex;
@@ -453,56 +401,26 @@ with st.sidebar:
                 <span style="font-weight: 500;">{i}. {class_name}</span>
             </div>
             """, unsafe_allow_html=True)
-    else:
-        # Fallback si pas d'info couleur
-        if model_info and 'class_names' in model_info:
-            for i, class_name in enumerate(model_info['class_names']):
-                st.write(f"{i}. {class_name}")
-        else:
-            st.write("Informations des classes non disponibles")
 
+# États de session
+if "sample_comparison_result" not in st.session_state:
+    st.session_state.sample_comparison_result = None
+if "upload_comparison_result" not in st.session_state:
+    st.session_state.upload_comparison_result = None
 
-# Initialisation des états - MODIFIÉE pour séparer les résultats par onglet
-if "sample_prediction_result" not in st.session_state:
-    st.session_state.sample_prediction_result = None
-if "sample_selected_image" not in st.session_state:
-    st.session_state.sample_selected_image = None
+def reset_sample_comparison():
+    st.session_state.sample_comparison_result = None
 
-if "upload_prediction_result" not in st.session_state:
-    st.session_state.upload_prediction_result = None
-if "upload_uploaded_file" not in st.session_state:
-    st.session_state.upload_uploaded_file = None
+def reset_upload_comparison():
+    st.session_state.upload_comparison_result = None
 
-# ANCIEN système (gardé pour compatibilité si besoin)
-if "prediction_result" not in st.session_state:
-    st.session_state.prediction_result = None
-if "selected_image" not in st.session_state:
-    st.session_state.selected_image = None
-if "uploaded_file" not in st.session_state:
-    st.session_state.uploaded_file = None
-
-def reset_sample_state():
-    """Reset des états de session pour l'onglet Sample"""
-    st.session_state.sample_prediction_result = None
-    st.session_state.sample_selected_image = None
-
-def reset_upload_state():
-    """Reset des états de session pour l'onglet Upload"""
-    st.session_state.upload_prediction_result = None
-    st.session_state.upload_uploaded_file = None
-
-def reset_state():
-    """Reset complet (garde l'ancienne fonction pour compatibilité)"""
-    reset_sample_state()
-    reset_upload_state()
-
-# Interface principale - MODIFIÉE pour résultats indépendants
+# Interface principale
 tab1, tab2 = st.tabs(["📸 Images d'Exemple", "📤 Upload Personnel"])
 
 # Tab 1: Images d'exemple
 with tab1:
-    st.markdown("## 📸 Sélectionner une Image d'Exemple")
-    st.markdown("Ces images proviennent du dataset Cityscapes avec masques de vérité disponibles.")
+    st.markdown("## 📸 Comparaison sur Images d'Exemple")
+    st.markdown("Comparez les performances des deux modèles sur des images Cityscapes.")
     
     # Récupérer les images d'exemple
     sample_data = get_sample_images()
@@ -530,85 +448,47 @@ with tab1:
         # Prévisualisation de l'image sélectionnée
         if selected_filename:
             try:
-                # Récupérer l'image via l'API
                 image_url = f"{API_BASE_URL}/sample-image/{selected_filename}"
                 response = requests.get(image_url, timeout=5)
                 
                 if response.status_code == 200:
-                    # Charger l'image depuis les bytes
                     sample_image = Image.open(io.BytesIO(response.content))
                     st.image(
                         sample_image, 
                         caption=f"Aperçu: {selected_display_name}", 
                         use_container_width=True
                     )
-                elif response.status_code == 404:
-                    st.warning("📷 Image non trouvée sur le serveur")
-                else:
-                    st.info("📷 Prévisualisation temporairement indisponible")
-                    
-            except requests.exceptions.Timeout:
-                st.warning("⏱️ Timeout lors du chargement de l'aperçu")
-            except Exception as e:
-                st.info("📷 Aperçu non disponible - L'image sera chargée lors de la prédiction")
+            except:
+                st.info("📷 Aperçu non disponible - L'image sera chargée lors de la comparaison")
         
-        # Bouton de prédiction
-        if st.button("🔮 Lancer la Segmentation", key="predict_sample", type="primary"):
+        # Bouton de comparaison
+        if st.button("🔄 Lancer la Comparaison", key="compare_sample", type="primary"):
             if selected_filename:
-                with st.spinner("Segmentation en cours..."):
-                    result = predict_sample_image(selected_filename)
+                with st.spinner("Comparaison en cours..."):
+                    result = compare_sample_image(selected_filename)
                     if result and result['success']:
-                        # ✨ NOUVEAU: Stocker dans les variables Sample
-                        st.session_state.sample_prediction_result = result
-                        st.session_state.sample_selected_image = selected_filename
-                        st.success("✅ Segmentation terminée!")
+                        st.session_state.sample_comparison_result = result
+                        st.success("✅ Comparaison terminée!")
                         st.rerun()
                     else:
-                        st.error("❌ Erreur lors de la segmentation")
+                        st.error("❌ Erreur lors de la comparaison")
     else:
         st.warning("Aucune image d'exemple disponible")
 
-    # ✨ NOUVEAU: Affichage des résultats Sample (indépendants)
-    if st.session_state.sample_prediction_result and st.session_state.sample_selected_image:
+    # Affichage des résultats de comparaison
+    if st.session_state.sample_comparison_result:
         st.markdown("---")
-        result = st.session_state.sample_prediction_result
-        
-        st.markdown(f"## 🎯 Résultats pour: {st.session_state.sample_selected_image}")
-        
-        # Afficher la segmentation
-        if result.get('figure_data'):
-            try:
-                # Désérialiser la figure
-                fig_bytes = base64.b64decode(result['figure_data'])
-                fig = pickle.loads(fig_bytes)
-                display_segmentation_plots(fig)
-            except Exception as e:
-                st.error(f"Erreur lors de l'affichage de la segmentation: {e}")
-                # Fallback sur la légende des classes
-                display_class_legend(model_info)
-        else:
-            # Fallback si pas de figure_data
-            display_class_legend(model_info)
-        
-        # Afficher les statistiques
-        if result.get('prediction_stats'):
-            display_prediction_stats_with_chart(result['prediction_stats'], model_info)
-        
-        # Informations additionnelles
-        if result.get('ground_truth_available'):
-            st.success("🎯 Masque de vérité terrain disponible pour comparaison")
-        else:
-            st.info("ℹ️ Pas de vérité terrain disponible")
+        display_comparison_results(st.session_state.sample_comparison_result, models_info)
         
         # Bouton de reset
-        if st.button("🔄 Nouvelle Analyse", key="reset_sample", type="secondary"):
-            reset_sample_state()
+        if st.button("🔄 Nouvelle Comparaison", key="reset_sample", type="secondary"):
+            reset_sample_comparison()
             st.rerun()
 
 # Tab 2: Upload personnel
 with tab2:
-    st.markdown("## 📤 Uploader Votre Image")
-    st.markdown("Uploadez une image de rue pour analyse. Formats acceptés: JPG, PNG")
+    st.markdown("## 📤 Comparaison sur Image Personnelle")
+    st.markdown("Uploadez une image pour comparer les performances des deux modèles.")
     
     uploaded_file = st.file_uploader(
         "Sélectionnez une image:",
@@ -625,62 +505,33 @@ with tab2:
         file_size_mb = len(uploaded_file.getvalue()) / 1024 / 1024
         st.info(f"📁 **Fichier:** {uploaded_file.name} | 📊 **Taille:** {file_size_mb:.2f} MB")
         
-        # Bouton de prédiction
-        if st.button("🔮 Lancer la Segmentation", key="predict_upload", type="primary"):
-            with st.spinner("Segmentation en cours..."):
-                result = predict_uploaded_image(uploaded_file)
+        # Bouton de comparaison
+        if st.button("🔄 Lancer la Comparaison", key="compare_upload", type="primary"):
+            with st.spinner("Comparaison en cours..."):
+                result = compare_uploaded_image(uploaded_file)
                 if result and result['success']:
-                    # ✨ NOUVEAU: Stocker dans les variables Upload
-                    st.session_state.upload_prediction_result = result
-                    st.session_state.upload_uploaded_file = uploaded_file.name
-                    st.success("✅ Segmentation terminée!")
+                    st.session_state.upload_comparison_result = result
+                    st.success("✅ Comparaison terminée!")
                     st.rerun()
                 else:
-                    st.error("❌ Erreur lors de la segmentation")
+                    st.error("❌ Erreur lors de la comparaison")
 
-    # ✨ NOUVEAU: Affichage des résultats Upload (indépendants)
-    if st.session_state.upload_prediction_result and st.session_state.upload_uploaded_file:
+    # Affichage des résultats de comparaison
+    if st.session_state.upload_comparison_result:
         st.markdown("---")
-        result = st.session_state.upload_prediction_result
-        
-        st.markdown(f"## 🎯 Résultats pour: {st.session_state.upload_uploaded_file}")
-        
-        # Afficher la segmentation
-        if result.get('figure_data'):
-            try:
-                # Désérialiser la figure
-                fig_bytes = base64.b64decode(result['figure_data'])
-                fig = pickle.loads(fig_bytes)
-                display_segmentation_plots(fig)
-            except Exception as e:
-                st.error(f"Erreur lors de l'affichage de la segmentation: {e}")
-                # Fallback sur la légende des classes
-                display_class_legend(model_info)
-        else:
-            # Fallback si pas de figure_data
-            display_class_legend(model_info)
-        
-        # Afficher les statistiques
-        if result.get('prediction_stats'):
-            display_prediction_stats_with_chart(result['prediction_stats'], model_info)
-        
-        # Informations additionnelles
-        if result.get('ground_truth_available'):
-            st.success("🎯 Masque de vérité terrain disponible pour comparaison")
-        else:
-            st.info("ℹ️ Pas de vérité terrain disponible (image personnelle)")
+        display_comparison_results(st.session_state.upload_comparison_result, models_info)
         
         # Bouton de reset
-        if st.button("🔄 Nouvelle Analyse", key="reset_upload", type="secondary"):
-            reset_upload_state()
+        if st.button("🔄 Nouvelle Comparaison", key="reset_upload", type="secondary"):
+            reset_upload_comparison()
             st.rerun()
 
 # Footer
 st.markdown("---")
 st.markdown("""
 <div style="text-align: center; color: #666; padding: 1rem;">
-    🚗 <strong>Multi-Class Segmentation API</strong> - 
-    Système de vision par ordinateur pour véhicules autonomes<br>
-    <small>Développé avec FastAPI, Streamlit et TensorFlow</small>
+    🔄 <strong>Multi-Class Segmentation Comparison API</strong> - 
+    Système de comparaison de modèles de vision par ordinateur<br>
+    <small>Développé avec FastAPI, Streamlit et PyTorch</small>
 </div>
 """, unsafe_allow_html=True)
\ No newline at end of file
diff --git a/notebooks/requirements.txt b/notebooks/requirements.txt
index d6dc46d..5bb3c2f 100644
--- a/notebooks/requirements.txt
+++ b/notebooks/requirements.txt
@@ -606,13 +606,13 @@ tokenizers==0.21.4
 toml==0.10.2
 tomlkit==0.13.3
 toolz==0.12.1
-torch==2.8.0+cu126
+torch==2.8.0
 torchao==0.10.0
-torchaudio==2.8.0+cu126
+torchaudio==2.8.0
 torchdata==0.11.0
 torchsummary==1.5.1
 torchtune==0.6.1
-torchvision==0.23.0+cu126
+torchvision==0.23.0
 tornado==6.4.2
 tqdm==4.67.1
 traitlets==5.7.1
diff --git a/utils/utils.py b/utils/utils.py
index c3cb60e..39047f5 100644
--- a/utils/utils.py
+++ b/utils/utils.py
@@ -3,7 +3,6 @@ import json
 import tempfile
 import numpy as np
 import torch
-import torch.nn.functional as F
 import matplotlib.pyplot as plt
 import segmentation_models_pytorch as smp
 import mlflow
@@ -16,8 +15,8 @@ import matplotlib.patches as mpatches
 from matplotlib.colors import ListedColormap
 
 # Preprocessing/Images
-from PIL import Image
-import torchvision.transforms as transforms
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
 
 
 def load_cityscapes_config(config_path="../cityscapes_config.json", verbose=True):
@@ -141,10 +140,6 @@ def get_preprocessing_fn(img_size=(512, 512), encoder=None):
     Retourne une fonction qui prétraite à la fois les images et les masques
     pour l'inférence, en reproduisant exactement le pipeline CityscapesDataset.
     """
-    import albumentations as A
-    from albumentations.pytorch import ToTensorV2
-    import cv2
-    import torch
     
     if encoder is None:
         mean, std = [0.0, 0.0, 0.0], [1.0, 1.0, 1.0]
@@ -371,7 +366,7 @@ def preprocess_image_and_mask(image_path, mask_path=None, img_size=(512, 512),
 
 # Version mise à jour de run_inference_and_visualize
 def run_inference_and_visualize(image_path, model, encoder_name, img_size, mapping_config, 
-                               mask_path=None, save_dir=None, show_stats=True):
+                               mask_path=None, save_dir=None, show_stats=True, show_plot=True):
     """
     Version simplifiée utilisant preprocess_image_and_mask
     """
@@ -461,7 +456,6 @@ def run_inference_and_visualize(image_path, model, encoder_name, img_size, mappi
         fig.savefig(fig_path, dpi=150, bbox_inches='tight')
         print(f"💾 Résultat sauvegardé: {fig_path}")
 
-    plt.show()
 
     # Statistiques
     if show_stats:
@@ -481,7 +475,11 @@ def run_inference_and_visualize(image_path, model, encoder_name, img_size, mappi
                     percentage = (count / mask_gt_np.size) * 100
                     print(f"   {mapping_config['group_names'][class_id]}: {count} pixels ({percentage:.1f}%)")
 
-    return predicted_mask, inference_time
+    if show_plot:
+        plt.show()
+        return predicted_mask, inference_time
+    else:
+        return predicted_mask, inference_time, fig
 
 
 def load_model(run_id=None, experiment_name="OC Projet 9", 
@@ -736,6 +734,6 @@ def load_second_best_model(experiment_name="OC Projet 9", metric="test_mean_iou"
 
 def load_model_by_run_id(run_id):
     """Charge un modèle par run_id spécifique"""
-    return load_model_by_param(run_id=run_id)
+    return load_model(run_id=run_id)