初始化换发型项目：3个微服务代码 + 部署脚本
包含： - hair_service_sd: 换发型/换发色算法服务 (端口 8801) - photo_service: LoRA 训练调度服务 (端口 32678) - stable-diffusion-webui: SD WebUI 推理服务 (端口 57860) - kohya_ss_home: 训练环境代码 - meidaojia: 监控测试脚本 - setup.sh: 一键部署脚本 (conda环境恢复 + 配置生成 + 完整性检查) - start_all_services.sh: 启动3个服务 - configure.ini.template: 路径模板化 (BASE_DIR自动推导) - conda_envs/py310.yml: py310 环境定义 大文件 (weights/, models/, data/, conda_envs/*.tar.gz 等) 通过 .gitignore 排除， 由网盘单独上传。
This commit is contained in:
colomi
+35
@@ -0,0 +1,35 @@
|
||||
# Python
|
||||
__pycache__/
|
||||
*.pyc
|
||||
*.pyo
|
||||
*.egg-info/
|
||||
.eggs/
|
||||
|
||||
# IDE
|
||||
.idea/
|
||||
.vscode/
|
||||
|
||||
# Logs
|
||||
*.log
|
||||
nohup.out
|
||||
/logs/
|
||||
|
||||
# 大文件目录（网盘上传）— 注意 /data/ 用前导斜杠仅匹配顶层
|
||||
hair_service_sd/weights/
|
||||
stable-diffusion-webui/models/
|
||||
stable-diffusion-webui/extensions/sd-webui-controlnet/
|
||||
stable-diffusion-webui/repositories/
|
||||
kohya_ss_home/.local/
|
||||
kohya_ss_home/.cache/
|
||||
kohya_ss_home/kohya_ss/
|
||||
/data/
|
||||
conda_envs/*.tar.gz
|
||||
docker/
|
||||
|
||||
# 运行时生成（setup.sh 生成）
|
||||
hair_service_sd/config/configure.ini
|
||||
*.pid
|
||||
|
||||
# 注意：hair_service_sd/data/ (19M 小静态文件) 不排除，需 git 管理
|
||||
# 注意：conda_envs/py310.yml 不排除，需 git 管理
|
||||
# 注意：kohya_ss_home/start_docker.sh 不排除，需 git 管理
|
||||
@@ -0,0 +1,330 @@
|
||||
# 换发型项目迁移收尾计划
|
||||
|
||||
## 背景
|
||||
|
||||
项目从 `/home/szlc/project` 迁移到 `/home/szlc/change_hair_3090`。前序工作已完成代码复制和大部分路径修改。本计划覆盖剩余的收尾工作。
|
||||
|
||||
## 当前状态
|
||||
|
||||
### 已完成
|
||||
- ✅ 代码复制：hair_service_sd (代码+weights 9.9G)、photo_service (1.9M)、stable-diffusion-webui (代码+extensions 4.4G+repositories 225M)、meidaojia (124K)、kohya_ss_home (kohya_ss 439M + .local 7.6G + .cache 3.3M)、data (4.9G)
|
||||
- ✅ configure.ini.template 创建 + configure.ini 生成
|
||||
- ✅ lora_train_service_1.py 修改（BASE_DIR + Docker 命令路径修正）
|
||||
- ✅ watch_delete.py 修改（从 config 读取路径）
|
||||
- ✅ gunicorn_config.py 修改（注释 chdir）
|
||||
- ✅ config.json 修改（删除 onediff 配置，改 sd_model_checkpoint）
|
||||
- ✅ git init（main 分支，无提交）
|
||||
|
||||
### 未完成（本计划覆盖）
|
||||
1. ❌ **stable-diffusion-webui/models/ 目录未复制**（41G，含 228 个 LoRA 33G + v1-5-pruned 4G + 辅助模型）
|
||||
2. ❌ **hairstyle_model_infer.py:1096 未注释**（活跃代码，换发色会崩溃）
|
||||
3. ❌ **.gitignore 未创建**
|
||||
4. ❌ **setup.sh 未创建**
|
||||
5. ❌ **start_all_services.sh 未创建**
|
||||
6. ❌ **README.md 未创建**
|
||||
7. ❌ **conda 环境未导出**（my_hair, sdwebui, py310）
|
||||
8. ❌ **git commit 未执行**
|
||||
|
||||
### 用户决策
|
||||
- Docker 镜像 `chinatszrn/ubuntu:kohya_ss` **不需要导出**（用户表示当前训练已不用 Docker）
|
||||
- majicmixRealistic_v7.safetensors 用户会从网上下载
|
||||
- onediff 不迁移
|
||||
|
||||
## 服务启动方式（从 .bash 脚本确认）
|
||||
|
||||
| 服务 | 端口 | conda 环境 | 启动命令 |
|
||||
|------|------|-----------|---------|
|
||||
| hair_service_sd | 8801 | my_hair | `python run_copy_cost_colorb64.py` |
|
||||
| stable-diffusion-webui | 57860 | sdwebui | `python webui.py --api --listen --xformers --port 57860` |
|
||||
| photo_service | 32678 | py310 | `python lora_train_service_1.py` |
|
||||
|
||||
---
|
||||
|
||||
## 执行步骤
|
||||
|
||||
### 步骤 1：复制 models/ 目录（41G）
|
||||
|
||||
源：`/home/szlc/project/onediff/stable-diffusion-webui/models/`
|
||||
目标：`/home/szlc/change_hair_3090/stable-diffusion-webui/models/`
|
||||
|
||||
```bash
|
||||
rsync -av --progress /home/szlc/project/onediff/stable-diffusion-webui/models/ /home/szlc/change_hair_3090/stable-diffusion-webui/models/
|
||||
```
|
||||
|
||||
包含：
|
||||
- Lora/ (33G, 228 个 LoRA 文件)
|
||||
- Stable-diffusion/ (4G, v1-5-pruned-emaonly.safetensors)
|
||||
- ESRGAN/ (1.2G)、BLIP/ (855M)、torch_deepdanbooru/ (615M)、roop/ (529M)、GFPGAN/ (519M)、Codeformer/ (360M)、RealESRGAN/ (64M) 等
|
||||
|
||||
注意：majicmixRealistic_v7.safetensors 不在源目录中，用户会单独下载。
|
||||
|
||||
### 步骤 2：注释 hairstyle_model_infer.py:1096
|
||||
|
||||
文件：`hair_service_sd/hairstyle_model_infer.py`
|
||||
|
||||
第 1096 行：
|
||||
```python
|
||||
# 修改前
|
||||
cv2.imwrite('/home/student/Desktop/tmp_color/need/face_base2.png', user_res_8uc3_orisize2)
|
||||
# 修改后
|
||||
# cv2.imwrite('/home/student/Desktop/tmp_color/need/face_base2.png', user_res_8uc3_orisize2) # 调试用，路径不存在会导致崩溃
|
||||
```
|
||||
|
||||
### 步骤 3：创建 .gitignore
|
||||
|
||||
文件：`/home/szlc/change_hair_3090/.gitignore`
|
||||
|
||||
```gitignore
|
||||
# Python
|
||||
__pycache__/
|
||||
*.pyc
|
||||
*.pyo
|
||||
*.egg-info/
|
||||
.eggs/
|
||||
|
||||
# IDE
|
||||
.idea/
|
||||
.vscode/
|
||||
|
||||
# Logs
|
||||
*.log
|
||||
nohup.out
|
||||
/logs/
|
||||
|
||||
# 大文件目录（网盘上传）— 注意 /data/ 用前导斜杠仅匹配顶层
|
||||
hair_service_sd/weights/
|
||||
stable-diffusion-webui/models/
|
||||
stable-diffusion-webui/extensions/sd-webui-controlnet/
|
||||
stable-diffusion-webui/repositories/
|
||||
kohya_ss_home/.local/
|
||||
kohya_ss_home/.cache/
|
||||
kohya_ss_home/kohya_ss/
|
||||
/data/
|
||||
conda_envs/*.tar.gz
|
||||
docker/
|
||||
|
||||
# 运行时生成（setup.sh 生成）
|
||||
hair_service_sd/config/configure.ini
|
||||
*.pid
|
||||
|
||||
# 注意：hair_service_sd/data/ (19M 小静态文件) 不排除，需 git 管理
|
||||
# 注意：conda_envs/py310.yml 不排除，需 git 管理
|
||||
# 注意：kohya_ss_home/start_docker.sh 不排除，需 git 管理
|
||||
```
|
||||
|
||||
### 步骤 4：创建 setup.sh
|
||||
|
||||
文件：`/home/szlc/change_hair_3090/setup.sh`
|
||||
|
||||
功能：
|
||||
1. 检查前提条件（conda, git, nvidia-docker）
|
||||
2. 从 template 生成 configure.ini（sed 替换 `__BASE_DIR__`）
|
||||
3. 用 conda-pack 恢复 my_hair 和 sdwebui 环境
|
||||
4. 用 yml 创建 py310 环境
|
||||
5. 检查模型/数据目录完整性
|
||||
6. 检查 majicmixRealistic_v7.safetensors 是否存在
|
||||
7. 创建运行时目录（tmp, res_dir, userImage 等）
|
||||
|
||||
```bash
|
||||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
BASE_DIR="$(cd "$(dirname "$0")" && pwd)"
|
||||
CONDA_BASE="${CONDA_BASE:-/home/szlc/miniconda3}"
|
||||
|
||||
# 初始化 conda（脚本中需要 source 才能用 conda activate）
|
||||
if [ -f "$CONDA_BASE/etc/profile.d/conda.sh" ]; then
|
||||
source "$CONDA_BASE/etc/profile.d/conda.sh"
|
||||
else
|
||||
echo "WARNING: 未找到 conda.sh，请确认 CONDA_BASE=$CONDA_BASE 正确"
|
||||
fi
|
||||
|
||||
echo "=== 换发型项目部署脚本 ==="
|
||||
echo "BASE_DIR: $BASE_DIR"
|
||||
|
||||
# 1. 检查前提条件
|
||||
echo "[1/7] 检查前提条件..."
|
||||
command -v git >/dev/null || { echo "ERROR: git 未安装"; exit 1; }
|
||||
command -v conda >/dev/null || { echo "ERROR: conda 未安装"; exit 1; }
|
||||
nvidia-smi >/dev/null 2>&1 || { echo "ERROR: NVIDIA 驱动未安装"; exit 1; }
|
||||
|
||||
# 2. 生成 configure.ini
|
||||
echo "[2/7] 生成 configure.ini..."
|
||||
sed "s|__BASE_DIR__|$BASE_DIR|g" "$BASE_DIR/hair_service_sd/config/configure.ini.template" > "$BASE_DIR/hair_service_sd/config/configure.ini"
|
||||
echo " configure.ini 已生成"
|
||||
|
||||
# 3. 恢复 conda 环境（conda-pack）
|
||||
echo "[3/7] 恢复 conda 环境..."
|
||||
for env_name in my_hair sdwebui; do
|
||||
if [ -f "$BASE_DIR/conda_envs/${env_name}.tar.gz" ]; then
|
||||
echo " 恢复 $env_name ..."
|
||||
rm -rf "$CONDA_BASE/envs/$env_name"
|
||||
mkdir -p "$CONDA_BASE/envs/$env_name"
|
||||
tar -xzf "$BASE_DIR/conda_envs/${env_name}.tar.gz" -C "$CONDA_BASE/envs/$env_name"
|
||||
conda activate "$env_name"
|
||||
conda-unpack # 修复打包后的硬编码路径
|
||||
conda deactivate
|
||||
echo " ✓ $env_name 已恢复"
|
||||
else
|
||||
echo " ✗ conda_envs/${env_name}.tar.gz 不存在，请从网盘下载"
|
||||
fi
|
||||
done
|
||||
|
||||
# 4. 创建 py310 环境（从 yml）
|
||||
echo "[4/7] 创建 py310 环境..."
|
||||
if [ -f "$BASE_DIR/conda_envs/py310.yml" ]; then
|
||||
conda env create -f "$BASE_DIR/conda_envs/py310.yml" -n py310 2>/dev/null || echo " py310 环境已存在，跳过"
|
||||
echo " ✓ py310 已就绪"
|
||||
else
|
||||
echo " ✗ conda_envs/py310.yml 不存在"
|
||||
fi
|
||||
|
||||
# 5. 检查模型/数据目录
|
||||
echo "[5/7] 检查模型和数据目录..."
|
||||
for dir in hair_service_sd/weights stable-diffusion-webui/models/Lora stable-diffusion-webui/models/Stable-diffusion stable-diffusion-webui/extensions/sd-webui-controlnet stable-diffusion-webui/repositories kohya_ss_home/.local kohya_ss_home/kohya_ss data/ref_hairstyle; do
|
||||
if [ -d "$BASE_DIR/$dir" ]; then
|
||||
echo " ✓ $dir"
|
||||
else
|
||||
echo " ✗ $dir 缺失，请从网盘下载"
|
||||
fi
|
||||
done
|
||||
|
||||
# 6. 检查训练底模
|
||||
echo "[6/7] 检查 majicmixRealistic_v7..."
|
||||
if [ -f "$BASE_DIR/stable-diffusion-webui/models/Stable-diffusion/majicmixRealistic_v7.safetensors" ]; then
|
||||
echo " ✓ majicmixRealistic_v7.safetensors 已存在"
|
||||
else
|
||||
echo " ✗ majicmixRealistic_v7.safetensors 缺失，请从网上下载并放置到 stable-diffusion-webui/models/Stable-diffusion/"
|
||||
fi
|
||||
|
||||
# 7. 创建运行时目录
|
||||
echo "[7/7] 创建运行时目录..."
|
||||
mkdir -p "$BASE_DIR/logs"
|
||||
mkdir -p "$BASE_DIR/data/tmp" "$BASE_DIR/data/res_dir" "$BASE_DIR/data/userImage" "$BASE_DIR/data/user_info"
|
||||
mkdir -p "$BASE_DIR/kohya_ss_home/train_material"
|
||||
|
||||
echo ""
|
||||
echo "=== 部署完成 ==="
|
||||
echo "启动服务: ./start_all_services.sh"
|
||||
```
|
||||
|
||||
### 步骤 5：创建 start_all_services.sh
|
||||
|
||||
文件：`/home/szlc/change_hair_3090/start_all_services.sh`
|
||||
|
||||
```bash
|
||||
#!/bin/bash
|
||||
BASE_DIR="$(cd "$(dirname "$0")" && pwd)"
|
||||
CONDA_BASE="${CONDA_BASE:-/home/szlc/miniconda3}"
|
||||
|
||||
# 1. hair_service_sd (端口 8801)
|
||||
cd "$BASE_DIR/hair_service_sd"
|
||||
nohup "$CONDA_BASE/envs/my_hair/bin/python" run_copy_cost_colorb64.py > "$BASE_DIR/logs/hair_service.log" 2>&1 &
|
||||
echo "hair_service_sd 已启动 (PID: $!)"
|
||||
|
||||
# 2. stable-diffusion-webui (端口 57860)
|
||||
cd "$BASE_DIR/stable-diffusion-webui"
|
||||
nohup "$CONDA_BASE/envs/sdwebui/bin/python" webui.py --api --listen --xformers --port 57860 > "$BASE_DIR/logs/webui.log" 2>&1 &
|
||||
echo "stable-diffusion-webui 已启动 (PID: $!)"
|
||||
|
||||
# 3. photo_service (端口 32678)
|
||||
cd "$BASE_DIR/photo_service"
|
||||
nohup "$CONDA_BASE/envs/py310/bin/python" lora_train_service_1.py > "$BASE_DIR/logs/photo_service.log" 2>&1 &
|
||||
echo "photo_service 已启动 (PID: $!)"
|
||||
|
||||
echo "所有服务已在后台启动，日志在 $BASE_DIR/logs/"
|
||||
```
|
||||
|
||||
### 步骤 6：创建 README.md
|
||||
|
||||
文件：`/home/szlc/change_hair_3090/README.md`
|
||||
|
||||
内容包含：
|
||||
- 项目简介（3 个微服务）
|
||||
- 目录结构
|
||||
- 前提条件（Ubuntu 22.04, 3090 GPU, conda, git）
|
||||
- 部署步骤（git clone → 下载网盘文件 → ./setup.sh → ./start_all_services.sh）
|
||||
- 网盘文件清单（models/, weights/, data/, conda_envs/, etc.）
|
||||
- 服务端口说明
|
||||
- 常见问题
|
||||
|
||||
### 步骤 7：导出 conda 环境
|
||||
|
||||
```bash
|
||||
# 安装 conda-pack
|
||||
conda install -n base -c conda-forge conda-pack -y
|
||||
|
||||
# 导出 my_hair (10G → ~5G 压缩)
|
||||
conda pack -n my_hair -o /home/szlc/change_hair_3090/conda_envs/my_hair.tar.gz --force
|
||||
|
||||
# 导出 sdwebui (7.3G → ~4G 压缩)
|
||||
conda pack -n sdwebui -o /home/szlc/change_hair_3090/conda_envs/sdwebui.tar.gz --force
|
||||
|
||||
# 导出 py310 (yml)
|
||||
conda env export -n py310 --no-builds > /home/szlc/change_hair_3090/conda_envs/py310.yml
|
||||
```
|
||||
|
||||
注意：conda-pack 可能需要较长时间（每个环境 5-10 分钟），建议后台执行。
|
||||
|
||||
### 步骤 8：Git 提交
|
||||
|
||||
```bash
|
||||
cd /home/szlc/change_hair_3090
|
||||
git add .
|
||||
# 验证大文件未被追踪
|
||||
git status --short | grep -E "\.(safetensors|pth|pt|tar\.gz|tar)$" # 应无输出
|
||||
git commit -m "初始化换发型项目：3个微服务代码 + 部署脚本"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 验证步骤
|
||||
|
||||
### 路径一致性验证
|
||||
```bash
|
||||
# 检查无残留旧路径（注释行除外）
|
||||
grep -rn "/home/szlc/project\|/gz-fs\|/mnt/nas_hdd\|/root/project" /home/szlc/change_hair_3090/ --include="*.py" --include="*.ini" --include="*.sh" | grep -v "^.*:#"
|
||||
# 应无输出或仅匹配注释行
|
||||
```
|
||||
|
||||
### Git 追踪验证
|
||||
```bash
|
||||
# 大文件不应被追踪
|
||||
git status --short | grep -E "weights/|models/.*safetensors|\.tar\.gz|\.tar$"
|
||||
# 应无输出
|
||||
```
|
||||
|
||||
### 目录完整性验证
|
||||
```bash
|
||||
# 关键目录存在
|
||||
for d in hair_service_sd/weights stable-diffusion-webui/models/Lora stable-diffusion-webui/models/Stable-diffusion stable-diffusion-webui/extensions/sd-webui-controlnet kohya_ss_home/.local data/ref_hairstyle; do
|
||||
[ -d "$d" ] && echo "✓ $d" || echo "✗ $d 缺失"
|
||||
done
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 网盘上传清单
|
||||
|
||||
| 路径 | 大小 | 说明 |
|
||||
|------|------|------|
|
||||
| hair_service_sd/weights/ | 10G | 换发算法模型 |
|
||||
| stable-diffusion-webui/models/ | 41G | SD模型+LoRA+辅助模型 |
|
||||
| stable-diffusion-webui/extensions/sd-webui-controlnet/ | 4.4G | ControlNet扩展 |
|
||||
| stable-diffusion-webui/repositories/ | 225M | SD依赖仓库 |
|
||||
| kohya_ss_home/.local/ | 7.6G | 训练pip包 |
|
||||
| kohya_ss_home/.cache/ | 3.3M | CLIP缓存 |
|
||||
| kohya_ss_home/kohya_ss/ | 439M | 训练代码 |
|
||||
| data/ | 4.9G | 业务数据 |
|
||||
| conda_envs/my_hair.tar.gz | ~5G | conda环境 |
|
||||
| conda_envs/sdwebui.tar.gz | ~4G | conda环境 |
|
||||
| **合计** | **约 78G** | |
|
||||
|
||||
---
|
||||
|
||||
## 关键风险
|
||||
|
||||
1. **hairstyle_model_infer.py:1096** — 活跃代码路径中的 `cv2.imwrite('/home/student/Desktop/...')`，不注释会导致换发色功能崩溃
|
||||
2. **majicmixRealistic_v7.safetensors** — 训练底模不在源目录中，用户需从网上下载
|
||||
3. **conda-pack glibc 兼容性** — 源机器和目标机器需同为 Ubuntu 且 glibc 版本兼容
|
||||
4. **Docker 训练** — 当前代码仍含 Docker 命令，用户表示不再需要 Docker 训练，代码可能需要后续修改
|
||||
@@ -0,0 +1,272 @@
|
||||
# 换发型项目迁移到 3090 部署计划
|
||||
|
||||
## Context
|
||||
|
||||
本机 `/home/szlc/project` 下有一个换发型/换发色/训练发型的 AI 项目，包含 3 个微服务协同工作。目标是将所有代码和依赖迁移到 `/home/szlc/change_hair_3090`，代码用 git 管理，模型/数据/环境等大文件放在同一目录但 .gitignore 排除（单独用网盘上传）。最终实现：在新 3090 Ubuntu 机器上 `git clone` + 下载网盘文件 + 运行 `setup.sh` 即可完成部署。
|
||||
|
||||
**用户决策**：包含全部 228 个 LoRA（33G）；Docker sudo 密码 asdfasdf；从网上下载 majicmixRealistic_v7；不迁移 onediff。
|
||||
|
||||
---
|
||||
|
||||
## 目标目录结构
|
||||
|
||||
```
|
||||
/home/szlc/change_hair_3090/ ← Git 仓库根
|
||||
├── .gitignore
|
||||
├── README.md ← 部署文档 [git]
|
||||
├── setup.sh ← 一键部署脚本 [git]
|
||||
├── start_all_services.sh ← 启动3个服务 [git]
|
||||
│
|
||||
├── hair_service_sd/ ← 换发算法服务 [git: 代码]
|
||||
│ ├── config/configure.ini.template ← 配置模板（含 __BASE_DIR__ 占位符）[git]
|
||||
│ ├── config/configure.ini ← setup.sh 生成 [gitignore]
|
||||
│ ├── core/ models/ utils/ common/ ← 代码 [git]
|
||||
│ ├── *.py ← 代码 [git]
|
||||
│ ├── data/ ← 小静态文件 9M [git]
|
||||
│ └── weights/ ← 10G 模型 [gitignore/网盘]
|
||||
│
|
||||
├── photo_service/ ← 训练服务 [git: 代码 4M]
|
||||
│ ├── lora_train_service_1.py ← 需修改路径 [git]
|
||||
│ ├── webui_im2im.py *.py utils/ ← [git]
|
||||
│ └── conda_envs/py310.yml ← py310 环境定义 [git]
|
||||
│
|
||||
├── stable-diffusion-webui/ ← SD WebUI [git: 代码 ~280M]
|
||||
│ ├── webui.py modules/ javascript/ ← 代码 [git]
|
||||
│ ├── config.json ← 清理onediff配置 [git]
|
||||
│ ├── extensions/sd-webui-controlnet/ ← 4.4G(annotator+models) [gitignore/网盘]
|
||||
│ ├── repositories/ ← 225M [gitignore/网盘]
|
||||
│ └── models/ ← 49G [gitignore/网盘]
|
||||
│ ├── Lora/ ← 33G 228个LoRA
|
||||
│ ├── Stable-diffusion/ ← v1-5 + majicmixRealistic_v7
|
||||
│ └── ESRGAN/ GFPGAN/ BLIP/ ...
|
||||
│
|
||||
├── kohya_ss_home/ ← 训练环境 [gitignore/网盘]
|
||||
│ ├── kohya_ss/ ← 训练代码 427M
|
||||
│ ├── .local/ ← pip包 7.6G
|
||||
│ ├── .cache/clip/ ← CLIP缓存 1.6M
|
||||
│ ├── .cache/huggingface/ ← 1.8M
|
||||
│ └── start_docker.sh ← [单独提取到git]
|
||||
│
|
||||
├── data/ ← 业务数据 4.9G [gitignore/网盘]
|
||||
│ ├── ref_hairstyle/ ref_color/ ref_online/ ...
|
||||
│
|
||||
├── conda_envs/ ← [gitignore/网盘]
|
||||
│ ├── my_hair.tar.gz ← conda-pack ~5G
|
||||
│ └── sdwebui.tar.gz ← conda-pack ~4G
|
||||
│
|
||||
├── docker/ ← [gitignore/网盘]
|
||||
│ └── kohya_ss_image.tar ← Docker镜像
|
||||
│
|
||||
├── logs/ ← [gitignore]
|
||||
└── meidaojia/ ← 监控测试脚本 [git]
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 路径迁移策略（核心）
|
||||
|
||||
### 策略：BASE_DIR 自动推导 + configure.ini 模板
|
||||
|
||||
所有服务通过脚本自身位置自动推导 `BASE_DIR`，不依赖环境变量，clone 到任意目录均可工作。
|
||||
|
||||
### 1. configure.ini → 模板化
|
||||
|
||||
创建 `hair_service_sd/config/configure.ini.template`（[git] 管理），所有 `/home/szlc/project/data/...` 替换为 `__BASE_DIR__/data/...`。
|
||||
|
||||
关键变更：
|
||||
- 所有数据路径：`/home/szlc/project/data/...` → `__BASE_DIR__/data/...`
|
||||
- `train_dir`：`/data/train_material` → `__BASE_DIR__/kohya_ss_home/train_material`（必须在 kohya_ss_home 下，Docker 容器才能访问）
|
||||
- `logpath`：`/home/szlc/project/logs` → `__BASE_DIR__/logs`
|
||||
- `configure.ini` 本身加入 .gitignore，由 setup.sh 用 sed 替换 `__BASE_DIR__` 生成
|
||||
|
||||
### 2. lora_train_service_1.py → BASE_DIR 自动推导
|
||||
|
||||
**第 38-43 行**（online 模式路径）：
|
||||
```python
|
||||
# 修改前
|
||||
kohya_ss_home_dir = '/root/project/kohya_ss_home'
|
||||
webui_lora_dir = '/gz-fs/Lora'
|
||||
# 修改后
|
||||
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
kohya_ss_home_dir = os.path.join(BASE_DIR, 'kohya_ss_home')
|
||||
webui_lora_dir = os.path.join(BASE_DIR, 'stable-diffusion-webui', 'models', 'Lora')
|
||||
```
|
||||
|
||||
### 3. Docker 训练命令（第 301-315 行）→ 路径映射修正
|
||||
|
||||
关键变更：
|
||||
- `--gpus "device=1"` → `--gpus "device=0"`（新机器单 GPU）
|
||||
- 删除 `-v /mnt:/mnt`（NAS 挂载不再需要）
|
||||
- 新增 `-v {sd_models_dir}:/mnt/sd_models`（挂载训练底模到容器）
|
||||
- `--pretrained_model_name_or_path` → `/mnt/sd_models/majicmixRealistic_v7.safetensors`
|
||||
- `--train_data_dir` 和 `--output_dir`：主机路径转换为容器路径（`kohya_ss_home_dir` → `/home/chinatszrn`）
|
||||
|
||||
```python
|
||||
sd_models_dir = os.path.join(BASE_DIR, 'stable-diffusion-webui', 'models', 'Stable-diffusion')
|
||||
container_images_dir = images_dir.replace(kohya_ss_home_dir, '/home/chinatszrn')
|
||||
container_model_dir = model_dir.replace(kohya_ss_home_dir, '/home/chinatszrn')
|
||||
```
|
||||
|
||||
### 4. 其他硬编码路径修复
|
||||
|
||||
| 文件 | 行号 | 修改 |
|
||||
|------|------|------|
|
||||
| `hairstyle_model_infer.py` | 1096 | 注释掉 `cv2.imwrite('/home/student/Desktop/...')`（**活跃代码路径，不注释会崩溃**） |
|
||||
| `watch_delete.py` | 8 | 从 configure.ini 读取 userDir 而非硬编码 |
|
||||
| `config.json` (webui) | - | 删除 `onediff_compiler_caches_path` 和 `onediff_compiler_backend` 两行 |
|
||||
| `gunicorn_config.py` | 21 | 注释掉 `chdir`（未使用） |
|
||||
| webui extensions | - | 删除失效的 onediff_sd_webui_extensions 符号链接 |
|
||||
|
||||
### 5. 安全提醒
|
||||
|
||||
`upload_oss.py` 第 8-9 行包含硬编码的阿里云 OSS AccessKey/Secret（作为 `os.getenv` 默认值）。这些密钥会随代码进入 git。如仓库非私有，应改为只从环境变量读取。
|
||||
|
||||
---
|
||||
|
||||
## Conda 环境迁移
|
||||
|
||||
| 环境 | 大小 | 策略 |
|
||||
|------|------|------|
|
||||
| my_hair | 10G/289包 | conda-pack 打包为 tar.gz（约 5G 压缩） |
|
||||
| sdwebui | 7.3G/193包 | conda-pack 打包为 tar.gz（约 4G 压缩） |
|
||||
| py310 | 701M/32包 | yml 导出（包少，重建可靠） |
|
||||
|
||||
setup.sh 中用 `conda pack` 恢复 tar.gz 环境，用 `conda env create -f` 恢复 py310。
|
||||
|
||||
---
|
||||
|
||||
## 执行步骤
|
||||
|
||||
### 阶段 1：创建目录结构 + 复制代码
|
||||
|
||||
1. 创建目标目录结构
|
||||
2. `rsync` 复制 hair_service_sd 代码（排除 weights, __pycache__, .git, *.log, nohup.out）
|
||||
3. `rsync` 复制 photo_service 代码（排除 __pycache__, .idea, *.log, 大图片）
|
||||
4. `rsync` 复制 stable-diffusion-webui 代码（排除 models/, extensions/sd-webui-controlnet/, repositories/, .git, cache/, outputs/）
|
||||
5. `rsync` 复制 meidaojia（排除 .git, nohup.out, 大图片）
|
||||
6. 复制 kohya_ss_home/start_docker.sh 到 git 管理
|
||||
|
||||
### 阶段 2：复制资源（网盘文件）
|
||||
|
||||
7. 复制 hair_service_sd/weights/（10G）
|
||||
8. 复制 hair_service_sd/data/（9M 静态文件）
|
||||
9. 复制 stable-diffusion-webui/models/（49G，含 33G LoRA）
|
||||
10. 复制 stable-diffusion-webui/extensions/sd-webui-controlnet/（4.4G）
|
||||
11. 复制 stable-diffusion-webui/repositories/（225M）
|
||||
12. 复制 data/ 业务数据（4.9G）
|
||||
13. 复制 kohya_ss_home/（kohya_ss 代码 + .local 7.6G + .cache/clip + .cache/huggingface，排除 pip/torch 缓存）
|
||||
|
||||
### 阶段 3：路径修改
|
||||
|
||||
14. 创建 `configure.ini.template`
|
||||
15. 修改 `lora_train_service_1.py`（BASE_DIR + Docker 命令）
|
||||
16. 注释 `hairstyle_model_infer.py:1096`
|
||||
17. 修改 `watch_delete.py`
|
||||
18. 清理 `config.json`（删除 onediff 配置）
|
||||
19. 注释 `gunicorn_config.py:21`
|
||||
20. 创建 `start_all_services.sh`
|
||||
21. 创建 `setup.sh`
|
||||
22. 创建 `README.md`
|
||||
23. 创建 `.gitignore`
|
||||
|
||||
### 阶段 4：导出环境和镜像
|
||||
|
||||
24. `conda install -n base -c conda-forge conda-pack`
|
||||
25. `conda pack -n my_hair -o conda_envs/my_hair.tar.gz`
|
||||
26. `conda pack -n sdwebui -o conda_envs/sdwebui.tar.gz`
|
||||
27. `conda env export -n py310 --no-builds > conda_envs/py310.yml`
|
||||
28. `sudo docker save chinatszrn/ubuntu:kohya_ss -o docker/kohya_ss_image.tar`
|
||||
|
||||
### 阶段 5：Git 提交
|
||||
|
||||
29. `git add .` + 验证 .gitignore 生效（大文件不被追踪）
|
||||
30. `git commit`
|
||||
|
||||
---
|
||||
|
||||
## setup.sh 部署脚本功能
|
||||
|
||||
新机器上执行 `./setup.sh` 完成：
|
||||
1. 检查前提条件（git, conda, docker, nvidia-docker）
|
||||
2. 用 sed 从 template 生成 configure.ini
|
||||
3. 用 conda-pack 恢复 my_hair 和 sdwebui 环境
|
||||
4. 用 yml 创建 py310 环境
|
||||
5. `docker load` 导入 kohya_ss 镜像
|
||||
6. 检查所有模型/数据目录是否存在（报告缺失项）
|
||||
7. 检查 majicmixRealistic_v7.safetensors 是否已下载
|
||||
8. 创建运行时目录（tmp, res_dir, userImage 等）
|
||||
|
||||
---
|
||||
|
||||
## 验证方案
|
||||
|
||||
### 环境验证
|
||||
```bash
|
||||
# Conda 环境
|
||||
/home/szlc/miniconda3/envs/my_hair/bin/python -c "import torch; print(torch.cuda.is_available())"
|
||||
/home/szlc/miniconda3/envs/sdwebui/bin/python -c "import torch; print(torch.cuda.is_available())"
|
||||
# Docker
|
||||
sudo docker run --rm --gpus all chinatszrn/ubuntu:kohya_ss nvidia-smi
|
||||
```
|
||||
|
||||
### 服务启动验证
|
||||
```bash
|
||||
# 1. hair_service_sd（端口 8801）
|
||||
cd /home/szlc/change_hair_3090/hair_service_sd
|
||||
/home/szlc/miniconda3/envs/my_hair/bin/python run_copy_cost_colorb64.py
|
||||
# 预期：看到 "[HAIR_INIT] All init done." 且端口 8801 可访问
|
||||
|
||||
# 2. webui（端口 57860）
|
||||
cd /home/szlc/change_hair_3090/stable-diffusion-webui
|
||||
/home/szlc/miniconda3/envs/sdwebui/bin/python webui.py --api --listen --xformers --port 57860
|
||||
# 预期：curl http://127.0.0.1:57860/sdapi/v1/options 返回 JSON
|
||||
|
||||
# 3. photo_service（端口 32678）
|
||||
cd /home/szlc/change_hair_3090/photo_service
|
||||
/home/szlc/miniconda3/envs/py310/bin/python lora_train_service_1.py
|
||||
# 预期：服务启动无报错
|
||||
```
|
||||
|
||||
### 路径一致性验证
|
||||
```bash
|
||||
# 无残留旧路径
|
||||
grep -r "/home/szlc/project" /home/szlc/change_hair_3090/ --include="*.py" --include="*.ini" --include="*.sh"
|
||||
grep -r "/gz-fs\|/mnt/nas_hdd\|/root/project" /home/szlc/change_hair_3090/ --include="*.py" --include="*.ini"
|
||||
# 以上应无输出（注释行除外）
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## .gitignore 中的例外
|
||||
|
||||
以下文件虽在大文件目录中，但需 git 管理（在 .gitignore 中用 `!` 添加例外）：
|
||||
- `!conda_envs/py310.yml` — py310 环境定义（小文本文件）
|
||||
- `!kohya_ss_home/start_docker.sh` — Docker 启动脚本（小文件）
|
||||
|
||||
---
|
||||
|
||||
## 关键风险
|
||||
|
||||
1. **hairstyle_model_infer.py:1096** — 活跃代码路径中的 `cv2.imwrite('/home/student/Desktop/...')`，不注释会导致换发色功能崩溃
|
||||
2. **Docker GPU 设备号** — 原代码 `--gpus "device=1"`，新机器需改为 `device=0`
|
||||
3. **train_dir 路径变更** — 原 `/data/train_material`（121G）不迁移，新路径在 `kohya_ss_home/train_material/`
|
||||
4. **OSS 密钥** — upload_oss.py 含硬编码密钥，会进入 git
|
||||
5. **majicmixRealistic_v7** — 需用户手动下载放置到 `stable-diffusion-webui/models/Stable-diffusion/`
|
||||
6. **conda-pack glibc 兼容性** — 源机器和目标机器需同为 Ubuntu 且 glibc 版本兼容
|
||||
|
||||
---
|
||||
|
||||
## 网盘需上传的文件清单
|
||||
|
||||
| 路径 | 大小 | 说明 |
|
||||
|------|------|------|
|
||||
| hair_service_sd/weights/ | 10G | 换发算法模型 |
|
||||
| stable-diffusion-webui/models/ | 49G | SD模型+LoRA+辅助模型 |
|
||||
| stable-diffusion-webui/extensions/sd-webui-controlnet/ | 4.4G | ControlNet扩展 |
|
||||
| stable-diffusion-webui/repositories/ | 225M | SD依赖仓库 |
|
||||
| kohya_ss_home/ | ~8G | 训练环境(代码+.local+clip缓存) |
|
||||
| data/ | 4.9G | 业务数据 |
|
||||
| conda_envs/my_hair.tar.gz | ~5G | conda环境 |
|
||||
| conda_envs/sdwebui.tar.gz | ~4G | conda环境 |
|
||||
| docker/kohya_ss_image.tar | ?G | Docker镜像 |
|
||||
| **合计** | **约 85G+** | |
|
||||
@@ -0,0 +1,123 @@
|
||||
# 换发型项目
|
||||
|
||||
AI 换发型/换发色/训练发型服务，基于 Stable Diffusion + ControlNet + LoRA。
|
||||
|
||||
## 项目结构
|
||||
|
||||
```
|
||||
change_hair_3090/
|
||||
├── hair_service_sd/ # 换发算法服务 (端口 8801)
|
||||
├── photo_service/ # LoRA 训练服务 (端口 32678)
|
||||
├── stable-diffusion-webui/ # SD WebUI 推理服务 (端口 57860)
|
||||
├── kohya_ss_home/ # kohya_ss 训练环境
|
||||
├── data/ # 业务数据（发型参考图、用户图片等）
|
||||
├── conda_envs/ # conda 环境打包文件
|
||||
├── setup.sh # 一键部署脚本
|
||||
├── start_all_services.sh # 启动所有服务
|
||||
└── README.md
|
||||
```
|
||||
|
||||
## 前提条件
|
||||
|
||||
- Ubuntu 22.04
|
||||
- NVIDIA 3090 GPU + 驱动
|
||||
- conda (Miniconda/Anaconda)
|
||||
- git
|
||||
|
||||
## 部署步骤
|
||||
|
||||
### 1. Clone 代码
|
||||
|
||||
```bash
|
||||
git clone <repo_url> ~/change_hair_3090
|
||||
cd ~/change_hair_3090
|
||||
```
|
||||
|
||||
### 2. 下载网盘文件
|
||||
|
||||
将网盘上的以下目录/文件下载到项目根目录（与代码合并）：
|
||||
|
||||
| 路径 | 大小 | 说明 |
|
||||
|------|------|------|
|
||||
| hair_service_sd/weights/ | 10G | 换发算法模型 |
|
||||
| stable-diffusion-webui/models/ | 41G | SD模型 + LoRA + 辅助模型 |
|
||||
| stable-diffusion-webui/extensions/sd-webui-controlnet/ | 4.4G | ControlNet 扩展 |
|
||||
| stable-diffusion-webui/repositories/ | 225M | SD 依赖仓库 |
|
||||
| kohya_ss_home/.local/ | 7.6G | 训练 pip 包 |
|
||||
| kohya_ss_home/.cache/ | 3.3M | CLIP 缓存 |
|
||||
| kohya_ss_home/kohya_ss/ | 439M | 训练代码 |
|
||||
| data/ | 4.9G | 业务数据 |
|
||||
| conda_envs/my_hair.tar.gz | ~5G | conda 环境 |
|
||||
| conda_envs/sdwebui.tar.gz | ~4G | conda 环境 |
|
||||
|
||||
### 3. 下载训练底模
|
||||
|
||||
从网上下载 `majicmixRealistic_v7.safetensors`，放置到：
|
||||
```
|
||||
stable-diffusion-webui/models/Stable-diffusion/majicmixRealistic_v7.safetensors
|
||||
```
|
||||
|
||||
### 4. 运行部署脚本
|
||||
|
||||
```bash
|
||||
chmod +x setup.sh start_all_services.sh
|
||||
./setup.sh
|
||||
```
|
||||
|
||||
setup.sh 会：
|
||||
1. 生成 configure.ini（自动替换路径）
|
||||
2. 恢复 conda 环境（my_hair, sdwebui, py310）
|
||||
3. 检查模型和数据目录完整性
|
||||
4. 创建运行时目录
|
||||
|
||||
### 5. 启动服务
|
||||
|
||||
```bash
|
||||
./start_all_services.sh
|
||||
```
|
||||
|
||||
## 服务说明
|
||||
|
||||
| 服务 | 端口 | conda 环境 | 说明 |
|
||||
|------|------|-----------|------|
|
||||
| hair_service_sd | 8801 | my_hair | 换发型/换发色核心算法 |
|
||||
| stable-diffusion-webui | 57860 | sdwebui | SD 图生图推理 |
|
||||
| photo_service | 32678 | py310 | LoRA 训练调度 |
|
||||
|
||||
## 查看日志
|
||||
|
||||
```bash
|
||||
tail -f logs/hair_service.log
|
||||
tail -f logs/webui.log
|
||||
tail -f logs/photo_service.log
|
||||
```
|
||||
|
||||
## 停止服务
|
||||
|
||||
```bash
|
||||
# 按端口查找并停止
|
||||
kill $(lsof -t -i:8801)
|
||||
kill $(lsof -t -i:57860)
|
||||
kill $(lsof -t -i:32678)
|
||||
```
|
||||
|
||||
## 常见问题
|
||||
|
||||
### conda 环境恢复失败
|
||||
确保 `CONDA_BASE` 环境变量指向你的 conda 安装路径（默认 `/home/szlc/miniconda3`）：
|
||||
```bash
|
||||
export CONDA_BASE=/your/conda/path
|
||||
./setup.sh
|
||||
```
|
||||
|
||||
### GPU 不可用
|
||||
检查 NVIDIA 驱动和 CUDA：
|
||||
```bash
|
||||
nvidia-smi
|
||||
```
|
||||
|
||||
### WebUI 启动慢
|
||||
首次启动需要加载模型（约 30-60 秒），请耐心等待。查看 `logs/webui.log` 确认状态。
|
||||
|
||||
### 路径问题
|
||||
所有路径通过 `BASE_DIR` 自动推导，`configure.ini` 由 `setup.sh` 自动生成。如需手动修改配置，编辑 `hair_service_sd/config/configure.ini.template` 后重新运行 `setup.sh`。
|
||||
@@ -0,0 +1,55 @@
|
||||
name: py310
|
||||
channels:
|
||||
- defaults
|
||||
dependencies:
|
||||
- _libgcc_mutex=0.1
|
||||
- _openmp_mutex=5.1
|
||||
- bzip2=1.0.8
|
||||
- ca-certificates=2024.7.2
|
||||
- ld_impl_linux-64=2.38
|
||||
- libffi=3.4.4
|
||||
- libgcc-ng=11.2.0
|
||||
- libgomp=11.2.0
|
||||
- libstdcxx-ng=11.2.0
|
||||
- libuuid=1.41.5
|
||||
- ncurses=6.4
|
||||
- openssl=3.0.14
|
||||
- pip=24.0
|
||||
- python=3.10.14
|
||||
- readline=8.2
|
||||
- setuptools=69.5.1
|
||||
- sqlite=3.45.3
|
||||
- tk=8.6.14
|
||||
- tqdm=4.66.4
|
||||
- tzdata=2024a
|
||||
- wheel=0.43.0
|
||||
- xz=5.4.6
|
||||
- zlib=1.2.13
|
||||
- pip:
|
||||
- blinker==1.8.2
|
||||
- certifi==2024.7.4
|
||||
- charset-normalizer==3.3.2
|
||||
- click==8.1.7
|
||||
- flask==3.0.3
|
||||
- gevent==24.2.1
|
||||
- greenlet==3.0.3
|
||||
- idna==3.7
|
||||
- imageio==2.34.2
|
||||
- itsdangerous==2.2.0
|
||||
- jinja2==3.1.4
|
||||
- lazy-loader==0.4
|
||||
- markupsafe==2.1.5
|
||||
- networkx==3.3
|
||||
- numpy==2.0.0
|
||||
- opencv-python==4.10.0.84
|
||||
- packaging==24.1
|
||||
- pillow==10.4.0
|
||||
- requests==2.32.3
|
||||
- scikit-image==0.24.0
|
||||
- scipy==1.14.0
|
||||
- tifffile==2024.7.2
|
||||
- urllib3==2.2.2
|
||||
- werkzeug==3.0.3
|
||||
- zope-event==5.0
|
||||
- zope-interface==6.4.post2
|
||||
prefix: /home/szlc/miniconda3/envs/py310
|
||||
@@ -0,0 +1,194 @@
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
share/python-wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
|
||||
# PyInstaller
|
||||
# Usually these files are written by a python script from a template
|
||||
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Installer logs
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.nox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
*.py,cover
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
cover/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
db.sqlite3
|
||||
db.sqlite3-journal
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
|
||||
# PyBuilder
|
||||
.pybuilder/
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# IPython
|
||||
profile_default/
|
||||
ipython_config.py
|
||||
|
||||
# pyenv
|
||||
# For a library or package, you might want to ignore these files since the code is
|
||||
# intended to run in multiple environments; otherwise, check them in:
|
||||
# .python-version
|
||||
|
||||
# pipenv
|
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||
# install all needed dependencies.
|
||||
#Pipfile.lock
|
||||
|
||||
# UV
|
||||
# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
#uv.lock
|
||||
|
||||
# poetry
|
||||
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
||||
#poetry.lock
|
||||
|
||||
# pdm
|
||||
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
||||
#pdm.lock
|
||||
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
|
||||
# in version control.
|
||||
# https://pdm.fming.dev/latest/usage/project/#working-with-version-control
|
||||
.pdm.toml
|
||||
.pdm-python
|
||||
.pdm-build/
|
||||
|
||||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
||||
__pypackages__/
|
||||
|
||||
# Celery stuff
|
||||
celerybeat-schedule
|
||||
celerybeat.pid
|
||||
|
||||
# SageMath parsed files
|
||||
*.sage.py
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
dmypy.json
|
||||
|
||||
# Pyre type checker
|
||||
.pyre/
|
||||
|
||||
# pytype static type analyzer
|
||||
.pytype/
|
||||
|
||||
# Cython debug symbols
|
||||
cython_debug/
|
||||
|
||||
# PyCharm
|
||||
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
||||
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
||||
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||
#.idea/
|
||||
|
||||
# Abstra
|
||||
# Abstra is an AI-powered process automation framework.
|
||||
# Ignore directories containing user credentials, local state, and settings.
|
||||
# Learn more at https://abstra.io/docs
|
||||
.abstra/
|
||||
|
||||
# Visual Studio Code
|
||||
# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
|
||||
# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. However, if you prefer,
|
||||
# you could uncomment the following to ignore the enitre vscode folder
|
||||
# .vscode/
|
||||
|
||||
# Ruff stuff:
|
||||
.ruff_cache/
|
||||
|
||||
# PyPI configuration file
|
||||
.pypirc
|
||||
|
||||
# Cursor
|
||||
# Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
|
||||
# exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
|
||||
# refer to https://docs.cursor.com/context/ignore-files
|
||||
.cursorignore
|
||||
.cursorindexingignore
|
||||
@@ -0,0 +1,28 @@
|
||||
换发算法服务代码
|
||||
|
||||
相关配置：config/configure.ini \
|
||||
模型:请存放于weights目录下 \
|
||||
服务启动：执行 gunicorn server:app -c gunicorn_config.py \
|
||||
测试相关代码请写在 unit_test.py里
|
||||
|
||||
建议环境：
|
||||
python3.7，pytorch1.8,cu11.1,torchvision10.0 \
|
||||
gunicorn需要在/usr/bin/gunicorn需要在里修改python地址
|
||||
|
||||
run run_copy_cost.py
|
||||
|
||||
|
||||
|
||||
需要启动的服务：
|
||||
环境变量设置
|
||||
export CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1
|
||||
1. 换发算法服务
|
||||
cd /root/project/hair_service_sd
|
||||
/usr/local/miniconda3/envs/condiff-train-hair/bin/python run_copy_cost_colorb64.py
|
||||
2. webUI服务推理 + onediff
|
||||
cd /root/project/onediff/stable-diffusion-webui
|
||||
/usr/local/miniconda3/envs/onediff/bin/python python webui.py --api --listen --xformers --port 57860
|
||||
./webui.sh --api --listen --disable-safe-unpickle --port 9038
|
||||
3. photo_service
|
||||
cd /root/project/photo_service
|
||||
/usr/local/miniconda3/envs/py310/bin/python lora_train_service_1.py
|
||||
@@ -0,0 +1,75 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
def get_norm(norm, out_channels=None):
|
||||
"""
|
||||
Args:
|
||||
norm (str or callable):
|
||||
|
||||
Returns:
|
||||
nn.Module or None: the normalization layer
|
||||
"""
|
||||
if isinstance(norm, str):
|
||||
if len(norm) == 0:
|
||||
return None
|
||||
norm = {
|
||||
"BN": nn.BatchNorm2d,
|
||||
"IN": nn.InstanceNorm2d,
|
||||
"GN": lambda channels: nn.GroupNorm(32, channels),
|
||||
"nnSyncBN": nn.SyncBatchNorm, # keep for debugging
|
||||
}[norm]
|
||||
if out_channels is not None:
|
||||
return norm(out_channels)
|
||||
else:
|
||||
return norm
|
||||
|
||||
class Conv2d(torch.nn.Conv2d):
|
||||
"""
|
||||
A wrapper around :class:`torch.nn.Conv2d` to support zero-size tensor and more features.
|
||||
"""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
"""
|
||||
Extra keyword arguments supported in addition to those in `torch.nn.Conv2d`:
|
||||
|
||||
Args:
|
||||
norm (nn.Module, optional): a normalization layer
|
||||
activation (callable(Tensor) -> Tensor): a callable activation function
|
||||
|
||||
It assumes that norm layer is used before activation.
|
||||
"""
|
||||
norm = kwargs.pop("norm", None)
|
||||
activation = kwargs.pop("activation", None)
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
self.norm = norm
|
||||
self.activation = activation
|
||||
|
||||
def forward(self, x):
|
||||
x = super().forward(x)
|
||||
if self.norm is not None:
|
||||
x = self.norm(x)
|
||||
if self.activation is not None:
|
||||
x = self.activation(x)
|
||||
return x
|
||||
|
||||
class Backbone(nn.Module):
|
||||
"""
|
||||
Abstract base class for network backbones.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""
|
||||
The `__init__` method of any subclass can specify its own set of arguments.
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
def forward(self):
|
||||
"""
|
||||
Subclasses must override this method, but adhere to the same return type.
|
||||
|
||||
Returns:
|
||||
dict[str: Tensor]: mapping from feature name (e.g., "res2") to tensor
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,268 @@
|
||||
import math
|
||||
import utils.weight_init as weight_init
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from bodyseg.backbone.backbone import Backbone, get_norm, Conv2d
|
||||
from bodyseg.backbone.resnet import build_resnet_backbone
|
||||
|
||||
class FPN(Backbone):
|
||||
"""
|
||||
This module implements Feature Pyramid Network.
|
||||
It creates pyramid features built on top of some input feature maps.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, bottom_up, in_features, out_channels, norm="", top_block=None, fuse_type="sum"
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
bottom_up (Backbone): module representing the bottom up subnetwork.
|
||||
Must be a subclass of :class:`Backbone`. The multi-scale feature
|
||||
maps generated by the bottom up network, and listed in `in_features`,
|
||||
are used to generate FPN levels.
|
||||
in_features (list[str]): names of the input feature maps coming
|
||||
from the backbone to which FPN is attached. For example, if the
|
||||
backbone produces ["res2", "res3", "res4"], any *contiguous* sublist
|
||||
of these may be used; order must be from high to low resolution.
|
||||
out_channels (int): number of channels in the output feature maps.
|
||||
norm (str): the normalization to use.
|
||||
top_block (nn.Module or None): if provided, an extra operation will
|
||||
be performed on the output of the last (smallest resolution)
|
||||
FPN output, and the result will extend the result list. The top_block
|
||||
further downsamples the feature map. It must have an attribute
|
||||
"num_levels", meaning the number of extra FPN levels added by
|
||||
this block, and "in_feature", which is a string representing
|
||||
its input feature (e.g., p5).
|
||||
fuse_type (str): types for fusing the top down features and the lateral
|
||||
ones. It can be "sum" (default), which sums up element-wise; or "avg",
|
||||
which takes the element-wise mean of the two.
|
||||
"""
|
||||
super(FPN, self).__init__()
|
||||
assert isinstance(bottom_up, Backbone)
|
||||
|
||||
# Feature map strides and channels from the bottom up network (e.g. ResNet)
|
||||
in_strides = [bottom_up._out_feature_strides[f] for f in in_features]
|
||||
in_channels = [bottom_up._out_feature_channels[f] for f in in_features]
|
||||
|
||||
_assert_strides_are_log2_contiguous(in_strides)
|
||||
lateral_convs = []
|
||||
output_convs = []
|
||||
|
||||
use_bias = norm == ""
|
||||
for idx, in_channels in enumerate(in_channels):
|
||||
lateral_norm = get_norm(norm, out_channels)
|
||||
output_norm = get_norm(norm, out_channels)
|
||||
|
||||
lateral_conv = Conv2d(
|
||||
in_channels, out_channels, kernel_size=1, bias=use_bias, norm=lateral_norm
|
||||
)
|
||||
output_conv = Conv2d(
|
||||
out_channels,
|
||||
out_channels,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
bias=use_bias,
|
||||
norm=output_norm,
|
||||
)
|
||||
weight_init.c2_xavier_fill(lateral_conv)
|
||||
weight_init.c2_xavier_fill(output_conv)
|
||||
stage = int(math.log2(in_strides[idx]))
|
||||
|
||||
lateral_convs.append(lateral_conv)
|
||||
output_convs.append(output_conv)
|
||||
# Place convs into top-down order (from low to high resolution)
|
||||
# to make the top-down computation in forward clearer.
|
||||
self.lateral_convs = nn.ModuleList(lateral_convs[::-1])
|
||||
self.output_convs = nn.ModuleList(output_convs[::-1])
|
||||
self.top_block = top_block
|
||||
self.in_features = in_features
|
||||
self.bottom_up = bottom_up
|
||||
# Return feature names are "p<stage>", like ["p2", "p3", ..., "p6"]
|
||||
self._out_feature_strides = {"p{}".format(int(math.log2(s))): s for s in in_strides}
|
||||
# top block output feature maps.
|
||||
if self.top_block is not None:
|
||||
for s in range(stage, stage + self.top_block.num_levels):
|
||||
self._out_feature_strides["p{}".format(s + 1)] = 2 ** (s + 1)
|
||||
|
||||
self._out_features = list(self._out_feature_strides.keys())
|
||||
self._out_feature_channels = {k: out_channels for k in self._out_features}
|
||||
assert fuse_type in {"avg", "sum"}
|
||||
self._fuse_type = fuse_type
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Args:
|
||||
input (dict[str: Tensor]): mapping feature map name (e.g., "res5") to
|
||||
feature map tensor for each feature level in high to low resolution order.
|
||||
|
||||
Returns:
|
||||
dict[str: Tensor]:
|
||||
mapping from feature map name to FPN feature map tensor
|
||||
in high to low resolution order. Returned feature names follow the FPN
|
||||
paper convention: "p<stage>", where stage has stride = 2 ** stage e.g.,
|
||||
["p2", "p3", ..., "p6"].
|
||||
"""
|
||||
# Reverse feature maps into top-down order (from low to high resolution)
|
||||
bottom_up_features = self.bottom_up(x)
|
||||
x = [bottom_up_features[f] for f in self.in_features[::-1]]
|
||||
results = []
|
||||
prev_features = self.lateral_convs[0](x[0])
|
||||
results.append(self.output_convs[0](prev_features))
|
||||
for features, lateral_conv, output_conv in zip(
|
||||
x[1:], self.lateral_convs[1:], self.output_convs[1:]
|
||||
):
|
||||
top_down_features = F.interpolate(prev_features, scale_factor=2, mode="nearest")
|
||||
lateral_features = lateral_conv(features)
|
||||
prev_features = lateral_features + top_down_features
|
||||
if self._fuse_type == "avg":
|
||||
prev_features /= 2
|
||||
results.insert(0, output_conv(prev_features))
|
||||
|
||||
if self.top_block is not None:
|
||||
top_block_in_feature = bottom_up_features.get(self.top_block.in_feature, None)
|
||||
if top_block_in_feature is None:
|
||||
top_block_in_feature = results[self._out_features.index(self.top_block.in_feature)]
|
||||
results.extend(self.top_block(top_block_in_feature))
|
||||
assert len(self._out_features) == len(results)
|
||||
return dict(zip(self._out_features, results))
|
||||
|
||||
def _assert_strides_are_log2_contiguous(strides):
|
||||
"""
|
||||
Assert that each stride is 2x times its preceding stride, i.e. "contiguous in log2".
|
||||
"""
|
||||
for i, stride in enumerate(strides[1:], 1):
|
||||
assert stride == 2 * strides[i - 1], "Strides {} {} are not log2 contiguous".format(
|
||||
stride, strides[i - 1]
|
||||
)
|
||||
|
||||
|
||||
class LastLevelMaxPool(nn.Module):
|
||||
"""
|
||||
This module is used in the original FPN to generate a downsampled
|
||||
P6 feature from P5.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.num_levels = 1
|
||||
self.in_feature = "p5"
|
||||
|
||||
def forward(self, x):
|
||||
return [F.max_pool2d(x, kernel_size=1, stride=2, padding=0)]
|
||||
|
||||
|
||||
class LastLevelP6P7(nn.Module):
|
||||
"""
|
||||
This module is used in RetinaNet to generate extra layers, P6 and P7 from
|
||||
C5 feature.
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels, out_channels):
|
||||
super().__init__()
|
||||
self.num_levels = 2
|
||||
self.in_feature = "res5"
|
||||
self.p6 = nn.Conv2d(in_channels, out_channels, 3, 2, 1)
|
||||
self.p7 = nn.Conv2d(out_channels, out_channels, 3, 2, 1)
|
||||
for module in [self.p6, self.p7]:
|
||||
weight_init.c2_xavier_fill(module)
|
||||
|
||||
def forward(self, c5):
|
||||
p6 = self.p6(c5)
|
||||
p7 = self.p7(F.relu(p6))
|
||||
return [p6, p7]
|
||||
|
||||
|
||||
def build_resnet_fpn_backbone(in_channels=3):
|
||||
"""
|
||||
Args:
|
||||
cfg: a detectron2 CfgNode
|
||||
|
||||
Returns:
|
||||
backbone (Backbone): backbone module, must be a subclass of :class:`Backbone`.
|
||||
"""
|
||||
bottom_up = build_resnet_backbone(in_channels)
|
||||
in_features = ["res2", "res3", "res4"]
|
||||
out_channels = 256
|
||||
backbone = FPN(
|
||||
bottom_up=bottom_up,
|
||||
in_features=in_features,
|
||||
out_channels=out_channels,
|
||||
norm="BN",
|
||||
# top_block=LastLevelMaxPool(),
|
||||
top_block=None,
|
||||
fuse_type="sum",
|
||||
)
|
||||
return backbone
|
||||
|
||||
def build_retinanet_resnet_fpn_backbone(cfg, in_channels=3):
|
||||
"""
|
||||
Args:
|
||||
cfg: a detectron2 CfgNode
|
||||
|
||||
Returns:
|
||||
backbone (Backbone): backbone module, must be a subclass of :class:`Backbone`.
|
||||
"""
|
||||
bottom_up = build_resnet_backbone(cfg, in_channels)
|
||||
in_features = cfg.MODEL.FPN.IN_FEATURES
|
||||
out_channels = cfg.MODEL.FPN.OUT_CHANNELS
|
||||
in_channels_p6p7 = bottom_up._out_feature_channels["res5"]
|
||||
backbone = FPN(
|
||||
bottom_up=bottom_up,
|
||||
in_features=in_features,
|
||||
out_channels=out_channels,
|
||||
norm=cfg.MODEL.FPN.NORM,
|
||||
top_block=LastLevelP6P7(in_channels_p6p7, out_channels),
|
||||
fuse_type=cfg.MODEL.FPN.FUSE_TYPE,
|
||||
)
|
||||
return backbone
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
from config.default import get_cfg
|
||||
|
||||
def setup(args):
|
||||
"""
|
||||
Create configs and perform basic setups.
|
||||
"""
|
||||
cfg = get_cfg()
|
||||
cfg.merge_from_file(args.cfg)
|
||||
cfg.merge_from_list(args.opts)
|
||||
cfg.freeze()
|
||||
return cfg
|
||||
|
||||
parser = argparse.ArgumentParser(description='Train ImageNet network')
|
||||
# general
|
||||
parser.add_argument('--cfg',
|
||||
help='experiment configure file name',
|
||||
required=True,
|
||||
type=str)
|
||||
|
||||
parser.add_argument('opts',
|
||||
help="Modify config options using the command-line",
|
||||
default=None,
|
||||
nargs=argparse.REMAINDER)
|
||||
|
||||
args = parser.parse_args()
|
||||
cfg = setup(args)
|
||||
print(cfg)
|
||||
|
||||
model = build_resnet_fpn_backbone(cfg, 3)
|
||||
# model = build_retinanet_resnet_fpn_backbone(cfg, 3)
|
||||
print(model)
|
||||
# model = torch.nn.DataParallel(model, list(range(2))).cuda()
|
||||
dummy_input = torch.randn(4, 3, 512, 512)
|
||||
|
||||
out = model(dummy_input)
|
||||
|
||||
for k, v in out.items():
|
||||
print(k, v.shape)
|
||||
|
||||
# torch.onnx.export(model, dummy_input, "tmp.onnx", verbose=True,
|
||||
# input_names=['input'],
|
||||
# output_names=['output'])
|
||||
|
||||
pass
|
||||
|
||||
@@ -0,0 +1,298 @@
|
||||
import numpy as np
|
||||
import utils.weight_init as weight_init
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
from bodyseg.backbone.backbone import Backbone, get_norm, Conv2d
|
||||
|
||||
class BasicStem(nn.Module):
|
||||
def __init__(self, in_channels=3, out_channels=64, norm="BN"):
|
||||
"""
|
||||
Args:
|
||||
norm (str or callable): a callable that takes the number of
|
||||
channels and return a `nn.Module`, or a pre-defined string
|
||||
(one of {"FrozenBN", "BN", "GN"}).
|
||||
"""
|
||||
super().__init__()
|
||||
self.conv1 = Conv2d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size=7,
|
||||
stride=2,
|
||||
padding=3,
|
||||
bias=False,
|
||||
norm=get_norm(norm, out_channels),
|
||||
)
|
||||
weight_init.c2_msra_fill(self.conv1)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = F.relu_(x)
|
||||
x = F.max_pool2d(x, kernel_size=3, stride=2, padding=1)
|
||||
return x
|
||||
|
||||
@property
|
||||
def out_channels(self):
|
||||
return self.conv1.out_channels
|
||||
|
||||
@property
|
||||
def stride(self):
|
||||
return 4 # = stride 2 conv -> stride 2 max pool
|
||||
|
||||
class ResNetBlockBase(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, stride):
|
||||
"""
|
||||
The `__init__` method of any subclass should also contain these arguments.
|
||||
|
||||
Args:
|
||||
in_channels (int):
|
||||
out_channels (int):
|
||||
stride (int):
|
||||
"""
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.stride = stride
|
||||
|
||||
class BottleneckBlock(ResNetBlockBase):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
*,
|
||||
bottleneck_channels,
|
||||
stride=1,
|
||||
num_groups=1,
|
||||
norm="BN",
|
||||
stride_in_1x1=False,
|
||||
dilation=1,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
norm (str or callable): a callable that takes the number of
|
||||
channels and return a `nn.Module`, or a pre-defined string
|
||||
(one of {"FrozenBN", "BN", "GN"}).
|
||||
stride_in_1x1 (bool): when stride==2, whether to put stride in the
|
||||
first 1x1 convolution or the bottleneck 3x3 convolution.
|
||||
"""
|
||||
super().__init__(in_channels, out_channels, stride)
|
||||
|
||||
if in_channels != out_channels:
|
||||
self.shortcut = Conv2d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size=1,
|
||||
stride=stride,
|
||||
bias=False,
|
||||
norm=get_norm(norm, out_channels),
|
||||
)
|
||||
else:
|
||||
self.shortcut = None
|
||||
|
||||
# The original MSRA ResNet models have stride in the first 1x1 conv
|
||||
# The subsequent fb.torch.resnet and Caffe2 ResNe[X]t implementations have
|
||||
# stride in the 3x3 conv
|
||||
stride_1x1, stride_3x3 = (stride, 1) if stride_in_1x1 else (1, stride)
|
||||
|
||||
self.conv1 = Conv2d(
|
||||
in_channels,
|
||||
bottleneck_channels,
|
||||
kernel_size=1,
|
||||
stride=stride_1x1,
|
||||
bias=False,
|
||||
norm=get_norm(norm, bottleneck_channels),
|
||||
)
|
||||
|
||||
self.conv2 = Conv2d(
|
||||
bottleneck_channels,
|
||||
bottleneck_channels,
|
||||
kernel_size=3,
|
||||
stride=stride_3x3,
|
||||
padding=1 * dilation,
|
||||
bias=False,
|
||||
groups=num_groups,
|
||||
dilation=dilation,
|
||||
norm=get_norm(norm, bottleneck_channels),
|
||||
)
|
||||
|
||||
self.conv3 = Conv2d(
|
||||
bottleneck_channels,
|
||||
out_channels,
|
||||
kernel_size=1,
|
||||
bias=False,
|
||||
norm=get_norm(norm, out_channels),
|
||||
)
|
||||
|
||||
for layer in [self.conv1, self.conv2, self.conv3, self.shortcut]:
|
||||
if layer is not None: # shortcut can be None
|
||||
weight_init.c2_msra_fill(layer)
|
||||
|
||||
# Zero-initialize the last normalization in each residual branch,
|
||||
# so that at the beginning, the residual branch starts with zeros,
|
||||
# and each residual block behaves like an identity.
|
||||
# See Sec 5.1 in "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour":
|
||||
# "For BN layers, the learnable scaling coefficient γ is initialized
|
||||
# to be 1, except for each residual block's last BN
|
||||
# where γ is initialized to be 0."
|
||||
|
||||
# nn.init.constant_(self.conv3.norm.weight, 0)
|
||||
# TODO this somehow hurts performance when training GN models from scratch.
|
||||
# Add it as an option when we need to use this code to train a backbone.
|
||||
|
||||
def forward(self, x):
|
||||
out = self.conv1(x)
|
||||
out = F.relu_(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = F.relu_(out)
|
||||
|
||||
out = self.conv3(out)
|
||||
|
||||
if self.shortcut is not None:
|
||||
shortcut = self.shortcut(x)
|
||||
else:
|
||||
shortcut = x
|
||||
|
||||
out += shortcut
|
||||
out = F.relu_(out)
|
||||
return out
|
||||
|
||||
def make_stage(block_class, num_blocks, first_stride, **kwargs):
|
||||
"""
|
||||
Create a resnet stage by creating many blocks.
|
||||
Args:
|
||||
block_class (class): a subclass of ResNetBlockBase
|
||||
num_blocks (int):
|
||||
first_stride (int): the stride of the first block. The other blocks will have stride=1.
|
||||
A `stride` argument will be passed to the block constructor.
|
||||
kwargs: other arguments passed to the block constructor.
|
||||
|
||||
Returns:
|
||||
list[nn.Module]: a list of block module.
|
||||
"""
|
||||
blocks = []
|
||||
for i in range(num_blocks):
|
||||
blocks.append(block_class(stride=first_stride if i == 0 else 1, **kwargs))
|
||||
kwargs["in_channels"] = kwargs["out_channels"]
|
||||
return blocks
|
||||
|
||||
class ResNet(Backbone):
|
||||
def __init__(self, stem, stages, num_classes=None, out_features=None):
|
||||
"""
|
||||
Args:
|
||||
stem (nn.Module): a stem module
|
||||
stages (list[list[ResNetBlock]]): several (typically 4) stages,
|
||||
each contains multiple :class:`ResNetBlockBase`.
|
||||
num_classes (None or int): if None, will not perform classification.
|
||||
out_features (list[str]): name of the layers whose outputs should
|
||||
be returned in forward. Can be anything in "stem", "linear", or "res2" ...
|
||||
If None, will return the output of the last layer.
|
||||
"""
|
||||
super(ResNet, self).__init__()
|
||||
self.stem = stem
|
||||
self.num_classes = num_classes
|
||||
|
||||
current_stride = self.stem.stride
|
||||
self._out_feature_strides = {"stem": current_stride}
|
||||
self._out_feature_channels = {"stem": self.stem.out_channels}
|
||||
|
||||
self.stages = []
|
||||
self.names = []
|
||||
for i, blocks in enumerate(stages):
|
||||
for block in blocks:
|
||||
assert isinstance(block, ResNetBlockBase), block
|
||||
curr_channels = block.out_channels
|
||||
stage = nn.Sequential(*blocks)
|
||||
name = "res" + str(i + 2)
|
||||
self.stages.append(stage)
|
||||
self.names.append(name)
|
||||
self._out_feature_strides[name] = current_stride = int(
|
||||
current_stride * np.prod([k.stride for k in blocks])
|
||||
)
|
||||
self._out_feature_channels[name] = blocks[-1].out_channels
|
||||
self.stages = nn.ModuleList(self.stages)
|
||||
|
||||
if num_classes is not None:
|
||||
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
|
||||
self.linear = nn.Linear(curr_channels, num_classes)
|
||||
|
||||
# Sec 5.1 in "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour":
|
||||
# "The 1000-way fully-connected layer is initialized by
|
||||
# drawing weights from a zero-mean Gaussian with standard deviation of 0.01."
|
||||
nn.init.normal_(self.linear.weight, stddev=0.01)
|
||||
name = "linear"
|
||||
|
||||
if out_features is None:
|
||||
out_features = [name]
|
||||
self._out_features = out_features
|
||||
assert len(self._out_features)
|
||||
for out_feature in self._out_features:
|
||||
assert out_feature in self.names, "Available children: {}".format(", ".join(self.names))
|
||||
|
||||
def forward(self, x):
|
||||
outputs = {}
|
||||
x = self.stem(x)
|
||||
if "stem" in self._out_features:
|
||||
outputs["stem"] = x
|
||||
for ix, stage in enumerate(self.stages):
|
||||
name = self.names[ix]
|
||||
x = stage(x)
|
||||
if name in self._out_features:
|
||||
outputs[name] = x
|
||||
if self.num_classes is not None:
|
||||
x = self.avgpool(x)
|
||||
x = self.linear(x)
|
||||
if "linear" in self._out_features:
|
||||
outputs["linear"] = x
|
||||
return outputs
|
||||
|
||||
def build_resnet_backbone(in_channels=3):
|
||||
norm = "BN"
|
||||
stem = BasicStem(
|
||||
in_channels=in_channels,
|
||||
out_channels=64,
|
||||
norm=norm,
|
||||
)
|
||||
|
||||
# fmt: off
|
||||
out_features = ["res2", "res3", "res4"]
|
||||
depth = 101
|
||||
num_groups = 1
|
||||
bottleneck_channels = 64
|
||||
in_channels = 64
|
||||
out_channels = 256
|
||||
stride_in_1x1 = True
|
||||
res5_dilation = 1
|
||||
# fmt: on
|
||||
assert res5_dilation in {1, 2}, "res5_dilation cannot be {}.".format(res5_dilation)
|
||||
|
||||
num_blocks_per_stage = {50: [3, 4, 6, 3], 101: [3, 4, 23, 3], 152: [3, 8, 36, 3]}[depth]
|
||||
|
||||
stages = []
|
||||
|
||||
# Avoid creating variables without gradients
|
||||
# It consumes extra memory and may cause allreduce to fail
|
||||
out_stage_idx = [{"res2": 2, "res3": 3, "res4": 4, "res5": 5}[f] for f in out_features]
|
||||
max_stage_idx = max(out_stage_idx)
|
||||
for idx, stage_idx in enumerate(range(2, max_stage_idx + 1)):
|
||||
dilation = res5_dilation if stage_idx == 5 else 1
|
||||
first_stride = 1 if idx == 0 or (stage_idx == 5 and dilation == 2) else 2
|
||||
stage_kargs = dict()
|
||||
stage_kargs.update({
|
||||
"num_blocks": num_blocks_per_stage[idx],
|
||||
"first_stride": first_stride,
|
||||
"in_channels": in_channels,
|
||||
"bottleneck_channels": bottleneck_channels,
|
||||
"out_channels": out_channels,
|
||||
"num_groups": num_groups,
|
||||
"norm": norm,
|
||||
"stride_in_1x1": stride_in_1x1,
|
||||
"dilation": dilation,
|
||||
})
|
||||
stage_kargs["block_class"] = BottleneckBlock
|
||||
blocks = make_stage(**stage_kargs)
|
||||
in_channels = out_channels
|
||||
out_channels *= 2
|
||||
bottleneck_channels *= 2
|
||||
stages.append(blocks)
|
||||
return ResNet(stem, stages, out_features=out_features)
|
||||
@@ -0,0 +1,244 @@
|
||||
import math
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.utils.model_zoo as model_zoo
|
||||
from bodyseg.backbone.backbone import Backbone
|
||||
|
||||
def fixed_padding(inputs, kernel_size, dilation):
|
||||
kernel_size_effective = kernel_size + (kernel_size - 1) * (dilation - 1)
|
||||
pad_total = kernel_size_effective - 1
|
||||
pad_beg = pad_total // 2
|
||||
pad_end = pad_total - pad_beg
|
||||
padded_inputs = F.pad(inputs, (pad_beg, pad_end, pad_beg, pad_end))
|
||||
return padded_inputs
|
||||
|
||||
class SeparableConv2d(nn.Module):
|
||||
def __init__(self, inplanes, planes, kernel_size=3, stride=1, dilation=1, bias=False, BatchNorm=None):
|
||||
super(SeparableConv2d, self).__init__()
|
||||
|
||||
self.conv1 = nn.Conv2d(inplanes, inplanes, kernel_size, stride, 0, dilation,
|
||||
groups=inplanes, bias=bias)
|
||||
self.bn = BatchNorm(inplanes)
|
||||
self.pointwise = nn.Conv2d(inplanes, planes, 1, 1, 0, 1, 1, bias=bias)
|
||||
|
||||
def forward(self, x):
|
||||
x = fixed_padding(x, self.conv1.kernel_size[0], dilation=self.conv1.dilation[0])
|
||||
x = self.conv1(x)
|
||||
x = self.bn(x)
|
||||
x = self.pointwise(x)
|
||||
return x
|
||||
|
||||
class Block(nn.Module):
|
||||
def __init__(self, inplanes, planes, reps, stride=1, dilation=1, BatchNorm=None,
|
||||
start_with_relu=True, grow_first=True, is_last=False):
|
||||
super(Block, self).__init__()
|
||||
|
||||
if planes != inplanes or stride != 1:
|
||||
self.skip = nn.Conv2d(inplanes, planes, 1, stride=stride, bias=False)
|
||||
self.skipbn = BatchNorm(planes)
|
||||
else:
|
||||
self.skip = None
|
||||
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
rep = []
|
||||
|
||||
filters = inplanes
|
||||
if grow_first:
|
||||
rep.append(self.relu)
|
||||
rep.append(SeparableConv2d(inplanes, planes, 3, 1, dilation, BatchNorm=BatchNorm))
|
||||
rep.append(BatchNorm(planes))
|
||||
filters = planes
|
||||
|
||||
for i in range(reps - 1):
|
||||
rep.append(self.relu)
|
||||
rep.append(SeparableConv2d(filters, filters, 3, 1, dilation, BatchNorm=BatchNorm))
|
||||
rep.append(BatchNorm(filters))
|
||||
|
||||
if not grow_first:
|
||||
rep.append(self.relu)
|
||||
rep.append(SeparableConv2d(inplanes, planes, 3, 1, dilation, BatchNorm=BatchNorm))
|
||||
rep.append(BatchNorm(planes))
|
||||
|
||||
if stride != 1:
|
||||
rep.append(self.relu)
|
||||
rep.append(SeparableConv2d(planes, planes, 3, 2, BatchNorm=BatchNorm))
|
||||
rep.append(BatchNorm(planes))
|
||||
|
||||
if stride == 1 and is_last:
|
||||
rep.append(self.relu)
|
||||
rep.append(SeparableConv2d(planes, planes, 3, 1, BatchNorm=BatchNorm))
|
||||
rep.append(BatchNorm(planes))
|
||||
|
||||
if not start_with_relu:
|
||||
rep = rep[1:]
|
||||
|
||||
self.rep = nn.Sequential(*rep)
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.rep(inp)
|
||||
|
||||
if self.skip is not None:
|
||||
skip = self.skip(inp)
|
||||
skip = self.skipbn(skip)
|
||||
else:
|
||||
skip = inp
|
||||
|
||||
x = x + skip
|
||||
|
||||
return x
|
||||
|
||||
class AlignedXception(Backbone):
|
||||
"""
|
||||
Modified Alighed Xception
|
||||
"""
|
||||
def __init__(self, output_stride, BatchNorm):
|
||||
super(AlignedXception, self).__init__()
|
||||
|
||||
if output_stride == 16:
|
||||
entry_block3_stride = 2
|
||||
middle_block_dilation = 1
|
||||
exit_block_dilations = (1, 2)
|
||||
elif output_stride == 8:
|
||||
entry_block3_stride = 1
|
||||
middle_block_dilation = 2
|
||||
exit_block_dilations = (2, 4)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
# Entry flow
|
||||
self.conv1 = nn.Conv2d(3, 32, 3, stride=2, padding=1, bias=False)
|
||||
self.bn1 = BatchNorm(32)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
self.conv2 = nn.Conv2d(32, 64, 3, stride=1, padding=1, bias=False)
|
||||
self.bn2 = BatchNorm(64)
|
||||
|
||||
self.block1 = Block(64, 128, reps=2, stride=2, BatchNorm=BatchNorm, start_with_relu=False)
|
||||
self.block2 = Block(128, 256, reps=2, stride=2, BatchNorm=BatchNorm, start_with_relu=False,
|
||||
grow_first=True)
|
||||
self.block3 = Block(256, 728, reps=2, stride=entry_block3_stride, BatchNorm=BatchNorm,
|
||||
start_with_relu=True, grow_first=True, is_last=True)
|
||||
|
||||
# Middle flow
|
||||
self.block4 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block5 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block6 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block7 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block8 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block9 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block10 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block11 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block12 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block13 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block14 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block15 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block16 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block17 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block18 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block19 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
|
||||
# Exit flow
|
||||
self.block20 = Block(728, 1024, reps=2, stride=1, dilation=exit_block_dilations[0],
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=False, is_last=True)
|
||||
|
||||
self.conv3 = SeparableConv2d(1024, 1536, 3, stride=1, dilation=exit_block_dilations[1], BatchNorm=BatchNorm)
|
||||
self.bn3 = BatchNorm(1536)
|
||||
|
||||
self.conv4 = SeparableConv2d(1536, 1536, 3, stride=1, dilation=exit_block_dilations[1], BatchNorm=BatchNorm)
|
||||
self.bn4 = BatchNorm(1536)
|
||||
|
||||
self.conv5 = SeparableConv2d(1536, 2048, 3, stride=1, dilation=exit_block_dilations[1], BatchNorm=BatchNorm)
|
||||
self.bn5 = BatchNorm(2048)
|
||||
|
||||
# Init weights
|
||||
self._init_weight()
|
||||
|
||||
def forward(self, x):
|
||||
# Entry flow
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
|
||||
x = self.conv2(x)
|
||||
x = self.bn2(x)
|
||||
x = self.relu(x)
|
||||
|
||||
x = self.block1(x)
|
||||
# add relu here
|
||||
x = self.relu(x)
|
||||
low_level_feat = x
|
||||
x = self.block2(x)
|
||||
x = self.block3(x)
|
||||
|
||||
# Middle flow
|
||||
x = self.block4(x)
|
||||
x = self.block5(x)
|
||||
x = self.block6(x)
|
||||
x = self.block7(x)
|
||||
x = self.block8(x)
|
||||
x = self.block9(x)
|
||||
x = self.block10(x)
|
||||
x = self.block11(x)
|
||||
x = self.block12(x)
|
||||
x = self.block13(x)
|
||||
x = self.block14(x)
|
||||
x = self.block15(x)
|
||||
x = self.block16(x)
|
||||
x = self.block17(x)
|
||||
x = self.block18(x)
|
||||
x = self.block19(x)
|
||||
|
||||
# Exit flow
|
||||
x = self.block20(x)
|
||||
x = self.relu(x)
|
||||
x = self.conv3(x)
|
||||
x = self.bn3(x)
|
||||
x = self.relu(x)
|
||||
|
||||
x = self.conv4(x)
|
||||
x = self.bn4(x)
|
||||
x = self.relu(x)
|
||||
|
||||
x = self.conv5(x)
|
||||
x = self.bn5(x)
|
||||
x = self.relu(x)
|
||||
|
||||
return x, low_level_feat
|
||||
|
||||
def _init_weight(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
|
||||
m.weight.data.normal_(0, math.sqrt(2. / n))
|
||||
elif isinstance(m, nn.SyncBatchNorm):
|
||||
m.weight.data.fill_(1)
|
||||
m.bias.data.zero_()
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
m.weight.data.fill_(1)
|
||||
m.bias.data.zero_()
|
||||
|
||||
if __name__ == "__main__":
|
||||
import torch
|
||||
model = AlignedXception(BatchNorm=nn.BatchNorm2d, output_stride=16)
|
||||
input = torch.rand(1, 3, 512, 512)
|
||||
output, low_level_feat = model(input)
|
||||
print(output.size())
|
||||
print(low_level_feat.size())
|
||||
@@ -0,0 +1,330 @@
|
||||
import sys
|
||||
# sys.path.append("/Users/momo/human_seg_train")
|
||||
# print(sys.path)
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from bodyseg.backbone.backbone import get_norm
|
||||
|
||||
class ConvBNReLU(nn.Sequential):
|
||||
def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1, norm_layer=None):
|
||||
padding = (kernel_size - 1) // 2
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
super(ConvBNReLU, self).__init__(
|
||||
nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
|
||||
norm_layer(out_planes),
|
||||
nn.ReLU6(inplace=True)
|
||||
)
|
||||
|
||||
class InvertedResidual(nn.Module):
|
||||
def __init__(self, inp, oup, stride, expand_ratio, norm_layer=None):
|
||||
super(InvertedResidual, self).__init__()
|
||||
self.stride = stride
|
||||
assert stride in [1, 2]
|
||||
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
|
||||
hidden_dim = int(round(inp * expand_ratio))
|
||||
self.use_res_connect = self.stride == 1 and inp == oup
|
||||
|
||||
layers = []
|
||||
if expand_ratio != 1:
|
||||
# pw
|
||||
layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1, norm_layer=norm_layer))
|
||||
layers.extend([
|
||||
# dw
|
||||
ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim, norm_layer=norm_layer),
|
||||
# pw-linear
|
||||
nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
|
||||
norm_layer(oup),
|
||||
])
|
||||
self.conv = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
if self.use_res_connect:
|
||||
return x + self.conv(x)
|
||||
else:
|
||||
return self.conv(x)
|
||||
|
||||
class UpSampleBlock(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, expand_ratio=6):
|
||||
super(UpSampleBlock, self).__init__()
|
||||
self.refine = InvertedResidual(in_channels, out_channels, 1, expand_ratio)
|
||||
|
||||
def forward(self, x0, x1):
|
||||
x = torch.cat([x0, x1], dim=1)
|
||||
x = self.refine(x)
|
||||
return x
|
||||
|
||||
class BodySegNet_32_thin_sigmod(nn.Module):
|
||||
def __init__(self, input_channels=3, class_nums=1, output_onnx=False):
|
||||
super(BodySegNet_32_thin_sigmod, self).__init__()
|
||||
self.class_nums = class_nums
|
||||
self.output_onnx = output_onnx
|
||||
|
||||
self.stage_1 = nn.Sequential(
|
||||
ConvBNReLU(input_channels, 16, kernel_size=3, stride=2)
|
||||
)
|
||||
self.stage_2 = nn.Sequential(
|
||||
ConvBNReLU(16, 16, kernel_size=3, stride=2, groups=16),
|
||||
ConvBNReLU(16, 16, kernel_size=1, stride=1),
|
||||
)
|
||||
self.stage_3 = nn.Sequential(
|
||||
InvertedResidual(16, 24, stride=2, expand_ratio=6),
|
||||
InvertedResidual(24, 24, stride=1, expand_ratio=6),
|
||||
InvertedResidual(24, 24, stride=1, expand_ratio=6),
|
||||
)
|
||||
self.stage_4 = nn.Sequential(
|
||||
InvertedResidual(24, 32, stride=2, expand_ratio=6),
|
||||
InvertedResidual(32, 32, stride=1, expand_ratio=6),
|
||||
InvertedResidual(32, 32, stride=1, expand_ratio=6),
|
||||
InvertedResidual(32, 32, stride=1, expand_ratio=6),
|
||||
)
|
||||
self.stage_5 = nn.Sequential(
|
||||
InvertedResidual(32, 48, stride=2, expand_ratio=6),
|
||||
InvertedResidual(48, 48, stride=1, expand_ratio=6),
|
||||
InvertedResidual(48, 48, stride=1, expand_ratio=6),
|
||||
InvertedResidual(48, 48, stride=1, expand_ratio=6)
|
||||
)
|
||||
self.up_to_4 = UpSampleBlock(48 + 32, 16)
|
||||
self.up_to_3 = UpSampleBlock(16 + 24, 16)
|
||||
self.up_to_2 = UpSampleBlock(16 + 16, 16)
|
||||
self.up_to_1 = UpSampleBlock(16 + 16, 16)
|
||||
self.last_layer = nn.Sequential(
|
||||
ConvBNReLU(16, 16, kernel_size=1, stride=1),
|
||||
nn.Conv2d(16, self.class_nums, kernel_size=1, stride=1)
|
||||
)
|
||||
self._initialize_weights()
|
||||
|
||||
def forward(self, x):
|
||||
feature_S = []
|
||||
x1 = self.stage_1(x)
|
||||
x2 = self.stage_2(x1)
|
||||
x3 = self.stage_3(x2)
|
||||
x4 = self.stage_4(x3)
|
||||
feature = self.stage_5(x4)
|
||||
|
||||
feature = F.interpolate(feature, size=x4.size()[2:], mode='bilinear', align_corners=True)
|
||||
feature = self.up_to_4(x4, feature)
|
||||
feature = F.interpolate(feature, size=x3.size()[2:], mode='bilinear', align_corners=True)
|
||||
feature = self.up_to_3(x3, feature)
|
||||
feature = F.interpolate(feature, size=x2.size()[2:], mode='bilinear', align_corners=True)
|
||||
feature = self.up_to_2(x2, feature)
|
||||
feature = F.interpolate(feature, size=x1.size()[2:], mode='bilinear', align_corners=True)
|
||||
feature = self.up_to_1(x1, feature)
|
||||
feature_S.append(feature)
|
||||
feature = self.last_layer(feature)
|
||||
feature_S.append(feature)
|
||||
output = F.interpolate(feature, size=x.size()[2:], mode='bilinear', align_corners=True)
|
||||
# output = torch.sigmoid(output)
|
||||
if self.output_onnx:
|
||||
output = torch.argmax(output, dim=1).to(torch.float32)
|
||||
|
||||
return output, feature_S
|
||||
|
||||
def _initialize_weights(self):
|
||||
for name, m in self.named_modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
if 'first' in name:
|
||||
nn.init.normal_(m.weight, 0, 0.01)
|
||||
else:
|
||||
nn.init.normal_(m.weight, 0, 1.0 / m.weight.shape[1])
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0.0001)
|
||||
nn.init.constant_(m.running_mean, 0)
|
||||
elif isinstance(m, nn.BatchNorm1d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0.0001)
|
||||
nn.init.constant_(m.running_mean, 0)
|
||||
elif isinstance(m, nn.Linear):
|
||||
nn.init.normal_(m.weight, 0, 0.01)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
class _ASPPModule(nn.Module):
|
||||
def __init__(self, inplanes, planes, kernel_size, padding, dilation, BatchNorm):
|
||||
super(_ASPPModule, self).__init__()
|
||||
self.atrous_conv = nn.Conv2d(inplanes, planes, kernel_size=kernel_size,
|
||||
stride=1, padding=padding, dilation=dilation, bias=False)
|
||||
self.bn = BatchNorm(planes)
|
||||
self.relu = nn.ReLU()
|
||||
|
||||
self._init_weight()
|
||||
|
||||
def forward(self, x):
|
||||
x = self.atrous_conv(x)
|
||||
x = self.bn(x)
|
||||
|
||||
return self.relu(x)
|
||||
|
||||
def _init_weight(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
torch.nn.init.kaiming_normal_(m.weight)
|
||||
|
||||
class ASPP(nn.Module):
|
||||
def __init__(self, backbone, output_stride, BatchNorm):
|
||||
super(ASPP, self).__init__()
|
||||
if backbone == 'drn':
|
||||
inplanes = 512
|
||||
elif backbone == 'mobilenet':
|
||||
inplanes = 320
|
||||
elif backbone == 'resnet_fpn':
|
||||
inplanes = 256
|
||||
else:
|
||||
inplanes = 2048
|
||||
if output_stride == 16:
|
||||
dilations = [1, 6, 12, 18]
|
||||
elif output_stride == 8:
|
||||
dilations = [1, 12, 24, 36]
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
self.aspp1 = _ASPPModule(inplanes, 256, 1, padding=0, dilation=dilations[0], BatchNorm=BatchNorm)
|
||||
self.aspp2 = _ASPPModule(inplanes, 256, 3, padding=dilations[1], dilation=dilations[1], BatchNorm=BatchNorm)
|
||||
self.aspp3 = _ASPPModule(inplanes, 256, 3, padding=dilations[2], dilation=dilations[2], BatchNorm=BatchNorm)
|
||||
self.aspp4 = _ASPPModule(inplanes, 256, 3, padding=dilations[3], dilation=dilations[3], BatchNorm=BatchNorm)
|
||||
|
||||
self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
|
||||
nn.Conv2d(inplanes, 256, 1, stride=1, bias=False),
|
||||
BatchNorm(256),
|
||||
nn.ReLU())
|
||||
self.conv1 = nn.Conv2d(1280, 256, 1, bias=False)
|
||||
self.bn1 = BatchNorm(256)
|
||||
self.relu = nn.ReLU()
|
||||
self.dropout = nn.Dropout(0.5)
|
||||
self._init_weight()
|
||||
|
||||
def forward(self, x):
|
||||
x1 = self.aspp1(x)
|
||||
x2 = self.aspp2(x)
|
||||
x3 = self.aspp3(x)
|
||||
x4 = self.aspp4(x)
|
||||
x5 = self.global_avg_pool(x)
|
||||
x5 = F.interpolate(x5, size=x4.size()[2:], mode='bilinear', align_corners=True)
|
||||
# x5 = F.interpolate(x5, size=x4.size()[2:], mode='nearest')
|
||||
x = torch.cat((x1, x2, x3, x4, x5), dim=1)
|
||||
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
|
||||
return self.dropout(x)
|
||||
|
||||
def _init_weight(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
torch.nn.init.kaiming_normal_(m.weight)
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(self,num_classes, backbone, BatchNorm):
|
||||
super(Decoder, self).__init__()
|
||||
if backbone == 'resnet_fpn' or backbone == 'drn':
|
||||
low_level_inplanes = 256
|
||||
elif backbone == 'xception':
|
||||
low_level_inplanes = 128
|
||||
elif backbone == 'mobilenet':
|
||||
low_level_inplanes = 24
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
self.conv1 = nn.Conv2d(low_level_inplanes, 16, 1, bias=False)
|
||||
self.bn1 = BatchNorm(16)
|
||||
self.relu = nn.ReLU()
|
||||
self.last_conv = nn.Sequential(nn.Conv2d(304, 256, kernel_size=3, stride=1, padding=1, bias=False),
|
||||
BatchNorm(256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.5),
|
||||
nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=False),
|
||||
BatchNorm(256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.1),
|
||||
nn.Conv2d(256, num_classes, kernel_size=1, stride=1))
|
||||
|
||||
self.up_to_1 = UpSampleBlock(16 + 16, 16)
|
||||
self.last_layer = nn.Sequential(
|
||||
ConvBNReLU(16, 16, kernel_size=1, stride=1),
|
||||
nn.Conv2d(16, 1, kernel_size=1, stride=1)
|
||||
)
|
||||
self._init_weight()
|
||||
|
||||
# x(1,256,8,6) low(1,256,32,24) -》 x(1,1,32,24)
|
||||
def forward(self, x, low_level_feat):
|
||||
feature_T = []
|
||||
# deeplab part
|
||||
#(1,256,32,24) -> (1,16,64,48)
|
||||
low_level_feat = self.conv1(low_level_feat)
|
||||
low_level_feat = self.bn1(low_level_feat)
|
||||
low_level_feat = self.relu(low_level_feat)
|
||||
|
||||
# x(1, 256, 8, 6)-> (1,16,32,24)
|
||||
x = F.interpolate(x, size=low_level_feat.size()[2:], mode='bilinear', align_corners=True)
|
||||
low_level_feat = F.interpolate(low_level_feat,
|
||||
size=[low_level_feat.size()[2] * 2, low_level_feat.size()[3] * 2],
|
||||
mode='bilinear', align_corners=True)
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
|
||||
# bodyseg part
|
||||
feature = F.interpolate(x, size=[x.size()[2] * 2,x.size()[3] * 2], mode='bilinear', align_corners=True)
|
||||
# 输入up_to_1 x(low)(1,16,64,48),上采样2倍后的feature(1,16,64,48)
|
||||
feature = self.up_to_1(low_level_feat, feature)
|
||||
feature_T.append(feature)
|
||||
# (1,1,64,48)
|
||||
feature = self.last_layer(feature)
|
||||
feature_T.append(feature)
|
||||
# (1,1,128,96)
|
||||
output = F.interpolate(feature, size=[128,96], mode='bilinear', align_corners=True)
|
||||
# output = torch.sigmoid(output)
|
||||
|
||||
return output, feature_T
|
||||
|
||||
def _init_weight(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
torch.nn.init.kaiming_normal_(m.weight)
|
||||
|
||||
def build_backbone(backbone, output_stride, BatchNorm, input_channel=3):
|
||||
from bodyseg.backbone.fpn import build_resnet_fpn_backbone
|
||||
from bodyseg.backbone.xception import AlignedXception
|
||||
|
||||
return build_resnet_fpn_backbone(input_channel)
|
||||
|
||||
def build_aspp(backbone, output_stride, BatchNorm):
|
||||
return ASPP(backbone, output_stride, BatchNorm)
|
||||
|
||||
def build_decoder(num_classes, backbone, BatchNorm):
|
||||
return Decoder(num_classes, backbone, BatchNorm)
|
||||
|
||||
class DeepLab(nn.Module):
|
||||
def __init__(self, input_channel=3, class_num=1):
|
||||
super(DeepLab, self).__init__()
|
||||
|
||||
BatchNorm = get_norm("BN")
|
||||
|
||||
self.backbone = build_backbone("resnet_fpn", 16, BatchNorm, input_channel=input_channel)
|
||||
self.aspp = build_aspp("resnet_fpn", 16, BatchNorm)
|
||||
self.decoder = build_decoder(class_num, "resnet_fpn", BatchNorm)
|
||||
|
||||
def forward(self, input):
|
||||
# input(1,3,128,96)
|
||||
#output: "p2"(1,256,32,24), "p3"(1,256,16,12), "p4"(1,256,8,6)
|
||||
output = self.backbone(input)
|
||||
# "p4"(1,256,8,6) "p2"(1,256,32,24)
|
||||
x, low_level_feat = output['p4'], output['p2']
|
||||
# x(1,256,8,6)-》(1,256,8,6)
|
||||
x = self.aspp(x)
|
||||
# x(1,256,8,6) low(1,256,32,24) -》 x(1,1,32,24)
|
||||
x, feature_T = self.decoder(x, low_level_feat)
|
||||
# x(1,1,128,96)
|
||||
x = F.interpolate(x, size=input.size()[2:], mode='bilinear', align_corners=True)
|
||||
# x = F.interpolate(x, size=input.size()[2:], mode='nearest')
|
||||
return x, feature_T
|
||||
@@ -0,0 +1,98 @@
|
||||
import time
|
||||
|
||||
import cv2
|
||||
from hairstyle_model_infer import HairStyle_Model_Infer
|
||||
from utils import enhance_hair
|
||||
import numpy as np
|
||||
from utils import landmark_processor
|
||||
|
||||
|
||||
hairstyle_process = HairStyle_Model_Infer(gpu=True, use_enhance=True)
|
||||
|
||||
|
||||
def resize_pre_webui(input_img, mask_img, out_h=None, out_w=None):
|
||||
# 获取头发处理的局部区域图像
|
||||
box_info = hairstyle_process.get_body_info(input_img)
|
||||
dst_size = (576, 768)
|
||||
|
||||
if out_h is not None and out_w is not None:
|
||||
dst_size = (out_w, out_h)
|
||||
|
||||
box_w, box_h = box_info[2] - box_info[0], box_info[3] - box_info[1]
|
||||
scale = min(dst_size[1] / box_h, dst_size[0] / box_w)
|
||||
rotate_center = [(box_info[2] + box_info[0]) * 0.5, (box_info[3] + box_info[1]) * 0.5]
|
||||
M = cv2.getRotationMatrix2D(rotate_center, 0, scale)
|
||||
M[:, 2] += np.float32([dst_size[0] * 0.5, dst_size[1] * 0.5]) - np.float32(rotate_center)
|
||||
input_result = landmark_processor.high_quality_warpAffine(input_img, M, dst_size)
|
||||
|
||||
crop_matting = cv2.warpAffine(mask_img, M, dst_size)
|
||||
mask_result = (crop_matting > 10).astype(np.float32)
|
||||
mask_dilate = cv2.dilate(mask_result, np.ones((3, 9), np.uint8))
|
||||
mask_dilate = np.clip(mask_dilate * 255, 0, 255).astype(np.uint8)
|
||||
# cv2.imshow("input_result", input_result)
|
||||
# cv2.imwrite("input_result.png", input_result)
|
||||
# cv2.imshow("mask_dilate", mask_dilate)
|
||||
# cv2.imwrite("mask_dilate.png", mask_dilate)
|
||||
# cv2.waitKey(0)
|
||||
|
||||
return input_result, mask_dilate, M
|
||||
|
||||
def process_infer(user_img_path, target_color, color_dir, dst_path, ratio=1.0):
|
||||
# 读取用户原始图片
|
||||
origin_img = cv2.imread(user_img_path)
|
||||
|
||||
# resize用户原始图片，防止图片过大
|
||||
user_scale = 1920 / max(origin_img.shape[0], origin_img.shape[1])
|
||||
if user_scale < 1.0:
|
||||
# 缩小图像应该使用INTER_AREA插值
|
||||
origin_img = cv2.resize(origin_img, (0, 0), fx=user_scale, fy=user_scale, interpolation=cv2.INTER_AREA)
|
||||
|
||||
# 首先完成第一阶段的换发色
|
||||
s1 = time.time()
|
||||
hair_color_changed_img, hair_mask, status = hairstyle_process.infer_haircolor_v4(origin_img, target_hair_color=target_color, haircolor_dir=color_dir)
|
||||
print("----------------------------infer_haircolor_v4", time.time()-s1)
|
||||
|
||||
hair_color_changed_img = (hair_color_changed_img * ratio + origin_img * (1 - ratio)).astype(np.uint8)
|
||||
|
||||
# 通过webui给照片加上标签
|
||||
s2 = time.time()
|
||||
# prompt = enhance_hair.webui_tag_by_clip(hair_color_changed_img)
|
||||
prompt = ""
|
||||
print("----------------------------webui_tag_by_clip", time.time() - s2)
|
||||
|
||||
# 处理图像尺寸
|
||||
s3 = time.time()
|
||||
crop_img, crop_mask, M = resize_pre_webui(hair_color_changed_img, hair_mask, 1024, 768)
|
||||
print(crop_img.shape)
|
||||
print(crop_mask.shape)
|
||||
print("----------------------------resize_pre_webui", time.time() - s3)
|
||||
|
||||
# 通过webui对照片进行增强
|
||||
s4 = time.time()
|
||||
# cv2.imwrite('/home/student/Downloads/12121.jpg',crop_img)
|
||||
enhanced_img = enhance_hair.webui_img2img(crop_img, crop_mask, prompt=prompt)
|
||||
print("----------------------------webui_img2img", time.time() - s4)
|
||||
|
||||
# 写回原图
|
||||
origin_img_final = origin_img.copy()
|
||||
# restore the sd_result_small to origin_img
|
||||
M_inv = cv2.invertAffineTransform(M)
|
||||
cv2.warpAffine(enhanced_img, M_inv, (origin_img_final.shape[1], origin_img_final.shape[0]),
|
||||
dst=origin_img_final,
|
||||
borderMode=cv2.BORDER_TRANSPARENT, flags=cv2.INTER_LANCZOS4)
|
||||
# cv2.imshow("origin_img_final", origin_img_final)
|
||||
# cv2.waitKey(0)
|
||||
|
||||
cv2.imwrite(dst_path, origin_img_final)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
user_img_path = '/home/student/Downloads/base_color_v2.jpg'
|
||||
# target_hair_color = [28, 55, 1] # RGB
|
||||
# target_hair_color = [76, 55, 36] # RGB
|
||||
# target_hair_color = [75, 1, 0] # RGB
|
||||
target_hair_color = [247,235,213] # RGB
|
||||
color_dir = "/home/data/hair/data/ref_color/247_235_213"
|
||||
dst = "/home/data/hair/data/tmp/4.png"
|
||||
|
||||
process_infer(user_img_path=user_img_path, target_color=target_hair_color, color_dir=color_dir, dst_path=dst)
|
||||
@@ -0,0 +1,140 @@
|
||||
#coding:utf-8
|
||||
import traceback
|
||||
from uuid import uuid4
|
||||
import imghdr
|
||||
|
||||
import torch
|
||||
from gevent import monkey
|
||||
monkey.patch_all()
|
||||
|
||||
import base64
|
||||
import os
|
||||
import random
|
||||
import shutil
|
||||
import time
|
||||
import json
|
||||
import os.path as osp
|
||||
import urllib.request
|
||||
import hashlib
|
||||
import cv2
|
||||
from datetime import datetime
|
||||
import glob
|
||||
import numpy as np
|
||||
from core.hairstyle_model import HairStyle_Model
|
||||
from hairstyle_model_infer import HairStyle_Model_Infer
|
||||
from prepare_ref_hairstyle_data import prepare_single, prepare_single_color
|
||||
|
||||
|
||||
from utils import enhance_hair
|
||||
import configparser
|
||||
from common.logger import config
|
||||
|
||||
from change_color import process_infer, resize_pre_webui
|
||||
|
||||
|
||||
|
||||
hairstyle_process = HairStyle_Model(gpu=True,use_enhance=True)
|
||||
user_img_save_dir = config.get('default', 'userDir')
|
||||
user_img_tmp_dir = config.get('default', 'tmp_dir')
|
||||
user_img_res_dir = config.get('default', 'res_dir')
|
||||
ref_user_dir = config.get('default', 'ref_user_dir')
|
||||
train_save_dir = config.get('default', 'train_dir')
|
||||
hair_template_material_dir = config.get('default', 'hair_template_material_dir')
|
||||
ref_color_dir = config.get('default', 'ref_color')
|
||||
ref_color_imgs_dir = config.get('default', 'ref_color_img')
|
||||
train_upload_dir = config.get('default', 'upload_train_dir')
|
||||
|
||||
def download_img(img_url, userId=None, isfix=False, ismask=False):
|
||||
try:
|
||||
img_name = img_url.split("/")[-1]
|
||||
tmp_dir = osp.join(user_img_tmp_dir, img_name)
|
||||
# if osp.exists(tmp_dir):
|
||||
# os.remove(tmp_dir)
|
||||
print(img_url)
|
||||
download_success = False
|
||||
for i in range(3):
|
||||
hairstyle_process.oss2.download_img(img_url, tmp_dir)
|
||||
if osp.exists(tmp_dir) and osp.getsize(tmp_dir) > 0:
|
||||
download_success = True
|
||||
break
|
||||
if download_success:
|
||||
img_type = imghdr.what(tmp_dir)
|
||||
new_tmp_dir = tmp_dir[:tmp_dir.rfind(".")+1] + img_type
|
||||
shutil.move(tmp_dir, new_tmp_dir)
|
||||
print("save path", new_tmp_dir)
|
||||
|
||||
else:
|
||||
return None, None
|
||||
return new_tmp_dir, None
|
||||
except Exception as e:
|
||||
print(e)
|
||||
return None,None
|
||||
|
||||
|
||||
def change_hair_colorv3():
|
||||
datanow = datetime.now()
|
||||
time_convert = datanow.strftime("%Y%m%d%H")
|
||||
taskid = ''.join(str(random.choice(range(10))) for _ in range(6))
|
||||
taskid = str(time_convert) + str(taskid)
|
||||
start_time0 = time.time()
|
||||
|
||||
|
||||
try:
|
||||
|
||||
img_url = "https://ydapp-1317132355.cos.ap-beijing.myqcloud.com/hair_mz/images/hairstyle/b8e75994-ec6c-4309-ac98-d9687ca5dfbf/ad28024a-2a98-4768-b1d3-846950f4dc57.jpg"
|
||||
userId = "17519"
|
||||
color = ''
|
||||
rgb = [103, 103, 103]
|
||||
|
||||
ratio = 0.9
|
||||
|
||||
img_path, _ = download_img(img_url, userId)
|
||||
print("---------------download img:", time.time()-start_time0)
|
||||
|
||||
start_time1 = time.time()
|
||||
# gen material
|
||||
color_name = img_path[img_path.rfind("/") + 1:]
|
||||
print("color_name", color_name)
|
||||
new_color_ref_img_path = os.path.join(ref_color_imgs_dir, color_name)
|
||||
print("new_color_ref_img_path: ", new_color_ref_img_path)
|
||||
shutil.copy(img_path, new_color_ref_img_path)
|
||||
|
||||
color_id = ""
|
||||
for color_item in rgb:
|
||||
color_id += str(color_item) + "_"
|
||||
if color_id[-1] == "_":
|
||||
color_id = color_id[:-1]
|
||||
|
||||
ref_color_save_dir = os.path.join(ref_color_dir, color_id)
|
||||
if not os.path.exists(ref_color_save_dir):
|
||||
os.mkdir(ref_color_save_dir)
|
||||
print("!!!gen color material:", ref_color_save_dir)
|
||||
prepare_single_color(new_color_ref_img_path, ref_color_save_dir)
|
||||
|
||||
req_id = str(uuid4())
|
||||
res_path = os.path.join(user_img_res_dir, req_id + ".png")
|
||||
|
||||
print("---------------gen material:", time.time() - start_time1)
|
||||
|
||||
start_time2 = time.time()
|
||||
process_infer(img_path, rgb, ref_color_save_dir, res_path, ratio)
|
||||
print("res path: ", res_path)
|
||||
print("---------------process_infer:", time.time() - start_time2)
|
||||
|
||||
start_time3 = time.time()
|
||||
|
||||
if os.path.exists(res_path):
|
||||
ret_url = hairstyle_process.oss2.upload_file(res_path,
|
||||
"hair_mz/images/hairstyle/{}/{}".format(color,
|
||||
|
||||
req_id + '.jpg'))
|
||||
print("---------------upload img:", time.time() - start_time3)
|
||||
|
||||
|
||||
|
||||
except Exception as e:
|
||||
print(e)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
change_hair_colorv3()
|
||||
@@ -0,0 +1,46 @@
|
||||
import requests
|
||||
import json
|
||||
|
||||
callback_hairstyle_url = "https://puton.meidaojia.com/api/cloth/callBack"
|
||||
|
||||
def callback_color(color_id, cover_img, success):
|
||||
url = "http://172.21.0.3:8080/ydapp/system/platform/complete_color_model"
|
||||
payload = json.dumps({
|
||||
"success": success,
|
||||
"colorId": color_id,
|
||||
"coverImg": cover_img
|
||||
})
|
||||
headers = {
|
||||
'X-MZ-API-TOKEN': '98bcf37c942a2240ba2d907c96ed1137',
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
response = requests.request("POST", url, headers=headers, data=payload)
|
||||
return response
|
||||
# print(response.text)
|
||||
|
||||
|
||||
def callback_hairstyle(req_id, state, message, clothId):
|
||||
url = callback_hairstyle_url
|
||||
payload = json.dumps({
|
||||
"taskId": req_id,
|
||||
"status": state,
|
||||
"clothId": clothId,
|
||||
"msg": message
|
||||
})
|
||||
|
||||
status = -1
|
||||
|
||||
print(f"url:{url},complete_hair_model payload:{payload}")
|
||||
|
||||
headers = {
|
||||
'X-MZ-API-TOKEN': '98bcf37c942a2240ba2d907c96ed1137',
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
try:
|
||||
response = requests.request("POST", url, headers=headers, data=payload)
|
||||
print("response:", response.text)
|
||||
status = 0
|
||||
except Exception as e:
|
||||
print(e)
|
||||
|
||||
return response.text, status
|
||||
@@ -0,0 +1,27 @@
|
||||
import pynvml
|
||||
threshold = 0.9
|
||||
def get_gpu(need_gpu_id):
|
||||
used = get_gpu_threshold(need_gpu_id)
|
||||
#小于一定的
|
||||
if used > threshold:
|
||||
need_use = get_use_gpu()
|
||||
return need_use[0]
|
||||
else:
|
||||
return need_gpu_id
|
||||
|
||||
def get_use_gpu():
|
||||
use=[]
|
||||
for index in range(pynvml.nvmlDeviceGetCount()):
|
||||
used = get_gpu_threshold(index)
|
||||
if used > threshold:
|
||||
use.append(index)
|
||||
return use
|
||||
|
||||
def get_gpu_count():
|
||||
return pynvml.nvmlDeviceGetCount()
|
||||
|
||||
def get_gpu_threshold(index):
|
||||
handle = pynvml.nvmlDeviceGetHandleByIndex(index)
|
||||
meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
|
||||
used = meminfo.used / meminfo.total
|
||||
return used
|
||||
@@ -0,0 +1,63 @@
|
||||
#!/usr/bin/python
|
||||
# coding=utf-8
|
||||
import os
|
||||
from logging.handlers import TimedRotatingFileHandler
|
||||
import logging
|
||||
import configparser
|
||||
config = configparser.ConfigParser() # 创建对象
|
||||
config.read("config/configure.ini", encoding="utf-8") # 读取配置文件，如果配置文件不存在则创建
|
||||
|
||||
"""
|
||||
自定义日志处理
|
||||
"""
|
||||
|
||||
|
||||
class LogFactory(object):
|
||||
|
||||
@staticmethod
|
||||
def getLogger(log_name,log_level=None):
|
||||
if log_level is None:
|
||||
log_level = LogFactory.getLogLevel(getLevel())
|
||||
logger = logging.getLogger(log_name)
|
||||
path = getPath()
|
||||
isExists=os.path.exists(path)
|
||||
if not isExists:
|
||||
os.makedirs(path)
|
||||
log_file = os.path.join(path, '{}.log'.format(log_name))
|
||||
if len(logger.handlers) <= 0:
|
||||
handler = TimedRotatingFileHandler(log_file, when='D', interval=1)
|
||||
else:
|
||||
handler = logger.handlers[0]
|
||||
formatter = logging.Formatter("%(asctime)s-%(thread)d-%(filename)s-%(levelname)s-%(message)s", "%Y-%m-%d %H:%M:%S")
|
||||
handler.setFormatter(formatter)
|
||||
logger.addHandler(handler)
|
||||
logger.setLevel(log_level)
|
||||
return logger
|
||||
|
||||
@staticmethod
|
||||
def getLogLevel(log_level):
|
||||
log_level = getattr(logging, log_level.upper(), None)
|
||||
if log_level is None:
|
||||
raise Exception("No such log level.")
|
||||
return log_level
|
||||
|
||||
|
||||
def getPath():
|
||||
import sys
|
||||
process_order = sys.argv[1] if len(sys.argv) > 1 else None
|
||||
logpath = config.get('logger', "logpath")
|
||||
logpath = logpath.rstrip("/")
|
||||
|
||||
if process_order is None:
|
||||
return logpath
|
||||
|
||||
return logpath
|
||||
|
||||
|
||||
def getLevel():
|
||||
return config.get('logger', "level")
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
[default]
|
||||
modelDir = weights
|
||||
hairstyleDir = __BASE_DIR__/data/ref_hairstyle
|
||||
haircolorDir= __BASE_DIR__/data/ref_haircolor
|
||||
userDir=__BASE_DIR__/data/userImage
|
||||
tmp_dir=__BASE_DIR__/data/tmp
|
||||
res_dir=__BASE_DIR__/data/res_dir
|
||||
userInfo_dir=__BASE_DIR__/data/user_info
|
||||
baseColor_ID=HDR10_443322
|
||||
Port = 11023
|
||||
refer_dir = __BASE_DIR__/data/ref_online
|
||||
ref_user_dir = __BASE_DIR__/data/ref_user_imgs
|
||||
train_dir = __BASE_DIR__/kohya_ss_home/train_material
|
||||
refImgDir=__BASE_DIR__/data/refImage
|
||||
hair_template_material_dir=__BASE_DIR__/data/hair_template_material
|
||||
ref_color=__BASE_DIR__/data/ref_color
|
||||
ref_color_img=__BASE_DIR__/data/ref_color_imgs
|
||||
upload_train_dir=__BASE_DIR__/data/upload_train_imgs
|
||||
;version=local
|
||||
version=online
|
||||
|
||||
[fix]
|
||||
strength=1
|
||||
|
||||
[logger]
|
||||
logpath = __BASE_DIR__/logs
|
||||
level=INFO
|
||||
|
||||
[timelogger]
|
||||
name=watch-time
|
||||
|
||||
[errorlogger]
|
||||
level=ERROR
|
||||
name=w-error
|
||||
@@ -0,0 +1,116 @@
|
||||
import os
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
def copy_files_by_id(id_list_file, source_dirs, output_dir):
|
||||
"""
|
||||
根据ID列表拷贝文件到新目录，保持目录结构
|
||||
|
||||
参数:
|
||||
id_list_file: 包含文件ID列表的文本文件
|
||||
source_dirs: 源目录列表 [hair_template_material, ref_hairstyle, train_material, upload_train_imgs]
|
||||
output_dir: 输出目录
|
||||
"""
|
||||
# 读取ID列表
|
||||
with open(id_list_file, 'r') as f:
|
||||
id_list = [line.strip() for line in f.readlines() if line.strip()]
|
||||
|
||||
# 确保输出目录存在
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# 遍历所有源目录
|
||||
for src_dir in source_dirs:
|
||||
if not os.path.exists(src_dir):
|
||||
print(f"警告: 源目录不存在 {src_dir}")
|
||||
continue
|
||||
|
||||
# 遍历源目录下的所有文件
|
||||
for root, _, files in os.walk(src_dir):
|
||||
for file in files:
|
||||
# 检查文件名是否包含ID列表中的任一ID
|
||||
if any(file_id in file for file_id in id_list):
|
||||
src_path = os.path.join(root, file)
|
||||
|
||||
# 计算相对路径
|
||||
rel_path = os.path.relpath(root, src_dir)
|
||||
dst_dir = os.path.join(output_dir, os.path.basename(src_dir), rel_path)
|
||||
|
||||
# 创建目标目录并拷贝文件
|
||||
os.makedirs(dst_dir, exist_ok=True)
|
||||
dst_path = os.path.join(dst_dir, file)
|
||||
shutil.copy2(src_path, dst_path)
|
||||
print(f"已拷贝: {src_path} -> {dst_path}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
# import argparse
|
||||
|
||||
# parser = argparse.ArgumentParser(description='根据ID列表拷贝文件')
|
||||
# parser.add_argument('id_list', help='包含文件ID列表的文本文件路径')
|
||||
# parser.add_argument('output_dir', help='输出目录路径')
|
||||
# parser.add_argument('--source_dirs', nargs='+',
|
||||
# default=['hair_template_material', 'ref_hairstyle', 'train_material', 'upload_train_imgs'],
|
||||
# help='源目录列表')
|
||||
|
||||
# args = parser.parse_args()
|
||||
id_list = [1939351512506802177,
|
||||
1939351392453238785,
|
||||
1939351319849836545,
|
||||
1939351238526476289,
|
||||
1939351219396255745,
|
||||
1939339435616608258,
|
||||
1939330299344560130,
|
||||
1938848782956732417,
|
||||
1938848730150445058,
|
||||
1938848709308948482,
|
||||
1938841049238970370,
|
||||
1938492830059438081,
|
||||
1938492790033195009,
|
||||
1938490058547240961,
|
||||
1938489994072399874,
|
||||
1938489915521474561,
|
||||
1938489840707674114,
|
||||
1938489739088076801,
|
||||
1938489674713899010,
|
||||
1938489620162781185]
|
||||
output_dir = '/home/data/hair/data/hairStyle_addin/06300845'
|
||||
hair_template_material = '/home/data/hair/data/hair_template_material'
|
||||
ref_hairstyle = '/home/data/hair/data/ref_hairstyle'
|
||||
train_material = '/mnt/database2/online-server/hair-online-tj-v2/train_material'
|
||||
upload_train_imgs = '/home/data/hair/data/upload_train_imgs'
|
||||
copy_files_by_id(id_list, source_dirs, args.output_dir)
|
||||
|
||||
def copy_files_by_id(id_list, hair_template_material, ref_hairstyle, train_material, upload_train_imgs, output_dir):
|
||||
"""
|
||||
根据ID列表从指定目录拷贝文件到输出目录
|
||||
:param id_list: 文件ID列表(字符串格式)
|
||||
:param hair_template_material: 发型模板材质目录
|
||||
:param ref_hairstyle: 参考发型目录
|
||||
:param train_material: 训练材质目录
|
||||
:param upload_train_imgs: 上传训练图片目录
|
||||
:param output_dir: 输出目录
|
||||
"""
|
||||
# 将ID转换为字符串格式
|
||||
id_list = [str(id) for id in id_list]
|
||||
# 创建源目录列表
|
||||
source_dirs = [hair_template_material, ref_hairstyle, train_material, upload_train_imgs]
|
||||
for src_dir in source_dirs:
|
||||
if not os.path.exists(src_dir):
|
||||
print(f"警告: 源目录不存在 {src_dir}")
|
||||
continue
|
||||
|
||||
# 遍历源目录下的所有文件
|
||||
for root, _, files in os.walk(src_dir):
|
||||
for file in files:
|
||||
# 检查文件名是否包含ID列表中的任一ID
|
||||
if any(file_id in file for file_id in id_list):
|
||||
src_path = os.path.join(root, file)
|
||||
|
||||
# 计算相对路径
|
||||
rel_path = os.path.relpath(root, src_dir)
|
||||
dst_dir = os.path.join(output_dir, os.path.basename(src_dir), rel_path)
|
||||
|
||||
# 创建目标目录并拷贝文件
|
||||
os.makedirs(dst_dir, exist_ok=True)
|
||||
dst_path = os.path.join(dst_dir, file)
|
||||
shutil.copy2(src_path, dst_path)
|
||||
print(f"已拷贝: {src_path} -> {dst_path}")
|
||||
@@ -0,0 +1,196 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import math
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
from core.utils import landmark_processor
|
||||
import glob
|
||||
# from models.BigResNetStable import MomocvFaceAlignment
|
||||
|
||||
|
||||
def op_name(op_name, m):
|
||||
m.op_name = op_name
|
||||
return m
|
||||
|
||||
class Flatten(nn.Module):
|
||||
def __init__(self, axis=1):
|
||||
super(Flatten, self).__init__()
|
||||
self.axis = axis
|
||||
|
||||
def forward(self, x):
|
||||
assert self.axis == 1
|
||||
x = x.reshape(x.shape[0], -1)
|
||||
return x
|
||||
|
||||
def flatten(name, axis=1):
|
||||
return op_name(name, Flatten(axis))
|
||||
|
||||
def conv_relu(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1):
|
||||
return nn.Sequential(
|
||||
op_name(name, nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, True)),
|
||||
op_name(name.replace('conv', 'relu'), nn.LeakyReLU(inplace=False, negative_slope=5e-11)),
|
||||
)
|
||||
|
||||
def conv(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1):
|
||||
return op_name(name, nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, True))
|
||||
|
||||
def bn_conv_relu(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, momentum = 0.9, track_running_stats=True):
|
||||
return nn.Sequential(
|
||||
op_name(name + '_bn1', nn.BatchNorm2d(in_channels, momentum=momentum, eps=2e-5, track_running_stats=track_running_stats)),
|
||||
op_name(name + '_conv1', nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, True)),
|
||||
op_name(name + '_relu', nn.LeakyReLU(inplace=False, negative_slope=5e-11)),
|
||||
)
|
||||
|
||||
def bn(name, in_channels, momentum = 0.9, track_running_stats=True):
|
||||
return op_name(name, nn.BatchNorm2d(in_channels, momentum=momentum, eps=2e-5, track_running_stats=track_running_stats))
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
def __init__(self, stage, unit, inplanes, outplanes, stride):
|
||||
super(BasicBlock, self).__init__()
|
||||
|
||||
self.bn = bn('stage{}_unit{}_bn1'.format(stage, unit), inplanes)
|
||||
self.conv1 = conv_relu('stage{}_unit{}_conv1'.format(stage, unit), inplanes, outplanes, kernel_size=3, stride=1, padding=1)
|
||||
self.conv2 = conv('stage{}_unit{}_conv2'.format(stage, unit), outplanes, outplanes, kernel_size=3, stride=stride, padding=1)
|
||||
|
||||
if inplanes != outplanes or stride != 1:
|
||||
self.conv_sc = conv('stage{}_unit{}_conv1sc'.format(stage, unit), inplanes, outplanes, kernel_size=1, stride=stride, padding=0)
|
||||
|
||||
self.inplanes = inplanes
|
||||
self.outplanes = outplanes
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
residule = x
|
||||
out = self.bn(x)
|
||||
out = self.conv1(out)
|
||||
out = self.conv2(out)
|
||||
if self.inplanes != self.outplanes or self.stride != 1:
|
||||
residule = self.conv_sc(residule)
|
||||
ret = out + residule
|
||||
return ret
|
||||
|
||||
class ResnetBlock(nn.Module):
|
||||
def __init__(self, stage, inplanes, outplanes, stride=2, n_blocks=1):
|
||||
super(ResnetBlock, self).__init__()
|
||||
|
||||
self.conv = []
|
||||
for m in range(n_blocks):
|
||||
if m == 0:
|
||||
self.conv.append(BasicBlock(stage, m + 1, inplanes, outplanes, stride))
|
||||
else:
|
||||
self.conv.append(BasicBlock(stage, m + 1, outplanes, outplanes, 1))
|
||||
self.conv = nn.Sequential(*self.conv)
|
||||
|
||||
def forward(self, x):
|
||||
ret = self.conv(x)
|
||||
return ret
|
||||
|
||||
class FaceRecognitionServer(nn.Module):
|
||||
def __init__(self):
|
||||
super(FaceRecognitionServer, self).__init__()
|
||||
ch_num = [64, 64, 128, 256, 512]
|
||||
# ch_num = [64, 64]
|
||||
strides = [2, 2, 2, 2]
|
||||
n_blocks = [3, 4, 14, 3]
|
||||
op_list = [conv_relu('conv0', 3, 64, 3, 1, 1)]
|
||||
op_list += [ResnetBlock(i + 1, ch_num[i], ch_num[i + 1], strides[i], n_blocks[i]) for i in range(len(ch_num) - 1)]
|
||||
op_list += [bn('bn1', 512)]
|
||||
op_list += [flatten('flatten', 1)]
|
||||
op_list += [op_name('pre_fc1', nn.Linear(25088, 512))]
|
||||
self.features = nn.Sequential(*op_list)
|
||||
|
||||
model_path = os.path.dirname(os.path.split(os.path.realpath(__file__))[0])
|
||||
weights = torch.load('weights/MMCVFaceRecognitionServer.pth',
|
||||
map_location=lambda storage, loc: storage)
|
||||
self.load_state_dict(weights)
|
||||
self.eval()
|
||||
|
||||
def forward(self, x):
|
||||
features = self.features(x)
|
||||
return features
|
||||
|
||||
class MomocvFaceRecognitionServer(object):
|
||||
def __init__(self, gpu_id=None):
|
||||
self.gpu_id = gpu_id
|
||||
self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||
self.face_recognition_net = FaceRecognitionServer()
|
||||
# self.mmcv = MomocvFaceAlignment()
|
||||
# self.model_path, _ = os.path.split(os.path.realpath(__file__))
|
||||
# weights = torch.load(os.path.join(self.model_path, 'MMCVFaceRecognitionServer.pth'),
|
||||
# map_location=lambda storage, loc: storage)
|
||||
# self.face_recognition_net.load_state_dict(weights)
|
||||
self.face_recognition_net.to(self.device)
|
||||
# self.face_recognition_net.eval()
|
||||
|
||||
def forward(self, images, landmarks):
|
||||
dst_size = 112
|
||||
|
||||
with torch.no_grad():
|
||||
input_numpy = np.zeros((len(images), 3, dst_size, dst_size), dtype=np.float32)
|
||||
assert len(images) == len(landmarks)
|
||||
for ix, img in enumerate(images):
|
||||
landmark = landmarks[ix]
|
||||
|
||||
|
||||
# get angel
|
||||
# M = landmark_processor.get_transform_mat_full_face(landmark, 576, scale=1, offset=(0, 0.3))
|
||||
# pt1k_crop = landmark_processor.transform_points(landmark, M)
|
||||
# crop_face = cv2.warpAffine(img, M, (576, 768),
|
||||
# flags=cv2.INTER_LANCZOS4, borderMode=cv2.BORDER_REFLECT)
|
||||
# t0 = time.time()
|
||||
# landmarks87, poselayer, tracking_probe, occlusion_probe = self.mmcv.detect(crop_face, [pt1k_crop])
|
||||
# # print('time:', time.time() - t0)
|
||||
# if tracking_probe[0] < 0.5:
|
||||
# continue
|
||||
|
||||
mat = landmark_processor.get_transform_mat_for_face_recognition(landmark, dst_size)[0:2]
|
||||
|
||||
tmp = cv2.warpAffine(img, mat, (dst_size, dst_size))
|
||||
|
||||
input_numpy[ix, :, :, :] = (tmp.transpose((2, 0, 1)).astype(np.float32) - 127.5) / 128
|
||||
|
||||
# cv2.imshow('tmp', tmp)
|
||||
# cv2.waitKey()
|
||||
|
||||
in_tensor = torch.from_numpy(input_numpy)
|
||||
in_tensor = in_tensor.to(self.device)
|
||||
features = self.face_recognition_net(in_tensor)
|
||||
features = features.cpu().numpy()
|
||||
norm_factor = np.sqrt(np.sum(features ** 2, axis=1))
|
||||
norm_factor = norm_factor.reshape(-1, 1)
|
||||
features /= norm_factor
|
||||
return features
|
||||
|
||||
def cos_sim(self, a, b):
|
||||
a_norm = np.linalg.norm(a)
|
||||
b_norm = np.linalg.norm(b)
|
||||
cos = np.dot(a, b) / (a_norm * b_norm)
|
||||
return cos
|
||||
|
||||
if __name__ == '__main__':
|
||||
all_jpegs = glob.glob(r'D:\data\deepface_example\data\02b59e75ce91bda300ff827a85a74687d633cdfb5d830e7d3855e31b75201fa8\*.jpg')
|
||||
for s_filename_path in all_jpegs:
|
||||
img = cv2.imread(s_filename_path)
|
||||
|
||||
# cv2.imshow('img', img)
|
||||
# cv2.waitKey()
|
||||
|
||||
# dflpng = DATAIMG(str(s_filename_path), print_on_no_embedded_data=True)
|
||||
# if dflpng is None:
|
||||
# print('ERROR')
|
||||
#
|
||||
# landmarks = dflpng.get_landmarks_mmcv_137()
|
||||
#
|
||||
# mmcv = MomocvFaceRecognitionServer()
|
||||
# features = mmcv.forward([img, img], [landmarks, landmarks])
|
||||
# print(features)
|
||||
# print('conansherry')
|
||||
# fullyconnected1 = fullyconnected1[0]
|
||||
# fullyconnected1 = (np.reshape(fullyconnected1, (2, 87)).transpose((1, 0))).astype(np.int32)
|
||||
# occlusion_probe = occlusion_probe[0]
|
||||
# for ix, pt in enumerate(fullyconnected1):
|
||||
# cv2.circle(img, (int(pt[0]), int(pt[1])), 1, (0, 255, 0) if occlusion_probe[ix] > 0.1 else (0, 0, 255), 2)
|
||||
# # cv2.putText(img, str(ix), (int(pt[0]), int(pt[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 255, 0), 1)
|
||||
# cv2.imshow('img', img)
|
||||
# cv2.waitKey()
|
||||
Binary file not shown.
@@ -0,0 +1,75 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
def get_norm(norm, out_channels=None):
|
||||
"""
|
||||
Args:
|
||||
norm (str or callable):
|
||||
|
||||
Returns:
|
||||
nn.Module or None: the normalization layer
|
||||
"""
|
||||
if isinstance(norm, str):
|
||||
if len(norm) == 0:
|
||||
return None
|
||||
norm = {
|
||||
"BN": nn.BatchNorm2d,
|
||||
"IN": nn.InstanceNorm2d,
|
||||
"GN": lambda channels: nn.GroupNorm(32, channels),
|
||||
"nnSyncBN": nn.SyncBatchNorm, # keep for debugging
|
||||
}[norm]
|
||||
if out_channels is not None:
|
||||
return norm(out_channels)
|
||||
else:
|
||||
return norm
|
||||
|
||||
class Conv2d(torch.nn.Conv2d):
|
||||
"""
|
||||
A wrapper around :class:`torch.nn.Conv2d` to support zero-size tensor and more features.
|
||||
"""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
"""
|
||||
Extra keyword arguments supported in addition to those in `torch.nn.Conv2d`:
|
||||
|
||||
Args:
|
||||
norm (nn.Module, optional): a normalization layer
|
||||
activation (callable(Tensor) -> Tensor): a callable activation function
|
||||
|
||||
It assumes that norm layer is used before activation.
|
||||
"""
|
||||
norm = kwargs.pop("norm", None)
|
||||
activation = kwargs.pop("activation", None)
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
self.norm = norm
|
||||
self.activation = activation
|
||||
|
||||
def forward(self, x):
|
||||
x = super().forward(x)
|
||||
if self.norm is not None:
|
||||
x = self.norm(x)
|
||||
if self.activation is not None:
|
||||
x = self.activation(x)
|
||||
return x
|
||||
|
||||
class Backbone(nn.Module):
|
||||
"""
|
||||
Abstract base class for network backbones.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""
|
||||
The `__init__` method of any subclass can specify its own set of arguments.
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
def forward(self):
|
||||
"""
|
||||
Subclasses must override this method, but adhere to the same return type.
|
||||
|
||||
Returns:
|
||||
dict[str: Tensor]: mapping from feature name (e.g., "res2") to tensor
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,268 @@
|
||||
import math
|
||||
import core.utils.weight_init as weight_init
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from core.bodyseg.backbone.backbone import Backbone, get_norm, Conv2d
|
||||
from core.bodyseg.backbone.resnet import build_resnet_backbone
|
||||
|
||||
class FPN(Backbone):
|
||||
"""
|
||||
This module implements Feature Pyramid Network.
|
||||
It creates pyramid features built on top of some input feature maps.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, bottom_up, in_features, out_channels, norm="", top_block=None, fuse_type="sum"
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
bottom_up (Backbone): module representing the bottom up subnetwork.
|
||||
Must be a subclass of :class:`Backbone`. The multi-scale feature
|
||||
maps generated by the bottom up network, and listed in `in_features`,
|
||||
are used to generate FPN levels.
|
||||
in_features (list[str]): names of the input feature maps coming
|
||||
from the backbone to which FPN is attached. For example, if the
|
||||
backbone produces ["res2", "res3", "res4"], any *contiguous* sublist
|
||||
of these may be used; order must be from high to low resolution.
|
||||
out_channels (int): number of channels in the output feature maps.
|
||||
norm (str): the normalization to use.
|
||||
top_block (nn.Module or None): if provided, an extra operation will
|
||||
be performed on the output of the last (smallest resolution)
|
||||
FPN output, and the result will extend the result list. The top_block
|
||||
further downsamples the feature map. It must have an attribute
|
||||
"num_levels", meaning the number of extra FPN levels added by
|
||||
this block, and "in_feature", which is a string representing
|
||||
its input feature (e.g., p5).
|
||||
fuse_type (str): types for fusing the top down features and the lateral
|
||||
ones. It can be "sum" (default), which sums up element-wise; or "avg",
|
||||
which takes the element-wise mean of the two.
|
||||
"""
|
||||
super(FPN, self).__init__()
|
||||
assert isinstance(bottom_up, Backbone)
|
||||
|
||||
# Feature map strides and channels from the bottom up network (e.g. ResNet)
|
||||
in_strides = [bottom_up._out_feature_strides[f] for f in in_features]
|
||||
in_channels = [bottom_up._out_feature_channels[f] for f in in_features]
|
||||
|
||||
_assert_strides_are_log2_contiguous(in_strides)
|
||||
lateral_convs = []
|
||||
output_convs = []
|
||||
|
||||
use_bias = norm == ""
|
||||
for idx, in_channels in enumerate(in_channels):
|
||||
lateral_norm = get_norm(norm, out_channels)
|
||||
output_norm = get_norm(norm, out_channels)
|
||||
|
||||
lateral_conv = Conv2d(
|
||||
in_channels, out_channels, kernel_size=1, bias=use_bias, norm=lateral_norm
|
||||
)
|
||||
output_conv = Conv2d(
|
||||
out_channels,
|
||||
out_channels,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
bias=use_bias,
|
||||
norm=output_norm,
|
||||
)
|
||||
weight_init.c2_xavier_fill(lateral_conv)
|
||||
weight_init.c2_xavier_fill(output_conv)
|
||||
stage = int(math.log2(in_strides[idx]))
|
||||
|
||||
lateral_convs.append(lateral_conv)
|
||||
output_convs.append(output_conv)
|
||||
# Place convs into top-down order (from low to high resolution)
|
||||
# to make the top-down computation in forward clearer.
|
||||
self.lateral_convs = nn.ModuleList(lateral_convs[::-1])
|
||||
self.output_convs = nn.ModuleList(output_convs[::-1])
|
||||
self.top_block = top_block
|
||||
self.in_features = in_features
|
||||
self.bottom_up = bottom_up
|
||||
# Return feature names are "p<stage>", like ["p2", "p3", ..., "p6"]
|
||||
self._out_feature_strides = {"p{}".format(int(math.log2(s))): s for s in in_strides}
|
||||
# top block output feature maps.
|
||||
if self.top_block is not None:
|
||||
for s in range(stage, stage + self.top_block.num_levels):
|
||||
self._out_feature_strides["p{}".format(s + 1)] = 2 ** (s + 1)
|
||||
|
||||
self._out_features = list(self._out_feature_strides.keys())
|
||||
self._out_feature_channels = {k: out_channels for k in self._out_features}
|
||||
assert fuse_type in {"avg", "sum"}
|
||||
self._fuse_type = fuse_type
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Args:
|
||||
input (dict[str: Tensor]): mapping feature map name (e.g., "res5") to
|
||||
feature map tensor for each feature level in high to low resolution order.
|
||||
|
||||
Returns:
|
||||
dict[str: Tensor]:
|
||||
mapping from feature map name to FPN feature map tensor
|
||||
in high to low resolution order. Returned feature names follow the FPN
|
||||
paper convention: "p<stage>", where stage has stride = 2 ** stage e.g.,
|
||||
["p2", "p3", ..., "p6"].
|
||||
"""
|
||||
# Reverse feature maps into top-down order (from low to high resolution)
|
||||
bottom_up_features = self.bottom_up(x)
|
||||
x = [bottom_up_features[f] for f in self.in_features[::-1]]
|
||||
results = []
|
||||
prev_features = self.lateral_convs[0](x[0])
|
||||
results.append(self.output_convs[0](prev_features))
|
||||
for features, lateral_conv, output_conv in zip(
|
||||
x[1:], self.lateral_convs[1:], self.output_convs[1:]
|
||||
):
|
||||
top_down_features = F.interpolate(prev_features, scale_factor=2, mode="nearest")
|
||||
lateral_features = lateral_conv(features)
|
||||
prev_features = lateral_features + top_down_features
|
||||
if self._fuse_type == "avg":
|
||||
prev_features /= 2
|
||||
results.insert(0, output_conv(prev_features))
|
||||
|
||||
if self.top_block is not None:
|
||||
top_block_in_feature = bottom_up_features.get(self.top_block.in_feature, None)
|
||||
if top_block_in_feature is None:
|
||||
top_block_in_feature = results[self._out_features.index(self.top_block.in_feature)]
|
||||
results.extend(self.top_block(top_block_in_feature))
|
||||
assert len(self._out_features) == len(results)
|
||||
return dict(zip(self._out_features, results))
|
||||
|
||||
def _assert_strides_are_log2_contiguous(strides):
|
||||
"""
|
||||
Assert that each stride is 2x times its preceding stride, i.e. "contiguous in log2".
|
||||
"""
|
||||
for i, stride in enumerate(strides[1:], 1):
|
||||
assert stride == 2 * strides[i - 1], "Strides {} {} are not log2 contiguous".format(
|
||||
stride, strides[i - 1]
|
||||
)
|
||||
|
||||
|
||||
class LastLevelMaxPool(nn.Module):
|
||||
"""
|
||||
This module is used in the original FPN to generate a downsampled
|
||||
P6 feature from P5.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.num_levels = 1
|
||||
self.in_feature = "p5"
|
||||
|
||||
def forward(self, x):
|
||||
return [F.max_pool2d(x, kernel_size=1, stride=2, padding=0)]
|
||||
|
||||
|
||||
class LastLevelP6P7(nn.Module):
|
||||
"""
|
||||
This module is used in RetinaNet to generate extra layers, P6 and P7 from
|
||||
C5 feature.
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels, out_channels):
|
||||
super().__init__()
|
||||
self.num_levels = 2
|
||||
self.in_feature = "res5"
|
||||
self.p6 = nn.Conv2d(in_channels, out_channels, 3, 2, 1)
|
||||
self.p7 = nn.Conv2d(out_channels, out_channels, 3, 2, 1)
|
||||
for module in [self.p6, self.p7]:
|
||||
weight_init.c2_xavier_fill(module)
|
||||
|
||||
def forward(self, c5):
|
||||
p6 = self.p6(c5)
|
||||
p7 = self.p7(F.relu(p6))
|
||||
return [p6, p7]
|
||||
|
||||
|
||||
def build_resnet_fpn_backbone(in_channels=3):
|
||||
"""
|
||||
Args:
|
||||
cfg: a detectron2 CfgNode
|
||||
|
||||
Returns:
|
||||
backbone (Backbone): backbone module, must be a subclass of :class:`Backbone`.
|
||||
"""
|
||||
bottom_up = build_resnet_backbone(in_channels)
|
||||
in_features = ["res2", "res3", "res4"]
|
||||
out_channels = 256
|
||||
backbone = FPN(
|
||||
bottom_up=bottom_up,
|
||||
in_features=in_features,
|
||||
out_channels=out_channels,
|
||||
norm="BN",
|
||||
# top_block=LastLevelMaxPool(),
|
||||
top_block=None,
|
||||
fuse_type="sum",
|
||||
)
|
||||
return backbone
|
||||
|
||||
def build_retinanet_resnet_fpn_backbone(cfg, in_channels=3):
|
||||
"""
|
||||
Args:
|
||||
cfg: a detectron2 CfgNode
|
||||
|
||||
Returns:
|
||||
backbone (Backbone): backbone module, must be a subclass of :class:`Backbone`.
|
||||
"""
|
||||
bottom_up = build_resnet_backbone(cfg, in_channels)
|
||||
in_features = cfg.MODEL.FPN.IN_FEATURES
|
||||
out_channels = cfg.MODEL.FPN.OUT_CHANNELS
|
||||
in_channels_p6p7 = bottom_up._out_feature_channels["res5"]
|
||||
backbone = FPN(
|
||||
bottom_up=bottom_up,
|
||||
in_features=in_features,
|
||||
out_channels=out_channels,
|
||||
norm=cfg.MODEL.FPN.NORM,
|
||||
top_block=LastLevelP6P7(in_channels_p6p7, out_channels),
|
||||
fuse_type=cfg.MODEL.FPN.FUSE_TYPE,
|
||||
)
|
||||
return backbone
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
from config.default import get_cfg
|
||||
|
||||
def setup(args):
|
||||
"""
|
||||
Create configs and perform basic setups.
|
||||
"""
|
||||
cfg = get_cfg()
|
||||
cfg.merge_from_file(args.cfg)
|
||||
cfg.merge_from_list(args.opts)
|
||||
cfg.freeze()
|
||||
return cfg
|
||||
|
||||
parser = argparse.ArgumentParser(description='Train ImageNet network')
|
||||
# general
|
||||
parser.add_argument('--cfg',
|
||||
help='experiment configure file name',
|
||||
required=True,
|
||||
type=str)
|
||||
|
||||
parser.add_argument('opts',
|
||||
help="Modify config options using the command-line",
|
||||
default=None,
|
||||
nargs=argparse.REMAINDER)
|
||||
|
||||
args = parser.parse_args()
|
||||
cfg = setup(args)
|
||||
print(cfg)
|
||||
|
||||
model = build_resnet_fpn_backbone(cfg, 3)
|
||||
# model = build_retinanet_resnet_fpn_backbone(cfg, 3)
|
||||
print(model)
|
||||
# model = torch.nn.DataParallel(model, list(range(2))).cuda()
|
||||
dummy_input = torch.randn(4, 3, 512, 512)
|
||||
|
||||
out = model(dummy_input)
|
||||
|
||||
for k, v in out.items():
|
||||
print(k, v.shape)
|
||||
|
||||
# torch.onnx.export(model, dummy_input, "tmp.onnx", verbose=True,
|
||||
# input_names=['input'],
|
||||
# output_names=['output'])
|
||||
|
||||
pass
|
||||
|
||||
@@ -0,0 +1,298 @@
|
||||
import numpy as np
|
||||
import core.utils.weight_init as weight_init
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
from core.bodyseg.backbone.backbone import Backbone, get_norm, Conv2d
|
||||
|
||||
class BasicStem(nn.Module):
|
||||
def __init__(self, in_channels=3, out_channels=64, norm="BN"):
|
||||
"""
|
||||
Args:
|
||||
norm (str or callable): a callable that takes the number of
|
||||
channels and return a `nn.Module`, or a pre-defined string
|
||||
(one of {"FrozenBN", "BN", "GN"}).
|
||||
"""
|
||||
super().__init__()
|
||||
self.conv1 = Conv2d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size=7,
|
||||
stride=2,
|
||||
padding=3,
|
||||
bias=False,
|
||||
norm=get_norm(norm, out_channels),
|
||||
)
|
||||
weight_init.c2_msra_fill(self.conv1)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = F.relu_(x)
|
||||
x = F.max_pool2d(x, kernel_size=3, stride=2, padding=1)
|
||||
return x
|
||||
|
||||
@property
|
||||
def out_channels(self):
|
||||
return self.conv1.out_channels
|
||||
|
||||
@property
|
||||
def stride(self):
|
||||
return 4 # = stride 2 conv -> stride 2 max pool
|
||||
|
||||
class ResNetBlockBase(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, stride):
|
||||
"""
|
||||
The `__init__` method of any subclass should also contain these arguments.
|
||||
|
||||
Args:
|
||||
in_channels (int):
|
||||
out_channels (int):
|
||||
stride (int):
|
||||
"""
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.stride = stride
|
||||
|
||||
class BottleneckBlock(ResNetBlockBase):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
*,
|
||||
bottleneck_channels,
|
||||
stride=1,
|
||||
num_groups=1,
|
||||
norm="BN",
|
||||
stride_in_1x1=False,
|
||||
dilation=1,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
norm (str or callable): a callable that takes the number of
|
||||
channels and return a `nn.Module`, or a pre-defined string
|
||||
(one of {"FrozenBN", "BN", "GN"}).
|
||||
stride_in_1x1 (bool): when stride==2, whether to put stride in the
|
||||
first 1x1 convolution or the bottleneck 3x3 convolution.
|
||||
"""
|
||||
super().__init__(in_channels, out_channels, stride)
|
||||
|
||||
if in_channels != out_channels:
|
||||
self.shortcut = Conv2d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size=1,
|
||||
stride=stride,
|
||||
bias=False,
|
||||
norm=get_norm(norm, out_channels),
|
||||
)
|
||||
else:
|
||||
self.shortcut = None
|
||||
|
||||
# The original MSRA ResNet models have stride in the first 1x1 conv
|
||||
# The subsequent fb.torch.resnet and Caffe2 ResNe[X]t implementations have
|
||||
# stride in the 3x3 conv
|
||||
stride_1x1, stride_3x3 = (stride, 1) if stride_in_1x1 else (1, stride)
|
||||
|
||||
self.conv1 = Conv2d(
|
||||
in_channels,
|
||||
bottleneck_channels,
|
||||
kernel_size=1,
|
||||
stride=stride_1x1,
|
||||
bias=False,
|
||||
norm=get_norm(norm, bottleneck_channels),
|
||||
)
|
||||
|
||||
self.conv2 = Conv2d(
|
||||
bottleneck_channels,
|
||||
bottleneck_channels,
|
||||
kernel_size=3,
|
||||
stride=stride_3x3,
|
||||
padding=1 * dilation,
|
||||
bias=False,
|
||||
groups=num_groups,
|
||||
dilation=dilation,
|
||||
norm=get_norm(norm, bottleneck_channels),
|
||||
)
|
||||
|
||||
self.conv3 = Conv2d(
|
||||
bottleneck_channels,
|
||||
out_channels,
|
||||
kernel_size=1,
|
||||
bias=False,
|
||||
norm=get_norm(norm, out_channels),
|
||||
)
|
||||
|
||||
for layer in [self.conv1, self.conv2, self.conv3, self.shortcut]:
|
||||
if layer is not None: # shortcut can be None
|
||||
weight_init.c2_msra_fill(layer)
|
||||
|
||||
# Zero-initialize the last normalization in each residual branch,
|
||||
# so that at the beginning, the residual branch starts with zeros,
|
||||
# and each residual block behaves like an identity.
|
||||
# See Sec 5.1 in "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour":
|
||||
# "For BN layers, the learnable scaling coefficient γ is initialized
|
||||
# to be 1, except for each residual block's last BN
|
||||
# where γ is initialized to be 0."
|
||||
|
||||
# nn.init.constant_(self.conv3.norm.weight, 0)
|
||||
# TODO this somehow hurts performance when training GN models from scratch.
|
||||
# Add it as an option when we need to use this code to train a backbone.
|
||||
|
||||
def forward(self, x):
|
||||
out = self.conv1(x)
|
||||
out = F.relu_(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = F.relu_(out)
|
||||
|
||||
out = self.conv3(out)
|
||||
|
||||
if self.shortcut is not None:
|
||||
shortcut = self.shortcut(x)
|
||||
else:
|
||||
shortcut = x
|
||||
|
||||
out += shortcut
|
||||
out = F.relu_(out)
|
||||
return out
|
||||
|
||||
def make_stage(block_class, num_blocks, first_stride, **kwargs):
|
||||
"""
|
||||
Create a resnet stage by creating many blocks.
|
||||
Args:
|
||||
block_class (class): a subclass of ResNetBlockBase
|
||||
num_blocks (int):
|
||||
first_stride (int): the stride of the first block. The other blocks will have stride=1.
|
||||
A `stride` argument will be passed to the block constructor.
|
||||
kwargs: other arguments passed to the block constructor.
|
||||
|
||||
Returns:
|
||||
list[nn.Module]: a list of block module.
|
||||
"""
|
||||
blocks = []
|
||||
for i in range(num_blocks):
|
||||
blocks.append(block_class(stride=first_stride if i == 0 else 1, **kwargs))
|
||||
kwargs["in_channels"] = kwargs["out_channels"]
|
||||
return blocks
|
||||
|
||||
class ResNet(Backbone):
|
||||
def __init__(self, stem, stages, num_classes=None, out_features=None):
|
||||
"""
|
||||
Args:
|
||||
stem (nn.Module): a stem module
|
||||
stages (list[list[ResNetBlock]]): several (typically 4) stages,
|
||||
each contains multiple :class:`ResNetBlockBase`.
|
||||
num_classes (None or int): if None, will not perform classification.
|
||||
out_features (list[str]): name of the layers whose outputs should
|
||||
be returned in forward. Can be anything in "stem", "linear", or "res2" ...
|
||||
If None, will return the output of the last layer.
|
||||
"""
|
||||
super(ResNet, self).__init__()
|
||||
self.stem = stem
|
||||
self.num_classes = num_classes
|
||||
|
||||
current_stride = self.stem.stride
|
||||
self._out_feature_strides = {"stem": current_stride}
|
||||
self._out_feature_channels = {"stem": self.stem.out_channels}
|
||||
|
||||
self.stages = []
|
||||
self.names = []
|
||||
for i, blocks in enumerate(stages):
|
||||
for block in blocks:
|
||||
assert isinstance(block, ResNetBlockBase), block
|
||||
curr_channels = block.out_channels
|
||||
stage = nn.Sequential(*blocks)
|
||||
name = "res" + str(i + 2)
|
||||
self.stages.append(stage)
|
||||
self.names.append(name)
|
||||
self._out_feature_strides[name] = current_stride = int(
|
||||
current_stride * np.prod([k.stride for k in blocks])
|
||||
)
|
||||
self._out_feature_channels[name] = blocks[-1].out_channels
|
||||
self.stages = nn.ModuleList(self.stages)
|
||||
|
||||
if num_classes is not None:
|
||||
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
|
||||
self.linear = nn.Linear(curr_channels, num_classes)
|
||||
|
||||
# Sec 5.1 in "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour":
|
||||
# "The 1000-way fully-connected layer is initialized by
|
||||
# drawing weights from a zero-mean Gaussian with standard deviation of 0.01."
|
||||
nn.init.normal_(self.linear.weight, stddev=0.01)
|
||||
name = "linear"
|
||||
|
||||
if out_features is None:
|
||||
out_features = [name]
|
||||
self._out_features = out_features
|
||||
assert len(self._out_features)
|
||||
for out_feature in self._out_features:
|
||||
assert out_feature in self.names, "Available children: {}".format(", ".join(self.names))
|
||||
|
||||
def forward(self, x):
|
||||
outputs = {}
|
||||
x = self.stem(x)
|
||||
if "stem" in self._out_features:
|
||||
outputs["stem"] = x
|
||||
for ix, stage in enumerate(self.stages):
|
||||
name = self.names[ix]
|
||||
x = stage(x)
|
||||
if name in self._out_features:
|
||||
outputs[name] = x
|
||||
if self.num_classes is not None:
|
||||
x = self.avgpool(x)
|
||||
x = self.linear(x)
|
||||
if "linear" in self._out_features:
|
||||
outputs["linear"] = x
|
||||
return outputs
|
||||
|
||||
def build_resnet_backbone(in_channels=3):
|
||||
norm = "BN"
|
||||
stem = BasicStem(
|
||||
in_channels=in_channels,
|
||||
out_channels=64,
|
||||
norm=norm,
|
||||
)
|
||||
|
||||
# fmt: off
|
||||
out_features = ["res2", "res3", "res4"]
|
||||
depth = 101
|
||||
num_groups = 1
|
||||
bottleneck_channels = 64
|
||||
in_channels = 64
|
||||
out_channels = 256
|
||||
stride_in_1x1 = True
|
||||
res5_dilation = 1
|
||||
# fmt: on
|
||||
assert res5_dilation in {1, 2}, "res5_dilation cannot be {}.".format(res5_dilation)
|
||||
|
||||
num_blocks_per_stage = {50: [3, 4, 6, 3], 101: [3, 4, 23, 3], 152: [3, 8, 36, 3]}[depth]
|
||||
|
||||
stages = []
|
||||
|
||||
# Avoid creating variables without gradients
|
||||
# It consumes extra memory and may cause allreduce to fail
|
||||
out_stage_idx = [{"res2": 2, "res3": 3, "res4": 4, "res5": 5}[f] for f in out_features]
|
||||
max_stage_idx = max(out_stage_idx)
|
||||
for idx, stage_idx in enumerate(range(2, max_stage_idx + 1)):
|
||||
dilation = res5_dilation if stage_idx == 5 else 1
|
||||
first_stride = 1 if idx == 0 or (stage_idx == 5 and dilation == 2) else 2
|
||||
stage_kargs = dict()
|
||||
stage_kargs.update({
|
||||
"num_blocks": num_blocks_per_stage[idx],
|
||||
"first_stride": first_stride,
|
||||
"in_channels": in_channels,
|
||||
"bottleneck_channels": bottleneck_channels,
|
||||
"out_channels": out_channels,
|
||||
"num_groups": num_groups,
|
||||
"norm": norm,
|
||||
"stride_in_1x1": stride_in_1x1,
|
||||
"dilation": dilation,
|
||||
})
|
||||
stage_kargs["block_class"] = BottleneckBlock
|
||||
blocks = make_stage(**stage_kargs)
|
||||
in_channels = out_channels
|
||||
out_channels *= 2
|
||||
bottleneck_channels *= 2
|
||||
stages.append(blocks)
|
||||
return ResNet(stem, stages, out_features=out_features)
|
||||
@@ -0,0 +1,244 @@
|
||||
import math
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.utils.model_zoo as model_zoo
|
||||
from core.bodyseg.backbone.backbone import Backbone
|
||||
|
||||
def fixed_padding(inputs, kernel_size, dilation):
|
||||
kernel_size_effective = kernel_size + (kernel_size - 1) * (dilation - 1)
|
||||
pad_total = kernel_size_effective - 1
|
||||
pad_beg = pad_total // 2
|
||||
pad_end = pad_total - pad_beg
|
||||
padded_inputs = F.pad(inputs, (pad_beg, pad_end, pad_beg, pad_end))
|
||||
return padded_inputs
|
||||
|
||||
class SeparableConv2d(nn.Module):
|
||||
def __init__(self, inplanes, planes, kernel_size=3, stride=1, dilation=1, bias=False, BatchNorm=None):
|
||||
super(SeparableConv2d, self).__init__()
|
||||
|
||||
self.conv1 = nn.Conv2d(inplanes, inplanes, kernel_size, stride, 0, dilation,
|
||||
groups=inplanes, bias=bias)
|
||||
self.bn = BatchNorm(inplanes)
|
||||
self.pointwise = nn.Conv2d(inplanes, planes, 1, 1, 0, 1, 1, bias=bias)
|
||||
|
||||
def forward(self, x):
|
||||
x = fixed_padding(x, self.conv1.kernel_size[0], dilation=self.conv1.dilation[0])
|
||||
x = self.conv1(x)
|
||||
x = self.bn(x)
|
||||
x = self.pointwise(x)
|
||||
return x
|
||||
|
||||
class Block(nn.Module):
|
||||
def __init__(self, inplanes, planes, reps, stride=1, dilation=1, BatchNorm=None,
|
||||
start_with_relu=True, grow_first=True, is_last=False):
|
||||
super(Block, self).__init__()
|
||||
|
||||
if planes != inplanes or stride != 1:
|
||||
self.skip = nn.Conv2d(inplanes, planes, 1, stride=stride, bias=False)
|
||||
self.skipbn = BatchNorm(planes)
|
||||
else:
|
||||
self.skip = None
|
||||
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
rep = []
|
||||
|
||||
filters = inplanes
|
||||
if grow_first:
|
||||
rep.append(self.relu)
|
||||
rep.append(SeparableConv2d(inplanes, planes, 3, 1, dilation, BatchNorm=BatchNorm))
|
||||
rep.append(BatchNorm(planes))
|
||||
filters = planes
|
||||
|
||||
for i in range(reps - 1):
|
||||
rep.append(self.relu)
|
||||
rep.append(SeparableConv2d(filters, filters, 3, 1, dilation, BatchNorm=BatchNorm))
|
||||
rep.append(BatchNorm(filters))
|
||||
|
||||
if not grow_first:
|
||||
rep.append(self.relu)
|
||||
rep.append(SeparableConv2d(inplanes, planes, 3, 1, dilation, BatchNorm=BatchNorm))
|
||||
rep.append(BatchNorm(planes))
|
||||
|
||||
if stride != 1:
|
||||
rep.append(self.relu)
|
||||
rep.append(SeparableConv2d(planes, planes, 3, 2, BatchNorm=BatchNorm))
|
||||
rep.append(BatchNorm(planes))
|
||||
|
||||
if stride == 1 and is_last:
|
||||
rep.append(self.relu)
|
||||
rep.append(SeparableConv2d(planes, planes, 3, 1, BatchNorm=BatchNorm))
|
||||
rep.append(BatchNorm(planes))
|
||||
|
||||
if not start_with_relu:
|
||||
rep = rep[1:]
|
||||
|
||||
self.rep = nn.Sequential(*rep)
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.rep(inp)
|
||||
|
||||
if self.skip is not None:
|
||||
skip = self.skip(inp)
|
||||
skip = self.skipbn(skip)
|
||||
else:
|
||||
skip = inp
|
||||
|
||||
x = x + skip
|
||||
|
||||
return x
|
||||
|
||||
class AlignedXception(Backbone):
|
||||
"""
|
||||
Modified Alighed Xception
|
||||
"""
|
||||
def __init__(self, output_stride, BatchNorm):
|
||||
super(AlignedXception, self).__init__()
|
||||
|
||||
if output_stride == 16:
|
||||
entry_block3_stride = 2
|
||||
middle_block_dilation = 1
|
||||
exit_block_dilations = (1, 2)
|
||||
elif output_stride == 8:
|
||||
entry_block3_stride = 1
|
||||
middle_block_dilation = 2
|
||||
exit_block_dilations = (2, 4)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
# Entry flow
|
||||
self.conv1 = nn.Conv2d(3, 32, 3, stride=2, padding=1, bias=False)
|
||||
self.bn1 = BatchNorm(32)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
self.conv2 = nn.Conv2d(32, 64, 3, stride=1, padding=1, bias=False)
|
||||
self.bn2 = BatchNorm(64)
|
||||
|
||||
self.block1 = Block(64, 128, reps=2, stride=2, BatchNorm=BatchNorm, start_with_relu=False)
|
||||
self.block2 = Block(128, 256, reps=2, stride=2, BatchNorm=BatchNorm, start_with_relu=False,
|
||||
grow_first=True)
|
||||
self.block3 = Block(256, 728, reps=2, stride=entry_block3_stride, BatchNorm=BatchNorm,
|
||||
start_with_relu=True, grow_first=True, is_last=True)
|
||||
|
||||
# Middle flow
|
||||
self.block4 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block5 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block6 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block7 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block8 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block9 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block10 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block11 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block12 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block13 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block14 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block15 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block16 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block17 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block18 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
self.block19 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
|
||||
|
||||
# Exit flow
|
||||
self.block20 = Block(728, 1024, reps=2, stride=1, dilation=exit_block_dilations[0],
|
||||
BatchNorm=BatchNorm, start_with_relu=True, grow_first=False, is_last=True)
|
||||
|
||||
self.conv3 = SeparableConv2d(1024, 1536, 3, stride=1, dilation=exit_block_dilations[1], BatchNorm=BatchNorm)
|
||||
self.bn3 = BatchNorm(1536)
|
||||
|
||||
self.conv4 = SeparableConv2d(1536, 1536, 3, stride=1, dilation=exit_block_dilations[1], BatchNorm=BatchNorm)
|
||||
self.bn4 = BatchNorm(1536)
|
||||
|
||||
self.conv5 = SeparableConv2d(1536, 2048, 3, stride=1, dilation=exit_block_dilations[1], BatchNorm=BatchNorm)
|
||||
self.bn5 = BatchNorm(2048)
|
||||
|
||||
# Init weights
|
||||
self._init_weight()
|
||||
|
||||
def forward(self, x):
|
||||
# Entry flow
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
|
||||
x = self.conv2(x)
|
||||
x = self.bn2(x)
|
||||
x = self.relu(x)
|
||||
|
||||
x = self.block1(x)
|
||||
# add relu here
|
||||
x = self.relu(x)
|
||||
low_level_feat = x
|
||||
x = self.block2(x)
|
||||
x = self.block3(x)
|
||||
|
||||
# Middle flow
|
||||
x = self.block4(x)
|
||||
x = self.block5(x)
|
||||
x = self.block6(x)
|
||||
x = self.block7(x)
|
||||
x = self.block8(x)
|
||||
x = self.block9(x)
|
||||
x = self.block10(x)
|
||||
x = self.block11(x)
|
||||
x = self.block12(x)
|
||||
x = self.block13(x)
|
||||
x = self.block14(x)
|
||||
x = self.block15(x)
|
||||
x = self.block16(x)
|
||||
x = self.block17(x)
|
||||
x = self.block18(x)
|
||||
x = self.block19(x)
|
||||
|
||||
# Exit flow
|
||||
x = self.block20(x)
|
||||
x = self.relu(x)
|
||||
x = self.conv3(x)
|
||||
x = self.bn3(x)
|
||||
x = self.relu(x)
|
||||
|
||||
x = self.conv4(x)
|
||||
x = self.bn4(x)
|
||||
x = self.relu(x)
|
||||
|
||||
x = self.conv5(x)
|
||||
x = self.bn5(x)
|
||||
x = self.relu(x)
|
||||
|
||||
return x, low_level_feat
|
||||
|
||||
def _init_weight(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
|
||||
m.weight.data.normal_(0, math.sqrt(2. / n))
|
||||
elif isinstance(m, nn.SyncBatchNorm):
|
||||
m.weight.data.fill_(1)
|
||||
m.bias.data.zero_()
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
m.weight.data.fill_(1)
|
||||
m.bias.data.zero_()
|
||||
|
||||
if __name__ == "__main__":
|
||||
import torch
|
||||
model = AlignedXception(BatchNorm=nn.BatchNorm2d, output_stride=16)
|
||||
input = torch.rand(1, 3, 512, 512)
|
||||
output, low_level_feat = model(input)
|
||||
print(output.size())
|
||||
print(low_level_feat.size())
|
||||
@@ -0,0 +1,330 @@
|
||||
import sys
|
||||
# sys.path.append("/Users/momo/human_seg_train")
|
||||
# print(sys.path)
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from core.bodyseg.backbone.backbone import get_norm
|
||||
|
||||
class ConvBNReLU(nn.Sequential):
|
||||
def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1, norm_layer=None):
|
||||
padding = (kernel_size - 1) // 2
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
super(ConvBNReLU, self).__init__(
|
||||
nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
|
||||
norm_layer(out_planes),
|
||||
nn.ReLU6(inplace=True)
|
||||
)
|
||||
|
||||
class InvertedResidual(nn.Module):
|
||||
def __init__(self, inp, oup, stride, expand_ratio, norm_layer=None):
|
||||
super(InvertedResidual, self).__init__()
|
||||
self.stride = stride
|
||||
assert stride in [1, 2]
|
||||
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
|
||||
hidden_dim = int(round(inp * expand_ratio))
|
||||
self.use_res_connect = self.stride == 1 and inp == oup
|
||||
|
||||
layers = []
|
||||
if expand_ratio != 1:
|
||||
# pw
|
||||
layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1, norm_layer=norm_layer))
|
||||
layers.extend([
|
||||
# dw
|
||||
ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim, norm_layer=norm_layer),
|
||||
# pw-linear
|
||||
nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
|
||||
norm_layer(oup),
|
||||
])
|
||||
self.conv = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
if self.use_res_connect:
|
||||
return x + self.conv(x)
|
||||
else:
|
||||
return self.conv(x)
|
||||
|
||||
class UpSampleBlock(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, expand_ratio=6):
|
||||
super(UpSampleBlock, self).__init__()
|
||||
self.refine = InvertedResidual(in_channels, out_channels, 1, expand_ratio)
|
||||
|
||||
def forward(self, x0, x1):
|
||||
x = torch.cat([x0, x1], dim=1)
|
||||
x = self.refine(x)
|
||||
return x
|
||||
|
||||
class BodySegNet_32_thin_sigmod(nn.Module):
|
||||
def __init__(self, input_channels=3, class_nums=1, output_onnx=False):
|
||||
super(BodySegNet_32_thin_sigmod, self).__init__()
|
||||
self.class_nums = class_nums
|
||||
self.output_onnx = output_onnx
|
||||
|
||||
self.stage_1 = nn.Sequential(
|
||||
ConvBNReLU(input_channels, 16, kernel_size=3, stride=2)
|
||||
)
|
||||
self.stage_2 = nn.Sequential(
|
||||
ConvBNReLU(16, 16, kernel_size=3, stride=2, groups=16),
|
||||
ConvBNReLU(16, 16, kernel_size=1, stride=1),
|
||||
)
|
||||
self.stage_3 = nn.Sequential(
|
||||
InvertedResidual(16, 24, stride=2, expand_ratio=6),
|
||||
InvertedResidual(24, 24, stride=1, expand_ratio=6),
|
||||
InvertedResidual(24, 24, stride=1, expand_ratio=6),
|
||||
)
|
||||
self.stage_4 = nn.Sequential(
|
||||
InvertedResidual(24, 32, stride=2, expand_ratio=6),
|
||||
InvertedResidual(32, 32, stride=1, expand_ratio=6),
|
||||
InvertedResidual(32, 32, stride=1, expand_ratio=6),
|
||||
InvertedResidual(32, 32, stride=1, expand_ratio=6),
|
||||
)
|
||||
self.stage_5 = nn.Sequential(
|
||||
InvertedResidual(32, 48, stride=2, expand_ratio=6),
|
||||
InvertedResidual(48, 48, stride=1, expand_ratio=6),
|
||||
InvertedResidual(48, 48, stride=1, expand_ratio=6),
|
||||
InvertedResidual(48, 48, stride=1, expand_ratio=6)
|
||||
)
|
||||
self.up_to_4 = UpSampleBlock(48 + 32, 16)
|
||||
self.up_to_3 = UpSampleBlock(16 + 24, 16)
|
||||
self.up_to_2 = UpSampleBlock(16 + 16, 16)
|
||||
self.up_to_1 = UpSampleBlock(16 + 16, 16)
|
||||
self.last_layer = nn.Sequential(
|
||||
ConvBNReLU(16, 16, kernel_size=1, stride=1),
|
||||
nn.Conv2d(16, self.class_nums, kernel_size=1, stride=1)
|
||||
)
|
||||
self._initialize_weights()
|
||||
|
||||
def forward(self, x):
|
||||
feature_S = []
|
||||
x1 = self.stage_1(x)
|
||||
x2 = self.stage_2(x1)
|
||||
x3 = self.stage_3(x2)
|
||||
x4 = self.stage_4(x3)
|
||||
feature = self.stage_5(x4)
|
||||
|
||||
feature = F.interpolate(feature, size=x4.size()[2:], mode='bilinear', align_corners=True)
|
||||
feature = self.up_to_4(x4, feature)
|
||||
feature = F.interpolate(feature, size=x3.size()[2:], mode='bilinear', align_corners=True)
|
||||
feature = self.up_to_3(x3, feature)
|
||||
feature = F.interpolate(feature, size=x2.size()[2:], mode='bilinear', align_corners=True)
|
||||
feature = self.up_to_2(x2, feature)
|
||||
feature = F.interpolate(feature, size=x1.size()[2:], mode='bilinear', align_corners=True)
|
||||
feature = self.up_to_1(x1, feature)
|
||||
feature_S.append(feature)
|
||||
feature = self.last_layer(feature)
|
||||
feature_S.append(feature)
|
||||
output = F.interpolate(feature, size=x.size()[2:], mode='bilinear', align_corners=True)
|
||||
# output = torch.sigmoid(output)
|
||||
if self.output_onnx:
|
||||
output = torch.argmax(output, dim=1).to(torch.float32)
|
||||
|
||||
return output, feature_S
|
||||
|
||||
def _initialize_weights(self):
|
||||
for name, m in self.named_modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
if 'first' in name:
|
||||
nn.init.normal_(m.weight, 0, 0.01)
|
||||
else:
|
||||
nn.init.normal_(m.weight, 0, 1.0 / m.weight.shape[1])
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0.0001)
|
||||
nn.init.constant_(m.running_mean, 0)
|
||||
elif isinstance(m, nn.BatchNorm1d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0.0001)
|
||||
nn.init.constant_(m.running_mean, 0)
|
||||
elif isinstance(m, nn.Linear):
|
||||
nn.init.normal_(m.weight, 0, 0.01)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
class _ASPPModule(nn.Module):
|
||||
def __init__(self, inplanes, planes, kernel_size, padding, dilation, BatchNorm):
|
||||
super(_ASPPModule, self).__init__()
|
||||
self.atrous_conv = nn.Conv2d(inplanes, planes, kernel_size=kernel_size,
|
||||
stride=1, padding=padding, dilation=dilation, bias=False)
|
||||
self.bn = BatchNorm(planes)
|
||||
self.relu = nn.ReLU()
|
||||
|
||||
self._init_weight()
|
||||
|
||||
def forward(self, x):
|
||||
x = self.atrous_conv(x)
|
||||
x = self.bn(x)
|
||||
|
||||
return self.relu(x)
|
||||
|
||||
def _init_weight(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
torch.nn.init.kaiming_normal_(m.weight)
|
||||
|
||||
class ASPP(nn.Module):
|
||||
def __init__(self, backbone, output_stride, BatchNorm):
|
||||
super(ASPP, self).__init__()
|
||||
if backbone == 'drn':
|
||||
inplanes = 512
|
||||
elif backbone == 'mobilenet':
|
||||
inplanes = 320
|
||||
elif backbone == 'resnet_fpn':
|
||||
inplanes = 256
|
||||
else:
|
||||
inplanes = 2048
|
||||
if output_stride == 16:
|
||||
dilations = [1, 6, 12, 18]
|
||||
elif output_stride == 8:
|
||||
dilations = [1, 12, 24, 36]
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
self.aspp1 = _ASPPModule(inplanes, 256, 1, padding=0, dilation=dilations[0], BatchNorm=BatchNorm)
|
||||
self.aspp2 = _ASPPModule(inplanes, 256, 3, padding=dilations[1], dilation=dilations[1], BatchNorm=BatchNorm)
|
||||
self.aspp3 = _ASPPModule(inplanes, 256, 3, padding=dilations[2], dilation=dilations[2], BatchNorm=BatchNorm)
|
||||
self.aspp4 = _ASPPModule(inplanes, 256, 3, padding=dilations[3], dilation=dilations[3], BatchNorm=BatchNorm)
|
||||
|
||||
self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
|
||||
nn.Conv2d(inplanes, 256, 1, stride=1, bias=False),
|
||||
BatchNorm(256),
|
||||
nn.ReLU())
|
||||
self.conv1 = nn.Conv2d(1280, 256, 1, bias=False)
|
||||
self.bn1 = BatchNorm(256)
|
||||
self.relu = nn.ReLU()
|
||||
self.dropout = nn.Dropout(0.5)
|
||||
self._init_weight()
|
||||
|
||||
def forward(self, x):
|
||||
x1 = self.aspp1(x)
|
||||
x2 = self.aspp2(x)
|
||||
x3 = self.aspp3(x)
|
||||
x4 = self.aspp4(x)
|
||||
x5 = self.global_avg_pool(x)
|
||||
x5 = F.interpolate(x5, size=x4.size()[2:], mode='bilinear', align_corners=True)
|
||||
# x5 = F.interpolate(x5, size=x4.size()[2:], mode='nearest')
|
||||
x = torch.cat((x1, x2, x3, x4, x5), dim=1)
|
||||
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
|
||||
return self.dropout(x)
|
||||
|
||||
def _init_weight(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
torch.nn.init.kaiming_normal_(m.weight)
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(self,num_classes, backbone, BatchNorm):
|
||||
super(Decoder, self).__init__()
|
||||
if backbone == 'resnet_fpn' or backbone == 'drn':
|
||||
low_level_inplanes = 256
|
||||
elif backbone == 'xception':
|
||||
low_level_inplanes = 128
|
||||
elif backbone == 'mobilenet':
|
||||
low_level_inplanes = 24
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
self.conv1 = nn.Conv2d(low_level_inplanes, 16, 1, bias=False)
|
||||
self.bn1 = BatchNorm(16)
|
||||
self.relu = nn.ReLU()
|
||||
self.last_conv = nn.Sequential(nn.Conv2d(304, 256, kernel_size=3, stride=1, padding=1, bias=False),
|
||||
BatchNorm(256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.5),
|
||||
nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=False),
|
||||
BatchNorm(256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.1),
|
||||
nn.Conv2d(256, num_classes, kernel_size=1, stride=1))
|
||||
|
||||
self.up_to_1 = UpSampleBlock(16 + 16, 16)
|
||||
self.last_layer = nn.Sequential(
|
||||
ConvBNReLU(16, 16, kernel_size=1, stride=1),
|
||||
nn.Conv2d(16, 1, kernel_size=1, stride=1)
|
||||
)
|
||||
self._init_weight()
|
||||
|
||||
# x(1,256,8,6) low(1,256,32,24) -》 x(1,1,32,24)
|
||||
def forward(self, x, low_level_feat):
|
||||
feature_T = []
|
||||
# deeplab part
|
||||
#(1,256,32,24) -> (1,16,64,48)
|
||||
low_level_feat = self.conv1(low_level_feat)
|
||||
low_level_feat = self.bn1(low_level_feat)
|
||||
low_level_feat = self.relu(low_level_feat)
|
||||
|
||||
# x(1, 256, 8, 6)-> (1,16,32,24)
|
||||
x = F.interpolate(x, size=low_level_feat.size()[2:], mode='bilinear', align_corners=True)
|
||||
low_level_feat = F.interpolate(low_level_feat,
|
||||
size=[low_level_feat.size()[2] * 2, low_level_feat.size()[3] * 2],
|
||||
mode='bilinear', align_corners=True)
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
|
||||
# bodyseg part
|
||||
feature = F.interpolate(x, size=[x.size()[2] * 2,x.size()[3] * 2], mode='bilinear', align_corners=True)
|
||||
# 输入up_to_1 x(low)(1,16,64,48),上采样2倍后的feature(1,16,64,48)
|
||||
feature = self.up_to_1(low_level_feat, feature)
|
||||
feature_T.append(feature)
|
||||
# (1,1,64,48)
|
||||
feature = self.last_layer(feature)
|
||||
feature_T.append(feature)
|
||||
# (1,1,128,96)
|
||||
output = F.interpolate(feature, size=[128,96], mode='bilinear', align_corners=True)
|
||||
# output = torch.sigmoid(output)
|
||||
|
||||
return output, feature_T
|
||||
|
||||
def _init_weight(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
torch.nn.init.kaiming_normal_(m.weight)
|
||||
|
||||
def build_backbone(backbone, output_stride, BatchNorm, input_channel=3):
|
||||
from core.bodyseg.backbone.fpn import build_resnet_fpn_backbone
|
||||
from core.bodyseg.backbone.xception import AlignedXception
|
||||
|
||||
return build_resnet_fpn_backbone(input_channel)
|
||||
|
||||
def build_aspp(backbone, output_stride, BatchNorm):
|
||||
return ASPP(backbone, output_stride, BatchNorm)
|
||||
|
||||
def build_decoder(num_classes, backbone, BatchNorm):
|
||||
return Decoder(num_classes, backbone, BatchNorm)
|
||||
|
||||
class DeepLab(nn.Module):
|
||||
def __init__(self, input_channel=3, class_num=1):
|
||||
super(DeepLab, self).__init__()
|
||||
|
||||
BatchNorm = get_norm("BN")
|
||||
|
||||
self.backbone = build_backbone("resnet_fpn", 16, BatchNorm, input_channel=input_channel)
|
||||
self.aspp = build_aspp("resnet_fpn", 16, BatchNorm)
|
||||
self.decoder = build_decoder(class_num, "resnet_fpn", BatchNorm)
|
||||
|
||||
def forward(self, input):
|
||||
# input(1,3,128,96)
|
||||
#output: "p2"(1,256,32,24), "p3"(1,256,16,12), "p4"(1,256,8,6)
|
||||
output = self.backbone(input)
|
||||
# "p4"(1,256,8,6) "p2"(1,256,32,24)
|
||||
x, low_level_feat = output['p4'], output['p2']
|
||||
# x(1,256,8,6)-》(1,256,8,6)
|
||||
x = self.aspp(x)
|
||||
# x(1,256,8,6) low(1,256,32,24) -》 x(1,1,32,24)
|
||||
x, feature_T = self.decoder(x, low_level_feat)
|
||||
# x(1,1,128,96)
|
||||
x = F.interpolate(x, size=input.size()[2:], mode='bilinear', align_corners=True)
|
||||
# x = F.interpolate(x, size=input.size()[2:], mode='nearest')
|
||||
return x, feature_T
|
||||
@@ -0,0 +1,57 @@
|
||||
from qcloud_cos import CosConfig
|
||||
from qcloud_cos import CosS3Client
|
||||
import sys
|
||||
import os.path as osp
|
||||
import time
|
||||
import logging
|
||||
from common.logger import config as confccc
|
||||
|
||||
|
||||
class COS_object():
|
||||
def __init__(self):
|
||||
# 正常情况日志级别使用 INFO，需要定位时可以修改为 DEBUG，此时 SDK 会打印和服务端的通信信息
|
||||
logging.basicConfig(level=logging.INFO, stream=sys.stdout)
|
||||
|
||||
# 1. 设置用户属性, 包括 secret_id, secret_key, region等。Appid 已在 CosConfig 中移除，请在参数 Bucket 中带上 Appid。Bucket 由 BucketName-Appid 组成
|
||||
secret_id = 'AKIDAIZ3MTLyNXRaWLvGAsLaZhgbzQjqkyGN' # 用户的 SecretId，建议使用子账号密钥，授权遵循最小权限指引，降低使用风险。子账号密钥获取可参见 https://cloud.tencent.com/document/product/598/37140
|
||||
secret_key = '73kEZgjRxIsGQTM4oaFjCXGp9IH7xgIm' # 用户的 SecretKey，建议使用子账号密钥，授权遵循最小权限指引，降低使用风险。子账号密钥获取可参见 https://cloud.tencent.com/document/product/598/37140
|
||||
self.region = 'ap-beijing' # 替换为用户的 region，已创建桶归属的 region 可以在控制台查看，https://console.cloud.tencent.com/cos5/bucket
|
||||
# COS 支持的所有 region 列表参见https://cloud.tencent.com/document/product/436/6224
|
||||
token = None # 如果使用永久密钥不需要填入 token，如果使用临时密钥需要填入，临时密钥生成和使用指引参见 https://cloud.tencent.com/document/product/436/14048
|
||||
self.scheme = 'https' # 指定使用 http/https 协议来访问 COS，默认为 https，可不填
|
||||
self.BucketName= 'ydapp-1317132355'
|
||||
config = CosConfig(Region=self.region, SecretId=secret_id, SecretKey=secret_key, Token=token, Scheme=self.scheme)
|
||||
self.client = CosS3Client(config)
|
||||
|
||||
def upload_file(self, file, target_name):
|
||||
t0 = time.time()
|
||||
with open(file, 'rb') as fp:
|
||||
response = self.client.put_object(
|
||||
Bucket=self.BucketName, # Bucket 由 BucketName-APPID 组成
|
||||
Body=fp,
|
||||
Key=target_name,
|
||||
StorageClass='STANDARD',
|
||||
ContentType='image/jpeg;image/jpg;image/png;image/gif'
|
||||
)
|
||||
ret_url = self.scheme + '://' + self.BucketName + '.cos.' + self.region + '.myqcloud.com/' + target_name
|
||||
print('time costs: {}'.format(time.time() - t0), ret_url)
|
||||
return ret_url
|
||||
|
||||
def download_img(self, img_url, tmp_dir):
|
||||
# user_img_save_dir = confccc.get('default', 'userDir')
|
||||
# user_img_tmp_dir = confccc.get('default', 'tmp_dir')
|
||||
fileName = img_url.split('.com/')[-1]
|
||||
print(fileName)
|
||||
response = self.client.get_object(
|
||||
Bucket=self.BucketName,
|
||||
Key=fileName,
|
||||
)
|
||||
response['Body'].get_stream_to_file(tmp_dir)
|
||||
|
||||
|
||||
if __name__=='__main__':
|
||||
t0 = time.time()
|
||||
cos = COS_object()
|
||||
cos.upload_file('/home/chinatszrn/Pictures/vaffflue12.png', 'hair_mz/images/vaffflue12.png')
|
||||
# cos.download_img('https://ydapp-1317132355.cos.ap-beijing.myqcloud.com/hair_mz/images/hairstyle/fb7d7a58-3228-40f2-84a6-337088ee31e2/2023031623425988.jpg', 'a0')
|
||||
print(time.time() - t0)
|
||||
@@ -0,0 +1,2 @@
|
||||
from . import mesh
|
||||
from . import morphable_model
|
||||
@@ -0,0 +1,7 @@
|
||||
# from .cython import mesh_core_cython
|
||||
from . import io
|
||||
from . import vis
|
||||
from . import transform
|
||||
from . import light
|
||||
from . import render
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
*.cpython-36m-x86_64-linux-gnu.so
|
||||
@@ -0,0 +1,375 @@
|
||||
/*
|
||||
functions that can not be optimazed by vertorization in python.
|
||||
1. rasterization.(need process each triangle)
|
||||
2. normal of each vertex.(use one-ring, need process each vertex)
|
||||
3. write obj(seems that it can be verctorized? anyway, writing it in c++ is simple, so also add function here. --> however, why writting in c++ is still slow?)
|
||||
|
||||
Author: Yao Feng
|
||||
Mail: yaofeng1995@gmail.com
|
||||
*/
|
||||
|
||||
#include "mesh_core.h"
|
||||
|
||||
|
||||
/* Judge whether the point is in the triangle
|
||||
Method:
|
||||
http://blackpawn.com/texts/pointinpoly/
|
||||
Args:
|
||||
point: [x, y]
|
||||
tri_points: three vertices(2d points) of a triangle. 2 coords x 3 vertices
|
||||
Returns:
|
||||
bool: true for in triangle
|
||||
*/
|
||||
bool isPointInTri(point p, point p0, point p1, point p2)
|
||||
{
|
||||
// vectors
|
||||
point v0, v1, v2;
|
||||
v0 = p2 - p0;
|
||||
v1 = p1 - p0;
|
||||
v2 = p - p0;
|
||||
|
||||
// dot products
|
||||
float dot00 = v0.dot(v0); //v0.x * v0.x + v0.y * v0.y //np.dot(v0.T, v0)
|
||||
float dot01 = v0.dot(v1); //v0.x * v1.x + v0.y * v1.y //np.dot(v0.T, v1)
|
||||
float dot02 = v0.dot(v2); //v0.x * v2.x + v0.y * v2.y //np.dot(v0.T, v2)
|
||||
float dot11 = v1.dot(v1); //v1.x * v1.x + v1.y * v1.y //np.dot(v1.T, v1)
|
||||
float dot12 = v1.dot(v2); //v1.x * v2.x + v1.y * v2.y//np.dot(v1.T, v2)
|
||||
|
||||
// barycentric coordinates
|
||||
float inverDeno;
|
||||
if(dot00*dot11 - dot01*dot01 == 0)
|
||||
inverDeno = 0;
|
||||
else
|
||||
inverDeno = 1/(dot00*dot11 - dot01*dot01);
|
||||
|
||||
float u = (dot11*dot02 - dot01*dot12)*inverDeno;
|
||||
float v = (dot00*dot12 - dot01*dot02)*inverDeno;
|
||||
|
||||
// check if point in triangle
|
||||
return (u >= 0) && (v >= 0) && (u + v < 1);
|
||||
}
|
||||
|
||||
|
||||
void get_point_weight(float* weight, point p, point p0, point p1, point p2)
|
||||
{
|
||||
// vectors
|
||||
point v0, v1, v2;
|
||||
v0 = p2 - p0;
|
||||
v1 = p1 - p0;
|
||||
v2 = p - p0;
|
||||
|
||||
// dot products
|
||||
float dot00 = v0.dot(v0); //v0.x * v0.x + v0.y * v0.y //np.dot(v0.T, v0)
|
||||
float dot01 = v0.dot(v1); //v0.x * v1.x + v0.y * v1.y //np.dot(v0.T, v1)
|
||||
float dot02 = v0.dot(v2); //v0.x * v2.x + v0.y * v2.y //np.dot(v0.T, v2)
|
||||
float dot11 = v1.dot(v1); //v1.x * v1.x + v1.y * v1.y //np.dot(v1.T, v1)
|
||||
float dot12 = v1.dot(v2); //v1.x * v2.x + v1.y * v2.y//np.dot(v1.T, v2)
|
||||
|
||||
// barycentric coordinates
|
||||
float inverDeno;
|
||||
if(dot00*dot11 - dot01*dot01 == 0)
|
||||
inverDeno = 0;
|
||||
else
|
||||
inverDeno = 1/(dot00*dot11 - dot01*dot01);
|
||||
|
||||
float u = (dot11*dot02 - dot01*dot12)*inverDeno;
|
||||
float v = (dot00*dot12 - dot01*dot02)*inverDeno;
|
||||
|
||||
// weight
|
||||
weight[0] = 1 - u - v;
|
||||
weight[1] = v;
|
||||
weight[2] = u;
|
||||
}
|
||||
|
||||
|
||||
void _get_normal_core(
|
||||
float* normal, float* tri_normal, int* triangles,
|
||||
int ntri)
|
||||
{
|
||||
int i, j;
|
||||
int tri_p0_ind, tri_p1_ind, tri_p2_ind;
|
||||
|
||||
for(i = 0; i < ntri; i++)
|
||||
{
|
||||
tri_p0_ind = triangles[3*i];
|
||||
tri_p1_ind = triangles[3*i + 1];
|
||||
tri_p2_ind = triangles[3*i + 2];
|
||||
|
||||
for(j = 0; j < 3; j++)
|
||||
{
|
||||
normal[3*tri_p0_ind + j] = normal[3*tri_p0_ind + j] + tri_normal[3*i + j];
|
||||
normal[3*tri_p1_ind + j] = normal[3*tri_p1_ind + j] + tri_normal[3*i + j];
|
||||
normal[3*tri_p2_ind + j] = normal[3*tri_p2_ind + j] + tri_normal[3*i + j];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
void _rasterize_triangles_core(
|
||||
float* vertices, int* triangles,
|
||||
float* depth_buffer, int* triangle_buffer, float* barycentric_weight,
|
||||
int nver, int ntri,
|
||||
int h, int w)
|
||||
{
|
||||
int i;
|
||||
int x, y, k;
|
||||
int tri_p0_ind, tri_p1_ind, tri_p2_ind;
|
||||
point p0, p1, p2, p;
|
||||
int x_min, x_max, y_min, y_max;
|
||||
float p_depth, p0_depth, p1_depth, p2_depth;
|
||||
float weight[3];
|
||||
|
||||
for(i = 0; i < ntri; i++)
|
||||
{
|
||||
tri_p0_ind = triangles[3*i];
|
||||
tri_p1_ind = triangles[3*i + 1];
|
||||
tri_p2_ind = triangles[3*i + 2];
|
||||
|
||||
p0.x = vertices[3*tri_p0_ind]; p0.y = vertices[3*tri_p0_ind + 1]; p0_depth = vertices[3*tri_p0_ind + 2];
|
||||
p1.x = vertices[3*tri_p1_ind]; p1.y = vertices[3*tri_p1_ind + 1]; p1_depth = vertices[3*tri_p1_ind + 2];
|
||||
p2.x = vertices[3*tri_p2_ind]; p2.y = vertices[3*tri_p2_ind + 1]; p2_depth = vertices[3*tri_p2_ind + 2];
|
||||
|
||||
x_min = max((int)ceil(min(p0.x, min(p1.x, p2.x))), 0);
|
||||
x_max = min((int)floor(max(p0.x, max(p1.x, p2.x))), w - 1);
|
||||
|
||||
y_min = max((int)ceil(min(p0.y, min(p1.y, p2.y))), 0);
|
||||
y_max = min((int)floor(max(p0.y, max(p1.y, p2.y))), h - 1);
|
||||
|
||||
if(x_max < x_min || y_max < y_min)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
for(y = y_min; y <= y_max; y++) //h
|
||||
{
|
||||
for(x = x_min; x <= x_max; x++) //w
|
||||
{
|
||||
p.x = x; p.y = y;
|
||||
if(p.x < 2 || p.x > w - 3 || p.y < 2 || p.y > h - 3 || isPointInTri(p, p0, p1, p2))
|
||||
{
|
||||
get_point_weight(weight, p, p0, p1, p2);
|
||||
p_depth = weight[0]*p0_depth + weight[1]*p1_depth + weight[2]*p2_depth;
|
||||
|
||||
if((p_depth > depth_buffer[y*w + x]))
|
||||
{
|
||||
depth_buffer[y*w + x] = p_depth;
|
||||
triangle_buffer[y*w + x] = i;
|
||||
for(k = 0; k < 3; k++)
|
||||
{
|
||||
barycentric_weight[y*w*3 + x*3 + k] = weight[k];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
void _render_colors_core(
|
||||
float* image, float* vertices, int* triangles,
|
||||
float* colors,
|
||||
float* depth_buffer,
|
||||
int nver, int ntri,
|
||||
int h, int w, int c)
|
||||
{
|
||||
int i;
|
||||
int x, y, k;
|
||||
int tri_p0_ind, tri_p1_ind, tri_p2_ind;
|
||||
point p0, p1, p2, p;
|
||||
int x_min, x_max, y_min, y_max;
|
||||
float p_depth, p0_depth, p1_depth, p2_depth;
|
||||
float p_color, p0_color, p1_color, p2_color;
|
||||
float weight[3];
|
||||
|
||||
for(i = 0; i < ntri; i++)
|
||||
{
|
||||
tri_p0_ind = triangles[3*i];
|
||||
tri_p1_ind = triangles[3*i + 1];
|
||||
tri_p2_ind = triangles[3*i + 2];
|
||||
|
||||
p0.x = vertices[3*tri_p0_ind]; p0.y = vertices[3*tri_p0_ind + 1]; p0_depth = vertices[3*tri_p0_ind + 2];
|
||||
p1.x = vertices[3*tri_p1_ind]; p1.y = vertices[3*tri_p1_ind + 1]; p1_depth = vertices[3*tri_p1_ind + 2];
|
||||
p2.x = vertices[3*tri_p2_ind]; p2.y = vertices[3*tri_p2_ind + 1]; p2_depth = vertices[3*tri_p2_ind + 2];
|
||||
|
||||
x_min = max((int)ceil(min(p0.x, min(p1.x, p2.x))), 0);
|
||||
x_max = min((int)floor(max(p0.x, max(p1.x, p2.x))), w - 1);
|
||||
|
||||
y_min = max((int)ceil(min(p0.y, min(p1.y, p2.y))), 0);
|
||||
y_max = min((int)floor(max(p0.y, max(p1.y, p2.y))), h - 1);
|
||||
|
||||
if(x_max < x_min || y_max < y_min)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
for(y = y_min; y <= y_max; y++) //h
|
||||
{
|
||||
for(x = x_min; x <= x_max; x++) //w
|
||||
{
|
||||
p.x = x; p.y = y;
|
||||
if(p.x < 2 || p.x > w - 3 || p.y < 2 || p.y > h - 3 || isPointInTri(p, p0, p1, p2))
|
||||
{
|
||||
get_point_weight(weight, p, p0, p1, p2);
|
||||
p_depth = weight[0]*p0_depth + weight[1]*p1_depth + weight[2]*p2_depth;
|
||||
|
||||
if((p_depth > depth_buffer[y*w + x]))
|
||||
{
|
||||
for(k = 0; k < c; k++) // c
|
||||
{
|
||||
p0_color = colors[c*tri_p0_ind + k];
|
||||
p1_color = colors[c*tri_p1_ind + k];
|
||||
p2_color = colors[c*tri_p2_ind + k];
|
||||
|
||||
p_color = weight[0]*p0_color + weight[1]*p1_color + weight[2]*p2_color;
|
||||
image[y*w*c + x*c + k] = p_color;
|
||||
}
|
||||
|
||||
depth_buffer[y*w + x] = p_depth;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
void _render_texture_core(
|
||||
float* image, float* vertices, int* triangles,
|
||||
float* texture, float* tex_coords, int* tex_triangles,
|
||||
float* depth_buffer,
|
||||
int nver, int tex_nver, int ntri,
|
||||
int h, int w, int c,
|
||||
int tex_h, int tex_w, int tex_c,
|
||||
int mapping_type)
|
||||
{
|
||||
int i;
|
||||
int x, y, k;
|
||||
int tri_p0_ind, tri_p1_ind, tri_p2_ind;
|
||||
int tex_tri_p0_ind, tex_tri_p1_ind, tex_tri_p2_ind;
|
||||
point p0, p1, p2, p;
|
||||
point tex_p0, tex_p1, tex_p2, tex_p;
|
||||
int x_min, x_max, y_min, y_max;
|
||||
float weight[3];
|
||||
float p_depth, p0_depth, p1_depth, p2_depth;
|
||||
float xd, yd;
|
||||
float ul, ur, dl, dr;
|
||||
for(i = 0; i < ntri; i++)
|
||||
{
|
||||
// mesh
|
||||
tri_p0_ind = triangles[3*i];
|
||||
tri_p1_ind = triangles[3*i + 1];
|
||||
tri_p2_ind = triangles[3*i + 2];
|
||||
|
||||
p0.x = vertices[3*tri_p0_ind]; p0.y = vertices[3*tri_p0_ind + 1]; p0_depth = vertices[3*tri_p0_ind + 2];
|
||||
p1.x = vertices[3*tri_p1_ind]; p1.y = vertices[3*tri_p1_ind + 1]; p1_depth = vertices[3*tri_p1_ind + 2];
|
||||
p2.x = vertices[3*tri_p2_ind]; p2.y = vertices[3*tri_p2_ind + 1]; p2_depth = vertices[3*tri_p2_ind + 2];
|
||||
|
||||
// texture
|
||||
tex_tri_p0_ind = tex_triangles[3*i];
|
||||
tex_tri_p1_ind = tex_triangles[3*i + 1];
|
||||
tex_tri_p2_ind = tex_triangles[3*i + 2];
|
||||
|
||||
tex_p0.x = tex_coords[3*tex_tri_p0_ind]; tex_p0.y = tex_coords[3*tri_p0_ind + 1];
|
||||
tex_p1.x = tex_coords[3*tex_tri_p1_ind]; tex_p1.y = tex_coords[3*tri_p1_ind + 1];
|
||||
tex_p2.x = tex_coords[3*tex_tri_p2_ind]; tex_p2.y = tex_coords[3*tri_p2_ind + 1];
|
||||
|
||||
|
||||
x_min = max((int)ceil(min(p0.x, min(p1.x, p2.x))), 0);
|
||||
x_max = min((int)floor(max(p0.x, max(p1.x, p2.x))), w - 1);
|
||||
|
||||
y_min = max((int)ceil(min(p0.y, min(p1.y, p2.y))), 0);
|
||||
y_max = min((int)floor(max(p0.y, max(p1.y, p2.y))), h - 1);
|
||||
|
||||
|
||||
if(x_max < x_min || y_max < y_min)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
for(y = y_min; y <= y_max; y++) //h
|
||||
{
|
||||
for(x = x_min; x <= x_max; x++) //w
|
||||
{
|
||||
p.x = x; p.y = y;
|
||||
if(p.x < 2 || p.x > w - 3 || p.y < 2 || p.y > h - 3 || isPointInTri(p, p0, p1, p2))
|
||||
{
|
||||
get_point_weight(weight, p, p0, p1, p2);
|
||||
p_depth = weight[0]*p0_depth + weight[1]*p1_depth + weight[2]*p2_depth;
|
||||
|
||||
if((p_depth > depth_buffer[y*w + x]))
|
||||
{
|
||||
// -- color from texture
|
||||
// cal weight in mesh tri
|
||||
get_point_weight(weight, p, p0, p1, p2);
|
||||
// cal coord in texture
|
||||
tex_p = tex_p0*weight[0] + tex_p1*weight[1] + tex_p2*weight[2];
|
||||
tex_p.x = max(min(tex_p.x, float(tex_w - 1)), float(0));
|
||||
tex_p.y = max(min(tex_p.y, float(tex_h - 1)), float(0));
|
||||
|
||||
yd = tex_p.y - floor(tex_p.y);
|
||||
xd = tex_p.x - floor(tex_p.x);
|
||||
for(k = 0; k < c; k++)
|
||||
{
|
||||
if(mapping_type==0)// nearest
|
||||
{
|
||||
image[y*w*c + x*c + k] = texture[int(round(tex_p.y))*tex_w*tex_c + int(round(tex_p.x))*tex_c + k];
|
||||
}
|
||||
else//bilinear interp
|
||||
{
|
||||
ul = texture[(int)floor(tex_p.y)*tex_w*tex_c + (int)floor(tex_p.x)*tex_c + k];
|
||||
ur = texture[(int)floor(tex_p.y)*tex_w*tex_c + (int)ceil(tex_p.x)*tex_c + k];
|
||||
dl = texture[(int)ceil(tex_p.y)*tex_w*tex_c + (int)floor(tex_p.x)*tex_c + k];
|
||||
dr = texture[(int)ceil(tex_p.y)*tex_w*tex_c + (int)ceil(tex_p.x)*tex_c + k];
|
||||
|
||||
image[y*w*c + x*c + k] = ul*(1-xd)*(1-yd) + ur*xd*(1-yd) + dl*(1-xd)*yd + dr*xd*yd;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
depth_buffer[y*w + x] = p_depth;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
// ------------------------------------------------- write
|
||||
// obj write
|
||||
// Ref: https://github.com/patrikhuber/eos/blob/master/include/eos/core/Mesh.hpp
|
||||
void _write_obj_with_colors_texture(string filename, string mtl_name,
|
||||
float* vertices, int* triangles, float* colors, float* uv_coords,
|
||||
int nver, int ntri, int ntexver)
|
||||
{
|
||||
int i;
|
||||
|
||||
ofstream obj_file(filename.c_str());
|
||||
|
||||
// first line of the obj file: the mtl name
|
||||
obj_file << "mtllib " << mtl_name << endl;
|
||||
|
||||
// write vertices
|
||||
for (i = 0; i < nver; ++i)
|
||||
{
|
||||
obj_file << "v " << vertices[3*i] << " " << vertices[3*i + 1] << " " << vertices[3*i + 2] << colors[3*i] << " " << colors[3*i + 1] << " " << colors[3*i + 2] << endl;
|
||||
}
|
||||
|
||||
// write uv coordinates
|
||||
for (i = 0; i < ntexver; ++i)
|
||||
{
|
||||
//obj_file << "vt " << uv_coords[2*i] << " " << (1 - uv_coords[2*i + 1]) << endl;
|
||||
obj_file << "vt " << uv_coords[2*i] << " " << uv_coords[2*i + 1] << endl;
|
||||
}
|
||||
|
||||
obj_file << "usemtl FaceTexture" << endl;
|
||||
// write triangles
|
||||
for (i = 0; i < ntri; ++i)
|
||||
{
|
||||
// obj_file << "f " << triangles[3*i] << "/" << triangles[3*i] << " " << triangles[3*i + 1] << "/" << triangles[3*i + 1] << " " << triangles[3*i + 2] << "/" << triangles[3*i + 2] << endl;
|
||||
obj_file << "f " << triangles[3*i + 2] << "/" << triangles[3*i + 2] << " " << triangles[3*i + 1] << "/" << triangles[3*i + 1] << " " << triangles[3*i] << "/" << triangles[3*i] << endl;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
#ifndef MESH_CORE_HPP_
|
||||
#define MESH_CORE_HPP_
|
||||
|
||||
#include <stdio.h>
|
||||
#include <cmath>
|
||||
#include <algorithm>
|
||||
#include <string>
|
||||
#include <iostream>
|
||||
#include <fstream>
|
||||
|
||||
using namespace std;
|
||||
|
||||
class point
|
||||
{
|
||||
public:
|
||||
float x;
|
||||
float y;
|
||||
|
||||
float dot(point p)
|
||||
{
|
||||
return this->x * p.x + this->y * p.y;
|
||||
}
|
||||
|
||||
point operator-(const point& p)
|
||||
{
|
||||
point np;
|
||||
np.x = this->x - p.x;
|
||||
np.y = this->y - p.y;
|
||||
return np;
|
||||
}
|
||||
|
||||
point operator+(const point& p)
|
||||
{
|
||||
point np;
|
||||
np.x = this->x + p.x;
|
||||
np.y = this->y + p.y;
|
||||
return np;
|
||||
}
|
||||
|
||||
point operator*(float s)
|
||||
{
|
||||
point np;
|
||||
np.x = s * this->x;
|
||||
np.y = s * this->y;
|
||||
return np;
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
bool isPointInTri(point p, point p0, point p1, point p2, int h, int w);
|
||||
void get_point_weight(float* weight, point p, point p0, point p1, point p2);
|
||||
|
||||
void _get_normal_core(
|
||||
float* normal, float* tri_normal, int* triangles,
|
||||
int ntri);
|
||||
|
||||
void _rasterize_triangles_core(
|
||||
float* vertices, int* triangles,
|
||||
float* depth_buffer, int* triangle_buffer, float* barycentric_weight,
|
||||
int nver, int ntri,
|
||||
int h, int w);
|
||||
|
||||
void _render_colors_core(
|
||||
float* image, float* vertices, int* triangles,
|
||||
float* colors,
|
||||
float* depth_buffer,
|
||||
int nver, int ntri,
|
||||
int h, int w, int c);
|
||||
|
||||
void _render_texture_core(
|
||||
float* image, float* vertices, int* triangles,
|
||||
float* texture, float* tex_coords, int* tex_triangles,
|
||||
float* depth_buffer,
|
||||
int nver, int tex_nver, int ntri,
|
||||
int h, int w, int c,
|
||||
int tex_h, int tex_w, int tex_c,
|
||||
int mapping_type);
|
||||
|
||||
void _write_obj_with_colors_texture(string filename, string mtl_name,
|
||||
float* vertices, int* triangles, float* colors, float* uv_coords,
|
||||
int nver, int ntri, int ntexver);
|
||||
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,22 @@
|
||||
'''
|
||||
python setup.py build_ext -i
|
||||
to compile
|
||||
'''
|
||||
|
||||
# setup.py
|
||||
from distutils.core import setup
|
||||
from setuptools import Extension
|
||||
|
||||
from Cython.Build import cythonize
|
||||
from Cython.Distutils import build_ext
|
||||
import numpy
|
||||
|
||||
setup(
|
||||
name = 'mesh_core_cython',
|
||||
cmdclass={'build_ext': build_ext},
|
||||
ext_modules=[Extension("mesh_core_cython",
|
||||
sources=["mesh_core_cython.pyx", "mesh_core.cpp"],
|
||||
language='c++',
|
||||
include_dirs=[numpy.get_include()])],
|
||||
)
|
||||
|
||||
@@ -0,0 +1,142 @@
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import numpy as np
|
||||
import os
|
||||
# from skimage import io
|
||||
from time import time
|
||||
|
||||
# from .cython import mesh_core_cython
|
||||
|
||||
## TODO
|
||||
## TODO: c++ version
|
||||
def read_obj(obj_name):
|
||||
''' read mesh
|
||||
'''
|
||||
return 0
|
||||
|
||||
# ------------------------- write
|
||||
def write_asc(path, vertices):
|
||||
'''
|
||||
Args:
|
||||
vertices: shape = (nver, 3)
|
||||
'''
|
||||
if path.split('.')[-1] == 'asc':
|
||||
np.savetxt(path, vertices)
|
||||
else:
|
||||
np.savetxt(path + '.asc', vertices)
|
||||
|
||||
def write_obj_with_colors(obj_name, vertices, triangles, colors):
|
||||
''' Save 3D face model with texture represented by colors.
|
||||
Args:
|
||||
obj_name: str
|
||||
vertices: shape = (nver, 3)
|
||||
triangles: shape = (ntri, 3)
|
||||
colors: shape = (nver, 3)
|
||||
'''
|
||||
triangles = triangles.copy()
|
||||
triangles += 1 # meshlab start with 1
|
||||
|
||||
if obj_name.split('.')[-1] != 'obj':
|
||||
obj_name = obj_name + '.obj'
|
||||
|
||||
# write obj
|
||||
with open(obj_name, 'w') as f:
|
||||
|
||||
# write vertices & colors
|
||||
for i in range(vertices.shape[0]):
|
||||
# s = 'v {} {} {} \n'.format(vertices[0,i], vertices[1,i], vertices[2,i])
|
||||
s = 'v {} {} {} {} {} {}\n'.format(vertices[i, 0], vertices[i, 1], vertices[i, 2], colors[i, 0], colors[i, 1], colors[i, 2])
|
||||
f.write(s)
|
||||
|
||||
# write f: ver ind/ uv ind
|
||||
[k, ntri] = triangles.shape
|
||||
for i in range(triangles.shape[0]):
|
||||
# s = 'f {} {} {}\n'.format(triangles[i, 0], triangles[i, 1], triangles[i, 2])
|
||||
s = 'f {} {} {}\n'.format(triangles[i, 2], triangles[i, 1], triangles[i, 0])
|
||||
f.write(s)
|
||||
|
||||
## TODO: c++ version
|
||||
def write_obj_with_texture(obj_name, vertices, triangles, texture, uv_coords):
|
||||
''' Save 3D face model with texture represented by texture map.
|
||||
Ref: https://github.com/patrikhuber/eos/blob/bd00155ebae4b1a13b08bf5a991694d682abbada/include/eos/core/Mesh.hpp
|
||||
Args:
|
||||
obj_name: str
|
||||
vertices: shape = (nver, 3)
|
||||
triangles: shape = (ntri, 3)
|
||||
texture: shape = (256,256,3)
|
||||
uv_coords: shape = (nver, 3) max value<=1
|
||||
'''
|
||||
if obj_name.split('.')[-1] != 'obj':
|
||||
obj_name = obj_name + '.obj'
|
||||
mtl_name = obj_name.replace('.obj', '.mtl')
|
||||
texture_name = obj_name.replace('.obj', '_texture.png')
|
||||
|
||||
triangles = triangles.copy()
|
||||
triangles += 1 # mesh lab start with 1
|
||||
|
||||
# write obj
|
||||
with open(obj_name, 'w') as f:
|
||||
# first line: write mtlib(material library)
|
||||
s = "mtllib {}\n".format(os.path.abspath(mtl_name))
|
||||
f.write(s)
|
||||
|
||||
# write vertices
|
||||
for i in range(vertices.shape[0]):
|
||||
s = 'v {} {} {}\n'.format(vertices[i, 0], vertices[i, 1], vertices[i, 2])
|
||||
f.write(s)
|
||||
|
||||
# write uv coords
|
||||
for i in range(uv_coords.shape[0]):
|
||||
s = 'vt {} {}\n'.format(uv_coords[i,0], 1 - uv_coords[i,1])
|
||||
f.write(s)
|
||||
|
||||
f.write("usemtl FaceTexture\n")
|
||||
|
||||
# write f: ver ind/ uv ind
|
||||
for i in range(triangles.shape[0]):
|
||||
s = 'f {}/{} {}/{} {}/{}\n'.format(triangles[i,2], triangles[i,2], triangles[i,1], triangles[i,1], triangles[i,0], triangles[i,0])
|
||||
f.write(s)
|
||||
|
||||
# write mtl
|
||||
with open(mtl_name, 'w') as f:
|
||||
f.write("newmtl FaceTexture\n")
|
||||
s = 'map_Kd {}\n'.format(os.path.abspath(texture_name)) # map to image
|
||||
f.write(s)
|
||||
|
||||
# write texture as png
|
||||
imsave(texture_name, texture)
|
||||
|
||||
# c++ version
|
||||
def write_obj_with_colors_texture(obj_name, vertices, triangles, colors, texture, uv_coords):
|
||||
''' Save 3D face model with texture.
|
||||
Ref: https://github.com/patrikhuber/eos/blob/bd00155ebae4b1a13b08bf5a991694d682abbada/include/eos/core/Mesh.hpp
|
||||
Args:
|
||||
obj_name: str
|
||||
vertices: shape = (nver, 3)
|
||||
triangles: shape = (ntri, 3)
|
||||
colors: shape = (nver, 3)
|
||||
texture: shape = (256,256,3)
|
||||
uv_coords: shape = (nver, 3) max value<=1
|
||||
'''
|
||||
if obj_name.split('.')[-1] != 'obj':
|
||||
obj_name = obj_name + '.obj'
|
||||
mtl_name = obj_name.replace('.obj', '.mtl')
|
||||
texture_name = obj_name.replace('.obj', '_texture.png')
|
||||
|
||||
triangles = triangles.copy()
|
||||
triangles += 1 # mesh lab start with 1
|
||||
|
||||
# write obj
|
||||
vertices, colors, uv_coords = vertices.astype(np.float32).copy(), colors.astype(np.float32).copy(), uv_coords.astype(np.float32).copy()
|
||||
mesh_core_cython.write_obj_with_colors_texture_core(str.encode(obj_name), str.encode(os.path.abspath(mtl_name)), vertices, triangles, colors, uv_coords, vertices.shape[0], triangles.shape[0], uv_coords.shape[0])
|
||||
|
||||
# write mtl
|
||||
with open(mtl_name, 'w') as f:
|
||||
f.write("newmtl FaceTexture\n")
|
||||
s = 'map_Kd {}\n'.format(os.path.abspath(texture_name)) # map to image
|
||||
f.write(s)
|
||||
|
||||
# write texture as png
|
||||
# io.imsave(texture_name, texture)
|
||||
@@ -0,0 +1,213 @@
|
||||
'''
|
||||
Functions about lighting mesh(changing colors/texture of mesh).
|
||||
1. add light to colors/texture (shade each vertex)
|
||||
2. fit light according to colors/texture & image.
|
||||
'''
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import numpy as np
|
||||
# from .cython import mesh_core_cython
|
||||
|
||||
def get_normal(vertices, triangles):
|
||||
''' calculate normal direction in each vertex
|
||||
Args:
|
||||
vertices: [nver, 3]
|
||||
triangles: [ntri, 3]
|
||||
Returns:
|
||||
normal: [nver, 3]
|
||||
'''
|
||||
pt0 = vertices[triangles[:, 0], :] # [ntri, 3]
|
||||
pt1 = vertices[triangles[:, 1], :] # [ntri, 3]
|
||||
pt2 = vertices[triangles[:, 2], :] # [ntri, 3]
|
||||
tri_normal = np.cross(pt0 - pt1, pt0 - pt2) # [ntri, 3]. normal of each triangle
|
||||
|
||||
normal = np.zeros_like(vertices, dtype = np.float32).copy() # [nver, 3]
|
||||
# for i in range(triangles.shape[0]):
|
||||
# normal[triangles[i, 0], :] = normal[triangles[i, 0], :] + tri_normal[i, :]
|
||||
# normal[triangles[i, 1], :] = normal[triangles[i, 1], :] + tri_normal[i, :]
|
||||
# normal[triangles[i, 2], :] = normal[triangles[i, 2], :] + tri_normal[i, :]
|
||||
mesh_core_cython.get_normal_core(normal, tri_normal.astype(np.float32).copy(), triangles.copy(), triangles.shape[0])
|
||||
|
||||
# normalize to unit length
|
||||
mag = np.sum(normal**2, 1) # [nver]
|
||||
zero_ind = (mag == 0)
|
||||
mag[zero_ind] = 1
|
||||
normal[zero_ind, 0] = np.ones((np.sum(zero_ind)))
|
||||
|
||||
normal = normal/np.sqrt(mag[:,np.newaxis])
|
||||
|
||||
return normal
|
||||
|
||||
# TODO: test
|
||||
def add_light_sh(vertices, triangles, colors, sh_coeff):
|
||||
'''
|
||||
In 3d face, usually assume:
|
||||
1. The surface of face is Lambertian(reflect only the low frequencies of lighting)
|
||||
2. Lighting can be an arbitrary combination of point sources
|
||||
--> can be expressed in terms of spherical harmonics(omit the lighting coefficients)
|
||||
I = albedo * (sh(n) x sh_coeff)
|
||||
|
||||
albedo: n x 1
|
||||
sh_coeff: 9 x 1
|
||||
Y(n) = (1, n_x, n_y, n_z, n_xn_y, n_xn_z, n_yn_z, n_x^2 - n_y^2, 3n_z^2 - 1)': n x 9
|
||||
# Y(n) = (1, n_x, n_y, n_z)': n x 4
|
||||
|
||||
Args:
|
||||
vertices: [nver, 3]
|
||||
triangles: [ntri, 3]
|
||||
colors: [nver, 3] albedo
|
||||
sh_coeff: [9, 1] spherical harmonics coefficients
|
||||
|
||||
Returns:
|
||||
lit_colors: [nver, 3]
|
||||
'''
|
||||
assert vertices.shape[0] == colors.shape[0]
|
||||
nver = vertices.shape[0]
|
||||
normal = get_normal(vertices, triangles) # [nver, 3]
|
||||
sh = np.array((np.ones(nver), n[:,0], n[:,1], n[:,2], n[:,0]*n[:,1], n[:,0]*n[:,2], n[:,1]*n[:,2], n[:,0]**2 - n[:,1]**2, 3*(n[:,2]**2) - 1)) # [nver, 9]
|
||||
ref = sh.dot(sh_coeff) #[nver, 1]
|
||||
lit_colors = colors*ref
|
||||
return lit_colors
|
||||
|
||||
|
||||
def add_light(vertices, triangles, colors, light_positions = 0, light_intensities = 0):
|
||||
''' Gouraud shading. add point lights.
|
||||
In 3d face, usually assume:
|
||||
1. The surface of face is Lambertian(reflect only the low frequencies of lighting)
|
||||
2. Lighting can be an arbitrary combination of point sources
|
||||
3. No specular (unless skin is oil, 23333)
|
||||
|
||||
Ref: https://cs184.eecs.berkeley.edu/lecture/pipeline
|
||||
Args:
|
||||
vertices: [nver, 3]
|
||||
triangles: [ntri, 3]
|
||||
light_positions: [nlight, 3]
|
||||
light_intensities: [nlight, 3]
|
||||
Returns:
|
||||
lit_colors: [nver, 3]
|
||||
'''
|
||||
nver = vertices.shape[0]
|
||||
normals = get_normal(vertices, triangles) # [nver, 3]
|
||||
|
||||
# ambient
|
||||
# La = ka*Ia
|
||||
|
||||
# diffuse
|
||||
# Ld = kd*(I/r^2)max(0, nxl)
|
||||
direction_to_lights = vertices[np.newaxis, :, :] - light_positions[:, np.newaxis, :] # [nlight, nver, 3]
|
||||
direction_to_lights_n = np.sqrt(np.sum(direction_to_lights**2, axis = 2)) # [nlight, nver]
|
||||
direction_to_lights = direction_to_lights/direction_to_lights_n[:, :, np.newaxis]
|
||||
normals_dot_lights = normals[np.newaxis, :, :]*direction_to_lights # [nlight, nver, 3]
|
||||
normals_dot_lights = np.sum(normals_dot_lights, axis = 2) # [nlight, nver]
|
||||
diffuse_output = colors[np.newaxis, :, :]*normals_dot_lights[:, :, np.newaxis]*light_intensities[:, np.newaxis, :]
|
||||
diffuse_output = np.sum(diffuse_output, axis = 0) # [nver, 3]
|
||||
|
||||
# specular
|
||||
# h = (v + l)/(|v + l|) bisector
|
||||
# Ls = ks*(I/r^2)max(0, nxh)^p
|
||||
# increasing p narrows the reflectionlob
|
||||
|
||||
lit_colors = diffuse_output # only diffuse part here.
|
||||
lit_colors = np.minimum(np.maximum(lit_colors, 0), 1)
|
||||
return lit_colors
|
||||
|
||||
|
||||
|
||||
## TODO. estimate light(sh coeff)
|
||||
## -------------------------------- estimate. can not use now.
|
||||
def fit_light(image, vertices, colors, triangles, vis_ind, lamb = 10, max_iter = 3):
|
||||
[h, w, c] = image.shape
|
||||
|
||||
# surface normal
|
||||
norm = get_normal(vertices, triangles)
|
||||
|
||||
nver = vertices.shape[1]
|
||||
|
||||
# vertices --> corresponding image pixel
|
||||
pt2d = vertices[:2, :]
|
||||
|
||||
pt2d[0,:] = np.minimum(np.maximum(pt2d[0,:], 0), w - 1)
|
||||
pt2d[1,:] = np.minimum(np.maximum(pt2d[1,:], 0), h - 1)
|
||||
pt2d = np.round(pt2d).astype(np.int32) # 2 x nver
|
||||
|
||||
image_pixel = image[pt2d[1,:], pt2d[0,:], :] # nver x 3
|
||||
image_pixel = image_pixel.T # 3 x nver
|
||||
|
||||
# vertices --> corresponding mean texture pixel with illumination
|
||||
# Spherical Harmonic Basis
|
||||
harmonic_dim = 9
|
||||
nx = norm[0,:];
|
||||
ny = norm[1,:];
|
||||
nz = norm[2,:];
|
||||
harmonic = np.zeros((nver, harmonic_dim))
|
||||
|
||||
pi = np.pi
|
||||
harmonic[:,0] = np.sqrt(1/(4*pi)) * np.ones((nver,));
|
||||
harmonic[:,1] = np.sqrt(3/(4*pi)) * nx;
|
||||
harmonic[:,2] = np.sqrt(3/(4*pi)) * ny;
|
||||
harmonic[:,3] = np.sqrt(3/(4*pi)) * nz;
|
||||
harmonic[:,4] = 1/2. * np.sqrt(3/(4*pi)) * (2*nz**2 - nx**2 - ny**2);
|
||||
harmonic[:,5] = 3 * np.sqrt(5/(12*pi)) * (ny*nz);
|
||||
harmonic[:,6] = 3 * np.sqrt(5/(12*pi)) * (nx*nz);
|
||||
harmonic[:,7] = 3 * np.sqrt(5/(12*pi)) * (nx*ny);
|
||||
harmonic[:,8] = 3/2. * np.sqrt(5/(12*pi)) * (nx*nx - ny*ny);
|
||||
|
||||
'''
|
||||
I' = sum(albedo * lj * hj) j = 0:9 (albedo = tex)
|
||||
set A = albedo*h (n x 9)
|
||||
alpha = lj (9 x 1)
|
||||
Y = I (n x 1)
|
||||
Y' = A.dot(alpha)
|
||||
|
||||
opt function:
|
||||
||Y - A*alpha|| + lambda*(alpha'*alpha)
|
||||
result:
|
||||
A'*(Y - A*alpha) + lambda*alpha = 0
|
||||
==>
|
||||
(A'*A*alpha - lambda)*alpha = A'*Y
|
||||
left: 9 x 9
|
||||
right: 9 x 1
|
||||
'''
|
||||
n_vis_ind = len(vis_ind)
|
||||
n = n_vis_ind*c
|
||||
|
||||
Y = np.zeros((n, 1))
|
||||
A = np.zeros((n, 9))
|
||||
light = np.zeros((3, 1))
|
||||
|
||||
for k in range(c):
|
||||
Y[k*n_vis_ind:(k+1)*n_vis_ind, :] = image_pixel[k, vis_ind][:, np.newaxis]
|
||||
A[k*n_vis_ind:(k+1)*n_vis_ind, :] = texture[k, vis_ind][:, np.newaxis] * harmonic[vis_ind, :]
|
||||
Ac = texture[k, vis_ind][:, np.newaxis]
|
||||
Yc = image_pixel[k, vis_ind][:, np.newaxis]
|
||||
light[k] = (Ac.T.dot(Yc))/(Ac.T.dot(Ac))
|
||||
|
||||
for i in range(max_iter):
|
||||
|
||||
Yc = Y.copy()
|
||||
for k in range(c):
|
||||
Yc[k*n_vis_ind:(k+1)*n_vis_ind, :] /= light[k]
|
||||
|
||||
# update alpha
|
||||
equation_left = np.dot(A.T, A) + lamb*np.eye(harmonic_dim); # why + ?
|
||||
equation_right = np.dot(A.T, Yc)
|
||||
alpha = np.dot(np.linalg.inv(equation_left), equation_right)
|
||||
|
||||
# update light
|
||||
for k in range(c):
|
||||
Ac = A[k*n_vis_ind:(k+1)*n_vis_ind, :].dot(alpha)
|
||||
Yc = Y[k*n_vis_ind:(k+1)*n_vis_ind, :]
|
||||
light[k] = (Ac.T.dot(Yc))/(Ac.T.dot(Ac))
|
||||
|
||||
appearance = np.zeros_like(texture)
|
||||
for k in range(c):
|
||||
tmp = np.dot(harmonic*texture[k, :][:, np.newaxis], alpha*light[k])
|
||||
appearance[k,:] = tmp.T
|
||||
|
||||
appearance = np.minimum(np.maximum(appearance, 0), 1)
|
||||
|
||||
return appearance
|
||||
|
||||
@@ -0,0 +1,135 @@
|
||||
'''
|
||||
functions about rendering mesh(from 3d obj to 2d image).
|
||||
only use rasterization render here.
|
||||
Note that:
|
||||
1. Generally, render func includes camera, light, raterize. Here no camera and light(I write these in other files)
|
||||
2. Generally, the input vertices are normalized to [-1,1] and cetered on [0, 0]. (in world space)
|
||||
Here, the vertices are using image coords, which centers on [w/2, h/2] with the y-axis pointing to oppisite direction.
|
||||
Means: render here only conducts interpolation.(I just want to make the input flexible)
|
||||
|
||||
Author: Yao Feng
|
||||
Mail: yaofeng1995@gmail.com
|
||||
'''
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import numpy as np
|
||||
from time import time
|
||||
|
||||
# from .cython import mesh_core_cython
|
||||
|
||||
def rasterize_triangles(vertices, triangles, h, w):
|
||||
'''
|
||||
Args:
|
||||
vertices: [nver, 3]
|
||||
triangles: [ntri, 3]
|
||||
h: height
|
||||
w: width
|
||||
Returns:
|
||||
depth_buffer: [h, w] saves the depth, here, the bigger the z, the fronter the point.
|
||||
triangle_buffer: [h, w] saves the tri id(-1 for no triangle).
|
||||
barycentric_weight: [h, w, 3] saves corresponding barycentric weight.
|
||||
|
||||
# Each triangle has 3 vertices & Each vertex has 3 coordinates x, y, z.
|
||||
# h, w is the size of rendering
|
||||
'''
|
||||
|
||||
# initial
|
||||
depth_buffer = np.zeros([h, w]) - 999999. #set the initial z to the farest position
|
||||
triangle_buffer = np.zeros([h, w], dtype = np.int32) - 1 # if tri id = -1, the pixel has no triangle correspondance
|
||||
barycentric_weight = np.zeros([h, w, 3], dtype = np.float32) #
|
||||
|
||||
vertices = vertices.astype(np.float32).copy()
|
||||
triangles = triangles.astype(np.int32).copy()
|
||||
|
||||
mesh_core_cython.rasterize_triangles_core(
|
||||
vertices, triangles,
|
||||
depth_buffer, triangle_buffer, barycentric_weight,
|
||||
vertices.shape[0], triangles.shape[0],
|
||||
h, w)
|
||||
|
||||
def render_colors(vertices, triangles, colors, h, w, c = 3, BG = None):
|
||||
''' render mesh with colors
|
||||
Args:
|
||||
vertices: [nver, 3]
|
||||
triangles: [ntri, 3]
|
||||
colors: [nver, 3]
|
||||
h: height
|
||||
w: width
|
||||
c: channel
|
||||
BG: background image
|
||||
Returns:
|
||||
image: [h, w, c]. rendered image./rendering.
|
||||
'''
|
||||
|
||||
# initial
|
||||
if BG is None:
|
||||
image = np.zeros((h, w, c), dtype = np.float32)
|
||||
else:
|
||||
assert BG.shape[0] == h and BG.shape[1] == w and BG.shape[2] == c
|
||||
image = BG
|
||||
depth_buffer = np.zeros([h, w], dtype = np.float32, order = 'C') - 999999.
|
||||
|
||||
# change orders. --> C-contiguous order(column major)
|
||||
vertices = vertices.astype(np.float32).copy()
|
||||
triangles = triangles.astype(np.int32).copy()
|
||||
colors = colors.astype(np.float32).copy()
|
||||
###
|
||||
st = time()
|
||||
mesh_core_cython.render_colors_core(
|
||||
image, vertices, triangles,
|
||||
colors,
|
||||
depth_buffer,
|
||||
vertices.shape[0], triangles.shape[0],
|
||||
h, w, c)
|
||||
return image
|
||||
|
||||
|
||||
def render_texture(vertices, triangles, texture, tex_coords, tex_triangles, h, w, c = 3, mapping_type = 'nearest', BG = None):
|
||||
''' render mesh with texture map
|
||||
Args:
|
||||
vertices: [3, nver]
|
||||
triangles: [3, ntri]
|
||||
texture: [tex_h, tex_w, 3]
|
||||
tex_coords: [ntexcoords, 3]
|
||||
tex_triangles: [ntri, 3]
|
||||
h: height of rendering
|
||||
w: width of rendering
|
||||
c: channel
|
||||
mapping_type: 'bilinear' or 'nearest'
|
||||
'''
|
||||
# initial
|
||||
if BG is None:
|
||||
image = np.zeros((h, w, c), dtype = np.float32)
|
||||
else:
|
||||
assert BG.shape[0] == h and BG.shape[1] == w and BG.shape[2] == c
|
||||
image = BG.astype(np.float32)
|
||||
|
||||
depth_buffer = np.zeros([h, w], dtype = np.float32, order = 'C') - 999999.
|
||||
|
||||
tex_h, tex_w, tex_c = texture.shape
|
||||
if mapping_type == 'nearest':
|
||||
mt = int(0)
|
||||
elif mapping_type == 'bilinear':
|
||||
mt = int(1)
|
||||
else:
|
||||
mt = int(0)
|
||||
|
||||
# -> C order
|
||||
vertices = vertices.astype(np.float32).copy()
|
||||
triangles = triangles.astype(np.int32).copy()
|
||||
texture = texture.astype(np.float32).copy()
|
||||
tex_coords = tex_coords.astype(np.float32).copy()
|
||||
tex_triangles = tex_triangles.astype(np.int32).copy()
|
||||
|
||||
mesh_core_cython.render_texture_core(
|
||||
image, vertices, triangles,
|
||||
texture, tex_coords, tex_triangles,
|
||||
depth_buffer,
|
||||
vertices.shape[0], tex_coords.shape[0], triangles.shape[0],
|
||||
h, w, c,
|
||||
tex_h, tex_w, tex_c,
|
||||
mt)
|
||||
return image
|
||||
|
||||
@@ -0,0 +1,383 @@
|
||||
'''
|
||||
Functions about transforming mesh(changing the position: modify vertices).
|
||||
1. forward: transform(transform, camera, project).
|
||||
2. backward: estimate transform matrix from correspondences.
|
||||
|
||||
Author: Yao Feng
|
||||
Mail: yaofeng1995@gmail.com
|
||||
'''
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import numpy as np
|
||||
import math
|
||||
from math import cos, sin
|
||||
|
||||
def angle2matrix(angles):
|
||||
''' get rotation matrix from three rotation angles(degree). right-handed.
|
||||
Args:
|
||||
angles: [3,]. x, y, z angles
|
||||
x: pitch. positive for looking down.
|
||||
y: yaw. positive for looking left.
|
||||
z: roll. positive for tilting head right.
|
||||
Returns:
|
||||
R: [3, 3]. rotation matrix.
|
||||
'''
|
||||
x, y, z = np.deg2rad(angles[0]), np.deg2rad(angles[1]), np.deg2rad(angles[2])
|
||||
# x
|
||||
Rx=np.array([[1, 0, 0],
|
||||
[0, cos(x), -sin(x)],
|
||||
[0, sin(x), cos(x)]])
|
||||
# y
|
||||
Ry=np.array([[ cos(y), 0, sin(y)],
|
||||
[ 0, 1, 0],
|
||||
[-sin(y), 0, cos(y)]])
|
||||
# z
|
||||
Rz=np.array([[cos(z), -sin(z), 0],
|
||||
[sin(z), cos(z), 0],
|
||||
[ 0, 0, 1]])
|
||||
|
||||
R=Rz.dot(Ry.dot(Rx))
|
||||
return R.astype(np.float32)
|
||||
|
||||
def angle2matrix_3ddfa(angles):
|
||||
''' get rotation matrix from three rotation angles(radian). The same as in 3DDFA.
|
||||
Args:
|
||||
angles: [3,]. x, y, z angles
|
||||
x: pitch.
|
||||
y: yaw.
|
||||
z: roll.
|
||||
Returns:
|
||||
R: 3x3. rotation matrix.
|
||||
'''
|
||||
# x, y, z = np.deg2rad(angles[0]), np.deg2rad(angles[1]), np.deg2rad(angles[2])
|
||||
x, y, z = angles[0], angles[1], angles[2]
|
||||
|
||||
# x
|
||||
Rx=np.array([[1, 0, 0],
|
||||
[0, cos(x), sin(x)],
|
||||
[0, -sin(x), cos(x)]])
|
||||
# y
|
||||
Ry=np.array([[ cos(y), 0, -sin(y)],
|
||||
[ 0, 1, 0],
|
||||
[sin(y), 0, cos(y)]])
|
||||
# z
|
||||
Rz=np.array([[cos(z), sin(z), 0],
|
||||
[-sin(z), cos(z), 0],
|
||||
[ 0, 0, 1]])
|
||||
R = Rx.dot(Ry).dot(Rz)
|
||||
return R.astype(np.float32)
|
||||
|
||||
|
||||
## ------------------------------------------ 1. transform(transform, project, camera).
|
||||
## ---------- 3d-3d transform. Transform obj in world space
|
||||
def rotate(vertices, angles):
|
||||
''' rotate vertices.
|
||||
X_new = R.dot(X). X: 3 x 1
|
||||
Args:
|
||||
vertices: [nver, 3].
|
||||
rx, ry, rz: degree angles
|
||||
rx: pitch. positive for looking down
|
||||
ry: yaw. positive for looking left
|
||||
rz: roll. positive for tilting head right
|
||||
Returns:
|
||||
rotated vertices: [nver, 3]
|
||||
'''
|
||||
R = angle2matrix(angles)
|
||||
rotated_vertices = vertices.dot(R.T)
|
||||
|
||||
return rotated_vertices
|
||||
|
||||
def similarity_transform(vertices, s, R, t3d):
|
||||
''' similarity transform. dof = 7.
|
||||
3D: s*R.dot(X) + t
|
||||
Homo: M = [[sR, t],[0^T, 1]]. M.dot(X)
|
||||
Args:(float32)
|
||||
vertices: [nver, 3].
|
||||
s: [1,]. scale factor.
|
||||
R: [3,3]. rotation matrix.
|
||||
t3d: [3,]. 3d translation vector.
|
||||
Returns:
|
||||
transformed vertices: [nver, 3]
|
||||
'''
|
||||
t3d = np.squeeze(np.array(t3d, dtype = np.float32))
|
||||
transformed_vertices = s * vertices.dot(R.T) + t3d[np.newaxis, :]
|
||||
|
||||
return transformed_vertices
|
||||
|
||||
|
||||
## -------------- Camera. from world space to camera space
|
||||
# Ref: https://cs184.eecs.berkeley.edu/lecture/transforms-2
|
||||
def normalize(x):
|
||||
epsilon = 1e-12
|
||||
norm = np.sqrt(np.sum(x**2, axis = 0))
|
||||
norm = np.maximum(norm, epsilon)
|
||||
return x/norm
|
||||
|
||||
def lookat_camera(vertices, eye, at = None, up = None):
|
||||
""" 'look at' transformation: from world space to camera space
|
||||
standard camera space:
|
||||
camera located at the origin.
|
||||
looking down negative z-axis.
|
||||
vertical vector is y-axis.
|
||||
Xcam = R(X - C)
|
||||
Homo: [[R, -RC], [0, 1]]
|
||||
Args:
|
||||
vertices: [nver, 3]
|
||||
eye: [3,] the XYZ world space position of the camera.
|
||||
at: [3,] a position along the center of the camera's gaze.
|
||||
up: [3,] up direction
|
||||
Returns:
|
||||
transformed_vertices: [nver, 3]
|
||||
"""
|
||||
if at is None:
|
||||
at = np.array([0, 0, 0], np.float32)
|
||||
if up is None:
|
||||
up = np.array([0, 1, 0], np.float32)
|
||||
|
||||
eye = np.array(eye).astype(np.float32)
|
||||
at = np.array(at).astype(np.float32)
|
||||
z_aixs = -normalize(at - eye) # look forward
|
||||
x_aixs = normalize(np.cross(up, z_aixs)) # look right
|
||||
y_axis = np.cross(z_aixs, x_aixs) # look up
|
||||
|
||||
R = np.stack((x_aixs, y_axis, z_aixs))#, axis = 0) # 3 x 3
|
||||
transformed_vertices = vertices - eye # translation
|
||||
transformed_vertices = transformed_vertices.dot(R.T) # rotation
|
||||
return transformed_vertices
|
||||
|
||||
## --------- 3d-2d project. from camera space to image plane
|
||||
# generally, image plane only keeps x,y channels, here reserve z channel for calculating z-buffer.
|
||||
def orthographic_project(vertices):
|
||||
''' scaled orthographic projection(just delete z)
|
||||
assumes: variations in depth over the object is small relative to the mean distance from camera to object
|
||||
x -> x*f/z, y -> x*f/z, z -> f.
|
||||
for point i,j. zi~=zj. so just delete z
|
||||
** often used in face
|
||||
Homo: P = [[1,0,0,0], [0,1,0,0], [0,0,1,0]]
|
||||
Args:
|
||||
vertices: [nver, 3]
|
||||
Returns:
|
||||
projected_vertices: [nver, 3] if isKeepZ=True. [nver, 2] if isKeepZ=False.
|
||||
'''
|
||||
return vertices.copy()
|
||||
|
||||
def perspective_project(vertices, fovy, aspect_ratio = 1., near = 0.1, far = 1000.):
|
||||
''' perspective projection.
|
||||
Args:
|
||||
vertices: [nver, 3]
|
||||
fovy: vertical angular field of view. degree.
|
||||
aspect_ratio : width / height of field of view
|
||||
near : depth of near clipping plane
|
||||
far : depth of far clipping plane
|
||||
Returns:
|
||||
projected_vertices: [nver, 3]
|
||||
'''
|
||||
fovy = np.deg2rad(fovy)
|
||||
top = near*np.tan(fovy)
|
||||
bottom = -top
|
||||
right = top*aspect_ratio
|
||||
left = -right
|
||||
|
||||
#-- homo
|
||||
P = np.array([[near/right, 0, 0, 0],
|
||||
[0, near/top, 0, 0],
|
||||
[0, 0, -(far+near)/(far-near), -2*far*near/(far-near)],
|
||||
[0, 0, -1, 0]])
|
||||
vertices_homo = np.hstack((vertices, np.ones((vertices.shape[0], 1)))) # [nver, 4]
|
||||
projected_vertices = vertices_homo.dot(P.T)
|
||||
projected_vertices = projected_vertices/projected_vertices[:,3:]
|
||||
projected_vertices = projected_vertices[:,:3]
|
||||
projected_vertices[:,2] = -projected_vertices[:,2]
|
||||
|
||||
#-- non homo. only fovy
|
||||
# projected_vertices = vertices.copy()
|
||||
# projected_vertices[:,0] = -(near/right)*vertices[:,0]/vertices[:,2]
|
||||
# projected_vertices[:,1] = -(near/top)*vertices[:,1]/vertices[:,2]
|
||||
return projected_vertices
|
||||
|
||||
|
||||
def to_image(vertices, h, w, is_perspective = False):
|
||||
''' change vertices to image coord system
|
||||
3d system: XYZ, center(0, 0, 0)
|
||||
2d image: x(u), y(v). center(w/2, h/2), flip y-axis.
|
||||
Args:
|
||||
vertices: [nver, 3]
|
||||
h: height of the rendering
|
||||
w : width of the rendering
|
||||
Returns:
|
||||
projected_vertices: [nver, 3]
|
||||
'''
|
||||
image_vertices = vertices.copy()
|
||||
if is_perspective:
|
||||
# if perspective, the projected vertices are normalized to [-1, 1]. so change it to image size first.
|
||||
image_vertices[:,0] = image_vertices[:,0]*w/2
|
||||
image_vertices[:,1] = image_vertices[:,1]*h/2
|
||||
# move to center of image
|
||||
image_vertices[:,0] = image_vertices[:,0] + w/2
|
||||
image_vertices[:,1] = image_vertices[:,1] + h/2
|
||||
# flip vertices along y-axis.
|
||||
image_vertices[:,1] = h - image_vertices[:,1] - 1
|
||||
return image_vertices
|
||||
|
||||
|
||||
#### -------------------------------------------2. estimate transform matrix from correspondences.
|
||||
def estimate_affine_matrix_3d23d(X, Y):
|
||||
''' Using least-squares solution
|
||||
Args:
|
||||
X: [n, 3]. 3d points(fixed)
|
||||
Y: [n, 3]. corresponding 3d points(moving). Y = PX
|
||||
Returns:
|
||||
P_Affine: (3, 4). Affine camera matrix (the third row is [0, 0, 0, 1]).
|
||||
'''
|
||||
X_homo = np.hstack((X, np.ones([X.shape[1],1]))) #n x 4
|
||||
P = np.linalg.lstsq(X_homo, Y)[0].T # Affine matrix. 3 x 4
|
||||
return P
|
||||
|
||||
def estimate_affine_matrix_3d22d(X, x):
|
||||
''' Using Golden Standard Algorithm for estimating an affine camera
|
||||
matrix P from world to image correspondences.
|
||||
See Alg.7.2. in MVGCV
|
||||
Code Ref: https://github.com/patrikhuber/eos/blob/master/include/eos/fitting/affine_camera_estimation.hpp
|
||||
x_homo = X_homo.dot(P_Affine)
|
||||
Args:
|
||||
X: [n, 3]. corresponding 3d points(fixed)
|
||||
x: [n, 2]. n>=4. 2d points(moving). x = PX
|
||||
Returns:
|
||||
P_Affine: [3, 4]. Affine camera matrix
|
||||
'''
|
||||
X = X.T; x = x.T
|
||||
assert(x.shape[1] == X.shape[1])
|
||||
n = x.shape[1]
|
||||
assert(n >= 4)
|
||||
|
||||
#--- 1. normalization
|
||||
# 2d points
|
||||
mean = np.mean(x, 1) # (2,)
|
||||
x = x - np.tile(mean[:, np.newaxis], [1, n])
|
||||
average_norm = np.mean(np.sqrt(np.sum(x**2, 0)))
|
||||
scale = np.sqrt(2) / average_norm
|
||||
x = scale * x
|
||||
|
||||
T = np.zeros((3,3), dtype = np.float32)
|
||||
T[0, 0] = T[1, 1] = scale
|
||||
T[:2, 2] = -mean*scale
|
||||
T[2, 2] = 1
|
||||
|
||||
# 3d points
|
||||
X_homo = np.vstack((X, np.ones((1, n))))
|
||||
mean = np.mean(X, 1) # (3,)
|
||||
X = X - np.tile(mean[:, np.newaxis], [1, n])
|
||||
m = X_homo[:3,:] - X
|
||||
average_norm = np.mean(np.sqrt(np.sum(X**2, 0)))
|
||||
scale = np.sqrt(3) / average_norm
|
||||
X = scale * X
|
||||
|
||||
U = np.zeros((4,4), dtype = np.float32)
|
||||
U[0, 0] = U[1, 1] = U[2, 2] = scale
|
||||
U[:3, 3] = -mean*scale
|
||||
U[3, 3] = 1
|
||||
|
||||
# --- 2. equations
|
||||
A = np.zeros((n*2, 8), dtype = np.float32);
|
||||
X_homo = np.vstack((X, np.ones((1, n)))).T
|
||||
A[:n, :4] = X_homo
|
||||
A[n:, 4:] = X_homo
|
||||
b = np.reshape(x, [-1, 1])
|
||||
|
||||
# --- 3. solution
|
||||
p_8 = np.linalg.pinv(A).dot(b)
|
||||
P = np.zeros((3, 4), dtype = np.float32)
|
||||
P[0, :] = p_8[:4, 0]
|
||||
P[1, :] = p_8[4:, 0]
|
||||
P[-1, -1] = 1
|
||||
|
||||
# --- 4. denormalization
|
||||
P_Affine = np.linalg.inv(T).dot(P.dot(U))
|
||||
return P_Affine
|
||||
|
||||
def P2sRt(P):
|
||||
''' decompositing camera matrix P
|
||||
Args:
|
||||
P: (3, 4). Affine Camera Matrix.
|
||||
Returns:
|
||||
s: scale factor.
|
||||
R: (3, 3). rotation matrix.
|
||||
t: (3,). translation.
|
||||
'''
|
||||
t = P[:, 3]
|
||||
R1 = P[0:1, :3]
|
||||
R2 = P[1:2, :3]
|
||||
s = (np.linalg.norm(R1) + np.linalg.norm(R2))/2.0
|
||||
r1 = R1/np.linalg.norm(R1)
|
||||
r2 = R2/np.linalg.norm(R2)
|
||||
r3 = np.cross(r1, r2)
|
||||
|
||||
R = np.concatenate((r1, r2, r3), 0)
|
||||
return s, R, t
|
||||
|
||||
#Ref: https://www.learnopencv.com/rotation-matrix-to-euler-angles/
|
||||
def isRotationMatrix(R):
|
||||
''' checks if a matrix is a valid rotation matrix(whether orthogonal or not)
|
||||
'''
|
||||
Rt = np.transpose(R)
|
||||
shouldBeIdentity = np.dot(Rt, R)
|
||||
I = np.identity(3, dtype = R.dtype)
|
||||
n = np.linalg.norm(I - shouldBeIdentity)
|
||||
return n < 1e-6
|
||||
|
||||
def matrix2angle(R):
|
||||
''' get three Euler angles from Rotation Matrix
|
||||
Args:
|
||||
R: (3,3). rotation matrix
|
||||
Returns:
|
||||
x: pitch
|
||||
y: yaw
|
||||
z: roll
|
||||
'''
|
||||
assert(isRotationMatrix)
|
||||
sy = math.sqrt(R[0,0] * R[0,0] + R[1,0] * R[1,0])
|
||||
|
||||
singular = sy < 1e-6
|
||||
|
||||
if not singular :
|
||||
x = math.atan2(R[2,1] , R[2,2])
|
||||
y = math.atan2(-R[2,0], sy)
|
||||
z = math.atan2(R[1,0], R[0,0])
|
||||
else :
|
||||
x = math.atan2(-R[1,2], R[1,1])
|
||||
y = math.atan2(-R[2,0], sy)
|
||||
z = 0
|
||||
|
||||
# rx, ry, rz = np.rad2deg(x), np.rad2deg(y), np.rad2deg(z)
|
||||
rx, ry, rz = x*180/np.pi, y*180/np.pi, z*180/np.pi
|
||||
return rx, ry, rz
|
||||
|
||||
# def matrix2angle(R):
|
||||
# ''' compute three Euler angles from a Rotation Matrix. Ref: http://www.gregslabaugh.net/publications/euler.pdf
|
||||
# Args:
|
||||
# R: (3,3). rotation matrix
|
||||
# Returns:
|
||||
# x: yaw
|
||||
# y: pitch
|
||||
# z: roll
|
||||
# '''
|
||||
# # assert(isRotationMatrix(R))
|
||||
|
||||
# if R[2,0] !=1 or R[2,0] != -1:
|
||||
# x = math.asin(R[2,0])
|
||||
# y = math.atan2(R[2,1]/cos(x), R[2,2]/cos(x))
|
||||
# z = math.atan2(R[1,0]/cos(x), R[0,0]/cos(x))
|
||||
|
||||
# else:# Gimbal lock
|
||||
# z = 0 #can be anything
|
||||
# if R[2,0] == -1:
|
||||
# x = np.pi/2
|
||||
# y = z + math.atan2(R[0,1], R[0,2])
|
||||
# else:
|
||||
# x = -np.pi/2
|
||||
# y = -z + math.atan2(-R[0,1], -R[0,2])
|
||||
|
||||
# return x, y, z
|
||||
@@ -0,0 +1,24 @@
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
# from skimage import measure
|
||||
from mpl_toolkits.mplot3d import Axes3D
|
||||
|
||||
def plot_mesh(vertices, triangles, subplot = [1,1,1], title = 'mesh', el = 90, az = -90, lwdt=.1, dist = 6, color = "grey"):
|
||||
'''
|
||||
plot the mesh
|
||||
Args:
|
||||
vertices: [nver, 3]
|
||||
triangles: [ntri, 3]
|
||||
'''
|
||||
ax = plt.subplot(subplot[0], subplot[1], subplot[2], projection = '3d')
|
||||
ax.plot_trisurf(vertices[:, 0], vertices[:, 1], vertices[:, 2], triangles = triangles, lw = lwdt, color = color, alpha = 1)
|
||||
ax.axis("off")
|
||||
ax.view_init(elev = el, azim = az)
|
||||
ax.dist = dist
|
||||
plt.title(title)
|
||||
|
||||
### -------------- Todo: use vtk to visualize mesh? or visvis? or VisPy?
|
||||
@@ -0,0 +1,7 @@
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
from .. import mesh
|
||||
from .morphabel_model import MorphabelModel
|
||||
from . import load
|
||||
@@ -0,0 +1,270 @@
|
||||
'''
|
||||
Estimating parameters about vertices: shape para, exp para, pose para(s, R, t)
|
||||
'''
|
||||
import numpy as np
|
||||
|
||||
''' TODO: a clear document.
|
||||
Given: image_points, 3D Model, Camera Matrix(s, R, t2d)
|
||||
Estimate: shape parameters, expression parameters
|
||||
|
||||
Inference:
|
||||
|
||||
projected_vertices = s*P*R(mu + shape + exp) + t2d --> image_points
|
||||
s*P*R*shape + s*P*R(mu + exp) + t2d --> image_poitns
|
||||
|
||||
# Define:
|
||||
X = vertices
|
||||
x_hat = projected_vertices
|
||||
x = image_points
|
||||
A = s*P*R
|
||||
b = s*P*R(mu + exp) + t2d
|
||||
==>
|
||||
x_hat = A*shape + b (2 x n)
|
||||
|
||||
A*shape (2 x n)
|
||||
shape = reshape(shapePC * sp) (3 x n)
|
||||
shapePC*sp : (3n x 1)
|
||||
|
||||
* flatten:
|
||||
x_hat_flatten = A*shape + b_flatten (2n x 1)
|
||||
A*shape (2n x 1)
|
||||
--> A*shapePC (2n x 199) sp: 199 x 1
|
||||
|
||||
# Define:
|
||||
pc_2d = A* reshape(shapePC)
|
||||
pc_2d_flatten = flatten(pc_2d) (2n x 199)
|
||||
|
||||
=====>
|
||||
x_hat_flatten = pc_2d_flatten * sp + b_flatten ---> x_flatten (2n x 1)
|
||||
|
||||
Goals:
|
||||
(ignore flatten, pc_2d-->pc)
|
||||
min E = || x_hat - x || + lambda*sum(sp/sigma)^2
|
||||
= || pc * sp + b - x || + lambda*sum(sp/sigma)^2
|
||||
|
||||
Solve:
|
||||
d(E)/d(sp) = 0
|
||||
2 * pc' * (pc * sp + b - x) + 2 * lambda * sp / (sigma' * sigma) = 0
|
||||
|
||||
Get:
|
||||
(pc' * pc + lambda / (sigma'* sigma)) * sp = pc' * (x - b)
|
||||
|
||||
'''
|
||||
|
||||
def estimate_shape(x, shapeMU, shapePC, shapeEV, expression, s, R, t2d, lamb = 3000):
|
||||
'''
|
||||
Args:
|
||||
x: (2, n). image points (to be fitted)
|
||||
shapeMU: (3n, 1)
|
||||
shapePC: (3n, n_sp)
|
||||
shapeEV: (n_sp, 1)
|
||||
expression: (3, n)
|
||||
s: scale
|
||||
R: (3, 3). rotation matrix
|
||||
t2d: (2,). 2d translation
|
||||
lambda: regulation coefficient
|
||||
|
||||
Returns:
|
||||
shape_para: (n_sp, 1) shape parameters(coefficients)
|
||||
'''
|
||||
x = x.copy()
|
||||
assert(shapeMU.shape[0] == shapePC.shape[0])
|
||||
assert(shapeMU.shape[0] == x.shape[1]*3)
|
||||
|
||||
dof = shapePC.shape[1]
|
||||
|
||||
n = x.shape[1]
|
||||
sigma = shapeEV
|
||||
t2d = np.array(t2d)
|
||||
P = np.array([[1, 0, 0], [0, 1, 0]], dtype = np.float32)
|
||||
A = s*P.dot(R)
|
||||
|
||||
# --- calc pc
|
||||
pc_3d = np.resize(shapePC.T, [dof, n, 3]) # 199 x n x 3
|
||||
pc_3d = np.reshape(pc_3d, [dof*n, 3])
|
||||
pc_2d = pc_3d.dot(A.T.copy()) # 199 x n x 2
|
||||
|
||||
pc = np.reshape(pc_2d, [dof, -1]).T # 2n x 199
|
||||
|
||||
# --- calc b
|
||||
# shapeMU
|
||||
mu_3d = np.resize(shapeMU, [n, 3]).T # 3 x n
|
||||
# expression
|
||||
exp_3d = expression
|
||||
#
|
||||
b = A.dot(mu_3d + exp_3d) + np.tile(t2d[:, np.newaxis], [1, n]) # 2 x n
|
||||
b = np.reshape(b.T, [-1, 1]) # 2n x 1
|
||||
|
||||
# --- solve
|
||||
equation_left = np.dot(pc.T, pc) + lamb * np.diagflat(1/sigma**2)
|
||||
x = np.reshape(x.T, [-1, 1])
|
||||
equation_right = np.dot(pc.T, x - b)
|
||||
|
||||
shape_para = np.dot(np.linalg.inv(equation_left), equation_right)
|
||||
|
||||
return shape_para
|
||||
|
||||
def estimate_expression(x, shapeMU, expPC, expEV, shape, s, R, t2d, lamb = 2000):
|
||||
'''
|
||||
Args:
|
||||
x: (2, n). image points (to be fitted)
|
||||
shapeMU: (3n, 1)
|
||||
expPC: (3n, n_ep)
|
||||
expEV: (n_ep, 1)
|
||||
shape: (3, n)
|
||||
s: scale
|
||||
R: (3, 3). rotation matrix
|
||||
t2d: (2,). 2d translation
|
||||
lambda: regulation coefficient
|
||||
|
||||
Returns:
|
||||
exp_para: (n_ep, 1) shape parameters(coefficients)
|
||||
'''
|
||||
x = x.copy()
|
||||
assert(shapeMU.shape[0] == expPC.shape[0])
|
||||
assert(shapeMU.shape[0] == x.shape[1]*3)
|
||||
|
||||
dof = expPC.shape[1]
|
||||
|
||||
n = x.shape[1]
|
||||
sigma = expEV
|
||||
t2d = np.array(t2d)
|
||||
P = np.array([[1, 0, 0], [0, 1, 0]], dtype = np.float32)
|
||||
A = s*P.dot(R)
|
||||
|
||||
# --- calc pc
|
||||
pc_3d = np.resize(expPC.T, [dof, n, 3])
|
||||
pc_3d = np.reshape(pc_3d, [dof*n, 3])
|
||||
pc_2d = pc_3d.dot(A.T)
|
||||
pc = np.reshape(pc_2d, [dof, -1]).T # 2n x 29
|
||||
|
||||
# --- calc b
|
||||
# shapeMU
|
||||
mu_3d = np.resize(shapeMU, [n, 3]).T # 3 x n
|
||||
# expression
|
||||
shape_3d = shape
|
||||
#
|
||||
b = A.dot(mu_3d + shape_3d) + np.tile(t2d[:, np.newaxis], [1, n]) # 2 x n
|
||||
b = np.reshape(b.T, [-1, 1]) # 2n x 1
|
||||
|
||||
# --- solve
|
||||
equation_left = np.dot(pc.T, pc) + lamb * np.diagflat(1/sigma**2)
|
||||
x = np.reshape(x.T, [-1, 1])
|
||||
equation_right = np.dot(pc.T, x - b)
|
||||
|
||||
exp_para = np.dot(np.linalg.inv(equation_left), equation_right)
|
||||
|
||||
return exp_para
|
||||
|
||||
|
||||
# ---------------- fit
|
||||
def fit_points(x, X_ind, model, n_sp, n_ep, max_iter = 4):
|
||||
'''
|
||||
Args:
|
||||
x: (n, 2) image points
|
||||
X_ind: (n,) corresponding Model vertex indices
|
||||
model: 3DMM
|
||||
max_iter: iteration
|
||||
Returns:
|
||||
sp: (n_sp, 1). shape parameters
|
||||
ep: (n_ep, 1). exp parameters
|
||||
s, R, t
|
||||
'''
|
||||
x = x.copy().T
|
||||
|
||||
#-- init
|
||||
sp = np.zeros((n_sp, 1), dtype = np.float32)
|
||||
ep = np.zeros((n_ep, 1), dtype = np.float32)
|
||||
|
||||
#-------------------- estimate
|
||||
X_ind_all = np.tile(X_ind[np.newaxis, :], [3, 1])*3
|
||||
X_ind_all[1, :] += 1
|
||||
X_ind_all[2, :] += 2
|
||||
valid_ind = X_ind_all.flatten('F')
|
||||
|
||||
shapeMU = model['shapeMU'][valid_ind, :]
|
||||
shapePC = model['shapePC'][valid_ind, :n_sp]
|
||||
expPC = model['expPC'][valid_ind, :n_ep]
|
||||
|
||||
for i in range(max_iter):
|
||||
X = shapeMU + shapePC.dot(sp) + expPC.dot(ep)
|
||||
X = np.reshape(X, [int(len(X)/3), 3]).T
|
||||
|
||||
#----- estimate pose
|
||||
P = Face.face3d.mesh.transform.estimate_affine_matrix_3d22d(X.T, x.T)
|
||||
s, R, t = Face.face3d.mesh.transform.P2sRt(P)
|
||||
rx, ry, rz = Face.face3d.mesh.transform.matrix2angle(R)
|
||||
# print('Iter:{}; estimated pose: s {}, rx {}, ry {}, rz {}, t1 {}, t2 {}'.format(i, s, rx, ry, rz, t[0], t[1]))
|
||||
|
||||
#----- estimate shape
|
||||
# expression
|
||||
shape = shapePC.dot(sp)
|
||||
shape = np.reshape(shape, [int(len(shape)/3), 3]).T
|
||||
ep = estimate_expression(x, shapeMU, expPC, model['expEV'][:n_ep,:], shape, s, R, t[:2], lamb = 20)
|
||||
|
||||
# shape
|
||||
expression = expPC.dot(ep)
|
||||
expression = np.reshape(expression, [int(len(expression)/3), 3]).T
|
||||
sp = estimate_shape(x, shapeMU, shapePC, model['shapeEV'][:n_sp,:], expression, s, R, t[:2], lamb = 40)
|
||||
|
||||
return sp, ep, s, R, t
|
||||
|
||||
|
||||
# ---------------- fitting process
|
||||
def fit_points_for_show(x, X_ind, model, n_sp, n_ep, max_iter = 4):
|
||||
'''
|
||||
Args:
|
||||
x: (n, 2) image points
|
||||
X_ind: (n,) corresponding Model vertex indices
|
||||
model: 3DMM
|
||||
max_iter: iteration
|
||||
Returns:
|
||||
sp: (n_sp, 1). shape parameters
|
||||
ep: (n_ep, 1). exp parameters
|
||||
s, R, t
|
||||
'''
|
||||
x = x.copy().T
|
||||
|
||||
#-- init
|
||||
sp = np.zeros((n_sp, 1), dtype = np.float32)
|
||||
ep = np.zeros((n_ep, 1), dtype = np.float32)
|
||||
|
||||
#-------------------- estimate
|
||||
X_ind_all = np.tile(X_ind[np.newaxis, :], [3, 1])*3
|
||||
X_ind_all[1, :] += 1
|
||||
X_ind_all[2, :] += 2
|
||||
valid_ind = X_ind_all.flatten('F')
|
||||
|
||||
shapeMU = model['shapeMU'][valid_ind, :]
|
||||
shapePC = model['shapePC'][valid_ind, :n_sp]
|
||||
expPC = model['expPC'][valid_ind, :n_ep]
|
||||
|
||||
s = 4e-04
|
||||
R = Face.face3d.mesh.transform.angle2matrix([0, 0, 0])
|
||||
t = [0, 0, 0]
|
||||
lsp = []; lep = []; ls = []; lR = []; lt = []
|
||||
for i in range(max_iter):
|
||||
X = shapeMU + shapePC.dot(sp) + expPC.dot(ep)
|
||||
X = np.reshape(X, [int(len(X)/3), 3]).T
|
||||
lsp.append(sp); lep.append(ep); ls.append(s), lR.append(R), lt.append(t)
|
||||
|
||||
#----- estimate pose
|
||||
P = Face.face3d.mesh.transform.estimate_affine_matrix_3d22d(X.T, x.T)
|
||||
s, R, t = Face.face3d.mesh.transform.P2sRt(P)
|
||||
lsp.append(sp); lep.append(ep); ls.append(s), lR.append(R), lt.append(t)
|
||||
|
||||
#----- estimate shape
|
||||
# expression
|
||||
shape = shapePC.dot(sp)
|
||||
shape = np.reshape(shape, [int(len(shape)/3), 3]).T
|
||||
ep = estimate_expression(x, shapeMU, expPC, model['expEV'][:n_ep,:], shape, s, R, t[:2], lamb = 20)
|
||||
lsp.append(sp); lep.append(ep); ls.append(s), lR.append(R), lt.append(t)
|
||||
|
||||
# shape
|
||||
expression = expPC.dot(ep)
|
||||
expression = np.reshape(expression, [int(len(expression)/3), 3]).T
|
||||
sp = estimate_shape(x, shapeMU, shapePC, model['shapeEV'][:n_sp,:], expression, s, R, t[:2], lamb = 40)
|
||||
|
||||
# print('ls', ls)
|
||||
# print('lR', lR)
|
||||
return np.array(lsp), np.array(lep), np.array(ls), np.array(lR), np.array(lt)
|
||||
@@ -0,0 +1,110 @@
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import numpy as np
|
||||
import scipy.io as sio
|
||||
|
||||
### --------------------------------- load BFM data
|
||||
def load_BFM(model_path):
|
||||
''' load BFM 3DMM model
|
||||
Args:
|
||||
model_path: path to BFM model.
|
||||
Returns:
|
||||
model: (nver = 53215, ntri = 105840). nver: number of vertices. ntri: number of triangles.
|
||||
'shapeMU': [3*nver, 1]
|
||||
'shapePC': [3*nver, 199]
|
||||
'shapeEV': [199, 1]
|
||||
'expMU': [3*nver, 1]
|
||||
'expPC': [3*nver, 29]
|
||||
'expEV': [29, 1]
|
||||
'texMU': [3*nver, 1]
|
||||
'texPC': [3*nver, 199]
|
||||
'texEV': [199, 1]
|
||||
'tri': [ntri, 3] (start from 1, should sub 1 in python and c++)
|
||||
'tri_mouth': [114, 3] (start from 1, as a supplement to mouth triangles)
|
||||
'kpt_ind': [68,] (start from 1)
|
||||
PS:
|
||||
You can change codes according to your own saved data.
|
||||
Just make sure the model has corresponding attributes.
|
||||
'''
|
||||
C = sio.loadmat(model_path)
|
||||
model = C['model']
|
||||
model = model[0,0]
|
||||
|
||||
# change dtype from double(np.float64) to np.float32,
|
||||
# since big matrix process(espetially matrix dot) is too slow in python.
|
||||
model['shapeMU'] = (model['shapeMU'] + model['expMU']).astype(np.float32)
|
||||
model['shapePC'] = model['shapePC'].astype(np.float32)
|
||||
model['shapeEV'] = model['shapeEV'].astype(np.float32)
|
||||
model['expEV'] = model['expEV'].astype(np.float32)
|
||||
model['expPC'] = model['expPC'].astype(np.float32)
|
||||
|
||||
# matlab start with 1. change to 0 in python.
|
||||
model['tri'] = model['tri'].T.copy(order = 'C').astype(np.int32) - 1
|
||||
model['tri_mouth'] = model['tri_mouth'].T.copy(order = 'C').astype(np.int32) - 1
|
||||
|
||||
# kpt ind
|
||||
model['kpt_ind'] = (np.squeeze(model['kpt_ind']) - 1).astype(np.int32)
|
||||
|
||||
return model
|
||||
|
||||
def load_BFM_info(path = 'BFM_info.mat'):
|
||||
''' load 3DMM model extra information
|
||||
Args:
|
||||
path: path to BFM info.
|
||||
Returns:
|
||||
model_info:
|
||||
'symlist': 2 x 26720
|
||||
'symlist_tri': 2 x 52937
|
||||
'segbin': 4 x n (0: nose, 1: eye, 2: mouth, 3: cheek)
|
||||
'segbin_tri': 4 x ntri
|
||||
'face_contour': 1 x 28
|
||||
'face_contour_line': 1 x 512
|
||||
'face_contour_front': 1 x 28
|
||||
'face_contour_front_line': 1 x 512
|
||||
'nose_hole': 1 x 142
|
||||
'nose_hole_right': 1 x 71
|
||||
'nose_hole_left': 1 x 71
|
||||
'parallel': 17 x 1 cell
|
||||
'parallel_face_contour': 28 x 1 cell
|
||||
'uv_coords': n x 2
|
||||
'''
|
||||
C = sio.loadmat(path)
|
||||
model_info = C['model_info']
|
||||
model_info = model_info[0,0]
|
||||
return model_info
|
||||
|
||||
def load_uv_coords(path = 'BFM_UV.mat'):
|
||||
''' load uv coords of BFM
|
||||
Args:
|
||||
path: path to data.
|
||||
Returns:
|
||||
uv_coords: [nver, 2]. range: 0-1
|
||||
'''
|
||||
C = sio.loadmat(path)
|
||||
uv_coords = C['UV'].copy(order = 'C')
|
||||
return uv_coords
|
||||
|
||||
def load_pncc_code(path = 'pncc_code.mat'):
|
||||
''' load pncc code of BFM
|
||||
PNCC code: Defined in 'Face Alignment Across Large Poses: A 3D Solution Xiangyu'
|
||||
download at http://www.cbsr.ia.ac.cn/users/xiangyuzhu/projects/3DDFA/main.htm.
|
||||
Args:
|
||||
path: path to data.
|
||||
Returns:
|
||||
pncc_code: [nver, 3]
|
||||
'''
|
||||
C = sio.loadmat(path)
|
||||
pncc_code = C['vertex_code'].T
|
||||
return pncc_code
|
||||
|
||||
##
|
||||
def get_organ_ind(model_info):
|
||||
''' get nose, eye, mouth index
|
||||
'''
|
||||
valid_bin = model_info['segbin'].astype(bool)
|
||||
organ_ind = np.nonzero(valid_bin[0,:])[0]
|
||||
for i in range(1, valid_bin.shape[0] - 1):
|
||||
organ_ind = np.union1d(organ_ind, np.nonzero(valid_bin[i,:])[0])
|
||||
return organ_ind.astype(np.int32)
|
||||
@@ -0,0 +1,141 @@
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import numpy as np
|
||||
from . import fit
|
||||
from . import load
|
||||
|
||||
class MorphabelModel(object):
|
||||
"""docstring for MorphabelModel
|
||||
model: nver: number of vertices. ntri: number of triangles. *: must have. ~: can generate ones array for place holder.
|
||||
'shapeMU': [3*nver, 1]. *
|
||||
'shapePC': [3*nver, n_shape_para]. *
|
||||
'shapeEV': [n_shape_para, 1]. ~
|
||||
'expMU': [3*nver, 1]. ~
|
||||
'expPC': [3*nver, n_exp_para]. ~
|
||||
'expEV': [n_exp_para, 1]. ~
|
||||
'texMU': [3*nver, 1]. ~
|
||||
'texPC': [3*nver, n_tex_para]. ~
|
||||
'texEV': [n_tex_para, 1]. ~
|
||||
'tri': [ntri, 3] (start from 1, should sub 1 in python and c++). *
|
||||
'tri_mouth': [114, 3] (start from 1, as a supplement to mouth triangles). ~
|
||||
'kpt_ind': [68,] (start from 1). ~
|
||||
"""
|
||||
def __init__(self, model_path, model_type = 'BFM'):
|
||||
super( MorphabelModel, self).__init__()
|
||||
if model_type=='BFM':
|
||||
self.model = load.load_BFM(model_path)
|
||||
else:
|
||||
print('sorry, not support other 3DMM model now')
|
||||
exit()
|
||||
|
||||
# fixed attributes
|
||||
self.nver = self.model['shapePC'].shape[0]/3
|
||||
self.ntri = self.model['tri'].shape[0]
|
||||
self.n_shape_para = self.model['shapePC'].shape[1]
|
||||
self.n_exp_para = self.model['expPC'].shape[1]
|
||||
self.n_tex_para = self.model['texMU'].shape[1]
|
||||
|
||||
self.kpt_ind = self.model['kpt_ind']
|
||||
self.triangles = self.model['tri']
|
||||
self.full_triangles = np.vstack((self.model['tri'], self.model['tri_mouth']))
|
||||
|
||||
# ------------------------------------- shape: represented with mesh(vertices & triangles(fixed))
|
||||
def get_shape_para(self, type = 'random'):
|
||||
if type == 'zero':
|
||||
sp = np.random.zeros((self.n_shape_para, 1))
|
||||
elif type == 'random':
|
||||
sp = np.random.rand(self.n_shape_para, 1)*1e04
|
||||
return sp
|
||||
|
||||
def get_exp_para(self, type = 'random'):
|
||||
if type == 'zero':
|
||||
ep = np.zeros((self.n_exp_para, 1))
|
||||
elif type == 'random':
|
||||
ep = -1.5 + 3*np.random.random([self.n_exp_para, 1])
|
||||
ep[6:, 0] = 0
|
||||
|
||||
return ep
|
||||
|
||||
def generate_vertices(self, shape_para, exp_para):
|
||||
'''
|
||||
Args:
|
||||
shape_para: (n_shape_para, 1)
|
||||
exp_para: (n_exp_para, 1)
|
||||
Returns:
|
||||
vertices: (nver, 3)
|
||||
'''
|
||||
vertices = self.model['shapeMU'] + self.model['shapePC'].dot(shape_para) + self.model['expPC'].dot(exp_para)
|
||||
vertices = np.reshape(vertices, [int(3), int(len(vertices)/3)], 'F').T
|
||||
|
||||
return vertices
|
||||
|
||||
# -------------------------------------- texture: here represented with rgb value(colors) in vertices.
|
||||
def get_tex_para(self, type = 'random'):
|
||||
if type == 'zero':
|
||||
tp = np.zeros((self.n_tex_para, 1))
|
||||
elif type == 'random':
|
||||
tp = np.random.rand(self.n_tex_para, 1)
|
||||
return tp
|
||||
|
||||
def generate_colors(self, tex_para):
|
||||
'''
|
||||
Args:
|
||||
tex_para: (n_tex_para, 1)
|
||||
Returns:
|
||||
colors: (nver, 3)
|
||||
'''
|
||||
colors = self.model['texMU'] + self.model['texPC'].dot(tex_para*self.model['texEV'])
|
||||
colors = np.reshape(colors, [int(3), int(len(colors)/3)], 'F').T/255.
|
||||
|
||||
return colors
|
||||
|
||||
|
||||
# ------------------------------------------- transformation
|
||||
# ------------- transform
|
||||
def rotate(self, vertices, angles):
|
||||
''' rotate face
|
||||
Args:
|
||||
vertices: [nver, 3]
|
||||
angles: [3] x, y, z rotation angle(degree)
|
||||
x: pitch. positive for looking down
|
||||
y: yaw. positive for looking left
|
||||
z: roll. positive for tilting head right
|
||||
Returns:
|
||||
vertices: rotated vertices
|
||||
'''
|
||||
return Face.face3d.mesh.transform.rotate(vertices, angles)
|
||||
|
||||
def transform(self, vertices, s, angles, t3d):
|
||||
R = Face.face3d.mesh.transform.angle2matrix(angles)
|
||||
return Face.face3d.mesh.transform.similarity_transform(vertices, s, R, t3d)
|
||||
|
||||
def transform_3ddfa(self, vertices, s, angles, t3d): # only used for processing 300W_LP data
|
||||
R = Face.face3d.mesh.transform.angle2matrix_3ddfa(angles)
|
||||
return Face.face3d.mesh.transform.similarity_transform(vertices, s, R, t3d)
|
||||
|
||||
# --------------------------------------------------- fitting
|
||||
def fit(self, x, X_ind, max_iter = 4, isShow = False):
|
||||
''' fit 3dmm & pose parameters
|
||||
Args:
|
||||
x: (n, 2) image points
|
||||
X_ind: (n,) corresponding Model vertex indices
|
||||
max_iter: iteration
|
||||
isShow: whether to reserve middle results for show
|
||||
Returns:
|
||||
fitted_sp: (n_sp, 1). shape parameters
|
||||
fitted_ep: (n_ep, 1). exp parameters
|
||||
s, angles, t
|
||||
'''
|
||||
if isShow:
|
||||
fitted_sp, fitted_ep, s, R, t = fit.fit_points_for_show(x, X_ind, self.model, n_sp = self.n_shape_para, n_ep = self.n_exp_para, max_iter = max_iter)
|
||||
angles = np.zeros((R.shape[0], 3))
|
||||
for i in range(R.shape[0]):
|
||||
angles[i] = Face.face3d.mesh.transform.matrix2angle(R[i])
|
||||
else:
|
||||
fitted_sp, fitted_ep, s, R, t = fit.fit_points(x, X_ind, self.model, n_sp = self.n_shape_para, n_ep = self.n_exp_para, max_iter = max_iter)
|
||||
angles = Face.face3d.mesh.transform.matrix2angle(R)
|
||||
return fitted_sp, fitted_ep, s, angles, t
|
||||
|
||||
|
||||
@@ -0,0 +1,96 @@
|
||||
import os
|
||||
import cv2
|
||||
import glob
|
||||
import numpy as np
|
||||
from core.utils import landmark_processor
|
||||
|
||||
from core.face_enhance.face_gan_pt import FaceGAN
|
||||
from time import time
|
||||
|
||||
class FaceEnhancement(object):
|
||||
def __init__(self, size=512, gpu_id=None):
|
||||
self.facegan = FaceGAN(size, gpu_id)
|
||||
self.size = size
|
||||
self.threshold = 0.9
|
||||
|
||||
# the mask for pasting restored faces back
|
||||
self.mask = np.zeros((512, 512), np.float32)
|
||||
cv2.rectangle(self.mask, (26, 26), (486, 486), (1, 1, 1), -1, cv2.LINE_AA)
|
||||
self.mask = cv2.GaussianBlur(self.mask, (101, 101), 11)
|
||||
self.mask = cv2.GaussianBlur(self.mask, (101, 101), 11)
|
||||
|
||||
self.kernel = np.array((
|
||||
[0.0625, 0.125, 0.0625],
|
||||
[0.125, 0.25, 0.125],
|
||||
[0.0625, 0.125, 0.0625]), dtype="float32")
|
||||
|
||||
def process(self, img, landmarks1k):
|
||||
|
||||
assert len(landmarks1k) == 1000
|
||||
|
||||
image_to_face_mat = landmark_processor.get_transform_mat_face_restore(landmarks1k, self.size)
|
||||
tfm_inv = cv2.invertAffineTransform(image_to_face_mat)
|
||||
|
||||
height, width = img.shape[:2]
|
||||
full_mask = np.zeros((height, width), dtype=np.float32)
|
||||
full_img = np.zeros(img.shape, dtype=np.uint8)
|
||||
|
||||
fh, fw = (landmarks1k[0][1]-landmarks1k[154][1]), (landmarks1k[95][0]-landmarks1k[215][0])
|
||||
|
||||
of = cv2.warpAffine(img, image_to_face_mat, (self.size, self.size), flags=cv2.INTER_LANCZOS4, borderMode=cv2.BORDER_CONSTANT, borderValue=[0, 0, 0])
|
||||
|
||||
# enhance the face
|
||||
t0 = time()
|
||||
|
||||
ef = self.facegan.process(of)
|
||||
# print('facegan process costs:', time() - t0)
|
||||
tmp_mask = self.mask
|
||||
tmp_mask = cv2.resize(tmp_mask, ef.shape[:2])
|
||||
tmp_mask = cv2.warpAffine(tmp_mask, tfm_inv, (width, height), flags=3)
|
||||
|
||||
if min(fh, fw)<100: # gaussian filter for small faces
|
||||
ef = cv2.filter2D(ef, -1, self.kernel)
|
||||
|
||||
# tmp_img = cv2.warpAffine(ef, tfm_inv, (width, height), flags=3)
|
||||
tmp_img = cv2.warpAffine(ef, tfm_inv, (width, height), dst=img.copy(), borderMode=cv2.BORDER_TRANSPARENT)
|
||||
|
||||
# cv2.imshow("tmp_img: ", tmp_img)
|
||||
|
||||
mask = tmp_mask - full_mask
|
||||
full_mask[np.where(mask>0)] = tmp_mask[np.where(mask>0)]
|
||||
full_img[np.where(mask>0)] = tmp_img[np.where(mask>0)]
|
||||
|
||||
full_mask = full_mask[:, :, np.newaxis]
|
||||
img = cv2.convertScaleAbs(img*(1-full_mask) + full_img*full_mask)
|
||||
|
||||
return img
|
||||
|
||||
if __name__=='__main__':
|
||||
|
||||
indir = '/media/DATA_4T/zao_data/tencent_con_0701/ref_c/res_contrast_blur13_notps_yunfu_res'
|
||||
outdir = '/media/DATA_4T/zao_data/tencent_con_0701/ref_c/res_contrast_blur13_notps_yunfu_res_outs2'
|
||||
os.makedirs(outdir, exist_ok=True)
|
||||
|
||||
faceenhancer = FaceEnhancement(base_dir="/", size=512, model="GPEN-512", channel_multiplier=2)
|
||||
|
||||
files = sorted(glob.glob(os.path.join(indir, '*.*g')))
|
||||
for n, file in enumerate(files[:]):
|
||||
filename = os.path.basename(file)
|
||||
txtname = file.replace(".jpg", "_landmark1k.txt")
|
||||
|
||||
im = cv2.imread(file, cv2.IMREAD_COLOR) # BGR
|
||||
print(txtname)
|
||||
landmark = np.loadtxt(txtname)
|
||||
if not isinstance(im, np.ndarray): print(filename, 'error'); continue
|
||||
|
||||
start = time()
|
||||
|
||||
img = faceenhancer.process(im, landmark)
|
||||
|
||||
end = time()
|
||||
|
||||
print("Time cost: {:.4f}".format(end - start))
|
||||
|
||||
cv2.imwrite(os.path.join(outdir, '.'.join(filename.split('.')[:-1])+'_2.jpg'), img)
|
||||
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
'''
|
||||
@paper: GAN Prior Embedded Network for Blind Face Restoration in the Wild (CVPR2021)
|
||||
@author: yangxy (yangtao9009@gmail.com)
|
||||
'''
|
||||
import torch
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
class FaceGAN(object):
|
||||
def __init__(self, size=512, gpu_id=0):
|
||||
# self.mfile = os.path.join(base_dir, model+'.pth')
|
||||
self.n_mlp = 8
|
||||
self.resolution = size
|
||||
self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||
|
||||
self.load_model()
|
||||
|
||||
def load_model(self):
|
||||
self.face_gan_model = os.path.join('weights', "face_enhance_0630.pt")
|
||||
self.model = torch.jit.load(self.face_gan_model, torch.device('cpu')).to(self.device)
|
||||
# self.model = torch.jit.load(self.face_gan_model).to(self.device)
|
||||
# self.model = torch.load(self.face_gan_model,map_location=self.device)
|
||||
|
||||
self.model.eval()
|
||||
|
||||
def process_o(self, img):
|
||||
img = cv2.resize(img, (self.resolution, self.resolution))
|
||||
img_t = self.img2tensor(img)
|
||||
|
||||
with torch.no_grad():
|
||||
out, __ = self.model(img_t)
|
||||
|
||||
out = self.tensor2img(out)
|
||||
|
||||
return out
|
||||
|
||||
def process(self, img):
|
||||
img = cv2.resize(img, (self.resolution, self.resolution))
|
||||
img_t = self.img2tensor(img)
|
||||
|
||||
with torch.no_grad():
|
||||
out = self.forward(img_t)
|
||||
|
||||
out = self.tensor2img(out)
|
||||
|
||||
return out
|
||||
|
||||
def forward(self, img_t):
|
||||
with torch.no_grad():
|
||||
out = self.model(img_t)
|
||||
|
||||
return out
|
||||
|
||||
def img2tensor(self, img):
|
||||
img_t = (torch.from_numpy(img).to(self.device)/255. - 0.5) / 0.5
|
||||
img_t = img_t.permute(2, 0, 1).unsqueeze(0).flip(1) # BGR->RGB
|
||||
return img_t
|
||||
|
||||
def tensor2img(self, image_tensor, pmax=255.0, imtype=np.uint8):
|
||||
image_tensor = image_tensor * 0.5 + 0.5
|
||||
image_tensor = image_tensor.squeeze(0).permute(1, 2, 0).flip(2) # RGB->BGR
|
||||
image_numpy = np.clip(image_tensor.float().cpu().numpy(), 0, 1) * pmax
|
||||
|
||||
return image_numpy.astype(imtype)
|
||||
@@ -0,0 +1,102 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
# @File : setup.py
|
||||
# @Time : 2020/1/15
|
||||
# @Author : yangchaojie (yangchaojie@immomo.com)
|
||||
|
||||
import os
|
||||
import sys
|
||||
import shutil
|
||||
import numpy
|
||||
import tempfile
|
||||
|
||||
from setuptools import setup
|
||||
from setuptools.extension import Extension
|
||||
|
||||
from Cython.Build import cythonize
|
||||
from Cython.Distutils import build_ext
|
||||
|
||||
import platform
|
||||
|
||||
|
||||
def get_root_path(root):
|
||||
if os.path.dirname(root) in ['', '.']:
|
||||
return os.path.basename(root)
|
||||
else:
|
||||
return get_root_path(os.path.dirname(root))
|
||||
|
||||
|
||||
def copy_file(src, dest):
|
||||
if os.path.exists(dest):
|
||||
return
|
||||
|
||||
if not os.path.exists(os.path.dirname(dest)):
|
||||
os.makedirs(os.path.dirname(dest))
|
||||
if os.path.isdir(src):
|
||||
shutil.copytree(src, dest)
|
||||
else:
|
||||
shutil.copyfile(src, dest)
|
||||
|
||||
|
||||
def touch_init_file():
|
||||
init_file_name = os.path.join(tempfile.mkdtemp(), '__init__.py')
|
||||
with open(init_file_name, 'w'):
|
||||
pass
|
||||
return init_file_name
|
||||
|
||||
|
||||
|
||||
|
||||
def compose_extensions(root='.'):
|
||||
for file_ in os.listdir(root):
|
||||
abs_file = os.path.join(root, file_)
|
||||
|
||||
if os.path.isfile(abs_file):
|
||||
if abs_file.endswith('.py'):
|
||||
extensions.append(Extension(get_root_path(abs_file) + '.*', [abs_file]))
|
||||
elif abs_file.endswith('.c') or abs_file.endswith('.pyc'):
|
||||
continue
|
||||
else:
|
||||
copy_file(abs_file, os.path.join(build_root_dir, abs_file))
|
||||
if abs_file.endswith('__init__.py'):
|
||||
copy_file(init_file, os.path.join(build_root_dir, abs_file))
|
||||
|
||||
else:
|
||||
if os.path.basename(abs_file) in ignore_folders :
|
||||
continue
|
||||
if os.path.basename(abs_file) in conf_folders:
|
||||
copy_file(abs_file, os.path.join(build_root_dir, abs_file))
|
||||
compose_extensions(abs_file)
|
||||
|
||||
|
||||
|
||||
# if __name__ == '__main__':
|
||||
build_root_dir = 'build/lib.' + platform.system().lower() + '-' + platform.machine() + '-' + str(
|
||||
sys.version_info.major) + '.' + str(sys.version_info.minor)
|
||||
|
||||
print(build_root_dir)
|
||||
|
||||
extensions = []
|
||||
ignore_folders = ['build', 'new_ref_zao_color_0818', 'ref_hair_online_0703_local', 'ref_分好类别', '.git']
|
||||
conf_folders = ['conf']
|
||||
|
||||
|
||||
init_file = touch_init_file()
|
||||
print(init_file)
|
||||
|
||||
|
||||
compose_extensions()
|
||||
os.remove(init_file)
|
||||
|
||||
setup(
|
||||
name='moxie_hairstyle',
|
||||
version='1.0',
|
||||
ext_modules=cythonize(
|
||||
extensions,
|
||||
nthreads=16,
|
||||
compiler_directives=dict(always_allow_keywords=True),
|
||||
include_path=[numpy.get_include()]),
|
||||
cmdclass=dict(build_ext=build_ext))
|
||||
|
||||
# python setup.py build_ext
|
||||
@@ -0,0 +1,186 @@
|
||||
import cv2
|
||||
|
||||
|
||||
class SingleFaceQualityInfo():
|
||||
def __int__(self, gpu_id):
|
||||
print('Model 3DDFA init success')
|
||||
|
||||
def getSpotProportion(self, grayImage):
|
||||
Proportion = 0
|
||||
hist = cv2.calcHist([grayImage], [0], None, [256], [0, 256])
|
||||
for i in range(253, 256):
|
||||
Proportion = Proportion + hist[i, 0]
|
||||
return Proportion
|
||||
|
||||
def GetSpotRatio(self, wrap_img, wrap_face_rect):
|
||||
scale = 0.9
|
||||
|
||||
x = max(0, int(wrap_face_rect[0]))
|
||||
y = max(0, int(wrap_face_rect[1]))
|
||||
w = max(0, int(wrap_face_rect[2]-wrap_face_rect[0]))
|
||||
h = max(0, int(wrap_face_rect[3]-wrap_face_rect[1]))
|
||||
|
||||
height, width, _ = wrap_img.shape
|
||||
if (x + w) > width:
|
||||
w = width - x
|
||||
if (y + h) > height:
|
||||
h = height - y
|
||||
p_cen = (x + w/2, y+h/2)
|
||||
new_w = w*scale
|
||||
new_h = h*scale
|
||||
na_box = [int(p_cen[0]-new_w/2), int(p_cen[1]) - int(new_h/2), int(new_w), int(new_h)]
|
||||
imgCropGray = cv2.cvtColor(wrap_img[na_box[1]:na_box[1] + na_box[3], na_box[0]:na_box[0] + na_box[2]],cv2.COLOR_RGB2GRAY)
|
||||
|
||||
Spot_Ratio = self.getSpotProportion(imgCropGray) / (imgCropGray.shape[0] * imgCropGray.shape[1])
|
||||
return Spot_Ratio
|
||||
|
||||
def BoxClarityAndBrightValue(self, wrap_img, wrap_face_rect):
|
||||
x = max(0, int(wrap_face_rect[0]))
|
||||
y = max(0, int(wrap_face_rect[1]))
|
||||
w = max(0, int(wrap_face_rect[2]-wrap_face_rect[0]))
|
||||
h = max(0, int(wrap_face_rect[3]-wrap_face_rect[1]))
|
||||
# cv2.imshow('face', wrap_img[y:y+h,x:x+w])
|
||||
# cv2.waitKey()
|
||||
height, width, _ = wrap_img.shape
|
||||
if (x + w) > width:
|
||||
w = width - x
|
||||
if (y + h) > height:
|
||||
h = height - y
|
||||
box_crop = [int(x), int(y), int(w), int(h)]
|
||||
imgCropGray = cv2.cvtColor(wrap_img[box_crop[1]:box_crop[1] + box_crop[3], box_crop[0]:box_crop[0] + box_crop[2]],cv2.COLOR_RGB2GRAY)
|
||||
left_ = [0,0,int(w/2.0),int(h)]
|
||||
right_ = [int(w/2.0),0,int(w/2.0),int(h)]
|
||||
imgR = imgCropGray[right_[1]:right_[1] + right_[3], right_[0]:right_[0] + right_[2]]
|
||||
imgL = imgCropGray[left_[1]:left_[1] + left_[3], left_[0]:left_[0] + left_[2]]
|
||||
mean_R, var_R = cv2.meanStdDev(imgR)
|
||||
mean_L, var_L = cv2.meanStdDev(imgL)
|
||||
|
||||
bV_R = mean_R[0]
|
||||
bV_L = mean_L[0]
|
||||
|
||||
imageCropSobel = cv2.Laplacian(imgCropGray, cv2.CV_64F, 3)
|
||||
mean, var = cv2.meanStdDev(imageCropSobel)
|
||||
clarityValue = var[0]
|
||||
result = [bV_R, bV_L, clarityValue]
|
||||
return result
|
||||
|
||||
def qualityTest(self, wrap_img, wrap_face_rect, single_eulers):
|
||||
single_face_quality_info = {
|
||||
'quality_score': -1,
|
||||
'face_quality_flag': -1,
|
||||
'SpotRatio': 0,
|
||||
'Brightness': 0,
|
||||
'Clarity': 0,
|
||||
'SpotScore': 0,
|
||||
'BrightnessScore': 0,
|
||||
'ClarityScore': 0
|
||||
}
|
||||
|
||||
if wrap_img is None:
|
||||
single_face_quality_info['face_quality_flag'] = 0
|
||||
qp_img_size = wrap_img.shape[1]
|
||||
Spot_Ratio = self.GetSpotRatio(wrap_img, wrap_face_rect)
|
||||
result = self.BoxClarityAndBrightValue(wrap_img, wrap_face_rect)
|
||||
bv_R, bv_L, clarityValue = result
|
||||
single_face_quality_info['SpotRatio'] = Spot_Ratio
|
||||
single_face_quality_info['Brightness'] = min(bv_L, bv_R)
|
||||
single_face_quality_info['Clarity'] = clarityValue
|
||||
if qp_img_size == 100:
|
||||
clarityValue_thresholdDown = 35
|
||||
clarityValue_thresholdUp = 300
|
||||
else:
|
||||
clarityValue_thresholdDown = 24
|
||||
clarityValue_thresholdUp = 300
|
||||
Bright_area_thre = 0.06
|
||||
brightnessValue_thresholddown = 40
|
||||
brightnessValue_thresholdUp = 230
|
||||
|
||||
whitespot = (Spot_Ratio > Bright_area_thre)
|
||||
brightness =(bv_R > brightnessValue_thresholddown and bv_R < brightnessValue_thresholdUp) and (bv_L > brightnessValue_thresholddown and bv_L < brightnessValue_thresholdUp)
|
||||
clarity = clarityValue < clarityValue_thresholdUp and clarityValue > clarityValue_thresholdDown
|
||||
if whitespot:
|
||||
single_face_quality_info['face_qality_flag'] = 2
|
||||
if brightness:
|
||||
single_face_quality_info['face_qality_flag'] = 4
|
||||
if not clarity:
|
||||
single_face_quality_info['face_qality_flag'] = 3
|
||||
if brightness and clarity and not whitespot:
|
||||
single_face_quality_info['face_qality_flag'] = 1
|
||||
|
||||
|
||||
if single_face_quality_info['face_qality_flag'] == 1:
|
||||
PerfectSpotRatio = 0.02
|
||||
if qp_img_size == 100:
|
||||
PerfectClarity = 130
|
||||
else:
|
||||
PerfectClarity = 60
|
||||
PerfectBrightness_down = 90
|
||||
PerfectBrightness_up = 200
|
||||
PerfectBrightness_standard = PerfectBrightness_down - brightnessValue_thresholddown
|
||||
PerfectClarity_standard = PerfectClarity - clarityValue_thresholdDown
|
||||
PerfectSpot_standard = Bright_area_thre - PerfectSpotRatio
|
||||
if single_face_quality_info['SpotRatio'] <= PerfectSpotRatio:
|
||||
single_face_quality_info['SpotScore'] = 1
|
||||
else:
|
||||
single_face_quality_info['SpotRatio'] = 1 - (single_face_quality_info['SpotRatio'] - PerfectSpotRatio) / PerfectSpot_standard
|
||||
if single_face_quality_info['Clarity'] >= PerfectClarity:
|
||||
single_face_quality_info['ClarityScore'] = 1
|
||||
else:
|
||||
single_face_quality_info['ClarityScore'] = 1 - (PerfectClarity - single_face_quality_info['Clarity']) / PerfectClarity_standard
|
||||
if single_face_quality_info['Brightness'] >= PerfectBrightness_down and single_face_quality_info['Brightness']<= PerfectBrightness_up:
|
||||
single_face_quality_info['BrightnessScore'] = 1
|
||||
elif single_face_quality_info['Brightness'] < PerfectBrightness_down:
|
||||
single_face_quality_info['BrightnessScore'] = 1 - (PerfectBrightness_down - single_face_quality_info['Brightness']) / PerfectBrightness_standard
|
||||
elif single_face_quality_info['Brightness'] > PerfectBrightness_up:
|
||||
single_face_quality_info['BrightnessScore'] = 1 - (single_face_quality_info['Brightness'] - PerfectBrightness_up)/PerfectBrightness_standard
|
||||
else:
|
||||
print("Wrong Brightness :%f",single_face_quality_info['Brightness'])
|
||||
else:
|
||||
single_face_quality_info['ClarityScore'] = 0
|
||||
single_face_quality_info['BrightnessScore'] = 0
|
||||
single_face_quality_info['SpotScore'] = 0
|
||||
|
||||
if len(single_eulers) > 2:
|
||||
euler_pitch_perfect = 10
|
||||
euler_yaw_perfect = 10
|
||||
|
||||
euler_pitch_max = 25
|
||||
euler_yaw_max = 25
|
||||
euler_roll_max = 30
|
||||
|
||||
image_ClarityScore = float(single_face_quality_info['ClarityScore'])
|
||||
image_BrightnessScore = single_face_quality_info['BrightnessScore']
|
||||
image_SpotScore = single_face_quality_info['SpotScore']
|
||||
|
||||
euler_pitch, euler_yaw, euler_roll = single_eulers
|
||||
|
||||
if abs(euler_pitch) > euler_pitch_perfect and abs(euler_pitch) < euler_pitch_max and single_face_quality_info['face_qality_flag'] == 1:
|
||||
image_PitchScore = 1 - (abs(euler_pitch) - euler_pitch_perfect) / (euler_pitch_max - euler_pitch_perfect)
|
||||
elif abs(euler_pitch) < euler_pitch_perfect and single_face_quality_info['face_qality_flag'] == 1:
|
||||
image_PitchScore = 1.0
|
||||
else:
|
||||
image_PitchScore = 0
|
||||
if abs(euler_yaw) > euler_yaw_perfect and abs(euler_yaw) < euler_yaw_max and single_face_quality_info['face_qality_flag']==1:
|
||||
image_YawScore = 1 - (abs(euler_yaw) - euler_yaw_perfect) / (euler_yaw_max - euler_yaw_perfect)
|
||||
elif abs(euler_yaw) < euler_yaw_perfect and single_face_quality_info['face_qality_flag']==1:
|
||||
image_YawScore = 1.0
|
||||
else:
|
||||
image_YawScore = 0
|
||||
|
||||
if image_PitchScore < 0.3 and image_YawScore < 0.3:
|
||||
image_spot_weight = 0
|
||||
image_clarity_weight = 0.1
|
||||
image_brightness_weight = 0.1
|
||||
euler_pitch_weight = 0.4
|
||||
euler_yaw_weight = 0.4
|
||||
else:
|
||||
image_spot_weight = 0.05
|
||||
image_clarity_weight = 0.35
|
||||
image_brightness_weight = 0.1
|
||||
euler_pitch_weight = 0.25
|
||||
euler_yaw_weight = 0.25
|
||||
if abs(euler_pitch) < euler_pitch_max and abs(euler_yaw) < euler_yaw_max and abs(euler_roll) < euler_roll_max and single_face_quality_info['face_qality_flag'] == 1:
|
||||
single_face_quality_info['quality_score'] = (image_spot_weight*image_SpotScore + image_clarity_weight * image_ClarityScore + image_brightness_weight * image_BrightnessScore + euler_pitch_weight * image_PitchScore + euler_yaw_weight* image_YawScore)
|
||||
else:
|
||||
single_face_quality_info['quality_score'] = 0
|
||||
return single_face_quality_info
|
||||
@@ -0,0 +1,75 @@
|
||||
import os
|
||||
import pickle
|
||||
|
||||
import torch
|
||||
from core.faceseg.u2net import U2NET
|
||||
from core.utils import landmark_processor
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
class FaceSeg:
|
||||
def __init__(self, gpu_id = 0):
|
||||
model = U2NET(in_ch=4, out_ch=1)
|
||||
weights = torch.load('weights/faceseg_20210927_01.pth', map_location='cpu')
|
||||
model_dict = model.state_dict()
|
||||
pretrained_dict = {}
|
||||
for ix, (k, v) in enumerate(model_dict.items()):
|
||||
if k in weights and weights[k].data.shape == v.data.shape:
|
||||
pretrained_dict[k] = weights[k]
|
||||
else:
|
||||
print('ignore {}'.format(k))
|
||||
model_dict.update(pretrained_dict)
|
||||
model.load_state_dict(model_dict)
|
||||
print('update success')
|
||||
model.cuda(gpu_id)
|
||||
model.eval()
|
||||
self.model = model
|
||||
self.last_mask = None
|
||||
self.output_img_size = 320
|
||||
self.gpu_id = gpu_id
|
||||
|
||||
def inference(self, frame, pt1k, video_mode = False):
|
||||
image_to_face_mat = landmark_processor.get_transform_mat_full_face(pt1k, self.output_img_size)
|
||||
face_image = cv2.warpAffine(frame, image_to_face_mat, (self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4)
|
||||
if face_image.dtype == np.uint8: face_image = face_image.astype(np.float32) / 255
|
||||
if video_mode and self.last_mask is not None:
|
||||
last_small_mask = cv2.warpAffine(self.last_mask, image_to_face_mat,
|
||||
(self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4)[:,:,np.newaxis]
|
||||
input_img = np.concatenate([face_image, last_small_mask], axis=2)
|
||||
else:
|
||||
zero_mask = np.zeros((face_image.shape[1], face_image.shape[0], 1), dtype=np.float32)
|
||||
input_img = np.concatenate([face_image, zero_mask], axis=2)
|
||||
|
||||
face_image_tensor = input_img.transpose((2, 0, 1))[np.newaxis]
|
||||
face_image_tensor = torch.from_numpy(face_image_tensor).cuda(self.gpu_id)
|
||||
with torch.no_grad():
|
||||
mask = self.model.test(face_image_tensor)
|
||||
mask = mask[0].detach().cpu().numpy().transpose((1, 2, 0))
|
||||
|
||||
origin_mask = cv2.warpAffine(mask, image_to_face_mat, (frame.shape[1], frame.shape[0]),
|
||||
flags=cv2.WARP_INVERSE_MAP|cv2.INTER_LANCZOS4)[:, :, np.newaxis]
|
||||
if video_mode: self.last_mask = origin_mask.copy()
|
||||
|
||||
return origin_mask
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
face_segmentor = FaceSeg(gpu_id=0)
|
||||
testdata_dir = "/home/yangchaojie/Desktop/faceswap_hd/datasets/origin/8171"
|
||||
for picname in os.listdir(testdata_dir):
|
||||
img_path = os.path.join(testdata_dir, picname)
|
||||
pkl_path = img_path[:-4]+".pkl"
|
||||
if not picname.endswith(".jpg"):
|
||||
continue
|
||||
if not os.path.exists(pkl_path):
|
||||
continue
|
||||
img = cv2.imread(img_path)
|
||||
with open(pkl_path, "rb") as fp:
|
||||
info = pickle.load(fp)
|
||||
pt1k = info["pt1k"]
|
||||
|
||||
face_seg_mask = face_segmentor.inference(img, pt1k, video_mode=False)
|
||||
# cv2.imshow("face_seg_mask", face_seg_mask)
|
||||
# cv2.imshow("img", img)
|
||||
# cv2.waitKey()
|
||||
@@ -0,0 +1,299 @@
|
||||
from __future__ import division
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.nn.init as init
|
||||
import torch.utils.model_zoo as model_zoo
|
||||
from torchvision import models
|
||||
|
||||
# general libs
|
||||
import cv2
|
||||
import matplotlib.pyplot as plt
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
import math
|
||||
import time
|
||||
import tqdm
|
||||
import os
|
||||
import argparse
|
||||
import copy
|
||||
import sys
|
||||
|
||||
from utils.helpers import *
|
||||
|
||||
|
||||
class ResBlock(nn.Module):
|
||||
def __init__(self, backbone, indim, outdim=None, stride=1):
|
||||
super(ResBlock, self).__init__()
|
||||
self.backbone = backbone
|
||||
if outdim == None:
|
||||
outdim = indim
|
||||
if indim == outdim and stride == 1:
|
||||
self.downsample = None
|
||||
else:
|
||||
self.downsample = nn.Conv2d(indim, outdim, kernel_size=3, padding=1, stride=stride)
|
||||
|
||||
self.conv1 = nn.Conv2d(indim, outdim, kernel_size=3, padding=1, stride=stride)
|
||||
self.conv2 = nn.Conv2d(outdim, outdim, kernel_size=3, padding=1)
|
||||
|
||||
def forward(self, x):
|
||||
if self.backbone == 'resnest101':
|
||||
r = self.conv1(F.relu(x, inplace=True))
|
||||
r = self.conv2(F.relu(r, inplace=True))
|
||||
else:
|
||||
r = self.conv1(F.relu(x))
|
||||
r = self.conv2(F.relu(r))
|
||||
|
||||
if self.downsample is not None:
|
||||
x = self.downsample(x)
|
||||
|
||||
return x + r
|
||||
|
||||
|
||||
class Encoder_M(nn.Module):
|
||||
def __init__(self, backbone):
|
||||
super(Encoder_M, self).__init__()
|
||||
if backbone == 'resnest101':
|
||||
self.conv1_m = nn.Conv2d(1, 128, kernel_size=7, stride=2, padding=3, bias=False)
|
||||
self.conv1_o = nn.Conv2d(1, 128, kernel_size=7, stride=2, padding=3, bias=False)
|
||||
else:
|
||||
self.conv1_m = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
|
||||
self.conv1_o = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
|
||||
|
||||
if backbone == 'resnet50':
|
||||
resnet = models.resnet50(pretrained=True)
|
||||
elif backbone == 'resnet18':
|
||||
resnet = models.resnet18(pretrained=True)
|
||||
|
||||
self.conv1 = resnet.conv1
|
||||
self.bn1 = resnet.bn1
|
||||
self.relu = resnet.relu # 1/2, 64
|
||||
self.maxpool = resnet.maxpool
|
||||
|
||||
self.res2 = resnet.layer1 # 1/4, 256
|
||||
self.res3 = resnet.layer2 # 1/8, 512
|
||||
self.res4 = resnet.layer3 # 1/8, 1024
|
||||
|
||||
self.register_buffer('mean', torch.FloatTensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
|
||||
self.register_buffer('std', torch.FloatTensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
|
||||
|
||||
def forward(self, in_f, in_m, in_o):
|
||||
f = (in_f - self.mean) / self.std
|
||||
m = torch.unsqueeze(in_m, dim=1).float() # add channel dim
|
||||
o = torch.unsqueeze(in_o, dim=1).float() # add channel dim
|
||||
|
||||
x = self.conv1(f) + self.conv1_m(m) + self.conv1_o(o)
|
||||
x = self.bn1(x)
|
||||
c1 = self.relu(x) # 1/2, 64
|
||||
x = self.maxpool(c1) # 1/4, 64
|
||||
r2 = self.res2(x) # 1/4, 256
|
||||
r3 = self.res3(r2) # 1/8, 512
|
||||
r4 = self.res4(r3) # 1/8, 1024
|
||||
return r4, r3, r2, c1, f
|
||||
|
||||
|
||||
class Encoder_Q(nn.Module):
|
||||
def __init__(self, backbone):
|
||||
super(Encoder_Q, self).__init__()
|
||||
|
||||
if backbone == 'resnet50':
|
||||
resnet = models.resnet50(pretrained=True)
|
||||
elif backbone == 'resnet18':
|
||||
resnet = models.resnet18(pretrained=True)
|
||||
|
||||
self.conv1 = resnet.conv1
|
||||
self.bn1 = resnet.bn1
|
||||
self.relu = resnet.relu # 1/2, 64
|
||||
self.maxpool = resnet.maxpool
|
||||
|
||||
self.res2 = resnet.layer1 # 1/4, 256
|
||||
self.res3 = resnet.layer2 # 1/8, 512
|
||||
self.res4 = resnet.layer3 # 1/8, 1024
|
||||
|
||||
self.register_buffer('mean', torch.FloatTensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
|
||||
self.register_buffer('std', torch.FloatTensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
|
||||
|
||||
def forward(self, in_f):
|
||||
f = (in_f - self.mean) / self.std
|
||||
|
||||
x = self.conv1(f)
|
||||
x = self.bn1(x)
|
||||
c1 = self.relu(x) # 1/2, 64
|
||||
x = self.maxpool(c1) # 1/4, 64
|
||||
r2 = self.res2(x) # 1/4, 256
|
||||
r3 = self.res3(r2) # 1/8, 512
|
||||
r4 = self.res4(r3) # 1/8, 1024
|
||||
return r4, r3, r2, c1, f
|
||||
|
||||
|
||||
class Refine(nn.Module):
|
||||
def __init__(self, backbone, inplanes, planes, scale_factor=2):
|
||||
super(Refine, self).__init__()
|
||||
self.convFS = nn.Conv2d(inplanes, planes, kernel_size=(3, 3), padding=(1, 1), stride=1)
|
||||
self.ResFS = ResBlock(backbone, planes, planes)
|
||||
self.ResMM = ResBlock(backbone, planes, planes)
|
||||
self.scale_factor = scale_factor
|
||||
|
||||
def forward(self, f, pm):
|
||||
s = self.ResFS(self.convFS(f))
|
||||
m = s + F.interpolate(pm, scale_factor=self.scale_factor, mode='bilinear', align_corners=False)
|
||||
m = self.ResMM(m)
|
||||
return m
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(self, mdim, scale_rate, backbone):
|
||||
super(Decoder, self).__init__()
|
||||
self.backbone = backbone
|
||||
if backbone == 'resnest101':
|
||||
self.convFM = nn.Conv2d(256, mdim, kernel_size=(3, 3), padding=(1, 1), stride=1)
|
||||
else:
|
||||
self.convFM = nn.Conv2d(1024 // scale_rate, mdim, kernel_size=(3, 3), padding=(1, 1), stride=1)
|
||||
self.ResMM = ResBlock(backbone, mdim, mdim)
|
||||
self.RF3 = Refine(backbone, 512 // scale_rate, mdim) # 1/8 -> 1/4
|
||||
self.RF2 = Refine(backbone, 256 // scale_rate, mdim) # 1/4 -> 1
|
||||
|
||||
self.pred2 = nn.Conv2d(mdim, 2, kernel_size=(3, 3), padding=(1, 1), stride=1)
|
||||
|
||||
def forward(self, r4, r3, r2):
|
||||
m4 = self.ResMM(self.convFM(r4))
|
||||
m3 = self.RF3(r3, m4) # out: 1/8, 256
|
||||
m2 = self.RF2(r2, m3) # out: 1/4, 256
|
||||
|
||||
if self.backbone == 'resnest101':
|
||||
p2 = self.pred2(F.relu(m2, inplace=True))
|
||||
else:
|
||||
p2 = self.pred2(F.relu(m2))
|
||||
|
||||
p = F.interpolate(p2, scale_factor=4, mode='bilinear', align_corners=False)
|
||||
return p # , p2, p3, p4
|
||||
|
||||
|
||||
class Memory(nn.Module):
|
||||
def __init__(self):
|
||||
super(Memory, self).__init__()
|
||||
|
||||
def forward(self, m_in, m_out, q_in, q_out): # m_in: o,c,t,h,w
|
||||
B, D_e, T, H, W = m_in.size()
|
||||
_, D_o, _, _, _ = m_out.size()
|
||||
|
||||
mi = m_in.view(B, D_e, T * H * W)
|
||||
mi = torch.transpose(mi, 1, 2) # b, THW, emb
|
||||
|
||||
qi = q_in.view(B, D_e, H * W) # b, emb, HW
|
||||
|
||||
p = torch.bmm(mi, qi) # b, THW, HW
|
||||
p = p / math.sqrt(D_e)
|
||||
p = F.softmax(p, dim=1) # b, THW, HW
|
||||
|
||||
mo = m_out.view(B, D_o, T * H * W)
|
||||
mem = torch.bmm(mo, p) # Weighted-sum B, D_o, HW
|
||||
mem = mem.view(B, D_o, H, W)
|
||||
|
||||
mem_out = torch.cat([mem, q_out], dim=1)
|
||||
|
||||
return mem_out, p
|
||||
|
||||
|
||||
class KeyValue(nn.Module):
|
||||
# Not using location
|
||||
def __init__(self, indim, keydim, valdim):
|
||||
super(KeyValue, self).__init__()
|
||||
self.Key = nn.Conv2d(indim, keydim, kernel_size=(3, 3), padding=(1, 1), stride=1)
|
||||
self.Value = nn.Conv2d(indim, valdim, kernel_size=(3, 3), padding=(1, 1), stride=1)
|
||||
|
||||
def forward(self, x):
|
||||
return self.Key(x), self.Value(x)
|
||||
|
||||
|
||||
class STM(nn.Module):
|
||||
def __init__(self, backbone='resnet50'):
|
||||
super(STM, self).__init__()
|
||||
self.backbone = backbone
|
||||
assert backbone == 'resnet50' or backbone == 'resnet18' or backbone == 'resnest101'
|
||||
scale_rate = (1 if (backbone == 'resnet50' or backbone == 'resnest101') else 4)
|
||||
|
||||
self.Encoder_M = Encoder_M(backbone)
|
||||
self.Encoder_Q = Encoder_Q(backbone)
|
||||
|
||||
self.KV_M_r4 = KeyValue(1024 // scale_rate, keydim=128 // scale_rate, valdim=512 // scale_rate)
|
||||
self.KV_Q_r4 = KeyValue(1024 // scale_rate, keydim=128 // scale_rate, valdim=512 // scale_rate)
|
||||
|
||||
self.Memory = Memory()
|
||||
self.Decoder = Decoder(256, scale_rate, backbone)
|
||||
|
||||
def Pad_memory(self, mems, num_objects, K):
|
||||
pad_mems = []
|
||||
for mem in mems:
|
||||
pad_mem = ToCuda(torch.zeros(1, K, mem.size()[1], 1, mem.size()[2], mem.size()[3]))
|
||||
pad_mem[0, 1:num_objects + 1, :, 0] = mem
|
||||
pad_mems.append(pad_mem)
|
||||
return pad_mems
|
||||
|
||||
def memorize(self, frame, masks, num_objects):
|
||||
# memorize a frame
|
||||
num_objects = num_objects[0].item()
|
||||
_, K, H, W = masks.shape # B = 1
|
||||
|
||||
(frame, masks), pad = pad_divide_by([frame, masks], 16, (frame.size()[2], frame.size()[3]))
|
||||
|
||||
# make batch arg list
|
||||
B_list = {'f': [], 'm': [], 'o': []}
|
||||
for o in range(1, num_objects + 1): # 1 - no
|
||||
B_list['f'].append(frame)
|
||||
B_list['m'].append(masks[:, o])
|
||||
B_list['o'].append((torch.sum(masks[:, 1:o], dim=1) + \
|
||||
torch.sum(masks[:, o + 1:num_objects + 1], dim=1)).clamp(0, 1))
|
||||
|
||||
# make Batch
|
||||
B_ = {}
|
||||
for arg in B_list.keys():
|
||||
B_[arg] = torch.cat(B_list[arg], dim=0)
|
||||
|
||||
r4, _, _, _, _ = self.Encoder_M(B_['f'], B_['m'], B_['o'])
|
||||
k4, v4 = self.KV_M_r4(r4) # num_objects, 128 and 512, H/16, W/16
|
||||
k4, v4 = self.Pad_memory([k4, v4], num_objects=num_objects, K=K)
|
||||
return k4, v4
|
||||
|
||||
def Soft_aggregation(self, ps, K):
|
||||
num_objects, H, W = ps.shape
|
||||
em = ToCuda(torch.zeros(1, K, H, W))
|
||||
em[0, 0] = torch.prod(1 - ps, dim=0) # bg prob
|
||||
em[0, 1:num_objects + 1] = ps # obj prob
|
||||
em = torch.clamp(em, 1e-7, 1 - 1e-7)
|
||||
logit = torch.log((em / (1 - em)))
|
||||
return logit
|
||||
|
||||
def segment(self, frame, keys, values, num_objects):
|
||||
num_objects = num_objects[0].item()
|
||||
_, K, keydim, T, H, W = keys.shape # B = 1
|
||||
# pad
|
||||
[frame], pad = pad_divide_by([frame], 16, (frame.size()[2], frame.size()[3]))
|
||||
|
||||
r4, r3, r2, _, _ = self.Encoder_Q(frame)
|
||||
k4, v4 = self.KV_Q_r4(r4) # 1, dim, H/16, W/16
|
||||
|
||||
# expand to --- no, c, h, w
|
||||
k4e, v4e = k4.expand(num_objects, -1, -1, -1), v4.expand(num_objects, -1, -1, -1)
|
||||
r3e, r2e = r3.expand(num_objects, -1, -1, -1), r2.expand(num_objects, -1, -1, -1)
|
||||
|
||||
# memory select kv:(1, K, C, T, H, W)
|
||||
m4, viz = self.Memory(keys[0, 1:num_objects + 1], values[0, 1:num_objects + 1], k4e, v4e)
|
||||
logits = self.Decoder(m4, r3e, r2e)
|
||||
ps = F.softmax(logits, dim=1)[:, 1] # no, h, w
|
||||
# ps = indipendant possibility to belong to each object
|
||||
|
||||
logit = self.Soft_aggregation(ps, K) # 1, K, H, W
|
||||
|
||||
if pad[2] + pad[3] > 0:
|
||||
logit = logit[:, :, pad[2]:-pad[3], :]
|
||||
if pad[0] + pad[1] > 0:
|
||||
logit = logit[:, :, :, pad[0]:-pad[1]]
|
||||
|
||||
return logit
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
if args[1].dim() > 4: # keys
|
||||
return self.segment(*args, **kwargs)
|
||||
else:
|
||||
return self.memorize(*args, **kwargs)
|
||||
@@ -0,0 +1,185 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
import numpy as np
|
||||
|
||||
class SequenceConv(nn.ModuleList):
|
||||
"""Sequence conv module.
|
||||
|
||||
Args:
|
||||
in_channels (int): input tensor channel.
|
||||
out_channels (int): output tensor channel.
|
||||
kernel_size (int): convolution kernel size.
|
||||
sequence_num (int): sequence length.
|
||||
conv_cfg (dict): convolution config dictionary.
|
||||
norm_cfg (dict): normalization config dictionary.
|
||||
act_cfg (dict): activation config dictionary.
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels, out_channels, kernel_size, sequence_num):
|
||||
super(SequenceConv, self).__init__()
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.kernel_size = kernel_size
|
||||
self.sequence_num = sequence_num
|
||||
for _ in range(sequence_num):
|
||||
self.append(
|
||||
nn.Sequential(
|
||||
nn.Conv2d(self.in_channels, self.out_channels, self.kernel_size, 1, self.kernel_size // 2, bias=False),
|
||||
nn.BatchNorm2d(self.out_channels),
|
||||
nn.ReLU()
|
||||
)
|
||||
)
|
||||
|
||||
def forward(self, sequence_imgs):
|
||||
"""
|
||||
|
||||
Args:
|
||||
sequence_imgs (Tensor): TxBxCxHxW
|
||||
|
||||
Returns:
|
||||
sequence conv output: TxBxCxHxW
|
||||
"""
|
||||
sequence_outs = []
|
||||
assert sequence_imgs.shape[0] == self.sequence_num
|
||||
for i, sequence_conv in enumerate(self):
|
||||
sequence_out = sequence_conv(sequence_imgs[i, ...])
|
||||
sequence_out = sequence_out.unsqueeze(0)
|
||||
sequence_outs.append(sequence_out)
|
||||
|
||||
sequence_outs = torch.cat(sequence_outs, dim=0) # TxBxCxHxW
|
||||
return sequence_outs
|
||||
|
||||
class MemoryModule(nn.Module):
|
||||
"""Memory read module.
|
||||
Args:
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
matmul_norm=False):
|
||||
super(MemoryModule, self).__init__()
|
||||
self.matmul_norm = matmul_norm
|
||||
|
||||
def forward(self, memory_keys, memory_values, query_key, query_value):
|
||||
"""
|
||||
Memory Module forward.
|
||||
Args:
|
||||
memory_keys (Tensor): memory keys tensor, shape: TxBxCxHxW
|
||||
memory_values (Tensor): memory values tensor, shape: TxBxCxHxW
|
||||
query_key (Tensor): query keys tensor, shape: BxCxHxW
|
||||
query_value (Tensor): query values tensor, shape: BxCxHxW
|
||||
|
||||
Returns:
|
||||
Concat query and memory tensor.
|
||||
"""
|
||||
sequence_num, batch_size, key_channels, height, width = memory_keys.shape
|
||||
_, _, value_channels, _, _ = memory_values.shape
|
||||
assert query_key.shape[1] == key_channels and query_value.shape[1] == value_channels
|
||||
memory_keys = memory_keys.permute(1, 2, 0, 3, 4).contiguous() # BxCxTxHxW
|
||||
memory_keys = memory_keys.view(batch_size, key_channels, sequence_num * height * width) # BxCxT*H*W
|
||||
|
||||
query_key = query_key.view(batch_size, key_channels, height * width).permute(0, 2, 1).contiguous() # BxH*WxCk
|
||||
key_attention = torch.bmm(query_key, memory_keys) # BxH*WxT*H*W
|
||||
if self.matmul_norm:
|
||||
key_attention = (key_channels ** -.5) * key_attention
|
||||
key_attention = F.softmax(key_attention, dim=-1) # BxH*WxT*H*W
|
||||
|
||||
memory_values = memory_values.permute(1, 2, 0, 3, 4).contiguous() # BxCxTxHxW
|
||||
memory_values = memory_values.view(batch_size, value_channels, sequence_num * height * width)
|
||||
memory_values = memory_values.permute(0, 2, 1).contiguous() # BxT*H*WxC
|
||||
memory = torch.bmm(key_attention, memory_values) # BxH*WxC
|
||||
memory = memory.permute(0, 2, 1).contiguous() # BxCxH*W
|
||||
memory = memory.view(batch_size, value_channels, height, width) # BxCxHxW
|
||||
|
||||
query_memory = torch.cat([query_value, memory], dim=1)
|
||||
return query_memory
|
||||
#
|
||||
# class TMAHead(nn.Module):
|
||||
# """TMAHead decoder for video semantic segmentation."""
|
||||
#
|
||||
# def __init__(self, sequence_num, key_channels, value_channels, num_classes=2, dropout_ratio=0):
|
||||
# super(TMAHead, self).__init__()
|
||||
#
|
||||
# self.sequence_num = sequence_num
|
||||
# self.memory_key_conv = nn.Sequential(
|
||||
# SequenceConv(self.in_channels, key_channels, 1, sequence_num),
|
||||
# SequenceConv(key_channels, key_channels, 3, sequence_num)
|
||||
# )
|
||||
# self.memory_value_conv = nn.Sequential(
|
||||
# SequenceConv(self.in_channels, value_channels, 1, sequence_num),
|
||||
# SequenceConv(value_channels, value_channels, 3, sequence_num)
|
||||
# )
|
||||
# self.query_key_conv = nn.Sequential(
|
||||
# nn.Sequential(
|
||||
# nn.Conv2d(self.in_channels, key_channels, 1, 1, 0, bias=False),
|
||||
# nn.BatchNorm2d(key_channels),
|
||||
# nn.ReLU()
|
||||
# ),
|
||||
# nn.Sequential(
|
||||
# nn.Conv2d(key_channels, key_channels, 3, 1, 1, bias=False),
|
||||
# nn.BatchNorm2d(key_channels),
|
||||
# nn.ReLU()
|
||||
# ),
|
||||
# )
|
||||
#
|
||||
# self.query_value_conv = nn.Sequential(
|
||||
# nn.Sequential(
|
||||
# nn.Conv2d(self.in_channels, value_channels, 1, 1, 0, bias=False),
|
||||
# nn.BatchNorm2d(value_channels),
|
||||
# nn.ReLU()
|
||||
# ),
|
||||
# nn.Sequential(
|
||||
# nn.Conv2d(value_channels, value_channels, 3, 1, 1, bias=False),
|
||||
# nn.BatchNorm2d(value_channels),
|
||||
# nn.ReLU()
|
||||
# ),
|
||||
# )
|
||||
# self.memory_module = MemoryModule(matmul_norm=False)
|
||||
# self.bottleneck = nn.Sequential(
|
||||
# nn.Conv2d(value_channels * 2, self.channels, 3, 1, 1, bias=False),
|
||||
# nn.BatchNorm2d(value_channels),
|
||||
# nn.ReLU()
|
||||
# )
|
||||
#
|
||||
# self.conv_seg = nn.Conv2d(self.channels, num_classes, kernel_size=1)
|
||||
# if dropout_ratio > 0:
|
||||
# self.dropout = nn.Dropout2d(dropout_ratio)
|
||||
# else:
|
||||
# self.dropout = None
|
||||
#
|
||||
# def cls_seg(self, feat):
|
||||
# """Classify each pixel."""
|
||||
# if self.dropout is not None:
|
||||
# feat = self.dropout(feat)
|
||||
# output = self.conv_seg(feat)
|
||||
# return output
|
||||
#
|
||||
# def forward(self, inputs, sequence_imgs):
|
||||
# """
|
||||
# Forward fuction.
|
||||
# Args:
|
||||
# inputs (list[Tensor]): backbone multi-level outputs.
|
||||
# sequence_imgs (list[Tensor]): len(sequence_imgs) is equal to batch_size,
|
||||
# each element is a Tensor with shape of TxCxHxW.
|
||||
#
|
||||
# Returns:
|
||||
# decoder logits.
|
||||
# """
|
||||
# x = inputs
|
||||
# sequence_imgs = [y.unsqueeze(0) for y in sequence_imgs] # T, BxCxHxW
|
||||
# sequence_imgs = torch.cat(sequence_imgs, dim=0) # TxBxCxHxW
|
||||
# sequence_num, batch_size, channels, height, width = sequence_imgs.shape
|
||||
#
|
||||
# assert sequence_num == self.sequence_num
|
||||
# memory_keys = self.memory_key_conv(sequence_imgs)
|
||||
# memory_values = self.memory_value_conv(sequence_imgs)
|
||||
# query_key = self.query_key_conv(x) # BxCxHxW
|
||||
# query_value = self.query_value_conv(x) # BxCxHxW
|
||||
#
|
||||
# # memory read
|
||||
# output = self.memory_module(memory_keys, memory_values, query_key, query_value)
|
||||
# output = self.bottleneck(output)
|
||||
# output = self.cls_seg(output)
|
||||
#
|
||||
# return output
|
||||
@@ -0,0 +1,624 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class REBNCONV(nn.Module):
|
||||
def __init__(self,in_ch=3,out_ch=3,dirate=1):
|
||||
super(REBNCONV,self).__init__()
|
||||
|
||||
self.conv_s1 = nn.Conv2d(in_ch,out_ch,3,padding=1*dirate,dilation=1*dirate)
|
||||
self.bn_s1 = nn.BatchNorm2d(out_ch)
|
||||
self.relu_s1 = nn.ReLU(inplace=True)
|
||||
|
||||
def forward(self,x):
|
||||
|
||||
hx = x
|
||||
xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
|
||||
|
||||
return xout
|
||||
|
||||
## upsample tensor 'src' to have the same spatial size with tensor 'tar'
|
||||
def _upsample_like(src,tar):
|
||||
|
||||
src = F.upsample(src,size=tar.shape[2:],mode='bilinear')
|
||||
|
||||
return src
|
||||
|
||||
|
||||
### RSU-7 ###
|
||||
class RSU7(nn.Module):#UNet07DRES(nn.Module):
|
||||
|
||||
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
|
||||
super(RSU7,self).__init__()
|
||||
|
||||
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
|
||||
|
||||
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
|
||||
self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
|
||||
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||
self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
|
||||
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||
self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
|
||||
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||
self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
|
||||
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||
self.pool5 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
|
||||
self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||
|
||||
self.rebnconv7 = REBNCONV(mid_ch,mid_ch,dirate=2)
|
||||
|
||||
self.rebnconv6d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||
self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
|
||||
|
||||
def forward(self,x):
|
||||
|
||||
hx = x
|
||||
hxin = self.rebnconvin(hx)
|
||||
|
||||
hx1 = self.rebnconv1(hxin)
|
||||
hx = self.pool1(hx1)
|
||||
|
||||
hx2 = self.rebnconv2(hx)
|
||||
hx = self.pool2(hx2)
|
||||
|
||||
hx3 = self.rebnconv3(hx)
|
||||
hx = self.pool3(hx3)
|
||||
|
||||
hx4 = self.rebnconv4(hx)
|
||||
hx = self.pool4(hx4)
|
||||
|
||||
hx5 = self.rebnconv5(hx)
|
||||
hx = self.pool5(hx5)
|
||||
|
||||
hx6 = self.rebnconv6(hx)
|
||||
|
||||
hx7 = self.rebnconv7(hx6)
|
||||
|
||||
hx6d = self.rebnconv6d(torch.cat((hx7,hx6),1))
|
||||
hx6dup = _upsample_like(hx6d,hx5)
|
||||
|
||||
hx5d = self.rebnconv5d(torch.cat((hx6dup,hx5),1))
|
||||
hx5dup = _upsample_like(hx5d,hx4)
|
||||
|
||||
hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1))
|
||||
hx4dup = _upsample_like(hx4d,hx3)
|
||||
|
||||
hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
|
||||
hx3dup = _upsample_like(hx3d,hx2)
|
||||
|
||||
hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
|
||||
hx2dup = _upsample_like(hx2d,hx1)
|
||||
|
||||
hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
|
||||
|
||||
return hx1d + hxin
|
||||
|
||||
### RSU-6 ###
|
||||
class RSU6(nn.Module):#UNet06DRES(nn.Module):
|
||||
|
||||
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
|
||||
super(RSU6,self).__init__()
|
||||
|
||||
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
|
||||
|
||||
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
|
||||
self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
|
||||
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||
self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
|
||||
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||
self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
|
||||
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||
self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
|
||||
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||
|
||||
self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=2)
|
||||
|
||||
self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
|
||||
|
||||
def forward(self,x):
|
||||
|
||||
hx = x
|
||||
|
||||
hxin = self.rebnconvin(hx)
|
||||
|
||||
hx1 = self.rebnconv1(hxin)
|
||||
hx = self.pool1(hx1)
|
||||
|
||||
hx2 = self.rebnconv2(hx)
|
||||
hx = self.pool2(hx2)
|
||||
|
||||
hx3 = self.rebnconv3(hx)
|
||||
hx = self.pool3(hx3)
|
||||
|
||||
hx4 = self.rebnconv4(hx)
|
||||
hx = self.pool4(hx4)
|
||||
|
||||
hx5 = self.rebnconv5(hx)
|
||||
|
||||
hx6 = self.rebnconv6(hx5)
|
||||
|
||||
|
||||
hx5d = self.rebnconv5d(torch.cat((hx6,hx5),1))
|
||||
hx5dup = _upsample_like(hx5d,hx4)
|
||||
|
||||
hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1))
|
||||
hx4dup = _upsample_like(hx4d,hx3)
|
||||
|
||||
hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
|
||||
hx3dup = _upsample_like(hx3d,hx2)
|
||||
|
||||
hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
|
||||
hx2dup = _upsample_like(hx2d,hx1)
|
||||
|
||||
hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
|
||||
|
||||
return hx1d + hxin
|
||||
|
||||
### RSU-5 ###
|
||||
class RSU5(nn.Module):#UNet05DRES(nn.Module):
|
||||
|
||||
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
|
||||
super(RSU5,self).__init__()
|
||||
|
||||
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
|
||||
|
||||
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
|
||||
self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
|
||||
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||
self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
|
||||
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||
self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
|
||||
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||
|
||||
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=2)
|
||||
|
||||
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
|
||||
|
||||
def forward(self,x):
|
||||
|
||||
hx = x
|
||||
|
||||
hxin = self.rebnconvin(hx)
|
||||
|
||||
hx1 = self.rebnconv1(hxin)
|
||||
hx = self.pool1(hx1)
|
||||
|
||||
hx2 = self.rebnconv2(hx)
|
||||
hx = self.pool2(hx2)
|
||||
|
||||
hx3 = self.rebnconv3(hx)
|
||||
hx = self.pool3(hx3)
|
||||
|
||||
hx4 = self.rebnconv4(hx)
|
||||
|
||||
hx5 = self.rebnconv5(hx4)
|
||||
|
||||
hx4d = self.rebnconv4d(torch.cat((hx5,hx4),1))
|
||||
hx4dup = _upsample_like(hx4d,hx3)
|
||||
|
||||
hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
|
||||
hx3dup = _upsample_like(hx3d,hx2)
|
||||
|
||||
hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
|
||||
hx2dup = _upsample_like(hx2d,hx1)
|
||||
|
||||
hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
|
||||
|
||||
return hx1d + hxin
|
||||
|
||||
### RSU-4 ###
|
||||
class RSU4(nn.Module):#UNet04DRES(nn.Module):
|
||||
|
||||
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
|
||||
super(RSU4,self).__init__()
|
||||
|
||||
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
|
||||
|
||||
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
|
||||
self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
|
||||
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||
self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
|
||||
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
|
||||
|
||||
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=2)
|
||||
|
||||
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
|
||||
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
|
||||
|
||||
def forward(self,x):
|
||||
|
||||
hx = x
|
||||
|
||||
hxin = self.rebnconvin(hx)
|
||||
|
||||
hx1 = self.rebnconv1(hxin)
|
||||
hx = self.pool1(hx1)
|
||||
|
||||
hx2 = self.rebnconv2(hx)
|
||||
hx = self.pool2(hx2)
|
||||
|
||||
hx3 = self.rebnconv3(hx)
|
||||
|
||||
hx4 = self.rebnconv4(hx3)
|
||||
|
||||
hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1))
|
||||
hx3dup = _upsample_like(hx3d,hx2)
|
||||
|
||||
hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
|
||||
hx2dup = _upsample_like(hx2d,hx1)
|
||||
|
||||
hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
|
||||
|
||||
return hx1d + hxin
|
||||
|
||||
### RSU-4F ###
|
||||
class RSU4F(nn.Module):#UNet04FRES(nn.Module):
|
||||
|
||||
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
|
||||
super(RSU4F,self).__init__()
|
||||
|
||||
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
|
||||
|
||||
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
|
||||
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=2)
|
||||
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=4)
|
||||
|
||||
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=8)
|
||||
|
||||
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=4)
|
||||
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=2)
|
||||
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
|
||||
|
||||
def forward(self,x):
|
||||
|
||||
hx = x
|
||||
|
||||
hxin = self.rebnconvin(hx)
|
||||
|
||||
hx1 = self.rebnconv1(hxin)
|
||||
hx2 = self.rebnconv2(hx1)
|
||||
hx3 = self.rebnconv3(hx2)
|
||||
|
||||
hx4 = self.rebnconv4(hx3)
|
||||
|
||||
hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1))
|
||||
hx2d = self.rebnconv2d(torch.cat((hx3d,hx2),1))
|
||||
hx1d = self.rebnconv1d(torch.cat((hx2d,hx1),1))
|
||||
|
||||
return hx1d + hxin
|
||||
|
||||
from core.faceseg.tma import SequenceConv, MemoryModule
|
||||
|
||||
##### U^2-Net ####
|
||||
class U2NET(nn.Module):
|
||||
|
||||
def __init__(self, in_ch=3, out_ch=1):
|
||||
super(U2NET, self).__init__()
|
||||
|
||||
self.stage1 = RSU7(in_ch,32,64)
|
||||
self.pool12 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
|
||||
self.stage2 = RSU6(64,32,128)
|
||||
self.pool23 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
|
||||
self.stage3 = RSU5(128,64,256)
|
||||
self.pool34 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
|
||||
self.stage4 = RSU4(256,128,512)
|
||||
self.pool45 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
|
||||
self.stage5 = RSU4F(512,256,512)
|
||||
self.pool56 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
|
||||
self.stage6 = RSU4F(512,256,512)
|
||||
|
||||
# decoder
|
||||
self.stage5d = RSU4F(1024,256,512)
|
||||
self.stage4d = RSU4(1024,128,256)
|
||||
self.stage3d = RSU5(512,64,128)
|
||||
self.stage2d = RSU6(256,32,64)
|
||||
self.stage1d = RSU7(128,16,64)
|
||||
|
||||
self.side1 = nn.Conv2d(64,out_ch,3,padding=1)
|
||||
self.side2 = nn.Conv2d(64,out_ch,3,padding=1)
|
||||
self.side3 = nn.Conv2d(128,out_ch,3,padding=1)
|
||||
self.side4 = nn.Conv2d(256,out_ch,3,padding=1)
|
||||
self.side5 = nn.Conv2d(512,out_ch,3,padding=1)
|
||||
self.side6 = nn.Conv2d(512,out_ch,3,padding=1)
|
||||
|
||||
self.outconv = nn.Conv2d(6*out_ch,out_ch,1)
|
||||
|
||||
self.in_channels = 512
|
||||
key_channels = 128
|
||||
value_channels = 512
|
||||
self.sequence_num = sequence_num = 2
|
||||
self.memory_key_conv = nn.Sequential(
|
||||
SequenceConv(self.in_channels, key_channels, 1, sequence_num),
|
||||
SequenceConv(key_channels, key_channels, 3, sequence_num)
|
||||
)
|
||||
self.memory_value_conv = nn.Sequential(
|
||||
SequenceConv(self.in_channels, value_channels, 1, sequence_num),
|
||||
SequenceConv(value_channels, value_channels, 3, sequence_num)
|
||||
)
|
||||
self.query_key_conv = nn.Sequential(
|
||||
nn.Sequential(
|
||||
nn.Conv2d(self.in_channels, key_channels, 1, 1, 0, bias=False),
|
||||
nn.BatchNorm2d(key_channels),
|
||||
nn.ReLU()
|
||||
),
|
||||
nn.Sequential(
|
||||
nn.Conv2d(key_channels, key_channels, 3, 1, 1, bias=False),
|
||||
nn.BatchNorm2d(key_channels),
|
||||
nn.ReLU()
|
||||
),
|
||||
)
|
||||
self.query_value_conv = nn.Sequential(
|
||||
nn.Sequential(
|
||||
nn.Conv2d(self.in_channels, value_channels, 1, 1, 0, bias=False),
|
||||
nn.BatchNorm2d(value_channels),
|
||||
nn.ReLU()
|
||||
),
|
||||
nn.Sequential(
|
||||
nn.Conv2d(value_channels, value_channels, 3, 1, 1, bias=False),
|
||||
nn.BatchNorm2d(value_channels),
|
||||
nn.ReLU()
|
||||
),
|
||||
)
|
||||
self.memory_module = MemoryModule(matmul_norm=False)
|
||||
self.bottleneck = nn.Sequential(
|
||||
nn.Conv2d(value_channels * 2, self.in_channels, 3, 1, 1, bias=False),
|
||||
nn.BatchNorm2d(value_channels),
|
||||
nn.ReLU()
|
||||
)
|
||||
|
||||
self.is_train = True
|
||||
|
||||
def extract_feature(self, x):
|
||||
hx = x
|
||||
|
||||
# stage 1
|
||||
hx1 = self.stage1(hx)
|
||||
hx = self.pool12(hx1)
|
||||
|
||||
# stage 2
|
||||
hx2 = self.stage2(hx)
|
||||
hx = self.pool23(hx2)
|
||||
|
||||
# stage 3
|
||||
hx3 = self.stage3(hx)
|
||||
hx = self.pool34(hx3)
|
||||
|
||||
# stage 4
|
||||
hx4 = self.stage4(hx)
|
||||
hx = self.pool45(hx4)
|
||||
|
||||
# stage 5
|
||||
hx5 = self.stage5(hx)
|
||||
hx = self.pool56(hx5)
|
||||
|
||||
# stage 6
|
||||
hx6 = self.stage6(hx)
|
||||
|
||||
return hx1, hx2, hx3, hx4, hx5, hx6
|
||||
|
||||
def decoder(self, hx1, hx2, hx3, hx4, hx5, hx6):
|
||||
hx6up = _upsample_like(hx6, hx5)
|
||||
|
||||
# -------------------- decoder --------------------
|
||||
hx5d = self.stage5d(torch.cat((hx6up, hx5), 1))
|
||||
hx5dup = _upsample_like(hx5d, hx4)
|
||||
|
||||
hx4d = self.stage4d(torch.cat((hx5dup, hx4), 1))
|
||||
hx4dup = _upsample_like(hx4d, hx3)
|
||||
|
||||
hx3d = self.stage3d(torch.cat((hx4dup, hx3), 1))
|
||||
hx3dup = _upsample_like(hx3d, hx2)
|
||||
|
||||
hx2d = self.stage2d(torch.cat((hx3dup, hx2), 1))
|
||||
hx2dup = _upsample_like(hx2d, hx1)
|
||||
|
||||
hx1d = self.stage1d(torch.cat((hx2dup, hx1), 1))
|
||||
|
||||
return hx1d, hx2d, hx3d, hx4d, hx5d
|
||||
|
||||
def side_output(self, hx1d, hx2d, hx3d, hx4d, hx5d, hx6):
|
||||
# side output
|
||||
d1 = self.side1(hx1d)
|
||||
|
||||
d2 = self.side2(hx2d)
|
||||
d2 = _upsample_like(d2, d1)
|
||||
|
||||
d3 = self.side3(hx3d)
|
||||
d3 = _upsample_like(d3, d1)
|
||||
|
||||
d4 = self.side4(hx4d)
|
||||
d4 = _upsample_like(d4, d1)
|
||||
|
||||
d5 = self.side5(hx5d)
|
||||
d5 = _upsample_like(d5, d1)
|
||||
|
||||
d6 = self.side6(hx6)
|
||||
d6 = _upsample_like(d6, d1)
|
||||
|
||||
d0 = self.outconv(torch.cat((d1, d2, d3, d4, d5, d6), 1))
|
||||
|
||||
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
|
||||
|
||||
def forward(self, x, memory_sequence=None):
|
||||
if self.is_train:
|
||||
hx1, hx2, hx3, hx4, hx5, hx6 = self.extract_feature(x)
|
||||
|
||||
if memory_sequence is None:
|
||||
memory_hx6 = [hx6 for _ in range(self.sequence_num)]
|
||||
else:
|
||||
memory_hx6 = []
|
||||
for single_memory in memory_sequence:
|
||||
_, _, _, _, _, hx6 = self.extract_feature(single_memory)
|
||||
memory_hx6.append(hx6)
|
||||
|
||||
memory_hx6 = [mhx6.unsqueeze(0) for mhx6 in memory_hx6] # T, BxCxHxW
|
||||
memory_hx6 = torch.cat(memory_hx6, dim=0)
|
||||
memory_keys = self.memory_key_conv(memory_hx6)
|
||||
memory_values = self.memory_value_conv(memory_hx6)
|
||||
query_key = self.query_key_conv(hx6)
|
||||
query_value = self.query_value_conv(hx6)
|
||||
merge_hx6 = self.memory_module(memory_keys, memory_values, query_key, query_value)
|
||||
merge_hx6 = self.bottleneck(merge_hx6)
|
||||
|
||||
hx1d, hx2d, hx3d, hx4d, hx5d = self.decoder(hx1, hx2, hx3, hx4, hx5, merge_hx6)
|
||||
|
||||
return self.side_output(hx1d, hx2d, hx3d, hx4d, hx5d, hx6)
|
||||
else:
|
||||
return self.test(x)
|
||||
|
||||
def test(self, x):
|
||||
with torch.no_grad():
|
||||
hx1, hx2, hx3, hx4, hx5, hx6 = self.extract_feature(x)
|
||||
|
||||
memory_hx6 = [hx6 for _ in range(self.sequence_num)]
|
||||
|
||||
memory_hx6 = [mhx6.unsqueeze(0) for mhx6 in memory_hx6] # T, BxCxHxW
|
||||
memory_hx6 = torch.cat(memory_hx6, dim=0)
|
||||
memory_keys = self.memory_key_conv(memory_hx6)
|
||||
memory_values = self.memory_value_conv(memory_hx6)
|
||||
query_key = self.query_key_conv(hx6)
|
||||
query_value = self.query_value_conv(hx6)
|
||||
merge_hx6 = self.memory_module(memory_keys, memory_values, query_key, query_value)
|
||||
merge_hx6 = self.bottleneck(merge_hx6)
|
||||
|
||||
hx1d, hx2d, hx3d, hx4d, hx5d = self.decoder(hx1, hx2, hx3, hx4, hx5, merge_hx6)
|
||||
|
||||
mask, _, _, _, _, _, _ = self.side_output(hx1d, hx2d, hx3d, hx4d, hx5d, hx6)
|
||||
|
||||
return mask
|
||||
#
|
||||
# ### U^2-Net small ###
|
||||
# class U2NETP(nn.Module):
|
||||
#
|
||||
# def __init__(self,in_ch=3,out_ch=1):
|
||||
# super(U2NETP,self).__init__()
|
||||
#
|
||||
# self.stage1 = RSU7(in_ch,16,64)
|
||||
# self.pool12 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
#
|
||||
# self.stage2 = RSU6(64,16,64)
|
||||
# self.pool23 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
#
|
||||
# self.stage3 = RSU5(64,16,64)
|
||||
# self.pool34 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
#
|
||||
# self.stage4 = RSU4(64,16,64)
|
||||
# self.pool45 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
#
|
||||
# self.stage5 = RSU4F(64,16,64)
|
||||
# self.pool56 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
|
||||
#
|
||||
# self.stage6 = RSU4F(64,16,64)
|
||||
#
|
||||
# # decoder
|
||||
# self.stage5d = RSU4F(128,16,64)
|
||||
# self.stage4d = RSU4(128,16,64)
|
||||
# self.stage3d = RSU5(128,16,64)
|
||||
# self.stage2d = RSU6(128,16,64)
|
||||
# self.stage1d = RSU7(128,16,64)
|
||||
#
|
||||
# self.side1 = nn.Conv2d(64,out_ch,3,padding=1)
|
||||
# self.side2 = nn.Conv2d(64,out_ch,3,padding=1)
|
||||
# self.side3 = nn.Conv2d(64,out_ch,3,padding=1)
|
||||
# self.side4 = nn.Conv2d(64,out_ch,3,padding=1)
|
||||
# self.side5 = nn.Conv2d(64,out_ch,3,padding=1)
|
||||
# self.side6 = nn.Conv2d(64,out_ch,3,padding=1)
|
||||
#
|
||||
# self.outconv = nn.Conv2d(6*out_ch,out_ch,1)
|
||||
#
|
||||
# def forward(self,x):
|
||||
#
|
||||
# hx = x
|
||||
#
|
||||
# #stage 1
|
||||
# hx1 = self.stage1(hx)
|
||||
# hx = self.pool12(hx1)
|
||||
#
|
||||
# #stage 2
|
||||
# hx2 = self.stage2(hx)
|
||||
# hx = self.pool23(hx2)
|
||||
#
|
||||
# #stage 3
|
||||
# hx3 = self.stage3(hx)
|
||||
# hx = self.pool34(hx3)
|
||||
#
|
||||
# #stage 4
|
||||
# hx4 = self.stage4(hx)
|
||||
# hx = self.pool45(hx4)
|
||||
#
|
||||
# #stage 5
|
||||
# hx5 = self.stage5(hx)
|
||||
# hx = self.pool56(hx5)
|
||||
#
|
||||
# #stage 6
|
||||
# hx6 = self.stage6(hx)
|
||||
# hx6up = _upsample_like(hx6,hx5)
|
||||
#
|
||||
# #decoder
|
||||
# hx5d = self.stage5d(torch.cat((hx6up,hx5),1))
|
||||
# hx5dup = _upsample_like(hx5d,hx4)
|
||||
#
|
||||
# hx4d = self.stage4d(torch.cat((hx5dup,hx4),1))
|
||||
# hx4dup = _upsample_like(hx4d,hx3)
|
||||
#
|
||||
# hx3d = self.stage3d(torch.cat((hx4dup,hx3),1))
|
||||
# hx3dup = _upsample_like(hx3d,hx2)
|
||||
#
|
||||
# hx2d = self.stage2d(torch.cat((hx3dup,hx2),1))
|
||||
# hx2dup = _upsample_like(hx2d,hx1)
|
||||
#
|
||||
# hx1d = self.stage1d(torch.cat((hx2dup,hx1),1))
|
||||
#
|
||||
#
|
||||
# #side output
|
||||
# d1 = self.side1(hx1d)
|
||||
#
|
||||
# d2 = self.side2(hx2d)
|
||||
# d2 = _upsample_like(d2,d1)
|
||||
#
|
||||
# d3 = self.side3(hx3d)
|
||||
# d3 = _upsample_like(d3,d1)
|
||||
#
|
||||
# d4 = self.side4(hx4d)
|
||||
# d4 = _upsample_like(d4,d1)
|
||||
#
|
||||
# d5 = self.side5(hx5d)
|
||||
# d5 = _upsample_like(d5,d1)
|
||||
#
|
||||
# d6 = self.side6(hx6)
|
||||
# d6 = _upsample_like(d6,d1)
|
||||
#
|
||||
# d0 = self.outconv(torch.cat((d1,d2,d3,d4,d5,d6),1))
|
||||
#
|
||||
# return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
|
||||
@@ -0,0 +1,210 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
import json
|
||||
import os
|
||||
import os.path as osp
|
||||
from common.logger import LogFactory
|
||||
from core.process_modules import Get_Landmark, Process_Data, localtranslationwarpfastwithstrength,\
|
||||
Generator_Hair, chinClass, Generator_Fusion_Res, Change_Hair_Color, GenderClassifyProcessor, BodySeg, \
|
||||
localtranslationwarpfastwithstrength_v2, localtranslationwarpfastwithstrength_v2_soft, updateEndPosition
|
||||
from core.face_enhance.face_enhancement import FaceEnhancement
|
||||
from datetime import datetime
|
||||
import random
|
||||
from core.utils import landmark_processor
|
||||
from core.cos_module import COS_object as OSS_object
|
||||
from core.faceseg.face_seg import FaceSeg
|
||||
from common.logger import config
|
||||
|
||||
class Prepare_Ref_HairColor_Data(object):
|
||||
def __init__(self, gpu, device_id):
|
||||
if gpu and torch.cuda.is_available():
|
||||
self.device = torch.device("cuda:%d" % device_id if device_id >= 0 else "cpu")
|
||||
self.cuda = True
|
||||
else:
|
||||
self.device = torch.device("cpu")
|
||||
self.cuda = False
|
||||
|
||||
self.process_data = Process_Data(gpu, device_id)
|
||||
self.get_landmark = Get_Landmark(gpu_id=device_id)
|
||||
self.color_output_size = 768
|
||||
|
||||
def get_prepare_ref_haircolor_768_data(self, ref_rgb_8uc3_orisize):
|
||||
|
||||
ref_landmark_1k2_f_orisize, _, _ = self.get_landmark.forward(ref_rgb_8uc3_orisize)
|
||||
|
||||
ref_matte_fg_8uc3_orisize, ref_matte_pred_8uc1_orisize, _ = self.process_data.generator_matte.matte_inference(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize)
|
||||
|
||||
# 光头分割
|
||||
ref_baldseg_8uc3_orisize = self.process_data.generator_baldseg.forward(ref_rgb_8uc3_orisize, ref_matte_pred_8uc1_orisize,
|
||||
ref_landmark_1k2_f_orisize)
|
||||
|
||||
color_hair_M = landmark_processor.get_transform_mat_hair_ratio_v1(ref_landmark_1k2_f_orisize, self.color_output_size, ratio=0.35, h_offset=0.45)
|
||||
|
||||
ref_matte_pred_8uc3_orisize = np.repeat(ref_matte_pred_8uc1_orisize[:, :, np.newaxis], 3, axis=2)
|
||||
|
||||
# show_concat = np.concatenate((ref_rgb_8uc3_orisize, ref_matte_pred_8uc3_orisize, ref_baldseg_8uc3_orisize), axis=1)
|
||||
# show_concat = cv2.resize(show_concat, (0, 0), fx=0.5, fy=0.5)
|
||||
# cv2.imshow("show_concat", show_concat)
|
||||
# cv2.waitKey()
|
||||
|
||||
ref_rgb_8uc3_768 = cv2.warpAffine(ref_rgb_8uc3_orisize, color_hair_M, (self.color_output_size, self.color_output_size))
|
||||
ref_matte_pred_8uc3_768 = cv2.warpAffine(ref_matte_pred_8uc3_orisize, color_hair_M, (self.color_output_size, self.color_output_size))
|
||||
ref_baldseg_8uc3_768 = cv2.warpAffine(ref_baldseg_8uc3_orisize, color_hair_M, (self.color_output_size, self.color_output_size), flags=cv2.INTER_NEAREST)
|
||||
ref_landmark_1k2_f_768 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, color_hair_M)
|
||||
|
||||
return ref_rgb_8uc3_768, ref_matte_pred_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_1k2_f_768
|
||||
|
||||
class Prepare_Ref_HairStyle_Data(object):
|
||||
def __init__(self, gpu, device_id):
|
||||
|
||||
if gpu and torch.cuda.is_available():
|
||||
self.device = torch.device("cuda:%d" % device_id if device_id >= 0 else "cpu")
|
||||
self.cuda = True
|
||||
else:
|
||||
self.device = torch.device("cpu")
|
||||
self.cuda = False
|
||||
|
||||
self.process_data = Process_Data(gpu, device_id)
|
||||
self.get_landmark = Get_Landmark(gpu_id=device_id)
|
||||
self.gender_classify = GenderClassifyProcessor(gpu_id=device_id)
|
||||
|
||||
def get_prepare_ref_768_color_data(self, ref_rgb_8uc3_orisize):
|
||||
ref_landmark_1k2_f_orisize = self.get_landmark.forward(ref_rgb_8uc3_orisize)
|
||||
ref_rgb_8uc3_color_768, ref_matting_8uc3_color_768, ref_baldseg_8uc3_color_768, ref_landmark_f1k2_color_768 = \
|
||||
self.process_data.get_prepare_ref_768_bald_data(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize)
|
||||
|
||||
return ref_rgb_8uc3_color_768, ref_matting_8uc3_color_768, ref_baldseg_8uc3_color_768, ref_landmark_f1k2_color_768
|
||||
|
||||
def get_prepare_ref_768_color_data_landmark1k(self, ref_rgb_8uc3_orisize,ref_landmark_1k2_f_orisize):
|
||||
# ref_landmark_1k2_f_orisize = self.get_landmark.inference(ref_rgb_8uc3_orisize)
|
||||
ref_rgb_8uc3_color_768, ref_matting_8uc3_color_768, ref_baldseg_8uc3_color_768, ref_landmark_f1k2_color_768 = \
|
||||
self.process_data.get_prepare_ref_768_bald_data(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize)
|
||||
|
||||
return ref_rgb_8uc3_color_768, ref_matting_8uc3_color_768, ref_baldseg_8uc3_color_768, ref_landmark_f1k2_color_768
|
||||
|
||||
def calculate_hair_ratio_after_align(self, hair_mask, origin_landmark1k, img_size=768):
|
||||
image_to_face_mat = landmark_processor.get_transform_mat_hair_ratio_v1(origin_landmark1k, 768, ratio=0.35, h_offset=0.32)
|
||||
hair_mask_align = cv2.warpAffine(hair_mask, image_to_face_mat, (img_size, img_size))
|
||||
|
||||
hair_rect = cv2.boundingRect(hair_mask_align[:, :, :1])
|
||||
hair_mask_ratio = hair_rect[2] * hair_rect[3] / (img_size * img_size)
|
||||
return hair_mask_ratio
|
||||
|
||||
def check_female_hair_ratio(self, origin_img_8uc3, landmark_1k2_f_orisize):
|
||||
ref_matte_fg_8uc3_orisize, ref_matte_pred_8uc1_orisize, _ = self.process_data.generator_matte.matte_inference(origin_img_8uc3, landmark_1k2_f_orisize)
|
||||
ref_matte_pred_8uc3_orisize = np.repeat(ref_matte_pred_8uc1_orisize[:, :, np.newaxis], 3, axis=2)
|
||||
hairstyle_M = self.process_data.get_hair_M_girl_v1(landmark_1k2_f_orisize)
|
||||
ref_rgb_8uc3_768 = cv2.warpAffine(ref_matte_pred_8uc3_orisize, hairstyle_M, (768, 768))
|
||||
|
||||
# cv2.imshow("ref_rgb_8uc3_768", ref_rgb_8uc3_768)
|
||||
# cv2.waitKey()
|
||||
|
||||
edge_width = 5
|
||||
if (ref_rgb_8uc3_768[-edge_width:, :, :]).max() > 0 or (ref_rgb_8uc3_768[:, -edge_width:, :]).max() > 0 or (ref_rgb_8uc3_768[:, :edge_width, :]).max() > 0:
|
||||
return False
|
||||
else:
|
||||
return True
|
||||
|
||||
def get_prepare_ref_768_data(self, ref_rgb_8uc3_orisize):
|
||||
|
||||
ref_landmark_1k2_f_orisize, _, _ = self.get_landmark.forward(ref_rgb_8uc3_orisize)
|
||||
# for i in range(1000):
|
||||
# cv2.circle(ref_rgb_8uc3_orisize, (int(ref_landmark_1k2_f_orisize[i][0]),int(ref_landmark_1k2_f_orisize[i][1])), 1,(255, 0,0), 1)
|
||||
# cv2.imshow('ffff', ref_rgb_8uc3_orisize)
|
||||
# cv2.waitKey()
|
||||
# ref_matte_fg_8uc3_orisize, ref_matte_pred_8uc1_orisize = self.process_data.generator_matte.matte_inference(
|
||||
# ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize)
|
||||
ref_landmark_137kpts_f_orisize = landmark_processor.pts_1k_to_137(ref_landmark_1k2_f_orisize)
|
||||
gender_res = self.gender_classify.forward(ref_rgb_8uc3_orisize, ref_landmark_137kpts_f_orisize)
|
||||
# hair_ratio = self.calculate_hair_ratio_after_align(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize)
|
||||
# if hair_ratio > 0.3:
|
||||
# ratio = 2
|
||||
# else:
|
||||
# if gender_res:
|
||||
# ratio = 1
|
||||
# else:
|
||||
# ratio = 0
|
||||
if not gender_res:
|
||||
ratio = 0
|
||||
else:
|
||||
check_res = self.check_female_hair_ratio(ref_rgb_8uc3_orisize, ref_landmark_137kpts_f_orisize)
|
||||
if check_res:
|
||||
ratio = 1
|
||||
else:
|
||||
ratio = 2
|
||||
|
||||
if gender_res:
|
||||
gender = "girl"
|
||||
else:
|
||||
gender = "boy"
|
||||
|
||||
# show_concat = np.concatenate((ref_rgb_8uc3_orisize, ref_matte_fg_8uc3_orisize), axis=1)
|
||||
# resize_ratio = 1024. / max(show_concat.shape[:2])
|
||||
# show_concat = cv2.resize(show_concat, (0, 0), fx=resize_ratio, fy=resize_ratio)
|
||||
# print("gender_res: ", gender_res, " hair_ratio: ", hair_ratio)
|
||||
# cv2.imshow("show_concat", show_concat)
|
||||
# cv2.imshow("ref_matte_pred_8uc1_orisize", ref_matte_pred_8uc1_orisize)
|
||||
# cv2.waitKey()
|
||||
|
||||
ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768 = \
|
||||
self.process_data.get_prepare_ref_768_data(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize, ratio)
|
||||
|
||||
# cv2.imshow("ref_rgb_8uc3_768", ref_rgb_8uc3_768)
|
||||
# cv2.waitKey()
|
||||
|
||||
return ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768, gender, ratio
|
||||
|
||||
def Generator_reftensor(self, ref_rgb_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768,
|
||||
ref_landmark_f1k2_768):
|
||||
|
||||
"""
|
||||
input：
|
||||
|
||||
图像尺寸基于： 人脸 512， 图像均为3通道
|
||||
ref_rgb_8uc3_512: 参考图， size 512， uint8 （0-255）
|
||||
ref_matting_8uc3_512：参考图 matting alpha 值， uint8 （0-255）
|
||||
ref_baldseg_8uc3_512: 参考图 光头分割， uint8 （0-255）
|
||||
ref_landmark_f1k2_512： 参考图 关键点 1k*2 float32
|
||||
|
||||
output：
|
||||
|
||||
input_another_pose_hair_image: 参考图 条件图， float32 （0-255）
|
||||
|
||||
"""
|
||||
|
||||
# 8 ********************** another pose hair_image **********************
|
||||
another_pose_image = ref_rgb_8uc3_768.copy()
|
||||
|
||||
another_nohair_pose_mask = ref_baldseg_8uc3_768.copy()
|
||||
|
||||
# cv2.imshow("another_nohair_pose_mask", another_nohair_pose_mask)
|
||||
# cv2.waitKey()
|
||||
|
||||
another_nohair_pose_mask[(np.abs(another_nohair_pose_mask - [0, 0, 255]) < 50).all(axis=2)] = [0, 255, 255]
|
||||
another_nohair_pose_mask[(another_nohair_pose_mask == [0, 0, 0]).all(axis=2)] = [255, 255, 255]
|
||||
another_pose_pts137 = landmark_processor.pts_1k_to_137(ref_landmark_f1k2_768).astype(np.int32)
|
||||
cv2.fillPoly(another_nohair_pose_mask,
|
||||
np.concatenate((another_pose_pts137[47:35:-1], another_pose_pts137[56:64],
|
||||
[another_pose_pts137[48], another_pose_pts137[22]]))[
|
||||
np.newaxis, :, :], (255, 255, 0))
|
||||
|
||||
cv2.fillPoly(another_nohair_pose_mask,
|
||||
np.concatenate((another_pose_pts137[22:37], another_pose_pts137[56:47:-1]))[np.newaxis, :, :],
|
||||
(255, 0, 128))
|
||||
|
||||
cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[88:104][np.newaxis, :, :], (255, 0, 255)) # eye
|
||||
cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[105:121][np.newaxis, :, :], (255, 0, 255)) # eye
|
||||
# Label nose
|
||||
cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[64:79][np.newaxis, :, :], (255, 255, 255))
|
||||
# Label eyebrow
|
||||
cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[121:129][np.newaxis, :, :], (255, 128, 0))
|
||||
cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[129:137][np.newaxis, :, :], (255, 128, 0))
|
||||
|
||||
another_pose_hair_alpha = ref_matting_8uc3_768.copy() / 255.
|
||||
another_pose_hair_image = (another_pose_image * another_pose_hair_alpha + another_nohair_pose_mask * (
|
||||
1 - another_pose_hair_alpha)).astype(np.uint8)
|
||||
|
||||
input_another_pose_hair_image = another_pose_hair_image.astype(np.float32) / 255
|
||||
|
||||
return input_another_pose_hair_image
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,208 @@
|
||||
import os
|
||||
import cv2
|
||||
import argparse
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
from torch.nn import functional as F
|
||||
|
||||
import utils
|
||||
from core.matting import networks
|
||||
from utils.data_preprocess import *
|
||||
from time import time
|
||||
|
||||
def single_inference(model, image_dict, device, return_offset=False):
|
||||
|
||||
with torch.no_grad():
|
||||
image, trimap = image_dict['image'], image_dict['trimap']
|
||||
alpha_shape = image_dict['alpha_shape']
|
||||
image = image.to(device)
|
||||
trimap = trimap.to(device)
|
||||
|
||||
matte_start = time()
|
||||
alpha_pred, info_dict = model(image, trimap)
|
||||
matte_end = time()
|
||||
# print("matte time cost : {:.4f}".format(matte_end-matte_start))
|
||||
|
||||
trimap_argmax = trimap.argmax(dim=1, keepdim=True)
|
||||
|
||||
alpha_pred[trimap_argmax == 2] = 1
|
||||
alpha_pred[trimap_argmax == 0] = 0
|
||||
|
||||
h, w = alpha_shape
|
||||
test_pred = alpha_pred[0, 0, ...].detach().cpu().numpy() * 255
|
||||
test_pred = test_pred.astype(np.uint8)
|
||||
test_pred = test_pred[32:h+32, 32:w+32]
|
||||
|
||||
if return_offset:
|
||||
short_side = h if h < w else w
|
||||
ratio = 512 / short_side
|
||||
offset_1 = utils.flow_to_image(info_dict['offset_1'][0][0, ...].data.cpu().numpy()).astype(np.uint8)
|
||||
# write softmax_scale to offset image
|
||||
scale = info_dict['offset_1'][1].cpu()
|
||||
offset_1 = cv2.resize(offset_1, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_NEAREST)
|
||||
text = 'unknown: {:.2f}, known: {:.2f}'.format(scale[-1, 0].item(), scale[-1,1].item())
|
||||
offset_1 = cv2.putText(offset_1, text, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2)
|
||||
|
||||
offset_2 = utils.flow_to_image(info_dict['offset_2'][0][0, ...].data.cpu().numpy()).astype(np.uint8)
|
||||
# write softmax_scale to offset image
|
||||
scale = info_dict['offset_2'][1].cpu()
|
||||
offset_2 = cv2.resize(offset_2, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_NEAREST)
|
||||
text = 'unknown: {:.2f}, known: {:.2f}'.format(scale[-1, 0].item(), scale[-1,1].item())
|
||||
offset_2 = cv2.putText(offset_2, text, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2)
|
||||
|
||||
return test_pred, (offset_1, offset_2)
|
||||
else:
|
||||
return test_pred, None
|
||||
|
||||
def generator_tensor_dict(image, trimap):
|
||||
|
||||
sample = {'image': image, 'trimap': trimap, 'alpha_shape': trimap.shape}
|
||||
|
||||
# reshape
|
||||
h, w = sample["alpha_shape"]
|
||||
|
||||
if h % 32 == 0 and w % 32 == 0:
|
||||
padded_image = np.pad(sample['image'], ((32, 32), (32, 32), (0, 0)), mode="reflect")
|
||||
padded_trimap = np.pad(sample['trimap'], ((32, 32), (32, 32)), mode="reflect")
|
||||
sample['image'] = padded_image
|
||||
sample['trimap'] = padded_trimap
|
||||
else:
|
||||
target_h = 32 * ((h - 1) // 32 + 1)
|
||||
target_w = 32 * ((w - 1) // 32 + 1)
|
||||
pad_h = target_h - h
|
||||
pad_w = target_w - w
|
||||
padded_image = np.pad(sample['image'], ((32, pad_h+32), (32, pad_w+32), (0, 0)), mode="reflect")
|
||||
padded_trimap = np.pad(sample['trimap'], ((32, pad_h+32), (32, pad_w+32)), mode="reflect")
|
||||
sample['image'] = padded_image
|
||||
sample['trimap'] = padded_trimap
|
||||
|
||||
# ImageNet mean & std
|
||||
mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
|
||||
std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
|
||||
# convert GBR images to RGB
|
||||
image, trimap = sample['image'][:, :, ::-1], sample['trimap']
|
||||
# swap color axis
|
||||
image = image.transpose((2, 0, 1)).astype(np.float32)
|
||||
trimap[trimap < 85] = 0
|
||||
trimap[trimap >= 170] = 2
|
||||
trimap[trimap >= 85] = 1
|
||||
# normalize image
|
||||
image /= 255.
|
||||
|
||||
# to tensor
|
||||
sample['image'], sample['trimap'] = torch.from_numpy(image), torch.from_numpy(trimap).to(torch.long)
|
||||
sample['image'] = sample['image'].sub_(mean).div_(std)
|
||||
|
||||
# if CONFIG.model.trimap_channel == 3:
|
||||
sample['trimap'] = F.one_hot(sample['trimap'], num_classes=3).permute(2, 0, 1).float()
|
||||
# elif CONFIG.model.trimap_channel == 1:
|
||||
# sample['trimap'] = sample['trimap'][None, ...].float()
|
||||
# else:
|
||||
# raise NotImplementedError("CONFIG.model.trimap_channel can only be 3 or 1")
|
||||
|
||||
# add first channel
|
||||
sample['image'], sample['trimap'] = sample['image'][None, ...], sample['trimap'][None, ...]
|
||||
|
||||
return sample
|
||||
|
||||
if __name__ == '__main__':
|
||||
print('Torch Version: ', torch.__version__)
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--checkpoint', type=str, default='/home/liyang/project/matting/GCA-Matting/checkpoints/gca-dist-align-truth-1101/best_model.pth',
|
||||
help="path of checkpoint")
|
||||
parser.add_argument('--image-dir', type=str, default='/home/liyang/project/matting/合格/origin', help="input image dir")
|
||||
parser.add_argument('--trimap-dir', type=str, default='/home/liyang/project/matting/合格/trimap', help="input trimap dir")
|
||||
parser.add_argument('--output', type=str, default='/home/liyang/project/matting/合格/pred6', help="output dir")
|
||||
|
||||
# Parse configuration
|
||||
args = parser.parse_args()
|
||||
|
||||
# # Check if toml config file is loaded
|
||||
# if CONFIG.is_default:
|
||||
# raise ValueError("No .toml config loaded.")
|
||||
|
||||
args.output = os.path.join(args.output, args.checkpoint.split('/')[-1])
|
||||
utils.make_dir(args.output)
|
||||
|
||||
# build model
|
||||
model = networks.get_generator(encoder="resnet_gca_encoder_29", decoder="res_gca_decoder_22", num_class=1)
|
||||
model.cuda()
|
||||
print("model: ", model)
|
||||
|
||||
# load checkpoint
|
||||
checkpoint = torch.load(args.checkpoint)
|
||||
model.load_state_dict(utils.remove_prefix_state_dict(checkpoint['state_dict']), strict=True)
|
||||
|
||||
# inference
|
||||
model = model.eval()
|
||||
|
||||
for image_name in os.listdir(args.image_dir):
|
||||
|
||||
if not is_image_file(image_name):
|
||||
continue
|
||||
# assume image and trimap have the same file name
|
||||
img_basename, ext = os.path.splitext(image_name)
|
||||
|
||||
image_path = os.path.join(args.image_dir, image_name)
|
||||
trimap_path = os.path.join(args.trimap_dir, image_name.replace(ext, "_alpha.png"))
|
||||
# trimap_path = os.path.join(args.trimap_dir, image_name)
|
||||
print('Image: ', image_path, ' Tirmap: ', trimap_path)
|
||||
|
||||
# read images
|
||||
img_basename, img_ext = os.path.splitext(image_name)
|
||||
# img_pt_path = image_path.replace(img_ext, "_landmark1k.txt")
|
||||
# img_landmark1k = np.loadtxt(img_pt_path)
|
||||
|
||||
image = cv2.imread(image_path)
|
||||
trimap = cv2.imread(trimap_path, 0)
|
||||
|
||||
ori_h, ori_w, _ = image.shape
|
||||
tri_h, tri_w = trimap.shape
|
||||
if image.shape[1] != trimap.shape[1] or image.shape[0] != trimap.shape[0]:
|
||||
image = cv2.resize(image, (trimap.shape[1], trimap.shape[0]), interpolation=cv2.INTER_CUBIC)
|
||||
# img_landmark1k = landmark_processor.resize_points(img_landmark1k, ori_w, ori_h,
|
||||
# tri_w, tri_h)
|
||||
|
||||
# hair_mat = landmark_processor.get_transform_mat_hair(img_landmark1k, 640, ratio=0.3, w_ratio=0.5,
|
||||
# h_ratio=0.4)
|
||||
# hair_img_landmark = landmark_processor.transform_points(img_landmark1k, hair_mat)
|
||||
#
|
||||
# image = cv2.warpAffine(image, hair_mat, (640, 640), flags=cv2.INTER_CUBIC)
|
||||
# trimap = cv2.warpAffine(trimap, hair_mat, (640, 640), flags=cv2.INTER_CUBIC)
|
||||
|
||||
if tri_h > 1920 or tri_w > 1920:
|
||||
if tri_h > 1920:
|
||||
new_tri_h = 1920
|
||||
new_tri_w = int(tri_w * 1920 / tri_h)
|
||||
else:
|
||||
new_tri_w = 1920
|
||||
new_tri_h = int(tri_h * 1920 / tri_w)
|
||||
|
||||
image_resize = cv2.resize(image, (new_tri_w, new_tri_h), interpolation=cv2.INTER_CUBIC)
|
||||
trimap_resize = cv2.resize(trimap, (new_tri_w, new_tri_h), interpolation=cv2.INTER_NEAREST)
|
||||
else:
|
||||
image_resize = image
|
||||
trimap_resize = trimap
|
||||
|
||||
# image_dict = generator_tensor_dict(image, trimap)
|
||||
image_dict = generator_tensor_dict(image_resize, trimap_resize)
|
||||
pred, offset = single_inference(model, image_dict)
|
||||
|
||||
# torch.cuda.empty_cache()
|
||||
|
||||
pred = cv2.resize(pred, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_CUBIC)
|
||||
# pred = cv2.warpAffine(pred, cv2.invertAffineTransform(hair_mat), (tri_w, tri_h), flags=cv2.INTER_CUBIC)
|
||||
|
||||
# offset[0] = cv2.resize(offset[0], (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
|
||||
# offset[1] = cv2.resize(offset[1], (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
# cv2.imshow("image_resize", image_resize)
|
||||
# cv2.imshow("trimap_resize", trimap_resize)
|
||||
# cv2.imshow("pred", pred)
|
||||
# cv2.waitKey()
|
||||
cv2.imwrite(os.path.join(args.output, image_name.replace(img_ext, "_noalign.png")), pred)
|
||||
# if offset is not None:
|
||||
# cv2.imwrite(os.path.join(args.output, os.path.splitext(image_name)[0]+'_offset1.png'), offset[0])
|
||||
# cv2.imwrite(os.path.join(args.output, os.path.splitext(image_name)[0]+'_offset2.png'), offset[1])
|
||||
@@ -0,0 +1,200 @@
|
||||
import os
|
||||
import cv2
|
||||
import argparse
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
from torch.nn import functional as F
|
||||
|
||||
import utils
|
||||
from core.matting import networks
|
||||
from utils.data_preprocess import *
|
||||
from time import time
|
||||
|
||||
def single_inference(model, image_dict, device, return_offset=False):
|
||||
|
||||
with torch.no_grad():
|
||||
image, trimap = image_dict['image'], image_dict['trimap']
|
||||
alpha_shape = image_dict['alpha_shape']
|
||||
image = image.to(device)
|
||||
trimap = trimap.to(device)
|
||||
|
||||
matte_start = time()
|
||||
|
||||
alpha_pred, info_dict = model(image, trimap)
|
||||
|
||||
fg_pred = alpha_pred[:, :-1, :, :]
|
||||
alpha_pred = alpha_pred[:, -1, :, :].unsqueeze(1)
|
||||
matte_end = time()
|
||||
# print("matte time cost : {:.4f}".format(matte_end-matte_start))
|
||||
|
||||
trimap_argmax = trimap.argmax(dim=1, keepdim=True)
|
||||
|
||||
alpha_pred[trimap_argmax == 2] = 1
|
||||
alpha_pred[trimap_argmax == 0] = 0
|
||||
|
||||
h, w = alpha_shape
|
||||
test_fg_pred = fg_pred[0].detach().cpu().numpy().transpose(1, 2, 0)[:, :, ::-1] * 255
|
||||
test_fg_pred = test_fg_pred.astype(np.uint8)
|
||||
test_fg_pred = test_fg_pred[32:h+32, 32:w+32]
|
||||
|
||||
test_pred = alpha_pred[0, 0, ...].detach().cpu().numpy() * 255
|
||||
test_pred = test_pred.astype(np.uint8)
|
||||
test_pred = test_pred[32:h+32, 32:w+32]
|
||||
|
||||
# cv2.imshow('test_fg_pred', test_fg_pred)
|
||||
# cv2.imshow('test_pred', test_pred)
|
||||
# cv2.waitKey()
|
||||
|
||||
if return_offset:
|
||||
short_side = h if h < w else w
|
||||
ratio = 512 / short_side
|
||||
offset_1 = utils.flow_to_image(info_dict['offset_1'][0][0, ...].data.cpu().numpy()).astype(np.uint8)
|
||||
# write softmax_scale to offset image
|
||||
scale = info_dict['offset_1'][1].cpu()
|
||||
offset_1 = cv2.resize(offset_1, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_NEAREST)
|
||||
text = 'unknown: {:.2f}, known: {:.2f}'.format(scale[-1, 0].item(), scale[-1,1].item())
|
||||
offset_1 = cv2.putText(offset_1, text, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2)
|
||||
|
||||
offset_2 = utils.flow_to_image(info_dict['offset_2'][0][0, ...].data.cpu().numpy()).astype(np.uint8)
|
||||
# write softmax_scale to offset image
|
||||
scale = info_dict['offset_2'][1].cpu()
|
||||
offset_2 = cv2.resize(offset_2, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_NEAREST)
|
||||
text = 'unknown: {:.2f}, known: {:.2f}'.format(scale[-1, 0].item(), scale[-1,1].item())
|
||||
offset_2 = cv2.putText(offset_2, text, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2)
|
||||
|
||||
return test_fg_pred, test_pred, (offset_1, offset_2)
|
||||
else:
|
||||
return test_fg_pred, test_pred, None
|
||||
|
||||
def generator_tensor_dict(image, trimap):
|
||||
|
||||
sample = {'image': image, 'trimap': trimap, 'alpha_shape': trimap.shape}
|
||||
|
||||
# reshape
|
||||
h, w = sample["alpha_shape"]
|
||||
|
||||
if h % 32 == 0 and w % 32 == 0:
|
||||
padded_image = np.pad(sample['image'], ((32, 32), (32, 32), (0, 0)), mode="reflect")
|
||||
padded_trimap = np.pad(sample['trimap'], ((32, 32), (32, 32)), mode="reflect")
|
||||
sample['image'] = padded_image
|
||||
sample['trimap'] = padded_trimap
|
||||
else:
|
||||
target_h = 32 * ((h - 1) // 32 + 1)
|
||||
target_w = 32 * ((w - 1) // 32 + 1)
|
||||
pad_h = target_h - h
|
||||
pad_w = target_w - w
|
||||
padded_image = np.pad(sample['image'], ((32, pad_h+32), (32, pad_w+32), (0, 0)), mode="reflect")
|
||||
padded_trimap = np.pad(sample['trimap'], ((32, pad_h+32), (32, pad_w+32)), mode="reflect")
|
||||
sample['image'] = padded_image
|
||||
sample['trimap'] = padded_trimap
|
||||
|
||||
# ImageNet mean & std
|
||||
mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
|
||||
std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
|
||||
# convert GBR images to RGB
|
||||
image, trimap = sample['image'][:, :, ::-1], sample['trimap']
|
||||
# swap color axis
|
||||
image = image.transpose((2, 0, 1)).astype(np.float32)
|
||||
trimap[trimap < 85] = 0
|
||||
trimap[trimap >= 170] = 2
|
||||
trimap[trimap >= 85] = 1
|
||||
# normalize image
|
||||
image /= 255.
|
||||
|
||||
# to tensor
|
||||
sample['image'], sample['trimap'] = torch.from_numpy(image), torch.from_numpy(trimap).to(torch.long)
|
||||
sample['image'] = sample['image'].sub_(mean).div_(std)
|
||||
|
||||
sample['trimap'] = F.one_hot(sample['trimap'], num_classes=3).permute(2, 0, 1).float()
|
||||
|
||||
# add first channel
|
||||
sample['image'], sample['trimap'] = sample['image'][None, ...], sample['trimap'][None, ...]
|
||||
|
||||
return sample
|
||||
|
||||
if __name__ == '__main__':
|
||||
print('Torch Version: ', torch.__version__)
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--checkpoint', type=str, default='/home/liyang/project/matting/GCA-Matting/checkpoints/gca-dist-align-truth-1101/best_model.pth',
|
||||
help="path of checkpoint")
|
||||
parser.add_argument('--image-dir', type=str, default='/home/liyang/project/matting/合格/origin', help="input image dir")
|
||||
parser.add_argument('--trimap-dir', type=str, default='/home/liyang/project/matting/合格/trimap', help="input trimap dir")
|
||||
parser.add_argument('--output', type=str, default='/home/liyang/project/matting/合格/pred6', help="output dir")
|
||||
|
||||
# Parse configuration
|
||||
args = parser.parse_args()
|
||||
|
||||
# # Check if toml config file is loaded
|
||||
# if CONFIG.is_default:
|
||||
# raise ValueError("No .toml config loaded.")
|
||||
|
||||
args.output = os.path.join(args.output, args.checkpoint.split('/')[-1])
|
||||
utils.make_dir(args.output)
|
||||
|
||||
# build model
|
||||
model = networks.get_generator(encoder=CONFIG.model.arch.encoder, decoder=CONFIG.model.arch.decoder)
|
||||
model.cuda()
|
||||
print("model: ", model)
|
||||
|
||||
# load checkpoint
|
||||
checkpoint = torch.load(args.checkpoint)
|
||||
model.load_state_dict(utils.remove_prefix_state_dict(checkpoint['state_dict']), strict=True)
|
||||
|
||||
# inference
|
||||
model = model.eval()
|
||||
export_onnx_file = "test.onnx"
|
||||
torch.onnx.export(model, x, export_onnx_file, opset_version=10, do_constant_folding=True, input_names=["image", "trimap"], # 输入名
|
||||
output_names=["fg_pred", "alpha_pred", "None"])
|
||||
|
||||
for image_name in os.listdir(args.image_dir):
|
||||
|
||||
if not is_image_file(image_name):
|
||||
continue
|
||||
# assume image and trimap have the same file name
|
||||
img_basename, ext = os.path.splitext(image_name)
|
||||
|
||||
image_path = os.path.join(args.image_dir, image_name)
|
||||
trimap_path = os.path.join(args.trimap_dir, image_name.replace(ext, "_alpha.png"))
|
||||
# trimap_path = os.path.join(args.trimap_dir, image_name)
|
||||
print('Image: ', image_path, ' Tirmap: ', trimap_path)
|
||||
|
||||
# read images
|
||||
img_basename, img_ext = os.path.splitext(image_name)
|
||||
|
||||
image = cv2.imread(image_path)
|
||||
trimap = cv2.imread(trimap_path, 0)
|
||||
|
||||
ori_h, ori_w, _ = image.shape
|
||||
tri_h, tri_w = trimap.shape
|
||||
if image.shape[1] != trimap.shape[1] or image.shape[0] != trimap.shape[0]:
|
||||
image = cv2.resize(image, (trimap.shape[1], trimap.shape[0]), interpolation=cv2.INTER_CUBIC)
|
||||
|
||||
|
||||
if tri_h > 1920 or tri_w > 1920:
|
||||
if tri_h > 1920:
|
||||
new_tri_h = 1920
|
||||
new_tri_w = int(tri_w * 1920 / tri_h)
|
||||
else:
|
||||
new_tri_w = 1920
|
||||
new_tri_h = int(tri_h * 1920 / tri_w)
|
||||
|
||||
image_resize = cv2.resize(image, (new_tri_w, new_tri_h), interpolation=cv2.INTER_CUBIC)
|
||||
trimap_resize = cv2.resize(trimap, (new_tri_w, new_tri_h), interpolation=cv2.INTER_NEAREST)
|
||||
else:
|
||||
image_resize = image
|
||||
trimap_resize = trimap
|
||||
|
||||
# image_dict = generator_tensor_dict(image, trimap)
|
||||
image_dict = generator_tensor_dict(image_resize, trimap_resize)
|
||||
fg_pred, pred, offset = single_inference(model, image_dict)
|
||||
|
||||
# torch.cuda.empty_cache()
|
||||
|
||||
fg_pred = cv2.resize(fg_pred, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_CUBIC)
|
||||
pred = cv2.resize(pred, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_CUBIC)
|
||||
|
||||
cv2.imwrite(os.path.join(args.output, image_name.replace(img_ext, "_noalign.png")), pred)
|
||||
cv2.imwrite(os.path.join(args.output, image_name.replace(img_ext, "_fg_p.png")), fg_pred)
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
from .generators import *
|
||||
@@ -0,0 +1,28 @@
|
||||
from .resnet_dec import ResNet_D_Dec, BasicBlock
|
||||
from .res_shortcut_dec import ResShortCut_D_Dec
|
||||
from .res_gca_dec import ResGuidedCxtAtten_Dec
|
||||
|
||||
|
||||
__all__ = ['res_shortcut_decoder_22', 'res_gca_decoder_22']
|
||||
|
||||
|
||||
def _res_shortcut_D_dec(block, layers, **kwargs):
|
||||
model = ResShortCut_D_Dec(block, layers, **kwargs)
|
||||
return model
|
||||
|
||||
|
||||
def _res_gca_D_dec(block, layers, num_class, **kwargs):
|
||||
model = ResGuidedCxtAtten_Dec(block, layers, num_class, **kwargs)
|
||||
return model
|
||||
|
||||
|
||||
def res_shortcut_decoder_22(**kwargs):
|
||||
"""Constructs a resnet_encoder_14 model.
|
||||
"""
|
||||
return _res_shortcut_D_dec(BasicBlock, [2, 3, 3, 2], **kwargs)
|
||||
|
||||
|
||||
def res_gca_decoder_22(num_class=1, **kwargs):
|
||||
"""Constructs a resnet_encoder_14 model.
|
||||
"""
|
||||
return _res_gca_D_dec(BasicBlock, [2, 3, 3, 2], num_class, **kwargs)
|
||||
@@ -0,0 +1,28 @@
|
||||
from core.matting.networks.ops import GuidedCxtAtten
|
||||
from core.matting.networks.decoders.res_shortcut_dec import ResShortCut_D_Dec
|
||||
|
||||
|
||||
class ResGuidedCxtAtten_Dec(ResShortCut_D_Dec):
|
||||
|
||||
def __init__(self, block, layers, num_class=1, norm_layer=None, large_kernel=False):
|
||||
super(ResGuidedCxtAtten_Dec, self).__init__(block, layers, num_class, norm_layer, large_kernel)
|
||||
self.gca = GuidedCxtAtten(128, 128)
|
||||
self.num_class = num_class
|
||||
|
||||
def forward(self, x, mid_fea):
|
||||
fea1, fea2, fea3, fea4, fea5 = mid_fea['shortcut']
|
||||
im = mid_fea['image_fea']
|
||||
x = self.layer1(x) + fea5 # N x 256 x 32 x 32
|
||||
x = self.layer2(x) + fea4 # N x 128 x 64 x 64
|
||||
x, offset = self.gca(im, x, mid_fea['unknown']) # contextual attention
|
||||
x = self.layer3(x) + fea3 # N x 64 x 128 x 128
|
||||
x = self.layer4(x) + fea2 # N x 32 x 256 x 256
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.leaky_relu(x) + fea1
|
||||
x = self.conv2(x)
|
||||
|
||||
alpha = (self.tanh(x) + 1.0) / 2.0
|
||||
|
||||
return alpha, {'offset_1': mid_fea['offset_1'], 'offset_2': offset}
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
from core.matting.networks.decoders.resnet_dec import ResNet_D_Dec
|
||||
|
||||
|
||||
class ResShortCut_D_Dec(ResNet_D_Dec):
|
||||
|
||||
def __init__(self, block, layers, num_class=1, norm_layer=None, large_kernel=False, late_downsample=False):
|
||||
super(ResShortCut_D_Dec, self).__init__(block, layers, num_class, norm_layer, large_kernel,
|
||||
late_downsample=late_downsample)
|
||||
|
||||
def forward(self, x, mid_fea):
|
||||
fea1, fea2, fea3, fea4, fea5 = mid_fea['shortcut']
|
||||
x = self.layer1(x) + fea5
|
||||
x = self.layer2(x) + fea4
|
||||
x = self.layer3(x) + fea3
|
||||
x = self.layer4(x) + fea2
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.leaky_relu(x) + fea1
|
||||
x = self.conv2(x)
|
||||
|
||||
alpha = (self.tanh(x) + 1.0) / 2.0
|
||||
|
||||
return alpha, None
|
||||
|
||||
@@ -0,0 +1,142 @@
|
||||
import logging
|
||||
import torch.nn as nn
|
||||
from core.matting.networks.ops import SpectralNorm
|
||||
|
||||
def conv5x5(in_planes, out_planes, stride=1, groups=1, dilation=1):
|
||||
"""5x5 convolution with padding"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=5, stride=stride,
|
||||
padding=2, groups=groups, bias=False, dilation=dilation)
|
||||
|
||||
|
||||
def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
|
||||
"""3x3 convolution with padding"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
|
||||
padding=dilation, groups=groups, bias=False, dilation=dilation)
|
||||
|
||||
|
||||
def conv1x1(in_planes, out_planes, stride=1):
|
||||
"""1x1 convolution"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
|
||||
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
expansion = 1
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, upsample=None, norm_layer=None, large_kernel=False):
|
||||
super(BasicBlock, self).__init__()
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
self.stride = stride
|
||||
conv = conv5x5 if large_kernel else conv3x3
|
||||
# Both self.conv1 and self.downsample layers downsample the input when stride != 1
|
||||
if self.stride > 1:
|
||||
self.conv1 = SpectralNorm(nn.ConvTranspose2d(inplanes, inplanes, kernel_size=4, stride=2, padding=1, bias=False))
|
||||
else:
|
||||
self.conv1 = SpectralNorm(conv(inplanes, inplanes))
|
||||
self.bn1 = norm_layer(inplanes)
|
||||
self.activation = nn.LeakyReLU(0.2, inplace=True)
|
||||
self.conv2 = SpectralNorm(conv(inplanes, planes))
|
||||
self.bn2 = norm_layer(planes)
|
||||
self.upsample = upsample
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.activation(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.upsample is not None:
|
||||
identity = self.upsample(x)
|
||||
|
||||
out += identity
|
||||
out = self.activation(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class ResNet_D_Dec(nn.Module):
|
||||
|
||||
def __init__(self, block, layers, num_class=1, norm_layer=None, large_kernel=False, late_downsample=False):
|
||||
super(ResNet_D_Dec, self).__init__()
|
||||
self.logger = logging.getLogger("Logger")
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
self._norm_layer = norm_layer
|
||||
self.large_kernel = large_kernel
|
||||
self.kernel_size = 5 if self.large_kernel else 3
|
||||
|
||||
self.inplanes = 512 if layers[0] > 0 else 256
|
||||
self.late_downsample = late_downsample
|
||||
self.midplanes = 64 if late_downsample else 32
|
||||
|
||||
self.conv1 = SpectralNorm(nn.ConvTranspose2d(self.midplanes, 32, kernel_size=4, stride=2, padding=1, bias=False))
|
||||
self.bn1 = norm_layer(32)
|
||||
self.leaky_relu = nn.LeakyReLU(0.2, inplace=True)
|
||||
self.conv2 = nn.Conv2d(32, num_class, kernel_size=self.kernel_size, stride=1, padding=self.kernel_size//2)
|
||||
self.upsample = nn.UpsamplingNearest2d(scale_factor=2)
|
||||
self.tanh = nn.Tanh()
|
||||
self.layer1 = self._make_layer(block, 256, layers[0], stride=2)
|
||||
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
|
||||
self.layer3 = self._make_layer(block, 64, layers[2], stride=2)
|
||||
self.layer4 = self._make_layer(block, self.midplanes, layers[3], stride=2)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
if hasattr(m, "weight_bar"):
|
||||
nn.init.xavier_uniform_(m.weight_bar)
|
||||
else:
|
||||
nn.init.xavier_uniform_(m.weight)
|
||||
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
# Zero-initialize the last BN in each residual branch,
|
||||
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
|
||||
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
|
||||
for m in self.modules():
|
||||
if isinstance(m, BasicBlock):
|
||||
nn.init.constant_(m.bn2.weight, 0)
|
||||
|
||||
self.logger.debug(self)
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1):
|
||||
if blocks == 0:
|
||||
return nn.Sequential(nn.Identity())
|
||||
norm_layer = self._norm_layer
|
||||
upsample = None
|
||||
if stride != 1:
|
||||
upsample = nn.Sequential(
|
||||
nn.UpsamplingNearest2d(scale_factor=2),
|
||||
SpectralNorm(conv1x1(self.inplanes, planes * block.expansion)),
|
||||
norm_layer(planes * block.expansion),
|
||||
)
|
||||
elif self.inplanes != planes * block.expansion:
|
||||
upsample = nn.Sequential(
|
||||
SpectralNorm(conv1x1(self.inplanes, planes * block.expansion)),
|
||||
norm_layer(planes * block.expansion),
|
||||
)
|
||||
|
||||
layers = [block(self.inplanes, planes, stride, upsample, norm_layer, self.large_kernel)]
|
||||
self.inplanes = planes * block.expansion
|
||||
for _ in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes, norm_layer=norm_layer, large_kernel=self.large_kernel))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x, mid_fea):
|
||||
x = self.layer1(x) # N x 256 x 32 x 32
|
||||
x = self.layer2(x) # N x 128 x 64 x 64
|
||||
x = self.layer3(x) # N x 64 x 128 x 128
|
||||
x = self.layer4(x) # N x 32 x 256 x 256
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.leaky_relu(x)
|
||||
x = self.conv2(x)
|
||||
|
||||
alpha = (self.tanh(x) + 1.0) / 2.0
|
||||
|
||||
return alpha, None
|
||||
@@ -0,0 +1,39 @@
|
||||
import logging
|
||||
from .resnet_enc import ResNet_D, BasicBlock
|
||||
from .res_shortcut_enc import ResShortCut_D
|
||||
from .res_gca_enc import ResGuidedCxtAtten
|
||||
|
||||
|
||||
__all__ = ['res_shortcut_encoder_29', 'resnet_gca_encoder_29']
|
||||
|
||||
|
||||
def _res_shortcut_D(block, layers, **kwargs):
|
||||
model = ResShortCut_D(block, layers, **kwargs)
|
||||
return model
|
||||
|
||||
|
||||
def _res_gca_D(block, layers, **kwargs):
|
||||
model = ResGuidedCxtAtten(block, layers, **kwargs)
|
||||
return model
|
||||
|
||||
|
||||
def resnet_gca_encoder_29(**kwargs):
|
||||
"""Constructs a resnet_encoder_29 model.
|
||||
"""
|
||||
return _res_gca_D(BasicBlock, [3, 4, 4, 2], **kwargs)
|
||||
|
||||
|
||||
def res_shortcut_encoder_29(**kwargs):
|
||||
"""Constructs a resnet_encoder_25 model.
|
||||
"""
|
||||
return _res_shortcut_D(BasicBlock, [3, 4, 4, 2], **kwargs)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import torch
|
||||
logging.basicConfig(level=logging.DEBUG, format='[%(asctime)s] %(levelname)s: %(message)s',
|
||||
datefmt='%m-%d %H:%M:%S')
|
||||
resnet_encoder = res_shortcut_encoder_29()
|
||||
x = torch.randn(4,6,512,512)
|
||||
z = resnet_encoder(x)
|
||||
print(z[0].shape)
|
||||
@@ -0,0 +1,97 @@
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
# from utils import CONFIG
|
||||
from core.matting.networks.encoders.resnet_enc import ResNet_D
|
||||
from core.matting.networks.ops import GuidedCxtAtten, SpectralNorm
|
||||
|
||||
class ResGuidedCxtAtten(ResNet_D):
|
||||
|
||||
def __init__(self, block, layers, norm_layer=None, late_downsample=False):
|
||||
super(ResGuidedCxtAtten, self).__init__(block, layers, norm_layer, late_downsample=late_downsample)
|
||||
first_inplane = 3 + 3
|
||||
self.shortcut_inplane = [first_inplane, self.midplanes, 64, 128, 256]
|
||||
self.shortcut_plane = [32, self.midplanes, 64, 128, 256]
|
||||
|
||||
self.shortcut = nn.ModuleList()
|
||||
for stage, inplane in enumerate(self.shortcut_inplane):
|
||||
self.shortcut.append(self._make_shortcut(inplane, self.shortcut_plane[stage]))
|
||||
|
||||
self.guidance_head = nn.Sequential(
|
||||
nn.ReflectionPad2d(1),
|
||||
SpectralNorm(nn.Conv2d(3, 16, kernel_size=3, padding=0, stride=2, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(16),
|
||||
nn.ReflectionPad2d(1),
|
||||
SpectralNorm(nn.Conv2d(16, 32, kernel_size=3, padding=0, stride=2, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(32),
|
||||
nn.ReflectionPad2d(1),
|
||||
SpectralNorm(nn.Conv2d(32, 128, kernel_size=3, padding=0, stride=2, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(128)
|
||||
)
|
||||
|
||||
self.gca = GuidedCxtAtten(128, 128)
|
||||
|
||||
# initialize guidance head
|
||||
for layers in range(len(self.guidance_head)):
|
||||
m = self.guidance_head[layers]
|
||||
if isinstance(m, nn.Conv2d):
|
||||
if hasattr(m, "weight_bar"):
|
||||
nn.init.xavier_uniform_(m.weight_bar)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def _make_shortcut(self, inplane, planes):
|
||||
return nn.Sequential(
|
||||
SpectralNorm(nn.Conv2d(inplane, planes, kernel_size=3, padding=1, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(planes),
|
||||
SpectralNorm(nn.Conv2d(planes, planes, kernel_size=3, padding=1, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(planes)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.activation(out)
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
x1 = self.activation(out) # N x 32 x 256 x 256
|
||||
out = self.conv3(x1)
|
||||
out = self.bn3(out)
|
||||
out = self.activation(out)
|
||||
|
||||
im_fea = self.guidance_head(x[:, :3, ...]) # downsample origin image and extract features
|
||||
# if CONFIG.model.trimap_channel == 3:
|
||||
unknown = F.interpolate(x[:, 4:5, ...], scale_factor=1/8, mode='nearest')
|
||||
# else:
|
||||
# unknown = F.interpolate(x[:,3:,...].eq(1.).float(), scale_factor=1/8, mode='nearest')
|
||||
|
||||
x2 = self.layer1(out) # N x 64 x 128 x 128
|
||||
x3= self.layer2(x2) # N x 128 x 64 x 64
|
||||
x3, offset = self.gca(im_fea, x3, unknown) # contextual attention
|
||||
x4 = self.layer3(x3) # N x 256 x 32 x 32
|
||||
out = self.layer_bottleneck(x4) # N x 512 x 16 x 16
|
||||
|
||||
fea1 = self.shortcut[0](x) # input image and trimap
|
||||
fea2 = self.shortcut[1](x1)
|
||||
fea3 = self.shortcut[2](x2)
|
||||
fea4 = self.shortcut[3](x3)
|
||||
fea5 = self.shortcut[4](x4)
|
||||
|
||||
return out, {'shortcut': (fea1, fea2, fea3, fea4, fea5),
|
||||
'image_fea': im_fea,
|
||||
'unknown': unknown,
|
||||
'offset_1': offset}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from core.matting.networks.encoders.resnet_enc import BasicBlock
|
||||
m = ResGuidedCxtAtten(BasicBlock, [3, 4, 4, 2])
|
||||
for m in m.modules():
|
||||
print(m)
|
||||
@@ -0,0 +1,51 @@
|
||||
import torch.nn as nn
|
||||
# from utils import CONFIG
|
||||
from core.matting.networks.encoders.resnet_enc import ResNet_D
|
||||
from core.matting.networks.ops import SpectralNorm
|
||||
|
||||
class ResShortCut_D(ResNet_D):
|
||||
|
||||
def __init__(self, block, layers, norm_layer=None, late_downsample=False):
|
||||
super(ResShortCut_D, self).__init__(block, layers, norm_layer, late_downsample=late_downsample)
|
||||
first_inplane = 3 + 3
|
||||
self.shortcut_inplane = [first_inplane, self.midplanes, 64, 128, 256]
|
||||
self.shortcut_plane = [32, self.midplanes, 64, 128, 256]
|
||||
|
||||
self.shortcut = nn.ModuleList()
|
||||
for stage, inplane in enumerate(self.shortcut_inplane):
|
||||
self.shortcut.append(self._make_shortcut(inplane, self.shortcut_plane[stage]))
|
||||
|
||||
def _make_shortcut(self, inplane, planes):
|
||||
return nn.Sequential(
|
||||
SpectralNorm(nn.Conv2d(inplane, planes, kernel_size=3, padding=1, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(planes),
|
||||
SpectralNorm(nn.Conv2d(planes, planes, kernel_size=3, padding=1, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(planes)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.activation(out)
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
x1 = self.activation(out) # N x 32 x 256 x 256
|
||||
out = self.conv3(x1)
|
||||
out = self.bn3(out)
|
||||
out = self.activation(out)
|
||||
|
||||
x2 = self.layer1(out) # N x 64 x 128 x 128
|
||||
x3= self.layer2(x2) # N x 128 x 64 x 64
|
||||
x4 = self.layer3(x3) # N x 256 x 32 x 32
|
||||
out = self.layer_bottleneck(x4) # N x 512 x 16 x 16
|
||||
|
||||
fea1 = self.shortcut[0](x) # input image and trimap
|
||||
fea2 = self.shortcut[1](x1)
|
||||
fea3 = self.shortcut[2](x2)
|
||||
fea4 = self.shortcut[3](x3)
|
||||
fea5 = self.shortcut[4](x4)
|
||||
|
||||
return out, {'shortcut':(fea1, fea2, fea3, fea4, fea5), 'image':x[:,:3,...]}
|
||||
@@ -0,0 +1,150 @@
|
||||
import logging
|
||||
import torch.nn as nn
|
||||
from core.matting.networks.ops import SpectralNorm
|
||||
|
||||
def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
|
||||
"""3x3 convolution with padding"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
|
||||
padding=dilation, groups=groups, bias=False, dilation=dilation)
|
||||
|
||||
|
||||
def conv1x1(in_planes, out_planes, stride=1):
|
||||
"""1x1 convolution"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
|
||||
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
expansion = 1
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None, norm_layer=None):
|
||||
super(BasicBlock, self).__init__()
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
# Both self.conv1 and self.downsample layers downsample the input when stride != 1
|
||||
self.conv1 = SpectralNorm(conv3x3(inplanes, planes, stride))
|
||||
self.bn1 = norm_layer(planes)
|
||||
self.activation = nn.ReLU(inplace=True)
|
||||
self.conv2 = SpectralNorm(conv3x3(planes, planes))
|
||||
self.bn2 = norm_layer(planes)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.activation(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
identity = self.downsample(x)
|
||||
|
||||
out += identity
|
||||
out = self.activation(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class ResNet_D(nn.Module):
|
||||
"""
|
||||
Implement and pre-train on ImageNet with the tricks from
|
||||
https://arxiv.org/abs/1812.01187
|
||||
without the mix-up part.
|
||||
"""
|
||||
|
||||
def __init__(self, block, layers, norm_layer=None, late_downsample=False):
|
||||
super(ResNet_D, self).__init__()
|
||||
self.logger = logging.getLogger("Logger")
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
self._norm_layer = norm_layer
|
||||
|
||||
self.inplanes = 64
|
||||
self.late_downsample = late_downsample
|
||||
self.midplanes = 64 if late_downsample else 32
|
||||
self.start_stride = [1, 2, 1, 2] if late_downsample else [2, 1, 2, 1]
|
||||
self.conv1 = SpectralNorm(nn.Conv2d(3 + 3, 32, kernel_size=3,
|
||||
stride=self.start_stride[0], padding=1, bias=False))
|
||||
self.conv2 = SpectralNorm(nn.Conv2d(32, self.midplanes, kernel_size=3, stride=self.start_stride[1], padding=1,
|
||||
bias=False))
|
||||
self.conv3 = SpectralNorm(nn.Conv2d(self.midplanes, self.inplanes, kernel_size=3, stride=self.start_stride[2],
|
||||
padding=1, bias=False))
|
||||
self.bn1 = norm_layer(32)
|
||||
self.bn2 = norm_layer(self.midplanes)
|
||||
self.bn3 = norm_layer(self.inplanes)
|
||||
self.activation = nn.ReLU(inplace=True)
|
||||
self.layer1 = self._make_layer(block, 64, layers[0], stride=self.start_stride[3])
|
||||
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
|
||||
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
|
||||
self.layer_bottleneck = self._make_layer(block, 512, layers[3], stride=2)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.xavier_uniform_(m.weight_bar)
|
||||
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
# Zero-initialize the last BN in each residual branch,
|
||||
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
|
||||
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, BasicBlock):
|
||||
nn.init.constant_(m.bn2.weight, 0)
|
||||
|
||||
self.logger.debug("encoder conv1 weight shape: {}".format(str(self.conv1.module.weight_bar.data.shape)))
|
||||
self.conv1.module.weight_bar.data[:,3:,:,:] = 0
|
||||
|
||||
self.logger.debug(self)
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1):
|
||||
if blocks == 0:
|
||||
return nn.Sequential(nn.Identity())
|
||||
norm_layer = self._norm_layer
|
||||
downsample = None
|
||||
if stride != 1:
|
||||
downsample = nn.Sequential(
|
||||
nn.AvgPool2d(2, stride),
|
||||
SpectralNorm(conv1x1(self.inplanes, planes * block.expansion)),
|
||||
norm_layer(planes * block.expansion),
|
||||
)
|
||||
elif self.inplanes != planes * block.expansion:
|
||||
downsample = nn.Sequential(
|
||||
SpectralNorm(conv1x1(self.inplanes, planes * block.expansion, stride)),
|
||||
norm_layer(planes * block.expansion),
|
||||
)
|
||||
|
||||
layers = [block(self.inplanes, planes, stride, downsample, norm_layer)]
|
||||
self.inplanes = planes * block.expansion
|
||||
for _ in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes, norm_layer=norm_layer))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.activation(x)
|
||||
x = self.conv2(x)
|
||||
x = self.bn2(x)
|
||||
x1 = self.activation(x) # N x 32 x 256 x 256
|
||||
x = self.conv3(x1)
|
||||
x = self.bn3(x)
|
||||
x2 = self.activation(x) # N x 64 x 128 x 128
|
||||
|
||||
x3 = self.layer1(x2) # N x 64 x 128 x 128
|
||||
x4 = self.layer2(x3) # N x 128 x 64 x 64
|
||||
x5 = self.layer3(x4) # N x 256 x 32 x 32
|
||||
x = self.layer_bottleneck(x5) # N x 512 x 16 x 16
|
||||
|
||||
return x, (x1, x2, x3, x4, x5)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
m = ResNet_D(BasicBlock, [3, 4, 4, 2])
|
||||
for m in m.modules():
|
||||
print(m._get_name())
|
||||
@@ -0,0 +1,60 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
# from utils import CONFIG
|
||||
from core.matting.networks import decoders, encoders
|
||||
|
||||
|
||||
class Generator(nn.Module):
|
||||
def __init__(self, encoder, decoder, num_class=1):
|
||||
|
||||
super(Generator, self).__init__()
|
||||
|
||||
if encoder not in encoders.__all__:
|
||||
raise NotImplementedError("Unknown Encoder {}".format(encoder))
|
||||
self.encoder = encoders.__dict__[encoder]()
|
||||
|
||||
if decoder not in decoders.__all__:
|
||||
raise NotImplementedError("Unknown Decoder {}".format(decoder))
|
||||
self.decoder = decoders.__dict__[decoder](num_class)
|
||||
|
||||
def forward(self, image, trimap):
|
||||
inp = torch.cat((image, trimap), dim=1)
|
||||
embedding, mid_fea = self.encoder(inp)
|
||||
alpha, info_dict = self.decoder(embedding, mid_fea)
|
||||
|
||||
return alpha, info_dict
|
||||
|
||||
|
||||
def get_generator(encoder, decoder, num_class=1):
|
||||
generator = Generator(encoder=encoder, decoder=decoder, num_class=num_class)
|
||||
return generator
|
||||
|
||||
|
||||
if __name__=="__main__":
|
||||
import time
|
||||
# generator = get_generator(encoder=CONFIG.model.arch.encoder, decoder=CONFIG.model.arch.decoder).cuda().train()
|
||||
batch_size = 12
|
||||
# generator.eval()
|
||||
n_eval = 10
|
||||
# pre run the model
|
||||
# with torch.no_grad():
|
||||
# for i in range(2):
|
||||
# x = torch.rand(batch_size, 3, 512, 512, device=device)
|
||||
# y = torch.rand(batch_size, 3, 512, 512, device=device)
|
||||
# z = generator(x,y)
|
||||
# test without GPU IO
|
||||
|
||||
# x = torch.zeros(batch_size, 3, 512, 512, device=device)
|
||||
# y = torch.zeros(batch_size, 1, 512, 512, device=device)
|
||||
x = torch.randn(batch_size, 3, 512, 512)
|
||||
y = torch.randn(batch_size, 3, 512, 512)
|
||||
t = time.time()
|
||||
# with torch.no_grad():
|
||||
# for i in range(n_eval):
|
||||
# a = generator(x.cuda(),y.cuda())
|
||||
# torch.cuda.synchronize()
|
||||
# print(generator.__class__.__name__, 'With IO \t', f'{(time.time() - t)/n_eval/batch_size:.5f} s')
|
||||
# print(generator.__class__.__name__, 'FPS \t\t', f'{1/((time.time() - t)/n_eval/batch_size):.5f} s')
|
||||
# for n, p in generator.named_parameters():
|
||||
# print(n)
|
||||
@@ -0,0 +1,256 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import Parameter
|
||||
from torch.autograd import Variable
|
||||
from torch.nn import functional as F
|
||||
|
||||
|
||||
def l2normalize(v, eps=1e-12):
|
||||
return v / (v.norm() + eps)
|
||||
|
||||
|
||||
class SpectralNorm(nn.Module):
|
||||
"""
|
||||
Based on https://github.com/heykeetae/Self-Attention-GAN/blob/master/spectral.py
|
||||
and add _noupdate_u_v() for evaluation
|
||||
"""
|
||||
def __init__(self, module, name='weight', power_iterations=1):
|
||||
super(SpectralNorm, self).__init__()
|
||||
self.module = module
|
||||
self.name = name
|
||||
self.power_iterations = power_iterations
|
||||
if not self._made_params():
|
||||
self._make_params()
|
||||
|
||||
def _update_u_v(self):
|
||||
u = getattr(self.module, self.name + "_u")
|
||||
v = getattr(self.module, self.name + "_v")
|
||||
w = getattr(self.module, self.name + "_bar")
|
||||
|
||||
height = w.data.shape[0]
|
||||
for _ in range(self.power_iterations):
|
||||
v.data = l2normalize(torch.mv(torch.t(w.view(height,-1).data), u.data))
|
||||
u.data = l2normalize(torch.mv(w.view(height,-1).data, v.data))
|
||||
|
||||
sigma = u.dot(w.view(height, -1).mv(v))
|
||||
setattr(self.module, self.name, w / sigma.expand_as(w))
|
||||
|
||||
def _noupdate_u_v(self):
|
||||
u = getattr(self.module, self.name + "_u")
|
||||
v = getattr(self.module, self.name + "_v")
|
||||
w = getattr(self.module, self.name + "_bar")
|
||||
|
||||
height = w.data.shape[0]
|
||||
sigma = u.dot(w.view(height, -1).mv(v))
|
||||
setattr(self.module, self.name, w / sigma.expand_as(w))
|
||||
|
||||
def _made_params(self):
|
||||
try:
|
||||
u = getattr(self.module, self.name + "_u")
|
||||
v = getattr(self.module, self.name + "_v")
|
||||
w = getattr(self.module, self.name + "_bar")
|
||||
return True
|
||||
except AttributeError:
|
||||
return False
|
||||
|
||||
def _make_params(self):
|
||||
w = getattr(self.module, self.name)
|
||||
|
||||
height = w.data.shape[0]
|
||||
width = w.view(height, -1).data.shape[1]
|
||||
|
||||
u = Parameter(w.data.new(height).normal_(0, 1), requires_grad=False)
|
||||
v = Parameter(w.data.new(width).normal_(0, 1), requires_grad=False)
|
||||
u.data = l2normalize(u.data)
|
||||
v.data = l2normalize(v.data)
|
||||
w_bar = Parameter(w.data)
|
||||
|
||||
del self.module._parameters[self.name]
|
||||
|
||||
self.module.register_parameter(self.name + "_u", u)
|
||||
self.module.register_parameter(self.name + "_v", v)
|
||||
self.module.register_parameter(self.name + "_bar", w_bar)
|
||||
|
||||
def forward(self, *args):
|
||||
# if torch.is_grad_enabled() and self.module.training:
|
||||
if self.module.training:
|
||||
self._update_u_v()
|
||||
else:
|
||||
self._noupdate_u_v()
|
||||
return self.module.forward(*args)
|
||||
|
||||
|
||||
class GuidedCxtAtten(nn.Module):
|
||||
# based on https://github.com/nbei/Deep-Flow-Guided-Video-Inpainting/blob/a6fe298fec502bfd9cbc64eb01e39f78a3262a59/models/DeepFill_Models/ops.py#L210
|
||||
def __init__(self, out_channels, guidance_channels, rate=2):
|
||||
super(GuidedCxtAtten, self).__init__()
|
||||
self.rate = rate
|
||||
self.padding = nn.ReflectionPad2d(1)
|
||||
self.up_sample = nn.Upsample(scale_factor=self.rate, mode='nearest')
|
||||
|
||||
self.guidance_conv = nn.Conv2d(in_channels=guidance_channels, out_channels=guidance_channels//2,
|
||||
kernel_size=1, stride=1, padding=0)
|
||||
|
||||
self.W = nn.Sequential(
|
||||
nn.Conv2d(in_channels=out_channels, out_channels=out_channels,
|
||||
kernel_size=1, stride=1, padding=0, bias=False),
|
||||
nn.BatchNorm2d(out_channels)
|
||||
)
|
||||
|
||||
nn.init.xavier_uniform_(self.guidance_conv.weight)
|
||||
nn.init.constant_(self.guidance_conv.bias, 0)
|
||||
nn.init.xavier_uniform_(self.W[0].weight)
|
||||
nn.init.constant_(self.W[1].weight, 1e-3)
|
||||
nn.init.constant_(self.W[1].bias, 0)
|
||||
|
||||
def forward(self, f, alpha, unknown=None, ksize=3, stride=1, fuse_k=3, softmax_scale=1., training=True):
|
||||
|
||||
f = self.guidance_conv(f)
|
||||
# get shapes
|
||||
raw_int_fs = list(f.size()) # N x 64 x 64 x 64
|
||||
raw_int_alpha = list(alpha.size()) # N x 128 x 64 x 64
|
||||
|
||||
# extract patches from background with stride and rate
|
||||
kernel = 2*self.rate
|
||||
alpha_w = self.extract_patches(alpha, kernel=kernel, stride=self.rate)
|
||||
alpha_w = alpha_w.permute(0, 2, 3, 4, 5, 1)
|
||||
alpha_w = alpha_w.contiguous().view(raw_int_alpha[0], raw_int_alpha[2] // self.rate, raw_int_alpha[3] // self.rate, -1)
|
||||
alpha_w = alpha_w.contiguous().view(raw_int_alpha[0], -1, kernel, kernel, raw_int_alpha[1])
|
||||
alpha_w = alpha_w.permute(0, 1, 4, 2, 3)
|
||||
|
||||
f = F.interpolate(f, scale_factor=1/self.rate, mode='nearest')
|
||||
|
||||
fs = f.size() # B x 64 x 32 x 32
|
||||
f_groups = torch.split(f, 1, dim=0) # Split tensors by batch dimension; tuple is returned
|
||||
|
||||
# from b(B*H*W*C) to w(b*k*k*c*h*w)
|
||||
int_fs = list(fs)
|
||||
w = self.extract_patches(f)
|
||||
w = w.permute(0, 2, 3, 4, 5, 1)
|
||||
w = w.contiguous().view(raw_int_fs[0], raw_int_fs[2] // self.rate, raw_int_fs[3] // self.rate, -1)
|
||||
w = w.contiguous().view(raw_int_fs[0], -1, ksize, ksize, raw_int_fs[1])
|
||||
w = w.permute(0, 1, 4, 2, 3)
|
||||
# process mask
|
||||
|
||||
if unknown is not None:
|
||||
unknown = unknown.clone()
|
||||
unknown = F.interpolate(unknown, scale_factor=1/self.rate, mode='nearest')
|
||||
assert unknown.size(2) == f.size(2), "mask should have same size as f at dim 2,3"
|
||||
unknown_mean = unknown.mean(dim=[2,3])
|
||||
known_mean = 1 - unknown_mean
|
||||
unknown_scale = torch.clamp(torch.sqrt(unknown_mean / known_mean), 0.1, 10).to(alpha)
|
||||
known_scale = torch.clamp(torch.sqrt(known_mean / unknown_mean), 0.1, 10).to(alpha)
|
||||
softmax_scale = torch.cat([unknown_scale, known_scale], dim=1)
|
||||
else:
|
||||
unknown = torch.ones([fs[0], 1, fs[2], fs[3]]).to(alpha)
|
||||
softmax_scale = torch.FloatTensor([softmax_scale, softmax_scale]).view(1,2).repeat(fs[0],1).to(alpha)
|
||||
|
||||
m = self.extract_patches(unknown)
|
||||
|
||||
m = m.permute(0, 2, 3, 4, 5, 1)
|
||||
m = m.contiguous().view(raw_int_fs[0], raw_int_fs[2]//self.rate, raw_int_fs[3]//self.rate, -1)
|
||||
m = m.contiguous().view(raw_int_fs[0], -1, ksize, ksize)
|
||||
|
||||
m = self.reduce_mean(m) # smoothing, maybe
|
||||
# mask out the
|
||||
mm = m.gt(0.).float() # (N, 32*32, 1, 1)
|
||||
|
||||
# the correlation with itself should be 0
|
||||
self_mask = F.one_hot(torch.arange(fs[2] * fs[3]).view(fs[2], fs[3]).contiguous().to(alpha).long(),
|
||||
num_classes=int_fs[2] * int_fs[3])
|
||||
self_mask = self_mask.permute(2, 0, 1).view(1, fs[2] * fs[3], fs[2], fs[3]).float() * (-1e4)
|
||||
|
||||
w_groups = torch.split(w, 1, dim=0) # Split tensors by batch dimension; tuple is returned
|
||||
alpha_w_groups = torch.split(alpha_w, 1, dim=0) # Split tensors by batch dimension; tuple is returned
|
||||
mm_groups = torch.split(mm, 1, dim=0)
|
||||
scale_group = torch.split(softmax_scale, 1, dim=0)
|
||||
y = []
|
||||
offsets = []
|
||||
k = fuse_k
|
||||
y_test = []
|
||||
for xi, wi, alpha_wi, mmi, scale in zip(f_groups, w_groups, alpha_w_groups, mm_groups, scale_group):
|
||||
# conv for compare
|
||||
wi = wi[0]
|
||||
escape_NaN = Variable(torch.FloatTensor([1e-4])).to(alpha)
|
||||
wi_normed = wi / torch.max(self.l2_norm(wi), escape_NaN)
|
||||
xi = F.pad(xi, (1,1,1,1), mode='reflect')
|
||||
yi = F.conv2d(xi, wi_normed, stride=1, padding=0) # yi => (B=1, C=32*32, H=32, W=32)
|
||||
y_test.append(yi)
|
||||
# conv implementation for fuse scores to encourage large patches
|
||||
yi = yi.permute(0, 2, 3, 1)
|
||||
yi = yi.contiguous().view(1, fs[2], fs[3], fs[2] * fs[3])
|
||||
yi = yi.permute(0, 3, 1, 2) # (B=1, C=32*32, H=32, W=32)
|
||||
|
||||
# softmax to match
|
||||
# scale the correlation with predicted scale factor for known and unknown area
|
||||
yi = yi * (scale[0,0] * mmi.gt(0.).float() + scale[0,1] * mmi.le(0.).float()) # mmi => (1, 32*32, 1, 1)
|
||||
# mask itself, self-mask only applied to unknown area
|
||||
yi = yi + self_mask * mmi # self_mask: (1, 32*32, 32, 32)
|
||||
# for small input inference
|
||||
yi = F.softmax(yi, dim=1)
|
||||
|
||||
_, offset = torch.max(yi, dim=1) # argmax; index
|
||||
offset = torch.stack([offset // fs[3], offset % fs[3]], dim=1)
|
||||
|
||||
wi_center = alpha_wi[0]
|
||||
|
||||
if self.rate == 1:
|
||||
left = (kernel) // 2
|
||||
right = (kernel - 1) // 2
|
||||
yi = F.pad(yi, (left, right, left, right), mode='reflect')
|
||||
wi_center = wi_center.permute(1, 0, 2, 3)
|
||||
yi = F.conv2d(yi, wi_center, padding=0) / 4. # (B=1, C=128, H=64, W=64)
|
||||
else:
|
||||
yi = F.conv_transpose2d(yi, wi_center, stride=self.rate, padding=1) / 4. # (B=1, C=128, H=64, W=64)
|
||||
y.append(yi)
|
||||
offsets.append(offset)
|
||||
|
||||
y = torch.cat(y, dim=0) # back to the mini-batch
|
||||
y.contiguous().view(raw_int_alpha)
|
||||
offsets = torch.cat(offsets, dim=0)
|
||||
offsets = offsets.view([int_fs[0]] + [2] + int_fs[2:])
|
||||
|
||||
# # case1: visualize optical flow: minus current position
|
||||
# h_add = Variable(torch.arange(0,float(fs[2]))).to(alpha).view([1, 1, fs[2], 1])
|
||||
# h_add = h_add.expand(fs[0], 1, fs[2], fs[3])
|
||||
# w_add = Variable(torch.arange(0,float(fs[3]))).to(alpha).view([1, 1, 1, fs[3]])
|
||||
# w_add = w_add.expand(fs[0], 1, fs[2], fs[3])
|
||||
#
|
||||
# offsets = offsets - torch.cat([h_add, w_add], dim=1).long()
|
||||
|
||||
# case2: visualize absolute position
|
||||
offsets = offsets - torch.Tensor([fs[2]//2, fs[3]//2]).view(1,2,1,1).to(alpha).long()
|
||||
|
||||
y = self.W(y) + alpha
|
||||
|
||||
return y, (offsets, softmax_scale)
|
||||
|
||||
@staticmethod
|
||||
def extract_patches(x, kernel=3, stride=1):
|
||||
left =(kernel - stride + 1) // 2
|
||||
right =(kernel - stride) // 2
|
||||
x = F.pad(x, (left, right, left, right), mode='reflect')
|
||||
all_patches = x.unfold(2, kernel, stride).unfold(3, kernel, stride)
|
||||
|
||||
return all_patches
|
||||
|
||||
@staticmethod
|
||||
def reduce_mean(x):
|
||||
for i in range(4):
|
||||
if i <= 1:
|
||||
continue
|
||||
x = torch.mean(x, dim=i, keepdim=True)
|
||||
return x
|
||||
|
||||
@staticmethod
|
||||
def l2_norm(x):
|
||||
def reduce_sum(x):
|
||||
for i in range(4):
|
||||
if i == 0:
|
||||
continue
|
||||
x = torch.sum(x, dim=i, keepdim=True)
|
||||
return x
|
||||
|
||||
x = x**2
|
||||
x = reduce_sum(x)
|
||||
return torch.sqrt(x)
|
||||
@@ -0,0 +1,102 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
# @File : setup.py
|
||||
# @Time : 2020/1/15
|
||||
# @Author : yangchaojie (yangchaojie@immomo.com)
|
||||
|
||||
import os
|
||||
import sys
|
||||
import shutil
|
||||
import numpy
|
||||
import tempfile
|
||||
|
||||
from setuptools import setup
|
||||
from setuptools.extension import Extension
|
||||
|
||||
from Cython.Build import cythonize
|
||||
from Cython.Distutils import build_ext
|
||||
|
||||
import platform
|
||||
|
||||
|
||||
def get_root_path(root):
|
||||
if os.path.dirname(root) in ['', '.']:
|
||||
return os.path.basename(root)
|
||||
else:
|
||||
return get_root_path(os.path.dirname(root))
|
||||
|
||||
|
||||
def copy_file(src, dest):
|
||||
if os.path.exists(dest):
|
||||
return
|
||||
|
||||
if not os.path.exists(os.path.dirname(dest)):
|
||||
os.makedirs(os.path.dirname(dest))
|
||||
if os.path.isdir(src):
|
||||
shutil.copytree(src, dest)
|
||||
else:
|
||||
shutil.copyfile(src, dest)
|
||||
|
||||
|
||||
def touch_init_file():
|
||||
init_file_name = os.path.join(tempfile.mkdtemp(), '__init__.py')
|
||||
with open(init_file_name, 'w'):
|
||||
pass
|
||||
return init_file_name
|
||||
|
||||
|
||||
|
||||
|
||||
def compose_extensions(root='.'):
|
||||
for file_ in os.listdir(root):
|
||||
abs_file = os.path.join(root, file_)
|
||||
|
||||
if os.path.isfile(abs_file):
|
||||
if abs_file.endswith('.py'):
|
||||
extensions.append(Extension(get_root_path(abs_file) + '.*', [abs_file]))
|
||||
elif abs_file.endswith('.c') or abs_file.endswith('.pyc'):
|
||||
continue
|
||||
else:
|
||||
copy_file(abs_file, os.path.join(build_root_dir, abs_file))
|
||||
if abs_file.endswith('__init__.py'):
|
||||
copy_file(init_file, os.path.join(build_root_dir, abs_file))
|
||||
|
||||
else:
|
||||
if os.path.basename(abs_file) in ignore_folders :
|
||||
continue
|
||||
if os.path.basename(abs_file) in conf_folders:
|
||||
copy_file(abs_file, os.path.join(build_root_dir, abs_file))
|
||||
compose_extensions(abs_file)
|
||||
|
||||
|
||||
|
||||
# if __name__ == '__main__':
|
||||
build_root_dir = 'build/lib.' + platform.system().lower() + '-' + platform.machine() + '-' + str(
|
||||
sys.version_info.major) + '.' + str(sys.version_info.minor)
|
||||
|
||||
print(build_root_dir)
|
||||
|
||||
extensions = []
|
||||
ignore_folders = ['build', 'new_ref_zao_color_0818', 'ref_hair_online_0703_local', 'ref_分好类别', '.git']
|
||||
conf_folders = ['conf']
|
||||
|
||||
|
||||
init_file = touch_init_file()
|
||||
print(init_file)
|
||||
|
||||
|
||||
compose_extensions()
|
||||
os.remove(init_file)
|
||||
|
||||
setup(
|
||||
name='moxie_hairstyle',
|
||||
version='1.0',
|
||||
ext_modules=cythonize(
|
||||
extensions,
|
||||
nthreads=16,
|
||||
compiler_directives=dict(always_allow_keywords=True),
|
||||
include_path=[numpy.get_include()]),
|
||||
cmdclass=dict(build_ext=build_ext))
|
||||
|
||||
# python setup.py build_ext
|
||||
@@ -0,0 +1,319 @@
|
||||
import torch.nn as nn
|
||||
import torch.utils.model_zoo as model_zoo
|
||||
import torch
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
from core.utils import utils_3ddfa,params_3ddfa,landmark_processor
|
||||
__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
|
||||
'resnet152']
|
||||
|
||||
|
||||
model_urls = {
|
||||
'resnet18': 'http://download.pytorch.org/models/resnet18-5c106cde.pth',
|
||||
'resnet34': 'http://download.pytorch.org/models/resnet34-333f7ec4.pth',
|
||||
'resnet50': 'http://download.pytorch.org/models/resnet50-19c8e357.pth',
|
||||
'resnet101': 'http://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
|
||||
'resnet152': 'http://download.pytorch.org/models/resnet152-b121ed2d.pth',
|
||||
}
|
||||
|
||||
|
||||
def conv3x3(in_planes, out_planes, stride=1):
|
||||
"""3x3 convolution with padding"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
|
||||
padding=1, bias=False)
|
||||
|
||||
|
||||
def conv1x1(in_planes, out_planes, stride=1):
|
||||
"""1x1 convolution"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
|
||||
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
expansion = 1
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None):
|
||||
super(BasicBlock, self).__init__()
|
||||
self.conv1 = conv3x3(inplanes, planes, stride)
|
||||
self.bn1 = nn.BatchNorm2d(planes)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.conv2 = conv3x3(planes, planes)
|
||||
self.bn2 = nn.BatchNorm2d(planes)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
identity = self.downsample(x)
|
||||
|
||||
out += identity
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class Bottleneck(nn.Module):
|
||||
expansion = 4
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None):
|
||||
super(Bottleneck, self).__init__()
|
||||
self.conv1 = conv1x1(inplanes, planes)
|
||||
self.bn1 = nn.BatchNorm2d(planes)
|
||||
self.conv2 = conv3x3(planes, planes, stride)
|
||||
self.bn2 = nn.BatchNorm2d(planes)
|
||||
self.conv3 = conv1x1(planes, planes * self.expansion)
|
||||
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
identity = self.downsample(x)
|
||||
|
||||
out += identity
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class ResNet(nn.Module):
|
||||
|
||||
def __init__(self, block, layers, num_classes=1000, no_branch=False, no_activate=False, zero_init_residual=False):
|
||||
super(ResNet, self).__init__()
|
||||
self.inplanes = 64
|
||||
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(64)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
|
||||
self.layer1 = self._make_layer(block, 64, layers[0])
|
||||
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
|
||||
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
|
||||
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
|
||||
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
|
||||
if no_activate:
|
||||
if no_branch:
|
||||
self.fc = nn.Sequential(*[nn.Linear(512 * block.expansion, num_classes)])
|
||||
else:
|
||||
self.fc_key = nn.Sequential(*[nn.Linear(256 * block.expansion, 45 * 2)])
|
||||
self.fc_ctrl = nn.Sequential(*[nn.Linear(256 * block.expansion, 48 * 2)])
|
||||
else:
|
||||
if no_branch:
|
||||
self.fc = nn.Sequential(*[nn.Linear(512 * block.expansion, num_classes), nn.Tanh()])
|
||||
else:
|
||||
self.fc_key = nn.Sequential(*[nn.Linear(256 * block.expansion, 45 * 2), nn.Tanh()])
|
||||
self.fc_ctrl = nn.Sequential(*[nn.Linear(256 * block.expansion, 48 * 2), nn.Tanh()])
|
||||
self.no_branch = no_branch
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
# Zero-initialize the last BN in each residual branch,
|
||||
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
|
||||
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
|
||||
if zero_init_residual:
|
||||
for m in self.modules():
|
||||
if isinstance(m, Bottleneck):
|
||||
nn.init.constant_(m.bn3.weight, 0)
|
||||
elif isinstance(m, BasicBlock):
|
||||
nn.init.constant_(m.bn2.weight, 0)
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1):
|
||||
downsample = None
|
||||
if stride != 1 or self.inplanes != planes * block.expansion:
|
||||
downsample = nn.Sequential(
|
||||
conv1x1(self.inplanes, planes * block.expansion, stride),
|
||||
nn.BatchNorm2d(planes * block.expansion),
|
||||
)
|
||||
|
||||
layers = []
|
||||
layers.append(block(self.inplanes, planes, stride, downsample))
|
||||
self.inplanes = planes * block.expansion
|
||||
for _ in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
x = self.maxpool(x)
|
||||
|
||||
x = self.layer1(x)
|
||||
x = self.layer2(x)
|
||||
x = self.layer3(x)
|
||||
x = self.layer4(x)
|
||||
|
||||
if self.no_branch:
|
||||
key = self.avgpool(x)
|
||||
key = key.view(key.size(0), -1)
|
||||
key = self.fc(key)
|
||||
return key
|
||||
else:
|
||||
key, ctrl = torch.chunk(x, 2, dim=1)
|
||||
|
||||
key = self.avgpool(key)
|
||||
key = key.view(key.size(0), -1)
|
||||
key = self.fc_key(key)
|
||||
|
||||
ctrl = self.avgpool(ctrl)
|
||||
ctrl = ctrl.view(ctrl.size(0), -1)
|
||||
ctrl = self.fc_ctrl(ctrl)
|
||||
|
||||
return key, ctrl
|
||||
|
||||
def resnet18(pretrained=False, **kwargs):
|
||||
"""Constructs a ResNet-18 model.
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
|
||||
if pretrained:
|
||||
model.load_state_dict(model_zoo.load_url(model_urls['resnet18']), strict=False)
|
||||
return model
|
||||
|
||||
|
||||
def resnet34(pretrained=False, **kwargs):
|
||||
"""Constructs a ResNet-34 model.
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
|
||||
if pretrained:
|
||||
model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
|
||||
return model
|
||||
|
||||
|
||||
def resnet50(pretrained=False, **kwargs):
|
||||
"""Constructs a ResNet-50 model.
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
|
||||
if pretrained:
|
||||
model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
|
||||
return model
|
||||
|
||||
|
||||
def resnet101(pretrained=False, **kwargs):
|
||||
"""Constructs a ResNet-101 model.
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
|
||||
if pretrained:
|
||||
model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
|
||||
return model
|
||||
|
||||
|
||||
def resnet152(pretrained=False, **kwargs):
|
||||
"""Constructs a ResNet-152 model.
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
|
||||
if pretrained:
|
||||
model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
|
||||
return model
|
||||
|
||||
|
||||
class Model_3DDFA(nn.Module):
|
||||
|
||||
|
||||
def __init__(self, gpu_id=None):
|
||||
super(Model_3DDFA, self).__init__()
|
||||
|
||||
self.init_status = False
|
||||
self.face_alignment_net = resnet18(pretrained=True, num_classes=76, no_branch=True, no_activate=True)
|
||||
self.model_path = 'weights'
|
||||
|
||||
# model_path, _ = os.path.split(os.path.realpath(__file__))
|
||||
weights = torch.load(os.path.join(self.model_path, 'face_3ddfa.pth'), map_location=lambda storage, loc: storage)
|
||||
self.load_state_dict(weights)
|
||||
self.eval()
|
||||
self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||
self.to(self.device)
|
||||
self.init_status = True
|
||||
|
||||
def forward(self, imgs):
|
||||
pred_pose_shape_exp = self.face_alignment_net(imgs)
|
||||
return pred_pose_shape_exp
|
||||
|
||||
def running(self):
|
||||
return self.init_status
|
||||
|
||||
def forward_np(self, imgs):
|
||||
pred_pose_shape_exp = self.face_alignment_net(imgs)
|
||||
return pred_pose_shape_exp.detach().cpu().numpy()
|
||||
|
||||
def detect(self, images, landmarks):
|
||||
dst_size = 256
|
||||
with torch.no_grad():
|
||||
input_numpy = np.zeros((len(images), 3, dst_size, dst_size), dtype=np.float32)
|
||||
assert len(images) == len(landmarks)
|
||||
all_mat = []
|
||||
all_res = []
|
||||
all_height = []
|
||||
for ix, img in enumerate(images):
|
||||
landmark = landmarks[ix]
|
||||
|
||||
mat = landmark_processor.get_transform_mat_full_face(landmark, dst_size)
|
||||
|
||||
all_mat.append(mat)
|
||||
all_height.append(img.shape[0])
|
||||
|
||||
tmp = cv2.warpAffine(img, mat, (dst_size, dst_size))
|
||||
|
||||
input_numpy[ix, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
|
||||
|
||||
# cv2.imshow('3ddfa_', tmp)
|
||||
# cv2.waitKey()
|
||||
|
||||
in_tensor = torch.from_numpy(input_numpy)
|
||||
in_tensor = in_tensor.to(self.device)
|
||||
params = self.face_alignment_net(in_tensor)
|
||||
params = params.cpu().numpy()
|
||||
|
||||
for ix, param in enumerate(params):
|
||||
param[0] = param[0] / params_3ddfa.SCALE_F
|
||||
param[1:4] = param[1:4] / params_3ddfa.SCALE_ROTATE
|
||||
param[4:6] = param[4:6] / params_3ddfa.SCALE_OFFSET
|
||||
param[6:56] = (param[6:56] / params_3ddfa.SCALE_SHAPE)
|
||||
param[56:] = (param[56:] / params_3ddfa.SCALE_EXP)
|
||||
new_param = utils_3ddfa.transform_params(param, cv2.invertAffineTransform(all_mat[ix]), all_height[ix],
|
||||
dst_size)
|
||||
all_res.append(new_param)
|
||||
|
||||
return all_res
|
||||
@@ -0,0 +1,67 @@
|
||||
import os
|
||||
import torch
|
||||
import torch.nn.parallel
|
||||
|
||||
import numpy as np
|
||||
from core.utils import landmark_processor
|
||||
import cv2
|
||||
|
||||
modelRoot = "weights"
|
||||
|
||||
label_map = [
|
||||
[0, 0, 0],
|
||||
[255, 0, 0],
|
||||
[0, 0, 255],
|
||||
[0, 255, 0],
|
||||
[0, 255, 255],
|
||||
# [0, 255, 0]
|
||||
]
|
||||
class Generator_BaldSeg_5c(object):
|
||||
def __init__(self, gpu_flag, gpu_id):
|
||||
|
||||
if not gpu_flag:
|
||||
self.device = torch.device("cpu")
|
||||
else:
|
||||
self.device = torch.device('cuda:{0}'.format(gpu_id))
|
||||
|
||||
# load seg model
|
||||
self.model_dir = modelRoot
|
||||
self.pre_trained_model = os.path.join('weights', "ori_hair_checkpoint_7660_0611.pt")
|
||||
self.net = torch.jit.load(self.pre_trained_model, map_location=self.device).to(self.device)
|
||||
self.net.eval()
|
||||
|
||||
self.output_img_size = 512
|
||||
self.img_ratio = 0.4
|
||||
|
||||
def label_to_mask(self, label_np):
|
||||
label_np = label_np.astype(np.int32)[:, :, np.newaxis]
|
||||
mask = np.zeros((label_np.shape[0], label_np.shape[1], 3), dtype=np.uint8)
|
||||
for id, color in enumerate(label_map):
|
||||
index = (label_np == id).all(axis=2)
|
||||
mask[index] = color
|
||||
return mask
|
||||
def forward(self, image, alpha, landmarks1k):
|
||||
image_to_face_mat = landmark_processor.get_transform_mat_full_face_ratio_deeplab(landmarks1k, self.output_img_size, self.img_ratio)
|
||||
img_ = (image * (1 - alpha[:, :, np.newaxis].astype(np.float32) / 255)).astype(np.uint8)
|
||||
img_cuted = cv2.warpAffine(img_, image_to_face_mat, (self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4, borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255])
|
||||
inv_image_to_face_mat = cv2.invertAffineTransform(image_to_face_mat)
|
||||
|
||||
img = (img_cuted.astype(np.float32) / 255).transpose((2, 0, 1))
|
||||
img = np.expand_dims(img, axis=0)
|
||||
img = torch.from_numpy(img).to(self.device) #cuda(self.gpu_id)
|
||||
|
||||
with torch.no_grad():
|
||||
output = self.net(img)
|
||||
|
||||
pred = output.detach().cpu().numpy().squeeze().astype(np.float32) #torch.max(output[:1], 1)[1].detach().cpu().numpy().squeeze().astype(np.float32)
|
||||
mask = self.label_to_mask(pred)
|
||||
|
||||
# cv2.imshow("img_cuted: ", img_cuted)
|
||||
# cv2.imshow("mask: ", mask)
|
||||
# cv2.waitKey()
|
||||
|
||||
black_img = np.zeros(image.shape).astype(np.uint8)
|
||||
cv2.warpAffine(mask, inv_image_to_face_mat, (image.shape[1], image.shape[0]),
|
||||
dst=black_img, flags=cv2.INTER_NEAREST, borderMode=cv2.BORDER_TRANSPARENT)
|
||||
|
||||
return black_img
|
||||
@@ -0,0 +1,441 @@
|
||||
import torch.nn as nn
|
||||
import torch.utils.model_zoo as model_zoo
|
||||
import torch
|
||||
import numpy as np
|
||||
import os
|
||||
from core.utils import landmark_processor
|
||||
from core.utils.umeyama import umeyama
|
||||
import cv2
|
||||
|
||||
modelRoot = "weights"
|
||||
|
||||
__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
|
||||
'resnet152']
|
||||
|
||||
|
||||
model_urls = {
|
||||
'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
|
||||
'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
|
||||
'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
|
||||
'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
|
||||
'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
|
||||
}
|
||||
|
||||
|
||||
def conv3x3(in_planes, out_planes, stride=1):
|
||||
"""3x3 convolution with padding"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
|
||||
padding=1, bias=False)
|
||||
|
||||
|
||||
def conv1x1(in_planes, out_planes, stride=1):
|
||||
"""1x1 convolution"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
|
||||
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
expansion = 1
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None):
|
||||
super(BasicBlock, self).__init__()
|
||||
self.conv1 = conv3x3(inplanes, planes, stride)
|
||||
self.bn1 = nn.BatchNorm2d(planes)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.conv2 = conv3x3(planes, planes)
|
||||
self.bn2 = nn.BatchNorm2d(planes)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
identity = self.downsample(x)
|
||||
|
||||
out += identity
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class Bottleneck(nn.Module):
|
||||
expansion = 4
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None):
|
||||
super(Bottleneck, self).__init__()
|
||||
self.conv1 = conv1x1(inplanes, planes)
|
||||
self.bn1 = nn.BatchNorm2d(planes)
|
||||
self.conv2 = conv3x3(planes, planes, stride)
|
||||
self.bn2 = nn.BatchNorm2d(planes)
|
||||
self.conv3 = conv1x1(planes, planes * self.expansion)
|
||||
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
identity = self.downsample(x)
|
||||
|
||||
out += identity
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class ResNet(nn.Module):
|
||||
|
||||
def __init__(self, block, layers, num_classes=1000, is_1k=False, zero_init_residual=False):
|
||||
super(ResNet, self).__init__()
|
||||
self.inplanes = 64
|
||||
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
|
||||
bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(64)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
|
||||
self.layer1 = self._make_layer(block, 64, layers[0])
|
||||
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
|
||||
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
|
||||
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
|
||||
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
|
||||
if is_1k:
|
||||
self.fc = nn.Sequential(*[nn.Linear(512 * block.expansion, num_classes), nn.Tanh()])
|
||||
else:
|
||||
self.fc_key = nn.Sequential(*[nn.Linear(256 * block.expansion, 45 * 2), nn.Tanh()])
|
||||
self.fc_ctrl = nn.Sequential(*[nn.Linear(256 * block.expansion, 48 * 2), nn.Tanh()])
|
||||
self.is_1k = is_1k
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
# Zero-initialize the last BN in each residual branch,
|
||||
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
|
||||
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
|
||||
if zero_init_residual:
|
||||
for m in self.modules():
|
||||
if isinstance(m, Bottleneck):
|
||||
nn.init.constant_(m.bn3.weight, 0)
|
||||
elif isinstance(m, BasicBlock):
|
||||
nn.init.constant_(m.bn2.weight, 0)
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1):
|
||||
downsample = None
|
||||
if stride != 1 or self.inplanes != planes * block.expansion:
|
||||
downsample = nn.Sequential(
|
||||
conv1x1(self.inplanes, planes * block.expansion, stride),
|
||||
nn.BatchNorm2d(planes * block.expansion),
|
||||
)
|
||||
|
||||
layers = []
|
||||
layers.append(block(self.inplanes, planes, stride, downsample))
|
||||
self.inplanes = planes * block.expansion
|
||||
for _ in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
x = self.maxpool(x)
|
||||
|
||||
x = self.layer1(x)
|
||||
x = self.layer2(x)
|
||||
x = self.layer3(x)
|
||||
x = self.layer4(x)
|
||||
|
||||
if self.is_1k:
|
||||
key = self.avgpool(x)
|
||||
key = key.view(key.size(0), -1)
|
||||
key = self.fc(key)
|
||||
return key
|
||||
else:
|
||||
key, ctrl = torch.chunk(x, 2, dim=1)
|
||||
|
||||
key = self.avgpool(key)
|
||||
key = key.view(key.size(0), -1)
|
||||
key = self.fc_key(key)
|
||||
|
||||
ctrl = self.avgpool(ctrl)
|
||||
ctrl = ctrl.view(ctrl.size(0), -1)
|
||||
ctrl = self.fc_ctrl(ctrl)
|
||||
|
||||
return key, ctrl
|
||||
|
||||
def resnet18(pretrained=False, **kwargs):
|
||||
"""Constructs a ResNet-18 model.
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
|
||||
if pretrained:
|
||||
model.load_state_dict(model_zoo.load_url(model_urls['resnet18']), strict=False)
|
||||
return model
|
||||
|
||||
|
||||
def resnet34(pretrained=False, **kwargs):
|
||||
"""Constructs a ResNet-34 model.
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
|
||||
if pretrained:
|
||||
model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
|
||||
return model
|
||||
|
||||
|
||||
def resnet50(pretrained=False, **kwargs):
|
||||
"""Constructs a ResNet-50 model.
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
|
||||
if pretrained:
|
||||
model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
|
||||
return model
|
||||
|
||||
|
||||
def resnet101(pretrained=False, **kwargs):
|
||||
"""Constructs a ResNet-101 model.
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
|
||||
if pretrained:
|
||||
model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
|
||||
return model
|
||||
|
||||
|
||||
def resnet152(pretrained=False, **kwargs):
|
||||
"""Constructs a ResNet-152 model.
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
|
||||
if pretrained:
|
||||
model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
|
||||
return model
|
||||
|
||||
class Model1k(nn.Module):
|
||||
def __init__(self, gpu_id=None):
|
||||
super(Model1k, self).__init__()
|
||||
|
||||
self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||
|
||||
self.face_alignment_net = resnet18(pretrained=False, num_classes=1000 * 2, is_1k=True)
|
||||
|
||||
self.model_path, _ = os.path.split(os.path.realpath(__file__))
|
||||
# weights = torch.load(os.path.join(self.model_path, 'face_alignment_1k.pth'), map_location=lambda storage, loc: storage)
|
||||
weights = torch.load(os.path.join(modelRoot, 'face_alignment_1k.pth'), map_location=lambda storage, loc: storage)
|
||||
self.load_state_dict(weights)
|
||||
self.to(self.device)
|
||||
self.eval()
|
||||
|
||||
def forward(self, imgs):
|
||||
pred_key_pts = self.face_alignment_net(imgs)
|
||||
pred_key_pts = pred_key_pts + 0.5
|
||||
return pred_key_pts
|
||||
|
||||
class MomocvFaceAlignment1K(object):
|
||||
def __init__(self, gpu_id=None):
|
||||
self.gpu_id = gpu_id
|
||||
|
||||
self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||
|
||||
self.face_alignment_net = Model1k(gpu_id)
|
||||
|
||||
self.trackingFaceRects = []
|
||||
|
||||
print('conansherry MomocvFaceAlignment1K')
|
||||
|
||||
def forward(self, img_tensor):
|
||||
fullyconnected1 = self.face_alignment_net(img_tensor).detach().cpu().numpy()
|
||||
return fullyconnected1
|
||||
|
||||
def detect(self, img, landmarks):
|
||||
dst_size = 256
|
||||
landmarks_res = []
|
||||
with torch.no_grad():
|
||||
input_numpy = np.zeros((len(landmarks), 3, dst_size, dst_size), dtype=np.float32)
|
||||
all_mat = []
|
||||
for ix, landmark in enumerate(landmarks):
|
||||
M = landmark_processor.get_transform_mat_full_face(landmark, dst_size)
|
||||
all_mat.append(M)
|
||||
tmp = cv2.warpAffine(img, M, (dst_size, dst_size))
|
||||
|
||||
# cv2.imshow('inp', tmp)
|
||||
# cv2.waitKey()
|
||||
|
||||
input_numpy[ix, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
|
||||
in_tensor = torch.from_numpy(input_numpy)
|
||||
in_tensor = in_tensor.to(self.device)
|
||||
fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy()
|
||||
for ix, pts in enumerate(fullyconnected1):
|
||||
orig_pts = (np.reshape(pts, (2, 1000)).transpose((1, 0)) * dst_size)
|
||||
orig_pts = landmark_processor.transform_points(orig_pts, all_mat[ix], invert=True)
|
||||
landmarks_res.append(orig_pts)
|
||||
return landmarks_res
|
||||
def detect_single_face(self, img):
|
||||
dst_size = 256
|
||||
with torch.no_grad():
|
||||
crop_img = img[104:img.shape[0] - 104, 104:img.shape[1] - 104, :]
|
||||
tmp = cv2.resize(crop_img, (dst_size, dst_size))
|
||||
input_numpy = np.zeros((1, 3, dst_size, dst_size), dtype=np.float32)
|
||||
input_numpy[0, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
|
||||
in_tensor = torch.from_numpy(input_numpy)
|
||||
in_tensor = in_tensor.to(self.device)
|
||||
fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy()
|
||||
orig_pts = (np.reshape(fullyconnected1[0], (2, 1000)).transpose((1, 0)) * crop_img.shape[0])
|
||||
orig_pts[:, 0] += 104
|
||||
orig_pts[:, 1] += 104
|
||||
return orig_pts
|
||||
|
||||
def detect_single_face_old(self, img):
|
||||
dst_size = 256
|
||||
with torch.no_grad():
|
||||
tmp = cv2.resize(img, (dst_size, dst_size))
|
||||
input_numpy = np.zeros((1, 3, dst_size, dst_size), dtype=np.float32)
|
||||
input_numpy[0, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
|
||||
in_tensor = torch.from_numpy(input_numpy)
|
||||
in_tensor = in_tensor.to(self.device)
|
||||
fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy()
|
||||
orig_pts = (np.reshape(fullyconnected1[0], (2, 1000)).transpose((1, 0)) * img.shape[0])
|
||||
return orig_pts
|
||||
def detect_according_5pts(self, img, pts5):
|
||||
dst_size = 256
|
||||
with torch.no_grad():
|
||||
input_numpy = np.zeros((1, 3, dst_size, dst_size), dtype=np.float32)
|
||||
eye_dis = 0.34
|
||||
mouth_dis = 0.34
|
||||
g_Average_5point_180 = np.array([
|
||||
eye_dis, 0.3,
|
||||
1 - eye_dis, 0.37,
|
||||
0.5, 0.6,
|
||||
mouth_dis, 0.63,
|
||||
1 - mouth_dis, 0.7
|
||||
])
|
||||
# print(g_Average_5point_180)
|
||||
# left_eye = np.array([pts5[0], pts5[5]])
|
||||
# right_eye = np.array([pts5[1], pts5[6]])
|
||||
# nose = np.array([pts5[2], pts5[7]])
|
||||
# left_mouth = np.array([pts5[3], pts5[8]])
|
||||
# right_mouth = np.array([pts5[4], pts5[9]])
|
||||
left_eye = np.array([pts5[0], pts5[1]])
|
||||
right_eye = np.array([pts5[2], pts5[3]])
|
||||
nose = np.array([pts5[4], pts5[5]])
|
||||
left_mouth = np.array([pts5[6], pts5[7]])
|
||||
right_mouth = np.array([pts5[8], pts5[9]])
|
||||
pts5_src = np.vstack((left_eye, right_eye,
|
||||
nose,
|
||||
left_mouth, right_mouth))
|
||||
pts5_src = np.array(pts5_src).astype(np.int32)
|
||||
pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size
|
||||
# print('pts5_src: ', pts5_src)
|
||||
# print('pts5_dst: ', pts5_dst)
|
||||
|
||||
# mat = cv2.estimateAffinePartial2D(pts5_src, pts5_dst, False)[0]
|
||||
mat = umeyama(pts5_src, pts5_dst, True)[0:2]
|
||||
print('mat: ', mat)
|
||||
tmp = cv2.warpAffine(img, mat, (dst_size, dst_size))
|
||||
|
||||
# tmp2 = cv2.warpAffine(img, mat2, (dst_size, dst_size))
|
||||
# cv2.imshow("tmp2", tmp2)
|
||||
# cv2.imshow("tmp", tmp)
|
||||
# cv2.waitKey()
|
||||
input_numpy[0, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
|
||||
in_tensor = torch.from_numpy(input_numpy)
|
||||
in_tensor = in_tensor.to(self.device)
|
||||
fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy()
|
||||
orig_pts = (np.reshape(fullyconnected1[0], (2, 1000)).transpose((1, 0)) * dst_size)
|
||||
orig_pts = landmark_processor.transform_points(orig_pts, mat, invert=True)
|
||||
return orig_pts
|
||||
|
||||
def stable_forward(self, image, detected_faces, reset=False):
|
||||
if reset is True:
|
||||
self.trackingFaceRects = []
|
||||
|
||||
if len(self.trackingFaceRects) == 0:
|
||||
for face_rect in detected_faces:
|
||||
new_tracking_rect = [face_rect, True, [0, 0], 0, None]
|
||||
self.trackingFaceRects.append(new_tracking_rect)
|
||||
|
||||
with torch.no_grad():
|
||||
landmarks = []
|
||||
for ix, tracking_face_rect in enumerate(self.trackingFaceRects):
|
||||
if tracking_face_rect[1] == True:
|
||||
d = tracking_face_rect[0]
|
||||
src_center = np.array([d[2] - (d[2] - d[0]) / 2.0, d[3] - (d[3] - d[1]) / 2.0])
|
||||
rotate_degree = tracking_face_rect[3]
|
||||
scale = 256 * 0.6 / min(d[2] - d[0], d[3] - d[1])
|
||||
dst_center = np.array([0.5, 0.5]) * 256
|
||||
offset = dst_center - src_center
|
||||
M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), rotate_degree, scale)
|
||||
M[:, 2] += offset
|
||||
else:
|
||||
rotate_degree = 0
|
||||
M = landmark_processor.get_transform_mat_mmcv_bigger(tracking_face_rect[4], 256)
|
||||
inp = cv2.warpAffine(image, M, (256, 256))
|
||||
|
||||
# cv2.imshow('inp_{}'.format(ix), inp)
|
||||
# cv2.waitKey()
|
||||
orig_inp = inp
|
||||
|
||||
inp = inp.transpose((2, 0, 1)).astype(np.float32)
|
||||
inp = inp[np.newaxis, :, :, :] / 255
|
||||
|
||||
in_tensor = torch.from_numpy(inp)
|
||||
in_tensor = in_tensor.cuda(0)
|
||||
fullyconnected1 = self.forward(in_tensor)
|
||||
fullyconnected1 = fullyconnected1[0]
|
||||
orig_pts = (np.reshape(fullyconnected1, (2, 1000)).transpose((1, 0))) * 256
|
||||
|
||||
t2 = cv2.getTickCount()
|
||||
orig_pts = landmark_processor.transform_points(orig_pts, M, invert=True)
|
||||
|
||||
# orig_pts = orig_pts.transpose((1, 0))
|
||||
fullyconnected1 = orig_pts
|
||||
|
||||
# update tracking infos
|
||||
tracking_face_rect[1] = False
|
||||
tracking_face_rect[2] = None
|
||||
tracking_face_rect[3] = rotate_degree
|
||||
tracking_face_rect[4] = fullyconnected1
|
||||
|
||||
# fullyconnected1 = landmark_processor.pts_1k_to_137(fullyconnected1)
|
||||
|
||||
# eye_landmark = self.detect_eye(image, fullyconnected1)
|
||||
# fullyconnected1[87:104] = eye_landmark[0]
|
||||
# fullyconnected1[104:121] = eye_landmark[1]
|
||||
|
||||
landmarks.append(fullyconnected1)
|
||||
return landmarks
|
||||
@@ -0,0 +1,133 @@
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
def nms(boxes, overlap_threshold=0.5, mode='union'):
|
||||
""" Pure Python NMS baseline. """
|
||||
x1 = boxes[:, 0]
|
||||
y1 = boxes[:, 1]
|
||||
x2 = boxes[:, 2]
|
||||
y2 = boxes[:, 3]
|
||||
scores = boxes[:, 4]
|
||||
|
||||
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
|
||||
order = scores.argsort()[::-1]
|
||||
|
||||
keep = []
|
||||
while order.size > 0:
|
||||
i = order[0]
|
||||
keep.append(i)
|
||||
xx1 = np.maximum(x1[i], x1[order[1:]])
|
||||
yy1 = np.maximum(y1[i], y1[order[1:]])
|
||||
xx2 = np.minimum(x2[i], x2[order[1:]])
|
||||
yy2 = np.minimum(y2[i], y2[order[1:]])
|
||||
|
||||
w = np.maximum(0.0, xx2 - xx1 + 1)
|
||||
h = np.maximum(0.0, yy2 - yy1 + 1)
|
||||
inter = w * h
|
||||
|
||||
if mode is 'min':
|
||||
ovr = inter / np.minimum(areas[i], areas[order[1:]])
|
||||
else:
|
||||
ovr = inter / (areas[i] + areas[order[1:]] - inter)
|
||||
|
||||
inds = np.where(ovr <= overlap_threshold)[0]
|
||||
order = order[inds + 1]
|
||||
|
||||
return keep
|
||||
|
||||
|
||||
def convert_to_square(bboxes):
|
||||
"""
|
||||
Convert bounding boxes to a square form.
|
||||
"""
|
||||
square_bboxes = np.zeros_like(bboxes)
|
||||
x1, y1, x2, y2 = [bboxes[:, i] for i in range(4)]
|
||||
h = y2 - y1 + 1.0
|
||||
w = x2 - x1 + 1.0
|
||||
max_side = np.maximum(h, w)
|
||||
square_bboxes[:, 0] = x1 + w*0.5 - max_side*0.5
|
||||
square_bboxes[:, 1] = y1 + h*0.5 - max_side*0.5
|
||||
square_bboxes[:, 2] = square_bboxes[:, 0] + max_side - 1.0
|
||||
square_bboxes[:, 3] = square_bboxes[:, 1] + max_side - 1.0
|
||||
return square_bboxes
|
||||
|
||||
|
||||
def calibrate_box(bboxes, offsets):
|
||||
"""Transform bounding boxes to be more like true bounding boxes.
|
||||
'offsets' is one of the outputs of the nets.
|
||||
"""
|
||||
x1, y1, x2, y2 = [bboxes[:, i] for i in range(4)]
|
||||
w = x2 - x1 + 1.0
|
||||
h = y2 - y1 + 1.0
|
||||
w = np.expand_dims(w, 1)
|
||||
h = np.expand_dims(h, 1)
|
||||
|
||||
translation = np.hstack([w, h, w, h])*offsets
|
||||
bboxes[:, 0:4] = bboxes[:, 0:4] + translation
|
||||
return bboxes
|
||||
|
||||
|
||||
def get_image_boxes(bounding_boxes, img, size=24):
|
||||
"""Cut out boxes from the image.
|
||||
"""
|
||||
num_boxes = len(bounding_boxes)
|
||||
(height, width, _) = img.shape
|
||||
|
||||
[dy, edy, dx, edx, y, ey, x, ex, w, h] = correct_bboxes(bounding_boxes, width, height)
|
||||
img_boxes = np.zeros((num_boxes, 3, size, size), 'float32')
|
||||
|
||||
for i in range(num_boxes):
|
||||
img_box = np.zeros((h[i], w[i], 3), 'uint8')
|
||||
|
||||
img_array = np.asarray(img, 'uint8')
|
||||
img_box[dy[i]:(edy[i] + 1), dx[i]:(edx[i] + 1), :] =\
|
||||
img_array[y[i]:(ey[i] + 1), x[i]:(ex[i] + 1), :]
|
||||
|
||||
img_box = cv2.resize(img_box, (size, size))
|
||||
img_box = np.asarray(img_box, 'float32')
|
||||
|
||||
img_boxes[i, :, :, :] = _preprocess(img_box)
|
||||
|
||||
return img_boxes
|
||||
|
||||
|
||||
def correct_bboxes(bboxes, width, height):
|
||||
"""Crop boxes that are too big and get coordinates
|
||||
with respect to cutouts.
|
||||
"""
|
||||
x1, y1, x2, y2 = [bboxes[:, i] for i in range(4)]
|
||||
w, h = x2 - x1 + 1.0, y2 - y1 + 1.0
|
||||
num_boxes = bboxes.shape[0]
|
||||
|
||||
x, y, ex, ey = x1, y1, x2, y2
|
||||
dx, dy = np.zeros((num_boxes,)), np.zeros((num_boxes,))
|
||||
edx, edy = w.copy() - 1.0, h.copy() - 1.0
|
||||
|
||||
ind = np.where(ex > width - 1.0)[0]
|
||||
edx[ind] = w[ind] + width - 2.0 - ex[ind]
|
||||
ex[ind] = width - 1.0
|
||||
|
||||
ind = np.where(ey > height - 1.0)[0]
|
||||
edy[ind] = h[ind] + height - 2.0 - ey[ind]
|
||||
ey[ind] = height - 1.0
|
||||
|
||||
ind = np.where(x < 0.0)[0]
|
||||
dx[ind] = 0.0 - x[ind]
|
||||
x[ind] = 0.0
|
||||
|
||||
ind = np.where(y < 0.0)[0]
|
||||
dy[ind] = 0.0 - y[ind]
|
||||
y[ind] = 0.0
|
||||
return_list = [dy, edy, dx, edx, y, ey, x, ex, w, h]
|
||||
return_list = [i.astype('int32') for i in return_list]
|
||||
|
||||
return return_list
|
||||
|
||||
|
||||
def _preprocess(img):
|
||||
"""Preprocessing step before feeding the network.
|
||||
"""
|
||||
img = img.transpose((2, 0, 1))
|
||||
img = np.expand_dims(img, 0)
|
||||
img = (img - 127.5)*0.0078125
|
||||
return img
|
||||
@@ -0,0 +1,42 @@
|
||||
# config.py
|
||||
|
||||
cfg_mnet = {
|
||||
'name': 'mobilenet0.25',
|
||||
'min_sizes': [[16, 32], [64, 128], [256, 512]],
|
||||
'steps': [8, 16, 32],
|
||||
'variance': [0.1, 0.2],
|
||||
'clip': False,
|
||||
'loc_weight': 2.0,
|
||||
'gpu_train': True,
|
||||
'batch_size': 32,
|
||||
'ngpu': 1,
|
||||
'epoch': 250,
|
||||
'decay1': 190,
|
||||
'decay2': 220,
|
||||
'image_size': 640,
|
||||
'pretrain': True,
|
||||
'return_layers': {'stage1': 1, 'stage2': 2, 'stage3': 3},
|
||||
'in_channel': 32,
|
||||
'out_channel': 64
|
||||
}
|
||||
|
||||
cfg_re50 = {
|
||||
'name': 'Resnet50',
|
||||
'min_sizes': [[16, 32], [64, 128], [256, 512]],
|
||||
'steps': [8, 16, 32],
|
||||
'variance': [0.1, 0.2],
|
||||
'clip': False,
|
||||
'loc_weight': 2.0,
|
||||
'gpu_train': True,
|
||||
'batch_size': 24,
|
||||
'ngpu': 4,
|
||||
'epoch': 100,
|
||||
'decay1': 70,
|
||||
'decay2': 90,
|
||||
'image_size': 840,
|
||||
'pretrain': True,
|
||||
'return_layers': {'layer2': 1, 'layer3': 2, 'layer4': 3},
|
||||
'in_channel': 256,
|
||||
'out_channel': 256
|
||||
}
|
||||
|
||||
@@ -0,0 +1,259 @@
|
||||
import math
|
||||
import numpy as np
|
||||
from .box_utils import nms, calibrate_box, get_image_boxes, convert_to_square, _preprocess
|
||||
import torch
|
||||
import cv2
|
||||
from .nms.py_cpu_nms import py_cpu_nms
|
||||
from core.utils import box_utils_Retina
|
||||
from .layers.functions.prior_box import PriorBox
|
||||
from .config import cfg_re50
|
||||
from .retinaface import RetinaFace
|
||||
from mtcnn.model import PNet, RNet, ONet
|
||||
|
||||
gpu_id = 0
|
||||
device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||
pnet, rnet, onet= PNet(), RNet(), ONet()
|
||||
pnet.to(device)
|
||||
rnet.to(device)
|
||||
onet.to(device)
|
||||
onet.eval()
|
||||
|
||||
|
||||
class RetinaFaceDetector(object):
|
||||
def __init__(self, gpu_id=None):
|
||||
self.gpu_id = gpu_id
|
||||
self.cfg = cfg_re50
|
||||
self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||
model = RetinaFace(cfg=self.cfg, phase='test')
|
||||
model = self.load_model(model, 'weights/Resnet50_Final.pth', True)
|
||||
model.eval()
|
||||
self.net = model.to(self.device)
|
||||
self.resize = 1
|
||||
self.confidence_threshold = 0.02
|
||||
self.top_k = 5000
|
||||
self.nms_threshold = 0.4
|
||||
self.keep_top_k = 750
|
||||
|
||||
def remove_prefix(self, state_dict, prefix):
|
||||
# print('remove prefix \'{}\''.format(prefix))
|
||||
f = lambda x: x.split(prefix, 1)[-1] if x.startswith(prefix) else x
|
||||
return {f(key): value for key, value in state_dict.items()}
|
||||
|
||||
def check_keys(self, model, pretrained_state_dict):
|
||||
ckpt_keys = set(pretrained_state_dict.keys())
|
||||
model_keys = set(model.state_dict().keys())
|
||||
used_pretrained_keys = model_keys & ckpt_keys
|
||||
# unused_pretrained_keys = ckpt_keys - model_keys
|
||||
# missing_keys = model_keys - ckpt_keys
|
||||
assert len(used_pretrained_keys) > 0, 'load NONE from pretrained checkpoint'
|
||||
return True
|
||||
|
||||
def load_model(self, model, pretrained_path, load_to_cpu):
|
||||
# print('Loading pretrained model from {}'.format(pretrained_path))
|
||||
if load_to_cpu:
|
||||
pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage)
|
||||
else:
|
||||
device = torch.cuda.current_device()
|
||||
pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage.cuda(device))
|
||||
if "state_dict" in pretrained_dict.keys():
|
||||
pretrained_dict = self.remove_prefix(pretrained_dict['state_dict'], 'module.')
|
||||
else:
|
||||
pretrained_dict = self.remove_prefix(pretrained_dict, 'module.')
|
||||
self.check_keys(model, pretrained_dict)
|
||||
model.load_state_dict(pretrained_dict, strict=False)
|
||||
return model
|
||||
|
||||
def forward(self, img_raw, min_face_size=50):
|
||||
img_scale = 640 / max(img_raw.shape[0], img_raw.shape[1])
|
||||
img = cv2.resize(img_raw, (0, 0), fx=img_scale, fy=img_scale)
|
||||
img = np.float32(img)
|
||||
|
||||
im_height, im_width, _ = img.shape
|
||||
scale = torch.Tensor([img.shape[1], img.shape[0], img.shape[1], img.shape[0]])
|
||||
img -= (104, 117, 123)
|
||||
img = img.transpose(2, 0, 1)
|
||||
img = torch.from_numpy(img).unsqueeze(0)
|
||||
img = img.to(self.device)
|
||||
scale = scale.to(self.device)
|
||||
|
||||
loc, conf, landms = self.net(img) # forward pass
|
||||
|
||||
priorbox = PriorBox(self.cfg, image_size=(im_height, im_width))
|
||||
priors = priorbox.forward()
|
||||
priors = priors.to(self.device)
|
||||
prior_data = priors.data
|
||||
boxes = box_utils_Retina.decode(loc.data.squeeze(0), prior_data, self.cfg['variance'])
|
||||
|
||||
boxes = boxes * scale / self.resize
|
||||
boxes = boxes.cpu().numpy()
|
||||
scores = conf.squeeze(0).data.cpu().numpy()[:, 1]
|
||||
landms = box_utils_Retina.decode_landm(landms.data.squeeze(0), prior_data, self.cfg['variance'])
|
||||
scale1 = torch.Tensor([img.shape[3], img.shape[2], img.shape[3], img.shape[2],
|
||||
img.shape[3], img.shape[2], img.shape[3], img.shape[2],
|
||||
img.shape[3], img.shape[2]])
|
||||
scale1 = scale1.to(self.device)
|
||||
landms = landms * scale1 / self.resize
|
||||
landms = landms.cpu().numpy()
|
||||
|
||||
# ignore low scores
|
||||
inds = np.where(scores > self.confidence_threshold)[0]
|
||||
boxes = boxes[inds]
|
||||
landms = landms[inds]
|
||||
scores = scores[inds]
|
||||
|
||||
# keep top-K before NMS
|
||||
order = scores.argsort()[::-1][:self.top_k]
|
||||
boxes = boxes[order]
|
||||
landms = landms[order]
|
||||
scores = scores[order]
|
||||
|
||||
# do NMS
|
||||
dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
|
||||
keep = py_cpu_nms(dets, self.nms_threshold, min_face_size = min_face_size * img_scale)
|
||||
# keep = nms(dets, args.nms_threshold,force_cpu=args.cpu)
|
||||
dets = dets[keep, :]
|
||||
landms = landms[keep]
|
||||
|
||||
# keep top-K faster NMS
|
||||
dets = dets[:self.keep_top_k, :]
|
||||
landms = landms[:self.keep_top_k, :]
|
||||
|
||||
dets[:, :4] = dets[:, :4] / img_scale
|
||||
landms /= img_scale
|
||||
return dets, landms
|
||||
|
||||
def forward_v2(self, image, min_face_size=20.0, thresholds=[0.6, 0.7, 0.8], nms_thresholds=[0.7, 0.7, 0.7]):
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
(height, width, _) = image.shape
|
||||
min_length = min(height, width)
|
||||
min_detection_size = 12
|
||||
factor = 0.707 # sqrt(0.5)
|
||||
|
||||
scales = []
|
||||
m = min_detection_size / min_face_size
|
||||
min_length *= m
|
||||
|
||||
factor_count = 0
|
||||
while min_length > min_detection_size:
|
||||
scales.append(m * factor ** factor_count)
|
||||
min_length *= factor
|
||||
factor_count += 1
|
||||
|
||||
# STAGE 1
|
||||
bounding_boxes = []
|
||||
for s in scales: # run P-Net on different scales
|
||||
boxes = run_first_stage(image, pnet, scale=s, threshold=thresholds[0], gpu_id=self.gpu_id)
|
||||
bounding_boxes.append(boxes)
|
||||
bounding_boxes = [i for i in bounding_boxes if i is not None]
|
||||
if len(bounding_boxes) == 0:
|
||||
return [], []
|
||||
bounding_boxes = np.vstack(bounding_boxes)
|
||||
|
||||
keep = nms(bounding_boxes[:, 0:5], nms_thresholds[0])
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes = calibrate_box(bounding_boxes[:, 0:5], bounding_boxes[:, 5:])
|
||||
bounding_boxes = convert_to_square(bounding_boxes)
|
||||
bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
|
||||
|
||||
# STAGE 2
|
||||
img_boxes = get_image_boxes(bounding_boxes, image, size=24)
|
||||
img_boxes = torch.from_numpy(img_boxes)
|
||||
img_boxes = img_boxes.to(self.device)
|
||||
output = rnet(img_boxes)
|
||||
offsets = output[0].to('cpu').data.numpy() # shape [n_boxes, 4]
|
||||
probs = output[1].to('cpu').data.numpy() # shape [n_boxes, 2]
|
||||
|
||||
keep = np.where(probs[:, 1] > thresholds[1])[0]
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
|
||||
offsets = offsets[keep]
|
||||
|
||||
keep = nms(bounding_boxes, nms_thresholds[1])
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes = calibrate_box(bounding_boxes, offsets[keep])
|
||||
bounding_boxes = convert_to_square(bounding_boxes)
|
||||
bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
|
||||
|
||||
# STAGE 3
|
||||
img_boxes = get_image_boxes(bounding_boxes, image, size=48)
|
||||
if len(img_boxes) == 0:
|
||||
return [], []
|
||||
img_boxes = torch.from_numpy(img_boxes)
|
||||
img_boxes = img_boxes.to(self.device)
|
||||
output = onet(img_boxes)
|
||||
landmarks = output[0].to('cpu').data.numpy() # shape [n_boxes, 10]
|
||||
offsets = output[1].to('cpu').data.numpy() # shape [n_boxes, 4]
|
||||
probs = output[2].to('cpu').data.numpy() # shape [n_boxes, 2]
|
||||
|
||||
keep = np.where(probs[:, 1] > thresholds[2])[0]
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
|
||||
offsets = offsets[keep]
|
||||
landmarks = landmarks[keep]
|
||||
|
||||
# compute landmark points
|
||||
width = bounding_boxes[:, 2] - bounding_boxes[:, 0] + 1.0
|
||||
height = bounding_boxes[:, 3] - bounding_boxes[:, 1] + 1.0
|
||||
xmin, ymin = bounding_boxes[:, 0], bounding_boxes[:, 1]
|
||||
landmarks[:, 0:5] = np.expand_dims(xmin, 1) + np.expand_dims(width, 1) * landmarks[:, 0:5]
|
||||
landmarks[:, 5:10] = np.expand_dims(ymin, 1) + np.expand_dims(height, 1) * landmarks[:, 5:10]
|
||||
|
||||
bounding_boxes = calibrate_box(bounding_boxes, offsets)
|
||||
keep = nms(bounding_boxes, nms_thresholds[2], mode='min')
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
landmarks = landmarks[keep]
|
||||
|
||||
return bounding_boxes, landmarks
|
||||
|
||||
def run_first_stage(image, net, scale, threshold, gpu_id=0):
|
||||
"""
|
||||
Run P-Net, generate bounding boxes, and do NMS.
|
||||
"""
|
||||
device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||
(height, width, _) = image.shape
|
||||
sw, sh = math.ceil(width * scale), math.ceil(height * scale)
|
||||
img = cv2.resize(image, (sw, sh))
|
||||
# img = image.resize((sw, sh), Image.BILINEAR)
|
||||
img = np.asarray(img, 'float32')
|
||||
img = torch.from_numpy(_preprocess(img))
|
||||
img = img.to(device)
|
||||
|
||||
output = net(img)
|
||||
probs = output[1].to('cpu').data.numpy()[0, 1, :, :]
|
||||
offsets = output[0].to('cpu').data.numpy()
|
||||
|
||||
boxes = _generate_bboxes(probs, offsets, scale, threshold)
|
||||
if len(boxes) == 0:
|
||||
return None
|
||||
|
||||
keep = nms(boxes[:, 0:5], overlap_threshold=0.5)
|
||||
return boxes[keep]
|
||||
|
||||
|
||||
def _generate_bboxes(probs, offsets, scale, threshold):
|
||||
"""
|
||||
Generate bounding boxes at places where there is probably a face.
|
||||
"""
|
||||
stride = 2
|
||||
cell_size = 12
|
||||
|
||||
inds = np.where(probs > threshold)
|
||||
|
||||
if inds[0].size == 0:
|
||||
return np.array([])
|
||||
|
||||
tx1, ty1, tx2, ty2 = [offsets[0, i, inds[0], inds[1]] for i in range(4)]
|
||||
|
||||
offsets = np.array([tx1, ty1, tx2, ty2])
|
||||
score = probs[inds[0], inds[1]]
|
||||
|
||||
# P-Net is applied to scaled images, so we need to rescale bounding boxes back
|
||||
bounding_boxes = np.vstack([
|
||||
np.round((stride * inds[1] + 1.0) / scale),
|
||||
np.round((stride * inds[0] + 1.0) / scale),
|
||||
np.round((stride * inds[1] + 1.0 + cell_size) / scale),
|
||||
np.round((stride * inds[0] + 1.0 + cell_size) / scale),
|
||||
score, offsets
|
||||
])
|
||||
|
||||
return bounding_boxes.T
|
||||
@@ -0,0 +1,2 @@
|
||||
from .functions import *
|
||||
from .modules import *
|
||||
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)
|
||||
@@ -0,0 +1,34 @@
|
||||
import torch
|
||||
from itertools import product as product
|
||||
import numpy as np
|
||||
from math import ceil
|
||||
|
||||
|
||||
class PriorBox(object):
|
||||
def __init__(self, cfg, image_size=None, phase='train'):
|
||||
super(PriorBox, self).__init__()
|
||||
self.min_sizes = cfg['min_sizes']
|
||||
self.steps = cfg['steps']
|
||||
self.clip = cfg['clip']
|
||||
self.image_size = image_size
|
||||
self.feature_maps = [[ceil(self.image_size[0]/step), ceil(self.image_size[1]/step)] for step in self.steps]
|
||||
self.name = "s"
|
||||
|
||||
def forward(self):
|
||||
anchors = []
|
||||
for k, f in enumerate(self.feature_maps):
|
||||
min_sizes = self.min_sizes[k]
|
||||
for i, j in product(range(f[0]), range(f[1])):
|
||||
for min_size in min_sizes:
|
||||
s_kx = min_size / self.image_size[1]
|
||||
s_ky = min_size / self.image_size[0]
|
||||
dense_cx = [x * self.steps[k] / self.image_size[1] for x in [j + 0.5]]
|
||||
dense_cy = [y * self.steps[k] / self.image_size[0] for y in [i + 0.5]]
|
||||
for cy, cx in product(dense_cy, dense_cx):
|
||||
anchors += [cx, cy, s_kx, s_ky]
|
||||
|
||||
# back to torch land
|
||||
output = torch.Tensor(anchors).view(-1, 4)
|
||||
if self.clip:
|
||||
output.clamp_(max=1, min=0)
|
||||
return output
|
||||
@@ -0,0 +1,3 @@
|
||||
from .multibox_loss import MultiBoxLoss
|
||||
|
||||
__all__ = ['MultiBoxLoss']
|
||||
@@ -0,0 +1,125 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.autograd import Variable
|
||||
from core.utils.box_utils_Retina import match, log_sum_exp
|
||||
from core.models.layers.data import cfg_mnet
|
||||
GPU = cfg_mnet['gpu_train']
|
||||
|
||||
class MultiBoxLoss(nn.Module):
|
||||
"""SSD Weighted Loss Function
|
||||
Compute Targets:
|
||||
1) Produce Confidence Target Indices by matching ground truth boxes
|
||||
with (default) 'priorboxes' that have jaccard index > threshold parameter
|
||||
(default threshold: 0.5).
|
||||
2) Produce localization target by 'encoding' variance into offsets of ground
|
||||
truth boxes and their matched 'priorboxes'.
|
||||
3) Hard negative mining to filter the excessive number of negative examples
|
||||
that comes with using a large number of default bounding boxes.
|
||||
(default negative:positive ratio 3:1)
|
||||
Objective Loss:
|
||||
L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N
|
||||
Where, Lconf is the CrossEntropy Loss and Lloc is the SmoothL1 Loss
|
||||
weighted by α which is set to 1 by cross val.
|
||||
Args:
|
||||
c: class confidences,
|
||||
l: predicted boxes,
|
||||
g: ground truth boxes
|
||||
N: number of matched default boxes
|
||||
See: https://arxiv.org/pdf/1512.02325.pdf for more details.
|
||||
"""
|
||||
|
||||
def __init__(self, num_classes, overlap_thresh, prior_for_matching, bkg_label, neg_mining, neg_pos, neg_overlap, encode_target):
|
||||
super(MultiBoxLoss, self).__init__()
|
||||
self.num_classes = num_classes
|
||||
self.threshold = overlap_thresh
|
||||
self.background_label = bkg_label
|
||||
self.encode_target = encode_target
|
||||
self.use_prior_for_matching = prior_for_matching
|
||||
self.do_neg_mining = neg_mining
|
||||
self.negpos_ratio = neg_pos
|
||||
self.neg_overlap = neg_overlap
|
||||
self.variance = [0.1, 0.2]
|
||||
|
||||
def forward(self, predictions, priors, targets):
|
||||
"""Multibox Loss
|
||||
Args:
|
||||
predictions (tuple): A tuple containing loc preds, conf preds,
|
||||
and prior boxes from SSD net.
|
||||
conf shape: torch.size(batch_size,num_priors,num_classes)
|
||||
loc shape: torch.size(batch_size,num_priors,4)
|
||||
priors shape: torch.size(num_priors,4)
|
||||
|
||||
ground_truth (tensor): Ground truth boxes and labels for a batch,
|
||||
shape: [batch_size,num_objs,5] (last idx is the label).
|
||||
"""
|
||||
|
||||
loc_data, conf_data, landm_data = predictions
|
||||
priors = priors
|
||||
num = loc_data.size(0)
|
||||
num_priors = (priors.size(0))
|
||||
|
||||
# match priors (default boxes) and ground truth boxes
|
||||
loc_t = torch.Tensor(num, num_priors, 4)
|
||||
landm_t = torch.Tensor(num, num_priors, 10)
|
||||
conf_t = torch.LongTensor(num, num_priors)
|
||||
for idx in range(num):
|
||||
truths = targets[idx][:, :4].data
|
||||
labels = targets[idx][:, -1].data
|
||||
landms = targets[idx][:, 4:14].data
|
||||
defaults = priors.data
|
||||
match(self.threshold, truths, defaults, self.variance, labels, landms, loc_t, conf_t, landm_t, idx)
|
||||
if GPU:
|
||||
loc_t = loc_t.cuda()
|
||||
conf_t = conf_t.cuda()
|
||||
landm_t = landm_t.cuda()
|
||||
|
||||
zeros = torch.tensor(0).cuda()
|
||||
# landm Loss (Smooth L1)
|
||||
# Shape: [batch,num_priors,10]
|
||||
pos1 = conf_t > zeros
|
||||
num_pos_landm = pos1.long().sum(1, keepdim=True)
|
||||
N1 = max(num_pos_landm.data.sum().float(), 1)
|
||||
pos_idx1 = pos1.unsqueeze(pos1.dim()).expand_as(landm_data)
|
||||
landm_p = landm_data[pos_idx1].view(-1, 10)
|
||||
landm_t = landm_t[pos_idx1].view(-1, 10)
|
||||
loss_landm = F.smooth_l1_loss(landm_p, landm_t, reduction='sum')
|
||||
|
||||
|
||||
pos = conf_t != zeros
|
||||
conf_t[pos] = 1
|
||||
|
||||
# Localization Loss (Smooth L1)
|
||||
# Shape: [batch,num_priors,4]
|
||||
pos_idx = pos.unsqueeze(pos.dim()).expand_as(loc_data)
|
||||
loc_p = loc_data[pos_idx].view(-1, 4)
|
||||
loc_t = loc_t[pos_idx].view(-1, 4)
|
||||
loss_l = F.smooth_l1_loss(loc_p, loc_t, reduction='sum')
|
||||
|
||||
# Compute max conf across batch for hard negative mining
|
||||
batch_conf = conf_data.view(-1, self.num_classes)
|
||||
loss_c = log_sum_exp(batch_conf) - batch_conf.gather(1, conf_t.view(-1, 1))
|
||||
|
||||
# Hard Negative Mining
|
||||
loss_c[pos.view(-1, 1)] = 0 # filter out pos boxes for now
|
||||
loss_c = loss_c.view(num, -1)
|
||||
_, loss_idx = loss_c.sort(1, descending=True)
|
||||
_, idx_rank = loss_idx.sort(1)
|
||||
num_pos = pos.long().sum(1, keepdim=True)
|
||||
num_neg = torch.clamp(self.negpos_ratio*num_pos, max=pos.size(1)-1)
|
||||
neg = idx_rank < num_neg.expand_as(idx_rank)
|
||||
|
||||
# Confidence Loss Including Positive and Negative Examples
|
||||
pos_idx = pos.unsqueeze(2).expand_as(conf_data)
|
||||
neg_idx = neg.unsqueeze(2).expand_as(conf_data)
|
||||
conf_p = conf_data[(pos_idx+neg_idx).gt(0)].view(-1,self.num_classes)
|
||||
targets_weighted = conf_t[(pos+neg).gt(0)]
|
||||
loss_c = F.cross_entropy(conf_p, targets_weighted, reduction='sum')
|
||||
|
||||
# Sum of losses: L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N
|
||||
N = max(num_pos.data.sum().float(), 1)
|
||||
loss_l /= N
|
||||
loss_c /= N
|
||||
loss_landm /= N1
|
||||
|
||||
return loss_l, loss_c, loss_landm
|
||||
@@ -0,0 +1,137 @@
|
||||
import time
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
# import torchvision.models._utils as _utils
|
||||
# import torchvision.models as models
|
||||
import torch.nn.functional as F
|
||||
# from torch.autograd import Variable
|
||||
|
||||
def conv_bn(inp, oup, stride = 1, leaky = 0):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
|
||||
nn.BatchNorm2d(oup),
|
||||
nn.LeakyReLU(negative_slope=leaky, inplace=True)
|
||||
)
|
||||
|
||||
def conv_bn_no_relu(inp, oup, stride):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
|
||||
nn.BatchNorm2d(oup),
|
||||
)
|
||||
|
||||
def conv_bn1X1(inp, oup, stride, leaky=0):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(inp, oup, 1, stride, padding=0, bias=False),
|
||||
nn.BatchNorm2d(oup),
|
||||
nn.LeakyReLU(negative_slope=leaky, inplace=True)
|
||||
)
|
||||
|
||||
def conv_dw(inp, oup, stride, leaky=0.1):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
|
||||
nn.BatchNorm2d(inp),
|
||||
nn.LeakyReLU(negative_slope= leaky,inplace=True),
|
||||
|
||||
nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
|
||||
nn.BatchNorm2d(oup),
|
||||
nn.LeakyReLU(negative_slope= leaky,inplace=True),
|
||||
)
|
||||
|
||||
class SSH(nn.Module):
|
||||
def __init__(self, in_channel, out_channel):
|
||||
super(SSH, self).__init__()
|
||||
assert out_channel % 4 == 0
|
||||
leaky = 0
|
||||
if (out_channel <= 64):
|
||||
leaky = 0.1
|
||||
self.conv3X3 = conv_bn_no_relu(in_channel, out_channel//2, stride=1)
|
||||
|
||||
self.conv5X5_1 = conv_bn(in_channel, out_channel//4, stride=1, leaky = leaky)
|
||||
self.conv5X5_2 = conv_bn_no_relu(out_channel//4, out_channel//4, stride=1)
|
||||
|
||||
self.conv7X7_2 = conv_bn(out_channel//4, out_channel//4, stride=1, leaky = leaky)
|
||||
self.conv7x7_3 = conv_bn_no_relu(out_channel//4, out_channel//4, stride=1)
|
||||
|
||||
def forward(self, input):
|
||||
conv3X3 = self.conv3X3(input)
|
||||
|
||||
conv5X5_1 = self.conv5X5_1(input)
|
||||
conv5X5 = self.conv5X5_2(conv5X5_1)
|
||||
|
||||
conv7X7_2 = self.conv7X7_2(conv5X5_1)
|
||||
conv7X7 = self.conv7x7_3(conv7X7_2)
|
||||
|
||||
out = torch.cat([conv3X3, conv5X5, conv7X7], dim=1)
|
||||
out = F.relu(out)
|
||||
return out
|
||||
|
||||
class FPN(nn.Module):
|
||||
def __init__(self,in_channels_list,out_channels):
|
||||
super(FPN,self).__init__()
|
||||
leaky = 0
|
||||
if (out_channels <= 64):
|
||||
leaky = 0.1
|
||||
self.output1 = conv_bn1X1(in_channels_list[0], out_channels, stride = 1, leaky = leaky)
|
||||
self.output2 = conv_bn1X1(in_channels_list[1], out_channels, stride = 1, leaky = leaky)
|
||||
self.output3 = conv_bn1X1(in_channels_list[2], out_channels, stride = 1, leaky = leaky)
|
||||
|
||||
self.merge1 = conv_bn(out_channels, out_channels, leaky = leaky)
|
||||
self.merge2 = conv_bn(out_channels, out_channels, leaky = leaky)
|
||||
|
||||
def forward(self, input):
|
||||
# names = list(input.keys())
|
||||
input = list(input.values())
|
||||
|
||||
output1 = self.output1(input[0])
|
||||
output2 = self.output2(input[1])
|
||||
output3 = self.output3(input[2])
|
||||
|
||||
up3 = F.interpolate(output3, size=[output2.size(2), output2.size(3)], mode="nearest")
|
||||
output2 = output2 + up3
|
||||
output2 = self.merge2(output2)
|
||||
|
||||
up2 = F.interpolate(output2, size=[output1.size(2), output1.size(3)], mode="nearest")
|
||||
output1 = output1 + up2
|
||||
output1 = self.merge1(output1)
|
||||
|
||||
out = [output1, output2, output3]
|
||||
return out
|
||||
|
||||
|
||||
|
||||
class MobileNetV1(nn.Module):
|
||||
def __init__(self):
|
||||
super(MobileNetV1, self).__init__()
|
||||
self.stage1 = nn.Sequential(
|
||||
conv_bn(3, 8, 2, leaky = 0.1), # 3
|
||||
conv_dw(8, 16, 1), # 7
|
||||
conv_dw(16, 32, 2), # 11
|
||||
conv_dw(32, 32, 1), # 19
|
||||
conv_dw(32, 64, 2), # 27
|
||||
conv_dw(64, 64, 1), # 43
|
||||
)
|
||||
self.stage2 = nn.Sequential(
|
||||
conv_dw(64, 128, 2), # 43 + 16 = 59
|
||||
conv_dw(128, 128, 1), # 59 + 32 = 91
|
||||
conv_dw(128, 128, 1), # 91 + 32 = 123
|
||||
conv_dw(128, 128, 1), # 123 + 32 = 155
|
||||
conv_dw(128, 128, 1), # 155 + 32 = 187
|
||||
conv_dw(128, 128, 1), # 187 + 32 = 219
|
||||
)
|
||||
self.stage3 = nn.Sequential(
|
||||
conv_dw(128, 256, 2), # 219 +3 2 = 241
|
||||
conv_dw(256, 256, 1), # 241 + 64 = 301
|
||||
)
|
||||
self.avg = nn.AdaptiveAvgPool2d((1,1))
|
||||
self.fc = nn.Linear(256, 1000)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.stage1(x)
|
||||
x = self.stage2(x)
|
||||
x = self.stage3(x)
|
||||
x = self.avg(x)
|
||||
# x = self.model(x)
|
||||
x = x.view(-1, 256)
|
||||
x = self.fc(x)
|
||||
return x
|
||||
|
||||
@@ -0,0 +1,46 @@
|
||||
# --------------------------------------------------------
|
||||
# Fast R-CNN
|
||||
# Copyright (c) 2015 Microsoft
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# Written by Ross Girshick
|
||||
# --------------------------------------------------------
|
||||
|
||||
import numpy as np
|
||||
|
||||
def py_cpu_nms(dets, thresh, min_face_size = 50):
|
||||
"""Pure Python NMS baseline."""
|
||||
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]
|
||||
|
||||
#filter_small_faces
|
||||
filter_list = []
|
||||
for idx in keep:
|
||||
w = np.abs(dets[idx, 2] - dets[idx, 0])
|
||||
h = np.abs(dets[idx, 3] - dets[idx, 1])
|
||||
if max(w, h) < min_face_size: continue
|
||||
filter_list.append(idx)
|
||||
|
||||
return filter_list
|
||||
@@ -0,0 +1,342 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
|
||||
'resnet152', 'resnext50_32x4d', 'resnext101_32x8d',
|
||||
'wide_resnet50_2', 'wide_resnet101_2']
|
||||
|
||||
|
||||
model_urls = {
|
||||
'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
|
||||
'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
|
||||
'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
|
||||
'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
|
||||
'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
|
||||
'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
|
||||
'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
|
||||
'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',
|
||||
'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth',
|
||||
}
|
||||
|
||||
|
||||
def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
|
||||
"""3x3 convolution with padding"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
|
||||
padding=dilation, groups=groups, bias=False, dilation=dilation)
|
||||
|
||||
|
||||
def conv1x1(in_planes, out_planes, stride=1):
|
||||
"""1x1 convolution"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
|
||||
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
expansion = 1
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
|
||||
base_width=64, dilation=1, norm_layer=None):
|
||||
super(BasicBlock, self).__init__()
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
if groups != 1 or base_width != 64:
|
||||
raise ValueError('BasicBlock only supports groups=1 and base_width=64')
|
||||
if dilation > 1:
|
||||
raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
|
||||
# Both self.conv1 and self.downsample layers downsample the input when stride != 1
|
||||
self.conv1 = conv3x3(inplanes, planes, stride)
|
||||
self.bn1 = norm_layer(planes)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.conv2 = conv3x3(planes, planes)
|
||||
self.bn2 = norm_layer(planes)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
identity = self.downsample(x)
|
||||
|
||||
out += identity
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class Bottleneck(nn.Module):
|
||||
# Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
|
||||
# while original implementation places the stride at the first 1x1 convolution(self.conv1)
|
||||
# according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
|
||||
# This variant is also known as ResNet V1.5 and improves accuracy according to
|
||||
# https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
|
||||
|
||||
expansion = 4
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
|
||||
base_width=64, dilation=1, norm_layer=None):
|
||||
super(Bottleneck, self).__init__()
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
width = int(planes * (base_width / 64.)) * groups
|
||||
# Both self.conv2 and self.downsample layers downsample the input when stride != 1
|
||||
self.conv1 = conv1x1(inplanes, width)
|
||||
self.bn1 = norm_layer(width)
|
||||
self.conv2 = conv3x3(width, width, stride, groups, dilation)
|
||||
self.bn2 = norm_layer(width)
|
||||
self.conv3 = conv1x1(width, planes * self.expansion)
|
||||
self.bn3 = norm_layer(planes * self.expansion)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
identity = self.downsample(x)
|
||||
|
||||
out += identity
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class ResNet(nn.Module):
|
||||
def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,
|
||||
groups=1, width_per_group=64, replace_stride_with_dilation=None,
|
||||
norm_layer=None):
|
||||
super(ResNet, self).__init__()
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
self._norm_layer = norm_layer
|
||||
|
||||
self.inplanes = 64
|
||||
self.dilation = 1
|
||||
if replace_stride_with_dilation is None:
|
||||
# each element in the tuple indicates if we should replace
|
||||
# the 2x2 stride with a dilated convolution instead
|
||||
replace_stride_with_dilation = [False, False, False]
|
||||
if len(replace_stride_with_dilation) != 3:
|
||||
raise ValueError("replace_stride_with_dilation should be None "
|
||||
"or a 3-element tuple, got {}".format(replace_stride_with_dilation))
|
||||
self.groups = groups
|
||||
self.base_width = width_per_group
|
||||
self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=2, padding=3, bias=False) ### ycj
|
||||
self.bn1 = norm_layer(self.inplanes)
|
||||
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,
|
||||
dilate=replace_stride_with_dilation[0])
|
||||
self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
|
||||
dilate=replace_stride_with_dilation[1])
|
||||
self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
|
||||
dilate=replace_stride_with_dilation[2])
|
||||
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
|
||||
self.fc = nn.Linear(512 * block.expansion, num_classes)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
||||
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
# Zero-initialize the last BN in each residual branch,
|
||||
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
|
||||
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
|
||||
if zero_init_residual:
|
||||
for m in self.modules():
|
||||
if isinstance(m, Bottleneck):
|
||||
nn.init.constant_(m.bn3.weight, 0)
|
||||
elif isinstance(m, BasicBlock):
|
||||
nn.init.constant_(m.bn2.weight, 0)
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
|
||||
norm_layer = self._norm_layer
|
||||
downsample = None
|
||||
previous_dilation = self.dilation
|
||||
if dilate:
|
||||
self.dilation *= stride
|
||||
stride = 1
|
||||
if stride != 1 or self.inplanes != planes * block.expansion:
|
||||
downsample = nn.Sequential(
|
||||
conv1x1(self.inplanes, planes * block.expansion, stride),
|
||||
norm_layer(planes * block.expansion),
|
||||
)
|
||||
|
||||
layers = []
|
||||
layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
|
||||
self.base_width, previous_dilation, norm_layer))
|
||||
self.inplanes = planes * block.expansion
|
||||
for _ in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes, groups=self.groups,
|
||||
base_width=self.base_width, dilation=self.dilation,
|
||||
norm_layer=norm_layer))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def _forward_impl(self, x):
|
||||
# See note [TorchScript super()]
|
||||
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.avgpool(x)
|
||||
x = torch.flatten(x, 1)
|
||||
x = self.fc(x)
|
||||
|
||||
return x
|
||||
|
||||
def forward(self, x):
|
||||
return self._forward_impl(x)
|
||||
|
||||
|
||||
def _resnet(arch, block, layers, pretrained, progress, n_class, **kwargs):
|
||||
model = ResNet(block, layers, num_classes=n_class, **kwargs)
|
||||
# if pretrained:
|
||||
# state_dict = load_state_dict_from_url(model_urls[arch], progress=progress)
|
||||
# model.load_state_dict(state_dict)
|
||||
return model
|
||||
|
||||
|
||||
def resnet18(n_class=1000, pretrained=False, progress=True, **kwargs):
|
||||
r"""ResNet-18 model from
|
||||
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress, n_class,
|
||||
**kwargs)
|
||||
|
||||
|
||||
def resnet34(pretrained=False, progress=True, **kwargs):
|
||||
r"""ResNet-34 model from
|
||||
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress,
|
||||
**kwargs)
|
||||
|
||||
|
||||
def resnet50(n_class=1000, pretrained=False, progress=True, **kwargs):
|
||||
r"""ResNet-50 model from
|
||||
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress, n_class, **kwargs)
|
||||
|
||||
def resnet101(n_class=1000, pretrained=False, progress=True, **kwargs):
|
||||
r"""ResNet-101 model from
|
||||
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress, n_class, **kwargs)
|
||||
|
||||
|
||||
def resnet152(n_class=1000, pretrained=False, progress=True, **kwargs):
|
||||
r"""ResNet-152 model from
|
||||
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress, n_class,
|
||||
**kwargs)
|
||||
|
||||
|
||||
def resnext50_32x4d(pretrained=False, progress=True, **kwargs):
|
||||
r"""ResNeXt-50 32x4d model from
|
||||
`"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
kwargs['groups'] = 32
|
||||
kwargs['width_per_group'] = 4
|
||||
return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3],
|
||||
pretrained, progress, **kwargs)
|
||||
|
||||
|
||||
def resnext101_32x8d(pretrained=False, progress=True, **kwargs):
|
||||
r"""ResNeXt-101 32x8d model from
|
||||
`"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
kwargs['groups'] = 32
|
||||
kwargs['width_per_group'] = 8
|
||||
return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3],
|
||||
pretrained, progress, **kwargs)
|
||||
|
||||
|
||||
def wide_resnet50_2(pretrained=False, progress=True, **kwargs):
|
||||
r"""Wide ResNet-50-2 model from
|
||||
`"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_
|
||||
The model is the same as ResNet except for the bottleneck number of channels
|
||||
which is twice larger in every block. The number of channels in outer 1x1
|
||||
convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
|
||||
channels, and in Wide ResNet-50-2 has 2048-1024-2048.
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
kwargs['width_per_group'] = 64 * 2
|
||||
return _resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3],
|
||||
pretrained, progress, **kwargs)
|
||||
|
||||
|
||||
def wide_resnet101_2(pretrained=False, progress=True, **kwargs):
|
||||
r"""Wide ResNet-101-2 model from
|
||||
`"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_
|
||||
The model is the same as ResNet except for the bottleneck number of channels
|
||||
which is twice larger in every block. The number of channels in outer 1x1
|
||||
convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
|
||||
channels, and in Wide ResNet-50-2 has 2048-1024-2048.
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
kwargs['width_per_group'] = 64 * 2
|
||||
return _resnet('wide_resnet101_2', Bottleneck, [3, 4, 23, 3],
|
||||
pretrained, progress, **kwargs)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
net = resnet18(n_class=33)
|
||||
input_size = (1, 3, 224, 224)
|
||||
x = torch.randn(input_size)
|
||||
out = net(x)
|
||||
print(out.shape)
|
||||
@@ -0,0 +1,127 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
# import torchvision.models.detection.backbone_utils as backbone_utils
|
||||
import torchvision.models._utils as _utils
|
||||
import torch.nn.functional as F
|
||||
# from collections import OrderedDict
|
||||
|
||||
from core.models.net import MobileNetV1 as MobileNetV1
|
||||
from core.models.net import FPN as FPN
|
||||
from core.models.net import SSH as SSH
|
||||
|
||||
|
||||
|
||||
class ClassHead(nn.Module):
|
||||
def __init__(self,inchannels=512,num_anchors=3):
|
||||
super(ClassHead,self).__init__()
|
||||
self.num_anchors = num_anchors
|
||||
self.conv1x1 = nn.Conv2d(inchannels,self.num_anchors*2,kernel_size=(1,1),stride=1,padding=0)
|
||||
|
||||
def forward(self,x):
|
||||
out = self.conv1x1(x)
|
||||
out = out.permute(0,2,3,1).contiguous()
|
||||
|
||||
return out.view(out.shape[0], -1, 2)
|
||||
|
||||
class BboxHead(nn.Module):
|
||||
def __init__(self,inchannels=512,num_anchors=3):
|
||||
super(BboxHead,self).__init__()
|
||||
self.conv1x1 = nn.Conv2d(inchannels,num_anchors*4,kernel_size=(1,1),stride=1,padding=0)
|
||||
|
||||
def forward(self,x):
|
||||
out = self.conv1x1(x)
|
||||
out = out.permute(0,2,3,1).contiguous()
|
||||
|
||||
return out.view(out.shape[0], -1, 4)
|
||||
|
||||
class LandmarkHead(nn.Module):
|
||||
def __init__(self,inchannels=512,num_anchors=3):
|
||||
super(LandmarkHead,self).__init__()
|
||||
self.conv1x1 = nn.Conv2d(inchannels,num_anchors*10,kernel_size=(1,1),stride=1,padding=0)
|
||||
|
||||
def forward(self,x):
|
||||
out = self.conv1x1(x)
|
||||
out = out.permute(0,2,3,1).contiguous()
|
||||
|
||||
return out.view(out.shape[0], -1, 10)
|
||||
|
||||
class RetinaFace(nn.Module):
|
||||
def __init__(self, cfg = None, phase = 'train'):
|
||||
"""
|
||||
:param cfg: Network related settings.
|
||||
:param phase: train or test.
|
||||
"""
|
||||
super(RetinaFace,self).__init__()
|
||||
self.phase = phase
|
||||
backbone = None
|
||||
if cfg['name'] == 'mobilenet0.25':
|
||||
backbone = MobileNetV1()
|
||||
if cfg['pretrain']:
|
||||
checkpoint = torch.load("./weights/mobilenetV1X0.25_pretrain.tar", map_location=torch.device('cpu'))
|
||||
from collections import OrderedDict
|
||||
new_state_dict = OrderedDict()
|
||||
for k, v in checkpoint['state_dict'].items():
|
||||
name = k[7:] # remove module.
|
||||
new_state_dict[name] = v
|
||||
# load params
|
||||
backbone.load_state_dict(new_state_dict)
|
||||
elif cfg['name'] == 'Resnet50':
|
||||
import torchvision.models as models
|
||||
backbone = models.resnet50(pretrained=cfg['pretrain'])
|
||||
|
||||
self.body = _utils.IntermediateLayerGetter(backbone, cfg['return_layers'])
|
||||
in_channels_stage2 = cfg['in_channel']
|
||||
in_channels_list = [
|
||||
in_channels_stage2 * 2,
|
||||
in_channels_stage2 * 4,
|
||||
in_channels_stage2 * 8,
|
||||
]
|
||||
out_channels = cfg['out_channel']
|
||||
self.fpn = FPN(in_channels_list,out_channels)
|
||||
self.ssh1 = SSH(out_channels, out_channels)
|
||||
self.ssh2 = SSH(out_channels, out_channels)
|
||||
self.ssh3 = SSH(out_channels, out_channels)
|
||||
|
||||
self.ClassHead = self._make_class_head(fpn_num=3, inchannels=cfg['out_channel'])
|
||||
self.BboxHead = self._make_bbox_head(fpn_num=3, inchannels=cfg['out_channel'])
|
||||
self.LandmarkHead = self._make_landmark_head(fpn_num=3, inchannels=cfg['out_channel'])
|
||||
|
||||
def _make_class_head(self,fpn_num=3,inchannels=64,anchor_num=2):
|
||||
classhead = nn.ModuleList()
|
||||
for i in range(fpn_num):
|
||||
classhead.append(ClassHead(inchannels,anchor_num))
|
||||
return classhead
|
||||
|
||||
def _make_bbox_head(self,fpn_num=3,inchannels=64,anchor_num=2):
|
||||
bboxhead = nn.ModuleList()
|
||||
for i in range(fpn_num):
|
||||
bboxhead.append(BboxHead(inchannels,anchor_num))
|
||||
return bboxhead
|
||||
|
||||
def _make_landmark_head(self,fpn_num=3,inchannels=64,anchor_num=2):
|
||||
landmarkhead = nn.ModuleList()
|
||||
for i in range(fpn_num):
|
||||
landmarkhead.append(LandmarkHead(inchannels,anchor_num))
|
||||
return landmarkhead
|
||||
|
||||
def forward(self,inputs):
|
||||
out = self.body(inputs)
|
||||
|
||||
# FPN
|
||||
fpn = self.fpn(out)
|
||||
|
||||
# SSH
|
||||
feature1 = self.ssh1(fpn[0])
|
||||
feature2 = self.ssh2(fpn[1])
|
||||
feature3 = self.ssh3(fpn[2])
|
||||
features = [feature1, feature2, feature3]
|
||||
|
||||
bbox_regressions = torch.cat([self.BboxHead[i](feature) for i, feature in enumerate(features)], dim=1)
|
||||
classifications = torch.cat([self.ClassHead[i](feature) for i, feature in enumerate(features)],dim=1)
|
||||
ldm_regressions = torch.cat([self.LandmarkHead[i](feature) for i, feature in enumerate(features)], dim=1)
|
||||
|
||||
if self.phase == 'train':
|
||||
output = (bbox_regressions, classifications, ldm_regressions)
|
||||
else:
|
||||
output = (bbox_regressions, F.softmax(classifications, dim=-1), ldm_regressions)
|
||||
return output
|
||||
@@ -0,0 +1 @@
|
||||
|
||||
@@ -0,0 +1,133 @@
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
def nms(boxes, overlap_threshold=0.5, mode='union'):
|
||||
""" Pure Python NMS baseline. """
|
||||
x1 = boxes[:, 0]
|
||||
y1 = boxes[:, 1]
|
||||
x2 = boxes[:, 2]
|
||||
y2 = boxes[:, 3]
|
||||
scores = boxes[:, 4]
|
||||
|
||||
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
|
||||
order = scores.argsort()[::-1]
|
||||
|
||||
keep = []
|
||||
while order.size > 0:
|
||||
i = order[0]
|
||||
keep.append(i)
|
||||
xx1 = np.maximum(x1[i], x1[order[1:]])
|
||||
yy1 = np.maximum(y1[i], y1[order[1:]])
|
||||
xx2 = np.minimum(x2[i], x2[order[1:]])
|
||||
yy2 = np.minimum(y2[i], y2[order[1:]])
|
||||
|
||||
w = np.maximum(0.0, xx2 - xx1 + 1)
|
||||
h = np.maximum(0.0, yy2 - yy1 + 1)
|
||||
inter = w * h
|
||||
|
||||
if mode is 'min':
|
||||
ovr = inter / np.minimum(areas[i], areas[order[1:]])
|
||||
else:
|
||||
ovr = inter / (areas[i] + areas[order[1:]] - inter)
|
||||
|
||||
inds = np.where(ovr <= overlap_threshold)[0]
|
||||
order = order[inds + 1]
|
||||
|
||||
return keep
|
||||
|
||||
|
||||
def convert_to_square(bboxes):
|
||||
"""
|
||||
Convert bounding boxes to a square form.
|
||||
"""
|
||||
square_bboxes = np.zeros_like(bboxes)
|
||||
x1, y1, x2, y2 = [bboxes[:, i] for i in range(4)]
|
||||
h = y2 - y1 + 1.0
|
||||
w = x2 - x1 + 1.0
|
||||
max_side = np.maximum(h, w)
|
||||
square_bboxes[:, 0] = x1 + w*0.5 - max_side*0.5
|
||||
square_bboxes[:, 1] = y1 + h*0.5 - max_side*0.5
|
||||
square_bboxes[:, 2] = square_bboxes[:, 0] + max_side - 1.0
|
||||
square_bboxes[:, 3] = square_bboxes[:, 1] + max_side - 1.0
|
||||
return square_bboxes
|
||||
|
||||
|
||||
def calibrate_box(bboxes, offsets):
|
||||
"""Transform bounding boxes to be more like true bounding boxes.
|
||||
'offsets' is one of the outputs of the nets.
|
||||
"""
|
||||
x1, y1, x2, y2 = [bboxes[:, i] for i in range(4)]
|
||||
w = x2 - x1 + 1.0
|
||||
h = y2 - y1 + 1.0
|
||||
w = np.expand_dims(w, 1)
|
||||
h = np.expand_dims(h, 1)
|
||||
|
||||
translation = np.hstack([w, h, w, h])*offsets
|
||||
bboxes[:, 0:4] = bboxes[:, 0:4] + translation
|
||||
return bboxes
|
||||
|
||||
|
||||
def get_image_boxes(bounding_boxes, img, size=24):
|
||||
"""Cut out boxes from the image.
|
||||
"""
|
||||
num_boxes = len(bounding_boxes)
|
||||
(height, width, _) = img.shape
|
||||
|
||||
[dy, edy, dx, edx, y, ey, x, ex, w, h] = correct_bboxes(bounding_boxes, width, height)
|
||||
img_boxes = np.zeros((num_boxes, 3, size, size), 'float32')
|
||||
|
||||
for i in range(num_boxes):
|
||||
img_box = np.zeros((h[i], w[i], 3), 'uint8')
|
||||
|
||||
img_array = np.asarray(img, 'uint8')
|
||||
img_box[dy[i]:(edy[i] + 1), dx[i]:(edx[i] + 1), :] =\
|
||||
img_array[y[i]:(ey[i] + 1), x[i]:(ex[i] + 1), :]
|
||||
|
||||
img_box = cv2.resize(img_box, (size, size))
|
||||
img_box = np.asarray(img_box, 'float32')
|
||||
|
||||
img_boxes[i, :, :, :] = _preprocess(img_box)
|
||||
|
||||
return img_boxes
|
||||
|
||||
|
||||
def correct_bboxes(bboxes, width, height):
|
||||
"""Crop boxes that are too big and get coordinates
|
||||
with respect to cutouts.
|
||||
"""
|
||||
x1, y1, x2, y2 = [bboxes[:, i] for i in range(4)]
|
||||
w, h = x2 - x1 + 1.0, y2 - y1 + 1.0
|
||||
num_boxes = bboxes.shape[0]
|
||||
|
||||
x, y, ex, ey = x1, y1, x2, y2
|
||||
dx, dy = np.zeros((num_boxes,)), np.zeros((num_boxes,))
|
||||
edx, edy = w.copy() - 1.0, h.copy() - 1.0
|
||||
|
||||
ind = np.where(ex > width - 1.0)[0]
|
||||
edx[ind] = w[ind] + width - 2.0 - ex[ind]
|
||||
ex[ind] = width - 1.0
|
||||
|
||||
ind = np.where(ey > height - 1.0)[0]
|
||||
edy[ind] = h[ind] + height - 2.0 - ey[ind]
|
||||
ey[ind] = height - 1.0
|
||||
|
||||
ind = np.where(x < 0.0)[0]
|
||||
dx[ind] = 0.0 - x[ind]
|
||||
x[ind] = 0.0
|
||||
|
||||
ind = np.where(y < 0.0)[0]
|
||||
dy[ind] = 0.0 - y[ind]
|
||||
y[ind] = 0.0
|
||||
return_list = [dy, edy, dx, edx, y, ey, x, ex, w, h]
|
||||
return_list = [i.astype('int32') for i in return_list]
|
||||
|
||||
return return_list
|
||||
|
||||
|
||||
def _preprocess(img):
|
||||
"""Preprocessing step before feeding the network.
|
||||
"""
|
||||
img = img.transpose((2, 0, 1))
|
||||
img = np.expand_dims(img, 0)
|
||||
img = (img - 127.5)*0.0078125
|
||||
return img
|
||||
@@ -0,0 +1,242 @@
|
||||
import math
|
||||
import numpy as np
|
||||
import torch
|
||||
from .model import PNet, RNet, ONet
|
||||
from .box_utils import nms, calibrate_box, get_image_boxes, convert_to_square, _preprocess
|
||||
import torch
|
||||
import cv2
|
||||
|
||||
def detect_faces(image, min_face_size=20.0, thresholds=[0.6, 0.7, 0.8],
|
||||
nms_thresholds=[0.7, 0.7, 0.7], gpu_id=0):
|
||||
device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||
pnet, rnet, onet= PNet(), RNet(), ONet()
|
||||
pnet.to(device)
|
||||
rnet.to(device)
|
||||
onet.to(device)
|
||||
onet.eval()
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
(height, width, _) = image.shape
|
||||
min_length = min(height, width)
|
||||
min_detection_size = 12
|
||||
factor = 0.707 # sqrt(0.5)
|
||||
|
||||
scales = []
|
||||
m = min_detection_size/min_face_size
|
||||
min_length *= m
|
||||
|
||||
factor_count = 0
|
||||
while min_length > min_detection_size:
|
||||
scales.append(m*factor**factor_count)
|
||||
min_length *= factor
|
||||
factor_count += 1
|
||||
|
||||
# STAGE 1
|
||||
bounding_boxes = []
|
||||
for s in scales: # run P-Net on different scales
|
||||
boxes = run_first_stage(image, pnet, scale=s, threshold=thresholds[0], gpu_id=gpu_id)
|
||||
bounding_boxes.append(boxes)
|
||||
bounding_boxes = [i for i in bounding_boxes if i is not None]
|
||||
bounding_boxes = np.vstack(bounding_boxes)
|
||||
|
||||
keep = nms(bounding_boxes[:, 0:5], nms_thresholds[0])
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes = calibrate_box(bounding_boxes[:, 0:5], bounding_boxes[:, 5:])
|
||||
bounding_boxes = convert_to_square(bounding_boxes)
|
||||
bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
|
||||
|
||||
# STAGE 2
|
||||
img_boxes = get_image_boxes(bounding_boxes, image, size=24)
|
||||
img_boxes = torch.from_numpy(img_boxes)
|
||||
img_boxes = img_boxes.to(device)
|
||||
output = rnet(img_boxes)
|
||||
offsets = output[0].to('cpu').data.numpy() # shape [n_boxes, 4]
|
||||
probs = output[1].to('cpu').data.numpy() # shape [n_boxes, 2]
|
||||
|
||||
keep = np.where(probs[:, 1] > thresholds[1])[0]
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
|
||||
offsets = offsets[keep]
|
||||
|
||||
keep = nms(bounding_boxes, nms_thresholds[1])
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes = calibrate_box(bounding_boxes, offsets[keep])
|
||||
bounding_boxes = convert_to_square(bounding_boxes)
|
||||
bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
|
||||
|
||||
# STAGE 3
|
||||
img_boxes = get_image_boxes(bounding_boxes, image, size=48)
|
||||
if len(img_boxes) == 0:
|
||||
return [], []
|
||||
img_boxes = torch.from_numpy(img_boxes)
|
||||
img_boxes = img_boxes.to(device)
|
||||
output = onet(img_boxes)
|
||||
landmarks = output[0].to('cpu').data.numpy() # shape [n_boxes, 10]
|
||||
offsets = output[1].to('cpu').data.numpy() # shape [n_boxes, 4]
|
||||
probs = output[2].to('cpu').data.numpy() # shape [n_boxes, 2]
|
||||
|
||||
keep = np.where(probs[:, 1] > thresholds[2])[0]
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
|
||||
offsets = offsets[keep]
|
||||
landmarks = landmarks[keep]
|
||||
|
||||
# compute landmark points
|
||||
width = bounding_boxes[:, 2] - bounding_boxes[:, 0] + 1.0
|
||||
height = bounding_boxes[:, 3] - bounding_boxes[:, 1] + 1.0
|
||||
xmin, ymin = bounding_boxes[:, 0], bounding_boxes[:, 1]
|
||||
landmarks[:, 0:5] = np.expand_dims(xmin, 1) + np.expand_dims(width, 1)*landmarks[:, 0:5]
|
||||
landmarks[:, 5:10] = np.expand_dims(ymin, 1) + np.expand_dims(height, 1)*landmarks[:, 5:10]
|
||||
|
||||
bounding_boxes = calibrate_box(bounding_boxes, offsets)
|
||||
keep = nms(bounding_boxes, nms_thresholds[2], mode='min')
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
landmarks = landmarks[keep]
|
||||
|
||||
return bounding_boxes, landmarks
|
||||
|
||||
class MTCNNFaceDetector(object):
|
||||
def __init__(self, gpu_id=None):
|
||||
self.gpu_id = gpu_id
|
||||
self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||
self.pnet, self.rnet, self.onet = PNet(), RNet(), ONet()
|
||||
self.pnet.to(self.device)
|
||||
self.rnet.to(self.device)
|
||||
self.onet.to(self.device)
|
||||
self.onet.eval()
|
||||
|
||||
def forward(self, image, min_face_size=20.0, thresholds=[0.6, 0.7, 0.8], nms_thresholds=[0.7, 0.7, 0.7]):
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
(height, width, _) = image.shape
|
||||
min_length = min(height, width)
|
||||
min_detection_size = 12
|
||||
factor = 0.707 # sqrt(0.5)
|
||||
|
||||
scales = []
|
||||
m = min_detection_size / min_face_size
|
||||
min_length *= m
|
||||
|
||||
factor_count = 0
|
||||
while min_length > min_detection_size:
|
||||
scales.append(m * factor ** factor_count)
|
||||
min_length *= factor
|
||||
factor_count += 1
|
||||
|
||||
# STAGE 1
|
||||
bounding_boxes = []
|
||||
for s in scales: # run P-Net on different scales
|
||||
boxes = run_first_stage(image, self.pnet, scale=s, threshold=thresholds[0], gpu_id=self.gpu_id)
|
||||
bounding_boxes.append(boxes)
|
||||
bounding_boxes = [i for i in bounding_boxes if i is not None]
|
||||
if len(bounding_boxes) == 0:
|
||||
return [], []
|
||||
bounding_boxes = np.vstack(bounding_boxes)
|
||||
|
||||
keep = nms(bounding_boxes[:, 0:5], nms_thresholds[0])
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes = calibrate_box(bounding_boxes[:, 0:5], bounding_boxes[:, 5:])
|
||||
bounding_boxes = convert_to_square(bounding_boxes)
|
||||
bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
|
||||
|
||||
# STAGE 2
|
||||
img_boxes = get_image_boxes(bounding_boxes, image, size=24)
|
||||
img_boxes = torch.from_numpy(img_boxes)
|
||||
img_boxes = img_boxes.to(self.device)
|
||||
output = self.rnet(img_boxes)
|
||||
offsets = output[0].to('cpu').data.numpy() # shape [n_boxes, 4]
|
||||
probs = output[1].to('cpu').data.numpy() # shape [n_boxes, 2]
|
||||
|
||||
keep = np.where(probs[:, 1] > thresholds[1])[0]
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
|
||||
offsets = offsets[keep]
|
||||
|
||||
keep = nms(bounding_boxes, nms_thresholds[1])
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes = calibrate_box(bounding_boxes, offsets[keep])
|
||||
bounding_boxes = convert_to_square(bounding_boxes)
|
||||
bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
|
||||
|
||||
# STAGE 3
|
||||
img_boxes = get_image_boxes(bounding_boxes, image, size=48)
|
||||
if len(img_boxes) == 0:
|
||||
return [], []
|
||||
img_boxes = torch.from_numpy(img_boxes)
|
||||
img_boxes = img_boxes.to(self.device)
|
||||
output = self.onet(img_boxes)
|
||||
landmarks = output[0].to('cpu').data.numpy() # shape [n_boxes, 10]
|
||||
offsets = output[1].to('cpu').data.numpy() # shape [n_boxes, 4]
|
||||
probs = output[2].to('cpu').data.numpy() # shape [n_boxes, 2]
|
||||
|
||||
keep = np.where(probs[:, 1] > thresholds[2])[0]
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
|
||||
offsets = offsets[keep]
|
||||
landmarks = landmarks[keep]
|
||||
|
||||
# compute landmark points
|
||||
width = bounding_boxes[:, 2] - bounding_boxes[:, 0] + 1.0
|
||||
height = bounding_boxes[:, 3] - bounding_boxes[:, 1] + 1.0
|
||||
xmin, ymin = bounding_boxes[:, 0], bounding_boxes[:, 1]
|
||||
landmarks[:, 0:5] = np.expand_dims(xmin, 1) + np.expand_dims(width, 1) * landmarks[:, 0:5]
|
||||
landmarks[:, 5:10] = np.expand_dims(ymin, 1) + np.expand_dims(height, 1) * landmarks[:, 5:10]
|
||||
|
||||
bounding_boxes = calibrate_box(bounding_boxes, offsets)
|
||||
keep = nms(bounding_boxes, nms_thresholds[2], mode='min')
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
landmarks = landmarks[keep]
|
||||
|
||||
return bounding_boxes, landmarks
|
||||
|
||||
def run_first_stage(image, net, scale, threshold, gpu_id=0):
|
||||
"""
|
||||
Run P-Net, generate bounding boxes, and do NMS.
|
||||
"""
|
||||
device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||
(height, width, _) = image.shape
|
||||
sw, sh = math.ceil(width*scale), math.ceil(height*scale)
|
||||
img = cv2.resize(image, (sw, sh))
|
||||
# img = image.resize((sw, sh), Image.BILINEAR)
|
||||
img = np.asarray(img, 'float32')
|
||||
img = torch.from_numpy(_preprocess(img))
|
||||
img = img.to(device)
|
||||
|
||||
output = net(img)
|
||||
probs = output[1].to('cpu').data.numpy()[0, 1, :, :]
|
||||
offsets = output[0].to('cpu').data.numpy()
|
||||
|
||||
boxes = _generate_bboxes(probs, offsets, scale, threshold)
|
||||
if len(boxes) == 0:
|
||||
return None
|
||||
|
||||
keep = nms(boxes[:, 0:5], overlap_threshold=0.5)
|
||||
return boxes[keep]
|
||||
|
||||
|
||||
def _generate_bboxes(probs, offsets, scale, threshold):
|
||||
"""
|
||||
Generate bounding boxes at places where there is probably a face.
|
||||
"""
|
||||
stride = 2
|
||||
cell_size = 12
|
||||
|
||||
inds = np.where(probs > threshold)
|
||||
|
||||
if inds[0].size == 0:
|
||||
return np.array([])
|
||||
|
||||
tx1, ty1, tx2, ty2 = [offsets[0, i, inds[0], inds[1]] for i in range(4)]
|
||||
|
||||
offsets = np.array([tx1, ty1, tx2, ty2])
|
||||
score = probs[inds[0], inds[1]]
|
||||
|
||||
# P-Net is applied to scaled images, so we need to rescale bounding boxes back
|
||||
bounding_boxes = np.vstack([
|
||||
np.round((stride*inds[1] + 1.0)/scale),
|
||||
np.round((stride*inds[0] + 1.0)/scale),
|
||||
np.round((stride*inds[1] + 1.0 + cell_size)/scale),
|
||||
np.round((stride*inds[0] + 1.0 + cell_size)/scale),
|
||||
score, offsets
|
||||
])
|
||||
|
||||
return bounding_boxes.T
|
||||
@@ -0,0 +1,242 @@
|
||||
import math
|
||||
import numpy as np
|
||||
import torch
|
||||
from .model import PNet, RNet, ONet
|
||||
from .box_utils import nms, calibrate_box, get_image_boxes, convert_to_square, _preprocess
|
||||
import torch
|
||||
import cv2
|
||||
|
||||
def detect_faces(image, min_face_size=20.0, thresholds=[0.6, 0.7, 0.8],
|
||||
nms_thresholds=[0.7, 0.7, 0.7], gpu_id=0):
|
||||
device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||
pnet, rnet, onet= PNet(), RNet(), ONet()
|
||||
pnet.to(device)
|
||||
rnet.to(device)
|
||||
onet.to(device)
|
||||
onet.eval()
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
(height, width, _) = image.shape
|
||||
min_length = min(height, width)
|
||||
min_detection_size = 12
|
||||
factor = 0.707 # sqrt(0.5)
|
||||
|
||||
scales = []
|
||||
m = min_detection_size/min_face_size
|
||||
min_length *= m
|
||||
|
||||
factor_count = 0
|
||||
while min_length > min_detection_size:
|
||||
scales.append(m*factor**factor_count)
|
||||
min_length *= factor
|
||||
factor_count += 1
|
||||
|
||||
# STAGE 1
|
||||
bounding_boxes = []
|
||||
for s in scales: # run P-Net on different scales
|
||||
boxes = run_first_stage(image, pnet, scale=s, threshold=thresholds[0], gpu_id=gpu_id)
|
||||
bounding_boxes.append(boxes)
|
||||
bounding_boxes = [i for i in bounding_boxes if i is not None]
|
||||
bounding_boxes = np.vstack(bounding_boxes)
|
||||
|
||||
keep = nms(bounding_boxes[:, 0:5], nms_thresholds[0])
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes = calibrate_box(bounding_boxes[:, 0:5], bounding_boxes[:, 5:])
|
||||
bounding_boxes = convert_to_square(bounding_boxes)
|
||||
bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
|
||||
|
||||
# STAGE 2
|
||||
img_boxes = get_image_boxes(bounding_boxes, image, size=24)
|
||||
img_boxes = torch.from_numpy(img_boxes)
|
||||
img_boxes = img_boxes.to(device)
|
||||
output = rnet(img_boxes)
|
||||
offsets = output[0].to('cpu').data.numpy() # shape [n_boxes, 4]
|
||||
probs = output[1].to('cpu').data.numpy() # shape [n_boxes, 2]
|
||||
|
||||
keep = np.where(probs[:, 1] > thresholds[1])[0]
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
|
||||
offsets = offsets[keep]
|
||||
|
||||
keep = nms(bounding_boxes, nms_thresholds[1])
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes = calibrate_box(bounding_boxes, offsets[keep])
|
||||
bounding_boxes = convert_to_square(bounding_boxes)
|
||||
bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
|
||||
|
||||
# STAGE 3
|
||||
img_boxes = get_image_boxes(bounding_boxes, image, size=48)
|
||||
if len(img_boxes) == 0:
|
||||
return [], []
|
||||
img_boxes = torch.from_numpy(img_boxes)
|
||||
img_boxes = img_boxes.to(device)
|
||||
output = onet(img_boxes)
|
||||
landmarks = output[0].to('cpu').data.numpy() # shape [n_boxes, 10]
|
||||
offsets = output[1].to('cpu').data.numpy() # shape [n_boxes, 4]
|
||||
probs = output[2].to('cpu').data.numpy() # shape [n_boxes, 2]
|
||||
|
||||
keep = np.where(probs[:, 1] > thresholds[2])[0]
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
|
||||
offsets = offsets[keep]
|
||||
landmarks = landmarks[keep]
|
||||
|
||||
# compute landmark points
|
||||
width = bounding_boxes[:, 2] - bounding_boxes[:, 0] + 1.0
|
||||
height = bounding_boxes[:, 3] - bounding_boxes[:, 1] + 1.0
|
||||
xmin, ymin = bounding_boxes[:, 0], bounding_boxes[:, 1]
|
||||
landmarks[:, 0:5] = np.expand_dims(xmin, 1) + np.expand_dims(width, 1)*landmarks[:, 0:5]
|
||||
landmarks[:, 5:10] = np.expand_dims(ymin, 1) + np.expand_dims(height, 1)*landmarks[:, 5:10]
|
||||
|
||||
bounding_boxes = calibrate_box(bounding_boxes, offsets)
|
||||
keep = nms(bounding_boxes, nms_thresholds[2], mode='min')
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
landmarks = landmarks[keep]
|
||||
|
||||
return bounding_boxes, landmarks
|
||||
|
||||
class MTCNNFaceDetector(object):
|
||||
def __init__(self, gpu_id=None):
|
||||
self.gpu_id = gpu_id
|
||||
self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||
self.pnet, self.rnet, self.onet = PNet(), RNet(), ONet()
|
||||
self.pnet.to(self.device)
|
||||
self.rnet.to(self.device)
|
||||
self.onet.to(self.device)
|
||||
self.onet.eval()
|
||||
|
||||
def forward(self, image, min_face_size=20.0, thresholds=[0.6, 0.7, 0.8], nms_thresholds=[0.7, 0.7, 0.7]):
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
(height, width, _) = image.shape
|
||||
min_length = min(height, width)
|
||||
min_detection_size = 12
|
||||
factor = 0.707 # sqrt(0.5)
|
||||
|
||||
scales = []
|
||||
m = min_detection_size / min_face_size
|
||||
min_length *= m
|
||||
|
||||
factor_count = 0
|
||||
while min_length > min_detection_size:
|
||||
scales.append(m * factor ** factor_count)
|
||||
min_length *= factor
|
||||
factor_count += 1
|
||||
|
||||
# STAGE 1
|
||||
bounding_boxes = []
|
||||
for s in scales: # run P-Net on different scales
|
||||
boxes = run_first_stage(image, self.pnet, scale=s, threshold=thresholds[0], gpu_id=self.gpu_id)
|
||||
bounding_boxes.append(boxes)
|
||||
bounding_boxes = [i for i in bounding_boxes if i is not None]
|
||||
if len(bounding_boxes) == 0:
|
||||
return [], []
|
||||
bounding_boxes = np.vstack(bounding_boxes)
|
||||
|
||||
keep = nms(bounding_boxes[:, 0:5], nms_thresholds[0])
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes = calibrate_box(bounding_boxes[:, 0:5], bounding_boxes[:, 5:])
|
||||
bounding_boxes = convert_to_square(bounding_boxes)
|
||||
bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
|
||||
|
||||
# STAGE 2
|
||||
img_boxes = get_image_boxes(bounding_boxes, image, size=24)
|
||||
img_boxes = torch.from_numpy(img_boxes)
|
||||
img_boxes = img_boxes.to(self.device)
|
||||
output = self.rnet(img_boxes)
|
||||
offsets = output[0].to('cpu').data.numpy() # shape [n_boxes, 4]
|
||||
probs = output[1].to('cpu').data.numpy() # shape [n_boxes, 2]
|
||||
|
||||
keep = np.where(probs[:, 1] > thresholds[1])[0]
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
|
||||
offsets = offsets[keep]
|
||||
|
||||
keep = nms(bounding_boxes, nms_thresholds[1])
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes = calibrate_box(bounding_boxes, offsets[keep])
|
||||
bounding_boxes = convert_to_square(bounding_boxes)
|
||||
bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
|
||||
|
||||
# STAGE 3
|
||||
img_boxes = get_image_boxes(bounding_boxes, image, size=48)
|
||||
if len(img_boxes) == 0:
|
||||
return [], []
|
||||
img_boxes = torch.from_numpy(img_boxes)
|
||||
img_boxes = img_boxes.to(self.device)
|
||||
output = self.onet(img_boxes)
|
||||
landmarks = output[0].to('cpu').data.numpy() # shape [n_boxes, 10]
|
||||
offsets = output[1].to('cpu').data.numpy() # shape [n_boxes, 4]
|
||||
probs = output[2].to('cpu').data.numpy() # shape [n_boxes, 2]
|
||||
|
||||
keep = np.where(probs[:, 1] > thresholds[2])[0]
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
|
||||
offsets = offsets[keep]
|
||||
landmarks = landmarks[keep]
|
||||
|
||||
# compute landmark points
|
||||
width = bounding_boxes[:, 2] - bounding_boxes[:, 0] + 1.0
|
||||
height = bounding_boxes[:, 3] - bounding_boxes[:, 1] + 1.0
|
||||
xmin, ymin = bounding_boxes[:, 0], bounding_boxes[:, 1]
|
||||
landmarks[:, 0:5] = np.expand_dims(xmin, 1) + np.expand_dims(width, 1) * landmarks[:, 0:5]
|
||||
landmarks[:, 5:10] = np.expand_dims(ymin, 1) + np.expand_dims(height, 1) * landmarks[:, 5:10]
|
||||
|
||||
bounding_boxes = calibrate_box(bounding_boxes, offsets)
|
||||
keep = nms(bounding_boxes, nms_thresholds[2], mode='min')
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
landmarks = landmarks[keep]
|
||||
|
||||
return bounding_boxes, landmarks
|
||||
|
||||
def run_first_stage(image, net, scale, threshold, gpu_id=0):
|
||||
"""
|
||||
Run P-Net, generate bounding boxes, and do NMS.
|
||||
"""
|
||||
device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||
(height, width, _) = image.shape
|
||||
sw, sh = math.ceil(width*scale), math.ceil(height*scale)
|
||||
img = cv2.resize(image, (sw, sh))
|
||||
# img = image.resize((sw, sh), Image.BILINEAR)
|
||||
img = np.asarray(img, 'float32')
|
||||
img = torch.from_numpy(_preprocess(img))
|
||||
img = img.to(device)
|
||||
|
||||
output = net(img)
|
||||
probs = output[1].to('cpu').data.numpy()[0, 1, :, :]
|
||||
offsets = output[0].to('cpu').data.numpy()
|
||||
|
||||
boxes = _generate_bboxes(probs, offsets, scale, threshold)
|
||||
if len(boxes) == 0:
|
||||
return None
|
||||
|
||||
keep = nms(boxes[:, 0:5], overlap_threshold=0.5)
|
||||
return boxes[keep]
|
||||
|
||||
|
||||
def _generate_bboxes(probs, offsets, scale, threshold):
|
||||
"""
|
||||
Generate bounding boxes at places where there is probably a face.
|
||||
"""
|
||||
stride = 2
|
||||
cell_size = 12
|
||||
|
||||
inds = np.where(probs > threshold)
|
||||
|
||||
if inds[0].size == 0:
|
||||
return np.array([])
|
||||
|
||||
tx1, ty1, tx2, ty2 = [offsets[0, i, inds[0], inds[1]] for i in range(4)]
|
||||
|
||||
offsets = np.array([tx1, ty1, tx2, ty2])
|
||||
score = probs[inds[0], inds[1]]
|
||||
|
||||
# P-Net is applied to scaled images, so we need to rescale bounding boxes back
|
||||
bounding_boxes = np.vstack([
|
||||
np.round((stride*inds[1] + 1.0)/scale),
|
||||
np.round((stride*inds[0] + 1.0)/scale),
|
||||
np.round((stride*inds[1] + 1.0 + cell_size)/scale),
|
||||
np.round((stride*inds[0] + 1.0 + cell_size)/scale),
|
||||
score, offsets
|
||||
])
|
||||
|
||||
return bounding_boxes.T
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user