import os
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.preprocessing import image
import gradio as gr

# Paintings' authors considered for this version 
artists = ['botero', 
          'davinci', 
          'elgreco', 
          'melzi', 
          'michelangelo',
          'modigliani', 
          'picasso', 
          'rembrandt', 
          'rubens', 
          'vermeer']

model = keras.models.load_model('models/model_2023-08-29T1856_ep40_bz32_img224_nc10.h5')

description = 'Welcome!!! This app was built based on Gradio. The aim of this App is to predict the author of a painting. In this first version of the App, we only considered 10 authors [botero, davinci, elgreco, melzi, michelangelo,modigliani, picasso, rembrandt, rubens, vermeer]'

def predicting_author(input):

    try:
        if input is None:
            return 'Please upload an image'
        
        x = image.img_to_array(input)
        x = np.expand_dims(x, axis=0)
        x = x.astype('float32') / 255.0
        
        prediction = model.predict(x)
        class_probabilities = prediction[0] 
        results = {artists[i]: float(class_probabilities[i]) for i in range(len(artists))}
        
        return results
    
    except Exception as e:
        print("An error occurred:", e)
        # Print traceback to see more details
        import traceback
        traceback.print_exc()
        return "An error occurred"

demo = gr.Interface(
    title='Predicting paintings authors',
    description=description,
    fn=predicting_author, 
    inputs=gr.Image(shape=(224, 224)), 
    outputs=gr.Label(num_top_classes=3),
    examples=['./test_images/image1.jpg', './test_images/image2.jpg', './test_images/image3.jpg']
    )
             
demo.launch()