MCP-1st-Birthday/Free-View_Expressive_Talking_Head_Video_Editing

1from __future__ import print_function2import os3import torch4from torch.utils.model_zoo import load_url5from enum import Enum6import numpy as np7import cv28try:9    import urllib.request as request_file10except BaseException:11    import urllib as request_file12 13from .models import FAN, ResNetDepth14from .utils import *15 16 17class LandmarksType(Enum):18    """Enum class defining the type of landmarks to detect.19 20    ``_2D`` - the detected points ``(x,y)`` are detected in a 2D space and follow the visible contour of the face21    ``_2halfD`` - this points represent the projection of the 3D points into 3D22    ``_3D`` - detect the points ``(x,y,z)``` in a 3D space23 24    """25    _2D = 126    _2halfD = 227    _3D = 328 29 30class NetworkSize(Enum):31    # TINY = 132    # SMALL = 233    # MEDIUM = 334    LARGE = 435 36    def __new__(cls, value):37        member = object.__new__(cls)38        member._value_ = value39        return member40 41    def __int__(self):42        return self.value43 44ROOT = os.path.dirname(os.path.abspath(__file__))45 46class FaceAlignment:47    def __init__(self, landmarks_type, network_size=NetworkSize.LARGE,48                 device='cuda', flip_input=False, face_detector='sfd', verbose=False):49        self.device = device50        self.flip_input = flip_input51        self.landmarks_type = landmarks_type52        self.verbose = verbose53 54        network_size = int(network_size)55 56        if 'cuda' in device:57            torch.backends.cudnn.benchmark = True58 59        # Get the face detector60        face_detector_module = __import__('face_detection.detection.' + face_detector,61                                          globals(), locals(), [face_detector], 0)62        self.face_detector = face_detector_module.FaceDetector(device=device, verbose=verbose)63 64    def get_detections_for_batch(self, images):65        images = images[..., ::-1]66        detected_faces = self.face_detector.detect_from_batch(images.copy())67        results = []68 69        for i, d in enumerate(detected_faces):70            if len(d) == 0:71                results.append(None)72                continue73            d = d[0]74            d = np.clip(d, 0, None)75            76            x1, y1, x2, y2 = map(int, d[:-1])77            results.append((x1, y1, x2, y2))78 79        return results