repo_description: Powerful & Easy-to-Use Video Face Swapping and Editing Software
repo_topics: ai computer-vision deepfake face-editor faceswap live-portrait video-editor
code: from itertools import product as product
  from typing import TYPE_CHECKING
  import pickle
  
  import torch
  import cv2
  import numpy as np
  from torchvision.transforms import v2
  
  if TYPE_CHECKING:
      from app.processors.models_processor import ModelsProcessor
  from app.processors.models_data import models_dir
  from app.processors.utils import faceutil
  
  class FaceLandmarkDetectors:
      def __init__(self, models_processor: 'ModelsProcessor'):
          self.models_processor = models_processor
  
      def run_detect_landmark(self, img, bbox, det_kpss, detect_mode='203', score=0.5, from_points=False):
          kpss_5 = []
          kpss = []
          scores = []
  
          if detect_mode=='5':
              if not self.models_processor.models['FaceLandmark5']:
                  self.models_processor.models['FaceLandmark5'] = self.models_processor.load_model('FaceLandmark5')
  
                  feature_maps = [[64, 64], [32, 32], [16, 16]]
                  min_sizes = [[16, 32], [64, 128], [256, 512]]
                  steps = [8, 16, 32]
                  image_size = 512
                  # re-initialize self.models_processor.anchors due to clear_mem function
                  self.models_processor.anchors  = []
  
                  for k, f in enumerate(feature_maps):
                      min_size_array = min_sizes[k]
                      for i, j in product(range(f[0]), range(f[1])):
                          for min_size in min_size_array:
                              s_kx = min_size / image_size
                              s_ky = min_size / image_size
                              dense_cx = [x * steps[k] / image_size for x in [j + 0.5]]
                              dense_cy = [y * steps[k] / image_size for y in [i + 0.5]]
                              for cy, cx in product(dense_cy, dense_cx):
                                  self.models_processor.anchors += [cx, cy, s_kx, s_ky]
  
              kpss_5, kpss, scores = self.detect_face_landmark_5(img, bbox=bbox, det_kpss=det_kpss, from_points=from_points)
  
          elif detect_mode=='68':
              if not self.models_processor.models['FaceLandmark68']:
                  self.models_processor.models['FaceLandmark68'] = self.models_processor.load_model('FaceLandmark68')
  
              kpss_5, kpss, scores = self.detect_face_landmark_68(img, bbox=bbox, det_kpss=det_kpss, from_points=from_points)
  
          elif detect_mode=='3d68':
              if not self.models_processor.models['FaceLandmark3d68']:
                  self.models_processor.models['FaceLandmark3d68'] = self.models_processor.load_model('FaceLandmark3d68')
                  with open(f'{models_dir}/meanshape_68.pkl', 'rb') as f:
                      self.models_processor.mean_lmk = pickle.load(f)
  
              kpss_5, kpss, scores = self.detect_face_landmark_3d68(img, bbox=bbox, det_kpss=det_kpss, from_points=from_points)
  
              return kpss_5, kpss, scores
  
          elif detect_mode=='98':
              if not self.models_processor.models['FaceLandmark98']:
                  self.models_processor.models['FaceLandmark98'] = self.models_processor.load_model('FaceLandmark98')
  
              kpss_5, kpss, scores = self.detect_face_landmark_98(img, bbox=bbox, det_kpss=det_kpss, from_points=from_points)
  
          elif detect_mode=='106':
              if not self.models_processor.models['FaceLandmark106']:
                  self.models_processor.models['FaceLandmark106'] = self.models_processor.load_model('FaceLandmark106')
  
              kpss_5, kpss, scores = self.detect_face_landmark_106(img, bbox=bbox, det_kpss=det_kpss, from_points=from_points)
  
              return kpss_5, kpss, scores
  
          elif detect_mode=='203':
              if not self.models_processor.models['FaceLandmark203']:
                  self.models_processor.models['FaceLandmark203'] = self.models_processor.load_model('FaceLandmark203')
  
              kpss_5, kpss, scores = self.detect_face_landmark_203(img, bbox=bbox, det_kpss=det_kpss, from_points=from_points)
  
              return kpss_5, kpss, scores
  
          elif detect_mode=='478':
              if not self.models_processor.models['FaceLandmark478']:
                  self.models_processor.models['FaceLandmark478'] = self.models_processor.load_model('FaceLandmark478')
  
              if not self.models_processor.models['FaceBlendShapes']:
                  self.models_processor.models['FaceBlendShapes'] = self.models_processor.load_model('FaceBlendShapes')
  
              kpss_5, kpss, scores = self.detect_face_landmark_478(img, bbox=bbox, det_kpss=det_kpss, from_points=from_points)
  
              return kpss_5, kpss, scores
  
          if len(kpss_5) > 0:
              if len(scores) > 0:
                  if np.mean(scores) >= score:
                      return kpss_5, kpss, scores
              else:
                  return kpss_5, kpss, scores
  
          return [], [], []
  
  
      def detect_face_landmark_5(self, img, bbox, det_kpss, from_points=False):
          if from_points == False:
              w, h = (bbox[2] - bbox[0]), (bbox[3] - bbox[1])
              center = (bbox[2] + bbox[0]) / 2, (bbox[3] + bbox[1]) / 2
              rotate = 0
              _scale = 512.0  / (max(w, h)*1.5)
              image, M = faceutil.transform(img, center, 512, _scale, rotate)
          else:
              image, M = faceutil.warp_face_by_face_landmark_5(img, det_kpss, 512, mode='arcface128', interpolation=v2.InterpolationMode.BILINEAR)
  
          image = image.permute(1,2,0)
  
          mean = torch.tensor([104, 117, 123], dtype=torch.float32, device=self.models_processor.device)
          image = torch.sub(image, mean)
  
          image = image.permute(2,0,1)
          image = torch.reshape(image, (1, 3, 512, 512))
  
          height, width = (512, 512)
          tmp = [width, height, width, height, width, height, width, height, width, height]
          scale1 = torch.tensor(tmp, dtype=torch.float32, device=self.models_processor.device)
  
          conf = torch.empty((1,10752,2), dtype=torch.float32, device=self.models_processor.device).contiguous()
          landmarks = torch.empty((1,10752,10), dtype=torch.float32, device=self.models_processor.device).contiguous()
  
          io_binding = self.models_processor.models['FaceLandmark5'].io_binding()
          io_binding.bind_input(name='input', device_type=self.models_processor.device, device_id=0, element_type=np.float32, shape=(1,3,512,512), buffer_ptr=image.data_ptr())
          io_binding.bind_output(name='conf', device_type=self.models_processor.device, device_id=0, element_type=np.float32, shape=(1,10752,2), buffer_ptr=conf.data_ptr())
          io_binding.bind_output(name='landmarks', device_type=self.models_processor.device, device_id=0, element_type=np.float32, shape=(1,10752,10), buffer_ptr=landmarks.data_ptr())
  
          if self.models_processor.device == "cuda":
              torch.cuda.synchronize()
          elif self.models_processor.device != "cpu":
              self.models_processor.syncvec.cpu()
          self.models_processor.models['FaceLandmark5'].run_with_iobinding(io_binding)
  
          scores = torch.squeeze(conf)[:, 1]
          priors = torch.tensor(self.models_processor.anchors).view(-1, 4)
          priors = priors.to(self.models_processor.device)
  
          pre = torch.squeeze(landmarks, 0)
  
          tmp = (priors[:, :2] + pre[:, :2] * 0.1 * priors[:, 2:], priors[:, :2] + pre[:, 2:4] * 0.1 * priors[:, 2:], priors[:, :2] + pre[:, 4:6] * 0.1 * priors[:, 2:], priors[:, :2] + pre[:, 6:8] * 0.1 * priors[:, 2:], priors[:, :2] + pre[:, 8:10] * 0.1 * priors[:, 2:])
          landmarks = torch.cat(tmp, dim=1)
          landmarks = torch.mul(landmarks, scale1)
  
          landmarks = landmarks.cpu().numpy()
  
          # ignore low scores
          score=.1
          inds = torch.where(scores>score)[0]
          inds = inds.cpu().numpy()
          scores = scores.cpu().numpy()
  
          landmarks, scores = landmarks[inds], scores[inds]
  
          # sort
          order = scores.argsort()[::-1]
  
          if len(order) > 0:
              landmarks = landmarks[order][0]
              scores = scores[order][0]
  
              landmarks = np.array([[landmarks[i], landmarks[i + 1]] for i in range(0,10,2)])
  
              IM = faceutil.invertAffineTransform(M)
              landmarks = faceutil.trans_points2d(landmarks, IM)
              scores = np.array([scores])
  
              return landmarks, landmarks, scores
  
          return [], [], []
  
      def detect_face_landmark_68(self, img, bbox, det_kpss, from_points=False):
          if from_points == False:
              crop_image, affine_matrix = faceutil.warp_face_by_bounding_box_for_landmark_68(img, bbox, (256, 256))
          else:
              crop_image, affine_matrix = faceutil.warp_face_by_face_landmark_5(img, det_kpss, 256, mode='arcface128', interpolation=v2.InterpolationMode.BILINEAR)
  
          crop_image = crop_image.to(dtype=torch.float32)
          crop_image = torch.div(crop_image, 255.0)
          crop_image = torch.unsqueeze(crop_image, 0).contiguous()
  
          io_binding = self.models_processor.models['FaceLandmark68'].io_binding()
          io_binding.bind_input(name='input', device_type=self.models_processor.device, device_id=0, element_type=np.float32,  shape=crop_image.size(), buffer_ptr=crop_image.data_ptr())
  
          io_binding.bind_output('landmarks_xyscore', self.models_processor.device)
          io_binding.bind_output('heatmaps', self.models_processor.device)
  
          # Sync and run model
          if self.models_processor.device == "cuda":
              torch.cuda.synchronize()
          elif self.models_processor.device != "cpu":
              self.models_processor.syncvec.cpu()
          self.models_processor.models['FaceLandmark68'].run_with_iobinding(io_binding)
          net_outs = io_binding.copy_outputs_to_cpu()
          face_landmark_68 = net_outs[0]
          face_heatmap = net_outs[1]
  
          face_landmark_68 = face_landmark_68[:, :, :2][0] / 64.0
          face_landmark_68 = face_landmark_68.reshape(1, -1, 2) * 256.0
          face_landmark_68 = cv2.transform(face_landmark_68, cv2.invertAffineTransform(affine_matrix))
  
          face_landmark_68 = face_landmark_68.reshape(-1, 2)
          face_landmark_68_score = np.amax(face_heatmap, axis = (2, 3))
          face_landmark_68_score = face_landmark_68_score.reshape(-1, 1)
  
          face_landmark_68_5, face_landmark_68_score = faceutil.convert_face_landmark_68_to_5(face_landmark_68, face_landmark_68_score)
  
          return face_landmark_68_5, face_landmark_68, face_landmark_68_score
  
      def detect_face_landmark_3d68(self, img, bbox, det_kpss, from_points=False):
          if from_points == False:
              w, h = (bbox[2] - bbox[0]), (bbox[3] - bbox[1])
              center = (bbox[2] + bbox[0]) / 2, (bbox[3] + bbox[1]) / 2
              rotate = 0
              _scale = 192  / (max(w, h)*1.5)
              aimg, M = faceutil.transform(img, center, 192, _scale, rotate)
          else:
              aimg, M = faceutil.warp_face_by_face_landmark_5(img, det_kpss, image_size=192, mode='arcface128', interpolation=v2.InterpolationMode.BILINEAR)
  
          aimg = torch.unsqueeze(aimg, 0).contiguous()
          aimg = aimg.to(dtype=torch.float32)
          aimg = self.models_processor.normalize(aimg)
          io_binding = self.models_processor.models['FaceLandmark3d68'].io_binding()
          io_binding.bind_input(name='data', device_type=self.models_processor.device, device_id=0, element_type=np.float32,  shape=aimg.size(), buffer_ptr=aimg.data_ptr())
  
          io_binding.bind_output('fc1', self.models_processor.device)
  
          # Sync and run model
          if self.models_processor.device == "cuda":
              torch.cuda.synchronize()
          elif self.models_processor.device != "cpu":
              self.models_processor.syncvec.cpu()
          self.models_processor.models['FaceLandmark3d68'].run_with_iobinding(io_binding)
          pred = io_binding.copy_outputs_to_cpu()[0][0]
  
          if pred.shape[0] >= 3000:
              pred = pred.reshape((-1, 3))
          else:
              pred = pred.reshape((-1, 2))
          if 68 < pred.shape[0]:
              pred = pred[68*-1:,:]
          pred[:, 0:2] += 1
          pred[:, 0:2] *= (192 // 2)
          if pred.shape[1] == 3:
              pred[:, 2] *= (192 // 2)
  
          IM = faceutil.invertAffineTransform(M)
          pred = faceutil.trans_points3d(pred, IM)
  
          # at moment we don't use 3d points
          
          #'''
          #P = faceutil.estimate_affine_matrix_3d23d(self.models_processor.mean_lmk, pred)
          #s, R, t = faceutil.P2sRt(P)
          #rx, ry, rz = faceutil.matrix2angle(R)
          #pose = np.array( [rx, ry, rz], dtype=np.float32 ) #pitch, yaw, roll
          #'''
  
          # convert from 3d68 to 2d68 keypoints
          landmark2d68 = np.array(pred[:, [0, 1]])
  
          # convert from 68 to 5 keypoints
          landmark2d68_5, _ = faceutil.convert_face_landmark_68_to_5(landmark2d68, [])
  
          return landmark2d68_5, landmark2d68, []
  
      def detect_face_landmark_98(self, img, bbox, det_kpss, from_points=False):
          if from_points == False:
              crop_image, detail = faceutil.warp_face_by_bounding_box_for_landmark_98(img, bbox, (256, 256))
          else:
              crop_image, M = faceutil.warp_face_by_face_landmark_5(img, det_kpss, image_size=256, mode='arcface128', interpolation=v2.InterpolationMode.BILINEAR)
              h, w = (crop_image.size(dim=1), crop_image.size(dim=2))
  
          landmark = []
          landmark_5 = []
          landmark_score = []
          if crop_image is not None:
              crop_image = crop_image.to(dtype=torch.float32)
              crop_image = torch.div(crop_image, 255.0)
              crop_image = torch.unsqueeze(crop_image, 0).contiguous()
  
              io_binding = self.models_processor.models['FaceLandmark98'].io_binding()
              io_binding.bind_input(name='input', device_type=self.models_processor.device, device_id=0, element_type=np.float32,  shape=crop_image.size(), buffer_ptr=crop_image.data_ptr())
  
              io_binding.bind_output('landmarks_xyscore', self.models_processor.device)
  
              # Sync and run model
              if self.models_processor.device == "cuda":
                  torch.cuda.synchronize()
              elif self.models_processor.device != "cpu":
                  self.models_processor.syncvec.cpu()
              self.models_processor.models['FaceLandmark98'].run_with_iobinding(io_binding)
              landmarks_xyscore = io_binding.copy_outputs_to_cpu()[0]
  
              if len(landmarks_xyscore) > 0:
                  for one_face_landmarks in landmarks_xyscore:
                      landmark_score = one_face_landmarks[:, [2]].reshape(-1)
                      landmark = one_face_landmarks[:, [0, 1]].reshape(-1,2)
  
                      ##recover, and grouped as [98,2]
                      if from_points == False:
                          landmark[:, 0] = landmark[:, 0] * detail[1] + detail[3] - detail[4]
                          landmark[:, 1] = landmark[:, 1] * detail[0] + detail[2] - detail[4]
                      else:
                          landmark[:, 0] = landmark[:, 0] * w
                          landmark[:, 1] = landmark[:, 1] * h
  
                          IM = faceutil.invertAffineTransform(M)
                          landmark = faceutil.trans_points2d(landmark, IM)
  
                      landmark_5, landmark_score = faceutil.convert_face_landmark_98_to_5(landmark, landmark_score)
  
          return landmark_5, landmark, landmark_score
  
      def detect_face_landmark_106(self, img, bbox, det_kpss, from_points=False):
          if from_points == False:
              w, h = (bbox[2] - bbox[0]), (bbox[3] - bbox[1])
              center = (bbox[2] + bbox[0]) / 2, (bbox[3] + bbox[1]) / 2
              rotate = 0
              _scale = 192  / (max(w, h)*1.5)
              #print('param:', img.size(), bbox, center, (192, 192), _scale, rotate)
              aimg, M = faceutil.transform(img, center, 192, _scale, rotate)
          else:
              aimg, M = faceutil.warp_face_by_face_landmark_5(img, det_kpss, image_size=192, mode='arcface128', interpolation=v2.InterpolationMode.BILINEAR)
  
          aimg = torch.unsqueeze(aimg, 0).contiguous()
          aimg = aimg.to(dtype=torch.float32)
          aimg = self.models_processor.normalize(aimg)
          io_binding = self.models_processor.models['FaceLandmark106'].io_binding()
          io_binding.bind_input(name='data', device_type=self.models_processor.device, device_id=0, element_type=np.float32,  shape=aimg.size(), buffer_ptr=aimg.data_ptr())
  
          io_binding.bind_output('fc1', self.models_processor.device)
  
          # Sync and run model
          if self.models_processor.device == "cuda":
              torch.cuda.synchronize()
          elif self.models_processor.device != "cpu":
              self.models_processor.syncvec.cpu()
          self.models_processor.models['FaceLandmark106'].run_with_iobinding(io_binding)
          pred = io_binding.copy_outputs_to_cpu()[0][0]
  
          if pred.shape[0] >= 3000:
              pred = pred.reshape((-1, 3))
          else:
              pred = pred.reshape((-1, 2))
  
          if 106 < pred.shape[0]:
              pred = pred[106*-1:,:]
  
          pred[:, 0:2] += 1
          pred[:, 0:2] *= (192 // 2)
          if pred.shape[1] == 3:
              pred[:, 2] *= (192 // 2)
  
          IM = faceutil.invertAffineTransform(M)
          pred = faceutil.trans_points(pred, IM)
  
          pred_5 = []
          if pred is not None:
              # convert from 106 to 5 keypoints
              pred_5 = faceutil.convert_face_landmark_106_to_5(pred)
  
          return pred_5, pred, []
  
      def detect_face_landmark_203(self, img, bbox, det_kpss, from_points=False):
          IM = None
          if from_points == False:
              w, h = (bbox[2] - bbox[0]), (bbox[3] - bbox[1])
              center = (bbox[2] + bbox[0]) / 2, (bbox[3] + bbox[1]) / 2
              rotate = 0
              _scale = 224  / (max(w, h)*1.5)
  
              aimg, M = faceutil.transform(img, center, 224, _scale, rotate)
          elif len(det_kpss) == 0:
              return [], [], []
          else:
              if det_kpss.shape[0] == 5:
                  aimg, M = faceutil.warp_face_by_face_landmark_5(img, det_kpss, image_size=224, mode='arcface128', interpolation=v2.InterpolationMode.BILINEAR)
              else:
                  aimg, M, IM = faceutil.warp_face_by_face_landmark_x(img, det_kpss, dsize=224, scale=1.5, vy_ratio=-0.1, interpolation=v2.InterpolationMode.BILINEAR)
  
          aimg = torch.unsqueeze(aimg, 0).contiguous()
          aimg = aimg.to(dtype=torch.float32)
          aimg = torch.div(aimg, 255.0)
          io_binding = self.models_processor.models['FaceLandmark203'].io_binding()
          io_binding.bind_input(name='input', device_type=self.models_processor.device, device_id=0, element_type=np.float32,  shape=aimg.size(), buffer_ptr=aimg.data_ptr())
  
          io_binding.bind_output('output', self.models_processor.device)
          io_binding.bind_output('853', self.models_processor.device)
          io_binding.bind_output('856', self.models_processor.device)
  
          # Sync and run model
          if self.models_processor.device == "cuda":
              torch.cuda.synchronize()
          elif self.models_processor.device != "cpu":
              self.models_processor.syncvec.cpu()
          self.models_processor.models['FaceLandmark203'].run_with_iobinding(io_binding)
          out_lst = io_binding.copy_outputs_to_cpu()
          out_pts = out_lst[2]
  
          out_pts = out_pts.reshape((-1, 2)) * 224.0
  
          if len(det_kpss) == 0 or det_kpss.shape[0] == 5:
              IM = faceutil.invertAffineTransform(M)
  
          out_pts = faceutil.trans_points(out_pts, IM)
  
          out_pts_5 = []
          if out_pts is not None:
              # convert from 203 to 5 keypoints
              out_pts_5 = faceutil.convert_face_landmark_203_to_5(out_pts)
  
          return out_pts_5, out_pts, []
  
      def detect_face_landmark_478(self, img, bbox, det_kpss, from_points=False):
          if from_points == False:
              w, h = (bbox[2] - bbox[0]), (bbox[3] - bbox[1])
              center = (bbox[2] + bbox[0]) / 2, (bbox[3] + bbox[1]) / 2
              rotate = 0
              _scale = 256.0  / (max(w, h)*1.5)
              #print('param:', img.size(), bbox, center, (192, 192), _scale, rotate)
              aimg, M = faceutil.transform(img, center, 256, _scale, rotate)
          else:
              aimg, M = faceutil.warp_face_by_face_landmark_5(img, det_kpss, 256, mode='arcfacemap', interpolation=v2.InterpolationMode.BILINEAR)
  
          aimg = torch.unsqueeze(aimg, 0).contiguous()
          aimg = aimg.to(dtype=torch.float32)
          aimg = torch.div(aimg, 255.0)
          io_binding = self.models_processor.models['FaceLandmark478'].io_binding()
          io_binding.bind_input(name='input_12', device_type=self.models_processor.device, device_id=0, element_type=np.float32,  shape=aimg.size(), buffer_ptr=aimg.data_ptr())
  
          io_binding.bind_output('Identity', self.models_processor.device)
          io_binding.bind_output('Identity_1', self.models_processor.device)
          io_binding.bind_output('Identity_2', self.models_processor.device)
  
          # Sync and run model
          if self.models_processor.device == "cuda":
              torch.cuda.synchronize()
          elif self.models_processor.device != "cpu":
              self.models_processor.syncvec.cpu()
          self.models_processor.models['FaceLandmark478'].run_with_iobinding(io_binding)
          landmarks, faceflag, blendshapes = io_binding.copy_outputs_to_cpu() # pylint: disable=unused-variable
          landmarks = landmarks.reshape( (1,478,3))
  
          landmark = []
          landmark_5 = []
          landmark_score = [] # pylint: disable=unused-variable
          if len(landmarks) > 0:
              for one_face_landmarks in landmarks:
                  landmark = one_face_landmarks
                  IM = faceutil.invertAffineTransform(M)
                  landmark = faceutil.trans_points3d(landmark, IM)
  
                  #'''
                  #P = faceutil.estimate_affine_matrix_3d23d(self.models_processor.mean_lmk, landmark)
                  #s, R, t = faceutil.P2sRt(P)
                  #rx, ry, rz = faceutil.matrix2angle(R)
                  #pose = np.array( [rx, ry, rz], dtype=np.float32 ) #pitch, yaw, roll
                  #'''
                  landmark = landmark[:, [0, 1]].reshape(-1,2)
  
                  #get scores
                  landmark_for_score = landmark[self.models_processor.LandmarksSubsetIdxs]
                  landmark_for_score = landmark_for_score[:, :2]
                  landmark_for_score = np.expand_dims(landmark_for_score, axis=0)
                  landmark_for_score = landmark_for_score.astype(np.float32)
                  landmark_for_score = torch.from_numpy(landmark_for_score).to(self.models_processor.device)
  
                  io_binding_bs = self.models_processor.models['FaceBlendShapes'].io_binding()
                  io_binding_bs.bind_input(name='input_points', device_type=self.models_processor.device, device_id=0, element_type=np.float32,  shape=tuple(landmark_for_score.shape), buffer_ptr=landmark_for_score.data_ptr())
                  io_binding_bs.bind_output('output', self.models_processor.device)
  
                  # Sync and run model
                  if self.models_processor.device == "cuda":
                      torch.cuda.synchronize()
                  elif self.models_processor.device != "cpu":
                      self.models_processor.syncvec.cpu()
                  self.models_processor.models['FaceBlendShapes'].run_with_iobinding(io_binding_bs)
                  landmark_score = io_binding_bs.copy_outputs_to_cpu()[0] # pylint: disable=unused-variable
  
                  # convert from 478 to 5 keypoints
                  landmark_5 = faceutil.convert_face_landmark_478_to_5(landmark)
  
          #return landmark, landmark_score
          return landmark_5, landmark, []