Update predict.py

predict.py

@@ -1,27 +1,29 @@
-
 import argparse
 import numpy as np
 from pathlib import Path
 from PIL import Image
+import tensorflow as tf
+import numpy as np
+import torch
+from torchvision import transforms
+from model_pytorch import SaraCNN
+
+
+
+
 
 CLASS_NAMES = ["buildings", "forest", "glacier", "mountain", "sea", "street"]
-IMAGE_SIZE  = (150, 150)
+IMAGE_SIZE = (150, 150)
 MEAN = [0.485, 0.456, 0.406]
-STD  = [0.229, 0.224, 0.225]
+STD = [0.229, 0.224, 0.225]
 
 
-# ════════════════════════════════════════════════════════════════════════════
-#  PyTorch Prediction
-# ════════════════════════════════════════════════════════════════════════════
 
+#  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
-
     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)
 
     model = SaraCNN(num_classes=len(class_names))
@@ -35,14 +37,14 @@ def predict_pytorch(image_path: str, model_path: str = "sara_model.pth") -> dict
         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])
 
@@ -56,25 +58,18 @@ def predict_pytorch(image_path: str, model_path: str = "sara_model.pth") -> dict
     }
 
 
-# ════════════════════════════════════════════════════════════════════════════
-#  TensorFlow Prediction
-# ════════════════════════════════════════════════════════════════════════════
 
+#  TensorFlow Prediction
 def predict_tensorflow(image_path: str, model_path: str = "sara_model.keras") -> dict:
-    import tensorflow as tf
-    import numpy as np
-
-    # Format .h5 = compatible toutes versions Keras, chargement simple
     model = tf.keras.models.load_model(model_path, compile=False)
 
-    # Prétraitement
     img = tf.keras.utils.load_img(image_path, target_size=(150, 150))
     arr = tf.keras.utils.img_to_array(img) / 255.0
     arr = np.expand_dims(arr, axis=0)
 
-    # Prédiction
-    probs      = model.predict(arr, verbose=0)[0]
-    idx        = int(np.argmax(probs))
+   
+    probs = model.predict(arr, verbose=0)[0]
+    idx = int(np.argmax(probs))
     pred_class = CLASS_NAMES[idx]
     confidence = float(probs[idx])
 
@@ -88,9 +83,7 @@ def predict_tensorflow(image_path: str, model_path: str = "sara_model.keras") ->
     }
 
 
-# ════════════════════════════════════════════════════════════════════════════
-#  CLI
-# ════════════════════════════════════════════════════════════════════════════
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -100,15 +93,15 @@ if __name__ == "__main__":
     args = parser.parse_args()
 
     if args.model == "pytorch":
-        path   = args.model_path or "sara_model.pth"
+        path = args.model_path or "sara_model.pth"
         result = predict_pytorch(args.image, model_path=path)
     else:
-        path   = args.model_path or "sara_model.keras"
+        path = args.model_path or "sara_model.keras"
         result = predict_tensorflow(args.image, model_path=path)
 
     print(f"\n Classe     : {result['predicted_class']}")
-    print(f" Confiance  : {result['confidence']}%")
-    print("\n Toutes les probabilités :")
+    print(f" Confidence  : {result['confidence']}%")
+    print("\n All probabilities :")
     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}")