0037-检测人脸是否佩戴口罩-基于keras的mobilenetV2的三分类模型

个人微信公众号

train.py

# -*- coding: utf-8 -*-
# @Time    : 2021/1/27 10:57
# @Author  : Johnson
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.layers import AveragePooling2D
from tensorflow.keras.layers import Dropout
from tensorflow.keras.layers import Flatten
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import Input
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
from tensorflow.keras.callbacks import EarlyStopping
from tensorflow.keras.callbacks import ModelCheckpoint
from sklearn.metrics import classification_report
import matplotlib.pyplot as plt
import numpy as np
import argparse
import os

# Instantiate an argument parser and parse the arguments
parser = argparse.ArgumentParser()
parser.add_argument("--dataset", "-d",
                    default='MFN', choices=['MFN', 'RMFD'],
                    help="dataset to train the model on")
args = parser.parse_args()

#validate  argument
if args.dataset != "MFN" and args.dataset != "RMFD":
    raise ValueError("Please provide a valid dataset choice: `MFN` or `RMFD`.")

# Change the working directory from src to root if needed
current_full_dir = os.getcwd()
print("Current working directory: " + current_full_dir)
if current_full_dir.split("/")[-1] == "src":
    root = current_full_dir[:-4]
    os.chdir(root)
    print("Changed working directory to: " + root)

# Initialize number of classes and labels
NUM_CLASS, class_names = None, None
if args.dataset == "MFN":
    NUM_CLASS = 3
    class_names = ['face_with_mask_incorrect', 'face_with_mask_correct', 'face_no_mask']
elif args.dataset == "RMFD":
    NUM_CLASS = 2
    class_names = ['face_with_mask', 'face_no_mask']

# Initialize the initial learning rate, number of epochs to train for, and batch size
LEARNING_RATE = 1e-4
EPOCHS = 20
BATCH_SIZE = 32
IMG_SIZE = 224
dataset_path = "./data/" + args.dataset + "/"
checkpoint_filepath = "./checkpoint_" + args.dataset + "/epoch-{epoch:02d}-val_acc-{val_accuracy:.4f}.h5"
model_save_path = "./mask_detector_models/mask_detector_" + args.dataset + ".h5"
figure_save_path = "./figures/train_plot_" + args.dataset + ".jpg"

print("Num of classes: " + str(NUM_CLASS))
print("Classes: " + str(class_names))
print("Dataset path: " + dataset_path)
print("Checkpoint: " + checkpoint_filepath)
print("Figure save path: " + figure_save_path)

# Construct the training/validation image generator for data augmentation
data_generator = ImageDataGenerator(
    rotation_range=20,
    zoom_range=0.15,
    width_shift_range=0.2,
    height_shift_range=0.2,
    shear_range=0.15,
    horizontal_flip=True,
    fill_mode="nearest",
    preprocessing_function=preprocess_input,
    validation_split=0.2)

# Set as training data
train_generator = data_generator.flow_from_directory(
    dataset_path,
    target_size=(IMG_SIZE, IMG_SIZE),
    batch_size=BATCH_SIZE,
    class_mode="categorical",
    shuffle=True,
    subset='training')

# Set as validation data
validation_generator = data_generator.flow_from_directory(
    dataset_path,
    target_size=(IMG_SIZE, IMG_SIZE),
    batch_size=BATCH_SIZE,
    class_mode="categorical",
    shuffle=False,
    subset='validation')

# Load the pre-trained model and remove the head FC layer
base_model = MobileNetV2(weights="imagenet", include_top=False, input_tensor=Input(shape=(IMG_SIZE, IMG_SIZE, 3)))

# Construct the head of the model that will be placed on top of the base model
head_model = base_model.output
head_model = AveragePooling2D(pool_size=(7, 7))(head_model)
head_model = Flatten(name="flatten")(head_model)
head_model = Dense(128, activation="relu")(head_model)
head_model = Dropout(0.5)(head_model)
head_model = Dense(NUM_CLASS, activation="softmax")(head_model)

# Place the head FC model on top of the base model (this will become the actual model we will train)
model = Model(inputs=base_model.input, outputs=head_model)

# Loop over all layers in the base model and freeze them so they will *not* be updated during the first training process
for layer in base_model.layers:
    layer.trainable = False

# Compile our model
print("[INFO] compiling model...")
opt = Adam(lr=LEARNING_RATE)
if args.dataset == "MFN":
    model.compile(loss="categorical_crossentropy", optimizer=opt, metrics=["accuracy"])
elif args.dataset == "RMFD":
    model.compile(loss="binary_crossentropy", optimizer=opt, metrics=["accuracy"])

# Add early stopping criterion
early_stopping = EarlyStopping(
    monitor='val_accuracy',
    min_delta=0.0001,
    patience=3,
    verbose=1,
    mode='auto',
    baseline=None,
    restore_best_weights=True)

