from fastai.vision.all import *
import string
import gradio as gr

def getClassName(fileName):
    fileName = fileName[:-4] # remove extension .jpg
    names = fileName.split('_')
    if(names[0] == "Albrecht"):
        return "Albrecht Dürer"
    
    artist = ''.join(name+" " for name in names if name[0] in string.ascii_letters)
    return artist[:-1] #Remove last space

def predict(img):
    img = PILImage.create(img)
    pred, pred_idx, probs = learn.predict(img)
    return {labels[i]: float(probs[i]) for i in range(len(labels))}


learn = load_learner('drawingstyle.pkl')
labels = learn.dls.vocab

title = "What's your drawing style?"
description = "An artistic style classifier trained on the <a href='https://www.kaggle.com/datasets/ikarus777/best-artworks-of-all-time'>best-artworks-of-all-time</a> dataset with fastai. Created as homework for lesson 2 of the course <a href='https://course.fast.ai/Lessons/lesson2.html'>Pratical Deep Learning for Coders</a>.<br>Do you feel like something is off? Check the <a href='https://www.kaggle.com/code/nerusskyhigh/artisticdrawingstyle' target='_blank'>Kaggle Notebook</a>! ;)"


article="<p style='text-align: center'><a href='https://github.com/NerusSkyhigh/practicaldeeplearning' target='_blank'>Check the repository!</a></br>You can find more test images in the original dataset. Only the first 100 images per artist were used to train the neural network. As the other images were discarded, you can use them to test the network.</p>"

examples = ['Vincent_van_Gogh_151.jpg', 'Leonardo_da_Vinci_121.jpg', 'Pablo_Picasso_164.jpg']


gr.Interface( fn=predict, inputs=gr.inputs.Image(shape=(512,512)), outputs=gr.outputs.Label(num_top_classes=3),
                title=title, description=description, article=article, examples=examples).launch(share=True)