Create app.py

app.py

@@ -0,0 +1,32 @@
+import gradio as gr
+from transformers import AutoModelForImageClassification, AutoFeatureExtractor
+import requests
+from PIL import Image
+import torch
+
+# Load the pre-trained model and preprocessor (feature extractor)
+model_name = "jjuarez/Vit_waste_image_class"
+model = AutoModelForImageClassification.from_pretrained(model_name)
+feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
+
+def classify_image(image):
+    # Preprocess the image
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    # Make prediction
+    with torch.no_grad():
+        logits = model(**inputs).logits
+    # Retrieve the highest probability class label
+    predicted_class_idx = logits.argmax(-1).item()
+    # Convert the index to the model's class label
+    label = model.config.id2label[predicted_class_idx]
+    return label
+
+# Create Gradio interface
+iface = gr.Interface(fn=classify_image, 
+                     inputs=gr.inputs.Image(shape=(224, 224)), 
+                     outputs="label",
+                     title="Waste Classification with ViT",
+                     description="Upload an image of waste, and the model will classify it.")
+
+# Launch the app
+iface.launch()