Update app.py

app.py

@@ -11,13 +11,18 @@ import torch
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # Load the main classifier (Detector_best_model.pth)
-main_model = models.resnet18(weights=None)  # Updated: weights=None
-num_ftrs = main_model.fc.in_features
+main_model = models.mobilenet_v3_large(weights=None)  # Updated: weights=None
+
+#num_ftrs = main_model.fc.in_features
 # main_model.fc = nn.Linear(num_ftrs, 2)  # 2 classes: AI-generated_Image, Real_Image
-main_model.fc = nn.Sequential(
-    nn.Dropout(p=0.5),  # Match the training architecture
-    nn.Linear(num_ftrs, 2)  # 2 classes: AI-generated Image, Real Image
-)
+
+num_ftrs = main_model.classifier[3].in_features
+main_model.classifier[3] = nn.Linear(num_ftrs, 2)
+
+# main_model.fc = nn.Sequential(
+#     nn.Dropout(p=0.5),  # Match the training architecture
+#     nn.Linear(num_ftrs, 2)  # 2 classes: AI-generated Image, Real Image
+# )
 
 main_model.load_state_dict(torch.load('best_model6for_RESNET18_After_First_Half_training_part.pth', map_location=device, weights_only=True))  # Updated: weights_only=True
 main_model = main_model.to(device)