summaryrefslogtreecommitdiff
path: root/megapixels/commands/demo/face_vector.py
blob: c7b5ef2e5d2b91f17fd0452cd5d912477ee2996d (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
"""
Tests if the feature vector generator works
"""

import click

from app.settings import types
from app.utils import click_utils
from app.settings import app_cfg as cfg


@click.command()
@click.option('-i', '--input', 'opt_fp_in', default=None, required=True,
  help='Image filepath')
@click.option('--size', 'opt_size', 
  type=(int, int), default=cfg.DEFAULT_SIZE_FACE_DETECT,
  help='Output image size') 
@click.option('-g', '--gpu', 'opt_gpu', default=0,
  help='GPU index')
@click.option('--display/--no-display', 'opt_display', is_flag=True, default=True,
  help='Display detections to debug')
@click.pass_context
def cli(ctx, opt_fp_in, opt_gpu, opt_size, opt_display):
  """Demo generating face vector"""
  
  import sys
  import os
  from os.path import join
  from pathlib import Path
  import time
  
  import numpy as np
  import cv2 as cv
  import dlib  # NB: keep a reference in main file if using dlib detector processors

  from app.utils import logger_utils, file_utils, im_utils, display_utils, draw_utils
  from app.utils import plot_utils
  from app.processors import face_detector
  from app.models.data_store import DataStore

  # -------------------------------------------------
  # init here

  log = logger_utils.Logger.getLogger()

  # ----------------------------------------------------------------------------
  # load image

  im = cv.imread(opt_fp_in)
  im_resized = im_utils.resize(im, width=opt_size[0], height=opt_size[1])

  
  # ----------------------------------------------------------------------------
  # detect face

  face_detector = face_detector.DetectorCVDNN()  # -1 for CPU
  bboxes = face_detector.detect(im_resized, largest=True)
  bbox_norm = bboxes[0]
  dim = im_resized.shape[:2][::-1]
  bbox_dim = bbox_norm.to_dim(dim)
  if not bbox_norm:
    log.error('no face detected')
    return


  # ----------------------------------------------------------------------------
  # generate face vectors, only to test if feature extraction works
  
  from app.processors import face_extractor
  extractor = face_extractor.ExtractorVGG()
  vec = extractor.extract(im_resized, bbox_norm)
  vec_str = extractor.to_str(vec)
  log.info(f'\n{vec_str}')

  if opt_display:
    im_resized = draw_utils.draw_bbox(im_resized, bbox_dim)
    cv.imshow('Original', im_resized)
    display_utils.handle_keyboard()