face_detect.py

import cv2
import numpy
import argparse
import time
import re

from edgetpu.detection.engine import DetectionEngine
from PIL import Image, ImageDraw, ImageFont

from imutils.video import FPS
from imutils.video import VideoStream


def ReadLabelFile(file_path):
    with open(file_path, 'r', encoding='utf-8') as f:
        lines = f.readlines()
    ret = {}
    for line in lines:
        pair = re.split(r'[:\s]+', line.strip(), maxsplit=1)
        ret[int(pair[0])] = pair[1].strip()
    return ret


def draw_image(image, results, labels):
    result_size = len(results)
    for idx, obj in enumerate(results):

        # Prepare image for drawing
        draw = ImageDraw.Draw(image)

        # Prepare boundary box
        box = obj.bounding_box.flatten().tolist()

        # Draw rectangle to desired thickness
        for x in range( 0, 4 ):
            draw.rectangle(box, outline=(255, 255, 0))

        # Annotate image with label and confidence score
        display_str = labels[obj.label_id] + ": " + str(round(obj.score*100, 2)) + "%"
        draw.text((box[0], box[1]), display_str, font=ImageFont.truetype("/usr/share/fonts/truetype/piboto/Piboto-Regular.ttf", 20))

        displayImage = numpy.asarray(image)
        cv2.imshow('Coral Live Object Detection', displayImage)


def main():
    model = "path to model"
    threshold = 0.3
    max_objects = 3

    # Prepare labels.
    labels = None
    # Initialize engine.
    engine = DetectionEngine(args.model)

    # Initialize video stream
    vs = VideoStream(usePiCamera=True, resolution=(640, 480)).start()
    time.sleep(1)

    fps = FPS().start()

    while True:
        try:
            # Read frame from video
            screenshot = vs.read()
            image = Image.fromarray(screenshot)

            # Perform inference
            results = engine.DetectWithImage(image, threshold=threshold, keep_aspect_ratio=True, relative_coord=False, top_k=maxobjects)

            # draw image
            draw_image(image, results, labels)

            # closing condition
            if cv2.waitKey(5) & 0xFF == ord('q'):
                fps.stop()
                break

            fps.update()
        except KeyboardInterrupt:
            fps.stop()
            break

    print("Elapsed time: " + str(fps.elapsed()))
    print("Approx FPS: :" + str(fps.fps()))

    cv2.destroyAllWindows()
    vs.stop()
    time.sleep(2)


if __name__ == '__main__':
    main()