'''PyTorch Food Classifier - FastAI 2022- Mostly Good For Pastries
and trained on ResNet 34'''

import streamlit as st
import os
from PIL import Image
import time
from fastai.vision.all import *
from fastai.learner import load_learner
    
#Load the Learner (Exported from ipnyb file with learn.export() )
learn = load_learner('export.pkl')

categories = ('penguin', 'puffin', 'pufferfish')

#Classify image
def classify_image(cl_img):
    img = Image.open(cl_img)
    st.image(img)
    pred, pred_idx, prob = learn.predict(img)
    confidence = prob[pred_idx].item() * 100
    return pred, prob
    
st.set_page_config(page_title="Penguin vs Puffin Classifier - FastAI 2025", page_icon=":robot:")
st.header("Penguin vs Puffin Classifier")

file_up = st.file_uploader("Upload Your Image Below", type=["jpg","png"])

if st.button('Run Model'):
    st.write("Button Pressed")
    pred_label, confidence = classify_image(file_up)
    st.write(f"The model predicts {pred_label} with {confidence} confidence.")

st.write('This classifier is trained on Resnet-34 and specializes in differentiating penguins from puffins).')

# from fastai.vision.all import *
# from fastai.learner import load_learner
# import gradio as gr

# learn = load_learner('model.pkl')

# categories = ('Penguin', 'Puffin')

# def classifyImage(img):
#     pred, idx, prob = learn.predict(img)
#     return dict(zip(categories, map(float, prob)))

# image = gr.Image(shape=(192, 192))
# label = gr.Label()
# examples = ['penguin.jpg', 'puffin.png', 'razorbill.jpg']

# intf = gr.Interface(fn=classifyImage, inputs=image, outputs=label, examples=examples)
# intf.launch(inline=False)