import gradio as gr
import pandas as pd
from PIL import Image
from transformers import pipeline

pipe = pipeline("image-classification", model="raffaelsiregar/dog-breeds-classification")

def dog_classifier(dog_image):
    image = Image.fromarray(dog_image)
    output = pipe(image)
    
    # creating pandas dataframe
    df = pd.DataFrame(output)
    
    # adjust score column
    df['score'] = df['score'] * 100
    df['score'] = df['score'].apply(lambda x: round(x, 4))
    
    # rename the columns to make it more user-friendly
    df.columns = ['Breed', 'Confidence (%)']
    
    return df

title = "Dog Breed Classification"
description = "Upload an image (jpg is recommended) of a dog to predict its breed. The model will provide the top predictions with the confidence levels."
article = """
### How It Works
- The model classifies the breed of the dog in the image.
- It returns a table of the top predictions along with their confidence levels.
- This tool is built using a pre-trained image classification model from Hugging Face.
"""
themes = gr.themes.Citrus()

input_image = gr.Image(type="numpy", label="Upload a dog image")
output_table = gr.DataFrame(headers=["Breed", "Confidence (%)"], type="pandas")

# gradio interface
interface = gr.Interface(fn=dog_classifier,
                         inputs=input_image,
                         outputs=output_table,
                         title=title,
                         description=description,
                         article=article,
                         theme=themes)
interface.launch()