90 lines
2.7 KiB
Python
90 lines
2.7 KiB
Python
|
"""trt_mtcnn.py
|
||
|
|
|
||
|
|
This script demonstrates how to do real-time face detection with
|
||
|
|
Cython wrapped TensorRT optimized MTCNN engine.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import time
|
||
|
|
import argparse
|
||
|
|
|
||
|
|
import cv2
|
||
|
|
from utils.camera import add_camera_args, Camera
|
||
|
|
from utils.display import open_window, set_display, show_fps
|
||
|
|
from utils.mtcnn import TrtMtcnn
|
||
|
|
|
||
|
|
|
||
|
|
WINDOW_NAME = 'TrtMtcnnDemo'
|
||
|
|
BBOX_COLOR = (0, 255, 0) # green
|
||
|
|
|
||
|
|
|
||
|
|
def parse_args():
|
||
|
|
"""Parse input arguments."""
|
||
|
|
desc = ('Capture and display live camera video, while doing '
|
||
|
|
'real-time face detection with TrtMtcnn on Jetson '
|
||
|
|
'Nano')
|
||
|
|
parser = argparse.ArgumentParser(description=desc)
|
||
|
|
parser = add_camera_args(parser)
|
||
|
|
parser.add_argument('--minsize', type=int, default=40,
|
||
|
|
help='minsize (in pixels) for detection [40]')
|
||
|
|
args = parser.parse_args()
|
||
|
|
return args
|
||
|
|
|
||
|
|
|
||
|
|
def show_faces(img, boxes, landmarks):
|
||
|
|
"""Draw bounding boxes and face landmarks on image."""
|
||
|
|
for bb, ll in zip(boxes, landmarks):
|
||
|
|
x1, y1, x2, y2 = int(bb[0]), int(bb[1]), int(bb[2]), int(bb[3])
|
||
|
|
cv2.rectangle(img, (x1, y1), (x2, y2), BBOX_COLOR, 2)
|
||
|
|
for j in range(5):
|
||
|
|
cv2.circle(img, (int(ll[j]), int(ll[j+5])), 2, BBOX_COLOR, 2)
|
||
|
|
return img
|
||
|
|
|
||
|
|
|
||
|
|
def loop_and_detect(cam, mtcnn, minsize):
|
||
|
|
"""Continuously capture images from camera and do face detection."""
|
||
|
|
full_scrn = False
|
||
|
|
fps = 0.0
|
||
|
|
tic = time.time()
|
||
|
|
while True:
|
||
|
|
if cv2.getWindowProperty(WINDOW_NAME, 0) < 0:
|
||
|
|
break
|
||
|
|
img = cam.read()
|
||
|
|
if img is not None:
|
||
|
|
dets, landmarks = mtcnn.detect(img, minsize=minsize)
|
||
|
|
print('{} face(s) found'.format(len(dets)))
|
||
|
|
img = show_faces(img, dets, landmarks)
|
||
|
|
img = show_fps(img, fps)
|
||
|
|
cv2.imshow(WINDOW_NAME, img)
|
||
|
|
toc = time.time()
|
||
|
|
curr_fps = 1.0 / (toc - tic)
|
||
|
|
# calculate an exponentially decaying average of fps number
|
||
|
|
fps = curr_fps if fps == 0.0 else (fps*0.95 + curr_fps*0.05)
|
||
|
|
tic = toc
|
||
|
|
key = cv2.waitKey(1)
|
||
|
|
if key == 27: # ESC key: quit program
|
||
|
|
break
|
||
|
|
elif key == ord('F') or key == ord('f'): # Toggle fullscreen
|
||
|
|
full_scrn = not full_scrn
|
||
|
|
set_display(WINDOW_NAME, full_scrn)
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
args = parse_args()
|
||
|
|
cam = Camera(args)
|
||
|
|
if not cam.isOpened():
|
||
|
|
raise SystemExit('ERROR: failed to open camera!')
|
||
|
|
|
||
|
|
mtcnn = TrtMtcnn()
|
||
|
|
|
||
|
|
open_window(
|
||
|
|
WINDOW_NAME, 'Camera TensorRT MTCNN Demo for Jetson Nano',
|
||
|
|
cam.img_width, cam.img_height)
|
||
|
|
loop_and_detect(cam, mtcnn, args.minsize)
|
||
|
|
|
||
|
|
cam.release()
|
||
|
|
cv2.destroyAllWindows()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
main()
|