Update model

app.py

@@ -13,12 +13,12 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torchvision import transforms
-from tensorflow.keras import layers, models
 
 app = Flask(__name__)
 
 CLASSES = ["buildings", "forest", "glacier", "mountain", "sea", "street"]
-IMG_SIZE = 150
+PYTORCH_IMG_SIZE = 150
+TF_IMG_SIZE = 228
 
 _pytorch_model = None
 _tf_model = None
@@ -26,78 +26,64 @@ _tf_model = None
 
 class CNN_Torch(nn.Module):
     """
-    CNN amélioré 4 blocs
-    Entrée : (B, 3, 150, 150)
-    Sortie : (B, num_classes)
+    CNN PyTorch 4 blocs pour images RGB (3 canaux, 150×150).
+    Entrée  : (B, 3, 150, 150) — normalisée ImageNet (mean/std)
+    Sortie  : (B, num_classes) — logits bruts (CrossEntropyLoss)
+
+    Architecture :
+        Block 1 : Conv(3→32)×2   + BN + ReLU + MaxPool(2)            150→75
+        Block 2 : Conv(32→64)×2  + BN + ReLU + MaxPool(2) + Drop2d   75→37
+        Block 3 : Conv(64→128)×2 + BN + ReLU + MaxPool(2) + Drop2d   37→18
+        Block 4 : Conv(128→256)×2+ BN + ReLU + MaxPool(2) + Drop2d   18→9
+        GAP     : AdaptiveAvgPool2d(1)                                  →(B,256)
+        Head    : Linear(256→256) + ReLU + Dropout + Linear(256→C)
     """
