import gradio as gr
import numpy as np
import tensorflow as tf
import cv2
import os

# Suppress TensorFlow logging
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

# Constants
IMAGE_H = 512
IMAGE_W = 512
NUM_CLASSES = 11

# RGB color codes for each class
RGB_CODES = [
    [0, 0, 0], [0, 153, 255], [102, 255, 153], [0, 204, 153],
    [255, 255, 102], [255, 255, 204], [255, 153, 0], [255, 102, 255],
    [102, 0, 51], [255, 204, 255], [255, 0, 102]
]

# Load the trained model
model = tf.keras.models.load_model("ten_epoch_model.h5")

# Function to convert grayscale mask to RGB mask
def grayscale_to_rgb(mask, rgb_codes):
    h, w = mask.shape[0], mask.shape[1]
    mask = mask.astype(np.int32)
    output = [rgb_codes[pixel] for pixel in mask.flatten()]
    return np.reshape(output, (h, w, 3)).astype(np.uint8)

# Gradio inference function
def segment_face(image):
    # Resize and normalize input image
    image_resized = cv2.resize(image, (IMAGE_W, IMAGE_H))
    image_input = image_resized / 255.0
    image_input = np.expand_dims(image_input, axis=0).astype(np.float32)

    # Predict the mask
    pred = model.predict(image_input, verbose=0)[0]
    pred_mask = np.argmax(pred, axis=-1).astype(np.uint8)

    # Convert predicted mask to RGB
    rgb_mask = grayscale_to_rgb(pred_mask, RGB_CODES)

    return rgb_mask

# Launch Gradio app
iface = gr.Interface(
    fn=segment_face,
    inputs=gr.Image(type="numpy", label="Upload Face Image"),
    outputs=gr.Image(type="numpy", label="Segmentation Mask"),
    title="Face Segmentation",
    description="Upload a face image to get a segmentation mask with different facial components marked in color."
)

if __name__ == "__main__":
    iface.launch()