From 6cdb0ebde530258623e22f4b11dbc8670197f92d Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Thu, 25 Jul 2024 11:54:42 +0300 Subject: [PATCH] Add face_alignment as alternative detector --- face_alignment/LICENSE | 29 ++ face_alignment/__init__.py | 7 + face_alignment/api.py | 305 +++++++++++ face_alignment/detection/__init__.py | 1 + .../detection/blazeface/__init__.py | 1 + .../detection/blazeface/blazeface_detector.py | 69 +++ face_alignment/detection/blazeface/detect.py | 74 +++ .../detection/blazeface/net_blazeface.py | 482 ++++++++++++++++++ face_alignment/detection/blazeface/utils.py | 65 +++ face_alignment/detection/core.py | 150 ++++++ face_alignment/detection/dlib/__init__.py | 1 + .../detection/dlib/dlib_detector.py | 48 ++ face_alignment/detection/folder/__init__.py | 1 + .../detection/folder/folder_detector.py | 53 ++ face_alignment/detection/sfd/__init__.py | 1 + face_alignment/detection/sfd/bbox.py | 71 +++ face_alignment/detection/sfd/detect.py | 85 +++ face_alignment/detection/sfd/net_s3fd.py | 127 +++++ face_alignment/detection/sfd/sfd_detector.py | 71 +++ face_alignment/folder_data.py | 38 ++ face_alignment/utils.py | 379 ++++++++++++++ liveportrait/utils/cropper.py | 69 ++- nodes.py | 58 +++ 23 files changed, 2184 insertions(+), 1 deletion(-) create mode 100644 face_alignment/LICENSE create mode 100644 face_alignment/__init__.py create mode 100644 face_alignment/api.py create mode 100644 face_alignment/detection/__init__.py create mode 100644 face_alignment/detection/blazeface/__init__.py create mode 100644 face_alignment/detection/blazeface/blazeface_detector.py create mode 100644 face_alignment/detection/blazeface/detect.py create mode 100644 face_alignment/detection/blazeface/net_blazeface.py create mode 100644 face_alignment/detection/blazeface/utils.py create mode 100644 face_alignment/detection/core.py create mode 100644 face_alignment/detection/dlib/__init__.py create mode 100644 face_alignment/detection/dlib/dlib_detector.py create mode 100644 face_alignment/detection/folder/__init__.py create mode 100644 face_alignment/detection/folder/folder_detector.py create mode 100644 face_alignment/detection/sfd/__init__.py create mode 100644 face_alignment/detection/sfd/bbox.py create mode 100644 face_alignment/detection/sfd/detect.py create mode 100644 face_alignment/detection/sfd/net_s3fd.py create mode 100644 face_alignment/detection/sfd/sfd_detector.py create mode 100644 face_alignment/folder_data.py create mode 100644 face_alignment/utils.py diff --git a/face_alignment/LICENSE b/face_alignment/LICENSE new file mode 100644 index 0000000..ed4c6d1 --- /dev/null +++ b/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. diff --git a/face_alignment/__init__.py b/face_alignment/__init__.py new file mode 100644 index 0000000..ba91140 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/face_alignment/api.py b/face_alignment/api.py new file mode 100644 index 0000000..6b4d767 --- /dev/null +++ b/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 diff --git a/face_alignment/detection/__init__.py b/face_alignment/detection/__init__.py new file mode 100644 index 0000000..1a6b040 --- /dev/null +++ b/face_alignment/detection/__init__.py @@ -0,0 +1 @@ +from .core import FaceDetector \ No newline at end of file diff --git a/face_alignment/detection/blazeface/__init__.py b/face_alignment/detection/blazeface/__init__.py new file mode 100644 index 0000000..e2a4708 --- /dev/null +++ b/face_alignment/detection/blazeface/__init__.py @@ -0,0 +1 @@ +from .blazeface_detector import BlazeFaceDetector as FaceDetector diff --git a/face_alignment/detection/blazeface/blazeface_detector.py b/face_alignment/detection/blazeface/blazeface_detector.py new file mode 100644 index 0000000..77c750f --- /dev/null +++ b/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 diff --git a/face_alignment/detection/blazeface/detect.py b/face_alignment/detection/blazeface/detect.py new file mode 100644 index 0000000..6da231a --- /dev/null +++ b/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]) diff --git a/face_alignment/detection/blazeface/net_blazeface.py b/face_alignment/detection/blazeface/net_blazeface.py new file mode 100644 index 0000000..e4da671 --- /dev/null +++ b/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) diff --git a/face_alignment/detection/blazeface/utils.py b/face_alignment/detection/blazeface/utils.py new file mode 100644 index 0000000..7596a8c --- /dev/null +++ b/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) diff --git a/face_alignment/detection/core.py b/face_alignment/detection/core.py new file mode 100644 index 0000000..46cec5f --- /dev/null +++ b/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 diff --git a/face_alignment/detection/dlib/__init__.py b/face_alignment/detection/dlib/__init__.py new file mode 100644 index 0000000..d8e5ee5 --- /dev/null +++ b/face_alignment/detection/dlib/__init__.py @@ -0,0 +1 @@ +from .dlib_detector import DlibDetector as FaceDetector \ No newline at end of file diff --git a/face_alignment/detection/dlib/dlib_detector.py b/face_alignment/detection/dlib/dlib_detector.py new file mode 100644 index 0000000..0af3316 --- /dev/null +++ b/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 diff --git a/face_alignment/detection/folder/__init__.py b/face_alignment/detection/folder/__init__.py new file mode 100644 index 0000000..0a9128e --- /dev/null +++ b/face_alignment/detection/folder/__init__.py @@ -0,0 +1 @@ +from .folder_detector import FolderDetector as FaceDetector \ No newline at end of file diff --git a/face_alignment/detection/folder/folder_detector.py b/face_alignment/detection/folder/folder_detector.py new file mode 100644 index 0000000..add19fa --- /dev/null +++ b/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 diff --git a/face_alignment/detection/sfd/__init__.py b/face_alignment/detection/sfd/__init__.py new file mode 100644 index 0000000..5a63ecd --- /dev/null +++ b/face_alignment/detection/sfd/__init__.py @@ -0,0 +1 @@ +from .sfd_detector import SFDDetector as FaceDetector \ No newline at end of file diff --git a/face_alignment/detection/sfd/bbox.py b/face_alignment/detection/sfd/bbox.py new file mode 100644 index 0000000..3d26696 --- /dev/null +++ b/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 diff --git a/face_alignment/detection/sfd/detect.py b/face_alignment/detection/sfd/detect.py new file mode 100644 index 0000000..63b8318 --- /dev/null +++ b/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]) diff --git a/face_alignment/detection/sfd/net_s3fd.py b/face_alignment/detection/sfd/net_s3fd.py new file mode 100644 index 0000000..ecfa0f6 --- /dev/null +++ b/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] diff --git a/face_alignment/detection/sfd/sfd_detector.py b/face_alignment/detection/sfd/sfd_detector.py new file mode 100644 index 0000000..3799c43 --- /dev/null +++ b/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 diff --git a/face_alignment/folder_data.py b/face_alignment/folder_data.py new file mode 100644 index 0000000..65bc295 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/face_alignment/utils.py b/face_alignment/utils.py new file mode 100644 index 0000000..2618770 --- /dev/null +++ b/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 diff --git a/liveportrait/utils/cropper.py b/liveportrait/utils/cropper.py index 3e0bb00..254bc3b 100644 --- a/liveportrait/utils/cropper.py +++ b/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 \ No newline at end of file + 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 + diff --git a/nodes.py b/nodes.py index 84a1a32..83cdfd8 100644 --- a/nodes.py +++ b/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", } \ No newline at end of file