# AUTOGENERATED! DO NOT EDIT! File to edit: ../drive/MyDrive/Codici/Python/Apps/Gradio_App/Image Classification/simlpe_image_models.ipynb.

# %% auto 0
__all__ = ['cat_dog_model', 'bear_model', 'input_model', 'input_image', 'output_label', 'bear_examples', 'cat_examples',
           'examples', 'intf', 'get_model', 'classify_image']

# %% ../drive/MyDrive/Codici/Python/Apps/Gradio_App/Image Classification/simlpe_image_models.ipynb 2
# !pip install -Uqq fastai
# !pip install -Uqq fastbook
# !pip install gradio

# %% ../drive/MyDrive/Codici/Python/Apps/Gradio_App/Image Classification/simlpe_image_models.ipynb 3
# import fastbook
# fastbook.setup_book()
# from fastbook import *
# from fastai.vision.widgets import *
from fastai.vision.all import *
import gradio as gr

# %% ../drive/MyDrive/Codici/Python/Apps/Gradio_App/Image Classification/simlpe_image_models.ipynb 4
cat_dog_model = load_learner('cat_dog.pkl')
bear_model = load_learner("bear_classifier.pkl")

# %% ../drive/MyDrive/Codici/Python/Apps/Gradio_App/Image Classification/simlpe_image_models.ipynb 7
# Define a function to load the appropriate model
def get_model(model_name):
  if model_name == 'Cat vs Dog Model':
    return cat_dog_model
  elif model_name == 'Bear Model':
    return bear_model
  else:
    raise ValueError("Model not found")

# Classification Function
def classify_image(model_name, img):
  model = get_model(model_name)
  categories = model.dls.vocab
  pred, idx, probs = model.predict(img)
  return dict(zip(categories, map(float, probs)))

# %% ../drive/MyDrive/Codici/Python/Apps/Gradio_App/Image Classification/simlpe_image_models.ipynb 10
#INPUTS
input_model = gr.Dropdown(['Cat vs Dog Model', 'Bear Model'])
input_image = gr.components.Image(width=192, height=192)

#OUTPUT
output_label = gr.components.Label()

#EXAMPLES
bear_examples = ['black.jpg', 'grizzly.jpg', 'teddy.jpg']
cat_examples =  ['cat.jpg', 'dog.jpg']

examples = [
    ['Bear Model',bear_examples[0]], # added model name
    ['Bear Model',bear_examples[1]], # added model name
    ['Bear Model',bear_examples[2]], # added model name
    ['Cat vs Dog Model', cat_examples[0]], # added model name
    ['Cat vs Dog Model', cat_examples[1]], # added model name
]

#INTERFACE
intf = gr.Interface(fn=classify_image,
                    inputs=[
                        input_model,
                        input_image],
                    outputs=output_label,
                    examples=examples,
                    theme=gr.themes.Ocean())

#LAUNCH APP
intf.launch(inline=False,
            # share=True
            )