Update app.py

app.py

@@ -1,4 +1,10 @@
 import gradio as gr
+
+from PIL import Image
+
+import spaces
+
+from typing import Dict
 import torch
 from transformers import ViTImageProcessor, AutoFeatureExtractor, AutoModelForImageClassification
 
@@ -9,20 +15,31 @@ model = AutoModelForImageClassification.from_pretrained("saved_model_files")
 
 labels = ['angular_leaf_spot', 'bean_rust', 'healthy'] 
 
+@spaces.GPU(duration=240)
+def classify(image: Image.Image) -> Dict[str, float]:
+    """
+    Classify an image of a bean plant leaf into one of several health categories.
 
-def classify(image):
-  features = image_processor(image, return_tensors='pt')
-  logits = model(features["pixel_values"])[-1]
-  probability = torch.nn.functional.softmax(logits, dim=-1)
-  probs = probability[0].detach().numpy()
-  confidences = {label: float(probs[i]) for i, label in enumerate(labels)} 
-  return confidences
+    Args:
+        image (Image.Image): The input image of the bean leaf to be classified.
 
-#### APP #####
+    Returns:
+        Dict[str, float]: A dictionary where the keys are the health labels 
+                          (e.g., 'angular_leaf_spot', 'bean_rust', 'healthy') and 
+                          the values are the confidence scores for each label.
+    """
+    features = image_processor(image, return_tensors='pt')
+    logits = model(features["pixel_values"])[-1]
+    probability = torch.nn.functional.softmax(logits, dim=-1)
+    probs = probability[0].detach().numpy()
+    confidences = {label: float(probs[i]) for i, label in enumerate(labels)} 
+    return confidences
+
+####### GRADIO APP #######
 title = """<h1 id="title">Bean plant health predictor through images of leaves using ViT image classifier</h1>"""
 
 description = """
-Use Case: A farming company that is having issues with diseases affecting their bean plants. The farmers have to constantly monitor the leaves of the plants so that they can immediately treat the leaves if they show any signs of disease. 
+Problem Statement: A farming company that is having issues with diseases affecting their bean plants. The farmers have to constantly monitor the leaves of the plants so that they can immediately treat the leaves if they show any signs of disease. 
 We are asked to build a machine learning-based app they can deploy on a drone to quickly identify diseased plants.
 
 
@@ -43,8 +60,9 @@ demo = gr.Blocks(css=css, theme=theme)
 with demo:
   gr.Markdown(title)
   gr.Markdown(description)
+  interface =  gr.Interface(fn=classify, 
+                            inputs="image",
+                            outputs="label",
+                           examples="images")
 
-
-  interface =  gr.Interface(fn=classify, inputs="image", outputs="label")
-
-demo.launch()
+demo.launch()