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

# Load model and haarcascade once
model = load_model("masknet.h5")
face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')

mask_label = {0: 'MASK', 1: 'NO MASK'}
dist_label = {0: (0, 255, 0), 1: (255, 0, 0)}
MIN_DISTANCE = 75

def predict(image):
    img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    orig_img = img.copy()

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(gray, 1.1, 4)

    label = [0 for _ in range(len(faces))]

    if len(faces) >= 2:
        centers = [(x + w // 2, y + h // 2) for (x, y, w, h) in faces]
        for i in range(len(centers) - 1):
            for j in range(i + 1, len(centers)):
                dist = distance.euclidean(centers[i], centers[j])
                if dist < MIN_DISTANCE:
                    label[i] = 1
                    label[j] = 1

    for i in range(len(faces)):
        (x, y, w, h) = faces[i]
        face_crop = img[y:y + h, x:x + w]
        face_crop = cv2.resize(face_crop, (128, 128))
        face_crop = np.reshape(face_crop, [1, 128, 128, 3]) / 255.0

        result = model.predict(face_crop)
        class_id = np.argmax(result)

        cv2.rectangle(orig_img, (x, y), (x + w, y + h), dist_label[label[i]], 2)
        cv2.putText(orig_img, mask_label[class_id], (x, y - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, dist_label[label[i]], 2)

    output = cv2.cvtColor(orig_img, cv2.COLOR_BGR2RGB)
    return output

iface = gr.Interface(fn=predict, inputs=gr.Image(type="numpy"), outputs="image", title="Face Mask Detection")
iface.launch()