face-mask-detection.py

import sys
import time

import cv2

import ailia

# import original modules
sys.path.append('../../util')
from utils import get_base_parser, update_parser  # noqa: E402
from model_utils import check_and_download_models  # noqa: E402
import webcamera_utils  # noqa: E402
from detector_utils import plot_results, load_image  # noqa: E402
from nms_utils import nms_between_categories  # noqa: E402


# ======================
# Parameters
# ======================

MODEL_LISTS = ['yolov3-tiny', 'yolov3', 'mb2-ssd']

IMAGE_PATH = 'ferry.jpg'
SAVE_IMAGE_PATH = 'output.png'
IMAGE_HEIGHT = 416
IMAGE_WIDTH = 416

FACE_CATEGORY = ['unmasked', 'masked']

IOU = 0.45


# ======================
# Arguemnt Parser Config
# ======================
parser = get_base_parser(
    'masked face detection model', IMAGE_PATH, SAVE_IMAGE_PATH
)
parser.add_argument(
    '-a', '--arch', metavar='ARCH',
    default='yolov3-tiny', choices=MODEL_LISTS,
    help='model lists: ' + ' | '.join(MODEL_LISTS)
)
args = update_parser(parser)


if args.arch == "yolov3-tiny":
    WEIGHT_PATH = 'face-mask-detection-yolov3-tiny.opt.obf.onnx'
    MODEL_PATH = 'face-mask-detection-yolov3-tiny.opt.onnx.prototxt'
    RANGE = ailia.NETWORK_IMAGE_RANGE_U_FP32
    ALGORITHM = ailia.DETECTOR_ALGORITHM_YOLOV3
    THRESHOLD = 0.4
elif args.arch == "yolov3":
    WEIGHT_PATH = 'face-mask-detection-yolov3.opt.obf.onnx'
    MODEL_PATH = 'face-mask-detection-yolov3.opt.onnx.prototxt'
    RANGE = ailia.NETWORK_IMAGE_RANGE_U_FP32
    ALGORITHM = ailia.DETECTOR_ALGORITHM_YOLOV3
    THRESHOLD = 0.4
else:
    WEIGHT_PATH = 'face-mask-detection-mb2-ssd-lite.obf.onnx'
    MODEL_PATH = 'face-mask-detection-mb2-ssd-lite.onnx.prototxt'
    RANGE = ailia.NETWORK_IMAGE_RANGE_S_FP32
    ALGORITHM = ailia.DETECTOR_ALGORITHM_SSD
    THRESHOLD = 0.2
REMOTE_PATH = \
    'https://storage.googleapis.com/ailia-models/face-mask-detection/'


# ======================
# Main functions
# ======================
def recognize_from_image():
    # prepare input data
    img = load_image(args.input)
    print(f'input image shape: {img.shape}')

    # net initialize
    detector = ailia.Detector(
        MODEL_PATH,
        WEIGHT_PATH,
        len(FACE_CATEGORY),
        format=ailia.NETWORK_IMAGE_FORMAT_RGB,
        channel=ailia.NETWORK_IMAGE_CHANNEL_FIRST,
        range=RANGE,
        algorithm=ALGORITHM,
        env_id=args.env_id
    )

    # inference
    print('Start inference...')
    if args.benchmark:
        print('BENCHMARK mode')
        for i in range(5):
            start = int(round(time.time() * 1000))
            detector.compute(img, THRESHOLD, IOU)
            end = int(round(time.time() * 1000))
            print(f'\tailia processing time {end - start} ms')
    else:
        detector.compute(img, THRESHOLD, IOU)

    # nms
    detections = []
    for idx in range(detector.get_object_count()):
        obj = detector.get_object(idx)
        detections.append(obj)
    detections = nms_between_categories(
        detections,
        img.shape[1],
        img.shape[0],
        categories=[0, 1],
        iou_threshold=IOU
    )

    # plot result
    res_img = plot_results(detections, img, FACE_CATEGORY)
    cv2.imwrite(args.savepath, res_img)
    print('Script finished successfully.')


def recognize_from_video():
    # net initialize
    detector = ailia.Detector(
        MODEL_PATH,
        WEIGHT_PATH,
        len(FACE_CATEGORY),
        format=ailia.NETWORK_IMAGE_FORMAT_RGB,
        channel=ailia.NETWORK_IMAGE_CHANNEL_FIRST,
        range=RANGE,
        algorithm=ALGORITHM,
        env_id=args.env_id
    )

    capture = webcamera_utils.get_capture(args.video)

    if args.savepath != SAVE_IMAGE_PATH:
        writer = webcamera_utils.get_writer(
            args.savepath,
            IMAGE_HEIGHT,
            IMAGE_WIDTH,
            fps=capture.get(cv2.CAP_PROP_FPS),
        )
    else:
        writer = None

    while(True):
        ret, frame = capture.read()
        if (cv2.waitKey(1) & 0xFF == ord('q')) or not ret:
            break

        _, resized_img = webcamera_utils.adjust_frame_size(
            frame, IMAGE_HEIGHT, IMAGE_WIDTH
        )

        img = cv2.cvtColor(resized_img, cv2.COLOR_BGR2BGRA)
        detector.compute(img, THRESHOLD, IOU)

        detections = []
        for idx in range(detector.get_object_count()):
            obj = detector.get_object(idx)
            detections.append(obj)
        detections = nms_between_categories(
            detections,
            frame.shape[1],
            frame.shape[0],
            categories=[0, 1],
            iou_threshold=IOU
        )

        res_img = plot_results(detections, resized_img, FACE_CATEGORY, False)
        cv2.imshow('frame', res_img)

        # save results
        if writer is not None:
            writer.write(res_img)

    capture.release()
    cv2.destroyAllWindows()
    if writer is not None:
        writer.release()
    print('Script finished successfully.')


def main():
    # model files check and download
    check_and_download_models(WEIGHT_PATH, MODEL_PATH, REMOTE_PATH)

    if args.video is not None:
        # video mode
        recognize_from_video()
    else:
        # image mode
        recognize_from_image()


if __name__ == '__main__':
    main()