Update app.py

app.py

@@ -3,14 +3,16 @@ import requests
 from PIL import Image
 from io import BytesIO
 from transformers import pipeline
+import torch
 
+# Cache the model loading
 model_id = "Honey-Bee-Society/honeybee_bumblebee_vespidae_resnet50"
-classifier = pipeline("image-classification", model=model_id)
+classifier = pipeline("image-classification", model=model_id, device=0 if torch.cuda.is_available() else -1)
 
 def classify_image_from_url(image_url: str):
     """
     Downloads an image from a public URL and runs it through
-    the ResNet-50 fine tuned image-classification pipeline, returning the top predictions.
+    the ResNet-50 fine-tuned image-classification pipeline, returning the top predictions.
     """
     try:
         response = requests.get(image_url)
@@ -22,12 +24,14 @@ def classify_image_from_url(image_url: str):
 
         # Format scores to remove scientific notation
         for r in results:
-            r["score"] = float(f"{r['score']:.2f}")
+            r["score"] = float(f"{r['score']:.8f}")
 
         return results
     
+    except requests.exceptions.RequestException as e:
+        return {"error": f"Failed to download image: {str(e)}"}
     except Exception as e:
-        return {"error": str(e)}
+        return {"error": f"An error occurred during classification: {str(e)}"}
 
 demo = gr.Interface(
     fn=classify_image_from_url,
@@ -38,4 +42,4 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()