import gradio as gr
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

def recognize_image(img):
    pred, idx, probs = learn.predict(img)
    
    return dict(zip(categories, map(float, probs)))

image = gr.Image()
label = gr.Label()

examples = [
    '1.jpeg',
    '2.jpg',
    '3.jpg',
    '4.jpeg'
]

iface = gr.Interface(
    fn=recognize_image,
    inputs=image,
    outputs=label,
    examples=examples,
    title="Daily Object Packaging Image Classification",
    description="Classifies packaging into 12 categories: Bottled Water, Milk Cartons, Snack Chips, etc."
)

if __name__ == "__main__":
    iface.launch(inline=False)