xsl/change_hair
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Compare commits

merge into: xsl/change_hair:feature/inpaint-and-feather
Branches Tags
xsl/change_hair:master xsl/change_hair:ubuntu3090_debug xsl/change_hair:ubuntu3090 xsl/change_hair:xsl5090 xsl/change_hair:feature/inpaint-and-feather
..
pull from: xsl/change_hair:ubuntu3090_debug
Branches Tags
xsl/change_hair:ubuntu3090_debug xsl/change_hair:ubuntu3090 xsl/change_hair:xsl5090 xsl/change_hair:master xsl/change_hair:feature/inpaint-and-feather
11 Commits
feature/inpaint-and-feather .. ubuntu3090_debug
Author SHA1 Message Date
xsl 326b206d08 feat(swapHair): webui_steps参数可外部传入(接口2调试页可调) 
- gen_super_image.py: webui_img2img + build_body_v2 接收可选 steps 参数
- run_copy_cost_colorb64.py: swapHair 接口接收 webui_steps，透传给 webui_img2img
- 未传时保持环境变量 WEBUI_STEPS 默认(15)，向后兼容
 2026-07-26 16:49:21 +08:00
xsl 990fae929f perf(swapHair): webui降步数 + 用户图预处理内存缓存 
优化1 - webui img2img steps 20→15 (gen_super_image.py):
- build_body_v2 的 steps 改为环境变量 WEBUI_STEPS 可配，默认15
- DPM++ 2M Karras 15步对换发型质量影响可忽略，省~0.1s

优化2 - infer_hairstyle_diy_jy 用户图预处理内存缓存 (hairstyle_model.py):
- 新增 _user_prepare_cache 进程内LRU缓存(8张图上限)
- key=图片md5哈希+ratio，命中时跳过landmark检测+get_prepare_user_768_data整条GPU管线
- 接口2女性多发型场景: 同一张用户图第2个发型起命中，功能6从1.8s降至1.1s(省0.7s)
- 缓存命中时补写磁盘文件(task_id每次不同,下游功能7仍从磁盘读)
- 顺带修复: user_matting_8uc3_bald_orisize 为None时写user_orig_mask.png的crash
- 加分步计时日志(主GAN/融合耗时),便于定位热点

实测: 同图连续请求 swapHair 从4.28s降至3.5s(省18%)
 2026-07-26 16:41:19 +08:00
xslandCursor c5e50de40a chore: 固化 WebUI 底模为 v1-5，并纳入仓库配置副本 
本机无 majicmix，将 sd_model_checkpoint 固定为 v1-5-pruned-emaonly；因 onediff/webui 是独立 git 仓无法直接跟踪，配置副本放在 configs/。

Co-authored-by: Cursor <cursoragent@cursor.com>
 2026-07-19 15:52:51 +08:00
xslandCursor 0926c61bd0 fix: 适配 ubuntu 路径并增强换发色/训练流程稳定性 
将配置与训练脚本从 /home/xsl 切到本机 /home/ubuntu；换发色在 webui 增强失败或缺色板时降级返回，训练结束后自动重启 hair 服务再回调。

Co-authored-by: Cursor <cursoragent@cursor.com>
 2026-07-19 00:49:16 +08:00
xsl fc66534a74 feat: 系统配置文件适配 RTX 3090 
添加系统服务配置文件：
- change_hair-hair.service: hair_service_sd swapHair (8801)
- change_hair-photo.service: photo_service LoRA scheduler (32678)
- change_hair-webui.service: SD WebUI (57860)
 2026-07-18 19:18:56 +08:00
xsl 6f34e8876c feat: 适配 RTX 3090 (24GB) 环境优化 
硬件迁移：从 RTX 5090 (32GB) 迁移到 RTX 3090 (24GB)

主要改动：
- 所有启动脚本和配置文件中的 /home/xsl/ 路径替换为 /home/ubuntu/
- 适配新的 24GB VRAM 环境
 2026-07-18 19:00:15 +08:00
xsl 9b7324cd26 chore: 适配本机 xsl 路径（RTX 5090 机器） 
将所有绝对路径从 /home/ubuntu 改回 /home/xsl，覆盖启动脚本、
configure.ini 及训练/测试脚本。纯路径替换，无功能性改动。

本分支(xsl5090)用于本机(RTX 5090)运行；master 维持 /home/ubuntu
供远程 ubuntu 机器使用。
 2026-07-18 16:26:16 +08:00
xsl b2738fba7d asdf 2026-07-17 22:13:56 +08:00
xsl 2606aca7c5 save 2026-07-09 22:02:35 +08:00
xsl ce445b64a6 完善部署并训练5个新发型 + 换发型集成文档 
部署修复：
- torch.load 增加 weights_only=False patch，兼容 PyTorch 2.6+ 加载旧权重
- OSS 改为懒加载，本地用 output_format=base64 无需配凭证即可启动
- 补全被 gitignore 误排除的必需代码：core/models/layers/data、models/layers/data、keypoints/lib
- webui 训练命令 --xformers 改 --sdpa（修复 xformers 无 CUDA 支持报错）

功能调整：
- hair_grow_service 端口改 8899、preview 路由修复（send_file）
- list_hairstyles 增加发型白名单，测试页只展示当前5个发型

新增脚本：
- train_lora_parallel.py：直接调 kohya 并行训练 LoRA（绕过 photo_service 串行限制）
- train_hairstyles_parallel.py / train_batch_stepC.py：批量训练辅助脚本
- scripts/sync_data_to_server.sh：大文件断点续传到云服务器

文档：
- docs/换发型集成文档.md：换发型完整流程、服务架构、资源依赖、训练方法、集成步骤
 2026-07-09 20:45:08 +08:00
xsl aaa1bf159f 添加训练数据 2026-07-08 20:35:12 +08:00
85 changed files with 11405 additions and 924 deletions
Expand all files Collapse all files
README.md
+2 -2
View File
@@ -296,8 +296,8 @@ curl -X POST http://127.0.0.1:32678/api/hair/train \
| 改动 | 说明 | | 改动 | 说明 |
|------|------| |------|------|
| `config.json` 的 onediff compiler 路径 | 指向本机 onediff 目录 | | `config.json` 的 onediff compiler 路径 | 指向本机 onediff 目录 |
| 启动加 `HF_HUB_OFFLINE=1` | 本机无法访问 huggingface.co，用本地缓存离线加载 CLIP | | 启动加 `HF_HUB_OFFLINE=1` | 本机无法稳定访问 huggingface.co，用本地缓存离线加载 CLIP（需预先缓存 `openai/clip-vit-large-patch14`） |
| SD 模型 | 实际加载 `v1-5-pruned-emaonly`（config 里写的 majicmix 不存在，webui 自动 fallback） | | SD 模型 | `sd_model_checkpoint` 固定为 `v1-5-pruned-emaonly.safetensors`（本机无 majicmix）。仓库内权威副本：`configs/stable-diffusion-webui.config.json`；部署时拷到 `project/onediff/stable-diffusion-webui/config.json`（该目录是独立 git 仓库，无法被本仓直接跟踪） |
### 3. 换发型推理改动（`hair_service_sd/gen_super_image.py`） ### 3. 换发型推理改动（`hair_service_sd/gen_super_image.py`）
batch_train_hairstyles.py
+12 -12
View File
@@ -8,8 +8,8 @@
3. 串行执行，记录每个发型成功/失败 3. 串行执行，记录每个发型成功/失败
用法: 用法:
cd /home/xsl/change_hair/project/hair_service_sd cd /home/ubuntu/change_hair/project/hair_service_sd
python /home/xsl/change_hair/batch_train_hairstyles.py --src /home/xsl/change_hair/hair_type_images --gender girl [--only 圆-心形] [--start-from 心形-心形] python /home/ubuntu/change_hair/batch_train_hairstyles.py --src /home/ubuntu/change_hair/hair_type_images --gender girl [--only 圆-心形] [--start-from 心形-心形]
注意： 注意：
- 中文 hair_id 直接用文件名（去扩展名）作为 ID - 中文 hair_id 直接用文件名（去扩展名）作为 ID
@@ -24,14 +24,14 @@ import argparse
import subprocess import subprocess
from datetime import datetime from datetime import datetime
SRC_DEFAULT = "/home/xsl/change_hair/hair_type_images" SRC_DEFAULT = "/home/ubuntu/change_hair/hair_type_images"
WORK_DIR = "/home/xsl/change_hair/data/batch_train_inputs" WORK_DIR = "/home/ubuntu/change_hair/data/batch_train_inputs"
LOG_FILE = "/home/xsl/change_hair/data/batch_train_log.txt" LOG_FILE = "/home/ubuntu/change_hair/data/batch_train_log.txt"
TRAIN_SCRIPT = "/home/xsl/change_hair/train_hairstyle_full.py" TRAIN_SCRIPT = "/home/ubuntu/change_hair/train_hairstyle_full.py"
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd" HAIR_SERVICE_DIR = "/home/ubuntu/change_hair/project/hair_service_sd"
LOCK_FILE = "/home/xsl/change_hair/data/batch_train.pid" LOCK_FILE = "/home/ubuntu/change_hair/data/batch_train.pid"
DONE_FILE = "/home/xsl/change_hair/data/batch_train_done.txt" # 已完成发型清单（断点续跑） DONE_FILE = "/home/ubuntu/change_hair/data/batch_train_done.txt" # 已完成发型清单（断点续跑）
LORA_DIR = "/home/xsl/change_hair/data/train_material" # LoRA 输出根目录 LORA_DIR = "/home/ubuntu/change_hair/data/train_material" # LoRA 输出根目录
def acquire_lock(): def acquire_lock():
@@ -125,7 +125,7 @@ def run_one(hair_id, src_img, gender, py, done_set, force=False):
log(f" 调用: {' '.join(cmd[:2])} ... --hair-id {hair_id}") log(f" 调用: {' '.join(cmd[:2])} ... --hair-id {hair_id}")
try: try:
# 子进程输出实时写到日志文件 # 子进程输出实时写到日志文件
proc_log = os.path.join("/home/xsl/change_hair/data", f"subprocess_{hair_id}.log") proc_log = os.path.join("/home/ubuntu/change_hair/data", f"subprocess_{hair_id}.log")
with open(proc_log, "w", encoding="utf-8") as f: with open(proc_log, "w", encoding="utf-8") as f:
r = subprocess.run( r = subprocess.run(
cmd, cwd=HAIR_SERVICE_DIR, cmd, cwd=HAIR_SERVICE_DIR,
@@ -152,7 +152,7 @@ def main():
ap = argparse.ArgumentParser(description="批量训练发型") ap = argparse.ArgumentParser(description="批量训练发型")
ap.add_argument("--src", default=SRC_DEFAULT, help="原图目录") ap.add_argument("--src", default=SRC_DEFAULT, help="原图目录")
ap.add_argument("--gender", default="girl", choices=["boy", "girl"]) ap.add_argument("--gender", default="girl", choices=["boy", "girl"])
ap.add_argument("--py", default="/home/xsl/miniconda3/envs/my_hair/bin/python", ap.add_argument("--py", default="/home/ubuntu/miniconda3/envs/my_hair/bin/python",
help="python 解释器") help="python 解释器")
ap.add_argument("--only", default=None, help="只训练指定 hair_id（调试用）") ap.add_argument("--only", default=None, help="只训练指定 hair_id（调试用）")
ap.add_argument("--start-from", default=None, ap.add_argument("--start-from", default=None,
change_hair-hair.service
+19
View File
@@ -0,0 +1,19 @@
[Unit]
Description=change_hair hair_service_sd swapHair (8801)
After=network-online.target change_hair-webui.service change_hair-photo.service
Wants=change_hair-webui.service change_hair-photo.service
[Service]
Type=simple
User=ubuntu
WorkingDirectory=/home/ubuntu/change_hair/project/hair_service_sd
Environment=CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1
Environment=CUDA_VISIBLE_DEVICES=0
Environment=APP_WORKER_ID=1
ExecStart=/home/ubuntu/miniconda3/envs/my_hair/bin/python run_copy_cost_colorb64.py
Restart=on-failure
RestartSec=10
TimeoutStartSec=300
[Install]
WantedBy=multi-user.target
change_hair-photo.service
+17
View File
@@ -0,0 +1,17 @@
[Unit]
Description=change_hair photo_service LoRA scheduler (32678)
After=network-online.target
Wants=network-online.target
[Service]
Type=simple
User=ubuntu
WorkingDirectory=/home/ubuntu/change_hair/project/photo_service
Environment=CUDA_VISIBLE_DEVICES=0
Environment=APP_WORKER_ID=1
ExecStart=/home/ubuntu/miniconda3/envs/py310/bin/python -u lora_train_service_1.py
Restart=on-failure
RestartSec=10
[Install]
WantedBy=multi-user.target
change_hair-webui.service
+20
View File
@@ -0,0 +1,20 @@
[Unit]
Description=change_hair SD WebUI (57860)
After=network-online.target
Wants=network-online.target
[Service]
Type=simple
User=ubuntu
WorkingDirectory=/home/ubuntu/change_hair/project/onediff/stable-diffusion-webui
Environment=CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1
Environment=CUDA_VISIBLE_DEVICES=0
Environment=HF_HUB_OFFLINE=1
Environment=TRANSFORMERS_OFFLINE=1
ExecStart=/home/ubuntu/miniconda3/envs/sdwebui/bin/python webui.py --api --listen --xformers --port 57860
Restart=on-failure
RestartSec=10
TimeoutStartSec=600
[Install]
WantedBy=multi-user.target
configs/stable-diffusion-webui.config.json
Executable
+40
View File
@@ -0,0 +1,40 @@
{
"ldsr_steps": 100,
"ldsr_cached": false,
"SCUNET_tile": 256,
"SCUNET_tile_overlap": 8,
"SWIN_tile": 192,
"SWIN_tile_overlap": 8,
"SWIN_torch_compile": false,
"hypertile_enable_unet": false,
"hypertile_enable_unet_secondpass": false,
"hypertile_max_depth_unet": 3,
"hypertile_max_tile_unet": 256,
"hypertile_swap_size_unet": 3,
"hypertile_enable_vae": false,
"hypertile_max_depth_vae": 3,
"hypertile_max_tile_vae": 128,
"hypertile_swap_size_vae": 3,
"onediff_compiler_caches_path": "/home/ubuntu/change_hair/project/onediff/stable-diffusion-webui/extensions/onediff_sd_webui_extensions/compiler_caches",
"onediff_compiler_backend": "oneflow",
"sd_model_checkpoint": "v1-5-pruned-emaonly.safetensors",
"sd_checkpoint_hash": "6ce0161689b3853acaa03779ec93eafe75a02f4ced659bee03f50797806fa2fa",
"control_net_detectedmap_dir": "detected_maps",
"control_net_models_path": "",
"control_net_modules_path": "",
"control_net_unit_count": 3,
"control_net_model_cache_size": 2,
"control_net_inpaint_blur_sigma": 7,
"control_net_no_detectmap": false,
"control_net_detectmap_autosaving": false,
"control_net_allow_script_control": false,
"control_net_sync_field_args": true,
"controlnet_show_batch_images_in_ui": false,
"controlnet_increment_seed_during_batch": false,
"controlnet_disable_openpose_edit": false,
"controlnet_disable_photopea_edit": false,
"controlnet_photopea_warning": true,
"controlnet_ignore_noninpaint_mask": false,
"controlnet_clip_detector_on_cpu": false,
"controlnet_control_type_dropdown": false
}
docs/HAIRSTYLES_AVAILABLE.md
+7 -9
View File
@@ -1,6 +1,6 @@
# 可用发型清单 # 可用发型清单
> 当前共 **180 个发型**可用于换发型功能。 > 当前共 **179 个发型**可用于换发型功能（hair_flow 训练完成后为 180 个）。
> 数据来源：`/api/hairstyles` 接口实时返回，要求发型同时具备「LoRA权重 + ref材质 + config.json」三件齐全。 > 数据来源：`/api/hairstyles` 接口实时返回，要求发型同时具备「LoRA权重 + ref材质 + config.json」三件齐全。
> 全量清单见同目录 `hairstyles_available.csv`。 > 全量清单见同目录 `hairstyles_available.csv`。
@@ -10,11 +10,11 @@
| 类别 | 数量 | 说明 | | 类别 | 数量 | 说明 |
|------|------|------| |------|------|------|
| **总可用发型** | **180** | 三件齐全，可直接换发型 | | **总可用发型** | **179** | 三件齐全，可直接换发型 |
| ├ 男款 (boy) | 93 | 适合男性脸型 | | ├ 男款 (boy) | 93 | 适合男性脸型 |
| ├ 女款 (girl) | 87 | 适合女性脸型（+hair_flow） | | ├ 女款 (girl) | 86 | 适合女性脸型 |
| ├ 中文名命名 | 7 | 按脸型训练的新发型 | | ├ 中文名命名 | 7 | 本次按脸型训练的新发型 |
| ├ 英文名命名 | 16 | 历史测试发型 + hair_flow | | ├ 英文名命名 | 15 | 历史测试发型 |
| └ 数字ID命名 | 157 | 生产环境发型（雪花算法ID） | | └ 数字ID命名 | 157 | 生产环境发型（雪花算法ID） |
**性别说明**：gender 字段决定发型被分到"男款/女款"，但换发型时**男女发型都能用**（只是预览图用对应性别的标准脸）。 **性别说明**：gender 字段决定发型被分到"男款/女款"，但换发型时**男女发型都能用**（只是预览图用对应性别的标准脸）。
@@ -55,7 +55,6 @@
| hair_id | gender | 备注 | | hair_id | gender | 备注 |
|---------|--------|------| |---------|--------|------|
| **hair_flow** | **girl** | **波浪卷发，最新训练（2026-07-07）** |
| huaban1 | girl | 画板1 | | huaban1 | girl | 画板1 |
| huaban2 | girl | 画板2 | | huaban2 | girl | 画板2 |
| new_test_001 | boy | 测试 | | new_test_001 | boy | 测试 |
@@ -105,9 +104,8 @@
### 2. 调试页（带可视化） ### 2. 调试页（带可视化）
- `http://localhost:8888/debug` — 全参数调试台 - `http://localhost:8888/debug` — 全参数调试台
- `http://localhost:8888/manual` — 手绘 mask 换发型（含 enhance 二次增强 + 羽化贴回） - `http://localhost:8888/manual` — 手绘 mask 重绘
- `http://localhost:8888/inpaint` — **mask区域重绘**（发丝自然化，enhance/pure_inpaint两种模式） - `http://localhost:8888/hairline` — 发际线 mask 实验
- `http://localhost:8888/hairline` — 发际线 mask 实验（4种方案对比）
### 3. API 调用 ### 3. API 调用
```bash ```bash
docs/hairstyles_available.csv
+180 -181
View File
@@ -1,181 +1,180 @@
hair_id,gender hair_id,gender,预览图,LoRA大小MB
1905785164224868354,girl 1905785164224868354,girl,有,144
1907641420518580226,boy 1907641420518580226,boy,有,144
1907641442341543937,boy 1907641442341543937,boy,有,144
1907641528173780994,boy 1907641528173780994,boy,有,144
1907650883958923266,boy 1907650883958923266,boy,有,144
1907650924350070785,boy 1907650924350070785,boy,有,144
1907651105032298498,boy 1907651105032298498,boy,有,144
1907651244006367233,boy 1907651244006367233,boy,有,144
1907651308460236802,girl 1907651308460236802,girl,有,144
1907651348528422914,girl 1907651348528422914,girl,有,144
1907651413477220354,girl 1907651413477220354,girl,有,144
1907651464794529794,girl 1907651464794529794,girl,有,144
1907651525024735234,girl 1907651525024735234,girl,有,144
1907651578753769473,girl 1907651578753769473,girl,有,144
1907651632629604354,girl 1907651632629604354,girl,有,144
1907651680352395265,girl 1907651680352395265,girl,有,144
1907651747310264321,girl 1907651747310264321,girl,有,144
1907805226322419713,boy 1907805226322419713,boy,有,144
1929905737381048322,boy 1929905737381048322,boy,有,144
1929905770088230914,boy 1929905770088230914,boy,有,144
1929906202869100546,boy 1929906202869100546,boy,有,144
1929906230224351234,boy 1929906230224351234,boy,有,144
1929906315553271809,boy 1929906315553271809,boy,有,144
1929906349401305089,boy 1929906349401305089,boy,有,144
1929906429839667202,boy 1929906429839667202,boy,有,144
1935354779623858177,boy 1935354779623858177,boy,有,144
1935354812175851522,boy 1935354812175851522,boy,有,144
1935354869738479617,boy 1935354869738479617,boy,有,144
1935354938051108865,boy 1935354938051108865,boy,有,144
1935355030892027905,boy 1935355030892027905,boy,有,144
1935355149217538050,girl 1935355149217538050,girl,有,144
1935355188652384257,girl 1935355188652384257,girl,有,144
1935355258143612929,girl 1935355258143612929,girl,有,144
1935355292960530433,girl 1935355292960530433,girl,有,144
1935355361726144514,girl 1935355361726144514,girl,有,144
1935355403207811074,girl 1935355403207811074,girl,有,144
1935355449991077889,girl 1935355449991077889,girl,有,144
1935355504156319745,girl 1935355504156319745,girl,有,144
1935355551300296705,girl 1935355551300296705,girl,有,144
1935355607394918402,girl 1935355607394918402,girl,有,144
1935355698319040514,girl 1935355698319040514,girl,有,144
1935355740559876097,girl 1935355740559876097,girl,有,144
1935355778526715905,girl 1935355778526715905,girl,有,144
1935355868347736065,girl 1935355868347736065,girl,有,144
1935355934743568386,girl 1935355934743568386,girl,有,144
1935355969141055490,girl 1935355969141055490,girl,有,144
1935356042201636865,girl 1935356042201636865,girl,有,144
1935356092684279809,girl 1935356092684279809,girl,有,144
1935356131338985473,girl 1935356131338985473,girl,有,144
1935356178227109890,girl 1935356178227109890,girl,有,144
1938489620162781185,boy 1938489620162781185,boy,有,144
1938489674713899010,boy 1938489674713899010,boy,有,144
1938489739088076801,boy 1938489739088076801,boy,有,144
1938489840707674114,boy 1938489840707674114,boy,有,144
1938489915521474561,boy 1938489915521474561,boy,有,144
1938489994072399874,boy 1938489994072399874,boy,有,144
1938490058547240961,boy 1938490058547240961,boy,有,144
1938492790033195009,girl 1938492790033195009,girl,有,144
1938492830059438081,girl 1938492830059438081,girl,有,144
1938841049238970370,girl 1938841049238970370,girl,有,144
1938848709308948482,girl 1938848709308948482,girl,有,144
1938848730150445058,girl 1938848730150445058,girl,有,144
1938848782956732417,girl 1938848782956732417,girl,有,144
1939330299344560130,girl 1939330299344560130,girl,有,144
1939339435616608258,girl 1939339435616608258,girl,有,144
1939351219396255745,boy 1939351219396255745,boy,有,144
1939351238526476289,boy 1939351238526476289,boy,有,144
1939351319849836545,boy 1939351319849836545,boy,有,144
1939351392453238785,girl 1939351392453238785,girl,有,144
1939351512506802177,girl 1939351512506802177,girl,有,144
1940007538260287490,boy 1940007538260287490,boy,有,144
1941163289636872193,boy 1941163289636872193,boy,有,144
1941163324332154882,boy 1941163324332154882,boy,有,144
1941163373069967362,boy 1941163373069967362,boy,有,144
1941163615962112002,boy 1941163615962112002,boy,有,144
1941163645464846338,boy 1941163645464846338,boy,有,144
1941163754357366785,boy 1941163754357366785,boy,有,144
1941163816533729282,boy 1941163816533729282,boy,有,144
1941163884078800898,boy 1941163884078800898,boy,有,144
1941163952840220673,boy 1941163952840220673,boy,有,144
1941164012827156482,boy 1941164012827156482,boy,有,144
1941352458711572482,boy 1941352458711572482,boy,有,144
1941352541028982786,boy 1941352541028982786,boy,有,144
1941352599665352705,boy 1941352599665352705,boy,有,144
1943143930804936705,boy 1943143930804936705,boy,有,144
1943144007388733441,boy 1943144007388733441,boy,有,144
1945789428732895233,boy 1945789428732895233,boy,有,144
1953267161121464322,boy 1953267161121464322,boy,有,144
1953267288284372994,boy 1953267288284372994,boy,有,144
1953270137953230849,girl 1953270137953230849,girl,有,144
1953270523594317826,girl 1953270523594317826,girl,有,144
1953270724765720577,boy 1953270724765720577,boy,有,144
1953271922151432194,boy 1953271922151432194,boy,有,144
1953272913265467394,boy 1953272913265467394,boy,有,144
1953288800932503554,boy 1953288800932503554,boy,有,144
1954536242109784066,boy 1954536242109784066,boy,有,144
1954536297776586754,boy 1954536297776586754,boy,有,144
1954536402516746241,boy 1954536402516746241,boy,有,144
1954536893304840194,boy 1954536893304840194,boy,有,144
1954536956919848962,boy 1954536956919848962,boy,有,144
1954537022875279362,boy 1954537022875279362,boy,有,144
1954537085882114050,boy 1954537085882114050,boy,有,144
1954537191729569793,boy 1954537191729569793,boy,有,144
1954537709436706817,girl 1954537709436706817,girl,有,144
1954537743532204034,girl 1954537743532204034,girl,有,144
1954537926731014146,girl 1954537926731014146,girl,有,144
1954537961069780994,girl 1954537961069780994,girl,有,144
1954538119216013313,girl 1954538119216013313,girl,有,144
1954538156499181569,girl 1954538156499181569,girl,有,144
1954538259330932738,girl 1954538259330932738,girl,有,144
1954538335306555393,girl 1954538335306555393,girl,有,144
1980312945301954561,boy 1980312945301954561,boy,有,144
1980313052302843905,girl 1980313052302843905,girl,有,144
1980313166576656385,girl 1980313166576656385,girl,有,144
1980313210109337601,girl 1980313210109337601,girl,有,144
1980313433502162946,boy 1980313433502162946,boy,有,144
1980313576506957825,girl 1980313576506957825,girl,有,144
1980313740403580930,girl 1980313740403580930,girl,有,144
1980313819915001858,girl 1980313819915001858,girl,有,144
1980314032369082369,girl 1980314032369082369,girl,有,144
1980314146483511297,girl 1980314146483511297,girl,有,144
1980314209322573825,girl 1980314209322573825,girl,有,144
1980314290419441665,girl 1980314290419441665,girl,有,144
1980314381058351106,girl 1980314381058351106,girl,有,144
1980314477720281090,girl 1980314477720281090,girl,有,144
1980314710101499906,boy 1980314710101499906,boy,有,144
1980314764421931009,boy 1980314764421931009,boy,有,144
1980314868713299969,boy 1980314868713299969,boy,有,144
1980315054516772866,boy 1980315054516772866,boy,有,144
1980315179934851074,girl 1980315179934851074,girl,有,144
1980315238260842497,girl 1980315238260842497,girl,有,144
1980361636742209537,girl 1980361636742209537,girl,有,144
1980367742952579074,boy 1980367742952579074,boy,有,144
1980367781510815746,boy 1980367781510815746,boy,有,144
1980367800787836929,boy 1980367800787836929,boy,有,144
1980367852847538177,boy 1980367852847538177,boy,有,144
1980367968627105793,boy 1980367968627105793,boy,有,144
1980368109866098690,boy 1980368109866098690,boy,有,144
1980368205752082434,boy 1980368205752082434,boy,有,144
1980368368847593473,boy 1980368368847593473,boy,有,144
1980368465631158274,boy 1980368465631158274,boy,有,144
1980368588293578754,boy 1980368588293578754,boy,有,144
1980368787732733953,boy 1980368787732733953,boy,有,144
1980368837162606594,boy 1980368837162606594,boy,有,144
1980368937435832321,boy 1980368937435832321,boy,有,144
1980413320612847618,girl 1980413320612847618,girl,有,144
1980415346071609345,girl 1980415346071609345,girl,有,144
1980415897056354306,girl 1980415897056354306,girl,有,144
1980422351658196994,boy 1980422351658196994,boy,有,144
1980422373388886018,boy 1980422373388886018,boy,有,144
1981177324700549122,girl 1981177324700549122,girl,有,144
1981177378714796034,boy 1981177378714796034,boy,有,144
1981177477423546370,boy 1981177477423546370,boy,有,144
1981177594750812162,boy 1981177594750812162,boy,有,144
1981177742742634497,boy 1981177742742634497,boy,有,144
1981177931859607554,boy 1981177931859607554,boy,有,144
3325266,girl 3325266,girl,有,144
hair_flow,girl huaban1,girl,有,144
huaban1,girl huaban2,girl,有,144
huaban2,girl new_test_001,boy,有,144
new_test_001,boy newf_1214,girl,有,144
newf_1214,girl newf_12141,girl,有,144
newf_12141,girl newf_121445553,boy,有,144
newf_121445553,boy newf_12314,girl,有,144
newf_12314,girl newf_1235314,girl,有,144
newf_1235314,girl newf_204,girl,有,144
newf_204,girl newf_216,girl,有,144
newf_216,girl test001,boy,有,144
test001,boy test002,girl,有,144
test002,girl test003,boy,有,144
test003,boy test0036,boy,有,144
test0036,boy test888,boy,有,144
test888,boy 圆-心形,girl,有,144
圆-心形,girl 圆-椭圆,girl,有,144
圆-椭圆,girl 圆-波浪,girl,有,144
圆-波浪,girl 圆-直线,girl,有,144
圆-直线,girl 圆-花瓣,girl,有,144
圆-花瓣,girl 心形-心形,girl,有,144
心形-心形,girl 心形-椭圆,girl,有,144
心形-椭圆,girl
1 hair_id gender 预览图 LoRA大小MB
2 1905785164224868354 girl 有 144
3 1907641420518580226 boy 有 144
4 1907641442341543937 boy 有 144
5 1907641528173780994 boy 有 144
6 1907650883958923266 boy 有 144
7 1907650924350070785 boy 有 144
8 1907651105032298498 boy 有 144
9 1907651244006367233 boy 有 144
10 1907651308460236802 girl 有 144
11 1907651348528422914 girl 有 144
12 1907651413477220354 girl 有 144
13 1907651464794529794 girl 有 144
14 1907651525024735234 girl 有 144
15 1907651578753769473 girl 有 144
16 1907651632629604354 girl 有 144
17 1907651680352395265 girl 有 144
18 1907651747310264321 girl 有 144
19 1907805226322419713 boy 有 144
20 1929905737381048322 boy 有 144
21 1929905770088230914 boy 有 144
22 1929906202869100546 boy 有 144
23 1929906230224351234 boy 有 144
24 1929906315553271809 boy 有 144
25 1929906349401305089 boy 有 144
26 1929906429839667202 boy 有 144
27 1935354779623858177 boy 有 144
28 1935354812175851522 boy 有 144
29 1935354869738479617 boy 有 144
30 1935354938051108865 boy 有 144
31 1935355030892027905 boy 有 144
32 1935355149217538050 girl 有 144
33 1935355188652384257 girl 有 144
34 1935355258143612929 girl 有 144
35 1935355292960530433 girl 有 144
36 1935355361726144514 girl 有 144
37 1935355403207811074 girl 有 144
38 1935355449991077889 girl 有 144
39 1935355504156319745 girl 有 144
40 1935355551300296705 girl 有 144
41 1935355607394918402 girl 有 144
42 1935355698319040514 girl 有 144
43 1935355740559876097 girl 有 144
44 1935355778526715905 girl 有 144
45 1935355868347736065 girl 有 144
46 1935355934743568386 girl 有 144
47 1935355969141055490 girl 有 144
48 1935356042201636865 girl 有 144
49 1935356092684279809 girl 有 144
50 1935356131338985473 girl 有 144
51 1935356178227109890 girl 有 144
52 1938489620162781185 boy 有 144
53 1938489674713899010 boy 有 144
54 1938489739088076801 boy 有 144
55 1938489840707674114 boy 有 144
56 1938489915521474561 boy 有 144
57 1938489994072399874 boy 有 144
58 1938490058547240961 boy 有 144
59 1938492790033195009 girl 有 144
60 1938492830059438081 girl 有 144
61 1938841049238970370 girl 有 144
62 1938848709308948482 girl 有 144
63 1938848730150445058 girl 有 144
64 1938848782956732417 girl 有 144
65 1939330299344560130 girl 有 144
66 1939339435616608258 girl 有 144
67 1939351219396255745 boy 有 144
68 1939351238526476289 boy 有 144
69 1939351319849836545 boy 有 144
70 1939351392453238785 girl 有 144
71 1939351512506802177 girl 有 144
72 1940007538260287490 boy 有 144
73 1941163289636872193 boy 有 144
74 1941163324332154882 boy 有 144
75 1941163373069967362 boy 有 144
76 1941163615962112002 boy 有 144
77 1941163645464846338 boy 有 144
78 1941163754357366785 boy 有 144
79 1941163816533729282 boy 有 144
80 1941163884078800898 boy 有 144
81 1941163952840220673 boy 有 144
82 1941164012827156482 boy 有 144
83 1941352458711572482 boy 有 144
84 1941352541028982786 boy 有 144
85 1941352599665352705 boy 有 144
86 1943143930804936705 boy 有 144
87 1943144007388733441 boy 有 144
88 1945789428732895233 boy 有 144
89 1953267161121464322 boy 有 144
90 1953267288284372994 boy 有 144
91 1953270137953230849 girl 有 144
92 1953270523594317826 girl 有 144
93 1953270724765720577 boy 有 144
94 1953271922151432194 boy 有 144
95 1953272913265467394 boy 有 144
96 1953288800932503554 boy 有 144
97 1954536242109784066 boy 有 144
98 1954536297776586754 boy 有 144
99 1954536402516746241 boy 有 144
100 1954536893304840194 boy 有 144
101 1954536956919848962 boy 有 144
102 1954537022875279362 boy 有 144
103 1954537085882114050 boy 有 144
104 1954537191729569793 boy 有 144
105 1954537709436706817 girl 有 144
106 1954537743532204034 girl 有 144
107 1954537926731014146 girl 有 144
108 1954537961069780994 girl 有 144
109 1954538119216013313 girl 有 144
110 1954538156499181569 girl 有 144
111 1954538259330932738 girl 有 144
112 1954538335306555393 girl 有 144
113 1980312945301954561 boy 有 144
114 1980313052302843905 girl 有 144
115 1980313166576656385 girl 有 144
116 1980313210109337601 girl 有 144
117 1980313433502162946 boy 有 144
118 1980313576506957825 girl 有 144
119 1980313740403580930 girl 有 144
120 1980313819915001858 girl 有 144
121 1980314032369082369 girl 有 144
122 1980314146483511297 girl 有 144
123 1980314209322573825 girl 有 144
124 1980314290419441665 girl 有 144
125 1980314381058351106 girl 有 144
126 1980314477720281090 girl 有 144
127 1980314710101499906 boy 有 144
128 1980314764421931009 boy 有 144
129 1980314868713299969 boy 有 144
130 1980315054516772866 boy 有 144
131 1980315179934851074 girl 有 144
132 1980315238260842497 girl 有 144
133 1980361636742209537 girl 有 144
134 1980367742952579074 boy 有 144
135 1980367781510815746 boy 有 144
136 1980367800787836929 boy 有 144
137 1980367852847538177 boy 有 144
138 1980367968627105793 boy 有 144
139 1980368109866098690 boy 有 144
140 1980368205752082434 boy 有 144
141 1980368368847593473 boy 有 144
142 1980368465631158274 boy 有 144
143 1980368588293578754 boy 有 144
144 1980368787732733953 boy 有 144
145 1980368837162606594 boy 有 144
146 1980368937435832321 boy 有 144
147 1980413320612847618 girl 有 144
148 1980415346071609345 girl 有 144
149 1980415897056354306 girl 有 144
150 1980422351658196994 boy 有 144
151 1980422373388886018 boy 有 144
152 1981177324700549122 girl 有 144
153 1981177378714796034 boy 有 144
154 1981177477423546370 boy 有 144
155 1981177594750812162 boy 有 144
156 1981177742742634497 boy 有 144
157 1981177931859607554 boy 有 144
158 3325266 girl 有 144
159 hair_flow huaban1 girl 有 144
160 huaban1 huaban2 girl 有 144
161 huaban2 new_test_001 girl boy 有 144
162 new_test_001 newf_1214 boy girl 有 144
163 newf_1214 newf_12141 girl 有 144
164 newf_12141 newf_121445553 girl boy 有 144
165 newf_121445553 newf_12314 boy girl 有 144
166 newf_12314 newf_1235314 girl 有 144
167 newf_1235314 newf_204 girl 有 144
168 newf_204 newf_216 girl 有 144
169 newf_216 test001 girl boy 有 144
170 test001 test002 boy girl 有 144
171 test002 test003 girl boy 有 144
172 test003 test0036 boy 有 144
173 test0036 test888 boy 有 144
174 test888 圆-心形 boy girl 有 144
175 圆-心形 圆-椭圆 girl 有 144
176 圆-椭圆 圆-波浪 girl 有 144
177 圆-波浪 圆-直线 girl 有 144
178 圆-直线 圆-花瓣 girl 有 144
179 圆-花瓣 心形-心形 girl 有 144
180 心形-心形 心形-椭圆 girl 有 144
心形-椭圆 girl
docs/换发型集成文档.md
+408
View File
@@ -0,0 +1,408 @@
# 换发型功能集成文档
本文档说明"换发型"功能的完整工作流程、服务依赖、资源结构，供其他项目集成参考。
---
## 一、功能概述
**换发型**：用户上传一张人像照片，选择一个目标发型（hair_id），系统把用户的头发替换成目标发型，保留人脸和身体。
- **输入**：用户人像图（base64 或 URL） + 目标发型 ID
- **输出**：换好发型的结果图（base64 或 OSS URL）
- **耗时**：单次约 10-12 秒（GPU: RTX 5090）
### 同类功能（同套服务支撑）
- **换发色**（`/hairColor/v2`）：把头发颜色换成指定 RGB，额外依赖 `ref_color` 颜色查找表
- **训练新发型**：用发型图片训练新的 LoRA，详见第七节
---
## 二、服务架构
换发型由 **3 个后端服务**协同完成（缺一不可），另有 1 个前端页面服务（可选）。
| 服务 | 端口 | conda 环境 | 职责 |
|------|------|-----------|------|
| **webui** | 57860 | `sdwebui` | Stable Diffusion img2img 推理引擎（基于 AUTOMATIC1111 webui） |
| **photo_service** | 32678 | `py310` | LoRA 调度：推理时把对应发型的 LoRA 复制到 webui 并注入 prompt；训练时调度 kohya |
| **hair_service_sd** | 8801 | `my_hair` | 换发型主服务：人脸检测、关键点、头发抠图、材质推理、图像合成 |
| hair_grow_service | 8899 | `my_hair` | （可选）前端测试页面，提供发型列表/预览/换发型 UI |
### 服务间调用链（换发型）
```
调用方
│ POST /api/swapHair/v1
▼
hair_service_sd (8801) ← 人脸检测/关键点/抠图/材质推理/图像合成
│ POST /api/hair/inference (转发 img2img 请求)
▼
photo_service (32678) ← 复制 LoRA 到 webui、prompt 注入 <lora:xxx:1.0>
│ POST /sdapi/v1/img2img
▼
webui (57860) ← SD img2img 真正推理
```
### 启动方式
```bash
cd /home/xsl/change_hair
bash start_all.sh start # 启动 webui + photo_service + hair_service_sd
# 等待 1-2 分钟（webui 加载 SD 模型）
# 可选：前端页面
nohup bash start_hairgrow.sh > project/logs/hair_grow_service.log 2>&1 &
```
> 三个核心服务有启动顺序依赖：webui 必须先启动并加载完模型（约 1-2 分钟），photo_service 和 hair_service_sd 才能正常工作。
---
## 三、换发型 API 接口
### 请求
```
POST http://<host>:8801/api/swapHair/v1
Content-Type: application/json
```
**请求体参数：**
| 参数 | 类型 | 必填 | 说明 |
|------|------|:---:|------|
| `hair_id` | string | ✓ | 目标发型 ID（如 `chang_tuoyuan`），对应发型资源目录名 |
| `task_id` | string | ✓ | 任务唯一标识，用于临时文件命名 |
| `user_img_path` | string | ✓ | 用户人像图。支持 `data:image/jpeg;base64,<base64>` 或图片 URL |
| `is_hr` | string | ✓ | 是否高清模式。`"true"` 输出 1152×1536，`"false"` 输出 576×768 |
| `output_format` | string | ✗ | 输出格式。`"base64"`（推荐，直接返回图）或 `"url"`（上传 OSS 返回链接，需配 OSS 凭证） |
**请求示例：**
```bash
curl -X POST http://127.0.0.1:8801/api/swapHair/v1 \
-H "Content-Type: application/json" \
-d '{
"hair_id": "chang_tuoyuan",
"task_id": "req_001",
"is_hr": "true",
"user_img_path": "data:image/jpeg;base64,/9j/4AAQ...",
"output_format": "base64"
}'
```
```python
# Python 示例
import base64, requests
with open("user.jpg", "rb") as f:
img_b64 = base64.b64encode(f.read()).decode()
r = requests.post("http://127.0.0.1:8801/api/swapHair/v1", json={
"hair_id": "chang_tuoyuan",
"task_id": "req_001",
"is_hr": "true",
"user_img_path": "data:image/jpeg;base64," + img_b64,
"output_format": "base64",
}, timeout=300)
result = r.json() # {"state":0, "data":"<base64>", "msg":"success"}
with open("out.jpg", "wb") as f:
f.write(base64.b64decode(result["data"]))
```
### 响应
**成功（HTTP 200）：**
```json
{
"state": 0,
"msg": "success",
"data": "<base64 编码的结果图>" // output_format=url 时为 OSS 图片链接
}
```
**失败（HTTP 400）：**
```json
{
"state": -1,
"msg": "推理发型失败", // 具体错误信息，如：参数错误/用户图像加载失败/发型模板加载失败/推理发型失败
"data": "",
"task_id": "req_001"
}
```
---
## 四、换发型内部流程（9 个功能步骤）
理解内部流程有助于排查问题和性能优化。`hair_service_sd` 收到请求后依次执行：
| 步骤 | 功能 | 说明 | 耗时占比 |
|------|------|------|---------|
| 1 | 获取请求参数 | 解析参数，base64/URL 图片解码落盘 | <1% |
| 2 | 加载用户图像 | 复制到工作目录 | <1% |
| 3 | 加载发型模板图像 | 从 `upload_train_imgs/<hair_id>/` 读模板原图 | <1% |
| 4 | 检查遮挡眼睛 | 从 `hair_template_material/<hair_id>/` 读 `*_matting.png` 和 `.pkl`（关键点） | <1% |
| 5 | 检查发型材质 | 确认 `ref_hairstyle/<hair_id>/` 存在 | <1% |
| 6 | **推理发型** | 用 GAN 模型把用户图和 `ref_rgb_8uc3_768.png` 融合，生成初步换发图 | ~30% |
| 7 | 判断遮挡眼睛 | 检查头发是否遮挡眼睛（决定后续处理路径） | <1% |
| 8 | 处理发型区域 | 头发区域裁剪、mask 生成、透视变换到 576×768/1152×1536 | ~5% |
| 9 | **webui img2img** | 经 photo_service 调 webui，加载 LoRA 做 SD 重绘（核心画质步骤） | ~60% |
**关键：LoRA 加载机制**（步骤 9）
1. `hair_service_sd` 把裁剪好的头发区域图 + mask 发给 `photo_service`
2. `photo_service` 把 `train_material/<hair_id>/model/hairstyle_hd_lora.safetensors` 复制到 webui 的 `models/Lora/<hair_id>_hd.safetensors`（已存在则跳过）
3. `photo_service` 在 prompt 前注入 `<lora:<hair_id>_hd:1.0>`
4. webui 收到 img2img 请求，自动加载该 LoRA 做局部重绘
5. 结果返回 `hair_service_sd`，贴回原图对应位置
---
## 五、必需资源清单
换发型功能依赖以下资源。集成时需保证这些文件到位。
### 5.1 核心模型权重（全局共享，约 10G）
位置：`project/hair_service_sd/weights/`
| 文件/目录 | 用途 |
|-----------|------|
| `yolov5l.pt` | RetinaFace/人脸检测 |
| `Resnet50_Final.pth` | 人脸检测辅助 |
| `79999_iter.pth` | 关键点检测（MomocvFaceAlignment1K） |
| `master_hair_v8_onlyhair_nowarp_all_0709.pt` | 头发迁移 GAN（发型融合） |
| `zxm_v7_gen_bald_add_changed_hair_pairdata_5seg_0611.pt` | 秃头分割 + 发型生成 GAN |
| `parser_model.pth` | 头发抠图（matting） |
| `gender_models/` | 性别识别 |
| `keypoints/` | 关键点模型库 |
### 5.2 Stable Diffusion 推理资源（约 5G）
位置：`project/onediff/stable-diffusion-webui/models/`
| 文件/目录 | 用途 |
|-----------|------|
| `Stable-diffusion/v1-5-pruned-emaonly.safetensors` | **SD 1.5 底模**（必需，4G） |
| `Lora/<hair_id>_hd.safetensors` | 各发型的 LoRA（推理时自动生成，见 5.3） |
| `ESRGAN/4x-UltraSharp.pth` | 超分辨率（高清模式用） |
| `GFPGAN/` | 人脸修复（可选） |
### 5.3 单个发型的资源（每个发型约 160M）
**一个可用的发型需要 4 个目录的资源配套：**
```
# 1. 发型材质（换发型必需，由训练 step3 生成）
project/data/ref_hairstyle/<hair_id>/
├── config.json # {"gender":"girl","version":"20221206","ratio":"1"}
├── ref_rgb_8uc3_768.png # 发型 RGB 图（GAN 融合参考图，必需）
├── ref_matting_8uc3_768.png # 发型 mask
├── ref_matting_fg_8uc3_768.png # 发型前景
├── ref_baldseg_8uc3_768.png # 秃头分割
├── ref_landmark_f1k2_768.txt # 1000 个关键点
└── input_another_pose_hair_image.npy
# 2. LoRA 权重 + 训练图（换发型必需，由训练 step1+step2 生成）
data/train_material/<hair_id>/
├── model/
│ └── hairstyle_hd_lora.safetensors # LoRA 权重（~145M，SD 重绘用）
└── images/
└── 1_hairstyle/ # 训练图 + caption
├── 0001.png
├── 0001.txt # caption tag（webui 推理 prompt 用）
└── ... (约20张)
# 3. 模板材质（换发型必需，由训练 step4 生成）
project/data/hair_template_material/<hair_id>/
├── first##<uuid>.png # 发型模板原图
├── first##<uuid>_matting.png # 发型 matting（遮挡检测用）
└── first##<uuid>.pkl # 关键点（遮挡检测用）
# 4. 发型模板原图（换发型必需，由训练准备）
project/data/upload_train_imgs/<hair_id>/
└── first##<hair_id>.jpg # 发型模板原图（步骤3读取）
```
> **注意**：`images/1_hairstyle/` 目录名必须以数字开头（`1_hairstyle`），代码用 `os.listdir(images)[0]` 取子目录并校验数字开头。
### 5.4 路径配置
所有路径在 `project/hair_service_sd/config/configure.ini` 集中配置，迁移时改这里的路径即可：
```ini
[default]
modelDir = weights # 核心模型目录（相对 hair_service_sd）
hairstyleDir = /home/xsl/change_hair/project/data/ref_hairstyle # 发型材质（5.3-1）
haircolorDir= /home/xsl/change_hair/project/data/ref_haircolor # 换发色用
userDir=/home/xsl/change_hair/project/data/userImage # 用户图工作目录
tmp_dir=/home/xsl/change_hair/project/data/tmp # 临时目录
res_dir=/home/xsl/change_hair/project/data/res_dir # 结果输出目录
userInfo_dir=/home/xsl/change_hair/project/data/user_info # 用户关键点缓存
train_dir = /home/xsl/change_hair/data/train_material # LoRA+训练图（5.3-2）
refImgDir=/home/xsl/change_hair/project/data/refImage # 参考图缓存
hair_template_material_dir=/home/xsl/change_hair/project/data/hair_template_material # 5.3-3
upload_train_dir=/home/xsl/change_hair/project/data/upload_train_imgs # 5.3-4
version=online # online 用 0.0.0.0 地址
```
`photo_service` 的 webui/LoRA 路径在 `project/photo_service/lora_train_service_1.py` 第 41 行：
```python
webui_lora_dir = '/home/xsl/change_hair/project/onediff/stable-diffusion-webui/models/Lora'
```
---
## 六、集成步骤（接入新项目）
### 1. 部署服务端
按第二节启动 webui + photo_service + hair_service_sd 三个服务，确认 `http://<host>:8801` 可达。
### 2. 调用换发型接口
新项目只需调 `hair_service_sd` 的 `/api/swapHair/v1`，无需关心内部服务调用。
### 3. 准备发型资源
每个要支持的发型，必须备齐第五节 5.3 的 4 个目录资源。两种获取方式：
- **训练新发型**（见第七节）：用发型图训练，自动生成全部资源
- **复制现有发型**：从已部署环境复制 `<hair_id>` 的 4 个目录
### 4. 查询可用发型列表
调 `hair_grow_service` 的接口（需启动 8899）：
```
GET http://<host>:8899/api/hairstyles
→ {"state":0, "count":5, "data":[{"hair_id":"chang_tuoyuan","gender":"girl"}, ...]}
```
> 注意：该接口在 `hair_grow_swap.py:list_hairstyles()` 内有白名单过滤（当前只返回 5 个 chang_* 发型），集成时按需调整或去掉。
### 5. 获取发型预览图
```
GET http://<host>:8899/preview/<hair_id>
→ 直接返回预览图（优先 static/previews/<hair_id>.jpg，回退到 ref_rgb_8uc3_768.png）
```
---
## 七、训练新发型
### 7.1 一键训练脚本
```bash
cd /home/xsl/change_hair/project/hair_service_sd
python /home/xsl/change_hair/train_hairstyle_full.py \
--hair-id <新发型ID> \
--input <训练图目录> \ # 目录里放该发型的 jpg/png（1张即可，多张效果更好）
--gender <girl|boy> \
--template-img <模板图> # 用于生成材质和预览，通常是该发型最代表性的图
```
### 7.2 训练 5 个阶段
| 阶段 | 功能 | 产物 | 耗时 |
|------|------|------|------|
| step1 | 准备训练数据 | 抠图 + 数据增广（水平翻转 × 5 分辨率）→ `train_material/<hair_id>/images/1_hairstyle/` 约 20 张 png+txt | ~1 分钟 |
| step2 | LoRA 训练 | kohya 训练 1500 步 → `train_material/<hair_id>/model/hairstyle_hd_lora.safetensors` | ~10 分钟（单卡） |
| step3 | 生成 ref 材质 | 调 hair_service_sd 回调 → `ref_hairstyle/<hair_id>/` 6 个文件 | ~1 秒 |
| step4 | 生成模板材质 | → `hair_template_material/<hair_id>/` 3 个文件 | ~3 秒 |
| step5 | 生成预览图 | → `hair_grow_service/static/previews/<hair_id>.jpg` | ~10 秒 |
### 7.3 批量并行训练（多个发型）
用 `train_lora_parallel.py` 直接调 kohya，支持并发：
```python
# 编辑 train_lora_parallel.py 顶部的 HAIRSTYLES 列表
HAIRSTYLES = ["新发型1", "新发型2", "新发型3"]
PARALLEL = 2 # 并发数（按显存调整，5090 32G 建议 2）
# step1 数据需先用 train_hairstyles_parallel.py 或逐个跑 step1 准备好
python train_lora_parallel.py
```
### 7.4 训练注意事项
- **hair_id 命名**：只能用英文/数字/下划线，不能含中文或特殊字符（会作目录名、文件名、LoRA 名）
- **单图训练**：1 张图经增广出 ~20 张训练图，LoRA 能用但泛化有限；多图（不同角度/光照）效果更好
- **GPU 显存**：训练峰值约 13G/卡，推理约 6G。并行训练需控制并发数避免 OOM
- **训练前提**：webui + photo_service + hair_service_sd 必须在运行（step3 回调依赖）
---
## 八、环境依赖
### Conda 环境
| 环境 | Python | 用途 | 关键依赖 |
|------|--------|------|---------|
| `my_hair` | 3.10 | hair_service_sd / hair_grow_service / 训练 step1 | torch, cv2, flask, gevent, oss2 |
| `sdwebui` | 3.10 | webui | torch, gradio, transformers（AUTOMATIC1111 全套） |
| `py310` | 3.10 | photo_service | torch, flask, requests |
| `kohya` | 3.10 | LoRA 训练（kohya_ss） | torch, accelerate, safetensors |
### 关键兼容性处理（迁移时注意）
1. **PyTorch 2.6+ 的 `weights_only`**：加载旧权重（yolov5l.pt 等）需 `torch.load(..., weights_only=False)`。代码已在 `run_copy_cost_colorb64.py` 和 `hair_grow_service/app.py` 顶部做了全局 patch。
2. **OSS 上传（可选）**：`output_format=url` 时需配 OSS 凭证环境变量：
```bash
export OSS_TEST_ACCESS_KEY_ID=<...>
export OSS_TEST_ACCESS_KEY_SECRET=<...>
export OSS_TEST_BUCKET=<...>
export OSS_TEST_ENDPOINT=<...>
```
用 `output_format=base64` 则完全不需要 OSS。
3. **webui xformers**：kohya 训练命令用 `--sdpa`（PyTorch 原生），不要用 `--xformers`（若 xformers 与 torch 版本不匹配会报 CUDA 错误）。
---
## 九、故障排查
| 现象 | 可能原因 | 排查 |
|------|---------|------|
| 换发型返回 `推理发型失败` | webui 没启动 / LoRA 文件缺失 | 检查 webui 57860 端口；检查 `train_material/<hair_id>/model/` 是否有 LoRA |
| 换发型返回 `发型模板加载失败` | `upload_train_imgs/<hair_id>/` 缺 `first##` 开头的文件 | 检查 5.3-4 目录 |
| 换发型返回 `检查遮挡失败` | `hair_template_material/<hair_id>/` 缺 matting/pkl | 检查 5.3-3 目录 |
| 预览图 404 | `ref_hairstyle/<hair_id>/ref_rgb_8uc3_768.png` 不存在 | 检查 5.3-1 目录 |
| 发型列表少发型 | `hair_grow_swap.py:list_hairstyles()` 白名单过滤 | 调整 `_VISIBLE_HAIRSTYLES` 集合 |
| webui 启动报 xformers 错 | xformers 与 torch 版本不匹配 | webui 启动去掉 `--xformers`，或用 `--sdpa` |
| 训练报 CUDA error | 显存不足或 xformers 问题 | 训练命令用 `--sdpa`；降低并发数 |
### 日志位置
- webui：`project/logs/webui.log`
- photo_service：`project/logs/photo_service.log`
- hair_service_sd：`project/logs/hair_service.log`
- hair_grow_service：`project/logs/hair_grow_service.log`
- 训练：`project/logs/train_<hair_id>.log`
---
## 十、目录结构总览
```
/home/xsl/change_hair/
├── project/
│ ├── hair_service_sd/ # 换发型主服务
│ │ ├── weights/ # 核心模型权重（10G）
│ │ ├── config/configure.ini # 路径配置
│ │ ├── run_copy_cost_colorb64.py # 主入口（API 定义）
│ │ ├── gen_super_image.py # webui img2img 调用
│ │ └── core/ # 发型推理核心逻辑
│ ├── photo_service/ # LoRA 调度服务
│ │ └── lora_train_service_1.py # 训练调度 + 推理 LoRA 切换
│ ├── onediff/ # SD webui（推理引擎，14G）
│ │ └── stable-diffusion-webui/models/
│ │ ├── Stable-diffusion/ # SD 底模
│ │ └── Lora/ # 各发型 LoRA（自动生成）
│ ├── kohya_ss_home/ # LoRA 训练框架（13G）
│ ├── data/ # 数据目录（由 configure.ini 指向）
│ │ ├── ref_hairstyle/ # 发型材质（5.3-1）
│ │ ├── hair_template_material/ # 模板材质（5.3-3）
│ │ ├── upload_train_imgs/ # 发型模板原图（5.3-4）
│ │ └── ...
│ └── logs/
├── data/
│ └── train_material/ # LoRA + 训练图（5.3-2）
├── hair_grow_service/ # 前端页面（可选）
├── train_hairstyle_full.py # 单发型一键训练
├── train_lora_parallel.py # 批量并行训练
├── start_all.sh # 服务启停脚本
└── docs/换发型集成文档.md # 本文档
```
docs/方案A-上传文件清单.md
+302
View File
@@ -0,0 +1,302 @@
# 方案 A：换发型服务本机部署 — 上传文件清单
> 适用场景：从原服务器打包上传 conda 环境 + SD WebUI 代码，在本机完成换发型推理部署。
>
> 本机路径：`/home/xsl/change_hair`（GPU：RTX 5090 32G，CUDA 12.8）
>
> 原服务器参考：`szlc@192.168.101.63`
>
> 更新日期：2026-07-09（核实本机已有资源，更新上传清单）
---
## 一、本机已有资源（已到位，无需再传）✅
以下资源已确认存在于本机项目中：
| 本机路径 | 大小 | 说明 |
|---------|------|------|
| `project/hair_service_sd/weights/` | 9.9G | 核心推理权重（28个文件，含 GAN/分割/关键点/人脸等全部模型） |
| `project/onediff/stable-diffusion-webui/models/` | 4.6G | SD 底模 `v1-5-pruned-emaonly.safetensors` + ESRGAN `4x-UltraSharp.pth` + GFPGAN 全套 |
| `project/data/ref_hairstyle/` | 40M | 5 个发型材质（config.json + npy + png + txt） |
| `project/data/hair_template_material/` | 24M | 5 个发型模板（pkl + png + matting） |
| `project/data/upload_train_imgs/` | 2.5M | 5 个发型原图（jpg） |
| `data/train_material/` | 739M | 5 个发型 LoRA 权重（`hairstyle_hd_lora.safetensors`）+ 训练图像 |
**当前可用发型（5 个）：**
- `chang_bolang`（长波浪）
- `chang_huaban`（长花瓣）
- `chang_tuoyuan`（长椭圆）
- `chang_xinxing`（长心形）
- `chang_zhixian`（长直线）
> 以上发型可做换发型推理。如需更多发型，见第四节可选上传。
**项目代码（已在 git 仓库中，无需上传）：**
- `project/hair_service_sd/`（代码，不含 weights）
- `project/photo_service/`
- `hair_grow_service/`
- 启动脚本、测试脚本等
---
## 二、必须上传（换发型推理最低配置）🔴
以下 4 项需要从原服务器打包上传，**总计约 18–30G**。
### 2.1 Conda 环境 × 3
从原服务器 `/home/szlc/miniconda3/envs/` 打包，上传到本机 `/home/xsl/miniconda3/envs/`。
| 环境名 | 原服务器路径 | 本机目标路径 | 用途 | 预估大小 |
|--------|-------------|-------------|------|---------|
| `my_hair` | `/home/szlc/miniconda3/envs/my_hair/` | `/home/xsl/miniconda3/envs/my_hair/` | hair_service_sd 主服务（端口 8801） | ~5–8G |
| `sdwebui` | `/home/szlc/miniconda3/envs/sdwebui/` | `/home/xsl/miniconda3/envs/sdwebui/` | SD WebUI 推理（端口 57860） | ~8–12G |
| `py310` | `/home/szlc/miniconda3/envs/py310/` | `/home/xsl/miniconda3/envs/py310/` | photo_service LoRA 调度（端口 32678） | ~2–4G |
> ⚠️ **为什么必须上传，不能从公网安装？** 这三个环境包含特定版本的 PyTorch + CUDA 扩展、历史版本的 pip 包、以及编译好的 C++/CUDA 算子，依赖关系复杂。项目中的 `env.yaml` 是旧的 Python 3.7 规格，与实际运行环境（Python 3.10）不一致，无法用于重建。
#### 原服务器打包示例
```bash
# 在原服务器（szlc@192.168.101.63）执行
cd /home/szlc/miniconda3/envs
tar czf my_hair.tar.gz my_hair/
tar czf sdwebui.tar.gz sdwebui/
tar czf py310.tar.gz py310/
```
#### 本机解压示例
```bash
# 在本机执行（上传完成后）
mkdir -p /home/xsl/miniconda3/envs
cd /home/xsl/miniconda3/envs
tar xzf /path/to/my_hair.tar.gz
tar xzf /path/to/sdwebui.tar.gz
tar xzf /path/to/py310.tar.gz
```
---
### 2.2 SD WebUI 代码（不含 models）
从原服务器打包 **webui 程序本体**，**排除 models 目录**（本机已有 4.6G 模型权重）。
| 原服务器路径 | 本机目标路径 | 说明 | 预估大小 |
|-------------|-------------|------|---------|
| `/home/szlc/project/onediff/stable-diffusion-webui/`（排除 `models/`） | `/home/xsl/change_hair/project/onediff/stable-diffusion-webui/` | webui.py、modules/、extensions/、repositories/ 等 | ~2–5G |
> ⚠️ 本机当前 `project/onediff/stable-diffusion-webui/` 目录只有 `models/` 子目录（4.6G 模型权重），缺少 `webui.py`、`modules/`、`extensions/` 等所有代码文件。
**必须包含的关键文件/目录：**
```
stable-diffusion-webui/
├── webui.py # 主入口（必需）
├── webui.sh # 启动脚本
├── webui-user.sh # 用户配置
├── launch.py
├── modules/ # webui 核心模块
├── extensions/ # 扩展（如有）
├── repositories/ # 依赖仓库（如有）
├── config.json # 配置文件
├── ui-config.json # UI 配置（如有）
└── style.css / javascript/ # 前端资源
```
**不要上传（本机已有，勿覆盖）：**
```
models/Stable-diffusion/v1-5-pruned-emaonly.safetensors # 4.0G
models/ESRGAN/4x-UltraSharp.pth # 64M
models/GFPGAN/GFPGANv1.4.pth # 333M
models/GFPGAN/detection_Resnet50_Final.pth # 105M
models/GFPGAN/parsing_parsenet.pth # 82M
```
#### 原服务器打包示例（排除 models）
```bash
# 在原服务器（szlc@192.168.101.63）执行
cd /home/szlc/project/onediff
tar czf webui_code.tar.gz \
--exclude='stable-diffusion-webui/models' \
stable-diffusion-webui/
```
#### 本机解压示例
```bash
# 在本机执行（保留已有 models 目录）
cd /home/xsl/change_hair/project/onediff
tar xzf /path/to/webui_code.tar.gz
# 解压后确认 models/ 仍在且未被覆盖
ls -lh stable-diffusion-webui/models/Stable-diffusion/
```
---
## 三、上传清单速查表
复制以下列表，在原服务器逐项确认后打包：
```
# ===== 必须上传（方案 A，共 4 项，约 18–30G）=====
# 1. Conda 环境（3 个）
/home/szlc/miniconda3/envs/my_hair/ → /home/xsl/miniconda3/envs/my_hair/
/home/szlc/miniconda3/envs/sdwebui/ → /home/xsl/miniconda3/envs/sdwebui/
/home/szlc/miniconda3/envs/py310/ → /home/xsl/miniconda3/envs/py310/
# 2. SD WebUI 代码（排除 models/）
/home/szlc/project/onediff/stable-diffusion-webui/（排除 models/）
→ /home/xsl/change_hair/project/onediff/stable-diffusion-webui/
```
---
## 四、可选上传（非换发型推理必需）🟡
以下功能暂不需要时可跳过，后续按需补传。
### 4.1 训练新发型
| 原服务器路径 | 本机目标路径 | 预估大小 |
|-------------|-------------|---------|
| `/home/szlc/project/kohya_ss_home/` | `/home/xsl/change_hair/project/kohya_ss_home/` | ~13G |
| `/home/szlc/miniconda3/envs/kohya/` | `/home/xsl/miniconda3/envs/kohya/` | ~5G |
### 4.2 换发色功能
| 原服务器路径 | 本机目标路径 |
|-------------|-------------|
| `/home/szlc/project/data/ref_haircolor/` | `/home/xsl/change_hair/project/data/ref_haircolor/` |
| `/home/szlc/project/data/ref_color/` | `/home/xsl/change_hair/project/data/ref_color/` |
| `/home/szlc/project/data/ref_color_imgs/` | `/home/xsl/change_hair/project/data/ref_color_imgs/` |
| `/home/szlc/.../weights/classify_3_2499.pth` | `project/hair_service_sd/weights/classify_3_2499.pth` |
### 4.3 更多发型
当前只有 5 个发型。如需更多，从原服务器追加：
```
/data/train_material/<hair_id>/ → data/train_material/<hair_id>/
/home/szlc/project/data/ref_hairstyle/<hair_id>/ → project/data/ref_hairstyle/<hair_id>/
/home/szlc/project/data/hair_template_material/<hair_id>/ → project/data/hair_template_material/<hair_id>/
/home/szlc/project/data/upload_train_imgs/<hair_id>/ → project/data/upload_train_imgs/<hair_id>/
```
原服务器约有 338 个发型，全量 `train_material` 约 121G。
### 4.4 小文件（可选）
| 文件 | 用途 | 大小 |
|------|------|------|
| `weights/mediapipe/face_landmarker.task` | 发际线检测 | ~3.6M |
---
## 五、上传方式建议
### 方式 1：rsync 直传（推荐，支持断点续传）
```bash
# 从原服务器直传到本机（在本机执行）
rsync -avP --partial \
szlc@192.168.101.63:/home/szlc/miniconda3/envs/my_hair/ \
/home/xsl/miniconda3/envs/my_hair/
rsync -avP --partial \
szlc@192.168.101.63:/home/szlc/miniconda3/envs/sdwebui/ \
/home/xsl/miniconda3/envs/sdwebui/
rsync -avP --partial \
szlc@192.168.101.63:/home/szlc/miniconda3/envs/py310/ \
/home/xsl/miniconda3/envs/py310/
rsync -avP --partial \
--exclude='models/' \
szlc@192.168.101.63:/home/szlc/project/onediff/stable-diffusion-webui/ \
/home/xsl/change_hair/project/onediff/stable-diffusion-webui/
```
### 方式 2：先 tar 再上传到 `/home/xsl/data`
若通过网盘中转，先 tar 再上传到 `/home/xsl/data/`，上传完成后告知我解压。
---
## 六、上传完成后的验证
上传完成后，在本机执行以下命令确认文件到位：
```bash
# 1. Conda 环境
ls /home/xsl/miniconda3/envs/my_hair/bin/python
ls /home/xsl/miniconda3/envs/sdwebui/bin/python
ls /home/xsl/miniconda3/envs/py310/bin/python
# 2. WebUI 代码
ls /home/xsl/change_hair/project/onediff/stable-diffusion-webui/webui.py
# 3. 模型未被覆盖
ls -lh /home/xsl/change_hair/project/onediff/stable-diffusion-webui/models/Stable-diffusion/v1-5-pruned-emaonly.safetensors
```
期望结果：4 项均存在，SD 底模约 4.0G。
---
## 七、上传完成后我会做的事
你上传完毕并告知后，我会依次执行：
1. **路径适配**：运行 `adapt_paths.sh`，将所有硬编码路径从 `/home/xsl/change_hair` 改为 `/home/xsl/change_hair`，`/home/xsl/miniconda3` 改为 `/home/xsl/miniconda3`
2. **Conda 路径修复**：运行 `fix_conda_paths.py`，修复 3 个环境中二进制文件的 shebang 和内部路径前缀
3. **安装 Miniconda 基础环境**（如尚未安装）：从 `repo.anaconda.com` 下载安装
4. **启动服务**：`bash start_all.sh start`（依赖顺序：webui → photo_service → hair_service_sd）
5. **接口测试**：用 `test_swaphair.py` 验证换发型功能
---
## 八、部署后服务端口
| 服务 | 端口 | 说明 |
|------|------|------|
| webui | 57860 | SD img2img 推理 |
| photo_service | 32678 | LoRA 加载与调度 |
| hair_service_sd | 8801 | 换发型主接口 |
| hair_grow_service | 8899 | 生发前端（独立 Flask） |
换发型接口：
```
POST http://127.0.0.1:8801/api/swapHair/v1
```
---
## 九、常见问题
**Q：能否只上传 miniconda3 整包，不传单个 env？**
可以。上传 `/home/szlc/miniconda3/` 整包到 `/home/xsl/miniconda3/` 也行，体积更大（约 20–30G）但更省事。
**Q：webui 上传时不小心覆盖了 models 怎么办？**
重新从 `/home/xsl/data/sdwebui/models/` 或原服务器补传 models 目录即可。
**Q：本机没有 miniconda 基础安装，只传 envs 够吗？**
不够。3 个 env 内的 python 依赖 miniconda 基础库（如 `libstdc++`、`libgcc` 等）。我会在上传完成后先安装 miniconda，再修复 env 内部路径。
**Q：上传后启动报 xformers 错误？**
webui 启动参数可能需要去掉 `--xformers`，改用 `--opt-sdp-no-mem-attention`（RTX 5090 支持，我会在路径适配时处理）。
**Q：本机磁盘空间是否足够？**
当前 `/home` 剩余约 30G。3 个 conda 环境约 20–25G，解压后可能吃紧。建议解压一个删一个 tar 包，或考虑清理 `/home/xsl/data/`（已有 16G 冗余数据，解压完成后可删除）。
gen_hairline_mask.py
+3 -3
View File
@@ -8,7 +8,7 @@ import sys
import cv2 import cv2
import numpy as np import numpy as np
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd" HAIR_SERVICE_DIR = "/home/ubuntu/change_hair/project/hair_service_sd"
sys.path.insert(0, HAIR_SERVICE_DIR) sys.path.insert(0, HAIR_SERVICE_DIR)
os.environ.setdefault("HF_HUB_OFFLINE", "1") os.environ.setdefault("HF_HUB_OFFLINE", "1")
@@ -87,7 +87,7 @@ def gen_hairline_mask(img_path, mask_path):
if __name__ == "__main__": if __name__ == "__main__":
img_path = sys.argv[1] if len(sys.argv) > 1 else \ img_path = sys.argv[1] if len(sys.argv) > 1 else \
"/home/xsl/change_hair/project/data/userImage/8488902485_20250630055547.jpg" "/home/ubuntu/change_hair/project/data/userImage/8488902485_20250630055547.jpg"
mask_path = sys.argv[2] if len(sys.argv) > 2 else \ mask_path = sys.argv[2] if len(sys.argv) > 2 else \
"/home/xsl/change_hair/project/logs/hairgrow_test_mask.png" "/home/ubuntu/change_hair/project/logs/hairgrow_test_mask.png"
gen_hairline_mask(img_path, mask_path) gen_hairline_mask(img_path, mask_path)
gen_hairstyle_previews.py
+6 -6
View File
@@ -14,15 +14,15 @@ import time
import base64 import base64
import requests import requests
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd" HAIR_SERVICE_DIR = "/home/ubuntu/change_hair/project/hair_service_sd"
sys.path.insert(0, HAIR_SERVICE_DIR) sys.path.insert(0, HAIR_SERVICE_DIR)
os.chdir(HAIR_SERVICE_DIR) os.chdir(HAIR_SERVICE_DIR)
BOY_IMG = "/home/xsl/change_hair/images/boy.png" BOY_IMG = "/home/ubuntu/change_hair/images/boy.png"
GIRL_IMG = "/home/xsl/change_hair/images/girl.png" GIRL_IMG = "/home/ubuntu/change_hair/images/girl.png"
HAIRSTYLE_DIR = "/home/xsl/change_hair/project/data/ref_hairstyle" HAIRSTYLE_DIR = "/home/ubuntu/change_hair/project/data/ref_hairstyle"
TRAIN_DIR = "/home/xsl/change_hair/data/train_material" TRAIN_DIR = "/home/ubuntu/change_hair/data/train_material"
PREVIEW_DIR = "/home/xsl/change_hair/hair_grow_service/static/previews" # 预览图存储目录 PREVIEW_DIR = "/home/ubuntu/change_hair/hair_grow_service/static/previews" # 预览图存储目录
SWAP_API = "http://127.0.0.1:8801/api/swapHair/v1" SWAP_API = "http://127.0.0.1:8801/api/swapHair/v1"
gen_template_material.py
+2 -2
View File
@@ -8,7 +8,7 @@
hair_template_material/<hair_id>/first##<name>.pkl 1k关键点 hair_template_material/<hair_id>/first##<name>.pkl 1k关键点
用法: 用法:
cd /home/xsl/change_hair/project/hair_service_sd cd /home/ubuntu/change_hair/project/hair_service_sd
python gen_template_material.py --hair-id new_test_001 --img /path/to/template.jpg python gen_template_material.py --hair-id new_test_001 --img /path/to/template.jpg
""" """
import os import os
@@ -19,7 +19,7 @@ import pickle
import argparse import argparse
ssl._create_default_https_context = ssl._create_unverified_context ssl._create_default_https_context = ssl._create_unverified_context
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd" HAIR_SERVICE_DIR = "/home/ubuntu/change_hair/project/hair_service_sd"
os.chdir(HAIR_SERVICE_DIR) os.chdir(HAIR_SERVICE_DIR)
sys.path.insert(0, HAIR_SERVICE_DIR) sys.path.insert(0, HAIR_SERVICE_DIR)
os.environ.setdefault("HF_HUB_OFFLINE", "1") os.environ.setdefault("HF_HUB_OFFLINE", "1")
hair_grow_cli.py
+4 -4
View File
@@ -6,12 +6,12 @@
需要在 hair_service_sd 目录下运行（依赖其模块导入）。 需要在 hair_service_sd 目录下运行（依赖其模块导入）。
用法: 用法:
cd /home/xsl/change_hair/project/hair_service_sd cd /home/ubuntu/change_hair/project/hair_service_sd
/home/xsl/miniconda3/envs/my_hair/bin/python /home/xsl/change_hair/hair_grow_cli.py \ /home/ubuntu/miniconda3/envs/my_hair/bin/python /home/ubuntu/change_hair/hair_grow_cli.py \
--img test.jpg --mask mask.png --strength 0.5 -o result.jpg --img test.jpg --mask mask.png --strength 0.5 -o result.jpg
# 批量跑三档强度对比 # 批量跑三档强度对比
/home/xsl/miniconda3/envs/my_hair/bin/python /home/xsl/change_hair/hair_grow_cli.py \ /home/ubuntu/miniconda3/envs/my_hair/bin/python /home/ubuntu/change_hair/hair_grow_cli.py \
--img test.jpg --mask mask.png --compare --img test.jpg --mask mask.png --compare
""" """
import os import os
@@ -20,7 +20,7 @@ import cv2
import argparse import argparse
# 把 hair_service_sd 加入路径，使其模块可被导入 # 把 hair_service_sd 加入路径，使其模块可被导入
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd" HAIR_SERVICE_DIR = "/home/ubuntu/change_hair/project/hair_service_sd"
sys.path.insert(0, HAIR_SERVICE_DIR) sys.path.insert(0, HAIR_SERVICE_DIR)
# 设置离线模式（本机无法访问 huggingface.co） # 设置离线模式（本机无法访问 huggingface.co）
hair_grow_service/app.py
+19 -81
View File
@@ -24,10 +24,19 @@ os.environ.setdefault("CRYPTOGRAPHY_OPENSSL_NO_LEGACY", "1")
import cv2 import cv2
import numpy as np import numpy as np
from flask import Flask, request, jsonify, send_from_directory from flask import Flask, request, jsonify, send_from_directory, send_file
from gevent import pywsgi from gevent import pywsgi
PORT = 8888 # PyTorch 2.6+ 默认 weights_only=True，无法加载含 numpy 对象的旧权重（yolov5l.pt 等）。
# 统一改回 False（权重均为本机自有可信文件）。
import torch
_orig_torch_load = torch.load
def _torch_load(*args, **kwargs):
kwargs.setdefault("weights_only", False)
return _orig_torch_load(*args, **kwargs)
torch.load = _torch_load
PORT = 8899
BASE_DIR = os.path.dirname(os.path.abspath(__file__)) BASE_DIR = os.path.dirname(os.path.abspath(__file__))
app = Flask(__name__, static_folder="static", static_url_path="/static") app = Flask(__name__, static_folder="static", static_url_path="/static")
@@ -87,12 +96,6 @@ def manual_page():
return send_from_directory(os.path.join(BASE_DIR, "static"), "manual.html") return send_from_directory(os.path.join(BASE_DIR, "static"), "manual.html")
@app.route("/inpaint")
def inpaint_page():
"""mask区域重绘测试页（纯发丝增强/自然重绘）"""
return send_from_directory(os.path.join(BASE_DIR, "static"), "inpaint.html")
@app.route("/api/swap", methods=["POST"]) @app.route("/api/swap", methods=["POST"])
def api_swap_proxy(): def api_swap_proxy():
"""代理换发型请求到 8801（避免前端跨域问题）""" """代理换发型请求到 8801（避免前端跨域问题）"""
@@ -346,7 +349,6 @@ def api_swap_manual():
blend_dilate=(int(bd[0]), int(bd[1])), blend_dilate=(int(bd[0]), int(bd[1])),
seamless_dilate=(int(sd[0]), int(sd[1])), seamless_dilate=(int(sd[0]), int(sd[1])),
feather_px=int(g("feather_px", 0)), feather_px=int(g("feather_px", 0)),
feather_alpha=float(g("feather_alpha", 1.0)),
enhance=bool(g("enhance", False)), enhance=bool(g("enhance", False)),
enhance_denoising=float(g("enhance_denoising", 0.35)), enhance_denoising=float(g("enhance_denoising", 0.35)),
) )
@@ -401,7 +403,7 @@ def preview_img(hair_id):
hairstyle_dir = config.get('default', 'hairstyleDir') hairstyle_dir = config.get('default', 'hairstyleDir')
fallback = os.path.join(hairstyle_dir, hair_id, "ref_rgb_8uc3_768.png") fallback = os.path.join(hairstyle_dir, hair_id, "ref_rgb_8uc3_768.png")
if os.path.exists(fallback): if os.path.exists(fallback):
return send_file_or_404(fallback) return send_file(fallback)
except Exception: except Exception:
pass pass
return ("", 404) return ("", 404)
@@ -409,25 +411,14 @@ def preview_img(hair_id):
@app.route("/train_src/<hair_id>") @app.route("/train_src/<hair_id>")
def train_src_img(hair_id): def train_src_img(hair_id):
"""返回发型的训练原图，用于测试页展示发型真实样子。 """返回发型的训练原图（hair_type_images/<hair_id>.jpg/.png）。
依次查找：hair_type_images/ → images/ → train_material/<id>/images 输入 → 预览图（兜底）。 用于测试页展示发型真实样子，而非套在标准脸上的效果图。
""" """
# 候选源目录（按优先级） src_dir = "/home/xsl/change_hair/hair_type_images"
candidate_dirs = [ for ext in (".jpg", ".jpeg", ".png"):
"/home/xsl/change_hair/hair_type_images", path = os.path.join(src_dir, hair_id + ext)
"/home/xsl/change_hair/images", if os.path.exists(path):
"/home/xsl/change_hair/data/batch_train_inputs/" + hair_id, return send_from_directory(src_dir, hair_id + ext)
]
for src_dir in candidate_dirs:
for ext in (".jpg", ".jpeg", ".png"):
path = os.path.join(src_dir, hair_id + ext)
if os.path.exists(path):
return send_from_directory(src_dir, hair_id + ext)
# 兜底：返回套脸预览图（previews/<hair_id>.jpg）
preview_dir = os.path.join(BASE_DIR, "static", "previews")
preview_path = os.path.join(preview_dir, hair_id + ".jpg")
if os.path.exists(preview_path):
return send_from_directory(preview_dir, hair_id + ".jpg")
return ("", 404) return ("", 404)
@@ -525,59 +516,6 @@ def api_grow():
return jsonify({"state": -1, "msg": f"生发失败: {e}"}), 500 return jsonify({"state": -1, "msg": f"生发失败: {e}"}), 500
@app.route("/api/inpaint", methods=["POST"])
def api_inpaint():
"""mask区域重绘接口（严格只在mask区做SD inpainting）。
入参 JSON:
img: 原图 base64
mask: 手绘mask base64（白色=重绘区）
prompt: 提示词（空则用模式默认值）
denoising_strength: 重绘强度 0~1（默认0.5）
mode: "enhance" 或 "pure_inpaint"（默认enhance）
返回: {state, result, mask_used, info}
"""
try:
d = request.json
img_b64 = d.get("img", "")
mask_b64 = d.get("mask", "")
if not img_b64 or not mask_b64:
return jsonify({"state": -1, "msg": "img 和 mask 不能为空（请先手绘重绘区域）"}), 400
img = _b64_to_ndarray(img_b64, color=True)
mask = _b64_to_ndarray(mask_b64, color=False)
if img is None or mask is None:
return jsonify({"state": -1, "msg": "img 或 mask 解析失败"}), 400
prompt = d.get("prompt", "")
denoising = float(d.get("denoising_strength", 0.5))
mode = d.get("mode", "enhance")
if mode not in ("enhance", "pure_inpaint"):
mode = "enhance"
print(f"[inpaint] mode={mode}, denoising={denoising}, prompt={prompt[:40]}, img={img.shape}")
import sys
sys.path.insert(0, "/home/xsl/change_hair/project/hair_service_sd")
from inpaint_mask import inpaint_mask as do_inpaint
result, mask_used, info = do_inpaint(
img, mask, prompt=prompt,
denoising_strength=denoising, mode=mode
)
_, buf = cv2.imencode(".jpg", result, [cv2.IMWRITE_JPEG_QUALITY, 95])
result_b64 = base64.b64encode(buf).decode("utf-8")
_, mbuf = cv2.imencode(".png", mask_used)
mask_used_b64 = "data:image/png;base64," + base64.b64encode(mbuf).decode("utf-8")
return jsonify({"state": 0, "result": result_b64,
"mask_used": mask_used_b64, "info": info})
except Exception as e:
print(f"[inpaint] 失败: {e}")
traceback.print_exc()
return jsonify({"state": -1, "msg": f"mask重绘失败: {e}"}), 500
if __name__ == "__main__": if __name__ == "__main__":
print(f"生发服务启动，端口 {PORT}") print(f"生发服务启动，端口 {PORT}")
print(f"测试页面: http://0.0.0.0:{PORT}") print(f"测试页面: http://0.0.0.0:{PORT}")
hair_grow_service/static/inpaint.html
-323
View File
@@ -1,323 +0,0 @@
<!DOCTYPE html>
<html lang="zh-CN">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>mask区域重绘</title>
<style>
* { margin: 0; padding: 0; box-sizing: border-box; }
body { font-family: -apple-system, "Segoe UI", "Microsoft YaHei", sans-serif; background: #1a1a2e; color: #eee; min-height: 100vh; font-size: 14px; }
.header { background: #16213e; padding: 16px 28px; border-bottom: 1px solid #0f3460; display: flex; justify-content: space-between; align-items: center; }
.header h1 { font-size: 19px; font-weight: 600; }
.header p { font-size: 12px; color: #888; margin-top: 3px; }
.header .links a { color: #4ecca3; font-size: 12px; text-decoration: none; margin-left: 12px; }
.container { display: flex; gap: 18px; padding: 18px; max-width: 1800px; margin: 0 auto; }
/* 左侧工具+参数 */
.panel { width: 300px; flex-shrink: 0; background: #16213e; border-radius: 10px; padding: 16px; height: fit-content; }
.panel h3 { font-size: 13px; color: #4ecca3; margin: 14px 0 8px; text-transform: uppercase; letter-spacing: 1px; padding-bottom: 5px; border-bottom: 1px solid #0f3460; }
.panel h3:first-child { margin-top: 0; }
.hint { font-size: 11px; color: #888; line-height: 1.6; margin: 6px 0 10px; }
.tool-row { display: flex; gap: 6px; margin-bottom: 10px; }
.btn-tool { flex: 1; padding: 8px; background: #0f3460; color: #ccc; border: none; border-radius: 5px; cursor: pointer; font-size: 12px; }
.btn-tool.active { background: #4ecca3; color: #16213e; font-weight: 600; }
.btn-tool:hover { background: #1a4a7a; }
.btn-tool.active:hover { background: #6ee0bd; }
.btn-danger { width: 100%; padding: 7px; background: #c0392b; color: #fff; border: none; border-radius: 5px; cursor: pointer; font-size: 12px; margin-bottom: 10px; }
.btn-danger:hover { background: #e74c3c; }
.field { margin-bottom: 12px; }
.field-label { display: flex; justify-content: space-between; align-items: center; font-size: 12px; color: #bbb; margin-bottom: 5px; }
.field-label .val { color: #4ecca3; font-family: monospace; font-size: 11px; }
.field input[type="range"] { width: 100%; height: 4px; -webkit-appearance: none; background: #0f3460; border-radius: 2px; outline: none; }
.field input[type="range"]::-webkit-slider-thumb { -webkit-appearance: none; width: 14px; height: 14px; border-radius: 50%; background: #4ecca3; cursor: pointer; }
.field input[type="text"], .field textarea { width: 100%; padding: 6px 8px; background: #0d1b3e; border: 1px solid #0f3460; border-radius: 4px; color: #eee; font-size: 12px; font-family: monospace; }
.field textarea { min-height: 60px; resize: vertical; }
.mode-tabs { display: flex; gap: 6px; margin-bottom: 10px; }
.mode-tab { flex: 1; padding: 8px; background: #0f3460; border: none; border-radius: 5px; color: #aaa; cursor: pointer; font-size: 12px; text-align: center; }
.mode-tab.active { background: #4ecca3; color: #16213e; font-weight: 600; }
.btn { display: block; width: 100%; padding: 12px; border: none; border-radius: 6px; font-size: 15px; cursor: pointer; transition: .15s; }
.btn-primary { background: #4ecca3; color: #16213e; font-weight: 600; margin-top: 8px; }
.btn-primary:hover { background: #6ee0bd; }
.btn-primary:disabled { background: #555; color: #999; cursor: not-allowed; }
/* 右侧画板+结果 */
.workspace { flex: 1; min-width: 0; }
.canvas-wrap { background: #0d0d1a; border-radius: 10px; padding: 18px; text-align: center; min-height: 280px; display: flex; align-items: center; justify-content: center; }
.upload-zone { width: 100%; max-width: 600px; border: 2px dashed #0f3460; border-radius: 10px; padding: 40px 20px; text-align: center; cursor: pointer; transition: .2s; }
.upload-zone:hover { border-color: #4ecca3; background: rgba(78,204,163,.05); }
.upload-zone svg { width: 42px; height: 42px; fill: #4ecca3; margin-bottom: 10px; }
.upload-zone p { color: #888; font-size: 13px; }
.upload-zone p.highlight { color: #4ecca3; }
#canvas-container { position: relative; display: inline-block; max-width: 100%; }
#img-canvas, #mask-canvas { display: block; max-width: 100%; border-radius: 6px; }
#mask-canvas { position: absolute; top: 0; left: 0; cursor: crosshair; touch-action: none; }
#mask-canvas.eraser { cursor: cell; }
.canvas-tip { font-size: 11px; color: #666; margin-top: 8px; }
#status { text-align: center; padding: 24px; color: #4ecca3; display: none; }
.spinner { display: inline-block; width: 18px; height: 18px; border: 3px solid #0f3460; border-top-color: #4ecca3; border-radius: 50%; animation: spin .8s linear infinite; margin-right: 8px; vertical-align: middle; }
@keyframes spin { to { transform: rotate(360deg); } }
#error-msg { color: #e74c3c; padding: 16px; display: none; font-size: 13px; background: rgba(231,76,60,.1); border-radius: 6px; margin-top: 12px; }
.result-section { margin-top: 16px; background: #16213e; border-radius: 10px; padding: 16px; display: none; }
.result-section.show { display: block; }
.result-section h3 { font-size: 14px; color: #4ecca3; margin-bottom: 12px; }
.result-grid { display: flex; gap: 16px; flex-wrap: wrap; }
.result-card { flex: 1; min-width: 240px; }
.result-card h4 { font-size: 12px; color: #aaa; margin-bottom: 6px; text-align: center; }
.result-card img { width: 100%; border-radius: 6px; display: block; cursor: zoom-in; }
.info-box { margin-top: 12px; padding: 10px 12px; background: #0d1b3e; border-radius: 6px; font-size: 11px; color: #888; line-height: 1.7; }
#lightbox { display: none; position: fixed; inset: 0; background: rgba(0,0,0,.92); z-index: 999; justify-content: center; align-items: center; cursor: zoom-out; padding: 20px; }
#lightbox img { max-width: 95%; max-height: 95%; border-radius: 8px; }
#lightbox.show { display: flex; }
</style>
</head>
<body>
<div class="header">
<div>
<h1>🖌️ mask区域重绘（发丝自然化）</h1>
<p>上传图片 → 涂抹要重绘的区域 → 调提示词和强度 → 生成。严格只在涂抹区重绘，其余保留原图</p>
</div>
<div class="links">
<a href="/manual">→ 手绘mask换发型</a>
<a href="/debug">→ 换发型调试</a>
</div>
</div>
<div class="container">
<!-- 左侧 -->
<div class="panel">
<h3>① 画笔工具</h3>
<div class="hint">在右侧图片上涂抹需要重绘的区域（白色=重绘，黑色=保留原图）</div>
<div class="tool-row">
<button class="btn-tool active" id="tool-brush">🖌️ 画笔</button>
<button class="btn-tool" id="tool-eraser">🧹 橡皮</button>
</div>
<div class="field">
<div class="field-label">画笔大小 <span class="val" id="size-val">25px</span></div>
<input type="range" id="brush-size" min="5" max="80" value="25">
</div>
<button class="btn-danger" id="btn-clear">✕ 清除全部涂抹</button>
<h3>② 重绘模式</h3>
<div class="mode-tabs">
<button class="mode-tab active" id="mode-enhance" onclick="setMode('enhance')">发丝增强<br><span style="font-size:10px;opacity:.7">enhance</span></button>
<button class="mode-tab" id="mode-pure" onclick="setMode('pure_inpaint')">自然重绘<br><span style="font-size:10px;opacity:.7">pure_inpaint</span></button>
</div>
<h3>③ 提示词</h3>
<div class="field">
<textarea id="prompt" placeholder="留空用模式默认提示词">high quality, detailed natural hair strands, realistic hair texture, sharp focus, photorealistic</textarea>
</div>
<h3>④ 重绘强度</h3>
<div class="field">
<div class="field-label">denoising_strength <span class="val" id="v-den">0.50</span><span style="color:#666;font-size:10px;margin-left:6px">越低越保留原图</span></div>
<input type="range" id="den" min="0.1" max="1.0" step="0.05" value="0.5" oninput="$('v-den').textContent=parseFloat(this.value).toFixed(2)">
</div>
<button class="btn btn-primary" id="btn-run" disabled>▶ 生成（mask区域重绘）</button>
<div class="hint" style="margin-top:10px">
<b>说明：</b><br>
• 严格只在涂抹区域重绘，mask外完全保留原图<br>
• enhance模式：发丝细节增强（默认0.35强度）<br>
• pure_inpaint模式：自然重绘（默认0.5强度）<br>
• 改提示词可控制重绘方向（如"detailed curly hair"）
</div>
</div>
<!-- 右侧画板+结果 -->
<div class="workspace">
<div class="canvas-wrap">
<div class="upload-zone" id="upload-zone">
<svg viewBox="0 0 24 24"><path d="M19 13v6H5v-6H3v8h18v-8zM6 9l1.41 1.41L11 6.83V18h2V6.83l3.59 3.58L18 9l-6-6z"/></svg>
<p class="highlight">点击上传图片</p>
<p>支持 jpg/png，建议人脸清晰照</p>
</div>
<div id="canvas-container" style="display:none">
<canvas id="img-canvas"></canvas>
<canvas id="mask-canvas"></canvas>
</div>
<input type="file" id="file-input" accept="image/*" style="display:none">
</div>
<div class="canvas-tip" id="canvas-tip" style="display:none;text-align:center">用左侧画笔在图片上涂抹要重绘的区域</div>
<div id="status"><span class="spinner"></span><span id="status-text">重绘中，约15-30秒...</span></div>
<div id="error-msg"></div>
<div class="result-section" id="result-section">
<h3>对比结果</h3>
<div class="result-grid">
<div class="result-card"><h4>原图</h4><img id="r-orig"></div>
<div class="result-card"><h4>使用的mask</h4><img id="r-mask"></div>
<div class="result-card"><h4>重绘结果</h4><img id="r-new"></div>
</div>
<div class="info-box" id="info-box"></div>
</div>
</div>
</div>
<div id="lightbox" onclick="this.classList.remove('show')"><img id="lightbox-img"></div>
<script>
(function(){
const $=id=>document.getElementById(id);
let mode="enhance", imgB64Orig=null;
const imgCanvas=$("img-canvas"),maskCanvas=$("mask-canvas");
const imgCtx=imgCanvas.getContext("2d"),maskCtx=maskCanvas.getContext("2d",{willReadFrequently:true});
let brushSize=25,drawing=false,lastX=0,lastY=0;
// 模式切换
window.setMode=(m)=>{
mode=m;
$("mode-enhance").classList.toggle("active",m==="enhance");
$("mode-pure").classList.toggle("active",m==="pure_inpaint");
// 切换模式时更新默认提示词（如果当前是空的或等于另一模式的默认值）
const defaults={
enhance:"high quality, detailed natural hair strands, realistic hair texture, sharp focus, photorealistic",
pure_inpaint:"natural realistic hair, soft hair strands, photorealistic, high quality, detailed"
};
const cur=$("prompt").value.trim();
if(!cur || Object.values(defaults).includes(cur)){
$("prompt").value=defaults[m];
}
// 切换模式时调整默认denoising
if(m==="enhance" && parseFloat($("den").value)===0.5){ $("den").value=0.35; $("v-den").textContent="0.35"; }
if(m==="pure_inpaint" && parseFloat($("den").value)===0.35){ $("den").value=0.5; $("v-den").textContent="0.50"; }
};
// 上传
$("upload-zone").onclick=()=>$("file-input").click();
$("file-input").onchange=e=>{ if(e.target.files[0]) loadImage(e.target.files[0]); };
$("upload-zone").ondragover=e=>{e.preventDefault();$("upload-zone").style.borderColor="#4ecca3";};
$("upload-zone").ondragleave=()=>$("upload-zone").style.borderColor="#0f3460";
$("upload-zone").ondrop=e=>{e.preventDefault();if(e.dataTransfer.files[0])loadImage(e.dataTransfer.files[0]);};
function loadImage(file){
if(!file.type.startsWith("image/")){alert("请上传图片");return;}
const rd=new FileReader();
rd.onload=e=>{
const im=new Image();
im.onload=()=>{
let w=im.width,h=im.height;
const MAX=1024;
if(Math.max(w,h)>MAX){const sc=MAX/Math.max(w,h);w=Math.round(w*sc);h=Math.round(h*sc);}
imgCanvas.width=w;imgCanvas.height=h;
maskCanvas.width=w;maskCanvas.height=h;
imgCtx.drawImage(im,0,0,w,h);
maskCtx.clearRect(0,0,w,h);
imgB64Orig=imgCanvas.toDataURL("image/jpeg",0.95);
$("canvas-container").style.display="inline-block";
$("upload-zone").style.display="none";
$("canvas-tip").style.display="block";
$("btn-run").disabled=false;
};
im.src=e.target.result;
};
rd.readAsDataURL(file);
}
// 画笔
let drawMode="brush";
$("tool-brush").onclick=()=>{drawMode="brush";$("tool-brush").classList.add("active");$("tool-eraser").classList.remove("active");maskCanvas.classList.remove("eraser");};
$("tool-eraser").onclick=()=>{drawMode="eraser";$("tool-eraser").classList.add("active");$("tool-brush").classList.remove("active");maskCanvas.classList.add("eraser");};
$("brush-size").oninput=e=>{brushSize=+e.target.value;$("size-val").textContent=brushSize+"px";};
$("btn-clear").onclick=()=>{maskCtx.globalCompositeOperation="source-over";maskCtx.clearRect(0,0,maskCanvas.width,maskCanvas.height);};
function getPos(e){
const r=maskCanvas.getBoundingClientRect();
const sc=maskCanvas.width/r.width;
const cx=(e.touches?e.touches[0].clientX:e.clientX)-r.left;
const cy=(e.touches?e.touches[0].clientY:e.clientY)-r.top;
return [cx*sc,cy*sc];
}
function startDraw(e){e.preventDefault();drawing=true;const[x,y]=getPos(e);lastX=x;lastY=y;drawDot(x,y);}
function moveDraw(e){if(!drawing)return;e.preventDefault();const[x,y]=getPos(e);drawLine(lastX,lastY,x,y);lastX=x;lastY=y;}
function endDraw(){drawing=false;}
function drawDot(x,y){
if(drawMode==="brush"){maskCtx.globalCompositeOperation="source-over";maskCtx.fillStyle="rgba(255,60,80,0.5)";}
else{maskCtx.globalCompositeOperation="destination-out";maskCtx.fillStyle="rgba(0,0,0,1)";}
maskCtx.beginPath();maskCtx.arc(x,y,brushSize/2,0,Math.PI*2);maskCtx.fill();
}
function drawLine(x1,y1,x2,y2){
if(drawMode==="brush"){maskCtx.globalCompositeOperation="source-over";maskCtx.strokeStyle="rgba(255,60,80,0.5)";}
else{maskCtx.globalCompositeOperation="destination-out";maskCtx.strokeStyle="rgba(0,0,0,1)";}
maskCtx.beginPath();maskCtx.moveTo(x1,y1);maskCtx.lineTo(x2,y2);
maskCtx.lineWidth=brushSize;maskCtx.lineCap="round";maskCtx.lineJoin="round";maskCtx.stroke();
}
maskCanvas.addEventListener("mousedown",startDraw);
maskCanvas.addEventListener("mousemove",moveDraw);
window.addEventListener("mouseup",endDraw);
maskCanvas.addEventListener("touchstart",startDraw,{passive:false});
maskCanvas.addEventListener("touchmove",moveDraw,{passive:false});
maskCanvas.addEventListener("touchend",endDraw);
// 提取mask为黑白二值PNG
function extractMaskB64(){
const tmp=document.createElement("canvas");
tmp.width=maskCanvas.width;tmp.height=maskCanvas.height;
const tc=tmp.getContext("2d");
const md=maskCtx.getImageData(0,0,maskCanvas.width,maskCanvas.height);
const out=tc.createImageData(tmp.width,tmp.height);
for(let i=0;i<md.data.length;i+=4){
const v=md.data[i+3]>30?255:0;
out.data[i]=v;out.data[i+1]=v;out.data[i+2]=v;out.data[i+3]=255;
}
tc.putImageData(out,0,0);
return tmp.toDataURL("image/png");
}
// 执行
$("btn-run").onclick=async()=>{
if(!imgB64Orig) return;
// 检查mask
const md=maskCtx.getImageData(0,0,maskCanvas.width,maskCanvas.height);
let hasMask=false;
for(let i=3;i<md.data.length;i+=4){if(md.data[i]>30){hasMask=true;break;}}
if(!hasMask){alert("请先用画笔涂抹需要重绘的区域");return;}
const maskB64=extractMaskB64();
$("btn-run").disabled=true;
$("status").style.display="block";
$("status-text").textContent="mask区域重绘中，约15-30秒...";
$("error-msg").style.display="none";
$("result-section").classList.remove("show");
try{
const resp=await fetch("/api/inpaint",{
method:"POST",headers:{"Content-Type":"application/json"},
body:JSON.stringify({
img:imgB64Orig, mask:maskB64,
prompt:$("prompt").value.trim(),
denoising_strength:parseFloat($("den").value),
mode:mode
})
});
const data=await resp.json();
$("status").style.display="none";
if(data.state===0){
$("r-orig").src=imgB64Orig;
$("r-mask").src=data.mask_used;
$("r-new").src="data:image/jpeg;base64,"+data.result;
$("info-box").textContent=data.info;
// 绑定点击放大
["r-orig","r-mask","r-new"].forEach(id=>{
$(id).onclick=(e)=>{ $("lightbox-img").src=e.target.src; $("lightbox").classList.add("show"); };
});
$("result-section").classList.add("show");
$("result-section").scrollIntoView({behavior:"smooth"});
}else{
$("error-msg").textContent="❌ "+(data.msg||"失败");
$("error-msg").style.display="block";
}
}catch(err){
$("status").style.display="none";
$("error-msg").textContent="❌ 请求失败: "+err.message;
$("error-msg").style.display="block";
}
$("btn-run").disabled=false;
};
})();
</script>
</body>
</html>
hair_grow_service/static/manual.html
+11 -13
View File
@@ -159,13 +159,9 @@ body { font-family: -apple-system, "Segoe UI", "Microsoft YaHei", sans-serif; ba
</div> </div>
</div> </div>
<div class="field"> <div class="field">
<div class="field-label">feather_px <span class="val" id="v-fp">0</span><span class="tip" title="羽化范围(像素)，边缘带宽度。0=无羽化，越大羽化范围越宽">ⓘ</span></div> <div class="field-label">feather_px <span class="val" id="v-fp">0</span><span class="tip" title="严格mask贴回的边缘羽化像素。0=硬边缘，>0=高斯模糊边缘平滑过渡（消除贴回接缝）">ⓘ</span></div>
<input type="range" id="fp" min="0" max="60" step="1" value="0" oninput="syncInt('fp')"> <input type="range" id="fp" min="0" max="60" step="1" value="0" oninput="syncInt('fp')">
</div> </div>
<div class="field">
<div class="field-label">feather_alpha <span class="val" id="v-fa">1.00</span><span class="tip" title="羽化强度(0~1)。0=硬切不羽化，1=最大羽化。控制边缘过渡柔和程度">ⓘ</span></div>
<input type="range" id="fa" min="0" max="1.0" step="0.05" value="1.0" oninput="syncFloat('fa')">
</div>
<div class="readonly-info"> <div class="readonly-info">
<b style="color:#888">webui端固定：</b><br> <b style="color:#888">webui端固定：</b><br>
<code>cfg=7</code> <code>steps=20</code> <code>sampler=DPM++ 2M Karras</code> <code>cfg=7</code> <code>steps=20</code> <code>sampler=DPM++ 2M Karras</code>
@@ -186,7 +182,7 @@ body { font-family: -apple-system, "Segoe UI", "Microsoft YaHei", sans-serif; ba
<!-- 中间发型选择 --> <!-- 中间发型选择 -->
<div class="mid-panel"> <div class="mid-panel">
<h3>③ 选发型 <span id="count-info" style="font-size:10px;color:#888;font-weight:normal"></span></h3> <h3>③ 选发型</h3>
<div class="hairstyle-grid" id="hairstyle-grid"><div style="color:#666;font-size:11px">加载中...</div></div> <div class="hairstyle-grid" id="hairstyle-grid"><div style="color:#666;font-size:11px">加载中...</div></div>
<div style="font-size:11px;color:#4ecca3;margin-top:8px;text-align:center" id="sel-info">未选择</div> <div style="font-size:11px;color:#4ecca3;margin-top:8px;text-align:center" id="sel-info">未选择</div>
<button class="btn-tool" style="margin-top:8px;width:100%" onclick="loadStatus()">↻ 刷新</button> <button class="btn-tool" style="margin-top:8px;width:100%" onclick="loadStatus()">↻ 刷新</button>
@@ -229,6 +225,14 @@ body { font-family: -apple-system, "Segoe UI", "Microsoft YaHei", sans-serif; ba
<script> <script>
(function(){ (function(){
const $=id=>document.getElementById(id); const $=id=>document.getElementById(id);
const HAIRSTYLES = [
{face:"圆",name:"圆-心形"},{face:"圆",name:"圆-椭圆"},{face:"圆",name:"圆-波浪"},{face:"圆",name:"圆-直线"},{face:"圆",name:"圆-花瓣"},
{face:"心形",name:"心形-心形"},{face:"心形",name:"心形-椭圆"},{face:"心形",name:"心形-波浪"},{face:"心形",name:"心形-直线"},{face:"心形",name:"心形-花瓣"},
{face:"方脸",name:"方脸-心形"},{face:"方脸",name:"方脸-椭圆"},{face:"方脸",name:"方脸-波浪"},{face:"方脸",name:"方脸-直线"},{face:"方脸",name:"方脸-花瓣"},
{face:"椭圆",name:"椭圆-心形"},{face:"椭圆",name:"椭圆-椭圆"},{face:"椭圆",name:"椭圆-波浪"},{face:"椭圆",name:"椭圆-直线"},{face:"椭圆",name:"椭圆-花瓣"},
{face:"菱形",name:"菱形-心"},{face:"菱形",name:"菱形-椭圆"},{face:"菱形",name:"菱形-波浪"},{face:"菱形",name:"菱形-直线"},{face:"菱形",name:"菱形-花瓣"},
{face:"长",name:"长-心形"},{face:"长",name:"长 -椭圆"},{face:"长",name:"长 -波浪"},{face:"长",name:"长 -直线"},{face:"长",name:"长 -花瓣"}
];
let readySet=new Set(), selectedHair=null, imgB64Orig=null; let readySet=new Set(), selectedHair=null, imgB64Orig=null;
// 画板状态 // 画板状态
@@ -236,11 +240,6 @@ body { font-family: -apple-system, "Segoe UI", "Microsoft YaHei", sans-serif; ba
const imgCtx=imgCanvas.getContext("2d"),maskCtx=maskCanvas.getContext("2d",{willReadFrequently:true}); const imgCtx=imgCanvas.getContext("2d"),maskCtx=maskCanvas.getContext("2d",{willReadFrequently:true});
let mode="brush",brushSize=25,drawing=false,lastX=0,lastY=0; let mode="brush",brushSize=25,drawing=false,lastX=0,lastY=0;
// 本次训练的发型（昨天6个 + 今天hair_flow）
const HAIRSTYLES = [
{name:"hair_flow"},{name:"圆-心形"},{name:"圆-椭圆"},{name:"圆-波浪"},
{name:"圆-直线"},{name:"圆-花瓣"},{name:"心形-心形"},{name:"心形-椭圆"}
];
async function loadStatus(){ async function loadStatus(){
try{ try{
const r=await fetch("/api/hairstyles");const d=await r.json(); const r=await fetch("/api/hairstyles");const d=await r.json();
@@ -255,7 +254,7 @@ body { font-family: -apple-system, "Segoe UI", "Microsoft YaHei", sans-serif; ba
const ready=readySet.has(h.name); const ready=readySet.has(h.name);
const it=document.createElement("div"); const it=document.createElement("div");
it.className="hairstyle-item"+(ready?" ready":" pending")+(h.name===selectedHair?" selected":""); it.className="hairstyle-item"+(ready?" ready":" pending")+(h.name===selectedHair?" selected":"");
it.innerHTML=`<img loading="lazy" src="/train_src/${encodeURIComponent(h.name)}" onerror="this.style.objectFit='contain';this.src='data:image/svg+xml;utf8,<svg xmlns=%22http://www.w3.org/2000/svg%22 width=%22100%22 height=%22100%22><text x=%2250%22 y=%2255%22 text-anchor=%22middle%22 fill=%22%23666%22 font-size=%2210%22>无图</text></svg>'"><div class="name">${h.name}</div>`; it.innerHTML=`<img loading="lazy" src="/train_src/${encodeURIComponent(h.name)}" onerror="this.style.objectFit='contain';this.src='data:image/svg+xml;utf8,<svg xmlns=%22http://www.w3.org/2000/svg%22 width=%22100%22 height=%22100%22><text x=%2250%22 y=%2255%22 text-anchor=%22middle%22 fill=%22%23666%22 font-size=%2210%22>训练中</text></svg>'"><div class="name">${h.name}</div>`;
if(ready) it.onclick=()=>{ if(ready) it.onclick=()=>{
selectedHair=h.name; selectedHair=h.name;
$("sel-info").textContent="已选: "+h.name; $("sel-info").textContent="已选: "+h.name;
@@ -367,7 +366,6 @@ body { font-family: -apple-system, "Segoe UI", "Microsoft YaHei", sans-serif; ba
is_hr:$("sw-is_hr").classList.contains("on"), is_hr:$("sw-is_hr").classList.contains("on"),
dilate_kernel:[parseInt($("dk0").value),parseInt($("dk1").value)], dilate_kernel:[parseInt($("dk0").value),parseInt($("dk1").value)],
feather_px:parseInt($("fp").value), feather_px:parseInt($("fp").value),
feather_alpha:parseFloat($("fa").value),
denoising_strength:parseFloat($("den").value), denoising_strength:parseFloat($("den").value),
enhance:$("sw-enhance").classList.contains("on"), enhance:$("sw-enhance").classList.contains("on"),
enhance_denoising:parseFloat($("ed").value), enhance_denoising:parseFloat($("ed").value),
images/girl/长 -椭圆.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 508 KiB

images/girl/长 -波浪.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 480 KiB

images/girl/长 -直线.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 513 KiB

images/girl/长 -花瓣.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 493 KiB

images/girl/长-心形.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 463 KiB

prepare_train_data.py
+4 -4
View File
@@ -6,9 +6,9 @@
原始图 → 人脸检测+1k关键点 → Generator_Matte头发抠图 → 白底合成 + 居中裁剪 → 打标签 原始图 → 人脸检测+1k关键点 → Generator_Matte头发抠图 → 白底合成 + 居中裁剪 → 打标签
用法: 用法:
cd /home/xsl/change_hair/project/hair_service_sd cd /home/ubuntu/change_hair/project/hair_service_sd
/home/xsl/miniconda3/envs/my_hair/bin/python /home/xsl/change_hair/prepare_train_data.py \ /home/ubuntu/miniconda3/envs/my_hair/bin/python /home/ubuntu/change_hair/prepare_train_data.py \
--input /home/xsl/change_hair/train_images \ --input /home/ubuntu/change_hair/train_images \
--hair-id new_hairstyle_001 \ --hair-id new_hairstyle_001 \
--gender boy --gender boy
""" """
@@ -22,7 +22,7 @@ import argparse
ssl._create_default_https_context = ssl._create_unverified_context ssl._create_default_https_context = ssl._create_unverified_context
# 必须在 hair_service_sd 目录运行（依赖其模块） # 必须在 hair_service_sd 目录运行（依赖其模块）
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd" HAIR_SERVICE_DIR = "/home/ubuntu/change_hair/project/hair_service_sd"
os.chdir(HAIR_SERVICE_DIR) os.chdir(HAIR_SERVICE_DIR)
sys.path.insert(0, HAIR_SERVICE_DIR) sys.path.insert(0, HAIR_SERVICE_DIR)
os.environ.setdefault("HF_HUB_OFFLINE", "1") os.environ.setdefault("HF_HUB_OFFLINE", "1")
project/hair_service_sd/change_color.py
+5 -1
View File
@@ -70,7 +70,11 @@ def process_infer(user_img_path, target_color, color_dir, dst_path, ratio=1.0):
# 通过webui对照片进行增强 # 通过webui对照片进行增强
s4 = time.time() s4 = time.time()
# cv2.imwrite('/home/student/Downloads/12121.jpg',crop_img) # cv2.imwrite('/home/student/Downloads/12121.jpg',crop_img)
enhanced_img = enhance_hair.webui_img2img(crop_img, crop_mask, prompt=prompt) try:
enhanced_img = enhance_hair.webui_img2img(crop_img, crop_mask, prompt=prompt)
except Exception as e:
print(f"[process_infer] webui enhance skipped: {e}")
enhanced_img = crop_img
print("----------------------------webui_img2img", time.time() - s4) print("----------------------------webui_img2img", time.time() - s4)
# 写回原图 # 写回原图
project/hair_service_sd/config/configure.ini
+15 -16
View File
@@ -1,21 +1,21 @@
[default] [default]
modelDir = weights modelDir = weights
hairstyleDir = /home/xsl/change_hair/project/data/ref_hairstyle hairstyleDir = /home/ubuntu/change_hair/project/data/ref_hairstyle
haircolorDir= /home/xsl/change_hair/project/data/ref_haircolor haircolorDir= /home/ubuntu/change_hair/project/data/ref_haircolor
userDir=/home/xsl/change_hair/project/data/userImage userDir=/home/ubuntu/change_hair/project/data/userImage
tmp_dir=/home/xsl/change_hair/project/data/tmp tmp_dir=/home/ubuntu/change_hair/project/data/tmp
res_dir=/home/xsl/change_hair/project/data/res_dir res_dir=/home/ubuntu/change_hair/project/data/res_dir
userInfo_dir=/home/xsl/change_hair/project/data/user_info userInfo_dir=/home/ubuntu/change_hair/project/data/user_info
baseColor_ID=HDR10_443322 baseColor_ID=HDR10_443322
Port = 11023 Port = 11023
refer_dir = /home/xsl/change_hair/project/data/ref_online refer_dir = /home/ubuntu/change_hair/project/data/ref_online
ref_user_dir = /home/xsl/change_hair/project/data/ref_user_imgs ref_user_dir = /home/ubuntu/change_hair/project/data/ref_user_imgs
train_dir = /home/xsl/change_hair/data/train_material train_dir = /home/ubuntu/change_hair/data/train_material
refImgDir=/home/xsl/change_hair/project/data/refImage refImgDir=/home/ubuntu/change_hair/project/data/refImage
hair_template_material_dir=/home/xsl/change_hair/project/data/hair_template_material hair_template_material_dir=/home/ubuntu/change_hair/project/data/hair_template_material
ref_color=/home/xsl/change_hair/project/data/ref_color ref_color=/home/ubuntu/change_hair/project/data/ref_color
ref_color_img=/home/xsl/change_hair/project/data/ref_color_imgs ref_color_img=/home/ubuntu/change_hair/project/data/ref_color_imgs
upload_train_dir=/home/xsl/change_hair/project/data/upload_train_imgs upload_train_dir=/home/ubuntu/change_hair/project/data/upload_train_imgs
;version=local ;version=local
version=online version=online
@@ -23,7 +23,7 @@ version=online
strength=1 strength=1
[logger] [logger]
logpath = /home/xsl/change_hair/project/logs logpath = /home/ubuntu/change_hair/project/logs
level=INFO level=INFO
[timelogger] [timelogger]
@@ -32,4 +32,3 @@ name=watch-time
[errorlogger] [errorlogger]
level=ERROR level=ERROR
name=w-error name=w-error
project/hair_service_sd/core/cos_module.py
+2 -5
View File
@@ -20,11 +20,8 @@ class COS_object():
token = None # 如果使用永久密钥不需要填入 token，如果使用临时密钥需要填入，临时密钥生成和使用指引参见 https://cloud.tencent.com/document/product/436/14048 token = None # 如果使用永久密钥不需要填入 token，如果使用临时密钥需要填入，临时密钥生成和使用指引参见 https://cloud.tencent.com/document/product/436/14048
self.scheme = 'https' # 指定使用 http/https 协议来访问 COS，默认为 https，可不填 self.scheme = 'https' # 指定使用 http/https 协议来访问 COS，默认为 https，可不填
self.BucketName= os.getenv('COS_BUCKET', '<your-cos-bucket>') self.BucketName= os.getenv('COS_BUCKET', '<your-cos-bucket>')
# 密钥未设置时标记不可用（不抛异常） for param in (secret_id, secret_key, self.region, self.BucketName):
if '<' in secret_id or '<' in secret_key: assert '<' not in param, '请设置环境变量 COS_SECRET_ID / COS_SECRET_KEY / COS_REGION / COS_BUCKET'
print('[COS] ⚠️ 未设置 COS 密钥环境变量，COS 上传功能不可用')
self.client = None
return
config = CosConfig(Region=self.region, SecretId=secret_id, SecretKey=secret_key, Token=token, Scheme=self.scheme) config = CosConfig(Region=self.region, SecretId=secret_id, SecretKey=secret_key, Token=token, Scheme=self.scheme)
self.client = CosS3Client(config) self.client = CosS3Client(config)
project/hair_service_sd/core/hairstyle_model.py
+106 -39
View File
@@ -90,6 +90,14 @@ class HairStyle_Model(object):
self.keypoints_processor = hair_init.keypoints_processor self.keypoints_processor = hair_init.keypoints_processor
self.human_keypoint = hair_init.human_keypoint self.human_keypoint = hair_init.human_keypoint
# 用户图预处理产物内存缓存（一级缓存，避免重复 GPU 推理 + 磁盘 IO）。
# key = 图片字节哈希 + ratio；value = 7 个产物 + landmarks_1k 的 dict。
# 接口2 女性多发型场景：同一张用户图连续请求，第二次起命中缓存省 ~1.6s。
import collections
self._user_prepare_cache = {} # {cache_key: {产物dict}}
self._user_prepare_cache_keys = collections.deque() # LRU 顺序
self._USER_CACHE_MAX = 8 # 最多缓存 8 张图（防止内存膨胀）
# worker_id = int(os.environ.get('APP_WORKER_ID', 1)) # worker_id = int(os.environ.get('APP_WORKER_ID', 1))
# rand_max = 9527 # rand_max = 9527
@@ -2268,48 +2276,99 @@ class HairStyle_Model(object):
ref_landmark_f1k2_768) ref_landmark_f1k2_768)
cv2.imwrite(another_pose_hair_image_dir, another_pose_hair_image*255) cv2.imwrite(another_pose_hair_image_dir, another_pose_hair_image*255)
landmark1k_dir = osp.join(userinfo_dir, 'kpt_1k.txt') landmark1k_dir = osp.join(userinfo_dir, 'kpt_1k.txt')
if not osp.exists(landmark1k_dir):
landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward_diy(user_rgb_8uc3_orisize) # ===== 内存缓存（一级）：key = 图片哈希 + ratio =====
# 命中则跳过 landmark 检测 + get_prepare_user_768_data 整条 GPU 管线（省 ~1.6s）
import hashlib as _hashlib
_img_hash = _hashlib.md5(user_rgb_8uc3_orisize.tobytes()).hexdigest()[:16]
_cache_key = f"{_img_hash}_r{ratio}"
_cached = self._user_prepare_cache.get(_cache_key)
if _cached is not None:
# 命中内存缓存：直接取所有产物，跳过 landmark 检测和预处理管线
landmarks_origin_img_1k = _cached['landmarks_1k']
user_bald_res_8uc3_orisize = _cached['bald_res']
user_baldseg_8uc3_orisize = _cached['baldseg_ori']
user_baldseg_8uc3_768 = _cached['baldseg_768']
user_bald_8uc3_768 = _cached['bald_768']
user_landmark_f1k2_768 = _cached['lmk_768']
user_hairstyle_M = _cached['M']
user_matting_8uc3_bald_orisize = _cached['matting_ori']
# 补写磁盘文件：下游代码（功能7等）仍从 userinfo_dir 读这些文件，
# 而 task_id 每次不同导致 userinfo_dir 不同，必须补写保证下游可用
os.makedirs(userinfo_dir, exist_ok=True)
np.savetxt(landmark1k_dir, landmarks_origin_img_1k)
cv2.imwrite(osp.join(userinfo_dir, 'bald_res_ori.png'), user_bald_res_8uc3_orisize)
cv2.imwrite(osp.join(userinfo_dir, 'bald_seg_ori.png'), user_baldseg_8uc3_orisize)
cv2.imwrite(osp.join(userinfo_dir, 'user_baldseg_768.png'), user_baldseg_8uc3_768)
cv2.imwrite(osp.join(userinfo_dir, 'bald_seg_768.png'), user_bald_8uc3_768)
np.savetxt(osp.join(userinfo_dir, 'landmark_f1k2_768.txt'), user_landmark_f1k2_768)
np.savetxt(osp.join(userinfo_dir, 'hairstyle_M.txt'), user_hairstyle_M)
if user_matting_8uc3_bald_orisize is not None:
cv2.imwrite(osp.join(userinfo_dir, 'user_orig_mask.png'), user_matting_8uc3_bald_orisize)
self.logger_process.info(f"内存缓存命中 key={_cache_key}，跳过用户图预处理(补写磁盘文件)")
else:
# 未命中：走原有逻辑（landmark 检测 + 预处理管线 + 磁盘缓存）
if not osp.exists(landmark1k_dir):
landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward_diy(user_rgb_8uc3_orisize)
if landmarks_origin_img_1k is None:
return None, 10001, None, None, None
np.savetxt(landmark1k_dir, landmarks_origin_img_1k)
else:
landmarks_origin_img_1k = np.loadtxt(landmark1k_dir)
# landmarks_origin_img_1k, _, _ = self.get_landmark.forward(user_rgb_8uc3_orisize)
if landmarks_origin_img_1k is None: if landmarks_origin_img_1k is None:
return None, 10001, None, None, None return None, 10001, None, None, None
np.savetxt(landmark1k_dir, landmarks_origin_img_1k)
else:
landmarks_origin_img_1k = np.loadtxt(landmark1k_dir)
# landmarks_origin_img_1k, _, _ = self.get_landmark.forward(user_rgb_8uc3_orisize)
if landmarks_origin_img_1k is None:
return None, 10001, None, None, None
user_bald_res_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_res_ori.png') user_bald_res_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_res_ori.png')
user_baldseg_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_seg_ori.png') user_baldseg_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_seg_ori.png')
user_baldseg_8uc3_768_dir = osp.join(userinfo_dir, 'user_baldseg_768.png') user_baldseg_8uc3_768_dir = osp.join(userinfo_dir, 'user_baldseg_768.png')
user_bald_8uc3_768_dir = osp.join(userinfo_dir, 'bald_seg_768.png') user_bald_8uc3_768_dir = osp.join(userinfo_dir, 'bald_seg_768.png')
user_landmark_f1k2_768_dir = osp.join(userinfo_dir, 'landmark_f1k2_768.txt') user_landmark_f1k2_768_dir = osp.join(userinfo_dir, 'landmark_f1k2_768.txt')
user_hairstyle_M_dir = osp.join(userinfo_dir, 'hairstyle_M.txt') user_hairstyle_M_dir = osp.join(userinfo_dir, 'hairstyle_M.txt')
condition_exist2 = True condition_exist2 = True
pre_list = [user_bald_res_8uc3_orisize_dir, user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_768_dir, user_bald_8uc3_768_dir, pre_list = [user_bald_res_8uc3_orisize_dir, user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_768_dir, user_bald_8uc3_768_dir,
user_landmark_f1k2_768_dir, user_hairstyle_M_dir] user_landmark_f1k2_768_dir, user_hairstyle_M_dir]
for tmp_dir in pre_list: for tmp_dir in pre_list:
if not osp.exists(tmp_dir): if not osp.exists(tmp_dir):
condition_exist2 = False condition_exist2 = False
user_matting_8uc3_bald_orisize = None user_matting_8uc3_bald_orisize = None
if not condition_exist2: if not condition_exist2:
user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, user_baldseg_8uc3_768, user_bald_8uc3_768, \ user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, user_baldseg_8uc3_768, user_bald_8uc3_768, \
user_landmark_f1k2_768, user_hairstyle_M, user_matting_8uc3_bald_orisize = self.process_data.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio) user_landmark_f1k2_768, user_hairstyle_M, user_matting_8uc3_bald_orisize = self.process_data.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio)
cv2.imwrite(user_bald_res_8uc3_orisize_dir, user_bald_res_8uc3_orisize) cv2.imwrite(user_bald_res_8uc3_orisize_dir, user_bald_res_8uc3_orisize)
cv2.imwrite(user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_orisize) cv2.imwrite(user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_orisize)
cv2.imwrite(user_baldseg_8uc3_768_dir, user_baldseg_8uc3_768) cv2.imwrite(user_baldseg_8uc3_768_dir, user_baldseg_8uc3_768)
cv2.imwrite(user_bald_8uc3_768_dir, user_bald_8uc3_768) cv2.imwrite(user_bald_8uc3_768_dir, user_bald_8uc3_768)
np.savetxt(user_landmark_f1k2_768_dir, user_landmark_f1k2_768) np.savetxt(user_landmark_f1k2_768_dir, user_landmark_f1k2_768)
np.savetxt(user_hairstyle_M_dir, user_hairstyle_M) np.savetxt(user_hairstyle_M_dir, user_hairstyle_M)
else: else:
user_bald_res_8uc3_orisize = cv2.imread(user_bald_res_8uc3_orisize_dir) user_bald_res_8uc3_orisize = cv2.imread(user_bald_res_8uc3_orisize_dir)
user_baldseg_8uc3_orisize = cv2.imread(user_baldseg_8uc3_orisize_dir) user_baldseg_8uc3_orisize = cv2.imread(user_baldseg_8uc3_orisize_dir)
user_baldseg_8uc3_768 = cv2.imread(user_baldseg_8uc3_768_dir) user_baldseg_8uc3_768 = cv2.imread(user_baldseg_8uc3_768_dir)
user_bald_8uc3_768 = cv2.imread(user_bald_8uc3_768_dir) user_bald_8uc3_768 = cv2.imread(user_bald_8uc3_768_dir)
user_landmark_f1k2_768 = np.loadtxt(user_landmark_f1k2_768_dir) user_landmark_f1k2_768 = np.loadtxt(user_landmark_f1k2_768_dir)
user_hairstyle_M = np.loadtxt(user_hairstyle_M_dir) user_hairstyle_M = np.loadtxt(user_hairstyle_M_dir)
# 写入内存缓存（含 landmarks_1k 和 matting，供后续同图请求命中）
self._user_prepare_cache[_cache_key] = {
'landmarks_1k': landmarks_origin_img_1k,
'bald_res': user_bald_res_8uc3_orisize,
'baldseg_ori': user_baldseg_8uc3_orisize,
'baldseg_768': user_baldseg_8uc3_768,
'bald_768': user_bald_8uc3_768,
'lmk_768': user_landmark_f1k2_768,
'M': user_hairstyle_M,
'matting_ori': user_matting_8uc3_bald_orisize,
}
self._user_prepare_cache_keys.append(_cache_key)
# LRU 淘汰：超过上限删除最老的
while len(self._user_prepare_cache_keys) > self._USER_CACHE_MAX:
_old = self._user_prepare_cache_keys.popleft()
self._user_prepare_cache.pop(_old, None)
self.logger_process.info(f"内存缓存写入 key={_cache_key}，当前缓存 {len(self._user_prepare_cache)} 张")
# show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1) # show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1)
# ratio = 1536. / max(show_concat.shape[:2]) # ratio = 1536. / max(show_concat.shape[:2])
@@ -2318,18 +2377,26 @@ class HairStyle_Model(object):
# cv2.waitKey() # cv2.waitKey()
user_orig_mask_path = os.path.join(userinfo_dir, "user_orig_mask.png") user_orig_mask_path = os.path.join(userinfo_dir, "user_orig_mask.png")
if not os.path.exists(user_orig_mask_path): if user_matting_8uc3_bald_orisize is not None and not os.path.exists(user_orig_mask_path):
cv2.imwrite(user_orig_mask_path, user_matting_8uc3_bald_orisize) cv2.imwrite(user_orig_mask_path, user_matting_8uc3_bald_orisize)
# 换发型 # 换发型（主 GAN）
import time as _time
_t_gan0 = _time.perf_counter()
hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image, hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image,
user_baldseg_8uc3_768, user_baldseg_8uc3_768,
user_bald_8uc3_768, user_bald_8uc3_768,
user_landmark_f1k2_768, gender) user_landmark_f1k2_768, gender)
_t_gan = _time.perf_counter() - _t_gan0
# 融合（第3次matte + 融合GAN）
_t_fuse0 = _time.perf_counter()
hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste( hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste(
user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M) user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M)
_t_fuse = _time.perf_counter() - _t_fuse0
self.logger_process.info(
f"功能6分步计时: 主GAN={_t_gan:.3f}s 融合={_t_fuse:.3f}s (缓存={'命中' if _cached is not None else '未命中'})")
gen_hair_mask_path_2 = os.path.join(userinfo_dir, "hair_mask_2.png") gen_hair_mask_path_2 = os.path.join(userinfo_dir, "hair_mask_2.png")
cv2.imwrite(gen_hair_mask_path_2, hair_gene_matte_8uc3_orisize) cv2.imwrite(gen_hair_mask_path_2, hair_gene_matte_8uc3_orisize)
project/hair_service_sd/core/models/layers/data/FDDB/img_list.txt
+2845
View File
File diff suppressed because it is too large Load Diff
project/hair_service_sd/core/models/layers/data/__init__.py
+3
View File
@@ -0,0 +1,3 @@
from .wider_face import WiderFaceDetection, detection_collate
from .data_augment import *
from .config import *
project/hair_service_sd/core/models/layers/data/config.py
+42
View File
@@ -0,0 +1,42 @@
# config.py
cfg_mnet = {
'name': 'mobilenet0.25',
'min_sizes': [[16, 32], [64, 128], [256, 512]],
'steps': [8, 16, 32],
'variance': [0.1, 0.2],
'clip': False,
'loc_weight': 2.0,
'gpu_train': True,
'batch_size': 32,
'ngpu': 1,
'epoch': 250,
'decay1': 190,
'decay2': 220,
'image_size': 640,
'pretrain': True,
'return_layers': {'stage1': 1, 'stage2': 2, 'stage3': 3},
'in_channel': 32,
'out_channel': 64
}
cfg_re50 = {
'name': 'Resnet50',
'min_sizes': [[16, 32], [64, 128], [256, 512]],
'steps': [8, 16, 32],
'variance': [0.1, 0.2],
'clip': False,
'loc_weight': 2.0,
'gpu_train': True,
'batch_size': 24,
'ngpu': 4,
'epoch': 100,
'decay1': 70,
'decay2': 90,
'image_size': 840,
'pretrain': True,
'return_layers': {'layer2': 1, 'layer3': 2, 'layer4': 3},
'in_channel': 256,
'out_channel': 256
}
project/hair_service_sd/core/models/layers/data/data_augment.py
+237
View File
@@ -0,0 +1,237 @@
import cv2
import numpy as np
import random
from core.utils.box_utils_Retina import matrix_iof
def _crop(image, boxes, labels, landm, img_dim):
height, width, _ = image.shape
pad_image_flag = True
for _ in range(250):
"""
if random.uniform(0, 1) <= 0.2:
scale = 1.0
else:
scale = random.uniform(0.3, 1.0)
"""
PRE_SCALES = [0.3, 0.45, 0.6, 0.8, 1.0]
scale = random.choice(PRE_SCALES)
short_side = min(width, height)
w = int(scale * short_side)
h = w
if width == w:
l = 0
else:
l = random.randrange(width - w)
if height == h:
t = 0
else:
t = random.randrange(height - h)
roi = np.array((l, t, l + w, t + h))
value = matrix_iof(boxes, roi[np.newaxis])
flag = (value >= 1)
if not flag.any():
continue
centers = (boxes[:, :2] + boxes[:, 2:]) / 2
mask_a = np.logical_and(roi[:2] < centers, centers < roi[2:]).all(axis=1)
boxes_t = boxes[mask_a].copy()
labels_t = labels[mask_a].copy()
landms_t = landm[mask_a].copy()
landms_t = landms_t.reshape([-1, 5, 2])
if boxes_t.shape[0] == 0:
continue
image_t = image[roi[1]:roi[3], roi[0]:roi[2]]
boxes_t[:, :2] = np.maximum(boxes_t[:, :2], roi[:2])
boxes_t[:, :2] -= roi[:2]
boxes_t[:, 2:] = np.minimum(boxes_t[:, 2:], roi[2:])
boxes_t[:, 2:] -= roi[:2]
# landm
landms_t[:, :, :2] = landms_t[:, :, :2] - roi[:2]
landms_t[:, :, :2] = np.maximum(landms_t[:, :, :2], np.array([0, 0]))
landms_t[:, :, :2] = np.minimum(landms_t[:, :, :2], roi[2:] - roi[:2])
landms_t = landms_t.reshape([-1, 10])
# make sure that the cropped image contains at least one face > 16 pixel at training image scale
b_w_t = (boxes_t[:, 2] - boxes_t[:, 0] + 1) / w * img_dim
b_h_t = (boxes_t[:, 3] - boxes_t[:, 1] + 1) / h * img_dim
mask_b = np.minimum(b_w_t, b_h_t) > 0.0
boxes_t = boxes_t[mask_b]
labels_t = labels_t[mask_b]
landms_t = landms_t[mask_b]
if boxes_t.shape[0] == 0:
continue
pad_image_flag = False
return image_t, boxes_t, labels_t, landms_t, pad_image_flag
return image, boxes, labels, landm, pad_image_flag
def _distort(image):
def _convert(image, alpha=1, beta=0):
tmp = image.astype(float) * alpha + beta
tmp[tmp < 0] = 0
tmp[tmp > 255] = 255
image[:] = tmp
image = image.copy()
if random.randrange(2):
#brightness distortion
if random.randrange(2):
_convert(image, beta=random.uniform(-32, 32))
#contrast distortion
if random.randrange(2):
_convert(image, alpha=random.uniform(0.5, 1.5))
image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
#saturation distortion
if random.randrange(2):
_convert(image[:, :, 1], alpha=random.uniform(0.5, 1.5))
#hue distortion
if random.randrange(2):
tmp = image[:, :, 0].astype(int) + random.randint(-18, 18)
tmp %= 180
image[:, :, 0] = tmp
image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
else:
#brightness distortion
if random.randrange(2):
_convert(image, beta=random.uniform(-32, 32))
image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
#saturation distortion
if random.randrange(2):
_convert(image[:, :, 1], alpha=random.uniform(0.5, 1.5))
#hue distortion
if random.randrange(2):
tmp = image[:, :, 0].astype(int) + random.randint(-18, 18)
tmp %= 180
image[:, :, 0] = tmp
image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
#contrast distortion
if random.randrange(2):
_convert(image, alpha=random.uniform(0.5, 1.5))
return image
def _expand(image, boxes, fill, p):
if random.randrange(2):
return image, boxes
height, width, depth = image.shape
scale = random.uniform(1, p)
w = int(scale * width)
h = int(scale * height)
left = random.randint(0, w - width)
top = random.randint(0, h - height)
boxes_t = boxes.copy()
boxes_t[:, :2] += (left, top)
boxes_t[:, 2:] += (left, top)
expand_image = np.empty(
(h, w, depth),
dtype=image.dtype)
expand_image[:, :] = fill
expand_image[top:top + height, left:left + width] = image
image = expand_image
return image, boxes_t
def _mirror(image, boxes, landms):
_, width, _ = image.shape
if random.randrange(2):
image = image[:, ::-1]
boxes = boxes.copy()
boxes[:, 0::2] = width - boxes[:, 2::-2]
# landm
landms = landms.copy()
landms = landms.reshape([-1, 5, 2])
landms[:, :, 0] = width - landms[:, :, 0]
tmp = landms[:, 1, :].copy()
landms[:, 1, :] = landms[:, 0, :]
landms[:, 0, :] = tmp
tmp1 = landms[:, 4, :].copy()
landms[:, 4, :] = landms[:, 3, :]
landms[:, 3, :] = tmp1
landms = landms.reshape([-1, 10])
return image, boxes, landms
def _pad_to_square(image, rgb_mean, pad_image_flag):
if not pad_image_flag:
return image
height, width, _ = image.shape
long_side = max(width, height)
image_t = np.empty((long_side, long_side, 3), dtype=image.dtype)
image_t[:, :] = rgb_mean
image_t[0:0 + height, 0:0 + width] = image
return image_t
def _resize_subtract_mean(image, insize, rgb_mean):
interp_methods = [cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_NEAREST, cv2.INTER_LANCZOS4]
interp_method = interp_methods[random.randrange(5)]
image = cv2.resize(image, (insize, insize), interpolation=interp_method)
image = image.astype(np.float32)
image -= rgb_mean
return image.transpose(2, 0, 1)
class preproc(object):
def __init__(self, img_dim, rgb_means):
self.img_dim = img_dim
self.rgb_means = rgb_means
def __call__(self, image, targets):
assert targets.shape[0] > 0, "this image does not have gt"
boxes = targets[:, :4].copy()
labels = targets[:, -1].copy()
landm = targets[:, 4:-1].copy()
image_t, boxes_t, labels_t, landm_t, pad_image_flag = _crop(image, boxes, labels, landm, self.img_dim)
image_t = _distort(image_t)
image_t = _pad_to_square(image_t,self.rgb_means, pad_image_flag)
image_t, boxes_t, landm_t = _mirror(image_t, boxes_t, landm_t)
height, width, _ = image_t.shape
image_t = _resize_subtract_mean(image_t, self.img_dim, self.rgb_means)
boxes_t[:, 0::2] /= width
boxes_t[:, 1::2] /= height
landm_t[:, 0::2] /= width
landm_t[:, 1::2] /= height
labels_t = np.expand_dims(labels_t, 1)
targets_t = np.hstack((boxes_t, landm_t, labels_t))
return image_t, targets_t
project/hair_service_sd/core/models/layers/data/wider_face.py
+101
View File
@@ -0,0 +1,101 @@
import os
import os.path
import sys
import torch
import torch.utils.data as data
import cv2
import numpy as np
class WiderFaceDetection(data.Dataset):
def __init__(self, txt_path, preproc=None):
self.preproc = preproc
self.imgs_path = []
self.words = []
f = open(txt_path,'r')
lines = f.readlines()
isFirst = True
labels = []
for line in lines:
line = line.rstrip()
if line.startswith('#'):
if isFirst is True:
isFirst = False
else:
labels_copy = labels.copy()
self.words.append(labels_copy)
labels.clear()
path = line[2:]
path = txt_path.replace('label.txt','images/') + path
self.imgs_path.append(path)
else:
line = line.split(' ')
label = [float(x) for x in line]
labels.append(label)
self.words.append(labels)
def __len__(self):
return len(self.imgs_path)
def __getitem__(self, index):
img = cv2.imread(self.imgs_path[index])
height, width, _ = img.shape
labels = self.words[index]
annotations = np.zeros((0, 15))
if len(labels) == 0:
return annotations
for idx, label in enumerate(labels):
annotation = np.zeros((1, 15))
# bbox
annotation[0, 0] = label[0] # x1
annotation[0, 1] = label[1] # y1
annotation[0, 2] = label[0] + label[2] # x2
annotation[0, 3] = label[1] + label[3] # y2
# landmarks
annotation[0, 4] = label[4] # l0_x
annotation[0, 5] = label[5] # l0_y
annotation[0, 6] = label[7] # l1_x
annotation[0, 7] = label[8] # l1_y
annotation[0, 8] = label[10] # l2_x
annotation[0, 9] = label[11] # l2_y
annotation[0, 10] = label[13] # l3_x
annotation[0, 11] = label[14] # l3_y
annotation[0, 12] = label[16] # l4_x
annotation[0, 13] = label[17] # l4_y
if (annotation[0, 4]<0):
annotation[0, 14] = -1
else:
annotation[0, 14] = 1
annotations = np.append(annotations, annotation, axis=0)
target = np.array(annotations)
if self.preproc is not None:
img, target = self.preproc(img, target)
return torch.from_numpy(img), target
def detection_collate(batch):
"""Custom collate fn for dealing with batches of images that have a different
number of associated object annotations (bounding boxes).
Arguments:
batch: (tuple) A tuple of tensor images and lists of annotations
Return:
A tuple containing:
1) (tensor) batch of images stacked on their 0 dim
2) (list of tensors) annotations for a given image are stacked on 0 dim
"""
targets = []
imgs = []
for _, sample in enumerate(batch):
for _, tup in enumerate(sample):
if torch.is_tensor(tup):
imgs.append(tup)
elif isinstance(tup, type(np.empty(0))):
annos = torch.from_numpy(tup).float()
targets.append(annos)
return (torch.stack(imgs, 0), targets)
project/hair_service_sd/core/oss_module.py
+17 -20
View File
@@ -1,31 +1,28 @@
import time import time
import os
try: import oss2
import oss2 import os
_OSS2_AVAILABLE = True
except Exception as _e:
print(f'[OSS] ⚠️ oss2 导入失败({_e})，OSS 上传功能不可用')
_OSS2_AVAILABLE = False
class OSS_object(): class OSS_object():
def __init__(self): def __init__(self):
if not _OSS2_AVAILABLE: # 懒加载：只读取凭证，不立即连接 OSS。
print('[OSS] oss2 不可用，OSS 功能关闭') # 这样本地用 output_format=base64 时不配 OSS 凭证也能启动服务；
self.bucket = None # 仅在真正调用 upload_file 时才校验并连接。
self.access_key_id = os.getenv('OSS_TEST_ACCESS_KEY_ID', '<your-access-key-id>')
self.access_key_secret = os.getenv('OSS_TEST_ACCESS_KEY_SECRET', '<your-access-key-secret>')
self.bucket_name = os.getenv('OSS_TEST_BUCKET', '<your-bucket-name>')
self.endpoint = os.getenv('OSS_TEST_ENDPOINT', '<your-endpoint>')
self.bucket = None
def _ensure_bucket(self):
if self.bucket is not None:
return return
access_key_id = os.getenv('OSS_TEST_ACCESS_KEY_ID', '<your-access-key-id>') for param in (self.access_key_id, self.access_key_secret, self.bucket_name, self.endpoint):
access_key_secret = os.getenv('OSS_TEST_ACCESS_KEY_SECRET', '<your-access-key-secret>') assert '<' not in param, '请设置环境变量 OSS_TEST_ACCESS_KEY_ID / OSS_TEST_ACCESS_KEY_SECRET / OSS_TEST_BUCKET / OSS_TEST_ENDPOINT'
bucket_name = os.getenv('OSS_TEST_BUCKET', '<your-bucket-name>') self.bucket = oss2.Bucket(oss2.Auth(self.access_key_id, self.access_key_secret), self.endpoint, self.bucket_name)
endpoint = os.getenv('OSS_TEST_ENDPOINT', '<your-endpoint>')
# 密钥未设置时标记不可用（不抛异常，避免服务启动/换发型时崩溃）
if '<' in access_key_id or '<' in access_key_secret:
print('[OSS] ⚠️ 未设置 OSS 密钥环境变量，OSS 上传功能不可用（本地测试用 base64 返回即可）')
self.bucket = None
return
self.bucket = oss2.Bucket(oss2.Auth(access_key_id, access_key_secret), endpoint, bucket_name)
def upload_file(self, file, target_name): def upload_file(self, file, target_name):
self._ensure_bucket()
t0 = time.time() t0 = time.time()
with open(oss2.to_unicode(file), 'rb') as f: with open(oss2.to_unicode(file), 'rb') as f:
ret = self.bucket.put_object(target_name, f) ret = self.bucket.put_object(target_name, f)
project/hair_service_sd/gen_super_image.py
+5 -4
View File
@@ -1,4 +1,5 @@
import io import io
import os
import os.path import os.path
import time import time
@@ -83,13 +84,13 @@ class ControlnetRequestImg2Img:
return encoded_image return encoded_image
def build_body_v2(self, dst_width, dst_height, cfg_scale, base_img, denoising_strength=0.7): def build_body_v2(self, dst_width, dst_height, cfg_scale, base_img, denoising_strength=0.7, steps=None):
self.body = { self.body = {
"prompt": self.prompt, "prompt": self.prompt,
"negative_prompt": self.neg_prompt, "negative_prompt": self.neg_prompt,
"sampler_name": "DPM++ 2M Karras", "sampler_name": "DPM++ 2M Karras",
"batch_size": 1, "batch_size": 1,
"steps": 20, "steps": int(steps) if steps is not None else int(os.environ.get("WEBUI_STEPS", "15")),
"width": dst_width, "width": dst_width,
"height": dst_height, "height": dst_height,
"cfg_scale": cfg_scale, "cfg_scale": cfg_scale,
@@ -413,7 +414,7 @@ def encode_numpy_to_base64(img):
encoded_image = base64.b64encode(bytes).decode('utf-8') encoded_image = base64.b64encode(bytes).decode('utf-8')
return encoded_image return encoded_image
def webui_img2img(img=None, mask_img=None, in_gender=None, task_id=None, hair_id=None, lora_material_path=None, tag="", is_hr=False, denoising_strength=0.7, inference_port="57860", refiner_switch_at=0.5): def webui_img2img(img=None, mask_img=None, in_gender=None, task_id=None, hair_id=None, lora_material_path=None, tag="", is_hr=False, denoising_strength=0.7, inference_port="57860", refiner_switch_at=0.5, webui_steps=None):
# url = "http://hairservice.tslead.net:57860/sdapi/v1/img2img" # url = "http://hairservice.tslead.net:57860/sdapi/v1/img2img"
neg_prompt = '(nsfw:1.5), ng_deepnegative_v1_75t, (badhandv4:1.2), (worst quality:2), (low quality:2), (normal quality:2), lowres, bad anatomy, bad hands, ((monochrome)), ((grayscale)) watermark, moles, large breast, big breast, bad_pictures,easynegative' neg_prompt = '(nsfw:1.5), ng_deepnegative_v1_75t, (badhandv4:1.2), (worst quality:2), (low quality:2), (normal quality:2), lowres, bad anatomy, bad hands, ((monochrome)), ((grayscale)) watermark, moles, large breast, big breast, bad_pictures,easynegative'
@@ -429,7 +430,7 @@ def webui_img2img(img=None, mask_img=None, in_gender=None, task_id=None, hair_id
control_net = ControlnetRequestImg2Img(prompt, neg_prompt, mask_img) control_net = ControlnetRequestImg2Img(prompt, neg_prompt, mask_img)
# control_net.build_body_hr(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength) # control_net.build_body_hr(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength)
if not is_hr: if not is_hr:
control_net.build_body_v2(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength) control_net.build_body_v2(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength, steps=webui_steps)
else: else:
control_net.build_body_hr(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength, refiner_switch_at=refiner_switch_at) control_net.build_body_hr(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength, refiner_switch_at=refiner_switch_at)
project/hair_service_sd/gunicorn_config.py
+1 -1
View File
@@ -1,4 +1,4 @@
#!/home/ubuntu/miniconda3/envs/yd/bin/python #!/home/xsl/miniconda3/envs/yd/bin/python
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
import os import os
# 监听本机的端口 # 监听本机的端口
project/hair_service_sd/hair_grow_swap.py
+7
View File
@@ -177,10 +177,17 @@ def list_hairstyles(hairstyle_dir, train_dir, upload_dir, limit=None):
返回: 返回:
[{hair_id, gender, has_lora}, ...] [{hair_id, gender, has_lora}, ...]
""" """
# 仅展示这些发型（其余数据保留在磁盘，但不在列表中显示）
_VISIBLE_HAIRSTYLES = {
"chang_tuoyuan", "chang_bolang", "chang_zhixian",
"chang_huaban", "chang_xinxing",
}
styles = [] styles = []
if not os.path.isdir(hairstyle_dir): if not os.path.isdir(hairstyle_dir):
return styles return styles
for hair_id in os.listdir(hairstyle_dir): for hair_id in os.listdir(hairstyle_dir):
if hair_id not in _VISIBLE_HAIRSTYLES:
continue
cfg_path = os.path.join(hairstyle_dir, hair_id, "config.json") cfg_path = os.path.join(hairstyle_dir, hair_id, "config.json")
ref_path = os.path.join(hairstyle_dir, hair_id, "ref_rgb_8uc3_768.png") ref_path = os.path.join(hairstyle_dir, hair_id, "ref_rgb_8uc3_768.png")
lora_path = os.path.join(train_dir, hair_id, "model", "hairstyle_hd_lora.safetensors") lora_path = os.path.join(train_dir, hair_id, "model", "hairstyle_hd_lora.safetensors")
project/hair_service_sd/hair_swap_manual.py
+12 -46
View File
@@ -36,8 +36,7 @@ def hair_swap_manual(origin_img, hand_mask, hair_id, hairstyle_process, landmark
denoising_strength=0.6, # 重绘强度（唯一能透传到webui的SD参数） denoising_strength=0.6, # 重绘强度（唯一能透传到webui的SD参数）
# ===== 贴回/融合（真正生效）===== # ===== 贴回/融合（真正生效）=====
blend_dilate=(5, 5), # 步骤⑥ strict mask贴回时mask膨胀核 blend_dilate=(5, 5), # 步骤⑥ strict mask贴回时mask膨胀核
feather_px=0, # ★ strict贴回边缘羽化范围（像素，边缘带宽度，0=无羽化） feather_px=0, # ★ strict贴回边缘羽化像素（0=无羽化硬边缘，>0=高斯模糊边缘）
feather_alpha=1.0, # ★ 羽化强度（0~1，0=硬切，1=最大羽化，控制边缘过渡柔和度）
seamless_dilate=(9, 9), # 步骤⑦ 泊松融合mask膨胀核 seamless_dilate=(9, 9), # 步骤⑦ 泊松融合mask膨胀核
# ===== enhance 二次增强重绘 ===== # ===== enhance 二次增强重绘 =====
enhance=False, # 步骤⑧ 是否对结果再做一次低强度SD重绘（让发丝更清晰） enhance=False, # 步骤⑧ 是否对结果再做一次低强度SD重绘（让发丝更清晰）
@@ -231,20 +230,10 @@ def hair_swap_manual(origin_img, hand_mask, hair_id, hairstyle_process, landmark
result_strict_hard[mask_back == 0] = origin_img[mask_back == 0] result_strict_hard[mask_back == 0] = origin_img[mask_back == 0]
if feather_px and feather_px > 0: if feather_px and feather_px > 0:
# ★ 只羽化边缘：mask 内部保持硬切(255)，只有边缘带做渐变
# 1. 对 mask_back 做高斯模糊（得到边缘的渐变）
k = max(1, int(feather_px)) * 2 + 1 # 高斯核必须是奇数 k = max(1, int(feather_px)) * 2 + 1 # 高斯核必须是奇数
mask_blur = cv2.GaussianBlur(mask_back, (k, k), 0) mask_back_blur = cv2.GaussianBlur(mask_back, (k, k), 0)
# 2. 取 mask 内部区域（离边缘 > feather_px 的部分），强制为255（不羽化） mask_blend = (mask_back_blur.astype(np.float32) / 255)[..., None]
erode_kernel = np.ones((k, k), np.uint8) feather_info = f"，边缘羽化{feather_px}px"
mask_inner = cv2.erode(mask_back, erode_kernel)
# 3. 合成：内部硬切(255) + 边缘用模糊值
mask_blend_mask = np.maximum(mask_inner, mask_blur).astype(np.float32)
# 4. ★ feather_alpha 控制羽化强度：在「全硬切(原mask)」和「边缘渐变(mask_blend_mask)」之间插值
# alpha=0 → 完全硬切（无羽化），alpha=1 → 最大羽化
mask_blend_mask = mask_back.astype(np.float32) * (1 - feather_alpha) + mask_blend_mask * feather_alpha
mask_blend = (mask_blend_mask / 255)[..., None]
feather_info = f"，边缘羽化{feather_px}px 强度{feather_alpha:.2f}"
else: else:
mask_blend = hard_blend mask_blend = hard_blend
result_strict = (sd_result_back.astype(np.float32) * mask_blend + result_strict = (sd_result_back.astype(np.float32) * mask_blend +
@@ -252,38 +241,17 @@ def hair_swap_manual(origin_img, hand_mask, hair_id, hairstyle_process, landmark
result_strict = np.clip(result_strict, 0, 255).astype(np.uint8) result_strict = np.clip(result_strict, 0, 255).astype(np.uint8)
result_strict[mask_back == 0] = origin_img[mask_back == 0] result_strict[mask_back == 0] = origin_img[mask_back == 0]
# ★ 羽化对比图（黑底版）：纯黑背景 + 贴回图，凸显羽化边缘过渡 # ★ 羽化对比图：左=硬边缘(feather=0) 右=当前羽化，裁剪mask边界附近放大
# 用黑色背景而非原图，避免原图干扰，纯粹对比羽化效果
feather_compare = None feather_compare = None
feather_compare_full = None
if feather_px and feather_px > 0: if feather_px and feather_px > 0:
# 计算 mask_blend（用于把贴回图叠加到黑底上）
ys, xs = np.where(mask_back > 10) ys, xs = np.where(mask_back > 10)
if len(ys) > 0: if len(ys) > 0:
# 黑底版：贴回图 * mask_blend + 黑色(0) * (1-mask_blend) x1 = max(0, xs.min()-20); x2 = min(origin_img.shape[1], xs.max()+20)
# result_strict_hard 已是「贴回图+原图」混合，这里重新算黑底版 y1 = max(0, ys.min()-20); y2 = min(origin_img.shape[0], ys.max()+20)
black_bg = np.zeros_like(origin_img) hard_crop = result_strict_hard[y1:y2, x1:x2]
# 无羽化版（黑底） blur_crop = result_strict[y1:y2, x1:x2]
hard_on_black = (sd_result_back.astype(np.float32) * hard_blend + sep = np.full((hard_crop.shape[0], 3, 3), 128, dtype=np.uint8)
black_bg.astype(np.float32) * (1 - hard_blend))
hard_on_black = np.clip(hard_on_black, 0, 255).astype(np.uint8)
hard_on_black[mask_back == 0] = 0 # mask外纯黑
# 有羽化版（黑底）
blur_on_black = (sd_result_back.astype(np.float32) * mask_blend +
black_bg.astype(np.float32) * (1 - mask_blend))
blur_on_black = np.clip(blur_on_black, 0, 255).astype(np.uint8)
blur_on_black[mask_back == 0] = 0 # mask外纯黑
# 局部放大版（裁剪mask边界附近）
x1 = max(0, xs.min()-30); x2 = min(origin_img.shape[1], xs.max()+30)
y1 = max(0, ys.min()-30); y2 = min(origin_img.shape[0], ys.max()+30)
hard_crop = hard_on_black[y1:y2, x1:x2]
blur_crop = blur_on_black[y1:y2, x1:x2]
sep = np.full((hard_crop.shape[0], 3, 3), 60, dtype=np.uint8) # 深灰分隔线
feather_compare = np.hstack([hard_crop, sep, blur_crop]) feather_compare = np.hstack([hard_crop, sep, blur_crop])
# 完整图版（不裁剪）
sep_full = np.full((origin_img.shape[0], 3, 3), 60, dtype=np.uint8)
feather_compare_full = np.hstack([hard_on_black, sep_full, blur_on_black])
result = result_strict if strict_mask else result_full result = result_strict if strict_mask else result_full
@@ -294,9 +262,7 @@ def hair_swap_manual(origin_img, hand_mask, hair_id, hairstyle_process, landmark
{"label": f"当前选用({'严格mask' if strict_mask else '整框覆盖'})", "b64": _enc(result)}, {"label": f"当前选用({'严格mask' if strict_mask else '整框覆盖'})", "b64": _enc(result)},
] ]
if feather_compare is not None: if feather_compare is not None:
step6_images.append({"label": f"★ 黑底羽化对比·边缘(左=无羽化 右=羽化{feather_px}pxα{feather_alpha:.1f})", "b64": _enc(feather_compare)}) step6_images.append({"label": f"★ 羽化对比(左=硬边 右=羽化{feather_px}px)", "b64": _enc(feather_compare)})
if feather_compare_full is not None:
step6_images.append({"label": f"★ 黑底羽化对比·完整(左=无羽化 右=羽化{feather_px}pxα{feather_alpha:.1f})", "b64": _enc(feather_compare_full)})
steps.append({ steps.append({
"title": "⑥ 贴回原图", "title": "⑥ 贴回原图",
"desc": f"把SD结果逆warpAffine贴回用户原图。可调：strict_mask={strict_mask}, blend_dilate={bd}{feather_info}。整框覆盖=整个裁剪框覆盖原图；严格mask=只在mask区域融合，mask外保留原图", "desc": f"把SD结果逆warpAffine贴回用户原图。可调：strict_mask={strict_mask}, blend_dilate={bd}{feather_info}。整框覆盖=整个裁剪框覆盖原图；严格mask=只在mask区域融合，mask外保留原图",
@@ -424,7 +390,7 @@ def hair_swap_manual(origin_img, hand_mask, hair_id, hairstyle_process, landmark
"strict_mask": strict_mask, "seamless_blend": seamless_blend, "strict_mask": strict_mask, "seamless_blend": seamless_blend,
"is_hr": is_hr, "dilate_kernel": list(dilate_kernel), "is_hr": is_hr, "dilate_kernel": list(dilate_kernel),
"denoising_strength": denoising_strength, "denoising_strength": denoising_strength,
"blend_dilate": list(blend_dilate), "feather_px": feather_px, "feather_alpha": feather_alpha, "blend_dilate": list(blend_dilate), "feather_px": feather_px,
"seamless_dilate": list(seamless_dilate), "seamless_dilate": list(seamless_dilate),
"enhance": enhance, "enhance_denoising": enhance_denoising, "enhance": enhance, "enhance_denoising": enhance_denoising,
"mask_area_ratio_pct": round(mask_area_ratio, 1), "mask_area_ratio_pct": round(mask_area_ratio, 1),
project/hair_service_sd/hairstyle_model_infer.py
+10 -5
View File
@@ -72,6 +72,8 @@ class HairStyle_Model_Infer(object):
# haircolor_dir = '/home/data/hair/data/ref_color/3628746832766' # haircolor_dir = '/home/data/hair/data/ref_color/3628746832766'
face_base, hair_matting, status = self.infer_haircolor_new(user_rgb_8uc3_orisize, haircolor_dir,target_hair_color, face_base, hair_matting, status = self.infer_haircolor_new(user_rgb_8uc3_orisize, haircolor_dir,target_hair_color,
return_matting=True) return_matting=True)
if status != 0:
return user_rgb_8uc3_orisize, None, status
user_rgb_8uc3_orisize = face_base user_rgb_8uc3_orisize = face_base
# 构建一个色板 # 构建一个色板
r, g, b = target_hair_color r, g, b = target_hair_color
@@ -1016,7 +1018,7 @@ class HairStyle_Model_Infer(object):
def infer_haircolor_new(self, user_rgb_8uc3_orisize, haircolor_dir, target_hair_color, return_matting=False): def infer_haircolor_new(self, user_rgb_8uc3_orisize, haircolor_dir, target_hair_color, return_matting=False):
landmarks_origin_img_1k= self.get_landmark.forward(user_rgb_8uc3_orisize) landmarks_origin_img_1k= self.get_landmark.forward(user_rgb_8uc3_orisize)
if landmarks_origin_img_1k is None: if landmarks_origin_img_1k is None:
return None, 10001 return None, None, 10001
need_process = (target_hair_color[0] * 0.299 + target_hair_color[1] * 0.587 + target_hair_color[2] * 0.114) > 150 need_process = (target_hair_color[0] * 0.299 + target_hair_color[1] * 0.587 + target_hair_color[2] * 0.114) > 150
_, user_matting_8uc1_bald_orisize = self.process_data_infer.generator_matte.matte_inference(user_rgb_8uc3_orisize, _, user_matting_8uc1_bald_orisize = self.process_data_infer.generator_matte.matte_inference(user_rgb_8uc3_orisize,
landmarks_origin_img_1k) landmarks_origin_img_1k)
@@ -1036,10 +1038,13 @@ class HairStyle_Model_Infer(object):
need_process=True need_process=True
if need_process: if need_process:
haircolor_dir_tmp = os.path.join(config.get('default', "haircolorDir"), config.get('default', "baseColor_ID")) haircolor_dir_tmp = os.path.join(config.get('default', "haircolorDir"), config.get('default', "baseColor_ID"))
face_base, status1 = self.infer_haircolor_tj(user_rgb_8uc3_orisize, haircolor_dir_tmp) if os.path.exists(haircolor_dir_tmp):
# cv2.imwrite('/home/student/Desktop/tmp_color/need/face_base.png', face_base) face_base, status1 = self.infer_haircolor_tj(user_rgb_8uc3_orisize, haircolor_dir_tmp)
if status1 == 0: # cv2.imwrite('/home/student/Desktop/tmp_color/need/face_base.png', face_base)
return face_base, user_matting_8uc1_bald_orisize, 0 if status1 == 0:
return face_base, user_matting_8uc1_bald_orisize, 0
else:
need_process = False
else: else:
need_process = False need_process = False
if not need_process: if not need_process:
project/hair_service_sd/inpaint_mask.py
-64
View File
@@ -1,64 +0,0 @@
# -*- coding: utf-8 -*-
"""mask 区域重绘（inpainting）服务
两种模式（底层都是 webui img2img inpainting，严格只在mask区重绘）：
1. enhance: 用 enhance_hair.webui_img2img（默认提示词=发丝增强，denoising默认0.35）
2. pure_inpaint: 纯SD inpainting（默认提示词=自然头发，denoising默认0.5）
严格只重绘 mask 区域：inpainting_fill=1（mask外保留原图）
"""
import cv2
import numpy as np
from utils import enhance_hair
def inpaint_mask(img, mask, prompt="", denoising_strength=0.5, mode="enhance"):
"""在 mask 区域做 SD inpainting 重绘。
参数:
img: 输入图 BGR
mask: 灰度mask，白色(255)=重绘区，与img同尺寸
prompt: 提示词（空则用模式默认值）
denoising_strength: 重绘强度 0~1
mode: "enhance" 或 "pure_inpaint"
返回:
result: 重绘后的图（mask区已重绘，其余保留原图）
info: 本次参数说明
"""
h, w = img.shape[:2]
# 归一化 mask（保证 0/255，尺寸一致）
if mask.shape[:2] != (h, w):
mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)
if mask.ndim == 3:
mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
_, mask_bin = cv2.threshold(mask, 30, 255, cv2.THRESH_BINARY)
# 模式默认值
if not prompt:
if mode == "enhance":
prompt = "high quality, detailed natural hair strands, realistic hair texture, sharp focus, photorealistic"
else:
prompt = "natural realistic hair, soft hair strands, photorealistic, high quality, detailed"
# 统计 mask 区域占比
mask_ratio = float((mask_bin > 0).sum()) / mask_bin.size * 100
# 调用 enhance_hair.webui_img2img（底层是 inpainting，inpainting_fill=1 严格只在mask区重绘）
result = enhance_hair.webui_img2img(
img, mask_bin, prompt=prompt,
denoising_strength=denoising_strength
)
# webui 可能调整尺寸（调整到8的倍数），resize 回原图尺寸
if result.shape[:2] != (h, w):
result = cv2.resize(result, (w, h), interpolation=cv2.INTER_LANCZOS4)
# 双重保险：强制 mask 外区域用原图（虽然 inpainting_fill=1 已保证，再加一层确保严格）
mask_3c = (mask_bin > 0)[..., None]
result = np.where(mask_3c, result, img).astype(np.uint8)
info = (f"模式={mode}，提示词=\"{prompt[:50]}...\"，denoising={denoising_strength}，"
f"mask区域={mask_ratio:.1f}%。严格只重绘mask区，其余保留原图。")
return result, mask_bin, info
project/hair_service_sd/keypoints/lib/config/__init__.py
+9
View File
@@ -0,0 +1,9 @@
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------
from .default import _C as cfg
from .default import update_config
from .models import MODEL_EXTRAS
project/hair_service_sd/keypoints/lib/config/default.py
+160
View File
@@ -0,0 +1,160 @@
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
from yacs.config import CfgNode as CN
_C = CN()
_C.OUTPUT_DIR = ''
_C.LOG_DIR = ''
_C.DATA_DIR = ''
_C.GPUS = (0,)
_C.WORKERS = 4
_C.PRINT_FREQ = 20
_C.AUTO_RESUME = False
_C.PIN_MEMORY = True
_C.RANK = 0
# Cudnn related params
_C.CUDNN = CN()
_C.CUDNN.BENCHMARK = True
_C.CUDNN.DETERMINISTIC = False
_C.CUDNN.ENABLED = True
# common params for NETWORK
_C.MODEL = CN()
_C.MODEL.NAME = 'pose_hrnet'
_C.MODEL.INIT_WEIGHTS = True
_C.MODEL.PRETRAINED = ''
_C.MODEL.NUM_JOINTS = 17
_C.MODEL.TAG_PER_JOINT = True
_C.MODEL.TARGET_TYPE = 'gaussian'
_C.MODEL.IMAGE_SIZE = [256, 256] # width * height, ex: 192 * 256
_C.MODEL.HEATMAP_SIZE = [64, 64] # width * height, ex: 24 * 32
_C.MODEL.SIGMA = 2
_C.MODEL.EXTRA = CN(new_allowed=True)
_C.LOSS = CN()
_C.LOSS.USE_OHKM = False
_C.LOSS.TOPK = 8
_C.LOSS.USE_TARGET_WEIGHT = True
_C.LOSS.USE_DIFFERENT_JOINTS_WEIGHT = False
# DATASET related params
_C.DATASET = CN()
_C.DATASET.ROOT = ''
_C.DATASET.DATASET = 'mpii'
_C.DATASET.TRAIN_SET = 'train'
_C.DATASET.TEST_SET = 'valid'
_C.DATASET.DATA_FORMAT = 'jpg'
_C.DATASET.HYBRID_JOINTS_TYPE = ''
_C.DATASET.SELECT_DATA = False
# training data augmentation
_C.DATASET.FLIP = True
_C.DATASET.SCALE_FACTOR = 0.25
_C.DATASET.ROT_FACTOR = 30
_C.DATASET.PROB_HALF_BODY = 0.0
_C.DATASET.NUM_JOINTS_HALF_BODY = 8
_C.DATASET.COLOR_RGB = False
# train
_C.TRAIN = CN()
_C.TRAIN.LR_FACTOR = 0.1
_C.TRAIN.LR_STEP = [90, 110]
_C.TRAIN.LR = 0.001
_C.TRAIN.OPTIMIZER = 'adam'
_C.TRAIN.MOMENTUM = 0.9
_C.TRAIN.WD = 0.0001
_C.TRAIN.NESTEROV = False
_C.TRAIN.GAMMA1 = 0.99
_C.TRAIN.GAMMA2 = 0.0
_C.TRAIN.BEGIN_EPOCH = 0
_C.TRAIN.END_EPOCH = 140
_C.TRAIN.RESUME = False
_C.TRAIN.CHECKPOINT = ''
_C.TRAIN.BATCH_SIZE_PER_GPU = 32
_C.TRAIN.SHUFFLE = True
# testing
_C.TEST = CN()
# size of images for each device
_C.TEST.BATCH_SIZE_PER_GPU = 32
# Test Model Epoch
_C.TEST.FLIP_TEST = False
_C.TEST.POST_PROCESS = False
_C.TEST.SHIFT_HEATMAP = False
_C.TEST.USE_GT_BBOX = False
# nms
_C.TEST.IMAGE_THRE = 0.1
_C.TEST.NMS_THRE = 0.6
_C.TEST.SOFT_NMS = False
_C.TEST.OKS_THRE = 0.5
_C.TEST.IN_VIS_THRE = 0.0
_C.TEST.COCO_BBOX_FILE = ''
_C.TEST.BBOX_THRE = 1.0
_C.TEST.MODEL_FILE = ''
# debug
_C.DEBUG = CN()
_C.DEBUG.DEBUG = False
_C.DEBUG.SAVE_BATCH_IMAGES_GT = False
_C.DEBUG.SAVE_BATCH_IMAGES_PRED = False
_C.DEBUG.SAVE_HEATMAPS_GT = False
_C.DEBUG.SAVE_HEATMAPS_PRED = False
def update_config(cfg, args):
cfg.defrost()
cfg.merge_from_file(args.cfg)
cfg.merge_from_list(args.opts)
if args.modelDir:
cfg.OUTPUT_DIR = args.modelDir
if args.logDir:
cfg.LOG_DIR = args.logDir
if args.dataDir:
cfg.DATA_DIR = args.dataDir
cfg.DATASET.ROOT = os.path.join(
cfg.DATA_DIR, cfg.DATASET.ROOT
)
cfg.MODEL.PRETRAINED = os.path.join(
cfg.DATA_DIR, cfg.MODEL.PRETRAINED
)
if cfg.TEST.MODEL_FILE:
cfg.TEST.MODEL_FILE = os.path.join(
cfg.DATA_DIR, cfg.TEST.MODEL_FILE
)
cfg.freeze()
if __name__ == '__main__':
import sys
with open(sys.argv[1], 'w') as f:
print(_C, file=f)
project/hair_service_sd/keypoints/lib/config/models.py
+58
View File
@@ -0,0 +1,58 @@
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from yacs.config import CfgNode as CN
# pose_resnet related params
POSE_RESNET = CN()
POSE_RESNET.NUM_LAYERS = 50
POSE_RESNET.DECONV_WITH_BIAS = False
POSE_RESNET.NUM_DECONV_LAYERS = 3
POSE_RESNET.NUM_DECONV_FILTERS = [256, 256, 256]
POSE_RESNET.NUM_DECONV_KERNELS = [4, 4, 4]
POSE_RESNET.FINAL_CONV_KERNEL = 1
POSE_RESNET.PRETRAINED_LAYERS = ['*']
# pose_multi_resoluton_net related params
POSE_HIGH_RESOLUTION_NET = CN()
POSE_HIGH_RESOLUTION_NET.PRETRAINED_LAYERS = ['*']
POSE_HIGH_RESOLUTION_NET.STEM_INPLANES = 64
POSE_HIGH_RESOLUTION_NET.FINAL_CONV_KERNEL = 1
POSE_HIGH_RESOLUTION_NET.STAGE2 = CN()
POSE_HIGH_RESOLUTION_NET.STAGE2.NUM_MODULES = 1
POSE_HIGH_RESOLUTION_NET.STAGE2.NUM_BRANCHES = 2
POSE_HIGH_RESOLUTION_NET.STAGE2.NUM_BLOCKS = [4, 4]
POSE_HIGH_RESOLUTION_NET.STAGE2.NUM_CHANNELS = [32, 64]
POSE_HIGH_RESOLUTION_NET.STAGE2.BLOCK = 'BASIC'
POSE_HIGH_RESOLUTION_NET.STAGE2.FUSE_METHOD = 'SUM'
POSE_HIGH_RESOLUTION_NET.STAGE3 = CN()
POSE_HIGH_RESOLUTION_NET.STAGE3.NUM_MODULES = 1
POSE_HIGH_RESOLUTION_NET.STAGE3.NUM_BRANCHES = 3
POSE_HIGH_RESOLUTION_NET.STAGE3.NUM_BLOCKS = [4, 4, 4]
POSE_HIGH_RESOLUTION_NET.STAGE3.NUM_CHANNELS = [32, 64, 128]
POSE_HIGH_RESOLUTION_NET.STAGE3.BLOCK = 'BASIC'
POSE_HIGH_RESOLUTION_NET.STAGE3.FUSE_METHOD = 'SUM'
POSE_HIGH_RESOLUTION_NET.STAGE4 = CN()
POSE_HIGH_RESOLUTION_NET.STAGE4.NUM_MODULES = 1
POSE_HIGH_RESOLUTION_NET.STAGE4.NUM_BRANCHES = 4
POSE_HIGH_RESOLUTION_NET.STAGE4.NUM_BLOCKS = [4, 4, 4, 4]
POSE_HIGH_RESOLUTION_NET.STAGE4.NUM_CHANNELS = [32, 64, 128, 256]
POSE_HIGH_RESOLUTION_NET.STAGE4.BLOCK = 'BASIC'
POSE_HIGH_RESOLUTION_NET.STAGE4.FUSE_METHOD = 'SUM'
MODEL_EXTRAS = {
'pose_resnet': POSE_RESNET,
'pose_high_resolution_net': POSE_HIGH_RESOLUTION_NET,
}
project/hair_service_sd/keypoints/lib/core/evaluate.py
+73
View File
@@ -0,0 +1,73 @@
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
from core.inference import get_max_preds
def calc_dists(preds, target, normalize):
preds = preds.astype(np.float32)
target = target.astype(np.float32)
dists = np.zeros((preds.shape[1], preds.shape[0]))
for n in range(preds.shape[0]):
for c in range(preds.shape[1]):
if target[n, c, 0] > 1 and target[n, c, 1] > 1:
normed_preds = preds[n, c, :] / normalize[n]
normed_targets = target[n, c, :] / normalize[n]
dists[c, n] = np.linalg.norm(normed_preds - normed_targets)
else:
dists[c, n] = -1
return dists
def dist_acc(dists, thr=0.5):
''' Return percentage below threshold while ignoring values with a -1 '''
dist_cal = np.not_equal(dists, -1)
num_dist_cal = dist_cal.sum()
if num_dist_cal > 0:
return np.less(dists[dist_cal], thr).sum() * 1.0 / num_dist_cal
else:
return -1
def accuracy(output, target, hm_type='gaussian', thr=0.5):
'''
Calculate accuracy according to PCK,
but uses ground truth heatmap rather than x,y locations
First value to be returned is average accuracy across 'idxs',
followed by individual accuracies
'''
idx = list(range(output.shape[1]))
norm = 1.0
if hm_type == 'gaussian':
pred, _ = get_max_preds(output)
target, _ = get_max_preds(target)
h = output.shape[2]
w = output.shape[3]
norm = np.ones((pred.shape[0], 2)) * np.array([h, w]) / 10
dists = calc_dists(pred, target, norm)
acc = np.zeros((len(idx) + 1))
avg_acc = 0
cnt = 0
for i in range(len(idx)):
acc[i + 1] = dist_acc(dists[idx[i]])
if acc[i + 1] >= 0:
avg_acc = avg_acc + acc[i + 1]
cnt += 1
avg_acc = avg_acc / cnt if cnt != 0 else 0
if cnt != 0:
acc[0] = avg_acc
return acc, avg_acc, cnt, pred
project/hair_service_sd/keypoints/lib/core/function.py
+279
View File
@@ -0,0 +1,279 @@
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import time
import logging
import os
import numpy as np
import torch
from core.evaluate import accuracy
from core.inference import get_final_preds
from utils.transforms import flip_back
from utils.vis import save_debug_images
logger = logging.getLogger(__name__)
def train(config, train_loader, model, criterion, optimizer, epoch,
output_dir, tb_log_dir, writer_dict):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acc = AverageMeter()
# switch to train mode
model.train()
end = time.time()
for i, (input, target, target_weight, meta) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
# compute output
outputs = model(input)
target = target.cuda(non_blocking=True)
target_weight = target_weight.cuda(non_blocking=True)
if isinstance(outputs, list):
loss = criterion(outputs[0], target, target_weight)
for output in outputs[1:]:
loss += criterion(output, target, target_weight)
else:
output = outputs
loss = criterion(output, target, target_weight)
# loss = criterion(output, target, target_weight)
# compute gradient and do update step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure accuracy and record loss
losses.update(loss.item(), input.size(0))
_, avg_acc, cnt, pred = accuracy(output.detach().cpu().numpy(),
target.detach().cpu().numpy())
acc.update(avg_acc, cnt)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % config.PRINT_FREQ == 0:
msg = 'Epoch: [{0}][{1}/{2}]\t' \
'Time {batch_time.val:.3f}s ({batch_time.avg:.3f}s)\t' \
'Speed {speed:.1f} samples/s\t' \
'Data {data_time.val:.3f}s ({data_time.avg:.3f}s)\t' \
'Loss {loss.val:.5f} ({loss.avg:.5f})\t' \
'Accuracy {acc.val:.3f} ({acc.avg:.3f})'.format(
epoch, i, len(train_loader), batch_time=batch_time,
speed=input.size(0)/batch_time.val,
data_time=data_time, loss=losses, acc=acc)
logger.info(msg)
writer = writer_dict['writer']
global_steps = writer_dict['train_global_steps']
writer.add_scalar('train_loss', losses.val, global_steps)
writer.add_scalar('train_acc', acc.val, global_steps)
writer_dict['train_global_steps'] = global_steps + 1
prefix = '{}_{}'.format(os.path.join(output_dir, 'train'), i)
save_debug_images(config, input, meta, target, pred*4, output,
prefix)
def validate(config, val_loader, val_dataset, model, criterion, output_dir,
tb_log_dir, writer_dict=None):
batch_time = AverageMeter()
losses = AverageMeter()
acc = AverageMeter()
# switch to evaluate mode
model.eval()
num_samples = len(val_dataset)
all_preds = np.zeros(
(num_samples, config.MODEL.NUM_JOINTS, 3),
dtype=np.float32
)
all_boxes = np.zeros((num_samples, 6))
image_path = []
filenames = []
imgnums = []
idx = 0
with torch.no_grad():
end = time.time()
for i, (input, target, target_weight, meta) in enumerate(val_loader):
# compute output
outputs = model(input)
if isinstance(outputs, list):
output = outputs[-1]
else:
output = outputs
if config.TEST.FLIP_TEST:
# this part is ugly, because pytorch has not supported negative index
# input_flipped = model(input[:, :, :, ::-1])
input_flipped = np.flip(input.cpu().numpy(), 3).copy()
input_flipped = torch.from_numpy(input_flipped).cuda()
outputs_flipped = model(input_flipped)
if isinstance(outputs_flipped, list):
output_flipped = outputs_flipped[-1]
else:
output_flipped = outputs_flipped
output_flipped = flip_back(output_flipped.cpu().numpy(),
val_dataset.flip_pairs)
output_flipped = torch.from_numpy(output_flipped.copy()).cuda()
# feature is not aligned, shift flipped heatmap for higher accuracy
if config.TEST.SHIFT_HEATMAP:
output_flipped[:, :, :, 1:] = \
output_flipped.clone()[:, :, :, 0:-1]
output = (output + output_flipped) * 0.5
target = target.cuda(non_blocking=True)
target_weight = target_weight.cuda(non_blocking=True)
loss = criterion(output, target, target_weight)
num_images = input.size(0)
# measure accuracy and record loss
losses.update(loss.item(), num_images)
_, avg_acc, cnt, pred = accuracy(output.cpu().numpy(),
target.cpu().numpy())
acc.update(avg_acc, cnt)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
c = meta['center'].numpy()
s = meta['scale'].numpy()
score = meta['score'].numpy()
preds, maxvals = get_final_preds(
config, output.clone().cpu().numpy(), c, s)
all_preds[idx:idx + num_images, :, 0:2] = preds[:, :, 0:2]
all_preds[idx:idx + num_images, :, 2:3] = maxvals
# double check this all_boxes parts
all_boxes[idx:idx + num_images, 0:2] = c[:, 0:2]
all_boxes[idx:idx + num_images, 2:4] = s[:, 0:2]
all_boxes[idx:idx + num_images, 4] = np.prod(s*200, 1)
all_boxes[idx:idx + num_images, 5] = score
image_path.extend(meta['image'])
idx += num_images
if i % config.PRINT_FREQ == 0:
msg = 'Test: [{0}/{1}]\t' \
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' \
'Loss {loss.val:.4f} ({loss.avg:.4f})\t' \
'Accuracy {acc.val:.3f} ({acc.avg:.3f})'.format(
i, len(val_loader), batch_time=batch_time,
loss=losses, acc=acc)
logger.info(msg)
prefix = '{}_{}'.format(
os.path.join(output_dir, 'val'), i
)
save_debug_images(config, input, meta, target, pred*4, output,
prefix)
name_values, perf_indicator = val_dataset.evaluate(
config, all_preds, output_dir, all_boxes, image_path,
filenames, imgnums
)
model_name = config.MODEL.NAME
if isinstance(name_values, list):
for name_value in name_values:
_print_name_value(name_value, model_name)
else:
_print_name_value(name_values, model_name)
if writer_dict:
writer = writer_dict['writer']
global_steps = writer_dict['valid_global_steps']
writer.add_scalar(
'valid_loss',
losses.avg,
global_steps
)
writer.add_scalar(
'valid_acc',
acc.avg,
global_steps
)
if isinstance(name_values, list):
for name_value in name_values:
writer.add_scalars(
'valid',
dict(name_value),
global_steps
)
else:
writer.add_scalars(
'valid',
dict(name_values),
global_steps
)
writer_dict['valid_global_steps'] = global_steps + 1
return perf_indicator
# markdown format output
def _print_name_value(name_value, full_arch_name):
names = name_value.keys()
values = name_value.values()
num_values = len(name_value)
logger.info(
'| Arch ' +
' '.join(['| {}'.format(name) for name in names]) +
' |'
)
logger.info('|---' * (num_values+1) + '|')
if len(full_arch_name) > 15:
full_arch_name = full_arch_name[:8] + '...'
logger.info(
'| ' + full_arch_name + ' ' +
' '.join(['| {:.3f}'.format(value) for value in values]) +
' |'
)
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count if self.count != 0 else 0
project/hair_service_sd/keypoints/lib/core/inference.py
+79
View File
@@ -0,0 +1,79 @@
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import math
import numpy as np
from keypoints.lib.utils.transforms import transform_preds
def get_max_preds(batch_heatmaps):
'''
get predictions from score maps
heatmaps: numpy.ndarray([batch_size, num_joints, height, width])
'''
assert isinstance(batch_heatmaps, np.ndarray), \
'batch_heatmaps should be numpy.ndarray'
assert batch_heatmaps.ndim == 4, 'batch_images should be 4-ndim'
batch_size = batch_heatmaps.shape[0]
num_joints = batch_heatmaps.shape[1]
width = batch_heatmaps.shape[3]
heatmaps_reshaped = batch_heatmaps.reshape((batch_size, num_joints, -1))
idx = np.argmax(heatmaps_reshaped, 2)
maxvals = np.amax(heatmaps_reshaped, 2)
maxvals = maxvals.reshape((batch_size, num_joints, 1))
idx = idx.reshape((batch_size, num_joints, 1))
preds = np.tile(idx, (1, 1, 2)).astype(np.float32)
preds[:, :, 0] = (preds[:, :, 0]) % width
preds[:, :, 1] = np.floor((preds[:, :, 1]) / width)
pred_mask = np.tile(np.greater(maxvals, 0.0), (1, 1, 2))
pred_mask = pred_mask.astype(np.float32)
preds *= pred_mask
return preds, maxvals
def get_final_preds(config, batch_heatmaps, center, scale):
coords, maxvals = get_max_preds(batch_heatmaps)
heatmap_height = batch_heatmaps.shape[2]
heatmap_width = batch_heatmaps.shape[3]
# post-processing
if config.TEST.POST_PROCESS:
for n in range(coords.shape[0]):
for p in range(coords.shape[1]):
hm = batch_heatmaps[n][p]
px = int(math.floor(coords[n][p][0] + 0.5))
py = int(math.floor(coords[n][p][1] + 0.5))
if 1 < px < heatmap_width-1 and 1 < py < heatmap_height-1:
diff = np.array(
[
hm[py][px+1] - hm[py][px-1],
hm[py+1][px]-hm[py-1][px]
]
)
coords[n][p] += np.sign(diff) * .25
preds = coords.copy()
# Transform back
for i in range(coords.shape[0]):
preds[i] = transform_preds(
coords[i], center[i], scale[i], [heatmap_width, heatmap_height]
)
return preds, maxvals
project/hair_service_sd/keypoints/lib/core/loss.py
+84
View File
@@ -0,0 +1,84 @@
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import torch
import torch.nn as nn
class JointsMSELoss(nn.Module):
def __init__(self, use_target_weight):
super(JointsMSELoss, self).__init__()
self.criterion = nn.MSELoss(reduction='mean')
self.use_target_weight = use_target_weight
def forward(self, output, target, target_weight):
batch_size = output.size(0)
num_joints = output.size(1)
heatmaps_pred = output.reshape((batch_size, num_joints, -1)).split(1, 1)
heatmaps_gt = target.reshape((batch_size, num_joints, -1)).split(1, 1)
loss = 0
for idx in range(num_joints):
heatmap_pred = heatmaps_pred[idx].squeeze()
heatmap_gt = heatmaps_gt[idx].squeeze()
if self.use_target_weight:
loss += 0.5 * self.criterion(
heatmap_pred.mul(target_weight[:, idx]),
heatmap_gt.mul(target_weight[:, idx])
)
else:
loss += 0.5 * self.criterion(heatmap_pred, heatmap_gt)
return loss / num_joints
class JointsOHKMMSELoss(nn.Module):
def __init__(self, use_target_weight, topk=8):
super(JointsOHKMMSELoss, self).__init__()
self.criterion = nn.MSELoss(reduction='none')
self.use_target_weight = use_target_weight
self.topk = topk
def ohkm(self, loss):
ohkm_loss = 0.
for i in range(loss.size()[0]):
sub_loss = loss[i]
topk_val, topk_idx = torch.topk(
sub_loss, k=self.topk, dim=0, sorted=False
)
tmp_loss = torch.gather(sub_loss, 0, topk_idx)
ohkm_loss += torch.sum(tmp_loss) / self.topk
ohkm_loss /= loss.size()[0]
return ohkm_loss
def forward(self, output, target, target_weight):
batch_size = output.size(0)
num_joints = output.size(1)
heatmaps_pred = output.reshape((batch_size, num_joints, -1)).split(1, 1)
heatmaps_gt = target.reshape((batch_size, num_joints, -1)).split(1, 1)
loss = []
for idx in range(num_joints):
heatmap_pred = heatmaps_pred[idx].squeeze()
heatmap_gt = heatmaps_gt[idx].squeeze()
if self.use_target_weight:
loss.append(0.5 * self.criterion(
heatmap_pred.mul(target_weight[:, idx]),
heatmap_gt.mul(target_weight[:, idx])
))
else:
loss.append(
0.5 * self.criterion(heatmap_pred, heatmap_gt)
)
loss = [l.mean(dim=1).unsqueeze(dim=1) for l in loss]
loss = torch.cat(loss, dim=1)
return self.ohkm(loss)
project/hair_service_sd/keypoints/lib/dataset/JointsDataset.py
+292
View File
@@ -0,0 +1,292 @@
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import copy
import logging
import random
import cv2
import numpy as np
import torch
from torch.utils.data import Dataset
from utils.transforms import get_affine_transform
from utils.transforms import affine_transform
from utils.transforms import fliplr_joints
logger = logging.getLogger(__name__)
class JointsDataset(Dataset):
def __init__(self, cfg, root, image_set, is_train, transform=None):
self.num_joints = 0
self.pixel_std = 200
self.flip_pairs = []
self.parent_ids = []
self.is_train = is_train
self.root = root
self.image_set = image_set
self.output_path = cfg.OUTPUT_DIR
self.data_format = cfg.DATASET.DATA_FORMAT
self.scale_factor = cfg.DATASET.SCALE_FACTOR
self.rotation_factor = cfg.DATASET.ROT_FACTOR
self.flip = cfg.DATASET.FLIP
self.num_joints_half_body = cfg.DATASET.NUM_JOINTS_HALF_BODY
self.prob_half_body = cfg.DATASET.PROB_HALF_BODY
self.color_rgb = cfg.DATASET.COLOR_RGB
self.target_type = cfg.MODEL.TARGET_TYPE
self.image_size = np.array(cfg.MODEL.IMAGE_SIZE)
self.heatmap_size = np.array(cfg.MODEL.HEATMAP_SIZE)
self.sigma = cfg.MODEL.SIGMA
self.use_different_joints_weight = cfg.LOSS.USE_DIFFERENT_JOINTS_WEIGHT
self.joints_weight = 1
self.transform = transform
self.db = []
def _get_db(self):
raise NotImplementedError
def evaluate(self, cfg, preds, output_dir, *args, **kwargs):
raise NotImplementedError
def half_body_transform(self, joints, joints_vis):
upper_joints = []
lower_joints = []
for joint_id in range(self.num_joints):
if joints_vis[joint_id][0] > 0:
if joint_id in self.upper_body_ids:
upper_joints.append(joints[joint_id])
else:
lower_joints.append(joints[joint_id])
if np.random.randn() < 0.5 and len(upper_joints) > 2:
selected_joints = upper_joints
else:
selected_joints = lower_joints \
if len(lower_joints) > 2 else upper_joints
if len(selected_joints) < 2:
return None, None
selected_joints = np.array(selected_joints, dtype=np.float32)
center = selected_joints.mean(axis=0)[:2]
left_top = np.amin(selected_joints, axis=0)
right_bottom = np.amax(selected_joints, axis=0)
w = right_bottom[0] - left_top[0]
h = right_bottom[1] - left_top[1]
if w > self.aspect_ratio * h:
h = w * 1.0 / self.aspect_ratio
elif w < self.aspect_ratio * h:
w = h * self.aspect_ratio
scale = np.array(
[
w * 1.0 / self.pixel_std,
h * 1.0 / self.pixel_std
],
dtype=np.float32
)
scale = scale * 1.5
return center, scale
def __len__(self,):
return len(self.db)
def __getitem__(self, idx):
db_rec = copy.deepcopy(self.db[idx])
image_file = db_rec['image']
filename = db_rec['filename'] if 'filename' in db_rec else ''
imgnum = db_rec['imgnum'] if 'imgnum' in db_rec else ''
if self.data_format == 'zip':
from utils import zipreader
data_numpy = zipreader.imread(
image_file, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION
)
else:
data_numpy = cv2.imread(
image_file, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION
)
if self.color_rgb:
data_numpy = cv2.cvtColor(data_numpy, cv2.COLOR_BGR2RGB)
if data_numpy is None:
logger.error('=> fail to read {}'.format(image_file))
raise ValueError('Fail to read {}'.format(image_file))
joints = db_rec['joints_3d']
joints_vis = db_rec['joints_3d_vis']
c = db_rec['center']
s = db_rec['scale']
score = db_rec['score'] if 'score' in db_rec else 1
r = 0
if self.is_train:
if (np.sum(joints_vis[:, 0]) > self.num_joints_half_body
and np.random.rand() < self.prob_half_body):
c_half_body, s_half_body = self.half_body_transform(
joints, joints_vis
)
if c_half_body is not None and s_half_body is not None:
c, s = c_half_body, s_half_body
sf = self.scale_factor
rf = self.rotation_factor
s = s * np.clip(np.random.randn()*sf + 1, 1 - sf, 1 + sf)
r = np.clip(np.random.randn()*rf, -rf*2, rf*2) \
if random.random() <= 0.6 else 0
if self.flip and random.random() <= 0.5:
data_numpy = data_numpy[:, ::-1, :]
joints, joints_vis = fliplr_joints(
joints, joints_vis, data_numpy.shape[1], self.flip_pairs)
c[0] = data_numpy.shape[1] - c[0] - 1
trans = get_affine_transform(c, s, r, self.image_size)
input = cv2.warpAffine(
data_numpy,
trans,
(int(self.image_size[0]), int(self.image_size[1])),
flags=cv2.INTER_LINEAR)
cv2.imshow('input', input)
cv2.waitKey()
if self.transform:
input = self.transform(input)
for i in range(self.num_joints):
if joints_vis[i, 0] > 0.0:
joints[i, 0:2] = affine_transform(joints[i, 0:2], trans)
target, target_weight = self.generate_target(joints, joints_vis)
target = torch.from_numpy(target)
target_weight = torch.from_numpy(target_weight)
meta = {
'image': image_file,
'filename': filename,
'imgnum': imgnum,
'joints': joints,
'joints_vis': joints_vis,
'center': c,
'scale': s,
'rotation': r,
'score': score
}
return input, target, target_weight, meta
def select_data(self, db):
db_selected = []
for rec in db:
num_vis = 0
joints_x = 0.0
joints_y = 0.0
for joint, joint_vis in zip(
rec['joints_3d'], rec['joints_3d_vis']):
if joint_vis[0] <= 0:
continue
num_vis += 1
joints_x += joint[0]
joints_y += joint[1]
if num_vis == 0:
continue
joints_x, joints_y = joints_x / num_vis, joints_y / num_vis
area = rec['scale'][0] * rec['scale'][1] * (self.pixel_std**2)
joints_center = np.array([joints_x, joints_y])
bbox_center = np.array(rec['center'])
diff_norm2 = np.linalg.norm((joints_center-bbox_center), 2)
ks = np.exp(-1.0*(diff_norm2**2) / ((0.2)**2*2.0*area))
metric = (0.2 / 16) * num_vis + 0.45 - 0.2 / 16
if ks > metric:
db_selected.append(rec)
logger.info('=> num db: {}'.format(len(db)))
logger.info('=> num selected db: {}'.format(len(db_selected)))
return db_selected
def generate_target(self, joints, joints_vis):
'''
:param joints: [num_joints, 3]
:param joints_vis: [num_joints, 3]
:return: target, target_weight(1: visible, 0: invisible)
'''
target_weight = np.ones((self.num_joints, 1), dtype=np.float32)
target_weight[:, 0] = joints_vis[:, 0]
assert self.target_type == 'gaussian', \
'Only support gaussian map now!'
if self.target_type == 'gaussian':
target = np.zeros((self.num_joints,
self.heatmap_size[1],
self.heatmap_size[0]),
dtype=np.float32)
tmp_size = self.sigma * 3
for joint_id in range(self.num_joints):
feat_stride = self.image_size / self.heatmap_size
mu_x = int(joints[joint_id][0] / feat_stride[0] + 0.5)
mu_y = int(joints[joint_id][1] / feat_stride[1] + 0.5)
# Check that any part of the gaussian is in-bounds
ul = [int(mu_x - tmp_size), int(mu_y - tmp_size)]
br = [int(mu_x + tmp_size + 1), int(mu_y + tmp_size + 1)]
if ul[0] >= self.heatmap_size[0] or ul[1] >= self.heatmap_size[1] \
or br[0] < 0 or br[1] < 0:
# If not, just return the image as is
target_weight[joint_id] = 0
continue
# # Generate gaussian
size = 2 * tmp_size + 1
x = np.arange(0, size, 1, np.float32)
y = x[:, np.newaxis]
x0 = y0 = size // 2
# The gaussian is not normalized, we want the center value to equal 1
g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * self.sigma ** 2))
# Usable gaussian range
g_x = max(0, -ul[0]), min(br[0], self.heatmap_size[0]) - ul[0]
g_y = max(0, -ul[1]), min(br[1], self.heatmap_size[1]) - ul[1]
# Image range
img_x = max(0, ul[0]), min(br[0], self.heatmap_size[0])
img_y = max(0, ul[1]), min(br[1], self.heatmap_size[1])
v = target_weight[joint_id]
if v > 0.5:
target[joint_id][img_y[0]:img_y[1], img_x[0]:img_x[1]] = \
g[g_y[0]:g_y[1], g_x[0]:g_x[1]]
if self.use_different_joints_weight:
target_weight = np.multiply(target_weight, self.joints_weight)
return target, target_weight
project/hair_service_sd/keypoints/lib/dataset/__init__.py
+12
View File
@@ -0,0 +1,12 @@
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from .mpii import MPIIDataset as mpii
from .coco import COCODataset as coco
project/hair_service_sd/keypoints/lib/dataset/coco.py
+445
View File
@@ -0,0 +1,445 @@
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from collections import defaultdict
from collections import OrderedDict
import logging
import os
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
import json_tricks as json
import numpy as np
from dataset.JointsDataset import JointsDataset
from nms.nms import oks_nms
from nms.nms import soft_oks_nms
logger = logging.getLogger(__name__)
class COCODataset(JointsDataset):
'''
"keypoints": {
0: "nose",
1: "left_eye",
2: "right_eye",
3: "left_ear",
4: "right_ear",
5: "left_shoulder",
6: "right_shoulder",
7: "left_elbow",
8: "right_elbow",
9: "left_wrist",
10: "right_wrist",
11: "left_hip",
12: "right_hip",
13: "left_knee",
14: "right_knee",
15: "left_ankle",
16: "right_ankle"
},
"skeleton": [
[16,14],[14,12],[17,15],[15,13],[12,13],[6,12],[7,13], [6,7],[6,8],
[7,9],[8,10],[9,11],[2,3],[1,2],[1,3],[2,4],[3,5],[4,6],[5,7]]
'''
def __init__(self, cfg, root, image_set, is_train, transform=None):
super().__init__(cfg, root, image_set, is_train, transform)
self.nms_thre = cfg.TEST.NMS_THRE
self.image_thre = cfg.TEST.IMAGE_THRE
self.soft_nms = cfg.TEST.SOFT_NMS
self.oks_thre = cfg.TEST.OKS_THRE
self.in_vis_thre = cfg.TEST.IN_VIS_THRE
self.bbox_file = cfg.TEST.COCO_BBOX_FILE
self.use_gt_bbox = cfg.TEST.USE_GT_BBOX
self.image_width = cfg.MODEL.IMAGE_SIZE[0]
self.image_height = cfg.MODEL.IMAGE_SIZE[1]
self.aspect_ratio = self.image_width * 1.0 / self.image_height
self.pixel_std = 200
self.coco = COCO(self._get_ann_file_keypoint())
# deal with class names
cats = [cat['name']
for cat in self.coco.loadCats(self.coco.getCatIds())]
self.classes = ['__background__'] + cats
logger.info('=> classes: {}'.format(self.classes))
self.num_classes = len(self.classes)
self._class_to_ind = dict(zip(self.classes, range(self.num_classes)))
self._class_to_coco_ind = dict(zip(cats, self.coco.getCatIds()))
self._coco_ind_to_class_ind = dict(
[
(self._class_to_coco_ind[cls], self._class_to_ind[cls])
for cls in self.classes[1:]
]
)
# load image file names
self.image_set_index = self._load_image_set_index()
self.num_images = len(self.image_set_index)
logger.info('=> num_images: {}'.format(self.num_images))
self.num_joints = 17
self.flip_pairs = [[1, 2], [3, 4], [5, 6], [7, 8],
[9, 10], [11, 12], [13, 14], [15, 16]]
self.parent_ids = None
self.upper_body_ids = (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10)
self.lower_body_ids = (11, 12, 13, 14, 15, 16)
self.joints_weight = np.array(
[
1., 1., 1., 1., 1., 1., 1., 1.2, 1.2,
1.5, 1.5, 1., 1., 1.2, 1.2, 1.5, 1.5
],
dtype=np.float32
).reshape((self.num_joints, 1))
self.db = self._get_db()
if is_train and cfg.DATASET.SELECT_DATA:
self.db = self.select_data(self.db)
logger.info('=> load {} samples'.format(len(self.db)))
def _get_ann_file_keypoint(self):
""" self.root / annotations / person_keypoints_train2017.json """
prefix = 'person_keypoints' \
if 'test' not in self.image_set else 'image_info'
return os.path.join(
self.root,
'annotations',
prefix + '_' + self.image_set + '.json'
)
def _load_image_set_index(self):
""" image id: int """
image_ids = self.coco.getImgIds()
return image_ids
def _get_db(self):
if self.is_train or self.use_gt_bbox:
# use ground truth bbox
gt_db = self._load_coco_keypoint_annotations()
else:
# use bbox from detection
gt_db = self._load_coco_person_detection_results()
return gt_db
def _load_coco_keypoint_annotations(self):
""" ground truth bbox and keypoints """
gt_db = []
for index in self.image_set_index:
gt_db.extend(self._load_coco_keypoint_annotation_kernal(index))
return gt_db
def _load_coco_keypoint_annotation_kernal(self, index):
"""
coco ann: [u'segmentation', u'area', u'iscrowd', u'image_id', u'bbox', u'category_id', u'id']
iscrowd:
crowd instances are handled by marking their overlaps with all categories to -1
and later excluded in training
bbox:
[x1, y1, w, h]
:param index: coco image id
:return: db entry
"""
im_ann = self.coco.loadImgs(index)[0]
width = im_ann['width']
height = im_ann['height']
annIds = self.coco.getAnnIds(imgIds=index, iscrowd=False)
objs = self.coco.loadAnns(annIds)
# sanitize bboxes
valid_objs = []
for obj in objs:
x, y, w, h = obj['bbox']
x1 = np.max((0, x))
y1 = np.max((0, y))
x2 = np.min((width - 1, x1 + np.max((0, w - 1))))
y2 = np.min((height - 1, y1 + np.max((0, h - 1))))
if obj['area'] > 0 and x2 >= x1 and y2 >= y1:
obj['clean_bbox'] = [x1, y1, x2-x1, y2-y1]
valid_objs.append(obj)
objs = valid_objs
rec = []
for obj in objs:
cls = self._coco_ind_to_class_ind[obj['category_id']]
if cls != 1:
continue
# ignore objs without keypoints annotation
if max(obj['keypoints']) == 0:
continue
joints_3d = np.zeros((self.num_joints, 3), dtype=np.float)
joints_3d_vis = np.zeros((self.num_joints, 3), dtype=np.float)
for ipt in range(self.num_joints):
joints_3d[ipt, 0] = obj['keypoints'][ipt * 3 + 0]
joints_3d[ipt, 1] = obj['keypoints'][ipt * 3 + 1]
joints_3d[ipt, 2] = 0
t_vis = obj['keypoints'][ipt * 3 + 2]
if t_vis > 1:
t_vis = 1
joints_3d_vis[ipt, 0] = t_vis
joints_3d_vis[ipt, 1] = t_vis
joints_3d_vis[ipt, 2] = 0
center, scale = self._box2cs(obj['clean_bbox'][:4])
rec.append({
'image': self.image_path_from_index(index),
'center': center,
'scale': scale,
'joints_3d': joints_3d,
'joints_3d_vis': joints_3d_vis,
'filename': '',
'imgnum': 0,
})
return rec
def _box2cs(self, box):
x, y, w, h = box[:4]
return self._xywh2cs(x, y, w, h)
def _xywh2cs(self, x, y, w, h):
center = np.zeros((2), dtype=np.float32)
center[0] = x + w * 0.5
center[1] = y + h * 0.5
if w > self.aspect_ratio * h:
h = w * 1.0 / self.aspect_ratio
elif w < self.aspect_ratio * h:
w = h * self.aspect_ratio
scale = np.array(
[w * 1.0 / self.pixel_std, h * 1.0 / self.pixel_std],
dtype=np.float32)
if center[0] != -1:
scale = scale * 1.25
return center, scale
def image_path_from_index(self, index):
""" example: images / train2017 / 000000119993.jpg """
file_name = '%012d.jpg' % index
if '2014' in self.image_set:
file_name = 'COCO_%s_' % self.image_set + file_name
prefix = 'test2017' if 'test' in self.image_set else self.image_set
data_name = prefix + '.zip@' if self.data_format == 'zip' else prefix
image_path = os.path.join(
self.root, 'images', data_name, file_name)
return image_path
def _load_coco_person_detection_results(self):
all_boxes = None
with open(self.bbox_file, 'r') as f:
all_boxes = json.load(f)[:10]
if not all_boxes:
logger.error('=> Load %s fail!' % self.bbox_file)
return None
logger.info('=> Total boxes: {}'.format(len(all_boxes)))
kpt_db = []
num_boxes = 0
for n_img in range(0, len(all_boxes)):
det_res = all_boxes[n_img]
if det_res['category_id'] != 1:
continue
img_name = self.image_path_from_index(det_res['image_id'])
box = det_res['bbox']
score = det_res['score']
if score < self.image_thre:
continue
num_boxes = num_boxes + 1
center, scale = self._box2cs(box)
joints_3d = np.zeros((self.num_joints, 3), dtype=np.float)
joints_3d_vis = np.ones(
(self.num_joints, 3), dtype=np.float)
kpt_db.append({
'image': img_name,
'center': center,
'scale': scale,
'score': score,
'joints_3d': joints_3d,
'joints_3d_vis': joints_3d_vis,
})
logger.info('=> Total boxes after fliter low score@{}: {}'.format(
self.image_thre, num_boxes))
return kpt_db
def evaluate(self, cfg, preds, output_dir, all_boxes, img_path,
*args, **kwargs):
rank = cfg.RANK
res_folder = os.path.join(output_dir, 'results')
if not os.path.exists(res_folder):
try:
os.makedirs(res_folder)
except Exception:
logger.error('Fail to make {}'.format(res_folder))
res_file = os.path.join(
res_folder, 'keypoints_{}_results_{}.json'.format(
self.image_set, rank)
)
# person x (keypoints)
_kpts = []
for idx, kpt in enumerate(preds):
_kpts.append({
'keypoints': kpt,
'center': all_boxes[idx][0:2],
'scale': all_boxes[idx][2:4],
'area': all_boxes[idx][4],
'score': all_boxes[idx][5],
'image': int(img_path[idx][-16:-4])
})
# image x person x (keypoints)
kpts = defaultdict(list)
for kpt in _kpts:
kpts[kpt['image']].append(kpt)
# rescoring and oks nms
num_joints = self.num_joints
in_vis_thre = self.in_vis_thre
oks_thre = self.oks_thre
oks_nmsed_kpts = []
for img in kpts.keys():
img_kpts = kpts[img]
for n_p in img_kpts:
box_score = n_p['score']
kpt_score = 0
valid_num = 0
for n_jt in range(0, num_joints):
t_s = n_p['keypoints'][n_jt][2]
if t_s > in_vis_thre:
kpt_score = kpt_score + t_s
valid_num = valid_num + 1
if valid_num != 0:
kpt_score = kpt_score / valid_num
# rescoring
n_p['score'] = kpt_score * box_score
if self.soft_nms:
keep = soft_oks_nms(
[img_kpts[i] for i in range(len(img_kpts))],
oks_thre
)
else:
keep = oks_nms(
[img_kpts[i] for i in range(len(img_kpts))],
oks_thre
)
if len(keep) == 0:
oks_nmsed_kpts.append(img_kpts)
else:
oks_nmsed_kpts.append([img_kpts[_keep] for _keep in keep])
self._write_coco_keypoint_results(
oks_nmsed_kpts, res_file)
if 'test' not in self.image_set:
info_str = self._do_python_keypoint_eval(
res_file, res_folder)
name_value = OrderedDict(info_str)
return name_value, name_value['AP']
else:
return {'Null': 0}, 0
def _write_coco_keypoint_results(self, keypoints, res_file):
data_pack = [
{
'cat_id': self._class_to_coco_ind[cls],
'cls_ind': cls_ind,
'cls': cls,
'ann_type': 'keypoints',
'keypoints': keypoints
}
for cls_ind, cls in enumerate(self.classes) if not cls == '__background__'
]
results = self._coco_keypoint_results_one_category_kernel(data_pack[0])
logger.info('=> writing results json to %s' % res_file)
with open(res_file, 'w') as f:
json.dump(results, f, sort_keys=True, indent=4)
try:
json.load(open(res_file))
except Exception:
content = []
with open(res_file, 'r') as f:
for line in f:
content.append(line)
content[-1] = ']'
with open(res_file, 'w') as f:
for c in content:
f.write(c)
def _coco_keypoint_results_one_category_kernel(self, data_pack):
cat_id = data_pack['cat_id']
keypoints = data_pack['keypoints']
cat_results = []
for img_kpts in keypoints:
if len(img_kpts) == 0:
continue
_key_points = np.array([img_kpts[k]['keypoints']
for k in range(len(img_kpts))])
key_points = np.zeros(
(_key_points.shape[0], self.num_joints * 3), dtype=np.float
)
for ipt in range(self.num_joints):
key_points[:, ipt * 3 + 0] = _key_points[:, ipt, 0]
key_points[:, ipt * 3 + 1] = _key_points[:, ipt, 1]
key_points[:, ipt * 3 + 2] = _key_points[:, ipt, 2] # keypoints score.
result = [
{
'image_id': img_kpts[k]['image'],
'category_id': cat_id,
'keypoints': list(key_points[k]),
'score': img_kpts[k]['score'],
'center': list(img_kpts[k]['center']),
'scale': list(img_kpts[k]['scale'])
}
for k in range(len(img_kpts))
]
cat_results.extend(result)
return cat_results
def _do_python_keypoint_eval(self, res_file, res_folder):
coco_dt = self.coco.loadRes(res_file)
coco_eval = COCOeval(self.coco, coco_dt, 'keypoints')
coco_eval.params.useSegm = None
coco_eval.evaluate()
coco_eval.accumulate()
coco_eval.summarize()
stats_names = ['AP', 'Ap .5', 'AP .75', 'AP (M)', 'AP (L)', 'AR', 'AR .5', 'AR .75', 'AR (M)', 'AR (L)']
info_str = []
for ind, name in enumerate(stats_names):
info_str.append((name, coco_eval.stats[ind]))
return info_str
project/hair_service_sd/keypoints/lib/dataset/mpii.py
+181
View File
@@ -0,0 +1,181 @@
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import logging
import os
import json_tricks as json
from collections import OrderedDict
import numpy as np
from scipy.io import loadmat, savemat
from dataset.JointsDataset import JointsDataset
logger = logging.getLogger(__name__)
class MPIIDataset(JointsDataset):
def __init__(self, cfg, root, image_set, is_train, transform=None):
super().__init__(cfg, root, image_set, is_train, transform)
self.num_joints = 16
self.flip_pairs = [[0, 5], [1, 4], [2, 3], [10, 15], [11, 14], [12, 13]]
self.parent_ids = [1, 2, 6, 6, 3, 4, 6, 6, 7, 8, 11, 12, 7, 7, 13, 14]
self.upper_body_ids = (7, 8, 9, 10, 11, 12, 13, 14, 15)
self.lower_body_ids = (0, 1, 2, 3, 4, 5, 6)
self.db = self._get_db()
if is_train and cfg.DATASET.SELECT_DATA:
self.db = self.select_data(self.db)
logger.info('=> load {} samples'.format(len(self.db)))
def _get_db(self):
# create train/val split
file_name = os.path.join(
self.root, 'annot', self.image_set+'.json'
)
with open(file_name) as anno_file:
anno = json.load(anno_file)
gt_db = []
for a in anno:
image_name = a['image']
c = np.array(a['center'], dtype=np.float)
s = np.array([a['scale'], a['scale']], dtype=np.float)
# Adjust center/scale slightly to avoid cropping limbs
if c[0] != -1:
c[1] = c[1] + 15 * s[1]
s = s * 1.25
# MPII uses matlab format, index is based 1,
# we should first convert to 0-based index
c = c - 1
joints_3d = np.zeros((self.num_joints, 3), dtype=np.float)
joints_3d_vis = np.zeros((self.num_joints, 3), dtype=np.float)
if self.image_set != 'test':
joints = np.array(a['joints'])
joints[:, 0:2] = joints[:, 0:2] - 1
joints_vis = np.array(a['joints_vis'])
assert len(joints) == self.num_joints, \
'joint num diff: {} vs {}'.format(len(joints),
self.num_joints)
joints_3d[:, 0:2] = joints[:, 0:2]
joints_3d_vis[:, 0] = joints_vis[:]
joints_3d_vis[:, 1] = joints_vis[:]
image_dir = 'images.zip@' if self.data_format == 'zip' else 'images'
gt_db.append(
{
'image': os.path.join(self.root, image_dir, image_name),
'center': c,
'scale': s,
'joints_3d': joints_3d,
'joints_3d_vis': joints_3d_vis,
'filename': '',
'imgnum': 0,
}
)
return gt_db
def evaluate(self, cfg, preds, output_dir, *args, **kwargs):
# convert 0-based index to 1-based index
preds = preds[:, :, 0:2] + 1.0
if output_dir:
pred_file = os.path.join(output_dir, 'pred.mat')
savemat(pred_file, mdict={'preds': preds})
if 'test' in cfg.DATASET.TEST_SET:
return {'Null': 0.0}, 0.0
SC_BIAS = 0.6
threshold = 0.5
gt_file = os.path.join(cfg.DATASET.ROOT,
'annot',
'gt_{}.mat'.format(cfg.DATASET.TEST_SET))
gt_dict = loadmat(gt_file)
dataset_joints = gt_dict['dataset_joints']
jnt_missing = gt_dict['jnt_missing']
pos_gt_src = gt_dict['pos_gt_src']
headboxes_src = gt_dict['headboxes_src']
pos_pred_src = np.transpose(preds, [1, 2, 0])
head = np.where(dataset_joints == 'head')[1][0]
lsho = np.where(dataset_joints == 'lsho')[1][0]
lelb = np.where(dataset_joints == 'lelb')[1][0]
lwri = np.where(dataset_joints == 'lwri')[1][0]
lhip = np.where(dataset_joints == 'lhip')[1][0]
lkne = np.where(dataset_joints == 'lkne')[1][0]
lank = np.where(dataset_joints == 'lank')[1][0]
rsho = np.where(dataset_joints == 'rsho')[1][0]
relb = np.where(dataset_joints == 'relb')[1][0]
rwri = np.where(dataset_joints == 'rwri')[1][0]
rkne = np.where(dataset_joints == 'rkne')[1][0]
rank = np.where(dataset_joints == 'rank')[1][0]
rhip = np.where(dataset_joints == 'rhip')[1][0]
jnt_visible = 1 - jnt_missing
uv_error = pos_pred_src - pos_gt_src
uv_err = np.linalg.norm(uv_error, axis=1)
headsizes = headboxes_src[1, :, :] - headboxes_src[0, :, :]
headsizes = np.linalg.norm(headsizes, axis=0)
headsizes *= SC_BIAS
scale = np.multiply(headsizes, np.ones((len(uv_err), 1)))
scaled_uv_err = np.divide(uv_err, scale)
scaled_uv_err = np.multiply(scaled_uv_err, jnt_visible)
jnt_count = np.sum(jnt_visible, axis=1)
less_than_threshold = np.multiply((scaled_uv_err <= threshold),
jnt_visible)
PCKh = np.divide(100.*np.sum(less_than_threshold, axis=1), jnt_count)
# save
rng = np.arange(0, 0.5+0.01, 0.01)
pckAll = np.zeros((len(rng), 16))
for r in range(len(rng)):
threshold = rng[r]
less_than_threshold = np.multiply(scaled_uv_err <= threshold,
jnt_visible)
pckAll[r, :] = np.divide(100.*np.sum(less_than_threshold, axis=1),
jnt_count)
PCKh = np.ma.array(PCKh, mask=False)
PCKh.mask[6:8] = True
jnt_count = np.ma.array(jnt_count, mask=False)
jnt_count.mask[6:8] = True
jnt_ratio = jnt_count / np.sum(jnt_count).astype(np.float64)
name_value = [
('Head', PCKh[head]),
('Shoulder', 0.5 * (PCKh[lsho] + PCKh[rsho])),
('Elbow', 0.5 * (PCKh[lelb] + PCKh[relb])),
('Wrist', 0.5 * (PCKh[lwri] + PCKh[rwri])),
('Hip', 0.5 * (PCKh[lhip] + PCKh[rhip])),
('Knee', 0.5 * (PCKh[lkne] + PCKh[rkne])),
('Ankle', 0.5 * (PCKh[lank] + PCKh[rank])),
('Mean', np.sum(PCKh * jnt_ratio)),
('Mean@0.1', np.sum(pckAll[11, :] * jnt_ratio))
]
name_value = OrderedDict(name_value)
return name_value, name_value['Mean']
project/hair_service_sd/keypoints/lib/models/__init__.py
+13
View File
@@ -0,0 +1,13 @@
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
project/hair_service_sd/keypoints/lib/models/pose_hrnet.py
+501
View File
@@ -0,0 +1,501 @@
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import logging
import torch
import torch.nn as nn
BN_MOMENTUM = 0.1
logger = logging.getLogger(__name__)
def conv3x3(in_planes, out_planes, stride=1):
"""3x3 convolution with padding"""
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
padding=1, bias=False)
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(BasicBlock, self).__init__()
self.conv1 = conv3x3(inplanes, planes, stride)
self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
self.relu = nn.ReLU(inplace=True)
self.conv2 = conv3x3(planes, planes)
self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
self.downsample = downsample
self.stride = stride
def forward(self, x):
residual = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
if self.downsample is not None:
residual = self.downsample(x)
out += residual
out = self.relu(out)
return out
class Bottleneck(nn.Module):
expansion = 4
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(Bottleneck, self).__init__()
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1,
bias=False)
self.bn3 = nn.BatchNorm2d(planes * self.expansion,
momentum=BN_MOMENTUM)
self.relu = nn.ReLU(inplace=True)
self.downsample = downsample
self.stride = stride
def forward(self, x):
residual = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
out = self.relu(out)
out = self.conv3(out)
out = self.bn3(out)
if self.downsample is not None:
residual = self.downsample(x)
out += residual
out = self.relu(out)
return out
class HighResolutionModule(nn.Module):
def __init__(self, num_branches, blocks, num_blocks, num_inchannels,
num_channels, fuse_method, multi_scale_output=True):
super(HighResolutionModule, self).__init__()
self._check_branches(
num_branches, blocks, num_blocks, num_inchannels, num_channels)
self.num_inchannels = num_inchannels
self.fuse_method = fuse_method
self.num_branches = num_branches
self.multi_scale_output = multi_scale_output
self.branches = self._make_branches(
num_branches, blocks, num_blocks, num_channels)
self.fuse_layers = self._make_fuse_layers()
self.relu = nn.ReLU(True)
def _check_branches(self, num_branches, blocks, num_blocks,
num_inchannels, num_channels):
if num_branches != len(num_blocks):
error_msg = 'NUM_BRANCHES({}) <> NUM_BLOCKS({})'.format(
num_branches, len(num_blocks))
logger.error(error_msg)
raise ValueError(error_msg)
if num_branches != len(num_channels):
error_msg = 'NUM_BRANCHES({}) <> NUM_CHANNELS({})'.format(
num_branches, len(num_channels))
logger.error(error_msg)
raise ValueError(error_msg)
if num_branches != len(num_inchannels):
error_msg = 'NUM_BRANCHES({}) <> NUM_INCHANNELS({})'.format(
num_branches, len(num_inchannels))
logger.error(error_msg)
raise ValueError(error_msg)
def _make_one_branch(self, branch_index, block, num_blocks, num_channels,
stride=1):
downsample = None
if stride != 1 or \
self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion:
downsample = nn.Sequential(
nn.Conv2d(
self.num_inchannels[branch_index],
num_channels[branch_index] * block.expansion,
kernel_size=1, stride=stride, bias=False
),
nn.BatchNorm2d(
num_channels[branch_index] * block.expansion,
momentum=BN_MOMENTUM
),
)
layers = []
layers.append(
block(
self.num_inchannels[branch_index],
num_channels[branch_index],
stride,
downsample
)
)
self.num_inchannels[branch_index] = \
num_channels[branch_index] * block.expansion
for i in range(1, num_blocks[branch_index]):
layers.append(
block(
self.num_inchannels[branch_index],
num_channels[branch_index]
)
)
return nn.Sequential(*layers)
def _make_branches(self, num_branches, block, num_blocks, num_channels):
branches = []
for i in range(num_branches):
branches.append(
self._make_one_branch(i, block, num_blocks, num_channels)
)
return nn.ModuleList(branches)
def _make_fuse_layers(self):
if self.num_branches == 1:
return None
num_branches = self.num_branches
num_inchannels = self.num_inchannels
fuse_layers = []
for i in range(num_branches if self.multi_scale_output else 1):
fuse_layer = []
for j in range(num_branches):
if j > i:
fuse_layer.append(
nn.Sequential(
nn.Conv2d(
num_inchannels[j],
num_inchannels[i],
1, 1, 0, bias=False
),
nn.BatchNorm2d(num_inchannels[i]),
nn.Upsample(scale_factor=2**(j-i), mode='nearest')
)
)
elif j == i:
fuse_layer.append(None)
else:
conv3x3s = []
for k in range(i-j):
if k == i - j - 1:
num_outchannels_conv3x3 = num_inchannels[i]
conv3x3s.append(
nn.Sequential(
nn.Conv2d(
num_inchannels[j],
num_outchannels_conv3x3,
3, 2, 1, bias=False
),
nn.BatchNorm2d(num_outchannels_conv3x3)
)
)
else:
num_outchannels_conv3x3 = num_inchannels[j]
conv3x3s.append(
nn.Sequential(
nn.Conv2d(
num_inchannels[j],
num_outchannels_conv3x3,
3, 2, 1, bias=False
),
nn.BatchNorm2d(num_outchannels_conv3x3),
nn.ReLU(True)
)
)
fuse_layer.append(nn.Sequential(*conv3x3s))
fuse_layers.append(nn.ModuleList(fuse_layer))
return nn.ModuleList(fuse_layers)
def get_num_inchannels(self):
return self.num_inchannels
def forward(self, x):
if self.num_branches == 1:
return [self.branches[0](x[0])]
for i in range(self.num_branches):
x[i] = self.branches[i](x[i])
x_fuse = []
for i in range(len(self.fuse_layers)):
y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
for j in range(1, self.num_branches):
if i == j:
y = y + x[j]
else:
y = y + self.fuse_layers[i][j](x[j])
x_fuse.append(self.relu(y))
return x_fuse
blocks_dict = {
'BASIC': BasicBlock,
'BOTTLENECK': Bottleneck
}
class PoseHighResolutionNet(nn.Module):
def __init__(self, cfg, **kwargs):
self.inplanes = 64
extra = cfg.MODEL.EXTRA
super(PoseHighResolutionNet, self).__init__()
# stem net
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1,
bias=False)
self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1,
bias=False)
self.bn2 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
self.relu = nn.ReLU(inplace=True)
self.layer1 = self._make_layer(Bottleneck, 64, 4)
self.stage2_cfg = cfg['MODEL']['EXTRA']['STAGE2']
num_channels = self.stage2_cfg['NUM_CHANNELS']
block = blocks_dict[self.stage2_cfg['BLOCK']]
num_channels = [
num_channels[i] * block.expansion for i in range(len(num_channels))
]
self.transition1 = self._make_transition_layer([256], num_channels)
self.stage2, pre_stage_channels = self._make_stage(
self.stage2_cfg, num_channels)
self.stage3_cfg = cfg['MODEL']['EXTRA']['STAGE3']
num_channels = self.stage3_cfg['NUM_CHANNELS']
block = blocks_dict[self.stage3_cfg['BLOCK']]
num_channels = [
num_channels[i] * block.expansion for i in range(len(num_channels))
]
self.transition2 = self._make_transition_layer(
pre_stage_channels, num_channels)
self.stage3, pre_stage_channels = self._make_stage(
self.stage3_cfg, num_channels)
self.stage4_cfg = cfg['MODEL']['EXTRA']['STAGE4']
num_channels = self.stage4_cfg['NUM_CHANNELS']
block = blocks_dict[self.stage4_cfg['BLOCK']]
num_channels = [
num_channels[i] * block.expansion for i in range(len(num_channels))
]
self.transition3 = self._make_transition_layer(
pre_stage_channels, num_channels)
self.stage4, pre_stage_channels = self._make_stage(
self.stage4_cfg, num_channels, multi_scale_output=False)
self.final_layer = nn.Conv2d(
in_channels=pre_stage_channels[0],
out_channels=cfg.MODEL.NUM_JOINTS,
kernel_size=extra.FINAL_CONV_KERNEL,
stride=1,
padding=1 if extra.FINAL_CONV_KERNEL == 3 else 0
)
self.pretrained_layers = cfg['MODEL']['EXTRA']['PRETRAINED_LAYERS']
def _make_transition_layer(
self, num_channels_pre_layer, num_channels_cur_layer):
num_branches_cur = len(num_channels_cur_layer)
num_branches_pre = len(num_channels_pre_layer)
transition_layers = []
for i in range(num_branches_cur):
if i < num_branches_pre:
if num_channels_cur_layer[i] != num_channels_pre_layer[i]:
transition_layers.append(
nn.Sequential(
nn.Conv2d(
num_channels_pre_layer[i],
num_channels_cur_layer[i],
3, 1, 1, bias=False
),
nn.BatchNorm2d(num_channels_cur_layer[i]),
nn.ReLU(inplace=True)
)
)
else:
transition_layers.append(None)
else:
conv3x3s = []
for j in range(i+1-num_branches_pre):
inchannels = num_channels_pre_layer[-1]
outchannels = num_channels_cur_layer[i] \
if j == i-num_branches_pre else inchannels
conv3x3s.append(
nn.Sequential(
nn.Conv2d(
inchannels, outchannels, 3, 2, 1, bias=False
),
nn.BatchNorm2d(outchannels),
nn.ReLU(inplace=True)
)
)
transition_layers.append(nn.Sequential(*conv3x3s))
return nn.ModuleList(transition_layers)
def _make_layer(self, block, planes, blocks, stride=1):
downsample = None
if stride != 1 or self.inplanes != planes * block.expansion:
downsample = nn.Sequential(
nn.Conv2d(
self.inplanes, planes * block.expansion,
kernel_size=1, stride=stride, bias=False
),
nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM),
)
layers = []
layers.append(block(self.inplanes, planes, stride, downsample))
self.inplanes = planes * block.expansion
for i in range(1, blocks):
layers.append(block(self.inplanes, planes))
return nn.Sequential(*layers)
def _make_stage(self, layer_config, num_inchannels,
multi_scale_output=True):
num_modules = layer_config['NUM_MODULES']
num_branches = layer_config['NUM_BRANCHES']
num_blocks = layer_config['NUM_BLOCKS']
num_channels = layer_config['NUM_CHANNELS']
block = blocks_dict[layer_config['BLOCK']]
fuse_method = layer_config['FUSE_METHOD']
modules = []
for i in range(num_modules):
# multi_scale_output is only used last module
if not multi_scale_output and i == num_modules - 1:
reset_multi_scale_output = False
else:
reset_multi_scale_output = True
modules.append(
HighResolutionModule(
num_branches,
block,
num_blocks,
num_inchannels,
num_channels,
fuse_method,
reset_multi_scale_output
)
)
num_inchannels = modules[-1].get_num_inchannels()
return nn.Sequential(*modules), num_inchannels
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.conv2(x)
x = self.bn2(x)
x = self.relu(x)
x = self.layer1(x)
x_list = []
for i in range(self.stage2_cfg['NUM_BRANCHES']):
if self.transition1[i] is not None:
x_list.append(self.transition1[i](x))
else:
x_list.append(x)
y_list = self.stage2(x_list)
x_list = []
for i in range(self.stage3_cfg['NUM_BRANCHES']):
if self.transition2[i] is not None:
x_list.append(self.transition2[i](y_list[-1]))
else:
x_list.append(y_list[i])
y_list = self.stage3(x_list)
x_list = []
for i in range(self.stage4_cfg['NUM_BRANCHES']):
if self.transition3[i] is not None:
x_list.append(self.transition3[i](y_list[-1]))
else:
x_list.append(y_list[i])
y_list = self.stage4(x_list)
x = self.final_layer(y_list[0])
return x
def init_weights(self, pretrained=''):
logger.info('=> init weights from normal distribution')
for m in self.modules():
if isinstance(m, nn.Conv2d):
# nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
nn.init.normal_(m.weight, std=0.001)
for name, _ in m.named_parameters():
if name in ['bias']:
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.ConvTranspose2d):
nn.init.normal_(m.weight, std=0.001)
for name, _ in m.named_parameters():
if name in ['bias']:
nn.init.constant_(m.bias, 0)
if os.path.isfile(pretrained):
pretrained_state_dict = torch.load(pretrained)
logger.info('=> loading pretrained model {}'.format(pretrained))
need_init_state_dict = {}
for name, m in pretrained_state_dict.items():
if name.split('.')[0] in self.pretrained_layers \
or self.pretrained_layers[0] is '*':
need_init_state_dict[name] = m
self.load_state_dict(need_init_state_dict, strict=False)
elif pretrained:
logger.error('=> please download pre-trained models first!')
raise ValueError('{} is not exist!'.format(pretrained))
def get_pose_net(cfg, is_train, **kwargs):
model = PoseHighResolutionNet(cfg, **kwargs)
if is_train and cfg.MODEL.INIT_WEIGHTS:
model.init_weights(cfg.MODEL.PRETRAINED)
return model
project/hair_service_sd/keypoints/lib/models/pose_resnet.py
+271
View File
@@ -0,0 +1,271 @@
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import logging
import torch
import torch.nn as nn
BN_MOMENTUM = 0.1
logger = logging.getLogger(__name__)
def conv3x3(in_planes, out_planes, stride=1):
"""3x3 convolution with padding"""
return nn.Conv2d(
in_planes, out_planes, kernel_size=3, stride=stride,
padding=1, bias=False
)
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(BasicBlock, self).__init__()
self.conv1 = conv3x3(inplanes, planes, stride)
self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
self.relu = nn.ReLU(inplace=True)
self.conv2 = conv3x3(planes, planes)
self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
self.downsample = downsample
self.stride = stride
def forward(self, x):
residual = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
if self.downsample is not None:
residual = self.downsample(x)
out += residual
out = self.relu(out)
return out
class Bottleneck(nn.Module):
expansion = 4
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(Bottleneck, self).__init__()
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1,
bias=False)
self.bn3 = nn.BatchNorm2d(planes * self.expansion,
momentum=BN_MOMENTUM)
self.relu = nn.ReLU(inplace=True)
self.downsample = downsample
self.stride = stride
def forward(self, x):
residual = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
out = self.relu(out)
out = self.conv3(out)
out = self.bn3(out)
if self.downsample is not None:
residual = self.downsample(x)
out += residual
out = self.relu(out)
return out
class PoseResNet(nn.Module):
def __init__(self, block, layers, cfg, **kwargs):
self.inplanes = 64
extra = cfg.MODEL.EXTRA
self.deconv_with_bias = extra.DECONV_WITH_BIAS
super(PoseResNet, self).__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
bias=False)
self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
# used for deconv layers
self.deconv_layers = self._make_deconv_layer(
extra.NUM_DECONV_LAYERS,
extra.NUM_DECONV_FILTERS,
extra.NUM_DECONV_KERNELS,
)
self.final_layer = nn.Conv2d(
in_channels=extra.NUM_DECONV_FILTERS[-1],
out_channels=cfg.MODEL.NUM_JOINTS,
kernel_size=extra.FINAL_CONV_KERNEL,
stride=1,
padding=1 if extra.FINAL_CONV_KERNEL == 3 else 0
)
def _make_layer(self, block, planes, blocks, stride=1):
downsample = None
if stride != 1 or self.inplanes != planes * block.expansion:
downsample = nn.Sequential(
nn.Conv2d(self.inplanes, planes * block.expansion,
kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM),
)
layers = []
layers.append(block(self.inplanes, planes, stride, downsample))
self.inplanes = planes * block.expansion
for i in range(1, blocks):
layers.append(block(self.inplanes, planes))
return nn.Sequential(*layers)
def _get_deconv_cfg(self, deconv_kernel, index):
if deconv_kernel == 4:
padding = 1
output_padding = 0
elif deconv_kernel == 3:
padding = 1
output_padding = 1
elif deconv_kernel == 2:
padding = 0
output_padding = 0
return deconv_kernel, padding, output_padding
def _make_deconv_layer(self, num_layers, num_filters, num_kernels):
assert num_layers == len(num_filters), \
'ERROR: num_deconv_layers is different len(num_deconv_filters)'
assert num_layers == len(num_kernels), \
'ERROR: num_deconv_layers is different len(num_deconv_filters)'
layers = []
for i in range(num_layers):
kernel, padding, output_padding = \
self._get_deconv_cfg(num_kernels[i], i)
planes = num_filters[i]
layers.append(
nn.ConvTranspose2d(
in_channels=self.inplanes,
out_channels=planes,
kernel_size=kernel,
stride=2,
padding=padding,
output_padding=output_padding,
bias=self.deconv_with_bias))
layers.append(nn.BatchNorm2d(planes, momentum=BN_MOMENTUM))
layers.append(nn.ReLU(inplace=True))
self.inplanes = planes
return nn.Sequential(*layers)
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.deconv_layers(x)
x = self.final_layer(x)
return x
def init_weights(self, pretrained=''):
if os.path.isfile(pretrained):
logger.info('=> init deconv weights from normal distribution')
for name, m in self.deconv_layers.named_modules():
if isinstance(m, nn.ConvTranspose2d):
logger.info('=> init {}.weight as normal(0, 0.001)'.format(name))
logger.info('=> init {}.bias as 0'.format(name))
nn.init.normal_(m.weight, std=0.001)
if self.deconv_with_bias:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.BatchNorm2d):
logger.info('=> init {}.weight as 1'.format(name))
logger.info('=> init {}.bias as 0'.format(name))
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
logger.info('=> init final conv weights from normal distribution')
for m in self.final_layer.modules():
if isinstance(m, nn.Conv2d):
# nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
logger.info('=> init {}.weight as normal(0, 0.001)'.format(name))
logger.info('=> init {}.bias as 0'.format(name))
nn.init.normal_(m.weight, std=0.001)
nn.init.constant_(m.bias, 0)
pretrained_state_dict = torch.load(pretrained)
logger.info('=> loading pretrained model {}'.format(pretrained))
self.load_state_dict(pretrained_state_dict, strict=False)
else:
logger.info('=> init weights from normal distribution')
for m in self.modules():
if isinstance(m, nn.Conv2d):
# nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
nn.init.normal_(m.weight, std=0.001)
# 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.ConvTranspose2d):
nn.init.normal_(m.weight, std=0.001)
if self.deconv_with_bias:
nn.init.constant_(m.bias, 0)
resnet_spec = {
18: (BasicBlock, [2, 2, 2, 2]),
34: (BasicBlock, [3, 4, 6, 3]),
50: (Bottleneck, [3, 4, 6, 3]),
101: (Bottleneck, [3, 4, 23, 3]),
152: (Bottleneck, [3, 8, 36, 3])
}
def get_pose_net(cfg, is_train, **kwargs):
num_layers = cfg.MODEL.EXTRA.NUM_LAYERS
block_class, layers = resnet_spec[num_layers]
model = PoseResNet(block_class, layers, cfg, **kwargs)
if is_train and cfg.MODEL.INIT_WEIGHTS:
model.init_weights(cfg.MODEL.PRETRAINED)
return model
project/hair_service_sd/keypoints/lib/nms/__init__.py
View File
project/hair_service_sd/keypoints/lib/nms/nms.py
+164
View File
@@ -0,0 +1,164 @@
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Modified from py-faster-rcnn (https://github.com/rbgirshick/py-faster-rcnn)
# ------------------------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
def py_nms_wrapper(thresh):
def _nms(dets):
return nms(dets, thresh)
return _nms
def nms(dets, thresh):
"""
greedily select boxes with high confidence and overlap with current maximum <= thresh
rule out overlap >= thresh
:param dets: [[x1, y1, x2, y2 score]]
:param thresh: retain overlap < thresh
:return: indexes to keep
"""
if dets.shape[0] == 0:
return []
x1 = dets[:, 0]
y1 = dets[:, 1]
x2 = dets[:, 2]
y2 = dets[:, 3]
scores = 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 = np.maximum(x1[i], x1[order[1:]])
yy1 = np.maximum(y1[i], y1[order[1:]])
xx2 = np.minimum(x2[i], x2[order[1:]])
yy2 = np.minimum(y2[i], y2[order[1:]])
w = np.maximum(0.0, xx2 - xx1 + 1)
h = np.maximum(0.0, yy2 - yy1 + 1)
inter = w * h
ovr = inter / (areas[i] + areas[order[1:]] - inter)
inds = np.where(ovr <= thresh)[0]
order = order[inds + 1]
return keep
def oks_iou(g, d, a_g, a_d, sigmas=None, in_vis_thre=None):
if not isinstance(sigmas, np.ndarray):
sigmas = np.array([.26, .25, .25, .35, .35, .79, .79, .72, .72, .62, .62, 1.07, 1.07, .87, .87, .89, .89]) / 10.0
vars = (sigmas * 2) ** 2
xg = g[0::3]
yg = g[1::3]
vg = g[2::3]
ious = np.zeros((d.shape[0]))
for n_d in range(0, d.shape[0]):
xd = d[n_d, 0::3]
yd = d[n_d, 1::3]
vd = d[n_d, 2::3]
dx = xd - xg
dy = yd - yg
e = (dx ** 2 + dy ** 2) / vars / ((a_g + a_d[n_d]) / 2 + np.spacing(1)) / 2
if in_vis_thre is not None:
ind = list(vg > in_vis_thre) and list(vd > in_vis_thre)
e = e[ind]
ious[n_d] = np.sum(np.exp(-e)) / e.shape[0] if e.shape[0] != 0 else 0.0
return ious
def oks_nms(kpts_db, thresh, sigmas=None, in_vis_thre=None):
"""
greedily select boxes with high confidence and overlap with current maximum <= thresh
rule out overlap >= thresh, overlap = oks
:param kpts_db
:param thresh: retain overlap < thresh
:return: indexes to keep
"""
if len(kpts_db) == 0:
return []
scores = np.array([kpts_db[i]['score'] for i in range(len(kpts_db))])
kpts = np.array([kpts_db[i]['keypoints'].flatten() for i in range(len(kpts_db))])
areas = np.array([kpts_db[i]['area'] for i in range(len(kpts_db))])
order = scores.argsort()[::-1]
keep = []
while order.size > 0:
i = order[0]
keep.append(i)
oks_ovr = oks_iou(kpts[i], kpts[order[1:]], areas[i], areas[order[1:]], sigmas, in_vis_thre)
inds = np.where(oks_ovr <= thresh)[0]
order = order[inds + 1]
return keep
def rescore(overlap, scores, thresh, type='gaussian'):
assert overlap.shape[0] == scores.shape[0]
if type == 'linear':
inds = np.where(overlap >= thresh)[0]
scores[inds] = scores[inds] * (1 - overlap[inds])
else:
scores = scores * np.exp(- overlap**2 / thresh)
return scores
def soft_oks_nms(kpts_db, thresh, sigmas=None, in_vis_thre=None):
"""
greedily select boxes with high confidence and overlap with current maximum <= thresh
rule out overlap >= thresh, overlap = oks
:param kpts_db
:param thresh: retain overlap < thresh
:return: indexes to keep
"""
if len(kpts_db) == 0:
return []
scores = np.array([kpts_db[i]['score'] for i in range(len(kpts_db))])
kpts = np.array([kpts_db[i]['keypoints'].flatten() for i in range(len(kpts_db))])
areas = np.array([kpts_db[i]['area'] for i in range(len(kpts_db))])
order = scores.argsort()[::-1]
scores = scores[order]
# max_dets = order.size
max_dets = 20
keep = np.zeros(max_dets, dtype=np.intp)
keep_cnt = 0
while order.size > 0 and keep_cnt < max_dets:
i = order[0]
oks_ovr = oks_iou(kpts[i], kpts[order[1:]], areas[i], areas[order[1:]], sigmas, in_vis_thre)
order = order[1:]
scores = rescore(oks_ovr, scores[1:], thresh)
tmp = scores.argsort()[::-1]
order = order[tmp]
scores = scores[tmp]
keep[keep_cnt] = i
keep_cnt += 1
keep = keep[:keep_cnt]
return keep
# kpts_db = kpts_db[:keep_cnt]
# return kpts_db
project/hair_service_sd/keypoints/lib/utils/__init__.py
View File
project/hair_service_sd/keypoints/lib/utils/transforms.py
+121
View File
@@ -0,0 +1,121 @@
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
import cv2
def flip_back(output_flipped, matched_parts):
'''
ouput_flipped: numpy.ndarray(batch_size, num_joints, height, width)
'''
assert output_flipped.ndim == 4,\
'output_flipped should be [batch_size, num_joints, height, width]'
output_flipped = output_flipped[:, :, :, ::-1]
for pair in matched_parts:
tmp = output_flipped[:, pair[0], :, :].copy()
output_flipped[:, pair[0], :, :] = output_flipped[:, pair[1], :, :]
output_flipped[:, pair[1], :, :] = tmp
return output_flipped
def fliplr_joints(joints, joints_vis, width, matched_parts):
"""
flip coords
"""
# Flip horizontal
joints[:, 0] = width - joints[:, 0] - 1
# Change left-right parts
for pair in matched_parts:
joints[pair[0], :], joints[pair[1], :] = \
joints[pair[1], :], joints[pair[0], :].copy()
joints_vis[pair[0], :], joints_vis[pair[1], :] = \
joints_vis[pair[1], :], joints_vis[pair[0], :].copy()
return joints*joints_vis, joints_vis
def transform_preds(coords, center, scale, output_size):
target_coords = np.zeros(coords.shape)
trans = get_affine_transform(center, scale, 0, output_size, inv=1)
for p in range(coords.shape[0]):
target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans)
return target_coords
def get_affine_transform(
center, scale, rot, output_size,
shift=np.array([0, 0], dtype=np.float32), inv=0
):
if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
print(scale)
scale = np.array([scale, scale])
scale_tmp = scale * 200.0
src_w = scale_tmp[0]
dst_w = output_size[0]
dst_h = output_size[1]
rot_rad = np.pi * rot / 180
src_dir = get_dir([0, src_w * -0.5], rot_rad)
dst_dir = np.array([0, dst_w * -0.5], np.float32)
src = np.zeros((3, 2), dtype=np.float32)
dst = np.zeros((3, 2), dtype=np.float32)
src[0, :] = center + scale_tmp * shift
src[1, :] = center + src_dir + scale_tmp * shift
dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
src[2:, :] = get_3rd_point(src[0, :], src[1, :])
dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])
if inv:
trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
else:
trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
return trans
def affine_transform(pt, t):
new_pt = np.array([pt[0], pt[1], 1.]).T
new_pt = np.dot(t, new_pt)
return new_pt[:2]
def get_3rd_point(a, b):
direct = a - b
return b + np.array([-direct[1], direct[0]], dtype=np.float32)
def get_dir(src_point, rot_rad):
sn, cs = np.sin(rot_rad), np.cos(rot_rad)
src_result = [0, 0]
src_result[0] = src_point[0] * cs - src_point[1] * sn
src_result[1] = src_point[0] * sn + src_point[1] * cs
return src_result
def crop(img, center, scale, output_size, rot=0):
trans = get_affine_transform(center, scale, rot, output_size)
dst_img = cv2.warpAffine(
img, trans, (int(output_size[0]), int(output_size[1])),
flags=cv2.INTER_LINEAR
)
return dst_img
project/hair_service_sd/keypoints/lib/utils/utils.py
+203
View File
@@ -0,0 +1,203 @@
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import logging
import time
from collections import namedtuple
from pathlib import Path
import torch
import torch.optim as optim
import torch.nn as nn
def create_logger(cfg, cfg_name, phase='train'):
root_output_dir = Path(cfg.OUTPUT_DIR)
# set up logger
if not root_output_dir.exists():
print('=> creating {}'.format(root_output_dir))
root_output_dir.mkdir()
dataset = cfg.DATASET.DATASET + '_' + cfg.DATASET.HYBRID_JOINTS_TYPE \
if cfg.DATASET.HYBRID_JOINTS_TYPE else cfg.DATASET.DATASET
dataset = dataset.replace(':', '_')
model = cfg.MODEL.NAME
cfg_name = os.path.basename(cfg_name).split('.')[0]
final_output_dir = root_output_dir / dataset / model / cfg_name
print('=> creating {}'.format(final_output_dir))
final_output_dir.mkdir(parents=True, exist_ok=True)
time_str = time.strftime('%Y-%m-%d-%H-%M')
log_file = '{}_{}_{}.log'.format(cfg_name, time_str, phase)
final_log_file = final_output_dir / log_file
head = '%(asctime)-15s %(message)s'
logging.basicConfig(filename=str(final_log_file),
format=head)
logger = logging.getLogger()
logger.setLevel(logging.INFO)
console = logging.StreamHandler()
logging.getLogger('').addHandler(console)
tensorboard_log_dir = Path(cfg.LOG_DIR) / dataset / model / \
(cfg_name + '_' + time_str)
print('=> creating {}'.format(tensorboard_log_dir))
tensorboard_log_dir.mkdir(parents=True, exist_ok=True)
return logger, str(final_output_dir), str(tensorboard_log_dir)
def get_optimizer(cfg, model):
optimizer = None
if cfg.TRAIN.OPTIMIZER == 'sgd':
optimizer = optim.SGD(
model.parameters(),
lr=cfg.TRAIN.LR,
momentum=cfg.TRAIN.MOMENTUM,
weight_decay=cfg.TRAIN.WD,
nesterov=cfg.TRAIN.NESTEROV
)
elif cfg.TRAIN.OPTIMIZER == 'adam':
optimizer = optim.Adam(
model.parameters(),
lr=cfg.TRAIN.LR
)
return optimizer
def save_checkpoint(states, is_best, output_dir,
filename='checkpoint.pth'):
torch.save(states, os.path.join(output_dir, filename))
if is_best and 'state_dict' in states:
torch.save(states['best_state_dict'],
os.path.join(output_dir, 'model_best.pth'))
def get_model_summary(model, *input_tensors, item_length=26, verbose=False):
"""
:param model:
:param input_tensors:
:param item_length:
:return:
"""
summary = []
ModuleDetails = namedtuple(
"Layer", ["name", "input_size", "output_size", "num_parameters", "multiply_adds"])
hooks = []
layer_instances = {}
def add_hooks(module):
def hook(module, input, output):
class_name = str(module.__class__.__name__)
instance_index = 1
if class_name not in layer_instances:
layer_instances[class_name] = instance_index
else:
instance_index = layer_instances[class_name] + 1
layer_instances[class_name] = instance_index
layer_name = class_name + "_" + str(instance_index)
params = 0
if class_name.find("Conv") != -1 or class_name.find("BatchNorm") != -1 or \
class_name.find("Linear") != -1:
for param_ in module.parameters():
params += param_.view(-1).size(0)
flops = "Not Available"
if class_name.find("Conv") != -1 and hasattr(module, "weight"):
flops = (
torch.prod(
torch.LongTensor(list(module.weight.data.size()))) *
torch.prod(
torch.LongTensor(list(output.size())[2:]))).item()
elif isinstance(module, nn.Linear):
flops = (torch.prod(torch.LongTensor(list(output.size()))) \
* input[0].size(1)).item()
if isinstance(input[0], list):
input = input[0]
if isinstance(output, list):
output = output[0]
summary.append(
ModuleDetails(
name=layer_name,
input_size=list(input[0].size()),
output_size=list(output.size()),
num_parameters=params,
multiply_adds=flops)
)
if not isinstance(module, nn.ModuleList) \
and not isinstance(module, nn.Sequential) \
and module != model:
hooks.append(module.register_forward_hook(hook))
model.eval()
model.apply(add_hooks)
space_len = item_length
model(*input_tensors)
for hook in hooks:
hook.remove()
details = ''
if verbose:
details = "Model Summary" + \
os.linesep + \
"Name{}Input Size{}Output Size{}Parameters{}Multiply Adds (Flops){}".format(
' ' * (space_len - len("Name")),
' ' * (space_len - len("Input Size")),
' ' * (space_len - len("Output Size")),
' ' * (space_len - len("Parameters")),
' ' * (space_len - len("Multiply Adds (Flops)"))) \
+ os.linesep + '-' * space_len * 5 + os.linesep
params_sum = 0
flops_sum = 0
for layer in summary:
params_sum += layer.num_parameters
if layer.multiply_adds != "Not Available":
flops_sum += layer.multiply_adds
if verbose:
details += "{}{}{}{}{}{}{}{}{}{}".format(
layer.name,
' ' * (space_len - len(layer.name)),
layer.input_size,
' ' * (space_len - len(str(layer.input_size))),
layer.output_size,
' ' * (space_len - len(str(layer.output_size))),
layer.num_parameters,
' ' * (space_len - len(str(layer.num_parameters))),
layer.multiply_adds,
' ' * (space_len - len(str(layer.multiply_adds)))) \
+ os.linesep + '-' * space_len * 5 + os.linesep
details += os.linesep \
+ "Total Parameters: {:,}".format(params_sum) \
+ os.linesep + '-' * space_len * 5 + os.linesep
details += "Total Multiply Adds (For Convolution and Linear Layers only): {:,} GFLOPs".format(flops_sum/(1024**3)) \
+ os.linesep + '-' * space_len * 5 + os.linesep
details += "Number of Layers" + os.linesep
for layer in layer_instances:
details += "{} : {} layers ".format(layer, layer_instances[layer])
return details
project/hair_service_sd/keypoints/lib/utils/vis.py
+141
View File
@@ -0,0 +1,141 @@
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import math
import numpy as np
import torchvision
import cv2
from core.inference import get_max_preds
def save_batch_image_with_joints(batch_image, batch_joints, batch_joints_vis,
file_name, nrow=8, padding=2):
'''
batch_image: [batch_size, channel, height, width]
batch_joints: [batch_size, num_joints, 3],
batch_joints_vis: [batch_size, num_joints, 1],
}
'''
grid = torchvision.utils.make_grid(batch_image, nrow, padding, True)
ndarr = grid.mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()
ndarr = ndarr.copy()
nmaps = batch_image.size(0)
xmaps = min(nrow, nmaps)
ymaps = int(math.ceil(float(nmaps) / xmaps))
height = int(batch_image.size(2) + padding)
width = int(batch_image.size(3) + padding)
k = 0
for y in range(ymaps):
for x in range(xmaps):
if k >= nmaps:
break
joints = batch_joints[k]
joints_vis = batch_joints_vis[k]
for joint, joint_vis in zip(joints, joints_vis):
joint[0] = x * width + padding + joint[0]
joint[1] = y * height + padding + joint[1]
if joint_vis[0]:
cv2.circle(ndarr, (int(joint[0]), int(joint[1])), 2, [255, 0, 0], 2)
k = k + 1
cv2.imwrite(file_name, ndarr)
def save_batch_heatmaps(batch_image, batch_heatmaps, file_name,
normalize=True):
'''
batch_image: [batch_size, channel, height, width]
batch_heatmaps: ['batch_size, num_joints, height, width]
file_name: saved file name
'''
if normalize:
batch_image = batch_image.clone()
min = float(batch_image.min())
max = float(batch_image.max())
batch_image.add_(-min).div_(max - min + 1e-5)
batch_size = batch_heatmaps.size(0)
num_joints = batch_heatmaps.size(1)
heatmap_height = batch_heatmaps.size(2)
heatmap_width = batch_heatmaps.size(3)
grid_image = np.zeros((batch_size*heatmap_height,
(num_joints+1)*heatmap_width,
3),
dtype=np.uint8)
preds, maxvals = get_max_preds(batch_heatmaps.detach().cpu().numpy())
for i in range(batch_size):
image = batch_image[i].mul(255)\
.clamp(0, 255)\
.byte()\
.permute(1, 2, 0)\
.cpu().numpy()
heatmaps = batch_heatmaps[i].mul(255)\
.clamp(0, 255)\
.byte()\
.cpu().numpy()
resized_image = cv2.resize(image,
(int(heatmap_width), int(heatmap_height)))
height_begin = heatmap_height * i
height_end = heatmap_height * (i + 1)
for j in range(num_joints):
cv2.circle(resized_image,
(int(preds[i][j][0]), int(preds[i][j][1])),
1, [0, 0, 255], 1)
heatmap = heatmaps[j, :, :]
colored_heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
masked_image = colored_heatmap*0.7 + resized_image*0.3
cv2.circle(masked_image,
(int(preds[i][j][0]), int(preds[i][j][1])),
1, [0, 0, 255], 1)
width_begin = heatmap_width * (j+1)
width_end = heatmap_width * (j+2)
grid_image[height_begin:height_end, width_begin:width_end, :] = \
masked_image
# grid_image[height_begin:height_end, width_begin:width_end, :] = \
# colored_heatmap*0.7 + resized_image*0.3
grid_image[height_begin:height_end, 0:heatmap_width, :] = resized_image
cv2.imwrite(file_name, grid_image)
def save_debug_images(config, input, meta, target, joints_pred, output,
prefix):
if not config.DEBUG.DEBUG:
return
if config.DEBUG.SAVE_BATCH_IMAGES_GT:
save_batch_image_with_joints(
input, meta['joints'], meta['joints_vis'],
'{}_gt.jpg'.format(prefix)
)
if config.DEBUG.SAVE_BATCH_IMAGES_PRED:
save_batch_image_with_joints(
input, joints_pred, meta['joints_vis'],
'{}_pred.jpg'.format(prefix)
)
if config.DEBUG.SAVE_HEATMAPS_GT:
save_batch_heatmaps(
input, target, '{}_hm_gt.jpg'.format(prefix)
)
if config.DEBUG.SAVE_HEATMAPS_PRED:
save_batch_heatmaps(
input, output, '{}_hm_pred.jpg'.format(prefix)
)
project/hair_service_sd/keypoints/lib/utils/zipreader.py
+70
View File
@@ -0,0 +1,70 @@
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import zipfile
import xml.etree.ElementTree as ET
import cv2
import numpy as np
_im_zfile = []
_xml_path_zip = []
_xml_zfile = []
def imread(filename, flags=cv2.IMREAD_COLOR):
global _im_zfile
path = filename
pos_at = path.index('@')
if pos_at == -1:
print("character '@' is not found from the given path '%s'"%(path))
assert 0
path_zip = path[0: pos_at]
path_img = path[pos_at + 2:]
if not os.path.isfile(path_zip):
print("zip file '%s' is not found"%(path_zip))
assert 0
for i in range(len(_im_zfile)):
if _im_zfile[i]['path'] == path_zip:
data = _im_zfile[i]['zipfile'].read(path_img)
return cv2.imdecode(np.frombuffer(data, np.uint8), flags)
_im_zfile.append({
'path': path_zip,
'zipfile': zipfile.ZipFile(path_zip, 'r')
})
data = _im_zfile[-1]['zipfile'].read(path_img)
return cv2.imdecode(np.frombuffer(data, np.uint8), flags)
def xmlread(filename):
global _xml_path_zip
global _xml_zfile
path = filename
pos_at = path.index('@')
if pos_at == -1:
print("character '@' is not found from the given path '%s'"%(path))
assert 0
path_zip = path[0: pos_at]
path_xml = path[pos_at + 2:]
if not os.path.isfile(path_zip):
print("zip file '%s' is not found"%(path_zip))
assert 0
for i in xrange(len(_xml_path_zip)):
if _xml_path_zip[i] == path_zip:
data = _xml_zfile[i].open(path_xml)
return ET.fromstring(data.read())
_xml_path_zip.append(path_zip)
print("read new xml file '%s'"%(path_zip))
_xml_zfile.append(zipfile.ZipFile(path_zip, 'r'))
data = _xml_zfile[-1].open(path_xml)
return ET.fromstring(data.read())
project/hair_service_sd/models/layers/data/FDDB/img_list.txt
+2845
View File
File diff suppressed because it is too large Load Diff
project/hair_service_sd/models/layers/data/__init__.py
+3
View File
@@ -0,0 +1,3 @@
from .wider_face import WiderFaceDetection, detection_collate
from .data_augment import *
from .config import *
project/hair_service_sd/models/layers/data/config.py
+42
View File
@@ -0,0 +1,42 @@
# config.py
cfg_mnet = {
'name': 'mobilenet0.25',
'min_sizes': [[16, 32], [64, 128], [256, 512]],
'steps': [8, 16, 32],
'variance': [0.1, 0.2],
'clip': False,
'loc_weight': 2.0,
'gpu_train': True,
'batch_size': 32,
'ngpu': 1,
'epoch': 250,
'decay1': 190,
'decay2': 220,
'image_size': 640,
'pretrain': True,
'return_layers': {'stage1': 1, 'stage2': 2, 'stage3': 3},
'in_channel': 32,
'out_channel': 64
}
cfg_re50 = {
'name': 'Resnet50',
'min_sizes': [[16, 32], [64, 128], [256, 512]],
'steps': [8, 16, 32],
'variance': [0.1, 0.2],
'clip': False,
'loc_weight': 2.0,
'gpu_train': True,
'batch_size': 24,
'ngpu': 4,
'epoch': 100,
'decay1': 70,
'decay2': 90,
'image_size': 840,
'pretrain': True,
'return_layers': {'layer2': 1, 'layer3': 2, 'layer4': 3},
'in_channel': 256,
'out_channel': 256
}
project/hair_service_sd/models/layers/data/data_augment.py
+237
View File
@@ -0,0 +1,237 @@
import cv2
import numpy as np
import random
from utils.box_utils_Retina import matrix_iof
def _crop(image, boxes, labels, landm, img_dim):
height, width, _ = image.shape
pad_image_flag = True
for _ in range(250):
"""
if random.uniform(0, 1) <= 0.2:
scale = 1.0
else:
scale = random.uniform(0.3, 1.0)
"""
PRE_SCALES = [0.3, 0.45, 0.6, 0.8, 1.0]
scale = random.choice(PRE_SCALES)
short_side = min(width, height)
w = int(scale * short_side)
h = w
if width == w:
l = 0
else:
l = random.randrange(width - w)
if height == h:
t = 0
else:
t = random.randrange(height - h)
roi = np.array((l, t, l + w, t + h))
value = matrix_iof(boxes, roi[np.newaxis])
flag = (value >= 1)
if not flag.any():
continue
centers = (boxes[:, :2] + boxes[:, 2:]) / 2
mask_a = np.logical_and(roi[:2] < centers, centers < roi[2:]).all(axis=1)
boxes_t = boxes[mask_a].copy()
labels_t = labels[mask_a].copy()
landms_t = landm[mask_a].copy()
landms_t = landms_t.reshape([-1, 5, 2])
if boxes_t.shape[0] == 0:
continue
image_t = image[roi[1]:roi[3], roi[0]:roi[2]]
boxes_t[:, :2] = np.maximum(boxes_t[:, :2], roi[:2])
boxes_t[:, :2] -= roi[:2]
boxes_t[:, 2:] = np.minimum(boxes_t[:, 2:], roi[2:])
boxes_t[:, 2:] -= roi[:2]
# landm
landms_t[:, :, :2] = landms_t[:, :, :2] - roi[:2]
landms_t[:, :, :2] = np.maximum(landms_t[:, :, :2], np.array([0, 0]))
landms_t[:, :, :2] = np.minimum(landms_t[:, :, :2], roi[2:] - roi[:2])
landms_t = landms_t.reshape([-1, 10])
# make sure that the cropped image contains at least one face > 16 pixel at training image scale
b_w_t = (boxes_t[:, 2] - boxes_t[:, 0] + 1) / w * img_dim
b_h_t = (boxes_t[:, 3] - boxes_t[:, 1] + 1) / h * img_dim
mask_b = np.minimum(b_w_t, b_h_t) > 0.0
boxes_t = boxes_t[mask_b]
labels_t = labels_t[mask_b]
landms_t = landms_t[mask_b]
if boxes_t.shape[0] == 0:
continue
pad_image_flag = False
return image_t, boxes_t, labels_t, landms_t, pad_image_flag
return image, boxes, labels, landm, pad_image_flag
def _distort(image):
def _convert(image, alpha=1, beta=0):
tmp = image.astype(float) * alpha + beta
tmp[tmp < 0] = 0
tmp[tmp > 255] = 255
image[:] = tmp
image = image.copy()
if random.randrange(2):
#brightness distortion
if random.randrange(2):
_convert(image, beta=random.uniform(-32, 32))
#contrast distortion
if random.randrange(2):
_convert(image, alpha=random.uniform(0.5, 1.5))
image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
#saturation distortion
if random.randrange(2):
_convert(image[:, :, 1], alpha=random.uniform(0.5, 1.5))
#hue distortion
if random.randrange(2):
tmp = image[:, :, 0].astype(int) + random.randint(-18, 18)
tmp %= 180
image[:, :, 0] = tmp
image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
else:
#brightness distortion
if random.randrange(2):
_convert(image, beta=random.uniform(-32, 32))
image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
#saturation distortion
if random.randrange(2):
_convert(image[:, :, 1], alpha=random.uniform(0.5, 1.5))
#hue distortion
if random.randrange(2):
tmp = image[:, :, 0].astype(int) + random.randint(-18, 18)
tmp %= 180
image[:, :, 0] = tmp
image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
#contrast distortion
if random.randrange(2):
_convert(image, alpha=random.uniform(0.5, 1.5))
return image
def _expand(image, boxes, fill, p):
if random.randrange(2):
return image, boxes
height, width, depth = image.shape
scale = random.uniform(1, p)
w = int(scale * width)
h = int(scale * height)
left = random.randint(0, w - width)
top = random.randint(0, h - height)
boxes_t = boxes.copy()
boxes_t[:, :2] += (left, top)
boxes_t[:, 2:] += (left, top)
expand_image = np.empty(
(h, w, depth),
dtype=image.dtype)
expand_image[:, :] = fill
expand_image[top:top + height, left:left + width] = image
image = expand_image
return image, boxes_t
def _mirror(image, boxes, landms):
_, width, _ = image.shape
if random.randrange(2):
image = image[:, ::-1]
boxes = boxes.copy()
boxes[:, 0::2] = width - boxes[:, 2::-2]
# landm
landms = landms.copy()
landms = landms.reshape([-1, 5, 2])
landms[:, :, 0] = width - landms[:, :, 0]
tmp = landms[:, 1, :].copy()
landms[:, 1, :] = landms[:, 0, :]
landms[:, 0, :] = tmp
tmp1 = landms[:, 4, :].copy()
landms[:, 4, :] = landms[:, 3, :]
landms[:, 3, :] = tmp1
landms = landms.reshape([-1, 10])
return image, boxes, landms
def _pad_to_square(image, rgb_mean, pad_image_flag):
if not pad_image_flag:
return image
height, width, _ = image.shape
long_side = max(width, height)
image_t = np.empty((long_side, long_side, 3), dtype=image.dtype)
image_t[:, :] = rgb_mean
image_t[0:0 + height, 0:0 + width] = image
return image_t
def _resize_subtract_mean(image, insize, rgb_mean):
interp_methods = [cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_NEAREST, cv2.INTER_LANCZOS4]
interp_method = interp_methods[random.randrange(5)]
image = cv2.resize(image, (insize, insize), interpolation=interp_method)
image = image.astype(np.float32)
image -= rgb_mean
return image.transpose(2, 0, 1)
class preproc(object):
def __init__(self, img_dim, rgb_means):
self.img_dim = img_dim
self.rgb_means = rgb_means
def __call__(self, image, targets):
assert targets.shape[0] > 0, "this image does not have gt"
boxes = targets[:, :4].copy()
labels = targets[:, -1].copy()
landm = targets[:, 4:-1].copy()
image_t, boxes_t, labels_t, landm_t, pad_image_flag = _crop(image, boxes, labels, landm, self.img_dim)
image_t = _distort(image_t)
image_t = _pad_to_square(image_t,self.rgb_means, pad_image_flag)
image_t, boxes_t, landm_t = _mirror(image_t, boxes_t, landm_t)
height, width, _ = image_t.shape
image_t = _resize_subtract_mean(image_t, self.img_dim, self.rgb_means)
boxes_t[:, 0::2] /= width
boxes_t[:, 1::2] /= height
landm_t[:, 0::2] /= width
landm_t[:, 1::2] /= height
labels_t = np.expand_dims(labels_t, 1)
targets_t = np.hstack((boxes_t, landm_t, labels_t))
return image_t, targets_t
project/hair_service_sd/models/layers/data/wider_face.py
+101
View File
@@ -0,0 +1,101 @@
import os
import os.path
import sys
import torch
import torch.utils.data as data
import cv2
import numpy as np
class WiderFaceDetection(data.Dataset):
def __init__(self, txt_path, preproc=None):
self.preproc = preproc
self.imgs_path = []
self.words = []
f = open(txt_path,'r')
lines = f.readlines()
isFirst = True
labels = []
for line in lines:
line = line.rstrip()
if line.startswith('#'):
if isFirst is True:
isFirst = False
else:
labels_copy = labels.copy()
self.words.append(labels_copy)
labels.clear()
path = line[2:]
path = txt_path.replace('label.txt','images/') + path
self.imgs_path.append(path)
else:
line = line.split(' ')
label = [float(x) for x in line]
labels.append(label)
self.words.append(labels)
def __len__(self):
return len(self.imgs_path)
def __getitem__(self, index):
img = cv2.imread(self.imgs_path[index])
height, width, _ = img.shape
labels = self.words[index]
annotations = np.zeros((0, 15))
if len(labels) == 0:
return annotations
for idx, label in enumerate(labels):
annotation = np.zeros((1, 15))
# bbox
annotation[0, 0] = label[0] # x1
annotation[0, 1] = label[1] # y1
annotation[0, 2] = label[0] + label[2] # x2
annotation[0, 3] = label[1] + label[3] # y2
# landmarks
annotation[0, 4] = label[4] # l0_x
annotation[0, 5] = label[5] # l0_y
annotation[0, 6] = label[7] # l1_x
annotation[0, 7] = label[8] # l1_y
annotation[0, 8] = label[10] # l2_x
annotation[0, 9] = label[11] # l2_y
annotation[0, 10] = label[13] # l3_x
annotation[0, 11] = label[14] # l3_y
annotation[0, 12] = label[16] # l4_x
annotation[0, 13] = label[17] # l4_y
if (annotation[0, 4]<0):
annotation[0, 14] = -1
else:
annotation[0, 14] = 1
annotations = np.append(annotations, annotation, axis=0)
target = np.array(annotations)
if self.preproc is not None:
img, target = self.preproc(img, target)
return torch.from_numpy(img), target
def detection_collate(batch):
"""Custom collate fn for dealing with batches of images that have a different
number of associated object annotations (bounding boxes).
Arguments:
batch: (tuple) A tuple of tensor images and lists of annotations
Return:
A tuple containing:
1) (tensor) batch of images stacked on their 0 dim
2) (list of tensors) annotations for a given image are stacked on 0 dim
"""
targets = []
imgs = []
for _, sample in enumerate(batch):
for _, tup in enumerate(sample):
if torch.is_tensor(tup):
imgs.append(tup)
elif isinstance(tup, type(np.empty(0))):
annos = torch.from_numpy(tup).float()
targets.append(annos)
return (torch.stack(imgs, 0), targets)
project/hair_service_sd/run_copy_cost_colorb64.py
+49 -18
View File
@@ -4,6 +4,14 @@ from uuid import uuid4
import imghdr import imghdr
import traceback import traceback
import torch import torch
# PyTorch 2.6+ 默认 weights_only=True，无法加载含 numpy 对象的旧权重
# （yolov5l.pt、各 *.pth 等）。统一改回 False（权重均为本机自有可信文件）。
_orig_torch_load = torch.load
def _torch_load(*args, **kwargs):
kwargs.setdefault("weights_only", False)
return _orig_torch_load(*args, **kwargs)
torch.load = _torch_load
from gevent import monkey from gevent import monkey
monkey.patch_all() monkey.patch_all()
@@ -247,7 +255,8 @@ def change_hair_colorv3():
jsonify({'msg': '算法解析错误', 'result': '', 'umd': '', 'state': -1}), 400) jsonify({'msg': '算法解析错误', 'result': '', 'umd': '', 'state': -1}), 400)
except Exception as e: except Exception as e:
print(e) print(f"[hairColor ERROR] {e}")
traceback.print_exc()
return make_response( return make_response(
jsonify({'msg': '算法解析错误', 'result': '', 'umd': '', 'state': -1}), 400) jsonify({'msg': '算法解析错误', 'result': '', 'umd': '', 'state': -1}), 400)
@@ -527,21 +536,41 @@ def change_hairstyle_v4():
crop_result = landmark_processor.high_quality_warpAffine(img_res, M, dst_size) crop_result = landmark_processor.high_quality_warpAffine(img_res, M, dst_size)
# crop_result = landmark_processor.high_quality_warpAffine(img_res, M, (w, h)) # crop_result = landmark_processor.high_quality_warpAffine(img_res, M, (w, h))
origin_matting = cv2.imread(user_orig_mask_path, cv2.IMREAD_GRAYSCALE) # 接口11：可选外部遮罩 ext_mask（接口9 头发遮罩，原图坐标、白=生发区）。
new_matting = cv2.imread(hair_matting_path, cv2.IMREAD_GRAYSCALE) # 传入则经 M 变换到 crop 坐标后替换内部 matting 合成遮罩，webui 精确重绘该区域；
left_point = user_landmarks_origin_img_137[14] # (x1, y1) 获取刘海区域的圆心和半径 # 不传则保持原 matting_merge 逻辑（向后兼容，其它功能不受影响）。
right_point = user_landmarks_origin_img_137[7] # (x2, y2) ext_mask_b64 = input_info.get('ext_mask', '')
center_x = int((left_point[0] + right_point[0]) // 2) if ext_mask_b64:
center_y = int((left_point[1] + right_point[1]) // 2) if 'data:image/' in ext_mask_b64:
radius = int(np.sqrt((right_point[0] - left_point[0]) ** 2 + (right_point[1] - left_point[1]) ** 2) / 2) ext_mask_b64 = ext_mask_b64.split(',')[1]
h, w = origin_matting.shape # 初始化刘海区域为空（全黑色） ext_mask_np = cv2.imdecode(
part_bangs = np.zeros((h, w), dtype=np.uint8) np.frombuffer(base64.b64decode(ext_mask_b64), np.uint8),
cv2.circle(part_bangs, (center_x, center_y), radius, 255, -1) # 在 part_bangs 上绘制圆作为刘海区域 圆形区域设为白色 cv2.IMREAD_GRAYSCALE)
no_bang_result = cv2.subtract(origin_matting, part_bangs) # 去除刘海区域 if ext_mask_np is None:
matting_merge = np.max( np.stack([no_bang_result, new_matting], axis=2), axis=2).astype(np.uint8)# 合并 origin_matting 和 new_matting raise RuntimeError("ext_mask 解析失败")
crop_matting = cv2.warpAffine(matting_merge, M, dst_size) if ext_mask_np.shape[:2] != origin_img.shape[:2]:
# crop_matting = cv2.warpAffine(matting_merge, M, (w, h)) ext_mask_np = cv2.resize(
mask = (crop_matting > 10).astype(np.float32) ext_mask_np, (origin_img.shape[1], origin_img.shape[0]),
interpolation=cv2.INTER_NEAREST)
print(f"[swapHair] 使用外部遮罩 ext_mask, shape={ext_mask_np.shape}, "
f"白像素={int((ext_mask_np > 10).sum())}", flush=True)
crop_matting = cv2.warpAffine(ext_mask_np, M, dst_size)
mask = (crop_matting > 10).astype(np.float32)
else:
origin_matting = cv2.imread(user_orig_mask_path, cv2.IMREAD_GRAYSCALE)
new_matting = cv2.imread(hair_matting_path, cv2.IMREAD_GRAYSCALE)
left_point = user_landmarks_origin_img_137[14] # (x1, y1) 获取刘海区域的圆心和半径
right_point = user_landmarks_origin_img_137[7] # (x2, y2)
center_x = int((left_point[0] + right_point[0]) // 2)
center_y = int((left_point[1] + right_point[1]) // 2)
radius = int(np.sqrt((right_point[0] - left_point[0]) ** 2 + (right_point[1] - left_point[1]) ** 2) / 2)
h, w = origin_matting.shape # 初始化刘海区域为空（全黑色）
part_bangs = np.zeros((h, w), dtype=np.uint8)
cv2.circle(part_bangs, (center_x, center_y), radius, 255, -1) # 在 part_bangs 上绘制圆作为刘海区域 圆形区域设为白色
no_bang_result = cv2.subtract(origin_matting, part_bangs) # 去除刘海区域
matting_merge = np.max( np.stack([no_bang_result, new_matting], axis=2), axis=2).astype(np.uint8)# 合并 origin_matting 和 new_matting
crop_matting = cv2.warpAffine(matting_merge, M, dst_size)
mask = (crop_matting > 10).astype(np.float32)
if not is_hr: if not is_hr:
mask_dilate = cv2.dilate(mask, np.ones((3, 9), np.uint8)) mask_dilate = cv2.dilate(mask, np.ones((3, 9), np.uint8))
else: else:
@@ -576,7 +605,9 @@ def change_hairstyle_v4():
p_tag = p_tag[p_tag.find("simple background, ") + len("simple background, "):] p_tag = p_tag[p_tag.find("simple background, ") + len("simple background, "):]
else: else:
p_tag = "" p_tag = ""
denoising_strength = 0.6 denoising_strength = float(input_info.get('denoising_strength', 0.6)) # 接口11 可调；默认 0.6
webui_steps_raw = input_info.get('webui_steps', None) # 可选：webui img2img 步数，None用服务端默认
webui_steps = int(webui_steps_raw) if webui_steps_raw is not None else None
print(f"功能8：处理发型区域，耗时：{time.time() - start_time:.3f}s") print(f"功能8：处理发型区域，耗时：{time.time() - start_time:.3f}s")
# cv2.imwrite(f"{task_id}_mask_dilate.jpg", mask_dilate) # cv2.imwrite(f"{task_id}_mask_dilate.jpg", mask_dilate)
# cv2.imwrite(f"{task_id}_final_img.jpg", final_img) # cv2.imwrite(f"{task_id}_final_img.jpg", final_img)
@@ -590,7 +621,7 @@ def change_hairstyle_v4():
# mask_dilate = cv2.imread("/root/project/hair_service_sd/gt_mask_dilate.jpg") # mask_dilate = cv2.imread("/root/project/hair_service_sd/gt_mask_dilate.jpg")
sd_result = webui_img2img(img=final_img, mask_img=mask_dilate, in_gender=in_gender, task_id=task_id, sd_result = webui_img2img(img=final_img, mask_img=mask_dilate, in_gender=in_gender, task_id=task_id,
hair_id=hair_id, lora_material_path=hair_material_dir, tag=p_tag, is_hr=is_hr, hair_id=hair_id, lora_material_path=hair_material_dir, tag=p_tag, is_hr=is_hr,
denoising_strength=denoising_strength, inference_port="57860") denoising_strength=denoising_strength, inference_port="57860", webui_steps=webui_steps)
print(f"功能：webui，耗时：{time.time() - start_time:.3f}s") print(f"功能：webui，耗时：{time.time() - start_time:.3f}s")
# 功能：后处理并上传结果 # 功能：后处理并上传结果
project/hair_service_sd/upload_oss.py
+17 -20
View File
@@ -1,31 +1,28 @@
import time import time
import os
try: import oss2
import oss2 import os
_OSS2_AVAILABLE = True
except Exception as _e:
print(f'[OSS upload_oss] ⚠️ oss2 导入失败({_e})，OSS 上传功能不可用')
_OSS2_AVAILABLE = False
class OSS_object(): class OSS_object():
def __init__(self): def __init__(self):
if not _OSS2_AVAILABLE: # 懒加载：只读取凭证，不立即连接 OSS。
print('[OSS upload_oss] oss2 不可用，OSS 功能关闭') # 本地用 output_format=base64 时不配 OSS 凭证也能启动服务；
self.bucket = None # 仅在真正调用 upload_file 时才校验并连接。
self.access_key_id = os.getenv('OSS_TEST_ACCESS_KEY_ID', '<your-access-key-id>')
self.access_key_secret = os.getenv('OSS_TEST_ACCESS_KEY_SECRET', '<your-access-key-secret>')
self.bucket_name = os.getenv('OSS_TEST_BUCKET', '<your-bucket-name>')
self.endpoint = os.getenv('OSS_TEST_ENDPOINT', '<your-endpoint>')
self.bucket = None
def _ensure_bucket(self):
if self.bucket is not None:
return return
access_key_id = os.getenv('OSS_TEST_ACCESS_KEY_ID', '<your-access-key-id>') for param in (self.access_key_id, self.access_key_secret, self.bucket_name, self.endpoint):
access_key_secret = os.getenv('OSS_TEST_ACCESS_KEY_SECRET', '<your-access-key-secret>') assert '<' not in param, '请设置环境变量 OSS_TEST_ACCESS_KEY_ID / OSS_TEST_ACCESS_KEY_SECRET / OSS_TEST_BUCKET / OSS_TEST_ENDPOINT'
bucket_name = os.getenv('OSS_TEST_BUCKET', '<your-bucket-name>') self.bucket = oss2.Bucket(oss2.Auth(self.access_key_id, self.access_key_secret), self.endpoint, self.bucket_name)
endpoint = os.getenv('OSS_TEST_ENDPOINT', '<your-endpoint>')
# 密钥未设置时标记不可用（不抛异常）
if '<' in access_key_id or '<' in access_key_secret:
print('[OSS upload_oss] ⚠️ 未设置 OSS 密钥环境变量，OSS 上传功能不可用')
self.bucket = None
return
self.bucket = oss2.Bucket(oss2.Auth(access_key_id, access_key_secret), endpoint, bucket_name)
def upload_file(self, file, target_name): def upload_file(self, file, target_name):
self._ensure_bucket()
t0 = time.time() t0 = time.time()
with open(oss2.to_unicode(file), 'rb') as f: with open(oss2.to_unicode(file), 'rb') as f:
ret = self.bucket.put_object(target_name, f) ret = self.bucket.put_object(target_name, f)
project/hair_service_sd/utils/enhance_hair.py
+3
View File
@@ -43,6 +43,9 @@ def webui_img2img(img, mask, prompt='', denoising_strength=0.35):
} }
response = requests.post(url=url, json=request_dict) response = requests.post(url=url, json=request_dict)
ret_json = response.json() ret_json = response.json()
if 'images' not in ret_json:
print(f"[webui_img2img ERROR] {ret_json}")
raise Exception(f"webui error: {ret_json.get('error', 'unknown')}")
result = ret_json['images'][0] result = ret_json['images'][0]
img = cv2.imdecode(np.frombuffer(base64.b64decode(result.split(",", 1)[0]), np.uint8), cv2.IMREAD_COLOR) img = cv2.imdecode(np.frombuffer(base64.b64decode(result.split(",", 1)[0]), np.uint8), cv2.IMREAD_COLOR)
return img return img
project/photo_service/lora_train_service_1.py
+6 -6
View File
@@ -37,8 +37,8 @@ if version == "local":
callback_url = 'http://service.aicloud.fit:7395/api/hair/trainCallBack' callback_url = 'http://service.aicloud.fit:7395/api/hair/trainCallBack'
else: else:
current_url = 'http://0.0.0.0:7393/' current_url = 'http://0.0.0.0:7393/'
kohya_ss_home_dir = '/home/xsl/change_hair/project/kohya_ss_home' kohya_ss_home_dir = '/home/ubuntu/change_hair/project/kohya_ss_home'
webui_lora_dir = '/home/xsl/change_hair/project/onediff/stable-diffusion-webui/models/Lora' webui_lora_dir = '/home/ubuntu/change_hair/project/onediff/stable-diffusion-webui/models/Lora'
inference_use_onediff = False inference_use_onediff = False
callback_url = 'http://0.0.0.0:8801/api/hair/trainCallBack' callback_url = 'http://0.0.0.0:8801/api/hair/trainCallBack'
base_webui_port = '57860' base_webui_port = '57860'
@@ -307,13 +307,13 @@ def train_thread(sq, gpu_id):
# 3. GPU 固定 device=0(单卡) # 3. GPU 固定 device=0(单卡)
# 4. 去掉 tokenizer_cache_dir(改用 HF 本地缓存 + 离线模式) # 4. 去掉 tokenizer_cache_dir(改用 HF 本地缓存 + 离线模式)
# 5. 设置 HF_HUB_OFFLINE 避免联网检查 # 5. 设置 HF_HUB_OFFLINE 避免联网检查
kohya_python = '/home/xsl/miniconda3/envs/kohya/bin/python' kohya_python = '/home/ubuntu/miniconda3/envs/my_hair/bin/python'
kohya_workdir = os.path.join(kohya_ss_home_dir, 'kohya_ss') kohya_workdir = os.path.join(kohya_ss_home_dir, 'kohya_ss')
base_model = '/home/xsl/change_hair/project/onediff/stable-diffusion-webui/models/Stable-diffusion/v1-5-pruned-emaonly.safetensors' base_model = '/home/ubuntu/change_hair/project/onediff/stable-diffusion-webui/models/Stable-diffusion/v1-5-pruned-emaonly.safetensors'
cmd_train = ( cmd_train = (
f'cd {kohya_workdir} && ' f'cd {kohya_workdir} && '
f'HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 CUDA_VISIBLE_DEVICES={device_id} ' f'HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 CUDA_VISIBLE_DEVICES={device_id} '
f'/home/xsl/miniconda3/envs/kohya/bin/accelerate launch --num_cpu_threads_per_process=2 "./train_network.py" --enable_bucket ' f'/home/ubuntu/miniconda3/envs/my_hair/bin/accelerate launch --num_cpu_threads_per_process=2 "./train_network.py" --enable_bucket '
f'--min_bucket_reso=256 --max_bucket_reso=2048 --pretrained_model_name_or_path="{base_model}" ' f'--min_bucket_reso=256 --max_bucket_reso=2048 --pretrained_model_name_or_path="{base_model}" '
f'--train_data_dir={images_dir} --resolution="2000,2000" ' f'--train_data_dir={images_dir} --resolution="2000,2000" '
f'--output_dir={model_dir} ' f'--output_dir={model_dir} '
@@ -324,7 +324,7 @@ def train_thread(sq, gpu_id):
'--caption_extension=".txt" --sample_sampler=ddim ' '--caption_extension=".txt" --sample_sampler=ddim '
f'--sample_prompts={sample_txt} --sample_every_n_epochs="1000" ' f'--sample_prompts={sample_txt} --sample_every_n_epochs="1000" '
'--seed="1234" --cache_latents --optimizer_type="AdamW" --max_data_loader_n_workers="0" --bucket_reso_steps=64 ' '--seed="1234" --cache_latents --optimizer_type="AdamW" --max_data_loader_n_workers="0" --bucket_reso_steps=64 '
'--xformers --bucket_no_upscale --noise_offset=0.0') '--sdpa --bucket_no_upscale --noise_offset=0.0')
print("cmd_train:", cmd_train) print("cmd_train:", cmd_train)
os.system(cmd_train) os.system(cmd_train)
scripts/sync_data_to_server.sh
Executable
+69
View File
@@ -0,0 +1,69 @@
#!/usr/bin/env bash
# 断点续传大文件到云服务器 /home/xsl/data
# 用法: bash scripts/sync_data_to_server.sh [all|weights|sd|hairstyles]
set -euo pipefail
REMOTE="ubuntu@117.50.213.111"
REMOTE_BASE="/home/xsl/data"
LOCAL_BASE="/home/xsl/change_hair"
LOG_DIR="${LOCAL_BASE}/project/logs"
mkdir -p "$LOG_DIR"
RSYNC_OPTS=(-avP --partial --info=progress2)
sync_weights() {
echo ">>> [1/3] 同步 hair_service_sd/weights (~10G) ..."
rsync "${RSYNC_OPTS[@]}" \
"${LOCAL_BASE}/project/hair_service_sd/weights/" \
"${REMOTE}:${REMOTE_BASE}/hair_service_sd/weights/"
}
sync_sd_models() {
echo ">>> [2/3] 同步 SD 底模 + ESRGAN + GFPGAN (~4.6G) ..."
rsync "${RSYNC_OPTS[@]}" \
"${LOCAL_BASE}/project/onediff/stable-diffusion-webui/models/Stable-diffusion/v1-5-pruned-emaonly.safetensors" \
"${REMOTE}:${REMOTE_BASE}/sdwebui/models/Stable-diffusion/"
rsync "${RSYNC_OPTS[@]}" \
"${LOCAL_BASE}/project/onediff/stable-diffusion-webui/models/ESRGAN/" \
"${REMOTE}:${REMOTE_BASE}/sdwebui/models/ESRGAN/"
rsync "${RSYNC_OPTS[@]}" \
"${LOCAL_BASE}/project/onediff/stable-diffusion-webui/models/GFPGAN/" \
"${REMOTE}:${REMOTE_BASE}/sdwebui/models/GFPGAN/"
}
sync_hairstyles() {
echo ">>> [3/3] 同步 5 个 chang_* 发型资源 (~0.8G) ..."
for id in chang_tuoyuan chang_bolang chang_zhixian chang_huaban chang_xinxing; do
echo " - $id"
rsync "${RSYNC_OPTS[@]}" \
"${LOCAL_BASE}/project/data/ref_hairstyle/${id}/" \
"${REMOTE}:${REMOTE_BASE}/project/data/ref_hairstyle/${id}/"
rsync "${RSYNC_OPTS[@]}" \
"${LOCAL_BASE}/project/data/hair_template_material/${id}/" \
"${REMOTE}:${REMOTE_BASE}/project/data/hair_template_material/${id}/"
rsync "${RSYNC_OPTS[@]}" \
"${LOCAL_BASE}/project/data/upload_train_imgs/${id}/" \
"${REMOTE}:${REMOTE_BASE}/project/data/upload_train_imgs/${id}/"
rsync "${RSYNC_OPTS[@]}" \
"${LOCAL_BASE}/data/train_material/${id}/" \
"${REMOTE}:${REMOTE_BASE}/train_material/${id}/"
done
}
TARGET="${1:-all}"
case "$TARGET" in
weights) sync_weights ;;
sd) sync_sd_models ;;
hairstyles) sync_hairstyles ;;
all)
sync_weights
sync_sd_models
sync_hairstyles
;;
*)
echo "用法: $0 [all|weights|sd|hairstyles]"
exit 1
;;
esac
echo ">>> 同步完成: ${TARGET}"
start_all.sh
+4 -4
View File
@@ -10,16 +10,16 @@
# bash start_all.sh status # 查看状态 # bash start_all.sh status # 查看状态
# ===================================================================== # =====================================================================
set -u set -u
BASE="/home/xsl/change_hair" BASE="/home/ubuntu/change_hair"
PROJ="$BASE/project" PROJ="$BASE/project"
LOGDIR="$PROJ/logs" LOGDIR="$PROJ/logs"
PIDD="$LOGDIR/pids" PIDD="$LOGDIR/pids"
mkdir -p "$LOGDIR" "$PIDD" mkdir -p "$LOGDIR" "$PIDD"
# conda 环境的 python 路径 # conda 环境的 python 路径
PY_HAIR="/home/xsl/miniconda3/envs/my_hair/bin/python" PY_HAIR="/home/ubuntu/miniconda3/envs/my_hair/bin/python"
PY_SD="/home/xsl/miniconda3/envs/sdwebui/bin/python" PY_SD="/home/ubuntu/miniconda3/envs/sdwebui/bin/python"
PY_PHOTO="/home/xsl/miniconda3/envs/py310/bin/python" PY_PHOTO="/home/ubuntu/miniconda3/envs/py310/bin/python"
# 公共环境变量 # 公共环境变量
export CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1 # 旧版 cryptography 兼容 export CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1 # 旧版 cryptography 兼容
start_hair.sh
+2 -2
View File
@@ -2,5 +2,5 @@
export CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1 export CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1
export CUDA_VISIBLE_DEVICES=0 export CUDA_VISIBLE_DEVICES=0
export APP_WORKER_ID=1 export APP_WORKER_ID=1
cd /home/xsl/change_hair/project/hair_service_sd cd /home/ubuntu/change_hair/project/hair_service_sd
exec /home/xsl/miniconda3/envs/my_hair/bin/python run_copy_cost_colorb64.py exec /home/ubuntu/miniconda3/envs/my_hair/bin/python run_copy_cost_colorb64.py
start_hairgrow.sh
+2 -2
View File
@@ -5,5 +5,5 @@ export CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1
export CUDA_VISIBLE_DEVICES=0 export CUDA_VISIBLE_DEVICES=0
export HF_HUB_OFFLINE=1 export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1 export TRANSFORMERS_OFFLINE=1
cd /home/xsl/change_hair/hair_grow_service cd /home/ubuntu/change_hair/hair_grow_service
exec /home/xsl/miniconda3/envs/my_hair/bin/python app.py exec /home/ubuntu/miniconda3/envs/my_hair/bin/python app.py
start_photo.sh
+2 -2
View File
@@ -1,5 +1,5 @@
#!/bin/bash #!/bin/bash
# 独立启动 photo_service，确保脱离会话 # 独立启动 photo_service，确保脱离会话
export CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1 export CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1
cd /home/xsl/change_hair/project/photo_service cd /home/ubuntu/change_hair/project/photo_service
exec /home/xsl/miniconda3/envs/py310/bin/python -u lora_train_service_1.py exec /home/ubuntu/miniconda3/envs/py310/bin/python -u lora_train_service_1.py
sync_all.sh
+10 -10
View File
@@ -4,7 +4,7 @@
# 已传数据自动续传（--partial） # 已传数据自动续传（--partial）
set -u set -u
LOGDIR="/home/xsl/change_hair/project/logs" LOGDIR="/home/ubuntu/change_hair/project/logs"
DONE_MARKER="$LOGDIR/sync_all.done" DONE_MARKER="$LOGDIR/sync_all.done"
mkdir -p "$LOGDIR" mkdir -p "$LOGDIR"
@@ -14,15 +14,15 @@ mkdir -p "$LOGDIR"
# 任务定义：名称|源|目标 # 任务定义：名称|源|目标
# 优先级排序：换发色关键链路优先 → 换发型 → 训练相关 → 大数据最后 # 优先级排序：换发色关键链路优先 → 换发型 → 训练相关 → 大数据最后
TASKS=( TASKS=(
"photo_service|szlc@192.168.101.63:/home/szlc/project/photo_service/|/home/xsl/change_hair/project/photo_service/" "photo_service|szlc@192.168.101.63:/home/szlc/project/photo_service/|/home/ubuntu/change_hair/project/photo_service/"
"project_data|szlc@192.168.101.63:/home/szlc/project/data/|/home/xsl/change_hair/project/data/" "project_data|szlc@192.168.101.63:/home/szlc/project/data/|/home/ubuntu/change_hair/project/data/"
"hair_service_sd|szlc@192.168.101.63:/home/szlc/project/hair_service_sd/|/home/xsl/change_hair/project/hair_service_sd/" "hair_service_sd|szlc@192.168.101.63:/home/szlc/project/hair_service_sd/|/home/ubuntu/change_hair/project/hair_service_sd/"
"conda_py310|szlc@192.168.101.63:/home/szlc/miniconda3/envs/py310/|/home/xsl/miniconda3/envs/py310/" "conda_py310|szlc@192.168.101.63:/home/szlc/miniconda3/envs/py310/|/home/ubuntu/miniconda3/envs/py310/"
"conda_my_hair|szlc@192.168.101.63:/home/szlc/miniconda3/envs/my_hair/|/home/xsl/miniconda3/envs/my_hair/" "conda_my_hair|szlc@192.168.101.63:/home/szlc/miniconda3/envs/my_hair/|/home/ubuntu/miniconda3/envs/my_hair/"
"conda_sdwebui|szlc@192.168.101.63:/home/szlc/miniconda3/envs/sdwebui/|/home/xsl/miniconda3/envs/sdwebui/" "conda_sdwebui|szlc@192.168.101.63:/home/szlc/miniconda3/envs/sdwebui/|/home/ubuntu/miniconda3/envs/sdwebui/"
"kohya_ss_home|szlc@192.168.101.63:/home/szlc/project/kohya_ss_home/|/home/xsl/change_hair/project/kohya_ss_home/" "kohya_ss_home|szlc@192.168.101.63:/home/szlc/project/kohya_ss_home/|/home/ubuntu/change_hair/project/kohya_ss_home/"
"onediff|szlc@192.168.101.63:/home/szlc/project/onediff/|/home/xsl/change_hair/project/onediff/" "onediff|szlc@192.168.101.63:/home/szlc/project/onediff/|/home/ubuntu/change_hair/project/onediff/"
"train_material|szlc@192.168.101.63:/data/train_material/|/home/xsl/change_hair/data/train_material/" "train_material|szlc@192.168.101.63:/data/train_material/|/home/ubuntu/change_hair/data/train_material/"
) )
echo "========================================" | tee -a "$LOGDIR/sync_main.log" echo "========================================" | tee -a "$LOGDIR/sync_main.log"
sync_conda.sh
+2 -2
View File
@@ -1,8 +1,8 @@
#!/bin/bash #!/bin/bash
# 重新同步 conda 环境（上次被 sed 损坏，这次纯净拷贝） # 重新同步 conda 环境（上次被 sed 损坏，这次纯净拷贝）
set -u set -u
CONDA_ENVS="/home/xsl/miniconda3/envs" CONDA_ENVS="/home/ubuntu/miniconda3/envs"
LOGDIR="/home/xsl/change_hair/project/logs" LOGDIR="/home/ubuntu/change_hair/project/logs"
mkdir -p "$CONDA_ENVS" "$LOGDIR" mkdir -p "$CONDA_ENVS" "$LOGDIR"
echo "conda 环境重传开始: $(date '+%H:%M:%S')" | tee -a "$LOGDIR/sync_conda.log" echo "conda 环境重传开始: $(date '+%H:%M:%S')" | tee -a "$LOGDIR/sync_conda.log"
test_grow_swap.py
+3 -3
View File
@@ -2,8 +2,8 @@
"""端到端测试：生发走换发型工作流""" """端到端测试：生发走换发型工作流"""
import base64, requests, time, sys import base64, requests, time, sys
img_path = "/home/xsl/change_hair/project/data/userImage/8488902485_20250630055547.jpg" img_path = "/home/ubuntu/change_hair/project/data/userImage/8488902485_20250630055547.jpg"
mask_path = "/home/xsl/change_hair/project/logs/hairgrow_test_mask.png" mask_path = "/home/ubuntu/change_hair/project/logs/hairgrow_test_mask.png"
with open(img_path, "rb") as f: with open(img_path, "rb") as f:
img_b64 = "data:image/jpeg;base64," + base64.b64encode(f.read()).decode() img_b64 = "data:image/jpeg;base64," + base64.b64encode(f.read()).decode()
@@ -24,7 +24,7 @@ try:
d = r.json() d = r.json()
print(f"HTTP {r.status_code} | state={d.get('state')} | msg={d.get('msg','')} | 耗时={time.time()-t0:.1f}s") print(f"HTTP {r.status_code} | state={d.get('state')} | msg={d.get('msg','')} | 耗时={time.time()-t0:.1f}s")
if d.get("state") == 0: if d.get("state") == 0:
out = "/home/xsl/change_hair/project/logs/test_grow_swap_result.jpg" out = "/home/ubuntu/change_hair/project/logs/test_grow_swap_result.jpg"
with open(out, "wb") as f: with open(out, "wb") as f:
f.write(base64.b64decode(d["result"])) f.write(base64.b64decode(d["result"]))
print(f"✅ 生发成功! 结果图: {out}") print(f"✅ 生发成功! 结果图: {out}")
test_haircolor.py
+2 -2
View File
@@ -2,7 +2,7 @@
"""换发色接口测试""" """换发色接口测试"""
import base64, json, requests, sys import base64, json, requests, sys
img_path = "/home/xsl/change_hair/project/data/userImage/8488902485_20250630055547.jpg" img_path = "/home/ubuntu/change_hair/project/data/userImage/8488902485_20250630055547.jpg"
with open(img_path, "rb") as f: with open(img_path, "rb") as f:
img_b64 = base64.b64encode(f.read()).decode() img_b64 = base64.b64encode(f.read()).decode()
@@ -21,7 +21,7 @@ try:
result = d.get("result", "") result = d.get("result", "")
if d.get("state") == 0 and result: if d.get("state") == 0 and result:
# 保存结果图 # 保存结果图
out = "/home/xsl/change_hair/project/logs/test_haircolor_result.jpg" out = "/home/ubuntu/change_hair/project/logs/test_haircolor_result.jpg"
with open(out, "wb") as f: with open(out, "wb") as f:
f.write(base64.b64decode(result)) f.write(base64.b64decode(result))
print(f"✅ 换发色成功! 结果图已保存: {out} ({len(result)} bytes base64)") print(f"✅ 换发色成功! 结果图已保存: {out} ({len(result)} bytes base64)")
test_hairgrow_api.py
+3 -3
View File
@@ -2,8 +2,8 @@
"""测试 /api/hairGrow/v1 接口""" """测试 /api/hairGrow/v1 接口"""
import base64, requests, sys import base64, requests, sys
img_path = "/home/xsl/change_hair/project/data/userImage/8488902485_20250630055547.jpg" img_path = "/home/ubuntu/change_hair/project/data/userImage/8488902485_20250630055547.jpg"
mask_path = "/home/xsl/change_hair/project/logs/hairgrow_test_mask.png" mask_path = "/home/ubuntu/change_hair/project/logs/hairgrow_test_mask.png"
with open(img_path, "rb") as f: with open(img_path, "rb") as f:
img_b64 = "data:image/jpeg;base64," + base64.b64encode(f.read()).decode() img_b64 = "data:image/jpeg;base64," + base64.b64encode(f.read()).decode()
@@ -26,7 +26,7 @@ try:
print(f"HTTP {r.status_code} | msg={d.get('msg')} | state={d.get('state')} | 耗时={time.time()-t0:.1f}s") print(f"HTTP {r.status_code} | msg={d.get('msg')} | state={d.get('state')} | 耗时={time.time()-t0:.1f}s")
result = d.get("result", "") result = d.get("result", "")
if d.get("state") == 0 and result: if d.get("state") == 0 and result:
out = "/home/xsl/change_hair/project/logs/test_hairgrow_api_result.jpg" out = "/home/ubuntu/change_hair/project/logs/test_hairgrow_api_result.jpg"
with open(out, "wb") as f: with open(out, "wb") as f:
f.write(base64.b64decode(result)) f.write(base64.b64decode(result))
print(f"✅ 生发成功! 结果图: {out}") print(f"✅ 生发成功! 结果图: {out}")
test_swaphair.py
+2 -2
View File
@@ -2,7 +2,7 @@
"""换发型接口测试""" """换发型接口测试"""
import base64, json, requests, sys, time import base64, json, requests, sys, time
img_path = "/home/xsl/change_hair/project/data/userImage/8488902485_20250630055547.jpg" img_path = "/home/ubuntu/change_hair/project/data/userImage/8488902485_20250630055547.jpg"
with open(img_path, "rb") as f: with open(img_path, "rb") as f:
img_b64 = base64.b64encode(f.read()).decode() img_b64 = base64.b64encode(f.read()).decode()
@@ -22,7 +22,7 @@ try:
print(f"HTTP {r.status_code} | msg={d.get('msg')} | state={d.get('state')} | 耗时={time.time()-t0:.1f}s") print(f"HTTP {r.status_code} | msg={d.get('msg')} | state={d.get('state')} | 耗时={time.time()-t0:.1f}s")
data = d.get("data", "") data = d.get("data", "")
if d.get("state") == 0 and data: if d.get("state") == 0 and data:
out = "/home/xsl/change_hair/project/logs/test_swaphair_result.jpg" out = "/home/ubuntu/change_hair/project/logs/test_swaphair_result.jpg"
with open(out, "wb") as f: with open(out, "wb") as f:
f.write(base64.b64decode(data)) f.write(base64.b64decode(data))
print(f"✅ 换发型成功! 结果图已保存: {out} ({len(data)} bytes base64)") print(f"✅ 换发型成功! 结果图已保存: {out} ({len(data)} bytes base64)")
test_train_api.py
+1 -1
View File
@@ -4,7 +4,7 @@ import requests, json
# 用一个已有完整训练素材的发型 # 用一个已有完整训练素材的发型
HAIR_ID = "1905785164224868354" HAIR_ID = "1905785164224868354"
MATERIAL_DIR = f"/home/xsl/change_hair/data/train_material/{HAIR_ID}" MATERIAL_DIR = f"/home/ubuntu/change_hair/data/train_material/{HAIR_ID}"
payload = { payload = {
"task_id": f"test_train_{HAIR_ID}", "task_id": f"test_train_{HAIR_ID}",
train_batch_stepC.py
+96
View File
@@ -0,0 +1,96 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""阶段C：为已训练的发型生成 ref材质 + 模板材质 + 预览图
前置：LoRA 已训练完成，webui/photo_service/hair_service_sd 已启动。
对每个发型依次：
step3: 调 hair_service_sd trainCallBack 回调生成 ref 材质
step4: 生成模板材质
step5: 生成预览图
"""
import os
import sys
import time
import json
import shutil
HAIR_SERVICE_DIR = "/home/ubuntu/change_hair/project/hair_service_sd"
os.chdir(HAIR_SERVICE_DIR)
sys.path.insert(0, HAIR_SERVICE_DIR)
os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
os.environ.setdefault("CRYPTOGRAPHY_OPENSSL_NO_LEGACY", "1")
import requests as req
# (hair_id, gender, template_img)
HAIRSTYLES = [
("chang_tuoyuan", "girl", "/home/ubuntu/change_hair/hair_type_images/chang_tuoyuan/chang_tuoyuan.jpg"),
("chang_bolang", "girl", "/home/ubuntu/change_hair/hair_type_images/chang_bolang/chang_bolang.jpg"),
("chang_zhixian", "girl", "/home/ubuntu/change_hair/hair_type_images/chang_zhixian/chang_zhixian.jpg"),
("chang_huaban", "girl", "/home/ubuntu/change_hair/hair_type_images/chang_huaban/chang_huaban.jpg"),
("chang_xinxing", "girl", "/home/ubuntu/change_hair/hair_type_images/chang_xinxing/chang_xinxing.jpg"),
]
HAIR_CALLBACK = "http://127.0.0.1:8801/api/hair/trainCallBack"
UPLOAD_TRAIN_DIR = "/home/ubuntu/change_hair/project/data/upload_train_imgs"
# 导入 step4/step5 函数（会触发模型加载）
from train_hairstyle_full import step4_template_material, step5_preview
from common.logger import config
def log(msg):
print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)
def step3_callback(hair_id, template_img):
"""调 trainCallBack 生成 ref 材质。需要先把模板图放到 upload_train_imgs"""
log(f" --- step3: {hair_id} 生成 ref 材质 ---")
# 准备 upload_train_imgs/<hair_id>/first##<hair_id>.jpg
src_dir = os.path.join(UPLOAD_TRAIN_DIR, hair_id)
os.makedirs(src_dir, exist_ok=True)
dst = os.path.join(src_dir, f"first##{hair_id}.jpg")
shutil.copy(template_img, dst)
try:
r = req.post(HAIR_CALLBACK, json={
"task_id": f"train_{hair_id}", "hair_id": hair_id,
"state": 0, "msg": "头发lora训练成功", "is_tj": "1"
}, timeout=300)
d = r.json()
ref_dir = os.path.join(config.get('default', 'hairstyleDir'), hair_id)
ok = d.get("state") == 0 and os.path.exists(os.path.join(ref_dir, "ref_rgb_8uc3_768.png"))
log(f" {'✓' if ok else '✗'} {hair_id} ref材质 {'成功' if ok else '失败: ' + str(d.get('msg'))}")
return ok
except Exception as e:
log(f" ✗ {hair_id} ref材质异常: {e}")
return False
def main():
t0 = time.time()
log(f"阶段C: 为 {len(HAIRSTYLES)} 个发型生成材质和预览")
results = {}
for hair_id, gender, template_img in HAIRSTYLES:
log(f"\n{'='*50}")
log(f"处理: {hair_id}")
log(f"{'='*50}")
ok3 = step3_callback(hair_id, template_img)
ok4 = step4_template_material(hair_id, template_img)
ok5 = step5_preview(hair_id, gender)
results[hair_id] = (ok3, ok4, ok5)
log(f" {hair_id}: ref={'✓' if ok3 else '✗'} 模板={'✓' if ok4 else '✗'} 预览={'✓' if ok5 else '✗'}")
log(f"\n{'='*60}")
log(f"阶段C完成，耗时 {time.time()-t0:.0f}s")
log(f"{'='*60}")
for hair_id, (ok3, ok4, ok5) in results.items():
status = "✓全部成功" if (ok3 and ok4 and ok5) else f"⚠️ ref={'✓' if ok3 else '✗'} 模板={'✓' if ok4 else '✗'} 预览={'✓' if ok5 else '✗'}"
log(f" {hair_id}: {status}")
if __name__ == "__main__":
main()
train_hairstyle_full.py
+10 -6
View File
@@ -6,8 +6,8 @@
串接 prepare_train_data / gen_template_material / photo_service训练 / 回调 / 预览图 串接 prepare_train_data / gen_template_material / photo_service训练 / 回调 / 预览图
用法: 用法:
cd /home/xsl/change_hair/project/hair_service_sd cd /home/ubuntu/change_hair/project/hair_service_sd
python /home/xsl/change_hair/train_hairstyle_full.py --hair-id huaban1 --input /home/xsl/change_hair/huaban1 --gender girl --template-img /home/xsl/change_hair/huaban1/huaban1.jpg python /home/ubuntu/change_hair/train_hairstyle_full.py --hair-id huaban1 --input /home/ubuntu/change_hair/huaban1 --gender girl --template-img /home/ubuntu/change_hair/huaban1/huaban1.jpg
""" """
import os import os
import sys import sys
@@ -21,7 +21,7 @@ import argparse
import requests as req import requests as req
ssl._create_default_https_context = ssl._create_unverified_context ssl._create_default_https_context = ssl._create_unverified_context
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd" HAIR_SERVICE_DIR = "/home/ubuntu/change_hair/project/hair_service_sd"
os.chdir(HAIR_SERVICE_DIR) os.chdir(HAIR_SERVICE_DIR)
sys.path.insert(0, HAIR_SERVICE_DIR) sys.path.insert(0, HAIR_SERVICE_DIR)
os.environ.setdefault("HF_HUB_OFFLINE", "1") os.environ.setdefault("HF_HUB_OFFLINE", "1")
@@ -48,9 +48,9 @@ RESOLUTIONS = [512, 768, 1024, 1280, 1536]
PHOTO_TRAIN = "http://127.0.0.1:32678/api/hair/train" PHOTO_TRAIN = "http://127.0.0.1:32678/api/hair/train"
HAIR_CALLBACK = "http://127.0.0.1:8801/api/hair/trainCallBack" HAIR_CALLBACK = "http://127.0.0.1:8801/api/hair/trainCallBack"
SWAP_API = "http://127.0.0.1:8801/api/swapHair/v1" SWAP_API = "http://127.0.0.1:8801/api/swapHair/v1"
GIRL_IMG = "/home/xsl/change_hair/images/girl.png" GIRL_IMG = "/home/ubuntu/change_hair/images/girl.png"
BOY_IMG = "/home/xsl/change_hair/images/boy.png" BOY_IMG = "/home/ubuntu/change_hair/images/boy.png"
PREVIEW_DIR = "/home/xsl/change_hair/hair_grow_service/static/previews" PREVIEW_DIR = "/home/ubuntu/change_hair/hair_grow_service/static/previews"
_detector = _aligner = _matte = None _detector = _aligner = _matte = None
def get_models(): def get_models():
@@ -158,6 +158,10 @@ def step3_wait_and_callback(hair_id, template_img):
else: else:
print(" ✗ 训练超时"); return False print(" ✗ 训练超时"); return False
time.sleep(5) # 等训练进程写完 time.sleep(5) # 等训练进程写完
# 训练时停掉了 hair 服务，回调前需要重启
print(" 重启 change_hair-hair 服务...")
os.system("sudo systemctl restart change_hair-hair")
time.sleep(20)
# 触发回调生成ref材质 # 触发回调生成ref材质
print(" 触发回调生成ref材质...") print(" 触发回调生成ref材质...")
r = req.post(HAIR_CALLBACK, json={"task_id":f"train_{hair_id}","hair_id":hair_id, r = req.post(HAIR_CALLBACK, json={"task_id":f"train_{hair_id}","hair_id":hair_id,
train_hairstyles_parallel.py
+150
View File
@@ -0,0 +1,150 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""批量并行训练多个发型
流程：
阶段A: 串行 step1（准备训练数据，共用GPU模型，避免重复加载）
阶段B: 并行 step2（LoRA训练，kohya独立进程，限制并发数）
阶段C: 启动服务后串行 step3(回调生成ref材质) + step4(模板) + step5(预览)
用法:
cd /home/ubuntu/change_hair/project/hair_service_sd
python /home/ubuntu/change_hair/train_hairstyles_parallel.py
"""
import os
import sys
import time
import json
import subprocess
from concurrent.futures import ThreadPoolExecutor, as_completed
HAIR_SERVICE_DIR = "/home/ubuntu/change_hair/project/hair_service_sd"
os.chdir(HAIR_SERVICE_DIR)
sys.path.insert(0, HAIR_SERVICE_DIR)
os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
os.environ.setdefault("CRYPTOGRAPHY_OPENSSL_NO_LEGACY", "1")
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0")
# ====== 配置：本次要训练的发型 ======
# (hair_id, input_dir, gender, template_img)
HAIRSTYLES = [
("chang_tuoyuan", "/home/ubuntu/change_hair/hair_type_images/chang_tuoyuan", "girl", "/home/ubuntu/change_hair/hair_type_images/chang_tuoyuan/chang_tuoyuan.jpg"),
("chang_bolang", "/home/ubuntu/change_hair/hair_type_images/chang_bolang", "girl", "/home/ubuntu/change_hair/hair_type_images/chang_bolang/chang_bolang.jpg"),
("chang_zhixian", "/home/ubuntu/change_hair/hair_type_images/chang_zhixian", "girl", "/home/ubuntu/change_hair/hair_type_images/chang_zhixian/chang_zhixian.jpg"),
("chang_huaban", "/home/ubuntu/change_hair/hair_type_images/chang_huaban", "girl", "/home/ubuntu/change_hair/hair_type_images/chang_huaban/chang_huaban.jpg"),
("chang_xinxing", "/home/ubuntu/change_hair/hair_type_images/chang_xinxing", "girl", "/home/ubuntu/change_hair/hair_type_images/chang_xinxing/chang_xinxing.jpg"),
]
PARALLEL = 2 # LoRA 训练并发数
# 导入训练模块（复用其 step 函数）
from train_hairstyle_full import (
step1_prepare_data, step2_train, step3_wait_and_callback,
step4_template_material, step5_preview, get_models
)
import train_hairstyle_full as TH
from common.logger import config
def log(msg):
print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)
def main():
t0 = time.time()
ids = [h[0] for h in HAIRSTYLES]
log(f"批量训练 {len(ids)} 个发型: {ids} (并发={PARALLEL})")
# ============ 阶段A: 串行 step1（准备数据）============
log("=" * 60)
log("阶段A: 串行准备训练数据 (step1)")
log("=" * 60)
log(" 预加载模型...")
get_models() # 预加载，5个发型共用
step1_ok, step1_fail = [], []
for hair_id, input_dir, gender, _ in HAIRSTYLES:
# 清理可能的历史训练数据，确保干净
train_dir = config.get('default', 'train_dir')
out_dir = os.path.join(train_dir, hair_id, "images", "1_hairstyle")
os.makedirs(out_dir, exist_ok=True)
os.makedirs(os.path.join(train_dir, hair_id, "model"), exist_ok=True)
log(f" --- step1: {hair_id} ---")
try:
ok = step1_prepare_data(hair_id, input_dir, gender)
if ok:
step1_ok.append(hair_id)
log(f" ✓ {hair_id} 数据准备完成")
else:
step1_fail.append(hair_id)
log(f" ✗ {hair_id} 数据准备失败")
except Exception as e:
step1_fail.append(hair_id)
log(f" ✗ {hair_id} 数据准备异常: {e}")
log(f"阶段A完成: 成功 {len(step1_ok)}/{len(ids)}，失败 {step1_fail}")
if not step1_ok:
log("全部 step1 失败，退出"); return
# 释放 step1 用的 GPU 模型，给 step2 训练腾显存
log("释放 step1 模型显存...")
import torch
try:
for m in [TH._detector, TH._aligner, TH._matte]:
if m is not None and hasattr(m, 'model'):
pass
torch.cuda.empty_cache()
except Exception:
pass
# ============ 阶段B: 并行 step2（LoRA 训练）============
log("=" * 60)
log(f"阶段B: 并行 LoRA 训练 (step2, 并发={PARALLEL})")
log("=" * 60)
step2_results = {} # hair_id -> bool
with ThreadPoolExecutor(max_workers=PARALLEL) as ex:
futures = {ex.submit(step2_train, hid): hid for hid in step1_ok}
for fut in as_completed(futures):
hid = futures[fut]
try:
ok = fut.result()
step2_results[hid] = ok
log(f" {'✓' if ok else '✗'} {hid} step2 (训练启动)={'成功' if ok else '失败'}")
except Exception as e:
step2_results[hid] = False
log(f" ✗ {hid} step2 异常: {e}")
# 等待所有 LoRA 训练真正完成（step2 只是启动，训练在 photo_service 后台跑）
log("等待所有 LoRA 训练完成（检查权重文件）...")
train_dir = config.get('default', 'train_dir')
pending = [hid for hid in step1_ok if step2_results.get(hid)]
completed = set()
deadline = time.time() + 3600 # 最多等1小时
while pending and time.time() < deadline:
still = []
for hid in pending:
lora = os.path.join(train_dir, hid, "model", "hairstyle_hd_lora.safetensors")
if os.path.exists(lora) and os.path.getsize(lora) > 100000000:
completed.add(hid)
log(f" ✓ {hid} LoRA 训练完成 ({os.path.getsize(lora)//1024//1024}MB)")
else:
still.append(hid)
pending = still
if pending:
log(f" ...等待中: {pending}")
time.sleep(30)
if pending:
log(f"⚠️ 超时未完成: {pending}")
log(f"阶段B完成: LoRA 训练完成 {len(completed)}/{len(step1_ok)}")
# 写一个完成清单文件，供阶段C脚本读取
done_file = "/home/ubuntu/change_hair/project/logs/train_batch_done.json"
with open(done_file, "w") as f:
json.dump({"completed": sorted(completed), "all": ids}, f)
log(f"已写入完成清单: {done_file}")
log(f"\n批量训练阶段A+B总耗时 {time.time()-t0:.0f}s")
log(f"完成训练的发型: {sorted(completed)}")
log("接下来需启动服务跑阶段C（step3/4/5 生成材质和预览）")
if __name__ == "__main__":
main()
train_lora_parallel.py
+124
View File
@@ -0,0 +1,124 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""直接调 kohya train_network.py 真正并行训练 LoRA
绕过 photo_service 的串行阻塞，用 subprocess 直接起 train_network.py 进程，
通过 ThreadPoolExecutor 控制并发数。
用法:
python /home/ubuntu/change_hair/train_lora_parallel.py
前提：step1 训练数据已准备好（data/train_material/<hid>/images/1_hairstyle/*.png）
"""
import os
import sys
import time
import subprocess
from concurrent.futures import ThreadPoolExecutor, as_completed
TRAIN_DIR = "/home/ubuntu/change_hair/data/train_material"
KOHYA_WORKDIR = "/home/ubuntu/change_hair/project/kohya_ss_home/kohya_ss"
BASE_MODEL = "/home/ubuntu/change_hair/project/onediff/stable-diffusion-webui/models/Stable-diffusion/v1-5-pruned-emaonly.safetensors"
KOHYA_ACCEL = "/home/ubuntu/miniconda3/envs/kohya/bin/accelerate"
LOG_DIR = "/home/ubuntu/change_hair/project/logs"
# 本次要训练的发型
HAIRSTYLES = ["chang_tuoyuan", "chang_zhixian", "chang_huaban", "chang_xinxing"]
# chang_bolang 已完成，跳过
PARALLEL = 2
def log(msg):
print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)
def train_one(hair_id):
"""训练单个发型的 LoRA，返回 (hair_id, success, log_path)"""
images_dir = os.path.join(TRAIN_DIR, hair_id, "images")
model_dir = os.path.join(TRAIN_DIR, hair_id, "model")
os.makedirs(model_dir, exist_ok=True)
# 准备 sample prompt 文件
sample_dir = os.path.join(model_dir, "sample")
os.makedirs(sample_dir, exist_ok=True)
sample_txt = os.path.join(sample_dir, "prompt.txt")
with open(sample_txt, "w") as f:
f.write('titor hairstyle, easyphoto, faceless, no human, white background, simple background, '
' --n low quality, worst quality, bad anatomy, bad composition, poor, low effort --h 768 '
'--w 768 --s 30 --l 7')
log_path = os.path.join(LOG_DIR, f"train_{hair_id}.log")
cmd = [
"bash", "-c",
f'cd {KOHYA_WORKDIR} && '
f'HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 CUDA_VISIBLE_DEVICES=0 '
f'{KOHYA_ACCEL} launch --num_cpu_threads_per_process=2 "./train_network.py" --enable_bucket '
f'--min_bucket_reso=256 --max_bucket_reso=2048 --pretrained_model_name_or_path="{BASE_MODEL}" '
f'--train_data_dir={images_dir} --resolution="2000,2000" '
f'--output_dir={model_dir} '
'--network_alpha="64" --save_model_as=safetensors --network_module=networks.lora --text_encoder_lr=5e-05 '
'--unet_lr=0.0001 --network_dim=128 --output_name="hairstyle_hd_lora" --lr_scheduler_num_cycles="20" '
'--no_half_vae --learning_rate="0.0001" --lr_scheduler="cosine" --lr_warmup_steps="650" --train_batch_size="1" '
'--max_train_steps="1500" --save_every_n_epochs="100" --mixed_precision="fp16" --save_precision="fp16" '
'--caption_extension=".txt" --sample_sampler=ddim '
f'--sample_prompts={sample_txt} --sample_every_n_epochs="1000" '
'--seed="1234" --cache_latents --optimizer_type="AdamW" --max_data_loader_n_workers="0" --bucket_reso_steps=64 '
'--sdpa --bucket_no_upscale --noise_offset=0.0'
]
log(f" ▶ {hair_id} 开始训练，日志: {log_path}")
t0 = time.time()
try:
with open(log_path, "w") as lf:
proc = subprocess.run(cmd, stdout=lf, stderr=subprocess.STDOUT, timeout=1800)
ok = proc.returncode == 0
lora = os.path.join(model_dir, "hairstyle_hd_lora.safetensors")
ok = ok and os.path.exists(lora) and os.path.getsize(lora) > 100000000
elapsed = time.time() - t0
sz = os.path.getsize(lora) // 1024 // 1024 if os.path.exists(lora) else 0
log(f" {'✓' if ok else '✗'} {hair_id} {'训练完成' if ok else '训练失败'} ({sz}MB, {elapsed:.0f}s)")
return (hair_id, ok, log_path)
except subprocess.TimeoutExpired:
log(f" ✗ {hair_id} 训练超时")
return (hair_id, False, log_path)
except Exception as e:
log(f" ✗ {hair_id} 训练异常: {e}")
return (hair_id, False, log_path)
def main():
t0 = time.time()
log(f"并行 LoRA 训练: {HAIRSTYLES} (并发={PARALLEL})")
results = {}
with ThreadPoolExecutor(max_workers=PARALLEL) as ex:
futures = {ex.submit(train_one, hid): hid for hid in HAIRSTYLES}
for fut in as_completed(futures):
hid = futures[fut]
try:
hair_id, ok, log_path = fut.result()
results[hair_id] = ok
except Exception as e:
log(f" ✗ {hid} 异常: {e}")
results[hid] = False
log(f"\n{'='*60}")
log(f"全部完成，耗时 {time.time()-t0:.0f}s")
for hid, ok in results.items():
lora = os.path.join(TRAIN_DIR, hid, "model", "hairstyle_hd_lora.safetensors")
sz = os.path.getsize(lora) // 1024 // 1024 if os.path.exists(lora) else 0
log(f" {'✓' if ok else '✗'} {hid} ({sz}MB)")
success = [h for h, ok in results.items() if ok]
log(f"成功: {success}")
# 写完成清单
done_file = os.path.join(LOG_DIR, "train_lora_done.json")
import json
with open(done_file, "w") as f:
json.dump({"completed": success, "all": HAIRSTYLES}, f)
log(f"完成清单: {done_file}")
if __name__ == "__main__":
main()