Add face_alignment as alternative detector

face_alignment/LICENSE

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2017, Adrian Bulat
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

face_alignment/__init__.py

@@ -0,0 +1,7 @@
+# -*- coding: utf-8 -*-
+
+__author__ = """Adrian Bulat"""
+__email__ = 'adrian@adrianbulat.com'
+__version__ = '1.4.1'
+
+from .api import FaceAlignment, LandmarksType, NetworkSize

face_alignment/api.py

@@ -0,0 +1,305 @@
+import torch
+import warnings
+from enum import IntEnum
+from skimage import io
+import numpy as np
+from packaging import version
+from tqdm import tqdm
+import importlib
+
+from .utils import *
+from .folder_data import FolderData
+
+
+class LandmarksType(IntEnum):
+    """Enum class defining the type of landmarks to detect.
+
+    ``TWO_D`` - the detected points ``(x,y)`` are detected in a 2D space and follow the visible contour of the face
+    ``TWO_HALF_D`` - this points represent the projection of the 3D points into 3D
+    ``THREE_D`` - detect the points ``(x,y,z)``` in a 3D space
+
+    """
+    TWO_D = 1
+    TWO_HALF_D = 2
+    THREE_D = 3
+
+
+class NetworkSize(IntEnum):
+    # TINY = 1
+    # SMALL = 2
+    # MEDIUM = 3
+    LARGE = 4
+
+
+default_model_urls = {
+    '2DFAN-4': 'https://www.adrianbulat.com/downloads/python-fan/2DFAN4-cd938726ad.zip',
+    '3DFAN-4': 'https://www.adrianbulat.com/downloads/python-fan/3DFAN4-4a694010b9.zip',
+    'depth': 'https://www.adrianbulat.com/downloads/python-fan/depth-6c4283c0e0.zip',
+}
+
+models_urls = {
+    '1.6': {
+        '2DFAN-4': 'https://www.adrianbulat.com/downloads/python-fan/2DFAN4_1.6-c827573f02.zip',
+        '3DFAN-4': 'https://www.adrianbulat.com/downloads/python-fan/3DFAN4_1.6-ec5cf40a1d.zip',
+        'depth': 'https://www.adrianbulat.com/downloads/python-fan/depth_1.6-2aa3f18772.zip',
+    },
+    '1.5': {
+        '2DFAN-4': 'https://www.adrianbulat.com/downloads/python-fan/2DFAN4_1.5-a60332318a.zip',
+        '3DFAN-4': 'https://www.adrianbulat.com/downloads/python-fan/3DFAN4_1.5-176570af4d.zip',
+        'depth': 'https://www.adrianbulat.com/downloads/python-fan/depth_1.5-bc10f98e39.zip',
+    },
+}
+
+
+class FaceAlignment:
+    def __init__(self, landmarks_type, network_size=NetworkSize.LARGE,
+                 device='cuda', dtype=torch.float32, flip_input=False, face_detector='sfd', face_detector_kwargs=None, verbose=False):
+        self.device = device
+        self.flip_input = flip_input
+        self.landmarks_type = landmarks_type
+        self.verbose = verbose
+        self.dtype = dtype
+
+        if version.parse(torch.__version__) < version.parse('1.5.0'):
+            raise ImportError(f'Unsupported pytorch version detected. Minimum supported version of pytorch: 1.5.0\
+                            Either upgrade (recommended) your pytorch setup, or downgrade to face-alignment 1.2.0')
+
+        network_size = int(network_size)
+        pytorch_version = torch.__version__
+        if 'dev' in pytorch_version:
+            pytorch_version = pytorch_version.rsplit('.', 2)[0]
+        else:
+            pytorch_version = pytorch_version.rsplit('.', 1)[0]
+
+        #if 'cuda' in device:
+        #    torch.backends.cudnn.benchmark = True
+
+        # Get the face detector
+        package_directory_name = os.path.basename(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+        face_detector_module = importlib.import_module('.face_alignment.detection.' + face_detector, package=package_directory_name)
+
+
+        face_detector_kwargs = face_detector_kwargs or {}
+        self.face_detector = face_detector_module.FaceDetector(device=device, verbose=verbose, **face_detector_kwargs)
+
+        # Initialise the face alignemnt networks
+        if landmarks_type == LandmarksType.TWO_D:
+            network_name = '2DFAN-' + str(network_size)
+        else:
+            network_name = '3DFAN-' + str(network_size)
+        self.face_alignment_net = torch.jit.load(
+            load_file_from_url(models_urls.get(pytorch_version, default_model_urls)[network_name]))
+
+        self.face_alignment_net.to(device, dtype=dtype)
+        self.face_alignment_net.eval()
+
+        # Initialiase the depth prediciton network
+        if landmarks_type == LandmarksType.THREE_D:
+            self.depth_prediciton_net = torch.jit.load(
+                load_file_from_url(models_urls.get(pytorch_version, default_model_urls)['depth']))
+
+            self.depth_prediciton_net.to(device, dtype=dtype)
+            self.depth_prediciton_net.eval()
+
+    def get_landmarks(self, image_or_path, detected_faces=None, return_bboxes=False, return_landmark_score=False):
+        """Deprecated, please use get_landmarks_from_image
+
+        Arguments:
+            image_or_path {string or numpy.array or torch.tensor} -- The input image or path to it
+
+        Keyword Arguments:
+            detected_faces {list of numpy.array} -- list of bounding boxes, one for each face found
+            in the image (default: {None})
+            return_bboxes {boolean} -- If True, return the face bounding boxes in addition to the keypoints.
+            return_landmark_score {boolean} -- If True, return the keypoint scores along with the keypoints.
+        """
+        return self.get_landmarks_from_image(image_or_path, detected_faces, return_bboxes, return_landmark_score)
+
+    @torch.no_grad()
+    def get_landmarks_from_image(self, image_or_path, detected_faces=None, return_bboxes=False,
+                                 return_landmark_score=False):
+        """Predict the landmarks for each face present in the image.
+
+        This function predicts a set of 68 2D or 3D images, one for each image present.
+        If detect_faces is None the method will also run a face detector.
+
+         Arguments:
+            image_or_path {string or numpy.array or torch.tensor} -- The input image or path to it.
+
+        Keyword Arguments:
+            detected_faces {list of numpy.array} -- list of bounding boxes, one for each face found
+            in the image (default: {None})
+            return_bboxes {boolean} -- If True, return the face bounding boxes in addition to the keypoints.
+            return_landmark_score {boolean} -- If True, return the keypoint scores along with the keypoints.
+
+        Return:
+            result:
+                1. if both return_bboxes and return_landmark_score are False, result will be:
+                    landmark
+                2. Otherwise, result will be one of the following, depending on the actual value of return_* arguments.
+                    (landmark, landmark_score, detected_face)
+                    (landmark, None,           detected_face)
+                    (landmark, landmark_score, None         )
+        """
+        image = get_image(image_or_path)
+
+        if detected_faces is None:
+            detected_faces = self.face_detector.detect_from_image(image.copy())
+
+        if len(detected_faces) == 0:
+            warnings.warn("No faces were detected.")
+            if return_bboxes or return_landmark_score:
+                return None, None, None
+            else:
+                return None
+
+        landmarks = []
+        landmarks_scores = []
+        for i, d in enumerate(detected_faces):
+            center = torch.tensor(
+                [d[2] - (d[2] - d[0]) / 2.0, d[3] - (d[3] - d[1]) / 2.0])
+            center[1] = center[1] - (d[3] - d[1]) * 0.12
+            scale = (d[2] - d[0] + d[3] - d[1]) / self.face_detector.reference_scale
+
+            inp = crop(image, center, scale)
+            inp = torch.from_numpy(inp.transpose(
+                (2, 0, 1))).float()
+
+            inp = inp.to(self.device, dtype=self.dtype)
+            inp.div_(255.0).unsqueeze_(0)
+
+            out = self.face_alignment_net(inp).detach()
+            if self.flip_input:
+                out += flip(self.face_alignment_net(flip(inp)).detach(), is_label=True)
+            out = out.to(device='cpu', dtype=torch.float32).numpy()
+
+            pts, pts_img, scores = get_preds_fromhm(out, center.numpy(), scale)
+            pts, pts_img = torch.from_numpy(pts), torch.from_numpy(pts_img)
+            pts, pts_img = pts.view(68, 2) * 4, pts_img.view(68, 2)
+            scores = scores.squeeze(0)
+
+            if self.landmarks_type == LandmarksType.THREE_D:
+                heatmaps = np.zeros((68, 256, 256), dtype=np.float32)
+                for i in range(68):
+                    if pts[i, 0] > 0 and pts[i, 1] > 0:
+                        heatmaps[i] = draw_gaussian(
+                            heatmaps[i], pts[i], 2)
+                heatmaps = torch.from_numpy(
+                    heatmaps).unsqueeze_(0)
+
+                heatmaps = heatmaps.to(self.device, dtype=self.dtype)
+                depth_pred = self.depth_prediciton_net(
+                    torch.cat((inp, heatmaps), 1)).data.cpu().view(68, 1).to(dtype=torch.float32)
+                pts_img = torch.cat(
+                    (pts_img, depth_pred * (1.0 / (256.0 / (200.0 * scale)))), 1)
+
+            landmarks.append(pts_img.numpy())
+            landmarks_scores.append(scores)
+
+        if not return_bboxes:
+            detected_faces = None
+        if not return_landmark_score:
+            landmarks_scores = None
+        if return_bboxes or return_landmark_score:
+            return landmarks, landmarks_scores, detected_faces
+        else:
+            return landmarks
+
+    @torch.no_grad()
+    def get_landmarks_from_batch(self, image_batch, detected_faces=None, return_bboxes=False,
+                                 return_landmark_score=False):
+        """Predict the landmarks for each face present in the image.
+
+        This function predicts a set of 68 2D or 3D images, one for each image in a batch in parallel.
+        If detect_faces is None the method will also run a face detector.
+
+         Arguments:
+            image_batch {torch.tensor} -- The input images batch
+
+        Keyword Arguments:
+            detected_faces {list of numpy.array} -- list of bounding boxes, one for each face found
+            in the image (default: {None})
+            return_bboxes {boolean} -- If True, return the face bounding boxes in addition to the keypoints.
+            return_landmark_score {boolean} -- If True, return the keypoint scores along with the keypoints.
+
+        Return:
+            result:
+                1. if both return_bboxes and return_landmark_score are False, result will be:
+                    landmarks
+                2. Otherwise, result will be one of the following, depending on the actual value of return_* arguments.
+                    (landmark, landmark_score, detected_face)
+                    (landmark, None,           detected_face)
+                    (landmark, landmark_score, None         )
+        """
+
+        if detected_faces is None:
+            detected_faces = self.face_detector.detect_from_batch(image_batch)
+
+        if len(detected_faces) == 0:
+            warnings.warn("No faces were detected.")
+            if return_bboxes or return_landmark_score:
+                return None, None, None
+            else:
+                return None
+
+        landmarks = []
+        landmarks_scores_list = []
+        # A batch for each frame
+        for i, faces in enumerate(detected_faces):
+            res = self.get_landmarks_from_image(
+                image_batch[i].cpu().numpy().transpose(1, 2, 0),
+                detected_faces=faces,
+                return_landmark_score=return_landmark_score,
+            )
+            if return_landmark_score:
+                landmark_set, landmarks_scores, _ = res
+                landmarks_scores_list.append(landmarks_scores)
+            else:
+                landmark_set = res
+            # Bacward compatibility
+            if landmark_set is not None:
+                landmark_set = np.concatenate(landmark_set, axis=0)
+            else:
+                landmark_set = []
+            landmarks.append(landmark_set)
+
+        if not return_bboxes:
+            detected_faces = None
+        if not return_landmark_score:
+            landmarks_scores_list = None
+        if return_bboxes or return_landmark_score:
+            return landmarks, landmarks_scores_list, detected_faces
+        else:
+            return landmarks
+
+    def get_landmarks_from_directory(self, path, extensions=['.jpg', '.png'], recursive=True, show_progress_bar=True,
+                                     return_bboxes=False, return_landmark_score=False):
+        """Scan a directory for images with a given extension type(s) and predict the landmarks for each
+            face present in the images found.
+
+         Arguments:
+            path {str} -- path to the target directory containing the images
+
+        Keyword Arguments:
+            extensions {list of str} -- list containing the image extensions considered (default: ['.jpg', '.png'])
+            recursive {boolean} -- If True, scans for images recursively (default: True)
+            show_progress_bar {boolean} -- If True displays a progress bar (default: True)
+            return_bboxes {boolean} -- If True, return the face bounding boxes in addition to the keypoints.
+            return_landmark_score {boolean} -- If True, return the keypoint scores along with the keypoints.
+        """
+        dataset = FolderData(path, self.face_detector.tensor_or_path_to_ndarray, extensions, recursive, self.verbose)
+        dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=2, prefetch_factor=4)
+      
+        predictions = {}
+        for (image_path, image) in tqdm(dataloader, disable=not show_progress_bar):
+            image_path, image = image_path[0], image[0]
+            bounding_boxes = self.face_detector.detect_from_image(image)
+            if return_bboxes or return_landmark_score:
+                preds, bbox, score = self.get_landmarks_from_image(
+                    image, bounding_boxes, return_bboxes=return_bboxes, return_landmark_score=return_landmark_score)
+                predictions[image_path] = (preds, bbox, score)
+            else:
+                preds = self.get_landmarks_from_image(image, bounding_boxes)
+                predictions[image_path] = preds
+
+        return predictions

