Update

app.py

@@ -8,6 +8,11 @@ import numpy as np
 from flask import Flask, jsonify, render_template, request
 from PIL import Image
 
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import transforms
+
 app = Flask(__name__)
 
 CLASSES  = ["buildings", "forest", "glacier", "mountain", "sea", "street"]
@@ -23,31 +28,62 @@ def load_pytorch():
     if _pytorch_model is not None:
         return _pytorch_model
 
-    import torch
-    import torch.nn as nn
-    import torch.nn.functional as F
-    from torchvision import transforms
-
-    class CNN_Torch(nn.Module):
-        def __init__(self, num_classes=6):
-            super().__init__()
-            self.conv1      = nn.Conv2d(3, 32, kernel_size=3, padding=1)
-            self.conv2      = nn.Conv2d(32, 64, kernel_size=3, padding=1)
-            self.conv3      = nn.Conv2d(64, 128, kernel_size=3, padding=1)
-            self.conv3_drop = nn.Dropout2d(p=0.25)
-            self.pool       = nn.MaxPool2d(2, 2)
-            self.fc1        = nn.Linear(128 * 18 * 18, 256)
-            self.fc2        = nn.Linear(256, num_classes)
-
-        def forward(self, x):
-            x = self.pool(F.relu(self.conv1(x)))
-            x = self.pool(F.relu(self.conv2(x)))
-            x = self.pool(F.relu(self.conv3_drop(self.conv3(x))))
-            x = x.view(-1, 128 * 18 * 18)
-            x = F.relu(self.fc1(x))
-            x = F.dropout(x, training=self.training)
-            x = self.fc2(x)
-            return F.log_softmax(x, dim=1)
+
+class CNN_Torch(nn.Module):
+    """
+    CNN PyTorch allégé pour images RGB (3 canaux).
+    Entrée  : (B, 3, 150, 150)
+    Sortie  : (B, num_classes) — log-softmax
+
+    Architecture :
+        Block 1 : Conv2d(3→32)  + BN + ReLU + MaxPool2d(2)          → 75×75
+        Block 2 : Conv2d(32→64) + BN + ReLU + MaxPool2d(2) + Drop2d → 37×37
+        Block 3 : Conv2d(64→128)+ BN + ReLU + MaxPool2d(2) + Drop2d → 18×18
+        GAP     : AdaptiveAvgPool2d(1)                               → (B,128)
+        Head    : Linear(128→256) + ReLU + Dropout + Linear(256→C)
+
+    Paramètre `dropout` contrôlé depuis l'extérieur → utilisé dans le CV.
+    """
+    def __init__(self, num_classes: int = 6, dropout: float = 0.5):
+        super().__init__()
+
+        self.features = nn.Sequential(
+            # Block 1 — 150 → 75
+            nn.Conv2d(3, 32, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(32),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),
+
+            # Block 2 — 75 → 37
+            nn.Conv2d(32, 64, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),
+            nn.Dropout2d(0.1),
+
+            # Block 3 — 37 → 18
+            nn.Conv2d(64, 128, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(128),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),
+            nn.Dropout2d(0.2),
+        )
+
+        self.gap = nn.AdaptiveAvgPool2d(1)   # (B, 128, 18, 18) → (B, 128, 1, 1)
+
+        self.classifier = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(128, 256),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dropout),
+            nn.Linear(256, num_classes),
+        )
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.gap(x)
+        x = self.classifier(x)
+        return F.log_softmax(x, dim=1)
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model  = CNN_Torch(6).to(device)