import os from os.path import join from pathlib import Path import cv2 as cv import numpy as np import dlib import imutils from app.utils import im_utils, logger_utils from app.models.bbox import BBox from app.settings import app_cfg as cfg from app.settings import types class RecognitionDLIB: # https://github.com/davisking/dlib/blob/master/python_examples/face_recognition.py # facerec.compute_face_descriptor(img, shape, 100, 0.25) def __init__(self, opt_gpu=0): self.log = logger_utils.Logger.getLogger() if opt_gpu > 0: cuda_visible_devices = os.getenv('CUDA_VISIBLE_DEVICES', '') os.environ['CUDA_VISIBLE_DEVICES'] = str(opt_gpu) self.predictor = dlib.shape_predictor(cfg.DIR_MODELS_DLIB_5PT) self.facerec = dlib.face_recognition_model_v1(cfg.DIR_MODELS_DLIB_FACEREC_RESNET) os.environ['CUDA_VISIBLE_DEVICES'] = cuda_visible_devices # reset GPU env def vec(self, im, bbox, width=100, jitters=cfg.DLIB_FACEREC_JITTERS, padding=cfg.DLIB_FACEREC_PADDING): # Converts image and bbox into 128d vector # scale the image so the face is always 100x100 pixels scale = width / bbox.width im = cv.resize(im, (scale, scale), interploation=cv.INTER_LANCZOS4) bbox_dlib = bbox.to_dlib() face_shape = self.predictor(im, bbox_dlib) vec = self.facerec.compute_face_descriptor(im, face_shape, jitters, padding) return vec def similarity(self, query_enc, known_enc): return np.linalg.norm(query_enc - known_enc, axis=1)