Update app.py

app.py

@@ -1,20 +1,20 @@
 import gradio as gr
-from transformers import AutoModelForImageClassification, pipeline, AutoImageProcessor
+from transformers import AutoModelForImageClassification, pipeline, AutoImageProcessor, SwinForImageClassification
 from torchvision import transforms
 
-model = AutoModelForImageClassification.from_pretrained("Nicole-M/Dataset1-SwinV2")
-image_processor = AutoImageProcessor.from_pretrained("Nicole-M/Dataset1-SwinV2")
+model = AutoModelForImageClassification.from_pretrained("haywoodsloan/ai-image-detector-deploy")
+image_processor = AutoImageProcessor.from_pretrained("haywoodsloan/ai-image-detector-deploy")
 clf = pipeline(model=model, task="image-classification", image_processor=image_processor)
 
-class_names = ['Benign', 'Malignant']
+class_names = ['artificial', 'real']
 
 def predict_image(img):
   img = transforms.ToPILImage()(img)
-  img = transforms.Resize((224,224))(img)
+  img = transforms.Resize((256,256))(img)
   prediction=clf.predict(img)
   return {class_names[i]: float(prediction[i]["score"]) for i in range(2)}
 
-image = gr.Image(label="Select a mammogram image", sources=['upload'])
+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="Mammogram classification").launch()
+gr.Interface(fn=predict_image, inputs=image, outputs=label, title="AI Generated Classification").launch()