# Add model checkpoint
checkpoint = ModelCheckpoint(
    filepath=checkpoint_filepath,
    save_best_only=False,
    save_weights_only=False,
    monitor='val_accuracy',
    mode='auto')

# Train the head of the network
print("[INFO] training head...")
H = model.fit_generator(
    train_generator,
    steps_per_epoch=train_generator.samples // train_generator.batch_size,
    callbacks=[early_stopping, checkpoint],
    validation_data=validation_generator,
    validation_steps=validation_generator.samples // validation_generator.batch_size,
    epochs=EPOCHS)

# Save best model
model.save(model_save_path)

# Create classification report
prediction = model.predict_generator(
    generator=validation_generator,
    verbose=1)
y_pred = np.argmax(prediction, axis=1)
print("Classification Report:")
print(classification_report(validation_generator.classes, y_pred, target_names=class_names))

# Plot the training loss and accuracy
plt.style.use("ggplot")
plt.figure()
plt.plot(np.arange(0, len(H.history["loss"])), H.history["loss"], label="train_loss")
plt.plot(np.arange(0, len(H.history["val_loss"])), H.history["val_loss"], label="val_loss")
plt.plot(np.arange(0, len(H.history["accuracy"])), H.history["accuracy"], label="train_acc")
plt.plot(np.arange(0, len(H.history["val_accuracy"])), H.history["val_accuracy"], label="val_acc")
plt.title("Training Loss and Accuracy")
plt.xlabel("Epoch #")
plt.ylabel("Loss/Accuracy")
plt.legend(loc="lower left")
plt.savefig(figure_save_path)

detect_mask_image.py

# -*- coding: utf-8 -*-
# @Time    : 2021/1/27 11:02
# @Author  : Johnson
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 argparse
import cv2
import os


def detect_mask(img, face_detector, mask_detector, confidence_threshold, image_show=True):
    # Initialize the labels and colors for bounding boxes
    num_class = mask_detector.layers[-1].get_output_at(0).get_shape()[-1]
    labels, colors = None, None
    if num_class == 3:
        labels = ["Face with Mask Incorrect", "Face with Mask Correct", "Face without Mask"]
        colors = [(0, 255, 255), (0, 255, 0), (0, 0, 255)]
    elif num_class == 2:
        labels = ["Face with Mask", "Face without Mask"]
        colors = [(0, 255, 0), (0, 0, 255)]

    # Load the input image from disk, clone it, and grab the image spatial dimensions
    (h, w) = img.shape[:2]

    # Construct a blob from the image
    blob = cv2.dnn.blobFromImage(img, 1.0, (300, 300), (104.0, 177.0, 123.0))

    # Pass the blob through the network and obtain the face detections
    print("[INFO] computing face detections...")
    face_detector.setInput(blob)
    detections = face_detector.forward()

    # Record status
    # MFN: 0 is "mask correctly" and 1 is "no mask"
    # RMFD: 0 is "mask correctly", 1 is "mask incorrectly", and 2 is "no mask"
    status = 0

    # Loop over the detections
    for i in range(0, detections.shape[2]):
        # Extract the confidence (i.e., probability) associated with the detection
        confidence = detections[0, 0, i, 2]

        # Filter out weak detections by ensuring the confidence is greater than the minimum confidence
        if confidence > confidence_threshold:
            # Compute the (x, y)-coordinates of the bounding box for the object
            box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
            (start_x, start_y, end_x, end_y) = box.astype("int")

            # Ensure the bounding boxes fall within the dimensions of the frame
            (start_x, start_y) = (max(0, start_x), max(0, start_y))
            (end_x, end_y) = (min(w - 1, end_x), min(h - 1, end_y))

            # Extract the face ROI, convert it from BGR to RGB channel ordering, resize it to 224x224, and preprocess it
            face = img[start_y:end_y, start_x:end_x]
            if face.shape[0] == 0 or face.shape[1] == 0:
                continue
            face = cv2.cvtColor(face, cv2.COLOR_BGR2RGB)
            face = cv2.resize(face, (224, 224))
            face = img_to_array(face)
            face = preprocess_input(face)
            face = np.expand_dims(face, axis=0)

            # Pass the face through the model to determine if the face has a mask or not
            prediction = mask_detector.predict(face)[0]
            label_idx = np.argmax(prediction)

            # Determine the class label and color we'll use to draw the bounding box and text
            label = labels[label_idx]
            color = colors[label_idx]

            # Update the status
            if num_class == 3:
                if label_idx == 0:
                    temp = 1
                elif label_idx == 1:
                    temp = 0
                else:
                    temp = 2
                status = max(status, temp)
            elif num_class == 2:
                status = max(status, label_idx)

            # Include the probability in the label
            label = "{}: {:.2f}%".format(label, max(prediction) * 100)

            # Display the label and bounding box rectangle on the output frame
            cv2.putText(img, label, (start_x, start_y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.45, color, 2)
            cv2.rectangle(img, (start_x, start_y), (end_x, end_y), color, 2)
        else:
            break

    if image_show is True:
        # Show the output image
        cv2.imshow("Output", img)
        cv2.waitKey(0)

    return status, img


def main():
    # Instantiate an argument parser and parse the arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("--image", "-i", required=True,
                        help="path to input image")
    parser.add_argument("--model", "-m", type=str,
                        default='MFN', choices=['MFN', 'RMFD'],
                        help="face mask detector model")
    parser.add_argument("--confidence", "-c", type=float, default=0.5,
                        help="minimum probability to filter weak face detections")
    args = parser.parse_args()

    # Change the working directory from src to root if needed
    current_full_dir = os.getcwd()
    print("Current working directory: " + current_full_dir)
    if current_full_dir.split("/")[-1] == "src":
        root = current_full_dir[:-4]
        os.chdir(root)
        print("Changed working directory to: " + root)

    # Validate arguments
    if not os.path.isfile(args.image):
        raise ValueError("Please provide a valid path relative to the root directory.")
    if args.model != "MFN" and args.model != "RMFD":
        raise ValueError("Please provide a valid model choice: `MFN` or `RMFD`.")
    if args.confidence > 1 or args.confidence < 0:
        raise ValueError("Please provide a valid confidence value between 0 and 1 (inclusive).")

    # Initialize model save path
    mask_detector_model_path = "./mask_detector_models/mask_detector_" + args.model + ".h5"
    confidence_threshold = args.confidence
    print("Mask detector save path: " + mask_detector_model_path)
    print("Face detector thresholding confidence: " + str(confidence_threshold))

    # Load the face detector model from disk
    print("[INFO] loading face detector model...")
    prototxt_path = "./face_detector_model/deploy.prototxt"
    weights_path = "./face_detector_model/res10_300x300_ssd_iter_140000.caffemodel"
    face_detector = cv2.dnn.readNet(prototxt_path, weights_path)

    # Load the face mask detector model from disk
    print("[INFO] loading face mask detector model...")
    mask_detector = load_model(mask_detector_model_path)

    # Read the image and detect mask
    img = cv2.imread(args.image)
    if img is None:
        raise ValueError("Your file type is not supported.")

    detect_mask(img, face_detector, mask_detector, confidence_threshold)


