Compare commits
33
Commits
tensorrt_testing
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
 kijai
|
4d9dc6205b | ||
|
kijai
|
c74fb1d479 | ||
|
kijai
|
2d9497348e | ||
|
kijai
|
c2bb34d4f8 | ||
|
kijai
|
f3916f522a | ||
|
Jukka Seppänen
|
ee48647370 | ||
|
snomiao
|
59a230bfd3 | ||
|
Jukka Seppänen
|
3493b6aab6 | ||
|
Jukka Seppänen
|
451c2f9f02 | ||
|
Mario Klingemann
|
af14e360c0 | ||
|
Mario Klingemann
|
906cbf6e2f | ||
|
Jukka Seppänen
|
01b9705a7d | ||
|
kijai
|
73951396b6 | ||
|
kijai
|
3de788b1f2 | ||
|
kijai
|
5d0f993a1c | ||
|
Jukka Seppänen
|
6017c095be | ||
|
kijai
|
cc19214c18 | ||
|
kijai
|
e1e445823f | ||
|
kijai
|
4688258534 | ||
|
Mario Klingemann
|
0e04480507 | ||
|
Jukka Seppänen
|
2b23af0508 | ||
|
kijai
|
6cdb0ebde5 | ||
|
kijai
|
4d33a47ff2 | ||
|
kijai
|
490ff149bb | ||
|
kijai
|
7b26292418 | ||
|
kijai
|
8b28441984 | ||
|
Jukka Seppänen
|
0d0929530d | ||
|
kijai
|
3336dad236 | ||
|
kijai
|
fe2afdb3a2 | ||
|
kijai
|
3d195208db | ||
|
kijai
|
cae686d921 | ||
|
kijai
|
36e46cb5d3 | ||
|
Jukka Seppänen
|
3508b80c8b |
File diff suppressed because it is too large
Load Diff
+388
-295
@@ -1,70 +1,96 @@
|
||||
{
|
||||
"last_node_id": 31,
|
||||
"last_link_id": 68,
|
||||
"last_node_id": 208,
|
||||
"last_link_id": 480,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 4,
|
||||
"type": "LoadImage",
|
||||
"id": 204,
|
||||
"type": "LivePortraitLoadMediaPipeCropper",
|
||||
"pos": [
|
||||
138,
|
||||
323
|
||||
-1059,
|
||||
-767
|
||||
],
|
||||
"size": {
|
||||
"0": 272.85791015625,
|
||||
"1": 331.60894775390625
|
||||
"0": 315,
|
||||
"1": 82
|
||||
},
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"name": "cropper",
|
||||
"type": "LPCROPPER",
|
||||
"links": [
|
||||
59
|
||||
478
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": null,
|
||||
"shape": 3
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
"Node name for S&R": "LivePortraitLoadMediaPipeCropper"
|
||||
},
|
||||
"widgets_values": [
|
||||
"oldman.jpg",
|
||||
"image"
|
||||
"CPU",
|
||||
true
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 19,
|
||||
"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": [
|
||||
507,
|
||||
675
|
||||
-715,
|
||||
-617
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 242
|
||||
},
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"order": 6,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"link": 30
|
||||
"link": 466
|
||||
},
|
||||
{
|
||||
"name": "get_image_size",
|
||||
"type": "IMAGE",
|
||||
"link": 68
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "width_input",
|
||||
@@ -88,7 +114,7 @@
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
32
|
||||
434
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
@@ -112,290 +138,282 @@
|
||||
"widgets_values": [
|
||||
512,
|
||||
512,
|
||||
"nearest-exact",
|
||||
false,
|
||||
"lanczos",
|
||||
true,
|
||||
2,
|
||||
0,
|
||||
0
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 18,
|
||||
"type": "ImageConcatMulti",
|
||||
"id": 78,
|
||||
"type": "GetImageSizeAndCount",
|
||||
"pos": [
|
||||
860,
|
||||
679
|
||||
-364,
|
||||
-619
|
||||
],
|
||||
"size": {
|
||||
"0": 210,
|
||||
"1": 150
|
||||
"1": 86
|
||||
},
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"order": 7,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "image_1",
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"link": 32
|
||||
},
|
||||
{
|
||||
"name": "image_2",
|
||||
"type": "IMAGE",
|
||||
"link": 67
|
||||
"link": 434
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
64
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
2,
|
||||
"right",
|
||||
false,
|
||||
null
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 23,
|
||||
"type": "VHS_VideoCombine",
|
||||
"pos": [
|
||||
1098,
|
||||
240
|
||||
],
|
||||
"size": [
|
||||
1253.234130859375,
|
||||
940.6170654296875
|
||||
],
|
||||
"flags": {},
|
||||
"order": 6,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 64
|
||||
},
|
||||
{
|
||||
"name": "audio",
|
||||
"type": "VHS_AUDIO",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "meta_batch",
|
||||
"type": "VHS_BatchManager",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "vae",
|
||||
"type": "VAE",
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "Filenames",
|
||||
"type": "VHS_FILENAMES",
|
||||
"links": null,
|
||||
"shape": 3
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VHS_VideoCombine"
|
||||
},
|
||||
"widgets_values": {
|
||||
"frame_rate": 24,
|
||||
"loop_count": 0,
|
||||
"filename_prefix": "LivePortrait",
|
||||
"format": "video/h264-mp4",
|
||||
"pix_fmt": "yuv420p",
|
||||
"crf": 19,
|
||||
"save_metadata": true,
|
||||
"pingpong": false,
|
||||
"save_output": false,
|
||||
"videopreview": {
|
||||
"hidden": false,
|
||||
"paused": false,
|
||||
"params": {
|
||||
"filename": "LivePortrait_00001.mp4",
|
||||
"subfolder": "",
|
||||
"type": "temp",
|
||||
"format": "video/h264-mp4",
|
||||
"frame_rate": 24
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 1,
|
||||
"type": "DownloadAndLoadLivePortraitModels",
|
||||
"pos": [
|
||||
142,
|
||||
205
|
||||
],
|
||||
"size": {
|
||||
"0": 252,
|
||||
"1": 58
|
||||
},
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "live_portrait_pipe",
|
||||
"type": "LIVEPORTRAITPIPE",
|
||||
"links": [
|
||||
58
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "DownloadAndLoadLivePortraitModels"
|
||||
},
|
||||
"widgets_values": [
|
||||
"fp16"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 8,
|
||||
"type": "VHS_LoadVideo",
|
||||
"pos": [
|
||||
161,
|
||||
714
|
||||
],
|
||||
"size": [
|
||||
235.1999969482422,
|
||||
491.1999969482422
|
||||
],
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "meta_batch",
|
||||
"type": "VHS_BatchManager",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "vae",
|
||||
"type": "VAE",
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
30,
|
||||
60
|
||||
445,
|
||||
475
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "frame_count",
|
||||
"name": "512 width",
|
||||
"type": "INT",
|
||||
"links": null,
|
||||
"shape": 3
|
||||
},
|
||||
{
|
||||
"name": "audio",
|
||||
"type": "VHS_AUDIO",
|
||||
"name": "384 height",
|
||||
"type": "INT",
|
||||
"links": null,
|
||||
"shape": 3
|
||||
},
|
||||
{
|
||||
"name": "video_info",
|
||||
"type": "VHS_VIDEOINFO",
|
||||
"name": "1 count",
|
||||
"type": "INT",
|
||||
"links": null,
|
||||
"shape": 3
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VHS_LoadVideo"
|
||||
},
|
||||
"widgets_values": {
|
||||
"video": "d3.mp4",
|
||||
"force_rate": 0,
|
||||
"force_size": "Disabled",
|
||||
"custom_width": 512,
|
||||
"custom_height": 512,
|
||||
"frame_load_cap": 0,
|
||||
"skip_first_frames": 0,
|
||||
"select_every_nth": 1,
|
||||
"choose video to upload": "image",
|
||||
"videopreview": {
|
||||
"hidden": false,
|
||||
"paused": false,
|
||||
"params": {
|
||||
"frame_load_cap": 0,
|
||||
"skip_first_frames": 0,
|
||||
"force_rate": 0,
|
||||
"filename": "d3.mp4",
|
||||
"type": "input",
|
||||
"format": "video/mp4",
|
||||
"select_every_nth": 1
|
||||
}
|
||||
}
|
||||
"Node name for S&R": "GetImageSizeAndCount"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 30,
|
||||
"type": "LivePortraitProcess",
|
||||
"id": 207,
|
||||
"type": "Note",
|
||||
"pos": [
|
||||
500,
|
||||
249
|
||||
-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": 367.79998779296875,
|
||||
"1": 362
|
||||
"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": 58
|
||||
"link": 448
|
||||
},
|
||||
{
|
||||
"name": "crop_info",
|
||||
"type": "CROPINFO",
|
||||
"link": 449
|
||||
},
|
||||
{
|
||||
"name": "source_image",
|
||||
"type": "IMAGE",
|
||||
"link": 59
|
||||
"link": 475
|
||||
},
|
||||
{
|
||||
"name": "driving_images",
|
||||
"type": "IMAGE",
|
||||
"link": 60
|
||||
"link": 479,
|
||||
"slot_index": 3
|
||||
},
|
||||
{
|
||||
"name": "opt_retargeting_info",
|
||||
"type": "RETARGETINGINFO",
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "cropped_images",
|
||||
"name": "cropped_image",
|
||||
"type": "IMAGE",
|
||||
"links": [],
|
||||
"links": [
|
||||
470
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "full_images",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
67,
|
||||
68
|
||||
],
|
||||
"name": "output",
|
||||
"type": "LP_OUT",
|
||||
"links": [],
|
||||
"shape": 3,
|
||||
"slot_index": 1
|
||||
}
|
||||
@@ -403,85 +421,160 @@
|
||||
"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.3,
|
||||
2.34,
|
||||
0.099,
|
||||
0.148,
|
||||
0,
|
||||
-0.11,
|
||||
true,
|
||||
false,
|
||||
1,
|
||||
false,
|
||||
1,
|
||||
true,
|
||||
true,
|
||||
"CPU"
|
||||
"large-small",
|
||||
false
|
||||
]
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
30,
|
||||
8,
|
||||
434,
|
||||
165,
|
||||
0,
|
||||
19,
|
||||
78,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
32,
|
||||
19,
|
||||
445,
|
||||
78,
|
||||
0,
|
||||
18,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
58,
|
||||
1,
|
||||
0,
|
||||
30,
|
||||
0,
|
||||
"LIVEPORTRAITPIPE"
|
||||
],
|
||||
[
|
||||
59,
|
||||
4,
|
||||
0,
|
||||
30,
|
||||
1,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
60,
|
||||
8,
|
||||
0,
|
||||
30,
|
||||
189,
|
||||
2,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
64,
|
||||
18,
|
||||
446,
|
||||
1,
|
||||
0,
|
||||
23,
|
||||
189,
|
||||
0,
|
||||
"LIVEPORTRAITPIPE"
|
||||
],
|
||||
[
|
||||
448,
|
||||
1,
|
||||
0,
|
||||
190,
|
||||
0,
|
||||
"LIVEPORTRAITPIPE"
|
||||
],
|
||||
[
|
||||
449,
|
||||
189,
|
||||
1,
|
||||
190,
|
||||
1,
|
||||
"CROPINFO"
|
||||
],
|
||||
[
|
||||
466,
|
||||
196,
|
||||
0,
|
||||
165,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
67,
|
||||
30,
|
||||
1,
|
||||
18,
|
||||
1,
|
||||
470,
|
||||
190,
|
||||
0,
|
||||
198,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
68,
|
||||
30,
|
||||
1,
|
||||
19,
|
||||
475,
|
||||
78,
|
||||
0,
|
||||
190,
|
||||
2,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
478,
|
||||
204,
|
||||
0,
|
||||
189,
|
||||
1,
|
||||
"LPCROPPER"
|
||||
],
|
||||
[
|
||||
479,
|
||||
205,
|
||||
0,
|
||||
190,
|
||||
3,
|
||||
"IMAGE"
|
||||
]
|
||||
],
|
||||
@@ -489,10 +582,10 @@
|
||||
"config": {},
|
||||
"extra": {
|
||||
"ds": {
|
||||
"scale": 0.8264462809917354,
|
||||
"scale": 0.7513148009015781,
|
||||
"offset": {
|
||||
"0": 173.40487670898438,
|
||||
"1": -0.9636010527610779
|
||||
"0": 1468.4081568988054,
|
||||
"1": 1224.8414164288351
|
||||
}
|
||||
}
|
||||
},
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,29 @@
|
||||
BSD 3-Clause License
|
||||
|
||||
Copyright (c) 2017, Adrian Bulat
|
||||
All rights reserved.
|
||||
|
||||
Redistribution and use in source and binary forms, with or without
|
||||
modification, are permitted provided that the following conditions are met:
|
||||
|
||||
* Redistributions of source code must retain the above copyright notice, this
|
||||
list of conditions and the following disclaimer.
|
||||
|
||||
* Redistributions in binary form must reproduce the above copyright notice,
|
||||
this list of conditions and the following disclaimer in the documentation
|
||||
and/or other materials provided with the distribution.
|
||||
|
||||
* Neither the name of the copyright holder nor the names of its
|
||||
contributors may be used to endorse or promote products derived from
|
||||
this software without specific prior written permission.
|
||||
|
||||
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
@@ -0,0 +1,7 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
__author__ = """Adrian Bulat"""
|
||||
__email__ = 'adrian@adrianbulat.com'
|
||||
__version__ = '1.4.1'
|
||||
|
||||
from .api import FaceAlignment, LandmarksType, NetworkSize
|
||||
@@ -0,0 +1,305 @@
|
||||
import torch
|
||||
import warnings
|
||||
from enum import IntEnum
|
||||
from skimage import io
|
||||
import numpy as np
|
||||
from packaging import version
|
||||
from tqdm import tqdm
|
||||
import importlib
|
||||
|
||||
from .utils import *
|
||||
from .folder_data import FolderData
|
||||
|
||||
|
||||
class LandmarksType(IntEnum):
|
||||
"""Enum class defining the type of landmarks to detect.
|
||||
|
||||
``TWO_D`` - the detected points ``(x,y)`` are detected in a 2D space and follow the visible contour of the face
|
||||
``TWO_HALF_D`` - this points represent the projection of the 3D points into 3D
|
||||
``THREE_D`` - detect the points ``(x,y,z)``` in a 3D space
|
||||
|
||||
"""
|
||||
TWO_D = 1
|
||||
TWO_HALF_D = 2
|
||||
THREE_D = 3
|
||||
|
||||
|
||||
class NetworkSize(IntEnum):
|
||||
# TINY = 1
|
||||
# SMALL = 2
|
||||
# MEDIUM = 3
|
||||
LARGE = 4
|
||||
|
||||
|
||||
default_model_urls = {
|
||||
'2DFAN-4': 'https://www.adrianbulat.com/downloads/python-fan/2DFAN4-cd938726ad.zip',
|
||||
'3DFAN-4': 'https://www.adrianbulat.com/downloads/python-fan/3DFAN4-4a694010b9.zip',
|
||||
'depth': 'https://www.adrianbulat.com/downloads/python-fan/depth-6c4283c0e0.zip',
|
||||
}
|
||||
|
||||
models_urls = {
|
||||
'1.6': {
|
||||
'2DFAN-4': 'https://www.adrianbulat.com/downloads/python-fan/2DFAN4_1.6-c827573f02.zip',
|
||||
'3DFAN-4': 'https://www.adrianbulat.com/downloads/python-fan/3DFAN4_1.6-ec5cf40a1d.zip',
|
||||
'depth': 'https://www.adrianbulat.com/downloads/python-fan/depth_1.6-2aa3f18772.zip',
|
||||
},
|
||||
'1.5': {
|
||||
'2DFAN-4': 'https://www.adrianbulat.com/downloads/python-fan/2DFAN4_1.5-a60332318a.zip',
|
||||
'3DFAN-4': 'https://www.adrianbulat.com/downloads/python-fan/3DFAN4_1.5-176570af4d.zip',
|
||||
'depth': 'https://www.adrianbulat.com/downloads/python-fan/depth_1.5-bc10f98e39.zip',
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class FaceAlignment:
|
||||
def __init__(self, landmarks_type, network_size=NetworkSize.LARGE,
|
||||
device='cuda', dtype=torch.float32, flip_input=False, face_detector='sfd', face_detector_kwargs=None, verbose=False):
|
||||
self.device = device
|
||||
self.flip_input = flip_input
|
||||
self.landmarks_type = landmarks_type
|
||||
self.verbose = verbose
|
||||
self.dtype = dtype
|
||||
|
||||
if version.parse(torch.__version__) < version.parse('1.5.0'):
|
||||
raise ImportError(f'Unsupported pytorch version detected. Minimum supported version of pytorch: 1.5.0\
|
||||
Either upgrade (recommended) your pytorch setup, or downgrade to face-alignment 1.2.0')
|
||||
|
||||
network_size = int(network_size)
|
||||
pytorch_version = torch.__version__
|
||||
if 'dev' in pytorch_version:
|
||||
pytorch_version = pytorch_version.rsplit('.', 2)[0]
|
||||
else:
|
||||
pytorch_version = pytorch_version.rsplit('.', 1)[0]
|
||||
|
||||
#if 'cuda' in device:
|
||||
# torch.backends.cudnn.benchmark = True
|
||||
|
||||
# Get the face detector
|
||||
package_directory_name = os.path.basename(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
face_detector_module = importlib.import_module('.face_alignment.detection.' + face_detector, package=package_directory_name)
|
||||
|
||||
|
||||
face_detector_kwargs = face_detector_kwargs or {}
|
||||
self.face_detector = face_detector_module.FaceDetector(device=device, verbose=verbose, **face_detector_kwargs)
|
||||
|
||||
# Initialise the face alignemnt networks
|
||||
if landmarks_type == LandmarksType.TWO_D:
|
||||
network_name = '2DFAN-' + str(network_size)
|
||||
else:
|
||||
network_name = '3DFAN-' + str(network_size)
|
||||
self.face_alignment_net = torch.jit.load(
|
||||
load_file_from_url(models_urls.get(pytorch_version, default_model_urls)[network_name]))
|
||||
|
||||
self.face_alignment_net.to(device, dtype=dtype)
|
||||
self.face_alignment_net.eval()
|
||||
|
||||
# Initialiase the depth prediciton network
|
||||
if landmarks_type == LandmarksType.THREE_D:
|
||||
self.depth_prediciton_net = torch.jit.load(
|
||||
load_file_from_url(models_urls.get(pytorch_version, default_model_urls)['depth']))
|
||||
|
||||
self.depth_prediciton_net.to(device, dtype=dtype)
|
||||
self.depth_prediciton_net.eval()
|
||||
|
||||
def get_landmarks(self, image_or_path, detected_faces=None, return_bboxes=False, return_landmark_score=False):
|
||||
"""Deprecated, please use get_landmarks_from_image
|
||||
|
||||
Arguments:
|
||||
image_or_path {string or numpy.array or torch.tensor} -- The input image or path to it
|
||||
|
||||
Keyword Arguments:
|
||||
detected_faces {list of numpy.array} -- list of bounding boxes, one for each face found
|
||||
in the image (default: {None})
|
||||
return_bboxes {boolean} -- If True, return the face bounding boxes in addition to the keypoints.
|
||||
return_landmark_score {boolean} -- If True, return the keypoint scores along with the keypoints.
|
||||
"""
|
||||
return self.get_landmarks_from_image(image_or_path, detected_faces, return_bboxes, return_landmark_score)
|
||||
|
||||
@torch.no_grad()
|
||||
def get_landmarks_from_image(self, image_or_path, detected_faces=None, return_bboxes=False,
|
||||
return_landmark_score=False):
|
||||
"""Predict the landmarks for each face present in the image.
|
||||
|
||||
This function predicts a set of 68 2D or 3D images, one for each image present.
|
||||
If detect_faces is None the method will also run a face detector.
|
||||
|
||||
Arguments:
|
||||
image_or_path {string or numpy.array or torch.tensor} -- The input image or path to it.
|
||||
|
||||
Keyword Arguments:
|
||||
detected_faces {list of numpy.array} -- list of bounding boxes, one for each face found
|
||||
in the image (default: {None})
|
||||
return_bboxes {boolean} -- If True, return the face bounding boxes in addition to the keypoints.
|
||||
return_landmark_score {boolean} -- If True, return the keypoint scores along with the keypoints.
|
||||
|
||||
Return:
|
||||
result:
|
||||
1. if both return_bboxes and return_landmark_score are False, result will be:
|
||||
landmark
|
||||
2. Otherwise, result will be one of the following, depending on the actual value of return_* arguments.
|
||||
(landmark, landmark_score, detected_face)
|
||||
(landmark, None, detected_face)
|
||||
(landmark, landmark_score, None )
|
||||
"""
|
||||
image = get_image(image_or_path)
|
||||
|
||||
if detected_faces is None:
|
||||
detected_faces = self.face_detector.detect_from_image(image.copy())
|
||||
|
||||
if len(detected_faces) == 0:
|
||||
warnings.warn("No faces were detected.")
|
||||
if return_bboxes or return_landmark_score:
|
||||
return None, None, None
|
||||
else:
|
||||
return None
|
||||
|
||||
landmarks = []
|
||||
landmarks_scores = []
|
||||
for i, d in enumerate(detected_faces):
|
||||
center = torch.tensor(
|
||||
[d[2] - (d[2] - d[0]) / 2.0, d[3] - (d[3] - d[1]) / 2.0])
|
||||
center[1] = center[1] - (d[3] - d[1]) * 0.12
|
||||
scale = (d[2] - d[0] + d[3] - d[1]) / self.face_detector.reference_scale
|
||||
|
||||
inp = crop(image, center, scale)
|
||||
inp = torch.from_numpy(inp.transpose(
|
||||
(2, 0, 1))).float()
|
||||
|
||||
inp = inp.to(self.device, dtype=self.dtype)
|
||||
inp.div_(255.0).unsqueeze_(0)
|
||||
|
||||
out = self.face_alignment_net(inp).detach()
|
||||
if self.flip_input:
|
||||
out += flip(self.face_alignment_net(flip(inp)).detach(), is_label=True)
|
||||
out = out.to(device='cpu', dtype=torch.float32).numpy()
|
||||
|
||||
pts, pts_img, scores = get_preds_fromhm(out, center.numpy(), scale)
|
||||
pts, pts_img = torch.from_numpy(pts), torch.from_numpy(pts_img)
|
||||
pts, pts_img = pts.view(68, 2) * 4, pts_img.view(68, 2)
|
||||
scores = scores.squeeze(0)
|
||||
|
||||
if self.landmarks_type == LandmarksType.THREE_D:
|
||||
heatmaps = np.zeros((68, 256, 256), dtype=np.float32)
|
||||
for i in range(68):
|
||||
if pts[i, 0] > 0 and pts[i, 1] > 0:
|
||||
heatmaps[i] = draw_gaussian(
|
||||
heatmaps[i], pts[i], 2)
|
||||
heatmaps = torch.from_numpy(
|
||||
heatmaps).unsqueeze_(0)
|
||||
|
||||
heatmaps = heatmaps.to(self.device, dtype=self.dtype)
|
||||
depth_pred = self.depth_prediciton_net(
|
||||
torch.cat((inp, heatmaps), 1)).data.cpu().view(68, 1).to(dtype=torch.float32)
|
||||
pts_img = torch.cat(
|
||||
(pts_img, depth_pred * (1.0 / (256.0 / (200.0 * scale)))), 1)
|
||||
|
||||
landmarks.append(pts_img.numpy())
|
||||
landmarks_scores.append(scores)
|
||||
|
||||
if not return_bboxes:
|
||||
detected_faces = None
|
||||
if not return_landmark_score:
|
||||
landmarks_scores = None
|
||||
if return_bboxes or return_landmark_score:
|
||||
return landmarks, landmarks_scores, detected_faces
|
||||
else:
|
||||
return landmarks
|
||||
|
||||
@torch.no_grad()
|
||||
def get_landmarks_from_batch(self, image_batch, detected_faces=None, return_bboxes=False,
|
||||
return_landmark_score=False):
|
||||
"""Predict the landmarks for each face present in the image.
|
||||
|
||||
This function predicts a set of 68 2D or 3D images, one for each image in a batch in parallel.
|
||||
If detect_faces is None the method will also run a face detector.
|
||||
|
||||
Arguments:
|
||||
image_batch {torch.tensor} -- The input images batch
|
||||
|
||||
Keyword Arguments:
|
||||
detected_faces {list of numpy.array} -- list of bounding boxes, one for each face found
|
||||
in the image (default: {None})
|
||||
return_bboxes {boolean} -- If True, return the face bounding boxes in addition to the keypoints.
|
||||
return_landmark_score {boolean} -- If True, return the keypoint scores along with the keypoints.
|
||||
|
||||
Return:
|
||||
result:
|
||||
1. if both return_bboxes and return_landmark_score are False, result will be:
|
||||
landmarks
|
||||
2. Otherwise, result will be one of the following, depending on the actual value of return_* arguments.
|
||||
(landmark, landmark_score, detected_face)
|
||||
(landmark, None, detected_face)
|
||||
(landmark, landmark_score, None )
|
||||
"""
|
||||
|
||||
if detected_faces is None:
|
||||
detected_faces = self.face_detector.detect_from_batch(image_batch)
|
||||
|
||||
if len(detected_faces) == 0:
|
||||
warnings.warn("No faces were detected.")
|
||||
if return_bboxes or return_landmark_score:
|
||||
return None, None, None
|
||||
else:
|
||||
return None
|
||||
|
||||
landmarks = []
|
||||
landmarks_scores_list = []
|
||||
# A batch for each frame
|
||||
for i, faces in enumerate(detected_faces):
|
||||
res = self.get_landmarks_from_image(
|
||||
image_batch[i].cpu().numpy().transpose(1, 2, 0),
|
||||
detected_faces=faces,
|
||||
return_landmark_score=return_landmark_score,
|
||||
)
|
||||
if return_landmark_score:
|
||||
landmark_set, landmarks_scores, _ = res
|
||||
landmarks_scores_list.append(landmarks_scores)
|
||||
else:
|
||||
landmark_set = res
|
||||
# Bacward compatibility
|
||||
if landmark_set is not None:
|
||||
landmark_set = np.concatenate(landmark_set, axis=0)
|
||||
else:
|
||||
landmark_set = []
|
||||
landmarks.append(landmark_set)
|
||||
|
||||
if not return_bboxes:
|
||||
detected_faces = None
|
||||
if not return_landmark_score:
|
||||
landmarks_scores_list = None
|
||||
if return_bboxes or return_landmark_score:
|
||||
return landmarks, landmarks_scores_list, detected_faces
|
||||
else:
|
||||
return landmarks
|
||||
|
||||
def get_landmarks_from_directory(self, path, extensions=['.jpg', '.png'], recursive=True, show_progress_bar=True,
|
||||
return_bboxes=False, return_landmark_score=False):
|
||||
"""Scan a directory for images with a given extension type(s) and predict the landmarks for each
|
||||
face present in the images found.
|
||||
|
||||
Arguments:
|
||||
path {str} -- path to the target directory containing the images
|
||||
|
||||
Keyword Arguments:
|
||||
extensions {list of str} -- list containing the image extensions considered (default: ['.jpg', '.png'])
|
||||
recursive {boolean} -- If True, scans for images recursively (default: True)
|
||||
show_progress_bar {boolean} -- If True displays a progress bar (default: True)
|
||||
return_bboxes {boolean} -- If True, return the face bounding boxes in addition to the keypoints.
|
||||
return_landmark_score {boolean} -- If True, return the keypoint scores along with the keypoints.
|
||||
"""
|
||||
dataset = FolderData(path, self.face_detector.tensor_or_path_to_ndarray, extensions, recursive, self.verbose)
|
||||
dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=2, prefetch_factor=4)
|
||||
|
||||
predictions = {}
|
||||
for (image_path, image) in tqdm(dataloader, disable=not show_progress_bar):
|
||||
image_path, image = image_path[0], image[0]
|
||||
bounding_boxes = self.face_detector.detect_from_image(image)
|
||||
if return_bboxes or return_landmark_score:
|
||||
preds, bbox, score = self.get_landmarks_from_image(
|
||||
image, bounding_boxes, return_bboxes=return_bboxes, return_landmark_score=return_landmark_score)
|
||||
predictions[image_path] = (preds, bbox, score)
|
||||
else:
|
||||
preds = self.get_landmarks_from_image(image, bounding_boxes)
|
||||
predictions[image_path] = preds
|
||||
|
||||
return predictions
|
||||
@@ -0,0 +1 @@
|
||||
from .core import FaceDetector
|
||||
@@ -0,0 +1 @@
|
||||
from .blazeface_detector import BlazeFaceDetector as FaceDetector
|
||||
@@ -0,0 +1,69 @@
|
||||
from torch.utils.model_zoo import load_url
|
||||
|
||||
from ..core import FaceDetector
|
||||
from ...utils import load_file_from_url
|
||||
|
||||
from .net_blazeface import BlazeFace
|
||||
from .detect import *
|
||||
|
||||
models_urls = {
|
||||
'blazeface_weights': 'https://github.com/hollance/BlazeFace-PyTorch/blob/master/blazeface.pth?raw=true',
|
||||
'blazeface_anchors': 'https://github.com/hollance/BlazeFace-PyTorch/blob/master/anchors.npy?raw=true',
|
||||
'blazeface_back_weights': 'https://github.com/hollance/BlazeFace-PyTorch/blob/master/blazefaceback.pth?raw=true',
|
||||
'blazeface_back_anchors': 'https://github.com/hollance/BlazeFace-PyTorch/blob/master/anchorsback.npy?raw=true'
|
||||
}
|
||||
|
||||
|
||||
class BlazeFaceDetector(FaceDetector):
|
||||
def __init__(self, device, back_model=False, path_to_detector=None, path_to_anchor=None, verbose=False,
|
||||
min_score_thresh=0.5, min_suppression_threshold=0.3):
|
||||
super(BlazeFaceDetector, self).__init__(device, verbose)
|
||||
|
||||
# Initialise the face detector
|
||||
self.back_model = back_model
|
||||
if path_to_detector is None:
|
||||
if back_model:
|
||||
model_weights = load_url(models_urls['blazeface_back_weights'])
|
||||
model_anchors = np.load(load_file_from_url(models_urls['blazeface_back_anchors']))
|
||||
else:
|
||||
model_weights = load_url(models_urls['blazeface_weights'])
|
||||
model_anchors = np.load(load_file_from_url(models_urls['blazeface_anchors']))
|
||||
else:
|
||||
model_weights = torch.load(path_to_detector)
|
||||
model_anchors = np.load(path_to_anchor)
|
||||
|
||||
self.face_detector = BlazeFace(back_model=back_model)
|
||||
self.face_detector.load_state_dict(model_weights)
|
||||
self.face_detector.load_anchors_from_npy(model_anchors, device)
|
||||
|
||||
# Optionally change the thresholds:
|
||||
self.face_detector.min_score_thresh = min_score_thresh
|
||||
self.face_detector.min_suppression_threshold = min_suppression_threshold
|
||||
|
||||
self.face_detector.to(device)
|
||||
self.face_detector.eval()
|
||||
|
||||
def detect_from_image(self, tensor_or_path):
|
||||
image = self.tensor_or_path_to_ndarray(tensor_or_path)
|
||||
|
||||
image_size = 256 if self.back_model else 128
|
||||
bboxlist = detect(self.face_detector, image, target_size=image_size, device=self.device)[0]
|
||||
|
||||
return bboxlist
|
||||
|
||||
def detect_from_batch(self, tensor):
|
||||
image_size = 256 if self.back_model else 128
|
||||
bboxlists = batch_detect(self.face_detector, tensor, target_size=image_size, device=self.device)
|
||||
return bboxlists
|
||||
|
||||
@property
|
||||
def reference_scale(self):
|
||||
return 195
|
||||
|
||||
@property
|
||||
def reference_x_shift(self):
|
||||
return 0
|
||||
|
||||
@property
|
||||
def reference_y_shift(self):
|
||||
return 0
|
||||
@@ -0,0 +1,74 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from .utils import *
|
||||
|
||||
|
||||
def detect(net, img, target_size=128, device='cuda'):
|
||||
H, W, C = img.shape
|
||||
orig_size = min(H, W)
|
||||
img, (xshift, yshift) = resize_and_crop_image(img, target_size)
|
||||
preds = net.predict_on_image(img)
|
||||
|
||||
if 0 == len(preds):
|
||||
return [[]]
|
||||
|
||||
shift = np.array([xshift, yshift] * 2)
|
||||
scores = preds[:, -1:]
|
||||
|
||||
# TODO: ugly
|
||||
# reverses, x and y to adapt with face-alignment code
|
||||
locs = np.concatenate((preds[:, 1:2], preds[:, 0:1], preds[:, 3:4], preds[:, 2:3]), axis=1)
|
||||
return [np.concatenate((locs * orig_size + shift, scores), axis=1)]
|
||||
|
||||
|
||||
def batch_detect(net, img_batch, target_size=128, device='cuda'):
|
||||
"""
|
||||
Inputs:
|
||||
- net: BlazeFace model
|
||||
- img_batch: a numpy array or tensor of shape (Batch size, Channels, Height, Width)
|
||||
- target_size: target size of the input image
|
||||
Outputs:
|
||||
- list of 2-dim numpy arrays with shape (faces_on_this_image, 5): x1, y1, x2, y2, confidence
|
||||
(x1, y1) - top left corner, (x2, y2) - bottom right corner
|
||||
"""
|
||||
B, C, H, W = img_batch.shape
|
||||
orig_size = min(H, W)
|
||||
|
||||
if isinstance(img_batch, torch.Tensor):
|
||||
img_batch = img_batch.cpu().numpy()
|
||||
|
||||
img_batch = img_batch.transpose((0, 2, 3, 1))
|
||||
|
||||
imgs, (xshift, yshift) = resize_and_crop_batch(img_batch, target_size)
|
||||
preds = net.predict_on_batch(imgs)
|
||||
bboxlists = []
|
||||
for pred in preds:
|
||||
shift = np.array([xshift, yshift] * 2)
|
||||
scores = pred[:, -1:]
|
||||
locs = np.concatenate((pred[:, 1:2], pred[:, 0:1], pred[:, 3:4], pred[:, 2:3]), axis=1)
|
||||
bboxlists.append(np.concatenate((locs * orig_size + shift, scores), axis=1))
|
||||
|
||||
return bboxlists
|
||||
|
||||
|
||||
def flip_detect(net, img, device):
|
||||
img = cv2.flip(img, 1)
|
||||
b = detect(net, img, device)
|
||||
|
||||
bboxlist = np.zeros(b.shape)
|
||||
bboxlist[:, 0] = img.shape[1] - b[:, 2]
|
||||
bboxlist[:, 1] = b[:, 1]
|
||||
bboxlist[:, 2] = img.shape[1] - b[:, 0]
|
||||
bboxlist[:, 3] = b[:, 3]
|
||||
bboxlist[:, 4] = b[:, 4]
|
||||
return bboxlist
|
||||
|
||||
|
||||
def pts_to_bb(pts):
|
||||
min_x, min_y = np.min(pts, axis=0)
|
||||
max_x, max_y = np.max(pts, axis=0)
|
||||
return np.array([min_x, min_y, max_x, max_y])
|
||||
@@ -0,0 +1,482 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class BlazeBlock(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1):
|
||||
super(BlazeBlock, self).__init__()
|
||||
|
||||
self.stride = stride
|
||||
self.channel_pad = out_channels - in_channels
|
||||
|
||||
# TFLite uses slightly different padding than PyTorch
|
||||
# on the depthwise conv layer when the stride is 2.
|
||||
if stride == 2:
|
||||
self.max_pool = nn.MaxPool2d(kernel_size=stride, stride=stride)
|
||||
padding = 0
|
||||
else:
|
||||
padding = (kernel_size - 1) // 2
|
||||
|
||||
self.convs = nn.Sequential(
|
||||
nn.Conv2d(in_channels=in_channels, out_channels=in_channels,
|
||||
kernel_size=kernel_size, stride=stride, padding=padding,
|
||||
groups=in_channels, bias=True),
|
||||
nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
|
||||
kernel_size=1, stride=1, padding=0, bias=True),
|
||||
)
|
||||
|
||||
self.act = nn.ReLU(inplace=True)
|
||||
|
||||
def forward(self, x):
|
||||
if self.stride == 2:
|
||||
h = F.pad(x, (0, 2, 0, 2), "constant", 0)
|
||||
x = self.max_pool(x)
|
||||
else:
|
||||
h = x
|
||||
|
||||
if self.channel_pad > 0:
|
||||
x = F.pad(x, (0, 0, 0, 0, 0, self.channel_pad), "constant", 0)
|
||||
|
||||
return self.act(self.convs(h) + x)
|
||||
|
||||
class FinalBlazeBlock(nn.Module):
|
||||
def __init__(self, channels, kernel_size=3):
|
||||
super(FinalBlazeBlock, self).__init__()
|
||||
# TFLite uses slightly different padding than PyTorch
|
||||
# on the depthwise conv layer when the stride is 2.
|
||||
self.convs = nn.Sequential(
|
||||
nn.Conv2d(in_channels=channels, out_channels=channels,
|
||||
kernel_size=kernel_size, stride=2, padding=0,
|
||||
groups=channels, bias=True),
|
||||
nn.Conv2d(in_channels=channels, out_channels=channels,
|
||||
kernel_size=1, stride=1, padding=0, bias=True),
|
||||
)
|
||||
|
||||
self.act = nn.ReLU(inplace=True)
|
||||
|
||||
def forward(self, x):
|
||||
h = F.pad(x, (0, 2, 0, 2), "constant", 0)
|
||||
|
||||
return self.act(self.convs(h))
|
||||
|
||||
|
||||
class BlazeFace(nn.Module):
|
||||
"""The BlazeFace face detection model from MediaPipe.
|
||||
|
||||
The version from MediaPipe is simpler than the one in the paper;
|
||||
it does not use the "double" BlazeBlocks.
|
||||
|
||||
Because we won't be training this model, it doesn't need to have
|
||||
batchnorm layers. These have already been "folded" into the conv
|
||||
weights by TFLite.
|
||||
|
||||
The conversion to PyTorch is fairly straightforward, but there are
|
||||
some small differences between TFLite and PyTorch in how they handle
|
||||
padding on conv layers with stride 2.
|
||||
|
||||
This version works on batches, while the MediaPipe version can only
|
||||
handle a single image at a time.
|
||||
|
||||
Based on code from https://github.com/tkat0/PyTorch_BlazeFace/ and
|
||||
https://github.com/google/mediapipe/
|
||||
"""
|
||||
|
||||
def __init__(self, back_model=False):
|
||||
super(BlazeFace, self).__init__()
|
||||
|
||||
# These are the settings from the MediaPipe example graph
|
||||
# mediapipe/graphs/face_detection/face_detection_mobile_gpu.pbtxt
|
||||
# and mediapipe/graphs/face_detection/face_detection_back_mobile_gpu.pbtxt
|
||||
self.num_classes = 1
|
||||
self.num_anchors = 896
|
||||
self.num_coords = 16
|
||||
self.score_clipping_thresh = 100.0
|
||||
self.back_model = back_model
|
||||
if back_model:
|
||||
self.x_scale = 256.0
|
||||
self.y_scale = 256.0
|
||||
self.h_scale = 256.0
|
||||
self.w_scale = 256.0
|
||||
self.min_score_thresh = 0.65
|
||||
else:
|
||||
self.x_scale = 128.0
|
||||
self.y_scale = 128.0
|
||||
self.h_scale = 128.0
|
||||
self.w_scale = 128.0
|
||||
self.min_score_thresh = 0.75
|
||||
self.min_suppression_threshold = 0.3
|
||||
|
||||
self._define_layers()
|
||||
|
||||
def _define_back_model_layers(self):
|
||||
self.backbone = nn.Sequential(
|
||||
nn.Conv2d(in_channels=3, out_channels=24, kernel_size=5, stride=2, padding=0, bias=True),
|
||||
nn.ReLU(inplace=True),
|
||||
*[BlazeBlock(24, 24) for _ in range(7)],
|
||||
BlazeBlock(24, 24, stride=2),
|
||||
*[BlazeBlock(24, 24) for _ in range(7)],
|
||||
BlazeBlock(24, 48, stride=2),
|
||||
*[BlazeBlock(48, 48) for _ in range(7)],
|
||||
BlazeBlock(48, 96, stride=2),
|
||||
*[BlazeBlock(96, 96) for _ in range(7)],
|
||||
)
|
||||
self.final = FinalBlazeBlock(96)
|
||||
self.classifier_8 = nn.Conv2d(96, 2, 1, bias=True)
|
||||
self.classifier_16 = nn.Conv2d(96, 6, 1, bias=True)
|
||||
|
||||
self.regressor_8 = nn.Conv2d(96, 32, 1, bias=True)
|
||||
self.regressor_16 = nn.Conv2d(96, 96, 1, bias=True)
|
||||
|
||||
def _define_front_model_layers(self):
|
||||
self.backbone1 = nn.Sequential(
|
||||
nn.Conv2d(in_channels=3, out_channels=24, kernel_size=5,
|
||||
stride=2, padding=0, bias=True),
|
||||
nn.ReLU(inplace=True),
|
||||
|
||||
BlazeBlock(24, 24),
|
||||
BlazeBlock(24, 28),
|
||||
BlazeBlock(28, 32, stride=2),
|
||||
BlazeBlock(32, 36),
|
||||
BlazeBlock(36, 42),
|
||||
BlazeBlock(42, 48, stride=2),
|
||||
BlazeBlock(48, 56),
|
||||
BlazeBlock(56, 64),
|
||||
BlazeBlock(64, 72),
|
||||
BlazeBlock(72, 80),
|
||||
BlazeBlock(80, 88),
|
||||
)
|
||||
|
||||
self.backbone2 = nn.Sequential(
|
||||
BlazeBlock(88, 96, stride=2),
|
||||
BlazeBlock(96, 96),
|
||||
BlazeBlock(96, 96),
|
||||
BlazeBlock(96, 96),
|
||||
BlazeBlock(96, 96),
|
||||
)
|
||||
|
||||
self.classifier_8 = nn.Conv2d(88, 2, 1, bias=True)
|
||||
self.classifier_16 = nn.Conv2d(96, 6, 1, bias=True)
|
||||
|
||||
self.regressor_8 = nn.Conv2d(88, 32, 1, bias=True)
|
||||
self.regressor_16 = nn.Conv2d(96, 96, 1, bias=True)
|
||||
|
||||
|
||||
def _define_layers(self):
|
||||
if self.back_model:
|
||||
self._define_back_model_layers()
|
||||
else:
|
||||
self._define_front_model_layers()
|
||||
|
||||
def forward(self, x):
|
||||
# TFLite uses slightly different padding on the first conv layer
|
||||
# than PyTorch, so do it manually.
|
||||
x = F.pad(x, (1, 2, 1, 2), "constant", 0)
|
||||
|
||||
b = x.shape[0] # batch size, needed for reshaping later
|
||||
|
||||
if self.back_model:
|
||||
x = self.backbone(x) # (b, 16, 16, 96)
|
||||
h = self.final(x) # (b, 8, 8, 96)
|
||||
else:
|
||||
x = self.backbone1(x) # (b, 88, 16, 16)
|
||||
h = self.backbone2(x) # (b, 96, 8, 8)
|
||||
|
||||
# Note: Because PyTorch is NCHW but TFLite is NHWC, we need to
|
||||
# permute the output from the conv layers before reshaping it.
|
||||
|
||||
c1 = self.classifier_8(x) # (b, 2, 16, 16)
|
||||
c1 = c1.permute(0, 2, 3, 1) # (b, 16, 16, 2)
|
||||
c1 = c1.reshape(b, -1, 1) # (b, 512, 1)
|
||||
|
||||
c2 = self.classifier_16(h) # (b, 6, 8, 8)
|
||||
c2 = c2.permute(0, 2, 3, 1) # (b, 8, 8, 6)
|
||||
c2 = c2.reshape(b, -1, 1) # (b, 384, 1)
|
||||
|
||||
c = torch.cat((c1, c2), dim=1) # (b, 896, 1)
|
||||
|
||||
r1 = self.regressor_8(x) # (b, 32, 16, 16)
|
||||
r1 = r1.permute(0, 2, 3, 1) # (b, 16, 16, 32)
|
||||
r1 = r1.reshape(b, -1, 16) # (b, 512, 16)
|
||||
|
||||
r2 = self.regressor_16(h) # (b, 96, 8, 8)
|
||||
r2 = r2.permute(0, 2, 3, 1) # (b, 8, 8, 96)
|
||||
r2 = r2.reshape(b, -1, 16) # (b, 384, 16)
|
||||
|
||||
r = torch.cat((r1, r2), dim=1) # (b, 896, 16)
|
||||
return [r, c]
|
||||
|
||||
def _device(self):
|
||||
"""Which device (CPU or GPU) is being used by this model?"""
|
||||
return self.classifier_8.weight.device
|
||||
|
||||
def load_weights(self, path):
|
||||
self.load_state_dict(torch.load(path))
|
||||
self.eval()
|
||||
|
||||
def load_anchors(self, path, device=None):
|
||||
device = device or self._device()
|
||||
self.anchors = torch.tensor(
|
||||
np.load(path), dtype=torch.float32, device=device)
|
||||
assert(self.anchors.ndimension() == 2)
|
||||
assert(self.anchors.shape[0] == self.num_anchors)
|
||||
assert(self.anchors.shape[1] == 4)
|
||||
|
||||
def load_anchors_from_npy(self, arr, device=None):
|
||||
device = device or self._device()
|
||||
self.anchors = torch.tensor(
|
||||
arr, dtype=torch.float32, device=device)
|
||||
assert(self.anchors.ndimension() == 2)
|
||||
assert(self.anchors.shape[0] == self.num_anchors)
|
||||
assert(self.anchors.shape[1] == 4)
|
||||
|
||||
def _preprocess(self, x):
|
||||
"""Converts the image pixels to the range [-1, 1]."""
|
||||
return x.float() / 127.5 - 1.0
|
||||
|
||||
def predict_on_image(self, img):
|
||||
"""Makes a prediction on a single image.
|
||||
|
||||
Arguments:
|
||||
img: a NumPy array of shape (H, W, 3) or a PyTorch tensor of
|
||||
shape (3, H, W). The image's height and width should be
|
||||
128 pixels.
|
||||
|
||||
Returns:
|
||||
A tensor with face detections.
|
||||
"""
|
||||
if isinstance(img, np.ndarray):
|
||||
img = torch.from_numpy(img).permute((2, 0, 1))
|
||||
|
||||
return self.predict_on_batch(img.unsqueeze(0))[0]
|
||||
|
||||
def predict_on_batch(self, x):
|
||||
"""Makes a prediction on a batch of images.
|
||||
|
||||
Arguments:
|
||||
x: a NumPy array of shape (b, H, W, 3) or a PyTorch tensor of
|
||||
shape (b, 3, H, W). The height and width should be 128 pixels.
|
||||
|
||||
Returns:
|
||||
A list containing a tensor of face detections for each image in
|
||||
the batch. If no faces are found for an image, returns a tensor
|
||||
of shape (0, 17).
|
||||
|
||||
Each face detection is a PyTorch tensor consisting of 17 numbers:
|
||||
- ymin, xmin, ymax, xmax
|
||||
- x,y-coordinates for the 6 keypoints
|
||||
- confidence score
|
||||
"""
|
||||
if isinstance(x, np.ndarray):
|
||||
x = torch.from_numpy(x).permute((0, 3, 1, 2))
|
||||
|
||||
assert x.shape[1] == 3
|
||||
if self.back_model:
|
||||
assert x.shape[2] == 256
|
||||
assert x.shape[3] == 256
|
||||
else:
|
||||
assert x.shape[2] == 128
|
||||
assert x.shape[3] == 128
|
||||
|
||||
# 1. Preprocess the images into tensors:
|
||||
x = x.to(self._device())
|
||||
x = self._preprocess(x)
|
||||
|
||||
# 2. Run the neural network:
|
||||
with torch.inference_mode():
|
||||
out = self.__call__(x)
|
||||
|
||||
# 3. Postprocess the raw predictions:
|
||||
detections = self._tensors_to_detections(out[0], out[1], self.anchors)
|
||||
|
||||
# 4. Non-maximum suppression to remove overlapping detections:
|
||||
filtered_detections = []
|
||||
for i in range(len(detections)):
|
||||
faces = self._weighted_non_max_suppression(detections[i])
|
||||
faces = torch.stack(faces) if len(
|
||||
faces) > 0 else torch.zeros((0, 17))
|
||||
filtered_detections.append(faces)
|
||||
|
||||
return filtered_detections
|
||||
|
||||
def _tensors_to_detections(self, raw_box_tensor, raw_score_tensor, anchors):
|
||||
"""The output of the neural network is a tensor of shape (b, 896, 16)
|
||||
containing the bounding box regressor predictions, as well as a tensor
|
||||
of shape (b, 896, 1) with the classification confidences.
|
||||
|
||||
This function converts these two "raw" tensors into proper detections.
|
||||
Returns a list of (num_detections, 17) tensors, one for each image in
|
||||
the batch.
|
||||
|
||||
This is based on the source code from:
|
||||
mediapipe/calculators/tflite/tflite_tensors_to_detections_calculator.cc
|
||||
mediapipe/calculators/tflite/tflite_tensors_to_detections_calculator.proto
|
||||
"""
|
||||
assert raw_box_tensor.ndimension() == 3
|
||||
assert raw_box_tensor.shape[1] == self.num_anchors
|
||||
assert raw_box_tensor.shape[2] == self.num_coords
|
||||
|
||||
assert raw_score_tensor.ndimension() == 3
|
||||
assert raw_score_tensor.shape[1] == self.num_anchors
|
||||
assert raw_score_tensor.shape[2] == self.num_classes
|
||||
|
||||
assert raw_box_tensor.shape[0] == raw_score_tensor.shape[0]
|
||||
|
||||
detection_boxes = self._decode_boxes(raw_box_tensor, anchors)
|
||||
|
||||
thresh = self.score_clipping_thresh
|
||||
raw_score_tensor = raw_score_tensor.clamp(-thresh, thresh)
|
||||
detection_scores = raw_score_tensor.sigmoid().squeeze(dim=-1)
|
||||
|
||||
# Note: we stripped off the last dimension from the scores tensor
|
||||
# because there is only has one class. Now we can simply use a mask
|
||||
# to filter out the boxes with too low confidence.
|
||||
mask = detection_scores >= self.min_score_thresh
|
||||
|
||||
# Because each image from the batch can have a different number of
|
||||
# detections, process them one at a time using a loop.
|
||||
output_detections = []
|
||||
for i in range(raw_box_tensor.shape[0]):
|
||||
boxes = detection_boxes[i, mask[i]]
|
||||
scores = detection_scores[i, mask[i]].unsqueeze(dim=-1)
|
||||
output_detections.append(torch.cat((boxes, scores), dim=-1).to('cpu'))
|
||||
|
||||
return output_detections
|
||||
|
||||
def _decode_boxes(self, raw_boxes, anchors):
|
||||
"""Converts the predictions into actual coordinates using
|
||||
the anchor boxes. Processes the entire batch at once.
|
||||
"""
|
||||
boxes = torch.zeros_like(raw_boxes)
|
||||
|
||||
x_center = raw_boxes[..., 0] / self.x_scale * \
|
||||
anchors[:, 2] + anchors[:, 0]
|
||||
y_center = raw_boxes[..., 1] / self.y_scale * \
|
||||
anchors[:, 3] + anchors[:, 1]
|
||||
|
||||
w = raw_boxes[..., 2] / self.w_scale * anchors[:, 2]
|
||||
h = raw_boxes[..., 3] / self.h_scale * anchors[:, 3]
|
||||
|
||||
boxes[..., 0] = y_center - h / 2. # ymin
|
||||
boxes[..., 1] = x_center - w / 2. # xmin
|
||||
boxes[..., 2] = y_center + h / 2. # ymax
|
||||
boxes[..., 3] = x_center + w / 2. # xmax
|
||||
|
||||
for k in range(6):
|
||||
offset = 4 + k * 2
|
||||
keypoint_x = raw_boxes[..., offset] / \
|
||||
self.x_scale * anchors[:, 2] + anchors[:, 0]
|
||||
keypoint_y = raw_boxes[..., offset + 1] / \
|
||||
self.y_scale * anchors[:, 3] + anchors[:, 1]
|
||||
boxes[..., offset] = keypoint_x
|
||||
boxes[..., offset + 1] = keypoint_y
|
||||
|
||||
return boxes
|
||||
|
||||
def _weighted_non_max_suppression(self, detections):
|
||||
"""The alternative NMS method as mentioned in the BlazeFace paper:
|
||||
|
||||
"We replace the suppression algorithm with a blending strategy that
|
||||
estimates the regression parameters of a bounding box as a weighted
|
||||
mean between the overlapping predictions."
|
||||
|
||||
The original MediaPipe code assigns the score of the most confident
|
||||
detection to the weighted detection, but we take the average score
|
||||
of the overlapping detections.
|
||||
|
||||
The input detections should be a Tensor of shape (count, 17).
|
||||
|
||||
Returns a list of PyTorch tensors, one for each detected face.
|
||||
|
||||
This is based on the source code from:
|
||||
mediapipe/calculators/util/non_max_suppression_calculator.cc
|
||||
mediapipe/calculators/util/non_max_suppression_calculator.proto
|
||||
"""
|
||||
if len(detections) == 0:
|
||||
return []
|
||||
|
||||
output_detections = []
|
||||
|
||||
# Sort the detections from highest to lowest score.
|
||||
remaining = torch.argsort(detections[:, 16], descending=True)
|
||||
|
||||
while len(remaining) > 0:
|
||||
detection = detections[remaining[0]]
|
||||
|
||||
# Compute the overlap between the first box and the other
|
||||
# remaining boxes. (Note that the other_boxes also include
|
||||
# the first_box.)
|
||||
first_box = detection[:4]
|
||||
other_boxes = detections[remaining, :4]
|
||||
ious = overlap_similarity(first_box, other_boxes)
|
||||
|
||||
# If two detections don't overlap enough, they are considered
|
||||
# to be from different faces.
|
||||
mask = ious > self.min_suppression_threshold
|
||||
overlapping = remaining[mask]
|
||||
remaining = remaining[~mask]
|
||||
|
||||
# Take an average of the coordinates from the overlapping
|
||||
# detections, weighted by their confidence scores.
|
||||
weighted_detection = detection.clone()
|
||||
if len(overlapping) > 1:
|
||||
coordinates = detections[overlapping, :16]
|
||||
scores = detections[overlapping, 16:17]
|
||||
total_score = scores.sum()
|
||||
weighted = (coordinates * scores).sum(dim=0) / total_score
|
||||
weighted_detection[:16] = weighted
|
||||
weighted_detection[16] = total_score / len(overlapping)
|
||||
|
||||
output_detections.append(weighted_detection)
|
||||
|
||||
return output_detections
|
||||
|
||||
|
||||
# IOU code from https://github.com/amdegroot/ssd.pytorch/blob/master/layers/box_utils.py
|
||||
|
||||
def intersect(box_a, box_b):
|
||||
""" We resize both tensors to [A,B,2] without new malloc:
|
||||
[A,2] -> [A,1,2] -> [A,B,2]
|
||||
[B,2] -> [1,B,2] -> [A,B,2]
|
||||
Then we compute the area of intersect between box_a and box_b.
|
||||
Args:
|
||||
box_a: (tensor) bounding boxes, Shape: [A,4].
|
||||
box_b: (tensor) bounding boxes, Shape: [B,4].
|
||||
Return:
|
||||
(tensor) intersection area, Shape: [A,B].
|
||||
"""
|
||||
A = box_a.size(0)
|
||||
B = box_b.size(0)
|
||||
max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2),
|
||||
box_b[:, 2:].unsqueeze(0).expand(A, B, 2))
|
||||
min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2),
|
||||
box_b[:, :2].unsqueeze(0).expand(A, B, 2))
|
||||
inter = torch.clamp((max_xy - min_xy), min=0)
|
||||
return inter[:, :, 0] * inter[:, :, 1]
|
||||
|
||||
|
||||
def jaccard(box_a, box_b):
|
||||
"""Compute the jaccard overlap of two sets of boxes. The jaccard overlap
|
||||
is simply the intersection over union of two boxes. Here we operate on
|
||||
ground truth boxes and default boxes.
|
||||
E.g.:
|
||||
A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B)
|
||||
Args:
|
||||
box_a: (tensor) Ground truth bounding boxes, Shape: [num_objects,4]
|
||||
box_b: (tensor) Prior boxes from priorbox layers, Shape: [num_priors,4]
|
||||
Return:
|
||||
jaccard overlap: (tensor) Shape: [box_a.size(0), box_b.size(0)]
|
||||
"""
|
||||
inter = intersect(box_a, box_b)
|
||||
area_a = ((box_a[:, 2] - box_a[:, 0])
|
||||
* (box_a[:, 3] - box_a[:, 1])).unsqueeze(1).expand_as(inter) # [A,B]
|
||||
area_b = ((box_b[:, 2] - box_b[:, 0])
|
||||
* (box_b[:, 3] - box_b[:, 1])).unsqueeze(0).expand_as(inter) # [A,B]
|
||||
union = area_a + area_b - inter
|
||||
return inter / union # [A,B]
|
||||
|
||||
|
||||
def overlap_similarity(box, other_boxes):
|
||||
"""Computes the IOU between a bounding box and set of other boxes."""
|
||||
return jaccard(box.unsqueeze(0), other_boxes).squeeze(0)
|
||||
@@ -0,0 +1,65 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
def image_resize(image, width=None, height=None, inter=cv2.INTER_AREA):
|
||||
# initialize the dimensions of the image to be resized and
|
||||
# grab the image size
|
||||
dim = None
|
||||
(h, w) = image.shape[:2]
|
||||
|
||||
# if both the width and height are None, then return the
|
||||
# original image
|
||||
if width is None and height is None:
|
||||
return image
|
||||
|
||||
# check to see if the width is None
|
||||
if width is None:
|
||||
# calculate the ratio of the height and construct the
|
||||
# dimensions
|
||||
r = height / float(h)
|
||||
dim = (int(w * r), height)
|
||||
|
||||
# otherwise, the height is None
|
||||
else:
|
||||
# calculate the ratio of the width and construct the
|
||||
# dimensions
|
||||
r = width / float(w)
|
||||
dim = (width, int(h * r))
|
||||
|
||||
# resize the image
|
||||
resized = cv2.resize(image, dim, interpolation=inter)
|
||||
|
||||
# return the resized image
|
||||
return resized
|
||||
|
||||
|
||||
def resize_and_crop_image(image, dim):
|
||||
if image.shape[0] > image.shape[1]:
|
||||
img = image_resize(image, width=dim)
|
||||
yshift, xshift = (image.shape[0] - image.shape[1]) // 2, 0
|
||||
y_start = (img.shape[0] - img.shape[1]) // 2
|
||||
y_end = y_start + dim
|
||||
return img[y_start:y_end, :, :], (xshift, yshift)
|
||||
else:
|
||||
img = image_resize(image, height=dim)
|
||||
yshift, xshift = 0, (image.shape[1] - image.shape[0]) // 2
|
||||
x_start = (img.shape[1] - img.shape[0]) // 2
|
||||
x_end = x_start + dim
|
||||
return img[:, x_start:x_end, :], (xshift, yshift)
|
||||
|
||||
|
||||
def resize_and_crop_batch(frames, dim):
|
||||
"""
|
||||
Center crop + resize to (dim x dim)
|
||||
inputs:
|
||||
- frames: list of images (numpy arrays)
|
||||
- dim: output dimension size
|
||||
"""
|
||||
smframes = []
|
||||
xshift, yshift = 0, 0
|
||||
for i in range(len(frames)):
|
||||
smframe, (xshift, yshift) = resize_and_crop_image(frames[i], dim)
|
||||
smframes.append(smframe)
|
||||
smframes = np.stack(smframes)
|
||||
return smframes, (xshift, yshift)
|
||||
@@ -0,0 +1,150 @@
|
||||
import logging
|
||||
import glob
|
||||
from tqdm import tqdm
|
||||
import numpy as np
|
||||
import torch
|
||||
from skimage import io
|
||||
|
||||
|
||||
class FaceDetector(object):
|
||||
"""An abstract class representing a face detector.
|
||||
|
||||
Any other face detection implementation must subclass it. All subclasses
|
||||
must implement ``detect_from_image``, that return a list of detected
|
||||
bounding boxes. Optionally, for speed considerations detect from path is
|
||||
recommended.
|
||||
"""
|
||||
|
||||
def __init__(self, device, verbose):
|
||||
self.device = device
|
||||
self.verbose = verbose
|
||||
|
||||
if verbose:
|
||||
if 'cpu' in device:
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.warning("Detection running on CPU, this may be potentially slow.")
|
||||
|
||||
if 'cpu' not in device and 'cuda' not in device and 'mps' not in device:
|
||||
if verbose:
|
||||
logger.error("Expected values for device are: {cpu, cuda, mps} but got: %s", device)
|
||||
raise ValueError
|
||||
|
||||
def detect_from_image(self, tensor_or_path):
|
||||
"""Detects faces in a given image.
|
||||
|
||||
This function detects the faces present in a provided BGR(usually)
|
||||
image. The input can be either the image itself or the path to it.
|
||||
|
||||
Arguments:
|
||||
tensor_or_path {numpy.ndarray, torch.tensor or string} -- the path
|
||||
to an image or the image itself.
|
||||
|
||||
Example::
|
||||
|
||||
>>> path_to_image = 'data/image_01.jpg'
|
||||
... detected_faces = detect_from_image(path_to_image)
|
||||
[A list of bounding boxes (x1, y1, x2, y2)]
|
||||
>>> image = cv2.imread(path_to_image)
|
||||
... detected_faces = detect_from_image(image)
|
||||
[A list of bounding boxes (x1, y1, x2, y2)]
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def detect_from_batch(self, tensor):
|
||||
"""Detects faces in a given image.
|
||||
|
||||
This function detects the faces present in a provided BGR(usually)
|
||||
image. The input can be either the image itself or the path to it.
|
||||
|
||||
Arguments:
|
||||
tensor {torch.tensor} -- image batch tensor.
|
||||
|
||||
Example::
|
||||
|
||||
>>> path_to_image = 'data/image_01.jpg'
|
||||
... detected_faces = detect_from_image(path_to_image)
|
||||
[A list of bounding boxes (x1, y1, x2, y2)]
|
||||
>>> image = cv2.imread(path_to_image)
|
||||
... detected_faces = detect_from_image(image)
|
||||
[A list of bounding boxes (x1, y1, x2, y2)]
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def detect_from_directory(self, path, extensions=['.jpg', '.png'], recursive=False, show_progress_bar=True):
|
||||
"""Detects faces from all the images present in a given directory.
|
||||
|
||||
Arguments:
|
||||
path {string} -- a string containing a path that points to the folder containing the images
|
||||
|
||||
Keyword Arguments:
|
||||
extensions {list} -- list of string containing the extensions to be
|
||||
consider in the following format: ``.extension_name`` (default:
|
||||
{['.jpg', '.png']}) recursive {bool} -- option wherever to scan the
|
||||
folder recursively (default: {False}) show_progress_bar {bool} --
|
||||
display a progressbar (default: {True})
|
||||
|
||||
Example:
|
||||
>>> directory = 'data'
|
||||
... detected_faces = detect_from_directory(directory)
|
||||
{A dictionary of [lists containing bounding boxes(x1, y1, x2, y2)]}
|
||||
|
||||
"""
|
||||
if self.verbose:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
if len(extensions) == 0:
|
||||
if self.verbose:
|
||||
logger.error("Expected at list one extension, but none was received.")
|
||||
raise ValueError
|
||||
|
||||
if self.verbose:
|
||||
logger.info("Constructing the list of images.")
|
||||
additional_pattern = '/**/*' if recursive else '/*'
|
||||
files = []
|
||||
for extension in extensions:
|
||||
files.extend(glob.glob(path + additional_pattern + extension, recursive=recursive))
|
||||
|
||||
if self.verbose:
|
||||
logger.info("Finished searching for images. %s images found", len(files))
|
||||
logger.info("Preparing to run the detection.")
|
||||
|
||||
predictions = {}
|
||||
for image_path in tqdm(files, disable=not show_progress_bar):
|
||||
if self.verbose:
|
||||
logger.info("Running the face detector on image: %s", image_path)
|
||||
predictions[image_path] = self.detect_from_image(image_path)
|
||||
|
||||
if self.verbose:
|
||||
logger.info("The detector was successfully run on all %s images", len(files))
|
||||
|
||||
return predictions
|
||||
|
||||
@property
|
||||
def reference_scale(self):
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def reference_x_shift(self):
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def reference_y_shift(self):
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def tensor_or_path_to_ndarray(tensor_or_path):
|
||||
"""Convert path (represented as a string) or torch.tensor to a numpy.ndarray
|
||||
|
||||
Arguments:
|
||||
tensor_or_path {numpy.ndarray, torch.tensor or string} -- path to the image, or the image itself
|
||||
"""
|
||||
if isinstance(tensor_or_path, str):
|
||||
return io.imread(tensor_or_path)
|
||||
elif torch.is_tensor(tensor_or_path):
|
||||
return tensor_or_path.cpu().numpy()
|
||||
elif isinstance(tensor_or_path, np.ndarray):
|
||||
return tensor_or_path
|
||||
else:
|
||||
raise TypeError
|
||||
@@ -0,0 +1 @@
|
||||
from .dlib_detector import DlibDetector as FaceDetector
|
||||
@@ -0,0 +1,48 @@
|
||||
import warnings
|
||||
import cv2
|
||||
import dlib
|
||||
|
||||
from ..core import FaceDetector
|
||||
from ...utils import load_file_from_url
|
||||
|
||||
|
||||
class DlibDetector(FaceDetector):
|
||||
def __init__(self, device, path_to_detector=None, verbose=False):
|
||||
super().__init__(device, verbose)
|
||||
|
||||
warnings.warn('Warning: this detector is deprecated. Please use a different one, i.e.: S3FD.')
|
||||
|
||||
# Initialise the face detector
|
||||
if 'cuda' in device:
|
||||
if path_to_detector is None:
|
||||
path_to_detector = load_file_from_url(
|
||||
"https://www.adrianbulat.com/downloads/dlib/mmod_human_face_detector.dat")
|
||||
|
||||
self.face_detector = dlib.cnn_face_detection_model_v1(path_to_detector)
|
||||
else:
|
||||
self.face_detector = dlib.get_frontal_face_detector()
|
||||
|
||||
def detect_from_image(self, tensor_or_path):
|
||||
image = self.tensor_or_path_to_ndarray(tensor_or_path)
|
||||
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
|
||||
|
||||
detected_faces = self.face_detector(cv2.cvtColor(image, cv2.COLOR_BGR2GRAY))
|
||||
|
||||
if 'cuda' not in self.device:
|
||||
detected_faces = [[d.left(), d.top(), d.right(), d.bottom()] for d in detected_faces]
|
||||
else:
|
||||
detected_faces = [[d.rect.left(), d.rect.top(), d.rect.right(), d.rect.bottom()] for d in detected_faces]
|
||||
|
||||
return detected_faces
|
||||
|
||||
@property
|
||||
def reference_scale(self):
|
||||
return 195
|
||||
|
||||
@property
|
||||
def reference_x_shift(self):
|
||||
return 0
|
||||
|
||||
@property
|
||||
def reference_y_shift(self):
|
||||
return 0
|
||||
@@ -0,0 +1 @@
|
||||
from .folder_detector import FolderDetector as FaceDetector
|
||||
@@ -0,0 +1,53 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from ..core import FaceDetector
|
||||
|
||||
|
||||
class FolderDetector(FaceDetector):
|
||||
'''This is a simple helper module that assumes the faces were detected already
|
||||
(either previously or are provided as ground truth).
|
||||
|
||||
The class expects to find the bounding boxes in the same format used by
|
||||
the rest of face detectors, mainly ``list[(x1,y1,x2,y2),...]``.
|
||||
For each image the detector will search for a file with the same name and with one of the
|
||||
following extensions: .npy, .t7 or .pth
|
||||
|
||||
'''
|
||||
|
||||
def __init__(self, device, path_to_detector=None, verbose=False):
|
||||
super(FolderDetector, self).__init__(device, verbose)
|
||||
|
||||
def detect_from_image(self, tensor_or_path):
|
||||
# Only strings supported
|
||||
if not isinstance(tensor_or_path, str):
|
||||
raise ValueError
|
||||
|
||||
base_name = os.path.splitext(tensor_or_path)[0]
|
||||
|
||||
if os.path.isfile(base_name + '.npy'):
|
||||
detected_faces = np.load(base_name + '.npy')
|
||||
elif os.path.isfile(base_name + '.t7'):
|
||||
detected_faces = torch.load(base_name + '.t7')
|
||||
elif os.path.isfile(base_name + '.pth'):
|
||||
detected_faces = torch.load(base_name + '.pth')
|
||||
else:
|
||||
raise FileNotFoundError
|
||||
|
||||
if not isinstance(detected_faces, list):
|
||||
raise TypeError
|
||||
|
||||
return detected_faces
|
||||
|
||||
@property
|
||||
def reference_scale(self):
|
||||
return 195
|
||||
|
||||
@property
|
||||
def reference_x_shift(self):
|
||||
return 0
|
||||
|
||||
@property
|
||||
def reference_y_shift(self):
|
||||
return 0
|
||||
@@ -0,0 +1 @@
|
||||
from .sfd_detector import SFDDetector as FaceDetector
|
||||
@@ -0,0 +1,71 @@
|
||||
import math
|
||||
import numpy as np
|
||||
|
||||
|
||||
def nms(dets, thresh):
|
||||
if 0 == len(dets):
|
||||
return []
|
||||
x1, y1, x2, y2, scores = dets[:, 0], dets[:, 1], dets[:, 2], dets[:, 3], dets[:, 4]
|
||||
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
|
||||
order = scores.argsort()[::-1]
|
||||
|
||||
keep = []
|
||||
while order.size > 0:
|
||||
i = order[0]
|
||||
keep.append(i)
|
||||
xx1, yy1 = np.maximum(x1[i], x1[order[1:]]), np.maximum(y1[i], y1[order[1:]])
|
||||
xx2, yy2 = np.minimum(x2[i], x2[order[1:]]), np.minimum(y2[i], y2[order[1:]])
|
||||
|
||||
w, h = np.maximum(0.0, xx2 - xx1 + 1), np.maximum(0.0, yy2 - yy1 + 1)
|
||||
ovr = w * h / (areas[i] + areas[order[1:]] - w * h)
|
||||
|
||||
inds = np.where(ovr <= thresh)[0]
|
||||
order = order[inds + 1]
|
||||
|
||||
return keep
|
||||
|
||||
|
||||
def encode(matched, priors, variances):
|
||||
"""Encode the variances from the priorbox layers into the ground truth boxes
|
||||
we have matched (based on jaccard overlap) with the prior boxes.
|
||||
Args:
|
||||
matched: (tensor) Coords of ground truth for each prior in point-form
|
||||
Shape: [num_priors, 4].
|
||||
priors: (tensor) Prior boxes in center-offset form
|
||||
Shape: [num_priors,4].
|
||||
variances: (list[float]) Variances of priorboxes
|
||||
Return:
|
||||
encoded boxes (tensor), Shape: [num_priors, 4]
|
||||
"""
|
||||
|
||||
# dist b/t match center and prior's center
|
||||
g_cxcy = (matched[:, :2] + matched[:, 2:]) / 2 - priors[:, :2]
|
||||
# encode variance
|
||||
g_cxcy /= (variances[0] * priors[:, 2:])
|
||||
# match wh / prior wh
|
||||
g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:]
|
||||
g_wh = np.log(g_wh) / variances[1]
|
||||
|
||||
# return target for smooth_l1_loss
|
||||
return np.concatenate([g_cxcy, g_wh], 1) # [num_priors,4]
|
||||
|
||||
|
||||
def decode(loc, priors, variances):
|
||||
"""Decode locations from predictions using priors to undo
|
||||
the encoding we did for offset regression at train time.
|
||||
Args:
|
||||
loc (tensor): location predictions for loc layers,
|
||||
Shape: [num_priors,4]
|
||||
priors (tensor): Prior boxes in center-offset form.
|
||||
Shape: [num_priors,4].
|
||||
variances: (list[float]) Variances of priorboxes
|
||||
Return:
|
||||
decoded bounding box predictions
|
||||
"""
|
||||
|
||||
boxes = np.concatenate((
|
||||
priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
|
||||
priors[:, 2:] * np.exp(loc[:, 2:] * variances[1])), 1)
|
||||
boxes[:, :2] -= boxes[:, 2:] / 2
|
||||
boxes[:, 2:] += boxes[:, :2]
|
||||
return boxes
|
||||
@@ -0,0 +1,85 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from .bbox import *
|
||||
|
||||
|
||||
def detect(net, img, device):
|
||||
img = img.transpose(2, 0, 1)
|
||||
# Creates a batch of 1
|
||||
img = np.expand_dims(img, 0)
|
||||
|
||||
img = torch.from_numpy(img.copy()).to(device, dtype=torch.float32)
|
||||
|
||||
return batch_detect(net, img, device)
|
||||
|
||||
|
||||
def batch_detect(net, img_batch, device):
|
||||
"""
|
||||
Inputs:
|
||||
- img_batch: a torch.Tensor of shape (Batch size, Channels, Height, Width)
|
||||
"""
|
||||
|
||||
if 'cuda' in device:
|
||||
torch.backends.cudnn.benchmark = True
|
||||
|
||||
batch_size = img_batch.size(0)
|
||||
img_batch = img_batch.to(device, dtype=torch.float32)
|
||||
|
||||
img_batch = img_batch.flip(-3) # RGB to BGR
|
||||
img_batch = img_batch - torch.tensor([104.0, 117.0, 123.0], device=device).view(1, 3, 1, 1)
|
||||
|
||||
with torch.no_grad():
|
||||
olist = net(img_batch) # patched uint8_t overflow error
|
||||
|
||||
for i in range(len(olist) // 2):
|
||||
olist[i * 2] = F.softmax(olist[i * 2], dim=1)
|
||||
|
||||
olist = [oelem.data.cpu().numpy() for oelem in olist]
|
||||
|
||||
bboxlists = get_predictions(olist, batch_size)
|
||||
return bboxlists
|
||||
|
||||
|
||||
def get_predictions(olist, batch_size):
|
||||
bboxlists = []
|
||||
variances = [0.1, 0.2]
|
||||
for i in range(len(olist) // 2):
|
||||
ocls, oreg = olist[i * 2], olist[i * 2 + 1]
|
||||
stride = 2**(i + 2) # 4,8,16,32,64,128
|
||||
poss = zip(*np.where(ocls[:, 1, :, :] > 0.05))
|
||||
for Iindex, hindex, windex in poss:
|
||||
axc, ayc = stride / 2 + windex * stride, stride / 2 + hindex * stride
|
||||
priors = np.array([[axc / 1.0, ayc / 1.0, stride * 4 / 1.0, stride * 4 / 1.0]])
|
||||
score = ocls[:, 1, hindex, windex][:,None]
|
||||
loc = oreg[:, :, hindex, windex].copy()
|
||||
boxes = decode(loc, priors, variances)
|
||||
bboxlists.append(np.concatenate((boxes, score), axis=1))
|
||||
|
||||
if len(bboxlists) == 0: # No candidates within given threshold
|
||||
bboxlists = np.array([[] for _ in range(batch_size)])
|
||||
else:
|
||||
bboxlists = np.stack(bboxlists, axis=1)
|
||||
return bboxlists
|
||||
|
||||
|
||||
def flip_detect(net, img, device):
|
||||
img = cv2.flip(img, 1)
|
||||
b = detect(net, img, device)
|
||||
|
||||
bboxlist = np.zeros(b.shape)
|
||||
bboxlist[:, 0] = img.shape[1] - b[:, 2]
|
||||
bboxlist[:, 1] = b[:, 1]
|
||||
bboxlist[:, 2] = img.shape[1] - b[:, 0]
|
||||
bboxlist[:, 3] = b[:, 3]
|
||||
bboxlist[:, 4] = b[:, 4]
|
||||
return bboxlist
|
||||
|
||||
|
||||
def pts_to_bb(pts):
|
||||
min_x, min_y = np.min(pts, axis=0)
|
||||
max_x, max_y = np.max(pts, axis=0)
|
||||
return np.array([min_x, min_y, max_x, max_y])
|
||||
@@ -0,0 +1,127 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class L2Norm(nn.Module):
|
||||
def __init__(self, n_channels, scale=1.0):
|
||||
super(L2Norm, self).__init__()
|
||||
self.n_channels = n_channels
|
||||
self.scale = scale
|
||||
self.eps = 1e-10
|
||||
self.weight = nn.Parameter(torch.empty(self.n_channels).fill_(self.scale))
|
||||
|
||||
def forward(self, x):
|
||||
norm = x.pow(2).sum(dim=1, keepdim=True).sqrt() + self.eps
|
||||
x = x / norm * self.weight.view(1, -1, 1, 1)
|
||||
return x
|
||||
|
||||
|
||||
class s3fd(nn.Module):
|
||||
def __init__(self):
|
||||
super(s3fd, self).__init__()
|
||||
self.conv1_1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
|
||||
self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
|
||||
self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)
|
||||
self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
|
||||
self.conv3_3 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1)
|
||||
self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
|
||||
self.conv4_3 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
self.conv5_1 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
|
||||
self.conv5_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
|
||||
self.conv5_3 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
self.fc6 = nn.Conv2d(512, 1024, kernel_size=3, stride=1, padding=3)
|
||||
self.fc7 = nn.Conv2d(1024, 1024, kernel_size=1, stride=1, padding=0)
|
||||
|
||||
self.conv6_1 = nn.Conv2d(1024, 256, kernel_size=1, stride=1, padding=0)
|
||||
self.conv6_2 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1)
|
||||
|
||||
self.conv7_1 = nn.Conv2d(512, 128, kernel_size=1, stride=1, padding=0)
|
||||
self.conv7_2 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
|
||||
|
||||
self.conv3_3_norm = L2Norm(256, scale=10)
|
||||
self.conv4_3_norm = L2Norm(512, scale=8)
|
||||
self.conv5_3_norm = L2Norm(512, scale=5)
|
||||
|
||||
self.conv3_3_norm_mbox_conf = nn.Conv2d(256, 4, kernel_size=3, stride=1, padding=1)
|
||||
self.conv3_3_norm_mbox_loc = nn.Conv2d(256, 4, kernel_size=3, stride=1, padding=1)
|
||||
self.conv4_3_norm_mbox_conf = nn.Conv2d(512, 2, kernel_size=3, stride=1, padding=1)
|
||||
self.conv4_3_norm_mbox_loc = nn.Conv2d(512, 4, kernel_size=3, stride=1, padding=1)
|
||||
self.conv5_3_norm_mbox_conf = nn.Conv2d(512, 2, kernel_size=3, stride=1, padding=1)
|
||||
self.conv5_3_norm_mbox_loc = nn.Conv2d(512, 4, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
self.fc7_mbox_conf = nn.Conv2d(1024, 2, kernel_size=3, stride=1, padding=1)
|
||||
self.fc7_mbox_loc = nn.Conv2d(1024, 4, kernel_size=3, stride=1, padding=1)
|
||||
self.conv6_2_mbox_conf = nn.Conv2d(512, 2, kernel_size=3, stride=1, padding=1)
|
||||
self.conv6_2_mbox_loc = nn.Conv2d(512, 4, kernel_size=3, stride=1, padding=1)
|
||||
self.conv7_2_mbox_conf = nn.Conv2d(256, 2, kernel_size=3, stride=1, padding=1)
|
||||
self.conv7_2_mbox_loc = nn.Conv2d(256, 4, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, x):
|
||||
h = F.relu(self.conv1_1(x), inplace=True)
|
||||
h = F.relu(self.conv1_2(h), inplace=True)
|
||||
h = F.max_pool2d(h, 2, 2)
|
||||
|
||||
h = F.relu(self.conv2_1(h), inplace=True)
|
||||
h = F.relu(self.conv2_2(h), inplace=True)
|
||||
h = F.max_pool2d(h, 2, 2)
|
||||
|
||||
h = F.relu(self.conv3_1(h), inplace=True)
|
||||
h = F.relu(self.conv3_2(h), inplace=True)
|
||||
h = F.relu(self.conv3_3(h), inplace=True)
|
||||
f3_3 = h
|
||||
h = F.max_pool2d(h, 2, 2)
|
||||
|
||||
h = F.relu(self.conv4_1(h), inplace=True)
|
||||
h = F.relu(self.conv4_2(h), inplace=True)
|
||||
h = F.relu(self.conv4_3(h), inplace=True)
|
||||
f4_3 = h
|
||||
h = F.max_pool2d(h, 2, 2)
|
||||
|
||||
h = F.relu(self.conv5_1(h), inplace=True)
|
||||
h = F.relu(self.conv5_2(h), inplace=True)
|
||||
h = F.relu(self.conv5_3(h), inplace=True)
|
||||
f5_3 = h
|
||||
h = F.max_pool2d(h, 2, 2)
|
||||
|
||||
h = F.relu(self.fc6(h), inplace=True)
|
||||
h = F.relu(self.fc7(h), inplace=True)
|
||||
ffc7 = h
|
||||
h = F.relu(self.conv6_1(h), inplace=True)
|
||||
h = F.relu(self.conv6_2(h), inplace=True)
|
||||
f6_2 = h
|
||||
h = F.relu(self.conv7_1(h), inplace=True)
|
||||
h = F.relu(self.conv7_2(h), inplace=True)
|
||||
f7_2 = h
|
||||
|
||||
f3_3 = self.conv3_3_norm(f3_3)
|
||||
f4_3 = self.conv4_3_norm(f4_3)
|
||||
f5_3 = self.conv5_3_norm(f5_3)
|
||||
|
||||
cls1 = self.conv3_3_norm_mbox_conf(f3_3)
|
||||
reg1 = self.conv3_3_norm_mbox_loc(f3_3)
|
||||
cls2 = self.conv4_3_norm_mbox_conf(f4_3)
|
||||
reg2 = self.conv4_3_norm_mbox_loc(f4_3)
|
||||
cls3 = self.conv5_3_norm_mbox_conf(f5_3)
|
||||
reg3 = self.conv5_3_norm_mbox_loc(f5_3)
|
||||
cls4 = self.fc7_mbox_conf(ffc7)
|
||||
reg4 = self.fc7_mbox_loc(ffc7)
|
||||
cls5 = self.conv6_2_mbox_conf(f6_2)
|
||||
reg5 = self.conv6_2_mbox_loc(f6_2)
|
||||
cls6 = self.conv7_2_mbox_conf(f7_2)
|
||||
reg6 = self.conv7_2_mbox_loc(f7_2)
|
||||
|
||||
# max-out background label
|
||||
chunk = torch.chunk(cls1, 4, 1)
|
||||
bmax = torch.max(torch.max(chunk[0], chunk[1]), chunk[2])
|
||||
cls1 = torch.cat([bmax, chunk[3]], dim=1)
|
||||
|
||||
return [cls1, reg1, cls2, reg2, cls3, reg3, cls4, reg4, cls5, reg5, cls6, reg6]
|
||||
@@ -0,0 +1,71 @@
|
||||
import torch
|
||||
from torch.utils.model_zoo import load_url
|
||||
|
||||
from ..core import FaceDetector
|
||||
|
||||
from .net_s3fd import s3fd
|
||||
from .bbox import nms
|
||||
from .detect import detect, batch_detect
|
||||
|
||||
models_urls = {
|
||||
's3fd': 'https://www.adrianbulat.com/downloads/python-fan/s3fd-619a316812.pth',
|
||||
}
|
||||
|
||||
|
||||
class SFDDetector(FaceDetector):
|
||||
'''SF3D Detector.
|
||||
'''
|
||||
|
||||
def __init__(self, device, path_to_detector=None, verbose=False, filter_threshold=0.5):
|
||||
super(SFDDetector, self).__init__(device, verbose)
|
||||
|
||||
# Initialise the face detector
|
||||
if path_to_detector is None:
|
||||
model_weights = load_url(models_urls['s3fd'])
|
||||
else:
|
||||
model_weights = torch.load(path_to_detector)
|
||||
|
||||
self.fiter_threshold = filter_threshold
|
||||
self.face_detector = s3fd()
|
||||
self.face_detector.load_state_dict(model_weights)
|
||||
self.face_detector.to(device)
|
||||
self.face_detector.eval()
|
||||
|
||||
def _filter_bboxes(self, bboxlist):
|
||||
if len(bboxlist) > 0:
|
||||
keep = nms(bboxlist, 0.3)
|
||||
bboxlist = bboxlist[keep, :]
|
||||
bboxlist = [x for x in bboxlist if x[-1] > self.fiter_threshold]
|
||||
|
||||
return bboxlist
|
||||
|
||||
def detect_from_image(self, tensor_or_path):
|
||||
image = self.tensor_or_path_to_ndarray(tensor_or_path)
|
||||
|
||||
bboxlist = detect(self.face_detector, image, device=self.device)[0]
|
||||
bboxlist = self._filter_bboxes(bboxlist)
|
||||
|
||||
return bboxlist
|
||||
|
||||
def detect_from_batch(self, tensor):
|
||||
bboxlists = batch_detect(self.face_detector, tensor, device=self.device)
|
||||
|
||||
new_bboxlists = []
|
||||
for i in range(bboxlists.shape[0]):
|
||||
bboxlist = bboxlists[i]
|
||||
bboxlist = self._filter_bboxes(bboxlist)
|
||||
new_bboxlists.append(bboxlist)
|
||||
|
||||
return new_bboxlists
|
||||
|
||||
@property
|
||||
def reference_scale(self):
|
||||
return 195
|
||||
|
||||
@property
|
||||
def reference_x_shift(self):
|
||||
return 0
|
||||
|
||||
@property
|
||||
def reference_y_shift(self):
|
||||
return 0
|
||||
@@ -0,0 +1,38 @@
|
||||
import logging
|
||||
import glob
|
||||
|
||||
import torch
|
||||
|
||||
class FolderData(torch.utils.data.Dataset):
|
||||
def __init__(self, path, transforms, extensions=['.jpg', '.png'], recursive=False, verbose=False):
|
||||
self.verbose = verbose
|
||||
if self.verbose:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
if len(extensions) == 0:
|
||||
if self.verbose:
|
||||
logger.error("Expected at list one extension, but none was received.")
|
||||
raise ValueError
|
||||
|
||||
if self.verbose:
|
||||
logger.info("Constructing the list of images.")
|
||||
additional_pattern = '/**/*' if recursive else '/*'
|
||||
files = []
|
||||
for extension in extensions:
|
||||
files.extend(glob.glob(path + additional_pattern + extension, recursive=recursive))
|
||||
|
||||
if self.verbose:
|
||||
logger.info("Finished searching for images. %s images found", len(files))
|
||||
logger.info("Preparing to run the detection.")
|
||||
|
||||
self.files = files
|
||||
self.transforms = transforms
|
||||
|
||||
def __getitem__(self, idx):
|
||||
image_path = self.files[idx]
|
||||
image = self.transforms(image_path)
|
||||
|
||||
return image_path, image
|
||||
|
||||
def __len__(self):
|
||||
return len(self.files)
|
||||
@@ -0,0 +1,379 @@
|
||||
import os
|
||||
import sys
|
||||
import errno
|
||||
import torch
|
||||
import math
|
||||
import numpy as np
|
||||
import cv2
|
||||
from skimage import io
|
||||
from skimage import color
|
||||
from numba import jit
|
||||
|
||||
from urllib.parse import urlparse
|
||||
from torch.hub import download_url_to_file, HASH_REGEX
|
||||
try:
|
||||
from torch.hub import get_dir
|
||||
except BaseException:
|
||||
from torch.hub import _get_torch_home as get_dir
|
||||
|
||||
gauss_kernel = None
|
||||
|
||||
|
||||
def _gaussian(
|
||||
size=3, sigma=0.25, amplitude=1, normalize=False, width=None,
|
||||
height=None, sigma_horz=None, sigma_vert=None, mean_horz=0.5,
|
||||
mean_vert=0.5):
|
||||
# handle some defaults
|
||||
if width is None:
|
||||
width = size
|
||||
if height is None:
|
||||
height = size
|
||||
if sigma_horz is None:
|
||||
sigma_horz = sigma
|
||||
if sigma_vert is None:
|
||||
sigma_vert = sigma
|
||||
center_x = mean_horz * width + 0.5
|
||||
center_y = mean_vert * height + 0.5
|
||||
gauss = np.empty((height, width), dtype=np.float32)
|
||||
# generate kernel
|
||||
for i in range(height):
|
||||
for j in range(width):
|
||||
gauss[i][j] = amplitude * math.exp(-(math.pow((j + 1 - center_x) / (
|
||||
sigma_horz * width), 2) / 2.0 + math.pow((i + 1 - center_y) / (sigma_vert * height), 2) / 2.0))
|
||||
if normalize:
|
||||
gauss = gauss / np.sum(gauss)
|
||||
return gauss
|
||||
|
||||
|
||||
def draw_gaussian(image, point, sigma):
|
||||
global gauss_kernel
|
||||
# Check if the gaussian is inside
|
||||
ul = [math.floor(point[0] - 3 * sigma), math.floor(point[1] - 3 * sigma)]
|
||||
br = [math.floor(point[0] + 3 * sigma), math.floor(point[1] + 3 * sigma)]
|
||||
if (ul[0] > image.shape[1] or ul[1] > image.shape[0] or br[0] < 1 or br[1] < 1):
|
||||
return image
|
||||
size = 6 * sigma + 1
|
||||
if gauss_kernel is None:
|
||||
g = _gaussian(size)
|
||||
gauss_kernel = g
|
||||
else:
|
||||
g = gauss_kernel
|
||||
g_x = [int(max(1, -ul[0])), int(min(br[0], image.shape[1])) - int(max(1, ul[0])) + int(max(1, -ul[0]))]
|
||||
g_y = [int(max(1, -ul[1])), int(min(br[1], image.shape[0])) - int(max(1, ul[1])) + int(max(1, -ul[1]))]
|
||||
img_x = [int(max(1, ul[0])), int(min(br[0], image.shape[1]))]
|
||||
img_y = [int(max(1, ul[1])), int(min(br[1], image.shape[0]))]
|
||||
assert (g_x[0] > 0 and g_y[1] > 0)
|
||||
image[img_y[0] - 1:img_y[1], img_x[0] - 1:img_x[1]
|
||||
] = image[img_y[0] - 1:img_y[1], img_x[0] - 1:img_x[1]] + g[g_y[0] - 1:g_y[1], g_x[0] - 1:g_x[1]]
|
||||
image[image > 1] = 1
|
||||
return image
|
||||
|
||||
|
||||
def transform(point, center, scale, resolution, invert=False):
|
||||
"""Generate and affine transformation matrix.
|
||||
|
||||
Given a set of points, a center, a scale and a targer resolution, the
|
||||
function generates and affine transformation matrix. If invert is ``True``
|
||||
it will produce the inverse transformation.
|
||||
|
||||
Arguments:
|
||||
point {torch.tensor} -- the input 2D point
|
||||
center {torch.tensor or numpy.array} -- the center around which to perform the transformations
|
||||
scale {float} -- the scale of the face/object
|
||||
resolution {float} -- the output resolution
|
||||
|
||||
Keyword Arguments:
|
||||
invert {bool} -- define wherever the function should produce the direct or the
|
||||
inverse transformation matrix (default: {False})
|
||||
"""
|
||||
_pt = torch.ones(3)
|
||||
_pt[0] = point[0]
|
||||
_pt[1] = point[1]
|
||||
|
||||
h = 200.0 * scale
|
||||
t = torch.eye(3)
|
||||
t[0, 0] = resolution / h
|
||||
t[1, 1] = resolution / h
|
||||
t[0, 2] = resolution * (-center[0] / h + 0.5)
|
||||
t[1, 2] = resolution * (-center[1] / h + 0.5)
|
||||
|
||||
if invert:
|
||||
t = torch.inverse(t)
|
||||
|
||||
new_point = (torch.matmul(t, _pt))[0:2]
|
||||
|
||||
return new_point.int()
|
||||
|
||||
|
||||
def crop(image, center, scale, resolution=256.0):
|
||||
"""Center crops an image or set of heatmaps
|
||||
|
||||
Arguments:
|
||||
image {numpy.array} -- an rgb image
|
||||
center {numpy.array} -- the center of the object, usually the same as of the bounding box
|
||||
scale {float} -- scale of the face
|
||||
|
||||
Keyword Arguments:
|
||||
resolution {float} -- the size of the output cropped image (default: {256.0})
|
||||
|
||||
Returns:
|
||||
[type] -- [description]
|
||||
""" # Crop around the center point
|
||||
""" Crops the image around the center. Input is expected to be an np.ndarray """
|
||||
ul = transform([1, 1], center, scale, resolution, True)
|
||||
br = transform([resolution, resolution], center, scale, resolution, True)
|
||||
# pad = math.ceil(torch.norm((ul - br).float()) / 2.0 - (br[0] - ul[0]) / 2.0)
|
||||
if image.ndim > 2:
|
||||
newDim = np.array([br[1] - ul[1], br[0] - ul[0],
|
||||
image.shape[2]], dtype=np.int32)
|
||||
newImg = np.zeros(newDim, dtype=np.uint8)
|
||||
else:
|
||||
newDim = np.array([br[1] - ul[1], br[0] - ul[0]], dtype=np.int)
|
||||
newImg = np.zeros(newDim, dtype=np.uint8)
|
||||
ht = image.shape[0]
|
||||
wd = image.shape[1]
|
||||
newX = np.array(
|
||||
[max(1, -ul[0] + 1), min(br[0], wd) - ul[0]], dtype=np.int32)
|
||||
newY = np.array(
|
||||
[max(1, -ul[1] + 1), min(br[1], ht) - ul[1]], dtype=np.int32)
|
||||
oldX = np.array([max(1, ul[0] + 1), min(br[0], wd)], dtype=np.int32)
|
||||
oldY = np.array([max(1, ul[1] + 1), min(br[1], ht)], dtype=np.int32)
|
||||
newImg[newY[0] - 1:newY[1], newX[0] - 1:newX[1]
|
||||
] = image[oldY[0] - 1:oldY[1], oldX[0] - 1:oldX[1], :]
|
||||
newImg = cv2.resize(newImg, dsize=(int(resolution), int(resolution)),
|
||||
interpolation=cv2.INTER_LINEAR)
|
||||
return newImg
|
||||
|
||||
|
||||
@jit(nopython=True)
|
||||
def transform_np(point, center, scale, resolution, invert=False):
|
||||
"""Generate and affine transformation matrix.
|
||||
|
||||
Given a set of points, a center, a scale and a targer resolution, the
|
||||
function generates and affine transformation matrix. If invert is ``True``
|
||||
it will produce the inverse transformation.
|
||||
|
||||
Arguments:
|
||||
point {numpy.array} -- the input 2D point
|
||||
center {numpy.array} -- the center around which to perform the transformations
|
||||
scale {float} -- the scale of the face/object
|
||||
resolution {float} -- the output resolution
|
||||
|
||||
Keyword Arguments:
|
||||
invert {bool} -- define wherever the function should produce the direct or the
|
||||
inverse transformation matrix (default: {False})
|
||||
"""
|
||||
_pt = np.ones(3)
|
||||
_pt[0] = point[0]
|
||||
_pt[1] = point[1]
|
||||
|
||||
h = 200.0 * scale
|
||||
t = np.eye(3)
|
||||
t[0, 0] = resolution / h
|
||||
t[1, 1] = resolution / h
|
||||
t[0, 2] = resolution * (-center[0] / h + 0.5)
|
||||
t[1, 2] = resolution * (-center[1] / h + 0.5)
|
||||
|
||||
if invert:
|
||||
t = np.ascontiguousarray(np.linalg.pinv(t))
|
||||
|
||||
new_point = np.dot(t, _pt)[0:2]
|
||||
|
||||
return new_point.astype(np.int32)
|
||||
|
||||
|
||||
def get_preds_fromhm(hm, center=None, scale=None):
|
||||
"""Obtain (x,y) coordinates given a set of N heatmaps. If the center
|
||||
and the scale is provided the function will return the points also in
|
||||
the original coordinate frame.
|
||||
|
||||
Arguments:
|
||||
hm {torch.tensor} -- the predicted heatmaps, of shape [B, N, W, H]
|
||||
|
||||
Keyword Arguments:
|
||||
center {torch.tensor} -- the center of the bounding box (default: {None})
|
||||
scale {float} -- face scale (default: {None})
|
||||
"""
|
||||
B, C, H, W = hm.shape
|
||||
hm_reshape = hm.reshape(B, C, H * W)
|
||||
idx = np.argmax(hm_reshape, axis=-1)
|
||||
scores = np.take_along_axis(hm_reshape, np.expand_dims(idx, axis=-1), axis=-1).squeeze(-1)
|
||||
preds, preds_orig = _get_preds_fromhm(hm, idx, center, scale)
|
||||
|
||||
return preds, preds_orig, scores
|
||||
|
||||
|
||||
@jit(nopython=True)
|
||||
def _get_preds_fromhm(hm, idx, center=None, scale=None):
|
||||
"""Obtain (x,y) coordinates given a set of N heatmaps and the
|
||||
coresponding locations of the maximums. If the center
|
||||
and the scale is provided the function will return the points also in
|
||||
the original coordinate frame.
|
||||
|
||||
Arguments:
|
||||
hm {torch.tensor} -- the predicted heatmaps, of shape [B, N, W, H]
|
||||
|
||||
Keyword Arguments:
|
||||
center {torch.tensor} -- the center of the bounding box (default: {None})
|
||||
scale {float} -- face scale (default: {None})
|
||||
"""
|
||||
B, C, H, W = hm.shape
|
||||
idx += 1
|
||||
preds = idx.repeat(2).reshape(B, C, 2).astype(np.float32)
|
||||
preds[:, :, 0] = (preds[:, :, 0] - 1) % W + 1
|
||||
preds[:, :, 1] = np.floor((preds[:, :, 1] - 1) / H) + 1
|
||||
|
||||
for i in range(B):
|
||||
for j in range(C):
|
||||
hm_ = hm[i, j, :]
|
||||
pX, pY = int(preds[i, j, 0]) - 1, int(preds[i, j, 1]) - 1
|
||||
if pX > 0 and pX < 63 and pY > 0 and pY < 63:
|
||||
diff = np.array(
|
||||
[hm_[pY, pX + 1] - hm_[pY, pX - 1],
|
||||
hm_[pY + 1, pX] - hm_[pY - 1, pX]])
|
||||
preds[i, j] += np.sign(diff) * 0.25
|
||||
|
||||
preds -= 0.5
|
||||
|
||||
preds_orig = np.zeros_like(preds)
|
||||
if center is not None and scale is not None:
|
||||
for i in range(B):
|
||||
for j in range(C):
|
||||
preds_orig[i, j] = transform_np(
|
||||
preds[i, j], center, scale, H, True)
|
||||
|
||||
return preds, preds_orig
|
||||
|
||||
|
||||
def create_target_heatmap(target_landmarks, centers, scales):
|
||||
heatmaps = np.zeros((target_landmarks.shape[0], 68, 64, 64), dtype=np.float32)
|
||||
for i in range(heatmaps.shape[0]):
|
||||
for p in range(68):
|
||||
landmark_cropped_coor = transform(target_landmarks[i, p] + 1, centers[i], scales[i], 64, invert=False)
|
||||
heatmaps[i, p] = draw_gaussian(heatmaps[i, p], landmark_cropped_coor + 1, 2)
|
||||
return torch.tensor(heatmaps)
|
||||
|
||||
|
||||
def create_bounding_box(target_landmarks, expansion_factor=0.0):
|
||||
"""
|
||||
gets a batch of landmarks and calculates a bounding box that includes all the landmarks per set of landmarks in
|
||||
the batch
|
||||
:param target_landmarks: batch of landmarks of dim (n x 68 x 2). Where n is the batch size
|
||||
:param expansion_factor: expands the bounding box by this factor. For example, a `expansion_factor` of 0.2 leads
|
||||
to 20% increase in width and height of the boxes
|
||||
:return: a batch of bounding boxes of dim (n x 4) where the second dim is (x1,y1,x2,y2)
|
||||
"""
|
||||
# Calc bounding box
|
||||
x_y_min, _ = target_landmarks.reshape(-1, 68, 2).min(dim=1)
|
||||
x_y_max, _ = target_landmarks.reshape(-1, 68, 2).max(dim=1)
|
||||
# expanding the bounding box
|
||||
expansion_factor /= 2
|
||||
bb_expansion_x = (x_y_max[:, 0] - x_y_min[:, 0]) * expansion_factor
|
||||
bb_expansion_y = (x_y_max[:, 1] - x_y_min[:, 1]) * expansion_factor
|
||||
x_y_min[:, 0] -= bb_expansion_x
|
||||
x_y_max[:, 0] += bb_expansion_x
|
||||
x_y_min[:, 1] -= bb_expansion_y
|
||||
x_y_max[:, 1] += bb_expansion_y
|
||||
return torch.cat([x_y_min, x_y_max], dim=1)
|
||||
|
||||
|
||||
def shuffle_lr(parts, pairs=None):
|
||||
"""Shuffle the points left-right according to the axis of symmetry
|
||||
of the object.
|
||||
|
||||
Arguments:
|
||||
parts {torch.tensor} -- a 3D or 4D object containing the
|
||||
heatmaps.
|
||||
|
||||
Keyword Arguments:
|
||||
pairs {list of integers} -- [order of the flipped points] (default: {None})
|
||||
"""
|
||||
if pairs is None:
|
||||
pairs = [16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0,
|
||||
26, 25, 24, 23, 22, 21, 20, 19, 18, 17, 27, 28, 29, 30, 35,
|
||||
34, 33, 32, 31, 45, 44, 43, 42, 47, 46, 39, 38, 37, 36, 41,
|
||||
40, 54, 53, 52, 51, 50, 49, 48, 59, 58, 57, 56, 55, 64, 63,
|
||||
62, 61, 60, 67, 66, 65]
|
||||
if parts.ndimension() == 3:
|
||||
parts = parts[pairs, ...]
|
||||
else:
|
||||
parts = parts[:, pairs, ...]
|
||||
|
||||
return parts
|
||||
|
||||
|
||||
def flip(tensor, is_label=False):
|
||||
"""Flip an image or a set of heatmaps left-right
|
||||
|
||||
Arguments:
|
||||
tensor {numpy.array or torch.tensor} -- [the input image or heatmaps]
|
||||
|
||||
Keyword Arguments:
|
||||
is_label {bool} -- [denote wherever the input is an image or a set of heatmaps ] (default: {False})
|
||||
"""
|
||||
if not torch.is_tensor(tensor):
|
||||
tensor = torch.from_numpy(tensor)
|
||||
|
||||
if is_label:
|
||||
tensor = shuffle_lr(tensor).flip(tensor.ndimension() - 1)
|
||||
else:
|
||||
tensor = tensor.flip(tensor.ndimension() - 1)
|
||||
|
||||
return tensor
|
||||
|
||||
|
||||
def get_image(image_or_path):
|
||||
"""Reads an image from file or array/tensor and converts it to RGB (H,W,3).
|
||||
|
||||
Arguments:
|
||||
tensor {Sstring, numpy.array or torch.tensor} -- [the input image or path to it]
|
||||
"""
|
||||
if isinstance(image_or_path, str):
|
||||
try:
|
||||
image = io.imread(image_or_path)
|
||||
except IOError:
|
||||
print("error opening file :: ", image_or_path)
|
||||
return None
|
||||
elif isinstance(image_or_path, torch.Tensor):
|
||||
image = image_or_path.detach().cpu().numpy()
|
||||
else:
|
||||
image = image_or_path
|
||||
|
||||
if image.ndim == 2:
|
||||
image = color.gray2rgb(image)
|
||||
elif image.ndim == 4:
|
||||
image = image[..., :3]
|
||||
|
||||
return image
|
||||
|
||||
|
||||
# Pytorch load supports only pytorch models
|
||||
def load_file_from_url(url, model_dir=None, progress=True, check_hash=False, file_name=None):
|
||||
if model_dir is None:
|
||||
hub_dir = get_dir()
|
||||
model_dir = os.path.join(hub_dir, 'checkpoints')
|
||||
|
||||
try:
|
||||
os.makedirs(model_dir)
|
||||
except OSError as e:
|
||||
if e.errno == errno.EEXIST:
|
||||
# Directory already exists, ignore.
|
||||
pass
|
||||
else:
|
||||
# Unexpected OSError, re-raise.
|
||||
raise
|
||||
|
||||
parts = urlparse(url)
|
||||
filename = os.path.basename(parts.path)
|
||||
if file_name is not None:
|
||||
filename = file_name
|
||||
cached_file = os.path.join(model_dir, filename)
|
||||
if not os.path.exists(cached_file):
|
||||
sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file))
|
||||
hash_prefix = None
|
||||
if check_hash:
|
||||
r = HASH_REGEX.search(filename) # r is Optional[Match[str]]
|
||||
hash_prefix = r.group(1) if r else None
|
||||
download_url_to_file(url, cached_file, hash_prefix, progress=progress)
|
||||
|
||||
return cached_file
|
||||
@@ -1,44 +0,0 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
config for user
|
||||
"""
|
||||
|
||||
import os.path as osp
|
||||
from dataclasses import dataclass
|
||||
#import tyro
|
||||
from typing_extensions import Annotated
|
||||
from .base_config import PrintableConfig, make_abs_path
|
||||
|
||||
|
||||
@dataclass(repr=False) # use repr from PrintableConfig
|
||||
class ArgumentConfig(PrintableConfig):
|
||||
########## input arguments ##########
|
||||
#source_image: Annotated[str, tyro.conf.arg(aliases=["-s"])] = make_abs_path('../../assets/examples/source/s6.jpg') # path to the reference portrait
|
||||
#driving_info: Annotated[str, tyro.conf.arg(aliases=["-d"])] = make_abs_path('../../assets/examples/driving/d0.mp4') # path to driving video or template (.pkl format)
|
||||
#output_dir: Annotated[str, tyro.conf.arg(aliases=["-o"])] = 'animations/' # directory to save output video
|
||||
#####################################
|
||||
|
||||
########## inference arguments ##########
|
||||
device_id: int = 0
|
||||
flag_lip_zero : bool = True # whether let the lip to close state before animation, only take effect when flag_eye_retargeting and flag_lip_retargeting is False
|
||||
flag_eye_retargeting: bool = False
|
||||
flag_lip_retargeting: bool = False
|
||||
flag_stitching: bool = True # we recommend setting it to True!
|
||||
flag_relative: bool = True # whether to use relative pose
|
||||
flag_pasteback: bool = True # whether to paste-back/stitch the animated face cropping from the face-cropping space to the original image space
|
||||
flag_do_crop: bool = True # 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
|
||||
#########################################
|
||||
|
||||
########## crop arguments ##########
|
||||
dsize: int = 512
|
||||
scale: float = 2.3
|
||||
vx_ratio: float = 0 # vx ratio
|
||||
vy_ratio: float = -0.125 # vy ratio +up, -down
|
||||
####################################
|
||||
|
||||
########## gradio arguments ##########
|
||||
#server_port: Annotated[int, tyro.conf.arg(aliases=["-p"])] = 8890
|
||||
#share: bool = False
|
||||
#server_name: str = "0.0.0.0"
|
||||
@@ -1,18 +0,0 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
parameters used for crop faces
|
||||
"""
|
||||
|
||||
import os.path as osp
|
||||
from dataclasses import dataclass
|
||||
from typing import Union, List
|
||||
from .base_config import PrintableConfig
|
||||
|
||||
|
||||
@dataclass(repr=False) # use repr from PrintableConfig
|
||||
class CropConfig(PrintableConfig):
|
||||
dsize: int = 512 # crop size
|
||||
scale: float = 2.3 # scale factor
|
||||
vx_ratio: float = 0 # vx ratio
|
||||
vy_ratio: float = -0.125 # vy ratio +up, -down
|
||||
@@ -29,21 +29,13 @@ class InferenceConfig(PrintableConfig):
|
||||
flag_stitching: bool = True # we recommend setting it to True!
|
||||
|
||||
flag_relative: bool = True # whether to use relative pose
|
||||
anchor_frame: int = 0 # set this value if find_best_frame is True
|
||||
|
||||
input_shape: Tuple[int, int] = (256, 256) # input shape
|
||||
output_format: Literal['mp4', 'gif'] = 'mp4' # output video format
|
||||
output_fps: int = 30 # fps for output video
|
||||
crf: int = 15 # crf for output video
|
||||
|
||||
flag_write_result: bool = True # whether to write 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
|
||||
size_gif: int = 256
|
||||
ref_max_shape: int = 1280
|
||||
ref_shape_n: int = 2
|
||||
|
||||
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
|
||||
|
||||
@@ -4,184 +4,287 @@
|
||||
Pipeline of LivePortrait
|
||||
"""
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import os.path as osp
|
||||
from rich.progress import track
|
||||
|
||||
from .config.inference_config import InferenceConfig
|
||||
|
||||
#from .utils.cropper import Cropper
|
||||
from .utils.camera import get_rotation_matrix
|
||||
#from .utils.video import images2video, concat_frames
|
||||
from .utils.crop import _transform_img
|
||||
#from .utils.retargeting_utils import calc_lip_close_ratio
|
||||
#from .utils.io import load_image_rgb, load_driving_info
|
||||
#from .utils.helper import mkdir, basename, dct2cuda, is_video, is_template, resize_to_limit
|
||||
from .utils.helper import resize_to_limit
|
||||
#from .utils.rprint import rlog as log
|
||||
from .live_portrait_wrapper import LivePortraitWrapper
|
||||
|
||||
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
|
||||
from .utils.camera import get_rotation_matrix
|
||||
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
|
||||
|
||||
def make_abs_path(fn):
|
||||
return osp.join(osp.dirname(osp.realpath(__file__)), fn)
|
||||
|
||||
import os
|
||||
script_directory = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
class LivePortraitPipeline(object):
|
||||
|
||||
def __init__(self, appearance_feature_extractor, motion_extractor, warping_module,
|
||||
spade_generator, stitching_retargeting_module, inference_cfg: InferenceConfig):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
appearance_feature_extractor,
|
||||
motion_extractor,
|
||||
warping_module,
|
||||
spade_generator,
|
||||
stitching_retargeting_module,
|
||||
inference_cfg: InferenceConfig,
|
||||
):
|
||||
self.live_portrait_wrapper: LivePortraitWrapper = LivePortraitWrapper(
|
||||
appearance_feature_extractor, motion_extractor, warping_module,
|
||||
spade_generator, stitching_retargeting_module, cfg=inference_cfg)
|
||||
appearance_feature_extractor,
|
||||
motion_extractor,
|
||||
warping_module,
|
||||
spade_generator,
|
||||
stitching_retargeting_module,
|
||||
cfg=inference_cfg,
|
||||
)
|
||||
|
||||
def execute(self, img_rgb, driving_images_np):
|
||||
inference_cfg = self.live_portrait_wrapper.cfg # for convenience
|
||||
######## process reference portrait ########
|
||||
#img_rgb = load_image_rgb(args.source_image)
|
||||
img_rgb = resize_to_limit(img_rgb, inference_cfg.ref_max_shape, inference_cfg.ref_shape_n)
|
||||
#log(f"Load source image from {args.source_image}")
|
||||
crop_info = self.cropper.crop_single_image(img_rgb)
|
||||
source_lmk = crop_info['lmk_crop']
|
||||
_, img_crop_256x256 = crop_info['img_crop'], crop_info['img_crop_256x256']
|
||||
if inference_cfg.flag_do_crop:
|
||||
I_s = self.live_portrait_wrapper.prepare_source(img_crop_256x256)
|
||||
else:
|
||||
I_s = self.live_portrait_wrapper.prepare_source(img_rgb)
|
||||
x_s_info = self.live_portrait_wrapper.get_kp_info(I_s)
|
||||
x_c_s = x_s_info['kp']
|
||||
R_s = get_rotation_matrix(x_s_info['pitch'], x_s_info['yaw'], x_s_info['roll'])
|
||||
f_s = self.live_portrait_wrapper.extract_feature_3d(I_s)
|
||||
x_s = self.live_portrait_wrapper.transform_keypoint(x_s_info)
|
||||
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,
|
||||
):
|
||||
inference_cfg = self.live_portrait_wrapper.cfg
|
||||
device = inference_cfg.device_id
|
||||
|
||||
if inference_cfg.flag_lip_zero:
|
||||
# let lip-open scalar to be 0 at first
|
||||
c_d_lip_before_animation = [0.]
|
||||
combined_lip_ratio_tensor_before_animation = self.live_portrait_wrapper.calc_combined_lip_ratio(c_d_lip_before_animation, source_lmk)
|
||||
if combined_lip_ratio_tensor_before_animation[0][0] < inference_cfg.lip_zero_threshold:
|
||||
inference_cfg.flag_lip_zero = False
|
||||
else:
|
||||
lip_delta_before_animation = self.live_portrait_wrapper.retarget_lip(x_s, combined_lip_ratio_tensor_before_animation)
|
||||
############################################
|
||||
|
||||
######## process driving info ########
|
||||
#if is_video(args.driving_info):
|
||||
#log(f"Load from video file (mp4 mov avi etc...): {args.driving_info}")
|
||||
# TODO: 这里track一下驱动视频 -> 构建模板
|
||||
#driving_rgb_lst = load_driving_info(args.driving_info)
|
||||
|
||||
driving_rgb_lst = driving_images_np
|
||||
|
||||
driving_rgb_lst_256 = [cv2.resize(_, (256, 256)) for _ in driving_rgb_lst]
|
||||
I_d_lst = self.live_portrait_wrapper.prepare_driving_videos(driving_rgb_lst_256)
|
||||
n_frames = I_d_lst.shape[0]
|
||||
if inference_cfg.flag_eye_retargeting or inference_cfg.flag_lip_retargeting:
|
||||
driving_lmk_lst = self.cropper.get_retargeting_lmk_info(driving_rgb_lst)
|
||||
input_eye_ratio_lst, input_lip_ratio_lst = self.live_portrait_wrapper.calc_retargeting_ratio(source_lmk, driving_lmk_lst)
|
||||
|
||||
# elif is_template(args.driving_info):
|
||||
# log(f"Load from video templates {args.driving_info}")
|
||||
# with open(args.driving_info, 'rb') as f:
|
||||
# template_lst, driving_lmk_lst = pickle.load(f)
|
||||
# n_frames = template_lst[0]['n_frames']
|
||||
# input_eye_ratio_lst, input_lip_ratio_lst = self.live_portrait_wrapper.calc_retargeting_ratio(source_lmk, driving_lmk_lst)
|
||||
# else:
|
||||
# raise Exception("Unsupported driving types!")
|
||||
#########################################
|
||||
|
||||
######## prepare for pasteback ########
|
||||
if inference_cfg.flag_pasteback:
|
||||
if inference_cfg.mask_crop is None:
|
||||
inference_cfg.mask_crop = cv2.imread(make_abs_path('./utils/resources/mask_template.png'), cv2.IMREAD_COLOR)
|
||||
mask_ori = _transform_img(inference_cfg.mask_crop, crop_info['M_c2o'], dsize=(img_rgb.shape[1], img_rgb.shape[0]))
|
||||
mask_ori = mask_ori.astype(np.float32) / 255.
|
||||
I_p_paste_lst = []
|
||||
#########################################
|
||||
|
||||
I_p_lst = []
|
||||
out_list = []
|
||||
R_d_0, x_d_0_info = None, None
|
||||
pbar = comfy.utils.ProgressBar(n_frames)
|
||||
for i in track(range(n_frames), description='Animating...', total=n_frames):
|
||||
#if is_video(args.driving_info):
|
||||
# extract kp info by M
|
||||
I_d_i = I_d_lst[i]
|
||||
x_d_i_info = self.live_portrait_wrapper.get_kp_info(I_d_i)
|
||||
R_d_i = get_rotation_matrix(x_d_i_info['pitch'], x_d_i_info['yaw'], x_d_i_info['roll'])
|
||||
# else:
|
||||
# # from template
|
||||
# x_d_i_info = template_lst[i]
|
||||
# x_d_i_info = dct2cuda(x_d_i_info, inference_cfg.device_id)
|
||||
# R_d_i = x_d_i_info['R_d']
|
||||
|
||||
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
|
||||
else:
|
||||
total_frames = driving_images.shape[0]
|
||||
|
||||
|
||||
disable_progress_bar = True if relative_motion_mode == "single_frame" else False
|
||||
|
||||
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"]
|
||||
|
||||
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):
|
||||
#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))
|
||||
|
||||
if i == 0:
|
||||
R_d_0 = R_d_i
|
||||
x_d_0_info = x_d_i_info
|
||||
first = x_d_info
|
||||
|
||||
if inference_cfg.flag_relative:
|
||||
R_new = (R_d_i @ R_d_0.permute(0, 2, 1)) @ R_s
|
||||
delta_new = x_s_info['exp'] + (x_d_i_info['exp'] - x_d_0_info['exp'])
|
||||
scale_new = x_s_info['scale'] * (x_d_i_info['scale'] / x_d_0_info['scale'])
|
||||
t_new = x_s_info['t'] + (x_d_i_info['t'] - x_d_0_info['t'])
|
||||
driving_info.append(x_d_info)
|
||||
|
||||
driving_exp = source_info[safe_index]["exp"] + x_d_info["exp"] - first["exp"]
|
||||
driving_exp_list.append(driving_exp.cpu())
|
||||
|
||||
R_d = get_rotation_matrix(
|
||||
x_d_info["pitch"], x_d_info["yaw"], x_d_info["roll"]
|
||||
)
|
||||
driving_rot_list.append(R_d)
|
||||
|
||||
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
|
||||
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
|
||||
x_d_r_lst.append(dot)
|
||||
|
||||
driving_exp_list_smooth = smooth(driving_exp_list, source_info[0]["exp"].shape, device, observation_variance=driving_smooth_observation_variance)
|
||||
driving_rot_list_smooth = smooth(x_d_r_lst, source_rot_list[0].shape, device, observation_variance=driving_smooth_observation_variance)
|
||||
|
||||
pbar = comfy.utils.ProgressBar(total_frames)
|
||||
|
||||
for i in tqdm(range(total_frames), desc='Animating...', total=total_frames, disable=disable_progress_bar):
|
||||
|
||||
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)
|
||||
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"]
|
||||
|
||||
#lip zero
|
||||
if inference_cfg.flag_lip_zero:
|
||||
c_d_lip_before_animation = [0.0]
|
||||
combined_lip_ratio_tensor_before_animation = (self.live_portrait_wrapper.calc_combined_lip_ratio(c_d_lip_before_animation, source_lmk))
|
||||
|
||||
if (combined_lip_ratio_tensor_before_animation[0][0] < inference_cfg.lip_zero_threshold):
|
||||
inference_cfg.flag_lip_zero = False
|
||||
else:
|
||||
lip_delta_before_animation = (self.live_portrait_wrapper.retarget_lip(x_s, combined_lip_ratio_tensor_before_animation))
|
||||
|
||||
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"])
|
||||
t_new = x_s_info["t"] + (x_d_info["t"] - x_d_0_info["t"])
|
||||
elif relative_motion_mode == "source_video_smoothed":
|
||||
R_new = driving_rot_list_smooth[i]
|
||||
delta_new = driving_exp_list_smooth[i]
|
||||
scale_new = x_s_info["scale"]
|
||||
t_new = x_d_info["t"]
|
||||
elif relative_motion_mode == "relative_rotation_only":
|
||||
R_new = R_s
|
||||
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_i
|
||||
delta_new = x_d_i_info['exp']
|
||||
scale_new = x_s_info['scale']
|
||||
t_new = x_d_i_info['t']
|
||||
R_new = R_d
|
||||
delta_new = x_s_info['exp']
|
||||
scale_new = x_s_info["scale"]
|
||||
t_new = x_d_info["t"]
|
||||
|
||||
t_new[..., 2].fill_(0) # zero tz
|
||||
t_new[..., 2].fill_(0) # zero tz
|
||||
|
||||
delta_new = delta_new * delta_multiplier
|
||||
|
||||
x_d_i_new = scale_new * (x_c_s @ R_new + delta_new) + t_new
|
||||
|
||||
# Algorithm 1:
|
||||
if not inference_cfg.flag_stitching and not inference_cfg.flag_eye_retargeting and not inference_cfg.flag_lip_retargeting:
|
||||
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)
|
||||
else:
|
||||
pass
|
||||
elif inference_cfg.flag_stitching and not inference_cfg.flag_eye_retargeting and not inference_cfg.flag_lip_retargeting:
|
||||
elif (
|
||||
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(x_s, x_d_i_new) + lip_delta_before_animation.reshape(-1, x_s.shape[1], 3)
|
||||
x_d_i_new = self.live_portrait_wrapper.stitching(
|
||||
x_s, x_d_i_new
|
||||
) + 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:
|
||||
c_d_eyes_i = input_eye_ratio_lst[i]
|
||||
combined_eye_ratio_tensor = self.live_portrait_wrapper.calc_combined_eye_ratio(c_d_eyes_i, source_lmk)
|
||||
combined_eye_ratio_tensor = combined_eye_ratio_tensor * inference_cfg.eyes_retargeting_multiplier
|
||||
c_d_eyes_i = calc_eye_close_ratio(driving_landmarks[i][None])
|
||||
combined_eye_ratio_tensor = (
|
||||
self.live_portrait_wrapper.calc_combined_eye_ratio(
|
||||
c_d_eyes_i, source_lmk
|
||||
)
|
||||
)
|
||||
combined_eye_ratio_tensor = (
|
||||
combined_eye_ratio_tensor
|
||||
* inference_cfg.eyes_retargeting_multiplier
|
||||
)
|
||||
# ∆_eyes,i = R_eyes(x_s; c_s,eyes, c_d,eyes,i)
|
||||
eyes_delta = self.live_portrait_wrapper.retarget_eye(x_s, combined_eye_ratio_tensor)
|
||||
eyes_delta = self.live_portrait_wrapper.retarget_eye(
|
||||
x_s, combined_eye_ratio_tensor
|
||||
)
|
||||
if inference_cfg.flag_lip_retargeting:
|
||||
c_d_lip_i = input_lip_ratio_lst[i]
|
||||
combined_lip_ratio_tensor = self.live_portrait_wrapper.calc_combined_lip_ratio(c_d_lip_i, source_lmk)
|
||||
combined_lip_ratio_tensor = combined_lip_ratio_tensor * inference_cfg.lip_retargeting_multiplier
|
||||
c_d_lip_i = calc_lip_close_ratio(driving_landmarks[i][None])
|
||||
combined_lip_ratio_tensor = (
|
||||
self.live_portrait_wrapper.calc_combined_lip_ratio(
|
||||
c_d_lip_i, source_lmk
|
||||
)
|
||||
)
|
||||
combined_lip_ratio_tensor = (
|
||||
combined_lip_ratio_tensor
|
||||
* inference_cfg.lip_retargeting_multiplier
|
||||
)
|
||||
# ∆_lip,i = R_lip(x_s; c_s,lip, c_d,lip,i)
|
||||
lip_delta = self.live_portrait_wrapper.retarget_lip(x_s, combined_lip_ratio_tensor)
|
||||
lip_delta = self.live_portrait_wrapper.retarget_lip(
|
||||
x_s, combined_lip_ratio_tensor
|
||||
)
|
||||
|
||||
if inference_cfg.flag_relative: # use x_s
|
||||
x_d_i_new = x_s + \
|
||||
(eyes_delta.reshape(-1, x_s.shape[1], 3) if eyes_delta is not None else 0) + \
|
||||
(lip_delta.reshape(-1, x_s.shape[1], 3) if lip_delta is not None else 0)
|
||||
if relative_motion_mode != "off": # use x_s
|
||||
x_d_i_new = (
|
||||
x_s
|
||||
+ (
|
||||
eyes_delta.reshape(-1, x_s.shape[1], 3)
|
||||
if eyes_delta is not None
|
||||
else 0
|
||||
)
|
||||
+ (
|
||||
lip_delta.reshape(-1, x_s.shape[1], 3)
|
||||
if lip_delta is not None
|
||||
else 0
|
||||
)
|
||||
)
|
||||
else: # use x_d,i
|
||||
x_d_i_new = x_d_i_new + \
|
||||
(eyes_delta.reshape(-1, x_s.shape[1], 3) if eyes_delta is not None else 0) + \
|
||||
(lip_delta.reshape(-1, x_s.shape[1], 3) if lip_delta is not None else 0)
|
||||
x_d_i_new = (
|
||||
x_d_i_new
|
||||
+ (
|
||||
eyes_delta.reshape(-1, x_s.shape[1], 3)
|
||||
if eyes_delta is not None
|
||||
else 0
|
||||
)
|
||||
+ (
|
||||
lip_delta.reshape(-1, x_s.shape[1], 3)
|
||||
if lip_delta is not None
|
||||
else 0
|
||||
)
|
||||
)
|
||||
|
||||
if inference_cfg.flag_stitching:
|
||||
x_d_i_new = self.live_portrait_wrapper.stitching(x_s, x_d_i_new)
|
||||
|
||||
if inference_cfg.flag_stitching:
|
||||
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)
|
||||
I_p_i = self.live_portrait_wrapper.parse_output(out['out'])[0]
|
||||
I_p_lst.append(I_p_i)
|
||||
|
||||
out_list.append(out)
|
||||
|
||||
pbar.update(1)
|
||||
|
||||
#if inference_cfg.flag_pasteback:
|
||||
I_p_i_to_ori = _transform_img(I_p_i, crop_info['M_c2o'], dsize=(img_rgb.shape[1], img_rgb.shape[0]))
|
||||
I_p_i_to_ori_blend = np.clip(mask_ori * I_p_i_to_ori + (1 - mask_ori) * img_rgb, 0, 255).astype(np.uint8)
|
||||
out = np.hstack([I_p_i_to_ori, I_p_i_to_ori_blend])
|
||||
I_p_paste_lst.append(I_p_i_to_ori_blend)
|
||||
out_dict = {
|
||||
"out_list": out_list,
|
||||
"crop_info": crop_info,
|
||||
"mismatch_method": mismatch_method,
|
||||
}
|
||||
|
||||
return I_p_lst, I_p_paste_lst
|
||||
return out_dict
|
||||
|
||||
@@ -13,6 +13,7 @@ from .utils.retargeting_utils import compute_eye_delta, compute_lip_delta
|
||||
from .utils.camera import headpose_pred_to_degree, get_rotation_matrix
|
||||
from .utils.retargeting_utils import calc_eye_close_ratio, calc_lip_close_ratio
|
||||
from .config.inference_config import InferenceConfig
|
||||
from contextlib import nullcontext
|
||||
|
||||
from comfy.model_management import get_autocast_device
|
||||
|
||||
@@ -31,11 +32,6 @@ class LivePortraitWrapper(object):
|
||||
self.device_id = cfg.device_id
|
||||
self.timer = Timer()
|
||||
|
||||
def update_config(self, user_args):
|
||||
for k, v in user_args.items():
|
||||
if hasattr(self.cfg, k):
|
||||
setattr(self.cfg, k, v)
|
||||
|
||||
def prepare_source(self, img: np.ndarray) -> torch.Tensor:
|
||||
""" construct the input as standard
|
||||
img: HxWx3, uint8, 256x256
|
||||
@@ -57,31 +53,12 @@ class LivePortraitWrapper(object):
|
||||
x = x.to(self.device_id)
|
||||
return x
|
||||
|
||||
def prepare_driving_videos(self, imgs) -> torch.Tensor:
|
||||
""" construct the input as standard
|
||||
imgs: NxBxHxWx3, uint8
|
||||
"""
|
||||
if isinstance(imgs, list):
|
||||
_imgs = np.array(imgs)[..., np.newaxis] # TxHxWx3x1
|
||||
elif isinstance(imgs, np.ndarray):
|
||||
_imgs = imgs
|
||||
else:
|
||||
raise ValueError(f'imgs type error: {type(imgs)}')
|
||||
|
||||
y = _imgs.astype(np.float32) / 255.
|
||||
y = np.clip(y, 0, 1) # clip to 0~1
|
||||
y = torch.from_numpy(y).permute(0, 4, 3, 1, 2) # TxHxWx3x1 -> Tx1x3xHxW
|
||||
y = y.to(self.device_id)
|
||||
|
||||
return y
|
||||
|
||||
def extract_feature_3d(self, x: torch.Tensor) -> torch.Tensor:
|
||||
""" get the appearance feature of the image by F
|
||||
x: Bx3xHxW, normalized to 0~1
|
||||
"""
|
||||
with torch.no_grad():
|
||||
with torch.autocast(device_type=get_autocast_device(self.device_id), dtype=torch.float16, enabled=self.cfg.flag_use_half_precision):
|
||||
feature_3d = self.appearance_feature_extractor(x)
|
||||
with torch.autocast(get_autocast_device(self.device_id), dtype=torch.float16) if self.cfg.flag_use_half_precision else nullcontext():
|
||||
feature_3d = self.appearance_feature_extractor(x)
|
||||
|
||||
return feature_3d.float()
|
||||
|
||||
@@ -91,15 +68,14 @@ class LivePortraitWrapper(object):
|
||||
flag_refine_info: whether to trandform the pose to degrees and the dimention of the reshape
|
||||
return: A dict contains keys: 'pitch', 'yaw', 'roll', 't', 'exp', 'scale', 'kp'
|
||||
"""
|
||||
with torch.no_grad():
|
||||
with torch.autocast(device_type=get_autocast_device(self.device_id), dtype=torch.float16, enabled=self.cfg.flag_use_half_precision):
|
||||
kp_info = self.motion_extractor(x)
|
||||
with torch.autocast(get_autocast_device(self.device_id), dtype=torch.float16) if self.cfg.flag_use_half_precision else nullcontext():
|
||||
kp_info = self.motion_extractor(x)
|
||||
|
||||
if self.cfg.flag_use_half_precision:
|
||||
# float the dict
|
||||
for k, v in kp_info.items():
|
||||
if isinstance(v, torch.Tensor):
|
||||
kp_info[k] = v.float()
|
||||
if self.cfg.flag_use_half_precision:
|
||||
# float the dict
|
||||
for k, v in kp_info.items():
|
||||
if isinstance(v, torch.Tensor):
|
||||
kp_info[k] = v.float()
|
||||
|
||||
flag_refine_info: bool = kwargs.get('flag_refine_info', True)
|
||||
if flag_refine_info:
|
||||
@@ -194,35 +170,6 @@ 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
|
||||
@@ -265,32 +212,22 @@ class LivePortraitWrapper(object):
|
||||
kp_source: BxNx3
|
||||
kp_driving: BxNx3
|
||||
"""
|
||||
# The line 18 in Algorithm 1: D(W(f_s; x_s, x′_d,i)）
|
||||
with torch.no_grad():
|
||||
with torch.autocast(device_type=get_autocast_device(self.device_id), dtype=torch.float16, enabled=self.cfg.flag_use_half_precision):
|
||||
# get decoder input
|
||||
ret_dct = self.warping_module(feature_3d, kp_source=kp_source, kp_driving=kp_driving)
|
||||
# decode
|
||||
ret_dct['out'] = self.spade_generator(feature=ret_dct['out'])
|
||||
# The line 18 in Algorithm 1: D(W(f_s; x_s, x′_d,i)
|
||||
with torch.autocast(get_autocast_device(self.device_id), dtype=torch.float16) if self.cfg.flag_use_half_precision else nullcontext():
|
||||
# get decoder input
|
||||
ret_dct = self.warping_module(feature_3d, kp_source=kp_source, kp_driving=kp_driving)
|
||||
# decode
|
||||
ret_dct['out'] = self.spade_generator(feature=ret_dct['out'])
|
||||
|
||||
# float the dict
|
||||
if self.cfg.flag_use_half_precision:
|
||||
for k, v in ret_dct.items():
|
||||
if isinstance(v, torch.Tensor):
|
||||
ret_dct[k] = v.float()
|
||||
# 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()
|
||||
|
||||
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 = []
|
||||
@@ -303,17 +240,19 @@ class LivePortraitWrapper(object):
|
||||
|
||||
def calc_combined_eye_ratio(self, input_eye_ratio, source_lmk):
|
||||
eye_close_ratio = calc_eye_close_ratio(source_lmk[None])
|
||||
eye_close_ratio_tensor = torch.from_numpy(eye_close_ratio).float().to(self.device_id)
|
||||
input_eye_ratio_tensor = torch.Tensor([input_eye_ratio[0][0]]).reshape(1, 1).to(self.device_id)
|
||||
eye_close_ratios_tensor = torch.from_numpy(eye_close_ratio).float().to(self.device_id)
|
||||
input_eye_ratio_array = np.array(input_eye_ratio[0][0]).reshape(1, 1)
|
||||
input_eye_ratio_tensor = torch.from_numpy(input_eye_ratio_array).float().to(self.device_id)
|
||||
# [c_s,eyes, c_d,eyes,i]
|
||||
combined_eye_ratio_tensor = torch.cat([eye_close_ratio_tensor, input_eye_ratio_tensor], dim=1)
|
||||
return combined_eye_ratio_tensor
|
||||
combined_eye_ratios_tensor = torch.cat([eye_close_ratios_tensor, input_eye_ratio_tensor], dim=1)
|
||||
return combined_eye_ratios_tensor
|
||||
|
||||
def calc_combined_lip_ratio(self, input_lip_ratio, source_lmk):
|
||||
lip_close_ratio = calc_lip_close_ratio(source_lmk[None])
|
||||
lip_close_ratio_tensor = torch.from_numpy(lip_close_ratio).float().to(self.device_id)
|
||||
# [c_s,lip, c_d,lip,i]
|
||||
input_lip_ratio_tensor = torch.Tensor([input_lip_ratio[0]]).to(self.device_id)
|
||||
input_lip_ratio_array = np.array([input_lip_ratio[0]])
|
||||
input_lip_ratio_tensor = torch.from_numpy(input_lip_ratio_array).float().to(self.device_id)
|
||||
if input_lip_ratio_tensor.shape != [1, 1]:
|
||||
input_lip_ratio_tensor = input_lip_ratio_tensor.reshape(1, 1)
|
||||
combined_lip_ratio_tensor = torch.cat([lip_close_ratio_tensor, input_lip_ratio_tensor], dim=1)
|
||||
|
||||
@@ -47,7 +47,13 @@ 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)
|
||||
sparse_deformed = F.grid_sample(feature_repeat, sparse_motions, align_corners=False)
|
||||
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 = sparse_deformed.view((bs, self.num_kp+1, -1, d, h, w)) # (bs, num_kp+1, c, d, h, w)
|
||||
|
||||
return sparse_deformed
|
||||
@@ -61,7 +67,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:
|
||||
except ValueError:
|
||||
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,7 +158,11 @@ class DownBlock3d(nn.Module):
|
||||
out = self.conv(x)
|
||||
out = self.norm(out)
|
||||
out = F.relu(out)
|
||||
out = self.pool(out)
|
||||
try:
|
||||
out = self.pool(out)
|
||||
except NotImplementedError:
|
||||
out_device = out.device # Store input device
|
||||
out = self.pool(out.to('cpu')).to(out_device)
|
||||
return out
|
||||
|
||||
|
||||
|
||||
@@ -44,7 +44,11 @@ class WarpingNetwork(nn.Module):
|
||||
self.estimate_occlusion_map = estimate_occlusion_map
|
||||
|
||||
def deform_input(self, inp, deformation):
|
||||
return F.grid_sample(inp, deformation, align_corners=False)
|
||||
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)
|
||||
|
||||
def forward(self, feature_3d, kp_driving, kp_source):
|
||||
if self.dense_motion_network is not None:
|
||||
|
||||
@@ -1,65 +0,0 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
Make video template
|
||||
"""
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
import pickle
|
||||
from rich.progress import track
|
||||
from .utils.cropper import Cropper
|
||||
|
||||
from .utils.io import load_driving_info
|
||||
from .utils.camera import get_rotation_matrix
|
||||
from .utils.helper import mkdir, basename
|
||||
from .utils.rprint import rlog as log
|
||||
from .config.crop_config import CropConfig
|
||||
from .config.inference_config import InferenceConfig
|
||||
from .live_portrait_wrapper import LivePortraitWrapper
|
||||
|
||||
class TemplateMaker:
|
||||
|
||||
def __init__(self, inference_cfg: InferenceConfig, crop_cfg: CropConfig):
|
||||
self.live_portrait_wrapper: LivePortraitWrapper = LivePortraitWrapper(cfg=inference_cfg)
|
||||
self.cropper = Cropper(crop_cfg=crop_cfg)
|
||||
|
||||
def make_motion_template(self, video_fp: str, output_path: str, **kwargs):
|
||||
""" make video template (.pkl format)
|
||||
video_fp: driving video file path
|
||||
output_path: where to save the pickle file
|
||||
"""
|
||||
|
||||
driving_rgb_lst = load_driving_info(video_fp)
|
||||
driving_rgb_lst = [cv2.resize(_, (256, 256)) for _ in driving_rgb_lst]
|
||||
driving_lmk_lst = self.cropper.get_retargeting_lmk_info(driving_rgb_lst)
|
||||
I_d_lst = self.live_portrait_wrapper.prepare_driving_videos(driving_rgb_lst)
|
||||
|
||||
n_frames = I_d_lst.shape[0]
|
||||
|
||||
templates = []
|
||||
|
||||
|
||||
for i in track(range(n_frames), description='Making templates...', total=n_frames):
|
||||
I_d_i = I_d_lst[i]
|
||||
x_d_i_info = self.live_portrait_wrapper.get_kp_info(I_d_i)
|
||||
R_d_i = get_rotation_matrix(x_d_i_info['pitch'], x_d_i_info['yaw'], x_d_i_info['roll'])
|
||||
# collect s_d, R_d, δ_d and t_d for inference
|
||||
template_dct = {
|
||||
'n_frames': n_frames,
|
||||
'frames_index': i,
|
||||
}
|
||||
template_dct['scale'] = x_d_i_info['scale'].cpu().numpy().astype(np.float32)
|
||||
template_dct['R_d'] = R_d_i.cpu().numpy().astype(np.float32)
|
||||
template_dct['exp'] = x_d_i_info['exp'].cpu().numpy().astype(np.float32)
|
||||
template_dct['t'] = x_d_i_info['t'].cpu().numpy().astype(np.float32)
|
||||
|
||||
templates.append(template_dct)
|
||||
|
||||
mkdir(output_path)
|
||||
# Save the dictionary as a pickle file
|
||||
pickle_fp = os.path.join(output_path, f'{basename(video_fp)}.pkl')
|
||||
with open(pickle_fp, 'wb') as f:
|
||||
pickle.dump([templates, driving_lmk_lst], f)
|
||||
log(f"Template saved at {pickle_fp}")
|
||||
+105
-24
@@ -4,14 +4,12 @@
|
||||
cropping function and the related preprocess functions for cropping
|
||||
"""
|
||||
|
||||
import cv2; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False) # NOTE: enforce single thread
|
||||
import cv2#; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False) # NOTE: enforce single thread
|
||||
import numpy as np
|
||||
from .rprint import rprint as print
|
||||
from math import sin, cos, acos, degrees
|
||||
|
||||
DTYPE = np.float32
|
||||
CV2_INTERP = cv2.INTER_LINEAR
|
||||
|
||||
import comfy.model_management as mm
|
||||
|
||||
def _transform_img(img, M, dsize, flags=CV2_INTERP, borderMode=None):
|
||||
""" conduct similarity or affine transformation to the image, do not do border operation!
|
||||
@@ -29,6 +27,37 @@ def _transform_img(img, M, dsize, flags=CV2_INTERP, borderMode=None):
|
||||
else:
|
||||
return cv2.warpAffine(img, M[:2, :], dsize=_dsize, flags=flags)
|
||||
|
||||
import torch
|
||||
import kornia.geometry.transform as KGT
|
||||
|
||||
def _transform_img_kornia(img, M, dsize, device, 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).
|
||||
"""
|
||||
|
||||
# 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, :]
|
||||
if M.shape == (2, 3):
|
||||
M = M.unsqueeze(0)
|
||||
|
||||
# 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)
|
||||
|
||||
# Apply the affine transformation
|
||||
img_warped = KGT.warp_affine(img, M, _dsize, mode=flags, padding_mode=borderMode)
|
||||
|
||||
return img_warped
|
||||
|
||||
def _transform_pts(pts, M):
|
||||
""" conduct similarity or affine transformation to the pts
|
||||
@@ -39,6 +68,23 @@ 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
|
||||
@@ -89,29 +135,60 @@ 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": []}
|
||||
|
||||
def parse_pt2_from_pt68(pt68, use_lip=True):
|
||||
"""
|
||||
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:
|
||||
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
|
||||
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([
|
||||
(pt5[0] + pt5[1]) / 2,
|
||||
(pt5[3] + pt5[4]) / 2
|
||||
(pt9[0] + pt9[1]) / 2, # eye
|
||||
pt9[3] # lip
|
||||
], axis=0)
|
||||
else:
|
||||
pt2 = np.stack([
|
||||
np.mean(pt68[lm_idx[[1, 2]], :], 0), # left eye
|
||||
np.mean(pt68[lm_idx[[3, 4]], :], 0), # right eye
|
||||
(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": []}
|
||||
|
||||
|
||||
'''
|
||||
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
|
||||
|
||||
], axis=0)
|
||||
pt2 = np.stack([
|
||||
(pt68[0] + pt68[1]) / 2,
|
||||
pt68[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
|
||||
], axis=0)
|
||||
|
||||
return pt2
|
||||
@@ -145,9 +222,13 @@ 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}')
|
||||
|
||||
@@ -350,13 +431,15 @@ def crop_image(img, pts: np.ndarray, **kwargs):
|
||||
dsize = kwargs.get('dsize', 224)
|
||||
scale = kwargs.get('scale', 1.5) # 1.5 | 1.6
|
||||
vy_ratio = kwargs.get('vy_ratio', -0.1) # -0.0625 | -0.1
|
||||
vx_ratio = kwargs.get('vx_ratio', 0)
|
||||
|
||||
M_INV, _ = _estimate_similar_transform_from_pts(
|
||||
pts,
|
||||
dsize=dsize,
|
||||
scale=scale,
|
||||
vy_ratio=vy_ratio,
|
||||
flag_do_rot=kwargs.get('flag_do_rot', True),
|
||||
vx_ratio=vx_ratio,
|
||||
flag_do_rot=kwargs.get('rotate', True),
|
||||
)
|
||||
|
||||
if img is None:
|
||||
@@ -379,15 +462,13 @@ 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
|
||||
return ret_dct, img_crop
|
||||
|
||||
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()
|
||||
|
||||
|
||||
+159
-90
@@ -1,72 +1,50 @@
|
||||
# coding: utf-8
|
||||
|
||||
import numpy as np
|
||||
import os.path as osp
|
||||
from typing import List, Union, Tuple
|
||||
from dataclasses import dataclass, field
|
||||
import cv2; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False)
|
||||
import cv2#; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False)
|
||||
|
||||
from .landmark_runner import LandmarkRunner
|
||||
from .face_analysis_diy import FaceAnalysisDIY
|
||||
#from .helper import prefix
|
||||
from .crop import crop_image, crop_image_by_bbox, parse_bbox_from_landmark, average_bbox_lst
|
||||
#from .timer import Timer
|
||||
from .rprint import rlog as log
|
||||
from .io import load_image_rgb
|
||||
#from .video import VideoWriter, get_fps, change_video_fps
|
||||
from .landmark_runner import LandmarkRunner, LandmarkRunnerTorch
|
||||
|
||||
from .crop import crop_image
|
||||
|
||||
import folder_paths
|
||||
import os
|
||||
script_directory = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
def make_abs_path(fn):
|
||||
return osp.join(osp.dirname(osp.realpath(__file__)), fn)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Trajectory:
|
||||
start: int = -1 # 起始帧 闭区间
|
||||
end: int = -1 # 结束帧 闭区间
|
||||
start: int = -1
|
||||
end: int = -1
|
||||
lmk_lst: Union[Tuple, List, np.ndarray] = field(default_factory=list) # lmk list
|
||||
bbox_lst: Union[Tuple, List, np.ndarray] = field(default_factory=list) # bbox list
|
||||
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 Cropper(object):
|
||||
def __init__(self, provider, **kwargs) -> None:
|
||||
class CropperInsightFace(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=make_abs_path('../../pretrained_weights/liveportrait/landmark.onnx'),
|
||||
ckpt_path=os.path.join(folder_paths.models_dir, 'liveportrait', 'landmark.onnx'),
|
||||
onnx_provider=provider,
|
||||
device_id=device_id
|
||||
)
|
||||
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))
|
||||
self.face_analysis_wrapper.prepare(ctx_id=device_id, det_size=(512, 512), det_thresh=detection_threshold)
|
||||
self.face_analysis_wrapper.warmup()
|
||||
|
||||
self.crop_cfg = kwargs.get('crop_cfg', None)
|
||||
|
||||
def update_config(self, user_args):
|
||||
for k, v in user_args.items():
|
||||
if hasattr(self.crop_cfg, k):
|
||||
setattr(self.crop_cfg, k, v)
|
||||
|
||||
def crop_single_image(self, obj, **kwargs):
|
||||
direction = kwargs.get('direction', 'large-small')
|
||||
|
||||
# crop and align a single image
|
||||
if isinstance(obj, str):
|
||||
img_rgb = load_image_rgb(obj)
|
||||
elif isinstance(obj, np.ndarray):
|
||||
img_rgb = obj
|
||||
def crop_single_image(self, img_rgb, dsize, scale, vy_ratio, vx_ratio, face_index, face_index_order, rotate):
|
||||
direction = face_index_order
|
||||
|
||||
src_face = self.face_analysis_wrapper.get(
|
||||
img_rgb,
|
||||
@@ -75,73 +53,164 @@ class Cropper(object):
|
||||
)
|
||||
|
||||
if len(src_face) == 0:
|
||||
log('No face detected in the source image.')
|
||||
raise Exception("No face detected in the source image!")
|
||||
elif len(src_face) > 1:
|
||||
log(f'More than one face detected in the image, only pick one face by rule {direction}.')
|
||||
ret_dct = {}
|
||||
cropped_image_256 = None
|
||||
return ret_dct, cropped_image_256
|
||||
|
||||
src_face = src_face[0]
|
||||
src_face = src_face[face_index] # choose the index if multiple faces detected
|
||||
pts = src_face.landmark_2d_106
|
||||
|
||||
|
||||
# crop the face
|
||||
ret_dct = crop_image(
|
||||
ret_dct, image_crop = crop_image(
|
||||
img_rgb, # ndarray
|
||||
pts, # 106x2 or Nx2
|
||||
dsize=kwargs.get('dsize', 512),
|
||||
scale=kwargs.get('scale', 2.3),
|
||||
vy_ratio=kwargs.get('vy_ratio', -0.15),
|
||||
dsize=dsize,
|
||||
scale=scale,
|
||||
vy_ratio=vy_ratio,
|
||||
vx_ratio=vx_ratio,
|
||||
rotate=rotate
|
||||
)
|
||||
# update a 256x256 version for network input or else
|
||||
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 / kwargs.get('dsize', 512)
|
||||
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
|
||||
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')
|
||||
|
||||
def get_retargeting_lmk_info(self, driving_rgb_lst):
|
||||
# TODO: implement a tracking-based version
|
||||
driving_lmk_lst = []
|
||||
for driving_image in driving_rgb_lst:
|
||||
ret_dct = self.crop_single_image(driving_image)
|
||||
driving_lmk_lst.append(ret_dct['lmk_crop'])
|
||||
return driving_lmk_lst
|
||||
|
||||
def make_video_clip(self, driving_rgb_lst, output_path, output_fps=30, **kwargs):
|
||||
trajectory = Trajectory()
|
||||
direction = kwargs.get('direction', 'large-small')
|
||||
for idx, driving_image in enumerate(driving_rgb_lst):
|
||||
if idx == 0 or trajectory.start == -1:
|
||||
src_face = self.face_analysis_wrapper.get(
|
||||
driving_image,
|
||||
flag_do_landmark_2d_106=True,
|
||||
direction=direction
|
||||
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
|
||||
)
|
||||
if len(src_face) == 0:
|
||||
# No face detected in the driving_image
|
||||
continue
|
||||
elif len(src_face) > 1:
|
||||
log(f'More than one face detected in the driving frame_{idx}, only pick one face by rule {direction}.')
|
||||
src_face = src_face[0]
|
||||
pts = src_face.landmark_2d_106
|
||||
lmk_203 = self.landmark_runner(driving_image, pts)['pts']
|
||||
trajectory.start, trajectory.end = idx, idx
|
||||
else:
|
||||
lmk_203 = self.face_recon_wrapper(driving_image, trajectory.lmk_lst[-1])['pts']
|
||||
trajectory.end = idx
|
||||
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()
|
||||
|
||||
trajectory.lmk_lst.append(lmk_203)
|
||||
ret_bbox = parse_bbox_from_landmark(lmk_203, scale=self.crop_cfg.globalscale, vy_ratio=elf.crop_cfg.vy_ratio)['bbox']
|
||||
bbox = [ret_bbox[0, 0], ret_bbox[0, 1], ret_bbox[2, 0], ret_bbox[2, 1]] # 4,
|
||||
trajectory.bbox_lst.append(bbox) # bbox
|
||||
trajectory.frame_rgb_lst.append(driving_image)
|
||||
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)
|
||||
|
||||
global_bbox = average_bbox_lst(trajectory.bbox_lst)
|
||||
for idx, (frame_rgb, lmk) in enumerate(zip(trajectory.frame_rgb_lst, trajectory.lmk_lst)):
|
||||
ret_dct = crop_image_by_bbox(
|
||||
frame_rgb, global_bbox, lmk=lmk,
|
||||
dsize=self.video_crop_cfg.dsize, flag_rot=self.video_crop_cfg.flag_rot, borderValue=self.video_crop_cfg.borderValue
|
||||
)
|
||||
frame_rgb_crop = ret_dct['img_crop']
|
||||
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
|
||||
|
||||
|
||||
@@ -1,16 +1,30 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
face detectoin and alignment using InsightFace
|
||||
face detection and alignment using InsightFace
|
||||
"""
|
||||
from insightface.utils import transform
|
||||
|
||||
#patch Insightface function to get rid of the annoying warnings
|
||||
def patched_estimate_affine_matrix_3d23d(X, Y):
|
||||
''' Using least-squares solution
|
||||
Args:
|
||||
X: [n, 3]. 3d points(fixed)
|
||||
Y: [n, 3]. corresponding 3d points(moving). Y = PX
|
||||
Returns:
|
||||
P_Affine: (3, 4). Affine camera matrix (the third row is [0, 0, 0, 1]).
|
||||
'''
|
||||
X_homo = np.hstack((X, np.ones([X.shape[0],1]))) # n x 4
|
||||
P = np.linalg.lstsq(X_homo, Y, rcond=None)[0].T # Affine matrix. 3 x 4
|
||||
return P
|
||||
|
||||
transform.estimate_affine_matrix_3d23d = patched_estimate_affine_matrix_3d23d
|
||||
|
||||
import numpy as np
|
||||
from .rprint import rlog as log
|
||||
from insightface.app import FaceAnalysis
|
||||
from insightface.app.common import Face
|
||||
from .timer import Timer
|
||||
|
||||
|
||||
def sort_by_direction(faces, direction: str = 'large-small', face_center=None):
|
||||
if len(faces) <= 0:
|
||||
return faces
|
||||
@@ -76,4 +90,4 @@ class FaceAnalysisDIY(FaceAnalysis):
|
||||
self.get(img_bgr)
|
||||
|
||||
elapse = self.timer.toc()
|
||||
log(f'FaceAnalysisDIY warmup time: {elapse:.3f}s')
|
||||
print(f'FaceAnalysisDIY warmup time: {elapse:.3f}s')
|
||||
|
||||
@@ -0,0 +1,30 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
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_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]
|
||||
|
||||
return x_d_lst_smooth
|
||||
@@ -4,58 +4,17 @@
|
||||
utility functions and classes to handle feature extraction and model loading
|
||||
"""
|
||||
|
||||
import os
|
||||
import os.path as osp
|
||||
import cv2
|
||||
import torch
|
||||
import numpy as np
|
||||
from typing import Union
|
||||
from collections import OrderedDict
|
||||
|
||||
def suffix(filename):
|
||||
"""a.jpg -> jpg"""
|
||||
pos = filename.rfind(".")
|
||||
if pos == -1:
|
||||
return ""
|
||||
return filename[pos + 1:]
|
||||
|
||||
|
||||
def prefix(filename):
|
||||
"""a.jpg -> a"""
|
||||
pos = filename.rfind(".")
|
||||
if pos == -1:
|
||||
return filename
|
||||
return filename[:pos]
|
||||
|
||||
|
||||
def basename(filename):
|
||||
"""a/b/c.jpg -> c"""
|
||||
return prefix(osp.basename(filename))
|
||||
|
||||
|
||||
def is_video(file_path):
|
||||
if file_path.lower().endswith((".mp4", ".mov", ".avi", ".webm")) or osp.isdir(file_path):
|
||||
return True
|
||||
return False
|
||||
|
||||
def is_template(file_path):
|
||||
if file_path.endswith(".pkl"):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def mkdir(d, log=False):
|
||||
# return self-assined `d`, for one line code
|
||||
if not osp.exists(d):
|
||||
os.makedirs(d, exist_ok=True)
|
||||
if log:
|
||||
print(f"Make dir: {d}")
|
||||
return d
|
||||
|
||||
from scipy.spatial import ConvexHull # pylint: disable=E0401,E0611
|
||||
|
||||
def squeeze_tensor_to_numpy(tensor):
|
||||
out = tensor.data.squeeze(0).cpu().numpy()
|
||||
return out
|
||||
|
||||
|
||||
def dct2cuda(dct: dict, device_id: int):
|
||||
for key in dct:
|
||||
dct[key] = torch.tensor(dct[key]).to(device_id)
|
||||
@@ -95,11 +54,6 @@ def calculate_transformation(config, s_kp_info, t_0_kp_info, t_i_kp_info, R_s, R
|
||||
new_scale = s_kp_info['scale'] * (t_i_kp_info['scale'] / t_0_kp_info['scale'])
|
||||
return new_rotation, new_expression, new_translation, new_scale
|
||||
|
||||
def load_description(fp):
|
||||
with open(fp, 'r', encoding='utf-8') as f:
|
||||
content = f.read()
|
||||
return content
|
||||
|
||||
|
||||
def resize_to_limit(img, max_dim=1280, n=2):
|
||||
h, w = img.shape[:2]
|
||||
@@ -119,3 +73,24 @@ 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
|
||||
@@ -1,97 +0,0 @@
|
||||
# coding: utf-8
|
||||
|
||||
import os
|
||||
from glob import glob
|
||||
import os.path as osp
|
||||
import imageio
|
||||
import numpy as np
|
||||
import cv2; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False)
|
||||
|
||||
|
||||
def load_image_rgb(image_path: str):
|
||||
if not osp.exists(image_path):
|
||||
raise FileNotFoundError(f"Image not found: {image_path}")
|
||||
img = cv2.imread(image_path, cv2.IMREAD_COLOR)
|
||||
return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
|
||||
|
||||
def load_driving_info(driving_info):
|
||||
driving_video_ori = []
|
||||
|
||||
def load_images_from_directory(directory):
|
||||
image_paths = sorted(glob(osp.join(directory, '*.png')) + glob(osp.join(directory, '*.jpg')))
|
||||
return [load_image_rgb(im_path) for im_path in image_paths]
|
||||
|
||||
def load_images_from_video(file_path):
|
||||
reader = imageio.get_reader(file_path)
|
||||
return [image for idx, image in enumerate(reader)]
|
||||
|
||||
if osp.isdir(driving_info):
|
||||
driving_video_ori = load_images_from_directory(driving_info)
|
||||
elif osp.isfile(driving_info):
|
||||
driving_video_ori = load_images_from_video(driving_info)
|
||||
|
||||
return driving_video_ori
|
||||
|
||||
|
||||
def contiguous(obj):
|
||||
if not obj.flags.c_contiguous:
|
||||
obj = obj.copy(order="C")
|
||||
return obj
|
||||
|
||||
|
||||
def _resize_to_limit(img: np.ndarray, max_dim=1920, n=2):
|
||||
"""
|
||||
ajust the size of the image so that the maximum dimension does not exceed max_dim, and the width and the height of the image are multiples of n.
|
||||
:param img: the image to be processed.
|
||||
:param max_dim: the maximum dimension constraint.
|
||||
:param n: the number that needs to be multiples of.
|
||||
:return: the adjusted image.
|
||||
"""
|
||||
h, w = img.shape[:2]
|
||||
|
||||
# ajust the size of the image according to the maximum dimension
|
||||
if max_dim > 0 and max(h, w) > max_dim:
|
||||
if h > w:
|
||||
new_h = max_dim
|
||||
new_w = int(w * (max_dim / h))
|
||||
else:
|
||||
new_w = max_dim
|
||||
new_h = int(h * (max_dim / w))
|
||||
img = cv2.resize(img, (new_w, new_h))
|
||||
|
||||
# ensure that the image dimensions are multiples of n
|
||||
n = max(n, 1)
|
||||
new_h = img.shape[0] - (img.shape[0] % n)
|
||||
new_w = img.shape[1] - (img.shape[1] % n)
|
||||
|
||||
if new_h == 0 or new_w == 0:
|
||||
# when the width or height is less than n, no need to process
|
||||
return img
|
||||
|
||||
if new_h != img.shape[0] or new_w != img.shape[1]:
|
||||
img = img[:new_h, :new_w]
|
||||
|
||||
return img
|
||||
|
||||
|
||||
def load_img_online(obj, mode="bgr", **kwargs):
|
||||
max_dim = kwargs.get("max_dim", 1920)
|
||||
n = kwargs.get("n", 2)
|
||||
if isinstance(obj, str):
|
||||
if mode.lower() == "gray":
|
||||
img = cv2.imread(obj, cv2.IMREAD_GRAYSCALE)
|
||||
else:
|
||||
img = cv2.imread(obj, cv2.IMREAD_COLOR)
|
||||
else:
|
||||
img = obj
|
||||
|
||||
# Resize image to satisfy constraints
|
||||
img = _resize_to_limit(img, max_dim=max_dim, n=n)
|
||||
|
||||
if mode.lower() == "bgr":
|
||||
return contiguous(img)
|
||||
elif mode.lower() == "rgb":
|
||||
return contiguous(img[..., ::-1])
|
||||
else:
|
||||
raise Exception(f"Unknown mode {mode}")
|
||||
@@ -1,19 +1,13 @@
|
||||
# coding: utf-8
|
||||
|
||||
import os.path as osp
|
||||
import cv2; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False)
|
||||
import cv2#; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False)
|
||||
import torch
|
||||
import numpy as np
|
||||
import onnxruntime
|
||||
|
||||
from .timer import Timer
|
||||
from .rprint import rlog
|
||||
from .crop import crop_image, _transform_pts
|
||||
|
||||
|
||||
def make_abs_path(fn):
|
||||
return osp.join(osp.dirname(osp.realpath(__file__)), fn)
|
||||
|
||||
|
||||
import folder_paths
|
||||
import os
|
||||
def to_ndarray(obj):
|
||||
if isinstance(obj, torch.Tensor):
|
||||
return obj.cpu().numpy()
|
||||
@@ -22,16 +16,17 @@ def to_ndarray(obj):
|
||||
else:
|
||||
return np.array(obj)
|
||||
|
||||
|
||||
class LandmarkRunner(object):
|
||||
"""landmark runner"""
|
||||
def __init__(self, **kwargs):
|
||||
ckpt_path = kwargs.get('ckpt_path')
|
||||
onnx_provider = kwargs.get('onnx_provider', 'cuda') # 默认用cuda
|
||||
onnx_provider = kwargs.get('onnx_provider', 'cuda')
|
||||
device_id = kwargs.get('device_id', 0)
|
||||
self.dsize = kwargs.get('dsize', 224)
|
||||
self.timer = Timer()
|
||||
|
||||
import onnxruntime
|
||||
|
||||
if onnx_provider.lower() == 'cuda':
|
||||
self.session = onnxruntime.InferenceSession(
|
||||
ckpt_path, providers=[
|
||||
@@ -40,7 +35,7 @@ class LandmarkRunner(object):
|
||||
)
|
||||
else:
|
||||
opts = onnxruntime.SessionOptions()
|
||||
opts.intra_op_num_threads = 4 # 默认线程数为 4
|
||||
opts.intra_op_num_threads = 4
|
||||
self.session = onnxruntime.InferenceSession(
|
||||
ckpt_path, providers=['CPUExecutionProvider'],
|
||||
sess_options=opts
|
||||
@@ -52,8 +47,7 @@ class LandmarkRunner(object):
|
||||
|
||||
def run(self, img_rgb: np.ndarray, lmk=None):
|
||||
if lmk is not None:
|
||||
crop_dct = crop_image(img_rgb, lmk, dsize=self.dsize, scale=1.5, vy_ratio=-0.1)
|
||||
img_crop_rgb = crop_dct['img_crop']
|
||||
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
|
||||
@@ -72,13 +66,12 @@ 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
|
||||
}
|
||||
|
||||
def warmup(self):
|
||||
# 构造dummy image进行warmup
|
||||
self.timer.tic()
|
||||
|
||||
dummy_image = np.zeros((1, 3, self.dsize, self.dsize), dtype=np.float32)
|
||||
@@ -86,4 +79,55 @@ class LandmarkRunner(object):
|
||||
_ = self._run(dummy_image)
|
||||
|
||||
elapse = self.timer.toc()
|
||||
rlog(f'LandmarkRunner warmup time: {elapse:.3f}s')
|
||||
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,16 +0,0 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
custom print and log functions
|
||||
"""
|
||||
|
||||
__all__ = ['rprint', 'rlog']
|
||||
|
||||
try:
|
||||
from rich.console import Console
|
||||
console = Console()
|
||||
rprint = console.print
|
||||
rlog = console.log
|
||||
except:
|
||||
rprint = print
|
||||
rlog = print
|
||||
@@ -1,139 +0,0 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
functions for processing video
|
||||
"""
|
||||
|
||||
import os.path as osp
|
||||
import numpy as np
|
||||
import subprocess
|
||||
import imageio
|
||||
import cv2
|
||||
|
||||
from rich.progress import track
|
||||
from .helper import prefix
|
||||
from .rprint import rprint as print
|
||||
|
||||
|
||||
def exec_cmd(cmd):
|
||||
subprocess.run(cmd, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
|
||||
|
||||
|
||||
def images2video(images, wfp, **kwargs):
|
||||
fps = kwargs.get('fps', 30)
|
||||
video_format = kwargs.get('format', 'mp4') # default is mp4 format
|
||||
codec = kwargs.get('codec', 'libx264') # default is libx264 encoding
|
||||
quality = kwargs.get('quality') # video quality
|
||||
pixelformat = kwargs.get('pixelformat', 'yuv420p') # video pixel format
|
||||
image_mode = kwargs.get('image_mode', 'rgb')
|
||||
macro_block_size = kwargs.get('macro_block_size', 2)
|
||||
ffmpeg_params = ['-crf', str(kwargs.get('crf', 18))]
|
||||
|
||||
writer = imageio.get_writer(
|
||||
wfp, fps=fps, format=video_format,
|
||||
codec=codec, quality=quality, ffmpeg_params=ffmpeg_params, pixelformat=pixelformat, macro_block_size=macro_block_size
|
||||
)
|
||||
|
||||
n = len(images)
|
||||
for i in track(range(n), description='writing', transient=True):
|
||||
if image_mode.lower() == 'bgr':
|
||||
writer.append_data(images[i][..., ::-1])
|
||||
else:
|
||||
writer.append_data(images[i])
|
||||
|
||||
writer.close()
|
||||
|
||||
# print(f':smiley: Dump to {wfp}\n', style="bold green")
|
||||
print(f'Dump to {wfp}\n')
|
||||
|
||||
|
||||
def video2gif(video_fp, fps=30, size=256):
|
||||
if osp.exists(video_fp):
|
||||
d = osp.split(video_fp)[0]
|
||||
fn = prefix(osp.basename(video_fp))
|
||||
palette_wfp = osp.join(d, 'palette.png')
|
||||
gif_wfp = osp.join(d, f'{fn}.gif')
|
||||
# generate the palette
|
||||
cmd = f'ffmpeg -i {video_fp} -vf "fps={fps},scale={size}:-1:flags=lanczos,palettegen" {palette_wfp} -y'
|
||||
exec_cmd(cmd)
|
||||
# use the palette to generate the gif
|
||||
cmd = f'ffmpeg -i {video_fp} -i {palette_wfp} -filter_complex "fps={fps},scale={size}:-1:flags=lanczos[x];[x][1:v]paletteuse" {gif_wfp} -y'
|
||||
exec_cmd(cmd)
|
||||
else:
|
||||
print(f'video_fp: {video_fp} not exists!')
|
||||
|
||||
|
||||
def merge_audio_video(video_fp, audio_fp, wfp):
|
||||
if osp.exists(video_fp) and osp.exists(audio_fp):
|
||||
cmd = f'ffmpeg -i {video_fp} -i {audio_fp} -c:v copy -c:a aac {wfp} -y'
|
||||
exec_cmd(cmd)
|
||||
print(f'merge {video_fp} and {audio_fp} to {wfp}')
|
||||
else:
|
||||
print(f'video_fp: {video_fp} or audio_fp: {audio_fp} not exists!')
|
||||
|
||||
|
||||
def blend(img: np.ndarray, mask: np.ndarray, background_color=(255, 255, 255)):
|
||||
mask_float = mask.astype(np.float32) / 255.
|
||||
background_color = np.array(background_color).reshape([1, 1, 3])
|
||||
bg = np.ones_like(img) * background_color
|
||||
img = np.clip(mask_float * img + (1 - mask_float) * bg, 0, 255).astype(np.uint8)
|
||||
return img
|
||||
|
||||
|
||||
def concat_frames(I_p_lst, driving_rgb_lst, img_rgb):
|
||||
# TODO: add more concat style, e.g., left-down corner driving
|
||||
out_lst = []
|
||||
for idx, _ in track(enumerate(I_p_lst), total=len(I_p_lst), description='Concatenating result...'):
|
||||
source_image_drived = I_p_lst[idx]
|
||||
image_drive = driving_rgb_lst[idx]
|
||||
|
||||
# resize images to match source_image_drived shape
|
||||
h, w, _ = source_image_drived.shape
|
||||
image_drive_resized = cv2.resize(image_drive, (w, h))
|
||||
img_rgb_resized = cv2.resize(img_rgb, (w, h))
|
||||
|
||||
# concatenate images horizontally
|
||||
frame = np.concatenate((image_drive_resized, img_rgb_resized, source_image_drived), axis=1)
|
||||
out_lst.append(frame)
|
||||
return out_lst
|
||||
|
||||
|
||||
class VideoWriter:
|
||||
def __init__(self, **kwargs):
|
||||
self.fps = kwargs.get('fps', 30)
|
||||
self.wfp = kwargs.get('wfp', 'video.mp4')
|
||||
self.video_format = kwargs.get('format', 'mp4')
|
||||
self.codec = kwargs.get('codec', 'libx264')
|
||||
self.quality = kwargs.get('quality')
|
||||
self.pixelformat = kwargs.get('pixelformat', 'yuv420p')
|
||||
self.image_mode = kwargs.get('image_mode', 'rgb')
|
||||
self.ffmpeg_params = kwargs.get('ffmpeg_params')
|
||||
|
||||
self.writer = imageio.get_writer(
|
||||
self.wfp, fps=self.fps, format=self.video_format,
|
||||
codec=self.codec, quality=self.quality,
|
||||
ffmpeg_params=self.ffmpeg_params, pixelformat=self.pixelformat
|
||||
)
|
||||
|
||||
def write(self, image):
|
||||
if self.image_mode.lower() == 'bgr':
|
||||
self.writer.append_data(image[..., ::-1])
|
||||
else:
|
||||
self.writer.append_data(image)
|
||||
|
||||
def close(self):
|
||||
if self.writer is not None:
|
||||
self.writer.close()
|
||||
|
||||
|
||||
def change_video_fps(input_file, output_file, fps=20, codec='libx264', crf=5):
|
||||
cmd = f"ffmpeg -i {input_file} -c:v {codec} -crf {crf} -r {fps} {output_file} -y"
|
||||
exec_cmd(cmd)
|
||||
|
||||
|
||||
def get_fps(filepath):
|
||||
import ffmpeg
|
||||
probe = ffmpeg.probe(filepath)
|
||||
video_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None)
|
||||
fps = eval(video_stream['avg_frame_rate'])
|
||||
return fps
|
||||
@@ -0,0 +1 @@
|
||||
from .mp_utils import LMKExtractor
|
||||
File diff suppressed because it is too large
Load Diff
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,38 @@
|
||||
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
|
||||
|
||||
+3
-3
@@ -1,8 +1,8 @@
|
||||
[project]
|
||||
name = "comfyui-liveportraitkj"
|
||||
description = "Nodes for [a/LivePortrait](https://github.com/KwaiVGI/LivePortrait), insightface is required"
|
||||
version = "1.0.0"
|
||||
license = "LICENSE"
|
||||
description = "Nodes for [a/LivePortrait](https://github.com/KwaiVGI/LivePortrait)"
|
||||
version = "1.1.0"
|
||||
license = { text = "MIT License" }
|
||||
dependencies = ["pyyaml", "numpy", "opencv-python", "rich"]
|
||||
|
||||
[project.urls]
|
||||
|
||||
@@ -1,15 +1,66 @@
|
||||
# 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
|
||||
|
||||
|
||||
|
||||
|
||||
https://github.com/kijai/ComfyUI-LivePortrait/assets/40791699/e55e10f6-af61-4d73-b162-af29eb847516
|
||||
# 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
|
||||
|
||||
|
||||
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
|
||||
|
||||
Insightface is also required.
|
||||
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:
|
||||
If you have a working compile environment, installing it can be as easy as:
|
||||
|
||||
`pip install insightface`
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
pyyaml
|
||||
numpy<=1.26.4
|
||||
opencv-python
|
||||
onnxruntime
|
||||
pykalman
|
||||
mediapipe
|
||||
onnx2torch
|
||||
+5
-3
@@ -1,5 +1,7 @@
|
||||
pyyaml
|
||||
numpy
|
||||
numpy<=1.26.4
|
||||
opencv-python
|
||||
rich
|
||||
onnxruntime-gpu
|
||||
onnxruntime-gpu
|
||||
pykalman
|
||||
mediapipe
|
||||
onnx2torch
|
||||
Reference in New Issue
Block a user