Compare commits
5
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 xsl
|
9b7324cd26 | ||
|
xsl
|
b2738fba7d | ||
|
xsl
|
2606aca7c5 | ||
|
xsl
|
ce445b64a6 | ||
|
xsl
|
aaa1bf159f |
+408
@@ -0,0 +1,408 @@
|
||||
# 换发型功能集成文档
|
||||
|
||||
本文档说明"换发型"功能的完整工作流程、服务依赖、资源结构，供其他项目集成参考。
|
||||
|
||||
---
|
||||
|
||||
## 一、功能概述
|
||||
|
||||
**换发型**：用户上传一张人像照片，选择一个目标发型（hair_id），系统把用户的头发替换成目标发型，保留人脸和身体。
|
||||
|
||||
- **输入**：用户人像图（base64 或 URL） + 目标发型 ID
|
||||
- **输出**：换好发型的结果图（base64 或 OSS URL）
|
||||
- **耗时**：单次约 10-12 秒（GPU: RTX 5090）
|
||||
|
||||
### 同类功能（同套服务支撑）
|
||||
- **换发色**（`/hairColor/v2`）：把头发颜色换成指定 RGB，额外依赖 `ref_color` 颜色查找表
|
||||
- **训练新发型**：用发型图片训练新的 LoRA，详见第七节
|
||||
|
||||
---
|
||||
|
||||
## 二、服务架构
|
||||
|
||||
换发型由 **3 个后端服务**协同完成（缺一不可），另有 1 个前端页面服务（可选）。
|
||||
|
||||
| 服务 | 端口 | conda 环境 | 职责 |
|
||||
|------|------|-----------|------|
|
||||
| **webui** | 57860 | `sdwebui` | Stable Diffusion img2img 推理引擎（基于 AUTOMATIC1111 webui） |
|
||||
| **photo_service** | 32678 | `py310` | LoRA 调度：推理时把对应发型的 LoRA 复制到 webui 并注入 prompt；训练时调度 kohya |
|
||||
| **hair_service_sd** | 8801 | `my_hair` | 换发型主服务：人脸检测、关键点、头发抠图、材质推理、图像合成 |
|
||||
| hair_grow_service | 8899 | `my_hair` | （可选）前端测试页面，提供发型列表/预览/换发型 UI |
|
||||
|
||||
### 服务间调用链（换发型）
|
||||
|
||||
```
|
||||
调用方
|
||||
│ POST /api/swapHair/v1
|
||||
▼
|
||||
hair_service_sd (8801) ← 人脸检测/关键点/抠图/材质推理/图像合成
|
||||
│ POST /api/hair/inference (转发 img2img 请求)
|
||||
▼
|
||||
photo_service (32678) ← 复制 LoRA 到 webui、prompt 注入 <lora:xxx:1.0>
|
||||
│ POST /sdapi/v1/img2img
|
||||
▼
|
||||
webui (57860) ← SD img2img 真正推理
|
||||
```
|
||||
|
||||
### 启动方式
|
||||
|
||||
```bash
|
||||
cd /home/xsl/change_hair
|
||||
bash start_all.sh start # 启动 webui + photo_service + hair_service_sd
|
||||
# 等待 1-2 分钟（webui 加载 SD 模型）
|
||||
|
||||
# 可选：前端页面
|
||||
nohup bash start_hairgrow.sh > project/logs/hair_grow_service.log 2>&1 &
|
||||
```
|
||||
|
||||
> 三个核心服务有启动顺序依赖：webui 必须先启动并加载完模型（约 1-2 分钟），photo_service 和 hair_service_sd 才能正常工作。
|
||||
|
||||
---
|
||||
|
||||
## 三、换发型 API 接口
|
||||
|
||||
### 请求
|
||||
|
||||
```
|
||||
POST http://<host>:8801/api/swapHair/v1
|
||||
Content-Type: application/json
|
||||
```
|
||||
|
||||
**请求体参数：**
|
||||
|
||||
| 参数 | 类型 | 必填 | 说明 |
|
||||
|------|------|:---:|------|
|
||||
| `hair_id` | string | ✓ | 目标发型 ID（如 `chang_tuoyuan`），对应发型资源目录名 |
|
||||
| `task_id` | string | ✓ | 任务唯一标识，用于临时文件命名 |
|
||||
| `user_img_path` | string | ✓ | 用户人像图。支持 `data:image/jpeg;base64,<base64>` 或图片 URL |
|
||||
| `is_hr` | string | ✓ | 是否高清模式。`"true"` 输出 1152×1536，`"false"` 输出 576×768 |
|
||||
| `output_format` | string | ✗ | 输出格式。`"base64"`（推荐，直接返回图）或 `"url"`（上传 OSS 返回链接，需配 OSS 凭证） |
|
||||
|
||||
**请求示例：**
|
||||
|
||||
```bash
|
||||
curl -X POST http://127.0.0.1:8801/api/swapHair/v1 \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"hair_id": "chang_tuoyuan",
|
||||
"task_id": "req_001",
|
||||
"is_hr": "true",
|
||||
"user_img_path": "data:image/jpeg;base64,/9j/4AAQ...",
|
||||
"output_format": "base64"
|
||||
}'
|
||||
```
|
||||
|
||||
```python
|
||||
# Python 示例
|
||||
import base64, requests
|
||||
with open("user.jpg", "rb") as f:
|
||||
img_b64 = base64.b64encode(f.read()).decode()
|
||||
r = requests.post("http://127.0.0.1:8801/api/swapHair/v1", json={
|
||||
"hair_id": "chang_tuoyuan",
|
||||
"task_id": "req_001",
|
||||
"is_hr": "true",
|
||||
"user_img_path": "data:image/jpeg;base64," + img_b64,
|
||||
"output_format": "base64",
|
||||
}, timeout=300)
|
||||
result = r.json() # {"state":0, "data":"<base64>", "msg":"success"}
|
||||
with open("out.jpg", "wb") as f:
|
||||
f.write(base64.b64decode(result["data"]))
|
||||
```
|
||||
|
||||
### 响应
|
||||
|
||||
**成功（HTTP 200）：**
|
||||
```json
|
||||
{
|
||||
"state": 0,
|
||||
"msg": "success",
|
||||
"data": "<base64 编码的结果图>" // output_format=url 时为 OSS 图片链接
|
||||
}
|
||||
```
|
||||
|
||||
**失败（HTTP 400）：**
|
||||
```json
|
||||
{
|
||||
"state": -1,
|
||||
"msg": "推理发型失败", // 具体错误信息，如：参数错误/用户图像加载失败/发型模板加载失败/推理发型失败
|
||||
"data": "",
|
||||
"task_id": "req_001"
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 四、换发型内部流程（9 个功能步骤）
|
||||
|
||||
理解内部流程有助于排查问题和性能优化。`hair_service_sd` 收到请求后依次执行：
|
||||
|
||||
| 步骤 | 功能 | 说明 | 耗时占比 |
|
||||
|------|------|------|---------|
|
||||
| 1 | 获取请求参数 | 解析参数，base64/URL 图片解码落盘 | <1% |
|
||||
| 2 | 加载用户图像 | 复制到工作目录 | <1% |
|
||||
| 3 | 加载发型模板图像 | 从 `upload_train_imgs/<hair_id>/` 读模板原图 | <1% |
|
||||
| 4 | 检查遮挡眼睛 | 从 `hair_template_material/<hair_id>/` 读 `*_matting.png` 和 `.pkl`（关键点） | <1% |
|
||||
| 5 | 检查发型材质 | 确认 `ref_hairstyle/<hair_id>/` 存在 | <1% |
|
||||
| 6 | **推理发型** | 用 GAN 模型把用户图和 `ref_rgb_8uc3_768.png` 融合，生成初步换发图 | ~30% |
|
||||
| 7 | 判断遮挡眼睛 | 检查头发是否遮挡眼睛（决定后续处理路径） | <1% |
|
||||
| 8 | 处理发型区域 | 头发区域裁剪、mask 生成、透视变换到 576×768/1152×1536 | ~5% |
|
||||
| 9 | **webui img2img** | 经 photo_service 调 webui，加载 LoRA 做 SD 重绘（核心画质步骤） | ~60% |
|
||||
|
||||
**关键：LoRA 加载机制**（步骤 9）
|
||||
1. `hair_service_sd` 把裁剪好的头发区域图 + mask 发给 `photo_service`
|
||||
2. `photo_service` 把 `train_material/<hair_id>/model/hairstyle_hd_lora.safetensors` 复制到 webui 的 `models/Lora/<hair_id>_hd.safetensors`（已存在则跳过）
|
||||
3. `photo_service` 在 prompt 前注入 `<lora:<hair_id>_hd:1.0>`
|
||||
4. webui 收到 img2img 请求，自动加载该 LoRA 做局部重绘
|
||||
5. 结果返回 `hair_service_sd`，贴回原图对应位置
|
||||
|
||||
---
|
||||
|
||||
## 五、必需资源清单
|
||||
|
||||
换发型功能依赖以下资源。集成时需保证这些文件到位。
|
||||
|
||||
### 5.1 核心模型权重（全局共享，约 10G）
|
||||
|
||||
位置：`project/hair_service_sd/weights/`
|
||||
|
||||
| 文件/目录 | 用途 |
|
||||
|-----------|------|
|
||||
| `yolov5l.pt` | RetinaFace/人脸检测 |
|
||||
| `Resnet50_Final.pth` | 人脸检测辅助 |
|
||||
| `79999_iter.pth` | 关键点检测（MomocvFaceAlignment1K） |
|
||||
| `master_hair_v8_onlyhair_nowarp_all_0709.pt` | 头发迁移 GAN（发型融合） |
|
||||
| `zxm_v7_gen_bald_add_changed_hair_pairdata_5seg_0611.pt` | 秃头分割 + 发型生成 GAN |
|
||||
| `parser_model.pth` | 头发抠图（matting） |
|
||||
| `gender_models/` | 性别识别 |
|
||||
| `keypoints/` | 关键点模型库 |
|
||||
|
||||
### 5.2 Stable Diffusion 推理资源（约 5G）
|
||||
|
||||
位置：`project/onediff/stable-diffusion-webui/models/`
|
||||
|
||||
| 文件/目录 | 用途 |
|
||||
|-----------|------|
|
||||
| `Stable-diffusion/v1-5-pruned-emaonly.safetensors` | **SD 1.5 底模**（必需，4G） |
|
||||
| `Lora/<hair_id>_hd.safetensors` | 各发型的 LoRA（推理时自动生成，见 5.3） |
|
||||
| `ESRGAN/4x-UltraSharp.pth` | 超分辨率（高清模式用） |
|
||||
| `GFPGAN/` | 人脸修复（可选） |
|
||||
|
||||
### 5.3 单个发型的资源（每个发型约 160M）
|
||||
|
||||
**一个可用的发型需要 4 个目录的资源配套：**
|
||||
|
||||
```
|
||||
# 1. 发型材质（换发型必需，由训练 step3 生成）
|
||||
project/data/ref_hairstyle/<hair_id>/
|
||||
├── config.json # {"gender":"girl","version":"20221206","ratio":"1"}
|
||||
├── ref_rgb_8uc3_768.png # 发型 RGB 图（GAN 融合参考图，必需）
|
||||
├── ref_matting_8uc3_768.png # 发型 mask
|
||||
├── ref_matting_fg_8uc3_768.png # 发型前景
|
||||
├── ref_baldseg_8uc3_768.png # 秃头分割
|
||||
├── ref_landmark_f1k2_768.txt # 1000 个关键点
|
||||
└── input_another_pose_hair_image.npy
|
||||
|
||||
# 2. LoRA 权重 + 训练图（换发型必需，由训练 step1+step2 生成）
|
||||
data/train_material/<hair_id>/
|
||||
├── model/
|
||||
│ └── hairstyle_hd_lora.safetensors # LoRA 权重（~145M，SD 重绘用）
|
||||
└── images/
|
||||
└── 1_hairstyle/ # 训练图 + caption
|
||||
├── 0001.png
|
||||
├── 0001.txt # caption tag（webui 推理 prompt 用）
|
||||
└── ... (约20张)
|
||||
|
||||
# 3. 模板材质（换发型必需，由训练 step4 生成）
|
||||
project/data/hair_template_material/<hair_id>/
|
||||
├── first##<uuid>.png # 发型模板原图
|
||||
├── first##<uuid>_matting.png # 发型 matting（遮挡检测用）
|
||||
└── first##<uuid>.pkl # 关键点（遮挡检测用）
|
||||
|
||||
# 4. 发型模板原图（换发型必需，由训练准备）
|
||||
project/data/upload_train_imgs/<hair_id>/
|
||||
└── first##<hair_id>.jpg # 发型模板原图（步骤3读取）
|
||||
```
|
||||
|
||||
> **注意**：`images/1_hairstyle/` 目录名必须以数字开头（`1_hairstyle`），代码用 `os.listdir(images)[0]` 取子目录并校验数字开头。
|
||||
|
||||
### 5.4 路径配置
|
||||
|
||||
所有路径在 `project/hair_service_sd/config/configure.ini` 集中配置，迁移时改这里的路径即可：
|
||||
|
||||
```ini
|
||||
[default]
|
||||
modelDir = weights # 核心模型目录（相对 hair_service_sd）
|
||||
hairstyleDir = /home/xsl/change_hair/project/data/ref_hairstyle # 发型材质（5.3-1）
|
||||
haircolorDir= /home/xsl/change_hair/project/data/ref_haircolor # 换发色用
|
||||
userDir=/home/xsl/change_hair/project/data/userImage # 用户图工作目录
|
||||
tmp_dir=/home/xsl/change_hair/project/data/tmp # 临时目录
|
||||
res_dir=/home/xsl/change_hair/project/data/res_dir # 结果输出目录
|
||||
userInfo_dir=/home/xsl/change_hair/project/data/user_info # 用户关键点缓存
|
||||
train_dir = /home/xsl/change_hair/data/train_material # LoRA+训练图（5.3-2）
|
||||
refImgDir=/home/xsl/change_hair/project/data/refImage # 参考图缓存
|
||||
hair_template_material_dir=/home/xsl/change_hair/project/data/hair_template_material # 5.3-3
|
||||
upload_train_dir=/home/xsl/change_hair/project/data/upload_train_imgs # 5.3-4
|
||||
version=online # online 用 0.0.0.0 地址
|
||||
```
|
||||
|
||||
`photo_service` 的 webui/LoRA 路径在 `project/photo_service/lora_train_service_1.py` 第 41 行：
|
||||
```python
|
||||
webui_lora_dir = '/home/xsl/change_hair/project/onediff/stable-diffusion-webui/models/Lora'
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 六、集成步骤（接入新项目）
|
||||
|
||||
### 1. 部署服务端
|
||||
按第二节启动 webui + photo_service + hair_service_sd 三个服务，确认 `http://<host>:8801` 可达。
|
||||
|
||||
### 2. 调用换发型接口
|
||||
新项目只需调 `hair_service_sd` 的 `/api/swapHair/v1`，无需关心内部服务调用。
|
||||
|
||||
### 3. 准备发型资源
|
||||
每个要支持的发型，必须备齐第五节 5.3 的 4 个目录资源。两种获取方式：
|
||||
- **训练新发型**（见第七节）：用发型图训练，自动生成全部资源
|
||||
- **复制现有发型**：从已部署环境复制 `<hair_id>` 的 4 个目录
|
||||
|
||||
### 4. 查询可用发型列表
|
||||
调 `hair_grow_service` 的接口（需启动 8899）：
|
||||
```
|
||||
GET http://<host>:8899/api/hairstyles
|
||||
→ {"state":0, "count":5, "data":[{"hair_id":"chang_tuoyuan","gender":"girl"}, ...]}
|
||||
```
|
||||
> 注意：该接口在 `hair_grow_swap.py:list_hairstyles()` 内有白名单过滤（当前只返回 5 个 chang_* 发型），集成时按需调整或去掉。
|
||||
|
||||
### 5. 获取发型预览图
|
||||
```
|
||||
GET http://<host>:8899/preview/<hair_id>
|
||||
→ 直接返回预览图（优先 static/previews/<hair_id>.jpg，回退到 ref_rgb_8uc3_768.png）
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 七、训练新发型
|
||||
|
||||
### 7.1 一键训练脚本
|
||||
|
||||
```bash
|
||||
cd /home/xsl/change_hair/project/hair_service_sd
|
||||
python /home/xsl/change_hair/train_hairstyle_full.py \
|
||||
--hair-id <新发型ID> \
|
||||
--input <训练图目录> \ # 目录里放该发型的 jpg/png（1张即可，多张效果更好）
|
||||
--gender <girl|boy> \
|
||||
--template-img <模板图> # 用于生成材质和预览，通常是该发型最代表性的图
|
||||
```
|
||||
|
||||
### 7.2 训练 5 个阶段
|
||||
|
||||
| 阶段 | 功能 | 产物 | 耗时 |
|
||||
|------|------|------|------|
|
||||
| step1 | 准备训练数据 | 抠图 + 数据增广（水平翻转 × 5 分辨率）→ `train_material/<hair_id>/images/1_hairstyle/` 约 20 张 png+txt | ~1 分钟 |
|
||||
| step2 | LoRA 训练 | kohya 训练 1500 步 → `train_material/<hair_id>/model/hairstyle_hd_lora.safetensors` | ~10 分钟（单卡） |
|
||||
| step3 | 生成 ref 材质 | 调 hair_service_sd 回调 → `ref_hairstyle/<hair_id>/` 6 个文件 | ~1 秒 |
|
||||
| step4 | 生成模板材质 | → `hair_template_material/<hair_id>/` 3 个文件 | ~3 秒 |
|
||||
| step5 | 生成预览图 | → `hair_grow_service/static/previews/<hair_id>.jpg` | ~10 秒 |
|
||||
|
||||
### 7.3 批量并行训练（多个发型）
|
||||
|
||||
用 `train_lora_parallel.py` 直接调 kohya，支持并发：
|
||||
|
||||
```python
|
||||
# 编辑 train_lora_parallel.py 顶部的 HAIRSTYLES 列表
|
||||
HAIRSTYLES = ["新发型1", "新发型2", "新发型3"]
|
||||
PARALLEL = 2 # 并发数（按显存调整，5090 32G 建议 2）
|
||||
|
||||
# step1 数据需先用 train_hairstyles_parallel.py 或逐个跑 step1 准备好
|
||||
python train_lora_parallel.py
|
||||
```
|
||||
|
||||
### 7.4 训练注意事项
|
||||
- **hair_id 命名**：只能用英文/数字/下划线，不能含中文或特殊字符（会作目录名、文件名、LoRA 名）
|
||||
- **单图训练**：1 张图经增广出 ~20 张训练图，LoRA 能用但泛化有限；多图（不同角度/光照）效果更好
|
||||
- **GPU 显存**：训练峰值约 13G/卡，推理约 6G。并行训练需控制并发数避免 OOM
|
||||
- **训练前提**：webui + photo_service + hair_service_sd 必须在运行（step3 回调依赖）
|
||||
|
||||
---
|
||||
|
||||
## 八、环境依赖
|
||||
|
||||
### Conda 环境
|
||||
|
||||
| 环境 | Python | 用途 | 关键依赖 |
|
||||
|------|--------|------|---------|
|
||||
| `my_hair` | 3.10 | hair_service_sd / hair_grow_service / 训练 step1 | torch, cv2, flask, gevent, oss2 |
|
||||
| `sdwebui` | 3.10 | webui | torch, gradio, transformers（AUTOMATIC1111 全套） |
|
||||
| `py310` | 3.10 | photo_service | torch, flask, requests |
|
||||
| `kohya` | 3.10 | LoRA 训练（kohya_ss） | torch, accelerate, safetensors |
|
||||
|
||||
### 关键兼容性处理（迁移时注意）
|
||||
|
||||
1. **PyTorch 2.6+ 的 `weights_only`**：加载旧权重（yolov5l.pt 等）需 `torch.load(..., weights_only=False)`。代码已在 `run_copy_cost_colorb64.py` 和 `hair_grow_service/app.py` 顶部做了全局 patch。
|
||||
|
||||
2. **OSS 上传（可选）**：`output_format=url` 时需配 OSS 凭证环境变量：
|
||||
```bash
|
||||
export OSS_TEST_ACCESS_KEY_ID=<...>
|
||||
export OSS_TEST_ACCESS_KEY_SECRET=<...>
|
||||
export OSS_TEST_BUCKET=<...>
|
||||
export OSS_TEST_ENDPOINT=<...>
|
||||
```
|
||||
用 `output_format=base64` 则完全不需要 OSS。
|
||||
|
||||
3. **webui xformers**：kohya 训练命令用 `--sdpa`（PyTorch 原生），不要用 `--xformers`（若 xformers 与 torch 版本不匹配会报 CUDA 错误）。
|
||||
|
||||
---
|
||||
|
||||
## 九、故障排查
|
||||
|
||||
| 现象 | 可能原因 | 排查 |
|
||||
|------|---------|------|
|
||||
| 换发型返回 `推理发型失败` | webui 没启动 / LoRA 文件缺失 | 检查 webui 57860 端口；检查 `train_material/<hair_id>/model/` 是否有 LoRA |
|
||||
| 换发型返回 `发型模板加载失败` | `upload_train_imgs/<hair_id>/` 缺 `first##` 开头的文件 | 检查 5.3-4 目录 |
|
||||
| 换发型返回 `检查遮挡失败` | `hair_template_material/<hair_id>/` 缺 matting/pkl | 检查 5.3-3 目录 |
|
||||
| 预览图 404 | `ref_hairstyle/<hair_id>/ref_rgb_8uc3_768.png` 不存在 | 检查 5.3-1 目录 |
|
||||
| 发型列表少发型 | `hair_grow_swap.py:list_hairstyles()` 白名单过滤 | 调整 `_VISIBLE_HAIRSTYLES` 集合 |
|
||||
| webui 启动报 xformers 错 | xformers 与 torch 版本不匹配 | webui 启动去掉 `--xformers`，或用 `--sdpa` |
|
||||
| 训练报 CUDA error | 显存不足或 xformers 问题 | 训练命令用 `--sdpa`；降低并发数 |
|
||||
|
||||
### 日志位置
|
||||
- webui：`project/logs/webui.log`
|
||||
- photo_service：`project/logs/photo_service.log`
|
||||
- hair_service_sd：`project/logs/hair_service.log`
|
||||
- hair_grow_service：`project/logs/hair_grow_service.log`
|
||||
- 训练：`project/logs/train_<hair_id>.log`
|
||||
|
||||
---
|
||||
|
||||
## 十、目录结构总览
|
||||
|
||||
```
|
||||
/home/xsl/change_hair/
|
||||
├── project/
|
||||
│ ├── hair_service_sd/ # 换发型主服务
|
||||
│ │ ├── weights/ # 核心模型权重（10G）
|
||||
│ │ ├── config/configure.ini # 路径配置
|
||||
│ │ ├── run_copy_cost_colorb64.py # 主入口（API 定义）
|
||||
│ │ ├── gen_super_image.py # webui img2img 调用
|
||||
│ │ └── core/ # 发型推理核心逻辑
|
||||
│ ├── photo_service/ # LoRA 调度服务
|
||||
│ │ └── lora_train_service_1.py # 训练调度 + 推理 LoRA 切换
|
||||
│ ├── onediff/ # SD webui（推理引擎，14G）
|
||||
│ │ └── stable-diffusion-webui/models/
|
||||
│ │ ├── Stable-diffusion/ # SD 底模
|
||||
│ │ └── Lora/ # 各发型 LoRA（自动生成）
|
||||
│ ├── kohya_ss_home/ # LoRA 训练框架（13G）
|
||||
│ ├── data/ # 数据目录（由 configure.ini 指向）
|
||||
│ │ ├── ref_hairstyle/ # 发型材质（5.3-1）
|
||||
│ │ ├── hair_template_material/ # 模板材质（5.3-3）
|
||||
│ │ ├── upload_train_imgs/ # 发型模板原图（5.3-4）
|
||||
│ │ └── ...
|
||||
│ └── logs/
|
||||
├── data/
|
||||
│ └── train_material/ # LoRA + 训练图（5.3-2）
|
||||
├── hair_grow_service/ # 前端页面（可选）
|
||||
├── train_hairstyle_full.py # 单发型一键训练
|
||||
├── train_lora_parallel.py # 批量并行训练
|
||||
├── start_all.sh # 服务启停脚本
|
||||
└── docs/换发型集成文档.md # 本文档
|
||||
```
|
||||
@@ -0,0 +1,302 @@
|
||||
# 方案 A：换发型服务本机部署 — 上传文件清单
|
||||
|
||||
> 适用场景：从原服务器打包上传 conda 环境 + SD WebUI 代码，在本机完成换发型推理部署。
|
||||
>
|
||||
> 本机路径：`/home/xsl/change_hair`（GPU：RTX 5090 32G，CUDA 12.8）
|
||||
>
|
||||
> 原服务器参考：`szlc@192.168.101.63`
|
||||
>
|
||||
> 更新日期：2026-07-09（核实本机已有资源，更新上传清单）
|
||||
|
||||
---
|
||||
|
||||
## 一、本机已有资源（已到位，无需再传）✅
|
||||
|
||||
以下资源已确认存在于本机项目中：
|
||||
|
||||
| 本机路径 | 大小 | 说明 |
|
||||
|---------|------|------|
|
||||
| `project/hair_service_sd/weights/` | 9.9G | 核心推理权重（28个文件，含 GAN/分割/关键点/人脸等全部模型） |
|
||||
| `project/onediff/stable-diffusion-webui/models/` | 4.6G | SD 底模 `v1-5-pruned-emaonly.safetensors` + ESRGAN `4x-UltraSharp.pth` + GFPGAN 全套 |
|
||||
| `project/data/ref_hairstyle/` | 40M | 5 个发型材质（config.json + npy + png + txt） |
|
||||
| `project/data/hair_template_material/` | 24M | 5 个发型模板（pkl + png + matting） |
|
||||
| `project/data/upload_train_imgs/` | 2.5M | 5 个发型原图（jpg） |
|
||||
| `data/train_material/` | 739M | 5 个发型 LoRA 权重（`hairstyle_hd_lora.safetensors`）+ 训练图像 |
|
||||
|
||||
**当前可用发型（5 个）：**
|
||||
|
||||
- `chang_bolang`（长波浪）
|
||||
- `chang_huaban`（长花瓣）
|
||||
- `chang_tuoyuan`（长椭圆）
|
||||
- `chang_xinxing`（长心形）
|
||||
- `chang_zhixian`（长直线）
|
||||
|
||||
> 以上发型可做换发型推理。如需更多发型，见第四节可选上传。
|
||||
|
||||
**项目代码（已在 git 仓库中，无需上传）：**
|
||||
|
||||
- `project/hair_service_sd/`（代码，不含 weights）
|
||||
- `project/photo_service/`
|
||||
- `hair_grow_service/`
|
||||
- 启动脚本、测试脚本等
|
||||
|
||||
---
|
||||
|
||||
## 二、必须上传（换发型推理最低配置）🔴
|
||||
|
||||
以下 4 项需要从原服务器打包上传，**总计约 18–30G**。
|
||||
|
||||
### 2.1 Conda 环境 × 3
|
||||
|
||||
从原服务器 `/home/szlc/miniconda3/envs/` 打包，上传到本机 `/home/xsl/miniconda3/envs/`。
|
||||
|
||||
| 环境名 | 原服务器路径 | 本机目标路径 | 用途 | 预估大小 |
|
||||
|--------|-------------|-------------|------|---------|
|
||||
| `my_hair` | `/home/szlc/miniconda3/envs/my_hair/` | `/home/xsl/miniconda3/envs/my_hair/` | hair_service_sd 主服务（端口 8801） | ~5–8G |
|
||||
| `sdwebui` | `/home/szlc/miniconda3/envs/sdwebui/` | `/home/xsl/miniconda3/envs/sdwebui/` | SD WebUI 推理（端口 57860） | ~8–12G |
|
||||
| `py310` | `/home/szlc/miniconda3/envs/py310/` | `/home/xsl/miniconda3/envs/py310/` | photo_service LoRA 调度（端口 32678） | ~2–4G |
|
||||
|
||||
> ⚠️ **为什么必须上传，不能从公网安装？** 这三个环境包含特定版本的 PyTorch + CUDA 扩展、历史版本的 pip 包、以及编译好的 C++/CUDA 算子，依赖关系复杂。项目中的 `env.yaml` 是旧的 Python 3.7 规格，与实际运行环境（Python 3.10）不一致，无法用于重建。
|
||||
|
||||
#### 原服务器打包示例
|
||||
|
||||
```bash
|
||||
# 在原服务器（szlc@192.168.101.63）执行
|
||||
cd /home/szlc/miniconda3/envs
|
||||
tar czf my_hair.tar.gz my_hair/
|
||||
tar czf sdwebui.tar.gz sdwebui/
|
||||
tar czf py310.tar.gz py310/
|
||||
```
|
||||
|
||||
#### 本机解压示例
|
||||
|
||||
```bash
|
||||
# 在本机执行（上传完成后）
|
||||
mkdir -p /home/xsl/miniconda3/envs
|
||||
cd /home/xsl/miniconda3/envs
|
||||
tar xzf /path/to/my_hair.tar.gz
|
||||
tar xzf /path/to/sdwebui.tar.gz
|
||||
tar xzf /path/to/py310.tar.gz
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### 2.2 SD WebUI 代码（不含 models）
|
||||
|
||||
从原服务器打包 **webui 程序本体**，**排除 models 目录**（本机已有 4.6G 模型权重）。
|
||||
|
||||
| 原服务器路径 | 本机目标路径 | 说明 | 预估大小 |
|
||||
|-------------|-------------|------|---------|
|
||||
| `/home/szlc/project/onediff/stable-diffusion-webui/`（排除 `models/`） | `/home/xsl/change_hair/project/onediff/stable-diffusion-webui/` | webui.py、modules/、extensions/、repositories/ 等 | ~2–5G |
|
||||
|
||||
> ⚠️ 本机当前 `project/onediff/stable-diffusion-webui/` 目录只有 `models/` 子目录（4.6G 模型权重），缺少 `webui.py`、`modules/`、`extensions/` 等所有代码文件。
|
||||
|
||||
**必须包含的关键文件/目录：**
|
||||
|
||||
```
|
||||
stable-diffusion-webui/
|
||||
├── webui.py # 主入口（必需）
|
||||
├── webui.sh # 启动脚本
|
||||
├── webui-user.sh # 用户配置
|
||||
├── launch.py
|
||||
├── modules/ # webui 核心模块
|
||||
├── extensions/ # 扩展（如有）
|
||||
├── repositories/ # 依赖仓库（如有）
|
||||
├── config.json # 配置文件
|
||||
├── ui-config.json # UI 配置（如有）
|
||||
└── style.css / javascript/ # 前端资源
|
||||
```
|
||||
|
||||
**不要上传（本机已有，勿覆盖）：**
|
||||
|
||||
```
|
||||
models/Stable-diffusion/v1-5-pruned-emaonly.safetensors # 4.0G
|
||||
models/ESRGAN/4x-UltraSharp.pth # 64M
|
||||
models/GFPGAN/GFPGANv1.4.pth # 333M
|
||||
models/GFPGAN/detection_Resnet50_Final.pth # 105M
|
||||
models/GFPGAN/parsing_parsenet.pth # 82M
|
||||
```
|
||||
|
||||
#### 原服务器打包示例（排除 models）
|
||||
|
||||
```bash
|
||||
# 在原服务器（szlc@192.168.101.63）执行
|
||||
cd /home/szlc/project/onediff
|
||||
tar czf webui_code.tar.gz \
|
||||
--exclude='stable-diffusion-webui/models' \
|
||||
stable-diffusion-webui/
|
||||
```
|
||||
|
||||
#### 本机解压示例
|
||||
|
||||
```bash
|
||||
# 在本机执行（保留已有 models 目录）
|
||||
cd /home/xsl/change_hair/project/onediff
|
||||
tar xzf /path/to/webui_code.tar.gz
|
||||
# 解压后确认 models/ 仍在且未被覆盖
|
||||
ls -lh stable-diffusion-webui/models/Stable-diffusion/
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 三、上传清单速查表
|
||||
|
||||
复制以下列表，在原服务器逐项确认后打包：
|
||||
|
||||
```
|
||||
# ===== 必须上传（方案 A，共 4 项，约 18–30G）=====
|
||||
|
||||
# 1. Conda 环境（3 个）
|
||||
/home/szlc/miniconda3/envs/my_hair/ → /home/xsl/miniconda3/envs/my_hair/
|
||||
/home/szlc/miniconda3/envs/sdwebui/ → /home/xsl/miniconda3/envs/sdwebui/
|
||||
/home/szlc/miniconda3/envs/py310/ → /home/xsl/miniconda3/envs/py310/
|
||||
|
||||
# 2. SD WebUI 代码（排除 models/）
|
||||
/home/szlc/project/onediff/stable-diffusion-webui/（排除 models/）
|
||||
→ /home/xsl/change_hair/project/onediff/stable-diffusion-webui/
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 四、可选上传（非换发型推理必需）🟡
|
||||
|
||||
以下功能暂不需要时可跳过，后续按需补传。
|
||||
|
||||
### 4.1 训练新发型
|
||||
|
||||
| 原服务器路径 | 本机目标路径 | 预估大小 |
|
||||
|-------------|-------------|---------|
|
||||
| `/home/szlc/project/kohya_ss_home/` | `/home/xsl/change_hair/project/kohya_ss_home/` | ~13G |
|
||||
| `/home/szlc/miniconda3/envs/kohya/` | `/home/xsl/miniconda3/envs/kohya/` | ~5G |
|
||||
|
||||
### 4.2 换发色功能
|
||||
|
||||
| 原服务器路径 | 本机目标路径 |
|
||||
|-------------|-------------|
|
||||
| `/home/szlc/project/data/ref_haircolor/` | `/home/xsl/change_hair/project/data/ref_haircolor/` |
|
||||
| `/home/szlc/project/data/ref_color/` | `/home/xsl/change_hair/project/data/ref_color/` |
|
||||
| `/home/szlc/project/data/ref_color_imgs/` | `/home/xsl/change_hair/project/data/ref_color_imgs/` |
|
||||
| `/home/szlc/.../weights/classify_3_2499.pth` | `project/hair_service_sd/weights/classify_3_2499.pth` |
|
||||
|
||||
### 4.3 更多发型
|
||||
|
||||
当前只有 5 个发型。如需更多，从原服务器追加：
|
||||
|
||||
```
|
||||
/data/train_material/<hair_id>/ → data/train_material/<hair_id>/
|
||||
/home/szlc/project/data/ref_hairstyle/<hair_id>/ → project/data/ref_hairstyle/<hair_id>/
|
||||
/home/szlc/project/data/hair_template_material/<hair_id>/ → project/data/hair_template_material/<hair_id>/
|
||||
/home/szlc/project/data/upload_train_imgs/<hair_id>/ → project/data/upload_train_imgs/<hair_id>/
|
||||
```
|
||||
|
||||
原服务器约有 338 个发型，全量 `train_material` 约 121G。
|
||||
|
||||
### 4.4 小文件（可选）
|
||||
|
||||
| 文件 | 用途 | 大小 |
|
||||
|------|------|------|
|
||||
| `weights/mediapipe/face_landmarker.task` | 发际线检测 | ~3.6M |
|
||||
|
||||
---
|
||||
|
||||
## 五、上传方式建议
|
||||
|
||||
### 方式 1：rsync 直传（推荐，支持断点续传）
|
||||
|
||||
```bash
|
||||
# 从原服务器直传到本机（在本机执行）
|
||||
rsync -avP --partial \
|
||||
szlc@192.168.101.63:/home/szlc/miniconda3/envs/my_hair/ \
|
||||
/home/xsl/miniconda3/envs/my_hair/
|
||||
|
||||
rsync -avP --partial \
|
||||
szlc@192.168.101.63:/home/szlc/miniconda3/envs/sdwebui/ \
|
||||
/home/xsl/miniconda3/envs/sdwebui/
|
||||
|
||||
rsync -avP --partial \
|
||||
szlc@192.168.101.63:/home/szlc/miniconda3/envs/py310/ \
|
||||
/home/xsl/miniconda3/envs/py310/
|
||||
|
||||
rsync -avP --partial \
|
||||
--exclude='models/' \
|
||||
szlc@192.168.101.63:/home/szlc/project/onediff/stable-diffusion-webui/ \
|
||||
/home/xsl/change_hair/project/onediff/stable-diffusion-webui/
|
||||
```
|
||||
|
||||
### 方式 2：先 tar 再上传到 `/home/xsl/data`
|
||||
|
||||
若通过网盘中转，先 tar 再上传到 `/home/xsl/data/`，上传完成后告知我解压。
|
||||
|
||||
---
|
||||
|
||||
## 六、上传完成后的验证
|
||||
|
||||
上传完成后，在本机执行以下命令确认文件到位：
|
||||
|
||||
```bash
|
||||
# 1. Conda 环境
|
||||
ls /home/xsl/miniconda3/envs/my_hair/bin/python
|
||||
ls /home/xsl/miniconda3/envs/sdwebui/bin/python
|
||||
ls /home/xsl/miniconda3/envs/py310/bin/python
|
||||
|
||||
# 2. WebUI 代码
|
||||
ls /home/xsl/change_hair/project/onediff/stable-diffusion-webui/webui.py
|
||||
|
||||
# 3. 模型未被覆盖
|
||||
ls -lh /home/xsl/change_hair/project/onediff/stable-diffusion-webui/models/Stable-diffusion/v1-5-pruned-emaonly.safetensors
|
||||
```
|
||||
|
||||
期望结果：4 项均存在，SD 底模约 4.0G。
|
||||
|
||||
---
|
||||
|
||||
## 七、上传完成后我会做的事
|
||||
|
||||
你上传完毕并告知后，我会依次执行：
|
||||
|
||||
1. **路径适配**：运行 `adapt_paths.sh`，将所有硬编码路径从 `/home/xsl/change_hair` 改为 `/home/xsl/change_hair`，`/home/xsl/miniconda3` 改为 `/home/xsl/miniconda3`
|
||||
2. **Conda 路径修复**：运行 `fix_conda_paths.py`，修复 3 个环境中二进制文件的 shebang 和内部路径前缀
|
||||
3. **安装 Miniconda 基础环境**（如尚未安装）：从 `repo.anaconda.com` 下载安装
|
||||
4. **启动服务**：`bash start_all.sh start`（依赖顺序：webui → photo_service → hair_service_sd）
|
||||
5. **接口测试**：用 `test_swaphair.py` 验证换发型功能
|
||||
|
||||
---
|
||||
|
||||
## 八、部署后服务端口
|
||||
|
||||
| 服务 | 端口 | 说明 |
|
||||
|------|------|------|
|
||||
| webui | 57860 | SD img2img 推理 |
|
||||
| photo_service | 32678 | LoRA 加载与调度 |
|
||||
| hair_service_sd | 8801 | 换发型主接口 |
|
||||
| hair_grow_service | 8899 | 生发前端（独立 Flask） |
|
||||
|
||||
换发型接口：
|
||||
|
||||
```
|
||||
POST http://127.0.0.1:8801/api/swapHair/v1
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 九、常见问题
|
||||
|
||||
**Q：能否只上传 miniconda3 整包，不传单个 env？**
|
||||
|
||||
可以。上传 `/home/szlc/miniconda3/` 整包到 `/home/xsl/miniconda3/` 也行，体积更大（约 20–30G）但更省事。
|
||||
|
||||
**Q：webui 上传时不小心覆盖了 models 怎么办？**
|
||||
|
||||
重新从 `/home/xsl/data/sdwebui/models/` 或原服务器补传 models 目录即可。
|
||||
|
||||
**Q：本机没有 miniconda 基础安装，只传 envs 够吗？**
|
||||
|
||||
不够。3 个 env 内的 python 依赖 miniconda 基础库（如 `libstdc++`、`libgcc` 等）。我会在上传完成后先安装 miniconda，再修复 env 内部路径。
|
||||
|
||||
**Q：上传后启动报 xformers 错误？**
|
||||
|
||||
webui 启动参数可能需要去掉 `--xformers`，改用 `--opt-sdp-no-mem-attention`（RTX 5090 支持，我会在路径适配时处理）。
|
||||
|
||||
**Q：本机磁盘空间是否足够？**
|
||||
|
||||
当前 `/home` 剩余约 30G。3 个 conda 环境约 20–25G，解压后可能吃紧。建议解压一个删一个 tar 包，或考虑清理 `/home/xsl/data/`（已有 16G 冗余数据，解压完成后可删除）。
|
||||
@@ -24,10 +24,19 @@ os.environ.setdefault("CRYPTOGRAPHY_OPENSSL_NO_LEGACY", "1")
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from flask import Flask, request, jsonify, send_from_directory
|
||||
from flask import Flask, request, jsonify, send_from_directory, send_file
|
||||
from gevent import pywsgi
|
||||
|
||||
PORT = 8888
|
||||
# PyTorch 2.6+ 默认 weights_only=True，无法加载含 numpy 对象的旧权重（yolov5l.pt 等）。
|
||||
# 统一改回 False（权重均为本机自有可信文件）。
|
||||
import torch
|
||||
_orig_torch_load = torch.load
|
||||
def _torch_load(*args, **kwargs):
|
||||
kwargs.setdefault("weights_only", False)
|
||||
return _orig_torch_load(*args, **kwargs)
|
||||
torch.load = _torch_load
|
||||
|
||||
PORT = 8899
|
||||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
app = Flask(__name__, static_folder="static", static_url_path="/static")
|
||||
@@ -394,7 +403,7 @@ def preview_img(hair_id):
|
||||
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)
|
||||
return send_file(fallback)
|
||||
except Exception:
|
||||
pass
|
||||
return ("", 404)
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 508 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 480 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 513 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 493 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 463 KiB |
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,3 @@
|
||||
from .wider_face import WiderFaceDetection, detection_collate
|
||||
from .data_augment import *
|
||||
from .config import *
|
||||
@@ -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,237 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
import random
|
||||
from core.utils.box_utils_Retina import matrix_iof
|
||||
|
||||
|
||||
def _crop(image, boxes, labels, landm, img_dim):
|
||||
height, width, _ = image.shape
|
||||
pad_image_flag = True
|
||||
|
||||
for _ in range(250):
|
||||
"""
|
||||
if random.uniform(0, 1) <= 0.2:
|
||||
scale = 1.0
|
||||
else:
|
||||
scale = random.uniform(0.3, 1.0)
|
||||
"""
|
||||
PRE_SCALES = [0.3, 0.45, 0.6, 0.8, 1.0]
|
||||
scale = random.choice(PRE_SCALES)
|
||||
short_side = min(width, height)
|
||||
w = int(scale * short_side)
|
||||
h = w
|
||||
|
||||
if width == w:
|
||||
l = 0
|
||||
else:
|
||||
l = random.randrange(width - w)
|
||||
if height == h:
|
||||
t = 0
|
||||
else:
|
||||
t = random.randrange(height - h)
|
||||
roi = np.array((l, t, l + w, t + h))
|
||||
|
||||
value = matrix_iof(boxes, roi[np.newaxis])
|
||||
flag = (value >= 1)
|
||||
if not flag.any():
|
||||
continue
|
||||
|
||||
centers = (boxes[:, :2] + boxes[:, 2:]) / 2
|
||||
mask_a = np.logical_and(roi[:2] < centers, centers < roi[2:]).all(axis=1)
|
||||
boxes_t = boxes[mask_a].copy()
|
||||
labels_t = labels[mask_a].copy()
|
||||
landms_t = landm[mask_a].copy()
|
||||
landms_t = landms_t.reshape([-1, 5, 2])
|
||||
|
||||
if boxes_t.shape[0] == 0:
|
||||
continue
|
||||
|
||||
image_t = image[roi[1]:roi[3], roi[0]:roi[2]]
|
||||
|
||||
boxes_t[:, :2] = np.maximum(boxes_t[:, :2], roi[:2])
|
||||
boxes_t[:, :2] -= roi[:2]
|
||||
boxes_t[:, 2:] = np.minimum(boxes_t[:, 2:], roi[2:])
|
||||
boxes_t[:, 2:] -= roi[:2]
|
||||
|
||||
# landm
|
||||
landms_t[:, :, :2] = landms_t[:, :, :2] - roi[:2]
|
||||
landms_t[:, :, :2] = np.maximum(landms_t[:, :, :2], np.array([0, 0]))
|
||||
landms_t[:, :, :2] = np.minimum(landms_t[:, :, :2], roi[2:] - roi[:2])
|
||||
landms_t = landms_t.reshape([-1, 10])
|
||||
|
||||
|
||||
# make sure that the cropped image contains at least one face > 16 pixel at training image scale
|
||||
b_w_t = (boxes_t[:, 2] - boxes_t[:, 0] + 1) / w * img_dim
|
||||
b_h_t = (boxes_t[:, 3] - boxes_t[:, 1] + 1) / h * img_dim
|
||||
mask_b = np.minimum(b_w_t, b_h_t) > 0.0
|
||||
boxes_t = boxes_t[mask_b]
|
||||
labels_t = labels_t[mask_b]
|
||||
landms_t = landms_t[mask_b]
|
||||
|
||||
if boxes_t.shape[0] == 0:
|
||||
continue
|
||||
|
||||
pad_image_flag = False
|
||||
|
||||
return image_t, boxes_t, labels_t, landms_t, pad_image_flag
|
||||
return image, boxes, labels, landm, pad_image_flag
|
||||
|
||||
|
||||
def _distort(image):
|
||||
|
||||
def _convert(image, alpha=1, beta=0):
|
||||
tmp = image.astype(float) * alpha + beta
|
||||
tmp[tmp < 0] = 0
|
||||
tmp[tmp > 255] = 255
|
||||
image[:] = tmp
|
||||
|
||||
image = image.copy()
|
||||
|
||||
if random.randrange(2):
|
||||
|
||||
#brightness distortion
|
||||
if random.randrange(2):
|
||||
_convert(image, beta=random.uniform(-32, 32))
|
||||
|
||||
#contrast distortion
|
||||
if random.randrange(2):
|
||||
_convert(image, alpha=random.uniform(0.5, 1.5))
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
|
||||
|
||||
#saturation distortion
|
||||
if random.randrange(2):
|
||||
_convert(image[:, :, 1], alpha=random.uniform(0.5, 1.5))
|
||||
|
||||
#hue distortion
|
||||
if random.randrange(2):
|
||||
tmp = image[:, :, 0].astype(int) + random.randint(-18, 18)
|
||||
tmp %= 180
|
||||
image[:, :, 0] = tmp
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
|
||||
|
||||
else:
|
||||
|
||||
#brightness distortion
|
||||
if random.randrange(2):
|
||||
_convert(image, beta=random.uniform(-32, 32))
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
|
||||
|
||||
#saturation distortion
|
||||
if random.randrange(2):
|
||||
_convert(image[:, :, 1], alpha=random.uniform(0.5, 1.5))
|
||||
|
||||
#hue distortion
|
||||
if random.randrange(2):
|
||||
tmp = image[:, :, 0].astype(int) + random.randint(-18, 18)
|
||||
tmp %= 180
|
||||
image[:, :, 0] = tmp
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
|
||||
|
||||
#contrast distortion
|
||||
if random.randrange(2):
|
||||
_convert(image, alpha=random.uniform(0.5, 1.5))
|
||||
|
||||
return image
|
||||
|
||||
|
||||
def _expand(image, boxes, fill, p):
|
||||
if random.randrange(2):
|
||||
return image, boxes
|
||||
|
||||
height, width, depth = image.shape
|
||||
|
||||
scale = random.uniform(1, p)
|
||||
w = int(scale * width)
|
||||
h = int(scale * height)
|
||||
|
||||
left = random.randint(0, w - width)
|
||||
top = random.randint(0, h - height)
|
||||
|
||||
boxes_t = boxes.copy()
|
||||
boxes_t[:, :2] += (left, top)
|
||||
boxes_t[:, 2:] += (left, top)
|
||||
expand_image = np.empty(
|
||||
(h, w, depth),
|
||||
dtype=image.dtype)
|
||||
expand_image[:, :] = fill
|
||||
expand_image[top:top + height, left:left + width] = image
|
||||
image = expand_image
|
||||
|
||||
return image, boxes_t
|
||||
|
||||
|
||||
def _mirror(image, boxes, landms):
|
||||
_, width, _ = image.shape
|
||||
if random.randrange(2):
|
||||
image = image[:, ::-1]
|
||||
boxes = boxes.copy()
|
||||
boxes[:, 0::2] = width - boxes[:, 2::-2]
|
||||
|
||||
# landm
|
||||
landms = landms.copy()
|
||||
landms = landms.reshape([-1, 5, 2])
|
||||
landms[:, :, 0] = width - landms[:, :, 0]
|
||||
tmp = landms[:, 1, :].copy()
|
||||
landms[:, 1, :] = landms[:, 0, :]
|
||||
landms[:, 0, :] = tmp
|
||||
tmp1 = landms[:, 4, :].copy()
|
||||
landms[:, 4, :] = landms[:, 3, :]
|
||||
landms[:, 3, :] = tmp1
|
||||
landms = landms.reshape([-1, 10])
|
||||
|
||||
return image, boxes, landms
|
||||
|
||||
|
||||
def _pad_to_square(image, rgb_mean, pad_image_flag):
|
||||
if not pad_image_flag:
|
||||
return image
|
||||
height, width, _ = image.shape
|
||||
long_side = max(width, height)
|
||||
image_t = np.empty((long_side, long_side, 3), dtype=image.dtype)
|
||||
image_t[:, :] = rgb_mean
|
||||
image_t[0:0 + height, 0:0 + width] = image
|
||||
return image_t
|
||||
|
||||
|
||||
def _resize_subtract_mean(image, insize, rgb_mean):
|
||||
interp_methods = [cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_NEAREST, cv2.INTER_LANCZOS4]
|
||||
interp_method = interp_methods[random.randrange(5)]
|
||||
image = cv2.resize(image, (insize, insize), interpolation=interp_method)
|
||||
image = image.astype(np.float32)
|
||||
image -= rgb_mean
|
||||
return image.transpose(2, 0, 1)
|
||||
|
||||
|
||||
class preproc(object):
|
||||
|
||||
def __init__(self, img_dim, rgb_means):
|
||||
self.img_dim = img_dim
|
||||
self.rgb_means = rgb_means
|
||||
|
||||
def __call__(self, image, targets):
|
||||
assert targets.shape[0] > 0, "this image does not have gt"
|
||||
|
||||
boxes = targets[:, :4].copy()
|
||||
labels = targets[:, -1].copy()
|
||||
landm = targets[:, 4:-1].copy()
|
||||
|
||||
image_t, boxes_t, labels_t, landm_t, pad_image_flag = _crop(image, boxes, labels, landm, self.img_dim)
|
||||
image_t = _distort(image_t)
|
||||
image_t = _pad_to_square(image_t,self.rgb_means, pad_image_flag)
|
||||
image_t, boxes_t, landm_t = _mirror(image_t, boxes_t, landm_t)
|
||||
height, width, _ = image_t.shape
|
||||
image_t = _resize_subtract_mean(image_t, self.img_dim, self.rgb_means)
|
||||
boxes_t[:, 0::2] /= width
|
||||
boxes_t[:, 1::2] /= height
|
||||
|
||||
landm_t[:, 0::2] /= width
|
||||
landm_t[:, 1::2] /= height
|
||||
|
||||
labels_t = np.expand_dims(labels_t, 1)
|
||||
targets_t = np.hstack((boxes_t, landm_t, labels_t))
|
||||
|
||||
return image_t, targets_t
|
||||
@@ -0,0 +1,101 @@
|
||||
import os
|
||||
import os.path
|
||||
import sys
|
||||
import torch
|
||||
import torch.utils.data as data
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
class WiderFaceDetection(data.Dataset):
|
||||
def __init__(self, txt_path, preproc=None):
|
||||
self.preproc = preproc
|
||||
self.imgs_path = []
|
||||
self.words = []
|
||||
f = open(txt_path,'r')
|
||||
lines = f.readlines()
|
||||
isFirst = True
|
||||
labels = []
|
||||
for line in lines:
|
||||
line = line.rstrip()
|
||||
if line.startswith('#'):
|
||||
if isFirst is True:
|
||||
isFirst = False
|
||||
else:
|
||||
labels_copy = labels.copy()
|
||||
self.words.append(labels_copy)
|
||||
labels.clear()
|
||||
path = line[2:]
|
||||
path = txt_path.replace('label.txt','images/') + path
|
||||
self.imgs_path.append(path)
|
||||
else:
|
||||
line = line.split(' ')
|
||||
label = [float(x) for x in line]
|
||||
labels.append(label)
|
||||
|
||||
self.words.append(labels)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.imgs_path)
|
||||
|
||||
def __getitem__(self, index):
|
||||
img = cv2.imread(self.imgs_path[index])
|
||||
height, width, _ = img.shape
|
||||
|
||||
labels = self.words[index]
|
||||
annotations = np.zeros((0, 15))
|
||||
if len(labels) == 0:
|
||||
return annotations
|
||||
for idx, label in enumerate(labels):
|
||||
annotation = np.zeros((1, 15))
|
||||
# bbox
|
||||
annotation[0, 0] = label[0] # x1
|
||||
annotation[0, 1] = label[1] # y1
|
||||
annotation[0, 2] = label[0] + label[2] # x2
|
||||
annotation[0, 3] = label[1] + label[3] # y2
|
||||
|
||||
# landmarks
|
||||
annotation[0, 4] = label[4] # l0_x
|
||||
annotation[0, 5] = label[5] # l0_y
|
||||
annotation[0, 6] = label[7] # l1_x
|
||||
annotation[0, 7] = label[8] # l1_y
|
||||
annotation[0, 8] = label[10] # l2_x
|
||||
annotation[0, 9] = label[11] # l2_y
|
||||
annotation[0, 10] = label[13] # l3_x
|
||||
annotation[0, 11] = label[14] # l3_y
|
||||
annotation[0, 12] = label[16] # l4_x
|
||||
annotation[0, 13] = label[17] # l4_y
|
||||
if (annotation[0, 4]<0):
|
||||
annotation[0, 14] = -1
|
||||
else:
|
||||
annotation[0, 14] = 1
|
||||
|
||||
annotations = np.append(annotations, annotation, axis=0)
|
||||
target = np.array(annotations)
|
||||
if self.preproc is not None:
|
||||
img, target = self.preproc(img, target)
|
||||
|
||||
return torch.from_numpy(img), target
|
||||
|
||||
def detection_collate(batch):
|
||||
"""Custom collate fn for dealing with batches of images that have a different
|
||||
number of associated object annotations (bounding boxes).
|
||||
|
||||
Arguments:
|
||||
batch: (tuple) A tuple of tensor images and lists of annotations
|
||||
|
||||
Return:
|
||||
A tuple containing:
|
||||
1) (tensor) batch of images stacked on their 0 dim
|
||||
2) (list of tensors) annotations for a given image are stacked on 0 dim
|
||||
"""
|
||||
targets = []
|
||||
imgs = []
|
||||
for _, sample in enumerate(batch):
|
||||
for _, tup in enumerate(sample):
|
||||
if torch.is_tensor(tup):
|
||||
imgs.append(tup)
|
||||
elif isinstance(tup, type(np.empty(0))):
|
||||
annos = torch.from_numpy(tup).float()
|
||||
targets.append(annos)
|
||||
|
||||
return (torch.stack(imgs, 0), targets)
|
||||
@@ -5,16 +5,24 @@ import os
|
||||
|
||||
class OSS_object():
|
||||
def __init__(self):
|
||||
access_key_id = os.getenv('OSS_TEST_ACCESS_KEY_ID', '<your-access-key-id>')
|
||||
access_key_secret = os.getenv('OSS_TEST_ACCESS_KEY_SECRET', '<your-access-key-secret>')
|
||||
bucket_name = os.getenv('OSS_TEST_BUCKET', '<your-bucket-name>')
|
||||
endpoint = os.getenv('OSS_TEST_ENDPOINT', '<your-endpoint>')
|
||||
for param in (access_key_id, access_key_secret, bucket_name, endpoint):
|
||||
assert '<' not in param, '请设置环境变量 OSS_TEST_ACCESS_KEY_ID / OSS_TEST_ACCESS_KEY_SECRET / OSS_TEST_BUCKET / OSS_TEST_ENDPOINT'
|
||||
# 懒加载：只读取凭证，不立即连接 OSS。
|
||||
# 这样本地用 output_format=base64 时不配 OSS 凭证也能启动服务；
|
||||
# 仅在真正调用 upload_file 时才校验并连接。
|
||||
self.access_key_id = os.getenv('OSS_TEST_ACCESS_KEY_ID', '<your-access-key-id>')
|
||||
self.access_key_secret = os.getenv('OSS_TEST_ACCESS_KEY_SECRET', '<your-access-key-secret>')
|
||||
self.bucket_name = os.getenv('OSS_TEST_BUCKET', '<your-bucket-name>')
|
||||
self.endpoint = os.getenv('OSS_TEST_ENDPOINT', '<your-endpoint>')
|
||||
self.bucket = None
|
||||
|
||||
self.bucket = oss2.Bucket(oss2.Auth(access_key_id, access_key_secret), endpoint, bucket_name)
|
||||
def _ensure_bucket(self):
|
||||
if self.bucket is not None:
|
||||
return
|
||||
for param in (self.access_key_id, self.access_key_secret, self.bucket_name, self.endpoint):
|
||||
assert '<' not in param, '请设置环境变量 OSS_TEST_ACCESS_KEY_ID / OSS_TEST_ACCESS_KEY_SECRET / OSS_TEST_BUCKET / OSS_TEST_ENDPOINT'
|
||||
self.bucket = oss2.Bucket(oss2.Auth(self.access_key_id, self.access_key_secret), self.endpoint, self.bucket_name)
|
||||
|
||||
def upload_file(self, file, target_name):
|
||||
self._ensure_bucket()
|
||||
t0 = time.time()
|
||||
with open(oss2.to_unicode(file), 'rb') as f:
|
||||
ret = self.bucket.put_object(target_name, f)
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
#!/home/ubuntu/miniconda3/envs/yd/bin/python
|
||||
#!/home/xsl/miniconda3/envs/yd/bin/python
|
||||
# -*- coding: utf-8 -*-
|
||||
import os
|
||||
# 监听本机的端口
|
||||
|
||||
@@ -177,10 +177,17 @@ def list_hairstyles(hairstyle_dir, train_dir, upload_dir, limit=None):
|
||||
返回:
|
||||
[{hair_id, gender, has_lora}, ...]
|
||||
"""
|
||||
# 仅展示这些发型（其余数据保留在磁盘，但不在列表中显示）
|
||||
_VISIBLE_HAIRSTYLES = {
|
||||
"chang_tuoyuan", "chang_bolang", "chang_zhixian",
|
||||
"chang_huaban", "chang_xinxing",
|
||||
}
|
||||
styles = []
|
||||
if not os.path.isdir(hairstyle_dir):
|
||||
return styles
|
||||
for hair_id in os.listdir(hairstyle_dir):
|
||||
if hair_id not in _VISIBLE_HAIRSTYLES:
|
||||
continue
|
||||
cfg_path = os.path.join(hairstyle_dir, hair_id, "config.json")
|
||||
ref_path = os.path.join(hairstyle_dir, hair_id, "ref_rgb_8uc3_768.png")
|
||||
lora_path = os.path.join(train_dir, hair_id, "model", "hairstyle_hd_lora.safetensors")
|
||||
|
||||
@@ -0,0 +1,9 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from .default import _C as cfg
|
||||
from .default import update_config
|
||||
from .models import MODEL_EXTRAS
|
||||
@@ -0,0 +1,160 @@
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import os
|
||||
|
||||
from yacs.config import CfgNode as CN
|
||||
|
||||
|
||||
_C = CN()
|
||||
|
||||
_C.OUTPUT_DIR = ''
|
||||
_C.LOG_DIR = ''
|
||||
_C.DATA_DIR = ''
|
||||
_C.GPUS = (0,)
|
||||
_C.WORKERS = 4
|
||||
_C.PRINT_FREQ = 20
|
||||
_C.AUTO_RESUME = False
|
||||
_C.PIN_MEMORY = True
|
||||
_C.RANK = 0
|
||||
|
||||
# Cudnn related params
|
||||
_C.CUDNN = CN()
|
||||
_C.CUDNN.BENCHMARK = True
|
||||
_C.CUDNN.DETERMINISTIC = False
|
||||
_C.CUDNN.ENABLED = True
|
||||
|
||||
# common params for NETWORK
|
||||
_C.MODEL = CN()
|
||||
_C.MODEL.NAME = 'pose_hrnet'
|
||||
_C.MODEL.INIT_WEIGHTS = True
|
||||
_C.MODEL.PRETRAINED = ''
|
||||
_C.MODEL.NUM_JOINTS = 17
|
||||
_C.MODEL.TAG_PER_JOINT = True
|
||||
_C.MODEL.TARGET_TYPE = 'gaussian'
|
||||
_C.MODEL.IMAGE_SIZE = [256, 256] # width * height, ex: 192 * 256
|
||||
_C.MODEL.HEATMAP_SIZE = [64, 64] # width * height, ex: 24 * 32
|
||||
_C.MODEL.SIGMA = 2
|
||||
_C.MODEL.EXTRA = CN(new_allowed=True)
|
||||
|
||||
_C.LOSS = CN()
|
||||
_C.LOSS.USE_OHKM = False
|
||||
_C.LOSS.TOPK = 8
|
||||
_C.LOSS.USE_TARGET_WEIGHT = True
|
||||
_C.LOSS.USE_DIFFERENT_JOINTS_WEIGHT = False
|
||||
|
||||
# DATASET related params
|
||||
_C.DATASET = CN()
|
||||
_C.DATASET.ROOT = ''
|
||||
_C.DATASET.DATASET = 'mpii'
|
||||
_C.DATASET.TRAIN_SET = 'train'
|
||||
_C.DATASET.TEST_SET = 'valid'
|
||||
_C.DATASET.DATA_FORMAT = 'jpg'
|
||||
_C.DATASET.HYBRID_JOINTS_TYPE = ''
|
||||
_C.DATASET.SELECT_DATA = False
|
||||
|
||||
# training data augmentation
|
||||
_C.DATASET.FLIP = True
|
||||
_C.DATASET.SCALE_FACTOR = 0.25
|
||||
_C.DATASET.ROT_FACTOR = 30
|
||||
_C.DATASET.PROB_HALF_BODY = 0.0
|
||||
_C.DATASET.NUM_JOINTS_HALF_BODY = 8
|
||||
_C.DATASET.COLOR_RGB = False
|
||||
|
||||
# train
|
||||
_C.TRAIN = CN()
|
||||
|
||||
_C.TRAIN.LR_FACTOR = 0.1
|
||||
_C.TRAIN.LR_STEP = [90, 110]
|
||||
_C.TRAIN.LR = 0.001
|
||||
|
||||
_C.TRAIN.OPTIMIZER = 'adam'
|
||||
_C.TRAIN.MOMENTUM = 0.9
|
||||
_C.TRAIN.WD = 0.0001
|
||||
_C.TRAIN.NESTEROV = False
|
||||
_C.TRAIN.GAMMA1 = 0.99
|
||||
_C.TRAIN.GAMMA2 = 0.0
|
||||
|
||||
_C.TRAIN.BEGIN_EPOCH = 0
|
||||
_C.TRAIN.END_EPOCH = 140
|
||||
|
||||
_C.TRAIN.RESUME = False
|
||||
_C.TRAIN.CHECKPOINT = ''
|
||||
|
||||
_C.TRAIN.BATCH_SIZE_PER_GPU = 32
|
||||
_C.TRAIN.SHUFFLE = True
|
||||
|
||||
# testing
|
||||
_C.TEST = CN()
|
||||
|
||||
# size of images for each device
|
||||
_C.TEST.BATCH_SIZE_PER_GPU = 32
|
||||
# Test Model Epoch
|
||||
_C.TEST.FLIP_TEST = False
|
||||
_C.TEST.POST_PROCESS = False
|
||||
_C.TEST.SHIFT_HEATMAP = False
|
||||
|
||||
_C.TEST.USE_GT_BBOX = False
|
||||
|
||||
# nms
|
||||
_C.TEST.IMAGE_THRE = 0.1
|
||||
_C.TEST.NMS_THRE = 0.6
|
||||
_C.TEST.SOFT_NMS = False
|
||||
_C.TEST.OKS_THRE = 0.5
|
||||
_C.TEST.IN_VIS_THRE = 0.0
|
||||
_C.TEST.COCO_BBOX_FILE = ''
|
||||
_C.TEST.BBOX_THRE = 1.0
|
||||
_C.TEST.MODEL_FILE = ''
|
||||
|
||||
# debug
|
||||
_C.DEBUG = CN()
|
||||
_C.DEBUG.DEBUG = False
|
||||
_C.DEBUG.SAVE_BATCH_IMAGES_GT = False
|
||||
_C.DEBUG.SAVE_BATCH_IMAGES_PRED = False
|
||||
_C.DEBUG.SAVE_HEATMAPS_GT = False
|
||||
_C.DEBUG.SAVE_HEATMAPS_PRED = False
|
||||
|
||||
|
||||
def update_config(cfg, args):
|
||||
cfg.defrost()
|
||||
cfg.merge_from_file(args.cfg)
|
||||
cfg.merge_from_list(args.opts)
|
||||
|
||||
if args.modelDir:
|
||||
cfg.OUTPUT_DIR = args.modelDir
|
||||
|
||||
if args.logDir:
|
||||
cfg.LOG_DIR = args.logDir
|
||||
|
||||
if args.dataDir:
|
||||
cfg.DATA_DIR = args.dataDir
|
||||
|
||||
cfg.DATASET.ROOT = os.path.join(
|
||||
cfg.DATA_DIR, cfg.DATASET.ROOT
|
||||
)
|
||||
|
||||
cfg.MODEL.PRETRAINED = os.path.join(
|
||||
cfg.DATA_DIR, cfg.MODEL.PRETRAINED
|
||||
)
|
||||
|
||||
if cfg.TEST.MODEL_FILE:
|
||||
cfg.TEST.MODEL_FILE = os.path.join(
|
||||
cfg.DATA_DIR, cfg.TEST.MODEL_FILE
|
||||
)
|
||||
|
||||
cfg.freeze()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
import sys
|
||||
with open(sys.argv[1], 'w') as f:
|
||||
print(_C, file=f)
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
from yacs.config import CfgNode as CN
|
||||
|
||||
|
||||
# pose_resnet related params
|
||||
POSE_RESNET = CN()
|
||||
POSE_RESNET.NUM_LAYERS = 50
|
||||
POSE_RESNET.DECONV_WITH_BIAS = False
|
||||
POSE_RESNET.NUM_DECONV_LAYERS = 3
|
||||
POSE_RESNET.NUM_DECONV_FILTERS = [256, 256, 256]
|
||||
POSE_RESNET.NUM_DECONV_KERNELS = [4, 4, 4]
|
||||
POSE_RESNET.FINAL_CONV_KERNEL = 1
|
||||
POSE_RESNET.PRETRAINED_LAYERS = ['*']
|
||||
|
||||
# pose_multi_resoluton_net related params
|
||||
POSE_HIGH_RESOLUTION_NET = CN()
|
||||
POSE_HIGH_RESOLUTION_NET.PRETRAINED_LAYERS = ['*']
|
||||
POSE_HIGH_RESOLUTION_NET.STEM_INPLANES = 64
|
||||
POSE_HIGH_RESOLUTION_NET.FINAL_CONV_KERNEL = 1
|
||||
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE2 = CN()
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE2.NUM_MODULES = 1
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE2.NUM_BRANCHES = 2
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE2.NUM_BLOCKS = [4, 4]
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE2.NUM_CHANNELS = [32, 64]
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE2.BLOCK = 'BASIC'
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE2.FUSE_METHOD = 'SUM'
|
||||
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE3 = CN()
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE3.NUM_MODULES = 1
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE3.NUM_BRANCHES = 3
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE3.NUM_BLOCKS = [4, 4, 4]
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE3.NUM_CHANNELS = [32, 64, 128]
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE3.BLOCK = 'BASIC'
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE3.FUSE_METHOD = 'SUM'
|
||||
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE4 = CN()
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE4.NUM_MODULES = 1
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE4.NUM_BRANCHES = 4
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE4.NUM_BLOCKS = [4, 4, 4, 4]
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE4.NUM_CHANNELS = [32, 64, 128, 256]
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE4.BLOCK = 'BASIC'
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE4.FUSE_METHOD = 'SUM'
|
||||
|
||||
|
||||
MODEL_EXTRAS = {
|
||||
'pose_resnet': POSE_RESNET,
|
||||
'pose_high_resolution_net': POSE_HIGH_RESOLUTION_NET,
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import numpy as np
|
||||
|
||||
from core.inference import get_max_preds
|
||||
|
||||
|
||||
def calc_dists(preds, target, normalize):
|
||||
preds = preds.astype(np.float32)
|
||||
target = target.astype(np.float32)
|
||||
dists = np.zeros((preds.shape[1], preds.shape[0]))
|
||||
for n in range(preds.shape[0]):
|
||||
for c in range(preds.shape[1]):
|
||||
if target[n, c, 0] > 1 and target[n, c, 1] > 1:
|
||||
normed_preds = preds[n, c, :] / normalize[n]
|
||||
normed_targets = target[n, c, :] / normalize[n]
|
||||
dists[c, n] = np.linalg.norm(normed_preds - normed_targets)
|
||||
else:
|
||||
dists[c, n] = -1
|
||||
return dists
|
||||
|
||||
|
||||
def dist_acc(dists, thr=0.5):
|
||||
''' Return percentage below threshold while ignoring values with a -1 '''
|
||||
dist_cal = np.not_equal(dists, -1)
|
||||
num_dist_cal = dist_cal.sum()
|
||||
if num_dist_cal > 0:
|
||||
return np.less(dists[dist_cal], thr).sum() * 1.0 / num_dist_cal
|
||||
else:
|
||||
return -1
|
||||
|
||||
|
||||
def accuracy(output, target, hm_type='gaussian', thr=0.5):
|
||||
'''
|
||||
Calculate accuracy according to PCK,
|
||||
but uses ground truth heatmap rather than x,y locations
|
||||
First value to be returned is average accuracy across 'idxs',
|
||||
followed by individual accuracies
|
||||
'''
|
||||
idx = list(range(output.shape[1]))
|
||||
norm = 1.0
|
||||
if hm_type == 'gaussian':
|
||||
pred, _ = get_max_preds(output)
|
||||
target, _ = get_max_preds(target)
|
||||
h = output.shape[2]
|
||||
w = output.shape[3]
|
||||
norm = np.ones((pred.shape[0], 2)) * np.array([h, w]) / 10
|
||||
dists = calc_dists(pred, target, norm)
|
||||
|
||||
acc = np.zeros((len(idx) + 1))
|
||||
avg_acc = 0
|
||||
cnt = 0
|
||||
|
||||
for i in range(len(idx)):
|
||||
acc[i + 1] = dist_acc(dists[idx[i]])
|
||||
if acc[i + 1] >= 0:
|
||||
avg_acc = avg_acc + acc[i + 1]
|
||||
cnt += 1
|
||||
|
||||
avg_acc = avg_acc / cnt if cnt != 0 else 0
|
||||
if cnt != 0:
|
||||
acc[0] = avg_acc
|
||||
return acc, avg_acc, cnt, pred
|
||||
|
||||
|
||||
@@ -0,0 +1,279 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import time
|
||||
import logging
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from core.evaluate import accuracy
|
||||
from core.inference import get_final_preds
|
||||
from utils.transforms import flip_back
|
||||
from utils.vis import save_debug_images
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def train(config, train_loader, model, criterion, optimizer, epoch,
|
||||
output_dir, tb_log_dir, writer_dict):
|
||||
batch_time = AverageMeter()
|
||||
data_time = AverageMeter()
|
||||
losses = AverageMeter()
|
||||
acc = AverageMeter()
|
||||
|
||||
# switch to train mode
|
||||
model.train()
|
||||
|
||||
end = time.time()
|
||||
for i, (input, target, target_weight, meta) in enumerate(train_loader):
|
||||
# measure data loading time
|
||||
data_time.update(time.time() - end)
|
||||
|
||||
# compute output
|
||||
outputs = model(input)
|
||||
|
||||
target = target.cuda(non_blocking=True)
|
||||
target_weight = target_weight.cuda(non_blocking=True)
|
||||
|
||||
if isinstance(outputs, list):
|
||||
loss = criterion(outputs[0], target, target_weight)
|
||||
for output in outputs[1:]:
|
||||
loss += criterion(output, target, target_weight)
|
||||
else:
|
||||
output = outputs
|
||||
loss = criterion(output, target, target_weight)
|
||||
|
||||
# loss = criterion(output, target, target_weight)
|
||||
|
||||
# compute gradient and do update step
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
# measure accuracy and record loss
|
||||
losses.update(loss.item(), input.size(0))
|
||||
|
||||
_, avg_acc, cnt, pred = accuracy(output.detach().cpu().numpy(),
|
||||
target.detach().cpu().numpy())
|
||||
acc.update(avg_acc, cnt)
|
||||
|
||||
# measure elapsed time
|
||||
batch_time.update(time.time() - end)
|
||||
end = time.time()
|
||||
|
||||
if i % config.PRINT_FREQ == 0:
|
||||
msg = 'Epoch: [{0}][{1}/{2}]\t' \
|
||||
'Time {batch_time.val:.3f}s ({batch_time.avg:.3f}s)\t' \
|
||||
'Speed {speed:.1f} samples/s\t' \
|
||||
'Data {data_time.val:.3f}s ({data_time.avg:.3f}s)\t' \
|
||||
'Loss {loss.val:.5f} ({loss.avg:.5f})\t' \
|
||||
'Accuracy {acc.val:.3f} ({acc.avg:.3f})'.format(
|
||||
epoch, i, len(train_loader), batch_time=batch_time,
|
||||
speed=input.size(0)/batch_time.val,
|
||||
data_time=data_time, loss=losses, acc=acc)
|
||||
logger.info(msg)
|
||||
|
||||
writer = writer_dict['writer']
|
||||
global_steps = writer_dict['train_global_steps']
|
||||
writer.add_scalar('train_loss', losses.val, global_steps)
|
||||
writer.add_scalar('train_acc', acc.val, global_steps)
|
||||
writer_dict['train_global_steps'] = global_steps + 1
|
||||
|
||||
prefix = '{}_{}'.format(os.path.join(output_dir, 'train'), i)
|
||||
save_debug_images(config, input, meta, target, pred*4, output,
|
||||
prefix)
|
||||
|
||||
|
||||
def validate(config, val_loader, val_dataset, model, criterion, output_dir,
|
||||
tb_log_dir, writer_dict=None):
|
||||
batch_time = AverageMeter()
|
||||
losses = AverageMeter()
|
||||
acc = AverageMeter()
|
||||
|
||||
# switch to evaluate mode
|
||||
model.eval()
|
||||
|
||||
num_samples = len(val_dataset)
|
||||
all_preds = np.zeros(
|
||||
(num_samples, config.MODEL.NUM_JOINTS, 3),
|
||||
dtype=np.float32
|
||||
)
|
||||
all_boxes = np.zeros((num_samples, 6))
|
||||
image_path = []
|
||||
filenames = []
|
||||
imgnums = []
|
||||
idx = 0
|
||||
with torch.no_grad():
|
||||
end = time.time()
|
||||
for i, (input, target, target_weight, meta) in enumerate(val_loader):
|
||||
# compute output
|
||||
outputs = model(input)
|
||||
if isinstance(outputs, list):
|
||||
output = outputs[-1]
|
||||
else:
|
||||
output = outputs
|
||||
|
||||
if config.TEST.FLIP_TEST:
|
||||
# this part is ugly, because pytorch has not supported negative index
|
||||
# input_flipped = model(input[:, :, :, ::-1])
|
||||
input_flipped = np.flip(input.cpu().numpy(), 3).copy()
|
||||
input_flipped = torch.from_numpy(input_flipped).cuda()
|
||||
outputs_flipped = model(input_flipped)
|
||||
|
||||
if isinstance(outputs_flipped, list):
|
||||
output_flipped = outputs_flipped[-1]
|
||||
else:
|
||||
output_flipped = outputs_flipped
|
||||
|
||||
output_flipped = flip_back(output_flipped.cpu().numpy(),
|
||||
val_dataset.flip_pairs)
|
||||
output_flipped = torch.from_numpy(output_flipped.copy()).cuda()
|
||||
|
||||
|
||||
# feature is not aligned, shift flipped heatmap for higher accuracy
|
||||
if config.TEST.SHIFT_HEATMAP:
|
||||
output_flipped[:, :, :, 1:] = \
|
||||
output_flipped.clone()[:, :, :, 0:-1]
|
||||
|
||||
output = (output + output_flipped) * 0.5
|
||||
|
||||
target = target.cuda(non_blocking=True)
|
||||
target_weight = target_weight.cuda(non_blocking=True)
|
||||
|
||||
loss = criterion(output, target, target_weight)
|
||||
|
||||
num_images = input.size(0)
|
||||
# measure accuracy and record loss
|
||||
losses.update(loss.item(), num_images)
|
||||
_, avg_acc, cnt, pred = accuracy(output.cpu().numpy(),
|
||||
target.cpu().numpy())
|
||||
|
||||
acc.update(avg_acc, cnt)
|
||||
|
||||
# measure elapsed time
|
||||
batch_time.update(time.time() - end)
|
||||
end = time.time()
|
||||
|
||||
c = meta['center'].numpy()
|
||||
s = meta['scale'].numpy()
|
||||
score = meta['score'].numpy()
|
||||
|
||||
preds, maxvals = get_final_preds(
|
||||
config, output.clone().cpu().numpy(), c, s)
|
||||
|
||||
all_preds[idx:idx + num_images, :, 0:2] = preds[:, :, 0:2]
|
||||
all_preds[idx:idx + num_images, :, 2:3] = maxvals
|
||||
# double check this all_boxes parts
|
||||
all_boxes[idx:idx + num_images, 0:2] = c[:, 0:2]
|
||||
all_boxes[idx:idx + num_images, 2:4] = s[:, 0:2]
|
||||
all_boxes[idx:idx + num_images, 4] = np.prod(s*200, 1)
|
||||
all_boxes[idx:idx + num_images, 5] = score
|
||||
image_path.extend(meta['image'])
|
||||
|
||||
idx += num_images
|
||||
|
||||
if i % config.PRINT_FREQ == 0:
|
||||
msg = 'Test: [{0}/{1}]\t' \
|
||||
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' \
|
||||
'Loss {loss.val:.4f} ({loss.avg:.4f})\t' \
|
||||
'Accuracy {acc.val:.3f} ({acc.avg:.3f})'.format(
|
||||
i, len(val_loader), batch_time=batch_time,
|
||||
loss=losses, acc=acc)
|
||||
logger.info(msg)
|
||||
|
||||
prefix = '{}_{}'.format(
|
||||
os.path.join(output_dir, 'val'), i
|
||||
)
|
||||
save_debug_images(config, input, meta, target, pred*4, output,
|
||||
prefix)
|
||||
|
||||
name_values, perf_indicator = val_dataset.evaluate(
|
||||
config, all_preds, output_dir, all_boxes, image_path,
|
||||
filenames, imgnums
|
||||
)
|
||||
|
||||
model_name = config.MODEL.NAME
|
||||
if isinstance(name_values, list):
|
||||
for name_value in name_values:
|
||||
_print_name_value(name_value, model_name)
|
||||
else:
|
||||
_print_name_value(name_values, model_name)
|
||||
|
||||
if writer_dict:
|
||||
writer = writer_dict['writer']
|
||||
global_steps = writer_dict['valid_global_steps']
|
||||
writer.add_scalar(
|
||||
'valid_loss',
|
||||
losses.avg,
|
||||
global_steps
|
||||
)
|
||||
writer.add_scalar(
|
||||
'valid_acc',
|
||||
acc.avg,
|
||||
global_steps
|
||||
)
|
||||
if isinstance(name_values, list):
|
||||
for name_value in name_values:
|
||||
writer.add_scalars(
|
||||
'valid',
|
||||
dict(name_value),
|
||||
global_steps
|
||||
)
|
||||
else:
|
||||
writer.add_scalars(
|
||||
'valid',
|
||||
dict(name_values),
|
||||
global_steps
|
||||
)
|
||||
writer_dict['valid_global_steps'] = global_steps + 1
|
||||
|
||||
return perf_indicator
|
||||
|
||||
|
||||
# markdown format output
|
||||
def _print_name_value(name_value, full_arch_name):
|
||||
names = name_value.keys()
|
||||
values = name_value.values()
|
||||
num_values = len(name_value)
|
||||
logger.info(
|
||||
'| Arch ' +
|
||||
' '.join(['| {}'.format(name) for name in names]) +
|
||||
' |'
|
||||
)
|
||||
logger.info('|---' * (num_values+1) + '|')
|
||||
|
||||
if len(full_arch_name) > 15:
|
||||
full_arch_name = full_arch_name[:8] + '...'
|
||||
logger.info(
|
||||
'| ' + full_arch_name + ' ' +
|
||||
' '.join(['| {:.3f}'.format(value) for value in values]) +
|
||||
' |'
|
||||
)
|
||||
|
||||
|
||||
class AverageMeter(object):
|
||||
"""Computes and stores the average and current value"""
|
||||
def __init__(self):
|
||||
self.reset()
|
||||
|
||||
def reset(self):
|
||||
self.val = 0
|
||||
self.avg = 0
|
||||
self.sum = 0
|
||||
self.count = 0
|
||||
|
||||
def update(self, val, n=1):
|
||||
self.val = val
|
||||
self.sum += val * n
|
||||
self.count += n
|
||||
self.avg = self.sum / self.count if self.count != 0 else 0
|
||||
@@ -0,0 +1,79 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
|
||||
from keypoints.lib.utils.transforms import transform_preds
|
||||
|
||||
|
||||
def get_max_preds(batch_heatmaps):
|
||||
'''
|
||||
get predictions from score maps
|
||||
heatmaps: numpy.ndarray([batch_size, num_joints, height, width])
|
||||
'''
|
||||
assert isinstance(batch_heatmaps, np.ndarray), \
|
||||
'batch_heatmaps should be numpy.ndarray'
|
||||
assert batch_heatmaps.ndim == 4, 'batch_images should be 4-ndim'
|
||||
|
||||
batch_size = batch_heatmaps.shape[0]
|
||||
num_joints = batch_heatmaps.shape[1]
|
||||
width = batch_heatmaps.shape[3]
|
||||
heatmaps_reshaped = batch_heatmaps.reshape((batch_size, num_joints, -1))
|
||||
idx = np.argmax(heatmaps_reshaped, 2)
|
||||
maxvals = np.amax(heatmaps_reshaped, 2)
|
||||
|
||||
maxvals = maxvals.reshape((batch_size, num_joints, 1))
|
||||
idx = idx.reshape((batch_size, num_joints, 1))
|
||||
|
||||
preds = np.tile(idx, (1, 1, 2)).astype(np.float32)
|
||||
|
||||
preds[:, :, 0] = (preds[:, :, 0]) % width
|
||||
preds[:, :, 1] = np.floor((preds[:, :, 1]) / width)
|
||||
|
||||
pred_mask = np.tile(np.greater(maxvals, 0.0), (1, 1, 2))
|
||||
pred_mask = pred_mask.astype(np.float32)
|
||||
|
||||
preds *= pred_mask
|
||||
return preds, maxvals
|
||||
|
||||
|
||||
def get_final_preds(config, batch_heatmaps, center, scale):
|
||||
coords, maxvals = get_max_preds(batch_heatmaps)
|
||||
|
||||
heatmap_height = batch_heatmaps.shape[2]
|
||||
heatmap_width = batch_heatmaps.shape[3]
|
||||
|
||||
# post-processing
|
||||
if config.TEST.POST_PROCESS:
|
||||
for n in range(coords.shape[0]):
|
||||
for p in range(coords.shape[1]):
|
||||
hm = batch_heatmaps[n][p]
|
||||
px = int(math.floor(coords[n][p][0] + 0.5))
|
||||
py = int(math.floor(coords[n][p][1] + 0.5))
|
||||
if 1 < px < heatmap_width-1 and 1 < py < heatmap_height-1:
|
||||
diff = np.array(
|
||||
[
|
||||
hm[py][px+1] - hm[py][px-1],
|
||||
hm[py+1][px]-hm[py-1][px]
|
||||
]
|
||||
)
|
||||
coords[n][p] += np.sign(diff) * .25
|
||||
|
||||
preds = coords.copy()
|
||||
|
||||
# Transform back
|
||||
for i in range(coords.shape[0]):
|
||||
preds[i] = transform_preds(
|
||||
coords[i], center[i], scale[i], [heatmap_width, heatmap_height]
|
||||
)
|
||||
|
||||
return preds, maxvals
|
||||
@@ -0,0 +1,84 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class JointsMSELoss(nn.Module):
|
||||
def __init__(self, use_target_weight):
|
||||
super(JointsMSELoss, self).__init__()
|
||||
self.criterion = nn.MSELoss(reduction='mean')
|
||||
self.use_target_weight = use_target_weight
|
||||
|
||||
def forward(self, output, target, target_weight):
|
||||
batch_size = output.size(0)
|
||||
num_joints = output.size(1)
|
||||
heatmaps_pred = output.reshape((batch_size, num_joints, -1)).split(1, 1)
|
||||
heatmaps_gt = target.reshape((batch_size, num_joints, -1)).split(1, 1)
|
||||
loss = 0
|
||||
|
||||
for idx in range(num_joints):
|
||||
heatmap_pred = heatmaps_pred[idx].squeeze()
|
||||
heatmap_gt = heatmaps_gt[idx].squeeze()
|
||||
if self.use_target_weight:
|
||||
loss += 0.5 * self.criterion(
|
||||
heatmap_pred.mul(target_weight[:, idx]),
|
||||
heatmap_gt.mul(target_weight[:, idx])
|
||||
)
|
||||
else:
|
||||
loss += 0.5 * self.criterion(heatmap_pred, heatmap_gt)
|
||||
|
||||
return loss / num_joints
|
||||
|
||||
|
||||
class JointsOHKMMSELoss(nn.Module):
|
||||
def __init__(self, use_target_weight, topk=8):
|
||||
super(JointsOHKMMSELoss, self).__init__()
|
||||
self.criterion = nn.MSELoss(reduction='none')
|
||||
self.use_target_weight = use_target_weight
|
||||
self.topk = topk
|
||||
|
||||
def ohkm(self, loss):
|
||||
ohkm_loss = 0.
|
||||
for i in range(loss.size()[0]):
|
||||
sub_loss = loss[i]
|
||||
topk_val, topk_idx = torch.topk(
|
||||
sub_loss, k=self.topk, dim=0, sorted=False
|
||||
)
|
||||
tmp_loss = torch.gather(sub_loss, 0, topk_idx)
|
||||
ohkm_loss += torch.sum(tmp_loss) / self.topk
|
||||
ohkm_loss /= loss.size()[0]
|
||||
return ohkm_loss
|
||||
|
||||
def forward(self, output, target, target_weight):
|
||||
batch_size = output.size(0)
|
||||
num_joints = output.size(1)
|
||||
heatmaps_pred = output.reshape((batch_size, num_joints, -1)).split(1, 1)
|
||||
heatmaps_gt = target.reshape((batch_size, num_joints, -1)).split(1, 1)
|
||||
|
||||
loss = []
|
||||
for idx in range(num_joints):
|
||||
heatmap_pred = heatmaps_pred[idx].squeeze()
|
||||
heatmap_gt = heatmaps_gt[idx].squeeze()
|
||||
if self.use_target_weight:
|
||||
loss.append(0.5 * self.criterion(
|
||||
heatmap_pred.mul(target_weight[:, idx]),
|
||||
heatmap_gt.mul(target_weight[:, idx])
|
||||
))
|
||||
else:
|
||||
loss.append(
|
||||
0.5 * self.criterion(heatmap_pred, heatmap_gt)
|
||||
)
|
||||
|
||||
loss = [l.mean(dim=1).unsqueeze(dim=1) for l in loss]
|
||||
loss = torch.cat(loss, dim=1)
|
||||
|
||||
return self.ohkm(loss)
|
||||
@@ -0,0 +1,292 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import copy
|
||||
import logging
|
||||
import random
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
from utils.transforms import get_affine_transform
|
||||
from utils.transforms import affine_transform
|
||||
from utils.transforms import fliplr_joints
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class JointsDataset(Dataset):
|
||||
def __init__(self, cfg, root, image_set, is_train, transform=None):
|
||||
self.num_joints = 0
|
||||
self.pixel_std = 200
|
||||
self.flip_pairs = []
|
||||
self.parent_ids = []
|
||||
|
||||
self.is_train = is_train
|
||||
self.root = root
|
||||
self.image_set = image_set
|
||||
|
||||
self.output_path = cfg.OUTPUT_DIR
|
||||
self.data_format = cfg.DATASET.DATA_FORMAT
|
||||
|
||||
self.scale_factor = cfg.DATASET.SCALE_FACTOR
|
||||
self.rotation_factor = cfg.DATASET.ROT_FACTOR
|
||||
self.flip = cfg.DATASET.FLIP
|
||||
self.num_joints_half_body = cfg.DATASET.NUM_JOINTS_HALF_BODY
|
||||
self.prob_half_body = cfg.DATASET.PROB_HALF_BODY
|
||||
self.color_rgb = cfg.DATASET.COLOR_RGB
|
||||
|
||||
self.target_type = cfg.MODEL.TARGET_TYPE
|
||||
self.image_size = np.array(cfg.MODEL.IMAGE_SIZE)
|
||||
self.heatmap_size = np.array(cfg.MODEL.HEATMAP_SIZE)
|
||||
self.sigma = cfg.MODEL.SIGMA
|
||||
self.use_different_joints_weight = cfg.LOSS.USE_DIFFERENT_JOINTS_WEIGHT
|
||||
self.joints_weight = 1
|
||||
|
||||
self.transform = transform
|
||||
self.db = []
|
||||
|
||||
def _get_db(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def evaluate(self, cfg, preds, output_dir, *args, **kwargs):
|
||||
raise NotImplementedError
|
||||
|
||||
def half_body_transform(self, joints, joints_vis):
|
||||
upper_joints = []
|
||||
lower_joints = []
|
||||
for joint_id in range(self.num_joints):
|
||||
if joints_vis[joint_id][0] > 0:
|
||||
if joint_id in self.upper_body_ids:
|
||||
upper_joints.append(joints[joint_id])
|
||||
else:
|
||||
lower_joints.append(joints[joint_id])
|
||||
|
||||
if np.random.randn() < 0.5 and len(upper_joints) > 2:
|
||||
selected_joints = upper_joints
|
||||
else:
|
||||
selected_joints = lower_joints \
|
||||
if len(lower_joints) > 2 else upper_joints
|
||||
|
||||
if len(selected_joints) < 2:
|
||||
return None, None
|
||||
|
||||
selected_joints = np.array(selected_joints, dtype=np.float32)
|
||||
center = selected_joints.mean(axis=0)[:2]
|
||||
|
||||
left_top = np.amin(selected_joints, axis=0)
|
||||
right_bottom = np.amax(selected_joints, axis=0)
|
||||
|
||||
w = right_bottom[0] - left_top[0]
|
||||
h = right_bottom[1] - left_top[1]
|
||||
|
||||
if w > self.aspect_ratio * h:
|
||||
h = w * 1.0 / self.aspect_ratio
|
||||
elif w < self.aspect_ratio * h:
|
||||
w = h * self.aspect_ratio
|
||||
|
||||
scale = np.array(
|
||||
[
|
||||
w * 1.0 / self.pixel_std,
|
||||
h * 1.0 / self.pixel_std
|
||||
],
|
||||
dtype=np.float32
|
||||
)
|
||||
|
||||
scale = scale * 1.5
|
||||
|
||||
return center, scale
|
||||
|
||||
def __len__(self,):
|
||||
return len(self.db)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
db_rec = copy.deepcopy(self.db[idx])
|
||||
|
||||
image_file = db_rec['image']
|
||||
filename = db_rec['filename'] if 'filename' in db_rec else ''
|
||||
imgnum = db_rec['imgnum'] if 'imgnum' in db_rec else ''
|
||||
|
||||
if self.data_format == 'zip':
|
||||
from utils import zipreader
|
||||
data_numpy = zipreader.imread(
|
||||
image_file, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION
|
||||
)
|
||||
else:
|
||||
data_numpy = cv2.imread(
|
||||
image_file, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION
|
||||
)
|
||||
|
||||
if self.color_rgb:
|
||||
data_numpy = cv2.cvtColor(data_numpy, cv2.COLOR_BGR2RGB)
|
||||
|
||||
if data_numpy is None:
|
||||
logger.error('=> fail to read {}'.format(image_file))
|
||||
raise ValueError('Fail to read {}'.format(image_file))
|
||||
|
||||
joints = db_rec['joints_3d']
|
||||
joints_vis = db_rec['joints_3d_vis']
|
||||
|
||||
c = db_rec['center']
|
||||
s = db_rec['scale']
|
||||
score = db_rec['score'] if 'score' in db_rec else 1
|
||||
r = 0
|
||||
|
||||
if self.is_train:
|
||||
if (np.sum(joints_vis[:, 0]) > self.num_joints_half_body
|
||||
and np.random.rand() < self.prob_half_body):
|
||||
c_half_body, s_half_body = self.half_body_transform(
|
||||
joints, joints_vis
|
||||
)
|
||||
|
||||
if c_half_body is not None and s_half_body is not None:
|
||||
c, s = c_half_body, s_half_body
|
||||
|
||||
sf = self.scale_factor
|
||||
rf = self.rotation_factor
|
||||
s = s * np.clip(np.random.randn()*sf + 1, 1 - sf, 1 + sf)
|
||||
r = np.clip(np.random.randn()*rf, -rf*2, rf*2) \
|
||||
if random.random() <= 0.6 else 0
|
||||
|
||||
if self.flip and random.random() <= 0.5:
|
||||
data_numpy = data_numpy[:, ::-1, :]
|
||||
joints, joints_vis = fliplr_joints(
|
||||
joints, joints_vis, data_numpy.shape[1], self.flip_pairs)
|
||||
c[0] = data_numpy.shape[1] - c[0] - 1
|
||||
|
||||
trans = get_affine_transform(c, s, r, self.image_size)
|
||||
input = cv2.warpAffine(
|
||||
data_numpy,
|
||||
trans,
|
||||
(int(self.image_size[0]), int(self.image_size[1])),
|
||||
flags=cv2.INTER_LINEAR)
|
||||
|
||||
cv2.imshow('input', input)
|
||||
cv2.waitKey()
|
||||
|
||||
if self.transform:
|
||||
input = self.transform(input)
|
||||
|
||||
for i in range(self.num_joints):
|
||||
if joints_vis[i, 0] > 0.0:
|
||||
joints[i, 0:2] = affine_transform(joints[i, 0:2], trans)
|
||||
|
||||
target, target_weight = self.generate_target(joints, joints_vis)
|
||||
|
||||
target = torch.from_numpy(target)
|
||||
target_weight = torch.from_numpy(target_weight)
|
||||
|
||||
meta = {
|
||||
'image': image_file,
|
||||
'filename': filename,
|
||||
'imgnum': imgnum,
|
||||
'joints': joints,
|
||||
'joints_vis': joints_vis,
|
||||
'center': c,
|
||||
'scale': s,
|
||||
'rotation': r,
|
||||
'score': score
|
||||
}
|
||||
|
||||
return input, target, target_weight, meta
|
||||
|
||||
def select_data(self, db):
|
||||
db_selected = []
|
||||
for rec in db:
|
||||
num_vis = 0
|
||||
joints_x = 0.0
|
||||
joints_y = 0.0
|
||||
for joint, joint_vis in zip(
|
||||
rec['joints_3d'], rec['joints_3d_vis']):
|
||||
if joint_vis[0] <= 0:
|
||||
continue
|
||||
num_vis += 1
|
||||
|
||||
joints_x += joint[0]
|
||||
joints_y += joint[1]
|
||||
if num_vis == 0:
|
||||
continue
|
||||
|
||||
joints_x, joints_y = joints_x / num_vis, joints_y / num_vis
|
||||
|
||||
area = rec['scale'][0] * rec['scale'][1] * (self.pixel_std**2)
|
||||
joints_center = np.array([joints_x, joints_y])
|
||||
bbox_center = np.array(rec['center'])
|
||||
diff_norm2 = np.linalg.norm((joints_center-bbox_center), 2)
|
||||
ks = np.exp(-1.0*(diff_norm2**2) / ((0.2)**2*2.0*area))
|
||||
|
||||
metric = (0.2 / 16) * num_vis + 0.45 - 0.2 / 16
|
||||
if ks > metric:
|
||||
db_selected.append(rec)
|
||||
|
||||
logger.info('=> num db: {}'.format(len(db)))
|
||||
logger.info('=> num selected db: {}'.format(len(db_selected)))
|
||||
return db_selected
|
||||
|
||||
def generate_target(self, joints, joints_vis):
|
||||
'''
|
||||
:param joints: [num_joints, 3]
|
||||
:param joints_vis: [num_joints, 3]
|
||||
:return: target, target_weight(1: visible, 0: invisible)
|
||||
'''
|
||||
target_weight = np.ones((self.num_joints, 1), dtype=np.float32)
|
||||
target_weight[:, 0] = joints_vis[:, 0]
|
||||
|
||||
assert self.target_type == 'gaussian', \
|
||||
'Only support gaussian map now!'
|
||||
|
||||
if self.target_type == 'gaussian':
|
||||
target = np.zeros((self.num_joints,
|
||||
self.heatmap_size[1],
|
||||
self.heatmap_size[0]),
|
||||
dtype=np.float32)
|
||||
|
||||
tmp_size = self.sigma * 3
|
||||
|
||||
for joint_id in range(self.num_joints):
|
||||
feat_stride = self.image_size / self.heatmap_size
|
||||
mu_x = int(joints[joint_id][0] / feat_stride[0] + 0.5)
|
||||
mu_y = int(joints[joint_id][1] / feat_stride[1] + 0.5)
|
||||
# Check that any part of the gaussian is in-bounds
|
||||
ul = [int(mu_x - tmp_size), int(mu_y - tmp_size)]
|
||||
br = [int(mu_x + tmp_size + 1), int(mu_y + tmp_size + 1)]
|
||||
if ul[0] >= self.heatmap_size[0] or ul[1] >= self.heatmap_size[1] \
|
||||
or br[0] < 0 or br[1] < 0:
|
||||
# If not, just return the image as is
|
||||
target_weight[joint_id] = 0
|
||||
continue
|
||||
|
||||
# # Generate gaussian
|
||||
size = 2 * tmp_size + 1
|
||||
x = np.arange(0, size, 1, np.float32)
|
||||
y = x[:, np.newaxis]
|
||||
x0 = y0 = size // 2
|
||||
# The gaussian is not normalized, we want the center value to equal 1
|
||||
g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * self.sigma ** 2))
|
||||
|
||||
# Usable gaussian range
|
||||
g_x = max(0, -ul[0]), min(br[0], self.heatmap_size[0]) - ul[0]
|
||||
g_y = max(0, -ul[1]), min(br[1], self.heatmap_size[1]) - ul[1]
|
||||
# Image range
|
||||
img_x = max(0, ul[0]), min(br[0], self.heatmap_size[0])
|
||||
img_y = max(0, ul[1]), min(br[1], self.heatmap_size[1])
|
||||
|
||||
v = target_weight[joint_id]
|
||||
if v > 0.5:
|
||||
target[joint_id][img_y[0]:img_y[1], img_x[0]:img_x[1]] = \
|
||||
g[g_y[0]:g_y[1], g_x[0]:g_x[1]]
|
||||
|
||||
if self.use_different_joints_weight:
|
||||
target_weight = np.multiply(target_weight, self.joints_weight)
|
||||
|
||||
return target, target_weight
|
||||
@@ -0,0 +1,12 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
from .mpii import MPIIDataset as mpii
|
||||
from .coco import COCODataset as coco
|
||||
@@ -0,0 +1,445 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
from collections import defaultdict
|
||||
from collections import OrderedDict
|
||||
import logging
|
||||
import os
|
||||
|
||||
from pycocotools.coco import COCO
|
||||
from pycocotools.cocoeval import COCOeval
|
||||
import json_tricks as json
|
||||
import numpy as np
|
||||
|
||||
from dataset.JointsDataset import JointsDataset
|
||||
from nms.nms import oks_nms
|
||||
from nms.nms import soft_oks_nms
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class COCODataset(JointsDataset):
|
||||
'''
|
||||
"keypoints": {
|
||||
0: "nose",
|
||||
1: "left_eye",
|
||||
2: "right_eye",
|
||||
3: "left_ear",
|
||||
4: "right_ear",
|
||||
5: "left_shoulder",
|
||||
6: "right_shoulder",
|
||||
7: "left_elbow",
|
||||
8: "right_elbow",
|
||||
9: "left_wrist",
|
||||
10: "right_wrist",
|
||||
11: "left_hip",
|
||||
12: "right_hip",
|
||||
13: "left_knee",
|
||||
14: "right_knee",
|
||||
15: "left_ankle",
|
||||
16: "right_ankle"
|
||||
},
|
||||
"skeleton": [
|
||||
[16,14],[14,12],[17,15],[15,13],[12,13],[6,12],[7,13], [6,7],[6,8],
|
||||
[7,9],[8,10],[9,11],[2,3],[1,2],[1,3],[2,4],[3,5],[4,6],[5,7]]
|
||||
'''
|
||||
def __init__(self, cfg, root, image_set, is_train, transform=None):
|
||||
super().__init__(cfg, root, image_set, is_train, transform)
|
||||
self.nms_thre = cfg.TEST.NMS_THRE
|
||||
self.image_thre = cfg.TEST.IMAGE_THRE
|
||||
self.soft_nms = cfg.TEST.SOFT_NMS
|
||||
self.oks_thre = cfg.TEST.OKS_THRE
|
||||
self.in_vis_thre = cfg.TEST.IN_VIS_THRE
|
||||
self.bbox_file = cfg.TEST.COCO_BBOX_FILE
|
||||
self.use_gt_bbox = cfg.TEST.USE_GT_BBOX
|
||||
self.image_width = cfg.MODEL.IMAGE_SIZE[0]
|
||||
self.image_height = cfg.MODEL.IMAGE_SIZE[1]
|
||||
self.aspect_ratio = self.image_width * 1.0 / self.image_height
|
||||
self.pixel_std = 200
|
||||
|
||||
self.coco = COCO(self._get_ann_file_keypoint())
|
||||
|
||||
# deal with class names
|
||||
cats = [cat['name']
|
||||
for cat in self.coco.loadCats(self.coco.getCatIds())]
|
||||
self.classes = ['__background__'] + cats
|
||||
logger.info('=> classes: {}'.format(self.classes))
|
||||
self.num_classes = len(self.classes)
|
||||
self._class_to_ind = dict(zip(self.classes, range(self.num_classes)))
|
||||
self._class_to_coco_ind = dict(zip(cats, self.coco.getCatIds()))
|
||||
self._coco_ind_to_class_ind = dict(
|
||||
[
|
||||
(self._class_to_coco_ind[cls], self._class_to_ind[cls])
|
||||
for cls in self.classes[1:]
|
||||
]
|
||||
)
|
||||
|
||||
# load image file names
|
||||
self.image_set_index = self._load_image_set_index()
|
||||
self.num_images = len(self.image_set_index)
|
||||
logger.info('=> num_images: {}'.format(self.num_images))
|
||||
|
||||
self.num_joints = 17
|
||||
self.flip_pairs = [[1, 2], [3, 4], [5, 6], [7, 8],
|
||||
[9, 10], [11, 12], [13, 14], [15, 16]]
|
||||
self.parent_ids = None
|
||||
self.upper_body_ids = (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10)
|
||||
self.lower_body_ids = (11, 12, 13, 14, 15, 16)
|
||||
|
||||
self.joints_weight = np.array(
|
||||
[
|
||||
1., 1., 1., 1., 1., 1., 1., 1.2, 1.2,
|
||||
1.5, 1.5, 1., 1., 1.2, 1.2, 1.5, 1.5
|
||||
],
|
||||
dtype=np.float32
|
||||
).reshape((self.num_joints, 1))
|
||||
|
||||
self.db = self._get_db()
|
||||
|
||||
if is_train and cfg.DATASET.SELECT_DATA:
|
||||
self.db = self.select_data(self.db)
|
||||
|
||||
logger.info('=> load {} samples'.format(len(self.db)))
|
||||
|
||||
def _get_ann_file_keypoint(self):
|
||||
""" self.root / annotations / person_keypoints_train2017.json """
|
||||
prefix = 'person_keypoints' \
|
||||
if 'test' not in self.image_set else 'image_info'
|
||||
return os.path.join(
|
||||
self.root,
|
||||
'annotations',
|
||||
prefix + '_' + self.image_set + '.json'
|
||||
)
|
||||
|
||||
def _load_image_set_index(self):
|
||||
""" image id: int """
|
||||
image_ids = self.coco.getImgIds()
|
||||
return image_ids
|
||||
|
||||
def _get_db(self):
|
||||
if self.is_train or self.use_gt_bbox:
|
||||
# use ground truth bbox
|
||||
gt_db = self._load_coco_keypoint_annotations()
|
||||
else:
|
||||
# use bbox from detection
|
||||
gt_db = self._load_coco_person_detection_results()
|
||||
return gt_db
|
||||
|
||||
def _load_coco_keypoint_annotations(self):
|
||||
""" ground truth bbox and keypoints """
|
||||
gt_db = []
|
||||
for index in self.image_set_index:
|
||||
gt_db.extend(self._load_coco_keypoint_annotation_kernal(index))
|
||||
return gt_db
|
||||
|
||||
def _load_coco_keypoint_annotation_kernal(self, index):
|
||||
"""
|
||||
coco ann: [u'segmentation', u'area', u'iscrowd', u'image_id', u'bbox', u'category_id', u'id']
|
||||
iscrowd:
|
||||
crowd instances are handled by marking their overlaps with all categories to -1
|
||||
and later excluded in training
|
||||
bbox:
|
||||
[x1, y1, w, h]
|
||||
:param index: coco image id
|
||||
:return: db entry
|
||||
"""
|
||||
im_ann = self.coco.loadImgs(index)[0]
|
||||
width = im_ann['width']
|
||||
height = im_ann['height']
|
||||
|
||||
annIds = self.coco.getAnnIds(imgIds=index, iscrowd=False)
|
||||
objs = self.coco.loadAnns(annIds)
|
||||
|
||||
# sanitize bboxes
|
||||
valid_objs = []
|
||||
for obj in objs:
|
||||
x, y, w, h = obj['bbox']
|
||||
x1 = np.max((0, x))
|
||||
y1 = np.max((0, y))
|
||||
x2 = np.min((width - 1, x1 + np.max((0, w - 1))))
|
||||
y2 = np.min((height - 1, y1 + np.max((0, h - 1))))
|
||||
if obj['area'] > 0 and x2 >= x1 and y2 >= y1:
|
||||
obj['clean_bbox'] = [x1, y1, x2-x1, y2-y1]
|
||||
valid_objs.append(obj)
|
||||
objs = valid_objs
|
||||
|
||||
rec = []
|
||||
for obj in objs:
|
||||
cls = self._coco_ind_to_class_ind[obj['category_id']]
|
||||
if cls != 1:
|
||||
continue
|
||||
|
||||
# ignore objs without keypoints annotation
|
||||
if max(obj['keypoints']) == 0:
|
||||
continue
|
||||
|
||||
joints_3d = np.zeros((self.num_joints, 3), dtype=np.float)
|
||||
joints_3d_vis = np.zeros((self.num_joints, 3), dtype=np.float)
|
||||
for ipt in range(self.num_joints):
|
||||
joints_3d[ipt, 0] = obj['keypoints'][ipt * 3 + 0]
|
||||
joints_3d[ipt, 1] = obj['keypoints'][ipt * 3 + 1]
|
||||
joints_3d[ipt, 2] = 0
|
||||
t_vis = obj['keypoints'][ipt * 3 + 2]
|
||||
if t_vis > 1:
|
||||
t_vis = 1
|
||||
joints_3d_vis[ipt, 0] = t_vis
|
||||
joints_3d_vis[ipt, 1] = t_vis
|
||||
joints_3d_vis[ipt, 2] = 0
|
||||
|
||||
center, scale = self._box2cs(obj['clean_bbox'][:4])
|
||||
rec.append({
|
||||
'image': self.image_path_from_index(index),
|
||||
'center': center,
|
||||
'scale': scale,
|
||||
'joints_3d': joints_3d,
|
||||
'joints_3d_vis': joints_3d_vis,
|
||||
'filename': '',
|
||||
'imgnum': 0,
|
||||
})
|
||||
|
||||
return rec
|
||||
|
||||
def _box2cs(self, box):
|
||||
x, y, w, h = box[:4]
|
||||
return self._xywh2cs(x, y, w, h)
|
||||
|
||||
def _xywh2cs(self, x, y, w, h):
|
||||
center = np.zeros((2), dtype=np.float32)
|
||||
center[0] = x + w * 0.5
|
||||
center[1] = y + h * 0.5
|
||||
|
||||
if w > self.aspect_ratio * h:
|
||||
h = w * 1.0 / self.aspect_ratio
|
||||
elif w < self.aspect_ratio * h:
|
||||
w = h * self.aspect_ratio
|
||||
scale = np.array(
|
||||
[w * 1.0 / self.pixel_std, h * 1.0 / self.pixel_std],
|
||||
dtype=np.float32)
|
||||
if center[0] != -1:
|
||||
scale = scale * 1.25
|
||||
|
||||
return center, scale
|
||||
|
||||
def image_path_from_index(self, index):
|
||||
""" example: images / train2017 / 000000119993.jpg """
|
||||
file_name = '%012d.jpg' % index
|
||||
if '2014' in self.image_set:
|
||||
file_name = 'COCO_%s_' % self.image_set + file_name
|
||||
|
||||
prefix = 'test2017' if 'test' in self.image_set else self.image_set
|
||||
|
||||
data_name = prefix + '.zip@' if self.data_format == 'zip' else prefix
|
||||
|
||||
image_path = os.path.join(
|
||||
self.root, 'images', data_name, file_name)
|
||||
|
||||
return image_path
|
||||
|
||||
def _load_coco_person_detection_results(self):
|
||||
all_boxes = None
|
||||
with open(self.bbox_file, 'r') as f:
|
||||
all_boxes = json.load(f)[:10]
|
||||
|
||||
if not all_boxes:
|
||||
logger.error('=> Load %s fail!' % self.bbox_file)
|
||||
return None
|
||||
|
||||
logger.info('=> Total boxes: {}'.format(len(all_boxes)))
|
||||
|
||||
kpt_db = []
|
||||
num_boxes = 0
|
||||
for n_img in range(0, len(all_boxes)):
|
||||
det_res = all_boxes[n_img]
|
||||
if det_res['category_id'] != 1:
|
||||
continue
|
||||
img_name = self.image_path_from_index(det_res['image_id'])
|
||||
box = det_res['bbox']
|
||||
score = det_res['score']
|
||||
|
||||
if score < self.image_thre:
|
||||
continue
|
||||
|
||||
num_boxes = num_boxes + 1
|
||||
|
||||
center, scale = self._box2cs(box)
|
||||
joints_3d = np.zeros((self.num_joints, 3), dtype=np.float)
|
||||
joints_3d_vis = np.ones(
|
||||
(self.num_joints, 3), dtype=np.float)
|
||||
kpt_db.append({
|
||||
'image': img_name,
|
||||
'center': center,
|
||||
'scale': scale,
|
||||
'score': score,
|
||||
'joints_3d': joints_3d,
|
||||
'joints_3d_vis': joints_3d_vis,
|
||||
})
|
||||
|
||||
logger.info('=> Total boxes after fliter low score@{}: {}'.format(
|
||||
self.image_thre, num_boxes))
|
||||
return kpt_db
|
||||
|
||||
def evaluate(self, cfg, preds, output_dir, all_boxes, img_path,
|
||||
*args, **kwargs):
|
||||
rank = cfg.RANK
|
||||
|
||||
res_folder = os.path.join(output_dir, 'results')
|
||||
if not os.path.exists(res_folder):
|
||||
try:
|
||||
os.makedirs(res_folder)
|
||||
except Exception:
|
||||
logger.error('Fail to make {}'.format(res_folder))
|
||||
|
||||
res_file = os.path.join(
|
||||
res_folder, 'keypoints_{}_results_{}.json'.format(
|
||||
self.image_set, rank)
|
||||
)
|
||||
|
||||
# person x (keypoints)
|
||||
_kpts = []
|
||||
for idx, kpt in enumerate(preds):
|
||||
_kpts.append({
|
||||
'keypoints': kpt,
|
||||
'center': all_boxes[idx][0:2],
|
||||
'scale': all_boxes[idx][2:4],
|
||||
'area': all_boxes[idx][4],
|
||||
'score': all_boxes[idx][5],
|
||||
'image': int(img_path[idx][-16:-4])
|
||||
})
|
||||
# image x person x (keypoints)
|
||||
kpts = defaultdict(list)
|
||||
for kpt in _kpts:
|
||||
kpts[kpt['image']].append(kpt)
|
||||
|
||||
# rescoring and oks nms
|
||||
num_joints = self.num_joints
|
||||
in_vis_thre = self.in_vis_thre
|
||||
oks_thre = self.oks_thre
|
||||
oks_nmsed_kpts = []
|
||||
for img in kpts.keys():
|
||||
img_kpts = kpts[img]
|
||||
for n_p in img_kpts:
|
||||
box_score = n_p['score']
|
||||
kpt_score = 0
|
||||
valid_num = 0
|
||||
for n_jt in range(0, num_joints):
|
||||
t_s = n_p['keypoints'][n_jt][2]
|
||||
if t_s > in_vis_thre:
|
||||
kpt_score = kpt_score + t_s
|
||||
valid_num = valid_num + 1
|
||||
if valid_num != 0:
|
||||
kpt_score = kpt_score / valid_num
|
||||
# rescoring
|
||||
n_p['score'] = kpt_score * box_score
|
||||
|
||||
if self.soft_nms:
|
||||
keep = soft_oks_nms(
|
||||
[img_kpts[i] for i in range(len(img_kpts))],
|
||||
oks_thre
|
||||
)
|
||||
else:
|
||||
keep = oks_nms(
|
||||
[img_kpts[i] for i in range(len(img_kpts))],
|
||||
oks_thre
|
||||
)
|
||||
|
||||
if len(keep) == 0:
|
||||
oks_nmsed_kpts.append(img_kpts)
|
||||
else:
|
||||
oks_nmsed_kpts.append([img_kpts[_keep] for _keep in keep])
|
||||
|
||||
self._write_coco_keypoint_results(
|
||||
oks_nmsed_kpts, res_file)
|
||||
if 'test' not in self.image_set:
|
||||
info_str = self._do_python_keypoint_eval(
|
||||
res_file, res_folder)
|
||||
name_value = OrderedDict(info_str)
|
||||
return name_value, name_value['AP']
|
||||
else:
|
||||
return {'Null': 0}, 0
|
||||
|
||||
def _write_coco_keypoint_results(self, keypoints, res_file):
|
||||
data_pack = [
|
||||
{
|
||||
'cat_id': self._class_to_coco_ind[cls],
|
||||
'cls_ind': cls_ind,
|
||||
'cls': cls,
|
||||
'ann_type': 'keypoints',
|
||||
'keypoints': keypoints
|
||||
}
|
||||
for cls_ind, cls in enumerate(self.classes) if not cls == '__background__'
|
||||
]
|
||||
|
||||
results = self._coco_keypoint_results_one_category_kernel(data_pack[0])
|
||||
logger.info('=> writing results json to %s' % res_file)
|
||||
with open(res_file, 'w') as f:
|
||||
json.dump(results, f, sort_keys=True, indent=4)
|
||||
try:
|
||||
json.load(open(res_file))
|
||||
except Exception:
|
||||
content = []
|
||||
with open(res_file, 'r') as f:
|
||||
for line in f:
|
||||
content.append(line)
|
||||
content[-1] = ']'
|
||||
with open(res_file, 'w') as f:
|
||||
for c in content:
|
||||
f.write(c)
|
||||
|
||||
def _coco_keypoint_results_one_category_kernel(self, data_pack):
|
||||
cat_id = data_pack['cat_id']
|
||||
keypoints = data_pack['keypoints']
|
||||
cat_results = []
|
||||
|
||||
for img_kpts in keypoints:
|
||||
if len(img_kpts) == 0:
|
||||
continue
|
||||
|
||||
_key_points = np.array([img_kpts[k]['keypoints']
|
||||
for k in range(len(img_kpts))])
|
||||
key_points = np.zeros(
|
||||
(_key_points.shape[0], self.num_joints * 3), dtype=np.float
|
||||
)
|
||||
|
||||
for ipt in range(self.num_joints):
|
||||
key_points[:, ipt * 3 + 0] = _key_points[:, ipt, 0]
|
||||
key_points[:, ipt * 3 + 1] = _key_points[:, ipt, 1]
|
||||
key_points[:, ipt * 3 + 2] = _key_points[:, ipt, 2] # keypoints score.
|
||||
|
||||
result = [
|
||||
{
|
||||
'image_id': img_kpts[k]['image'],
|
||||
'category_id': cat_id,
|
||||
'keypoints': list(key_points[k]),
|
||||
'score': img_kpts[k]['score'],
|
||||
'center': list(img_kpts[k]['center']),
|
||||
'scale': list(img_kpts[k]['scale'])
|
||||
}
|
||||
for k in range(len(img_kpts))
|
||||
]
|
||||
cat_results.extend(result)
|
||||
|
||||
return cat_results
|
||||
|
||||
def _do_python_keypoint_eval(self, res_file, res_folder):
|
||||
coco_dt = self.coco.loadRes(res_file)
|
||||
coco_eval = COCOeval(self.coco, coco_dt, 'keypoints')
|
||||
coco_eval.params.useSegm = None
|
||||
coco_eval.evaluate()
|
||||
coco_eval.accumulate()
|
||||
coco_eval.summarize()
|
||||
|
||||
stats_names = ['AP', 'Ap .5', 'AP .75', 'AP (M)', 'AP (L)', 'AR', 'AR .5', 'AR .75', 'AR (M)', 'AR (L)']
|
||||
|
||||
info_str = []
|
||||
for ind, name in enumerate(stats_names):
|
||||
info_str.append((name, coco_eval.stats[ind]))
|
||||
|
||||
return info_str
|
||||
@@ -0,0 +1,181 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import logging
|
||||
import os
|
||||
import json_tricks as json
|
||||
from collections import OrderedDict
|
||||
|
||||
import numpy as np
|
||||
from scipy.io import loadmat, savemat
|
||||
|
||||
from dataset.JointsDataset import JointsDataset
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MPIIDataset(JointsDataset):
|
||||
def __init__(self, cfg, root, image_set, is_train, transform=None):
|
||||
super().__init__(cfg, root, image_set, is_train, transform)
|
||||
|
||||
self.num_joints = 16
|
||||
self.flip_pairs = [[0, 5], [1, 4], [2, 3], [10, 15], [11, 14], [12, 13]]
|
||||
self.parent_ids = [1, 2, 6, 6, 3, 4, 6, 6, 7, 8, 11, 12, 7, 7, 13, 14]
|
||||
|
||||
self.upper_body_ids = (7, 8, 9, 10, 11, 12, 13, 14, 15)
|
||||
self.lower_body_ids = (0, 1, 2, 3, 4, 5, 6)
|
||||
|
||||
self.db = self._get_db()
|
||||
|
||||
if is_train and cfg.DATASET.SELECT_DATA:
|
||||
self.db = self.select_data(self.db)
|
||||
|
||||
logger.info('=> load {} samples'.format(len(self.db)))
|
||||
|
||||
def _get_db(self):
|
||||
# create train/val split
|
||||
file_name = os.path.join(
|
||||
self.root, 'annot', self.image_set+'.json'
|
||||
)
|
||||
with open(file_name) as anno_file:
|
||||
anno = json.load(anno_file)
|
||||
|
||||
gt_db = []
|
||||
for a in anno:
|
||||
image_name = a['image']
|
||||
|
||||
c = np.array(a['center'], dtype=np.float)
|
||||
s = np.array([a['scale'], a['scale']], dtype=np.float)
|
||||
|
||||
# Adjust center/scale slightly to avoid cropping limbs
|
||||
if c[0] != -1:
|
||||
c[1] = c[1] + 15 * s[1]
|
||||
s = s * 1.25
|
||||
|
||||
# MPII uses matlab format, index is based 1,
|
||||
# we should first convert to 0-based index
|
||||
c = c - 1
|
||||
|
||||
joints_3d = np.zeros((self.num_joints, 3), dtype=np.float)
|
||||
joints_3d_vis = np.zeros((self.num_joints, 3), dtype=np.float)
|
||||
if self.image_set != 'test':
|
||||
joints = np.array(a['joints'])
|
||||
joints[:, 0:2] = joints[:, 0:2] - 1
|
||||
joints_vis = np.array(a['joints_vis'])
|
||||
assert len(joints) == self.num_joints, \
|
||||
'joint num diff: {} vs {}'.format(len(joints),
|
||||
self.num_joints)
|
||||
|
||||
joints_3d[:, 0:2] = joints[:, 0:2]
|
||||
joints_3d_vis[:, 0] = joints_vis[:]
|
||||
joints_3d_vis[:, 1] = joints_vis[:]
|
||||
|
||||
image_dir = 'images.zip@' if self.data_format == 'zip' else 'images'
|
||||
gt_db.append(
|
||||
{
|
||||
'image': os.path.join(self.root, image_dir, image_name),
|
||||
'center': c,
|
||||
'scale': s,
|
||||
'joints_3d': joints_3d,
|
||||
'joints_3d_vis': joints_3d_vis,
|
||||
'filename': '',
|
||||
'imgnum': 0,
|
||||
}
|
||||
)
|
||||
|
||||
return gt_db
|
||||
|
||||
def evaluate(self, cfg, preds, output_dir, *args, **kwargs):
|
||||
# convert 0-based index to 1-based index
|
||||
preds = preds[:, :, 0:2] + 1.0
|
||||
|
||||
if output_dir:
|
||||
pred_file = os.path.join(output_dir, 'pred.mat')
|
||||
savemat(pred_file, mdict={'preds': preds})
|
||||
|
||||
if 'test' in cfg.DATASET.TEST_SET:
|
||||
return {'Null': 0.0}, 0.0
|
||||
|
||||
SC_BIAS = 0.6
|
||||
threshold = 0.5
|
||||
|
||||
gt_file = os.path.join(cfg.DATASET.ROOT,
|
||||
'annot',
|
||||
'gt_{}.mat'.format(cfg.DATASET.TEST_SET))
|
||||
gt_dict = loadmat(gt_file)
|
||||
dataset_joints = gt_dict['dataset_joints']
|
||||
jnt_missing = gt_dict['jnt_missing']
|
||||
pos_gt_src = gt_dict['pos_gt_src']
|
||||
headboxes_src = gt_dict['headboxes_src']
|
||||
|
||||
pos_pred_src = np.transpose(preds, [1, 2, 0])
|
||||
|
||||
head = np.where(dataset_joints == 'head')[1][0]
|
||||
lsho = np.where(dataset_joints == 'lsho')[1][0]
|
||||
lelb = np.where(dataset_joints == 'lelb')[1][0]
|
||||
lwri = np.where(dataset_joints == 'lwri')[1][0]
|
||||
lhip = np.where(dataset_joints == 'lhip')[1][0]
|
||||
lkne = np.where(dataset_joints == 'lkne')[1][0]
|
||||
lank = np.where(dataset_joints == 'lank')[1][0]
|
||||
|
||||
rsho = np.where(dataset_joints == 'rsho')[1][0]
|
||||
relb = np.where(dataset_joints == 'relb')[1][0]
|
||||
rwri = np.where(dataset_joints == 'rwri')[1][0]
|
||||
rkne = np.where(dataset_joints == 'rkne')[1][0]
|
||||
rank = np.where(dataset_joints == 'rank')[1][0]
|
||||
rhip = np.where(dataset_joints == 'rhip')[1][0]
|
||||
|
||||
jnt_visible = 1 - jnt_missing
|
||||
uv_error = pos_pred_src - pos_gt_src
|
||||
uv_err = np.linalg.norm(uv_error, axis=1)
|
||||
headsizes = headboxes_src[1, :, :] - headboxes_src[0, :, :]
|
||||
headsizes = np.linalg.norm(headsizes, axis=0)
|
||||
headsizes *= SC_BIAS
|
||||
scale = np.multiply(headsizes, np.ones((len(uv_err), 1)))
|
||||
scaled_uv_err = np.divide(uv_err, scale)
|
||||
scaled_uv_err = np.multiply(scaled_uv_err, jnt_visible)
|
||||
jnt_count = np.sum(jnt_visible, axis=1)
|
||||
less_than_threshold = np.multiply((scaled_uv_err <= threshold),
|
||||
jnt_visible)
|
||||
PCKh = np.divide(100.*np.sum(less_than_threshold, axis=1), jnt_count)
|
||||
|
||||
# save
|
||||
rng = np.arange(0, 0.5+0.01, 0.01)
|
||||
pckAll = np.zeros((len(rng), 16))
|
||||
|
||||
for r in range(len(rng)):
|
||||
threshold = rng[r]
|
||||
less_than_threshold = np.multiply(scaled_uv_err <= threshold,
|
||||
jnt_visible)
|
||||
pckAll[r, :] = np.divide(100.*np.sum(less_than_threshold, axis=1),
|
||||
jnt_count)
|
||||
|
||||
PCKh = np.ma.array(PCKh, mask=False)
|
||||
PCKh.mask[6:8] = True
|
||||
|
||||
jnt_count = np.ma.array(jnt_count, mask=False)
|
||||
jnt_count.mask[6:8] = True
|
||||
jnt_ratio = jnt_count / np.sum(jnt_count).astype(np.float64)
|
||||
|
||||
name_value = [
|
||||
('Head', PCKh[head]),
|
||||
('Shoulder', 0.5 * (PCKh[lsho] + PCKh[rsho])),
|
||||
('Elbow', 0.5 * (PCKh[lelb] + PCKh[relb])),
|
||||
('Wrist', 0.5 * (PCKh[lwri] + PCKh[rwri])),
|
||||
('Hip', 0.5 * (PCKh[lhip] + PCKh[rhip])),
|
||||
('Knee', 0.5 * (PCKh[lkne] + PCKh[rkne])),
|
||||
('Ankle', 0.5 * (PCKh[lank] + PCKh[rank])),
|
||||
('Mean', np.sum(PCKh * jnt_ratio)),
|
||||
('Mean@0.1', np.sum(pckAll[11, :] * jnt_ratio))
|
||||
]
|
||||
name_value = OrderedDict(name_value)
|
||||
|
||||
return name_value, name_value['Mean']
|
||||
@@ -0,0 +1,13 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
@@ -0,0 +1,501 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import os
|
||||
import logging
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
BN_MOMENTUM = 0.1
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def conv3x3(in_planes, out_planes, stride=1):
|
||||
"""3x3 convolution with padding"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
|
||||
padding=1, bias=False)
|
||||
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
expansion = 1
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None):
|
||||
super(BasicBlock, self).__init__()
|
||||
self.conv1 = conv3x3(inplanes, planes, stride)
|
||||
self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.conv2 = conv3x3(planes, planes)
|
||||
self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
residual = self.downsample(x)
|
||||
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class Bottleneck(nn.Module):
|
||||
expansion = 4
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None):
|
||||
super(Bottleneck, self).__init__()
|
||||
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
|
||||
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
|
||||
padding=1, bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
|
||||
self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1,
|
||||
bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes * self.expansion,
|
||||
momentum=BN_MOMENTUM)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
residual = self.downsample(x)
|
||||
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class HighResolutionModule(nn.Module):
|
||||
def __init__(self, num_branches, blocks, num_blocks, num_inchannels,
|
||||
num_channels, fuse_method, multi_scale_output=True):
|
||||
super(HighResolutionModule, self).__init__()
|
||||
self._check_branches(
|
||||
num_branches, blocks, num_blocks, num_inchannels, num_channels)
|
||||
|
||||
self.num_inchannels = num_inchannels
|
||||
self.fuse_method = fuse_method
|
||||
self.num_branches = num_branches
|
||||
|
||||
self.multi_scale_output = multi_scale_output
|
||||
|
||||
self.branches = self._make_branches(
|
||||
num_branches, blocks, num_blocks, num_channels)
|
||||
self.fuse_layers = self._make_fuse_layers()
|
||||
self.relu = nn.ReLU(True)
|
||||
|
||||
def _check_branches(self, num_branches, blocks, num_blocks,
|
||||
num_inchannels, num_channels):
|
||||
if num_branches != len(num_blocks):
|
||||
error_msg = 'NUM_BRANCHES({}) <> NUM_BLOCKS({})'.format(
|
||||
num_branches, len(num_blocks))
|
||||
logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
if num_branches != len(num_channels):
|
||||
error_msg = 'NUM_BRANCHES({}) <> NUM_CHANNELS({})'.format(
|
||||
num_branches, len(num_channels))
|
||||
logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
if num_branches != len(num_inchannels):
|
||||
error_msg = 'NUM_BRANCHES({}) <> NUM_INCHANNELS({})'.format(
|
||||
num_branches, len(num_inchannels))
|
||||
logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
def _make_one_branch(self, branch_index, block, num_blocks, num_channels,
|
||||
stride=1):
|
||||
downsample = None
|
||||
if stride != 1 or \
|
||||
self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion:
|
||||
downsample = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
self.num_inchannels[branch_index],
|
||||
num_channels[branch_index] * block.expansion,
|
||||
kernel_size=1, stride=stride, bias=False
|
||||
),
|
||||
nn.BatchNorm2d(
|
||||
num_channels[branch_index] * block.expansion,
|
||||
momentum=BN_MOMENTUM
|
||||
),
|
||||
)
|
||||
|
||||
layers = []
|
||||
layers.append(
|
||||
block(
|
||||
self.num_inchannels[branch_index],
|
||||
num_channels[branch_index],
|
||||
stride,
|
||||
downsample
|
||||
)
|
||||
)
|
||||
self.num_inchannels[branch_index] = \
|
||||
num_channels[branch_index] * block.expansion
|
||||
for i in range(1, num_blocks[branch_index]):
|
||||
layers.append(
|
||||
block(
|
||||
self.num_inchannels[branch_index],
|
||||
num_channels[branch_index]
|
||||
)
|
||||
)
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def _make_branches(self, num_branches, block, num_blocks, num_channels):
|
||||
branches = []
|
||||
|
||||
for i in range(num_branches):
|
||||
branches.append(
|
||||
self._make_one_branch(i, block, num_blocks, num_channels)
|
||||
)
|
||||
|
||||
return nn.ModuleList(branches)
|
||||
|
||||
def _make_fuse_layers(self):
|
||||
if self.num_branches == 1:
|
||||
return None
|
||||
|
||||
num_branches = self.num_branches
|
||||
num_inchannels = self.num_inchannels
|
||||
fuse_layers = []
|
||||
for i in range(num_branches if self.multi_scale_output else 1):
|
||||
fuse_layer = []
|
||||
for j in range(num_branches):
|
||||
if j > i:
|
||||
fuse_layer.append(
|
||||
nn.Sequential(
|
||||
nn.Conv2d(
|
||||
num_inchannels[j],
|
||||
num_inchannels[i],
|
||||
1, 1, 0, bias=False
|
||||
),
|
||||
nn.BatchNorm2d(num_inchannels[i]),
|
||||
nn.Upsample(scale_factor=2**(j-i), mode='nearest')
|
||||
)
|
||||
)
|
||||
elif j == i:
|
||||
fuse_layer.append(None)
|
||||
else:
|
||||
conv3x3s = []
|
||||
for k in range(i-j):
|
||||
if k == i - j - 1:
|
||||
num_outchannels_conv3x3 = num_inchannels[i]
|
||||
conv3x3s.append(
|
||||
nn.Sequential(
|
||||
nn.Conv2d(
|
||||
num_inchannels[j],
|
||||
num_outchannels_conv3x3,
|
||||
3, 2, 1, bias=False
|
||||
),
|
||||
nn.BatchNorm2d(num_outchannels_conv3x3)
|
||||
)
|
||||
)
|
||||
else:
|
||||
num_outchannels_conv3x3 = num_inchannels[j]
|
||||
conv3x3s.append(
|
||||
nn.Sequential(
|
||||
nn.Conv2d(
|
||||
num_inchannels[j],
|
||||
num_outchannels_conv3x3,
|
||||
3, 2, 1, bias=False
|
||||
),
|
||||
nn.BatchNorm2d(num_outchannels_conv3x3),
|
||||
nn.ReLU(True)
|
||||
)
|
||||
)
|
||||
fuse_layer.append(nn.Sequential(*conv3x3s))
|
||||
fuse_layers.append(nn.ModuleList(fuse_layer))
|
||||
|
||||
return nn.ModuleList(fuse_layers)
|
||||
|
||||
def get_num_inchannels(self):
|
||||
return self.num_inchannels
|
||||
|
||||
def forward(self, x):
|
||||
if self.num_branches == 1:
|
||||
return [self.branches[0](x[0])]
|
||||
|
||||
for i in range(self.num_branches):
|
||||
x[i] = self.branches[i](x[i])
|
||||
|
||||
x_fuse = []
|
||||
|
||||
for i in range(len(self.fuse_layers)):
|
||||
y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
|
||||
for j in range(1, self.num_branches):
|
||||
if i == j:
|
||||
y = y + x[j]
|
||||
else:
|
||||
y = y + self.fuse_layers[i][j](x[j])
|
||||
x_fuse.append(self.relu(y))
|
||||
|
||||
return x_fuse
|
||||
|
||||
|
||||
blocks_dict = {
|
||||
'BASIC': BasicBlock,
|
||||
'BOTTLENECK': Bottleneck
|
||||
}
|
||||
|
||||
|
||||
class PoseHighResolutionNet(nn.Module):
|
||||
|
||||
def __init__(self, cfg, **kwargs):
|
||||
self.inplanes = 64
|
||||
extra = cfg.MODEL.EXTRA
|
||||
super(PoseHighResolutionNet, self).__init__()
|
||||
|
||||
# stem net
|
||||
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1,
|
||||
bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
|
||||
self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1,
|
||||
bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.layer1 = self._make_layer(Bottleneck, 64, 4)
|
||||
|
||||
self.stage2_cfg = cfg['MODEL']['EXTRA']['STAGE2']
|
||||
num_channels = self.stage2_cfg['NUM_CHANNELS']
|
||||
block = blocks_dict[self.stage2_cfg['BLOCK']]
|
||||
num_channels = [
|
||||
num_channels[i] * block.expansion for i in range(len(num_channels))
|
||||
]
|
||||
self.transition1 = self._make_transition_layer([256], num_channels)
|
||||
self.stage2, pre_stage_channels = self._make_stage(
|
||||
self.stage2_cfg, num_channels)
|
||||
|
||||
self.stage3_cfg = cfg['MODEL']['EXTRA']['STAGE3']
|
||||
num_channels = self.stage3_cfg['NUM_CHANNELS']
|
||||
block = blocks_dict[self.stage3_cfg['BLOCK']]
|
||||
num_channels = [
|
||||
num_channels[i] * block.expansion for i in range(len(num_channels))
|
||||
]
|
||||
self.transition2 = self._make_transition_layer(
|
||||
pre_stage_channels, num_channels)
|
||||
self.stage3, pre_stage_channels = self._make_stage(
|
||||
self.stage3_cfg, num_channels)
|
||||
|
||||
self.stage4_cfg = cfg['MODEL']['EXTRA']['STAGE4']
|
||||
num_channels = self.stage4_cfg['NUM_CHANNELS']
|
||||
block = blocks_dict[self.stage4_cfg['BLOCK']]
|
||||
num_channels = [
|
||||
num_channels[i] * block.expansion for i in range(len(num_channels))
|
||||
]
|
||||
self.transition3 = self._make_transition_layer(
|
||||
pre_stage_channels, num_channels)
|
||||
self.stage4, pre_stage_channels = self._make_stage(
|
||||
self.stage4_cfg, num_channels, multi_scale_output=False)
|
||||
|
||||
self.final_layer = nn.Conv2d(
|
||||
in_channels=pre_stage_channels[0],
|
||||
out_channels=cfg.MODEL.NUM_JOINTS,
|
||||
kernel_size=extra.FINAL_CONV_KERNEL,
|
||||
stride=1,
|
||||
padding=1 if extra.FINAL_CONV_KERNEL == 3 else 0
|
||||
)
|
||||
|
||||
self.pretrained_layers = cfg['MODEL']['EXTRA']['PRETRAINED_LAYERS']
|
||||
|
||||
def _make_transition_layer(
|
||||
self, num_channels_pre_layer, num_channels_cur_layer):
|
||||
num_branches_cur = len(num_channels_cur_layer)
|
||||
num_branches_pre = len(num_channels_pre_layer)
|
||||
|
||||
transition_layers = []
|
||||
for i in range(num_branches_cur):
|
||||
if i < num_branches_pre:
|
||||
if num_channels_cur_layer[i] != num_channels_pre_layer[i]:
|
||||
transition_layers.append(
|
||||
nn.Sequential(
|
||||
nn.Conv2d(
|
||||
num_channels_pre_layer[i],
|
||||
num_channels_cur_layer[i],
|
||||
3, 1, 1, bias=False
|
||||
),
|
||||
nn.BatchNorm2d(num_channels_cur_layer[i]),
|
||||
nn.ReLU(inplace=True)
|
||||
)
|
||||
)
|
||||
else:
|
||||
transition_layers.append(None)
|
||||
else:
|
||||
conv3x3s = []
|
||||
for j in range(i+1-num_branches_pre):
|
||||
inchannels = num_channels_pre_layer[-1]
|
||||
outchannels = num_channels_cur_layer[i] \
|
||||
if j == i-num_branches_pre else inchannels
|
||||
conv3x3s.append(
|
||||
nn.Sequential(
|
||||
nn.Conv2d(
|
||||
inchannels, outchannels, 3, 2, 1, bias=False
|
||||
),
|
||||
nn.BatchNorm2d(outchannels),
|
||||
nn.ReLU(inplace=True)
|
||||
)
|
||||
)
|
||||
transition_layers.append(nn.Sequential(*conv3x3s))
|
||||
|
||||
return nn.ModuleList(transition_layers)
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1):
|
||||
downsample = None
|
||||
if stride != 1 or self.inplanes != planes * block.expansion:
|
||||
downsample = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
self.inplanes, planes * block.expansion,
|
||||
kernel_size=1, stride=stride, bias=False
|
||||
),
|
||||
nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM),
|
||||
)
|
||||
|
||||
layers = []
|
||||
layers.append(block(self.inplanes, planes, stride, downsample))
|
||||
self.inplanes = planes * block.expansion
|
||||
for i in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def _make_stage(self, layer_config, num_inchannels,
|
||||
multi_scale_output=True):
|
||||
num_modules = layer_config['NUM_MODULES']
|
||||
num_branches = layer_config['NUM_BRANCHES']
|
||||
num_blocks = layer_config['NUM_BLOCKS']
|
||||
num_channels = layer_config['NUM_CHANNELS']
|
||||
block = blocks_dict[layer_config['BLOCK']]
|
||||
fuse_method = layer_config['FUSE_METHOD']
|
||||
|
||||
modules = []
|
||||
for i in range(num_modules):
|
||||
# multi_scale_output is only used last module
|
||||
if not multi_scale_output and i == num_modules - 1:
|
||||
reset_multi_scale_output = False
|
||||
else:
|
||||
reset_multi_scale_output = True
|
||||
|
||||
modules.append(
|
||||
HighResolutionModule(
|
||||
num_branches,
|
||||
block,
|
||||
num_blocks,
|
||||
num_inchannels,
|
||||
num_channels,
|
||||
fuse_method,
|
||||
reset_multi_scale_output
|
||||
)
|
||||
)
|
||||
num_inchannels = modules[-1].get_num_inchannels()
|
||||
|
||||
return nn.Sequential(*modules), num_inchannels
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
x = self.conv2(x)
|
||||
x = self.bn2(x)
|
||||
x = self.relu(x)
|
||||
x = self.layer1(x)
|
||||
|
||||
x_list = []
|
||||
for i in range(self.stage2_cfg['NUM_BRANCHES']):
|
||||
if self.transition1[i] is not None:
|
||||
x_list.append(self.transition1[i](x))
|
||||
else:
|
||||
x_list.append(x)
|
||||
y_list = self.stage2(x_list)
|
||||
|
||||
x_list = []
|
||||
for i in range(self.stage3_cfg['NUM_BRANCHES']):
|
||||
if self.transition2[i] is not None:
|
||||
x_list.append(self.transition2[i](y_list[-1]))
|
||||
else:
|
||||
x_list.append(y_list[i])
|
||||
y_list = self.stage3(x_list)
|
||||
|
||||
x_list = []
|
||||
for i in range(self.stage4_cfg['NUM_BRANCHES']):
|
||||
if self.transition3[i] is not None:
|
||||
x_list.append(self.transition3[i](y_list[-1]))
|
||||
else:
|
||||
x_list.append(y_list[i])
|
||||
y_list = self.stage4(x_list)
|
||||
|
||||
x = self.final_layer(y_list[0])
|
||||
|
||||
return x
|
||||
|
||||
def init_weights(self, pretrained=''):
|
||||
logger.info('=> init weights from normal distribution')
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
# nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
||||
nn.init.normal_(m.weight, std=0.001)
|
||||
for name, _ in m.named_parameters():
|
||||
if name in ['bias']:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.ConvTranspose2d):
|
||||
nn.init.normal_(m.weight, std=0.001)
|
||||
for name, _ in m.named_parameters():
|
||||
if name in ['bias']:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
if os.path.isfile(pretrained):
|
||||
pretrained_state_dict = torch.load(pretrained)
|
||||
logger.info('=> loading pretrained model {}'.format(pretrained))
|
||||
|
||||
need_init_state_dict = {}
|
||||
for name, m in pretrained_state_dict.items():
|
||||
if name.split('.')[0] in self.pretrained_layers \
|
||||
or self.pretrained_layers[0] is '*':
|
||||
need_init_state_dict[name] = m
|
||||
self.load_state_dict(need_init_state_dict, strict=False)
|
||||
elif pretrained:
|
||||
logger.error('=> please download pre-trained models first!')
|
||||
raise ValueError('{} is not exist!'.format(pretrained))
|
||||
|
||||
|
||||
def get_pose_net(cfg, is_train, **kwargs):
|
||||
model = PoseHighResolutionNet(cfg, **kwargs)
|
||||
|
||||
if is_train and cfg.MODEL.INIT_WEIGHTS:
|
||||
model.init_weights(cfg.MODEL.PRETRAINED)
|
||||
|
||||
return model
|
||||
@@ -0,0 +1,271 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import os
|
||||
import logging
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
BN_MOMENTUM = 0.1
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def conv3x3(in_planes, out_planes, stride=1):
|
||||
"""3x3 convolution with padding"""
|
||||
return nn.Conv2d(
|
||||
in_planes, out_planes, kernel_size=3, stride=stride,
|
||||
padding=1, bias=False
|
||||
)
|
||||
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
expansion = 1
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None):
|
||||
super(BasicBlock, self).__init__()
|
||||
self.conv1 = conv3x3(inplanes, planes, stride)
|
||||
self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.conv2 = conv3x3(planes, planes)
|
||||
self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
residual = self.downsample(x)
|
||||
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class Bottleneck(nn.Module):
|
||||
expansion = 4
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None):
|
||||
super(Bottleneck, self).__init__()
|
||||
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
|
||||
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
|
||||
padding=1, bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
|
||||
self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1,
|
||||
bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes * self.expansion,
|
||||
momentum=BN_MOMENTUM)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
residual = self.downsample(x)
|
||||
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class PoseResNet(nn.Module):
|
||||
|
||||
def __init__(self, block, layers, cfg, **kwargs):
|
||||
self.inplanes = 64
|
||||
extra = cfg.MODEL.EXTRA
|
||||
self.deconv_with_bias = extra.DECONV_WITH_BIAS
|
||||
|
||||
super(PoseResNet, self).__init__()
|
||||
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
|
||||
bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
|
||||
self.layer1 = self._make_layer(block, 64, layers[0])
|
||||
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
|
||||
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
|
||||
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
|
||||
|
||||
# used for deconv layers
|
||||
self.deconv_layers = self._make_deconv_layer(
|
||||
extra.NUM_DECONV_LAYERS,
|
||||
extra.NUM_DECONV_FILTERS,
|
||||
extra.NUM_DECONV_KERNELS,
|
||||
)
|
||||
|
||||
self.final_layer = nn.Conv2d(
|
||||
in_channels=extra.NUM_DECONV_FILTERS[-1],
|
||||
out_channels=cfg.MODEL.NUM_JOINTS,
|
||||
kernel_size=extra.FINAL_CONV_KERNEL,
|
||||
stride=1,
|
||||
padding=1 if extra.FINAL_CONV_KERNEL == 3 else 0
|
||||
)
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1):
|
||||
downsample = None
|
||||
if stride != 1 or self.inplanes != planes * block.expansion:
|
||||
downsample = nn.Sequential(
|
||||
nn.Conv2d(self.inplanes, planes * block.expansion,
|
||||
kernel_size=1, stride=stride, bias=False),
|
||||
nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM),
|
||||
)
|
||||
|
||||
layers = []
|
||||
layers.append(block(self.inplanes, planes, stride, downsample))
|
||||
self.inplanes = planes * block.expansion
|
||||
for i in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def _get_deconv_cfg(self, deconv_kernel, index):
|
||||
if deconv_kernel == 4:
|
||||
padding = 1
|
||||
output_padding = 0
|
||||
elif deconv_kernel == 3:
|
||||
padding = 1
|
||||
output_padding = 1
|
||||
elif deconv_kernel == 2:
|
||||
padding = 0
|
||||
output_padding = 0
|
||||
|
||||
return deconv_kernel, padding, output_padding
|
||||
|
||||
def _make_deconv_layer(self, num_layers, num_filters, num_kernels):
|
||||
assert num_layers == len(num_filters), \
|
||||
'ERROR: num_deconv_layers is different len(num_deconv_filters)'
|
||||
assert num_layers == len(num_kernels), \
|
||||
'ERROR: num_deconv_layers is different len(num_deconv_filters)'
|
||||
|
||||
layers = []
|
||||
for i in range(num_layers):
|
||||
kernel, padding, output_padding = \
|
||||
self._get_deconv_cfg(num_kernels[i], i)
|
||||
|
||||
planes = num_filters[i]
|
||||
layers.append(
|
||||
nn.ConvTranspose2d(
|
||||
in_channels=self.inplanes,
|
||||
out_channels=planes,
|
||||
kernel_size=kernel,
|
||||
stride=2,
|
||||
padding=padding,
|
||||
output_padding=output_padding,
|
||||
bias=self.deconv_with_bias))
|
||||
layers.append(nn.BatchNorm2d(planes, momentum=BN_MOMENTUM))
|
||||
layers.append(nn.ReLU(inplace=True))
|
||||
self.inplanes = planes
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
x = self.maxpool(x)
|
||||
|
||||
x = self.layer1(x)
|
||||
x = self.layer2(x)
|
||||
x = self.layer3(x)
|
||||
x = self.layer4(x)
|
||||
|
||||
x = self.deconv_layers(x)
|
||||
x = self.final_layer(x)
|
||||
|
||||
return x
|
||||
|
||||
def init_weights(self, pretrained=''):
|
||||
if os.path.isfile(pretrained):
|
||||
logger.info('=> init deconv weights from normal distribution')
|
||||
for name, m in self.deconv_layers.named_modules():
|
||||
if isinstance(m, nn.ConvTranspose2d):
|
||||
logger.info('=> init {}.weight as normal(0, 0.001)'.format(name))
|
||||
logger.info('=> init {}.bias as 0'.format(name))
|
||||
nn.init.normal_(m.weight, std=0.001)
|
||||
if self.deconv_with_bias:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
logger.info('=> init {}.weight as 1'.format(name))
|
||||
logger.info('=> init {}.bias as 0'.format(name))
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
logger.info('=> init final conv weights from normal distribution')
|
||||
for m in self.final_layer.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
# nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
||||
logger.info('=> init {}.weight as normal(0, 0.001)'.format(name))
|
||||
logger.info('=> init {}.bias as 0'.format(name))
|
||||
nn.init.normal_(m.weight, std=0.001)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
pretrained_state_dict = torch.load(pretrained)
|
||||
logger.info('=> loading pretrained model {}'.format(pretrained))
|
||||
self.load_state_dict(pretrained_state_dict, strict=False)
|
||||
else:
|
||||
logger.info('=> init weights from normal distribution')
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
# nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
||||
nn.init.normal_(m.weight, std=0.001)
|
||||
# nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.ConvTranspose2d):
|
||||
nn.init.normal_(m.weight, std=0.001)
|
||||
if self.deconv_with_bias:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
|
||||
resnet_spec = {
|
||||
18: (BasicBlock, [2, 2, 2, 2]),
|
||||
34: (BasicBlock, [3, 4, 6, 3]),
|
||||
50: (Bottleneck, [3, 4, 6, 3]),
|
||||
101: (Bottleneck, [3, 4, 23, 3]),
|
||||
152: (Bottleneck, [3, 8, 36, 3])
|
||||
}
|
||||
|
||||
|
||||
def get_pose_net(cfg, is_train, **kwargs):
|
||||
num_layers = cfg.MODEL.EXTRA.NUM_LAYERS
|
||||
|
||||
block_class, layers = resnet_spec[num_layers]
|
||||
|
||||
model = PoseResNet(block_class, layers, cfg, **kwargs)
|
||||
|
||||
if is_train and cfg.MODEL.INIT_WEIGHTS:
|
||||
model.init_weights(cfg.MODEL.PRETRAINED)
|
||||
|
||||
return model
|
||||
@@ -0,0 +1,164 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Modified from py-faster-rcnn (https://github.com/rbgirshick/py-faster-rcnn)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import numpy as np
|
||||
|
||||
def py_nms_wrapper(thresh):
|
||||
def _nms(dets):
|
||||
return nms(dets, thresh)
|
||||
return _nms
|
||||
|
||||
|
||||
def nms(dets, thresh):
|
||||
"""
|
||||
greedily select boxes with high confidence and overlap with current maximum <= thresh
|
||||
rule out overlap >= thresh
|
||||
:param dets: [[x1, y1, x2, y2 score]]
|
||||
:param thresh: retain overlap < thresh
|
||||
:return: indexes to keep
|
||||
"""
|
||||
if dets.shape[0] == 0:
|
||||
return []
|
||||
|
||||
x1 = dets[:, 0]
|
||||
y1 = dets[:, 1]
|
||||
x2 = dets[:, 2]
|
||||
y2 = dets[:, 3]
|
||||
scores = dets[:, 4]
|
||||
|
||||
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
|
||||
order = scores.argsort()[::-1]
|
||||
|
||||
keep = []
|
||||
while order.size > 0:
|
||||
i = order[0]
|
||||
keep.append(i)
|
||||
xx1 = np.maximum(x1[i], x1[order[1:]])
|
||||
yy1 = np.maximum(y1[i], y1[order[1:]])
|
||||
xx2 = np.minimum(x2[i], x2[order[1:]])
|
||||
yy2 = np.minimum(y2[i], y2[order[1:]])
|
||||
|
||||
w = np.maximum(0.0, xx2 - xx1 + 1)
|
||||
h = np.maximum(0.0, yy2 - yy1 + 1)
|
||||
inter = w * h
|
||||
ovr = inter / (areas[i] + areas[order[1:]] - inter)
|
||||
|
||||
inds = np.where(ovr <= thresh)[0]
|
||||
order = order[inds + 1]
|
||||
|
||||
return keep
|
||||
|
||||
|
||||
def oks_iou(g, d, a_g, a_d, sigmas=None, in_vis_thre=None):
|
||||
if not isinstance(sigmas, np.ndarray):
|
||||
sigmas = np.array([.26, .25, .25, .35, .35, .79, .79, .72, .72, .62, .62, 1.07, 1.07, .87, .87, .89, .89]) / 10.0
|
||||
vars = (sigmas * 2) ** 2
|
||||
xg = g[0::3]
|
||||
yg = g[1::3]
|
||||
vg = g[2::3]
|
||||
ious = np.zeros((d.shape[0]))
|
||||
for n_d in range(0, d.shape[0]):
|
||||
xd = d[n_d, 0::3]
|
||||
yd = d[n_d, 1::3]
|
||||
vd = d[n_d, 2::3]
|
||||
dx = xd - xg
|
||||
dy = yd - yg
|
||||
e = (dx ** 2 + dy ** 2) / vars / ((a_g + a_d[n_d]) / 2 + np.spacing(1)) / 2
|
||||
if in_vis_thre is not None:
|
||||
ind = list(vg > in_vis_thre) and list(vd > in_vis_thre)
|
||||
e = e[ind]
|
||||
ious[n_d] = np.sum(np.exp(-e)) / e.shape[0] if e.shape[0] != 0 else 0.0
|
||||
return ious
|
||||
|
||||
|
||||
def oks_nms(kpts_db, thresh, sigmas=None, in_vis_thre=None):
|
||||
"""
|
||||
greedily select boxes with high confidence and overlap with current maximum <= thresh
|
||||
rule out overlap >= thresh, overlap = oks
|
||||
:param kpts_db
|
||||
:param thresh: retain overlap < thresh
|
||||
:return: indexes to keep
|
||||
"""
|
||||
if len(kpts_db) == 0:
|
||||
return []
|
||||
|
||||
scores = np.array([kpts_db[i]['score'] for i in range(len(kpts_db))])
|
||||
kpts = np.array([kpts_db[i]['keypoints'].flatten() for i in range(len(kpts_db))])
|
||||
areas = np.array([kpts_db[i]['area'] for i in range(len(kpts_db))])
|
||||
|
||||
order = scores.argsort()[::-1]
|
||||
|
||||
keep = []
|
||||
while order.size > 0:
|
||||
i = order[0]
|
||||
keep.append(i)
|
||||
|
||||
oks_ovr = oks_iou(kpts[i], kpts[order[1:]], areas[i], areas[order[1:]], sigmas, in_vis_thre)
|
||||
|
||||
inds = np.where(oks_ovr <= thresh)[0]
|
||||
order = order[inds + 1]
|
||||
|
||||
return keep
|
||||
|
||||
|
||||
def rescore(overlap, scores, thresh, type='gaussian'):
|
||||
assert overlap.shape[0] == scores.shape[0]
|
||||
if type == 'linear':
|
||||
inds = np.where(overlap >= thresh)[0]
|
||||
scores[inds] = scores[inds] * (1 - overlap[inds])
|
||||
else:
|
||||
scores = scores * np.exp(- overlap**2 / thresh)
|
||||
|
||||
return scores
|
||||
|
||||
|
||||
def soft_oks_nms(kpts_db, thresh, sigmas=None, in_vis_thre=None):
|
||||
"""
|
||||
greedily select boxes with high confidence and overlap with current maximum <= thresh
|
||||
rule out overlap >= thresh, overlap = oks
|
||||
:param kpts_db
|
||||
:param thresh: retain overlap < thresh
|
||||
:return: indexes to keep
|
||||
"""
|
||||
if len(kpts_db) == 0:
|
||||
return []
|
||||
|
||||
scores = np.array([kpts_db[i]['score'] for i in range(len(kpts_db))])
|
||||
kpts = np.array([kpts_db[i]['keypoints'].flatten() for i in range(len(kpts_db))])
|
||||
areas = np.array([kpts_db[i]['area'] for i in range(len(kpts_db))])
|
||||
|
||||
order = scores.argsort()[::-1]
|
||||
scores = scores[order]
|
||||
|
||||
# max_dets = order.size
|
||||
max_dets = 20
|
||||
keep = np.zeros(max_dets, dtype=np.intp)
|
||||
keep_cnt = 0
|
||||
while order.size > 0 and keep_cnt < max_dets:
|
||||
i = order[0]
|
||||
|
||||
oks_ovr = oks_iou(kpts[i], kpts[order[1:]], areas[i], areas[order[1:]], sigmas, in_vis_thre)
|
||||
|
||||
order = order[1:]
|
||||
scores = rescore(oks_ovr, scores[1:], thresh)
|
||||
|
||||
tmp = scores.argsort()[::-1]
|
||||
order = order[tmp]
|
||||
scores = scores[tmp]
|
||||
|
||||
keep[keep_cnt] = i
|
||||
keep_cnt += 1
|
||||
|
||||
keep = keep[:keep_cnt]
|
||||
|
||||
return keep
|
||||
# kpts_db = kpts_db[:keep_cnt]
|
||||
|
||||
# return kpts_db
|
||||
@@ -0,0 +1,121 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
|
||||
def flip_back(output_flipped, matched_parts):
|
||||
'''
|
||||
ouput_flipped: numpy.ndarray(batch_size, num_joints, height, width)
|
||||
'''
|
||||
assert output_flipped.ndim == 4,\
|
||||
'output_flipped should be [batch_size, num_joints, height, width]'
|
||||
|
||||
output_flipped = output_flipped[:, :, :, ::-1]
|
||||
|
||||
for pair in matched_parts:
|
||||
tmp = output_flipped[:, pair[0], :, :].copy()
|
||||
output_flipped[:, pair[0], :, :] = output_flipped[:, pair[1], :, :]
|
||||
output_flipped[:, pair[1], :, :] = tmp
|
||||
|
||||
return output_flipped
|
||||
|
||||
|
||||
def fliplr_joints(joints, joints_vis, width, matched_parts):
|
||||
"""
|
||||
flip coords
|
||||
"""
|
||||
# Flip horizontal
|
||||
joints[:, 0] = width - joints[:, 0] - 1
|
||||
|
||||
# Change left-right parts
|
||||
for pair in matched_parts:
|
||||
joints[pair[0], :], joints[pair[1], :] = \
|
||||
joints[pair[1], :], joints[pair[0], :].copy()
|
||||
joints_vis[pair[0], :], joints_vis[pair[1], :] = \
|
||||
joints_vis[pair[1], :], joints_vis[pair[0], :].copy()
|
||||
|
||||
return joints*joints_vis, joints_vis
|
||||
|
||||
|
||||
def transform_preds(coords, center, scale, output_size):
|
||||
target_coords = np.zeros(coords.shape)
|
||||
trans = get_affine_transform(center, scale, 0, output_size, inv=1)
|
||||
for p in range(coords.shape[0]):
|
||||
target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans)
|
||||
return target_coords
|
||||
|
||||
|
||||
def get_affine_transform(
|
||||
center, scale, rot, output_size,
|
||||
shift=np.array([0, 0], dtype=np.float32), inv=0
|
||||
):
|
||||
if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
|
||||
print(scale)
|
||||
scale = np.array([scale, scale])
|
||||
|
||||
scale_tmp = scale * 200.0
|
||||
src_w = scale_tmp[0]
|
||||
dst_w = output_size[0]
|
||||
dst_h = output_size[1]
|
||||
|
||||
rot_rad = np.pi * rot / 180
|
||||
src_dir = get_dir([0, src_w * -0.5], rot_rad)
|
||||
dst_dir = np.array([0, dst_w * -0.5], np.float32)
|
||||
|
||||
src = np.zeros((3, 2), dtype=np.float32)
|
||||
dst = np.zeros((3, 2), dtype=np.float32)
|
||||
src[0, :] = center + scale_tmp * shift
|
||||
src[1, :] = center + src_dir + scale_tmp * shift
|
||||
dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
|
||||
dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
|
||||
|
||||
src[2:, :] = get_3rd_point(src[0, :], src[1, :])
|
||||
dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])
|
||||
|
||||
if inv:
|
||||
trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
|
||||
else:
|
||||
trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
|
||||
|
||||
return trans
|
||||
|
||||
|
||||
def affine_transform(pt, t):
|
||||
new_pt = np.array([pt[0], pt[1], 1.]).T
|
||||
new_pt = np.dot(t, new_pt)
|
||||
return new_pt[:2]
|
||||
|
||||
|
||||
def get_3rd_point(a, b):
|
||||
direct = a - b
|
||||
return b + np.array([-direct[1], direct[0]], dtype=np.float32)
|
||||
|
||||
|
||||
def get_dir(src_point, rot_rad):
|
||||
sn, cs = np.sin(rot_rad), np.cos(rot_rad)
|
||||
|
||||
src_result = [0, 0]
|
||||
src_result[0] = src_point[0] * cs - src_point[1] * sn
|
||||
src_result[1] = src_point[0] * sn + src_point[1] * cs
|
||||
|
||||
return src_result
|
||||
|
||||
|
||||
def crop(img, center, scale, output_size, rot=0):
|
||||
trans = get_affine_transform(center, scale, rot, output_size)
|
||||
|
||||
dst_img = cv2.warpAffine(
|
||||
img, trans, (int(output_size[0]), int(output_size[1])),
|
||||
flags=cv2.INTER_LINEAR
|
||||
)
|
||||
|
||||
return dst_img
|
||||
@@ -0,0 +1,203 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import os
|
||||
import logging
|
||||
import time
|
||||
from collections import namedtuple
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.optim as optim
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
def create_logger(cfg, cfg_name, phase='train'):
|
||||
root_output_dir = Path(cfg.OUTPUT_DIR)
|
||||
# set up logger
|
||||
if not root_output_dir.exists():
|
||||
print('=> creating {}'.format(root_output_dir))
|
||||
root_output_dir.mkdir()
|
||||
|
||||
dataset = cfg.DATASET.DATASET + '_' + cfg.DATASET.HYBRID_JOINTS_TYPE \
|
||||
if cfg.DATASET.HYBRID_JOINTS_TYPE else cfg.DATASET.DATASET
|
||||
dataset = dataset.replace(':', '_')
|
||||
model = cfg.MODEL.NAME
|
||||
cfg_name = os.path.basename(cfg_name).split('.')[0]
|
||||
|
||||
final_output_dir = root_output_dir / dataset / model / cfg_name
|
||||
|
||||
print('=> creating {}'.format(final_output_dir))
|
||||
final_output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
time_str = time.strftime('%Y-%m-%d-%H-%M')
|
||||
log_file = '{}_{}_{}.log'.format(cfg_name, time_str, phase)
|
||||
final_log_file = final_output_dir / log_file
|
||||
head = '%(asctime)-15s %(message)s'
|
||||
logging.basicConfig(filename=str(final_log_file),
|
||||
format=head)
|
||||
logger = logging.getLogger()
|
||||
logger.setLevel(logging.INFO)
|
||||
console = logging.StreamHandler()
|
||||
logging.getLogger('').addHandler(console)
|
||||
|
||||
tensorboard_log_dir = Path(cfg.LOG_DIR) / dataset / model / \
|
||||
(cfg_name + '_' + time_str)
|
||||
|
||||
print('=> creating {}'.format(tensorboard_log_dir))
|
||||
tensorboard_log_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
return logger, str(final_output_dir), str(tensorboard_log_dir)
|
||||
|
||||
|
||||
def get_optimizer(cfg, model):
|
||||
optimizer = None
|
||||
if cfg.TRAIN.OPTIMIZER == 'sgd':
|
||||
optimizer = optim.SGD(
|
||||
model.parameters(),
|
||||
lr=cfg.TRAIN.LR,
|
||||
momentum=cfg.TRAIN.MOMENTUM,
|
||||
weight_decay=cfg.TRAIN.WD,
|
||||
nesterov=cfg.TRAIN.NESTEROV
|
||||
)
|
||||
elif cfg.TRAIN.OPTIMIZER == 'adam':
|
||||
optimizer = optim.Adam(
|
||||
model.parameters(),
|
||||
lr=cfg.TRAIN.LR
|
||||
)
|
||||
|
||||
return optimizer
|
||||
|
||||
|
||||
def save_checkpoint(states, is_best, output_dir,
|
||||
filename='checkpoint.pth'):
|
||||
torch.save(states, os.path.join(output_dir, filename))
|
||||
if is_best and 'state_dict' in states:
|
||||
torch.save(states['best_state_dict'],
|
||||
os.path.join(output_dir, 'model_best.pth'))
|
||||
|
||||
|
||||
def get_model_summary(model, *input_tensors, item_length=26, verbose=False):
|
||||
"""
|
||||
:param model:
|
||||
:param input_tensors:
|
||||
:param item_length:
|
||||
:return:
|
||||
"""
|
||||
|
||||
summary = []
|
||||
|
||||
ModuleDetails = namedtuple(
|
||||
"Layer", ["name", "input_size", "output_size", "num_parameters", "multiply_adds"])
|
||||
hooks = []
|
||||
layer_instances = {}
|
||||
|
||||
def add_hooks(module):
|
||||
|
||||
def hook(module, input, output):
|
||||
class_name = str(module.__class__.__name__)
|
||||
|
||||
instance_index = 1
|
||||
if class_name not in layer_instances:
|
||||
layer_instances[class_name] = instance_index
|
||||
else:
|
||||
instance_index = layer_instances[class_name] + 1
|
||||
layer_instances[class_name] = instance_index
|
||||
|
||||
layer_name = class_name + "_" + str(instance_index)
|
||||
|
||||
params = 0
|
||||
|
||||
if class_name.find("Conv") != -1 or class_name.find("BatchNorm") != -1 or \
|
||||
class_name.find("Linear") != -1:
|
||||
for param_ in module.parameters():
|
||||
params += param_.view(-1).size(0)
|
||||
|
||||
flops = "Not Available"
|
||||
if class_name.find("Conv") != -1 and hasattr(module, "weight"):
|
||||
flops = (
|
||||
torch.prod(
|
||||
torch.LongTensor(list(module.weight.data.size()))) *
|
||||
torch.prod(
|
||||
torch.LongTensor(list(output.size())[2:]))).item()
|
||||
elif isinstance(module, nn.Linear):
|
||||
flops = (torch.prod(torch.LongTensor(list(output.size()))) \
|
||||
* input[0].size(1)).item()
|
||||
|
||||
if isinstance(input[0], list):
|
||||
input = input[0]
|
||||
if isinstance(output, list):
|
||||
output = output[0]
|
||||
|
||||
summary.append(
|
||||
ModuleDetails(
|
||||
name=layer_name,
|
||||
input_size=list(input[0].size()),
|
||||
output_size=list(output.size()),
|
||||
num_parameters=params,
|
||||
multiply_adds=flops)
|
||||
)
|
||||
|
||||
if not isinstance(module, nn.ModuleList) \
|
||||
and not isinstance(module, nn.Sequential) \
|
||||
and module != model:
|
||||
hooks.append(module.register_forward_hook(hook))
|
||||
|
||||
model.eval()
|
||||
model.apply(add_hooks)
|
||||
|
||||
space_len = item_length
|
||||
|
||||
model(*input_tensors)
|
||||
for hook in hooks:
|
||||
hook.remove()
|
||||
|
||||
details = ''
|
||||
if verbose:
|
||||
details = "Model Summary" + \
|
||||
os.linesep + \
|
||||
"Name{}Input Size{}Output Size{}Parameters{}Multiply Adds (Flops){}".format(
|
||||
' ' * (space_len - len("Name")),
|
||||
' ' * (space_len - len("Input Size")),
|
||||
' ' * (space_len - len("Output Size")),
|
||||
' ' * (space_len - len("Parameters")),
|
||||
' ' * (space_len - len("Multiply Adds (Flops)"))) \
|
||||
+ os.linesep + '-' * space_len * 5 + os.linesep
|
||||
|
||||
params_sum = 0
|
||||
flops_sum = 0
|
||||
for layer in summary:
|
||||
params_sum += layer.num_parameters
|
||||
if layer.multiply_adds != "Not Available":
|
||||
flops_sum += layer.multiply_adds
|
||||
if verbose:
|
||||
details += "{}{}{}{}{}{}{}{}{}{}".format(
|
||||
layer.name,
|
||||
' ' * (space_len - len(layer.name)),
|
||||
layer.input_size,
|
||||
' ' * (space_len - len(str(layer.input_size))),
|
||||
layer.output_size,
|
||||
' ' * (space_len - len(str(layer.output_size))),
|
||||
layer.num_parameters,
|
||||
' ' * (space_len - len(str(layer.num_parameters))),
|
||||
layer.multiply_adds,
|
||||
' ' * (space_len - len(str(layer.multiply_adds)))) \
|
||||
+ os.linesep + '-' * space_len * 5 + os.linesep
|
||||
|
||||
details += os.linesep \
|
||||
+ "Total Parameters: {:,}".format(params_sum) \
|
||||
+ os.linesep + '-' * space_len * 5 + os.linesep
|
||||
details += "Total Multiply Adds (For Convolution and Linear Layers only): {:,} GFLOPs".format(flops_sum/(1024**3)) \
|
||||
+ os.linesep + '-' * space_len * 5 + os.linesep
|
||||
details += "Number of Layers" + os.linesep
|
||||
for layer in layer_instances:
|
||||
details += "{} : {} layers ".format(layer, layer_instances[layer])
|
||||
|
||||
return details
|
||||
@@ -0,0 +1,141 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
import torchvision
|
||||
import cv2
|
||||
|
||||
from core.inference import get_max_preds
|
||||
|
||||
|
||||
def save_batch_image_with_joints(batch_image, batch_joints, batch_joints_vis,
|
||||
file_name, nrow=8, padding=2):
|
||||
'''
|
||||
batch_image: [batch_size, channel, height, width]
|
||||
batch_joints: [batch_size, num_joints, 3],
|
||||
batch_joints_vis: [batch_size, num_joints, 1],
|
||||
}
|
||||
'''
|
||||
grid = torchvision.utils.make_grid(batch_image, nrow, padding, True)
|
||||
ndarr = grid.mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()
|
||||
ndarr = ndarr.copy()
|
||||
|
||||
nmaps = batch_image.size(0)
|
||||
xmaps = min(nrow, nmaps)
|
||||
ymaps = int(math.ceil(float(nmaps) / xmaps))
|
||||
height = int(batch_image.size(2) + padding)
|
||||
width = int(batch_image.size(3) + padding)
|
||||
k = 0
|
||||
for y in range(ymaps):
|
||||
for x in range(xmaps):
|
||||
if k >= nmaps:
|
||||
break
|
||||
joints = batch_joints[k]
|
||||
joints_vis = batch_joints_vis[k]
|
||||
|
||||
for joint, joint_vis in zip(joints, joints_vis):
|
||||
joint[0] = x * width + padding + joint[0]
|
||||
joint[1] = y * height + padding + joint[1]
|
||||
if joint_vis[0]:
|
||||
cv2.circle(ndarr, (int(joint[0]), int(joint[1])), 2, [255, 0, 0], 2)
|
||||
k = k + 1
|
||||
cv2.imwrite(file_name, ndarr)
|
||||
|
||||
|
||||
def save_batch_heatmaps(batch_image, batch_heatmaps, file_name,
|
||||
normalize=True):
|
||||
'''
|
||||
batch_image: [batch_size, channel, height, width]
|
||||
batch_heatmaps: ['batch_size, num_joints, height, width]
|
||||
file_name: saved file name
|
||||
'''
|
||||
if normalize:
|
||||
batch_image = batch_image.clone()
|
||||
min = float(batch_image.min())
|
||||
max = float(batch_image.max())
|
||||
|
||||
batch_image.add_(-min).div_(max - min + 1e-5)
|
||||
|
||||
batch_size = batch_heatmaps.size(0)
|
||||
num_joints = batch_heatmaps.size(1)
|
||||
heatmap_height = batch_heatmaps.size(2)
|
||||
heatmap_width = batch_heatmaps.size(3)
|
||||
|
||||
grid_image = np.zeros((batch_size*heatmap_height,
|
||||
(num_joints+1)*heatmap_width,
|
||||
3),
|
||||
dtype=np.uint8)
|
||||
|
||||
preds, maxvals = get_max_preds(batch_heatmaps.detach().cpu().numpy())
|
||||
|
||||
for i in range(batch_size):
|
||||
image = batch_image[i].mul(255)\
|
||||
.clamp(0, 255)\
|
||||
.byte()\
|
||||
.permute(1, 2, 0)\
|
||||
.cpu().numpy()
|
||||
heatmaps = batch_heatmaps[i].mul(255)\
|
||||
.clamp(0, 255)\
|
||||
.byte()\
|
||||
.cpu().numpy()
|
||||
|
||||
resized_image = cv2.resize(image,
|
||||
(int(heatmap_width), int(heatmap_height)))
|
||||
|
||||
height_begin = heatmap_height * i
|
||||
height_end = heatmap_height * (i + 1)
|
||||
for j in range(num_joints):
|
||||
cv2.circle(resized_image,
|
||||
(int(preds[i][j][0]), int(preds[i][j][1])),
|
||||
1, [0, 0, 255], 1)
|
||||
heatmap = heatmaps[j, :, :]
|
||||
colored_heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
|
||||
masked_image = colored_heatmap*0.7 + resized_image*0.3
|
||||
cv2.circle(masked_image,
|
||||
(int(preds[i][j][0]), int(preds[i][j][1])),
|
||||
1, [0, 0, 255], 1)
|
||||
|
||||
width_begin = heatmap_width * (j+1)
|
||||
width_end = heatmap_width * (j+2)
|
||||
grid_image[height_begin:height_end, width_begin:width_end, :] = \
|
||||
masked_image
|
||||
# grid_image[height_begin:height_end, width_begin:width_end, :] = \
|
||||
# colored_heatmap*0.7 + resized_image*0.3
|
||||
|
||||
grid_image[height_begin:height_end, 0:heatmap_width, :] = resized_image
|
||||
|
||||
cv2.imwrite(file_name, grid_image)
|
||||
|
||||
|
||||
def save_debug_images(config, input, meta, target, joints_pred, output,
|
||||
prefix):
|
||||
if not config.DEBUG.DEBUG:
|
||||
return
|
||||
|
||||
if config.DEBUG.SAVE_BATCH_IMAGES_GT:
|
||||
save_batch_image_with_joints(
|
||||
input, meta['joints'], meta['joints_vis'],
|
||||
'{}_gt.jpg'.format(prefix)
|
||||
)
|
||||
if config.DEBUG.SAVE_BATCH_IMAGES_PRED:
|
||||
save_batch_image_with_joints(
|
||||
input, joints_pred, meta['joints_vis'],
|
||||
'{}_pred.jpg'.format(prefix)
|
||||
)
|
||||
if config.DEBUG.SAVE_HEATMAPS_GT:
|
||||
save_batch_heatmaps(
|
||||
input, target, '{}_hm_gt.jpg'.format(prefix)
|
||||
)
|
||||
if config.DEBUG.SAVE_HEATMAPS_PRED:
|
||||
save_batch_heatmaps(
|
||||
input, output, '{}_hm_pred.jpg'.format(prefix)
|
||||
)
|
||||
@@ -0,0 +1,70 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import os
|
||||
import zipfile
|
||||
import xml.etree.ElementTree as ET
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
_im_zfile = []
|
||||
_xml_path_zip = []
|
||||
_xml_zfile = []
|
||||
|
||||
|
||||
def imread(filename, flags=cv2.IMREAD_COLOR):
|
||||
global _im_zfile
|
||||
path = filename
|
||||
pos_at = path.index('@')
|
||||
if pos_at == -1:
|
||||
print("character '@' is not found from the given path '%s'"%(path))
|
||||
assert 0
|
||||
path_zip = path[0: pos_at]
|
||||
path_img = path[pos_at + 2:]
|
||||
if not os.path.isfile(path_zip):
|
||||
print("zip file '%s' is not found"%(path_zip))
|
||||
assert 0
|
||||
for i in range(len(_im_zfile)):
|
||||
if _im_zfile[i]['path'] == path_zip:
|
||||
data = _im_zfile[i]['zipfile'].read(path_img)
|
||||
return cv2.imdecode(np.frombuffer(data, np.uint8), flags)
|
||||
|
||||
_im_zfile.append({
|
||||
'path': path_zip,
|
||||
'zipfile': zipfile.ZipFile(path_zip, 'r')
|
||||
})
|
||||
data = _im_zfile[-1]['zipfile'].read(path_img)
|
||||
|
||||
return cv2.imdecode(np.frombuffer(data, np.uint8), flags)
|
||||
|
||||
|
||||
def xmlread(filename):
|
||||
global _xml_path_zip
|
||||
global _xml_zfile
|
||||
path = filename
|
||||
pos_at = path.index('@')
|
||||
if pos_at == -1:
|
||||
print("character '@' is not found from the given path '%s'"%(path))
|
||||
assert 0
|
||||
path_zip = path[0: pos_at]
|
||||
path_xml = path[pos_at + 2:]
|
||||
if not os.path.isfile(path_zip):
|
||||
print("zip file '%s' is not found"%(path_zip))
|
||||
assert 0
|
||||
for i in xrange(len(_xml_path_zip)):
|
||||
if _xml_path_zip[i] == path_zip:
|
||||
data = _xml_zfile[i].open(path_xml)
|
||||
return ET.fromstring(data.read())
|
||||
_xml_path_zip.append(path_zip)
|
||||
print("read new xml file '%s'"%(path_zip))
|
||||
_xml_zfile.append(zipfile.ZipFile(path_zip, 'r'))
|
||||
data = _xml_zfile[-1].open(path_xml)
|
||||
return ET.fromstring(data.read())
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,3 @@
|
||||
from .wider_face import WiderFaceDetection, detection_collate
|
||||
from .data_augment import *
|
||||
from .config import *
|
||||
@@ -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,237 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
import random
|
||||
from utils.box_utils_Retina import matrix_iof
|
||||
|
||||
|
||||
def _crop(image, boxes, labels, landm, img_dim):
|
||||
height, width, _ = image.shape
|
||||
pad_image_flag = True
|
||||
|
||||
for _ in range(250):
|
||||
"""
|
||||
if random.uniform(0, 1) <= 0.2:
|
||||
scale = 1.0
|
||||
else:
|
||||
scale = random.uniform(0.3, 1.0)
|
||||
"""
|
||||
PRE_SCALES = [0.3, 0.45, 0.6, 0.8, 1.0]
|
||||
scale = random.choice(PRE_SCALES)
|
||||
short_side = min(width, height)
|
||||
w = int(scale * short_side)
|
||||
h = w
|
||||
|
||||
if width == w:
|
||||
l = 0
|
||||
else:
|
||||
l = random.randrange(width - w)
|
||||
if height == h:
|
||||
t = 0
|
||||
else:
|
||||
t = random.randrange(height - h)
|
||||
roi = np.array((l, t, l + w, t + h))
|
||||
|
||||
value = matrix_iof(boxes, roi[np.newaxis])
|
||||
flag = (value >= 1)
|
||||
if not flag.any():
|
||||
continue
|
||||
|
||||
centers = (boxes[:, :2] + boxes[:, 2:]) / 2
|
||||
mask_a = np.logical_and(roi[:2] < centers, centers < roi[2:]).all(axis=1)
|
||||
boxes_t = boxes[mask_a].copy()
|
||||
labels_t = labels[mask_a].copy()
|
||||
landms_t = landm[mask_a].copy()
|
||||
landms_t = landms_t.reshape([-1, 5, 2])
|
||||
|
||||
if boxes_t.shape[0] == 0:
|
||||
continue
|
||||
|
||||
image_t = image[roi[1]:roi[3], roi[0]:roi[2]]
|
||||
|
||||
boxes_t[:, :2] = np.maximum(boxes_t[:, :2], roi[:2])
|
||||
boxes_t[:, :2] -= roi[:2]
|
||||
boxes_t[:, 2:] = np.minimum(boxes_t[:, 2:], roi[2:])
|
||||
boxes_t[:, 2:] -= roi[:2]
|
||||
|
||||
# landm
|
||||
landms_t[:, :, :2] = landms_t[:, :, :2] - roi[:2]
|
||||
landms_t[:, :, :2] = np.maximum(landms_t[:, :, :2], np.array([0, 0]))
|
||||
landms_t[:, :, :2] = np.minimum(landms_t[:, :, :2], roi[2:] - roi[:2])
|
||||
landms_t = landms_t.reshape([-1, 10])
|
||||
|
||||
|
||||
# make sure that the cropped image contains at least one face > 16 pixel at training image scale
|
||||
b_w_t = (boxes_t[:, 2] - boxes_t[:, 0] + 1) / w * img_dim
|
||||
b_h_t = (boxes_t[:, 3] - boxes_t[:, 1] + 1) / h * img_dim
|
||||
mask_b = np.minimum(b_w_t, b_h_t) > 0.0
|
||||
boxes_t = boxes_t[mask_b]
|
||||
labels_t = labels_t[mask_b]
|
||||
landms_t = landms_t[mask_b]
|
||||
|
||||
if boxes_t.shape[0] == 0:
|
||||
continue
|
||||
|
||||
pad_image_flag = False
|
||||
|
||||
return image_t, boxes_t, labels_t, landms_t, pad_image_flag
|
||||
return image, boxes, labels, landm, pad_image_flag
|
||||
|
||||
|
||||
def _distort(image):
|
||||
|
||||
def _convert(image, alpha=1, beta=0):
|
||||
tmp = image.astype(float) * alpha + beta
|
||||
tmp[tmp < 0] = 0
|
||||
tmp[tmp > 255] = 255
|
||||
image[:] = tmp
|
||||
|
||||
image = image.copy()
|
||||
|
||||
if random.randrange(2):
|
||||
|
||||
#brightness distortion
|
||||
if random.randrange(2):
|
||||
_convert(image, beta=random.uniform(-32, 32))
|
||||
|
||||
#contrast distortion
|
||||
if random.randrange(2):
|
||||
_convert(image, alpha=random.uniform(0.5, 1.5))
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
|
||||
|
||||
#saturation distortion
|
||||
if random.randrange(2):
|
||||
_convert(image[:, :, 1], alpha=random.uniform(0.5, 1.5))
|
||||
|
||||
#hue distortion
|
||||
if random.randrange(2):
|
||||
tmp = image[:, :, 0].astype(int) + random.randint(-18, 18)
|
||||
tmp %= 180
|
||||
image[:, :, 0] = tmp
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
|
||||
|
||||
else:
|
||||
|
||||
#brightness distortion
|
||||
if random.randrange(2):
|
||||
_convert(image, beta=random.uniform(-32, 32))
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
|
||||
|
||||
#saturation distortion
|
||||
if random.randrange(2):
|
||||
_convert(image[:, :, 1], alpha=random.uniform(0.5, 1.5))
|
||||
|
||||
#hue distortion
|
||||
if random.randrange(2):
|
||||
tmp = image[:, :, 0].astype(int) + random.randint(-18, 18)
|
||||
tmp %= 180
|
||||
image[:, :, 0] = tmp
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
|
||||
|
||||
#contrast distortion
|
||||
if random.randrange(2):
|
||||
_convert(image, alpha=random.uniform(0.5, 1.5))
|
||||
|
||||
return image
|
||||
|
||||
|
||||
def _expand(image, boxes, fill, p):
|
||||
if random.randrange(2):
|
||||
return image, boxes
|
||||
|
||||
height, width, depth = image.shape
|
||||
|
||||
scale = random.uniform(1, p)
|
||||
w = int(scale * width)
|
||||
h = int(scale * height)
|
||||
|
||||
left = random.randint(0, w - width)
|
||||
top = random.randint(0, h - height)
|
||||
|
||||
boxes_t = boxes.copy()
|
||||
boxes_t[:, :2] += (left, top)
|
||||
boxes_t[:, 2:] += (left, top)
|
||||
expand_image = np.empty(
|
||||
(h, w, depth),
|
||||
dtype=image.dtype)
|
||||
expand_image[:, :] = fill
|
||||
expand_image[top:top + height, left:left + width] = image
|
||||
image = expand_image
|
||||
|
||||
return image, boxes_t
|
||||
|
||||
|
||||
def _mirror(image, boxes, landms):
|
||||
_, width, _ = image.shape
|
||||
if random.randrange(2):
|
||||
image = image[:, ::-1]
|
||||
boxes = boxes.copy()
|
||||
boxes[:, 0::2] = width - boxes[:, 2::-2]
|
||||
|
||||
# landm
|
||||
landms = landms.copy()
|
||||
landms = landms.reshape([-1, 5, 2])
|
||||
landms[:, :, 0] = width - landms[:, :, 0]
|
||||
tmp = landms[:, 1, :].copy()
|
||||
landms[:, 1, :] = landms[:, 0, :]
|
||||
landms[:, 0, :] = tmp
|
||||
tmp1 = landms[:, 4, :].copy()
|
||||
landms[:, 4, :] = landms[:, 3, :]
|
||||
landms[:, 3, :] = tmp1
|
||||
landms = landms.reshape([-1, 10])
|
||||
|
||||
return image, boxes, landms
|
||||
|
||||
|
||||
def _pad_to_square(image, rgb_mean, pad_image_flag):
|
||||
if not pad_image_flag:
|
||||
return image
|
||||
height, width, _ = image.shape
|
||||
long_side = max(width, height)
|
||||
image_t = np.empty((long_side, long_side, 3), dtype=image.dtype)
|
||||
image_t[:, :] = rgb_mean
|
||||
image_t[0:0 + height, 0:0 + width] = image
|
||||
return image_t
|
||||
|
||||
|
||||
def _resize_subtract_mean(image, insize, rgb_mean):
|
||||
interp_methods = [cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_NEAREST, cv2.INTER_LANCZOS4]
|
||||
interp_method = interp_methods[random.randrange(5)]
|
||||
image = cv2.resize(image, (insize, insize), interpolation=interp_method)
|
||||
image = image.astype(np.float32)
|
||||
image -= rgb_mean
|
||||
return image.transpose(2, 0, 1)
|
||||
|
||||
|
||||
class preproc(object):
|
||||
|
||||
def __init__(self, img_dim, rgb_means):
|
||||
self.img_dim = img_dim
|
||||
self.rgb_means = rgb_means
|
||||
|
||||
def __call__(self, image, targets):
|
||||
assert targets.shape[0] > 0, "this image does not have gt"
|
||||
|
||||
boxes = targets[:, :4].copy()
|
||||
labels = targets[:, -1].copy()
|
||||
landm = targets[:, 4:-1].copy()
|
||||
|
||||
image_t, boxes_t, labels_t, landm_t, pad_image_flag = _crop(image, boxes, labels, landm, self.img_dim)
|
||||
image_t = _distort(image_t)
|
||||
image_t = _pad_to_square(image_t,self.rgb_means, pad_image_flag)
|
||||
image_t, boxes_t, landm_t = _mirror(image_t, boxes_t, landm_t)
|
||||
height, width, _ = image_t.shape
|
||||
image_t = _resize_subtract_mean(image_t, self.img_dim, self.rgb_means)
|
||||
boxes_t[:, 0::2] /= width
|
||||
boxes_t[:, 1::2] /= height
|
||||
|
||||
landm_t[:, 0::2] /= width
|
||||
landm_t[:, 1::2] /= height
|
||||
|
||||
labels_t = np.expand_dims(labels_t, 1)
|
||||
targets_t = np.hstack((boxes_t, landm_t, labels_t))
|
||||
|
||||
return image_t, targets_t
|
||||
@@ -0,0 +1,101 @@
|
||||
import os
|
||||
import os.path
|
||||
import sys
|
||||
import torch
|
||||
import torch.utils.data as data
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
class WiderFaceDetection(data.Dataset):
|
||||
def __init__(self, txt_path, preproc=None):
|
||||
self.preproc = preproc
|
||||
self.imgs_path = []
|
||||
self.words = []
|
||||
f = open(txt_path,'r')
|
||||
lines = f.readlines()
|
||||
isFirst = True
|
||||
labels = []
|
||||
for line in lines:
|
||||
line = line.rstrip()
|
||||
if line.startswith('#'):
|
||||
if isFirst is True:
|
||||
isFirst = False
|
||||
else:
|
||||
labels_copy = labels.copy()
|
||||
self.words.append(labels_copy)
|
||||
labels.clear()
|
||||
path = line[2:]
|
||||
path = txt_path.replace('label.txt','images/') + path
|
||||
self.imgs_path.append(path)
|
||||
else:
|
||||
line = line.split(' ')
|
||||
label = [float(x) for x in line]
|
||||
labels.append(label)
|
||||
|
||||
self.words.append(labels)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.imgs_path)
|
||||
|
||||
def __getitem__(self, index):
|
||||
img = cv2.imread(self.imgs_path[index])
|
||||
height, width, _ = img.shape
|
||||
|
||||
labels = self.words[index]
|
||||
annotations = np.zeros((0, 15))
|
||||
if len(labels) == 0:
|
||||
return annotations
|
||||
for idx, label in enumerate(labels):
|
||||
annotation = np.zeros((1, 15))
|
||||
# bbox
|
||||
annotation[0, 0] = label[0] # x1
|
||||
annotation[0, 1] = label[1] # y1
|
||||
annotation[0, 2] = label[0] + label[2] # x2
|
||||
annotation[0, 3] = label[1] + label[3] # y2
|
||||
|
||||
# landmarks
|
||||
annotation[0, 4] = label[4] # l0_x
|
||||
annotation[0, 5] = label[5] # l0_y
|
||||
annotation[0, 6] = label[7] # l1_x
|
||||
annotation[0, 7] = label[8] # l1_y
|
||||
annotation[0, 8] = label[10] # l2_x
|
||||
annotation[0, 9] = label[11] # l2_y
|
||||
annotation[0, 10] = label[13] # l3_x
|
||||
annotation[0, 11] = label[14] # l3_y
|
||||
annotation[0, 12] = label[16] # l4_x
|
||||
annotation[0, 13] = label[17] # l4_y
|
||||
if (annotation[0, 4]<0):
|
||||
annotation[0, 14] = -1
|
||||
else:
|
||||
annotation[0, 14] = 1
|
||||
|
||||
annotations = np.append(annotations, annotation, axis=0)
|
||||
target = np.array(annotations)
|
||||
if self.preproc is not None:
|
||||
img, target = self.preproc(img, target)
|
||||
|
||||
return torch.from_numpy(img), target
|
||||
|
||||
def detection_collate(batch):
|
||||
"""Custom collate fn for dealing with batches of images that have a different
|
||||
number of associated object annotations (bounding boxes).
|
||||
|
||||
Arguments:
|
||||
batch: (tuple) A tuple of tensor images and lists of annotations
|
||||
|
||||
Return:
|
||||
A tuple containing:
|
||||
1) (tensor) batch of images stacked on their 0 dim
|
||||
2) (list of tensors) annotations for a given image are stacked on 0 dim
|
||||
"""
|
||||
targets = []
|
||||
imgs = []
|
||||
for _, sample in enumerate(batch):
|
||||
for _, tup in enumerate(sample):
|
||||
if torch.is_tensor(tup):
|
||||
imgs.append(tup)
|
||||
elif isinstance(tup, type(np.empty(0))):
|
||||
annos = torch.from_numpy(tup).float()
|
||||
targets.append(annos)
|
||||
|
||||
return (torch.stack(imgs, 0), targets)
|
||||
@@ -4,6 +4,14 @@ from uuid import uuid4
|
||||
import imghdr
|
||||
import traceback
|
||||
import torch
|
||||
# PyTorch 2.6+ 默认 weights_only=True，无法加载含 numpy 对象的旧权重
|
||||
# （yolov5l.pt、各 *.pth 等）。统一改回 False（权重均为本机自有可信文件）。
|
||||
_orig_torch_load = torch.load
|
||||
def _torch_load(*args, **kwargs):
|
||||
kwargs.setdefault("weights_only", False)
|
||||
return _orig_torch_load(*args, **kwargs)
|
||||
torch.load = _torch_load
|
||||
|
||||
from gevent import monkey
|
||||
monkey.patch_all()
|
||||
|
||||
@@ -527,6 +535,27 @@ def change_hairstyle_v4():
|
||||
crop_result = landmark_processor.high_quality_warpAffine(img_res, M, dst_size)
|
||||
# crop_result = landmark_processor.high_quality_warpAffine(img_res, M, (w, h))
|
||||
|
||||
# 接口11：可选外部遮罩 ext_mask（接口9 头发遮罩，原图坐标、白=生发区）。
|
||||
# 传入则经 M 变换到 crop 坐标后替换内部 matting 合成遮罩，webui 精确重绘该区域；
|
||||
# 不传则保持原 matting_merge 逻辑（向后兼容，其它功能不受影响）。
|
||||
ext_mask_b64 = input_info.get('ext_mask', '')
|
||||
if ext_mask_b64:
|
||||
if 'data:image/' in ext_mask_b64:
|
||||
ext_mask_b64 = ext_mask_b64.split(',')[1]
|
||||
ext_mask_np = cv2.imdecode(
|
||||
np.frombuffer(base64.b64decode(ext_mask_b64), np.uint8),
|
||||
cv2.IMREAD_GRAYSCALE)
|
||||
if ext_mask_np is None:
|
||||
raise RuntimeError("ext_mask 解析失败")
|
||||
if ext_mask_np.shape[:2] != origin_img.shape[:2]:
|
||||
ext_mask_np = cv2.resize(
|
||||
ext_mask_np, (origin_img.shape[1], origin_img.shape[0]),
|
||||
interpolation=cv2.INTER_NEAREST)
|
||||
print(f"[swapHair] 使用外部遮罩 ext_mask, shape={ext_mask_np.shape}, "
|
||||
f"白像素={int((ext_mask_np > 10).sum())}", flush=True)
|
||||
crop_matting = cv2.warpAffine(ext_mask_np, M, dst_size)
|
||||
mask = (crop_matting > 10).astype(np.float32)
|
||||
else:
|
||||
origin_matting = cv2.imread(user_orig_mask_path, cv2.IMREAD_GRAYSCALE)
|
||||
new_matting = cv2.imread(hair_matting_path, cv2.IMREAD_GRAYSCALE)
|
||||
left_point = user_landmarks_origin_img_137[14] # (x1, y1) 获取刘海区域的圆心和半径
|
||||
@@ -540,7 +569,6 @@ def change_hairstyle_v4():
|
||||
no_bang_result = cv2.subtract(origin_matting, part_bangs) # 去除刘海区域
|
||||
matting_merge = np.max( np.stack([no_bang_result, new_matting], axis=2), axis=2).astype(np.uint8)# 合并 origin_matting 和 new_matting
|
||||
crop_matting = cv2.warpAffine(matting_merge, M, dst_size)
|
||||
# crop_matting = cv2.warpAffine(matting_merge, M, (w, h))
|
||||
mask = (crop_matting > 10).astype(np.float32)
|
||||
if not is_hr:
|
||||
mask_dilate = cv2.dilate(mask, np.ones((3, 9), np.uint8))
|
||||
@@ -576,7 +604,7 @@ def change_hairstyle_v4():
|
||||
p_tag = p_tag[p_tag.find("simple background, ") + len("simple background, "):]
|
||||
else:
|
||||
p_tag = ""
|
||||
denoising_strength = 0.6
|
||||
denoising_strength = float(input_info.get('denoising_strength', 0.6)) # 接口11 可调；默认 0.6
|
||||
print(f"功能8：处理发型区域，耗时：{time.time() - start_time:.3f}s")
|
||||
# cv2.imwrite(f"{task_id}_mask_dilate.jpg", mask_dilate)
|
||||
# cv2.imwrite(f"{task_id}_final_img.jpg", final_img)
|
||||
|
||||
@@ -5,16 +5,24 @@ import os
|
||||
|
||||
class OSS_object():
|
||||
def __init__(self):
|
||||
access_key_id = os.getenv('OSS_TEST_ACCESS_KEY_ID', '<your-access-key-id>')
|
||||
access_key_secret = os.getenv('OSS_TEST_ACCESS_KEY_SECRET', '<your-access-key-secret>')
|
||||
bucket_name = os.getenv('OSS_TEST_BUCKET', '<your-bucket-name>')
|
||||
endpoint = os.getenv('OSS_TEST_ENDPOINT', '<your-endpoint>')
|
||||
for param in (access_key_id, access_key_secret, bucket_name, endpoint):
|
||||
assert '<' not in param, '请设置环境变量 OSS_TEST_ACCESS_KEY_ID / OSS_TEST_ACCESS_KEY_SECRET / OSS_TEST_BUCKET / OSS_TEST_ENDPOINT'
|
||||
# 懒加载：只读取凭证，不立即连接 OSS。
|
||||
# 本地用 output_format=base64 时不配 OSS 凭证也能启动服务；
|
||||
# 仅在真正调用 upload_file 时才校验并连接。
|
||||
self.access_key_id = os.getenv('OSS_TEST_ACCESS_KEY_ID', '<your-access-key-id>')
|
||||
self.access_key_secret = os.getenv('OSS_TEST_ACCESS_KEY_SECRET', '<your-access-key-secret>')
|
||||
self.bucket_name = os.getenv('OSS_TEST_BUCKET', '<your-bucket-name>')
|
||||
self.endpoint = os.getenv('OSS_TEST_ENDPOINT', '<your-endpoint>')
|
||||
self.bucket = None
|
||||
|
||||
self.bucket = oss2.Bucket(oss2.Auth(access_key_id, access_key_secret), endpoint, bucket_name)
|
||||
def _ensure_bucket(self):
|
||||
if self.bucket is not None:
|
||||
return
|
||||
for param in (self.access_key_id, self.access_key_secret, self.bucket_name, self.endpoint):
|
||||
assert '<' not in param, '请设置环境变量 OSS_TEST_ACCESS_KEY_ID / OSS_TEST_ACCESS_KEY_SECRET / OSS_TEST_BUCKET / OSS_TEST_ENDPOINT'
|
||||
self.bucket = oss2.Bucket(oss2.Auth(self.access_key_id, self.access_key_secret), self.endpoint, self.bucket_name)
|
||||
|
||||
def upload_file(self, file, target_name):
|
||||
self._ensure_bucket()
|
||||
t0 = time.time()
|
||||
with open(oss2.to_unicode(file), 'rb') as f:
|
||||
ret = self.bucket.put_object(target_name, f)
|
||||
|
||||
@@ -324,7 +324,7 @@ def train_thread(sq, gpu_id):
|
||||
'--caption_extension=".txt" --sample_sampler=ddim '
|
||||
f'--sample_prompts={sample_txt} --sample_every_n_epochs="1000" '
|
||||
'--seed="1234" --cache_latents --optimizer_type="AdamW" --max_data_loader_n_workers="0" --bucket_reso_steps=64 '
|
||||
'--xformers --bucket_no_upscale --noise_offset=0.0')
|
||||
'--sdpa --bucket_no_upscale --noise_offset=0.0')
|
||||
print("cmd_train:", cmd_train)
|
||||
os.system(cmd_train)
|
||||
|
||||
|
||||
Executable
+69
@@ -0,0 +1,69 @@
|
||||
#!/usr/bin/env bash
|
||||
# 断点续传大文件到云服务器 /home/xsl/data
|
||||
# 用法: bash scripts/sync_data_to_server.sh [all|weights|sd|hairstyles]
|
||||
set -euo pipefail
|
||||
|
||||
REMOTE="ubuntu@117.50.213.111"
|
||||
REMOTE_BASE="/home/xsl/data"
|
||||
LOCAL_BASE="/home/xsl/change_hair"
|
||||
LOG_DIR="${LOCAL_BASE}/project/logs"
|
||||
mkdir -p "$LOG_DIR"
|
||||
|
||||
RSYNC_OPTS=(-avP --partial --info=progress2)
|
||||
|
||||
sync_weights() {
|
||||
echo ">>> [1/3] 同步 hair_service_sd/weights (~10G) ..."
|
||||
rsync "${RSYNC_OPTS[@]}" \
|
||||
"${LOCAL_BASE}/project/hair_service_sd/weights/" \
|
||||
"${REMOTE}:${REMOTE_BASE}/hair_service_sd/weights/"
|
||||
}
|
||||
|
||||
sync_sd_models() {
|
||||
echo ">>> [2/3] 同步 SD 底模 + ESRGAN + GFPGAN (~4.6G) ..."
|
||||
rsync "${RSYNC_OPTS[@]}" \
|
||||
"${LOCAL_BASE}/project/onediff/stable-diffusion-webui/models/Stable-diffusion/v1-5-pruned-emaonly.safetensors" \
|
||||
"${REMOTE}:${REMOTE_BASE}/sdwebui/models/Stable-diffusion/"
|
||||
rsync "${RSYNC_OPTS[@]}" \
|
||||
"${LOCAL_BASE}/project/onediff/stable-diffusion-webui/models/ESRGAN/" \
|
||||
"${REMOTE}:${REMOTE_BASE}/sdwebui/models/ESRGAN/"
|
||||
rsync "${RSYNC_OPTS[@]}" \
|
||||
"${LOCAL_BASE}/project/onediff/stable-diffusion-webui/models/GFPGAN/" \
|
||||
"${REMOTE}:${REMOTE_BASE}/sdwebui/models/GFPGAN/"
|
||||
}
|
||||
|
||||
sync_hairstyles() {
|
||||
echo ">>> [3/3] 同步 5 个 chang_* 发型资源 (~0.8G) ..."
|
||||
for id in chang_tuoyuan chang_bolang chang_zhixian chang_huaban chang_xinxing; do
|
||||
echo " - $id"
|
||||
rsync "${RSYNC_OPTS[@]}" \
|
||||
"${LOCAL_BASE}/project/data/ref_hairstyle/${id}/" \
|
||||
"${REMOTE}:${REMOTE_BASE}/project/data/ref_hairstyle/${id}/"
|
||||
rsync "${RSYNC_OPTS[@]}" \
|
||||
"${LOCAL_BASE}/project/data/hair_template_material/${id}/" \
|
||||
"${REMOTE}:${REMOTE_BASE}/project/data/hair_template_material/${id}/"
|
||||
rsync "${RSYNC_OPTS[@]}" \
|
||||
"${LOCAL_BASE}/project/data/upload_train_imgs/${id}/" \
|
||||
"${REMOTE}:${REMOTE_BASE}/project/data/upload_train_imgs/${id}/"
|
||||
rsync "${RSYNC_OPTS[@]}" \
|
||||
"${LOCAL_BASE}/data/train_material/${id}/" \
|
||||
"${REMOTE}:${REMOTE_BASE}/train_material/${id}/"
|
||||
done
|
||||
}
|
||||
|
||||
TARGET="${1:-all}"
|
||||
case "$TARGET" in
|
||||
weights) sync_weights ;;
|
||||
sd) sync_sd_models ;;
|
||||
hairstyles) sync_hairstyles ;;
|
||||
all)
|
||||
sync_weights
|
||||
sync_sd_models
|
||||
sync_hairstyles
|
||||
;;
|
||||
*)
|
||||
echo "用法: $0 [all|weights|sd|hairstyles]"
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
|
||||
echo ">>> 同步完成: ${TARGET}"
|
||||
@@ -0,0 +1,96 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""阶段C：为已训练的发型生成 ref材质 + 模板材质 + 预览图
|
||||
|
||||
前置：LoRA 已训练完成，webui/photo_service/hair_service_sd 已启动。
|
||||
|
||||
对每个发型依次：
|
||||
step3: 调 hair_service_sd trainCallBack 回调生成 ref 材质
|
||||
step4: 生成模板材质
|
||||
step5: 生成预览图
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import json
|
||||
import shutil
|
||||
|
||||
HAIR_SERVICE_DIR = "/home/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")
|
||||
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
|
||||
os.environ.setdefault("CRYPTOGRAPHY_OPENSSL_NO_LEGACY", "1")
|
||||
|
||||
import requests as req
|
||||
|
||||
# (hair_id, gender, template_img)
|
||||
HAIRSTYLES = [
|
||||
("chang_tuoyuan", "girl", "/home/xsl/change_hair/hair_type_images/chang_tuoyuan/chang_tuoyuan.jpg"),
|
||||
("chang_bolang", "girl", "/home/xsl/change_hair/hair_type_images/chang_bolang/chang_bolang.jpg"),
|
||||
("chang_zhixian", "girl", "/home/xsl/change_hair/hair_type_images/chang_zhixian/chang_zhixian.jpg"),
|
||||
("chang_huaban", "girl", "/home/xsl/change_hair/hair_type_images/chang_huaban/chang_huaban.jpg"),
|
||||
("chang_xinxing", "girl", "/home/xsl/change_hair/hair_type_images/chang_xinxing/chang_xinxing.jpg"),
|
||||
]
|
||||
|
||||
HAIR_CALLBACK = "http://127.0.0.1:8801/api/hair/trainCallBack"
|
||||
UPLOAD_TRAIN_DIR = "/home/xsl/change_hair/project/data/upload_train_imgs"
|
||||
|
||||
# 导入 step4/step5 函数（会触发模型加载）
|
||||
from train_hairstyle_full import step4_template_material, step5_preview
|
||||
from common.logger import config
|
||||
|
||||
|
||||
def log(msg):
|
||||
print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)
|
||||
|
||||
|
||||
def step3_callback(hair_id, template_img):
|
||||
"""调 trainCallBack 生成 ref 材质。需要先把模板图放到 upload_train_imgs"""
|
||||
log(f" --- step3: {hair_id} 生成 ref 材质 ---")
|
||||
# 准备 upload_train_imgs/<hair_id>/first##<hair_id>.jpg
|
||||
src_dir = os.path.join(UPLOAD_TRAIN_DIR, hair_id)
|
||||
os.makedirs(src_dir, exist_ok=True)
|
||||
dst = os.path.join(src_dir, f"first##{hair_id}.jpg")
|
||||
shutil.copy(template_img, dst)
|
||||
|
||||
try:
|
||||
r = req.post(HAIR_CALLBACK, json={
|
||||
"task_id": f"train_{hair_id}", "hair_id": hair_id,
|
||||
"state": 0, "msg": "头发lora训练成功", "is_tj": "1"
|
||||
}, timeout=300)
|
||||
d = r.json()
|
||||
ref_dir = os.path.join(config.get('default', 'hairstyleDir'), hair_id)
|
||||
ok = d.get("state") == 0 and os.path.exists(os.path.join(ref_dir, "ref_rgb_8uc3_768.png"))
|
||||
log(f" {'✓' if ok else '✗'} {hair_id} ref材质 {'成功' if ok else '失败: ' + str(d.get('msg'))}")
|
||||
return ok
|
||||
except Exception as e:
|
||||
log(f" ✗ {hair_id} ref材质异常: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def main():
|
||||
t0 = time.time()
|
||||
log(f"阶段C: 为 {len(HAIRSTYLES)} 个发型生成材质和预览")
|
||||
|
||||
results = {}
|
||||
for hair_id, gender, template_img in HAIRSTYLES:
|
||||
log(f"\n{'='*50}")
|
||||
log(f"处理: {hair_id}")
|
||||
log(f"{'='*50}")
|
||||
ok3 = step3_callback(hair_id, template_img)
|
||||
ok4 = step4_template_material(hair_id, template_img)
|
||||
ok5 = step5_preview(hair_id, gender)
|
||||
results[hair_id] = (ok3, ok4, ok5)
|
||||
log(f" {hair_id}: ref={'✓' if ok3 else '✗'} 模板={'✓' if ok4 else '✗'} 预览={'✓' if ok5 else '✗'}")
|
||||
|
||||
log(f"\n{'='*60}")
|
||||
log(f"阶段C完成，耗时 {time.time()-t0:.0f}s")
|
||||
log(f"{'='*60}")
|
||||
for hair_id, (ok3, ok4, ok5) in results.items():
|
||||
status = "✓全部成功" if (ok3 and ok4 and ok5) else f"⚠️ ref={'✓' if ok3 else '✗'} 模板={'✓' if ok4 else '✗'} 预览={'✓' if ok5 else '✗'}"
|
||||
log(f" {hair_id}: {status}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,150 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""批量并行训练多个发型
|
||||
|
||||
流程：
|
||||
阶段A: 串行 step1（准备训练数据，共用GPU模型，避免重复加载）
|
||||
阶段B: 并行 step2（LoRA训练，kohya独立进程，限制并发数）
|
||||
阶段C: 启动服务后串行 step3(回调生成ref材质) + step4(模板) + step5(预览)
|
||||
|
||||
用法:
|
||||
cd /home/xsl/change_hair/project/hair_service_sd
|
||||
python /home/xsl/change_hair/train_hairstyles_parallel.py
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import json
|
||||
import subprocess
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
HAIR_SERVICE_DIR = "/home/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")
|
||||
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
|
||||
os.environ.setdefault("CRYPTOGRAPHY_OPENSSL_NO_LEGACY", "1")
|
||||
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0")
|
||||
|
||||
# ====== 配置：本次要训练的发型 ======
|
||||
# (hair_id, input_dir, gender, template_img)
|
||||
HAIRSTYLES = [
|
||||
("chang_tuoyuan", "/home/xsl/change_hair/hair_type_images/chang_tuoyuan", "girl", "/home/xsl/change_hair/hair_type_images/chang_tuoyuan/chang_tuoyuan.jpg"),
|
||||
("chang_bolang", "/home/xsl/change_hair/hair_type_images/chang_bolang", "girl", "/home/xsl/change_hair/hair_type_images/chang_bolang/chang_bolang.jpg"),
|
||||
("chang_zhixian", "/home/xsl/change_hair/hair_type_images/chang_zhixian", "girl", "/home/xsl/change_hair/hair_type_images/chang_zhixian/chang_zhixian.jpg"),
|
||||
("chang_huaban", "/home/xsl/change_hair/hair_type_images/chang_huaban", "girl", "/home/xsl/change_hair/hair_type_images/chang_huaban/chang_huaban.jpg"),
|
||||
("chang_xinxing", "/home/xsl/change_hair/hair_type_images/chang_xinxing", "girl", "/home/xsl/change_hair/hair_type_images/chang_xinxing/chang_xinxing.jpg"),
|
||||
]
|
||||
PARALLEL = 2 # LoRA 训练并发数
|
||||
|
||||
# 导入训练模块（复用其 step 函数）
|
||||
from train_hairstyle_full import (
|
||||
step1_prepare_data, step2_train, step3_wait_and_callback,
|
||||
step4_template_material, step5_preview, get_models
|
||||
)
|
||||
import train_hairstyle_full as TH
|
||||
from common.logger import config
|
||||
|
||||
|
||||
def log(msg):
|
||||
print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)
|
||||
|
||||
|
||||
def main():
|
||||
t0 = time.time()
|
||||
ids = [h[0] for h in HAIRSTYLES]
|
||||
log(f"批量训练 {len(ids)} 个发型: {ids} (并发={PARALLEL})")
|
||||
|
||||
# ============ 阶段A: 串行 step1（准备数据）============
|
||||
log("=" * 60)
|
||||
log("阶段A: 串行准备训练数据 (step1)")
|
||||
log("=" * 60)
|
||||
log(" 预加载模型...")
|
||||
get_models() # 预加载，5个发型共用
|
||||
step1_ok, step1_fail = [], []
|
||||
for hair_id, input_dir, gender, _ in HAIRSTYLES:
|
||||
# 清理可能的历史训练数据，确保干净
|
||||
train_dir = config.get('default', 'train_dir')
|
||||
out_dir = os.path.join(train_dir, hair_id, "images", "1_hairstyle")
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
os.makedirs(os.path.join(train_dir, hair_id, "model"), exist_ok=True)
|
||||
|
||||
log(f" --- step1: {hair_id} ---")
|
||||
try:
|
||||
ok = step1_prepare_data(hair_id, input_dir, gender)
|
||||
if ok:
|
||||
step1_ok.append(hair_id)
|
||||
log(f" ✓ {hair_id} 数据准备完成")
|
||||
else:
|
||||
step1_fail.append(hair_id)
|
||||
log(f" ✗ {hair_id} 数据准备失败")
|
||||
except Exception as e:
|
||||
step1_fail.append(hair_id)
|
||||
log(f" ✗ {hair_id} 数据准备异常: {e}")
|
||||
log(f"阶段A完成: 成功 {len(step1_ok)}/{len(ids)}，失败 {step1_fail}")
|
||||
if not step1_ok:
|
||||
log("全部 step1 失败，退出"); return
|
||||
|
||||
# 释放 step1 用的 GPU 模型，给 step2 训练腾显存
|
||||
log("释放 step1 模型显存...")
|
||||
import torch
|
||||
try:
|
||||
for m in [TH._detector, TH._aligner, TH._matte]:
|
||||
if m is not None and hasattr(m, 'model'):
|
||||
pass
|
||||
torch.cuda.empty_cache()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ============ 阶段B: 并行 step2（LoRA 训练）============
|
||||
log("=" * 60)
|
||||
log(f"阶段B: 并行 LoRA 训练 (step2, 并发={PARALLEL})")
|
||||
log("=" * 60)
|
||||
step2_results = {} # hair_id -> bool
|
||||
with ThreadPoolExecutor(max_workers=PARALLEL) as ex:
|
||||
futures = {ex.submit(step2_train, hid): hid for hid in step1_ok}
|
||||
for fut in as_completed(futures):
|
||||
hid = futures[fut]
|
||||
try:
|
||||
ok = fut.result()
|
||||
step2_results[hid] = ok
|
||||
log(f" {'✓' if ok else '✗'} {hid} step2 (训练启动)={'成功' if ok else '失败'}")
|
||||
except Exception as e:
|
||||
step2_results[hid] = False
|
||||
log(f" ✗ {hid} step2 异常: {e}")
|
||||
|
||||
# 等待所有 LoRA 训练真正完成（step2 只是启动，训练在 photo_service 后台跑）
|
||||
log("等待所有 LoRA 训练完成（检查权重文件）...")
|
||||
train_dir = config.get('default', 'train_dir')
|
||||
pending = [hid for hid in step1_ok if step2_results.get(hid)]
|
||||
completed = set()
|
||||
deadline = time.time() + 3600 # 最多等1小时
|
||||
while pending and time.time() < deadline:
|
||||
still = []
|
||||
for hid in pending:
|
||||
lora = os.path.join(train_dir, hid, "model", "hairstyle_hd_lora.safetensors")
|
||||
if os.path.exists(lora) and os.path.getsize(lora) > 100000000:
|
||||
completed.add(hid)
|
||||
log(f" ✓ {hid} LoRA 训练完成 ({os.path.getsize(lora)//1024//1024}MB)")
|
||||
else:
|
||||
still.append(hid)
|
||||
pending = still
|
||||
if pending:
|
||||
log(f" ...等待中: {pending}")
|
||||
time.sleep(30)
|
||||
if pending:
|
||||
log(f"⚠️ 超时未完成: {pending}")
|
||||
log(f"阶段B完成: LoRA 训练完成 {len(completed)}/{len(step1_ok)}")
|
||||
|
||||
# 写一个完成清单文件，供阶段C脚本读取
|
||||
done_file = "/home/xsl/change_hair/project/logs/train_batch_done.json"
|
||||
with open(done_file, "w") as f:
|
||||
json.dump({"completed": sorted(completed), "all": ids}, f)
|
||||
log(f"已写入完成清单: {done_file}")
|
||||
log(f"\n批量训练阶段A+B总耗时 {time.time()-t0:.0f}s")
|
||||
log(f"完成训练的发型: {sorted(completed)}")
|
||||
log("接下来需启动服务跑阶段C（step3/4/5 生成材质和预览）")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,124 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""直接调 kohya train_network.py 真正并行训练 LoRA
|
||||
|
||||
绕过 photo_service 的串行阻塞，用 subprocess 直接起 train_network.py 进程，
|
||||
通过 ThreadPoolExecutor 控制并发数。
|
||||
|
||||
用法:
|
||||
python /home/xsl/change_hair/train_lora_parallel.py
|
||||
|
||||
前提：step1 训练数据已准备好（data/train_material/<hid>/images/1_hairstyle/*.png）
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import subprocess
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
TRAIN_DIR = "/home/xsl/change_hair/data/train_material"
|
||||
KOHYA_WORKDIR = "/home/xsl/change_hair/project/kohya_ss_home/kohya_ss"
|
||||
BASE_MODEL = "/home/xsl/change_hair/project/onediff/stable-diffusion-webui/models/Stable-diffusion/v1-5-pruned-emaonly.safetensors"
|
||||
KOHYA_ACCEL = "/home/xsl/miniconda3/envs/kohya/bin/accelerate"
|
||||
LOG_DIR = "/home/xsl/change_hair/project/logs"
|
||||
|
||||
# 本次要训练的发型
|
||||
HAIRSTYLES = ["chang_tuoyuan", "chang_zhixian", "chang_huaban", "chang_xinxing"]
|
||||
# chang_bolang 已完成，跳过
|
||||
PARALLEL = 2
|
||||
|
||||
|
||||
def log(msg):
|
||||
print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)
|
||||
|
||||
|
||||
def train_one(hair_id):
|
||||
"""训练单个发型的 LoRA，返回 (hair_id, success, log_path)"""
|
||||
images_dir = os.path.join(TRAIN_DIR, hair_id, "images")
|
||||
model_dir = os.path.join(TRAIN_DIR, hair_id, "model")
|
||||
os.makedirs(model_dir, exist_ok=True)
|
||||
|
||||
# 准备 sample prompt 文件
|
||||
sample_dir = os.path.join(model_dir, "sample")
|
||||
os.makedirs(sample_dir, exist_ok=True)
|
||||
sample_txt = os.path.join(sample_dir, "prompt.txt")
|
||||
with open(sample_txt, "w") as f:
|
||||
f.write('titor hairstyle, easyphoto, faceless, no human, white background, simple background, '
|
||||
' --n low quality, worst quality, bad anatomy, bad composition, poor, low effort --h 768 '
|
||||
'--w 768 --s 30 --l 7')
|
||||
|
||||
log_path = os.path.join(LOG_DIR, f"train_{hair_id}.log")
|
||||
|
||||
cmd = [
|
||||
"bash", "-c",
|
||||
f'cd {KOHYA_WORKDIR} && '
|
||||
f'HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 CUDA_VISIBLE_DEVICES=0 '
|
||||
f'{KOHYA_ACCEL} launch --num_cpu_threads_per_process=2 "./train_network.py" --enable_bucket '
|
||||
f'--min_bucket_reso=256 --max_bucket_reso=2048 --pretrained_model_name_or_path="{BASE_MODEL}" '
|
||||
f'--train_data_dir={images_dir} --resolution="2000,2000" '
|
||||
f'--output_dir={model_dir} '
|
||||
'--network_alpha="64" --save_model_as=safetensors --network_module=networks.lora --text_encoder_lr=5e-05 '
|
||||
'--unet_lr=0.0001 --network_dim=128 --output_name="hairstyle_hd_lora" --lr_scheduler_num_cycles="20" '
|
||||
'--no_half_vae --learning_rate="0.0001" --lr_scheduler="cosine" --lr_warmup_steps="650" --train_batch_size="1" '
|
||||
'--max_train_steps="1500" --save_every_n_epochs="100" --mixed_precision="fp16" --save_precision="fp16" '
|
||||
'--caption_extension=".txt" --sample_sampler=ddim '
|
||||
f'--sample_prompts={sample_txt} --sample_every_n_epochs="1000" '
|
||||
'--seed="1234" --cache_latents --optimizer_type="AdamW" --max_data_loader_n_workers="0" --bucket_reso_steps=64 '
|
||||
'--sdpa --bucket_no_upscale --noise_offset=0.0'
|
||||
]
|
||||
|
||||
log(f" ▶ {hair_id} 开始训练，日志: {log_path}")
|
||||
t0 = time.time()
|
||||
try:
|
||||
with open(log_path, "w") as lf:
|
||||
proc = subprocess.run(cmd, stdout=lf, stderr=subprocess.STDOUT, timeout=1800)
|
||||
ok = proc.returncode == 0
|
||||
lora = os.path.join(model_dir, "hairstyle_hd_lora.safetensors")
|
||||
ok = ok and os.path.exists(lora) and os.path.getsize(lora) > 100000000
|
||||
elapsed = time.time() - t0
|
||||
sz = os.path.getsize(lora) // 1024 // 1024 if os.path.exists(lora) else 0
|
||||
log(f" {'✓' if ok else '✗'} {hair_id} {'训练完成' if ok else '训练失败'} ({sz}MB, {elapsed:.0f}s)")
|
||||
return (hair_id, ok, log_path)
|
||||
except subprocess.TimeoutExpired:
|
||||
log(f" ✗ {hair_id} 训练超时")
|
||||
return (hair_id, False, log_path)
|
||||
except Exception as e:
|
||||
log(f" ✗ {hair_id} 训练异常: {e}")
|
||||
return (hair_id, False, log_path)
|
||||
|
||||
|
||||
def main():
|
||||
t0 = time.time()
|
||||
log(f"并行 LoRA 训练: {HAIRSTYLES} (并发={PARALLEL})")
|
||||
|
||||
results = {}
|
||||
with ThreadPoolExecutor(max_workers=PARALLEL) as ex:
|
||||
futures = {ex.submit(train_one, hid): hid for hid in HAIRSTYLES}
|
||||
for fut in as_completed(futures):
|
||||
hid = futures[fut]
|
||||
try:
|
||||
hair_id, ok, log_path = fut.result()
|
||||
results[hair_id] = ok
|
||||
except Exception as e:
|
||||
log(f" ✗ {hid} 异常: {e}")
|
||||
results[hid] = False
|
||||
|
||||
log(f"\n{'='*60}")
|
||||
log(f"全部完成，耗时 {time.time()-t0:.0f}s")
|
||||
for hid, ok in results.items():
|
||||
lora = os.path.join(TRAIN_DIR, hid, "model", "hairstyle_hd_lora.safetensors")
|
||||
sz = os.path.getsize(lora) // 1024 // 1024 if os.path.exists(lora) else 0
|
||||
log(f" {'✓' if ok else '✗'} {hid} ({sz}MB)")
|
||||
success = [h for h, ok in results.items() if ok]
|
||||
log(f"成功: {success}")
|
||||
|
||||
# 写完成清单
|
||||
done_file = os.path.join(LOG_DIR, "train_lora_done.json")
|
||||
import json
|
||||
with open(done_file, "w") as f:
|
||||
json.dump({"completed": success, "all": HAIRSTYLES}, f)
|
||||
log(f"完成清单: {done_file}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user