Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from transformers import pipeline, AutoImageProcessor, Swinv2ForImageClassification
 from torchvision import transforms
 import torch
+from PIL import Image
 
 # Ensure using GPU if available
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -17,12 +18,11 @@ class_names = ['artificial', 'real']
 
 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
+    img_pil = Image.fromarray(img).convert('RGB')  # Convert NumPy array to PIL Image
+    img_pil = transforms.Resize((256, 256))(img_pil)
     
     # Get the prediction
-    prediction = clf(img)
+    prediction = clf(img_pil)
     
     # Process the prediction to match the class names
     result = {pred['label']: pred['score'] for pred in prediction}
@@ -41,7 +41,7 @@ def predict_image(img, confidence_threshold):
         return "Uncertain Classification"
 
 # Define the Gradio interface
-image = gr.Image(label="Image to Analyze", sources=['upload'])
+image = gr.Image(label="Image to Analyze", sources=['upload'], type='pil')  # Ensure the image type is PIL
 confidence_slider = gr.Slider(0.0, 1.0, value=0.5, step=0.01, label="Confidence Threshold")
 label = gr.Label(num_top_classes=2)