Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
 toto
|
fde067c9f5 | ||
|
toto
|
f233cae900 | ||
|
tototianhao
|
94dcf5b2d9 | ||
|
tototianhao
|
db728d3b70 | ||
|
toto
|
9a7cdbd756 | ||
|
toto
|
74dce3296d | ||
|
toto
|
e30f28d20e | ||
|
TaylorGoulding
|
baa3724ae2 | ||
|
TaylorGoulding
|
9453492f02 | ||
|
toto
|
b2d44fb139 | ||
|
toto
|
1d09d46f16 | ||
|
toto
|
bac4d69fba | ||
|
toto
|
0c06d3e3bf | ||
|
toto
|
7861f289c5 | ||
|
toto
|
fd5f4737fe | ||
|
toto
|
fb2a1e3283 | ||
|
toto
|
75cb3303a9 | ||
|
toto
|
2763361b55 | ||
|
toto
|
3f70c7eb6b | ||
|
toto
|
5beb93db0a | ||
|
toto
|
98d3b8ac13 | ||
|
toto
|
c871825137 | ||
|
toto
|
26245524ea | ||
|
toto
|
2292031e5f | ||
|
toto
|
017dcee81b | ||
|
toto
|
6d6eb679d5 | ||
|
toto
|
2abacd7fde | ||
|
toto
|
6a2d356ed7 | ||
|
toto
|
97c93eff3d | ||
|
toto
|
2598a6efe4 | ||
|
toto
|
dedd32d29e | ||
|
toto
|
a839966e9f | ||
|
toto
|
c8e3535869 | ||
|
tototianhao
|
359e72cd2e | ||
|
tototianhao
|
a11d54294f | ||
|
TaylorGoulding
|
6013eb426c | ||
|
TaylorGoulding
|
0b005db217 | ||
|
TaylorGoulding
|
658f057ccd | ||
|
TaylorGoulding
|
5b77e1cf3d | ||
|
toto
|
a6dbede699 | ||
|
toto
|
dc01281d85 | ||
|
toto
|
948b679e70 | ||
|
toto
|
6335f3c3d4 | ||
|
toto
|
3744780d0b | ||
|
tototianhao
|
425924bcbb | ||
|
toto
|
d364f2874a | ||
|
toto
|
e75c9a402a | ||
|
toto
|
e45b89d111 | ||
|
toto
|
3c61c01ba1 | ||
|
toto
|
a41dfa668c | ||
|
toto
|
e54e9ea579 | ||
|
toto
|
d9ade2e531 | ||
|
toto
|
4abaf4cd92 | ||
|
toto
|
60bebf7b02 | ||
|
toto
|
4874aa13ae | ||
|
toto
|
a9e70be153 | ||
|
toto
|
27582c9b83 | ||
|
toto
|
71e609657a | ||
|
toto
|
5fe61af9b9 | ||
|
toto
|
982613bd65 | ||
|
TaylorGoulding
|
8eb6c7713b | ||
|
toto
|
8964fc1222 | ||
|
toto
|
e53dff4b45 |
@@ -1,2 +1,3 @@
|
||||
__pycache__
|
||||
.idea
|
||||
*.DS_Store
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# [Portrait-Maker](https://github.com/THtianhao/ComfyUI-Portrait-Maker)
|
||||
This project is an adaptation of [EasyPhoto](https://github.com/aigc-apps/sd-webui-EasyPhoto), which breaks down the process of [EasyPhoto](https://github.com/aigc-apps/sd-webui-EasyPhoto) and will add a series of operations on human portraits in the future.
|
||||
|
||||
|
||||
|
||||
English | [简体中文](./README_zh-CN.md)
|
||||
|
||||
@@ -11,10 +11,23 @@ If you have any questions or suggestions, you can reach us through:
|
||||
- Email: tototianhao@gmail.com
|
||||
- telegram: https://t.me/+JoFE2vqHU4phZjg1
|
||||
- QQ Group: 10419777
|
||||
- WeChat Group: <img src="./images/wechat.jpg" width="200">
|
||||
- WeChat Group: <img src="./images/wechat.jpg" width="300">
|
||||
|
||||
## V1.2.0 Update
|
||||
1. Add PM_SuperColorTransfer node to simplify the color transfer process
|
||||
2. Add PM_SuperMakeUpTransfer node to simplify the process of makeup transfer
|
||||
3. Add v1.2.0 workflow
|
||||
|
||||
|
||||
## Recent Updates
|
||||
## V1.1.0 Update
|
||||
1. faceskin adds blur option
|
||||
2. Add PM_FaceShapMatch node. same as easyphot faceshap match
|
||||
3. Add PM_MakeUpTransfer node. same as easyphoto makeup transfer.
|
||||
4. Add a super-resolution model to the PM_PortraitEnhancement node. This super-resolution model can not highlight faces.
|
||||
5. Add v1.1.0 workflow
|
||||
6. RetinaFace supports face selection
|
||||
|
||||
## V1.0.0 Update
|
||||
1. Added log for model downloads.
|
||||
2. Renamed nodes to resolve conflicts with other plugins.
|
||||
3. Added "roop" model to the Facefusion PM node.
|
||||
@@ -28,10 +41,6 @@ If you have any questions or suggestions, you can reach us through:
|
||||
## Installation
|
||||
**Note: When you start the plugin for the first time, it will download all the models required by EasyPhoto. You can see the download progress in the terminal. Please do not interrupt the download (no hash verification for startup speed). If the download is interrupted, you need to manually delete the files downloaded halfway last time and download them again.**
|
||||
|
||||
### For Windows users
|
||||
|
||||
If you are using a project extracted from a zip package while using ComfyUI, you won't be able to use this plugin. This project relies on ModelScope, but the virtual environment provided in the official ComfyUI zip package cannot install ModelScope. Furthermore, the ComfyUI author has responded, stating that this issue cannot be resolved.[aliyunsdkcor error](https://github.com/ltdrdata/ComfyUI-Impact-Pack/issues/223) If Windows users wish to use this plugin for analyzing and composing ComfyUI workflows, they will need to create their own virtual environment. (I am using Python 3.10.6.).Of course, if you know a solution, feel free to submit a pull request (PR).
|
||||
|
||||
### install step
|
||||
1. First, install ComfyUI.
|
||||
|
||||
@@ -59,42 +68,68 @@ Click "Load" in the right panel of ComfyUI and select the ./workflow/easyphoto_w
|
||||
|
||||
## Node Introduction
|
||||
|
||||
* RetainFace PM: Processes images using the pipeline `damo/cv_resnet50_face-detection_retinaface` from Model Scope
|
||||
* image: Input image
|
||||
* multi_user_facecrop_ratio: Multiple for extracting the face area
|
||||
* FaceFusion PM: Fuses two face in the image using the pipeline `damo/cv_unet-image-face-fusion_damo` from Model Scope
|
||||
* image: Input image
|
||||
* user_image: Image to be fused
|
||||
* model: use ali model or roop model for fusion
|
||||
* RatioMerge2Image PM: Merges two images according to a ratio
|
||||
* image1: Input image
|
||||
* Image2: Input image
|
||||
* fusion_rate: Fusion ratio, maximum is 1, larger values lean towards image1
|
||||
* MaskMerge2Image PM: Merges images using a mask
|
||||
* image1: Input image
|
||||
* image2: Input image
|
||||
* mask: Mask to be replaced
|
||||
* ReplaceBoxImg PM: Replaces the image in a box area
|
||||
* origin_image: Original image
|
||||
* box_area: Area
|
||||
* replace_image: Image to be replaced in the area (resolution of box_area and replace_image must match)
|
||||
* ExpandMaskFaceWidth PM: Proportionally expands the width of the mask
|
||||
* mask: Input mask
|
||||
* box: Box corresponding to the mask
|
||||
* expand_width: Width expansion ratio based on the width of the box
|
||||
* BoxCropImage PM: Crops images using a box
|
||||
* ColorTransfer PM: Color transfer for images
|
||||
* FaceSkin PM: Extracts the mask of the facial part of the image
|
||||
* MaskDilateErode PM: Dilates and erodes the mask
|
||||
* SkinRetouching PM: Processes images using the pipeline `damo/cv_gpen_image-portrait-enhancement` from Model Scope
|
||||
* PortraitEnhancement PM: Processes images using the pipeline `damo/cv_gpen_image-portrait-enhancement` from Model Scope
|
||||
* ImageResizeTarget PM: Resizes images to a target width and height
|
||||
* ImageScaleShort PM: Reduces the width and height of the image's shorter side
|
||||
* image: Input image
|
||||
* size: Length to be scaled (proportionally scaled based on the shorter side of width and height)
|
||||
* crop_face: Width and height must be multiples of 32 after scaling
|
||||
* GetImageInfo PM: Extracts the width and height of the image
|
||||
* RetainFace PM: Perform matting using models from Model Scope. [Link](https://www.modelscope.cn/models/damo/cv_resnet50_face-detection_retinaface/summary)
|
||||
* image: Input image
|
||||
* multi_user_facecrop_ratio: Multiplicative factor for extracting the head region.
|
||||
* face_index : Choose which face
|
||||
|
||||
* FaceFusion PM: Merge faces from two images.
|
||||
* image: Input image
|
||||
* user_image: The image with the face to be merged.
|
||||
* model: Choose between Ali's model or Roop's model for merging.
|
||||
* ali: [Link](https://www.modelscope.cn/models/damo/cv_unet-image-face-fusion_damo/summary)
|
||||
* roop: [Link](https://github.com/deepinsight/insightface)
|
||||
|
||||
* RatioMerge2Image PM: Merge two images according to a specified ratio.
|
||||
* image1: First input image
|
||||
* image2: Second input image
|
||||
* fusion_rate: Fusion ratio, ranging from 0 to 1, where higher values favor image1.
|
||||
|
||||
* MaskMerge2Image PM: Merge images using a mask.
|
||||
* image1: First input image
|
||||
* image2: Second input image
|
||||
* mask: The mask to be applied for replacement.
|
||||
|
||||
* ReplaceBoxImg PM: Replace the image inside a specified box area.
|
||||
* origin_image: The original image
|
||||
* box_area: The area to be replaced
|
||||
* replace_image: The image to replace (ensure the resolution matches box_area)
|
||||
|
||||
* ExpandMaskFaceWidth PM: Proportionally expand the width of the mask.
|
||||
* mask: Input mask
|
||||
* box: Corresponding box of the mask
|
||||
* expand_width: The width expansion ratio, based on the box's width.
|
||||
|
||||
* BoxCropImage PM: Crop an image using a box.
|
||||
|
||||
* ColorTransfer PM: Perform color transfer on images.
|
||||
|
||||
* FaceSkin PM: Extract the mask of the facial region from an image.
|
||||
|
||||
* MaskDilateErode PM: Dilate and erode masks.
|
||||
|
||||
* Skin Retouching PM: Apply skin retouching using the following model.
|
||||
* [Link](https://www.modelscope.cn/models/damo/cv_unet_skin-retouching/summary)
|
||||
|
||||
* Portrait Enhancement PM: Process images using the following model.
|
||||
* model
|
||||
* gpen: [Link](https://www.modelscope.cn/models/damo/cv_gpen_image-portrait-enhancement/summary)
|
||||
* real_gan: [Link](https://www.modelscope.cn/models/bubbliiiing/cv_rrdb_image-super-resolution_x2/summary)
|
||||
|
||||
* ImageResizeTarget PM: Resize images to a target width and height.
|
||||
|
||||
* ImageScaleShort PM: Reduce the smaller dimension of an image proportionally.
|
||||
* image: Input image
|
||||
* size: Desired length for resizing (maintains the aspect ratio)
|
||||
* crop_face: Ensure the resulting width and height are multiples of 32.
|
||||
|
||||
* GetImageInfo PM: Extract the width and height of an image.
|
||||
|
||||
* Face Shape Match PM: Apply a certain level of fusion between the diffused image and the original image to reduce differences around the face.
|
||||
|
||||
* Makeup Transfer PM: Use a GAN network model to perform makeup transfer.
|
||||
* SuperMakeUpTransfer PM:(Multi-node integration) makeup by merging two pictures
|
||||
* SuperColorTransfer PM：(Multi-node integration) transfer the colors of two pictures
|
||||
## Contribution
|
||||
|
||||
If you find any issues or have suggestions for improvement, feel free to contribute. Follow these steps:
|
||||
|
||||
+30
-11
@@ -2,7 +2,7 @@
|
||||
|
||||
这个项目改编于[EasyPhoto](https://github.com/aigc-apps/sd-webui-EasyPhoto)，对于[EasyPhoto](https://github.com/aigc-apps/sd-webui-EasyPhoto)进行了流程上的拆解，后续会加入其他项目处理人物头像上的系列操作。
|
||||
|
||||
|
||||
|
||||
|
||||
English | [简体中文](./README_zh-CN.md)
|
||||
|
||||
@@ -13,10 +13,22 @@ English | [简体中文](./README_zh-CN.md)
|
||||
- 电子邮件：tototianhao@gmail.com
|
||||
- telegram: https://t.me/+JoFE2vqHU4phZjg1
|
||||
- QQ 群：10419777
|
||||
- 微信群： <img src="./images/wechat.jpg" width="200">
|
||||
- 微信群： <img src="./images/wechat.jpg" width="300">
|
||||
|
||||
## V1.2.0 Update
|
||||
1. 增加PM_SuperColorTransfer 节点，简化了颜色迁移的流程
|
||||
2. 增加PM_SuperMakeUpTransfer 节点，简化了进行装扮迁移的流程
|
||||
3. 增加v1.2.0 workflow
|
||||
|
||||
## 近期更新
|
||||
## v1.1.0 更新
|
||||
1. faceskin 增加模糊选项
|
||||
2. 增加 PM_FaceShapMatch节点 与easyphoto的FaceshapMatch一致
|
||||
3. 增加 PM_MakeUpTransfer节点 与easyphoto的MakeupTransfer一致
|
||||
4. PM_PortraitEnhancement节点增加一种超分模型，此超分模型可以对人脸不做高光
|
||||
5. 增加v1.1.0 workflow
|
||||
6. RetinaFace 支持选择人脸
|
||||
|
||||
## v1.0.0 更新
|
||||
|
||||
1. 增加模型下载的log
|
||||
2. 节点重命名解决与其他插件冲突问题
|
||||
@@ -32,10 +44,6 @@ English | [简体中文](./README_zh-CN.md)
|
||||
|
||||
**注意：初次启动插件的时候会下载EasyPhoto所需要的所有模型，在terminal中可以看到下载进度，请不要中断下载，(为了启动速度，没有做hash校验)，如果中断下载，需要手动删除上次下载一半的文件，重新下载。**
|
||||
|
||||
### windows用户
|
||||
|
||||
如果在使用ComfyUI的时候使用zip包解压后的项目，是无法使用本插件的，本项目依赖modelscope，但是ComfyUI官方zip包中的虚拟环境无法安装modelscope，并且ComfyUI作者已经回复了表示无法解决此问题[aliyunsdkcor error](https://github.com/ltdrdata/ComfyUI-Impact-Pack/issues/223)如果windows用户想使用本插件来分析、组合ComfyUI的流程，请自己创建虚拟环境。(我使用的是python3.10.6)，当然如果您知道解决的方法，欢迎提交pr
|
||||
|
||||
### 步骤
|
||||
1. 首先安装ComfyUI
|
||||
|
||||
@@ -62,13 +70,16 @@ Easyphoto工作位置: [./workflow/easyphoto.json](./workflows/easyphoto.json )
|
||||
|
||||
## 节点介绍
|
||||
|
||||
* RetainFace PM:使用Model Scope中的pipleline `damo/cv_resnet50_face-detection_retinaface`处理图像
|
||||
* RetainFace PM:使用Model Scope中的模型进行抠图 [链接](https://www.modelscope.cn/models/damo/cv_resnet50_face-detection_retinaface/summary)
|
||||
* image：输入图像
|
||||
* multi_user_facecrop_ratio：提取头像区域的倍数
|
||||
* FaceFusion PM:使用Model Scope中的pipleline `damo/cv_unet-image-face-fusion_damo`将两张图像的人脸进行融合
|
||||
* face_index : 选择第几个人脸
|
||||
* FaceFusion PM:将两张图像的人脸进行融合
|
||||
* image：输入图像
|
||||
* user_image：要融合的头像
|
||||
* model: 使用ali的模型还是roop模型进行融合
|
||||
* ali：[链接](https://www.modelscope.cn/models/damo/cv_unet-image-face-fusion_damo/summary)
|
||||
* roop: [链接](https://github.com/deepinsight/insightface)
|
||||
* RatioMerge2Image PM: 按照比例融合两张图片
|
||||
* image1：输入的图像
|
||||
* Image2：输入的图像
|
||||
@@ -89,14 +100,22 @@ Easyphoto工作位置: [./workflow/easyphoto.json](./workflows/easyphoto.json )
|
||||
* ColorTransfer PM:对图片进行颜色迁移
|
||||
* FaceSkin PM:提取图片中人脸的部分的Mask
|
||||
* MaskDilateErode PM: 对Mask进行膨胀与腐蚀
|
||||
* SkinRetouching PM:使用Model Scope中的pipleline `damo/cv_gpen_image-portrait-enhancement`处理图像
|
||||
* PortraitEnhancement PM:使用Model Scope中的pipleline `damo/cv_gpen_image-portrait-enhancement`处理图像
|
||||
* SkinRetouching PM:使用以下模型进行皮肤美化
|
||||
* [链接](https://www.modelscope.cn/models/damo/cv_unet_skin-retouching/summary)
|
||||
* PortraitEnhancement PM:使用以下模型处理图像
|
||||
* model
|
||||
* gpen : [链接](https://www.modelscope.cn/models/damo/cv_gpen_image-portrait-enhancement/summary)
|
||||
* real_gan:[链接](https://www.modelscope.cn/models/bubbliiiing/cv_rrdb_image-super-resolution_x2/summary)
|
||||
* ImageResizeTarget PM:将图片缩放到目标宽高
|
||||
* ImageScaleShort PM: 将图片的宽高中小的部分缩减到
|
||||
* image：输入图像
|
||||
* size：要缩放的长度（按照宽高中最短的一边进行比例缩放）
|
||||
* crop_face：缩放后宽高要以32为倍数
|
||||
* GetImageInfo PM: 提取图片的宽高
|
||||
* FaceShapMatchPM: 扩散后的图片和原图片进行一定的融合，减少脸旁边的差异
|
||||
* MakeUpTransferPM: 使用gan网络模型对妆容进行一定的迁移
|
||||
* SuperMakeUpTransferPM:(多节点的整合)融合两张图片的装扮
|
||||
* SuperColorTransferPM：(多节点的整合)迁移两张图片的颜色
|
||||
|
||||
## 贡献
|
||||
|
||||
|
||||
+14
-29
@@ -1,11 +1,9 @@
|
||||
import sys
|
||||
import os
|
||||
import os, sys
|
||||
|
||||
main_path = os.path.dirname(__file__)
|
||||
sys.path.append(main_path)
|
||||
|
||||
import subprocess
|
||||
import threading
|
||||
|
||||
import portrait.install
|
||||
import requests
|
||||
from tqdm import tqdm
|
||||
from portrait.nodes import *
|
||||
@@ -13,30 +11,6 @@ from portrait.nodes import *
|
||||
# import pydevd_pycharm
|
||||
# pydevd_pycharm.settrace('49.7.62.197', port=10090, stdoutToServer=True, stderrToServer=True)
|
||||
|
||||
|
||||
def handle_stream(stream, prefix):
|
||||
for line in stream:
|
||||
print(prefix, line, end="")
|
||||
|
||||
def run_script(cmd, cwd='.'):
|
||||
process = subprocess.Popen(cmd, cwd=cwd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1)
|
||||
|
||||
stdout_thread = threading.Thread(target=handle_stream, args=(process.stdout, ""))
|
||||
stderr_thread = threading.Thread(target=handle_stream, args=(process.stderr, "[!]"))
|
||||
|
||||
stdout_thread.start()
|
||||
stderr_thread.start()
|
||||
|
||||
stdout_thread.join()
|
||||
stderr_thread.join()
|
||||
|
||||
return process.wait()
|
||||
|
||||
print("## installing dependencies")
|
||||
|
||||
requirements_path = os.path.join(main_path, "requirements.txt")
|
||||
run_script([sys.executable, '-s', '-m', 'pip', 'install', '-q', '-r', requirements_path])
|
||||
|
||||
def urldownload_progressbar(url, file_path):
|
||||
response = requests.get(url, stream=True)
|
||||
total_size = int(response.headers.get('content-length', 0))
|
||||
@@ -74,6 +48,11 @@ NODE_CLASS_MAPPINGS = {
|
||||
"PM_ImageScaleShort": ImageScaleShortPM,
|
||||
"PM_ImageResizeTarget": ImageResizeTargetPM,
|
||||
"PM_GetImageInfo": GetImageInfoPM,
|
||||
"PM_MakeUpTransfer": MakeUpTransferPM,
|
||||
"PM_FaceShapMatch": FaceShapMatchPM,
|
||||
"PM_SuperColorTransfer": SuperColorTransferPM,
|
||||
"PM_SuperMakeUpTransfer": SuperMakeUpTransferPM,
|
||||
"PM_Similarity": SimilarityPM,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PM_RetinaFace": "RetinaFace PM",
|
||||
@@ -91,6 +70,12 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PM_ImageScaleShort": "ImageScaleShort PM",
|
||||
"PM_ImageResizeTarget": "ImageResizeTarget PM",
|
||||
"PM_GetImageInfo": "GetImageInfo PM",
|
||||
"PM_MakeUpTransfer": "MakeUpTransfer PM",
|
||||
"PM_FaceShapMatch": "FaceShapMatch PM",
|
||||
"PM_SuperColorTransfer": "SuperColorTransfer PM",
|
||||
"PM_SuperMakeUpTransfer": "SuperMakeUpTransfer PM",
|
||||
"PM_Similarity": "Similarity PM",
|
||||
|
||||
}
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 156 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 754 KiB |
Regular → Executable
BIN
Binary file not shown.
|
Before Width: | Height: | Size: 105 KiB After Width: | Height: | Size: 679 KiB |
Binary file not shown.
+9
-3
@@ -6,7 +6,7 @@ utils_path = os.path.join(root_path, "utils")
|
||||
models_path = os.path.join(root_path, "models")
|
||||
# save_dirs
|
||||
urls = [
|
||||
"https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/webui/ChilloutMix-ni-fp16.safetensors",
|
||||
# "https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/webui/ChilloutMix-ni-fp16.safetensors",
|
||||
"https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/webui/control_v11p_sd15_openpose.pth",
|
||||
"https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/webui/control_v11p_sd15_canny.pth",
|
||||
"https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/webui/control_v11f1e_sd15_tile.pth",
|
||||
@@ -17,10 +17,13 @@ urls = [
|
||||
"https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/webui/hand_pose_model.pth",
|
||||
"https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/webui/vae-ft-mse-840000-ema-pruned.ckpt",
|
||||
"https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/webui/face_skin.pth",
|
||||
"https://huggingface.co/ezioruan/inswapper_128.onnx/resolve/main/inswapper_128.onnx"
|
||||
"https://huggingface.co/ezioruan/inswapper_128.onnx/resolve/main/inswapper_128.onnx",
|
||||
"https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/webui/face_landmarks.pth",
|
||||
"https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/webui/makeup_transfer.pth",
|
||||
|
||||
]
|
||||
filenames = [
|
||||
os.path.join(folder_names_and_paths['checkpoints'][0][0], "Chilloutmix-Ni-pruned-fp16-fix.safetensors"),
|
||||
# os.path.join(folder_names_and_paths['checkpoints'][0][0], "Chilloutmix-Ni-pruned-fp16-fix.safetensors"),
|
||||
os.path.join(folder_names_and_paths['controlnet'][0][0], "control_v11p_sd15_openpose.pth"),
|
||||
os.path.join(folder_names_and_paths['controlnet'][0][0], "control_v11p_sd15_canny.pth"),
|
||||
os.path.join(folder_names_and_paths['controlnet'][0][0], "control_v11f1e_sd15_tile.pth"),
|
||||
@@ -32,6 +35,9 @@ filenames = [
|
||||
os.path.join(folder_names_and_paths['vae'][0][0], "vae-ft-mse-840000-ema-pruned.ckpt"),
|
||||
os.path.join(models_path, "face_skin.pth"),
|
||||
os.path.join(models_path, "inswapper_128.onnx"),
|
||||
os.path.join(models_path, "face_landmarks.pth"),
|
||||
os.path.join(models_path, "makeup_transfer.pth"),
|
||||
|
||||
]
|
||||
# prompts
|
||||
validation_prompt = "easyphoto_face, easyphoto, 1person"
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
import os
|
||||
import sys
|
||||
import subprocess
|
||||
import threading
|
||||
import locale
|
||||
import traceback
|
||||
import re
|
||||
|
||||
from .config import root_path
|
||||
|
||||
plugin_name = os.path.basename(root_path)
|
||||
windows_not_install = ['mmcv_full\n']
|
||||
|
||||
def log(msg, end=None, file=None):
|
||||
print(f'{plugin_name} :', msg, end=end, file=file)
|
||||
|
||||
def handle_stream(stream, is_stdout):
|
||||
stream.reconfigure(encoding=locale.getpreferredencoding(), errors='replace')
|
||||
|
||||
for msg in stream:
|
||||
if is_stdout:
|
||||
log(msg, end="", file=sys.stdout)
|
||||
else:
|
||||
log(msg, end="", file=sys.stderr)
|
||||
|
||||
def process_wrap(cmd_str, cwd=None, handler=None):
|
||||
log(f"EXECUTE: {cmd_str} in '{cwd}'")
|
||||
process = subprocess.Popen(cmd_str, cwd=cwd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1)
|
||||
|
||||
if handler is None:
|
||||
handler = handle_stream
|
||||
|
||||
stdout_thread = threading.Thread(target=handler, args=(process.stdout, True))
|
||||
stderr_thread = threading.Thread(target=handler, args=(process.stderr, False))
|
||||
|
||||
stdout_thread.start()
|
||||
stderr_thread.start()
|
||||
|
||||
stdout_thread.join()
|
||||
stderr_thread.join()
|
||||
|
||||
return process.wait()
|
||||
|
||||
# ---
|
||||
pip_list = None
|
||||
|
||||
def get_installed_packages():
|
||||
global pip_list
|
||||
if pip_list is None:
|
||||
try:
|
||||
result = subprocess.check_output([sys.executable, '-m', 'pip', 'list'], universal_newlines=True)
|
||||
pip_list = set([line.split()[0].lower() for line in result.split('\n') if line.strip()])
|
||||
except subprocess.CalledProcessError as e:
|
||||
log(f"Failed to retrieve the information of installed pip packages.")
|
||||
return set()
|
||||
|
||||
return pip_list
|
||||
|
||||
def mmcv_install():
|
||||
process_wrap(pip_install + ['-U', 'openmim'], cwd=root_path)
|
||||
process_wrap(mim_install + ['mmcv-full'], cwd=root_path)
|
||||
pass
|
||||
|
||||
def is_installed(name):
|
||||
name = name.strip()
|
||||
pattern = r'([^<>!=]+)([<>!=]=?)'
|
||||
match = re.search(pattern, name)
|
||||
|
||||
if match:
|
||||
name = match.group(1)
|
||||
|
||||
result = name.lower() in get_installed_packages()
|
||||
return result
|
||||
|
||||
def check_and_install_requirements(file_path):
|
||||
log(file_path)
|
||||
version = sys.version_info[:2]
|
||||
if os.path.exists(file_path):
|
||||
with open(file_path, 'r') as file:
|
||||
lines = file.readlines()
|
||||
for line in lines:
|
||||
log(line)
|
||||
if not is_installed(line):
|
||||
if platform.system() == "Windows" and version[1] == 11 and 'insightface' in line:
|
||||
process_wrap(pip_install + ['insightface-0.7.3-cp311-cp311-win_amd64.whl'], cwd=root_path)
|
||||
continue
|
||||
if platform.system() == "Windows" and line in windows_not_install:
|
||||
log(f"windows skip {line}")
|
||||
continue
|
||||
log(f"install {line}")
|
||||
process_wrap(pip_install + [line], cwd=root_path)
|
||||
return False
|
||||
return True
|
||||
|
||||
try:
|
||||
import platform
|
||||
|
||||
log("### : Check dependencies")
|
||||
if "python_embed" in sys.executable or "python_embedded" in sys.executable:
|
||||
pip_install = [sys.executable, '-s', '-m', 'pip', 'install', '-q']
|
||||
mim_install = [sys.executable, '-s', '-m', 'mim', 'install', '-q']
|
||||
else:
|
||||
pip_install = [sys.executable, '-m', 'pip', 'install', '-q']
|
||||
mim_install = [sys.executable, '-m', 'mim', 'install', '-q']
|
||||
|
||||
subpack_req = os.path.join(root_path, "requirements.txt")
|
||||
# mmcv_install()
|
||||
check_and_install_requirements(subpack_req)
|
||||
if sys.argv[0] == 'install.py':
|
||||
sys.path.append('..') # for portable version
|
||||
|
||||
except Exception as e:
|
||||
log("Dependency installation has failed. Please install manually.")
|
||||
traceback.print_exc()
|
||||
@@ -3,6 +3,7 @@ from modelscope.utils.constant import Tasks
|
||||
import insightface
|
||||
from insightface.app import FaceAnalysis
|
||||
from .utils.face_process_utils import Face_Skin
|
||||
from .utils.psgan_utils import PSGAN_Inference
|
||||
|
||||
from .config import *
|
||||
|
||||
@@ -13,6 +14,9 @@ face_skin = None
|
||||
roop = None
|
||||
skin_retouching = None
|
||||
portrait_enhancement = None
|
||||
psgan_interface = None
|
||||
real_gan_sr = None
|
||||
face_recognition = None
|
||||
|
||||
def get_retinaface_detection():
|
||||
global retinaface_detection
|
||||
@@ -55,3 +59,24 @@ def get_portrait_enhancement():
|
||||
if portrait_enhancement is None:
|
||||
portrait_enhancement = pipeline(Tasks.image_portrait_enhancement, model='damo/cv_gpen_image-portrait-enhancement', model_revision='v1.0.0')
|
||||
return portrait_enhancement
|
||||
|
||||
def get_real_gan_sr():
|
||||
global real_gan_sr
|
||||
if real_gan_sr is None:
|
||||
real_gan_sr = pipeline('image-super-resolution-x2', model='bubbliiiing/cv_rrdb_image-super-resolution_x2', model_revision="v1.0.2")
|
||||
return real_gan_sr
|
||||
|
||||
def get_pagan_interface():
|
||||
global psgan_interface
|
||||
if psgan_interface is None:
|
||||
face_landmarks_model_path = os.path.join(models_path, "face_landmarks.pth")
|
||||
makeup_transfer_model_path = os.path.join(models_path, "makeup_transfer.pth")
|
||||
psgan_interface = PSGAN_Inference("cuda", makeup_transfer_model_path, get_retinaface_detection(), get_face_skin(), face_landmarks_model_path)
|
||||
return psgan_interface
|
||||
|
||||
def get_face_recognition():
|
||||
global face_recognition
|
||||
if face_recognition is None:
|
||||
face_recognition = pipeline("face_recognition", model="bubbliiiing/cv_retinafce_recognition", model_revision="v1.0.3")
|
||||
return face_recognition
|
||||
|
||||
|
||||
+238
-33
@@ -3,17 +3,19 @@ import numpy as np
|
||||
from PIL import Image
|
||||
from modelscope.outputs import OutputKeys
|
||||
from .utils.face_process_utils import call_face_crop, color_transfer, Face_Skin
|
||||
from .utils.img_utils import img_to_tensor, tensor_to_img, tensor_to_np, np_to_tensor, np_to_mask, img_to_mask
|
||||
from .utils.img_utils import img_to_tensor, tensor_to_img, tensor_to_np, np_to_tensor, np_to_mask, img_to_mask, img_to_np
|
||||
from .model_holder import *
|
||||
|
||||
# import pydevd_pycharm
|
||||
#
|
||||
# pydevd_pycharm.settrace('49.7.62.197', port=10090, stdoutToServer=True, stderrToServer=True)
|
||||
|
||||
class RetinaFacePM:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {"image": ("IMAGE",),
|
||||
"multi_user_facecrop_ratio": ("FLOAT", {"default": 1, "min": 0, "max": 10, "step": 0.1})
|
||||
"multi_user_facecrop_ratio": ("FLOAT", {"default": 1, "min": 0, "max": 10, "step": 0.01}),
|
||||
"face_index": ("INT", {"default": 0, "min": 0, "max": 10, "step": 1})
|
||||
}}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK", "BOX")
|
||||
@@ -21,12 +23,14 @@ class RetinaFacePM:
|
||||
FUNCTION = "retain_face"
|
||||
CATEGORY = "protrait/model"
|
||||
|
||||
def retain_face(self, image, multi_user_facecrop_ratio):
|
||||
def retain_face(self, image, multi_user_facecrop_ratio, face_index):
|
||||
np_image = np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
|
||||
image = Image.fromarray(np_image)
|
||||
retinaface_boxes, retinaface_keypoints, retinaface_masks, retinaface_tensor = call_face_crop(get_retinaface_detection(), image, multi_user_facecrop_ratio)
|
||||
crop_image = image.crop(retinaface_boxes[0])
|
||||
return (img_to_tensor(crop_image), retinaface_tensor, retinaface_boxes[0])
|
||||
retinaface_boxes, retinaface_keypoints, retinaface_masks, retinaface_mask_nps = call_face_crop(get_retinaface_detection(), image, multi_user_facecrop_ratio)
|
||||
crop_image = image.crop(retinaface_boxes[face_index])
|
||||
retinaface_mask = np_to_mask(retinaface_mask_nps[face_index])
|
||||
retinaface_boxe = retinaface_boxes[face_index]
|
||||
return (img_to_tensor(crop_image), retinaface_mask, retinaface_boxe)
|
||||
|
||||
class FaceFusionPM:
|
||||
|
||||
@@ -42,21 +46,46 @@ class FaceFusionPM:
|
||||
|
||||
CATEGORY = "protrait/model"
|
||||
|
||||
def resize(self, tensor):
|
||||
image = tensor_to_img(tensor)
|
||||
short_side = max(image.width, image.height)
|
||||
resize = float(short_side / 640)
|
||||
new_size = (int(image.width // resize), int(image.height // resize))
|
||||
resize_image = image.resize(new_size, Image.Resampling.LANCZOS)
|
||||
return img_to_np(resize_image)
|
||||
|
||||
def img_face_fusion(self, source_image, swap_image, mode):
|
||||
if mode == "ali":
|
||||
source_image = tensor_to_img(source_image)
|
||||
swap_image = tensor_to_img(swap_image)
|
||||
fusion_image = get_image_face_fusion()(dict(template=source_image, user=swap_image))[
|
||||
source_image_pil = tensor_to_img(source_image)
|
||||
swap_image_pil = tensor_to_img(swap_image)
|
||||
fusion_image = get_image_face_fusion()(dict(template=source_image_pil, user=swap_image_pil))[
|
||||
OutputKeys.OUTPUT_IMG]
|
||||
result_image = Image.fromarray(cv2.cvtColor(fusion_image, cv2.COLOR_BGR2RGB))
|
||||
return (img_to_tensor(result_image),)
|
||||
else:
|
||||
width, height = source_image.shape[2], source_image.shape[1]
|
||||
need_resize = False
|
||||
source_np = tensor_to_np(source_image)
|
||||
swap_np = tensor_to_np(swap_image)
|
||||
if source_image.shape[2] > 640 or source_image.shape[1] > 640:
|
||||
source_np = self.resize(source_image)
|
||||
need_resize = True
|
||||
if swap_image.shape[2] > 640 or swap_image.shape[1] > 640:
|
||||
swap_np = self.resize(swap_image)
|
||||
get_face_analysis().prepare(ctx_id=0, det_size=(640, 640))
|
||||
source_image = tensor_to_np(source_image)
|
||||
faces = get_face_analysis().get(source_image)
|
||||
swap_image = tensor_to_np(swap_image)
|
||||
swap_face = get_face_analysis().get(swap_image)
|
||||
result_image = get_roop().get(source_image, faces[0], swap_face[0], paste_back=True)
|
||||
faces = get_face_analysis().get(source_np)
|
||||
swap_faces = get_face_analysis().get(swap_np)
|
||||
if len(faces) == 0:
|
||||
raise RuntimeError("No face was recognized in the source image / source image 没有识别到人脸")
|
||||
if len(swap_faces) == 0:
|
||||
print("No face was recognized in the swap faces / swap faces没有识别到人脸, 用原脸替换!!!!!!!!!")
|
||||
return (source_image,)
|
||||
result_image = get_roop().get(source_np, faces[0], swap_faces[0], paste_back=True)
|
||||
if need_resize:
|
||||
image = Image.fromarray(result_image)
|
||||
new_size = width, height
|
||||
result_image = image.resize(new_size, Image.Resampling.LANCZOS)
|
||||
result_image = img_to_np(result_image)
|
||||
return (np_to_tensor(result_image),)
|
||||
|
||||
class RatioMerge2ImagePM:
|
||||
@@ -73,7 +102,7 @@ class RatioMerge2ImagePM:
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "image_ratio_merge"
|
||||
|
||||
CATEGORY = "protrait/model"
|
||||
CATEGORY = "protrait/other"
|
||||
|
||||
def image_ratio_merge(self, image1, image2, fusion_rate):
|
||||
rate_fusion_image = image1 * (1 - fusion_rate) + image2 * fusion_rate
|
||||
@@ -132,7 +161,7 @@ class ExpandMaskFaceWidthPM:
|
||||
RETURN_TYPES = ("MASK", "BOX")
|
||||
FUNCTION = "expand_mask_face_width"
|
||||
|
||||
CATEGORY = "protrait/model"
|
||||
CATEGORY = "protrait/other"
|
||||
|
||||
def expand_mask_face_width(self, mask, box, expand_width):
|
||||
h, w = mask.shape[1], mask.shape[2]
|
||||
@@ -160,7 +189,7 @@ class BoxCropImagePM:
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("crop_image",)
|
||||
FUNCTION = "box_crop_image"
|
||||
CATEGORY = "protrait/model"
|
||||
CATEGORY = "protrait/other"
|
||||
|
||||
def box_crop_image(self, image, box):
|
||||
image = image[:, box[1]:box[3], box[0]:box[2], :]
|
||||
@@ -178,7 +207,7 @@ class ColorTransferPM:
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "color_transfer"
|
||||
|
||||
CATEGORY = "protrait/model"
|
||||
CATEGORY = "protrait/other"
|
||||
|
||||
def color_transfer(self, transfer_from, transfer_to):
|
||||
transfer_result = color_transfer(tensor_to_np(transfer_from), tensor_to_np(transfer_to)) # 进行颜色迁移
|
||||
@@ -189,17 +218,24 @@ class FaceSkinPM:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{"image": ("IMAGE",), }
|
||||
}
|
||||
{
|
||||
"image": ("IMAGE",),
|
||||
"blur_edge": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}),
|
||||
"blur_threshold": ("INT", {"default": 32, "min": 0, "max": 64, "step": 1}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "face_skin_mask"
|
||||
|
||||
CATEGORY = "protrait/model"
|
||||
|
||||
def face_skin_mask(self, image):
|
||||
face_skin_one = get_face_skin().detect(tensor_to_img(image), get_retinaface_detection(), [1, 2, 3, 4, 5, 10, 12, 13])
|
||||
return (face_skin_one,)
|
||||
def face_skin_mask(self, image, blur_edge, blur_threshold):
|
||||
face_skin_img = get_face_skin()(tensor_to_img(image), get_retinaface_detection(), [[1, 2, 3, 4, 5, 10, 12, 13]])[0]
|
||||
face_skin_np = img_to_np(face_skin_img)
|
||||
if blur_edge:
|
||||
face_skin_np = cv2.blur(face_skin_np, (blur_threshold, blur_threshold))
|
||||
return (np_to_mask(face_skin_np),)
|
||||
|
||||
class MaskDilateErodePM:
|
||||
|
||||
@@ -212,7 +248,7 @@ class MaskDilateErodePM:
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "mask_dilate_erode"
|
||||
|
||||
CATEGORY = "protrait/model"
|
||||
CATEGORY = "protrait/other"
|
||||
|
||||
def mask_dilate_erode(self, mask):
|
||||
out_mask = Image.fromarray(np.uint8(cv2.dilate(tensor_to_np(mask), np.ones((96, 96), np.uint8), iterations=1) - cv2.erode(tensor_to_np(mask), np.ones((48, 48), np.uint8), iterations=1)))
|
||||
@@ -236,20 +272,25 @@ class SkinRetouchingPM:
|
||||
|
||||
class PortraitEnhancementPM:
|
||||
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{"image": ("IMAGE",), }
|
||||
}
|
||||
{
|
||||
"image": ("IMAGE",),
|
||||
"model": (["pgen", "real_gan"],),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "protrait_enhancement_pass"
|
||||
|
||||
CATEGORY = "protrait/model"
|
||||
|
||||
def protrait_enhancement_pass(self, image):
|
||||
output_image = cv2.cvtColor(get_portrait_enhancement()(tensor_to_img(image))[OutputKeys.OUTPUT_IMG], cv2.COLOR_BGR2RGB)
|
||||
def protrait_enhancement_pass(self, image, model):
|
||||
if model == "pgen":
|
||||
output_image = cv2.cvtColor(get_portrait_enhancement()(tensor_to_img(image))[OutputKeys.OUTPUT_IMG], cv2.COLOR_BGR2RGB)
|
||||
elif model == "real_gan":
|
||||
output_image = cv2.cvtColor(get_real_gan_sr()(tensor_to_img(image))[OutputKeys.OUTPUT_IMG], cv2.COLOR_BGR2RGB)
|
||||
return (np_to_tensor(output_image),)
|
||||
|
||||
class ImageScaleShortPM:
|
||||
@@ -259,13 +300,13 @@ class ImageScaleShortPM:
|
||||
return {"required": {
|
||||
"image": ("IMAGE",),
|
||||
"size": ("INT", {"default": 512, "min": 0, "max": 2048, "step": 1}),
|
||||
"crop_face": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}),
|
||||
"crop_face": ("BOOLEAN", {"default": False}),
|
||||
}}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "image_scale_short"
|
||||
|
||||
CATEGORY = "protrait/model"
|
||||
CATEGORY = "protrait/other"
|
||||
|
||||
def image_scale_short(self, image, size, crop_face):
|
||||
input_image = tensor_to_img(image)
|
||||
@@ -292,7 +333,7 @@ class ImageResizeTargetPM:
|
||||
|
||||
FUNCTION = "image_resize_target"
|
||||
|
||||
CATEGORY = "protrait/model"
|
||||
CATEGORY = "protrait/other"
|
||||
|
||||
def image_resize_target(self, image, width, height):
|
||||
imagepi = tensor_to_img(image)
|
||||
@@ -311,9 +352,173 @@ class GetImageInfoPM:
|
||||
|
||||
FUNCTION = "get_image_info"
|
||||
|
||||
CATEGORY = "protrait/model"
|
||||
CATEGORY = "protrait/other"
|
||||
|
||||
def get_image_info(self, image):
|
||||
width = image.shape[2]
|
||||
height = image.shape[1]
|
||||
return (width, height)
|
||||
|
||||
class MakeUpTransferPM:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"source_image": ("IMAGE",),
|
||||
"makeup_image": ("IMAGE",),
|
||||
}}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
|
||||
FUNCTION = "makeup_transfer"
|
||||
|
||||
CATEGORY = "protrait/model"
|
||||
|
||||
def makeup_transfer(self, source_image, makeup_image):
|
||||
source_image = tensor_to_img(source_image)
|
||||
makeup_image = tensor_to_img(makeup_image)
|
||||
result = get_pagan_interface().transfer(source_image, makeup_image)
|
||||
return (img_to_tensor(result),)
|
||||
|
||||
class FaceShapMatchPM:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"source_image": ("IMAGE",),
|
||||
"match_image": ("IMAGE",),
|
||||
"face_box": ("BOX",),
|
||||
}}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
|
||||
FUNCTION = "faceshap_match"
|
||||
|
||||
CATEGORY = "protrait/model"
|
||||
|
||||
def faceshap_match(self, source_image, match_image, face_box):
|
||||
# detect face area
|
||||
source_image_copy = tensor_to_img(source_image)
|
||||
match_image_copy = tensor_to_img(match_image)
|
||||
face_skin_mask = get_face_skin()(source_image_copy, get_retinaface_detection(), needs_index=[[1, 2, 3, 4, 5, 7, 8, 10, 11, 12, 13]])[0]
|
||||
face_width = face_box[2] - face_box[0]
|
||||
kernel_size = np.ones((int(face_width // 10), int(face_width // 10)), np.uint8)
|
||||
|
||||
# Fill small holes with a close operation
|
||||
face_skin_mask = Image.fromarray(np.uint8(cv2.morphologyEx(np.array(face_skin_mask), cv2.MORPH_CLOSE, kernel_size)))
|
||||
|
||||
# Use dilate to reconstruct the surrounding area of the face
|
||||
face_skin_mask = Image.fromarray(np.uint8(cv2.dilate(np.array(face_skin_mask), kernel_size, iterations=1)))
|
||||
face_skin_mask = cv2.blur(np.float32(face_skin_mask), (32, 32)) / 255
|
||||
|
||||
# paste back to photo, Using I2I generation controlled solely by OpenPose, even with a very small denoise amplitude,
|
||||
# still carries the risk of introducing NSFW and global incoherence.!!! important!!!
|
||||
input_image_uint8 = np.array(source_image_copy) * face_skin_mask + np.array(match_image_copy) * (1 - face_skin_mask)
|
||||
|
||||
return (np_to_tensor(input_image_uint8),)
|
||||
|
||||
class SuperColorTransferPM:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return \
|
||||
{
|
||||
"required": {
|
||||
"main_image": ("IMAGE",),
|
||||
"transfer_image": ("IMAGE",),
|
||||
},
|
||||
"optional": {
|
||||
"avatar_box": ("BOX",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
|
||||
FUNCTION = "super_color_transfer"
|
||||
CATEGORY = "protrait/super"
|
||||
|
||||
def super_color_transfer(self, main_image, transfer_image, avatar_box=None):
|
||||
origin_np = tensor_to_np(main_image)
|
||||
result_np = None
|
||||
if avatar_box is not None:
|
||||
main_image = main_image[:, avatar_box[1]:avatar_box[3], avatar_box[0]:avatar_box[2], :]
|
||||
transfer_image = transfer_image[:, avatar_box[1]:avatar_box[3], avatar_box[0]:avatar_box[2], :]
|
||||
|
||||
transfer_result = color_transfer(tensor_to_np(main_image), tensor_to_np(transfer_image)) # 进行颜色迁移
|
||||
|
||||
face_skin_img = get_face_skin()(Image.fromarray(transfer_result), get_retinaface_detection(), [[1, 2, 3, 4, 5, 10, 12, 13]])[0]
|
||||
face_skin_np = img_to_np(face_skin_img)
|
||||
face_skin_np = cv2.blur(face_skin_np, (32, 32)) / 255
|
||||
|
||||
masked_img_np = tensor_to_np(main_image) * (1 - face_skin_np) + transfer_result * face_skin_np
|
||||
result_np = masked_img_np
|
||||
|
||||
if avatar_box is not None:
|
||||
origin_np[avatar_box[1]:avatar_box[3], avatar_box[0]:avatar_box[2], :] = masked_img_np
|
||||
result_np = origin_np
|
||||
|
||||
return (np_to_tensor(result_np),)
|
||||
|
||||
class SuperMakeUpTransferPM:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return \
|
||||
{
|
||||
"required": {
|
||||
"main_image": ("IMAGE",),
|
||||
"makeup_image": ("IMAGE",),
|
||||
},
|
||||
"optional": {
|
||||
"avatar_box": ("BOX",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "super_makeup_transfer"
|
||||
CATEGORY = "protrait/super"
|
||||
|
||||
def super_makeup_transfer(self, main_image, makeup_image, avatar_box=None):
|
||||
box_width, box_height = avatar_box[2] - avatar_box[0], avatar_box[3] - avatar_box[1]
|
||||
origin_np = tensor_to_np(main_image)
|
||||
if avatar_box is not None:
|
||||
main_image = main_image[:, avatar_box[1]:avatar_box[3], avatar_box[0]:avatar_box[2], :]
|
||||
makeup_image = makeup_image[:, avatar_box[1]:avatar_box[3], avatar_box[0]:avatar_box[2], :]
|
||||
resize_source_box_image = tensor_to_img(main_image).resize([256, 256])
|
||||
resize_makeup_box_image = tensor_to_img(makeup_image).resize([256, 256])
|
||||
transfer_image = get_pagan_interface().transfer(resize_source_box_image, resize_makeup_box_image)
|
||||
box_size_transfer = transfer_image.resize([box_width, box_height], Image.Resampling.LANCZOS)
|
||||
origin_np[avatar_box[1]:avatar_box[3], avatar_box[0]:avatar_box[2], :] = img_to_np(box_size_transfer)
|
||||
return (np_to_tensor(origin_np),)
|
||||
|
||||
class SimilarityPM:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return \
|
||||
{
|
||||
"required": {
|
||||
"main_image": ("IMAGE",),
|
||||
"compare_image": ("IMAGE",),
|
||||
"model": (["sim"],),
|
||||
"result_prefix": ("STRING", {"default": ""}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
|
||||
FUNCTION = "similarity_compare"
|
||||
CATEGORY = "protrait/model"
|
||||
|
||||
def similarity_compare(self, main_image, compare_image, model, result_prefix):
|
||||
main_image_copy = tensor_to_img(main_image)
|
||||
compare_image_copy = tensor_to_img(compare_image)
|
||||
score = None
|
||||
result = None
|
||||
if model == "sim":
|
||||
root_embedding = get_face_recognition()(dict(user=Image.fromarray(np.uint8(main_image_copy))))[OutputKeys.IMG_EMBEDDING]
|
||||
compare_embedding = get_face_recognition()(dict(user=Image.fromarray(np.uint8(compare_image_copy))))[OutputKeys.IMG_EMBEDDING]
|
||||
score = float(np.dot(root_embedding, np.transpose(compare_embedding))[0][0])
|
||||
if result_prefix == "":
|
||||
result = str(round(score, 2))
|
||||
else:
|
||||
result = f"{result_prefix}_{round(score, 2)}"
|
||||
return (result,)
|
||||
|
||||
@@ -65,9 +65,8 @@ def safe_get_box_mask_keypoints(image, retinaface_result, crop_ratio, face_seg,
|
||||
retinaface_boxs = [retinaface_boxs[index] for index in argindex]
|
||||
retinaface_keypoints = [retinaface_keypoints[index] for index in argindex]
|
||||
retinaface_mask_pils = [retinaface_mask_pils[index] for index in argindex]
|
||||
retinaface_mask_np = [retinaface_masks[index] for index in argindex]
|
||||
mask_tensor = np_to_mask(retinaface_mask_np[0])
|
||||
return retinaface_boxs, retinaface_keypoints, retinaface_mask_pils, mask_tensor
|
||||
retinaface_mask_nps = [retinaface_masks[index] for index in argindex]
|
||||
return retinaface_boxs, retinaface_keypoints, retinaface_mask_pils, retinaface_mask_nps
|
||||
|
||||
else:
|
||||
retinaface_box = np.array([])
|
||||
@@ -120,9 +119,9 @@ def call_face_crop(retinaface_detection, image, crop_ratio, prefix="tmp"):
|
||||
# retinaface detect
|
||||
retinaface_result = retinaface_detection(image)
|
||||
# get mask and keypoints
|
||||
retinaface_box, retinaface_keypoints, retinaface_mask_pil, retinaface_mask_tensor = safe_get_box_mask_keypoints(image, retinaface_result, crop_ratio, None, "crop")
|
||||
retinaface_box, retinaface_keypoints, retinaface_mask_pil, retinaface_mask_nps = safe_get_box_mask_keypoints(image, retinaface_result, crop_ratio, None, "crop")
|
||||
|
||||
return retinaface_box, retinaface_keypoints, retinaface_mask_pil, retinaface_mask_tensor
|
||||
return retinaface_box, retinaface_keypoints, retinaface_mask_pil, retinaface_mask_nps
|
||||
|
||||
def color_transfer(sc, dc):
|
||||
"""
|
||||
@@ -496,19 +495,27 @@ class Face_Skin(object):
|
||||
self.model.load_state_dict(torch.load(model_path, map_location='cpu'))
|
||||
self.model.eval()
|
||||
|
||||
self.cuda = torch.cuda.is_available()
|
||||
if self.cuda:
|
||||
self.model.cuda()
|
||||
|
||||
# transform for input image
|
||||
self.trans = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
|
||||
])
|
||||
|
||||
def detect(self, image, retinaface_detection, needs_index=[12, 13]):
|
||||
# index => label
|
||||
# 1:'skin', 2:'left_brow', 3:'right_brow', 4:'left_eye', 5:'right_eye', 6:'eye_g', 7:'left_ear', 8:'right_ear',
|
||||
# 9:'ear_r', 10:'nose', 11:'mouth', 12:'upper_lip', 13:'low_lip', 14:'neck', 15:'neck_l', 16:'cloth',
|
||||
# 17:'hair', 18:'hat'
|
||||
def __call__(self, image, retinaface_detection, needs_index=[[12, 13]]):
|
||||
# needs_index 12, 13 means seg the lip
|
||||
with torch.no_grad():
|
||||
total_mask = np.zeros_like(np.uint8(image))
|
||||
|
||||
# detect image
|
||||
retinaface_boxes, _, _, _ = call_face_crop(retinaface_detection, image, 13, prefix="tmp")
|
||||
retinaface_boxes, _, _, _ = call_face_crop(retinaface_detection, image, 1.5, prefix="tmp")
|
||||
retinaface_box = retinaface_boxes[0]
|
||||
|
||||
# sub_face for seg skin
|
||||
@@ -520,17 +527,21 @@ class Face_Skin(object):
|
||||
|
||||
torch_img = self.trans(PIL_img)
|
||||
torch_img = torch.unsqueeze(torch_img, 0)
|
||||
|
||||
if self.cuda:
|
||||
torch_img = torch_img.cuda()
|
||||
out = self.model(torch_img)[0]
|
||||
model_mask = out.squeeze(0).cpu().numpy().argmax(0)
|
||||
|
||||
sub_mask = np.zeros_like(model_mask)
|
||||
for index in needs_index:
|
||||
sub_mask += np.uint8(model_mask == index)
|
||||
masks = []
|
||||
for _needs_index in needs_index:
|
||||
total_mask = np.zeros_like(np.uint8(image))
|
||||
sub_mask = np.zeros_like(model_mask)
|
||||
for index in _needs_index:
|
||||
sub_mask += np.uint8(model_mask == index)
|
||||
|
||||
sub_mask = np.clip(sub_mask, 0, 1) * 255
|
||||
sub_mask = np.tile(np.expand_dims(cv2.resize(np.uint8(sub_mask), (image_w, image_h)), -1), [1, 1, 3])
|
||||
sub_mask = np.clip(sub_mask, 0, 1) * 255
|
||||
sub_mask = np.tile(np.expand_dims(cv2.resize(np.uint8(sub_mask), (image_w, image_h)), -1), [1, 1, 3])
|
||||
total_mask[retinaface_box[1]:retinaface_box[3], retinaface_box[0]:retinaface_box[2], :] = sub_mask
|
||||
masks.append(Image.fromarray(np.uint8(total_mask)))
|
||||
|
||||
# detect image
|
||||
total_mask[retinaface_box[1]:retinaface_box[3], retinaface_box[0]:retinaface_box[2], :] = sub_mask
|
||||
return np_to_mask(total_mask)
|
||||
return masks
|
||||
@@ -13,6 +13,12 @@ def img_to_tensor(input):
|
||||
tensor = torch.from_numpy(image)[None,]
|
||||
return tensor
|
||||
|
||||
def img_to_np(input):
|
||||
i = ImageOps.exif_transpose(input)
|
||||
image = i.convert("RGB")
|
||||
image_np = np.array(image).astype(np.float32)
|
||||
return image_np
|
||||
|
||||
def img_to_mask(input):
|
||||
i = ImageOps.exif_transpose(input)
|
||||
image = i.convert("RGB")
|
||||
|
||||
Executable
+884
@@ -0,0 +1,884 @@
|
||||
#!/usr/bin/python
|
||||
# -*- encoding: utf-8 -*-
|
||||
import math
|
||||
import os.path as osp
|
||||
from collections import OrderedDict
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from PIL import Image
|
||||
from torch import nn
|
||||
from torch.autograd import Variable
|
||||
from torch.nn import Parameter, functional
|
||||
from torchvision import transforms
|
||||
from torchvision.transforms import ToPILImage
|
||||
|
||||
pwd = osp.split(osp.realpath(__file__))[0]
|
||||
|
||||
# Preprocess part
|
||||
def to_var(x, requires_grad=True):
|
||||
if requires_grad:
|
||||
return Variable(x).float()
|
||||
else:
|
||||
return Variable(x, requires_grad=requires_grad).float()
|
||||
|
||||
def copy_area(tar, src, lms):
|
||||
rect = [int(min(lms[:, 1])) - PreProcess.eye_margin,
|
||||
int(min(lms[:, 0])) - PreProcess.eye_margin,
|
||||
int(max(lms[:, 1])) + PreProcess.eye_margin + 1,
|
||||
int(max(lms[:, 0])) + PreProcess.eye_margin + 1]
|
||||
tar[:, :, rect[1]:rect[3], rect[0]:rect[2]] = \
|
||||
src[:, :, rect[1]:rect[3], rect[0]:rect[2]]
|
||||
src[:, :, rect[1]:rect[3], rect[0]:rect[2]] = 0
|
||||
|
||||
class rectangle():
|
||||
def __init__(self, left, top, right, bottom):
|
||||
self.left_num = left
|
||||
self.top_num = top
|
||||
self.right_num = right
|
||||
self.bottom_num = bottom
|
||||
|
||||
def left(self):
|
||||
return self.left_num
|
||||
|
||||
def top(self):
|
||||
return self.top_num
|
||||
|
||||
def right(self):
|
||||
return self.right_num
|
||||
|
||||
def bottom(self):
|
||||
return self.bottom_num
|
||||
|
||||
def height(self):
|
||||
return self.bottom_num - self.top_num
|
||||
|
||||
def width(self):
|
||||
return self.right_num - self.left_num
|
||||
|
||||
def crop(image: Image, face, up_ratio, down_ratio, width_ratio) -> (Image, 'face'):
|
||||
width, height = image.size
|
||||
face_height = face.height()
|
||||
face_width = face.width()
|
||||
delta_up = up_ratio * face_height
|
||||
delta_down = down_ratio * face_height
|
||||
delta_width = width_ratio * width
|
||||
|
||||
img_left = int(max(0, face.left() - delta_width))
|
||||
img_top = int(max(0, face.top() - delta_up))
|
||||
img_right = int(min(width, face.right() + delta_width))
|
||||
img_bottom = int(min(height, face.bottom() + delta_down))
|
||||
image = image.crop((img_left, img_top, img_right, img_bottom))
|
||||
|
||||
face = rectangle(face.left() - img_left, face.top() - img_top,
|
||||
face.right() - img_left, face.bottom() - img_top)
|
||||
|
||||
center = [(img_right - img_left) / 2, (img_bottom - img_top) / 2]
|
||||
width, height = image.size
|
||||
# import ipdb; ipdb.set_trace()
|
||||
crop_left = img_left
|
||||
crop_top = img_top
|
||||
crop_right = img_right
|
||||
crop_bottom = img_bottom
|
||||
if width > height:
|
||||
left = int(center[0] - height / 2)
|
||||
right = int(center[0] + height / 2)
|
||||
if left < 0:
|
||||
left, right = 0, height
|
||||
elif right > width:
|
||||
left, right = width - height, width
|
||||
image = image.crop((left, 0, right, height))
|
||||
face = rectangle(face.left() - left, face.top(),
|
||||
face.right() - left, face.bottom())
|
||||
crop_left += left
|
||||
crop_right = crop_left + height
|
||||
elif width < height:
|
||||
top = int(center[1] - width / 2)
|
||||
bottom = int(center[1] + width / 2)
|
||||
if top < 0:
|
||||
top, bottom = 0, width
|
||||
elif bottom > height:
|
||||
top, bottom = height - width, height
|
||||
image = image.crop((0, top, width, bottom))
|
||||
face = rectangle(face.left(), face.top() - top,
|
||||
face.right(), face.bottom() - top)
|
||||
crop_top += top
|
||||
crop_bottom = crop_top + width
|
||||
crop_face = rectangle(crop_left, crop_top, crop_right, crop_bottom)
|
||||
return image, face, crop_face
|
||||
|
||||
class FaceParser:
|
||||
def __init__(self, device="cpu", face_skin=None):
|
||||
mapper = [0, 1, 2, 3, 4, 5, 0, 11, 12, 0, 6, 8, 7, 9, 13, 0, 0, 10, 0]
|
||||
self.device = device
|
||||
self.dic = torch.tensor(mapper, device=device)
|
||||
self.net = face_skin.model
|
||||
self.to_tensor = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
|
||||
])
|
||||
|
||||
def parse(self, image: Image):
|
||||
assert image.shape[:2] == (512, 512)
|
||||
with torch.no_grad():
|
||||
image = self.to_tensor(image).to(self.device)
|
||||
image = torch.unsqueeze(image, 0)
|
||||
out = self.net(image)[0]
|
||||
parsing = out.squeeze(0).argmax(0)
|
||||
mask = torch.zeros_like(parsing)
|
||||
for index, num in enumerate(self.dic):
|
||||
mask[parsing == index] = num
|
||||
return mask.float()
|
||||
|
||||
class WingLoss(nn.Module):
|
||||
|
||||
def __init__(self, wing_w=10.0, wing_epsilon=2.0):
|
||||
super(WingLoss, self).__init__()
|
||||
self.wing_w = wing_w
|
||||
self.wing_epsilon = wing_epsilon
|
||||
self.wing_c = self.wing_w * (1.0 - math.log(1.0 + self.wing_w / self.wing_epsilon))
|
||||
|
||||
def forward(self, targets, predictions, euler_angle_weights=None):
|
||||
abs_error = torch.abs(targets - predictions)
|
||||
loss = torch.where(torch.le(abs_error, self.wing_w),
|
||||
self.wing_w * torch.log(1.0 + abs_error / self.wing_epsilon), abs_error - self.wing_c)
|
||||
loss_sum = torch.sum(loss, 1)
|
||||
if euler_angle_weights is not None:
|
||||
loss_sum *= euler_angle_weights
|
||||
return torch.mean(loss_sum)
|
||||
|
||||
class LinearBottleneck(nn.Module):
|
||||
def __init__(self, input_channels, out_channels, expansion, stride=1, activation=nn.ReLU6):
|
||||
super(LinearBottleneck, self).__init__()
|
||||
self.expansion_channels = input_channels * expansion
|
||||
|
||||
self.conv1 = nn.Conv2d(input_channels, self.expansion_channels, stride=1, kernel_size=1)
|
||||
self.bn1 = nn.BatchNorm2d(self.expansion_channels)
|
||||
|
||||
self.depth_conv2 = nn.Conv2d(self.expansion_channels, self.expansion_channels, stride=stride, kernel_size=3,
|
||||
groups=self.expansion_channels, padding=1)
|
||||
self.bn2 = nn.BatchNorm2d(self.expansion_channels)
|
||||
|
||||
self.conv3 = nn.Conv2d(self.expansion_channels, out_channels, stride=1, kernel_size=1)
|
||||
self.bn3 = nn.BatchNorm2d(out_channels)
|
||||
|
||||
self.activation = activation(inplace=True) # inplace=True
|
||||
self.stride = stride
|
||||
self.input_channels = input_channels
|
||||
self.out_channels = out_channels
|
||||
|
||||
def forward(self, input):
|
||||
residual = input
|
||||
out = self.conv1(input)
|
||||
out = self.bn1(out)
|
||||
# out = self.activation(out)
|
||||
|
||||
out = self.depth_conv2(out)
|
||||
out = self.bn2(out)
|
||||
out = self.activation(out)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
if self.stride == 1 and self.input_channels == self.out_channels:
|
||||
out += residual
|
||||
return out
|
||||
|
||||
class AuxiliaryNet(nn.Module):
|
||||
def __init__(self, input_channels, nums_class=3, activation=nn.ReLU, first_conv_stride=2):
|
||||
super(AuxiliaryNet, self).__init__()
|
||||
self.input_channels = input_channels
|
||||
# self.num_channels = [128, 128, 32, 128, 32]
|
||||
self.num_channels = [512, 512, 512, 512, 1024]
|
||||
self.conv1 = nn.Conv2d(self.input_channels, self.num_channels[0], kernel_size=3, stride=first_conv_stride,
|
||||
padding=1)
|
||||
self.bn1 = nn.BatchNorm2d(self.num_channels[0])
|
||||
|
||||
self.conv2 = nn.Conv2d(self.num_channels[0], self.num_channels[1], kernel_size=3, stride=1, padding=1)
|
||||
self.bn2 = nn.BatchNorm2d(self.num_channels[1])
|
||||
|
||||
self.conv3 = nn.Conv2d(self.num_channels[1], self.num_channels[2], kernel_size=3, stride=2, padding=1)
|
||||
self.bn3 = nn.BatchNorm2d(self.num_channels[2])
|
||||
|
||||
self.conv4 = nn.Conv2d(self.num_channels[2], self.num_channels[3], kernel_size=7, stride=1, padding=3)
|
||||
self.bn4 = nn.BatchNorm2d(self.num_channels[3])
|
||||
|
||||
self.fc1 = nn.Linear(in_features=self.num_channels[3], out_features=self.num_channels[4])
|
||||
self.fc2 = nn.Linear(in_features=self.num_channels[4], out_features=nums_class)
|
||||
|
||||
self.activation = activation(inplace=True)
|
||||
|
||||
self.init_params()
|
||||
|
||||
def init_params(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, mode='fan_out')
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.Linear):
|
||||
nn.init.normal_(m.weight, std=0.01)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def forward(self, input):
|
||||
out = self.conv1(input)
|
||||
out = self.bn1(out)
|
||||
out = self.activation(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
out = self.activation(out)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
out = self.activation(out)
|
||||
|
||||
out = self.conv4(out)
|
||||
out = self.bn4(out)
|
||||
out = self.activation(out)
|
||||
|
||||
out = functional.adaptive_avg_pool2d(out, 1).squeeze(-1).squeeze(-1)
|
||||
out = self.fc1(out)
|
||||
euler_angles_pre = self.fc2(out)
|
||||
|
||||
return euler_angles_pre
|
||||
|
||||
class MobileNetV2(nn.Module):
|
||||
def __init__(self, input_channels=3, num_of_channels=None, nums_class=136, activation=nn.ReLU6):
|
||||
super(MobileNetV2, self).__init__()
|
||||
assert num_of_channels is not None
|
||||
self.num_of_channels = num_of_channels
|
||||
self.conv1 = nn.Conv2d(input_channels, self.num_of_channels[0], kernel_size=3, stride=2, padding=1)
|
||||
self.bn1 = nn.BatchNorm2d(self.num_of_channels[0])
|
||||
|
||||
self.depth_conv2 = nn.Conv2d(self.num_of_channels[0], self.num_of_channels[0], kernel_size=3, stride=1,
|
||||
padding=1, groups=self.num_of_channels[0])
|
||||
self.bn2 = nn.BatchNorm2d(self.num_of_channels[0])
|
||||
|
||||
self.stage0 = self.make_stage(self.num_of_channels[0], self.num_of_channels[0], stride=2, stage=0, times=5,
|
||||
expansion=2, activation=activation)
|
||||
|
||||
self.stage1 = self.make_stage(self.num_of_channels[0], self.num_of_channels[1], stride=2, stage=1, times=7,
|
||||
expansion=4, activation=activation)
|
||||
|
||||
self.linear_bottleneck_end = nn.Sequential(LinearBottleneck(self.num_of_channels[1], self.num_of_channels[2],
|
||||
expansion=2, stride=1, activation=activation))
|
||||
|
||||
self.conv3 = nn.Conv2d(self.num_of_channels[2], self.num_of_channels[3], kernel_size=3, stride=2, padding=1)
|
||||
self.bn3 = nn.BatchNorm2d(self.num_of_channels[3])
|
||||
|
||||
self.conv4 = nn.Conv2d(self.num_of_channels[3], self.num_of_channels[4], kernel_size=7, stride=1)
|
||||
self.bn4 = nn.BatchNorm2d(self.num_of_channels[4])
|
||||
|
||||
self.activation = activation(inplace=True)
|
||||
|
||||
self.in_features = 14 * 14 * self.num_of_channels[2] + 7 * 7 * self.num_of_channels[3] + 1 * 1 * self.num_of_channels[4]
|
||||
self.fc = nn.Linear(in_features=self.in_features, out_features=nums_class)
|
||||
|
||||
self.init_params()
|
||||
|
||||
def init_params(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, mode='fan_out')
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.Linear):
|
||||
nn.init.normal_(m.weight, std=0.01)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def make_stage(self, input_channels, out_channels, stride, stage, times, expansion, activation=nn.ReLU6):
|
||||
modules = OrderedDict()
|
||||
stage_name = 'LinearBottleneck{}'.format(stage)
|
||||
|
||||
module = LinearBottleneck(input_channels, out_channels, expansion=2,
|
||||
stride=stride, activation=activation)
|
||||
modules[stage_name + '_0'] = module
|
||||
|
||||
for i in range(times - 1):
|
||||
module = LinearBottleneck(out_channels, out_channels, expansion=expansion, stride=1,
|
||||
activation=activation)
|
||||
module_name = stage_name + '_{}'.format(i + 1)
|
||||
modules[module_name] = module
|
||||
|
||||
return nn.Sequential(modules)
|
||||
|
||||
def forward(self, input):
|
||||
with torch.no_grad():
|
||||
out = self.conv1(input)
|
||||
out = self.bn1(out)
|
||||
out = self.activation(out)
|
||||
|
||||
out = self.depth_conv2(out)
|
||||
out = self.bn2(out)
|
||||
out = self.activation(out)
|
||||
|
||||
out = self.stage0(out)
|
||||
out1 = self.stage1(out)
|
||||
|
||||
out1 = self.linear_bottleneck_end(out1)
|
||||
|
||||
out2 = self.conv3(out1)
|
||||
out2 = self.bn3(out2)
|
||||
out2 = self.activation(out2)
|
||||
|
||||
out3 = self.conv4(out2)
|
||||
out3 = self.bn4(out3)
|
||||
out3 = self.activation(out3)
|
||||
|
||||
out1 = out1.contiguous().view(out1.size(0), -1)
|
||||
out2 = out2.contiguous().view(out2.size(0), -1)
|
||||
out3 = out3.contiguous().view(out3.size(0), -1)
|
||||
|
||||
multi_scale = torch.cat([out1, out2, out3], 1)
|
||||
assert multi_scale.size(1) == self.in_features
|
||||
pre_landmarks = self.fc(multi_scale)
|
||||
return pre_landmarks, out
|
||||
|
||||
class PreProcess:
|
||||
eye_margin = 16
|
||||
diff_size = (64, 64)
|
||||
|
||||
def __init__(self, device="cpu", need_parser=True, retinaface_detection=None, face_skin=None, landmark_path=None):
|
||||
self.device = device
|
||||
self.img_size = 256
|
||||
|
||||
xs, ys = np.meshgrid(
|
||||
np.linspace(
|
||||
0, self.img_size - 1,
|
||||
self.img_size
|
||||
),
|
||||
np.linspace(
|
||||
0, self.img_size - 1,
|
||||
self.img_size
|
||||
)
|
||||
)
|
||||
xs = xs[None].repeat(68, axis=0)
|
||||
ys = ys[None].repeat(68, axis=0)
|
||||
fix = np.concatenate([ys, xs], axis=0)
|
||||
self.fix = torch.Tensor(fix).to(self.device)
|
||||
self.retinaface_detection = retinaface_detection
|
||||
if need_parser:
|
||||
self.face_parse = FaceParser(device=device, face_skin=face_skin)
|
||||
|
||||
self.landmark = MobileNetV2(num_of_channels=[64, 128, 16, 32, 128], nums_class=136)
|
||||
self.landmark.load_state_dict(torch.load(landmark_path))
|
||||
self.landmark.eval().to(self.device)
|
||||
self.up_ratio = 0.6 / 0.85
|
||||
self.down_ratio = 0.2 / 0.85
|
||||
self.width_ratio = 0.2 / 0.85
|
||||
self.lip_class = [7, 9]
|
||||
self.face_class = [1, 6]
|
||||
|
||||
self.transform = transforms.Compose(
|
||||
[
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
|
||||
]
|
||||
)
|
||||
|
||||
def relative2absolute(self, lms):
|
||||
return lms * self.img_size
|
||||
|
||||
def process(self, mask, lms, device="cpu"):
|
||||
diff = to_var(
|
||||
(self.fix.double() - torch.tensor(lms.transpose((1, 0)
|
||||
).reshape(-1, 1, 1)).to(self.device)
|
||||
).unsqueeze(0), requires_grad=False).to(self.device)
|
||||
|
||||
lms_eye_left = lms[42:48]
|
||||
lms_eye_right = lms[36:42]
|
||||
lms = lms.transpose((1, 0)).reshape(-1, 1, 1) # transpose to (y-x)
|
||||
# lms = np.tile(lms, (1, 256, 256)) # (136, h, w)
|
||||
diff = to_var((self.fix.double() - torch.tensor(lms).to(self.device)).unsqueeze(0), requires_grad=False).to(self.device)
|
||||
|
||||
mask_lip = (mask == self.lip_class[0]).float() + (mask == self.lip_class[1]).float()
|
||||
mask_face = (mask == self.face_class[0]).float() + (mask == self.face_class[1]).float()
|
||||
|
||||
mask_eyes = torch.zeros_like(mask, device=device)
|
||||
copy_area(mask_eyes, mask_face, lms_eye_left)
|
||||
copy_area(mask_eyes, mask_face, lms_eye_right)
|
||||
mask_eyes = to_var(mask_eyes, requires_grad=False).to(device)
|
||||
|
||||
mask_list = [mask_lip, mask_face, mask_eyes]
|
||||
mask_aug = torch.cat(mask_list, 0) # (3, 1, h, w)
|
||||
mask_re = F.interpolate(mask_aug, size=self.diff_size).repeat(1, diff.shape[1], 1, 1) # (3, 136, 64, 64)
|
||||
diff_re = F.interpolate(diff, size=self.diff_size).repeat(3, 1, 1, 1) # (3, 136, 64, 64)
|
||||
diff_re = diff_re * mask_re # (3, 136, 32, 32)
|
||||
norm = torch.norm(diff_re, dim=1, keepdim=True).repeat(1, diff_re.shape[1], 1, 1)
|
||||
norm = torch.where(norm == 0, torch.tensor(1e10, device=device), norm)
|
||||
diff_re /= norm
|
||||
|
||||
return mask_aug, diff_re
|
||||
|
||||
def __call__(self, image: Image):
|
||||
retinaface_result = self.retinaface_detection(image)
|
||||
face = []
|
||||
for box in retinaface_result['boxes']:
|
||||
face.append(rectangle(*np.int32(box)))
|
||||
|
||||
if len(face) == 0:
|
||||
return None, None, None
|
||||
|
||||
face_on_image = face[0]
|
||||
image, face, crop_face = crop(image, face_on_image, self.up_ratio, self.down_ratio, self.width_ratio)
|
||||
np_image = np.array(image)
|
||||
mask = self.face_parse.parse(cv2.resize(np_image, (512, 512)))
|
||||
|
||||
# obtain face parsing result
|
||||
mask = F.interpolate(
|
||||
mask.view(1, 1, 512, 512),
|
||||
(self.img_size, self.img_size),
|
||||
mode="nearest")
|
||||
mask = mask.type(torch.uint8)
|
||||
mask = to_var(mask, requires_grad=False).to(self.device)
|
||||
|
||||
input = image.crop([face.left(), face.top(), face.right(), face.bottom()])
|
||||
input = input.resize([112, 112])
|
||||
input = np.expand_dims(np.array(input, np.float32) / 255.0, 0)
|
||||
input = torch.Tensor(input.transpose((0, 3, 1, 2))).to(self.device)
|
||||
|
||||
pre_landmarks, _ = self.landmark(input)
|
||||
lms = pre_landmarks[0].cpu().detach().numpy()
|
||||
lms = lms.reshape(-1, 2) * [face.width(), face.height()] + np.int32([face.left(), face.top()])
|
||||
lms = lms / [np.shape(image)[0], np.shape(image)[1]] * self.img_size
|
||||
lms = lms[:, ::-1]
|
||||
|
||||
mask, diff = self.process(mask, lms, device=self.device)
|
||||
image = image.resize((self.img_size, self.img_size), Image.Resampling.NEAREST)
|
||||
image = self.transform(image)
|
||||
real = to_var(image.unsqueeze(0))
|
||||
return [real, mask, diff], face_on_image, crop_face
|
||||
|
||||
# Solver part (GAN part)
|
||||
def l2normalize(v, eps=1e-12):
|
||||
return v / (v.norm() + eps)
|
||||
|
||||
class SpectralNorm(object):
|
||||
def __init__(self):
|
||||
self.name = "weight"
|
||||
self.power_iterations = 1
|
||||
|
||||
def compute_weight(self, module):
|
||||
u = getattr(module, self.name + "_u")
|
||||
v = getattr(module, self.name + "_v")
|
||||
w = getattr(module, self.name + "_bar")
|
||||
|
||||
height = w.data.shape[0]
|
||||
for _ in range(self.power_iterations):
|
||||
v.data = l2normalize(torch.mv(torch.t(w.view(height, -1).data), u.data))
|
||||
u.data = l2normalize(torch.mv(w.view(height, -1).data, v.data))
|
||||
# sigma = torch.dot(u.data, torch.mv(w.view(height,-1).data, v.data))
|
||||
sigma = u.dot(w.view(height, -1).mv(v))
|
||||
return w / sigma.expand_as(w)
|
||||
|
||||
@staticmethod
|
||||
def apply(module):
|
||||
name = "weight"
|
||||
fn = SpectralNorm()
|
||||
|
||||
try:
|
||||
u = getattr(module, name + "_u")
|
||||
v = getattr(module, name + "_v")
|
||||
w = getattr(module, name + "_bar")
|
||||
except AttributeError:
|
||||
w = getattr(module, name)
|
||||
height = w.data.shape[0]
|
||||
width = w.view(height, -1).data.shape[1]
|
||||
u = Parameter(w.data.new(height).normal_(0, 1), requires_grad=False)
|
||||
v = Parameter(w.data.new(width).normal_(0, 1), requires_grad=False)
|
||||
w_bar = Parameter(w.data)
|
||||
|
||||
# del module._parameters[name]
|
||||
|
||||
module.register_parameter(name + "_u", u)
|
||||
module.register_parameter(name + "_v", v)
|
||||
module.register_parameter(name + "_bar", w_bar)
|
||||
|
||||
# remove w from parameter list
|
||||
del module._parameters[name]
|
||||
|
||||
setattr(module, name, fn.compute_weight(module))
|
||||
|
||||
# recompute weight before every forward()
|
||||
module.register_forward_pre_hook(fn)
|
||||
|
||||
return fn
|
||||
|
||||
def remove(self, module):
|
||||
weight = self.compute_weight(module)
|
||||
delattr(module, self.name)
|
||||
del module._parameters[self.name + '_u']
|
||||
del module._parameters[self.name + '_v']
|
||||
del module._parameters[self.name + '_bar']
|
||||
module.register_parameter(self.name, Parameter(weight.data))
|
||||
|
||||
def __call__(self, module, inputs):
|
||||
setattr(module, self.name, self.compute_weight(module))
|
||||
|
||||
def spectral_norm(module):
|
||||
SpectralNorm.apply(module)
|
||||
return module
|
||||
|
||||
def remove_spectral_norm(module):
|
||||
name = 'weight'
|
||||
for k, hook in module._forward_pre_hooks.items():
|
||||
if isinstance(hook, SpectralNorm) and hook.name == name:
|
||||
hook.remove(module)
|
||||
del module._forward_pre_hooks[k]
|
||||
return module
|
||||
|
||||
raise ValueError("spectral_norm of '{}' not found in {}"
|
||||
.format(name, module))
|
||||
|
||||
# Defines the GAN loss which uses either LSGAN or the regular GAN.
|
||||
# When LSGAN is used, it is basically same as MSELoss,
|
||||
# but it abstracts away the need to create the target label tensor
|
||||
# that has the same size as the input
|
||||
class ResidualBlock(nn.Module):
|
||||
"""Residual Block."""
|
||||
|
||||
def __init__(self, dim_in, dim_out, net_mode=None):
|
||||
if net_mode == 'p' or (net_mode is None):
|
||||
use_affine = True
|
||||
elif net_mode == 't':
|
||||
use_affine = False
|
||||
super(ResidualBlock, self).__init__()
|
||||
self.main = nn.Sequential(
|
||||
nn.Conv2d(dim_in, dim_out, kernel_size=3, stride=1, padding=1, bias=False),
|
||||
nn.InstanceNorm2d(dim_out, affine=use_affine),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(dim_out, dim_out, kernel_size=3, stride=1, padding=1, bias=False),
|
||||
nn.InstanceNorm2d(dim_out, affine=use_affine)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return x + self.main(x)
|
||||
|
||||
class GetMatrix(nn.Module):
|
||||
def __init__(self, dim_in, dim_out):
|
||||
super(GetMatrix, self).__init__()
|
||||
self.get_gamma = nn.Conv2d(dim_in, dim_out, kernel_size=1, stride=1, padding=0, bias=False)
|
||||
self.get_beta = nn.Conv2d(dim_in, dim_out, kernel_size=1, stride=1, padding=0, bias=False)
|
||||
|
||||
def forward(self, x):
|
||||
gamma = self.get_gamma(x)
|
||||
beta = self.get_beta(x)
|
||||
return x, gamma, beta
|
||||
|
||||
class NONLocalBlock2D(nn.Module):
|
||||
def __init__(self):
|
||||
super(NONLocalBlock2D, self).__init__()
|
||||
self.g = nn.Conv2d(in_channels=1, out_channels=1,
|
||||
kernel_size=1, stride=1, padding=0)
|
||||
|
||||
def forward(self, source, weight):
|
||||
"""(b, c, h, w)
|
||||
src_diff: (3, 136, 32, 32)
|
||||
"""
|
||||
batch_size = source.size(0)
|
||||
|
||||
g_source = source.view(batch_size, 1, -1) # (N, C, H*W)
|
||||
g_source = g_source.permute(0, 2, 1) # (N, H*W, C)
|
||||
|
||||
y = torch.bmm(weight.to_dense(), g_source)
|
||||
y = y.permute(0, 2, 1).contiguous() # (N, C, H*W)
|
||||
y = y.view(batch_size, 1, *source.size()[2:])
|
||||
return y
|
||||
|
||||
class Generator(nn.Module):
|
||||
"""Generator. Encoder-Decoder Architecture."""
|
||||
|
||||
def __init__(self):
|
||||
super(Generator, self).__init__()
|
||||
|
||||
# -------------------------- PNet(MDNet) for obtaining makeup matrices --------------------------
|
||||
|
||||
layers = nn.Sequential(
|
||||
nn.Conv2d(3, 64, kernel_size=7, stride=1, padding=3, bias=False),
|
||||
nn.InstanceNorm2d(64, affine=True),
|
||||
nn.ReLU(inplace=True)
|
||||
)
|
||||
self.pnet_in = layers
|
||||
|
||||
# Down-Sampling
|
||||
curr_dim = 64
|
||||
for i in range(2):
|
||||
layers = nn.Sequential(
|
||||
nn.Conv2d(curr_dim, curr_dim * 2, kernel_size=4, stride=2, padding=1, bias=False),
|
||||
nn.InstanceNorm2d(curr_dim * 2, affine=True),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
|
||||
setattr(self, f'pnet_down_{i + 1}', layers)
|
||||
curr_dim = curr_dim * 2
|
||||
|
||||
# Bottleneck. All bottlenecks share the same attention module
|
||||
self.atten_bottleneck_g = NONLocalBlock2D()
|
||||
self.atten_bottleneck_b = NONLocalBlock2D()
|
||||
self.simple_spade = GetMatrix(curr_dim, 1) # get the makeup matrix
|
||||
|
||||
for i in range(3):
|
||||
setattr(self, f'pnet_bottleneck_{i + 1}', ResidualBlock(dim_in=curr_dim, dim_out=curr_dim, net_mode='p'))
|
||||
|
||||
# --------------------------- TNet(MANet) for applying makeup transfer ----------------------------
|
||||
|
||||
self.tnet_in_conv = nn.Conv2d(3, 64, kernel_size=7, stride=1, padding=3, bias=False)
|
||||
self.tnet_in_spade = nn.InstanceNorm2d(64, affine=False)
|
||||
self.tnet_in_relu = nn.ReLU(inplace=True)
|
||||
|
||||
# Down-Sampling
|
||||
curr_dim = 64
|
||||
for i in range(2):
|
||||
setattr(self, f'tnet_down_conv_{i + 1}', nn.Conv2d(curr_dim, curr_dim * 2, kernel_size=4, stride=2, padding=1, bias=False))
|
||||
setattr(self, f'tnet_down_spade_{i + 1}', nn.InstanceNorm2d(curr_dim * 2, affine=False))
|
||||
setattr(self, f'tnet_down_relu_{i + 1}', nn.ReLU(inplace=True))
|
||||
curr_dim = curr_dim * 2
|
||||
|
||||
# Bottleneck
|
||||
for i in range(6):
|
||||
setattr(self, f'tnet_bottleneck_{i + 1}', ResidualBlock(dim_in=curr_dim, dim_out=curr_dim, net_mode='t'))
|
||||
|
||||
# Up-Sampling
|
||||
for i in range(2):
|
||||
setattr(self, f'tnet_up_conv_{i + 1}', nn.ConvTranspose2d(curr_dim, curr_dim // 2, kernel_size=4, stride=2, padding=1, bias=False))
|
||||
setattr(self, f'tnet_up_spade_{i + 1}', nn.InstanceNorm2d(curr_dim // 2, affine=False))
|
||||
setattr(self, f'tnet_up_relu_{i + 1}', nn.ReLU(inplace=True))
|
||||
curr_dim = curr_dim // 2
|
||||
|
||||
layers = nn.Sequential(
|
||||
nn.Conv2d(curr_dim, 3, kernel_size=7, stride=1, padding=3, bias=False),
|
||||
nn.Tanh()
|
||||
)
|
||||
self.tnet_out = layers
|
||||
|
||||
@staticmethod
|
||||
def atten_feature(mask_s, weight, gamma_s, beta_s, atten_module_g, atten_module_b):
|
||||
"""
|
||||
feature size: (1, c, h, w)
|
||||
mask_c(s): (3, 1, h, w)
|
||||
diff_c: (1, 138, 256, 256)
|
||||
return: (1, c, h, w)
|
||||
"""
|
||||
channel_num = gamma_s.shape[1]
|
||||
|
||||
mask_s_re = F.interpolate(mask_s, size=gamma_s.shape[2:]).repeat(1, channel_num, 1, 1)
|
||||
gamma_s_re = gamma_s.repeat(3, 1, 1, 1)
|
||||
gamma_s = gamma_s_re * mask_s_re # (3, c, h, w)
|
||||
beta_s_re = beta_s.repeat(3, 1, 1, 1)
|
||||
beta_s = beta_s_re * mask_s_re
|
||||
|
||||
gamma = atten_module_g(gamma_s, weight) # (3, c, h, w)
|
||||
beta = atten_module_b(beta_s, weight)
|
||||
|
||||
gamma = (gamma[0] + gamma[1] + gamma[2]).unsqueeze(0) # (c, h, w) combine the three parts
|
||||
beta = (beta[0] + beta[1] + beta[2]).unsqueeze(0)
|
||||
return gamma, beta
|
||||
|
||||
def get_weight(self, mask_c, mask_s, fea_c, fea_s, diff_c, diff_s):
|
||||
""" s --> source; c --> target
|
||||
feature size: (1, 256, 64, 64)
|
||||
diff: (3, 136, 32, 32)
|
||||
"""
|
||||
HW = 64 * 64
|
||||
batch_size = 3
|
||||
assert fea_s is not None # fea_s when i==3
|
||||
# get 3 part fea using mask
|
||||
channel_num = fea_s.shape[1]
|
||||
|
||||
mask_c_re = F.interpolate(mask_c, size=64).repeat(1, channel_num, 1, 1) # (3, c, h, w)
|
||||
fea_c = fea_c.repeat(3, 1, 1, 1) # (3, c, h, w)
|
||||
fea_c = fea_c * mask_c_re # (3, c, h, w) 3 stands for 3 parts
|
||||
|
||||
mask_s_re = F.interpolate(mask_s, size=64).repeat(1, channel_num, 1, 1)
|
||||
fea_s = fea_s.repeat(3, 1, 1, 1)
|
||||
fea_s = fea_s * mask_s_re
|
||||
|
||||
theta_input = torch.cat((fea_c * 0.01, diff_c), dim=1)
|
||||
phi_input = torch.cat((fea_s * 0.01, diff_s), dim=1)
|
||||
|
||||
theta_target = theta_input.view(batch_size, -1, HW) # (N, C+136, H*W)
|
||||
theta_target = theta_target.permute(0, 2, 1) # (N, H*W, C+136)
|
||||
|
||||
phi_source = phi_input.view(batch_size, -1, HW) # (N, C+136, H*W)
|
||||
|
||||
weight = torch.bmm(theta_target, phi_source) # (3, HW, HW)
|
||||
with torch.no_grad():
|
||||
v = weight.detach().nonzero().long().permute(1, 0)
|
||||
# This clone is required to correctly release cuda memory.
|
||||
weight_ind = v.clone()
|
||||
del v
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
weight *= 200 # hyper parameters for visual feature
|
||||
weight = F.softmax(weight, dim=-1)
|
||||
weight = weight[weight_ind[0], weight_ind[1], weight_ind[2]]
|
||||
ret = torch.sparse.FloatTensor(weight_ind, weight, torch.Size([3, HW, HW]))
|
||||
return ret
|
||||
|
||||
def forward(self, c, s, mask_c, mask_s, diff_c, diff_s, gamma=None, beta=None, ret=False):
|
||||
c, s, mask_c, mask_s, diff_c, diff_s = [x.squeeze(0) if x.ndim == 5 else x for x in [c, s, mask_c, mask_s, diff_c, diff_s]]
|
||||
"""attention version
|
||||
c: content, stands for source image. shape: (b, c, h, w)
|
||||
s: style, stands for reference image. shape: (b, c, h, w)
|
||||
mask_list_c: lip, skin, eye. (b, 1, h, w)
|
||||
"""
|
||||
|
||||
# forward c in tnet(MANet)
|
||||
c_tnet = self.tnet_in_conv(c)
|
||||
s = self.pnet_in(s)
|
||||
c_tnet = self.tnet_in_spade(c_tnet)
|
||||
c_tnet = self.tnet_in_relu(c_tnet)
|
||||
|
||||
# down-sampling
|
||||
for i in range(2):
|
||||
if gamma is None:
|
||||
cur_pnet_down = getattr(self, f'pnet_down_{i + 1}')
|
||||
s = cur_pnet_down(s)
|
||||
|
||||
cur_tnet_down_conv = getattr(self, f'tnet_down_conv_{i + 1}')
|
||||
cur_tnet_down_spade = getattr(self, f'tnet_down_spade_{i + 1}')
|
||||
cur_tnet_down_relu = getattr(self, f'tnet_down_relu_{i + 1}')
|
||||
c_tnet = cur_tnet_down_conv(c_tnet)
|
||||
c_tnet = cur_tnet_down_spade(c_tnet)
|
||||
c_tnet = cur_tnet_down_relu(c_tnet)
|
||||
|
||||
# bottleneck
|
||||
for i in range(6):
|
||||
if gamma is None and i <= 2:
|
||||
cur_pnet_bottleneck = getattr(self, f'pnet_bottleneck_{i + 1}')
|
||||
cur_tnet_bottleneck = getattr(self, f'tnet_bottleneck_{i + 1}')
|
||||
|
||||
# get s_pnet from p and transform
|
||||
if i == 3:
|
||||
if gamma is None: # not in test_mix
|
||||
s, gamma, beta = self.simple_spade(s)
|
||||
weight = self.get_weight(mask_c, mask_s, c_tnet, s, diff_c, diff_s)
|
||||
gamma, beta = self.atten_feature(mask_s, weight, gamma, beta, self.atten_bottleneck_g, self.atten_bottleneck_b)
|
||||
if ret:
|
||||
return [gamma, beta]
|
||||
# else: # in test mode
|
||||
# gamma, beta = param_A[0]*w + param_B[0]*(1-w), param_A[1]*w + param_B[1]*(1-w)
|
||||
|
||||
c_tnet = c_tnet * (1 + gamma) + beta # apply makeup transfer using makeup matrices
|
||||
|
||||
if gamma is None and i <= 2:
|
||||
s = cur_pnet_bottleneck(s)
|
||||
c_tnet = cur_tnet_bottleneck(c_tnet)
|
||||
|
||||
# up-sampling
|
||||
for i in range(2):
|
||||
cur_tnet_up_conv = getattr(self, f'tnet_up_conv_{i + 1}')
|
||||
cur_tnet_up_spade = getattr(self, f'tnet_up_spade_{i + 1}')
|
||||
cur_tnet_up_relu = getattr(self, f'tnet_up_relu_{i + 1}')
|
||||
c_tnet = cur_tnet_up_conv(c_tnet)
|
||||
c_tnet = cur_tnet_up_spade(c_tnet)
|
||||
c_tnet = cur_tnet_up_relu(c_tnet)
|
||||
|
||||
c_tnet = self.tnet_out(c_tnet)
|
||||
return c_tnet
|
||||
|
||||
# Gan Solver
|
||||
class Solver():
|
||||
def __init__(self, device="cpu", inference=None):
|
||||
self.G = Generator()
|
||||
self.G.load_state_dict(torch.load(inference, map_location=torch.device(device)))
|
||||
self.G = self.G.to(device).eval()
|
||||
return
|
||||
|
||||
def generate(self, org_A, ref_B, lms_A=None, lms_B=None, mask_A=None, mask_B=None,
|
||||
diff_A=None, diff_B=None, gamma=None, beta=None, ret=False):
|
||||
"""org_A is content, ref_B is style"""
|
||||
res = self.G(org_A, ref_B, mask_A, mask_B, diff_A, diff_B, gamma, beta, ret)
|
||||
return res
|
||||
|
||||
def test(self, real_A, mask_A, diff_A, real_B, mask_B, diff_B):
|
||||
cur_prama = None
|
||||
with torch.no_grad():
|
||||
cur_prama = self.generate(real_A, real_B, None, None, mask_A, mask_B,
|
||||
diff_A, diff_B, ret=True)
|
||||
fake_A = self.generate(real_A, real_B, None, None, mask_A, mask_B,
|
||||
diff_A, diff_B, gamma=cur_prama[0], beta=cur_prama[1])
|
||||
fake_A = fake_A.squeeze(0)
|
||||
|
||||
# normalize
|
||||
min_, max_ = fake_A.min(), fake_A.max()
|
||||
fake_A.add_(-min_).div_(max_ - min_ + 1e-5)
|
||||
|
||||
return ToPILImage()(fake_A.cpu())
|
||||
|
||||
# PostProcess part
|
||||
class PostProcess:
|
||||
def __init__(self):
|
||||
self.denoise = False
|
||||
self.img_size = 256
|
||||
|
||||
def __call__(self, source: Image, result: Image):
|
||||
source = np.array(source)
|
||||
result = np.array(result)
|
||||
|
||||
height, width = source.shape[:2]
|
||||
small_source = cv2.resize(source, (self.img_size, self.img_size))
|
||||
laplacian_diff = source.astype(np.float64) - cv2.resize(small_source, (width, height)).astype(np.float64)
|
||||
result = (cv2.resize(result, (width, height)) + laplacian_diff).round().clip(0, 255).astype(np.uint8)
|
||||
if self.denoise:
|
||||
result = cv2.fastNlMeansDenoisingColored(result)
|
||||
result = Image.fromarray(result).convert('RGB')
|
||||
return result
|
||||
|
||||
class PSGAN_Inference:
|
||||
"""
|
||||
An inference wrapper for makeup transfer.
|
||||
It takes two image `source` and `reference` in,
|
||||
and transfers the makeup of reference to source.
|
||||
"""
|
||||
|
||||
def __init__(self, device="cpu", model_path="assets/models/G.pth", retinaface_detection=None, face_skin=None, landmark_path=None):
|
||||
"""
|
||||
Args:
|
||||
device (str): Device type and index, such as "cpu" or "cuda:2".
|
||||
device_id (int): Specifying which device index
|
||||
will be used for inference.
|
||||
"""
|
||||
self.device = device
|
||||
self.solver = Solver(device, inference=model_path)
|
||||
self.preprocess = PreProcess(device, retinaface_detection=retinaface_detection, face_skin=face_skin, landmark_path=landmark_path)
|
||||
self.postprocess = PostProcess()
|
||||
|
||||
def transfer(self, source: Image, reference: Image):
|
||||
"""
|
||||
Args:
|
||||
source (Image): The image where makeup will be transferred to.
|
||||
reference (Image): Image containing targeted makeup.
|
||||
Return:
|
||||
Image: Transferred image.
|
||||
"""
|
||||
source_input, face, crop_face = self.preprocess(source)
|
||||
reference_input, _, _ = self.preprocess(reference)
|
||||
|
||||
if not (source_input and reference_input):
|
||||
return source
|
||||
|
||||
for i in range(len(source_input)):
|
||||
source_input[i] = source_input[i].to(self.device)
|
||||
|
||||
for i in range(len(reference_input)):
|
||||
reference_input[i] = reference_input[i].to(self.device)
|
||||
|
||||
# TODO: Abridge the parameter list.
|
||||
result = self.solver.test(*source_input, *reference_input)
|
||||
|
||||
source_crop = source.crop((crop_face.left(), crop_face.top(), crop_face.right(), crop_face.bottom()))
|
||||
result = self.postprocess(source_crop, result)
|
||||
return result
|
||||
+3
-1
@@ -1,3 +1,4 @@
|
||||
aliyun-python-sdk-core-v3==2.13.10
|
||||
opencv-python
|
||||
tensorflow-cpu
|
||||
tensorflow
|
||||
@@ -7,4 +8,5 @@ modelscope
|
||||
scikit-image
|
||||
matplotlib
|
||||
insightface
|
||||
diffusers==0.18.2
|
||||
diffusers==0.18.2
|
||||
sentencepiece
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user