diff options
Diffstat (limited to 'megapixels/app/processors/face_detector.py')
| -rw-r--r-- | megapixels/app/processors/face_detector.py | 11 |
1 files changed, 8 insertions, 3 deletions
diff --git a/megapixels/app/processors/face_detector.py b/megapixels/app/processors/face_detector.py index c0762564..0e194f7d 100644 --- a/megapixels/app/processors/face_detector.py +++ b/megapixels/app/processors/face_detector.py @@ -21,9 +21,14 @@ class DetectorMTCNN: dnn_size = (300, 300) - def __init__(self, size=(400,400)): + def __init__(self, size=(400,400), gpu=0): + self.log = logger_utils.Logger.getLogger() + device_cur = os.getenv('CUDA_VISIBLE_DEVICES', '') + self.log.info(f'Change CUDA_VISIBLE_DEVICES from "{device_cur}" to "{gpu}"') + os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu) from mtcnn.mtcnn import MTCNN self.detector = MTCNN() + os.environ['CUDA_VISIBLE_DEVICES'] = device_cur # reset def detect(self, im, size=(400,400), conf_thresh=None, pyramids=None, largest=False, zone=None): '''Detects face using MTCNN and returns (list) of BBox @@ -70,14 +75,14 @@ class DetectorDLIBCNN: def __init__(self, gpu=0): import dlib self.log = logger_utils.Logger.getLogger() - cuda_visible_devices = os.getenv('CUDA_VISIBLE_DEVICES', '') + device_cur = os.getenv('CUDA_VISIBLE_DEVICES', '') if dlib.DLIB_USE_CUDA and gpu < 0: self.log.error('dlib was compiled with CUDA but you selected CPU. Use GPU >= 0 if dlib.DLIB_USE_CUDA') sys.exit() os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu) self.log.info('load model: {}'.format(cfg.DIR_MODELS_DLIB_CNN)) self.detector = dlib.cnn_face_detection_model_v1(cfg.DIR_MODELS_DLIB_CNN) - os.environ['CUDA_VISIBLE_DEVICES'] = cuda_visible_devices # reset + os.environ['CUDA_VISIBLE_DEVICES'] = device_cur # reset def detect(self, im, conf_thresh=None, pyramids=None, largest=False, zone=None): bboxes = [] |