Compare commits
47
Commits
main
..
kornia_test
| Author | SHA1 | Date | |
|---|---|---|---|
 kijai
|
5e3c92d55c | ||
|
kijai
|
cc0501a2db | ||
|
kijai
|
3dc822fd2f | ||
|
kijai
|
697b9a78e6 | ||
|
kijai
|
2a7bd6116f | ||
|
kijai
|
a7d09f5d49 | ||
|
kijai
|
8e85d5b96d | ||
|
kijai
|
30989a9d37 | ||
|
kijai
|
eecf645603 | ||
|
kijai
|
1b080706df | ||
|
kijai
|
336f3f7c23 | ||
|
kijai
|
f27e1cca13 | ||
|
kijai
|
86e91a6e9d | ||
|
kijai
|
92529f7ca8 | ||
|
kijai
|
c0959056ae | ||
|
kijai
|
2e40fe3820 | ||
|
kijai
|
0a5e187637 | ||
|
kijai
|
4e19dbd6d1 | ||
|
kijai
|
9c190804a7 | ||
|
kijai
|
d9ca40e1d6 | ||
|
kijai
|
b68cf8788c | ||
|
kijai
|
e702b26895 | ||
|
kijai
|
c21705edb5 | ||
|
kijai
|
a284bb52b2 | ||
|
kijai
|
9884aac18a | ||
|
kijai
|
f8aada81db | ||
|
kijai
|
6735771664 | ||
|
kijai
|
857ddbc6d7 | ||
|
kijai
|
a6edcda97d | ||
|
kijai
|
ca01d706d0 | ||
|
kijai
|
0dc9a8a695 | ||
|
Mel Massadian
|
ba6b3f5f68 | ||
|
kijai
|
ee7d5b4241 | ||
|
kijai
|
03df9f35cd | ||
|
kijai
|
ec6b5c8c85 | ||
|
Mel Massadian
|
8509d9a551 | ||
|
Mel Massadian
|
e724da1161 | ||
|
Mel Massadian
|
68d0ddf72a | ||
|
kijai
|
6f9dba7777 | ||
|
kijai
|
811ca557fb | ||
|
kijai
|
6d790bdcc3 | ||
|
kijai
|
ef8b4263b4 | ||
|
Mel Massadian
|
eb5fddf4de | ||
|
Mel Massadian
|
9c7db3c59a | ||
|
Mel Massadian
|
bf3410cd0d | ||
|
Mel Massadian
|
24c65627db | ||
|
Mel Massadian
|
72bb6910e9 |
File diff suppressed because it is too large
Load Diff
@@ -1,593 +0,0 @@
|
||||
{
|
||||
"last_node_id": 208,
|
||||
"last_link_id": 480,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 204,
|
||||
"type": "LivePortraitLoadMediaPipeCropper",
|
||||
"pos": [
|
||||
-1059,
|
||||
-767
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 82
|
||||
},
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "cropper",
|
||||
"type": "LPCROPPER",
|
||||
"links": [
|
||||
478
|
||||
],
|
||||
"shape": 3
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LivePortraitLoadMediaPipeCropper"
|
||||
},
|
||||
"widgets_values": [
|
||||
"CPU",
|
||||
true
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 1,
|
||||
"type": "DownloadAndLoadLivePortraitModels",
|
||||
"pos": [
|
||||
-1046,
|
||||
-904
|
||||
],
|
||||
"size": {
|
||||
"0": 302.43463134765625,
|
||||
"1": 58
|
||||
},
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "live_portrait_pipe",
|
||||
"type": "LIVEPORTRAITPIPE",
|
||||
"links": [
|
||||
446,
|
||||
448
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "DownloadAndLoadLivePortraitModels"
|
||||
},
|
||||
"widgets_values": [
|
||||
"auto"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 165,
|
||||
"type": "ImageResizeKJ",
|
||||
"pos": [
|
||||
-715,
|
||||
-617
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 242
|
||||
},
|
||||
"flags": {},
|
||||
"order": 6,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"link": 466
|
||||
},
|
||||
{
|
||||
"name": "get_image_size",
|
||||
"type": "IMAGE",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "width_input",
|
||||
"type": "INT",
|
||||
"link": null,
|
||||
"widget": {
|
||||
"name": "width_input"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "height_input",
|
||||
"type": "INT",
|
||||
"link": null,
|
||||
"widget": {
|
||||
"name": "height_input"
|
||||
}
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
434
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "width",
|
||||
"type": "INT",
|
||||
"links": null,
|
||||
"shape": 3
|
||||
},
|
||||
{
|
||||
"name": "height",
|
||||
"type": "INT",
|
||||
"links": null,
|
||||
"shape": 3
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "ImageResizeKJ"
|
||||
},
|
||||
"widgets_values": [
|
||||
512,
|
||||
512,
|
||||
"lanczos",
|
||||
true,
|
||||
2,
|
||||
0,
|
||||
0
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 78,
|
||||
"type": "GetImageSizeAndCount",
|
||||
"pos": [
|
||||
-364,
|
||||
-619
|
||||
],
|
||||
"size": {
|
||||
"0": 210,
|
||||
"1": 86
|
||||
},
|
||||
"flags": {},
|
||||
"order": 7,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"link": 434
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
445,
|
||||
475
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "512 width",
|
||||
"type": "INT",
|
||||
"links": null,
|
||||
"shape": 3
|
||||
},
|
||||
{
|
||||
"name": "384 height",
|
||||
"type": "INT",
|
||||
"links": null,
|
||||
"shape": 3
|
||||
},
|
||||
{
|
||||
"name": "1 count",
|
||||
"type": "INT",
|
||||
"links": null,
|
||||
"shape": 3
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "GetImageSizeAndCount"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 207,
|
||||
"type": "Note",
|
||||
"pos": [
|
||||
-850,
|
||||
-202
|
||||
],
|
||||
"size": [
|
||||
230.3331222759989,
|
||||
105.34622315856666
|
||||
],
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"properties": {
|
||||
"text": ""
|
||||
},
|
||||
"widgets_values": [
|
||||
"Example live inputs, direct webcam capture using cv2 or screencapture using mss. Both are about same speed."
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#653"
|
||||
},
|
||||
{
|
||||
"id": 206,
|
||||
"type": "Screencap_mss",
|
||||
"pos": [
|
||||
-568,
|
||||
-34
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 178
|
||||
},
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"links": null,
|
||||
"shape": 3
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "Screencap_mss"
|
||||
},
|
||||
"widgets_values": [
|
||||
0,
|
||||
0,
|
||||
512,
|
||||
512,
|
||||
1,
|
||||
0.1
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 198,
|
||||
"type": "PreviewImage",
|
||||
"pos": [
|
||||
413,
|
||||
-826
|
||||
],
|
||||
"size": {
|
||||
"0": 521.2196044921875,
|
||||
"1": 566.1187133789062
|
||||
},
|
||||
"flags": {},
|
||||
"order": 10,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 470
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "PreviewImage"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 196,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
-1058,
|
||||
-623
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 314
|
||||
},
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
466
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": null,
|
||||
"shape": 3
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"Mona-Lisa-oil-wood-panel-Leonardo-da.webp",
|
||||
"image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 205,
|
||||
"type": "WebcamCaptureCV2",
|
||||
"pos": [
|
||||
-577,
|
||||
-283
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 178
|
||||
},
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
479
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "WebcamCaptureCV2"
|
||||
},
|
||||
"widgets_values": [
|
||||
0,
|
||||
0,
|
||||
512,
|
||||
512,
|
||||
0,
|
||||
false
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 190,
|
||||
"type": "LivePortraitProcess",
|
||||
"pos": [
|
||||
-79,
|
||||
-552
|
||||
],
|
||||
"size": {
|
||||
"0": 430.8000183105469,
|
||||
"1": 282
|
||||
},
|
||||
"flags": {},
|
||||
"order": 9,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "pipeline",
|
||||
"type": "LIVEPORTRAITPIPE",
|
||||
"link": 448
|
||||
},
|
||||
{
|
||||
"name": "crop_info",
|
||||
"type": "CROPINFO",
|
||||
"link": 449
|
||||
},
|
||||
{
|
||||
"name": "source_image",
|
||||
"type": "IMAGE",
|
||||
"link": 475
|
||||
},
|
||||
{
|
||||
"name": "driving_images",
|
||||
"type": "IMAGE",
|
||||
"link": 479,
|
||||
"slot_index": 3
|
||||
},
|
||||
{
|
||||
"name": "opt_retargeting_info",
|
||||
"type": "RETARGETINGINFO",
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "cropped_image",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
470
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "output",
|
||||
"type": "LP_OUT",
|
||||
"links": [],
|
||||
"shape": 3,
|
||||
"slot_index": 1
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LivePortraitProcess"
|
||||
},
|
||||
"widgets_values": [
|
||||
false,
|
||||
0.03,
|
||||
true,
|
||||
1,
|
||||
"constant",
|
||||
"single_frame",
|
||||
0.000003
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 189,
|
||||
"type": "LivePortraitCropper",
|
||||
"pos": [
|
||||
-73,
|
||||
-876
|
||||
],
|
||||
"size": {
|
||||
"0": 330,
|
||||
"1": 242
|
||||
},
|
||||
"flags": {},
|
||||
"order": 8,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "pipeline",
|
||||
"type": "LIVEPORTRAITPIPE",
|
||||
"link": 446,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "cropper",
|
||||
"type": "LPCROPPER",
|
||||
"link": 478,
|
||||
"slot_index": 1
|
||||
},
|
||||
{
|
||||
"name": "source_image",
|
||||
"type": "IMAGE",
|
||||
"link": 445
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "cropped_image",
|
||||
"type": "IMAGE",
|
||||
"links": null,
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "crop_info",
|
||||
"type": "CROPINFO",
|
||||
"links": [
|
||||
449
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 1
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LivePortraitCropper"
|
||||
},
|
||||
"widgets_values": [
|
||||
512,
|
||||
2.34,
|
||||
0.099,
|
||||
0.148,
|
||||
0,
|
||||
"large-small",
|
||||
false
|
||||
]
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
434,
|
||||
165,
|
||||
0,
|
||||
78,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
445,
|
||||
78,
|
||||
0,
|
||||
189,
|
||||
2,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
446,
|
||||
1,
|
||||
0,
|
||||
189,
|
||||
0,
|
||||
"LIVEPORTRAITPIPE"
|
||||
],
|
||||
[
|
||||
448,
|
||||
1,
|
||||
0,
|
||||
190,
|
||||
0,
|
||||
"LIVEPORTRAITPIPE"
|
||||
],
|
||||
[
|
||||
449,
|
||||
189,
|
||||
1,
|
||||
190,
|
||||
1,
|
||||
"CROPINFO"
|
||||
],
|
||||
[
|
||||
466,
|
||||
196,
|
||||
0,
|
||||
165,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
470,
|
||||
190,
|
||||
0,
|
||||
198,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
475,
|
||||
78,
|
||||
0,
|
||||
190,
|
||||
2,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
478,
|
||||
204,
|
||||
0,
|
||||
189,
|
||||
1,
|
||||
"LPCROPPER"
|
||||
],
|
||||
[
|
||||
479,
|
||||
205,
|
||||
0,
|
||||
190,
|
||||
3,
|
||||
"IMAGE"
|
||||
]
|
||||
],
|
||||
"groups": [],
|
||||
"config": {},
|
||||
"extra": {
|
||||
"ds": {
|
||||
"scale": 0.7513148009015781,
|
||||
"offset": {
|
||||
"0": 1468.4081568988054,
|
||||
"1": 1224.8414164288351
|
||||
}
|
||||
}
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,29 +0,0 @@
|
||||
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.
|
||||
@@ -1,7 +0,0 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
__author__ = """Adrian Bulat"""
|
||||
__email__ = 'adrian@adrianbulat.com'
|
||||
__version__ = '1.4.1'
|
||||
|
||||
from .api import FaceAlignment, LandmarksType, NetworkSize
|
||||
@@ -1,305 +0,0 @@
|
||||
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
|
||||
@@ -1 +0,0 @@
|
||||
from .core import FaceDetector
|
||||
@@ -1 +0,0 @@
|
||||
from .blazeface_detector import BlazeFaceDetector as FaceDetector
|
||||
@@ -1,69 +0,0 @@
|
||||
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
|
||||
@@ -1,74 +0,0 @@
|
||||
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])
|
||||
@@ -1,482 +0,0 @@
|
||||
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)
|
||||
@@ -1,65 +0,0 @@
|
||||
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)
|
||||
@@ -1,150 +0,0 @@
|
||||
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
|
||||
@@ -1 +0,0 @@
|
||||
from .dlib_detector import DlibDetector as FaceDetector
|
||||
@@ -1,48 +0,0 @@
|
||||
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
|
||||
@@ -1 +0,0 @@
|
||||
from .folder_detector import FolderDetector as FaceDetector
|
||||
@@ -1,53 +0,0 @@
|
||||
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
|
||||
@@ -1 +0,0 @@
|
||||
from .sfd_detector import SFDDetector as FaceDetector
|
||||
@@ -1,71 +0,0 @@
|
||||
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
|
||||
@@ -1,85 +0,0 @@
|
||||
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])
|
||||
@@ -1,127 +0,0 @@
|
||||
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]
|
||||
@@ -1,71 +0,0 @@
|
||||
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
|
||||
@@ -1,38 +0,0 @@
|
||||
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)
|
||||
@@ -1,379 +0,0 @@
|
||||
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
|
||||
@@ -35,7 +35,10 @@ class InferenceConfig(PrintableConfig):
|
||||
output_fps: int = 30 # fps for output video
|
||||
crf: int = 15 # crf for output video
|
||||
|
||||
flag_pasteback: bool = True # whether to paste-back/stitch the animated face cropping from the face-cropping space to the original image space
|
||||
mask_crop = None
|
||||
flag_write_gif: bool = False
|
||||
|
||||
device_id: int = 0
|
||||
flag_do_crop: bool = False # whether to crop the reference portrait to the face-cropping space
|
||||
flag_do_rot: bool = True # whether to conduct the rotation when flag_do_crop is True
|
||||
|
||||
@@ -5,16 +5,15 @@ Pipeline of LivePortrait
|
||||
"""
|
||||
|
||||
import comfy.utils
|
||||
import comfy.model_management as mm
|
||||
import gc
|
||||
from tqdm import tqdm
|
||||
import numpy as np
|
||||
from .config.inference_config import InferenceConfig
|
||||
import torch
|
||||
from .utils.camera import get_rotation_matrix
|
||||
from .utils.crop import _transform_img, _transform_img_kornia
|
||||
from .live_portrait_wrapper import LivePortraitWrapper
|
||||
from .utils.retargeting_utils import calc_eye_close_ratio, calc_lip_close_ratio
|
||||
from .utils.filter import smooth
|
||||
from .utils.helper import calc_motion_multiplier
|
||||
|
||||
import os
|
||||
script_directory = os.path.dirname(os.path.abspath(__file__))
|
||||
@@ -38,46 +37,63 @@ class LivePortraitPipeline(object):
|
||||
cfg=inference_cfg,
|
||||
)
|
||||
|
||||
def _get_source_frame(self, source_np, idx, method):
|
||||
if source_np.shape[0] == 1:
|
||||
return source_np[0]
|
||||
|
||||
if method == "constant":
|
||||
return source_np[min(idx, source_np.shape[0] - 1)]
|
||||
elif method == "cycle":
|
||||
return source_np[idx % source_np.shape[0]]
|
||||
elif method == "mirror":
|
||||
cycle_length = 2 * source_np.shape[0] - 2
|
||||
mirror_idx = idx % cycle_length
|
||||
if mirror_idx >= source_np.shape[0]:
|
||||
mirror_idx = cycle_length - mirror_idx
|
||||
return source_np[mirror_idx]
|
||||
|
||||
def execute(
|
||||
self, driving_images, crop_info, driving_landmarks, delta_multiplier, relative_motion_mode, driving_smooth_observation_variance, mismatch_method="constant", expression_friendly=False, driving_multiplier=1.0,
|
||||
self, source_np, driving_images, crop_info, driving_landmarks, delta_multiplier, relative_motion_mode, driving_smooth_observation_variance, mismatch_method="constant",
|
||||
):
|
||||
inference_cfg = self.live_portrait_wrapper.cfg
|
||||
device = inference_cfg.device_id
|
||||
|
||||
out_list = []
|
||||
cropped_image_list = []
|
||||
composited_image_list = []
|
||||
out_mask_list = []
|
||||
R_d_0, x_d_0_info = None, None
|
||||
|
||||
source_images_num = len(crop_info["crop_info_list"])
|
||||
|
||||
if mismatch_method == "cut" or relative_motion_mode == "source_video_smoothed":
|
||||
total_frames = source_images_num
|
||||
if mismatch_method == "cut":
|
||||
total_frames = source_np.shape[0]
|
||||
else:
|
||||
total_frames = driving_images.shape[0]
|
||||
|
||||
|
||||
disable_progress_bar = True if relative_motion_mode == "single_frame" else False
|
||||
source_info = []
|
||||
source_rot_list = []
|
||||
f_s_list = []
|
||||
for i in tqdm(range(source_np.shape[0]), desc='Processing source images...', total=source_np.shape[0]):
|
||||
#get source keypoints info
|
||||
img_crop_256x256 = crop_info["crop_info_list"][i]["img_crop_256x256"]
|
||||
I_s = self.live_portrait_wrapper.prepare_source(img_crop_256x256)
|
||||
x_s_info = self.live_portrait_wrapper.get_kp_info(I_s)
|
||||
f_s = self.live_portrait_wrapper.extract_feature_3d(I_s)
|
||||
f_s_list.append(f_s)
|
||||
source_info.append(x_s_info)
|
||||
|
||||
source_info = crop_info["source_info"]
|
||||
source_rot_list = crop_info["source_rot_list"]
|
||||
f_s_list = crop_info["f_s_list"]
|
||||
x_s_list = crop_info["x_s_list"]
|
||||
R_s = get_rotation_matrix(
|
||||
x_s_info["pitch"], x_s_info["yaw"], x_s_info["roll"]
|
||||
)
|
||||
source_rot_list.append(R_s)
|
||||
|
||||
driving_info = []
|
||||
driving_exp_list = []
|
||||
driving_rot_list = []
|
||||
|
||||
for i in tqdm(range(driving_images.shape[0]), desc='Processing driving images...', total=driving_images.shape[0], disable=disable_progress_bar):
|
||||
for i in tqdm(range(driving_images.shape[0]), desc='Processing driving images...', total=driving_images.shape[0]):
|
||||
#get driving keypoints info
|
||||
safe_index = min(i, source_images_num - 1)
|
||||
if crop_info["crop_info_list"][safe_index] is None:
|
||||
driving_info.append(None)
|
||||
driving_rot_list.append(None)
|
||||
driving_exp_list.append(None)
|
||||
if i == 0:
|
||||
raise ValueError("No face detected in FIRST source image")
|
||||
continue
|
||||
x_d_info = self.live_portrait_wrapper.get_kp_info(driving_images[i].unsqueeze(0).to(device))
|
||||
|
||||
safe_index = min(i, len(crop_info["crop_info_list"]) - 1)
|
||||
if i == 0:
|
||||
first = x_d_info
|
||||
|
||||
@@ -94,10 +110,7 @@ class LivePortraitPipeline(object):
|
||||
if relative_motion_mode == "source_video_smoothed":
|
||||
x_d_r_lst = []
|
||||
first_driving_rot = driving_rot_list[0].cpu().numpy().astype(np.float32).transpose(0, 2, 1)
|
||||
for i in tqdm(range(source_images_num), desc='Smoothing...', total=source_images_num):
|
||||
if driving_rot_list[i] is None:
|
||||
x_d_r_lst.append(None)
|
||||
continue
|
||||
for i in tqdm(range(source_np.shape[0]), desc='Smoothing...', total=source_np.shape[0]):
|
||||
driving_rot = driving_rot_list[i].cpu().numpy().astype(np.float32)
|
||||
source_rot = source_rot_list[i].cpu().numpy().astype(np.float32)
|
||||
dot = np.dot(driving_rot, first_driving_rot) @ source_rot
|
||||
@@ -108,28 +121,28 @@ class LivePortraitPipeline(object):
|
||||
|
||||
pbar = comfy.utils.ProgressBar(total_frames)
|
||||
|
||||
for i in tqdm(range(total_frames), desc='Animating...', total=total_frames, disable=disable_progress_bar):
|
||||
for i in tqdm(range(total_frames), desc='Animating...', total=total_frames):
|
||||
|
||||
safe_index = min(i, len(crop_info["crop_info_list"]) - 1)
|
||||
|
||||
# skip and return empty frames if no crop due to no face detected
|
||||
if crop_info["crop_info_list"][safe_index] is None:
|
||||
out_list.append({})
|
||||
pbar.update(1)
|
||||
if not crop_info["crop_info_list"][safe_index]:
|
||||
composited_image_list.append(source_np[safe_index])
|
||||
cropped_image_list.append(torch.zeros(1, 512, 512, 3, dtype=torch.float32, device = device))
|
||||
out_mask_list.append(np.zeros((source_np.shape[1], source_np.shape[2], 3), dtype=np.uint8))
|
||||
continue
|
||||
|
||||
source_lmk = crop_info["crop_info_list"][safe_index]["lmk_crop"]
|
||||
|
||||
x_d_info = driving_info[i]
|
||||
R_d = driving_rot_list[i]
|
||||
|
||||
x_s_info = source_info[safe_index]
|
||||
R_s = source_rot_list[safe_index]
|
||||
f_s = f_s_list[safe_index]
|
||||
x_s = x_s_list[safe_index]
|
||||
|
||||
x_c_s = x_s_info["kp"]
|
||||
|
||||
R_s = source_rot_list[safe_index]
|
||||
f_s = f_s_list[safe_index]
|
||||
x_s = self.live_portrait_wrapper.transform_keypoint(x_s_info)
|
||||
|
||||
#lip zero
|
||||
if inference_cfg.flag_lip_zero:
|
||||
c_d_lip_before_animation = [0.0]
|
||||
@@ -140,10 +153,13 @@ class LivePortraitPipeline(object):
|
||||
else:
|
||||
lip_delta_before_animation = (self.live_portrait_wrapper.retarget_lip(x_s, combined_lip_ratio_tensor_before_animation))
|
||||
|
||||
R_d = driving_rot_list[i]
|
||||
|
||||
if i == 0:
|
||||
R_d_0 = R_d
|
||||
x_d_0_info = x_d_info
|
||||
|
||||
if relative_motion_mode == "relative":
|
||||
if i == 0:
|
||||
R_d_0 = R_d
|
||||
x_d_0_info = x_d_info
|
||||
R_new = (R_d @ R_d_0.permute(0, 2, 1)) @ R_s
|
||||
delta_new = x_s_info["exp"] + (x_d_info["exp"] - x_d_0_info["exp"])
|
||||
scale_new = x_s_info["scale"] * (x_d_info["scale"] / x_d_0_info["scale"])
|
||||
@@ -158,11 +174,6 @@ class LivePortraitPipeline(object):
|
||||
delta_new = x_s_info['exp']
|
||||
scale_new = x_s_info["scale"]
|
||||
t_new = x_d_info["t"]
|
||||
elif relative_motion_mode == "single_frame":
|
||||
R_new = R_d
|
||||
delta_new = x_d_info['exp']
|
||||
scale_new = x_s_info["scale"]
|
||||
t_new = x_d_info["t"]
|
||||
else:
|
||||
R_new = R_d
|
||||
delta_new = x_s_info['exp']
|
||||
@@ -174,20 +185,11 @@ class LivePortraitPipeline(object):
|
||||
delta_new = delta_new * delta_multiplier
|
||||
|
||||
x_d_i_new = scale_new * (x_c_s @ R_new + delta_new) + t_new
|
||||
|
||||
if expression_friendly:
|
||||
if i == 0:
|
||||
x_d_0_new = x_d_i_new
|
||||
motion_multiplier = calc_motion_multiplier(x_s, x_d_0_new)
|
||||
motion_multiplier *= driving_multiplier
|
||||
x_d_diff = (x_d_i_new - x_d_0_new) * motion_multiplier
|
||||
x_d_i_new = x_d_diff + x_s
|
||||
|
||||
if (
|
||||
not inference_cfg.flag_stitching
|
||||
and not inference_cfg.flag_eye_retargeting
|
||||
and not inference_cfg.flag_lip_retargeting
|
||||
):
|
||||
):
|
||||
# without stitching or retargeting
|
||||
if inference_cfg.flag_lip_zero:
|
||||
x_d_i_new += lip_delta_before_animation.reshape(-1, x_s.shape[1], 3)
|
||||
@@ -197,7 +199,7 @@ class LivePortraitPipeline(object):
|
||||
inference_cfg.flag_stitching
|
||||
and not inference_cfg.flag_eye_retargeting
|
||||
and not inference_cfg.flag_lip_retargeting
|
||||
):
|
||||
):
|
||||
# with stitching and without retargeting
|
||||
if inference_cfg.flag_lip_zero:
|
||||
x_d_i_new = self.live_portrait_wrapper.stitching(
|
||||
@@ -205,8 +207,6 @@ class LivePortraitPipeline(object):
|
||||
) + lip_delta_before_animation.reshape(-1, x_s.shape[1], 3)
|
||||
else:
|
||||
x_d_i_new = self.live_portrait_wrapper.stitching(x_s, x_d_i_new)
|
||||
|
||||
#with eye/lip retargeting
|
||||
else:
|
||||
eyes_delta, lip_delta = None, None
|
||||
if inference_cfg.flag_eye_retargeting:
|
||||
@@ -240,7 +240,7 @@ class LivePortraitPipeline(object):
|
||||
x_s, combined_lip_ratio_tensor
|
||||
)
|
||||
|
||||
if relative_motion_mode != "off": # use x_s
|
||||
if inference_cfg.flag_relative: # use x_s
|
||||
x_d_i_new = (
|
||||
x_s
|
||||
+ (
|
||||
@@ -276,15 +276,38 @@ class LivePortraitPipeline(object):
|
||||
x_d_i_new = self.live_portrait_wrapper.stitching(x_s, x_d_i_new)
|
||||
|
||||
out = self.live_portrait_wrapper.warp_decode(f_s, x_s, x_d_i_new)
|
||||
|
||||
out_list.append(out)
|
||||
|
||||
cropped_image = torch.clamp(out["out"], 0, 1).permute(0, 2, 3, 1)
|
||||
|
||||
cropped_image_list.append(cropped_image)
|
||||
|
||||
if mismatch_method == "cut" or inference_cfg.flag_eye_retargeting or inference_cfg.flag_lip_retargeting:
|
||||
source_frame_rgb = source_np[safe_index]
|
||||
else:
|
||||
source_frame_rgb = self._get_source_frame(source_np, i, mismatch_method)
|
||||
|
||||
# Transform and blend
|
||||
if inference_cfg.flag_pasteback:
|
||||
cropped_image_to_original = _transform_img_kornia(
|
||||
cropped_image,
|
||||
crop_info["crop_info_list"][safe_index]["M_c2o"],
|
||||
dsize=(source_frame_rgb.shape[1], source_frame_rgb.shape[0]),
|
||||
)
|
||||
|
||||
mask_ori = _transform_img_kornia(
|
||||
inference_cfg.mask_crop,
|
||||
crop_info["crop_info_list"][safe_index]["M_c2o"],
|
||||
dsize=(source_frame_rgb.shape[1], source_frame_rgb.shape[0]),
|
||||
)
|
||||
|
||||
source_frame_torch = torch.from_numpy(source_frame_rgb).unsqueeze(0).permute(0, 3, 1, 2).to(mask_ori.device) / 255
|
||||
|
||||
cropped_image_to_original_blend = torch.clip(
|
||||
mask_ori * cropped_image_to_original + (1 - mask_ori) * source_frame_torch, 0, 1
|
||||
)
|
||||
|
||||
composited_image_list.append(cropped_image_to_original_blend)
|
||||
out_mask_list.append(mask_ori)
|
||||
pbar.update(1)
|
||||
|
||||
out_dict = {
|
||||
"out_list": out_list,
|
||||
"crop_info": crop_info,
|
||||
"mismatch_method": mismatch_method,
|
||||
}
|
||||
|
||||
return out_dict
|
||||
return cropped_image_list, composited_image_list, out_mask_list
|
||||
|
||||
@@ -170,6 +170,35 @@ class LivePortraitWrapper(object):
|
||||
|
||||
return delta
|
||||
|
||||
def retarget_keypoints(self, frame_idx, num_keypoints, input_eye_ratios, input_lip_ratios, source_landmarks, portrait_wrapper, kp_source, driving_transformed_kp):
|
||||
# TODO: GPT style, refactor it...
|
||||
if self.cfg.flag_eye_retargeting:
|
||||
print("Retargeting eye...")
|
||||
# ∆_eyes,i = R_eyes(x_s; c_s,eyes, c_d,eyes,i)
|
||||
eye_delta = compute_eye_delta(frame_idx, input_eye_ratios, source_landmarks, portrait_wrapper, kp_source)
|
||||
else:
|
||||
# α_eyes = 0
|
||||
eye_delta = None
|
||||
|
||||
if self.cfg.flag_lip_retargeting:
|
||||
print("Retargeting lip...")
|
||||
# ∆_lip,i = R_lip(x_s; c_s,lip, c_d,lip,i)
|
||||
lip_delta = compute_lip_delta(frame_idx, input_lip_ratios, source_landmarks, portrait_wrapper, kp_source)
|
||||
else:
|
||||
# α_lip = 0
|
||||
lip_delta = None
|
||||
|
||||
if self.cfg.flag_relative: # use x_s
|
||||
new_driving_kp = kp_source + \
|
||||
(eye_delta.reshape(-1, num_keypoints, 3) if eye_delta is not None else 0) + \
|
||||
(lip_delta.reshape(-1, num_keypoints, 3) if lip_delta is not None else 0)
|
||||
else: # use x_d,i
|
||||
new_driving_kp = driving_transformed_kp + \
|
||||
(eye_delta.reshape(-1, num_keypoints, 3) if eye_delta is not None else 0) + \
|
||||
(lip_delta.reshape(-1, num_keypoints, 3) if lip_delta is not None else 0)
|
||||
|
||||
return new_driving_kp
|
||||
|
||||
def stitch(self, kp_source: torch.Tensor, kp_driving: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
kp_source: BxNx3
|
||||
@@ -220,14 +249,23 @@ class LivePortraitWrapper(object):
|
||||
ret_dct['out'] = self.spade_generator(feature=ret_dct['out'])
|
||||
|
||||
# float the dict
|
||||
for k, v in ret_dct.items():
|
||||
if isinstance(v, torch.Tensor):
|
||||
ret_dct[k] = v.cpu()
|
||||
if self.cfg.flag_use_half_precision:
|
||||
ret_dct[k] = ret_dct[k].float()
|
||||
if self.cfg.flag_use_half_precision:
|
||||
for k, v in ret_dct.items():
|
||||
if isinstance(v, torch.Tensor):
|
||||
ret_dct[k] = v.float()
|
||||
|
||||
return ret_dct
|
||||
|
||||
def parse_output(self, out: torch.Tensor) -> np.ndarray:
|
||||
""" construct the output as standard
|
||||
return: 1xHxWx3, uint8
|
||||
"""
|
||||
out = np.transpose(out.data.cpu().numpy(), [0, 2, 3, 1]) # 1x3xHxW -> 1xHxWx3
|
||||
out = np.clip(out, 0, 1) # clip to 0~1
|
||||
out = np.clip(out * 255, 0, 255).astype(np.uint8) # 0~1 -> 0~255
|
||||
|
||||
return out
|
||||
|
||||
def calc_retargeting_ratio(self, source_lmk, driving_lmk_lst):
|
||||
input_eye_ratio_lst = []
|
||||
input_lip_ratio_lst = []
|
||||
|
||||
@@ -47,13 +47,7 @@ class DenseMotionNetwork(nn.Module):
|
||||
feature_repeat = feature.unsqueeze(1).unsqueeze(1).repeat(1, self.num_kp+1, 1, 1, 1, 1, 1) # (bs, num_kp+1, 1, c, d, h, w)
|
||||
feature_repeat = feature_repeat.view(bs * (self.num_kp+1), -1, d, h, w) # (bs*(num_kp+1), c, d, h, w)
|
||||
sparse_motions = sparse_motions.view((bs * (self.num_kp+1), d, h, w, -1)) # (bs*(num_kp+1), d, h, w, 3)
|
||||
try:
|
||||
sparse_deformed = F.grid_sample(feature_repeat, sparse_motions, align_corners=False)
|
||||
except NotImplementedError: #MPS fallback
|
||||
out_device = feature_repeat.device # Store input device
|
||||
feature_repeat = feature_repeat.to('cpu')
|
||||
sparse_motions = sparse_motions.to('cpu')
|
||||
sparse_deformed = F.grid_sample(feature_repeat, sparse_motions, align_corners=False).to(out_device)
|
||||
sparse_deformed = F.grid_sample(feature_repeat, sparse_motions, align_corners=False)
|
||||
sparse_deformed = sparse_deformed.view((bs, self.num_kp+1, -1, d, h, w)) # (bs, num_kp+1, c, d, h, w)
|
||||
|
||||
return sparse_deformed
|
||||
@@ -67,7 +61,7 @@ class DenseMotionNetwork(nn.Module):
|
||||
# adding background feature
|
||||
try:
|
||||
zeros = torch.zeros(heatmap.shape[0], 1, spatial_size[0], spatial_size[1], spatial_size[2]).type(heatmap.type()).to(heatmap.device)
|
||||
except ValueError:
|
||||
except:
|
||||
zeros = torch.zeros(heatmap.shape[0], 1, spatial_size[0], spatial_size[1], spatial_size[2]).to(heatmap.device)
|
||||
heatmap = torch.cat([zeros, heatmap], dim=1)
|
||||
heatmap = heatmap.unsqueeze(2) # (bs, 1+num_kp, 1, d, h, w)
|
||||
|
||||
@@ -158,11 +158,7 @@ class DownBlock3d(nn.Module):
|
||||
out = self.conv(x)
|
||||
out = self.norm(out)
|
||||
out = F.relu(out)
|
||||
try:
|
||||
out = self.pool(out)
|
||||
except NotImplementedError:
|
||||
out_device = out.device # Store input device
|
||||
out = self.pool(out.to('cpu')).to(out_device)
|
||||
out = self.pool(out)
|
||||
return out
|
||||
|
||||
|
||||
|
||||
@@ -44,11 +44,7 @@ class WarpingNetwork(nn.Module):
|
||||
self.estimate_occlusion_map = estimate_occlusion_map
|
||||
|
||||
def deform_input(self, inp, deformation):
|
||||
try:
|
||||
return F.grid_sample(inp, deformation, align_corners=False)
|
||||
except NotImplementedError:
|
||||
out_device = inp.device # Store input device
|
||||
return F.grid_sample(inp.to('cpu'), deformation.to('cpu'), align_corners=False).to(out_device)
|
||||
return F.grid_sample(inp, deformation, align_corners=False)
|
||||
|
||||
def forward(self, feature_3d, kp_driving, kp_source):
|
||||
if self.dense_motion_network is not None:
|
||||
|
||||
+31
-75
@@ -30,30 +30,36 @@ def _transform_img(img, M, dsize, flags=CV2_INTERP, borderMode=None):
|
||||
import torch
|
||||
import kornia.geometry.transform as KGT
|
||||
|
||||
def _transform_img_kornia(img, M, dsize, device, flags='bilinear', borderMode='zeros'):
|
||||
def _transform_img_kornia(img, M, dsize, flags='bilinear', borderMode='zeros'):
|
||||
"""Conduct similarity or affine transformation to the image using Kornia.
|
||||
|
||||
img: Input image as a PyTorch tensor of shape (C, H, W).
|
||||
M: 2x3 transformation matrix as a PyTorch tensor.
|
||||
dsize: Target shape (width, height).
|
||||
"""
|
||||
|
||||
device = mm.get_torch_device()
|
||||
# Convert dsize to tensor shape (H, W)
|
||||
_dsize = torch.tensor([dsize[1], dsize[0]]) # Kornia expects (H, W)
|
||||
|
||||
# Convert M from numpy.ndarray to PyTorch tensor
|
||||
M = torch.from_numpy(M).float().to(device)
|
||||
|
||||
if M.shape == (3, 3):
|
||||
M = M[:2, :]
|
||||
M = M[:2, :].unsqueeze(0) # Adjust M to the expected shape Bx2x3
|
||||
elif M.shape == (2, 3):
|
||||
M = M.unsqueeze(0) # Add batch dimension if not present
|
||||
|
||||
# Reshape M for Kornia (1, 2, 3) and upscale to 3D affine matrix if not already
|
||||
if M.shape == (2, 3):
|
||||
M = M.unsqueeze(0)
|
||||
M = M.unsqueeze(0) # Add batch dimension
|
||||
|
||||
# Convert image to floating point tensor if not already
|
||||
if img.dtype != torch.float32:
|
||||
img = img.float()
|
||||
img = img.permute(0, 3, 1, 2).to(device) # Reshape img for Kornia (B, C, H, W)
|
||||
|
||||
img = img.to(device)
|
||||
|
||||
# Reshape img for Kornia (B, C, H, W)
|
||||
img = img.permute(0, 3, 1, 2)
|
||||
|
||||
# Apply the affine transformation
|
||||
img_warped = KGT.warp_affine(img, M, _dsize, mode=flags, padding_mode=borderMode)
|
||||
|
||||
@@ -68,23 +74,6 @@ def _transform_pts(pts, M):
|
||||
return pts @ M[:2, :2].T + M[:2, 2]
|
||||
|
||||
|
||||
def parse_pt2_from_pt478(pt478, use_lip=True):
|
||||
"""
|
||||
parsing the 2 points according to the 101 points, which cancels the roll
|
||||
"""
|
||||
# the former version use the eye center, but it is not robust, now use interpolation
|
||||
pt_left_eye = pt478[468] # left eye center
|
||||
pt_right_eye = pt478[473] # right eye center
|
||||
|
||||
if use_lip:
|
||||
# use lip
|
||||
pt_center_eye = (pt_left_eye + pt_right_eye) / 2
|
||||
pt_center_lip = pt478[14]
|
||||
pt2 = np.stack([pt_center_eye, pt_center_lip], axis=0)
|
||||
else:
|
||||
pt2 = np.stack([pt_left_eye, pt_right_eye], axis=0)
|
||||
return pt2
|
||||
|
||||
def parse_pt2_from_pt101(pt101, use_lip=True):
|
||||
"""
|
||||
parsing the 2 points according to the 101 points, which cancels the roll
|
||||
@@ -135,60 +124,29 @@ def parse_pt2_from_pt203(pt203, use_lip=True):
|
||||
pt2 = np.stack([pt_left_eye, pt_right_eye], axis=0)
|
||||
return pt2
|
||||
|
||||
def parse_pt2_from_pt9(pt9, use_lip=True):
|
||||
'''
|
||||
animal_face = {"keypoints": ['right eye right', 'right eye left', 'left eye right', 'left eye left', 'nose tip', 'lip right', 'lip left', 'upper lip', 'lower lip'], "skeleton": []}
|
||||
|
||||
|
||||
'''
|
||||
if use_lip:
|
||||
pt9 = np.stack([
|
||||
(pt9[2]+pt9[3])/2, # left eye
|
||||
(pt9[0]+pt9[1])/2, # right eye
|
||||
pt9[4],
|
||||
# (pt9[5]+pt9[6]+pt9[7]+pt9[8])/4 # lip
|
||||
(pt9[5] + pt9[6] ) / 2 # lip
|
||||
], axis=0)
|
||||
pt2 = np.stack([
|
||||
(pt9[0] + pt9[1]) / 2, # eye
|
||||
pt9[3] # lip
|
||||
], axis=0)
|
||||
else:
|
||||
pt2 = np.stack([
|
||||
(pt9[2] + pt9[3]) / 2,
|
||||
(pt9[0] + pt9[1]) / 2,
|
||||
], axis=0)
|
||||
|
||||
return pt2
|
||||
|
||||
def parse_pt2_from_pt68(pt68, use_lip=True):
|
||||
'''
|
||||
face = {"keypoints": ['right cheekbone 1', 'right cheekbone 2', 'right cheek 1', 'right cheek 2', 'right cheek 3', 'right cheek 4', 'right cheek 5', 'right chin', 'chin center',
|
||||
'left chin', 'left cheek 5', 'left cheek 4', 'left cheek 3', 'left cheek 2', 'left cheek 1', 'left cheekbone 2', 'left cheekbone 1', 'right eyebrow 1', 'right eyebrow 2', 'right eyebrow 3',
|
||||
'right eyebrow 4', 'right eyebrow 5', 'left eyebrow 1', 'left eyebrow 2', 'left eyebrow 3', 'left eyebrow 4', 'left eyebrow 5', 'nasal bridge 1', 'nasal bridge 2', 'nasal bridge 3', 'nasal bridge 4',
|
||||
'right nasal wing 1', 'right nasal wing 2', 'nasal wing center', 'left nasal wing 1', 'left nasal wing 2', 'right eye eye corner 1', 'right eye upper eyelid 1', 'right eye upper eyelid 2',
|
||||
'right eye eye corner 2', 'right eye lower eyelid 2', 'right eye lower eyelid 1', 'left eye eye corner 1', 'left eye upper eyelid 1', 'left eye upper eyelid 2', 'left eye eye corner 2', 'left eye lower eyelid 2',
|
||||
'left eye lower eyelid 1', 'right mouth corner', 'upper lip outer edge 1', 'upper lip outer edge 2', 'upper lip outer edge 3', 'upper lip outer edge 4', 'upper lip outer edge 5', 'left mouth corner',
|
||||
'lower lip outer edge 5', 'lower lip outer edge 4', 'lower lip outer edge 3', 'lower lip outer edge 2', 'lower lip outer edge 1', 'upper lip inter edge 1', 'upper lip inter edge 2', 'upper lip inter edge 3',
|
||||
'upper lip inter edge 4', 'upper lip inter edge 5', 'lower lip inter edge 3', 'lower lip inter edge 2', 'lower lip inter edge 1'], "skeleton": []}
|
||||
|
||||
|
||||
'''
|
||||
"""
|
||||
parsing the 2 points according to the 68 points, which cancels the roll
|
||||
"""
|
||||
lm_idx = np.array([31, 37, 40, 43, 46, 49, 55], dtype=np.int32) - 1
|
||||
if use_lip:
|
||||
pt68 = np.stack([
|
||||
(pt68[42] + pt68[43] + pt68[44] + pt68[45] + pt68[46]+ pt68[47])/6, # left eye
|
||||
(pt68[36] + pt68[37] + pt68[38] + pt68[39] + pt68[40] + pt68[41]) / 6, # right eye
|
||||
(pt68[48] + pt68[54])/2
|
||||
|
||||
pt5 = np.stack([
|
||||
np.mean(pt68[lm_idx[[1, 2]], :], 0), # left eye
|
||||
np.mean(pt68[lm_idx[[3, 4]], :], 0), # right eye
|
||||
pt68[lm_idx[0], :], # nose
|
||||
pt68[lm_idx[5], :], # lip
|
||||
pt68[lm_idx[6], :] # lip
|
||||
], axis=0)
|
||||
|
||||
pt2 = np.stack([
|
||||
(pt68[0] + pt68[1]) / 2,
|
||||
pt68[2]
|
||||
(pt5[0] + pt5[1]) / 2,
|
||||
(pt5[3] + pt5[4]) / 2
|
||||
], axis=0)
|
||||
else:
|
||||
pt2 = np.stack([
|
||||
(pt68[42] + pt68[43] + pt68[44] + pt68[45] + pt68[46] + pt68[47]) / 6, # left eye
|
||||
(pt68[36] + pt68[37] + pt68[38] + pt68[39] + pt68[40] + pt68[41]) / 6, # right eye
|
||||
np.mean(pt68[lm_idx[[1, 2]], :], 0), # left eye
|
||||
np.mean(pt68[lm_idx[[3, 4]], :], 0), # right eye
|
||||
], axis=0)
|
||||
|
||||
return pt2
|
||||
@@ -222,13 +180,9 @@ def parse_pt2_from_pt_x(pts, use_lip=True):
|
||||
pt2 = parse_pt2_from_pt5(pts, use_lip=use_lip)
|
||||
elif pts.shape[0] == 203:
|
||||
pt2 = parse_pt2_from_pt203(pts, use_lip=use_lip)
|
||||
elif pts.shape[0] == 478:
|
||||
pt2 = parse_pt2_from_pt478(pts, use_lip=use_lip)
|
||||
elif pts.shape[0] > 101:
|
||||
# take the first 101 points
|
||||
pt2 = parse_pt2_from_pt101(pts[:101], use_lip=use_lip)
|
||||
elif pts.shape[0] == 9:
|
||||
pt2 = parse_pt2_from_pt9(pts, use_lip=use_lip)
|
||||
else:
|
||||
raise Exception(f'Unknow shape: {pts.shape}')
|
||||
|
||||
@@ -462,13 +416,15 @@ def crop_image(img, pts: np.ndarray, **kwargs):
|
||||
ret_dct = {
|
||||
'M_o2c': M_o2c, # from the original image to the cropped image 3x3
|
||||
'M_c2o': M_c2o, # from the cropped image to the original image 3x3
|
||||
'img_crop': img_crop, # the cropped image
|
||||
'pt_crop': pt_crop, # the landmarks of the cropped image
|
||||
}
|
||||
|
||||
return ret_dct, img_crop
|
||||
return ret_dct
|
||||
|
||||
def average_bbox_lst(bbox_lst):
|
||||
if len(bbox_lst) == 0:
|
||||
return None
|
||||
bbox_arr = np.array(bbox_lst)
|
||||
return np.mean(bbox_arr, axis=0).tolist()
|
||||
|
||||
|
||||
+12
-143
@@ -5,8 +5,8 @@ from typing import List, Union, Tuple
|
||||
from dataclasses import dataclass, field
|
||||
import cv2#; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False)
|
||||
|
||||
from .landmark_runner import LandmarkRunner, LandmarkRunnerTorch
|
||||
|
||||
from .landmark_runner import LandmarkRunner
|
||||
from .face_analysis_diy import FaceAnalysisDIY
|
||||
from .crop import crop_image
|
||||
|
||||
import folder_paths
|
||||
@@ -22,11 +22,11 @@ class Trajectory:
|
||||
frame_rgb_lst: Union[Tuple, List, np.ndarray] = field(default_factory=list) # frame list
|
||||
frame_rgb_crop_lst: Union[Tuple, List, np.ndarray] = field(default_factory=list) # frame crop list
|
||||
|
||||
class CropperInsightFace(object):
|
||||
|
||||
class Cropper(object):
|
||||
def __init__(self, **kwargs) -> None:
|
||||
device_id = kwargs.get('device_id', 0)
|
||||
provider = kwargs.get('onnx_device', 'CPU')
|
||||
detection_threshold = kwargs.get('detection_threshold', 0.5)
|
||||
self.landmark_runner = LandmarkRunner(
|
||||
ckpt_path=os.path.join(folder_paths.models_dir, 'liveportrait', 'landmark.onnx'),
|
||||
onnx_provider=provider,
|
||||
@@ -34,13 +34,12 @@ class CropperInsightFace(object):
|
||||
)
|
||||
self.landmark_runner.warmup()
|
||||
|
||||
from .face_analysis_diy import FaceAnalysisDIY
|
||||
self.face_analysis_wrapper = FaceAnalysisDIY(
|
||||
name='buffalo_l',
|
||||
root=os.path.join(folder_paths.models_dir, 'insightface'),
|
||||
providers=[provider + 'ExecutionProvider',]
|
||||
)
|
||||
self.face_analysis_wrapper.prepare(ctx_id=device_id, det_size=(512, 512), det_thresh=detection_threshold)
|
||||
self.face_analysis_wrapper.prepare(ctx_id=device_id, det_size=(512, 512))
|
||||
self.face_analysis_wrapper.warmup()
|
||||
|
||||
def crop_single_image(self, img_rgb, dsize, scale, vy_ratio, vx_ratio, face_index, face_index_order, rotate):
|
||||
@@ -54,14 +53,16 @@ class CropperInsightFace(object):
|
||||
|
||||
if len(src_face) == 0:
|
||||
ret_dct = {}
|
||||
cropped_image_256 = None
|
||||
return ret_dct, cropped_image_256
|
||||
return ret_dct
|
||||
#raise Exception("No face detected in the source image!")
|
||||
#elif len(src_face) > 1:
|
||||
# print(f'More than one face detected in the image, only pick one face by rule {direction}.')
|
||||
|
||||
src_face = src_face[face_index] # choose the index if multiple faces detected
|
||||
pts = src_face.landmark_2d_106
|
||||
|
||||
# crop the face
|
||||
ret_dct, image_crop = crop_image(
|
||||
ret_dct = crop_image(
|
||||
img_rgb, # ndarray
|
||||
pts, # 106x2 or Nx2
|
||||
dsize=dsize,
|
||||
@@ -71,8 +72,7 @@ class CropperInsightFace(object):
|
||||
rotate=rotate
|
||||
)
|
||||
# update a 256x256 version for network input or else
|
||||
cropped_image_256 = cv2.resize(image_crop, (256, 256), interpolation=cv2.INTER_AREA)
|
||||
del image_crop
|
||||
ret_dct['img_crop_256x256'] = cv2.resize(ret_dct['img_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]
|
||||
@@ -82,135 +82,4 @@ class CropperInsightFace(object):
|
||||
lmk = recon_ret['pts']
|
||||
ret_dct['lmk_crop'] = lmk
|
||||
|
||||
return ret_dct, cropped_image_256
|
||||
|
||||
class CropperMediaPipe(object):
|
||||
def __init__(self, **kwargs) -> None:
|
||||
device_id = kwargs.get('device_id', 0)
|
||||
provider = kwargs.get('onnx_device', 'CPU')
|
||||
|
||||
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 ...media_pipe.mp_utils import LMKExtractor
|
||||
self.lmk_extractor = LMKExtractor()
|
||||
|
||||
def crop_single_image(self, img_rgb, dsize, scale, vy_ratio, vx_ratio, face_index, face_index_order, rotate):
|
||||
|
||||
face_result = self.lmk_extractor(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]
|
||||
|
||||
lmks = []
|
||||
for index in range(len(face_landmarks)):
|
||||
x = face_landmarks[index].x * img_rgb.shape[1]
|
||||
y = face_landmarks[index].y * img_rgb.shape[0]
|
||||
lmks.append([x, y])
|
||||
pts = np.array(lmks)
|
||||
|
||||
# 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)
|
||||
del image_crop
|
||||
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
|
||||
|
||||
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)
|
||||
del image_crop
|
||||
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
|
||||
|
||||
return ret_dct
|
||||
@@ -4,27 +4,14 @@ from pykalman import KalmanFilter
|
||||
|
||||
|
||||
def smooth(x_d_lst, shape, device, observation_variance=3e-6, process_variance=1e-5):
|
||||
# Reshape x_d_lst, skipping None values
|
||||
x_d_lst_reshape = [x.reshape(-1) for x in x_d_lst if x is not None]
|
||||
|
||||
if not x_d_lst_reshape: # Check if x_d_lst_reshape is empty after filtering
|
||||
return [None] * len(x_d_lst) # Return a list of Nones with the same length as x_d_lst
|
||||
|
||||
x_d_lst_reshape = [x.reshape(-1) for x in x_d_lst]
|
||||
x_d_stacked = np.vstack(x_d_lst_reshape)
|
||||
|
||||
kf = KalmanFilter(
|
||||
initial_state_mean=x_d_stacked[0],
|
||||
n_dim_obs=x_d_stacked.shape[1],
|
||||
transition_covariance=process_variance * np.eye(x_d_stacked.shape[1]),
|
||||
observation_covariance=observation_variance * np.eye(x_d_stacked.shape[1])
|
||||
)
|
||||
|
||||
smoothed_state_means, _ = kf.smooth(x_d_stacked)
|
||||
|
||||
# Initialize an iterator for smoothed_state_means
|
||||
smoothed_states_iter = iter(smoothed_state_means)
|
||||
|
||||
# Create x_d_lst_smooth, inserting None for each None encountered in the original list
|
||||
x_d_lst_smooth = [torch.tensor(next(smoothed_states_iter).reshape(shape[-2:]), dtype=torch.float32, device=device) if x is not None else None for x in x_d_lst]
|
||||
|
||||
x_d_lst_smooth = [torch.tensor(state_mean.reshape(shape[-2:]), dtype=torch.float32, device=device) for state_mean in smoothed_state_means]
|
||||
return x_d_lst_smooth
|
||||
@@ -4,12 +4,10 @@
|
||||
utility functions and classes to handle feature extraction and model loading
|
||||
"""
|
||||
|
||||
import os.path as osp
|
||||
import cv2
|
||||
import torch
|
||||
import numpy as np
|
||||
from typing import Union
|
||||
from collections import OrderedDict
|
||||
from scipy.spatial import ConvexHull # pylint: disable=E0401,E0611
|
||||
|
||||
def squeeze_tensor_to_numpy(tensor):
|
||||
out = tensor.data.squeeze(0).cpu().numpy()
|
||||
@@ -73,24 +71,3 @@ def resize_to_limit(img, max_dim=1280, n=2):
|
||||
if new_h != img.shape[0] or new_w != img.shape[1]:
|
||||
img = img[:new_h, :new_w]
|
||||
return img
|
||||
|
||||
def tensor_to_numpy(data: Union[np.ndarray, torch.Tensor]) -> np.ndarray:
|
||||
"""transform torch.Tensor into numpy.ndarray"""
|
||||
if isinstance(data, torch.Tensor):
|
||||
return data.data.cpu().numpy()
|
||||
return data
|
||||
|
||||
def calc_motion_multiplier(
|
||||
kp_source: Union[np.ndarray, torch.Tensor],
|
||||
kp_driving_initial: Union[np.ndarray, torch.Tensor]
|
||||
) -> float:
|
||||
"""calculate motion_multiplier based on the source image and the first driving frame"""
|
||||
kp_source_np = tensor_to_numpy(kp_source)
|
||||
kp_driving_initial_np = tensor_to_numpy(kp_driving_initial)
|
||||
|
||||
source_area = ConvexHull(kp_source_np.squeeze(0)).volume
|
||||
driving_area = ConvexHull(kp_driving_initial_np.squeeze(0)).volume
|
||||
motion_multiplier = np.sqrt(source_area) / np.sqrt(driving_area)
|
||||
# motion_multiplier = np.cbrt(source_area) / np.cbrt(driving_area)
|
||||
|
||||
return motion_multiplier
|
||||
@@ -3,11 +3,10 @@
|
||||
import cv2#; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False)
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
import onnxruntime
|
||||
from .timer import Timer
|
||||
from .crop import crop_image, _transform_pts
|
||||
import folder_paths
|
||||
import os
|
||||
|
||||
def to_ndarray(obj):
|
||||
if isinstance(obj, torch.Tensor):
|
||||
return obj.cpu().numpy()
|
||||
@@ -25,8 +24,6 @@ class LandmarkRunner(object):
|
||||
self.dsize = kwargs.get('dsize', 224)
|
||||
self.timer = Timer()
|
||||
|
||||
import onnxruntime
|
||||
|
||||
if onnx_provider.lower() == 'cuda':
|
||||
self.session = onnxruntime.InferenceSession(
|
||||
ckpt_path, providers=[
|
||||
@@ -47,7 +44,8 @@ class LandmarkRunner(object):
|
||||
|
||||
def run(self, img_rgb: np.ndarray, lmk=None):
|
||||
if lmk is not None:
|
||||
crop_dct, img_crop_rgb = crop_image(img_rgb, lmk, dsize=self.dsize, scale=1.5, vy_ratio=-0.1)
|
||||
crop_dct = crop_image(img_rgb, lmk, dsize=self.dsize, scale=1.5, vy_ratio=-0.1)
|
||||
img_crop_rgb = crop_dct['img_crop']
|
||||
else:
|
||||
img_crop_rgb = cv2.resize(img_rgb, (self.dsize, self.dsize))
|
||||
scale = max(img_rgb.shape[:2]) / self.dsize
|
||||
@@ -66,7 +64,7 @@ class LandmarkRunner(object):
|
||||
|
||||
pts = to_ndarray(out_pts[0]).reshape(-1, 2) * self.dsize # scale to 0-224
|
||||
pts = _transform_pts(pts, M=crop_dct['M_c2o'])
|
||||
del crop_dct, img_crop_rgb
|
||||
|
||||
return {
|
||||
'pts': pts, # 2d landmarks 203 points
|
||||
}
|
||||
@@ -80,54 +78,3 @@ class LandmarkRunner(object):
|
||||
|
||||
elapse = self.timer.toc()
|
||||
print(f'LandmarkRunner warmup time: {elapse:.3f}s')
|
||||
|
||||
class LandmarkRunnerTorch(object):
|
||||
"""landmark runner torch version"""
|
||||
def __init__(self, **kwargs):
|
||||
self.device = kwargs.get('device_id', 0)
|
||||
self.dsize = kwargs.get('dsize', 224)
|
||||
ckpt_path = kwargs.get('ckpt_path')
|
||||
|
||||
if not os.path.exists(ckpt_path):
|
||||
download_path = os.path.join(folder_paths.models_dir, "liveportrait")
|
||||
print(f"Downloading model to: {ckpt_path}")
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
snapshot_download(
|
||||
repo_id="Kijai/LivePortrait_safetensors",
|
||||
allow_patterns="*landmark_model.pth*",
|
||||
local_dir=download_path,
|
||||
local_dir_use_symlinks=False,
|
||||
)
|
||||
self.model = torch.load(ckpt_path).to(self.device)
|
||||
|
||||
def _run(self, inp):
|
||||
input = torch.from_numpy(inp).to(self.device)
|
||||
out = self.model(input)
|
||||
return out
|
||||
|
||||
def run(self, img_rgb: np.ndarray, lmk=None):
|
||||
if lmk is not None:
|
||||
crop_dct, img_crop_rgb = crop_image(img_rgb, lmk, dsize=self.dsize, scale=1.5, vy_ratio=-0.1)
|
||||
else:
|
||||
img_crop_rgb = cv2.resize(img_rgb, (self.dsize, self.dsize))
|
||||
scale = max(img_rgb.shape[:2]) / self.dsize
|
||||
crop_dct = {
|
||||
'M_c2o': np.array([
|
||||
[scale, 0., 0.],
|
||||
[0., scale, 0.],
|
||||
[0., 0., 1.],
|
||||
], dtype=np.float32),
|
||||
}
|
||||
|
||||
inp = (img_crop_rgb.astype(np.float32) / 255.).transpose(2, 0, 1)[None, ...] # HxWx3 (BGR) -> 1x3xHxW (RGB!)
|
||||
|
||||
out_lst = self._run(inp)
|
||||
out_pts = out_lst[2]
|
||||
|
||||
pts = to_ndarray(out_pts[0]).reshape(-1, 2) * self.dsize # scale to 0-224
|
||||
pts = _transform_pts(pts, M=crop_dct['M_c2o'])
|
||||
del crop_dct, img_crop_rgb
|
||||
return {
|
||||
'pts': pts, # 2d landmarks 203 points
|
||||
}
|
||||
@@ -1 +0,0 @@
|
||||
from .mp_utils import LMKExtractor
|
||||
File diff suppressed because it is too large
Load Diff
Binary file not shown.
Binary file not shown.
@@ -1,38 +0,0 @@
|
||||
import os
|
||||
import mediapipe as mp
|
||||
|
||||
from mediapipe.tasks import python
|
||||
from mediapipe.tasks.python import vision
|
||||
from . import face_landmark
|
||||
|
||||
CUR_DIR = os.path.dirname(__file__)
|
||||
|
||||
class LMKExtractor():
|
||||
def __init__(self):
|
||||
# Create an FaceLandmarker object.
|
||||
self.mode = mp.tasks.vision.FaceDetectorOptions.running_mode.IMAGE
|
||||
base_options = python.BaseOptions(model_asset_path=os.path.join(CUR_DIR, 'mp_models','face_landmarker_v2_with_blendshapes.task'))
|
||||
base_options.delegate = mp.tasks.BaseOptions.Delegate.CPU
|
||||
options = vision.FaceLandmarkerOptions(base_options=base_options,
|
||||
running_mode=self.mode,
|
||||
output_face_blendshapes=False,
|
||||
output_facial_transformation_matrixes=True,
|
||||
num_faces=1,
|
||||
min_face_detection_confidence=0.5,
|
||||
min_face_presence_confidence=0.5,
|
||||
min_tracking_confidence=0.5)
|
||||
self.detector = face_landmark.FaceLandmarker.create_from_options(options)
|
||||
|
||||
det_base_options = python.BaseOptions(model_asset_path=os.path.join(CUR_DIR, 'mp_models','blaze_face_short_range.tflite'))
|
||||
det_options = vision.FaceDetectorOptions(base_options=det_base_options)
|
||||
self.det_detector = vision.FaceDetector.create_from_options(det_options)
|
||||
|
||||
def __call__(self, img):
|
||||
image = mp.Image(image_format=mp.ImageFormat.SRGB, data=img)
|
||||
try:
|
||||
detection_result, _ = self.detector.detect(image)
|
||||
except:
|
||||
return None
|
||||
|
||||
return detection_result.face_landmarks
|
||||
|
||||
@@ -7,28 +7,11 @@ import comfy.utils
|
||||
import numpy as np
|
||||
import cv2
|
||||
from tqdm import tqdm
|
||||
import gc
|
||||
|
||||
import logging
|
||||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
script_directory = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
from .liveportrait.live_portrait_pipeline import LivePortraitPipeline
|
||||
try:
|
||||
from .liveportrait.utils.cropper import CropperMediaPipe
|
||||
except:
|
||||
log.warning("Can't load MediaPipe, MediaPipeCropper not available")
|
||||
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.utils.cropper import Cropper
|
||||
from .liveportrait.modules.spade_generator import SPADEDecoder
|
||||
from .liveportrait.modules.warping_network import WarpingNetwork
|
||||
from .liveportrait.modules.motion_extractor import MotionExtractor
|
||||
@@ -38,21 +21,27 @@ from .liveportrait.modules.appearance_feature_extractor import (
|
||||
from .liveportrait.modules.stitching_retargeting_network import (
|
||||
StitchingRetargetingNetwork,
|
||||
)
|
||||
from .liveportrait.utils.camera import get_rotation_matrix
|
||||
from .liveportrait.utils.crop import _transform_img_kornia
|
||||
|
||||
import logging
|
||||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
class InferenceConfig:
|
||||
def __init__(
|
||||
self,
|
||||
mask_crop=None,
|
||||
flag_use_half_precision=True,
|
||||
flag_lip_zero=True,
|
||||
lip_zero_threshold=0.03,
|
||||
flag_eye_retargeting=False,
|
||||
flag_lip_retargeting=False,
|
||||
flag_stitching=True,
|
||||
flag_relative=True,
|
||||
flag_relative_rotation_only=False,
|
||||
input_shape=(256, 256),
|
||||
flag_pasteback=True,
|
||||
device_id=0,
|
||||
flag_do_crop=True,
|
||||
flag_do_rot=True,
|
||||
):
|
||||
self.flag_use_half_precision = flag_use_half_precision
|
||||
@@ -61,10 +50,15 @@ class InferenceConfig:
|
||||
self.flag_eye_retargeting = flag_eye_retargeting
|
||||
self.flag_lip_retargeting = flag_lip_retargeting
|
||||
self.flag_stitching = flag_stitching
|
||||
self.flag_relative = flag_relative
|
||||
self.flag_relative_rotation_only = flag_relative_rotation_only
|
||||
self.input_shape = input_shape
|
||||
self.flag_pasteback = flag_pasteback
|
||||
self.device_id = device_id
|
||||
self.flag_do_crop = flag_do_crop
|
||||
self.flag_do_rot = flag_do_rot
|
||||
|
||||
self.mask_crop = mask_crop
|
||||
|
||||
class DownloadAndLoadLivePortraitModels:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
@@ -79,12 +73,6 @@ class DownloadAndLoadLivePortraitModels:
|
||||
],
|
||||
{"default": "auto"},
|
||||
),
|
||||
"mode": (
|
||||
[
|
||||
"human",
|
||||
"animal",
|
||||
],
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
@@ -93,43 +81,39 @@ class DownloadAndLoadLivePortraitModels:
|
||||
FUNCTION = "loadmodel"
|
||||
CATEGORY = "LivePortrait"
|
||||
|
||||
def loadmodel(self, precision="fp16", mode="human"):
|
||||
def loadmodel(self, precision="fp16"):
|
||||
device = mm.get_torch_device()
|
||||
mm.soft_empty_cache()
|
||||
|
||||
if precision == 'auto':
|
||||
try:
|
||||
if mm.is_device_mps(device):
|
||||
log.info("LivePortrait using fp32 for MPS")
|
||||
print("LivePortrait using fp32 for MPS")
|
||||
dtype = 'fp32'
|
||||
elif mm.should_use_fp16():
|
||||
log.info("LivePortrait using fp16")
|
||||
print("LivePortrait using fp16")
|
||||
dtype = 'fp16'
|
||||
else:
|
||||
log.info("LivePortrait using fp32")
|
||||
print("LivePortrait using fp32")
|
||||
dtype = 'fp32'
|
||||
except:
|
||||
raise AttributeError("ComfyUI version too old, can't autodetect properly. Set your dtypes manually.")
|
||||
else:
|
||||
dtype = precision
|
||||
log.info(f"LivePortrait using {dtype}")
|
||||
print(f"LivePortrait using {dtype}")
|
||||
|
||||
pbar = comfy.utils.ProgressBar(3)
|
||||
|
||||
base_bath = os.path.join(folder_paths.models_dir, "liveportrait")
|
||||
if mode == "human":
|
||||
model_path = base_bath
|
||||
else:
|
||||
model_path = os.path.join(base_bath, "animal")
|
||||
|
||||
download_path = os.path.join(folder_paths.models_dir, "liveportrait")
|
||||
model_path = os.path.join(download_path)
|
||||
|
||||
if not os.path.exists(model_path):
|
||||
log.info(f"Downloading model to: {model_path}")
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
snapshot_download(
|
||||
repo_id="Kijai/LivePortrait_safetensors",
|
||||
ignore_patterns=["*landmark_model.pth*","*animal*"] if mode == "human" else ["*landmark_model.pth*"],
|
||||
local_dir=base_bath,
|
||||
local_dir=download_path,
|
||||
local_dir_use_symlinks=False,
|
||||
)
|
||||
|
||||
@@ -207,19 +191,22 @@ class DownloadAndLoadLivePortraitModels:
|
||||
stitcher_checkpoint = filter_checkpoint_for_model(checkpoint, stitcher_prefix)
|
||||
stitcher = StitchingRetargetingNetwork(**config.get("stitching"))
|
||||
stitcher.load_state_dict(stitcher_checkpoint)
|
||||
stitcher = stitcher.to(device).eval()
|
||||
stitcher = stitcher.to(device)
|
||||
stitcher.eval()
|
||||
|
||||
lip_prefix = "retarget_mouth"
|
||||
lip_checkpoint = filter_checkpoint_for_model(checkpoint, lip_prefix)
|
||||
retargetor_lip = StitchingRetargetingNetwork(**config.get("lip"))
|
||||
retargetor_lip.load_state_dict(lip_checkpoint)
|
||||
retargetor_lip = retargetor_lip.to(device).eval()
|
||||
retargetor_lip = retargetor_lip.to(device)
|
||||
retargetor_lip.eval()
|
||||
|
||||
eye_prefix = "retarget_eye"
|
||||
eye_checkpoint = filter_checkpoint_for_model(checkpoint, eye_prefix)
|
||||
retargetor_eye = StitchingRetargetingNetwork(**config.get("eye"))
|
||||
retargetor_eye.load_state_dict(eye_checkpoint)
|
||||
retargetor_eye = retargetor_eye.to(device).eval()
|
||||
retargetor_eye = retargetor_eye.to(device)
|
||||
retargetor_eye.eval()
|
||||
log.info("Load stitching_retargeting_module done.")
|
||||
|
||||
self.stich_retargeting_module = {
|
||||
@@ -255,7 +242,7 @@ class LivePortraitProcess:
|
||||
"lip_zero": ("BOOLEAN", {"default": False}),
|
||||
"lip_zero_threshold": ("FLOAT", {"default": 0.03, "min": 0.001, "max": 4.0, "step": 0.001}),
|
||||
"stitching": ("BOOLEAN", {"default": True}),
|
||||
"delta_multiplier": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.001}),
|
||||
"delta_multiplier": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.001}),
|
||||
"mismatch_method": (
|
||||
[
|
||||
"constant",
|
||||
@@ -271,7 +258,6 @@ class LivePortraitProcess:
|
||||
"relative",
|
||||
"source_video_smoothed",
|
||||
"relative_rotation_only",
|
||||
"single_frame",
|
||||
"off"
|
||||
],
|
||||
),
|
||||
@@ -279,19 +265,20 @@ class LivePortraitProcess:
|
||||
},
|
||||
|
||||
"optional": {
|
||||
"mask": ("MASK", {"default": None}),
|
||||
"opt_retargeting_info": ("RETARGETINGINFO", {"default": None}),
|
||||
"expression_friendly": ("BOOLEAN", {"default": False}),
|
||||
"expression_friendly_multiplier": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 100.0, "step": 0.001}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"IMAGE",
|
||||
"LP_OUT",
|
||||
"IMAGE",
|
||||
"MASK",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"cropped_image",
|
||||
"output",
|
||||
"cropped_images",
|
||||
"full_images",
|
||||
"mask",
|
||||
)
|
||||
FUNCTION = "process"
|
||||
CATEGORY = "LivePortrait"
|
||||
@@ -309,14 +296,12 @@ class LivePortraitProcess:
|
||||
driving_smooth_observation_variance: float,
|
||||
delta_multiplier: float = 1.0,
|
||||
mismatch_method: str = "constant",
|
||||
mask: torch.Tensor = None,
|
||||
opt_retargeting_info: dict = None,
|
||||
expression_friendly: bool = False,
|
||||
expression_friendly_multiplier: float = 1.0,
|
||||
):
|
||||
if driving_images.shape[0] < source_image.shape[0]:
|
||||
raise ValueError("The number of driving images should be larger than the number of source images.")
|
||||
if expression_friendly and source_image.shape[0] > 1:
|
||||
raise ValueError("expression_friendly works only with single source image")
|
||||
source_np = (source_image * 255).byte().numpy()
|
||||
|
||||
if opt_retargeting_info is not None:
|
||||
pipeline.live_portrait_wrapper.cfg.flag_eye_retargeting = opt_retargeting_info["eye_retargeting"]
|
||||
@@ -335,163 +320,52 @@ class LivePortraitProcess:
|
||||
pipeline.live_portrait_wrapper.cfg.flag_lip_zero = lip_zero
|
||||
pipeline.live_portrait_wrapper.cfg.lip_zero_threshold = lip_zero_threshold
|
||||
|
||||
if relative_motion_mode != "off":
|
||||
pipeline.live_portrait_wrapper.cfg.flag_relative = True
|
||||
else:
|
||||
pipeline.live_portrait_wrapper.cfg.flag_relative = False
|
||||
|
||||
if lip_zero and opt_retargeting_info is not None:
|
||||
log.warning("Warning: lip_zero only has an effect with lip or eye retargeting")
|
||||
|
||||
if driving_images.shape[1] != 256 or driving_images.shape[2] != 256:
|
||||
driving_images_256 = comfy.utils.common_upscale(driving_images.permute(0, 3, 1, 2), 256, 256, "lanczos", "disabled")
|
||||
else:
|
||||
driving_images_256 = driving_images.permute(0, 3, 1, 2)
|
||||
|
||||
if mask is not None:
|
||||
crop_mask = mask.unsqueeze(-1).expand(-1, -1, -1, 3)
|
||||
pipeline.live_portrait_wrapper.cfg.mask_crop = crop_mask
|
||||
else:
|
||||
log.info("Using default mask template")
|
||||
pipeline.live_portrait_wrapper.cfg.mask_crop = cv2.imread(os.path.join(script_directory, "liveportrait", "utils", "resources", "mask_template.png"), cv2.IMREAD_COLOR)
|
||||
|
||||
driving_images_256 = comfy.utils.common_upscale(driving_images.permute(0, 3, 1, 2), 256, 256, "lanczos", "disabled")
|
||||
if pipeline.live_portrait_wrapper.cfg.flag_use_half_precision:
|
||||
driving_images_256 = driving_images_256.to(torch.float16)
|
||||
|
||||
out = pipeline.execute(
|
||||
cropped_out_list = []
|
||||
full_out_list = []
|
||||
|
||||
cropped_out_list, full_out_list, out_mask_list = pipeline.execute(
|
||||
source_np,
|
||||
driving_images_256,
|
||||
crop_info,
|
||||
driving_landmarks,
|
||||
delta_multiplier,
|
||||
relative_motion_mode,
|
||||
driving_smooth_observation_variance,
|
||||
mismatch_method,
|
||||
expression_friendly=expression_friendly,
|
||||
driving_multiplier=expression_friendly_multiplier,
|
||||
mismatch_method
|
||||
)
|
||||
|
||||
total_frames = len(out["out_list"])
|
||||
|
||||
if total_frames > 1:
|
||||
cropped_image_list = []
|
||||
for i in (range(total_frames)):
|
||||
if not out["out_list"][i]:
|
||||
cropped_image_list.append(torch.zeros(1, 512, 512, 3, dtype=torch.float32, device = "cpu"))
|
||||
else:
|
||||
cropped_image = torch.clamp(out["out_list"][i]["out"], 0, 1).permute(0, 2, 3, 1).cpu()
|
||||
cropped_image_list.append(cropped_image)
|
||||
cropped_out_tensors = torch.cat(cropped_out_list, dim=0)
|
||||
|
||||
cropped_out_tensors = torch.cat(cropped_image_list, dim=0)
|
||||
else:
|
||||
cropped_out_tensors = torch.clamp(out["out_list"][0]["out"], 0, 1).permute(0, 2, 3, 1)
|
||||
|
||||
return (cropped_out_tensors, out,)
|
||||
|
||||
class LivePortraitComposite:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
|
||||
"source_image": ("IMAGE",),
|
||||
"cropped_image": ("IMAGE",),
|
||||
"liveportrait_out": ("LP_OUT", ),
|
||||
},
|
||||
"optional": {
|
||||
"mask": ("MASK", {"default": None}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"IMAGE",
|
||||
"MASK",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"full_images",
|
||||
"mask",
|
||||
)
|
||||
FUNCTION = "process"
|
||||
CATEGORY = "LivePortrait"
|
||||
|
||||
def process(self, source_image, cropped_image, liveportrait_out, mask=None):
|
||||
mm.soft_empty_cache()
|
||||
gc.collect()
|
||||
device = mm.get_torch_device()
|
||||
if mm.is_device_mps(device):
|
||||
device = torch.device('cpu') #this function returns NaNs on MPS, defaulting to CPU
|
||||
|
||||
B, H, W, C = source_image.shape
|
||||
source_image = source_image.permute(0, 3, 1, 2) # B,H,W,C -> B,C,H,W
|
||||
cropped_image = cropped_image.permute(0, 3, 1, 2)
|
||||
|
||||
if mask is not None:
|
||||
if len(mask.size())==2:
|
||||
crop_mask = mask.unsqueeze(0).unsqueeze(-1).expand(-1, -1, -1, 3)
|
||||
else:
|
||||
crop_mask = mask.unsqueeze(-1).expand(-1, -1, -1, 3)
|
||||
else:
|
||||
log.info("Using default mask template")
|
||||
crop_mask = cv2.imread(os.path.join(script_directory, "liveportrait", "utils", "resources", "mask_template.png"), cv2.IMREAD_COLOR)
|
||||
crop_mask = torch.from_numpy(crop_mask)
|
||||
crop_mask = crop_mask.unsqueeze(0).float() / 255.0
|
||||
|
||||
crop_info = liveportrait_out["crop_info"]
|
||||
composited_image_list = []
|
||||
out_mask_list = []
|
||||
|
||||
total_frames = len(liveportrait_out["out_list"])
|
||||
log.info(f"Total frames: {total_frames}")
|
||||
|
||||
pbar = comfy.utils.ProgressBar(total_frames)
|
||||
for i in tqdm(range(total_frames), desc='Compositing..', total=total_frames):
|
||||
safe_index = min(i, len(crop_info["crop_info_list"]) - 1)
|
||||
|
||||
if liveportrait_out["mismatch_method"] == "cut":
|
||||
source_frame = source_image[safe_index].unsqueeze(0).to(device)
|
||||
else:
|
||||
source_frame = _get_source_frame(source_image, i, liveportrait_out["mismatch_method"]).unsqueeze(0).to(device)
|
||||
|
||||
if not liveportrait_out["out_list"][i]:
|
||||
composited_image_list.append(source_frame.cpu())
|
||||
out_mask_list.append(torch.zeros((1, 3, H, W), device="cpu"))
|
||||
else:
|
||||
cropped_image = torch.clamp(liveportrait_out["out_list"][i]["out"], 0, 1).permute(0, 2, 3, 1)
|
||||
|
||||
# Transform and blend
|
||||
cropped_image_to_original = _transform_img_kornia(
|
||||
cropped_image,
|
||||
crop_info["crop_info_list"][safe_index]["M_c2o"],
|
||||
dsize=(W, H),
|
||||
device=device
|
||||
)
|
||||
|
||||
mask_ori = _transform_img_kornia(
|
||||
crop_mask[min(i,len(crop_mask)-1)].unsqueeze(0),
|
||||
crop_info["crop_info_list"][safe_index]["M_c2o"],
|
||||
dsize=(W, H),
|
||||
device=device
|
||||
)
|
||||
|
||||
cropped_image_to_original_blend = torch.clip(
|
||||
mask_ori * cropped_image_to_original + (1 - mask_ori) * source_frame, 0, 1
|
||||
)
|
||||
|
||||
composited_image_list.append(cropped_image_to_original_blend.cpu())
|
||||
out_mask_list.append(mask_ori.cpu())
|
||||
pbar.update(1)
|
||||
|
||||
full_tensors_out = torch.cat(composited_image_list, dim=0)
|
||||
full_tensors_out = torch.cat(full_out_list, dim=0)
|
||||
full_tensors_out = full_tensors_out.permute(0, 2, 3, 1)
|
||||
|
||||
mask_tensors_out = torch.cat(out_mask_list, dim=0)
|
||||
mask_tensors_out = mask_tensors_out[:, 0, :, :]
|
||||
mask_tensors_out = mask_tensors_out[:, :, :, 0]
|
||||
|
||||
return (
|
||||
full_tensors_out.float(),
|
||||
mask_tensors_out.float()
|
||||
cropped_out_tensors.cpu().float(),
|
||||
full_tensors_out.cpu().float(),
|
||||
mask_tensors_out.cpu().float()
|
||||
)
|
||||
|
||||
def _get_source_frame(source, idx, method):
|
||||
if source.shape[0] == 1:
|
||||
return source[0]
|
||||
|
||||
if method == "constant":
|
||||
return source[min(idx, source.shape[0] - 1)]
|
||||
elif method == "cycle":
|
||||
return source[idx % source.shape[0]]
|
||||
elif method == "mirror":
|
||||
cycle_length = 2 * source.shape[0] - 2
|
||||
mirror_idx = idx % cycle_length
|
||||
if mirror_idx >= source.shape[0]:
|
||||
mirror_idx = cycle_length - mirror_idx
|
||||
return source[mirror_idx]
|
||||
|
||||
class LivePortraitLoadCropper:
|
||||
@classmethod
|
||||
@@ -503,40 +377,6 @@ class LivePortraitLoadCropper:
|
||||
"default": 'CPU'
|
||||
}),
|
||||
"keep_model_loaded": ("BOOLEAN", {"default": True})
|
||||
},
|
||||
"optional": {
|
||||
"detection_threshold": ("FLOAT", {"default": 0.5, "min": 0.05, "max": 1.0, "step": 0.01}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LPCROPPER",)
|
||||
RETURN_NAMES = ("cropper",)
|
||||
FUNCTION = "crop"
|
||||
CATEGORY = "LivePortrait"
|
||||
|
||||
def crop(self, onnx_device, keep_model_loaded, detection_threshold=0.5):
|
||||
cropper_init_config = {
|
||||
'keep_model_loaded': keep_model_loaded,
|
||||
'onnx_device': onnx_device,
|
||||
'detection_threshold': detection_threshold
|
||||
}
|
||||
|
||||
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 = CropperInsightFace(**cropper_init_config)
|
||||
|
||||
return (self.cropper,)
|
||||
|
||||
class LivePortraitLoadMediaPipeCropper:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
|
||||
"landmarkrunner_onnx_device": (
|
||||
['CPU', 'CUDA', 'ROCM', 'CoreML', 'torch_gpu'], {
|
||||
"default": 'CPU'
|
||||
}),
|
||||
"keep_model_loaded": ("BOOLEAN", {"default": True})
|
||||
},
|
||||
}
|
||||
|
||||
@@ -545,67 +385,15 @@ class LivePortraitLoadMediaPipeCropper:
|
||||
FUNCTION = "crop"
|
||||
CATEGORY = "LivePortrait"
|
||||
|
||||
def crop(self, landmarkrunner_onnx_device, keep_model_loaded):
|
||||
def crop(self, onnx_device, keep_model_loaded):
|
||||
cropper_init_config = {
|
||||
'keep_model_loaded': keep_model_loaded,
|
||||
'onnx_device': landmarkrunner_onnx_device
|
||||
'onnx_device': onnx_device
|
||||
}
|
||||
|
||||
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 = CropperMediaPipe(**cropper_init_config)
|
||||
|
||||
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)
|
||||
self.cropper = Cropper(**cropper_init_config)
|
||||
|
||||
return (self.cropper,)
|
||||
|
||||
@@ -613,7 +401,6 @@ class LivePortraitCropper:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"pipeline": ("LIVEPORTRAITPIPE",),
|
||||
"cropper": ("LPCROPPER",),
|
||||
"source_image": ("IMAGE",),
|
||||
"dsize": ("INT", {"default": 512, "min": 64, "max": 2048}),
|
||||
@@ -641,68 +428,31 @@ class LivePortraitCropper:
|
||||
FUNCTION = "process"
|
||||
CATEGORY = "LivePortrait"
|
||||
|
||||
def process(self, pipeline, cropper, source_image, dsize, scale, vx_ratio, vy_ratio, face_index, face_index_order, rotate):
|
||||
source_image_np = (source_image.contiguous() * 255).byte().numpy()
|
||||
def process(self, cropper, source_image, dsize, scale, vx_ratio, vy_ratio, face_index, face_index_order, rotate):
|
||||
source_image_np = (source_image * 255).byte().numpy()
|
||||
|
||||
# Initialize lists
|
||||
crop_info_list = []
|
||||
cropped_images_list = []
|
||||
source_info = []
|
||||
source_rot_list = []
|
||||
f_s_list = []
|
||||
x_s_list = []
|
||||
|
||||
# Initialize a progress bar for the combined operation
|
||||
|
||||
pbar = comfy.utils.ProgressBar(len(source_image_np))
|
||||
for i in tqdm(range(len(source_image_np)), desc='Detecting, cropping, and processing..', total=len(source_image_np)):
|
||||
# Cropping operation
|
||||
crop_info, cropped_image_256 = cropper.crop_single_image(source_image_np[i], dsize, scale, vy_ratio, vx_ratio, face_index, face_index_order, rotate)
|
||||
|
||||
# Processing source images
|
||||
for i in tqdm(range(len(source_image_np)), desc='Detecting and cropping..', total=len(source_image_np)):
|
||||
crop_info = cropper.crop_single_image(source_image_np[i], dsize, scale, vy_ratio, vx_ratio, face_index, face_index_order, rotate)
|
||||
crop_info_list.append(crop_info)
|
||||
if crop_info:
|
||||
crop_info_list.append(crop_info)
|
||||
|
||||
cropped_images_list.append(cropped_image_256)
|
||||
|
||||
I_s = pipeline.live_portrait_wrapper.prepare_source(cropped_image_256)
|
||||
|
||||
x_s_info = pipeline.live_portrait_wrapper.get_kp_info(I_s)
|
||||
source_info.append(x_s_info)
|
||||
|
||||
x_s = pipeline.live_portrait_wrapper.transform_keypoint(x_s_info)
|
||||
x_s_list.append(x_s)
|
||||
|
||||
R_s = get_rotation_matrix(x_s_info["pitch"], x_s_info["yaw"], x_s_info["roll"])
|
||||
source_rot_list.append(R_s)
|
||||
|
||||
f_s = pipeline.live_portrait_wrapper.extract_feature_3d(I_s)
|
||||
f_s_list.append(f_s)
|
||||
|
||||
del I_s
|
||||
|
||||
cropped_image = crop_info['img_crop_256x256']
|
||||
else:
|
||||
log.warning(f"Warning: No face detected on frame {str(i)}, skipping")
|
||||
cropped_images_list.append(np.zeros((256, 256, 3), dtype=np.uint8))
|
||||
crop_info_list.append(None)
|
||||
f_s_list.append(None)
|
||||
x_s_list.append(None)
|
||||
source_info.append(None)
|
||||
source_rot_list.append(None)
|
||||
|
||||
# Update progress bar
|
||||
cropped_image = np.zeros((256, 256, 3), dtype=np.uint8)
|
||||
cropped_images_list.append(cropped_image)
|
||||
|
||||
pbar.update(1)
|
||||
|
||||
cropped_tensors_out = (
|
||||
torch.stack([torch.from_numpy(np_array) for np_array in cropped_images_list])
|
||||
/ 255
|
||||
)
|
||||
|
||||
|
||||
crop_info_dict = {
|
||||
'crop_info_list': crop_info_list,
|
||||
'source_rot_list': source_rot_list,
|
||||
'f_s_list': f_s_list,
|
||||
'x_s_list': x_s_list,
|
||||
'source_info': source_info
|
||||
'crop_info_list': crop_info_list
|
||||
}
|
||||
|
||||
return (cropped_tensors_out, crop_info_dict)
|
||||
@@ -746,7 +496,6 @@ class KeypointsToImage:
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"crop_info": ("CROPINFO", {"default": []}),
|
||||
"draw_lines": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
}
|
||||
|
||||
@@ -755,45 +504,20 @@ class KeypointsToImage:
|
||||
FUNCTION = "drawkeypoints"
|
||||
CATEGORY = "LivePortrait"
|
||||
|
||||
def drawkeypoints(self, crop_info, draw_lines):
|
||||
# left upper eye | left lower eye | right upper eye | right lower eye | upper lip top | lower lip bottom | upper lip bottom | lower lip top | jawline | left eyebrow | right eyebrow | nose | left pupil | right pupil | nose center
|
||||
indices = [ 12, 24, 37, 48, 66, 85, 96, 108, 145, 165, 185, 197, 198, 199, 203]
|
||||
colorlut = [(0, 0, 255), (0, 255, 0), (0, 0, 255), (0, 255, 0), (255, 0, 0), (255, 0, 255), (255, 255, 0), (0, 255, 255), (128, 128, 128), (128, 128, 0), (128, 128, 0), (0,128,128), (255, 255,255), (255, 255,255), (255,255,255)]
|
||||
colors = []
|
||||
c = 0
|
||||
for i in range(203):
|
||||
if i == indices[c]:
|
||||
c+=1
|
||||
colors.append(colorlut[c])
|
||||
try:
|
||||
height, width = crop_info["crop_info_list"][0]['input_image_size']
|
||||
except:
|
||||
height, width = 512, 512
|
||||
def drawkeypoints(self, crop_info):
|
||||
height, width = crop_info["crop_info_list"][0]['input_image_size']
|
||||
keypoints_img_list = []
|
||||
pbar = comfy.utils.ProgressBar(len(crop_info))
|
||||
for crop in crop_info["crop_info_list"]:
|
||||
if crop:
|
||||
keypoints = crop['lmk_crop'].copy()
|
||||
# Draw each landmark as a circle
|
||||
blank_image = np.zeros((height, width, 3), dtype=np.uint8) * 255
|
||||
|
||||
if draw_lines:
|
||||
start_idx = 0
|
||||
for end_idx in indices:
|
||||
color = colors[start_idx]
|
||||
for i in range(start_idx, end_idx - 1):
|
||||
pt1 = tuple(map(int, keypoints[i]))
|
||||
pt2 = tuple(map(int, keypoints[i+1]))
|
||||
if all(0 <= c < d for c, d in zip(pt1 + pt2, (width, height) * 2)):
|
||||
cv2.line(blank_image, pt1, pt2, color, thickness=1)
|
||||
if end_idx == start_idx +1:
|
||||
x,y = keypoints[start_idx]
|
||||
cv2.circle(blank_image, (int(x), int(y)), radius=1, thickness=-1, color=colors[start_idx])
|
||||
|
||||
start_idx = end_idx
|
||||
else:
|
||||
for index, (x, y) in enumerate(keypoints):
|
||||
cv2.circle(blank_image, (int(x), int(y)), radius=1, thickness=-1, color=colors[index])
|
||||
|
||||
for (x, y) in keypoints:
|
||||
# Ensure the coordinates are within the dimensions of the blank image
|
||||
if 0 <= x < width and 0 <= y < height:
|
||||
cv2.circle(blank_image, (int(x), int(y)), radius=2, color=(0, 0, 255))
|
||||
|
||||
keypoints_image = cv2.cvtColor(blank_image, cv2.COLOR_BGR2RGB)
|
||||
else:
|
||||
keypoints_image = np.zeros((height, width, 3), dtype=np.uint8) * 255
|
||||
@@ -803,6 +527,7 @@ class KeypointsToImage:
|
||||
keypoints_img_tensor = (
|
||||
torch.stack([torch.from_numpy(np_array) for np_array in keypoints_img_list]) / 255).float()
|
||||
|
||||
|
||||
return (keypoints_img_tensor,)
|
||||
|
||||
class KeypointScaler:
|
||||
@@ -862,20 +587,14 @@ NODE_CLASS_MAPPINGS = {
|
||||
"LivePortraitRetargeting": LivePortraitRetargeting,
|
||||
#"KeypointScaler": KeypointScaler,
|
||||
"KeypointsToImage": KeypointsToImage,
|
||||
"LivePortraitLoadCropper": LivePortraitLoadCropper,
|
||||
"LivePortraitLoadMediaPipeCropper": LivePortraitLoadMediaPipeCropper,
|
||||
"LivePortraitLoadFaceAlignmentCropper": LivePortraitLoadFaceAlignmentCropper,
|
||||
"LivePortraitComposite": LivePortraitComposite,
|
||||
"LivePortraitLoadCropper": LivePortraitLoadCropper
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DownloadAndLoadLivePortraitModels": "(Down)Load LivePortraitModels",
|
||||
"LivePortraitProcess": "LivePortrait Process",
|
||||
"LivePortraitCropper": "LivePortrait Cropper",
|
||||
"LivePortraitRetargeting": "LivePortrait Retargeting",
|
||||
"LivePortraitProcess": "LivePortraitProcess",
|
||||
"LivePortraitCropper": "LivePortraitCropper",
|
||||
"LivePortraitRetargeting": "LivePortraitRetargeting",
|
||||
#"KeypointScaler": "KeypointScaler",
|
||||
"KeypointsToImage": "LivePortrait KeypointsToImage",
|
||||
"LivePortraitLoadCropper": "LivePortrait Load InsightFaceCropper",
|
||||
"LivePortraitLoadMediaPipeCropper": "LivePortrait Load MediaPipeCropper",
|
||||
"LivePortraitLoadFaceAlignmentCropper": "LivePortrait Load FaceAlignmentCropper",
|
||||
"LivePortraitComposite": "LivePortrait Composite",
|
||||
}
|
||||
"LivePortraitLoadCropper": "LivePortrait LoadCropper"
|
||||
}
|
||||
+3
-3
@@ -1,8 +1,8 @@
|
||||
[project]
|
||||
name = "comfyui-liveportraitkj"
|
||||
description = "Nodes for [a/LivePortrait](https://github.com/KwaiVGI/LivePortrait)"
|
||||
version = "1.1.0"
|
||||
license = { text = "MIT License" }
|
||||
description = "Nodes for [a/LivePortrait](https://github.com/KwaiVGI/LivePortrait), insightface is required"
|
||||
version = "1.0.0"
|
||||
license = "LICENSE"
|
||||
dependencies = ["pyyaml", "numpy", "opencv-python", "rich"]
|
||||
|
||||
[project.urls]
|
||||
|
||||
@@ -1,66 +1,15 @@
|
||||
# ComfyUI nodes to use [LivePortrait](https://github.com/KwaiVGI/LivePortrait)
|
||||
## Update 2
|
||||
|
||||
Added another alternative face detector: https://github.com/1adrianb/face-alignment
|
||||
|
||||
|
||||
|
||||
As this can use blazeface back camera model (or SFD), it's far better for smaller faces than MediaPipe, that only can use the blazeface short -model.
|
||||
The warmup on the first run when using this can take a long time, but subsequent runs are quick.
|
||||
|
||||
Example detection using the blazeface_back_camera:
|
||||
|
||||
https://github.com/user-attachments/assets/40b1fdb4-0b1f-4ea8-8322-aa9151055db0
|
||||
|
||||
## Update
|
||||
|
||||
Rework of almost the whole thing that's been in develop is now merged into main, this means old workflows will not work, but everything should be faster and there's lots of new features.
|
||||
For legacy purposes the old main branch is moved to the legacy -branch
|
||||
|
||||
Changes
|
||||
- Added MediaPipe as alternative to Insightface, everything should now be covered under MIT and Apache-2.0 licenses when using it.
|
||||
- Proper Vid2vid including smoothing algorhitm (thanks @melMass)
|
||||
- Improved speed and efficiency, allows for near realtime view even in Comfy (~80-100ms delay)
|
||||
- Restructured nodes for more options
|
||||
- Auto skipping frames with no face detected
|
||||
- Numerous other things I have forgotten about at this point, it's been a lot
|
||||
- Better Mac support on MPS (thanks @Grant-CP
|
||||
|
||||
update to this update:
|
||||
- converted the landmark runner onnx model to torch model, not something I have done before and I didn't manage to do anything but make it .pth file, so you'll just have to trust me on it.
|
||||
This allows running all this without even having onnxruntime, it runs on GPU and is about just as fast. It's available on the MediaPipe cropper node as option:
|
||||
When selected it's automatically downloaded from here: https://huggingface.co/Kijai/LivePortrait_safetensors/blob/main/landmark_model.pth
|
||||
|
||||
|
||||
|
||||
|
||||
# Examples:
|
||||
|
||||
Realtime with webcam feed:
|
||||
|
||||
https://github.com/user-attachments/assets/31f77c10-b757-44ae-bb26-39e45ec0b2d9
|
||||
|
||||
Image2vid:
|
||||
|
||||
https://github.com/user-attachments/assets/cfec0419-d1eb-4e67-8913-890eeb155eef
|
||||
|
||||
Vid2Vid:
|
||||
|
||||
https://github.com/user-attachments/assets/28438fcb-fbb0-4e4e-baf4-00fe06c455de
|
||||
https://github.com/kijai/ComfyUI-LivePortrait/assets/40791699/e55e10f6-af61-4d73-b162-af29eb847516
|
||||
|
||||
|
||||
I have converted all the pickle files to safetensors: https://huggingface.co/Kijai/LivePortrait_safetensors/tree/main
|
||||
|
||||
They go here (and are automatically downloaded if the folder is not present) `ComfyUI/models/liveportrait`
|
||||
|
||||
# Face detectors
|
||||
|
||||
You can either use the original default Insightface, or Google's MediaPipe.
|
||||
|
||||
Biggest difference is the license: Insightface is strictly for NON-COMMERCIAL use.
|
||||
MediaPipe is a bit worse at detection, and can't run on GPU in Windows, though it's much faster on CPU compared to Insightface
|
||||
|
||||
Insightface is not automatically installed, if you wish to use it follow these instructions:
|
||||
Insightface is also required.
|
||||
If you have a working compile environment, installing it can be as easy as:
|
||||
|
||||
`pip install insightface`
|
||||
|
||||
@@ -1,7 +0,0 @@
|
||||
pyyaml
|
||||
numpy<=1.26.4
|
||||
opencv-python
|
||||
onnxruntime
|
||||
pykalman
|
||||
mediapipe
|
||||
onnx2torch
|
||||
+2
-5
@@ -1,7 +1,4 @@
|
||||
pyyaml
|
||||
numpy<=1.26.4
|
||||
numpy
|
||||
opencv-python
|
||||
onnxruntime-gpu
|
||||
pykalman
|
||||
mediapipe
|
||||
onnx2torch
|
||||
onnxruntime-gpu
|
||||
Reference in New Issue
Block a user