import tensorflow as tf
import pandas as pd
import gradio as gr

authors_df = pd.read_csv('authors.csv')
labels = sorted(list(authors_df.name))

model = tf.keras.models.load_model('efficientnetb0.h5')

description = 'This is a DEMO that attempts to recognize the inspirations used by the AI art generator. After uploading a picture of an image, the application displays the predicted artist along with the probability of predicting the top three authors.The DEMO uses EfficientNetB0 convolutional neural network as a base model whose classifier was modified and trained the 8,000+ paintings from [Kaggle](https://www.kaggle.com/datasets/ikarus777/best-artworks-of-all-time) dataset. Model trained by osydorchuk89. Given the dataset limitations, the model only recognizes paintings of [50 artists](https://huggingface.co/spaces/osydorchuk/painting_authors/blob/main/authors.csv).'

def predict_author(input):
    if input is None:
        return 'Please upload an image'
    input = input.reshape((-1, 224, 224, 3))
    prediction = model.predict(input).flatten()
    confidences = {labels[i]: float(prediction[i]) for i in range(50)}
    return confidences

demo = gr.Interface(
    title='the AI art generator sources of inspiration',
    description=description,
    fn=predict_author, 
    inputs=gr.Image(shape=(224, 224)), 
    outputs=gr.Label(num_top_classes=3),
    examples=['test_pics/eva_miro.jpg', 'test_pics/eva_bosch.jpg',  'test_pics/eva_miro_2.jpg', 'test_pics/eva_rtology.jpg']
    )
             
demo.launch()