import cv2
import numpy as np
import gradio as gr
from tensorflow.keras.models import load_model

# Load the trained Keras model
model = load_model("face_mask_cnn_model.keras")

# Prediction function — same logic as Python script!
def predict_mask(image):
    # Convert Gradio's RGB image to BGR like OpenCV
    img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    
    # Resize (BGR format as used in training)
    img_resized = cv2.resize(img_bgr, (64, 64))

    # Normalize and reshape
    img_normalized = img_resized.astype('float32') / 255.0
    img_input = np.expand_dims(img_normalized, axis=0)  # (1, 64, 64, 3)

    # Predict
    prob_with_mask = float(model.predict(img_input)[0][0])
    prob_without_mask = 1.0 - prob_with_mask

    # Debug print (optional)
    print(f"[DEBUG] With Mask: {prob_with_mask:.3f}, Without Mask: {prob_without_mask:.3f}")

    # Return probability dict
    return {
        "With Mask": round(prob_with_mask, 3),
        "Without Mask": round(prob_without_mask, 3)
    }

# Gradio Interface
interface = gr.Interface(
    fn=predict_mask,
    inputs=gr.Image(label="Upload Image or Use Webcam", type="numpy", sources=["upload", "webcam"]),
    outputs=gr.Label(num_top_classes=2),
    live=True,
    title="Real-Time Face Mask Detection",
    description="Upload an image or use your webcam to detect mask-wearing probability!"
)

interface.launch()