Update app.py

app.py

@@ -1,27 +1,34 @@
 import gradio as gr
-# from transformers import AutoBackbone, AutoModelForImageClassification, AutoImageProcessor, Swinv2ForImageClassification
 from transformers import pipeline, AutoImageProcessor, Swinv2ForImageClassification
 from torchvision import transforms
 
-# model = AutoModelForImageClassification.from_pretrained("haywoodsloan/ai-image-detector-deploy")
-# image_processor = AutoImageProcessor.from_pretrained("haywoodsloan/ai-image-detector-deploy")
+# Load the model and processor
 image_processor = AutoImageProcessor.from_pretrained("haywoodsloan/ai-image-detector-deploy")
-# image_processor = Swinv2ForImageClassification.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)
 
+# Define class names
 class_names = ['artificial', 'real']
 
 def predict_image(img):
-  img = transforms.ToPILImage()(img)
-  img = transforms.Resize((256,256))(img)
-  prediction=clf.predict(img)
-  print(prediction)
-
-  return {class_names[i]: float(prediction[i]["score"]) for i in range(2)}
+    # Convert the image to a PIL Image and resize it
+    img = transforms.ToPILImage()(img)
+    img = transforms.Resize((256, 256))(img)
+    
+    # Get the prediction
+    prediction = clf(img)
+    
+    # Process the prediction to match the class names
+    result = {pred['label']: pred['score'] for pred in prediction}
+    
+    # Ensure the result dictionary contains both class names
+    for class_name in class_names:
+        if class_name not in result:
+            result[class_name] = 0.0
+    
+    return result
 
+# Define the Gradio interface
 image = gr.Image(label="Image to Analyze", sources=['upload'])
 label = gr.Label(num_top_classes=2)
-
 gr.Interface(fn=predict_image, inputs=image, outputs=label, title="AI Generated Classification").launch()