from utils import create_model, get_optimal_font_scale
from mtcnn import MTCNN
import cv2
import gradio as gr
from PIL import Image
import gradio

model = create_model()
checkpoint_path = "best_model_weights/checkpoint.ckpt"
model.load_weights(checkpoint_path)

detector = MTCNN()

def face_mask_detection(image):
    img = image.copy()
    (h,w) = img.shape[:2]
    objects = detector.detect_faces(img) # https://github.com/ipazc/mtcnn/blob/master/example.ipynb
    
    for obj in objects:
        x, y, width, height = obj['box']
        startX, startY, endX, endY = x, y, x+width, y+height

        #ensure the bounding boxes fall within the dimensions of the frame
        (startX,startY)=(max(0,startX),max(0,startY))
        (endX,endY)=(min(w-1,endX), min(h-1,endY))


        #extract the face ROI, convert it from BGR to RGB channel, resize it to 224,224 and preprocess it
        face=img[startY:endY, startX:endX]
        face=cv2.resize(face,(224,224))



        probs = model.predict(face.reshape(1,224,224,3))[0][0] # example: result of model.predict: [[0.823] then < 0.5 -> mask, > 0.5 -> without mask
        is_masked = probs < 0.5

        #determine the class label and color we will use to draw the bounding box and text
        label='Mask' if is_masked else 'No Mask'
        color=(0,255,0) if label=='Mask' else (0,0,255)

        #include the probability in the label
        confidence = probs if probs > 0.5 else (1-probs)
        label="{}: {:.2f}%".format(label,confidence*100)

        #display the label and bounding boxes
        fontScale = get_optimal_font_scale(label, endX-startX)
        cv2.putText(img,label,(startX+5,startY+15),cv2.FONT_HERSHEY_SIMPLEX,fontScale,color,1)
        cv2.rectangle(img,(startX,startY),(endX,endY),color,1)

    return Image.fromarray(img)


title = "Face-Masked Detection"
description = "This is an easy-to-use demo for detecting faces inside an image and then decide if they are wearing mask or not. MTCNN will be used for face-detection and EfficientNetB0 will be used for masked-classification."

demo = gr.Interface(fn=face_mask_detection,
             inputs=gr.Image(shape=(224, 224)),
             outputs=gr.outputs.Image(type = "pil", label = "Output Image"),
             title=title,
             description=description)

if __name__ == "__main__":
    demo.queue(max_size=10).launch()