Update app.py

app.py

@@ -1,21 +1,38 @@
+# Install necessary packages
+!pip install torch transformers gradio
+
+# Import libraries
 import torch
 from transformers import AutoModelForImageClassification, AutoFeatureExtractor
 import gradio as gr
 
-model_id = f'Caniya/vit-base-patch16-224-finetuned-flower'
+# Define model ID and labels
+model_id = 'Caniya/vit-base-patch16-224-finetuned-flower'
 labels = ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']
 
+# Define classify_image function
 def classify_image(image):
-  model = AutoModelForImageClassification.from_pretrained(model_id)
-  feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)
-  inp = feature_extractor(image, return_tensors='pt')
-  outp = model(**inp)
-  pred = torch.nn.functional.softmax(outp.logits, dim=-1)
-  preds = pred[0].cpu().detach().numpy()
-  confidence = {label: float(preds[i]) for i, label in enumerate(labels)}
-  return confidence
+    model = AutoModelForImageClassification.from_pretrained(model_id)
+    feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)
+    inp = feature_extractor(image, return_tensors='pt')
+    outp = model(**inp)
+    pred = torch.nn.functional.softmax(outp.logits, dim=-1)
+    preds = pred[0].cpu().detach().numpy()
+    confidence = {label: float(preds[i]) for i, label in enumerate(labels)}
+    return confidence
+
+# Create interface
+interface = gr.Interface(
+    fn=classify_image,
+    inputs='image',
+    outputs='label',
+    title='Flower Image Classifier',
+    description='Classify images of flowers into different categories.',
+    examples=[
+        ['flower-1.jpeg'],
+        ['flower-2.jpeg']
+    ]
+)
 
-interface = gr.Interface(fn=classify_image,
-                         inputs='image',
-                         examples=['Dandelion_2.jpg', 'dandelion_1.jpg'],
-                         outputs='label').launch()
+# Launch the interface
+interface.launch()