Update predict.py

predict.py

@@ -1,43 +1,36 @@
+"""
+predict.py
+----------
+Charge un modèle sauvegardé et prédit la classe d'une image.
+Corrigé pour gérer les incompatibilités de version Keras.
+"""
+
 import argparse
 import numpy as np
 from pathlib import Path
 from PIL import Image
-import torch
-from torchvision import transforms
-from model_pytorch import SaraCNN
-import tensorflow as tf
-
 
 CLASS_NAMES = ["buildings", "forest", "glacier", "mountain", "sea", "street"]
-
-IMAGE_SIZE = (150, 150)  
-MEAN = [0.485, 0.456, 0.406] 
-STD = [0.229, 0.224, 0.225]
+IMAGE_SIZE  = (150, 150)
+MEAN = [0.485, 0.456, 0.406]
+STD  = [0.229, 0.224, 0.225]
 
 
+# ════════════════════════════════════════════════════════════════════════════
+#  PyTorch Prediction
+# ════════════════════════════════════════════════════════════════════════════
 
+def predict_pytorch(image_path: str, model_path: str = "sara_model.pth") -> dict:
+    import torch
+    from torchvision import transforms
+    from model_pytorch import SaraCNN
 
-#  PyTorch Prediction
-def predict_pytorch(image_path: str,
-                    model_path: str = "sara_model.pth") -> dict:
-    """
-    Load the PyTorch checkpoint and return class probabilities.
-
-    Args:
-        image_path : path to the input image
-        model_path : path to the saved .pth file
-
-    Returns:
-        dict with keys: predicted_class, confidence, all_probabilities
-    """
-    
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    checkpoint = torch.load(model_path, map_location=device)
+    checkpoint  = torch.load(model_path, map_location=device)
     class_names = checkpoint.get("class_names", CLASS_NAMES)
-    num_classes = len(class_names)
 
-    model = SaraCNN(num_classes=num_classes)
+    model = SaraCNN(num_classes=len(class_names))
     model.load_state_dict(checkpoint["model_state"])
     model.to(device)
     model.eval()
@@ -48,62 +41,116 @@ def predict_pytorch(image_path: str,
         transforms.Normalize(MEAN, STD),
     ])
 
-    img = Image.open(image_path).convert("RGB")
+    img    = Image.open(image_path).convert("RGB")
     tensor = transform(img).unsqueeze(0).to(device)
 
     with torch.no_grad():
         logits = model(tensor)
-        probs = torch.softmax(logits, dim=1).squeeze().cpu().numpy()
+        probs  = torch.softmax(logits, dim=1).squeeze().cpu().numpy()
 
