App.py revision 2

app.py

@@ -6,7 +6,7 @@ import torch
 model = ViTForImageClassification.from_pretrained('sreeramajay/pollution')
 transforms = ViTFeatureExtractor.from_pretrained('sreeramajay/pollution')
 
-def predict(image): 
+def polln_classify(image): 
     labels = {0:"Air Pollution", 1: "Land Pollution" , 2: "Water Pollution"}
     inputs = transforms(image, return_tensors='pt')
     output = model(**inputs)
@@ -15,9 +15,9 @@ def predict(image):
     return {labels[i.item()]: v.item() for i, v in zip(indices.numpy()[0], values.detach().numpy()[0])}
 
 
-gr.Interface(
-    predict,
+gr.Interface(polln_classify,
     inputs = gr.inputs.Image(type="pil", label="Chosen Image"),
     outputs = 'label',
+    examples = ["air_pollution.jpg","land_pollution.jpg","water_pollution.jpg"],
     theme="seafoam",
 ).launch(debug=True)