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

# Load the SavedModel
model = tf.keras.models.load_model("mask_mobilenet_savedmodel")

def predict_mask(image):
    try:
        # Convert RGBA to RGB if needed
        if image.shape[2] == 4:
            image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)

        # Resize and normalize
        image = cv2.resize(image, (224, 224))
        image = image / 255.0
        image = np.expand_dims(image, axis=0)

        # Use the 'call' method instead of calling the object
        preds = model.call(tf.convert_to_tensor(image), training=False).numpy()

        result = "Mask" if preds[0][0] > 0.5 else "No Mask"
        return result

    except Exception as e:
        print("Error:", e)
        return f"Error: {e}"

# Gradio interface
iface = gr.Interface(
    fn=predict_mask,
    inputs=gr.Image(type="numpy"),
    outputs="text",
    title="Mask Detection",
    description="Upload an image to check if a person is wearing a mask."
)

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