79 lines
2.2 KiB
Python
79 lines
2.2 KiB
Python
|
from __future__ import print_function
|
||
|
|
import os
|
||
|
|
import torch
|
||
|
|
from torch.utils.model_zoo import load_url
|
||
|
|
from enum import Enum
|
||
|
|
import numpy as np
|
||
|
|
import cv2
|
||
|
|
try:
|
||
|
|
import urllib.request as request_file
|
||
|
|
except BaseException:
|
||
|
|
import urllib as request_file
|
||
|
|
|
||
|
|
from .models import FAN, ResNetDepth
|
||
|
|
from .utils import *
|
||
|
|
|
||
|
|
|
||
|
|
class LandmarksType(Enum):
|
||
|
|
"""Enum class defining the type of landmarks to detect.
|
||
|
|
|
||
|
|
``_2D`` - the detected points ``(x,y)`` are detected in a 2D space and follow the visible contour of the face
|
||
|
|
``_2halfD`` - this points represent the projection of the 3D points into 3D
|
||
|
|
``_3D`` - detect the points ``(x,y,z)``` in a 3D space
|
||
|
|
|
||
|
|
"""
|
||
|
|
_2D = 1
|
||
|
|
_2halfD = 2
|
||
|
|
_3D = 3
|
||
|
|
|
||
|
|
|
||
|
|
class NetworkSize(Enum):
|
||
|
|
# TINY = 1
|
||
|
|
# SMALL = 2
|
||
|
|
# MEDIUM = 3
|
||
|
|
LARGE = 4
|
||
|
|
|
||
|
|
def __new__(cls, value):
|
||
|
|
member = object.__new__(cls)
|
||
|
|
member._value_ = value
|
||
|
|
return member
|
||
|
|
|
||
|
|
def __int__(self):
|
||
|
|
return self.value
|
||
|
|
|
||
|
|
ROOT = os.path.dirname(os.path.abspath(__file__))
|
||
|
|
|
||
|
|
class FaceAlignment:
|
||
|
|
def __init__(self, landmarks_type, network_size=NetworkSize.LARGE,
|
||
|
|
device='cuda', flip_input=False, face_detector='sfd', verbose=False):
|
||
|
|
self.device = device
|
||
|
|
self.flip_input = flip_input
|
||
|
|
self.landmarks_type = landmarks_type
|
||
|
|
self.verbose = verbose
|
||
|
|
|
||
|
|
network_size = int(network_size)
|
||
|
|
|
||
|
|
if 'cuda' in device:
|
||
|
|
torch.backends.cudnn.benchmark = True
|
||
|
|
|
||
|
|
# Get the face detector
|
||
|
|
face_detector_module = __import__('face_detection.detection.' + face_detector,
|
||
|
|
globals(), locals(), [face_detector], 0)
|
||
|
|
self.face_detector = face_detector_module.FaceDetector(device=device, verbose=verbose)
|
||
|
|
|
||
|
|
def get_detections_for_batch(self, images):
|
||
|
|
images = images[..., ::-1]
|
||
|
|
detected_faces = self.face_detector.detect_from_batch(images.copy())
|
||
|
|
results = []
|
||
|
|
|
||
|
|
for i, d in enumerate(detected_faces):
|
||
|
|
if len(d) == 0:
|
||
|
|
results.append(None)
|
||
|
|
continue
|
||
|
|
d = d[0]
|
||
|
|
d = np.clip(d, 0, None)
|
||
|
|
|
||
|
|
x1, y1, x2, y2 = map(int, d[:-1])
|
||
|
|
results.append((x1, y1, x2, y2))
|
||
|
|
|
||
|
|
return results
|