face_alignment/detection/__init__.py

@@ -0,0 +1 @@
+from .core import FaceDetector

face_alignment/detection/blazeface/__init__.py

@@ -0,0 +1 @@
+from .blazeface_detector import BlazeFaceDetector as FaceDetector

face_alignment/detection/blazeface/blazeface_detector.py

@@ -0,0 +1,69 @@
+from torch.utils.model_zoo import load_url
+
+from ..core import FaceDetector
+from ...utils import load_file_from_url
+
+from .net_blazeface import BlazeFace
+from .detect import *
+
+models_urls = {
+    'blazeface_weights': 'https://github.com/hollance/BlazeFace-PyTorch/blob/master/blazeface.pth?raw=true',
+    'blazeface_anchors': 'https://github.com/hollance/BlazeFace-PyTorch/blob/master/anchors.npy?raw=true',
+    'blazeface_back_weights': 'https://github.com/hollance/BlazeFace-PyTorch/blob/master/blazefaceback.pth?raw=true',
+    'blazeface_back_anchors': 'https://github.com/hollance/BlazeFace-PyTorch/blob/master/anchorsback.npy?raw=true'
+}
+
+
+class BlazeFaceDetector(FaceDetector):
+    def __init__(self, device, back_model=False, path_to_detector=None, path_to_anchor=None, verbose=False,
+                 min_score_thresh=0.5, min_suppression_threshold=0.3):
+        super(BlazeFaceDetector, self).__init__(device, verbose)
+
+        # Initialise the face detector
+        self.back_model = back_model
+        if path_to_detector is None:
+            if back_model:
+                model_weights = load_url(models_urls['blazeface_back_weights'])
+                model_anchors = np.load(load_file_from_url(models_urls['blazeface_back_anchors']))
+            else:
+                model_weights = load_url(models_urls['blazeface_weights'])
+                model_anchors = np.load(load_file_from_url(models_urls['blazeface_anchors']))
+        else:
+            model_weights = torch.load(path_to_detector)
+            model_anchors = np.load(path_to_anchor)
+
+        self.face_detector = BlazeFace(back_model=back_model)
+        self.face_detector.load_state_dict(model_weights)
+        self.face_detector.load_anchors_from_npy(model_anchors, device)
+
+        # Optionally change the thresholds:
+        self.face_detector.min_score_thresh = min_score_thresh
+        self.face_detector.min_suppression_threshold = min_suppression_threshold
+
+        self.face_detector.to(device)
+        self.face_detector.eval()
+
+    def detect_from_image(self, tensor_or_path):
+        image = self.tensor_or_path_to_ndarray(tensor_or_path)
+
+        image_size = 256 if self.back_model else 128
+        bboxlist = detect(self.face_detector, image, target_size=image_size, device=self.device)[0]
+
+        return bboxlist
+
+    def detect_from_batch(self, tensor):
+        image_size = 256 if self.back_model else 128
+        bboxlists = batch_detect(self.face_detector, tensor, target_size=image_size, device=self.device)
+        return bboxlists
+
+    @property
+    def reference_scale(self):
+        return 195
+
+    @property
+    def reference_x_shift(self):
+        return 0
+
+    @property
+    def reference_y_shift(self):
+        return 0

face_alignment/detection/blazeface/detect.py

@@ -0,0 +1,74 @@
+import torch
+import torch.nn.functional as F
+
+import cv2
+import numpy as np
+
+from .utils import *
+
+
+def detect(net, img, target_size=128, device='cuda'):
+    H, W, C = img.shape
+    orig_size = min(H, W)
+    img, (xshift, yshift) = resize_and_crop_image(img, target_size)
+    preds = net.predict_on_image(img)
+
+    if 0 == len(preds):
+        return [[]]
+
+    shift = np.array([xshift, yshift] * 2)
+    scores = preds[:, -1:]
+
+    # TODO: ugly
+    # reverses, x and y to adapt with face-alignment code
+    locs = np.concatenate((preds[:, 1:2], preds[:, 0:1], preds[:, 3:4], preds[:, 2:3]), axis=1)
+    return [np.concatenate((locs * orig_size + shift, scores), axis=1)]
+
+
+def batch_detect(net, img_batch, target_size=128, device='cuda'):
+    """
+    Inputs:
+        - net: BlazeFace model
+        - img_batch: a numpy array or tensor of shape (Batch size, Channels, Height, Width)
+        - target_size: target size of the input image
+    Outputs:
+        - list of 2-dim numpy arrays with shape (faces_on_this_image, 5): x1, y1, x2, y2, confidence
+          (x1, y1) - top left corner, (x2, y2) - bottom right corner
+    """
+    B, C, H, W = img_batch.shape
+    orig_size = min(H, W)
+
+    if isinstance(img_batch, torch.Tensor):
+        img_batch = img_batch.cpu().numpy()
+
+    img_batch = img_batch.transpose((0, 2, 3, 1))
+
+    imgs, (xshift, yshift) = resize_and_crop_batch(img_batch, target_size)
+    preds = net.predict_on_batch(imgs)
+    bboxlists = []
+    for pred in preds:
+        shift = np.array([xshift, yshift] * 2)
+        scores = pred[:, -1:]
+        locs = np.concatenate((pred[:, 1:2], pred[:, 0:1], pred[:, 3:4], pred[:, 2:3]), axis=1)
+        bboxlists.append(np.concatenate((locs * orig_size + shift, scores), axis=1))
+
+    return bboxlists
+
+
+def flip_detect(net, img, device):
+    img = cv2.flip(img, 1)
+    b = detect(net, img, device)
+
+    bboxlist = np.zeros(b.shape)
+    bboxlist[:, 0] = img.shape[1] - b[:, 2]
+    bboxlist[:, 1] = b[:, 1]
+    bboxlist[:, 2] = img.shape[1] - b[:, 0]
+    bboxlist[:, 3] = b[:, 3]
+    bboxlist[:, 4] = b[:, 4]
+    return bboxlist
+
+
+def pts_to_bb(pts):
+    min_x, min_y = np.min(pts, axis=0)
+    max_x, max_y = np.max(pts, axis=0)
+    return np.array([min_x, min_y, max_x, max_y])

face_alignment/detection/blazeface/net_blazeface.py