-    idx = int(np.argmax(probs))
+    idx        = int(np.argmax(probs))
     pred_class = class_names[idx]
     confidence = float(probs[idx])
 
     return {
-        "predicted_class": pred_class,
-        "confidence":      round(confidence * 100, 2),
-        "all_probabilities": {
+        "predicted_class":    pred_class,
+        "confidence":         round(confidence * 100, 2),
+        "all_probabilities":  {
             cls: round(float(p) * 100, 2)
             for cls, p in zip(class_names, probs)
         },
     }
 
 
+# ════════════════════════════════════════════════════════════════════════════
+#  TensorFlow Prediction — corrigé pour incompatibilité Keras 2.x / 3.x
+# ════════════════════════════════════════════════════════════════════════════
+
+def predict_tensorflow(image_path: str, model_path: str = "sara_model.keras") -> dict:
+    import tensorflow as tf
+
+    # ── Couche Dense patchée qui ignore quantization_config ───────────────
+    # Problème : Keras 3.x sauvegarde quantization_config=None dans le .keras
+    # mais Keras 2.x ne connaît pas cet argument → erreur de désérialisation
+    # Solution : on surcharge Dense pour accepter et ignorer cet argument
+    class PatchedDense(tf.keras.layers.Dense):
+        def __init__(self, *args, **kwargs):
+            # Supprimer les arguments inconnus avant d'appeler le vrai Dense
+            kwargs.pop("quantization_config", None)
+            super().__init__(*args, **kwargs)
+
+    # ── Couche SeparableConv2D patchée (même raison) ───────────────────────
+    class PatchedSeparableConv2D(tf.keras.layers.SeparableConv2D):
+        def __init__(self, *args, **kwargs):
+            kwargs.pop("quantization_config", None)
+            super().__init__(*args, **kwargs)
+
+    # ── Couche BatchNormalization patchée ──────────────────────────────────
+    class PatchedBatchNorm(tf.keras.layers.BatchNormalization):
+        def __init__(self, *args, **kwargs):
+            kwargs.pop("quantization_config", None)
+            super().__init__(*args, **kwargs)
+
+    # ── Chargement avec les couches patchées ──────────────────────────────
+    custom_objects = {
+        "Dense":               PatchedDense,
+        "SeparableConv2D":     PatchedSeparableConv2D,
+        "BatchNormalization":  PatchedBatchNorm,
+    }
 
-
-
-
-
-#  TensorFlow Prediction
-def predict_tensorflow(image_path: str,
-                       model_path: str = "sara_model.keras") -> dict:
-    """
-    Load the Keras model and return class probabilities.
-
-    Args:
-        image_path : path to the input image
-        model_path : path to the saved .keras file
-
-    Returns:
-        dict with keys: predicted_class, confidence, all_probabilities
-    """
-    
-    model = tf.keras.models.load_model(model_path)
-
+    try:
+        # Essai 1 : chargement normal (fonctionne si versions compatibles)
+        model = tf.keras.models.load_model(model_path)
+        print("[TF] Modèle chargé normalement.")
+
+    except Exception as e1:
+        print(f"[TF] Chargement normal échoué ({e1}), tentative avec custom_objects...")
+
+        try:
+            # Essai 2 : avec les couches patchées
+            model = tf.keras.models.load_model(
+                model_path,
+                custom_objects=custom_objects
+            )
+            print("[TF] Modèle chargé avec custom_objects.")
+
+        except Exception as e2:
+            print(f"[TF] Echec avec custom_objects ({e2}), tentative safe_mode=False...")
+
+            try:
+                # Essai 3 : safe_mode=False (Keras 3.x uniquement)
+                model = tf.keras.models.load_model(
+                    model_path,
+                    safe_mode=False,
+                    custom_objects=custom_objects
+                )
+                print("[TF] Modèle chargé avec safe_mode=False.")
+
+            except Exception as e3:
+                raise RuntimeError(
+                    f"Impossible de charger le modèle TensorFlow.\n"
+                    f"Essai 1 : {e1}\n"
+                    f"Essai 2 : {e2}\n"
+                    f"Essai 3 : {e3}\n\n"
+                    f"Solution : Ré-entraîne et sauvegarde le modèle avec "
+                    f"la même version de TensorFlow que celle du serveur."
+                )
+
+    # ── Prétraitement de l'image ──────────────────────────────────────────
     img = tf.keras.utils.load_img(image_path, target_size=IMAGE_SIZE)
-    arr = tf.keras.utils.img_to_array(img) 
-    arr = arr / 255.0                         
-    arr = np.expand_dims(arr, axis=0)  
+    arr = tf.keras.utils.img_to_array(img)   # (H, W, 3), valeurs 0-255
+    arr = arr / 255.0                         # normalisation → 0-1
+    arr = np.expand_dims(arr, axis=0)         # ajout dimension batch → (1, H, W, 3)
 
-   
-    probs = model.predict(arr, verbose=0)[0] 
-    idx = int(np.argmax(probs))
+    # ── Inférence ─────────────────────────────────────────────────────────
+    probs      = model.predict(arr, verbose=0)[0]
+    idx        = int(np.argmax(probs))
     pred_class = CLASS_NAMES[idx]
     confidence = float(probs[idx])
 
     return {
-        "predicted_class": pred_class,
-        "confidence":      round(confidence * 100, 2),
+        "predicted_class":   pred_class,
+        "confidence":        round(confidence * 100, 2),
         "all_probabilities": {
             cls: round(float(p) * 100, 2)
             for cls, p in zip(CLASS_NAMES, probs)
@@ -111,20 +158,16 @@ def predict_tensorflow(image_path: str,
     }
 
 
-
-
-
-#  CLI Entry Point
-def parse_args():
-    p = argparse.ArgumentParser(description="Predict image class")
-    p.add_argument("--model", required=True, choices=["pytorch", "tensorflow"])
-    p.add_argument("--image", required=True, help="Path to the image file")
-    p.add_argument("--model_path", default=None, help="Override default model file path")
-    return p.parse_args()
-
+# ════════════════════════════════════════════════════════════════════════════
+#  CLI
+# ════════════════════════════════════════════════════════════════════════════
 
 if __name__ == "__main__":
-    args = parse_args()
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model",      required=True, choices=["pytorch", "tensorflow"])
+    parser.add_argument("--image",      required=True)
+    parser.add_argument("--model_path", default=None)
+    args = parser.parse_args()
 
     if args.model == "pytorch":
         path   = args.model_path or "sara_model.pth"
@@ -133,9 +176,9 @@ if __name__ == "__main__":
         path   = args.model_path or "sara_model.keras"
         result = predict_tensorflow(args.image, model_path=path)
 
-    print(f"\n Predicted class : {result['predicted_class']}")
-    print(f" Confidence      : {result['confidence']}%")
-    print("\n All probabilities:")
+    print(f"\n Classe     : {result['predicted_class']}")
+    print(f" Confiance  : {result['confidence']}%")
+    print("\n Toutes les probabilités :")
     for cls, prob in sorted(result["all_probabilities"].items(), key=lambda x: -x[1]):
         bar = " " * int(prob / 5)
-        print(f"{cls:<12} {prob:6.2f}%  {bar}")
+        print(f"   {cls:<12} {prob:6.2f}%  {bar}")