Update app.py

app.py

@@ -1,19 +1,25 @@
 import gradio as gr
 from transformers import pipeline, AutoImageProcessor, Swinv2ForImageClassification
 from torchvision import transforms
+import torch
+
+# Ensure using GPU if available
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # Load the model and processor
 image_processor = AutoImageProcessor.from_pretrained("haywoodsloan/ai-image-detector-deploy")
 model = Swinv2ForImageClassification.from_pretrained("haywoodsloan/ai-image-detector-deploy")
-clf = pipeline(model=model, task="image-classification", image_processor=image_processor)
+model = model.to(device)
+clf = pipeline(model=model, task="image-classification", image_processor=image_processor, device=device)
 
 # Define class names
 class_names = ['artificial', 'real']
 
-def predict_image(img):
+def predict_image(img, confidence_threshold):
     # Convert the image to a PIL Image and resize it
     img = transforms.ToPILImage()(img)
     img = transforms.Resize((256, 256))(img)
+    img = transforms.ToTensor()(img).unsqueeze(0).to(device)  # Add batch dimension and move to GPU
     
     # Get the prediction
     prediction = clf(img)
@@ -26,9 +32,22 @@ def predict_image(img):
         if class_name not in result:
             result[class_name] = 0.0
     
-    return result
+    # Check if either class meets the confidence threshold
+    if result['artificial'] >= confidence_threshold:
+        return f"Label: artificial, Confidence: {result['artificial']:.4f}"
+    elif result['real'] >= confidence_threshold:
+        return f"Label: real, Confidence: {result['real']:.4f}"
+    else:
+        return "Uncertain Classification"
 
 # Define the Gradio interface
 image = gr.Image(label="Image to Analyze", sources=['upload'])
+confidence_slider = gr.Slider(0.0, 1.0, value=0.5, step=0.01, label="Confidence Threshold")
 label = gr.Label(num_top_classes=2)
-gr.Interface(fn=predict_image, inputs=image, outputs=label, title="AI Generated Classification").launch()
+
+gr.Interface(
+    fn=predict_image,
+    inputs=[image, confidence_slider],
+    outputs=label,
+    title="AI Generated Classification"
+).launch()