Create app.py

app.py

@@ -0,0 +1,38 @@
+import torch
+from transformers import ViTForImageClassification, ViTFeatureExtractor
+import gradio as gr
+from PIL import Image
+
+# Check if GPU is available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load pre-trained ViT model from Hugging Face
+model = ViTForImageClassification.from_pretrained('Dhahlan2000/banana_ripeness_level_detection', num_labels=20)
+model.to(device)
+model.eval()
+
+# Load ViT feature extractor
+feature_extractor = ViTFeatureExtractor.from_pretrained('Dhahlan2000/banana_ripeness_level_detection')
+
+# Class labels
+predicted_classes = ['Overripe', 'ripe', 'rotten', 'unripe']
+
+# Function for inference
+def classify_fruit(image):
+    inputs = feature_extractor(images=image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = model(**inputs)
+        logits = outputs.logits
+        predicted_class_idx = logits.argmax(-1).item()
+    return predicted_classes[predicted_class_idx]
+
+# Gradio UI
+demo = gr.Interface(
+    fn=classify_fruit,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Label(),
+    title="Fruit Ripeness Detection",
+    description="Upload an image of a fruit to determine whether it's fresh or rotten."
+)
+
+demo.launch()