@@ -0,0 +1,482 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BlazeBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1):
+        super(BlazeBlock, self).__init__()
+
+        self.stride = stride
+        self.channel_pad = out_channels - in_channels
+
+        # TFLite uses slightly different padding than PyTorch
+        # on the depthwise conv layer when the stride is 2.
+        if stride == 2:
+            self.max_pool = nn.MaxPool2d(kernel_size=stride, stride=stride)
+            padding = 0
+        else:
+            padding = (kernel_size - 1) // 2
+
+        self.convs = nn.Sequential(
+            nn.Conv2d(in_channels=in_channels, out_channels=in_channels,
+                      kernel_size=kernel_size, stride=stride, padding=padding,
+                      groups=in_channels, bias=True),
+            nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
+                      kernel_size=1, stride=1, padding=0, bias=True),
+        )
+
+        self.act = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        if self.stride == 2:
+            h = F.pad(x, (0, 2, 0, 2), "constant", 0)
+            x = self.max_pool(x)
+        else:
+            h = x
+
+        if self.channel_pad > 0:
+            x = F.pad(x, (0, 0, 0, 0, 0, self.channel_pad), "constant", 0)
+
+        return self.act(self.convs(h) + x)
+    
+class FinalBlazeBlock(nn.Module):
+    def __init__(self, channels, kernel_size=3):
+        super(FinalBlazeBlock, self).__init__()
+        # TFLite uses slightly different padding than PyTorch
+        # on the depthwise conv layer when the stride is 2.
+        self.convs = nn.Sequential(
+            nn.Conv2d(in_channels=channels, out_channels=channels,
+                      kernel_size=kernel_size, stride=2, padding=0,
+                      groups=channels, bias=True),
+            nn.Conv2d(in_channels=channels, out_channels=channels,
+                      kernel_size=1, stride=1, padding=0, bias=True),
+        )
+
+        self.act = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        h = F.pad(x, (0, 2, 0, 2), "constant", 0)
+
+        return self.act(self.convs(h))
+
+
+class BlazeFace(nn.Module):
+    """The BlazeFace face detection model from MediaPipe.
+
+    The version from MediaPipe is simpler than the one in the paper;
+    it does not use the "double" BlazeBlocks.
+
+    Because we won't be training this model, it doesn't need to have
+    batchnorm layers. These have already been "folded" into the conv
+    weights by TFLite.
+
+    The conversion to PyTorch is fairly straightforward, but there are
+    some small differences between TFLite and PyTorch in how they handle
+    padding on conv layers with stride 2.
+
+    This version works on batches, while the MediaPipe version can only
+    handle a single image at a time.
+
+    Based on code from https://github.com/tkat0/PyTorch_BlazeFace/ and
+    https://github.com/google/mediapipe/
+    """
+
+    def __init__(self, back_model=False):
+        super(BlazeFace, self).__init__()
+
+        # These are the settings from the MediaPipe example graph
+        # mediapipe/graphs/face_detection/face_detection_mobile_gpu.pbtxt
+        # and mediapipe/graphs/face_detection/face_detection_back_mobile_gpu.pbtxt
+        self.num_classes = 1
+        self.num_anchors = 896
+        self.num_coords = 16
+        self.score_clipping_thresh = 100.0
+        self.back_model = back_model
+        if back_model:
+            self.x_scale = 256.0
+            self.y_scale = 256.0
+            self.h_scale = 256.0
+            self.w_scale = 256.0
+            self.min_score_thresh = 0.65
+        else:
+            self.x_scale = 128.0
+            self.y_scale = 128.0
+            self.h_scale = 128.0
+            self.w_scale = 128.0
+            self.min_score_thresh = 0.75
+        self.min_suppression_threshold = 0.3
+
+        self._define_layers()
+
+    def _define_back_model_layers(self):
+        self.backbone = nn.Sequential(
+            nn.Conv2d(in_channels=3, out_channels=24, kernel_size=5, stride=2, padding=0, bias=True),
+            nn.ReLU(inplace=True),
+            *[BlazeBlock(24, 24) for _ in range(7)],
+            BlazeBlock(24, 24, stride=2),
+            *[BlazeBlock(24, 24) for _ in range(7)],
+            BlazeBlock(24, 48, stride=2),
+            *[BlazeBlock(48, 48) for _ in range(7)],
+            BlazeBlock(48, 96, stride=2),
+            *[BlazeBlock(96, 96) for _ in range(7)],
+        )
+        self.final = FinalBlazeBlock(96)
+        self.classifier_8 = nn.Conv2d(96, 2, 1, bias=True)
+        self.classifier_16 = nn.Conv2d(96, 6, 1, bias=True)
+
+        self.regressor_8 = nn.Conv2d(96, 32, 1, bias=True)
+        self.regressor_16 = nn.Conv2d(96, 96, 1, bias=True)
+
+    def _define_front_model_layers(self):
+        self.backbone1 = nn.Sequential(
+            nn.Conv2d(in_channels=3, out_channels=24, kernel_size=5,
+                      stride=2, padding=0, bias=True),
+            nn.ReLU(inplace=True),
+
+            BlazeBlock(24, 24),
+            BlazeBlock(24, 28),
+            BlazeBlock(28, 32, stride=2),
+            BlazeBlock(32, 36),
+            BlazeBlock(36, 42),
+            BlazeBlock(42, 48, stride=2),
+            BlazeBlock(48, 56),
+            BlazeBlock(56, 64),
+            BlazeBlock(64, 72),
+            BlazeBlock(72, 80),
+            BlazeBlock(80, 88),
+        )
+
+        self.backbone2 = nn.Sequential(
+            BlazeBlock(88, 96, stride=2),
+            BlazeBlock(96, 96),
+            BlazeBlock(96, 96),
+            BlazeBlock(96, 96),
+            BlazeBlock(96, 96),
+        )
+
+        self.classifier_8 = nn.Conv2d(88, 2, 1, bias=True)
+        self.classifier_16 = nn.Conv2d(96, 6, 1, bias=True)
+
+        self.regressor_8 = nn.Conv2d(88, 32, 1, bias=True)
+        self.regressor_16 = nn.Conv2d(96, 96, 1, bias=True)
+
+
+    def _define_layers(self):
+        if self.back_model:
+            self._define_back_model_layers()
+        else:
+            self._define_front_model_layers()
+
+    def forward(self, x):
+        # TFLite uses slightly different padding on the first conv layer
+        # than PyTorch, so do it manually.
+        x = F.pad(x, (1, 2, 1, 2), "constant", 0)
+
+        b = x.shape[0]      # batch size, needed for reshaping later
+
+        if self.back_model:
+            x = self.backbone(x)           # (b, 16, 16, 96)
+            h = self.final(x)              # (b, 8, 8, 96)
+        else:
+            x = self.backbone1(x)           # (b, 88, 16, 16)
+            h = self.backbone2(x)           # (b, 96, 8, 8)
+
+        # Note: Because PyTorch is NCHW but TFLite is NHWC, we need to
+        # permute the output from the conv layers before reshaping it.
+
+        c1 = self.classifier_8(x)       # (b, 2, 16, 16)
+        c1 = c1.permute(0, 2, 3, 1)     # (b, 16, 16, 2)
+        c1 = c1.reshape(b, -1, 1)       # (b, 512, 1)
+
+        c2 = self.classifier_16(h)      # (b, 6, 8, 8)
+        c2 = c2.permute(0, 2, 3, 1)     # (b, 8, 8, 6)
+        c2 = c2.reshape(b, -1, 1)       # (b, 384, 1)
+
+        c = torch.cat((c1, c2), dim=1)  # (b, 896, 1)
+
+        r1 = self.regressor_8(x)        # (b, 32, 16, 16)
+        r1 = r1.permute(0, 2, 3, 1)     # (b, 16, 16, 32)
+        r1 = r1.reshape(b, -1, 16)      # (b, 512, 16)
+
+        r2 = self.regressor_16(h)       # (b, 96, 8, 8)
+        r2 = r2.permute(0, 2, 3, 1)     # (b, 8, 8, 96)
+        r2 = r2.reshape(b, -1, 16)      # (b, 384, 16)
+
+        r = torch.cat((r1, r2), dim=1)  # (b, 896, 16)
+        return [r, c]
+
+    def _device(self):
+        """Which device (CPU or GPU) is being used by this model?"""
+        return self.classifier_8.weight.device
+
+    def load_weights(self, path):
+        self.load_state_dict(torch.load(path))
+        self.eval()
+
+    def load_anchors(self, path, device=None):
+        device = device or self._device()
+        self.anchors = torch.tensor(
+            np.load(path), dtype=torch.float32, device=device)
+        assert(self.anchors.ndimension() == 2)
+        assert(self.anchors.shape[0] == self.num_anchors)
+        assert(self.anchors.shape[1] == 4)
+
+    def load_anchors_from_npy(self, arr, device=None):
+        device = device or self._device()
+        self.anchors = torch.tensor(
+            arr, dtype=torch.float32, device=device)
+        assert(self.anchors.ndimension() == 2)
+        assert(self.anchors.shape[0] == self.num_anchors)
+        assert(self.anchors.shape[1] == 4)
+
+    def _preprocess(self, x):
+        """Converts the image pixels to the range [-1, 1]."""
+        return x.float() / 127.5 - 1.0
+
+    def predict_on_image(self, img):
+        """Makes a prediction on a single image.
+
+        Arguments:
+            img: a NumPy array of shape (H, W, 3) or a PyTorch tensor of
+                 shape (3, H, W). The image's height and width should be
+                 128 pixels.
+
+        Returns:
+            A tensor with face detections.
+        """
+        if isinstance(img, np.ndarray):
+            img = torch.from_numpy(img).permute((2, 0, 1))
+
+        return self.predict_on_batch(img.unsqueeze(0))[0]
+
+    def predict_on_batch(self, x):
+        """Makes a prediction on a batch of images.
+
+        Arguments:
+            x: a NumPy array of shape (b, H, W, 3) or a PyTorch tensor of
+               shape (b, 3, H, W). The height and width should be 128 pixels.
+
+        Returns:
+            A list containing a tensor of face detections for each image in
+            the batch. If no faces are found for an image, returns a tensor
+            of shape (0, 17).
+
+        Each face detection is a PyTorch tensor consisting of 17 numbers:
+            - ymin, xmin, ymax, xmax
+            - x,y-coordinates for the 6 keypoints
+            - confidence score
+        """
+        if isinstance(x, np.ndarray):
+            x = torch.from_numpy(x).permute((0, 3, 1, 2))
+
+        assert x.shape[1] == 3
+        if self.back_model:
+            assert x.shape[2] == 256
+            assert x.shape[3] == 256
+        else:
+            assert x.shape[2] == 128
+            assert x.shape[3] == 128
+
+        # 1. Preprocess the images into tensors:
+        x = x.to(self._device())
+        x = self._preprocess(x)
+
+        # 2. Run the neural network:
+        with torch.inference_mode():
+            out = self.__call__(x)
+
+        # 3. Postprocess the raw predictions:
+        detections = self._tensors_to_detections(out[0], out[1], self.anchors)
+
+        # 4. Non-maximum suppression to remove overlapping detections:
+        filtered_detections = []
+        for i in range(len(detections)):
+            faces = self._weighted_non_max_suppression(detections[i])
+            faces = torch.stack(faces) if len(
+                faces) > 0 else torch.zeros((0, 17))
+            filtered_detections.append(faces)
+
+        return filtered_detections
+
+    def _tensors_to_detections(self, raw_box_tensor, raw_score_tensor, anchors):
+        """The output of the neural network is a tensor of shape (b, 896, 16)
+        containing the bounding box regressor predictions, as well as a tensor
+        of shape (b, 896, 1) with the classification confidences.
+
+        This function converts these two "raw" tensors into proper detections.
+        Returns a list of (num_detections, 17) tensors, one for each image in
+        the batch.
+
+        This is based on the source code from:
+        mediapipe/calculators/tflite/tflite_tensors_to_detections_calculator.cc
+        mediapipe/calculators/tflite/tflite_tensors_to_detections_calculator.proto
+        """
+        assert raw_box_tensor.ndimension() == 3
+        assert raw_box_tensor.shape[1] == self.num_anchors
+        assert raw_box_tensor.shape[2] == self.num_coords
+
+        assert raw_score_tensor.ndimension() == 3
+        assert raw_score_tensor.shape[1] == self.num_anchors
+        assert raw_score_tensor.shape[2] == self.num_classes
+
+        assert raw_box_tensor.shape[0] == raw_score_tensor.shape[0]
+
+        detection_boxes = self._decode_boxes(raw_box_tensor, anchors)
+
+        thresh = self.score_clipping_thresh
+        raw_score_tensor = raw_score_tensor.clamp(-thresh, thresh)
+        detection_scores = raw_score_tensor.sigmoid().squeeze(dim=-1)
+
+        # Note: we stripped off the last dimension from the scores tensor
+        # because there is only has one class. Now we can simply use a mask
+        # to filter out the boxes with too low confidence.
+        mask = detection_scores >= self.min_score_thresh
+
+        # Because each image from the batch can have a different number of
+        # detections, process them one at a time using a loop.
+        output_detections = []
+        for i in range(raw_box_tensor.shape[0]):
+            boxes = detection_boxes[i, mask[i]]
+            scores = detection_scores[i, mask[i]].unsqueeze(dim=-1)
+            output_detections.append(torch.cat((boxes, scores), dim=-1).to('cpu'))
+
+        return output_detections
+
+    def _decode_boxes(self, raw_boxes, anchors):
+        """Converts the predictions into actual coordinates using
+        the anchor boxes. Processes the entire batch at once.
+        """
+        boxes = torch.zeros_like(raw_boxes)
+
+        x_center = raw_boxes[..., 0] / self.x_scale * \
+            anchors[:, 2] + anchors[:, 0]
+        y_center = raw_boxes[..., 1] / self.y_scale * \
+            anchors[:, 3] + anchors[:, 1]
+
+        w = raw_boxes[..., 2] / self.w_scale * anchors[:, 2]
+        h = raw_boxes[..., 3] / self.h_scale * anchors[:, 3]
+
+        boxes[..., 0] = y_center - h / 2.  # ymin
+        boxes[..., 1] = x_center - w / 2.  # xmin
+        boxes[..., 2] = y_center + h / 2.  # ymax
+        boxes[..., 3] = x_center + w / 2.  # xmax
+
+        for k in range(6):
+            offset = 4 + k * 2
+            keypoint_x = raw_boxes[..., offset] / \
+                self.x_scale * anchors[:, 2] + anchors[:, 0]
+            keypoint_y = raw_boxes[..., offset + 1] / \
+                self.y_scale * anchors[:, 3] + anchors[:, 1]
+            boxes[..., offset] = keypoint_x
+            boxes[..., offset + 1] = keypoint_y
+
+        return boxes
+
+    def _weighted_non_max_suppression(self, detections):
+        """The alternative NMS method as mentioned in the BlazeFace paper:
+
+        "We replace the suppression algorithm with a blending strategy that
+        estimates the regression parameters of a bounding box as a weighted
+        mean between the overlapping predictions."
+
+        The original MediaPipe code assigns the score of the most confident
+        detection to the weighted detection, but we take the average score
+        of the overlapping detections.
+
+        The input detections should be a Tensor of shape (count, 17).
+
+        Returns a list of PyTorch tensors, one for each detected face.
+
+        This is based on the source code from:
+        mediapipe/calculators/util/non_max_suppression_calculator.cc
+        mediapipe/calculators/util/non_max_suppression_calculator.proto
+        """
+        if len(detections) == 0:
+            return []
+
+        output_detections = []
+
+        # Sort the detections from highest to lowest score.
+        remaining = torch.argsort(detections[:, 16], descending=True)
+
+        while len(remaining) > 0:
+            detection = detections[remaining[0]]
+
+            # Compute the overlap between the first box and the other
+            # remaining boxes. (Note that the other_boxes also include
+            # the first_box.)
+            first_box = detection[:4]
+            other_boxes = detections[remaining, :4]
+            ious = overlap_similarity(first_box, other_boxes)
+
+            # If two detections don't overlap enough, they are considered
+            # to be from different faces.
+            mask = ious > self.min_suppression_threshold
+            overlapping = remaining[mask]
+            remaining = remaining[~mask]
+
+            # Take an average of the coordinates from the overlapping
+            # detections, weighted by their confidence scores.
+            weighted_detection = detection.clone()
+            if len(overlapping) > 1:
+                coordinates = detections[overlapping, :16]
+                scores = detections[overlapping, 16:17]
+                total_score = scores.sum()
+                weighted = (coordinates * scores).sum(dim=0) / total_score
+                weighted_detection[:16] = weighted
+                weighted_detection[16] = total_score / len(overlapping)
+
+            output_detections.append(weighted_detection)
+
+        return output_detections
+
+
+# IOU code from https://github.com/amdegroot/ssd.pytorch/blob/master/layers/box_utils.py
+
+def intersect(box_a, box_b):
+    """ We resize both tensors to [A,B,2] without new malloc:
+    [A,2] -> [A,1,2] -> [A,B,2]
+    [B,2] -> [1,B,2] -> [A,B,2]
+    Then we compute the area of intersect between box_a and box_b.
+    Args:
+      box_a: (tensor) bounding boxes, Shape: [A,4].
+      box_b: (tensor) bounding boxes, Shape: [B,4].
+    Return:
+      (tensor) intersection area, Shape: [A,B].
+    """
+    A = box_a.size(0)
+    B = box_b.size(0)
+    max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2),
+                       box_b[:, 2:].unsqueeze(0).expand(A, B, 2))
+    min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2),
+                       box_b[:, :2].unsqueeze(0).expand(A, B, 2))
+    inter = torch.clamp((max_xy - min_xy), min=0)
+    return inter[:, :, 0] * inter[:, :, 1]
+
+
+def jaccard(box_a, box_b):
+    """Compute the jaccard overlap of two sets of boxes.  The jaccard overlap
+    is simply the intersection over union of two boxes.  Here we operate on
+    ground truth boxes and default boxes.
+    E.g.:
+        A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B)
+    Args:
+        box_a: (tensor) Ground truth bounding boxes, Shape: [num_objects,4]
+        box_b: (tensor) Prior boxes from priorbox layers, Shape: [num_priors,4]
+    Return:
+        jaccard overlap: (tensor) Shape: [box_a.size(0), box_b.size(0)]
+    """
+    inter = intersect(box_a, box_b)
+    area_a = ((box_a[:, 2] - box_a[:, 0])
+              * (box_a[:, 3] - box_a[:, 1])).unsqueeze(1).expand_as(inter)  # [A,B]
+    area_b = ((box_b[:, 2] - box_b[:, 0])
+              * (box_b[:, 3] - box_b[:, 1])).unsqueeze(0).expand_as(inter)  # [A,B]
+    union = area_a + area_b - inter
+    return inter / union  # [A,B]
+
+
+def overlap_similarity(box, other_boxes):
+    """Computes the IOU between a bounding box and set of other boxes."""
+    return jaccard(box.unsqueeze(0), other_boxes).squeeze(0)