-
-    def __init__(self, num_classes=6):
+    def __init__(self, num_classes: int = 6):
         super().__init__()
 
         self.features = nn.Sequential(
-            # Block 1: 150 -> 75
-            nn.Conv2d(3, 32, kernel_size=3, padding=1),
-            nn.BatchNorm2d(32),
-            nn.ReLU(inplace=True),
-
-            nn.Conv2d(32, 32, kernel_size=3, padding=1),
-            nn.BatchNorm2d(32),
-            nn.ReLU(inplace=True),
-
-            nn.MaxPool2d(2),
-
-            # Block 2: 75 -> 37
-            nn.Conv2d(32, 64, kernel_size=3, padding=1),
-            nn.BatchNorm2d(64),
-            nn.ReLU(inplace=True),
-
-            nn.Conv2d(64, 64, kernel_size=3, padding=1),
-            nn.BatchNorm2d(64),
-            nn.ReLU(inplace=True),
-
-            nn.MaxPool2d(2),
-            nn.Dropout2d(0.10),
-
-            # Block 3: 37 -> 18
-            nn.Conv2d(64, 128, kernel_size=3, padding=1),
-            nn.BatchNorm2d(128),
-            nn.ReLU(inplace=True),
-
-            nn.Conv2d(128, 128, kernel_size=3, padding=1),
-            nn.BatchNorm2d(128),
-            nn.ReLU(inplace=True),
-
+            # Block 1 — 150×150 → 75×75
+            nn.Conv2d(3, 32, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(32), nn.ReLU(inplace=True),
+            nn.Conv2d(32, 32, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(32), nn.ReLU(inplace=True),
             nn.MaxPool2d(2),
-            nn.Dropout2d(0.15),
 
-            # Block 4: 18 -> 9
-            nn.Conv2d(128, 256, kernel_size=3, padding=1),
-            nn.BatchNorm2d(256),
-            nn.ReLU(inplace=True),
-
-            nn.Conv2d(256, 256, kernel_size=3, padding=1),
-            nn.BatchNorm2d(256),
-            nn.ReLU(inplace=True),
-
-            nn.MaxPool2d(2),
-            nn.Dropout2d(0.20),
+            # Block 2 — 75×75 → 37×37
+            nn.Conv2d(32, 64, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(64), nn.ReLU(inplace=True),
+            nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(64), nn.ReLU(inplace=True),
+            nn.MaxPool2d(2), nn.Dropout2d(0.10),
+
+            # Block 3 — 37×37 → 18×18
+            nn.Conv2d(64, 128, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(128), nn.ReLU(inplace=True),
+            nn.Conv2d(128, 128, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(128), nn.ReLU(inplace=True),
+            nn.MaxPool2d(2), nn.Dropout2d(0.15),
+
+            # Block 4 — 18×18 → 9×9
+            nn.Conv2d(128, 256, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(256), nn.ReLU(inplace=True),
+            nn.Conv2d(256, 256, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(256), nn.ReLU(inplace=True),
+            nn.MaxPool2d(2), nn.Dropout2d(0.20),
         )
 
+        # (B,256,9,9) → (B,256,1,1) → (B,256)
         self.gap = nn.AdaptiveAvgPool2d(1)
 
         self.classifier = nn.Sequential(
             nn.Flatten(),
             nn.Linear(256, 256),
             nn.ReLU(inplace=True),
-            nn.Dropout(0.3),
-            nn.Linear(256, num_classes)
+            nn.Dropout(0.30),
+            nn.Linear(256, num_classes),
         )
 
     def forward(self, x):
-        x = self.features(x)
-        x = self.gap(x)
-        x = self.classifier(x)
-        return x
+        return self.classifier(self.gap(self.features(x)))
 
 
 def load_pytorch():
@@ -113,7 +99,7 @@ def load_pytorch():
     model.eval()
 
     tf_transform = transforms.Compose([
-        transforms.Resize((IMG_SIZE, IMG_SIZE)),
+        transforms.Resize((PYTORCH_IMG_SIZE, PYTORCH_IMG_SIZE)),
         transforms.ToTensor(),
         transforms.Normalize(
             [0.485, 0.456, 0.406],
@@ -125,44 +111,11 @@ def load_pytorch():
     return _pytorch_model
 
 
-
-def build_cnn_tf(num_classes: int = 6, input_shape: tuple = (228, 228, 3)):
-    return models.Sequential([
-        layers.Input(shape=input_shape),
-
-        layers.Conv2D(32, (5, 5), activation="relu"),
-        layers.MaxPooling2D(2, 2),
-
-        layers.Conv2D(32, (5, 5), activation="relu"),
-        layers.MaxPooling2D(2, 2),
-
-        layers.Conv2D(32, (3, 3), activation="relu"),
-        layers.MaxPooling2D(2, 2),
-
-        layers.Conv2D(64, (3, 3), activation="relu"),
-        layers.MaxPooling2D(2, 2),
-
-        layers.Conv2D(64, (3, 3), activation="relu"),
-        layers.MaxPooling2D(2, 2),
-
-        layers.Flatten(),
-        layers.Dense(1024, activation="relu"),
-        layers.Dropout(0.20),
-        layers.Dense(124, activation="relu"),
-        layers.Dropout(0.20),
-        layers.Dense(num_classes, activation="softmax"),
-    ])
-
-
 def load_tensorflow():
     global _tf_model
-    if _tf_model is not None:
-        return _tf_model
-
-    model = build_cnn_tf(num_classes=6, input_shape=(228, 228, 3))
-    model.load_weights("parfait_model.keras")
-
-    _tf_model = model
+    if _tf_model is None:
+        import tensorflow as tf
+        _tf_model = tf.keras.models.load_model("parfait_model.keras", compile=False)
     return _tf_model
     
 
@@ -205,9 +158,9 @@ def predict():
         elif framework == "tensorflow":
             model = load_tensorflow()
             arr = np.array(
-                pil_img.resize((IMG_SIZE, IMG_SIZE)),
+                pil_img.resize((TF_IMG_SIZE, TF_IMG_SIZE)),
                 dtype=np.float32
-            )
+            ) / 255.0
             arr = np.expand_dims(arr, axis=0)
             probs = model.predict(arr, verbose=0)[0]
 
@@ -234,4 +187,4 @@ def predict():
 
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", 5000))
-    app.run(host="0.0.0.0", port=port, debug=False)
+    app.run(host="0.0.0.0", port=port, debug=False)