DD0101's picture
Update app.py
47f1d49 verified
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()