face_alignment/detection/blazeface/utils.py

@@ -0,0 +1,65 @@
+import cv2
+import numpy as np
+
+
+def image_resize(image, width=None, height=None, inter=cv2.INTER_AREA):
+    # initialize the dimensions of the image to be resized and
+    # grab the image size
+    dim = None
+    (h, w) = image.shape[:2]
+
+    # if both the width and height are None, then return the
+    # original image
+    if width is None and height is None:
+        return image
+
+    # check to see if the width is None
+    if width is None:
+        # calculate the ratio of the height and construct the
+        # dimensions
+        r = height / float(h)
+        dim = (int(w * r), height)
+
+    # otherwise, the height is None
+    else:
+        # calculate the ratio of the width and construct the
+        # dimensions
+        r = width / float(w)
+        dim = (width, int(h * r))
+
+    # resize the image
+    resized = cv2.resize(image, dim, interpolation=inter)
+
+    # return the resized image
+    return resized
+
+
+def resize_and_crop_image(image, dim):
+    if image.shape[0] > image.shape[1]:
+        img = image_resize(image, width=dim)
+        yshift, xshift = (image.shape[0] - image.shape[1]) // 2, 0
+        y_start = (img.shape[0] - img.shape[1]) // 2
+        y_end = y_start + dim
+        return img[y_start:y_end, :, :], (xshift, yshift)
+    else:
+        img = image_resize(image, height=dim)
+        yshift, xshift = 0, (image.shape[1] - image.shape[0]) // 2
+        x_start = (img.shape[1] - img.shape[0]) // 2
+        x_end = x_start + dim
+        return img[:, x_start:x_end, :], (xshift, yshift)
+
+
+def resize_and_crop_batch(frames, dim):
+    """
+    Center crop + resize to (dim x dim)
+    inputs:
+        - frames: list of images (numpy arrays)
+        - dim: output dimension size
+    """
+    smframes = []
+    xshift, yshift = 0, 0
+    for i in range(len(frames)):
+        smframe, (xshift, yshift) = resize_and_crop_image(frames[i], dim)
+        smframes.append(smframe)
+    smframes = np.stack(smframes)
+    return smframes, (xshift, yshift)

face_alignment/detection/core.py

@@ -0,0 +1,150 @@
+import logging
+import glob
+from tqdm import tqdm
+import numpy as np
+import torch
+from skimage import io
+
+
+class FaceDetector(object):
+    """An abstract class representing a face detector.
+
+    Any other face detection implementation must subclass it. All subclasses
+    must implement ``detect_from_image``, that return a list of detected
+    bounding boxes. Optionally, for speed considerations detect from path is
+    recommended.
+    """
+
+    def __init__(self, device, verbose):
+        self.device = device
+        self.verbose = verbose
+
+        if verbose:
+            if 'cpu' in device:
+                logger = logging.getLogger(__name__)
+                logger.warning("Detection running on CPU, this may be potentially slow.")
+
+        if 'cpu' not in device and 'cuda' not in device and 'mps' not in device:
+            if verbose:
+                logger.error("Expected values for device are: {cpu, cuda, mps} but got: %s", device)
+            raise ValueError
+
+    def detect_from_image(self, tensor_or_path):
+        """Detects faces in a given image.
+
+        This function detects the faces present in a provided BGR(usually)
+        image. The input can be either the image itself or the path to it.
+
+        Arguments:
+            tensor_or_path {numpy.ndarray, torch.tensor or string} -- the path
+            to an image or the image itself.
+
+        Example::
+
+            >>> path_to_image = 'data/image_01.jpg'
+            ...   detected_faces = detect_from_image(path_to_image)
+            [A list of bounding boxes (x1, y1, x2, y2)]
+            >>> image = cv2.imread(path_to_image)
+            ...   detected_faces = detect_from_image(image)
+            [A list of bounding boxes (x1, y1, x2, y2)]
+
+        """
+        raise NotImplementedError
+
+    def detect_from_batch(self, tensor):
+        """Detects faces in a given image.
+
+        This function detects the faces present in a provided BGR(usually)
+        image. The input can be either the image itself or the path to it.
+
+        Arguments:
+            tensor {torch.tensor} -- image batch tensor.
+
+        Example::
+
+            >>> path_to_image = 'data/image_01.jpg'
+            ...   detected_faces = detect_from_image(path_to_image)
+            [A list of bounding boxes (x1, y1, x2, y2)]
+            >>> image = cv2.imread(path_to_image)
+            ...   detected_faces = detect_from_image(image)
+            [A list of bounding boxes (x1, y1, x2, y2)]
+
+        """
+        raise NotImplementedError
+
+    def detect_from_directory(self, path, extensions=['.jpg', '.png'], recursive=False, show_progress_bar=True):
+        """Detects faces from all the images present in a given directory.
+
+        Arguments:
+            path {string} -- a string containing a path that points to the folder containing the images
+
+        Keyword Arguments:
+            extensions {list} -- list of string containing the extensions to be
+            consider in the following format: ``.extension_name`` (default:
+            {['.jpg', '.png']}) recursive {bool} -- option wherever to scan the
+            folder recursively (default: {False}) show_progress_bar {bool} --
+            display a progressbar (default: {True})
+
+        Example:
+        >>> directory = 'data'
+        ...   detected_faces = detect_from_directory(directory)
+        {A dictionary of [lists containing bounding boxes(x1, y1, x2, y2)]}
+
+        """
+        if self.verbose:
+            logger = logging.getLogger(__name__)
+
+        if len(extensions) == 0:
+            if self.verbose:
+                logger.error("Expected at list one extension, but none was received.")
+            raise ValueError
+
+        if self.verbose:
+            logger.info("Constructing the list of images.")
+        additional_pattern = '/**/*' if recursive else '/*'
+        files = []
+        for extension in extensions:
+            files.extend(glob.glob(path + additional_pattern + extension, recursive=recursive))
+
+        if self.verbose:
+            logger.info("Finished searching for images. %s images found", len(files))
+            logger.info("Preparing to run the detection.")
+
+        predictions = {}
+        for image_path in tqdm(files, disable=not show_progress_bar):
+            if self.verbose:
+                logger.info("Running the face detector on image: %s", image_path)
+            predictions[image_path] = self.detect_from_image(image_path)
+
+        if self.verbose:
+            logger.info("The detector was successfully run on all %s images", len(files))
+
+        return predictions
+
+    @property
+    def reference_scale(self):
+        raise NotImplementedError
+
+    @property
+    def reference_x_shift(self):
+        raise NotImplementedError
+
+    @property
+    def reference_y_shift(self):
+        raise NotImplementedError
+
+    @staticmethod
+    def tensor_or_path_to_ndarray(tensor_or_path):
+        """Convert path (represented as a string) or torch.tensor to a numpy.ndarray
+
+        Arguments:
+            tensor_or_path {numpy.ndarray, torch.tensor or string} -- path to the image, or the image itself
+        """
+        if isinstance(tensor_or_path, str):
+            return io.imread(tensor_or_path)
+        elif torch.is_tensor(tensor_or_path):
+            return tensor_or_path.cpu().numpy()
+        elif isinstance(tensor_or_path, np.ndarray):
+            return tensor_or_path
+        else:
+            raise TypeError

