From 0889cde7600b88a5470ede65cf5b791bec1a23d7 Mon Sep 17 00:00:00 2001 From: Channingss Date: Mon, 2 Mar 2020 10:33:50 +0800 Subject: [PATCH] support show detection result, support get video resolution automatically --- demo/mask_detection/python/README.md | 2 +- demo/mask_detection/python/infer.py | 108 +++++++++++++-------------- 2 files changed, 55 insertions(+), 55 deletions(-) diff --git a/demo/mask_detection/python/README.md b/demo/mask_detection/python/README.md index 042e8a42..b7314362 100644 --- a/demo/mask_detection/python/README.md +++ b/demo/mask_detection/python/README.md @@ -54,7 +54,7 @@ python infer.py --models_dir=/path/to/models --img_paths=/path/to/images --video | img_paths |img_paths/video_path 二选一|需要预测的图片目录 | | video_path |img_paths/video_path 二选一|需要预测的视频目录| | use_camera |No|是否打开摄像头进行预测，默认为False | -| open_imshow |No|是否进行检测结果实时绘图，默认为False | +| open_imshow |No|是否对视频的检测结果实时绘图，默认为False | | use_gpu |No|是否GPU，默认为False| 说明： diff --git a/demo/mask_detection/python/infer.py b/demo/mask_detection/python/infer.py index 074b1646..710377d6 100644 --- a/demo/mask_detection/python/infer.py +++ b/demo/mask_detection/python/infer.py @@ -33,27 +33,30 @@ import argparse def parse_args(): parser = argparse.ArgumentParser('mask detection.') - parser.add_argument( - '--models_dir', type=str, default='', help='path of models.') - parser.add_argument( - '--img_paths', type=str, default='', help='path of images') - parser.add_argument( - '--video_path', type=str, default='', help='path of video.') - parser.add_argument( - '--use_camera', - type=bool, - default=False, - help='switch detect video or camera, default:video.') - parser.add_argument( - '--open_imshow', - type=bool, - default=False, - help='visualize results in real time.') - parser.add_argument( - '--use_gpu', - type=bool, - default=False, - help='switch cpu/gpu, default:cpu.') + parser.add_argument('--models_dir', + type=str, + default='', + help='path of models.') + parser.add_argument('--img_paths', + type=str, + default='', + help='path of images') + parser.add_argument('--video_path', + type=str, + default='', + help='path of video.') + parser.add_argument('--use_camera', + type=bool, + default=False, + help='switch detect video or camera, default:video.') + parser.add_argument('--open_imshow', + type=bool, + default=False, + help='visualize video detection results in real time.') + parser.add_argument('--use_gpu', + type=bool, + default=False, + help='switch cpu/gpu, default:cpu.') args = parser.parse_args() return args @@ -105,11 +108,10 @@ class MaskClassifier: h, w = self.EVAL_SIZE[1], self.EVAL_SIZE[0] inputs = [] for face in faces: - im = cv2.resize( - face.rect_data, (128, 128), - fx=0, - fy=0, - interpolation=cv2.INTER_CUBIC) + im = cv2.resize(face.rect_data, (128, 128), + fx=0, + fy=0, + interpolation=cv2.INTER_CUBIC) # HWC -> CHW im = im.swapaxes(1, 2) im = im.swapaxes(0, 1) @@ -149,8 +151,10 @@ class FaceDetector: def Preprocess(self, image, shrink): h, w = int(image.shape[1] * shrink), int(image.shape[0] * shrink) - im = cv2.resize( - image, (h, w), fx=0, fy=0, interpolation=cv2.INTER_CUBIC) + im = cv2.resize(image, (h, w), + fx=0, + fy=0, + interpolation=cv2.INTER_CUBIC) # HWC -> CHW im = im.swapaxes(1, 2) im = im.swapaxes(0, 1) @@ -190,18 +194,16 @@ class FaceDetector: def predict_images(args): - detector = FaceDetector( - model_dir=args.models_dir + '/pyramidbox_lite/', - mean=[104.0, 177.0, 123.0], - scale=[0.007843, 0.007843, 0.007843], - use_gpu=args.use_gpu, - threshold=0.7) - - classifier = MaskClassifier( - model_dir=args.models_dir + '/mask_detector/', - mean=[0.5, 0.5, 0.5], - scale=[1.0, 1.0, 1.0], - use_gpu=args.use_gpu) + detector = FaceDetector(model_dir=args.models_dir + '/pyramidbox_lite/', + mean=[104.0, 177.0, 123.0], + scale=[0.007843, 0.007843, 0.007843], + use_gpu=args.use_gpu, + threshold=0.7) + + classifier = MaskClassifier(model_dir=args.models_dir + '/mask_detector/', + mean=[0.5, 0.5, 0.5], + scale=[1.0, 1.0, 1.0], + use_gpu=args.use_gpu) names = [] image_paths = [] for name in os.listdir(args.img_paths): @@ -227,18 +229,16 @@ def predict_video(args, im_shape=(1920, 1080), use_camera=False): capture = cv2.VideoCapture(0) else: capture = cv2.VideoCapture(args.video_path) - detector = FaceDetector( - model_dir=args.models_dir + '/pyramidbox_lite/', - mean=[104.0, 177.0, 123.0], - scale=[0.007843, 0.007843, 0.007843], - use_gpu=args.use_gpu, - threshold=0.7) - - classifier = MaskClassifier( - model_dir=args.models_dir + '/mask_detector/', - mean=[0.5, 0.5, 0.5], - scale=[1.0, 1.0, 1.0], - use_gpu=args.use_gpu) + detector = FaceDetector(model_dir=args.models_dir + '/pyramidbox_lite/', + mean=[104.0, 177.0, 123.0], + scale=[0.007843, 0.007843, 0.007843], + use_gpu=args.use_gpu, + threshold=0.7) + + classifier = MaskClassifier(model_dir=args.models_dir + '/mask_detector/', + mean=[0.5, 0.5, 0.5], + scale=[1.0, 1.0, 1.0], + use_gpu=args.use_gpu) path = './result' isExists = os.path.exists(path) @@ -248,8 +248,8 @@ def predict_video(args, im_shape=(1920, 1080), use_camera=False): width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) fourcc = cv2.VideoWriter_fourcc(*'mp4v') - writer = cv2.VideoWriter( - os.path.join(path, 'result.mp4'), fourcc, fps, (width, height)) + writer = cv2.VideoWriter(os.path.join(path, 'result.mp4'), fourcc, fps, + (width, height)) import time start_time = time.time() index = 0 -- GitLab