add age, gender, emotion stubs

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 = []