face_alignment/detection/dlib/__init__.py

@@ -0,0 +1 @@
+from .dlib_detector import DlibDetector as FaceDetector

face_alignment/detection/dlib/dlib_detector.py

@@ -0,0 +1,48 @@
+import warnings
+import cv2
+import dlib
+
+from ..core import FaceDetector
+from ...utils import load_file_from_url
+
+
+class DlibDetector(FaceDetector):
+    def __init__(self, device, path_to_detector=None, verbose=False):
+        super().__init__(device, verbose)
+
+        warnings.warn('Warning: this detector is deprecated. Please use a different one, i.e.: S3FD.')
+
+        # Initialise the face detector
+        if 'cuda' in device:
+            if path_to_detector is None:
+                path_to_detector = load_file_from_url(
+                    "https://www.adrianbulat.com/downloads/dlib/mmod_human_face_detector.dat")
+
+            self.face_detector = dlib.cnn_face_detection_model_v1(path_to_detector)
+        else:
+            self.face_detector = dlib.get_frontal_face_detector()
+
+    def detect_from_image(self, tensor_or_path):
+        image = self.tensor_or_path_to_ndarray(tensor_or_path)
+        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+
+        detected_faces = self.face_detector(cv2.cvtColor(image, cv2.COLOR_BGR2GRAY))
+
+        if 'cuda' not in self.device:
+            detected_faces = [[d.left(), d.top(), d.right(), d.bottom()] for d in detected_faces]
+        else:
+            detected_faces = [[d.rect.left(), d.rect.top(), d.rect.right(), d.rect.bottom()] for d in detected_faces]
+
+        return detected_faces
+
+    @property
+    def reference_scale(self):
+        return 195
+
+    @property
+    def reference_x_shift(self):
+        return 0
+
+    @property
+    def reference_y_shift(self):
+        return 0

face_alignment/detection/folder/__init__.py

@@ -0,0 +1 @@
+from .folder_detector import FolderDetector as FaceDetector

face_alignment/detection/folder/folder_detector.py

@@ -0,0 +1,53 @@
+import os
+import numpy as np
+import torch
+
+from ..core import FaceDetector
+
+
+class FolderDetector(FaceDetector):
+    '''This is a simple helper module that assumes the faces were detected already
+        (either previously or are provided as ground truth).
+
+        The class expects to find the bounding boxes in the same format used by
+        the rest of face detectors, mainly ``list[(x1,y1,x2,y2),...]``.
+        For each image the detector will search for a file with the same name and with one of the
+        following extensions: .npy, .t7 or .pth
+
+    '''
+
+    def __init__(self, device, path_to_detector=None, verbose=False):
+        super(FolderDetector, self).__init__(device, verbose)
+
+    def detect_from_image(self, tensor_or_path):
+        # Only strings supported
+        if not isinstance(tensor_or_path, str):
+            raise ValueError
+
+        base_name = os.path.splitext(tensor_or_path)[0]
+
+        if os.path.isfile(base_name + '.npy'):
+            detected_faces = np.load(base_name + '.npy')
+        elif os.path.isfile(base_name + '.t7'):
+            detected_faces = torch.load(base_name + '.t7')
+        elif os.path.isfile(base_name + '.pth'):
+            detected_faces = torch.load(base_name + '.pth')
+        else:
+            raise FileNotFoundError
+
+        if not isinstance(detected_faces, list):
+            raise TypeError
+
+        return detected_faces
+
+    @property
+    def reference_scale(self):
+        return 195
+
+    @property
+    def reference_x_shift(self):
+        return 0
+
+    @property
+    def reference_y_shift(self):
+        return 0

face_alignment/detection/sfd/__init__.py

@@ -0,0 +1 @@
+from .sfd_detector import SFDDetector as FaceDetector

face_alignment/detection/sfd/bbox.py

@@ -0,0 +1,71 @@
+import math
+import numpy as np
+
+
+def nms(dets, thresh):
+    if 0 == len(dets):
+        return []
+    x1, y1, x2, y2, scores = dets[:, 0], dets[:, 1], dets[:, 2], dets[:, 3], dets[:, 4]
+    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
+    order = scores.argsort()[::-1]
+
+    keep = []
+    while order.size > 0:
+        i = order[0]
+        keep.append(i)
+        xx1, yy1 = np.maximum(x1[i], x1[order[1:]]), np.maximum(y1[i], y1[order[1:]])
+        xx2, yy2 = np.minimum(x2[i], x2[order[1:]]), np.minimum(y2[i], y2[order[1:]])
+
+        w, h = np.maximum(0.0, xx2 - xx1 + 1), np.maximum(0.0, yy2 - yy1 + 1)
+        ovr = w * h / (areas[i] + areas[order[1:]] - w * h)
+
+        inds = np.where(ovr <= thresh)[0]
+        order = order[inds + 1]
+
+    return keep
+
+
+def encode(matched, priors, variances):
+    """Encode the variances from the priorbox layers into the ground truth boxes
+    we have matched (based on jaccard overlap) with the prior boxes.
+    Args:
+        matched: (tensor) Coords of ground truth for each prior in point-form
+            Shape: [num_priors, 4].
+        priors: (tensor) Prior boxes in center-offset form
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        encoded boxes (tensor), Shape: [num_priors, 4]
+    """
+
+    # dist b/t match center and prior's center
+    g_cxcy = (matched[:, :2] + matched[:, 2:]) / 2 - priors[:, :2]
+    # encode variance
+    g_cxcy /= (variances[0] * priors[:, 2:])
+    # match wh / prior wh
+    g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:]
+    g_wh = np.log(g_wh) / variances[1]
+
+    # return target for smooth_l1_loss
+    return np.concatenate([g_cxcy, g_wh], 1)  # [num_priors,4]
+
+
+def decode(loc, priors, variances):
+    """Decode locations from predictions using priors to undo
+    the encoding we did for offset regression at train time.
+    Args:
+        loc (tensor): location predictions for loc layers,
+            Shape: [num_priors,4]
+        priors (tensor): Prior boxes in center-offset form.
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        decoded bounding box predictions
+    """
+
+    boxes = np.concatenate((
+        priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
+        priors[:, 2:] * np.exp(loc[:, 2:] * variances[1])), 1)
+    boxes[:, :2] -= boxes[:, 2:] / 2
+    boxes[:, 2:] += boxes[:, :2]
+    return boxes

face_alignment/detection/sfd/detect.py

