Update app.py

app.py

@@ -5,68 +5,66 @@ import torch.nn as nn
 import torchvision.transforms as transforms
 import torchvision.models as models
 import os
-import torch
 
 # Set device
 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.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
-)
-
-main_model.load_state_dict(torch.load('best_model9.pth', map_location=device, weights_only=True))  # Updated: weights_only=True
-main_model = main_model.to(device)
-main_model.eval()
+# Define available models and their corresponding file names
+model_options = {
+    "ResNet-18": (models.resnet18, "resnet18_model.pth"),
+    "MobileNetV3 Large": (models.mobilenet_v3_large, "mobilenet_v3_large_model.pth"),
+    "MobileNetV3 Small": (models.mobilenet_v3_small, "mobilenet_v3_small_model.pth")
+}
 
-# Define class names for the classifier based on the Folder structure
 classes_name = ['AI-generated Image', 'Real Image']
 
-def convert_to_rgb(image):
-    """
-    Converts 'P' mode images with transparency to 'RGBA', and then to 'RGB'.
-    This is to avoid transparency issues during model training.
-    """
-    if image.mode in ('P', 'RGBA'):
-        return image.convert('RGB')
-    return 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():
+        num_ftrs = model.fc.in_features
+        model.fc = nn.Sequential(
+            nn.Dropout(p=0.5),
+            nn.Linear(num_ftrs, 2)
+        )
+    else:  # For MobileNetV3
+        num_ftrs = model.classifier[-1].in_features
+        model.classifier[-1] = nn.Linear(num_ftrs, 2)
+    
+    model.load_state_dict(torch.load(model_path, map_location=device))
+    model = model.to(device)
+    model.eval()
+    
+    return model
 
-# Define preprocessing transformations (same used during training)
+# Define preprocessing transformations
 preprocess = transforms.Compose([
-    transforms.Lambda(convert_to_rgb),
-    transforms.Resize((224, 224)),  # Resize here, no need for shape argument in gr.Image
+    transforms.Lambda(lambda img: img.convert('RGB') if img.mode in ('P', 'RGBA') else img),
+    transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # ImageNet normalization
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 ])
 
-def classify_image(image):
-    # Open the image using PIL
-    image = Image.fromarray(image)
+def classify_image(image, model_name):
+    model = load_model(model_name)
     
-    # Preprocess the image
+    image = Image.fromarray(image)
     input_image = preprocess(image).unsqueeze(0).to(device)
     
-    # Perform inference with the main classifier
     with torch.no_grad():
-        output = main_model(input_image)
+        output = model(input_image)
         probabilities = torch.nn.functional.softmax(output[0], dim=0)
         confidence, predicted_class = torch.max(probabilities, 0)
     
-    # Main classifier result
-    main_prediction = classes_name[predicted_class]
-    main_confidence = confidence.item()
-    
-    return f"Image is : {main_prediction} (Confidence: {main_confidence:.4f})"
+    return f"Image is: {classes_name[predicted_class]} (Confidence: {confidence.item():.4f})"
 
-# Gradio interface (updated)
-image_input = gr.Image(image_mode="RGB")  # Removed shape argument
+# Gradio interface
+image_input = gr.Image(image_mode="RGB")
+model_choice = gr.Radio(choices=list(model_options.keys()), label="Choose Model", value="ResNet-18")
 output_text = gr.Textbox()
 
-gr.Interface(fn=classify_image, inputs=image_input, outputs=[output_text], 
-             title="Detect AI-generated Image ",
-             description="Upload an image to Detected AI-generated Image .",
-             theme="default").launch()
+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()