summaryrefslogtreecommitdiff
path: root/megapixels/app/processors/face_extractor.py
blob: 2666e090c5a9702f014db2c8d3ceb97d704804a7 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
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

def similarity(self, query_enc, known_enc):
    return np.linalg.norm(query_enc - known_enc, axis=1)

def flatten(vec):
    '''Converts N-D vector into a flattened list for CSV
    :param points: (list) a feature vector as list of floats
    :returns dict item for each point (eg {'d1':0.28442156, 'd1': 0.1868632})
    '''
    vec_flat = {}
    for idx, val in enumerate(vec, 1):
      vec_flat[f'd{idx}'] = val
    return vec_flat



class Extractor:

  n_dim = None  # override

  def __init__(self):
    self.log = logger_utils.Logger.getLogger()

  def flatten(self, vec):
    '''Converts N-D vector into a flattened list for CSV
    :param points: (list) a feature vector as list of floats
    :returns dict item for each point (eg {'d1':0.28442156, 'd1': 0.1868632})
    '''
    vec_flat = {}
    for idx, val in enumerate(vec, 1):
      vec_flat[f'd{idx}'] = val
    return vec_flat

  def unflatten_df(self, df):
    # convert from
    return [df[f'd{i}'] for i in range(1,257)]


class ExtractorVGG(Extractor):

  # https://github.com/ox-vgg/vgg_face2
  # Uses OpenCV DNN to extract feature vector for VGG Face 2 models
  n_dim = 256
  dnn_dim = (224,224)
  dnn_mean = (91.4953, 103.8827, 131.0912)

  def __init__(self):
    super().__init__()
    fp_model = '/data_store_hdd/apps/megapixels/models/caffe/vgg_face2/resnet50_256_caffe/resnet50_256.caffemodel'
    fp_prototxt = '/data_store_hdd/apps/megapixels/models/caffe/vgg_face2/resnet50_256_caffe/resnet50_256.prototxt'
    self.dnn = cv.dnn.readNetFromCaffe(fp_prototxt, fp_model)
    self.feat_layer = self.dnn.getLayerNames()[-2]

  def extract(self, im, bbox_norm, padding=0.3):
    '''Extracts feature vector for face crop
    :param im:
    :param bbox_norm: (BBox) normalized
    :param padding: (float) percent to extend ROI
    :param jitters: not used here
    :returns (list) of (float)'''
    
    bbox_ext = bbox_norm.expand(padding)
    dim = im.shape[:2][::-1]
    bbox_ext_dim = bbox_ext.to_dim(dim)
    x1,y1,x2,y2 = bbox_ext_dim.to_xyxy()
    im = im[y1:y2, x1:x2]
    # According to VGG, model trained using Bilinear interpolation (INTER_LINEAR)
    im = cv.resize(im, self.dnn_dim, interpolation=cv.INTER_LINEAR)
    blob = cv.dnn.blobFromImage(im, 1.0, self.dnn_dim, self.dnn_mean)
    self.dnn.setInput(blob)
    vec = np.array(self.dnn.forward(self.feat_layer)[0])
    vec_norm = np.array(vec)/np.linalg.norm(vec)  # normalize
    return vec_norm


class ExtractorDLIB(Extractor):

  # https://github.com/davisking/dlib/blob/master/python_examples/face_recognition.py
  # facerec.compute_face_descriptor(img, shape, 100, 0.25)
  # padding=opt_padding not yet implemented in dlib===19.16 but merged in master
  n_dim = 128
  process_width = 100

  def __init__(self, gpu=0, jitters=cfg.DLIB_FACEREC_JITTERS):
    super().__init__()
    self.num_jitters = cfg.DLIB_FACEREC_JITTERS
    # set and swap GPU visibility
    if gpu > -1:
      cuda_visible_devices = os.getenv('CUDA_VISIBLE_DEVICES', '')
      os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
    self.predictor = dlib.shape_predictor(cfg.DIR_MODELS_DLIB_5PT)
    self.facerec = dlib.face_recognition_model_v1(cfg.DIR_MODELS_DLIB_FACEREC_RESNET)
    # unset and swap GPU visibility
    if gpu > -1:
      os.environ['CUDA_VISIBLE_DEVICES'] = cuda_visible_devices  # reset GPU env


  def extract(self, im, bbox_norm):
    '''Converts image and bbox into 128d vector
    :param im: (numpy.ndarray) BGR image
    :param bbox_norm: (BBox) normalized
    '''
    # scale the image so the face is always 100x100 pixels
    dim = im.shape[:2][::-1]
    bbox_dim = bbox_norm.to_dim(dim)
    scale = self.process_width / bbox_dim.width
    cv.resize(im, None, fx=scale, fy=scale, interpolation=cv.INTER_LANCZOS4) 
    bbox_dim_dlib = bbox_dim.to_dlib()
    face_shape = self.predictor(im, bbox_dim_dlib)
    # this is only in dlib version 19.6++?
    # vec = self.facerec.compute_face_descriptor(im, face_shape, self.num_jitters, self.padding)
    # vectors are already normalized
    vec = self.facerec.compute_face_descriptor(im, face_shape, self.num_jitters)
    return vec