@@ -0,0 +1,85 @@
+import torch
+import torch.nn.functional as F
+
+import cv2
+import numpy as np
+
+from .bbox import *
+
+
+def detect(net, img, device):
+    img = img.transpose(2, 0, 1)
+    # Creates a batch of 1
+    img = np.expand_dims(img, 0)
+
+    img = torch.from_numpy(img.copy()).to(device, dtype=torch.float32)
+
+    return batch_detect(net, img, device)
+
+
+def batch_detect(net, img_batch, device):
+    """
+    Inputs:
+        - img_batch: a torch.Tensor of shape (Batch size, Channels, Height, Width)
+    """
+
+    if 'cuda' in device:
+        torch.backends.cudnn.benchmark = True
+
+    batch_size = img_batch.size(0)
+    img_batch = img_batch.to(device, dtype=torch.float32)
+
+    img_batch = img_batch.flip(-3)  # RGB to BGR
+    img_batch = img_batch - torch.tensor([104.0, 117.0, 123.0], device=device).view(1, 3, 1, 1)
+
+    with torch.no_grad():
+        olist = net(img_batch)  # patched uint8_t overflow error
+
+    for i in range(len(olist) // 2):
+        olist[i * 2] = F.softmax(olist[i * 2], dim=1)
+
+    olist = [oelem.data.cpu().numpy() for oelem in olist]
+
+    bboxlists = get_predictions(olist, batch_size)
+    return bboxlists
+
+
+def get_predictions(olist, batch_size):
+    bboxlists = []
+    variances = [0.1, 0.2]
+    for i in range(len(olist) // 2):
+        ocls, oreg = olist[i * 2], olist[i * 2 + 1]
+        stride = 2**(i + 2)    # 4,8,16,32,64,128
+        poss = zip(*np.where(ocls[:, 1, :, :] > 0.05))
+        for Iindex, hindex, windex in poss:
+            axc, ayc = stride / 2 + windex * stride, stride / 2 + hindex * stride
+            priors = np.array([[axc / 1.0, ayc / 1.0, stride * 4 / 1.0, stride * 4 / 1.0]])
+            score = ocls[:, 1, hindex, windex][:,None]
+            loc = oreg[:, :, hindex, windex].copy()
+            boxes = decode(loc, priors, variances)
+            bboxlists.append(np.concatenate((boxes, score), axis=1))
+    
+    if len(bboxlists) == 0: # No candidates within given threshold
+        bboxlists = np.array([[] for _ in range(batch_size)])
+    else:
+        bboxlists = np.stack(bboxlists, axis=1)
+    return bboxlists
+
+
+def flip_detect(net, img, device):
+    img = cv2.flip(img, 1)
+    b = detect(net, img, device)
+
+    bboxlist = np.zeros(b.shape)
+    bboxlist[:, 0] = img.shape[1] - b[:, 2]
+    bboxlist[:, 1] = b[:, 1]
+    bboxlist[:, 2] = img.shape[1] - b[:, 0]
+    bboxlist[:, 3] = b[:, 3]
+    bboxlist[:, 4] = b[:, 4]
+    return bboxlist
+
+
+def pts_to_bb(pts):
+    min_x, min_y = np.min(pts, axis=0)
+    max_x, max_y = np.max(pts, axis=0)
+    return np.array([min_x, min_y, max_x, max_y])

face_alignment/detection/sfd/net_s3fd.py

@@ -0,0 +1,127 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class L2Norm(nn.Module):
+    def __init__(self, n_channels, scale=1.0):
+        super(L2Norm, self).__init__()
+        self.n_channels = n_channels
+        self.scale = scale
+        self.eps = 1e-10
+        self.weight = nn.Parameter(torch.empty(self.n_channels).fill_(self.scale))
+
+    def forward(self, x):
+        norm = x.pow(2).sum(dim=1, keepdim=True).sqrt() + self.eps
+        x = x / norm * self.weight.view(1, -1, 1, 1)
+        return x
+
+
+class s3fd(nn.Module):
+    def __init__(self):
+        super(s3fd, self).__init__()
+        self.conv1_1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
+        self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
+
+        self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
+        self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
+
+        self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)
+        self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
+        self.conv3_3 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
+
+        self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1)
+        self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
+        self.conv4_3 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
+
+        self.conv5_1 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
+        self.conv5_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
+        self.conv5_3 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
+
+        self.fc6 = nn.Conv2d(512, 1024, kernel_size=3, stride=1, padding=3)
+        self.fc7 = nn.Conv2d(1024, 1024, kernel_size=1, stride=1, padding=0)
+
+        self.conv6_1 = nn.Conv2d(1024, 256, kernel_size=1, stride=1, padding=0)
+        self.conv6_2 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1)
+
+        self.conv7_1 = nn.Conv2d(512, 128, kernel_size=1, stride=1, padding=0)
+        self.conv7_2 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
+
+        self.conv3_3_norm = L2Norm(256, scale=10)
+        self.conv4_3_norm = L2Norm(512, scale=8)
+        self.conv5_3_norm = L2Norm(512, scale=5)
+
+        self.conv3_3_norm_mbox_conf = nn.Conv2d(256, 4, kernel_size=3, stride=1, padding=1)
+        self.conv3_3_norm_mbox_loc = nn.Conv2d(256, 4, kernel_size=3, stride=1, padding=1)
+        self.conv4_3_norm_mbox_conf = nn.Conv2d(512, 2, kernel_size=3, stride=1, padding=1)
+        self.conv4_3_norm_mbox_loc = nn.Conv2d(512, 4, kernel_size=3, stride=1, padding=1)
+        self.conv5_3_norm_mbox_conf = nn.Conv2d(512, 2, kernel_size=3, stride=1, padding=1)
+        self.conv5_3_norm_mbox_loc = nn.Conv2d(512, 4, kernel_size=3, stride=1, padding=1)
+
+        self.fc7_mbox_conf = nn.Conv2d(1024, 2, kernel_size=3, stride=1, padding=1)
+        self.fc7_mbox_loc = nn.Conv2d(1024, 4, kernel_size=3, stride=1, padding=1)
+        self.conv6_2_mbox_conf = nn.Conv2d(512, 2, kernel_size=3, stride=1, padding=1)
+        self.conv6_2_mbox_loc = nn.Conv2d(512, 4, kernel_size=3, stride=1, padding=1)
+        self.conv7_2_mbox_conf = nn.Conv2d(256, 2, kernel_size=3, stride=1, padding=1)
+        self.conv7_2_mbox_loc = nn.Conv2d(256, 4, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x):
+        h = F.relu(self.conv1_1(x), inplace=True)
+        h = F.relu(self.conv1_2(h), inplace=True)
+        h = F.max_pool2d(h, 2, 2)
+
+        h = F.relu(self.conv2_1(h), inplace=True)
+        h = F.relu(self.conv2_2(h), inplace=True)
+        h = F.max_pool2d(h, 2, 2)
+
+        h = F.relu(self.conv3_1(h), inplace=True)
+        h = F.relu(self.conv3_2(h), inplace=True)
+        h = F.relu(self.conv3_3(h), inplace=True)
+        f3_3 = h
+        h = F.max_pool2d(h, 2, 2)
+
+        h = F.relu(self.conv4_1(h), inplace=True)
+        h = F.relu(self.conv4_2(h), inplace=True)
+        h = F.relu(self.conv4_3(h), inplace=True)
+        f4_3 = h
+        h = F.max_pool2d(h, 2, 2)
+
+        h = F.relu(self.conv5_1(h), inplace=True)
+        h = F.relu(self.conv5_2(h), inplace=True)
+        h = F.relu(self.conv5_3(h), inplace=True)
+        f5_3 = h
+        h = F.max_pool2d(h, 2, 2)
+
+        h = F.relu(self.fc6(h), inplace=True)
+        h = F.relu(self.fc7(h), inplace=True)
+        ffc7 = h
+        h = F.relu(self.conv6_1(h), inplace=True)
+        h = F.relu(self.conv6_2(h), inplace=True)
+        f6_2 = h
+        h = F.relu(self.conv7_1(h), inplace=True)
+        h = F.relu(self.conv7_2(h), inplace=True)
+        f7_2 = h
+
+        f3_3 = self.conv3_3_norm(f3_3)
+        f4_3 = self.conv4_3_norm(f4_3)
+        f5_3 = self.conv5_3_norm(f5_3)
+
+        cls1 = self.conv3_3_norm_mbox_conf(f3_3)
+        reg1 = self.conv3_3_norm_mbox_loc(f3_3)
+        cls2 = self.conv4_3_norm_mbox_conf(f4_3)
+        reg2 = self.conv4_3_norm_mbox_loc(f4_3)
+        cls3 = self.conv5_3_norm_mbox_conf(f5_3)
+        reg3 = self.conv5_3_norm_mbox_loc(f5_3)
+        cls4 = self.fc7_mbox_conf(ffc7)
+        reg4 = self.fc7_mbox_loc(ffc7)
+        cls5 = self.conv6_2_mbox_conf(f6_2)
+        reg5 = self.conv6_2_mbox_loc(f6_2)
+        cls6 = self.conv7_2_mbox_conf(f7_2)
+        reg6 = self.conv7_2_mbox_loc(f7_2)
+
+        # max-out background label
+        chunk = torch.chunk(cls1, 4, 1)
+        bmax = torch.max(torch.max(chunk[0], chunk[1]), chunk[2])
+        cls1 = torch.cat([bmax, chunk[3]], dim=1)
+
+        return [cls1, reg1, cls2, reg2, cls3, reg3, cls4, reg4, cls5, reg5, cls6, reg6]

face_alignment/detection/sfd/sfd_detector.py

@@ -0,0 +1,71 @@
+import torch
+from torch.utils.model_zoo import load_url
+
+from ..core import FaceDetector
+
+from .net_s3fd import s3fd
+from .bbox import nms
+from .detect import detect, batch_detect
+
+models_urls = {
+    's3fd': 'https://www.adrianbulat.com/downloads/python-fan/s3fd-619a316812.pth',
+}
+
+
+class SFDDetector(FaceDetector):
+    '''SF3D Detector.
+    '''
+
+    def __init__(self, device, path_to_detector=None, verbose=False, filter_threshold=0.5):
+        super(SFDDetector, self).__init__(device, verbose)
+
+        # Initialise the face detector
+        if path_to_detector is None:
+            model_weights = load_url(models_urls['s3fd'])
+        else:
+            model_weights = torch.load(path_to_detector)
+
+        self.fiter_threshold = filter_threshold
+        self.face_detector = s3fd()
+        self.face_detector.load_state_dict(model_weights)
+        self.face_detector.to(device)
+        self.face_detector.eval()
+
+    def _filter_bboxes(self, bboxlist):
+        if len(bboxlist) > 0:
+            keep = nms(bboxlist, 0.3)
+            bboxlist = bboxlist[keep, :]
+            bboxlist = [x for x in bboxlist if x[-1] > self.fiter_threshold]
+
+        return bboxlist
+
+    def detect_from_image(self, tensor_or_path):
+        image = self.tensor_or_path_to_ndarray(tensor_or_path)
+
+        bboxlist = detect(self.face_detector, image, device=self.device)[0]
+        bboxlist = self._filter_bboxes(bboxlist)
+
+        return bboxlist
+
+    def detect_from_batch(self, tensor):
+        bboxlists = batch_detect(self.face_detector, tensor, device=self.device)
+
+        new_bboxlists = []
+        for i in range(bboxlists.shape[0]):
+            bboxlist = bboxlists[i]
+            bboxlist = self._filter_bboxes(bboxlist)
+            new_bboxlists.append(bboxlist)
+
+        return new_bboxlists
+
+    @property
+    def reference_scale(self):
+        return 195
+
+    @property
+    def reference_x_shift(self):
+        return 0
+
+    @property
+    def reference_y_shift(self):
+        return 0

face_alignment/folder_data.py

@@ -0,0 +1,38 @@
+import logging
+import glob
+
+import torch
+
+class FolderData(torch.utils.data.Dataset):
+    def __init__(self, path, transforms, extensions=['.jpg', '.png'], recursive=False, verbose=False):
+        self.verbose = verbose
+        if self.verbose:
+            logger = logging.getLogger(__name__)
+        
+        if len(extensions) == 0:
+            if self.verbose:
+                logger.error("Expected at list one extension, but none was received.")
+            raise ValueError
+
+        if self.verbose:
+            logger.info("Constructing the list of images.")
+        additional_pattern = '/**/*' if recursive else '/*'
+        files = []
+        for extension in extensions:
+            files.extend(glob.glob(path + additional_pattern + extension, recursive=recursive))
+
+        if self.verbose:
+            logger.info("Finished searching for images. %s images found", len(files))
+            logger.info("Preparing to run the detection.")
+        
+        self.files = files
+        self.transforms = transforms
+
+    def __getitem__(self, idx):
+        image_path = self.files[idx]
+        image = self.transforms(image_path)
+
+        return image_path, image
+
+    def __len__(self):
+        return len(self.files)

face_alignment/utils.py

