初始化：换发型/换发色/训练发型服务
包含： - hair_service_sd: 主服务（换发型/换发色/生发，端口8801） - photo_service: LoRA调度+训练（端口32678） - hair_grow_service: 调试测试页（端口8888，含4个测试页） - 批量训练脚本（batch_train_hairstyles.py） - 发际线mask自动识别（hairline_mask.py，4种方案） - 手绘mask换发型（hair_swap_manual.py） - 文档：README.md + LARGE_FILES.md + docs/ 大文件（模型权重200G、训练数据123G）已排除，见 LARGE_FILES.md OSS/COS密钥已脱敏为环境变量，原文件备份在本地
This commit is contained in:
xsl
+95
@@ -0,0 +1,95 @@
|
|||||||
|
# ==============================================================
|
||||||
|
# 换发型服务 .gitignore
|
||||||
|
# 大文件（模型权重/训练数据/日志/缓存）全部排除，后续上传网盘
|
||||||
|
# ==============================================================
|
||||||
|
|
||||||
|
# ===== 模型权重（共~200G，网盘保存）=====
|
||||||
|
*.safetensors
|
||||||
|
*.pth
|
||||||
|
*.pt
|
||||||
|
*.ckpt
|
||||||
|
*.bin
|
||||||
|
*.onnx
|
||||||
|
*.tflite
|
||||||
|
*.task
|
||||||
|
# so/dll/pack 等编译产物
|
||||||
|
*.so
|
||||||
|
*.dll
|
||||||
|
*.pack
|
||||||
|
|
||||||
|
# ===== 大数据目录（共~123G，网盘保存）=====
|
||||||
|
# 训练素材（341个发型的LoRA+训练图）
|
||||||
|
data/
|
||||||
|
# webui（SD模型+228个LoRA+ControlNet，~52G，第三方）
|
||||||
|
project/onediff/
|
||||||
|
# kohya训练框架（~13G，第三方）
|
||||||
|
project/kohya_ss_home/
|
||||||
|
|
||||||
|
# hair_service_sd 内的模型权重目录（~9.9G）
|
||||||
|
project/hair_service_sd/weights/
|
||||||
|
# photo_service 临时数据
|
||||||
|
project/photo_service/data/
|
||||||
|
# ref/材质/上传图等运行时生成数据（~5G）
|
||||||
|
project/data/ref_hairstyle/
|
||||||
|
project/data/ref_haircolor/
|
||||||
|
project/data/hair_template_material/
|
||||||
|
project/data/upload_train_imgs/
|
||||||
|
project/data/tmp/
|
||||||
|
project/data/res_dir/
|
||||||
|
project/data/userImage/
|
||||||
|
project/data/user_info/
|
||||||
|
project/data/refImage/
|
||||||
|
project/data/ref_color*/
|
||||||
|
project/data/ref_user_imgs/
|
||||||
|
project/data/ref_online/
|
||||||
|
|
||||||
|
# 训练原图（本次30张发型图，网盘保存）
|
||||||
|
hair_type_images/
|
||||||
|
train_images/
|
||||||
|
huaban1/
|
||||||
|
huaban2/
|
||||||
|
gz-fs/
|
||||||
|
|
||||||
|
# ===== Python 缓存 =====
|
||||||
|
__pycache__/
|
||||||
|
*.pyc
|
||||||
|
*.pyo
|
||||||
|
*.pyd
|
||||||
|
*.egg-info/
|
||||||
|
.Python
|
||||||
|
build/
|
||||||
|
dist/
|
||||||
|
|
||||||
|
# ===== 日志/临时文件 =====
|
||||||
|
*.log
|
||||||
|
nohup.out
|
||||||
|
*.pid
|
||||||
|
project/logs/
|
||||||
|
project/hair_service_sd/hair_service.log
|
||||||
|
project/photo_service/photo_service.log
|
||||||
|
|
||||||
|
# ===== 编辑器/系统 =====
|
||||||
|
.vscode/
|
||||||
|
.idea/
|
||||||
|
*.swp
|
||||||
|
*.swo
|
||||||
|
*~
|
||||||
|
.DS_Store
|
||||||
|
Thumbs.db
|
||||||
|
|
||||||
|
# ===== 测试残留/运行时样图（非代码）=====
|
||||||
|
project/photo_service/crop_origin.png
|
||||||
|
project/photo_service/crop_result.png
|
||||||
|
project/photo_service/image.png
|
||||||
|
project/photo_service/mask.png
|
||||||
|
project/photo_service/request_json.json
|
||||||
|
project/photo_service/photo_service.log
|
||||||
|
project/photo_service/webd/
|
||||||
|
project/hair_service_sd/nohup.out
|
||||||
|
project/hair_service_sd/hair_service.log
|
||||||
|
hair_grow_service/static/previews/
|
||||||
|
|
||||||
|
# ===== 其他运行时临时 =====
|
||||||
|
*.tmp
|
||||||
|
tmp/
|
||||||
|
.cache/
|
||||||
@@ -0,0 +1,344 @@
|
|||||||
|
# 换发型服务 - 本地部署文档
|
||||||
|
|
||||||
|
> 本文档说明换发型服务（hair_service_sd）在本机的部署架构、接口用法、训练功能与日常运维操作。
|
||||||
|
> 该服务从 `szlc@192.168.101.63:/home/szlc/project` 克隆部署至本机 `/home/xsl/change_hair`。
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 一、总体架构
|
||||||
|
|
||||||
|
服务由三个独立进程组成，存在调用依赖关系：
|
||||||
|
|
||||||
|
```
|
||||||
|
┌─────────────────────────────┐
|
||||||
|
客户端 ──HTTP──▶ │ hair_service_sd (:8801) │ 主服务，对外提供所有接口
|
||||||
|
│ conda env: my_hair │
|
||||||
|
└──────────┬──────────────────┘
|
||||||
|
│ 换发型/训练时内部调用
|
||||||
|
┌──────────────┴───────────────┐
|
||||||
|
▼ ▼
|
||||||
|
┌──────────────────────┐ ┌──────────────────────┐
|
||||||
|
│ photo_service (:32678)│ │ webui (:57860) │
|
||||||
|
│ conda env: py310 │ │ conda env: sdwebui │
|
||||||
|
│ LoRA加载/训练调度 │────▶ │ SD模型 + ControlNet │
|
||||||
|
│ 训练: conda env:kohya│ 推理 │ LoRA权重 │
|
||||||
|
└──────────────────────┘ └──────────────────────┘
|
||||||
|
```
|
||||||
|
|
||||||
|
| 进程 | 端口 | conda 环境 | 职责 |
|
||||||
|
|------|------|-----------|------|
|
||||||
|
| **hair_service_sd** | 8801 | my_hair | 主服务，对外提供换发色/换发型/上传发型/训练回调接口 |
|
||||||
|
| **photo_service** | 32678 | py310 (推理) / kohya (训练) | LoRA 权重加载、webui 推理调度、LoRA 训练执行 |
|
||||||
|
| **webui** | 57860 | sdwebui | Stable Diffusion img2img 推理（SD模型 + ControlNet + LoRA） |
|
||||||
|
|
||||||
|
**启动顺序**（存在依赖）：webui → photo_service → hair_service_sd
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 二、目录结构
|
||||||
|
|
||||||
|
所有数据部署在 `/home/xsl/change_hair`（对应原服务器 `/home/szlc/project` + `/data`）：
|
||||||
|
|
||||||
|
```
|
||||||
|
/home/xsl/change_hair/
|
||||||
|
├── README.md # 本文档
|
||||||
|
├── start_all.sh # 统一启动/停止/状态脚本
|
||||||
|
├── start_photo.sh # photo_service 启动包装脚本
|
||||||
|
├── start_hair.sh # hair_service_sd 启动包装脚本
|
||||||
|
├── adapt_paths.sh # 路径适配脚本（部署时一次性执行）
|
||||||
|
├── fix_conda_paths.py # conda 环境路径修复脚本（部署时一次性执行）
|
||||||
|
├── test_haircolor.py # 换发色接口测试脚本
|
||||||
|
├── test_swaphair.py # 换发型接口测试脚本
|
||||||
|
├── test_train_api.py # 训练接口测试脚本
|
||||||
|
│
|
||||||
|
├── project/ # ← 原服务器 /home/szlc/project
|
||||||
|
│ ├── hair_service_sd/ # 主服务（9.9G，含 weights 模型权重）
|
||||||
|
│ ├── onediff/ # webui（49G，含SD模型+228个LoRA+ControlNet）
|
||||||
|
│ │ └── stable-diffusion-webui/
|
||||||
|
│ ├── photo_service/ # LoRA调度+训练服务（4.1M）
|
||||||
|
│ ├── kohya_ss_home/ # kohya LoRA训练框架（13G）
|
||||||
|
│ │ └── kohya_ss/ # train_network.py 所在目录
|
||||||
|
│ ├── data/ # 发型素材（4.9G）
|
||||||
|
│ │ ├── ref_hairstyle/ # 换发型发型材质（183个发型）
|
||||||
|
│ │ ├── ref_haircolor/ # 换发色参考
|
||||||
|
│ │ ├── hair_template_material/ # 发型模板（matting+pkl）
|
||||||
|
│ │ ├── upload_train_imgs/ # 上传发型训练图
|
||||||
|
│ │ ├── ref_user_imgs/ # ⚠️ 训练素材生成用的参考图（当前为空，见下文说明）
|
||||||
|
│ │ └── ...
|
||||||
|
│ └── logs/ # 所有服务日志
|
||||||
|
│
|
||||||
|
└── data/ # ← 原服务器 /data
|
||||||
|
└── train_material/ # 338个发型的LoRA训练素材（121G）
|
||||||
|
└── <hair_id>/
|
||||||
|
├── images/<N>_hairstyle/*.png # 训练图片
|
||||||
|
└── model/hairstyle_hd_lora.safetensors # 训练产出的LoRA权重
|
||||||
|
```
|
||||||
|
|
||||||
|
conda 环境位于 `/home/xsl/miniconda3/envs/`：
|
||||||
|
|
||||||
|
| 环境 | Python | 用途 |
|
||||||
|
|------|--------|------|
|
||||||
|
| my_hair | 3.10.13 | hair_service_sd 主服务 |
|
||||||
|
| sdwebui | 3.10.19 | webui SD 推理 |
|
||||||
|
| py310 | 3.10.14 | photo_service LoRA调度 |
|
||||||
|
| kohya | 3.10.20 | LoRA 训练（**本次部署新建**，原服务用 Docker） |
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 三、日常运维操作
|
||||||
|
|
||||||
|
### 启动 / 停止 / 查看状态
|
||||||
|
|
||||||
|
```bash
|
||||||
|
cd /home/xsl/change_hair
|
||||||
|
|
||||||
|
bash start_all.sh start # 启动全部服务（按依赖顺序）
|
||||||
|
bash start_all.sh stop # 停止全部服务
|
||||||
|
bash start_all.sh restart # 重启全部服务
|
||||||
|
bash start_all.sh status # 查看运行状态与端口
|
||||||
|
```
|
||||||
|
|
||||||
|
启动后 webui 加载 SD 模型需约 1-2 分钟，`status` 显示端口就绪即表示可用。
|
||||||
|
|
||||||
|
### 查看日志
|
||||||
|
|
||||||
|
```bash
|
||||||
|
tail -f project/logs/webui.log # webui 日志
|
||||||
|
tail -f project/logs/photo_service.log # photo_service 日志
|
||||||
|
tail -f project/logs/hair_service.log # 主服务日志
|
||||||
|
```
|
||||||
|
|
||||||
|
### 运行接口测试
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# 换发色测试（生成红色头发）
|
||||||
|
/home/xsl/miniconda3/envs/my_hair/bin/python test_haircolor.py
|
||||||
|
|
||||||
|
# 换发型测试（用已有发型）
|
||||||
|
/home/xsl/miniconda3/envs/my_hair/bin/python test_swaphair.py
|
||||||
|
```
|
||||||
|
|
||||||
|
测试结果图保存在 `project/logs/test_*_result.jpg`。
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 四、接口文档
|
||||||
|
|
||||||
|
主服务地址：`http://127.0.0.1:8801`
|
||||||
|
|
||||||
|
### 1. 换发色 `POST /hairColor/v2`
|
||||||
|
|
||||||
|
把照片中头发的颜色改为指定 RGB。
|
||||||
|
|
||||||
|
**请求参数**（JSON）：
|
||||||
|
|
||||||
|
| 参数 | 类型 | 必填 | 说明 |
|
||||||
|
|------|------|------|------|
|
||||||
|
| img | string | 是 | 图片 base64，需带 `data:image/jpeg;base64,` 前缀 |
|
||||||
|
| userId | string | 是 | 用户标识 |
|
||||||
|
| rgb | [int,int,int] | 是 | 目标颜色，长度3的数组，如 `[255,0,0]` 红色 |
|
||||||
|
| output_format | string | 否 | `base64`（返回base64图片）或 `url`（上传OSS返回链接，默认） |
|
||||||
|
|
||||||
|
**请求示例**：
|
||||||
|
|
||||||
|
```python
|
||||||
|
import base64, requests
|
||||||
|
|
||||||
|
with open("photo.jpg", "rb") as f:
|
||||||
|
img_b64 = "data:image/jpeg;base64," + base64.b64encode(f.read()).decode()
|
||||||
|
|
||||||
|
resp = requests.post("http://127.0.0.1:8801/hairColor/v2", json={
|
||||||
|
"img": img_b64,
|
||||||
|
"userId": "user001",
|
||||||
|
"rgb": [255, 0, 0], # 红色
|
||||||
|
"output_format": "base64"
|
||||||
|
})
|
||||||
|
result = resp.json()
|
||||||
|
# {"msg": "success", "state": 0, "result": "<base64图片>", "umd": ""}
|
||||||
|
```
|
||||||
|
|
||||||
|
### 2. 换发型 `POST /api/swapHair/v1`
|
||||||
|
|
||||||
|
把照片中人物的发型换成指定发型（需该发型已训练好 LoRA）。
|
||||||
|
|
||||||
|
**请求参数**（JSON）：
|
||||||
|
|
||||||
|
| 参数 | 类型 | 必填 | 说明 |
|
||||||
|
|------|------|------|------|
|
||||||
|
| hair_id | string | 是 | 发型ID（对应 train_material 和 ref_hairstyle 中的目录名） |
|
||||||
|
| task_id | string | 是 | 任务ID（任意唯一字符串） |
|
||||||
|
| is_hr | string | 是 | `"true"` 高清版 / `"false"` 普通版 |
|
||||||
|
| user_img_path | string | 是 | 用户照片 base64，需带 `data:image/jpeg;base64,` 前缀 |
|
||||||
|
| output_format | string | 否 | `base64` 或 `url`（默认） |
|
||||||
|
|
||||||
|
**请求示例**：
|
||||||
|
|
||||||
|
```python
|
||||||
|
resp = requests.post("http://127.0.0.1:8801/api/swapHair/v1", json={
|
||||||
|
"hair_id": "1907641420518580226",
|
||||||
|
"task_id": "task_001",
|
||||||
|
"is_hr": "true",
|
||||||
|
"user_img_path": img_b64, # data:image/jpeg;base64,...
|
||||||
|
"output_format": "base64"
|
||||||
|
})
|
||||||
|
result = resp.json()
|
||||||
|
# {"state": 0, "msg": "success", "data": "<base64图片>", "task_id": "task_001"}
|
||||||
|
```
|
||||||
|
|
||||||
|
**可用发型**：发型需同时具备 `ref_hairstyle/<hair_id>` 材质 + `train_material/<hair_id>/model/hairstyle_hd_lora.safetensors` 权重。查看可用发型：
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# 列出材质+权重都齐全的发型
|
||||||
|
for hid in $(ls project/data/ref_hairstyle); do
|
||||||
|
[ -f data/train_material/$hid/model/hairstyle_hd_lora.safetensors ] && echo "$hid"
|
||||||
|
done
|
||||||
|
```
|
||||||
|
|
||||||
|
### 3. 新增发型训练 `POST /api/uploadHair/v1`
|
||||||
|
|
||||||
|
上传新发型图片，异步训练 LoRA。训练完成后发型即可用于换发型。
|
||||||
|
|
||||||
|
**请求参数**（JSON）：
|
||||||
|
|
||||||
|
| 参数 | 类型 | 必填 | 说明 |
|
||||||
|
|------|------|------|------|
|
||||||
|
| hair_id | string | 是 | 新发型ID（自定义唯一ID） |
|
||||||
|
| img_lists | [string] | 是 | 发型图片 URL 列表（至少1张，需 HTTP 可访问的URL） |
|
||||||
|
|
||||||
|
**注意**：此接口 `img_lists` 必须是 HTTP URL（内部用 requests 下载），不支持 base64。如用本地图片，需先启动一个 HTTP 服务托管。
|
||||||
|
|
||||||
|
**响应**：立即返回 `{"msg": "OK", "state": 0}`，训练在后台异步进行（约 30-40 分钟）。
|
||||||
|
|
||||||
|
### 4. 训练回调 `POST /api/hair/trainCallBack`
|
||||||
|
|
||||||
|
photo_service 训练完成后自动调用此接口通知主服务，一般不由外部直接调用。
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 五、新增发型训练（LoRA Training）
|
||||||
|
|
||||||
|
### 工作原理
|
||||||
|
|
||||||
|
原服务通过 Docker 容器（`chinatszrn/ubuntu:kohya_ss`）运行 kohya 的 `train_network.py` 训练 LoRA。**本机部署已改为本地 conda 环境直接执行**，无需 Docker。
|
||||||
|
|
||||||
|
完整链路：
|
||||||
|
```
|
||||||
|
uploadHair → 下载图片 → 生成训练素材(matting/pkl/训练图)
|
||||||
|
→ photo_service 入队
|
||||||
|
→ 后台 train_thread 执行:
|
||||||
|
accelerate launch train_network.py (1500步, 约30-40分钟)
|
||||||
|
→ 产出 hairstyle_hd_lora.safetensors
|
||||||
|
→ 回调 trainCallBack
|
||||||
|
→ 该发型即可用于换发型
|
||||||
|
```
|
||||||
|
|
||||||
|
### 训练配置（已适配本机）
|
||||||
|
|
||||||
|
photo_service 的训练命令（`lora_train_service_1.py` 的 `train_thread` 函数）关键参数：
|
||||||
|
|
||||||
|
| 配置项 | 值 | 说明 |
|
||||||
|
|--------|-----|------|
|
||||||
|
| 训练脚本 | `kohya_ss_home/kohya_ss/train_network.py` | kohya LoRA 训练 |
|
||||||
|
| conda 环境 | `/home/xsl/miniconda3/envs/kohya` | torch 2.0.1+cu118 |
|
||||||
|
| 基础模型 | `v1-5-pruned-emaonly.safetensors` | SD底模（原为 majicmix，本机不存在，已替换） |
|
||||||
|
| 优化器 | `AdamW` | 原为 AdamW8bit（依赖 bitsandbytes，有兼容问题，已改） |
|
||||||
|
| 训练步数 | 1500 | |
|
||||||
|
| LoRA维度 | 128 (network_dim) | |
|
||||||
|
| GPU | device 0 | 单卡 |
|
||||||
|
|
||||||
|
### 手动触发训练（不经过 uploadHair）
|
||||||
|
|
||||||
|
如果已有训练素材（`train_material/<hair_id>/images/<N>_hairstyle/*.png`），可直接调用 photo_service 训练：
|
||||||
|
|
||||||
|
```bash
|
||||||
|
curl -X POST http://127.0.0.1:32678/api/hair/train \
|
||||||
|
-H "Content-Type: application/json" \
|
||||||
|
-d '{
|
||||||
|
"task_id": "manual_train_001",
|
||||||
|
"hair_id": "1905785164224868354",
|
||||||
|
"hair_material_dir": "/home/xsl/change_hair/data/train_material/1905785164224868354",
|
||||||
|
"is_tj": "1",
|
||||||
|
"device_id": "0",
|
||||||
|
"webui_addr": "http://0.0.0.0:32678/"
|
||||||
|
}'
|
||||||
|
```
|
||||||
|
|
||||||
|
训练进度查看：`tail -f project/logs/photo_service.log`（搜索 `steps:`）
|
||||||
|
|
||||||
|
### ⚠️ 训练素材生成依赖说明
|
||||||
|
|
||||||
|
`uploadHair` 的完整流程中，**生成训练素材这一步依赖 `project/data/ref_user_imgs/` 参考图目录**（用于人脸对齐生成训练对）。该目录在原服务器上即为空（这是原项目状态）。
|
||||||
|
|
||||||
|
- 若 `ref_user_imgs` 为空，uploadHair 能下载图片并触发 LoRA 训练，但**生成的新训练素材质量会受限**
|
||||||
|
- 要完整使用 uploadHair 上传全新发型训练，需往 `ref_user_imgs/` 补充参考人像图 + 对应的 `.pkl` 关键点文件
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 六、部署改动记录
|
||||||
|
|
||||||
|
以下是对原服务代码/配置的改动，供维护参考：
|
||||||
|
|
||||||
|
### 1. 路径适配（所有配置和代码）
|
||||||
|
|
||||||
|
原服务器路径 → 本机路径的映射：
|
||||||
|
|
||||||
|
| 原路径 | 本机路径 |
|
||||||
|
|--------|---------|
|
||||||
|
| `/home/szlc/project/...` | `/home/xsl/change_hair/project/...` |
|
||||||
|
| `/data/train_material` | `/home/xsl/change_hair/data/train_material` |
|
||||||
|
| `/home/szlc/miniconda3` | `/home/xsl/miniconda3` |
|
||||||
|
|
||||||
|
涉及文件：`hair_service_sd/config/configure.ini`、`hair_service_sd/*.py`、`photo_service/*.py`、`onediff/.../config.json` 等。
|
||||||
|
执行脚本：`adapt_paths.sh`（业务代码）+ `fix_conda_paths.py`（conda环境）。
|
||||||
|
|
||||||
|
### 2. webui 配置改动
|
||||||
|
|
||||||
|
| 改动 | 说明 |
|
||||||
|
|------|------|
|
||||||
|
| `config.json` 的 onediff compiler 路径 | 指向本机 onediff 目录 |
|
||||||
|
| 启动加 `HF_HUB_OFFLINE=1` | 本机无法访问 huggingface.co，用本地缓存离线加载 CLIP |
|
||||||
|
| SD 模型 | 实际加载 `v1-5-pruned-emaonly`（config 里写的 majicmix 不存在，webui 自动 fallback） |
|
||||||
|
|
||||||
|
### 3. 换发型推理改动（`hair_service_sd/gen_super_image.py`）
|
||||||
|
|
||||||
|
`refiner_checkpoint` 从不存在的 `majicmixRealistic_v7.safetensors` 改为 `v1-5-pruned-emaonly.safetensors`（第185、213行）。否则 webui 报 `Could not find checkpoint` 错误。
|
||||||
|
|
||||||
|
### 4. photo_service 训练改动（`photo_service/lora_train_service_1.py`）
|
||||||
|
|
||||||
|
| 改动 | 原值 | 新值 |
|
||||||
|
|------|------|------|
|
||||||
|
| LoRA 输出目录 | `/gz-fs/Lora` | webui 的 models/Lora 目录（实际路径） |
|
||||||
|
| kohya 训练目录 | `/root/project/kohya_ss_home` | 本机 kohya_ss_home |
|
||||||
|
| 训练方式 | `sudo docker run chinatszrn/ubuntu:kohya_ss accelerate launch ...` | 本地 `accelerate launch ...`（conda kohya 环境） |
|
||||||
|
| 基础模型 | `/mnt/nas_hdd/.../majicmixRealistic_v7.safetensors` | 本机 `v1-5-pruned-emaonly.safetensors` |
|
||||||
|
| 优化器 | AdamW8bit（bitsandbytes） | AdamW（torch 原生） |
|
||||||
|
| tokenizer | `--tokenizer_cache_dir=/home/chinatszrn/.cache/clip` | 删除（用 HF 缓存 + 离线模式） |
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 七、已知限制与注意事项
|
||||||
|
|
||||||
|
1. **OSS 上传**：换发色/换发型默认 `output_format=url` 会把结果图上传到**生产阿里云 OSS**（密钥硬编码在代码中）。仅本地测试时建议用 `output_format=base64` 避免污染生产环境。
|
||||||
|
|
||||||
|
2. **网络限制**：本机无法访问 `huggingface.co`。webui 和训练都依赖本地 HF 缓存（`~/.cache/huggingface`），已配置离线模式。如需新增 HF 模型需手动准备缓存。
|
||||||
|
|
||||||
|
3. **ref_user_imgs 为空**：影响 uploadHair 生成训练素材的完整性（见第五节说明）。
|
||||||
|
|
||||||
|
4. **单卡部署**：GPU 固定使用 device 0。`CUDA_VISIBLE_DEVICES=0` 已写入启动脚本。如有多卡需求需调整配置。
|
||||||
|
|
||||||
|
5. **服务进程保活**：服务通过 nohup + wrapper 脚本启动。**切勿直接 `python xxx.py &` 启动**，会被 shell 会话退出时清理。请始终用 `start_all.sh start` 或 `start_*.sh` 脚本。
|
||||||
|
|
||||||
|
6. **webui 启动较慢**：加载 SD 模型需 1-2 分钟，启动后需等待 `start_all.sh status` 显示端口就绪才能正常服务。
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 八、快速排障
|
||||||
|
|
||||||
|
| 现象 | 可能原因 | 解决 |
|
||||||
|
|------|---------|------|
|
||||||
|
| 换发型报 `推理发型失败` | webui 未就绪或 LoRA 缺失 | `start_all.sh status` 检查 webui 端口；确认发型 LoRA 存在 |
|
||||||
|
| webui 报 `Could not find checkpoint` | refiner 模型名不对 | 确认 `gen_super_image.py` 的 refiner_checkpoint 为 `v1-5-pruned-emaonly.safetensors` |
|
||||||
|
| 训练报 bitsandbytes 错误 | bnb 版本/优化器问题 | 确认用 `AdamW`（非 AdamW8bit） |
|
||||||
|
| webui 卡在 huggingface 重试 | HF 离线模式未启用 | 启动脚本已设 `HF_HUB_OFFLINE=1`，确认未丢失 |
|
||||||
|
| 服务启动后立即退出 | 未用 wrapper 脚本启动 | 用 `start_all.sh start` 启动 |
|
||||||
|
| conda 环境 python 段错误 | 环境二进制损坏 | 重新 rsync 环境并用 `fix_conda_paths.py` 修复（勿用 sed） |
|
||||||
+100
@@ -0,0 +1,100 @@
|
|||||||
|
#!/bin/bash
|
||||||
|
# 路径适配脚本：修正所有硬编码路径 + conda环境修复 + 端口适配
|
||||||
|
# 在数据同步完成后执行
|
||||||
|
|
||||||
|
set -u
|
||||||
|
BASE="/home/xsl/change_hair"
|
||||||
|
PROJ="$BASE/project"
|
||||||
|
LOG="$PROJ/logs/adapt.log"
|
||||||
|
mkdir -p "$PROJ/logs"
|
||||||
|
echo "路径适配开始: $(date '+%Y-%m-%d %H:%M:%S')" | tee "$LOG"
|
||||||
|
|
||||||
|
# ============================================================
|
||||||
|
# 第一部分：conda 环境硬编码路径修复
|
||||||
|
# ============================================================
|
||||||
|
echo "=== [1/4] 修复 conda 环境 shebang/路径前缀 ===" | tee -a "$LOG"
|
||||||
|
for ENV in py310 my_hair sdwebui; do
|
||||||
|
ENVDIR="/home/xsl/miniconda3/envs/$ENV"
|
||||||
|
if [ -d "$ENVDIR" ]; then
|
||||||
|
# 替换所有文本文件里的旧前缀为新前缀
|
||||||
|
# 旧前缀有两种：/home/szlc/miniconda3 和 /usr/local/miniconda3
|
||||||
|
grep -rl '/home/szlc/miniconda3\|/usr/local/miniconda3' "$ENVDIR" 2>/dev/null | while read -r f; do
|
||||||
|
sed -i 's|/home/szlc/miniconda3|/home/xsl/miniconda3|g; s|/usr/local/miniconda3|/home/xsl/miniconda3|g' "$f"
|
||||||
|
done
|
||||||
|
echo " [$ENV] 修复完成" | tee -a "$LOG"
|
||||||
|
else
|
||||||
|
echo " [$ENV] 目录不存在，跳过" | tee -a "$LOG"
|
||||||
|
fi
|
||||||
|
done
|
||||||
|
|
||||||
|
# ============================================================
|
||||||
|
# 第二部分：hair_service_sd 配置和代码路径适配
|
||||||
|
# ============================================================
|
||||||
|
echo "=== [2/4] 适配 hair_service_sd 路径 ===" | tee -a "$LOG"
|
||||||
|
HS="$PROJ/hair_service_sd"
|
||||||
|
|
||||||
|
# 安全替换函数：只处理 .py / .ini / .conf 文件，跳过权重/图片/二进制
|
||||||
|
adapt_hair_paths() {
|
||||||
|
local f="$1"
|
||||||
|
[ -f "$f" ] || return 0
|
||||||
|
case "$f" in
|
||||||
|
*.py|*.ini|*.conf|*.sh|*.json|*.yaml|*.yml|*.txt) ;;
|
||||||
|
*) return 0 ;;
|
||||||
|
esac
|
||||||
|
sed -i \
|
||||||
|
-e 's|/home/szlc/project|/home/xsl/change_hair/project|g' \
|
||||||
|
-e 's|/data/train_material|/home/xsl/change_hair/data/train_material|g' \
|
||||||
|
-e 's|/home/szlc/miniconda3|/home/xsl/miniconda3|g' \
|
||||||
|
-e 's|/usr/local/miniconda3|/home/xsl/miniconda3|g' \
|
||||||
|
"$f"
|
||||||
|
}
|
||||||
|
|
||||||
|
# 处理 hair_service_sd 顶层文件和 common/config 子目录
|
||||||
|
for f in "$HS"/*.py "$HS"/*.ini "$HS"/*.sh "$HS"/*.conf "$HS"/common/*.py "$HS"/config/*.ini "$HS"/config/*.py "$HS"/core/*.py "$HS"/utils/*.py; do
|
||||||
|
adapt_hair_paths "$f"
|
||||||
|
done
|
||||||
|
echo " hair_service_sd 配置/代码路径已适配" | tee -a "$LOG"
|
||||||
|
|
||||||
|
# ============================================================
|
||||||
|
# 第三部分：photo_service 路径适配
|
||||||
|
# 注意：photo_service 的 chinatszrn/kohya 路径属于 Docker 训练链路，
|
||||||
|
# 本机不配 Docker，这些路径保留不动（不影响换发色/换发型）
|
||||||
|
# 但基础路径和 callback_url 要适配
|
||||||
|
# ============================================================
|
||||||
|
echo "=== [3/4] 适配 photo_service 路径 ===" | tee -a "$LOG"
|
||||||
|
PS="$PROJ/photo_service"
|
||||||
|
for f in "$PS"/*.py; do
|
||||||
|
[ -f "$f" ] || continue
|
||||||
|
sed -i \
|
||||||
|
-e 's|/home/szlc/project|/home/xsl/change_hair/project|g' \
|
||||||
|
-e 's|/home/szlc/miniconda3|/home/xsl/miniconda3|g' \
|
||||||
|
"$f"
|
||||||
|
done
|
||||||
|
echo " photo_service 基础路径已适配（kohya/Docker路径保留）" | tee -a "$LOG"
|
||||||
|
|
||||||
|
# ============================================================
|
||||||
|
# 第四部分：onediff webui 路径适配
|
||||||
|
# ============================================================
|
||||||
|
echo "=== [4/4] 适配 onediff webui 路径 ===" | tee -a "$LOG"
|
||||||
|
OD="$PROJ/onediff/stable-diffusion-webui"
|
||||||
|
# webui config.json 里的 onediff compiler_caches 路径
|
||||||
|
if [ -f "$OD/config.json" ]; then
|
||||||
|
sed -i 's|/home/student/Documents/workspace_cxt_tianjing_hair/onediff|/home/xsl/change_hair/project/onediff|g; s|/home/szlc/project/onediff|/home/xsl/change_hair/project/onediff|g' "$OD/config.json"
|
||||||
|
echo " webui config.json 已适配" | tee -a "$LOG"
|
||||||
|
fi
|
||||||
|
# webui 启动脚本
|
||||||
|
for f in "$OD"/webui.sh "$OD"/webui-user.sh "$PROJ"/onediff/onediff/*.sh; do
|
||||||
|
[ -f "$f" ] && sed -i 's|/home/szlc/miniconda3|/home/xsl/miniconda3|g' "$f"
|
||||||
|
done
|
||||||
|
|
||||||
|
# ============================================================
|
||||||
|
# 第五部分：Lora 软链接（photo_service 训练产物输出位置）
|
||||||
|
# 本机用真实目录，不走 /gz-fs 软链
|
||||||
|
# ============================================================
|
||||||
|
echo "=== [5/5] 配置 Lora 目录 ===" | tee -a "$LOG"
|
||||||
|
mkdir -p "$BASE/gz-fs/Lora"
|
||||||
|
# webui 的 models/Lora 已随 onediff 拷贝过来（真实目录），无需软链
|
||||||
|
echo " Lora 目录就绪: $BASE/gz-fs/Lora" | tee -a "$LOG"
|
||||||
|
|
||||||
|
echo "路径适配完成: $(date '+%Y-%m-%d %H:%M:%S')" | tee -a "$LOG"
|
||||||
|
echo "请检查 $LOG 确认无误" | tee -a "$LOG"
|
||||||
@@ -0,0 +1,232 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
"""批量训练发型 LoRA
|
||||||
|
|
||||||
|
对 hair_type_images/ 下每张图（一种发型）：
|
||||||
|
1. 拷贝到独立输入目录（train_hairstyle_full.py 要求 --input 是目录）
|
||||||
|
2. 调用 train_hairstyle_full.py 完整流程：数据增强→训练→材质→女生预览
|
||||||
|
3. 串行执行，记录每个发型成功/失败
|
||||||
|
|
||||||
|
用法:
|
||||||
|
cd /home/xsl/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 心形-心形]
|
||||||
|
|
||||||
|
注意：
|
||||||
|
- 中文 hair_id 直接用文件名（去扩展名）作为 ID
|
||||||
|
- 每个发型约 30-40 分钟，30 个串行约 15-20 小时
|
||||||
|
- 失败的发型记录到失败清单，最后可单独重跑
|
||||||
|
"""
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
import time
|
||||||
|
import shutil
|
||||||
|
import argparse
|
||||||
|
import subprocess
|
||||||
|
from datetime import datetime
|
||||||
|
|
||||||
|
SRC_DEFAULT = "/home/xsl/change_hair/hair_type_images"
|
||||||
|
WORK_DIR = "/home/xsl/change_hair/data/batch_train_inputs"
|
||||||
|
LOG_FILE = "/home/xsl/change_hair/data/batch_train_log.txt"
|
||||||
|
TRAIN_SCRIPT = "/home/xsl/change_hair/train_hairstyle_full.py"
|
||||||
|
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd"
|
||||||
|
LOCK_FILE = "/home/xsl/change_hair/data/batch_train.pid"
|
||||||
|
DONE_FILE = "/home/xsl/change_hair/data/batch_train_done.txt" # 已完成发型清单（断点续跑）
|
||||||
|
LORA_DIR = "/home/xsl/change_hair/data/train_material" # LoRA 输出根目录
|
||||||
|
|
||||||
|
|
||||||
|
def acquire_lock():
|
||||||
|
"""PID 锁，防止重复启动多个实例互相干扰"""
|
||||||
|
import psutil
|
||||||
|
if os.path.exists(LOCK_FILE):
|
||||||
|
try:
|
||||||
|
old_pid = int(open(LOCK_FILE).read().strip())
|
||||||
|
if psutil.pid_exists(old_pid):
|
||||||
|
# 检查是不是真的本脚本进程
|
||||||
|
try:
|
||||||
|
proc = psutil.Process(old_pid)
|
||||||
|
if "batch_train_hairstyles" in " ".join(proc.cmdline()):
|
||||||
|
print(f"✗ 已有批量训练实例在运行 (PID={old_pid})，请先停止它再启动。")
|
||||||
|
print(f" 锁文件: {LOCK_FILE}")
|
||||||
|
sys.exit(2)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
except (ValueError, OSError):
|
||||||
|
pass
|
||||||
|
os.makedirs(os.path.dirname(LOCK_FILE), exist_ok=True)
|
||||||
|
open(LOCK_FILE, "w").write(str(os.getpid()))
|
||||||
|
|
||||||
|
|
||||||
|
def release_lock():
|
||||||
|
try:
|
||||||
|
os.remove(LOCK_FILE)
|
||||||
|
except OSError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
def load_done():
|
||||||
|
"""读取已完成的发型集合"""
|
||||||
|
if not os.path.exists(DONE_FILE):
|
||||||
|
return set()
|
||||||
|
return set(l.strip() for l in open(DONE_FILE, encoding="utf-8") if l.strip())
|
||||||
|
|
||||||
|
|
||||||
|
def mark_done(hair_id):
|
||||||
|
with open(DONE_FILE, "a", encoding="utf-8") as f:
|
||||||
|
f.write(hair_id + "\n")
|
||||||
|
|
||||||
|
|
||||||
|
def log(msg):
|
||||||
|
line = f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] {msg}"
|
||||||
|
print(line, flush=True)
|
||||||
|
with open(LOG_FILE, "a", encoding="utf-8") as f:
|
||||||
|
f.write(line + "\n")
|
||||||
|
|
||||||
|
|
||||||
|
def prepare_input_dir(hair_id, src_img, work_dir):
|
||||||
|
"""为单个发型创建独立输入目录，里面只放这一张图"""
|
||||||
|
d = os.path.join(work_dir, hair_id)
|
||||||
|
if os.path.exists(d):
|
||||||
|
shutil.rmtree(d)
|
||||||
|
os.makedirs(d, exist_ok=True)
|
||||||
|
# 统一用 png 拷入，避免扩展名问题
|
||||||
|
dst = os.path.join(d, hair_id + ".png")
|
||||||
|
shutil.copy(src_img, dst)
|
||||||
|
return d, dst
|
||||||
|
|
||||||
|
|
||||||
|
def run_one(hair_id, src_img, gender, py, done_set, force=False):
|
||||||
|
"""跑单个发型的完整训练流程，返回 (ok, elapsed, skipped)"""
|
||||||
|
# 断点续跑：LoRA 已存在视为完成
|
||||||
|
lora_path = os.path.join(LORA_DIR, hair_id, "model", "hairstyle_hd_lora.safetensors")
|
||||||
|
if not force and (hair_id in done_set or
|
||||||
|
(os.path.exists(lora_path) and os.path.getsize(lora_path) > 100000000)):
|
||||||
|
log(f" ⊙ {hair_id} 已完成（LoRA 存在），跳过")
|
||||||
|
return True, 0, True
|
||||||
|
|
||||||
|
t0 = time.time()
|
||||||
|
log(f"==== 开始训练: {hair_id} ====")
|
||||||
|
|
||||||
|
# 步骤1：准备输入目录
|
||||||
|
try:
|
||||||
|
in_dir, tpl_img = prepare_input_dir(hair_id, src_img, WORK_DIR)
|
||||||
|
log(f" 输入目录就绪: {in_dir}")
|
||||||
|
except Exception as e:
|
||||||
|
log(f" ✗ 输入目录准备失败: {e}")
|
||||||
|
return False, time.time() - t0, False
|
||||||
|
|
||||||
|
# 步骤2-5：调用 train_hairstyle_full.py
|
||||||
|
cmd = [
|
||||||
|
py, TRAIN_SCRIPT,
|
||||||
|
"--hair-id", hair_id,
|
||||||
|
"--input", in_dir,
|
||||||
|
"--gender", gender,
|
||||||
|
"--template-img", tpl_img,
|
||||||
|
]
|
||||||
|
log(f" 调用: {' '.join(cmd[:2])} ... --hair-id {hair_id}")
|
||||||
|
try:
|
||||||
|
# 子进程输出实时写到日志文件
|
||||||
|
proc_log = os.path.join("/home/xsl/change_hair/data", f"subprocess_{hair_id}.log")
|
||||||
|
with open(proc_log, "w", encoding="utf-8") as f:
|
||||||
|
r = subprocess.run(
|
||||||
|
cmd, cwd=HAIR_SERVICE_DIR,
|
||||||
|
stdout=f, stderr=subprocess.STDOUT,
|
||||||
|
timeout=5400, # 单发型最长 90 分钟
|
||||||
|
)
|
||||||
|
elapsed = time.time() - t0
|
||||||
|
if r.returncode == 0:
|
||||||
|
log(f" ✓ 训练成功: {hair_id} (耗时 {elapsed/60:.1f} 分钟)")
|
||||||
|
mark_done(hair_id)
|
||||||
|
return True, elapsed, False
|
||||||
|
else:
|
||||||
|
log(f" ✗ 训练失败(returncode={r.returncode}): {hair_id}, 详见 {proc_log}")
|
||||||
|
return False, elapsed, False
|
||||||
|
except subprocess.TimeoutExpired:
|
||||||
|
log(f" ✗ 训练超时(>90分钟): {hair_id}")
|
||||||
|
return False, time.time() - t0, False
|
||||||
|
except Exception as e:
|
||||||
|
log(f" ✗ 异常: {hair_id}: {e}")
|
||||||
|
return False, time.time() - t0, False
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
ap = argparse.ArgumentParser(description="批量训练发型")
|
||||||
|
ap.add_argument("--src", default=SRC_DEFAULT, help="原图目录")
|
||||||
|
ap.add_argument("--gender", default="girl", choices=["boy", "girl"])
|
||||||
|
ap.add_argument("--py", default="/home/xsl/miniconda3/envs/my_hair/bin/python",
|
||||||
|
help="python 解释器")
|
||||||
|
ap.add_argument("--only", default=None, help="只训练指定 hair_id（调试用）")
|
||||||
|
ap.add_argument("--start-from", default=None,
|
||||||
|
help="从指定 hair_id 开始（含），跳过之前的（断点续跑）")
|
||||||
|
ap.add_argument("--force", action="store_true",
|
||||||
|
help="强制重训（忽略已完成清单和已存在的 LoRA）")
|
||||||
|
args = ap.parse_args()
|
||||||
|
|
||||||
|
acquire_lock()
|
||||||
|
done_set = load_done() if not args.force else set()
|
||||||
|
|
||||||
|
# 收集发型清单（按文件名排序）
|
||||||
|
files = []
|
||||||
|
for f in sorted(os.listdir(args.src)):
|
||||||
|
if f.lower().endswith((".jpg", ".jpeg", ".png")):
|
||||||
|
files.append(f)
|
||||||
|
if not files:
|
||||||
|
log(f"✗ 源目录无图片: {args.src}")
|
||||||
|
release_lock(); sys.exit(1)
|
||||||
|
|
||||||
|
# 构造 (hair_id, src_img)
|
||||||
|
tasks = []
|
||||||
|
for f in files:
|
||||||
|
hair_id = os.path.splitext(f)[0]
|
||||||
|
tasks.append((hair_id, os.path.join(args.src, f)))
|
||||||
|
|
||||||
|
# 过滤
|
||||||
|
if args.only:
|
||||||
|
tasks = [t for t in tasks if t[0] == args.only]
|
||||||
|
if not tasks:
|
||||||
|
log(f"✗ --only {args.only} 未匹配到发型"); release_lock(); sys.exit(1)
|
||||||
|
if args.start_from:
|
||||||
|
idx = next((i for i, t in enumerate(tasks) if t[0] == args.start_from), 0)
|
||||||
|
tasks = tasks[idx:]
|
||||||
|
|
||||||
|
log(f"\n{'#'*60}")
|
||||||
|
log(f"# 批量训练启动 (PID={os.getpid()})")
|
||||||
|
log(f"# 共 {len(tasks)} 个发型 | gender={args.gender} | force={args.force}")
|
||||||
|
if done_set:
|
||||||
|
log(f"# 已完成 {len(done_set)} 个（将跳过）")
|
||||||
|
log(f"# 预计总耗时约 {len(tasks)*35} 分钟")
|
||||||
|
log(f"{'#'*60}\n")
|
||||||
|
|
||||||
|
os.makedirs(WORK_DIR, exist_ok=True)
|
||||||
|
os.makedirs(os.path.dirname(LOG_FILE), exist_ok=True)
|
||||||
|
|
||||||
|
results = [] # (hair_id, ok, elapsed, skipped)
|
||||||
|
t_total = time.time()
|
||||||
|
try:
|
||||||
|
for i, (hair_id, src) in enumerate(tasks, 1):
|
||||||
|
log(f"\n[进度 {i}/{len(tasks)}] === {hair_id} ===")
|
||||||
|
ok, el, skipped = run_one(hair_id, src, args.gender, args.py, done_set, args.force)
|
||||||
|
results.append((hair_id, ok, el, skipped))
|
||||||
|
finally:
|
||||||
|
release_lock()
|
||||||
|
|
||||||
|
# 汇总
|
||||||
|
log(f"\n{'#'*60}")
|
||||||
|
log(f"# 批量训练完成汇总")
|
||||||
|
log(f"{'#'*60}")
|
||||||
|
succ = [r for r in results if r[1]]
|
||||||
|
fail = [r for r in results if not r[1]]
|
||||||
|
skip = [r for r in results if r[3]]
|
||||||
|
total_time = time.time() - t_total
|
||||||
|
for r in results:
|
||||||
|
hid, ok, el = r[0], r[1], r[2]
|
||||||
|
flag = "SKIP" if r[3] else ("OK" if ok else "FAIL")
|
||||||
|
log(f" [{flag}] {hid:20s} ({el/60:.1f} 分钟)")
|
||||||
|
log(f"\n成功 {len(succ)}/{len(results)}（含跳过 {len(skip)}），失败 {len(fail)}，总耗时 {total_time/3600:.1f} 小时")
|
||||||
|
if fail:
|
||||||
|
log(f"失败清单（可重跑）: {' '.join(f[0] for f in fail)}")
|
||||||
|
log(f"重跑命令: python {__file__} --gender {args.gender} --only <hair_id>")
|
||||||
|
log(f"详细日志: {LOG_FILE}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -0,0 +1,301 @@
|
|||||||
|
# 区域生发功能 - 技术调研报告
|
||||||
|
|
||||||
|
> **功能定义**：输入一张头发稀少/发际线偏高的人头照片，在照片上指定一个区域或遮罩，输出在该区域内"生发"（补全/填充头发）的结果图。
|
||||||
|
>
|
||||||
|
> 调研日期：2026-06-21
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 一、需求分析与技术挑战
|
||||||
|
|
||||||
|
### 1.1 需求拆解
|
||||||
|
|
||||||
|
本功能本质是一个 **"受限区域头发补全"**（constrained hair inpainting）任务，可拆为三个子问题：
|
||||||
|
|
||||||
|
| 子问题 | 描述 | 难度 |
|
||||||
|
|--------|------|------|
|
||||||
|
| ① 在哪生发 | 确定缺发区域（发际线后退区、头顶稀疏区），即 mask | 中（可人工画，也可自动检测） |
|
||||||
|
| ② 生成什么 | 在 mask 区域内生成与本人发色、发质、走向自然衔接的头发 | 高（核心难点） |
|
||||||
|
| ③ 如何融合 | 把生成结果无缝贴回原图，消除边缘拼接痕迹 | 低（已有成熟工具） |
|
||||||
|
|
||||||
|
### 1.2 关键挑战
|
||||||
|
|
||||||
|
1. **生成真实感**：新生头发必须符合该人的发色/卷曲度/走向，不能突兀。SD1.5 对细密发丝的生成能力有限，容易出现"塑料感"或糊。
|
||||||
|
2. **边缘自然衔接**：新生头发与原有头发、额头皮质的交界处不能有硬边或色差。
|
||||||
|
3. **发丝纹理一致性**：新头发应延续原头发的纹理方向（从发旋/发根向外）。
|
||||||
|
4. **不破坏五官**：生发区域靠近额头/眉毛，生成时不能改变五官。
|
||||||
|
5. **mask 边界合理性**：mask 区域应只覆盖"该长头发但目前没长"的地方，不能覆盖到脸上或背景。
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 二、技术方案全景
|
||||||
|
|
||||||
|
综合代码库已有能力和业界方案，归纳出 **4 种实现路径**，按推荐度排序。
|
||||||
|
|
||||||
|
### 方案 A：复用现有 webui Inpainting（★★★★★ 强烈推荐）
|
||||||
|
|
||||||
|
**核心思路**：把"生发"当作一次 SD inpainting——mask 区域重绘为头发，非 mask 区域保持不变。这正是换发型服务里 `webui_img2img` 已经在做的事，几乎可原样复用。
|
||||||
|
|
||||||
|
```
|
||||||
|
[原图] → [人检+1k关键点] → [构造生发mask] → [webui_img2img(mask内重绘为头发)] → [泊松融合] → [结果]
|
||||||
|
```
|
||||||
|
|
||||||
|
**可复用的现成组件**（均在代码库内，已验证可用）：
|
||||||
|
|
||||||
|
| 环节 | 复用组件 | 文件:行号 | 说明 |
|
||||||
|
|------|---------|-----------|------|
|
||||||
|
| 人脸检测 | `RetinaFaceDetector.forward` | models/detector.py:148 | 定位人脸框 |
|
||||||
|
| 1k关键点 | `MomocvFaceAlignment1K.stable_forward` | utils/MomocvFaceAlignment1K.py:404 | 发际线定位的关键 |
|
||||||
|
| 头发mask | `Generator_Matte.matte_inference` | hair_matting/Generator_Matte.py:81 | 现有头发alpha |
|
||||||
|
| **发际线/头皮分割** | `Generator_BaldSeg_5c.forward` | models/Generator_BaldSeg.py:45 | **自动定位缺发区，权重存在** |
|
||||||
|
| 额头凸包mask | `draw_bigger_hull_mask` | utils/landmark_processor.py:1079 | 关键点→额头区mask |
|
||||||
|
| **SD生成** | `webui_img2img(img, mask_img, ...)` | gen_super_image.py:416 | **核心生成，已验证可用** |
|
||||||
|
| 泊松融合 | `cv2.seamlessClone` | core/hairstyle_model.py:1380 | 消除拼接缝 |
|
||||||
|
|
||||||
|
**mask 构造方法（两种交互模式）**：
|
||||||
|
|
||||||
|
- **手动模式**（用户画遮罩）：前端提供画笔工具，用户在发际线/稀疏区涂抹 → 直接作为 inpainting mask。最简单可靠。
|
||||||
|
- **自动模式**（自动检测缺发区）：`BaldSeg` 的 5 类分割直接区分"有发区/头皮区"，或用 `额头凸包mask - 现有头发mask = 缺发带`。无需用户画，但需调参。
|
||||||
|
|
||||||
|
**生成参数建议**（基于 `build_body_v2`，gen_super_image.py:86）：
|
||||||
|
- `denoising_strength = 0.5~0.7`（重绘强度，过低生不出新发，过高破坏原有头发）
|
||||||
|
- `mask_blur = 8~12`（边缘羽化，避免硬边）
|
||||||
|
- `prompt = "thick natural hair, detailed hair strands, realistic hairline, high detail"`
|
||||||
|
- `negative_prompt = "bald, receding hairline, thinning hair, bald spot, skin, low quality"`
|
||||||
|
- `inpainting_fill = 1`（mask 区域填噪声重绘）
|
||||||
|
- `steps = 20~30`
|
||||||
|
|
||||||
|
**优势**：
|
||||||
|
- ✅ 90% 代码已存在且验证过，开发量最小（主要是拼装 + mask 构造逻辑）
|
||||||
|
- ✅ 与现有服务架构完全一致，可作为一个新接口 `/api/hairGrow/v1` 接入
|
||||||
|
- ✅ 无需新增任何模型权重
|
||||||
|
|
||||||
|
**劣势**：
|
||||||
|
- ⚠️ SD1.5（v1-5-pruned）生成细密发丝的真实感有限，发丝可能偏糊
|
||||||
|
- ⚠️ 同一个人多次生发结果可能不一致（生成模型随机性）
|
||||||
|
|
||||||
|
**预估工作量**：2-3 天（mask 构造逻辑 + 接口封装 + 调参）
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### 方案 B：ControlNet Canny 约束发丝走向（★★★★ 推荐）
|
||||||
|
|
||||||
|
**核心思路**：在方案 A 基础上，引入 **ControlNet Canny** 约束新生头发的走向，使其与已有头发的纹理方向一致，提升真实感。
|
||||||
|
|
||||||
|
```
|
||||||
|
[原图] → 提取已有头发的 canny 边缘 → [mask内: SD inpainting + ControlNet canny约束] → [结果]
|
||||||
|
```
|
||||||
|
|
||||||
|
**关键改动**（相对方案A）：
|
||||||
|
- 启用 webui 的 ControlNet（代码库已装 `control_v11p_sd15_canny` 模型，可用）
|
||||||
|
- 参照 `gen_super_image.py:244 build_body` 的 `alwayson_scripts` 结构，把 openpose 改成 canny
|
||||||
|
- canny 输入图：对原图做 `cv2.Canny`，或在 mask 外保留原头发边缘作为约束
|
||||||
|
|
||||||
|
**ControlNet 配置示例**（需新增到 build_body）：
|
||||||
|
```python
|
||||||
|
"alwayson_scripts": {
|
||||||
|
"controlnet": {
|
||||||
|
"args": [{
|
||||||
|
"model": "control_v11p_sd15_canny",
|
||||||
|
"module": "canny",
|
||||||
|
"weight": 0.5, # 约束强度，0.3-0.7
|
||||||
|
"input_image": <原图base64>, # webui自动做canny
|
||||||
|
"resize_mode": "just_resize"
|
||||||
|
}]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
**优势**：
|
||||||
|
- ✅ 发丝走向更自然、与原发衔接更好
|
||||||
|
- ✅ ControlNet 模型和扩展都已在 webui 里，无需额外部署
|
||||||
|
|
||||||
|
**劣势**：
|
||||||
|
- ⚠️ canny 约束在"完全无发的光头区"无法提供有效边缘（那里本来就没头发），主要对"稀疏→加密"有效
|
||||||
|
- ⚠️ 需要调 `weight`，过强会让新生头发僵硬
|
||||||
|
|
||||||
|
**预估工作量**：3-4 天（方案A + ControlNet 调试）
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### 方案 C：SAM 自动分割 + SD Inpainting（★★★ 可选）
|
||||||
|
|
||||||
|
**核心思路**：用 **Segment Anything (SAM)** 替代手动画 mask——用户点击缺发区域，SAM 自动分割出精确边界，再做 inpainting。
|
||||||
|
|
||||||
|
```
|
||||||
|
[原图] → [用户点击发际线区域] → [SAM分割出该区域mask] → [SD inpainting] → [结果]
|
||||||
|
```
|
||||||
|
|
||||||
|
**业界实践**：这是目前社区主流的"点击式 inpainting"范式，已有成熟工具：
|
||||||
|
- **sd-webui-inpaint-anything** 扩展（[GitHub](https://github.com/Uminosachi/sd-webui-inpaint-anything)）：webui 集成 SAM + inpainting
|
||||||
|
- 学术：[Segment Anything Meets Image Inpainting (arXiv:2304.06790)](https://arxiv.org/abs/2304.06790) 提出 click-and-fill 范式
|
||||||
|
|
||||||
|
**是否需要 SAM**：
|
||||||
|
- 若用户**手动画 mask**（方案A）→ 不需要 SAM
|
||||||
|
- 若要**点击式自动分割** → 需要 SAM（代码库没有，需新增 ~2.5G 模型）
|
||||||
|
|
||||||
|
**优势**：
|
||||||
|
- ✅ 交互体验好（点一下即可），mask 边界比手画精确
|
||||||
|
|
||||||
|
**劣势**：
|
||||||
|
- ⚠️ 需新增 SAM 模型（vit_h 约 2.5G），增加部署复杂度
|
||||||
|
- ⚠️ SAM 对"发际线缺发带"这种渐变区域的分割不一定准确（头发与皮肤对比度低时）
|
||||||
|
- ⚠️ 代码库已有 `Generator_BaldSeg`（头皮分割）+ 1k关键点（发际线定位），**可达到类似自动 mask 效果，不一定非要用 SAM**
|
||||||
|
|
||||||
|
**预估工作量**：5-7 天（含 SAM 集成）
|
||||||
|
|
||||||
|
**结论**：**不优先推荐**。代码库的 BaldSeg + 关键点方案已能实现"自动定位缺发区"，SAM 的边际价值不大，除非未来需要通用点击分割。
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### 方案 D：专项 LoRA 微调 / 训练生发模型（★★ 备选）
|
||||||
|
|
||||||
|
**核心思路**：收集"稀疏发/高发际线 → 浓密发"的配对数据，训练一个生发专用 LoRA 或微调模型，专门优化生发效果。
|
||||||
|
|
||||||
|
**业界相关研究**：
|
||||||
|
- **HairDiffusion (NeurIPS 2024)**：把发型/发色编辑重新定义为 latent diffusion inpainting 任务
|
||||||
|
- **AnyBald (WACV 2026)**：diffusion inpainting + 可学习文本提示做头发去除（反向任务，架构可借鉴）
|
||||||
|
- **Stable-Hair (arXiv 2024)**：diffusion 真实世界发型迁移
|
||||||
|
|
||||||
|
**适用场景**：当方案 A/B 用通用 SD 的生发效果不够好（发丝太糊/不真实），且业务对质量要求高时。
|
||||||
|
|
||||||
|
**优势**：
|
||||||
|
- ✅ 可针对性优化生发质量，效果上限最高
|
||||||
|
|
||||||
|
**劣势**：
|
||||||
|
- ❌ **需要大量配对训练数据**（稀疏发→浓密发），数据收集是最大瓶颈
|
||||||
|
- ❌ 训练+调优周期长（2-4 周）
|
||||||
|
- ❌ 代码库的 LoRA 训练流程（kohya）是针对"发型模板"的，生发是不同任务，需重新组织数据
|
||||||
|
|
||||||
|
**预估工作量**：3-4 周（含数据收集、训练、调优）
|
||||||
|
|
||||||
|
**结论**：**作为后续优化手段备选**，不作为首期方案。先用方案 A/B 跑通，效果不达标再考虑。
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 三、推荐方案与实施路径
|
||||||
|
|
||||||
|
### 3.1 推荐：方案 A 为主，方案 B 为增强
|
||||||
|
|
||||||
|
**首期实现方案 A**（复用 webui inpainting），最快验证可行性；**若效果不足，叠加方案 B**（ControlNet canny）。方案 C/D 暂不引入。
|
||||||
|
|
||||||
|
理由：
|
||||||
|
1. 方案 A 的 90% 组件已存在且验证过，风险最低、见效最快
|
||||||
|
2. 代码库的 `Generator_BaldSeg` 能自动定位缺发区，免去了 SAM 的额外部署
|
||||||
|
3. 与现有服务架构一致，可作为 hair_service_sd 的新接口无缝接入
|
||||||
|
|
||||||
|
### 3.2 推荐的功能管线（方案A）
|
||||||
|
|
||||||
|
```
|
||||||
|
输入: 人头图 + (可选)手绘mask
|
||||||
|
│
|
||||||
|
├─ 若无手绘mask(自动模式):
|
||||||
|
│ 1. RetinaFace 检测人脸 → 1k关键点
|
||||||
|
│ 2. Generator_BaldSeg 分割 → 提取"头皮/发际线"区
|
||||||
|
│ 3. 或: 额头凸包mask(关键点) - 现有头发mask(matting) = 缺发带
|
||||||
|
│ 4. mask 形态学处理(dilate/erode 羽化)
|
||||||
|
│
|
||||||
|
├─ 若有手绘mask(手动模式):
|
||||||
|
│ 直接使用，轻微 dilate 羽化
|
||||||
|
│
|
||||||
|
▼
|
||||||
|
mask_blur=10 的 inpainting:
|
||||||
|
webui_img2img(
|
||||||
|
img = 原图,
|
||||||
|
mask_img = 生发mask,
|
||||||
|
tag = "thick natural hair, detailed hair strands, hairline",
|
||||||
|
denoising_strength = 0.6,
|
||||||
|
is_hr = True
|
||||||
|
)
|
||||||
|
│
|
||||||
|
▼
|
||||||
|
泊松融合贴回(cv2.seamlessClone + GaussianBlur)
|
||||||
|
│
|
||||||
|
▼
|
||||||
|
输出: 生发结果图
|
||||||
|
```
|
||||||
|
|
||||||
|
### 3.3 接口设计建议
|
||||||
|
|
||||||
|
```
|
||||||
|
POST /api/hairGrow/v1
|
||||||
|
{
|
||||||
|
"img": "data:image/jpeg;base64,...", # 人头图
|
||||||
|
"mask": "data:image/png;base64,...", # 可选: 手绘mask(白色=生发区)
|
||||||
|
"auto_mask": true, # 可选: 无mask时自动检测缺发区
|
||||||
|
"userId": "xxx",
|
||||||
|
"output_format": "base64"
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
返回（同换发型）：
|
||||||
|
```json
|
||||||
|
{"state": 0, "msg": "success", "data": "<base64结果图>", "task_id": "..."}
|
||||||
|
```
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 四、代码库可复用能力清单（详细）
|
||||||
|
|
||||||
|
以下能力**全部已在本机部署中验证可用**，无需新增模型：
|
||||||
|
|
||||||
|
| 能力 | 模块 | 文件:行号 | 权重 | 对生发的作用 |
|
||||||
|
|------|------|-----------|------|-------------|
|
||||||
|
| 人脸检测 | RetinaFace | models/detector.py:148 | Resnet50_Final.pth (109M) | 入口，定位人脸 |
|
||||||
|
| 1k关键点 | MomocvFaceAlignment1K | utils/MomocvFaceAlignment1K.py:404 | face_alignment_1k.pth (49M) | **发际线定位** |
|
||||||
|
| 头发抠图 | Generator_Matte | hair_matting/Generator_Matte.py:81 | deeplabv3+gca (535M) | 现有头发边界 |
|
||||||
|
| **头皮分割** | **Generator_BaldSeg_5c** | **models/Generator_BaldSeg.py:45** | **ori_hair_checkpoint (134M)** | **自动定位缺发区** |
|
||||||
|
| 人脸分割 | FaceSeg | core/faceseg/face_seg.py:32 | faceseg_20210927 (229M) | 排除五官区 |
|
||||||
|
| SD inpainting | webui_img2img | gen_super_image.py:416 | v1-5 SD模型 (4.2G) | **核心生成** |
|
||||||
|
| 额头mask构造 | draw_bigger_hull_mask | utils/landmark_processor.py:1079 | - | 关键点→额头区 |
|
||||||
|
| 泊松融合 | cv2.seamlessClone | core/hairstyle_model.py:1380 | - | 边缘融合 |
|
||||||
|
| 形态学处理 | cv2.dilate/erode | utils/landmark_processor.py 多处 | - | mask羽化 |
|
||||||
|
| ControlNet canny | sd-webui-controlnet | webui扩展(已装) | control_v11p_sd15_canny | 发丝走向约束(方案B) |
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 五、风险与质量评估
|
||||||
|
|
||||||
|
| 风险点 | 严重度 | 应对 |
|
||||||
|
|--------|--------|------|
|
||||||
|
| SD1.5生发发丝偏糊 | 中 | 方案B加canny约束；或后续换SDXL inpainting模型(需升级webui) |
|
||||||
|
| 自动mask不准(误判缺发区) | 中 | 首期优先支持手绘mask；自动模式作为可选增强 |
|
||||||
|
| 生发破坏五官 | 高 | mask构造时必须减去FaceSeg的人脸区；inpainting的mask严格限定 |
|
||||||
|
| 多次结果不一致 | 低 | 固定seed(代码已用seed=123456789)；可返回多张供选 |
|
||||||
|
| 边缘拼接痕迹 | 低 | 泊松融合+GaussianBlur已有成熟方案 |
|
||||||
|
|
||||||
|
**质量预期**：方案A对手绘mask场景（用户明确指定生发区）效果应较好；对稀疏加密（发际线整体下移）效果中等；对完全光头大区域生发效果有限（建议用方案D专项优化）。
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 六、实施建议（分阶段）
|
||||||
|
|
||||||
|
**阶段一（2-3天）：MVP 手动 mask 生发**
|
||||||
|
- 实现 `/api/hairGrow/v1` 接口，接收手绘 mask
|
||||||
|
- 复用 `webui_img2img`，prompt 写生发描述
|
||||||
|
- 加泊松融合后处理
|
||||||
|
- 用几张发际线高的测试图调参（denoising/mask_blur/prompt）
|
||||||
|
|
||||||
|
**阶段二（2-3天）：自动 mask + ControlNet**
|
||||||
|
- 接入 `Generator_BaldSeg` 自动检测缺发区
|
||||||
|
- 叠加 ControlNet canny 约束发丝走向（方案B）
|
||||||
|
- 对比有无 canny 的效果差异
|
||||||
|
|
||||||
|
**阶段三（按需）：质量优化**
|
||||||
|
- 若效果不足：收集生发配对数据，训练专项 LoRA（方案D）
|
||||||
|
- 或升级到 SDXL inpainting 模型提升发丝真实感
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 七、参考资料
|
||||||
|
|
||||||
|
**学术论文**：
|
||||||
|
- [HairDiffusion: Vivid Multi-Colored Hair Editing via Latent Diffusion (NeurIPS 2024)](https://openreview.net/forum?id=UQflshLbZv) — 发型编辑的 inpainting 范式
|
||||||
|
- [Stable-Hair: Real-World Hair Transfer via Diffusion Model (arXiv 2024)](https://arxiv.org/html/2407.14078v1) — 扩散发型迁移
|
||||||
|
- [Segment Anything Meets Image Inpainting (arXiv:2304.06790)](https://arxiv.org/abs/2304.06790) — SAM click-and-fill 范式
|
||||||
|
|
||||||
|
**实践工具/社区**：
|
||||||
|
- [sd-webui-inpaint-anything (GitHub)](https://github.com/Uminosachi/sd-webui-inpaint-anything) — SAM + inpainting webui 扩展
|
||||||
|
- [Stable Diffusion Inpainting with SAM (HuggingFace Space)](https://huggingface.co/spaces/Sanshruth/Stable-Diffusion-Inpainting_with_SAM) — 在线 demo
|
||||||
|
- [Reddit: 用 inpainting 添加头发的技巧](https://www.reddit.com/r/StableDiffusion/comments/15x9yvs/) — 社区经验：先预填底色再 inpainting
|
||||||
|
- [Evoto Hair Editor](https://www.evoto.ai/features/hair-part) — 商业产品效果参考（稀疏加密/发际线调整）
|
||||||
@@ -0,0 +1,69 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""安全修复 conda 环境内的硬编码路径前缀。
|
||||||
|
只处理文本文件（通过二进制检测），跳过 .so/.pyc/权重等二进制文件。
|
||||||
|
"""
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
|
||||||
|
OLD_PREFIXES = [
|
||||||
|
b"/home/szlc/miniconda3",
|
||||||
|
b"/usr/local/miniconda3",
|
||||||
|
]
|
||||||
|
NEW_PREFIX = b"/home/xsl/miniconda3"
|
||||||
|
|
||||||
|
def is_probably_binary(filepath):
|
||||||
|
"""通过文件扩展名和内容判断是否为二进制文件"""
|
||||||
|
# 明确的文本扩展名直接处理
|
||||||
|
text_exts = ('.py', '.sh', '.txt', '.cfg', '.ini', '.json', '.yaml', '.yml',
|
||||||
|
'.cmake', '.make', '.pc', '.la', '.template', '.desktop',
|
||||||
|
'.service', '.conf', '.md', '.rst', '.in', '.bashrc', '.profile')
|
||||||
|
binary_exts = ('.so', '.pyc', '.pyo', '.a', '.o', '.dylib', '.dll',
|
||||||
|
'.pth', '.pt', '.bin', '.safetensors', '.ckpt', '.npz',
|
||||||
|
'.npy', '.pkl', '.zip', '.tar', '.gz', '.bz2', '.png',
|
||||||
|
'.jpg', '.jpeg', '.woff', '.ttf', '.eot', '.ico')
|
||||||
|
if filepath.endswith(binary_exts):
|
||||||
|
return True
|
||||||
|
if filepath.endswith(text_exts):
|
||||||
|
return False
|
||||||
|
# 无明确扩展名：读取前 8KB 检测是否含 null 字节
|
||||||
|
try:
|
||||||
|
with open(filepath, 'rb') as f:
|
||||||
|
chunk = f.read(8192)
|
||||||
|
return b'\x00' in chunk
|
||||||
|
except Exception:
|
||||||
|
return True
|
||||||
|
|
||||||
|
def fix_env(env_dir):
|
||||||
|
fixed = 0
|
||||||
|
scanned = 0
|
||||||
|
for root, dirs, files in os.walk(env_dir):
|
||||||
|
# 跳过明显的二进制/缓存目录
|
||||||
|
dirs[:] = [d for d in dirs if d not in ('__pycache__', '.git')]
|
||||||
|
for fname in files:
|
||||||
|
fpath = os.path.join(root, fname)
|
||||||
|
scanned += 1
|
||||||
|
try:
|
||||||
|
if is_probably_binary(fpath):
|
||||||
|
continue
|
||||||
|
with open(fpath, 'rb') as f:
|
||||||
|
content = f.read()
|
||||||
|
original = content
|
||||||
|
for old in OLD_PREFIXES:
|
||||||
|
content = content.replace(old, NEW_PREFIX)
|
||||||
|
if content != original:
|
||||||
|
with open(fpath, 'wb') as f:
|
||||||
|
f.write(content)
|
||||||
|
fixed += 1
|
||||||
|
except (PermissionError, OSError):
|
||||||
|
continue
|
||||||
|
print(f" 扫描 {scanned} 文件, 修改 {fixed} 文件")
|
||||||
|
return fixed
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
for env in sys.argv[1:]:
|
||||||
|
env_path = f"/home/xsl/miniconda3/envs/{env}"
|
||||||
|
if os.path.isdir(env_path):
|
||||||
|
print(f"[{env}] 修复中...")
|
||||||
|
fix_env(env_path)
|
||||||
|
else:
|
||||||
|
print(f"[{env}] 目录不存在，跳过")
|
||||||
@@ -0,0 +1,93 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
"""为测试图自动生成发际线/额头区域的 mask（模拟用户手画遮罩）。
|
||||||
|
用 RetinaFace + 1k关键点定位额头，生成该区域的白色 mask。
|
||||||
|
"""
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd"
|
||||||
|
sys.path.insert(0, HAIR_SERVICE_DIR)
|
||||||
|
os.environ.setdefault("HF_HUB_OFFLINE", "1")
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from utils.landmark_processor import get_max_rect
|
||||||
|
|
||||||
|
|
||||||
|
def gen_hairline_mask(img_path, mask_path):
|
||||||
|
"""用关键点定位额头，生成发际线区域 mask"""
|
||||||
|
from models.detector import RetinaFaceDetector
|
||||||
|
from utils.MomocvFaceAlignment1K import MomocvFaceAlignment1K
|
||||||
|
|
||||||
|
img = cv2.imread(img_path)
|
||||||
|
if img is None:
|
||||||
|
print(f"❌ 读取失败: {img_path}")
|
||||||
|
return False
|
||||||
|
|
||||||
|
print(f"原图: {img_path} shape={img.shape}")
|
||||||
|
gpu_id = 0
|
||||||
|
|
||||||
|
# 1. RetinaFace 检测
|
||||||
|
detector = RetinaFaceDetector(gpu_id=gpu_id)
|
||||||
|
dets, landms = detector.forward(img)
|
||||||
|
if len(dets) == 0:
|
||||||
|
print("❌ 未检测到人脸")
|
||||||
|
return False
|
||||||
|
max_idx = get_max_rect(dets)
|
||||||
|
det = dets[max_idx]
|
||||||
|
print(f"检测到人脸: {det}")
|
||||||
|
|
||||||
|
# 2. 1k 关键点
|
||||||
|
aligner = MomocvFaceAlignment1K(gpu_id=gpu_id)
|
||||||
|
landmarks_1k = aligner.stable_forward(img.copy(), [det])
|
||||||
|
if landmarks_1k is None or len(landmarks_1k) == 0:
|
||||||
|
print("❌ 关键点检测失败")
|
||||||
|
return False
|
||||||
|
pts = landmarks_1k[0] # (1000, 2)
|
||||||
|
print(f"1k关键点 shape: {pts.shape}")
|
||||||
|
|
||||||
|
# 3. 构造发际线/额头 mask
|
||||||
|
# 1k点的前312点是整脸外轮廓（含发际线）
|
||||||
|
# 额头区：用眉骨以上的外轮廓点 + 发际线点构造凸包
|
||||||
|
h, w = img.shape[:2]
|
||||||
|
mask = np.zeros((h, w), dtype=np.uint8)
|
||||||
|
|
||||||
|
# 取脸部上半部分轮廓点（发际线到太阳穴）构造额头区域
|
||||||
|
# 索引 0-311 是脸部外轮廓，取上半部分（额头/太阳穴）
|
||||||
|
forehead_pts = []
|
||||||
|
# 发际线区域：取轮廓的上半段 + 稍微往下扩展
|
||||||
|
for i in range(0, 312, 3): # 隔点采样，减少密度
|
||||||
|
pt = pts[i]
|
||||||
|
# 只取上半脸（y < 图像中线偏上）
|
||||||
|
if pt[1] < h * 0.55:
|
||||||
|
forehead_pts.append(pt)
|
||||||
|
# 加入额头中心区域一些点，确保覆盖额头
|
||||||
|
forehead_pts = np.array(forehead_pts, dtype=np.int32)
|
||||||
|
|
||||||
|
if len(forehead_pts) < 3:
|
||||||
|
print("❌ 额头点不足")
|
||||||
|
return False
|
||||||
|
|
||||||
|
# 凸包填充
|
||||||
|
hull = cv2.convexHull(forehead_pts)
|
||||||
|
cv2.fillConvexPoly(mask, hull, 255)
|
||||||
|
|
||||||
|
# 膨胀让区域稍大一点（模拟用户涂抹）
|
||||||
|
mask = cv2.dilate(mask, np.ones((5, 5), np.uint8), iterations=2)
|
||||||
|
|
||||||
|
# 统计 mask 区域
|
||||||
|
nonzero = cv2.countNonZero(mask)
|
||||||
|
print(f"mask 生成完成: 白色区域={nonzero}px ({nonzero/(h*w)*100:.1f}%)")
|
||||||
|
cv2.imwrite(mask_path, mask)
|
||||||
|
print(f"✅ mask 已保存: {mask_path}")
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
img_path = sys.argv[1] if len(sys.argv) > 1 else \
|
||||||
|
"/home/xsl/change_hair/project/data/userImage/8488902485_20250630055547.jpg"
|
||||||
|
mask_path = sys.argv[2] if len(sys.argv) > 2 else \
|
||||||
|
"/home/xsl/change_hair/project/logs/hairgrow_test_mask.png"
|
||||||
|
gen_hairline_mask(img_path, mask_path)
|
||||||
@@ -0,0 +1,122 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
"""批量生成发型预览图
|
||||||
|
|
||||||
|
用 boy.png 给男款发型、girl.png 给女款发型跑换发型服务，
|
||||||
|
生成每个发型套在标准脸型上的效果图，存为缩略图供前端展示。
|
||||||
|
|
||||||
|
用法: python gen_hairstyle_previews.py [--limit N] (limit用于测试)
|
||||||
|
"""
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
import json
|
||||||
|
import time
|
||||||
|
import base64
|
||||||
|
import requests
|
||||||
|
|
||||||
|
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd"
|
||||||
|
sys.path.insert(0, HAIR_SERVICE_DIR)
|
||||||
|
os.chdir(HAIR_SERVICE_DIR)
|
||||||
|
|
||||||
|
BOY_IMG = "/home/xsl/change_hair/images/boy.png"
|
||||||
|
GIRL_IMG = "/home/xsl/change_hair/images/girl.png"
|
||||||
|
HAIRSTYLE_DIR = "/home/xsl/change_hair/project/data/ref_hairstyle"
|
||||||
|
TRAIN_DIR = "/home/xsl/change_hair/data/train_material"
|
||||||
|
PREVIEW_DIR = "/home/xsl/change_hair/hair_grow_service/static/previews" # 预览图存储目录
|
||||||
|
SWAP_API = "http://127.0.0.1:8801/api/swapHair/v1"
|
||||||
|
|
||||||
|
|
||||||
|
def load_hairstyles():
|
||||||
|
"""返回 {hair_id: gender}"""
|
||||||
|
styles = {}
|
||||||
|
for hid in os.listdir(HAIRSTYLE_DIR):
|
||||||
|
cfg = os.path.join(HAIRSTYLE_DIR, hid, "config.json")
|
||||||
|
lora = os.path.join(TRAIN_DIR, hid, "model", "hairstyle_hd_lora.safetensors")
|
||||||
|
if not (os.path.exists(cfg) and os.path.exists(lora)):
|
||||||
|
continue
|
||||||
|
g = json.load(open(cfg)).get("gender", "?")
|
||||||
|
styles[hid] = g
|
||||||
|
return styles
|
||||||
|
|
||||||
|
|
||||||
|
def gen_one(hair_id, gender, img_path):
|
||||||
|
"""给一个发型生成预览图，返回是否成功"""
|
||||||
|
out_path = os.path.join(PREVIEW_DIR, f"{hair_id}.jpg")
|
||||||
|
if os.path.exists(out_path):
|
||||||
|
return True, "已存在跳过"
|
||||||
|
|
||||||
|
with open(img_path, "rb") as f:
|
||||||
|
img_b64 = "data:image/png;base64," + base64.b64encode(f.read()).decode()
|
||||||
|
|
||||||
|
payload = {
|
||||||
|
"hair_id": hair_id,
|
||||||
|
"task_id": f"preview_{hair_id}",
|
||||||
|
"is_hr": "false", # 非高清，快一点
|
||||||
|
"user_img_path": img_b64,
|
||||||
|
"output_format": "base64"
|
||||||
|
}
|
||||||
|
try:
|
||||||
|
r = requests.post(SWAP_API, json=payload, timeout=180)
|
||||||
|
d = r.json()
|
||||||
|
if d.get("state") == 0 and d.get("data"):
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
img = cv2.imdecode(np.frombuffer(base64.b64decode(d["data"]), np.uint8), cv2.IMREAD_COLOR)
|
||||||
|
# 缩小到 300x400 当缩略图
|
||||||
|
if img is not None:
|
||||||
|
h, w = img.shape[:2]
|
||||||
|
scale = min(300 / w, 400 / h)
|
||||||
|
img = cv2.resize(img, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA)
|
||||||
|
cv2.imwrite(out_path, img, [cv2.IMWRITE_JPEG_QUALITY, 85])
|
||||||
|
return True, "生成成功"
|
||||||
|
return False, f"state={d.get('state')} msg={d.get('msg','')}"
|
||||||
|
except Exception as e:
|
||||||
|
return False, str(e)
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
import cv2
|
||||||
|
limit = None
|
||||||
|
if "--limit" in sys.argv:
|
||||||
|
limit = int(sys.argv[sys.argv.index("--limit") + 1])
|
||||||
|
|
||||||
|
os.makedirs(PREVIEW_DIR, exist_ok=True)
|
||||||
|
styles = load_hairstyles()
|
||||||
|
print(f"共 {len(styles)} 个发型需要生成预览图")
|
||||||
|
|
||||||
|
# 按性别选基准图
|
||||||
|
boy_styles = [(h, g) for h, g in styles.items() if g == "boy"]
|
||||||
|
girl_styles = [(h, g) for h, g in styles.items() if g == "girl"]
|
||||||
|
print(f" boy: {len(boy_styles)} (用 boy.png)")
|
||||||
|
print(f" girl: {len(girl_styles)} (用 girl.png)")
|
||||||
|
|
||||||
|
tasks = boy_styles + girl_styles
|
||||||
|
if limit:
|
||||||
|
tasks = tasks[:limit]
|
||||||
|
print(f" [测试模式] 只处理前 {limit} 个")
|
||||||
|
|
||||||
|
ok, fail, skip = 0, 0, 0
|
||||||
|
t0 = time.time()
|
||||||
|
for i, (hair_id, gender) in enumerate(tasks):
|
||||||
|
img_path = BOY_IMG if gender == "boy" else GIRL_IMG
|
||||||
|
out_path = os.path.join(PREVIEW_DIR, f"{hair_id}.jpg")
|
||||||
|
if os.path.exists(out_path):
|
||||||
|
skip += 1
|
||||||
|
continue
|
||||||
|
elapsed = time.time() - t0
|
||||||
|
eta = elapsed / max(ok + fail, 1) * (len(tasks) - i) if (ok + fail) > 0 else 0
|
||||||
|
print(f"[{i+1}/{len(tasks)}] {hair_id}({gender})...", end=" ", flush=True)
|
||||||
|
success, msg = gen_one(hair_id, gender, img_path)
|
||||||
|
if success:
|
||||||
|
ok += 1
|
||||||
|
print(f"✓ ({msg})")
|
||||||
|
else:
|
||||||
|
fail += 1
|
||||||
|
print(f"✗ ({msg})")
|
||||||
|
|
||||||
|
print(f"\n=== 完成: 成功{ok} 失败{fail} 跳过{skip} 总耗时{time.time()-t0:.0f}s ===")
|
||||||
|
print(f"预览图目录: {PREVIEW_DIR}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -0,0 +1,99 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
"""为新发型生成 hair_template_material（first##_matting.png + .pkl）
|
||||||
|
|
||||||
|
换发型功能4 需要：
|
||||||
|
hair_template_material/<hair_id>/first##<name>.png 原图
|
||||||
|
hair_template_material/<hair_id>/first##<name>_matting.png 头发抠图mask
|
||||||
|
hair_template_material/<hair_id>/first##<name>.pkl 1k关键点
|
||||||
|
|
||||||
|
用法:
|
||||||
|
cd /home/xsl/change_hair/project/hair_service_sd
|
||||||
|
python gen_template_material.py --hair-id new_test_001 --img /path/to/template.jpg
|
||||||
|
"""
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
import ssl
|
||||||
|
import uuid
|
||||||
|
import pickle
|
||||||
|
import argparse
|
||||||
|
|
||||||
|
ssl._create_default_https_context = ssl._create_unverified_context
|
||||||
|
HAIR_SERVICE_DIR = "/home/xsl/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")
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
|
||||||
|
_orig_torch_load = torch.load
|
||||||
|
def _patched_torch_load(*args, **kwargs):
|
||||||
|
if 'map_location' in kwargs and callable(kwargs['map_location']) and not isinstance(kwargs['map_location'], str):
|
||||||
|
kwargs['map_location'] = 'cpu'
|
||||||
|
return _orig_torch_load(*args, **kwargs)
|
||||||
|
torch.load = _patched_torch_load
|
||||||
|
|
||||||
|
from models.detector import RetinaFaceDetector
|
||||||
|
from models.MomocvFaceAlignment1K import MomocvFaceAlignment1K
|
||||||
|
from utils.landmark_processor import get_max_rect
|
||||||
|
from core.process_modules import Generator_Matte
|
||||||
|
from common.logger import config
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser()
|
||||||
|
parser.add_argument("--hair-id", required=True)
|
||||||
|
parser.add_argument("--img", required=True, help="模板图(发型图)")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
template_dir = config.get('default', 'hair_template_material_dir')
|
||||||
|
out_dir = os.path.join(template_dir, args.hair_id)
|
||||||
|
os.makedirs(out_dir, exist_ok=True)
|
||||||
|
|
||||||
|
img = cv2.imread(args.img)
|
||||||
|
if img is None:
|
||||||
|
print(f"✗ 读取失败: {args.img}"); sys.exit(1)
|
||||||
|
print(f"模板图: {img.shape}")
|
||||||
|
|
||||||
|
# 1. 人脸检测 + 1k关键点
|
||||||
|
print("加载模型...")
|
||||||
|
detector = RetinaFaceDetector(gpu_id=0)
|
||||||
|
aligner = MomocvFaceAlignment1K(gpu_id=0)
|
||||||
|
matte_gen = Generator_Matte(gpu=True, device_id=0)
|
||||||
|
print("模型加载完成")
|
||||||
|
|
||||||
|
dets, landms = detector.forward(img)
|
||||||
|
if len(dets) == 0:
|
||||||
|
print("✗ 未检测到人脸"); sys.exit(1)
|
||||||
|
det = dets[get_max_rect(dets)]
|
||||||
|
landmarks_1k = aligner.stable_forward(img.copy(), [det])
|
||||||
|
pt1k = landmarks_1k[0]
|
||||||
|
|
||||||
|
# 2. 头发抠图
|
||||||
|
with torch.no_grad():
|
||||||
|
_, alpha, _ = matte_gen.matte_inference(img, pt1k)
|
||||||
|
if alpha.shape != img.shape[:2]:
|
||||||
|
alpha = cv2.resize(alpha, (img.shape[1], img.shape[0]))
|
||||||
|
|
||||||
|
# 3. 保存三个文件（first##前缀）
|
||||||
|
file_uuid = f"first##{uuid.uuid4()}"
|
||||||
|
png_path = os.path.join(out_dir, f"{file_uuid}.png")
|
||||||
|
matting_path = os.path.join(out_dir, f"{file_uuid}_matting.png")
|
||||||
|
pkl_path = os.path.join(out_dir, f"{file_uuid}.pkl")
|
||||||
|
|
||||||
|
cv2.imwrite(png_path, img)
|
||||||
|
cv2.imwrite(matting_path, alpha)
|
||||||
|
with open(pkl_path, 'wb') as f:
|
||||||
|
# 存 numpy array（不能用 tolist()，代码用 landmarks[index, :] 索引需要 ndarray）
|
||||||
|
pickle.dump({'human_pt1k': np.asarray(pt1k, dtype=np.float32)}, f)
|
||||||
|
|
||||||
|
print(f"✓ 生成完成:")
|
||||||
|
print(f" {png_path}")
|
||||||
|
print(f" {matting_path}")
|
||||||
|
print(f" {pkl_path}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -0,0 +1,89 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
"""区域生发 - 独立验证脚本
|
||||||
|
|
||||||
|
不依赖 hair_service 主服务，直接调用 hair_grow.hair_grow()。
|
||||||
|
需要在 hair_service_sd 目录下运行（依赖其模块导入）。
|
||||||
|
|
||||||
|
用法:
|
||||||
|
cd /home/xsl/change_hair/project/hair_service_sd
|
||||||
|
/home/xsl/miniconda3/envs/my_hair/bin/python /home/xsl/change_hair/hair_grow_cli.py \
|
||||||
|
--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 \
|
||||||
|
--img test.jpg --mask mask.png --compare
|
||||||
|
"""
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
import cv2
|
||||||
|
import argparse
|
||||||
|
|
||||||
|
# 把 hair_service_sd 加入路径，使其模块可被导入
|
||||||
|
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd"
|
||||||
|
sys.path.insert(0, HAIR_SERVICE_DIR)
|
||||||
|
|
||||||
|
# 设置离线模式（本机无法访问 huggingface.co）
|
||||||
|
os.environ.setdefault("HF_HUB_OFFLINE", "1")
|
||||||
|
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
|
||||||
|
|
||||||
|
|
||||||
|
def run_once(img_path, mask_path, strength, out_path):
|
||||||
|
"""单次生发"""
|
||||||
|
from hair_grow import hair_grow
|
||||||
|
|
||||||
|
img = cv2.imread(img_path)
|
||||||
|
mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
|
||||||
|
if img is None:
|
||||||
|
print(f"❌ 无法读取图片: {img_path}")
|
||||||
|
return False
|
||||||
|
if mask is None:
|
||||||
|
print(f"❌ 无法读取mask: {mask_path}")
|
||||||
|
return False
|
||||||
|
|
||||||
|
print(f"输入: img={img.shape}, mask={mask.shape}, strength={strength}")
|
||||||
|
result = hair_grow(img, mask, strength=strength)
|
||||||
|
|
||||||
|
cv2.imwrite(out_path, result)
|
||||||
|
print(f"✅ 生发结果已保存: {out_path}")
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
def run_compare(img_path, mask_path, out_dir):
|
||||||
|
"""三档强度对比"""
|
||||||
|
os.makedirs(out_dir, exist_ok=True)
|
||||||
|
for strength in [0.2, 0.5, 0.8]:
|
||||||
|
out_path = os.path.join(out_dir, f"result_s{strength}.jpg")
|
||||||
|
print(f"\n{'='*50}")
|
||||||
|
print(f"生发强度 strength={strength}")
|
||||||
|
print(f"{'='*50}")
|
||||||
|
if not run_once(img_path, mask_path, strength, out_path):
|
||||||
|
return False
|
||||||
|
print(f"\n✅ 三档对比完成，结果在: {out_dir}")
|
||||||
|
print(" - result_s0.2.jpg (低强度，轻微生发)")
|
||||||
|
print(" - result_s0.5.jpg (中强度，明显生发)")
|
||||||
|
print(" - result_s0.8.jpg (高强度，浓密生发)")
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser(description="区域生发验证脚本")
|
||||||
|
parser.add_argument("--img", required=True, help="人头像图片路径")
|
||||||
|
parser.add_argument("--mask", required=True, help="遮罩图片路径(白=生发区)")
|
||||||
|
parser.add_argument("--strength", type=float, default=0.5,
|
||||||
|
help="生发强度 0.1~1.0 (默认0.5)")
|
||||||
|
parser.add_argument("-o", "--output", default="result_hairgrow.jpg",
|
||||||
|
help="输出图片路径")
|
||||||
|
parser.add_argument("--compare", action="store_true",
|
||||||
|
help="跑三档强度(0.2/0.5/0.8)对比")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
if args.compare:
|
||||||
|
out_dir = os.path.dirname(args.output) or "."
|
||||||
|
run_compare(args.img, args.mask, out_dir)
|
||||||
|
else:
|
||||||
|
run_once(args.img, args.mask, args.strength, args.output)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -0,0 +1,68 @@
|
|||||||
|
# 区域生发测试服务
|
||||||
|
|
||||||
|
提供带「画笔涂抹遮罩」的 Web 测试页面，在用户涂抹的区域内生发。
|
||||||
|
|
||||||
|
## 快速开始
|
||||||
|
|
||||||
|
### 1. 确认 webui 已启动（生发依赖 SD inpainting）
|
||||||
|
```bash
|
||||||
|
ss -tlnp | grep 57860 # 确认 webui 在跑
|
||||||
|
bash /home/xsl/change_hair/start_all.sh status # 看 webui 是否就绪
|
||||||
|
```
|
||||||
|
|
||||||
|
### 2. 启动生发服务（端口 8899）
|
||||||
|
```bash
|
||||||
|
nohup /home/xsl/change_hair/start_hairgrow.sh > /home/xsl/change_hair/project/logs/hair_grow_service.log 2>&1 &
|
||||||
|
```
|
||||||
|
|
||||||
|
### 3. 打开测试页面
|
||||||
|
浏览器访问：
|
||||||
|
```
|
||||||
|
http://172.21.246.70:8899
|
||||||
|
```
|
||||||
|
（如从 Windows 宿主机访问 WSL，用上面地址；或用 localhost:8899）
|
||||||
|
|
||||||
|
### 4. 使用流程
|
||||||
|
1. **上传图片**：点击或拖拽一张人头像照片（正面、发际线/稀疏区清晰）
|
||||||
|
2. **涂抹遮罩**：用画笔在需要生发的区域涂抹（白色 = 生发区）
|
||||||
|
- 左侧可调画笔粗细、切换橡皮、清空
|
||||||
|
3. **调强度**：生发强度滑块（0.1~1.0，默认 0.5）
|
||||||
|
- 低强度：轻微生发，严格保持遮罩边界
|
||||||
|
- 高强度：浓密生发，边缘自然过渡
|
||||||
|
4. **点「生成」**：等待约 5 秒
|
||||||
|
5. **看结果**：原图与生发结果并排对比，可下载
|
||||||
|
|
||||||
|
## 文件说明
|
||||||
|
```
|
||||||
|
hair_grow_service/
|
||||||
|
├── app.py # Flask 服务（端口 8899）
|
||||||
|
├── static/index.html # 测试页面（画笔遮罩 + 结果对比）
|
||||||
|
└── README.md
|
||||||
|
|
||||||
|
start_hairgrow.sh # 启动脚本
|
||||||
|
```
|
||||||
|
|
||||||
|
## API 接口（可单独调用）
|
||||||
|
```
|
||||||
|
POST http://<host>:8899/api/grow
|
||||||
|
Content-Type: application/json
|
||||||
|
|
||||||
|
{
|
||||||
|
"img": "data:image/jpeg;base64,...", # 原图 base64
|
||||||
|
"mask": "data:image/png;base64,...", # 遮罩 base64，白色(255)=生发区，与img同分辨率
|
||||||
|
"strength": 0.5 # 生发强度 0.1~1.0
|
||||||
|
}
|
||||||
|
|
||||||
|
→ {"state": 0, "result": "<base64结果图>"}
|
||||||
|
{"state": -1, "msg": "错误信息"}
|
||||||
|
```
|
||||||
|
|
||||||
|
## 停止服务
|
||||||
|
```bash
|
||||||
|
pkill -f "hair_grow_service/app.py"
|
||||||
|
```
|
||||||
|
|
||||||
|
## 注意事项
|
||||||
|
- 生发服务**依赖 webui(57860)**，webui 未启动会报连接失败
|
||||||
|
- 图片过大（>1024px）会自动缩放并对齐到 8 的倍数
|
||||||
|
- 生发结果是确定性的（固定 seed=123456789），同输入同输出
|
||||||
@@ -0,0 +1,515 @@
|
|||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
"""区域生发服务（走换发型工作流，端口 8888）
|
||||||
|
|
||||||
|
工作流与换发型一致：infer_hairstyle_diy_jy → warpAffine → photo_service+LoRA → webui → 贴回。
|
||||||
|
mask = origin_matting ∪ new_matting ∪ 手绘mask（不减刘海）。
|
||||||
|
|
||||||
|
启动: python app.py
|
||||||
|
页面: http://<host>:8888
|
||||||
|
"""
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
import time
|
||||||
|
import base64
|
||||||
|
import traceback
|
||||||
|
|
||||||
|
# 把 hair_service_sd 加入 path，使其模块可被导入
|
||||||
|
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd"
|
||||||
|
sys.path.insert(0, HAIR_SERVICE_DIR)
|
||||||
|
# 切换到 hair_service_sd 目录，使 common.logger 能读到 config/configure.ini（相对路径）
|
||||||
|
os.chdir(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 cv2
|
||||||
|
import numpy as np
|
||||||
|
from flask import Flask, request, jsonify, send_from_directory
|
||||||
|
from gevent import pywsgi
|
||||||
|
|
||||||
|
PORT = 8888
|
||||||
|
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||||
|
|
||||||
|
app = Flask(__name__, static_folder="static", static_url_path="/static")
|
||||||
|
|
||||||
|
# ===== 模型懒加载（首次请求时初始化，避免启动卡死）=====
|
||||||
|
_hairstyle_process = None
|
||||||
|
_landmark_processor = None
|
||||||
|
|
||||||
|
|
||||||
|
def _get_models():
|
||||||
|
"""懒加载换发型全套模型"""
|
||||||
|
global _hairstyle_process, _landmark_processor
|
||||||
|
if _hairstyle_process is None:
|
||||||
|
print("[init] 加载换发型模型（首次，约30-90秒）...")
|
||||||
|
t0 = time.time()
|
||||||
|
import torch
|
||||||
|
from core.hairstyle_model import HairStyle_Model
|
||||||
|
from utils import landmark_processor
|
||||||
|
_hairstyle_process = HairStyle_Model(gpu=True, use_enhance=True)
|
||||||
|
_landmark_processor = landmark_processor
|
||||||
|
print(f"[init] 模型加载完成，耗时 {time.time()-t0:.1f}s")
|
||||||
|
return _hairstyle_process, _landmark_processor
|
||||||
|
|
||||||
|
|
||||||
|
@app.route("/")
|
||||||
|
def index():
|
||||||
|
return send_from_directory(os.path.join(BASE_DIR, "static"), "index.html")
|
||||||
|
|
||||||
|
|
||||||
|
@app.route("/swap")
|
||||||
|
def swap_page():
|
||||||
|
"""换发型测试页（独立，无遮罩）"""
|
||||||
|
return send_from_directory(os.path.join(BASE_DIR, "static"), "swap.html")
|
||||||
|
|
||||||
|
|
||||||
|
@app.route("/test_new")
|
||||||
|
def test_new_page():
|
||||||
|
"""新发型效果测试页（聚焦本次训练的30款发型）"""
|
||||||
|
return send_from_directory(os.path.join(BASE_DIR, "static"), "test_new.html")
|
||||||
|
|
||||||
|
|
||||||
|
@app.route("/debug")
|
||||||
|
def debug_page():
|
||||||
|
"""换发型调试页（全参数可视化）"""
|
||||||
|
return send_from_directory(os.path.join(BASE_DIR, "static"), "debug.html")
|
||||||
|
|
||||||
|
|
||||||
|
@app.route("/hairline")
|
||||||
|
def hairline_page():
|
||||||
|
"""发际线带重绘实验页"""
|
||||||
|
return send_from_directory(os.path.join(BASE_DIR, "static"), "hairline.html")
|
||||||
|
|
||||||
|
|
||||||
|
@app.route("/manual")
|
||||||
|
def manual_page():
|
||||||
|
"""手绘mask重绘测试页"""
|
||||||
|
return send_from_directory(os.path.join(BASE_DIR, "static"), "manual.html")
|
||||||
|
|
||||||
|
|
||||||
|
@app.route("/api/swap", methods=["POST"])
|
||||||
|
def api_swap_proxy():
|
||||||
|
"""代理换发型请求到 8801（避免前端跨域问题）"""
|
||||||
|
import requests as req
|
||||||
|
try:
|
||||||
|
data = request.json
|
||||||
|
resp = req.post("http://127.0.0.1:8801/api/swapHair/v1",
|
||||||
|
json=data, timeout=600)
|
||||||
|
return jsonify(resp.json())
|
||||||
|
except Exception as e:
|
||||||
|
return jsonify({"state": -1, "msg": f"换发型代理失败: {e}"}), 500
|
||||||
|
|
||||||
|
|
||||||
|
@app.route("/api/swap_viz", methods=["POST"])
|
||||||
|
def api_swap_viz():
|
||||||
|
"""换发型（带可视化中间产物）
|
||||||
|
|
||||||
|
入参 JSON:
|
||||||
|
img: 原图 base64
|
||||||
|
hair_id: 发型ID
|
||||||
|
is_hr: "true"/"false"
|
||||||
|
返回:
|
||||||
|
{state, result: 最终图b64, steps: [{title, desc, images:[{label, b64}]}]}
|
||||||
|
"""
|
||||||
|
t0 = time.time()
|
||||||
|
try:
|
||||||
|
data = request.json
|
||||||
|
img_b64 = data.get("img", "")
|
||||||
|
hair_id = data.get("hair_id", "")
|
||||||
|
is_hr = str(data.get("is_hr", "false")).lower() == "true"
|
||||||
|
strict_mask = data.get("strict_mask", False)
|
||||||
|
|
||||||
|
if not img_b64 or not hair_id:
|
||||||
|
return jsonify({"state": -1, "msg": "img 和 hair_id 不能为空"}), 400
|
||||||
|
|
||||||
|
img = _b64_to_ndarray(img_b64, color=True)
|
||||||
|
if img is None:
|
||||||
|
return jsonify({"state": -1, "msg": "img 解析失败"}), 400
|
||||||
|
|
||||||
|
print(f"[swap_viz] img={img.shape}, hair_id={hair_id}, is_hr={is_hr}, strict_mask={strict_mask}")
|
||||||
|
hairstyle_process, landmark_processor = _get_models()
|
||||||
|
|
||||||
|
from hair_swap_viz import hair_swap_viz
|
||||||
|
task_id = f"swapviz_{int(time.time()*1000)}"
|
||||||
|
result, steps = hair_swap_viz(
|
||||||
|
origin_img=img, hair_id=hair_id,
|
||||||
|
hairstyle_process=hairstyle_process,
|
||||||
|
landmark_processor=landmark_processor,
|
||||||
|
task_id=task_id, is_hr=is_hr, strict_mask=strict_mask)
|
||||||
|
|
||||||
|
_, buf = cv2.imencode(".jpg", result, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||||||
|
result_b64 = base64.b64encode(buf).decode("utf-8")
|
||||||
|
print(f"[swap_viz] 完成，总耗时 {time.time()-t0:.1f}s，{len(steps)}个步骤")
|
||||||
|
return jsonify({"state": 0, "result": result_b64, "steps": steps})
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
print(f"[swap_viz] 失败: {e}")
|
||||||
|
traceback.print_exc()
|
||||||
|
return jsonify({"state": -1, "msg": f"换发型失败: {e}"}), 500
|
||||||
|
|
||||||
|
|
||||||
|
@app.route("/api/swap_debug", methods=["POST"])
|
||||||
|
def api_swap_debug():
|
||||||
|
"""换发型调试接口（全参数可调 + 每步可视化）。
|
||||||
|
|
||||||
|
入参 JSON:
|
||||||
|
img: 原图 base64
|
||||||
|
hair_id: 发型ID
|
||||||
|
# 流程开关
|
||||||
|
cut_bang: bool, 是否减刘海圆（默认 true）
|
||||||
|
strict_mask: bool, 严格按mask贴回（默认 false）
|
||||||
|
seamless_blend: bool, 泊松融合消除接缝（默认 true，仅 strict_mask 时生效）
|
||||||
|
# 尺寸/对齐
|
||||||
|
is_hr: bool, 高清模式（默认 true）
|
||||||
|
dilate_kernel: [x,y], mask膨胀核（默认 [6,18]）
|
||||||
|
# SD 推理（仅 denoising 可调）
|
||||||
|
denoising_strength: float 0~1, 重绘强度（默认 0.6）
|
||||||
|
# 贴回/融合
|
||||||
|
blend_dilate: [x,y], strict贴回mask膨胀（默认 [5,5]）
|
||||||
|
seamless_dilate: [x,y], 泊松融合mask膨胀（默认 [9,9]）
|
||||||
|
返回:
|
||||||
|
{state, result, steps, params}
|
||||||
|
"""
|
||||||
|
t0 = time.time()
|
||||||
|
try:
|
||||||
|
d = request.json
|
||||||
|
img_b64 = d.get("img", "")
|
||||||
|
hair_id = d.get("hair_id", "")
|
||||||
|
if not img_b64 or not hair_id:
|
||||||
|
return jsonify({"state": -1, "msg": "img 和 hair_id 不能为空"}), 400
|
||||||
|
|
||||||
|
img = _b64_to_ndarray(img_b64, color=True)
|
||||||
|
if img is None:
|
||||||
|
return jsonify({"state": -1, "msg": "img 解析失败"}), 400
|
||||||
|
|
||||||
|
# 解析参数（带默认值 + 类型转换）
|
||||||
|
def g(k, default): return d.get(k, default)
|
||||||
|
dk = g("dilate_kernel", [6, 18])
|
||||||
|
bd = g("blend_dilate", [5, 5])
|
||||||
|
sd = g("seamless_dilate", [9, 9])
|
||||||
|
|
||||||
|
print(f"[swap_debug] hair_id={hair_id}, cut_bang={g('cut_bang',True)}, "
|
||||||
|
f"strict_mask={g('strict_mask',False)}, seamless_blend={g('seamless_blend',True)}, "
|
||||||
|
f"is_hr={g('is_hr',True)}, denoising={g('denoising_strength',0.6)}")
|
||||||
|
|
||||||
|
hairstyle_process, landmark_processor = _get_models()
|
||||||
|
from hair_swap_debug import hair_swap_debug
|
||||||
|
task_id = f"swapdbg_{int(time.time()*1000)}"
|
||||||
|
result, steps, params = hair_swap_debug(
|
||||||
|
origin_img=img, hair_id=hair_id,
|
||||||
|
hairstyle_process=hairstyle_process,
|
||||||
|
landmark_processor=landmark_processor,
|
||||||
|
task_id=task_id,
|
||||||
|
cut_bang=bool(g("cut_bang", True)),
|
||||||
|
strict_mask=bool(g("strict_mask", False)),
|
||||||
|
seamless_blend=bool(g("seamless_blend", True)),
|
||||||
|
is_hr=bool(g("is_hr", True)),
|
||||||
|
dilate_kernel=(int(dk[0]), int(dk[1])),
|
||||||
|
denoising_strength=float(g("denoising_strength", 0.6)),
|
||||||
|
blend_dilate=(int(bd[0]), int(bd[1])),
|
||||||
|
seamless_dilate=(int(sd[0]), int(sd[1])),
|
||||||
|
)
|
||||||
|
|
||||||
|
_, buf = cv2.imencode(".jpg", result, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||||||
|
result_b64 = base64.b64encode(buf).decode("utf-8")
|
||||||
|
print(f"[swap_debug] 完成，总耗时 {time.time()-t0:.1f}s")
|
||||||
|
return jsonify({"state": 0, "result": result_b64, "steps": steps, "params": params})
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
print(f"[swap_debug] 失败: {e}")
|
||||||
|
traceback.print_exc()
|
||||||
|
return jsonify({"state": -1, "msg": f"换发型失败: {e}"}), 500
|
||||||
|
|
||||||
|
|
||||||
|
@app.route("/api/swap_hairline", methods=["POST"])
|
||||||
|
def api_swap_hairline():
|
||||||
|
"""发际线带重绘接口（实验性，流程同 swap_debug，仅步骤③改为边界带）。
|
||||||
|
|
||||||
|
比 swap_debug 多一个参数：
|
||||||
|
method: str, 发际线mask方案，可选 boundary_band/mediapipe/landmark_1k/deeplab（默认 mediapipe）
|
||||||
|
band_width: int, 边界带宽度（形态学核大小，默认15，带宽≈2*band_width）
|
||||||
|
其余参数同 swap_debug。
|
||||||
|
"""
|
||||||
|
t0 = time.time()
|
||||||
|
try:
|
||||||
|
d = request.json
|
||||||
|
img_b64 = d.get("img", "")
|
||||||
|
hair_id = d.get("hair_id", "")
|
||||||
|
if not img_b64 or not hair_id:
|
||||||
|
return jsonify({"state": -1, "msg": "img 和 hair_id 不能为空"}), 400
|
||||||
|
|
||||||
|
img = _b64_to_ndarray(img_b64, color=True)
|
||||||
|
if img is None:
|
||||||
|
return jsonify({"state": -1, "msg": "img 解析失败"}), 400
|
||||||
|
|
||||||
|
def g(k, default): return d.get(k, default)
|
||||||
|
dk = g("dilate_kernel", [6, 18])
|
||||||
|
bd = g("blend_dilate", [5, 5])
|
||||||
|
sd = g("seamless_dilate", [9, 9])
|
||||||
|
|
||||||
|
print(f"[swap_hairline] hair_id={hair_id}, method={g('method','mediapipe')}, "
|
||||||
|
f"band_width={g('band_width',15)}, denoising={g('denoising_strength',0.6)}")
|
||||||
|
|
||||||
|
hairstyle_process, landmark_processor = _get_models()
|
||||||
|
from hair_swap_hairline import hair_swap_hairline
|
||||||
|
task_id = f"hairline_{int(time.time()*1000)}"
|
||||||
|
result, steps, params = hair_swap_hairline(
|
||||||
|
origin_img=img, hair_id=hair_id,
|
||||||
|
hairstyle_process=hairstyle_process,
|
||||||
|
landmark_processor=landmark_processor,
|
||||||
|
task_id=task_id,
|
||||||
|
method=str(g("method", "mediapipe")),
|
||||||
|
strict_mask=bool(g("strict_mask", False)),
|
||||||
|
seamless_blend=bool(g("seamless_blend", True)),
|
||||||
|
is_hr=bool(g("is_hr", True)),
|
||||||
|
dilate_kernel=(int(dk[0]), int(dk[1])),
|
||||||
|
denoising_strength=float(g("denoising_strength", 0.6)),
|
||||||
|
blend_dilate=(int(bd[0]), int(bd[1])),
|
||||||
|
seamless_dilate=(int(sd[0]), int(sd[1])),
|
||||||
|
band_width=int(g("band_width", 15)),
|
||||||
|
preview_only=bool(g("preview_only", False)),
|
||||||
|
height_ratio=float(g("height_ratio", 0.432)),
|
||||||
|
width_ratio=float(g("width_ratio", 0.144)),
|
||||||
|
corner_ratio=float(g("corner_ratio", 0.25)),
|
||||||
|
vertical_offset=float(g("vertical_offset", 0.0)),
|
||||||
|
refiner_switch_at=float(g("refiner_switch_at", 0.5)),
|
||||||
|
)
|
||||||
|
|
||||||
|
_, buf = cv2.imencode(".jpg", result, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||||||
|
result_b64 = base64.b64encode(buf).decode("utf-8")
|
||||||
|
print(f"[swap_hairline] 完成，总耗时 {time.time()-t0:.1f}s")
|
||||||
|
return jsonify({"state": 0, "result": result_b64, "steps": steps, "params": params})
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
print(f"[swap_hairline] 失败: {e}")
|
||||||
|
traceback.print_exc()
|
||||||
|
return jsonify({"state": -1, "msg": f"发际线带重绘失败: {e}"}), 500
|
||||||
|
|
||||||
|
|
||||||
|
@app.route("/api/swap_manual", methods=["POST"])
|
||||||
|
def api_swap_manual():
|
||||||
|
"""手绘mask版换发型接口（重绘区域完全由用户手绘mask决定）。
|
||||||
|
|
||||||
|
入参 JSON:
|
||||||
|
img: 原图 base64
|
||||||
|
mask: 用户手绘mask base64（白色=重绘区，与img同分辨率）
|
||||||
|
hair_id: 发型ID
|
||||||
|
其余参数同 swap_debug（strict_mask/seamless_blend/is_hr/dilate_kernel/
|
||||||
|
denoising_strength/blend_dilate/seamless_dilate）
|
||||||
|
返回: {state, result, steps, params}
|
||||||
|
"""
|
||||||
|
t0 = time.time()
|
||||||
|
try:
|
||||||
|
d = request.json
|
||||||
|
img_b64 = d.get("img", "")
|
||||||
|
mask_b64 = d.get("mask", "")
|
||||||
|
hair_id = d.get("hair_id", "")
|
||||||
|
if not img_b64 or not mask_b64:
|
||||||
|
return jsonify({"state": -1, "msg": "img 和 mask 不能为空（请先手绘重绘区域）"}), 400
|
||||||
|
if not hair_id:
|
||||||
|
return jsonify({"state": -1, "msg": "hair_id 不能为空"}), 400
|
||||||
|
|
||||||
|
img = _b64_to_ndarray(img_b64, color=True)
|
||||||
|
hand_mask = _b64_to_ndarray(mask_b64, color=False)
|
||||||
|
if img is None:
|
||||||
|
return jsonify({"state": -1, "msg": "img 解析失败"}), 400
|
||||||
|
if hand_mask is None:
|
||||||
|
return jsonify({"state": -1, "msg": "mask 解析失败"}), 400
|
||||||
|
|
||||||
|
def g(k, default): return d.get(k, default)
|
||||||
|
dk = g("dilate_kernel", [6, 18])
|
||||||
|
bd = g("blend_dilate", [5, 5])
|
||||||
|
sd = g("seamless_dilate", [9, 9])
|
||||||
|
|
||||||
|
print(f"[swap_manual] hair_id={hair_id}, img={img.shape}, mask={hand_mask.shape}, "
|
||||||
|
f"strict_mask={g('strict_mask',False)}, denoising={g('denoising_strength',0.6)}")
|
||||||
|
|
||||||
|
hairstyle_process, landmark_processor = _get_models()
|
||||||
|
from hair_swap_manual import hair_swap_manual
|
||||||
|
task_id = f"manual_{int(time.time()*1000)}"
|
||||||
|
result, steps, params = hair_swap_manual(
|
||||||
|
origin_img=img, hand_mask=hand_mask, hair_id=hair_id,
|
||||||
|
hairstyle_process=hairstyle_process,
|
||||||
|
landmark_processor=landmark_processor,
|
||||||
|
task_id=task_id,
|
||||||
|
strict_mask=bool(g("strict_mask", False)),
|
||||||
|
seamless_blend=bool(g("seamless_blend", True)),
|
||||||
|
is_hr=bool(g("is_hr", True)),
|
||||||
|
dilate_kernel=(int(dk[0]), int(dk[1])),
|
||||||
|
denoising_strength=float(g("denoising_strength", 0.6)),
|
||||||
|
blend_dilate=(int(bd[0]), int(bd[1])),
|
||||||
|
seamless_dilate=(int(sd[0]), int(sd[1])),
|
||||||
|
feather_px=int(g("feather_px", 0)),
|
||||||
|
enhance=bool(g("enhance", False)),
|
||||||
|
enhance_denoising=float(g("enhance_denoising", 0.35)),
|
||||||
|
)
|
||||||
|
|
||||||
|
_, buf = cv2.imencode(".jpg", result, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||||||
|
result_b64 = base64.b64encode(buf).decode("utf-8")
|
||||||
|
print(f"[swap_manual] 完成，总耗时 {time.time()-t0:.1f}s")
|
||||||
|
return jsonify({"state": 0, "result": result_b64, "steps": steps, "params": params})
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
print(f"[swap_manual] 失败: {e}")
|
||||||
|
traceback.print_exc()
|
||||||
|
return jsonify({"state": -1, "msg": f"手绘mask换发型失败: {e}"}), 500
|
||||||
|
|
||||||
|
|
||||||
|
def _b64_to_ndarray(b64_str, color=True):
|
||||||
|
if "," in b64_str and b64_str.startswith("data:"):
|
||||||
|
b64_str = b64_str.split(",", 1)[1]
|
||||||
|
data = base64.b64decode(b64_str)
|
||||||
|
flag = cv2.IMREAD_COLOR if color else cv2.IMREAD_GRAYSCALE
|
||||||
|
return cv2.imdecode(np.frombuffer(data, np.uint8), flag)
|
||||||
|
|
||||||
|
|
||||||
|
@app.route("/api/hairstyles")
|
||||||
|
def api_hairstyles():
|
||||||
|
"""返回可用发型列表（含性别，用于前端下拉框分组）"""
|
||||||
|
try:
|
||||||
|
from common.logger import config
|
||||||
|
hairstyle_dir = config.get('default', 'hairstyleDir')
|
||||||
|
train_dir = config.get('default', 'train_dir')
|
||||||
|
upload_dir = config.get('default', 'upload_train_dir')
|
||||||
|
from hair_grow_swap import list_hairstyles
|
||||||
|
styles = list_hairstyles(hairstyle_dir, train_dir, upload_dir)
|
||||||
|
boy = [s for s in styles if s["gender"] == "boy"]
|
||||||
|
girl = [s for s in styles if s["gender"] == "girl"]
|
||||||
|
print(f"[hairstyles] 共 {len(styles)} 个可用发型 (boy={len(boy)}, girl={len(girl)})")
|
||||||
|
return jsonify({"state": 0, "data": styles, "count": len(styles)})
|
||||||
|
except Exception as e:
|
||||||
|
traceback.print_exc()
|
||||||
|
return jsonify({"state": -1, "msg": str(e)}), 500
|
||||||
|
|
||||||
|
|
||||||
|
@app.route("/preview/<hair_id>")
|
||||||
|
def preview_img(hair_id):
|
||||||
|
"""直接返回预览图文件（比 base64 API 快，浏览器可缓存）"""
|
||||||
|
effect_path = os.path.join(BASE_DIR, "static", "previews", f"{hair_id}.jpg")
|
||||||
|
if os.path.exists(effect_path):
|
||||||
|
return send_from_directory(os.path.join(BASE_DIR, "static", "previews"), f"{hair_id}.jpg")
|
||||||
|
# 回退到 ref_rgb
|
||||||
|
try:
|
||||||
|
from common.logger import config
|
||||||
|
hairstyle_dir = config.get('default', 'hairstyleDir')
|
||||||
|
fallback = os.path.join(hairstyle_dir, hair_id, "ref_rgb_8uc3_768.png")
|
||||||
|
if os.path.exists(fallback):
|
||||||
|
return send_file_or_404(fallback)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
return ("", 404)
|
||||||
|
|
||||||
|
|
||||||
|
@app.route("/train_src/<hair_id>")
|
||||||
|
def train_src_img(hair_id):
|
||||||
|
"""返回发型的训练原图（hair_type_images/<hair_id>.jpg/.png）。
|
||||||
|
用于测试页展示发型真实样子，而非套在标准脸上的效果图。
|
||||||
|
"""
|
||||||
|
src_dir = "/home/xsl/change_hair/hair_type_images"
|
||||||
|
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)
|
||||||
|
return ("", 404)
|
||||||
|
|
||||||
|
|
||||||
|
@app.route("/api/hairstyle_preview/<hair_id>")
|
||||||
|
def api_hairstyle_preview(hair_id):
|
||||||
|
"""返回某发型的预览图。
|
||||||
|
优先用生成的效果图 static/previews/<hair_id>.jpg（发型套在标准脸上的样子），
|
||||||
|
没有则回退到 first## 原始上传图。
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
# 1. 优先：生成的效果图（发型套在 boy/girl 标准脸上）
|
||||||
|
preview = None
|
||||||
|
effect_path = os.path.join(BASE_DIR, "static", "previews", f"{hair_id}.jpg")
|
||||||
|
if os.path.exists(effect_path):
|
||||||
|
preview = effect_path
|
||||||
|
else:
|
||||||
|
# 2. 回退：first## 原始上传图
|
||||||
|
from common.logger import config
|
||||||
|
upload_dir = config.get('default', 'upload_train_dir')
|
||||||
|
save_dir = os.path.join(upload_dir, hair_id)
|
||||||
|
if os.path.isdir(save_dir):
|
||||||
|
for name in os.listdir(save_dir):
|
||||||
|
if "first##" in name:
|
||||||
|
preview = os.path.join(save_dir, name)
|
||||||
|
break
|
||||||
|
# 3. 再回退：ref_rgb
|
||||||
|
if not preview or not os.path.exists(preview):
|
||||||
|
hairstyle_dir = config.get('default', 'hairstyleDir')
|
||||||
|
fallback = os.path.join(hairstyle_dir, hair_id, "ref_rgb_8uc3_768.png")
|
||||||
|
if os.path.exists(fallback):
|
||||||
|
preview = fallback
|
||||||
|
if not preview:
|
||||||
|
return jsonify({"state": -1, "msg": "无预览图"}), 404
|
||||||
|
img = cv2.imread(preview)
|
||||||
|
img = cv2.resize(img, (256, 256), interpolation=cv2.INTER_AREA)
|
||||||
|
_, buf = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, 85])
|
||||||
|
b64 = base64.b64encode(buf).decode()
|
||||||
|
return jsonify({"state": 0, "preview": b64})
|
||||||
|
except Exception as e:
|
||||||
|
return jsonify({"state": -1, "msg": str(e)}), 500
|
||||||
|
|
||||||
|
|
||||||
|
@app.route("/api/grow", methods=["POST"])
|
||||||
|
def api_grow():
|
||||||
|
"""生发接口（走换发型工作流）
|
||||||
|
|
||||||
|
入参 JSON:
|
||||||
|
img: 原图 base64
|
||||||
|
mask: 手绘遮罩 base64（白=生发区，与img同分辨率）
|
||||||
|
hair_id: 选择的发型ID（必填）
|
||||||
|
is_hr: 是否高清 "true"/"false"，默认 "true"
|
||||||
|
"""
|
||||||
|
t0 = time.time()
|
||||||
|
try:
|
||||||
|
data = request.json
|
||||||
|
img_b64 = data.get("img", "")
|
||||||
|
mask_b64 = data.get("mask", "")
|
||||||
|
hair_id = data.get("hair_id", "")
|
||||||
|
is_hr = str(data.get("is_hr", "true")).lower() == "true"
|
||||||
|
|
||||||
|
if not img_b64 or not mask_b64:
|
||||||
|
return jsonify({"state": -1, "msg": "img 和 mask 不能为空"}), 400
|
||||||
|
if not hair_id:
|
||||||
|
return jsonify({"state": -1, "msg": "hair_id 不能为空（请先选择发型）"}), 400
|
||||||
|
|
||||||
|
img = _b64_to_ndarray(img_b64, color=True)
|
||||||
|
mask = _b64_to_ndarray(mask_b64, color=False)
|
||||||
|
if img is None:
|
||||||
|
return jsonify({"state": -1, "msg": "img 解析失败"}), 400
|
||||||
|
if mask is None:
|
||||||
|
return jsonify({"state": -1, "msg": "mask 解析失败"}), 400
|
||||||
|
|
||||||
|
print(f"[grow] img={img.shape}, mask={mask.shape}, hair_id={hair_id}, is_hr={is_hr}")
|
||||||
|
|
||||||
|
# 加载模型（首次慢）
|
||||||
|
hairstyle_process, landmark_processor = _get_models()
|
||||||
|
|
||||||
|
# 调用生发（走换发型工作流）
|
||||||
|
from hair_grow_swap import hair_grow_swap
|
||||||
|
task_id = f"grow_{int(time.time()*1000)}"
|
||||||
|
result = hair_grow_swap(
|
||||||
|
origin_img=img, hand_mask=mask, hair_id=hair_id,
|
||||||
|
hairstyle_process=hairstyle_process,
|
||||||
|
landmark_processor=landmark_processor,
|
||||||
|
task_id=task_id, is_hr=is_hr)
|
||||||
|
|
||||||
|
_, buf = cv2.imencode(".jpg", result, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||||||
|
result_b64 = base64.b64encode(buf).decode("utf-8")
|
||||||
|
print(f"[grow] 完成，总耗时 {time.time()-t0:.1f}s")
|
||||||
|
return jsonify({"state": 0, "result": result_b64})
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
print(f"[grow] 失败: {e}")
|
||||||
|
traceback.print_exc()
|
||||||
|
return jsonify({"state": -1, "msg": f"生发失败: {e}"}), 500
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
print(f"生发服务启动，端口 {PORT}")
|
||||||
|
print(f"测试页面: http://0.0.0.0:{PORT}")
|
||||||
|
print(f"注意：首次请求会加载换发型模型（约30-60秒）")
|
||||||
|
server = pywsgi.WSGIServer(("0.0.0.0", PORT), app)
|
||||||
|
server.serve_forever()
|
||||||
@@ -0,0 +1,379 @@
|
|||||||
|
<!DOCTYPE html>
|
||||||
|
<html lang="zh-CN">
|
||||||
|
<head>
|
||||||
|
<meta charset="UTF-8">
|
||||||
|
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||||
|
<title>生发调试台</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: 340px; flex-shrink: 0; background: #16213e; border-radius: 10px; padding: 16px; height: fit-content; max-height: 92vh; overflow-y: auto; }
|
||||||
|
.panel::-webkit-scrollbar { width: 6px; }
|
||||||
|
.panel::-webkit-scrollbar-thumb { background: #0f3460; border-radius: 3px; }
|
||||||
|
.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; }
|
||||||
|
.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-label .tip { color: #666; font-size: 10px; margin-left: 4px; cursor: help; }
|
||||||
|
.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="number"] { width: 100%; padding: 5px 8px; background: #0d1b3e; border: 1px solid #0f3460; border-radius: 4px; color: #eee; font-size: 12px; }
|
||||||
|
.field select { width: 100%; padding: 5px 8px; background: #0d1b3e; border: 1px solid #0f3460; border-radius: 4px; color: #eee; font-size: 12px; }
|
||||||
|
.switch { display: flex; align-items: center; justify-content: space-between; padding: 7px 0; font-size: 13px; color: #ccc; cursor: pointer; }
|
||||||
|
.switch:hover { color: #eee; }
|
||||||
|
.toggle { position: relative; width: 36px; height: 20px; background: #0f3460; border-radius: 10px; transition: .2s; flex-shrink: 0; }
|
||||||
|
.toggle::after { content: ''; position: absolute; top: 2px; left: 2px; width: 16px; height: 16px; background: #888; border-radius: 50%; transition: .2s; }
|
||||||
|
.switch.on .toggle { background: #4ecca3; }
|
||||||
|
.switch.on .toggle::after { left: 18px; background: #16213e; }
|
||||||
|
.pair { display: flex; gap: 8px; }
|
||||||
|
.pair .field { flex: 1; }
|
||||||
|
.hint { font-size: 11px; color: #666; line-height: 1.5; margin-top: 8px; padding: 8px 10px; background: #0d1b3e; border-radius: 5px; }
|
||||||
|
.hint b { color: #aaa; }
|
||||||
|
.readonly-info { font-size: 11px; color: #777; padding: 8px 10px; background: #0d1b3e; border-radius: 5px; line-height: 1.6; }
|
||||||
|
.readonly-info code { color: #aaa; background: #16213e; padding: 1px 5px; border-radius: 3px; }
|
||||||
|
.btn { display: block; width: 100%; padding: 10px; border: none; border-radius: 6px; font-size: 14px; cursor: pointer; transition: .15s; }
|
||||||
|
.btn-primary { background: #4ecca3; color: #16213e; font-weight: 600; font-size: 15px; padding: 12px; margin-top: 8px; }
|
||||||
|
.btn-primary:hover { background: #6ee0bd; }
|
||||||
|
.btn-primary:disabled { background: #555; color: #999; cursor: not-allowed; }
|
||||||
|
.btn-ghost { background: #0f3460; color: #eee; margin-top: 6px; }
|
||||||
|
.btn-ghost:hover { background: #1a4a80; }
|
||||||
|
|
||||||
|
/* 中间发型选择 */
|
||||||
|
.mid-panel { width: 200px; flex-shrink: 0; background: #16213e; border-radius: 10px; padding: 14px; height: fit-content; max-height: 92vh; overflow-y: auto; }
|
||||||
|
.mid-panel h3 { font-size: 12px; color: #4ecca3; margin-bottom: 10px; }
|
||||||
|
.hairstyle-grid { display: grid; grid-template-columns: 1fr 1fr; gap: 6px; }
|
||||||
|
.hairstyle-item { position: relative; cursor: pointer; border-radius: 5px; overflow: hidden; border: 2px solid transparent; aspect-ratio: 1; background: #0d1b3e; }
|
||||||
|
.hairstyle-item img { width: 100%; height: 100%; object-fit: cover; display: block; }
|
||||||
|
.hairstyle-item.ready { border-color: #2a5a4a; }
|
||||||
|
.hairstyle-item.selected { border-color: #4ecca3; }
|
||||||
|
.hairstyle-item.selected::after { content: '✓'; position: absolute; top: 1px; right: 3px; color: #4ecca3; font-weight: bold; text-shadow: 0 0 3px #000; font-size: 11px; }
|
||||||
|
.hairstyle-item .name { position: absolute; bottom: 0; left: 0; right: 0; background: linear-gradient(transparent, rgba(0,0,0,.85)); color: #fff; font-size: 10px; padding: 8px 2px 2px; text-align: center; }
|
||||||
|
.hairstyle-item.pending { opacity: 0.4; 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: 460px; 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; }
|
||||||
|
#preview-img { display: none; max-width: 100%; max-height: 460px; border-radius: 6px; }
|
||||||
|
#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-top { margin-top: 16px; background: #16213e; border-radius: 10px; padding: 16px; display: none; }
|
||||||
|
.result-top.show { display: block; }
|
||||||
|
.result-top h3 { font-size: 14px; color: #4ecca3; margin-bottom: 12px; }
|
||||||
|
.result-pair { 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; }
|
||||||
|
|
||||||
|
.params-used { margin-top: 12px; padding: 10px 12px; background: #0d1b3e; border-radius: 6px; font-size: 11px; color: #888; line-height: 1.7; font-family: monospace; }
|
||||||
|
|
||||||
|
.steps-section { margin-top: 16px; }
|
||||||
|
.step-item { background: #16213e; border-radius: 10px; padding: 16px; margin-bottom: 14px; }
|
||||||
|
.step-head { display: flex; justify-content: space-between; align-items: baseline; margin-bottom: 4px; }
|
||||||
|
.step-title { font-size: 15px; font-weight: 600; color: #4ecca3; }
|
||||||
|
.step-time { font-size: 11px; color: #666; font-family: monospace; }
|
||||||
|
.step-desc { font-size: 12px; color: #999; margin-bottom: 12px; line-height: 1.6; }
|
||||||
|
.step-images { display: flex; gap: 12px; flex-wrap: wrap; }
|
||||||
|
.step-img-card { flex: 1; min-width: 180px; max-width: 280px; }
|
||||||
|
.step-img-card .lbl { font-size: 11px; color: #888; margin-bottom: 4px; text-align: center; }
|
||||||
|
.step-img-card img { width: 100%; border-radius: 6px; display: block; border: 1px solid #0f3460; cursor: zoom-in; transition: .15s; }
|
||||||
|
.step-img-card img:hover { border-color: #4ecca3; }
|
||||||
|
#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; }
|
||||||
|
#lightbox img { max-width: 92%; max-height: 92%; border-radius: 8px; }
|
||||||
|
#lightbox.show { display: flex; }
|
||||||
|
</style>
|
||||||
|
</head>
|
||||||
|
<body>
|
||||||
|
<div class="header">
|
||||||
|
<div>
|
||||||
|
<h1>🔬 生发调试台（逐步可视化 + 全参数）</h1>
|
||||||
|
<p>调参数 → 看每一步中间产物如何变化。绿色框=可编辑参数，灰色=webui端固定</p>
|
||||||
|
</div>
|
||||||
|
<div class="links">
|
||||||
|
<a href="/test_new">→ 简易测试页</a>
|
||||||
|
<a href="/swap">→ 全部发型页</a>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div class="container">
|
||||||
|
<!-- 左侧：参数面板 -->
|
||||||
|
<div class="panel">
|
||||||
|
<h3>① 流程开关</h3>
|
||||||
|
<div class="switch on" id="sw-cut_bang" onclick="toggleSwitch('cut_bang', this)">
|
||||||
|
<span>减刘海圆 <span class="tip" title="从原头发mask中扣除眼睛周围的圆形区域，避免重绘眼睛">ⓘ</span></span>
|
||||||
|
<div class="toggle"></div>
|
||||||
|
</div>
|
||||||
|
<div class="switch" id="sw-strict_mask" onclick="toggleSwitch('strict_mask', this)">
|
||||||
|
<span>严格mask贴回 <span class="tip" title="只在mask区域覆盖，mask外保留原图（否则整个裁剪框覆盖）">ⓘ</span></span>
|
||||||
|
<div class="toggle"></div>
|
||||||
|
</div>
|
||||||
|
<div class="switch on" id="sw-seamless_blend" onclick="toggleSwitch('seamless_blend', this)">
|
||||||
|
<span>泊松融合消接缝 <span class="tip" title="严格mask模式下用seamlessClone消除边缘色差（仅strict_mask开启时生效）">ⓘ</span></span>
|
||||||
|
<div class="toggle"></div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<h3>② 尺寸 / 对齐</h3>
|
||||||
|
<div class="switch on" id="sw-is_hr" onclick="toggleSwitch('is_hr', this)">
|
||||||
|
<span>高清模式 <span class="tip" title="1152x1536 vs 576x768，影响显存和速度">ⓘ</span></span>
|
||||||
|
<div class="toggle"></div>
|
||||||
|
</div>
|
||||||
|
<div class="pair">
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">dilate_kernel <span class="val" id="v-dk0">6</span></div>
|
||||||
|
<input type="number" id="dk0" value="6" min="1" max="40" oninput="syncNum('dk0')">
|
||||||
|
</div>
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">mask膨胀核(x,y) <span class="val" id="v-dk1">18</span></div>
|
||||||
|
<input type="number" id="dk1" value="18" min="1" max="40" oninput="syncNum('dk1')">
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<h3>③ SD 推理参数</h3>
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">denoising_strength <span class="val" id="v-den">0.60</span><span class="tip" title="重绘强度。越高越偏离原图，发丝细节越多但可能失真">ⓘ</span></div>
|
||||||
|
<input type="range" id="den" min="0.1" max="1.0" step="0.05" value="0.6" oninput="syncRange('den',2)">
|
||||||
|
</div>
|
||||||
|
<div class="readonly-info">
|
||||||
|
<b style="color:#888">webui端固定（不可调）：</b><br>
|
||||||
|
<code>cfg_scale=7</code> <code>steps=20</code><br>
|
||||||
|
<code>sampler=DPM++ 2M Karras</code><br>
|
||||||
|
<code>mask_blur=11</code> <code>seed=123456789</code>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<h3>④ 贴回 / 融合</h3>
|
||||||
|
<div class="pair">
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">blend_dilate <span class="val" id="v-bd0">5</span></div>
|
||||||
|
<input type="number" id="bd0" value="5" min="1" max="40" oninput="syncNum('bd0')">
|
||||||
|
</div>
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">(strict贴回) <span class="val" id="v-bd1">5</span></div>
|
||||||
|
<input type="number" id="bd1" value="5" min="1" max="40" oninput="syncNum('bd1')">
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
<div class="pair">
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">seamless_dilate <span class="val" id="v-sd0">9</span></div>
|
||||||
|
<input type="number" id="sd0" value="9" min="1" max="40" oninput="syncNum('sd0')">
|
||||||
|
</div>
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">(泊松融合) <span class="val" id="v-sd1">9</span></div>
|
||||||
|
<input type="number" id="sd1" value="9" min="1" max="40" oninput="syncNum('sd1')">
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<button class="btn btn-primary" id="btn-run" disabled>▶ 执行生发</button>
|
||||||
|
<button class="btn btn-ghost" id="btn-reset" onclick="resetParams()">↺ 恢复默认参数</button>
|
||||||
|
<div class="hint">
|
||||||
|
<b>使用说明</b><br>
|
||||||
|
1. 选发型（中间）→ 上传人像（右侧）<br>
|
||||||
|
2. 调整左侧参数<br>
|
||||||
|
3. 点"执行"，看7步中间产物<br>
|
||||||
|
4. 改参数重跑，对比效果差异
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<!-- 中间：发型选择 -->
|
||||||
|
<div class="mid-panel">
|
||||||
|
<h3>选择发型</h3>
|
||||||
|
<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>
|
||||||
|
<button class="btn btn-ghost" style="margin-top:8px;font-size:12px" onclick="loadStatus()">↻ 刷新</button>
|
||||||
|
</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>正面清晰照效果最佳</p>
|
||||||
|
</div>
|
||||||
|
<input type="file" id="file-input" accept="image/*" style="display:none">
|
||||||
|
<img id="preview-img">
|
||||||
|
</div>
|
||||||
|
<div id="status"><span class="spinner"></span><span id="status-text">执行中，约40-90秒...</span></div>
|
||||||
|
<div id="error-msg"></div>
|
||||||
|
|
||||||
|
<div class="result-top" id="result-top">
|
||||||
|
<h3>最终对比</h3>
|
||||||
|
<div class="result-pair">
|
||||||
|
<div class="result-card"><h4>原图</h4><img id="r-orig"></div>
|
||||||
|
<div class="result-card"><h4>生发结果</h4><img id="r-new"></div>
|
||||||
|
</div>
|
||||||
|
<div class="params-used" id="params-used"></div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div class="steps-section" id="steps-section"></div>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div id="lightbox" onclick="this.classList.remove('show')"><img id="lightbox-img"></div>
|
||||||
|
|
||||||
|
<script>
|
||||||
|
(function(){
|
||||||
|
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:"长 -花瓣"}
|
||||||
|
];
|
||||||
|
const defaults = {cut_bang:true,strict_mask:false,seamless_blend:true,is_hr:true,
|
||||||
|
dk0:6,dk1:18,den:0.6,bd0:5,bd1:5,sd0:9,sd1:9};
|
||||||
|
|
||||||
|
let readySet=new Set(), selectedHair=null, imgB64=null;
|
||||||
|
|
||||||
|
async function loadStatus(){
|
||||||
|
try{
|
||||||
|
const r=await fetch("/api/hairstyles");const d=await r.json();
|
||||||
|
if(d.state!==0)throw new Error(d.msg);
|
||||||
|
readySet=new Set(d.data.map(x=>x.hair_id));
|
||||||
|
}catch(e){readySet=new Set();}
|
||||||
|
renderGrid();
|
||||||
|
}
|
||||||
|
function renderGrid(){
|
||||||
|
const g=$("hairstyle-grid");g.innerHTML="";
|
||||||
|
HAIRSTYLES.forEach(h=>{
|
||||||
|
const ready=readySet.has(h.name);
|
||||||
|
const it=document.createElement("div");
|
||||||
|
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>`;
|
||||||
|
if(ready) it.onclick=()=>{
|
||||||
|
selectedHair=h.name;
|
||||||
|
$("sel-info").textContent="已选: "+h.name;
|
||||||
|
renderGrid(); updateBtn();
|
||||||
|
};
|
||||||
|
g.appendChild(it);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// 参数控件
|
||||||
|
window.toggleSwitch=(key,el)=>{
|
||||||
|
el.classList.toggle("on");
|
||||||
|
updateBtn();
|
||||||
|
};
|
||||||
|
// range 滑块绑定显示
|
||||||
|
["den"].forEach(id=>{
|
||||||
|
$(id).addEventListener("input",e=>$("v-"+id).textContent=parseFloat(e.target.value).toFixed(2));
|
||||||
|
});
|
||||||
|
window.syncNum=(id)=>{ $("v-"+id).textContent=$(id).value; };
|
||||||
|
window.resetParams=()=>{
|
||||||
|
Object.keys(defaults).forEach(k=>{
|
||||||
|
if(typeof defaults[k]==="boolean"){
|
||||||
|
const el=$("sw-"+k); el.classList.toggle("on",defaults[k]);
|
||||||
|
}else if(k==="den"){ $("den").value=defaults[k]; $("v-den").textContent=defaults[k].toFixed(2);
|
||||||
|
}else{ $(k).value=defaults[k]; $("v-"+k).textContent=defaults[k]; }
|
||||||
|
});
|
||||||
|
};
|
||||||
|
|
||||||
|
function collectParams(){
|
||||||
|
return {
|
||||||
|
cut_bang:$("sw-cut_bang").classList.contains("on"),
|
||||||
|
strict_mask:$("sw-strict_mask").classList.contains("on"),
|
||||||
|
seamless_blend:$("sw-seamless_blend").classList.contains("on"),
|
||||||
|
is_hr:$("sw-is_hr").classList.contains("on"),
|
||||||
|
dilate_kernel:[parseInt($("dk0").value),parseInt($("dk1").value)],
|
||||||
|
denoising_strength:parseFloat($("den").value),
|
||||||
|
blend_dilate:[parseInt($("bd0").value),parseInt($("bd1").value)],
|
||||||
|
seamless_dilate:[parseInt($("sd0").value),parseInt($("sd1").value)],
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
// 上传
|
||||||
|
$("upload-zone").onclick=()=>$("file-input").click();
|
||||||
|
$("file-input").onchange=e=>{ if(e.target.files[0]) loadImg(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])loadImg(e.dataTransfer.files[0]);};
|
||||||
|
function loadImg(file){
|
||||||
|
if(!file.type.startsWith("image/")){alert("请上传图片");return;}
|
||||||
|
const rd=new FileReader();
|
||||||
|
rd.onload=e=>{
|
||||||
|
imgB64=e.target.result;
|
||||||
|
$("preview-img").src=imgB64;
|
||||||
|
$("preview-img").style.display="block";
|
||||||
|
$("upload-zone").style.display="none";
|
||||||
|
updateBtn();
|
||||||
|
};
|
||||||
|
rd.readAsDataURL(file);
|
||||||
|
}
|
||||||
|
function updateBtn(){ $("btn-run").disabled=!(selectedHair&&imgB64); }
|
||||||
|
|
||||||
|
// 执行
|
||||||
|
$("btn-run").onclick=async()=>{
|
||||||
|
if(!selectedHair||!imgB64) return;
|
||||||
|
const p=collectParams();
|
||||||
|
$("btn-run").disabled=true;
|
||||||
|
$("status").style.display="block";
|
||||||
|
$("status-text").textContent="执行中（粗推理→mask→warpAffine→SD→贴回），约40-90秒...";
|
||||||
|
$("error-msg").style.display="none";
|
||||||
|
$("result-top").classList.remove("show");
|
||||||
|
$("steps-section").innerHTML="";
|
||||||
|
try{
|
||||||
|
const resp=await fetch("/api/swap_debug",{
|
||||||
|
method:"POST",headers:{"Content-Type":"application/json"},
|
||||||
|
body:JSON.stringify(Object.assign({img:imgB64,hair_id:selectedHair},p))
|
||||||
|
});
|
||||||
|
const data=await resp.json();
|
||||||
|
$("status").style.display="none";
|
||||||
|
if(data.state===0){
|
||||||
|
const url="data:image/jpeg;base64,"+data.result;
|
||||||
|
$("r-orig").src=imgB64;$("r-new").src=url;
|
||||||
|
// 参数回显
|
||||||
|
let ph=""; Object.keys(data.params).forEach(k=>{
|
||||||
|
let v=data.params[k]; if(Array.isArray(v))v="["+v.join(",")+"]";
|
||||||
|
if(typeof v==="boolean")v=v?"开":"关";
|
||||||
|
ph+=k+" = "+v+" ";
|
||||||
|
});
|
||||||
|
$("params-used").textContent="本次参数: "+ph;
|
||||||
|
$("result-top").classList.add("show");
|
||||||
|
// 步骤画廊
|
||||||
|
const sc=$("steps-section");
|
||||||
|
data.steps.forEach(step=>{
|
||||||
|
const div=document.createElement("div");div.className="step-item";
|
||||||
|
let imgs='<div class="step-images">';
|
||||||
|
step.images.forEach(im=>{imgs+='<div class="step-img-card"><div class="lbl">'+im.label+'</div><img src="data:image/jpeg;base64,'+im.b64+'" data-full="data:image/jpeg;base64,'+im.b64+'"></div>';});
|
||||||
|
imgs+='</div>';
|
||||||
|
div.innerHTML='<div class="step-head"><div class="step-title">'+step.title+'</div></div><div class="step-desc">'+step.desc+'</div>'+imgs;
|
||||||
|
sc.appendChild(div);
|
||||||
|
});
|
||||||
|
sc.querySelectorAll("img").forEach(img=>img.onclick=()=>showLightbox(img.dataset.full));
|
||||||
|
$("result-top").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";
|
||||||
|
}
|
||||||
|
updateBtn();
|
||||||
|
};
|
||||||
|
|
||||||
|
function showLightbox(src){$("lightbox-img").src=src;$("lightbox").classList.add("show");}
|
||||||
|
|
||||||
|
loadStatus();
|
||||||
|
})();
|
||||||
|
</script>
|
||||||
|
</body>
|
||||||
|
</html>
|
||||||
@@ -0,0 +1,434 @@
|
|||||||
|
<!DOCTYPE html>
|
||||||
|
<html lang="zh-CN">
|
||||||
|
<head>
|
||||||
|
<meta charset="UTF-8">
|
||||||
|
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||||
|
<title>发际线带重绘实验</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: 340px; flex-shrink: 0; background: #16213e; border-radius: 10px; padding: 16px; height: fit-content; max-height: 92vh; overflow-y: auto; }
|
||||||
|
.panel::-webkit-scrollbar { width: 6px; }
|
||||||
|
.panel::-webkit-scrollbar-thumb { background: #0f3460; border-radius: 3px; }
|
||||||
|
.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; }
|
||||||
|
.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-label .tip { color: #666; font-size: 10px; margin-left: 4px; cursor: help; }
|
||||||
|
.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="number"] { width: 100%; padding: 5px 8px; background: #0d1b3e; border: 1px solid #0f3460; border-radius: 4px; color: #eee; font-size: 12px; }
|
||||||
|
.field select { width: 100%; padding: 5px 8px; background: #0d1b3e; border: 1px solid #0f3460; border-radius: 4px; color: #eee; font-size: 12px; }
|
||||||
|
.switch { display: flex; align-items: center; justify-content: space-between; padding: 7px 0; font-size: 13px; color: #ccc; cursor: pointer; }
|
||||||
|
.switch:hover { color: #eee; }
|
||||||
|
.toggle { position: relative; width: 36px; height: 20px; background: #0f3460; border-radius: 10px; transition: .2s; flex-shrink: 0; }
|
||||||
|
.toggle::after { content: ''; position: absolute; top: 2px; left: 2px; width: 16px; height: 16px; background: #888; border-radius: 50%; transition: .2s; }
|
||||||
|
.switch.on .toggle { background: #4ecca3; }
|
||||||
|
.switch.on .toggle::after { left: 18px; background: #16213e; }
|
||||||
|
.pair { display: flex; gap: 8px; }
|
||||||
|
.pair .field { flex: 1; }
|
||||||
|
.hint { font-size: 11px; color: #666; line-height: 1.5; margin-top: 8px; padding: 8px 10px; background: #0d1b3e; border-radius: 5px; }
|
||||||
|
.hint b { color: #aaa; }
|
||||||
|
.readonly-info { font-size: 11px; color: #777; padding: 8px 10px; background: #0d1b3e; border-radius: 5px; line-height: 1.6; }
|
||||||
|
.readonly-info code { color: #aaa; background: #16213e; padding: 1px 5px; border-radius: 3px; }
|
||||||
|
.btn { display: block; width: 100%; padding: 10px; border: none; border-radius: 6px; font-size: 14px; cursor: pointer; transition: .15s; }
|
||||||
|
.btn-primary { background: #4ecca3; color: #16213e; font-weight: 600; font-size: 15px; padding: 12px; margin-top: 8px; }
|
||||||
|
.btn-primary:hover { background: #6ee0bd; }
|
||||||
|
.btn-primary:disabled { background: #555; color: #999; cursor: not-allowed; }
|
||||||
|
.btn-ghost { background: #0f3460; color: #eee; margin-top: 6px; }
|
||||||
|
.btn-ghost:hover { background: #1a4a80; }
|
||||||
|
|
||||||
|
/* 中间发型选择 */
|
||||||
|
.mid-panel { width: 200px; flex-shrink: 0; background: #16213e; border-radius: 10px; padding: 14px; height: fit-content; max-height: 92vh; overflow-y: auto; }
|
||||||
|
.mid-panel h3 { font-size: 12px; color: #4ecca3; margin-bottom: 10px; }
|
||||||
|
.hairstyle-grid { display: grid; grid-template-columns: 1fr 1fr; gap: 6px; }
|
||||||
|
.hairstyle-item { position: relative; cursor: pointer; border-radius: 5px; overflow: hidden; border: 2px solid transparent; aspect-ratio: 1; background: #0d1b3e; }
|
||||||
|
.hairstyle-item img { width: 100%; height: 100%; object-fit: cover; display: block; }
|
||||||
|
.hairstyle-item.ready { border-color: #2a5a4a; }
|
||||||
|
.hairstyle-item.selected { border-color: #4ecca3; }
|
||||||
|
.hairstyle-item.selected::after { content: '✓'; position: absolute; top: 1px; right: 3px; color: #4ecca3; font-weight: bold; text-shadow: 0 0 3px #000; font-size: 11px; }
|
||||||
|
.hairstyle-item .name { position: absolute; bottom: 0; left: 0; right: 0; background: linear-gradient(transparent, rgba(0,0,0,.85)); color: #fff; font-size: 10px; padding: 8px 2px 2px; text-align: center; }
|
||||||
|
.hairstyle-item.pending { opacity: 0.4; 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: 460px; 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; }
|
||||||
|
#preview-img { display: none; max-width: 100%; max-height: 460px; border-radius: 6px; }
|
||||||
|
#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-top { margin-top: 16px; background: #16213e; border-radius: 10px; padding: 16px; display: none; }
|
||||||
|
.result-top.show { display: block; }
|
||||||
|
.result-top h3 { font-size: 14px; color: #4ecca3; margin-bottom: 12px; }
|
||||||
|
.result-pair { 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; }
|
||||||
|
|
||||||
|
.params-used { margin-top: 12px; padding: 10px 12px; background: #0d1b3e; border-radius: 6px; font-size: 11px; color: #888; line-height: 1.7; font-family: monospace; }
|
||||||
|
|
||||||
|
.steps-section { margin-top: 16px; }
|
||||||
|
.step-item { background: #16213e; border-radius: 10px; padding: 16px; margin-bottom: 14px; }
|
||||||
|
.step-head { display: flex; justify-content: space-between; align-items: baseline; margin-bottom: 4px; }
|
||||||
|
.step-title { font-size: 15px; font-weight: 600; color: #4ecca3; }
|
||||||
|
.step-time { font-size: 11px; color: #666; font-family: monospace; }
|
||||||
|
.step-desc { font-size: 12px; color: #999; margin-bottom: 12px; line-height: 1.6; }
|
||||||
|
.step-images { display: flex; gap: 12px; flex-wrap: wrap; }
|
||||||
|
.step-img-card { flex: 1; min-width: 180px; max-width: 280px; }
|
||||||
|
.step-img-card .lbl { font-size: 11px; color: #888; margin-bottom: 4px; text-align: center; }
|
||||||
|
.step-img-card img { width: 100%; border-radius: 6px; display: block; border: 1px solid #0f3460; cursor: zoom-in; transition: .15s; }
|
||||||
|
.step-img-card img:hover { border-color: #4ecca3; }
|
||||||
|
#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; }
|
||||||
|
#lightbox img { max-width: 92%; max-height: 92%; border-radius: 8px; }
|
||||||
|
#lightbox.show { display: flex; }
|
||||||
|
</style>
|
||||||
|
</head>
|
||||||
|
<body>
|
||||||
|
<div class="header">
|
||||||
|
<div>
|
||||||
|
<h1>🔬 发际线带重绘实验（只重绘发际线，主体保留原图）</h1>
|
||||||
|
<p>实验性：重绘区域只取发际线边界带，观察"只改发际线、保留原发型主体"的效果</p>
|
||||||
|
</div>
|
||||||
|
<div class="links">
|
||||||
|
<a href="/test_new">→ 简易测试页</a>
|
||||||
|
<a href="/swap">→ 全部发型页</a>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div class="container">
|
||||||
|
<!-- 左侧：参数面板 -->
|
||||||
|
<div class="panel">
|
||||||
|
<h3>① 流程开关</h3>
|
||||||
|
<div class="switch" id="sw-strict_mask" onclick="toggleSwitch('strict_mask', this)">
|
||||||
|
<span>严格mask贴回 <span class="tip" title="只在mask区域覆盖，mask外保留原图（否则整个裁剪框覆盖）">ⓘ</span></span>
|
||||||
|
<div class="toggle"></div>
|
||||||
|
</div>
|
||||||
|
<div class="switch on" id="sw-seamless_blend" onclick="toggleSwitch('seamless_blend', this)">
|
||||||
|
<span>泊松融合消接缝 <span class="tip" title="严格mask模式下用seamlessClone消除边缘色差（仅strict_mask开启时生效）">ⓘ</span></span>
|
||||||
|
<div class="toggle"></div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<h3>② 尺寸 / 对齐</h3>
|
||||||
|
<div class="switch on" id="sw-is_hr" onclick="toggleSwitch('is_hr', this)">
|
||||||
|
<span>高清模式 <span class="tip" title="1152x1536 vs 576x768，影响显存和速度">ⓘ</span></span>
|
||||||
|
<div class="toggle"></div>
|
||||||
|
</div>
|
||||||
|
<div class="pair">
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">dilate_kernel <span class="val" id="v-dk0">6</span></div>
|
||||||
|
<input type="number" id="dk0" value="6" min="1" max="40" oninput="syncNum('dk0')">
|
||||||
|
</div>
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">mask膨胀核(x,y) <span class="val" id="v-dk1">18</span></div>
|
||||||
|
<input type="number" id="dk1" value="18" min="1" max="40" oninput="syncNum('dk1')">
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<h3>③ 贴回 / 融合</h3>
|
||||||
|
<div class="pair">
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">blend_dilate <span class="val" id="v-bd0">5</span></div>
|
||||||
|
<input type="number" id="bd0" value="5" min="1" max="40" oninput="syncNum('bd0')">
|
||||||
|
</div>
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">(strict贴回) <span class="val" id="v-bd1">5</span></div>
|
||||||
|
<input type="number" id="bd1" value="5" min="1" max="40" oninput="syncNum('bd1')">
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
<div class="pair">
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">seamless_dilate <span class="val" id="v-sd0">9</span></div>
|
||||||
|
<input type="number" id="sd0" value="9" min="1" max="40" oninput="syncNum('sd0')">
|
||||||
|
</div>
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">(泊松融合) <span class="val" id="v-sd1">9</span></div>
|
||||||
|
<input type="number" id="sd1" value="9" min="1" max="40" oninput="syncNum('sd1')">
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<h3 style="color:#e9b949">★ 发际线mask参数（landmark_1k 圆角矩形）</h3>
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">band_width <span class="val" id="v-bw">15</span><span class="tip" title="膨胀核宽度(像素)。越大重绘区域越宽">ⓘ</span></div>
|
||||||
|
<input type="range" id="bw" min="3" max="50" step="1" value="15" oninput="syncRangeInt('bw')">
|
||||||
|
</div>
|
||||||
|
<div class="hint" style="font-size:10px;color:#e9b949;margin:6px 0 4px">▼ 圆角矩形形状</div>
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">height_ratio <span class="val" id="v-hr">0.43</span><span class="tip" title="高度缩放比(相对原额头高度)。越小上下越窄">ⓘ</span></div>
|
||||||
|
<input type="range" id="hr" min="0.1" max="1.0" step="0.02" value="0.43" oninput="syncRangeFloat('hr',2)">
|
||||||
|
</div>
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">width_ratio <span class="val" id="v-wr">0.14</span><span class="tip" title="左右各扩展比(相对眉宽)。越大左右越宽">ⓘ</span></div>
|
||||||
|
<input type="range" id="wr" min="0" max="0.6" step="0.02" value="0.14" oninput="syncRangeFloat('wr',2)">
|
||||||
|
</div>
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">corner_ratio <span class="val" id="v-cr">0.25</span><span class="tip" title="圆角半径比(相对短边)。越大圆角越大">ⓘ</span></div>
|
||||||
|
<input type="range" id="cr" min="0" max="0.5" step="0.05" value="0.25" oninput="syncRangeFloat('cr',2)">
|
||||||
|
</div>
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">vertical_offset <span class="val" id="v-vo">0.00</span><span class="tip" title="上下平移比(相对额头高度)。正值下移，负值上移">ⓘ</span></div>
|
||||||
|
<input type="range" id="vo" min="-0.5" max="0.5" step="0.02" value="0" oninput="syncRangeFloat('vo',2)">
|
||||||
|
</div>
|
||||||
|
<div class="hint" style="font-size:10px;color:#4ecca3;margin:6px 0 4px">▼ SD 推理（仅执行换发型时生效）</div>
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">denoising <span class="val" id="v-den">0.60</span><span class="tip" title="重绘强度。越高越偏离原图">ⓘ</span></div>
|
||||||
|
<input type="range" id="den" min="0.1" max="1.0" step="0.05" value="0.6" oninput="syncRangeFloat('den',2)">
|
||||||
|
</div>
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">refiner_switch_at <span class="val" id="v-rsa">0.50</span><span class="tip" title="refiner切换点(仅高清模式)。0=不用refiner，0.5=中途切换，1=全程refiner">ⓘ</span></div>
|
||||||
|
<input type="range" id="rsa" min="0" max="1.0" step="0.05" value="0.5" oninput="syncRangeFloat('rsa',2)">
|
||||||
|
</div>
|
||||||
|
<div class="switch on" id="sw-preview_only" onclick="toggleSwitch('preview_only', this)">
|
||||||
|
<span>仅验证mask <span class="tip" title="开=只生成发际线mask(几秒)，不做换发型重绘。关=执行完整换发型(40-90秒)">ⓘ</span></span>
|
||||||
|
<div class="toggle"></div>
|
||||||
|
</div>
|
||||||
|
<div class="hint" style="border-left:3px solid #e9b949">
|
||||||
|
<b style="color:#e9b949">说明：</b><br>
|
||||||
|
步骤③用 <b>landmark_1k 关键点</b>定位额头，生成<b>圆角矩形</b>重绘区。头发主体保留原图，LoRA 仅在发际线区生效。<br>
|
||||||
|
看「★ 叠加用户原图」直观判断重绘区位置。
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<button class="btn btn-primary" id="btn-run" disabled>▶ 执行换发型</button>
|
||||||
|
<button class="btn btn-ghost" id="btn-reset" onclick="resetParams()">↺ 恢复默认参数</button>
|
||||||
|
<div class="hint">
|
||||||
|
<b>使用说明</b><br>
|
||||||
|
1. 选发型（中间）→ 上传人像（右侧）<br>
|
||||||
|
2. 调整左侧参数<br>
|
||||||
|
3. 点"执行"，看7步中间产物<br>
|
||||||
|
4. 改参数重跑，对比效果差异
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<!-- 中间：发型选择 -->
|
||||||
|
<div class="mid-panel">
|
||||||
|
<h3>选择发型</h3>
|
||||||
|
<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>
|
||||||
|
<button class="btn btn-ghost" style="margin-top:8px;font-size:12px" onclick="loadStatus()">↻ 刷新</button>
|
||||||
|
</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>正面清晰照效果最佳</p>
|
||||||
|
</div>
|
||||||
|
<input type="file" id="file-input" accept="image/*" style="display:none">
|
||||||
|
<img id="preview-img">
|
||||||
|
</div>
|
||||||
|
<div id="status"><span class="spinner"></span><span id="status-text">执行中，约40-90秒...</span></div>
|
||||||
|
<div id="error-msg"></div>
|
||||||
|
|
||||||
|
<div class="result-top" id="result-top">
|
||||||
|
<h3>最终对比</h3>
|
||||||
|
<div class="result-pair">
|
||||||
|
<div class="result-card"><h4>原图</h4><img id="r-orig"></div>
|
||||||
|
<div class="result-card"><h4>换发型结果</h4><img id="r-new"></div>
|
||||||
|
</div>
|
||||||
|
<div class="params-used" id="params-used"></div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div class="steps-section" id="steps-section"></div>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div id="lightbox" onclick="this.classList.remove('show')"><img id="lightbox-img"></div>
|
||||||
|
|
||||||
|
<script>
|
||||||
|
(function(){
|
||||||
|
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:"长 -花瓣"}
|
||||||
|
];
|
||||||
|
const defaults = {method:"landmark_1k",strict_mask:false,seamless_blend:true,is_hr:true,
|
||||||
|
dk0:6,dk1:18,bd0:5,bd1:5,sd0:9,sd1:9,bw:15,preview_only:true,
|
||||||
|
hr:0.43,wr:0.14,cr:0.25,vo:0,den:0.6,rsa:0.5};
|
||||||
|
|
||||||
|
let readySet=new Set(), selectedHair=null, imgB64=null;
|
||||||
|
|
||||||
|
async function loadStatus(){
|
||||||
|
try{
|
||||||
|
const r=await fetch("/api/hairstyles");const d=await r.json();
|
||||||
|
if(d.state!==0)throw new Error(d.msg);
|
||||||
|
readySet=new Set(d.data.map(x=>x.hair_id));
|
||||||
|
}catch(e){readySet=new Set();}
|
||||||
|
renderGrid();
|
||||||
|
}
|
||||||
|
function renderGrid(){
|
||||||
|
const g=$("hairstyle-grid");g.innerHTML="";
|
||||||
|
HAIRSTYLES.forEach(h=>{
|
||||||
|
const ready=readySet.has(h.name);
|
||||||
|
const it=document.createElement("div");
|
||||||
|
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>`;
|
||||||
|
if(ready) it.onclick=()=>{
|
||||||
|
selectedHair=h.name;
|
||||||
|
$("sel-info").textContent="已选: "+h.name;
|
||||||
|
renderGrid(); updateBtn();
|
||||||
|
};
|
||||||
|
g.appendChild(it);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// 参数控件
|
||||||
|
window.toggleSwitch=(key,el)=>{
|
||||||
|
el.classList.toggle("on");
|
||||||
|
if(key==="preview_only") updateRunBtn();
|
||||||
|
updateBtn();
|
||||||
|
};
|
||||||
|
function updateRunBtn(){
|
||||||
|
const po=$("sw-preview_only").classList.contains("on");
|
||||||
|
$("btn-run").innerHTML = po ? "▶ 仅生成mask（快速预览）" : "▶ 执行完整换发型（40-90秒）";
|
||||||
|
}
|
||||||
|
// range 滑块绑定显示
|
||||||
|
["den"].forEach(id=>{
|
||||||
|
$(id).addEventListener("input",e=>$("v-"+id).textContent=parseFloat(e.target.value).toFixed(2));
|
||||||
|
});
|
||||||
|
// band_width 滑块（整数，HTML 里 oninput 直接调用全局函数）
|
||||||
|
window.syncRangeInt=(id)=>{ $("v-"+id).textContent=$(id).value; };
|
||||||
|
// 浮点滑块（HTML oninput 调用，联动显示）
|
||||||
|
window.syncRangeFloat=(id,fix)=>{ $("v-"+id).textContent=parseFloat($(id).value).toFixed(fix); };
|
||||||
|
window.syncNum=(id)=>{ $("v-"+id).textContent=$(id).value; };
|
||||||
|
window.resetParams=()=>{
|
||||||
|
Object.keys(defaults).forEach(k=>{
|
||||||
|
if(typeof defaults[k]==="boolean"){
|
||||||
|
const el=$("sw-"+k); el.classList.toggle("on",defaults[k]);
|
||||||
|
}else if(k==="method"){
|
||||||
|
// method 固定 landmark_1k，无 UI
|
||||||
|
}else{
|
||||||
|
const v=defaults[k];
|
||||||
|
$(k).value=v;
|
||||||
|
const lbl=$("v-"+k);
|
||||||
|
if(lbl){
|
||||||
|
// 浮点显示2位小数，整数原样
|
||||||
|
lbl.textContent = (typeof v==="number" && !Number.isInteger(v)) ? v.toFixed(2) : v;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
updateRunBtn();
|
||||||
|
};
|
||||||
|
|
||||||
|
function collectParams(){
|
||||||
|
return {
|
||||||
|
method:"landmark_1k",
|
||||||
|
strict_mask:$("sw-strict_mask").classList.contains("on"),
|
||||||
|
seamless_blend:$("sw-seamless_blend").classList.contains("on"),
|
||||||
|
is_hr:$("sw-is_hr").classList.contains("on"),
|
||||||
|
dilate_kernel:[parseInt($("dk0").value),parseInt($("dk1").value)],
|
||||||
|
denoising_strength:parseFloat($("den").value),
|
||||||
|
blend_dilate:[parseInt($("bd0").value),parseInt($("bd1").value)],
|
||||||
|
seamless_dilate:[parseInt($("sd0").value),parseInt($("sd1").value)],
|
||||||
|
band_width:parseInt($("bw").value),
|
||||||
|
preview_only:$("sw-preview_only").classList.contains("on"),
|
||||||
|
height_ratio:parseFloat($("hr").value),
|
||||||
|
width_ratio:parseFloat($("wr").value),
|
||||||
|
corner_ratio:parseFloat($("cr").value),
|
||||||
|
vertical_offset:parseFloat($("vo").value),
|
||||||
|
refiner_switch_at:parseFloat($("rsa").value),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
// 上传
|
||||||
|
$("upload-zone").onclick=()=>$("file-input").click();
|
||||||
|
$("file-input").onchange=e=>{ if(e.target.files[0]) loadImg(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])loadImg(e.dataTransfer.files[0]);};
|
||||||
|
function loadImg(file){
|
||||||
|
if(!file.type.startsWith("image/")){alert("请上传图片");return;}
|
||||||
|
const rd=new FileReader();
|
||||||
|
rd.onload=e=>{
|
||||||
|
imgB64=e.target.result;
|
||||||
|
$("preview-img").src=imgB64;
|
||||||
|
$("preview-img").style.display="block";
|
||||||
|
$("upload-zone").style.display="none";
|
||||||
|
updateBtn();
|
||||||
|
};
|
||||||
|
rd.readAsDataURL(file);
|
||||||
|
}
|
||||||
|
function updateBtn(){ $("btn-run").disabled=!(selectedHair&&imgB64); }
|
||||||
|
|
||||||
|
// 执行
|
||||||
|
$("btn-run").onclick=async()=>{
|
||||||
|
if(!selectedHair||!imgB64) return;
|
||||||
|
const p=collectParams();
|
||||||
|
$("btn-run").disabled=true;
|
||||||
|
$("status").style.display="block";
|
||||||
|
$("status-text").textContent = p.preview_only
|
||||||
|
? "生成发际线mask中（粗推理+方案），约5-15秒..."
|
||||||
|
: "执行完整换发型中（粗推理→mask→warpAffine→SD→贴回），约40-90秒...";
|
||||||
|
$("error-msg").style.display="none";
|
||||||
|
$("result-top").classList.remove("show");
|
||||||
|
$("steps-section").innerHTML="";
|
||||||
|
try{
|
||||||
|
const resp=await fetch("/api/swap_hairline",{
|
||||||
|
method:"POST",headers:{"Content-Type":"application/json"},
|
||||||
|
body:JSON.stringify(Object.assign({img:imgB64,hair_id:selectedHair},p))
|
||||||
|
});
|
||||||
|
const data=await resp.json();
|
||||||
|
$("status").style.display="none";
|
||||||
|
if(data.state===0){
|
||||||
|
const url="data:image/jpeg;base64,"+data.result;
|
||||||
|
$("r-orig").src=imgB64;$("r-new").src=url;
|
||||||
|
// 参数回显
|
||||||
|
let ph=""; Object.keys(data.params).forEach(k=>{
|
||||||
|
let v=data.params[k]; if(Array.isArray(v))v="["+v.join(",")+"]";
|
||||||
|
if(typeof v==="boolean")v=v?"开":"关";
|
||||||
|
ph+=k+" = "+v+" ";
|
||||||
|
});
|
||||||
|
$("params-used").textContent="本次参数: "+ph;
|
||||||
|
$("result-top").classList.add("show");
|
||||||
|
// 步骤画廊
|
||||||
|
const sc=$("steps-section");
|
||||||
|
data.steps.forEach(step=>{
|
||||||
|
const div=document.createElement("div");div.className="step-item";
|
||||||
|
let imgs='<div class="step-images">';
|
||||||
|
step.images.forEach(im=>{imgs+='<div class="step-img-card"><div class="lbl">'+im.label+'</div><img src="data:image/jpeg;base64,'+im.b64+'" data-full="data:image/jpeg;base64,'+im.b64+'"></div>';});
|
||||||
|
imgs+='</div>';
|
||||||
|
div.innerHTML='<div class="step-head"><div class="step-title">'+step.title+'</div></div><div class="step-desc">'+step.desc+'</div>'+imgs;
|
||||||
|
sc.appendChild(div);
|
||||||
|
});
|
||||||
|
sc.querySelectorAll("img").forEach(img=>img.onclick=()=>showLightbox(img.dataset.full));
|
||||||
|
$("result-top").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";
|
||||||
|
}
|
||||||
|
updateBtn();
|
||||||
|
};
|
||||||
|
|
||||||
|
function showLightbox(src){$("lightbox-img").src=src;$("lightbox").classList.add("show");}
|
||||||
|
|
||||||
|
updateRunBtn();
|
||||||
|
loadStatus();
|
||||||
|
})();
|
||||||
|
</script>
|
||||||
|
</body>
|
||||||
|
</html>
|
||||||
@@ -0,0 +1,242 @@
|
|||||||
|
<!DOCTYPE html>
|
||||||
|
<html lang="zh-CN">
|
||||||
|
<head>
|
||||||
|
<meta charset="UTF-8">
|
||||||
|
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||||
|
<title>区域生发测试</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; }
|
||||||
|
.header { background: #16213e; padding: 18px 28px; border-bottom: 1px solid #0f3460; }
|
||||||
|
.header h1 { font-size: 20px; font-weight: 600; }
|
||||||
|
.header p { font-size: 13px; color: #888; margin-top: 4px; }
|
||||||
|
.container { display: flex; gap: 20px; padding: 20px; max-width: 1700px; margin: 0 auto; }
|
||||||
|
.panel { width: 280px; flex-shrink: 0; background: #16213e; border-radius: 10px; padding: 18px; height: fit-content; max-height: 90vh; overflow-y: auto; }
|
||||||
|
.panel h3 { font-size: 14px; color: #4ecca3; margin-bottom: 12px; text-transform: uppercase; letter-spacing: 1px; }
|
||||||
|
.panel::-webkit-scrollbar { width: 6px; }
|
||||||
|
.panel::-webkit-scrollbar-thumb { background: #0f3460; border-radius: 3px; }
|
||||||
|
.hairstyle-tabs { display: flex; gap: 6px; margin-bottom: 10px; }
|
||||||
|
.hairstyle-tab { flex: 1; padding: 6px; text-align: center; background: #0f3460; border: none; border-radius: 4px; color: #aaa; cursor: pointer; font-size: 12px; }
|
||||||
|
.hairstyle-tab.active { background: #4ecca3; color: #16213e; font-weight: 600; }
|
||||||
|
.search-box { width: 100%; padding: 6px 10px; background: #0d1b3e; border: 1px solid #0f3460; border-radius: 4px; color: #eee; font-size: 12px; margin-bottom: 10px; }
|
||||||
|
.hairstyle-grid { display: grid; grid-template-columns: repeat(3, 1fr); gap: 6px; max-height: 280px; overflow-y: auto; }
|
||||||
|
.hairstyle-item { position: relative; cursor: pointer; border-radius: 4px; overflow: hidden; border: 2px solid transparent; aspect-ratio: 1; background: #0d1b3e; }
|
||||||
|
.hairstyle-item img { width: 100%; height: 100%; object-fit: cover; display: block; }
|
||||||
|
.hairstyle-item.selected { border-color: #4ecca3; }
|
||||||
|
.hairstyle-item.selected::after { content: '✓'; position: absolute; top: 2px; right: 4px; color: #4ecca3; font-weight: bold; text-shadow: 0 0 3px #000; }
|
||||||
|
.selected-hairstyle-info { font-size: 11px; color: #888; margin-top: 8px; text-align: center; word-break: break-all; }
|
||||||
|
.tool-group { margin-bottom: 16px; }
|
||||||
|
.tool-group label { display: block; font-size: 13px; color: #aaa; margin-bottom: 6px; }
|
||||||
|
.btn { display: block; width: 100%; padding: 9px; border: none; border-radius: 6px; font-size: 13px; cursor: pointer; margin-bottom: 8px; transition: .15s; }
|
||||||
|
.btn-tool { background: #0f3460; color: #eee; }
|
||||||
|
.btn-tool:hover { background: #1a4a7a; }
|
||||||
|
.btn-tool.active { background: #4ecca3; color: #16213e; font-weight: 600; }
|
||||||
|
.btn-primary { background: #4ecca3; color: #16213e; font-weight: 600; font-size: 15px; padding: 12px; }
|
||||||
|
.btn-primary:hover { background: #6ee0bd; }
|
||||||
|
.btn-primary:disabled { background: #555; color: #999; cursor: not-allowed; }
|
||||||
|
.btn-danger { background: #c0392b; color: #fff; font-size: 12px; }
|
||||||
|
.btn-danger:hover { background: #e74c3c; }
|
||||||
|
.hint { font-size: 12px; color: #666; line-height: 1.6; margin-top: 14px; padding: 10px; background: #0d1b3e; border-radius: 6px; }
|
||||||
|
.workspace { flex: 1; min-width: 0; }
|
||||||
|
.canvas-wrap { position: relative; background: #0d0d1a; border-radius: 10px; padding: 20px; text-align: center; min-height: 400px; display: flex; align-items: center; justify-content: center; }
|
||||||
|
.upload-zone { width: 100%; max-width: 500px; border: 2px dashed #0f3460; border-radius: 10px; padding: 50px 20px; text-align: center; cursor: pointer; transition: .2s; }
|
||||||
|
.upload-zone:hover, .upload-zone.dragover { border-color: #4ecca3; background: rgba(78,204,163,.05); }
|
||||||
|
.upload-zone svg { width: 48px; height: 48px; fill: #4ecca3; margin-bottom: 12px; }
|
||||||
|
.upload-zone p { color: #888; font-size: 14px; }
|
||||||
|
.upload-zone p.highlight { color: #4ecca3; margin-bottom: 4px; }
|
||||||
|
#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; }
|
||||||
|
.result-section { margin-top: 20px; background: #16213e; border-radius: 10px; padding: 18px; display: none; }
|
||||||
|
.result-section.show { display: block; }
|
||||||
|
.result-section h3 { font-size: 14px; color: #4ecca3; margin-bottom: 14px; }
|
||||||
|
.result-grid { display: flex; gap: 16px; flex-wrap: wrap; }
|
||||||
|
.result-card { flex: 1; min-width: 250px; }
|
||||||
|
.result-card h4 { font-size: 13px; color: #aaa; margin-bottom: 8px; text-align: center; }
|
||||||
|
.result-card img { width: 100%; border-radius: 6px; display: block; }
|
||||||
|
.download-btn { display: inline-block; margin-top: 8px; padding: 6px 14px; background: #0f3460; color: #eee; border-radius: 4px; font-size: 12px; text-decoration: none; }
|
||||||
|
#status { text-align: center; padding: 30px; color: #4ecca3; font-size: 15px; display: none; }
|
||||||
|
.spinner { display: inline-block; width: 20px; height: 20px; 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; text-align: center; padding: 20px; display: none; font-size: 14px; }
|
||||||
|
.loading-hairstyles { text-align: center; padding: 20px; color: #666; font-size: 12px; }
|
||||||
|
</style>
|
||||||
|
</head>
|
||||||
|
<body>
|
||||||
|
<div class="header">
|
||||||
|
<h1>🟢 区域生发（换发型工作流）</h1>
|
||||||
|
<p>选发型 → 上传人像 → 在缺发区涂抹遮罩 → 生成（走 LoRA + 换发型链路）</p>
|
||||||
|
</div>
|
||||||
|
<div class="container">
|
||||||
|
<div class="panel">
|
||||||
|
<h3>① 选择发型</h3>
|
||||||
|
<div class="hairstyle-tabs">
|
||||||
|
<button class="hairstyle-tab active" data-gender="all">全部</button>
|
||||||
|
<button class="hairstyle-tab" data-gender="girl">女款</button>
|
||||||
|
<button class="hairstyle-tab" data-gender="boy">男款</button>
|
||||||
|
</div>
|
||||||
|
<input type="text" class="search-box" id="search" placeholder="搜索发型ID...">
|
||||||
|
<div class="hairstyle-grid" id="hairstyle-grid"><div class="loading-hairstyles">加载发型中...</div></div>
|
||||||
|
<div class="selected-hairstyle-info" id="selected-info">未选择发型</div>
|
||||||
|
<h3 style="margin-top:20px">② 涂抹遮罩</h3>
|
||||||
|
<div class="tool-group">
|
||||||
|
<button class="btn btn-tool active" id="btn-brush">🖌️ 画笔</button>
|
||||||
|
<button class="btn btn-tool" id="btn-eraser">🧽 橡皮</button>
|
||||||
|
</div>
|
||||||
|
<div class="tool-group">
|
||||||
|
<label>画笔粗细 <span style="float:right;color:#4ecca3" id="size-val">25px</span></label>
|
||||||
|
<input type="range" id="brush-size" min="5" max="60" value="25" style="width:100%">
|
||||||
|
</div>
|
||||||
|
<button class="btn btn-danger" id="btn-clear">🗑️ 清空遮罩</button>
|
||||||
|
<h3 style="margin-top:20px">③ 生成</h3>
|
||||||
|
<div class="tool-group">
|
||||||
|
<label><input type="checkbox" id="is-hr" checked> 高清模式 (较慢)</label>
|
||||||
|
</div>
|
||||||
|
<button class="btn btn-primary" id="btn-generate" disabled>✨ 生成生发</button>
|
||||||
|
<div class="hint">
|
||||||
|
<b>使用说明</b><br>
|
||||||
|
1. 左侧选择一个发型<br>
|
||||||
|
2. 上传一张人头像照片<br>
|
||||||
|
3. 用画笔在头发稀少/发际线区域涂抹（红色=生发区）<br>
|
||||||
|
4. 点"生成"，等待30-90秒<br>
|
||||||
|
工作流：换发型粗推理 + LoRA + mask并集(原头发∪新发型∪手绘)
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
<div class="workspace">
|
||||||
|
<div class="canvas-wrap" id="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>建议人脸清晰正面照</p>
|
||||||
|
</div>
|
||||||
|
<input type="file" id="file-input" accept="image/*" style="display:none">
|
||||||
|
<div id="canvas-container" style="display:none">
|
||||||
|
<canvas id="img-canvas"></canvas>
|
||||||
|
<canvas id="mask-canvas"></canvas>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
<div id="status"><span class="spinner"></span><span id="status-text">正在生成，请稍候...</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="result-orig" alt="原图"></div>
|
||||||
|
<div class="result-card"><h4>生发结果</h4><img id="result-new" alt="结果"><a class="download-btn" id="download-link" download="hairgrow_result.jpg">⬇️ 下载结果</a></div>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
<script>
|
||||||
|
(function(){
|
||||||
|
const $=id=>document.getElementById(id);
|
||||||
|
const uploadZone=$("upload-zone"),fileInput=$("file-input"),canvasContainer=$("canvas-container");
|
||||||
|
const imgCanvas=$("img-canvas"),maskCanvas=$("mask-canvas");
|
||||||
|
const imgCtx=imgCanvas.getContext("2d"),maskCtx=maskCanvas.getContext("2d",{willReadFrequently:true});
|
||||||
|
let mode="brush",brushSize=25,drawing=false,lastX=0,lastY=0;
|
||||||
|
let selectedHairId=null,allHairstyles=[],imgB64Orig=null;
|
||||||
|
|
||||||
|
async function loadHairstyles(){
|
||||||
|
try{
|
||||||
|
const r=await fetch("/api/hairstyles");const d=await r.json();
|
||||||
|
if(d.state!==0)throw new Error(d.msg);
|
||||||
|
allHairstyles=d.data;renderHairstyles("all");
|
||||||
|
}catch(e){$("hairstyle-grid").innerHTML='<div class="loading-hairstyles">加载失败: '+e.message+'</div>';}
|
||||||
|
}
|
||||||
|
function renderHairstyles(gf){
|
||||||
|
const s=$("search").value.trim().toLowerCase();
|
||||||
|
const f=allHairstyles.filter(x=>{if(gf!=="all"&&x.gender!==gf)return false;if(s&&!x.hair_id.toLowerCase().includes(s))return false;return true;});
|
||||||
|
const g=$("hairstyle-grid");
|
||||||
|
if(f.length===0){g.innerHTML='<div class="loading-hairstyles">无匹配发型</div>';return;}
|
||||||
|
g.innerHTML="";
|
||||||
|
f.slice(0,60).forEach(x=>{
|
||||||
|
const it=document.createElement("div");
|
||||||
|
it.className="hairstyle-item"+(x.hair_id===selectedHairId?" selected":"");
|
||||||
|
it.dataset.hairId=x.hair_id;
|
||||||
|
const im=document.createElement("img");im.loading="lazy";im.src="/preview/"+x.hair_id;
|
||||||
|
im.onerror=()=>{im.style.visibility="hidden";};
|
||||||
|
it.appendChild(im);it.onclick=()=>selectHairstyle(x.hair_id);g.appendChild(it);
|
||||||
|
});
|
||||||
|
if(f.length>60){const m=document.createElement("div");m.className="loading-hairstyles";m.style.gridColumn="1/-1";m.textContent="还有 "+(f.length-60)+" 个，请搜索";g.appendChild(m);}
|
||||||
|
}
|
||||||
|
function selectHairstyle(id){
|
||||||
|
selectedHairId=id;
|
||||||
|
document.querySelectorAll(".hairstyle-item").forEach(el=>el.classList.toggle("selected",el.dataset.hairId===id));
|
||||||
|
$("selected-info").textContent="已选: "+id;updateGenerateBtn();
|
||||||
|
}
|
||||||
|
document.querySelectorAll(".hairstyle-tab").forEach(t=>{t.onclick=()=>{document.querySelectorAll(".hairstyle-tab").forEach(x=>x.classList.remove("active"));t.classList.add("active");renderHairstyles(t.dataset.gender);};});
|
||||||
|
$("search").addEventListener("input",()=>{const a=document.querySelector(".hairstyle-tab.active");renderHairstyles(a?a.dataset.gender:"all");});
|
||||||
|
loadHairstyles();
|
||||||
|
|
||||||
|
$("brush-size").addEventListener("input",e=>{brushSize=+e.target.value;$("size-val").textContent=brushSize+"px";});
|
||||||
|
$("btn-brush").addEventListener("click",()=>{mode="brush";$("btn-brush").classList.add("active");$("btn-eraser").classList.remove("active");maskCanvas.classList.remove("eraser");});
|
||||||
|
$("btn-eraser").addEventListener("click",()=>{mode="eraser";$("btn-eraser").classList.add("active");$("btn-brush").classList.remove("active");maskCanvas.classList.add("eraser");});
|
||||||
|
$("btn-clear").addEventListener("click",()=>{maskCtx.globalCompositeOperation="source-over";maskCtx.clearRect(0,0,maskCanvas.width,maskCanvas.height);});
|
||||||
|
|
||||||
|
uploadZone.addEventListener("click",()=>fileInput.click());
|
||||||
|
uploadZone.addEventListener("dragover",e=>{e.preventDefault();uploadZone.classList.add("dragover");});
|
||||||
|
uploadZone.addEventListener("dragleave",()=>uploadZone.classList.remove("dragover"));
|
||||||
|
uploadZone.addEventListener("drop",e=>{e.preventDefault();uploadZone.classList.remove("dragover");if(e.dataTransfer.files[0])loadImage(e.dataTransfer.files[0]);});
|
||||||
|
fileInput.addEventListener("change",e=>{if(e.target.files[0])loadImage(e.target.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(w>MAX||h>MAX){const r=Math.min(MAX/w,MAX/h);w=Math.round(w*r);h=Math.round(h*r);}
|
||||||
|
w=w-(w%8);h=h-(h%8);
|
||||||
|
imgCanvas.width=maskCanvas.width=w;imgCanvas.height=maskCanvas.height=h;
|
||||||
|
imgCtx.drawImage(im,0,0,w,h);imgB64Orig=imgCanvas.toDataURL("image/jpeg",0.95);
|
||||||
|
maskCtx.clearRect(0,0,w,h);
|
||||||
|
uploadZone.style.display="none";canvasContainer.style.display="inline-block";
|
||||||
|
updateGenerateBtn();$("result-section").classList.remove("show");$("error-msg").style.display="none";
|
||||||
|
};im.src=e.target.result;
|
||||||
|
};rd.readAsDataURL(file);
|
||||||
|
}
|
||||||
|
function updateGenerateBtn(){$("btn-generate").disabled=!(selectedHairId&&imgB64Orig);}
|
||||||
|
|
||||||
|
function getPos(e){const r=maskCanvas.getBoundingClientRect();const sx=maskCanvas.width/r.width,sy=maskCanvas.height/r.height;const p=e.touches?e.touches[0]:e;return{x:(p.clientX-r.left)*sx,y:(p.clientY-r.top)*sy};}
|
||||||
|
function startDraw(e){e.preventDefault();drawing=true;const p=getPos(e);lastX=p.x;lastY=p.y;drawDot(p.x,p.y);}
|
||||||
|
function draw(e){if(!drawing)return;e.preventDefault();const p=getPos(e);drawLine(lastX,lastY,p.x,p.y);lastX=p.x;lastY=p.y;}
|
||||||
|
function endDraw(){drawing=false;}
|
||||||
|
function drawDot(x,y){if(mode==="brush"){maskCtx.globalCompositeOperation="source-over";maskCtx.fillStyle="rgba(255,50,50,0.45)";}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(mode==="brush"){maskCtx.globalCompositeOperation="source-over";maskCtx.strokeStyle="rgba(255,50,50,0.45)";}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",draw);
|
||||||
|
maskCanvas.addEventListener("mouseup",endDraw);maskCanvas.addEventListener("mouseleave",endDraw);
|
||||||
|
maskCanvas.addEventListener("touchstart",startDraw,{passive:false});maskCanvas.addEventListener("touchmove",draw,{passive:false});
|
||||||
|
maskCanvas.addEventListener("touchend",endDraw);
|
||||||
|
|
||||||
|
$("btn-generate").addEventListener("click",async()=>{
|
||||||
|
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 isHr=$("is-hr").checked;
|
||||||
|
const tmp=document.createElement("canvas");tmp.width=maskCanvas.width;tmp.height=maskCanvas.height;
|
||||||
|
const tmpCtx=tmp.getContext("2d");
|
||||||
|
const md2=maskCtx.getImageData(0,0,maskCanvas.width,maskCanvas.height);
|
||||||
|
const out=tmpCtx.createImageData(tmp.width,tmp.height);
|
||||||
|
for(let i=0;i<md2.data.length;i+=4){const v=md2.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;}
|
||||||
|
tmpCtx.putImageData(out,0,0);
|
||||||
|
const maskB64=tmp.toDataURL("image/png");
|
||||||
|
$("btn-generate").disabled=true;$("status").style.display="block";
|
||||||
|
$("status-text").textContent="正在生成（换发型工作流+LoRA），约30-90秒...";
|
||||||
|
$("error-msg").style.display="none";$("result-section").classList.remove("show");
|
||||||
|
try{
|
||||||
|
const resp=await fetch("/api/grow",{method:"POST",headers:{"Content-Type":"application/json"},body:JSON.stringify({img:imgB64Orig,mask:maskB64,hair_id:selectedHairId,is_hr:String(isHr)})});
|
||||||
|
const data=await resp.json();$("status").style.display="none";
|
||||||
|
if(data.state===0){
|
||||||
|
const url="data:image/jpeg;base64,"+data.result;
|
||||||
|
$("result-orig").src=imgB64Orig;$("result-new").src=url;$("download-link").href=url;
|
||||||
|
$("result-section").classList.add("show");$("result-section").scrollIntoView({behavior:"smooth"});
|
||||||
|
maskCtx.clearRect(0,0,maskCanvas.width,maskCanvas.height);
|
||||||
|
}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";}
|
||||||
|
updateGenerateBtn();
|
||||||
|
});
|
||||||
|
})();
|
||||||
|
</script>
|
||||||
|
</body>
|
||||||
|
</html>
|
||||||
@@ -0,0 +1,427 @@
|
|||||||
|
<!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: 290px; flex-shrink: 0; background: #16213e; border-radius: 10px; padding: 16px; height: fit-content; max-height: 92vh; overflow-y: auto; }
|
||||||
|
.panel::-webkit-scrollbar { width: 6px; }
|
||||||
|
.panel::-webkit-scrollbar-thumb { background: #0f3460; border-radius: 3px; }
|
||||||
|
.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; }
|
||||||
|
.switch { display: flex; align-items: center; justify-content: space-between; padding: 6px 0; font-size: 13px; color: #ccc; cursor: pointer; }
|
||||||
|
.switch:hover { color: #eee; }
|
||||||
|
.toggle { position: relative; width: 36px; height: 20px; background: #0f3460; border-radius: 10px; transition: .2s; flex-shrink: 0; }
|
||||||
|
.toggle::after { content: ''; position: absolute; top: 2px; left: 2px; width: 16px; height: 16px; background: #888; border-radius: 50%; transition: .2s; }
|
||||||
|
.switch.on .toggle { background: #4ecca3; }
|
||||||
|
.switch.on .toggle::after { left: 18px; background: #16213e; }
|
||||||
|
.pair { display: flex; gap: 8px; }
|
||||||
|
.pair .field { flex: 1; }
|
||||||
|
.readonly-info { font-size: 11px; color: #777; padding: 8px 10px; background: #0d1b3e; border-radius: 5px; line-height: 1.6; }
|
||||||
|
.readonly-info code { color: #aaa; background: #16213e; padding: 1px 5px; border-radius: 3px; }
|
||||||
|
|
||||||
|
.btn { display: block; width: 100%; padding: 10px; border: none; border-radius: 6px; font-size: 14px; cursor: pointer; transition: .15s; }
|
||||||
|
.btn-primary { background: #4ecca3; color: #16213e; font-weight: 600; font-size: 15px; padding: 12px; margin-top: 8px; }
|
||||||
|
.btn-primary:hover { background: #6ee0bd; }
|
||||||
|
.btn-primary:disabled { background: #555; color: #999; cursor: not-allowed; }
|
||||||
|
|
||||||
|
/* 中间发型选择 */
|
||||||
|
.mid-panel { width: 200px; flex-shrink: 0; background: #16213e; border-radius: 10px; padding: 14px; height: fit-content; max-height: 92vh; overflow-y: auto; }
|
||||||
|
.mid-panel h3 { font-size: 12px; color: #4ecca3; margin-bottom: 10px; }
|
||||||
|
.hairstyle-grid { display: grid; grid-template-columns: 1fr 1fr; gap: 6px; }
|
||||||
|
.hairstyle-item { position: relative; cursor: pointer; border-radius: 5px; overflow: hidden; border: 2px solid transparent; aspect-ratio: 1; background: #0d1b3e; }
|
||||||
|
.hairstyle-item img { width: 100%; height: 100%; object-fit: cover; display: block; }
|
||||||
|
.hairstyle-item.ready { border-color: #2a5a4a; }
|
||||||
|
.hairstyle-item.selected { border-color: #4ecca3; }
|
||||||
|
.hairstyle-item.selected::after { content: '✓'; position: absolute; top: 1px; right: 3px; color: #4ecca3; font-weight: bold; text-shadow: 0 0 3px #000; font-size: 11px; }
|
||||||
|
.hairstyle-item .name { position: absolute; bottom: 0; left: 0; right: 0; background: linear-gradient(transparent, rgba(0,0,0,.85)); color: #fff; font-size: 10px; padding: 8px 2px 2px; text-align: center; }
|
||||||
|
.hairstyle-item.pending { opacity: 0.4; 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-top { margin-top: 16px; background: #16213e; border-radius: 10px; padding: 16px; display: none; }
|
||||||
|
.result-top.show { display: block; }
|
||||||
|
.result-top h3 { font-size: 14px; color: #4ecca3; margin-bottom: 12px; }
|
||||||
|
.result-pair { 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; }
|
||||||
|
.params-used { margin-top: 12px; padding: 10px 12px; background: #0d1b3e; border-radius: 6px; font-size: 11px; color: #888; line-height: 1.7; font-family: monospace; }
|
||||||
|
|
||||||
|
.steps-section { margin-top: 16px; }
|
||||||
|
.step-item { background: #16213e; border-radius: 10px; padding: 16px; margin-bottom: 14px; }
|
||||||
|
.step-title { font-size: 15px; font-weight: 600; color: #4ecca3; margin-bottom: 4px; }
|
||||||
|
.step-desc { font-size: 12px; color: #999; margin-bottom: 12px; line-height: 1.6; }
|
||||||
|
.step-images { display: flex; gap: 12px; flex-wrap: wrap; }
|
||||||
|
.step-img-card { flex: 1; min-width: 180px; max-width: 280px; }
|
||||||
|
.step-img-card .lbl { font-size: 11px; color: #888; margin-bottom: 4px; text-align: center; }
|
||||||
|
.step-img-card img { width: 100%; border-radius: 6px; display: block; border: 1px solid #0f3460; cursor: zoom-in; transition: .15s; }
|
||||||
|
.step-img-card img:hover { border-color: #4ecca3; }
|
||||||
|
#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; }
|
||||||
|
#lightbox img { max-width: 92%; max-height: 92%; border-radius: 8px; }
|
||||||
|
#lightbox.show { display: flex; }
|
||||||
|
</style>
|
||||||
|
</head>
|
||||||
|
<body>
|
||||||
|
<div class="header">
|
||||||
|
<div>
|
||||||
|
<h1>🖌️ 手绘mask换发型（重绘区由你画）</h1>
|
||||||
|
<p>上传人像 → 用画笔涂出要重绘的区域 → 选发型 → 生成。SD 只在你涂的区域重绘</p>
|
||||||
|
</div>
|
||||||
|
<div class="links">
|
||||||
|
<a href="/debug">→ 全自动调试台</a>
|
||||||
|
<a href="/hairline">→ 发际线带实验</a>
|
||||||
|
<a href="/test_new">→ 简易测试</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="switch on" id="sw-is_hr" onclick="toggleSwitch('is_hr', this)">
|
||||||
|
<span>高清模式 (1152x1536)</span>
|
||||||
|
<div class="toggle"></div>
|
||||||
|
</div>
|
||||||
|
<div class="switch" id="sw-strict_mask" onclick="toggleSwitch('strict_mask', this)">
|
||||||
|
<span>严格mask贴回 <span style="color:#666;font-size:10px">(mask外保留原图)</span></span>
|
||||||
|
<div class="toggle"></div>
|
||||||
|
</div>
|
||||||
|
<div class="switch on" id="sw-seamless_blend" onclick="toggleSwitch('seamless_blend', this)">
|
||||||
|
<span>泊松融合消接缝 <span style="color:#e9b949;font-size:10px">⚠需同时开「严格mask」</span></span>
|
||||||
|
<div class="toggle"></div>
|
||||||
|
</div>
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">denoising_strength <span class="val" id="v-den">0.60</span></div>
|
||||||
|
<input type="range" id="den" min="0.1" max="1.0" step="0.05" value="0.6" oninput="syncFloat('den')">
|
||||||
|
</div>
|
||||||
|
<div class="pair">
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">dilate_kernel <span class="val" id="v-dk0">6</span></div>
|
||||||
|
<input type="number" id="dk0" value="6" min="1" max="40" oninput="$('v-dk0').textContent=this.value">
|
||||||
|
</div>
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">(膨胀) <span class="val" id="v-dk1">18</span></div>
|
||||||
|
<input type="number" id="dk1" value="18" min="1" max="40" oninput="$('v-dk1').textContent=this.value">
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
<div class="field">
|
||||||
|
<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')">
|
||||||
|
</div>
|
||||||
|
<div class="readonly-info">
|
||||||
|
<b style="color:#888">webui端固定：</b><br>
|
||||||
|
<code>cfg=7</code> <code>steps=20</code> <code>sampler=DPM++ 2M Karras</code>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<h3 style="color:#4ecca3;margin-top:14px">③ enhance 二次增强（可选）</h3>
|
||||||
|
<div class="switch" id="sw-enhance" onclick="toggleSwitch('enhance', this)">
|
||||||
|
<span>开启enhance <span class="tip" title="换发型后对结果再做一次低强度SD重绘，让发丝更清晰（+15-20秒）">ⓘ</span></span>
|
||||||
|
<div class="toggle"></div>
|
||||||
|
</div>
|
||||||
|
<div class="field">
|
||||||
|
<div class="field-label">enhance_denoising <span class="val" id="v-ed">0.35</span><span class="tip" title="enhance重绘强度。越低越保留原图，越高发丝变化越大">ⓘ</span></div>
|
||||||
|
<input type="range" id="ed" min="0.1" max="0.8" step="0.05" value="0.35" oninput="syncFloat('ed')">
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<button class="btn btn-primary" id="btn-run" disabled>▶ 生成（手绘区换发型）</button>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<!-- 中间发型选择 -->
|
||||||
|
<div class="mid-panel">
|
||||||
|
<h3>③ 选发型</h3>
|
||||||
|
<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>
|
||||||
|
<button class="btn-tool" style="margin-top:8px;width:100%" onclick="loadStatus()">↻ 刷新</button>
|
||||||
|
</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>正面清晰照效果最佳</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">生成中，约40-90秒...</span></div>
|
||||||
|
<div id="error-msg"></div>
|
||||||
|
|
||||||
|
<div class="result-top" id="result-top">
|
||||||
|
<h3>最终对比</h3>
|
||||||
|
<div class="result-pair">
|
||||||
|
<div class="result-card"><h4>原图</h4><img id="r-orig"></div>
|
||||||
|
<div class="result-card"><h4>换发型结果（仅手绘区）</h4><img id="r-new"></div>
|
||||||
|
</div>
|
||||||
|
<div class="params-used" id="params-used"></div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div class="steps-section" id="steps-section"></div>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div id="lightbox" onclick="this.classList.remove('show')"><img id="lightbox-img"></div>
|
||||||
|
|
||||||
|
<script>
|
||||||
|
(function(){
|
||||||
|
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;
|
||||||
|
// 画板状态
|
||||||
|
const imgCanvas=$("img-canvas"),maskCanvas=$("mask-canvas");
|
||||||
|
const imgCtx=imgCanvas.getContext("2d"),maskCtx=maskCanvas.getContext("2d",{willReadFrequently:true});
|
||||||
|
let mode="brush",brushSize=25,drawing=false,lastX=0,lastY=0;
|
||||||
|
|
||||||
|
async function loadStatus(){
|
||||||
|
try{
|
||||||
|
const r=await fetch("/api/hairstyles");const d=await r.json();
|
||||||
|
if(d.state!==0)throw new Error(d.msg);
|
||||||
|
readySet=new Set(d.data.map(x=>x.hair_id));
|
||||||
|
}catch(e){readySet=new Set();}
|
||||||
|
renderGrid();
|
||||||
|
}
|
||||||
|
function renderGrid(){
|
||||||
|
const g=$("hairstyle-grid");g.innerHTML="";
|
||||||
|
HAIRSTYLES.forEach(h=>{
|
||||||
|
const ready=readySet.has(h.name);
|
||||||
|
const it=document.createElement("div");
|
||||||
|
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>`;
|
||||||
|
if(ready) it.onclick=()=>{
|
||||||
|
selectedHair=h.name;
|
||||||
|
$("sel-info").textContent="已选: "+h.name;
|
||||||
|
renderGrid(); updateBtn();
|
||||||
|
};
|
||||||
|
g.appendChild(it);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
window.toggleSwitch=(key,el)=>{ el.classList.toggle("on"); };
|
||||||
|
// 浮点滑块联动显示
|
||||||
|
window.syncFloat=(id)=>{ $("v-"+id).textContent=parseFloat($(id).value).toFixed(2); };
|
||||||
|
// 整数滑块联动显示
|
||||||
|
window.syncInt=(id)=>{ $("v-"+id).textContent=$(id).value; };
|
||||||
|
|
||||||
|
// 上传图片
|
||||||
|
$("upload-zone").onclick=()=>$("file-input").click();
|
||||||
|
$("file-input").onchange=e=>{ if(e.target.files[0]) loadImage(e.target.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";
|
||||||
|
updateBtn();
|
||||||
|
};
|
||||||
|
im.src=e.target.result;
|
||||||
|
};
|
||||||
|
rd.readAsDataURL(file);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 画笔工具切换
|
||||||
|
$("tool-brush").onclick=()=>{mode="brush";$("tool-brush").classList.add("active");$("tool-eraser").classList.remove("active");maskCanvas.classList.remove("eraser");};
|
||||||
|
$("tool-eraser").onclick=()=>{mode="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(mode==="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(mode==="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);
|
||||||
|
|
||||||
|
function updateBtn(){ $("btn-run").disabled=!(selectedHair&&imgB64Orig); }
|
||||||
|
|
||||||
|
// 提取mask为黑白二值PNG base64
|
||||||
|
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; // 用alpha判断是否涂抹过
|
||||||
|
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(!selectedHair||!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();
|
||||||
|
const params={
|
||||||
|
strict_mask:$("sw-strict_mask").classList.contains("on"),
|
||||||
|
seamless_blend:$("sw-seamless_blend").classList.contains("on"),
|
||||||
|
is_hr:$("sw-is_hr").classList.contains("on"),
|
||||||
|
dilate_kernel:[parseInt($("dk0").value),parseInt($("dk1").value)],
|
||||||
|
feather_px:parseInt($("fp").value),
|
||||||
|
denoising_strength:parseFloat($("den").value),
|
||||||
|
enhance:$("sw-enhance").classList.contains("on"),
|
||||||
|
enhance_denoising:parseFloat($("ed").value),
|
||||||
|
};
|
||||||
|
$("btn-run").disabled=true;
|
||||||
|
$("status").style.display="block";
|
||||||
|
$("status-text").textContent = params.enhance
|
||||||
|
? "生成中（换发型 + enhance增强），约60-110秒..."
|
||||||
|
: "生成中（粗推理→手绘mask→warpAffine→SD→贴回），约40-90秒...";
|
||||||
|
$("error-msg").style.display="none";
|
||||||
|
$("result-top").classList.remove("show");
|
||||||
|
$("steps-section").innerHTML="";
|
||||||
|
try{
|
||||||
|
const resp=await fetch("/api/swap_manual",{
|
||||||
|
method:"POST",headers:{"Content-Type":"application/json"},
|
||||||
|
body:JSON.stringify(Object.assign({img:imgB64Orig,mask:maskB64,hair_id:selectedHair},params))
|
||||||
|
});
|
||||||
|
const data=await resp.json();
|
||||||
|
$("status").style.display="none";
|
||||||
|
if(data.state===0){
|
||||||
|
const url="data:image/jpeg;base64,"+data.result;
|
||||||
|
$("r-orig").src=imgB64Orig;$("r-new").src=url;
|
||||||
|
let ph="";Object.keys(data.params).forEach(k=>{
|
||||||
|
let v=data.params[k];if(Array.isArray(v))v="["+v.join(",")+"]";
|
||||||
|
if(typeof v==="boolean")v=v?"开":"关";
|
||||||
|
ph+=k+" = "+v+" ";
|
||||||
|
});
|
||||||
|
$("params-used").textContent="本次参数: "+ph;
|
||||||
|
$("result-top").classList.add("show");
|
||||||
|
// 步骤画廊
|
||||||
|
const sc=$("steps-section");
|
||||||
|
data.steps.forEach(step=>{
|
||||||
|
const div=document.createElement("div");div.className="step-item";
|
||||||
|
let imgs='<div class="step-images">';
|
||||||
|
step.images.forEach(im=>{imgs+='<div class="step-img-card"><div class="lbl">'+im.label+'</div><img src="data:image/jpeg;base64,'+im.b64+'" data-full="data:image/jpeg;base64,'+im.b64+'"></div>';});
|
||||||
|
imgs+='</div>';
|
||||||
|
div.innerHTML='<div class="step-title">'+step.title+'</div><div class="step-desc">'+step.desc+'</div>'+imgs;
|
||||||
|
sc.appendChild(div);
|
||||||
|
});
|
||||||
|
sc.querySelectorAll("img").forEach(img=>img.onclick=()=>showLightbox(img.dataset.full));
|
||||||
|
$("result-top").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";
|
||||||
|
}
|
||||||
|
updateBtn();
|
||||||
|
};
|
||||||
|
|
||||||
|
function showLightbox(src){$("lightbox-img").src=src;$("lightbox").classList.add("show");}
|
||||||
|
loadStatus();
|
||||||
|
})();
|
||||||
|
</script>
|
||||||
|
</body>
|
||||||
|
</html>
|
||||||
@@ -0,0 +1,210 @@
|
|||||||
|
<!DOCTYPE html>
|
||||||
|
<html lang="zh-CN">
|
||||||
|
<head>
|
||||||
|
<meta charset="UTF-8">
|
||||||
|
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||||
|
<title>换发型测试</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; }
|
||||||
|
.header { background: #16213e; padding: 18px 28px; border-bottom: 1px solid #0f3460; }
|
||||||
|
.header h1 { font-size: 20px; font-weight: 600; }
|
||||||
|
.header p { font-size: 13px; color: #888; margin-top: 4px; }
|
||||||
|
.container { display: flex; gap: 20px; padding: 20px; max-width: 1700px; margin: 0 auto; }
|
||||||
|
.panel { width: 280px; flex-shrink: 0; background: #16213e; border-radius: 10px; padding: 18px; height: fit-content; max-height: 90vh; overflow-y: auto; }
|
||||||
|
.panel h3 { font-size: 14px; color: #4ecca3; margin-bottom: 12px; text-transform: uppercase; letter-spacing: 1px; }
|
||||||
|
.panel::-webkit-scrollbar { width: 6px; }
|
||||||
|
.panel::-webkit-scrollbar-thumb { background: #0f3460; border-radius: 3px; }
|
||||||
|
.hairstyle-tabs { display: flex; gap: 6px; margin-bottom: 10px; }
|
||||||
|
.hairstyle-tab { flex: 1; padding: 6px; text-align: center; background: #0f3460; border: none; border-radius: 4px; color: #aaa; cursor: pointer; font-size: 12px; }
|
||||||
|
.hairstyle-tab.active { background: #4ecca3; color: #16213e; font-weight: 600; }
|
||||||
|
.search-box { width: 100%; padding: 6px 10px; background: #0d1b3e; border: 1px solid #0f3460; border-radius: 4px; color: #eee; font-size: 12px; margin-bottom: 10px; }
|
||||||
|
.hairstyle-grid { display: grid; grid-template-columns: repeat(3, 1fr); gap: 6px; max-height: 420px; overflow-y: auto; }
|
||||||
|
.hairstyle-item { position: relative; cursor: pointer; border-radius: 4px; overflow: hidden; border: 2px solid transparent; aspect-ratio: 1; background: #0d1b3e; }
|
||||||
|
.hairstyle-item img { width: 100%; height: 100%; object-fit: cover; display: block; }
|
||||||
|
.hairstyle-item.selected { border-color: #4ecca3; }
|
||||||
|
.hairstyle-item.selected::after { content: '✓'; position: absolute; top: 2px; right: 4px; color: #4ecca3; font-weight: bold; text-shadow: 0 0 3px #000; }
|
||||||
|
.selected-hairstyle-info { font-size: 11px; color: #888; margin-top: 8px; text-align: center; word-break: break-all; }
|
||||||
|
.tool-group { margin-bottom: 16px; }
|
||||||
|
.tool-group label { display: block; font-size: 13px; color: #aaa; margin-bottom: 6px; }
|
||||||
|
.btn { display: block; width: 100%; padding: 9px; border: none; border-radius: 6px; font-size: 13px; cursor: pointer; margin-bottom: 8px; transition: .15s; }
|
||||||
|
.btn-primary { background: #4ecca3; color: #16213e; font-weight: 600; font-size: 15px; padding: 12px; }
|
||||||
|
.btn-primary:hover { background: #6ee0bd; }
|
||||||
|
.btn-primary:disabled { background: #555; color: #999; cursor: not-allowed; }
|
||||||
|
.hint { font-size: 12px; color: #666; line-height: 1.6; margin-top: 14px; padding: 10px; background: #0d1b3e; border-radius: 6px; }
|
||||||
|
.workspace { flex: 1; min-width: 0; }
|
||||||
|
.canvas-wrap { position: relative; background: #0d0d1a; border-radius: 10px; padding: 20px; text-align: center; min-height: 400px; display: flex; align-items: center; justify-content: center; }
|
||||||
|
.upload-zone { width: 100%; max-width: 500px; border: 2px dashed #0f3460; border-radius: 10px; padding: 50px 20px; text-align: center; cursor: pointer; transition: .2s; }
|
||||||
|
.upload-zone:hover, .upload-zone.dragover { border-color: #4ecca3; background: rgba(78,204,163,.05); }
|
||||||
|
.upload-zone svg { width: 48px; height: 48px; fill: #4ecca3; margin-bottom: 12px; }
|
||||||
|
.upload-zone p { color: #888; font-size: 14px; }
|
||||||
|
.upload-zone p.highlight { color: #4ecca3; margin-bottom: 4px; }
|
||||||
|
#preview-img { display: none; max-width: 100%; max-height: 500px; border-radius: 6px; }
|
||||||
|
.result-section { margin-top: 20px; background: #16213e; border-radius: 10px; padding: 18px; display: none; }
|
||||||
|
.result-section.show { display: block; }
|
||||||
|
.result-section h3 { font-size: 14px; color: #4ecca3; margin-bottom: 14px; }
|
||||||
|
.result-grid { display: flex; gap: 16px; flex-wrap: wrap; }
|
||||||
|
.result-card { flex: 1; min-width: 250px; }
|
||||||
|
.result-card h4 { font-size: 13px; color: #aaa; margin-bottom: 8px; text-align: center; }
|
||||||
|
.result-card img { width: 100%; border-radius: 6px; display: block; }
|
||||||
|
.download-btn { display: inline-block; margin-top: 8px; padding: 6px 14px; background: #0f3460; color: #eee; border-radius: 4px; font-size: 12px; text-decoration: none; }
|
||||||
|
/* 步骤画廊 */
|
||||||
|
.steps-section { margin-top: 20px; background: #16213e; border-radius: 10px; padding: 18px; display: none; }
|
||||||
|
.steps-section.show { display: block; }
|
||||||
|
.steps-section h3 { font-size: 14px; color: #4ecca3; margin-bottom: 6px; }
|
||||||
|
.steps-section .sub { font-size: 12px; color: #666; margin-bottom: 16px; }
|
||||||
|
.step-item { margin-bottom: 24px; padding-bottom: 20px; border-bottom: 1px solid #0f3460; }
|
||||||
|
.step-item:last-child { border-bottom: none; }
|
||||||
|
.step-title { font-size: 15px; font-weight: 600; color: #4ecca3; margin-bottom: 4px; }
|
||||||
|
.step-desc { font-size: 12px; color: #999; margin-bottom: 12px; line-height: 1.5; }
|
||||||
|
.step-images { display: flex; gap: 12px; flex-wrap: wrap; }
|
||||||
|
.step-img-card { flex: 1; min-width: 180px; max-width: 280px; }
|
||||||
|
.step-img-card .lbl { font-size: 11px; color: #888; margin-bottom: 4px; text-align: center; }
|
||||||
|
.step-img-card img { width: 100%; border-radius: 6px; display: block; border: 1px solid #0f3460; cursor: zoom-in; }
|
||||||
|
.step-img-card img:hover { border-color: #4ecca3; }
|
||||||
|
#status { text-align: center; padding: 30px; color: #4ecca3; font-size: 15px; display: none; }
|
||||||
|
.spinner { display: inline-block; width: 20px; height: 20px; 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; text-align: center; padding: 20px; display: none; font-size: 14px; }
|
||||||
|
.loading-hairstyles { text-align: center; padding: 20px; color: #666; font-size: 12px; }
|
||||||
|
.link-row { margin-top: 12px; text-align: center; }
|
||||||
|
.link-row a { color: #4ecca3; font-size: 12px; text-decoration: none; }
|
||||||
|
/* 图片放大查看 */
|
||||||
|
#lightbox { display: none; position: fixed; top: 0; left: 0; width: 100%; height: 100%; background: rgba(0,0,0,.9); z-index: 999; justify-content: center; align-items: center; cursor: zoom-out; }
|
||||||
|
#lightbox img { max-width: 90%; max-height: 90%; border-radius: 8px; }
|
||||||
|
#lightbox.show { display: flex; }
|
||||||
|
</style>
|
||||||
|
</head>
|
||||||
|
<body>
|
||||||
|
<div class="header">
|
||||||
|
<h1>💇 换发型测试（过程可视化）</h1>
|
||||||
|
<p>选发型 → 上传人像 → 生成（展示换发型每个步骤的中间产物）</p>
|
||||||
|
</div>
|
||||||
|
<div class="container">
|
||||||
|
<div class="panel">
|
||||||
|
<h3>① 选择发型</h3>
|
||||||
|
<div class="hairstyle-tabs">
|
||||||
|
<button class="hairstyle-tab active" data-gender="all">全部</button>
|
||||||
|
<button class="hairstyle-tab" data-gender="girl">女款</button>
|
||||||
|
<button class="hairstyle-tab" data-gender="boy">男款</button>
|
||||||
|
</div>
|
||||||
|
<input type="text" class="search-box" id="search" placeholder="搜索发型ID...">
|
||||||
|
<div class="hairstyle-grid" id="hairstyle-grid"><div class="loading-hairstyles">加载发型中...</div></div>
|
||||||
|
<div class="selected-hairstyle-info" id="selected-info">未选择发型</div>
|
||||||
|
<h3 style="margin-top:20px">② 上传人像</h3>
|
||||||
|
<div class="hint" style="margin-top:0">点击右侧上传区选择人头像照片</div>
|
||||||
|
<h3 style="margin-top:20px">③ 生成</h3>
|
||||||
|
<div class="tool-group">
|
||||||
|
<label><input type="checkbox" id="is-hr"> 高清模式 (较慢)</label>
|
||||||
|
<label><input type="checkbox" id="strict-mask"> 严格按mask贴回 (mask外不变)</label>
|
||||||
|
</div>
|
||||||
|
<button class="btn btn-primary" id="btn-generate" disabled>✨ 换发型</button>
|
||||||
|
<div class="hint">
|
||||||
|
<b>说明</b><br>
|
||||||
|
换发型后会展示6个步骤：粗推理→生成mask→warpAffine→SD推理→贴回原图
|
||||||
|
</div>
|
||||||
|
<div class="link-row"><a href="/test_new">→ 去新发型测试页（30款）</a></div>
|
||||||
|
<div class="link-row"><a href="/">→ 去生发测试页</a></div>
|
||||||
|
</div>
|
||||||
|
<div class="workspace">
|
||||||
|
<div class="canvas-wrap" id="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>建议人脸清晰正面照</p>
|
||||||
|
</div>
|
||||||
|
<input type="file" id="file-input" accept="image/*" style="display:none">
|
||||||
|
<img id="preview-img">
|
||||||
|
</div>
|
||||||
|
<div id="status"><span class="spinner"></span><span id="status-text">正在换发型（含可视化），约60秒...</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="result-orig"></div>
|
||||||
|
<div class="result-card"><h4>换发型结果</h4><img id="result-new"><a class="download-btn" id="download-link" download="swaphair_result.jpg">⬇️ 下载结果</a></div>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
<div class="steps-section" id="steps-section">
|
||||||
|
<h3>🔍 换发型过程（逐步可视化）</h3>
|
||||||
|
<div class="sub">每个步骤用到的图片和说明，点击图片可放大查看</div>
|
||||||
|
<div id="steps-container"></div>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
<div id="lightbox" onclick="this.classList.remove('show')"><img id="lightbox-img"></div>
|
||||||
|
<script>
|
||||||
|
(function(){
|
||||||
|
const $=id=>document.getElementById(id);
|
||||||
|
const uploadZone=$("upload-zone"),fileInput=$("file-input"),previewImg=$("preview-img");
|
||||||
|
let selectedHairId=null,allHairstyles=[],imgB64Orig=null;
|
||||||
|
|
||||||
|
async function loadHairstyles(){
|
||||||
|
try{const r=await fetch("/api/hairstyles");const d=await r.json();
|
||||||
|
if(d.state!==0)throw new Error(d.msg);allHairstyles=d.data;renderHairstyles("all");
|
||||||
|
}catch(e){$("hairstyle-grid").innerHTML='<div class="loading-hairstyles">加载失败: '+e.message+'</div>';}
|
||||||
|
}
|
||||||
|
function renderHairstyles(gf){
|
||||||
|
const s=$("search").value.trim().toLowerCase();
|
||||||
|
const f=allHairstyles.filter(x=>{if(gf!=="all"&&x.gender!==gf)return false;if(s&&!x.hair_id.toLowerCase().includes(s))return false;return true;});
|
||||||
|
const g=$("hairstyle-grid");if(f.length===0){g.innerHTML='<div class="loading-hairstyles">无匹配发型</div>';return;}
|
||||||
|
g.innerHTML="";f.slice(0,90).forEach(x=>{
|
||||||
|
const it=document.createElement("div");it.className="hairstyle-item"+(x.hair_id===selectedHairId?" selected":"");it.dataset.hairId=x.hair_id;
|
||||||
|
const im=document.createElement("img");im.loading="lazy";im.src="/preview/"+x.hair_id;im.onerror=()=>{im.style.visibility="hidden";};
|
||||||
|
it.appendChild(im);it.onclick=()=>selectHairstyle(x.hair_id);g.appendChild(it);
|
||||||
|
});
|
||||||
|
if(f.length>90){const m=document.createElement("div");m.className="loading-hairstyles";m.style.gridColumn="1/-1";m.textContent="还有 "+(f.length-90)+" 个，请搜索";g.appendChild(m);}
|
||||||
|
}
|
||||||
|
function selectHairstyle(id){selectedHairId=id;document.querySelectorAll(".hairstyle-item").forEach(el=>el.classList.toggle("selected",el.dataset.hairId===id));$("selected-info").textContent="已选: "+id;updateBtn();}
|
||||||
|
document.querySelectorAll(".hairstyle-tab").forEach(t=>{t.onclick=()=>{document.querySelectorAll(".hairstyle-tab").forEach(x=>x.classList.remove("active"));t.classList.add("active");renderHairstyles(t.dataset.gender);};});
|
||||||
|
$("search").addEventListener("input",()=>{const a=document.querySelector(".hairstyle-tab.active");renderHairstyles(a?a.dataset.gender:"all");});
|
||||||
|
loadHairstyles();
|
||||||
|
|
||||||
|
uploadZone.addEventListener("click",()=>fileInput.click());
|
||||||
|
uploadZone.addEventListener("dragover",e=>{e.preventDefault();uploadZone.classList.add("dragover");});
|
||||||
|
uploadZone.addEventListener("dragleave",()=>uploadZone.classList.remove("dragover"));
|
||||||
|
uploadZone.addEventListener("drop",e=>{e.preventDefault();uploadZone.classList.remove("dragover");if(e.dataTransfer.files[0])loadImage(e.dataTransfer.files[0]);});
|
||||||
|
fileInput.addEventListener("change",e=>{if(e.target.files[0])loadImage(e.target.files[0]);});
|
||||||
|
function loadImage(file){if(!file.type.startsWith("image/")){alert("请上传图片文件");return;}
|
||||||
|
const rd=new FileReader();rd.onload=e=>{imgB64Orig=e.target.result;previewImg.src=imgB64Orig;previewImg.style.display="block";uploadZone.style.display="none";updateBtn();$("result-section").classList.remove("show");$("steps-section").classList.remove("show");$("error-msg").style.display="none";};rd.readAsDataURL(file);}
|
||||||
|
function updateBtn(){$("btn-generate").disabled=!(selectedHairId&&imgB64Orig);}
|
||||||
|
|
||||||
|
function showLightbox(src){$("lightbox-img").src=src;$("lightbox").classList.add("show");}
|
||||||
|
|
||||||
|
$("btn-generate").addEventListener("click",async()=>{
|
||||||
|
const isHr=$("is-hr").checked;
|
||||||
|
$("btn-generate").disabled=true;$("status").style.display="block";
|
||||||
|
$("status-text").textContent="正在换发型（含过程可视化），约60秒...";
|
||||||
|
$("error-msg").style.display="none";$("result-section").classList.remove("show");$("steps-section").classList.remove("show");
|
||||||
|
try{
|
||||||
|
const resp=await fetch("/api/swap_viz",{method:"POST",headers:{"Content-Type":"application/json"},
|
||||||
|
body:JSON.stringify({img:imgB64Orig,hair_id:selectedHairId,is_hr:String(isHr),strict_mask:$("strict-mask").checked})});
|
||||||
|
const data=await resp.json();$("status").style.display="none";
|
||||||
|
if(data.state===0){
|
||||||
|
const url="data:image/jpeg;base64,"+data.result;
|
||||||
|
$("result-orig").src=imgB64Orig;$("result-new").src=url;$("download-link").href=url;
|
||||||
|
$("result-section").classList.add("show");
|
||||||
|
// 渲染步骤画廊
|
||||||
|
if(data.steps&&data.steps.length){
|
||||||
|
const c=$("steps-container");c.innerHTML="";
|
||||||
|
data.steps.forEach(step=>{
|
||||||
|
const div=document.createElement("div");div.className="step-item";
|
||||||
|
let imgsHtml='<div class="step-images">';
|
||||||
|
step.images.forEach(im=>{imgsHtml+='<div class="step-img-card"><div class="lbl">'+im.label+'</div><img src="data:image/jpeg;base64,'+im.b64+'" data-full="data:image/jpeg;base64,'+im.b64+'"></div>';});
|
||||||
|
imgsHtml+='</div>';
|
||||||
|
div.innerHTML='<div class="step-title">'+step.title+'</div><div class="step-desc">'+step.desc+'</div>'+imgsHtml;
|
||||||
|
c.appendChild(div);
|
||||||
|
});
|
||||||
|
// 点击放大
|
||||||
|
c.querySelectorAll("img").forEach(img=>img.onclick=()=>showLightbox(img.dataset.full));
|
||||||
|
$("steps-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";}
|
||||||
|
updateBtn();
|
||||||
|
});
|
||||||
|
})();
|
||||||
|
</script>
|
||||||
|
</body>
|
||||||
|
</html>
|
||||||
@@ -0,0 +1,298 @@
|
|||||||
|
<!DOCTYPE html>
|
||||||
|
<html lang="zh-CN">
|
||||||
|
<head>
|
||||||
|
<meta charset="UTF-8">
|
||||||
|
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||||
|
<title>新发型效果测试</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; }
|
||||||
|
.header { background: #16213e; padding: 18px 28px; border-bottom: 1px solid #0f3460; }
|
||||||
|
.header h1 { font-size: 20px; font-weight: 600; }
|
||||||
|
.header p { font-size: 13px; color: #888; margin-top: 4px; }
|
||||||
|
.container { display: flex; gap: 20px; padding: 20px; max-width: 1700px; margin: 0 auto; }
|
||||||
|
.panel { width: 320px; flex-shrink: 0; background: #16213e; border-radius: 10px; padding: 18px; height: fit-content; max-height: 90vh; overflow-y: auto; }
|
||||||
|
.panel h3 { font-size: 14px; color: #4ecca3; margin-bottom: 12px; text-transform: uppercase; letter-spacing: 1px; }
|
||||||
|
.panel::-webkit-scrollbar { width: 6px; }
|
||||||
|
.panel::-webkit-scrollbar-thumb { background: #0f3460; border-radius: 3px; }
|
||||||
|
.face-tabs { display: flex; flex-wrap: wrap; gap: 6px; margin-bottom: 12px; }
|
||||||
|
.face-tab { padding: 5px 12px; background: #0f3460; border: none; border-radius: 4px; color: #aaa; cursor: pointer; font-size: 12px; }
|
||||||
|
.face-tab.active { background: #4ecca3; color: #16213e; font-weight: 600; }
|
||||||
|
.status-row { font-size: 12px; color: #888; margin-bottom: 10px; padding: 6px 10px; background: #0d1b3e; border-radius: 4px; }
|
||||||
|
.status-row b { color: #4ecca3; }
|
||||||
|
.hairstyle-grid { display: grid; grid-template-columns: repeat(2, 1fr); gap: 8px; max-height: 500px; overflow-y: auto; }
|
||||||
|
.hairstyle-item { position: relative; cursor: pointer; border-radius: 6px; overflow: hidden; border: 2px solid transparent; aspect-ratio: 1; background: #0d1b3e; }
|
||||||
|
.hairstyle-item img { width: 100%; height: 100%; object-fit: cover; display: block; }
|
||||||
|
.hairstyle-item.ready { border-color: #2a5a4a; }
|
||||||
|
.hairstyle-item.selected { border-color: #4ecca3; }
|
||||||
|
.hairstyle-item.selected::after { content: '✓'; position: absolute; top: 2px; right: 4px; color: #4ecca3; font-weight: bold; text-shadow: 0 0 3px #000; }
|
||||||
|
.hairstyle-item .name { position: absolute; bottom: 0; left: 0; right: 0; background: linear-gradient(transparent, rgba(0,0,0,.85)); color: #fff; font-size: 11px; padding: 12px 4px 3px; text-align: center; }
|
||||||
|
.hairstyle-item .badge { position: absolute; top: 3px; left: 3px; background: #4ecca3; color: #16213e; font-size: 9px; padding: 1px 5px; border-radius: 3px; font-weight: 600; }
|
||||||
|
.hairstyle-item .badge.training { background: #e9b949; }
|
||||||
|
.hairstyle-item .badge.pending { background: #555; color: #aaa; }
|
||||||
|
.hairstyle-item.pending { opacity: 0.45; cursor: not-allowed; }
|
||||||
|
.selected-hairstyle-info { font-size: 12px; color: #4ecca3; margin-top: 10px; text-align: center; padding: 6px; background: #0d1b3e; border-radius: 4px; }
|
||||||
|
.tool-group { margin-bottom: 16px; }
|
||||||
|
.tool-group label { display: block; font-size: 13px; color: #aaa; margin-bottom: 6px; cursor: pointer; }
|
||||||
|
.tool-group label input { margin-right: 6px; }
|
||||||
|
.btn { display: block; width: 100%; padding: 9px; border: none; border-radius: 6px; font-size: 13px; cursor: pointer; margin-bottom: 8px; transition: .15s; }
|
||||||
|
.btn-primary { background: #4ecca3; color: #16213e; font-weight: 600; font-size: 15px; padding: 12px; }
|
||||||
|
.btn-primary:hover { background: #6ee0bd; }
|
||||||
|
.btn-primary:disabled { background: #555; color: #999; cursor: not-allowed; }
|
||||||
|
.btn-ghost { background: #0f3460; color: #eee; }
|
||||||
|
.btn-ghost:hover { background: #1a4a80; }
|
||||||
|
.hint { font-size: 12px; color: #666; line-height: 1.6; margin-top: 14px; padding: 10px; background: #0d1b3e; border-radius: 6px; }
|
||||||
|
.workspace { flex: 1; min-width: 0; }
|
||||||
|
.canvas-wrap { position: relative; background: #0d0d1a; border-radius: 10px; padding: 20px; text-align: center; min-height: 320px; display: flex; align-items: center; justify-content: center; }
|
||||||
|
.upload-zone { width: 100%; max-width: 500px; border: 2px dashed #0f3460; border-radius: 10px; padding: 50px 20px; text-align: center; cursor: pointer; transition: .2s; }
|
||||||
|
.upload-zone:hover, .upload-zone.dragover { border-color: #4ecca3; background: rgba(78,204,163,.05); }
|
||||||
|
.upload-zone svg { width: 48px; height: 48px; fill: #4ecca3; margin-bottom: 12px; }
|
||||||
|
.upload-zone p { color: #888; font-size: 14px; }
|
||||||
|
.upload-zone p.highlight { color: #4ecca3; margin-bottom: 4px; }
|
||||||
|
#preview-img { display: none; max-width: 100%; max-height: 500px; border-radius: 6px; }
|
||||||
|
.result-section { margin-top: 20px; background: #16213e; border-radius: 10px; padding: 18px; display: none; }
|
||||||
|
.result-section.show { display: block; }
|
||||||
|
.result-section h3 { font-size: 14px; color: #4ecca3; margin-bottom: 14px; }
|
||||||
|
.result-grid { display: flex; gap: 16px; flex-wrap: wrap; }
|
||||||
|
.result-card { flex: 1; min-width: 250px; }
|
||||||
|
.result-card h4 { font-size: 13px; color: #aaa; margin-bottom: 8px; text-align: center; }
|
||||||
|
.result-card img { width: 100%; border-radius: 6px; display: block; }
|
||||||
|
.download-btn { display: inline-block; margin-top: 8px; padding: 6px 14px; background: #0f3460; color: #eee; border-radius: 4px; font-size: 12px; text-decoration: none; }
|
||||||
|
.steps-section { margin-top: 20px; background: #16213e; border-radius: 10px; padding: 18px; display: none; }
|
||||||
|
.steps-section.show { display: block; }
|
||||||
|
.steps-section h3 { font-size: 14px; color: #4ecca3; margin-bottom: 6px; }
|
||||||
|
.steps-section .sub { font-size: 12px; color: #666; margin-bottom: 16px; }
|
||||||
|
.step-item { margin-bottom: 24px; padding-bottom: 20px; border-bottom: 1px solid #0f3460; }
|
||||||
|
.step-item:last-child { border-bottom: none; }
|
||||||
|
.step-title { font-size: 15px; font-weight: 600; color: #4ecca3; margin-bottom: 4px; }
|
||||||
|
.step-desc { font-size: 12px; color: #999; margin-bottom: 12px; line-height: 1.5; }
|
||||||
|
.step-images { display: flex; gap: 12px; flex-wrap: wrap; }
|
||||||
|
.step-img-card { flex: 1; min-width: 180px; max-width: 280px; }
|
||||||
|
.step-img-card .lbl { font-size: 11px; color: #888; margin-bottom: 4px; text-align: center; }
|
||||||
|
.step-img-card img { width: 100%; border-radius: 6px; display: block; border: 1px solid #0f3460; cursor: zoom-in; }
|
||||||
|
.step-img-card img:hover { border-color: #4ecca3; }
|
||||||
|
#status { text-align: center; padding: 30px; color: #4ecca3; font-size: 15px; display: none; }
|
||||||
|
.spinner { display: inline-block; width: 20px; height: 20px; 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; text-align: center; padding: 20px; display: none; font-size: 14px; }
|
||||||
|
.loading-hairstyles { text-align: center; padding: 20px; color: #666; font-size: 12px; }
|
||||||
|
.link-row { margin-top: 12px; text-align: center; }
|
||||||
|
.link-row a { color: #4ecca3; font-size: 12px; text-decoration: none; }
|
||||||
|
#lightbox { display: none; position: fixed; top: 0; left: 0; width: 100%; height: 100%; background: rgba(0,0,0,.9); z-index: 999; justify-content: center; align-items: center; cursor: zoom-out; }
|
||||||
|
#lightbox img { max-width: 90%; max-height: 90%; border-radius: 8px; }
|
||||||
|
#lightbox.show { display: flex; }
|
||||||
|
</style>
|
||||||
|
</head>
|
||||||
|
<body>
|
||||||
|
<div class="header">
|
||||||
|
<h1>💇 新发型效果测试（30款 · 按脸型适配）</h1>
|
||||||
|
<p>选择发型 → 上传人像 → 生成对比。绿色徽章=可测试，黄色=训练中，灰色=待训练</p>
|
||||||
|
</div>
|
||||||
|
<div class="container">
|
||||||
|
<div class="panel">
|
||||||
|
<h3>① 选择发型</h3>
|
||||||
|
<div class="status-row" id="status-row">加载中...</div>
|
||||||
|
<div class="face-tabs" id="face-tabs"></div>
|
||||||
|
<div class="hairstyle-grid" id="hairstyle-grid"><div class="loading-hairstyles">加载中...</div></div>
|
||||||
|
<div class="selected-hairstyle-info" id="selected-info">未选择发型</div>
|
||||||
|
|
||||||
|
<h3 style="margin-top:20px">② 上传人像</h3>
|
||||||
|
<div class="hint" style="margin-top:0">点击右侧上传区选择一张人头像照片（正面、清晰）</div>
|
||||||
|
|
||||||
|
<h3 style="margin-top:20px">③ 生成</h3>
|
||||||
|
<div class="tool-group">
|
||||||
|
<label><input type="checkbox" id="is-hr"> 高清模式（更清晰但更慢）</label>
|
||||||
|
<label><input type="checkbox" id="strict-mask"> 严格按mask贴回（mask外不变）</label>
|
||||||
|
</div>
|
||||||
|
<button class="btn btn-primary" id="btn-generate" disabled>✨ 换发型</button>
|
||||||
|
<button class="btn btn-ghost" id="btn-refresh">↻ 刷新训练状态</button>
|
||||||
|
<div class="hint">
|
||||||
|
<b>说明</b><br>
|
||||||
|
• 只能选择「可测试」状态的发型<br>
|
||||||
|
• 换发型约需 30-60 秒<br>
|
||||||
|
• 会展示换发型的中间步骤便于评估效果
|
||||||
|
</div>
|
||||||
|
<div class="link-row"><a href="/swap">→ 去全部发型测试页</a></div>
|
||||||
|
</div>
|
||||||
|
<div class="workspace">
|
||||||
|
<div class="canvas-wrap" id="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>建议人脸清晰正面照</p>
|
||||||
|
</div>
|
||||||
|
<input type="file" id="file-input" accept="image/*" style="display:none">
|
||||||
|
<img id="preview-img">
|
||||||
|
</div>
|
||||||
|
<div id="status"><span class="spinner"></span><span id="status-text">正在换发型（含可视化），约60秒...</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="result-orig"></div>
|
||||||
|
<div class="result-card"><h4>换发型结果</h4><img id="result-new"><a class="download-btn" id="download-link" download="swaphair_result.jpg">⬇️ 下载结果</a></div>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
<div class="steps-section" id="steps-section">
|
||||||
|
<h3>🔍 换发型过程（逐步可视化）</h3>
|
||||||
|
<div class="sub">每个步骤用到的图片和说明，点击图片可放大查看</div>
|
||||||
|
<div id="steps-container"></div>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
<div id="lightbox" onclick="this.classList.remove('show')"><img id="lightbox-img"></div>
|
||||||
|
<script>
|
||||||
|
(function(){
|
||||||
|
const $=id=>document.getElementById(id);
|
||||||
|
|
||||||
|
// 30 个新发型 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:"长 -花瓣"}
|
||||||
|
];
|
||||||
|
|
||||||
|
const uploadZone=$("upload-zone"),fileInput=$("file-input"),previewImg=$("preview-img");
|
||||||
|
let selectedHairId=null, readySet=new Set(), imgB64Orig=null;
|
||||||
|
|
||||||
|
async function loadStatus(){
|
||||||
|
try{
|
||||||
|
const r=await fetch("/api/hairstyles");
|
||||||
|
const d=await r.json();
|
||||||
|
if(d.state!==0) throw new Error(d.msg||"加载失败");
|
||||||
|
readySet=new Set(d.data.map(x=>x.hair_id));
|
||||||
|
}catch(e){
|
||||||
|
$("status-row").innerHTML='<span style="color:#e74c3c">加载可用发型失败: '+e.message+'</span>';
|
||||||
|
}
|
||||||
|
renderStatus();
|
||||||
|
renderTabs();
|
||||||
|
renderGrid("全部");
|
||||||
|
}
|
||||||
|
function renderStatus(){
|
||||||
|
const ready=HAIRSTYLES.filter(h=>readySet.has(h.name)).length;
|
||||||
|
$("status-row").innerHTML='共 30 款发型｜<b style="color:#4ecca3">可测试 '+ready+'</b>｜训练中/待训练 '+(30-ready);
|
||||||
|
}
|
||||||
|
function renderTabs(){
|
||||||
|
const faces=["全部",...new Set(HAIRSTYLES.map(h=>h.face))];
|
||||||
|
$("face-tabs").innerHTML=faces.map((f,i)=>
|
||||||
|
`<button class="face-tab${i===0?' active':''}" data-face="${f}">${f}</button>`).join('');
|
||||||
|
$("face-tabs").querySelectorAll(".face-tab").forEach(t=>{
|
||||||
|
t.onclick=()=>{
|
||||||
|
$("face-tabs").querySelectorAll(".face-tab").forEach(x=>x.classList.remove("active"));
|
||||||
|
t.classList.add("active");
|
||||||
|
renderGrid(t.dataset.face);
|
||||||
|
};
|
||||||
|
});
|
||||||
|
}
|
||||||
|
function renderGrid(faceFilter){
|
||||||
|
let list=HAIRSTYLES;
|
||||||
|
if(faceFilter!=="全部") list=list.filter(h=>h.face===faceFilter);
|
||||||
|
const g=$("hairstyle-grid");
|
||||||
|
if(list.length===0){ g.innerHTML='<div class="loading-hairstyles">无发型</div>'; return; }
|
||||||
|
g.innerHTML="";
|
||||||
|
list.forEach(h=>{
|
||||||
|
const ready=readySet.has(h.name);
|
||||||
|
const it=document.createElement("div");
|
||||||
|
it.className="hairstyle-item"+(ready?" ready":" pending");
|
||||||
|
if(h.name===selectedHairId) it.classList.add("selected");
|
||||||
|
const badge=ready
|
||||||
|
? '<span class="badge">可测试</span>'
|
||||||
|
: '<span class="badge training">训练中</span>';
|
||||||
|
it.innerHTML=badge+
|
||||||
|
`<img loading="lazy" src="/train_src/${encodeURIComponent(h.name)}" onerror="this.style.background='#0d1b3e';this.style.objectFit='contain';this.src='data:image/svg+xml;utf8,<svg xmlns=%22http://www.w3.org/2000/svg%22 width=%22120%22 height=%22120%22><text x=%2260%22 y=%2265%22 text-anchor=%22middle%22 fill=%22%23666%22 font-size=%2212%22>训练中</text></svg>'">`+
|
||||||
|
`<div class="name">${h.name}</div>`;
|
||||||
|
if(ready){
|
||||||
|
it.onclick=()=>{
|
||||||
|
selectedHairId=h.name;
|
||||||
|
$("hairstyle-grid").querySelectorAll(".hairstyle-item").forEach(el=>el.classList.remove("selected"));
|
||||||
|
it.classList.add("selected");
|
||||||
|
$("selected-info").textContent="已选: "+h.name;
|
||||||
|
updateBtn();
|
||||||
|
};
|
||||||
|
}
|
||||||
|
g.appendChild(it);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
uploadZone.addEventListener("click",()=>fileInput.click());
|
||||||
|
uploadZone.addEventListener("dragover",e=>{e.preventDefault();uploadZone.classList.add("dragover");});
|
||||||
|
uploadZone.addEventListener("dragleave",()=>uploadZone.classList.remove("dragover"));
|
||||||
|
uploadZone.addEventListener("drop",e=>{e.preventDefault();uploadZone.classList.remove("dragover");if(e.dataTransfer.files[0])loadImage(e.dataTransfer.files[0]);});
|
||||||
|
fileInput.addEventListener("change",e=>{if(e.target.files[0])loadImage(e.target.files[0]);});
|
||||||
|
function loadImage(file){
|
||||||
|
if(!file.type.startsWith("image/")){alert("请上传图片文件");return;}
|
||||||
|
const rd=new FileReader();
|
||||||
|
rd.onload=e=>{
|
||||||
|
imgB64Orig=e.target.result;
|
||||||
|
previewImg.src=imgB64Orig;previewImg.style.display="block";
|
||||||
|
uploadZone.style.display="none";updateBtn();
|
||||||
|
$("result-section").classList.remove("show");
|
||||||
|
$("steps-section").classList.remove("show");
|
||||||
|
$("error-msg").style.display="none";
|
||||||
|
};
|
||||||
|
rd.readAsDataURL(file);
|
||||||
|
}
|
||||||
|
function updateBtn(){ $("btn-generate").disabled=!(selectedHairId&&imgB64Orig&&readySet.has(selectedHairId)); }
|
||||||
|
|
||||||
|
function showLightbox(src){ $("lightbox-img").src=src; $("lightbox").classList.add("show"); }
|
||||||
|
|
||||||
|
$("btn-refresh").onclick=()=>{ loadStatus(); };
|
||||||
|
|
||||||
|
$("btn-generate").addEventListener("click",async()=>{
|
||||||
|
if(!selectedHairId||!imgB64Orig) return;
|
||||||
|
const isHr=$("is-hr").checked;
|
||||||
|
$("btn-generate").disabled=true;$("status").style.display="block";
|
||||||
|
$("status-text").textContent="正在换发型（含过程可视化），约30-60秒...";
|
||||||
|
$("error-msg").style.display="none";
|
||||||
|
$("result-section").classList.remove("show");$("steps-section").classList.remove("show");
|
||||||
|
try{
|
||||||
|
const resp=await fetch("/api/swap_viz",{
|
||||||
|
method:"POST",headers:{"Content-Type":"application/json"},
|
||||||
|
body:JSON.stringify({img:imgB64Orig,hair_id:selectedHairId,is_hr:String(isHr),strict_mask:$("strict-mask").checked})
|
||||||
|
});
|
||||||
|
const data=await resp.json();
|
||||||
|
$("status").style.display="none";
|
||||||
|
if(data.state===0){
|
||||||
|
const url="data:image/jpeg;base64,"+data.result;
|
||||||
|
$("result-orig").src=imgB64Orig;$("result-new").src=url;$("download-link").href=url;
|
||||||
|
$("result-section").classList.add("show");
|
||||||
|
if(data.steps&&data.steps.length){
|
||||||
|
const c=$("steps-container");c.innerHTML="";
|
||||||
|
data.steps.forEach(step=>{
|
||||||
|
const div=document.createElement("div");div.className="step-item";
|
||||||
|
let imgsHtml='<div class="step-images">';
|
||||||
|
step.images.forEach(im=>{imgsHtml+='<div class="step-img-card"><div class="lbl">'+im.label+'</div><img src="data:image/jpeg;base64,'+im.b64+'" data-full="data:image/jpeg;base64,'+im.b64+'"></div>';});
|
||||||
|
imgsHtml+='</div>';
|
||||||
|
div.innerHTML='<div class="step-title">'+step.title+'</div><div class="step-desc">'+step.desc+'</div>'+imgsHtml;
|
||||||
|
c.appendChild(div);
|
||||||
|
});
|
||||||
|
c.querySelectorAll("img").forEach(img=>img.onclick=()=>showLightbox(img.dataset.full));
|
||||||
|
$("steps-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";
|
||||||
|
}
|
||||||
|
updateBtn();
|
||||||
|
});
|
||||||
|
|
||||||
|
loadStatus();
|
||||||
|
})();
|
||||||
|
</script>
|
||||||
|
</body>
|
||||||
|
</html>
|
||||||
Binary file not shown.
|
After Width: | Height: | Size: 921 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 410 KiB |
@@ -0,0 +1,182 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
"""自动准备 LoRA 训练数据
|
||||||
|
|
||||||
|
把 train_images/ 里的原始人像图处理成 kohya LoRA 训练格式：
|
||||||
|
原始图 → 人脸检测+1k关键点 → Generator_Matte头发抠图 → 白底合成 + 居中裁剪 → 打标签
|
||||||
|
|
||||||
|
用法:
|
||||||
|
cd /home/xsl/change_hair/project/hair_service_sd
|
||||||
|
/home/xsl/miniconda3/envs/my_hair/bin/python /home/xsl/change_hair/prepare_train_data.py \
|
||||||
|
--input /home/xsl/change_hair/train_images \
|
||||||
|
--hair-id new_hairstyle_001 \
|
||||||
|
--gender boy
|
||||||
|
"""
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
import ssl
|
||||||
|
import uuid
|
||||||
|
import argparse
|
||||||
|
|
||||||
|
# 禁用 SSL 验证（解决 resnet18 权重缓存的 SSL 证书问题）
|
||||||
|
ssl._create_default_https_context = ssl._create_unverified_context
|
||||||
|
|
||||||
|
# 必须在 hair_service_sd 目录运行（依赖其模块）
|
||||||
|
HAIR_SERVICE_DIR = "/home/xsl/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")
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
|
||||||
|
# Monkey-patch torch.load: 部分第三方代码用旧式 map_location lambda 触发 legacy_load,
|
||||||
|
# 但权重文件是新版 zip 格式，导致 UnpicklingError。统一用 map_location='cpu'。
|
||||||
|
_orig_torch_load = torch.load
|
||||||
|
def _patched_torch_load(*args, **kwargs):
|
||||||
|
if 'map_location' in kwargs and callable(kwargs['map_location']) and not isinstance(kwargs['map_location'], str):
|
||||||
|
kwargs['map_location'] = 'cpu'
|
||||||
|
return _orig_torch_load(*args, **kwargs)
|
||||||
|
torch.load = _patched_torch_load
|
||||||
|
|
||||||
|
from models.detector import RetinaFaceDetector
|
||||||
|
from models.MomocvFaceAlignment1K import MomocvFaceAlignment1K
|
||||||
|
from utils.landmark_processor import get_max_rect
|
||||||
|
# 用 core/process_modules 的 Generator_Matte（从 weights/ 加载真实权重，
|
||||||
|
# 而非 hair_matting/ 目录的 LFS 指针文件）
|
||||||
|
from core.process_modules import Generator_Matte
|
||||||
|
|
||||||
|
# kohya 标准训练标签（白底版，和已有发型完全一致）
|
||||||
|
CAPTION = "titor hairstyle, easyphoto, faceless, no human, white background, simple background, "
|
||||||
|
|
||||||
|
# 训练图分辨率列表（和 prepare_single_hairstyle_v2 一致）
|
||||||
|
RESOLUTIONS = [512, 768, 1024, 1280, 1536]
|
||||||
|
|
||||||
|
|
||||||
|
def process_one(img_path, detector, aligner, matte_gen, out_dir, idx):
|
||||||
|
"""处理单张原始图：检测→抠图→白底合成→多分辨率保存"""
|
||||||
|
img = cv2.imread(img_path)
|
||||||
|
if img is None:
|
||||||
|
print(f" [{idx}] ✗ 读取失败: {img_path}")
|
||||||
|
return 0
|
||||||
|
h, w = img.shape[:2]
|
||||||
|
|
||||||
|
# 1. 人脸检测
|
||||||
|
dets, landms = detector.forward(img)
|
||||||
|
if len(dets) == 0:
|
||||||
|
print(f" [{idx}] ✗ 未检测到人脸: {os.path.basename(img_path)}")
|
||||||
|
return 0
|
||||||
|
max_idx = get_max_rect(dets)
|
||||||
|
det = dets[max_idx]
|
||||||
|
|
||||||
|
# 2. 1k 关键点
|
||||||
|
landmarks_1k = aligner.stable_forward(img.copy(), [det])
|
||||||
|
if landmarks_1k is None or len(landmarks_1k) == 0:
|
||||||
|
print(f" [{idx}] ✗ 关键点检测失败: {os.path.basename(img_path)}")
|
||||||
|
return 0
|
||||||
|
pt1k = landmarks_1k[0]
|
||||||
|
|
||||||
|
# 3. 头发抠图（matte_inference 返回 pred_fg, pred, image_resize）
|
||||||
|
with torch.no_grad():
|
||||||
|
pred_fg, pred, image_resize = matte_gen.matte_inference(img, pt1k)
|
||||||
|
# pred 是 alpha mask，已 resize 回原图尺寸
|
||||||
|
alpha = pred
|
||||||
|
h, w = img.shape[:2]
|
||||||
|
if alpha.shape != (h, w):
|
||||||
|
alpha = cv2.resize(alpha, (w, h))
|
||||||
|
|
||||||
|
# 4. 白底合成（公式同 hairstyle_model.py:3666-3669）
|
||||||
|
alpha_f = alpha.astype(np.float32) / 255.0
|
||||||
|
white_bg = np.ones_like(img, dtype=np.float32) * 255.0
|
||||||
|
img_f = img.astype(np.float32)
|
||||||
|
composite = (img_f * alpha_f[:, :, None] + white_bg * (1 - alpha_f[:, :, None]))
|
||||||
|
composite = np.clip(composite, 0, 255).astype(np.uint8)
|
||||||
|
|
||||||
|
# 5. 居中裁剪到正方形（以人脸为中心）
|
||||||
|
pts = pt1k[:312] # 前312点是脸部外轮廓
|
||||||
|
cx, cy = int(pts[:, 0].mean()), int(pts[:, 1].mean())
|
||||||
|
# 正方形边长 = 图像短边（保证不裁掉头发）
|
||||||
|
side = min(h, w)
|
||||||
|
x1 = max(0, cx - side // 2)
|
||||||
|
y1 = max(0, cy - side // 2)
|
||||||
|
x2 = min(w, x1 + side)
|
||||||
|
y2 = min(h, y1 + side)
|
||||||
|
# 调整确保正方形
|
||||||
|
actual_side = min(x2 - x1, y2 - y1)
|
||||||
|
x2, y2 = x1 + actual_side, y1 + actual_side
|
||||||
|
composite_square = composite[y1:y2, x1:x2]
|
||||||
|
if composite_square.shape[0] < 100:
|
||||||
|
# fallback: 直接用整图短边
|
||||||
|
side = min(h, w)
|
||||||
|
composite_square = composite[:side, :side]
|
||||||
|
|
||||||
|
# 6. 多分辨率保存 + 同名标签
|
||||||
|
count = 0
|
||||||
|
for res in RESOLUTIONS:
|
||||||
|
if composite_square.shape[0] >= res // 2: # 只生成 >= res/2 的（避免过度放大）
|
||||||
|
if composite_square.shape[0] >= res:
|
||||||
|
resized = cv2.resize(composite_square, (res, res), interpolation=cv2.INTER_AREA)
|
||||||
|
else:
|
||||||
|
resized = cv2.resize(composite_square, (res, res), interpolation=cv2.INTER_LANCZOS4)
|
||||||
|
|
||||||
|
file_uuid = str(uuid.uuid4())
|
||||||
|
png_path = os.path.join(out_dir, f"{file_uuid}.png")
|
||||||
|
txt_path = os.path.join(out_dir, f"{file_uuid}.txt")
|
||||||
|
cv2.imwrite(png_path, resized)
|
||||||
|
with open(txt_path, "w") as f:
|
||||||
|
f.write(CAPTION)
|
||||||
|
count += 1
|
||||||
|
|
||||||
|
print(f" [{idx}] ✓ {os.path.basename(img_path)} → {count}张训练图")
|
||||||
|
return count
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser(description="准备 LoRA 训练数据")
|
||||||
|
parser.add_argument("--input", required=True, help="原始图目录")
|
||||||
|
parser.add_argument("--hair-id", required=True, help="新发型ID")
|
||||||
|
parser.add_argument("--gender", default="boy", choices=["boy", "girl"], help="性别")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
from common.logger import config
|
||||||
|
train_dir = config.get('default', 'train_dir')
|
||||||
|
out_base = os.path.join(train_dir, args.hair_id, "images", "1_hairstyle")
|
||||||
|
os.makedirs(out_base, exist_ok=True)
|
||||||
|
|
||||||
|
print(f"=== 准备训练数据 ===")
|
||||||
|
print(f"输入: {args.input}")
|
||||||
|
print(f"发型ID: {args.hair_id}, 性别: {args.gender}")
|
||||||
|
print(f"输出: {out_base}")
|
||||||
|
|
||||||
|
# 加载模型
|
||||||
|
print("加载模型（RetinaFace + 1k关键点 + Generator_Matte）...")
|
||||||
|
gpu_id = 0
|
||||||
|
detector = RetinaFaceDetector(gpu_id=gpu_id)
|
||||||
|
aligner = MomocvFaceAlignment1K(gpu_id=gpu_id)
|
||||||
|
matte_gen = Generator_Matte(gpu=True, device_id=gpu_id)
|
||||||
|
print("模型加载完成")
|
||||||
|
|
||||||
|
# 处理每张图
|
||||||
|
import glob
|
||||||
|
imgs = sorted(glob.glob(os.path.join(args.input, "*.jpg")) +
|
||||||
|
glob.glob(os.path.join(args.input, "*.png")) +
|
||||||
|
glob.glob(os.path.join(args.input, "*.jpeg")))
|
||||||
|
print(f"共 {len(imgs)} 张原始图")
|
||||||
|
|
||||||
|
total = 0
|
||||||
|
for i, img_path in enumerate(imgs):
|
||||||
|
total += process_one(img_path, detector, aligner, matte_gen, out_base, i + 1)
|
||||||
|
|
||||||
|
print(f"\n=== 完成: {len(imgs)}张原始图 → {total}张训练图 ===")
|
||||||
|
print(f"训练数据目录: {out_base}")
|
||||||
|
print(f"\n下一步训练命令:")
|
||||||
|
print(f"curl -X POST http://127.0.0.1:32678/api/hair/train \\")
|
||||||
|
print(f' -H "Content-Type: application/json" \\')
|
||||||
|
print(f' -d \'{{"task_id":"train_{args.hair_id}","hair_id":"{args.hair_id}",'
|
||||||
|
f'"hair_material_dir":"{os.path.join(train_dir, args.hair_id)}",'
|
||||||
|
f'"is_tj":"1","device_id":"0","webui_addr":"http://0.0.0.0:32678/"}}\'')
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -0,0 +1,128 @@
|
|||||||
|
# 大文件与模型权重说明
|
||||||
|
|
||||||
|
> 本仓库只托管代码和配置。模型权重、训练数据等大文件（共约 **200G**）已排除，需另行准备。
|
||||||
|
> 后续会上传网盘，地址见文末。
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 被排除的大目录清单
|
||||||
|
|
||||||
|
| 目录 | 大小 | 内容 | 如何恢复 |
|
||||||
|
|------|------|------|---------|
|
||||||
|
| `data/` | 123G | 341 个发型的训练素材 + LoRA 权重 | 网盘下载后放到 `data/train_material/` |
|
||||||
|
| `project/onediff/` | 52G | Stable Diffusion webui + SD模型 + 228个LoRA + ControlNet | 见下方「onediff 恢复」 |
|
||||||
|
| `project/kohya_ss_home/` | 13G | kohya_ss LoRA 训练框架（第三方） | 见下方「kohya 恢复」 |
|
||||||
|
| `project/hair_service_sd/weights/` | 9.9G | 人脸检测/关键点/抠图/分割等核心模型权重 | 网盘下载 |
|
||||||
|
| `project/data/` | 5.4G | 运行时材质数据（ref_hairstyle/hair_template_material 等） | 训练完成后自动生成 |
|
||||||
|
| `hair_type_images/` | 53M | 本次 30 张发型训练原图 | 网盘下载 |
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 1. `project/hair_service_sd/weights/` 核心模型（9.9G）
|
||||||
|
|
||||||
|
主服务推理必需。从网盘下载后解压到本目录。包含：
|
||||||
|
|
||||||
|
| 文件 | 大小 | 用途 |
|
||||||
|
|------|------|------|
|
||||||
|
| `Resnet50_Final.pth` | 109M | RetinaFace 人脸检测 |
|
||||||
|
| `face_alignment_1k.pth` | 49M | 1000点人脸关键点 |
|
||||||
|
| `deeplabv3_hair512_360_0520_wl.pth` | 233M | 头发分割 DeepLabV3+ (3分类) |
|
||||||
|
| `gca-dist-fg-0430-latest_model.pth` | 303M | GCA 头发抠图 |
|
||||||
|
| `20210927_01.pth` / `faceseg_20210927_01.pth` | 各229M | 人脸分割 U2NET |
|
||||||
|
| `master_hair_v8_onlyhair_nowarp_all_0709.pt` | 1.1G | SPADE 发型迁移（粗推理） |
|
||||||
|
| `hair_fusion_0427.pt` | 779M | 发型融合 |
|
||||||
|
| `ycj_hairstyle_gan_*.pt` / `ycj_haircolor_gan_*.pt` | 各~500M | 发型/发色 GAN |
|
||||||
|
| `gfpgan_hairline_512_inversion_ratio.pt` | 308M | 人脸增强 GFPGAN |
|
||||||
|
| `face_enhance_0630.pt` | 284M | 人脸增强 |
|
||||||
|
| `MMCVFaceRecognitionServer.pth` | 174M | 人脸识别 |
|
||||||
|
| `kpnts_detect_lyq_20200720.pth` | 255M | 关键点检测 |
|
||||||
|
| `human_seg_deeplabv3_288_384_sigmoid_msc.pth` | 134M | 人体分割 |
|
||||||
|
| `ori_hair_checkpoint_7660_0611.pt` | ~ | BaldSeg 光头分割 |
|
||||||
|
| `mediapipe/face_landmarker.task` | 3.6M | MediaPipe 478点关键点（发际线检测用） |
|
||||||
|
| 其余 `*.pth/*.pt` | — | 各功能模块权重 |
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 2. `project/onediff/` SD webui（52G，第三方）
|
||||||
|
|
||||||
|
换发型 SD 推理必需。这是 [AUTOMATIC1111/stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui) 的部署。
|
||||||
|
|
||||||
|
**恢复方式（二选一）**：
|
||||||
|
- **A. 网盘下载**：直接下载整个 onediff 目录解压（最快）
|
||||||
|
- **B. 重新部署**：
|
||||||
|
```bash
|
||||||
|
cd project
|
||||||
|
git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git onediff/stable-diffusion-webui
|
||||||
|
# 然后下载以下模型放到对应目录：
|
||||||
|
# models/Stable-diffusion/v1-5-pruned-emaonly.safetensors (4G) — SD底模
|
||||||
|
# models/Lora/*.safetensors — 各发型LoRA（由训练生成）
|
||||||
|
# models/ControlNet/control_v11p_sd15_canny.pth — ControlNet（可选）
|
||||||
|
```
|
||||||
|
|
||||||
|
**关键配置**：
|
||||||
|
- SD 底模：`v1-5-pruned-emaonly.safetensors`（原项目用 majicmix，本机替换为 v1-5）
|
||||||
|
- LoRA：训练完成的发型会自动生成到 `models/Lora/`
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 3. `project/kohya_ss_home/` 训练框架（13G，第三方）
|
||||||
|
|
||||||
|
LoRA 训练用。[kohya-ss/sd-scripts](https://github.com/kohya-ss/sd-scripts) 的部署。
|
||||||
|
|
||||||
|
**恢复方式**：
|
||||||
|
```bash
|
||||||
|
cd project
|
||||||
|
git clone https://github.com/bmaltais/kohya_ss.git kohya_ss_home/kohya_ss
|
||||||
|
# 安装依赖（用 kohya conda 环境）
|
||||||
|
/home/xsl/miniconda3/envs/kohya/bin/pip install -r kohya_ss_home/kohya_ss/requirements.txt
|
||||||
|
```
|
||||||
|
|
||||||
|
训练入口：`kohya_ss_home/kohya_ss/train_network.py`
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 4. `data/train_material/` 训练数据（123G）
|
||||||
|
|
||||||
|
341 个发型的 LoRA 训练素材和权重。每个发型一个目录：
|
||||||
|
```
|
||||||
|
data/train_material/<hair_id>/
|
||||||
|
├── images/1_hairstyle/*.png # 训练图（数据增强后约20张）
|
||||||
|
└── model/
|
||||||
|
└── hairstyle_hd_lora.safetensors # 训练产出的LoRA权重(~150M)
|
||||||
|
```
|
||||||
|
|
||||||
|
**恢复方式**：
|
||||||
|
- **网盘下载**：下载后解压到 `data/train_material/`
|
||||||
|
- **重新训练**：用 `batch_train_hairstyles.py` 重新训练需要的发型（每个约30分钟）
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 5. `project/data/` 运行时材质（5.4G）
|
||||||
|
|
||||||
|
换发型功能运行时需要的预生成材质。**训练完成发型后会自动生成**，无需手动准备：
|
||||||
|
- `ref_hairstyle/<hair_id>/` — 每个发型的参考材质（粗推理生成）
|
||||||
|
- `hair_template_material/<hair_id>/` — 模板材质（抠图+关键点）
|
||||||
|
- `upload_train_imgs/<hair_id>/` — 上传的训练原图
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 训练原图
|
||||||
|
|
||||||
|
`hair_type_images/`（30张发型原图，按 `<脸型>-<发型>.jpg` 命名）是本次新增发型的训练源。可从网盘下载，或自行准备同格式图片后用 `batch_train_hairstyles.py` 训练。
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 网盘地址
|
||||||
|
|
||||||
|
（待补充）—— 大文件上传网盘后，把分享链接填到这里。
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 快速检查清单（部署后）
|
||||||
|
|
||||||
|
部署完成后，确认以下文件存在：
|
||||||
|
- [ ] `project/hair_service_sd/weights/Resnet50_Final.pth`（人脸检测）
|
||||||
|
- [ ] `project/hair_service_sd/weights/face_alignment_1k.pth`（关键点）
|
||||||
|
- [ ] `project/hair_service_sd/weights/master_hair_v8_onlyhair_nowarp_all_0709.pt`（发型迁移）
|
||||||
|
- [ ] `project/onediff/stable-diffusion-webui/models/Stable-diffusion/v1-5-pruned-emaonly.safetensors`（SD底模）
|
||||||
|
- [ ] `data/train_material/<hair_id>/model/hairstyle_hd_lora.safetensors`（至少一个发型的LoRA）
|
||||||
@@ -0,0 +1,3 @@
|
|||||||
|
*.pth filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.pt filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.ttc filter=lfs diff=lfs merge=lfs -text
|
||||||
@@ -0,0 +1,194 @@
|
|||||||
|
# Byte-compiled / optimized / DLL files
|
||||||
|
__pycache__/
|
||||||
|
*.py[cod]
|
||||||
|
*$py.class
|
||||||
|
|
||||||
|
# C extensions
|
||||||
|
*.so
|
||||||
|
|
||||||
|
# Distribution / packaging
|
||||||
|
.Python
|
||||||
|
build/
|
||||||
|
develop-eggs/
|
||||||
|
dist/
|
||||||
|
downloads/
|
||||||
|
eggs/
|
||||||
|
.eggs/
|
||||||
|
lib/
|
||||||
|
lib64/
|
||||||
|
parts/
|
||||||
|
sdist/
|
||||||
|
var/
|
||||||
|
wheels/
|
||||||
|
share/python-wheels/
|
||||||
|
*.egg-info/
|
||||||
|
.installed.cfg
|
||||||
|
*.egg
|
||||||
|
MANIFEST
|
||||||
|
|
||||||
|
# PyInstaller
|
||||||
|
# Usually these files are written by a python script from a template
|
||||||
|
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||||
|
*.manifest
|
||||||
|
*.spec
|
||||||
|
|
||||||
|
# Installer logs
|
||||||
|
pip-log.txt
|
||||||
|
pip-delete-this-directory.txt
|
||||||
|
|
||||||
|
# Unit test / coverage reports
|
||||||
|
htmlcov/
|
||||||
|
.tox/
|
||||||
|
.nox/
|
||||||
|
.coverage
|
||||||
|
.coverage.*
|
||||||
|
.cache
|
||||||
|
nosetests.xml
|
||||||
|
coverage.xml
|
||||||
|
*.cover
|
||||||
|
*.py,cover
|
||||||
|
.hypothesis/
|
||||||
|
.pytest_cache/
|
||||||
|
cover/
|
||||||
|
|
||||||
|
# Translations
|
||||||
|
*.mo
|
||||||
|
*.pot
|
||||||
|
|
||||||
|
# Django stuff:
|
||||||
|
*.log
|
||||||
|
local_settings.py
|
||||||
|
db.sqlite3
|
||||||
|
db.sqlite3-journal
|
||||||
|
|
||||||
|
# Flask stuff:
|
||||||
|
instance/
|
||||||
|
.webassets-cache
|
||||||
|
|
||||||
|
# Scrapy stuff:
|
||||||
|
.scrapy
|
||||||
|
|
||||||
|
# Sphinx documentation
|
||||||
|
docs/_build/
|
||||||
|
|
||||||
|
# PyBuilder
|
||||||
|
.pybuilder/
|
||||||
|
target/
|
||||||
|
|
||||||
|
# Jupyter Notebook
|
||||||
|
.ipynb_checkpoints
|
||||||
|
|
||||||
|
# IPython
|
||||||
|
profile_default/
|
||||||
|
ipython_config.py
|
||||||
|
|
||||||
|
# pyenv
|
||||||
|
# For a library or package, you might want to ignore these files since the code is
|
||||||
|
# intended to run in multiple environments; otherwise, check them in:
|
||||||
|
# .python-version
|
||||||
|
|
||||||
|
# pipenv
|
||||||
|
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||||
|
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||||
|
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||||
|
# install all needed dependencies.
|
||||||
|
#Pipfile.lock
|
||||||
|
|
||||||
|
# UV
|
||||||
|
# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
|
||||||
|
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||||
|
# commonly ignored for libraries.
|
||||||
|
#uv.lock
|
||||||
|
|
||||||
|
# poetry
|
||||||
|
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
||||||
|
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||||
|
# commonly ignored for libraries.
|
||||||
|
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
||||||
|
#poetry.lock
|
||||||
|
|
||||||
|
# pdm
|
||||||
|
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
||||||
|
#pdm.lock
|
||||||
|
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
|
||||||
|
# in version control.
|
||||||
|
# https://pdm.fming.dev/latest/usage/project/#working-with-version-control
|
||||||
|
.pdm.toml
|
||||||
|
.pdm-python
|
||||||
|
.pdm-build/
|
||||||
|
|
||||||
|
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
||||||
|
__pypackages__/
|
||||||
|
|
||||||
|
# Celery stuff
|
||||||
|
celerybeat-schedule
|
||||||
|
celerybeat.pid
|
||||||
|
|
||||||
|
# SageMath parsed files
|
||||||
|
*.sage.py
|
||||||
|
|
||||||
|
# Environments
|
||||||
|
.env
|
||||||
|
.venv
|
||||||
|
env/
|
||||||
|
venv/
|
||||||
|
ENV/
|
||||||
|
env.bak/
|
||||||
|
venv.bak/
|
||||||
|
|
||||||
|
# Spyder project settings
|
||||||
|
.spyderproject
|
||||||
|
.spyproject
|
||||||
|
|
||||||
|
# Rope project settings
|
||||||
|
.ropeproject
|
||||||
|
|
||||||
|
# mkdocs documentation
|
||||||
|
/site
|
||||||
|
|
||||||
|
# mypy
|
||||||
|
.mypy_cache/
|
||||||
|
.dmypy.json
|
||||||
|
dmypy.json
|
||||||
|
|
||||||
|
# Pyre type checker
|
||||||
|
.pyre/
|
||||||
|
|
||||||
|
# pytype static type analyzer
|
||||||
|
.pytype/
|
||||||
|
|
||||||
|
# Cython debug symbols
|
||||||
|
cython_debug/
|
||||||
|
|
||||||
|
# PyCharm
|
||||||
|
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
||||||
|
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
||||||
|
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
||||||
|
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||||
|
#.idea/
|
||||||
|
|
||||||
|
# Abstra
|
||||||
|
# Abstra is an AI-powered process automation framework.
|
||||||
|
# Ignore directories containing user credentials, local state, and settings.
|
||||||
|
# Learn more at https://abstra.io/docs
|
||||||
|
.abstra/
|
||||||
|
|
||||||
|
# Visual Studio Code
|
||||||
|
# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
|
||||||
|
# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
|
||||||
|
# and can be added to the global gitignore or merged into this file. However, if you prefer,
|
||||||
|
# you could uncomment the following to ignore the enitre vscode folder
|
||||||
|
# .vscode/
|
||||||
|
|
||||||
|
# Ruff stuff:
|
||||||
|
.ruff_cache/
|
||||||
|
|
||||||
|
# PyPI configuration file
|
||||||
|
.pypirc
|
||||||
|
|
||||||
|
# Cursor
|
||||||
|
# Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
|
||||||
|
# exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
|
||||||
|
# refer to https://docs.cursor.com/context/ignore-files
|
||||||
|
.cursorignore
|
||||||
|
.cursorindexingignore
|
||||||
@@ -0,0 +1,28 @@
|
|||||||
|
换发算法服务代码
|
||||||
|
|
||||||
|
相关配置：config/configure.ini \
|
||||||
|
模型:请存放于weights目录下 \
|
||||||
|
服务启动：执行 gunicorn server:app -c gunicorn_config.py \
|
||||||
|
测试相关代码请写在 unit_test.py里
|
||||||
|
|
||||||
|
建议环境：
|
||||||
|
python3.7，pytorch1.8,cu11.1,torchvision10.0 \
|
||||||
|
gunicorn需要在/usr/bin/gunicorn需要在里修改python地址
|
||||||
|
|
||||||
|
run run_copy_cost.py
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
需要启动的服务：
|
||||||
|
环境变量设置
|
||||||
|
export CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1
|
||||||
|
1. 换发算法服务
|
||||||
|
cd /root/project/hair_service_sd
|
||||||
|
/usr/local/miniconda3/envs/condiff-train-hair/bin/python run_copy_cost_colorb64.py
|
||||||
|
2. webUI服务推理 + onediff
|
||||||
|
cd /root/project/onediff/stable-diffusion-webui
|
||||||
|
/usr/local/miniconda3/envs/onediff/bin/python python webui.py --api --listen --xformers --port 57860
|
||||||
|
./webui.sh --api --listen --disable-safe-unpickle --port 9038
|
||||||
|
3. photo_service
|
||||||
|
cd /root/project/photo_service
|
||||||
|
/usr/local/miniconda3/envs/py310/bin/python lora_train_service_1.py
|
||||||
@@ -0,0 +1,75 @@
|
|||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
from torch import nn
|
||||||
|
|
||||||
|
def get_norm(norm, out_channels=None):
|
||||||
|
"""
|
||||||
|
Args:
|
||||||
|
norm (str or callable):
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
nn.Module or None: the normalization layer
|
||||||
|
"""
|
||||||
|
if isinstance(norm, str):
|
||||||
|
if len(norm) == 0:
|
||||||
|
return None
|
||||||
|
norm = {
|
||||||
|
"BN": nn.BatchNorm2d,
|
||||||
|
"IN": nn.InstanceNorm2d,
|
||||||
|
"GN": lambda channels: nn.GroupNorm(32, channels),
|
||||||
|
"nnSyncBN": nn.SyncBatchNorm, # keep for debugging
|
||||||
|
}[norm]
|
||||||
|
if out_channels is not None:
|
||||||
|
return norm(out_channels)
|
||||||
|
else:
|
||||||
|
return norm
|
||||||
|
|
||||||
|
class Conv2d(torch.nn.Conv2d):
|
||||||
|
"""
|
||||||
|
A wrapper around :class:`torch.nn.Conv2d` to support zero-size tensor and more features.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, *args, **kwargs):
|
||||||
|
"""
|
||||||
|
Extra keyword arguments supported in addition to those in `torch.nn.Conv2d`:
|
||||||
|
|
||||||
|
Args:
|
||||||
|
norm (nn.Module, optional): a normalization layer
|
||||||
|
activation (callable(Tensor) -> Tensor): a callable activation function
|
||||||
|
|
||||||
|
It assumes that norm layer is used before activation.
|
||||||
|
"""
|
||||||
|
norm = kwargs.pop("norm", None)
|
||||||
|
activation = kwargs.pop("activation", None)
|
||||||
|
super().__init__(*args, **kwargs)
|
||||||
|
|
||||||
|
self.norm = norm
|
||||||
|
self.activation = activation
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x = super().forward(x)
|
||||||
|
if self.norm is not None:
|
||||||
|
x = self.norm(x)
|
||||||
|
if self.activation is not None:
|
||||||
|
x = self.activation(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
class Backbone(nn.Module):
|
||||||
|
"""
|
||||||
|
Abstract base class for network backbones.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
"""
|
||||||
|
The `__init__` method of any subclass can specify its own set of arguments.
|
||||||
|
"""
|
||||||
|
super().__init__()
|
||||||
|
|
||||||
|
def forward(self):
|
||||||
|
"""
|
||||||
|
Subclasses must override this method, but adhere to the same return type.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict[str: Tensor]: mapping from feature name (e.g., "res2") to tensor
|
||||||
|
"""
|
||||||
|
pass
|
||||||
@@ -0,0 +1,268 @@
|
|||||||
|
import math
|
||||||
|
import utils.weight_init as weight_init
|
||||||
|
import torch
|
||||||
|
import torch.nn.functional as F
|
||||||
|
from torch import nn
|
||||||
|
|
||||||
|
from bodyseg.backbone.backbone import Backbone, get_norm, Conv2d
|
||||||
|
from bodyseg.backbone.resnet import build_resnet_backbone
|
||||||
|
|
||||||
|
class FPN(Backbone):
|
||||||
|
"""
|
||||||
|
This module implements Feature Pyramid Network.
|
||||||
|
It creates pyramid features built on top of some input feature maps.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self, bottom_up, in_features, out_channels, norm="", top_block=None, fuse_type="sum"
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Args:
|
||||||
|
bottom_up (Backbone): module representing the bottom up subnetwork.
|
||||||
|
Must be a subclass of :class:`Backbone`. The multi-scale feature
|
||||||
|
maps generated by the bottom up network, and listed in `in_features`,
|
||||||
|
are used to generate FPN levels.
|
||||||
|
in_features (list[str]): names of the input feature maps coming
|
||||||
|
from the backbone to which FPN is attached. For example, if the
|
||||||
|
backbone produces ["res2", "res3", "res4"], any *contiguous* sublist
|
||||||
|
of these may be used; order must be from high to low resolution.
|
||||||
|
out_channels (int): number of channels in the output feature maps.
|
||||||
|
norm (str): the normalization to use.
|
||||||
|
top_block (nn.Module or None): if provided, an extra operation will
|
||||||
|
be performed on the output of the last (smallest resolution)
|
||||||
|
FPN output, and the result will extend the result list. The top_block
|
||||||
|
further downsamples the feature map. It must have an attribute
|
||||||
|
"num_levels", meaning the number of extra FPN levels added by
|
||||||
|
this block, and "in_feature", which is a string representing
|
||||||
|
its input feature (e.g., p5).
|
||||||
|
fuse_type (str): types for fusing the top down features and the lateral
|
||||||
|
ones. It can be "sum" (default), which sums up element-wise; or "avg",
|
||||||
|
which takes the element-wise mean of the two.
|
||||||
|
"""
|
||||||
|
super(FPN, self).__init__()
|
||||||
|
assert isinstance(bottom_up, Backbone)
|
||||||
|
|
||||||
|
# Feature map strides and channels from the bottom up network (e.g. ResNet)
|
||||||
|
in_strides = [bottom_up._out_feature_strides[f] for f in in_features]
|
||||||
|
in_channels = [bottom_up._out_feature_channels[f] for f in in_features]
|
||||||
|
|
||||||
|
_assert_strides_are_log2_contiguous(in_strides)
|
||||||
|
lateral_convs = []
|
||||||
|
output_convs = []
|
||||||
|
|
||||||
|
use_bias = norm == ""
|
||||||
|
for idx, in_channels in enumerate(in_channels):
|
||||||
|
lateral_norm = get_norm(norm, out_channels)
|
||||||
|
output_norm = get_norm(norm, out_channels)
|
||||||
|
|
||||||
|
lateral_conv = Conv2d(
|
||||||
|
in_channels, out_channels, kernel_size=1, bias=use_bias, norm=lateral_norm
|
||||||
|
)
|
||||||
|
output_conv = Conv2d(
|
||||||
|
out_channels,
|
||||||
|
out_channels,
|
||||||
|
kernel_size=3,
|
||||||
|
stride=1,
|
||||||
|
padding=1,
|
||||||
|
bias=use_bias,
|
||||||
|
norm=output_norm,
|
||||||
|
)
|
||||||
|
weight_init.c2_xavier_fill(lateral_conv)
|
||||||
|
weight_init.c2_xavier_fill(output_conv)
|
||||||
|
stage = int(math.log2(in_strides[idx]))
|
||||||
|
|
||||||
|
lateral_convs.append(lateral_conv)
|
||||||
|
output_convs.append(output_conv)
|
||||||
|
# Place convs into top-down order (from low to high resolution)
|
||||||
|
# to make the top-down computation in forward clearer.
|
||||||
|
self.lateral_convs = nn.ModuleList(lateral_convs[::-1])
|
||||||
|
self.output_convs = nn.ModuleList(output_convs[::-1])
|
||||||
|
self.top_block = top_block
|
||||||
|
self.in_features = in_features
|
||||||
|
self.bottom_up = bottom_up
|
||||||
|
# Return feature names are "p<stage>", like ["p2", "p3", ..., "p6"]
|
||||||
|
self._out_feature_strides = {"p{}".format(int(math.log2(s))): s for s in in_strides}
|
||||||
|
# top block output feature maps.
|
||||||
|
if self.top_block is not None:
|
||||||
|
for s in range(stage, stage + self.top_block.num_levels):
|
||||||
|
self._out_feature_strides["p{}".format(s + 1)] = 2 ** (s + 1)
|
||||||
|
|
||||||
|
self._out_features = list(self._out_feature_strides.keys())
|
||||||
|
self._out_feature_channels = {k: out_channels for k in self._out_features}
|
||||||
|
assert fuse_type in {"avg", "sum"}
|
||||||
|
self._fuse_type = fuse_type
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
"""
|
||||||
|
Args:
|
||||||
|
input (dict[str: Tensor]): mapping feature map name (e.g., "res5") to
|
||||||
|
feature map tensor for each feature level in high to low resolution order.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict[str: Tensor]:
|
||||||
|
mapping from feature map name to FPN feature map tensor
|
||||||
|
in high to low resolution order. Returned feature names follow the FPN
|
||||||
|
paper convention: "p<stage>", where stage has stride = 2 ** stage e.g.,
|
||||||
|
["p2", "p3", ..., "p6"].
|
||||||
|
"""
|
||||||
|
# Reverse feature maps into top-down order (from low to high resolution)
|
||||||
|
bottom_up_features = self.bottom_up(x)
|
||||||
|
x = [bottom_up_features[f] for f in self.in_features[::-1]]
|
||||||
|
results = []
|
||||||
|
prev_features = self.lateral_convs[0](x[0])
|
||||||
|
results.append(self.output_convs[0](prev_features))
|
||||||
|
for features, lateral_conv, output_conv in zip(
|
||||||
|
x[1:], self.lateral_convs[1:], self.output_convs[1:]
|
||||||
|
):
|
||||||
|
top_down_features = F.interpolate(prev_features, scale_factor=2, mode="nearest")
|
||||||
|
lateral_features = lateral_conv(features)
|
||||||
|
prev_features = lateral_features + top_down_features
|
||||||
|
if self._fuse_type == "avg":
|
||||||
|
prev_features /= 2
|
||||||
|
results.insert(0, output_conv(prev_features))
|
||||||
|
|
||||||
|
if self.top_block is not None:
|
||||||
|
top_block_in_feature = bottom_up_features.get(self.top_block.in_feature, None)
|
||||||
|
if top_block_in_feature is None:
|
||||||
|
top_block_in_feature = results[self._out_features.index(self.top_block.in_feature)]
|
||||||
|
results.extend(self.top_block(top_block_in_feature))
|
||||||
|
assert len(self._out_features) == len(results)
|
||||||
|
return dict(zip(self._out_features, results))
|
||||||
|
|
||||||
|
def _assert_strides_are_log2_contiguous(strides):
|
||||||
|
"""
|
||||||
|
Assert that each stride is 2x times its preceding stride, i.e. "contiguous in log2".
|
||||||
|
"""
|
||||||
|
for i, stride in enumerate(strides[1:], 1):
|
||||||
|
assert stride == 2 * strides[i - 1], "Strides {} {} are not log2 contiguous".format(
|
||||||
|
stride, strides[i - 1]
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class LastLevelMaxPool(nn.Module):
|
||||||
|
"""
|
||||||
|
This module is used in the original FPN to generate a downsampled
|
||||||
|
P6 feature from P5.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__()
|
||||||
|
self.num_levels = 1
|
||||||
|
self.in_feature = "p5"
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
return [F.max_pool2d(x, kernel_size=1, stride=2, padding=0)]
|
||||||
|
|
||||||
|
|
||||||
|
class LastLevelP6P7(nn.Module):
|
||||||
|
"""
|
||||||
|
This module is used in RetinaNet to generate extra layers, P6 and P7 from
|
||||||
|
C5 feature.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, in_channels, out_channels):
|
||||||
|
super().__init__()
|
||||||
|
self.num_levels = 2
|
||||||
|
self.in_feature = "res5"
|
||||||
|
self.p6 = nn.Conv2d(in_channels, out_channels, 3, 2, 1)
|
||||||
|
self.p7 = nn.Conv2d(out_channels, out_channels, 3, 2, 1)
|
||||||
|
for module in [self.p6, self.p7]:
|
||||||
|
weight_init.c2_xavier_fill(module)
|
||||||
|
|
||||||
|
def forward(self, c5):
|
||||||
|
p6 = self.p6(c5)
|
||||||
|
p7 = self.p7(F.relu(p6))
|
||||||
|
return [p6, p7]
|
||||||
|
|
||||||
|
|
||||||
|
def build_resnet_fpn_backbone(in_channels=3):
|
||||||
|
"""
|
||||||
|
Args:
|
||||||
|
cfg: a detectron2 CfgNode
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
backbone (Backbone): backbone module, must be a subclass of :class:`Backbone`.
|
||||||
|
"""
|
||||||
|
bottom_up = build_resnet_backbone(in_channels)
|
||||||
|
in_features = ["res2", "res3", "res4"]
|
||||||
|
out_channels = 256
|
||||||
|
backbone = FPN(
|
||||||
|
bottom_up=bottom_up,
|
||||||
|
in_features=in_features,
|
||||||
|
out_channels=out_channels,
|
||||||
|
norm="BN",
|
||||||
|
# top_block=LastLevelMaxPool(),
|
||||||
|
top_block=None,
|
||||||
|
fuse_type="sum",
|
||||||
|
)
|
||||||
|
return backbone
|
||||||
|
|
||||||
|
def build_retinanet_resnet_fpn_backbone(cfg, in_channels=3):
|
||||||
|
"""
|
||||||
|
Args:
|
||||||
|
cfg: a detectron2 CfgNode
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
backbone (Backbone): backbone module, must be a subclass of :class:`Backbone`.
|
||||||
|
"""
|
||||||
|
bottom_up = build_resnet_backbone(cfg, in_channels)
|
||||||
|
in_features = cfg.MODEL.FPN.IN_FEATURES
|
||||||
|
out_channels = cfg.MODEL.FPN.OUT_CHANNELS
|
||||||
|
in_channels_p6p7 = bottom_up._out_feature_channels["res5"]
|
||||||
|
backbone = FPN(
|
||||||
|
bottom_up=bottom_up,
|
||||||
|
in_features=in_features,
|
||||||
|
out_channels=out_channels,
|
||||||
|
norm=cfg.MODEL.FPN.NORM,
|
||||||
|
top_block=LastLevelP6P7(in_channels_p6p7, out_channels),
|
||||||
|
fuse_type=cfg.MODEL.FPN.FUSE_TYPE,
|
||||||
|
)
|
||||||
|
return backbone
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
import argparse
|
||||||
|
from config.default import get_cfg
|
||||||
|
|
||||||
|
def setup(args):
|
||||||
|
"""
|
||||||
|
Create configs and perform basic setups.
|
||||||
|
"""
|
||||||
|
cfg = get_cfg()
|
||||||
|
cfg.merge_from_file(args.cfg)
|
||||||
|
cfg.merge_from_list(args.opts)
|
||||||
|
cfg.freeze()
|
||||||
|
return cfg
|
||||||
|
|
||||||
|
parser = argparse.ArgumentParser(description='Train ImageNet network')
|
||||||
|
# general
|
||||||
|
parser.add_argument('--cfg',
|
||||||
|
help='experiment configure file name',
|
||||||
|
required=True,
|
||||||
|
type=str)
|
||||||
|
|
||||||
|
parser.add_argument('opts',
|
||||||
|
help="Modify config options using the command-line",
|
||||||
|
default=None,
|
||||||
|
nargs=argparse.REMAINDER)
|
||||||
|
|
||||||
|
args = parser.parse_args()
|
||||||
|
cfg = setup(args)
|
||||||
|
print(cfg)
|
||||||
|
|
||||||
|
model = build_resnet_fpn_backbone(cfg, 3)
|
||||||
|
# model = build_retinanet_resnet_fpn_backbone(cfg, 3)
|
||||||
|
print(model)
|
||||||
|
# model = torch.nn.DataParallel(model, list(range(2))).cuda()
|
||||||
|
dummy_input = torch.randn(4, 3, 512, 512)
|
||||||
|
|
||||||
|
out = model(dummy_input)
|
||||||
|
|
||||||
|
for k, v in out.items():
|
||||||
|
print(k, v.shape)
|
||||||
|
|
||||||
|
# torch.onnx.export(model, dummy_input, "tmp.onnx", verbose=True,
|
||||||
|
# input_names=['input'],
|
||||||
|
# output_names=['output'])
|
||||||
|
|
||||||
|
pass
|
||||||
|
|
||||||
@@ -0,0 +1,298 @@
|
|||||||
|
import numpy as np
|
||||||
|
import utils.weight_init as weight_init
|
||||||
|
import torch
|
||||||
|
import torch.nn.functional as F
|
||||||
|
from torch import nn
|
||||||
|
from bodyseg.backbone.backbone import Backbone, get_norm, Conv2d
|
||||||
|
|
||||||
|
class BasicStem(nn.Module):
|
||||||
|
def __init__(self, in_channels=3, out_channels=64, norm="BN"):
|
||||||
|
"""
|
||||||
|
Args:
|
||||||
|
norm (str or callable): a callable that takes the number of
|
||||||
|
channels and return a `nn.Module`, or a pre-defined string
|
||||||
|
(one of {"FrozenBN", "BN", "GN"}).
|
||||||
|
"""
|
||||||
|
super().__init__()
|
||||||
|
self.conv1 = Conv2d(
|
||||||
|
in_channels,
|
||||||
|
out_channels,
|
||||||
|
kernel_size=7,
|
||||||
|
stride=2,
|
||||||
|
padding=3,
|
||||||
|
bias=False,
|
||||||
|
norm=get_norm(norm, out_channels),
|
||||||
|
)
|
||||||
|
weight_init.c2_msra_fill(self.conv1)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x = self.conv1(x)
|
||||||
|
x = F.relu_(x)
|
||||||
|
x = F.max_pool2d(x, kernel_size=3, stride=2, padding=1)
|
||||||
|
return x
|
||||||
|
|
||||||
|
@property
|
||||||
|
def out_channels(self):
|
||||||
|
return self.conv1.out_channels
|
||||||
|
|
||||||
|
@property
|
||||||
|
def stride(self):
|
||||||
|
return 4 # = stride 2 conv -> stride 2 max pool
|
||||||
|
|
||||||
|
class ResNetBlockBase(nn.Module):
|
||||||
|
def __init__(self, in_channels, out_channels, stride):
|
||||||
|
"""
|
||||||
|
The `__init__` method of any subclass should also contain these arguments.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
in_channels (int):
|
||||||
|
out_channels (int):
|
||||||
|
stride (int):
|
||||||
|
"""
|
||||||
|
super().__init__()
|
||||||
|
self.in_channels = in_channels
|
||||||
|
self.out_channels = out_channels
|
||||||
|
self.stride = stride
|
||||||
|
|
||||||
|
class BottleneckBlock(ResNetBlockBase):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
in_channels,
|
||||||
|
out_channels,
|
||||||
|
*,
|
||||||
|
bottleneck_channels,
|
||||||
|
stride=1,
|
||||||
|
num_groups=1,
|
||||||
|
norm="BN",
|
||||||
|
stride_in_1x1=False,
|
||||||
|
dilation=1,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Args:
|
||||||
|
norm (str or callable): a callable that takes the number of
|
||||||
|
channels and return a `nn.Module`, or a pre-defined string
|
||||||
|
(one of {"FrozenBN", "BN", "GN"}).
|
||||||
|
stride_in_1x1 (bool): when stride==2, whether to put stride in the
|
||||||
|
first 1x1 convolution or the bottleneck 3x3 convolution.
|
||||||
|
"""
|
||||||
|
super().__init__(in_channels, out_channels, stride)
|
||||||
|
|
||||||
|
if in_channels != out_channels:
|
||||||
|
self.shortcut = Conv2d(
|
||||||
|
in_channels,
|
||||||
|
out_channels,
|
||||||
|
kernel_size=1,
|
||||||
|
stride=stride,
|
||||||
|
bias=False,
|
||||||
|
norm=get_norm(norm, out_channels),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
self.shortcut = None
|
||||||
|
|
||||||
|
# The original MSRA ResNet models have stride in the first 1x1 conv
|
||||||
|
# The subsequent fb.torch.resnet and Caffe2 ResNe[X]t implementations have
|
||||||
|
# stride in the 3x3 conv
|
||||||
|
stride_1x1, stride_3x3 = (stride, 1) if stride_in_1x1 else (1, stride)
|
||||||
|
|
||||||
|
self.conv1 = Conv2d(
|
||||||
|
in_channels,
|
||||||
|
bottleneck_channels,
|
||||||
|
kernel_size=1,
|
||||||
|
stride=stride_1x1,
|
||||||
|
bias=False,
|
||||||
|
norm=get_norm(norm, bottleneck_channels),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.conv2 = Conv2d(
|
||||||
|
bottleneck_channels,
|
||||||
|
bottleneck_channels,
|
||||||
|
kernel_size=3,
|
||||||
|
stride=stride_3x3,
|
||||||
|
padding=1 * dilation,
|
||||||
|
bias=False,
|
||||||
|
groups=num_groups,
|
||||||
|
dilation=dilation,
|
||||||
|
norm=get_norm(norm, bottleneck_channels),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.conv3 = Conv2d(
|
||||||
|
bottleneck_channels,
|
||||||
|
out_channels,
|
||||||
|
kernel_size=1,
|
||||||
|
bias=False,
|
||||||
|
norm=get_norm(norm, out_channels),
|
||||||
|
)
|
||||||
|
|
||||||
|
for layer in [self.conv1, self.conv2, self.conv3, self.shortcut]:
|
||||||
|
if layer is not None: # shortcut can be None
|
||||||
|
weight_init.c2_msra_fill(layer)
|
||||||
|
|
||||||
|
# Zero-initialize the last normalization in each residual branch,
|
||||||
|
# so that at the beginning, the residual branch starts with zeros,
|
||||||
|
# and each residual block behaves like an identity.
|
||||||
|
# See Sec 5.1 in "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour":
|
||||||
|
# "For BN layers, the learnable scaling coefficient γ is initialized
|
||||||
|
# to be 1, except for each residual block's last BN
|
||||||
|
# where γ is initialized to be 0."
|
||||||
|
|
||||||
|
# nn.init.constant_(self.conv3.norm.weight, 0)
|
||||||
|
# TODO this somehow hurts performance when training GN models from scratch.
|
||||||
|
# Add it as an option when we need to use this code to train a backbone.
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
out = self.conv1(x)
|
||||||
|
out = F.relu_(out)
|
||||||
|
|
||||||
|
out = self.conv2(out)
|
||||||
|
out = F.relu_(out)
|
||||||
|
|
||||||
|
out = self.conv3(out)
|
||||||
|
|
||||||
|
if self.shortcut is not None:
|
||||||
|
shortcut = self.shortcut(x)
|
||||||
|
else:
|
||||||
|
shortcut = x
|
||||||
|
|
||||||
|
out += shortcut
|
||||||
|
out = F.relu_(out)
|
||||||
|
return out
|
||||||
|
|
||||||
|
def make_stage(block_class, num_blocks, first_stride, **kwargs):
|
||||||
|
"""
|
||||||
|
Create a resnet stage by creating many blocks.
|
||||||
|
Args:
|
||||||
|
block_class (class): a subclass of ResNetBlockBase
|
||||||
|
num_blocks (int):
|
||||||
|
first_stride (int): the stride of the first block. The other blocks will have stride=1.
|
||||||
|
A `stride` argument will be passed to the block constructor.
|
||||||
|
kwargs: other arguments passed to the block constructor.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
list[nn.Module]: a list of block module.
|
||||||
|
"""
|
||||||
|
blocks = []
|
||||||
|
for i in range(num_blocks):
|
||||||
|
blocks.append(block_class(stride=first_stride if i == 0 else 1, **kwargs))
|
||||||
|
kwargs["in_channels"] = kwargs["out_channels"]
|
||||||
|
return blocks
|
||||||
|
|
||||||
|
class ResNet(Backbone):
|
||||||
|
def __init__(self, stem, stages, num_classes=None, out_features=None):
|
||||||
|
"""
|
||||||
|
Args:
|
||||||
|
stem (nn.Module): a stem module
|
||||||
|
stages (list[list[ResNetBlock]]): several (typically 4) stages,
|
||||||
|
each contains multiple :class:`ResNetBlockBase`.
|
||||||
|
num_classes (None or int): if None, will not perform classification.
|
||||||
|
out_features (list[str]): name of the layers whose outputs should
|
||||||
|
be returned in forward. Can be anything in "stem", "linear", or "res2" ...
|
||||||
|
If None, will return the output of the last layer.
|
||||||
|
"""
|
||||||
|
super(ResNet, self).__init__()
|
||||||
|
self.stem = stem
|
||||||
|
self.num_classes = num_classes
|
||||||
|
|
||||||
|
current_stride = self.stem.stride
|
||||||
|
self._out_feature_strides = {"stem": current_stride}
|
||||||
|
self._out_feature_channels = {"stem": self.stem.out_channels}
|
||||||
|
|
||||||
|
self.stages = []
|
||||||
|
self.names = []
|
||||||
|
for i, blocks in enumerate(stages):
|
||||||
|
for block in blocks:
|
||||||
|
assert isinstance(block, ResNetBlockBase), block
|
||||||
|
curr_channels = block.out_channels
|
||||||
|
stage = nn.Sequential(*blocks)
|
||||||
|
name = "res" + str(i + 2)
|
||||||
|
self.stages.append(stage)
|
||||||
|
self.names.append(name)
|
||||||
|
self._out_feature_strides[name] = current_stride = int(
|
||||||
|
current_stride * np.prod([k.stride for k in blocks])
|
||||||
|
)
|
||||||
|
self._out_feature_channels[name] = blocks[-1].out_channels
|
||||||
|
self.stages = nn.ModuleList(self.stages)
|
||||||
|
|
||||||
|
if num_classes is not None:
|
||||||
|
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
|
||||||
|
self.linear = nn.Linear(curr_channels, num_classes)
|
||||||
|
|
||||||
|
# Sec 5.1 in "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour":
|
||||||
|
# "The 1000-way fully-connected layer is initialized by
|
||||||
|
# drawing weights from a zero-mean Gaussian with standard deviation of 0.01."
|
||||||
|
nn.init.normal_(self.linear.weight, stddev=0.01)
|
||||||
|
name = "linear"
|
||||||
|
|
||||||
|
if out_features is None:
|
||||||
|
out_features = [name]
|
||||||
|
self._out_features = out_features
|
||||||
|
assert len(self._out_features)
|
||||||
|
for out_feature in self._out_features:
|
||||||
|
assert out_feature in self.names, "Available children: {}".format(", ".join(self.names))
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
outputs = {}
|
||||||
|
x = self.stem(x)
|
||||||
|
if "stem" in self._out_features:
|
||||||
|
outputs["stem"] = x
|
||||||
|
for ix, stage in enumerate(self.stages):
|
||||||
|
name = self.names[ix]
|
||||||
|
x = stage(x)
|
||||||
|
if name in self._out_features:
|
||||||
|
outputs[name] = x
|
||||||
|
if self.num_classes is not None:
|
||||||
|
x = self.avgpool(x)
|
||||||
|
x = self.linear(x)
|
||||||
|
if "linear" in self._out_features:
|
||||||
|
outputs["linear"] = x
|
||||||
|
return outputs
|
||||||
|
|
||||||
|
def build_resnet_backbone(in_channels=3):
|
||||||
|
norm = "BN"
|
||||||
|
stem = BasicStem(
|
||||||
|
in_channels=in_channels,
|
||||||
|
out_channels=64,
|
||||||
|
norm=norm,
|
||||||
|
)
|
||||||
|
|
||||||
|
# fmt: off
|
||||||
|
out_features = ["res2", "res3", "res4"]
|
||||||
|
depth = 101
|
||||||
|
num_groups = 1
|
||||||
|
bottleneck_channels = 64
|
||||||
|
in_channels = 64
|
||||||
|
out_channels = 256
|
||||||
|
stride_in_1x1 = True
|
||||||
|
res5_dilation = 1
|
||||||
|
# fmt: on
|
||||||
|
assert res5_dilation in {1, 2}, "res5_dilation cannot be {}.".format(res5_dilation)
|
||||||
|
|
||||||
|
num_blocks_per_stage = {50: [3, 4, 6, 3], 101: [3, 4, 23, 3], 152: [3, 8, 36, 3]}[depth]
|
||||||
|
|
||||||
|
stages = []
|
||||||
|
|
||||||
|
# Avoid creating variables without gradients
|
||||||
|
# It consumes extra memory and may cause allreduce to fail
|
||||||
|
out_stage_idx = [{"res2": 2, "res3": 3, "res4": 4, "res5": 5}[f] for f in out_features]
|
||||||
|
max_stage_idx = max(out_stage_idx)
|
||||||
|
for idx, stage_idx in enumerate(range(2, max_stage_idx + 1)):
|
||||||
|
dilation = res5_dilation if stage_idx == 5 else 1
|
||||||
|
first_stride = 1 if idx == 0 or (stage_idx == 5 and dilation == 2) else 2
|
||||||
|
stage_kargs = dict()
|
||||||
|
stage_kargs.update({
|
||||||
|
"num_blocks": num_blocks_per_stage[idx],
|
||||||
|
"first_stride": first_stride,
|
||||||
|
"in_channels": in_channels,
|
||||||
|
"bottleneck_channels": bottleneck_channels,
|
||||||
|
"out_channels": out_channels,
|
||||||
|
"num_groups": num_groups,
|
||||||
|
"norm": norm,
|
||||||
|
"stride_in_1x1": stride_in_1x1,
|
||||||
|
"dilation": dilation,
|
||||||
|
})
|
||||||
|
stage_kargs["block_class"] = BottleneckBlock
|
||||||
|
blocks = make_stage(**stage_kargs)
|
||||||
|
in_channels = out_channels
|
||||||
|
out_channels *= 2
|
||||||
|
bottleneck_channels *= 2
|
||||||
|
stages.append(blocks)
|
||||||
|
return ResNet(stem, stages, out_features=out_features)
|
||||||
@@ -0,0 +1,244 @@
|
|||||||
|
import math
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import torch.nn.functional as F
|
||||||
|
import torch.utils.model_zoo as model_zoo
|
||||||
|
from bodyseg.backbone.backbone import Backbone
|
||||||
|
|
||||||
|
def fixed_padding(inputs, kernel_size, dilation):
|
||||||
|
kernel_size_effective = kernel_size + (kernel_size - 1) * (dilation - 1)
|
||||||
|
pad_total = kernel_size_effective - 1
|
||||||
|
pad_beg = pad_total // 2
|
||||||
|
pad_end = pad_total - pad_beg
|
||||||
|
padded_inputs = F.pad(inputs, (pad_beg, pad_end, pad_beg, pad_end))
|
||||||
|
return padded_inputs
|
||||||
|
|
||||||
|
class SeparableConv2d(nn.Module):
|
||||||
|
def __init__(self, inplanes, planes, kernel_size=3, stride=1, dilation=1, bias=False, BatchNorm=None):
|
||||||
|
super(SeparableConv2d, self).__init__()
|
||||||
|
|
||||||
|
self.conv1 = nn.Conv2d(inplanes, inplanes, kernel_size, stride, 0, dilation,
|
||||||
|
groups=inplanes, bias=bias)
|
||||||
|
self.bn = BatchNorm(inplanes)
|
||||||
|
self.pointwise = nn.Conv2d(inplanes, planes, 1, 1, 0, 1, 1, bias=bias)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x = fixed_padding(x, self.conv1.kernel_size[0], dilation=self.conv1.dilation[0])
|
||||||
|
x = self.conv1(x)
|
||||||
|
x = self.bn(x)
|
||||||
|
x = self.pointwise(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
class Block(nn.Module):
|
||||||
|
def __init__(self, inplanes, planes, reps, stride=1, dilation=1, BatchNorm=None,
|
||||||
|
start_with_relu=True, grow_first=True, is_last=False):
|
||||||
|
super(Block, self).__init__()
|
||||||
|
|
||||||
|
if planes != inplanes or stride != 1:
|
||||||
|
self.skip = nn.Conv2d(inplanes, planes, 1, stride=stride, bias=False)
|
||||||
|
self.skipbn = BatchNorm(planes)
|
||||||
|
else:
|
||||||
|
self.skip = None
|
||||||
|
|
||||||
|
self.relu = nn.ReLU(inplace=True)
|
||||||
|
rep = []
|
||||||
|
|
||||||
|
filters = inplanes
|
||||||
|
if grow_first:
|
||||||
|
rep.append(self.relu)
|
||||||
|
rep.append(SeparableConv2d(inplanes, planes, 3, 1, dilation, BatchNorm=BatchNorm))
|
||||||
|
rep.append(BatchNorm(planes))
|
||||||
|
filters = planes
|
||||||
|
|
||||||
|
for i in range(reps - 1):
|
||||||
|
rep.append(self.relu)
|
||||||
|
rep.append(SeparableConv2d(filters, filters, 3, 1, dilation, BatchNorm=BatchNorm))
|
||||||
|
rep.append(BatchNorm(filters))
|
||||||
|
|
||||||
|
if not grow_first:
|
||||||
|
rep.append(self.relu)
|
||||||
|
rep.append(SeparableConv2d(inplanes, planes, 3, 1, dilation, BatchNorm=BatchNorm))
|
||||||
|
rep.append(BatchNorm(planes))
|
||||||
|
|
||||||
|
if stride != 1:
|
||||||
|
rep.append(self.relu)
|
||||||
|
rep.append(SeparableConv2d(planes, planes, 3, 2, BatchNorm=BatchNorm))
|
||||||
|
rep.append(BatchNorm(planes))
|
||||||
|
|
||||||
|
if stride == 1 and is_last:
|
||||||
|
rep.append(self.relu)
|
||||||
|
rep.append(SeparableConv2d(planes, planes, 3, 1, BatchNorm=BatchNorm))
|
||||||
|
rep.append(BatchNorm(planes))
|
||||||
|
|
||||||
|
if not start_with_relu:
|
||||||
|
rep = rep[1:]
|
||||||
|
|
||||||
|
self.rep = nn.Sequential(*rep)
|
||||||
|
|
||||||
|
def forward(self, inp):
|
||||||
|
x = self.rep(inp)
|
||||||
|
|
||||||
|
if self.skip is not None:
|
||||||
|
skip = self.skip(inp)
|
||||||
|
skip = self.skipbn(skip)
|
||||||
|
else:
|
||||||
|
skip = inp
|
||||||
|
|
||||||
|
x = x + skip
|
||||||
|
|
||||||
|
return x
|
||||||
|
|
||||||
|
class AlignedXception(Backbone):
|
||||||
|
"""
|
||||||
|
Modified Alighed Xception
|
||||||
|
"""
|
||||||
|
def __init__(self, output_stride, BatchNorm):
|
||||||
|
super(AlignedXception, self).__init__()
|
||||||
|
|
||||||
|
if output_stride == 16:
|
||||||
|
entry_block3_stride = 2
|
||||||
|
middle_block_dilation = 1
|
||||||
|
exit_block_dilations = (1, 2)
|
||||||
|
elif output_stride == 8:
|
||||||
|
entry_block3_stride = 1
|
||||||
|
middle_block_dilation = 2
|
||||||
|
exit_block_dilations = (2, 4)
|
||||||
|
else:
|
||||||
|
raise NotImplementedError
|
||||||
|
|
||||||
|
|
||||||
|
# Entry flow
|
||||||
|
self.conv1 = nn.Conv2d(3, 32, 3, stride=2, padding=1, bias=False)
|
||||||
|
self.bn1 = BatchNorm(32)
|
||||||
|
self.relu = nn.ReLU(inplace=True)
|
||||||
|
|
||||||
|
self.conv2 = nn.Conv2d(32, 64, 3, stride=1, padding=1, bias=False)
|
||||||
|
self.bn2 = BatchNorm(64)
|
||||||
|
|
||||||
|
self.block1 = Block(64, 128, reps=2, stride=2, BatchNorm=BatchNorm, start_with_relu=False)
|
||||||
|
self.block2 = Block(128, 256, reps=2, stride=2, BatchNorm=BatchNorm, start_with_relu=False,
|
||||||
|
grow_first=True)
|
||||||
|
self.block3 = Block(256, 728, reps=2, stride=entry_block3_stride, BatchNorm=BatchNorm,
|
||||||
|
start_with_relu=True, grow_first=True, is_last=True)
|
||||||
|
|
||||||
|
# Middle flow
|
||||||
|
self.block4 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block5 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block6 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block7 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block8 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block9 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block10 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block11 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block12 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block13 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block14 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block15 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block16 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block17 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block18 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block19 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
|
||||||
|
# Exit flow
|
||||||
|
self.block20 = Block(728, 1024, reps=2, stride=1, dilation=exit_block_dilations[0],
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=False, is_last=True)
|
||||||
|
|
||||||
|
self.conv3 = SeparableConv2d(1024, 1536, 3, stride=1, dilation=exit_block_dilations[1], BatchNorm=BatchNorm)
|
||||||
|
self.bn3 = BatchNorm(1536)
|
||||||
|
|
||||||
|
self.conv4 = SeparableConv2d(1536, 1536, 3, stride=1, dilation=exit_block_dilations[1], BatchNorm=BatchNorm)
|
||||||
|
self.bn4 = BatchNorm(1536)
|
||||||
|
|
||||||
|
self.conv5 = SeparableConv2d(1536, 2048, 3, stride=1, dilation=exit_block_dilations[1], BatchNorm=BatchNorm)
|
||||||
|
self.bn5 = BatchNorm(2048)
|
||||||
|
|
||||||
|
# Init weights
|
||||||
|
self._init_weight()
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
# Entry flow
|
||||||
|
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.block1(x)
|
||||||
|
# add relu here
|
||||||
|
x = self.relu(x)
|
||||||
|
low_level_feat = x
|
||||||
|
x = self.block2(x)
|
||||||
|
x = self.block3(x)
|
||||||
|
|
||||||
|
# Middle flow
|
||||||
|
x = self.block4(x)
|
||||||
|
x = self.block5(x)
|
||||||
|
x = self.block6(x)
|
||||||
|
x = self.block7(x)
|
||||||
|
x = self.block8(x)
|
||||||
|
x = self.block9(x)
|
||||||
|
x = self.block10(x)
|
||||||
|
x = self.block11(x)
|
||||||
|
x = self.block12(x)
|
||||||
|
x = self.block13(x)
|
||||||
|
x = self.block14(x)
|
||||||
|
x = self.block15(x)
|
||||||
|
x = self.block16(x)
|
||||||
|
x = self.block17(x)
|
||||||
|
x = self.block18(x)
|
||||||
|
x = self.block19(x)
|
||||||
|
|
||||||
|
# Exit flow
|
||||||
|
x = self.block20(x)
|
||||||
|
x = self.relu(x)
|
||||||
|
x = self.conv3(x)
|
||||||
|
x = self.bn3(x)
|
||||||
|
x = self.relu(x)
|
||||||
|
|
||||||
|
x = self.conv4(x)
|
||||||
|
x = self.bn4(x)
|
||||||
|
x = self.relu(x)
|
||||||
|
|
||||||
|
x = self.conv5(x)
|
||||||
|
x = self.bn5(x)
|
||||||
|
x = self.relu(x)
|
||||||
|
|
||||||
|
return x, low_level_feat
|
||||||
|
|
||||||
|
def _init_weight(self):
|
||||||
|
for m in self.modules():
|
||||||
|
if isinstance(m, nn.Conv2d):
|
||||||
|
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
|
||||||
|
m.weight.data.normal_(0, math.sqrt(2. / n))
|
||||||
|
elif isinstance(m, nn.SyncBatchNorm):
|
||||||
|
m.weight.data.fill_(1)
|
||||||
|
m.bias.data.zero_()
|
||||||
|
elif isinstance(m, nn.BatchNorm2d):
|
||||||
|
m.weight.data.fill_(1)
|
||||||
|
m.bias.data.zero_()
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
import torch
|
||||||
|
model = AlignedXception(BatchNorm=nn.BatchNorm2d, output_stride=16)
|
||||||
|
input = torch.rand(1, 3, 512, 512)
|
||||||
|
output, low_level_feat = model(input)
|
||||||
|
print(output.size())
|
||||||
|
print(low_level_feat.size())
|
||||||
@@ -0,0 +1,330 @@
|
|||||||
|
import sys
|
||||||
|
# sys.path.append("/Users/momo/human_seg_train")
|
||||||
|
# print(sys.path)
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import torch.nn.functional as F
|
||||||
|
from bodyseg.backbone.backbone import get_norm
|
||||||
|
|
||||||
|
class ConvBNReLU(nn.Sequential):
|
||||||
|
def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1, norm_layer=None):
|
||||||
|
padding = (kernel_size - 1) // 2
|
||||||
|
if norm_layer is None:
|
||||||
|
norm_layer = nn.BatchNorm2d
|
||||||
|
super(ConvBNReLU, self).__init__(
|
||||||
|
nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
|
||||||
|
norm_layer(out_planes),
|
||||||
|
nn.ReLU6(inplace=True)
|
||||||
|
)
|
||||||
|
|
||||||
|
class InvertedResidual(nn.Module):
|
||||||
|
def __init__(self, inp, oup, stride, expand_ratio, norm_layer=None):
|
||||||
|
super(InvertedResidual, self).__init__()
|
||||||
|
self.stride = stride
|
||||||
|
assert stride in [1, 2]
|
||||||
|
|
||||||
|
if norm_layer is None:
|
||||||
|
norm_layer = nn.BatchNorm2d
|
||||||
|
|
||||||
|
hidden_dim = int(round(inp * expand_ratio))
|
||||||
|
self.use_res_connect = self.stride == 1 and inp == oup
|
||||||
|
|
||||||
|
layers = []
|
||||||
|
if expand_ratio != 1:
|
||||||
|
# pw
|
||||||
|
layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1, norm_layer=norm_layer))
|
||||||
|
layers.extend([
|
||||||
|
# dw
|
||||||
|
ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim, norm_layer=norm_layer),
|
||||||
|
# pw-linear
|
||||||
|
nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
|
||||||
|
norm_layer(oup),
|
||||||
|
])
|
||||||
|
self.conv = nn.Sequential(*layers)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
if self.use_res_connect:
|
||||||
|
return x + self.conv(x)
|
||||||
|
else:
|
||||||
|
return self.conv(x)
|
||||||
|
|
||||||
|
class UpSampleBlock(nn.Module):
|
||||||
|
def __init__(self, in_channels, out_channels, expand_ratio=6):
|
||||||
|
super(UpSampleBlock, self).__init__()
|
||||||
|
self.refine = InvertedResidual(in_channels, out_channels, 1, expand_ratio)
|
||||||
|
|
||||||
|
def forward(self, x0, x1):
|
||||||
|
x = torch.cat([x0, x1], dim=1)
|
||||||
|
x = self.refine(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
class BodySegNet_32_thin_sigmod(nn.Module):
|
||||||
|
def __init__(self, input_channels=3, class_nums=1, output_onnx=False):
|
||||||
|
super(BodySegNet_32_thin_sigmod, self).__init__()
|
||||||
|
self.class_nums = class_nums
|
||||||
|
self.output_onnx = output_onnx
|
||||||
|
|
||||||
|
self.stage_1 = nn.Sequential(
|
||||||
|
ConvBNReLU(input_channels, 16, kernel_size=3, stride=2)
|
||||||
|
)
|
||||||
|
self.stage_2 = nn.Sequential(
|
||||||
|
ConvBNReLU(16, 16, kernel_size=3, stride=2, groups=16),
|
||||||
|
ConvBNReLU(16, 16, kernel_size=1, stride=1),
|
||||||
|
)
|
||||||
|
self.stage_3 = nn.Sequential(
|
||||||
|
InvertedResidual(16, 24, stride=2, expand_ratio=6),
|
||||||
|
InvertedResidual(24, 24, stride=1, expand_ratio=6),
|
||||||
|
InvertedResidual(24, 24, stride=1, expand_ratio=6),
|
||||||
|
)
|
||||||
|
self.stage_4 = nn.Sequential(
|
||||||
|
InvertedResidual(24, 32, stride=2, expand_ratio=6),
|
||||||
|
InvertedResidual(32, 32, stride=1, expand_ratio=6),
|
||||||
|
InvertedResidual(32, 32, stride=1, expand_ratio=6),
|
||||||
|
InvertedResidual(32, 32, stride=1, expand_ratio=6),
|
||||||
|
)
|
||||||
|
self.stage_5 = nn.Sequential(
|
||||||
|
InvertedResidual(32, 48, stride=2, expand_ratio=6),
|
||||||
|
InvertedResidual(48, 48, stride=1, expand_ratio=6),
|
||||||
|
InvertedResidual(48, 48, stride=1, expand_ratio=6),
|
||||||
|
InvertedResidual(48, 48, stride=1, expand_ratio=6)
|
||||||
|
)
|
||||||
|
self.up_to_4 = UpSampleBlock(48 + 32, 16)
|
||||||
|
self.up_to_3 = UpSampleBlock(16 + 24, 16)
|
||||||
|
self.up_to_2 = UpSampleBlock(16 + 16, 16)
|
||||||
|
self.up_to_1 = UpSampleBlock(16 + 16, 16)
|
||||||
|
self.last_layer = nn.Sequential(
|
||||||
|
ConvBNReLU(16, 16, kernel_size=1, stride=1),
|
||||||
|
nn.Conv2d(16, self.class_nums, kernel_size=1, stride=1)
|
||||||
|
)
|
||||||
|
self._initialize_weights()
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
feature_S = []
|
||||||
|
x1 = self.stage_1(x)
|
||||||
|
x2 = self.stage_2(x1)
|
||||||
|
x3 = self.stage_3(x2)
|
||||||
|
x4 = self.stage_4(x3)
|
||||||
|
feature = self.stage_5(x4)
|
||||||
|
|
||||||
|
feature = F.interpolate(feature, size=x4.size()[2:], mode='bilinear', align_corners=True)
|
||||||
|
feature = self.up_to_4(x4, feature)
|
||||||
|
feature = F.interpolate(feature, size=x3.size()[2:], mode='bilinear', align_corners=True)
|
||||||
|
feature = self.up_to_3(x3, feature)
|
||||||
|
feature = F.interpolate(feature, size=x2.size()[2:], mode='bilinear', align_corners=True)
|
||||||
|
feature = self.up_to_2(x2, feature)
|
||||||
|
feature = F.interpolate(feature, size=x1.size()[2:], mode='bilinear', align_corners=True)
|
||||||
|
feature = self.up_to_1(x1, feature)
|
||||||
|
feature_S.append(feature)
|
||||||
|
feature = self.last_layer(feature)
|
||||||
|
feature_S.append(feature)
|
||||||
|
output = F.interpolate(feature, size=x.size()[2:], mode='bilinear', align_corners=True)
|
||||||
|
# output = torch.sigmoid(output)
|
||||||
|
if self.output_onnx:
|
||||||
|
output = torch.argmax(output, dim=1).to(torch.float32)
|
||||||
|
|
||||||
|
return output, feature_S
|
||||||
|
|
||||||
|
def _initialize_weights(self):
|
||||||
|
for name, m in self.named_modules():
|
||||||
|
if isinstance(m, nn.Conv2d):
|
||||||
|
if 'first' in name:
|
||||||
|
nn.init.normal_(m.weight, 0, 0.01)
|
||||||
|
else:
|
||||||
|
nn.init.normal_(m.weight, 0, 1.0 / m.weight.shape[1])
|
||||||
|
if m.bias is not None:
|
||||||
|
nn.init.constant_(m.bias, 0)
|
||||||
|
elif isinstance(m, nn.BatchNorm2d):
|
||||||
|
nn.init.constant_(m.weight, 1)
|
||||||
|
if m.bias is not None:
|
||||||
|
nn.init.constant_(m.bias, 0.0001)
|
||||||
|
nn.init.constant_(m.running_mean, 0)
|
||||||
|
elif isinstance(m, nn.BatchNorm1d):
|
||||||
|
nn.init.constant_(m.weight, 1)
|
||||||
|
if m.bias is not None:
|
||||||
|
nn.init.constant_(m.bias, 0.0001)
|
||||||
|
nn.init.constant_(m.running_mean, 0)
|
||||||
|
elif isinstance(m, nn.Linear):
|
||||||
|
nn.init.normal_(m.weight, 0, 0.01)
|
||||||
|
if m.bias is not None:
|
||||||
|
nn.init.constant_(m.bias, 0)
|
||||||
|
|
||||||
|
class _ASPPModule(nn.Module):
|
||||||
|
def __init__(self, inplanes, planes, kernel_size, padding, dilation, BatchNorm):
|
||||||
|
super(_ASPPModule, self).__init__()
|
||||||
|
self.atrous_conv = nn.Conv2d(inplanes, planes, kernel_size=kernel_size,
|
||||||
|
stride=1, padding=padding, dilation=dilation, bias=False)
|
||||||
|
self.bn = BatchNorm(planes)
|
||||||
|
self.relu = nn.ReLU()
|
||||||
|
|
||||||
|
self._init_weight()
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x = self.atrous_conv(x)
|
||||||
|
x = self.bn(x)
|
||||||
|
|
||||||
|
return self.relu(x)
|
||||||
|
|
||||||
|
def _init_weight(self):
|
||||||
|
for m in self.modules():
|
||||||
|
if isinstance(m, nn.Conv2d):
|
||||||
|
torch.nn.init.kaiming_normal_(m.weight)
|
||||||
|
|
||||||
|
class ASPP(nn.Module):
|
||||||
|
def __init__(self, backbone, output_stride, BatchNorm):
|
||||||
|
super(ASPP, self).__init__()
|
||||||
|
if backbone == 'drn':
|
||||||
|
inplanes = 512
|
||||||
|
elif backbone == 'mobilenet':
|
||||||
|
inplanes = 320
|
||||||
|
elif backbone == 'resnet_fpn':
|
||||||
|
inplanes = 256
|
||||||
|
else:
|
||||||
|
inplanes = 2048
|
||||||
|
if output_stride == 16:
|
||||||
|
dilations = [1, 6, 12, 18]
|
||||||
|
elif output_stride == 8:
|
||||||
|
dilations = [1, 12, 24, 36]
|
||||||
|
else:
|
||||||
|
raise NotImplementedError
|
||||||
|
|
||||||
|
self.aspp1 = _ASPPModule(inplanes, 256, 1, padding=0, dilation=dilations[0], BatchNorm=BatchNorm)
|
||||||
|
self.aspp2 = _ASPPModule(inplanes, 256, 3, padding=dilations[1], dilation=dilations[1], BatchNorm=BatchNorm)
|
||||||
|
self.aspp3 = _ASPPModule(inplanes, 256, 3, padding=dilations[2], dilation=dilations[2], BatchNorm=BatchNorm)
|
||||||
|
self.aspp4 = _ASPPModule(inplanes, 256, 3, padding=dilations[3], dilation=dilations[3], BatchNorm=BatchNorm)
|
||||||
|
|
||||||
|
self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
|
||||||
|
nn.Conv2d(inplanes, 256, 1, stride=1, bias=False),
|
||||||
|
BatchNorm(256),
|
||||||
|
nn.ReLU())
|
||||||
|
self.conv1 = nn.Conv2d(1280, 256, 1, bias=False)
|
||||||
|
self.bn1 = BatchNorm(256)
|
||||||
|
self.relu = nn.ReLU()
|
||||||
|
self.dropout = nn.Dropout(0.5)
|
||||||
|
self._init_weight()
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x1 = self.aspp1(x)
|
||||||
|
x2 = self.aspp2(x)
|
||||||
|
x3 = self.aspp3(x)
|
||||||
|
x4 = self.aspp4(x)
|
||||||
|
x5 = self.global_avg_pool(x)
|
||||||
|
x5 = F.interpolate(x5, size=x4.size()[2:], mode='bilinear', align_corners=True)
|
||||||
|
# x5 = F.interpolate(x5, size=x4.size()[2:], mode='nearest')
|
||||||
|
x = torch.cat((x1, x2, x3, x4, x5), dim=1)
|
||||||
|
|
||||||
|
x = self.conv1(x)
|
||||||
|
x = self.bn1(x)
|
||||||
|
x = self.relu(x)
|
||||||
|
|
||||||
|
return self.dropout(x)
|
||||||
|
|
||||||
|
def _init_weight(self):
|
||||||
|
for m in self.modules():
|
||||||
|
if isinstance(m, nn.Conv2d):
|
||||||
|
torch.nn.init.kaiming_normal_(m.weight)
|
||||||
|
|
||||||
|
class Decoder(nn.Module):
|
||||||
|
def __init__(self,num_classes, backbone, BatchNorm):
|
||||||
|
super(Decoder, self).__init__()
|
||||||
|
if backbone == 'resnet_fpn' or backbone == 'drn':
|
||||||
|
low_level_inplanes = 256
|
||||||
|
elif backbone == 'xception':
|
||||||
|
low_level_inplanes = 128
|
||||||
|
elif backbone == 'mobilenet':
|
||||||
|
low_level_inplanes = 24
|
||||||
|
else:
|
||||||
|
raise NotImplementedError
|
||||||
|
|
||||||
|
self.conv1 = nn.Conv2d(low_level_inplanes, 16, 1, bias=False)
|
||||||
|
self.bn1 = BatchNorm(16)
|
||||||
|
self.relu = nn.ReLU()
|
||||||
|
self.last_conv = nn.Sequential(nn.Conv2d(304, 256, kernel_size=3, stride=1, padding=1, bias=False),
|
||||||
|
BatchNorm(256),
|
||||||
|
nn.ReLU(),
|
||||||
|
nn.Dropout(0.5),
|
||||||
|
nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=False),
|
||||||
|
BatchNorm(256),
|
||||||
|
nn.ReLU(),
|
||||||
|
nn.Dropout(0.1),
|
||||||
|
nn.Conv2d(256, num_classes, kernel_size=1, stride=1))
|
||||||
|
|
||||||
|
self.up_to_1 = UpSampleBlock(16 + 16, 16)
|
||||||
|
self.last_layer = nn.Sequential(
|
||||||
|
ConvBNReLU(16, 16, kernel_size=1, stride=1),
|
||||||
|
nn.Conv2d(16, 1, kernel_size=1, stride=1)
|
||||||
|
)
|
||||||
|
self._init_weight()
|
||||||
|
|
||||||
|
# x(1,256,8,6) low(1,256,32,24) -》 x(1,1,32,24)
|
||||||
|
def forward(self, x, low_level_feat):
|
||||||
|
feature_T = []
|
||||||
|
# deeplab part
|
||||||
|
#(1,256,32,24) -> (1,16,64,48)
|
||||||
|
low_level_feat = self.conv1(low_level_feat)
|
||||||
|
low_level_feat = self.bn1(low_level_feat)
|
||||||
|
low_level_feat = self.relu(low_level_feat)
|
||||||
|
|
||||||
|
# x(1, 256, 8, 6)-> (1,16,32,24)
|
||||||
|
x = F.interpolate(x, size=low_level_feat.size()[2:], mode='bilinear', align_corners=True)
|
||||||
|
low_level_feat = F.interpolate(low_level_feat,
|
||||||
|
size=[low_level_feat.size()[2] * 2, low_level_feat.size()[3] * 2],
|
||||||
|
mode='bilinear', align_corners=True)
|
||||||
|
x = self.conv1(x)
|
||||||
|
x = self.bn1(x)
|
||||||
|
x = self.relu(x)
|
||||||
|
|
||||||
|
# bodyseg part
|
||||||
|
feature = F.interpolate(x, size=[x.size()[2] * 2,x.size()[3] * 2], mode='bilinear', align_corners=True)
|
||||||
|
# 输入up_to_1 x(low)(1,16,64,48),上采样2倍后的feature(1,16,64,48)
|
||||||
|
feature = self.up_to_1(low_level_feat, feature)
|
||||||
|
feature_T.append(feature)
|
||||||
|
# (1,1,64,48)
|
||||||
|
feature = self.last_layer(feature)
|
||||||
|
feature_T.append(feature)
|
||||||
|
# (1,1,128,96)
|
||||||
|
output = F.interpolate(feature, size=[128,96], mode='bilinear', align_corners=True)
|
||||||
|
# output = torch.sigmoid(output)
|
||||||
|
|
||||||
|
return output, feature_T
|
||||||
|
|
||||||
|
def _init_weight(self):
|
||||||
|
for m in self.modules():
|
||||||
|
if isinstance(m, nn.Conv2d):
|
||||||
|
torch.nn.init.kaiming_normal_(m.weight)
|
||||||
|
|
||||||
|
def build_backbone(backbone, output_stride, BatchNorm, input_channel=3):
|
||||||
|
from bodyseg.backbone.fpn import build_resnet_fpn_backbone
|
||||||
|
from bodyseg.backbone.xception import AlignedXception
|
||||||
|
|
||||||
|
return build_resnet_fpn_backbone(input_channel)
|
||||||
|
|
||||||
|
def build_aspp(backbone, output_stride, BatchNorm):
|
||||||
|
return ASPP(backbone, output_stride, BatchNorm)
|
||||||
|
|
||||||
|
def build_decoder(num_classes, backbone, BatchNorm):
|
||||||
|
return Decoder(num_classes, backbone, BatchNorm)
|
||||||
|
|
||||||
|
class DeepLab(nn.Module):
|
||||||
|
def __init__(self, input_channel=3, class_num=1):
|
||||||
|
super(DeepLab, self).__init__()
|
||||||
|
|
||||||
|
BatchNorm = get_norm("BN")
|
||||||
|
|
||||||
|
self.backbone = build_backbone("resnet_fpn", 16, BatchNorm, input_channel=input_channel)
|
||||||
|
self.aspp = build_aspp("resnet_fpn", 16, BatchNorm)
|
||||||
|
self.decoder = build_decoder(class_num, "resnet_fpn", BatchNorm)
|
||||||
|
|
||||||
|
def forward(self, input):
|
||||||
|
# input(1,3,128,96)
|
||||||
|
#output: "p2"(1,256,32,24), "p3"(1,256,16,12), "p4"(1,256,8,6)
|
||||||
|
output = self.backbone(input)
|
||||||
|
# "p4"(1,256,8,6) "p2"(1,256,32,24)
|
||||||
|
x, low_level_feat = output['p4'], output['p2']
|
||||||
|
# x(1,256,8,6)-》(1,256,8,6)
|
||||||
|
x = self.aspp(x)
|
||||||
|
# x(1,256,8,6) low(1,256,32,24) -》 x(1,1,32,24)
|
||||||
|
x, feature_T = self.decoder(x, low_level_feat)
|
||||||
|
# x(1,1,128,96)
|
||||||
|
x = F.interpolate(x, size=input.size()[2:], mode='bilinear', align_corners=True)
|
||||||
|
# x = F.interpolate(x, size=input.size()[2:], mode='nearest')
|
||||||
|
return x, feature_T
|
||||||
@@ -0,0 +1,98 @@
|
|||||||
|
import time
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
from hairstyle_model_infer import HairStyle_Model_Infer
|
||||||
|
from utils import enhance_hair
|
||||||
|
import numpy as np
|
||||||
|
from utils import landmark_processor
|
||||||
|
|
||||||
|
|
||||||
|
hairstyle_process = HairStyle_Model_Infer(gpu=True, use_enhance=True)
|
||||||
|
|
||||||
|
|
||||||
|
def resize_pre_webui(input_img, mask_img, out_h=None, out_w=None):
|
||||||
|
# 获取头发处理的局部区域图像
|
||||||
|
box_info = hairstyle_process.get_body_info(input_img)
|
||||||
|
dst_size = (576, 768)
|
||||||
|
|
||||||
|
if out_h is not None and out_w is not None:
|
||||||
|
dst_size = (out_w, out_h)
|
||||||
|
|
||||||
|
box_w, box_h = box_info[2] - box_info[0], box_info[3] - box_info[1]
|
||||||
|
scale = min(dst_size[1] / box_h, dst_size[0] / box_w)
|
||||||
|
rotate_center = [(box_info[2] + box_info[0]) * 0.5, (box_info[3] + box_info[1]) * 0.5]
|
||||||
|
M = cv2.getRotationMatrix2D(rotate_center, 0, scale)
|
||||||
|
M[:, 2] += np.float32([dst_size[0] * 0.5, dst_size[1] * 0.5]) - np.float32(rotate_center)
|
||||||
|
input_result = landmark_processor.high_quality_warpAffine(input_img, M, dst_size)
|
||||||
|
|
||||||
|
crop_matting = cv2.warpAffine(mask_img, M, dst_size)
|
||||||
|
mask_result = (crop_matting > 10).astype(np.float32)
|
||||||
|
mask_dilate = cv2.dilate(mask_result, np.ones((3, 9), np.uint8))
|
||||||
|
mask_dilate = np.clip(mask_dilate * 255, 0, 255).astype(np.uint8)
|
||||||
|
# cv2.imshow("input_result", input_result)
|
||||||
|
# cv2.imwrite("input_result.png", input_result)
|
||||||
|
# cv2.imshow("mask_dilate", mask_dilate)
|
||||||
|
# cv2.imwrite("mask_dilate.png", mask_dilate)
|
||||||
|
# cv2.waitKey(0)
|
||||||
|
|
||||||
|
return input_result, mask_dilate, M
|
||||||
|
|
||||||
|
def process_infer(user_img_path, target_color, color_dir, dst_path, ratio=1.0):
|
||||||
|
# 读取用户原始图片
|
||||||
|
origin_img = cv2.imread(user_img_path)
|
||||||
|
|
||||||
|
# resize用户原始图片，防止图片过大
|
||||||
|
user_scale = 1920 / max(origin_img.shape[0], origin_img.shape[1])
|
||||||
|
if user_scale < 1.0:
|
||||||
|
# 缩小图像应该使用INTER_AREA插值
|
||||||
|
origin_img = cv2.resize(origin_img, (0, 0), fx=user_scale, fy=user_scale, interpolation=cv2.INTER_AREA)
|
||||||
|
|
||||||
|
# 首先完成第一阶段的换发色
|
||||||
|
s1 = time.time()
|
||||||
|
hair_color_changed_img, hair_mask, status = hairstyle_process.infer_haircolor_v4(origin_img, target_hair_color=target_color, haircolor_dir=color_dir)
|
||||||
|
print("----------------------------infer_haircolor_v4", time.time()-s1)
|
||||||
|
|
||||||
|
hair_color_changed_img = (hair_color_changed_img * ratio + origin_img * (1 - ratio)).astype(np.uint8)
|
||||||
|
|
||||||
|
# 通过webui给照片加上标签
|
||||||
|
s2 = time.time()
|
||||||
|
# prompt = enhance_hair.webui_tag_by_clip(hair_color_changed_img)
|
||||||
|
prompt = ""
|
||||||
|
print("----------------------------webui_tag_by_clip", time.time() - s2)
|
||||||
|
|
||||||
|
# 处理图像尺寸
|
||||||
|
s3 = time.time()
|
||||||
|
crop_img, crop_mask, M = resize_pre_webui(hair_color_changed_img, hair_mask, 1024, 768)
|
||||||
|
print(crop_img.shape)
|
||||||
|
print(crop_mask.shape)
|
||||||
|
print("----------------------------resize_pre_webui", time.time() - s3)
|
||||||
|
|
||||||
|
# 通过webui对照片进行增强
|
||||||
|
s4 = time.time()
|
||||||
|
# cv2.imwrite('/home/student/Downloads/12121.jpg',crop_img)
|
||||||
|
enhanced_img = enhance_hair.webui_img2img(crop_img, crop_mask, prompt=prompt)
|
||||||
|
print("----------------------------webui_img2img", time.time() - s4)
|
||||||
|
|
||||||
|
# 写回原图
|
||||||
|
origin_img_final = origin_img.copy()
|
||||||
|
# restore the sd_result_small to origin_img
|
||||||
|
M_inv = cv2.invertAffineTransform(M)
|
||||||
|
cv2.warpAffine(enhanced_img, M_inv, (origin_img_final.shape[1], origin_img_final.shape[0]),
|
||||||
|
dst=origin_img_final,
|
||||||
|
borderMode=cv2.BORDER_TRANSPARENT, flags=cv2.INTER_LANCZOS4)
|
||||||
|
# cv2.imshow("origin_img_final", origin_img_final)
|
||||||
|
# cv2.waitKey(0)
|
||||||
|
|
||||||
|
cv2.imwrite(dst_path, origin_img_final)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
user_img_path = '/home/student/Downloads/base_color_v2.jpg'
|
||||||
|
# target_hair_color = [28, 55, 1] # RGB
|
||||||
|
# target_hair_color = [76, 55, 36] # RGB
|
||||||
|
# target_hair_color = [75, 1, 0] # RGB
|
||||||
|
target_hair_color = [247,235,213] # RGB
|
||||||
|
color_dir = "/home/data/hair/data/ref_color/247_235_213"
|
||||||
|
dst = "/home/data/hair/data/tmp/4.png"
|
||||||
|
|
||||||
|
process_infer(user_img_path=user_img_path, target_color=target_hair_color, color_dir=color_dir, dst_path=dst)
|
||||||
@@ -0,0 +1,140 @@
|
|||||||
|
#coding:utf-8
|
||||||
|
import traceback
|
||||||
|
from uuid import uuid4
|
||||||
|
import imghdr
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from gevent import monkey
|
||||||
|
monkey.patch_all()
|
||||||
|
|
||||||
|
import base64
|
||||||
|
import os
|
||||||
|
import random
|
||||||
|
import shutil
|
||||||
|
import time
|
||||||
|
import json
|
||||||
|
import os.path as osp
|
||||||
|
import urllib.request
|
||||||
|
import hashlib
|
||||||
|
import cv2
|
||||||
|
from datetime import datetime
|
||||||
|
import glob
|
||||||
|
import numpy as np
|
||||||
|
from core.hairstyle_model import HairStyle_Model
|
||||||
|
from hairstyle_model_infer import HairStyle_Model_Infer
|
||||||
|
from prepare_ref_hairstyle_data import prepare_single, prepare_single_color
|
||||||
|
|
||||||
|
|
||||||
|
from utils import enhance_hair
|
||||||
|
import configparser
|
||||||
|
from common.logger import config
|
||||||
|
|
||||||
|
from change_color import process_infer, resize_pre_webui
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
hairstyle_process = HairStyle_Model(gpu=True,use_enhance=True)
|
||||||
|
user_img_save_dir = config.get('default', 'userDir')
|
||||||
|
user_img_tmp_dir = config.get('default', 'tmp_dir')
|
||||||
|
user_img_res_dir = config.get('default', 'res_dir')
|
||||||
|
ref_user_dir = config.get('default', 'ref_user_dir')
|
||||||
|
train_save_dir = config.get('default', 'train_dir')
|
||||||
|
hair_template_material_dir = config.get('default', 'hair_template_material_dir')
|
||||||
|
ref_color_dir = config.get('default', 'ref_color')
|
||||||
|
ref_color_imgs_dir = config.get('default', 'ref_color_img')
|
||||||
|
train_upload_dir = config.get('default', 'upload_train_dir')
|
||||||
|
|
||||||
|
def download_img(img_url, userId=None, isfix=False, ismask=False):
|
||||||
|
try:
|
||||||
|
img_name = img_url.split("/")[-1]
|
||||||
|
tmp_dir = osp.join(user_img_tmp_dir, img_name)
|
||||||
|
# if osp.exists(tmp_dir):
|
||||||
|
# os.remove(tmp_dir)
|
||||||
|
print(img_url)
|
||||||
|
download_success = False
|
||||||
|
for i in range(3):
|
||||||
|
hairstyle_process.oss2.download_img(img_url, tmp_dir)
|
||||||
|
if osp.exists(tmp_dir) and osp.getsize(tmp_dir) > 0:
|
||||||
|
download_success = True
|
||||||
|
break
|
||||||
|
if download_success:
|
||||||
|
img_type = imghdr.what(tmp_dir)
|
||||||
|
new_tmp_dir = tmp_dir[:tmp_dir.rfind(".")+1] + img_type
|
||||||
|
shutil.move(tmp_dir, new_tmp_dir)
|
||||||
|
print("save path", new_tmp_dir)
|
||||||
|
|
||||||
|
else:
|
||||||
|
return None, None
|
||||||
|
return new_tmp_dir, None
|
||||||
|
except Exception as e:
|
||||||
|
print(e)
|
||||||
|
return None,None
|
||||||
|
|
||||||
|
|
||||||
|
def change_hair_colorv3():
|
||||||
|
datanow = datetime.now()
|
||||||
|
time_convert = datanow.strftime("%Y%m%d%H")
|
||||||
|
taskid = ''.join(str(random.choice(range(10))) for _ in range(6))
|
||||||
|
taskid = str(time_convert) + str(taskid)
|
||||||
|
start_time0 = time.time()
|
||||||
|
|
||||||
|
|
||||||
|
try:
|
||||||
|
|
||||||
|
img_url = "https://ydapp-1317132355.cos.ap-beijing.myqcloud.com/hair_mz/images/hairstyle/b8e75994-ec6c-4309-ac98-d9687ca5dfbf/ad28024a-2a98-4768-b1d3-846950f4dc57.jpg"
|
||||||
|
userId = "17519"
|
||||||
|
color = ''
|
||||||
|
rgb = [103, 103, 103]
|
||||||
|
|
||||||
|
ratio = 0.9
|
||||||
|
|
||||||
|
img_path, _ = download_img(img_url, userId)
|
||||||
|
print("---------------download img:", time.time()-start_time0)
|
||||||
|
|
||||||
|
start_time1 = time.time()
|
||||||
|
# gen material
|
||||||
|
color_name = img_path[img_path.rfind("/") + 1:]
|
||||||
|
print("color_name", color_name)
|
||||||
|
new_color_ref_img_path = os.path.join(ref_color_imgs_dir, color_name)
|
||||||
|
print("new_color_ref_img_path: ", new_color_ref_img_path)
|
||||||
|
shutil.copy(img_path, new_color_ref_img_path)
|
||||||
|
|
||||||
|
color_id = ""
|
||||||
|
for color_item in rgb:
|
||||||
|
color_id += str(color_item) + "_"
|
||||||
|
if color_id[-1] == "_":
|
||||||
|
color_id = color_id[:-1]
|
||||||
|
|
||||||
|
ref_color_save_dir = os.path.join(ref_color_dir, color_id)
|
||||||
|
if not os.path.exists(ref_color_save_dir):
|
||||||
|
os.mkdir(ref_color_save_dir)
|
||||||
|
print("!!!gen color material:", ref_color_save_dir)
|
||||||
|
prepare_single_color(new_color_ref_img_path, ref_color_save_dir)
|
||||||
|
|
||||||
|
req_id = str(uuid4())
|
||||||
|
res_path = os.path.join(user_img_res_dir, req_id + ".png")
|
||||||
|
|
||||||
|
print("---------------gen material:", time.time() - start_time1)
|
||||||
|
|
||||||
|
start_time2 = time.time()
|
||||||
|
process_infer(img_path, rgb, ref_color_save_dir, res_path, ratio)
|
||||||
|
print("res path: ", res_path)
|
||||||
|
print("---------------process_infer:", time.time() - start_time2)
|
||||||
|
|
||||||
|
start_time3 = time.time()
|
||||||
|
|
||||||
|
if os.path.exists(res_path):
|
||||||
|
ret_url = hairstyle_process.oss2.upload_file(res_path,
|
||||||
|
"hair_mz/images/hairstyle/{}/{}".format(color,
|
||||||
|
|
||||||
|
req_id + '.jpg'))
|
||||||
|
print("---------------upload img:", time.time() - start_time3)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
print(e)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
change_hair_colorv3()
|
||||||
@@ -0,0 +1,46 @@
|
|||||||
|
import requests
|
||||||
|
import json
|
||||||
|
|
||||||
|
callback_hairstyle_url = "https://puton.meidaojia.com/api/cloth/callBack"
|
||||||
|
|
||||||
|
def callback_color(color_id, cover_img, success):
|
||||||
|
url = "http://172.21.0.3:8080/ydapp/system/platform/complete_color_model"
|
||||||
|
payload = json.dumps({
|
||||||
|
"success": success,
|
||||||
|
"colorId": color_id,
|
||||||
|
"coverImg": cover_img
|
||||||
|
})
|
||||||
|
headers = {
|
||||||
|
'X-MZ-API-TOKEN': '98bcf37c942a2240ba2d907c96ed1137',
|
||||||
|
'Content-Type': 'application/json'
|
||||||
|
}
|
||||||
|
response = requests.request("POST", url, headers=headers, data=payload)
|
||||||
|
return response
|
||||||
|
# print(response.text)
|
||||||
|
|
||||||
|
|
||||||
|
def callback_hairstyle(req_id, state, message, clothId):
|
||||||
|
url = callback_hairstyle_url
|
||||||
|
payload = json.dumps({
|
||||||
|
"taskId": req_id,
|
||||||
|
"status": state,
|
||||||
|
"clothId": clothId,
|
||||||
|
"msg": message
|
||||||
|
})
|
||||||
|
|
||||||
|
status = -1
|
||||||
|
|
||||||
|
print(f"url:{url},complete_hair_model payload:{payload}")
|
||||||
|
|
||||||
|
headers = {
|
||||||
|
'X-MZ-API-TOKEN': '98bcf37c942a2240ba2d907c96ed1137',
|
||||||
|
'Content-Type': 'application/json'
|
||||||
|
}
|
||||||
|
try:
|
||||||
|
response = requests.request("POST", url, headers=headers, data=payload)
|
||||||
|
print("response:", response.text)
|
||||||
|
status = 0
|
||||||
|
except Exception as e:
|
||||||
|
print(e)
|
||||||
|
|
||||||
|
return response.text, status
|
||||||
@@ -0,0 +1,27 @@
|
|||||||
|
import pynvml
|
||||||
|
threshold = 0.9
|
||||||
|
def get_gpu(need_gpu_id):
|
||||||
|
used = get_gpu_threshold(need_gpu_id)
|
||||||
|
#小于一定的
|
||||||
|
if used > threshold:
|
||||||
|
need_use = get_use_gpu()
|
||||||
|
return need_use[0]
|
||||||
|
else:
|
||||||
|
return need_gpu_id
|
||||||
|
|
||||||
|
def get_use_gpu():
|
||||||
|
use=[]
|
||||||
|
for index in range(pynvml.nvmlDeviceGetCount()):
|
||||||
|
used = get_gpu_threshold(index)
|
||||||
|
if used > threshold:
|
||||||
|
use.append(index)
|
||||||
|
return use
|
||||||
|
|
||||||
|
def get_gpu_count():
|
||||||
|
return pynvml.nvmlDeviceGetCount()
|
||||||
|
|
||||||
|
def get_gpu_threshold(index):
|
||||||
|
handle = pynvml.nvmlDeviceGetHandleByIndex(index)
|
||||||
|
meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
|
||||||
|
used = meminfo.used / meminfo.total
|
||||||
|
return used
|
||||||
@@ -0,0 +1,63 @@
|
|||||||
|
#!/usr/bin/python
|
||||||
|
# coding=utf-8
|
||||||
|
import os
|
||||||
|
from logging.handlers import TimedRotatingFileHandler
|
||||||
|
import logging
|
||||||
|
import configparser
|
||||||
|
config = configparser.ConfigParser() # 创建对象
|
||||||
|
config.read("config/configure.ini", encoding="utf-8") # 读取配置文件，如果配置文件不存在则创建
|
||||||
|
|
||||||
|
"""
|
||||||
|
自定义日志处理
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
class LogFactory(object):
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def getLogger(log_name,log_level=None):
|
||||||
|
if log_level is None:
|
||||||
|
log_level = LogFactory.getLogLevel(getLevel())
|
||||||
|
logger = logging.getLogger(log_name)
|
||||||
|
path = getPath()
|
||||||
|
isExists=os.path.exists(path)
|
||||||
|
if not isExists:
|
||||||
|
os.makedirs(path)
|
||||||
|
log_file = os.path.join(path, '{}.log'.format(log_name))
|
||||||
|
if len(logger.handlers) <= 0:
|
||||||
|
handler = TimedRotatingFileHandler(log_file, when='D', interval=1)
|
||||||
|
else:
|
||||||
|
handler = logger.handlers[0]
|
||||||
|
formatter = logging.Formatter("%(asctime)s-%(thread)d-%(filename)s-%(levelname)s-%(message)s", "%Y-%m-%d %H:%M:%S")
|
||||||
|
handler.setFormatter(formatter)
|
||||||
|
logger.addHandler(handler)
|
||||||
|
logger.setLevel(log_level)
|
||||||
|
return logger
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def getLogLevel(log_level):
|
||||||
|
log_level = getattr(logging, log_level.upper(), None)
|
||||||
|
if log_level is None:
|
||||||
|
raise Exception("No such log level.")
|
||||||
|
return log_level
|
||||||
|
|
||||||
|
|
||||||
|
def getPath():
|
||||||
|
import sys
|
||||||
|
process_order = sys.argv[1] if len(sys.argv) > 1 else None
|
||||||
|
logpath = config.get('logger', "logpath")
|
||||||
|
logpath = logpath.rstrip("/")
|
||||||
|
|
||||||
|
if process_order is None:
|
||||||
|
return logpath
|
||||||
|
|
||||||
|
return logpath
|
||||||
|
|
||||||
|
|
||||||
|
def getLevel():
|
||||||
|
return config.get('logger', "level")
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
@@ -0,0 +1,35 @@
|
|||||||
|
[default]
|
||||||
|
modelDir = weights
|
||||||
|
hairstyleDir = /home/xsl/change_hair/project/data/ref_hairstyle
|
||||||
|
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
|
||||||
|
baseColor_ID=HDR10_443322
|
||||||
|
Port = 11023
|
||||||
|
refer_dir = /home/xsl/change_hair/project/data/ref_online
|
||||||
|
ref_user_dir = /home/xsl/change_hair/project/data/ref_user_imgs
|
||||||
|
train_dir = /home/xsl/change_hair/data/train_material
|
||||||
|
refImgDir=/home/xsl/change_hair/project/data/refImage
|
||||||
|
hair_template_material_dir=/home/xsl/change_hair/project/data/hair_template_material
|
||||||
|
ref_color=/home/xsl/change_hair/project/data/ref_color
|
||||||
|
ref_color_img=/home/xsl/change_hair/project/data/ref_color_imgs
|
||||||
|
upload_train_dir=/home/xsl/change_hair/project/data/upload_train_imgs
|
||||||
|
;version=local
|
||||||
|
version=online
|
||||||
|
|
||||||
|
[fix]
|
||||||
|
strength=1
|
||||||
|
|
||||||
|
[logger]
|
||||||
|
logpath = /home/xsl/change_hair/project/logs
|
||||||
|
level=INFO
|
||||||
|
|
||||||
|
[timelogger]
|
||||||
|
name=watch-time
|
||||||
|
|
||||||
|
[errorlogger]
|
||||||
|
level=ERROR
|
||||||
|
name=w-error
|
||||||
|
|
||||||
@@ -0,0 +1,116 @@
|
|||||||
|
import os
|
||||||
|
import shutil
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
def copy_files_by_id(id_list_file, source_dirs, output_dir):
|
||||||
|
"""
|
||||||
|
根据ID列表拷贝文件到新目录，保持目录结构
|
||||||
|
|
||||||
|
参数:
|
||||||
|
id_list_file: 包含文件ID列表的文本文件
|
||||||
|
source_dirs: 源目录列表 [hair_template_material, ref_hairstyle, train_material, upload_train_imgs]
|
||||||
|
output_dir: 输出目录
|
||||||
|
"""
|
||||||
|
# 读取ID列表
|
||||||
|
with open(id_list_file, 'r') as f:
|
||||||
|
id_list = [line.strip() for line in f.readlines() if line.strip()]
|
||||||
|
|
||||||
|
# 确保输出目录存在
|
||||||
|
os.makedirs(output_dir, exist_ok=True)
|
||||||
|
|
||||||
|
# 遍历所有源目录
|
||||||
|
for src_dir in source_dirs:
|
||||||
|
if not os.path.exists(src_dir):
|
||||||
|
print(f"警告: 源目录不存在 {src_dir}")
|
||||||
|
continue
|
||||||
|
|
||||||
|
# 遍历源目录下的所有文件
|
||||||
|
for root, _, files in os.walk(src_dir):
|
||||||
|
for file in files:
|
||||||
|
# 检查文件名是否包含ID列表中的任一ID
|
||||||
|
if any(file_id in file for file_id in id_list):
|
||||||
|
src_path = os.path.join(root, file)
|
||||||
|
|
||||||
|
# 计算相对路径
|
||||||
|
rel_path = os.path.relpath(root, src_dir)
|
||||||
|
dst_dir = os.path.join(output_dir, os.path.basename(src_dir), rel_path)
|
||||||
|
|
||||||
|
# 创建目标目录并拷贝文件
|
||||||
|
os.makedirs(dst_dir, exist_ok=True)
|
||||||
|
dst_path = os.path.join(dst_dir, file)
|
||||||
|
shutil.copy2(src_path, dst_path)
|
||||||
|
print(f"已拷贝: {src_path} -> {dst_path}")
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
# import argparse
|
||||||
|
|
||||||
|
# parser = argparse.ArgumentParser(description='根据ID列表拷贝文件')
|
||||||
|
# parser.add_argument('id_list', help='包含文件ID列表的文本文件路径')
|
||||||
|
# parser.add_argument('output_dir', help='输出目录路径')
|
||||||
|
# parser.add_argument('--source_dirs', nargs='+',
|
||||||
|
# default=['hair_template_material', 'ref_hairstyle', 'train_material', 'upload_train_imgs'],
|
||||||
|
# help='源目录列表')
|
||||||
|
|
||||||
|
# args = parser.parse_args()
|
||||||
|
id_list = [1939351512506802177,
|
||||||
|
1939351392453238785,
|
||||||
|
1939351319849836545,
|
||||||
|
1939351238526476289,
|
||||||
|
1939351219396255745,
|
||||||
|
1939339435616608258,
|
||||||
|
1939330299344560130,
|
||||||
|
1938848782956732417,
|
||||||
|
1938848730150445058,
|
||||||
|
1938848709308948482,
|
||||||
|
1938841049238970370,
|
||||||
|
1938492830059438081,
|
||||||
|
1938492790033195009,
|
||||||
|
1938490058547240961,
|
||||||
|
1938489994072399874,
|
||||||
|
1938489915521474561,
|
||||||
|
1938489840707674114,
|
||||||
|
1938489739088076801,
|
||||||
|
1938489674713899010,
|
||||||
|
1938489620162781185]
|
||||||
|
output_dir = '/home/data/hair/data/hairStyle_addin/06300845'
|
||||||
|
hair_template_material = '/home/data/hair/data/hair_template_material'
|
||||||
|
ref_hairstyle = '/home/data/hair/data/ref_hairstyle'
|
||||||
|
train_material = '/mnt/database2/online-server/hair-online-tj-v2/train_material'
|
||||||
|
upload_train_imgs = '/home/data/hair/data/upload_train_imgs'
|
||||||
|
copy_files_by_id(id_list, source_dirs, args.output_dir)
|
||||||
|
|
||||||
|
def copy_files_by_id(id_list, hair_template_material, ref_hairstyle, train_material, upload_train_imgs, output_dir):
|
||||||
|
"""
|
||||||
|
根据ID列表从指定目录拷贝文件到输出目录
|
||||||
|
:param id_list: 文件ID列表(字符串格式)
|
||||||
|
:param hair_template_material: 发型模板材质目录
|
||||||
|
:param ref_hairstyle: 参考发型目录
|
||||||
|
:param train_material: 训练材质目录
|
||||||
|
:param upload_train_imgs: 上传训练图片目录
|
||||||
|
:param output_dir: 输出目录
|
||||||
|
"""
|
||||||
|
# 将ID转换为字符串格式
|
||||||
|
id_list = [str(id) for id in id_list]
|
||||||
|
# 创建源目录列表
|
||||||
|
source_dirs = [hair_template_material, ref_hairstyle, train_material, upload_train_imgs]
|
||||||
|
for src_dir in source_dirs:
|
||||||
|
if not os.path.exists(src_dir):
|
||||||
|
print(f"警告: 源目录不存在 {src_dir}")
|
||||||
|
continue
|
||||||
|
|
||||||
|
# 遍历源目录下的所有文件
|
||||||
|
for root, _, files in os.walk(src_dir):
|
||||||
|
for file in files:
|
||||||
|
# 检查文件名是否包含ID列表中的任一ID
|
||||||
|
if any(file_id in file for file_id in id_list):
|
||||||
|
src_path = os.path.join(root, file)
|
||||||
|
|
||||||
|
# 计算相对路径
|
||||||
|
rel_path = os.path.relpath(root, src_dir)
|
||||||
|
dst_dir = os.path.join(output_dir, os.path.basename(src_dir), rel_path)
|
||||||
|
|
||||||
|
# 创建目标目录并拷贝文件
|
||||||
|
os.makedirs(dst_dir, exist_ok=True)
|
||||||
|
dst_path = os.path.join(dst_dir, file)
|
||||||
|
shutil.copy2(src_path, dst_path)
|
||||||
|
print(f"已拷贝: {src_path} -> {dst_path}")
|
||||||
@@ -0,0 +1,196 @@
|
|||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import math
|
||||||
|
import os
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
from core.utils import landmark_processor
|
||||||
|
import glob
|
||||||
|
# from models.BigResNetStable import MomocvFaceAlignment
|
||||||
|
|
||||||
|
|
||||||
|
def op_name(op_name, m):
|
||||||
|
m.op_name = op_name
|
||||||
|
return m
|
||||||
|
|
||||||
|
class Flatten(nn.Module):
|
||||||
|
def __init__(self, axis=1):
|
||||||
|
super(Flatten, self).__init__()
|
||||||
|
self.axis = axis
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
assert self.axis == 1
|
||||||
|
x = x.reshape(x.shape[0], -1)
|
||||||
|
return x
|
||||||
|
|
||||||
|
def flatten(name, axis=1):
|
||||||
|
return op_name(name, Flatten(axis))
|
||||||
|
|
||||||
|
def conv_relu(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1):
|
||||||
|
return nn.Sequential(
|
||||||
|
op_name(name, nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, True)),
|
||||||
|
op_name(name.replace('conv', 'relu'), nn.LeakyReLU(inplace=False, negative_slope=5e-11)),
|
||||||
|
)
|
||||||
|
|
||||||
|
def conv(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1):
|
||||||
|
return op_name(name, nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, True))
|
||||||
|
|
||||||
|
def bn_conv_relu(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, momentum = 0.9, track_running_stats=True):
|
||||||
|
return nn.Sequential(
|
||||||
|
op_name(name + '_bn1', nn.BatchNorm2d(in_channels, momentum=momentum, eps=2e-5, track_running_stats=track_running_stats)),
|
||||||
|
op_name(name + '_conv1', nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, True)),
|
||||||
|
op_name(name + '_relu', nn.LeakyReLU(inplace=False, negative_slope=5e-11)),
|
||||||
|
)
|
||||||
|
|
||||||
|
def bn(name, in_channels, momentum = 0.9, track_running_stats=True):
|
||||||
|
return op_name(name, nn.BatchNorm2d(in_channels, momentum=momentum, eps=2e-5, track_running_stats=track_running_stats))
|
||||||
|
|
||||||
|
class BasicBlock(nn.Module):
|
||||||
|
def __init__(self, stage, unit, inplanes, outplanes, stride):
|
||||||
|
super(BasicBlock, self).__init__()
|
||||||
|
|
||||||
|
self.bn = bn('stage{}_unit{}_bn1'.format(stage, unit), inplanes)
|
||||||
|
self.conv1 = conv_relu('stage{}_unit{}_conv1'.format(stage, unit), inplanes, outplanes, kernel_size=3, stride=1, padding=1)
|
||||||
|
self.conv2 = conv('stage{}_unit{}_conv2'.format(stage, unit), outplanes, outplanes, kernel_size=3, stride=stride, padding=1)
|
||||||
|
|
||||||
|
if inplanes != outplanes or stride != 1:
|
||||||
|
self.conv_sc = conv('stage{}_unit{}_conv1sc'.format(stage, unit), inplanes, outplanes, kernel_size=1, stride=stride, padding=0)
|
||||||
|
|
||||||
|
self.inplanes = inplanes
|
||||||
|
self.outplanes = outplanes
|
||||||
|
self.stride = stride
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
residule = x
|
||||||
|
out = self.bn(x)
|
||||||
|
out = self.conv1(out)
|
||||||
|
out = self.conv2(out)
|
||||||
|
if self.inplanes != self.outplanes or self.stride != 1:
|
||||||
|
residule = self.conv_sc(residule)
|
||||||
|
ret = out + residule
|
||||||
|
return ret
|
||||||
|
|
||||||
|
class ResnetBlock(nn.Module):
|
||||||
|
def __init__(self, stage, inplanes, outplanes, stride=2, n_blocks=1):
|
||||||
|
super(ResnetBlock, self).__init__()
|
||||||
|
|
||||||
|
self.conv = []
|
||||||
|
for m in range(n_blocks):
|
||||||
|
if m == 0:
|
||||||
|
self.conv.append(BasicBlock(stage, m + 1, inplanes, outplanes, stride))
|
||||||
|
else:
|
||||||
|
self.conv.append(BasicBlock(stage, m + 1, outplanes, outplanes, 1))
|
||||||
|
self.conv = nn.Sequential(*self.conv)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
ret = self.conv(x)
|
||||||
|
return ret
|
||||||
|
|
||||||
|
class FaceRecognitionServer(nn.Module):
|
||||||
|
def __init__(self):
|
||||||
|
super(FaceRecognitionServer, self).__init__()
|
||||||
|
ch_num = [64, 64, 128, 256, 512]
|
||||||
|
# ch_num = [64, 64]
|
||||||
|
strides = [2, 2, 2, 2]
|
||||||
|
n_blocks = [3, 4, 14, 3]
|
||||||
|
op_list = [conv_relu('conv0', 3, 64, 3, 1, 1)]
|
||||||
|
op_list += [ResnetBlock(i + 1, ch_num[i], ch_num[i + 1], strides[i], n_blocks[i]) for i in range(len(ch_num) - 1)]
|
||||||
|
op_list += [bn('bn1', 512)]
|
||||||
|
op_list += [flatten('flatten', 1)]
|
||||||
|
op_list += [op_name('pre_fc1', nn.Linear(25088, 512))]
|
||||||
|
self.features = nn.Sequential(*op_list)
|
||||||
|
|
||||||
|
model_path = os.path.dirname(os.path.split(os.path.realpath(__file__))[0])
|
||||||
|
weights = torch.load('weights/MMCVFaceRecognitionServer.pth',
|
||||||
|
map_location=lambda storage, loc: storage)
|
||||||
|
self.load_state_dict(weights)
|
||||||
|
self.eval()
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
features = self.features(x)
|
||||||
|
return features
|
||||||
|
|
||||||
|
class MomocvFaceRecognitionServer(object):
|
||||||
|
def __init__(self, gpu_id=None):
|
||||||
|
self.gpu_id = gpu_id
|
||||||
|
self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||||
|
self.face_recognition_net = FaceRecognitionServer()
|
||||||
|
# self.mmcv = MomocvFaceAlignment()
|
||||||
|
# self.model_path, _ = os.path.split(os.path.realpath(__file__))
|
||||||
|
# weights = torch.load(os.path.join(self.model_path, 'MMCVFaceRecognitionServer.pth'),
|
||||||
|
# map_location=lambda storage, loc: storage)
|
||||||
|
# self.face_recognition_net.load_state_dict(weights)
|
||||||
|
self.face_recognition_net.to(self.device)
|
||||||
|
# self.face_recognition_net.eval()
|
||||||
|
|
||||||
|
def forward(self, images, landmarks):
|
||||||
|
dst_size = 112
|
||||||
|
|
||||||
|
with torch.no_grad():
|
||||||
|
input_numpy = np.zeros((len(images), 3, dst_size, dst_size), dtype=np.float32)
|
||||||
|
assert len(images) == len(landmarks)
|
||||||
|
for ix, img in enumerate(images):
|
||||||
|
landmark = landmarks[ix]
|
||||||
|
|
||||||
|
|
||||||
|
# get angel
|
||||||
|
# M = landmark_processor.get_transform_mat_full_face(landmark, 576, scale=1, offset=(0, 0.3))
|
||||||
|
# pt1k_crop = landmark_processor.transform_points(landmark, M)
|
||||||
|
# crop_face = cv2.warpAffine(img, M, (576, 768),
|
||||||
|
# flags=cv2.INTER_LANCZOS4, borderMode=cv2.BORDER_REFLECT)
|
||||||
|
# t0 = time.time()
|
||||||
|
# landmarks87, poselayer, tracking_probe, occlusion_probe = self.mmcv.detect(crop_face, [pt1k_crop])
|
||||||
|
# # print('time:', time.time() - t0)
|
||||||
|
# if tracking_probe[0] < 0.5:
|
||||||
|
# continue
|
||||||
|
|
||||||
|
mat = landmark_processor.get_transform_mat_for_face_recognition(landmark, dst_size)[0:2]
|
||||||
|
|
||||||
|
tmp = cv2.warpAffine(img, mat, (dst_size, dst_size))
|
||||||
|
|
||||||
|
input_numpy[ix, :, :, :] = (tmp.transpose((2, 0, 1)).astype(np.float32) - 127.5) / 128
|
||||||
|
|
||||||
|
# cv2.imshow('tmp', tmp)
|
||||||
|
# cv2.waitKey()
|
||||||
|
|
||||||
|
in_tensor = torch.from_numpy(input_numpy)
|
||||||
|
in_tensor = in_tensor.to(self.device)
|
||||||
|
features = self.face_recognition_net(in_tensor)
|
||||||
|
features = features.cpu().numpy()
|
||||||
|
norm_factor = np.sqrt(np.sum(features ** 2, axis=1))
|
||||||
|
norm_factor = norm_factor.reshape(-1, 1)
|
||||||
|
features /= norm_factor
|
||||||
|
return features
|
||||||
|
|
||||||
|
def cos_sim(self, a, b):
|
||||||
|
a_norm = np.linalg.norm(a)
|
||||||
|
b_norm = np.linalg.norm(b)
|
||||||
|
cos = np.dot(a, b) / (a_norm * b_norm)
|
||||||
|
return cos
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
all_jpegs = glob.glob(r'D:\data\deepface_example\data\02b59e75ce91bda300ff827a85a74687d633cdfb5d830e7d3855e31b75201fa8\*.jpg')
|
||||||
|
for s_filename_path in all_jpegs:
|
||||||
|
img = cv2.imread(s_filename_path)
|
||||||
|
|
||||||
|
# cv2.imshow('img', img)
|
||||||
|
# cv2.waitKey()
|
||||||
|
|
||||||
|
# dflpng = DATAIMG(str(s_filename_path), print_on_no_embedded_data=True)
|
||||||
|
# if dflpng is None:
|
||||||
|
# print('ERROR')
|
||||||
|
#
|
||||||
|
# landmarks = dflpng.get_landmarks_mmcv_137()
|
||||||
|
#
|
||||||
|
# mmcv = MomocvFaceRecognitionServer()
|
||||||
|
# features = mmcv.forward([img, img], [landmarks, landmarks])
|
||||||
|
# print(features)
|
||||||
|
# print('conansherry')
|
||||||
|
# fullyconnected1 = fullyconnected1[0]
|
||||||
|
# fullyconnected1 = (np.reshape(fullyconnected1, (2, 87)).transpose((1, 0))).astype(np.int32)
|
||||||
|
# occlusion_probe = occlusion_probe[0]
|
||||||
|
# for ix, pt in enumerate(fullyconnected1):
|
||||||
|
# cv2.circle(img, (int(pt[0]), int(pt[1])), 1, (0, 255, 0) if occlusion_probe[ix] > 0.1 else (0, 0, 255), 2)
|
||||||
|
# # cv2.putText(img, str(ix), (int(pt[0]), int(pt[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 255, 0), 1)
|
||||||
|
# cv2.imshow('img', img)
|
||||||
|
# cv2.waitKey()
|
||||||
Binary file not shown.
@@ -0,0 +1,75 @@
|
|||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
from torch import nn
|
||||||
|
|
||||||
|
def get_norm(norm, out_channels=None):
|
||||||
|
"""
|
||||||
|
Args:
|
||||||
|
norm (str or callable):
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
nn.Module or None: the normalization layer
|
||||||
|
"""
|
||||||
|
if isinstance(norm, str):
|
||||||
|
if len(norm) == 0:
|
||||||
|
return None
|
||||||
|
norm = {
|
||||||
|
"BN": nn.BatchNorm2d,
|
||||||
|
"IN": nn.InstanceNorm2d,
|
||||||
|
"GN": lambda channels: nn.GroupNorm(32, channels),
|
||||||
|
"nnSyncBN": nn.SyncBatchNorm, # keep for debugging
|
||||||
|
}[norm]
|
||||||
|
if out_channels is not None:
|
||||||
|
return norm(out_channels)
|
||||||
|
else:
|
||||||
|
return norm
|
||||||
|
|
||||||
|
class Conv2d(torch.nn.Conv2d):
|
||||||
|
"""
|
||||||
|
A wrapper around :class:`torch.nn.Conv2d` to support zero-size tensor and more features.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, *args, **kwargs):
|
||||||
|
"""
|
||||||
|
Extra keyword arguments supported in addition to those in `torch.nn.Conv2d`:
|
||||||
|
|
||||||
|
Args:
|
||||||
|
norm (nn.Module, optional): a normalization layer
|
||||||
|
activation (callable(Tensor) -> Tensor): a callable activation function
|
||||||
|
|
||||||
|
It assumes that norm layer is used before activation.
|
||||||
|
"""
|
||||||
|
norm = kwargs.pop("norm", None)
|
||||||
|
activation = kwargs.pop("activation", None)
|
||||||
|
super().__init__(*args, **kwargs)
|
||||||
|
|
||||||
|
self.norm = norm
|
||||||
|
self.activation = activation
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x = super().forward(x)
|
||||||
|
if self.norm is not None:
|
||||||
|
x = self.norm(x)
|
||||||
|
if self.activation is not None:
|
||||||
|
x = self.activation(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
class Backbone(nn.Module):
|
||||||
|
"""
|
||||||
|
Abstract base class for network backbones.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
"""
|
||||||
|
The `__init__` method of any subclass can specify its own set of arguments.
|
||||||
|
"""
|
||||||
|
super().__init__()
|
||||||
|
|
||||||
|
def forward(self):
|
||||||
|
"""
|
||||||
|
Subclasses must override this method, but adhere to the same return type.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict[str: Tensor]: mapping from feature name (e.g., "res2") to tensor
|
||||||
|
"""
|
||||||
|
pass
|
||||||
@@ -0,0 +1,268 @@
|
|||||||
|
import math
|
||||||
|
import core.utils.weight_init as weight_init
|
||||||
|
import torch
|
||||||
|
import torch.nn.functional as F
|
||||||
|
from torch import nn
|
||||||
|
|
||||||
|
from core.bodyseg.backbone.backbone import Backbone, get_norm, Conv2d
|
||||||
|
from core.bodyseg.backbone.resnet import build_resnet_backbone
|
||||||
|
|
||||||
|
class FPN(Backbone):
|
||||||
|
"""
|
||||||
|
This module implements Feature Pyramid Network.
|
||||||
|
It creates pyramid features built on top of some input feature maps.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self, bottom_up, in_features, out_channels, norm="", top_block=None, fuse_type="sum"
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Args:
|
||||||
|
bottom_up (Backbone): module representing the bottom up subnetwork.
|
||||||
|
Must be a subclass of :class:`Backbone`. The multi-scale feature
|
||||||
|
maps generated by the bottom up network, and listed in `in_features`,
|
||||||
|
are used to generate FPN levels.
|
||||||
|
in_features (list[str]): names of the input feature maps coming
|
||||||
|
from the backbone to which FPN is attached. For example, if the
|
||||||
|
backbone produces ["res2", "res3", "res4"], any *contiguous* sublist
|
||||||
|
of these may be used; order must be from high to low resolution.
|
||||||
|
out_channels (int): number of channels in the output feature maps.
|
||||||
|
norm (str): the normalization to use.
|
||||||
|
top_block (nn.Module or None): if provided, an extra operation will
|
||||||
|
be performed on the output of the last (smallest resolution)
|
||||||
|
FPN output, and the result will extend the result list. The top_block
|
||||||
|
further downsamples the feature map. It must have an attribute
|
||||||
|
"num_levels", meaning the number of extra FPN levels added by
|
||||||
|
this block, and "in_feature", which is a string representing
|
||||||
|
its input feature (e.g., p5).
|
||||||
|
fuse_type (str): types for fusing the top down features and the lateral
|
||||||
|
ones. It can be "sum" (default), which sums up element-wise; or "avg",
|
||||||
|
which takes the element-wise mean of the two.
|
||||||
|
"""
|
||||||
|
super(FPN, self).__init__()
|
||||||
|
assert isinstance(bottom_up, Backbone)
|
||||||
|
|
||||||
|
# Feature map strides and channels from the bottom up network (e.g. ResNet)
|
||||||
|
in_strides = [bottom_up._out_feature_strides[f] for f in in_features]
|
||||||
|
in_channels = [bottom_up._out_feature_channels[f] for f in in_features]
|
||||||
|
|
||||||
|
_assert_strides_are_log2_contiguous(in_strides)
|
||||||
|
lateral_convs = []
|
||||||
|
output_convs = []
|
||||||
|
|
||||||
|
use_bias = norm == ""
|
||||||
|
for idx, in_channels in enumerate(in_channels):
|
||||||
|
lateral_norm = get_norm(norm, out_channels)
|
||||||
|
output_norm = get_norm(norm, out_channels)
|
||||||
|
|
||||||
|
lateral_conv = Conv2d(
|
||||||
|
in_channels, out_channels, kernel_size=1, bias=use_bias, norm=lateral_norm
|
||||||
|
)
|
||||||
|
output_conv = Conv2d(
|
||||||
|
out_channels,
|
||||||
|
out_channels,
|
||||||
|
kernel_size=3,
|
||||||
|
stride=1,
|
||||||
|
padding=1,
|
||||||
|
bias=use_bias,
|
||||||
|
norm=output_norm,
|
||||||
|
)
|
||||||
|
weight_init.c2_xavier_fill(lateral_conv)
|
||||||
|
weight_init.c2_xavier_fill(output_conv)
|
||||||
|
stage = int(math.log2(in_strides[idx]))
|
||||||
|
|
||||||
|
lateral_convs.append(lateral_conv)
|
||||||
|
output_convs.append(output_conv)
|
||||||
|
# Place convs into top-down order (from low to high resolution)
|
||||||
|
# to make the top-down computation in forward clearer.
|
||||||
|
self.lateral_convs = nn.ModuleList(lateral_convs[::-1])
|
||||||
|
self.output_convs = nn.ModuleList(output_convs[::-1])
|
||||||
|
self.top_block = top_block
|
||||||
|
self.in_features = in_features
|
||||||
|
self.bottom_up = bottom_up
|
||||||
|
# Return feature names are "p<stage>", like ["p2", "p3", ..., "p6"]
|
||||||
|
self._out_feature_strides = {"p{}".format(int(math.log2(s))): s for s in in_strides}
|
||||||
|
# top block output feature maps.
|
||||||
|
if self.top_block is not None:
|
||||||
|
for s in range(stage, stage + self.top_block.num_levels):
|
||||||
|
self._out_feature_strides["p{}".format(s + 1)] = 2 ** (s + 1)
|
||||||
|
|
||||||
|
self._out_features = list(self._out_feature_strides.keys())
|
||||||
|
self._out_feature_channels = {k: out_channels for k in self._out_features}
|
||||||
|
assert fuse_type in {"avg", "sum"}
|
||||||
|
self._fuse_type = fuse_type
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
"""
|
||||||
|
Args:
|
||||||
|
input (dict[str: Tensor]): mapping feature map name (e.g., "res5") to
|
||||||
|
feature map tensor for each feature level in high to low resolution order.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict[str: Tensor]:
|
||||||
|
mapping from feature map name to FPN feature map tensor
|
||||||
|
in high to low resolution order. Returned feature names follow the FPN
|
||||||
|
paper convention: "p<stage>", where stage has stride = 2 ** stage e.g.,
|
||||||
|
["p2", "p3", ..., "p6"].
|
||||||
|
"""
|
||||||
|
# Reverse feature maps into top-down order (from low to high resolution)
|
||||||
|
bottom_up_features = self.bottom_up(x)
|
||||||
|
x = [bottom_up_features[f] for f in self.in_features[::-1]]
|
||||||
|
results = []
|
||||||
|
prev_features = self.lateral_convs[0](x[0])
|
||||||
|
results.append(self.output_convs[0](prev_features))
|
||||||
|
for features, lateral_conv, output_conv in zip(
|
||||||
|
x[1:], self.lateral_convs[1:], self.output_convs[1:]
|
||||||
|
):
|
||||||
|
top_down_features = F.interpolate(prev_features, scale_factor=2, mode="nearest")
|
||||||
|
lateral_features = lateral_conv(features)
|
||||||
|
prev_features = lateral_features + top_down_features
|
||||||
|
if self._fuse_type == "avg":
|
||||||
|
prev_features /= 2
|
||||||
|
results.insert(0, output_conv(prev_features))
|
||||||
|
|
||||||
|
if self.top_block is not None:
|
||||||
|
top_block_in_feature = bottom_up_features.get(self.top_block.in_feature, None)
|
||||||
|
if top_block_in_feature is None:
|
||||||
|
top_block_in_feature = results[self._out_features.index(self.top_block.in_feature)]
|
||||||
|
results.extend(self.top_block(top_block_in_feature))
|
||||||
|
assert len(self._out_features) == len(results)
|
||||||
|
return dict(zip(self._out_features, results))
|
||||||
|
|
||||||
|
def _assert_strides_are_log2_contiguous(strides):
|
||||||
|
"""
|
||||||
|
Assert that each stride is 2x times its preceding stride, i.e. "contiguous in log2".
|
||||||
|
"""
|
||||||
|
for i, stride in enumerate(strides[1:], 1):
|
||||||
|
assert stride == 2 * strides[i - 1], "Strides {} {} are not log2 contiguous".format(
|
||||||
|
stride, strides[i - 1]
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class LastLevelMaxPool(nn.Module):
|
||||||
|
"""
|
||||||
|
This module is used in the original FPN to generate a downsampled
|
||||||
|
P6 feature from P5.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__()
|
||||||
|
self.num_levels = 1
|
||||||
|
self.in_feature = "p5"
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
return [F.max_pool2d(x, kernel_size=1, stride=2, padding=0)]
|
||||||
|
|
||||||
|
|
||||||
|
class LastLevelP6P7(nn.Module):
|
||||||
|
"""
|
||||||
|
This module is used in RetinaNet to generate extra layers, P6 and P7 from
|
||||||
|
C5 feature.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, in_channels, out_channels):
|
||||||
|
super().__init__()
|
||||||
|
self.num_levels = 2
|
||||||
|
self.in_feature = "res5"
|
||||||
|
self.p6 = nn.Conv2d(in_channels, out_channels, 3, 2, 1)
|
||||||
|
self.p7 = nn.Conv2d(out_channels, out_channels, 3, 2, 1)
|
||||||
|
for module in [self.p6, self.p7]:
|
||||||
|
weight_init.c2_xavier_fill(module)
|
||||||
|
|
||||||
|
def forward(self, c5):
|
||||||
|
p6 = self.p6(c5)
|
||||||
|
p7 = self.p7(F.relu(p6))
|
||||||
|
return [p6, p7]
|
||||||
|
|
||||||
|
|
||||||
|
def build_resnet_fpn_backbone(in_channels=3):
|
||||||
|
"""
|
||||||
|
Args:
|
||||||
|
cfg: a detectron2 CfgNode
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
backbone (Backbone): backbone module, must be a subclass of :class:`Backbone`.
|
||||||
|
"""
|
||||||
|
bottom_up = build_resnet_backbone(in_channels)
|
||||||
|
in_features = ["res2", "res3", "res4"]
|
||||||
|
out_channels = 256
|
||||||
|
backbone = FPN(
|
||||||
|
bottom_up=bottom_up,
|
||||||
|
in_features=in_features,
|
||||||
|
out_channels=out_channels,
|
||||||
|
norm="BN",
|
||||||
|
# top_block=LastLevelMaxPool(),
|
||||||
|
top_block=None,
|
||||||
|
fuse_type="sum",
|
||||||
|
)
|
||||||
|
return backbone
|
||||||
|
|
||||||
|
def build_retinanet_resnet_fpn_backbone(cfg, in_channels=3):
|
||||||
|
"""
|
||||||
|
Args:
|
||||||
|
cfg: a detectron2 CfgNode
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
backbone (Backbone): backbone module, must be a subclass of :class:`Backbone`.
|
||||||
|
"""
|
||||||
|
bottom_up = build_resnet_backbone(cfg, in_channels)
|
||||||
|
in_features = cfg.MODEL.FPN.IN_FEATURES
|
||||||
|
out_channels = cfg.MODEL.FPN.OUT_CHANNELS
|
||||||
|
in_channels_p6p7 = bottom_up._out_feature_channels["res5"]
|
||||||
|
backbone = FPN(
|
||||||
|
bottom_up=bottom_up,
|
||||||
|
in_features=in_features,
|
||||||
|
out_channels=out_channels,
|
||||||
|
norm=cfg.MODEL.FPN.NORM,
|
||||||
|
top_block=LastLevelP6P7(in_channels_p6p7, out_channels),
|
||||||
|
fuse_type=cfg.MODEL.FPN.FUSE_TYPE,
|
||||||
|
)
|
||||||
|
return backbone
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
import argparse
|
||||||
|
from config.default import get_cfg
|
||||||
|
|
||||||
|
def setup(args):
|
||||||
|
"""
|
||||||
|
Create configs and perform basic setups.
|
||||||
|
"""
|
||||||
|
cfg = get_cfg()
|
||||||
|
cfg.merge_from_file(args.cfg)
|
||||||
|
cfg.merge_from_list(args.opts)
|
||||||
|
cfg.freeze()
|
||||||
|
return cfg
|
||||||
|
|
||||||
|
parser = argparse.ArgumentParser(description='Train ImageNet network')
|
||||||
|
# general
|
||||||
|
parser.add_argument('--cfg',
|
||||||
|
help='experiment configure file name',
|
||||||
|
required=True,
|
||||||
|
type=str)
|
||||||
|
|
||||||
|
parser.add_argument('opts',
|
||||||
|
help="Modify config options using the command-line",
|
||||||
|
default=None,
|
||||||
|
nargs=argparse.REMAINDER)
|
||||||
|
|
||||||
|
args = parser.parse_args()
|
||||||
|
cfg = setup(args)
|
||||||
|
print(cfg)
|
||||||
|
|
||||||
|
model = build_resnet_fpn_backbone(cfg, 3)
|
||||||
|
# model = build_retinanet_resnet_fpn_backbone(cfg, 3)
|
||||||
|
print(model)
|
||||||
|
# model = torch.nn.DataParallel(model, list(range(2))).cuda()
|
||||||
|
dummy_input = torch.randn(4, 3, 512, 512)
|
||||||
|
|
||||||
|
out = model(dummy_input)
|
||||||
|
|
||||||
|
for k, v in out.items():
|
||||||
|
print(k, v.shape)
|
||||||
|
|
||||||
|
# torch.onnx.export(model, dummy_input, "tmp.onnx", verbose=True,
|
||||||
|
# input_names=['input'],
|
||||||
|
# output_names=['output'])
|
||||||
|
|
||||||
|
pass
|
||||||
|
|
||||||
@@ -0,0 +1,298 @@
|
|||||||
|
import numpy as np
|
||||||
|
import core.utils.weight_init as weight_init
|
||||||
|
import torch
|
||||||
|
import torch.nn.functional as F
|
||||||
|
from torch import nn
|
||||||
|
from core.bodyseg.backbone.backbone import Backbone, get_norm, Conv2d
|
||||||
|
|
||||||
|
class BasicStem(nn.Module):
|
||||||
|
def __init__(self, in_channels=3, out_channels=64, norm="BN"):
|
||||||
|
"""
|
||||||
|
Args:
|
||||||
|
norm (str or callable): a callable that takes the number of
|
||||||
|
channels and return a `nn.Module`, or a pre-defined string
|
||||||
|
(one of {"FrozenBN", "BN", "GN"}).
|
||||||
|
"""
|
||||||
|
super().__init__()
|
||||||
|
self.conv1 = Conv2d(
|
||||||
|
in_channels,
|
||||||
|
out_channels,
|
||||||
|
kernel_size=7,
|
||||||
|
stride=2,
|
||||||
|
padding=3,
|
||||||
|
bias=False,
|
||||||
|
norm=get_norm(norm, out_channels),
|
||||||
|
)
|
||||||
|
weight_init.c2_msra_fill(self.conv1)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x = self.conv1(x)
|
||||||
|
x = F.relu_(x)
|
||||||
|
x = F.max_pool2d(x, kernel_size=3, stride=2, padding=1)
|
||||||
|
return x
|
||||||
|
|
||||||
|
@property
|
||||||
|
def out_channels(self):
|
||||||
|
return self.conv1.out_channels
|
||||||
|
|
||||||
|
@property
|
||||||
|
def stride(self):
|
||||||
|
return 4 # = stride 2 conv -> stride 2 max pool
|
||||||
|
|
||||||
|
class ResNetBlockBase(nn.Module):
|
||||||
|
def __init__(self, in_channels, out_channels, stride):
|
||||||
|
"""
|
||||||
|
The `__init__` method of any subclass should also contain these arguments.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
in_channels (int):
|
||||||
|
out_channels (int):
|
||||||
|
stride (int):
|
||||||
|
"""
|
||||||
|
super().__init__()
|
||||||
|
self.in_channels = in_channels
|
||||||
|
self.out_channels = out_channels
|
||||||
|
self.stride = stride
|
||||||
|
|
||||||
|
class BottleneckBlock(ResNetBlockBase):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
in_channels,
|
||||||
|
out_channels,
|
||||||
|
*,
|
||||||
|
bottleneck_channels,
|
||||||
|
stride=1,
|
||||||
|
num_groups=1,
|
||||||
|
norm="BN",
|
||||||
|
stride_in_1x1=False,
|
||||||
|
dilation=1,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Args:
|
||||||
|
norm (str or callable): a callable that takes the number of
|
||||||
|
channels and return a `nn.Module`, or a pre-defined string
|
||||||
|
(one of {"FrozenBN", "BN", "GN"}).
|
||||||
|
stride_in_1x1 (bool): when stride==2, whether to put stride in the
|
||||||
|
first 1x1 convolution or the bottleneck 3x3 convolution.
|
||||||
|
"""
|
||||||
|
super().__init__(in_channels, out_channels, stride)
|
||||||
|
|
||||||
|
if in_channels != out_channels:
|
||||||
|
self.shortcut = Conv2d(
|
||||||
|
in_channels,
|
||||||
|
out_channels,
|
||||||
|
kernel_size=1,
|
||||||
|
stride=stride,
|
||||||
|
bias=False,
|
||||||
|
norm=get_norm(norm, out_channels),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
self.shortcut = None
|
||||||
|
|
||||||
|
# The original MSRA ResNet models have stride in the first 1x1 conv
|
||||||
|
# The subsequent fb.torch.resnet and Caffe2 ResNe[X]t implementations have
|
||||||
|
# stride in the 3x3 conv
|
||||||
|
stride_1x1, stride_3x3 = (stride, 1) if stride_in_1x1 else (1, stride)
|
||||||
|
|
||||||
|
self.conv1 = Conv2d(
|
||||||
|
in_channels,
|
||||||
|
bottleneck_channels,
|
||||||
|
kernel_size=1,
|
||||||
|
stride=stride_1x1,
|
||||||
|
bias=False,
|
||||||
|
norm=get_norm(norm, bottleneck_channels),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.conv2 = Conv2d(
|
||||||
|
bottleneck_channels,
|
||||||
|
bottleneck_channels,
|
||||||
|
kernel_size=3,
|
||||||
|
stride=stride_3x3,
|
||||||
|
padding=1 * dilation,
|
||||||
|
bias=False,
|
||||||
|
groups=num_groups,
|
||||||
|
dilation=dilation,
|
||||||
|
norm=get_norm(norm, bottleneck_channels),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.conv3 = Conv2d(
|
||||||
|
bottleneck_channels,
|
||||||
|
out_channels,
|
||||||
|
kernel_size=1,
|
||||||
|
bias=False,
|
||||||
|
norm=get_norm(norm, out_channels),
|
||||||
|
)
|
||||||
|
|
||||||
|
for layer in [self.conv1, self.conv2, self.conv3, self.shortcut]:
|
||||||
|
if layer is not None: # shortcut can be None
|
||||||
|
weight_init.c2_msra_fill(layer)
|
||||||
|
|
||||||
|
# Zero-initialize the last normalization in each residual branch,
|
||||||
|
# so that at the beginning, the residual branch starts with zeros,
|
||||||
|
# and each residual block behaves like an identity.
|
||||||
|
# See Sec 5.1 in "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour":
|
||||||
|
# "For BN layers, the learnable scaling coefficient γ is initialized
|
||||||
|
# to be 1, except for each residual block's last BN
|
||||||
|
# where γ is initialized to be 0."
|
||||||
|
|
||||||
|
# nn.init.constant_(self.conv3.norm.weight, 0)
|
||||||
|
# TODO this somehow hurts performance when training GN models from scratch.
|
||||||
|
# Add it as an option when we need to use this code to train a backbone.
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
out = self.conv1(x)
|
||||||
|
out = F.relu_(out)
|
||||||
|
|
||||||
|
out = self.conv2(out)
|
||||||
|
out = F.relu_(out)
|
||||||
|
|
||||||
|
out = self.conv3(out)
|
||||||
|
|
||||||
|
if self.shortcut is not None:
|
||||||
|
shortcut = self.shortcut(x)
|
||||||
|
else:
|
||||||
|
shortcut = x
|
||||||
|
|
||||||
|
out += shortcut
|
||||||
|
out = F.relu_(out)
|
||||||
|
return out
|
||||||
|
|
||||||
|
def make_stage(block_class, num_blocks, first_stride, **kwargs):
|
||||||
|
"""
|
||||||
|
Create a resnet stage by creating many blocks.
|
||||||
|
Args:
|
||||||
|
block_class (class): a subclass of ResNetBlockBase
|
||||||
|
num_blocks (int):
|
||||||
|
first_stride (int): the stride of the first block. The other blocks will have stride=1.
|
||||||
|
A `stride` argument will be passed to the block constructor.
|
||||||
|
kwargs: other arguments passed to the block constructor.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
list[nn.Module]: a list of block module.
|
||||||
|
"""
|
||||||
|
blocks = []
|
||||||
|
for i in range(num_blocks):
|
||||||
|
blocks.append(block_class(stride=first_stride if i == 0 else 1, **kwargs))
|
||||||
|
kwargs["in_channels"] = kwargs["out_channels"]
|
||||||
|
return blocks
|
||||||
|
|
||||||
|
class ResNet(Backbone):
|
||||||
|
def __init__(self, stem, stages, num_classes=None, out_features=None):
|
||||||
|
"""
|
||||||
|
Args:
|
||||||
|
stem (nn.Module): a stem module
|
||||||
|
stages (list[list[ResNetBlock]]): several (typically 4) stages,
|
||||||
|
each contains multiple :class:`ResNetBlockBase`.
|
||||||
|
num_classes (None or int): if None, will not perform classification.
|
||||||
|
out_features (list[str]): name of the layers whose outputs should
|
||||||
|
be returned in forward. Can be anything in "stem", "linear", or "res2" ...
|
||||||
|
If None, will return the output of the last layer.
|
||||||
|
"""
|
||||||
|
super(ResNet, self).__init__()
|
||||||
|
self.stem = stem
|
||||||
|
self.num_classes = num_classes
|
||||||
|
|
||||||
|
current_stride = self.stem.stride
|
||||||
|
self._out_feature_strides = {"stem": current_stride}
|
||||||
|
self._out_feature_channels = {"stem": self.stem.out_channels}
|
||||||
|
|
||||||
|
self.stages = []
|
||||||
|
self.names = []
|
||||||
|
for i, blocks in enumerate(stages):
|
||||||
|
for block in blocks:
|
||||||
|
assert isinstance(block, ResNetBlockBase), block
|
||||||
|
curr_channels = block.out_channels
|
||||||
|
stage = nn.Sequential(*blocks)
|
||||||
|
name = "res" + str(i + 2)
|
||||||
|
self.stages.append(stage)
|
||||||
|
self.names.append(name)
|
||||||
|
self._out_feature_strides[name] = current_stride = int(
|
||||||
|
current_stride * np.prod([k.stride for k in blocks])
|
||||||
|
)
|
||||||
|
self._out_feature_channels[name] = blocks[-1].out_channels
|
||||||
|
self.stages = nn.ModuleList(self.stages)
|
||||||
|
|
||||||
|
if num_classes is not None:
|
||||||
|
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
|
||||||
|
self.linear = nn.Linear(curr_channels, num_classes)
|
||||||
|
|
||||||
|
# Sec 5.1 in "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour":
|
||||||
|
# "The 1000-way fully-connected layer is initialized by
|
||||||
|
# drawing weights from a zero-mean Gaussian with standard deviation of 0.01."
|
||||||
|
nn.init.normal_(self.linear.weight, stddev=0.01)
|
||||||
|
name = "linear"
|
||||||
|
|
||||||
|
if out_features is None:
|
||||||
|
out_features = [name]
|
||||||
|
self._out_features = out_features
|
||||||
|
assert len(self._out_features)
|
||||||
|
for out_feature in self._out_features:
|
||||||
|
assert out_feature in self.names, "Available children: {}".format(", ".join(self.names))
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
outputs = {}
|
||||||
|
x = self.stem(x)
|
||||||
|
if "stem" in self._out_features:
|
||||||
|
outputs["stem"] = x
|
||||||
|
for ix, stage in enumerate(self.stages):
|
||||||
|
name = self.names[ix]
|
||||||
|
x = stage(x)
|
||||||
|
if name in self._out_features:
|
||||||
|
outputs[name] = x
|
||||||
|
if self.num_classes is not None:
|
||||||
|
x = self.avgpool(x)
|
||||||
|
x = self.linear(x)
|
||||||
|
if "linear" in self._out_features:
|
||||||
|
outputs["linear"] = x
|
||||||
|
return outputs
|
||||||
|
|
||||||
|
def build_resnet_backbone(in_channels=3):
|
||||||
|
norm = "BN"
|
||||||
|
stem = BasicStem(
|
||||||
|
in_channels=in_channels,
|
||||||
|
out_channels=64,
|
||||||
|
norm=norm,
|
||||||
|
)
|
||||||
|
|
||||||
|
# fmt: off
|
||||||
|
out_features = ["res2", "res3", "res4"]
|
||||||
|
depth = 101
|
||||||
|
num_groups = 1
|
||||||
|
bottleneck_channels = 64
|
||||||
|
in_channels = 64
|
||||||
|
out_channels = 256
|
||||||
|
stride_in_1x1 = True
|
||||||
|
res5_dilation = 1
|
||||||
|
# fmt: on
|
||||||
|
assert res5_dilation in {1, 2}, "res5_dilation cannot be {}.".format(res5_dilation)
|
||||||
|
|
||||||
|
num_blocks_per_stage = {50: [3, 4, 6, 3], 101: [3, 4, 23, 3], 152: [3, 8, 36, 3]}[depth]
|
||||||
|
|
||||||
|
stages = []
|
||||||
|
|
||||||
|
# Avoid creating variables without gradients
|
||||||
|
# It consumes extra memory and may cause allreduce to fail
|
||||||
|
out_stage_idx = [{"res2": 2, "res3": 3, "res4": 4, "res5": 5}[f] for f in out_features]
|
||||||
|
max_stage_idx = max(out_stage_idx)
|
||||||
|
for idx, stage_idx in enumerate(range(2, max_stage_idx + 1)):
|
||||||
|
dilation = res5_dilation if stage_idx == 5 else 1
|
||||||
|
first_stride = 1 if idx == 0 or (stage_idx == 5 and dilation == 2) else 2
|
||||||
|
stage_kargs = dict()
|
||||||
|
stage_kargs.update({
|
||||||
|
"num_blocks": num_blocks_per_stage[idx],
|
||||||
|
"first_stride": first_stride,
|
||||||
|
"in_channels": in_channels,
|
||||||
|
"bottleneck_channels": bottleneck_channels,
|
||||||
|
"out_channels": out_channels,
|
||||||
|
"num_groups": num_groups,
|
||||||
|
"norm": norm,
|
||||||
|
"stride_in_1x1": stride_in_1x1,
|
||||||
|
"dilation": dilation,
|
||||||
|
})
|
||||||
|
stage_kargs["block_class"] = BottleneckBlock
|
||||||
|
blocks = make_stage(**stage_kargs)
|
||||||
|
in_channels = out_channels
|
||||||
|
out_channels *= 2
|
||||||
|
bottleneck_channels *= 2
|
||||||
|
stages.append(blocks)
|
||||||
|
return ResNet(stem, stages, out_features=out_features)
|
||||||
@@ -0,0 +1,244 @@
|
|||||||
|
import math
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import torch.nn.functional as F
|
||||||
|
import torch.utils.model_zoo as model_zoo
|
||||||
|
from core.bodyseg.backbone.backbone import Backbone
|
||||||
|
|
||||||
|
def fixed_padding(inputs, kernel_size, dilation):
|
||||||
|
kernel_size_effective = kernel_size + (kernel_size - 1) * (dilation - 1)
|
||||||
|
pad_total = kernel_size_effective - 1
|
||||||
|
pad_beg = pad_total // 2
|
||||||
|
pad_end = pad_total - pad_beg
|
||||||
|
padded_inputs = F.pad(inputs, (pad_beg, pad_end, pad_beg, pad_end))
|
||||||
|
return padded_inputs
|
||||||
|
|
||||||
|
class SeparableConv2d(nn.Module):
|
||||||
|
def __init__(self, inplanes, planes, kernel_size=3, stride=1, dilation=1, bias=False, BatchNorm=None):
|
||||||
|
super(SeparableConv2d, self).__init__()
|
||||||
|
|
||||||
|
self.conv1 = nn.Conv2d(inplanes, inplanes, kernel_size, stride, 0, dilation,
|
||||||
|
groups=inplanes, bias=bias)
|
||||||
|
self.bn = BatchNorm(inplanes)
|
||||||
|
self.pointwise = nn.Conv2d(inplanes, planes, 1, 1, 0, 1, 1, bias=bias)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x = fixed_padding(x, self.conv1.kernel_size[0], dilation=self.conv1.dilation[0])
|
||||||
|
x = self.conv1(x)
|
||||||
|
x = self.bn(x)
|
||||||
|
x = self.pointwise(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
class Block(nn.Module):
|
||||||
|
def __init__(self, inplanes, planes, reps, stride=1, dilation=1, BatchNorm=None,
|
||||||
|
start_with_relu=True, grow_first=True, is_last=False):
|
||||||
|
super(Block, self).__init__()
|
||||||
|
|
||||||
|
if planes != inplanes or stride != 1:
|
||||||
|
self.skip = nn.Conv2d(inplanes, planes, 1, stride=stride, bias=False)
|
||||||
|
self.skipbn = BatchNorm(planes)
|
||||||
|
else:
|
||||||
|
self.skip = None
|
||||||
|
|
||||||
|
self.relu = nn.ReLU(inplace=True)
|
||||||
|
rep = []
|
||||||
|
|
||||||
|
filters = inplanes
|
||||||
|
if grow_first:
|
||||||
|
rep.append(self.relu)
|
||||||
|
rep.append(SeparableConv2d(inplanes, planes, 3, 1, dilation, BatchNorm=BatchNorm))
|
||||||
|
rep.append(BatchNorm(planes))
|
||||||
|
filters = planes
|
||||||
|
|
||||||
|
for i in range(reps - 1):
|
||||||
|
rep.append(self.relu)
|
||||||
|
rep.append(SeparableConv2d(filters, filters, 3, 1, dilation, BatchNorm=BatchNorm))
|
||||||
|
rep.append(BatchNorm(filters))
|
||||||
|
|
||||||
|
if not grow_first:
|
||||||
|
rep.append(self.relu)
|
||||||
|
rep.append(SeparableConv2d(inplanes, planes, 3, 1, dilation, BatchNorm=BatchNorm))
|
||||||
|
rep.append(BatchNorm(planes))
|
||||||
|
|
||||||
|
if stride != 1:
|
||||||
|
rep.append(self.relu)
|
||||||
|
rep.append(SeparableConv2d(planes, planes, 3, 2, BatchNorm=BatchNorm))
|
||||||
|
rep.append(BatchNorm(planes))
|
||||||
|
|
||||||
|
if stride == 1 and is_last:
|
||||||
|
rep.append(self.relu)
|
||||||
|
rep.append(SeparableConv2d(planes, planes, 3, 1, BatchNorm=BatchNorm))
|
||||||
|
rep.append(BatchNorm(planes))
|
||||||
|
|
||||||
|
if not start_with_relu:
|
||||||
|
rep = rep[1:]
|
||||||
|
|
||||||
|
self.rep = nn.Sequential(*rep)
|
||||||
|
|
||||||
|
def forward(self, inp):
|
||||||
|
x = self.rep(inp)
|
||||||
|
|
||||||
|
if self.skip is not None:
|
||||||
|
skip = self.skip(inp)
|
||||||
|
skip = self.skipbn(skip)
|
||||||
|
else:
|
||||||
|
skip = inp
|
||||||
|
|
||||||
|
x = x + skip
|
||||||
|
|
||||||
|
return x
|
||||||
|
|
||||||
|
class AlignedXception(Backbone):
|
||||||
|
"""
|
||||||
|
Modified Alighed Xception
|
||||||
|
"""
|
||||||
|
def __init__(self, output_stride, BatchNorm):
|
||||||
|
super(AlignedXception, self).__init__()
|
||||||
|
|
||||||
|
if output_stride == 16:
|
||||||
|
entry_block3_stride = 2
|
||||||
|
middle_block_dilation = 1
|
||||||
|
exit_block_dilations = (1, 2)
|
||||||
|
elif output_stride == 8:
|
||||||
|
entry_block3_stride = 1
|
||||||
|
middle_block_dilation = 2
|
||||||
|
exit_block_dilations = (2, 4)
|
||||||
|
else:
|
||||||
|
raise NotImplementedError
|
||||||
|
|
||||||
|
|
||||||
|
# Entry flow
|
||||||
|
self.conv1 = nn.Conv2d(3, 32, 3, stride=2, padding=1, bias=False)
|
||||||
|
self.bn1 = BatchNorm(32)
|
||||||
|
self.relu = nn.ReLU(inplace=True)
|
||||||
|
|
||||||
|
self.conv2 = nn.Conv2d(32, 64, 3, stride=1, padding=1, bias=False)
|
||||||
|
self.bn2 = BatchNorm(64)
|
||||||
|
|
||||||
|
self.block1 = Block(64, 128, reps=2, stride=2, BatchNorm=BatchNorm, start_with_relu=False)
|
||||||
|
self.block2 = Block(128, 256, reps=2, stride=2, BatchNorm=BatchNorm, start_with_relu=False,
|
||||||
|
grow_first=True)
|
||||||
|
self.block3 = Block(256, 728, reps=2, stride=entry_block3_stride, BatchNorm=BatchNorm,
|
||||||
|
start_with_relu=True, grow_first=True, is_last=True)
|
||||||
|
|
||||||
|
# Middle flow
|
||||||
|
self.block4 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block5 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block6 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block7 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block8 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block9 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block10 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block11 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block12 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block13 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block14 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block15 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block16 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block17 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block18 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
self.block19 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||||
|
|
||||||
|
# Exit flow
|
||||||
|
self.block20 = Block(728, 1024, reps=2, stride=1, dilation=exit_block_dilations[0],
|
||||||
|
BatchNorm=BatchNorm, start_with_relu=True, grow_first=False, is_last=True)
|
||||||
|
|
||||||
|
self.conv3 = SeparableConv2d(1024, 1536, 3, stride=1, dilation=exit_block_dilations[1], BatchNorm=BatchNorm)
|
||||||
|
self.bn3 = BatchNorm(1536)
|
||||||
|
|
||||||
|
self.conv4 = SeparableConv2d(1536, 1536, 3, stride=1, dilation=exit_block_dilations[1], BatchNorm=BatchNorm)
|
||||||
|
self.bn4 = BatchNorm(1536)
|
||||||
|
|
||||||
|
self.conv5 = SeparableConv2d(1536, 2048, 3, stride=1, dilation=exit_block_dilations[1], BatchNorm=BatchNorm)
|
||||||
|
self.bn5 = BatchNorm(2048)
|
||||||
|
|
||||||
|
# Init weights
|
||||||
|
self._init_weight()
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
# Entry flow
|
||||||
|
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.block1(x)
|
||||||
|
# add relu here
|
||||||
|
x = self.relu(x)
|
||||||
|
low_level_feat = x
|
||||||
|
x = self.block2(x)
|
||||||
|
x = self.block3(x)
|
||||||
|
|
||||||
|
# Middle flow
|
||||||
|
x = self.block4(x)
|
||||||
|
x = self.block5(x)
|
||||||
|
x = self.block6(x)
|
||||||
|
x = self.block7(x)
|
||||||
|
x = self.block8(x)
|
||||||
|
x = self.block9(x)
|
||||||
|
x = self.block10(x)
|
||||||
|
x = self.block11(x)
|
||||||
|
x = self.block12(x)
|
||||||
|
x = self.block13(x)
|
||||||
|
x = self.block14(x)
|
||||||
|
x = self.block15(x)
|
||||||
|
x = self.block16(x)
|
||||||
|
x = self.block17(x)
|
||||||
|
x = self.block18(x)
|
||||||
|
x = self.block19(x)
|
||||||
|
|
||||||
|
# Exit flow
|
||||||
|
x = self.block20(x)
|
||||||
|
x = self.relu(x)
|
||||||
|
x = self.conv3(x)
|
||||||
|
x = self.bn3(x)
|
||||||
|
x = self.relu(x)
|
||||||
|
|
||||||
|
x = self.conv4(x)
|
||||||
|
x = self.bn4(x)
|
||||||
|
x = self.relu(x)
|
||||||
|
|
||||||
|
x = self.conv5(x)
|
||||||
|
x = self.bn5(x)
|
||||||
|
x = self.relu(x)
|
||||||
|
|
||||||
|
return x, low_level_feat
|
||||||
|
|
||||||
|
def _init_weight(self):
|
||||||
|
for m in self.modules():
|
||||||
|
if isinstance(m, nn.Conv2d):
|
||||||
|
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
|
||||||
|
m.weight.data.normal_(0, math.sqrt(2. / n))
|
||||||
|
elif isinstance(m, nn.SyncBatchNorm):
|
||||||
|
m.weight.data.fill_(1)
|
||||||
|
m.bias.data.zero_()
|
||||||
|
elif isinstance(m, nn.BatchNorm2d):
|
||||||
|
m.weight.data.fill_(1)
|
||||||
|
m.bias.data.zero_()
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
import torch
|
||||||
|
model = AlignedXception(BatchNorm=nn.BatchNorm2d, output_stride=16)
|
||||||
|
input = torch.rand(1, 3, 512, 512)
|
||||||
|
output, low_level_feat = model(input)
|
||||||
|
print(output.size())
|
||||||
|
print(low_level_feat.size())
|
||||||
@@ -0,0 +1,330 @@
|
|||||||
|
import sys
|
||||||
|
# sys.path.append("/Users/momo/human_seg_train")
|
||||||
|
# print(sys.path)
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import torch.nn.functional as F
|
||||||
|
from core.bodyseg.backbone.backbone import get_norm
|
||||||
|
|
||||||
|
class ConvBNReLU(nn.Sequential):
|
||||||
|
def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1, norm_layer=None):
|
||||||
|
padding = (kernel_size - 1) // 2
|
||||||
|
if norm_layer is None:
|
||||||
|
norm_layer = nn.BatchNorm2d
|
||||||
|
super(ConvBNReLU, self).__init__(
|
||||||
|
nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
|
||||||
|
norm_layer(out_planes),
|
||||||
|
nn.ReLU6(inplace=True)
|
||||||
|
)
|
||||||
|
|
||||||
|
class InvertedResidual(nn.Module):
|
||||||
|
def __init__(self, inp, oup, stride, expand_ratio, norm_layer=None):
|
||||||
|
super(InvertedResidual, self).__init__()
|
||||||
|
self.stride = stride
|
||||||
|
assert stride in [1, 2]
|
||||||
|
|
||||||
|
if norm_layer is None:
|
||||||
|
norm_layer = nn.BatchNorm2d
|
||||||
|
|
||||||
|
hidden_dim = int(round(inp * expand_ratio))
|
||||||
|
self.use_res_connect = self.stride == 1 and inp == oup
|
||||||
|
|
||||||
|
layers = []
|
||||||
|
if expand_ratio != 1:
|
||||||
|
# pw
|
||||||
|
layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1, norm_layer=norm_layer))
|
||||||
|
layers.extend([
|
||||||
|
# dw
|
||||||
|
ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim, norm_layer=norm_layer),
|
||||||
|
# pw-linear
|
||||||
|
nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
|
||||||
|
norm_layer(oup),
|
||||||
|
])
|
||||||
|
self.conv = nn.Sequential(*layers)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
if self.use_res_connect:
|
||||||
|
return x + self.conv(x)
|
||||||
|
else:
|
||||||
|
return self.conv(x)
|
||||||
|
|
||||||
|
class UpSampleBlock(nn.Module):
|
||||||
|
def __init__(self, in_channels, out_channels, expand_ratio=6):
|
||||||
|
super(UpSampleBlock, self).__init__()
|
||||||
|
self.refine = InvertedResidual(in_channels, out_channels, 1, expand_ratio)
|
||||||
|
|
||||||
|
def forward(self, x0, x1):
|
||||||
|
x = torch.cat([x0, x1], dim=1)
|
||||||
|
x = self.refine(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
class BodySegNet_32_thin_sigmod(nn.Module):
|
||||||
|
def __init__(self, input_channels=3, class_nums=1, output_onnx=False):
|
||||||
|
super(BodySegNet_32_thin_sigmod, self).__init__()
|
||||||
|
self.class_nums = class_nums
|
||||||
|
self.output_onnx = output_onnx
|
||||||
|
|
||||||
|
self.stage_1 = nn.Sequential(
|
||||||
|
ConvBNReLU(input_channels, 16, kernel_size=3, stride=2)
|
||||||
|
)
|
||||||
|
self.stage_2 = nn.Sequential(
|
||||||
|
ConvBNReLU(16, 16, kernel_size=3, stride=2, groups=16),
|
||||||
|
ConvBNReLU(16, 16, kernel_size=1, stride=1),
|
||||||
|
)
|
||||||
|
self.stage_3 = nn.Sequential(
|
||||||
|
InvertedResidual(16, 24, stride=2, expand_ratio=6),
|
||||||
|
InvertedResidual(24, 24, stride=1, expand_ratio=6),
|
||||||
|
InvertedResidual(24, 24, stride=1, expand_ratio=6),
|
||||||
|
)
|
||||||
|
self.stage_4 = nn.Sequential(
|
||||||
|
InvertedResidual(24, 32, stride=2, expand_ratio=6),
|
||||||
|
InvertedResidual(32, 32, stride=1, expand_ratio=6),
|
||||||
|
InvertedResidual(32, 32, stride=1, expand_ratio=6),
|
||||||
|
InvertedResidual(32, 32, stride=1, expand_ratio=6),
|
||||||
|
)
|
||||||
|
self.stage_5 = nn.Sequential(
|
||||||
|
InvertedResidual(32, 48, stride=2, expand_ratio=6),
|
||||||
|
InvertedResidual(48, 48, stride=1, expand_ratio=6),
|
||||||
|
InvertedResidual(48, 48, stride=1, expand_ratio=6),
|
||||||
|
InvertedResidual(48, 48, stride=1, expand_ratio=6)
|
||||||
|
)
|
||||||
|
self.up_to_4 = UpSampleBlock(48 + 32, 16)
|
||||||
|
self.up_to_3 = UpSampleBlock(16 + 24, 16)
|
||||||
|
self.up_to_2 = UpSampleBlock(16 + 16, 16)
|
||||||
|
self.up_to_1 = UpSampleBlock(16 + 16, 16)
|
||||||
|
self.last_layer = nn.Sequential(
|
||||||
|
ConvBNReLU(16, 16, kernel_size=1, stride=1),
|
||||||
|
nn.Conv2d(16, self.class_nums, kernel_size=1, stride=1)
|
||||||
|
)
|
||||||
|
self._initialize_weights()
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
feature_S = []
|
||||||
|
x1 = self.stage_1(x)
|
||||||
|
x2 = self.stage_2(x1)
|
||||||
|
x3 = self.stage_3(x2)
|
||||||
|
x4 = self.stage_4(x3)
|
||||||
|
feature = self.stage_5(x4)
|
||||||
|
|
||||||
|
feature = F.interpolate(feature, size=x4.size()[2:], mode='bilinear', align_corners=True)
|
||||||
|
feature = self.up_to_4(x4, feature)
|
||||||
|
feature = F.interpolate(feature, size=x3.size()[2:], mode='bilinear', align_corners=True)
|
||||||
|
feature = self.up_to_3(x3, feature)
|
||||||
|
feature = F.interpolate(feature, size=x2.size()[2:], mode='bilinear', align_corners=True)
|
||||||
|
feature = self.up_to_2(x2, feature)
|
||||||
|
feature = F.interpolate(feature, size=x1.size()[2:], mode='bilinear', align_corners=True)
|
||||||
|
feature = self.up_to_1(x1, feature)
|
||||||
|
feature_S.append(feature)
|
||||||
|
feature = self.last_layer(feature)
|
||||||
|
feature_S.append(feature)
|
||||||
|
output = F.interpolate(feature, size=x.size()[2:], mode='bilinear', align_corners=True)
|
||||||
|
# output = torch.sigmoid(output)
|
||||||
|
if self.output_onnx:
|
||||||
|
output = torch.argmax(output, dim=1).to(torch.float32)
|
||||||
|
|
||||||
|
return output, feature_S
|
||||||
|
|
||||||
|
def _initialize_weights(self):
|
||||||
|
for name, m in self.named_modules():
|
||||||
|
if isinstance(m, nn.Conv2d):
|
||||||
|
if 'first' in name:
|
||||||
|
nn.init.normal_(m.weight, 0, 0.01)
|
||||||
|
else:
|
||||||
|
nn.init.normal_(m.weight, 0, 1.0 / m.weight.shape[1])
|
||||||
|
if m.bias is not None:
|
||||||
|
nn.init.constant_(m.bias, 0)
|
||||||
|
elif isinstance(m, nn.BatchNorm2d):
|
||||||
|
nn.init.constant_(m.weight, 1)
|
||||||
|
if m.bias is not None:
|
||||||
|
nn.init.constant_(m.bias, 0.0001)
|
||||||
|
nn.init.constant_(m.running_mean, 0)
|
||||||
|
elif isinstance(m, nn.BatchNorm1d):
|
||||||
|
nn.init.constant_(m.weight, 1)
|
||||||
|
if m.bias is not None:
|
||||||
|
nn.init.constant_(m.bias, 0.0001)
|
||||||
|
nn.init.constant_(m.running_mean, 0)
|
||||||
|
elif isinstance(m, nn.Linear):
|
||||||
|
nn.init.normal_(m.weight, 0, 0.01)
|
||||||
|
if m.bias is not None:
|
||||||
|
nn.init.constant_(m.bias, 0)
|
||||||
|
|
||||||
|
class _ASPPModule(nn.Module):
|
||||||
|
def __init__(self, inplanes, planes, kernel_size, padding, dilation, BatchNorm):
|
||||||
|
super(_ASPPModule, self).__init__()
|
||||||
|
self.atrous_conv = nn.Conv2d(inplanes, planes, kernel_size=kernel_size,
|
||||||
|
stride=1, padding=padding, dilation=dilation, bias=False)
|
||||||
|
self.bn = BatchNorm(planes)
|
||||||
|
self.relu = nn.ReLU()
|
||||||
|
|
||||||
|
self._init_weight()
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x = self.atrous_conv(x)
|
||||||
|
x = self.bn(x)
|
||||||
|
|
||||||
|
return self.relu(x)
|
||||||
|
|
||||||
|
def _init_weight(self):
|
||||||
|
for m in self.modules():
|
||||||
|
if isinstance(m, nn.Conv2d):
|
||||||
|
torch.nn.init.kaiming_normal_(m.weight)
|
||||||
|
|
||||||
|
class ASPP(nn.Module):
|
||||||
|
def __init__(self, backbone, output_stride, BatchNorm):
|
||||||
|
super(ASPP, self).__init__()
|
||||||
|
if backbone == 'drn':
|
||||||
|
inplanes = 512
|
||||||
|
elif backbone == 'mobilenet':
|
||||||
|
inplanes = 320
|
||||||
|
elif backbone == 'resnet_fpn':
|
||||||
|
inplanes = 256
|
||||||
|
else:
|
||||||
|
inplanes = 2048
|
||||||
|
if output_stride == 16:
|
||||||
|
dilations = [1, 6, 12, 18]
|
||||||
|
elif output_stride == 8:
|
||||||
|
dilations = [1, 12, 24, 36]
|
||||||
|
else:
|
||||||
|
raise NotImplementedError
|
||||||
|
|
||||||
|
self.aspp1 = _ASPPModule(inplanes, 256, 1, padding=0, dilation=dilations[0], BatchNorm=BatchNorm)
|
||||||
|
self.aspp2 = _ASPPModule(inplanes, 256, 3, padding=dilations[1], dilation=dilations[1], BatchNorm=BatchNorm)
|
||||||
|
self.aspp3 = _ASPPModule(inplanes, 256, 3, padding=dilations[2], dilation=dilations[2], BatchNorm=BatchNorm)
|
||||||
|
self.aspp4 = _ASPPModule(inplanes, 256, 3, padding=dilations[3], dilation=dilations[3], BatchNorm=BatchNorm)
|
||||||
|
|
||||||
|
self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
|
||||||
|
nn.Conv2d(inplanes, 256, 1, stride=1, bias=False),
|
||||||
|
BatchNorm(256),
|
||||||
|
nn.ReLU())
|
||||||
|
self.conv1 = nn.Conv2d(1280, 256, 1, bias=False)
|
||||||
|
self.bn1 = BatchNorm(256)
|
||||||
|
self.relu = nn.ReLU()
|
||||||
|
self.dropout = nn.Dropout(0.5)
|
||||||
|
self._init_weight()
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x1 = self.aspp1(x)
|
||||||
|
x2 = self.aspp2(x)
|
||||||
|
x3 = self.aspp3(x)
|
||||||
|
x4 = self.aspp4(x)
|
||||||
|
x5 = self.global_avg_pool(x)
|
||||||
|
x5 = F.interpolate(x5, size=x4.size()[2:], mode='bilinear', align_corners=True)
|
||||||
|
# x5 = F.interpolate(x5, size=x4.size()[2:], mode='nearest')
|
||||||
|
x = torch.cat((x1, x2, x3, x4, x5), dim=1)
|
||||||
|
|
||||||
|
x = self.conv1(x)
|
||||||
|
x = self.bn1(x)
|
||||||
|
x = self.relu(x)
|
||||||
|
|
||||||
|
return self.dropout(x)
|
||||||
|
|
||||||
|
def _init_weight(self):
|
||||||
|
for m in self.modules():
|
||||||
|
if isinstance(m, nn.Conv2d):
|
||||||
|
torch.nn.init.kaiming_normal_(m.weight)
|
||||||
|
|
||||||
|
class Decoder(nn.Module):
|
||||||
|
def __init__(self,num_classes, backbone, BatchNorm):
|
||||||
|
super(Decoder, self).__init__()
|
||||||
|
if backbone == 'resnet_fpn' or backbone == 'drn':
|
||||||
|
low_level_inplanes = 256
|
||||||
|
elif backbone == 'xception':
|
||||||
|
low_level_inplanes = 128
|
||||||
|
elif backbone == 'mobilenet':
|
||||||
|
low_level_inplanes = 24
|
||||||
|
else:
|
||||||
|
raise NotImplementedError
|
||||||
|
|
||||||
|
self.conv1 = nn.Conv2d(low_level_inplanes, 16, 1, bias=False)
|
||||||
|
self.bn1 = BatchNorm(16)
|
||||||
|
self.relu = nn.ReLU()
|
||||||
|
self.last_conv = nn.Sequential(nn.Conv2d(304, 256, kernel_size=3, stride=1, padding=1, bias=False),
|
||||||
|
BatchNorm(256),
|
||||||
|
nn.ReLU(),
|
||||||
|
nn.Dropout(0.5),
|
||||||
|
nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=False),
|
||||||
|
BatchNorm(256),
|
||||||
|
nn.ReLU(),
|
||||||
|
nn.Dropout(0.1),
|
||||||
|
nn.Conv2d(256, num_classes, kernel_size=1, stride=1))
|
||||||
|
|
||||||
|
self.up_to_1 = UpSampleBlock(16 + 16, 16)
|
||||||
|
self.last_layer = nn.Sequential(
|
||||||
|
ConvBNReLU(16, 16, kernel_size=1, stride=1),
|
||||||
|
nn.Conv2d(16, 1, kernel_size=1, stride=1)
|
||||||
|
)
|
||||||
|
self._init_weight()
|
||||||
|
|
||||||
|
# x(1,256,8,6) low(1,256,32,24) -》 x(1,1,32,24)
|
||||||
|
def forward(self, x, low_level_feat):
|
||||||
|
feature_T = []
|
||||||
|
# deeplab part
|
||||||
|
#(1,256,32,24) -> (1,16,64,48)
|
||||||
|
low_level_feat = self.conv1(low_level_feat)
|
||||||
|
low_level_feat = self.bn1(low_level_feat)
|
||||||
|
low_level_feat = self.relu(low_level_feat)
|
||||||
|
|
||||||
|
# x(1, 256, 8, 6)-> (1,16,32,24)
|
||||||
|
x = F.interpolate(x, size=low_level_feat.size()[2:], mode='bilinear', align_corners=True)
|
||||||
|
low_level_feat = F.interpolate(low_level_feat,
|
||||||
|
size=[low_level_feat.size()[2] * 2, low_level_feat.size()[3] * 2],
|
||||||
|
mode='bilinear', align_corners=True)
|
||||||
|
x = self.conv1(x)
|
||||||
|
x = self.bn1(x)
|
||||||
|
x = self.relu(x)
|
||||||
|
|
||||||
|
# bodyseg part
|
||||||
|
feature = F.interpolate(x, size=[x.size()[2] * 2,x.size()[3] * 2], mode='bilinear', align_corners=True)
|
||||||
|
# 输入up_to_1 x(low)(1,16,64,48),上采样2倍后的feature(1,16,64,48)
|
||||||
|
feature = self.up_to_1(low_level_feat, feature)
|
||||||
|
feature_T.append(feature)
|
||||||
|
# (1,1,64,48)
|
||||||
|
feature = self.last_layer(feature)
|
||||||
|
feature_T.append(feature)
|
||||||
|
# (1,1,128,96)
|
||||||
|
output = F.interpolate(feature, size=[128,96], mode='bilinear', align_corners=True)
|
||||||
|
# output = torch.sigmoid(output)
|
||||||
|
|
||||||
|
return output, feature_T
|
||||||
|
|
||||||
|
def _init_weight(self):
|
||||||
|
for m in self.modules():
|
||||||
|
if isinstance(m, nn.Conv2d):
|
||||||
|
torch.nn.init.kaiming_normal_(m.weight)
|
||||||
|
|
||||||
|
def build_backbone(backbone, output_stride, BatchNorm, input_channel=3):
|
||||||
|
from core.bodyseg.backbone.fpn import build_resnet_fpn_backbone
|
||||||
|
from core.bodyseg.backbone.xception import AlignedXception
|
||||||
|
|
||||||
|
return build_resnet_fpn_backbone(input_channel)
|
||||||
|
|
||||||
|
def build_aspp(backbone, output_stride, BatchNorm):
|
||||||
|
return ASPP(backbone, output_stride, BatchNorm)
|
||||||
|
|
||||||
|
def build_decoder(num_classes, backbone, BatchNorm):
|
||||||
|
return Decoder(num_classes, backbone, BatchNorm)
|
||||||
|
|
||||||
|
class DeepLab(nn.Module):
|
||||||
|
def __init__(self, input_channel=3, class_num=1):
|
||||||
|
super(DeepLab, self).__init__()
|
||||||
|
|
||||||
|
BatchNorm = get_norm("BN")
|
||||||
|
|
||||||
|
self.backbone = build_backbone("resnet_fpn", 16, BatchNorm, input_channel=input_channel)
|
||||||
|
self.aspp = build_aspp("resnet_fpn", 16, BatchNorm)
|
||||||
|
self.decoder = build_decoder(class_num, "resnet_fpn", BatchNorm)
|
||||||
|
|
||||||
|
def forward(self, input):
|
||||||
|
# input(1,3,128,96)
|
||||||
|
#output: "p2"(1,256,32,24), "p3"(1,256,16,12), "p4"(1,256,8,6)
|
||||||
|
output = self.backbone(input)
|
||||||
|
# "p4"(1,256,8,6) "p2"(1,256,32,24)
|
||||||
|
x, low_level_feat = output['p4'], output['p2']
|
||||||
|
# x(1,256,8,6)-》(1,256,8,6)
|
||||||
|
x = self.aspp(x)
|
||||||
|
# x(1,256,8,6) low(1,256,32,24) -》 x(1,1,32,24)
|
||||||
|
x, feature_T = self.decoder(x, low_level_feat)
|
||||||
|
# x(1,1,128,96)
|
||||||
|
x = F.interpolate(x, size=input.size()[2:], mode='bilinear', align_corners=True)
|
||||||
|
# x = F.interpolate(x, size=input.size()[2:], mode='nearest')
|
||||||
|
return x, feature_T
|
||||||
@@ -0,0 +1,59 @@
|
|||||||
|
from qcloud_cos import CosConfig
|
||||||
|
from qcloud_cos import CosS3Client
|
||||||
|
import sys
|
||||||
|
import os
|
||||||
|
import os.path as osp
|
||||||
|
import time
|
||||||
|
import logging
|
||||||
|
from common.logger import config as confccc
|
||||||
|
|
||||||
|
|
||||||
|
class COS_object():
|
||||||
|
def __init__(self):
|
||||||
|
# 正常情况日志级别使用 INFO，需要定位时可以修改为 DEBUG，此时 SDK 会打印和服务端的通信信息
|
||||||
|
logging.basicConfig(level=logging.INFO, stream=sys.stdout)
|
||||||
|
|
||||||
|
# 1. 设置用户属性, 包括 secret_id, secret_key, region等。从环境变量读取，避免硬编码密钥。
|
||||||
|
secret_id = os.getenv('COS_SECRET_ID', '<your-cos-secret-id>')
|
||||||
|
secret_key = os.getenv('COS_SECRET_KEY', '<your-cos-secret-key>')
|
||||||
|
self.region = os.getenv('COS_REGION', '<your-cos-region>') # 例: ap-beijing
|
||||||
|
token = None # 如果使用永久密钥不需要填入 token，如果使用临时密钥需要填入，临时密钥生成和使用指引参见 https://cloud.tencent.com/document/product/436/14048
|
||||||
|
self.scheme = 'https' # 指定使用 http/https 协议来访问 COS，默认为 https，可不填
|
||||||
|
self.BucketName= os.getenv('COS_BUCKET', '<your-cos-bucket>')
|
||||||
|
for param in (secret_id, secret_key, self.region, self.BucketName):
|
||||||
|
assert '<' not in param, '请设置环境变量 COS_SECRET_ID / COS_SECRET_KEY / COS_REGION / COS_BUCKET'
|
||||||
|
config = CosConfig(Region=self.region, SecretId=secret_id, SecretKey=secret_key, Token=token, Scheme=self.scheme)
|
||||||
|
self.client = CosS3Client(config)
|
||||||
|
|
||||||
|
def upload_file(self, file, target_name):
|
||||||
|
t0 = time.time()
|
||||||
|
with open(file, 'rb') as fp:
|
||||||
|
response = self.client.put_object(
|
||||||
|
Bucket=self.BucketName, # Bucket 由 BucketName-APPID 组成
|
||||||
|
Body=fp,
|
||||||
|
Key=target_name,
|
||||||
|
StorageClass='STANDARD',
|
||||||
|
ContentType='image/jpeg;image/jpg;image/png;image/gif'
|
||||||
|
)
|
||||||
|
ret_url = self.scheme + '://' + self.BucketName + '.cos.' + self.region + '.myqcloud.com/' + target_name
|
||||||
|
print('time costs: {}'.format(time.time() - t0), ret_url)
|
||||||
|
return ret_url
|
||||||
|
|
||||||
|
def download_img(self, img_url, tmp_dir):
|
||||||
|
# user_img_save_dir = confccc.get('default', 'userDir')
|
||||||
|
# user_img_tmp_dir = confccc.get('default', 'tmp_dir')
|
||||||
|
fileName = img_url.split('.com/')[-1]
|
||||||
|
print(fileName)
|
||||||
|
response = self.client.get_object(
|
||||||
|
Bucket=self.BucketName,
|
||||||
|
Key=fileName,
|
||||||
|
)
|
||||||
|
response['Body'].get_stream_to_file(tmp_dir)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__=='__main__':
|
||||||
|
t0 = time.time()
|
||||||
|
cos = COS_object()
|
||||||
|
cos.upload_file('/home/chinatszrn/Pictures/vaffflue12.png', 'hair_mz/images/vaffflue12.png')
|
||||||
|
# cos.download_img('https://ydapp-1317132355.cos.ap-beijing.myqcloud.com/hair_mz/images/hairstyle/fb7d7a58-3228-40f2-84a6-337088ee31e2/2023031623425988.jpg', 'a0')
|
||||||
|
print(time.time() - t0)
|
||||||
@@ -0,0 +1,2 @@
|
|||||||
|
from . import mesh
|
||||||
|
from . import morphable_model
|
||||||
@@ -0,0 +1,7 @@
|
|||||||
|
# from .cython import mesh_core_cython
|
||||||
|
from . import io
|
||||||
|
from . import vis
|
||||||
|
from . import transform
|
||||||
|
from . import light
|
||||||
|
from . import render
|
||||||
|
|
||||||
@@ -0,0 +1 @@
|
|||||||
|
*.cpython-36m-x86_64-linux-gnu.so
|
||||||
@@ -0,0 +1,375 @@
|
|||||||
|
/*
|
||||||
|
functions that can not be optimazed by vertorization in python.
|
||||||
|
1. rasterization.(need process each triangle)
|
||||||
|
2. normal of each vertex.(use one-ring, need process each vertex)
|
||||||
|
3. write obj(seems that it can be verctorized? anyway, writing it in c++ is simple, so also add function here. --> however, why writting in c++ is still slow?)
|
||||||
|
|
||||||
|
Author: Yao Feng
|
||||||
|
Mail: yaofeng1995@gmail.com
|
||||||
|
*/
|
||||||
|
|
||||||
|
#include "mesh_core.h"
|
||||||
|
|
||||||
|
|
||||||
|
/* Judge whether the point is in the triangle
|
||||||
|
Method:
|
||||||
|
http://blackpawn.com/texts/pointinpoly/
|
||||||
|
Args:
|
||||||
|
point: [x, y]
|
||||||
|
tri_points: three vertices(2d points) of a triangle. 2 coords x 3 vertices
|
||||||
|
Returns:
|
||||||
|
bool: true for in triangle
|
||||||
|
*/
|
||||||
|
bool isPointInTri(point p, point p0, point p1, point p2)
|
||||||
|
{
|
||||||
|
// vectors
|
||||||
|
point v0, v1, v2;
|
||||||
|
v0 = p2 - p0;
|
||||||
|
v1 = p1 - p0;
|
||||||
|
v2 = p - p0;
|
||||||
|
|
||||||
|
// dot products
|
||||||
|
float dot00 = v0.dot(v0); //v0.x * v0.x + v0.y * v0.y //np.dot(v0.T, v0)
|
||||||
|
float dot01 = v0.dot(v1); //v0.x * v1.x + v0.y * v1.y //np.dot(v0.T, v1)
|
||||||
|
float dot02 = v0.dot(v2); //v0.x * v2.x + v0.y * v2.y //np.dot(v0.T, v2)
|
||||||
|
float dot11 = v1.dot(v1); //v1.x * v1.x + v1.y * v1.y //np.dot(v1.T, v1)
|
||||||
|
float dot12 = v1.dot(v2); //v1.x * v2.x + v1.y * v2.y//np.dot(v1.T, v2)
|
||||||
|
|
||||||
|
// barycentric coordinates
|
||||||
|
float inverDeno;
|
||||||
|
if(dot00*dot11 - dot01*dot01 == 0)
|
||||||
|
inverDeno = 0;
|
||||||
|
else
|
||||||
|
inverDeno = 1/(dot00*dot11 - dot01*dot01);
|
||||||
|
|
||||||
|
float u = (dot11*dot02 - dot01*dot12)*inverDeno;
|
||||||
|
float v = (dot00*dot12 - dot01*dot02)*inverDeno;
|
||||||
|
|
||||||
|
// check if point in triangle
|
||||||
|
return (u >= 0) && (v >= 0) && (u + v < 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
void get_point_weight(float* weight, point p, point p0, point p1, point p2)
|
||||||
|
{
|
||||||
|
// vectors
|
||||||
|
point v0, v1, v2;
|
||||||
|
v0 = p2 - p0;
|
||||||
|
v1 = p1 - p0;
|
||||||
|
v2 = p - p0;
|
||||||
|
|
||||||
|
// dot products
|
||||||
|
float dot00 = v0.dot(v0); //v0.x * v0.x + v0.y * v0.y //np.dot(v0.T, v0)
|
||||||
|
float dot01 = v0.dot(v1); //v0.x * v1.x + v0.y * v1.y //np.dot(v0.T, v1)
|
||||||
|
float dot02 = v0.dot(v2); //v0.x * v2.x + v0.y * v2.y //np.dot(v0.T, v2)
|
||||||
|
float dot11 = v1.dot(v1); //v1.x * v1.x + v1.y * v1.y //np.dot(v1.T, v1)
|
||||||
|
float dot12 = v1.dot(v2); //v1.x * v2.x + v1.y * v2.y//np.dot(v1.T, v2)
|
||||||
|
|
||||||
|
// barycentric coordinates
|
||||||
|
float inverDeno;
|
||||||
|
if(dot00*dot11 - dot01*dot01 == 0)
|
||||||
|
inverDeno = 0;
|
||||||
|
else
|
||||||
|
inverDeno = 1/(dot00*dot11 - dot01*dot01);
|
||||||
|
|
||||||
|
float u = (dot11*dot02 - dot01*dot12)*inverDeno;
|
||||||
|
float v = (dot00*dot12 - dot01*dot02)*inverDeno;
|
||||||
|
|
||||||
|
// weight
|
||||||
|
weight[0] = 1 - u - v;
|
||||||
|
weight[1] = v;
|
||||||
|
weight[2] = u;
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
void _get_normal_core(
|
||||||
|
float* normal, float* tri_normal, int* triangles,
|
||||||
|
int ntri)
|
||||||
|
{
|
||||||
|
int i, j;
|
||||||
|
int tri_p0_ind, tri_p1_ind, tri_p2_ind;
|
||||||
|
|
||||||
|
for(i = 0; i < ntri; i++)
|
||||||
|
{
|
||||||
|
tri_p0_ind = triangles[3*i];
|
||||||
|
tri_p1_ind = triangles[3*i + 1];
|
||||||
|
tri_p2_ind = triangles[3*i + 2];
|
||||||
|
|
||||||
|
for(j = 0; j < 3; j++)
|
||||||
|
{
|
||||||
|
normal[3*tri_p0_ind + j] = normal[3*tri_p0_ind + j] + tri_normal[3*i + j];
|
||||||
|
normal[3*tri_p1_ind + j] = normal[3*tri_p1_ind + j] + tri_normal[3*i + j];
|
||||||
|
normal[3*tri_p2_ind + j] = normal[3*tri_p2_ind + j] + tri_normal[3*i + j];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
void _rasterize_triangles_core(
|
||||||
|
float* vertices, int* triangles,
|
||||||
|
float* depth_buffer, int* triangle_buffer, float* barycentric_weight,
|
||||||
|
int nver, int ntri,
|
||||||
|
int h, int w)
|
||||||
|
{
|
||||||
|
int i;
|
||||||
|
int x, y, k;
|
||||||
|
int tri_p0_ind, tri_p1_ind, tri_p2_ind;
|
||||||
|
point p0, p1, p2, p;
|
||||||
|
int x_min, x_max, y_min, y_max;
|
||||||
|
float p_depth, p0_depth, p1_depth, p2_depth;
|
||||||
|
float weight[3];
|
||||||
|
|
||||||
|
for(i = 0; i < ntri; i++)
|
||||||
|
{
|
||||||
|
tri_p0_ind = triangles[3*i];
|
||||||
|
tri_p1_ind = triangles[3*i + 1];
|
||||||
|
tri_p2_ind = triangles[3*i + 2];
|
||||||
|
|
||||||
|
p0.x = vertices[3*tri_p0_ind]; p0.y = vertices[3*tri_p0_ind + 1]; p0_depth = vertices[3*tri_p0_ind + 2];
|
||||||
|
p1.x = vertices[3*tri_p1_ind]; p1.y = vertices[3*tri_p1_ind + 1]; p1_depth = vertices[3*tri_p1_ind + 2];
|
||||||
|
p2.x = vertices[3*tri_p2_ind]; p2.y = vertices[3*tri_p2_ind + 1]; p2_depth = vertices[3*tri_p2_ind + 2];
|
||||||
|
|
||||||
|
x_min = max((int)ceil(min(p0.x, min(p1.x, p2.x))), 0);
|
||||||
|
x_max = min((int)floor(max(p0.x, max(p1.x, p2.x))), w - 1);
|
||||||
|
|
||||||
|
y_min = max((int)ceil(min(p0.y, min(p1.y, p2.y))), 0);
|
||||||
|
y_max = min((int)floor(max(p0.y, max(p1.y, p2.y))), h - 1);
|
||||||
|
|
||||||
|
if(x_max < x_min || y_max < y_min)
|
||||||
|
{
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
for(y = y_min; y <= y_max; y++) //h
|
||||||
|
{
|
||||||
|
for(x = x_min; x <= x_max; x++) //w
|
||||||
|
{
|
||||||
|
p.x = x; p.y = y;
|
||||||
|
if(p.x < 2 || p.x > w - 3 || p.y < 2 || p.y > h - 3 || isPointInTri(p, p0, p1, p2))
|
||||||
|
{
|
||||||
|
get_point_weight(weight, p, p0, p1, p2);
|
||||||
|
p_depth = weight[0]*p0_depth + weight[1]*p1_depth + weight[2]*p2_depth;
|
||||||
|
|
||||||
|
if((p_depth > depth_buffer[y*w + x]))
|
||||||
|
{
|
||||||
|
depth_buffer[y*w + x] = p_depth;
|
||||||
|
triangle_buffer[y*w + x] = i;
|
||||||
|
for(k = 0; k < 3; k++)
|
||||||
|
{
|
||||||
|
barycentric_weight[y*w*3 + x*3 + k] = weight[k];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
void _render_colors_core(
|
||||||
|
float* image, float* vertices, int* triangles,
|
||||||
|
float* colors,
|
||||||
|
float* depth_buffer,
|
||||||
|
int nver, int ntri,
|
||||||
|
int h, int w, int c)
|
||||||
|
{
|
||||||
|
int i;
|
||||||
|
int x, y, k;
|
||||||
|
int tri_p0_ind, tri_p1_ind, tri_p2_ind;
|
||||||
|
point p0, p1, p2, p;
|
||||||
|
int x_min, x_max, y_min, y_max;
|
||||||
|
float p_depth, p0_depth, p1_depth, p2_depth;
|
||||||
|
float p_color, p0_color, p1_color, p2_color;
|
||||||
|
float weight[3];
|
||||||
|
|
||||||
|
for(i = 0; i < ntri; i++)
|
||||||
|
{
|
||||||
|
tri_p0_ind = triangles[3*i];
|
||||||
|
tri_p1_ind = triangles[3*i + 1];
|
||||||
|
tri_p2_ind = triangles[3*i + 2];
|
||||||
|
|
||||||
|
p0.x = vertices[3*tri_p0_ind]; p0.y = vertices[3*tri_p0_ind + 1]; p0_depth = vertices[3*tri_p0_ind + 2];
|
||||||
|
p1.x = vertices[3*tri_p1_ind]; p1.y = vertices[3*tri_p1_ind + 1]; p1_depth = vertices[3*tri_p1_ind + 2];
|
||||||
|
p2.x = vertices[3*tri_p2_ind]; p2.y = vertices[3*tri_p2_ind + 1]; p2_depth = vertices[3*tri_p2_ind + 2];
|
||||||
|
|
||||||
|
x_min = max((int)ceil(min(p0.x, min(p1.x, p2.x))), 0);
|
||||||
|
x_max = min((int)floor(max(p0.x, max(p1.x, p2.x))), w - 1);
|
||||||
|
|
||||||
|
y_min = max((int)ceil(min(p0.y, min(p1.y, p2.y))), 0);
|
||||||
|
y_max = min((int)floor(max(p0.y, max(p1.y, p2.y))), h - 1);
|
||||||
|
|
||||||
|
if(x_max < x_min || y_max < y_min)
|
||||||
|
{
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
for(y = y_min; y <= y_max; y++) //h
|
||||||
|
{
|
||||||
|
for(x = x_min; x <= x_max; x++) //w
|
||||||
|
{
|
||||||
|
p.x = x; p.y = y;
|
||||||
|
if(p.x < 2 || p.x > w - 3 || p.y < 2 || p.y > h - 3 || isPointInTri(p, p0, p1, p2))
|
||||||
|
{
|
||||||
|
get_point_weight(weight, p, p0, p1, p2);
|
||||||
|
p_depth = weight[0]*p0_depth + weight[1]*p1_depth + weight[2]*p2_depth;
|
||||||
|
|
||||||
|
if((p_depth > depth_buffer[y*w + x]))
|
||||||
|
{
|
||||||
|
for(k = 0; k < c; k++) // c
|
||||||
|
{
|
||||||
|
p0_color = colors[c*tri_p0_ind + k];
|
||||||
|
p1_color = colors[c*tri_p1_ind + k];
|
||||||
|
p2_color = colors[c*tri_p2_ind + k];
|
||||||
|
|
||||||
|
p_color = weight[0]*p0_color + weight[1]*p1_color + weight[2]*p2_color;
|
||||||
|
image[y*w*c + x*c + k] = p_color;
|
||||||
|
}
|
||||||
|
|
||||||
|
depth_buffer[y*w + x] = p_depth;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
void _render_texture_core(
|
||||||
|
float* image, float* vertices, int* triangles,
|
||||||
|
float* texture, float* tex_coords, int* tex_triangles,
|
||||||
|
float* depth_buffer,
|
||||||
|
int nver, int tex_nver, int ntri,
|
||||||
|
int h, int w, int c,
|
||||||
|
int tex_h, int tex_w, int tex_c,
|
||||||
|
int mapping_type)
|
||||||
|
{
|
||||||
|
int i;
|
||||||
|
int x, y, k;
|
||||||
|
int tri_p0_ind, tri_p1_ind, tri_p2_ind;
|
||||||
|
int tex_tri_p0_ind, tex_tri_p1_ind, tex_tri_p2_ind;
|
||||||
|
point p0, p1, p2, p;
|
||||||
|
point tex_p0, tex_p1, tex_p2, tex_p;
|
||||||
|
int x_min, x_max, y_min, y_max;
|
||||||
|
float weight[3];
|
||||||
|
float p_depth, p0_depth, p1_depth, p2_depth;
|
||||||
|
float xd, yd;
|
||||||
|
float ul, ur, dl, dr;
|
||||||
|
for(i = 0; i < ntri; i++)
|
||||||
|
{
|
||||||
|
// mesh
|
||||||
|
tri_p0_ind = triangles[3*i];
|
||||||
|
tri_p1_ind = triangles[3*i + 1];
|
||||||
|
tri_p2_ind = triangles[3*i + 2];
|
||||||
|
|
||||||
|
p0.x = vertices[3*tri_p0_ind]; p0.y = vertices[3*tri_p0_ind + 1]; p0_depth = vertices[3*tri_p0_ind + 2];
|
||||||
|
p1.x = vertices[3*tri_p1_ind]; p1.y = vertices[3*tri_p1_ind + 1]; p1_depth = vertices[3*tri_p1_ind + 2];
|
||||||
|
p2.x = vertices[3*tri_p2_ind]; p2.y = vertices[3*tri_p2_ind + 1]; p2_depth = vertices[3*tri_p2_ind + 2];
|
||||||
|
|
||||||
|
// texture
|
||||||
|
tex_tri_p0_ind = tex_triangles[3*i];
|
||||||
|
tex_tri_p1_ind = tex_triangles[3*i + 1];
|
||||||
|
tex_tri_p2_ind = tex_triangles[3*i + 2];
|
||||||
|
|
||||||
|
tex_p0.x = tex_coords[3*tex_tri_p0_ind]; tex_p0.y = tex_coords[3*tri_p0_ind + 1];
|
||||||
|
tex_p1.x = tex_coords[3*tex_tri_p1_ind]; tex_p1.y = tex_coords[3*tri_p1_ind + 1];
|
||||||
|
tex_p2.x = tex_coords[3*tex_tri_p2_ind]; tex_p2.y = tex_coords[3*tri_p2_ind + 1];
|
||||||
|
|
||||||
|
|
||||||
|
x_min = max((int)ceil(min(p0.x, min(p1.x, p2.x))), 0);
|
||||||
|
x_max = min((int)floor(max(p0.x, max(p1.x, p2.x))), w - 1);
|
||||||
|
|
||||||
|
y_min = max((int)ceil(min(p0.y, min(p1.y, p2.y))), 0);
|
||||||
|
y_max = min((int)floor(max(p0.y, max(p1.y, p2.y))), h - 1);
|
||||||
|
|
||||||
|
|
||||||
|
if(x_max < x_min || y_max < y_min)
|
||||||
|
{
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
for(y = y_min; y <= y_max; y++) //h
|
||||||
|
{
|
||||||
|
for(x = x_min; x <= x_max; x++) //w
|
||||||
|
{
|
||||||
|
p.x = x; p.y = y;
|
||||||
|
if(p.x < 2 || p.x > w - 3 || p.y < 2 || p.y > h - 3 || isPointInTri(p, p0, p1, p2))
|
||||||
|
{
|
||||||
|
get_point_weight(weight, p, p0, p1, p2);
|
||||||
|
p_depth = weight[0]*p0_depth + weight[1]*p1_depth + weight[2]*p2_depth;
|
||||||
|
|
||||||
|
if((p_depth > depth_buffer[y*w + x]))
|
||||||
|
{
|
||||||
|
// -- color from texture
|
||||||
|
// cal weight in mesh tri
|
||||||
|
get_point_weight(weight, p, p0, p1, p2);
|
||||||
|
// cal coord in texture
|
||||||
|
tex_p = tex_p0*weight[0] + tex_p1*weight[1] + tex_p2*weight[2];
|
||||||
|
tex_p.x = max(min(tex_p.x, float(tex_w - 1)), float(0));
|
||||||
|
tex_p.y = max(min(tex_p.y, float(tex_h - 1)), float(0));
|
||||||
|
|
||||||
|
yd = tex_p.y - floor(tex_p.y);
|
||||||
|
xd = tex_p.x - floor(tex_p.x);
|
||||||
|
for(k = 0; k < c; k++)
|
||||||
|
{
|
||||||
|
if(mapping_type==0)// nearest
|
||||||
|
{
|
||||||
|
image[y*w*c + x*c + k] = texture[int(round(tex_p.y))*tex_w*tex_c + int(round(tex_p.x))*tex_c + k];
|
||||||
|
}
|
||||||
|
else//bilinear interp
|
||||||
|
{
|
||||||
|
ul = texture[(int)floor(tex_p.y)*tex_w*tex_c + (int)floor(tex_p.x)*tex_c + k];
|
||||||
|
ur = texture[(int)floor(tex_p.y)*tex_w*tex_c + (int)ceil(tex_p.x)*tex_c + k];
|
||||||
|
dl = texture[(int)ceil(tex_p.y)*tex_w*tex_c + (int)floor(tex_p.x)*tex_c + k];
|
||||||
|
dr = texture[(int)ceil(tex_p.y)*tex_w*tex_c + (int)ceil(tex_p.x)*tex_c + k];
|
||||||
|
|
||||||
|
image[y*w*c + x*c + k] = ul*(1-xd)*(1-yd) + ur*xd*(1-yd) + dl*(1-xd)*yd + dr*xd*yd;
|
||||||
|
}
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
|
depth_buffer[y*w + x] = p_depth;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
// ------------------------------------------------- write
|
||||||
|
// obj write
|
||||||
|
// Ref: https://github.com/patrikhuber/eos/blob/master/include/eos/core/Mesh.hpp
|
||||||
|
void _write_obj_with_colors_texture(string filename, string mtl_name,
|
||||||
|
float* vertices, int* triangles, float* colors, float* uv_coords,
|
||||||
|
int nver, int ntri, int ntexver)
|
||||||
|
{
|
||||||
|
int i;
|
||||||
|
|
||||||
|
ofstream obj_file(filename.c_str());
|
||||||
|
|
||||||
|
// first line of the obj file: the mtl name
|
||||||
|
obj_file << "mtllib " << mtl_name << endl;
|
||||||
|
|
||||||
|
// write vertices
|
||||||
|
for (i = 0; i < nver; ++i)
|
||||||
|
{
|
||||||
|
obj_file << "v " << vertices[3*i] << " " << vertices[3*i + 1] << " " << vertices[3*i + 2] << colors[3*i] << " " << colors[3*i + 1] << " " << colors[3*i + 2] << endl;
|
||||||
|
}
|
||||||
|
|
||||||
|
// write uv coordinates
|
||||||
|
for (i = 0; i < ntexver; ++i)
|
||||||
|
{
|
||||||
|
//obj_file << "vt " << uv_coords[2*i] << " " << (1 - uv_coords[2*i + 1]) << endl;
|
||||||
|
obj_file << "vt " << uv_coords[2*i] << " " << uv_coords[2*i + 1] << endl;
|
||||||
|
}
|
||||||
|
|
||||||
|
obj_file << "usemtl FaceTexture" << endl;
|
||||||
|
// write triangles
|
||||||
|
for (i = 0; i < ntri; ++i)
|
||||||
|
{
|
||||||
|
// obj_file << "f " << triangles[3*i] << "/" << triangles[3*i] << " " << triangles[3*i + 1] << "/" << triangles[3*i + 1] << " " << triangles[3*i + 2] << "/" << triangles[3*i + 2] << endl;
|
||||||
|
obj_file << "f " << triangles[3*i + 2] << "/" << triangles[3*i + 2] << " " << triangles[3*i + 1] << "/" << triangles[3*i + 1] << " " << triangles[3*i] << "/" << triangles[3*i] << endl;
|
||||||
|
}
|
||||||
|
|
||||||
|
}
|
||||||
@@ -0,0 +1,83 @@
|
|||||||
|
#ifndef MESH_CORE_HPP_
|
||||||
|
#define MESH_CORE_HPP_
|
||||||
|
|
||||||
|
#include <stdio.h>
|
||||||
|
#include <cmath>
|
||||||
|
#include <algorithm>
|
||||||
|
#include <string>
|
||||||
|
#include <iostream>
|
||||||
|
#include <fstream>
|
||||||
|
|
||||||
|
using namespace std;
|
||||||
|
|
||||||
|
class point
|
||||||
|
{
|
||||||
|
public:
|
||||||
|
float x;
|
||||||
|
float y;
|
||||||
|
|
||||||
|
float dot(point p)
|
||||||
|
{
|
||||||
|
return this->x * p.x + this->y * p.y;
|
||||||
|
}
|
||||||
|
|
||||||
|
point operator-(const point& p)
|
||||||
|
{
|
||||||
|
point np;
|
||||||
|
np.x = this->x - p.x;
|
||||||
|
np.y = this->y - p.y;
|
||||||
|
return np;
|
||||||
|
}
|
||||||
|
|
||||||
|
point operator+(const point& p)
|
||||||
|
{
|
||||||
|
point np;
|
||||||
|
np.x = this->x + p.x;
|
||||||
|
np.y = this->y + p.y;
|
||||||
|
return np;
|
||||||
|
}
|
||||||
|
|
||||||
|
point operator*(float s)
|
||||||
|
{
|
||||||
|
point np;
|
||||||
|
np.x = s * this->x;
|
||||||
|
np.y = s * this->y;
|
||||||
|
return np;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
|
||||||
|
bool isPointInTri(point p, point p0, point p1, point p2, int h, int w);
|
||||||
|
void get_point_weight(float* weight, point p, point p0, point p1, point p2);
|
||||||
|
|
||||||
|
void _get_normal_core(
|
||||||
|
float* normal, float* tri_normal, int* triangles,
|
||||||
|
int ntri);
|
||||||
|
|
||||||
|
void _rasterize_triangles_core(
|
||||||
|
float* vertices, int* triangles,
|
||||||
|
float* depth_buffer, int* triangle_buffer, float* barycentric_weight,
|
||||||
|
int nver, int ntri,
|
||||||
|
int h, int w);
|
||||||
|
|
||||||
|
void _render_colors_core(
|
||||||
|
float* image, float* vertices, int* triangles,
|
||||||
|
float* colors,
|
||||||
|
float* depth_buffer,
|
||||||
|
int nver, int ntri,
|
||||||
|
int h, int w, int c);
|
||||||
|
|
||||||
|
void _render_texture_core(
|
||||||
|
float* image, float* vertices, int* triangles,
|
||||||
|
float* texture, float* tex_coords, int* tex_triangles,
|
||||||
|
float* depth_buffer,
|
||||||
|
int nver, int tex_nver, int ntri,
|
||||||
|
int h, int w, int c,
|
||||||
|
int tex_h, int tex_w, int tex_c,
|
||||||
|
int mapping_type);
|
||||||
|
|
||||||
|
void _write_obj_with_colors_texture(string filename, string mtl_name,
|
||||||
|
float* vertices, int* triangles, float* colors, float* uv_coords,
|
||||||
|
int nver, int ntri, int ntexver);
|
||||||
|
|
||||||
|
#endif
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,22 @@
|
|||||||
|
'''
|
||||||
|
python setup.py build_ext -i
|
||||||
|
to compile
|
||||||
|
'''
|
||||||
|
|
||||||
|
# setup.py
|
||||||
|
from distutils.core import setup
|
||||||
|
from setuptools import Extension
|
||||||
|
|
||||||
|
from Cython.Build import cythonize
|
||||||
|
from Cython.Distutils import build_ext
|
||||||
|
import numpy
|
||||||
|
|
||||||
|
setup(
|
||||||
|
name = 'mesh_core_cython',
|
||||||
|
cmdclass={'build_ext': build_ext},
|
||||||
|
ext_modules=[Extension("mesh_core_cython",
|
||||||
|
sources=["mesh_core_cython.pyx", "mesh_core.cpp"],
|
||||||
|
language='c++',
|
||||||
|
include_dirs=[numpy.get_include()])],
|
||||||
|
)
|
||||||
|
|
||||||
@@ -0,0 +1,142 @@
|
|||||||
|
from __future__ import absolute_import
|
||||||
|
from __future__ import division
|
||||||
|
from __future__ import print_function
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import os
|
||||||
|
# from skimage import io
|
||||||
|
from time import time
|
||||||
|
|
||||||
|
# from .cython import mesh_core_cython
|
||||||
|
|
||||||
|
## TODO
|
||||||
|
## TODO: c++ version
|
||||||
|
def read_obj(obj_name):
|
||||||
|
''' read mesh
|
||||||
|
'''
|
||||||
|
return 0
|
||||||
|
|
||||||
|
# ------------------------- write
|
||||||
|
def write_asc(path, vertices):
|
||||||
|
'''
|
||||||
|
Args:
|
||||||
|
vertices: shape = (nver, 3)
|
||||||
|
'''
|
||||||
|
if path.split('.')[-1] == 'asc':
|
||||||
|
np.savetxt(path, vertices)
|
||||||
|
else:
|
||||||
|
np.savetxt(path + '.asc', vertices)
|
||||||
|
|
||||||
|
def write_obj_with_colors(obj_name, vertices, triangles, colors):
|
||||||
|
''' Save 3D face model with texture represented by colors.
|
||||||
|
Args:
|
||||||
|
obj_name: str
|
||||||
|
vertices: shape = (nver, 3)
|
||||||
|
triangles: shape = (ntri, 3)
|
||||||
|
colors: shape = (nver, 3)
|
||||||
|
'''
|
||||||
|
triangles = triangles.copy()
|
||||||
|
triangles += 1 # meshlab start with 1
|
||||||
|
|
||||||
|
if obj_name.split('.')[-1] != 'obj':
|
||||||
|
obj_name = obj_name + '.obj'
|
||||||
|
|
||||||
|
# write obj
|
||||||
|
with open(obj_name, 'w') as f:
|
||||||
|
|
||||||
|
# write vertices & colors
|
||||||
|
for i in range(vertices.shape[0]):
|
||||||
|
# s = 'v {} {} {} \n'.format(vertices[0,i], vertices[1,i], vertices[2,i])
|
||||||
|
s = 'v {} {} {} {} {} {}\n'.format(vertices[i, 0], vertices[i, 1], vertices[i, 2], colors[i, 0], colors[i, 1], colors[i, 2])
|
||||||
|
f.write(s)
|
||||||
|
|
||||||
|
# write f: ver ind/ uv ind
|
||||||
|
[k, ntri] = triangles.shape
|
||||||
|
for i in range(triangles.shape[0]):
|
||||||
|
# s = 'f {} {} {}\n'.format(triangles[i, 0], triangles[i, 1], triangles[i, 2])
|
||||||
|
s = 'f {} {} {}\n'.format(triangles[i, 2], triangles[i, 1], triangles[i, 0])
|
||||||
|
f.write(s)
|
||||||
|
|
||||||
|
## TODO: c++ version
|
||||||
|
def write_obj_with_texture(obj_name, vertices, triangles, texture, uv_coords):
|
||||||
|
''' Save 3D face model with texture represented by texture map.
|
||||||
|
Ref: https://github.com/patrikhuber/eos/blob/bd00155ebae4b1a13b08bf5a991694d682abbada/include/eos/core/Mesh.hpp
|
||||||
|
Args:
|
||||||
|
obj_name: str
|
||||||
|
vertices: shape = (nver, 3)
|
||||||
|
triangles: shape = (ntri, 3)
|
||||||
|
texture: shape = (256,256,3)
|
||||||
|
uv_coords: shape = (nver, 3) max value<=1
|
||||||
|
'''
|
||||||
|
if obj_name.split('.')[-1] != 'obj':
|
||||||
|
obj_name = obj_name + '.obj'
|
||||||
|
mtl_name = obj_name.replace('.obj', '.mtl')
|
||||||
|
texture_name = obj_name.replace('.obj', '_texture.png')
|
||||||
|
|
||||||
|
triangles = triangles.copy()
|
||||||
|
triangles += 1 # mesh lab start with 1
|
||||||
|
|
||||||
|
# write obj
|
||||||
|
with open(obj_name, 'w') as f:
|
||||||
|
# first line: write mtlib(material library)
|
||||||
|
s = "mtllib {}\n".format(os.path.abspath(mtl_name))
|
||||||
|
f.write(s)
|
||||||
|
|
||||||
|
# write vertices
|
||||||
|
for i in range(vertices.shape[0]):
|
||||||
|
s = 'v {} {} {}\n'.format(vertices[i, 0], vertices[i, 1], vertices[i, 2])
|
||||||
|
f.write(s)
|
||||||
|
|
||||||
|
# write uv coords
|
||||||
|
for i in range(uv_coords.shape[0]):
|
||||||
|
s = 'vt {} {}\n'.format(uv_coords[i,0], 1 - uv_coords[i,1])
|
||||||
|
f.write(s)
|
||||||
|
|
||||||
|
f.write("usemtl FaceTexture\n")
|
||||||
|
|
||||||
|
# write f: ver ind/ uv ind
|
||||||
|
for i in range(triangles.shape[0]):
|
||||||
|
s = 'f {}/{} {}/{} {}/{}\n'.format(triangles[i,2], triangles[i,2], triangles[i,1], triangles[i,1], triangles[i,0], triangles[i,0])
|
||||||
|
f.write(s)
|
||||||
|
|
||||||
|
# write mtl
|
||||||
|
with open(mtl_name, 'w') as f:
|
||||||
|
f.write("newmtl FaceTexture\n")
|
||||||
|
s = 'map_Kd {}\n'.format(os.path.abspath(texture_name)) # map to image
|
||||||
|
f.write(s)
|
||||||
|
|
||||||
|
# write texture as png
|
||||||
|
imsave(texture_name, texture)
|
||||||
|
|
||||||
|
# c++ version
|
||||||
|
def write_obj_with_colors_texture(obj_name, vertices, triangles, colors, texture, uv_coords):
|
||||||
|
''' Save 3D face model with texture.
|
||||||
|
Ref: https://github.com/patrikhuber/eos/blob/bd00155ebae4b1a13b08bf5a991694d682abbada/include/eos/core/Mesh.hpp
|
||||||
|
Args:
|
||||||
|
obj_name: str
|
||||||
|
vertices: shape = (nver, 3)
|
||||||
|
triangles: shape = (ntri, 3)
|
||||||
|
colors: shape = (nver, 3)
|
||||||
|
texture: shape = (256,256,3)
|
||||||
|
uv_coords: shape = (nver, 3) max value<=1
|
||||||
|
'''
|
||||||
|
if obj_name.split('.')[-1] != 'obj':
|
||||||
|
obj_name = obj_name + '.obj'
|
||||||
|
mtl_name = obj_name.replace('.obj', '.mtl')
|
||||||
|
texture_name = obj_name.replace('.obj', '_texture.png')
|
||||||
|
|
||||||
|
triangles = triangles.copy()
|
||||||
|
triangles += 1 # mesh lab start with 1
|
||||||
|
|
||||||
|
# write obj
|
||||||
|
vertices, colors, uv_coords = vertices.astype(np.float32).copy(), colors.astype(np.float32).copy(), uv_coords.astype(np.float32).copy()
|
||||||
|
mesh_core_cython.write_obj_with_colors_texture_core(str.encode(obj_name), str.encode(os.path.abspath(mtl_name)), vertices, triangles, colors, uv_coords, vertices.shape[0], triangles.shape[0], uv_coords.shape[0])
|
||||||
|
|
||||||
|
# write mtl
|
||||||
|
with open(mtl_name, 'w') as f:
|
||||||
|
f.write("newmtl FaceTexture\n")
|
||||||
|
s = 'map_Kd {}\n'.format(os.path.abspath(texture_name)) # map to image
|
||||||
|
f.write(s)
|
||||||
|
|
||||||
|
# write texture as png
|
||||||
|
# io.imsave(texture_name, texture)
|
||||||
@@ -0,0 +1,213 @@
|
|||||||
|
'''
|
||||||
|
Functions about lighting mesh(changing colors/texture of mesh).
|
||||||
|
1. add light to colors/texture (shade each vertex)
|
||||||
|
2. fit light according to colors/texture & image.
|
||||||
|
'''
|
||||||
|
|
||||||
|
from __future__ import absolute_import
|
||||||
|
from __future__ import division
|
||||||
|
from __future__ import print_function
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
# from .cython import mesh_core_cython
|
||||||
|
|
||||||
|
def get_normal(vertices, triangles):
|
||||||
|
''' calculate normal direction in each vertex
|
||||||
|
Args:
|
||||||
|
vertices: [nver, 3]
|
||||||
|
triangles: [ntri, 3]
|
||||||
|
Returns:
|
||||||
|
normal: [nver, 3]
|
||||||
|
'''
|
||||||
|
pt0 = vertices[triangles[:, 0], :] # [ntri, 3]
|
||||||
|
pt1 = vertices[triangles[:, 1], :] # [ntri, 3]
|
||||||
|
pt2 = vertices[triangles[:, 2], :] # [ntri, 3]
|
||||||
|
tri_normal = np.cross(pt0 - pt1, pt0 - pt2) # [ntri, 3]. normal of each triangle
|
||||||
|
|
||||||
|
normal = np.zeros_like(vertices, dtype = np.float32).copy() # [nver, 3]
|
||||||
|
# for i in range(triangles.shape[0]):
|
||||||
|
# normal[triangles[i, 0], :] = normal[triangles[i, 0], :] + tri_normal[i, :]
|
||||||
|
# normal[triangles[i, 1], :] = normal[triangles[i, 1], :] + tri_normal[i, :]
|
||||||
|
# normal[triangles[i, 2], :] = normal[triangles[i, 2], :] + tri_normal[i, :]
|
||||||
|
mesh_core_cython.get_normal_core(normal, tri_normal.astype(np.float32).copy(), triangles.copy(), triangles.shape[0])
|
||||||
|
|
||||||
|
# normalize to unit length
|
||||||
|
mag = np.sum(normal**2, 1) # [nver]
|
||||||
|
zero_ind = (mag == 0)
|
||||||
|
mag[zero_ind] = 1
|
||||||
|
normal[zero_ind, 0] = np.ones((np.sum(zero_ind)))
|
||||||
|
|
||||||
|
normal = normal/np.sqrt(mag[:,np.newaxis])
|
||||||
|
|
||||||
|
return normal
|
||||||
|
|
||||||
|
# TODO: test
|
||||||
|
def add_light_sh(vertices, triangles, colors, sh_coeff):
|
||||||
|
'''
|
||||||
|
In 3d face, usually assume:
|
||||||
|
1. The surface of face is Lambertian(reflect only the low frequencies of lighting)
|
||||||
|
2. Lighting can be an arbitrary combination of point sources
|
||||||
|
--> can be expressed in terms of spherical harmonics(omit the lighting coefficients)
|
||||||
|
I = albedo * (sh(n) x sh_coeff)
|
||||||
|
|
||||||
|
albedo: n x 1
|
||||||
|
sh_coeff: 9 x 1
|
||||||
|
Y(n) = (1, n_x, n_y, n_z, n_xn_y, n_xn_z, n_yn_z, n_x^2 - n_y^2, 3n_z^2 - 1)': n x 9
|
||||||
|
# Y(n) = (1, n_x, n_y, n_z)': n x 4
|
||||||
|
|
||||||
|
Args:
|
||||||
|
vertices: [nver, 3]
|
||||||
|
triangles: [ntri, 3]
|
||||||
|
colors: [nver, 3] albedo
|
||||||
|
sh_coeff: [9, 1] spherical harmonics coefficients
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
lit_colors: [nver, 3]
|
||||||
|
'''
|
||||||
|
assert vertices.shape[0] == colors.shape[0]
|
||||||
|
nver = vertices.shape[0]
|
||||||
|
normal = get_normal(vertices, triangles) # [nver, 3]
|
||||||
|
sh = np.array((np.ones(nver), n[:,0], n[:,1], n[:,2], n[:,0]*n[:,1], n[:,0]*n[:,2], n[:,1]*n[:,2], n[:,0]**2 - n[:,1]**2, 3*(n[:,2]**2) - 1)) # [nver, 9]
|
||||||
|
ref = sh.dot(sh_coeff) #[nver, 1]
|
||||||
|
lit_colors = colors*ref
|
||||||
|
return lit_colors
|
||||||
|
|
||||||
|
|
||||||
|
def add_light(vertices, triangles, colors, light_positions = 0, light_intensities = 0):
|
||||||
|
''' Gouraud shading. add point lights.
|
||||||
|
In 3d face, usually assume:
|
||||||
|
1. The surface of face is Lambertian(reflect only the low frequencies of lighting)
|
||||||
|
2. Lighting can be an arbitrary combination of point sources
|
||||||
|
3. No specular (unless skin is oil, 23333)
|
||||||
|
|
||||||
|
Ref: https://cs184.eecs.berkeley.edu/lecture/pipeline
|
||||||
|
Args:
|
||||||
|
vertices: [nver, 3]
|
||||||
|
triangles: [ntri, 3]
|
||||||
|
light_positions: [nlight, 3]
|
||||||
|
light_intensities: [nlight, 3]
|
||||||
|
Returns:
|
||||||
|
lit_colors: [nver, 3]
|
||||||
|
'''
|
||||||
|
nver = vertices.shape[0]
|
||||||
|
normals = get_normal(vertices, triangles) # [nver, 3]
|
||||||
|
|
||||||
|
# ambient
|
||||||
|
# La = ka*Ia
|
||||||
|
|
||||||
|
# diffuse
|
||||||
|
# Ld = kd*(I/r^2)max(0, nxl)
|
||||||
|
direction_to_lights = vertices[np.newaxis, :, :] - light_positions[:, np.newaxis, :] # [nlight, nver, 3]
|
||||||
|
direction_to_lights_n = np.sqrt(np.sum(direction_to_lights**2, axis = 2)) # [nlight, nver]
|
||||||
|
direction_to_lights = direction_to_lights/direction_to_lights_n[:, :, np.newaxis]
|
||||||
|
normals_dot_lights = normals[np.newaxis, :, :]*direction_to_lights # [nlight, nver, 3]
|
||||||
|
normals_dot_lights = np.sum(normals_dot_lights, axis = 2) # [nlight, nver]
|
||||||
|
diffuse_output = colors[np.newaxis, :, :]*normals_dot_lights[:, :, np.newaxis]*light_intensities[:, np.newaxis, :]
|
||||||
|
diffuse_output = np.sum(diffuse_output, axis = 0) # [nver, 3]
|
||||||
|
|
||||||
|
# specular
|
||||||
|
# h = (v + l)/(|v + l|) bisector
|
||||||
|
# Ls = ks*(I/r^2)max(0, nxh)^p
|
||||||
|
# increasing p narrows the reflectionlob
|
||||||
|
|
||||||
|
lit_colors = diffuse_output # only diffuse part here.
|
||||||
|
lit_colors = np.minimum(np.maximum(lit_colors, 0), 1)
|
||||||
|
return lit_colors
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
## TODO. estimate light(sh coeff)
|
||||||
|
## -------------------------------- estimate. can not use now.
|
||||||
|
def fit_light(image, vertices, colors, triangles, vis_ind, lamb = 10, max_iter = 3):
|
||||||
|
[h, w, c] = image.shape
|
||||||
|
|
||||||
|
# surface normal
|
||||||
|
norm = get_normal(vertices, triangles)
|
||||||
|
|
||||||
|
nver = vertices.shape[1]
|
||||||
|
|
||||||
|
# vertices --> corresponding image pixel
|
||||||
|
pt2d = vertices[:2, :]
|
||||||
|
|
||||||
|
pt2d[0,:] = np.minimum(np.maximum(pt2d[0,:], 0), w - 1)
|
||||||
|
pt2d[1,:] = np.minimum(np.maximum(pt2d[1,:], 0), h - 1)
|
||||||
|
pt2d = np.round(pt2d).astype(np.int32) # 2 x nver
|
||||||
|
|
||||||
|
image_pixel = image[pt2d[1,:], pt2d[0,:], :] # nver x 3
|
||||||
|
image_pixel = image_pixel.T # 3 x nver
|
||||||
|
|
||||||
|
# vertices --> corresponding mean texture pixel with illumination
|
||||||
|
# Spherical Harmonic Basis
|
||||||
|
harmonic_dim = 9
|
||||||
|
nx = norm[0,:];
|
||||||
|
ny = norm[1,:];
|
||||||
|
nz = norm[2,:];
|
||||||
|
harmonic = np.zeros((nver, harmonic_dim))
|
||||||
|
|
||||||
|
pi = np.pi
|
||||||
|
harmonic[:,0] = np.sqrt(1/(4*pi)) * np.ones((nver,));
|
||||||
|
harmonic[:,1] = np.sqrt(3/(4*pi)) * nx;
|
||||||
|
harmonic[:,2] = np.sqrt(3/(4*pi)) * ny;
|
||||||
|
harmonic[:,3] = np.sqrt(3/(4*pi)) * nz;
|
||||||
|
harmonic[:,4] = 1/2. * np.sqrt(3/(4*pi)) * (2*nz**2 - nx**2 - ny**2);
|
||||||
|
harmonic[:,5] = 3 * np.sqrt(5/(12*pi)) * (ny*nz);
|
||||||
|
harmonic[:,6] = 3 * np.sqrt(5/(12*pi)) * (nx*nz);
|
||||||
|
harmonic[:,7] = 3 * np.sqrt(5/(12*pi)) * (nx*ny);
|
||||||
|
harmonic[:,8] = 3/2. * np.sqrt(5/(12*pi)) * (nx*nx - ny*ny);
|
||||||
|
|
||||||
|
'''
|
||||||
|
I' = sum(albedo * lj * hj) j = 0:9 (albedo = tex)
|
||||||
|
set A = albedo*h (n x 9)
|
||||||
|
alpha = lj (9 x 1)
|
||||||
|
Y = I (n x 1)
|
||||||
|
Y' = A.dot(alpha)
|
||||||
|
|
||||||
|
opt function:
|
||||||
|
||Y - A*alpha|| + lambda*(alpha'*alpha)
|
||||||
|
result:
|
||||||
|
A'*(Y - A*alpha) + lambda*alpha = 0
|
||||||
|
==>
|
||||||
|
(A'*A*alpha - lambda)*alpha = A'*Y
|
||||||
|
left: 9 x 9
|
||||||
|
right: 9 x 1
|
||||||
|
'''
|
||||||
|
n_vis_ind = len(vis_ind)
|
||||||
|
n = n_vis_ind*c
|
||||||
|
|
||||||
|
Y = np.zeros((n, 1))
|
||||||
|
A = np.zeros((n, 9))
|
||||||
|
light = np.zeros((3, 1))
|
||||||
|
|
||||||
|
for k in range(c):
|
||||||
|
Y[k*n_vis_ind:(k+1)*n_vis_ind, :] = image_pixel[k, vis_ind][:, np.newaxis]
|
||||||
|
A[k*n_vis_ind:(k+1)*n_vis_ind, :] = texture[k, vis_ind][:, np.newaxis] * harmonic[vis_ind, :]
|
||||||
|
Ac = texture[k, vis_ind][:, np.newaxis]
|
||||||
|
Yc = image_pixel[k, vis_ind][:, np.newaxis]
|
||||||
|
light[k] = (Ac.T.dot(Yc))/(Ac.T.dot(Ac))
|
||||||
|
|
||||||
|
for i in range(max_iter):
|
||||||
|
|
||||||
|
Yc = Y.copy()
|
||||||
|
for k in range(c):
|
||||||
|
Yc[k*n_vis_ind:(k+1)*n_vis_ind, :] /= light[k]
|
||||||
|
|
||||||
|
# update alpha
|
||||||
|
equation_left = np.dot(A.T, A) + lamb*np.eye(harmonic_dim); # why + ?
|
||||||
|
equation_right = np.dot(A.T, Yc)
|
||||||
|
alpha = np.dot(np.linalg.inv(equation_left), equation_right)
|
||||||
|
|
||||||
|
# update light
|
||||||
|
for k in range(c):
|
||||||
|
Ac = A[k*n_vis_ind:(k+1)*n_vis_ind, :].dot(alpha)
|
||||||
|
Yc = Y[k*n_vis_ind:(k+1)*n_vis_ind, :]
|
||||||
|
light[k] = (Ac.T.dot(Yc))/(Ac.T.dot(Ac))
|
||||||
|
|
||||||
|
appearance = np.zeros_like(texture)
|
||||||
|
for k in range(c):
|
||||||
|
tmp = np.dot(harmonic*texture[k, :][:, np.newaxis], alpha*light[k])
|
||||||
|
appearance[k,:] = tmp.T
|
||||||
|
|
||||||
|
appearance = np.minimum(np.maximum(appearance, 0), 1)
|
||||||
|
|
||||||
|
return appearance
|
||||||
|
|
||||||
@@ -0,0 +1,135 @@
|
|||||||
|
'''
|
||||||
|
functions about rendering mesh(from 3d obj to 2d image).
|
||||||
|
only use rasterization render here.
|
||||||
|
Note that:
|
||||||
|
1. Generally, render func includes camera, light, raterize. Here no camera and light(I write these in other files)
|
||||||
|
2. Generally, the input vertices are normalized to [-1,1] and cetered on [0, 0]. (in world space)
|
||||||
|
Here, the vertices are using image coords, which centers on [w/2, h/2] with the y-axis pointing to oppisite direction.
|
||||||
|
Means: render here only conducts interpolation.(I just want to make the input flexible)
|
||||||
|
|
||||||
|
Author: Yao Feng
|
||||||
|
Mail: yaofeng1995@gmail.com
|
||||||
|
'''
|
||||||
|
from __future__ import absolute_import
|
||||||
|
from __future__ import division
|
||||||
|
from __future__ import print_function
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
from time import time
|
||||||
|
|
||||||
|
# from .cython import mesh_core_cython
|
||||||
|
|
||||||
|
def rasterize_triangles(vertices, triangles, h, w):
|
||||||
|
'''
|
||||||
|
Args:
|
||||||
|
vertices: [nver, 3]
|
||||||
|
triangles: [ntri, 3]
|
||||||
|
h: height
|
||||||
|
w: width
|
||||||
|
Returns:
|
||||||
|
depth_buffer: [h, w] saves the depth, here, the bigger the z, the fronter the point.
|
||||||
|
triangle_buffer: [h, w] saves the tri id(-1 for no triangle).
|
||||||
|
barycentric_weight: [h, w, 3] saves corresponding barycentric weight.
|
||||||
|
|
||||||
|
# Each triangle has 3 vertices & Each vertex has 3 coordinates x, y, z.
|
||||||
|
# h, w is the size of rendering
|
||||||
|
'''
|
||||||
|
|
||||||
|
# initial
|
||||||
|
depth_buffer = np.zeros([h, w]) - 999999. #set the initial z to the farest position
|
||||||
|
triangle_buffer = np.zeros([h, w], dtype = np.int32) - 1 # if tri id = -1, the pixel has no triangle correspondance
|
||||||
|
barycentric_weight = np.zeros([h, w, 3], dtype = np.float32) #
|
||||||
|
|
||||||
|
vertices = vertices.astype(np.float32).copy()
|
||||||
|
triangles = triangles.astype(np.int32).copy()
|
||||||
|
|
||||||
|
mesh_core_cython.rasterize_triangles_core(
|
||||||
|
vertices, triangles,
|
||||||
|
depth_buffer, triangle_buffer, barycentric_weight,
|
||||||
|
vertices.shape[0], triangles.shape[0],
|
||||||
|
h, w)
|
||||||
|
|
||||||
|
def render_colors(vertices, triangles, colors, h, w, c = 3, BG = None):
|
||||||
|
''' render mesh with colors
|
||||||
|
Args:
|
||||||
|
vertices: [nver, 3]
|
||||||
|
triangles: [ntri, 3]
|
||||||
|
colors: [nver, 3]
|
||||||
|
h: height
|
||||||
|
w: width
|
||||||
|
c: channel
|
||||||
|
BG: background image
|
||||||
|
Returns:
|
||||||
|
image: [h, w, c]. rendered image./rendering.
|
||||||
|
'''
|
||||||
|
|
||||||
|
# initial
|
||||||
|
if BG is None:
|
||||||
|
image = np.zeros((h, w, c), dtype = np.float32)
|
||||||
|
else:
|
||||||
|
assert BG.shape[0] == h and BG.shape[1] == w and BG.shape[2] == c
|
||||||
|
image = BG
|
||||||
|
depth_buffer = np.zeros([h, w], dtype = np.float32, order = 'C') - 999999.
|
||||||
|
|
||||||
|
# change orders. --> C-contiguous order(column major)
|
||||||
|
vertices = vertices.astype(np.float32).copy()
|
||||||
|
triangles = triangles.astype(np.int32).copy()
|
||||||
|
colors = colors.astype(np.float32).copy()
|
||||||
|
###
|
||||||
|
st = time()
|
||||||
|
mesh_core_cython.render_colors_core(
|
||||||
|
image, vertices, triangles,
|
||||||
|
colors,
|
||||||
|
depth_buffer,
|
||||||
|
vertices.shape[0], triangles.shape[0],
|
||||||
|
h, w, c)
|
||||||
|
return image
|
||||||
|
|
||||||
|
|
||||||
|
def render_texture(vertices, triangles, texture, tex_coords, tex_triangles, h, w, c = 3, mapping_type = 'nearest', BG = None):
|
||||||
|
''' render mesh with texture map
|
||||||
|
Args:
|
||||||
|
vertices: [3, nver]
|
||||||
|
triangles: [3, ntri]
|
||||||
|
texture: [tex_h, tex_w, 3]
|
||||||
|
tex_coords: [ntexcoords, 3]
|
||||||
|
tex_triangles: [ntri, 3]
|
||||||
|
h: height of rendering
|
||||||
|
w: width of rendering
|
||||||
|
c: channel
|
||||||
|
mapping_type: 'bilinear' or 'nearest'
|
||||||
|
'''
|
||||||
|
# initial
|
||||||
|
if BG is None:
|
||||||
|
image = np.zeros((h, w, c), dtype = np.float32)
|
||||||
|
else:
|
||||||
|
assert BG.shape[0] == h and BG.shape[1] == w and BG.shape[2] == c
|
||||||
|
image = BG.astype(np.float32)
|
||||||
|
|
||||||
|
depth_buffer = np.zeros([h, w], dtype = np.float32, order = 'C') - 999999.
|
||||||
|
|
||||||
|
tex_h, tex_w, tex_c = texture.shape
|
||||||
|
if mapping_type == 'nearest':
|
||||||
|
mt = int(0)
|
||||||
|
elif mapping_type == 'bilinear':
|
||||||
|
mt = int(1)
|
||||||
|
else:
|
||||||
|
mt = int(0)
|
||||||
|
|
||||||
|
# -> C order
|
||||||
|
vertices = vertices.astype(np.float32).copy()
|
||||||
|
triangles = triangles.astype(np.int32).copy()
|
||||||
|
texture = texture.astype(np.float32).copy()
|
||||||
|
tex_coords = tex_coords.astype(np.float32).copy()
|
||||||
|
tex_triangles = tex_triangles.astype(np.int32).copy()
|
||||||
|
|
||||||
|
mesh_core_cython.render_texture_core(
|
||||||
|
image, vertices, triangles,
|
||||||
|
texture, tex_coords, tex_triangles,
|
||||||
|
depth_buffer,
|
||||||
|
vertices.shape[0], tex_coords.shape[0], triangles.shape[0],
|
||||||
|
h, w, c,
|
||||||
|
tex_h, tex_w, tex_c,
|
||||||
|
mt)
|
||||||
|
return image
|
||||||
|
|
||||||
@@ -0,0 +1,383 @@
|
|||||||
|
'''
|
||||||
|
Functions about transforming mesh(changing the position: modify vertices).
|
||||||
|
1. forward: transform(transform, camera, project).
|
||||||
|
2. backward: estimate transform matrix from correspondences.
|
||||||
|
|
||||||
|
Author: Yao Feng
|
||||||
|
Mail: yaofeng1995@gmail.com
|
||||||
|
'''
|
||||||
|
|
||||||
|
from __future__ import absolute_import
|
||||||
|
from __future__ import division
|
||||||
|
from __future__ import print_function
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import math
|
||||||
|
from math import cos, sin
|
||||||
|
|
||||||
|
def angle2matrix(angles):
|
||||||
|
''' get rotation matrix from three rotation angles(degree). right-handed.
|
||||||
|
Args:
|
||||||
|
angles: [3,]. x, y, z angles
|
||||||
|
x: pitch. positive for looking down.
|
||||||
|
y: yaw. positive for looking left.
|
||||||
|
z: roll. positive for tilting head right.
|
||||||
|
Returns:
|
||||||
|
R: [3, 3]. rotation matrix.
|
||||||
|
'''
|
||||||
|
x, y, z = np.deg2rad(angles[0]), np.deg2rad(angles[1]), np.deg2rad(angles[2])
|
||||||
|
# x
|
||||||
|
Rx=np.array([[1, 0, 0],
|
||||||
|
[0, cos(x), -sin(x)],
|
||||||
|
[0, sin(x), cos(x)]])
|
||||||
|
# y
|
||||||
|
Ry=np.array([[ cos(y), 0, sin(y)],
|
||||||
|
[ 0, 1, 0],
|
||||||
|
[-sin(y), 0, cos(y)]])
|
||||||
|
# z
|
||||||
|
Rz=np.array([[cos(z), -sin(z), 0],
|
||||||
|
[sin(z), cos(z), 0],
|
||||||
|
[ 0, 0, 1]])
|
||||||
|
|
||||||
|
R=Rz.dot(Ry.dot(Rx))
|
||||||
|
return R.astype(np.float32)
|
||||||
|
|
||||||
|
def angle2matrix_3ddfa(angles):
|
||||||
|
''' get rotation matrix from three rotation angles(radian). The same as in 3DDFA.
|
||||||
|
Args:
|
||||||
|
angles: [3,]. x, y, z angles
|
||||||
|
x: pitch.
|
||||||
|
y: yaw.
|
||||||
|
z: roll.
|
||||||
|
Returns:
|
||||||
|
R: 3x3. rotation matrix.
|
||||||
|
'''
|
||||||
|
# x, y, z = np.deg2rad(angles[0]), np.deg2rad(angles[1]), np.deg2rad(angles[2])
|
||||||
|
x, y, z = angles[0], angles[1], angles[2]
|
||||||
|
|
||||||
|
# x
|
||||||
|
Rx=np.array([[1, 0, 0],
|
||||||
|
[0, cos(x), sin(x)],
|
||||||
|
[0, -sin(x), cos(x)]])
|
||||||
|
# y
|
||||||
|
Ry=np.array([[ cos(y), 0, -sin(y)],
|
||||||
|
[ 0, 1, 0],
|
||||||
|
[sin(y), 0, cos(y)]])
|
||||||
|
# z
|
||||||
|
Rz=np.array([[cos(z), sin(z), 0],
|
||||||
|
[-sin(z), cos(z), 0],
|
||||||
|
[ 0, 0, 1]])
|
||||||
|
R = Rx.dot(Ry).dot(Rz)
|
||||||
|
return R.astype(np.float32)
|
||||||
|
|
||||||
|
|
||||||
|
## ------------------------------------------ 1. transform(transform, project, camera).
|
||||||
|
## ---------- 3d-3d transform. Transform obj in world space
|
||||||
|
def rotate(vertices, angles):
|
||||||
|
''' rotate vertices.
|
||||||
|
X_new = R.dot(X). X: 3 x 1
|
||||||
|
Args:
|
||||||
|
vertices: [nver, 3].
|
||||||
|
rx, ry, rz: degree angles
|
||||||
|
rx: pitch. positive for looking down
|
||||||
|
ry: yaw. positive for looking left
|
||||||
|
rz: roll. positive for tilting head right
|
||||||
|
Returns:
|
||||||
|
rotated vertices: [nver, 3]
|
||||||
|
'''
|
||||||
|
R = angle2matrix(angles)
|
||||||
|
rotated_vertices = vertices.dot(R.T)
|
||||||
|
|
||||||
|
return rotated_vertices
|
||||||
|
|
||||||
|
def similarity_transform(vertices, s, R, t3d):
|
||||||
|
''' similarity transform. dof = 7.
|
||||||
|
3D: s*R.dot(X) + t
|
||||||
|
Homo: M = [[sR, t],[0^T, 1]]. M.dot(X)
|
||||||
|
Args:(float32)
|
||||||
|
vertices: [nver, 3].
|
||||||
|
s: [1,]. scale factor.
|
||||||
|
R: [3,3]. rotation matrix.
|
||||||
|
t3d: [3,]. 3d translation vector.
|
||||||
|
Returns:
|
||||||
|
transformed vertices: [nver, 3]
|
||||||
|
'''
|
||||||
|
t3d = np.squeeze(np.array(t3d, dtype = np.float32))
|
||||||
|
transformed_vertices = s * vertices.dot(R.T) + t3d[np.newaxis, :]
|
||||||
|
|
||||||
|
return transformed_vertices
|
||||||
|
|
||||||
|
|
||||||
|
## -------------- Camera. from world space to camera space
|
||||||
|
# Ref: https://cs184.eecs.berkeley.edu/lecture/transforms-2
|
||||||
|
def normalize(x):
|
||||||
|
epsilon = 1e-12
|
||||||
|
norm = np.sqrt(np.sum(x**2, axis = 0))
|
||||||
|
norm = np.maximum(norm, epsilon)
|
||||||
|
return x/norm
|
||||||
|
|
||||||
|
def lookat_camera(vertices, eye, at = None, up = None):
|
||||||
|
""" 'look at' transformation: from world space to camera space
|
||||||
|
standard camera space:
|
||||||
|
camera located at the origin.
|
||||||
|
looking down negative z-axis.
|
||||||
|
vertical vector is y-axis.
|
||||||
|
Xcam = R(X - C)
|
||||||
|
Homo: [[R, -RC], [0, 1]]
|
||||||
|
Args:
|
||||||
|
vertices: [nver, 3]
|
||||||
|
eye: [3,] the XYZ world space position of the camera.
|
||||||
|
at: [3,] a position along the center of the camera's gaze.
|
||||||
|
up: [3,] up direction
|
||||||
|
Returns:
|
||||||
|
transformed_vertices: [nver, 3]
|
||||||
|
"""
|
||||||
|
if at is None:
|
||||||
|
at = np.array([0, 0, 0], np.float32)
|
||||||
|
if up is None:
|
||||||
|
up = np.array([0, 1, 0], np.float32)
|
||||||
|
|
||||||
|
eye = np.array(eye).astype(np.float32)
|
||||||
|
at = np.array(at).astype(np.float32)
|
||||||
|
z_aixs = -normalize(at - eye) # look forward
|
||||||
|
x_aixs = normalize(np.cross(up, z_aixs)) # look right
|
||||||
|
y_axis = np.cross(z_aixs, x_aixs) # look up
|
||||||
|
|
||||||
|
R = np.stack((x_aixs, y_axis, z_aixs))#, axis = 0) # 3 x 3
|
||||||
|
transformed_vertices = vertices - eye # translation
|
||||||
|
transformed_vertices = transformed_vertices.dot(R.T) # rotation
|
||||||
|
return transformed_vertices
|
||||||
|
|
||||||
|
## --------- 3d-2d project. from camera space to image plane
|
||||||
|
# generally, image plane only keeps x,y channels, here reserve z channel for calculating z-buffer.
|
||||||
|
def orthographic_project(vertices):
|
||||||
|
''' scaled orthographic projection(just delete z)
|
||||||
|
assumes: variations in depth over the object is small relative to the mean distance from camera to object
|
||||||
|
x -> x*f/z, y -> x*f/z, z -> f.
|
||||||
|
for point i,j. zi~=zj. so just delete z
|
||||||
|
** often used in face
|
||||||
|
Homo: P = [[1,0,0,0], [0,1,0,0], [0,0,1,0]]
|
||||||
|
Args:
|
||||||
|
vertices: [nver, 3]
|
||||||
|
Returns:
|
||||||
|
projected_vertices: [nver, 3] if isKeepZ=True. [nver, 2] if isKeepZ=False.
|
||||||
|
'''
|
||||||
|
return vertices.copy()
|
||||||
|
|
||||||
|
def perspective_project(vertices, fovy, aspect_ratio = 1., near = 0.1, far = 1000.):
|
||||||
|
''' perspective projection.
|
||||||
|
Args:
|
||||||
|
vertices: [nver, 3]
|
||||||
|
fovy: vertical angular field of view. degree.
|
||||||
|
aspect_ratio : width / height of field of view
|
||||||
|
near : depth of near clipping plane
|
||||||
|
far : depth of far clipping plane
|
||||||
|
Returns:
|
||||||
|
projected_vertices: [nver, 3]
|
||||||
|
'''
|
||||||
|
fovy = np.deg2rad(fovy)
|
||||||
|
top = near*np.tan(fovy)
|
||||||
|
bottom = -top
|
||||||
|
right = top*aspect_ratio
|
||||||
|
left = -right
|
||||||
|
|
||||||
|
#-- homo
|
||||||
|
P = np.array([[near/right, 0, 0, 0],
|
||||||
|
[0, near/top, 0, 0],
|
||||||
|
[0, 0, -(far+near)/(far-near), -2*far*near/(far-near)],
|
||||||
|
[0, 0, -1, 0]])
|
||||||
|
vertices_homo = np.hstack((vertices, np.ones((vertices.shape[0], 1)))) # [nver, 4]
|
||||||
|
projected_vertices = vertices_homo.dot(P.T)
|
||||||
|
projected_vertices = projected_vertices/projected_vertices[:,3:]
|
||||||
|
projected_vertices = projected_vertices[:,:3]
|
||||||
|
projected_vertices[:,2] = -projected_vertices[:,2]
|
||||||
|
|
||||||
|
#-- non homo. only fovy
|
||||||
|
# projected_vertices = vertices.copy()
|
||||||
|
# projected_vertices[:,0] = -(near/right)*vertices[:,0]/vertices[:,2]
|
||||||
|
# projected_vertices[:,1] = -(near/top)*vertices[:,1]/vertices[:,2]
|
||||||
|
return projected_vertices
|
||||||
|
|
||||||
|
|
||||||
|
def to_image(vertices, h, w, is_perspective = False):
|
||||||
|
''' change vertices to image coord system
|
||||||
|
3d system: XYZ, center(0, 0, 0)
|
||||||
|
2d image: x(u), y(v). center(w/2, h/2), flip y-axis.
|
||||||
|
Args:
|
||||||
|
vertices: [nver, 3]
|
||||||
|
h: height of the rendering
|
||||||
|
w : width of the rendering
|
||||||
|
Returns:
|
||||||
|
projected_vertices: [nver, 3]
|
||||||
|
'''
|
||||||
|
image_vertices = vertices.copy()
|
||||||
|
if is_perspective:
|
||||||
|
# if perspective, the projected vertices are normalized to [-1, 1]. so change it to image size first.
|
||||||
|
image_vertices[:,0] = image_vertices[:,0]*w/2
|
||||||
|
image_vertices[:,1] = image_vertices[:,1]*h/2
|
||||||
|
# move to center of image
|
||||||
|
image_vertices[:,0] = image_vertices[:,0] + w/2
|
||||||
|
image_vertices[:,1] = image_vertices[:,1] + h/2
|
||||||
|
# flip vertices along y-axis.
|
||||||
|
image_vertices[:,1] = h - image_vertices[:,1] - 1
|
||||||
|
return image_vertices
|
||||||
|
|
||||||
|
|
||||||
|
#### -------------------------------------------2. estimate transform matrix from correspondences.
|
||||||
|
def estimate_affine_matrix_3d23d(X, Y):
|
||||||
|
''' Using least-squares solution
|
||||||
|
Args:
|
||||||
|
X: [n, 3]. 3d points(fixed)
|
||||||
|
Y: [n, 3]. corresponding 3d points(moving). Y = PX
|
||||||
|
Returns:
|
||||||
|
P_Affine: (3, 4). Affine camera matrix (the third row is [0, 0, 0, 1]).
|
||||||
|
'''
|
||||||
|
X_homo = np.hstack((X, np.ones([X.shape[1],1]))) #n x 4
|
||||||
|
P = np.linalg.lstsq(X_homo, Y)[0].T # Affine matrix. 3 x 4
|
||||||
|
return P
|
||||||
|
|
||||||
|
def estimate_affine_matrix_3d22d(X, x):
|
||||||
|
''' Using Golden Standard Algorithm for estimating an affine camera
|
||||||
|
matrix P from world to image correspondences.
|
||||||
|
See Alg.7.2. in MVGCV
|
||||||
|
Code Ref: https://github.com/patrikhuber/eos/blob/master/include/eos/fitting/affine_camera_estimation.hpp
|
||||||
|
x_homo = X_homo.dot(P_Affine)
|
||||||
|
Args:
|
||||||
|
X: [n, 3]. corresponding 3d points(fixed)
|
||||||
|
x: [n, 2]. n>=4. 2d points(moving). x = PX
|
||||||
|
Returns:
|
||||||
|
P_Affine: [3, 4]. Affine camera matrix
|
||||||
|
'''
|
||||||
|
X = X.T; x = x.T
|
||||||
|
assert(x.shape[1] == X.shape[1])
|
||||||
|
n = x.shape[1]
|
||||||
|
assert(n >= 4)
|
||||||
|
|
||||||
|
#--- 1. normalization
|
||||||
|
# 2d points
|
||||||
|
mean = np.mean(x, 1) # (2,)
|
||||||
|
x = x - np.tile(mean[:, np.newaxis], [1, n])
|
||||||
|
average_norm = np.mean(np.sqrt(np.sum(x**2, 0)))
|
||||||
|
scale = np.sqrt(2) / average_norm
|
||||||
|
x = scale * x
|
||||||
|
|
||||||
|
T = np.zeros((3,3), dtype = np.float32)
|
||||||
|
T[0, 0] = T[1, 1] = scale
|
||||||
|
T[:2, 2] = -mean*scale
|
||||||
|
T[2, 2] = 1
|
||||||
|
|
||||||
|
# 3d points
|
||||||
|
X_homo = np.vstack((X, np.ones((1, n))))
|
||||||
|
mean = np.mean(X, 1) # (3,)
|
||||||
|
X = X - np.tile(mean[:, np.newaxis], [1, n])
|
||||||
|
m = X_homo[:3,:] - X
|
||||||
|
average_norm = np.mean(np.sqrt(np.sum(X**2, 0)))
|
||||||
|
scale = np.sqrt(3) / average_norm
|
||||||
|
X = scale * X
|
||||||
|
|
||||||
|
U = np.zeros((4,4), dtype = np.float32)
|
||||||
|
U[0, 0] = U[1, 1] = U[2, 2] = scale
|
||||||
|
U[:3, 3] = -mean*scale
|
||||||
|
U[3, 3] = 1
|
||||||
|
|
||||||
|
# --- 2. equations
|
||||||
|
A = np.zeros((n*2, 8), dtype = np.float32);
|
||||||
|
X_homo = np.vstack((X, np.ones((1, n)))).T
|
||||||
|
A[:n, :4] = X_homo
|
||||||
|
A[n:, 4:] = X_homo
|
||||||
|
b = np.reshape(x, [-1, 1])
|
||||||
|
|
||||||
|
# --- 3. solution
|
||||||
|
p_8 = np.linalg.pinv(A).dot(b)
|
||||||
|
P = np.zeros((3, 4), dtype = np.float32)
|
||||||
|
P[0, :] = p_8[:4, 0]
|
||||||
|
P[1, :] = p_8[4:, 0]
|
||||||
|
P[-1, -1] = 1
|
||||||
|
|
||||||
|
# --- 4. denormalization
|
||||||
|
P_Affine = np.linalg.inv(T).dot(P.dot(U))
|
||||||
|
return P_Affine
|
||||||
|
|
||||||
|
def P2sRt(P):
|
||||||
|
''' decompositing camera matrix P
|
||||||
|
Args:
|
||||||
|
P: (3, 4). Affine Camera Matrix.
|
||||||
|
Returns:
|
||||||
|
s: scale factor.
|
||||||
|
R: (3, 3). rotation matrix.
|
||||||
|
t: (3,). translation.
|
||||||
|
'''
|
||||||
|
t = P[:, 3]
|
||||||
|
R1 = P[0:1, :3]
|
||||||
|
R2 = P[1:2, :3]
|
||||||
|
s = (np.linalg.norm(R1) + np.linalg.norm(R2))/2.0
|
||||||
|
r1 = R1/np.linalg.norm(R1)
|
||||||
|
r2 = R2/np.linalg.norm(R2)
|
||||||
|
r3 = np.cross(r1, r2)
|
||||||
|
|
||||||
|
R = np.concatenate((r1, r2, r3), 0)
|
||||||
|
return s, R, t
|
||||||
|
|
||||||
|
#Ref: https://www.learnopencv.com/rotation-matrix-to-euler-angles/
|
||||||
|
def isRotationMatrix(R):
|
||||||
|
''' checks if a matrix is a valid rotation matrix(whether orthogonal or not)
|
||||||
|
'''
|
||||||
|
Rt = np.transpose(R)
|
||||||
|
shouldBeIdentity = np.dot(Rt, R)
|
||||||
|
I = np.identity(3, dtype = R.dtype)
|
||||||
|
n = np.linalg.norm(I - shouldBeIdentity)
|
||||||
|
return n < 1e-6
|
||||||
|
|
||||||
|
def matrix2angle(R):
|
||||||
|
''' get three Euler angles from Rotation Matrix
|
||||||
|
Args:
|
||||||
|
R: (3,3). rotation matrix
|
||||||
|
Returns:
|
||||||
|
x: pitch
|
||||||
|
y: yaw
|
||||||
|
z: roll
|
||||||
|
'''
|
||||||
|
assert(isRotationMatrix)
|
||||||
|
sy = math.sqrt(R[0,0] * R[0,0] + R[1,0] * R[1,0])
|
||||||
|
|
||||||
|
singular = sy < 1e-6
|
||||||
|
|
||||||
|
if not singular :
|
||||||
|
x = math.atan2(R[2,1] , R[2,2])
|
||||||
|
y = math.atan2(-R[2,0], sy)
|
||||||
|
z = math.atan2(R[1,0], R[0,0])
|
||||||
|
else :
|
||||||
|
x = math.atan2(-R[1,2], R[1,1])
|
||||||
|
y = math.atan2(-R[2,0], sy)
|
||||||
|
z = 0
|
||||||
|
|
||||||
|
# rx, ry, rz = np.rad2deg(x), np.rad2deg(y), np.rad2deg(z)
|
||||||
|
rx, ry, rz = x*180/np.pi, y*180/np.pi, z*180/np.pi
|
||||||
|
return rx, ry, rz
|
||||||
|
|
||||||
|
# def matrix2angle(R):
|
||||||
|
# ''' compute three Euler angles from a Rotation Matrix. Ref: http://www.gregslabaugh.net/publications/euler.pdf
|
||||||
|
# Args:
|
||||||
|
# R: (3,3). rotation matrix
|
||||||
|
# Returns:
|
||||||
|
# x: yaw
|
||||||
|
# y: pitch
|
||||||
|
# z: roll
|
||||||
|
# '''
|
||||||
|
# # assert(isRotationMatrix(R))
|
||||||
|
|
||||||
|
# if R[2,0] !=1 or R[2,0] != -1:
|
||||||
|
# x = math.asin(R[2,0])
|
||||||
|
# y = math.atan2(R[2,1]/cos(x), R[2,2]/cos(x))
|
||||||
|
# z = math.atan2(R[1,0]/cos(x), R[0,0]/cos(x))
|
||||||
|
|
||||||
|
# else:# Gimbal lock
|
||||||
|
# z = 0 #can be anything
|
||||||
|
# if R[2,0] == -1:
|
||||||
|
# x = np.pi/2
|
||||||
|
# y = z + math.atan2(R[0,1], R[0,2])
|
||||||
|
# else:
|
||||||
|
# x = -np.pi/2
|
||||||
|
# y = -z + math.atan2(-R[0,1], -R[0,2])
|
||||||
|
|
||||||
|
# return x, y, z
|
||||||
@@ -0,0 +1,24 @@
|
|||||||
|
from __future__ import absolute_import
|
||||||
|
from __future__ import division
|
||||||
|
from __future__ import print_function
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import matplotlib.pyplot as plt
|
||||||
|
# from skimage import measure
|
||||||
|
from mpl_toolkits.mplot3d import Axes3D
|
||||||
|
|
||||||
|
def plot_mesh(vertices, triangles, subplot = [1,1,1], title = 'mesh', el = 90, az = -90, lwdt=.1, dist = 6, color = "grey"):
|
||||||
|
'''
|
||||||
|
plot the mesh
|
||||||
|
Args:
|
||||||
|
vertices: [nver, 3]
|
||||||
|
triangles: [ntri, 3]
|
||||||
|
'''
|
||||||
|
ax = plt.subplot(subplot[0], subplot[1], subplot[2], projection = '3d')
|
||||||
|
ax.plot_trisurf(vertices[:, 0], vertices[:, 1], vertices[:, 2], triangles = triangles, lw = lwdt, color = color, alpha = 1)
|
||||||
|
ax.axis("off")
|
||||||
|
ax.view_init(elev = el, azim = az)
|
||||||
|
ax.dist = dist
|
||||||
|
plt.title(title)
|
||||||
|
|
||||||
|
### -------------- Todo: use vtk to visualize mesh? or visvis? or VisPy?
|
||||||
@@ -0,0 +1,7 @@
|
|||||||
|
from __future__ import absolute_import
|
||||||
|
from __future__ import division
|
||||||
|
from __future__ import print_function
|
||||||
|
|
||||||
|
from .. import mesh
|
||||||
|
from .morphabel_model import MorphabelModel
|
||||||
|
from . import load
|
||||||
@@ -0,0 +1,270 @@
|
|||||||
|
'''
|
||||||
|
Estimating parameters about vertices: shape para, exp para, pose para(s, R, t)
|
||||||
|
'''
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
''' TODO: a clear document.
|
||||||
|
Given: image_points, 3D Model, Camera Matrix(s, R, t2d)
|
||||||
|
Estimate: shape parameters, expression parameters
|
||||||
|
|
||||||
|
Inference:
|
||||||
|
|
||||||
|
projected_vertices = s*P*R(mu + shape + exp) + t2d --> image_points
|
||||||
|
s*P*R*shape + s*P*R(mu + exp) + t2d --> image_poitns
|
||||||
|
|
||||||
|
# Define:
|
||||||
|
X = vertices
|
||||||
|
x_hat = projected_vertices
|
||||||
|
x = image_points
|
||||||
|
A = s*P*R
|
||||||
|
b = s*P*R(mu + exp) + t2d
|
||||||
|
==>
|
||||||
|
x_hat = A*shape + b (2 x n)
|
||||||
|
|
||||||
|
A*shape (2 x n)
|
||||||
|
shape = reshape(shapePC * sp) (3 x n)
|
||||||
|
shapePC*sp : (3n x 1)
|
||||||
|
|
||||||
|
* flatten:
|
||||||
|
x_hat_flatten = A*shape + b_flatten (2n x 1)
|
||||||
|
A*shape (2n x 1)
|
||||||
|
--> A*shapePC (2n x 199) sp: 199 x 1
|
||||||
|
|
||||||
|
# Define:
|
||||||
|
pc_2d = A* reshape(shapePC)
|
||||||
|
pc_2d_flatten = flatten(pc_2d) (2n x 199)
|
||||||
|
|
||||||
|
=====>
|
||||||
|
x_hat_flatten = pc_2d_flatten * sp + b_flatten ---> x_flatten (2n x 1)
|
||||||
|
|
||||||
|
Goals:
|
||||||
|
(ignore flatten, pc_2d-->pc)
|
||||||
|
min E = || x_hat - x || + lambda*sum(sp/sigma)^2
|
||||||
|
= || pc * sp + b - x || + lambda*sum(sp/sigma)^2
|
||||||
|
|
||||||
|
Solve:
|
||||||
|
d(E)/d(sp) = 0
|
||||||
|
2 * pc' * (pc * sp + b - x) + 2 * lambda * sp / (sigma' * sigma) = 0
|
||||||
|
|
||||||
|
Get:
|
||||||
|
(pc' * pc + lambda / (sigma'* sigma)) * sp = pc' * (x - b)
|
||||||
|
|
||||||
|
'''
|
||||||
|
|
||||||
|
def estimate_shape(x, shapeMU, shapePC, shapeEV, expression, s, R, t2d, lamb = 3000):
|
||||||
|
'''
|
||||||
|
Args:
|
||||||
|
x: (2, n). image points (to be fitted)
|
||||||
|
shapeMU: (3n, 1)
|
||||||
|
shapePC: (3n, n_sp)
|
||||||
|
shapeEV: (n_sp, 1)
|
||||||
|
expression: (3, n)
|
||||||
|
s: scale
|
||||||
|
R: (3, 3). rotation matrix
|
||||||
|
t2d: (2,). 2d translation
|
||||||
|
lambda: regulation coefficient
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
shape_para: (n_sp, 1) shape parameters(coefficients)
|
||||||
|
'''
|
||||||
|
x = x.copy()
|
||||||
|
assert(shapeMU.shape[0] == shapePC.shape[0])
|
||||||
|
assert(shapeMU.shape[0] == x.shape[1]*3)
|
||||||
|
|
||||||
|
dof = shapePC.shape[1]
|
||||||
|
|
||||||
|
n = x.shape[1]
|
||||||
|
sigma = shapeEV
|
||||||
|
t2d = np.array(t2d)
|
||||||
|
P = np.array([[1, 0, 0], [0, 1, 0]], dtype = np.float32)
|
||||||
|
A = s*P.dot(R)
|
||||||
|
|
||||||
|
# --- calc pc
|
||||||
|
pc_3d = np.resize(shapePC.T, [dof, n, 3]) # 199 x n x 3
|
||||||
|
pc_3d = np.reshape(pc_3d, [dof*n, 3])
|
||||||
|
pc_2d = pc_3d.dot(A.T.copy()) # 199 x n x 2
|
||||||
|
|
||||||
|
pc = np.reshape(pc_2d, [dof, -1]).T # 2n x 199
|
||||||
|
|
||||||
|
# --- calc b
|
||||||
|
# shapeMU
|
||||||
|
mu_3d = np.resize(shapeMU, [n, 3]).T # 3 x n
|
||||||
|
# expression
|
||||||
|
exp_3d = expression
|
||||||
|
#
|
||||||
|
b = A.dot(mu_3d + exp_3d) + np.tile(t2d[:, np.newaxis], [1, n]) # 2 x n
|
||||||
|
b = np.reshape(b.T, [-1, 1]) # 2n x 1
|
||||||
|
|
||||||
|
# --- solve
|
||||||
|
equation_left = np.dot(pc.T, pc) + lamb * np.diagflat(1/sigma**2)
|
||||||
|
x = np.reshape(x.T, [-1, 1])
|
||||||
|
equation_right = np.dot(pc.T, x - b)
|
||||||
|
|
||||||
|
shape_para = np.dot(np.linalg.inv(equation_left), equation_right)
|
||||||
|
|
||||||
|
return shape_para
|
||||||
|
|
||||||
|
def estimate_expression(x, shapeMU, expPC, expEV, shape, s, R, t2d, lamb = 2000):
|
||||||
|
'''
|
||||||
|
Args:
|
||||||
|
x: (2, n). image points (to be fitted)
|
||||||
|
shapeMU: (3n, 1)
|
||||||
|
expPC: (3n, n_ep)
|
||||||
|
expEV: (n_ep, 1)
|
||||||
|
shape: (3, n)
|
||||||
|
s: scale
|
||||||
|
R: (3, 3). rotation matrix
|
||||||
|
t2d: (2,). 2d translation
|
||||||
|
lambda: regulation coefficient
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
exp_para: (n_ep, 1) shape parameters(coefficients)
|
||||||
|
'''
|
||||||
|
x = x.copy()
|
||||||
|
assert(shapeMU.shape[0] == expPC.shape[0])
|
||||||
|
assert(shapeMU.shape[0] == x.shape[1]*3)
|
||||||
|
|
||||||
|
dof = expPC.shape[1]
|
||||||
|
|
||||||
|
n = x.shape[1]
|
||||||
|
sigma = expEV
|
||||||
|
t2d = np.array(t2d)
|
||||||
|
P = np.array([[1, 0, 0], [0, 1, 0]], dtype = np.float32)
|
||||||
|
A = s*P.dot(R)
|
||||||
|
|
||||||
|
# --- calc pc
|
||||||
|
pc_3d = np.resize(expPC.T, [dof, n, 3])
|
||||||
|
pc_3d = np.reshape(pc_3d, [dof*n, 3])
|
||||||
|
pc_2d = pc_3d.dot(A.T)
|
||||||
|
pc = np.reshape(pc_2d, [dof, -1]).T # 2n x 29
|
||||||
|
|
||||||
|
# --- calc b
|
||||||
|
# shapeMU
|
||||||
|
mu_3d = np.resize(shapeMU, [n, 3]).T # 3 x n
|
||||||
|
# expression
|
||||||
|
shape_3d = shape
|
||||||
|
#
|
||||||
|
b = A.dot(mu_3d + shape_3d) + np.tile(t2d[:, np.newaxis], [1, n]) # 2 x n
|
||||||
|
b = np.reshape(b.T, [-1, 1]) # 2n x 1
|
||||||
|
|
||||||
|
# --- solve
|
||||||
|
equation_left = np.dot(pc.T, pc) + lamb * np.diagflat(1/sigma**2)
|
||||||
|
x = np.reshape(x.T, [-1, 1])
|
||||||
|
equation_right = np.dot(pc.T, x - b)
|
||||||
|
|
||||||
|
exp_para = np.dot(np.linalg.inv(equation_left), equation_right)
|
||||||
|
|
||||||
|
return exp_para
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------- fit
|
||||||
|
def fit_points(x, X_ind, model, n_sp, n_ep, max_iter = 4):
|
||||||
|
'''
|
||||||
|
Args:
|
||||||
|
x: (n, 2) image points
|
||||||
|
X_ind: (n,) corresponding Model vertex indices
|
||||||
|
model: 3DMM
|
||||||
|
max_iter: iteration
|
||||||
|
Returns:
|
||||||
|
sp: (n_sp, 1). shape parameters
|
||||||
|
ep: (n_ep, 1). exp parameters
|
||||||
|
s, R, t
|
||||||
|
'''
|
||||||
|
x = x.copy().T
|
||||||
|
|
||||||
|
#-- init
|
||||||
|
sp = np.zeros((n_sp, 1), dtype = np.float32)
|
||||||
|
ep = np.zeros((n_ep, 1), dtype = np.float32)
|
||||||
|
|
||||||
|
#-------------------- estimate
|
||||||
|
X_ind_all = np.tile(X_ind[np.newaxis, :], [3, 1])*3
|
||||||
|
X_ind_all[1, :] += 1
|
||||||
|
X_ind_all[2, :] += 2
|
||||||
|
valid_ind = X_ind_all.flatten('F')
|
||||||
|
|
||||||
|
shapeMU = model['shapeMU'][valid_ind, :]
|
||||||
|
shapePC = model['shapePC'][valid_ind, :n_sp]
|
||||||
|
expPC = model['expPC'][valid_ind, :n_ep]
|
||||||
|
|
||||||
|
for i in range(max_iter):
|
||||||
|
X = shapeMU + shapePC.dot(sp) + expPC.dot(ep)
|
||||||
|
X = np.reshape(X, [int(len(X)/3), 3]).T
|
||||||
|
|
||||||
|
#----- estimate pose
|
||||||
|
P = Face.face3d.mesh.transform.estimate_affine_matrix_3d22d(X.T, x.T)
|
||||||
|
s, R, t = Face.face3d.mesh.transform.P2sRt(P)
|
||||||
|
rx, ry, rz = Face.face3d.mesh.transform.matrix2angle(R)
|
||||||
|
# print('Iter:{}; estimated pose: s {}, rx {}, ry {}, rz {}, t1 {}, t2 {}'.format(i, s, rx, ry, rz, t[0], t[1]))
|
||||||
|
|
||||||
|
#----- estimate shape
|
||||||
|
# expression
|
||||||
|
shape = shapePC.dot(sp)
|
||||||
|
shape = np.reshape(shape, [int(len(shape)/3), 3]).T
|
||||||
|
ep = estimate_expression(x, shapeMU, expPC, model['expEV'][:n_ep,:], shape, s, R, t[:2], lamb = 20)
|
||||||
|
|
||||||
|
# shape
|
||||||
|
expression = expPC.dot(ep)
|
||||||
|
expression = np.reshape(expression, [int(len(expression)/3), 3]).T
|
||||||
|
sp = estimate_shape(x, shapeMU, shapePC, model['shapeEV'][:n_sp,:], expression, s, R, t[:2], lamb = 40)
|
||||||
|
|
||||||
|
return sp, ep, s, R, t
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------- fitting process
|
||||||
|
def fit_points_for_show(x, X_ind, model, n_sp, n_ep, max_iter = 4):
|
||||||
|
'''
|
||||||
|
Args:
|
||||||
|
x: (n, 2) image points
|
||||||
|
X_ind: (n,) corresponding Model vertex indices
|
||||||
|
model: 3DMM
|
||||||
|
max_iter: iteration
|
||||||
|
Returns:
|
||||||
|
sp: (n_sp, 1). shape parameters
|
||||||
|
ep: (n_ep, 1). exp parameters
|
||||||
|
s, R, t
|
||||||
|
'''
|
||||||
|
x = x.copy().T
|
||||||
|
|
||||||
|
#-- init
|
||||||
|
sp = np.zeros((n_sp, 1), dtype = np.float32)
|
||||||
|
ep = np.zeros((n_ep, 1), dtype = np.float32)
|
||||||
|
|
||||||
|
#-------------------- estimate
|
||||||
|
X_ind_all = np.tile(X_ind[np.newaxis, :], [3, 1])*3
|
||||||
|
X_ind_all[1, :] += 1
|
||||||
|
X_ind_all[2, :] += 2
|
||||||
|
valid_ind = X_ind_all.flatten('F')
|
||||||
|
|
||||||
|
shapeMU = model['shapeMU'][valid_ind, :]
|
||||||
|
shapePC = model['shapePC'][valid_ind, :n_sp]
|
||||||
|
expPC = model['expPC'][valid_ind, :n_ep]
|
||||||
|
|
||||||
|
s = 4e-04
|
||||||
|
R = Face.face3d.mesh.transform.angle2matrix([0, 0, 0])
|
||||||
|
t = [0, 0, 0]
|
||||||
|
lsp = []; lep = []; ls = []; lR = []; lt = []
|
||||||
|
for i in range(max_iter):
|
||||||
|
X = shapeMU + shapePC.dot(sp) + expPC.dot(ep)
|
||||||
|
X = np.reshape(X, [int(len(X)/3), 3]).T
|
||||||
|
lsp.append(sp); lep.append(ep); ls.append(s), lR.append(R), lt.append(t)
|
||||||
|
|
||||||
|
#----- estimate pose
|
||||||
|
P = Face.face3d.mesh.transform.estimate_affine_matrix_3d22d(X.T, x.T)
|
||||||
|
s, R, t = Face.face3d.mesh.transform.P2sRt(P)
|
||||||
|
lsp.append(sp); lep.append(ep); ls.append(s), lR.append(R), lt.append(t)
|
||||||
|
|
||||||
|
#----- estimate shape
|
||||||
|
# expression
|
||||||
|
shape = shapePC.dot(sp)
|
||||||
|
shape = np.reshape(shape, [int(len(shape)/3), 3]).T
|
||||||
|
ep = estimate_expression(x, shapeMU, expPC, model['expEV'][:n_ep,:], shape, s, R, t[:2], lamb = 20)
|
||||||
|
lsp.append(sp); lep.append(ep); ls.append(s), lR.append(R), lt.append(t)
|
||||||
|
|
||||||
|
# shape
|
||||||
|
expression = expPC.dot(ep)
|
||||||
|
expression = np.reshape(expression, [int(len(expression)/3), 3]).T
|
||||||
|
sp = estimate_shape(x, shapeMU, shapePC, model['shapeEV'][:n_sp,:], expression, s, R, t[:2], lamb = 40)
|
||||||
|
|
||||||
|
# print('ls', ls)
|
||||||
|
# print('lR', lR)
|
||||||
|
return np.array(lsp), np.array(lep), np.array(ls), np.array(lR), np.array(lt)
|
||||||
@@ -0,0 +1,110 @@
|
|||||||
|
from __future__ import absolute_import
|
||||||
|
from __future__ import division
|
||||||
|
from __future__ import print_function
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import scipy.io as sio
|
||||||
|
|
||||||
|
### --------------------------------- load BFM data
|
||||||
|
def load_BFM(model_path):
|
||||||
|
''' load BFM 3DMM model
|
||||||
|
Args:
|
||||||
|
model_path: path to BFM model.
|
||||||
|
Returns:
|
||||||
|
model: (nver = 53215, ntri = 105840). nver: number of vertices. ntri: number of triangles.
|
||||||
|
'shapeMU': [3*nver, 1]
|
||||||
|
'shapePC': [3*nver, 199]
|
||||||
|
'shapeEV': [199, 1]
|
||||||
|
'expMU': [3*nver, 1]
|
||||||
|
'expPC': [3*nver, 29]
|
||||||
|
'expEV': [29, 1]
|
||||||
|
'texMU': [3*nver, 1]
|
||||||
|
'texPC': [3*nver, 199]
|
||||||
|
'texEV': [199, 1]
|
||||||
|
'tri': [ntri, 3] (start from 1, should sub 1 in python and c++)
|
||||||
|
'tri_mouth': [114, 3] (start from 1, as a supplement to mouth triangles)
|
||||||
|
'kpt_ind': [68,] (start from 1)
|
||||||
|
PS:
|
||||||
|
You can change codes according to your own saved data.
|
||||||
|
Just make sure the model has corresponding attributes.
|
||||||
|
'''
|
||||||
|
C = sio.loadmat(model_path)
|
||||||
|
model = C['model']
|
||||||
|
model = model[0,0]
|
||||||
|
|
||||||
|
# change dtype from double(np.float64) to np.float32,
|
||||||
|
# since big matrix process(espetially matrix dot) is too slow in python.
|
||||||
|
model['shapeMU'] = (model['shapeMU'] + model['expMU']).astype(np.float32)
|
||||||
|
model['shapePC'] = model['shapePC'].astype(np.float32)
|
||||||
|
model['shapeEV'] = model['shapeEV'].astype(np.float32)
|
||||||
|
model['expEV'] = model['expEV'].astype(np.float32)
|
||||||
|
model['expPC'] = model['expPC'].astype(np.float32)
|
||||||
|
|
||||||
|
# matlab start with 1. change to 0 in python.
|
||||||
|
model['tri'] = model['tri'].T.copy(order = 'C').astype(np.int32) - 1
|
||||||
|
model['tri_mouth'] = model['tri_mouth'].T.copy(order = 'C').astype(np.int32) - 1
|
||||||
|
|
||||||
|
# kpt ind
|
||||||
|
model['kpt_ind'] = (np.squeeze(model['kpt_ind']) - 1).astype(np.int32)
|
||||||
|
|
||||||
|
return model
|
||||||
|
|
||||||
|
def load_BFM_info(path = 'BFM_info.mat'):
|
||||||
|
''' load 3DMM model extra information
|
||||||
|
Args:
|
||||||
|
path: path to BFM info.
|
||||||
|
Returns:
|
||||||
|
model_info:
|
||||||
|
'symlist': 2 x 26720
|
||||||
|
'symlist_tri': 2 x 52937
|
||||||
|
'segbin': 4 x n (0: nose, 1: eye, 2: mouth, 3: cheek)
|
||||||
|
'segbin_tri': 4 x ntri
|
||||||
|
'face_contour': 1 x 28
|
||||||
|
'face_contour_line': 1 x 512
|
||||||
|
'face_contour_front': 1 x 28
|
||||||
|
'face_contour_front_line': 1 x 512
|
||||||
|
'nose_hole': 1 x 142
|
||||||
|
'nose_hole_right': 1 x 71
|
||||||
|
'nose_hole_left': 1 x 71
|
||||||
|
'parallel': 17 x 1 cell
|
||||||
|
'parallel_face_contour': 28 x 1 cell
|
||||||
|
'uv_coords': n x 2
|
||||||
|
'''
|
||||||
|
C = sio.loadmat(path)
|
||||||
|
model_info = C['model_info']
|
||||||
|
model_info = model_info[0,0]
|
||||||
|
return model_info
|
||||||
|
|
||||||
|
def load_uv_coords(path = 'BFM_UV.mat'):
|
||||||
|
''' load uv coords of BFM
|
||||||
|
Args:
|
||||||
|
path: path to data.
|
||||||
|
Returns:
|
||||||
|
uv_coords: [nver, 2]. range: 0-1
|
||||||
|
'''
|
||||||
|
C = sio.loadmat(path)
|
||||||
|
uv_coords = C['UV'].copy(order = 'C')
|
||||||
|
return uv_coords
|
||||||
|
|
||||||
|
def load_pncc_code(path = 'pncc_code.mat'):
|
||||||
|
''' load pncc code of BFM
|
||||||
|
PNCC code: Defined in 'Face Alignment Across Large Poses: A 3D Solution Xiangyu'
|
||||||
|
download at http://www.cbsr.ia.ac.cn/users/xiangyuzhu/projects/3DDFA/main.htm.
|
||||||
|
Args:
|
||||||
|
path: path to data.
|
||||||
|
Returns:
|
||||||
|
pncc_code: [nver, 3]
|
||||||
|
'''
|
||||||
|
C = sio.loadmat(path)
|
||||||
|
pncc_code = C['vertex_code'].T
|
||||||
|
return pncc_code
|
||||||
|
|
||||||
|
##
|
||||||
|
def get_organ_ind(model_info):
|
||||||
|
''' get nose, eye, mouth index
|
||||||
|
'''
|
||||||
|
valid_bin = model_info['segbin'].astype(bool)
|
||||||
|
organ_ind = np.nonzero(valid_bin[0,:])[0]
|
||||||
|
for i in range(1, valid_bin.shape[0] - 1):
|
||||||
|
organ_ind = np.union1d(organ_ind, np.nonzero(valid_bin[i,:])[0])
|
||||||
|
return organ_ind.astype(np.int32)
|
||||||
@@ -0,0 +1,141 @@
|
|||||||
|
from __future__ import absolute_import
|
||||||
|
from __future__ import division
|
||||||
|
from __future__ import print_function
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
from . import fit
|
||||||
|
from . import load
|
||||||
|
|
||||||
|
class MorphabelModel(object):
|
||||||
|
"""docstring for MorphabelModel
|
||||||
|
model: nver: number of vertices. ntri: number of triangles. *: must have. ~: can generate ones array for place holder.
|
||||||
|
'shapeMU': [3*nver, 1]. *
|
||||||
|
'shapePC': [3*nver, n_shape_para]. *
|
||||||
|
'shapeEV': [n_shape_para, 1]. ~
|
||||||
|
'expMU': [3*nver, 1]. ~
|
||||||
|
'expPC': [3*nver, n_exp_para]. ~
|
||||||
|
'expEV': [n_exp_para, 1]. ~
|
||||||
|
'texMU': [3*nver, 1]. ~
|
||||||
|
'texPC': [3*nver, n_tex_para]. ~
|
||||||
|
'texEV': [n_tex_para, 1]. ~
|
||||||
|
'tri': [ntri, 3] (start from 1, should sub 1 in python and c++). *
|
||||||
|
'tri_mouth': [114, 3] (start from 1, as a supplement to mouth triangles). ~
|
||||||
|
'kpt_ind': [68,] (start from 1). ~
|
||||||
|
"""
|
||||||
|
def __init__(self, model_path, model_type = 'BFM'):
|
||||||
|
super( MorphabelModel, self).__init__()
|
||||||
|
if model_type=='BFM':
|
||||||
|
self.model = load.load_BFM(model_path)
|
||||||
|
else:
|
||||||
|
print('sorry, not support other 3DMM model now')
|
||||||
|
exit()
|
||||||
|
|
||||||
|
# fixed attributes
|
||||||
|
self.nver = self.model['shapePC'].shape[0]/3
|
||||||
|
self.ntri = self.model['tri'].shape[0]
|
||||||
|
self.n_shape_para = self.model['shapePC'].shape[1]
|
||||||
|
self.n_exp_para = self.model['expPC'].shape[1]
|
||||||
|
self.n_tex_para = self.model['texMU'].shape[1]
|
||||||
|
|
||||||
|
self.kpt_ind = self.model['kpt_ind']
|
||||||
|
self.triangles = self.model['tri']
|
||||||
|
self.full_triangles = np.vstack((self.model['tri'], self.model['tri_mouth']))
|
||||||
|
|
||||||
|
# ------------------------------------- shape: represented with mesh(vertices & triangles(fixed))
|
||||||
|
def get_shape_para(self, type = 'random'):
|
||||||
|
if type == 'zero':
|
||||||
|
sp = np.random.zeros((self.n_shape_para, 1))
|
||||||
|
elif type == 'random':
|
||||||
|
sp = np.random.rand(self.n_shape_para, 1)*1e04
|
||||||
|
return sp
|
||||||
|
|
||||||
|
def get_exp_para(self, type = 'random'):
|
||||||
|
if type == 'zero':
|
||||||
|
ep = np.zeros((self.n_exp_para, 1))
|
||||||
|
elif type == 'random':
|
||||||
|
ep = -1.5 + 3*np.random.random([self.n_exp_para, 1])
|
||||||
|
ep[6:, 0] = 0
|
||||||
|
|
||||||
|
return ep
|
||||||
|
|
||||||
|
def generate_vertices(self, shape_para, exp_para):
|
||||||
|
'''
|
||||||
|
Args:
|
||||||
|
shape_para: (n_shape_para, 1)
|
||||||
|
exp_para: (n_exp_para, 1)
|
||||||
|
Returns:
|
||||||
|
vertices: (nver, 3)
|
||||||
|
'''
|
||||||
|
vertices = self.model['shapeMU'] + self.model['shapePC'].dot(shape_para) + self.model['expPC'].dot(exp_para)
|
||||||
|
vertices = np.reshape(vertices, [int(3), int(len(vertices)/3)], 'F').T
|
||||||
|
|
||||||
|
return vertices
|
||||||
|
|
||||||
|
# -------------------------------------- texture: here represented with rgb value(colors) in vertices.
|
||||||
|
def get_tex_para(self, type = 'random'):
|
||||||
|
if type == 'zero':
|
||||||
|
tp = np.zeros((self.n_tex_para, 1))
|
||||||
|
elif type == 'random':
|
||||||
|
tp = np.random.rand(self.n_tex_para, 1)
|
||||||
|
return tp
|
||||||
|
|
||||||
|
def generate_colors(self, tex_para):
|
||||||
|
'''
|
||||||
|
Args:
|
||||||
|
tex_para: (n_tex_para, 1)
|
||||||
|
Returns:
|
||||||
|
colors: (nver, 3)
|
||||||
|
'''
|
||||||
|
colors = self.model['texMU'] + self.model['texPC'].dot(tex_para*self.model['texEV'])
|
||||||
|
colors = np.reshape(colors, [int(3), int(len(colors)/3)], 'F').T/255.
|
||||||
|
|
||||||
|
return colors
|
||||||
|
|
||||||
|
|
||||||
|
# ------------------------------------------- transformation
|
||||||
|
# ------------- transform
|
||||||
|
def rotate(self, vertices, angles):
|
||||||
|
''' rotate face
|
||||||
|
Args:
|
||||||
|
vertices: [nver, 3]
|
||||||
|
angles: [3] x, y, z rotation angle(degree)
|
||||||
|
x: pitch. positive for looking down
|
||||||
|
y: yaw. positive for looking left
|
||||||
|
z: roll. positive for tilting head right
|
||||||
|
Returns:
|
||||||
|
vertices: rotated vertices
|
||||||
|
'''
|
||||||
|
return Face.face3d.mesh.transform.rotate(vertices, angles)
|
||||||
|
|
||||||
|
def transform(self, vertices, s, angles, t3d):
|
||||||
|
R = Face.face3d.mesh.transform.angle2matrix(angles)
|
||||||
|
return Face.face3d.mesh.transform.similarity_transform(vertices, s, R, t3d)
|
||||||
|
|
||||||
|
def transform_3ddfa(self, vertices, s, angles, t3d): # only used for processing 300W_LP data
|
||||||
|
R = Face.face3d.mesh.transform.angle2matrix_3ddfa(angles)
|
||||||
|
return Face.face3d.mesh.transform.similarity_transform(vertices, s, R, t3d)
|
||||||
|
|
||||||
|
# --------------------------------------------------- fitting
|
||||||
|
def fit(self, x, X_ind, max_iter = 4, isShow = False):
|
||||||
|
''' fit 3dmm & pose parameters
|
||||||
|
Args:
|
||||||
|
x: (n, 2) image points
|
||||||
|
X_ind: (n,) corresponding Model vertex indices
|
||||||
|
max_iter: iteration
|
||||||
|
isShow: whether to reserve middle results for show
|
||||||
|
Returns:
|
||||||
|
fitted_sp: (n_sp, 1). shape parameters
|
||||||
|
fitted_ep: (n_ep, 1). exp parameters
|
||||||
|
s, angles, t
|
||||||
|
'''
|
||||||
|
if isShow:
|
||||||
|
fitted_sp, fitted_ep, s, R, t = fit.fit_points_for_show(x, X_ind, self.model, n_sp = self.n_shape_para, n_ep = self.n_exp_para, max_iter = max_iter)
|
||||||
|
angles = np.zeros((R.shape[0], 3))
|
||||||
|
for i in range(R.shape[0]):
|
||||||
|
angles[i] = Face.face3d.mesh.transform.matrix2angle(R[i])
|
||||||
|
else:
|
||||||
|
fitted_sp, fitted_ep, s, R, t = fit.fit_points(x, X_ind, self.model, n_sp = self.n_shape_para, n_ep = self.n_exp_para, max_iter = max_iter)
|
||||||
|
angles = Face.face3d.mesh.transform.matrix2angle(R)
|
||||||
|
return fitted_sp, fitted_ep, s, angles, t
|
||||||
|
|
||||||
|
|
||||||
@@ -0,0 +1,96 @@
|
|||||||
|
import os
|
||||||
|
import cv2
|
||||||
|
import glob
|
||||||
|
import numpy as np
|
||||||
|
from core.utils import landmark_processor
|
||||||
|
|
||||||
|
from core.face_enhance.face_gan_pt import FaceGAN
|
||||||
|
from time import time
|
||||||
|
|
||||||
|
class FaceEnhancement(object):
|
||||||
|
def __init__(self, size=512, gpu_id=None):
|
||||||
|
self.facegan = FaceGAN(size, gpu_id)
|
||||||
|
self.size = size
|
||||||
|
self.threshold = 0.9
|
||||||
|
|
||||||
|
# the mask for pasting restored faces back
|
||||||
|
self.mask = np.zeros((512, 512), np.float32)
|
||||||
|
cv2.rectangle(self.mask, (26, 26), (486, 486), (1, 1, 1), -1, cv2.LINE_AA)
|
||||||
|
self.mask = cv2.GaussianBlur(self.mask, (101, 101), 11)
|
||||||
|
self.mask = cv2.GaussianBlur(self.mask, (101, 101), 11)
|
||||||
|
|
||||||
|
self.kernel = np.array((
|
||||||
|
[0.0625, 0.125, 0.0625],
|
||||||
|
[0.125, 0.25, 0.125],
|
||||||
|
[0.0625, 0.125, 0.0625]), dtype="float32")
|
||||||
|
|
||||||
|
def process(self, img, landmarks1k):
|
||||||
|
|
||||||
|
assert len(landmarks1k) == 1000
|
||||||
|
|
||||||
|
image_to_face_mat = landmark_processor.get_transform_mat_face_restore(landmarks1k, self.size)
|
||||||
|
tfm_inv = cv2.invertAffineTransform(image_to_face_mat)
|
||||||
|
|
||||||
|
height, width = img.shape[:2]
|
||||||
|
full_mask = np.zeros((height, width), dtype=np.float32)
|
||||||
|
full_img = np.zeros(img.shape, dtype=np.uint8)
|
||||||
|
|
||||||
|
fh, fw = (landmarks1k[0][1]-landmarks1k[154][1]), (landmarks1k[95][0]-landmarks1k[215][0])
|
||||||
|
|
||||||
|
of = cv2.warpAffine(img, image_to_face_mat, (self.size, self.size), flags=cv2.INTER_LANCZOS4, borderMode=cv2.BORDER_CONSTANT, borderValue=[0, 0, 0])
|
||||||
|
|
||||||
|
# enhance the face
|
||||||
|
t0 = time()
|
||||||
|
|
||||||
|
ef = self.facegan.process(of)
|
||||||
|
# print('facegan process costs:', time() - t0)
|
||||||
|
tmp_mask = self.mask
|
||||||
|
tmp_mask = cv2.resize(tmp_mask, ef.shape[:2])
|
||||||
|
tmp_mask = cv2.warpAffine(tmp_mask, tfm_inv, (width, height), flags=3)
|
||||||
|
|
||||||
|
if min(fh, fw)<100: # gaussian filter for small faces
|
||||||
|
ef = cv2.filter2D(ef, -1, self.kernel)
|
||||||
|
|
||||||
|
# tmp_img = cv2.warpAffine(ef, tfm_inv, (width, height), flags=3)
|
||||||
|
tmp_img = cv2.warpAffine(ef, tfm_inv, (width, height), dst=img.copy(), borderMode=cv2.BORDER_TRANSPARENT)
|
||||||
|
|
||||||
|
# cv2.imshow("tmp_img: ", tmp_img)
|
||||||
|
|
||||||
|
mask = tmp_mask - full_mask
|
||||||
|
full_mask[np.where(mask>0)] = tmp_mask[np.where(mask>0)]
|
||||||
|
full_img[np.where(mask>0)] = tmp_img[np.where(mask>0)]
|
||||||
|
|
||||||
|
full_mask = full_mask[:, :, np.newaxis]
|
||||||
|
img = cv2.convertScaleAbs(img*(1-full_mask) + full_img*full_mask)
|
||||||
|
|
||||||
|
return img
|
||||||
|
|
||||||
|
if __name__=='__main__':
|
||||||
|
|
||||||
|
indir = '/media/DATA_4T/zao_data/tencent_con_0701/ref_c/res_contrast_blur13_notps_yunfu_res'
|
||||||
|
outdir = '/media/DATA_4T/zao_data/tencent_con_0701/ref_c/res_contrast_blur13_notps_yunfu_res_outs2'
|
||||||
|
os.makedirs(outdir, exist_ok=True)
|
||||||
|
|
||||||
|
faceenhancer = FaceEnhancement(base_dir="/", size=512, model="GPEN-512", channel_multiplier=2)
|
||||||
|
|
||||||
|
files = sorted(glob.glob(os.path.join(indir, '*.*g')))
|
||||||
|
for n, file in enumerate(files[:]):
|
||||||
|
filename = os.path.basename(file)
|
||||||
|
txtname = file.replace(".jpg", "_landmark1k.txt")
|
||||||
|
|
||||||
|
im = cv2.imread(file, cv2.IMREAD_COLOR) # BGR
|
||||||
|
print(txtname)
|
||||||
|
landmark = np.loadtxt(txtname)
|
||||||
|
if not isinstance(im, np.ndarray): print(filename, 'error'); continue
|
||||||
|
|
||||||
|
start = time()
|
||||||
|
|
||||||
|
img = faceenhancer.process(im, landmark)
|
||||||
|
|
||||||
|
end = time()
|
||||||
|
|
||||||
|
print("Time cost: {:.4f}".format(end - start))
|
||||||
|
|
||||||
|
cv2.imwrite(os.path.join(outdir, '.'.join(filename.split('.')[:-1])+'_2.jpg'), img)
|
||||||
|
|
||||||
|
|
||||||
@@ -0,0 +1,64 @@
|
|||||||
|
'''
|
||||||
|
@paper: GAN Prior Embedded Network for Blind Face Restoration in the Wild (CVPR2021)
|
||||||
|
@author: yangxy (yangtao9009@gmail.com)
|
||||||
|
'''
|
||||||
|
import torch
|
||||||
|
import os
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
class FaceGAN(object):
|
||||||
|
def __init__(self, size=512, gpu_id=0):
|
||||||
|
# self.mfile = os.path.join(base_dir, model+'.pth')
|
||||||
|
self.n_mlp = 8
|
||||||
|
self.resolution = size
|
||||||
|
self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||||
|
|
||||||
|
self.load_model()
|
||||||
|
|
||||||
|
def load_model(self):
|
||||||
|
self.face_gan_model = os.path.join('weights', "face_enhance_0630.pt")
|
||||||
|
self.model = torch.jit.load(self.face_gan_model, torch.device('cpu')).to(self.device)
|
||||||
|
# self.model = torch.jit.load(self.face_gan_model).to(self.device)
|
||||||
|
# self.model = torch.load(self.face_gan_model,map_location=self.device)
|
||||||
|
|
||||||
|
self.model.eval()
|
||||||
|
|
||||||
|
def process_o(self, img):
|
||||||
|
img = cv2.resize(img, (self.resolution, self.resolution))
|
||||||
|
img_t = self.img2tensor(img)
|
||||||
|
|
||||||
|
with torch.no_grad():
|
||||||
|
out, __ = self.model(img_t)
|
||||||
|
|
||||||
|
out = self.tensor2img(out)
|
||||||
|
|
||||||
|
return out
|
||||||
|
|
||||||
|
def process(self, img):
|
||||||
|
img = cv2.resize(img, (self.resolution, self.resolution))
|
||||||
|
img_t = self.img2tensor(img)
|
||||||
|
|
||||||
|
with torch.no_grad():
|
||||||
|
out = self.forward(img_t)
|
||||||
|
|
||||||
|
out = self.tensor2img(out)
|
||||||
|
|
||||||
|
return out
|
||||||
|
|
||||||
|
def forward(self, img_t):
|
||||||
|
with torch.no_grad():
|
||||||
|
out = self.model(img_t)
|
||||||
|
|
||||||
|
return out
|
||||||
|
|
||||||
|
def img2tensor(self, img):
|
||||||
|
img_t = (torch.from_numpy(img).to(self.device)/255. - 0.5) / 0.5
|
||||||
|
img_t = img_t.permute(2, 0, 1).unsqueeze(0).flip(1) # BGR->RGB
|
||||||
|
return img_t
|
||||||
|
|
||||||
|
def tensor2img(self, image_tensor, pmax=255.0, imtype=np.uint8):
|
||||||
|
image_tensor = image_tensor * 0.5 + 0.5
|
||||||
|
image_tensor = image_tensor.squeeze(0).permute(1, 2, 0).flip(2) # RGB->BGR
|
||||||
|
image_numpy = np.clip(image_tensor.float().cpu().numpy(), 0, 1) * pmax
|
||||||
|
|
||||||
|
return image_numpy.astype(imtype)
|
||||||
@@ -0,0 +1,102 @@
|
|||||||
|
#!/usr/bin/env python
|
||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
|
||||||
|
# @File : setup.py
|
||||||
|
# @Time : 2020/1/15
|
||||||
|
# @Author : yangchaojie (yangchaojie@immomo.com)
|
||||||
|
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
import shutil
|
||||||
|
import numpy
|
||||||
|
import tempfile
|
||||||
|
|
||||||
|
from setuptools import setup
|
||||||
|
from setuptools.extension import Extension
|
||||||
|
|
||||||
|
from Cython.Build import cythonize
|
||||||
|
from Cython.Distutils import build_ext
|
||||||
|
|
||||||
|
import platform
|
||||||
|
|
||||||
|
|
||||||
|
def get_root_path(root):
|
||||||
|
if os.path.dirname(root) in ['', '.']:
|
||||||
|
return os.path.basename(root)
|
||||||
|
else:
|
||||||
|
return get_root_path(os.path.dirname(root))
|
||||||
|
|
||||||
|
|
||||||
|
def copy_file(src, dest):
|
||||||
|
if os.path.exists(dest):
|
||||||
|
return
|
||||||
|
|
||||||
|
if not os.path.exists(os.path.dirname(dest)):
|
||||||
|
os.makedirs(os.path.dirname(dest))
|
||||||
|
if os.path.isdir(src):
|
||||||
|
shutil.copytree(src, dest)
|
||||||
|
else:
|
||||||
|
shutil.copyfile(src, dest)
|
||||||
|
|
||||||
|
|
||||||
|
def touch_init_file():
|
||||||
|
init_file_name = os.path.join(tempfile.mkdtemp(), '__init__.py')
|
||||||
|
with open(init_file_name, 'w'):
|
||||||
|
pass
|
||||||
|
return init_file_name
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
def compose_extensions(root='.'):
|
||||||
|
for file_ in os.listdir(root):
|
||||||
|
abs_file = os.path.join(root, file_)
|
||||||
|
|
||||||
|
if os.path.isfile(abs_file):
|
||||||
|
if abs_file.endswith('.py'):
|
||||||
|
extensions.append(Extension(get_root_path(abs_file) + '.*', [abs_file]))
|
||||||
|
elif abs_file.endswith('.c') or abs_file.endswith('.pyc'):
|
||||||
|
continue
|
||||||
|
else:
|
||||||
|
copy_file(abs_file, os.path.join(build_root_dir, abs_file))
|
||||||
|
if abs_file.endswith('__init__.py'):
|
||||||
|
copy_file(init_file, os.path.join(build_root_dir, abs_file))
|
||||||
|
|
||||||
|
else:
|
||||||
|
if os.path.basename(abs_file) in ignore_folders :
|
||||||
|
continue
|
||||||
|
if os.path.basename(abs_file) in conf_folders:
|
||||||
|
copy_file(abs_file, os.path.join(build_root_dir, abs_file))
|
||||||
|
compose_extensions(abs_file)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
# if __name__ == '__main__':
|
||||||
|
build_root_dir = 'build/lib.' + platform.system().lower() + '-' + platform.machine() + '-' + str(
|
||||||
|
sys.version_info.major) + '.' + str(sys.version_info.minor)
|
||||||
|
|
||||||
|
print(build_root_dir)
|
||||||
|
|
||||||
|
extensions = []
|
||||||
|
ignore_folders = ['build', 'new_ref_zao_color_0818', 'ref_hair_online_0703_local', 'ref_分好类别', '.git']
|
||||||
|
conf_folders = ['conf']
|
||||||
|
|
||||||
|
|
||||||
|
init_file = touch_init_file()
|
||||||
|
print(init_file)
|
||||||
|
|
||||||
|
|
||||||
|
compose_extensions()
|
||||||
|
os.remove(init_file)
|
||||||
|
|
||||||
|
setup(
|
||||||
|
name='moxie_hairstyle',
|
||||||
|
version='1.0',
|
||||||
|
ext_modules=cythonize(
|
||||||
|
extensions,
|
||||||
|
nthreads=16,
|
||||||
|
compiler_directives=dict(always_allow_keywords=True),
|
||||||
|
include_path=[numpy.get_include()]),
|
||||||
|
cmdclass=dict(build_ext=build_ext))
|
||||||
|
|
||||||
|
# python setup.py build_ext
|
||||||
@@ -0,0 +1,186 @@
|
|||||||
|
import cv2
|
||||||
|
|
||||||
|
|
||||||
|
class SingleFaceQualityInfo():
|
||||||
|
def __int__(self, gpu_id):
|
||||||
|
print('Model 3DDFA init success')
|
||||||
|
|
||||||
|
def getSpotProportion(self, grayImage):
|
||||||
|
Proportion = 0
|
||||||
|
hist = cv2.calcHist([grayImage], [0], None, [256], [0, 256])
|
||||||
|
for i in range(253, 256):
|
||||||
|
Proportion = Proportion + hist[i, 0]
|
||||||
|
return Proportion
|
||||||
|
|
||||||
|
def GetSpotRatio(self, wrap_img, wrap_face_rect):
|
||||||
|
scale = 0.9
|
||||||
|
|
||||||
|
x = max(0, int(wrap_face_rect[0]))
|
||||||
|
y = max(0, int(wrap_face_rect[1]))
|
||||||
|
w = max(0, int(wrap_face_rect[2]-wrap_face_rect[0]))
|
||||||
|
h = max(0, int(wrap_face_rect[3]-wrap_face_rect[1]))
|
||||||
|
|
||||||
|
height, width, _ = wrap_img.shape
|
||||||
|
if (x + w) > width:
|
||||||
|
w = width - x
|
||||||
|
if (y + h) > height:
|
||||||
|
h = height - y
|
||||||
|
p_cen = (x + w/2, y+h/2)
|
||||||
|
new_w = w*scale
|
||||||
|
new_h = h*scale
|
||||||
|
na_box = [int(p_cen[0]-new_w/2), int(p_cen[1]) - int(new_h/2), int(new_w), int(new_h)]
|
||||||
|
imgCropGray = cv2.cvtColor(wrap_img[na_box[1]:na_box[1] + na_box[3], na_box[0]:na_box[0] + na_box[2]],cv2.COLOR_RGB2GRAY)
|
||||||
|
|
||||||
|
Spot_Ratio = self.getSpotProportion(imgCropGray) / (imgCropGray.shape[0] * imgCropGray.shape[1])
|
||||||
|
return Spot_Ratio
|
||||||
|
|
||||||
|
def BoxClarityAndBrightValue(self, wrap_img, wrap_face_rect):
|
||||||
|
x = max(0, int(wrap_face_rect[0]))
|
||||||
|
y = max(0, int(wrap_face_rect[1]))
|
||||||
|
w = max(0, int(wrap_face_rect[2]-wrap_face_rect[0]))
|
||||||
|
h = max(0, int(wrap_face_rect[3]-wrap_face_rect[1]))
|
||||||
|
# cv2.imshow('face', wrap_img[y:y+h,x:x+w])
|
||||||
|
# cv2.waitKey()
|
||||||
|
height, width, _ = wrap_img.shape
|
||||||
|
if (x + w) > width:
|
||||||
|
w = width - x
|
||||||
|
if (y + h) > height:
|
||||||
|
h = height - y
|
||||||
|
box_crop = [int(x), int(y), int(w), int(h)]
|
||||||
|
imgCropGray = cv2.cvtColor(wrap_img[box_crop[1]:box_crop[1] + box_crop[3], box_crop[0]:box_crop[0] + box_crop[2]],cv2.COLOR_RGB2GRAY)
|
||||||
|
left_ = [0,0,int(w/2.0),int(h)]
|
||||||
|
right_ = [int(w/2.0),0,int(w/2.0),int(h)]
|
||||||
|
imgR = imgCropGray[right_[1]:right_[1] + right_[3], right_[0]:right_[0] + right_[2]]
|
||||||
|
imgL = imgCropGray[left_[1]:left_[1] + left_[3], left_[0]:left_[0] + left_[2]]
|
||||||
|
mean_R, var_R = cv2.meanStdDev(imgR)
|
||||||
|
mean_L, var_L = cv2.meanStdDev(imgL)
|
||||||
|
|
||||||
|
bV_R = mean_R[0]
|
||||||
|
bV_L = mean_L[0]
|
||||||
|
|
||||||
|
imageCropSobel = cv2.Laplacian(imgCropGray, cv2.CV_64F, 3)
|
||||||
|
mean, var = cv2.meanStdDev(imageCropSobel)
|
||||||
|
clarityValue = var[0]
|
||||||
|
result = [bV_R, bV_L, clarityValue]
|
||||||
|
return result
|
||||||
|
|
||||||
|
def qualityTest(self, wrap_img, wrap_face_rect, single_eulers):
|
||||||
|
single_face_quality_info = {
|
||||||
|
'quality_score': -1,
|
||||||
|
'face_quality_flag': -1,
|
||||||
|
'SpotRatio': 0,
|
||||||
|
'Brightness': 0,
|
||||||
|
'Clarity': 0,
|
||||||
|
'SpotScore': 0,
|
||||||
|
'BrightnessScore': 0,
|
||||||
|
'ClarityScore': 0
|
||||||
|
}
|
||||||
|
|
||||||
|
if wrap_img is None:
|
||||||
|
single_face_quality_info['face_quality_flag'] = 0
|
||||||
|
qp_img_size = wrap_img.shape[1]
|
||||||
|
Spot_Ratio = self.GetSpotRatio(wrap_img, wrap_face_rect)
|
||||||
|
result = self.BoxClarityAndBrightValue(wrap_img, wrap_face_rect)
|
||||||
|
bv_R, bv_L, clarityValue = result
|
||||||
|
single_face_quality_info['SpotRatio'] = Spot_Ratio
|
||||||
|
single_face_quality_info['Brightness'] = min(bv_L, bv_R)
|
||||||
|
single_face_quality_info['Clarity'] = clarityValue
|
||||||
|
if qp_img_size == 100:
|
||||||
|
clarityValue_thresholdDown = 35
|
||||||
|
clarityValue_thresholdUp = 300
|
||||||
|
else:
|
||||||
|
clarityValue_thresholdDown = 24
|
||||||
|
clarityValue_thresholdUp = 300
|
||||||
|
Bright_area_thre = 0.06
|
||||||
|
brightnessValue_thresholddown = 40
|
||||||
|
brightnessValue_thresholdUp = 230
|
||||||
|
|
||||||
|
whitespot = (Spot_Ratio > Bright_area_thre)
|
||||||
|
brightness =(bv_R > brightnessValue_thresholddown and bv_R < brightnessValue_thresholdUp) and (bv_L > brightnessValue_thresholddown and bv_L < brightnessValue_thresholdUp)
|
||||||
|
clarity = clarityValue < clarityValue_thresholdUp and clarityValue > clarityValue_thresholdDown
|
||||||
|
if whitespot:
|
||||||
|
single_face_quality_info['face_qality_flag'] = 2
|
||||||
|
if brightness:
|
||||||
|
single_face_quality_info['face_qality_flag'] = 4
|
||||||
|
if not clarity:
|
||||||
|
single_face_quality_info['face_qality_flag'] = 3
|
||||||
|
if brightness and clarity and not whitespot:
|
||||||
|
single_face_quality_info['face_qality_flag'] = 1
|
||||||
|
|
||||||
|
|
||||||
|
if single_face_quality_info['face_qality_flag'] == 1:
|
||||||
|
PerfectSpotRatio = 0.02
|
||||||
|
if qp_img_size == 100:
|
||||||
|
PerfectClarity = 130
|
||||||
|
else:
|
||||||
|
PerfectClarity = 60
|
||||||
|
PerfectBrightness_down = 90
|
||||||
|
PerfectBrightness_up = 200
|
||||||
|
PerfectBrightness_standard = PerfectBrightness_down - brightnessValue_thresholddown
|
||||||
|
PerfectClarity_standard = PerfectClarity - clarityValue_thresholdDown
|
||||||
|
PerfectSpot_standard = Bright_area_thre - PerfectSpotRatio
|
||||||
|
if single_face_quality_info['SpotRatio'] <= PerfectSpotRatio:
|
||||||
|
single_face_quality_info['SpotScore'] = 1
|
||||||
|
else:
|
||||||
|
single_face_quality_info['SpotRatio'] = 1 - (single_face_quality_info['SpotRatio'] - PerfectSpotRatio) / PerfectSpot_standard
|
||||||
|
if single_face_quality_info['Clarity'] >= PerfectClarity:
|
||||||
|
single_face_quality_info['ClarityScore'] = 1
|
||||||
|
else:
|
||||||
|
single_face_quality_info['ClarityScore'] = 1 - (PerfectClarity - single_face_quality_info['Clarity']) / PerfectClarity_standard
|
||||||
|
if single_face_quality_info['Brightness'] >= PerfectBrightness_down and single_face_quality_info['Brightness']<= PerfectBrightness_up:
|
||||||
|
single_face_quality_info['BrightnessScore'] = 1
|
||||||
|
elif single_face_quality_info['Brightness'] < PerfectBrightness_down:
|
||||||
|
single_face_quality_info['BrightnessScore'] = 1 - (PerfectBrightness_down - single_face_quality_info['Brightness']) / PerfectBrightness_standard
|
||||||
|
elif single_face_quality_info['Brightness'] > PerfectBrightness_up:
|
||||||
|
single_face_quality_info['BrightnessScore'] = 1 - (single_face_quality_info['Brightness'] - PerfectBrightness_up)/PerfectBrightness_standard
|
||||||
|
else:
|
||||||
|
print("Wrong Brightness :%f",single_face_quality_info['Brightness'])
|
||||||
|
else:
|
||||||
|
single_face_quality_info['ClarityScore'] = 0
|
||||||
|
single_face_quality_info['BrightnessScore'] = 0
|
||||||
|
single_face_quality_info['SpotScore'] = 0
|
||||||
|
|
||||||
|
if len(single_eulers) > 2:
|
||||||
|
euler_pitch_perfect = 10
|
||||||
|
euler_yaw_perfect = 10
|
||||||
|
|
||||||
|
euler_pitch_max = 25
|
||||||
|
euler_yaw_max = 25
|
||||||
|
euler_roll_max = 30
|
||||||
|
|
||||||
|
image_ClarityScore = float(single_face_quality_info['ClarityScore'])
|
||||||
|
image_BrightnessScore = single_face_quality_info['BrightnessScore']
|
||||||
|
image_SpotScore = single_face_quality_info['SpotScore']
|
||||||
|
|
||||||
|
euler_pitch, euler_yaw, euler_roll = single_eulers
|
||||||
|
|
||||||
|
if abs(euler_pitch) > euler_pitch_perfect and abs(euler_pitch) < euler_pitch_max and single_face_quality_info['face_qality_flag'] == 1:
|
||||||
|
image_PitchScore = 1 - (abs(euler_pitch) - euler_pitch_perfect) / (euler_pitch_max - euler_pitch_perfect)
|
||||||
|
elif abs(euler_pitch) < euler_pitch_perfect and single_face_quality_info['face_qality_flag'] == 1:
|
||||||
|
image_PitchScore = 1.0
|
||||||
|
else:
|
||||||
|
image_PitchScore = 0
|
||||||
|
if abs(euler_yaw) > euler_yaw_perfect and abs(euler_yaw) < euler_yaw_max and single_face_quality_info['face_qality_flag']==1:
|
||||||
|
image_YawScore = 1 - (abs(euler_yaw) - euler_yaw_perfect) / (euler_yaw_max - euler_yaw_perfect)
|
||||||
|
elif abs(euler_yaw) < euler_yaw_perfect and single_face_quality_info['face_qality_flag']==1:
|
||||||
|
image_YawScore = 1.0
|
||||||
|
else:
|
||||||
|
image_YawScore = 0
|
||||||
|
|
||||||
|
if image_PitchScore < 0.3 and image_YawScore < 0.3:
|
||||||
|
image_spot_weight = 0
|
||||||
|
image_clarity_weight = 0.1
|
||||||
|
image_brightness_weight = 0.1
|
||||||
|
euler_pitch_weight = 0.4
|
||||||
|
euler_yaw_weight = 0.4
|
||||||
|
else:
|
||||||
|
image_spot_weight = 0.05
|
||||||
|
image_clarity_weight = 0.35
|
||||||
|
image_brightness_weight = 0.1
|
||||||
|
euler_pitch_weight = 0.25
|
||||||
|
euler_yaw_weight = 0.25
|
||||||
|
if abs(euler_pitch) < euler_pitch_max and abs(euler_yaw) < euler_yaw_max and abs(euler_roll) < euler_roll_max and single_face_quality_info['face_qality_flag'] == 1:
|
||||||
|
single_face_quality_info['quality_score'] = (image_spot_weight*image_SpotScore + image_clarity_weight * image_ClarityScore + image_brightness_weight * image_BrightnessScore + euler_pitch_weight * image_PitchScore + euler_yaw_weight* image_YawScore)
|
||||||
|
else:
|
||||||
|
single_face_quality_info['quality_score'] = 0
|
||||||
|
return single_face_quality_info
|
||||||
@@ -0,0 +1,75 @@
|
|||||||
|
import os
|
||||||
|
import pickle
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from core.faceseg.u2net import U2NET
|
||||||
|
from core.utils import landmark_processor
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
|
||||||
|
class FaceSeg:
|
||||||
|
def __init__(self, gpu_id = 0):
|
||||||
|
model = U2NET(in_ch=4, out_ch=1)
|
||||||
|
weights = torch.load('weights/faceseg_20210927_01.pth', map_location='cpu')
|
||||||
|
model_dict = model.state_dict()
|
||||||
|
pretrained_dict = {}
|
||||||
|
for ix, (k, v) in enumerate(model_dict.items()):
|
||||||
|
if k in weights and weights[k].data.shape == v.data.shape:
|
||||||
|
pretrained_dict[k] = weights[k]
|
||||||
|
else:
|
||||||
|
print('ignore {}'.format(k))
|
||||||
|
model_dict.update(pretrained_dict)
|
||||||
|
model.load_state_dict(model_dict)
|
||||||
|
print('update success')
|
||||||
|
model.cuda(gpu_id)
|
||||||
|
model.eval()
|
||||||
|
self.model = model
|
||||||
|
self.last_mask = None
|
||||||
|
self.output_img_size = 320
|
||||||
|
self.gpu_id = gpu_id
|
||||||
|
|
||||||
|
def inference(self, frame, pt1k, video_mode = False):
|
||||||
|
image_to_face_mat = landmark_processor.get_transform_mat_full_face(pt1k, self.output_img_size)
|
||||||
|
face_image = cv2.warpAffine(frame, image_to_face_mat, (self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4)
|
||||||
|
if face_image.dtype == np.uint8: face_image = face_image.astype(np.float32) / 255
|
||||||
|
if video_mode and self.last_mask is not None:
|
||||||
|
last_small_mask = cv2.warpAffine(self.last_mask, image_to_face_mat,
|
||||||
|
(self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4)[:,:,np.newaxis]
|
||||||
|
input_img = np.concatenate([face_image, last_small_mask], axis=2)
|
||||||
|
else:
|
||||||
|
zero_mask = np.zeros((face_image.shape[1], face_image.shape[0], 1), dtype=np.float32)
|
||||||
|
input_img = np.concatenate([face_image, zero_mask], axis=2)
|
||||||
|
|
||||||
|
face_image_tensor = input_img.transpose((2, 0, 1))[np.newaxis]
|
||||||
|
face_image_tensor = torch.from_numpy(face_image_tensor).cuda(self.gpu_id)
|
||||||
|
with torch.no_grad():
|
||||||
|
mask = self.model.test(face_image_tensor)
|
||||||
|
mask = mask[0].detach().cpu().numpy().transpose((1, 2, 0))
|
||||||
|
|
||||||
|
origin_mask = cv2.warpAffine(mask, image_to_face_mat, (frame.shape[1], frame.shape[0]),
|
||||||
|
flags=cv2.WARP_INVERSE_MAP|cv2.INTER_LANCZOS4)[:, :, np.newaxis]
|
||||||
|
if video_mode: self.last_mask = origin_mask.copy()
|
||||||
|
|
||||||
|
return origin_mask
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
face_segmentor = FaceSeg(gpu_id=0)
|
||||||
|
testdata_dir = "/home/yangchaojie/Desktop/faceswap_hd/datasets/origin/8171"
|
||||||
|
for picname in os.listdir(testdata_dir):
|
||||||
|
img_path = os.path.join(testdata_dir, picname)
|
||||||
|
pkl_path = img_path[:-4]+".pkl"
|
||||||
|
if not picname.endswith(".jpg"):
|
||||||
|
continue
|
||||||
|
if not os.path.exists(pkl_path):
|
||||||
|
continue
|
||||||
|
img = cv2.imread(img_path)
|
||||||
|
with open(pkl_path, "rb") as fp:
|
||||||
|
info = pickle.load(fp)
|
||||||
|
pt1k = info["pt1k"]
|
||||||
|
|
||||||
|
face_seg_mask = face_segmentor.inference(img, pt1k, video_mode=False)
|
||||||
|
# cv2.imshow("face_seg_mask", face_seg_mask)
|
||||||
|
# cv2.imshow("img", img)
|
||||||
|
# cv2.waitKey()
|
||||||
@@ -0,0 +1,299 @@
|
|||||||
|
from __future__ import division
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import torch.nn.functional as F
|
||||||
|
import torch.nn.init as init
|
||||||
|
import torch.utils.model_zoo as model_zoo
|
||||||
|
from torchvision import models
|
||||||
|
|
||||||
|
# general libs
|
||||||
|
import cv2
|
||||||
|
import matplotlib.pyplot as plt
|
||||||
|
from PIL import Image
|
||||||
|
import numpy as np
|
||||||
|
import math
|
||||||
|
import time
|
||||||
|
import tqdm
|
||||||
|
import os
|
||||||
|
import argparse
|
||||||
|
import copy
|
||||||
|
import sys
|
||||||
|
|
||||||
|
from utils.helpers import *
|
||||||
|
|
||||||
|
|
||||||
|
class ResBlock(nn.Module):
|
||||||
|
def __init__(self, backbone, indim, outdim=None, stride=1):
|
||||||
|
super(ResBlock, self).__init__()
|
||||||
|
self.backbone = backbone
|
||||||
|
if outdim == None:
|
||||||
|
outdim = indim
|
||||||
|
if indim == outdim and stride == 1:
|
||||||
|
self.downsample = None
|
||||||
|
else:
|
||||||
|
self.downsample = nn.Conv2d(indim, outdim, kernel_size=3, padding=1, stride=stride)
|
||||||
|
|
||||||
|
self.conv1 = nn.Conv2d(indim, outdim, kernel_size=3, padding=1, stride=stride)
|
||||||
|
self.conv2 = nn.Conv2d(outdim, outdim, kernel_size=3, padding=1)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
if self.backbone == 'resnest101':
|
||||||
|
r = self.conv1(F.relu(x, inplace=True))
|
||||||
|
r = self.conv2(F.relu(r, inplace=True))
|
||||||
|
else:
|
||||||
|
r = self.conv1(F.relu(x))
|
||||||
|
r = self.conv2(F.relu(r))
|
||||||
|
|
||||||
|
if self.downsample is not None:
|
||||||
|
x = self.downsample(x)
|
||||||
|
|
||||||
|
return x + r
|
||||||
|
|
||||||
|
|
||||||
|
class Encoder_M(nn.Module):
|
||||||
|
def __init__(self, backbone):
|
||||||
|
super(Encoder_M, self).__init__()
|
||||||
|
if backbone == 'resnest101':
|
||||||
|
self.conv1_m = nn.Conv2d(1, 128, kernel_size=7, stride=2, padding=3, bias=False)
|
||||||
|
self.conv1_o = nn.Conv2d(1, 128, kernel_size=7, stride=2, padding=3, bias=False)
|
||||||
|
else:
|
||||||
|
self.conv1_m = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
|
||||||
|
self.conv1_o = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
|
||||||
|
|
||||||
|
if backbone == 'resnet50':
|
||||||
|
resnet = models.resnet50(pretrained=True)
|
||||||
|
elif backbone == 'resnet18':
|
||||||
|
resnet = models.resnet18(pretrained=True)
|
||||||
|
|
||||||
|
self.conv1 = resnet.conv1
|
||||||
|
self.bn1 = resnet.bn1
|
||||||
|
self.relu = resnet.relu # 1/2, 64
|
||||||
|
self.maxpool = resnet.maxpool
|
||||||
|
|
||||||
|
self.res2 = resnet.layer1 # 1/4, 256
|
||||||
|
self.res3 = resnet.layer2 # 1/8, 512
|
||||||
|
self.res4 = resnet.layer3 # 1/8, 1024
|
||||||
|
|
||||||
|
self.register_buffer('mean', torch.FloatTensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
|
||||||
|
self.register_buffer('std', torch.FloatTensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
|
||||||
|
|
||||||
|
def forward(self, in_f, in_m, in_o):
|
||||||
|
f = (in_f - self.mean) / self.std
|
||||||
|
m = torch.unsqueeze(in_m, dim=1).float() # add channel dim
|
||||||
|
o = torch.unsqueeze(in_o, dim=1).float() # add channel dim
|
||||||
|
|
||||||
|
x = self.conv1(f) + self.conv1_m(m) + self.conv1_o(o)
|
||||||
|
x = self.bn1(x)
|
||||||
|
c1 = self.relu(x) # 1/2, 64
|
||||||
|
x = self.maxpool(c1) # 1/4, 64
|
||||||
|
r2 = self.res2(x) # 1/4, 256
|
||||||
|
r3 = self.res3(r2) # 1/8, 512
|
||||||
|
r4 = self.res4(r3) # 1/8, 1024
|
||||||
|
return r4, r3, r2, c1, f
|
||||||
|
|
||||||
|
|
||||||
|
class Encoder_Q(nn.Module):
|
||||||
|
def __init__(self, backbone):
|
||||||
|
super(Encoder_Q, self).__init__()
|
||||||
|
|
||||||
|
if backbone == 'resnet50':
|
||||||
|
resnet = models.resnet50(pretrained=True)
|
||||||
|
elif backbone == 'resnet18':
|
||||||
|
resnet = models.resnet18(pretrained=True)
|
||||||
|
|
||||||
|
self.conv1 = resnet.conv1
|
||||||
|
self.bn1 = resnet.bn1
|
||||||
|
self.relu = resnet.relu # 1/2, 64
|
||||||
|
self.maxpool = resnet.maxpool
|
||||||
|
|
||||||
|
self.res2 = resnet.layer1 # 1/4, 256
|
||||||
|
self.res3 = resnet.layer2 # 1/8, 512
|
||||||
|
self.res4 = resnet.layer3 # 1/8, 1024
|
||||||
|
|
||||||
|
self.register_buffer('mean', torch.FloatTensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
|
||||||
|
self.register_buffer('std', torch.FloatTensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
|
||||||
|
|
||||||
|
def forward(self, in_f):
|
||||||
|
f = (in_f - self.mean) / self.std
|
||||||
|
|
||||||
|
x = self.conv1(f)
|
||||||
|
x = self.bn1(x)
|
||||||
|
c1 = self.relu(x) # 1/2, 64
|
||||||
|
x = self.maxpool(c1) # 1/4, 64
|
||||||
|
r2 = self.res2(x) # 1/4, 256
|
||||||
|
r3 = self.res3(r2) # 1/8, 512
|
||||||
|
r4 = self.res4(r3) # 1/8, 1024
|
||||||
|
return r4, r3, r2, c1, f
|
||||||
|
|
||||||
|
|
||||||
|
class Refine(nn.Module):
|
||||||
|
def __init__(self, backbone, inplanes, planes, scale_factor=2):
|
||||||
|
super(Refine, self).__init__()
|
||||||
|
self.convFS = nn.Conv2d(inplanes, planes, kernel_size=(3, 3), padding=(1, 1), stride=1)
|
||||||
|
self.ResFS = ResBlock(backbone, planes, planes)
|
||||||
|
self.ResMM = ResBlock(backbone, planes, planes)
|
||||||
|
self.scale_factor = scale_factor
|
||||||
|
|
||||||
|
def forward(self, f, pm):
|
||||||
|
s = self.ResFS(self.convFS(f))
|
||||||
|
m = s + F.interpolate(pm, scale_factor=self.scale_factor, mode='bilinear', align_corners=False)
|
||||||
|
m = self.ResMM(m)
|
||||||
|
return m
|
||||||
|
|
||||||
|
|
||||||
|
class Decoder(nn.Module):
|
||||||
|
def __init__(self, mdim, scale_rate, backbone):
|
||||||
|
super(Decoder, self).__init__()
|
||||||
|
self.backbone = backbone
|
||||||
|
if backbone == 'resnest101':
|
||||||
|
self.convFM = nn.Conv2d(256, mdim, kernel_size=(3, 3), padding=(1, 1), stride=1)
|
||||||
|
else:
|
||||||
|
self.convFM = nn.Conv2d(1024 // scale_rate, mdim, kernel_size=(3, 3), padding=(1, 1), stride=1)
|
||||||
|
self.ResMM = ResBlock(backbone, mdim, mdim)
|
||||||
|
self.RF3 = Refine(backbone, 512 // scale_rate, mdim) # 1/8 -> 1/4
|
||||||
|
self.RF2 = Refine(backbone, 256 // scale_rate, mdim) # 1/4 -> 1
|
||||||
|
|
||||||
|
self.pred2 = nn.Conv2d(mdim, 2, kernel_size=(3, 3), padding=(1, 1), stride=1)
|
||||||
|
|
||||||
|
def forward(self, r4, r3, r2):
|
||||||
|
m4 = self.ResMM(self.convFM(r4))
|
||||||
|
m3 = self.RF3(r3, m4) # out: 1/8, 256
|
||||||
|
m2 = self.RF2(r2, m3) # out: 1/4, 256
|
||||||
|
|
||||||
|
if self.backbone == 'resnest101':
|
||||||
|
p2 = self.pred2(F.relu(m2, inplace=True))
|
||||||
|
else:
|
||||||
|
p2 = self.pred2(F.relu(m2))
|
||||||
|
|
||||||
|
p = F.interpolate(p2, scale_factor=4, mode='bilinear', align_corners=False)
|
||||||
|
return p # , p2, p3, p4
|
||||||
|
|
||||||
|
|
||||||
|
class Memory(nn.Module):
|
||||||
|
def __init__(self):
|
||||||
|
super(Memory, self).__init__()
|
||||||
|
|
||||||
|
def forward(self, m_in, m_out, q_in, q_out): # m_in: o,c,t,h,w
|
||||||
|
B, D_e, T, H, W = m_in.size()
|
||||||
|
_, D_o, _, _, _ = m_out.size()
|
||||||
|
|
||||||
|
mi = m_in.view(B, D_e, T * H * W)
|
||||||
|
mi = torch.transpose(mi, 1, 2) # b, THW, emb
|
||||||
|
|
||||||
|
qi = q_in.view(B, D_e, H * W) # b, emb, HW
|
||||||
|
|
||||||
|
p = torch.bmm(mi, qi) # b, THW, HW
|
||||||
|
p = p / math.sqrt(D_e)
|
||||||
|
p = F.softmax(p, dim=1) # b, THW, HW
|
||||||
|
|
||||||
|
mo = m_out.view(B, D_o, T * H * W)
|
||||||
|
mem = torch.bmm(mo, p) # Weighted-sum B, D_o, HW
|
||||||
|
mem = mem.view(B, D_o, H, W)
|
||||||
|
|
||||||
|
mem_out = torch.cat([mem, q_out], dim=1)
|
||||||
|
|
||||||
|
return mem_out, p
|
||||||
|
|
||||||
|
|
||||||
|
class KeyValue(nn.Module):
|
||||||
|
# Not using location
|
||||||
|
def __init__(self, indim, keydim, valdim):
|
||||||
|
super(KeyValue, self).__init__()
|
||||||
|
self.Key = nn.Conv2d(indim, keydim, kernel_size=(3, 3), padding=(1, 1), stride=1)
|
||||||
|
self.Value = nn.Conv2d(indim, valdim, kernel_size=(3, 3), padding=(1, 1), stride=1)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
return self.Key(x), self.Value(x)
|
||||||
|
|
||||||
|
|
||||||
|
class STM(nn.Module):
|
||||||
|
def __init__(self, backbone='resnet50'):
|
||||||
|
super(STM, self).__init__()
|
||||||
|
self.backbone = backbone
|
||||||
|
assert backbone == 'resnet50' or backbone == 'resnet18' or backbone == 'resnest101'
|
||||||
|
scale_rate = (1 if (backbone == 'resnet50' or backbone == 'resnest101') else 4)
|
||||||
|
|
||||||
|
self.Encoder_M = Encoder_M(backbone)
|
||||||
|
self.Encoder_Q = Encoder_Q(backbone)
|
||||||
|
|
||||||
|
self.KV_M_r4 = KeyValue(1024 // scale_rate, keydim=128 // scale_rate, valdim=512 // scale_rate)
|
||||||
|
self.KV_Q_r4 = KeyValue(1024 // scale_rate, keydim=128 // scale_rate, valdim=512 // scale_rate)
|
||||||
|
|
||||||
|
self.Memory = Memory()
|
||||||
|
self.Decoder = Decoder(256, scale_rate, backbone)
|
||||||
|
|
||||||
|
def Pad_memory(self, mems, num_objects, K):
|
||||||
|
pad_mems = []
|
||||||
|
for mem in mems:
|
||||||
|
pad_mem = ToCuda(torch.zeros(1, K, mem.size()[1], 1, mem.size()[2], mem.size()[3]))
|
||||||
|
pad_mem[0, 1:num_objects + 1, :, 0] = mem
|
||||||
|
pad_mems.append(pad_mem)
|
||||||
|
return pad_mems
|
||||||
|
|
||||||
|
def memorize(self, frame, masks, num_objects):
|
||||||
|
# memorize a frame
|
||||||
|
num_objects = num_objects[0].item()
|
||||||
|
_, K, H, W = masks.shape # B = 1
|
||||||
|
|
||||||
|
(frame, masks), pad = pad_divide_by([frame, masks], 16, (frame.size()[2], frame.size()[3]))
|
||||||
|
|
||||||
|
# make batch arg list
|
||||||
|
B_list = {'f': [], 'm': [], 'o': []}
|
||||||
|
for o in range(1, num_objects + 1): # 1 - no
|
||||||
|
B_list['f'].append(frame)
|
||||||
|
B_list['m'].append(masks[:, o])
|
||||||
|
B_list['o'].append((torch.sum(masks[:, 1:o], dim=1) + \
|
||||||
|
torch.sum(masks[:, o + 1:num_objects + 1], dim=1)).clamp(0, 1))
|
||||||
|
|
||||||
|
# make Batch
|
||||||
|
B_ = {}
|
||||||
|
for arg in B_list.keys():
|
||||||
|
B_[arg] = torch.cat(B_list[arg], dim=0)
|
||||||
|
|
||||||
|
r4, _, _, _, _ = self.Encoder_M(B_['f'], B_['m'], B_['o'])
|
||||||
|
k4, v4 = self.KV_M_r4(r4) # num_objects, 128 and 512, H/16, W/16
|
||||||
|
k4, v4 = self.Pad_memory([k4, v4], num_objects=num_objects, K=K)
|
||||||
|
return k4, v4
|
||||||
|
|
||||||
|
def Soft_aggregation(self, ps, K):
|
||||||
|
num_objects, H, W = ps.shape
|
||||||
|
em = ToCuda(torch.zeros(1, K, H, W))
|
||||||
|
em[0, 0] = torch.prod(1 - ps, dim=0) # bg prob
|
||||||
|
em[0, 1:num_objects + 1] = ps # obj prob
|
||||||
|
em = torch.clamp(em, 1e-7, 1 - 1e-7)
|
||||||
|
logit = torch.log((em / (1 - em)))
|
||||||
|
return logit
|
||||||
|
|
||||||
|
def segment(self, frame, keys, values, num_objects):
|
||||||
|
num_objects = num_objects[0].item()
|
||||||
|
_, K, keydim, T, H, W = keys.shape # B = 1
|
||||||
|
# pad
|
||||||
|
[frame], pad = pad_divide_by([frame], 16, (frame.size()[2], frame.size()[3]))
|
||||||
|
|
||||||
|
r4, r3, r2, _, _ = self.Encoder_Q(frame)
|
||||||
|
k4, v4 = self.KV_Q_r4(r4) # 1, dim, H/16, W/16
|
||||||
|
|
||||||
|
# expand to --- no, c, h, w
|
||||||
|
k4e, v4e = k4.expand(num_objects, -1, -1, -1), v4.expand(num_objects, -1, -1, -1)
|
||||||
|
r3e, r2e = r3.expand(num_objects, -1, -1, -1), r2.expand(num_objects, -1, -1, -1)
|
||||||
|
|
||||||
|
# memory select kv:(1, K, C, T, H, W)
|
||||||
|
m4, viz = self.Memory(keys[0, 1:num_objects + 1], values[0, 1:num_objects + 1], k4e, v4e)
|
||||||
|
logits = self.Decoder(m4, r3e, r2e)
|
||||||
|
ps = F.softmax(logits, dim=1)[:, 1] # no, h, w
|
||||||
|
# ps = indipendant possibility to belong to each object
|
||||||
|
|
||||||
|
logit = self.Soft_aggregation(ps, K) # 1, K, H, W
|
||||||
|
|
||||||
|
if pad[2] + pad[3] > 0:
|
||||||
|
logit = logit[:, :, pad[2]:-pad[3], :]
|
||||||
|
if pad[0] + pad[1] > 0:
|
||||||
|
logit = logit[:, :, :, pad[0]:-pad[1]]
|
||||||
|
|
||||||
|
return logit
|
||||||
|
|
||||||
|
def forward(self, *args, **kwargs):
|
||||||
|
if args[1].dim() > 4: # keys
|
||||||
|
return self.segment(*args, **kwargs)
|
||||||
|
else:
|
||||||
|
return self.memorize(*args, **kwargs)
|
||||||
@@ -0,0 +1,185 @@
|
|||||||
|
import torch
|
||||||
|
import torch.nn.functional as F
|
||||||
|
from torch import nn
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
class SequenceConv(nn.ModuleList):
|
||||||
|
"""Sequence conv module.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
in_channels (int): input tensor channel.
|
||||||
|
out_channels (int): output tensor channel.
|
||||||
|
kernel_size (int): convolution kernel size.
|
||||||
|
sequence_num (int): sequence length.
|
||||||
|
conv_cfg (dict): convolution config dictionary.
|
||||||
|
norm_cfg (dict): normalization config dictionary.
|
||||||
|
act_cfg (dict): activation config dictionary.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, in_channels, out_channels, kernel_size, sequence_num):
|
||||||
|
super(SequenceConv, self).__init__()
|
||||||
|
self.in_channels = in_channels
|
||||||
|
self.out_channels = out_channels
|
||||||
|
self.kernel_size = kernel_size
|
||||||
|
self.sequence_num = sequence_num
|
||||||
|
for _ in range(sequence_num):
|
||||||
|
self.append(
|
||||||
|
nn.Sequential(
|
||||||
|
nn.Conv2d(self.in_channels, self.out_channels, self.kernel_size, 1, self.kernel_size // 2, bias=False),
|
||||||
|
nn.BatchNorm2d(self.out_channels),
|
||||||
|
nn.ReLU()
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, sequence_imgs):
|
||||||
|
"""
|
||||||
|
|
||||||
|
Args:
|
||||||
|
sequence_imgs (Tensor): TxBxCxHxW
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
sequence conv output: TxBxCxHxW
|
||||||
|
"""
|
||||||
|
sequence_outs = []
|
||||||
|
assert sequence_imgs.shape[0] == self.sequence_num
|
||||||
|
for i, sequence_conv in enumerate(self):
|
||||||
|
sequence_out = sequence_conv(sequence_imgs[i, ...])
|
||||||
|
sequence_out = sequence_out.unsqueeze(0)
|
||||||
|
sequence_outs.append(sequence_out)
|
||||||
|
|
||||||
|
sequence_outs = torch.cat(sequence_outs, dim=0) # TxBxCxHxW
|
||||||
|
return sequence_outs
|
||||||
|
|
||||||
|
class MemoryModule(nn.Module):
|
||||||
|
"""Memory read module.
|
||||||
|
Args:
|
||||||
|
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self,
|
||||||
|
matmul_norm=False):
|
||||||
|
super(MemoryModule, self).__init__()
|
||||||
|
self.matmul_norm = matmul_norm
|
||||||
|
|
||||||
|
def forward(self, memory_keys, memory_values, query_key, query_value):
|
||||||
|
"""
|
||||||
|
Memory Module forward.
|
||||||
|
Args:
|
||||||
|
memory_keys (Tensor): memory keys tensor, shape: TxBxCxHxW
|
||||||
|
memory_values (Tensor): memory values tensor, shape: TxBxCxHxW
|
||||||
|
query_key (Tensor): query keys tensor, shape: BxCxHxW
|
||||||
|
query_value (Tensor): query values tensor, shape: BxCxHxW
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Concat query and memory tensor.
|
||||||
|
"""
|
||||||
|
sequence_num, batch_size, key_channels, height, width = memory_keys.shape
|
||||||
|
_, _, value_channels, _, _ = memory_values.shape
|
||||||
|
assert query_key.shape[1] == key_channels and query_value.shape[1] == value_channels
|
||||||
|
memory_keys = memory_keys.permute(1, 2, 0, 3, 4).contiguous() # BxCxTxHxW
|
||||||
|
memory_keys = memory_keys.view(batch_size, key_channels, sequence_num * height * width) # BxCxT*H*W
|
||||||
|
|
||||||
|
query_key = query_key.view(batch_size, key_channels, height * width).permute(0, 2, 1).contiguous() # BxH*WxCk
|
||||||
|
key_attention = torch.bmm(query_key, memory_keys) # BxH*WxT*H*W
|
||||||
|
if self.matmul_norm:
|
||||||
|
key_attention = (key_channels ** -.5) * key_attention
|
||||||
|
key_attention = F.softmax(key_attention, dim=-1) # BxH*WxT*H*W
|
||||||
|
|
||||||
|
memory_values = memory_values.permute(1, 2, 0, 3, 4).contiguous() # BxCxTxHxW
|
||||||
|
memory_values = memory_values.view(batch_size, value_channels, sequence_num * height * width)
|
||||||
|
memory_values = memory_values.permute(0, 2, 1).contiguous() # BxT*H*WxC
|
||||||
|
memory = torch.bmm(key_attention, memory_values) # BxH*WxC
|
||||||
|
memory = memory.permute(0, 2, 1).contiguous() # BxCxH*W
|
||||||
|
memory = memory.view(batch_size, value_channels, height, width) # BxCxHxW
|
||||||
|
|
||||||
|
query_memory = torch.cat([query_value, memory], dim=1)
|
||||||
|
return query_memory
|
||||||
|
#
|
||||||
|
# class TMAHead(nn.Module):
|
||||||
|
# """TMAHead decoder for video semantic segmentation."""
|
||||||
|
#
|
||||||
|
# def __init__(self, sequence_num, key_channels, value_channels, num_classes=2, dropout_ratio=0):
|
||||||
|
# super(TMAHead, self).__init__()
|
||||||
|
#
|
||||||
|
# self.sequence_num = sequence_num
|
||||||
|
# self.memory_key_conv = nn.Sequential(
|
||||||
|
# SequenceConv(self.in_channels, key_channels, 1, sequence_num),
|
||||||
|
# SequenceConv(key_channels, key_channels, 3, sequence_num)
|
||||||
|
# )
|
||||||
|
# self.memory_value_conv = nn.Sequential(
|
||||||
|
# SequenceConv(self.in_channels, value_channels, 1, sequence_num),
|
||||||
|
# SequenceConv(value_channels, value_channels, 3, sequence_num)
|
||||||
|
# )
|
||||||
|
# self.query_key_conv = nn.Sequential(
|
||||||
|
# nn.Sequential(
|
||||||
|
# nn.Conv2d(self.in_channels, key_channels, 1, 1, 0, bias=False),
|
||||||
|
# nn.BatchNorm2d(key_channels),
|
||||||
|
# nn.ReLU()
|
||||||
|
# ),
|
||||||
|
# nn.Sequential(
|
||||||
|
# nn.Conv2d(key_channels, key_channels, 3, 1, 1, bias=False),
|
||||||
|
# nn.BatchNorm2d(key_channels),
|
||||||
|
# nn.ReLU()
|
||||||
|
# ),
|
||||||
|
# )
|
||||||
|
#
|
||||||
|
# self.query_value_conv = nn.Sequential(
|
||||||
|
# nn.Sequential(
|
||||||
|
# nn.Conv2d(self.in_channels, value_channels, 1, 1, 0, bias=False),
|
||||||
|
# nn.BatchNorm2d(value_channels),
|
||||||
|
# nn.ReLU()
|
||||||
|
# ),
|
||||||
|
# nn.Sequential(
|
||||||
|
# nn.Conv2d(value_channels, value_channels, 3, 1, 1, bias=False),
|
||||||
|
# nn.BatchNorm2d(value_channels),
|
||||||
|
# nn.ReLU()
|
||||||
|
# ),
|
||||||
|
# )
|
||||||
|
# self.memory_module = MemoryModule(matmul_norm=False)
|
||||||
|
# self.bottleneck = nn.Sequential(
|
||||||
|
# nn.Conv2d(value_channels * 2, self.channels, 3, 1, 1, bias=False),
|
||||||
|
# nn.BatchNorm2d(value_channels),
|
||||||
|
# nn.ReLU()
|
||||||
|
# )
|
||||||
|
#
|
||||||
|
# self.conv_seg = nn.Conv2d(self.channels, num_classes, kernel_size=1)
|
||||||
|
# if dropout_ratio > 0:
|
||||||
|
# self.dropout = nn.Dropout2d(dropout_ratio)
|
||||||
|
# else:
|
||||||
|
# self.dropout = None
|
||||||
|
#
|
||||||
|
# def cls_seg(self, feat):
|
||||||
|
# """Classify each pixel."""
|
||||||
|
# if self.dropout is not None:
|
||||||
|
# feat = self.dropout(feat)
|
||||||
|
# output = self.conv_seg(feat)
|
||||||
|
# return output
|
||||||
|
#
|
||||||
|
# def forward(self, inputs, sequence_imgs):
|
||||||
|
# """
|
||||||
|
# Forward fuction.
|
||||||
|
# Args:
|
||||||
|
# inputs (list[Tensor]): backbone multi-level outputs.
|
||||||
|
# sequence_imgs (list[Tensor]): len(sequence_imgs) is equal to batch_size,
|
||||||
|
# each element is a Tensor with shape of TxCxHxW.
|
||||||
|
#
|
||||||
|
# Returns:
|
||||||
|
# decoder logits.
|
||||||
|
# """
|
||||||
|
# x = inputs
|
||||||
|
# sequence_imgs = [y.unsqueeze(0) for y in sequence_imgs] # T, BxCxHxW
|
||||||
|
# sequence_imgs = torch.cat(sequence_imgs, dim=0) # TxBxCxHxW
|
||||||
|
# sequence_num, batch_size, channels, height, width = sequence_imgs.shape
|
||||||
|
#
|
||||||
|
# assert sequence_num == self.sequence_num
|
||||||
|
# memory_keys = self.memory_key_conv(sequence_imgs)
|
||||||
|
# memory_values = self.memory_value_conv(sequence_imgs)
|
||||||
|
# query_key = self.query_key_conv(x) # BxCxHxW
|
||||||
|
# query_value = self.query_value_conv(x) # BxCxHxW
|
||||||
|
#
|
||||||
|
# # memory read
|
||||||
|
# output = self.memory_module(memory_keys, memory_values, query_key, query_value)
|
||||||
|
# output = self.bottleneck(output)
|
||||||
|
# output = self.cls_seg(output)
|
||||||
|
#
|
||||||
|
# return output
|
||||||
@@ -0,0 +1,624 @@
|
|||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import torch.nn.functional as F
|
||||||
|
|
||||||
|
class REBNCONV(nn.Module):
|
||||||
|
def __init__(self,in_ch=3,out_ch=3,dirate=1):
|
||||||
|
super(REBNCONV,self).__init__()
|
||||||
|
|
||||||
|
self.conv_s1 = nn.Conv2d(in_ch,out_ch,3,padding=1*dirate,dilation=1*dirate)
|
||||||
|
self.bn_s1 = nn.BatchNorm2d(out_ch)
|
||||||
|
self.relu_s1 = nn.ReLU(inplace=True)
|
||||||
|
|
||||||
|
def forward(self,x):
|
||||||
|
|
||||||
|
hx = x
|
||||||
|
xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
|
||||||
|
|
||||||
|
return xout
|
||||||
|
|
||||||
|
## upsample tensor 'src' to have the same spatial size with tensor 'tar'
|
||||||
|
def _upsample_like(src,tar):
|
||||||
|
|
||||||
|
src = F.upsample(src,size=tar.shape[2:],mode='bilinear')
|
||||||
|
|
||||||
|
return src
|
||||||
|
|
||||||
|
|
||||||
|
### RSU-7 ###
|
||||||
|
class RSU7(nn.Module):#UNet07DRES(nn.Module):
|
||||||
|
|
||||||
|
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
|
||||||
|
super(RSU7,self).__init__()
|
||||||
|
|
||||||
|
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
|
||||||
|
|
||||||
|
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
|
||||||
|
self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
|
||||||
|
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||||
|
self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
|
||||||
|
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||||
|
self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
|
||||||
|
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||||
|
self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
|
||||||
|
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||||
|
self.pool5 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
|
||||||
|
self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||||
|
|
||||||
|
self.rebnconv7 = REBNCONV(mid_ch,mid_ch,dirate=2)
|
||||||
|
|
||||||
|
self.rebnconv6d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||||
|
self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||||
|
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||||
|
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||||
|
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||||
|
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
|
||||||
|
|
||||||
|
def forward(self,x):
|
||||||
|
|
||||||
|
hx = x
|
||||||
|
hxin = self.rebnconvin(hx)
|
||||||
|
|
||||||
|
hx1 = self.rebnconv1(hxin)
|
||||||
|
hx = self.pool1(hx1)
|
||||||
|
|
||||||
|
hx2 = self.rebnconv2(hx)
|
||||||
|
hx = self.pool2(hx2)
|
||||||
|
|
||||||
|
hx3 = self.rebnconv3(hx)
|
||||||
|
hx = self.pool3(hx3)
|
||||||
|
|
||||||
|
hx4 = self.rebnconv4(hx)
|
||||||
|
hx = self.pool4(hx4)
|
||||||
|
|
||||||
|
hx5 = self.rebnconv5(hx)
|
||||||
|
hx = self.pool5(hx5)
|
||||||
|
|
||||||
|
hx6 = self.rebnconv6(hx)
|
||||||
|
|
||||||
|
hx7 = self.rebnconv7(hx6)
|
||||||
|
|
||||||
|
hx6d = self.rebnconv6d(torch.cat((hx7,hx6),1))
|
||||||
|
hx6dup = _upsample_like(hx6d,hx5)
|
||||||
|
|
||||||
|
hx5d = self.rebnconv5d(torch.cat((hx6dup,hx5),1))
|
||||||
|
hx5dup = _upsample_like(hx5d,hx4)
|
||||||
|
|
||||||
|
hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1))
|
||||||
|
hx4dup = _upsample_like(hx4d,hx3)
|
||||||
|
|
||||||
|
hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
|
||||||
|
hx3dup = _upsample_like(hx3d,hx2)
|
||||||
|
|
||||||
|
hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
|
||||||
|
hx2dup = _upsample_like(hx2d,hx1)
|
||||||
|
|
||||||
|
hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
|
||||||
|
|
||||||
|
return hx1d + hxin
|
||||||
|
|
||||||
|
### RSU-6 ###
|
||||||
|
class RSU6(nn.Module):#UNet06DRES(nn.Module):
|
||||||
|
|
||||||
|
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
|
||||||
|
super(RSU6,self).__init__()
|
||||||
|
|
||||||
|
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
|
||||||
|
|
||||||
|
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
|
||||||
|
self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
|
||||||
|
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||||
|
self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
|
||||||
|
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||||
|
self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
|
||||||
|
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||||
|
self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
|
||||||
|
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||||
|
|
||||||
|
self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=2)
|
||||||
|
|
||||||
|
self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||||
|
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||||
|
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||||
|
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||||
|
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
|
||||||
|
|
||||||
|
def forward(self,x):
|
||||||
|
|
||||||
|
hx = x
|
||||||
|
|
||||||
|
hxin = self.rebnconvin(hx)
|
||||||
|
|
||||||
|
hx1 = self.rebnconv1(hxin)
|
||||||
|
hx = self.pool1(hx1)
|
||||||
|
|
||||||
|
hx2 = self.rebnconv2(hx)
|
||||||
|
hx = self.pool2(hx2)
|
||||||
|
|
||||||
|
hx3 = self.rebnconv3(hx)
|
||||||
|
hx = self.pool3(hx3)
|
||||||
|
|
||||||
|
hx4 = self.rebnconv4(hx)
|
||||||
|
hx = self.pool4(hx4)
|
||||||
|
|
||||||
|
hx5 = self.rebnconv5(hx)
|
||||||
|
|
||||||
|
hx6 = self.rebnconv6(hx5)
|
||||||
|
|
||||||
|
|
||||||
|
hx5d = self.rebnconv5d(torch.cat((hx6,hx5),1))
|
||||||
|
hx5dup = _upsample_like(hx5d,hx4)
|
||||||
|
|
||||||
|
hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1))
|
||||||
|
hx4dup = _upsample_like(hx4d,hx3)
|
||||||
|
|
||||||
|
hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
|
||||||
|
hx3dup = _upsample_like(hx3d,hx2)
|
||||||
|
|
||||||
|
hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
|
||||||
|
hx2dup = _upsample_like(hx2d,hx1)
|
||||||
|
|
||||||
|
hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
|
||||||
|
|
||||||
|
return hx1d + hxin
|
||||||
|
|
||||||
|
### RSU-5 ###
|
||||||
|
class RSU5(nn.Module):#UNet05DRES(nn.Module):
|
||||||
|
|
||||||
|
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
|
||||||
|
super(RSU5,self).__init__()
|
||||||
|
|
||||||
|
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
|
||||||
|
|
||||||
|
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
|
||||||
|
self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
|
||||||
|
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||||
|
self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
|
||||||
|
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||||
|
self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
|
||||||
|
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||||
|
|
||||||
|
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=2)
|
||||||
|
|
||||||
|
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||||
|
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||||
|
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||||
|
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
|
||||||
|
|
||||||
|
def forward(self,x):
|
||||||
|
|
||||||
|
hx = x
|
||||||
|
|
||||||
|
hxin = self.rebnconvin(hx)
|
||||||
|
|
||||||
|
hx1 = self.rebnconv1(hxin)
|
||||||
|
hx = self.pool1(hx1)
|
||||||
|
|
||||||
|
hx2 = self.rebnconv2(hx)
|
||||||
|
hx = self.pool2(hx2)
|
||||||
|
|
||||||
|
hx3 = self.rebnconv3(hx)
|
||||||
|
hx = self.pool3(hx3)
|
||||||
|
|
||||||
|
hx4 = self.rebnconv4(hx)
|
||||||
|
|
||||||
|
hx5 = self.rebnconv5(hx4)
|
||||||
|
|
||||||
|
hx4d = self.rebnconv4d(torch.cat((hx5,hx4),1))
|
||||||
|
hx4dup = _upsample_like(hx4d,hx3)
|
||||||
|
|
||||||
|
hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
|
||||||
|
hx3dup = _upsample_like(hx3d,hx2)
|
||||||
|
|
||||||
|
hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
|
||||||
|
hx2dup = _upsample_like(hx2d,hx1)
|
||||||
|
|
||||||
|
hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
|
||||||
|
|
||||||
|
return hx1d + hxin
|
||||||
|
|
||||||
|
### RSU-4 ###
|
||||||
|
class RSU4(nn.Module):#UNet04DRES(nn.Module):
|
||||||
|
|
||||||
|
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
|
||||||
|
super(RSU4,self).__init__()
|
||||||
|
|
||||||
|
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
|
||||||
|
|
||||||
|
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
|
||||||
|
self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
|
||||||
|
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||||
|
self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
|
||||||
|
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||||
|
|
||||||
|
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=2)
|
||||||
|
|
||||||
|
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||||
|
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||||
|
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
|
||||||
|
|
||||||
|
def forward(self,x):
|
||||||
|
|
||||||
|
hx = x
|
||||||
|
|
||||||
|
hxin = self.rebnconvin(hx)
|
||||||
|
|
||||||
|
hx1 = self.rebnconv1(hxin)
|
||||||
|
hx = self.pool1(hx1)
|
||||||
|
|
||||||
|
hx2 = self.rebnconv2(hx)
|
||||||
|
hx = self.pool2(hx2)
|
||||||
|
|
||||||
|
hx3 = self.rebnconv3(hx)
|
||||||
|
|
||||||
|
hx4 = self.rebnconv4(hx3)
|
||||||
|
|
||||||
|
hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1))
|
||||||
|
hx3dup = _upsample_like(hx3d,hx2)
|
||||||
|
|
||||||
|
hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
|
||||||
|
hx2dup = _upsample_like(hx2d,hx1)
|
||||||
|
|
||||||
|
hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
|
||||||
|
|
||||||
|
return hx1d + hxin
|
||||||
|
|
||||||
|
### RSU-4F ###
|
||||||
|
class RSU4F(nn.Module):#UNet04FRES(nn.Module):
|
||||||
|
|
||||||
|
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
|
||||||
|
super(RSU4F,self).__init__()
|
||||||
|
|
||||||
|
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
|
||||||
|
|
||||||
|
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
|
||||||
|
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=2)
|
||||||
|
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=4)
|
||||||
|
|
||||||
|
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=8)
|
||||||
|
|
||||||
|
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=4)
|
||||||
|
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=2)
|
||||||
|
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
|
||||||
|
|
||||||
|
def forward(self,x):
|
||||||
|
|
||||||
|
hx = x
|
||||||
|
|
||||||
|
hxin = self.rebnconvin(hx)
|
||||||
|
|
||||||
|
hx1 = self.rebnconv1(hxin)
|
||||||
|
hx2 = self.rebnconv2(hx1)
|
||||||
|
hx3 = self.rebnconv3(hx2)
|
||||||
|
|
||||||
|
hx4 = self.rebnconv4(hx3)
|
||||||
|
|
||||||
|
hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1))
|
||||||
|
hx2d = self.rebnconv2d(torch.cat((hx3d,hx2),1))
|
||||||
|
hx1d = self.rebnconv1d(torch.cat((hx2d,hx1),1))
|
||||||
|
|
||||||
|
return hx1d + hxin
|
||||||
|
|
||||||
|
from core.faceseg.tma import SequenceConv, MemoryModule
|
||||||
|
|
||||||
|
##### U^2-Net ####
|
||||||
|
class U2NET(nn.Module):
|
||||||
|
|
||||||
|
def __init__(self, in_ch=3, out_ch=1):
|
||||||
|
super(U2NET, self).__init__()
|
||||||
|
|
||||||
|
self.stage1 = RSU7(in_ch,32,64)
|
||||||
|
self.pool12 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
|
||||||
|
self.stage2 = RSU6(64,32,128)
|
||||||
|
self.pool23 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
|
||||||
|
self.stage3 = RSU5(128,64,256)
|
||||||
|
self.pool34 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
|
||||||
|
self.stage4 = RSU4(256,128,512)
|
||||||
|
self.pool45 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
|
||||||
|
self.stage5 = RSU4F(512,256,512)
|
||||||
|
self.pool56 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
|
||||||
|
self.stage6 = RSU4F(512,256,512)
|
||||||
|
|
||||||
|
# decoder
|
||||||
|
self.stage5d = RSU4F(1024,256,512)
|
||||||
|
self.stage4d = RSU4(1024,128,256)
|
||||||
|
self.stage3d = RSU5(512,64,128)
|
||||||
|
self.stage2d = RSU6(256,32,64)
|
||||||
|
self.stage1d = RSU7(128,16,64)
|
||||||
|
|
||||||
|
self.side1 = nn.Conv2d(64,out_ch,3,padding=1)
|
||||||
|
self.side2 = nn.Conv2d(64,out_ch,3,padding=1)
|
||||||
|
self.side3 = nn.Conv2d(128,out_ch,3,padding=1)
|
||||||
|
self.side4 = nn.Conv2d(256,out_ch,3,padding=1)
|
||||||
|
self.side5 = nn.Conv2d(512,out_ch,3,padding=1)
|
||||||
|
self.side6 = nn.Conv2d(512,out_ch,3,padding=1)
|
||||||
|
|
||||||
|
self.outconv = nn.Conv2d(6*out_ch,out_ch,1)
|
||||||
|
|
||||||
|
self.in_channels = 512
|
||||||
|
key_channels = 128
|
||||||
|
value_channels = 512
|
||||||
|
self.sequence_num = sequence_num = 2
|
||||||
|
self.memory_key_conv = nn.Sequential(
|
||||||
|
SequenceConv(self.in_channels, key_channels, 1, sequence_num),
|
||||||
|
SequenceConv(key_channels, key_channels, 3, sequence_num)
|
||||||
|
)
|
||||||
|
self.memory_value_conv = nn.Sequential(
|
||||||
|
SequenceConv(self.in_channels, value_channels, 1, sequence_num),
|
||||||
|
SequenceConv(value_channels, value_channels, 3, sequence_num)
|
||||||
|
)
|
||||||
|
self.query_key_conv = nn.Sequential(
|
||||||
|
nn.Sequential(
|
||||||
|
nn.Conv2d(self.in_channels, key_channels, 1, 1, 0, bias=False),
|
||||||
|
nn.BatchNorm2d(key_channels),
|
||||||
|
nn.ReLU()
|
||||||
|
),
|
||||||
|
nn.Sequential(
|
||||||
|
nn.Conv2d(key_channels, key_channels, 3, 1, 1, bias=False),
|
||||||
|
nn.BatchNorm2d(key_channels),
|
||||||
|
nn.ReLU()
|
||||||
|
),
|
||||||
|
)
|
||||||
|
self.query_value_conv = nn.Sequential(
|
||||||
|
nn.Sequential(
|
||||||
|
nn.Conv2d(self.in_channels, value_channels, 1, 1, 0, bias=False),
|
||||||
|
nn.BatchNorm2d(value_channels),
|
||||||
|
nn.ReLU()
|
||||||
|
),
|
||||||
|
nn.Sequential(
|
||||||
|
nn.Conv2d(value_channels, value_channels, 3, 1, 1, bias=False),
|
||||||
|
nn.BatchNorm2d(value_channels),
|
||||||
|
nn.ReLU()
|
||||||
|
),
|
||||||
|
)
|
||||||
|
self.memory_module = MemoryModule(matmul_norm=False)
|
||||||
|
self.bottleneck = nn.Sequential(
|
||||||
|
nn.Conv2d(value_channels * 2, self.in_channels, 3, 1, 1, bias=False),
|
||||||
|
nn.BatchNorm2d(value_channels),
|
||||||
|
nn.ReLU()
|
||||||
|
)
|
||||||
|
|
||||||
|
self.is_train = True
|
||||||
|
|
||||||
|
def extract_feature(self, x):
|
||||||
|
hx = x
|
||||||
|
|
||||||
|
# stage 1
|
||||||
|
hx1 = self.stage1(hx)
|
||||||
|
hx = self.pool12(hx1)
|
||||||
|
|
||||||
|
# stage 2
|
||||||
|
hx2 = self.stage2(hx)
|
||||||
|
hx = self.pool23(hx2)
|
||||||
|
|
||||||
|
# stage 3
|
||||||
|
hx3 = self.stage3(hx)
|
||||||
|
hx = self.pool34(hx3)
|
||||||
|
|
||||||
|
# stage 4
|
||||||
|
hx4 = self.stage4(hx)
|
||||||
|
hx = self.pool45(hx4)
|
||||||
|
|
||||||
|
# stage 5
|
||||||
|
hx5 = self.stage5(hx)
|
||||||
|
hx = self.pool56(hx5)
|
||||||
|
|
||||||
|
# stage 6
|
||||||
|
hx6 = self.stage6(hx)
|
||||||
|
|
||||||
|
return hx1, hx2, hx3, hx4, hx5, hx6
|
||||||
|
|
||||||
|
def decoder(self, hx1, hx2, hx3, hx4, hx5, hx6):
|
||||||
|
hx6up = _upsample_like(hx6, hx5)
|
||||||
|
|
||||||
|
# -------------------- decoder --------------------
|
||||||
|
hx5d = self.stage5d(torch.cat((hx6up, hx5), 1))
|
||||||
|
hx5dup = _upsample_like(hx5d, hx4)
|
||||||
|
|
||||||
|
hx4d = self.stage4d(torch.cat((hx5dup, hx4), 1))
|
||||||
|
hx4dup = _upsample_like(hx4d, hx3)
|
||||||
|
|
||||||
|
hx3d = self.stage3d(torch.cat((hx4dup, hx3), 1))
|
||||||
|
hx3dup = _upsample_like(hx3d, hx2)
|
||||||
|
|
||||||
|
hx2d = self.stage2d(torch.cat((hx3dup, hx2), 1))
|
||||||
|
hx2dup = _upsample_like(hx2d, hx1)
|
||||||
|
|
||||||
|
hx1d = self.stage1d(torch.cat((hx2dup, hx1), 1))
|
||||||
|
|
||||||
|
return hx1d, hx2d, hx3d, hx4d, hx5d
|
||||||
|
|
||||||
|
def side_output(self, hx1d, hx2d, hx3d, hx4d, hx5d, hx6):
|
||||||
|
# side output
|
||||||
|
d1 = self.side1(hx1d)
|
||||||
|
|
||||||
|
d2 = self.side2(hx2d)
|
||||||
|
d2 = _upsample_like(d2, d1)
|
||||||
|
|
||||||
|
d3 = self.side3(hx3d)
|
||||||
|
d3 = _upsample_like(d3, d1)
|
||||||
|
|
||||||
|
d4 = self.side4(hx4d)
|
||||||
|
d4 = _upsample_like(d4, d1)
|
||||||
|
|
||||||
|
d5 = self.side5(hx5d)
|
||||||
|
d5 = _upsample_like(d5, d1)
|
||||||
|
|
||||||
|
d6 = self.side6(hx6)
|
||||||
|
d6 = _upsample_like(d6, d1)
|
||||||
|
|
||||||
|
d0 = self.outconv(torch.cat((d1, d2, d3, d4, d5, d6), 1))
|
||||||
|
|
||||||
|
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
|
||||||
|
|
||||||
|
def forward(self, x, memory_sequence=None):
|
||||||
|
if self.is_train:
|
||||||
|
hx1, hx2, hx3, hx4, hx5, hx6 = self.extract_feature(x)
|
||||||
|
|
||||||
|
if memory_sequence is None:
|
||||||
|
memory_hx6 = [hx6 for _ in range(self.sequence_num)]
|
||||||
|
else:
|
||||||
|
memory_hx6 = []
|
||||||
|
for single_memory in memory_sequence:
|
||||||
|
_, _, _, _, _, hx6 = self.extract_feature(single_memory)
|
||||||
|
memory_hx6.append(hx6)
|
||||||
|
|
||||||
|
memory_hx6 = [mhx6.unsqueeze(0) for mhx6 in memory_hx6] # T, BxCxHxW
|
||||||
|
memory_hx6 = torch.cat(memory_hx6, dim=0)
|
||||||
|
memory_keys = self.memory_key_conv(memory_hx6)
|
||||||
|
memory_values = self.memory_value_conv(memory_hx6)
|
||||||
|
query_key = self.query_key_conv(hx6)
|
||||||
|
query_value = self.query_value_conv(hx6)
|
||||||
|
merge_hx6 = self.memory_module(memory_keys, memory_values, query_key, query_value)
|
||||||
|
merge_hx6 = self.bottleneck(merge_hx6)
|
||||||
|
|
||||||
|
hx1d, hx2d, hx3d, hx4d, hx5d = self.decoder(hx1, hx2, hx3, hx4, hx5, merge_hx6)
|
||||||
|
|
||||||
|
return self.side_output(hx1d, hx2d, hx3d, hx4d, hx5d, hx6)
|
||||||
|
else:
|
||||||
|
return self.test(x)
|
||||||
|
|
||||||
|
def test(self, x):
|
||||||
|
with torch.no_grad():
|
||||||
|
hx1, hx2, hx3, hx4, hx5, hx6 = self.extract_feature(x)
|
||||||
|
|
||||||
|
memory_hx6 = [hx6 for _ in range(self.sequence_num)]
|
||||||
|
|
||||||
|
memory_hx6 = [mhx6.unsqueeze(0) for mhx6 in memory_hx6] # T, BxCxHxW
|
||||||
|
memory_hx6 = torch.cat(memory_hx6, dim=0)
|
||||||
|
memory_keys = self.memory_key_conv(memory_hx6)
|
||||||
|
memory_values = self.memory_value_conv(memory_hx6)
|
||||||
|
query_key = self.query_key_conv(hx6)
|
||||||
|
query_value = self.query_value_conv(hx6)
|
||||||
|
merge_hx6 = self.memory_module(memory_keys, memory_values, query_key, query_value)
|
||||||
|
merge_hx6 = self.bottleneck(merge_hx6)
|
||||||
|
|
||||||
|
hx1d, hx2d, hx3d, hx4d, hx5d = self.decoder(hx1, hx2, hx3, hx4, hx5, merge_hx6)
|
||||||
|
|
||||||
|
mask, _, _, _, _, _, _ = self.side_output(hx1d, hx2d, hx3d, hx4d, hx5d, hx6)
|
||||||
|
|
||||||
|
return mask
|
||||||
|
#
|
||||||
|
# ### U^2-Net small ###
|
||||||
|
# class U2NETP(nn.Module):
|
||||||
|
#
|
||||||
|
# def __init__(self,in_ch=3,out_ch=1):
|
||||||
|
# super(U2NETP,self).__init__()
|
||||||
|
#
|
||||||
|
# self.stage1 = RSU7(in_ch,16,64)
|
||||||
|
# self.pool12 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
#
|
||||||
|
# self.stage2 = RSU6(64,16,64)
|
||||||
|
# self.pool23 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
#
|
||||||
|
# self.stage3 = RSU5(64,16,64)
|
||||||
|
# self.pool34 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
#
|
||||||
|
# self.stage4 = RSU4(64,16,64)
|
||||||
|
# self.pool45 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
#
|
||||||
|
# self.stage5 = RSU4F(64,16,64)
|
||||||
|
# self.pool56 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||||
|
#
|
||||||
|
# self.stage6 = RSU4F(64,16,64)
|
||||||
|
#
|
||||||
|
# # decoder
|
||||||
|
# self.stage5d = RSU4F(128,16,64)
|
||||||
|
# self.stage4d = RSU4(128,16,64)
|
||||||
|
# self.stage3d = RSU5(128,16,64)
|
||||||
|
# self.stage2d = RSU6(128,16,64)
|
||||||
|
# self.stage1d = RSU7(128,16,64)
|
||||||
|
#
|
||||||
|
# self.side1 = nn.Conv2d(64,out_ch,3,padding=1)
|
||||||
|
# self.side2 = nn.Conv2d(64,out_ch,3,padding=1)
|
||||||
|
# self.side3 = nn.Conv2d(64,out_ch,3,padding=1)
|
||||||
|
# self.side4 = nn.Conv2d(64,out_ch,3,padding=1)
|
||||||
|
# self.side5 = nn.Conv2d(64,out_ch,3,padding=1)
|
||||||
|
# self.side6 = nn.Conv2d(64,out_ch,3,padding=1)
|
||||||
|
#
|
||||||
|
# self.outconv = nn.Conv2d(6*out_ch,out_ch,1)
|
||||||
|
#
|
||||||
|
# def forward(self,x):
|
||||||
|
#
|
||||||
|
# hx = x
|
||||||
|
#
|
||||||
|
# #stage 1
|
||||||
|
# hx1 = self.stage1(hx)
|
||||||
|
# hx = self.pool12(hx1)
|
||||||
|
#
|
||||||
|
# #stage 2
|
||||||
|
# hx2 = self.stage2(hx)
|
||||||
|
# hx = self.pool23(hx2)
|
||||||
|
#
|
||||||
|
# #stage 3
|
||||||
|
# hx3 = self.stage3(hx)
|
||||||
|
# hx = self.pool34(hx3)
|
||||||
|
#
|
||||||
|
# #stage 4
|
||||||
|
# hx4 = self.stage4(hx)
|
||||||
|
# hx = self.pool45(hx4)
|
||||||
|
#
|
||||||
|
# #stage 5
|
||||||
|
# hx5 = self.stage5(hx)
|
||||||
|
# hx = self.pool56(hx5)
|
||||||
|
#
|
||||||
|
# #stage 6
|
||||||
|
# hx6 = self.stage6(hx)
|
||||||
|
# hx6up = _upsample_like(hx6,hx5)
|
||||||
|
#
|
||||||
|
# #decoder
|
||||||
|
# hx5d = self.stage5d(torch.cat((hx6up,hx5),1))
|
||||||
|
# hx5dup = _upsample_like(hx5d,hx4)
|
||||||
|
#
|
||||||
|
# hx4d = self.stage4d(torch.cat((hx5dup,hx4),1))
|
||||||
|
# hx4dup = _upsample_like(hx4d,hx3)
|
||||||
|
#
|
||||||
|
# hx3d = self.stage3d(torch.cat((hx4dup,hx3),1))
|
||||||
|
# hx3dup = _upsample_like(hx3d,hx2)
|
||||||
|
#
|
||||||
|
# hx2d = self.stage2d(torch.cat((hx3dup,hx2),1))
|
||||||
|
# hx2dup = _upsample_like(hx2d,hx1)
|
||||||
|
#
|
||||||
|
# hx1d = self.stage1d(torch.cat((hx2dup,hx1),1))
|
||||||
|
#
|
||||||
|
#
|
||||||
|
# #side output
|
||||||
|
# d1 = self.side1(hx1d)
|
||||||
|
#
|
||||||
|
# d2 = self.side2(hx2d)
|
||||||
|
# d2 = _upsample_like(d2,d1)
|
||||||
|
#
|
||||||
|
# d3 = self.side3(hx3d)
|
||||||
|
# d3 = _upsample_like(d3,d1)
|
||||||
|
#
|
||||||
|
# d4 = self.side4(hx4d)
|
||||||
|
# d4 = _upsample_like(d4,d1)
|
||||||
|
#
|
||||||
|
# d5 = self.side5(hx5d)
|
||||||
|
# d5 = _upsample_like(d5,d1)
|
||||||
|
#
|
||||||
|
# d6 = self.side6(hx6)
|
||||||
|
# d6 = _upsample_like(d6,d1)
|
||||||
|
#
|
||||||
|
# d0 = self.outconv(torch.cat((d1,d2,d3,d4,d5,d6),1))
|
||||||
|
#
|
||||||
|
# return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
|
||||||
@@ -0,0 +1,210 @@
|
|||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import os.path as osp
|
||||||
|
from common.logger import LogFactory
|
||||||
|
from core.process_modules import Get_Landmark, Process_Data, localtranslationwarpfastwithstrength,\
|
||||||
|
Generator_Hair, chinClass, Generator_Fusion_Res, Change_Hair_Color, GenderClassifyProcessor, BodySeg, \
|
||||||
|
localtranslationwarpfastwithstrength_v2, localtranslationwarpfastwithstrength_v2_soft, updateEndPosition
|
||||||
|
from core.face_enhance.face_enhancement import FaceEnhancement
|
||||||
|
from datetime import datetime
|
||||||
|
import random
|
||||||
|
from core.utils import landmark_processor
|
||||||
|
from core.cos_module import COS_object as OSS_object
|
||||||
|
from core.faceseg.face_seg import FaceSeg
|
||||||
|
from common.logger import config
|
||||||
|
|
||||||
|
class Prepare_Ref_HairColor_Data(object):
|
||||||
|
def __init__(self, gpu, device_id):
|
||||||
|
if gpu and torch.cuda.is_available():
|
||||||
|
self.device = torch.device("cuda:%d" % device_id if device_id >= 0 else "cpu")
|
||||||
|
self.cuda = True
|
||||||
|
else:
|
||||||
|
self.device = torch.device("cpu")
|
||||||
|
self.cuda = False
|
||||||
|
|
||||||
|
self.process_data = Process_Data(gpu, device_id)
|
||||||
|
self.get_landmark = Get_Landmark(gpu_id=device_id)
|
||||||
|
self.color_output_size = 768
|
||||||
|
|
||||||
|
def get_prepare_ref_haircolor_768_data(self, ref_rgb_8uc3_orisize):
|
||||||
|
|
||||||
|
ref_landmark_1k2_f_orisize, _, _ = self.get_landmark.forward(ref_rgb_8uc3_orisize)
|
||||||
|
|
||||||
|
ref_matte_fg_8uc3_orisize, ref_matte_pred_8uc1_orisize, _ = self.process_data.generator_matte.matte_inference(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize)
|
||||||
|
|
||||||
|
# 光头分割
|
||||||
|
ref_baldseg_8uc3_orisize = self.process_data.generator_baldseg.forward(ref_rgb_8uc3_orisize, ref_matte_pred_8uc1_orisize,
|
||||||
|
ref_landmark_1k2_f_orisize)
|
||||||
|
|
||||||
|
color_hair_M = landmark_processor.get_transform_mat_hair_ratio_v1(ref_landmark_1k2_f_orisize, self.color_output_size, ratio=0.35, h_offset=0.45)
|
||||||
|
|
||||||
|
ref_matte_pred_8uc3_orisize = np.repeat(ref_matte_pred_8uc1_orisize[:, :, np.newaxis], 3, axis=2)
|
||||||
|
|
||||||
|
# show_concat = np.concatenate((ref_rgb_8uc3_orisize, ref_matte_pred_8uc3_orisize, ref_baldseg_8uc3_orisize), axis=1)
|
||||||
|
# show_concat = cv2.resize(show_concat, (0, 0), fx=0.5, fy=0.5)
|
||||||
|
# cv2.imshow("show_concat", show_concat)
|
||||||
|
# cv2.waitKey()
|
||||||
|
|
||||||
|
ref_rgb_8uc3_768 = cv2.warpAffine(ref_rgb_8uc3_orisize, color_hair_M, (self.color_output_size, self.color_output_size))
|
||||||
|
ref_matte_pred_8uc3_768 = cv2.warpAffine(ref_matte_pred_8uc3_orisize, color_hair_M, (self.color_output_size, self.color_output_size))
|
||||||
|
ref_baldseg_8uc3_768 = cv2.warpAffine(ref_baldseg_8uc3_orisize, color_hair_M, (self.color_output_size, self.color_output_size), flags=cv2.INTER_NEAREST)
|
||||||
|
ref_landmark_1k2_f_768 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, color_hair_M)
|
||||||
|
|
||||||
|
return ref_rgb_8uc3_768, ref_matte_pred_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_1k2_f_768
|
||||||
|
|
||||||
|
class Prepare_Ref_HairStyle_Data(object):
|
||||||
|
def __init__(self, gpu, device_id):
|
||||||
|
|
||||||
|
if gpu and torch.cuda.is_available():
|
||||||
|
self.device = torch.device("cuda:%d" % device_id if device_id >= 0 else "cpu")
|
||||||
|
self.cuda = True
|
||||||
|
else:
|
||||||
|
self.device = torch.device("cpu")
|
||||||
|
self.cuda = False
|
||||||
|
|
||||||
|
self.process_data = Process_Data(gpu, device_id)
|
||||||
|
self.get_landmark = Get_Landmark(gpu_id=device_id)
|
||||||
|
self.gender_classify = GenderClassifyProcessor(gpu_id=device_id)
|
||||||
|
|
||||||
|
def get_prepare_ref_768_color_data(self, ref_rgb_8uc3_orisize):
|
||||||
|
ref_landmark_1k2_f_orisize = self.get_landmark.forward(ref_rgb_8uc3_orisize)
|
||||||
|
ref_rgb_8uc3_color_768, ref_matting_8uc3_color_768, ref_baldseg_8uc3_color_768, ref_landmark_f1k2_color_768 = \
|
||||||
|
self.process_data.get_prepare_ref_768_bald_data(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize)
|
||||||
|
|
||||||
|
return ref_rgb_8uc3_color_768, ref_matting_8uc3_color_768, ref_baldseg_8uc3_color_768, ref_landmark_f1k2_color_768
|
||||||
|
|
||||||
|
def get_prepare_ref_768_color_data_landmark1k(self, ref_rgb_8uc3_orisize,ref_landmark_1k2_f_orisize):
|
||||||
|
# ref_landmark_1k2_f_orisize = self.get_landmark.inference(ref_rgb_8uc3_orisize)
|
||||||
|
ref_rgb_8uc3_color_768, ref_matting_8uc3_color_768, ref_baldseg_8uc3_color_768, ref_landmark_f1k2_color_768 = \
|
||||||
|
self.process_data.get_prepare_ref_768_bald_data(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize)
|
||||||
|
|
||||||
|
return ref_rgb_8uc3_color_768, ref_matting_8uc3_color_768, ref_baldseg_8uc3_color_768, ref_landmark_f1k2_color_768
|
||||||
|
|
||||||
|
def calculate_hair_ratio_after_align(self, hair_mask, origin_landmark1k, img_size=768):
|
||||||
|
image_to_face_mat = landmark_processor.get_transform_mat_hair_ratio_v1(origin_landmark1k, 768, ratio=0.35, h_offset=0.32)
|
||||||
|
hair_mask_align = cv2.warpAffine(hair_mask, image_to_face_mat, (img_size, img_size))
|
||||||
|
|
||||||
|
hair_rect = cv2.boundingRect(hair_mask_align[:, :, :1])
|
||||||
|
hair_mask_ratio = hair_rect[2] * hair_rect[3] / (img_size * img_size)
|
||||||
|
return hair_mask_ratio
|
||||||
|
|
||||||
|
def check_female_hair_ratio(self, origin_img_8uc3, landmark_1k2_f_orisize):
|
||||||
|
ref_matte_fg_8uc3_orisize, ref_matte_pred_8uc1_orisize, _ = self.process_data.generator_matte.matte_inference(origin_img_8uc3, landmark_1k2_f_orisize)
|
||||||
|
ref_matte_pred_8uc3_orisize = np.repeat(ref_matte_pred_8uc1_orisize[:, :, np.newaxis], 3, axis=2)
|
||||||
|
hairstyle_M = self.process_data.get_hair_M_girl_v1(landmark_1k2_f_orisize)
|
||||||
|
ref_rgb_8uc3_768 = cv2.warpAffine(ref_matte_pred_8uc3_orisize, hairstyle_M, (768, 768))
|
||||||
|
|
||||||
|
# cv2.imshow("ref_rgb_8uc3_768", ref_rgb_8uc3_768)
|
||||||
|
# cv2.waitKey()
|
||||||
|
|
||||||
|
edge_width = 5
|
||||||
|
if (ref_rgb_8uc3_768[-edge_width:, :, :]).max() > 0 or (ref_rgb_8uc3_768[:, -edge_width:, :]).max() > 0 or (ref_rgb_8uc3_768[:, :edge_width, :]).max() > 0:
|
||||||
|
return False
|
||||||
|
else:
|
||||||
|
return True
|
||||||
|
|
||||||
|
def get_prepare_ref_768_data(self, ref_rgb_8uc3_orisize):
|
||||||
|
|
||||||
|
ref_landmark_1k2_f_orisize, _, _ = self.get_landmark.forward(ref_rgb_8uc3_orisize)
|
||||||
|
# for i in range(1000):
|
||||||
|
# cv2.circle(ref_rgb_8uc3_orisize, (int(ref_landmark_1k2_f_orisize[i][0]),int(ref_landmark_1k2_f_orisize[i][1])), 1,(255, 0,0), 1)
|
||||||
|
# cv2.imshow('ffff', ref_rgb_8uc3_orisize)
|
||||||
|
# cv2.waitKey()
|
||||||
|
# ref_matte_fg_8uc3_orisize, ref_matte_pred_8uc1_orisize = self.process_data.generator_matte.matte_inference(
|
||||||
|
# ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize)
|
||||||
|
ref_landmark_137kpts_f_orisize = landmark_processor.pts_1k_to_137(ref_landmark_1k2_f_orisize)
|
||||||
|
gender_res = self.gender_classify.forward(ref_rgb_8uc3_orisize, ref_landmark_137kpts_f_orisize)
|
||||||
|
# hair_ratio = self.calculate_hair_ratio_after_align(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize)
|
||||||
|
# if hair_ratio > 0.3:
|
||||||
|
# ratio = 2
|
||||||
|
# else:
|
||||||
|
# if gender_res:
|
||||||
|
# ratio = 1
|
||||||
|
# else:
|
||||||
|
# ratio = 0
|
||||||
|
if not gender_res:
|
||||||
|
ratio = 0
|
||||||
|
else:
|
||||||
|
check_res = self.check_female_hair_ratio(ref_rgb_8uc3_orisize, ref_landmark_137kpts_f_orisize)
|
||||||
|
if check_res:
|
||||||
|
ratio = 1
|
||||||
|
else:
|
||||||
|
ratio = 2
|
||||||
|
|
||||||
|
if gender_res:
|
||||||
|
gender = "girl"
|
||||||
|
else:
|
||||||
|
gender = "boy"
|
||||||
|
|
||||||
|
# show_concat = np.concatenate((ref_rgb_8uc3_orisize, ref_matte_fg_8uc3_orisize), axis=1)
|
||||||
|
# resize_ratio = 1024. / max(show_concat.shape[:2])
|
||||||
|
# show_concat = cv2.resize(show_concat, (0, 0), fx=resize_ratio, fy=resize_ratio)
|
||||||
|
# print("gender_res: ", gender_res, " hair_ratio: ", hair_ratio)
|
||||||
|
# cv2.imshow("show_concat", show_concat)
|
||||||
|
# cv2.imshow("ref_matte_pred_8uc1_orisize", ref_matte_pred_8uc1_orisize)
|
||||||
|
# cv2.waitKey()
|
||||||
|
|
||||||
|
ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768 = \
|
||||||
|
self.process_data.get_prepare_ref_768_data(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize, ratio)
|
||||||
|
|
||||||
|
# cv2.imshow("ref_rgb_8uc3_768", ref_rgb_8uc3_768)
|
||||||
|
# cv2.waitKey()
|
||||||
|
|
||||||
|
return ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768, gender, ratio
|
||||||
|
|
||||||
|
def Generator_reftensor(self, ref_rgb_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768,
|
||||||
|
ref_landmark_f1k2_768):
|
||||||
|
|
||||||
|
"""
|
||||||
|
input：
|
||||||
|
|
||||||
|
图像尺寸基于： 人脸 512， 图像均为3通道
|
||||||
|
ref_rgb_8uc3_512: 参考图， size 512， uint8 （0-255）
|
||||||
|
ref_matting_8uc3_512：参考图 matting alpha 值， uint8 （0-255）
|
||||||
|
ref_baldseg_8uc3_512: 参考图 光头分割， uint8 （0-255）
|
||||||
|
ref_landmark_f1k2_512： 参考图 关键点 1k*2 float32
|
||||||
|
|
||||||
|
output：
|
||||||
|
|
||||||
|
input_another_pose_hair_image: 参考图 条件图， float32 （0-255）
|
||||||
|
|
||||||
|
"""
|
||||||
|
|
||||||
|
# 8 ********************** another pose hair_image **********************
|
||||||
|
another_pose_image = ref_rgb_8uc3_768.copy()
|
||||||
|
|
||||||
|
another_nohair_pose_mask = ref_baldseg_8uc3_768.copy()
|
||||||
|
|
||||||
|
# cv2.imshow("another_nohair_pose_mask", another_nohair_pose_mask)
|
||||||
|
# cv2.waitKey()
|
||||||
|
|
||||||
|
another_nohair_pose_mask[(np.abs(another_nohair_pose_mask - [0, 0, 255]) < 50).all(axis=2)] = [0, 255, 255]
|
||||||
|
another_nohair_pose_mask[(another_nohair_pose_mask == [0, 0, 0]).all(axis=2)] = [255, 255, 255]
|
||||||
|
another_pose_pts137 = landmark_processor.pts_1k_to_137(ref_landmark_f1k2_768).astype(np.int32)
|
||||||
|
cv2.fillPoly(another_nohair_pose_mask,
|
||||||
|
np.concatenate((another_pose_pts137[47:35:-1], another_pose_pts137[56:64],
|
||||||
|
[another_pose_pts137[48], another_pose_pts137[22]]))[
|
||||||
|
np.newaxis, :, :], (255, 255, 0))
|
||||||
|
|
||||||
|
cv2.fillPoly(another_nohair_pose_mask,
|
||||||
|
np.concatenate((another_pose_pts137[22:37], another_pose_pts137[56:47:-1]))[np.newaxis, :, :],
|
||||||
|
(255, 0, 128))
|
||||||
|
|
||||||
|
cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[88:104][np.newaxis, :, :], (255, 0, 255)) # eye
|
||||||
|
cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[105:121][np.newaxis, :, :], (255, 0, 255)) # eye
|
||||||
|
# Label nose
|
||||||
|
cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[64:79][np.newaxis, :, :], (255, 255, 255))
|
||||||
|
# Label eyebrow
|
||||||
|
cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[121:129][np.newaxis, :, :], (255, 128, 0))
|
||||||
|
cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[129:137][np.newaxis, :, :], (255, 128, 0))
|
||||||
|
|
||||||
|
another_pose_hair_alpha = ref_matting_8uc3_768.copy() / 255.
|
||||||
|
another_pose_hair_image = (another_pose_image * another_pose_hair_alpha + another_nohair_pose_mask * (
|
||||||
|
1 - another_pose_hair_alpha)).astype(np.uint8)
|
||||||
|
|
||||||
|
input_another_pose_hair_image = another_pose_hair_image.astype(np.float32) / 255
|
||||||
|
|
||||||
|
return input_another_pose_hair_image
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,208 @@
|
|||||||
|
import os
|
||||||
|
import cv2
|
||||||
|
import argparse
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from torch.nn import functional as F
|
||||||
|
|
||||||
|
import utils
|
||||||
|
from core.matting import networks
|
||||||
|
from utils.data_preprocess import *
|
||||||
|
from time import time
|
||||||
|
|
||||||
|
def single_inference(model, image_dict, device, return_offset=False):
|
||||||
|
|
||||||
|
with torch.no_grad():
|
||||||
|
image, trimap = image_dict['image'], image_dict['trimap']
|
||||||
|
alpha_shape = image_dict['alpha_shape']
|
||||||
|
image = image.to(device)
|
||||||
|
trimap = trimap.to(device)
|
||||||
|
|
||||||
|
matte_start = time()
|
||||||
|
alpha_pred, info_dict = model(image, trimap)
|
||||||
|
matte_end = time()
|
||||||
|
# print("matte time cost : {:.4f}".format(matte_end-matte_start))
|
||||||
|
|
||||||
|
trimap_argmax = trimap.argmax(dim=1, keepdim=True)
|
||||||
|
|
||||||
|
alpha_pred[trimap_argmax == 2] = 1
|
||||||
|
alpha_pred[trimap_argmax == 0] = 0
|
||||||
|
|
||||||
|
h, w = alpha_shape
|
||||||
|
test_pred = alpha_pred[0, 0, ...].detach().cpu().numpy() * 255
|
||||||
|
test_pred = test_pred.astype(np.uint8)
|
||||||
|
test_pred = test_pred[32:h+32, 32:w+32]
|
||||||
|
|
||||||
|
if return_offset:
|
||||||
|
short_side = h if h < w else w
|
||||||
|
ratio = 512 / short_side
|
||||||
|
offset_1 = utils.flow_to_image(info_dict['offset_1'][0][0, ...].data.cpu().numpy()).astype(np.uint8)
|
||||||
|
# write softmax_scale to offset image
|
||||||
|
scale = info_dict['offset_1'][1].cpu()
|
||||||
|
offset_1 = cv2.resize(offset_1, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_NEAREST)
|
||||||
|
text = 'unknown: {:.2f}, known: {:.2f}'.format(scale[-1, 0].item(), scale[-1,1].item())
|
||||||
|
offset_1 = cv2.putText(offset_1, text, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2)
|
||||||
|
|
||||||
|
offset_2 = utils.flow_to_image(info_dict['offset_2'][0][0, ...].data.cpu().numpy()).astype(np.uint8)
|
||||||
|
# write softmax_scale to offset image
|
||||||
|
scale = info_dict['offset_2'][1].cpu()
|
||||||
|
offset_2 = cv2.resize(offset_2, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_NEAREST)
|
||||||
|
text = 'unknown: {:.2f}, known: {:.2f}'.format(scale[-1, 0].item(), scale[-1,1].item())
|
||||||
|
offset_2 = cv2.putText(offset_2, text, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2)
|
||||||
|
|
||||||
|
return test_pred, (offset_1, offset_2)
|
||||||
|
else:
|
||||||
|
return test_pred, None
|
||||||
|
|
||||||
|
def generator_tensor_dict(image, trimap):
|
||||||
|
|
||||||
|
sample = {'image': image, 'trimap': trimap, 'alpha_shape': trimap.shape}
|
||||||
|
|
||||||
|
# reshape
|
||||||
|
h, w = sample["alpha_shape"]
|
||||||
|
|
||||||
|
if h % 32 == 0 and w % 32 == 0:
|
||||||
|
padded_image = np.pad(sample['image'], ((32, 32), (32, 32), (0, 0)), mode="reflect")
|
||||||
|
padded_trimap = np.pad(sample['trimap'], ((32, 32), (32, 32)), mode="reflect")
|
||||||
|
sample['image'] = padded_image
|
||||||
|
sample['trimap'] = padded_trimap
|
||||||
|
else:
|
||||||
|
target_h = 32 * ((h - 1) // 32 + 1)
|
||||||
|
target_w = 32 * ((w - 1) // 32 + 1)
|
||||||
|
pad_h = target_h - h
|
||||||
|
pad_w = target_w - w
|
||||||
|
padded_image = np.pad(sample['image'], ((32, pad_h+32), (32, pad_w+32), (0, 0)), mode="reflect")
|
||||||
|
padded_trimap = np.pad(sample['trimap'], ((32, pad_h+32), (32, pad_w+32)), mode="reflect")
|
||||||
|
sample['image'] = padded_image
|
||||||
|
sample['trimap'] = padded_trimap
|
||||||
|
|
||||||
|
# ImageNet mean & std
|
||||||
|
mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
|
||||||
|
std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
|
||||||
|
# convert GBR images to RGB
|
||||||
|
image, trimap = sample['image'][:, :, ::-1], sample['trimap']
|
||||||
|
# swap color axis
|
||||||
|
image = image.transpose((2, 0, 1)).astype(np.float32)
|
||||||
|
trimap[trimap < 85] = 0
|
||||||
|
trimap[trimap >= 170] = 2
|
||||||
|
trimap[trimap >= 85] = 1
|
||||||
|
# normalize image
|
||||||
|
image /= 255.
|
||||||
|
|
||||||
|
# to tensor
|
||||||
|
sample['image'], sample['trimap'] = torch.from_numpy(image), torch.from_numpy(trimap).to(torch.long)
|
||||||
|
sample['image'] = sample['image'].sub_(mean).div_(std)
|
||||||
|
|
||||||
|
# if CONFIG.model.trimap_channel == 3:
|
||||||
|
sample['trimap'] = F.one_hot(sample['trimap'], num_classes=3).permute(2, 0, 1).float()
|
||||||
|
# elif CONFIG.model.trimap_channel == 1:
|
||||||
|
# sample['trimap'] = sample['trimap'][None, ...].float()
|
||||||
|
# else:
|
||||||
|
# raise NotImplementedError("CONFIG.model.trimap_channel can only be 3 or 1")
|
||||||
|
|
||||||
|
# add first channel
|
||||||
|
sample['image'], sample['trimap'] = sample['image'][None, ...], sample['trimap'][None, ...]
|
||||||
|
|
||||||
|
return sample
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
print('Torch Version: ', torch.__version__)
|
||||||
|
|
||||||
|
parser = argparse.ArgumentParser()
|
||||||
|
parser.add_argument('--checkpoint', type=str, default='/home/liyang/project/matting/GCA-Matting/checkpoints/gca-dist-align-truth-1101/best_model.pth',
|
||||||
|
help="path of checkpoint")
|
||||||
|
parser.add_argument('--image-dir', type=str, default='/home/liyang/project/matting/合格/origin', help="input image dir")
|
||||||
|
parser.add_argument('--trimap-dir', type=str, default='/home/liyang/project/matting/合格/trimap', help="input trimap dir")
|
||||||
|
parser.add_argument('--output', type=str, default='/home/liyang/project/matting/合格/pred6', help="output dir")
|
||||||
|
|
||||||
|
# Parse configuration
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
# # Check if toml config file is loaded
|
||||||
|
# if CONFIG.is_default:
|
||||||
|
# raise ValueError("No .toml config loaded.")
|
||||||
|
|
||||||
|
args.output = os.path.join(args.output, args.checkpoint.split('/')[-1])
|
||||||
|
utils.make_dir(args.output)
|
||||||
|
|
||||||
|
# build model
|
||||||
|
model = networks.get_generator(encoder="resnet_gca_encoder_29", decoder="res_gca_decoder_22", num_class=1)
|
||||||
|
model.cuda()
|
||||||
|
print("model: ", model)
|
||||||
|
|
||||||
|
# load checkpoint
|
||||||
|
checkpoint = torch.load(args.checkpoint)
|
||||||
|
model.load_state_dict(utils.remove_prefix_state_dict(checkpoint['state_dict']), strict=True)
|
||||||
|
|
||||||
|
# inference
|
||||||
|
model = model.eval()
|
||||||
|
|
||||||
|
for image_name in os.listdir(args.image_dir):
|
||||||
|
|
||||||
|
if not is_image_file(image_name):
|
||||||
|
continue
|
||||||
|
# assume image and trimap have the same file name
|
||||||
|
img_basename, ext = os.path.splitext(image_name)
|
||||||
|
|
||||||
|
image_path = os.path.join(args.image_dir, image_name)
|
||||||
|
trimap_path = os.path.join(args.trimap_dir, image_name.replace(ext, "_alpha.png"))
|
||||||
|
# trimap_path = os.path.join(args.trimap_dir, image_name)
|
||||||
|
print('Image: ', image_path, ' Tirmap: ', trimap_path)
|
||||||
|
|
||||||
|
# read images
|
||||||
|
img_basename, img_ext = os.path.splitext(image_name)
|
||||||
|
# img_pt_path = image_path.replace(img_ext, "_landmark1k.txt")
|
||||||
|
# img_landmark1k = np.loadtxt(img_pt_path)
|
||||||
|
|
||||||
|
image = cv2.imread(image_path)
|
||||||
|
trimap = cv2.imread(trimap_path, 0)
|
||||||
|
|
||||||
|
ori_h, ori_w, _ = image.shape
|
||||||
|
tri_h, tri_w = trimap.shape
|
||||||
|
if image.shape[1] != trimap.shape[1] or image.shape[0] != trimap.shape[0]:
|
||||||
|
image = cv2.resize(image, (trimap.shape[1], trimap.shape[0]), interpolation=cv2.INTER_CUBIC)
|
||||||
|
# img_landmark1k = landmark_processor.resize_points(img_landmark1k, ori_w, ori_h,
|
||||||
|
# tri_w, tri_h)
|
||||||
|
|
||||||
|
# hair_mat = landmark_processor.get_transform_mat_hair(img_landmark1k, 640, ratio=0.3, w_ratio=0.5,
|
||||||
|
# h_ratio=0.4)
|
||||||
|
# hair_img_landmark = landmark_processor.transform_points(img_landmark1k, hair_mat)
|
||||||
|
#
|
||||||
|
# image = cv2.warpAffine(image, hair_mat, (640, 640), flags=cv2.INTER_CUBIC)
|
||||||
|
# trimap = cv2.warpAffine(trimap, hair_mat, (640, 640), flags=cv2.INTER_CUBIC)
|
||||||
|
|
||||||
|
if tri_h > 1920 or tri_w > 1920:
|
||||||
|
if tri_h > 1920:
|
||||||
|
new_tri_h = 1920
|
||||||
|
new_tri_w = int(tri_w * 1920 / tri_h)
|
||||||
|
else:
|
||||||
|
new_tri_w = 1920
|
||||||
|
new_tri_h = int(tri_h * 1920 / tri_w)
|
||||||
|
|
||||||
|
image_resize = cv2.resize(image, (new_tri_w, new_tri_h), interpolation=cv2.INTER_CUBIC)
|
||||||
|
trimap_resize = cv2.resize(trimap, (new_tri_w, new_tri_h), interpolation=cv2.INTER_NEAREST)
|
||||||
|
else:
|
||||||
|
image_resize = image
|
||||||
|
trimap_resize = trimap
|
||||||
|
|
||||||
|
# image_dict = generator_tensor_dict(image, trimap)
|
||||||
|
image_dict = generator_tensor_dict(image_resize, trimap_resize)
|
||||||
|
pred, offset = single_inference(model, image_dict)
|
||||||
|
|
||||||
|
# torch.cuda.empty_cache()
|
||||||
|
|
||||||
|
pred = cv2.resize(pred, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_CUBIC)
|
||||||
|
# pred = cv2.warpAffine(pred, cv2.invertAffineTransform(hair_mat), (tri_w, tri_h), flags=cv2.INTER_CUBIC)
|
||||||
|
|
||||||
|
# offset[0] = cv2.resize(offset[0], (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
|
||||||
|
# offset[1] = cv2.resize(offset[1], (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
|
||||||
|
|
||||||
|
# cv2.imshow("image_resize", image_resize)
|
||||||
|
# cv2.imshow("trimap_resize", trimap_resize)
|
||||||
|
# cv2.imshow("pred", pred)
|
||||||
|
# cv2.waitKey()
|
||||||
|
cv2.imwrite(os.path.join(args.output, image_name.replace(img_ext, "_noalign.png")), pred)
|
||||||
|
# if offset is not None:
|
||||||
|
# cv2.imwrite(os.path.join(args.output, os.path.splitext(image_name)[0]+'_offset1.png'), offset[0])
|
||||||
|
# cv2.imwrite(os.path.join(args.output, os.path.splitext(image_name)[0]+'_offset2.png'), offset[1])
|
||||||
@@ -0,0 +1,200 @@
|
|||||||
|
import os
|
||||||
|
import cv2
|
||||||
|
import argparse
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from torch.nn import functional as F
|
||||||
|
|
||||||
|
import utils
|
||||||
|
from core.matting import networks
|
||||||
|
from utils.data_preprocess import *
|
||||||
|
from time import time
|
||||||
|
|
||||||
|
def single_inference(model, image_dict, device, return_offset=False):
|
||||||
|
|
||||||
|
with torch.no_grad():
|
||||||
|
image, trimap = image_dict['image'], image_dict['trimap']
|
||||||
|
alpha_shape = image_dict['alpha_shape']
|
||||||
|
image = image.to(device)
|
||||||
|
trimap = trimap.to(device)
|
||||||
|
|
||||||
|
matte_start = time()
|
||||||
|
|
||||||
|
alpha_pred, info_dict = model(image, trimap)
|
||||||
|
|
||||||
|
fg_pred = alpha_pred[:, :-1, :, :]
|
||||||
|
alpha_pred = alpha_pred[:, -1, :, :].unsqueeze(1)
|
||||||
|
matte_end = time()
|
||||||
|
# print("matte time cost : {:.4f}".format(matte_end-matte_start))
|
||||||
|
|
||||||
|
trimap_argmax = trimap.argmax(dim=1, keepdim=True)
|
||||||
|
|
||||||
|
alpha_pred[trimap_argmax == 2] = 1
|
||||||
|
alpha_pred[trimap_argmax == 0] = 0
|
||||||
|
|
||||||
|
h, w = alpha_shape
|
||||||
|
test_fg_pred = fg_pred[0].detach().cpu().numpy().transpose(1, 2, 0)[:, :, ::-1] * 255
|
||||||
|
test_fg_pred = test_fg_pred.astype(np.uint8)
|
||||||
|
test_fg_pred = test_fg_pred[32:h+32, 32:w+32]
|
||||||
|
|
||||||
|
test_pred = alpha_pred[0, 0, ...].detach().cpu().numpy() * 255
|
||||||
|
test_pred = test_pred.astype(np.uint8)
|
||||||
|
test_pred = test_pred[32:h+32, 32:w+32]
|
||||||
|
|
||||||
|
# cv2.imshow('test_fg_pred', test_fg_pred)
|
||||||
|
# cv2.imshow('test_pred', test_pred)
|
||||||
|
# cv2.waitKey()
|
||||||
|
|
||||||
|
if return_offset:
|
||||||
|
short_side = h if h < w else w
|
||||||
|
ratio = 512 / short_side
|
||||||
|
offset_1 = utils.flow_to_image(info_dict['offset_1'][0][0, ...].data.cpu().numpy()).astype(np.uint8)
|
||||||
|
# write softmax_scale to offset image
|
||||||
|
scale = info_dict['offset_1'][1].cpu()
|
||||||
|
offset_1 = cv2.resize(offset_1, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_NEAREST)
|
||||||
|
text = 'unknown: {:.2f}, known: {:.2f}'.format(scale[-1, 0].item(), scale[-1,1].item())
|
||||||
|
offset_1 = cv2.putText(offset_1, text, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2)
|
||||||
|
|
||||||
|
offset_2 = utils.flow_to_image(info_dict['offset_2'][0][0, ...].data.cpu().numpy()).astype(np.uint8)
|
||||||
|
# write softmax_scale to offset image
|
||||||
|
scale = info_dict['offset_2'][1].cpu()
|
||||||
|
offset_2 = cv2.resize(offset_2, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_NEAREST)
|
||||||
|
text = 'unknown: {:.2f}, known: {:.2f}'.format(scale[-1, 0].item(), scale[-1,1].item())
|
||||||
|
offset_2 = cv2.putText(offset_2, text, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2)
|
||||||
|
|
||||||
|
return test_fg_pred, test_pred, (offset_1, offset_2)
|
||||||
|
else:
|
||||||
|
return test_fg_pred, test_pred, None
|
||||||
|
|
||||||
|
def generator_tensor_dict(image, trimap):
|
||||||
|
|
||||||
|
sample = {'image': image, 'trimap': trimap, 'alpha_shape': trimap.shape}
|
||||||
|
|
||||||
|
# reshape
|
||||||
|
h, w = sample["alpha_shape"]
|
||||||
|
|
||||||
|
if h % 32 == 0 and w % 32 == 0:
|
||||||
|
padded_image = np.pad(sample['image'], ((32, 32), (32, 32), (0, 0)), mode="reflect")
|
||||||
|
padded_trimap = np.pad(sample['trimap'], ((32, 32), (32, 32)), mode="reflect")
|
||||||
|
sample['image'] = padded_image
|
||||||
|
sample['trimap'] = padded_trimap
|
||||||
|
else:
|
||||||
|
target_h = 32 * ((h - 1) // 32 + 1)
|
||||||
|
target_w = 32 * ((w - 1) // 32 + 1)
|
||||||
|
pad_h = target_h - h
|
||||||
|
pad_w = target_w - w
|
||||||
|
padded_image = np.pad(sample['image'], ((32, pad_h+32), (32, pad_w+32), (0, 0)), mode="reflect")
|
||||||
|
padded_trimap = np.pad(sample['trimap'], ((32, pad_h+32), (32, pad_w+32)), mode="reflect")
|
||||||
|
sample['image'] = padded_image
|
||||||
|
sample['trimap'] = padded_trimap
|
||||||
|
|
||||||
|
# ImageNet mean & std
|
||||||
|
mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
|
||||||
|
std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
|
||||||
|
# convert GBR images to RGB
|
||||||
|
image, trimap = sample['image'][:, :, ::-1], sample['trimap']
|
||||||
|
# swap color axis
|
||||||
|
image = image.transpose((2, 0, 1)).astype(np.float32)
|
||||||
|
trimap[trimap < 85] = 0
|
||||||
|
trimap[trimap >= 170] = 2
|
||||||
|
trimap[trimap >= 85] = 1
|
||||||
|
# normalize image
|
||||||
|
image /= 255.
|
||||||
|
|
||||||
|
# to tensor
|
||||||
|
sample['image'], sample['trimap'] = torch.from_numpy(image), torch.from_numpy(trimap).to(torch.long)
|
||||||
|
sample['image'] = sample['image'].sub_(mean).div_(std)
|
||||||
|
|
||||||
|
sample['trimap'] = F.one_hot(sample['trimap'], num_classes=3).permute(2, 0, 1).float()
|
||||||
|
|
||||||
|
# add first channel
|
||||||
|
sample['image'], sample['trimap'] = sample['image'][None, ...], sample['trimap'][None, ...]
|
||||||
|
|
||||||
|
return sample
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
print('Torch Version: ', torch.__version__)
|
||||||
|
|
||||||
|
parser = argparse.ArgumentParser()
|
||||||
|
parser.add_argument('--checkpoint', type=str, default='/home/liyang/project/matting/GCA-Matting/checkpoints/gca-dist-align-truth-1101/best_model.pth',
|
||||||
|
help="path of checkpoint")
|
||||||
|
parser.add_argument('--image-dir', type=str, default='/home/liyang/project/matting/合格/origin', help="input image dir")
|
||||||
|
parser.add_argument('--trimap-dir', type=str, default='/home/liyang/project/matting/合格/trimap', help="input trimap dir")
|
||||||
|
parser.add_argument('--output', type=str, default='/home/liyang/project/matting/合格/pred6', help="output dir")
|
||||||
|
|
||||||
|
# Parse configuration
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
# # Check if toml config file is loaded
|
||||||
|
# if CONFIG.is_default:
|
||||||
|
# raise ValueError("No .toml config loaded.")
|
||||||
|
|
||||||
|
args.output = os.path.join(args.output, args.checkpoint.split('/')[-1])
|
||||||
|
utils.make_dir(args.output)
|
||||||
|
|
||||||
|
# build model
|
||||||
|
model = networks.get_generator(encoder=CONFIG.model.arch.encoder, decoder=CONFIG.model.arch.decoder)
|
||||||
|
model.cuda()
|
||||||
|
print("model: ", model)
|
||||||
|
|
||||||
|
# load checkpoint
|
||||||
|
checkpoint = torch.load(args.checkpoint)
|
||||||
|
model.load_state_dict(utils.remove_prefix_state_dict(checkpoint['state_dict']), strict=True)
|
||||||
|
|
||||||
|
# inference
|
||||||
|
model = model.eval()
|
||||||
|
export_onnx_file = "test.onnx"
|
||||||
|
torch.onnx.export(model, x, export_onnx_file, opset_version=10, do_constant_folding=True, input_names=["image", "trimap"], # 输入名
|
||||||
|
output_names=["fg_pred", "alpha_pred", "None"])
|
||||||
|
|
||||||
|
for image_name in os.listdir(args.image_dir):
|
||||||
|
|
||||||
|
if not is_image_file(image_name):
|
||||||
|
continue
|
||||||
|
# assume image and trimap have the same file name
|
||||||
|
img_basename, ext = os.path.splitext(image_name)
|
||||||
|
|
||||||
|
image_path = os.path.join(args.image_dir, image_name)
|
||||||
|
trimap_path = os.path.join(args.trimap_dir, image_name.replace(ext, "_alpha.png"))
|
||||||
|
# trimap_path = os.path.join(args.trimap_dir, image_name)
|
||||||
|
print('Image: ', image_path, ' Tirmap: ', trimap_path)
|
||||||
|
|
||||||
|
# read images
|
||||||
|
img_basename, img_ext = os.path.splitext(image_name)
|
||||||
|
|
||||||
|
image = cv2.imread(image_path)
|
||||||
|
trimap = cv2.imread(trimap_path, 0)
|
||||||
|
|
||||||
|
ori_h, ori_w, _ = image.shape
|
||||||
|
tri_h, tri_w = trimap.shape
|
||||||
|
if image.shape[1] != trimap.shape[1] or image.shape[0] != trimap.shape[0]:
|
||||||
|
image = cv2.resize(image, (trimap.shape[1], trimap.shape[0]), interpolation=cv2.INTER_CUBIC)
|
||||||
|
|
||||||
|
|
||||||
|
if tri_h > 1920 or tri_w > 1920:
|
||||||
|
if tri_h > 1920:
|
||||||
|
new_tri_h = 1920
|
||||||
|
new_tri_w = int(tri_w * 1920 / tri_h)
|
||||||
|
else:
|
||||||
|
new_tri_w = 1920
|
||||||
|
new_tri_h = int(tri_h * 1920 / tri_w)
|
||||||
|
|
||||||
|
image_resize = cv2.resize(image, (new_tri_w, new_tri_h), interpolation=cv2.INTER_CUBIC)
|
||||||
|
trimap_resize = cv2.resize(trimap, (new_tri_w, new_tri_h), interpolation=cv2.INTER_NEAREST)
|
||||||
|
else:
|
||||||
|
image_resize = image
|
||||||
|
trimap_resize = trimap
|
||||||
|
|
||||||
|
# image_dict = generator_tensor_dict(image, trimap)
|
||||||
|
image_dict = generator_tensor_dict(image_resize, trimap_resize)
|
||||||
|
fg_pred, pred, offset = single_inference(model, image_dict)
|
||||||
|
|
||||||
|
# torch.cuda.empty_cache()
|
||||||
|
|
||||||
|
fg_pred = cv2.resize(fg_pred, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_CUBIC)
|
||||||
|
pred = cv2.resize(pred, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_CUBIC)
|
||||||
|
|
||||||
|
cv2.imwrite(os.path.join(args.output, image_name.replace(img_ext, "_noalign.png")), pred)
|
||||||
|
cv2.imwrite(os.path.join(args.output, image_name.replace(img_ext, "_fg_p.png")), fg_pred)
|
||||||
|
|
||||||
@@ -0,0 +1 @@
|
|||||||
|
from .generators import *
|
||||||
@@ -0,0 +1,28 @@
|
|||||||
|
from .resnet_dec import ResNet_D_Dec, BasicBlock
|
||||||
|
from .res_shortcut_dec import ResShortCut_D_Dec
|
||||||
|
from .res_gca_dec import ResGuidedCxtAtten_Dec
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = ['res_shortcut_decoder_22', 'res_gca_decoder_22']
|
||||||
|
|
||||||
|
|
||||||
|
def _res_shortcut_D_dec(block, layers, **kwargs):
|
||||||
|
model = ResShortCut_D_Dec(block, layers, **kwargs)
|
||||||
|
return model
|
||||||
|
|
||||||
|
|
||||||
|
def _res_gca_D_dec(block, layers, num_class, **kwargs):
|
||||||
|
model = ResGuidedCxtAtten_Dec(block, layers, num_class, **kwargs)
|
||||||
|
return model
|
||||||
|
|
||||||
|
|
||||||
|
def res_shortcut_decoder_22(**kwargs):
|
||||||
|
"""Constructs a resnet_encoder_14 model.
|
||||||
|
"""
|
||||||
|
return _res_shortcut_D_dec(BasicBlock, [2, 3, 3, 2], **kwargs)
|
||||||
|
|
||||||
|
|
||||||
|
def res_gca_decoder_22(num_class=1, **kwargs):
|
||||||
|
"""Constructs a resnet_encoder_14 model.
|
||||||
|
"""
|
||||||
|
return _res_gca_D_dec(BasicBlock, [2, 3, 3, 2], num_class, **kwargs)
|
||||||
@@ -0,0 +1,28 @@
|
|||||||
|
from core.matting.networks.ops import GuidedCxtAtten
|
||||||
|
from core.matting.networks.decoders.res_shortcut_dec import ResShortCut_D_Dec
|
||||||
|
|
||||||
|
|
||||||
|
class ResGuidedCxtAtten_Dec(ResShortCut_D_Dec):
|
||||||
|
|
||||||
|
def __init__(self, block, layers, num_class=1, norm_layer=None, large_kernel=False):
|
||||||
|
super(ResGuidedCxtAtten_Dec, self).__init__(block, layers, num_class, norm_layer, large_kernel)
|
||||||
|
self.gca = GuidedCxtAtten(128, 128)
|
||||||
|
self.num_class = num_class
|
||||||
|
|
||||||
|
def forward(self, x, mid_fea):
|
||||||
|
fea1, fea2, fea3, fea4, fea5 = mid_fea['shortcut']
|
||||||
|
im = mid_fea['image_fea']
|
||||||
|
x = self.layer1(x) + fea5 # N x 256 x 32 x 32
|
||||||
|
x = self.layer2(x) + fea4 # N x 128 x 64 x 64
|
||||||
|
x, offset = self.gca(im, x, mid_fea['unknown']) # contextual attention
|
||||||
|
x = self.layer3(x) + fea3 # N x 64 x 128 x 128
|
||||||
|
x = self.layer4(x) + fea2 # N x 32 x 256 x 256
|
||||||
|
x = self.conv1(x)
|
||||||
|
x = self.bn1(x)
|
||||||
|
x = self.leaky_relu(x) + fea1
|
||||||
|
x = self.conv2(x)
|
||||||
|
|
||||||
|
alpha = (self.tanh(x) + 1.0) / 2.0
|
||||||
|
|
||||||
|
return alpha, {'offset_1': mid_fea['offset_1'], 'offset_2': offset}
|
||||||
|
|
||||||
@@ -0,0 +1,24 @@
|
|||||||
|
from core.matting.networks.decoders.resnet_dec import ResNet_D_Dec
|
||||||
|
|
||||||
|
|
||||||
|
class ResShortCut_D_Dec(ResNet_D_Dec):
|
||||||
|
|
||||||
|
def __init__(self, block, layers, num_class=1, norm_layer=None, large_kernel=False, late_downsample=False):
|
||||||
|
super(ResShortCut_D_Dec, self).__init__(block, layers, num_class, norm_layer, large_kernel,
|
||||||
|
late_downsample=late_downsample)
|
||||||
|
|
||||||
|
def forward(self, x, mid_fea):
|
||||||
|
fea1, fea2, fea3, fea4, fea5 = mid_fea['shortcut']
|
||||||
|
x = self.layer1(x) + fea5
|
||||||
|
x = self.layer2(x) + fea4
|
||||||
|
x = self.layer3(x) + fea3
|
||||||
|
x = self.layer4(x) + fea2
|
||||||
|
x = self.conv1(x)
|
||||||
|
x = self.bn1(x)
|
||||||
|
x = self.leaky_relu(x) + fea1
|
||||||
|
x = self.conv2(x)
|
||||||
|
|
||||||
|
alpha = (self.tanh(x) + 1.0) / 2.0
|
||||||
|
|
||||||
|
return alpha, None
|
||||||
|
|
||||||
@@ -0,0 +1,142 @@
|
|||||||
|
import logging
|
||||||
|
import torch.nn as nn
|
||||||
|
from core.matting.networks.ops import SpectralNorm
|
||||||
|
|
||||||
|
def conv5x5(in_planes, out_planes, stride=1, groups=1, dilation=1):
|
||||||
|
"""5x5 convolution with padding"""
|
||||||
|
return nn.Conv2d(in_planes, out_planes, kernel_size=5, stride=stride,
|
||||||
|
padding=2, groups=groups, bias=False, dilation=dilation)
|
||||||
|
|
||||||
|
|
||||||
|
def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
|
||||||
|
"""3x3 convolution with padding"""
|
||||||
|
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
|
||||||
|
padding=dilation, groups=groups, bias=False, dilation=dilation)
|
||||||
|
|
||||||
|
|
||||||
|
def conv1x1(in_planes, out_planes, stride=1):
|
||||||
|
"""1x1 convolution"""
|
||||||
|
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
|
||||||
|
|
||||||
|
|
||||||
|
class BasicBlock(nn.Module):
|
||||||
|
expansion = 1
|
||||||
|
|
||||||
|
def __init__(self, inplanes, planes, stride=1, upsample=None, norm_layer=None, large_kernel=False):
|
||||||
|
super(BasicBlock, self).__init__()
|
||||||
|
if norm_layer is None:
|
||||||
|
norm_layer = nn.BatchNorm2d
|
||||||
|
self.stride = stride
|
||||||
|
conv = conv5x5 if large_kernel else conv3x3
|
||||||
|
# Both self.conv1 and self.downsample layers downsample the input when stride != 1
|
||||||
|
if self.stride > 1:
|
||||||
|
self.conv1 = SpectralNorm(nn.ConvTranspose2d(inplanes, inplanes, kernel_size=4, stride=2, padding=1, bias=False))
|
||||||
|
else:
|
||||||
|
self.conv1 = SpectralNorm(conv(inplanes, inplanes))
|
||||||
|
self.bn1 = norm_layer(inplanes)
|
||||||
|
self.activation = nn.LeakyReLU(0.2, inplace=True)
|
||||||
|
self.conv2 = SpectralNorm(conv(inplanes, planes))
|
||||||
|
self.bn2 = norm_layer(planes)
|
||||||
|
self.upsample = upsample
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
identity = x
|
||||||
|
|
||||||
|
out = self.conv1(x)
|
||||||
|
out = self.bn1(out)
|
||||||
|
out = self.activation(out)
|
||||||
|
|
||||||
|
out = self.conv2(out)
|
||||||
|
out = self.bn2(out)
|
||||||
|
|
||||||
|
if self.upsample is not None:
|
||||||
|
identity = self.upsample(x)
|
||||||
|
|
||||||
|
out += identity
|
||||||
|
out = self.activation(out)
|
||||||
|
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
class ResNet_D_Dec(nn.Module):
|
||||||
|
|
||||||
|
def __init__(self, block, layers, num_class=1, norm_layer=None, large_kernel=False, late_downsample=False):
|
||||||
|
super(ResNet_D_Dec, self).__init__()
|
||||||
|
self.logger = logging.getLogger("Logger")
|
||||||
|
if norm_layer is None:
|
||||||
|
norm_layer = nn.BatchNorm2d
|
||||||
|
self._norm_layer = norm_layer
|
||||||
|
self.large_kernel = large_kernel
|
||||||
|
self.kernel_size = 5 if self.large_kernel else 3
|
||||||
|
|
||||||
|
self.inplanes = 512 if layers[0] > 0 else 256
|
||||||
|
self.late_downsample = late_downsample
|
||||||
|
self.midplanes = 64 if late_downsample else 32
|
||||||
|
|
||||||
|
self.conv1 = SpectralNorm(nn.ConvTranspose2d(self.midplanes, 32, kernel_size=4, stride=2, padding=1, bias=False))
|
||||||
|
self.bn1 = norm_layer(32)
|
||||||
|
self.leaky_relu = nn.LeakyReLU(0.2, inplace=True)
|
||||||
|
self.conv2 = nn.Conv2d(32, num_class, kernel_size=self.kernel_size, stride=1, padding=self.kernel_size//2)
|
||||||
|
self.upsample = nn.UpsamplingNearest2d(scale_factor=2)
|
||||||
|
self.tanh = nn.Tanh()
|
||||||
|
self.layer1 = self._make_layer(block, 256, layers[0], stride=2)
|
||||||
|
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
|
||||||
|
self.layer3 = self._make_layer(block, 64, layers[2], stride=2)
|
||||||
|
self.layer4 = self._make_layer(block, self.midplanes, layers[3], stride=2)
|
||||||
|
|
||||||
|
for m in self.modules():
|
||||||
|
if isinstance(m, nn.Conv2d):
|
||||||
|
if hasattr(m, "weight_bar"):
|
||||||
|
nn.init.xavier_uniform_(m.weight_bar)
|
||||||
|
else:
|
||||||
|
nn.init.xavier_uniform_(m.weight)
|
||||||
|
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
|
||||||
|
nn.init.constant_(m.weight, 1)
|
||||||
|
nn.init.constant_(m.bias, 0)
|
||||||
|
|
||||||
|
# Zero-initialize the last BN in each residual branch,
|
||||||
|
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
|
||||||
|
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
|
||||||
|
for m in self.modules():
|
||||||
|
if isinstance(m, BasicBlock):
|
||||||
|
nn.init.constant_(m.bn2.weight, 0)
|
||||||
|
|
||||||
|
self.logger.debug(self)
|
||||||
|
|
||||||
|
def _make_layer(self, block, planes, blocks, stride=1):
|
||||||
|
if blocks == 0:
|
||||||
|
return nn.Sequential(nn.Identity())
|
||||||
|
norm_layer = self._norm_layer
|
||||||
|
upsample = None
|
||||||
|
if stride != 1:
|
||||||
|
upsample = nn.Sequential(
|
||||||
|
nn.UpsamplingNearest2d(scale_factor=2),
|
||||||
|
SpectralNorm(conv1x1(self.inplanes, planes * block.expansion)),
|
||||||
|
norm_layer(planes * block.expansion),
|
||||||
|
)
|
||||||
|
elif self.inplanes != planes * block.expansion:
|
||||||
|
upsample = nn.Sequential(
|
||||||
|
SpectralNorm(conv1x1(self.inplanes, planes * block.expansion)),
|
||||||
|
norm_layer(planes * block.expansion),
|
||||||
|
)
|
||||||
|
|
||||||
|
layers = [block(self.inplanes, planes, stride, upsample, norm_layer, self.large_kernel)]
|
||||||
|
self.inplanes = planes * block.expansion
|
||||||
|
for _ in range(1, blocks):
|
||||||
|
layers.append(block(self.inplanes, planes, norm_layer=norm_layer, large_kernel=self.large_kernel))
|
||||||
|
|
||||||
|
return nn.Sequential(*layers)
|
||||||
|
|
||||||
|
def forward(self, x, mid_fea):
|
||||||
|
x = self.layer1(x) # N x 256 x 32 x 32
|
||||||
|
x = self.layer2(x) # N x 128 x 64 x 64
|
||||||
|
x = self.layer3(x) # N x 64 x 128 x 128
|
||||||
|
x = self.layer4(x) # N x 32 x 256 x 256
|
||||||
|
x = self.conv1(x)
|
||||||
|
x = self.bn1(x)
|
||||||
|
x = self.leaky_relu(x)
|
||||||
|
x = self.conv2(x)
|
||||||
|
|
||||||
|
alpha = (self.tanh(x) + 1.0) / 2.0
|
||||||
|
|
||||||
|
return alpha, None
|
||||||
@@ -0,0 +1,39 @@
|
|||||||
|
import logging
|
||||||
|
from .resnet_enc import ResNet_D, BasicBlock
|
||||||
|
from .res_shortcut_enc import ResShortCut_D
|
||||||
|
from .res_gca_enc import ResGuidedCxtAtten
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = ['res_shortcut_encoder_29', 'resnet_gca_encoder_29']
|
||||||
|
|
||||||
|
|
||||||
|
def _res_shortcut_D(block, layers, **kwargs):
|
||||||
|
model = ResShortCut_D(block, layers, **kwargs)
|
||||||
|
return model
|
||||||
|
|
||||||
|
|
||||||
|
def _res_gca_D(block, layers, **kwargs):
|
||||||
|
model = ResGuidedCxtAtten(block, layers, **kwargs)
|
||||||
|
return model
|
||||||
|
|
||||||
|
|
||||||
|
def resnet_gca_encoder_29(**kwargs):
|
||||||
|
"""Constructs a resnet_encoder_29 model.
|
||||||
|
"""
|
||||||
|
return _res_gca_D(BasicBlock, [3, 4, 4, 2], **kwargs)
|
||||||
|
|
||||||
|
|
||||||
|
def res_shortcut_encoder_29(**kwargs):
|
||||||
|
"""Constructs a resnet_encoder_25 model.
|
||||||
|
"""
|
||||||
|
return _res_shortcut_D(BasicBlock, [3, 4, 4, 2], **kwargs)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
import torch
|
||||||
|
logging.basicConfig(level=logging.DEBUG, format='[%(asctime)s] %(levelname)s: %(message)s',
|
||||||
|
datefmt='%m-%d %H:%M:%S')
|
||||||
|
resnet_encoder = res_shortcut_encoder_29()
|
||||||
|
x = torch.randn(4,6,512,512)
|
||||||
|
z = resnet_encoder(x)
|
||||||
|
print(z[0].shape)
|
||||||
@@ -0,0 +1,97 @@
|
|||||||
|
import torch.nn as nn
|
||||||
|
import torch.nn.functional as F
|
||||||
|
|
||||||
|
# from utils import CONFIG
|
||||||
|
from core.matting.networks.encoders.resnet_enc import ResNet_D
|
||||||
|
from core.matting.networks.ops import GuidedCxtAtten, SpectralNorm
|
||||||
|
|
||||||
|
class ResGuidedCxtAtten(ResNet_D):
|
||||||
|
|
||||||
|
def __init__(self, block, layers, norm_layer=None, late_downsample=False):
|
||||||
|
super(ResGuidedCxtAtten, self).__init__(block, layers, norm_layer, late_downsample=late_downsample)
|
||||||
|
first_inplane = 3 + 3
|
||||||
|
self.shortcut_inplane = [first_inplane, self.midplanes, 64, 128, 256]
|
||||||
|
self.shortcut_plane = [32, self.midplanes, 64, 128, 256]
|
||||||
|
|
||||||
|
self.shortcut = nn.ModuleList()
|
||||||
|
for stage, inplane in enumerate(self.shortcut_inplane):
|
||||||
|
self.shortcut.append(self._make_shortcut(inplane, self.shortcut_plane[stage]))
|
||||||
|
|
||||||
|
self.guidance_head = nn.Sequential(
|
||||||
|
nn.ReflectionPad2d(1),
|
||||||
|
SpectralNorm(nn.Conv2d(3, 16, kernel_size=3, padding=0, stride=2, bias=False)),
|
||||||
|
nn.ReLU(inplace=True),
|
||||||
|
self._norm_layer(16),
|
||||||
|
nn.ReflectionPad2d(1),
|
||||||
|
SpectralNorm(nn.Conv2d(16, 32, kernel_size=3, padding=0, stride=2, bias=False)),
|
||||||
|
nn.ReLU(inplace=True),
|
||||||
|
self._norm_layer(32),
|
||||||
|
nn.ReflectionPad2d(1),
|
||||||
|
SpectralNorm(nn.Conv2d(32, 128, kernel_size=3, padding=0, stride=2, bias=False)),
|
||||||
|
nn.ReLU(inplace=True),
|
||||||
|
self._norm_layer(128)
|
||||||
|
)
|
||||||
|
|
||||||
|
self.gca = GuidedCxtAtten(128, 128)
|
||||||
|
|
||||||
|
# initialize guidance head
|
||||||
|
for layers in range(len(self.guidance_head)):
|
||||||
|
m = self.guidance_head[layers]
|
||||||
|
if isinstance(m, nn.Conv2d):
|
||||||
|
if hasattr(m, "weight_bar"):
|
||||||
|
nn.init.xavier_uniform_(m.weight_bar)
|
||||||
|
elif isinstance(m, nn.BatchNorm2d):
|
||||||
|
nn.init.constant_(m.weight, 1)
|
||||||
|
nn.init.constant_(m.bias, 0)
|
||||||
|
|
||||||
|
def _make_shortcut(self, inplane, planes):
|
||||||
|
return nn.Sequential(
|
||||||
|
SpectralNorm(nn.Conv2d(inplane, planes, kernel_size=3, padding=1, bias=False)),
|
||||||
|
nn.ReLU(inplace=True),
|
||||||
|
self._norm_layer(planes),
|
||||||
|
SpectralNorm(nn.Conv2d(planes, planes, kernel_size=3, padding=1, bias=False)),
|
||||||
|
nn.ReLU(inplace=True),
|
||||||
|
self._norm_layer(planes)
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
|
||||||
|
out = self.conv1(x)
|
||||||
|
out = self.bn1(out)
|
||||||
|
out = self.activation(out)
|
||||||
|
out = self.conv2(out)
|
||||||
|
out = self.bn2(out)
|
||||||
|
x1 = self.activation(out) # N x 32 x 256 x 256
|
||||||
|
out = self.conv3(x1)
|
||||||
|
out = self.bn3(out)
|
||||||
|
out = self.activation(out)
|
||||||
|
|
||||||
|
im_fea = self.guidance_head(x[:, :3, ...]) # downsample origin image and extract features
|
||||||
|
# if CONFIG.model.trimap_channel == 3:
|
||||||
|
unknown = F.interpolate(x[:, 4:5, ...], scale_factor=1/8, mode='nearest')
|
||||||
|
# else:
|
||||||
|
# unknown = F.interpolate(x[:,3:,...].eq(1.).float(), scale_factor=1/8, mode='nearest')
|
||||||
|
|
||||||
|
x2 = self.layer1(out) # N x 64 x 128 x 128
|
||||||
|
x3= self.layer2(x2) # N x 128 x 64 x 64
|
||||||
|
x3, offset = self.gca(im_fea, x3, unknown) # contextual attention
|
||||||
|
x4 = self.layer3(x3) # N x 256 x 32 x 32
|
||||||
|
out = self.layer_bottleneck(x4) # N x 512 x 16 x 16
|
||||||
|
|
||||||
|
fea1 = self.shortcut[0](x) # input image and trimap
|
||||||
|
fea2 = self.shortcut[1](x1)
|
||||||
|
fea3 = self.shortcut[2](x2)
|
||||||
|
fea4 = self.shortcut[3](x3)
|
||||||
|
fea5 = self.shortcut[4](x4)
|
||||||
|
|
||||||
|
return out, {'shortcut': (fea1, fea2, fea3, fea4, fea5),
|
||||||
|
'image_fea': im_fea,
|
||||||
|
'unknown': unknown,
|
||||||
|
'offset_1': offset}
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
from core.matting.networks.encoders.resnet_enc import BasicBlock
|
||||||
|
m = ResGuidedCxtAtten(BasicBlock, [3, 4, 4, 2])
|
||||||
|
for m in m.modules():
|
||||||
|
print(m)
|
||||||
@@ -0,0 +1,51 @@
|
|||||||
|
import torch.nn as nn
|
||||||
|
# from utils import CONFIG
|
||||||
|
from core.matting.networks.encoders.resnet_enc import ResNet_D
|
||||||
|
from core.matting.networks.ops import SpectralNorm
|
||||||
|
|
||||||
|
class ResShortCut_D(ResNet_D):
|
||||||
|
|
||||||
|
def __init__(self, block, layers, norm_layer=None, late_downsample=False):
|
||||||
|
super(ResShortCut_D, self).__init__(block, layers, norm_layer, late_downsample=late_downsample)
|
||||||
|
first_inplane = 3 + 3
|
||||||
|
self.shortcut_inplane = [first_inplane, self.midplanes, 64, 128, 256]
|
||||||
|
self.shortcut_plane = [32, self.midplanes, 64, 128, 256]
|
||||||
|
|
||||||
|
self.shortcut = nn.ModuleList()
|
||||||
|
for stage, inplane in enumerate(self.shortcut_inplane):
|
||||||
|
self.shortcut.append(self._make_shortcut(inplane, self.shortcut_plane[stage]))
|
||||||
|
|
||||||
|
def _make_shortcut(self, inplane, planes):
|
||||||
|
return nn.Sequential(
|
||||||
|
SpectralNorm(nn.Conv2d(inplane, planes, kernel_size=3, padding=1, bias=False)),
|
||||||
|
nn.ReLU(inplace=True),
|
||||||
|
self._norm_layer(planes),
|
||||||
|
SpectralNorm(nn.Conv2d(planes, planes, kernel_size=3, padding=1, bias=False)),
|
||||||
|
nn.ReLU(inplace=True),
|
||||||
|
self._norm_layer(planes)
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
|
||||||
|
out = self.conv1(x)
|
||||||
|
out = self.bn1(out)
|
||||||
|
out = self.activation(out)
|
||||||
|
out = self.conv2(out)
|
||||||
|
out = self.bn2(out)
|
||||||
|
x1 = self.activation(out) # N x 32 x 256 x 256
|
||||||
|
out = self.conv3(x1)
|
||||||
|
out = self.bn3(out)
|
||||||
|
out = self.activation(out)
|
||||||
|
|
||||||
|
x2 = self.layer1(out) # N x 64 x 128 x 128
|
||||||
|
x3= self.layer2(x2) # N x 128 x 64 x 64
|
||||||
|
x4 = self.layer3(x3) # N x 256 x 32 x 32
|
||||||
|
out = self.layer_bottleneck(x4) # N x 512 x 16 x 16
|
||||||
|
|
||||||
|
fea1 = self.shortcut[0](x) # input image and trimap
|
||||||
|
fea2 = self.shortcut[1](x1)
|
||||||
|
fea3 = self.shortcut[2](x2)
|
||||||
|
fea4 = self.shortcut[3](x3)
|
||||||
|
fea5 = self.shortcut[4](x4)
|
||||||
|
|
||||||
|
return out, {'shortcut':(fea1, fea2, fea3, fea4, fea5), 'image':x[:,:3,...]}
|
||||||
@@ -0,0 +1,150 @@
|
|||||||
|
import logging
|
||||||
|
import torch.nn as nn
|
||||||
|
from core.matting.networks.ops import SpectralNorm
|
||||||
|
|
||||||
|
def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
|
||||||
|
"""3x3 convolution with padding"""
|
||||||
|
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
|
||||||
|
padding=dilation, groups=groups, bias=False, dilation=dilation)
|
||||||
|
|
||||||
|
|
||||||
|
def conv1x1(in_planes, out_planes, stride=1):
|
||||||
|
"""1x1 convolution"""
|
||||||
|
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
|
||||||
|
|
||||||
|
|
||||||
|
class BasicBlock(nn.Module):
|
||||||
|
expansion = 1
|
||||||
|
|
||||||
|
def __init__(self, inplanes, planes, stride=1, downsample=None, norm_layer=None):
|
||||||
|
super(BasicBlock, self).__init__()
|
||||||
|
if norm_layer is None:
|
||||||
|
norm_layer = nn.BatchNorm2d
|
||||||
|
# Both self.conv1 and self.downsample layers downsample the input when stride != 1
|
||||||
|
self.conv1 = SpectralNorm(conv3x3(inplanes, planes, stride))
|
||||||
|
self.bn1 = norm_layer(planes)
|
||||||
|
self.activation = nn.ReLU(inplace=True)
|
||||||
|
self.conv2 = SpectralNorm(conv3x3(planes, planes))
|
||||||
|
self.bn2 = norm_layer(planes)
|
||||||
|
self.downsample = downsample
|
||||||
|
self.stride = stride
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
identity = x
|
||||||
|
|
||||||
|
out = self.conv1(x)
|
||||||
|
out = self.bn1(out)
|
||||||
|
out = self.activation(out)
|
||||||
|
|
||||||
|
out = self.conv2(out)
|
||||||
|
out = self.bn2(out)
|
||||||
|
|
||||||
|
if self.downsample is not None:
|
||||||
|
identity = self.downsample(x)
|
||||||
|
|
||||||
|
out += identity
|
||||||
|
out = self.activation(out)
|
||||||
|
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
class ResNet_D(nn.Module):
|
||||||
|
"""
|
||||||
|
Implement and pre-train on ImageNet with the tricks from
|
||||||
|
https://arxiv.org/abs/1812.01187
|
||||||
|
without the mix-up part.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, block, layers, norm_layer=None, late_downsample=False):
|
||||||
|
super(ResNet_D, self).__init__()
|
||||||
|
self.logger = logging.getLogger("Logger")
|
||||||
|
if norm_layer is None:
|
||||||
|
norm_layer = nn.BatchNorm2d
|
||||||
|
self._norm_layer = norm_layer
|
||||||
|
|
||||||
|
self.inplanes = 64
|
||||||
|
self.late_downsample = late_downsample
|
||||||
|
self.midplanes = 64 if late_downsample else 32
|
||||||
|
self.start_stride = [1, 2, 1, 2] if late_downsample else [2, 1, 2, 1]
|
||||||
|
self.conv1 = SpectralNorm(nn.Conv2d(3 + 3, 32, kernel_size=3,
|
||||||
|
stride=self.start_stride[0], padding=1, bias=False))
|
||||||
|
self.conv2 = SpectralNorm(nn.Conv2d(32, self.midplanes, kernel_size=3, stride=self.start_stride[1], padding=1,
|
||||||
|
bias=False))
|
||||||
|
self.conv3 = SpectralNorm(nn.Conv2d(self.midplanes, self.inplanes, kernel_size=3, stride=self.start_stride[2],
|
||||||
|
padding=1, bias=False))
|
||||||
|
self.bn1 = norm_layer(32)
|
||||||
|
self.bn2 = norm_layer(self.midplanes)
|
||||||
|
self.bn3 = norm_layer(self.inplanes)
|
||||||
|
self.activation = nn.ReLU(inplace=True)
|
||||||
|
self.layer1 = self._make_layer(block, 64, layers[0], stride=self.start_stride[3])
|
||||||
|
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
|
||||||
|
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
|
||||||
|
self.layer_bottleneck = self._make_layer(block, 512, layers[3], stride=2)
|
||||||
|
|
||||||
|
for m in self.modules():
|
||||||
|
if isinstance(m, nn.Conv2d):
|
||||||
|
nn.init.xavier_uniform_(m.weight_bar)
|
||||||
|
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
|
||||||
|
nn.init.constant_(m.weight, 1)
|
||||||
|
nn.init.constant_(m.bias, 0)
|
||||||
|
|
||||||
|
# Zero-initialize the last BN in each residual branch,
|
||||||
|
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
|
||||||
|
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
|
||||||
|
|
||||||
|
for m in self.modules():
|
||||||
|
if isinstance(m, BasicBlock):
|
||||||
|
nn.init.constant_(m.bn2.weight, 0)
|
||||||
|
|
||||||
|
self.logger.debug("encoder conv1 weight shape: {}".format(str(self.conv1.module.weight_bar.data.shape)))
|
||||||
|
self.conv1.module.weight_bar.data[:,3:,:,:] = 0
|
||||||
|
|
||||||
|
self.logger.debug(self)
|
||||||
|
|
||||||
|
def _make_layer(self, block, planes, blocks, stride=1):
|
||||||
|
if blocks == 0:
|
||||||
|
return nn.Sequential(nn.Identity())
|
||||||
|
norm_layer = self._norm_layer
|
||||||
|
downsample = None
|
||||||
|
if stride != 1:
|
||||||
|
downsample = nn.Sequential(
|
||||||
|
nn.AvgPool2d(2, stride),
|
||||||
|
SpectralNorm(conv1x1(self.inplanes, planes * block.expansion)),
|
||||||
|
norm_layer(planes * block.expansion),
|
||||||
|
)
|
||||||
|
elif self.inplanes != planes * block.expansion:
|
||||||
|
downsample = nn.Sequential(
|
||||||
|
SpectralNorm(conv1x1(self.inplanes, planes * block.expansion, stride)),
|
||||||
|
norm_layer(planes * block.expansion),
|
||||||
|
)
|
||||||
|
|
||||||
|
layers = [block(self.inplanes, planes, stride, downsample, norm_layer)]
|
||||||
|
self.inplanes = planes * block.expansion
|
||||||
|
for _ in range(1, blocks):
|
||||||
|
layers.append(block(self.inplanes, planes, norm_layer=norm_layer))
|
||||||
|
|
||||||
|
return nn.Sequential(*layers)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x = self.conv1(x)
|
||||||
|
x = self.bn1(x)
|
||||||
|
x = self.activation(x)
|
||||||
|
x = self.conv2(x)
|
||||||
|
x = self.bn2(x)
|
||||||
|
x1 = self.activation(x) # N x 32 x 256 x 256
|
||||||
|
x = self.conv3(x1)
|
||||||
|
x = self.bn3(x)
|
||||||
|
x2 = self.activation(x) # N x 64 x 128 x 128
|
||||||
|
|
||||||
|
x3 = self.layer1(x2) # N x 64 x 128 x 128
|
||||||
|
x4 = self.layer2(x3) # N x 128 x 64 x 64
|
||||||
|
x5 = self.layer3(x4) # N x 256 x 32 x 32
|
||||||
|
x = self.layer_bottleneck(x5) # N x 512 x 16 x 16
|
||||||
|
|
||||||
|
return x, (x1, x2, x3, x4, x5)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
m = ResNet_D(BasicBlock, [3, 4, 4, 2])
|
||||||
|
for m in m.modules():
|
||||||
|
print(m._get_name())
|
||||||
@@ -0,0 +1,60 @@
|
|||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
|
||||||
|
# from utils import CONFIG
|
||||||
|
from core.matting.networks import decoders, encoders
|
||||||
|
|
||||||
|
|
||||||
|
class Generator(nn.Module):
|
||||||
|
def __init__(self, encoder, decoder, num_class=1):
|
||||||
|
|
||||||
|
super(Generator, self).__init__()
|
||||||
|
|
||||||
|
if encoder not in encoders.__all__:
|
||||||
|
raise NotImplementedError("Unknown Encoder {}".format(encoder))
|
||||||
|
self.encoder = encoders.__dict__[encoder]()
|
||||||
|
|
||||||
|
if decoder not in decoders.__all__:
|
||||||
|
raise NotImplementedError("Unknown Decoder {}".format(decoder))
|
||||||
|
self.decoder = decoders.__dict__[decoder](num_class)
|
||||||
|
|
||||||
|
def forward(self, image, trimap):
|
||||||
|
inp = torch.cat((image, trimap), dim=1)
|
||||||
|
embedding, mid_fea = self.encoder(inp)
|
||||||
|
alpha, info_dict = self.decoder(embedding, mid_fea)
|
||||||
|
|
||||||
|
return alpha, info_dict
|
||||||
|
|
||||||
|
|
||||||
|
def get_generator(encoder, decoder, num_class=1):
|
||||||
|
generator = Generator(encoder=encoder, decoder=decoder, num_class=num_class)
|
||||||
|
return generator
|
||||||
|
|
||||||
|
|
||||||
|
if __name__=="__main__":
|
||||||
|
import time
|
||||||
|
# generator = get_generator(encoder=CONFIG.model.arch.encoder, decoder=CONFIG.model.arch.decoder).cuda().train()
|
||||||
|
batch_size = 12
|
||||||
|
# generator.eval()
|
||||||
|
n_eval = 10
|
||||||
|
# pre run the model
|
||||||
|
# with torch.no_grad():
|
||||||
|
# for i in range(2):
|
||||||
|
# x = torch.rand(batch_size, 3, 512, 512, device=device)
|
||||||
|
# y = torch.rand(batch_size, 3, 512, 512, device=device)
|
||||||
|
# z = generator(x,y)
|
||||||
|
# test without GPU IO
|
||||||
|
|
||||||
|
# x = torch.zeros(batch_size, 3, 512, 512, device=device)
|
||||||
|
# y = torch.zeros(batch_size, 1, 512, 512, device=device)
|
||||||
|
x = torch.randn(batch_size, 3, 512, 512)
|
||||||
|
y = torch.randn(batch_size, 3, 512, 512)
|
||||||
|
t = time.time()
|
||||||
|
# with torch.no_grad():
|
||||||
|
# for i in range(n_eval):
|
||||||
|
# a = generator(x.cuda(),y.cuda())
|
||||||
|
# torch.cuda.synchronize()
|
||||||
|
# print(generator.__class__.__name__, 'With IO \t', f'{(time.time() - t)/n_eval/batch_size:.5f} s')
|
||||||
|
# print(generator.__class__.__name__, 'FPS \t\t', f'{1/((time.time() - t)/n_eval/batch_size):.5f} s')
|
||||||
|
# for n, p in generator.named_parameters():
|
||||||
|
# print(n)
|
||||||
@@ -0,0 +1,256 @@
|
|||||||
|
import torch
|
||||||
|
from torch import nn
|
||||||
|
from torch.nn import Parameter
|
||||||
|
from torch.autograd import Variable
|
||||||
|
from torch.nn import functional as F
|
||||||
|
|
||||||
|
|
||||||
|
def l2normalize(v, eps=1e-12):
|
||||||
|
return v / (v.norm() + eps)
|
||||||
|
|
||||||
|
|
||||||
|
class SpectralNorm(nn.Module):
|
||||||
|
"""
|
||||||
|
Based on https://github.com/heykeetae/Self-Attention-GAN/blob/master/spectral.py
|
||||||
|
and add _noupdate_u_v() for evaluation
|
||||||
|
"""
|
||||||
|
def __init__(self, module, name='weight', power_iterations=1):
|
||||||
|
super(SpectralNorm, self).__init__()
|
||||||
|
self.module = module
|
||||||
|
self.name = name
|
||||||
|
self.power_iterations = power_iterations
|
||||||
|
if not self._made_params():
|
||||||
|
self._make_params()
|
||||||
|
|
||||||
|
def _update_u_v(self):
|
||||||
|
u = getattr(self.module, self.name + "_u")
|
||||||
|
v = getattr(self.module, self.name + "_v")
|
||||||
|
w = getattr(self.module, self.name + "_bar")
|
||||||
|
|
||||||
|
height = w.data.shape[0]
|
||||||
|
for _ in range(self.power_iterations):
|
||||||
|
v.data = l2normalize(torch.mv(torch.t(w.view(height,-1).data), u.data))
|
||||||
|
u.data = l2normalize(torch.mv(w.view(height,-1).data, v.data))
|
||||||
|
|
||||||
|
sigma = u.dot(w.view(height, -1).mv(v))
|
||||||
|
setattr(self.module, self.name, w / sigma.expand_as(w))
|
||||||
|
|
||||||
|
def _noupdate_u_v(self):
|
||||||
|
u = getattr(self.module, self.name + "_u")
|
||||||
|
v = getattr(self.module, self.name + "_v")
|
||||||
|
w = getattr(self.module, self.name + "_bar")
|
||||||
|
|
||||||
|
height = w.data.shape[0]
|
||||||
|
sigma = u.dot(w.view(height, -1).mv(v))
|
||||||
|
setattr(self.module, self.name, w / sigma.expand_as(w))
|
||||||
|
|
||||||
|
def _made_params(self):
|
||||||
|
try:
|
||||||
|
u = getattr(self.module, self.name + "_u")
|
||||||
|
v = getattr(self.module, self.name + "_v")
|
||||||
|
w = getattr(self.module, self.name + "_bar")
|
||||||
|
return True
|
||||||
|
except AttributeError:
|
||||||
|
return False
|
||||||
|
|
||||||
|
def _make_params(self):
|
||||||
|
w = getattr(self.module, self.name)
|
||||||
|
|
||||||
|
height = w.data.shape[0]
|
||||||
|
width = w.view(height, -1).data.shape[1]
|
||||||
|
|
||||||
|
u = Parameter(w.data.new(height).normal_(0, 1), requires_grad=False)
|
||||||
|
v = Parameter(w.data.new(width).normal_(0, 1), requires_grad=False)
|
||||||
|
u.data = l2normalize(u.data)
|
||||||
|
v.data = l2normalize(v.data)
|
||||||
|
w_bar = Parameter(w.data)
|
||||||
|
|
||||||
|
del self.module._parameters[self.name]
|
||||||
|
|
||||||
|
self.module.register_parameter(self.name + "_u", u)
|
||||||
|
self.module.register_parameter(self.name + "_v", v)
|
||||||
|
self.module.register_parameter(self.name + "_bar", w_bar)
|
||||||
|
|
||||||
|
def forward(self, *args):
|
||||||
|
# if torch.is_grad_enabled() and self.module.training:
|
||||||
|
if self.module.training:
|
||||||
|
self._update_u_v()
|
||||||
|
else:
|
||||||
|
self._noupdate_u_v()
|
||||||
|
return self.module.forward(*args)
|
||||||
|
|
||||||
|
|
||||||
|
class GuidedCxtAtten(nn.Module):
|
||||||
|
# based on https://github.com/nbei/Deep-Flow-Guided-Video-Inpainting/blob/a6fe298fec502bfd9cbc64eb01e39f78a3262a59/models/DeepFill_Models/ops.py#L210
|
||||||
|
def __init__(self, out_channels, guidance_channels, rate=2):
|
||||||
|
super(GuidedCxtAtten, self).__init__()
|
||||||
|
self.rate = rate
|
||||||
|
self.padding = nn.ReflectionPad2d(1)
|
||||||
|
self.up_sample = nn.Upsample(scale_factor=self.rate, mode='nearest')
|
||||||
|
|
||||||
|
self.guidance_conv = nn.Conv2d(in_channels=guidance_channels, out_channels=guidance_channels//2,
|
||||||
|
kernel_size=1, stride=1, padding=0)
|
||||||
|
|
||||||
|
self.W = nn.Sequential(
|
||||||
|
nn.Conv2d(in_channels=out_channels, out_channels=out_channels,
|
||||||
|
kernel_size=1, stride=1, padding=0, bias=False),
|
||||||
|
nn.BatchNorm2d(out_channels)
|
||||||
|
)
|
||||||
|
|
||||||
|
nn.init.xavier_uniform_(self.guidance_conv.weight)
|
||||||
|
nn.init.constant_(self.guidance_conv.bias, 0)
|
||||||
|
nn.init.xavier_uniform_(self.W[0].weight)
|
||||||
|
nn.init.constant_(self.W[1].weight, 1e-3)
|
||||||
|
nn.init.constant_(self.W[1].bias, 0)
|
||||||
|
|
||||||
|
def forward(self, f, alpha, unknown=None, ksize=3, stride=1, fuse_k=3, softmax_scale=1., training=True):
|
||||||
|
|
||||||
|
f = self.guidance_conv(f)
|
||||||
|
# get shapes
|
||||||
|
raw_int_fs = list(f.size()) # N x 64 x 64 x 64
|
||||||
|
raw_int_alpha = list(alpha.size()) # N x 128 x 64 x 64
|
||||||
|
|
||||||
|
# extract patches from background with stride and rate
|
||||||
|
kernel = 2*self.rate
|
||||||
|
alpha_w = self.extract_patches(alpha, kernel=kernel, stride=self.rate)
|
||||||
|
alpha_w = alpha_w.permute(0, 2, 3, 4, 5, 1)
|
||||||
|
alpha_w = alpha_w.contiguous().view(raw_int_alpha[0], raw_int_alpha[2] // self.rate, raw_int_alpha[3] // self.rate, -1)
|
||||||
|
alpha_w = alpha_w.contiguous().view(raw_int_alpha[0], -1, kernel, kernel, raw_int_alpha[1])
|
||||||
|
alpha_w = alpha_w.permute(0, 1, 4, 2, 3)
|
||||||
|
|
||||||
|
f = F.interpolate(f, scale_factor=1/self.rate, mode='nearest')
|
||||||
|
|
||||||
|
fs = f.size() # B x 64 x 32 x 32
|
||||||
|
f_groups = torch.split(f, 1, dim=0) # Split tensors by batch dimension; tuple is returned
|
||||||
|
|
||||||
|
# from b(B*H*W*C) to w(b*k*k*c*h*w)
|
||||||
|
int_fs = list(fs)
|
||||||
|
w = self.extract_patches(f)
|
||||||
|
w = w.permute(0, 2, 3, 4, 5, 1)
|
||||||
|
w = w.contiguous().view(raw_int_fs[0], raw_int_fs[2] // self.rate, raw_int_fs[3] // self.rate, -1)
|
||||||
|
w = w.contiguous().view(raw_int_fs[0], -1, ksize, ksize, raw_int_fs[1])
|
||||||
|
w = w.permute(0, 1, 4, 2, 3)
|
||||||
|
# process mask
|
||||||
|
|
||||||
|
if unknown is not None:
|
||||||
|
unknown = unknown.clone()
|
||||||
|
unknown = F.interpolate(unknown, scale_factor=1/self.rate, mode='nearest')
|
||||||
|
assert unknown.size(2) == f.size(2), "mask should have same size as f at dim 2,3"
|
||||||
|
unknown_mean = unknown.mean(dim=[2,3])
|
||||||
|
known_mean = 1 - unknown_mean
|
||||||
|
unknown_scale = torch.clamp(torch.sqrt(unknown_mean / known_mean), 0.1, 10).to(alpha)
|
||||||
|
known_scale = torch.clamp(torch.sqrt(known_mean / unknown_mean), 0.1, 10).to(alpha)
|
||||||
|
softmax_scale = torch.cat([unknown_scale, known_scale], dim=1)
|
||||||
|
else:
|
||||||
|
unknown = torch.ones([fs[0], 1, fs[2], fs[3]]).to(alpha)
|
||||||
|
softmax_scale = torch.FloatTensor([softmax_scale, softmax_scale]).view(1,2).repeat(fs[0],1).to(alpha)
|
||||||
|
|
||||||
|
m = self.extract_patches(unknown)
|
||||||
|
|
||||||
|
m = m.permute(0, 2, 3, 4, 5, 1)
|
||||||
|
m = m.contiguous().view(raw_int_fs[0], raw_int_fs[2]//self.rate, raw_int_fs[3]//self.rate, -1)
|
||||||
|
m = m.contiguous().view(raw_int_fs[0], -1, ksize, ksize)
|
||||||
|
|
||||||
|
m = self.reduce_mean(m) # smoothing, maybe
|
||||||
|
# mask out the
|
||||||
|
mm = m.gt(0.).float() # (N, 32*32, 1, 1)
|
||||||
|
|
||||||
|
# the correlation with itself should be 0
|
||||||
|
self_mask = F.one_hot(torch.arange(fs[2] * fs[3]).view(fs[2], fs[3]).contiguous().to(alpha).long(),
|
||||||
|
num_classes=int_fs[2] * int_fs[3])
|
||||||
|
self_mask = self_mask.permute(2, 0, 1).view(1, fs[2] * fs[3], fs[2], fs[3]).float() * (-1e4)
|
||||||
|
|
||||||
|
w_groups = torch.split(w, 1, dim=0) # Split tensors by batch dimension; tuple is returned
|
||||||
|
alpha_w_groups = torch.split(alpha_w, 1, dim=0) # Split tensors by batch dimension; tuple is returned
|
||||||
|
mm_groups = torch.split(mm, 1, dim=0)
|
||||||
|
scale_group = torch.split(softmax_scale, 1, dim=0)
|
||||||
|
y = []
|
||||||
|
offsets = []
|
||||||
|
k = fuse_k
|
||||||
|
y_test = []
|
||||||
|
for xi, wi, alpha_wi, mmi, scale in zip(f_groups, w_groups, alpha_w_groups, mm_groups, scale_group):
|
||||||
|
# conv for compare
|
||||||
|
wi = wi[0]
|
||||||
|
escape_NaN = Variable(torch.FloatTensor([1e-4])).to(alpha)
|
||||||
|
wi_normed = wi / torch.max(self.l2_norm(wi), escape_NaN)
|
||||||
|
xi = F.pad(xi, (1,1,1,1), mode='reflect')
|
||||||
|
yi = F.conv2d(xi, wi_normed, stride=1, padding=0) # yi => (B=1, C=32*32, H=32, W=32)
|
||||||
|
y_test.append(yi)
|
||||||
|
# conv implementation for fuse scores to encourage large patches
|
||||||
|
yi = yi.permute(0, 2, 3, 1)
|
||||||
|
yi = yi.contiguous().view(1, fs[2], fs[3], fs[2] * fs[3])
|
||||||
|
yi = yi.permute(0, 3, 1, 2) # (B=1, C=32*32, H=32, W=32)
|
||||||
|
|
||||||
|
# softmax to match
|
||||||
|
# scale the correlation with predicted scale factor for known and unknown area
|
||||||
|
yi = yi * (scale[0,0] * mmi.gt(0.).float() + scale[0,1] * mmi.le(0.).float()) # mmi => (1, 32*32, 1, 1)
|
||||||
|
# mask itself, self-mask only applied to unknown area
|
||||||
|
yi = yi + self_mask * mmi # self_mask: (1, 32*32, 32, 32)
|
||||||
|
# for small input inference
|
||||||
|
yi = F.softmax(yi, dim=1)
|
||||||
|
|
||||||
|
_, offset = torch.max(yi, dim=1) # argmax; index
|
||||||
|
offset = torch.stack([offset // fs[3], offset % fs[3]], dim=1)
|
||||||
|
|
||||||
|
wi_center = alpha_wi[0]
|
||||||
|
|
||||||
|
if self.rate == 1:
|
||||||
|
left = (kernel) // 2
|
||||||
|
right = (kernel - 1) // 2
|
||||||
|
yi = F.pad(yi, (left, right, left, right), mode='reflect')
|
||||||
|
wi_center = wi_center.permute(1, 0, 2, 3)
|
||||||
|
yi = F.conv2d(yi, wi_center, padding=0) / 4. # (B=1, C=128, H=64, W=64)
|
||||||
|
else:
|
||||||
|
yi = F.conv_transpose2d(yi, wi_center, stride=self.rate, padding=1) / 4. # (B=1, C=128, H=64, W=64)
|
||||||
|
y.append(yi)
|
||||||
|
offsets.append(offset)
|
||||||
|
|
||||||
|
y = torch.cat(y, dim=0) # back to the mini-batch
|
||||||
|
y.contiguous().view(raw_int_alpha)
|
||||||
|
offsets = torch.cat(offsets, dim=0)
|
||||||
|
offsets = offsets.view([int_fs[0]] + [2] + int_fs[2:])
|
||||||
|
|
||||||
|
# # case1: visualize optical flow: minus current position
|
||||||
|
# h_add = Variable(torch.arange(0,float(fs[2]))).to(alpha).view([1, 1, fs[2], 1])
|
||||||
|
# h_add = h_add.expand(fs[0], 1, fs[2], fs[3])
|
||||||
|
# w_add = Variable(torch.arange(0,float(fs[3]))).to(alpha).view([1, 1, 1, fs[3]])
|
||||||
|
# w_add = w_add.expand(fs[0], 1, fs[2], fs[3])
|
||||||
|
#
|
||||||
|
# offsets = offsets - torch.cat([h_add, w_add], dim=1).long()
|
||||||
|
|
||||||
|
# case2: visualize absolute position
|
||||||
|
offsets = offsets - torch.Tensor([fs[2]//2, fs[3]//2]).view(1,2,1,1).to(alpha).long()
|
||||||
|
|
||||||
|
y = self.W(y) + alpha
|
||||||
|
|
||||||
|
return y, (offsets, softmax_scale)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def extract_patches(x, kernel=3, stride=1):
|
||||||
|
left =(kernel - stride + 1) // 2
|
||||||
|
right =(kernel - stride) // 2
|
||||||
|
x = F.pad(x, (left, right, left, right), mode='reflect')
|
||||||
|
all_patches = x.unfold(2, kernel, stride).unfold(3, kernel, stride)
|
||||||
|
|
||||||
|
return all_patches
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def reduce_mean(x):
|
||||||
|
for i in range(4):
|
||||||
|
if i <= 1:
|
||||||
|
continue
|
||||||
|
x = torch.mean(x, dim=i, keepdim=True)
|
||||||
|
return x
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def l2_norm(x):
|
||||||
|
def reduce_sum(x):
|
||||||
|
for i in range(4):
|
||||||
|
if i == 0:
|
||||||
|
continue
|
||||||
|
x = torch.sum(x, dim=i, keepdim=True)
|
||||||
|
return x
|
||||||
|
|
||||||
|
x = x**2
|
||||||
|
x = reduce_sum(x)
|
||||||
|
return torch.sqrt(x)
|
||||||
@@ -0,0 +1,102 @@
|
|||||||
|
#!/usr/bin/env python
|
||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
|
||||||
|
# @File : setup.py
|
||||||
|
# @Time : 2020/1/15
|
||||||
|
# @Author : yangchaojie (yangchaojie@immomo.com)
|
||||||
|
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
import shutil
|
||||||
|
import numpy
|
||||||
|
import tempfile
|
||||||
|
|
||||||
|
from setuptools import setup
|
||||||
|
from setuptools.extension import Extension
|
||||||
|
|
||||||
|
from Cython.Build import cythonize
|
||||||
|
from Cython.Distutils import build_ext
|
||||||
|
|
||||||
|
import platform
|
||||||
|
|
||||||
|
|
||||||
|
def get_root_path(root):
|
||||||
|
if os.path.dirname(root) in ['', '.']:
|
||||||
|
return os.path.basename(root)
|
||||||
|
else:
|
||||||
|
return get_root_path(os.path.dirname(root))
|
||||||
|
|
||||||
|
|
||||||
|
def copy_file(src, dest):
|
||||||
|
if os.path.exists(dest):
|
||||||
|
return
|
||||||
|
|
||||||
|
if not os.path.exists(os.path.dirname(dest)):
|
||||||
|
os.makedirs(os.path.dirname(dest))
|
||||||
|
if os.path.isdir(src):
|
||||||
|
shutil.copytree(src, dest)
|
||||||
|
else:
|
||||||
|
shutil.copyfile(src, dest)
|
||||||
|
|
||||||
|
|
||||||
|
def touch_init_file():
|
||||||
|
init_file_name = os.path.join(tempfile.mkdtemp(), '__init__.py')
|
||||||
|
with open(init_file_name, 'w'):
|
||||||
|
pass
|
||||||
|
return init_file_name
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
def compose_extensions(root='.'):
|
||||||
|
for file_ in os.listdir(root):
|
||||||
|
abs_file = os.path.join(root, file_)
|
||||||
|
|
||||||
|
if os.path.isfile(abs_file):
|
||||||
|
if abs_file.endswith('.py'):
|
||||||
|
extensions.append(Extension(get_root_path(abs_file) + '.*', [abs_file]))
|
||||||
|
elif abs_file.endswith('.c') or abs_file.endswith('.pyc'):
|
||||||
|
continue
|
||||||
|
else:
|
||||||
|
copy_file(abs_file, os.path.join(build_root_dir, abs_file))
|
||||||
|
if abs_file.endswith('__init__.py'):
|
||||||
|
copy_file(init_file, os.path.join(build_root_dir, abs_file))
|
||||||
|
|
||||||
|
else:
|
||||||
|
if os.path.basename(abs_file) in ignore_folders :
|
||||||
|
continue
|
||||||
|
if os.path.basename(abs_file) in conf_folders:
|
||||||
|
copy_file(abs_file, os.path.join(build_root_dir, abs_file))
|
||||||
|
compose_extensions(abs_file)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
# if __name__ == '__main__':
|
||||||
|
build_root_dir = 'build/lib.' + platform.system().lower() + '-' + platform.machine() + '-' + str(
|
||||||
|
sys.version_info.major) + '.' + str(sys.version_info.minor)
|
||||||
|
|
||||||
|
print(build_root_dir)
|
||||||
|
|
||||||
|
extensions = []
|
||||||
|
ignore_folders = ['build', 'new_ref_zao_color_0818', 'ref_hair_online_0703_local', 'ref_分好类别', '.git']
|
||||||
|
conf_folders = ['conf']
|
||||||
|
|
||||||
|
|
||||||
|
init_file = touch_init_file()
|
||||||
|
print(init_file)
|
||||||
|
|
||||||
|
|
||||||
|
compose_extensions()
|
||||||
|
os.remove(init_file)
|
||||||
|
|
||||||
|
setup(
|
||||||
|
name='moxie_hairstyle',
|
||||||
|
version='1.0',
|
||||||
|
ext_modules=cythonize(
|
||||||
|
extensions,
|
||||||
|
nthreads=16,
|
||||||
|
compiler_directives=dict(always_allow_keywords=True),
|
||||||
|
include_path=[numpy.get_include()]),
|
||||||
|
cmdclass=dict(build_ext=build_ext))
|
||||||
|
|
||||||
|
# python setup.py build_ext
|
||||||
@@ -0,0 +1,319 @@
|
|||||||
|
import torch.nn as nn
|
||||||
|
import torch.utils.model_zoo as model_zoo
|
||||||
|
import torch
|
||||||
|
import os
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
from core.utils import utils_3ddfa,params_3ddfa,landmark_processor
|
||||||
|
__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
|
||||||
|
'resnet152']
|
||||||
|
|
||||||
|
|
||||||
|
model_urls = {
|
||||||
|
'resnet18': 'http://download.pytorch.org/models/resnet18-5c106cde.pth',
|
||||||
|
'resnet34': 'http://download.pytorch.org/models/resnet34-333f7ec4.pth',
|
||||||
|
'resnet50': 'http://download.pytorch.org/models/resnet50-19c8e357.pth',
|
||||||
|
'resnet101': 'http://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
|
||||||
|
'resnet152': 'http://download.pytorch.org/models/resnet152-b121ed2d.pth',
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
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)
|
||||||
|
|
||||||
|
|
||||||
|
def conv1x1(in_planes, out_planes, stride=1):
|
||||||
|
"""1x1 convolution"""
|
||||||
|
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, 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)
|
||||||
|
self.relu = nn.ReLU(inplace=True)
|
||||||
|
self.conv2 = conv3x3(planes, planes)
|
||||||
|
self.bn2 = nn.BatchNorm2d(planes)
|
||||||
|
self.downsample = downsample
|
||||||
|
self.stride = stride
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
identity = 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:
|
||||||
|
identity = self.downsample(x)
|
||||||
|
|
||||||
|
out += identity
|
||||||
|
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 = conv1x1(inplanes, planes)
|
||||||
|
self.bn1 = nn.BatchNorm2d(planes)
|
||||||
|
self.conv2 = conv3x3(planes, planes, stride)
|
||||||
|
self.bn2 = nn.BatchNorm2d(planes)
|
||||||
|
self.conv3 = conv1x1(planes, planes * self.expansion)
|
||||||
|
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
|
||||||
|
self.relu = nn.ReLU(inplace=True)
|
||||||
|
self.downsample = downsample
|
||||||
|
self.stride = stride
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
identity = 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:
|
||||||
|
identity = self.downsample(x)
|
||||||
|
|
||||||
|
out += identity
|
||||||
|
out = self.relu(out)
|
||||||
|
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
class ResNet(nn.Module):
|
||||||
|
|
||||||
|
def __init__(self, block, layers, num_classes=1000, no_branch=False, no_activate=False, zero_init_residual=False):
|
||||||
|
super(ResNet, self).__init__()
|
||||||
|
self.inplanes = 64
|
||||||
|
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
|
||||||
|
self.bn1 = nn.BatchNorm2d(64)
|
||||||
|
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)
|
||||||
|
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
|
||||||
|
if no_activate:
|
||||||
|
if no_branch:
|
||||||
|
self.fc = nn.Sequential(*[nn.Linear(512 * block.expansion, num_classes)])
|
||||||
|
else:
|
||||||
|
self.fc_key = nn.Sequential(*[nn.Linear(256 * block.expansion, 45 * 2)])
|
||||||
|
self.fc_ctrl = nn.Sequential(*[nn.Linear(256 * block.expansion, 48 * 2)])
|
||||||
|
else:
|
||||||
|
if no_branch:
|
||||||
|
self.fc = nn.Sequential(*[nn.Linear(512 * block.expansion, num_classes), nn.Tanh()])
|
||||||
|
else:
|
||||||
|
self.fc_key = nn.Sequential(*[nn.Linear(256 * block.expansion, 45 * 2), nn.Tanh()])
|
||||||
|
self.fc_ctrl = nn.Sequential(*[nn.Linear(256 * block.expansion, 48 * 2), nn.Tanh()])
|
||||||
|
self.no_branch = no_branch
|
||||||
|
|
||||||
|
for m in self.modules():
|
||||||
|
if isinstance(m, nn.Conv2d):
|
||||||
|
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
||||||
|
elif isinstance(m, nn.BatchNorm2d):
|
||||||
|
nn.init.constant_(m.weight, 1)
|
||||||
|
nn.init.constant_(m.bias, 0)
|
||||||
|
|
||||||
|
# Zero-initialize the last BN in each residual branch,
|
||||||
|
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
|
||||||
|
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
|
||||||
|
if zero_init_residual:
|
||||||
|
for m in self.modules():
|
||||||
|
if isinstance(m, Bottleneck):
|
||||||
|
nn.init.constant_(m.bn3.weight, 0)
|
||||||
|
elif isinstance(m, BasicBlock):
|
||||||
|
nn.init.constant_(m.bn2.weight, 0)
|
||||||
|
|
||||||
|
def _make_layer(self, block, planes, blocks, stride=1):
|
||||||
|
downsample = None
|
||||||
|
if stride != 1 or self.inplanes != planes * block.expansion:
|
||||||
|
downsample = nn.Sequential(
|
||||||
|
conv1x1(self.inplanes, planes * block.expansion, stride),
|
||||||
|
nn.BatchNorm2d(planes * block.expansion),
|
||||||
|
)
|
||||||
|
|
||||||
|
layers = []
|
||||||
|
layers.append(block(self.inplanes, planes, stride, downsample))
|
||||||
|
self.inplanes = planes * block.expansion
|
||||||
|
for _ in range(1, blocks):
|
||||||
|
layers.append(block(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)
|
||||||
|
|
||||||
|
if self.no_branch:
|
||||||
|
key = self.avgpool(x)
|
||||||
|
key = key.view(key.size(0), -1)
|
||||||
|
key = self.fc(key)
|
||||||
|
return key
|
||||||
|
else:
|
||||||
|
key, ctrl = torch.chunk(x, 2, dim=1)
|
||||||
|
|
||||||
|
key = self.avgpool(key)
|
||||||
|
key = key.view(key.size(0), -1)
|
||||||
|
key = self.fc_key(key)
|
||||||
|
|
||||||
|
ctrl = self.avgpool(ctrl)
|
||||||
|
ctrl = ctrl.view(ctrl.size(0), -1)
|
||||||
|
ctrl = self.fc_ctrl(ctrl)
|
||||||
|
|
||||||
|
return key, ctrl
|
||||||
|
|
||||||
|
def resnet18(pretrained=False, **kwargs):
|
||||||
|
"""Constructs a ResNet-18 model.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||||
|
"""
|
||||||
|
model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
|
||||||
|
if pretrained:
|
||||||
|
model.load_state_dict(model_zoo.load_url(model_urls['resnet18']), strict=False)
|
||||||
|
return model
|
||||||
|
|
||||||
|
|
||||||
|
def resnet34(pretrained=False, **kwargs):
|
||||||
|
"""Constructs a ResNet-34 model.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||||
|
"""
|
||||||
|
model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
|
||||||
|
if pretrained:
|
||||||
|
model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
|
||||||
|
return model
|
||||||
|
|
||||||
|
|
||||||
|
def resnet50(pretrained=False, **kwargs):
|
||||||
|
"""Constructs a ResNet-50 model.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||||
|
"""
|
||||||
|
model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
|
||||||
|
if pretrained:
|
||||||
|
model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
|
||||||
|
return model
|
||||||
|
|
||||||
|
|
||||||
|
def resnet101(pretrained=False, **kwargs):
|
||||||
|
"""Constructs a ResNet-101 model.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||||
|
"""
|
||||||
|
model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
|
||||||
|
if pretrained:
|
||||||
|
model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
|
||||||
|
return model
|
||||||
|
|
||||||
|
|
||||||
|
def resnet152(pretrained=False, **kwargs):
|
||||||
|
"""Constructs a ResNet-152 model.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||||
|
"""
|
||||||
|
model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
|
||||||
|
if pretrained:
|
||||||
|
model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
|
||||||
|
return model
|
||||||
|
|
||||||
|
|
||||||
|
class Model_3DDFA(nn.Module):
|
||||||
|
|
||||||
|
|
||||||
|
def __init__(self, gpu_id=None):
|
||||||
|
super(Model_3DDFA, self).__init__()
|
||||||
|
|
||||||
|
self.init_status = False
|
||||||
|
self.face_alignment_net = resnet18(pretrained=True, num_classes=76, no_branch=True, no_activate=True)
|
||||||
|
self.model_path = 'weights'
|
||||||
|
|
||||||
|
# model_path, _ = os.path.split(os.path.realpath(__file__))
|
||||||
|
weights = torch.load(os.path.join(self.model_path, 'face_3ddfa.pth'), map_location=lambda storage, loc: storage)
|
||||||
|
self.load_state_dict(weights)
|
||||||
|
self.eval()
|
||||||
|
self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||||
|
self.to(self.device)
|
||||||
|
self.init_status = True
|
||||||
|
|
||||||
|
def forward(self, imgs):
|
||||||
|
pred_pose_shape_exp = self.face_alignment_net(imgs)
|
||||||
|
return pred_pose_shape_exp
|
||||||
|
|
||||||
|
def running(self):
|
||||||
|
return self.init_status
|
||||||
|
|
||||||
|
def forward_np(self, imgs):
|
||||||
|
pred_pose_shape_exp = self.face_alignment_net(imgs)
|
||||||
|
return pred_pose_shape_exp.detach().cpu().numpy()
|
||||||
|
|
||||||
|
def detect(self, images, landmarks):
|
||||||
|
dst_size = 256
|
||||||
|
with torch.no_grad():
|
||||||
|
input_numpy = np.zeros((len(images), 3, dst_size, dst_size), dtype=np.float32)
|
||||||
|
assert len(images) == len(landmarks)
|
||||||
|
all_mat = []
|
||||||
|
all_res = []
|
||||||
|
all_height = []
|
||||||
|
for ix, img in enumerate(images):
|
||||||
|
landmark = landmarks[ix]
|
||||||
|
|
||||||
|
mat = landmark_processor.get_transform_mat_full_face(landmark, dst_size)
|
||||||
|
|
||||||
|
all_mat.append(mat)
|
||||||
|
all_height.append(img.shape[0])
|
||||||
|
|
||||||
|
tmp = cv2.warpAffine(img, mat, (dst_size, dst_size))
|
||||||
|
|
||||||
|
input_numpy[ix, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
|
||||||
|
|
||||||
|
# cv2.imshow('3ddfa_', tmp)
|
||||||
|
# cv2.waitKey()
|
||||||
|
|
||||||
|
in_tensor = torch.from_numpy(input_numpy)
|
||||||
|
in_tensor = in_tensor.to(self.device)
|
||||||
|
params = self.face_alignment_net(in_tensor)
|
||||||
|
params = params.cpu().numpy()
|
||||||
|
|
||||||
|
for ix, param in enumerate(params):
|
||||||
|
param[0] = param[0] / params_3ddfa.SCALE_F
|
||||||
|
param[1:4] = param[1:4] / params_3ddfa.SCALE_ROTATE
|
||||||
|
param[4:6] = param[4:6] / params_3ddfa.SCALE_OFFSET
|
||||||
|
param[6:56] = (param[6:56] / params_3ddfa.SCALE_SHAPE)
|
||||||
|
param[56:] = (param[56:] / params_3ddfa.SCALE_EXP)
|
||||||
|
new_param = utils_3ddfa.transform_params(param, cv2.invertAffineTransform(all_mat[ix]), all_height[ix],
|
||||||
|
dst_size)
|
||||||
|
all_res.append(new_param)
|
||||||
|
|
||||||
|
return all_res
|
||||||
@@ -0,0 +1,67 @@
|
|||||||
|
import os
|
||||||
|
import torch
|
||||||
|
import torch.nn.parallel
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
from core.utils import landmark_processor
|
||||||
|
import cv2
|
||||||
|
|
||||||
|
modelRoot = "weights"
|
||||||
|
|
||||||
|
label_map = [
|
||||||
|
[0, 0, 0],
|
||||||
|
[255, 0, 0],
|
||||||
|
[0, 0, 255],
|
||||||
|
[0, 255, 0],
|
||||||
|
[0, 255, 255],
|
||||||
|
# [0, 255, 0]
|
||||||
|
]
|
||||||
|
class Generator_BaldSeg_5c(object):
|
||||||
|
def __init__(self, gpu_flag, gpu_id):
|
||||||
|
|
||||||
|
if not gpu_flag:
|
||||||
|
self.device = torch.device("cpu")
|
||||||
|
else:
|
||||||
|
self.device = torch.device('cuda:{0}'.format(gpu_id))
|
||||||
|
|
||||||
|
# load seg model
|
||||||
|
self.model_dir = modelRoot
|
||||||
|
self.pre_trained_model = os.path.join('weights', "ori_hair_checkpoint_7660_0611.pt")
|
||||||
|
self.net = torch.jit.load(self.pre_trained_model, map_location=self.device).to(self.device)
|
||||||
|
self.net.eval()
|
||||||
|
|
||||||
|
self.output_img_size = 512
|
||||||
|
self.img_ratio = 0.4
|
||||||
|
|
||||||
|
def label_to_mask(self, label_np):
|
||||||
|
label_np = label_np.astype(np.int32)[:, :, np.newaxis]
|
||||||
|
mask = np.zeros((label_np.shape[0], label_np.shape[1], 3), dtype=np.uint8)
|
||||||
|
for id, color in enumerate(label_map):
|
||||||
|
index = (label_np == id).all(axis=2)
|
||||||
|
mask[index] = color
|
||||||
|
return mask
|
||||||
|
def forward(self, image, alpha, landmarks1k):
|
||||||
|
image_to_face_mat = landmark_processor.get_transform_mat_full_face_ratio_deeplab(landmarks1k, self.output_img_size, self.img_ratio)
|
||||||
|
img_ = (image * (1 - alpha[:, :, np.newaxis].astype(np.float32) / 255)).astype(np.uint8)
|
||||||
|
img_cuted = cv2.warpAffine(img_, image_to_face_mat, (self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4, borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255])
|
||||||
|
inv_image_to_face_mat = cv2.invertAffineTransform(image_to_face_mat)
|
||||||
|
|
||||||
|
img = (img_cuted.astype(np.float32) / 255).transpose((2, 0, 1))
|
||||||
|
img = np.expand_dims(img, axis=0)
|
||||||
|
img = torch.from_numpy(img).to(self.device) #cuda(self.gpu_id)
|
||||||
|
|
||||||
|
with torch.no_grad():
|
||||||
|
output = self.net(img)
|
||||||
|
|
||||||
|
pred = output.detach().cpu().numpy().squeeze().astype(np.float32) #torch.max(output[:1], 1)[1].detach().cpu().numpy().squeeze().astype(np.float32)
|
||||||
|
mask = self.label_to_mask(pred)
|
||||||
|
|
||||||
|
# cv2.imshow("img_cuted: ", img_cuted)
|
||||||
|
# cv2.imshow("mask: ", mask)
|
||||||
|
# cv2.waitKey()
|
||||||
|
|
||||||
|
black_img = np.zeros(image.shape).astype(np.uint8)
|
||||||
|
cv2.warpAffine(mask, inv_image_to_face_mat, (image.shape[1], image.shape[0]),
|
||||||
|
dst=black_img, flags=cv2.INTER_NEAREST, borderMode=cv2.BORDER_TRANSPARENT)
|
||||||
|
|
||||||
|
return black_img
|
||||||
@@ -0,0 +1,441 @@
|
|||||||
|
import torch.nn as nn
|
||||||
|
import torch.utils.model_zoo as model_zoo
|
||||||
|
import torch
|
||||||
|
import numpy as np
|
||||||
|
import os
|
||||||
|
from core.utils import landmark_processor
|
||||||
|
from core.utils.umeyama import umeyama
|
||||||
|
import cv2
|
||||||
|
|
||||||
|
modelRoot = "weights"
|
||||||
|
|
||||||
|
__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
|
||||||
|
'resnet152']
|
||||||
|
|
||||||
|
|
||||||
|
model_urls = {
|
||||||
|
'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
|
||||||
|
'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
|
||||||
|
'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
|
||||||
|
'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
|
||||||
|
'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
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)
|
||||||
|
|
||||||
|
|
||||||
|
def conv1x1(in_planes, out_planes, stride=1):
|
||||||
|
"""1x1 convolution"""
|
||||||
|
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, 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)
|
||||||
|
self.relu = nn.ReLU(inplace=True)
|
||||||
|
self.conv2 = conv3x3(planes, planes)
|
||||||
|
self.bn2 = nn.BatchNorm2d(planes)
|
||||||
|
self.downsample = downsample
|
||||||
|
self.stride = stride
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
identity = 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:
|
||||||
|
identity = self.downsample(x)
|
||||||
|
|
||||||
|
out += identity
|
||||||
|
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 = conv1x1(inplanes, planes)
|
||||||
|
self.bn1 = nn.BatchNorm2d(planes)
|
||||||
|
self.conv2 = conv3x3(planes, planes, stride)
|
||||||
|
self.bn2 = nn.BatchNorm2d(planes)
|
||||||
|
self.conv3 = conv1x1(planes, planes * self.expansion)
|
||||||
|
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
|
||||||
|
self.relu = nn.ReLU(inplace=True)
|
||||||
|
self.downsample = downsample
|
||||||
|
self.stride = stride
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
identity = 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:
|
||||||
|
identity = self.downsample(x)
|
||||||
|
|
||||||
|
out += identity
|
||||||
|
out = self.relu(out)
|
||||||
|
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
class ResNet(nn.Module):
|
||||||
|
|
||||||
|
def __init__(self, block, layers, num_classes=1000, is_1k=False, zero_init_residual=False):
|
||||||
|
super(ResNet, self).__init__()
|
||||||
|
self.inplanes = 64
|
||||||
|
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
|
||||||
|
bias=False)
|
||||||
|
self.bn1 = nn.BatchNorm2d(64)
|
||||||
|
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)
|
||||||
|
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
|
||||||
|
if is_1k:
|
||||||
|
self.fc = nn.Sequential(*[nn.Linear(512 * block.expansion, num_classes), nn.Tanh()])
|
||||||
|
else:
|
||||||
|
self.fc_key = nn.Sequential(*[nn.Linear(256 * block.expansion, 45 * 2), nn.Tanh()])
|
||||||
|
self.fc_ctrl = nn.Sequential(*[nn.Linear(256 * block.expansion, 48 * 2), nn.Tanh()])
|
||||||
|
self.is_1k = is_1k
|
||||||
|
|
||||||
|
for m in self.modules():
|
||||||
|
if isinstance(m, nn.Conv2d):
|
||||||
|
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
||||||
|
elif isinstance(m, nn.BatchNorm2d):
|
||||||
|
nn.init.constant_(m.weight, 1)
|
||||||
|
nn.init.constant_(m.bias, 0)
|
||||||
|
|
||||||
|
# Zero-initialize the last BN in each residual branch,
|
||||||
|
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
|
||||||
|
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
|
||||||
|
if zero_init_residual:
|
||||||
|
for m in self.modules():
|
||||||
|
if isinstance(m, Bottleneck):
|
||||||
|
nn.init.constant_(m.bn3.weight, 0)
|
||||||
|
elif isinstance(m, BasicBlock):
|
||||||
|
nn.init.constant_(m.bn2.weight, 0)
|
||||||
|
|
||||||
|
def _make_layer(self, block, planes, blocks, stride=1):
|
||||||
|
downsample = None
|
||||||
|
if stride != 1 or self.inplanes != planes * block.expansion:
|
||||||
|
downsample = nn.Sequential(
|
||||||
|
conv1x1(self.inplanes, planes * block.expansion, stride),
|
||||||
|
nn.BatchNorm2d(planes * block.expansion),
|
||||||
|
)
|
||||||
|
|
||||||
|
layers = []
|
||||||
|
layers.append(block(self.inplanes, planes, stride, downsample))
|
||||||
|
self.inplanes = planes * block.expansion
|
||||||
|
for _ in range(1, blocks):
|
||||||
|
layers.append(block(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)
|
||||||
|
|
||||||
|
if self.is_1k:
|
||||||
|
key = self.avgpool(x)
|
||||||
|
key = key.view(key.size(0), -1)
|
||||||
|
key = self.fc(key)
|
||||||
|
return key
|
||||||
|
else:
|
||||||
|
key, ctrl = torch.chunk(x, 2, dim=1)
|
||||||
|
|
||||||
|
key = self.avgpool(key)
|
||||||
|
key = key.view(key.size(0), -1)
|
||||||
|
key = self.fc_key(key)
|
||||||
|
|
||||||
|
ctrl = self.avgpool(ctrl)
|
||||||
|
ctrl = ctrl.view(ctrl.size(0), -1)
|
||||||
|
ctrl = self.fc_ctrl(ctrl)
|
||||||
|
|
||||||
|
return key, ctrl
|
||||||
|
|
||||||
|
def resnet18(pretrained=False, **kwargs):
|
||||||
|
"""Constructs a ResNet-18 model.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||||
|
"""
|
||||||
|
model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
|
||||||
|
if pretrained:
|
||||||
|
model.load_state_dict(model_zoo.load_url(model_urls['resnet18']), strict=False)
|
||||||
|
return model
|
||||||
|
|
||||||
|
|
||||||
|
def resnet34(pretrained=False, **kwargs):
|
||||||
|
"""Constructs a ResNet-34 model.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||||
|
"""
|
||||||
|
model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
|
||||||
|
if pretrained:
|
||||||
|
model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
|
||||||
|
return model
|
||||||
|
|
||||||
|
|
||||||
|
def resnet50(pretrained=False, **kwargs):
|
||||||
|
"""Constructs a ResNet-50 model.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||||
|
"""
|
||||||
|
model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
|
||||||
|
if pretrained:
|
||||||
|
model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
|
||||||
|
return model
|
||||||
|
|
||||||
|
|
||||||
|
def resnet101(pretrained=False, **kwargs):
|
||||||
|
"""Constructs a ResNet-101 model.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||||
|
"""
|
||||||
|
model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
|
||||||
|
if pretrained:
|
||||||
|
model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
|
||||||
|
return model
|
||||||
|
|
||||||
|
|
||||||
|
def resnet152(pretrained=False, **kwargs):
|
||||||
|
"""Constructs a ResNet-152 model.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||||
|
"""
|
||||||
|
model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
|
||||||
|
if pretrained:
|
||||||
|
model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
|
||||||
|
return model
|
||||||
|
|
||||||
|
class Model1k(nn.Module):
|
||||||
|
def __init__(self, gpu_id=None):
|
||||||
|
super(Model1k, self).__init__()
|
||||||
|
|
||||||
|
self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||||
|
|
||||||
|
self.face_alignment_net = resnet18(pretrained=False, num_classes=1000 * 2, is_1k=True)
|
||||||
|
|
||||||
|
self.model_path, _ = os.path.split(os.path.realpath(__file__))
|
||||||
|
# weights = torch.load(os.path.join(self.model_path, 'face_alignment_1k.pth'), map_location=lambda storage, loc: storage)
|
||||||
|
weights = torch.load(os.path.join(modelRoot, 'face_alignment_1k.pth'), map_location=lambda storage, loc: storage)
|
||||||
|
self.load_state_dict(weights)
|
||||||
|
self.to(self.device)
|
||||||
|
self.eval()
|
||||||
|
|
||||||
|
def forward(self, imgs):
|
||||||
|
pred_key_pts = self.face_alignment_net(imgs)
|
||||||
|
pred_key_pts = pred_key_pts + 0.5
|
||||||
|
return pred_key_pts
|
||||||
|
|
||||||
|
class MomocvFaceAlignment1K(object):
|
||||||
|
def __init__(self, gpu_id=None):
|
||||||
|
self.gpu_id = gpu_id
|
||||||
|
|
||||||
|
self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||||
|
|
||||||
|
self.face_alignment_net = Model1k(gpu_id)
|
||||||
|
|
||||||
|
self.trackingFaceRects = []
|
||||||
|
|
||||||
|
print('conansherry MomocvFaceAlignment1K')
|
||||||
|
|
||||||
|
def forward(self, img_tensor):
|
||||||
|
fullyconnected1 = self.face_alignment_net(img_tensor).detach().cpu().numpy()
|
||||||
|
return fullyconnected1
|
||||||
|
|
||||||
|
def detect(self, img, landmarks):
|
||||||
|
dst_size = 256
|
||||||
|
landmarks_res = []
|
||||||
|
with torch.no_grad():
|
||||||
|
input_numpy = np.zeros((len(landmarks), 3, dst_size, dst_size), dtype=np.float32)
|
||||||
|
all_mat = []
|
||||||
|
for ix, landmark in enumerate(landmarks):
|
||||||
|
M = landmark_processor.get_transform_mat_full_face(landmark, dst_size)
|
||||||
|
all_mat.append(M)
|
||||||
|
tmp = cv2.warpAffine(img, M, (dst_size, dst_size))
|
||||||
|
|
||||||
|
# cv2.imshow('inp', tmp)
|
||||||
|
# cv2.waitKey()
|
||||||
|
|
||||||
|
input_numpy[ix, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
|
||||||
|
in_tensor = torch.from_numpy(input_numpy)
|
||||||
|
in_tensor = in_tensor.to(self.device)
|
||||||
|
fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy()
|
||||||
|
for ix, pts in enumerate(fullyconnected1):
|
||||||
|
orig_pts = (np.reshape(pts, (2, 1000)).transpose((1, 0)) * dst_size)
|
||||||
|
orig_pts = landmark_processor.transform_points(orig_pts, all_mat[ix], invert=True)
|
||||||
|
landmarks_res.append(orig_pts)
|
||||||
|
return landmarks_res
|
||||||
|
def detect_single_face(self, img):
|
||||||
|
dst_size = 256
|
||||||
|
with torch.no_grad():
|
||||||
|
crop_img = img[104:img.shape[0] - 104, 104:img.shape[1] - 104, :]
|
||||||
|
tmp = cv2.resize(crop_img, (dst_size, dst_size))
|
||||||
|
input_numpy = np.zeros((1, 3, dst_size, dst_size), dtype=np.float32)
|
||||||
|
input_numpy[0, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
|
||||||
|
in_tensor = torch.from_numpy(input_numpy)
|
||||||
|
in_tensor = in_tensor.to(self.device)
|
||||||
|
fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy()
|
||||||
|
orig_pts = (np.reshape(fullyconnected1[0], (2, 1000)).transpose((1, 0)) * crop_img.shape[0])
|
||||||
|
orig_pts[:, 0] += 104
|
||||||
|
orig_pts[:, 1] += 104
|
||||||
|
return orig_pts
|
||||||
|
|
||||||
|
def detect_single_face_old(self, img):
|
||||||
|
dst_size = 256
|
||||||
|
with torch.no_grad():
|
||||||
|
tmp = cv2.resize(img, (dst_size, dst_size))
|
||||||
|
input_numpy = np.zeros((1, 3, dst_size, dst_size), dtype=np.float32)
|
||||||
|
input_numpy[0, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
|
||||||
|
in_tensor = torch.from_numpy(input_numpy)
|
||||||
|
in_tensor = in_tensor.to(self.device)
|
||||||
|
fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy()
|
||||||
|
orig_pts = (np.reshape(fullyconnected1[0], (2, 1000)).transpose((1, 0)) * img.shape[0])
|
||||||
|
return orig_pts
|
||||||
|
def detect_according_5pts(self, img, pts5):
|
||||||
|
dst_size = 256
|
||||||
|
with torch.no_grad():
|
||||||
|
input_numpy = np.zeros((1, 3, dst_size, dst_size), dtype=np.float32)
|
||||||
|
eye_dis = 0.34
|
||||||
|
mouth_dis = 0.34
|
||||||
|
g_Average_5point_180 = np.array([
|
||||||
|
eye_dis, 0.3,
|
||||||
|
1 - eye_dis, 0.37,
|
||||||
|
0.5, 0.6,
|
||||||
|
mouth_dis, 0.63,
|
||||||
|
1 - mouth_dis, 0.7
|
||||||
|
])
|
||||||
|
# print(g_Average_5point_180)
|
||||||
|
# left_eye = np.array([pts5[0], pts5[5]])
|
||||||
|
# right_eye = np.array([pts5[1], pts5[6]])
|
||||||
|
# nose = np.array([pts5[2], pts5[7]])
|
||||||
|
# left_mouth = np.array([pts5[3], pts5[8]])
|
||||||
|
# right_mouth = np.array([pts5[4], pts5[9]])
|
||||||
|
left_eye = np.array([pts5[0], pts5[1]])
|
||||||
|
right_eye = np.array([pts5[2], pts5[3]])
|
||||||
|
nose = np.array([pts5[4], pts5[5]])
|
||||||
|
left_mouth = np.array([pts5[6], pts5[7]])
|
||||||
|
right_mouth = np.array([pts5[8], pts5[9]])
|
||||||
|
pts5_src = np.vstack((left_eye, right_eye,
|
||||||
|
nose,
|
||||||
|
left_mouth, right_mouth))
|
||||||
|
pts5_src = np.array(pts5_src).astype(np.int32)
|
||||||
|
pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size
|
||||||
|
# print('pts5_src: ', pts5_src)
|
||||||
|
# print('pts5_dst: ', pts5_dst)
|
||||||
|
|
||||||
|
# mat = cv2.estimateAffinePartial2D(pts5_src, pts5_dst, False)[0]
|
||||||
|
mat = umeyama(pts5_src, pts5_dst, True)[0:2]
|
||||||
|
print('mat: ', mat)
|
||||||
|
tmp = cv2.warpAffine(img, mat, (dst_size, dst_size))
|
||||||
|
|
||||||
|
# tmp2 = cv2.warpAffine(img, mat2, (dst_size, dst_size))
|
||||||
|
# cv2.imshow("tmp2", tmp2)
|
||||||
|
# cv2.imshow("tmp", tmp)
|
||||||
|
# cv2.waitKey()
|
||||||
|
input_numpy[0, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
|
||||||
|
in_tensor = torch.from_numpy(input_numpy)
|
||||||
|
in_tensor = in_tensor.to(self.device)
|
||||||
|
fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy()
|
||||||
|
orig_pts = (np.reshape(fullyconnected1[0], (2, 1000)).transpose((1, 0)) * dst_size)
|
||||||
|
orig_pts = landmark_processor.transform_points(orig_pts, mat, invert=True)
|
||||||
|
return orig_pts
|
||||||
|
|
||||||
|
def stable_forward(self, image, detected_faces, reset=False):
|
||||||
|
if reset is True:
|
||||||
|
self.trackingFaceRects = []
|
||||||
|
|
||||||
|
if len(self.trackingFaceRects) == 0:
|
||||||
|
for face_rect in detected_faces:
|
||||||
|
new_tracking_rect = [face_rect, True, [0, 0], 0, None]
|
||||||
|
self.trackingFaceRects.append(new_tracking_rect)
|
||||||
|
|
||||||
|
with torch.no_grad():
|
||||||
|
landmarks = []
|
||||||
|
for ix, tracking_face_rect in enumerate(self.trackingFaceRects):
|
||||||
|
if tracking_face_rect[1] == True:
|
||||||
|
d = tracking_face_rect[0]
|
||||||
|
src_center = np.array([d[2] - (d[2] - d[0]) / 2.0, d[3] - (d[3] - d[1]) / 2.0])
|
||||||
|
rotate_degree = tracking_face_rect[3]
|
||||||
|
scale = 256 * 0.6 / min(d[2] - d[0], d[3] - d[1])
|
||||||
|
dst_center = np.array([0.5, 0.5]) * 256
|
||||||
|
offset = dst_center - src_center
|
||||||
|
M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), rotate_degree, scale)
|
||||||
|
M[:, 2] += offset
|
||||||
|
else:
|
||||||
|
rotate_degree = 0
|
||||||
|
M = landmark_processor.get_transform_mat_mmcv_bigger(tracking_face_rect[4], 256)
|
||||||
|
inp = cv2.warpAffine(image, M, (256, 256))
|
||||||
|
|
||||||
|
# cv2.imshow('inp_{}'.format(ix), inp)
|
||||||
|
# cv2.waitKey()
|
||||||
|
orig_inp = inp
|
||||||
|
|
||||||
|
inp = inp.transpose((2, 0, 1)).astype(np.float32)
|
||||||
|
inp = inp[np.newaxis, :, :, :] / 255
|
||||||
|
|
||||||
|
in_tensor = torch.from_numpy(inp)
|
||||||
|
in_tensor = in_tensor.cuda(0)
|
||||||
|
fullyconnected1 = self.forward(in_tensor)
|
||||||
|
fullyconnected1 = fullyconnected1[0]
|
||||||
|
orig_pts = (np.reshape(fullyconnected1, (2, 1000)).transpose((1, 0))) * 256
|
||||||
|
|
||||||
|
t2 = cv2.getTickCount()
|
||||||
|
orig_pts = landmark_processor.transform_points(orig_pts, M, invert=True)
|
||||||
|
|
||||||
|
# orig_pts = orig_pts.transpose((1, 0))
|
||||||
|
fullyconnected1 = orig_pts
|
||||||
|
|
||||||
|
# update tracking infos
|
||||||
|
tracking_face_rect[1] = False
|
||||||
|
tracking_face_rect[2] = None
|
||||||
|
tracking_face_rect[3] = rotate_degree
|
||||||
|
tracking_face_rect[4] = fullyconnected1
|
||||||
|
|
||||||
|
# fullyconnected1 = landmark_processor.pts_1k_to_137(fullyconnected1)
|
||||||
|
|
||||||
|
# eye_landmark = self.detect_eye(image, fullyconnected1)
|
||||||
|
# fullyconnected1[87:104] = eye_landmark[0]
|
||||||
|
# fullyconnected1[104:121] = eye_landmark[1]
|
||||||
|
|
||||||
|
landmarks.append(fullyconnected1)
|
||||||
|
return landmarks
|
||||||
@@ -0,0 +1,133 @@
|
|||||||
|
import numpy as np
|
||||||
|
import cv2
|
||||||
|
|
||||||
|
def nms(boxes, overlap_threshold=0.5, mode='union'):
|
||||||
|
""" Pure Python NMS baseline. """
|
||||||
|
x1 = boxes[:, 0]
|
||||||
|
y1 = boxes[:, 1]
|
||||||
|
x2 = boxes[:, 2]
|
||||||
|
y2 = boxes[:, 3]
|
||||||
|
scores = boxes[:, 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
|
||||||
|
|
||||||
|
if mode is 'min':
|
||||||
|
ovr = inter / np.minimum(areas[i], areas[order[1:]])
|
||||||
|
else:
|
||||||
|
ovr = inter / (areas[i] + areas[order[1:]] - inter)
|
||||||
|
|
||||||
|
inds = np.where(ovr <= overlap_threshold)[0]
|
||||||
|
order = order[inds + 1]
|
||||||
|
|
||||||
|
return keep
|
||||||
|
|
||||||
|
|
||||||
|
def convert_to_square(bboxes):
|
||||||
|
"""
|
||||||
|
Convert bounding boxes to a square form.
|
||||||
|
"""
|
||||||
|
square_bboxes = np.zeros_like(bboxes)
|
||||||
|
x1, y1, x2, y2 = [bboxes[:, i] for i in range(4)]
|
||||||
|
h = y2 - y1 + 1.0
|
||||||
|
w = x2 - x1 + 1.0
|
||||||
|
max_side = np.maximum(h, w)
|
||||||
|
square_bboxes[:, 0] = x1 + w*0.5 - max_side*0.5
|
||||||
|
square_bboxes[:, 1] = y1 + h*0.5 - max_side*0.5
|
||||||
|
square_bboxes[:, 2] = square_bboxes[:, 0] + max_side - 1.0
|
||||||
|
square_bboxes[:, 3] = square_bboxes[:, 1] + max_side - 1.0
|
||||||
|
return square_bboxes
|
||||||
|
|
||||||
|
|
||||||
|
def calibrate_box(bboxes, offsets):
|
||||||
|
"""Transform bounding boxes to be more like true bounding boxes.
|
||||||
|
'offsets' is one of the outputs of the nets.
|
||||||
|
"""
|
||||||
|
x1, y1, x2, y2 = [bboxes[:, i] for i in range(4)]
|
||||||
|
w = x2 - x1 + 1.0
|
||||||
|
h = y2 - y1 + 1.0
|
||||||
|
w = np.expand_dims(w, 1)
|
||||||
|
h = np.expand_dims(h, 1)
|
||||||
|
|
||||||
|
translation = np.hstack([w, h, w, h])*offsets
|
||||||
|
bboxes[:, 0:4] = bboxes[:, 0:4] + translation
|
||||||
|
return bboxes
|
||||||
|
|
||||||
|
|
||||||
|
def get_image_boxes(bounding_boxes, img, size=24):
|
||||||
|
"""Cut out boxes from the image.
|
||||||
|
"""
|
||||||
|
num_boxes = len(bounding_boxes)
|
||||||
|
(height, width, _) = img.shape
|
||||||
|
|
||||||
|
[dy, edy, dx, edx, y, ey, x, ex, w, h] = correct_bboxes(bounding_boxes, width, height)
|
||||||
|
img_boxes = np.zeros((num_boxes, 3, size, size), 'float32')
|
||||||
|
|
||||||
|
for i in range(num_boxes):
|
||||||
|
img_box = np.zeros((h[i], w[i], 3), 'uint8')
|
||||||
|
|
||||||
|
img_array = np.asarray(img, 'uint8')
|
||||||
|
img_box[dy[i]:(edy[i] + 1), dx[i]:(edx[i] + 1), :] =\
|
||||||
|
img_array[y[i]:(ey[i] + 1), x[i]:(ex[i] + 1), :]
|
||||||
|
|
||||||
|
img_box = cv2.resize(img_box, (size, size))
|
||||||
|
img_box = np.asarray(img_box, 'float32')
|
||||||
|
|
||||||
|
img_boxes[i, :, :, :] = _preprocess(img_box)
|
||||||
|
|
||||||
|
return img_boxes
|
||||||
|
|
||||||
|
|
||||||
|
def correct_bboxes(bboxes, width, height):
|
||||||
|
"""Crop boxes that are too big and get coordinates
|
||||||
|
with respect to cutouts.
|
||||||
|
"""
|
||||||
|
x1, y1, x2, y2 = [bboxes[:, i] for i in range(4)]
|
||||||
|
w, h = x2 - x1 + 1.0, y2 - y1 + 1.0
|
||||||
|
num_boxes = bboxes.shape[0]
|
||||||
|
|
||||||
|
x, y, ex, ey = x1, y1, x2, y2
|
||||||
|
dx, dy = np.zeros((num_boxes,)), np.zeros((num_boxes,))
|
||||||
|
edx, edy = w.copy() - 1.0, h.copy() - 1.0
|
||||||
|
|
||||||
|
ind = np.where(ex > width - 1.0)[0]
|
||||||
|
edx[ind] = w[ind] + width - 2.0 - ex[ind]
|
||||||
|
ex[ind] = width - 1.0
|
||||||
|
|
||||||
|
ind = np.where(ey > height - 1.0)[0]
|
||||||
|
edy[ind] = h[ind] + height - 2.0 - ey[ind]
|
||||||
|
ey[ind] = height - 1.0
|
||||||
|
|
||||||
|
ind = np.where(x < 0.0)[0]
|
||||||
|
dx[ind] = 0.0 - x[ind]
|
||||||
|
x[ind] = 0.0
|
||||||
|
|
||||||
|
ind = np.where(y < 0.0)[0]
|
||||||
|
dy[ind] = 0.0 - y[ind]
|
||||||
|
y[ind] = 0.0
|
||||||
|
return_list = [dy, edy, dx, edx, y, ey, x, ex, w, h]
|
||||||
|
return_list = [i.astype('int32') for i in return_list]
|
||||||
|
|
||||||
|
return return_list
|
||||||
|
|
||||||
|
|
||||||
|
def _preprocess(img):
|
||||||
|
"""Preprocessing step before feeding the network.
|
||||||
|
"""
|
||||||
|
img = img.transpose((2, 0, 1))
|
||||||
|
img = np.expand_dims(img, 0)
|
||||||
|
img = (img - 127.5)*0.0078125
|
||||||
|
return img
|
||||||
@@ -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
|
||||||
|
}
|
||||||
|
|
||||||
@@ -0,0 +1,259 @@
|
|||||||
|
import math
|
||||||
|
import numpy as np
|
||||||
|
from .box_utils import nms, calibrate_box, get_image_boxes, convert_to_square, _preprocess
|
||||||
|
import torch
|
||||||
|
import cv2
|
||||||
|
from .nms.py_cpu_nms import py_cpu_nms
|
||||||
|
from core.utils import box_utils_Retina
|
||||||
|
from .layers.functions.prior_box import PriorBox
|
||||||
|
from .config import cfg_re50
|
||||||
|
from .retinaface import RetinaFace
|
||||||
|
from mtcnn.model import PNet, RNet, ONet
|
||||||
|
|
||||||
|
gpu_id = 0
|
||||||
|
device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||||
|
pnet, rnet, onet= PNet(), RNet(), ONet()
|
||||||
|
pnet.to(device)
|
||||||
|
rnet.to(device)
|
||||||
|
onet.to(device)
|
||||||
|
onet.eval()
|
||||||
|
|
||||||
|
|
||||||
|
class RetinaFaceDetector(object):
|
||||||
|
def __init__(self, gpu_id=None):
|
||||||
|
self.gpu_id = gpu_id
|
||||||
|
self.cfg = cfg_re50
|
||||||
|
self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||||
|
model = RetinaFace(cfg=self.cfg, phase='test')
|
||||||
|
model = self.load_model(model, 'weights/Resnet50_Final.pth', True)
|
||||||
|
model.eval()
|
||||||
|
self.net = model.to(self.device)
|
||||||
|
self.resize = 1
|
||||||
|
self.confidence_threshold = 0.02
|
||||||
|
self.top_k = 5000
|
||||||
|
self.nms_threshold = 0.4
|
||||||
|
self.keep_top_k = 750
|
||||||
|
|
||||||
|
def remove_prefix(self, state_dict, prefix):
|
||||||
|
# print('remove prefix \'{}\''.format(prefix))
|
||||||
|
f = lambda x: x.split(prefix, 1)[-1] if x.startswith(prefix) else x
|
||||||
|
return {f(key): value for key, value in state_dict.items()}
|
||||||
|
|
||||||
|
def check_keys(self, model, pretrained_state_dict):
|
||||||
|
ckpt_keys = set(pretrained_state_dict.keys())
|
||||||
|
model_keys = set(model.state_dict().keys())
|
||||||
|
used_pretrained_keys = model_keys & ckpt_keys
|
||||||
|
# unused_pretrained_keys = ckpt_keys - model_keys
|
||||||
|
# missing_keys = model_keys - ckpt_keys
|
||||||
|
assert len(used_pretrained_keys) > 0, 'load NONE from pretrained checkpoint'
|
||||||
|
return True
|
||||||
|
|
||||||
|
def load_model(self, model, pretrained_path, load_to_cpu):
|
||||||
|
# print('Loading pretrained model from {}'.format(pretrained_path))
|
||||||
|
if load_to_cpu:
|
||||||
|
pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage)
|
||||||
|
else:
|
||||||
|
device = torch.cuda.current_device()
|
||||||
|
pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage.cuda(device))
|
||||||
|
if "state_dict" in pretrained_dict.keys():
|
||||||
|
pretrained_dict = self.remove_prefix(pretrained_dict['state_dict'], 'module.')
|
||||||
|
else:
|
||||||
|
pretrained_dict = self.remove_prefix(pretrained_dict, 'module.')
|
||||||
|
self.check_keys(model, pretrained_dict)
|
||||||
|
model.load_state_dict(pretrained_dict, strict=False)
|
||||||
|
return model
|
||||||
|
|
||||||
|
def forward(self, img_raw, min_face_size=50):
|
||||||
|
img_scale = 640 / max(img_raw.shape[0], img_raw.shape[1])
|
||||||
|
img = cv2.resize(img_raw, (0, 0), fx=img_scale, fy=img_scale)
|
||||||
|
img = np.float32(img)
|
||||||
|
|
||||||
|
im_height, im_width, _ = img.shape
|
||||||
|
scale = torch.Tensor([img.shape[1], img.shape[0], img.shape[1], img.shape[0]])
|
||||||
|
img -= (104, 117, 123)
|
||||||
|
img = img.transpose(2, 0, 1)
|
||||||
|
img = torch.from_numpy(img).unsqueeze(0)
|
||||||
|
img = img.to(self.device)
|
||||||
|
scale = scale.to(self.device)
|
||||||
|
|
||||||
|
loc, conf, landms = self.net(img) # forward pass
|
||||||
|
|
||||||
|
priorbox = PriorBox(self.cfg, image_size=(im_height, im_width))
|
||||||
|
priors = priorbox.forward()
|
||||||
|
priors = priors.to(self.device)
|
||||||
|
prior_data = priors.data
|
||||||
|
boxes = box_utils_Retina.decode(loc.data.squeeze(0), prior_data, self.cfg['variance'])
|
||||||
|
|
||||||
|
boxes = boxes * scale / self.resize
|
||||||
|
boxes = boxes.cpu().numpy()
|
||||||
|
scores = conf.squeeze(0).data.cpu().numpy()[:, 1]
|
||||||
|
landms = box_utils_Retina.decode_landm(landms.data.squeeze(0), prior_data, self.cfg['variance'])
|
||||||
|
scale1 = torch.Tensor([img.shape[3], img.shape[2], img.shape[3], img.shape[2],
|
||||||
|
img.shape[3], img.shape[2], img.shape[3], img.shape[2],
|
||||||
|
img.shape[3], img.shape[2]])
|
||||||
|
scale1 = scale1.to(self.device)
|
||||||
|
landms = landms * scale1 / self.resize
|
||||||
|
landms = landms.cpu().numpy()
|
||||||
|
|
||||||
|
# ignore low scores
|
||||||
|
inds = np.where(scores > self.confidence_threshold)[0]
|
||||||
|
boxes = boxes[inds]
|
||||||
|
landms = landms[inds]
|
||||||
|
scores = scores[inds]
|
||||||
|
|
||||||
|
# keep top-K before NMS
|
||||||
|
order = scores.argsort()[::-1][:self.top_k]
|
||||||
|
boxes = boxes[order]
|
||||||
|
landms = landms[order]
|
||||||
|
scores = scores[order]
|
||||||
|
|
||||||
|
# do NMS
|
||||||
|
dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
|
||||||
|
keep = py_cpu_nms(dets, self.nms_threshold, min_face_size = min_face_size * img_scale)
|
||||||
|
# keep = nms(dets, args.nms_threshold,force_cpu=args.cpu)
|
||||||
|
dets = dets[keep, :]
|
||||||
|
landms = landms[keep]
|
||||||
|
|
||||||
|
# keep top-K faster NMS
|
||||||
|
dets = dets[:self.keep_top_k, :]
|
||||||
|
landms = landms[:self.keep_top_k, :]
|
||||||
|
|
||||||
|
dets[:, :4] = dets[:, :4] / img_scale
|
||||||
|
landms /= img_scale
|
||||||
|
return dets, landms
|
||||||
|
|
||||||
|
def forward_v2(self, image, min_face_size=20.0, thresholds=[0.6, 0.7, 0.8], nms_thresholds=[0.7, 0.7, 0.7]):
|
||||||
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||||
|
(height, width, _) = image.shape
|
||||||
|
min_length = min(height, width)
|
||||||
|
min_detection_size = 12
|
||||||
|
factor = 0.707 # sqrt(0.5)
|
||||||
|
|
||||||
|
scales = []
|
||||||
|
m = min_detection_size / min_face_size
|
||||||
|
min_length *= m
|
||||||
|
|
||||||
|
factor_count = 0
|
||||||
|
while min_length > min_detection_size:
|
||||||
|
scales.append(m * factor ** factor_count)
|
||||||
|
min_length *= factor
|
||||||
|
factor_count += 1
|
||||||
|
|
||||||
|
# STAGE 1
|
||||||
|
bounding_boxes = []
|
||||||
|
for s in scales: # run P-Net on different scales
|
||||||
|
boxes = run_first_stage(image, pnet, scale=s, threshold=thresholds[0], gpu_id=self.gpu_id)
|
||||||
|
bounding_boxes.append(boxes)
|
||||||
|
bounding_boxes = [i for i in bounding_boxes if i is not None]
|
||||||
|
if len(bounding_boxes) == 0:
|
||||||
|
return [], []
|
||||||
|
bounding_boxes = np.vstack(bounding_boxes)
|
||||||
|
|
||||||
|
keep = nms(bounding_boxes[:, 0:5], nms_thresholds[0])
|
||||||
|
bounding_boxes = bounding_boxes[keep]
|
||||||
|
bounding_boxes = calibrate_box(bounding_boxes[:, 0:5], bounding_boxes[:, 5:])
|
||||||
|
bounding_boxes = convert_to_square(bounding_boxes)
|
||||||
|
bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
|
||||||
|
|
||||||
|
# STAGE 2
|
||||||
|
img_boxes = get_image_boxes(bounding_boxes, image, size=24)
|
||||||
|
img_boxes = torch.from_numpy(img_boxes)
|
||||||
|
img_boxes = img_boxes.to(self.device)
|
||||||
|
output = rnet(img_boxes)
|
||||||
|
offsets = output[0].to('cpu').data.numpy() # shape [n_boxes, 4]
|
||||||
|
probs = output[1].to('cpu').data.numpy() # shape [n_boxes, 2]
|
||||||
|
|
||||||
|
keep = np.where(probs[:, 1] > thresholds[1])[0]
|
||||||
|
bounding_boxes = bounding_boxes[keep]
|
||||||
|
bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
|
||||||
|
offsets = offsets[keep]
|
||||||
|
|
||||||
|
keep = nms(bounding_boxes, nms_thresholds[1])
|
||||||
|
bounding_boxes = bounding_boxes[keep]
|
||||||
|
bounding_boxes = calibrate_box(bounding_boxes, offsets[keep])
|
||||||
|
bounding_boxes = convert_to_square(bounding_boxes)
|
||||||
|
bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
|
||||||
|
|
||||||
|
# STAGE 3
|
||||||
|
img_boxes = get_image_boxes(bounding_boxes, image, size=48)
|
||||||
|
if len(img_boxes) == 0:
|
||||||
|
return [], []
|
||||||
|
img_boxes = torch.from_numpy(img_boxes)
|
||||||
|
img_boxes = img_boxes.to(self.device)
|
||||||
|
output = onet(img_boxes)
|
||||||
|
landmarks = output[0].to('cpu').data.numpy() # shape [n_boxes, 10]
|
||||||
|
offsets = output[1].to('cpu').data.numpy() # shape [n_boxes, 4]
|
||||||
|
probs = output[2].to('cpu').data.numpy() # shape [n_boxes, 2]
|
||||||
|
|
||||||
|
keep = np.where(probs[:, 1] > thresholds[2])[0]
|
||||||
|
bounding_boxes = bounding_boxes[keep]
|
||||||
|
bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
|
||||||
|
offsets = offsets[keep]
|
||||||
|
landmarks = landmarks[keep]
|
||||||
|
|
||||||
|
# compute landmark points
|
||||||
|
width = bounding_boxes[:, 2] - bounding_boxes[:, 0] + 1.0
|
||||||
|
height = bounding_boxes[:, 3] - bounding_boxes[:, 1] + 1.0
|
||||||
|
xmin, ymin = bounding_boxes[:, 0], bounding_boxes[:, 1]
|
||||||
|
landmarks[:, 0:5] = np.expand_dims(xmin, 1) + np.expand_dims(width, 1) * landmarks[:, 0:5]
|
||||||
|
landmarks[:, 5:10] = np.expand_dims(ymin, 1) + np.expand_dims(height, 1) * landmarks[:, 5:10]
|
||||||
|
|
||||||
|
bounding_boxes = calibrate_box(bounding_boxes, offsets)
|
||||||
|
keep = nms(bounding_boxes, nms_thresholds[2], mode='min')
|
||||||
|
bounding_boxes = bounding_boxes[keep]
|
||||||
|
landmarks = landmarks[keep]
|
||||||
|
|
||||||
|
return bounding_boxes, landmarks
|
||||||
|
|
||||||
|
def run_first_stage(image, net, scale, threshold, gpu_id=0):
|
||||||
|
"""
|
||||||
|
Run P-Net, generate bounding boxes, and do NMS.
|
||||||
|
"""
|
||||||
|
device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||||
|
(height, width, _) = image.shape
|
||||||
|
sw, sh = math.ceil(width * scale), math.ceil(height * scale)
|
||||||
|
img = cv2.resize(image, (sw, sh))
|
||||||
|
# img = image.resize((sw, sh), Image.BILINEAR)
|
||||||
|
img = np.asarray(img, 'float32')
|
||||||
|
img = torch.from_numpy(_preprocess(img))
|
||||||
|
img = img.to(device)
|
||||||
|
|
||||||
|
output = net(img)
|
||||||
|
probs = output[1].to('cpu').data.numpy()[0, 1, :, :]
|
||||||
|
offsets = output[0].to('cpu').data.numpy()
|
||||||
|
|
||||||
|
boxes = _generate_bboxes(probs, offsets, scale, threshold)
|
||||||
|
if len(boxes) == 0:
|
||||||
|
return None
|
||||||
|
|
||||||
|
keep = nms(boxes[:, 0:5], overlap_threshold=0.5)
|
||||||
|
return boxes[keep]
|
||||||
|
|
||||||
|
|
||||||
|
def _generate_bboxes(probs, offsets, scale, threshold):
|
||||||
|
"""
|
||||||
|
Generate bounding boxes at places where there is probably a face.
|
||||||
|
"""
|
||||||
|
stride = 2
|
||||||
|
cell_size = 12
|
||||||
|
|
||||||
|
inds = np.where(probs > threshold)
|
||||||
|
|
||||||
|
if inds[0].size == 0:
|
||||||
|
return np.array([])
|
||||||
|
|
||||||
|
tx1, ty1, tx2, ty2 = [offsets[0, i, inds[0], inds[1]] for i in range(4)]
|
||||||
|
|
||||||
|
offsets = np.array([tx1, ty1, tx2, ty2])
|
||||||
|
score = probs[inds[0], inds[1]]
|
||||||
|
|
||||||
|
# P-Net is applied to scaled images, so we need to rescale bounding boxes back
|
||||||
|
bounding_boxes = np.vstack([
|
||||||
|
np.round((stride * inds[1] + 1.0) / scale),
|
||||||
|
np.round((stride * inds[0] + 1.0) / scale),
|
||||||
|
np.round((stride * inds[1] + 1.0 + cell_size) / scale),
|
||||||
|
np.round((stride * inds[0] + 1.0 + cell_size) / scale),
|
||||||
|
score, offsets
|
||||||
|
])
|
||||||
|
|
||||||
|
return bounding_boxes.T
|
||||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user