import gradio as gr
from tensorflow.keras.models import load_model
from tensorflow.keras.layers import DepthwiseConv2D
from PIL import Image, ImageOps
import numpy as np

# Disable scientific notation for clarity
np.set_printoptions(suppress=True)

# Custom object for DepthwiseConv2D
custom_objects = {'DepthwiseConv2D': DepthwiseConv2D}

# Load the model with custom objects
model = load_model("model/pleasuredomes_image_model.h5", custom_objects=custom_objects, compile=False)

# Load the labels
class_names = open("model/labels.txt", "r").readlines()

def predict_image(image):
    """
    Function to process the image and make a prediction using the loaded model.
    """
    # Resize the image to be at least 224x224 and then crop from the center
    size = (224, 224)
    image = ImageOps.fit(image, size, Image.Resampling.LANCZOS)

    # Turn the image into a numpy array
    image_array = np.asarray(image)

    # Normalize the image
    normalized_image_array = (image_array.astype(np.float32) / 127.5) - 1

    # Create the array of the right shape to feed into the keras model
    data = np.ndarray(shape=(1, 224, 224, 3), dtype=np.float32)
    data[0] = normalized_image_array

    # Predict the model
    prediction = model.predict(data)
    index = np.argmax(prediction)
    class_name = class_names[index].strip()
    confidence_score = prediction[0][index]

    return f"Class: {class_name}, Confidence Score: {confidence_score:.2f}"

# Create a Gradio interface
interface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil"),  # Updated to gr.Image
    outputs="text",
    title="Image Classification",
    description="Upload an image to classify it using the pre-trained model.",
    flagging_options=None 
)

# Launch the interface
if __name__ == "__main__":
    interface.launch(share=False)