import gradio as gr
import tensorflow as tf
from tensorflow.compat.v2.experimental import dtensor

import numpy as np
from PIL import Image

# Load pre-trained MobileNetV2 model
model = tf.keras.applications.MobileNetV2(weights='imagenet')

def predict_difficulty_score(image):
    # Load image and preprocess it for the model
    img = Image.fromarray(image.astype('uint8'), 'RGB')
    img = img.resize((224, 224))
    img_array = tf.keras.preprocessing.image.img_to_array(img)
    img_array = tf.keras.applications.mobilenet_v2.preprocess_input(img_array[np.newaxis,...])

    # Use the model to predict the image class probabilities
    preds = model.predict(img_array)

    # Get the index of the top predicted class
    class_idx = np.argmax(preds[0])

    # Get the difficulty score based on the class index
    difficulty_score = round((class_idx / 999) * 99000) + 1000

    # Return the difficulty score
    return difficulty_score

# Create a Gradio interface
inputs = gr.inputs.Image(shape=(224, 224))
outputs = gr.outputs.Textbox(label="Difficulty Score")
interface = gr.Interface(fn=predict_difficulty_score, inputs=inputs, outputs=outputs, 
                         title="AI Art Difficulty Score", description="Upload an AI art image and get its difficulty score.")

# Launch the interface
interface.launch()