Spaces:
Sleeping
Sleeping
| 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() | |