Update app.py

app.py

@@ -13,28 +13,37 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 # Model Architecture מדויקת
 # -------------------------------
 model = nn.Sequential(
-    nn.Conv2d(3, 16, kernel_size=3, padding=1),
+    nn.Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
+    nn.BatchNorm2d(16),
     nn.ReLU(),
-    nn.MaxPool2d(2, 2),
-    
-    nn.Conv2d(16, 32, kernel_size=3, padding=1),
+    nn.MaxPool2d(kernel_size=2, stride=2),
+
+    nn.Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
+    nn.BatchNorm2d(32),
     nn.ReLU(),
-    nn.MaxPool2d(2, 2),
-    
-    nn.Conv2d(32, 64, kernel_size=3, padding=1),
+    nn.MaxPool2d(kernel_size=2, stride=2),
+
+    nn.Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
+    nn.BatchNorm2d(64),
     nn.ReLU(),
-    nn.MaxPool2d(2, 2),
-    
-    nn.Conv2d(64, 128, kernel_size=3, padding=1),
+    nn.MaxPool2d(kernel_size=2, stride=2),
+
+    nn.Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
+    nn.BatchNorm2d(128),
     nn.ReLU(),
-    nn.MaxPool2d(2, 2),
-    
-    nn.Flatten(),
-    nn.Dropout(0.6),
-    nn.Linear(128*14*14, 128),  # מתאים ל-224x224 קלט
+    nn.MaxPool2d(kernel_size=2, stride=2),
+
+    nn.Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
+    nn.BatchNorm2d(256),
+    nn.ReLU(),
+    nn.MaxPool2d(kernel_size=2, stride=2),
+
+    nn.Flatten(start_dim=1, end_dim=-1),
+    nn.Dropout(p=0.5),
+    nn.Linear(in_features=12544, out_features=256),
     nn.ReLU(),
-    nn.Dropout(0.4),
-    nn.Linear(128, 2)
+    nn.Dropout(p=0.5),
+    nn.Linear(in_features=256, out_features=2)
 ).to(device)
 
 model.load_state_dict(torch.load("cnn_model.pth", map_location=device))