@@ -0,0 +1,379 @@
+import os
+import sys
+import errno
+import torch
+import math
+import numpy as np
+import cv2
+from skimage import io
+from skimage import color
+from numba import jit
+
+from urllib.parse import urlparse
+from torch.hub import download_url_to_file, HASH_REGEX
+try:
+    from torch.hub import get_dir
+except BaseException:
+    from torch.hub import _get_torch_home as get_dir
+
+gauss_kernel = None
+
+
+def _gaussian(
+        size=3, sigma=0.25, amplitude=1, normalize=False, width=None,
+        height=None, sigma_horz=None, sigma_vert=None, mean_horz=0.5,
+        mean_vert=0.5):
+    # handle some defaults
+    if width is None:
+        width = size
+    if height is None:
+        height = size
+    if sigma_horz is None:
+        sigma_horz = sigma
+    if sigma_vert is None:
+        sigma_vert = sigma
+    center_x = mean_horz * width + 0.5
+    center_y = mean_vert * height + 0.5
+    gauss = np.empty((height, width), dtype=np.float32)
+    # generate kernel
+    for i in range(height):
+        for j in range(width):
+            gauss[i][j] = amplitude * math.exp(-(math.pow((j + 1 - center_x) / (
+                sigma_horz * width), 2) / 2.0 + math.pow((i + 1 - center_y) / (sigma_vert * height), 2) / 2.0))
+    if normalize:
+        gauss = gauss / np.sum(gauss)
+    return gauss
+
+
+def draw_gaussian(image, point, sigma):
+    global gauss_kernel
+    # Check if the gaussian is inside
+    ul = [math.floor(point[0] - 3 * sigma), math.floor(point[1] - 3 * sigma)]
+    br = [math.floor(point[0] + 3 * sigma), math.floor(point[1] + 3 * sigma)]
+    if (ul[0] > image.shape[1] or ul[1] > image.shape[0] or br[0] < 1 or br[1] < 1):
+        return image
+    size = 6 * sigma + 1
+    if gauss_kernel is None:
+        g = _gaussian(size)
+        gauss_kernel = g
+    else:
+        g = gauss_kernel
+    g_x = [int(max(1, -ul[0])), int(min(br[0], image.shape[1])) - int(max(1, ul[0])) + int(max(1, -ul[0]))]
+    g_y = [int(max(1, -ul[1])), int(min(br[1], image.shape[0])) - int(max(1, ul[1])) + int(max(1, -ul[1]))]
+    img_x = [int(max(1, ul[0])), int(min(br[0], image.shape[1]))]
+    img_y = [int(max(1, ul[1])), int(min(br[1], image.shape[0]))]
+    assert (g_x[0] > 0 and g_y[1] > 0)
+    image[img_y[0] - 1:img_y[1], img_x[0] - 1:img_x[1]
+          ] = image[img_y[0] - 1:img_y[1], img_x[0] - 1:img_x[1]] + g[g_y[0] - 1:g_y[1], g_x[0] - 1:g_x[1]]
+    image[image > 1] = 1
+    return image
+
+
+def transform(point, center, scale, resolution, invert=False):
+    """Generate and affine transformation matrix.
+
+    Given a set of points, a center, a scale and a targer resolution, the
+    function generates and affine transformation matrix. If invert is ``True``
+    it will produce the inverse transformation.
+
+    Arguments:
+        point {torch.tensor} -- the input 2D point
+        center {torch.tensor or numpy.array} -- the center around which to perform the transformations
+        scale {float} -- the scale of the face/object
+        resolution {float} -- the output resolution
+
+    Keyword Arguments:
+        invert {bool} -- define wherever the function should produce the direct or the
+        inverse transformation matrix (default: {False})
+    """
+    _pt = torch.ones(3)
+    _pt[0] = point[0]
+    _pt[1] = point[1]
+
+    h = 200.0 * scale
+    t = torch.eye(3)
+    t[0, 0] = resolution / h
+    t[1, 1] = resolution / h
+    t[0, 2] = resolution * (-center[0] / h + 0.5)
+    t[1, 2] = resolution * (-center[1] / h + 0.5)
+
+    if invert:
+        t = torch.inverse(t)
+
+    new_point = (torch.matmul(t, _pt))[0:2]
+
+    return new_point.int()
+
+
+def crop(image, center, scale, resolution=256.0):
+    """Center crops an image or set of heatmaps
+
+    Arguments:
+        image {numpy.array} -- an rgb image
+        center {numpy.array} -- the center of the object, usually the same as of the bounding box
+        scale {float} -- scale of the face
+
+    Keyword Arguments:
+        resolution {float} -- the size of the output cropped image (default: {256.0})
+
+    Returns:
+        [type] -- [description]
+    """  # Crop around the center point
+    """ Crops the image around the center. Input is expected to be an np.ndarray """
+    ul = transform([1, 1], center, scale, resolution, True)
+    br = transform([resolution, resolution], center, scale, resolution, True)
+    # pad = math.ceil(torch.norm((ul - br).float()) / 2.0 - (br[0] - ul[0]) / 2.0)
+    if image.ndim > 2:
+        newDim = np.array([br[1] - ul[1], br[0] - ul[0],
+                           image.shape[2]], dtype=np.int32)
+        newImg = np.zeros(newDim, dtype=np.uint8)
+    else:
+        newDim = np.array([br[1] - ul[1], br[0] - ul[0]], dtype=np.int)
+        newImg = np.zeros(newDim, dtype=np.uint8)
+    ht = image.shape[0]
+    wd = image.shape[1]
+    newX = np.array(
+        [max(1, -ul[0] + 1), min(br[0], wd) - ul[0]], dtype=np.int32)
+    newY = np.array(
+        [max(1, -ul[1] + 1), min(br[1], ht) - ul[1]], dtype=np.int32)
+    oldX = np.array([max(1, ul[0] + 1), min(br[0], wd)], dtype=np.int32)
+    oldY = np.array([max(1, ul[1] + 1), min(br[1], ht)], dtype=np.int32)
+    newImg[newY[0] - 1:newY[1], newX[0] - 1:newX[1]
+           ] = image[oldY[0] - 1:oldY[1], oldX[0] - 1:oldX[1], :]
+    newImg = cv2.resize(newImg, dsize=(int(resolution), int(resolution)),
+                        interpolation=cv2.INTER_LINEAR)
+    return newImg
+
+
+@jit(nopython=True)
+def transform_np(point, center, scale, resolution, invert=False):
+    """Generate and affine transformation matrix.
+
+    Given a set of points, a center, a scale and a targer resolution, the
+    function generates and affine transformation matrix. If invert is ``True``
+    it will produce the inverse transformation.
+
+    Arguments:
+        point {numpy.array} -- the input 2D point
+        center {numpy.array} -- the center around which to perform the transformations
+        scale {float} -- the scale of the face/object
+        resolution {float} -- the output resolution
+
+    Keyword Arguments:
+        invert {bool} -- define wherever the function should produce the direct or the
+        inverse transformation matrix (default: {False})
+    """
+    _pt = np.ones(3)
+    _pt[0] = point[0]
+    _pt[1] = point[1]
+
+    h = 200.0 * scale
+    t = np.eye(3)
+    t[0, 0] = resolution / h
+    t[1, 1] = resolution / h
+    t[0, 2] = resolution * (-center[0] / h + 0.5)
+    t[1, 2] = resolution * (-center[1] / h + 0.5)
+
+    if invert:
+        t = np.ascontiguousarray(np.linalg.pinv(t))
+
+    new_point = np.dot(t, _pt)[0:2]
+
+    return new_point.astype(np.int32)
+
+
+def get_preds_fromhm(hm, center=None, scale=None):
+    """Obtain (x,y) coordinates given a set of N heatmaps. If the center
+    and the scale is provided the function will return the points also in
+    the original coordinate frame.
+
+    Arguments:
+        hm {torch.tensor} -- the predicted heatmaps, of shape [B, N, W, H]
+
+    Keyword Arguments:
+        center {torch.tensor} -- the center of the bounding box (default: {None})
+        scale {float} -- face scale (default: {None})
+    """
+    B, C, H, W = hm.shape
+    hm_reshape = hm.reshape(B, C, H * W)
+    idx = np.argmax(hm_reshape, axis=-1)
+    scores = np.take_along_axis(hm_reshape, np.expand_dims(idx, axis=-1), axis=-1).squeeze(-1)
+    preds, preds_orig = _get_preds_fromhm(hm, idx, center, scale)
+
+    return preds, preds_orig, scores
+
+
+@jit(nopython=True)
+def _get_preds_fromhm(hm, idx, center=None, scale=None):
+    """Obtain (x,y) coordinates given a set of N heatmaps and the
+    coresponding locations of the maximums. If the center
+    and the scale is provided the function will return the points also in
+    the original coordinate frame.
+
+    Arguments:
+        hm {torch.tensor} -- the predicted heatmaps, of shape [B, N, W, H]
+
+    Keyword Arguments:
+        center {torch.tensor} -- the center of the bounding box (default: {None})
+        scale {float} -- face scale (default: {None})
+    """
+    B, C, H, W = hm.shape
+    idx += 1
+    preds = idx.repeat(2).reshape(B, C, 2).astype(np.float32)
+    preds[:, :, 0] = (preds[:, :, 0] - 1) % W + 1
+    preds[:, :, 1] = np.floor((preds[:, :, 1] - 1) / H) + 1
+
+    for i in range(B):
+        for j in range(C):
+            hm_ = hm[i, j, :]
+            pX, pY = int(preds[i, j, 0]) - 1, int(preds[i, j, 1]) - 1
+            if pX > 0 and pX < 63 and pY > 0 and pY < 63:
+                diff = np.array(
+                    [hm_[pY, pX + 1] - hm_[pY, pX - 1],
+                     hm_[pY + 1, pX] - hm_[pY - 1, pX]])
+                preds[i, j] += np.sign(diff) * 0.25
+
+    preds -= 0.5
+
+    preds_orig = np.zeros_like(preds)
+    if center is not None and scale is not None:
+        for i in range(B):
+            for j in range(C):
+                preds_orig[i, j] = transform_np(
+                    preds[i, j], center, scale, H, True)
+
+    return preds, preds_orig
+
+
+def create_target_heatmap(target_landmarks, centers, scales):
+    heatmaps = np.zeros((target_landmarks.shape[0], 68, 64, 64), dtype=np.float32)
+    for i in range(heatmaps.shape[0]):
+        for p in range(68):
+            landmark_cropped_coor = transform(target_landmarks[i, p] + 1, centers[i], scales[i], 64, invert=False)
+            heatmaps[i, p] = draw_gaussian(heatmaps[i, p], landmark_cropped_coor + 1, 2)
+    return torch.tensor(heatmaps)
+
+
+def create_bounding_box(target_landmarks, expansion_factor=0.0):
+    """
+    gets a batch of landmarks and calculates a bounding box that includes all the landmarks per set of landmarks in
+    the batch
+    :param target_landmarks: batch of landmarks of dim (n x 68 x 2). Where n is the batch size
+    :param expansion_factor: expands the bounding box by this factor. For example, a `expansion_factor` of 0.2 leads
+    to 20% increase in width and height of the boxes
+    :return: a batch of bounding boxes of dim (n x 4) where the second dim is (x1,y1,x2,y2)
+    """
+    # Calc bounding box
+    x_y_min, _ = target_landmarks.reshape(-1, 68, 2).min(dim=1)
+    x_y_max, _ = target_landmarks.reshape(-1, 68, 2).max(dim=1)
+    # expanding the bounding box
+    expansion_factor /= 2
+    bb_expansion_x = (x_y_max[:, 0] - x_y_min[:, 0]) * expansion_factor
+    bb_expansion_y = (x_y_max[:, 1] - x_y_min[:, 1]) * expansion_factor
+    x_y_min[:, 0] -= bb_expansion_x
+    x_y_max[:, 0] += bb_expansion_x
+    x_y_min[:, 1] -= bb_expansion_y
+    x_y_max[:, 1] += bb_expansion_y
+    return torch.cat([x_y_min, x_y_max], dim=1)
+
+
+def shuffle_lr(parts, pairs=None):
+    """Shuffle the points left-right according to the axis of symmetry
+    of the object.
+
+    Arguments:
+        parts {torch.tensor} -- a 3D or 4D object containing the
+        heatmaps.
+
+    Keyword Arguments:
+        pairs {list of integers} -- [order of the flipped points] (default: {None})
+    """
+    if pairs is None:
+        pairs = [16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0,
+                 26, 25, 24, 23, 22, 21, 20, 19, 18, 17, 27, 28, 29, 30, 35,
+                 34, 33, 32, 31, 45, 44, 43, 42, 47, 46, 39, 38, 37, 36, 41,
+                 40, 54, 53, 52, 51, 50, 49, 48, 59, 58, 57, 56, 55, 64, 63,
+                 62, 61, 60, 67, 66, 65]
+    if parts.ndimension() == 3:
+        parts = parts[pairs, ...]
+    else:
+        parts = parts[:, pairs, ...]
+
+    return parts
+
+
+def flip(tensor, is_label=False):
+    """Flip an image or a set of heatmaps left-right
+
+    Arguments:
+        tensor {numpy.array or torch.tensor} -- [the input image or heatmaps]
+
+    Keyword Arguments:
+        is_label {bool} -- [denote wherever the input is an image or a set of heatmaps ] (default: {False})
+    """
+    if not torch.is_tensor(tensor):
+        tensor = torch.from_numpy(tensor)
+
+    if is_label:
+        tensor = shuffle_lr(tensor).flip(tensor.ndimension() - 1)
+    else:
+        tensor = tensor.flip(tensor.ndimension() - 1)
+
+    return tensor
+
+
+def get_image(image_or_path):
+    """Reads an image from file or array/tensor and converts it to RGB (H,W,3).
+
+    Arguments:
+        tensor {Sstring, numpy.array or torch.tensor} -- [the input image or path to it]
+    """
+    if isinstance(image_or_path, str):
+        try:
+            image = io.imread(image_or_path)
+        except IOError:
+            print("error opening file :: ", image_or_path)
+            return None
+    elif isinstance(image_or_path, torch.Tensor):
+        image = image_or_path.detach().cpu().numpy()
+    else:
+        image = image_or_path
+
+    if image.ndim == 2:
+        image = color.gray2rgb(image)
+    elif image.ndim == 4:
+        image = image[..., :3]
+
+    return image
+
+
+# Pytorch load supports only pytorch models
+def load_file_from_url(url, model_dir=None, progress=True, check_hash=False, file_name=None):
+    if model_dir is None:
+        hub_dir = get_dir()
+        model_dir = os.path.join(hub_dir, 'checkpoints')
+
+    try:
+        os.makedirs(model_dir)
+    except OSError as e:
+        if e.errno == errno.EEXIST:
+            # Directory already exists, ignore.
+            pass
+        else:
+            # Unexpected OSError, re-raise.
+            raise
+
+    parts = urlparse(url)
+    filename = os.path.basename(parts.path)
+    if file_name is not None:
+        filename = file_name
+    cached_file = os.path.join(model_dir, filename)
+    if not os.path.exists(cached_file):
+        sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file))
+        hash_prefix = None
+        if check_hash:
+            r = HASH_REGEX.search(filename)  # r is Optional[Match[str]]
+            hash_prefix = r.group(1) if r else None
+        download_url_to_file(url, cached_file, hash_prefix, progress=progress)
+
+    return cached_file

