Update app.py

app.py

@@ -1,8 +1,20 @@
 import gradio as gr
-from fastai.vision.all import load_learner, PILImage
-
-
-learn = load_learner('pkg_recognizer_v1.pkl')
+from fastai.vision.all import load_learner, Learner
+import os
+
+# Try to load the model using the safer Learner.load method first
+# If that fails, fall back to load_learner with explicit file handling
+model_path = 'pkg_recognizer_v1'
+
+try:
+    # First try: Use the safer Learner.load method (expects .pth file)
+    learn = Learner.load(model_path)
+except (FileNotFoundError, RuntimeError):
+    try:
+        # Second try: Use load_learner for pickle files (.pkl)
+        learn = load_learner(model_path + '.pkl')
+    except Exception as e:
+        raise RuntimeError(f"Failed to load model. Ensure 'pkg_recognizer_v1.pth' or 'pkg_recognizer_v1.pkl' exists. Error: {e}")
 
 categories = learn.dls.vocab
 
@@ -11,7 +23,7 @@ def recognize_image(img):
     
     return dict(zip(categories, map(float, probs)))
 
-image = gr.Image(type="pil")
+image = gr.Image()
 label = gr.Label()
 
 examples = [