提交 0889cde7 编写于 作者: C Channingss

support show detection result, support get video resolution automatically

上级 07f287d7
...@@ -54,7 +54,7 @@ python infer.py --models_dir=/path/to/models --img_paths=/path/to/images --video ...@@ -54,7 +54,7 @@ python infer.py --models_dir=/path/to/models --img_paths=/path/to/images --video
| img_paths |img_paths/video_path 二选一|需要预测的图片目录 | | img_paths |img_paths/video_path 二选一|需要预测的图片目录 |
| video_path |img_paths/video_path 二选一|需要预测的视频目录| | video_path |img_paths/video_path 二选一|需要预测的视频目录|
| use_camera |No|是否打开摄像头进行预测,默认为False | | use_camera |No|是否打开摄像头进行预测,默认为False |
| open_imshow |No|是否进行检测结果实时绘图,默认为False | | open_imshow |No|是否对视频的检测结果实时绘图,默认为False |
| use_gpu |No|是否GPU,默认为False| | use_gpu |No|是否GPU,默认为False|
说明: 说明:
......
...@@ -33,27 +33,30 @@ import argparse ...@@ -33,27 +33,30 @@ import argparse
def parse_args(): def parse_args():
parser = argparse.ArgumentParser('mask detection.') parser = argparse.ArgumentParser('mask detection.')
parser.add_argument( parser.add_argument('--models_dir',
'--models_dir', type=str, default='', help='path of models.') type=str,
parser.add_argument( default='',
'--img_paths', type=str, default='', help='path of images') help='path of models.')
parser.add_argument( parser.add_argument('--img_paths',
'--video_path', type=str, default='', help='path of video.') type=str,
parser.add_argument( default='',
'--use_camera', help='path of images')
type=bool, parser.add_argument('--video_path',
default=False, type=str,
help='switch detect video or camera, default:video.') default='',
parser.add_argument( help='path of video.')
'--open_imshow', parser.add_argument('--use_camera',
type=bool, type=bool,
default=False, default=False,
help='visualize results in real time.') help='switch detect video or camera, default:video.')
parser.add_argument( parser.add_argument('--open_imshow',
'--use_gpu', type=bool,
type=bool, default=False,
default=False, help='visualize video detection results in real time.')
help='switch cpu/gpu, default:cpu.') parser.add_argument('--use_gpu',
type=bool,
default=False,
help='switch cpu/gpu, default:cpu.')
args = parser.parse_args() args = parser.parse_args()
return args return args
...@@ -105,11 +108,10 @@ class MaskClassifier: ...@@ -105,11 +108,10 @@ class MaskClassifier:
h, w = self.EVAL_SIZE[1], self.EVAL_SIZE[0] h, w = self.EVAL_SIZE[1], self.EVAL_SIZE[0]
inputs = [] inputs = []
for face in faces: for face in faces:
im = cv2.resize( im = cv2.resize(face.rect_data, (128, 128),
face.rect_data, (128, 128), fx=0,
fx=0, fy=0,
fy=0, interpolation=cv2.INTER_CUBIC)
interpolation=cv2.INTER_CUBIC)
# HWC -> CHW # HWC -> CHW
im = im.swapaxes(1, 2) im = im.swapaxes(1, 2)
im = im.swapaxes(0, 1) im = im.swapaxes(0, 1)
...@@ -149,8 +151,10 @@ class FaceDetector: ...@@ -149,8 +151,10 @@ class FaceDetector:
def Preprocess(self, image, shrink): def Preprocess(self, image, shrink):
h, w = int(image.shape[1] * shrink), int(image.shape[0] * shrink) h, w = int(image.shape[1] * shrink), int(image.shape[0] * shrink)
im = cv2.resize( im = cv2.resize(image, (h, w),
image, (h, w), fx=0, fy=0, interpolation=cv2.INTER_CUBIC) fx=0,
fy=0,
interpolation=cv2.INTER_CUBIC)
# HWC -> CHW # HWC -> CHW
im = im.swapaxes(1, 2) im = im.swapaxes(1, 2)
im = im.swapaxes(0, 1) im = im.swapaxes(0, 1)
...@@ -190,18 +194,16 @@ class FaceDetector: ...@@ -190,18 +194,16 @@ class FaceDetector:
def predict_images(args): def predict_images(args):
detector = FaceDetector( detector = FaceDetector(model_dir=args.models_dir + '/pyramidbox_lite/',
model_dir=args.models_dir + '/pyramidbox_lite/', mean=[104.0, 177.0, 123.0],
mean=[104.0, 177.0, 123.0], scale=[0.007843, 0.007843, 0.007843],
scale=[0.007843, 0.007843, 0.007843], use_gpu=args.use_gpu,
use_gpu=args.use_gpu, threshold=0.7)
threshold=0.7)
classifier = MaskClassifier(model_dir=args.models_dir + '/mask_detector/',
classifier = MaskClassifier( mean=[0.5, 0.5, 0.5],
model_dir=args.models_dir + '/mask_detector/', scale=[1.0, 1.0, 1.0],
mean=[0.5, 0.5, 0.5], use_gpu=args.use_gpu)
scale=[1.0, 1.0, 1.0],
use_gpu=args.use_gpu)
names = [] names = []
image_paths = [] image_paths = []
for name in os.listdir(args.img_paths): for name in os.listdir(args.img_paths):
...@@ -227,18 +229,16 @@ def predict_video(args, im_shape=(1920, 1080), use_camera=False): ...@@ -227,18 +229,16 @@ def predict_video(args, im_shape=(1920, 1080), use_camera=False):
capture = cv2.VideoCapture(0) capture = cv2.VideoCapture(0)
else: else:
capture = cv2.VideoCapture(args.video_path) capture = cv2.VideoCapture(args.video_path)
detector = FaceDetector( detector = FaceDetector(model_dir=args.models_dir + '/pyramidbox_lite/',
model_dir=args.models_dir + '/pyramidbox_lite/', mean=[104.0, 177.0, 123.0],
mean=[104.0, 177.0, 123.0], scale=[0.007843, 0.007843, 0.007843],
scale=[0.007843, 0.007843, 0.007843], use_gpu=args.use_gpu,
use_gpu=args.use_gpu, threshold=0.7)
threshold=0.7)
classifier = MaskClassifier(model_dir=args.models_dir + '/mask_detector/',
classifier = MaskClassifier( mean=[0.5, 0.5, 0.5],
model_dir=args.models_dir + '/mask_detector/', scale=[1.0, 1.0, 1.0],
mean=[0.5, 0.5, 0.5], use_gpu=args.use_gpu)
scale=[1.0, 1.0, 1.0],
use_gpu=args.use_gpu)
path = './result' path = './result'
isExists = os.path.exists(path) isExists = os.path.exists(path)
...@@ -248,8 +248,8 @@ def predict_video(args, im_shape=(1920, 1080), use_camera=False): ...@@ -248,8 +248,8 @@ def predict_video(args, im_shape=(1920, 1080), use_camera=False):
width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)) width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
fourcc = cv2.VideoWriter_fourcc(*'mp4v') fourcc = cv2.VideoWriter_fourcc(*'mp4v')
writer = cv2.VideoWriter( writer = cv2.VideoWriter(os.path.join(path, 'result.mp4'), fourcc, fps,
os.path.join(path, 'result.mp4'), fourcc, fps, (width, height)) (width, height))
import time import time
start_time = time.time() start_time = time.time()
index = 0 index = 0
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册