Update app.py

app.py

@@ -10,34 +10,43 @@ import os
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # Define available models and their corresponding file names
-model_options = {
-    "ResNet-18": (models.resnet18, "resnet18_model.pth"),
-    "MobileNetV3 Large": (models.mobilenet_v3_large, "best_model3_mobilenetv3_large.pth"),
-    "MobileNetV3 Small": (models.mobilenet_v3_small, "best_model13_mobilenetv3_small.pth")
+model_paths = {
+    "ResNet-18": "resnet18_model.pth",
+    "MobileNetV3 Large": "best_model3_mobilenetv3_large.pth",
+    "MobileNetV3 Small": "best_model13_mobilenetv3_small.pth"
 }
 
-classes_name = ['AI-generated Image', 'Real Image']
-
-def load_model(model_name):
-    model_func, model_path = model_options[model_name]
-    model = model_func(weights=None)  # Load model without pretrained weights
-    
-    if "resnet" in model_name.lower():
+def initialize_model(model_name):
+    if model_name == "ResNet-18":
+        model = models.resnet18(weights=None)
         num_ftrs = model.fc.in_features
         model.fc = nn.Sequential(
             nn.Dropout(p=0.5),
             nn.Linear(num_ftrs, 2)
         )
-    else:  # For MobileNetV3
+    elif model_name == "MobileNetV3 Large":
+        model = models.mobilenet_v3_large(weights=None)
+        num_ftrs = model.classifier[-1].in_features
+        model.classifier[-1] = nn.Linear(num_ftrs, 2)
+    elif model_name == "MobileNetV3 Small":
+        model = models.mobilenet_v3_small(weights=None)
         num_ftrs = model.classifier[-1].in_features
         model.classifier[-1] = nn.Linear(num_ftrs, 2)
+    else:
+        raise ValueError("Invalid model name")
     
+    return model
+
+def load_model(model_name):
+    model = initialize_model(model_name)
+    model_path = model_paths[model_name]
     model.load_state_dict(torch.load(model_path, map_location=device))
     model = model.to(device)
     model.eval()
-    
     return model
 
+classes_name = ['AI-generated Image', 'Real Image']
+
 # Define preprocessing transformations
 preprocess = transforms.Compose([
     transforms.Lambda(lambda img: img.convert('RGB') if img.mode in ('P', 'RGBA') else img),
@@ -63,10 +72,10 @@ def classify_image(image, model_name):
 
 # Gradio interface
 image_input = gr.Image(image_mode="RGB")
-model_choice = gr.Dropdown(choices=list(model_options.keys()), label="Choose Model", value="ResNet-18")
+model_choice = gr.Dropdown(choices=list(model_paths.keys()), label="Choose Model", value="ResNet-18")
 output_text = gr.Textbox()
 
 gr.Interface(fn=classify_image, inputs=[image_input, model_choice], outputs=[output_text],
              title="AI-Generated Image Detector",
              description="Upload an image and choose a model to detect AI-generated images.",
-             theme="default").launch()
+             theme="default").launch()