import gradio as gr
import numpy as np
import cv2
from tensorflow.keras.preprocessing.image import img_to_array
from tensorflow.keras.models import load_model
import tensorflow as tf

# Load your our trained Model
model = load_model('Crack-segmentation.h5')


def preprocess_image(image, target_size=(128, 128)):
    # Convert the PIL image to a NumPy array
    img_array = np.array(image)
    # Resize the image using OpenCV
    img_resized = cv2.resize(img_array, target_size)
    # Normalize the image
    img_resized = img_resized.astype('float32') / 255.0
    # Expand dimensions to match the input shape for the model
    img_resized = np.expand_dims(img_resized, axis=0)
    return img_resized

def predict_mask(image):
    preprocessed_image = preprocess_image(image)
    # Make predictions
    predictions = model.predict(preprocessed_image)
    # Squeeze the prediction to remove the batch dimension
    predicted_mask = predictions.squeeze()
    # Normalize the mask to [0, 255] and convert to uint8
    predicted_mask = (predicted_mask * 255).astype(np.uint8)
    return predicted_mask

# Define the Gradio interface
iface = gr.Interface(
    fn=predict_mask,
    inputs=gr.Image(label='Upload Image of Wall'),
    outputs=gr.Image(type="numpy",label='Segmented Image 🚀'),
    title="Crack Segmentation",
    description="Upload an Image 📥"
)

# Launch the Gradio interface
iface.launch()