import gradio as gr
import tensorflow as tf
from tensorflow.keras.applications.inception_resnet_v2 import preprocess_input
from tensorflow.keras.preprocessing import image
import numpy as np

# กำหนดเลเยอร์ที่กำหนดเอง (CustomScaleLayer)
class CustomScaleLayer(tf.keras.layers.Layer):
    def __init__(self, scale=1.0, **kwargs):
        super(CustomScaleLayer, self).__init__(**kwargs)
        self.scale = scale

    def call(self, inputs, *args, **kwargs):
        if isinstance(inputs, list):
            return [input_tensor * self.scale for input_tensor in inputs]
        else:
            return inputs * self.scale

# ใช้ custom_object_scope เพื่อทำให้เลเยอร์ที่กำหนดเองสามารถใช้งานได้
with tf.keras.utils.custom_object_scope({'CustomScaleLayer': CustomScaleLayer}):
    model = tf.keras.models.load_model("jo33_model_v222.h5")

# Function for prediction
def predict(img):
    img = img.resize((224, 224))  # Resize image to the target size
    img_array = image.img_to_array(img)  # Convert image to array
    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
    img_array = preprocess_input(img_array)  # Preprocess image according to model requirements

    predictions = model.predict(img_array)
    class_idx = np.argmax(predictions, axis=1)[0]
    class_label = list(train_generator.class_indices.keys())[class_idx]
    confidence = predictions[0][class_idx]

    return {class_label: confidence}

# Create Gradio Interface
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload an Image"),
    outputs=gr.Label(num_top_classes=2, label="Predicted Class"),
    title="Image Classification with InceptionResNetV2",
    description="Upload an image to classify it into one of the classes."
)

# Launch the interface
interface.launch()