import gradio as gr
from transformers import pipeline, AutoModel

model = AutoModel.from_pretrained("ericxlima/DogsClassifierModel")

 
dogs = { 
    'Zwergspitz Dog': [],
    'Bouledogue Français Dog': [],
    'Shih Tzu Dog': [],
    'Rottweiler Dog': [],
    'Pug Dog': [],
    'Golden Retriever Dog': [],
    'Deutscher Schäferhund Dog': [],
    'Yorkshire Terrier Dog': [],
    'Border Collie Dog': [],
    'Dachshund Dog': [],
    'Poodle Dog': [],
    'Labrador Retriever Dog': [],
    'Pinscher Dog': [],
    'Golden Retriever': [],
    }

pipeline = pipeline(model=model)


def predict(image):
  predictions = pipeline(image)
  return {p["label"]: p["score"] for p in predictions}



def list_breeds():
    global dogs
    html = "<div class='row'>"
    html += "<div class='column'>"
    html += "<h2>List of breed dogs trained:</h2>"
    html += "<ol>" + "".join([f"<li>{breed}</li>" for breed in list(dogs.keys())]) + "</ol>"
    html += "</div>"
    html += "<div class='column'>"
    html += "<h2>Author:</h2>"
    html += "<a href='https://github.com/ericxlima'><img src='https://avatars.githubusercontent.com/u/58092119?v=4' alt='profile image' style='width:40%' /></a>"
    html += "<h2><a href='https://github.com/ericxlima'>Eric de Lima</a></h2>"
    html += "</div>"
    html += "</div>"
    return html 


image = gr.Image(shape=(224, 224))
label = gr.Label(num_top_classes=3)
# breeds_list = list_breeds()

demo = gr.Interface(
    fn=predict,
    inputs=image,
    outputs=label,
    title="🐶 Dog Breed Classifier",
    interpretation="default",
    description="Upload an image of a dog and the model will predict its breed.",
    # article=breeds_list,
    css=".row { display: flex; } .column { flex: 50%; }",
)
 
demo.launch(share=True, debug=True)