liveportrait/utils/cropper.py

@@ -1,6 +1,7 @@
 # coding: utf-8
 
 import numpy as np
+import torch
 from typing import List, Union, Tuple
 from dataclasses import dataclass, field
 import cv2#; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False)
@@ -142,4 +143,70 @@ class CropperMediaPipe(object):
         lmk = recon_ret['pts']
         ret_dct['lmk_crop'] = lmk
 
-        return ret_dct, cropped_image_256
+        return ret_dct, cropped_image_256
+    
+class CropperFaceAlignment(object):
+    def __init__(self, **kwargs) -> None:
+        device_id = kwargs.get('device_id', 0)
+        provider = kwargs.get('onnx_device', 'CPU')
+        face_detector_device = kwargs.get('face_detector_device', 'cuda')
+        face_detector = kwargs.get('face_detector', 'blazeface')
+        face_detector_dtype = kwargs.get('face_detector_dtype', 'fp16')
+
+        if provider != "torch_gpu":
+            self.landmark_runner = LandmarkRunner(
+                ckpt_path=os.path.join(folder_paths.models_dir, 'liveportrait', 'landmark.onnx'),
+                onnx_provider=provider,
+                device_id=device_id
+                )
+            self.landmark_runner.warmup()
+        else:
+            self.landmark_runner = LandmarkRunnerTorch(
+                    ckpt_path=os.path.join(folder_paths.models_dir, 'liveportrait', 'landmark_model.pth'),
+                    onnx_provider=provider,
+                    device_id=device_id
+                )
+            
+        from ...face_alignment import FaceAlignment, LandmarksType
+        if 'blazeface' in face_detector:
+            face_detector_kwargs = {'back_model': face_detector == 'blazeface_back_camera'}
+            self.fa = FaceAlignment(LandmarksType.TWO_D, flip_input=False, device=face_detector_device, dtype=face_detector_dtype, face_detector='blazeface', face_detector_kwargs=face_detector_kwargs)
+        else:
+            self.fa = FaceAlignment(LandmarksType.TWO_D, flip_input=False, device=face_detector_device, dtype=face_detector_dtype, face_detector=face_detector)
+
+    def crop_single_image(self, img_rgb, dsize, scale, vy_ratio, vx_ratio, face_index, face_index_order, rotate):
+       
+        face_result = self.fa.get_landmarks_from_image(img_rgb)
+
+        if face_result is None:
+            ret_dct = {}
+            cropped_image_256 = None
+            return ret_dct, cropped_image_256
+
+        face_landmarks = face_result[face_index]
+
+        pts = np.array(face_landmarks)
+       
+        # crop the face
+        ret_dct, image_crop = crop_image(
+            img_rgb,  # ndarray
+            pts,  # 106x2 or Nx2
+            dsize=dsize,
+            scale=scale,
+            vy_ratio=vy_ratio,
+            vx_ratio=vx_ratio,
+            rotate=rotate
+        )
+        # update a 256x256 version for network input or else
+        cropped_image_256 = cv2.resize(image_crop, (256, 256), interpolation=cv2.INTER_AREA)
+        ret_dct['pt_crop_256x256'] = ret_dct['pt_crop'] * 256 / dsize
+
+        input_image_size = img_rgb.shape[:2]
+        ret_dct['input_image_size'] = input_image_size
+    
+        recon_ret = self.landmark_runner.run(img_rgb, pts)
+        lmk = recon_ret['pts']
+        ret_dct['lmk_crop'] = lmk
+
+        return ret_dct, cropped_image_256
+    

nodes.py

@@ -24,6 +24,10 @@ try:
     from .liveportrait.utils.cropper import CropperInsightFace
 except:
     log.warning("Can't load MediaPipe, MediaPipeCropper not available")
+try:
+    from .liveportrait.utils.cropper import CropperFaceAlignment
+except:
+    log.warning("Can't load FaceAlignment, CropperFaceAlignment not available")
 
 from .liveportrait.modules.spade_generator import SPADEDecoder
 from .liveportrait.modules.warping_network import WarpingNetwork
@@ -529,6 +533,58 @@ class LivePortraitLoadMediaPipeCropper:
 
         return (self.cropper,)
     
+class LivePortraitLoadFaceAlignmentCropper:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+            "face_detector": (
+                    ['blazeface', 'blazeface_back_camera', 'sfd'], {
+                        "default": 'blazeface_back_camera'
+                    }),
+
+            "landmarkrunner_device": (
+                    ['CPU', 'CUDA', 'ROCM', 'CoreML', 'torch_gpu'], {
+                        "default": 'torch_gpu'
+                    }),
+            "face_detector_device": (
+                    ['cuda', 'cpu', 'mps'], {
+                        "default": 'cuda'
+                    }),
+
+            "face_detector_dtype": (
+                    [
+                        "fp16",
+                        "bf16",
+                        "fp32",
+                    ],
+                    {"default": "fp16"},
+                ),
+            "keep_model_loaded": ("BOOLEAN", {"default": True})
+
+            },           
+        }
+
+    RETURN_TYPES = ("LPCROPPER",)
+    RETURN_NAMES = ("cropper",)
+    FUNCTION = "crop"
+    CATEGORY = "LivePortrait"
+
+    def crop(self, landmarkrunner_device, keep_model_loaded, face_detector, face_detector_device, face_detector_dtype):
+        dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[face_detector_dtype]
+        cropper_init_config = {
+            'keep_model_loaded': keep_model_loaded,
+            'onnx_device': landmarkrunner_device,
+            'face_detector_device': face_detector_device,
+            'face_detector': face_detector,
+            'face_detector_dtype': dtype
+        }
+        
+        if not hasattr(self, 'cropper') or self.cropper is None or self.current_config != cropper_init_config:
+            self.current_config = cropper_init_config
+            self.cropper = CropperFaceAlignment(**cropper_init_config)
+
+        return (self.cropper,)
+    
 class LivePortraitCropper:
     @classmethod
     def INPUT_TYPES(s):
@@ -759,6 +815,7 @@ NODE_CLASS_MAPPINGS = {
     "KeypointsToImage": KeypointsToImage,
     "LivePortraitLoadCropper": LivePortraitLoadCropper,
     "LivePortraitLoadMediaPipeCropper": LivePortraitLoadMediaPipeCropper,
+    "LivePortraitLoadFaceAlignmentCropper": LivePortraitLoadFaceAlignmentCropper,
     "LivePortraitComposite": LivePortraitComposite,
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -770,5 +827,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "KeypointsToImage": "LivePortrait KeypointsToImage",
     "LivePortraitLoadCropper": "LivePortrait Load InsightFaceCropper",
     "LivePortraitLoadMediaPipeCropper": "LivePortrait Load MediaPipeCropper",
+    "LivePortraitLoadFaceAlignmentCropper": "LivePortrait Load FaceAlignmentCropper",
     "LivePortraitComposite": "LivePortrait Composite",
     }