Update app.py

app.py

@@ -4,12 +4,16 @@ from PIL import Image
 from io import BytesIO
 from transformers import pipeline
 
-#new
-# 1. Load a pretrained ResNet-50 from the Hugging Face Hub
+# Our label mapping:
+label_map = {
+    "LABEL_0": "honeybee",
+    "LABEL_1": "bumblebee",
+    "LABEL_2": "vespidae"
+}
+
 model_id = "Honey-Bee-Society/honeybee_bumblebee_vespidae_resnet50"
 classifier = pipeline("image-classification", model=model_id)
 
-# 2. Define an inference function
 def classify_image_from_url(image_url: str):
     """
     Downloads an image from a public URL and runs it through
@@ -24,15 +28,16 @@ def classify_image_from_url(image_url: str):
         # Run inference
         results = classifier(image)
         
-        # You can return raw results or format them as desired
+        # Post-process to replace "LABEL_0" etc. with "honeybee", "bumblebee", "vespidae"
+        for r in results:
+            if r["label"] in label_map:
+                r["label"] = label_map[r["label"]]
+        
         return results
     
     except Exception as e:
         return {"error": str(e)}
 
-# 3. Create a Gradio interface
-#    - We accept a single Textbox input (the public image URL)
-#    - We return the classification results in JSON format
 demo = gr.Interface(
     fn=classify_image_from_url,
     inputs=gr.Textbox(lines=1, label="Image URL"),
@@ -41,7 +46,5 @@ demo = gr.Interface(
     description="Enter a public image URL to get top predictions."
 )
 
-# 4. Launch the app
 if __name__ == "__main__":
     demo.launch()
-