summaryrefslogtreecommitdiff
path: root/megapixels/commands/demo/face_landmarks_2d.py
blob: 22e09297524a336c414f15bd4e25f55f72ebe843 (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
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
"""
Crop images to prepare for training
"""

import click
# from PIL import Image, ImageOps, ImageFilter, ImageDraw

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('-o', '--output', 'opt_fp_out', default=None,
  help='GIF output path')
@click.option('--size', 'opt_size', 
  type=(int, int), default=(300, 300),
  help='Output image size')
@click.option('--gif-size', 'opt_gif_size', 
  type=(int, int), default=(480, 480),
  help='GIF output size')
@click.option('--gif-frames', 'opt_gif_frames', default=15,
  help='GIF frames')
@click.option('-g', '--gpu', 'opt_gpu', default=0,
  help='GPU index')
@click.option('-f', '--force', 'opt_force', is_flag=True,
  help='Force overwrite file')
@click.option('--display/--no-display', 'opt_display', is_flag=True, default=False,
  help='Display detections to debug')
@click.pass_context
def cli(ctx, opt_fp_in, opt_fp_out, opt_gpu, opt_gif_frames, 
  opt_size, opt_gif_size, opt_force, opt_display):
  """Generates 3D landmark animations from CSV files"""
  
  import sys
  import os
  from os.path import join
  from pathlib import Path
  import time
  
  from tqdm import tqdm
  import numpy as np
  import pandas as pd
  import cv2 as cv
  import dlib
  from PIL import Image
  import matplotlib.pyplot as plt

  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, face_landmarks
  from app.models.data_store import DataStore
  
  # TOOD add selective testing
  opt_run_pose = True
  opt_run_2d_68 = True
  opt_run_3d_68 = True
  opt_run_3d_68 = True


  # -------------------------------------------------
  # 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.DetectorDLIBCNN(gpu=opt_gpu)  # -1 for CPU
  log.info('detecting face...')
  st = time.time()
  bboxes = face_detector.detect(im_resized, largest=True)
  bbox = bboxes[0]
  dim = im_resized.shape[:2][::-1]
  bbox_dim = bbox.to_dim(dim)
  if not bbox:
    log.error('no face detected')
    return
  else:
    log.info(f'Detected face in {(time.time() - st):.2f}s')
  log.info('')


  # ----------------------------------------------------------------------------
  # detect 3D landmarks
  
  log.info('loading 3D landmark generator files...')
  landmark_detector_3d_68 = face_landmarks.FaceAlignment3D_68(gpu=opt_gpu)  # -1 for CPU
  log.info('generating 3D landmarks...')
  st = time.time()
  points_3d_68 = landmark_detector_3d_68.landmarks(im_resized, bbox_dim.to_xyxy())
  log.info(f'generated 3D landmarks in {(time.time() - st):.2f}s')
  log.info('')


  # ----------------------------------------------------------------------------
  # generate 3D GIF animation
  
  log.info('generating 3D animation...')
  if not opt_fp_out:
    fpp_im = Path(opt_fp_in)
    fp_out = join(fpp_im.parent, f'{fpp_im.stem}_anim.gif')
  else:
    fp_out = opt_fp_out
  st = time.time()
  plot_utils.generate_3d_landmark_anim(np.array(points_3d_68), fp_out, 
    size=opt_gif_size, num_frames=opt_gif_frames)
  log.info(f'Generated animation in {(time.time() - st):.2f}s')
  log.info(f'Saved to: {fp_out}')
  log.info('')


  # ----------------------------------------------------------------------------
  # generate face vectors, only to test if feature extraction works
  
  log.info('initialize face recognition model...')
  from app.processors import face_recognition
  face_rec = face_recognition.RecognitionDLIB()
  st = time.time()
  log.info('generating face vector...')
  vec = face_rec.vec(im_resized, bbox_dim)
  log.info(f'generated face vector in {(time.time() - st):.2f}s')
  log.info('')


  # ----------------------------------------------------------------------------
  # generate 68 point landmarks using dlib
  
  log.info('initializing face landmarks 68 dlib...')
  from app.processors import face_landmarks
  landmark_detector_2d_68 = face_landmarks.Dlib2D_68()
  log.info('generating 2D 68PT landmarks...')
  st = time.time()
  points_2d_68 = landmark_detector_2d_68.landmarks(im_resized, bbox_dim)
  log.info(f'generated 2D 68PT face landmarks in {(time.time() - st):.2f}s')
  log.info('')


  # ----------------------------------------------------------------------------
  # generate pose from 68 point 2D landmarks
  
  if opt_run_pose:
    log.info('initialize pose...')
    from app.processors import face_pose
    pose_detector = face_pose.FacePoseDLIB()
    log.info('generating pose...')
    st = time.time()
    pose_data = pose_detector.pose(points_2d_68, dim)
    log.info(f'generated pose {(time.time() - st):.2f}s')
    log.info('')


  # x



  # display
  if opt_display:
    
    # draw bbox
    
    # draw 3d landmarks
    im_landmarks_3d_68 = im_resized.copy()
    draw_utils.draw_landmarks3D(im_landmarks_3d_68, points_3d_68)
    draw_utils.draw_bbox(im_landmarks_3d_68, bbox_dim)

    # draw 2d landmarks
    im_landmarks_2d_68 = im_resized.copy()
    draw_utils.draw_landmarks2D(im_landmarks_2d_68, points_2d_68)
    draw_utils.draw_bbox(im_landmarks_2d_68, bbox_dim)

    # draw pose
    if opt_run_pose:
      im_pose = im_resized.copy()
      draw_utils.draw_pose(im_pose, pose_data['point_nose'], pose_data['points'])
      draw_utils.draw_degrees(im_pose, pose_data)

    # draw animated GIF
    im = Image.open(opt_fp_out)
    im_frames = []
    duration = im.info['duration']
    try:
      while True:
        im.seek(len(im_frames))
        mypalette = im.getpalette()
        im.putpalette(mypalette)
        im_jpg = Image.new("RGB", im.size)
        im_jpg.paste(im)
        im_np = im_utils.pil2np(im_jpg.copy())
        im_frames.append(im_np)
    except EOFError:
      pass # end of GIF sequence

    n_frames = len(im_frames)
    frame_number = 0

    while True:
      # show all images here
      cv.imshow('Original', im_resized)
      cv.imshow('2D 68PT Landmarks', im_landmarks_2d_68)
      cv.imshow('3D 68PT Landmarks', im_landmarks_3d_68)
      cv.imshow('Pose', im_pose)
      cv.imshow('3D 68pt GIF', im_frames[frame_number])
      frame_number = (frame_number + 1) % n_frames
      k = cv.waitKey(duration) & 0xFF
      if k == 27 or k == ord('q'):  # ESC
        cv.destroyAllWindows()
        sys.exit()
      elif k != 255:
        # any key to continue
        break