102 lines
2.8 KiB
Python
102 lines
2.8 KiB
Python
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
|
|
from tensorflow.keras.preprocessing.image import img_to_array
|
|
from tensorflow.keras.models import load_model
|
|
import numpy as np
|
|
import cv2
|
|
|
|
labels = {0: 'mask_incorrectly_worn', 1: 'with_mask', 2: 'without_mask'}
|
|
|
|
|
|
def most_common(lst):
|
|
return max(set(lst), key=lst.count)
|
|
|
|
|
|
def detect_mask(img, face_detector, face_mask_detector):
|
|
(h, w) = img.shape[:2]
|
|
blob = cv2.dnn.blobFromImage(img, 1.0, (128, 128), (104.0, 177.0, 123.0))
|
|
|
|
face_detector.setInput(blob)
|
|
detections = face_detector.forward()
|
|
|
|
faces = []
|
|
locs = []
|
|
preds = []
|
|
|
|
for i in range(0, detections.shape[2]):
|
|
confidence = detections[0, 0, i, 2]
|
|
|
|
if confidence > 0.5:
|
|
box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
|
|
(startX, startY, endX, endY) = box.astype("int")
|
|
|
|
(startX, startY) = (max(0, startX), max(0, startY))
|
|
(endX, endY) = (min(w - 1, endX), min(h - 1, endY))
|
|
|
|
face = img[startY:endY, startX:endX]
|
|
face = cv2.cvtColor(face, cv2.COLOR_BGR2RGB)
|
|
face = cv2.resize(face, (128, 128))
|
|
face = img_to_array(face)
|
|
face = preprocess_input(face)
|
|
|
|
faces.append(face)
|
|
locs.append((startX, startY, endX, endY))
|
|
|
|
if len(faces) > 0:
|
|
faces = np.array(faces, dtype="float32")
|
|
preds = face_mask_detector.predict(faces, batch_size=32)
|
|
|
|
return locs, preds
|
|
|
|
|
|
prototxtPath = r"face_detector\deploy.prototxt"
|
|
weightsPath = r"face_detector\res10_300x300_ssd_iter_140000.caffemodel"
|
|
faceNet = cv2.dnn.readNet(prototxtPath, weightsPath)
|
|
|
|
mask_detection_model = load_model("face_mask_detection2.h5")
|
|
|
|
states = []
|
|
current_label = ''
|
|
current_color = (0, 0, 0)
|
|
|
|
cam = cv2.VideoCapture(0)
|
|
|
|
while True:
|
|
ret, frame = cam.read()
|
|
if not ret:
|
|
print("failed to grab frame")
|
|
break
|
|
|
|
(locs, preds) = detect_mask(frame, faceNet, mask_detection_model)
|
|
|
|
for (box, pred) in zip(locs, preds):
|
|
(startX, startY, endX, endY) = box
|
|
|
|
label_index = np.argmax(pred)
|
|
|
|
states.append(label_index)
|
|
|
|
if len(states) == 10:
|
|
index = most_common(states)
|
|
current_label = labels[index]
|
|
if index == 2:
|
|
current_color = (0, 0, 255)
|
|
elif index == 1:
|
|
current_color = (0, 255, 0)
|
|
else:
|
|
current_color = (0, 127, 255)
|
|
states.clear()
|
|
|
|
if current_label != '':
|
|
cv2.putText(frame, current_label, (startX, startY - 10),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.45, current_color, 2)
|
|
cv2.rectangle(frame, (startX, startY), (endX, endY), current_color, 2)
|
|
|
|
cv2.imshow("Face Mask Detection", frame)
|
|
key = cv2.waitKey(1) & 0xFF
|
|
|
|
if key == ord("q"):
|
|
break
|
|
|
|
cv2.destroyAllWindows()
|
|
cam.release()
|