if __name__ == '__main__':
    main()

detect_mask_video.py

# -*- coding: utf-8 -*-
# @Time    : 2021/1/27 11:03
# @Author  : Johnson
from tensorflow.keras.models import load_model
from detect_mask_image import detect_mask
import argparse
import cv2
import os


def main():
    # Instantiate an argument parser and parse the arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", "-m", type=str,
                        default='MFN', choices=['MFN', 'RMFD'],
                        help="face mask detector model")
    parser.add_argument("--confidence", "-c", type=float, default=0.5,
                        help="minimum probability to filter weak face detections")
    args = parser.parse_args()

    # Change the working directory from src to root if needed
    current_full_dir = os.getcwd()
    print("Current working directory: " + current_full_dir)
    if current_full_dir.split("/")[-1] == "src":
        root = current_full_dir[:-4]
        os.chdir(root)
        print("Changed working directory to: " + root)

    # Validate arguments
    if args.model != "MFN" and args.model != "RMFD":
        raise ValueError("Please provide a valid model choice: `MFN` or `RMFD`.")
    if args.confidence > 1 or args.confidence < 0:
        raise ValueError("Please provide a valid confidence value between 0 and 1 (inclusive).")

    # Initialize model save path
    mask_detector_model_path = "./mask_detector_models/mask_detector_" + args.model + ".h5"
    confidence_threshold = args.confidence
    print("Mask detector save path: " + mask_detector_model_path)
    print("Face detector thresholding confidence: " + str(confidence_threshold))

    # Load the face detector model from disk
    print("[INFO] loading face detector model...")
    prototxt_path = "./face_detector_model/deploy.prototxt"
    weights_path = "./face_detector_model/res10_300x300_ssd_iter_140000.caffemodel"
    face_detector = cv2.dnn.readNet(prototxt_path, weights_path)

    # Load the face mask detector model from disk
    print("[INFO] loading face mask detector model...")
    mask_detector = load_model(mask_detector_model_path)

    # Initialize the video stream and allow the camera sensor to warm up
    print("[INFO] starting video stream...")
    capture = cv2.VideoCapture(0)

    # Loop over the frames from the video stream
    while capture.isOpened():
        # Grab the frame from the threaded video stream and resize it to have a maximum width of 400 pixels
        flags, frame = capture.read()

        # Detect faces in the frame and determine if they are wearing a face mask or not
        detect_mask(frame, face_detector, mask_detector, confidence_threshold)

        # Show the output frame
        cv2.imshow("Frame", frame)
        key = cv2.waitKey(1) & 0xFF

        # If the `q` key was pressed, break from the loop
        if key == ord("q"):
            break

    capture.release()
    cv2.destroyAllWindows()


if __name__ == '__main__':
    main()

模型地址

百度网盘链接，提取码：e0jt