Create app.py

app.py

@@ -0,0 +1,20 @@
+import gradio as gr
+from transformers import AutoModelForImageClassification, pipeline, AutoImageProcessor
+from torchvision import transforms
+
+model = AutoModelForImageClassification.from_pretrained("Nicole-M/Dataset1-SwinV2")
+image_processor = AutoImageProcessor.from_pretrained("Nicole-M/Dataset1-SwinV2")
+clf = pipeline(model=model, task="image-classification", image_processor=image_processor)
+
+class_names = ['Benign', 'Malignant']
+
+def predict_image(img):
+  img = transforms.ToPILImage()(img)
+  img = transforms.Resize((224,224))(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'])
+label = gr.Label(num_top_classes=2)
+
+gr.Interface(fn=predict_image, inputs=image, outputs=label, title="Mammogram classification").launch()