commit 0eb61f3e604ea02519bb256be2e5b844c8e0241d Author: colomi <1421901449@qq.com> Date: Sat Jul 11 18:11:49 2026 +0800 初始化换发型项目：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 排除， 由网盘单独上传。 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..9139c91 --- /dev/null +++ b/.gitignore @@ -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 管理 diff --git a/.trae/documents/迁移收尾计划.md b/.trae/documents/迁移收尾计划.md new file mode 100644 index 0000000..637ba1e --- /dev/null +++ b/.trae/documents/迁移收尾计划.md @@ -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 训练，代码可能需要后续修改 diff --git a/.trae/documents/项目迁移到3090部署计划.md b/.trae/documents/项目迁移到3090部署计划.md new file mode 100644 index 0000000..59806b1 --- /dev/null +++ b/.trae/documents/项目迁移到3090部署计划.md @@ -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+** | | diff --git a/README.md b/README.md new file mode 100644 index 0000000..f982ce6 --- /dev/null +++ b/README.md @@ -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 ~/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`。 diff --git a/conda_envs/py310.yml b/conda_envs/py310.yml new file mode 100644 index 0000000..5277234 --- /dev/null +++ b/conda_envs/py310.yml @@ -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 diff --git a/hair_service_sd/.gitignore b/hair_service_sd/.gitignore new file mode 100644 index 0000000..7b004e5 --- /dev/null +++ b/hair_service_sd/.gitignore @@ -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 \ No newline at end of file diff --git a/hair_service_sd/README.md b/hair_service_sd/README.md new file mode 100644 index 0000000..dafd1d3 --- /dev/null +++ b/hair_service_sd/README.md @@ -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 \ No newline at end of file diff --git a/hair_service_sd/bodyseg/backbone/backbone.py b/hair_service_sd/bodyseg/backbone/backbone.py new file mode 100644 index 0000000..151c9f8 --- /dev/null +++ b/hair_service_sd/bodyseg/backbone/backbone.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 diff --git a/hair_service_sd/bodyseg/backbone/fpn.py b/hair_service_sd/bodyseg/backbone/fpn.py new file mode 100644 index 0000000..32bf205 --- /dev/null +++ b/hair_service_sd/bodyseg/backbone/fpn.py @@ -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", 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", 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 + diff --git a/hair_service_sd/bodyseg/backbone/resnet.py b/hair_service_sd/bodyseg/backbone/resnet.py new file mode 100644 index 0000000..d82c31d --- /dev/null +++ b/hair_service_sd/bodyseg/backbone/resnet.py @@ -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) diff --git a/hair_service_sd/bodyseg/backbone/xception.py b/hair_service_sd/bodyseg/backbone/xception.py new file mode 100644 index 0000000..d69f407 --- /dev/null +++ b/hair_service_sd/bodyseg/backbone/xception.py @@ -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()) diff --git a/hair_service_sd/bodyseg/msc_distilling.py b/hair_service_sd/bodyseg/msc_distilling.py new file mode 100644 index 0000000..face96b --- /dev/null +++ b/hair_service_sd/bodyseg/msc_distilling.py @@ -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 \ No newline at end of file diff --git a/hair_service_sd/change_color.py b/hair_service_sd/change_color.py new file mode 100644 index 0000000..30055ba --- /dev/null +++ b/hair_service_sd/change_color.py @@ -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) diff --git a/hair_service_sd/color_test.py b/hair_service_sd/color_test.py new file mode 100644 index 0000000..28dbc27 --- /dev/null +++ b/hair_service_sd/color_test.py @@ -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() \ No newline at end of file diff --git a/hair_service_sd/common/__init__.py b/hair_service_sd/common/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hair_service_sd/common/callback.py b/hair_service_sd/common/callback.py new file mode 100644 index 0000000..79a73b3 --- /dev/null +++ b/hair_service_sd/common/callback.py @@ -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 diff --git a/hair_service_sd/common/gpu_process.py b/hair_service_sd/common/gpu_process.py new file mode 100644 index 0000000..0443228 --- /dev/null +++ b/hair_service_sd/common/gpu_process.py @@ -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 diff --git a/hair_service_sd/common/logger.py b/hair_service_sd/common/logger.py new file mode 100644 index 0000000..dd60358 --- /dev/null +++ b/hair_service_sd/common/logger.py @@ -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") + + + + + diff --git a/hair_service_sd/config/ErrorCode.py b/hair_service_sd/config/ErrorCode.py new file mode 100644 index 0000000..e69de29 diff --git a/hair_service_sd/config/__init__.py b/hair_service_sd/config/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hair_service_sd/config/configure.ini.template b/hair_service_sd/config/configure.ini.template new file mode 100644 index 0000000..919571e --- /dev/null +++ b/hair_service_sd/config/configure.ini.template @@ -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 diff --git a/hair_service_sd/copy_files_by_id.py b/hair_service_sd/copy_files_by_id.py new file mode 100644 index 0000000..dc8350b --- /dev/null +++ b/hair_service_sd/copy_files_by_id.py @@ -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}") \ No newline at end of file diff --git a/hair_service_sd/core/MMCVFaceRecognitionServer.py b/hair_service_sd/core/MMCVFaceRecognitionServer.py new file mode 100644 index 0000000..1aa48b6 --- /dev/null +++ b/hair_service_sd/core/MMCVFaceRecognitionServer.py @@ -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() diff --git a/hair_service_sd/core/PingFangMedium.ttf b/hair_service_sd/core/PingFangMedium.ttf new file mode 100644 index 0000000..982661b Binary files /dev/null and b/hair_service_sd/core/PingFangMedium.ttf differ diff --git a/hair_service_sd/core/__init__.py b/hair_service_sd/core/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hair_service_sd/core/bodyseg/backbone/backbone.py b/hair_service_sd/core/bodyseg/backbone/backbone.py new file mode 100644 index 0000000..151c9f8 --- /dev/null +++ b/hair_service_sd/core/bodyseg/backbone/backbone.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 diff --git a/hair_service_sd/core/bodyseg/backbone/fpn.py b/hair_service_sd/core/bodyseg/backbone/fpn.py new file mode 100644 index 0000000..5adcd95 --- /dev/null +++ b/hair_service_sd/core/bodyseg/backbone/fpn.py @@ -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", 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", 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 + diff --git a/hair_service_sd/core/bodyseg/backbone/resnet.py b/hair_service_sd/core/bodyseg/backbone/resnet.py new file mode 100644 index 0000000..4318c80 --- /dev/null +++ b/hair_service_sd/core/bodyseg/backbone/resnet.py @@ -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) diff --git a/hair_service_sd/core/bodyseg/backbone/xception.py b/hair_service_sd/core/bodyseg/backbone/xception.py new file mode 100644 index 0000000..ae53726 --- /dev/null +++ b/hair_service_sd/core/bodyseg/backbone/xception.py @@ -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()) diff --git a/hair_service_sd/core/bodyseg/msc_distilling.py b/hair_service_sd/core/bodyseg/msc_distilling.py new file mode 100644 index 0000000..a06fbee --- /dev/null +++ b/hair_service_sd/core/bodyseg/msc_distilling.py @@ -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 \ No newline at end of file diff --git a/hair_service_sd/core/cos_module.py b/hair_service_sd/core/cos_module.py new file mode 100644 index 0000000..1711ae8 --- /dev/null +++ b/hair_service_sd/core/cos_module.py @@ -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) \ No newline at end of file diff --git a/hair_service_sd/core/face3d/__init__.py b/hair_service_sd/core/face3d/__init__.py new file mode 100644 index 0000000..6db1992 --- /dev/null +++ b/hair_service_sd/core/face3d/__init__.py @@ -0,0 +1,2 @@ +from . import mesh +from . import morphable_model diff --git a/hair_service_sd/core/face3d/mesh/__init__.py b/hair_service_sd/core/face3d/mesh/__init__.py new file mode 100644 index 0000000..af38c13 --- /dev/null +++ b/hair_service_sd/core/face3d/mesh/__init__.py @@ -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 + diff --git a/hair_service_sd/core/face3d/mesh/cython/.gitignore b/hair_service_sd/core/face3d/mesh/cython/.gitignore new file mode 100644 index 0000000..17dcbad --- /dev/null +++ b/hair_service_sd/core/face3d/mesh/cython/.gitignore @@ -0,0 +1 @@ +*.cpython-36m-x86_64-linux-gnu.so diff --git a/hair_service_sd/core/face3d/mesh/cython/__init__.py b/hair_service_sd/core/face3d/mesh/cython/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hair_service_sd/core/face3d/mesh/cython/mesh_core.cpp b/hair_service_sd/core/face3d/mesh/cython/mesh_core.cpp new file mode 100644 index 0000000..221f55a --- /dev/null +++ b/hair_service_sd/core/face3d/mesh/cython/mesh_core.cpp @@ -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; + } + +} \ No newline at end of file diff --git a/hair_service_sd/core/face3d/mesh/cython/mesh_core.h b/hair_service_sd/core/face3d/mesh/cython/mesh_core.h new file mode 100644 index 0000000..3eb0029 --- /dev/null +++ b/hair_service_sd/core/face3d/mesh/cython/mesh_core.h @@ -0,0 +1,83 @@ +#ifndef MESH_CORE_HPP_ +#define MESH_CORE_HPP_ + +#include +#include +#include +#include +#include +#include + +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 \ No newline at end of file diff --git a/hair_service_sd/core/face3d/mesh/cython/mesh_core_cython.cpp b/hair_service_sd/core/face3d/mesh/cython/mesh_core_cython.cpp new file mode 100644 index 0000000..5d52baf --- /dev/null +++ b/hair_service_sd/core/face3d/mesh/cython/mesh_core_cython.cpp @@ -0,0 +1,2618 @@ +/* Generated by Cython 0.29.26 */ + +#ifndef PY_SSIZE_T_CLEAN +#define PY_SSIZE_T_CLEAN +#endif /* PY_SSIZE_T_CLEAN */ +#include "Python.h" +#ifndef Py_PYTHON_H + #error Python headers needed to compile C extensions, please install development version of Python. +#elif PY_VERSION_HEX < 0x02060000 || (0x03000000 <= PY_VERSION_HEX && PY_VERSION_HEX < 0x03030000) + #error Cython requires Python 2.6+ or Python 3.3+. +#else +#define CYTHON_ABI "0_29_26" +#define CYTHON_HEX_VERSION 0x001D1AF0 +#define CYTHON_FUTURE_DIVISION 0 +#include +#ifndef offsetof + #define offsetof(type, member) ( (size_t) & ((type*)0) -> member ) +#endif +#if !defined(WIN32) && !defined(MS_WINDOWS) + #ifndef __stdcall + #define __stdcall + #endif + #ifndef __cdecl + #define __cdecl + #endif + #ifndef __fastcall + #define __fastcall + #endif +#endif +#ifndef DL_IMPORT + #define DL_IMPORT(t) t +#endif +#ifndef DL_EXPORT + #define DL_EXPORT(t) t +#endif +#define __PYX_COMMA , +#ifndef HAVE_LONG_LONG + #if PY_VERSION_HEX >= 0x02070000 + #define HAVE_LONG_LONG + #endif +#endif +#ifndef PY_LONG_LONG + #define PY_LONG_LONG LONG_LONG +#endif +#ifndef Py_HUGE_VAL + #define Py_HUGE_VAL HUGE_VAL +#endif +#ifdef PYPY_VERSION + #define CYTHON_COMPILING_IN_PYPY 1 + #define CYTHON_COMPILING_IN_PYSTON 0 + #define CYTHON_COMPILING_IN_CPYTHON 0 + #undef CYTHON_USE_TYPE_SLOTS + #define CYTHON_USE_TYPE_SLOTS 0 + #undef CYTHON_USE_PYTYPE_LOOKUP + #define CYTHON_USE_PYTYPE_LOOKUP 0 + #if PY_VERSION_HEX < 0x03050000 + #undef CYTHON_USE_ASYNC_SLOTS + #define CYTHON_USE_ASYNC_SLOTS 0 + #elif !defined(CYTHON_USE_ASYNC_SLOTS) + #define CYTHON_USE_ASYNC_SLOTS 1 + #endif + #undef CYTHON_USE_PYLIST_INTERNALS + #define CYTHON_USE_PYLIST_INTERNALS 0 + #undef CYTHON_USE_UNICODE_INTERNALS + #define CYTHON_USE_UNICODE_INTERNALS 0 + #undef CYTHON_USE_UNICODE_WRITER + #define CYTHON_USE_UNICODE_WRITER 0 + #undef CYTHON_USE_PYLONG_INTERNALS + #define CYTHON_USE_PYLONG_INTERNALS 0 + #undef CYTHON_AVOID_BORROWED_REFS + #define CYTHON_AVOID_BORROWED_REFS 1 + #undef CYTHON_ASSUME_SAFE_MACROS + #define CYTHON_ASSUME_SAFE_MACROS 0 + #undef CYTHON_UNPACK_METHODS + #define CYTHON_UNPACK_METHODS 0 + #undef CYTHON_FAST_THREAD_STATE + #define CYTHON_FAST_THREAD_STATE 0 + #undef CYTHON_FAST_PYCALL + #define CYTHON_FAST_PYCALL 0 + #undef CYTHON_PEP489_MULTI_PHASE_INIT + #define CYTHON_PEP489_MULTI_PHASE_INIT 0 + #undef CYTHON_USE_TP_FINALIZE + #define CYTHON_USE_TP_FINALIZE 0 + #undef CYTHON_USE_DICT_VERSIONS + #define CYTHON_USE_DICT_VERSIONS 0 + #undef CYTHON_USE_EXC_INFO_STACK + #define CYTHON_USE_EXC_INFO_STACK 0 +#elif defined(PYSTON_VERSION) + #define CYTHON_COMPILING_IN_PYPY 0 + #define CYTHON_COMPILING_IN_PYSTON 1 + #define CYTHON_COMPILING_IN_CPYTHON 0 + #ifndef CYTHON_USE_TYPE_SLOTS + #define CYTHON_USE_TYPE_SLOTS 1 + #endif + #undef CYTHON_USE_PYTYPE_LOOKUP + #define CYTHON_USE_PYTYPE_LOOKUP 0 + #undef CYTHON_USE_ASYNC_SLOTS + #define CYTHON_USE_ASYNC_SLOTS 0 + #undef CYTHON_USE_PYLIST_INTERNALS + #define CYTHON_USE_PYLIST_INTERNALS 0 + #ifndef CYTHON_USE_UNICODE_INTERNALS + #define CYTHON_USE_UNICODE_INTERNALS 1 + #endif + #undef CYTHON_USE_UNICODE_WRITER + #define CYTHON_USE_UNICODE_WRITER 0 + #undef CYTHON_USE_PYLONG_INTERNALS + #define CYTHON_USE_PYLONG_INTERNALS 0 + #ifndef CYTHON_AVOID_BORROWED_REFS + #define CYTHON_AVOID_BORROWED_REFS 0 + #endif + #ifndef CYTHON_ASSUME_SAFE_MACROS + #define CYTHON_ASSUME_SAFE_MACROS 1 + #endif + #ifndef CYTHON_UNPACK_METHODS + #define CYTHON_UNPACK_METHODS 1 + #endif + #undef CYTHON_FAST_THREAD_STATE + #define CYTHON_FAST_THREAD_STATE 0 + #undef CYTHON_FAST_PYCALL + #define CYTHON_FAST_PYCALL 0 + #undef CYTHON_PEP489_MULTI_PHASE_INIT + #define CYTHON_PEP489_MULTI_PHASE_INIT 0 + #undef CYTHON_USE_TP_FINALIZE + #define CYTHON_USE_TP_FINALIZE 0 + #undef CYTHON_USE_DICT_VERSIONS + #define CYTHON_USE_DICT_VERSIONS 0 + #undef CYTHON_USE_EXC_INFO_STACK + #define CYTHON_USE_EXC_INFO_STACK 0 +#else + #define CYTHON_COMPILING_IN_PYPY 0 + #define CYTHON_COMPILING_IN_PYSTON 0 + #define CYTHON_COMPILING_IN_CPYTHON 1 + #ifndef CYTHON_USE_TYPE_SLOTS + #define CYTHON_USE_TYPE_SLOTS 1 + #endif + #if PY_VERSION_HEX < 0x02070000 + #undef CYTHON_USE_PYTYPE_LOOKUP + #define CYTHON_USE_PYTYPE_LOOKUP 0 + #elif !defined(CYTHON_USE_PYTYPE_LOOKUP) + #define CYTHON_USE_PYTYPE_LOOKUP 1 + #endif + #if PY_MAJOR_VERSION < 3 + #undef CYTHON_USE_ASYNC_SLOTS + #define CYTHON_USE_ASYNC_SLOTS 0 + #elif !defined(CYTHON_USE_ASYNC_SLOTS) + #define CYTHON_USE_ASYNC_SLOTS 1 + #endif + #if PY_VERSION_HEX < 0x02070000 + #undef CYTHON_USE_PYLONG_INTERNALS + #define CYTHON_USE_PYLONG_INTERNALS 0 + #elif !defined(CYTHON_USE_PYLONG_INTERNALS) + #define CYTHON_USE_PYLONG_INTERNALS 1 + #endif + #ifndef CYTHON_USE_PYLIST_INTERNALS + #define CYTHON_USE_PYLIST_INTERNALS 1 + #endif + #ifndef CYTHON_USE_UNICODE_INTERNALS + #define CYTHON_USE_UNICODE_INTERNALS 1 + #endif + #if PY_VERSION_HEX < 0x030300F0 || PY_VERSION_HEX >= 0x030B00A2 + #undef CYTHON_USE_UNICODE_WRITER + #define CYTHON_USE_UNICODE_WRITER 0 + #elif !defined(CYTHON_USE_UNICODE_WRITER) + #define CYTHON_USE_UNICODE_WRITER 1 + #endif + #ifndef CYTHON_AVOID_BORROWED_REFS + #define CYTHON_AVOID_BORROWED_REFS 0 + #endif + #ifndef CYTHON_ASSUME_SAFE_MACROS + #define CYTHON_ASSUME_SAFE_MACROS 1 + #endif + #ifndef CYTHON_UNPACK_METHODS + #define CYTHON_UNPACK_METHODS 1 + #endif + #ifndef CYTHON_FAST_THREAD_STATE + #define CYTHON_FAST_THREAD_STATE 1 + #endif + #ifndef CYTHON_FAST_PYCALL + #define CYTHON_FAST_PYCALL (PY_VERSION_HEX < 0x030B00A1) + #endif + #ifndef CYTHON_PEP489_MULTI_PHASE_INIT + #define CYTHON_PEP489_MULTI_PHASE_INIT (PY_VERSION_HEX >= 0x03050000) + #endif + #ifndef CYTHON_USE_TP_FINALIZE + #define CYTHON_USE_TP_FINALIZE (PY_VERSION_HEX >= 0x030400a1) + #endif + #ifndef CYTHON_USE_DICT_VERSIONS + #define CYTHON_USE_DICT_VERSIONS (PY_VERSION_HEX >= 0x030600B1) + #endif + #ifndef CYTHON_USE_EXC_INFO_STACK + #define CYTHON_USE_EXC_INFO_STACK (PY_VERSION_HEX >= 0x030700A3) + #endif +#endif +#if !defined(CYTHON_FAST_PYCCALL) +#define CYTHON_FAST_PYCCALL (CYTHON_FAST_PYCALL && PY_VERSION_HEX >= 0x030600B1) +#endif +#if CYTHON_USE_PYLONG_INTERNALS + #if PY_MAJOR_VERSION < 3 + #include "longintrepr.h" + #endif + #undef SHIFT + #undef BASE + #undef MASK + #ifdef SIZEOF_VOID_P + enum { __pyx_check_sizeof_voidp = 1 / (int)(SIZEOF_VOID_P == sizeof(void*)) }; + #endif +#endif +#ifndef __has_attribute + #define __has_attribute(x) 0 +#endif +#ifndef __has_cpp_attribute + #define __has_cpp_attribute(x) 0 +#endif +#ifndef CYTHON_RESTRICT + #if defined(__GNUC__) + #define CYTHON_RESTRICT __restrict__ + #elif defined(_MSC_VER) && _MSC_VER >= 1400 + #define CYTHON_RESTRICT __restrict + #elif defined (__STDC_VERSION__) && __STDC_VERSION__ >= 199901L + #define CYTHON_RESTRICT restrict + #else + #define CYTHON_RESTRICT + #endif +#endif +#ifndef CYTHON_UNUSED +# if defined(__GNUC__) +# if !(defined(__cplusplus)) || (__GNUC__ > 3 || (__GNUC__ == 3 && __GNUC_MINOR__ >= 4)) +# define CYTHON_UNUSED __attribute__ ((__unused__)) +# else +# define CYTHON_UNUSED +# endif +# elif defined(__ICC) || (defined(__INTEL_COMPILER) && !defined(_MSC_VER)) +# define CYTHON_UNUSED __attribute__ ((__unused__)) +# else +# define CYTHON_UNUSED +# endif +#endif +#ifndef CYTHON_MAYBE_UNUSED_VAR +# if defined(__cplusplus) + template void CYTHON_MAYBE_UNUSED_VAR( const T& ) { } +# else +# define CYTHON_MAYBE_UNUSED_VAR(x) (void)(x) +# endif +#endif +#ifndef CYTHON_NCP_UNUSED +# if CYTHON_COMPILING_IN_CPYTHON +# define CYTHON_NCP_UNUSED +# else +# define CYTHON_NCP_UNUSED CYTHON_UNUSED +# endif +#endif +#define __Pyx_void_to_None(void_result) ((void)(void_result), Py_INCREF(Py_None), Py_None) +#ifdef _MSC_VER + #ifndef _MSC_STDINT_H_ + #if _MSC_VER < 1300 + typedef unsigned char uint8_t; + typedef unsigned int uint32_t; + #else + typedef unsigned __int8 uint8_t; + typedef unsigned __int32 uint32_t; + #endif + #endif +#else + #include +#endif +#ifndef CYTHON_FALLTHROUGH + #if defined(__cplusplus) && __cplusplus >= 201103L + #if __has_cpp_attribute(fallthrough) + #define CYTHON_FALLTHROUGH [[fallthrough]] + #elif __has_cpp_attribute(clang::fallthrough) + #define CYTHON_FALLTHROUGH [[clang::fallthrough]] + #elif __has_cpp_attribute(gnu::fallthrough) + #define CYTHON_FALLTHROUGH [[gnu::fallthrough]] + #endif + #endif + #ifndef CYTHON_FALLTHROUGH + #if __has_attribute(fallthrough) + #define CYTHON_FALLTHROUGH __attribute__((fallthrough)) + #else + #define CYTHON_FALLTHROUGH + #endif + #endif + #if defined(__clang__ ) && defined(__apple_build_version__) + #if __apple_build_version__ < 7000000 + #undef CYTHON_FALLTHROUGH + #define CYTHON_FALLTHROUGH + #endif + #endif +#endif + +#ifndef __cplusplus + #error "Cython files generated with the C++ option must be compiled with a C++ compiler." +#endif +#ifndef CYTHON_INLINE + #if defined(__clang__) + #define CYTHON_INLINE __inline__ __attribute__ ((__unused__)) + #else + #define CYTHON_INLINE inline + #endif +#endif +template +void __Pyx_call_destructor(T& x) { + x.~T(); +} +template +class __Pyx_FakeReference { + public: + __Pyx_FakeReference() : ptr(NULL) { } + __Pyx_FakeReference(const T& ref) : ptr(const_cast(&ref)) { } + T *operator->() { return ptr; } + T *operator&() { return ptr; } + operator T&() { return *ptr; } + template bool operator ==(U other) { return *ptr == other; } + template bool operator !=(U other) { return *ptr != other; } + private: + T *ptr; +}; + +#if CYTHON_COMPILING_IN_PYPY && PY_VERSION_HEX < 0x02070600 && !defined(Py_OptimizeFlag) + #define Py_OptimizeFlag 0 +#endif +#define __PYX_BUILD_PY_SSIZE_T "n" +#define CYTHON_FORMAT_SSIZE_T "z" +#if PY_MAJOR_VERSION < 3 + #define __Pyx_BUILTIN_MODULE_NAME "__builtin__" + #define __Pyx_PyCode_New(a, k, l, s, f, code, c, n, v, fv, cell, fn, name, fline, lnos)\ + PyCode_New(a+k, l, s, f, code, c, n, v, fv, cell, fn, name, fline, lnos) + #define __Pyx_DefaultClassType PyClass_Type +#else + #define __Pyx_BUILTIN_MODULE_NAME "builtins" + #define __Pyx_DefaultClassType PyType_Type +#if PY_VERSION_HEX >= 0x030B00A1 + static CYTHON_INLINE PyCodeObject* __Pyx_PyCode_New(int a, int k, int l, int s, int f, + PyObject *code, PyObject *c, PyObject* n, PyObject *v, + PyObject *fv, PyObject *cell, PyObject* fn, + PyObject *name, int fline, PyObject *lnos) { + PyObject *kwds=NULL, *argcount=NULL, *posonlyargcount=NULL, *kwonlyargcount=NULL; + PyObject *nlocals=NULL, *stacksize=NULL, *flags=NULL, *replace=NULL, *call_result=NULL, *empty=NULL; + const char *fn_cstr=NULL; + const char *name_cstr=NULL; + PyCodeObject* co=NULL; + PyObject *type, *value, *traceback; + PyErr_Fetch(&type, &value, &traceback); + if (!(kwds=PyDict_New())) goto end; + if (!(argcount=PyLong_FromLong(a))) goto end; + if (PyDict_SetItemString(kwds, "co_argcount", argcount) != 0) goto end; + if (!(posonlyargcount=PyLong_FromLong(0))) goto end; + if (PyDict_SetItemString(kwds, "co_posonlyargcount", posonlyargcount) != 0) goto end; + if (!(kwonlyargcount=PyLong_FromLong(k))) goto end; + if (PyDict_SetItemString(kwds, "co_kwonlyargcount", kwonlyargcount) != 0) goto end; + if (!(nlocals=PyLong_FromLong(l))) goto end; + if (PyDict_SetItemString(kwds, "co_nlocals", nlocals) != 0) goto end; + if (!(stacksize=PyLong_FromLong(s))) goto end; + if (PyDict_SetItemString(kwds, "co_stacksize", stacksize) != 0) goto end; + if (!(flags=PyLong_FromLong(f))) goto end; + if (PyDict_SetItemString(kwds, "co_flags", flags) != 0) goto end; + if (PyDict_SetItemString(kwds, "co_code", code) != 0) goto end; + if (PyDict_SetItemString(kwds, "co_consts", c) != 0) goto end; + if (PyDict_SetItemString(kwds, "co_names", n) != 0) goto end; + if (PyDict_SetItemString(kwds, "co_varnames", v) != 0) goto end; + if (PyDict_SetItemString(kwds, "co_freevars", fv) != 0) goto end; + if (PyDict_SetItemString(kwds, "co_cellvars", cell) != 0) goto end; + if (PyDict_SetItemString(kwds, "co_linetable", lnos) != 0) goto end; + if (!(fn_cstr=PyUnicode_AsUTF8AndSize(fn, NULL))) goto end; + if (!(name_cstr=PyUnicode_AsUTF8AndSize(name, NULL))) goto end; + if (!(co = PyCode_NewEmpty(fn_cstr, name_cstr, fline))) goto end; + if (!(replace = PyObject_GetAttrString((PyObject*)co, "replace"))) goto cleanup_code_too; + if (!(empty = PyTuple_New(0))) goto cleanup_code_too; // unfortunately __pyx_empty_tuple isn't available here + if (!(call_result = PyObject_Call(replace, empty, kwds))) goto cleanup_code_too; + Py_XDECREF((PyObject*)co); + co = (PyCodeObject*)call_result; + call_result = NULL; + if (0) { + cleanup_code_too: + Py_XDECREF((PyObject*)co); + co = NULL; + } + end: + Py_XDECREF(kwds); + Py_XDECREF(argcount); + Py_XDECREF(posonlyargcount); + Py_XDECREF(kwonlyargcount); + Py_XDECREF(nlocals); + Py_XDECREF(stacksize); + Py_XDECREF(replace); + Py_XDECREF(call_result); + Py_XDECREF(empty); + if (type) { + PyErr_Restore(type, value, traceback); + } + return co; + } +#else + #define __Pyx_PyCode_New(a, k, l, s, f, code, c, n, v, fv, cell, fn, name, fline, lnos)\ + PyCode_New(a, k, l, s, f, code, c, n, v, fv, cell, fn, name, fline, lnos) +#endif + #define __Pyx_DefaultClassType PyType_Type +#endif +#ifndef Py_TPFLAGS_CHECKTYPES + #define Py_TPFLAGS_CHECKTYPES 0 +#endif +#ifndef Py_TPFLAGS_HAVE_INDEX + #define Py_TPFLAGS_HAVE_INDEX 0 +#endif +#ifndef Py_TPFLAGS_HAVE_NEWBUFFER + #define Py_TPFLAGS_HAVE_NEWBUFFER 0 +#endif +#ifndef Py_TPFLAGS_HAVE_FINALIZE + #define Py_TPFLAGS_HAVE_FINALIZE 0 +#endif +#ifndef METH_STACKLESS + #define METH_STACKLESS 0 +#endif +#if PY_VERSION_HEX <= 0x030700A3 || !defined(METH_FASTCALL) + #ifndef METH_FASTCALL + #define METH_FASTCALL 0x80 + #endif + typedef PyObject *(*__Pyx_PyCFunctionFast) (PyObject *self, PyObject *const *args, Py_ssize_t nargs); + typedef PyObject *(*__Pyx_PyCFunctionFastWithKeywords) (PyObject *self, PyObject *const *args, + Py_ssize_t nargs, PyObject *kwnames); +#else + #define __Pyx_PyCFunctionFast _PyCFunctionFast + #define __Pyx_PyCFunctionFastWithKeywords _PyCFunctionFastWithKeywords +#endif +#if CYTHON_FAST_PYCCALL +#define __Pyx_PyFastCFunction_Check(func)\ + ((PyCFunction_Check(func) && (METH_FASTCALL == (PyCFunction_GET_FLAGS(func) & ~(METH_CLASS | METH_STATIC | METH_COEXIST | METH_KEYWORDS | METH_STACKLESS))))) +#else +#define __Pyx_PyFastCFunction_Check(func) 0 +#endif +#if CYTHON_COMPILING_IN_PYPY && !defined(PyObject_Malloc) + #define PyObject_Malloc(s) PyMem_Malloc(s) + #define PyObject_Free(p) PyMem_Free(p) + #define PyObject_Realloc(p) PyMem_Realloc(p) +#endif +#if CYTHON_COMPILING_IN_CPYTHON && PY_VERSION_HEX < 0x030400A1 + #define PyMem_RawMalloc(n) PyMem_Malloc(n) + #define PyMem_RawRealloc(p, n) PyMem_Realloc(p, n) + #define PyMem_RawFree(p) PyMem_Free(p) +#endif +#if CYTHON_COMPILING_IN_PYSTON + #define __Pyx_PyCode_HasFreeVars(co) PyCode_HasFreeVars(co) + #define __Pyx_PyFrame_SetLineNumber(frame, lineno) PyFrame_SetLineNumber(frame, lineno) +#else + #define __Pyx_PyCode_HasFreeVars(co) (PyCode_GetNumFree(co) > 0) + #define __Pyx_PyFrame_SetLineNumber(frame, lineno) (frame)->f_lineno = (lineno) +#endif +#if !CYTHON_FAST_THREAD_STATE || PY_VERSION_HEX < 0x02070000 + #define __Pyx_PyThreadState_Current PyThreadState_GET() +#elif PY_VERSION_HEX >= 0x03060000 + #define __Pyx_PyThreadState_Current _PyThreadState_UncheckedGet() +#elif PY_VERSION_HEX >= 0x03000000 + #define __Pyx_PyThreadState_Current PyThreadState_GET() +#else + #define __Pyx_PyThreadState_Current _PyThreadState_Current +#endif +#if PY_VERSION_HEX < 0x030700A2 && !defined(PyThread_tss_create) && !defined(Py_tss_NEEDS_INIT) +#include "pythread.h" +#define Py_tss_NEEDS_INIT 0 +typedef int Py_tss_t; +static CYTHON_INLINE int PyThread_tss_create(Py_tss_t *key) { + *key = PyThread_create_key(); + return 0; +} +static CYTHON_INLINE Py_tss_t * PyThread_tss_alloc(void) { + Py_tss_t *key = (Py_tss_t *)PyObject_Malloc(sizeof(Py_tss_t)); + *key = Py_tss_NEEDS_INIT; + return key; +} +static CYTHON_INLINE void PyThread_tss_free(Py_tss_t *key) { + PyObject_Free(key); +} +static CYTHON_INLINE int PyThread_tss_is_created(Py_tss_t *key) { + return *key != Py_tss_NEEDS_INIT; +} +static CYTHON_INLINE void PyThread_tss_delete(Py_tss_t *key) { + PyThread_delete_key(*key); + *key = Py_tss_NEEDS_INIT; +} +static CYTHON_INLINE int PyThread_tss_set(Py_tss_t *key, void *value) { + return PyThread_set_key_value(*key, value); +} +static CYTHON_INLINE void * PyThread_tss_get(Py_tss_t *key) { + return PyThread_get_key_value(*key); +} +#endif +#if CYTHON_COMPILING_IN_CPYTHON || defined(_PyDict_NewPresized) +#define __Pyx_PyDict_NewPresized(n) ((n <= 8) ? PyDict_New() : _PyDict_NewPresized(n)) +#else +#define __Pyx_PyDict_NewPresized(n) PyDict_New() +#endif +#if PY_MAJOR_VERSION >= 3 || CYTHON_FUTURE_DIVISION + #define __Pyx_PyNumber_Divide(x,y) PyNumber_TrueDivide(x,y) + #define __Pyx_PyNumber_InPlaceDivide(x,y) PyNumber_InPlaceTrueDivide(x,y) +#else + #define __Pyx_PyNumber_Divide(x,y) PyNumber_Divide(x,y) + #define __Pyx_PyNumber_InPlaceDivide(x,y) PyNumber_InPlaceDivide(x,y) +#endif +#if CYTHON_COMPILING_IN_CPYTHON && PY_VERSION_HEX >= 0x030500A1 && CYTHON_USE_UNICODE_INTERNALS +#define __Pyx_PyDict_GetItemStr(dict, name) _PyDict_GetItem_KnownHash(dict, name, ((PyASCIIObject *) name)->hash) +#else +#define __Pyx_PyDict_GetItemStr(dict, name) PyDict_GetItem(dict, name) +#endif +#if PY_VERSION_HEX > 0x03030000 && defined(PyUnicode_KIND) + #define CYTHON_PEP393_ENABLED 1 + #if defined(PyUnicode_IS_READY) + #define __Pyx_PyUnicode_READY(op) (likely(PyUnicode_IS_READY(op)) ?\ + 0 : _PyUnicode_Ready((PyObject *)(op))) + #else + #define __Pyx_PyUnicode_READY(op) (0) + #endif + #define __Pyx_PyUnicode_GET_LENGTH(u) PyUnicode_GET_LENGTH(u) + #define __Pyx_PyUnicode_READ_CHAR(u, i) PyUnicode_READ_CHAR(u, i) + #define __Pyx_PyUnicode_MAX_CHAR_VALUE(u) PyUnicode_MAX_CHAR_VALUE(u) + #define __Pyx_PyUnicode_KIND(u) PyUnicode_KIND(u) + #define __Pyx_PyUnicode_DATA(u) PyUnicode_DATA(u) + #define __Pyx_PyUnicode_READ(k, d, i) PyUnicode_READ(k, d, i) + #define __Pyx_PyUnicode_WRITE(k, d, i, ch) PyUnicode_WRITE(k, d, i, ch) + #if defined(PyUnicode_IS_READY) && defined(PyUnicode_GET_SIZE) + #if CYTHON_COMPILING_IN_CPYTHON && PY_VERSION_HEX >= 0x03090000 + #define __Pyx_PyUnicode_IS_TRUE(u) (0 != (likely(PyUnicode_IS_READY(u)) ? PyUnicode_GET_LENGTH(u) : ((PyCompactUnicodeObject *)(u))->wstr_length)) + #else + #define __Pyx_PyUnicode_IS_TRUE(u) (0 != (likely(PyUnicode_IS_READY(u)) ? PyUnicode_GET_LENGTH(u) : PyUnicode_GET_SIZE(u))) + #endif + #else + #define __Pyx_PyUnicode_IS_TRUE(u) (0 != PyUnicode_GET_LENGTH(u)) + #endif +#else + #define CYTHON_PEP393_ENABLED 0 + #define PyUnicode_1BYTE_KIND 1 + #define PyUnicode_2BYTE_KIND 2 + #define PyUnicode_4BYTE_KIND 4 + #define __Pyx_PyUnicode_READY(op) (0) + #define __Pyx_PyUnicode_GET_LENGTH(u) PyUnicode_GET_SIZE(u) + #define __Pyx_PyUnicode_READ_CHAR(u, i) ((Py_UCS4)(PyUnicode_AS_UNICODE(u)[i])) + #define __Pyx_PyUnicode_MAX_CHAR_VALUE(u) ((sizeof(Py_UNICODE) == 2) ? 65535 : 1114111) + #define __Pyx_PyUnicode_KIND(u) (sizeof(Py_UNICODE)) + #define __Pyx_PyUnicode_DATA(u) ((void*)PyUnicode_AS_UNICODE(u)) + #define __Pyx_PyUnicode_READ(k, d, i) ((void)(k), (Py_UCS4)(((Py_UNICODE*)d)[i])) + #define __Pyx_PyUnicode_WRITE(k, d, i, ch) (((void)(k)), ((Py_UNICODE*)d)[i] = ch) + #define __Pyx_PyUnicode_IS_TRUE(u) (0 != PyUnicode_GET_SIZE(u)) +#endif +#if CYTHON_COMPILING_IN_PYPY + #define __Pyx_PyUnicode_Concat(a, b) PyNumber_Add(a, b) + #define __Pyx_PyUnicode_ConcatSafe(a, b) PyNumber_Add(a, b) +#else + #define __Pyx_PyUnicode_Concat(a, b) PyUnicode_Concat(a, b) + #define __Pyx_PyUnicode_ConcatSafe(a, b) ((unlikely((a) == Py_None) || unlikely((b) == Py_None)) ?\ + PyNumber_Add(a, b) : __Pyx_PyUnicode_Concat(a, b)) +#endif +#if CYTHON_COMPILING_IN_PYPY && !defined(PyUnicode_Contains) + #define PyUnicode_Contains(u, s) PySequence_Contains(u, s) +#endif +#if CYTHON_COMPILING_IN_PYPY && !defined(PyByteArray_Check) + #define PyByteArray_Check(obj) PyObject_TypeCheck(obj, &PyByteArray_Type) +#endif +#if CYTHON_COMPILING_IN_PYPY && !defined(PyObject_Format) + #define PyObject_Format(obj, fmt) PyObject_CallMethod(obj, "__format__", "O", fmt) +#endif +#define __Pyx_PyString_FormatSafe(a, b) ((unlikely((a) == Py_None || (PyString_Check(b) && !PyString_CheckExact(b)))) ? PyNumber_Remainder(a, b) : __Pyx_PyString_Format(a, b)) +#define __Pyx_PyUnicode_FormatSafe(a, b) ((unlikely((a) == Py_None || (PyUnicode_Check(b) && !PyUnicode_CheckExact(b)))) ? PyNumber_Remainder(a, b) : PyUnicode_Format(a, b)) +#if PY_MAJOR_VERSION >= 3 + #define __Pyx_PyString_Format(a, b) PyUnicode_Format(a, b) +#else + #define __Pyx_PyString_Format(a, b) PyString_Format(a, b) +#endif +#if PY_MAJOR_VERSION < 3 && !defined(PyObject_ASCII) + #define PyObject_ASCII(o) PyObject_Repr(o) +#endif +#if PY_MAJOR_VERSION >= 3 + #define PyBaseString_Type PyUnicode_Type + #define PyStringObject PyUnicodeObject + #define PyString_Type PyUnicode_Type + #define PyString_Check PyUnicode_Check + #define PyString_CheckExact PyUnicode_CheckExact +#ifndef PyObject_Unicode + #define PyObject_Unicode PyObject_Str +#endif +#endif +#if PY_MAJOR_VERSION >= 3 + #define __Pyx_PyBaseString_Check(obj) PyUnicode_Check(obj) + #define __Pyx_PyBaseString_CheckExact(obj) PyUnicode_CheckExact(obj) +#else + #define __Pyx_PyBaseString_Check(obj) (PyString_Check(obj) || PyUnicode_Check(obj)) + #define __Pyx_PyBaseString_CheckExact(obj) (PyString_CheckExact(obj) || PyUnicode_CheckExact(obj)) +#endif +#ifndef PySet_CheckExact + #define PySet_CheckExact(obj) (Py_TYPE(obj) == &PySet_Type) +#endif +#if PY_VERSION_HEX >= 0x030900A4 + #define __Pyx_SET_REFCNT(obj, refcnt) Py_SET_REFCNT(obj, refcnt) + #define __Pyx_SET_SIZE(obj, size) Py_SET_SIZE(obj, size) +#else + #define __Pyx_SET_REFCNT(obj, refcnt) Py_REFCNT(obj) = (refcnt) + #define __Pyx_SET_SIZE(obj, size) Py_SIZE(obj) = (size) +#endif +#if CYTHON_ASSUME_SAFE_MACROS + #define __Pyx_PySequence_SIZE(seq) Py_SIZE(seq) +#else + #define __Pyx_PySequence_SIZE(seq) PySequence_Size(seq) +#endif +#if PY_MAJOR_VERSION >= 3 + #define PyIntObject PyLongObject + #define PyInt_Type PyLong_Type + #define PyInt_Check(op) PyLong_Check(op) + #define PyInt_CheckExact(op) PyLong_CheckExact(op) + #define PyInt_FromString PyLong_FromString + #define PyInt_FromUnicode PyLong_FromUnicode + #define PyInt_FromLong PyLong_FromLong + #define PyInt_FromSize_t PyLong_FromSize_t + #define PyInt_FromSsize_t PyLong_FromSsize_t + #define PyInt_AsLong PyLong_AsLong + #define PyInt_AS_LONG PyLong_AS_LONG + #define PyInt_AsSsize_t PyLong_AsSsize_t + #define PyInt_AsUnsignedLongMask PyLong_AsUnsignedLongMask + #define PyInt_AsUnsignedLongLongMask PyLong_AsUnsignedLongLongMask + #define PyNumber_Int PyNumber_Long +#endif +#if PY_MAJOR_VERSION >= 3 + #define PyBoolObject PyLongObject +#endif +#if PY_MAJOR_VERSION >= 3 && CYTHON_COMPILING_IN_PYPY + #ifndef PyUnicode_InternFromString + #define PyUnicode_InternFromString(s) PyUnicode_FromString(s) + #endif +#endif +#if PY_VERSION_HEX < 0x030200A4 + typedef long Py_hash_t; + #define __Pyx_PyInt_FromHash_t PyInt_FromLong + #define __Pyx_PyInt_AsHash_t __Pyx_PyIndex_AsHash_t +#else + #define __Pyx_PyInt_FromHash_t PyInt_FromSsize_t + #define __Pyx_PyInt_AsHash_t __Pyx_PyIndex_AsSsize_t +#endif +#if PY_MAJOR_VERSION >= 3 + #define __Pyx_PyMethod_New(func, self, klass) ((self) ? ((void)(klass), PyMethod_New(func, self)) : __Pyx_NewRef(func)) +#else + #define __Pyx_PyMethod_New(func, self, klass) PyMethod_New(func, self, klass) +#endif +#if CYTHON_USE_ASYNC_SLOTS + #if PY_VERSION_HEX >= 0x030500B1 + #define __Pyx_PyAsyncMethodsStruct PyAsyncMethods + #define __Pyx_PyType_AsAsync(obj) (Py_TYPE(obj)->tp_as_async) + #else + #define __Pyx_PyType_AsAsync(obj) ((__Pyx_PyAsyncMethodsStruct*) (Py_TYPE(obj)->tp_reserved)) + #endif +#else + #define __Pyx_PyType_AsAsync(obj) NULL +#endif +#ifndef __Pyx_PyAsyncMethodsStruct + typedef struct { + unaryfunc am_await; + unaryfunc am_aiter; + unaryfunc am_anext; + } __Pyx_PyAsyncMethodsStruct; +#endif + +#if defined(WIN32) || defined(MS_WINDOWS) + #define _USE_MATH_DEFINES +#endif +#include +#ifdef NAN +#define __PYX_NAN() ((float) NAN) +#else +static CYTHON_INLINE float __PYX_NAN() { + float value; + memset(&value, 0xFF, sizeof(value)); + return value; +} +#endif +#if defined(__CYGWIN__) && defined(_LDBL_EQ_DBL) +#define __Pyx_truncl trunc +#else +#define __Pyx_truncl truncl +#endif + +#define __PYX_MARK_ERR_POS(f_index, lineno) \ + { __pyx_filename = __pyx_f[f_index]; (void)__pyx_filename; __pyx_lineno = lineno; (void)__pyx_lineno; __pyx_clineno = __LINE__; (void)__pyx_clineno; } +#define __PYX_ERR(f_index, lineno, Ln_error) \ + { __PYX_MARK_ERR_POS(f_index, lineno) goto Ln_error; } + +#ifndef __PYX_EXTERN_C + #ifdef __cplusplus + #define __PYX_EXTERN_C extern "C" + #else + #define __PYX_EXTERN_C extern + #endif +#endif + +#define __PYX_HAVE__mesh_core_cython +#define __PYX_HAVE_API__mesh_core_cython +/* Early includes */ +#ifdef _OPENMP +#include +#endif /* _OPENMP */ + +#if defined(PYREX_WITHOUT_ASSERTIONS) && !defined(CYTHON_WITHOUT_ASSERTIONS) +#define CYTHON_WITHOUT_ASSERTIONS +#endif + +typedef struct {PyObject **p; const char *s; const Py_ssize_t n; const char* encoding; + const char is_unicode; const char is_str; const char intern; } __Pyx_StringTabEntry; + +#define __PYX_DEFAULT_STRING_ENCODING_IS_ASCII 0 +#define __PYX_DEFAULT_STRING_ENCODING_IS_UTF8 0 +#define __PYX_DEFAULT_STRING_ENCODING_IS_DEFAULT (PY_MAJOR_VERSION >= 3 && __PYX_DEFAULT_STRING_ENCODING_IS_UTF8) +#define __PYX_DEFAULT_STRING_ENCODING "" +#define __Pyx_PyObject_FromString __Pyx_PyBytes_FromString +#define __Pyx_PyObject_FromStringAndSize __Pyx_PyBytes_FromStringAndSize +#define __Pyx_uchar_cast(c) ((unsigned char)c) +#define __Pyx_long_cast(x) ((long)x) +#define __Pyx_fits_Py_ssize_t(v, type, is_signed) (\ + (sizeof(type) < sizeof(Py_ssize_t)) ||\ + (sizeof(type) > sizeof(Py_ssize_t) &&\ + likely(v < (type)PY_SSIZE_T_MAX ||\ + v == (type)PY_SSIZE_T_MAX) &&\ + (!is_signed || likely(v > (type)PY_SSIZE_T_MIN ||\ + v == (type)PY_SSIZE_T_MIN))) ||\ + (sizeof(type) == sizeof(Py_ssize_t) &&\ + (is_signed || likely(v < (type)PY_SSIZE_T_MAX ||\ + v == (type)PY_SSIZE_T_MAX))) ) +static CYTHON_INLINE int __Pyx_is_valid_index(Py_ssize_t i, Py_ssize_t limit) { + return (size_t) i < (size_t) limit; +} +#if defined (__cplusplus) && __cplusplus >= 201103L + #include + #define __Pyx_sst_abs(value) std::abs(value) +#elif SIZEOF_INT >= SIZEOF_SIZE_T + #define __Pyx_sst_abs(value) abs(value) +#elif SIZEOF_LONG >= SIZEOF_SIZE_T + #define __Pyx_sst_abs(value) labs(value) +#elif defined (_MSC_VER) + #define __Pyx_sst_abs(value) ((Py_ssize_t)_abs64(value)) +#elif defined (__STDC_VERSION__) && __STDC_VERSION__ >= 199901L + #define __Pyx_sst_abs(value) llabs(value) +#elif defined (__GNUC__) + #define __Pyx_sst_abs(value) __builtin_llabs(value) +#else + #define __Pyx_sst_abs(value) ((value<0) ? -value : value) +#endif +static CYTHON_INLINE const char* __Pyx_PyObject_AsString(PyObject*); +static CYTHON_INLINE const char* __Pyx_PyObject_AsStringAndSize(PyObject*, Py_ssize_t* length); +#define __Pyx_PyByteArray_FromString(s) PyByteArray_FromStringAndSize((const char*)s, strlen((const char*)s)) +#define __Pyx_PyByteArray_FromStringAndSize(s, l) PyByteArray_FromStringAndSize((const char*)s, l) +#define __Pyx_PyBytes_FromString PyBytes_FromString +#define __Pyx_PyBytes_FromStringAndSize PyBytes_FromStringAndSize +static CYTHON_INLINE PyObject* __Pyx_PyUnicode_FromString(const char*); +#if PY_MAJOR_VERSION < 3 + #define __Pyx_PyStr_FromString __Pyx_PyBytes_FromString + #define __Pyx_PyStr_FromStringAndSize __Pyx_PyBytes_FromStringAndSize +#else + #define __Pyx_PyStr_FromString __Pyx_PyUnicode_FromString + #define __Pyx_PyStr_FromStringAndSize __Pyx_PyUnicode_FromStringAndSize +#endif +#define __Pyx_PyBytes_AsWritableString(s) ((char*) PyBytes_AS_STRING(s)) +#define __Pyx_PyBytes_AsWritableSString(s) ((signed char*) PyBytes_AS_STRING(s)) +#define __Pyx_PyBytes_AsWritableUString(s) ((unsigned char*) PyBytes_AS_STRING(s)) +#define __Pyx_PyBytes_AsString(s) ((const char*) PyBytes_AS_STRING(s)) +#define __Pyx_PyBytes_AsSString(s) ((const signed char*) PyBytes_AS_STRING(s)) +#define __Pyx_PyBytes_AsUString(s) ((const unsigned char*) PyBytes_AS_STRING(s)) +#define __Pyx_PyObject_AsWritableString(s) ((char*) __Pyx_PyObject_AsString(s)) +#define __Pyx_PyObject_AsWritableSString(s) ((signed char*) __Pyx_PyObject_AsString(s)) +#define __Pyx_PyObject_AsWritableUString(s) ((unsigned char*) __Pyx_PyObject_AsString(s)) +#define __Pyx_PyObject_AsSString(s) ((const signed char*) __Pyx_PyObject_AsString(s)) +#define __Pyx_PyObject_AsUString(s) ((const unsigned char*) __Pyx_PyObject_AsString(s)) +#define __Pyx_PyObject_FromCString(s) __Pyx_PyObject_FromString((const char*)s) +#define __Pyx_PyBytes_FromCString(s) __Pyx_PyBytes_FromString((const char*)s) +#define __Pyx_PyByteArray_FromCString(s) __Pyx_PyByteArray_FromString((const char*)s) +#define __Pyx_PyStr_FromCString(s) __Pyx_PyStr_FromString((const char*)s) +#define __Pyx_PyUnicode_FromCString(s) __Pyx_PyUnicode_FromString((const char*)s) +static CYTHON_INLINE size_t __Pyx_Py_UNICODE_strlen(const Py_UNICODE *u) { + const Py_UNICODE *u_end = u; + while (*u_end++) ; + return (size_t)(u_end - u - 1); +} +#define __Pyx_PyUnicode_FromUnicode(u) PyUnicode_FromUnicode(u, __Pyx_Py_UNICODE_strlen(u)) +#define __Pyx_PyUnicode_FromUnicodeAndLength PyUnicode_FromUnicode +#define __Pyx_PyUnicode_AsUnicode PyUnicode_AsUnicode +#define __Pyx_NewRef(obj) (Py_INCREF(obj), obj) +#define __Pyx_Owned_Py_None(b) __Pyx_NewRef(Py_None) +static CYTHON_INLINE PyObject * __Pyx_PyBool_FromLong(long b); +static CYTHON_INLINE int __Pyx_PyObject_IsTrue(PyObject*); +static CYTHON_INLINE int __Pyx_PyObject_IsTrueAndDecref(PyObject*); +static CYTHON_INLINE PyObject* __Pyx_PyNumber_IntOrLong(PyObject* x); +#define __Pyx_PySequence_Tuple(obj)\ + (likely(PyTuple_CheckExact(obj)) ? __Pyx_NewRef(obj) : PySequence_Tuple(obj)) +static CYTHON_INLINE Py_ssize_t __Pyx_PyIndex_AsSsize_t(PyObject*); +static CYTHON_INLINE PyObject * __Pyx_PyInt_FromSize_t(size_t); +static CYTHON_INLINE Py_hash_t __Pyx_PyIndex_AsHash_t(PyObject*); +#if CYTHON_ASSUME_SAFE_MACROS +#define __pyx_PyFloat_AsDouble(x) (PyFloat_CheckExact(x) ? PyFloat_AS_DOUBLE(x) : PyFloat_AsDouble(x)) +#else +#define __pyx_PyFloat_AsDouble(x) PyFloat_AsDouble(x) +#endif +#define __pyx_PyFloat_AsFloat(x) ((float) __pyx_PyFloat_AsDouble(x)) +#if PY_MAJOR_VERSION >= 3 +#define __Pyx_PyNumber_Int(x) (PyLong_CheckExact(x) ? __Pyx_NewRef(x) : PyNumber_Long(x)) +#else +#define __Pyx_PyNumber_Int(x) (PyInt_CheckExact(x) ? __Pyx_NewRef(x) : PyNumber_Int(x)) +#endif +#define __Pyx_PyNumber_Float(x) (PyFloat_CheckExact(x) ? __Pyx_NewRef(x) : PyNumber_Float(x)) +#if PY_MAJOR_VERSION < 3 && __PYX_DEFAULT_STRING_ENCODING_IS_ASCII +static int __Pyx_sys_getdefaultencoding_not_ascii; +static int __Pyx_init_sys_getdefaultencoding_params(void) { + PyObject* sys; + PyObject* default_encoding = NULL; + PyObject* ascii_chars_u = NULL; + PyObject* ascii_chars_b = NULL; + const char* default_encoding_c; + sys = PyImport_ImportModule("sys"); + if (!sys) goto bad; + default_encoding = PyObject_CallMethod(sys, (char*) "getdefaultencoding", NULL); + Py_DECREF(sys); + if (!default_encoding) goto bad; + default_encoding_c = PyBytes_AsString(default_encoding); + if (!default_encoding_c) goto bad; + if (strcmp(default_encoding_c, "ascii") == 0) { + __Pyx_sys_getdefaultencoding_not_ascii = 0; + } else { + char ascii_chars[128]; + int c; + for (c = 0; c < 128; c++) { + ascii_chars[c] = c; + } + __Pyx_sys_getdefaultencoding_not_ascii = 1; + ascii_chars_u = PyUnicode_DecodeASCII(ascii_chars, 128, NULL); + if (!ascii_chars_u) goto bad; + ascii_chars_b = PyUnicode_AsEncodedString(ascii_chars_u, default_encoding_c, NULL); + if (!ascii_chars_b || !PyBytes_Check(ascii_chars_b) || memcmp(ascii_chars, PyBytes_AS_STRING(ascii_chars_b), 128) != 0) { + PyErr_Format( + PyExc_ValueError, + "This module compiled with c_string_encoding=ascii, but default encoding '%.200s' is not a superset of ascii.", + default_encoding_c); + goto bad; + } + Py_DECREF(ascii_chars_u); + Py_DECREF(ascii_chars_b); + } + Py_DECREF(default_encoding); + return 0; +bad: + Py_XDECREF(default_encoding); + Py_XDECREF(ascii_chars_u); + Py_XDECREF(ascii_chars_b); + return -1; +} +#endif +#if __PYX_DEFAULT_STRING_ENCODING_IS_DEFAULT && PY_MAJOR_VERSION >= 3 +#define __Pyx_PyUnicode_FromStringAndSize(c_str, size) PyUnicode_DecodeUTF8(c_str, size, NULL) +#else +#define __Pyx_PyUnicode_FromStringAndSize(c_str, size) PyUnicode_Decode(c_str, size, __PYX_DEFAULT_STRING_ENCODING, NULL) +#if __PYX_DEFAULT_STRING_ENCODING_IS_DEFAULT +static char* __PYX_DEFAULT_STRING_ENCODING; +static int __Pyx_init_sys_getdefaultencoding_params(void) { + PyObject* sys; + PyObject* default_encoding = NULL; + char* default_encoding_c; + sys = PyImport_ImportModule("sys"); + if (!sys) goto bad; + default_encoding = PyObject_CallMethod(sys, (char*) (const char*) "getdefaultencoding", NULL); + Py_DECREF(sys); + if (!default_encoding) goto bad; + default_encoding_c = PyBytes_AsString(default_encoding); + if (!default_encoding_c) goto bad; + __PYX_DEFAULT_STRING_ENCODING = (char*) malloc(strlen(default_encoding_c) + 1); + if (!__PYX_DEFAULT_STRING_ENCODING) goto bad; + strcpy(__PYX_DEFAULT_STRING_ENCODING, default_encoding_c); + Py_DECREF(default_encoding); + return 0; +bad: + Py_XDECREF(default_encoding); + return -1; +} +#endif +#endif + + +/* Test for GCC > 2.95 */ +#if defined(__GNUC__) && (__GNUC__ > 2 || (__GNUC__ == 2 && (__GNUC_MINOR__ > 95))) + #define likely(x) __builtin_expect(!!(x), 1) + #define unlikely(x) __builtin_expect(!!(x), 0) +#else /* !__GNUC__ or GCC < 2.95 */ + #define likely(x) (x) + #define unlikely(x) (x) +#endif /* __GNUC__ */ +static CYTHON_INLINE void __Pyx_pretend_to_initialize(void* ptr) { (void)ptr; } + +static PyObject *__pyx_m = NULL; +static PyObject *__pyx_d; +static PyObject *__pyx_b; +static PyObject *__pyx_cython_runtime = NULL; +static PyObject *__pyx_empty_tuple; +static PyObject *__pyx_empty_bytes; +static PyObject *__pyx_empty_unicode; +static int __pyx_lineno; +static int __pyx_clineno = 0; +static const char * __pyx_cfilenm= __FILE__; +static const char *__pyx_filename; + + +static const char *__pyx_f[] = { + "mesh_core_cython.pyx", +}; + +/*--- Type declarations ---*/ + +/* --- Runtime support code (head) --- */ +/* Refnanny.proto */ +#ifndef CYTHON_REFNANNY + #define CYTHON_REFNANNY 0 +#endif +#if CYTHON_REFNANNY + typedef struct { + void (*INCREF)(void*, PyObject*, int); + void (*DECREF)(void*, PyObject*, int); + void (*GOTREF)(void*, PyObject*, int); + void (*GIVEREF)(void*, PyObject*, int); + void* (*SetupContext)(const char*, int, const char*); + void (*FinishContext)(void**); + } __Pyx_RefNannyAPIStruct; + static __Pyx_RefNannyAPIStruct *__Pyx_RefNanny = NULL; + static __Pyx_RefNannyAPIStruct *__Pyx_RefNannyImportAPI(const char *modname); + #define __Pyx_RefNannyDeclarations void *__pyx_refnanny = NULL; +#ifdef WITH_THREAD + #define __Pyx_RefNannySetupContext(name, acquire_gil)\ + if (acquire_gil) {\ + PyGILState_STATE __pyx_gilstate_save = PyGILState_Ensure();\ + __pyx_refnanny = __Pyx_RefNanny->SetupContext((name), __LINE__, __FILE__);\ + PyGILState_Release(__pyx_gilstate_save);\ + } else {\ + __pyx_refnanny = __Pyx_RefNanny->SetupContext((name), __LINE__, __FILE__);\ + } +#else + #define __Pyx_RefNannySetupContext(name, acquire_gil)\ + __pyx_refnanny = __Pyx_RefNanny->SetupContext((name), __LINE__, __FILE__) +#endif + #define __Pyx_RefNannyFinishContext()\ + __Pyx_RefNanny->FinishContext(&__pyx_refnanny) + #define __Pyx_INCREF(r) __Pyx_RefNanny->INCREF(__pyx_refnanny, (PyObject *)(r), __LINE__) + #define __Pyx_DECREF(r) __Pyx_RefNanny->DECREF(__pyx_refnanny, (PyObject *)(r), __LINE__) + #define __Pyx_GOTREF(r) __Pyx_RefNanny->GOTREF(__pyx_refnanny, (PyObject *)(r), __LINE__) + #define __Pyx_GIVEREF(r) __Pyx_RefNanny->GIVEREF(__pyx_refnanny, (PyObject *)(r), __LINE__) + #define __Pyx_XINCREF(r) do { if((r) != NULL) {__Pyx_INCREF(r); }} while(0) + #define __Pyx_XDECREF(r) do { if((r) != NULL) {__Pyx_DECREF(r); }} while(0) + #define __Pyx_XGOTREF(r) do { if((r) != NULL) {__Pyx_GOTREF(r); }} while(0) + #define __Pyx_XGIVEREF(r) do { if((r) != NULL) {__Pyx_GIVEREF(r);}} while(0) +#else + #define __Pyx_RefNannyDeclarations + #define __Pyx_RefNannySetupContext(name, acquire_gil) + #define __Pyx_RefNannyFinishContext() + #define __Pyx_INCREF(r) Py_INCREF(r) + #define __Pyx_DECREF(r) Py_DECREF(r) + #define __Pyx_GOTREF(r) + #define __Pyx_GIVEREF(r) + #define __Pyx_XINCREF(r) Py_XINCREF(r) + #define __Pyx_XDECREF(r) Py_XDECREF(r) + #define __Pyx_XGOTREF(r) + #define __Pyx_XGIVEREF(r) +#endif +#define __Pyx_XDECREF_SET(r, v) do {\ + PyObject *tmp = (PyObject *) r;\ + r = v; __Pyx_XDECREF(tmp);\ + } while (0) +#define __Pyx_DECREF_SET(r, v) do {\ + PyObject *tmp = (PyObject *) r;\ + r = v; __Pyx_DECREF(tmp);\ + } while (0) +#define __Pyx_CLEAR(r) do { PyObject* tmp = ((PyObject*)(r)); r = NULL; __Pyx_DECREF(tmp);} while(0) +#define __Pyx_XCLEAR(r) do { if((r) != NULL) {PyObject* tmp = ((PyObject*)(r)); r = NULL; __Pyx_DECREF(tmp);}} while(0) + +/* PyDictVersioning.proto */ +#if CYTHON_USE_DICT_VERSIONS && CYTHON_USE_TYPE_SLOTS +#define __PYX_DICT_VERSION_INIT ((PY_UINT64_T) -1) +#define __PYX_GET_DICT_VERSION(dict) (((PyDictObject*)(dict))->ma_version_tag) +#define __PYX_UPDATE_DICT_CACHE(dict, value, cache_var, version_var)\ + (version_var) = __PYX_GET_DICT_VERSION(dict);\ + (cache_var) = (value); +#define __PYX_PY_DICT_LOOKUP_IF_MODIFIED(VAR, DICT, LOOKUP) {\ + static PY_UINT64_T __pyx_dict_version = 0;\ + static PyObject *__pyx_dict_cached_value = NULL;\ + if (likely(__PYX_GET_DICT_VERSION(DICT) == __pyx_dict_version)) {\ + (VAR) = __pyx_dict_cached_value;\ + } else {\ + (VAR) = __pyx_dict_cached_value = (LOOKUP);\ + __pyx_dict_version = __PYX_GET_DICT_VERSION(DICT);\ + }\ +} +static CYTHON_INLINE PY_UINT64_T __Pyx_get_tp_dict_version(PyObject *obj); +static CYTHON_INLINE PY_UINT64_T __Pyx_get_object_dict_version(PyObject *obj); +static CYTHON_INLINE int __Pyx_object_dict_version_matches(PyObject* obj, PY_UINT64_T tp_dict_version, PY_UINT64_T obj_dict_version); +#else +#define __PYX_GET_DICT_VERSION(dict) (0) +#define __PYX_UPDATE_DICT_CACHE(dict, value, cache_var, version_var) +#define __PYX_PY_DICT_LOOKUP_IF_MODIFIED(VAR, DICT, LOOKUP) (VAR) = (LOOKUP); +#endif + +/* PyObjectGetAttrStr.proto */ +#if CYTHON_USE_TYPE_SLOTS +static CYTHON_INLINE PyObject* __Pyx_PyObject_GetAttrStr(PyObject* obj, PyObject* attr_name); +#else +#define __Pyx_PyObject_GetAttrStr(o,n) PyObject_GetAttr(o,n) +#endif + +/* PyThreadStateGet.proto */ +#if CYTHON_FAST_THREAD_STATE +#define __Pyx_PyThreadState_declare PyThreadState *__pyx_tstate; +#define __Pyx_PyThreadState_assign __pyx_tstate = __Pyx_PyThreadState_Current; +#define __Pyx_PyErr_Occurred() __pyx_tstate->curexc_type +#else +#define __Pyx_PyThreadState_declare +#define __Pyx_PyThreadState_assign +#define __Pyx_PyErr_Occurred() PyErr_Occurred() +#endif + +/* PyErrFetchRestore.proto */ +#if CYTHON_FAST_THREAD_STATE +#define __Pyx_PyErr_Clear() __Pyx_ErrRestore(NULL, NULL, NULL) +#define __Pyx_ErrRestoreWithState(type, value, tb) __Pyx_ErrRestoreInState(PyThreadState_GET(), type, value, tb) +#define __Pyx_ErrFetchWithState(type, value, tb) __Pyx_ErrFetchInState(PyThreadState_GET(), type, value, tb) +#define __Pyx_ErrRestore(type, value, tb) __Pyx_ErrRestoreInState(__pyx_tstate, type, value, tb) +#define __Pyx_ErrFetch(type, value, tb) __Pyx_ErrFetchInState(__pyx_tstate, type, value, tb) +static CYTHON_INLINE void __Pyx_ErrRestoreInState(PyThreadState *tstate, PyObject *type, PyObject *value, PyObject *tb); +static CYTHON_INLINE void __Pyx_ErrFetchInState(PyThreadState *tstate, PyObject **type, PyObject **value, PyObject **tb); +#if CYTHON_COMPILING_IN_CPYTHON +#define __Pyx_PyErr_SetNone(exc) (Py_INCREF(exc), __Pyx_ErrRestore((exc), NULL, NULL)) +#else +#define __Pyx_PyErr_SetNone(exc) PyErr_SetNone(exc) +#endif +#else +#define __Pyx_PyErr_Clear() PyErr_Clear() +#define __Pyx_PyErr_SetNone(exc) PyErr_SetNone(exc) +#define __Pyx_ErrRestoreWithState(type, value, tb) PyErr_Restore(type, value, tb) +#define __Pyx_ErrFetchWithState(type, value, tb) PyErr_Fetch(type, value, tb) +#define __Pyx_ErrRestoreInState(tstate, type, value, tb) PyErr_Restore(type, value, tb) +#define __Pyx_ErrFetchInState(tstate, type, value, tb) PyErr_Fetch(type, value, tb) +#define __Pyx_ErrRestore(type, value, tb) PyErr_Restore(type, value, tb) +#define __Pyx_ErrFetch(type, value, tb) PyErr_Fetch(type, value, tb) +#endif + +/* CLineInTraceback.proto */ +#ifdef CYTHON_CLINE_IN_TRACEBACK +#define __Pyx_CLineForTraceback(tstate, c_line) (((CYTHON_CLINE_IN_TRACEBACK)) ? c_line : 0) +#else +static int __Pyx_CLineForTraceback(PyThreadState *tstate, int c_line); +#endif + +/* CodeObjectCache.proto */ +typedef struct { + PyCodeObject* code_object; + int code_line; +} __Pyx_CodeObjectCacheEntry; +struct __Pyx_CodeObjectCache { + int count; + int max_count; + __Pyx_CodeObjectCacheEntry* entries; +}; +static struct __Pyx_CodeObjectCache __pyx_code_cache = {0,0,NULL}; +static int __pyx_bisect_code_objects(__Pyx_CodeObjectCacheEntry* entries, int count, int code_line); +static PyCodeObject *__pyx_find_code_object(int code_line); +static void __pyx_insert_code_object(int code_line, PyCodeObject* code_object); + +/* AddTraceback.proto */ +static void __Pyx_AddTraceback(const char *funcname, int c_line, + int py_line, const char *filename); + +/* GCCDiagnostics.proto */ +#if defined(__GNUC__) && (__GNUC__ > 4 || (__GNUC__ == 4 && __GNUC_MINOR__ >= 6)) +#define __Pyx_HAS_GCC_DIAGNOSTIC +#endif + +/* CIntToPy.proto */ +static CYTHON_INLINE PyObject* __Pyx_PyInt_From_long(long value); + +/* CIntFromPy.proto */ +static CYTHON_INLINE long __Pyx_PyInt_As_long(PyObject *); + +/* CIntFromPy.proto */ +static CYTHON_INLINE int __Pyx_PyInt_As_int(PyObject *); + +/* FastTypeChecks.proto */ +#if CYTHON_COMPILING_IN_CPYTHON +#define __Pyx_TypeCheck(obj, type) __Pyx_IsSubtype(Py_TYPE(obj), (PyTypeObject *)type) +static CYTHON_INLINE int __Pyx_IsSubtype(PyTypeObject *a, PyTypeObject *b); +static CYTHON_INLINE int __Pyx_PyErr_GivenExceptionMatches(PyObject *err, PyObject *type); +static CYTHON_INLINE int __Pyx_PyErr_GivenExceptionMatches2(PyObject *err, PyObject *type1, PyObject *type2); +#else +#define __Pyx_TypeCheck(obj, type) PyObject_TypeCheck(obj, (PyTypeObject *)type) +#define __Pyx_PyErr_GivenExceptionMatches(err, type) PyErr_GivenExceptionMatches(err, type) +#define __Pyx_PyErr_GivenExceptionMatches2(err, type1, type2) (PyErr_GivenExceptionMatches(err, type1) || PyErr_GivenExceptionMatches(err, type2)) +#endif +#define __Pyx_PyException_Check(obj) __Pyx_TypeCheck(obj, PyExc_Exception) + +/* CheckBinaryVersion.proto */ +static int __Pyx_check_binary_version(void); + +/* InitStrings.proto */ +static int __Pyx_InitStrings(__Pyx_StringTabEntry *t); + + +/* Module declarations from 'mesh_core_cython' */ +#define __Pyx_MODULE_NAME "mesh_core_cython" +extern int __pyx_module_is_main_mesh_core_cython; +int __pyx_module_is_main_mesh_core_cython = 0; + +/* Implementation of 'mesh_core_cython' */ +static const char __pyx_k_main[] = "__main__"; +static const char __pyx_k_name[] = "__name__"; +static const char __pyx_k_test[] = "__test__"; +static const char __pyx_k_cline_in_traceback[] = "cline_in_traceback"; +static PyObject *__pyx_n_s_cline_in_traceback; +static PyObject *__pyx_n_s_main; +static PyObject *__pyx_n_s_name; +static PyObject *__pyx_n_s_test; +/* Late includes */ + +static PyMethodDef __pyx_methods[] = { + {0, 0, 0, 0} +}; + +#if PY_MAJOR_VERSION >= 3 +#if CYTHON_PEP489_MULTI_PHASE_INIT +static PyObject* __pyx_pymod_create(PyObject *spec, PyModuleDef *def); /*proto*/ +static int __pyx_pymod_exec_mesh_core_cython(PyObject* module); /*proto*/ +static PyModuleDef_Slot __pyx_moduledef_slots[] = { + {Py_mod_create, (void*)__pyx_pymod_create}, + {Py_mod_exec, (void*)__pyx_pymod_exec_mesh_core_cython}, + {0, NULL} +}; +#endif + +static struct PyModuleDef __pyx_moduledef = { + PyModuleDef_HEAD_INIT, + "mesh_core_cython", + 0, /* m_doc */ + #if CYTHON_PEP489_MULTI_PHASE_INIT + 0, /* m_size */ + #else + -1, /* m_size */ + #endif + __pyx_methods /* m_methods */, + #if CYTHON_PEP489_MULTI_PHASE_INIT + __pyx_moduledef_slots, /* m_slots */ + #else + NULL, /* m_reload */ + #endif + NULL, /* m_traverse */ + NULL, /* m_clear */ + NULL /* m_free */ +}; +#endif +#ifndef CYTHON_SMALL_CODE +#if defined(__clang__) + #define CYTHON_SMALL_CODE +#elif defined(__GNUC__) && (__GNUC__ > 4 || (__GNUC__ == 4 && __GNUC_MINOR__ >= 3)) + #define CYTHON_SMALL_CODE __attribute__((cold)) +#else + #define CYTHON_SMALL_CODE +#endif +#endif + +static __Pyx_StringTabEntry __pyx_string_tab[] = { + {&__pyx_n_s_cline_in_traceback, __pyx_k_cline_in_traceback, sizeof(__pyx_k_cline_in_traceback), 0, 0, 1, 1}, + {&__pyx_n_s_main, __pyx_k_main, sizeof(__pyx_k_main), 0, 0, 1, 1}, + {&__pyx_n_s_name, __pyx_k_name, sizeof(__pyx_k_name), 0, 0, 1, 1}, + {&__pyx_n_s_test, __pyx_k_test, sizeof(__pyx_k_test), 0, 0, 1, 1}, + {0, 0, 0, 0, 0, 0, 0} +}; +static CYTHON_SMALL_CODE int __Pyx_InitCachedBuiltins(void) { + return 0; +} + +static CYTHON_SMALL_CODE int __Pyx_InitCachedConstants(void) { + __Pyx_RefNannyDeclarations + __Pyx_RefNannySetupContext("__Pyx_InitCachedConstants", 0); + __Pyx_RefNannyFinishContext(); + return 0; +} + +static CYTHON_SMALL_CODE int __Pyx_InitGlobals(void) { + if (__Pyx_InitStrings(__pyx_string_tab) < 0) __PYX_ERR(0, 1, __pyx_L1_error); + return 0; + __pyx_L1_error:; + return -1; +} + +static CYTHON_SMALL_CODE int __Pyx_modinit_global_init_code(void); /*proto*/ +static CYTHON_SMALL_CODE int __Pyx_modinit_variable_export_code(void); /*proto*/ +static CYTHON_SMALL_CODE int __Pyx_modinit_function_export_code(void); /*proto*/ +static CYTHON_SMALL_CODE int __Pyx_modinit_type_init_code(void); /*proto*/ +static CYTHON_SMALL_CODE int __Pyx_modinit_type_import_code(void); /*proto*/ +static CYTHON_SMALL_CODE int __Pyx_modinit_variable_import_code(void); /*proto*/ +static CYTHON_SMALL_CODE int __Pyx_modinit_function_import_code(void); /*proto*/ + +static int __Pyx_modinit_global_init_code(void) { + __Pyx_RefNannyDeclarations + __Pyx_RefNannySetupContext("__Pyx_modinit_global_init_code", 0); + /*--- Global init code ---*/ + __Pyx_RefNannyFinishContext(); + return 0; +} + +static int __Pyx_modinit_variable_export_code(void) { + __Pyx_RefNannyDeclarations + __Pyx_RefNannySetupContext("__Pyx_modinit_variable_export_code", 0); + /*--- Variable export code ---*/ + __Pyx_RefNannyFinishContext(); + return 0; +} + +static int __Pyx_modinit_function_export_code(void) { + __Pyx_RefNannyDeclarations + __Pyx_RefNannySetupContext("__Pyx_modinit_function_export_code", 0); + /*--- Function export code ---*/ + __Pyx_RefNannyFinishContext(); + return 0; +} + +static int __Pyx_modinit_type_init_code(void) { + __Pyx_RefNannyDeclarations + __Pyx_RefNannySetupContext("__Pyx_modinit_type_init_code", 0); + /*--- Type init code ---*/ + __Pyx_RefNannyFinishContext(); + return 0; +} + +static int __Pyx_modinit_type_import_code(void) { + __Pyx_RefNannyDeclarations + __Pyx_RefNannySetupContext("__Pyx_modinit_type_import_code", 0); + /*--- Type import code ---*/ + __Pyx_RefNannyFinishContext(); + return 0; +} + +static int __Pyx_modinit_variable_import_code(void) { + __Pyx_RefNannyDeclarations + __Pyx_RefNannySetupContext("__Pyx_modinit_variable_import_code", 0); + /*--- Variable import code ---*/ + __Pyx_RefNannyFinishContext(); + return 0; +} + +static int __Pyx_modinit_function_import_code(void) { + __Pyx_RefNannyDeclarations + __Pyx_RefNannySetupContext("__Pyx_modinit_function_import_code", 0); + /*--- Function import code ---*/ + __Pyx_RefNannyFinishContext(); + return 0; +} + + +#ifndef CYTHON_NO_PYINIT_EXPORT +#define __Pyx_PyMODINIT_FUNC PyMODINIT_FUNC +#elif PY_MAJOR_VERSION < 3 +#ifdef __cplusplus +#define __Pyx_PyMODINIT_FUNC extern "C" void +#else +#define __Pyx_PyMODINIT_FUNC void +#endif +#else +#ifdef __cplusplus +#define __Pyx_PyMODINIT_FUNC extern "C" PyObject * +#else +#define __Pyx_PyMODINIT_FUNC PyObject * +#endif +#endif + + +#if PY_MAJOR_VERSION < 3 +__Pyx_PyMODINIT_FUNC initmesh_core_cython(void) CYTHON_SMALL_CODE; /*proto*/ +__Pyx_PyMODINIT_FUNC initmesh_core_cython(void) +#else +__Pyx_PyMODINIT_FUNC PyInit_mesh_core_cython(void) CYTHON_SMALL_CODE; /*proto*/ +__Pyx_PyMODINIT_FUNC PyInit_mesh_core_cython(void) +#if CYTHON_PEP489_MULTI_PHASE_INIT +{ + return PyModuleDef_Init(&__pyx_moduledef); +} +static CYTHON_SMALL_CODE int __Pyx_check_single_interpreter(void) { + #if PY_VERSION_HEX >= 0x030700A1 + static PY_INT64_T main_interpreter_id = -1; + PY_INT64_T current_id = PyInterpreterState_GetID(PyThreadState_Get()->interp); + if (main_interpreter_id == -1) { + main_interpreter_id = current_id; + return (unlikely(current_id == -1)) ? -1 : 0; + } else if (unlikely(main_interpreter_id != current_id)) + #else + static PyInterpreterState *main_interpreter = NULL; + PyInterpreterState *current_interpreter = PyThreadState_Get()->interp; + if (!main_interpreter) { + main_interpreter = current_interpreter; + } else if (unlikely(main_interpreter != current_interpreter)) + #endif + { + PyErr_SetString( + PyExc_ImportError, + "Interpreter change detected - this module can only be loaded into one interpreter per process."); + return -1; + } + return 0; +} +static CYTHON_SMALL_CODE int __Pyx_copy_spec_to_module(PyObject *spec, PyObject *moddict, const char* from_name, const char* to_name, int allow_none) { + PyObject *value = PyObject_GetAttrString(spec, from_name); + int result = 0; + if (likely(value)) { + if (allow_none || value != Py_None) { + result = PyDict_SetItemString(moddict, to_name, value); + } + Py_DECREF(value); + } else if (PyErr_ExceptionMatches(PyExc_AttributeError)) { + PyErr_Clear(); + } else { + result = -1; + } + return result; +} +static CYTHON_SMALL_CODE PyObject* __pyx_pymod_create(PyObject *spec, CYTHON_UNUSED PyModuleDef *def) { + PyObject *module = NULL, *moddict, *modname; + if (__Pyx_check_single_interpreter()) + return NULL; + if (__pyx_m) + return __Pyx_NewRef(__pyx_m); + modname = PyObject_GetAttrString(spec, "name"); + if (unlikely(!modname)) goto bad; + module = PyModule_NewObject(modname); + Py_DECREF(modname); + if (unlikely(!module)) goto bad; + moddict = PyModule_GetDict(module); + if (unlikely(!moddict)) goto bad; + if (unlikely(__Pyx_copy_spec_to_module(spec, moddict, "loader", "__loader__", 1) < 0)) goto bad; + if (unlikely(__Pyx_copy_spec_to_module(spec, moddict, "origin", "__file__", 1) < 0)) goto bad; + if (unlikely(__Pyx_copy_spec_to_module(spec, moddict, "parent", "__package__", 1) < 0)) goto bad; + if (unlikely(__Pyx_copy_spec_to_module(spec, moddict, "submodule_search_locations", "__path__", 0) < 0)) goto bad; + return module; +bad: + Py_XDECREF(module); + return NULL; +} + + +static CYTHON_SMALL_CODE int __pyx_pymod_exec_mesh_core_cython(PyObject *__pyx_pyinit_module) +#endif +#endif +{ + PyObject *__pyx_t_1 = NULL; + int __pyx_lineno = 0; + const char *__pyx_filename = NULL; + int __pyx_clineno = 0; + __Pyx_RefNannyDeclarations + #if CYTHON_PEP489_MULTI_PHASE_INIT + if (__pyx_m) { + if (__pyx_m == __pyx_pyinit_module) return 0; + PyErr_SetString(PyExc_RuntimeError, "Module 'mesh_core_cython' has already been imported. Re-initialisation is not supported."); + return -1; + } + #elif PY_MAJOR_VERSION >= 3 + if (__pyx_m) return __Pyx_NewRef(__pyx_m); + #endif + #if CYTHON_REFNANNY +__Pyx_RefNanny = __Pyx_RefNannyImportAPI("refnanny"); +if (!__Pyx_RefNanny) { + PyErr_Clear(); + __Pyx_RefNanny = __Pyx_RefNannyImportAPI("Cython.Runtime.refnanny"); + if (!__Pyx_RefNanny) + Py_FatalError("failed to import 'refnanny' module"); +} +#endif + __Pyx_RefNannySetupContext("__Pyx_PyMODINIT_FUNC PyInit_mesh_core_cython(void)", 0); + if (__Pyx_check_binary_version() < 0) __PYX_ERR(0, 1, __pyx_L1_error) + #ifdef __Pxy_PyFrame_Initialize_Offsets + __Pxy_PyFrame_Initialize_Offsets(); + #endif + __pyx_empty_tuple = PyTuple_New(0); if (unlikely(!__pyx_empty_tuple)) __PYX_ERR(0, 1, __pyx_L1_error) + __pyx_empty_bytes = PyBytes_FromStringAndSize("", 0); if (unlikely(!__pyx_empty_bytes)) __PYX_ERR(0, 1, __pyx_L1_error) + __pyx_empty_unicode = PyUnicode_FromStringAndSize("", 0); if (unlikely(!__pyx_empty_unicode)) __PYX_ERR(0, 1, __pyx_L1_error) + #ifdef __Pyx_CyFunction_USED + if (__pyx_CyFunction_init() < 0) __PYX_ERR(0, 1, __pyx_L1_error) + #endif + #ifdef __Pyx_FusedFunction_USED + if (__pyx_FusedFunction_init() < 0) __PYX_ERR(0, 1, __pyx_L1_error) + #endif + #ifdef __Pyx_Coroutine_USED + if (__pyx_Coroutine_init() < 0) __PYX_ERR(0, 1, __pyx_L1_error) + #endif + #ifdef __Pyx_Generator_USED + if (__pyx_Generator_init() < 0) __PYX_ERR(0, 1, __pyx_L1_error) + #endif + #ifdef __Pyx_AsyncGen_USED + if (__pyx_AsyncGen_init() < 0) __PYX_ERR(0, 1, __pyx_L1_error) + #endif + #ifdef __Pyx_StopAsyncIteration_USED + if (__pyx_StopAsyncIteration_init() < 0) __PYX_ERR(0, 1, __pyx_L1_error) + #endif + /*--- Library function declarations ---*/ + /*--- Threads initialization code ---*/ + #if defined(WITH_THREAD) && PY_VERSION_HEX < 0x030700F0 && defined(__PYX_FORCE_INIT_THREADS) && __PYX_FORCE_INIT_THREADS + PyEval_InitThreads(); + #endif + /*--- Module creation code ---*/ + #if CYTHON_PEP489_MULTI_PHASE_INIT + __pyx_m = __pyx_pyinit_module; + Py_INCREF(__pyx_m); + #else + #if PY_MAJOR_VERSION < 3 + __pyx_m = Py_InitModule4("mesh_core_cython", __pyx_methods, 0, 0, PYTHON_API_VERSION); Py_XINCREF(__pyx_m); + #else + __pyx_m = PyModule_Create(&__pyx_moduledef); + #endif + if (unlikely(!__pyx_m)) __PYX_ERR(0, 1, __pyx_L1_error) + #endif + __pyx_d = PyModule_GetDict(__pyx_m); if (unlikely(!__pyx_d)) __PYX_ERR(0, 1, __pyx_L1_error) + Py_INCREF(__pyx_d); + __pyx_b = PyImport_AddModule(__Pyx_BUILTIN_MODULE_NAME); if (unlikely(!__pyx_b)) __PYX_ERR(0, 1, __pyx_L1_error) + Py_INCREF(__pyx_b); + __pyx_cython_runtime = PyImport_AddModule((char *) "cython_runtime"); if (unlikely(!__pyx_cython_runtime)) __PYX_ERR(0, 1, __pyx_L1_error) + Py_INCREF(__pyx_cython_runtime); + if (PyObject_SetAttrString(__pyx_m, "__builtins__", __pyx_b) < 0) __PYX_ERR(0, 1, __pyx_L1_error); + /*--- Initialize various global constants etc. ---*/ + if (__Pyx_InitGlobals() < 0) __PYX_ERR(0, 1, __pyx_L1_error) + #if PY_MAJOR_VERSION < 3 && (__PYX_DEFAULT_STRING_ENCODING_IS_ASCII || __PYX_DEFAULT_STRING_ENCODING_IS_DEFAULT) + if (__Pyx_init_sys_getdefaultencoding_params() < 0) __PYX_ERR(0, 1, __pyx_L1_error) + #endif + if (__pyx_module_is_main_mesh_core_cython) { + if (PyObject_SetAttr(__pyx_m, __pyx_n_s_name, __pyx_n_s_main) < 0) __PYX_ERR(0, 1, __pyx_L1_error) + } + #if PY_MAJOR_VERSION >= 3 + { + PyObject *modules = PyImport_GetModuleDict(); if (unlikely(!modules)) __PYX_ERR(0, 1, __pyx_L1_error) + if (!PyDict_GetItemString(modules, "mesh_core_cython")) { + if (unlikely(PyDict_SetItemString(modules, "mesh_core_cython", __pyx_m) < 0)) __PYX_ERR(0, 1, __pyx_L1_error) + } + } + #endif + /*--- Builtin init code ---*/ + if (__Pyx_InitCachedBuiltins() < 0) __PYX_ERR(0, 1, __pyx_L1_error) + /*--- Constants init code ---*/ + if (__Pyx_InitCachedConstants() < 0) __PYX_ERR(0, 1, __pyx_L1_error) + /*--- Global type/function init code ---*/ + (void)__Pyx_modinit_global_init_code(); + (void)__Pyx_modinit_variable_export_code(); + (void)__Pyx_modinit_function_export_code(); + (void)__Pyx_modinit_type_init_code(); + (void)__Pyx_modinit_type_import_code(); + (void)__Pyx_modinit_variable_import_code(); + (void)__Pyx_modinit_function_import_code(); + /*--- Execution code ---*/ + #if defined(__Pyx_Generator_USED) || defined(__Pyx_Coroutine_USED) + if (__Pyx_patch_abc() < 0) __PYX_ERR(0, 1, __pyx_L1_error) + #endif + + /* "mesh_core_cython.pyx":1 + # <<<<<<<<<<<<<< + */ + __pyx_t_1 = __Pyx_PyDict_NewPresized(0); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 1, __pyx_L1_error) + __Pyx_GOTREF(__pyx_t_1); + if (PyDict_SetItem(__pyx_d, __pyx_n_s_test, __pyx_t_1) < 0) __PYX_ERR(0, 1, __pyx_L1_error) + __Pyx_DECREF(__pyx_t_1); __pyx_t_1 = 0; + + /*--- Wrapped vars code ---*/ + + goto __pyx_L0; + __pyx_L1_error:; + __Pyx_XDECREF(__pyx_t_1); + if (__pyx_m) { + if (__pyx_d) { + __Pyx_AddTraceback("init mesh_core_cython", __pyx_clineno, __pyx_lineno, __pyx_filename); + } + Py_CLEAR(__pyx_m); + } else if (!PyErr_Occurred()) { + PyErr_SetString(PyExc_ImportError, "init mesh_core_cython"); + } + __pyx_L0:; + __Pyx_RefNannyFinishContext(); + #if CYTHON_PEP489_MULTI_PHASE_INIT + return (__pyx_m != NULL) ? 0 : -1; + #elif PY_MAJOR_VERSION >= 3 + return __pyx_m; + #else + return; + #endif +} + +/* --- Runtime support code --- */ +/* Refnanny */ +#if CYTHON_REFNANNY +static __Pyx_RefNannyAPIStruct *__Pyx_RefNannyImportAPI(const char *modname) { + PyObject *m = NULL, *p = NULL; + void *r = NULL; + m = PyImport_ImportModule(modname); + if (!m) goto end; + p = PyObject_GetAttrString(m, "RefNannyAPI"); + if (!p) goto end; + r = PyLong_AsVoidPtr(p); +end: + Py_XDECREF(p); + Py_XDECREF(m); + return (__Pyx_RefNannyAPIStruct *)r; +} +#endif + +/* PyDictVersioning */ +#if CYTHON_USE_DICT_VERSIONS && CYTHON_USE_TYPE_SLOTS +static CYTHON_INLINE PY_UINT64_T __Pyx_get_tp_dict_version(PyObject *obj) { + PyObject *dict = Py_TYPE(obj)->tp_dict; + return likely(dict) ? __PYX_GET_DICT_VERSION(dict) : 0; +} +static CYTHON_INLINE PY_UINT64_T __Pyx_get_object_dict_version(PyObject *obj) { + PyObject **dictptr = NULL; + Py_ssize_t offset = Py_TYPE(obj)->tp_dictoffset; + if (offset) { +#if CYTHON_COMPILING_IN_CPYTHON + dictptr = (likely(offset > 0)) ? (PyObject **) ((char *)obj + offset) : _PyObject_GetDictPtr(obj); +#else + dictptr = _PyObject_GetDictPtr(obj); +#endif + } + return (dictptr && *dictptr) ? __PYX_GET_DICT_VERSION(*dictptr) : 0; +} +static CYTHON_INLINE int __Pyx_object_dict_version_matches(PyObject* obj, PY_UINT64_T tp_dict_version, PY_UINT64_T obj_dict_version) { + PyObject *dict = Py_TYPE(obj)->tp_dict; + if (unlikely(!dict) || unlikely(tp_dict_version != __PYX_GET_DICT_VERSION(dict))) + return 0; + return obj_dict_version == __Pyx_get_object_dict_version(obj); +} +#endif + +/* PyObjectGetAttrStr */ +#if CYTHON_USE_TYPE_SLOTS +static CYTHON_INLINE PyObject* __Pyx_PyObject_GetAttrStr(PyObject* obj, PyObject* attr_name) { + PyTypeObject* tp = Py_TYPE(obj); + if (likely(tp->tp_getattro)) + return tp->tp_getattro(obj, attr_name); +#if PY_MAJOR_VERSION < 3 + if (likely(tp->tp_getattr)) + return tp->tp_getattr(obj, PyString_AS_STRING(attr_name)); +#endif + return PyObject_GetAttr(obj, attr_name); +} +#endif + +/* PyErrFetchRestore */ +#if CYTHON_FAST_THREAD_STATE +static CYTHON_INLINE void __Pyx_ErrRestoreInState(PyThreadState *tstate, PyObject *type, PyObject *value, PyObject *tb) { + PyObject *tmp_type, *tmp_value, *tmp_tb; + tmp_type = tstate->curexc_type; + tmp_value = tstate->curexc_value; + tmp_tb = tstate->curexc_traceback; + tstate->curexc_type = type; + tstate->curexc_value = value; + tstate->curexc_traceback = tb; + Py_XDECREF(tmp_type); + Py_XDECREF(tmp_value); + Py_XDECREF(tmp_tb); +} +static CYTHON_INLINE void __Pyx_ErrFetchInState(PyThreadState *tstate, PyObject **type, PyObject **value, PyObject **tb) { + *type = tstate->curexc_type; + *value = tstate->curexc_value; + *tb = tstate->curexc_traceback; + tstate->curexc_type = 0; + tstate->curexc_value = 0; + tstate->curexc_traceback = 0; +} +#endif + +/* CLineInTraceback */ +#ifndef CYTHON_CLINE_IN_TRACEBACK +static int __Pyx_CLineForTraceback(CYTHON_NCP_UNUSED PyThreadState *tstate, int c_line) { + PyObject *use_cline; + PyObject *ptype, *pvalue, *ptraceback; +#if CYTHON_COMPILING_IN_CPYTHON + PyObject **cython_runtime_dict; +#endif + if (unlikely(!__pyx_cython_runtime)) { + return c_line; + } + __Pyx_ErrFetchInState(tstate, &ptype, &pvalue, &ptraceback); +#if CYTHON_COMPILING_IN_CPYTHON + cython_runtime_dict = _PyObject_GetDictPtr(__pyx_cython_runtime); + if (likely(cython_runtime_dict)) { + __PYX_PY_DICT_LOOKUP_IF_MODIFIED( + use_cline, *cython_runtime_dict, + __Pyx_PyDict_GetItemStr(*cython_runtime_dict, __pyx_n_s_cline_in_traceback)) + } else +#endif + { + PyObject *use_cline_obj = __Pyx_PyObject_GetAttrStr(__pyx_cython_runtime, __pyx_n_s_cline_in_traceback); + if (use_cline_obj) { + use_cline = PyObject_Not(use_cline_obj) ? Py_False : Py_True; + Py_DECREF(use_cline_obj); + } else { + PyErr_Clear(); + use_cline = NULL; + } + } + if (!use_cline) { + c_line = 0; + (void) PyObject_SetAttr(__pyx_cython_runtime, __pyx_n_s_cline_in_traceback, Py_False); + } + else if (use_cline == Py_False || (use_cline != Py_True && PyObject_Not(use_cline) != 0)) { + c_line = 0; + } + __Pyx_ErrRestoreInState(tstate, ptype, pvalue, ptraceback); + return c_line; +} +#endif + +/* CodeObjectCache */ +static int __pyx_bisect_code_objects(__Pyx_CodeObjectCacheEntry* entries, int count, int code_line) { + int start = 0, mid = 0, end = count - 1; + if (end >= 0 && code_line > entries[end].code_line) { + return count; + } + while (start < end) { + mid = start + (end - start) / 2; + if (code_line < entries[mid].code_line) { + end = mid; + } else if (code_line > entries[mid].code_line) { + start = mid + 1; + } else { + return mid; + } + } + if (code_line <= entries[mid].code_line) { + return mid; + } else { + return mid + 1; + } +} +static PyCodeObject *__pyx_find_code_object(int code_line) { + PyCodeObject* code_object; + int pos; + if (unlikely(!code_line) || unlikely(!__pyx_code_cache.entries)) { + return NULL; + } + pos = __pyx_bisect_code_objects(__pyx_code_cache.entries, __pyx_code_cache.count, code_line); + if (unlikely(pos >= __pyx_code_cache.count) || unlikely(__pyx_code_cache.entries[pos].code_line != code_line)) { + return NULL; + } + code_object = __pyx_code_cache.entries[pos].code_object; + Py_INCREF(code_object); + return code_object; +} +static void __pyx_insert_code_object(int code_line, PyCodeObject* code_object) { + int pos, i; + __Pyx_CodeObjectCacheEntry* entries = __pyx_code_cache.entries; + if (unlikely(!code_line)) { + return; + } + if (unlikely(!entries)) { + entries = (__Pyx_CodeObjectCacheEntry*)PyMem_Malloc(64*sizeof(__Pyx_CodeObjectCacheEntry)); + if (likely(entries)) { + __pyx_code_cache.entries = entries; + __pyx_code_cache.max_count = 64; + __pyx_code_cache.count = 1; + entries[0].code_line = code_line; + entries[0].code_object = code_object; + Py_INCREF(code_object); + } + return; + } + pos = __pyx_bisect_code_objects(__pyx_code_cache.entries, __pyx_code_cache.count, code_line); + if ((pos < __pyx_code_cache.count) && unlikely(__pyx_code_cache.entries[pos].code_line == code_line)) { + PyCodeObject* tmp = entries[pos].code_object; + entries[pos].code_object = code_object; + Py_DECREF(tmp); + return; + } + if (__pyx_code_cache.count == __pyx_code_cache.max_count) { + int new_max = __pyx_code_cache.max_count + 64; + entries = (__Pyx_CodeObjectCacheEntry*)PyMem_Realloc( + __pyx_code_cache.entries, ((size_t)new_max) * sizeof(__Pyx_CodeObjectCacheEntry)); + if (unlikely(!entries)) { + return; + } + __pyx_code_cache.entries = entries; + __pyx_code_cache.max_count = new_max; + } + for (i=__pyx_code_cache.count; i>pos; i--) { + entries[i] = entries[i-1]; + } + entries[pos].code_line = code_line; + entries[pos].code_object = code_object; + __pyx_code_cache.count++; + Py_INCREF(code_object); +} + +/* AddTraceback */ +#include "compile.h" +#include "frameobject.h" +#include "traceback.h" +static PyCodeObject* __Pyx_CreateCodeObjectForTraceback( + const char *funcname, int c_line, + int py_line, const char *filename) { + PyCodeObject *py_code = NULL; + PyObject *py_funcname = NULL; + #if PY_MAJOR_VERSION < 3 + PyObject *py_srcfile = NULL; + py_srcfile = PyString_FromString(filename); + if (!py_srcfile) goto bad; + #endif + if (c_line) { + #if PY_MAJOR_VERSION < 3 + py_funcname = PyString_FromFormat( "%s (%s:%d)", funcname, __pyx_cfilenm, c_line); + if (!py_funcname) goto bad; + #else + py_funcname = PyUnicode_FromFormat( "%s (%s:%d)", funcname, __pyx_cfilenm, c_line); + if (!py_funcname) goto bad; + funcname = PyUnicode_AsUTF8(py_funcname); + if (!funcname) goto bad; + #endif + } + else { + #if PY_MAJOR_VERSION < 3 + py_funcname = PyString_FromString(funcname); + if (!py_funcname) goto bad; + #endif + } + #if PY_MAJOR_VERSION < 3 + py_code = __Pyx_PyCode_New( + 0, + 0, + 0, + 0, + 0, + __pyx_empty_bytes, /*PyObject *code,*/ + __pyx_empty_tuple, /*PyObject *consts,*/ + __pyx_empty_tuple, /*PyObject *names,*/ + __pyx_empty_tuple, /*PyObject *varnames,*/ + __pyx_empty_tuple, /*PyObject *freevars,*/ + __pyx_empty_tuple, /*PyObject *cellvars,*/ + py_srcfile, /*PyObject *filename,*/ + py_funcname, /*PyObject *name,*/ + py_line, + __pyx_empty_bytes /*PyObject *lnotab*/ + ); + Py_DECREF(py_srcfile); + #else + py_code = PyCode_NewEmpty(filename, funcname, py_line); + #endif + Py_XDECREF(py_funcname); // XDECREF since it's only set on Py3 if cline + return py_code; +bad: + Py_XDECREF(py_funcname); + #if PY_MAJOR_VERSION < 3 + Py_XDECREF(py_srcfile); + #endif + return NULL; +} +static void __Pyx_AddTraceback(const char *funcname, int c_line, + int py_line, const char *filename) { + PyCodeObject *py_code = 0; + PyFrameObject *py_frame = 0; + PyThreadState *tstate = __Pyx_PyThreadState_Current; + if (c_line) { + c_line = __Pyx_CLineForTraceback(tstate, c_line); + } + py_code = __pyx_find_code_object(c_line ? -c_line : py_line); + if (!py_code) { + py_code = __Pyx_CreateCodeObjectForTraceback( + funcname, c_line, py_line, filename); + if (!py_code) goto bad; + __pyx_insert_code_object(c_line ? -c_line : py_line, py_code); + } + py_frame = PyFrame_New( + tstate, /*PyThreadState *tstate,*/ + py_code, /*PyCodeObject *code,*/ + __pyx_d, /*PyObject *globals,*/ + 0 /*PyObject *locals*/ + ); + if (!py_frame) goto bad; + __Pyx_PyFrame_SetLineNumber(py_frame, py_line); + PyTraceBack_Here(py_frame); +bad: + Py_XDECREF(py_code); + Py_XDECREF(py_frame); +} + +/* CIntToPy */ +static CYTHON_INLINE PyObject* __Pyx_PyInt_From_long(long value) { +#ifdef __Pyx_HAS_GCC_DIAGNOSTIC +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wconversion" +#endif + const long neg_one = (long) -1, const_zero = (long) 0; +#ifdef __Pyx_HAS_GCC_DIAGNOSTIC +#pragma GCC diagnostic pop +#endif + const int is_unsigned = neg_one > const_zero; + if (is_unsigned) { + if (sizeof(long) < sizeof(long)) { + return PyInt_FromLong((long) value); + } else if (sizeof(long) <= sizeof(unsigned long)) { + return PyLong_FromUnsignedLong((unsigned long) value); +#ifdef HAVE_LONG_LONG + } else if (sizeof(long) <= sizeof(unsigned PY_LONG_LONG)) { + return PyLong_FromUnsignedLongLong((unsigned PY_LONG_LONG) value); +#endif + } + } else { + if (sizeof(long) <= sizeof(long)) { + return PyInt_FromLong((long) value); +#ifdef HAVE_LONG_LONG + } else if (sizeof(long) <= sizeof(PY_LONG_LONG)) { + return PyLong_FromLongLong((PY_LONG_LONG) value); +#endif + } + } + { + int one = 1; int little = (int)*(unsigned char *)&one; + unsigned char *bytes = (unsigned char *)&value; + return _PyLong_FromByteArray(bytes, sizeof(long), + little, !is_unsigned); + } +} + +/* CIntFromPyVerify */ +#define __PYX_VERIFY_RETURN_INT(target_type, func_type, func_value)\ + __PYX__VERIFY_RETURN_INT(target_type, func_type, func_value, 0) +#define __PYX_VERIFY_RETURN_INT_EXC(target_type, func_type, func_value)\ + __PYX__VERIFY_RETURN_INT(target_type, func_type, func_value, 1) +#define __PYX__VERIFY_RETURN_INT(target_type, func_type, func_value, exc)\ + {\ + func_type value = func_value;\ + if (sizeof(target_type) < sizeof(func_type)) {\ + if (unlikely(value != (func_type) (target_type) value)) {\ + func_type zero = 0;\ + if (exc && unlikely(value == (func_type)-1 && PyErr_Occurred()))\ + return (target_type) -1;\ + if (is_unsigned && unlikely(value < zero))\ + goto raise_neg_overflow;\ + else\ + goto raise_overflow;\ + }\ + }\ + return (target_type) value;\ + } + +/* CIntFromPy */ +static CYTHON_INLINE long __Pyx_PyInt_As_long(PyObject *x) { +#ifdef __Pyx_HAS_GCC_DIAGNOSTIC +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wconversion" +#endif + const long neg_one = (long) -1, const_zero = (long) 0; +#ifdef __Pyx_HAS_GCC_DIAGNOSTIC +#pragma GCC diagnostic pop +#endif + const int is_unsigned = neg_one > const_zero; +#if PY_MAJOR_VERSION < 3 + if (likely(PyInt_Check(x))) { + if (sizeof(long) < sizeof(long)) { + __PYX_VERIFY_RETURN_INT(long, long, PyInt_AS_LONG(x)) + } else { + long val = PyInt_AS_LONG(x); + if (is_unsigned && unlikely(val < 0)) { + goto raise_neg_overflow; + } + return (long) val; + } + } else +#endif + if (likely(PyLong_Check(x))) { + if (is_unsigned) { +#if CYTHON_USE_PYLONG_INTERNALS + const digit* digits = ((PyLongObject*)x)->ob_digit; + switch (Py_SIZE(x)) { + case 0: return (long) 0; + case 1: __PYX_VERIFY_RETURN_INT(long, digit, digits[0]) + case 2: + if (8 * sizeof(long) > 1 * PyLong_SHIFT) { + if (8 * sizeof(unsigned long) > 2 * PyLong_SHIFT) { + __PYX_VERIFY_RETURN_INT(long, unsigned long, (((((unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) + } else if (8 * sizeof(long) >= 2 * PyLong_SHIFT) { + return (long) (((((long)digits[1]) << PyLong_SHIFT) | (long)digits[0])); + } + } + break; + case 3: + if (8 * sizeof(long) > 2 * PyLong_SHIFT) { + if (8 * sizeof(unsigned long) > 3 * PyLong_SHIFT) { + __PYX_VERIFY_RETURN_INT(long, unsigned long, (((((((unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) + } else if (8 * sizeof(long) >= 3 * PyLong_SHIFT) { + return (long) (((((((long)digits[2]) << PyLong_SHIFT) | (long)digits[1]) << PyLong_SHIFT) | (long)digits[0])); + } + } + break; + case 4: + if (8 * sizeof(long) > 3 * PyLong_SHIFT) { + if (8 * sizeof(unsigned long) > 4 * PyLong_SHIFT) { + __PYX_VERIFY_RETURN_INT(long, unsigned long, (((((((((unsigned long)digits[3]) << PyLong_SHIFT) | (unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) + } else if (8 * sizeof(long) >= 4 * PyLong_SHIFT) { + return (long) (((((((((long)digits[3]) << PyLong_SHIFT) | (long)digits[2]) << PyLong_SHIFT) | (long)digits[1]) << PyLong_SHIFT) | (long)digits[0])); + } + } + break; + } +#endif +#if CYTHON_COMPILING_IN_CPYTHON + if (unlikely(Py_SIZE(x) < 0)) { + goto raise_neg_overflow; + } +#else + { + int result = PyObject_RichCompareBool(x, Py_False, Py_LT); + if (unlikely(result < 0)) + return (long) -1; + if (unlikely(result == 1)) + goto raise_neg_overflow; + } +#endif + if (sizeof(long) <= sizeof(unsigned long)) { + __PYX_VERIFY_RETURN_INT_EXC(long, unsigned long, PyLong_AsUnsignedLong(x)) +#ifdef HAVE_LONG_LONG + } else if (sizeof(long) <= sizeof(unsigned PY_LONG_LONG)) { + __PYX_VERIFY_RETURN_INT_EXC(long, unsigned PY_LONG_LONG, PyLong_AsUnsignedLongLong(x)) +#endif + } + } else { +#if CYTHON_USE_PYLONG_INTERNALS + const digit* digits = ((PyLongObject*)x)->ob_digit; + switch (Py_SIZE(x)) { + case 0: return (long) 0; + case -1: __PYX_VERIFY_RETURN_INT(long, sdigit, (sdigit) (-(sdigit)digits[0])) + case 1: __PYX_VERIFY_RETURN_INT(long, digit, +digits[0]) + case -2: + if (8 * sizeof(long) - 1 > 1 * PyLong_SHIFT) { + if (8 * sizeof(unsigned long) > 2 * PyLong_SHIFT) { + __PYX_VERIFY_RETURN_INT(long, long, -(long) (((((unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) + } else if (8 * sizeof(long) - 1 > 2 * PyLong_SHIFT) { + return (long) (((long)-1)*(((((long)digits[1]) << PyLong_SHIFT) | (long)digits[0]))); + } + } + break; + case 2: + if (8 * sizeof(long) > 1 * PyLong_SHIFT) { + if (8 * sizeof(unsigned long) > 2 * PyLong_SHIFT) { + __PYX_VERIFY_RETURN_INT(long, unsigned long, (((((unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) + } else if (8 * sizeof(long) - 1 > 2 * PyLong_SHIFT) { + return (long) ((((((long)digits[1]) << PyLong_SHIFT) | (long)digits[0]))); + } + } + break; + case -3: + if (8 * sizeof(long) - 1 > 2 * PyLong_SHIFT) { + if (8 * sizeof(unsigned long) > 3 * PyLong_SHIFT) { + __PYX_VERIFY_RETURN_INT(long, long, -(long) (((((((unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) + } else if (8 * sizeof(long) - 1 > 3 * PyLong_SHIFT) { + return (long) (((long)-1)*(((((((long)digits[2]) << PyLong_SHIFT) | (long)digits[1]) << PyLong_SHIFT) | (long)digits[0]))); + } + } + break; + case 3: + if (8 * sizeof(long) > 2 * PyLong_SHIFT) { + if (8 * sizeof(unsigned long) > 3 * PyLong_SHIFT) { + __PYX_VERIFY_RETURN_INT(long, unsigned long, (((((((unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) + } else if (8 * sizeof(long) - 1 > 3 * PyLong_SHIFT) { + return (long) ((((((((long)digits[2]) << PyLong_SHIFT) | (long)digits[1]) << PyLong_SHIFT) | (long)digits[0]))); + } + } + break; + case -4: + if (8 * sizeof(long) - 1 > 3 * PyLong_SHIFT) { + if (8 * sizeof(unsigned long) > 4 * PyLong_SHIFT) { + __PYX_VERIFY_RETURN_INT(long, long, -(long) (((((((((unsigned long)digits[3]) << PyLong_SHIFT) | (unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) + } else if (8 * sizeof(long) - 1 > 4 * PyLong_SHIFT) { + return (long) (((long)-1)*(((((((((long)digits[3]) << PyLong_SHIFT) | (long)digits[2]) << PyLong_SHIFT) | (long)digits[1]) << PyLong_SHIFT) | (long)digits[0]))); + } + } + break; + case 4: + if (8 * sizeof(long) > 3 * PyLong_SHIFT) { + if (8 * sizeof(unsigned long) > 4 * PyLong_SHIFT) { + __PYX_VERIFY_RETURN_INT(long, unsigned long, (((((((((unsigned long)digits[3]) << PyLong_SHIFT) | (unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) + } else if (8 * sizeof(long) - 1 > 4 * PyLong_SHIFT) { + return (long) ((((((((((long)digits[3]) << PyLong_SHIFT) | (long)digits[2]) << PyLong_SHIFT) | (long)digits[1]) << PyLong_SHIFT) | (long)digits[0]))); + } + } + break; + } +#endif + if (sizeof(long) <= sizeof(long)) { + __PYX_VERIFY_RETURN_INT_EXC(long, long, PyLong_AsLong(x)) +#ifdef HAVE_LONG_LONG + } else if (sizeof(long) <= sizeof(PY_LONG_LONG)) { + __PYX_VERIFY_RETURN_INT_EXC(long, PY_LONG_LONG, PyLong_AsLongLong(x)) +#endif + } + } + { +#if CYTHON_COMPILING_IN_PYPY && !defined(_PyLong_AsByteArray) + PyErr_SetString(PyExc_RuntimeError, + "_PyLong_AsByteArray() not available in PyPy, cannot convert large numbers"); +#else + long val; + PyObject *v = __Pyx_PyNumber_IntOrLong(x); + #if PY_MAJOR_VERSION < 3 + if (likely(v) && !PyLong_Check(v)) { + PyObject *tmp = v; + v = PyNumber_Long(tmp); + Py_DECREF(tmp); + } + #endif + if (likely(v)) { + int one = 1; int is_little = (int)*(unsigned char *)&one; + unsigned char *bytes = (unsigned char *)&val; + int ret = _PyLong_AsByteArray((PyLongObject *)v, + bytes, sizeof(val), + is_little, !is_unsigned); + Py_DECREF(v); + if (likely(!ret)) + return val; + } +#endif + return (long) -1; + } + } else { + long val; + PyObject *tmp = __Pyx_PyNumber_IntOrLong(x); + if (!tmp) return (long) -1; + val = __Pyx_PyInt_As_long(tmp); + Py_DECREF(tmp); + return val; + } +raise_overflow: + PyErr_SetString(PyExc_OverflowError, + "value too large to convert to long"); + return (long) -1; +raise_neg_overflow: + PyErr_SetString(PyExc_OverflowError, + "can't convert negative value to long"); + return (long) -1; +} + +/* CIntFromPy */ +static CYTHON_INLINE int __Pyx_PyInt_As_int(PyObject *x) { +#ifdef __Pyx_HAS_GCC_DIAGNOSTIC +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wconversion" +#endif + const int neg_one = (int) -1, const_zero = (int) 0; +#ifdef __Pyx_HAS_GCC_DIAGNOSTIC +#pragma GCC diagnostic pop +#endif + const int is_unsigned = neg_one > const_zero; +#if PY_MAJOR_VERSION < 3 + if (likely(PyInt_Check(x))) { + if (sizeof(int) < sizeof(long)) { + __PYX_VERIFY_RETURN_INT(int, long, PyInt_AS_LONG(x)) + } else { + long val = PyInt_AS_LONG(x); + if (is_unsigned && unlikely(val < 0)) { + goto raise_neg_overflow; + } + return (int) val; + } + } else +#endif + if (likely(PyLong_Check(x))) { + if (is_unsigned) { +#if CYTHON_USE_PYLONG_INTERNALS + const digit* digits = ((PyLongObject*)x)->ob_digit; + switch (Py_SIZE(x)) { + case 0: return (int) 0; + case 1: __PYX_VERIFY_RETURN_INT(int, digit, digits[0]) + case 2: + if (8 * sizeof(int) > 1 * PyLong_SHIFT) { + if (8 * sizeof(unsigned long) > 2 * PyLong_SHIFT) { + __PYX_VERIFY_RETURN_INT(int, unsigned long, (((((unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) + } else if (8 * sizeof(int) >= 2 * PyLong_SHIFT) { + return (int) (((((int)digits[1]) << PyLong_SHIFT) | (int)digits[0])); + } + } + break; + case 3: + if (8 * sizeof(int) > 2 * PyLong_SHIFT) { + if (8 * sizeof(unsigned long) > 3 * PyLong_SHIFT) { + __PYX_VERIFY_RETURN_INT(int, unsigned long, (((((((unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) + } else if (8 * sizeof(int) >= 3 * PyLong_SHIFT) { + return (int) (((((((int)digits[2]) << PyLong_SHIFT) | (int)digits[1]) << PyLong_SHIFT) | (int)digits[0])); + } + } + break; + case 4: + if (8 * sizeof(int) > 3 * PyLong_SHIFT) { + if (8 * sizeof(unsigned long) > 4 * PyLong_SHIFT) { + __PYX_VERIFY_RETURN_INT(int, unsigned long, (((((((((unsigned long)digits[3]) << PyLong_SHIFT) | (unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) + } else if (8 * sizeof(int) >= 4 * PyLong_SHIFT) { + return (int) (((((((((int)digits[3]) << PyLong_SHIFT) | (int)digits[2]) << PyLong_SHIFT) | (int)digits[1]) << PyLong_SHIFT) | (int)digits[0])); + } + } + break; + } +#endif +#if CYTHON_COMPILING_IN_CPYTHON + if (unlikely(Py_SIZE(x) < 0)) { + goto raise_neg_overflow; + } +#else + { + int result = PyObject_RichCompareBool(x, Py_False, Py_LT); + if (unlikely(result < 0)) + return (int) -1; + if (unlikely(result == 1)) + goto raise_neg_overflow; + } +#endif + if (sizeof(int) <= sizeof(unsigned long)) { + __PYX_VERIFY_RETURN_INT_EXC(int, unsigned long, PyLong_AsUnsignedLong(x)) +#ifdef HAVE_LONG_LONG + } else if (sizeof(int) <= sizeof(unsigned PY_LONG_LONG)) { + __PYX_VERIFY_RETURN_INT_EXC(int, unsigned PY_LONG_LONG, PyLong_AsUnsignedLongLong(x)) +#endif + } + } else { +#if CYTHON_USE_PYLONG_INTERNALS + const digit* digits = ((PyLongObject*)x)->ob_digit; + switch (Py_SIZE(x)) { + case 0: return (int) 0; + case -1: __PYX_VERIFY_RETURN_INT(int, sdigit, (sdigit) (-(sdigit)digits[0])) + case 1: __PYX_VERIFY_RETURN_INT(int, digit, +digits[0]) + case -2: + if (8 * sizeof(int) - 1 > 1 * PyLong_SHIFT) { + if (8 * sizeof(unsigned long) > 2 * PyLong_SHIFT) { + __PYX_VERIFY_RETURN_INT(int, long, -(long) (((((unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) + } else if (8 * sizeof(int) - 1 > 2 * PyLong_SHIFT) { + return (int) (((int)-1)*(((((int)digits[1]) << PyLong_SHIFT) | (int)digits[0]))); + } + } + break; + case 2: + if (8 * sizeof(int) > 1 * PyLong_SHIFT) { + if (8 * sizeof(unsigned long) > 2 * PyLong_SHIFT) { + __PYX_VERIFY_RETURN_INT(int, unsigned long, (((((unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) + } else if (8 * sizeof(int) - 1 > 2 * PyLong_SHIFT) { + return (int) ((((((int)digits[1]) << PyLong_SHIFT) | (int)digits[0]))); + } + } + break; + case -3: + if (8 * sizeof(int) - 1 > 2 * PyLong_SHIFT) { + if (8 * sizeof(unsigned long) > 3 * PyLong_SHIFT) { + __PYX_VERIFY_RETURN_INT(int, long, -(long) (((((((unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) + } else if (8 * sizeof(int) - 1 > 3 * PyLong_SHIFT) { + return (int) (((int)-1)*(((((((int)digits[2]) << PyLong_SHIFT) | (int)digits[1]) << PyLong_SHIFT) | (int)digits[0]))); + } + } + break; + case 3: + if (8 * sizeof(int) > 2 * PyLong_SHIFT) { + if (8 * sizeof(unsigned long) > 3 * PyLong_SHIFT) { + __PYX_VERIFY_RETURN_INT(int, unsigned long, (((((((unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) + } else if (8 * sizeof(int) - 1 > 3 * PyLong_SHIFT) { + return (int) ((((((((int)digits[2]) << PyLong_SHIFT) | (int)digits[1]) << PyLong_SHIFT) | (int)digits[0]))); + } + } + break; + case -4: + if (8 * sizeof(int) - 1 > 3 * PyLong_SHIFT) { + if (8 * sizeof(unsigned long) > 4 * PyLong_SHIFT) { + __PYX_VERIFY_RETURN_INT(int, long, -(long) (((((((((unsigned long)digits[3]) << PyLong_SHIFT) | (unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) + } else if (8 * sizeof(int) - 1 > 4 * PyLong_SHIFT) { + return (int) (((int)-1)*(((((((((int)digits[3]) << PyLong_SHIFT) | (int)digits[2]) << PyLong_SHIFT) | (int)digits[1]) << PyLong_SHIFT) | (int)digits[0]))); + } + } + break; + case 4: + if (8 * sizeof(int) > 3 * PyLong_SHIFT) { + if (8 * sizeof(unsigned long) > 4 * PyLong_SHIFT) { + __PYX_VERIFY_RETURN_INT(int, unsigned long, (((((((((unsigned long)digits[3]) << PyLong_SHIFT) | (unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) + } else if (8 * sizeof(int) - 1 > 4 * PyLong_SHIFT) { + return (int) ((((((((((int)digits[3]) << PyLong_SHIFT) | (int)digits[2]) << PyLong_SHIFT) | (int)digits[1]) << PyLong_SHIFT) | (int)digits[0]))); + } + } + break; + } +#endif + if (sizeof(int) <= sizeof(long)) { + __PYX_VERIFY_RETURN_INT_EXC(int, long, PyLong_AsLong(x)) +#ifdef HAVE_LONG_LONG + } else if (sizeof(int) <= sizeof(PY_LONG_LONG)) { + __PYX_VERIFY_RETURN_INT_EXC(int, PY_LONG_LONG, PyLong_AsLongLong(x)) +#endif + } + } + { +#if CYTHON_COMPILING_IN_PYPY && !defined(_PyLong_AsByteArray) + PyErr_SetString(PyExc_RuntimeError, + "_PyLong_AsByteArray() not available in PyPy, cannot convert large numbers"); +#else + int val; + PyObject *v = __Pyx_PyNumber_IntOrLong(x); + #if PY_MAJOR_VERSION < 3 + if (likely(v) && !PyLong_Check(v)) { + PyObject *tmp = v; + v = PyNumber_Long(tmp); + Py_DECREF(tmp); + } + #endif + if (likely(v)) { + int one = 1; int is_little = (int)*(unsigned char *)&one; + unsigned char *bytes = (unsigned char *)&val; + int ret = _PyLong_AsByteArray((PyLongObject *)v, + bytes, sizeof(val), + is_little, !is_unsigned); + Py_DECREF(v); + if (likely(!ret)) + return val; + } +#endif + return (int) -1; + } + } else { + int val; + PyObject *tmp = __Pyx_PyNumber_IntOrLong(x); + if (!tmp) return (int) -1; + val = __Pyx_PyInt_As_int(tmp); + Py_DECREF(tmp); + return val; + } +raise_overflow: + PyErr_SetString(PyExc_OverflowError, + "value too large to convert to int"); + return (int) -1; +raise_neg_overflow: + PyErr_SetString(PyExc_OverflowError, + "can't convert negative value to int"); + return (int) -1; +} + +/* FastTypeChecks */ +#if CYTHON_COMPILING_IN_CPYTHON +static int __Pyx_InBases(PyTypeObject *a, PyTypeObject *b) { + while (a) { + a = a->tp_base; + if (a == b) + return 1; + } + return b == &PyBaseObject_Type; +} +static CYTHON_INLINE int __Pyx_IsSubtype(PyTypeObject *a, PyTypeObject *b) { + PyObject *mro; + if (a == b) return 1; + mro = a->tp_mro; + if (likely(mro)) { + Py_ssize_t i, n; + n = PyTuple_GET_SIZE(mro); + for (i = 0; i < n; i++) { + if (PyTuple_GET_ITEM(mro, i) == (PyObject *)b) + return 1; + } + return 0; + } + return __Pyx_InBases(a, b); +} +#if PY_MAJOR_VERSION == 2 +static int __Pyx_inner_PyErr_GivenExceptionMatches2(PyObject *err, PyObject* exc_type1, PyObject* exc_type2) { + PyObject *exception, *value, *tb; + int res; + __Pyx_PyThreadState_declare + __Pyx_PyThreadState_assign + __Pyx_ErrFetch(&exception, &value, &tb); + res = exc_type1 ? PyObject_IsSubclass(err, exc_type1) : 0; + if (unlikely(res == -1)) { + PyErr_WriteUnraisable(err); + res = 0; + } + if (!res) { + res = PyObject_IsSubclass(err, exc_type2); + if (unlikely(res == -1)) { + PyErr_WriteUnraisable(err); + res = 0; + } + } + __Pyx_ErrRestore(exception, value, tb); + return res; +} +#else +static CYTHON_INLINE int __Pyx_inner_PyErr_GivenExceptionMatches2(PyObject *err, PyObject* exc_type1, PyObject *exc_type2) { + int res = exc_type1 ? __Pyx_IsSubtype((PyTypeObject*)err, (PyTypeObject*)exc_type1) : 0; + if (!res) { + res = __Pyx_IsSubtype((PyTypeObject*)err, (PyTypeObject*)exc_type2); + } + return res; +} +#endif +static int __Pyx_PyErr_GivenExceptionMatchesTuple(PyObject *exc_type, PyObject *tuple) { + Py_ssize_t i, n; + assert(PyExceptionClass_Check(exc_type)); + n = PyTuple_GET_SIZE(tuple); +#if PY_MAJOR_VERSION >= 3 + for (i=0; ip) { + #if PY_MAJOR_VERSION < 3 + if (t->is_unicode) { + *t->p = PyUnicode_DecodeUTF8(t->s, t->n - 1, NULL); + } else if (t->intern) { + *t->p = PyString_InternFromString(t->s); + } else { + *t->p = PyString_FromStringAndSize(t->s, t->n - 1); + } + #else + if (t->is_unicode | t->is_str) { + if (t->intern) { + *t->p = PyUnicode_InternFromString(t->s); + } else if (t->encoding) { + *t->p = PyUnicode_Decode(t->s, t->n - 1, t->encoding, NULL); + } else { + *t->p = PyUnicode_FromStringAndSize(t->s, t->n - 1); + } + } else { + *t->p = PyBytes_FromStringAndSize(t->s, t->n - 1); + } + #endif + if (!*t->p) + return -1; + if (PyObject_Hash(*t->p) == -1) + return -1; + ++t; + } + return 0; +} + +static CYTHON_INLINE PyObject* __Pyx_PyUnicode_FromString(const char* c_str) { + return __Pyx_PyUnicode_FromStringAndSize(c_str, (Py_ssize_t)strlen(c_str)); +} +static CYTHON_INLINE const char* __Pyx_PyObject_AsString(PyObject* o) { + Py_ssize_t ignore; + return __Pyx_PyObject_AsStringAndSize(o, &ignore); +} +#if __PYX_DEFAULT_STRING_ENCODING_IS_ASCII || __PYX_DEFAULT_STRING_ENCODING_IS_DEFAULT +#if !CYTHON_PEP393_ENABLED +static const char* __Pyx_PyUnicode_AsStringAndSize(PyObject* o, Py_ssize_t *length) { + char* defenc_c; + PyObject* defenc = _PyUnicode_AsDefaultEncodedString(o, NULL); + if (!defenc) return NULL; + defenc_c = PyBytes_AS_STRING(defenc); +#if __PYX_DEFAULT_STRING_ENCODING_IS_ASCII + { + char* end = defenc_c + PyBytes_GET_SIZE(defenc); + char* c; + for (c = defenc_c; c < end; c++) { + if ((unsigned char) (*c) >= 128) { + PyUnicode_AsASCIIString(o); + return NULL; + } + } + } +#endif + *length = PyBytes_GET_SIZE(defenc); + return defenc_c; +} +#else +static CYTHON_INLINE const char* __Pyx_PyUnicode_AsStringAndSize(PyObject* o, Py_ssize_t *length) { + if (unlikely(__Pyx_PyUnicode_READY(o) == -1)) return NULL; +#if __PYX_DEFAULT_STRING_ENCODING_IS_ASCII + if (likely(PyUnicode_IS_ASCII(o))) { + *length = PyUnicode_GET_LENGTH(o); + return PyUnicode_AsUTF8(o); + } else { + PyUnicode_AsASCIIString(o); + return NULL; + } +#else + return PyUnicode_AsUTF8AndSize(o, length); +#endif +} +#endif +#endif +static CYTHON_INLINE const char* __Pyx_PyObject_AsStringAndSize(PyObject* o, Py_ssize_t *length) { +#if __PYX_DEFAULT_STRING_ENCODING_IS_ASCII || __PYX_DEFAULT_STRING_ENCODING_IS_DEFAULT + if ( +#if PY_MAJOR_VERSION < 3 && __PYX_DEFAULT_STRING_ENCODING_IS_ASCII + __Pyx_sys_getdefaultencoding_not_ascii && +#endif + PyUnicode_Check(o)) { + return __Pyx_PyUnicode_AsStringAndSize(o, length); + } else +#endif +#if (!CYTHON_COMPILING_IN_PYPY) || (defined(PyByteArray_AS_STRING) && defined(PyByteArray_GET_SIZE)) + if (PyByteArray_Check(o)) { + *length = PyByteArray_GET_SIZE(o); + return PyByteArray_AS_STRING(o); + } else +#endif + { + char* result; + int r = PyBytes_AsStringAndSize(o, &result, length); + if (unlikely(r < 0)) { + return NULL; + } else { + return result; + } + } +} +static CYTHON_INLINE int __Pyx_PyObject_IsTrue(PyObject* x) { + int is_true = x == Py_True; + if (is_true | (x == Py_False) | (x == Py_None)) return is_true; + else return PyObject_IsTrue(x); +} +static CYTHON_INLINE int __Pyx_PyObject_IsTrueAndDecref(PyObject* x) { + int retval; + if (unlikely(!x)) return -1; + retval = __Pyx_PyObject_IsTrue(x); + Py_DECREF(x); + return retval; +} +static PyObject* __Pyx_PyNumber_IntOrLongWrongResultType(PyObject* result, const char* type_name) { +#if PY_MAJOR_VERSION >= 3 + if (PyLong_Check(result)) { + if (PyErr_WarnFormat(PyExc_DeprecationWarning, 1, + "__int__ returned non-int (type %.200s). " + "The ability to return an instance of a strict subclass of int " + "is deprecated, and may be removed in a future version of Python.", + Py_TYPE(result)->tp_name)) { + Py_DECREF(result); + return NULL; + } + return result; + } +#endif + PyErr_Format(PyExc_TypeError, + "__%.4s__ returned non-%.4s (type %.200s)", + type_name, type_name, Py_TYPE(result)->tp_name); + Py_DECREF(result); + return NULL; +} +static CYTHON_INLINE PyObject* __Pyx_PyNumber_IntOrLong(PyObject* x) { +#if CYTHON_USE_TYPE_SLOTS + PyNumberMethods *m; +#endif + const char *name = NULL; + PyObject *res = NULL; +#if PY_MAJOR_VERSION < 3 + if (likely(PyInt_Check(x) || PyLong_Check(x))) +#else + if (likely(PyLong_Check(x))) +#endif + return __Pyx_NewRef(x); +#if CYTHON_USE_TYPE_SLOTS + m = Py_TYPE(x)->tp_as_number; + #if PY_MAJOR_VERSION < 3 + if (m && m->nb_int) { + name = "int"; + res = m->nb_int(x); + } + else if (m && m->nb_long) { + name = "long"; + res = m->nb_long(x); + } + #else + if (likely(m && m->nb_int)) { + name = "int"; + res = m->nb_int(x); + } + #endif +#else + if (!PyBytes_CheckExact(x) && !PyUnicode_CheckExact(x)) { + res = PyNumber_Int(x); + } +#endif + if (likely(res)) { +#if PY_MAJOR_VERSION < 3 + if (unlikely(!PyInt_Check(res) && !PyLong_Check(res))) { +#else + if (unlikely(!PyLong_CheckExact(res))) { +#endif + return __Pyx_PyNumber_IntOrLongWrongResultType(res, name); + } + } + else if (!PyErr_Occurred()) { + PyErr_SetString(PyExc_TypeError, + "an integer is required"); + } + return res; +} +static CYTHON_INLINE Py_ssize_t __Pyx_PyIndex_AsSsize_t(PyObject* b) { + Py_ssize_t ival; + PyObject *x; +#if PY_MAJOR_VERSION < 3 + if (likely(PyInt_CheckExact(b))) { + if (sizeof(Py_ssize_t) >= sizeof(long)) + return PyInt_AS_LONG(b); + else + return PyInt_AsSsize_t(b); + } +#endif + if (likely(PyLong_CheckExact(b))) { + #if CYTHON_USE_PYLONG_INTERNALS + const digit* digits = ((PyLongObject*)b)->ob_digit; + const Py_ssize_t size = Py_SIZE(b); + if (likely(__Pyx_sst_abs(size) <= 1)) { + ival = likely(size) ? digits[0] : 0; + if (size == -1) ival = -ival; + return ival; + } else { + switch (size) { + case 2: + if (8 * sizeof(Py_ssize_t) > 2 * PyLong_SHIFT) { + return (Py_ssize_t) (((((size_t)digits[1]) << PyLong_SHIFT) | (size_t)digits[0])); + } + break; + case -2: + if (8 * sizeof(Py_ssize_t) > 2 * PyLong_SHIFT) { + return -(Py_ssize_t) (((((size_t)digits[1]) << PyLong_SHIFT) | (size_t)digits[0])); + } + break; + case 3: + if (8 * sizeof(Py_ssize_t) > 3 * PyLong_SHIFT) { + return (Py_ssize_t) (((((((size_t)digits[2]) << PyLong_SHIFT) | (size_t)digits[1]) << PyLong_SHIFT) | (size_t)digits[0])); + } + break; + case -3: + if (8 * sizeof(Py_ssize_t) > 3 * PyLong_SHIFT) { + return -(Py_ssize_t) (((((((size_t)digits[2]) << PyLong_SHIFT) | (size_t)digits[1]) << PyLong_SHIFT) | (size_t)digits[0])); + } + break; + case 4: + if (8 * sizeof(Py_ssize_t) > 4 * PyLong_SHIFT) { + return (Py_ssize_t) (((((((((size_t)digits[3]) << PyLong_SHIFT) | (size_t)digits[2]) << PyLong_SHIFT) | (size_t)digits[1]) << PyLong_SHIFT) | (size_t)digits[0])); + } + break; + case -4: + if (8 * sizeof(Py_ssize_t) > 4 * PyLong_SHIFT) { + return -(Py_ssize_t) (((((((((size_t)digits[3]) << PyLong_SHIFT) | (size_t)digits[2]) << PyLong_SHIFT) | (size_t)digits[1]) << PyLong_SHIFT) | (size_t)digits[0])); + } + break; + } + } + #endif + return PyLong_AsSsize_t(b); + } + x = PyNumber_Index(b); + if (!x) return -1; + ival = PyInt_AsSsize_t(x); + Py_DECREF(x); + return ival; +} +static CYTHON_INLINE Py_hash_t __Pyx_PyIndex_AsHash_t(PyObject* o) { + if (sizeof(Py_hash_t) == sizeof(Py_ssize_t)) { + return (Py_hash_t) __Pyx_PyIndex_AsSsize_t(o); +#if PY_MAJOR_VERSION < 3 + } else if (likely(PyInt_CheckExact(o))) { + return PyInt_AS_LONG(o); +#endif + } else { + Py_ssize_t ival; + PyObject *x; + x = PyNumber_Index(o); + if (!x) return -1; + ival = PyInt_AsLong(x); + Py_DECREF(x); + return ival; + } +} +static CYTHON_INLINE PyObject * __Pyx_PyBool_FromLong(long b) { + return b ? __Pyx_NewRef(Py_True) : __Pyx_NewRef(Py_False); +} +static CYTHON_INLINE PyObject * __Pyx_PyInt_FromSize_t(size_t ival) { + return PyInt_FromSize_t(ival); +} + + +#endif /* Py_PYTHON_H */ diff --git a/hair_service_sd/core/face3d/mesh/cython/mesh_core_cython.pyx b/hair_service_sd/core/face3d/mesh/cython/mesh_core_cython.pyx new file mode 100644 index 0000000..e69de29 diff --git a/hair_service_sd/core/face3d/mesh/cython/setup.py b/hair_service_sd/core/face3d/mesh/cython/setup.py new file mode 100644 index 0000000..1422a2c --- /dev/null +++ b/hair_service_sd/core/face3d/mesh/cython/setup.py @@ -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()])], +) + diff --git a/hair_service_sd/core/face3d/mesh/io.py b/hair_service_sd/core/face3d/mesh/io.py new file mode 100644 index 0000000..9c90c7e --- /dev/null +++ b/hair_service_sd/core/face3d/mesh/io.py @@ -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) \ No newline at end of file diff --git a/hair_service_sd/core/face3d/mesh/light.py b/hair_service_sd/core/face3d/mesh/light.py new file mode 100644 index 0000000..17c49ee --- /dev/null +++ b/hair_service_sd/core/face3d/mesh/light.py @@ -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 + diff --git a/hair_service_sd/core/face3d/mesh/render.py b/hair_service_sd/core/face3d/mesh/render.py new file mode 100644 index 0000000..8951ac8 --- /dev/null +++ b/hair_service_sd/core/face3d/mesh/render.py @@ -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 + diff --git a/hair_service_sd/core/face3d/mesh/transform.py b/hair_service_sd/core/face3d/mesh/transform.py new file mode 100644 index 0000000..d91b09f --- /dev/null +++ b/hair_service_sd/core/face3d/mesh/transform.py @@ -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 \ No newline at end of file diff --git a/hair_service_sd/core/face3d/mesh/vis.py b/hair_service_sd/core/face3d/mesh/vis.py new file mode 100644 index 0000000..6fffe14 --- /dev/null +++ b/hair_service_sd/core/face3d/mesh/vis.py @@ -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? diff --git a/hair_service_sd/core/face3d/morphable_model/__init__.py b/hair_service_sd/core/face3d/morphable_model/__init__.py new file mode 100644 index 0000000..19fc6a6 --- /dev/null +++ b/hair_service_sd/core/face3d/morphable_model/__init__.py @@ -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 \ No newline at end of file diff --git a/hair_service_sd/core/face3d/morphable_model/fit.py b/hair_service_sd/core/face3d/morphable_model/fit.py new file mode 100644 index 0000000..d1621bb --- /dev/null +++ b/hair_service_sd/core/face3d/morphable_model/fit.py @@ -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) \ No newline at end of file diff --git a/hair_service_sd/core/face3d/morphable_model/load.py b/hair_service_sd/core/face3d/morphable_model/load.py new file mode 100644 index 0000000..0b80665 --- /dev/null +++ b/hair_service_sd/core/face3d/morphable_model/load.py @@ -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) diff --git a/hair_service_sd/core/face3d/morphable_model/morphabel_model.py b/hair_service_sd/core/face3d/morphable_model/morphabel_model.py new file mode 100644 index 0000000..e2f9b84 --- /dev/null +++ b/hair_service_sd/core/face3d/morphable_model/morphabel_model.py @@ -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 + + diff --git a/hair_service_sd/core/face_enhance/__init__.py b/hair_service_sd/core/face_enhance/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hair_service_sd/core/face_enhance/face_enhancement.py b/hair_service_sd/core/face_enhance/face_enhancement.py new file mode 100644 index 0000000..6714ca7 --- /dev/null +++ b/hair_service_sd/core/face_enhance/face_enhancement.py @@ -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) + + diff --git a/hair_service_sd/core/face_enhance/face_gan_pt.py b/hair_service_sd/core/face_enhance/face_gan_pt.py new file mode 100644 index 0000000..877d029 --- /dev/null +++ b/hair_service_sd/core/face_enhance/face_gan_pt.py @@ -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) diff --git a/hair_service_sd/core/face_enhance/setup.py b/hair_service_sd/core/face_enhance/setup.py new file mode 100644 index 0000000..ca6bf5a --- /dev/null +++ b/hair_service_sd/core/face_enhance/setup.py @@ -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 \ No newline at end of file diff --git a/hair_service_sd/core/face_quality.py b/hair_service_sd/core/face_quality.py new file mode 100644 index 0000000..cd28471 --- /dev/null +++ b/hair_service_sd/core/face_quality.py @@ -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 diff --git a/hair_service_sd/core/faceseg/face_seg.py b/hair_service_sd/core/faceseg/face_seg.py new file mode 100644 index 0000000..dfabced --- /dev/null +++ b/hair_service_sd/core/faceseg/face_seg.py @@ -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() \ No newline at end of file diff --git a/hair_service_sd/core/faceseg/stm.py b/hair_service_sd/core/faceseg/stm.py new file mode 100644 index 0000000..c97834e --- /dev/null +++ b/hair_service_sd/core/faceseg/stm.py @@ -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) \ No newline at end of file diff --git a/hair_service_sd/core/faceseg/tma.py b/hair_service_sd/core/faceseg/tma.py new file mode 100644 index 0000000..db0b674 --- /dev/null +++ b/hair_service_sd/core/faceseg/tma.py @@ -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 diff --git a/hair_service_sd/core/faceseg/u2net.py b/hair_service_sd/core/faceseg/u2net.py new file mode 100644 index 0000000..228aabc --- /dev/null +++ b/hair_service_sd/core/faceseg/u2net.py @@ -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) diff --git a/hair_service_sd/core/hairstyle_local.py b/hair_service_sd/core/hairstyle_local.py new file mode 100644 index 0000000..fa701db --- /dev/null +++ b/hair_service_sd/core/hairstyle_local.py @@ -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 \ No newline at end of file diff --git a/hair_service_sd/core/hairstyle_model.py b/hair_service_sd/core/hairstyle_model.py new file mode 100644 index 0000000..6016491 --- /dev/null +++ b/hair_service_sd/core/hairstyle_model.py @@ -0,0 +1,3992 @@ +import math +import shutil +import time + +from step05_detect_fa_hairmatting_inplace import pkl_process +from PIL import Image, ImageFont, ImageDraw +import cv2 +import numpy as np +import torch +import pickle +import json +import glob +import uuid +from gpt4v_caption import caption_image +from utils.call_hair_train import call_hair_train +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 process_modules import Get_Landmark as Get_Landmark_mtcnn +from core.face_enhance.face_enhancement import FaceEnhancement +from datetime import datetime +import random +from 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 +from process_modules import PersonProcessor_yolov5,KeypointsProcessor,Human_Keypoints,pt_conv_25_to_17 +from gen_super_image import get_high_train_img +from core.MMCVFaceRecognitionServer import MomocvFaceRecognitionServer +from hair_init import HairInit +from multiprocessing import Process, Queue + +resolution_list = [512, 768, 1024, 1280, 1536] +hair_template_material_dir = config.get('default', 'hair_template_material_dir') + +version = config.get('default', 'version') +if version == "local": + train_services = { + 0: "http://192.168.1.57:32678/", + + } +else: + train_services = { + 0: "http://0.0.0.0:32678/", + #1: "http://0.0.0.0:32679/", + #2: "http://0.0.0.0:32680/" + } + +def getM(center, angle, sx, sy): + angle = math.radians(angle) + alpha = math.cos(angle) + beta = math.sin(angle) + M = [[sx * alpha, sx * beta, (1 - sx * alpha) * center[0] - sx * beta * center[1]], + [-sy * beta, sy * alpha, sy * beta * center[0] + (1 - sy * alpha) * center[1]]] + return np.array(M) + +class HairStyle_Model(object): + def __init__(self, gpu, use_enhance=False): + hair_init = HairInit(gpu, use_enhance) + self.use_enhance = use_enhance + self.gpu_index = hair_init.gpu_index + self.get_landmark = hair_init.get_landmark + self.get_landmark_mtcnn = hair_init.get_landmark_mtcnn + self.face_recognition = hair_init.face_recognition + self.hair_size = hair_init.hair_size + self.color_output_size = hair_init.color_output_size + self.process_data = hair_init.process_data + self.generator_hair = hair_init.generator_hair + self.hair_fusion = hair_init.hair_fusion + self.face_enhance = hair_init.face_enhance + self.change_haircolor = hair_init.change_haircolor + self.gender_classify = hair_init.gender_classify + + self.logger_init = hair_init.logger_init + self.logger_process = hair_init.logger_process + self.logger_call = hair_init.logger_call + self.oss2 = hair_init.oss2 + + self.face_seg = hair_init.face_seg + self.chin_cls = hair_init.chin_cls + self.output_img_size = hair_init.output_img_size + self.gender_model = hair_init.gender_model + self.baseColor_dir = hair_init.baseColor_dir + self.effect_prepare_mask_fc32 = hair_init.effect_prepare_mask_fc32 + + self.person_processor = hair_init.person_processor + self.keypoints_processor = hair_init.keypoints_processor + self.human_keypoint = hair_init.human_keypoint + + + # worker_id = int(os.environ.get('APP_WORKER_ID', 1)) + # rand_max = 9527 + # self.gpu_index = (worker_id + rand_max) % train_gpu_nums + # os.environ['CUDA_VISIBLE_DEVICES'] = str(self.gpu_index) + # print('current worker id {} set the gpu id :{}'.format(worker_id, self.gpu_index)) + # device_id = self.gpu_index + # self.get_landmark = Get_Landmark(gpu_id=device_id) + # self.get_landmark_mtcnn = Get_Landmark_mtcnn(gpu_id=device_id) + # self.face_recognition = MomocvFaceRecognitionServer(gpu_id=device_id) + # + # self.hair_size = 768 + # self.color_output_size = 768 + # self.process_data = Process_Data(gpu, device_id) + # self.generator_hair = Generator_Hair(gpu, device_id) + # self.hair_fusion = Generator_Fusion_Res(gpu, device_id) + # self.use_enhance = use_enhance + # if self.use_enhance: + # self.face_enhance = FaceEnhancement(512, device_id) + # self.change_haircolor = Change_Hair_Color(gpu, device_id) + # self.logger_init = LogFactory.getLogger("init") + # self.logger_process = LogFactory.getLogger("process") + # self.logger_call = LogFactory.getLogger("call") + # self.oss2 = OSS_object() + # self.face_seg = FaceSeg(device_id) + # self.chin_cls = chinClass(device_id) + # model_path = "./weights/gender_models" + # self.output_img_size = 128 + # if not os.path.exists(model_path): + # print("GenderClassifyProcessor don't have model!") + # self.gender_model = GenderClassifyProcessor(gpu_id=device_id) + # print("Load model finish ... ") + # self.baseColor_dir = os.path.join(config.get('default', "haircolorDir"), config.get('default', "baseColor_ID")) + # + # for i in range(2): + # image = cv2.imread('data/front.jpg') + # with torch.no_grad(): + # self.infer_hairstyle(image, 'data/template', config.get('default', "tmp_dir"), 'test.jpg') + # + # self.effect_prepare_mask_fc32 = cv2.imread("./data/mask.png").astype(np.float32) / 255 + # # body process + # self.person_processor = PersonProcessor_yolov5(gpu_id=device_id) + # self.keypoints_processor = KeypointsProcessor(gpu_id=device_id) + # self.human_keypoint = Human_Keypoints(gpu=True, device_id=device_id) + # print("Load model finish ... ") + self.train_gpu_sq = Queue() + self.train_gpu_sq.put(0) + + def detect_same_person(self, img1, img2): + with torch.no_grad(): + landmarks_origin_img_1k_1 = self.get_landmark_mtcnn.forward(img1) + landmarks_origin_img_1k_2 = self.get_landmark_mtcnn.forward(img2) + if landmarks_origin_img_1k_1 is None or landmarks_origin_img_1k_2 is None: + return 10001 + features_1 = self.face_recognition.forward([img1], [landmarks_origin_img_1k_1]) + features_2 = self.face_recognition.forward([img2], [landmarks_origin_img_1k_2]) + feat_score = self.face_recognition.cos_sim(features_1[0], features_2[0]) + print('feat_score:',feat_score) + if feat_score < 0.37: + return 10002 + else: + return 10000 + + def judge_hair_can_be_used(self, hair_img): + bounding_boxes, landmarks = self.get_landmark_mtcnn.face_detector.forward(hair_img, min_face_size=50) + if len(bounding_boxes) != 1: + return 10001, 'hair image has no face or has more than one face' + box_index = self.get_landmark.get_max_rect(bounding_boxes) + pts5 = landmarks[box_index] + landmarks1k = self.get_landmark.face_alignmenter_1k.detect_according_5pts(hair_img, pts5) + + pts137 = landmark_processor.pts_1k_to_137(landmarks1k) + movie_params = self.get_landmark.model_3d.detect([hair_img], [pts137])[0] + pitch, yaw, roll = movie_params[1:4] + print("pitch: ", pitch, " yaw: ", yaw, " roll: ", roll) + if abs(pitch) > 0.5 or abs(yaw) > 0.5: + return 10002, 'hair image pitch or yaw is too large' + return 10000, 'hair image can be used' + # return + + + def infer_hairstyle_test(self, origin_img, hairstyle_dir, userinfo_dir): + config_path = os.path.join(hairstyle_dir, "config.json") + if not os.path.exists(config_path): + return None, None, 10001 + config_info = json.load(open(config_path, "rb")) + gender = config_info["gender"] + ratio = int(config_info["ratio"]) + # print("gender: ", gender, " ratio: ", ratio) + + user_rgb_8uc3_orisize = origin_img.copy() + + landmark1k_dir = osp.join(userinfo_dir, 'kpt_1k.txt') + if not osp.exists(landmark1k_dir): + landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward(origin_img) + if landmarks_origin_img_1k is None: + return None, None, 10001 + np.savetxt(landmark1k_dir, landmarks_origin_img_1k) + else: + landmarks_origin_img_1k = np.loadtxt(landmark1k_dir) + + user_bald_res_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_res_ori.png') + # res_matting_8uc3_bald_orisize_dir = osp.join(userinfo_dir, 'res_matting_mask_ori.png') + user_baldseg_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_seg_ori.png') + user_baldseg_8uc3_768_dir = osp.join(userinfo_dir, 'user_baldseg_768.png') + user_bald_8uc3_768_dir = osp.join(userinfo_dir, 'bald_seg_768.png') + user_landmark_f1k2_768_dir = osp.join(userinfo_dir, 'landmark_f1k2_768.txt') + user_hairstyle_M_dir = osp.join(userinfo_dir, 'hairstyle_M.txt') + + pre_list = [user_bald_res_8uc3_orisize_dir, user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_768_dir, user_bald_8uc3_768_dir, + user_landmark_f1k2_768_dir, user_hairstyle_M_dir] + + condition_exist = True + for tmp_file in pre_list: + if not osp.exists(tmp_file): + condition_exist = False + if not condition_exist: + user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, user_baldseg_8uc3_768, user_bald_8uc3_768, \ + user_landmark_f1k2_768, user_hairstyle_M,_ = self.process_data.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio) + cv2.imwrite(user_bald_res_8uc3_orisize_dir, user_bald_res_8uc3_orisize) + cv2.imwrite(user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_orisize) + cv2.imwrite(user_baldseg_8uc3_768_dir, user_baldseg_8uc3_768) + cv2.imwrite(user_bald_8uc3_768_dir, user_bald_8uc3_768) + np.savetxt(user_landmark_f1k2_768_dir, user_landmark_f1k2_768) + np.savetxt(user_hairstyle_M_dir, user_hairstyle_M) + else: + user_bald_res_8uc3_orisize = cv2.imread(user_bald_res_8uc3_orisize_dir) + user_baldseg_8uc3_orisize = cv2.imread(user_baldseg_8uc3_orisize_dir) + user_baldseg_8uc3_768 = cv2.imread(user_baldseg_8uc3_768_dir) + user_bald_8uc3_768 = cv2.imread(user_bald_8uc3_768_dir) + user_landmark_f1k2_768 = np.loadtxt(user_landmark_f1k2_768_dir) + user_hairstyle_M = np.loadtxt(user_hairstyle_M_dir) + + # show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("show_concat", show_concat) + # cv2.waitKey() + + another_pose_hair_img_path = os.path.join(hairstyle_dir, "input_another_pose_hair_image.npy") + if not os.path.exists(another_pose_hair_img_path): + return None, None, 10002 + another_pose_hair_image = np.load(another_pose_hair_img_path) + # cv2.imshow('user_baldseg_8uc3_768', user_baldseg_8uc3_768) + # cv2.imshow('user_bald_8uc3_768', user_bald_8uc3_768) + + t0 = time.time() + # 换发型 + + hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image, + user_baldseg_8uc3_768, + user_bald_8uc3_768, + user_landmark_f1k2_768, gender) + self.logger_process.info('Generator_Hair_inference_use_pref costs:{}'.format(time.time() - t0)) + # print('Generator_Hair_inference_use_pref costs:', time.time() - t0) + # cv2.imshow('hair_gene_8uc3_768', hair_gene_8uc3_768) + # cv2.imshow('user_baldseg_8uc3_768', user_baldseg_8uc3_768) + # cv2.imshow('user_bald_8uc3_768', user_bald_8uc3_768) + # cv2.waitKey() + t1 = time.time() + + hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste( + user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M) + + # res_hairstyle_before_8uc3 = origin_img.copy() + # hair_fix_img_mask_8uc4 = np.concatenate((res_hairstyle_before_8uc3, hair_gene_matte_8uc3_orisize[:, :, :1]), axis=2) + # cv2.imwrite(matting_before_path, hair_fix_img_mask_8uc4) + + + res_matting_mask_ori = osp.join(userinfo_dir, 'res_matting_mask_ori_fix.png') + res_matting_mask_ori_raw = osp.join(userinfo_dir, 'res_matting_mask_ori_raw.png') + if osp.exists(res_matting_mask_ori): + os.remove(res_matting_mask_ori) + if osp.exists(res_matting_mask_ori_raw): + os.remove(res_matting_mask_ori_raw) + # cv2.imwrite(res_matting_mask_ori, hair_gene_matte_8uc3_orisize) + cv2.imwrite(res_matting_mask_ori_raw, hair_gene_matte_8uc3_orisize) + # cv2.imshow('hair_gene_fusion_8uc3_orisize', hair_gene_fusion_8uc3_orisize) + # cv2.imshow('hair_gene_matte_8uc3_orisize', hair_gene_matte_8uc3_orisize) + # cv2.waitKey() + # print('get_fusion_res_hairpaste costs:', time.time() - t1) + self.logger_process.info('get_fusion_res_hairpaste costs:{}'.format(time.time() - t1)) + + user_res_8uc3_orisize = self.hair_fusion.inference(hair_gene_fusion_8uc3_orisize, + hair_gene_matte_8uc3_orisize, + landmarks_origin_img_1k, + user_baldseg_8uc3_orisize, + ratio) + # show_concat = np.concatenate((hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize, user_baldseg_8uc3_orisize, user_res_8uc3_orisize), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_concat", show_concat) + # cv2.waitKey() + t2 = time.time() + + user_res_8uc3_orisize_for_haircolor = user_res_8uc3_orisize.copy() + + # haircolr_res_dir = osp.join(userinfo_dir, 'user_res_8uc3_orisize_for_haircolor.png') + # cv2.imwrite(haircolr_res_dir, user_res_8uc3_orisize_for_haircolor) + + hair_gene_fusion_8uc3_orisize_LAB = cv2.cvtColor(hair_gene_fusion_8uc3_orisize, cv2.COLOR_BGR2LAB) + user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) + mix_ratio = 0.2 + user_res_8uc3_orisize_LAB[:, :, 0] = hair_gene_fusion_8uc3_orisize_LAB[:, :, 0] * mix_ratio + user_res_8uc3_orisize_LAB[:, :, 0] * (1 - mix_ratio) + user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_LAB, cv2.COLOR_LAB2BGR) + + + # mid_show = np.concatenate((hair_gene_matte_fc32_orisize, hair_gene_matte_fc32_orisize_erode, hair_gene_matte_fc32_orisize_circle), axis=1) + # ratio = 1536. / max(mid_show.shape[:2]) + # mid_show = cv2.resize(mid_show, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_show", mid_show) + # cv2.imshow("user_res_8uc3_orisize", user_res_8uc3_orisize) + # cv2.waitKey() + + if self.use_enhance: + user_res_8uc3_orisize_enhance = self.face_enhance.process(user_res_8uc3_orisize, landmarks_origin_img_1k) + # user_res_8uc3_orisize = user_bald_res_8uc3_orisize.astype(np.float32)/255. * (1. - hair_gene_matte_8uc3_orisize.astype(np.float32) / 255.) + user_res_8uc3_orisize_enhance.astype(np.float32)/255. * hair_gene_matte_8uc3_orisize.astype(np.float32) / 255. + # user_res_8uc3_orisize = (user_res_8uc3_orisize*255).astype(np.uint8) + self.logger_process.info('use_enhance costs:{}'.format(time.time() - t2)) + + res_fix_img_mask_8uc4_orisize = np.concatenate((user_res_8uc3_orisize, hair_gene_matte_8uc3_orisize[:, :, :1]), axis=2) + cv2.imwrite(res_matting_mask_ori, res_fix_img_mask_8uc4_orisize) + + return user_res_8uc3_orisize_enhance, user_res_8uc3_orisize_for_haircolor, 0 + + + def infer_hairstyle(self, origin_img, hairstyle_dir, userinfo_dir, mask_newname): + config_path = os.path.join(hairstyle_dir, "config.json") + if not os.path.exists(config_path): + return None, None, 10001 + config_info = json.load(open(config_path, "rb")) + gender = config_info["gender"] + ratio = int(config_info["ratio"]) + # print("gender: ", gender, " ratio: ", ratio) + + user_rgb_8uc3_orisize = origin_img.copy() + tt = time.time() + landmark1k_dir = osp.join(userinfo_dir, 'kpt_1k.txt') + if not osp.exists(landmark1k_dir): + with torch.no_grad(): + landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward(origin_img) + if landmarks_origin_img_1k is None: + return None, None, 10001 + np.savetxt(landmark1k_dir, landmarks_origin_img_1k) + else: + landmarks_origin_img_1k = np.loadtxt(landmark1k_dir) + + user_bald_res_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_res_ori_raw.png') + # res_matting_8uc3_bald_orisize_dir = osp.join(userinfo_dir, 'res_matting_mask_ori.png') + user_baldseg_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_seg_ori.png') + user_baldseg_8uc3_768_dir = osp.join(userinfo_dir, 'user_baldseg_768.png') + user_bald_8uc3_768_dir = osp.join(userinfo_dir, 'bald_seg_768.png') + user_landmark_f1k2_768_dir = osp.join(userinfo_dir, 'landmark_f1k2_768.txt') + user_hairstyle_M_dir = osp.join(userinfo_dir, 'hairstyle_M.txt') + + pre_list = [user_bald_res_8uc3_orisize_dir, user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_768_dir, user_bald_8uc3_768_dir, + user_landmark_f1k2_768_dir, user_hairstyle_M_dir] + + condition_exist = True + for tmp_file in pre_list: + if not osp.exists(tmp_file): + condition_exist = False + if not condition_exist: + user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, user_baldseg_8uc3_768, user_bald_8uc3_768, \ + user_landmark_f1k2_768, user_hairstyle_M,_ = self.process_data.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio) + cv2.imwrite(user_bald_res_8uc3_orisize_dir, user_bald_res_8uc3_orisize) + cv2.imwrite(user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_orisize) + cv2.imwrite(user_baldseg_8uc3_768_dir, user_baldseg_8uc3_768) + cv2.imwrite(user_bald_8uc3_768_dir, user_bald_8uc3_768) + np.savetxt(user_landmark_f1k2_768_dir, user_landmark_f1k2_768) + np.savetxt(user_hairstyle_M_dir, user_hairstyle_M) + else: + user_bald_res_8uc3_orisize = cv2.imread(user_bald_res_8uc3_orisize_dir) + user_baldseg_8uc3_orisize = cv2.imread(user_baldseg_8uc3_orisize_dir) + # user_baldseg_8uc3_768 = cv2.imread(user_baldseg_8uc3_768_dir) + # user_bald_8uc3_768 = cv2.imread(user_bald_8uc3_768_dir) + # user_landmark_f1k2_768 = np.loadtxt(user_landmark_f1k2_768_dir) + # user_hairstyle_M = np.loadtxt(user_hairstyle_M_dir) + if ratio == 0: + user_hairstyle_M = self.process_data.get_hair_M_boy_v1(landmarks_origin_img_1k) + elif ratio == 1: + user_hairstyle_M = self.process_data.get_hair_M_girl_v1(landmarks_origin_img_1k) + elif ratio == 2: + user_hairstyle_M = self.process_data.get_hair_M_girl_v2(landmarks_origin_img_1k) + else: + user_hairstyle_M = self.process_data.get_hair_M_girl_v1(landmarks_origin_img_1k) + user_landmark_f1k2_768 = landmark_processor.transform_points(landmarks_origin_img_1k, user_hairstyle_M) + user_bald_8uc3_768 = cv2.warpAffine(user_bald_res_8uc3_orisize, user_hairstyle_M, + (self.hair_size, self.hair_size)) + user_baldseg_8uc3_768 = cv2.warpAffine(user_baldseg_8uc3_orisize, user_hairstyle_M, + (self.hair_size, self.hair_size)) + print('condition cosst:', time.time() - tt ) + # show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("show_concat", show_concat) + # cv2.waitKey() + + another_pose_hair_img_path = os.path.join(hairstyle_dir, "input_another_pose_hair_image.npy") + if not os.path.exists(another_pose_hair_img_path): + return None, None, 10002 + another_pose_hair_image = np.load(another_pose_hair_img_path) + # cv2.imshow('user_baldseg_8uc3_768', user_baldseg_8uc3_768) + # cv2.imshow('user_bald_8uc3_768', user_bald_8uc3_768) + + t0 = time.time() + # 换发型 + + hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image, + user_baldseg_8uc3_768, + user_bald_8uc3_768, + user_landmark_f1k2_768, gender) + self.logger_process.info('Generator_Hair_inference_use_pref costs:{}'.format(time.time() - t0)) + + t1 = time.time() + + hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste( + user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M) + + + res_matting_mask_ori = osp.join(userinfo_dir, mask_newname) + res_matting_mask_ori_raw = osp.join(userinfo_dir, 'res_matting_mask_ori_raw.png') + if osp.exists(res_matting_mask_ori): + os.remove(res_matting_mask_ori) + if osp.exists(res_matting_mask_ori_raw): + os.remove(res_matting_mask_ori_raw) + + self.logger_process.info('get_fusion_res_hairpaste costs:{}'.format(time.time() - t1)) + + user_res_8uc3_orisize = self.hair_fusion.inference(hair_gene_fusion_8uc3_orisize, + hair_gene_matte_8uc3_orisize, + landmarks_origin_img_1k, + user_baldseg_8uc3_orisize, + ratio) + hair_gene_matte_8uc3_orisize_cp = hair_gene_matte_8uc3_orisize.copy() + # show_concat = np.concatenate((hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize, user_baldseg_8uc3_orisize, user_res_8uc3_orisize), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_concat", show_concat) + # cv2.waitKey() + t2 = time.time() + + user_res_8uc3_orisize_for_haircolor = user_res_8uc3_orisize.copy() + + # if self.use_enhance: + user_res_8uc3_orisize_enhance = self.face_enhance.process(user_res_8uc3_orisize, landmarks_origin_img_1k) + t3 = time.time() + self.logger_process.info('use_enhance costs:{}'.format(t3 - t2)) + + user_bald_res_8uc3_orisize_enhance, user_baldseg_8uc3_orisize_enhance, user_baldseg_8uc3_768_enhance, user_bald_8uc3_768_enhance, \ + user_landmark_f1k2_768_enhance, user_hairstyle_M_enhance, user_matting_8uc3_bald_orisize = self.process_data.get_prepare_user_768_data( + user_res_8uc3_orisize_enhance, landmarks_origin_img_1k, ratio=ratio) + + t4 = time.time() + self.logger_process.info('gen bald costs:{}'.format(t4 - t3)) + + # 重新提取matting + _, hair_gene_matte_8uC0_orisize, _ = self.process_data.generator_matte.matte_inference(user_res_8uc3_orisize_enhance, landmarks_origin_img_1k) + t5 = time.time() + self.logger_process.info('gen matte_inference costs:{}'.format(t5 - t4)) + hair_gene_matte_8uc3_orisize = np.repeat(hair_gene_matte_8uC0_orisize[:, :, np.newaxis], 3, axis=2) + hair_gene_matte_8uc3_orisize = cv2.blur(hair_gene_matte_8uc3_orisize, (3, 3)) + + # user_res_8uc3_orisize = user_bald_res_8uc3_orisize_enhance.astype(np.float32)/255. * (1. - hair_gene_matte_8uc3_orisize.astype(np.float32) / 255.) + user_res_8uc3_orisize_enhance.astype(np.float32)/255. * hair_gene_matte_8uc3_orisize.astype(np.float32) / 255. + # user_res_8uc3_orisize = (user_res_8uc3_orisize*255).astype(np.uint8) + + user_res_fc32_orisize_enhance = user_res_8uc3_orisize_enhance.astype(np.float32)/255. + user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32)/255. + hair_gene_matte_fc32_orisize = hair_gene_matte_8uc3_orisize.astype(np.float32)/255. + hair_gene_matte_fc32_orisize = cv2.GaussianBlur(hair_gene_matte_fc32_orisize, (11, 11), 0, 0) + user_res_fc32_orisize_enhance = user_res_fc32_orisize_enhance * hair_gene_matte_fc32_orisize + user_res_fc32_orisize * (1 - hair_gene_matte_fc32_orisize) + user_res_8uc3_orisize_enhance = (user_res_fc32_orisize_enhance * 255).astype(np.uint8) + + # mid_show = np.concatenate((user_res_fc32_orisize_enhance, user_res_8uc3_orisize_enhance, hair_gene_matte_8uc3_orisize, user_res_8uc3_orisize), axis=1) + # ratio = 1536. / max(mid_show.shape[:2]) + # mid_show = cv2.resize(mid_show, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_show", mid_show) + # cv2.imshow("user_res_8uc3_orisize_enhance_fg", user_res_8uc3_orisize_enhance.astype(np.float32)/255. * hair_gene_matte_8uc3_orisize.astype(np.float32) / 255) + # cv2.waitKey() + + user_bald_res_8uc3_orisize_for_fusion_dir = osp.join(userinfo_dir, 'bald_res_ori.png') + # if osp.exists(user_bald_res_8uc3_orisize_for_fusion_dir): + # os.remove(user_bald_res_8uc3_orisize_for_fusion_dir) + cv2.imwrite(user_bald_res_8uc3_orisize_for_fusion_dir, user_bald_res_8uc3_orisize_enhance) + # if osp.exists(res_matting_mask_ori_raw): + # os.remove(res_matting_mask_ori_raw) + cv2.imwrite(res_matting_mask_ori_raw, hair_gene_matte_8uc3_orisize) + + # user_bald_res_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_res_ori.png') + # cv2.imwrite(user_bald_res_8uc3_orisize_dir, user_res_8uc3_orisize) + + + res_fix_img_mask_8uc4_orisize = np.concatenate((user_res_8uc3_orisize_enhance, hair_gene_matte_8uc3_orisize_cp[:, :, :1]), axis=2) + # cv2.imshow('') + cv2.imwrite(res_matting_mask_ori, res_fix_img_mask_8uc4_orisize) + print('costs:', time.time() - t5) + return user_res_8uc3_orisize_enhance, user_res_8uc3_orisize_for_haircolor, 0 + + + def infer_hairstyle_v2(self, origin_img, hairstyle_dir, return_pt1k=False,use_enhance=False): + config_path = os.path.join(hairstyle_dir, "config.json") + if not os.path.exists(config_path): + return None, 10001 + config_info = json.load(open(config_path, "rb")) + gender = config_info["gender"] + ratio = int(config_info["ratio"]) + if ratio == 2: + ratio = 3 + print("gender: ", gender, " ratio: ", ratio) + + user_rgb_8uc3_orisize = origin_img + landmarks_origin_img_1k = self.get_landmark.forward_v2(origin_img) + + # 根据人脸关键点画出人脸区域的mask + hull_mask = self.draw_hull_mask(landmarks_origin_img_1k.astype(np.int32), w=origin_img.shape[1], h=origin_img.shape[0]) + user_mask_save_path = os.path.join(hairstyle_dir, "hull_mask.png") + cv2.imwrite(user_mask_save_path, hull_mask * 255) + + # cv2.imshow('hull_mask', hull_mask) + # cv2.waitKey(0) + + if landmarks_origin_img_1k is None: + return None, 10001 + + user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, \ + user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M, user_matting_8uc3_bald_orisize = self.process_data.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio) + + # show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("show_concat", show_concat) + # cv2.waitKey() + + # cv2.imshow("user_matting_8uc3_bald_orisize", user_matting_8uc3_bald_orisize) + # cv2.waitKey(0) + + user_orig_mask_path = os.path.join(hairstyle_dir, "user_orig_mask.png") + if not os.path.exists(user_orig_mask_path): + cv2.imwrite(user_orig_mask_path, user_matting_8uc3_bald_orisize) + + another_pose_hair_img_path = os.path.join(hairstyle_dir, "input_another_pose_hair_image.npy") + if not os.path.exists(another_pose_hair_img_path): + return None, 10002 + another_pose_hair_image = np.load(another_pose_hair_img_path) + + # 换发型 + hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image, + user_baldseg_8uc3_768, + user_bald_8uc3_768, + user_landmark_f1k2_768, gender) + + hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste( + user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M) + + # cv2.imshow("hair_gene_fusion_8uc3_orisize:", hair_gene_fusion_8uc3_orisize) + # cv2.imshow("hair_gene_matte_8uc3_orisize:", hair_gene_matte_8uc3_orisize) + # cv2.waitKey(0) + + gen_hair_mask_path = os.path.join(hairstyle_dir, "hair_mask.png") + cv2.imwrite(gen_hair_mask_path, hair_gene_matte_8uc3_orisize * 255) + + gen_hair_mask_path_2 = os.path.join(hairstyle_dir, "hair_mask_2.png") + cv2.imwrite(gen_hair_mask_path_2, hair_gene_matte_8uc3_orisize) + + if use_enhance: + user_res_8uc3_orisize = self.hair_fusion.inference(hair_gene_fusion_8uc3_orisize, + hair_gene_matte_8uc3_orisize, + landmarks_origin_img_1k, + user_baldseg_8uc3_orisize, + ratio) + hair_gene_fusion_8uc3_orisize_LAB = cv2.cvtColor(hair_gene_fusion_8uc3_orisize, cv2.COLOR_BGR2LAB) + user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) + mix_ratio = 0.2 + user_res_8uc3_orisize_LAB[:, :, 0] = hair_gene_fusion_8uc3_orisize_LAB[:, :, 0] * mix_ratio + user_res_8uc3_orisize_LAB[:, :, 0] * (1 - mix_ratio) + user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_LAB, cv2.COLOR_LAB2BGR) + + # show_concat = np.concatenate((hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize, + # user_baldseg_8uc3_orisize, user_res_8uc3_orisize), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_concat", show_concat) + # cv2.waitKey() + else: + user_res_8uc3_orisize = hair_gene_fusion_8uc3_orisize + + + # face_bbox = cv2.boundingRect(landmarks_origin_img_1k[np.newaxis, :, :]) + # face_max_len = max(face_bbox[2], face_bbox[3]) + # erode_kernel_size = int(face_max_len * 0.11) + # if erode_kernel_size % 2 == 0: + # erode_kernel_size += 1 + # blur_kernel_size = int(face_max_len * 0.05) + # if blur_kernel_size % 2 == 0: + # blur_kernel_size += 1 + # # print("blur_kernel_size: ", blur_kernels + # hair_gene_matte_8uc3_hard = hair_gene_matte_8uc3_orisize.copy() + # hair_gene_matte_8uc3_hard[hair_gene_matte_8uc3_orisize[:, :, 0] > 0] = 255 + # hair_gene_matte_8uc3_orisize_erode = cv2.erode(hair_gene_matte_8uc3_hard, np.ones((erode_kernel_size, erode_kernel_size), np.uint8), iterations=1) + # hair_gene_matte_fc32_orisize_erode = hair_gene_matte_8uc3_orisize_erode.astype(np.float32) / 255 + # hair_gene_matte_fc32_orisize = hair_gene_matte_8uc3_orisize.astype(np.float32) / 255 + # hair_gene_matte_fc32_orisize_circle = hair_gene_matte_fc32_orisize * (1 - hair_gene_matte_fc32_orisize_erode) + # hair_gene_matte_fc32_orisize_circle = cv2.GaussianBlur(hair_gene_matte_fc32_orisize_circle, (blur_kernel_size, blur_kernel_size), 0, 0) + # + # hair_gene_fusion_8uc3_orisize_LAB = cv2.cvtColor(hair_gene_fusion_8uc3_orisize, cv2.COLOR_BGR2LAB) + # user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) + # mix_ratio = 0.5 + # user_res_8uc3_orisize_LAB_new = user_res_8uc3_orisize_LAB.copy() + # user_res_8uc3_orisize_LAB_new[:, :, 0] = hair_gene_fusion_8uc3_orisize_LAB[:, :, 0] * mix_ratio + user_res_8uc3_orisize_LAB[:, :, 0] * (1 - mix_ratio) + # user_res_8uc3_orisize_LAB[:, :, 0] = user_res_8uc3_orisize_LAB_new[:, :, 0] * hair_gene_matte_fc32_orisize_circle[:, :, 0] + user_res_8uc3_orisize_LAB[:, :, 0] * (1 - hair_gene_matte_fc32_orisize_circle[:, :, 0]) + # user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_LAB, cv2.COLOR_LAB2BGR) + + # mid_show = np.concatenate((hair_gene_matte_fc32_orisize, hair_gene_matte_fc32_orisize_erode, hair_gene_matte_fc32_orisize_circle), axis=1) + # ratio = 1536. / max(mid_show.shape[:2]) + # mid_show = cv2.resize(mid_show, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_show", mid_show) + # # cv2.imshow("user_res_8uc3_orisize", user_res_8uc3_orisize) + # cv2.waitKey() + + if use_enhance: + user_res_8uc3_orisize_enhance = self.face_enhance.process(user_res_8uc3_orisize, landmarks_origin_img_1k) + user_res_8uc3_orisize = user_bald_res_8uc3_orisize.astype(np.float32) / 255. * ( + 1. - hair_gene_matte_8uc3_orisize.astype( + np.float32) / 255.) + user_res_8uc3_orisize_enhance.astype( + np.float32) / 255. * hair_gene_matte_8uc3_orisize.astype(np.float32) / 255. + user_res_8uc3_orisize = (user_res_8uc3_orisize * 255).astype(np.uint8) + + # for pt in landmarks_origin_img_1k.astype(np.int32): + # cv2.circle(origin_img, (pt[0], pt[1]), 1, (0, 255, 0), -1) + # show_concat_orisize = np.concatenate((origin_img, user_res_8uc3_orisize), axis=1) + # ratio = 1024. / max(show_concat_orisize.shape[:2]) + # show_concat_orisize = cv2.resize(show_concat_orisize, (0, 0), fx=ratio, fy=ratio) + # ratio = 1024. / max(show_concat_768.shape[:2]) + # show_concat_768 = cv2.resize(show_concat_768, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("show_concat_orisize", show_concat_orisize) + # cv2.imshow("show_concat_768", show_concat_768) + # cv2.waitKey() + + if return_pt1k: + ret_dict = dict(user_res_8uc3_orisize=user_res_8uc3_orisize, + landmarks_origin_img_1k=landmarks_origin_img_1k) + + return ret_dict, 0 + else: + return user_res_8uc3_orisize, 0 + + + def draw_hull_mask(self, fc_landmark, w=256, h=256, is_gray=False): + hull_mask = np.zeros((h, w), dtype=np.float32) + if len(fc_landmark) == 1000: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (1)) + else: + raise Exception('landmark should be 1000') + return hull_mask + + + def affine(self, img, a, b, c, d, tx, ty): + H, W, C = img.shape + + # temporary image + tem = img.copy() + img = np.zeros((H + 2, W + 2, C), dtype=np.float32) + img[1:H + 1, 1:W + 1] = tem + + # get new image shape + H_new = np.round(H * d).astype(np.int) + W_new = np.round(W * a).astype(np.int) + out = np.zeros((H_new + 1, W_new + 1, C), dtype=np.float32) + + # get position of new image + x_new = np.tile(np.arange(W_new), (H_new, 1)) + y_new = np.arange(H_new).repeat(W_new).reshape(H_new, -1) + + # get position of original image by affine + adbc = a * d - b * c + x = np.round((d * x_new - b * y_new) / adbc).astype(np.int) - tx + 1 + y = np.round((-c * x_new + a * y_new) / adbc).astype(np.int) - ty + 1 + + # 避免目标图像对应的原图像中的坐标溢出 + x = np.minimum(np.maximum(x, 0), W + 1).astype(np.int) + y = np.minimum(np.maximum(y, 0), H + 1).astype(np.int) + + # assgin pixcel to new image + out[y_new, x_new] = img[y, x] + + out = out[:H_new, :W_new] + out = out.astype(np.uint8) + + return out + + def fix_hairstyle(self, origin_img, hairstyle_dir, is_src, brush, taskid, userinfo_dir): + if is_src: + return None, 10002 + else: + user_rgb_8uc3_orisize = origin_img.copy() + landmark1k_dir = osp.join(userinfo_dir, 'kpt_1k.txt') + if not osp.exists(landmark1k_dir): + landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward(origin_img) + if landmarks_origin_img_1k is None: + return None, 10001 + np.savetxt(landmark1k_dir, landmarks_origin_img_1k) + else: + landmarks_origin_img_1k = np.loadtxt(landmark1k_dir) + [ori_x, ori_y, dst_x, dst_y, degree] = brush + + landmark137 = landmark_processor.pts_1k_to_137(landmarks_origin_img_1k) + face_left = landmark137[16] + face_right = landmark137[6] + face_length = face_right[0] - face_left[0] + radis = face_length/10 + + t0 = time.time() + ret_img = localtranslationwarpfastwithstrength(user_rgb_8uc3_orisize, landmark137, ori_x, ori_y, dst_x, dst_y, degree, radis) + # print('fix hair costs:', time.time() - t0) + cv2.circle(ret_img, (int(ori_x), int(ori_y)), 2, (255, 0, 0)) + cv2.circle(ret_img, (int(dst_x), int(dst_y)), 2, (255, 0, 0)) + + self.logger_process.info('fix hair costs:{}'.format(time.time() - t0)) + + fix_image_dir = osp.join(userinfo_dir, taskid + '.jpg') + cv2.imwrite(fix_image_dir, ret_img) + return fix_image_dir, 0 + + def get_mask_area(self, input_mask): + ret, thresh = cv2.threshold(input_mask, 0, 255, cv2.THRESH_BINARY) + ave = cv2.mean(thresh)[0] / 255 + area = ave * input_mask.shape[0] * input_mask.shape[1] + return area + def check_hair_fintune_out_face(self, hair_gene_matte_8uc3_orisize, landmarks_origin_img_1k, user_baldseg_8uc3_orisize, hair_gene_matte_fc32_orisize_raw, userinfo_dir): + # user_baldseg_8uc3_orisize = cv2.imread(os.path.join(userinfo_dir, "bald_seg_ori.png")) + # start_time = time.time() + # user_head_mask_fc32_orisize = (user_baldseg_8uc3_orisize == [255, 0, 0]).all(axis=2).astype(np.float32)[:, :, np.newaxis] + # user_head_mask_fc32_orisize = np.repeat(user_head_mask_fc32_orisize, 3, axis=2) + # face_kpts1k_mask = np.zeros_like(user_baldseg_8uc3_orisize) + # cv2.fillConvexPoly(face_kpts1k_mask, cv2.convexHull((landmarks_origin_img_1k[0:312]).astype(np.int32)), (255, 255, 255)) + + user_head_mask_orisize_path = osp.join(userinfo_dir, 'user_head_mask_orisize.png') + if os.path.exists(user_head_mask_orisize_path): + user_head_mask_fc32_orisize = cv2.imread(user_head_mask_orisize_path).astype(np.float32) / 255 + else: + user_head_mask_fc32_orisize = (user_baldseg_8uc3_orisize == [255, 0, 0]).all(axis=2).astype(np.float32)[:, :, np.newaxis] + user_head_mask_fc32_orisize = np.repeat(user_head_mask_fc32_orisize, 3, axis=2) + cv2.imwrite(user_head_mask_orisize_path, (user_head_mask_fc32_orisize*255).astype(np.uint8)) + face_kpts1k_mask_path = osp.join(userinfo_dir, 'face_kpts1k_mask.png') + if os.path.exists(face_kpts1k_mask_path): + face_kpts1k_mask = cv2.imread(face_kpts1k_mask_path) + else: + face_kpts1k_mask = np.zeros_like(user_baldseg_8uc3_orisize) + cv2.fillConvexPoly(face_kpts1k_mask, cv2.convexHull((landmarks_origin_img_1k[0:312]).astype(np.int32)), (255, 255, 255)) + cv2.imwrite(face_kpts1k_mask_path, face_kpts1k_mask) + + # print("check_hair_fintune_out_face prepare two mask cost: ", time.time() - start_time) + # hair_gene_matte_fc32_orisize_raw = cv2.imread(os.path.join(userinfo_dir, "res_matting_mask_ori_raw.png")).astype(np.float32) / 255 + hair_gene_matte_fc32_orisize = hair_gene_matte_8uc3_orisize.astype(np.float32) / 255 + face_kpts1k_mask_fc32 = face_kpts1k_mask.astype(np.float32) / 255 + + face_overlay_hair_mask_raw = face_kpts1k_mask * hair_gene_matte_fc32_orisize_raw + face_overlay_hair_mask = face_kpts1k_mask * hair_gene_matte_fc32_orisize + face_overlay_hair_mask_raw_area = self.get_mask_area(face_overlay_hair_mask_raw) + face_overlay_hair_mask_area = self.get_mask_area(face_overlay_hair_mask) + # print("face mask: ", face_overlay_hair_mask_raw_area, " ", face_overlay_hair_mask_area) + if face_overlay_hair_mask_raw_area > 0: + if face_overlay_hair_mask_area / face_overlay_hair_mask_raw_area < 0.75: + # print("face overlay mask not allow fintune") + return True + + bald_mask = user_head_mask_fc32_orisize * (1 - face_kpts1k_mask_fc32) + bald_overlay_hair_mask_raw = bald_mask * hair_gene_matte_fc32_orisize_raw + bald_overlay_hair_mask = bald_mask * hair_gene_matte_fc32_orisize + bald_overlay_hair_mask_raw_area = self.get_mask_area(bald_overlay_hair_mask_raw) + bald_overlay_hair_mask_area = self.get_mask_area(bald_overlay_hair_mask) + # print("hair mask: ", bald_overlay_hair_mask_raw_area, " ", bald_overlay_hair_mask_area) + if bald_overlay_hair_mask_raw_area > 0: + if bald_overlay_hair_mask_area / bald_overlay_hair_mask_raw_area < 0.9: + # print("bald mask not allow fintune") + return True + if bald_overlay_hair_mask_raw_area - bald_overlay_hair_mask_area > 200: + # print("bald mask not allow fintune") + return True + + hair_gene_matte_fc32_orisize_raw_area = self.get_mask_area(hair_gene_matte_fc32_orisize_raw) + hair_gene_matte_fc32_orisize_area = self.get_mask_area(hair_gene_matte_fc32_orisize) + if hair_gene_matte_fc32_orisize_raw_area > 0 and hair_gene_matte_fc32_orisize_area / hair_gene_matte_fc32_orisize_raw_area < 0.8: + # print("hair mask not allow fintune") + return True + return False + + def fix_hairstyle_v2(self, origin_img, brush, taskid, userinfo_dir, maskImgName): + t0 = time.time() + landmark1k_path = osp.join(userinfo_dir, 'kpt_1k.txt') + if os.path.exists(landmark1k_path): + landmarks_origin_img_1k = np.loadtxt(landmark1k_path) + else: + landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward(origin_img) + + if landmarks_origin_img_1k is None: + return None, 10001 + if maskImgName is None or maskImgName == "": + maskImgName = '{}_fix.png'.format(taskid) + matting_before_path = osp.join(userinfo_dir, maskImgName) + # res_hairstyle_before_path = osp.join(userinfo_dir, "user_res_8uc3_orisize_for_finetune.png") + bald_res_ori_path = osp.join(userinfo_dir, 'bald_res_ori.png') + bald_seg_ori_path = osp.join(userinfo_dir, 'bald_seg_ori.png') + matting_raw_path = osp.join(userinfo_dir, 'res_matting_mask_ori_raw.png') + + if os.path.exists(matting_before_path) and osp.exists(matting_raw_path): + hair_fix_img_mask_8uc4 = cv2.imread(matting_before_path, -1) + res_hairstyle_before_8uc3 = hair_fix_img_mask_8uc4[:, :, :3] + hair_matting_mask_8uc3 = np.repeat(hair_fix_img_mask_8uc4[:, :, 3:], 3, axis=2) + hair_gene_matte_fc32_orisize_raw = cv2.imread(matting_raw_path).astype(np.float32) / 255 + else: + _, user_matting_8uc1_bald_orisize, _ = self.process_data.generator_matte.matte_inference(origin_img, landmarks_origin_img_1k) + hair_matting_mask_8uc3 = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) + if not os.path.exists(userinfo_dir): + os.makedirs(userinfo_dir) + hair_gene_matte_fc32_orisize_raw = hair_matting_mask_8uc3.copy() + cv2.imwrite(matting_raw_path, hair_matting_mask_8uc3) + res_hairstyle_before_8uc3 = origin_img.copy() + hair_fix_img_mask_8uc4 = np.concatenate((res_hairstyle_before_8uc3, hair_matting_mask_8uc3[:, :, :1]), axis=2) + cv2.imwrite(matting_before_path, hair_fix_img_mask_8uc4) + + t1 = time.time() + print('hairfix v2, matte_inference costs', t1 - t0) + if osp.exists(bald_res_ori_path) and osp.exists(bald_seg_ori_path): + bald_res_ori_fc32 = cv2.imread(bald_res_ori_path).astype(np.float32)/255 + user_baldseg_8uc3_orisize = cv2.imread(bald_seg_ori_path).astype(np.float32)/255 + + else: + bald_res_ori_8uc3, user_baldseg_8uc3_orisize, _, _, _, _,_ = self.process_data.get_prepare_user_768_data(origin_img, landmarks_origin_img_1k, ratio=1) + bald_res_ori_fc32 = bald_res_ori_8uc3.astype(np.float32)/255 + # bald_res_ori_fc32 = self.process_data.generator_baldseg.forward(origin_img, user_matting_8uc1_bald_orisize, + # landmarks_origin_img_1k) + cv2.imwrite(bald_res_ori_path, bald_res_ori_8uc3) + cv2.imwrite(bald_seg_ori_path, user_baldseg_8uc3_orisize) + + t2 = time.time() + print('hairfix v2, bald costs', t2 - t1) + + # cv2.imshow('matting_before', hair_matting_mask_8uc3) + # cv2.imshow('bald_res_ori_fc32', bald_res_ori_fc32) + # cv2.waitKey() + check_res = self.check_hair_fintune_out_face(hair_matting_mask_8uc3, landmarks_origin_img_1k, user_baldseg_8uc3_orisize,hair_gene_matte_fc32_orisize_raw, userinfo_dir) + if check_res: + # print("hair out of face!!!") + fix_image_dir = osp.join(userinfo_dir, taskid + '.jpg') + cv2.imwrite(fix_image_dir, origin_img) + return fix_image_dir, 0 + + t3 = time.time() + print('hairfix v2, check_res costs', t3 - t2) + + # landmark137 = landmark_processor.pts_1k_to_137(landmarks_origin_img_1k) + user_rgb_8uc3_orisize = res_hairstyle_before_8uc3.copy() + [ori_x, ori_y, dst_x, dst_y, radis] = brush + t0 = time.time() + ret_img = localtranslationwarpfastwithstrength_v2(user_rgb_8uc3_orisize.copy(), ori_x, ori_y, dst_x, dst_y, radis) + # ret_img = user_rgb_8uc3_orisize.copy() + ret_matting_mask = localtranslationwarpfastwithstrength_v2(hair_matting_mask_8uc3.copy(), ori_x, ori_y, dst_x, dst_y, radis) + # ret_matting_mask = cv2.GaussianBlur(ret_matting_mask, (15, 15), 0, 0) + t4 = time.time() + print('hairfix v2, localtranslationwarpfastwithstrength_v2 costs', t4 - t0) + + ret_matting_mask_fc32 = ret_matting_mask.astype(np.float32) / 255 + ret_img_fc32 = ret_img.astype(np.float32) / 255 + hair_matting_mask_fc32 = hair_matting_mask_8uc3.astype(np.float32) / 255 + user_rgb_fc32_orisize = user_rgb_8uc3_orisize.astype(np.float32) / 255 + + effect_mask_fc32 = np.zeros_like(ret_matting_mask_fc32) + cv2.circle(effect_mask_fc32, (int(ori_x), int(ori_y)), radis, (1.0, 1.0, 1.0), -1) + effect_mask_fc32 = cv2.GaussianBlur(effect_mask_fc32, (21, 21), 0, 0) + # ret_img_fc32 = ret_img_fc32 * effect_mask_fc32 + user_rgb_fc32_orisize * (1 - effect_mask_fc32) + # ret_matting_mask_fc32 = ret_matting_mask_fc32 * effect_mask_fc32 + hair_matting_mask_fc32 * (1 - effect_mask_fc32) + # ret_img_fc32_new = ret_img_fc32 * ret_matting_mask_fc32 + bald_res_ori_fc32 * (1 - ret_matting_mask_fc32) + + origin_img_fc32 = origin_img.astype(np.float32) / 255 + ret_img_fc32_new = ret_img_fc32 * ret_matting_mask_fc32 + bald_res_ori_fc32 * (1 - ret_matting_mask_fc32) + + # show_concat = np.concatenate((ret_img_fc32, bald_res_ori_fc32, ret_matting_mask_fc32, ret_img_fc32_new), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_show1", show_concat) + # cv2.waitKey() + + # effect_mask_fc32 = effect_mask_fc32 * ret_matting_mask_fc32 + # effect_mask_fc32 = effect_mask_fc32 * hair_matting_mask_fc32 + ret_img_fc32_new2 = ret_img_fc32_new * effect_mask_fc32 + origin_img_fc32 * (1 - effect_mask_fc32) + + # show_concat = np.concatenate((ret_img_fc32_new, origin_img_fc32, effect_mask_fc32, ret_img_fc32_new2), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_show", show_concat) + # cv2.waitKey() + + ret_img = (ret_img_fc32_new2*255).astype(np.uint8) + self.logger_process.info('fix hair costs:{}'.format(time.time() - t0)) + fix_image_dir = osp.join(userinfo_dir, taskid + '.jpg') + matting_after_path = osp.join(userinfo_dir, taskid + '.png') + cv2.imwrite(fix_image_dir, ret_img) + ret_fix_img_mask_fc32_4chans = np.concatenate((ret_img_fc32, ret_matting_mask_fc32[:, :, :1]), axis=2) + cv2.imwrite(matting_after_path, (ret_fix_img_mask_fc32_4chans*255).astype(np.uint8)) + t4 = time.time() + print('hairfix v2, check_res costs', t4 - t3) + return fix_image_dir, 0 + + def fix_hairstyle_v3(self, origin_img, brush, taskid, userinfo_dir, maskImgName): + t0 = time.time() + landmark1k_path = osp.join(userinfo_dir, 'kpt_1k.txt') + if os.path.exists(landmark1k_path): + landmarks_origin_img_1k = np.loadtxt(landmark1k_path) + else: + + with torch.no_grad(): + landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward(origin_img) + + + if landmarks_origin_img_1k is None: + return None, 10001 + if maskImgName is None or maskImgName == "": + maskImgName = '{}_fix.png'.format(taskid) + matting_before_path = osp.join(userinfo_dir, maskImgName) + # res_hairstyle_before_path = osp.join(userinfo_dir, "user_res_8uc3_orisize_for_finetune.png") + bald_res_ori_path = osp.join(userinfo_dir, 'bald_res_ori.png') + bald_seg_ori_path = osp.join(userinfo_dir, 'bald_seg_ori.png') + matting_raw_path = osp.join(userinfo_dir, 'res_matting_mask_ori_raw.png') + + if os.path.exists(matting_before_path) and osp.exists(matting_raw_path): + hair_fix_img_mask_8uc4 = cv2.imread(matting_before_path, -1) + res_hairstyle_before_8uc3 = hair_fix_img_mask_8uc4[:, :, :3] + hair_matting_mask_8uc3 = np.repeat(hair_fix_img_mask_8uc4[:, :, 3:], 3, axis=2) + hair_gene_matte_fc32_orisize_raw = cv2.imread(matting_raw_path).astype(np.float32) / 255 + else: + _, user_matting_8uc1_bald_orisize, _ = self.process_data.generator_matte.matte_inference(origin_img, landmarks_origin_img_1k) + hair_matting_mask_8uc3 = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) + if not os.path.exists(userinfo_dir): + os.makedirs(userinfo_dir) + hair_gene_matte_fc32_orisize_raw = hair_matting_mask_8uc3.copy() + cv2.imwrite(matting_raw_path, hair_matting_mask_8uc3) + res_hairstyle_before_8uc3 = origin_img.copy() + hair_fix_img_mask_8uc4 = np.concatenate((res_hairstyle_before_8uc3, hair_matting_mask_8uc3[:, :, :1]), axis=2) + cv2.imwrite(matting_before_path, hair_fix_img_mask_8uc4) + + t1 = time.time() + print('hairfix v2, matte_inference costs', t1 - t0) + if osp.exists(bald_res_ori_path) and osp.exists(bald_seg_ori_path): + bald_res_ori_fc32 = cv2.imread(bald_res_ori_path).astype(np.float32)/255 + user_baldseg_8uc3_orisize = cv2.imread(bald_seg_ori_path).astype(np.float32)/255 + + else: + bald_res_ori_8uc3, user_baldseg_8uc3_orisize, _, _, _, _._ = self.process_data.get_prepare_user_768_data(origin_img, landmarks_origin_img_1k, ratio=1) + bald_res_ori_fc32 = bald_res_ori_8uc3.astype(np.float32)/255 + # bald_res_ori_fc32 = self.process_data.generator_baldseg.forward(origin_img, user_matting_8uc1_bald_orisize, + # landmarks_origin_img_1k) + cv2.imwrite(bald_res_ori_path, bald_res_ori_8uc3) + cv2.imwrite(bald_seg_ori_path, user_baldseg_8uc3_orisize) + + t2 = time.time() + print('hairfix v2, bald costs', t2 - t1) + + # cv2.imshow('matting_before', hair_matting_mask_8uc3) + # cv2.imshow('bald_res_ori_fc32', bald_res_ori_fc32) + # cv2.waitKey() + check_res = self.check_hair_fintune_out_face(hair_matting_mask_8uc3, landmarks_origin_img_1k, user_baldseg_8uc3_orisize,hair_gene_matte_fc32_orisize_raw, userinfo_dir) + if check_res: + # print("hair out of face!!!") + fix_image_dir = osp.join(userinfo_dir, taskid + '.jpg') + cv2.imwrite(fix_image_dir, origin_img) + return fix_image_dir, 0 + + t3 = time.time() + print('hairfix v2, check_res costs', t3 - t2) + + # landmark137 = landmark_processor.pts_1k_to_137(landmarks_origin_img_1k) + user_rgb_8uc3_orisize = res_hairstyle_before_8uc3.copy() + [ori_x, ori_y, dst_x, dst_y, radis] = brush + + # hair_matting_mask_8uc3 + # user_baldseg_8uc3_orisize + + def getProtectMask(headMask, hair): + # headMask = cv2.imread("/home/chinatszrn/Desktop/tmp/hair/bald_seg_ori.png").astype(np.float32) / 255. + # hair = cv2.imread("/home/chinatszrn/Desktop/tmp/hair/res_matting_mask_ori_fix.png", + # cv2.IMREAD_UNCHANGED).astype(np.float32) / 255. + # configs + dilateKernelSize = 35 + blurKernelSize = 10 + + # hairMask = hair[:, :, 3:4] + hairMask = hair[:, :, 2:3].astype(np.float32) / 255. + faceMask = (headMask == (1., 0, 0)).all(axis=2).astype(np.float32) + + faceDilateKernel = np.ones([dilateKernelSize, dilateKernelSize], dtype=np.uint8) # dilate to get margin + faceMaskDilate = cv2.dilate(faceMask, faceDilateKernel) + margin = faceMaskDilate - faceMask + margin = cv2.blur(margin, (blurKernelSize, blurKernelSize)) # blur to get soft edge + + protectZone = margin[:, :, np.newaxis] * hairMask + return protectZone + + t0 = time.time() + steps = updateEndPosition(ori_x, ori_y, dst_x, dst_y, radis) + + protectZone_32fc1 = getProtectMask(user_baldseg_8uc3_orisize, hair_matting_mask_8uc3) # fix seam issue + protectZone_32fc1 = np.zeros([user_baldseg_8uc3_orisize.shape[0], user_baldseg_8uc3_orisize.shape[1], 1], dtype=np.float32) # fix seam issue + ret_img = user_rgb_8uc3_orisize.copy() + ret_matting_mask = hair_matting_mask_8uc3.copy() + ret_matting_mask = (ret_matting_mask.astype(np.float32) / 255.) * (1 - protectZone_32fc1) + # cv2.imshow("before warp", cv2.resize(ret_matting_mask, (ret_matting_mask.shape[1] // 2, ret_matting_mask.shape[0] // 2))) + + for curStartx, curStarty, curEndx, curEndy, radius in steps: + radius *= 2 # magic number enlarge area + # ret = localTranslationWarp(canvas, startx, starty, endx, endy, radius) + ret_img = localtranslationwarpfastwithstrength_v2(ret_img, curStartx, curStarty, curEndx, curEndy, radius) + ret_matting_mask = localtranslationwarpfastwithstrength_v2(ret_matting_mask, curStartx, curStarty, curEndx, curEndy, radius) + + # add protected zone image + # cv2.imshow("ret img", cv2.resize(ret_img, (ret_img.shape[1] // 2, ret_img.shape[0] // 2))) + # cv2.imshow("after warp", cv2.resize(ret_matting_mask, (ret_matting_mask.shape[1] // 2, ret_matting_mask.shape[0] // 2))) + ret_img_fc32 = (ret_img.astype(np.float32) / 255.) * (1 - protectZone_32fc1) + (user_rgb_8uc3_orisize.astype(np.float32) / 255.) * protectZone_32fc1 + # ret_matting_mask_fc32 = (ret_matting_mask.astype(np.float32) / 255.) * (1 - protectZone_32fc1) + protectZone_32fc1 # (hair_matting_mask_8uc3.astype(np.float32) / 255.) + + ret_matting_mask_fc32 = ret_matting_mask + protectZone_32fc1 # (hair_matting_mask_8uc3.astype(np.float32) / 255.) + # ret_matting_mask_fc32 = ret_matting_mask + + # ret_matting_dilate_kernel = np.ones([25, 25], dtype=np.uint8) + # ret_matting_mask_fc32 = cv2.dilate(ret_matting_mask_fc32, ret_matting_dilate_kernel) + # ret_matting_mask_fc32 = cv2.erode(ret_matting_mask_fc32, ret_matting_dilate_kernel) + # cv2.imshow("protectZone_32fc1", cv2.resize(protectZone_32fc1, (protectZone_32fc1.shape[1] // 2, protectZone_32fc1.shape[0] // 2))) + # cv2.imshow("ret_img_fc32", cv2.resize(ret_img_fc32, (ret_img_fc32.shape[1] // 2, ret_img_fc32.shape[0] // 2))) + # cv2.imshow("ret_matting_mask_fc32", cv2.resize(ret_matting_mask_fc32, (ret_matting_mask_fc32.shape[1] // 2, ret_matting_mask_fc32.shape[0] // 2))) + # cv2.waitKey() + t4 = time.time() + print('hairfix v2, localtranslationwarpfastwithstrength_v2 costs', t4 - t0) + + # originMask_32fc1 = 1. - (hair_matting_mask_8uc3==[0, 0, 0]).all(axis=2).astype(np.float32) + # originMask_32fc1 = originMask_32fc1[:, :, np.newaxis] # not select has, but select larger + originMask_32fc1 = (hair_matting_mask_8uc3[:, :, 0].astype(np.float32) / 255. > ret_matting_mask_fc32[:, :, 0]).astype(np.float32) + originMask_32fc1 = originMask_32fc1[:, :, np.newaxis] + + # cv2.imshow("originMask_32fc1", cv2.resize(originMask_32fc1, (originMask_32fc1.shape[1] // 2, originMask_32fc1.shape[0] // 2))) + # cv2.imshow("hair_matting_mask_8uc3", cv2.resize(hair_matting_mask_8uc3, (hair_matting_mask_8uc3.shape[1] // 2, hair_matting_mask_8uc3.shape[0] // 2))) + # cv2.imshow("ret_matting_mask_fc32", cv2.resize(ret_matting_mask_fc32, (ret_matting_mask_fc32.shape[1] // 2, ret_matting_mask_fc32.shape[0] // 2))) + ret_matting_mask_fc32 = hair_matting_mask_8uc3.astype(np.float32) / 255. * originMask_32fc1 + ret_matting_mask_fc32 * (1 - originMask_32fc1) + + ret_img_fc32_new = ret_img_fc32 * ret_matting_mask_fc32 + bald_res_ori_fc32 * (1 - ret_matting_mask_fc32) # merge to bald, cloth impainting artifacts + # cv2.imshow("outputMerge", cv2.resize(ret_matting_mask_fc32, (ret_matting_mask_fc32.shape[1] // 2, ret_matting_mask_fc32.shape[0] // 2))) + # cv2.imshow("output1", cv2.resize(ret_img_fc32_new, (ret_img_fc32_new.shape[1] // 2, ret_img_fc32_new.shape[0] // 2))) + + # keep origin area + effect_mask_fc32 = np.zeros_like(ret_matting_mask_fc32) + cv2.circle(effect_mask_fc32, (int(ori_x), int(ori_y)), radis * 3, (1.0, 1.0, 1.0), -1) + blurKernel = radis if radis % 2 == 1 else radis - 1 + effect_mask_fc32 = cv2.GaussianBlur(effect_mask_fc32, (blurKernel, blurKernel), 0, 0) + ret_img_fc32_new = ret_img_fc32_new * effect_mask_fc32 + origin_img.astype(np.float32) / 255. * (1 - effect_mask_fc32) + # cv2.imshow("output2", cv2.resize(ret_img_fc32_new, (ret_img_fc32_new.shape[1] // 2, ret_img_fc32_new.shape[0] // 2))) + # cv2.waitKey() + + ret_img_fc32_new = np.clip(ret_img_fc32_new, 0., 1.0) + ret_img = (ret_img_fc32_new*255).astype(np.uint8) + + self.logger_process.info('fix hair costs:{}'.format(time.time() - t0)) + fix_image_dir = osp.join(userinfo_dir, taskid + '.jpg') + matting_after_path = osp.join(userinfo_dir, taskid + '.png') + cv2.imwrite(fix_image_dir, ret_img) + ret_fix_img_mask_fc32_4chans = np.concatenate((ret_img_fc32, ret_matting_mask_fc32[:, :, :1]), axis=2) + cv2.imwrite(matting_after_path, (ret_fix_img_mask_fc32_4chans*255).astype(np.uint8)) + t4 = time.time() + print('hairfix v2, check_res costs', t4 - t3) + return fix_image_dir, 0 + + def fix_hairstyle_v4(self, origin_img, brush, taskid, userinfo_dir, maskImgName, hairId): + t0 = time.time() + userinfo_dir = osp.join(userinfo_dir, hairId) + os.makedirs(userinfo_dir, exist_ok=True) + landmark1k_path = osp.join(userinfo_dir, 'kpt_1k.txt') + if os.path.exists(landmark1k_path): + landmarks_origin_img_1k = np.loadtxt(landmark1k_path) + else: + landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward(origin_img) + + if landmarks_origin_img_1k is None: + return None, 10001 + if maskImgName is None or maskImgName == "": + maskImgName = '{}_fix.png'.format(taskid) + matting_before_path = osp.join(userinfo_dir, maskImgName) + # res_hairstyle_before_path = osp.join(userinfo_dir, "user_res_8uc3_orisize_for_finetune.png") + bald_res_ori_path = osp.join(userinfo_dir, 'bald_res_ori.png') + bald_seg_ori_path = osp.join(userinfo_dir, 'bald_seg_ori.png') + matting_raw_path = osp.join(userinfo_dir, 'res_matting_mask_ori_raw.png') + + # face_bbox = cv2.boundingRect(landmarks_origin_img_1k[np.newaxis].astype(np.int32)) + # face_short_wide = min(face_bbox[2], face_bbox[3]) + # if brush[4] > int(0.5 * face_short_wide): + # brush[4] = int(0.5 * face_short_wide) + # if brush[4] < int(0.1 * face_short_wide): + # brush[4] = int(0.1 * face_short_wide) + + if os.path.exists(matting_before_path) and osp.exists(matting_raw_path): + hair_fix_img_mask_8uc4 = cv2.imread(matting_before_path, -1) + res_hairstyle_before_8uc3 = hair_fix_img_mask_8uc4[:, :, :3] + hair_matting_mask_8uc3 = np.repeat(hair_fix_img_mask_8uc4[:, :, 3:], 3, axis=2) + hair_gene_matte_fc32_orisize_raw = cv2.imread(matting_raw_path).astype(np.float32) / 255 + else: + _, user_matting_8uc1_bald_orisize, _ = self.process_data.generator_matte.matte_inference(origin_img, landmarks_origin_img_1k) + hair_matting_mask_8uc3 = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) + hair_gene_matte_fc32_orisize_raw = hair_matting_mask_8uc3.copy() + cv2.imwrite(matting_raw_path, hair_matting_mask_8uc3) + res_hairstyle_before_8uc3 = origin_img.copy() + hair_fix_img_mask_8uc4 = np.concatenate((res_hairstyle_before_8uc3, hair_matting_mask_8uc3[:, :, :1]), axis=2) + cv2.imwrite(matting_before_path, hair_fix_img_mask_8uc4) + print('@@@@@@@@@@ new matting') + + t1 = time.time() + print('hairfix v4, matte_inference costs', t1 - t0) + if osp.exists(bald_res_ori_path) and osp.exists(bald_seg_ori_path): + bald_res_ori_fc32 = cv2.imread(bald_res_ori_path).astype(np.float32)/255 + user_baldseg_8uc3_orisize = cv2.imread(bald_seg_ori_path).astype(np.float32)/255 + else: + bald_res_ori_8uc3, user_baldseg_8uc3_orisize, _, _, _, _,_ = self.process_data.get_prepare_user_768_data(origin_img, landmarks_origin_img_1k, ratio=1) + bald_res_ori_fc32 = bald_res_ori_8uc3.astype(np.float32)/255 + # bald_res_ori_fc32 = self.process_data.generator_baldseg.forward(origin_img, user_matting_8uc1_bald_orisize, + # landmarks_origin_img_1k) + cv2.imwrite(bald_res_ori_path, bald_res_ori_8uc3) + cv2.imwrite(bald_seg_ori_path, user_baldseg_8uc3_orisize) + + t2 = time.time() + print('hairfix v4, get_prepare_user_768_data costs', t2 - t1) + + # cv2.imshow('matting_before', hair_matting_mask_8uc3) + # cv2.imshow('bald_res_ori_fc32', bald_res_ori_fc32) + # cv2.waitKey() + check_res = self.check_hair_fintune_out_face(hair_matting_mask_8uc3, landmarks_origin_img_1k, user_baldseg_8uc3_orisize, hair_gene_matte_fc32_orisize_raw, userinfo_dir) + if check_res: + # print("hair out of face!!!") + fix_image_dir = osp.join(userinfo_dir, taskid + '.jpg') + cv2.imwrite(fix_image_dir, origin_img) + return fix_image_dir, 0 + + t3 = time.time() + print('hairfix v4, check_hair_fintune_out_face costs', t3 - t2) + + # landmark137 = landmark_processor.pts_1k_to_137(landmarks_origin_img_1k) + # user_rgb_8uc3_orisize = res_hairstyle_before_8uc3.copy() + [ori_x, ori_y, dst_x, dst_y, radis] = brush + + steps = updateEndPosition(ori_x, ori_y, dst_x, dst_y, radis) + ret_img = res_hairstyle_before_8uc3.copy() + ret_matting_mask = hair_matting_mask_8uc3.copy() + for curStartx, curStarty, curEndx, curEndy, radius in steps: + radius *= 2 # magic number enlarge area + # print("steps params: ", curStartx, " ", curStarty, " ", curEndx, " ", curEndy, " ", radius) + ret_img = localtranslationwarpfastwithstrength_v2_soft(ret_img, curStartx, curStarty, curEndx, curEndy, + radius) + ret_matting_mask = localtranslationwarpfastwithstrength_v2_soft(ret_matting_mask, curStartx, curStarty, + curEndx, curEndy, radius) + + # show_concat = np.concatenate((res_hairstyle_before_8uc3, ret_img, ret_matting_mask), axis=1) + # ratio = 1236. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("warp_comp", show_concat) + # cv2.waitKey() + + t4 = time.time() + print('hairfix v4, localtranslationwarpfastwithstrength_v2 costs', t4 - t3) + + ret_matting_mask_fc32 = ret_matting_mask.astype(np.float32) / 255 + ret_img_fc32 = ret_img.astype(np.float32) / 255 + # hair_matting_mask_fc32 = hair_matting_mask_8uc3.astype(np.float32) / 255 + # user_rgb_fc32_orisize = user_rgb_8uc3_orisize.astype(np.float32) / 255 + + # effect_mask_fc32 = np.zeros_like(ret_matting_mask_fc32) + # cv2.circle(effect_mask_fc32, (int(ori_x), int(ori_y)), radis * 1, (1.0, 1.0, 1.0), -1) + # blurKernel = radis if radis % 2 == 1 else radis - 1 + # effect_mask_fc32 = cv2.GaussianBlur(effect_mask_fc32, (blurKernel, blurKernel), 0, 0) + + center_x = self.effect_prepare_mask_fc32.shape[1] // 2 + center_y = self.effect_prepare_mask_fc32.shape[0] // 2 + blurKernel = radis if radis % 2 == 1 else radis - 1 + dst_circle_w = blurKernel + radis + random_scalex = dst_circle_w / (150 + 149) + random_scaley = random_scalex + M_warp = getM((center_x, center_y), 0, random_scalex, random_scaley) + M_warp[:, 2] += [int(ori_x) - center_x, int(ori_y) - center_y] + effect_mask_fc32 = cv2.warpAffine(self.effect_prepare_mask_fc32, M_warp, (ret_img_fc32.shape[1], ret_img_fc32.shape[0]), + flags=cv2.INTER_NEAREST) + # cv2.imshow("effect_mask_fc32", effect_mask_fc32) + # cv2.waitKey() + + origin_img_fc32 = origin_img.astype(np.float32) / 255 + ret_img_fc32_new = ret_img_fc32 * ret_matting_mask_fc32 + bald_res_ori_fc32 * (1 - ret_matting_mask_fc32) + + # show_concat = np.concatenate((ret_img_fc32, bald_res_ori_fc32, ret_matting_mask_fc32, ret_img_fc32_new), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_show1", show_concat) + # cv2.waitKey() + + hair_matting_mask_fc32 = hair_matting_mask_8uc3.astype(np.float32) / 255 + # dif_mask = hair_matting_mask_fc32 * (1 - ret_matting_mask_fc32) + # cv2.imshow("dif_mask", dif_mask) + + # effect_mask_fc32 = effect_mask_fc32 * ret_matting_mask_fc32 + effect_mask_fc32 = effect_mask_fc32 * hair_matting_mask_fc32 + ret_img_fc32_new2 = ret_img_fc32_new * effect_mask_fc32 + origin_img_fc32 * (1 - effect_mask_fc32) + + t5 = time.time() + print('hairfix v4, GaussianBlur and fusion costs', t5 - t4) + + # show_concat = np.concatenate((ret_img_fc32_new, origin_img_fc32, effect_mask_fc32, ret_img_fc32_new2), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_show2", show_concat) + # cv2.waitKey() + + ret_img = (ret_img_fc32_new2*255).astype(np.uint8) + fix_image_dir = osp.join(userinfo_dir, taskid + '.jpg') + matting_after_path = osp.join(userinfo_dir, taskid + '.png') + cv2.imwrite(fix_image_dir, ret_img) + ret_fix_img_mask_fc32_4chans = np.concatenate((ret_img_fc32_new2, ret_matting_mask_fc32[:, :, :1]), axis=2) + cv2.imwrite(matting_after_path, (ret_fix_img_mask_fc32_4chans*255).astype(np.uint8)) + self.logger_process.info('last step hair costs:{}'.format(time.time() - t5)) + return fix_image_dir, 0 + + def fix_hairstyle_v4_new(self, origin_img, brush, taskid, userinfo_dir, maskImgName, hairId): + t0 = time.time() + userinfo_dir = osp.join(userinfo_dir, hairId) + os.makedirs(userinfo_dir, exist_ok=True) + landmark1k_path = osp.join(userinfo_dir, 'kpt_1k.txt') + if os.path.exists(landmark1k_path): + landmarks_origin_img_1k = np.loadtxt(landmark1k_path) + else: + landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward(origin_img) + + if landmarks_origin_img_1k is None: + return None, 10001 + if maskImgName is None or maskImgName == "": + maskImgName = '{}_fix.png'.format(taskid) + matting_before_path = osp.join(userinfo_dir, maskImgName) + # res_hairstyle_before_path = osp.join(userinfo_dir, "user_res_8uc3_orisize_for_finetune.png") + bald_res_ori_path = osp.join(userinfo_dir, 'bald_res_ori.png') + bald_seg_ori_path = osp.join(userinfo_dir, 'bald_seg_ori.png') + matting_raw_path = osp.join(userinfo_dir, 'res_matting_mask_ori_raw.png') + + face_bbox = cv2.boundingRect(landmarks_origin_img_1k[np.newaxis].astype(np.int32)) + face_short_wide = min(face_bbox[2], face_bbox[3]) + if brush[4] > int(0.5 * face_short_wide): + brush[4] = int(0.5 * face_short_wide) + if brush[4] < int(0.15 * face_short_wide): + brush[4] = int(0.15 * face_short_wide) + + if os.path.exists(matting_before_path) and osp.exists(matting_raw_path): + hair_fix_img_mask_8uc4 = cv2.imread(matting_before_path, -1) + res_hairstyle_before_8uc3 = hair_fix_img_mask_8uc4[:, :, :3] + hair_matting_mask_8uc3 = np.repeat(hair_fix_img_mask_8uc4[:, :, 3:], 3, axis=2) + hair_gene_matte_fc32_orisize_raw = cv2.imread(matting_raw_path).astype(np.float32) / 255 + else: + _, user_matting_8uc1_bald_orisize, _ = self.process_data.generator_matte.matte_inference(origin_img, landmarks_origin_img_1k) + hair_matting_mask_8uc3 = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) + hair_gene_matte_fc32_orisize_raw = hair_matting_mask_8uc3.copy() + cv2.imwrite(matting_raw_path, hair_matting_mask_8uc3) + res_hairstyle_before_8uc3 = origin_img.copy() + hair_fix_img_mask_8uc4 = np.concatenate((res_hairstyle_before_8uc3, hair_matting_mask_8uc3[:, :, :1]), axis=2) + cv2.imwrite(matting_before_path, hair_fix_img_mask_8uc4) + print('@@@@@@@@@@ new matting') + + t1 = time.time() + print('hairfix v4, matte_inference costs', t1 - t0) + if osp.exists(bald_res_ori_path) and osp.exists(bald_seg_ori_path): + bald_res_ori_fc32 = cv2.imread(bald_res_ori_path).astype(np.float32)/255 + user_baldseg_8uc3_orisize = cv2.imread(bald_seg_ori_path).astype(np.float32)/255 + else: + bald_res_ori_8uc3, user_baldseg_8uc3_orisize, _, _, _, _,_ = self.process_data.get_prepare_user_768_data(origin_img, landmarks_origin_img_1k, ratio=1) + bald_res_ori_fc32 = bald_res_ori_8uc3.astype(np.float32)/255 + # bald_res_ori_fc32 = self.process_data.generator_baldseg.forward(origin_img, user_matting_8uc1_bald_orisize, + # landmarks_origin_img_1k) + cv2.imwrite(bald_res_ori_path, bald_res_ori_8uc3) + cv2.imwrite(bald_seg_ori_path, user_baldseg_8uc3_orisize) + + t2 = time.time() + print('hairfix v4, get_prepare_user_768_data costs', t2 - t1) + + # cv2.imshow('matting_before', hair_matting_mask_8uc3) + # cv2.imshow('bald_res_ori_fc32', bald_res_ori_fc32) + # cv2.waitKey() + check_res = self.check_hair_fintune_out_face(hair_matting_mask_8uc3, landmarks_origin_img_1k, user_baldseg_8uc3_orisize, hair_gene_matte_fc32_orisize_raw, userinfo_dir) + if check_res: + # print("hair out of face!!!") + fix_image_dir = osp.join(userinfo_dir, taskid + '.jpg') + cv2.imwrite(fix_image_dir, origin_img) + return fix_image_dir, 0 + + t3 = time.time() + print('hairfix v4, check_hair_fintune_out_face costs', t3 - t2) + + # landmark137 = landmark_processor.pts_1k_to_137(landmarks_origin_img_1k) + # user_rgb_8uc3_orisize = res_hairstyle_before_8uc3.copy() + [ori_x, ori_y, dst_x, dst_y, radis] = brush + + steps = updateEndPosition(ori_x, ori_y, dst_x, dst_y, radis) + ret_img = res_hairstyle_before_8uc3.copy() + ret_matting_mask = hair_matting_mask_8uc3.copy() + for curStartx, curStarty, curEndx, curEndy, radius in steps: + radius *= 2 # magic number enlarge area + # print("steps params: ", curStartx, " ", curStarty, " ", curEndx, " ", curEndy, " ", radius) + ret_img = localtranslationwarpfastwithstrength_v2_soft(ret_img, curStartx, curStarty, curEndx, curEndy, + radius) + ret_matting_mask = localtranslationwarpfastwithstrength_v2_soft(ret_matting_mask, curStartx, curStarty, + curEndx, curEndy, radius) + + # show_concat = np.concatenate((res_hairstyle_before_8uc3, ret_img, ret_matting_mask), axis=1) + # ratio = 1236. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("warp_comp", show_concat) + # cv2.waitKey() + + t4 = time.time() + print('hairfix v4, localtranslationwarpfastwithstrength_v2 costs', t4 - t3) + + ret_matting_mask_fc32 = ret_matting_mask.astype(np.float32) / 255 + ret_img_fc32 = ret_img.astype(np.float32) / 255 + # hair_matting_mask_fc32 = hair_matting_mask_8uc3.astype(np.float32) / 255 + # user_rgb_fc32_orisize = user_rgb_8uc3_orisize.astype(np.float32) / 255 + + # effect_mask_fc32 = np.zeros_like(ret_matting_mask_fc32) + # cv2.circle(effect_mask_fc32, (int(ori_x), int(ori_y)), radis * 1, (1.0, 1.0, 1.0), -1) + # blurKernel = radis if radis % 2 == 1 else radis - 1 + # effect_mask_fc32 = cv2.GaussianBlur(effect_mask_fc32, (blurKernel, blurKernel), 0, 0) + + center_x = self.effect_prepare_mask_fc32.shape[1] // 2 + center_y = self.effect_prepare_mask_fc32.shape[0] // 2 + blurKernel = radis if radis % 2 == 1 else radis - 1 + dst_circle_w = blurKernel + radis + random_scalex = dst_circle_w / (150 + 149) + random_scaley = random_scalex + M_warp = getM((center_x, center_y), 0, random_scalex, random_scaley) + M_warp[:, 2] += [int(ori_x) - center_x, int(ori_y) - center_y] + effect_mask_fc32 = cv2.warpAffine(self.effect_prepare_mask_fc32, M_warp, (ret_img_fc32.shape[1], ret_img_fc32.shape[0]), + flags=cv2.INTER_NEAREST) + # cv2.imshow("effect_mask_fc32", effect_mask_fc32) + # cv2.waitKey() + + origin_img_fc32 = origin_img.astype(np.float32) / 255 + ret_img_fc32_new = ret_img_fc32 * ret_matting_mask_fc32 + bald_res_ori_fc32 * (1 - ret_matting_mask_fc32) + + # show_concat = np.concatenate((ret_img_fc32, bald_res_ori_fc32, ret_matting_mask_fc32, ret_img_fc32_new), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_show1", show_concat) + # cv2.waitKey() + + hair_matting_mask_fc32 = hair_matting_mask_8uc3.astype(np.float32) / 255 + # dif_mask = hair_matting_mask_fc32 * (1 - ret_matting_mask_fc32) + # cv2.imshow("dif_mask", dif_mask) + + # effect_mask_fc32 = effect_mask_fc32 * ret_matting_mask_fc32 + effect_mask_fc32 = effect_mask_fc32 * hair_matting_mask_fc32 + ret_img_fc32_new2 = ret_img_fc32_new * effect_mask_fc32 + origin_img_fc32 * (1 - effect_mask_fc32) + + # show_concat = np.concatenate((ret_img_fc32_new, origin_img_fc32, effect_mask_fc32, ret_img_fc32_new2), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_show1", show_concat) + # cv2.waitKey() + + # effect_mask_8uc3_clip = np.zeros_like(origin_img) + # ret_matting_mask_8uc3_clip = np.zeros_like(origin_img) + # effect_mask_8uc3_clip[effect_mask_fc32[:, :, 0] > 0.01] = 255 + # ret_matting_mask_8uc3_clip[hair_matting_mask_8uc3[:, :, 0] > 1] = 255 + # erode_mask_fc32 = (1 - ret_matting_mask_8uc3_clip.astype(np.float32)/255) * effect_mask_8uc3_clip + # mask_concat = np.concatenate((ret_matting_mask_8uc3_clip, effect_mask_8uc3_clip, erode_mask_fc32), axis=1) + # ratio = 1536. / max(mask_concat.shape[:2]) + # mask_concat = cv2.resize(mask_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mask_concat", mask_concat) + # cv2.waitKey() + + erode_mask_fc32 = (1 - ret_matting_mask_fc32) * effect_mask_fc32 + erode_mask_8uc3_clip = np.zeros_like(origin_img) + erode_mask_8uc3_clip[erode_mask_fc32[:, :, 0] > 0.01] = 255 + # cv2.imshow("erode_mask_fc32", erode_mask_fc32) + # cv2.imshow("erode_mask_8uc3_clip", erode_mask_8uc3_clip) + ret_img_8uc3_new = (ret_img_fc32_new * 255).astype(np.uint8) + # effect_mask_fc32 = np.clip(effect_mask_fc32, 0, 1) + # fusion_mask_8uc3 = (effect_mask_fc32 * 255).astype(np.uint8)[:, :, 0] + fusion_mask_8uc3 = (erode_mask_8uc3_clip).astype(np.uint8)[:, :, 0] + # fusion_mask = effect_mask_fc32[:, :, 0] + box_past = cv2.boundingRect(fusion_mask_8uc3) # 外接矩形 + # print("box_past: ", box_past) + cx = int(box_past[0] + box_past[2] / 2) + cy = int(box_past[1] + box_past[3] / 2) + ret_img_8uc3_new_possion = cv2.seamlessClone(ret_img_8uc3_new, origin_img, fusion_mask_8uc3, (cx, cy), cv2.NORMAL_CLONE) + + ret_img_fc32_new_possion = ret_img_8uc3_new_possion.astype(np.float32) / 255 + # effect_mask_fc32_possion = cv2.GaussianBlur(effect_mask_fc32, (11, 11), 0, 0) + # erode_mask_fc32 = erode_mask_8uc3_clip.astype(np.float32) / 255 + erode_mask_fc32 = np.clip(erode_mask_fc32*10, 0, 1.0) + effect_mask_fc32_possion = cv2.GaussianBlur(erode_mask_fc32, (11, 11), 0, 0) + ret_img_fc32_new_possion2 = ret_img_fc32_new_possion * effect_mask_fc32_possion + ret_img_fc32_new2 * (1 - effect_mask_fc32_possion) + + # show_concat = np.concatenate((ret_img_fc32_new_possion, ret_img_fc32_new2, effect_mask_fc32_possion, ret_img_fc32_new_possion2), axis=1) + # ratio = 1436. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("possion_show", show_concat) + # cv2.waitKey() + + # cv2.imwrite(osp.join(userinfo_dir, "ret_img_8uc3_new_possion2_res.png"), (ret_img_fc32_new_possion2*255).astype(np.uint8)) + # show_concat = np.concatenate((ret_img_8uc3_new, origin_img, np.repeat(fusion_mask_8uc3[:, :, np.newaxis], 3, axis=2), + # ret_img_8uc3_new_possion, (ret_img_fc32_new_possion2*255).astype(np.uint8)), axis=1) + # ratio = 1436. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("possion_show", show_concat) + # cv2.waitKey() + + ret_img_fc32_new2 = ret_img_fc32_new_possion2 + t5 = time.time() + print('hairfix v4, GaussianBlur and fusion costs', t5 - t4) + + # show_concat = np.concatenate((ret_img_fc32_new, origin_img_fc32, effect_mask_fc32, ret_img_fc32_new2), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_show2", show_concat) + # cv2.waitKey() + + ret_img = (ret_img_fc32_new2*255).astype(np.uint8) + fix_image_dir = osp.join(userinfo_dir, taskid + '.jpg') + matting_after_path = osp.join(userinfo_dir, taskid + '.png') + cv2.imwrite(fix_image_dir, ret_img) + ret_fix_img_mask_fc32_4chans = np.concatenate((ret_img_fc32_new2, ret_matting_mask_fc32[:, :, :1]), axis=2) + cv2.imwrite(matting_after_path, (ret_fix_img_mask_fc32_4chans*255).astype(np.uint8)) + self.logger_process.info('last step hair costs:{}'.format(time.time() - t5)) + return fix_image_dir, 0 + + def fix_hairstyle_v5(self, origin_img, brush, taskid, userinfo_dir, maskImgName, hairId): + t0 = time.time() + userinfo_dir = osp.join(userinfo_dir, hairId) + os.makedirs(userinfo_dir, exist_ok=True) + landmark1k_path = osp.join(userinfo_dir, 'kpt_1k.txt') + if os.path.exists(landmark1k_path): + landmarks_origin_img_1k = np.loadtxt(landmark1k_path) + else: + landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward(origin_img) + + if landmarks_origin_img_1k is None: + return None, 10001 + if maskImgName is None or maskImgName == "": + maskImgName = '{}_fix.png'.format(taskid) + matting_before_path = osp.join(userinfo_dir, maskImgName) + # res_hairstyle_before_path = osp.join(userinfo_dir, "user_res_8uc3_orisize_for_finetune.png") + bald_res_ori_path = osp.join(userinfo_dir, 'bald_res_ori.png') + bald_seg_ori_path = osp.join(userinfo_dir, 'bald_seg_ori.png') + matting_raw_path = osp.join(userinfo_dir, 'res_matting_mask_ori_raw.png') + + if os.path.exists(matting_before_path) and osp.exists(matting_raw_path): + hair_fix_img_mask_8uc4 = cv2.imread(matting_before_path, -1) + res_hairstyle_before_8uc3 = hair_fix_img_mask_8uc4[:, :, :3] + hair_matting_mask_8uc3 = np.repeat(hair_fix_img_mask_8uc4[:, :, 3:], 3, axis=2) + hair_gene_matte_fc32_orisize_raw = cv2.imread(matting_raw_path).astype(np.float32) / 255 + else: + _, user_matting_8uc1_bald_orisize, _ = self.process_data.generator_matte.matte_inference(origin_img, landmarks_origin_img_1k) + hair_matting_mask_8uc3 = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) + hair_gene_matte_fc32_orisize_raw = hair_matting_mask_8uc3.copy() + cv2.imwrite(matting_raw_path, hair_matting_mask_8uc3) + res_hairstyle_before_8uc3 = origin_img.copy() + hair_fix_img_mask_8uc4 = np.concatenate((res_hairstyle_before_8uc3, hair_matting_mask_8uc3[:, :, :1]), axis=2) + cv2.imwrite(matting_before_path, hair_fix_img_mask_8uc4) + print('@@@@@@@@@@ new matting') + + t1 = time.time() + print('hairfix v5, matte_inference costs', t1 - t0) + if osp.exists(bald_res_ori_path) and osp.exists(bald_seg_ori_path): + bald_res_ori_fc32 = cv2.imread(bald_res_ori_path).astype(np.float32)/255 + user_baldseg_8uc3_orisize = cv2.imread(bald_seg_ori_path).astype(np.float32)/255 + else: + bald_res_ori_8uc3, user_baldseg_8uc3_orisize, _, _, _, _,_ = self.process_data.get_prepare_user_768_data(origin_img, landmarks_origin_img_1k, ratio=1) + bald_res_ori_fc32 = bald_res_ori_8uc3.astype(np.float32)/255 + # bald_res_ori_fc32 = self.process_data.generator_baldseg.forward(origin_img, user_matting_8uc1_bald_orisize, + # landmarks_origin_img_1k) + cv2.imwrite(bald_res_ori_path, bald_res_ori_8uc3) + cv2.imwrite(bald_seg_ori_path, user_baldseg_8uc3_orisize) + + t2 = time.time() + print('hairfix v5, get_prepare_user_768_data costs', t2 - t1) + + # cv2.imshow('matting_before', hair_matting_mask_8uc3) + # cv2.imshow('bald_res_ori_fc32', bald_res_ori_fc32) + # cv2.waitKey() + check_res = self.check_hair_fintune_out_face(hair_matting_mask_8uc3, landmarks_origin_img_1k, user_baldseg_8uc3_orisize, hair_gene_matte_fc32_orisize_raw, userinfo_dir) + if check_res: + # print("hair out of face!!!") + fix_image_dir = osp.join(userinfo_dir, taskid + '.jpg') + cv2.imwrite(fix_image_dir, origin_img) + return fix_image_dir, 0 + + t3 = time.time() + print('hairfix v5, check_hair_fintune_out_face costs', t3 - t2) + + # landmark137 = landmark_processor.pts_1k_to_137(landmarks_origin_img_1k) + # user_rgb_8uc3_orisize = res_hairstyle_before_8uc3.copy() + [ori_x, ori_y, dst_x, dst_y, radis] = brush + + steps = updateEndPosition(ori_x, ori_y, dst_x, dst_y, radis) + ret_img = res_hairstyle_before_8uc3.copy() + ret_matting_mask = hair_matting_mask_8uc3.copy() + for curStartx, curStarty, curEndx, curEndy, radius in steps: + # radius *= 2 # magic number enlarge area + print("steps params: ", curStartx, " ", curStarty, " ", curEndx, " ", curEndy, " ", radius) + ret_img = localtranslationwarpfastwithstrength_v2_soft(ret_img, curStartx, curStarty, curEndx, curEndy, radius) + ret_matting_mask = localtranslationwarpfastwithstrength_v2_soft(ret_matting_mask, curStartx, curStarty, curEndx, curEndy, radius) + + show_concat = np.concatenate((res_hairstyle_before_8uc3, ret_img, ret_matting_mask), axis=1) + ratio = 1236. / max(show_concat.shape[:2]) + show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + cv2.imshow("mid_show2", show_concat) + cv2.waitKey() + + t4 = time.time() + print('hairfix v5, localtranslationwarpfastwithstrength_v2 costs', t4 - t3) + + ret_matting_mask_fc32 = ret_matting_mask.astype(np.float32) / 255 + ret_img_fc32 = ret_img.astype(np.float32) / 255 + # hair_matting_mask_fc32 = hair_matting_mask_8uc3.astype(np.float32) / 255 + # user_rgb_fc32_orisize = user_rgb_8uc3_orisize.astype(np.float32) / 255 + + center_x = self.effect_prepare_mask_fc32.shape[1] // 2 + center_y = self.effect_prepare_mask_fc32.shape[0] // 2 + blurKernel = radis if radis % 2 == 1 else radis - 1 + dst_circle_w = blurKernel + radis + random_scalex = dst_circle_w / (150 + 149) + random_scaley = random_scalex + M_warp = getM((center_x, center_y), 0, random_scalex, random_scaley) + M_warp[:, 2] += [int(ori_x) - center_x, int(ori_y) - center_y] + effect_mask_fc32 = cv2.warpAffine(self.effect_prepare_mask_fc32, M_warp, (ret_img_fc32.shape[1], ret_img_fc32.shape[0]), + flags=cv2.INTER_NEAREST) + # cv2.imshow("effect_mask_fc32", effect_mask_fc32) + # cv2.waitKey() + + origin_img_fc32 = origin_img.astype(np.float32) / 255 + ret_img_fc32_new = ret_img_fc32 * ret_matting_mask_fc32 + bald_res_ori_fc32 * (1 - ret_matting_mask_fc32) + + show_concat = np.concatenate((ret_img_fc32, bald_res_ori_fc32, ret_matting_mask_fc32, ret_img_fc32_new), axis=1) + ratio = 1536. / max(show_concat.shape[:2]) + show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + cv2.imshow("mid_show1", show_concat) + cv2.waitKey() + + # effect_mask_fc32 = effect_mask_fc32 * ret_matting_mask_fc32 + effect_mask_fc32 = effect_mask_fc32 * hair_gene_matte_fc32_orisize_raw + # adapt_mask_fc32 = hair_gene_matte_fc32_orisize_raw * (1 - ret_matting_mask_fc32) * effect_mask_fc32 + + # ret_img_fc32_new2 = ret_img_fc32_new * effect_mask_fc32 + origin_img_fc32 * (1 - effect_mask_fc32) + + hair_matting_mask_fc32 = hair_matting_mask_8uc3.astype(np.float32) / 255 + fusion_mask_fc32 = hair_matting_mask_fc32 * (1 - ret_matting_mask_fc32) * effect_mask_fc32 + fusion_mask_8uc3 = np.zeros_like(hair_matting_mask_8uc3) + fusion_mask_8uc3[fusion_mask_fc32[:, :, 0] > 0.01] = 255 + # cv2.imshow("fusion_mask_fc32", fusion_mask_fc32) + # cv2.imshow("fusion_mask_8uc3", fusion_mask_8uc3) + + ret_img_8uc3_new = (ret_img_fc32_new*255).astype(np.uint8) + effect_mask_fc32 = np.clip(effect_mask_fc32, 0, 1) + # fusion_mask = (effect_mask_fc32*255).astype(np.uint8)[:, :, 0] + fusion_mask = fusion_mask_8uc3[:, :, 0] + box_past = cv2.boundingRect(fusion_mask) # 外接矩形 + if box_past[2] > 2 and box_past[3] > 2: + cx = int(box_past[0] + box_past[2] / 2) + cy = int(box_past[1] + box_past[3] / 2) + ret_img_8uc3_new2 = cv2.seamlessClone(ret_img_8uc3_new, origin_img, fusion_mask, (cx, cy), cv2.NORMAL_CLONE) + + show_concat = np.concatenate((ret_img_8uc3_new, origin_img, np.repeat(fusion_mask[:, :, np.newaxis], 3, axis=2), ret_img_8uc3_new2), axis=1) + ratio = 1436. / max(show_concat.shape[:2]) + show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + cv2.imshow("possion_show", show_concat) + cv2.waitKey() + + ret_img_fc32_new = ret_img_8uc3_new2.astype(np.float32) / 255 + effect_mask_fc32 = cv2.GaussianBlur(effect_mask_fc32, (11, 11), 0, 0) + ret_img_fc32_new2 = ret_img_fc32_new * effect_mask_fc32 + origin_img_fc32 * (1 - effect_mask_fc32) + else: + ret_img_fc32_new2 = ret_img_fc32_new * effect_mask_fc32 + origin_img_fc32 * (1 - effect_mask_fc32) + + t5 = time.time() + print('hairfix v5, GaussianBlur and fusion costs', t5 - t4) + + show_concat = np.concatenate((ret_img_fc32_new, origin_img_fc32, effect_mask_fc32, ret_img_fc32_new2), axis=1) + ratio = 1436. / max(show_concat.shape[:2]) + show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + cv2.imshow("mid_show2", show_concat) + cv2.waitKey() + + ret_img = (ret_img_fc32_new2*255).astype(np.uint8) + fix_image_dir = osp.join(userinfo_dir, taskid + '.jpg') + matting_after_path = osp.join(userinfo_dir, taskid + '.png') + cv2.imwrite(fix_image_dir, ret_img) + ret_fix_img_mask_fc32_4chans = np.concatenate((ret_img_fc32_new2, ret_matting_mask_fc32[:, :, :1]), axis=2) + cv2.imwrite(matting_after_path, (ret_fix_img_mask_fc32_4chans*255).astype(np.uint8)) + self.logger_process.info('last step hair costs:{}'.format(time.time() - t5)) + return fix_image_dir, 0 + + def fix_hairstyle_v6(self, origin_img, brush, taskid, userinfo_dir, maskImgName, hairId): + t0 = time.time() + userinfo_dir = osp.join(userinfo_dir, hairId) + os.makedirs(userinfo_dir, exist_ok=True) + landmark1k_path = osp.join(userinfo_dir, 'kpt_1k.txt') + if os.path.exists(landmark1k_path): + landmarks_origin_img_1k = np.loadtxt(landmark1k_path) + else: + landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward(origin_img) + + if landmarks_origin_img_1k is None: + return None, 10001 + if maskImgName is None or maskImgName == "": + maskImgName = '{}_fix.png'.format(taskid) + matting_before_path = osp.join(userinfo_dir, maskImgName) + # res_hairstyle_before_path = osp.join(userinfo_dir, "user_res_8uc3_orisize_for_finetune.png") + bald_res_ori_path = osp.join(userinfo_dir, 'bald_res_ori.png') + bald_seg_ori_path = osp.join(userinfo_dir, 'bald_seg_ori.png') + matting_raw_path = osp.join(userinfo_dir, 'res_matting_mask_ori_raw.png') + + if os.path.exists(matting_before_path) and osp.exists(matting_raw_path): + hair_fix_img_mask_8uc4 = cv2.imread(matting_before_path, -1) + res_hairstyle_before_8uc3 = hair_fix_img_mask_8uc4[:, :, :3] + hair_matting_mask_8uc3 = np.repeat(hair_fix_img_mask_8uc4[:, :, 3:], 3, axis=2) + hair_gene_matte_fc32_orisize_raw = cv2.imread(matting_raw_path).astype(np.float32) / 255 + else: + _, user_matting_8uc1_bald_orisize, _ = self.process_data.generator_matte.matte_inference(origin_img, landmarks_origin_img_1k) + hair_matting_mask_8uc3 = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) + hair_gene_matte_fc32_orisize_raw = hair_matting_mask_8uc3.copy() + cv2.imwrite(matting_raw_path, hair_matting_mask_8uc3) + res_hairstyle_before_8uc3 = origin_img.copy() + hair_fix_img_mask_8uc4 = np.concatenate((res_hairstyle_before_8uc3, hair_matting_mask_8uc3[:, :, :1]), axis=2) + cv2.imwrite(matting_before_path, hair_fix_img_mask_8uc4) + print('@@@@@@@@@@ new matting') + + t1 = time.time() + print('hairfix v6, matte_inference costs', t1 - t0) + if osp.exists(bald_res_ori_path) and osp.exists(bald_seg_ori_path): + bald_res_ori_fc32 = cv2.imread(bald_res_ori_path).astype(np.float32)/255 + user_baldseg_8uc3_orisize = cv2.imread(bald_seg_ori_path).astype(np.float32)/255 + else: + bald_res_ori_8uc3, user_baldseg_8uc3_orisize, _, _, _, _,_ = self.process_data.get_prepare_user_768_data(origin_img, landmarks_origin_img_1k, ratio=1) + bald_res_ori_fc32 = bald_res_ori_8uc3.astype(np.float32)/255 + # bald_res_ori_fc32 = self.process_data.generator_baldseg.forward(origin_img, user_matting_8uc1_bald_orisize, + # landmarks_origin_img_1k) + cv2.imwrite(bald_res_ori_path, bald_res_ori_8uc3) + cv2.imwrite(bald_seg_ori_path, user_baldseg_8uc3_orisize) + + t2 = time.time() + print('hairfix v6, get_prepare_user_768_data costs', t2 - t1) + + # cv2.imshow('matting_before', hair_matting_mask_8uc3) + # cv2.imshow('bald_res_ori_fc32', bald_res_ori_fc32) + # cv2.waitKey() + check_res = self.check_hair_fintune_out_face(hair_matting_mask_8uc3, landmarks_origin_img_1k, user_baldseg_8uc3_orisize, hair_gene_matte_fc32_orisize_raw, userinfo_dir) + if check_res: + # print("hair out of face!!!") + fix_image_dir = osp.join(userinfo_dir, taskid + '.jpg') + cv2.imwrite(fix_image_dir, origin_img) + return fix_image_dir, 0 + + t3 = time.time() + print('hairfix v6, check_hair_fintune_out_face costs', t3 - t2) + + # landmark137 = landmark_processor.pts_1k_to_137(landmarks_origin_img_1k) + # user_rgb_8uc3_orisize = res_hairstyle_before_8uc3.copy() + [ori_x, ori_y, dst_x, dst_y, radis] = brush + + steps = updateEndPosition(ori_x, ori_y, dst_x, dst_y, radis) + ret_img = res_hairstyle_before_8uc3.copy() + ret_matting_mask = hair_matting_mask_8uc3.copy() + for curStartx, curStarty, curEndx, curEndy, radius in steps: + radius *= 2 # magic number enlarge area + # print("steps params: ", curStartx, " ", curStarty, " ", curEndx, " ", curEndy, " ", radius) + + tmp_mask = np.zeros_like(ret_img) + cv2.circle(tmp_mask, (int(curStartx), int(curStarty)), radius, (255, 255, 255), -1) + cv2.circle(tmp_mask, (int(curEndx), int(curEndy)), radius, (255, 255, 255), -1) + tmp_bbox = cv2.boundingRect(tmp_mask[:, :, 0]) + # print("tmp_bbox: ", tmp_bbox) + # cv2.imshow("tmp_mask", tmp_mask) + # cv2.waitKey() + + crop_ret_img = ret_img[tmp_bbox[1]:(tmp_bbox[1]+tmp_bbox[3]), tmp_bbox[0]:(tmp_bbox[0]+tmp_bbox[2]), :] + crop_ret_matting_mask = ret_matting_mask[tmp_bbox[1]:(tmp_bbox[1]+tmp_bbox[3]), tmp_bbox[0]:(tmp_bbox[0]+tmp_bbox[2]), :] + + # crop_ret_img = + crop_ret_img = localtranslationwarpfastwithstrength_v2_soft(crop_ret_img, curStartx, curStarty, curEndx, curEndy, radius) + crop_ret_matting_mask = localtranslationwarpfastwithstrength_v2_soft(crop_ret_matting_mask, curStartx, curStarty, curEndx, curEndy, radius) + + ret_img[tmp_bbox[1]:(tmp_bbox[1] + tmp_bbox[3]), tmp_bbox[0]:(tmp_bbox[0] + tmp_bbox[2]), :] = crop_ret_img + ret_matting_mask[tmp_bbox[1]:(tmp_bbox[1] + tmp_bbox[3]), tmp_bbox[0]:(tmp_bbox[0] + tmp_bbox[2]), :] = crop_ret_matting_mask + + # show_concat = np.concatenate((res_hairstyle_before_8uc3, ret_img, ret_matting_mask), axis=1) + # ratio = 1236. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_show2", show_concat) + # cv2.waitKey() + + t4 = time.time() + print('hairfix v6, localtranslationwarpfastwithstrength_v2 costs', t4 - t3) + + ret_matting_mask_fc32 = ret_matting_mask.astype(np.float32) / 255 + ret_img_fc32 = ret_img.astype(np.float32) / 255 + hair_matting_mask_fc32 = hair_matting_mask_8uc3.astype(np.float32) / 255 + # user_rgb_fc32_orisize = user_rgb_8uc3_orisize.astype(np.float32) / 255 + + center_x = self.effect_prepare_mask_fc32.shape[1] // 2 + center_y = self.effect_prepare_mask_fc32.shape[0] // 2 + blurKernel = radis if radis % 2 == 1 else radis - 1 + dst_circle_w = blurKernel + radis + random_scalex = dst_circle_w / (150 + 149) + random_scaley = random_scalex + M_warp = getM((center_x, center_y), 0, random_scalex, random_scaley) + M_warp[:, 2] += [int(ori_x) - center_x, int(ori_y) - center_y] + effect_mask_fc32 = cv2.warpAffine(self.effect_prepare_mask_fc32, M_warp, (ret_img_fc32.shape[1], ret_img_fc32.shape[0]), + flags=cv2.INTER_NEAREST) + # cv2.imshow("effect_mask_fc32", effect_mask_fc32) + # cv2.waitKey() + + origin_img_fc32 = origin_img.astype(np.float32) / 255 + ret_img_fc32_new = ret_img_fc32 * ret_matting_mask_fc32 + bald_res_ori_fc32 * (1 - ret_matting_mask_fc32) + + # show_concat = np.concatenate((ret_img_fc32, bald_res_ori_fc32, ret_matting_mask_fc32, ret_img_fc32_new), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_show1", show_concat) + # cv2.waitKey() + + effect_mask_fc32 = effect_mask_fc32 * hair_matting_mask_fc32 + # effect_mask_fc32 = effect_mask_fc32 * hair_gene_matte_fc32_orisize_raw + # adapt_mask_fc32 = hair_gene_matte_fc32_orisize_raw * (1 - ret_matting_mask_fc32) * effect_mask_fc32 + + # dif_mask_fc32 = hair_matting_mask_fc32 * (1 - ret_matting_mask_fc32) + # cv2.imshow("dif_mask_fc32", dif_mask_fc32) + # cv2.waitKey() + + # ret_img_fc32_new2 = ret_img_fc32_new * effect_mask_fc32 + origin_img_fc32 * (1 - effect_mask_fc32) + + ret_img_8uc3_new = (ret_img_fc32_new*255).astype(np.uint8) + effect_mask_fc32 = np.clip(effect_mask_fc32, 0, 1) + fusion_mask_8uc3 = (effect_mask_fc32*255).astype(np.uint8)[:, :, 0] + # fusion_mask = effect_mask_fc32[:, :, 0] + box_past = cv2.boundingRect(fusion_mask_8uc3) # 外接矩形 + # print("box_past: ", box_past) + if box_past[2] > 2 and box_past[3] > 2: + cx = int(box_past[0] + box_past[2] / 2) + cy = int(box_past[1] + box_past[3] / 2) + ret_img_8uc3_new2 = cv2.seamlessClone(ret_img_8uc3_new, origin_img, fusion_mask_8uc3, (cx, cy), cv2.NORMAL_CLONE) + + show_concat = np.concatenate((ret_img_8uc3_new, origin_img, np.repeat(fusion_mask_8uc3[:, :, np.newaxis], 3, axis=2), ret_img_8uc3_new2), axis=1) + ratio = 1436. / max(show_concat.shape[:2]) + show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + cv2.imshow("possion_show", show_concat) + cv2.waitKey() + + ret_img_fc32_new = ret_img_8uc3_new2.astype(np.float32) / 255 + effect_mask_fc32 = cv2.GaussianBlur(effect_mask_fc32, (11, 11), 0, 0) + ret_img_fc32_new2 = ret_img_fc32_new * effect_mask_fc32 + origin_img_fc32 * (1 - effect_mask_fc32) + else: + print("fusion inter else！！！！！") + ret_img_fc32_new2 = ret_img_fc32_new * effect_mask_fc32 + origin_img_fc32 * (1 - effect_mask_fc32) + + t5 = time.time() + print('hairfix v6, GaussianBlur and fusion costs', t5 - t4) + + show_concat = np.concatenate((ret_img_fc32_new, origin_img_fc32, effect_mask_fc32, ret_img_fc32_new2), axis=1) + ratio = 1436. / max(show_concat.shape[:2]) + show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + cv2.imshow("mid_show", show_concat) + cv2.waitKey() + + ret_img = (ret_img_fc32_new2*255).astype(np.uint8) + fix_image_dir = osp.join(userinfo_dir, taskid + '.jpg') + matting_after_path = osp.join(userinfo_dir, taskid + '.png') + cv2.imwrite(fix_image_dir, ret_img) + ret_fix_img_mask_fc32_4chans = np.concatenate((ret_img_fc32_new2, ret_matting_mask_fc32[:, :, :1]), axis=2) + cv2.imwrite(matting_after_path, (ret_fix_img_mask_fc32_4chans*255).astype(np.uint8)) + self.logger_process.info('last step hair costs:{}'.format(time.time() - t5)) + return fix_image_dir, 0 + + def fix_hairstyle_v7(self, origin_img, brush, taskid, userinfo_dir, maskImgName, hairId): + t0 = time.time() + userinfo_dir = osp.join(userinfo_dir, hairId) + os.makedirs(userinfo_dir, exist_ok=True) + landmark1k_path = osp.join(userinfo_dir, 'kpt_1k.txt') + if os.path.exists(landmark1k_path): + landmarks_origin_img_1k = np.loadtxt(landmark1k_path) + else: + landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward(origin_img) + + if landmarks_origin_img_1k is None: + return None, 10001 + if maskImgName is None or maskImgName == "": + maskImgName = '{}_fix.png'.format(taskid) + # matting_before_path = osp.join(userinfo_dir, maskImgName) + # # res_hairstyle_before_path = osp.join(userinfo_dir, "user_res_8uc3_orisize_for_finetune.png") + # bald_res_ori_path = osp.join(userinfo_dir, 'bald_res_ori.png') + # bald_seg_ori_path = osp.join(userinfo_dir, 'bald_seg_ori.png') + # matting_raw_path = osp.join(userinfo_dir, 'res_matting_mask_ori_raw.png') + + # landmark137 = landmark_processor.pts_1k_to_137(landmarks_origin_img_1k) + # user_rgb_8uc3_orisize = res_hairstyle_before_8uc3.copy() + [ori_x, ori_y, dst_x, dst_y, radis] = brush + + t3 = time.time() + + steps = updateEndPosition(ori_x, ori_y, dst_x, dst_y, radis) + ret_img = origin_img.copy() + for curStartx, curStarty, curEndx, curEndy, radius in steps: + radius *= 2 # magic number enlarge area + # print("steps params: ", curStartx, " ", curStarty, " ", curEndx, " ", curEndy, " ", radius) + + ret_img = localtranslationwarpfastwithstrength_v2_soft(ret_img, curStartx, curStarty, curEndx, curEndy, radius) + + # show_concat = np.concatenate((origin_img, ret_img), axis=1) + # ratio = 1236. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_show2", show_concat) + # cv2.waitKey() + + t4 = time.time() + print('hairfix v7, localtranslationwarpfastwithstrength_v2 costs', t4 - t3) + + # center_x = self.effect_prepare_mask_fc32.shape[1] // 2 + # center_y = self.effect_prepare_mask_fc32.shape[0] // 2 + # blurKernel = radis if radis % 2 == 1 else radis - 1 + # dst_circle_w = blurKernel + radis + # random_scalex = dst_circle_w / (150 + 149) + # random_scaley = random_scalex + # M_warp = getM((center_x, center_y), 0, random_scalex, random_scaley) + # M_warp[:, 2] += [int(ori_x) - center_x, int(ori_y) - center_y] + # effect_mask_fc32 = cv2.warpAffine(self.effect_prepare_mask_fc32, M_warp, (origin_img.shape[1], origin_img.shape[0]), flags=cv2.INTER_NEAREST) + + # cv2.imshow("effect_mask_fc32", effect_mask_fc32) + # cv2.waitKey() + + # origin_img_fc32 = origin_img.astype(np.float32) / 255 + # ret_img_fc32 = ret_img.astype(np.float32) / 255 + # ret_img_fc32_new = ret_img_fc32 * effect_mask_fc32 + origin_img_fc32 * (1 - effect_mask_fc32) + + ret_img_fc32 = ret_img.astype(np.float32) / 255 + ret_img_fc32_new = ret_img_fc32 + + # show_concat = np.concatenate((ret_img_fc32, origin_img_fc32, effect_mask_fc32, ret_img_fc32_new), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_show1", show_concat) + # cv2.waitKey() + + ret_img = (ret_img_fc32_new*255).astype(np.uint8) + fix_image_dir = osp.join(userinfo_dir, taskid + '.jpg') + # matting_after_path = osp.join(userinfo_dir, taskid + '.png') + cv2.imwrite(fix_image_dir, ret_img) + # ret_fix_img_mask_fc32_4chans = np.concatenate((ret_img_fc32_new, ret_matting_mask_fc32[:, :, :1]), axis=2) + # cv2.imwrite(matting_after_path, (ret_fix_img_mask_fc32_4chans*255).astype(np.uint8)) + # self.logger_process.info('last step hair costs:{}'.format(time.time() - t5)) + return fix_image_dir, 0 + + def fix_hairstyle_v2_ycj(self, origin_img, brush, taskid, userinfo_dir, maskImgName): + landmark1k_path = osp.join(userinfo_dir, 'kpt_1k.txt') + if os.path.exists(landmark1k_path): + landmarks_origin_img_1k = np.loadtxt(landmark1k_path) + else: + landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward(origin_img) + + if landmarks_origin_img_1k is None: + return None, 10001 + if maskImgName is None or maskImgName == "": + maskImgName = 'res_matting_mask_ori_fix.png' + matting_before_path = osp.join(userinfo_dir, maskImgName) + bald_res_ori_path = osp.join(userinfo_dir, 'bald_res_ori.png') + bald_seg_ori_path = osp.join(userinfo_dir, 'bald_seg_ori.png') + matting_raw_path = osp.join(userinfo_dir, 'res_matting_mask_ori_raw.png') + + + if os.path.exists(matting_before_path) and osp.exists(matting_raw_path): + hair_matting_mask_8uc3 = cv2.imread(matting_before_path).astype(np.float32) + hair_gene_matte_fc32_orisize_raw = cv2.imread(matting_raw_path).astype(np.float32) / 255 + else: + _, user_matting_8uc1_bald_orisize, _ = self.process_data.generator_matte.matte_inference(origin_img, landmarks_origin_img_1k) + hair_matting_mask_8uc3 = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) + if not os.path.exists(userinfo_dir): + os.makedirs(userinfo_dir) + hair_gene_matte_fc32_orisize_raw = hair_matting_mask_8uc3.copy() + cv2.imwrite(matting_before_path, hair_matting_mask_8uc3) + cv2.imwrite(matting_raw_path, hair_matting_mask_8uc3) + + + if osp.exists(bald_res_ori_path) and osp.exists(bald_seg_ori_path): + bald_res_ori_fc32 = cv2.imread(bald_res_ori_path).astype(np.float32)/255 + user_baldseg_8uc3_orisize = cv2.imread(bald_seg_ori_path).astype(np.float32)/255 + + else: + bald_res_ori_8uc3, user_baldseg_8uc3_orisize, _, _, _, _,_ = self.process_data.get_prepare_user_768_data(origin_img, landmarks_origin_img_1k, ratio=1) + bald_res_ori_fc32 = bald_res_ori_8uc3.astype(np.float32)/255 + # bald_res_ori_fc32 = self.process_data.generator_baldseg.forward(origin_img, user_matting_8uc1_bald_orisize, + # landmarks_origin_img_1k) + cv2.imwrite(bald_res_ori_path, bald_res_ori_8uc3) + cv2.imwrite(bald_seg_ori_path, user_baldseg_8uc3_orisize) + # cv2.imshow('matting_before', hair_matting_mask_8uc3) + # cv2.imshow('bald_res_ori_fc32', bald_res_ori_fc32) + # cv2.waitKey() + check_res = self.check_hair_fintune_out_face(hair_matting_mask_8uc3, landmarks_origin_img_1k, user_baldseg_8uc3_orisize,hair_gene_matte_fc32_orisize_raw, userinfo_dir) + if check_res: + # print("hair out of face!!!") + fix_image_dir = osp.join(userinfo_dir, taskid + '.jpg') + cv2.imwrite(fix_image_dir, origin_img) + return fix_image_dir, 0 + + # landmark137 = landmark_processor.pts_1k_to_137(landmarks_origin_img_1k) + user_rgb_8uc3_orisize = origin_img.copy() + [ori_x, ori_y, dst_x, dst_y, radis] = brush + t0 = time.time() + hair_matting_mask_8uc3[hair_matting_mask_8uc3[:, :, 0] > 0] = 255 + ret_img = localtranslationwarpfastwithstrength_v2(user_rgb_8uc3_orisize.copy(), ori_x, ori_y, dst_x, dst_y, radis) + ret_matting_mask = localtranslationwarpfastwithstrength_v2(hair_matting_mask_8uc3.copy(), ori_x, ori_y, dst_x, dst_y, radis) + # ret_matting_mask = cv2.GaussianBlur(ret_matting_mask, (15, 15), 0, 0) + + # cv2.circle(ret_img, (int(ori_x), int(ori_y)), radis, (255, 0, 0), 3) + # cv2.circle(ret_img, (int(dst_x), int(dst_y)), radis, (255, 0, 0), 3) + # show_concat = np.concatenate((user_rgb_8uc3_orisize, ret_img, hair_matting_mask_8uc3, ret_matting_mask), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_show", show_concat.astype(np.uint8)) + # # cv2.imshow("ret_img", ret_img) + # cv2.waitKey(0) + + ret_matting_mask_fc32 = ret_matting_mask.astype(np.float32) / 255 + ret_img_fc32 = ret_img.astype(np.float32) / 255 + user_rgb_fc32_orisize = user_rgb_8uc3_orisize.astype(np.float32) / 255 + hair_matting_mask_fc32 = hair_matting_mask_8uc3.astype(np.float32) / 255 + + change_hair_mask_fc32 = hair_matting_mask_fc32 * (1 - ret_matting_mask_fc32) + change_hair_mask_fc32 = cv2.GaussianBlur(change_hair_mask_fc32, (15, 15), 0, 0) + ret_matting_mask_fc32 = ret_matting_mask_fc32 * (1 - change_hair_mask_fc32) + + # show_concat = np.concatenate((hair_matting_mask_fc32, change_hair_mask_fc32, ret_matting_mask_fc32), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mask_show", show_concat) + # cv2.waitKey(0) + + effect_mask_fc32 = np.zeros_like(hair_matting_mask_fc32) + cv2.circle(effect_mask_fc32, (int(ori_x), int(ori_y)), radis, (1.0, 1.0, 1.0), -1) + effect_mask_fc32 = cv2.GaussianBlur(effect_mask_fc32, (21, 21), 0, 0) + + ret_matting_mask_fc32_new = ret_matting_mask_fc32 * effect_mask_fc32 + hair_matting_mask_fc32 * (1 - effect_mask_fc32) + ret_img_fc32_new = ret_img_fc32 * effect_mask_fc32 + user_rgb_fc32_orisize * (1 - effect_mask_fc32) + + # show_concat = np.concatenate((hair_matting_mask_fc32, ret_matting_mask_fc32, effect_mask_fc32, ret_matting_mask_fc32_new), axis=1) + # show_concat2 = np.concatenate((user_rgb_fc32_orisize, ret_img_fc32, effect_mask_fc32, ret_img_fc32_new), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # show_concat2 = cv2.resize(show_concat2, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_show", show_concat) + # cv2.imshow("mid_show2", show_concat2) + # cv2.waitKey(0) + + ret_img_fc32_new_old = ret_img_fc32_new * ret_matting_mask_fc32_new + bald_res_ori_fc32 * (1 - ret_matting_mask_fc32_new) + # ret_img_fc32_new_old = ret_img_fc32_new * ret_matting_mask_fc32 + bald_res_ori_fc32 * (1 - ret_matting_mask_fc32) + ret_img_fc32_new = ret_img_fc32_new_old * effect_mask_fc32 + user_rgb_fc32_orisize * (1 - effect_mask_fc32) + + # show_concat = np.concatenate((ret_img_fc32_new_old, effect_mask_fc32, user_rgb_fc32_orisize, ret_img_fc32_new), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_show", show_concat) + # cv2.waitKey() + + # show_concat = np.concatenate((ret_img_fc32, bald_res_ori_fc32, ret_matting_mask_fc32, ret_img_fc32_new), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_show", show_concat) + # cv2.waitKey() + + ret_img = (ret_img_fc32_new*255) + self.logger_process.info('fix hair costs:{}'.format(time.time() - t0)) + fix_image_dir = osp.join(userinfo_dir, taskid + '.jpg') + matting_after_path = osp.join(userinfo_dir, taskid + '.png') + cv2.imwrite(fix_image_dir, ret_img) + cv2.imwrite(matting_after_path, (ret_matting_mask_fc32*255).astype(np.uint8)) + + return fix_image_dir, 0 + + def fix_hairstyle_v2_old(self, origin_img, brush, taskid, userinfo_dir): + landmark1k_path = osp.join(userinfo_dir, 'kpt_1k.txt') + if os.path.exists(landmark1k_path): + landmarks_origin_img_1k = np.loadtxt(landmark1k_path) + else: + landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward(origin_img) + if landmarks_origin_img_1k is None: + return None, 10001 + + landmark137 = landmark_processor.pts_1k_to_137(landmarks_origin_img_1k) + user_rgb_8uc3_orisize = origin_img.copy() + + [ori_x, ori_y, dst_x, dst_y, radis] = brush + t0 = time.time() + ret_img = localtranslationwarpfastwithstrength_v2(user_rgb_8uc3_orisize, landmark137,ori_x, ori_y, dst_x, dst_y, radis) + cv2.circle(ret_img, (int(ori_x), int(ori_y)), 3, (0, 255, 0)) + cv2.circle(ret_img, (int(dst_x), int(dst_y)), 5, (0, 255, 0)) + self.logger_process.info('fix hair costs:{}'.format(time.time() - t0)) + fix_image_dir = osp.join(userinfo_dir, taskid + '.jpg') + cv2.imwrite(fix_image_dir, ret_img) + return fix_image_dir, 0 + + 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 get_gender(self, img_ori): + ref_landmark_1k2_f_orisize = self.get_landmark_mtcnn.forward(img_ori) + if ref_landmark_1k2_f_orisize is None: + return None + ref_landmark_137kpts_f_orisize = landmark_processor.pts_1k_to_137(ref_landmark_1k2_f_orisize) + is_female = self.gender_model.forward(img_ori, ref_landmark_137kpts_f_orisize) + return is_female + + def get_prepare_ref_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_landmark_137kpts_f_orisize = landmark_processor.pts_1k_to_137(ref_landmark_1k2_f_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, 255, 255), -1) + # cv2.imshow('input_img', ref_rgb_8uc3_orisize) + # cv2.waitKey() + gender_res = self.gender_model.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 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 infer_hairstyle_diy(self, user_rgb_8uc3_orisize, ref_rgb_8uc3_orisize, userinfo_dir, refer_name): + refer_dir = config.get('default', 'refer_dir') + tmp_refer_dir = osp.join(refer_dir, refer_name.split('.')[0]) + os.makedirs(tmp_refer_dir, exist_ok=True) + # userinfo_dir = osp.join(userinfo_dir, 'diy') + os.makedirs(userinfo_dir, exist_ok=True) + with torch.no_grad(): + ref_landmark_1k2_f_orisize, _, _ = self.get_landmark.forward_diy(ref_rgb_8uc3_orisize) + ref_landmark_137kpts_f_orisize = landmark_processor.pts_1k_to_137(ref_landmark_1k2_f_orisize) + gender_res = self.gender_model.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 gender_res: + gender = "girl" + else: + gender = "boy" + # ref_rgb_8uc3_768_dir = osp.join(tmp_refer_dir, 'ref_rgb_8uc3_768.png') + # ref_matting_fg_8uc3_768_dir = osp.join(tmp_refer_dir, 'ref_matting_fg_8uc3_768.png') + # ref_matting_8uc3_768_dir = osp.join(tmp_refer_dir, 'ref_matting_8uc3_768.png') + # ref_baldseg_8uc3_768_dir = osp.join(tmp_refer_dir, 'ref_baldseg_8uc3_768.png') + # ref_landmark_f1k2_768_dir = osp.join(tmp_refer_dir, 'ref_landmark_f1k2_768.txt') + another_pose_hair_image_dir = osp.join(tmp_refer_dir, 'another_pose_hair_image.png') + preLists = [another_pose_hair_image_dir] + condition_exist = True + for onedir in preLists: + if not osp.exists(onedir): + condition_exist = False + if condition_exist: + # ref_rgb_8uc3_768 = cv2.imread(ref_rgb_8uc3_768_dir) + # ref_matting_fg_8uc3_768 = cv2.imread(ref_matting_fg_8uc3_768_dir) + # ref_matting_8uc3_768 = cv2.imread(ref_matting_8uc3_768_dir) + # ref_baldseg_8uc3_768 = cv2.imread(ref_baldseg_8uc3_768_dir) + # ref_landmark_f1k2_768 = np.loadtxt(ref_landmark_f1k2_768_dir) + another_pose_hair_image = cv2.imread(another_pose_hair_image_dir) + another_pose_hair_image = another_pose_hair_image/255. + else: + 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) + another_pose_hair_image = self.Generator_reftensor(ref_rgb_8uc3_768, ref_matting_8uc3_768, + ref_baldseg_8uc3_768, + ref_landmark_f1k2_768) + cv2.imwrite(another_pose_hair_image_dir, another_pose_hair_image*255) + landmark1k_dir = osp.join(userinfo_dir, 'kpt_1k.txt') + if not osp.exists(landmark1k_dir): + landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward_diy(user_rgb_8uc3_orisize) + if landmarks_origin_img_1k is None: + return None, 10001 + np.savetxt(landmark1k_dir, landmarks_origin_img_1k) + else: + landmarks_origin_img_1k = np.loadtxt(landmark1k_dir) + # landmarks_origin_img_1k, _, _ = self.get_landmark.forward(user_rgb_8uc3_orisize) + if landmarks_origin_img_1k is None: + return None, 10001 + + user_bald_res_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_res_ori.png') + user_baldseg_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_seg_ori.png') + user_baldseg_8uc3_768_dir = osp.join(userinfo_dir, 'user_baldseg_768.png') + user_bald_8uc3_768_dir = osp.join(userinfo_dir, 'bald_seg_768.png') + user_landmark_f1k2_768_dir = osp.join(userinfo_dir, 'landmark_f1k2_768.txt') + user_hairstyle_M_dir = osp.join(userinfo_dir, 'hairstyle_M.txt') + + condition_exist2 = True + pre_list = [user_bald_res_8uc3_orisize_dir, user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_768_dir, user_bald_8uc3_768_dir, + user_landmark_f1k2_768_dir, user_hairstyle_M_dir] + for tmp_dir in pre_list: + if not osp.exists(tmp_dir): + condition_exist2 = False + + user_matting_8uc3_bald_orisize = None + if not condition_exist2: + user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, user_baldseg_8uc3_768, user_bald_8uc3_768, \ + user_landmark_f1k2_768, user_hairstyle_M, user_matting_8uc3_bald_orisize = self.process_data.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio) + cv2.imwrite(user_bald_res_8uc3_orisize_dir, user_bald_res_8uc3_orisize) + cv2.imwrite(user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_orisize) + cv2.imwrite(user_baldseg_8uc3_768_dir, user_baldseg_8uc3_768) + cv2.imwrite(user_bald_8uc3_768_dir, user_bald_8uc3_768) + np.savetxt(user_landmark_f1k2_768_dir, user_landmark_f1k2_768) + np.savetxt(user_hairstyle_M_dir, user_hairstyle_M) + else: + user_bald_res_8uc3_orisize = cv2.imread(user_bald_res_8uc3_orisize_dir) + user_baldseg_8uc3_orisize = cv2.imread(user_baldseg_8uc3_orisize_dir) + user_baldseg_8uc3_768 = cv2.imread(user_baldseg_8uc3_768_dir) + user_bald_8uc3_768 = cv2.imread(user_bald_8uc3_768_dir) + user_landmark_f1k2_768 = np.loadtxt(user_landmark_f1k2_768_dir) + user_hairstyle_M = np.loadtxt(user_hairstyle_M_dir) + + # show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("show_concat", show_concat) + # cv2.waitKey() + + user_orig_mask_path = os.path.join(userinfo_dir, "user_orig_mask.png") + if not os.path.exists(user_orig_mask_path): + cv2.imwrite(user_orig_mask_path, user_matting_8uc3_bald_orisize) + + # 换发型 + hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image, + user_baldseg_8uc3_768, + user_bald_8uc3_768, + user_landmark_f1k2_768, gender) + + + hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste( + user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M) + + gen_hair_mask_path_2 = os.path.join(userinfo_dir, "hair_mask_2.png") + cv2.imwrite(gen_hair_mask_path_2, hair_gene_matte_8uc3_orisize) + + res_matting_mask_ori = osp.join(userinfo_dir, 'res_matting_mask_ori_fix.png') + res_matting_mask_ori_raw = osp.join(userinfo_dir, 'res_matting_mask_ori_raw.png') + if osp.exists(res_matting_mask_ori): + os.remove(res_matting_mask_ori) + if osp.exists(res_matting_mask_ori_raw): + os.remove(res_matting_mask_ori_raw) + # cv2.imwrite(res_matting_mask_ori, hair_gene_matte_8uc3_orisize) + cv2.imwrite(res_matting_mask_ori_raw, hair_gene_matte_8uc3_orisize) + + res_hairstyle_before_8uc3 = user_rgb_8uc3_orisize.copy() + # res_hairstyle_before_8uc3 = cv2.imread('/home/data/hair/data/userImage/zrn/d77957536f6b688bc678c3bb4e5095c3/d77957536f6b688bc678c3bb4e5095c3.jpg') + hair_fix_img_mask_8uc4 = np.concatenate((res_hairstyle_before_8uc3, hair_gene_matte_8uc3_orisize[:, :, :1]), axis=2) + + # cv2.imshow('res_hairstyle_before_8uc3', res_hairstyle_before_8uc3) + # cv2.imshow('hair_gene_matte_8uc3_orisize', hair_gene_matte_8uc3_orisize[:, :, :1]) + # cv2.waitKey() + + cv2.imwrite(res_matting_mask_ori, hair_fix_img_mask_8uc4) + + user_res_8uc3_orisize = self.hair_fusion.inference(hair_gene_fusion_8uc3_orisize, + hair_gene_matte_8uc3_orisize, + landmarks_origin_img_1k, + user_baldseg_8uc3_orisize, + ratio) + + if self.use_enhance: + user_res_8uc3_orisize = self.face_enhance.process(user_res_8uc3_orisize, landmarks_origin_img_1k) + + return user_res_8uc3_orisize, 0, gender + + def infer_hairstyle_diy_jy(self, user_rgb_8uc3_orisize, ref_rgb_8uc3_orisize, userinfo_dir, refer_name): + refer_dir = config.get('default', 'refer_dir') + tmp_refer_dir = osp.join(refer_dir, refer_name.split('.')[0]) + os.makedirs(tmp_refer_dir, exist_ok=True) + # userinfo_dir = osp.join(userinfo_dir, 'diy') + os.makedirs(userinfo_dir, exist_ok=True) + with torch.no_grad(): + ref_landmark_1k2_f_orisize, _, _ = self.get_landmark.forward_diy(ref_rgb_8uc3_orisize) + ref_landmark_137kpts_f_orisize = landmark_processor.pts_1k_to_137(ref_landmark_1k2_f_orisize) + gender_res = self.gender_model.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 gender_res: + gender = "girl" + else: + gender = "boy" + # ref_rgb_8uc3_768_dir = osp.join(tmp_refer_dir, 'ref_rgb_8uc3_768.png') + # ref_matting_fg_8uc3_768_dir = osp.join(tmp_refer_dir, 'ref_matting_fg_8uc3_768.png') + # ref_matting_8uc3_768_dir = osp.join(tmp_refer_dir, 'ref_matting_8uc3_768.png') + # ref_baldseg_8uc3_768_dir = osp.join(tmp_refer_dir, 'ref_baldseg_8uc3_768.png') + # ref_landmark_f1k2_768_dir = osp.join(tmp_refer_dir, 'ref_landmark_f1k2_768.txt') + another_pose_hair_image_dir = osp.join(tmp_refer_dir, 'another_pose_hair_image.png') + preLists = [another_pose_hair_image_dir] + condition_exist = True + for onedir in preLists: + if not osp.exists(onedir): + condition_exist = False + if condition_exist: + # ref_rgb_8uc3_768 = cv2.imread(ref_rgb_8uc3_768_dir) + # ref_matting_fg_8uc3_768 = cv2.imread(ref_matting_fg_8uc3_768_dir) + # ref_matting_8uc3_768 = cv2.imread(ref_matting_8uc3_768_dir) + # ref_baldseg_8uc3_768 = cv2.imread(ref_baldseg_8uc3_768_dir) + # ref_landmark_f1k2_768 = np.loadtxt(ref_landmark_f1k2_768_dir) + another_pose_hair_image = cv2.imread(another_pose_hair_image_dir) + another_pose_hair_image = another_pose_hair_image/255. + else: + 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) + another_pose_hair_image = self.Generator_reftensor(ref_rgb_8uc3_768, ref_matting_8uc3_768, + ref_baldseg_8uc3_768, + ref_landmark_f1k2_768) + cv2.imwrite(another_pose_hair_image_dir, another_pose_hair_image*255) + landmark1k_dir = osp.join(userinfo_dir, 'kpt_1k.txt') + if not osp.exists(landmark1k_dir): + landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward_diy(user_rgb_8uc3_orisize) + if landmarks_origin_img_1k is None: + return None, 10001, None, None, None + np.savetxt(landmark1k_dir, landmarks_origin_img_1k) + else: + landmarks_origin_img_1k = np.loadtxt(landmark1k_dir) + # landmarks_origin_img_1k, _, _ = self.get_landmark.forward(user_rgb_8uc3_orisize) + if landmarks_origin_img_1k is None: + return None, 10001, None, None, None + + user_bald_res_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_res_ori.png') + user_baldseg_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_seg_ori.png') + user_baldseg_8uc3_768_dir = osp.join(userinfo_dir, 'user_baldseg_768.png') + user_bald_8uc3_768_dir = osp.join(userinfo_dir, 'bald_seg_768.png') + user_landmark_f1k2_768_dir = osp.join(userinfo_dir, 'landmark_f1k2_768.txt') + user_hairstyle_M_dir = osp.join(userinfo_dir, 'hairstyle_M.txt') + + condition_exist2 = True + pre_list = [user_bald_res_8uc3_orisize_dir, user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_768_dir, user_bald_8uc3_768_dir, + user_landmark_f1k2_768_dir, user_hairstyle_M_dir] + for tmp_dir in pre_list: + if not osp.exists(tmp_dir): + condition_exist2 = False + + user_matting_8uc3_bald_orisize = None + if not condition_exist2: + user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, user_baldseg_8uc3_768, user_bald_8uc3_768, \ + user_landmark_f1k2_768, user_hairstyle_M, user_matting_8uc3_bald_orisize = self.process_data.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio) + cv2.imwrite(user_bald_res_8uc3_orisize_dir, user_bald_res_8uc3_orisize) + cv2.imwrite(user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_orisize) + cv2.imwrite(user_baldseg_8uc3_768_dir, user_baldseg_8uc3_768) + cv2.imwrite(user_bald_8uc3_768_dir, user_bald_8uc3_768) + np.savetxt(user_landmark_f1k2_768_dir, user_landmark_f1k2_768) + np.savetxt(user_hairstyle_M_dir, user_hairstyle_M) + else: + user_bald_res_8uc3_orisize = cv2.imread(user_bald_res_8uc3_orisize_dir) + user_baldseg_8uc3_orisize = cv2.imread(user_baldseg_8uc3_orisize_dir) + user_baldseg_8uc3_768 = cv2.imread(user_baldseg_8uc3_768_dir) + user_bald_8uc3_768 = cv2.imread(user_bald_8uc3_768_dir) + user_landmark_f1k2_768 = np.loadtxt(user_landmark_f1k2_768_dir) + user_hairstyle_M = np.loadtxt(user_hairstyle_M_dir) + + # show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("show_concat", show_concat) + # cv2.waitKey() + + user_orig_mask_path = os.path.join(userinfo_dir, "user_orig_mask.png") + if not os.path.exists(user_orig_mask_path): + cv2.imwrite(user_orig_mask_path, user_matting_8uc3_bald_orisize) + + # 换发型 + hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image, + user_baldseg_8uc3_768, + user_bald_8uc3_768, + user_landmark_f1k2_768, gender) + + + hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste( + user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M) + + gen_hair_mask_path_2 = os.path.join(userinfo_dir, "hair_mask_2.png") + cv2.imwrite(gen_hair_mask_path_2, hair_gene_matte_8uc3_orisize) + + res_matting_mask_ori = osp.join(userinfo_dir, 'res_matting_mask_ori_fix.png') + res_matting_mask_ori_raw = osp.join(userinfo_dir, 'res_matting_mask_ori_raw.png') + if osp.exists(res_matting_mask_ori): + os.remove(res_matting_mask_ori) + if osp.exists(res_matting_mask_ori_raw): + os.remove(res_matting_mask_ori_raw) + # cv2.imwrite(res_matting_mask_ori, hair_gene_matte_8uc3_orisize) + cv2.imwrite(res_matting_mask_ori_raw, hair_gene_matte_8uc3_orisize) + + res_hairstyle_before_8uc3 = user_rgb_8uc3_orisize.copy() + # res_hairstyle_before_8uc3 = cv2.imread('/home/data/hair/data/userImage/zrn/d77957536f6b688bc678c3bb4e5095c3/d77957536f6b688bc678c3bb4e5095c3.jpg') + hair_fix_img_mask_8uc4 = np.concatenate((res_hairstyle_before_8uc3, hair_gene_matte_8uc3_orisize[:, :, :1]), axis=2) + + # cv2.imshow('res_hairstyle_before_8uc3', res_hairstyle_before_8uc3) + # cv2.imshow('hair_gene_matte_8uc3_orisize', hair_gene_matte_8uc3_orisize[:, :, :1]) + # cv2.waitKey() + + cv2.imwrite(res_matting_mask_ori, hair_fix_img_mask_8uc4) + + user_res_8uc3_orisize = self.hair_fusion.inference(hair_gene_fusion_8uc3_orisize, + hair_gene_matte_8uc3_orisize, + landmarks_origin_img_1k, + user_baldseg_8uc3_orisize, + ratio) + + if self.use_enhance: + user_res_8uc3_orisize = self.face_enhance.process(user_res_8uc3_orisize, landmarks_origin_img_1k) + + return user_res_8uc3_orisize, 0, gender, landmarks_origin_img_1k, False + + def change_image_color_for_tiaoran(self, image, hair_mask_8uc3, color): + # b, g, r = color # [10, 50, 250] # [10, 250, 10] + r, g, b = color # [10, 50, 250] # [10, 250, 10] + + tar_color = np.zeros_like(image) + tar_color[:, :, 0] = b + tar_color[:, :, 1] = g + tar_color[:, :, 2] = r + image_hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) + tar_hsv = cv2.cvtColor(tar_color, cv2.COLOR_BGR2HSV) + + loc_index = hair_mask_8uc3[:, :, 0].nonzero() + color_val = image_hsv[loc_index] + if len(color_val) == 0: return None + # mean_v = np.mean(color_val[:, 2]) + mean_v = np.mean(color_val, axis=0) + # print("mean_v: ", mean_v) + # print("tar_hsv: ", tar_hsv[0, 0, :]) + + image_hsv[:, :, 1:2] = mean_v[1] + (image_hsv[:, :, 1:2] - mean_v[1]) * 0.5 + image_hsv[:, :, 2:3] = mean_v[2] + (image_hsv[:, :, 2:3] - mean_v[2]) * 0.2 + + image_hsv[:, :, 0:1] = tar_hsv[:, :, 0:1] + ratio_s = 0.3 + image_hsv[:, :, 1:2] = tar_hsv[:, :, 1:2] * ratio_s + image_hsv[:, :, 1:2] * (1 - ratio_s) + image_hsv[:, :, 2:3] = tar_hsv[:, :, 2:3] * ratio_s + image_hsv[:, :, 2:3] * (1 - ratio_s) + image_hsv[:, :, 1:2] = np.clip(image_hsv[:, :, 1:2] * 1.1, 0, 255) + image_hsv[:, :, 2:3] = np.clip(image_hsv[:, :, 2:3], 10, 245) + + img_res = cv2.cvtColor(image_hsv.astype(np.uint8), cv2.COLOR_HSV2BGR) + # changed_sharpen = sharpen(changed) + return img_res + + def infer_hairtiaoran_v3(self, user_rgb_8uc3_orisize, dst_color, draw_mask, userinfo_dir, userMask, taskid, hairId): + + if hairId == '' or hairId is None: + hairId = 'source' + user_base_color_img_path = os.path.join(userinfo_dir, hairId, "user_base_color_8uc3_orisize.png") + mask_newname = '{}_t.png'.format(taskid) + if userMask is None or userMask == '': + userMask = mask_newname + if os.path.exists(user_base_color_img_path): + face_base = cv2.imread(user_base_color_img_path) + status1 = 0 + else: + baseColor_dir = os.path.join(config.get('default', "haircolorDir"), config.get('default', "baseColor_ID")) + # baseColor_dir = os.path.join("/home/yangchaojie/Desktop/hairstyle/hairstyle_infer/data/ref_haircolor", config.get('default', "baseColor_ID")) + face_base, mask_newname, status1 = self.infer_haircolor(user_rgb_8uc3_orisize, baseColor_dir, userinfo_dir, userMask) + if not os.path.exists(os.path.dirname(user_base_color_img_path)): + os.makedirs(os.path.dirname(user_base_color_img_path)) + cv2.imwrite(user_base_color_img_path, face_base) + # user_matting_mask_path = os.path.join(userinfo_dir, "user_matting_mask_8uc3_orisize.png") + user_matting_mask_path = os.path.join(userinfo_dir, userMask) + user_matting_mask_path_new = os.path.join(userinfo_dir, mask_newname) + + landmark1k_path = osp.join(userinfo_dir, 'kpt_1k.txt') + if os.path.exists(landmark1k_path): + landmarks_origin_img_1k = np.loadtxt(landmark1k_path) + else: + landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward(user_rgb_8uc3_orisize) + if not os.path.exists(os.path.dirname(landmark1k_path)): + os.makedirs(os.path.dirname(landmark1k_path)) + np.savetxt(landmark1k_path, landmarks_origin_img_1k) + + if os.path.exists(user_matting_mask_path): + hair_fix_img_mask_8uc4 = cv2.imread(user_matting_mask_path, -1) + user_matting_mask_8uc1_orisize = hair_fix_img_mask_8uc4[:, :, 3:] + user_matting_mask_8uc3_orisize = np.repeat(user_matting_mask_8uc1_orisize, 3, axis=2) + # user_matting_mask_8uc3_orisize = cv2.imread(user_matting_mask_path, -1) + else: + _, user_matting_8uc1_bald_orisize, _ = self.process_data.generator_matte.matte_inference(user_rgb_8uc3_orisize, + landmarks_origin_img_1k) + user_matting_mask_8uc3_orisize = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) + if not os.path.exists(os.path.dirname(user_matting_mask_path)): + os.makedirs(os.path.dirname(user_matting_mask_path)) + + # res_hairstyle_before_8uc3 = user_rgb_8uc3_orisize.copy() + # hair_fix_img_mask_8uc4 = np.concatenate((res_hairstyle_before_8uc3, user_matting_mask_8uc3_orisize[:, :, :1]), axis=2) + # # cv2.imshow('res_hairstyle_before_8uc3', res_hairstyle_before_8uc3) + # # cv2.imshow('user_matting_mask_8uc3_orisize', user_matting_mask_8uc3_orisize) + # # cv2.waitKey() + # cv2.imwrite(user_matting_mask_path, hair_fix_img_mask_8uc4) + + if status1 == 0: + face_base_HSV = cv2.cvtColor(face_base, cv2.COLOR_BGR2HSV) + user_rgb_8uc3_orisize_HSV = cv2.cvtColor(user_rgb_8uc3_orisize, cv2.COLOR_BGR2HSV) + mix_ratio = 0.5 + face_base_HSV[:, :, 2:3] = user_rgb_8uc3_orisize_HSV[:, :, 2:3] * mix_ratio + face_base_HSV[:, :, 2:3] * (1 - mix_ratio) + face_base_new = cv2.cvtColor(face_base_HSV, cv2.COLOR_HSV2BGR) + + ret_color_img = self.change_image_color_for_tiaoran(face_base_new, user_matting_mask_8uc3_orisize, dst_color) + if ret_color_img is None: + return user_rgb_8uc3_orisize, 1 + user_matting_mask_fc32_orisize = user_matting_mask_8uc3_orisize.astype(np.float32) / 255 + + ret_color_img_LAB = cv2.cvtColor(ret_color_img, cv2.COLOR_BGR2LAB) + user_rgb_8uc3_orisize_LAB = cv2.cvtColor(user_rgb_8uc3_orisize, cv2.COLOR_BGR2LAB) + face_base_new_LAB = cv2.cvtColor(face_base_new, cv2.COLOR_BGR2LAB) + ret_color_img_LAB[:, :, 0] = face_base_new_LAB[:, :, 0] * user_matting_mask_fc32_orisize[:, :, 0] + user_rgb_8uc3_orisize_LAB[:, :, 0] * ( + 1 - user_matting_mask_fc32_orisize[:, :, 0]) + ret_color_img = cv2.cvtColor(ret_color_img_LAB, cv2.COLOR_LAB2BGR) + + face_bbox = cv2.boundingRect(landmarks_origin_img_1k.astype(np.int32)[np.newaxis, :, :]) + face_max_len = max(face_bbox[2], face_bbox[3]) + dilate_kernel_size = int(0.010 * face_max_len) + if dilate_kernel_size % 2 == 0: + dilate_kernel_size += 1 + # print("dilate_kernel_size: ", dilate_kernel_size) + blur_kernel_size = int(0.061 * face_max_len) + if blur_kernel_size % 2 == 0: + blur_kernel_size += 1 + # print("blur_kernel_size: ", blur_kernel_size) + gaussionblur_kernel_size = int(0.051 * face_max_len) + if gaussionblur_kernel_size % 2 == 0: + gaussionblur_kernel_size += 1 + # print("gaussionblur_kernel_size: ", gaussionblur_kernel_size) + ret_color_img_fc32 = ret_color_img.astype(np.float32) / 255 + user_rgb_8uc3_orisize_fc32 = user_rgb_8uc3_orisize.astype(np.float32) / 255 + + new_color = [0, 0, 0] + color_lists = [[0, 0, 0], [255, 255, 255]] + t0 = time.time() + new_mask = np.ones_like(draw_mask) * 255 + r_img, g_img, b_img = draw_mask[:, :, 0].copy(), draw_mask[:, :, 1].copy(), draw_mask[:, :, 2].copy() + for scr_color in color_lists: + new_mask[r_img == scr_color[0]] = new_color[0] + new_mask[g_img == scr_color[1]] = new_color[1] + new_mask[b_img == scr_color[2]] = new_color[2] + + draw_mask_fc32 = new_mask.astype(np.float32) / 255 + # cv2.imshow('draw_mask', draw_mask) + # + # cv2.imshow('new_mask', new_mask) + # cv2.imshow('draw_mask_fc32', draw_mask_fc32) + # cv2.waitKey() + # draw_mask_fc32 = cv2.dilate(draw_mask_fc32, np.ones((dilate_kernel_size, dilate_kernel_size), np.uint8), iterations=1) + draw_mask_fc32_erode = cv2.erode(draw_mask_fc32, np.ones((dilate_kernel_size, dilate_kernel_size), np.uint8), iterations=1) + draw_mask_fc32_erode = cv2.GaussianBlur(draw_mask_fc32_erode, (3, 3), 0, 0) + + # cv2.imshow("draw_mask_fc32_erode", draw_mask_fc32_erode) + # cv2.waitKey() + draw_mask_fc32_dilate = cv2.dilate(draw_mask_fc32_erode, np.ones((7, 7), np.uint8), iterations=1) + add_in_mask = draw_mask_fc32_dilate - draw_mask_fc32_erode + + draw_mask_fc32_blur = cv2.GaussianBlur(add_in_mask, (blur_kernel_size, blur_kernel_size), 0, 0) + draw_mask_fc32 = draw_mask_fc32_blur + draw_mask_fc32_erode + draw_mask_fc32 = cv2.GaussianBlur(draw_mask_fc32, (blur_kernel_size, blur_kernel_size), 0, 0) + + # draw_mask_fc32_blur = cv2.GaussianBlur(draw_mask_fc32_erode, (gaussionblur_kernel_size, gaussionblur_kernel_size), 0, 0) + # draw_mask_fc32 = cv2.GaussianBlur(draw_mask_fc32_blur, (blur_kernel_size, blur_kernel_size), 0, 0) + # draw_mask_fc32 = draw_mask_fc32 + draw_mask_fc32_blur + draw_mask_fc32 = np.clip(draw_mask_fc32, 0, 1) + # cv2.imshow("add_in_mask", add_in_mask) + # cv2.waitKey() + draw_mask_fc32 = draw_mask_fc32 * user_matting_mask_fc32_orisize + + ret_img_fc32 = ret_color_img_fc32 * draw_mask_fc32 + user_rgb_8uc3_orisize_fc32 * (1 - draw_mask_fc32) + + # mid_show = np.concatenate((ret_color_img_fc32, user_rgb_8uc3_orisize_fc32, draw_mask_fc32, ret_img_fc32), axis=1) + # resize_ratio = 1280. / max(mid_show.shape[:2]) + # mid_show = cv2.resize(mid_show, (0, 0), fx=resize_ratio, fy=resize_ratio) + # cv2.imshow("mid_show", mid_show) + # cv2.waitKey(0) + + ret_img = (ret_img_fc32*255).astype(np.uint8) + if osp.exists(user_matting_mask_path): + hair_fix_img_mask_8uc4[:,:,:3] = ret_img + else: + hair_fix_img_mask_8uc4 = np.concatenate((ret_img, user_matting_mask_8uc3_orisize[:, :, :1]), axis=2) + + + cv2.imwrite(user_matting_mask_path, hair_fix_img_mask_8uc4) + return ret_img, 0 + else: + return user_rgb_8uc3_orisize, 1 + + def infer_matting(self, user_rgb_8uc3_orisize): + landmarks_origin_img_1k, bounding_boxes, euler_info = self.get_landmark.forward(user_rgb_8uc3_orisize) + if landmarks_origin_img_1k is None: + return None, 10001 + user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, user_hair_color_M, user_matting_8uc3_orisize = \ + self.process_data.get_prepare_hair_color_user_data(user_rgb_8uc3_orisize, + landmarks_origin_img_1k) + + user_matting_8uc3_change_color_orisize = cv2.warpAffine(user_matting_8uc3_change_color_768, + cv2.invertAffineTransform(user_hair_color_M), + (user_rgb_8uc3_orisize.shape[1], + user_rgb_8uc3_orisize.shape[0]), flags=cv2.INTER_CUBIC) + + # return (user_rgb_8uc3_orisize * (user_matting_8uc3_change_color_orisize/255)).astype(np.uint8), user_rgb_8uc3_orisize + + return user_matting_8uc3_change_color_orisize, user_rgb_8uc3_orisize + + def getavgstd(self, image, mask): + if mask.shape[2] == 1: + mask = np.repeat(mask, 3, axis=2) + mask_index = np.flatnonzero((mask > 0.1).any(axis=2)) + lab_layer = image.reshape(-1, 3)[mask_index] + if len(lab_layer) < 100: return None + lab_layer = np.float32(lab_layer) + avg = [] + std = [] + image_avg_l = np.mean(lab_layer[:, 0]) + image_std_l = np.std(lab_layer[:, 0]) + image_avg_a = np.mean(lab_layer[:, 1]) + image_std_a = np.std(lab_layer[:, 1]) + image_avg_b = np.mean(lab_layer[:, 2]) + image_std_b = np.std(lab_layer[:, 2]) + avg.append(image_avg_l) + avg.append(image_avg_a) + avg.append(image_avg_b) + std.append(image_std_l) + std.append(image_std_a) + std.append(image_std_b) + return avg, std + def reinhard_rgb(self, origin_img, mask, tar_avg_std, ratio=0.5): + # origin_img_lab = cv2.cvtColor(origin_img, cv2.COLOR_BGR2LAB) + origin_img_lab = origin_img + src_avg_std = self.getavgstd(origin_img_lab, mask) + src_avg_std = np.float32(src_avg_std) + + origin_img_lab[:, :, 0] = origin_img_lab[:, :, 0] + (tar_avg_std[0][0] - src_avg_std[0][0]) * ratio + origin_img_lab[:, :, 1] = origin_img_lab[:, :, 1] + (tar_avg_std[0][1] - src_avg_std[0][1]) * ratio + origin_img_lab[:, :, 2] = origin_img_lab[:, :, 2] + (tar_avg_std[0][2] - src_avg_std[0][2]) * ratio + # img_ret = cv2.cvtColor(origin_img_lab, cv2.COLOR_LAB2BGR) + img_ret = origin_img_lab + img_ret = np.clip(img_ret, 0.0, 1.0) + return img_ret, src_avg_std + + def adjust_color_toolbar(self,source_image, userinfo_dir, hairId, ratio_v, userMask, mask_newname, facebase): + user_base_color_img_path = os.path.join(userinfo_dir, hairId, facebase) + if os.path.exists(user_base_color_img_path): + face_base = cv2.imread(user_base_color_img_path) + else: + print('!!!!!!!!!! not exists:{}'.format(user_base_color_img_path)) + return source_image, 1 + # else: + # baseColor_dir = os.path.join(config.get('default', "haircolorDir"), colorId) + # face_base, mask_newname, status1 = self.infer_haircolor(source_image, baseColor_dir, userinfo_dir, userMask) + # if not os.path.exists(os.path.dirname(user_base_color_img_path)): + # os.makedirs(os.path.dirname(user_base_color_img_path)) + # cv2.imwrite(user_base_color_img_path, face_base) + ret_color_img = source_image * (1 - ratio_v) + face_base * ratio_v + # cv2.imshow('source_image', source_image) + # cv2.imshow('face_base', face_base) + # cv2.imshow('source_image', ret_color_img.astype(np.uint8)) + # cv2.waitKey() + user_matting_mask_path = os.path.join(userinfo_dir, userMask) + user_matting_mask_path_new = os.path.join(userinfo_dir, mask_newname) + if os.path.exists(user_matting_mask_path): + hair_fix_img_mask_8uc4 = cv2.imread(user_matting_mask_path, -1) + hair_fix_img_mask_8uc4[:, :, :3] = ret_color_img + cv2.imwrite(user_matting_mask_path_new, hair_fix_img_mask_8uc4) + return ret_color_img, 0 + else: + print('!!!!!!!!!! not exists:{}'.format(user_matting_mask_path)) + return source_image, 1 + + def infer_haircolor_v3(self, user_rgb_8uc3_orisize, dst_color, userinfo_dir, userMask, mask_newname, hairId, + ratio_v=0.5): + if hairId == "" or hairId is None: + hairId = 'source' + user_base_color_img_path = os.path.join(userinfo_dir, hairId, "user_base_color_8uc3_orisize.png") + if userMask is None or userMask == '': + userMask = mask_newname + if os.path.exists(user_base_color_img_path): + face_base = cv2.imread(user_base_color_img_path) + status1 = 0 + else: + baseColor_dir = os.path.join(config.get('default', "haircolorDir"), config.get('default', "baseColor_ID")) + face_base, mask_newname, status1 = self.infer_haircolor(user_rgb_8uc3_orisize, baseColor_dir, userinfo_dir, + userMask) + if not os.path.exists(os.path.dirname(user_base_color_img_path)): + os.makedirs(os.path.dirname(user_base_color_img_path)) + cv2.imwrite(user_base_color_img_path, face_base) + user_matting_mask_path = os.path.join(userinfo_dir, userMask) + user_matting_mask_path_new = os.path.join(userinfo_dir, mask_newname) + + landmark1k_path = osp.join(userinfo_dir, 'kpt_1k.txt') + if os.path.exists(landmark1k_path): + landmarks_origin_img_1k = np.loadtxt(landmark1k_path) + else: + landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward(user_rgb_8uc3_orisize) + if not os.path.exists(os.path.dirname(landmark1k_path)): + os.makedirs(os.path.dirname(landmark1k_path)) + np.savetxt(landmark1k_path, landmarks_origin_img_1k) + + if os.path.exists(user_matting_mask_path): + hair_fix_img_mask_8uc4 = cv2.imread(user_matting_mask_path, -1) + user_matting_mask_8uc1_orisize = hair_fix_img_mask_8uc4[:, :, 3:] + user_matting_mask_8uc3_orisize = np.repeat(user_matting_mask_8uc1_orisize, 3, axis=2) + else: + _, user_matting_8uc1_bald_orisize, _ = self.process_data.generator_matte.matte_inference( + user_rgb_8uc3_orisize, + landmarks_origin_img_1k) + user_matting_mask_8uc3_orisize = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) + if not os.path.exists(os.path.dirname(user_matting_mask_path)): + os.makedirs(os.path.dirname(user_matting_mask_path)) + + if status1 == 0: + r, g, b = dst_color + # newColor = [b,g,r] + # b, g, r = dst_color # [10, 50, 250] # [10, 250, 10] + tar_color = np.zeros_like(user_rgb_8uc3_orisize) + tar_color[:, :, 0] = b + tar_color[:, :, 1] = g + tar_color[:, :, 2] = r + image_hsv = cv2.cvtColor(user_rgb_8uc3_orisize, cv2.COLOR_BGR2HSV) + tar_hsv = cv2.cvtColor(tar_color, cv2.COLOR_BGR2HSV) + + image_hsv[:, :, 0:1] = tar_hsv[:, :, 0:1] + image_hsv[:, :, 1:2] = tar_hsv[:, :, 1:2] + image_hsv[:, :, 2:3] = tar_hsv[:, :, 2:3] + + img_res = cv2.cvtColor(image_hsv.astype(np.uint8), cv2.COLOR_HSV2BGR) + origin_img_LAB = cv2.cvtColor(user_rgb_8uc3_orisize, cv2.COLOR_BGR2LAB) + img_res_LAB = cv2.cvtColor(img_res, cv2.COLOR_BGR2LAB) + img_res_LAB[:, :, 0] = origin_img_LAB[:, :, 0] + ret_color_img = cv2.cvtColor(img_res_LAB, cv2.COLOR_LAB2BGR) + + face_base_HSV = cv2.cvtColor(face_base, cv2.COLOR_BGR2HSV) + user_rgb_8uc3_orisize_HSV = cv2.cvtColor(user_rgb_8uc3_orisize, cv2.COLOR_BGR2HSV) + mix_ratio = 1 + face_base_HSV[:, :, 2:3] = user_rgb_8uc3_orisize_HSV[:, :, 2:3] * mix_ratio + face_base_HSV[:, :, 2:3] * ( + 1 - mix_ratio) + face_base_new = cv2.cvtColor(face_base_HSV, cv2.COLOR_HSV2BGR) + + # ret_color_img = self.change_image_color_for_tiaoran(face_base_new, user_matting_mask_8uc3_orisize, dst_color) + if ret_color_img is None: + return user_rgb_8uc3_orisize, 1 + user_matting_mask_fc32_orisize = user_matting_mask_8uc3_orisize.astype(np.float32) / 255 + + ret_color_img_LAB = cv2.cvtColor(ret_color_img, cv2.COLOR_BGR2LAB) + user_rgb_8uc3_orisize_LAB = cv2.cvtColor(user_rgb_8uc3_orisize, cv2.COLOR_BGR2LAB) + face_base_new_LAB = cv2.cvtColor(face_base_new, cv2.COLOR_BGR2LAB) + ret_color_img_LAB[:, :, 0] = face_base_new_LAB[:, :, 0] * user_matting_mask_fc32_orisize[:, :, + 0] + user_rgb_8uc3_orisize_LAB[:, :, 0] * ( + 1 - user_matting_mask_fc32_orisize[:, :, 0]) + ret_color_img = cv2.cvtColor(ret_color_img_LAB, cv2.COLOR_LAB2BGR) + + ret_img_fc32 = ret_color_img.astype( + np.float32) * user_matting_mask_fc32_orisize + user_rgb_8uc3_orisize.astype(np.float32) * ( + 1 - user_matting_mask_fc32_orisize) + ret_img = ret_img_fc32.astype(np.uint8) + + if osp.exists(user_matting_mask_path): + hair_fix_img_mask_8uc4[:, :, :3] = ret_img + else: + hair_fix_img_mask_8uc4 = np.concatenate((ret_img, user_matting_mask_8uc3_orisize[:, :, :1]), axis=2) + # cv2.imshow('res_hairstyle_before_8uc3', res_hairstyle_before_8uc3) + # cv2.imshow('user_matting_mask_8uc3_orisize', user_matting_mask_8uc3_orisize) + # cv2.waitKey() + # cv2.imshow('user_rgb_8uc3_orisize_fc32', user_rgb_8uc3_orisize) + # cv2.imshow('face_base_new', face_base_new) + # cv2.imshow('user_matting_mask_8uc3_orisize', user_matting_mask_8uc3_orisize) + # cv2.waitKey() + cv2.imwrite(user_matting_mask_path_new, hair_fix_img_mask_8uc4) + return ret_img, 0 + else: + return user_rgb_8uc3_orisize, 1 + def infer_haircolor_v4(self, user_rgb_8uc3_orisize, dst_color, userinfo_dir, userMask, mask_newname, hairId, ratio_v=0.9): + if hairId == "" or hairId is None: + hairId = 'source' + user_base_color_img_path = os.path.join(userinfo_dir, hairId, "user_base_color_8uc3_orisize.png") + user_now_color_img_path = os.path.join(userinfo_dir, hairId, mask_newname[:-4] + '_ac.jpg') + + if userMask is None or userMask == '': + userMask = mask_newname + if os.path.exists(user_base_color_img_path): + face_base = cv2.imread(user_base_color_img_path) + status1 = 0 + else: + baseColor_dir = os.path.join(config.get('default', "haircolorDir"), config.get('default', "baseColor_ID")) + face_base, mask_newname, status1 = self.infer_haircolor(user_rgb_8uc3_orisize, baseColor_dir, userinfo_dir, mask_newname) + if not os.path.exists(os.path.dirname(user_base_color_img_path)): + os.makedirs(os.path.dirname(user_base_color_img_path)) + cv2.imwrite(user_base_color_img_path, face_base) + user_matting_mask_path = os.path.join(userinfo_dir, userMask) + user_matting_mask_path_new = os.path.join(userinfo_dir, mask_newname) + + landmark1k_path = osp.join(userinfo_dir, 'kpt_1k.txt') + if os.path.exists(landmark1k_path): + landmarks_origin_img_1k = np.loadtxt(landmark1k_path) + else: + landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward(user_rgb_8uc3_orisize) + if not os.path.exists(os.path.dirname(landmark1k_path)): + os.makedirs(os.path.dirname(landmark1k_path)) + np.savetxt(landmark1k_path, landmarks_origin_img_1k) + + if os.path.exists(user_matting_mask_path): + hair_fix_img_mask_8uc4 = cv2.imread(user_matting_mask_path, -1) + user_matting_mask_8uc1_orisize = hair_fix_img_mask_8uc4[:, :, 3:] + user_matting_mask_8uc3_orisize = np.repeat(user_matting_mask_8uc1_orisize, 3, axis=2) + else: + _, user_matting_8uc1_bald_orisize, _ = self.process_data.generator_matte.matte_inference(user_rgb_8uc3_orisize, + landmarks_origin_img_1k) + user_matting_mask_8uc3_orisize = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) + if not os.path.exists(os.path.dirname(user_matting_mask_path)): + os.makedirs(os.path.dirname(user_matting_mask_path)) + + if status1 == 0: + r, g, b = dst_color + # newColor = [b,g,r] + # b, g, r = dst_color # [10, 50, 250] # [10, 250, 10] + tar_color = np.zeros_like(user_rgb_8uc3_orisize) + tar_color[:, :, 0] = b + tar_color[:, :, 1] = g + tar_color[:, :, 2] = r + image_hsv = cv2.cvtColor(user_rgb_8uc3_orisize, cv2.COLOR_BGR2HSV) + tar_hsv = cv2.cvtColor(tar_color, cv2.COLOR_BGR2HSV) + + image_hsv[:, :, 0:1] = tar_hsv[:, :, 0:1] + image_hsv[:, :, 1:2] = tar_hsv[:, :, 1:2] + image_hsv[:, :, 2:3] = tar_hsv[:, :, 2:3] + + img_res_change_hsv = cv2.cvtColor(image_hsv.astype(np.uint8), cv2.COLOR_HSV2BGR) + origin_img_LAB = cv2.cvtColor(user_rgb_8uc3_orisize, cv2.COLOR_BGR2LAB) + img_res_LAB = cv2.cvtColor(img_res_change_hsv, cv2.COLOR_BGR2LAB) + img_res_LAB[:, :, 0] = origin_img_LAB[:, :, 0] + ret_color_img = cv2.cvtColor(img_res_LAB, cv2.COLOR_LAB2BGR) + + face_base_HSV = cv2.cvtColor(face_base, cv2.COLOR_BGR2HSV) + user_rgb_8uc3_orisize_HSV = cv2.cvtColor(user_rgb_8uc3_orisize, cv2.COLOR_BGR2HSV) + mix_ratio = 1 + face_base_HSV[:, :, 2:3] = user_rgb_8uc3_orisize_HSV[:, :, 2:3] * mix_ratio + face_base_HSV[:, :, 2:3] * (1 - mix_ratio) + face_base_new = cv2.cvtColor(face_base_HSV, cv2.COLOR_HSV2BGR) + + user_matting_mask_fc32_orisize = user_matting_mask_8uc3_orisize.astype(np.float32) / 255 + + ### change ycj + user_rgb_8uc3_orisize_LAB = cv2.cvtColor(user_rgb_8uc3_orisize, cv2.COLOR_BGR2LAB) + img_res_change_hsv_LAB = cv2.cvtColor(img_res_change_hsv, cv2.COLOR_BGR2LAB) + mix_ratio = 0.98 + img_res_change_hsv_LAB[:, :, 0] = user_rgb_8uc3_orisize_LAB[:, :, 0] * mix_ratio + img_res_change_hsv_LAB[:, :, 0] * (1 - mix_ratio) + img_res_change_hsv_change_lab = cv2.cvtColor(img_res_change_hsv_LAB, cv2.COLOR_LAB2BGR) + + img_res_change_hsv_change_lab_fc32 = img_res_change_hsv_change_lab.astype(np.float32) / 255 + tar_color_fc32 = tar_color.astype(np.float32) / 255 + ref_rgb_avg_std = self.getavgstd(tar_color_fc32, user_matting_mask_fc32_orisize) + face_base_new_fc32_orisize_new, src_avg_std = self.reinhard_rgb(img_res_change_hsv_change_lab_fc32, user_matting_mask_fc32_orisize, ref_rgb_avg_std, + ratio=1.0) + face_base_new_fc32_orisize_new = face_base_new_fc32_orisize_new * user_matting_mask_fc32_orisize + \ + (user_rgb_8uc3_orisize.astype(np.float32)/255) * (1 - user_matting_mask_fc32_orisize) + face_base_new_new = (face_base_new_fc32_orisize_new * 255).astype(np.uint8) + cv2.imwrite(user_now_color_img_path, face_base_new_new) + # middle_show = np.concatenate((user_rgb_8uc3_orisize, face_base, face_base_new, face_base_new_new, tar_color, img_res_change_hsv_change_lab), axis=1) + # ratio = 1536. / max(middle_show.shape[:2]) + # middle_show = cv2.resize(middle_show, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow('middle_show', middle_show) + # cv2.waitKey(0) + ### change ycj + + # ret_color_img = self.change_image_color_for_tiaoran(face_base_new, user_matting_mask_8uc3_orisize, dst_color) + if ret_color_img is None: + return user_rgb_8uc3_orisize, user_now_color_img_path, None, 1 + + ret_color_img_LAB = cv2.cvtColor(ret_color_img, cv2.COLOR_BGR2LAB) + user_rgb_8uc3_orisize_LAB = cv2.cvtColor(user_rgb_8uc3_orisize, cv2.COLOR_BGR2LAB) + face_base_new_LAB = cv2.cvtColor(face_base_new, cv2.COLOR_BGR2LAB) + ret_color_img_LAB[:, :, 0] = face_base_new_LAB[:, :, 0] * user_matting_mask_fc32_orisize[:, :, 0] + user_rgb_8uc3_orisize_LAB[:, :, 0] * ( + 1 - user_matting_mask_fc32_orisize[:, :, 0]) + ret_color_img = cv2.cvtColor(ret_color_img_LAB, cv2.COLOR_LAB2BGR) + + ret_color_img_fc32 = ret_color_img.astype(np.float32) * user_matting_mask_fc32_orisize + user_rgb_8uc3_orisize.astype(np.float32) * (1 - user_matting_mask_fc32_orisize) + face_base_new_new_fc32 = face_base_new_new.astype(np.float32) * user_matting_mask_fc32_orisize + user_rgb_8uc3_orisize.astype(np.float32) * (1 - user_matting_mask_fc32_orisize) + ret_color_img = ret_color_img_fc32.astype(np.float32) * (1 - ratio_v) + face_base_new_new_fc32.astype(np.float32) * ratio_v + + # ret_color_img = ret_color_img.astype(np.float32) * (1 - ratio_v) + face_base_new_new.astype(np.float32) * ratio_v + # ret_img_fc32 = ret_color_img.astype(np.float32) * user_matting_mask_fc32_orisize + user_rgb_8uc3_orisize.astype(np.float32) * (1 - user_matting_mask_fc32_orisize) + # ret_img = ret_img_fc32.astype(np.uint8) + ret_img = ret_color_img.astype(np.uint8) + if osp.exists(user_matting_mask_path): + hair_fix_img_mask_8uc4[:,:,:3] = ret_img + else: + hair_fix_img_mask_8uc4 = np.concatenate((ret_img, user_matting_mask_8uc3_orisize[:, :, :1]), axis=2) + # cv2.imshow('res_hairstyle_before_8uc3', res_hairstyle_before_8uc3) + # cv2.imshow('user_matting_mask_8uc3_orisize', user_matting_mask_8uc3_orisize) + # cv2.waitKey() + # cv2.imshow('user_rgb_8uc3_orisize_fc32', user_rgb_8uc3_orisize) + # cv2.imshow('face_base_new', face_base_new) + # cv2.imshow('user_matting_mask_8uc3_orisize', user_matting_mask_8uc3_orisize) + # cv2.waitKey() + cv2.imwrite(user_matting_mask_path_new, hair_fix_img_mask_8uc4) + return ret_img, user_now_color_img_path, ret_color_img_fc32, 0 + else: + return user_rgb_8uc3_orisize, user_now_color_img_path, None, 1 + + def infer_haircolor_v2(self, user_rgb_8uc3_orisize, haircolor_dir, user_folder, maskImgName, mask_newname): + landmarks_origin_img_1k, bounding_boxes, euler_info = self.get_landmark.forward(user_rgb_8uc3_orisize) + if landmarks_origin_img_1k is None: + return None, 10001 + t0 = time.time() + _, user_matting_8uc1_bald_orisize, _ = self.process_data.generator_matte.matte_inference(user_rgb_8uc3_orisize, + landmarks_origin_img_1k) + user_matting_mask_8uc3_orisize = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) + + + ref_rgb_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_rgb_8uc3_color_768.png")) + ref_matting_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_matting_8uc3_color_768.png")) + + user_res_8uc3_orisize = user_rgb_8uc3_orisize.copy() + + # user_rgb_8uc3_orisize_LAB = cv2.cvtColor(user_rgb_8uc3_orisize, cv2.COLOR_BGR2LAB) + # user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) + # user_res_8uc3_orisize_LAB[:, :, 0] = user_rgb_8uc3_orisize_LAB[:, :, 0] + # user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_LAB, cv2.COLOR_LAB2BGR) + + user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32) / 255 + # user_res_fc32_orisize_hsv = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2HSV) + ref_rgb_fc32_change_color_768 = ref_rgb_8uc3_change_color_768.astype(np.float32) / 255 + ref_matting_fc32_change_color_768 = ref_matting_8uc3_change_color_768.astype(np.float32) / 255 + + ref_rgb_avg_std = self.getavgstd(ref_rgb_fc32_change_color_768, ref_matting_fc32_change_color_768) + user_matting_mask_fc32_orisize = user_matting_mask_8uc3_orisize.astype(np.float32) / 255 + user_res_fc32_orisize, src_avg_std = self.reinhard_rgb(user_res_fc32_orisize, user_matting_mask_fc32_orisize, + ref_rgb_avg_std, ratio=1.0) + + user_res_8uc3_orisize = (user_res_fc32_orisize * 255).astype(np.uint8) + + user_res_8uc3_orisize2 = user_rgb_8uc3_orisize * (1 - user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize * ( + user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize2 = (np.clip(user_res_8uc3_orisize2, 0, 255)).astype(np.uint8) + + # middle_show = np.concatenate((user_rgb_8uc3_orisize, user_matting_8uc3_orisize, user_res_8uc3_orisize, user_res_8uc3_orisize2), axis=1) + # ratio = 1536. / max(middle_show.shape[:2]) + # middle_show = cv2.resize(middle_show, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow('middle_show', middle_show) + # cv2.waitKey(0) + + if maskImgName is None or maskImgName == '': + res_hairstyle_before_8uc3 = user_res_8uc3_orisize2.copy() + hair_fix_img_mask_8uc4 = np.concatenate((res_hairstyle_before_8uc3, user_matting_mask_8uc3_orisize[:, :, :1]), + axis=2) + else: + umg_img_dir = osp.join(user_folder, maskImgName) + if osp.exists(umg_img_dir): + hair_fix_img_mask_8uc4 = cv2.imread(umg_img_dir, -1) + hair_fix_img_mask_8uc4[:,:,:3] = user_res_8uc3_orisize2 + else: + res_hairstyle_before_8uc3 = user_res_8uc3_orisize2.copy() + hair_fix_img_mask_8uc4 = np.concatenate((res_hairstyle_before_8uc3, user_matting_mask_8uc3_orisize[:, :, :1]), + axis=2) + # user_res_8uc3_orisize_enhance = self.face_enhance.process(user_res_8uc3_orisize2, landmarks_origin_img_1k) + + cv2.imwrite(os.path.join(user_folder, mask_newname), hair_fix_img_mask_8uc4) + # print('change color, last proces,', time.time() - t3) + return user_res_8uc3_orisize2, 0 + + def infer_haircolor(self, user_rgb_8uc3_orisize, haircolor_dir, user_folder, maskImgName): + landmarks_origin_img_1k, bounding_boxes, euler_info = self.get_landmark.forward(user_rgb_8uc3_orisize) + if landmarks_origin_img_1k is None: + return None, maskImgName, 10001 + t0 = time.time() + _, user_matting_8uc1_bald_orisize, _ = self.process_data.generator_matte.matte_inference(user_rgb_8uc3_orisize, + landmarks_origin_img_1k) + user_matting_mask_8uc3_orisize = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) + + + ref_rgb_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_rgb_8uc3_color_768.png")) + ref_matting_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_matting_8uc3_color_768.png")) + + user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, user_hair_color_M, user_matting_8uc3_orisize = \ + self.process_data.get_prepare_hair_color_user_data(user_rgb_8uc3_orisize, + landmarks_origin_img_1k) + + # start_time = time.time() + hair_gene_color_8uc3_768, user_matting_mask_8uc3_768 = self.change_haircolor.Change_Hair_inference( + user_rgb_8uc3_change_color_768, + user_matting_8uc3_change_color_768, + ref_rgb_8uc3_change_color_768, + ref_matting_8uc3_change_color_768) + + hair_gene_color_8uc3_orisize = cv2.warpAffine(hair_gene_color_8uc3_768, + cv2.invertAffineTransform(user_hair_color_M), + (user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), + dst=user_rgb_8uc3_orisize.copy(), + flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_TRANSPARENT) + + user_res_8uc3_orisize = self.face_enhance.process(hair_gene_color_8uc3_orisize, landmarks_origin_img_1k) + # user_res_8uc3_orisize = hair_gene_color_8uc3_orisize.copy() + + # user_rgb_8uc3_orisize_LAB = cv2.cvtColor(user_rgb_8uc3_orisize, cv2.COLOR_BGR2LAB) + # user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) + # user_res_8uc3_orisize_LAB[:, :, 0] = user_rgb_8uc3_orisize_LAB[:, :, 0] + # user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_LAB, cv2.COLOR_LAB2BGR) + + user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32) / 255 + # user_res_fc32_orisize_hsv = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2HSV) + ref_rgb_fc32_change_color_768 = ref_rgb_8uc3_change_color_768.astype(np.float32) / 255 + ref_matting_fc32_change_color_768 = ref_matting_8uc3_change_color_768.astype(np.float32) / 255 + + ref_rgb_avg_std = self.getavgstd(ref_rgb_fc32_change_color_768, ref_matting_fc32_change_color_768) + user_matting_mask_fc32_orisize = user_matting_mask_8uc3_orisize.astype(np.float32) / 255 + user_res_fc32_orisize, src_avg_std = self.reinhard_rgb(user_res_fc32_orisize, user_matting_mask_fc32_orisize, + ref_rgb_avg_std, ratio=1.0) + + user_res_8uc3_orisize_reinhard = (user_res_fc32_orisize * 255).astype(np.uint8) + + user_res_8uc3_orisize_reinhard_LAB = cv2.cvtColor(user_res_8uc3_orisize_reinhard, cv2.COLOR_BGR2LAB) + user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) + mix_ratio = 0.1 + user_res_8uc3_orisize_reinhard_LAB[:, :, 0] = user_res_8uc3_orisize_LAB[:, :, 0] * mix_ratio + user_res_8uc3_orisize_reinhard_LAB[:, :, 0] * ( + 1 - mix_ratio) + user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_reinhard_LAB, cv2.COLOR_LAB2BGR) + + user_res_8uc3_orisize2 = user_rgb_8uc3_orisize * (1 - user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize * ( + user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize2 = (np.clip(user_res_8uc3_orisize2, 0, 255)).astype(np.uint8) + + # middle_show = np.concatenate((user_rgb_8uc3_orisize, user_matting_8uc3_orisize, user_res_8uc3_orisize, user_res_8uc3_orisize2), axis=1) + # ratio = 1536. / max(middle_show.shape[:2]) + # middle_show = cv2.resize(middle_show, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow('middle_show', middle_show) + # cv2.waitKey(0) + + if maskImgName is None or maskImgName == '': + maskImgName = 'res_matting_mask_ori_fix.png' + res_hairstyle_before_8uc3 = user_res_8uc3_orisize2.copy() + hair_fix_img_mask_8uc4 = np.concatenate((res_hairstyle_before_8uc3, user_matting_mask_8uc3_orisize[:, :, :1]), + axis=2) + else: + umg_img_dir = osp.join(user_folder, maskImgName) + if osp.exists(umg_img_dir): + hair_fix_img_mask_8uc4 = cv2.imread(umg_img_dir, -1) + hair_fix_img_mask_8uc4[:,:,:3] = user_res_8uc3_orisize2 + else: + res_hairstyle_before_8uc3 = user_res_8uc3_orisize2.copy() + hair_fix_img_mask_8uc4 = np.concatenate((res_hairstyle_before_8uc3, user_matting_mask_8uc3_orisize[:, :, :1]), + axis=2) + # user_res_8uc3_orisize_enhance = self.face_enhance.process(user_res_8uc3_orisize2, landmarks_origin_img_1k) + + cv2.imwrite(os.path.join(user_folder, maskImgName), hair_fix_img_mask_8uc4) + # print('change color, last proces,', time.time() - t3) + return user_res_8uc3_orisize2, maskImgName, 0 + + def infer_bald(self, user_rgb_8uc3_orisize): + landmarks_origin_img_1k = self.get_landmark.forward(user_rgb_8uc3_orisize) + if landmarks_origin_img_1k is None: + return None, 10001 + # 返回光头 + user_res_8uc3_orisize, user_inter_res_8uc3_orisize, user_matting_8uc3_bald_768, user_baldseg_8uc3_bald_768 = self.process_data.get_user_blad( + user_rgb_8uc3_orisize, + landmarks_origin_img_1k) + if self.use_enhance: + user_res_8uc3_orisize = self.face_enhance.process(user_res_8uc3_orisize, landmarks_origin_img_1k) + return user_res_8uc3_orisize, 0 + + import cv2 + import numpy as np + def drawline(self, img, pt1, pt2, color, thickness=1, style='dotted', gap=10): + dist = ((pt1[0] - pt2[0]) ** 2 + (pt1[1] - pt2[1]) ** 2) ** .5 + pts = [] + for i in np.arange(0, dist, gap): + r = i / dist + x = int((pt1[0] * (1 - r) + pt2[0] * r) + .5) + y = int((pt1[1] * (1 - r) + pt2[1] * r) + .5) + p = (x, y) + pts.append(p) + + if style == 'dotted': + for p in pts: + cv2.circle(img, p, thickness, color, -1) + else: + s = pts[0] + e = pts[0] + i = 0 + for p in pts: + s = e + e = p + if i % 2 == 1: + cv2.line(img, s, e, color, thickness) + i += 1 + + def define_line(self, kpt1, kpt2, loc_x): + scale = (kpt2[1] - kpt1[1])/(kpt2[0] - kpt1[0]) + b = kpt1[1] - scale*kpt1[0] + loca_y = scale * loc_x + b + return loca_y + + def get_3t_5y(self, in_frame, landmark_137): + infoLists = {} + kpt137 = landmark_137.copy() + landmark_137 = landmark_137.astype('int').tolist() + left_eye_in = landmark_137[88] + left_eye_out = landmark_137[96] + right_eye_in = landmark_137[106] + right_eye_out = landmark_137[113] + face_left = landmark_137[16] + face_right = landmark_137[6] + face_top = landmark_137[11] + face_bottom = landmark_137[0] + eyebrow_left = landmark_137[131] + eyebrow_right = landmark_137[123] + nose_left = landmark_137[69] + nose_right = landmark_137[73] + eyebrow_height = int(eyebrow_left[1]/2 + eyebrow_right[1]/2) + length_face_high = face_bottom[1] - face_top[1] + length_face_width = face_right[0] - face_left[0] + eye_5_1_val = round((left_eye_out[0] - face_left[0])/length_face_width, 2) + eye_5_2_val = round((left_eye_in[0] - left_eye_out[0])/length_face_width, 2) + eye_5_3_val = round((right_eye_in[0] - left_eye_in[0])/length_face_width, 2) + eye_5_4_val = round((right_eye_out[0] - right_eye_in[0])/length_face_width, 2) + eye_5_5_val = round((1 - eye_5_1_val - eye_5_2_val - eye_5_3_val - eye_5_4_val), 2) + + ting_3_1_val = round((eyebrow_right[1] - face_top[1])/length_face_high, 2) + ting_3_2_val = round((nose_right[1] - eyebrow_right[1])/length_face_high, 2) + ting_3_3_val = round(1 - ting_3_1_val - ting_3_2_val, 2) + infoLists['ratios_st'] = [] + infoLists['ratios_wy'] = [] + infoLists['lines_st'] = [] + infoLists['lines_wy'] = [] + + infoLists['ratios_st'].append([ting_3_1_val, ting_3_2_val, ting_3_3_val]) + infoLists['ratios_wy'].append([eye_5_1_val, eye_5_2_val, eye_5_3_val, eye_5_4_val, eye_5_5_val]) + + + + # infoLists['lines_wy'].append([face_left[0], face_bottom[1], face_left[0], face_top[1]]) + infoLists['lines_wy'].append([left_eye_out[0], face_bottom[1], left_eye_out[0], face_top[1]]) + infoLists['lines_wy'].append([left_eye_in[0], face_bottom[1], left_eye_in[0], face_top[1]]) + infoLists['lines_wy'].append([right_eye_in[0], face_bottom[1], right_eye_in[0], face_top[1]]) + infoLists['lines_wy'].append([right_eye_out[0], face_bottom[1], right_eye_out[0], face_top[1]]) + # infoLists['lines_wy'].append([face_right[0], face_bottom[1], face_right[0], face_top[1]]) + + + black_image = np.zeros_like(in_frame) + in_frame_black = (black_image*0.4 + in_frame*0.6).astype(np.uint8) + # 五眼划虚线 + # self.drawline(in_frame, [face_right[0], face_bottom[1]], [face_right[0], face_top[1]], (255,0,0)) + cv2.line(in_frame_black, [right_eye_out[0], face_bottom[1]], [right_eye_out[0], face_top[1]], (255,255,255), thickness=1) + # self.drawline(in_frame, [face_left[0], face_bottom[1]], [face_left[0], face_top[1]], (255,0,0)) + cv2.line(in_frame_black, [left_eye_out[0], face_bottom[1]], [left_eye_out[0], face_top[1]], (255,255,255), thickness=1) + cv2.line(in_frame_black, [left_eye_in[0], face_bottom[1]], [left_eye_in[0], face_top[1]], (255,255,255), thickness=1) + cv2.line(in_frame_black, [right_eye_in[0], face_bottom[1]], [right_eye_in[0], face_top[1]], (255,255,255), thickness=1) + + expand_val = 50 + # 三庭画虚线 + loc_y1 = self.define_line([nose_left[0], nose_left[1]],[nose_right[0], nose_right[1]], face_left[0]-expand_val) + loc_y2 = self.define_line([nose_left[0], nose_left[1]],[nose_right[0], nose_right[1]], face_right[0] + expand_val) + + loc_y3 = self.define_line([eyebrow_left[0], eyebrow_left[1]],[eyebrow_right[0], eyebrow_right[1]], face_left[0]-expand_val) + loc_y4 = self.define_line([eyebrow_left[0], eyebrow_left[1]], [eyebrow_right[0], eyebrow_right[1]], face_right[0] + expand_val) + # self.drawline(in_frame, [face_left[0], face_top[1]], [face_right[0], face_top[1]], (255, 0, 0)) + self.drawline(in_frame_black, [face_left[0]-expand_val, loc_y1], [face_right[0]+expand_val, loc_y2], (255,255,255) , thickness=1,style='line') + self.drawline(in_frame_black, [face_left[0]-expand_val, loc_y3], [face_right[0]+expand_val, loc_y4], (255,255,255), thickness=1 ,style='line') + # self.drawline(in_frame, [face_left[0], face_bottom[1]], [face_right[0], face_bottom[1]], (255, 0, 0)) + + # infoLists['lines_st'].append([face_left[0], face_top[1], face_right[0], face_top[1]]) # 3t1 + infoLists['lines_st'].append([face_left[0]-expand_val, loc_y1, face_right[0]+expand_val, loc_y2]) # 3t2 + infoLists['lines_st'].append([face_left[0]-expand_val, loc_y3, face_right[0]+expand_val, loc_y4]) # 3t3 + # infoLists['lines_st'].append([face_left[0], face_bottom[1], face_right[0], face_bottom[1]]) # 3t4 + + + img_PIL = Image.fromarray(cv2.cvtColor(in_frame_black, cv2.COLOR_BGR2RGB)) + # font = ImageFont.truetype('./core/simsun.ttc', 15, encoding="UTF-8") + font = ImageFont.truetype('./core/PingFangMedium.ttf', 15, encoding="UTF-8") + + draw = ImageDraw.Draw(img_PIL) + # draw.text((65, 130), "上庭: {}".format(ting_3_1_val), font=font, fill=(255, 255,255)) + # draw.text((65, 230), "中庭: {}".format(ting_3_2_val), font=font, fill=(255, 255,255)) + # draw.text((65, 330), "下庭: {}".format(ting_3_3_val), font=font, fill=(255, 255,255)) + draw.text((65, 120), "上庭: {}".format(round(ting_3_1_val/ting_3_2_val, 1)), font=font, fill=(255, 255,255)) + draw.text((65, 230), "中庭: {}".format(1), font=font, fill=(255, 255,255)) + draw.text((65, 330), "下庭: {}".format(round(ting_3_3_val/ting_3_2_val, 1)), font=font, fill=(255, 255,255)) + + # draw.text((50, 390), "五眼比例: {}".format(eye_5_1_val), font=font, fill=(255, 255,255)) + # draw.text((int(landmark_137[87][0]-10), 390), "{}".format(eye_5_2_val), font=font, fill=(255, 255,255)) + # draw.text((int(landmark_137[86][0]-10), 390), "{}".format(eye_5_3_val), font=font, fill=(255, 255,255)) + # draw.text((int(landmark_137[104][0]-10), 390), "{}".format(eye_5_4_val), font=font, fill=(255, 255,255)) + # draw.text((face_right[0], 390), "{}".format(eye_5_5_val), font=font, fill=(255, 255,255)) + draw.text((50, 390), "五眼比例: {}".format(round(eye_5_1_val/eye_5_4_val, 1)), font=font, fill=(255, 255,255)) + draw.text((int(landmark_137[87][0]-10), 390), "{}".format(1), font=font, fill=(255, 255,255)) + draw.text((int(landmark_137[86][0]-10), 390), "{}".format(round(ting_3_3_val/eye_5_4_val, 1)), font=font, fill=(255, 255,255)) + draw.text((int(landmark_137[104][0]-10), 390), "{}".format(1), font=font, fill=(255, 255,255)) + draw.text((face_right[0], 390), "{}".format(round(eye_5_5_val/eye_5_4_val, 1)), font=font, fill=(255, 255,255)) + + chin_cls = self.chin_cls.process_face(in_frame_black, kpt137) + ret = self.get_landmark.get_face_shape(chin_cls, kpt137) + + face_shape_map = {'chang':"长形", 'fang':"方形", 'yuan':"圆形", 'tuoyuan':"椭圆形", 'xin':"心形"} + + draw.text((210, 420), "脸型: {}".format(face_shape_map[ret]), font=font, fill=(255, 255,255)) + + final_show = cv2.cvtColor(np.asarray(img_PIL), cv2.COLOR_RGB2BGR) + # cv2.imshow('final_show', final_show) + # cv2.waitKey() + return infoLists, final_show, ret + + def quality_control(self, input_img, landmark137, seg_mask): + ret_info = { + 'code': 200, + 'data': '', + 'msg': '' + } + clear_mask = np.zeros_like(input_img) + cv2.fillPoly(clear_mask, np.concatenate( + (landmark137[96:88:-1], landmark137[105:114], landmark137[2::-1], landmark137[21:19:-1]))[np.newaxis, :, :], + (255, 255, 255)) + mask_del = (np.clip(seg_mask * clear_mask, 0, 255)).astype(np.uint8) + # for i in range(137): + # cv2.circle(input_img, (int(landmark137[i][0]), int(landmark137[i][1])), 1, (255, 255, 255), -1) + # cv2.imshow('input_img', input_img) + # + # cv2.imshow('fff', mask_del) + # cv2.imshow('seg_mask', seg_mask) + # cv2.imshow('clear_mask', clear_mask) + # + # cv2.waitKey() + scale = (mask_del / 255).sum() / (clear_mask / 255).sum() + if scale < 0.9: + ret_info['code'] = '400' + ret_info['msg'] = '请勿遮挡人脸五官区域' + return ret_info + h, w, _ = input_img.shape + face_left = landmark137[16] + face_right = landmark137[6] + face_top = landmark137[11] + face_bottom = landmark137[0] + face_length = face_right[0] - face_left[0] + face_height = face_bottom[1] - face_top[1] + if face_length/w > 0.7 or face_height/h > 0.7: + ret_info['code'] = '400' + ret_info['msg'] = '请适当远离镜头' + return ret_info + if face_length/w < 0.1 or face_height/h < 0.1: + ret_info['code'] = '400' + ret_info['msg'] = '请适当靠近镜头' + return ret_info + return ret_info + + def buff(self, img, img_skin, value1, value2): + img = img.astype(np.float32) + dx = value1 * 5 + fc = value1 * 12.5 + p = 80 + temp1 = cv2.bilateralFilter(img, dx, fc, fc) + temp2 = (temp1 - img + 128) + temp2 = np.clip(temp2, 0, 255) + temp3 = cv2.GaussianBlur(temp2, (2 * value2 - 1, 2 * value2 - 1), 0, 0) + temp4 = img + 2 * temp3 - 255 + temp4 = np.clip(temp4, 0, 255) + dst = img * ((100 - p) / 100) + temp4 * (p / 100) + img_skin_c = 1-img_skin + dst = dst * img_skin + img * img_skin_c + return dst.astype(np.uint8) + + def face_dermabrasion(self, cut_image, mask): + # 创建待处理图像 + t0 = time.time() + pending_image = cv2.bitwise_and(cut_image, cut_image, mask=mask) + # 双边滤波 + blur_img = cv2.bilateralFilter(pending_image, 15, 60, 60) + # 图像融合 + fusion_img = cv2.addWeighted(pending_image, 0.5, blur_img, 0.5, 0) + + not_mask = 255 - mask + back_img = cv2.bitwise_and(fusion_img, fusion_img, mask=mask) + # cv2.imshow('back_img',back_img) + front_img = cv2.bitwise_and(cut_image, cut_image, mask=not_mask) + # cv2.imshow('front_img',front_img) + result_img = cv2.add(back_img, front_img) + return result_img + + def whitening(self, img, img_skin, value): + midtones_add = np.zeros(256) + for i in range(256): + midtones_add[i] = 0.667 * (1 - ((i - 127) / 127) * ((i - 127) / 127)) + lookup = np.zeros(256, dtype='uint8') + for i in range(256): + red = i + red += value * midtones_add[red] + red = max(0, red) + lookup[i] = np.uint(red) + w, h, c = img.shape + + img_skin = img_skin[:,:,-1] + index = np.where(img_skin == 1) + for i in range(index[0].shape[0]): + img[index[0][i], index[1][i], 0] = lookup[int(img[index[0][i], index[1][i], 0])] + img[index[0][i], index[1][i], 1] = lookup[int(img[index[0][i], index[1][i], 1])] + img[index[0][i], index[1][i], 2] = lookup[int(img[index[0][i], index[1][i], 2])] + + # for i in range(w): + # for j in range(h): + # if img_skin[i, j, 0] == 1: + # img[i, j, 0] = lookup[img[i, j, 0]] + # img[i, j, 1] = lookup[img[i, j, 1]] + # img[i, j, 2] = lookup[img[i, j, 2]] + return img + + def infer_face(self, origin_img, umd, savePath, quality_control=False, isLocal=1): + + tmp_info = { + 'umd':umd, + 'maskinfo':'' + } + + ret_info = { + 'state': 0, + 'face': {}, + 'msg': '', + 'umd':json.dumps(tmp_info) + } + get_facecolor = False + landmarks_origin_img_1k = self.get_landmark_mtcnn.forward(origin_img) + + if landmarks_origin_img_1k is None: + ret_info['state'] = 1 + ret_info['msg'] = '无法检测到人脸' + return ret_info + landmark_137 = landmark_processor.pts_1k_to_137(landmarks_origin_img_1k).astype(np.int) + # for i in range(137): + # + # cv2.circle(origin_img, (int(landmark_137[i][0]), int(landmark_137[i][1])), 1, (255, 255, 255), -1) + # cv2.imshow('srcimg', origin_img) + # cv2.waitKey() + isfemale = self.gender_model.forward(origin_img, landmark_137) + if isfemale: + gender_type = 'female' + else: + gender_type = 'male' + # if quality_control: + t01 = time.time() + face_mask_ori = self.face_seg.inference(origin_img, landmarks_origin_img_1k) + self.logger_process.info("face info , face_seg costs: {}s".format(time.time() - t01)) + + t02 = time.time() + ret = self.quality_control(origin_img, landmark_137, face_mask_ori) + self.logger_process.info("face info , quality_control costs: {}s".format(time.time() - t01)) + if ret['code'] != 200: + ret_info['state'] = 1 + ret_info['msg'] = ret['msg'] + return ret_info + + datanow = datetime.now() + time_convert = datanow.strftime("%Y%m%d%H") + taskid = ''.join(str(random.choice(range(10))) for _ in range(10)) + taskid = str(time_convert) + str(taskid) + res_dir = os.path.join(savePath, taskid + '_1.jpg') + res_dir2 = os.path.join(savePath, taskid + '_2.jpg') + + if int(isLocal) == 0: + # if get_facecolor: + h, w, _ = origin_img.shape + inpaint_mask = np.zeros((h, w), dtype=np.uint8) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[129:137]), (255,)) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[121:129]), (255,)) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[22:48]), (255,)) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[88:104]), (255,)) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[105:121]), (255,)) + inpaint_mask = cv2.cvtColor(inpaint_mask, cv2.COLOR_GRAY2BGR) + face_mask_del = np.clip(face_mask_ori - (inpaint_mask/255).astype(np.float32), 0, 1) + # face_image_del = (origin_img*face_mask_del).astype(np.uint8) + # face_mask_del = np.around(face_mask_del) + t0 = time.time() + dst_img_white = self.whitening(origin_img, face_mask_del, 10) + print('whitening', time.time() - t0) + t1 = time.time() + # bounding_boxs = bounding_boxs.astype(np.int) + face_left = landmark_137[16] + face_right = landmark_137[6] + face_top = landmark_137[11] + face_bottom = landmark_137[0] + x1, y1, x2, y2 = max(face_left[0]-20, 0), max(face_top[1]-20, 0), min(face_right[0]+20, w), min(face_bottom[1]+20, h) + img_beauty = dst_img_white.copy() + face_mask_c2 = np.ones((h, w), dtype=np.uint8) + mask_tmp = face_mask_c2[y1:y2, x1:x2] + img_tmp = origin_img[y1:y2, x1:x2] + + dst_img_white = self.face_dermabrasion(img_tmp, mask_tmp * 255) + + # dst_img_buff = self.buff(dst_img_white[y1:y2, x1:x2], face_mask_del[y1:y2, x1:x2], 2, 3) + img_beauty[y1:y2,x1:x2] = dst_img_white + print('buff', time.time() - t1) + + h, w, _ = img_beauty.shape + + + # dst_img_white = cv2.resize(dst_img_white, (int(w/2), int(h/2))) + # + # resize_dst = cv2.resize(dst_img_buff, (int(w/2), int(h/2))) + # resize_img = cv2.resize(origin_img, (int(w/2), int(h/2))) + # img_beauty = cv2.resize(img_beauty, (int(w/2), int(h/2))) + + # cv2.imshow('dst_img_white', dst_img_white) + # cv2.imshow('resize_img', resize_img) + # cv2.imshow('face_mask_del', resize_mask*255) + # cv2.imshow('img_beauty', img_beauty) + # cv2.waitKey() + else: + img_beauty = origin_img + M_hair = landmark_processor.get_transform_mat_full_face_592(landmark_137, 512, 0.7) + face_image_hair = cv2.warpAffine(origin_img, M_hair, (512, 512)) + kpt137_hair = landmark_processor.transform_points(landmark_137, M_hair) + t3y5_info, show_image, face_shape = self.get_3t_5y(face_image_hair, kpt137_hair) + + cv2.imwrite(res_dir, show_image) + cv2.imwrite(res_dir2, img_beauty) + t0 = time.time() + userId = osp.basename(savePath) + ret_url = self.oss2.upload_file(res_dir, "hair_mz/images/{}/{}".format(userId, taskid + '_1.jpg')) + ret_url2 = self.oss2.upload_file(res_dir2, "hair_mz/images/{}/{}".format(userId, taskid + '_2.jpg')) + + self.logger_process.info("face info upload oss costs: {}s".format(time.time() - t0)) + print('upload img costs:', time.time() - t0) + ret_info['face']['sanTing'] = t3y5_info + ret_info['face']['drawImg'] = ret_url + ret_info['face']['beautyImg'] = ret_url2 + ret_info['face']['faceType'] = face_shape + ret_info['face']['gender'] = gender_type + ret_info['face']['kpt137'] = landmark_137.tolist() + return ret_info + + 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 + + def prepare_single_hairstyle(self, input_info): + version = "20221206" + hairstyle_dir = '' + ret_data = {} + ret_data['code'] = 200 + ret_data['data'] = [] + ret_data['msg'] = 'success' + if len(input_info) == 0: + ret_data['code'] = 300 + ret_data['data'] = [] + ret_data['msg'] = 'no input info' + return ret_data + for single_info in input_info: + source_image, cover_image, degree, dst_dir = single_info + + if str(degree) == '1': + try: + ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768, gender, ratio = \ + self.get_prepare_ref_768_data(source_image) + + cv2.imwrite(os.path.join(dst_dir, "ref_rgb_8uc3_768.png"), ref_rgb_8uc3_768) + cv2.imwrite(os.path.join(dst_dir, "ref_matting_fg_8uc3_768.png"), ref_matting_fg_8uc3_768) + cv2.imwrite(os.path.join(dst_dir, "ref_matting_8uc3_768.png"), ref_matting_8uc3_768) + cv2.imwrite(os.path.join(dst_dir, "ref_baldseg_8uc3_768.png"), ref_baldseg_8uc3_768) + np.savetxt(os.path.join(dst_dir, "ref_landmark_f1k2_768.txt"), ref_landmark_f1k2_768) + input_another_pose_hair_image = self.Generator_reftensor(ref_rgb_8uc3_768, + ref_matting_8uc3_768, + ref_baldseg_8uc3_768, + ref_landmark_f1k2_768) + + np.save(os.path.join(dst_dir, "input_another_pose_hair_image.npy"), input_another_pose_hair_image) + + config_dict = {} + config_dict['gender'] = gender + config_dict['version'] = version + config_dict['ratio'] = str(ratio) + with open(os.path.join(dst_dir, "config.json"), "w") as f: + json.dump(config_dict, f) + print("写入文件完成...", dst_dir) + hairstyle_dir = dst_dir + + user_rgb_8uc3_orisize = cover_image.copy() + landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward(cover_image) + try: + user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, \ + user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M,_ = self.process_data.get_prepare_user_768_data( + user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio) + except Exception as e: + print(e) + ret_data['code'] = 300 + ret_data['data'] = [] + ret_data['msg'] = 'you should give an front face' + return ret_data + + another_pose_hair_img_path = os.path.join(hairstyle_dir, "input_another_pose_hair_image.npy") + if not os.path.exists(another_pose_hair_img_path): + ret_data['code'] = 300 + ret_data['data'] = [] + ret_data['msg'] = 'input_another_pose_hair_image not exists' + return ret_data + another_pose_hair_image = np.load(another_pose_hair_img_path) + + # 换发型 + hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image, + user_baldseg_8uc3_768, + user_bald_8uc3_768, + user_landmark_f1k2_768, + gender) + + hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste( + user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M) + + user_res_8uc3_orisize = self.hair_fusion.inference(hair_gene_fusion_8uc3_orisize, + hair_gene_matte_8uc3_orisize, + landmarks_origin_img_1k, + user_baldseg_8uc3_orisize, + ratio) + + if self.use_enhance: + user_res_8uc3_orisize = self.face_enhance.process(user_res_8uc3_orisize, + landmarks_origin_img_1k) + + img_name = ''.join(str(random.choice(range(10))) for _ in range(7)) + img_name = img_name + '.jpg' + cv2.imwrite(os.path.join(dst_dir, img_name), user_res_8uc3_orisize) + ret_url = self.oss2.upload_file(os.path.join(dst_dir, img_name), + "hair_mz/images/{}/{}".format(os.path.basename(dst_dir), img_name)) + ret_data['data'].append(ret_url) + + except Exception as e: + print(e) + ret_data['code'] = 300 + ret_data['data'] = [] + ret_data['msg'] = 'degree 1 deal bad' + return ret_data + return ret_data + + + def prepare_single_hairstyle_v2(self, input_info): + ret_data = {} + ret_data['code'] = 200 + ret_data['data'] = [] + ret_data['msg'] = 'success' + if len(input_info) == 0: + ret_data['code'] = 300 + ret_data['data'] = [] + ret_data['msg'] = 'no input info' + return ret_data + + save_cover_img = None + ref_user_save_dir = None + train_material_save_dir = None + hair_input_img_lists = [] + + k = 0 + + for single_info in input_info: + source_image, cover_image, degree, dst_dir, ref_user_dir, train_material_dir = single_info + if ref_user_save_dir == None: + ref_user_save_dir = ref_user_dir + + if train_material_save_dir == None: + train_material_save_dir = train_material_dir + + if k == 0: + suffix = source_image[:source_image.rfind("/")] + file_name = source_image[source_image.rfind("/") + 1:] + file_name = "first##" + file_name + save_new_path = os.path.join(suffix, file_name) + os.rename(source_image, save_new_path) + source_image = save_new_path + k += 1 + + if save_cover_img == None: + save_cover_img = cover_image + hair_input_img_lists.append(source_image) + print("hair_input_img_lists:", hair_input_img_lists) + + # if save_cover_img != None: + # hair_input_img_lists.append(save_cover_img) + + ref_user_dir = ref_user_save_dir + + # tag + tag_prompt = "" + + try: + # 1. 遍历用户输入的发型图，产出 pkl 和 matting 图 + hairId = osp.basename(input_info[0][-1]) + hair_template_save_dir = os.path.join(hair_template_material_dir, hairId) + if not os.path.exists(hair_template_save_dir): + os.makedirs(hair_template_save_dir) + + for i in range(0, len(hair_input_img_lists)): + img_path = hair_input_img_lists[i] + file_name = img_path[img_path.rfind("/") + 1:] + new_file_path = os.path.join(hair_template_save_dir, file_name) + new_file_path = new_file_path.replace("webp", "png") + + # 获取性别 + img_ori = cv2.imread(img_path) + gender = "girl" + is_female = self.get_gender(img_ori) + if not is_female: + gender = "boy" + + # 高分辨率图 + high_img_path = get_high_train_img(img_path, gender) + tmp_img = cv2.imread(high_img_path) + # cv2.imshow("tmp img", tmp_img) + # cv2.waitKey(0) + + shutil.copy(high_img_path, new_file_path) + + pkl_process(hair_template_save_dir) + print("hair_template_save_dir:", hair_template_save_dir) + + ref_user_img_list = glob.glob(os.path.join(ref_user_dir, '*.png')) + print("--------ref_user_img_list:", ref_user_img_list) + + len_user_imgs = len(ref_user_img_list) + len_hair_imgs = len(hair_input_img_lists) + index = 200 // (len_user_imgs * len_hair_imgs * 5) + if index == 0: + index = 1 + + # 2. 产出训练数据 + for root, dirs, files in os.walk(hair_template_save_dir): + for file in files: + if file.endswith('.png') or file.endswith('.jpg') or file.endswith('.jpeg') or file.endswith('.PNG') or file.endswith('.JPG') or file.endswith('.webp'): + img_path = os.path.join(root, file) + img_name = img_path.split(".")[0] + print("img_name:", img_name) + matting_path = img_name + '_matting.png' + if not os.path.exists(matting_path): + continue + pkl_path = img_name + '.pkl' + if not os.path.exists(pkl_path): + continue + hair_img = cv2.imread(img_path).astype(np.float32) / 255 + matting = cv2.imread(matting_path, cv2.IMREAD_GRAYSCALE).astype(np.float32) / 255 + + with open(pkl_path, 'rb') as fp: + data = pickle.load(fp) + pt1k = data['human_pt1k'] + + for user_img_path in ref_user_img_list: + ref_pt1k_path = user_img_path[:-4] + '.pkl' + if not os.path.exists(ref_pt1k_path): + continue + with open(ref_pt1k_path, 'rb') as fp: + ref_data = pickle.load(fp) + ref_pt1k = ref_data['human_pt1k'] + + ref_img = cv2.imread(user_img_path) + print("user_img_path ref_img:", user_img_path) + + max_size = max(ref_img.shape[0], ref_img.shape[1]) + M_fill = cv2.getRotationMatrix2D((ref_img.shape[1] / 2, ref_img.shape[0] / 2), 0, 1) + M_fill[:, 2] += np.float32([max_size, max_size]) / 2 - np.float32( + [ref_img.shape[1], ref_img.shape[0]]) / 2 + + M = landmark_processor.get_transform_mat_full_face_to_target(pt1k, ref_pt1k) + + # 将两个M矩阵级联 + M_final = M_fill @ np.concatenate([M, np.array([[0, 0, 1]])], axis=0) + hair_img_new = landmark_processor.high_quality_warpAffine(hair_img, M_final, + (max_size, max_size)) + matting_new = landmark_processor.high_quality_warpAffine(matting, M_final, + (max_size, max_size), + const_value=(0, 0, 0)) + # matting_new = cv2.warpAffine(matting, M_final, (max_size, max_size), flags=cv2.INTER_LANCZOS4) + bg = np.ones_like(hair_img_new) * 1.0 + final_img = hair_img_new * matting_new[:, :, None] + bg * (1 - matting_new[:, :, None]) + final_img = np.clip(final_img, 0, 1) + final_img = (final_img * 255).astype(np.uint8) + + # cv2.imshow('final_img', final_img) + # cv2.imshow('matting_new', matting_new) + # cv2.waitKey() + + # 创建输出目录 + data_path = train_material_save_dir + if not os.path.exists(data_path): + os.makedirs(data_path) + + images_path = os.path.join(data_path, "images") + if not os.path.exists(images_path): + os.makedirs(images_path) + + model_path = os.path.join(data_path, "model") + if not os.path.exists(model_path): + os.makedirs(model_path) + + hair_style_name = str(index) + "_hairstyle" + hair_style_path = os.path.join(images_path, hair_style_name) + if not os.path.exists(hair_style_path): + os.makedirs(hair_style_path) + print("hair_style_path:", hair_style_path) + + # res_list = random.sample(resolution_list, 5) + res_list = resolution_list + for res in res_list: + src_res = final_img.shape[0] + if src_res < res: + img2save = cv2.resize(final_img, (res, res), interpolation=cv2.INTER_LANCZOS4) + else: + img2save = cv2.resize(final_img, (res, res), interpolation=cv2.INTER_AREA) + cv2.imwrite(os.path.join(hair_style_path, str(uuid.uuid4()) + '.png'), img2save) + + if tag_prompt == "" and "first##" in img_path: + tag_prompt = caption_image(img2save) + print("tag_prompt:", tag_prompt) + + except Exception as e: + print(e) + ret_data['code'] = 300 + ret_data['data'] = [] + ret_data['msg'] = 'gen material failed' + return ret_data + + + # 3. 调用算法接口 + try: + hair_id = osp.basename(input_info[0][-1]) + task_id = hair_id + is_tj = "1" + + # 计算调用 webui 服务的地址 + # test + # # self.gpu_index = 1 + # webui_addr = train_services[self.gpu_index] + + try: + train_gpu_num = self.train_gpu_sq.get() + train_services_index = train_gpu_num % len(train_services) + train_gpu_id = list(train_services.keys())[train_services_index] + webui_addr = train_services[train_gpu_id] + + train_gpu_num += 1 + self.train_gpu_sq.put(train_gpu_num) + except Exception as e: + print(Exception) + train_gpu_id = next(iter(train_services)) + + # test, todo del + train_gpu_id = 0 + webui_addr = train_services[train_gpu_id] + + print(f"train_gpu_id: {train_gpu_id}, webui_addr: {webui_addr}") + + + + + # call_hair_train(task_id, hair_id, data_path, tag_prompt, is_tj, webui_addr=webui_addr, device_id=self.gpu_index) + call_hair_train(task_id, hair_id, data_path, tag_prompt, is_tj, webui_addr=webui_addr, device_id=train_gpu_id) + + except Exception as e: + print('call hair train failed, ConnectionError') + print(e) + ret_data['code'] = 300 + ret_data['data'] = [] + ret_data['msg'] = 'train failed' + return ret_data + + return ret_data + + + def get_body_info(self, img_res): + with torch.no_grad(): + img_w, img_h = 768, 1024 + ori_h, ori_w, _ = img_res.shape + print(f"ori_w, ori_h,: {ori_w},{ori_h}") + res = self.person_processor.forward(img_res) + if len(res['boxes']) == 0: + print('No person detected') + return [0,0,ori_w, ori_h] + filter_boxes = [] + + for ix, box_score in enumerate(res['scores']): + box_ = res['boxes'][ix] + box_w = box_[2] - box_[0] + box_h = box_[3] - box_[1] + max_box_len = max(box_w, box_h) + if box_score > 0.5 and max_box_len > 150: + filter_boxes.append([box_, box_w * box_h, box_h / img_res.shape[0]]) + + filter_boxes.sort(key=lambda x: x[1], reverse=True) + filter_boxes = list(filter(lambda x: x[2] > 0.2, filter_boxes)) + if len(filter_boxes) == 0: + return [0, 0, ori_w, ori_h] + filter_boxes = [item[0] for item in filter_boxes] + # return filter_boxes, res + # step2 person keypoints (version: hrnet 17pt) + keypoints = self.keypoints_processor.forward(img_res, filter_boxes) + + # step3 human keypoints: hands, face, body keypoints(25pt) + person_keypoints = keypoints[0] + human_box = filter_boxes[0].reshape((-1, 2)) + human_box[0][1] = max(0, human_box[0][1] - ori_h * 0.1) + # human_box[0][0] = max(0, human_box[0][0]-ori_w*0.1) + # human_box[1][0] = min(ori_w, human_box[1][0]+ori_w*0.1) + + knee_keypoints_left = person_keypoints[13] + knee_keypoints_right = person_keypoints[14] + if knee_keypoints_left[2] > 0.4 or knee_keypoints_right[2] > 0.4: + src_body_kpnts_select = person_keypoints[np.where(person_keypoints[:, 2] > 0.2)][:12][:, :2] + src_body_box = cv2.boundingRect(src_body_kpnts_select[np.newaxis, :, :]) + kpnts_bbox_tlx, kpnts_bbox_tly = src_body_box[:2] + kpnts_bbox_brx, kpnts_bbox_bry = kpnts_bbox_tlx + src_body_box[2] - 1, kpnts_bbox_tly + src_body_box[ + 3] - 1 + + kpnts_bbox_tlx, kpnts_bbox_tly = min(kpnts_bbox_tlx, human_box[0][0]), min(kpnts_bbox_tly, + human_box[0][1]) + kpnts_bbox_brx, kpnts_bbox_bry = min(kpnts_bbox_brx, human_box[1][0]), min(kpnts_bbox_bry, + human_box[1][1]) + box_center_x = int(kpnts_bbox_tlx / 2 + kpnts_bbox_brx / 2) + box_hight_now = kpnts_bbox_bry - kpnts_bbox_tly + kpnts_bbox_tlx = max(int(box_center_x - box_hight_now * 3 / 8), 0) + kpnts_bbox_brx = min(int(box_center_x + box_hight_now * 3 / 8), ori_w) + # cv2.rectangle(img_res, (kpnts_bbox_tlx, kpnts_bbox_tly), (kpnts_bbox_brx, kpnts_bbox_bry), (0, 255, 0), + # 2) + waist_keypoints_mean = (person_keypoints[11] + person_keypoints[12]) / 2 + + crop_img = img_res[kpnts_bbox_tly:kpnts_bbox_bry, kpnts_bbox_tlx:kpnts_bbox_brx] + print(f"a,b,c,d: {kpnts_bbox_tlx}, {kpnts_bbox_tly},{kpnts_bbox_brx}, {kpnts_bbox_bry}") + return [kpnts_bbox_tlx, kpnts_bbox_tly,kpnts_bbox_brx, kpnts_bbox_bry] + return [0, 0, ori_w, ori_h] + + + def haircolor_for_example_phone(self, image, color=[20, 20, 200]): + b, g, r = color # [10, 50, 250] # [10, 250, 10] + tar_color = np.zeros_like(image) + tar_color[:, :, 0] = b + tar_color[:, :, 1] = g + tar_color[:, :, 2] = r + image_hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) + tar_hsv = cv2.cvtColor(tar_color, cv2.COLOR_BGR2HSV) + + # hair_mask_8uc3 = np.ones_like(image) + # loc_index = hair_mask_8uc3[:, :, 0].nonzero() + # color_val = image_hsv[loc_index] + # mean_hsv = np.mean(color_val, axis=0) + # print("mean_hsv: ", mean_hsv) + # print("tar_hsv: ", tar_hsv[0, 0, :]) + + image_hsv[:, :, 0:1] = tar_hsv[:, :, 0:1] + image_hsv[:, :, 1:2] = tar_hsv[:, :, 1:2] + image_hsv[:, :, 2:3] = tar_hsv[:, :, 2:3] + + img_res = cv2.cvtColor(image_hsv.astype(np.uint8), cv2.COLOR_HSV2BGR) + origin_img_LAB = cv2.cvtColor(image, cv2.COLOR_BGR2LAB) + img_res_LAB = cv2.cvtColor(img_res, cv2.COLOR_BGR2LAB) + img_res_LAB[:, :, 0] = origin_img_LAB[:, :, 0] + img_res = cv2.cvtColor(img_res_LAB, cv2.COLOR_LAB2BGR) + + # mid_show = np.concatenate((image, img_res), axis=1) + # resize_ratio = 1280. / max(mid_show.shape[:2]) + # mid_show = cv2.resize(mid_show, (0, 0), fx=resize_ratio, fy=resize_ratio) + # cv2.imshow("mid_show", mid_show) + # cv2.waitKey() + return img_res + + def haircolor_for_example_phone_v2(self, image, color=[20, 20, 200]): + b, g, r = color # [10, 50, 250] # [10, 250, 10] + tar_color = np.zeros_like(image) + tar_color[:, :, 0] = b + tar_color[:, :, 1] = g + tar_color[:, :, 2] = r + image_hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) + tar_hsv = cv2.cvtColor(tar_color, cv2.COLOR_BGR2HSV) + + image_hsv[:, :, 0:1] = tar_hsv[:, :, 0:1] + image_hsv[:, :, 1:2] = tar_hsv[:, :, 1:2] + image_hsv[:, :, 2:3] = tar_hsv[:, :, 2:3] + + img_res = cv2.cvtColor(image_hsv.astype(np.uint8), cv2.COLOR_HSV2BGR) + origin_img_LAB = cv2.cvtColor(image, cv2.COLOR_BGR2LAB) + img_res_LAB = cv2.cvtColor(img_res, cv2.COLOR_BGR2LAB) + mix_ratio = 0.5 + img_res_LAB[:, :, 0] = origin_img_LAB[:, :, 0] * mix_ratio + img_res_LAB[:, :, 0] * (1 - mix_ratio) + img_res = cv2.cvtColor(img_res_LAB, cv2.COLOR_LAB2BGR) + + img_res_fc32 = img_res.astype(np.float32) / 255 + tar_color_fc32 = tar_color.astype(np.float32) / 255 + ref_rgb_avg_std = self.getavgstd(tar_color_fc32, np.ones_like(img_res_fc32)) + img_res_fc32, src_avg_std = self.reinhard_rgb(img_res_fc32, np.ones_like(img_res_fc32), ref_rgb_avg_std, ratio=1.0) + img_res = (img_res_fc32 * 255).astype(np.uint8) + # mid_show = np.concatenate((image, img_res), axis=1) + # resize_ratio = 1280. / max(mid_show.shape[:2]) + # mid_show = cv2.resize(mid_show, (0, 0), fx=resize_ratio, fy=resize_ratio) + # cv2.imshow("mid_show", mid_show) + # cv2.waitKey() + return img_res + def get_prepare_ref_haircolor_768_data(self, ref_rgb_8uc3_orisize): + + ref_landmark_1k2_f_orisize, boundboxs, euler_info = self.get_landmark.forward(ref_rgb_8uc3_orisize) + + ref_matte_fg_8uc3_orisize, ref_matte_pred_8uc1_orisize, ref_rgb_8uc3_resize= self.process_data.generator_matte.matte_inference( + ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize) + # cv2.imshow('ref_rgb_8uc3_orisize', ref_rgb_8uc3_orisize) + # cv2.imshow('ref_matte_pred_8uc1_orisize', ref_matte_pred_8uc1_orisize) + # cv2.imshow('ref_matte_fg_8uc3_orisize', ref_matte_fg_8uc3_orisize) + # cv2.waitKey() + + # 光头分割 + 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 + + def prepare_single_color_v2(self, source_image, dst_dir, rgb): + ret_data = {} + debug = False + ret_data['code'] = 200 + ret_data['data'] = "" + ret_data['msg'] = 'success' + try: + r,g,b = rgb + newColor = [b,g,r] + img_res = self.haircolor_for_example_phone_v2(source_image, color=newColor) + color_dir = os.path.join(dst_dir, "color_{}_{}_{}.png".format(r,g,b)) + cv2.imwrite(color_dir, img_res) + ret_url = self.oss2.upload_file(color_dir, "hair_mz/images/color/{}".format("colors_{}_{}_{}.png".format(r,g,b))) + ret_data['data'] = ret_url + except Exception as e: + print(e) + ret_data['code'] = 400 + return ret_data + return ret_data + + def check_hair_covering_eyes(self, hair_pt1k, hair_mask_img_path, new_user_img_path, user_pt1k, user_mask_img_path): + hair_mask_img = cv2.imread(hair_mask_img_path) + _, binary_hair_mask = cv2.threshold(hair_mask_img, 200, 255, cv2.THRESH_BINARY) + landmarks_origin_img_137 = landmark_processor.pts_1k_to_137(hair_pt1k) + # right_eye_coords = landmarks_origin_img_137[88:104].astype(int) + # left_eye_coords = landmarks_origin_img_137[105:121].astype(int) + eye_coords = landmarks_origin_img_137[87:121].astype(int) + points_in_white = [] + for coord in eye_coords: + x, y = coord + pixel_value = binary_hair_mask[y, x] + # print(f"眼坐标({x}, {y})处的颜色为：{pixel_value}") + if np.array_equal(pixel_value, np.array([255, 255, 255])): + points_in_white.append((x, y)) + break + # 在图像上标注眼睛关键点并打印坐标 + # colored_hair_mask_img = binary_hair_mask.copy() + # for coord in points_in_white: + # x, y = coord + # cv2.circle(colored_hair_mask_img, (x, y), 2, (0, 0, 255), -1) + # cv2.imwrite('a.png', colored_hair_mask_img) + + # user_gender = "girl" + # user_img_ori = cv2.imread(new_user_img_path) + # is_female = self.get_gender(user_img_ori) + # if not is_female: + # user_gender = "boy" + # user_mask_img = cv2.imread(user_mask_img_path) + # _, user_binary_hair_mask = cv2.threshold(user_mask_img, 200, 255, cv2.THRESH_BINARY) + user_landmarks_origin_img_137 = landmark_processor.pts_1k_to_137(user_pt1k) + # eyebrows_coords = user_landmarks_origin_img_137[121:137].astype(int) + # eyebrows_points_in_white = [] + # for eyebrows_coord in eyebrows_coords: + # eyebrows_x, eyebrows_y = eyebrows_coord + # eyebrows_pixel_value = user_binary_hair_mask[eyebrows_y, eyebrows_x] + # # print(f"眼坐标({x}, {y})处的颜色为：{pixel_value}") + # if np.array_equal(eyebrows_pixel_value, np.array([255, 255, 255])): + # eyebrows_points_in_white.append((eyebrows_x, eyebrows_y)) + # break + # # 在图像上标注眼睛关键点并打印坐标 + # # user_colored_hair_mask_img = user_binary_hair_mask.copy() + # # for coord in eyebrows_points_in_white: + # # x, y = coord + # # cv2.circle(user_colored_hair_mask_img, (x, y), 2, (0, 0, 255), -1) + # # cv2.imwrite('a.png', user_colored_hair_mask_img) + # + + if len(points_in_white) > 0: + result = 1 + # elif user_gender == 'boy' and len(eyebrows_points_in_white) > 0: + # result = 2 + else: + result = 0 + return result, user_landmarks_origin_img_137 diff --git a/hair_service_sd/core/matting/__init__.py b/hair_service_sd/core/matting/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hair_service_sd/core/matting/gca_matting_hair_single.py b/hair_service_sd/core/matting/gca_matting_hair_single.py new file mode 100644 index 0000000..c42d3f0 --- /dev/null +++ b/hair_service_sd/core/matting/gca_matting_hair_single.py @@ -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]) diff --git a/hair_service_sd/core/matting/gca_matting_hair_single_fg.py b/hair_service_sd/core/matting/gca_matting_hair_single_fg.py new file mode 100644 index 0000000..c7d7369 --- /dev/null +++ b/hair_service_sd/core/matting/gca_matting_hair_single_fg.py @@ -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) + diff --git a/hair_service_sd/core/matting/networks/__init__.py b/hair_service_sd/core/matting/networks/__init__.py new file mode 100644 index 0000000..aee2746 --- /dev/null +++ b/hair_service_sd/core/matting/networks/__init__.py @@ -0,0 +1 @@ +from .generators import * \ No newline at end of file diff --git a/hair_service_sd/core/matting/networks/decoders/__init__.py b/hair_service_sd/core/matting/networks/decoders/__init__.py new file mode 100644 index 0000000..a7eac0f --- /dev/null +++ b/hair_service_sd/core/matting/networks/decoders/__init__.py @@ -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) \ No newline at end of file diff --git a/hair_service_sd/core/matting/networks/decoders/res_gca_dec.py b/hair_service_sd/core/matting/networks/decoders/res_gca_dec.py new file mode 100644 index 0000000..b8b6027 --- /dev/null +++ b/hair_service_sd/core/matting/networks/decoders/res_gca_dec.py @@ -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} + diff --git a/hair_service_sd/core/matting/networks/decoders/res_shortcut_dec.py b/hair_service_sd/core/matting/networks/decoders/res_shortcut_dec.py new file mode 100644 index 0000000..fcb574a --- /dev/null +++ b/hair_service_sd/core/matting/networks/decoders/res_shortcut_dec.py @@ -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 + diff --git a/hair_service_sd/core/matting/networks/decoders/resnet_dec.py b/hair_service_sd/core/matting/networks/decoders/resnet_dec.py new file mode 100644 index 0000000..9f5b7b9 --- /dev/null +++ b/hair_service_sd/core/matting/networks/decoders/resnet_dec.py @@ -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 diff --git a/hair_service_sd/core/matting/networks/encoders/__init__.py b/hair_service_sd/core/matting/networks/encoders/__init__.py new file mode 100644 index 0000000..71693a6 --- /dev/null +++ b/hair_service_sd/core/matting/networks/encoders/__init__.py @@ -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) diff --git a/hair_service_sd/core/matting/networks/encoders/res_gca_enc.py b/hair_service_sd/core/matting/networks/encoders/res_gca_enc.py new file mode 100644 index 0000000..44025a1 --- /dev/null +++ b/hair_service_sd/core/matting/networks/encoders/res_gca_enc.py @@ -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) diff --git a/hair_service_sd/core/matting/networks/encoders/res_shortcut_enc.py b/hair_service_sd/core/matting/networks/encoders/res_shortcut_enc.py new file mode 100644 index 0000000..9fd66a6 --- /dev/null +++ b/hair_service_sd/core/matting/networks/encoders/res_shortcut_enc.py @@ -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,...]} \ No newline at end of file diff --git a/hair_service_sd/core/matting/networks/encoders/resnet_enc.py b/hair_service_sd/core/matting/networks/encoders/resnet_enc.py new file mode 100644 index 0000000..38ced43 --- /dev/null +++ b/hair_service_sd/core/matting/networks/encoders/resnet_enc.py @@ -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()) diff --git a/hair_service_sd/core/matting/networks/generators.py b/hair_service_sd/core/matting/networks/generators.py new file mode 100644 index 0000000..f64f259 --- /dev/null +++ b/hair_service_sd/core/matting/networks/generators.py @@ -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) \ No newline at end of file diff --git a/hair_service_sd/core/matting/networks/ops.py b/hair_service_sd/core/matting/networks/ops.py new file mode 100644 index 0000000..a2df034 --- /dev/null +++ b/hair_service_sd/core/matting/networks/ops.py @@ -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) \ No newline at end of file diff --git a/hair_service_sd/core/matting/setup.py b/hair_service_sd/core/matting/setup.py new file mode 100644 index 0000000..ca6bf5a --- /dev/null +++ b/hair_service_sd/core/matting/setup.py @@ -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 \ No newline at end of file diff --git a/hair_service_sd/core/model_3ddfa/model_3ddfa.py b/hair_service_sd/core/model_3ddfa/model_3ddfa.py new file mode 100644 index 0000000..f43aadf --- /dev/null +++ b/hair_service_sd/core/model_3ddfa/model_3ddfa.py @@ -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 \ No newline at end of file diff --git a/hair_service_sd/core/models/Generator_BaldSeg.py b/hair_service_sd/core/models/Generator_BaldSeg.py new file mode 100644 index 0000000..a9b61ce --- /dev/null +++ b/hair_service_sd/core/models/Generator_BaldSeg.py @@ -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 \ No newline at end of file diff --git a/hair_service_sd/core/models/MomocvFaceAlignment1K.py b/hair_service_sd/core/models/MomocvFaceAlignment1K.py new file mode 100644 index 0000000..cfc6c88 --- /dev/null +++ b/hair_service_sd/core/models/MomocvFaceAlignment1K.py @@ -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 \ No newline at end of file diff --git a/hair_service_sd/core/models/__init__.py b/hair_service_sd/core/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hair_service_sd/core/models/box_utils.py b/hair_service_sd/core/models/box_utils.py new file mode 100644 index 0000000..d7a076f --- /dev/null +++ b/hair_service_sd/core/models/box_utils.py @@ -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 diff --git a/hair_service_sd/core/models/config.py b/hair_service_sd/core/models/config.py new file mode 100644 index 0000000..591f349 --- /dev/null +++ b/hair_service_sd/core/models/config.py @@ -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 +} + diff --git a/hair_service_sd/core/models/detector.py b/hair_service_sd/core/models/detector.py new file mode 100644 index 0000000..1996cee --- /dev/null +++ b/hair_service_sd/core/models/detector.py @@ -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 diff --git a/hair_service_sd/core/models/layers/__init__.py b/hair_service_sd/core/models/layers/__init__.py new file mode 100644 index 0000000..53a3f4b --- /dev/null +++ b/hair_service_sd/core/models/layers/__init__.py @@ -0,0 +1,2 @@ +from .functions import * +from .modules import * diff --git a/hair_service_sd/core/models/layers/data/FDDB/img_list.txt b/hair_service_sd/core/models/layers/data/FDDB/img_list.txt new file mode 100644 index 0000000..5cf3d31 --- /dev/null +++ b/hair_service_sd/core/models/layers/data/FDDB/img_list.txt @@ -0,0 +1,2845 @@ +2002/08/11/big/img_591 +2002/08/26/big/img_265 +2002/07/19/big/img_423 +2002/08/24/big/img_490 +2002/08/31/big/img_17676 +2002/07/31/big/img_228 +2002/07/24/big/img_402 +2002/08/04/big/img_769 +2002/07/19/big/img_581 +2002/08/13/big/img_723 +2002/08/12/big/img_821 +2003/01/17/big/img_610 +2002/08/13/big/img_1116 +2002/08/28/big/img_19238 +2002/08/21/big/img_660 +2002/08/14/big/img_607 +2002/08/05/big/img_3708 +2002/08/19/big/img_511 +2002/08/07/big/img_1316 +2002/07/25/big/img_1047 +2002/07/23/big/img_474 +2002/07/27/big/img_970 +2002/09/02/big/img_15752 +2002/09/01/big/img_16378 +2002/09/01/big/img_16189 +2002/08/26/big/img_276 +2002/07/24/big/img_518 +2002/08/14/big/img_1027 +2002/08/24/big/img_733 +2002/08/15/big/img_249 +2003/01/15/big/img_1371 +2002/08/07/big/img_1348 +2003/01/01/big/img_331 +2002/08/23/big/img_536 +2002/07/30/big/img_224 +2002/08/10/big/img_763 +2002/08/21/big/img_293 +2002/08/15/big/img_1211 +2002/08/15/big/img_1194 +2003/01/15/big/img_390 +2002/08/06/big/img_2893 +2002/08/17/big/img_691 +2002/08/07/big/img_1695 +2002/08/16/big/img_829 +2002/07/25/big/img_201 +2002/08/23/big/img_36 +2003/01/15/big/img_763 +2003/01/15/big/img_637 +2002/08/22/big/img_592 +2002/07/25/big/img_817 +2003/01/15/big/img_1219 +2002/08/05/big/img_3508 +2002/08/15/big/img_1108 +2002/07/19/big/img_488 +2003/01/16/big/img_704 +2003/01/13/big/img_1087 +2002/08/10/big/img_670 +2002/07/24/big/img_104 +2002/08/27/big/img_19823 +2002/09/01/big/img_16229 +2003/01/13/big/img_846 +2002/08/04/big/img_412 +2002/07/22/big/img_554 +2002/08/12/big/img_331 +2002/08/02/big/img_533 +2002/08/12/big/img_259 +2002/08/18/big/img_328 +2003/01/14/big/img_630 +2002/08/05/big/img_3541 +2002/08/06/big/img_2390 +2002/08/20/big/img_150 +2002/08/02/big/img_1231 +2002/08/16/big/img_710 +2002/08/19/big/img_591 +2002/07/22/big/img_725 +2002/07/24/big/img_820 +2003/01/13/big/img_568 +2002/08/22/big/img_853 +2002/08/09/big/img_648 +2002/08/23/big/img_528 +2003/01/14/big/img_888 +2002/08/30/big/img_18201 +2002/08/13/big/img_965 +2003/01/14/big/img_660 +2002/07/19/big/img_517 +2003/01/14/big/img_406 +2002/08/30/big/img_18433 +2002/08/07/big/img_1630 +2002/08/06/big/img_2717 +2002/08/21/big/img_470 +2002/07/23/big/img_633 +2002/08/20/big/img_915 +2002/08/16/big/img_893 +2002/07/29/big/img_644 +2002/08/15/big/img_529 +2002/08/16/big/img_668 +2002/08/07/big/img_1871 +2002/07/25/big/img_192 +2002/07/31/big/img_961 +2002/08/19/big/img_738 +2002/07/31/big/img_382 +2002/08/19/big/img_298 +2003/01/17/big/img_608 +2002/08/21/big/img_514 +2002/07/23/big/img_183 +2003/01/17/big/img_536 +2002/07/24/big/img_478 +2002/08/06/big/img_2997 +2002/09/02/big/img_15380 +2002/08/07/big/img_1153 +2002/07/31/big/img_967 +2002/07/31/big/img_711 +2002/08/26/big/img_664 +2003/01/01/big/img_326 +2002/08/24/big/img_775 +2002/08/08/big/img_961 +2002/08/16/big/img_77 +2002/08/12/big/img_296 +2002/07/22/big/img_905 +2003/01/13/big/img_284 +2002/08/13/big/img_887 +2002/08/24/big/img_849 +2002/07/30/big/img_345 +2002/08/18/big/img_419 +2002/08/01/big/img_1347 +2002/08/05/big/img_3670 +2002/07/21/big/img_479 +2002/08/08/big/img_913 +2002/09/02/big/img_15828 +2002/08/30/big/img_18194 +2002/08/08/big/img_471 +2002/08/22/big/img_734 +2002/08/09/big/img_586 +2002/08/09/big/img_454 +2002/07/29/big/img_47 +2002/07/19/big/img_381 +2002/07/29/big/img_733 +2002/08/20/big/img_327 +2002/07/21/big/img_96 +2002/08/06/big/img_2680 +2002/07/25/big/img_919 +2002/07/21/big/img_158 +2002/07/22/big/img_801 +2002/07/22/big/img_567 +2002/07/24/big/img_804 +2002/07/24/big/img_690 +2003/01/15/big/img_576 +2002/08/14/big/img_335 +2003/01/13/big/img_390 +2002/08/11/big/img_258 +2002/07/23/big/img_917 +2002/08/15/big/img_525 +2003/01/15/big/img_505 +2002/07/30/big/img_886 +2003/01/16/big/img_640 +2003/01/14/big/img_642 +2003/01/17/big/img_844 +2002/08/04/big/img_571 +2002/08/29/big/img_18702 +2003/01/15/big/img_240 +2002/07/29/big/img_553 +2002/08/10/big/img_354 +2002/08/18/big/img_17 +2003/01/15/big/img_782 +2002/07/27/big/img_382 +2002/08/14/big/img_970 +2003/01/16/big/img_70 +2003/01/16/big/img_625 +2002/08/18/big/img_341 +2002/08/26/big/img_188 +2002/08/09/big/img_405 +2002/08/02/big/img_37 +2002/08/13/big/img_748 +2002/07/22/big/img_399 +2002/07/25/big/img_844 +2002/08/12/big/img_340 +2003/01/13/big/img_815 +2002/08/26/big/img_5 +2002/08/10/big/img_158 +2002/08/18/big/img_95 +2002/07/29/big/img_1297 +2003/01/13/big/img_508 +2002/09/01/big/img_16680 +2003/01/16/big/img_338 +2002/08/13/big/img_517 +2002/07/22/big/img_626 +2002/08/06/big/img_3024 +2002/07/26/big/img_499 +2003/01/13/big/img_387 +2002/08/31/big/img_18025 +2002/08/13/big/img_520 +2003/01/16/big/img_576 +2002/07/26/big/img_121 +2002/08/25/big/img_703 +2002/08/26/big/img_615 +2002/08/17/big/img_434 +2002/08/02/big/img_677 +2002/08/18/big/img_276 +2002/08/05/big/img_3672 +2002/07/26/big/img_700 +2002/07/31/big/img_277 +2003/01/14/big/img_220 +2002/08/23/big/img_232 +2002/08/31/big/img_17422 +2002/07/22/big/img_508 +2002/08/13/big/img_681 +2003/01/15/big/img_638 +2002/08/30/big/img_18408 +2003/01/14/big/img_533 +2003/01/17/big/img_12 +2002/08/28/big/img_19388 +2002/08/08/big/img_133 +2002/07/26/big/img_885 +2002/08/19/big/img_387 +2002/08/27/big/img_19976 +2002/08/26/big/img_118 +2002/08/28/big/img_19146 +2002/08/05/big/img_3259 +2002/08/15/big/img_536 +2002/07/22/big/img_279 +2002/07/22/big/img_9 +2002/08/13/big/img_301 +2002/08/15/big/img_974 +2002/08/06/big/img_2355 +2002/08/01/big/img_1526 +2002/08/03/big/img_417 +2002/08/04/big/img_407 +2002/08/15/big/img_1029 +2002/07/29/big/img_700 +2002/08/01/big/img_1463 +2002/08/31/big/img_17365 +2002/07/28/big/img_223 +2002/07/19/big/img_827 +2002/07/27/big/img_531 +2002/07/19/big/img_845 +2002/08/20/big/img_382 +2002/07/31/big/img_268 +2002/08/27/big/img_19705 +2002/08/02/big/img_830 +2002/08/23/big/img_250 +2002/07/20/big/img_777 +2002/08/21/big/img_879 +2002/08/26/big/img_20146 +2002/08/23/big/img_789 +2002/08/06/big/img_2683 +2002/08/25/big/img_576 +2002/08/09/big/img_498 +2002/08/08/big/img_384 +2002/08/26/big/img_592 +2002/07/29/big/img_1470 +2002/08/21/big/img_452 +2002/08/30/big/img_18395 +2002/08/15/big/img_215 +2002/07/21/big/img_643 +2002/07/22/big/img_209 +2003/01/17/big/img_346 +2002/08/25/big/img_658 +2002/08/21/big/img_221 +2002/08/14/big/img_60 +2003/01/17/big/img_885 +2003/01/16/big/img_482 +2002/08/19/big/img_593 +2002/08/08/big/img_233 +2002/07/30/big/img_458 +2002/07/23/big/img_384 +2003/01/15/big/img_670 +2003/01/15/big/img_267 +2002/08/26/big/img_540 +2002/07/29/big/img_552 +2002/07/30/big/img_997 +2003/01/17/big/img_377 +2002/08/21/big/img_265 +2002/08/09/big/img_561 +2002/07/31/big/img_945 +2002/09/02/big/img_15252 +2002/08/11/big/img_276 +2002/07/22/big/img_491 +2002/07/26/big/img_517 +2002/08/14/big/img_726 +2002/08/08/big/img_46 +2002/08/28/big/img_19458 +2002/08/06/big/img_2935 +2002/07/29/big/img_1392 +2002/08/13/big/img_776 +2002/08/24/big/img_616 +2002/08/14/big/img_1065 +2002/07/29/big/img_889 +2002/08/18/big/img_188 +2002/08/07/big/img_1453 +2002/08/02/big/img_760 +2002/07/28/big/img_416 +2002/08/07/big/img_1393 +2002/08/26/big/img_292 +2002/08/26/big/img_301 +2003/01/13/big/img_195 +2002/07/26/big/img_532 +2002/08/20/big/img_550 +2002/08/05/big/img_3658 +2002/08/26/big/img_738 +2002/09/02/big/img_15750 +2003/01/17/big/img_451 +2002/07/23/big/img_339 +2002/08/16/big/img_637 +2002/08/14/big/img_748 +2002/08/06/big/img_2739 +2002/07/25/big/img_482 +2002/08/19/big/img_191 +2002/08/26/big/img_537 +2003/01/15/big/img_716 +2003/01/15/big/img_767 +2002/08/02/big/img_452 +2002/08/08/big/img_1011 +2002/08/10/big/img_144 +2003/01/14/big/img_122 +2002/07/24/big/img_586 +2002/07/24/big/img_762 +2002/08/20/big/img_369 +2002/07/30/big/img_146 +2002/08/23/big/img_396 +2003/01/15/big/img_200 +2002/08/15/big/img_1183 +2003/01/14/big/img_698 +2002/08/09/big/img_792 +2002/08/06/big/img_2347 +2002/07/31/big/img_911 +2002/08/26/big/img_722 +2002/08/23/big/img_621 +2002/08/05/big/img_3790 +2003/01/13/big/img_633 +2002/08/09/big/img_224 +2002/07/24/big/img_454 +2002/07/21/big/img_202 +2002/08/02/big/img_630 +2002/08/30/big/img_18315 +2002/07/19/big/img_491 +2002/09/01/big/img_16456 +2002/08/09/big/img_242 +2002/07/25/big/img_595 +2002/07/22/big/img_522 +2002/08/01/big/img_1593 +2002/07/29/big/img_336 +2002/08/15/big/img_448 +2002/08/28/big/img_19281 +2002/07/29/big/img_342 +2002/08/12/big/img_78 +2003/01/14/big/img_525 +2002/07/28/big/img_147 +2002/08/11/big/img_353 +2002/08/22/big/img_513 +2002/08/04/big/img_721 +2002/08/17/big/img_247 +2003/01/14/big/img_891 +2002/08/20/big/img_853 +2002/07/19/big/img_414 +2002/08/01/big/img_1530 +2003/01/14/big/img_924 +2002/08/22/big/img_468 +2002/08/18/big/img_354 +2002/08/30/big/img_18193 +2002/08/23/big/img_492 +2002/08/15/big/img_871 +2002/08/12/big/img_494 +2002/08/06/big/img_2470 +2002/07/23/big/img_923 +2002/08/26/big/img_155 +2002/08/08/big/img_669 +2002/07/23/big/img_404 +2002/08/28/big/img_19421 +2002/08/29/big/img_18993 +2002/08/25/big/img_416 +2003/01/17/big/img_434 +2002/07/29/big/img_1370 +2002/07/28/big/img_483 +2002/08/11/big/img_50 +2002/08/10/big/img_404 +2002/09/02/big/img_15057 +2003/01/14/big/img_911 +2002/09/01/big/img_16697 +2003/01/16/big/img_665 +2002/09/01/big/img_16708 +2002/08/22/big/img_612 +2002/08/28/big/img_19471 +2002/08/02/big/img_198 +2003/01/16/big/img_527 +2002/08/22/big/img_209 +2002/08/30/big/img_18205 +2003/01/14/big/img_114 +2003/01/14/big/img_1028 +2003/01/16/big/img_894 +2003/01/14/big/img_837 +2002/07/30/big/img_9 +2002/08/06/big/img_2821 +2002/08/04/big/img_85 +2003/01/13/big/img_884 +2002/07/22/big/img_570 +2002/08/07/big/img_1773 +2002/07/26/big/img_208 +2003/01/17/big/img_946 +2002/07/19/big/img_930 +2003/01/01/big/img_698 +2003/01/17/big/img_612 +2002/07/19/big/img_372 +2002/07/30/big/img_721 +2003/01/14/big/img_649 +2002/08/19/big/img_4 +2002/07/25/big/img_1024 +2003/01/15/big/img_601 +2002/08/30/big/img_18470 +2002/07/22/big/img_29 +2002/08/07/big/img_1686 +2002/07/20/big/img_294 +2002/08/14/big/img_800 +2002/08/19/big/img_353 +2002/08/19/big/img_350 +2002/08/05/big/img_3392 +2002/08/09/big/img_622 +2003/01/15/big/img_236 +2002/08/11/big/img_643 +2002/08/05/big/img_3458 +2002/08/12/big/img_413 +2002/08/22/big/img_415 +2002/08/13/big/img_635 +2002/08/07/big/img_1198 +2002/08/04/big/img_873 +2002/08/12/big/img_407 +2003/01/15/big/img_346 +2002/08/02/big/img_275 +2002/08/17/big/img_997 +2002/08/21/big/img_958 +2002/08/20/big/img_579 +2002/07/29/big/img_142 +2003/01/14/big/img_1115 +2002/08/16/big/img_365 +2002/07/29/big/img_1414 +2002/08/17/big/img_489 +2002/08/13/big/img_1010 +2002/07/31/big/img_276 +2002/07/25/big/img_1000 +2002/08/23/big/img_524 +2002/08/28/big/img_19147 +2003/01/13/big/img_433 +2002/08/20/big/img_205 +2003/01/01/big/img_458 +2002/07/29/big/img_1449 +2003/01/16/big/img_696 +2002/08/28/big/img_19296 +2002/08/29/big/img_18688 +2002/08/21/big/img_767 +2002/08/20/big/img_532 +2002/08/26/big/img_187 +2002/07/26/big/img_183 +2002/07/27/big/img_890 +2003/01/13/big/img_576 +2002/07/30/big/img_15 +2002/07/31/big/img_889 +2002/08/31/big/img_17759 +2003/01/14/big/img_1114 +2002/07/19/big/img_445 +2002/08/03/big/img_593 +2002/07/24/big/img_750 +2002/07/30/big/img_133 +2002/08/25/big/img_671 +2002/07/20/big/img_351 +2002/08/31/big/img_17276 +2002/08/05/big/img_3231 +2002/09/02/big/img_15882 +2002/08/14/big/img_115 +2002/08/02/big/img_1148 +2002/07/25/big/img_936 +2002/07/31/big/img_639 +2002/08/04/big/img_427 +2002/08/22/big/img_843 +2003/01/17/big/img_17 +2003/01/13/big/img_690 +2002/08/13/big/img_472 +2002/08/09/big/img_425 +2002/08/05/big/img_3450 +2003/01/17/big/img_439 +2002/08/13/big/img_539 +2002/07/28/big/img_35 +2002/08/16/big/img_241 +2002/08/06/big/img_2898 +2003/01/16/big/img_429 +2002/08/05/big/img_3817 +2002/08/27/big/img_19919 +2002/07/19/big/img_422 +2002/08/15/big/img_560 +2002/07/23/big/img_750 +2002/07/30/big/img_353 +2002/08/05/big/img_43 +2002/08/23/big/img_305 +2002/08/01/big/img_2137 +2002/08/30/big/img_18097 +2002/08/01/big/img_1389 +2002/08/02/big/img_308 +2003/01/14/big/img_652 +2002/08/01/big/img_1798 +2003/01/14/big/img_732 +2003/01/16/big/img_294 +2002/08/26/big/img_213 +2002/07/24/big/img_842 +2003/01/13/big/img_630 +2003/01/13/big/img_634 +2002/08/06/big/img_2285 +2002/08/01/big/img_2162 +2002/08/30/big/img_18134 +2002/08/02/big/img_1045 +2002/08/01/big/img_2143 +2002/07/25/big/img_135 +2002/07/20/big/img_645 +2002/08/05/big/img_3666 +2002/08/14/big/img_523 +2002/08/04/big/img_425 +2003/01/14/big/img_137 +2003/01/01/big/img_176 +2002/08/15/big/img_505 +2002/08/24/big/img_386 +2002/08/05/big/img_3187 +2002/08/15/big/img_419 +2003/01/13/big/img_520 +2002/08/04/big/img_444 +2002/08/26/big/img_483 +2002/08/05/big/img_3449 +2002/08/30/big/img_18409 +2002/08/28/big/img_19455 +2002/08/27/big/img_20090 +2002/07/23/big/img_625 +2002/08/24/big/img_205 +2002/08/08/big/img_938 +2003/01/13/big/img_527 +2002/08/07/big/img_1712 +2002/07/24/big/img_801 +2002/08/09/big/img_579 +2003/01/14/big/img_41 +2003/01/15/big/img_1130 +2002/07/21/big/img_672 +2002/08/07/big/img_1590 +2003/01/01/big/img_532 +2002/08/02/big/img_529 +2002/08/05/big/img_3591 +2002/08/23/big/img_5 +2003/01/14/big/img_882 +2002/08/28/big/img_19234 +2002/07/24/big/img_398 +2003/01/14/big/img_592 +2002/08/22/big/img_548 +2002/08/12/big/img_761 +2003/01/16/big/img_497 +2002/08/18/big/img_133 +2002/08/08/big/img_874 +2002/07/19/big/img_247 +2002/08/15/big/img_170 +2002/08/27/big/img_19679 +2002/08/20/big/img_246 +2002/08/24/big/img_358 +2002/07/29/big/img_599 +2002/08/01/big/img_1555 +2002/07/30/big/img_491 +2002/07/30/big/img_371 +2003/01/16/big/img_682 +2002/07/25/big/img_619 +2003/01/15/big/img_587 +2002/08/02/big/img_1212 +2002/08/01/big/img_2152 +2002/07/25/big/img_668 +2003/01/16/big/img_574 +2002/08/28/big/img_19464 +2002/08/11/big/img_536 +2002/07/24/big/img_201 +2002/08/05/big/img_3488 +2002/07/25/big/img_887 +2002/07/22/big/img_789 +2002/07/30/big/img_432 +2002/08/16/big/img_166 +2002/09/01/big/img_16333 +2002/07/26/big/img_1010 +2002/07/21/big/img_793 +2002/07/22/big/img_720 +2002/07/31/big/img_337 +2002/07/27/big/img_185 +2002/08/23/big/img_440 +2002/07/31/big/img_801 +2002/07/25/big/img_478 +2003/01/14/big/img_171 +2002/08/07/big/img_1054 +2002/09/02/big/img_15659 +2002/07/29/big/img_1348 +2002/08/09/big/img_337 +2002/08/26/big/img_684 +2002/07/31/big/img_537 +2002/08/15/big/img_808 +2003/01/13/big/img_740 +2002/08/07/big/img_1667 +2002/08/03/big/img_404 +2002/08/06/big/img_2520 +2002/07/19/big/img_230 +2002/07/19/big/img_356 +2003/01/16/big/img_627 +2002/08/04/big/img_474 +2002/07/29/big/img_833 +2002/07/25/big/img_176 +2002/08/01/big/img_1684 +2002/08/21/big/img_643 +2002/08/27/big/img_19673 +2002/08/02/big/img_838 +2002/08/06/big/img_2378 +2003/01/15/big/img_48 +2002/07/30/big/img_470 +2002/08/15/big/img_963 +2002/08/24/big/img_444 +2002/08/16/big/img_662 +2002/08/15/big/img_1209 +2002/07/24/big/img_25 +2002/08/06/big/img_2740 +2002/07/29/big/img_996 +2002/08/31/big/img_18074 +2002/08/04/big/img_343 +2003/01/17/big/img_509 +2003/01/13/big/img_726 +2002/08/07/big/img_1466 +2002/07/26/big/img_307 +2002/08/10/big/img_598 +2002/08/13/big/img_890 +2002/08/14/big/img_997 +2002/07/19/big/img_392 +2002/08/02/big/img_475 +2002/08/29/big/img_19038 +2002/07/29/big/img_538 +2002/07/29/big/img_502 +2002/08/02/big/img_364 +2002/08/31/big/img_17353 +2002/08/08/big/img_539 +2002/08/01/big/img_1449 +2002/07/22/big/img_363 +2002/08/02/big/img_90 +2002/09/01/big/img_16867 +2002/08/05/big/img_3371 +2002/07/30/big/img_342 +2002/08/07/big/img_1363 +2002/08/22/big/img_790 +2003/01/15/big/img_404 +2002/08/05/big/img_3447 +2002/09/01/big/img_16167 +2003/01/13/big/img_840 +2002/08/22/big/img_1001 +2002/08/09/big/img_431 +2002/07/27/big/img_618 +2002/07/31/big/img_741 +2002/07/30/big/img_964 +2002/07/25/big/img_86 +2002/07/29/big/img_275 +2002/08/21/big/img_921 +2002/07/26/big/img_892 +2002/08/21/big/img_663 +2003/01/13/big/img_567 +2003/01/14/big/img_719 +2002/07/28/big/img_251 +2003/01/15/big/img_1123 +2002/07/29/big/img_260 +2002/08/24/big/img_337 +2002/08/01/big/img_1914 +2002/08/13/big/img_373 +2003/01/15/big/img_589 +2002/08/13/big/img_906 +2002/07/26/big/img_270 +2002/08/26/big/img_313 +2002/08/25/big/img_694 +2003/01/01/big/img_327 +2002/07/23/big/img_261 +2002/08/26/big/img_642 +2002/07/29/big/img_918 +2002/07/23/big/img_455 +2002/07/24/big/img_612 +2002/07/23/big/img_534 +2002/07/19/big/img_534 +2002/07/19/big/img_726 +2002/08/01/big/img_2146 +2002/08/02/big/img_543 +2003/01/16/big/img_777 +2002/07/30/big/img_484 +2002/08/13/big/img_1161 +2002/07/21/big/img_390 +2002/08/06/big/img_2288 +2002/08/21/big/img_677 +2002/08/13/big/img_747 +2002/08/15/big/img_1248 +2002/07/31/big/img_416 +2002/09/02/big/img_15259 +2002/08/16/big/img_781 +2002/08/24/big/img_754 +2002/07/24/big/img_803 +2002/08/20/big/img_609 +2002/08/28/big/img_19571 +2002/09/01/big/img_16140 +2002/08/26/big/img_769 +2002/07/20/big/img_588 +2002/08/02/big/img_898 +2002/07/21/big/img_466 +2002/08/14/big/img_1046 +2002/07/25/big/img_212 +2002/08/26/big/img_353 +2002/08/19/big/img_810 +2002/08/31/big/img_17824 +2002/08/12/big/img_631 +2002/07/19/big/img_828 +2002/07/24/big/img_130 +2002/08/25/big/img_580 +2002/07/31/big/img_699 +2002/07/23/big/img_808 +2002/07/31/big/img_377 +2003/01/16/big/img_570 +2002/09/01/big/img_16254 +2002/07/21/big/img_471 +2002/08/01/big/img_1548 +2002/08/18/big/img_252 +2002/08/19/big/img_576 +2002/08/20/big/img_464 +2002/07/27/big/img_735 +2002/08/21/big/img_589 +2003/01/15/big/img_1192 +2002/08/09/big/img_302 +2002/07/31/big/img_594 +2002/08/23/big/img_19 +2002/08/29/big/img_18819 +2002/08/19/big/img_293 +2002/07/30/big/img_331 +2002/08/23/big/img_607 +2002/07/30/big/img_363 +2002/08/16/big/img_766 +2003/01/13/big/img_481 +2002/08/06/big/img_2515 +2002/09/02/big/img_15913 +2002/09/02/big/img_15827 +2002/09/02/big/img_15053 +2002/08/07/big/img_1576 +2002/07/23/big/img_268 +2002/08/21/big/img_152 +2003/01/15/big/img_578 +2002/07/21/big/img_589 +2002/07/20/big/img_548 +2002/08/27/big/img_19693 +2002/08/31/big/img_17252 +2002/07/31/big/img_138 +2002/07/23/big/img_372 +2002/08/16/big/img_695 +2002/07/27/big/img_287 +2002/08/15/big/img_315 +2002/08/10/big/img_361 +2002/07/29/big/img_899 +2002/08/13/big/img_771 +2002/08/21/big/img_92 +2003/01/15/big/img_425 +2003/01/16/big/img_450 +2002/09/01/big/img_16942 +2002/08/02/big/img_51 +2002/09/02/big/img_15379 +2002/08/24/big/img_147 +2002/08/30/big/img_18122 +2002/07/26/big/img_950 +2002/08/07/big/img_1400 +2002/08/17/big/img_468 +2002/08/15/big/img_470 +2002/07/30/big/img_318 +2002/07/22/big/img_644 +2002/08/27/big/img_19732 +2002/07/23/big/img_601 +2002/08/26/big/img_398 +2002/08/21/big/img_428 +2002/08/06/big/img_2119 +2002/08/29/big/img_19103 +2003/01/14/big/img_933 +2002/08/11/big/img_674 +2002/08/28/big/img_19420 +2002/08/03/big/img_418 +2002/08/17/big/img_312 +2002/07/25/big/img_1044 +2003/01/17/big/img_671 +2002/08/30/big/img_18297 +2002/07/25/big/img_755 +2002/07/23/big/img_471 +2002/08/21/big/img_39 +2002/07/26/big/img_699 +2003/01/14/big/img_33 +2002/07/31/big/img_411 +2002/08/16/big/img_645 +2003/01/17/big/img_116 +2002/09/02/big/img_15903 +2002/08/20/big/img_120 +2002/08/22/big/img_176 +2002/07/29/big/img_1316 +2002/08/27/big/img_19914 +2002/07/22/big/img_719 +2002/08/28/big/img_19239 +2003/01/13/big/img_385 +2002/08/08/big/img_525 +2002/07/19/big/img_782 +2002/08/13/big/img_843 +2002/07/30/big/img_107 +2002/08/11/big/img_752 +2002/07/29/big/img_383 +2002/08/26/big/img_249 +2002/08/29/big/img_18860 +2002/07/30/big/img_70 +2002/07/26/big/img_194 +2002/08/15/big/img_530 +2002/08/08/big/img_816 +2002/07/31/big/img_286 +2003/01/13/big/img_294 +2002/07/31/big/img_251 +2002/07/24/big/img_13 +2002/08/31/big/img_17938 +2002/07/22/big/img_642 +2003/01/14/big/img_728 +2002/08/18/big/img_47 +2002/08/22/big/img_306 +2002/08/20/big/img_348 +2002/08/15/big/img_764 +2002/08/08/big/img_163 +2002/07/23/big/img_531 +2002/07/23/big/img_467 +2003/01/16/big/img_743 +2003/01/13/big/img_535 +2002/08/02/big/img_523 +2002/08/22/big/img_120 +2002/08/11/big/img_496 +2002/08/29/big/img_19075 +2002/08/08/big/img_465 +2002/08/09/big/img_790 +2002/08/19/big/img_588 +2002/08/23/big/img_407 +2003/01/17/big/img_435 +2002/08/24/big/img_398 +2002/08/27/big/img_19899 +2003/01/15/big/img_335 +2002/08/13/big/img_493 +2002/09/02/big/img_15460 +2002/07/31/big/img_470 +2002/08/05/big/img_3550 +2002/07/28/big/img_123 +2002/08/01/big/img_1498 +2002/08/04/big/img_504 +2003/01/17/big/img_427 +2002/08/27/big/img_19708 +2002/07/27/big/img_861 +2002/07/25/big/img_685 +2002/07/31/big/img_207 +2003/01/14/big/img_745 +2002/08/31/big/img_17756 +2002/08/24/big/img_288 +2002/08/18/big/img_181 +2002/08/10/big/img_520 +2002/08/25/big/img_705 +2002/08/23/big/img_226 +2002/08/04/big/img_727 +2002/07/24/big/img_625 +2002/08/28/big/img_19157 +2002/08/23/big/img_586 +2002/07/31/big/img_232 +2003/01/13/big/img_240 +2003/01/14/big/img_321 +2003/01/15/big/img_533 +2002/07/23/big/img_480 +2002/07/24/big/img_371 +2002/08/21/big/img_702 +2002/08/31/big/img_17075 +2002/09/02/big/img_15278 +2002/07/29/big/img_246 +2003/01/15/big/img_829 +2003/01/15/big/img_1213 +2003/01/16/big/img_441 +2002/08/14/big/img_921 +2002/07/23/big/img_425 +2002/08/15/big/img_296 +2002/07/19/big/img_135 +2002/07/26/big/img_402 +2003/01/17/big/img_88 +2002/08/20/big/img_872 +2002/08/13/big/img_1110 +2003/01/16/big/img_1040 +2002/07/23/big/img_9 +2002/08/13/big/img_700 +2002/08/16/big/img_371 +2002/08/27/big/img_19966 +2003/01/17/big/img_391 +2002/08/18/big/img_426 +2002/08/01/big/img_1618 +2002/07/21/big/img_754 +2003/01/14/big/img_1101 +2003/01/16/big/img_1022 +2002/07/22/big/img_275 +2002/08/24/big/img_86 +2002/08/17/big/img_582 +2003/01/15/big/img_765 +2003/01/17/big/img_449 +2002/07/28/big/img_265 +2003/01/13/big/img_552 +2002/07/28/big/img_115 +2003/01/16/big/img_56 +2002/08/02/big/img_1232 +2003/01/17/big/img_925 +2002/07/22/big/img_445 +2002/07/25/big/img_957 +2002/07/20/big/img_589 +2002/08/31/big/img_17107 +2002/07/29/big/img_483 +2002/08/14/big/img_1063 +2002/08/07/big/img_1545 +2002/08/14/big/img_680 +2002/09/01/big/img_16694 +2002/08/14/big/img_257 +2002/08/11/big/img_726 +2002/07/26/big/img_681 +2002/07/25/big/img_481 +2003/01/14/big/img_737 +2002/08/28/big/img_19480 +2003/01/16/big/img_362 +2002/08/27/big/img_19865 +2003/01/01/big/img_547 +2002/09/02/big/img_15074 +2002/08/01/big/img_1453 +2002/08/22/big/img_594 +2002/08/28/big/img_19263 +2002/08/13/big/img_478 +2002/07/29/big/img_1358 +2003/01/14/big/img_1022 +2002/08/16/big/img_450 +2002/08/02/big/img_159 +2002/07/26/big/img_781 +2003/01/13/big/img_601 +2002/08/20/big/img_407 +2002/08/15/big/img_468 +2002/08/31/big/img_17902 +2002/08/16/big/img_81 +2002/07/25/big/img_987 +2002/07/25/big/img_500 +2002/08/02/big/img_31 +2002/08/18/big/img_538 +2002/08/08/big/img_54 +2002/07/23/big/img_686 +2002/07/24/big/img_836 +2003/01/17/big/img_734 +2002/08/16/big/img_1055 +2003/01/16/big/img_521 +2002/07/25/big/img_612 +2002/08/22/big/img_778 +2002/08/03/big/img_251 +2002/08/12/big/img_436 +2002/08/23/big/img_705 +2002/07/28/big/img_243 +2002/07/25/big/img_1029 +2002/08/20/big/img_287 +2002/08/29/big/img_18739 +2002/08/05/big/img_3272 +2002/07/27/big/img_214 +2003/01/14/big/img_5 +2002/08/01/big/img_1380 +2002/08/29/big/img_19097 +2002/07/30/big/img_486 +2002/08/29/big/img_18707 +2002/08/10/big/img_559 +2002/08/15/big/img_365 +2002/08/09/big/img_525 +2002/08/10/big/img_689 +2002/07/25/big/img_502 +2002/08/03/big/img_667 +2002/08/10/big/img_855 +2002/08/10/big/img_706 +2002/08/18/big/img_603 +2003/01/16/big/img_1055 +2002/08/31/big/img_17890 +2002/08/15/big/img_761 +2003/01/15/big/img_489 +2002/08/26/big/img_351 +2002/08/01/big/img_1772 +2002/08/31/big/img_17729 +2002/07/25/big/img_609 +2003/01/13/big/img_539 +2002/07/27/big/img_686 +2002/07/31/big/img_311 +2002/08/22/big/img_799 +2003/01/16/big/img_936 +2002/08/31/big/img_17813 +2002/08/04/big/img_862 +2002/08/09/big/img_332 +2002/07/20/big/img_148 +2002/08/12/big/img_426 +2002/07/24/big/img_69 +2002/07/27/big/img_685 +2002/08/02/big/img_480 +2002/08/26/big/img_154 +2002/07/24/big/img_598 +2002/08/01/big/img_1881 +2002/08/20/big/img_667 +2003/01/14/big/img_495 +2002/07/21/big/img_744 +2002/07/30/big/img_150 +2002/07/23/big/img_924 +2002/08/08/big/img_272 +2002/07/23/big/img_310 +2002/07/25/big/img_1011 +2002/09/02/big/img_15725 +2002/07/19/big/img_814 +2002/08/20/big/img_936 +2002/07/25/big/img_85 +2002/08/24/big/img_662 +2002/08/09/big/img_495 +2003/01/15/big/img_196 +2002/08/16/big/img_707 +2002/08/28/big/img_19370 +2002/08/06/big/img_2366 +2002/08/06/big/img_3012 +2002/08/01/big/img_1452 +2002/07/31/big/img_742 +2002/07/27/big/img_914 +2003/01/13/big/img_290 +2002/07/31/big/img_288 +2002/08/02/big/img_171 +2002/08/22/big/img_191 +2002/07/27/big/img_1066 +2002/08/12/big/img_383 +2003/01/17/big/img_1018 +2002/08/01/big/img_1785 +2002/08/11/big/img_390 +2002/08/27/big/img_20037 +2002/08/12/big/img_38 +2003/01/15/big/img_103 +2002/08/26/big/img_31 +2002/08/18/big/img_660 +2002/07/22/big/img_694 +2002/08/15/big/img_24 +2002/07/27/big/img_1077 +2002/08/01/big/img_1943 +2002/07/22/big/img_292 +2002/09/01/big/img_16857 +2002/07/22/big/img_892 +2003/01/14/big/img_46 +2002/08/09/big/img_469 +2002/08/09/big/img_414 +2003/01/16/big/img_40 +2002/08/28/big/img_19231 +2002/07/27/big/img_978 +2002/07/23/big/img_475 +2002/07/25/big/img_92 +2002/08/09/big/img_799 +2002/07/25/big/img_491 +2002/08/03/big/img_654 +2003/01/15/big/img_687 +2002/08/11/big/img_478 +2002/08/07/big/img_1664 +2002/08/20/big/img_362 +2002/08/01/big/img_1298 +2003/01/13/big/img_500 +2002/08/06/big/img_2896 +2002/08/30/big/img_18529 +2002/08/16/big/img_1020 +2002/07/29/big/img_892 +2002/08/29/big/img_18726 +2002/07/21/big/img_453 +2002/08/17/big/img_437 +2002/07/19/big/img_665 +2002/07/22/big/img_440 +2002/07/19/big/img_582 +2002/07/21/big/img_233 +2003/01/01/big/img_82 +2002/07/25/big/img_341 +2002/07/29/big/img_864 +2002/08/02/big/img_276 +2002/08/29/big/img_18654 +2002/07/27/big/img_1024 +2002/08/19/big/img_373 +2003/01/15/big/img_241 +2002/07/25/big/img_84 +2002/08/13/big/img_834 +2002/08/10/big/img_511 +2002/08/01/big/img_1627 +2002/08/08/big/img_607 +2002/08/06/big/img_2083 +2002/08/01/big/img_1486 +2002/08/08/big/img_700 +2002/08/01/big/img_1954 +2002/08/21/big/img_54 +2002/07/30/big/img_847 +2002/08/28/big/img_19169 +2002/07/21/big/img_549 +2002/08/03/big/img_693 +2002/07/31/big/img_1002 +2003/01/14/big/img_1035 +2003/01/16/big/img_622 +2002/07/30/big/img_1201 +2002/08/10/big/img_444 +2002/07/31/big/img_374 +2002/08/21/big/img_301 +2002/08/13/big/img_1095 +2003/01/13/big/img_288 +2002/07/25/big/img_232 +2003/01/13/big/img_967 +2002/08/26/big/img_360 +2002/08/05/big/img_67 +2002/08/29/big/img_18969 +2002/07/28/big/img_16 +2002/08/16/big/img_515 +2002/07/20/big/img_708 +2002/08/18/big/img_178 +2003/01/15/big/img_509 +2002/07/25/big/img_430 +2002/08/21/big/img_738 +2002/08/16/big/img_886 +2002/09/02/big/img_15605 +2002/09/01/big/img_16242 +2002/08/24/big/img_711 +2002/07/25/big/img_90 +2002/08/09/big/img_491 +2002/07/30/big/img_534 +2003/01/13/big/img_474 +2002/08/25/big/img_510 +2002/08/15/big/img_555 +2002/08/02/big/img_775 +2002/07/23/big/img_975 +2002/08/19/big/img_229 +2003/01/17/big/img_860 +2003/01/02/big/img_10 +2002/07/23/big/img_542 +2002/08/06/big/img_2535 +2002/07/22/big/img_37 +2002/08/06/big/img_2342 +2002/08/25/big/img_515 +2002/08/25/big/img_336 +2002/08/18/big/img_837 +2002/08/21/big/img_616 +2003/01/17/big/img_24 +2002/07/26/big/img_936 +2002/08/14/big/img_896 +2002/07/29/big/img_465 +2002/07/31/big/img_543 +2002/08/01/big/img_1411 +2002/08/02/big/img_423 +2002/08/21/big/img_44 +2002/07/31/big/img_11 +2003/01/15/big/img_628 +2003/01/15/big/img_605 +2002/07/30/big/img_571 +2002/07/23/big/img_428 +2002/08/15/big/img_942 +2002/07/26/big/img_531 +2003/01/16/big/img_59 +2002/08/02/big/img_410 +2002/07/31/big/img_230 +2002/08/19/big/img_806 +2003/01/14/big/img_462 +2002/08/16/big/img_370 +2002/08/13/big/img_380 +2002/08/16/big/img_932 +2002/07/19/big/img_393 +2002/08/20/big/img_764 +2002/08/15/big/img_616 +2002/07/26/big/img_267 +2002/07/27/big/img_1069 +2002/08/14/big/img_1041 +2003/01/13/big/img_594 +2002/09/01/big/img_16845 +2002/08/09/big/img_229 +2003/01/16/big/img_639 +2002/08/19/big/img_398 +2002/08/18/big/img_978 +2002/08/24/big/img_296 +2002/07/29/big/img_415 +2002/07/30/big/img_923 +2002/08/18/big/img_575 +2002/08/22/big/img_182 +2002/07/25/big/img_806 +2002/07/22/big/img_49 +2002/07/29/big/img_989 +2003/01/17/big/img_789 +2003/01/15/big/img_503 +2002/09/01/big/img_16062 +2003/01/17/big/img_794 +2002/08/15/big/img_564 +2003/01/15/big/img_222 +2002/08/01/big/img_1656 +2003/01/13/big/img_432 +2002/07/19/big/img_426 +2002/08/17/big/img_244 +2002/08/13/big/img_805 +2002/09/02/big/img_15067 +2002/08/11/big/img_58 +2002/08/22/big/img_636 +2002/07/22/big/img_416 +2002/08/13/big/img_836 +2002/08/26/big/img_363 +2002/07/30/big/img_917 +2003/01/14/big/img_206 +2002/08/12/big/img_311 +2002/08/31/big/img_17623 +2002/07/29/big/img_661 +2003/01/13/big/img_417 +2002/08/02/big/img_463 +2002/08/02/big/img_669 +2002/08/26/big/img_670 +2002/08/02/big/img_375 +2002/07/19/big/img_209 +2002/08/08/big/img_115 +2002/08/21/big/img_399 +2002/08/20/big/img_911 +2002/08/07/big/img_1212 +2002/08/20/big/img_578 +2002/08/22/big/img_554 +2002/08/21/big/img_484 +2002/07/25/big/img_450 +2002/08/03/big/img_542 +2002/08/15/big/img_561 +2002/07/23/big/img_360 +2002/08/30/big/img_18137 +2002/07/25/big/img_250 +2002/08/03/big/img_647 +2002/08/20/big/img_375 +2002/08/14/big/img_387 +2002/09/01/big/img_16990 +2002/08/28/big/img_19341 +2003/01/15/big/img_239 +2002/08/20/big/img_528 +2002/08/12/big/img_130 +2002/09/02/big/img_15108 +2003/01/15/big/img_372 +2002/08/16/big/img_678 +2002/08/04/big/img_623 +2002/07/23/big/img_477 +2002/08/28/big/img_19590 +2003/01/17/big/img_978 +2002/09/01/big/img_16692 +2002/07/20/big/img_109 +2002/08/06/big/img_2660 +2003/01/14/big/img_464 +2002/08/09/big/img_618 +2002/07/22/big/img_722 +2002/08/25/big/img_419 +2002/08/03/big/img_314 +2002/08/25/big/img_40 +2002/07/27/big/img_430 +2002/08/10/big/img_569 +2002/08/23/big/img_398 +2002/07/23/big/img_893 +2002/08/16/big/img_261 +2002/08/06/big/img_2668 +2002/07/22/big/img_835 +2002/09/02/big/img_15093 +2003/01/16/big/img_65 +2002/08/21/big/img_448 +2003/01/14/big/img_351 +2003/01/17/big/img_133 +2002/07/28/big/img_493 +2003/01/15/big/img_640 +2002/09/01/big/img_16880 +2002/08/15/big/img_350 +2002/08/20/big/img_624 +2002/08/25/big/img_604 +2002/08/06/big/img_2200 +2002/08/23/big/img_290 +2002/08/13/big/img_1152 +2003/01/14/big/img_251 +2002/08/02/big/img_538 +2002/08/22/big/img_613 +2003/01/13/big/img_351 +2002/08/18/big/img_368 +2002/07/23/big/img_392 +2002/07/25/big/img_198 +2002/07/25/big/img_418 +2002/08/26/big/img_614 +2002/07/23/big/img_405 +2003/01/14/big/img_445 +2002/07/25/big/img_326 +2002/08/10/big/img_734 +2003/01/14/big/img_530 +2002/08/08/big/img_561 +2002/08/29/big/img_18990 +2002/08/10/big/img_576 +2002/07/29/big/img_1494 +2002/07/19/big/img_198 +2002/08/10/big/img_562 +2002/07/22/big/img_901 +2003/01/14/big/img_37 +2002/09/02/big/img_15629 +2003/01/14/big/img_58 +2002/08/01/big/img_1364 +2002/07/27/big/img_636 +2003/01/13/big/img_241 +2002/09/01/big/img_16988 +2003/01/13/big/img_560 +2002/08/09/big/img_533 +2002/07/31/big/img_249 +2003/01/17/big/img_1007 +2002/07/21/big/img_64 +2003/01/13/big/img_537 +2003/01/15/big/img_606 +2002/08/18/big/img_651 +2002/08/24/big/img_405 +2002/07/26/big/img_837 +2002/08/09/big/img_562 +2002/08/01/big/img_1983 +2002/08/03/big/img_514 +2002/07/29/big/img_314 +2002/08/12/big/img_493 +2003/01/14/big/img_121 +2003/01/14/big/img_479 +2002/08/04/big/img_410 +2002/07/22/big/img_607 +2003/01/17/big/img_417 +2002/07/20/big/img_547 +2002/08/13/big/img_396 +2002/08/31/big/img_17538 +2002/08/13/big/img_187 +2002/08/12/big/img_328 +2003/01/14/big/img_569 +2002/07/27/big/img_1081 +2002/08/14/big/img_504 +2002/08/23/big/img_785 +2002/07/26/big/img_339 +2002/08/07/big/img_1156 +2002/08/07/big/img_1456 +2002/08/23/big/img_378 +2002/08/27/big/img_19719 +2002/07/31/big/img_39 +2002/07/31/big/img_883 +2003/01/14/big/img_676 +2002/07/29/big/img_214 +2002/07/26/big/img_669 +2002/07/25/big/img_202 +2002/08/08/big/img_259 +2003/01/17/big/img_943 +2003/01/15/big/img_512 +2002/08/05/big/img_3295 +2002/08/27/big/img_19685 +2002/08/08/big/img_277 +2002/08/30/big/img_18154 +2002/07/22/big/img_663 +2002/08/29/big/img_18914 +2002/07/31/big/img_908 +2002/08/27/big/img_19926 +2003/01/13/big/img_791 +2003/01/15/big/img_827 +2002/08/18/big/img_878 +2002/08/14/big/img_670 +2002/07/20/big/img_182 +2002/08/15/big/img_291 +2002/08/06/big/img_2600 +2002/07/23/big/img_587 +2002/08/14/big/img_577 +2003/01/15/big/img_585 +2002/07/30/big/img_310 +2002/08/03/big/img_658 +2002/08/10/big/img_157 +2002/08/19/big/img_811 +2002/07/29/big/img_1318 +2002/08/04/big/img_104 +2002/07/30/big/img_332 +2002/07/24/big/img_789 +2002/07/29/big/img_516 +2002/07/23/big/img_843 +2002/08/01/big/img_1528 +2002/08/13/big/img_798 +2002/08/07/big/img_1729 +2002/08/28/big/img_19448 +2003/01/16/big/img_95 +2002/08/12/big/img_473 +2002/07/27/big/img_269 +2003/01/16/big/img_621 +2002/07/29/big/img_772 +2002/07/24/big/img_171 +2002/07/19/big/img_429 +2002/08/07/big/img_1933 +2002/08/27/big/img_19629 +2002/08/05/big/img_3688 +2002/08/07/big/img_1691 +2002/07/23/big/img_600 +2002/07/29/big/img_666 +2002/08/25/big/img_566 +2002/08/06/big/img_2659 +2002/08/29/big/img_18929 +2002/08/16/big/img_407 +2002/08/18/big/img_774 +2002/08/19/big/img_249 +2002/08/06/big/img_2427 +2002/08/29/big/img_18899 +2002/08/01/big/img_1818 +2002/07/31/big/img_108 +2002/07/29/big/img_500 +2002/08/11/big/img_115 +2002/07/19/big/img_521 +2002/08/02/big/img_1163 +2002/07/22/big/img_62 +2002/08/13/big/img_466 +2002/08/21/big/img_956 +2002/08/23/big/img_602 +2002/08/20/big/img_858 +2002/07/25/big/img_690 +2002/07/19/big/img_130 +2002/08/04/big/img_874 +2002/07/26/big/img_489 +2002/07/22/big/img_548 +2002/08/10/big/img_191 +2002/07/25/big/img_1051 +2002/08/18/big/img_473 +2002/08/12/big/img_755 +2002/08/18/big/img_413 +2002/08/08/big/img_1044 +2002/08/17/big/img_680 +2002/08/26/big/img_235 +2002/08/20/big/img_330 +2002/08/22/big/img_344 +2002/08/09/big/img_593 +2002/07/31/big/img_1006 +2002/08/14/big/img_337 +2002/08/16/big/img_728 +2002/07/24/big/img_834 +2002/08/04/big/img_552 +2002/09/02/big/img_15213 +2002/07/25/big/img_725 +2002/08/30/big/img_18290 +2003/01/01/big/img_475 +2002/07/27/big/img_1083 +2002/08/29/big/img_18955 +2002/08/31/big/img_17232 +2002/08/08/big/img_480 +2002/08/01/big/img_1311 +2002/07/30/big/img_745 +2002/08/03/big/img_649 +2002/08/12/big/img_193 +2002/07/29/big/img_228 +2002/07/25/big/img_836 +2002/08/20/big/img_400 +2002/07/30/big/img_507 +2002/09/02/big/img_15072 +2002/07/26/big/img_658 +2002/07/28/big/img_503 +2002/08/05/big/img_3814 +2002/08/24/big/img_745 +2003/01/13/big/img_817 +2002/08/08/big/img_579 +2002/07/22/big/img_251 +2003/01/13/big/img_689 +2002/07/25/big/img_407 +2002/08/13/big/img_1050 +2002/08/14/big/img_733 +2002/07/24/big/img_82 +2003/01/17/big/img_288 +2003/01/15/big/img_475 +2002/08/14/big/img_620 +2002/08/21/big/img_167 +2002/07/19/big/img_300 +2002/07/26/big/img_219 +2002/08/01/big/img_1468 +2002/07/23/big/img_260 +2002/08/09/big/img_555 +2002/07/19/big/img_160 +2002/08/02/big/img_1060 +2003/01/14/big/img_149 +2002/08/15/big/img_346 +2002/08/24/big/img_597 +2002/08/22/big/img_502 +2002/08/30/big/img_18228 +2002/07/21/big/img_766 +2003/01/15/big/img_841 +2002/07/24/big/img_516 +2002/08/02/big/img_265 +2002/08/15/big/img_1243 +2003/01/15/big/img_223 +2002/08/04/big/img_236 +2002/07/22/big/img_309 +2002/07/20/big/img_656 +2002/07/31/big/img_412 +2002/09/01/big/img_16462 +2003/01/16/big/img_431 +2002/07/22/big/img_793 +2002/08/15/big/img_877 +2002/07/26/big/img_282 +2002/07/25/big/img_529 +2002/08/24/big/img_613 +2003/01/17/big/img_700 +2002/08/06/big/img_2526 +2002/08/24/big/img_394 +2002/08/21/big/img_521 +2002/08/25/big/img_560 +2002/07/29/big/img_966 +2002/07/25/big/img_448 +2003/01/13/big/img_782 +2002/08/21/big/img_296 +2002/09/01/big/img_16755 +2002/08/05/big/img_3552 +2002/09/02/big/img_15823 +2003/01/14/big/img_193 +2002/07/21/big/img_159 +2002/08/02/big/img_564 +2002/08/16/big/img_300 +2002/07/19/big/img_269 +2002/08/13/big/img_676 +2002/07/28/big/img_57 +2002/08/05/big/img_3318 +2002/07/31/big/img_218 +2002/08/21/big/img_898 +2002/07/29/big/img_109 +2002/07/19/big/img_854 +2002/08/23/big/img_311 +2002/08/14/big/img_318 +2002/07/25/big/img_523 +2002/07/21/big/img_678 +2003/01/17/big/img_690 +2002/08/28/big/img_19503 +2002/08/18/big/img_251 +2002/08/22/big/img_672 +2002/08/20/big/img_663 +2002/08/02/big/img_148 +2002/09/02/big/img_15580 +2002/07/25/big/img_778 +2002/08/14/big/img_565 +2002/08/12/big/img_374 +2002/08/13/big/img_1018 +2002/08/20/big/img_474 +2002/08/25/big/img_33 +2002/08/02/big/img_1190 +2002/08/08/big/img_864 +2002/08/14/big/img_1071 +2002/08/30/big/img_18103 +2002/08/18/big/img_533 +2003/01/16/big/img_650 +2002/07/25/big/img_108 +2002/07/26/big/img_81 +2002/07/27/big/img_543 +2002/07/29/big/img_521 +2003/01/13/big/img_434 +2002/08/26/big/img_674 +2002/08/06/big/img_2932 +2002/08/07/big/img_1262 +2003/01/15/big/img_201 +2003/01/16/big/img_673 +2002/09/02/big/img_15988 +2002/07/29/big/img_1306 +2003/01/14/big/img_1072 +2002/08/30/big/img_18232 +2002/08/05/big/img_3711 +2002/07/23/big/img_775 +2002/08/01/big/img_16 +2003/01/16/big/img_630 +2002/08/22/big/img_695 +2002/08/14/big/img_51 +2002/08/14/big/img_782 +2002/08/24/big/img_742 +2003/01/14/big/img_512 +2003/01/15/big/img_1183 +2003/01/15/big/img_714 +2002/08/01/big/img_2078 +2002/07/31/big/img_682 +2002/09/02/big/img_15687 +2002/07/26/big/img_518 +2002/08/27/big/img_19676 +2002/09/02/big/img_15969 +2002/08/02/big/img_931 +2002/08/25/big/img_508 +2002/08/29/big/img_18616 +2002/07/22/big/img_839 +2002/07/28/big/img_313 +2003/01/14/big/img_155 +2002/08/02/big/img_1105 +2002/08/09/big/img_53 +2002/08/16/big/img_469 +2002/08/15/big/img_502 +2002/08/20/big/img_575 +2002/07/25/big/img_138 +2003/01/16/big/img_579 +2002/07/19/big/img_352 +2003/01/14/big/img_762 +2003/01/01/big/img_588 +2002/08/02/big/img_981 +2002/08/21/big/img_447 +2002/09/01/big/img_16151 +2003/01/14/big/img_769 +2002/08/23/big/img_461 +2002/08/17/big/img_240 +2002/09/02/big/img_15220 +2002/07/19/big/img_408 +2002/09/02/big/img_15496 +2002/07/29/big/img_758 +2002/08/28/big/img_19392 +2002/08/06/big/img_2723 +2002/08/31/big/img_17752 +2002/08/23/big/img_469 +2002/08/13/big/img_515 +2002/09/02/big/img_15551 +2002/08/03/big/img_462 +2002/07/24/big/img_613 +2002/07/22/big/img_61 +2002/08/08/big/img_171 +2002/08/21/big/img_177 +2003/01/14/big/img_105 +2002/08/02/big/img_1017 +2002/08/22/big/img_106 +2002/07/27/big/img_542 +2002/07/21/big/img_665 +2002/07/23/big/img_595 +2002/08/04/big/img_657 +2002/08/29/big/img_19002 +2003/01/15/big/img_550 +2002/08/14/big/img_662 +2002/07/20/big/img_425 +2002/08/30/big/img_18528 +2002/07/26/big/img_611 +2002/07/22/big/img_849 +2002/08/07/big/img_1655 +2002/08/21/big/img_638 +2003/01/17/big/img_732 +2003/01/01/big/img_496 +2002/08/18/big/img_713 +2002/08/08/big/img_109 +2002/07/27/big/img_1008 +2002/07/20/big/img_559 +2002/08/16/big/img_699 +2002/08/31/big/img_17702 +2002/07/31/big/img_1013 +2002/08/01/big/img_2027 +2002/08/02/big/img_1001 +2002/08/03/big/img_210 +2002/08/01/big/img_2087 +2003/01/14/big/img_199 +2002/07/29/big/img_48 +2002/07/19/big/img_727 +2002/08/09/big/img_249 +2002/08/04/big/img_632 +2002/08/22/big/img_620 +2003/01/01/big/img_457 +2002/08/05/big/img_3223 +2002/07/27/big/img_240 +2002/07/25/big/img_797 +2002/08/13/big/img_430 +2002/07/25/big/img_615 +2002/08/12/big/img_28 +2002/07/30/big/img_220 +2002/07/24/big/img_89 +2002/08/21/big/img_357 +2002/08/09/big/img_590 +2003/01/13/big/img_525 +2002/08/17/big/img_818 +2003/01/02/big/img_7 +2002/07/26/big/img_636 +2003/01/13/big/img_1122 +2002/07/23/big/img_810 +2002/08/20/big/img_888 +2002/07/27/big/img_3 +2002/08/15/big/img_451 +2002/09/02/big/img_15787 +2002/07/31/big/img_281 +2002/08/05/big/img_3274 +2002/08/07/big/img_1254 +2002/07/31/big/img_27 +2002/08/01/big/img_1366 +2002/07/30/big/img_182 +2002/08/27/big/img_19690 +2002/07/29/big/img_68 +2002/08/23/big/img_754 +2002/07/30/big/img_540 +2002/08/27/big/img_20063 +2002/08/14/big/img_471 +2002/08/02/big/img_615 +2002/07/30/big/img_186 +2002/08/25/big/img_150 +2002/07/27/big/img_626 +2002/07/20/big/img_225 +2003/01/15/big/img_1252 +2002/07/19/big/img_367 +2003/01/15/big/img_582 +2002/08/09/big/img_572 +2002/08/08/big/img_428 +2003/01/15/big/img_639 +2002/08/28/big/img_19245 +2002/07/24/big/img_321 +2002/08/02/big/img_662 +2002/08/08/big/img_1033 +2003/01/17/big/img_867 +2002/07/22/big/img_652 +2003/01/14/big/img_224 +2002/08/18/big/img_49 +2002/07/26/big/img_46 +2002/08/31/big/img_18021 +2002/07/25/big/img_151 +2002/08/23/big/img_540 +2002/08/25/big/img_693 +2002/07/23/big/img_340 +2002/07/28/big/img_117 +2002/09/02/big/img_15768 +2002/08/26/big/img_562 +2002/07/24/big/img_480 +2003/01/15/big/img_341 +2002/08/10/big/img_783 +2002/08/20/big/img_132 +2003/01/14/big/img_370 +2002/07/20/big/img_720 +2002/08/03/big/img_144 +2002/08/20/big/img_538 +2002/08/01/big/img_1745 +2002/08/11/big/img_683 +2002/08/03/big/img_328 +2002/08/10/big/img_793 +2002/08/14/big/img_689 +2002/08/02/big/img_162 +2003/01/17/big/img_411 +2002/07/31/big/img_361 +2002/08/15/big/img_289 +2002/08/08/big/img_254 +2002/08/15/big/img_996 +2002/08/20/big/img_785 +2002/07/24/big/img_511 +2002/08/06/big/img_2614 +2002/08/29/big/img_18733 +2002/08/17/big/img_78 +2002/07/30/big/img_378 +2002/08/31/big/img_17947 +2002/08/26/big/img_88 +2002/07/30/big/img_558 +2002/08/02/big/img_67 +2003/01/14/big/img_325 +2002/07/29/big/img_1357 +2002/07/19/big/img_391 +2002/07/30/big/img_307 +2003/01/13/big/img_219 +2002/07/24/big/img_807 +2002/08/23/big/img_543 +2002/08/29/big/img_18620 +2002/07/22/big/img_769 +2002/08/26/big/img_503 +2002/07/30/big/img_78 +2002/08/14/big/img_1036 +2002/08/09/big/img_58 +2002/07/24/big/img_616 +2002/08/02/big/img_464 +2002/07/26/big/img_576 +2002/07/22/big/img_273 +2003/01/16/big/img_470 +2002/07/29/big/img_329 +2002/07/30/big/img_1086 +2002/07/31/big/img_353 +2002/09/02/big/img_15275 +2003/01/17/big/img_555 +2002/08/26/big/img_212 +2002/08/01/big/img_1692 +2003/01/15/big/img_600 +2002/07/29/big/img_825 +2002/08/08/big/img_68 +2002/08/10/big/img_719 +2002/07/31/big/img_636 +2002/07/29/big/img_325 +2002/07/21/big/img_515 +2002/07/22/big/img_705 +2003/01/13/big/img_818 +2002/08/09/big/img_486 +2002/08/22/big/img_141 +2002/07/22/big/img_303 +2002/08/09/big/img_393 +2002/07/29/big/img_963 +2002/08/02/big/img_1215 +2002/08/19/big/img_674 +2002/08/12/big/img_690 +2002/08/21/big/img_637 +2002/08/21/big/img_841 +2002/08/24/big/img_71 +2002/07/25/big/img_596 +2002/07/24/big/img_864 +2002/08/18/big/img_293 +2003/01/14/big/img_657 +2002/08/15/big/img_411 +2002/08/16/big/img_348 +2002/08/05/big/img_3157 +2002/07/20/big/img_663 +2003/01/13/big/img_654 +2003/01/16/big/img_433 +2002/08/30/big/img_18200 +2002/08/12/big/img_226 +2003/01/16/big/img_491 +2002/08/08/big/img_666 +2002/07/19/big/img_576 +2003/01/15/big/img_776 +2003/01/16/big/img_899 +2002/07/19/big/img_397 +2002/08/14/big/img_44 +2003/01/15/big/img_762 +2002/08/02/big/img_982 +2002/09/02/big/img_15234 +2002/08/17/big/img_556 +2002/08/21/big/img_410 +2002/08/21/big/img_386 +2002/07/19/big/img_690 +2002/08/05/big/img_3052 +2002/08/14/big/img_219 +2002/08/16/big/img_273 +2003/01/15/big/img_752 +2002/08/08/big/img_184 +2002/07/31/big/img_743 +2002/08/23/big/img_338 +2003/01/14/big/img_1055 +2002/08/05/big/img_3405 +2003/01/15/big/img_17 +2002/08/03/big/img_141 +2002/08/14/big/img_549 +2002/07/27/big/img_1034 +2002/07/31/big/img_932 +2002/08/30/big/img_18487 +2002/09/02/big/img_15814 +2002/08/01/big/img_2086 +2002/09/01/big/img_16535 +2002/07/22/big/img_500 +2003/01/13/big/img_400 +2002/08/25/big/img_607 +2002/08/30/big/img_18384 +2003/01/14/big/img_951 +2002/08/13/big/img_1150 +2002/08/08/big/img_1022 +2002/08/10/big/img_428 +2002/08/28/big/img_19242 +2002/08/05/big/img_3098 +2002/07/23/big/img_400 +2002/08/26/big/img_365 +2002/07/20/big/img_318 +2002/08/13/big/img_740 +2003/01/16/big/img_37 +2002/08/26/big/img_274 +2002/08/02/big/img_205 +2002/08/21/big/img_695 +2002/08/06/big/img_2289 +2002/08/20/big/img_794 +2002/08/18/big/img_438 +2002/08/07/big/img_1380 +2002/08/02/big/img_737 +2002/08/07/big/img_1651 +2002/08/15/big/img_1238 +2002/08/01/big/img_1681 +2002/08/06/big/img_3017 +2002/07/23/big/img_706 +2002/07/31/big/img_392 +2002/08/09/big/img_539 +2002/07/29/big/img_835 +2002/08/26/big/img_723 +2002/08/28/big/img_19235 +2003/01/16/big/img_353 +2002/08/10/big/img_150 +2002/08/29/big/img_19025 +2002/08/21/big/img_310 +2002/08/10/big/img_823 +2002/07/26/big/img_981 +2002/08/11/big/img_288 +2002/08/19/big/img_534 +2002/08/21/big/img_300 +2002/07/31/big/img_49 +2002/07/30/big/img_469 +2002/08/28/big/img_19197 +2002/08/25/big/img_205 +2002/08/10/big/img_390 +2002/08/23/big/img_291 +2002/08/26/big/img_230 +2002/08/18/big/img_76 +2002/07/23/big/img_409 +2002/08/14/big/img_1053 +2003/01/14/big/img_291 +2002/08/10/big/img_503 +2002/08/27/big/img_19928 +2002/08/03/big/img_563 +2002/08/17/big/img_250 +2002/08/06/big/img_2381 +2002/08/17/big/img_948 +2002/08/06/big/img_2710 +2002/07/22/big/img_696 +2002/07/31/big/img_670 +2002/08/12/big/img_594 +2002/07/29/big/img_624 +2003/01/17/big/img_934 +2002/08/03/big/img_584 +2002/08/22/big/img_1003 +2002/08/05/big/img_3396 +2003/01/13/big/img_570 +2002/08/02/big/img_219 +2002/09/02/big/img_15774 +2002/08/16/big/img_818 +2002/08/23/big/img_402 +2003/01/14/big/img_552 +2002/07/29/big/img_71 +2002/08/05/big/img_3592 +2002/08/16/big/img_80 +2002/07/27/big/img_672 +2003/01/13/big/img_470 +2003/01/16/big/img_702 +2002/09/01/big/img_16130 +2002/08/08/big/img_240 +2002/09/01/big/img_16338 +2002/07/26/big/img_312 +2003/01/14/big/img_538 +2002/07/20/big/img_695 +2002/08/30/big/img_18098 +2002/08/25/big/img_259 +2002/08/16/big/img_1042 +2002/08/09/big/img_837 +2002/08/31/big/img_17760 +2002/07/31/big/img_14 +2002/08/09/big/img_361 +2003/01/16/big/img_107 +2002/08/14/big/img_124 +2002/07/19/big/img_463 +2003/01/15/big/img_275 +2002/07/25/big/img_1151 +2002/07/29/big/img_1501 +2002/08/27/big/img_19889 +2002/08/29/big/img_18603 +2003/01/17/big/img_601 +2002/08/25/big/img_355 +2002/08/08/big/img_297 +2002/08/20/big/img_290 +2002/07/31/big/img_195 +2003/01/01/big/img_336 +2002/08/18/big/img_369 +2002/07/25/big/img_621 +2002/08/11/big/img_508 +2003/01/14/big/img_458 +2003/01/15/big/img_795 +2002/08/12/big/img_498 +2002/08/01/big/img_1734 +2002/08/02/big/img_246 +2002/08/16/big/img_565 +2002/08/11/big/img_475 +2002/08/22/big/img_408 +2002/07/28/big/img_78 +2002/07/21/big/img_81 +2003/01/14/big/img_697 +2002/08/14/big/img_661 +2002/08/15/big/img_507 +2002/08/19/big/img_55 +2002/07/22/big/img_152 +2003/01/14/big/img_470 +2002/08/03/big/img_379 +2002/08/22/big/img_506 +2003/01/16/big/img_966 +2002/08/18/big/img_698 +2002/08/24/big/img_528 +2002/08/23/big/img_10 +2002/08/01/big/img_1655 +2002/08/22/big/img_953 +2002/07/19/big/img_630 +2002/07/22/big/img_889 +2002/08/16/big/img_351 +2003/01/16/big/img_83 +2002/07/19/big/img_805 +2002/08/14/big/img_704 +2002/07/19/big/img_389 +2002/08/31/big/img_17765 +2002/07/29/big/img_606 +2003/01/17/big/img_939 +2002/09/02/big/img_15081 +2002/08/21/big/img_181 +2002/07/29/big/img_1321 +2002/07/21/big/img_497 +2002/07/20/big/img_539 +2002/08/24/big/img_119 +2002/08/01/big/img_1281 +2002/07/26/big/img_207 +2002/07/26/big/img_432 +2002/07/27/big/img_1006 +2002/08/05/big/img_3087 +2002/08/14/big/img_252 +2002/08/14/big/img_798 +2002/07/24/big/img_538 +2002/09/02/big/img_15507 +2002/08/08/big/img_901 +2003/01/14/big/img_557 +2002/08/07/big/img_1819 +2002/08/04/big/img_470 +2002/08/01/big/img_1504 +2002/08/16/big/img_1070 +2002/08/16/big/img_372 +2002/08/23/big/img_416 +2002/08/30/big/img_18208 +2002/08/01/big/img_2043 +2002/07/22/big/img_385 +2002/08/22/big/img_466 +2002/08/21/big/img_869 +2002/08/28/big/img_19429 +2002/08/02/big/img_770 +2002/07/23/big/img_433 +2003/01/14/big/img_13 +2002/07/27/big/img_953 +2002/09/02/big/img_15728 +2002/08/01/big/img_1361 +2002/08/29/big/img_18897 +2002/08/26/big/img_534 +2002/08/11/big/img_121 +2002/08/26/big/img_20130 +2002/07/31/big/img_363 +2002/08/13/big/img_978 +2002/07/25/big/img_835 +2002/08/02/big/img_906 +2003/01/14/big/img_548 +2002/07/30/big/img_80 +2002/07/26/big/img_982 +2003/01/16/big/img_99 +2002/08/19/big/img_362 +2002/08/24/big/img_376 +2002/08/07/big/img_1264 +2002/07/27/big/img_938 +2003/01/17/big/img_535 +2002/07/26/big/img_457 +2002/08/08/big/img_848 +2003/01/15/big/img_859 +2003/01/15/big/img_622 +2002/07/30/big/img_403 +2002/07/29/big/img_217 +2002/07/26/big/img_891 +2002/07/24/big/img_70 +2002/08/25/big/img_619 +2002/08/05/big/img_3375 +2002/08/01/big/img_2160 +2002/08/06/big/img_2227 +2003/01/14/big/img_117 +2002/08/14/big/img_227 +2002/08/13/big/img_565 +2002/08/19/big/img_625 +2002/08/03/big/img_812 +2002/07/24/big/img_41 +2002/08/16/big/img_235 +2002/07/29/big/img_759 +2002/07/21/big/img_433 +2002/07/29/big/img_190 +2003/01/16/big/img_435 +2003/01/13/big/img_708 +2002/07/30/big/img_57 +2002/08/22/big/img_162 +2003/01/01/big/img_558 +2003/01/15/big/img_604 +2002/08/16/big/img_935 +2002/08/20/big/img_394 +2002/07/28/big/img_465 +2002/09/02/big/img_15534 +2002/08/16/big/img_87 +2002/07/22/big/img_469 +2002/08/12/big/img_245 +2003/01/13/big/img_236 +2002/08/06/big/img_2736 +2002/08/03/big/img_348 +2003/01/14/big/img_218 +2002/07/26/big/img_232 +2003/01/15/big/img_244 +2002/07/25/big/img_1121 +2002/08/01/big/img_1484 +2002/07/26/big/img_541 +2002/08/07/big/img_1244 +2002/07/31/big/img_3 +2002/08/30/big/img_18437 +2002/08/29/big/img_19094 +2002/08/01/big/img_1355 +2002/08/19/big/img_338 +2002/07/19/big/img_255 +2002/07/21/big/img_76 +2002/08/25/big/img_199 +2002/08/12/big/img_740 +2002/07/30/big/img_852 +2002/08/15/big/img_599 +2002/08/23/big/img_254 +2002/08/19/big/img_125 +2002/07/24/big/img_2 +2002/08/04/big/img_145 +2002/08/05/big/img_3137 +2002/07/28/big/img_463 +2003/01/14/big/img_801 +2002/07/23/big/img_366 +2002/08/26/big/img_600 +2002/08/26/big/img_649 +2002/09/02/big/img_15849 +2002/07/26/big/img_248 +2003/01/13/big/img_200 +2002/08/07/big/img_1794 +2002/08/31/big/img_17270 +2002/08/23/big/img_608 +2003/01/13/big/img_837 +2002/08/23/big/img_581 +2002/08/20/big/img_754 +2002/08/18/big/img_183 +2002/08/20/big/img_328 +2002/07/22/big/img_494 +2002/07/29/big/img_399 +2002/08/28/big/img_19284 +2002/08/08/big/img_566 +2002/07/25/big/img_376 +2002/07/23/big/img_138 +2002/07/25/big/img_435 +2002/08/17/big/img_685 +2002/07/19/big/img_90 +2002/07/20/big/img_716 +2002/08/31/big/img_17458 +2002/08/26/big/img_461 +2002/07/25/big/img_355 +2002/08/06/big/img_2152 +2002/07/27/big/img_932 +2002/07/23/big/img_232 +2002/08/08/big/img_1020 +2002/07/31/big/img_366 +2002/08/06/big/img_2667 +2002/08/21/big/img_465 +2002/08/15/big/img_305 +2002/08/02/big/img_247 +2002/07/28/big/img_46 +2002/08/27/big/img_19922 +2002/08/23/big/img_643 +2003/01/13/big/img_624 +2002/08/23/big/img_625 +2002/08/05/big/img_3787 +2003/01/13/big/img_627 +2002/09/01/big/img_16381 +2002/08/05/big/img_3668 +2002/07/21/big/img_535 +2002/08/27/big/img_19680 +2002/07/22/big/img_413 +2002/07/29/big/img_481 +2003/01/15/big/img_496 +2002/07/23/big/img_701 +2002/08/29/big/img_18670 +2002/07/28/big/img_319 +2003/01/14/big/img_517 +2002/07/26/big/img_256 +2003/01/16/big/img_593 +2002/07/30/big/img_956 +2002/07/30/big/img_667 +2002/07/25/big/img_100 +2002/08/11/big/img_570 +2002/07/26/big/img_745 +2002/08/04/big/img_834 +2002/08/25/big/img_521 +2002/08/01/big/img_2148 +2002/09/02/big/img_15183 +2002/08/22/big/img_514 +2002/08/23/big/img_477 +2002/07/23/big/img_336 +2002/07/26/big/img_481 +2002/08/20/big/img_409 +2002/07/23/big/img_918 +2002/08/09/big/img_474 +2002/08/02/big/img_929 +2002/08/31/big/img_17932 +2002/08/19/big/img_161 +2002/08/09/big/img_667 +2002/07/31/big/img_805 +2002/09/02/big/img_15678 +2002/08/31/big/img_17509 +2002/08/29/big/img_18998 +2002/07/23/big/img_301 +2002/08/07/big/img_1612 +2002/08/06/big/img_2472 +2002/07/23/big/img_466 +2002/08/27/big/img_19634 +2003/01/16/big/img_16 +2002/08/14/big/img_193 +2002/08/21/big/img_340 +2002/08/27/big/img_19799 +2002/08/01/big/img_1345 +2002/08/07/big/img_1448 +2002/08/11/big/img_324 +2003/01/16/big/img_754 +2002/08/13/big/img_418 +2003/01/16/big/img_544 +2002/08/19/big/img_135 +2002/08/10/big/img_455 +2002/08/10/big/img_693 +2002/08/31/big/img_17967 +2002/08/28/big/img_19229 +2002/08/04/big/img_811 +2002/09/01/big/img_16225 +2003/01/16/big/img_428 +2002/09/02/big/img_15295 +2002/07/26/big/img_108 +2002/07/21/big/img_477 +2002/08/07/big/img_1354 +2002/08/23/big/img_246 +2002/08/16/big/img_652 +2002/07/27/big/img_553 +2002/07/31/big/img_346 +2002/08/04/big/img_537 +2002/08/08/big/img_498 +2002/08/29/big/img_18956 +2003/01/13/big/img_922 +2002/08/31/big/img_17425 +2002/07/26/big/img_438 +2002/08/19/big/img_185 +2003/01/16/big/img_33 +2002/08/10/big/img_252 +2002/07/29/big/img_598 +2002/08/27/big/img_19820 +2002/08/06/big/img_2664 +2002/08/20/big/img_705 +2003/01/14/big/img_816 +2002/08/03/big/img_552 +2002/07/25/big/img_561 +2002/07/25/big/img_934 +2002/08/01/big/img_1893 +2003/01/14/big/img_746 +2003/01/16/big/img_519 +2002/08/03/big/img_681 +2002/07/24/big/img_808 +2002/08/14/big/img_803 +2002/08/25/big/img_155 +2002/07/30/big/img_1107 +2002/08/29/big/img_18882 +2003/01/15/big/img_598 +2002/08/19/big/img_122 +2002/07/30/big/img_428 +2002/07/24/big/img_684 +2002/08/22/big/img_192 +2002/08/22/big/img_543 +2002/08/07/big/img_1318 +2002/08/18/big/img_25 +2002/07/26/big/img_583 +2002/07/20/big/img_464 +2002/08/19/big/img_664 +2002/08/24/big/img_861 +2002/09/01/big/img_16136 +2002/08/22/big/img_400 +2002/08/12/big/img_445 +2003/01/14/big/img_174 +2002/08/27/big/img_19677 +2002/08/31/big/img_17214 +2002/08/30/big/img_18175 +2003/01/17/big/img_402 +2002/08/06/big/img_2396 +2002/08/18/big/img_448 +2002/08/21/big/img_165 +2002/08/31/big/img_17609 +2003/01/01/big/img_151 +2002/08/26/big/img_372 +2002/09/02/big/img_15994 +2002/07/26/big/img_660 +2002/09/02/big/img_15197 +2002/07/29/big/img_258 +2002/08/30/big/img_18525 +2003/01/13/big/img_368 +2002/07/29/big/img_1538 +2002/07/21/big/img_787 +2002/08/18/big/img_152 +2002/08/06/big/img_2379 +2003/01/17/big/img_864 +2002/08/27/big/img_19998 +2002/08/01/big/img_1634 +2002/07/25/big/img_414 +2002/08/22/big/img_627 +2002/08/07/big/img_1669 +2002/08/16/big/img_1052 +2002/08/31/big/img_17796 +2002/08/18/big/img_199 +2002/09/02/big/img_15147 +2002/08/09/big/img_460 +2002/08/14/big/img_581 +2002/08/30/big/img_18286 +2002/07/26/big/img_337 +2002/08/18/big/img_589 +2003/01/14/big/img_866 +2002/07/20/big/img_624 +2002/08/01/big/img_1801 +2002/07/24/big/img_683 +2002/08/09/big/img_725 +2003/01/14/big/img_34 +2002/07/30/big/img_144 +2002/07/30/big/img_706 +2002/08/08/big/img_394 +2002/08/19/big/img_619 +2002/08/06/big/img_2703 +2002/08/29/big/img_19034 +2002/07/24/big/img_67 +2002/08/27/big/img_19841 +2002/08/19/big/img_427 +2003/01/14/big/img_333 +2002/09/01/big/img_16406 +2002/07/19/big/img_882 +2002/08/17/big/img_238 +2003/01/14/big/img_739 +2002/07/22/big/img_151 +2002/08/21/big/img_743 +2002/07/25/big/img_1048 +2002/07/30/big/img_395 +2003/01/13/big/img_584 +2002/08/13/big/img_742 +2002/08/13/big/img_1168 +2003/01/14/big/img_147 +2002/07/26/big/img_803 +2002/08/05/big/img_3298 +2002/08/07/big/img_1451 +2002/08/16/big/img_424 +2002/07/29/big/img_1069 +2002/09/01/big/img_16735 +2002/07/21/big/img_637 +2003/01/14/big/img_585 +2002/08/02/big/img_358 +2003/01/13/big/img_358 +2002/08/14/big/img_198 +2002/08/17/big/img_935 +2002/08/04/big/img_42 +2002/08/30/big/img_18245 +2002/07/25/big/img_158 +2002/08/22/big/img_744 +2002/08/06/big/img_2291 +2002/08/05/big/img_3044 +2002/07/30/big/img_272 +2002/08/23/big/img_641 +2002/07/24/big/img_797 +2002/07/30/big/img_392 +2003/01/14/big/img_447 +2002/07/31/big/img_898 +2002/08/06/big/img_2812 +2002/08/13/big/img_564 +2002/07/22/big/img_43 +2002/07/26/big/img_634 +2002/07/19/big/img_843 +2002/08/26/big/img_58 +2002/07/21/big/img_375 +2002/08/25/big/img_729 +2002/07/19/big/img_561 +2003/01/15/big/img_884 +2002/07/25/big/img_891 +2002/08/09/big/img_558 +2002/08/26/big/img_587 +2002/08/13/big/img_1146 +2002/09/02/big/img_15153 +2002/07/26/big/img_316 +2002/08/01/big/img_1940 +2002/08/26/big/img_90 +2003/01/13/big/img_347 +2002/07/25/big/img_520 +2002/08/29/big/img_18718 +2002/08/28/big/img_19219 +2002/08/13/big/img_375 +2002/07/20/big/img_719 +2002/08/31/big/img_17431 +2002/07/28/big/img_192 +2002/08/26/big/img_259 +2002/08/18/big/img_484 +2002/07/29/big/img_580 +2002/07/26/big/img_84 +2002/08/02/big/img_302 +2002/08/31/big/img_17007 +2003/01/15/big/img_543 +2002/09/01/big/img_16488 +2002/08/22/big/img_798 +2002/07/30/big/img_383 +2002/08/04/big/img_668 +2002/08/13/big/img_156 +2002/08/07/big/img_1353 +2002/07/25/big/img_281 +2003/01/14/big/img_587 +2003/01/15/big/img_524 +2002/08/19/big/img_726 +2002/08/21/big/img_709 +2002/08/26/big/img_465 +2002/07/31/big/img_658 +2002/08/28/big/img_19148 +2002/07/23/big/img_423 +2002/08/16/big/img_758 +2002/08/22/big/img_523 +2002/08/16/big/img_591 +2002/08/23/big/img_845 +2002/07/26/big/img_678 +2002/08/09/big/img_806 +2002/08/06/big/img_2369 +2002/07/29/big/img_457 +2002/07/19/big/img_278 +2002/08/30/big/img_18107 +2002/07/26/big/img_444 +2002/08/20/big/img_278 +2002/08/26/big/img_92 +2002/08/26/big/img_257 +2002/07/25/big/img_266 +2002/08/05/big/img_3829 +2002/07/26/big/img_757 +2002/07/29/big/img_1536 +2002/08/09/big/img_472 +2003/01/17/big/img_480 +2002/08/28/big/img_19355 +2002/07/26/big/img_97 +2002/08/06/big/img_2503 +2002/07/19/big/img_254 +2002/08/01/big/img_1470 +2002/08/21/big/img_42 +2002/08/20/big/img_217 +2002/08/06/big/img_2459 +2002/07/19/big/img_552 +2002/08/13/big/img_717 +2002/08/12/big/img_586 +2002/08/20/big/img_411 +2003/01/13/big/img_768 +2002/08/07/big/img_1747 +2002/08/15/big/img_385 +2002/08/01/big/img_1648 +2002/08/15/big/img_311 +2002/08/21/big/img_95 +2002/08/09/big/img_108 +2002/08/21/big/img_398 +2002/08/17/big/img_340 +2002/08/14/big/img_474 +2002/08/13/big/img_294 +2002/08/24/big/img_840 +2002/08/09/big/img_808 +2002/08/23/big/img_491 +2002/07/28/big/img_33 +2003/01/13/big/img_664 +2002/08/02/big/img_261 +2002/08/09/big/img_591 +2002/07/26/big/img_309 +2003/01/14/big/img_372 +2002/08/19/big/img_581 +2002/08/19/big/img_168 +2002/08/26/big/img_422 +2002/07/24/big/img_106 +2002/08/01/big/img_1936 +2002/08/05/big/img_3764 +2002/08/21/big/img_266 +2002/08/31/big/img_17968 +2002/08/01/big/img_1941 +2002/08/15/big/img_550 +2002/08/14/big/img_13 +2002/07/30/big/img_171 +2003/01/13/big/img_490 +2002/07/25/big/img_427 +2002/07/19/big/img_770 +2002/08/12/big/img_759 +2003/01/15/big/img_1360 +2002/08/05/big/img_3692 +2003/01/16/big/img_30 +2002/07/25/big/img_1026 +2002/07/22/big/img_288 +2002/08/29/big/img_18801 +2002/07/24/big/img_793 +2002/08/13/big/img_178 +2002/08/06/big/img_2322 +2003/01/14/big/img_560 +2002/08/18/big/img_408 +2003/01/16/big/img_915 +2003/01/16/big/img_679 +2002/08/07/big/img_1552 +2002/08/29/big/img_19050 +2002/08/01/big/img_2172 +2002/07/31/big/img_30 +2002/07/30/big/img_1019 +2002/07/30/big/img_587 +2003/01/13/big/img_773 +2002/07/30/big/img_410 +2002/07/28/big/img_65 +2002/08/05/big/img_3138 +2002/07/23/big/img_541 +2002/08/22/big/img_963 +2002/07/27/big/img_657 +2002/07/30/big/img_1051 +2003/01/16/big/img_150 +2002/07/31/big/img_519 +2002/08/01/big/img_1961 +2002/08/05/big/img_3752 +2002/07/23/big/img_631 +2003/01/14/big/img_237 +2002/07/28/big/img_21 +2002/07/22/big/img_813 +2002/08/05/big/img_3563 +2003/01/17/big/img_620 +2002/07/19/big/img_523 +2002/07/30/big/img_904 +2002/08/29/big/img_18642 +2002/08/11/big/img_492 +2002/08/01/big/img_2130 +2002/07/25/big/img_618 +2002/08/17/big/img_305 +2003/01/16/big/img_520 +2002/07/26/big/img_495 +2002/08/17/big/img_164 +2002/08/03/big/img_440 +2002/07/24/big/img_441 +2002/08/06/big/img_2146 +2002/08/11/big/img_558 +2002/08/02/big/img_545 +2002/08/31/big/img_18090 +2003/01/01/big/img_136 +2002/07/25/big/img_1099 +2003/01/13/big/img_728 +2003/01/16/big/img_197 +2002/07/26/big/img_651 +2002/08/11/big/img_676 +2003/01/15/big/img_10 +2002/08/21/big/img_250 +2002/08/14/big/img_325 +2002/08/04/big/img_390 +2002/07/24/big/img_554 +2003/01/16/big/img_333 +2002/07/31/big/img_922 +2002/09/02/big/img_15586 +2003/01/16/big/img_184 +2002/07/22/big/img_766 +2002/07/21/big/img_608 +2002/08/07/big/img_1578 +2002/08/17/big/img_961 +2002/07/27/big/img_324 +2002/08/05/big/img_3765 +2002/08/23/big/img_462 +2003/01/16/big/img_382 +2002/08/27/big/img_19838 +2002/08/01/big/img_1505 +2002/08/21/big/img_662 +2002/08/14/big/img_605 +2002/08/19/big/img_816 +2002/07/29/big/img_136 +2002/08/20/big/img_719 +2002/08/06/big/img_2826 +2002/08/10/big/img_630 +2003/01/17/big/img_973 +2002/08/14/big/img_116 +2002/08/02/big/img_666 +2002/08/21/big/img_710 +2002/08/05/big/img_55 +2002/07/31/big/img_229 +2002/08/01/big/img_1549 +2002/07/23/big/img_432 +2002/07/21/big/img_430 +2002/08/21/big/img_549 +2002/08/08/big/img_985 +2002/07/20/big/img_610 +2002/07/23/big/img_978 +2002/08/23/big/img_219 +2002/07/25/big/img_175 +2003/01/15/big/img_230 +2002/08/23/big/img_385 +2002/07/31/big/img_879 +2002/08/12/big/img_495 +2002/08/22/big/img_499 +2002/08/30/big/img_18322 +2002/08/15/big/img_795 +2002/08/13/big/img_835 +2003/01/17/big/img_930 +2002/07/30/big/img_873 +2002/08/11/big/img_257 +2002/07/31/big/img_593 +2002/08/21/big/img_916 +2003/01/13/big/img_814 +2002/07/25/big/img_722 +2002/08/16/big/img_379 +2002/07/31/big/img_497 +2002/07/22/big/img_602 +2002/08/21/big/img_642 +2002/08/21/big/img_614 +2002/08/23/big/img_482 +2002/07/29/big/img_603 +2002/08/13/big/img_705 +2002/07/23/big/img_833 +2003/01/14/big/img_511 +2002/07/24/big/img_376 +2002/08/17/big/img_1030 +2002/08/05/big/img_3576 +2002/08/16/big/img_540 +2002/07/22/big/img_630 +2002/08/10/big/img_180 +2002/08/14/big/img_905 +2002/08/29/big/img_18777 +2002/08/22/big/img_693 +2003/01/16/big/img_933 +2002/08/20/big/img_555 +2002/08/15/big/img_549 +2003/01/14/big/img_830 +2003/01/16/big/img_64 +2002/08/27/big/img_19670 +2002/08/22/big/img_729 +2002/07/27/big/img_981 +2002/08/09/big/img_458 +2003/01/17/big/img_884 +2002/07/25/big/img_639 +2002/08/31/big/img_18008 +2002/08/22/big/img_249 +2002/08/17/big/img_971 +2002/08/04/big/img_308 +2002/07/28/big/img_362 +2002/08/12/big/img_142 +2002/08/26/big/img_61 +2002/08/14/big/img_422 +2002/07/19/big/img_607 +2003/01/15/big/img_717 +2002/08/01/big/img_1475 +2002/08/29/big/img_19061 +2003/01/01/big/img_346 +2002/07/20/big/img_315 +2003/01/15/big/img_756 +2002/08/15/big/img_879 +2002/08/08/big/img_615 +2003/01/13/big/img_431 +2002/08/05/big/img_3233 +2002/08/24/big/img_526 +2003/01/13/big/img_717 +2002/09/01/big/img_16408 +2002/07/22/big/img_217 +2002/07/31/big/img_960 +2002/08/21/big/img_610 +2002/08/05/big/img_3753 +2002/08/03/big/img_151 +2002/08/21/big/img_267 +2002/08/01/big/img_2175 +2002/08/04/big/img_556 +2002/08/21/big/img_527 +2002/09/02/big/img_15800 +2002/07/27/big/img_156 +2002/07/20/big/img_590 +2002/08/15/big/img_700 +2002/08/08/big/img_444 +2002/07/25/big/img_94 +2002/07/24/big/img_778 +2002/08/14/big/img_694 +2002/07/20/big/img_666 +2002/08/02/big/img_200 +2002/08/02/big/img_578 +2003/01/17/big/img_332 +2002/09/01/big/img_16352 +2002/08/27/big/img_19668 +2002/07/23/big/img_823 +2002/08/13/big/img_431 +2003/01/16/big/img_463 +2002/08/27/big/img_19711 +2002/08/23/big/img_154 +2002/07/31/big/img_360 +2002/08/23/big/img_555 +2002/08/10/big/img_561 +2003/01/14/big/img_550 +2002/08/07/big/img_1370 +2002/07/30/big/img_1184 +2002/08/01/big/img_1445 +2002/08/23/big/img_22 +2002/07/30/big/img_606 +2003/01/17/big/img_271 +2002/08/31/big/img_17316 +2002/08/16/big/img_973 +2002/07/26/big/img_77 +2002/07/20/big/img_788 +2002/08/06/big/img_2426 +2002/08/07/big/img_1498 +2002/08/16/big/img_358 +2002/08/06/big/img_2851 +2002/08/12/big/img_359 +2002/08/01/big/img_1521 +2002/08/02/big/img_709 +2002/08/20/big/img_935 +2002/08/12/big/img_188 +2002/08/24/big/img_411 +2002/08/22/big/img_680 +2002/08/06/big/img_2480 +2002/07/20/big/img_627 +2002/07/30/big/img_214 +2002/07/25/big/img_354 +2002/08/02/big/img_636 +2003/01/15/big/img_661 +2002/08/07/big/img_1327 +2002/08/01/big/img_2108 +2002/08/31/big/img_17919 +2002/08/29/big/img_18768 +2002/08/05/big/img_3840 +2002/07/26/big/img_242 +2003/01/14/big/img_451 +2002/08/20/big/img_923 +2002/08/27/big/img_19908 +2002/08/16/big/img_282 +2002/08/19/big/img_440 +2003/01/01/big/img_230 +2002/08/08/big/img_212 +2002/07/20/big/img_443 +2002/08/25/big/img_635 +2003/01/13/big/img_1169 +2002/07/26/big/img_998 +2002/08/15/big/img_995 +2002/08/06/big/img_3002 +2002/07/29/big/img_460 +2003/01/14/big/img_925 +2002/07/23/big/img_539 +2002/08/16/big/img_694 +2003/01/13/big/img_459 +2002/07/23/big/img_249 +2002/08/20/big/img_539 +2002/08/04/big/img_186 +2002/08/26/big/img_264 +2002/07/22/big/img_704 +2002/08/25/big/img_277 +2002/08/22/big/img_988 +2002/07/29/big/img_504 +2002/08/05/big/img_3600 +2002/08/30/big/img_18380 +2003/01/14/big/img_937 +2002/08/21/big/img_254 +2002/08/10/big/img_130 +2002/08/20/big/img_339 +2003/01/14/big/img_428 +2002/08/20/big/img_889 +2002/08/31/big/img_17637 +2002/07/26/big/img_644 +2002/09/01/big/img_16776 +2002/08/06/big/img_2239 +2002/08/06/big/img_2646 +2003/01/13/big/img_491 +2002/08/10/big/img_579 +2002/08/21/big/img_713 +2002/08/22/big/img_482 +2002/07/22/big/img_167 +2002/07/24/big/img_539 +2002/08/14/big/img_721 +2002/07/25/big/img_389 +2002/09/01/big/img_16591 +2002/08/13/big/img_543 +2003/01/14/big/img_432 +2002/08/09/big/img_287 +2002/07/26/big/img_126 +2002/08/23/big/img_412 +2002/08/15/big/img_1034 +2002/08/28/big/img_19485 +2002/07/31/big/img_236 +2002/07/30/big/img_523 +2002/07/19/big/img_141 +2003/01/17/big/img_957 +2002/08/04/big/img_81 +2002/07/25/big/img_206 +2002/08/15/big/img_716 +2002/08/13/big/img_403 +2002/08/15/big/img_685 +2002/07/26/big/img_884 +2002/07/19/big/img_499 +2002/07/23/big/img_772 +2002/07/27/big/img_752 +2003/01/14/big/img_493 +2002/08/25/big/img_664 +2002/07/31/big/img_334 +2002/08/26/big/img_678 +2002/09/01/big/img_16541 +2003/01/14/big/img_347 +2002/07/23/big/img_187 +2002/07/30/big/img_1163 +2002/08/05/big/img_35 +2002/08/22/big/img_944 +2002/08/07/big/img_1239 +2002/07/29/big/img_1215 +2002/08/03/big/img_312 +2002/08/05/big/img_3523 +2002/07/29/big/img_218 +2002/08/13/big/img_672 +2002/08/16/big/img_205 +2002/08/17/big/img_594 +2002/07/29/big/img_1411 +2002/07/30/big/img_942 +2003/01/16/big/img_312 +2002/08/08/big/img_312 +2002/07/25/big/img_15 +2002/08/09/big/img_839 +2002/08/01/big/img_2069 +2002/08/31/big/img_17512 +2002/08/01/big/img_3 +2002/07/31/big/img_320 +2003/01/15/big/img_1265 +2002/08/14/big/img_563 +2002/07/31/big/img_167 +2002/08/20/big/img_374 +2002/08/13/big/img_406 +2002/08/08/big/img_625 +2002/08/02/big/img_314 +2002/08/27/big/img_19964 +2002/09/01/big/img_16670 +2002/07/31/big/img_599 +2002/08/29/big/img_18906 +2002/07/24/big/img_373 +2002/07/26/big/img_513 +2002/09/02/big/img_15497 +2002/08/19/big/img_117 +2003/01/01/big/img_158 +2002/08/24/big/img_178 +2003/01/13/big/img_935 +2002/08/13/big/img_609 +2002/08/30/big/img_18341 +2002/08/25/big/img_674 +2003/01/13/big/img_209 +2002/08/13/big/img_258 +2002/08/05/big/img_3543 +2002/08/07/big/img_1970 +2002/08/06/big/img_3004 +2003/01/17/big/img_487 +2002/08/24/big/img_873 +2002/08/29/big/img_18730 +2002/08/09/big/img_375 +2003/01/16/big/img_751 +2002/08/02/big/img_603 +2002/08/19/big/img_325 +2002/09/01/big/img_16420 +2002/08/05/big/img_3633 +2002/08/21/big/img_516 +2002/07/19/big/img_501 +2002/07/26/big/img_688 +2002/07/24/big/img_256 +2002/07/25/big/img_438 +2002/07/31/big/img_1017 +2002/08/22/big/img_512 +2002/07/21/big/img_543 +2002/08/08/big/img_223 +2002/08/19/big/img_189 +2002/08/12/big/img_630 +2002/07/30/big/img_958 +2002/07/28/big/img_208 +2002/08/31/big/img_17691 +2002/07/22/big/img_542 +2002/07/19/big/img_741 +2002/07/19/big/img_158 +2002/08/15/big/img_399 +2002/08/01/big/img_2159 +2002/08/14/big/img_455 +2002/08/17/big/img_1011 +2002/08/26/big/img_744 +2002/08/12/big/img_624 +2003/01/17/big/img_821 +2002/08/16/big/img_980 +2002/07/28/big/img_281 +2002/07/25/big/img_171 +2002/08/03/big/img_116 +2002/07/22/big/img_467 +2002/07/31/big/img_750 +2002/07/26/big/img_435 +2002/07/19/big/img_822 +2002/08/13/big/img_626 +2002/08/11/big/img_344 +2002/08/02/big/img_473 +2002/09/01/big/img_16817 +2002/08/01/big/img_1275 +2002/08/28/big/img_19270 +2002/07/23/big/img_607 +2002/08/09/big/img_316 +2002/07/29/big/img_626 +2002/07/24/big/img_824 +2002/07/22/big/img_342 +2002/08/08/big/img_794 +2002/08/07/big/img_1209 +2002/07/19/big/img_18 +2002/08/25/big/img_634 +2002/07/24/big/img_730 +2003/01/17/big/img_356 +2002/07/23/big/img_305 +2002/07/30/big/img_453 +2003/01/13/big/img_972 +2002/08/06/big/img_2610 +2002/08/29/big/img_18920 +2002/07/31/big/img_123 +2002/07/26/big/img_979 +2002/08/24/big/img_635 +2002/08/05/big/img_3704 +2002/08/07/big/img_1358 +2002/07/22/big/img_306 +2002/08/13/big/img_619 +2002/08/02/big/img_366 diff --git a/hair_service_sd/core/models/layers/data/__init__.py b/hair_service_sd/core/models/layers/data/__init__.py new file mode 100644 index 0000000..ea50eba --- /dev/null +++ b/hair_service_sd/core/models/layers/data/__init__.py @@ -0,0 +1,3 @@ +from .wider_face import WiderFaceDetection, detection_collate +from .data_augment import * +from .config import * diff --git a/hair_service_sd/core/models/layers/data/config.py b/hair_service_sd/core/models/layers/data/config.py new file mode 100644 index 0000000..591f349 --- /dev/null +++ b/hair_service_sd/core/models/layers/data/config.py @@ -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 +} + diff --git a/hair_service_sd/core/models/layers/data/data_augment.py b/hair_service_sd/core/models/layers/data/data_augment.py new file mode 100644 index 0000000..2980fe8 --- /dev/null +++ b/hair_service_sd/core/models/layers/data/data_augment.py @@ -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 diff --git a/hair_service_sd/core/models/layers/data/wider_face.py b/hair_service_sd/core/models/layers/data/wider_face.py new file mode 100644 index 0000000..22f56ef --- /dev/null +++ b/hair_service_sd/core/models/layers/data/wider_face.py @@ -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) diff --git a/hair_service_sd/core/models/layers/functions/prior_box.py b/hair_service_sd/core/models/layers/functions/prior_box.py new file mode 100644 index 0000000..80c7f85 --- /dev/null +++ b/hair_service_sd/core/models/layers/functions/prior_box.py @@ -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 diff --git a/hair_service_sd/core/models/layers/modules/__init__.py b/hair_service_sd/core/models/layers/modules/__init__.py new file mode 100644 index 0000000..cf24bdd --- /dev/null +++ b/hair_service_sd/core/models/layers/modules/__init__.py @@ -0,0 +1,3 @@ +from .multibox_loss import MultiBoxLoss + +__all__ = ['MultiBoxLoss'] diff --git a/hair_service_sd/core/models/layers/modules/multibox_loss.py b/hair_service_sd/core/models/layers/modules/multibox_loss.py new file mode 100644 index 0000000..8a311a3 --- /dev/null +++ b/hair_service_sd/core/models/layers/modules/multibox_loss.py @@ -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 diff --git a/hair_service_sd/core/models/net.py b/hair_service_sd/core/models/net.py new file mode 100644 index 0000000..add71a2 --- /dev/null +++ b/hair_service_sd/core/models/net.py @@ -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 + diff --git a/hair_service_sd/core/models/nms/__init__.py b/hair_service_sd/core/models/nms/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hair_service_sd/core/models/nms/py_cpu_nms.py b/hair_service_sd/core/models/nms/py_cpu_nms.py new file mode 100644 index 0000000..260c5ba --- /dev/null +++ b/hair_service_sd/core/models/nms/py_cpu_nms.py @@ -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 diff --git a/hair_service_sd/core/models/resnet.py b/hair_service_sd/core/models/resnet.py new file mode 100644 index 0000000..f82a980 --- /dev/null +++ b/hair_service_sd/core/models/resnet.py @@ -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" `_ + 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" `_ + 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" `_ + 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" `_ + 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" `_ + 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" `_ + 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" `_ + 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" `_ + 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" `_ + 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) \ No newline at end of file diff --git a/hair_service_sd/core/models/retinaface.py b/hair_service_sd/core/models/retinaface.py new file mode 100644 index 0000000..f8e2336 --- /dev/null +++ b/hair_service_sd/core/models/retinaface.py @@ -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 \ No newline at end of file diff --git a/hair_service_sd/core/mtcnn/__init__.py b/hair_service_sd/core/mtcnn/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/hair_service_sd/core/mtcnn/__init__.py @@ -0,0 +1 @@ + diff --git a/hair_service_sd/core/mtcnn/box_utils.py b/hair_service_sd/core/mtcnn/box_utils.py new file mode 100644 index 0000000..d7a076f --- /dev/null +++ b/hair_service_sd/core/mtcnn/box_utils.py @@ -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 diff --git a/hair_service_sd/core/mtcnn/detector.py b/hair_service_sd/core/mtcnn/detector.py new file mode 100644 index 0000000..407c0b8 --- /dev/null +++ b/hair_service_sd/core/mtcnn/detector.py @@ -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 \ No newline at end of file diff --git a/hair_service_sd/core/mtcnn/detector_ly.py b/hair_service_sd/core/mtcnn/detector_ly.py new file mode 100644 index 0000000..407c0b8 --- /dev/null +++ b/hair_service_sd/core/mtcnn/detector_ly.py @@ -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 \ No newline at end of file diff --git a/hair_service_sd/core/mtcnn/model.py b/hair_service_sd/core/mtcnn/model.py new file mode 100644 index 0000000..d9a7e3b --- /dev/null +++ b/hair_service_sd/core/mtcnn/model.py @@ -0,0 +1,109 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from collections import OrderedDict +import numpy as np +import os +# from hairstyle_model import modelRoot +modelRoot = "weights" +class Flatten(nn.Module): + def __init__(self): + super(Flatten, self).__init__() + def forward(self, x): + x = x.transpose(3, 2).contiguous() + return x.view(x.size(0), -1) + +class PNet(nn.Module): + def __init__(self): + super(PNet, self).__init__() + self.model_path = modelRoot + + self.features = nn.Sequential(OrderedDict([ + ('conv1', nn.Conv2d(3, 10, 3, 1)), + ('prelu1', nn.PReLU(10)), + ('pool1', nn.MaxPool2d(2, 2, ceil_mode=True)), + ('conv2', nn.Conv2d(10, 16, 3, 1)), + ('prelu2', nn.PReLU(16)), + ('conv3', nn.Conv2d(16, 32, 3, 1)), + ('prelu3', nn.PReLU(32)) + ])) + self.conv4_1 = nn.Conv2d(32, 2, 1, 1) + self.conv4_2 = nn.Conv2d(32, 4, 1, 1) + weights = np.load(os.path.join(self.model_path, 'pnet.npy'), allow_pickle=True)[()] + for n, p in self.named_parameters(): + p.data = torch.FloatTensor(weights[n]) + + def forward(self, x): + x = self.features(x) + a = self.conv4_1(x) + b = self.conv4_2(x) + a = F.softmax(a, dim=1) + return b, a + +class RNet(nn.Module): + def __init__(self): + super(RNet, self).__init__() + self.model_path = modelRoot + + self.features = nn.Sequential(OrderedDict([ + ('conv1', nn.Conv2d(3, 28, 3, 1)), + ('prelu1', nn.PReLU(28)), + ('pool1', nn.MaxPool2d(3, 2, ceil_mode=True)), + ('conv2', nn.Conv2d(28, 48, 3, 1)), + ('prelu2', nn.PReLU(48)), + ('pool2', nn.MaxPool2d(3, 2, ceil_mode=True)), + ('conv3', nn.Conv2d(48, 64, 2, 1)), + ('prelu3', nn.PReLU(64)), + ('flatten', Flatten()), + ('conv4', nn.Linear(576, 128)), + ('prelu4', nn.PReLU(128)) + ])) + self.conv5_1 = nn.Linear(128, 2) + self.conv5_2 = nn.Linear(128, 4) + weights = np.load(os.path.join(self.model_path, 'rnet.npy'), allow_pickle=True)[()] + for n, p in self.named_parameters(): + p.data = torch.FloatTensor(weights[n]) + + def forward(self, x): + x = self.features(x) + a = self.conv5_1(x) + b = self.conv5_2(x) + a = F.softmax(a, dim=1) + return b, a + +class ONet(nn.Module): + def __init__(self): + super(ONet, self).__init__() + self.model_path = modelRoot + + self.features = nn.Sequential(OrderedDict([ + ('conv1', nn.Conv2d(3, 32, 3, 1)), + ('prelu1', nn.PReLU(32)), + ('pool1', nn.MaxPool2d(3, 2, ceil_mode=True)), + ('conv2', nn.Conv2d(32, 64, 3, 1)), + ('prelu2', nn.PReLU(64)), + ('pool2', nn.MaxPool2d(3, 2, ceil_mode=True)), + ('conv3', nn.Conv2d(64, 64, 3, 1)), + ('prelu3', nn.PReLU(64)), + ('pool3', nn.MaxPool2d(2, 2, ceil_mode=True)), + ('conv4', nn.Conv2d(64, 128, 2, 1)), + ('prelu4', nn.PReLU(128)), + ('flatten', Flatten()), + ('conv5', nn.Linear(1152, 256)), + ('drop5', nn.Dropout(0.25)), + ('prelu5', nn.PReLU(256)), + ])) + self.conv6_1 = nn.Linear(256, 2) + self.conv6_2 = nn.Linear(256, 4) + self.conv6_3 = nn.Linear(256, 10) + weights = np.load(os.path.join(self.model_path, 'onet.npy'), allow_pickle=True)[()] + for n, p in self.named_parameters(): + p.data = torch.FloatTensor(weights[n]) + + def forward(self, x): + x = self.features(x) + a = self.conv6_1(x) + b = self.conv6_2(x) + c = self.conv6_3(x) + a = F.softmax(a, dim=1) + return c, b, a diff --git a/hair_service_sd/core/oss_module.py b/hair_service_sd/core/oss_module.py new file mode 100644 index 0000000..4bd3161 --- /dev/null +++ b/hair_service_sd/core/oss_module.py @@ -0,0 +1,41 @@ +import time + +import oss2 +import os + +class OSS_object(): + def __init__(self): + access_key_id = os.getenv('OSS_TEST_ACCESS_KEY_ID', 'LTAI5tPZA6M67YRoxGPdJw1v') + access_key_secret = os.getenv('OSS_TEST_ACCESS_KEY_SECRET', 'BPydkvYFrXsbOYj3ix8UHQSLS4ivNP') + bucket_name = os.getenv('OSS_TEST_BUCKET', 'oss-aidigitalfield') + endpoint = os.getenv('OSS_TEST_ENDPOINT', 'oss-cn-beijing.aliyuncs.com') + # access_key_id = os.getenv('OSS_TEST_ACCESS_KEY_ID', 'LTAI5tMq9DivPYYkcpc6qhNP') + # access_key_secret = os.getenv('OSS_TEST_ACCESS_KEY_SECRET', 'XIkDAq7r4U9BVf7fkKECP46FXDLF9l') + # bucket_name = os.getenv('OSS_TEST_BUCKET', 'digit-person') + # endpoint = os.getenv('OSS_TEST_ENDPOINT', 'oss-cn-beijing.aliyuncs.com') + # access_key_id = os.getenv('OSS_TEST_ACCESS_KEY_ID', 'LTAI5tByNnrV4vRVioW66uq2') + # access_key_secret = os.getenv('OSS_TEST_ACCESS_KEY_SECRET', 'bLGYCKUNtQTDfzGfXqw06MWYomr5lw') + # bucket_name = os.getenv('OSS_TEST_BUCKET', 'mzyidong-tmp') + # endpoint = os.getenv('OSS_TEST_ENDPOINT', 'oss-cn-zhangjiakou-internal.aliyuncs.com') + for param in (access_key_id, access_key_secret, bucket_name, endpoint): + assert '<' not in param, '请设置参数：' + param + + self.bucket = oss2.Bucket(oss2.Auth(access_key_id, access_key_secret), endpoint, bucket_name) + + def upload_file(self, file, target_name): + t0 = time.time() + with open(oss2.to_unicode(file), 'rb') as f: + ret = self.bucket.put_object(target_name, f) + print(ret.headers['x-oss-request-id']) + + url = "https://oss-aidigitalfield.oss-cn-beijing.aliyuncs.com/{}".format(target_name) + print('耗时：{},签名url的地址为：{}'.format(time.time() - t0, url)) + return url + + + +if __name__ == '__main__': + oss_2 = OSS_object() + t0 = time.time() + oss_2.upload_file('../data/left.jpg', 'hair_mz/images/test.jpg') + print(time.time() - t0) \ No newline at end of file diff --git a/hair_service_sd/core/process_modules.py b/hair_service_sd/core/process_modules.py new file mode 100644 index 0000000..b81497d --- /dev/null +++ b/hair_service_sd/core/process_modules.py @@ -0,0 +1,2850 @@ +import os +import pickle +import sys +import time + +import torch +from torch.nn import functional as F +from torch import nn +import numpy as np +import cv2 +import math +from core.utils import landmark_processor, model_io +import torchvision +from core.models.MomocvFaceAlignment1K import MomocvFaceAlignment1K +# from core.mtcnn.detector import MTCNNFaceDetector +from core.models.detector import RetinaFaceDetector +from core.utils import util +from core.matting.networks import generators +from core.seg.hairseg_single_model import Evaluator +from core.models.Generator_BaldSeg import Generator_BaldSeg_5c +from core.model_3ddfa.model_3ddfa import Model_3DDFA +from core.bodyseg.msc_distilling import DeepLab +from core.models.resnet import resnet18 +from common.logger import config +modelRoot = "weights" + +class Get_Landmark(object): + def __init__(self, gpu_id=None): + self.img_size = 640 + self.face_alignmenter_1k = MomocvFaceAlignment1K(gpu_id=gpu_id) + print('face_alignmenter_1k load done') + # self.face_detector = MTCNNFaceDetector(gpu_id=gpu_id) + self.face_detector = RetinaFaceDetector(gpu_id=gpu_id) + self.model_3d = Model_3DDFA(gpu_id=gpu_id) + + def get_max_rect(self, bounding_boxes): + max_area = 0 + index = 0 + for i, box in enumerate(bounding_boxes): + width = box[2] - box[0] + height = box[3] - box[1] + if width * height > max_area: + index = i + max_area = width * height + return index + + def forward(self, img): + # bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=100,thresholds=[0.6, 0.8, 0.8]) # mtcnn face detect input + bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=50) + + if len(bounding_boxes) == 0: + return None, None, None + + if len(bounding_boxes) >= 1: + box_index = self.get_max_rect(bounding_boxes) + pts5 = landmarks[box_index] + else: + return None, None, None + # for i in range(5): + # cv2.circle(img, (int(pts5[i*2]), int(pts5[i*2+1])), 1, (255, 0,0), 1) + # cv2.imshow('img', img) + # cv2.waitKey() + landmarks1k = self.face_alignmenter_1k.detect_according_5pts(img, pts5) + # pts137 = landmark_processor.pts_1k_to_137(landmarks1k) + # movie_params = self.model_3d.detect([img], [pts137])[0] + # pitch, yaw, roll = movie_params[1:4] + ret_info = [landmarks1k, bounding_boxes[box_index], None] + return ret_info + + + def forward_color(self, img): + bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=50) + + if len(bounding_boxes) == 0: + return None + + if len(bounding_boxes) >= 1: + box_index = self.get_max_rect(bounding_boxes) + pts5 = landmarks[box_index] + else: + return None + + landmarks1k = self.face_alignmenter_1k.detect_according_5pts(img, pts5) + return landmarks1k + + + def forward_infer(self, img): + bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=50) + + if len(bounding_boxes) == 0: + return None, None, None + + if len(bounding_boxes) >= 1: + box_index = self.get_max_rect(bounding_boxes) + pts5 = landmarks[box_index] + else: + return None, None, None + + landmarks1k = self.face_alignmenter_1k.detect_according_5pts(img, pts5) + return landmarks1k + + def forward_diy(self, img): + # bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=100,thresholds=[0.6, 0.8, 0.8]) # mtcnn face detect input + bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=50) + + if len(bounding_boxes) == 0: + return None, None, None + + if len(bounding_boxes) >=1 : + box_index = self.get_max_rect(bounding_boxes) + pts5 = landmarks[box_index] + else: + return None, None, None + # for i in range(5): + # cv2.circle(img, (int(pts5[i*2]), int(pts5[i*2+1])), 1, (255, 0,0), 1) + # cv2.imshow('img', img) + # cv2.waitKey() + landmarks1k = self.face_alignmenter_1k.detect_according_5pts(img, pts5) + # pts137 = landmark_processor.pts_1k_to_137(landmarks1k) + # movie_params = self.model_3d.detect([img], [pts137])[0] + # pitch, yaw, roll = movie_params[1:4] + ret_info = [landmarks1k, bounding_boxes[box_index], None] + return ret_info + + def forward_v2(self, img): + # bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=50, + # thresholds=[0.6, 0.7, 0.9]) + bounding_boxes, landmarks = self.face_detector.forward_v2(img, min_face_size=30, + thresholds=[0.6, 0.6, 0.6]) + if len(bounding_boxes) == 0: + return None + + if len(bounding_boxes) > 0: + box_index = self.get_max_rect(bounding_boxes) + pts5 = landmarks[box_index] + else: + return None + + landmarks1k = self.face_alignmenter_1k.detect_according_5pts(img, pts5) + + return landmarks1k + + def get_face_shape(self, xiaba_type, landmark137): + + face_shape = ['chang', 'fang', 'yuan', 'tuoyuan', 'xin'] + forhead_idxs = [12, 13, 14] + mid_idxs = [15, 16, 17] + lowwer_idxs = [18, 19, 20] + face_width = max(landmark137[:,0]) - min(landmark137[:,0]) + face_widthest_idx = list(landmark137[:,0]).index(min(landmark137[:,0])) + face_scale_lw = 1.4 + face_scale_qx_top = 1.4 + face_scale_qx_bot = 1.2 + + face_height = landmark137[0][1] - landmark137[11][1] + face_width_qiane = landmark137[9][0] - landmark137[13][0] + face_width_xiahe = landmark137[3][0] - landmark137[19][0] + check_scale = face_width_qiane / face_width_xiahe + print('face_height', face_height) + print('face_width', face_width) + print('face_scale_lw', face_scale_lw) + print('face_width_qiane', face_width_qiane) + print('face_width_xiahe', face_width_xiahe) + print('face_scale_qx', check_scale) + if face_widthest_idx in forhead_idxs: + if face_height / face_width >= face_scale_lw: + return face_shape[0] + else: + if check_scale > face_scale_qx_bot and face_scale_qx_bot < face_scale_qx_top: + if xiaba_type == 'jian': + return face_shape[0] + elif xiaba_type == 'fang': + return face_shape[1] + else: + return face_shape[1] + elif check_scale > face_scale_qx_top: + if xiaba_type == 'jian': + return face_shape[4] + elif xiaba_type == 'fang': + return face_shape[4] + else: + return face_shape[1] + else: + return face_shape[1] + + elif face_widthest_idx in mid_idxs: + if face_height / face_width >= face_scale_lw: + return face_shape[0] + else: + if check_scale > face_scale_qx_bot and face_scale_qx_bot < face_scale_qx_top: + if xiaba_type == 'jian': + return face_shape[3] + elif xiaba_type == 'fang': + return face_shape[1] + else: + return face_shape[2] + elif check_scale > face_scale_qx_top: + if xiaba_type == 'jian': + return face_shape[3] + elif xiaba_type == 'fang': + return face_shape[3] + else: + return face_shape[2] + else: + return face_shape[1] + else: + if face_height / face_width >= face_scale_lw: + return face_shape[0] + else: + if check_scale > face_scale_qx_bot and face_scale_qx_bot < face_scale_qx_top: + if xiaba_type == 'jian': + return face_shape[0] + elif xiaba_type == 'fang': + return face_shape[1] + else: + return face_shape[1] + else: + return face_shape[1] + + +class GenTrimap(object): + def __init__(self): + self.erosion_kernels = [None] + [cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size, size)) for size in range(1,30)] + + def __call__(self, alpha): + + fg_mask = np.zeros_like(alpha) + bg_mask = np.zeros_like(alpha) + fg_mask[alpha == 255] = 1 + bg_mask[alpha == 0] = 1 + + fg_mask = fg_mask.astype(np.int).astype(np.uint8) + bg_mask = bg_mask.astype(np.int).astype(np.uint8) + + fg_mask = cv2.erode(fg_mask, self.erosion_kernels[15]) + bg_mask = cv2.erode(bg_mask, self.erosion_kernels[29]) + + trimap = np.ones_like(alpha, dtype=np.uint8) * 128 + trimap[fg_mask == 1] = 255 + trimap[bg_mask == 1] = 0 + return trimap +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) + alpha_pred, info_dict = model(image, trimap) + + fg_pred = alpha_pred[:, :-1, :, :] + alpha_pred = alpha_pred[:, -1, :, :].unsqueeze(1) + + 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] + + + if return_offset: + short_side = h if h < w else w + ratio = 512 / short_side + offset_1 = util.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 = util.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 + +class Generator_Matte(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.output_img_size = 512 + + triseg_model_path = os.path.join(modelRoot, 'deeplabv3_hair512_360_0520_wl.pth') + self.triseg_model = Evaluator(gpu_id=device_id, output_img_size=self.output_img_size, nclass=3, seg_model_path=triseg_model_path) + + hair_matte_model_path = os.path.join(modelRoot, 'gca-dist-fg-0430-latest_model.pth') # gca-dist-fg-0203-latest_model gca-dist-fg-0430-latest_model + self.matte_model = self.load_hair_matte_model(hair_matte_model_path) + + self.gen_trimap = GenTrimap() + + def load_hair_matte_model(self, hair_matte_model_path): + # build model + model = generators.get_generator(encoder="resnet_gca_encoder_29", decoder="res_gca_decoder_22", num_class=4) + + # load checkpoint + checkpoint = torch.load(hair_matte_model_path, map_location=lambda storage, loc: storage) + model.load_state_dict(util.remove_prefix_state_dict(checkpoint['state_dict']), strict=True) + model.to(self.device) + # print("matte_model: ", model) + + # inference + model = model.eval() + return model + + def matte_inference(self, image, landmark1k): + with torch.no_grad(): + trimap = self.triseg_model.eval(image, landmark1k) + + ori_h, ori_w, _ = image.shape + + limit_size = 1600 + if ori_h > limit_size or ori_w > limit_size: + if ori_h > ori_w: + new_tri_h = limit_size + new_tri_w = int(ori_w * limit_size / ori_h) + else: + new_tri_w = limit_size + new_tri_h = int(ori_h * limit_size / ori_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 + + trimap = self.gen_trimap(trimap_resize[:, :, 0]) + + # cv2.imwrite(os.path.join(args.output, image_name.replace(ext, "_trimap.png")), trimap) + + image_dict = generator_tensor_dict(image_resize, trimap) + + pred_fg, pred, offset = single_inference(self.matte_model, image_dict, device=self.device) + + pred_fg[trimap == 1] = image_resize[trimap == 1] + if pred.shape[1] != image.shape[1] or pred.shape[0] != image.shape[0]: + pred = cv2.resize(pred, (image.shape[1], image.shape[0])) + return pred_fg, pred, image_resize +class Generator_Bald(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.model_dir = modelRoot + generator_bald_model_path = os.path.join(self.model_dir, 'zxm_v7_gen_bald_add_changed_hair_pairdata_5seg_0611.pt') + print("load generator_bald_model_path", generator_bald_model_path) + self.load_generator_bald768_model(generator_bald_model_path) + self.output_size = 768 + + def load_generator_bald768_model(self, generator_bald_model_path): + self.generator_bald_model = torch.jit.load(generator_bald_model_path, map_location='cpu').to(self.device) + + def Geneator_Bald_inference(self, user_rgb_8uc3_bald_768, user_matting_8uc3_bald_768, user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768): + + """ + input： + 图像尺寸基于： 人脸 512， 图像均为3通道 + user_rgb_8uc3_bald_512: 输入图， size 512， uint8 （0-255） + user_matting_8uc3_bald_512： 输入图 matting alpha 值， uint8 （0-255） + user_baldseg_8uc3_bald_512： 输入图 光头分割， uint8 （0-255） + + user_landmark_f1k2_bald_512：输入图 关键点 1k*2 float32 + + output： + bald_gene_8uc3_bald_512: 生成 光头图， uint8 （0-255） + + """ + + # user_baldseg_8uc3_bald_768_cp = user_baldseg_8uc3_bald_768.copy() + # cood_y = int((user_landmark_f1k2_bald_768[214, 1] + user_landmark_f1k2_bald_768[99, 1]) // 2) + # + # random_int = 30 # random.randint(0, 30) + # kernel = np.ones((random_int, 1), np.uint8) + # user_mask_erode = cv2.erode(user_baldseg_8uc3_bald_768, kernel, iterations=1) + # user_baldseg_8uc3_bald_768[:cood_y, :, :] = user_mask_erode[:cood_y, :, :] + + # bald_con = np.concatenate((user_baldseg_8uc3_bald_768_cp, user_baldseg_8uc3_bald_768), axis=1) + + + # cv2.imwrite("/media/DATA_4T/project/hairstyle/分割测试/debug_line_res/user_rgb_8uc3_bald_768.png", user_rgb_8uc3_bald_768) + # cv2.imwrite("/media/DATA_4T/project/hairstyle/分割测试/debug_line_res/user_matting_8uc3_bald_768.png", user_matting_8uc3_bald_768) + # cv2.imwrite("/media/DATA_4T/project/hairstyle/分割测试/debug_line_res/user_baldseg_8uc3_bald_768.png", user_baldseg_8uc3_bald_768) + # np.savetxt("/media/DATA_4T/project/hairstyle/分割测试/debug_line_res/user_landmark_f1k2_bald_768.txt", user_landmark_f1k2_bald_768) + + real_res = user_rgb_8uc3_bald_768.copy() + + user_landmark_f1k2_bald_768 = user_landmark_f1k2_bald_768.astype(np.int32) + # kernel = np.ones((8, 8), np.uint8) + kernel_1 = np.ones((20, 20), np.uint8) + kernel_2 = np.ones((40, 40), np.uint8) + user_matting_8uc3_bald_768_dilate_1 = cv2.dilate(user_matting_8uc3_bald_768[:, :, 0], kernel_1, iterations=1) + user_matting_8uc3_bald_768_dilate_2 = cv2.dilate(user_matting_8uc3_bald_768[:, :, 0], kernel_2, iterations=1) + + user_matting_8uc3_bald_768_dilate_2[user_matting_8uc3_bald_768_dilate_2 > 0] = 255 + + b = user_baldseg_8uc3_bald_768[:, :, 0] > 125 + g = user_baldseg_8uc3_bald_768[:, :, 1] > 0 + r = user_baldseg_8uc3_bald_768[:, :, 2] > 125 + + user_matting_8uc3_bald_768_blur = user_matting_8uc3_bald_768_dilate_2.copy() + + # TODO change noface mask + # user_matting_8uc3_bald_768_blur[g] = 255 + user_matting_8uc3_bald_768_blur = cv2.blur(user_matting_8uc3_bald_768_blur, (30, 30)) + # user_matting_8uc3_bald_768_blur_2 = cv2.blur(user_matting_8uc3_bald_768_blur, (20, 20)) + # user_matting_8uc3_bald_768_blur[(~b) & (~g) & (~r)] = user_matting_8uc3_bald_768_blur_2[(~b) & (~g) & (~r)] + + + # user_matting_8uc3_bald_768_dilate_1[(~b) * (~g)] = user_matting_8uc3_bald_768[:, :, 0][(~b) * (~g)] + user_matting_8uc3_bald_768_dilate_1[(~b) * (~g) * (~r)] = user_matting_8uc3_bald_768[:, :, 0][(~b) * (~g) * (~r)] + # user_matting_8uc3_bald_768_dilate_1[(~b) * g * r] = user_matting_8uc3_bald_768[:, :, 0][(~b) * g * r] + user_rgb_8uc3_bald_768[user_matting_8uc3_bald_768_dilate_1.astype(np.bool)] = 255 + + ######################################## + # cv2.fillPoly(user_baldseg_8uc3_bald_512, user_pts1k[:311][np.newaxis, :, :], (200, 175, 0)) # face + cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[928:999][np.newaxis, :, :], + (0, 125, 0)) # left eyebrow + cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[856:927][np.newaxis, :, :], + (125, 0, 0)) # right eyebrow + cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[691:754][np.newaxis, :, :], + (0, 0, 125)) # left eye + cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[792:855][np.newaxis, :, :], + (125, 125, 0)) # right eye + cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[548:616][np.newaxis, :, :], + (0, 125, 125)) # rose + cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[312:547][np.newaxis, :, :], + (125, 125, 125)) # mouth + cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[655:690][np.newaxis, :, :], + (0, 0, 175)) # left black_eyeball + cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[756:791][np.newaxis, :, :], + (125, 0, 125)) # right black_eyeball + + eyes_mask = np.zeros(user_baldseg_8uc3_bald_768[:, :, 0].shape) + cv2.fillPoly(eyes_mask, user_landmark_f1k2_bald_768[691:754][np.newaxis, :, :], 1) # left eye + cv2.fillPoly(eyes_mask, user_landmark_f1k2_bald_768[792:855][np.newaxis, :, :], 1) # right eye + kernel = np.ones((30, 30), np.uint8) + eyes_mask = cv2.dilate(eyes_mask, kernel, iterations=1) + cv2.fillPoly(eyes_mask, user_landmark_f1k2_bald_768[312:467][np.newaxis, :, :], 1) # mouth + user_rgb_8uc3_bald_768[eyes_mask.astype(np.bool)] = real_res[eyes_mask.astype(np.bool)] + + ######################### add by zxm + # face_erode_mask = np.zeros(user_baldseg_8uc3_bald_768[:, :, 0].shape) + # cv2.fillPoly(face_erode_mask, user_landmark_f1k2_bald_768[:311][np.newaxis, :, :], 1) # face + # kernel = np.ones((60, 60), np.uint8) + # face_erode_mask = cv2.erode(face_erode_mask, kernel, iterations=1) + # user_rgb_8uc3_bald_768[face_erode_mask.astype(np.bool)] = real_res[face_erode_mask.astype(np.bool)] + # user_rgb_8uc3_bald_768[user_matting_8uc3_bald_768_dilate_copy.astype(np.bool)] = 255 + ######################### add by zxm + + # repeat + # user_rgb_8uc3_bald_512[eyes_mask.astype(np.bool)] = real_res[eyes_mask.astype(np.bool)] + + + condition = user_baldseg_8uc3_bald_768 # np.concatenate((user_baldseg_8uc3_bald_512, user_hair_mask), axis=2) + condition = (condition.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] + input_paf = (user_rgb_8uc3_bald_768.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] + + + # ******************* forward *******************$ + # test_fake, _, _ = model.preview(input_paf, condition) + with torch.no_grad(): + condition = torch.from_numpy(condition).to(self.device) + input_paf = torch.from_numpy(input_paf).to(self.device) + test_fake = self.generator_bald_model(input_paf, condition) + + test_res = (test_fake.detach()[0].to('cpu').numpy().transpose((1, 2, 0)) * 255).astype(np.uint8).copy() + + # kernel = np.ones((25, 25), np.uint8) + # user_hair_mask_dilate = cv2.dilate(user_matting_8uc3_bald_768[:, :, 0], kernel) + # user_hair_mask_dilate[user_hair_mask_dilate > 0] = 255 + # user_hair_mask_blur = cv2.blur(user_hair_mask_dilate, (20, 20)) + test_res_clear = real_res * (1 - user_matting_8uc3_bald_768_blur[:, :, np.newaxis] / 255) + test_res * ( + user_matting_8uc3_bald_768_blur[:, :, np.newaxis] / 255) + bald_gene_8uc3_bald_768 = (np.clip(test_res_clear, 0, 255)).astype(np.uint8) + + + # cv2.imwrite("/media/DATA_4T/project/hairstyle/分割测试/debug_line_res/test_res_768.png", test_res) + # cv2.imwrite("/media/DATA_4T/project/hairstyle/分割测试/debug_line_res/bald_gene_8uc3_bald_768.png", bald_gene_8uc3_bald_768) + # cv2.imwrite("/media/DATA_4T/project/hairstyle/分割测试/debug_line_res/user_matting_8uc3_bald_768_blur.png", user_matting_8uc3_bald_768_blur) + + return bald_gene_8uc3_bald_768, user_matting_8uc3_bald_768_blur +class Process_Data(object): + def __init__(self, gpu, device_id, save_name=None): + 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.generator_matte = Generator_Matte(gpu, device_id) + + self.generator_baldseg = Generator_BaldSeg_5c(gpu, device_id) + self.generator_bald = Generator_Bald(gpu, device_id) + + self.hair_size = 768 + self.bald_output_size = 768 + self.color_output_size = 768 + + def get_hair_M(self, landmark1k): + + hairstyle_M = landmark_processor.get_transform_mat_hair_ratio(landmark1k, 512, 0.5) + + return hairstyle_M + + def get_hair_M_girl_v1(self, landmark1k): + + hairstyle_M = landmark_processor.get_transform_mat_hair_ratio_v1(landmark1k, 768, ratio=0.5, + h_offset=0.45) # default 0.5 0.45 ratio=0.35, h_offset=0.32 + + return hairstyle_M + + def get_hair_M_girl_v2(self, landmark1k): + + hairstyle_M = landmark_processor.get_transform_mat_hair_ratio_v1(landmark1k, 768, ratio=0.35, + h_offset=0.32) # default 0.5 0.45 ratio=0.35, h_offset=0.32 + return hairstyle_M + + def get_hair_M_boy_v1(self, landmark1k): + hairstyle_M = landmark_processor.get_transform_mat_hair_ratio_v1(landmark1k, 768, ratio=0.6, h_offset=0.65) + return hairstyle_M + + def get_color_hair_M(self, landmark1k): + + # color_hair_M = landmark_processor.get_transform_mat_full_face_ratio_stylegan(landmark1k, self.color_output_size, 0.35) # 0.4 + color_hair_M = landmark_processor.get_transform_mat_hair_ratio_v1(landmark1k, self.color_output_size, + ratio=0.35, h_offset=0.45) + + return color_hair_M + + def get_prepare_data_bald(self, user_rgb_8uc3_orisize, user_baldseg_8uc3_orisize, + user_landmark_1k2_f_orisize): + user_pts1k = user_landmark_1k2_f_orisize.astype(np.int32) + user_bald_M = landmark_processor.get_transform_mat_full_face_ratio_stylegan(user_pts1k, self.bald_output_size, + 0.4) + + user_rgb_8uc3_bald_512 = cv2.warpAffine(user_rgb_8uc3_orisize, user_bald_M, + (self.bald_output_size, self.bald_output_size)) + + user_baldseg_8uc3_bald_512 = cv2.warpAffine(user_baldseg_8uc3_orisize, user_bald_M, + (self.bald_output_size, self.bald_output_size), + flags=cv2.INTER_NEAREST) + + user_landmark_f1k2_bald_512 = landmark_processor.transform_points(user_pts1k, user_bald_M) + + return user_rgb_8uc3_bald_512, user_baldseg_8uc3_bald_512, user_landmark_f1k2_bald_512, user_bald_M + + def get_prepare_data_bald_768(self, user_baldseg_8uc3_orisize, + user_landmark_1k2_f_orisize, user_color_M): + user_pts1k = user_landmark_1k2_f_orisize.astype(np.int32) + user_bald_M = landmark_processor.get_transform_mat_full_face_ratio_stylegan(user_pts1k, self.hair_size, 0.4) + + # user_rgb_8uc3_bald_768 = cv2.warpAffine(user_rgb_8uc3_orisize, user_color_M, (self.bald_output_size, self.bald_output_size)) + + user_baldseg_8uc3_bald_768 = cv2.warpAffine(user_baldseg_8uc3_orisize, user_color_M, + (self.bald_output_size, self.bald_output_size), + flags=cv2.INTER_NEAREST) + + user_baldseg_8uc3_bald_512 = cv2.warpAffine(user_baldseg_8uc3_orisize, user_bald_M, + (self.hair_size, self.hair_size), + flags=cv2.INTER_NEAREST) + + # user_baldseg_8uc3_bald_768_cp = user_baldseg_8uc3_bald_768.copy() + # + # user_pts1k_768 = landmark_processor.transform_points(user_pts1k, user_color_M) + # cood_y = int((user_pts1k_768[214, 1] + user_pts1k_768[99, 1]) // 2) + + # random_int = 30 # random.randint(0, 30) + # kernel = np.ones((random_int, 1), np.uint8) + # user_mask_erode = cv2.erode(user_baldseg_8uc3_bald_768, kernel, iterations=1) + # user_baldseg_8uc3_bald_768[:cood_y, :, :] = user_mask_erode[:cood_y, :, :] + # + # inter_res = np.concatenate((user_baldseg_8uc3_bald_768_cp, user_baldseg_8uc3_bald_768), axis=1) + + + # user_baldseg_8uc3_orisize = cv2.warpAffine(user_baldseg_8uc3_bald_768, cv2.invertAffineTransform(user_color_M), + # dsize=(user_baldseg_8uc3_orisize.shape[1], user_baldseg_8uc3_orisize.shape[0]), flags=cv2.INTER_NEAREST) + + return user_baldseg_8uc3_bald_768, user_baldseg_8uc3_bald_512, user_bald_M # , user_baldseg_8uc3_orisize + + def get_user_blad(self, user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize): + """ + input： + + user_rgb_8uc3_orisize: 用户图 原始尺寸 + user_landmark_1k2_f_orisize: 用户图 1k 点 + + output: + + 图像尺寸基于： 人脸 512， 图像均为3通道 + + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + + """ + + # 得到512 小图 M 矩阵 + user_color_M = self.get_color_hair_M(user_landmark_1k2_f_orisize) + # 得到512 小图 M 矩阵 + # user_hairstyle_M = self.get_hair_M(user_landmark_1k2_f_orisize) + + # 1k 点 转换到 512 尺寸 + # user_landmark_f1k2_512 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_hairstyle_M) + user_landmark_f1k2_bald_768 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_color_M) + + # # 用户图 头发毛躁 warpaffine + usr_color_ratio = np.sqrt((user_color_M[0][0] * user_color_M[0][0]) + (user_color_M[1][0] * user_color_M[1][0])) + user_color_M_tmp = user_color_M / usr_color_ratio + user_rgb_8uc3_bald_768 = cv2.warpAffine(user_rgb_8uc3_orisize, user_color_M_tmp, ( + int(self.bald_output_size / usr_color_ratio), int(self.bald_output_size / usr_color_ratio)), + borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) + user_rgb_8uc3_bald_768 = cv2.resize(user_rgb_8uc3_bald_768, (self.bald_output_size, self.bald_output_size), + fx=usr_color_ratio, fy=usr_color_ratio, interpolation=cv2.INTER_AREA) + + user_rgb_8uc3_bald_768_bl_bg = cv2.warpAffine(user_rgb_8uc3_orisize, user_color_M_tmp, ( + int(self.bald_output_size / usr_color_ratio), int(self.bald_output_size / usr_color_ratio))) + user_rgb_8uc3_bald_768_bl_bg = cv2.resize(user_rgb_8uc3_bald_768_bl_bg, + (self.bald_output_size, self.bald_output_size), fx=usr_color_ratio, + fy=usr_color_ratio, interpolation=cv2.INTER_AREA) + + # 用户图得到 头发 matting 小图 + _, user_matting_8uc1_bald_768, ref_rgb_resize = self.generator_matte.matte_inference(user_rgb_8uc3_bald_768, + user_landmark_f1k2_bald_768) + + + if user_matting_8uc1_bald_768.shape[0] != 768 or user_matting_8uc1_bald_768.shape[1] != 768: + user_matting_8uc1_bald_768 = cv2.resize(user_matting_8uc1_bald_768, + (user_rgb_8uc3_bald_768.shape[1], user_rgb_8uc3_bald_768.shape[0]), + interpolation=cv2.INTER_CUBIC) + user_matting_8uc3_bald_768 = np.repeat(user_matting_8uc1_bald_768[:, :, np.newaxis], 3, axis=2) + + user_matting_8uc3_bald_orisize = cv2.warpAffine(user_matting_8uc1_bald_768, + cv2.invertAffineTransform(user_color_M), + ( + user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), + flags=cv2.INTER_CUBIC) + + # 光头分割 + user_baldseg_8uc3_orisize = self.generator_baldseg.forward(user_rgb_8uc3_orisize, + user_matting_8uc3_bald_orisize, + user_landmark_1k2_f_orisize) + + # 用户图 生成光头 + user_baldseg_8uc3_bald_768, user_baldseg_8uc3_bald_512, user_bald_M \ + = self.get_prepare_data_bald_768(user_baldseg_8uc3_orisize, + user_landmark_1k2_f_orisize, user_color_M) + + orig_h, orig_w, _ = user_rgb_8uc3_orisize.shape + + # TODO: 测试直接采用M 矩阵warp 到 user_hairstyle_M 后的512 尺寸 + # bald_gene_8uc3_bald_768, user_hair_mask_blur_768 = self.generator_bald.Geneator_Bald_inference(user_rgb_8uc3_bald_768.copy(), + # user_matting_8uc3_bald_768, + # user_baldseg_8uc3_bald_768, + # user_landmark_f1k2_bald_768) + + bald_gene_8uc3_bald_768, user_hair_mask_blur_768 = self.generator_bald.Geneator_Bald_inference( + user_rgb_8uc3_bald_768_bl_bg.copy(), + user_matting_8uc3_bald_768, + user_baldseg_8uc3_bald_768, + user_landmark_f1k2_bald_768) + + inv_M_user = cv2.invertAffineTransform(user_color_M) + bald_gene_8uc3_orisize = user_rgb_8uc3_orisize.copy() + + user_bald_gene_8uc3_orisize = cv2.warpAffine(bald_gene_8uc3_bald_768, inv_M_user, (orig_w, orig_h), + dst=bald_gene_8uc3_orisize.copy(), + borderMode=cv2.BORDER_TRANSPARENT) + + user_hair_mask_blur_orisize = cv2.warpAffine(user_hair_mask_blur_768, inv_M_user, (orig_w, orig_h), + flags=cv2.INTER_CUBIC) + + user_bald_res_8uc3_orisize = ( + (1 - user_hair_mask_blur_orisize[:, :, np.newaxis] / 255) * user_rgb_8uc3_orisize + \ + (user_hair_mask_blur_orisize[:, :, np.newaxis] / 255) * user_bald_gene_8uc3_orisize).astype( + np.uint8) + + return user_bald_res_8uc3_orisize, user_bald_gene_8uc3_orisize, user_matting_8uc3_bald_768, user_baldseg_8uc3_bald_768 + + def hair512_to_768(self, user_color_M, user_hairstyle_M): + + M_ori = np.zeros((3, 3), dtype=np.float32) + M_ori[:2, :] = cv2.invertAffineTransform(user_hairstyle_M) + M_ori[2:, :] = [0, 0, 1] + + matAffine_ori = np.zeros((3, 3), dtype=np.float32) + matAffine_ori[:2, :] = user_color_M + matAffine_ori[2:, :] = [0, 0, 1] + + new_mat = matAffine_ori.dot(M_ori) + return new_mat[:2, :] + + def get_prepare_data(self, user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize, ref_rgb_8uc3_orisize, + ref_landmark_1k2_f_orisize): + """ + input： + + user_rgb_8uc3_orisize: 用户图 原始尺寸 + user_landmark_1k2_f_orisize: 用户图 1k 点 + ref_rgb_8uc3_orisize: 参考图 原始尺寸 + ref_landmark_1k2_f_orisize: 参考图 1k 点 + + output: + + 图像尺寸基于： 人脸 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 + + user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 + user_hairstyle_M: 用户图 原图到512小图 M 矩阵 + """ + + # 得到512 小图 M 矩阵 + user_hairstyle_M = self.get_hair_M(user_landmark_1k2_f_orisize) + ref_hairstyle_M = self.get_hair_M(ref_landmark_1k2_f_orisize) + + # 1k 点 转换到 512 尺寸 + user_landmark_f1k2_512 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_hairstyle_M) + ref_landmark_f1k2_512 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, ref_hairstyle_M) + + # 用户图 头发毛躁 warpaffine + usr_ratio = np.sqrt( + (user_hairstyle_M[0][0] * user_hairstyle_M[0][0]) + (user_hairstyle_M[1][0] * user_hairstyle_M[1][0])) + user_hairstyle_M_tmp = user_hairstyle_M / usr_ratio + user_rgb_8uc3_512 = cv2.warpAffine(user_rgb_8uc3_orisize, user_hairstyle_M_tmp, + (int(self.hair_size / usr_ratio), int(self.hair_size / usr_ratio)), + borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) + user_rgb_8uc3_512 = cv2.resize(user_rgb_8uc3_512, (self.hair_size, self.hair_size), fx=usr_ratio, fy=usr_ratio, + interpolation=cv2.INTER_AREA) + + # 参考图 头发毛躁 warpaffine + ratio = np.sqrt( + (ref_hairstyle_M[0][0] * ref_hairstyle_M[0][0]) + (ref_hairstyle_M[1][0] * ref_hairstyle_M[1][0])) + ref_hairstyle_M_tmp = ref_hairstyle_M / ratio + ref_rgb_8uc3_512 = cv2.warpAffine(ref_rgb_8uc3_orisize, ref_hairstyle_M_tmp, + (int(self.hair_size / ratio), int(self.hair_size / ratio)), + borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) + ref_rgb_8uc3_512 = cv2.resize(ref_rgb_8uc3_512, (self.hair_size, self.hair_size), fx=ratio, fy=ratio, + interpolation=cv2.INTER_AREA) + + # ref_rgb_8uc3_512 = cv2.warpAffine(ref_rgb_8uc3_orisize, ref_hairstyle_M, (self.hair_size, self.hair_size), borderValue=[255, 255, 255]) + + # 参考图直接得到 头发 matting 小图 + _, ref_matte_pred_8uc1_512, ref_rgb_resize = self.generator_matte.matte_inference(ref_rgb_8uc3_512, ref_landmark_f1k2_512) + + torch.cuda.empty_cache() + + ref_matting_8uc1_512 = cv2.resize(ref_matte_pred_8uc1_512, + (ref_rgb_8uc3_512.shape[1], ref_rgb_8uc3_512.shape[0]), + interpolation=cv2.INTER_CUBIC) + + ref_matting_8uc3_512 = np.repeat(ref_matting_8uc1_512[:, :, np.newaxis], 3, axis=2) + + # 光头分割 + user_baldseg_8uc3_orisize = self.generator_baldseg.forward(user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize) + ref_baldseg_8uc3_orisize = self.generator_baldseg.forward(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize) + + # 用户图 生成光头 + + user_rgb_8uc3_bald_512, user_baldseg_8uc3_bald_512, user_landmark_f1k2_bald_512, \ + user_bald_M = self.get_prepare_data_bald(user_rgb_8uc3_orisize, user_baldseg_8uc3_orisize, + user_landmark_1k2_f_orisize) + + # 用户图得到 头发 matting 小图 for blad + _, user_matting_8uc1_bald_512, ref_rgb_resize = self.generator_matte.matte_inference(user_rgb_8uc3_bald_512, + user_landmark_f1k2_bald_512) + if user_matting_8uc1_bald_512.shape[0] != 512 or user_matting_8uc1_bald_512.shape[1] != 512: + user_matting_8uc1_bald_512 = cv2.resize(user_matting_8uc1_bald_512, + (user_rgb_8uc3_bald_512.shape[1], user_rgb_8uc3_bald_512.shape[0]), + interpolation=cv2.INTER_CUBIC) + user_matting_8uc3_bald_512 = np.repeat(user_matting_8uc1_bald_512[:, :, np.newaxis], 3, axis=2) + + orig_h, orig_w, _ = user_rgb_8uc3_orisize.shape + + # TODO: 测试直接采用M 矩阵warp 到 user_hairstyle_M 后的512 尺寸 + bald_gene_8uc3_bald_512 = self.generator_bald.Geneator_Bald_inference(user_rgb_8uc3_bald_512, + user_matting_8uc3_bald_512, + user_baldseg_8uc3_bald_512, + user_landmark_f1k2_bald_512) + + inv_M_user = cv2.invertAffineTransform(user_bald_M) + bald_gene_8uc3_orisize = user_rgb_8uc3_orisize.copy() + + user_bald_gene_8uc3_orisize = cv2.warpAffine(bald_gene_8uc3_bald_512, inv_M_user, (orig_w, orig_h), + dst=bald_gene_8uc3_orisize, borderMode=cv2.BORDER_TRANSPARENT) + + user_baldseg_8uc3_512 = cv2.warpAffine(user_baldseg_8uc3_orisize, user_hairstyle_M, + (self.hair_size, self.hair_size)) + user_bald_8uc3_512 = cv2.warpAffine(user_bald_gene_8uc3_orisize, user_hairstyle_M, + (self.hair_size, self.hair_size), borderValue=[255, 255, 255]) + + ref_baldseg_8uc3_512 = cv2.warpAffine(ref_baldseg_8uc3_orisize, ref_hairstyle_M, + (self.hair_size, self.hair_size)) + + return user_rgb_8uc3_512, user_baldseg_8uc3_512, user_bald_8uc3_512, user_landmark_f1k2_512, user_hairstyle_M, ref_rgb_8uc3_512, ref_matting_8uc3_512, ref_baldseg_8uc3_512, ref_landmark_f1k2_512 + + def get_prepare_user_data(self, user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize): + """ + input：F + + user_rgb_8uc3_orisize: 用户图 原始尺寸 + user_landmark_1k2_f_orisize: 用户图 1k 点 + + output: + + 图像尺寸基于： 人脸 512， 图像均为3通道 + + user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 + user_hairstyle_M: 用户图 原图到512小图 M 矩阵 + """ + + # 得到512 小图 M 矩阵 + user_hairstyle_M = self.get_hair_M(user_landmark_1k2_f_orisize) + + # 1k 点 转换到 512 尺寸 + user_landmark_f1k2_512 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_hairstyle_M) + + # 用户图 头发毛躁 warpaffine + usr_ratio = np.sqrt( + (user_hairstyle_M[0][0] * user_hairstyle_M[0][0]) + (user_hairstyle_M[1][0] * user_hairstyle_M[1][0])) + user_hairstyle_M_tmp = user_hairstyle_M / usr_ratio + user_rgb_8uc3_512 = cv2.warpAffine(user_rgb_8uc3_orisize, user_hairstyle_M_tmp, + (int(self.hair_size / usr_ratio), int(self.hair_size / usr_ratio)), + borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) + user_rgb_8uc3_512 = cv2.resize(user_rgb_8uc3_512, (self.hair_size, self.hair_size), fx=usr_ratio, fy=usr_ratio, + interpolation=cv2.INTER_AREA) + + # 光头分割 + user_baldseg_8uc3_orisize = self.generator_baldseg.forward(user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize) + + # 用户图 生成光头 + user_rgb_8uc3_bald_512, user_baldseg_8uc3_bald_512, user_landmark_f1k2_bald_512, \ + user_bald_M = self.get_prepare_data_bald(user_rgb_8uc3_orisize, user_baldseg_8uc3_orisize, + user_landmark_1k2_f_orisize) + + # 用户图得到 头发 matting 小图 for blad + _, user_matting_8uc1_bald_512 , ref_rgb_resize= self.generator_matte.matte_inference(user_rgb_8uc3_bald_512, + user_landmark_f1k2_bald_512) + if user_matting_8uc1_bald_512.shape[0] != 512 or user_matting_8uc1_bald_512.shape[1] != 512: + user_matting_8uc1_bald_512 = cv2.resize(user_matting_8uc1_bald_512, + (user_rgb_8uc3_bald_512.shape[1], user_rgb_8uc3_bald_512.shape[0]), + interpolation=cv2.INTER_CUBIC) + user_matting_8uc3_bald_512 = np.repeat(user_matting_8uc1_bald_512[:, :, np.newaxis], 3, axis=2) + + orig_h, orig_w, _ = user_rgb_8uc3_orisize.shape + + # TODO: 测试直接采用M 矩阵warp 到 user_hairstyle_M 后的512 尺寸 + bald_gene_8uc3_bald_512 = self.generator_bald.Geneator_Bald_inference( + user_rgb_8uc3_bald_512, + user_matting_8uc3_bald_512, + user_baldseg_8uc3_bald_512, + user_landmark_f1k2_bald_512) + + inv_M_user = cv2.invertAffineTransform(user_bald_M) + bald_gene_8uc3_orisize = user_rgb_8uc3_orisize.copy() + + user_bald_gene_8uc3_orisize = cv2.warpAffine(bald_gene_8uc3_bald_512, inv_M_user, (orig_w, orig_h), + dst=bald_gene_8uc3_orisize, borderMode=cv2.BORDER_TRANSPARENT) + + user_baldseg_8uc3_512 = cv2.warpAffine(user_baldseg_8uc3_orisize, user_hairstyle_M, + (self.hair_size, self.hair_size)) + user_bald_8uc3_512 = cv2.warpAffine(user_bald_gene_8uc3_orisize, user_hairstyle_M, + (self.hair_size, self.hair_size), borderValue=[255, 255, 255]) + + return user_rgb_8uc3_512, user_baldseg_8uc3_512, user_bald_8uc3_512, user_landmark_f1k2_512, user_hairstyle_M + + def get_prepare_user_768_data(self, user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize, ratio=1): + """ + input： + + user_rgb_8uc3_orisize: 用户图 原始尺寸 + user_landmark_1k2_f_orisize: 用户图 1k 点 + + output: + + 图像尺寸基于： 人脸 512， 图像均为3通道 + + user_baldseg_8uc3_768：用户图 光头分割 mask， uint8 （0-255） + user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + user_landmark_f1k2_768： 用户图 关键点 1k*2 float32 + user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 + user_color_M: 用户图 原图到768小图 M 矩阵 + user_hairstyle_M: 用户图 原图到512小图 M 矩阵 + """ + # 得到768 小图 M 矩阵 + + user_color_M = self.get_color_hair_M(user_landmark_1k2_f_orisize) + # 得到768 小图 M 矩阵 发型 + + if ratio == 0: + user_hairstyle_M = self.get_hair_M_boy_v1(user_landmark_1k2_f_orisize) + elif ratio == 1: + user_hairstyle_M = self.get_hair_M_girl_v1(user_landmark_1k2_f_orisize) + elif ratio == 2: + user_hairstyle_M = self.get_hair_M_girl_v2(user_landmark_1k2_f_orisize) + else: + user_hairstyle_M = self.get_hair_M_girl_v1(user_landmark_1k2_f_orisize) + + # user_color_M = user_hairstyle_M + + # 1k 点 转换到 768 尺寸 + user_landmark_f1k2_bald_768 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_color_M) + # 1k 点 转换到 768 尺寸 发型 + user_landmark_f1k2_768 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_hairstyle_M) + + # 用户图 头发毛躁 warpaffine 768 + usr_color_ratio = np.sqrt( + (user_color_M[0][0] * user_color_M[0][0]) + (user_color_M[1][0] * user_color_M[1][0])) + user_color_M_tmp = user_color_M / usr_color_ratio + user_rgb_8uc3_768 = cv2.warpAffine(user_rgb_8uc3_orisize, user_color_M_tmp, + (int(self.bald_output_size / usr_color_ratio), + int(self.bald_output_size / usr_color_ratio)), + borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) + user_rgb_8uc3_768 = cv2.resize(user_rgb_8uc3_768, (self.bald_output_size, self.bald_output_size), + fx=usr_color_ratio, fy=usr_color_ratio, + interpolation=cv2.INTER_AREA) + + user_rgb_8uc3_768_bald = cv2.warpAffine(user_rgb_8uc3_orisize, user_color_M_tmp, + (int(self.bald_output_size / usr_color_ratio), + int(self.bald_output_size / usr_color_ratio))) + + user_rgb_8uc3_768_bald = cv2.resize(user_rgb_8uc3_768_bald, (self.bald_output_size, self.bald_output_size), + fx=usr_color_ratio, fy=usr_color_ratio, + interpolation=cv2.INTER_AREA) + + # 用户图得到 头发 matting 小图 for blad + user_matting_fg_8uc3_bald_768, user_matting_8uc1_bald_768, resize_source = self.generator_matte.matte_inference(user_rgb_8uc3_768, + user_landmark_f1k2_bald_768) + # cv2.imshow('user_matting_fg_8uc3_bald_768', user_matting_fg_8uc3_bald_768) + # cv2.imshow('user_matting_8uc1_bald_768', user_matting_8uc1_bald_768) + # cv2.imshow('user_rgb_8uc3_768', user_rgb_8uc3_768) + # cv2.waitKey() + if user_matting_8uc1_bald_768.shape[0] != 768 or user_matting_8uc1_bald_768.shape[1] != 768: + user_matting_8uc1_bald_768 = cv2.resize(user_matting_8uc1_bald_768, + (user_rgb_8uc3_768.shape[1], user_rgb_8uc3_768.shape[0]), + interpolation=cv2.INTER_CUBIC) + + # user_matting_fg_8uc3_bald_768 = cv2.resize(user_matting_fg_8uc3_bald_768, + # (user_rgb_8uc3_768.shape[1], user_rgb_8uc3_768.shape[0]), + # interpolation=cv2.INTER_CUBIC) + + user_matting_8uc3_bald_768 = np.repeat(user_matting_8uc1_bald_768[:, :, np.newaxis], 3, axis=2) + + user_matting_8uc3_bald_orisize = cv2.warpAffine(user_matting_8uc1_bald_768, + cv2.invertAffineTransform(user_color_M), + ( + user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), + flags=cv2.INTER_CUBIC) + + # 光头分割 + with torch.no_grad(): + user_baldseg_8uc3_orisize = self.generator_baldseg.forward(user_rgb_8uc3_orisize, + user_matting_8uc3_bald_orisize, + user_landmark_1k2_f_orisize) + # 用户图 生成光头 需准备 768 尺寸； 512 尺寸各一 + + user_baldseg_8uc3_bald_768, user_baldseg_8uc3_bald_512, user_bald_M \ + = self.get_prepare_data_bald_768(user_baldseg_8uc3_orisize, + user_landmark_1k2_f_orisize, user_color_M) + + orig_h, orig_w, _ = user_rgb_8uc3_orisize.shape + + bald_gene_8uc3_bald_768, user_hair_mask_blur_768 = self.generator_bald.Geneator_Bald_inference( + user_rgb_8uc3_768_bald.copy(), + user_matting_8uc3_bald_768, + user_baldseg_8uc3_bald_768, + user_landmark_f1k2_bald_768) + + inv_M_user = cv2.invertAffineTransform(user_color_M) + bald_gene_8uc3_orisize = user_rgb_8uc3_orisize.copy() + + user_bald_gene_8uc3_orisize = cv2.warpAffine(bald_gene_8uc3_bald_768, inv_M_user, (orig_w, orig_h), + dst=bald_gene_8uc3_orisize.copy(), + borderMode=cv2.BORDER_TRANSPARENT) + + user_hair_mask_blur_orisize = cv2.warpAffine(user_hair_mask_blur_768, inv_M_user, (orig_w, orig_h), + flags=cv2.INTER_CUBIC) + + user_bald_res_8uc3_orisize = ( + (1 - user_hair_mask_blur_orisize[:, :, np.newaxis] / 255) * user_rgb_8uc3_orisize + \ + (user_hair_mask_blur_orisize[:, :, np.newaxis] / 255) * user_bald_gene_8uc3_orisize).astype( + np.uint8) + + user_baldseg_8uc3_768 = cv2.warpAffine(user_baldseg_8uc3_orisize, user_hairstyle_M, + (self.hair_size, self.hair_size)) # , borderValue=[255, 255, 255] + user_bald_8uc3_768 = cv2.warpAffine(user_bald_res_8uc3_orisize, user_hairstyle_M, + (self.hair_size, self.hair_size)) # , borderValue=[255, 255, 255] + + return user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, \ + user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M, user_matting_8uc3_bald_orisize + + def get_prepare_ref_data(self, ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize): + """ + input： + ref_rgb_8uc3_orisize: 参考图 原始尺寸 + ref_landmark_1k2_f_orisize: 参考图 1k 点 + + output: + + 图像尺寸基于： 人脸 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 + + """ + + # 得到512 小图 M 矩阵 + ref_hairstyle_M = self.get_hair_M(ref_landmark_1k2_f_orisize) + # ref_color_hair_M = self.get_color_hair_M(ref_landmark_1k2_f_orisize) + + # 1k 点 转换到 512 尺寸 + ref_landmark_f1k2_512 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, ref_hairstyle_M) + # ref_landmark_color_f1k2_512 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, ref_color_hair_M) + + # 参考图 头发毛躁 warpaffine + ratio = np.sqrt( + (ref_hairstyle_M[0][0] * ref_hairstyle_M[0][0]) + (ref_hairstyle_M[1][0] * ref_hairstyle_M[1][0])) + ref_hairstyle_M_tmp = ref_hairstyle_M / ratio + ref_rgb_8uc3_512 = cv2.warpAffine(ref_rgb_8uc3_orisize, ref_hairstyle_M_tmp, + (int(self.hair_size / ratio), int(self.hair_size / ratio)), + borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) + ref_rgb_8uc3_512 = cv2.resize(ref_rgb_8uc3_512, (self.hair_size, self.hair_size), fx=ratio, fy=ratio, + interpolation=cv2.INTER_AREA) + + # # 参考图 头发毛躁 warpaffine + # ratio_color = np.sqrt( + # (ref_color_hair_M[0][0] * ref_color_hair_M[0][0]) + (ref_color_hair_M[1][0] * ref_color_hair_M[1][0])) + # ref_color_hair_M_tmp = ref_color_hair_M / ratio_color + # ref_rgb_8uc3_768 = cv2.warpAffine(ref_rgb_8uc3_orisize, ref_color_hair_M_tmp, + # (int(self.bald_output_size / ratio_color), int(self.bald_output_size / ratio_color)), + # borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) + # ref_rgb_8uc3_768 = cv2.resize(ref_rgb_8uc3_768, (self.bald_output_size, self.bald_output_size), fx=ratio_color, fy=ratio_color, + # interpolation=cv2.INTER_AREA) + + # 参考图直接得到 头发 matting 小图 + _, ref_matte_pred_8uc1_512,ref_rgb_resize = self.generator_matte.matte_inference(ref_rgb_8uc3_512, ref_landmark_f1k2_512) + torch.cuda.empty_cache() + + ref_matting_8uc1_512 = cv2.resize(ref_matte_pred_8uc1_512, + (ref_rgb_8uc3_512.shape[1], ref_rgb_8uc3_512.shape[0]), + interpolation=cv2.INTER_CUBIC) + + ref_matting_8uc3_512 = np.repeat(ref_matting_8uc1_512[:, :, np.newaxis], 3, axis=2) + + # 光头分割 + ref_baldseg_8uc3_orisize = self.generator_baldseg.forward(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize) + + + ref_baldseg_8uc3_512 = cv2.warpAffine(ref_baldseg_8uc3_orisize, ref_hairstyle_M, + (self.hair_size, self.hair_size), borderValue=[255, 255, 255]) + + return ref_rgb_8uc3_512, ref_matting_8uc3_512, ref_baldseg_8uc3_512, ref_landmark_f1k2_512 + + def get_prepare_ref_768_data(self, ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize, ratio=1): + """ + input： + ref_rgb_8uc3_orisize: 参考图 原始尺寸 + ref_landmark_1k2_f_orisize: 参考图 1k 点 + + output: + + 图像尺寸基于： 人脸 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 + + """ + # 得到512 小图 M 矩阵 + + if ratio == 0: + ref_hairstyle_M = self.get_hair_M_boy_v1(ref_landmark_1k2_f_orisize) + elif ratio == 1: + ref_hairstyle_M = self.get_hair_M_girl_v1(ref_landmark_1k2_f_orisize) + elif ratio == 2: + ref_hairstyle_M = self.get_hair_M_girl_v2(ref_landmark_1k2_f_orisize) + else: + ref_hairstyle_M = self.get_hair_M_girl_v1(ref_landmark_1k2_f_orisize) + # ref_color_hair_M = self.get_color_hair_M(ref_landmark_1k2_f_orisize) + + # print("ref_hairstyle_M: ", ref_hairstyle_M) + + # 1k 点 转换到 512 尺寸 + ref_landmark_f1k2_768 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, ref_hairstyle_M) + + # 参考图 头发毛躁 warpaffine + ratio = np.sqrt( + (ref_hairstyle_M[0][0] * ref_hairstyle_M[0][0]) + (ref_hairstyle_M[1][0] * ref_hairstyle_M[1][0])) + ref_hairstyle_M_tmp = ref_hairstyle_M / ratio + ref_rgb_8uc3_768 = cv2.warpAffine(ref_rgb_8uc3_orisize, ref_hairstyle_M_tmp, + (int(self.bald_output_size / ratio), int(self.bald_output_size / ratio)), + borderMode=cv2.BORDER_CONSTANT) # , borderValue=[255, 255, 255] + ref_rgb_8uc3_768 = cv2.resize(ref_rgb_8uc3_768, (self.bald_output_size, self.bald_output_size), fx=ratio, + fy=ratio, + interpolation=cv2.INTER_AREA) + # cv2.imshow('fff', ref_rgb_8uc3_768) + # cv2.waitKey() + # 参考图直接得到 头发 matting 小图 + with torch.no_grad(): + ref_matte_fg_8uc3_768, ref_matte_pred_8uc1_768, ref_rgb_resize = self.generator_matte.matte_inference(ref_rgb_8uc3_768, + ref_landmark_f1k2_768) + + torch.cuda.empty_cache() + + ref_matting_fg_8uc3_768 = cv2.resize(ref_matte_fg_8uc3_768, + (ref_rgb_8uc3_768.shape[1], ref_rgb_8uc3_768.shape[0]), + interpolation=cv2.INTER_CUBIC) + + ref_matting_8uc1_768 = cv2.resize(ref_matte_pred_8uc1_768, + (ref_rgb_8uc3_768.shape[1], ref_rgb_8uc3_768.shape[0]), + interpolation=cv2.INTER_CUBIC) + + ref_matting_8uc3_768 = np.repeat(ref_matting_8uc1_768[:, :, np.newaxis], 3, axis=2) + + ref_matting_8uc1_orisize = cv2.warpAffine(ref_matting_8uc1_768, cv2.invertAffineTransform(ref_hairstyle_M), + (ref_rgb_8uc3_orisize.shape[1], ref_rgb_8uc3_orisize.shape[0]), + flags=cv2.INTER_CUBIC) + + # 光头分割 + with torch.no_grad(): + ref_baldseg_8uc3_orisize = self.generator_baldseg.forward(ref_rgb_8uc3_orisize, ref_matting_8uc1_orisize, + ref_landmark_1k2_f_orisize) + + ref_baldseg_8uc3_768 = cv2.warpAffine(ref_baldseg_8uc3_orisize, ref_hairstyle_M, + (self.bald_output_size, self.bald_output_size), + flags=cv2.INTER_NEAREST) # , borderValue=[255, 255, 255] + + return ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768 + + def get_prepare_ref_768_bald_data(self, ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize): + """ + input： + ref_rgb_8uc3_orisize: 参考图 原始尺寸 + ref_landmark_1k2_f_orisize: 参考图 1k 点 + + output: + + 图像尺寸基于： 人脸 512， 图像均为3通道 + ref_rgb_8uc3_768: 参考图， size 512， uint8 （0-255） + ref_matting_8uc3_768：参考图 matting alpha 值， uint8 （0-255） + ref_baldseg_8uc3_768: 参考图 光头分割， uint8 （0-255） + ref_landmark_f1k2_768： 参考图 关键点 1k*2 float32 + + """ + + # 得到512 小图 M 矩阵 + # ref_hairstyle_M = self.get_hair_M(ref_landmark_1k2_f_orisize) + ref_color_hair_M = self.get_color_hair_M(ref_landmark_1k2_f_orisize) + + # 1k 点 转换到 512 尺寸 + # ref_landmark_f1k2_512 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, ref_hairstyle_M) + ref_landmark_color_f1k2_768 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, ref_color_hair_M) + + # 参考图 头发毛躁 warpaffine + ratio = np.sqrt( + (ref_color_hair_M[0][0] * ref_color_hair_M[0][0]) + (ref_color_hair_M[1][0] * ref_color_hair_M[1][0])) + ref_color_hair_M_tmp = ref_color_hair_M / ratio + ref_rgb_8uc3_768 = cv2.warpAffine(ref_rgb_8uc3_orisize, ref_color_hair_M_tmp, + (int(self.color_output_size / ratio), int(self.color_output_size / ratio)), + borderMode=cv2.BORDER_CONSTANT) # , borderValue=[255, 255, 255] + ref_rgb_8uc3_768 = cv2.resize(ref_rgb_8uc3_768, (self.color_output_size, self.color_output_size), fx=ratio, + fy=ratio, + interpolation=cv2.INTER_AREA) + + # 参考图直接得到 头发 matting 小图 + _, ref_matte_pred_8uc1_768 , ref_rgb_resize= self.generator_matte.matte_inference(ref_rgb_8uc3_768, ref_landmark_color_f1k2_768) + torch.cuda.empty_cache() + + ref_matting_8uc1_768 = cv2.resize(ref_matte_pred_8uc1_768, + (ref_rgb_8uc3_768.shape[1], ref_rgb_8uc3_768.shape[0]), + interpolation=cv2.INTER_CUBIC) + + ref_matting_8uc3_768 = np.repeat(ref_matting_8uc1_768[:, :, np.newaxis], 3, axis=2) + + ref_matting_8uc1_orisize = cv2.warpAffine(ref_matting_8uc1_768, cv2.invertAffineTransform(ref_color_hair_M), + (ref_rgb_8uc3_orisize.shape[1], ref_rgb_8uc3_orisize.shape[0]), + flags=cv2.INTER_CUBIC) + + # 光头分割 + with torch.no_grad(): + ref_baldseg_8uc3_orisize = self.generator_baldseg.forward(ref_rgb_8uc3_orisize, ref_matting_8uc1_orisize, + ref_landmark_1k2_f_orisize) + + ref_baldseg_8uc3_768 = cv2.warpAffine(ref_baldseg_8uc3_orisize, ref_color_hair_M, + (self.color_output_size, + self.color_output_size)) # , borderValue=[255, 255, 255] + + return ref_rgb_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_color_f1k2_768 + + def get_prepare_hair_color_data(self, user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize, ref_rgb_8uc3_orisize, + ref_landmark_1k2_f_orisize): + + """ + + input： + + user_rgb_8uc3_orisize: 用户图 原始尺寸 + user_landmark_1k2_f_orisize: 用户图 1k 点 + ref_rgb_8uc3_orisize: 参考图 原始尺寸 + ref_landmark_1k2_f_orisize: 参考图 1k 点 + + output: + + 图像尺寸基于： 人脸 768， 图像均为3通道 + user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768 + + user_rgb_8uc3_change_color_768: 用户图 align， size: 768 uint8 （0-255） + user_matting_8uc3_change_color_768: 用户图 matting， size: 768 uint8 （0-255） + + ref_rgb_8uc3_change_color_768: 参考图， align, size 512， uint8 （0-255） + ref_matting_8uc3_change_color_768: 参考图 matting alpha 值， uint8 （0-255） + + """ + + # 得到768 小图 M 矩阵 + user_hair_color_M = self.get_color_hair_M(user_landmark_1k2_f_orisize) + ref_hair_color_M = self.get_color_hair_M(ref_landmark_1k2_f_orisize) + + # 1k 点 转换到 768 尺寸 + user_landmark_f1k2_768 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_hair_color_M) + ref_landmark_f1k2_768 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, ref_hair_color_M) + + # 用户图 头发毛躁 warpaffine + usr_ratio = np.sqrt( + (user_hair_color_M[0][0] * user_hair_color_M[0][0]) + (user_hair_color_M[1][0] * user_hair_color_M[1][0])) + user_hair_color_M_tmp = user_hair_color_M / usr_ratio + user_rgb_8uc3_768 = cv2.warpAffine(user_rgb_8uc3_orisize, user_hair_color_M_tmp, + (int(self.color_output_size / usr_ratio), + int(self.color_output_size / usr_ratio))) + user_rgb_8uc3_768 = cv2.resize(user_rgb_8uc3_768, (self.color_output_size, self.color_output_size), + fx=usr_ratio, fy=usr_ratio, + interpolation=cv2.INTER_AREA) + + # 参考图 头发毛躁 warpaffine + ratio = np.sqrt( + (ref_hair_color_M[0][0] * ref_hair_color_M[0][0]) + (ref_hair_color_M[1][0] * ref_hair_color_M[1][0])) + ref_hair_color_M_tmp = ref_hair_color_M / ratio + ref_rgb_8uc3_768 = cv2.warpAffine(ref_rgb_8uc3_orisize, ref_hair_color_M_tmp, + (int(self.color_output_size / ratio), int(self.color_output_size / ratio))) + + ref_rgb_8uc3_768 = cv2.resize(ref_rgb_8uc3_768, (self.color_output_size, self.color_output_size), fx=ratio, + fy=ratio, + interpolation=cv2.INTER_AREA) + + # 参考图直接得到 头发 matting 小图 + with torch.no_grad(): + _, ref_matte_pred_8uc1_768 , ref_rgb_resize= self.generator_matte.matte_inference(ref_rgb_8uc3_768, ref_landmark_f1k2_768) + + # torch.cuda.empty_cache() + + ref_matting_8uc3_768 = np.repeat(ref_matte_pred_8uc1_768[:, :, np.newaxis], 3, axis=2) + + # 用户图得到 头发 matting 小图 + with torch.no_grad(): + _, user_matting_8uc1_bald_768, ref_rgb_resize = self.generator_matte.matte_inference(user_rgb_8uc3_768, user_landmark_f1k2_768) + + user_matting_8uc3_768 = np.repeat(user_matting_8uc1_bald_768[:, :, np.newaxis], 3, axis=2) + + return user_rgb_8uc3_768, user_matting_8uc3_768, ref_rgb_8uc3_768, ref_matting_8uc3_768 + + def get_prepare_hair_color_user_data(self, user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize): + + """ + + input： + + user_rgb_8uc3_orisize: 用户图 原始尺寸 + user_landmark_1k2_f_orisize: 用户图 1k 点 + + output: + + 图像尺寸基于： 人脸 768， 图像均为3通道 + user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768 + + user_rgb_8uc3_change_color_768: 用户图 align， size: 768 uint8 （0-255） + user_matting_8uc3_change_color_768: 用户图 matting， size: 768 uint8 （0-255） + + """ + + # 得到768 小图 M 矩阵 + user_hair_color_M = self.get_color_hair_M(user_landmark_1k2_f_orisize) + with torch.no_grad(): + _, user_matting_8uc1_orisize, _ = self.generator_matte.matte_inference(user_rgb_8uc3_orisize, + user_landmark_1k2_f_orisize) + user_matting_8uc3_orisize = np.repeat(user_matting_8uc1_orisize[:, :, np.newaxis], 3, axis=2) + + # 1k 点 转换到 768 尺寸 + user_landmark_f1k2_768 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_hair_color_M) + + # # 用户图 头发毛躁 warpaffine + # usr_ratio = np.sqrt( + # (user_hair_color_M[0][0] * user_hair_color_M[0][0]) + (user_hair_color_M[1][0] * user_hair_color_M[1][0])) + # user_hair_color_M_tmp = user_hair_color_M / usr_ratio + # user_rgb_8uc3_768 = cv2.warpAffine(user_rgb_8uc3_orisize, user_hair_color_M_tmp, + # (int(self.color_output_size / usr_ratio), int(self.color_output_size / usr_ratio))) + # user_rgb_8uc3_768 = cv2.resize(user_rgb_8uc3_768, (self.color_output_size, self.color_output_size), fx=usr_ratio, fy=usr_ratio, + # interpolation=cv2.INTER_AREA) + + user_rgb_8uc3_768 = cv2.warpAffine(user_rgb_8uc3_orisize, user_hair_color_M, + (self.color_output_size, self.color_output_size)) + user_matting_8uc3_768 = cv2.warpAffine(user_matting_8uc3_orisize, user_hair_color_M, + (self.color_output_size, self.color_output_size)) + + + # 用户图得到 头发 matting 小图 + # _, user_matting_8uc1_bald_768 , ref_rgb_resize= self.generator_matte.matte_inference(user_rgb_8uc3_768, user_landmark_f1k2_768) + + # user_matting_8uc3_768 = np.repeat(user_matting_8uc1_bald_768[:, :, np.newaxis], 3, axis=2) + + return user_rgb_8uc3_768, user_matting_8uc3_768, user_hair_color_M, user_matting_8uc3_orisize + + def get_matte_img(self, img, landmark1k): + + matte_fg, matte_img , ref_rgb_resize= self.generator_matte.matte_inference(img, landmark1k) + + return matte_fg, matte_img + + def get_max_countour(self, contours): + index_contour = 0 + max_num = 0 + for i, contour in enumerate(contours): + if len(contour) > max_num: + max_num = len(contour) + index_contour = i + return index_contour + + def judge_hair_pos(self, user_matting_8uc3_bald_768, user_baldseg_8uc3_bald_768, retData): + cloth_mask = np.zeros_like(user_baldseg_8uc3_bald_768, dtype=np.uint8) + cloth_pos = (user_baldseg_8uc3_bald_768[:, :, 0] < 10) & (user_baldseg_8uc3_bald_768[:, :, 1] < 10) & ( + user_baldseg_8uc3_bald_768[:, :, 2] > 250) + cloth_mask[cloth_pos] = 255 + + cloth_withhair_mask = (cloth_mask * user_matting_8uc3_bald_768.astype(np.float32) / 255).astype(np.uint8) + cloth_withhair_count = len(cloth_withhair_mask[cloth_withhair_mask[:, :, 0] > 0]) + + retData["cloth_withhair_count"] = str(cloth_withhair_count) + + def get_fusion_res(self, user_baldseg_8uc3_512, user_bald_8uc3_512, hair_gene_8uc3_512, user_landmark_f1k2_512): + """ + input： + + 图像尺寸基于： 人脸 512， 图像均为3通道 + + user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + hair_gene_8uc3_512: 用户图 换发型 贴合 光头， uint8 （0-255） + user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 + + output： + + hair_gene_fusion_8uc3_512: 生成 发型图， uint8 （0-255） + + """ + # "user_baldseg_8uc3_512, user_bald_8uc3_512, hair_gene_8uc3_512" + + # 换发型后的 matting 图 + hair_gene_matte_fg_8uc3_512, hair_gene_matte_8uc1_512 = self.get_matte_img(hair_gene_8uc3_512, + user_landmark_f1k2_512) + hair_gene_matte_32fc1_512 = hair_gene_matte_8uc1_512.astype(np.float32) / 255 + hair_gene_matte_32fc3_512 = np.repeat(hair_gene_matte_32fc1_512[:, :, np.newaxis], 3, axis=2) + + hair_bg = (user_bald_8uc3_512 * (1 - hair_gene_matte_32fc3_512)).astype(np.uint8) + + hair_bg_face = np.zeros_like(user_bald_8uc3_512) + + face_index = (user_baldseg_8uc3_512[:, :, 1] == 255) & (user_baldseg_8uc3_512[:, :, 0] == 0) & ( + user_baldseg_8uc3_512[:, :, 2] == 0) + + hair_bg_face[face_index] = user_bald_8uc3_512[face_index] + + hair_face_bg = np.zeros_like(user_bald_8uc3_512) + hair_face_bg[face_index] = hair_gene_8uc3_512[face_index] + + hair_face_mask = np.zeros_like(user_bald_8uc3_512, dtype=np.uint8) + hair_face_mask[face_index] = 255 + kernel_size = 7 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + hair_face_mask = cv2.erode(hair_face_mask, kernel, iterations=1).astype(np.uint8) + hair_face_mask_blur = cv2.blur(hair_face_mask, (5, 5)) + + hair_gene_fusion_8uc3_512 = (user_bald_8uc3_512 * ( + 1 - hair_gene_matte_32fc3_512) + hair_gene_matte_fg_8uc3_512 * hair_gene_matte_32fc3_512).astype( + np.uint8) + + hair_gene_fusion_8uc3_512 = (hair_gene_fusion_8uc3_512 * ( + 1 - hair_face_mask_blur.astype(np.float32) / 255) + hair_face_bg * hair_face_mask_blur.astype( + np.float32) / 255).astype(np.uint8) + + + return hair_gene_fusion_8uc3_512 + + def get_fusion_res_hairpaste(self, user_bald_8uc3_orisize, hair_gene_8uc3_768, + user_landmark_f1k2_768, + user_hairstyle_M): + """ + input： + + 图像尺寸基于： 人脸 512， 图像均为3通道 + + user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + hair_gene_8uc3_512: 用户图 换发型 贴合 光头， uint8 （0-255） + user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 + + output： + + hair_gene_fusion_8uc3_512: 生成 发型图， uint8 （0-255） + + """ + hair_gene_matte_fg_8uc3_768, hair_gene_matte_8uc1_768 = self.get_matte_img(hair_gene_8uc3_768, + user_landmark_f1k2_768) + + # hair_gene_8uc3_orisize = cv2.warpAffine(hair_gene_8uc3_768, cv2.invertAffineTransform(user_hairstyle_M), + # (user_bald_8uc3_orisize.shape[1], + # user_bald_8uc3_orisize.shape[0]), + # dst=user_bald_8uc3_orisize.copy(), + # borderMode=cv2.BORDER_TRANSPARENT) + + + hair_gene_matte_8uc1_orisize = cv2.warpAffine(hair_gene_matte_8uc1_768, + cv2.invertAffineTransform(user_hairstyle_M), + (user_bald_8uc3_orisize.shape[1], + user_bald_8uc3_orisize.shape[0]), + flags=cv2.INTER_CUBIC) + + + hair_gene_matte_fg_8uc3_orisize = cv2.warpAffine(hair_gene_matte_fg_8uc3_768, + cv2.invertAffineTransform(user_hairstyle_M), + (user_bald_8uc3_orisize.shape[1], + user_bald_8uc3_orisize.shape[0]), + flags=cv2.INTER_CUBIC) + + hair_gene_matte_8uc3_orisize = np.repeat(hair_gene_matte_8uc1_orisize[:, :, np.newaxis], 3, axis=2) + + hair_gene_matte_32fc1_orisize = hair_gene_matte_8uc1_orisize.astype(np.float32) / 255 + hair_gene_matte_32fc3_orisize = np.repeat(hair_gene_matte_32fc1_orisize[:, :, np.newaxis], 3, axis=2) + + hair_gene_fusion_8uc3_orisize = (user_bald_8uc3_orisize * ( + 1 - hair_gene_matte_32fc3_orisize) + hair_gene_matte_fg_8uc3_orisize * hair_gene_matte_32fc3_orisize).astype( + np.uint8) + + return hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize + + def get_fusion_res_hairblur(self, user_baldseg_8uc3_orisize, user_bald_8uc3_orisize, hair_gene_8uc3_768, + user_landmark_f1k2_768, user_hairstyle_M, ref_hairstyle_dir): + """ + input： + + 图像尺寸基于： 人脸 512， 图像均为3通道 + + user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + hair_gene_8uc3_512: 用户图 换发型 贴合 光头， uint8 （0-255） + user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 + + output： + + hair_gene_fusion_8uc3_512: 生成 发型图， uint8 （0-255） + + """ + + + # 换发型后的 matting 图 + + hair_gene_matte_fg_8uc3_768, hair_gene_matte_8uc1_768 = self.get_matte_img(hair_gene_8uc3_768, + user_landmark_f1k2_768) + + + hair_gene_8uc3_orisize = cv2.warpAffine(hair_gene_8uc3_768, cv2.invertAffineTransform(user_hairstyle_M), + (user_bald_8uc3_orisize.shape[1], + user_bald_8uc3_orisize.shape[0]), + dst=user_bald_8uc3_orisize.copy(), + borderMode=cv2.BORDER_TRANSPARENT) + + default_hair_tail_mask = cv2.imread(os.path.join(ref_hairstyle_dir, "default_hair_tail_mask.png")) + + hair_gene_matte_8uc1_768 = ( + hair_gene_matte_8uc1_768 * default_hair_tail_mask[:, :, 0].astype(np.float32) / 255).astype( + np.uint8) + + hair_gene_matte_8uc1_orisize = cv2.warpAffine(hair_gene_matte_8uc1_768, + cv2.invertAffineTransform(user_hairstyle_M), + (user_bald_8uc3_orisize.shape[1], + user_bald_8uc3_orisize.shape[0]), + flags=cv2.INTER_CUBIC) + + hair_gene_matte_fg_8uc3_768 = ( + hair_gene_matte_fg_8uc3_768 * default_hair_tail_mask.astype(np.float32) / 255).astype(np.uint8) + + + hair_gene_matte_fg_8uc3_orisize = cv2.warpAffine(hair_gene_matte_fg_8uc3_768, + cv2.invertAffineTransform(user_hairstyle_M), + (user_bald_8uc3_orisize.shape[1], + user_bald_8uc3_orisize.shape[0]), + flags=cv2.INTER_CUBIC) + + hair_gene_matte_32fc1_orisize = hair_gene_matte_8uc1_orisize.astype(np.float32) / 255 + hair_gene_matte_32fc3_orisize = np.repeat(hair_gene_matte_32fc1_orisize[:, :, np.newaxis], 3, axis=2) + + + hair_bg_face = np.zeros_like(user_bald_8uc3_orisize) + + face_index = (user_baldseg_8uc3_orisize[:, :, 1] == 255) & (user_baldseg_8uc3_orisize[:, :, 0] == 0) & ( + user_baldseg_8uc3_orisize[:, :, 2] == 0) + + hair_bg_face[face_index] = user_bald_8uc3_orisize[face_index] + + hair_face_bg = np.zeros_like(user_bald_8uc3_orisize) + hair_face_bg[face_index] = hair_gene_8uc3_orisize[face_index] + + hair_face_mask = np.zeros_like(user_bald_8uc3_orisize, dtype=np.uint8) + hair_face_mask[face_index] = 255 + + hair_face_rec_mask = (((hair_gene_matte_8uc1_orisize > 0) & (hair_face_mask[:, :, 0] > 0)).astype( + np.float32) * 255).astype(np.uint8) + + hair_face_rec_mask = cv2.resize(hair_face_rec_mask, (768, 768), interpolation=cv2.INTER_CUBIC) + + kernel_size = 11 + kernel_dilate = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + hair_face_rec_mask_dilate = cv2.erode(hair_face_rec_mask, kernel_dilate, iterations=1) + + hair_face_rec_mask_blur = cv2.blur(hair_face_rec_mask_dilate, (31, 31)) + + hair_face_rec_mask_blur = cv2.resize(hair_face_rec_mask_blur, (user_bald_8uc3_orisize.shape[1], + user_bald_8uc3_orisize.shape[0]), + interpolation=cv2.INTER_CUBIC) + + + kernel_size_2 = 7 + kernel_erode_2 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size_2, kernel_size_2)) + + hair_face_mask_erode_2 = cv2.erode(hair_face_mask, kernel_erode_2, iterations=1).astype(np.uint8) + + hair_face_mask_blur = cv2.blur(hair_face_mask_erode_2, (5, 5)) + hair_face_mask_blur = cv2.resize(hair_face_mask_blur, (user_bald_8uc3_orisize.shape[1], + user_bald_8uc3_orisize.shape[0]), + interpolation=cv2.INTER_CUBIC) + + hair_gene_fusion_8uc3_orisize = (user_bald_8uc3_orisize * ( + 1 - hair_gene_matte_32fc3_orisize) + hair_gene_matte_fg_8uc3_orisize * hair_gene_matte_32fc3_orisize).astype( + np.uint8) + + hair_gene_fusion_8uc3_orisize_gene = hair_gene_fusion_8uc3_orisize.copy() + hair_gene_fusion_8uc3_orisize_gene = (hair_gene_fusion_8uc3_orisize_gene * ( + 1 - hair_face_mask_blur.astype(np.float32) / 255) + hair_face_bg * hair_face_mask_blur.astype( + np.float32) / 255).astype(np.uint8) + + hair_gene_fusion_8uc3_512 = (hair_gene_fusion_8uc3_orisize_gene * ( + hair_face_rec_mask_blur[:, :, np.newaxis].astype(np.float32) / 255) + + hair_gene_fusion_8uc3_orisize * ( + 1 - hair_face_rec_mask_blur[:, :, np.newaxis].astype( + np.float32) / 255)).astype(np.uint8) + + + return hair_gene_fusion_8uc3_512, hair_gene_8uc3_orisize + + def get_fusion_res_onlyhair(self, user_baldseg_8uc3_512, user_bald_8uc3_512, hair_gene_8uc3_512, + user_landmark_f1k2_512): + """ + input： + + 图像尺寸基于： 人脸 512， 图像均为3通道 + + user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + hair_gene_8uc3_512: 用户图 换发型 贴合 光头， uint8 （0-255） + user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 + + output： + + hair_gene_fusion_8uc3_512: 生成 发型图， uint8 （0-255） + + """ + # "user_baldseg_8uc3_512, user_bald_8uc3_512, hair_gene_8uc3_512" + + + # 换发型后的 matting 图 + hair_gene_matte_fg_8uc3_512, hair_gene_matte_8uc1_512 = self.get_matte_img(hair_gene_8uc3_512, + user_landmark_f1k2_512) + hair_gene_matte_32fc1_512 = hair_gene_matte_8uc1_512.astype(np.float32) / 255 + hair_gene_matte_32fc3_512 = np.repeat(hair_gene_matte_32fc1_512[:, :, np.newaxis], 3, axis=2) + + # hair_bg = (user_bald_8uc3_512 * (1 - hair_gene_matte_32fc3_512)).astype(np.uint8) + + hair_bg_face = np.zeros_like(user_bald_8uc3_512) + + face_index = (user_baldseg_8uc3_512[:, :, 1] == 255) & (user_baldseg_8uc3_512[:, :, 0] == 0) & ( + user_baldseg_8uc3_512[:, :, 2] == 0) + + hair_bg_face[face_index] = user_bald_8uc3_512[face_index] + + hair_face_bg = np.zeros_like(user_bald_8uc3_512) + hair_face_bg[face_index] = hair_gene_8uc3_512[face_index] + + hair_face_mask = np.zeros_like(user_bald_8uc3_512, dtype=np.uint8) + hair_face_mask[face_index] = 255 + kernel_size_1 = 15 + kernel_size_2 = 7 + kernel_erode_1 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size_1, kernel_size_1)) + kernel_erode_2 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size_2, kernel_size_2)) + + hair_face_mask_erode_1 = cv2.erode(hair_face_mask, kernel_erode_1, iterations=1).astype(np.uint8) + hair_face_mask_erode_2 = cv2.erode(hair_face_mask, kernel_erode_2, iterations=1).astype(np.uint8) + # hair_face_mask_dilate = cv2.dilate(hair_face_mask, kernel, iterations=1).astype(np.uint8) + + hair_face_mask_rolling = hair_face_mask_erode_2 - hair_face_mask_erode_1 + hair_face_mask_use = 255 - hair_face_mask_rolling + + hair_face_mask_blur = cv2.blur(hair_face_mask_erode_2, (5, 5)) + hair_face_mask_use_blur = cv2.blur(hair_face_mask_use, (11, 11)) + + + # hair_gene_matte_32fc3_512_blur = cv2.blur(hair_gene_matte_32fc3_512, (3, 3)) + + # user_bald_8uc3_white = np.full_like(user_bald_8uc3_512, (187, 202, 240), dtype=(np.uint8)) + # user_bald_8uc3_fg = (user_bald_8uc3_white * ( + # 1 - hair_gene_matte_32fc3_512) + hair_gene_matte_fg_8uc3_512 * hair_gene_matte_32fc3_512).astype( + # np.uint8) + + # paste_hair = np.zeros_like(user_bald_8uc3_512) + hair_gene_fusion_8uc3_512 = (user_bald_8uc3_512 * ( + 1 - hair_gene_matte_32fc3_512) + hair_gene_matte_fg_8uc3_512 * hair_gene_matte_32fc3_512).astype( + np.uint8) + + hair_gene_fusion_8uc3_512_gene = hair_gene_fusion_8uc3_512.copy() + hair_gene_fusion_8uc3_512_gene = (hair_gene_fusion_8uc3_512_gene * ( + 1 - hair_face_mask_blur.astype(np.float32) / 255) + hair_face_bg * hair_face_mask_blur.astype( + np.float32) / 255).astype(np.uint8) + + hair_gene_fusion_8uc3_512 = (hair_gene_fusion_8uc3_512_gene * (1 - hair_face_mask_use_blur.astype( + np.float32) / 255) + hair_gene_fusion_8uc3_512 * hair_face_mask_use_blur.astype(np.float32) / 255).astype( + np.uint8) + + return hair_gene_fusion_8uc3_512 + + def get_fusion_res_forehead(self, user_baldseg_8uc3_512, user_bald_8uc3_512, hair_gene_8uc3_512, + user_landmark_f1k2_512): + """ + input： + + 图像尺寸基于： 人脸 512， 图像均为3通道 + + user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + hair_gene_8uc3_512: 用户图 换发型 贴合 光头， uint8 （0-255） + user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 + + output： + + hair_gene_fusion_8uc3_512: 生成 发型图， uint8 （0-255） + + """ + + # 换发型后的 matting 图 + hair_gene_matte_fg_8uc3_512, hair_gene_matte_8uc1_512 = self.get_matte_img(hair_gene_8uc3_512, + user_landmark_f1k2_512) + hair_gene_matte_32fc1_512 = hair_gene_matte_8uc1_512.astype(np.float32) / 255 + hair_gene_matte_32fc3_512 = np.repeat(hair_gene_matte_32fc1_512[:, :, np.newaxis], 3, axis=2) + + hair_bg = (user_bald_8uc3_512 * (1 - hair_gene_matte_32fc3_512)).astype(np.uint8) + + hair_bg_face = np.zeros_like(user_bald_8uc3_512) + + face_index = (user_baldseg_8uc3_512[:, :, 1] == 255) & (user_baldseg_8uc3_512[:, :, 0] == 0) & ( + user_baldseg_8uc3_512[:, :, 2] == 0) + + hair_bg_face[face_index] = user_bald_8uc3_512[face_index] + + hair_face_bg = np.zeros_like(user_bald_8uc3_512) + hair_face_bg[face_index] = hair_gene_8uc3_512[face_index] + + hair_face_mask = np.zeros_like(user_bald_8uc3_512, dtype=np.uint8) + hair_face_mask[face_index] = 255 + + user_landmark_f137_512 = landmark_processor.pts_1k_to_137(user_landmark_f1k2_512) + pts_leye_up = user_landmark_f137_512[89:96, :] + pts_reye_up = user_landmark_f137_512[106:113, :] + pts_leye_up_low = pts_leye_up.min(axis=0)[1] + pts_reye_up_low = pts_reye_up.min(axis=0)[1] + + pts_eye_up_low = min(pts_leye_up_low, pts_reye_up_low) + hair_face_mask[int(pts_eye_up_low):, :, :] = 0 + + kernel_size = 7 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + hair_face_mask = cv2.erode(hair_face_mask, kernel, iterations=1).astype(np.uint8) + hair_face_mask_blur = cv2.blur(hair_face_mask, (5, 5)) + + # paste_hair = np.zeros_like(user_bald_8uc3_512) + hair_gene_fusion_8uc3_512 = (user_bald_8uc3_512 * ( + 1 - hair_gene_matte_32fc3_512) + hair_gene_matte_fg_8uc3_512 * hair_gene_matte_32fc3_512).astype( + np.uint8) + + hair_gene_fusion_8uc3_512 = (hair_gene_fusion_8uc3_512 * ( + 1 - hair_face_mask_blur.astype(np.float32) / 255) + hair_face_bg * hair_face_mask_blur.astype( + np.float32) / 255).astype(np.uint8) + + # cv2.waitKey() + + return hair_gene_fusion_8uc3_512 + +class Generator_Hair(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.model_dir = modelRoot + + self.dst_height = 768 + self.dst_width = 768 + + generator_hair_path = os.path.join(self.model_dir, "master_hair_v8_onlyhair_nowarp_all_0709.pt") # default master_hair_v8_onlyhair_blur_aug_0413.pt master_hair_v8_onlyhair_min_0223 + # "master_hair_v8_onlyhair_nowarp_all_0709" # master_hair_v8_onlyhair_blur_aug_all_0703 + self.net = torch.jit.load(generator_hair_path, torch.device('cpu')).to(self.device) + # print("generate hair net: ", self.net) + self.net.eval() + + def Generator_Hair_inference(self, ref_rgb_8uc3_512, ref_matting_8uc3_512, ref_baldseg_8uc3_512, + ref_landmark_f1k2_512, + user_baldseg_8uc3_512, user_bald_8uc3_512, user_landmark_f1k2_512): + + """ + 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 + + user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 + + output： + + hair_gene_8uc3_512: 生成 发型图， uint8 （0-255） + + """ + nohair_mask = user_baldseg_8uc3_512.copy() + user_bald_mask = (user_baldseg_8uc3_512 == [0, 255, 0]).all(axis=2) + user_pts137 = landmark_processor.pts_1k_to_137(user_landmark_f1k2_512).astype(np.int32) + # 2 ********************** nohair_mask ********************** + # Label mouth + cv2.fillPoly(nohair_mask, + np.concatenate((user_pts137[47:35:-1], user_pts137[56:64], [user_pts137[48], user_pts137[22]]))[ + np.newaxis, :, :], (255, 255, 0)) + cv2.fillPoly(nohair_mask, np.concatenate((user_pts137[22:37], user_pts137[56:47:-1]))[np.newaxis, :, :], + (255, 0, 128)) + + cv2.fillPoly(nohair_mask, user_pts137[88:104][np.newaxis, :, :], (255, 0, 255)) # eye + cv2.fillPoly(nohair_mask, user_pts137[105:121][np.newaxis, :, :], (255, 0, 255)) # eye + + # Label nose + cv2.fillPoly(nohair_mask, user_pts137[64:79][np.newaxis, :, :], (255, 255, 255)) + + # Label eyebrow + cv2.fillPoly(nohair_mask, user_pts137[121:129][np.newaxis, :, :], (255, 128, 0)) + cv2.fillPoly(nohair_mask, user_pts137[129:137][np.newaxis, :, :], (255, 128, 0)) + + # 5 ********************** bald_img ********************** + bald_img = (user_bald_mask[:, :, np.newaxis] * user_bald_8uc3_512).astype(np.uint8) + + # 8 ********************** another pose hair_image ********************** + another_pose_image = ref_rgb_8uc3_512.copy() + + another_nohair_pose_mask = ref_baldseg_8uc3_512.copy() + 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_512).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_512.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_nohair_mask = torch.from_numpy( + (nohair_mask.astype(np.float32) / 255).transpose(2, 0, 1)[np.newaxis, :, :, :]).to(self.device) + # input_nohair_image = torch.from_numpy( (nohair_image.astype(np.float32) / 255).transpose(2, 0, 1)[np.newaxis, :, :, :]).to(self.device) + input_bald_img = torch.from_numpy( + (bald_img.astype(np.float32) / 255).transpose(2, 0, 1)[np.newaxis, :, :, :]).to(self.device) + input_another_pose_hair_image = torch.from_numpy( + (another_pose_hair_image.astype(np.float32) / 255).transpose(2, 0, 1)[np.newaxis, :, :, :]).to(self.device) + + zero_tensor = torch.from_numpy(np.zeros((1, 3, self.dst_height, self.dst_width), dtype=np.float32)).to( + self.device) + + with torch.no_grad(): + fake_image = self.net(input_nohair_mask, input_bald_img, input_another_pose_hair_image, zero_tensor) + fake_image_numpy = (fake_image[0].cpu().detach().numpy().transpose((1, 2, 0)) * 255).astype(np.uint8).copy() + + return fake_image_numpy + + def Generator_Hair_inference_use_pref(self, another_pose_hair_image, user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, gender="boy"): + + """ + input： + + 图像尺寸基于： 人脸 512， 图像均为3通道 + another_pose_hair_image: 参考图 处理好 numpy，float32 + user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 + + output： + + hair_gene_8uc3_512: 生成 发型图， uint8 （0-255） + + """ + + nohair_mask = user_baldseg_8uc3_768.copy() + # user_bald_mask = (user_baldseg_8uc3_768 == [0, 255, 0]).all(axis=2) + user_pts137 = landmark_processor.pts_1k_to_137(user_landmark_f1k2_768).astype(np.int32) + # 2 ********************** nohair_mask ********************** + # Label mouth + cv2.fillPoly(nohair_mask, np.concatenate((user_pts137[47:35:-1], user_pts137[56:64], [user_pts137[48], user_pts137[22]]))[ + np.newaxis, :, :], (255, 255, 0)) + cv2.fillPoly(nohair_mask, np.concatenate((user_pts137[22:37], user_pts137[56:47:-1]))[np.newaxis, :, :], + (255, 0, 128)) + + cv2.fillPoly(nohair_mask, user_pts137[88:104][np.newaxis, :, :], (255, 0, 255)) # eye + cv2.fillPoly(nohair_mask, user_pts137[105:121][np.newaxis, :, :], (255, 0, 255)) # eye + + # Label nose + cv2.fillPoly(nohair_mask, user_pts137[64:79][np.newaxis, :, :], (255, 255, 255)) + + # Label eyebrow + cv2.fillPoly(nohair_mask, user_pts137[121:129][np.newaxis, :, :], (255, 128, 0)) + cv2.fillPoly(nohair_mask, user_pts137[129:137][np.newaxis, :, :], (255, 128, 0)) + + # # 4 ********************** nohair_image ********************** + nohair_image_orig_blur = cv2.blur(user_bald_8uc3_768, (self.dst_width // 5, self.dst_width // 5)) + # bald_img = nohair_image_orig_blur.copy() + clear_mask = np.zeros((self.dst_height, self.dst_width, 3), dtype=np.uint8) + cv2.fillPoly(clear_mask, np.concatenate((user_pts137[96:88:-1], user_pts137[105:114], user_pts137[2::-1], user_pts137[21:19:-1]))[np.newaxis, :, :], (1, 1, 1)) + + bald_img = nohair_image_orig_blur * (1 - clear_mask) + user_bald_8uc3_768 * clear_mask + + # inter_res = np.concatenate((user_baldseg_8uc3_768, user_bald_8uc3_768, bald_img), axis=1) + # cv2.imshow("gen hair inter_res: ", inter_res) + # cv2.waitKey() + + # 改动 for blur version + # nohair_image_orig_blur = cv2.blur(user_bald_8uc3_768, (self.dst_width // 5, self.dst_width // 5)) + # bald_img = nohair_image_orig_blur.copy() + + input_nohair_mask = torch.from_numpy((nohair_mask.astype(np.float32) / 255).transpose(2, 0, 1)[np.newaxis, :, :, :]).to(self.device) + input_bald_img = torch.from_numpy((bald_img.astype(np.float32) / 255).transpose(2, 0, 1)[np.newaxis, :, :, :]).to(self.device) + input_another_pose_hair_image = torch.from_numpy(another_pose_hair_image.astype(np.float32).transpose(2, 0, 1)[np.newaxis, :, :, :]).to(self.device) + + zero_tensor = torch.from_numpy(np.zeros((1, 3, self.dst_height, self.dst_width), dtype=np.float32)).to(self.device) + print(self.device) + t0 = time.time() + with torch.no_grad(): + fake_image = self.net(input_nohair_mask, input_bald_img, input_another_pose_hair_image, zero_tensor) + print('fake_image gen costs:', time.time() - t0) + + fake_image_numpy = (fake_image[0].cpu().detach().numpy().transpose((1, 2, 0)) * 255).astype(np.uint8).copy() + + # hair__align_concat_show = np.concatenate((nohair_mask, bald_img, (another_pose_hair_image*255).astype(np.uint8), fake_image_numpy), axis=1) + # hair__align_concat_show = cv2.resize(hair__align_concat_show, (0, 0), fx=0.5, fy=0.5) + # cv2.imshow("hair__align_concat_show", hair__align_concat_show) + # cv2.waitKey(0) + + return fake_image_numpy + +class Generator_Fusion_Res(Process_Data): + 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.model_dir = modelRoot + self.output_img_size = 768 + + fusion_model_path = os.path.join(self.model_dir, 'hair_fusion_0427.pt') # v1: 0412 v2:zxm_v2_use_paired_data_remap_add_encode_04_12_768 + self.load_fusion_model(fusion_model_path) + + + def load_fusion_model(self, hair_fusion_model_path): + self.fusion_model = torch.jit.load(hair_fusion_model_path, map_location='cpu').to(self.device) + + + def inference_girl(self, user_generator_hair_8uc3_orisize, user_generator_matte_8uc3_orisize, + user_generator_landmark_8uc3_orisize, user_bald_mask_8uc3_orisize): + """ + input: + user_generator_hair_8uc3_orisize: 换发型结果 原图尺寸 8uc3 + user_generator_matte_8uc3_orisize: 换发型结果matting alpha图 原图尺寸 8uc3 + + output: + fusion_res + + """ + ratio = 0.5 + h_offset = 0.45 + user_pts1k = user_generator_landmark_8uc3_orisize.astype(np.int32) + user_real = user_generator_hair_8uc3_orisize + user_hair_mask = user_generator_matte_8uc3_orisize + + # user_bald_mask_8uc3_orisize = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/5c_mask_xxq/00B4D4B8-8F26-7DF9-51D0-41492D30DDF720201230_boy_20_None_0.png") + + + user_bald_mask = user_bald_mask_8uc3_orisize + user_hair_mask_orig = user_hair_mask[:, :, 0].copy() + + user_real_save = user_real.copy() + orig_h, orig_w, _ = user_real_save.shape + + M_user = landmark_processor.get_transform_mat_hair_ratio_v1(user_pts1k, self.output_img_size, ratio, + h_offset) + inv_M_user = cv2.invertAffineTransform(M_user) + user_real = cv2.warpAffine(user_real, M_user, (self.output_img_size, self.output_img_size)) + user_hair_mask = cv2.warpAffine(user_hair_mask, M_user, (self.output_img_size, self.output_img_size), + flags=cv2.INTER_CUBIC) + user_bald_mask = cv2.warpAffine(user_bald_mask, M_user, (self.output_img_size, self.output_img_size), + flags=cv2.INTER_NEAREST) + + random_v = 15 + kernel_e = np.ones((random_v, random_v), np.uint8) + kernel_d = np.ones((random_v + 10, random_v + 10), np.uint8) + user_hair_mask_erode = cv2.erode(user_hair_mask[:, :, 0], kernel_e, iterations=1) + user_hair_mask = cv2.dilate(user_hair_mask[:, :, 0], kernel_d, iterations=1) + + blend = np.clip(user_real * (user_hair_mask_erode[:, :, np.newaxis] / 255), 0, 255) + lab_img_paf = cv2.cvtColor(user_real, cv2.COLOR_BGR2LAB) + lab_img_paf[:, :, 1] = 0 + lab_img_paf[:, :, 2] = 0 + # img_paf = np.clip(lab_img_paf * (user_hair_mask[:, :, np.newaxis] / 255) + user_real * ( + # 1 - user_hair_mask[:, :, np.newaxis] / 255), 0, 255) + img_paf = user_real.copy() + img_paf[user_hair_mask > 0] = lab_img_paf[user_hair_mask > 0] + + + # img_paf = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/user_debug/fusion_test/img_paf.png") + # blend = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/user_debug/fusion_test/blend.png") + + condition = blend # np.concatenate((user_mask, user_hair_mask), axis=2) + condition = (condition.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] + input_paf = (img_paf.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] + condition = torch.from_numpy(condition).to(self.device) + input_paf = torch.from_numpy(input_paf).to(self.device) + + # ******************* forward ******************* + with torch.no_grad(): + test_fake = self.fusion_model(input_paf, condition) + test_res = (test_fake.detach()[0].to('cpu').numpy().transpose((1, 2, 0)) * 255).astype(np.uint8).copy() + + # input_out_show = np.concatenate((img_paf, user_bald_mask, (blend).astype(np.uint8), test_res), axis=1) + + + # cv2.imwrite("/media/DATA_4T/test_hair/debug_fusion/user_debug/test_new_res.png", test_res) + + test_res_orig = user_real_save.copy() + user_hair_mask_e_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) + user_hair_mask_d_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) + cv2.warpAffine(test_res, inv_M_user, (orig_w, orig_h), + dst=test_res_orig, borderMode=cv2.BORDER_TRANSPARENT) + cv2.warpAffine(user_hair_mask_erode, inv_M_user, (orig_w, orig_h), + dst=user_hair_mask_e_orig) + cv2.warpAffine(user_hair_mask, inv_M_user, (orig_w, orig_h), + dst=user_hair_mask_d_orig) + + + user_bald_mask_b_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) + user_bald_mask_b = np.zeros(user_bald_mask[:, :, 0].shape, dtype=np.uint8) + user_bald_mask_b[user_bald_mask[:, :, 0] > 0] = 1 + user_bald_mask_b[user_bald_mask[:, :, 1] > 0] = 1 + user_bald_mask_b[user_bald_mask[:, :, 2] > 0] = 1 + user_bald_mask_b = cv2.dilate(user_bald_mask_b, np.ones((25, 25))) + cv2.warpAffine(user_bald_mask_b, inv_M_user, (orig_w, orig_h), + dst=user_bald_mask_b_orig) + + # user_hair_mask_d_orig_b = cv2.blur(user_hair_mask_d_orig, (30, 30)) + user_hair_mask_d_orig_b = user_hair_mask_d_orig + user_hair_mask_e_orig_b = cv2.blur(user_hair_mask_e_orig, (15, 15)) + + # user_hair_mask_d_orig[~(user_bald_mask_b_orig.astype(np.bool))] = user_hair_mask_e_orig[~(user_bald_mask_b_orig.astype(np.bool))] + # # user_hair_mask_d_orig_b = cv2.blur(user_hair_mask_d_orig, (15, 15)) + user_hair_mask_d_orig_b[~(user_bald_mask_b_orig.astype(np.bool))] = user_hair_mask_e_orig[ + ~(user_bald_mask_b_orig.astype(np.bool))] + user_hair_mask_d_orig_b = cv2.blur(user_hair_mask_d_orig_b, (15, 15)) + + test_res_clear_orig = user_real_save * ( + 1 - user_hair_mask_d_orig_b[:, :, np.newaxis] / 255) + test_res_orig * ( + user_hair_mask_d_orig_b[:, :, np.newaxis] / 255) + test_res_clear_orig = (np.clip(test_res_clear_orig, 0, 255)).astype(np.uint8) + + test_res_orig_lab = cv2.cvtColor(test_res_clear_orig, cv2.COLOR_BGR2LAB) + user_real_save_lab = cv2.cvtColor(user_real_save, cv2.COLOR_BGR2LAB) + test_res_orig_lab[:, :, 0] = user_real_save_lab[:, :, 0] + fusion_res = cv2.cvtColor(test_res_orig_lab, cv2.COLOR_LAB2BGR) + + user_hair_mask_e_orig_b[~(user_bald_mask_b_orig.astype(np.bool))] = user_hair_mask_d_orig_b[~(user_bald_mask_b_orig.astype(np.bool))] + fusion_res = test_res_clear_orig * ( + 1 - user_hair_mask_e_orig_b[:, :, np.newaxis] / 255) + fusion_res * ( + user_hair_mask_e_orig_b[:, :, np.newaxis] / 255) + fusion_res = (np.clip(fusion_res, 0, 255)).astype(np.uint8) + + # con = np.concatenate((user_hair_mask_orig, user_hair_mask_d_orig_b), axis=1) + # con = cv2.resize(con, (0, 0), fx=1 / 4, fy=1 / 4, interpolation=cv2.INTER_CUBIC) + # + con_res = np.concatenate(( user_real_save, fusion_res), axis=1) + # con_res = cv2.resize(con_res, (0, 0), fx=1 / 4, fy=1 / 4, interpolation=cv2.INTER_CUBIC) + + return fusion_res + + def inference_boy(self, user_generator_hair_8uc3_orisize, user_generator_matte_8uc3_orisize, + user_generator_landmark_8uc3_orisize, user_bald_mask_8uc3_orisize): + """ + input: + user_generator_hair_8uc3_orisize: 换发型结果 原图尺寸 8uc3 + user_generator_matte_8uc3_orisize: 换发型结果matting alpha图 原图尺寸 8uc3 + + output: + fusion_res + + """ + ratio = 0.6 + h_offset = 0.65 + user_pts1k = user_generator_landmark_8uc3_orisize.astype(np.int32) + user_real = user_generator_hair_8uc3_orisize + user_hair_mask = user_generator_matte_8uc3_orisize + + # user_bald_mask_8uc3_orisize = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/5c_mask_xxq/00B4D4B8-8F26-7DF9-51D0-41492D30DDF720201230_boy_20_None_0.png") + + + user_bald_mask = user_bald_mask_8uc3_orisize + user_hair_mask_orig = user_hair_mask[:, :, 0].copy() + + user_real_save = user_real.copy() + orig_h, orig_w, _ = user_real_save.shape + + M_user = landmark_processor.get_transform_mat_hair_ratio_v1(user_pts1k, self.output_img_size, ratio, + h_offset) + inv_M_user = cv2.invertAffineTransform(M_user) + user_real = cv2.warpAffine(user_real, M_user, (self.output_img_size, self.output_img_size)) + user_hair_mask = cv2.warpAffine(user_hair_mask, M_user, (self.output_img_size, self.output_img_size), + flags=cv2.INTER_CUBIC) + user_bald_mask = cv2.warpAffine(user_bald_mask, M_user, (self.output_img_size, self.output_img_size), + flags=cv2.INTER_NEAREST) + + random_v = 15 + kernel_e = np.ones((random_v, random_v), np.uint8) + kernel_d = np.ones((random_v + 10, random_v + 10), np.uint8) + user_hair_mask_erode = cv2.erode(user_hair_mask[:, :, 0], kernel_e, iterations=1) + user_hair_mask = cv2.dilate(user_hair_mask[:, :, 0], kernel_d, iterations=1) + + blend = np.clip(user_real * (user_hair_mask_erode[:, :, np.newaxis] / 255), 0, 255) + lab_img_paf = cv2.cvtColor(user_real, cv2.COLOR_BGR2LAB) + lab_img_paf[:, :, 1] = 0 + lab_img_paf[:, :, 2] = 0 + # img_paf = np.clip(lab_img_paf * (user_hair_mask[:, :, np.newaxis] / 255) + user_real * ( + # 1 - user_hair_mask[:, :, np.newaxis] / 255), 0, 255) + img_paf = user_real.copy() + img_paf[user_hair_mask > 0] = lab_img_paf[user_hair_mask > 0] + + + # img_paf = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/user_debug/fusion_test/img_paf.png") + # blend = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/user_debug/fusion_test/blend.png") + + condition = blend # np.concatenate((user_mask, user_hair_mask), axis=2) + condition = (condition.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] + input_paf = (img_paf.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] + condition = torch.from_numpy(condition).to(self.device) + input_paf = torch.from_numpy(input_paf).to(self.device) + + # ******************* forward ******************* + with torch.no_grad(): + test_fake = self.fusion_model(input_paf, condition) + test_res = (test_fake.detach()[0].to('cpu').numpy().transpose((1, 2, 0)) * 255).astype(np.uint8).copy() + + # input_out_show = np.concatenate((img_paf, user_bald_mask, (blend).astype(np.uint8), test_res), axis=1) + + + test_res_orig = user_real_save.copy() + user_hair_mask_e_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) + user_hair_mask_d_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) + cv2.warpAffine(test_res, inv_M_user, (orig_w, orig_h), + dst=test_res_orig, borderMode=cv2.BORDER_TRANSPARENT) + cv2.warpAffine(user_hair_mask_erode, inv_M_user, (orig_w, orig_h), + dst=user_hair_mask_e_orig) + cv2.warpAffine(user_hair_mask, inv_M_user, (orig_w, orig_h), + dst=user_hair_mask_d_orig) + + + user_bald_mask_b_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) + user_bald_mask_b = np.zeros(user_bald_mask[:, :, 0].shape, dtype=np.uint8) + user_bald_mask_b[user_bald_mask[:, :, 0] > 0] = 1 + user_bald_mask_b[user_bald_mask[:, :, 1] > 0] = 1 + user_bald_mask_b[user_bald_mask[:, :, 2] > 0] = 1 + user_bald_mask_b = cv2.dilate(user_bald_mask_b, np.ones((10, 10))) + cv2.warpAffine(user_bald_mask_b, inv_M_user, (orig_w, orig_h), + dst=user_bald_mask_b_orig) + + user_hair_mask_d_orig_b = user_hair_mask_d_orig + user_hair_mask_d_orig_b[~(user_bald_mask_b_orig.astype(np.bool))] = user_hair_mask_e_orig[ + ~(user_bald_mask_b_orig.astype(np.bool))] + user_hair_mask_d_orig_b = cv2.blur(user_hair_mask_d_orig_b, (15, 15)) + user_hair_mask_e_orig_b = cv2.blur(user_hair_mask_e_orig, (15, 15)) + + test_res_clear_orig = user_real_save * ( + 1 - user_hair_mask_d_orig_b[:, :, np.newaxis] / 255) + test_res_orig * ( + user_hair_mask_d_orig_b[:, :, np.newaxis] / 255) + test_res_clear_orig = (np.clip(test_res_clear_orig, 0, 255)).astype(np.uint8) + + test_res_orig_lab = cv2.cvtColor(test_res_clear_orig, cv2.COLOR_BGR2LAB) + user_real_save_lab = cv2.cvtColor(user_real_save, cv2.COLOR_BGR2LAB) + test_res_orig_lab[:, :, 0] = user_real_save_lab[:, :, 0] + fusion_res = cv2.cvtColor(test_res_orig_lab, cv2.COLOR_LAB2BGR) + + user_hair_mask_e_orig_b[~(user_bald_mask_b_orig.astype(np.bool))] = user_hair_mask_d_orig_b[ + ~(user_bald_mask_b_orig.astype(np.bool))] + fusion_res = test_res_clear_orig * ( + 1 - user_hair_mask_e_orig_b[:, :, np.newaxis] / 255) + fusion_res * ( + user_hair_mask_e_orig_b[:, :, np.newaxis] / 255) + fusion_res = (np.clip(fusion_res, 0, 255)).astype(np.uint8) + + # con = np.concatenate((user_hair_mask_orig, user_hair_mask_d_orig_b), axis=1) + # con = cv2.resize(con, (0, 0), fx=1 / 4, fy=1 / 4, interpolation=cv2.INTER_CUBIC) + # + # con_res = np.concatenate((user_real_save, fusion_res), axis=1) + # con_res = cv2.resize(con_res, (0, 0), fx=1 / 4, fy=1 / 4, interpolation=cv2.INTER_CUBIC) + # + + return fusion_res + + def inference(self, user_generator_hair_8uc3_orisize, user_generator_matte_8uc3_orisize, + user_generator_landmark_8uc3_orisize, user_bald_mask_8uc3_orisize, ratio): + """ + input: + user_generator_hair_8uc3_orisize: 换发型结果 原图尺寸 8uc3 + user_generator_matte_8uc3_orisize: 换发型结果matting alpha图 原图尺寸 8uc3 + + output: + fusion_res + + """ + + user_pts1k = user_generator_landmark_8uc3_orisize.astype(np.int32) + user_real = user_generator_hair_8uc3_orisize + user_hair_mask = user_generator_matte_8uc3_orisize + + user_bald_mask = user_bald_mask_8uc3_orisize + user_hair_mask_orig = user_hair_mask[:, :, 0].copy() + + user_real_save = user_real.copy() + orig_h, orig_w, _ = user_real_save.shape + + if ratio == 0: + M_user = self.get_hair_M_boy_v1(user_pts1k) + elif ratio == 1: + M_user = self.get_hair_M_girl_v1(user_pts1k) + elif ratio == 2: + M_user = self.get_hair_M_girl_v2(user_pts1k) + else: + M_user = self.get_hair_M_girl_v1(user_pts1k) + + # M_user = landmark_processor.get_transform_mat_hair_ratio_v1(user_pts1k, self.output_img_size, ratio, + # h_offset) + inv_M_user = cv2.invertAffineTransform(M_user) + user_real = cv2.warpAffine(user_real, M_user, (self.output_img_size, self.output_img_size)) + user_hair_mask = cv2.warpAffine(user_hair_mask, M_user, (self.output_img_size, self.output_img_size), + flags=cv2.INTER_CUBIC) + user_bald_mask = cv2.warpAffine(user_bald_mask, M_user, (self.output_img_size, self.output_img_size), + flags=cv2.INTER_NEAREST) + + random_v = 15 + kernel_e = np.ones((random_v, random_v), np.uint8) + kernel_d = np.ones((random_v + 10, random_v + 10), np.uint8) + user_hair_mask_erode = cv2.erode(user_hair_mask[:, :, 0], kernel_e, iterations=1) + user_hair_mask = cv2.dilate(user_hair_mask[:, :, 0], kernel_d, iterations=1) + + blend = np.clip(user_real * (user_hair_mask_erode[:, :, np.newaxis] / 255), 0, 255) + lab_img_paf = cv2.cvtColor(user_real, cv2.COLOR_BGR2LAB) + lab_img_paf[:, :, 1] = 0 + lab_img_paf[:, :, 2] = 0 + # img_paf = np.clip(lab_img_paf * (user_hair_mask[:, :, np.newaxis] / 255) + user_real * ( + # 1 - user_hair_mask[:, :, np.newaxis] / 255), 0, 255) + img_paf = user_real.copy() + img_paf[user_hair_mask > 0] = lab_img_paf[user_hair_mask > 0] + + + + # img_paf = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/user_debug/fusion_test/img_paf.png") + # blend = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/user_debug/fusion_test/blend.png") + + condition = blend # np.concatenate((user_mask, user_hair_mask), axis=2) + condition = (condition.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] + input_paf = (img_paf.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] + condition = torch.from_numpy(condition).to(self.device) + input_paf = torch.from_numpy(input_paf).to(self.device) + + # ******************* forward ******************* + with torch.no_grad(): + test_fake = self.fusion_model(input_paf, condition) + test_res = (test_fake.detach()[0].to('cpu').numpy().transpose((1, 2, 0)) * 255).astype(np.uint8).copy() + + + + test_res_orig = user_real_save.copy() + user_hair_mask_e_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) + user_hair_mask_d_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) + cv2.warpAffine(test_res, inv_M_user, (orig_w, orig_h), + dst=test_res_orig, borderMode=cv2.BORDER_TRANSPARENT) + cv2.warpAffine(user_hair_mask_erode, inv_M_user, (orig_w, orig_h), + dst=user_hair_mask_e_orig) + cv2.warpAffine(user_hair_mask, inv_M_user, (orig_w, orig_h), + dst=user_hair_mask_d_orig) + + user_bald_mask_b_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) + user_bald_mask_b = np.zeros(user_bald_mask[:, :, 0].shape, dtype=np.uint8) + user_bald_mask_b[user_bald_mask[:, :, 0] > 0] = 1 + user_bald_mask_b[user_bald_mask[:, :, 1] > 0] = 1 + user_bald_mask_b[user_bald_mask[:, :, 2] > 0] = 1 + + if ratio != 0: + user_bald_mask_b = cv2.dilate(user_bald_mask_b, np.ones((25, 25))) # default 25 + cv2.warpAffine(user_bald_mask_b, inv_M_user, (orig_w, orig_h), + dst=user_bald_mask_b_orig) + + user_hair_mask_d_orig_b = user_hair_mask_d_orig + user_hair_mask_e_orig_b = cv2.blur(user_hair_mask_e_orig, (15, 15)) + + user_hair_mask_d_orig_b[~(user_bald_mask_b_orig.astype(np.bool))] = user_hair_mask_e_orig[ + ~(user_bald_mask_b_orig.astype(np.bool))] + user_hair_mask_d_orig_b = cv2.blur(user_hair_mask_d_orig_b, (15, 15)) + else: + user_bald_mask_b = cv2.dilate(user_bald_mask_b, np.ones((10, 10))) + cv2.warpAffine(user_bald_mask_b, inv_M_user, (orig_w, orig_h), + dst=user_bald_mask_b_orig) + + user_hair_mask_d_orig_b = user_hair_mask_d_orig + user_hair_mask_d_orig_b[~(user_bald_mask_b_orig.astype(np.bool))] = user_hair_mask_e_orig[ + ~(user_bald_mask_b_orig.astype(np.bool))] + user_hair_mask_d_orig_b = cv2.blur(user_hair_mask_d_orig_b, (15, 15)) + user_hair_mask_e_orig_b = cv2.blur(user_hair_mask_e_orig, (15, 15)) + # cv2.imshow('user_hair_mask_e_orig_b', user_hair_mask_e_orig_b) + # cv2.waitKey() + test_res_clear_orig = user_real_save * ( + 1 - user_hair_mask_d_orig_b[:, :, np.newaxis] / 255) + test_res_orig * ( + user_hair_mask_d_orig_b[:, :, np.newaxis] / 255) + test_res_clear_orig = (np.clip(test_res_clear_orig, 0, 255)).astype(np.uint8) + + test_res_orig_lab = cv2.cvtColor(test_res_clear_orig, cv2.COLOR_BGR2LAB) + user_real_save_lab = cv2.cvtColor(user_real_save, cv2.COLOR_BGR2LAB) + test_res_orig_lab[:, :, 0] = user_real_save_lab[:, :, 0] + fusion_res = cv2.cvtColor(test_res_orig_lab, cv2.COLOR_LAB2BGR) + + user_hair_mask_e_orig_b[~(user_bald_mask_b_orig.astype(np.bool))] = user_hair_mask_d_orig_b[~(user_bald_mask_b_orig.astype(np.bool))] + fusion_res = test_res_clear_orig * ( + 1 - user_hair_mask_e_orig_b[:, :, np.newaxis] / 255) + fusion_res * ( + user_hair_mask_e_orig_b[:, :, np.newaxis] / 255) + fusion_res = (np.clip(fusion_res, 0, 255)).astype(np.uint8) + return fusion_res + +class Change_Hair_Color(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.model_dir = modelRoot + + self.dst_height = 768 + self.dst_width = 768 + + generator_hair_path = os.path.join(self.model_dir, "zxm_v1_encode_refer_add_body_0107_use_gendata_from_hisd_0721.pt") #last: master_gen_hair_zxm_change_hair_color_v1_encode_refer_add_body + self.net = torch.jit.load(generator_hair_path, torch.device('cpu')).to(self.device) + self.net.eval() + + def Change_Hair_inference(self, user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768): + + """ + input： + + 图像尺寸基于： 人脸 768， 图像均为3通道 + + user_rgb_8uc3_change_color_768: 用户图， size 768， uint8 （0-255） + user_matting_8uc3_change_color_768: 用户图 matting结果 size 768， uint8 （0-255） + user_landmark_f1k2_change_color_768: 用户图 关键点 1k*2 size 768， float32 + + ref_rgb_8uc3_change_color_768: 参考图， size 768， uint8 （0-255） + ref_matting_8uc3_change_color_768: 参考图， size 768， uint8 （0-255） + ref_landmark_f1k2_change_color_768: 参考图， 关键点 1k*2 size 768， float32 + + output： + + hair_gene_color_8uc3_768: 头发换颜色生成图， size 768， uint8 （0-255） + + """ + + + + user_lab_user = cv2.cvtColor(user_rgb_8uc3_change_color_768, cv2.COLOR_BGR2LAB) + user_lab_user[:, :, 1] = 0 + user_lab_user[:, :, 2] = 0 + img_paf = user_rgb_8uc3_change_color_768.copy() + img_paf[user_matting_8uc3_change_color_768[:, :, 0].astype(np.bool)] = user_lab_user[user_matting_8uc3_change_color_768[:, :, 0].astype(np.bool)] + img_paf = (np.clip(img_paf, 0, 255)).astype(np.uint8) + + tmp = ref_matting_8uc3_change_color_768[:, :, 0] + tmp[tmp < 125] = 0 # tmp[tmp < 1] = 0 + + # mean_color = cv2.mean(refer_real, tmp) + condit_hair = (ref_rgb_8uc3_change_color_768 * (tmp[:, :, np.newaxis] / 255)).astype(np.uint8) + + + + # cv2.imwrite("/media/DATA_4T/test_hair/online_data/new_color/img_paf.png", img_paf) + # cv2.imwrite("/media/DATA_4T/test_hair/online_data/new_color/condition.png", condit_hair) + + condition = condit_hair # np.concatenate((user_mask, user_hair_mask), axis=2) + condition = (condition.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] + input_paf = (img_paf.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] + condition = torch.from_numpy(condition).to(self.device) + input_paf = torch.from_numpy(input_paf).to(self.device) + + # ******************* forward ******************* + with torch.no_grad(): + test_fake = self.net(input_paf, condition) + + hair_gene_color_8uc3_768 = (test_fake.detach()[0].to('cpu').numpy().transpose((1, 2, 0)) * 255).astype(np.uint8).copy() + + + # user_matting_mask_d = cv2.dilate(user_matting_8uc3_change_color_768, np.ones((30, 30), np.uint8)) + # user_matting_mask_d_b = cv2.blur(user_matting_mask_d, (10, 10)) + + # hair_gene_color_8uc3_768 = user_rgb_8uc3_change_color_768 * (1 - user_matting_8uc3_change_color_768 / 255) + test_res * (user_matting_8uc3_change_color_768 / 255) + # hair_gene_color_8uc3_768 = (np.clip(hair_gene_color_8uc3_768, 0, 255)).astype(np.uint8) + + return hair_gene_color_8uc3_768, user_matting_8uc3_change_color_768 + + +class GenderClassifyProcessor(object): + def __init__(self, gpu_id=0): + model_path = "./weights/gender_models" + self.output_img_size = 128 + if not os.path.exists(model_path): + print("GenderClassifyProcessor don't have model!") + + if gpu_id == 'cpu': + self.device = torch.device(gpu_id) + else: + self.device = torch.device('cuda:{0}'.format(gpu_id)) + + self.gender_model = cv2.dnn.readNetFromCaffe(os.path.join(model_path, "gender.prototxt"), os.path.join(model_path, "gender.caffemodel")) + + def forward(self, img, landmark_137): + + image_to_face_mat = landmark_processor.get_transform_mat_sex(landmark_137, self.output_img_size) + gender_img = cv2.warpAffine(img, image_to_face_mat, (self.output_img_size, self.output_img_size), cv2.INTER_LANCZOS4) + with torch.no_grad(): + inpBlob = cv2.dnn.blobFromImage(gender_img, 1.0, (self.output_img_size, self.output_img_size), (0, 0, 0), swapRB=False, + crop=False) + self.gender_model.setInput(inpBlob) + output = self.gender_model.forward() + + is_female = True + if output[0][0] > output[0][1]: + is_female = False + return is_female + +class BodySeg(): + def __init__(self, gpu_id=0): + self.gpu_id = gpu_id + self.device = torch.device(f"cuda:{gpu_id}") + + self.model = DeepLab() + model_io.load_model_by_path("./weights/human_seg_deeplabv3_288_384_sigmoid_msc.pth", self.model, gpu_id=gpu_id) + self.model.to(self.device) + self.model.eval() + self.output_img_size = [288, 384] + + self.resize_ratio = 2 + + def getM(self, center, angle, sx, sy): + angle = math.radians(angle) + alpha = math.cos(angle) + beta = math.sin(angle) + M = [[sx * alpha, sx * beta, (1 - sx * alpha) * center[0] - sx * beta * center[1]], + [-sy * beta, sy * alpha, sy * beta * center[0] + (1 - sy * alpha) * center[1]]] + return np.array(M) + def forward(self, frame): + # frame = cv2.imread(img_path).astype(np.float32) / 255 + + height, width = frame.shape[:2] + + x0, y0, x1, y1 = 0, 0, frame.shape[1], frame.shape[0] + + center_x, center_y = (x0 + x1) / 2, (y0 + y1) / 2 + random_scalex = min(self.output_img_size[0] * self.resize_ratio * 1.0 / (x1 - x0), + self.output_img_size[1] * self.resize_ratio * 1.0 / (y1 - y0)) + random_scaley = random_scalex + M = self.getM((center_x, center_y), 0, random_scalex, random_scaley) + M[:, 2] += [self.output_img_size[0] * self.resize_ratio / 2 - center_x, + self.output_img_size[1] * self.resize_ratio / 2 - center_y] + crop_img = cv2.warpAffine(frame, M, (self.output_img_size[0] * self.resize_ratio, self.output_img_size[1] * self.resize_ratio)) + crop_img_resize = cv2.resize(crop_img, (0, 0), fx=1 / self.resize_ratio, fy=1 / self.resize_ratio) + + input_tensor = torch.from_numpy(crop_img_resize.transpose([2, 0, 1])[np.newaxis]).to(self.device) + input_tensor = input_tensor.float() + with torch.no_grad(): + output_tensor, _ = self.model(input_tensor) + output_tensor = F.interpolate(output_tensor, + size=[self.output_img_size[1] * self.resize_ratio, self.output_img_size[0] * self.resize_ratio], + mode='bilinear', align_corners=True) + # out_mask = torch.max(output_tensor[:1], 1)[1].detach().cpu().numpy().squeeze().astype(np.float32) + output_tensor = torch.sigmoid(output_tensor) + out_mask = output_tensor[0, 0].detach().cpu().numpy().squeeze().astype(np.float32) + input_output_show = np.concatenate((crop_img, np.repeat(out_mask[:, :, np.newaxis], 3, axis=2)), axis=1) + # cv2.imshow("input_output_show", input_output_show) + + M_inv = cv2.invertAffineTransform(M) + # mask_rawsize = cv2.warpAffine(out_mask, M_inv, (frame.shape[1], frame.shape[0])) + mask_rawsize = cv2.warpAffine(out_mask, M_inv, (frame.shape[1], frame.shape[0]), flags=cv2.INTER_NEAREST) + + mask_rawsize = np.repeat(mask_rawsize[:, :, np.newaxis], 3, axis=2) + return mask_rawsize + + +def draw_protect_mask(sourceImage, landmark_137): + h, w, _ = sourceImage.shape + faceLists = [] + + inpaint_mask = np.zeros((h, w), dtype=np.uint8) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[129:137]), (255,)) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[121:129]), (255,)) + # cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[22:48]), (255,)) + # cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[88:104]), (255,)) + # cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[105:121]), (255,)) + for i in range(0,9): + faceLists.append(landmark_137[i]) + for i in range(14, 22): + faceLists.append(landmark_137[i]) + point_lists = np.array(faceLists) + + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(point_lists), (255,)) + + return inpaint_mask + + +def localtranslationwarpfastwithstrength(srcimg, kpt137, startx, starty, endx, endy, degree, radius): + startx = int(startx) + starty = int(starty) + endx = int(endx) + endy = int(endy) + radius = int(radius * degree/2) + strength = int(config.get('fix', 'strength')) * degree + ddradius = float(radius * radius) + mask_keep = draw_protect_mask(srcimg, kpt137[0].astype(np.int)) + # mask_keep = cv2.imread('/home/colo/Pictures/test/102.png') + # mask_keep = cv2.cvtColor(mask_keep, cv2.COLOR_BGR2GRAY) + # copyimg = np.zeros(srcimg.shape, np.uint8) + # copyimg = srcimg.copy() + maskimg = np.zeros(srcimg.shape[:2], np.uint8) + cv2.circle(maskimg, (startx, starty), math.ceil(radius), (255, 255, 255), -1) + # cv2.imshow('maskimg_before', maskimg) + # cv2.imshow('maskimg', maskimg) + # cv2.imshow('mask_keep', mask_keep) + # cv2.waitKey() + maskimg = maskimg * (1-mask_keep/255).astype(np.uint8) + + k0 = 100 / strength # 计算公式中的|m-c|^2 + ddmc_x = (endx - startx) * (endx - startx) + ddmc_y = (endy - starty) * (endy - starty) + h, w, c = srcimg.shape + mapx = np.vstack([np.arange(w).astype(np.float32).reshape(1, -1)] * h) + mapy = np.hstack([np.arange(h).astype(np.float32).reshape(-1, 1)] * w) + distance_x = (mapx - startx) * (mapx - startx) + distance_y = (mapy - starty) * (mapy - starty) + distance = distance_x + distance_y + k1 = np.sqrt(distance) + ratio_x = (ddradius - distance_x) / (ddradius - distance_x + k0 * ddmc_x) + ratio_y = (ddradius - distance_y) / (ddradius - distance_y + k0 * ddmc_y) + ratio_x = ratio_x * ratio_x + ratio_y = ratio_y * ratio_y + ux = mapx - ratio_x * (endx - startx) * (1 - k1/radius) + uy = mapy - ratio_y * (endy - starty) * (1 - k1/radius) + np.copyto(ux, mapx, where=maskimg == 0) + np.copyto(uy, mapy, where=maskimg == 0) + ux = ux.astype(np.float32) + uy = uy.astype(np.float32) + copyimg = cv2.remap(srcimg, ux, uy, interpolation=cv2.INTER_LINEAR) + return copyimg + +def localtranslationwarpfastwithstrength_v2(srcimg, startx, starty, endx, endy, radius): + startx = int(startx) + starty = int(starty) + endx = int(endx) + endy = int(endy) + strength = int(config.get('fix', 'strength')) + ddradius = float(radius * radius) + # for i in range(137): + # cv2.putText(srcimg, str(i), (int(kpt137[i][0]), int(kpt137[i][1])), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2) + + # cv2.circle(srcimg, (int(kpt137[i][0]), int(kpt137[i][1])), 1, (255, 255, 255), -1) + # cv2.imshow('srcimg', srcimg) + # cv2.imshow('mask_keep', mask_keep) + # cv2.waitKey() + maskimg = np.zeros(srcimg.shape[:2], np.uint8) + cv2.circle(maskimg, (startx, starty), math.ceil(radius), (255, 255, 255), -1) + # maskimg = maskimg * (1-mask_keep/255).astype(np.uint8) + + # mask_keep = cv2.imread('/home/colo/Pictures/test/102.png') + # mask_keep = cv2.cvtColor(mask_keep, cv2.COLOR_BGR2GRAY) + # copyimg = np.zeros(srcimg.shape, np.uint8) + # copyimg = srcimg.copy() + # maskimg = np.zeros(srcimg.shape[:2], np.uint8) + # cv2.imshow('srcimg', srcimg) + # cv2.imshow('maskimg', maskimg) + # cv2.imshow('mask_keep', mask_keep) + # cv2.waitKey() + + k0 = 100 / strength # 计算公式中的|m-c|^2 + ddmc_x = (endx - startx) * (endx - startx) + ddmc_y = (endy - starty) * (endy - starty) + h, w, c = srcimg.shape + mapx = np.vstack([np.arange(w).astype(np.float32).reshape(1, -1)] * h) + mapy = np.hstack([np.arange(h).astype(np.float32).reshape(-1, 1)] * w) + distance_x = (mapx - startx) * (mapx - startx) + distance_y = (mapy - starty) * (mapy - starty) + distance = distance_x + distance_y + k1 = np.sqrt(distance) + ratio_x = (ddradius - distance_x) / (ddradius - distance_x + k0 * ddmc_x) + ratio_y = (ddradius - distance_y) / (ddradius - distance_y + k0 * ddmc_y) + ratio_x = ratio_x * ratio_x + ratio_y = ratio_y * ratio_y + ux = mapx - ratio_x * (endx - startx) * (1 - k1/radius) + uy = mapy - ratio_y * (endy - starty) * (1 - k1/radius) + np.copyto(ux, mapx, where=maskimg == 0) + np.copyto(uy, mapy, where=maskimg == 0) + ux = ux.astype(np.float32) + uy = uy.astype(np.float32) + copyimg = cv2.remap(srcimg, ux, uy, interpolation=cv2.INTER_LINEAR) + return copyimg + +def localtranslationwarpfastwithstrength_v2_old(srcimg, startx, starty, endx, endy, radius): + startx = int(startx) + starty = int(starty) + endx = int(endx) + endy = int(endy) + strength = int(config.get('fix', 'strength')) + ddradius = float(radius * radius) + # for i in range(137): + # cv2.putText(srcimg, str(i), (int(kpt137[i][0]), int(kpt137[i][1])), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2) + + # cv2.circle(srcimg, (int(kpt137[i][0]), int(kpt137[i][1])), 1, (255, 255, 255), -1) + # cv2.imshow('srcimg', srcimg) + # cv2.imshow('mask_keep', mask_keep) + # cv2.waitKey() + maskimg = np.zeros(srcimg.shape[:2], np.uint8) + cv2.circle(maskimg, (startx, starty), math.ceil(radius), (255, 255, 255), -1) + # maskimg = maskimg * (1-mask_keep/255).astype(np.uint8) + + # mask_keep = cv2.imread('/home/colo/Pictures/test/102.png') + # mask_keep = cv2.cvtColor(mask_keep, cv2.COLOR_BGR2GRAY) + # copyimg = np.zeros(srcimg.shape, np.uint8) + # copyimg = srcimg.copy() + # maskimg = np.zeros(srcimg.shape[:2], np.uint8) + # cv2.imshow('srcimg', srcimg) + # cv2.imshow('maskimg', maskimg) + # cv2.imshow('mask_keep', mask_keep) + # cv2.waitKey() + + k0 = 100 / strength # 计算公式中的|m-c|^2 + ddmc_x = (endx - startx) * (endx - startx) + ddmc_y = (endy - starty) * (endy - starty) + h, w, c = srcimg.shape + mapx = np.vstack([np.arange(w).astype(np.float32).reshape(1, -1)] * h) + mapy = np.hstack([np.arange(h).astype(np.float32).reshape(-1, 1)] * w) + distance_x = (mapx - startx) * (mapx - startx) + distance_y = (mapy - starty) * (mapy - starty) + distance = distance_x + distance_y + k1 = np.sqrt(distance) + ratio_x = (ddradius - distance_x) / (ddradius - distance_x + k0 * ddmc_x) + ratio_y = (ddradius - distance_y) / (ddradius - distance_y + k0 * ddmc_y) + ratio_x = ratio_x * ratio_x + ratio_y = ratio_y * ratio_y + ux = mapx - ratio_x * (endx - startx) * (1 - k1/radius) + uy = mapy - ratio_y * (endy - starty) * (1 - k1/radius) + np.copyto(ux, mapx, where=maskimg == 0) + np.copyto(uy, mapy, where=maskimg == 0) + ux = ux.astype(np.float32) + uy = uy.astype(np.float32) + copyimg = cv2.remap(srcimg, ux, uy, interpolation=cv2.INTER_LINEAR) + return copyimg + +def localtranslationwarpfastwithstrength_v2_soft(srcimg, startx, starty, endx, endy, radius): + startx = int(startx) + starty = int(starty) + endx = int(endx) + endy = int(endy) + strength = int(config.get('fix', 'strength')) + strength = int(config.get('fix', 'strength')) + ddradius = float(radius * radius) + # for i in range(137): + # cv2.putText(srcimg, str(i), (int(kpt137[i][0]), int(kpt137[i][1])), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2) + + # cv2.circle(srcimg, (int(kpt137[i][0]), int(kpt137[i][1])), 1, (255, 255, 255), -1) + # cv2.imshow('srcimg', srcimg) + # cv2.imshow('mask_keep', mask_keep) + # cv2.waitKey() + maskimg = np.zeros(srcimg.shape[:2], np.uint8) + cv2.circle(maskimg, (startx, starty), math.ceil(radius), (255, 255, 255), -1) + # maskimg = maskimg * (1-mask_keep/255).astype(np.uint8) + + # mask_keep = cv2.imread('/home/colo/Pictures/test/102.png') + # mask_keep = cv2.cvtColor(mask_keep, cv2.COLOR_BGR2GRAY) + # copyimg = np.zeros(srcimg.shape, np.uint8) + # copyimg = srcimg.copy() + # maskimg = np.zeros(srcimg.shape[:2], np.uint8) + # cv2.imshow('srcimg', srcimg) + # cv2.imshow('maskimg', maskimg) + # cv2.imshow('mask_keep', mask_keep) + # cv2.waitKey() + + k0 = 100 / strength # 计算公式中的|m-c|^2 + ddmc_x = (endx - startx) * (endx - startx) + ddmc_y = (endy - starty) * (endy - starty) + h, w, c = srcimg.shape + mapx = np.vstack([np.arange(w).astype(np.float32).reshape(1, -1)] * h) + mapy = np.hstack([np.arange(h).astype(np.float32).reshape(-1, 1)] * w) + distance_x = (mapx - startx) * (mapx - startx) + distance_y = (mapy - starty) * (mapy - starty) + distance = distance_x + distance_y + k1 = np.sqrt(distance) + ratio_x = (ddradius - distance_x) / (ddradius - distance_x + k0 * ddmc_x) + ratio_y = (ddradius - distance_y) / (ddradius - distance_y + k0 * ddmc_y) + ratio_x = ratio_x * ratio_x + ratio_y = ratio_y * ratio_y + f = (np.sin((k1 / radius - 0.5) * np.pi) + 1) / 2. + ux = mapx - ratio_x * (endx - startx) * (1 - f) + uy = mapy - ratio_y * (endy - starty) * (1 - f) + np.copyto(ux, mapx, where=maskimg == 0) + np.copyto(uy, mapy, where=maskimg == 0) + ux = ux.astype(np.float32) + uy = uy.astype(np.float32) + copyimg = cv2.remap(srcimg, ux, uy, interpolation=cv2.INTER_LINEAR) + return copyimg + +def updateEndPosition(startx, starty, endx, endy, radius): + maxMargin = radius / 2. # need to calculate + # maxMargin = radius * 30. # need to calculate + + steps = [] + tmp_scale = 0.2 + disAll = np.sqrt((endx - startx) ** 2 + (endy - starty) ** 2) + + # mini_ratio = 0.1 + # if disAll < mini_ratio * radius: + # endx = startx + (endx - startx) * radius * mini_ratio / disAll + # endy = starty + (endy - starty) * radius * mini_ratio * tmp_scale / disAll + # disAll = radius * mini_ratio ### ycj + + if disAll > (radius * 1.0 * tmp_scale): + endx = startx + (endx - startx) * radius * 0.98 * tmp_scale / disAll + endy = starty + (endy - starty) * radius * 0.98 * tmp_scale / disAll + disAll = radius * 1.0 * tmp_scale ### ycj + + curStartX = int(startx) + curStartY = int(starty) + curEndX = int(endx) + curEndY = int(endy) + + while True: + dis = np.sqrt((endx - curStartX)**2 + (endy - curStartY)**2) + if dis < maxMargin: + # no need to clip max distance. + steps.append([curStartX, curStartY, min(curEndX, endx), min(curEndY, endy), radius]) + return steps + else: + curEndX = int(curStartX + (endx - startx) / disAll * maxMargin) + curEndY = int(curStartY + (endy - starty) / disAll * maxMargin) + + steps.append([curStartX, curStartY, curEndX, curEndY, radius]) + curStartX = int(curStartX + (endx - startx) / disAll * maxMargin) + curStartY = int(curStartY + (endy - starty) / disAll * maxMargin) + curEndX = int(curEndX + (endx - startx) / disAll * maxMargin) + curEndY = int(curEndY + (endy - starty) / disAll * maxMargin) + +def updateEndPosition_old(startx, starty, endx, endy, radius): + maxMargin = radius / 3. # need to calculate + # maxMargin = radius * 30. # need to calculate + + steps = [] + curStartX = int(startx) + curStartY = int(starty) + curEndX = int(endx) + curEndY = int(endy) + + disAll = np.sqrt((endx - startx) ** 2 + (endy - starty) ** 2) + + while True: + dis = np.sqrt((endx - curStartX)**2 + (endy - curStartY)**2) + if dis < maxMargin: + # no need to clip max distance. + steps.append([curStartX, curStartY, min(curEndX, endx), min(curEndY, endy), radius]) + return steps + else: + curEndX = int(curStartX + (endx - startx) / disAll * maxMargin) + curEndY = int(curStartY + (endy - starty) / disAll * maxMargin) + + steps.append([curStartX, curStartY, curEndX, curEndY, radius]) + curStartX = int(curStartX + (endx - startx) / disAll * maxMargin) + curStartY = int(curStartY + (endy - starty) / disAll * maxMargin) + curEndX = int(curEndX + (endx - startx) / disAll * maxMargin) + curEndY = int(curEndY + (endy - starty) / disAll * maxMargin) + +class chinClass(): + def __init__(self, gpuid): + self.model = resnet18(n_class=3) + model_io.load_model_by_path('./weights/classify_3_2499.pth',self.model, gpu_id=gpuid) + self.device = torch.device("cuda:%d" % gpuid) + self.model.to(self.device) + self.model.eval() + + self.label_map = { + 0:"ping", + 1:"yuan", + 2:"jian" + } + + def getM(self, center, angle, sx, sy): + angle = math.radians(angle) + alpha = math.cos(angle) + beta = math.sin(angle) + M = [[sx * alpha, sx * beta, (1 - sx * alpha) * center[0] - sx * beta * center[1]], + [-sy * beta, sy * alpha, sy * beta * center[0] + (1 - sy * alpha) * center[1]]] + return np.array(M) + + def process_face(self, in_frame, pts137_tmp_in_source): + crop_size = 256 + output_img_size = [224, 224] + image_to_face_mat = landmark_processor.get_transform_mat_mmcv_bigger(pts137_tmp_in_source, crop_size) + face_image = cv2.warpAffine(in_frame, image_to_face_mat, (crop_size, crop_size), cv2.INTER_LANCZOS4) + crop_face = face_image[128:, :] + Image = crop_face.copy() + center_x, center_y = Image.shape[1] / 2, Image.shape[0] / 2 + random_offset = (np.random.random([2]).astype(np.float32) - 0.5) * 0 + random_rotate = 0 + random_scalex = 1 + random_scaley = random_scalex + M = self.getM((center_x, center_y), random_rotate, random_scalex, random_scaley) + M[:, 2] += [output_img_size[0] / 2 - center_x, output_img_size[1] / 2 - center_y] + random_offset + Image = cv2.warpAffine(Image, M, output_img_size) + input_data = torch.from_numpy((Image.astype(np.float32) / 255.).transpose([2, 0, 1])).unsqueeze(0).to(self.device) + with torch.no_grad(): + res = self.model(input_data) + pred_res = int(torch.argmax(res[0]).detach().cpu().numpy()) + return self.label_map[pred_res] \ No newline at end of file diff --git a/hair_service_sd/core/seg/__init__.py b/hair_service_sd/core/seg/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hair_service_sd/core/seg/hairseg_single_model.py b/hair_service_sd/core/seg/hairseg_single_model.py new file mode 100644 index 0000000..51086bf --- /dev/null +++ b/hair_service_sd/core/seg/hairseg_single_model.py @@ -0,0 +1,86 @@ +import os +import torch + +from core.seg.networks.deeplabv3_plus import get_deeplabv3_plus +import numpy as np +import cv2 +from core.utils import landmark_processor + + +def label_to_mask(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 + +label_map = [ + [0, 0, 0], # + [128, 128, 128], + [255, 255, 255], +] + +class Evaluator(object): + def __init__(self, gpu_id, output_img_size, nclass, seg_model_path=None): + self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu') + # print("gpu_id: ", gpu_id) + + # create network + self.model = get_deeplabv3_plus(backbone='xception', nclass=nclass) + model_path = os.path.join(seg_model_path) + self.model.load_state_dict(torch.load(model_path, map_location=lambda storage, loc: storage)) + # print("seg device: ", self.model.device) + self.model.to(self.device) + self.model.eval() + + # images = torch.randn((1, 3, 512, 512)).to(self.device) + # torch.onnx.export(self.model, images, + # "deeplabv3_hair512_360_0520_wl.onnx", + # verbose=True, + # opset_version=11, + # input_names=['data'], + # do_constant_folding=True, + # output_names=['output']) + + # exit() + + self.output_img_size = output_img_size + self.nclass = nclass + + + def process_data(self, img): + img = (img.astype(np.float32) / 255).transpose((2, 0, 1)) + img = torch.from_numpy(img).unsqueeze(0) + + return img + + def eval(self, img, pts1k): + orig_h, orig_w, _ = img.shape + + M1 = landmark_processor.get_transform_mat_hair(pts1k, self.output_img_size, ratio=0.3) + crop_img = cv2.warpAffine(img, M1, (self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4) + crop_img = self.process_data(crop_img) + crop_img = crop_img.to(self.device) + + with torch.no_grad(): + # torch.cuda.synchronize() + outputs = self.model(crop_img) + pred = torch.argmax(outputs[0], 1) + + pred = pred[0].detach().cpu().numpy() + predict = pred.astype(np.float32) + + pred_mask = label_to_mask(predict) + + M1_invert = cv2.invertAffineTransform(M1) + img_pred = cv2.warpAffine(pred_mask, M1_invert, (orig_w, orig_h), flags=cv2.INTER_CUBIC) #flags=cv2.INTER_NEAREST + orig_mask = img_pred.copy() + + # show_concat = np.concatenate((img, orig_mask), axis=1) + # cv2.imshow("show_concat", show_concat) + # cv2.waitKey() + return orig_mask + + + diff --git a/hair_service_sd/core/seg/networks/__init__.py b/hair_service_sd/core/seg/networks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hair_service_sd/core/seg/networks/basic.py b/hair_service_sd/core/seg/networks/basic.py new file mode 100644 index 0000000..b241fad --- /dev/null +++ b/hair_service_sd/core/seg/networks/basic.py @@ -0,0 +1,462 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +__all__ = ['_ConvBNReLU', '_DWConvBNReLU', 'InvertedResidual', '_ASPP', '_FCNHead', + '_Hswish', '_ConvBNHswish', 'SEModule', 'Bottleneck', 'ShuffleNetUnit', + 'ShuffleNetV2Unit', 'InvertedIGCV3', 'MBConvBlock'] + + +class _ConvBNReLU(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, + dilation=1, groups=1, relu6=False, norm_layer=nn.BatchNorm2d, **kwargs): + super(_ConvBNReLU, self).__init__() + self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias=False) + self.bn = norm_layer(out_channels) + self.relu = nn.ReLU6(True) if relu6 else nn.ReLU(True) + + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + x = self.relu(x) + return x + + +class _FCNHead(nn.Module): + def __init__(self, in_channels, channels, norm_layer=nn.BatchNorm2d, norm_kwargs=None, **kwargs): + super(_FCNHead, self).__init__() + inter_channels = in_channels // 4 + self.block = nn.Sequential( + nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False), + norm_layer(inter_channels, **({} if norm_kwargs is None else norm_kwargs)), + nn.ReLU(True), + nn.Dropout(0.1), + nn.Conv2d(inter_channels, channels, 1) + ) + + def forward(self, x): + return self.block(x) + + +# ----------------------------------------------------------------- +# For MobileNet +# ----------------------------------------------------------------- +class _DWConvBNReLU(nn.Module): + """Depthwise Separable Convolution in MobileNet. + depthwise convolution + pointwise convolution + """ + + def __init__(self, in_channels, dw_channels, out_channels, stride, dilation=1, norm_layer=nn.BatchNorm2d, **kwargs): + super(_DWConvBNReLU, self).__init__() + self.conv = nn.Sequential( + _ConvBNReLU(in_channels, dw_channels, 3, stride, dilation, dilation, in_channels, norm_layer=norm_layer), + _ConvBNReLU(dw_channels, out_channels, 1, norm_layer=norm_layer)) + + def forward(self, x): + return self.conv(x) + + +# ----------------------------------------------------------------- +# For MobileNetV2 +# ----------------------------------------------------------------- +class InvertedResidual(nn.Module): + def __init__(self, in_channels, out_channels, stride, expand_ratio, + dilation=1, norm_layer=nn.BatchNorm2d, **kwargs): + super(InvertedResidual, self).__init__() + assert stride in [1, 2] + self.use_res_connect = stride == 1 and in_channels == out_channels + + layers = list() + inter_channels = int(round(in_channels * expand_ratio)) + if expand_ratio != 1: + # pw + layers.append(_ConvBNReLU(in_channels, inter_channels, 1, relu6=True, norm_layer=norm_layer)) + layers.extend([ + # dw + _ConvBNReLU(inter_channels, inter_channels, 3, stride, dilation, dilation, + groups=inter_channels, relu6=True, norm_layer=norm_layer), + # pw-linear + nn.Conv2d(inter_channels, out_channels, 1, bias=False), + norm_layer(out_channels)]) + self.conv = nn.Sequential(*layers) + + def forward(self, x): + if self.use_res_connect: + return x + self.conv(x) + else: + return self.conv(x) + + +# ----------------------------------------------------------------- +# ASPP: For MobileNetV2 +# ----------------------------------------------------------------- +class _AsppPooling(nn.Module): + def __init__(self, in_channels, out_channels, norm_layer, **kwargs): + super(_AsppPooling, self).__init__() + self.gap = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(in_channels, out_channels, 1, bias=False), + norm_layer(out_channels), + nn.ReLU(True) + ) + + def forward(self, x): +# size = x.size()[2:] + size = (48, 48) +# print("size: ", size) + pool = self.gap(x) +# out = F.interpolate(pool, size, mode='bilinear', align_corners=True) + out = F.interpolate(pool, size, mode='nearest') + return out + + +class _ASPP(nn.Module): + def __init__(self, in_channels, atrous_rates, norm_layer=nn.BatchNorm2d, **kwargs): + super(_ASPP, self).__init__() + out_channels = 256 + self.b0 = nn.Sequential( + nn.Conv2d(in_channels, out_channels, 1, bias=False), + norm_layer(out_channels), + nn.ReLU(True) + ) + + rate1, rate2, rate3 = tuple(atrous_rates) + self.b1 = _ConvBNReLU(in_channels, out_channels, 3, padding=rate1, dilation=rate1, norm_layer=norm_layer) + self.b2 = _ConvBNReLU(in_channels, out_channels, 3, padding=rate2, dilation=rate2, norm_layer=norm_layer) + self.b3 = _ConvBNReLU(in_channels, out_channels, 3, padding=rate3, dilation=rate3, norm_layer=norm_layer) + self.b4 = _AsppPooling(in_channels, out_channels, norm_layer=norm_layer) + + self.project = nn.Sequential( + nn.Conv2d(5 * out_channels, out_channels, 1, bias=False), + norm_layer(out_channels), + nn.ReLU(True), + nn.Dropout2d(0.5) + ) + + def forward(self, x): + feat1 = self.b0(x) + feat2 = self.b1(x) + feat3 = self.b2(x) + feat4 = self.b3(x) + feat5 = self.b4(x) + x = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1) + x = self.project(x) + return x + + +# ----------------------------------------------------------------- +# For MobileNetV3 +# ----------------------------------------------------------------- +class _Hswish(nn.Module): + def __init__(self, inplace=True): + super(_Hswish, self).__init__() + self.relu6 = nn.ReLU6(inplace) + + def forward(self, x): + return x * self.relu6(x + 3.) / 6. + + +class _Hsigmoid(nn.Module): + def __init__(self, inplace=True): + super(_Hsigmoid, self).__init__() + self.relu6 = nn.ReLU6(inplace) + + def forward(self, x): + return self.relu6(x + 3.) / 6. + + +class _ConvBNHswish(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, + dilation=1, groups=1, norm_layer=nn.BatchNorm2d, **kwargs): + super(_ConvBNHswish, self).__init__() + self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias=False) + self.bn = norm_layer(out_channels) + self.act = _Hswish(True) + + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + x = self.act(x) + return x + + +class SEModule(nn.Module): + def __init__(self, in_channels, reduction=4): + super(SEModule, self).__init__() + self.avg_pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(in_channels, in_channels // reduction, bias=False), + nn.ReLU(True), + nn.Linear(in_channels // reduction, in_channels, bias=False), + _Hsigmoid(True) + ) + + def forward(self, x): + n, c, _, _ = x.size() + out = self.avg_pool(x).view(n, c) + out = self.fc(out).view(n, c, 1, 1) + return x * out.expand_as(x) + + +class Identity(nn.Module): + def __init__(self, in_channels): + super(Identity, self).__init__() + + def forward(self, x): + return x + + +class Bottleneck(nn.Module): + def __init__(self, in_channels, out_channels, exp_size, kernel_size, stride, dilation=1, se=False, nl='RE', + norm_layer=nn.BatchNorm2d, **kwargs): + super(Bottleneck, self).__init__() + assert stride in [1, 2] + self.use_res_connect = stride == 1 and in_channels == out_channels + if nl == 'HS': + act = _Hswish + else: + act = nn.ReLU + if se: + SELayer = SEModule + else: + SELayer = Identity + + self.conv = nn.Sequential( + # pw + nn.Conv2d(in_channels, exp_size, 1, bias=False), + norm_layer(exp_size), + act(True), + # dw + nn.Conv2d(exp_size, exp_size, kernel_size, stride, (kernel_size - 1) // 2 * dilation, + dilation, groups=exp_size, bias=False), + norm_layer(exp_size), + SELayer(exp_size), + act(True), + # pw-linear + nn.Conv2d(exp_size, out_channels, 1, bias=False), + norm_layer(out_channels) + ) + + def forward(self, x): + if self.use_res_connect: + return x + self.conv(x) + else: + return self.conv(x) + + +# ----------------------------------------------------------------- +# For ShuffleNet +# ----------------------------------------------------------------- +def channel_shuffle(x, groups): + n, c, h, w = x.size() + + channels_per_group = c // groups + x = x.view(n, groups, channels_per_group, h, w) + x = torch.transpose(x, 1, 2).contiguous() + x = x.view(n, -1, h, w) + + return x + + +class ShuffleNetUnit(nn.Module): + def __init__(self, in_channels, out_channels, stride, groups, dilation=1, norm_layer=nn.BatchNorm2d, **kwargs): + super(ShuffleNetUnit, self).__init__() + self.stride = stride + self.groups = groups + self.dilation = dilation + assert stride in [1, 2, 3] + + inter_channels = out_channels // 4 + + if stride > 1: + self.shortcut = nn.AvgPool2d(3, stride, 1) + out_channels -= in_channels + elif dilation > 1: + out_channels -= in_channels + + g = 1 if in_channels == 24 else groups + self.conv1 = _ConvBNReLU(in_channels, inter_channels, 1, groups=g, norm_layer=norm_layer) + self.conv2 = _ConvBNReLU(inter_channels, inter_channels, 3, stride, dilation, + dilation, groups, norm_layer=norm_layer) + self.conv3 = nn.Sequential( + nn.Conv2d(inter_channels, out_channels, 1, groups=groups, bias=False), + norm_layer(out_channels)) + + def forward(self, x): + out = self.conv1(x) + out = channel_shuffle(out, self.groups) + out = self.conv2(out) + out = self.conv3(out) + if self.stride > 1: + x = self.shortcut(x) + out = torch.cat([out, x], dim=1) + elif self.dilation > 1: + out = torch.cat([out, x], dim=1) + else: + out = out + x + out = F.relu(out) + + return out + + +# ----------------------------------------------------------------- +# For ShuffleNetV2 +# ----------------------------------------------------------------- +class _DWConv(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, bias=False): + super(_DWConv, self).__init__() + self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, + padding, dilation, groups=in_channels, bias=bias) + + def forward(self, x): + return self.conv(x) + + +class ShuffleNetV2Unit(nn.Module): + def __init__(self, in_channels, out_channels, stride, dilation=1, norm_layer=nn.BatchNorm2d, **kwargs): + super(ShuffleNetV2Unit, self).__init__() + assert stride in [1, 2, 3] + self.stride = stride + self.dilation = dilation + + inter_channels = out_channels // 2 + + if (stride > 1) or (dilation > 1): + self.branch1 = nn.Sequential( + _DWConv(in_channels, in_channels, 3, stride, dilation, dilation), + norm_layer(in_channels), + _ConvBNReLU(in_channels, inter_channels, 1, norm_layer=norm_layer)) + self.branch2 = nn.Sequential( + _ConvBNReLU(in_channels if (stride > 1) else inter_channels, inter_channels, 1, norm_layer=norm_layer), + _DWConv(inter_channels, inter_channels, 3, stride, dilation, dilation), + norm_layer(inter_channels), + _ConvBNReLU(inter_channels, inter_channels, 1, norm_layer=norm_layer)) + + def forward(self, x): + if (self.stride == 1) and (self.dilation == 1): + x1, x2 = x.chunk(2, dim=1) + out = torch.cat((x1, self.branch2(x2)), dim=1) + else: + out = torch.cat((self.branch1(x), self.branch2(x)), dim=1) + out = channel_shuffle(out, 2) + + return out + + +# ----------------------------------------------------------------- +# For IGCV3 +# ----------------------------------------------------------------- +class PermutationBlock(nn.Module): + def __init__(self, groups): + super(PermutationBlock, self).__init__() + self.groups = groups + + def forward(self, x): + n, c, h, w = x.size() + x = x.view(n, self.groups, c // self.groups, h, w).permute(0, 2, 1, 3, 4).contiguous().view(n, c, h, w) + return x + + +class InvertedIGCV3(nn.Module): + def __init__(self, in_channels, out_channels, stride, expand_ratio, + dilation=1, norm_layer=nn.BatchNorm2d, **kwargs): + super(InvertedIGCV3, self).__init__() + assert stride in [1, 2] + self.use_res_connect = stride == 1 and in_channels == out_channels + + layers = list() + inter_channels = int(round(in_channels * expand_ratio)) + if expand_ratio != 1: + # pw + layers.append(_ConvBNReLU(in_channels, inter_channels, 1, + groups=2, relu6=True, norm_layer=norm_layer)) + # permutation + layers.append(PermutationBlock(groups=2)) + layers.extend([ + # dw + _ConvBNReLU(inter_channels, inter_channels, 3, stride, dilation, dilation, + groups=inter_channels, relu6=True, norm_layer=norm_layer), + # pw-linear + nn.Conv2d(inter_channels, out_channels, 1, groups=2, bias=False), + norm_layer(out_channels), + # permutation + PermutationBlock(groups=int(round(out_channels / 2))) + ]) + self.conv = nn.Sequential(*layers) + + def forward(self, x): + if self.use_res_connect: + return x + self.conv(x) + else: + return self.conv(x) + + +# ----------------------------------------------------------------- +# For EfficientNet +# ----------------------------------------------------------------- +class _Swish(nn.Module): + def __init__(self): + super(_Swish, self).__init__() + self.sigmoid = nn.Sigmoid() + + def forward(self, x): + return x * self.sigmoid(x) + + +class SEModuleV2(nn.Module): + def __init__(self, in_channels, se_ratio=0.25): + super(SEModuleV2, self).__init__() + self.avg_pool = nn.AdaptiveAvgPool2d(1) + se_channels = max(1, int(in_channels * se_ratio)) + self.fc = nn.Sequential( + nn.Conv2d(in_channels, se_channels, 1, bias=False), + _Swish(), + nn.Conv2d(se_channels, in_channels, 1, bias=False), + nn.Sigmoid() + ) + + def forward(self, x): + n, c, _, _ = x.size() + out = self.avg_pool(x) + out = self.fc(out) + return x * out.expand_as(x) + + +class MBConvBlock(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride, expand_ratio, + dilation=1, se_ratio=0.25, drop_connect_rate=0.2, norm_layer=nn.BatchNorm2d, **kwargs): + super(MBConvBlock, self).__init__() + assert stride in [1, 2] + self.use_res_connect = stride == 1 and in_channels == out_channels + self.drop_connect_rate = drop_connect_rate + use_se = (se_ratio is not None) and (0 < se_ratio <= 1.) + if use_se: + SELayer = SEModuleV2 + else: + SELayer = Identity + + layers = list() + inter_channels = int(round(in_channels * expand_ratio)) + if expand_ratio != 1: + layers.append(_ConvBNHswish(in_channels, inter_channels, 1, norm_layer=norm_layer)) + layers.extend([ + # dw + _ConvBNHswish(inter_channels, inter_channels, kernel_size, stride, kernel_size // 2 * dilation, dilation, + groups=inter_channels, norm_layer=norm_layer), # check act function + SELayer(inter_channels, se_ratio), + # pw-linear + nn.Conv2d(inter_channels, out_channels, 1, bias=False), + norm_layer(out_channels) + ]) + self.conv = nn.Sequential(*layers) + + if drop_connect_rate: + self.dropout = nn.Dropout2d(drop_connect_rate) + + def forward(self, x): + out = self.conv(x) + if self.use_res_connect: + if self.drop_connect_rate: + out = self.dropout(out) + out = x + out + return out diff --git a/hair_service_sd/core/seg/networks/deeplabv3.py b/hair_service_sd/core/seg/networks/deeplabv3.py new file mode 100644 index 0000000..559a6df --- /dev/null +++ b/hair_service_sd/core/seg/networks/deeplabv3.py @@ -0,0 +1,187 @@ +"""Pyramid Scene Parsing Network""" +import torch +import torch.nn as nn +import torch.nn.functional as F + +from core.seg.networks.segbase import SegBaseModel +from core.seg.networks.fcn import _FCNHead + +__all__ = ['DeepLabV3', 'get_deeplabv3', 'get_deeplabv3_resnet50_voc', 'get_deeplabv3_resnet101_voc', + 'get_deeplabv3_resnet152_voc', 'get_deeplabv3_resnet50_ade', 'get_deeplabv3_resnet101_ade', + 'get_deeplabv3_resnet152_ade'] + + +class DeepLabV3(SegBaseModel): + r"""DeepLabV3 + + Parameters + ---------- + nclass : int + Number of categories for the training dataset. + backbone : string + Pre-trained dilated backbone network type (default:'resnet50'; 'resnet50', + 'resnet101' or 'resnet152'). + norm_layer : object + Normalization layer used in backbone network (default: :class:`nn.BatchNorm`; + for Synchronized Cross-GPU BachNormalization). + aux : bool + Auxiliary loss. + + Reference: + Chen, Liang-Chieh, et al. "Rethinking atrous convolution for semantic image segmentation." + arXiv preprint arXiv:1706.05587 (2017). + """ + + def __init__(self, nclass, backbone='resnet50', aux=False, pretrained_base=True, **kwargs): + super(DeepLabV3, self).__init__(nclass, aux, backbone, pretrained_base=pretrained_base, **kwargs) + self.head = _DeepLabHead(nclass, **kwargs) + if self.aux: + self.auxlayer = _FCNHead(1024, nclass, **kwargs) + + self.__setattr__('exclusive', ['head', 'auxlayer'] if aux else ['head']) + + def forward(self, x): + size = x.size()[2:] + _, _, c3, c4 = self.base_forward(x) + outputs = [] + x = self.head(c4) + x = F.interpolate(x, size, mode='bilinear', align_corners=True) + outputs.append(x) + + if self.aux: + auxout = self.auxlayer(c3) + auxout = F.interpolate(auxout, size, mode='bilinear', align_corners=True) + outputs.append(auxout) + return tuple(outputs) + + +class _DeepLabHead(nn.Module): + def __init__(self, nclass, norm_layer=nn.BatchNorm2d, norm_kwargs=None, **kwargs): + super(_DeepLabHead, self).__init__() + self.aspp = _ASPP(2048, [12, 24, 36], norm_layer=norm_layer, norm_kwargs=norm_kwargs, **kwargs) + self.block = nn.Sequential( + nn.Conv2d(256, 256, 3, padding=1, bias=False), + norm_layer(256, **({} if norm_kwargs is None else norm_kwargs)), + nn.ReLU(True), + nn.Dropout(0.1), + nn.Conv2d(256, nclass, 1) + ) + + def forward(self, x): + x = self.aspp(x) + return self.block(x) + + +class _ASPPConv(nn.Module): + def __init__(self, in_channels, out_channels, atrous_rate, norm_layer, norm_kwargs): + super(_ASPPConv, self).__init__() + self.block = nn.Sequential( + nn.Conv2d(in_channels, out_channels, 3, padding=atrous_rate, dilation=atrous_rate, bias=False), + norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)), + nn.ReLU(True) + ) + + def forward(self, x): + return self.block(x) + + +class _AsppPooling(nn.Module): + def __init__(self, in_channels, out_channels, norm_layer, norm_kwargs, **kwargs): + super(_AsppPooling, self).__init__() + self.gap = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(in_channels, out_channels, 1, bias=False), + norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)), + nn.ReLU(True) + ) + + def forward(self, x): + size = x.size()[2:] +# print("before gap: ", x.size()) + pool = self.gap(x) + out = F.interpolate(pool, size, mode='bilinear', align_corners=True) + return out + + +class _ASPP(nn.Module): + def __init__(self, in_channels, atrous_rates, norm_layer, norm_kwargs=None, **kwargs): + super(_ASPP, self).__init__() + out_channels = 256 + self.b0 = nn.Sequential( + nn.Conv2d(in_channels, out_channels, 1, bias=False), + norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)), + nn.ReLU(True) + ) + + rate1, rate2, rate3 = tuple(atrous_rates) + self.b1 = _ASPPConv(in_channels, out_channels, rate1, norm_layer, norm_kwargs) + self.b2 = _ASPPConv(in_channels, out_channels, rate2, norm_layer, norm_kwargs) + self.b3 = _ASPPConv(in_channels, out_channels, rate3, norm_layer, norm_kwargs) + self.b4 = _AsppPooling(in_channels, out_channels, norm_layer=norm_layer, norm_kwargs=norm_kwargs) + + self.project = nn.Sequential( + nn.Conv2d(5 * out_channels, out_channels, 1, bias=False), + norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)), + nn.ReLU(True), + nn.Dropout(0.5) + ) + + def forward(self, x): + feat1 = self.b0(x) + feat2 = self.b1(x) + feat3 = self.b2(x) + feat4 = self.b3(x) +# print("before b4: ", x.size()) + feat5 = self.b4(x) + x = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1) + x = self.project(x) + return x + + +def get_deeplabv3(dataset='pascal_voc', backbone='resnet50', pretrained=False, root='~/.torch/models', + pretrained_base=True, **kwargs): + acronyms = { + 'pascal_voc': 'pascal_voc', + 'pascal_aug': 'pascal_aug', + 'ade20k': 'ade', + 'coco': 'coco', + 'citys': 'citys', + } + from ..data.dataloader import datasets + model = DeepLabV3(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs) + if pretrained: + from .model_store import get_model_file + device = torch.device(kwargs['local_rank']) + model.load_state_dict(torch.load(get_model_file('deeplabv3_%s_%s' % (backbone, acronyms[dataset]), root=root), + map_location=device)) + return model + + +def get_deeplabv3_resnet50_voc(**kwargs): + return get_deeplabv3('pascal_voc', 'resnet50', **kwargs) + + +def get_deeplabv3_resnet101_voc(**kwargs): + return get_deeplabv3('pascal_voc', 'resnet101', **kwargs) + + +def get_deeplabv3_resnet152_voc(**kwargs): + return get_deeplabv3('pascal_voc', 'resnet152', **kwargs) + + +def get_deeplabv3_resnet50_ade(**kwargs): + return get_deeplabv3('ade20k', 'resnet50', **kwargs) + + +def get_deeplabv3_resnet101_ade(**kwargs): + return get_deeplabv3('ade20k', 'resnet101', **kwargs) + + +def get_deeplabv3_resnet152_ade(**kwargs): + return get_deeplabv3('ade20k', 'resnet152', **kwargs) + + +if __name__ == '__main__': + model = get_deeplabv3_resnet50_voc() + img = torch.randn(2, 3, 480, 480) + output = model(img) diff --git a/hair_service_sd/core/seg/networks/deeplabv3_plus.py b/hair_service_sd/core/seg/networks/deeplabv3_plus.py new file mode 100644 index 0000000..a532fb8 --- /dev/null +++ b/hair_service_sd/core/seg/networks/deeplabv3_plus.py @@ -0,0 +1,160 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +from core.seg.networks.xception import get_xception +from core.seg.networks.deeplabv3 import _ASPP +from core.seg.networks.fcn import _FCNHead +from core.seg.networks.basic import _ConvBNReLU + +__all__ = ['DeepLabV3Plus', 'get_deeplabv3_plus', 'get_deeplabv3_plus_xception_voc'] + + +class DeepLabV3Plus(nn.Module): + r"""DeepLabV3Plus + Parameters + ---------- + nclass : int + Number of categories for the training dataset. + backbone : string + Pre-trained dilated backbone network type (default:'xception'). + norm_layer : object + Normalization layer used in backbone network (default: :class:`nn.BatchNorm`; + for Synchronized Cross-GPU BachNormalization). + aux : bool + Auxiliary loss. + + Reference: + Chen, Liang-Chieh, et al. "Encoder-Decoder with Atrous Separable Convolution for Semantic + Image Segmentation." + """ + + def __init__(self, nclass, backbone='xception', aux=True, pretrained_base=True, dilated=True, **kwargs): + super(DeepLabV3Plus, self).__init__() + self.aux = aux + self.nclass = nclass + output_stride = 8 if dilated else 32 + + self.pretrained = get_xception(pretrained=pretrained_base, output_stride=output_stride, **kwargs) + + # deeplabv3 plus + self.head = _DeepLabHead(nclass, **kwargs) + if aux: + self.auxlayer = _FCNHead(728, nclass, **kwargs) + + def base_forward(self, x): + # Entry flow + x = self.pretrained.conv1(x) + x = self.pretrained.bn1(x) + x = self.pretrained.relu(x) + + x = self.pretrained.conv2(x) + x = self.pretrained.bn2(x) + x = self.pretrained.relu(x) + + x = self.pretrained.block1(x) + # add relu here + x = self.pretrained.relu(x) + low_level_feat = x + + x = self.pretrained.block2(x) + x = self.pretrained.block3(x) + + # Middle flow + x = self.pretrained.midflow(x) + mid_level_feat = x + + # Exit flow + x = self.pretrained.block20(x) + x = self.pretrained.relu(x) + x = self.pretrained.conv3(x) + x = self.pretrained.bn3(x) + x = self.pretrained.relu(x) + + x = self.pretrained.conv4(x) + x = self.pretrained.bn4(x) + x = self.pretrained.relu(x) + + x = self.pretrained.conv5(x) + x = self.pretrained.bn5(x) + x = self.pretrained.relu(x) + return low_level_feat, mid_level_feat, x + + def forward(self, x): +# print("x size: ", x.size()) + size = x.size()[2:] + c1, c3, c4 = self.base_forward(x) +# print("c1 size: ", c1.size()) +# print("c4 size: ", c4.size()) + outputs = list() + x = self.head(c4, c1) + x = F.interpolate(x, size, mode='bilinear', align_corners=True) + outputs.append(x) + if self.aux: + auxout = self.auxlayer(c3) + auxout = F.interpolate(auxout, size, mode='bilinear', align_corners=True) + outputs.append(auxout) + + # for save onnx + # y = torch.max(x, 1)[1].to(torch.float32) + # return y + + return tuple(outputs) + + +class _DeepLabHead(nn.Module): + def __init__(self, nclass, c1_channels=128, norm_layer=nn.BatchNorm2d, **kwargs): + super(_DeepLabHead, self).__init__() + self.aspp = _ASPP(2048, [12, 24, 36], norm_layer=norm_layer, **kwargs) + self.c1_block = _ConvBNReLU(c1_channels, 48, 3, padding=1, norm_layer=norm_layer) + self.block = nn.Sequential( + _ConvBNReLU(304, 256, 3, padding=1, norm_layer=norm_layer), + nn.Dropout(0.5), + _ConvBNReLU(256, 256, 3, padding=1, norm_layer=norm_layer), + nn.Dropout(0.1), + nn.Conv2d(256, nclass, 1)) + + def forward(self, x, c1): + size = c1.size()[2:] + c1 = self.c1_block(c1) +# print("c1", c1.size()) +# print("before aspp: ", x.size()) + x = self.aspp(x) +# print("after aspp: ", x.size()) + x = F.interpolate(x, size, mode='bilinear', align_corners=True) + return self.block(torch.cat([x, c1], dim=1)) + + +def get_deeplabv3_plus(dataset='pascal_voc', backbone='xception', pretrained=False, root='../ckpt', + pretrained_base=False, nclass=3, **kwargs): + acronyms = { + 'pascal_voc': 'pascal_voc', + 'pascal_aug': 'pascal_aug', + 'ade20k': 'ade', + 'coco': 'coco', + 'citys': 'citys', + } + #from light.data import datasets + + model = DeepLabV3Plus(nclass, backbone=backbone, pretrained_base=pretrained_base, **kwargs) + if pretrained: + pass + # if dataset not in acronyms.keys(): + # print("root:", root) + # model_path = os.path.join(root, "deeplabv3_plus_28.pth") + # model.load_state_dict(torch.load(model_path), strict=False) + # else: + # from .model_store import get_model_file + # device = torch.device(kwargs['local_rank']) + # model.load_state_dict( + # torch.load(get_model_file('deeplabv3_plus_%s_%s' % (backbone, acronyms[dataset]), root=root), + # map_location=device)) + return model + + +def get_deeplabv3_plus_xception_voc(**kwargs): + return get_deeplabv3_plus('pascal_voc', 'xception', **kwargs) + + +if __name__ == '__main__': + model = get_deeplabv3_plus_xception_voc() diff --git a/hair_service_sd/core/seg/networks/fcn.py b/hair_service_sd/core/seg/networks/fcn.py new file mode 100644 index 0000000..100a601 --- /dev/null +++ b/hair_service_sd/core/seg/networks/fcn.py @@ -0,0 +1,221 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +from core.seg.networks.vgg import vgg16 + +__all__ = ['get_fcn32s', 'get_fcn16s', 'get_fcn8s', + 'get_fcn32s_vgg16_voc', 'get_fcn16s_vgg16_voc', 'get_fcn8s_vgg16_voc'] + + +class FCN32s(nn.Module): + """There are some difference from original fcn""" + + def __init__(self, nclass, backbone='vgg16', aux=False, pretrained_base=True, + norm_layer=nn.BatchNorm2d, **kwargs): + super(FCN32s, self).__init__() + self.aux = aux + if backbone == 'vgg16': + self.pretrained = vgg16(pretrained=pretrained_base).features + else: + raise RuntimeError('unknown backbone: {}'.format(backbone)) + self.head = _FCNHead(512, nclass, norm_layer) + if aux: + self.auxlayer = _FCNHead(512, nclass, norm_layer) + + self.__setattr__('exclusive', ['head', 'auxlayer'] if aux else ['head']) + + def forward(self, x): + size = x.size()[2:] + pool5 = self.pretrained(x) + + outputs = [] + out = self.head(pool5) + out = F.interpolate(out, size, mode='bilinear', align_corners=True) + outputs.append(out) + + if self.aux: + auxout = self.auxlayer(pool5) + auxout = F.interpolate(auxout, size, mode='bilinear', align_corners=True) + outputs.append(auxout) + + return tuple(outputs) + + +class FCN16s(nn.Module): + def __init__(self, nclass, backbone='vgg16', aux=False, pretrained_base=True, norm_layer=nn.BatchNorm2d, **kwargs): + super(FCN16s, self).__init__() + self.aux = aux + if backbone == 'vgg16': + self.pretrained = vgg16(pretrained=pretrained_base).features + else: + raise RuntimeError('unknown backbone: {}'.format(backbone)) + self.pool4 = nn.Sequential(*self.pretrained[:24]) + self.pool5 = nn.Sequential(*self.pretrained[24:]) + self.head = _FCNHead(512, nclass, norm_layer) + self.score_pool4 = nn.Conv2d(512, nclass, 1) + if aux: + self.auxlayer = _FCNHead(512, nclass, norm_layer) + + self.__setattr__('exclusive', ['head', 'score_pool4', 'auxlayer'] if aux else ['head', 'score_pool4']) + + def forward(self, x): + pool4 = self.pool4(x) + pool5 = self.pool5(pool4) + + outputs = [] + score_fr = self.head(pool5) + + score_pool4 = self.score_pool4(pool4) + + upscore2 = F.interpolate(score_fr, score_pool4.size()[2:], mode='bilinear', align_corners=True) + fuse_pool4 = upscore2 + score_pool4 + + out = F.interpolate(fuse_pool4, x.size()[2:], mode='bilinear', align_corners=True) + outputs.append(out) + + if self.aux: + auxout = self.auxlayer(pool5) + auxout = F.interpolate(auxout, x.size()[2:], mode='bilinear', align_corners=True) + outputs.append(auxout) + + return tuple(outputs) + + +class FCN8s(nn.Module): + def __init__(self, nclass, backbone='vgg16', aux=False, pretrained_base=True, norm_layer=nn.BatchNorm2d, **kwargs): + super(FCN8s, self).__init__() + self.aux = aux + if backbone == 'vgg16': + self.pretrained = vgg16(pretrained=pretrained_base).features + else: + raise RuntimeError('unknown backbone: {}'.format(backbone)) + self.pool3 = nn.Sequential(*self.pretrained[:17]) + self.pool4 = nn.Sequential(*self.pretrained[17:24]) + self.pool5 = nn.Sequential(*self.pretrained[24:]) + self.head = _FCNHead(512, nclass, norm_layer) + self.score_pool3 = nn.Conv2d(256, nclass, 1) + self.score_pool4 = nn.Conv2d(512, nclass, 1) + if aux: + self.auxlayer = _FCNHead(512, nclass, norm_layer) + + self.__setattr__('exclusive', + ['head', 'score_pool3', 'score_pool4', 'auxlayer'] if aux else ['head', 'score_pool3', + 'score_pool4']) + + def forward(self, x): + pool3 = self.pool3(x) + pool4 = self.pool4(pool3) + pool5 = self.pool5(pool4) + + outputs = [] + score_fr = self.head(pool5) + + score_pool4 = self.score_pool4(pool4) + score_pool3 = self.score_pool3(pool3) + + upscore2 = F.interpolate(score_fr, score_pool4.size()[2:], mode='bilinear', align_corners=True) + fuse_pool4 = upscore2 + score_pool4 + + upscore_pool4 = F.interpolate(fuse_pool4, score_pool3.size()[2:], mode='bilinear', align_corners=True) + fuse_pool3 = upscore_pool4 + score_pool3 + + out = F.interpolate(fuse_pool3, x.size()[2:], mode='bilinear', align_corners=True) + outputs.append(out) + + if self.aux: + auxout = self.auxlayer(pool5) + auxout = F.interpolate(auxout, x.size()[2:], mode='bilinear', align_corners=True) + outputs.append(auxout) + + return tuple(outputs) + + +class _FCNHead(nn.Module): + def __init__(self, in_channels, channels, norm_layer=nn.BatchNorm2d, **kwargs): + super(_FCNHead, self).__init__() + inter_channels = in_channels // 4 + self.block = nn.Sequential( + nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False), + norm_layer(inter_channels), + nn.ReLU(inplace=True), + nn.Dropout(0.1), + nn.Conv2d(inter_channels, channels, 1) + ) + + def forward(self, x): + return self.block(x) + + +def get_fcn32s(dataset='pascal_voc', backbone='vgg16', pretrained=False, root='~/.torch/models', + pretrained_base=True, **kwargs): + acronyms = { + 'pascal_voc': 'pascal_voc', + 'pascal_aug': 'pascal_aug', + 'ade20k': 'ade', + 'coco': 'coco', + 'citys': 'citys', + } + from ..data.dataloader import datasets + model = FCN32s(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs) + if pretrained: + from .model_store import get_model_file + device = torch.device(kwargs['local_rank']) + model.load_state_dict(torch.load(get_model_file('fcn32s_%s_%s' % (backbone, acronyms[dataset]), root=root), + map_location=device)) + return model + + +def get_fcn16s(dataset='pascal_voc', backbone='vgg16', pretrained=False, root='~/.torch/models', + pretrained_base=True, **kwargs): + acronyms = { + 'pascal_voc': 'pascal_voc', + 'pascal_aug': 'pascal_aug', + 'ade20k': 'ade', + 'coco': 'coco', + 'citys': 'citys', + } + from ..data.dataloader import datasets + model = FCN16s(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs) + if pretrained: + from .model_store import get_model_file + device = torch.device(kwargs['local_rank']) + model.load_state_dict(torch.load(get_model_file('fcn16s_%s_%s' % (backbone, acronyms[dataset]), root=root), + map_location=device)) + return model + + +def get_fcn8s(dataset='pascal_voc', backbone='vgg16', pretrained=False, root='~/.torch/models', + pretrained_base=True, **kwargs): + acronyms = { + 'pascal_voc': 'pascal_voc', + 'pascal_aug': 'pascal_aug', + 'ade20k': 'ade', + 'coco': 'coco', + 'citys': 'citys', + } + from ..data.dataloader import datasets + model = FCN8s(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs) + if pretrained: + from .model_store import get_model_file + device = torch.device(kwargs['local_rank']) + model.load_state_dict(torch.load(get_model_file('fcn8s_%s_%s' % (backbone, acronyms[dataset]), root=root), + map_location=device)) + return model + + +def get_fcn32s_vgg16_voc(**kwargs): + return get_fcn32s('pascal_voc', 'vgg16', **kwargs) + + +def get_fcn16s_vgg16_voc(**kwargs): + return get_fcn16s('pascal_voc', 'vgg16', **kwargs) + + +def get_fcn8s_vgg16_voc(**kwargs): + return get_fcn8s('pascal_voc', 'vgg16', **kwargs) + + +if __name__ == '__main__': + model = FCN16s(21) + print(model) diff --git a/hair_service_sd/core/seg/networks/jpu.py b/hair_service_sd/core/seg/networks/jpu.py new file mode 100644 index 0000000..db23bab --- /dev/null +++ b/hair_service_sd/core/seg/networks/jpu.py @@ -0,0 +1,68 @@ +"""Joint Pyramid Upsampling""" +import torch +import torch.nn as nn +import torch.nn.functional as F + +__all__ = ['JPU'] + + +class SeparableConv2d(nn.Module): + def __init__(self, inplanes, planes, kernel_size=3, stride=1, padding=1, + dilation=1, bias=False, norm_layer=nn.BatchNorm2d): + super(SeparableConv2d, self).__init__() + self.conv = nn.Conv2d(inplanes, inplanes, kernel_size, stride, padding, dilation, groups=inplanes, bias=bias) + self.bn = norm_layer(inplanes) + self.pointwise = nn.Conv2d(inplanes, planes, 1, bias=bias) + + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + x = self.pointwise(x) + return x + + +# copy from: https://github.com/wuhuikai/FastFCN/blob/master/encoding/nn/customize.py +class JPU(nn.Module): + def __init__(self, in_channels, width=512, norm_layer=nn.BatchNorm2d, **kwargs): + super(JPU, self).__init__() + + self.conv5 = nn.Sequential( + nn.Conv2d(in_channels[-1], width, 3, padding=1, bias=False), + norm_layer(width), + nn.ReLU(True)) + self.conv4 = nn.Sequential( + nn.Conv2d(in_channels[-2], width, 3, padding=1, bias=False), + norm_layer(width), + nn.ReLU(True)) + self.conv3 = nn.Sequential( + nn.Conv2d(in_channels[-3], width, 3, padding=1, bias=False), + norm_layer(width), + nn.ReLU(True)) + + self.dilation1 = nn.Sequential( + SeparableConv2d(3 * width, width, 3, padding=1, dilation=1, bias=False), + norm_layer(width), + nn.ReLU(True)) + self.dilation2 = nn.Sequential( + SeparableConv2d(3 * width, width, 3, padding=2, dilation=2, bias=False), + norm_layer(width), + nn.ReLU(True)) + self.dilation3 = nn.Sequential( + SeparableConv2d(3 * width, width, 3, padding=4, dilation=4, bias=False), + norm_layer(width), + nn.ReLU(True)) + self.dilation4 = nn.Sequential( + SeparableConv2d(3 * width, width, 3, padding=8, dilation=8, bias=False), + norm_layer(width), + nn.ReLU(True)) + + def forward(self, *inputs): + feats = [self.conv5(inputs[-1]), self.conv4(inputs[-2]), self.conv3(inputs[-3])] + size = feats[-1].size()[2:] + feats[-2] = F.interpolate(feats[-2], size, mode='bilinear', align_corners=True) + feats[-3] = F.interpolate(feats[-3], size, mode='bilinear', align_corners=True) + feat = torch.cat(feats, dim=1) + feat = torch.cat([self.dilation1(feat), self.dilation2(feat), self.dilation3(feat), self.dilation4(feat)], + dim=1) + + return inputs[0], inputs[1], inputs[2], feat diff --git a/hair_service_sd/core/seg/networks/resnetv1b.py b/hair_service_sd/core/seg/networks/resnetv1b.py new file mode 100644 index 0000000..21d67b7 --- /dev/null +++ b/hair_service_sd/core/seg/networks/resnetv1b.py @@ -0,0 +1,264 @@ +import torch +import torch.nn as nn +import torch.utils.model_zoo as model_zoo + +__all__ = ['ResNetV1b', 'resnet18_v1b', 'resnet34_v1b', 'resnet50_v1b', + 'resnet101_v1b', 'resnet152_v1b', 'resnet152_v1s', 'resnet101_v1s', 'resnet50_v1s'] + +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', +} + + +class BasicBlockV1b(nn.Module): + expansion = 1 + + def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None, + previous_dilation=1, norm_layer=nn.BatchNorm2d): + super(BasicBlockV1b, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, 3, stride, + dilation, dilation, bias=False) + self.bn1 = norm_layer(planes) + self.relu = nn.ReLU(True) + self.conv2 = nn.Conv2d(planes, planes, 3, 1, previous_dilation, + dilation=previous_dilation, bias=False) + 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 BottleneckV1b(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None, + previous_dilation=1, norm_layer=nn.BatchNorm2d): + super(BottleneckV1b, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False) + self.bn1 = norm_layer(planes) + self.conv2 = nn.Conv2d(planes, planes, 3, stride, + dilation, dilation, bias=False) + self.bn2 = norm_layer(planes) + self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False) + self.bn3 = norm_layer(planes * self.expansion) + self.relu = nn.ReLU(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 ResNetV1b(nn.Module): + + def __init__(self, block, layers, num_classes=1000, dilated=True, deep_stem=False, + zero_init_residual=False, norm_layer=nn.BatchNorm2d): + self.inplanes = 128 if deep_stem else 64 + super(ResNetV1b, self).__init__() + if deep_stem: + self.conv1 = nn.Sequential( + nn.Conv2d(3, 64, 3, 2, 1, bias=False), + norm_layer(64), + nn.ReLU(True), + nn.Conv2d(64, 64, 3, 1, 1, bias=False), + norm_layer(64), + nn.ReLU(True), + nn.Conv2d(64, 128, 3, 1, 1, bias=False) + ) + else: + self.conv1 = nn.Conv2d(3, 64, 7, 2, 3, bias=False) + self.bn1 = norm_layer(self.inplanes) + self.relu = nn.ReLU(True) + self.maxpool = nn.MaxPool2d(3, 2, 1) + self.layer1 = self._make_layer(block, 64, layers[0], norm_layer=norm_layer) + self.layer2 = self._make_layer(block, 128, layers[1], stride=2, norm_layer=norm_layer) + if dilated: + self.layer3 = self._make_layer(block, 256, layers[2], stride=1, dilation=2, norm_layer=norm_layer) + self.layer4 = self._make_layer(block, 512, layers[3], stride=1, dilation=4, norm_layer=norm_layer) + else: + self.layer3 = self._make_layer(block, 256, layers[2], stride=2, norm_layer=norm_layer) + self.layer4 = self._make_layer(block, 512, layers[3], stride=2, norm_layer=norm_layer) + 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.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + if zero_init_residual: + for m in self.modules(): + if isinstance(m, BottleneckV1b): + nn.init.constant_(m.bn3.weight, 0) + elif isinstance(m, BasicBlockV1b): + nn.init.constant_(m.bn2.weight, 0) + + def _make_layer(self, block, planes, blocks, stride=1, dilation=1, norm_layer=nn.BatchNorm2d): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.inplanes, planes * block.expansion, 1, stride, bias=False), + norm_layer(planes * block.expansion), + ) + + layers = [] + if dilation in (1, 2): + layers.append(block(self.inplanes, planes, stride, dilation=1, downsample=downsample, + previous_dilation=dilation, norm_layer=norm_layer)) + elif dilation == 4: + layers.append(block(self.inplanes, planes, stride, dilation=2, downsample=downsample, + previous_dilation=dilation, norm_layer=norm_layer)) + else: + raise RuntimeError("=> unknown dilation size: {}".format(dilation)) + self.inplanes = planes * block.expansion + for _ in range(1, blocks): + layers.append(block(self.inplanes, planes, dilation=dilation, + previous_dilation=dilation, norm_layer=norm_layer)) + + return nn.Sequential(*layers) + + def forward(self, x): + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + x = self.maxpool(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + + x = self.avgpool(x) + x = x.view(x.size(0), -1) + x = self.fc(x) + + return x + + +def resnet18_v1b(pretrained=False, **kwargs): + model = ResNetV1b(BasicBlockV1b, [2, 2, 2, 2], **kwargs) + if pretrained: + old_dict = model_zoo.load_url(model_urls['resnet18']) + model_dict = model.state_dict() + old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)} + model_dict.update(old_dict) + model.load_state_dict(model_dict) + return model + + +def resnet34_v1b(pretrained=False, **kwargs): + model = ResNetV1b(BasicBlockV1b, [3, 4, 6, 3], **kwargs) + if pretrained: + old_dict = model_zoo.load_url(model_urls['resnet34']) + model_dict = model.state_dict() + old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)} + model_dict.update(old_dict) + model.load_state_dict(model_dict) + return model + + +def resnet50_v1b(pretrained=False, **kwargs): + model = ResNetV1b(BottleneckV1b, [3, 4, 6, 3], **kwargs) + if pretrained: + old_dict = model_zoo.load_url(model_urls['resnet50']) + model_dict = model.state_dict() + old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)} + model_dict.update(old_dict) + model.load_state_dict(model_dict) + return model + + +def resnet101_v1b(pretrained=False, **kwargs): + model = ResNetV1b(BottleneckV1b, [3, 4, 23, 3], **kwargs) + if pretrained: + old_dict = model_zoo.load_url(model_urls['resnet101']) + model_dict = model.state_dict() + old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)} + model_dict.update(old_dict) + model.load_state_dict(model_dict) + return model + + +def resnet152_v1b(pretrained=False, **kwargs): + model = ResNetV1b(BottleneckV1b, [3, 8, 36, 3], **kwargs) + if pretrained: + old_dict = model_zoo.load_url(model_urls['resnet152']) + model_dict = model.state_dict() + old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)} + model_dict.update(old_dict) + model.load_state_dict(model_dict) + return model + + +def resnet50_v1s(pretrained=False, root='~/.torch/models', **kwargs): + model = ResNetV1b(BottleneckV1b, [3, 4, 6, 3], deep_stem=True, **kwargs) + if pretrained: + from ..model_store import get_resnet_file + model.load_state_dict(torch.load(get_resnet_file('resnet50', root=root)), strict=False) + return model + + +def resnet101_v1s(pretrained=False, root='~/.torch/models', **kwargs): + model = ResNetV1b(BottleneckV1b, [3, 4, 23, 3], deep_stem=True, **kwargs) + if pretrained: + from ..model_store import get_resnet_file + model.load_state_dict(torch.load(get_resnet_file('resnet101', root=root)), strict=False) + return model + + +def resnet152_v1s(pretrained=False, root='~/.torch/models', **kwargs): + model = ResNetV1b(BottleneckV1b, [3, 8, 36, 3], deep_stem=True, **kwargs) + if pretrained: + from ..model_store import get_resnet_file + model.load_state_dict(torch.load(get_resnet_file('resnet152', root=root)), strict=False) + return model + + +if __name__ == '__main__': + import torch + + img = torch.randn(4, 3, 224, 224) + model = resnet50_v1b(True) + output = model(img) diff --git a/hair_service_sd/core/seg/networks/segbase.py b/hair_service_sd/core/seg/networks/segbase.py new file mode 100644 index 0000000..c27c277 --- /dev/null +++ b/hair_service_sd/core/seg/networks/segbase.py @@ -0,0 +1,60 @@ +"""Base Model for Semantic Segmentation""" +import torch.nn as nn + +from core.seg.networks.jpu import JPU +from core.seg.networks.resnetv1b import resnet50_v1s, resnet101_v1s, resnet152_v1s + +__all__ = ['SegBaseModel'] + + +class SegBaseModel(nn.Module): + r"""Base Model for Semantic Segmentation + + Parameters + ---------- + backbone : string + Pre-trained dilated backbone network type (default:'resnet50'; 'resnet50', + 'resnet101' or 'resnet152'). + """ + + def __init__(self, nclass, aux, backbone='resnet50', jpu=False, pretrained_base=True, **kwargs): + super(SegBaseModel, self).__init__() + dilated = False if jpu else True + self.aux = aux + self.nclass = nclass + if backbone == 'resnet50': + self.pretrained = resnet50_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs) + elif backbone == 'resnet101': + self.pretrained = resnet101_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs) + elif backbone == 'resnet152': + self.pretrained = resnet152_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs) + else: + raise RuntimeError('unknown backbone: {}'.format(backbone)) + + self.jpu = JPU([512, 1024, 2048], width=512, **kwargs) if jpu else None + + def base_forward(self, x): + """forwarding pre-trained network""" + x = self.pretrained.conv1(x) + x = self.pretrained.bn1(x) + x = self.pretrained.relu(x) + x = self.pretrained.maxpool(x) + c1 = self.pretrained.layer1(x) + c2 = self.pretrained.layer2(c1) + c3 = self.pretrained.layer3(c2) + c4 = self.pretrained.layer4(c3) + + if self.jpu: + return self.jpu(c1, c2, c3, c4) + else: + return c1, c2, c3, c4 + + def evaluate(self, x): + """evaluating network with inputs and targets""" + return self.forward(x)[0] + + def demo(self, x): + pred = self.forward(x) + if self.aux: + pred = pred[0] + return pred diff --git a/hair_service_sd/core/seg/networks/vgg.py b/hair_service_sd/core/seg/networks/vgg.py new file mode 100644 index 0000000..fe5c163 --- /dev/null +++ b/hair_service_sd/core/seg/networks/vgg.py @@ -0,0 +1,191 @@ +import torch +import torch.nn as nn +import torch.utils.model_zoo as model_zoo + +__all__ = [ + 'VGG', 'vgg11', 'vgg11_bn', 'vgg13', 'vgg13_bn', 'vgg16', 'vgg16_bn', + 'vgg19_bn', 'vgg19', +] + +model_urls = { + 'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth', + 'vgg13': 'https://download.pytorch.org/models/vgg13-c768596a.pth', + 'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth', + 'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth', + 'vgg11_bn': 'https://download.pytorch.org/models/vgg11_bn-6002323d.pth', + 'vgg13_bn': 'https://download.pytorch.org/models/vgg13_bn-abd245e5.pth', + 'vgg16_bn': 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth', + 'vgg19_bn': 'https://download.pytorch.org/models/vgg19_bn-c79401a0.pth', +} + + +class VGG(nn.Module): + def __init__(self, features, num_classes=1000, init_weights=True): + super(VGG, self).__init__() + self.features = features + self.avgpool = nn.AdaptiveAvgPool2d((7, 7)) + self.classifier = nn.Sequential( + nn.Linear(512 * 7 * 7, 4096), + nn.ReLU(True), + nn.Dropout(), + nn.Linear(4096, 4096), + nn.ReLU(True), + nn.Dropout(), + nn.Linear(4096, num_classes) + ) + if init_weights: + self._initialize_weights() + + def forward(self, x): + x = self.features(x) + x = self.avgpool(x) + x = x.view(x.size(0), -1) + x = self.classifier(x) + return x + + def _initialize_weights(self): + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.BatchNorm2d): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.Linear): + nn.init.normal_(m.weight, 0, 0.01) + nn.init.constant_(m.bias, 0) + + +def make_layers(cfg, batch_norm=False): + layers = [] + in_channels = 3 + for v in cfg: + if v == 'M': + layers += [nn.MaxPool2d(kernel_size=2, stride=2)] + else: + conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1) + if batch_norm: + layers += (conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)) + else: + layers += [conv2d, nn.ReLU(inplace=True)] + in_channels = v + return nn.Sequential(*layers) + + +cfg = { + 'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], + 'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], + 'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'], + 'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'], +} + + +def vgg11(pretrained=False, **kwargs): + """VGG 11-layer model (configuration "A") + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['A']), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg11'])) + return model + + +def vgg11_bn(pretrained=False, **kwargs): + """VGG 11-layer model (configuration "A") with batch normalization + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['A'], batch_norm=True), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg11_bn'])) + return model + + +def vgg13(pretrained=False, **kwargs): + """VGG 13-layer model (configuration "B") + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['B']), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg13'])) + return model + + +def vgg13_bn(pretrained=False, **kwargs): + """VGG 13-layer model (configuration "B") with batch normalization + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['B'], batch_norm=True), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg13_bn'])) + return model + + +def vgg16(pretrained=False, **kwargs): + """VGG 16-layer model (configuration "D") + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['D']), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg16'])) + return model + + +def vgg16_bn(pretrained=False, **kwargs): + """VGG 16-layer model (configuration "D") with batch normalization + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['D'], batch_norm=True), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg16_bn'])) + return model + + +def vgg19(pretrained=False, **kwargs): + """VGG 19-layer model (configuration "E") + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['E']), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg19'])) + return model + + +def vgg19_bn(pretrained=False, **kwargs): + """VGG 19-layer model (configuration 'E') with batch normalization + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['E'], batch_norm=True), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg19_bn'])) + return model + + +if __name__ == '__main__': + img = torch.randn((4, 3, 480, 480)) + model = vgg16(pretrained=False) + out = model(img) diff --git a/hair_service_sd/core/seg/networks/xception.py b/hair_service_sd/core/seg/networks/xception.py new file mode 100644 index 0000000..52dc0b9 --- /dev/null +++ b/hair_service_sd/core/seg/networks/xception.py @@ -0,0 +1,411 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +__all__ = ['Enc', 'FCAttention', 'Xception65', 'Xception71', 'get_xception', 'get_xception_71', 'get_xception_a'] + + +class SeparableConv2d(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, dilation=1, bias=False, norm_layer=None): + super(SeparableConv2d, self).__init__() + self.kernel_size = kernel_size + self.dilation = dilation + + self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size, stride, 0, dilation, groups=in_channels, + bias=bias) + self.bn = norm_layer(in_channels) + self.pointwise = nn.Conv2d(in_channels, out_channels, 1, bias=bias) + + def forward(self, x): + x = self.fix_padding(x, self.kernel_size, self.dilation) + x = self.conv1(x) + x = self.bn(x) + x = self.pointwise(x) + + return x + + def fix_padding(self, x, 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(x, (pad_beg, pad_end, pad_beg, pad_end)) + return padded_inputs + + +class Block(nn.Module): + def __init__(self, in_channels, out_channels, reps, stride=1, dilation=1, norm_layer=None, + start_with_relu=True, grow_first=True, is_last=False): + super(Block, self).__init__() + if out_channels != in_channels or stride != 1: + self.skip = nn.Conv2d(in_channels, out_channels, 1, stride, bias=False) + self.skipbn = norm_layer(out_channels) + else: + self.skip = None + self.relu = nn.ReLU(True) + rep = list() + filters = in_channels + if grow_first: + if start_with_relu: + rep.append(self.relu) + rep.append(SeparableConv2d(in_channels, out_channels, 3, 1, dilation, norm_layer=norm_layer)) + rep.append(norm_layer(out_channels)) + filters = out_channels + for i in range(reps - 1): + if grow_first or start_with_relu: + rep.append(self.relu) + rep.append(SeparableConv2d(filters, filters, 3, 1, dilation, norm_layer=norm_layer)) + rep.append(norm_layer(filters)) + if not grow_first: + rep.append(self.relu) + rep.append(SeparableConv2d(in_channels, out_channels, 3, 1, dilation, norm_layer=norm_layer)) + if stride != 1: + rep.append(self.relu) + rep.append(SeparableConv2d(out_channels, out_channels, 3, stride, norm_layer=norm_layer)) + rep.append(norm_layer(out_channels)) + elif is_last: + rep.append(self.relu) + rep.append(SeparableConv2d(out_channels, out_channels, 3, 1, dilation, norm_layer=norm_layer)) + rep.append(norm_layer(out_channels)) + self.rep = nn.Sequential(*rep) + + def forward(self, x): + out = self.rep(x) + if self.skip is not None: + skip = self.skipbn(self.skip(x)) + else: + skip = x + out = out + skip + return out + + +class Xception65(nn.Module): + """Modified Aligned Xception + """ + + def __init__(self, num_classes=1000, output_stride=32, norm_layer=nn.BatchNorm2d): + super(Xception65, self).__init__() + if output_stride == 32: + entry_block3_stride = 2 + exit_block20_stride = 2 + middle_block_dilation = 1 + exit_block_dilations = (1, 1) + elif output_stride == 16: + entry_block3_stride = 2 + exit_block20_stride = 1 + middle_block_dilation = 1 + exit_block_dilations = (1, 2) + elif output_stride == 8: + entry_block3_stride = 1 + exit_block20_stride = 1 + middle_block_dilation = 2 + exit_block_dilations = (2, 4) + else: + raise NotImplementedError + # Entry flow + self.conv1 = nn.Conv2d(3, 32, 3, 2, 1, bias=False) + self.bn1 = norm_layer(32) + self.relu = nn.ReLU(True) + + self.conv2 = nn.Conv2d(32, 64, 3, 1, 1, bias=False) + self.bn2 = norm_layer(64) + + self.block1 = Block(64, 128, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False) + self.block2 = Block(128, 256, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False, grow_first=True) + self.block3 = Block(256, 728, reps=2, stride=entry_block3_stride, norm_layer=norm_layer, + start_with_relu=True, grow_first=True, is_last=True) + + # Middle flow + midflow = list() + for i in range(4, 20): + midflow.append(Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation, norm_layer=norm_layer, + start_with_relu=True, grow_first=True)) + self.midflow = nn.Sequential(*midflow) + + # Exit flow + self.block20 = Block(728, 1024, reps=2, stride=exit_block20_stride, dilation=exit_block_dilations[0], + norm_layer=norm_layer, start_with_relu=True, grow_first=False, is_last=True) + self.conv3 = SeparableConv2d(1024, 1536, 3, 1, dilation=exit_block_dilations[1], norm_layer=norm_layer) + self.bn3 = norm_layer(1536) + self.conv4 = SeparableConv2d(1536, 1536, 3, stride=1, dilation=exit_block_dilations[1], norm_layer=norm_layer) + self.bn4 = norm_layer(1536) + self.conv5 = SeparableConv2d(1536, 2048, 3, 1, dilation=exit_block_dilations[1], norm_layer=norm_layer) + self.bn5 = norm_layer(2048) + self.avgpool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Linear(2048, num_classes) + + 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) + x = self.relu(x) + # c1 = x + x = self.block2(x) + # c2 = x + x = self.block3(x) + + # Middle flow + x = self.midflow(x) + # c3 = 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) + + x = self.avgpool(x) + x = x.view(x.size(0), -1) + x = self.fc(x) + + return x + + +class Xception71(nn.Module): + """Modified Aligned Xception + """ + + def __init__(self, num_classes=1000, output_stride=32, norm_layer=nn.BatchNorm2d): + super(Xception71, self).__init__() + if output_stride == 32: + entry_block3_stride = 2 + exit_block20_stride = 2 + middle_block_dilation = 1 + exit_block_dilations = (1, 1) + elif output_stride == 16: + entry_block3_stride = 2 + exit_block20_stride = 1 + middle_block_dilation = 1 + exit_block_dilations = (1, 2) + elif output_stride == 8: + entry_block3_stride = 1 + exit_block20_stride = 1 + middle_block_dilation = 2 + exit_block_dilations = (2, 4) + else: + raise NotImplementedError + # Entry flow + self.conv1 = nn.Conv2d(3, 32, 3, 2, 1, bias=False) + self.bn1 = norm_layer(32) + self.relu = nn.ReLU(True) + + self.conv2 = nn.Conv2d(32, 64, 3, 1, 1, bias=False) + self.bn2 = norm_layer(64) + + self.block1 = Block(64, 128, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False) + self.block2 = nn.Sequential( + Block(128, 256, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False, grow_first=True), + Block(256, 728, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False, grow_first=True)) + self.block3 = Block(728, 728, reps=2, stride=entry_block3_stride, norm_layer=norm_layer, + start_with_relu=True, grow_first=True, is_last=True) + + # Middle flow + midflow = list() + for i in range(4, 20): + midflow.append(Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation, norm_layer=norm_layer, + start_with_relu=True, grow_first=True)) + self.midflow = nn.Sequential(*midflow) + + # Exit flow + self.block20 = Block(728, 1024, reps=2, stride=exit_block20_stride, dilation=exit_block_dilations[0], + norm_layer=norm_layer, start_with_relu=True, grow_first=False, is_last=True) + self.conv3 = SeparableConv2d(1024, 1536, 3, 1, dilation=exit_block_dilations[1], norm_layer=norm_layer) + self.bn3 = norm_layer(1536) + self.conv4 = SeparableConv2d(1536, 1536, 3, stride=1, dilation=exit_block_dilations[1], norm_layer=norm_layer) + self.bn4 = norm_layer(1536) + self.conv5 = SeparableConv2d(1536, 2048, 3, 1, dilation=exit_block_dilations[1], norm_layer=norm_layer) + self.bn5 = norm_layer(2048) + self.avgpool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Linear(2048, num_classes) + + 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) + x = self.relu(x) + # c1 = x + x = self.block2(x) + # c2 = x + x = self.block3(x) + + # Middle flow + x = self.midflow(x) + # c3 = 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) + + x = self.avgpool(x) + x = x.view(x.size(0), -1) + x = self.fc(x) + + return x + + +# ------------------------------------------------- +# For DFANet +# ------------------------------------------------- +class BlockA(nn.Module): + def __init__(self, in_channels, out_channels, stride=1, dilation=1, norm_layer=None, start_with_relu=True): + super(BlockA, self).__init__() + if out_channels != in_channels or stride != 1: + self.skip = nn.Conv2d(in_channels, out_channels, 1, stride, bias=False) + self.skipbn = norm_layer(out_channels) + else: + self.skip = None + self.relu = nn.ReLU(True) + rep = list() + inter_channels = out_channels // 4 + + if start_with_relu: + rep.append(self.relu) + rep.append(SeparableConv2d(in_channels, inter_channels, 3, 1, dilation, norm_layer=norm_layer)) + rep.append(norm_layer(inter_channels)) + + rep.append(self.relu) + rep.append(SeparableConv2d(inter_channels, inter_channels, 3, 1, dilation, norm_layer=norm_layer)) + rep.append(norm_layer(inter_channels)) + + if stride != 1: + rep.append(self.relu) + rep.append(SeparableConv2d(inter_channels, out_channels, 3, stride, norm_layer=norm_layer)) + rep.append(norm_layer(out_channels)) + else: + rep.append(self.relu) + rep.append(SeparableConv2d(inter_channels, out_channels, 3, 1, norm_layer=norm_layer)) + rep.append(norm_layer(out_channels)) + self.rep = nn.Sequential(*rep) + + def forward(self, x): + out = self.rep(x) + if self.skip is not None: + skip = self.skipbn(self.skip(x)) + else: + skip = x + out = out + skip + return out + + +class Enc(nn.Module): + def __init__(self, in_channels, out_channels, blocks, norm_layer=None): + super(Enc, self).__init__() + block = list() + block.append(BlockA(in_channels, out_channels, 2, norm_layer=norm_layer)) + for i in range(blocks - 1): + block.append(BlockA(out_channels, out_channels, 1, norm_layer=norm_layer)) + self.block = nn.Sequential(*block) + + def forward(self, x): + return self.block(x) + + +class FCAttention(nn.Module): + def __init__(self, in_channels, norm_layer=None): + super(FCAttention, self).__init__() + self.avgpool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Linear(in_channels, 1000) + self.conv = nn.Sequential( + nn.Conv2d(1000, in_channels, 1, bias=False), + norm_layer(in_channels), + nn.ReLU(True)) + + def forward(self, x): + n, c, _, _ = x.size() + att = self.avgpool(x).view(n, c) + att = self.fc(att).view(n, 1000, 1, 1) + att = self.conv(att) + return x * att.expand_as(x) + + +class XceptionA(nn.Module): + def __init__(self, num_classes=1000, norm_layer=nn.BatchNorm2d): + super(XceptionA, self).__init__() + self.conv1 = nn.Sequential(nn.Conv2d(3, 8, 3, 2, 1, bias=False), + norm_layer(8), + nn.ReLU(True)) + + self.enc2 = Enc(8, 48, 4, norm_layer=norm_layer) + self.enc3 = Enc(48, 96, 6, norm_layer=norm_layer) + self.enc4 = Enc(96, 192, 4, norm_layer=norm_layer) + + self.fca = FCAttention(192, norm_layer=norm_layer) + self.avgpool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Linear(192, num_classes) + + def forward(self, x): + x = self.conv1(x) + + x = self.enc2(x) + x = self.enc3(x) + x = self.enc4(x) + x = self.fca(x) + + x = self.avgpool(x) + x = x.view(x.size(0), -1) + x = self.fc(x) + + return x + + +# Constructor +def get_xception(pretrained=False, root='~/.torch/models', **kwargs): + model = Xception65(**kwargs) + if pretrained: + from ..model_store import get_model_file + model.load_state_dict(torch.load(get_model_file('xception', root=root))) + return model + + +def get_xception_71(pretrained=False, root='~/.torch/models', **kwargs): + model = Xception71(**kwargs) + if pretrained: + from ..model_store import get_model_file + model.load_state_dict(torch.load(get_model_file('xception71', root=root))) + return model + + +def get_xception_a(pretrained=False, root='~/.torch/models', **kwargs): + model = XceptionA(**kwargs) + if pretrained: + from ..model_store import get_model_file + model.load_state_dict(torch.load(get_model_file('xception_a', root=root))) + return model + + +if __name__ == '__main__': + model = get_xception_a() diff --git a/hair_service_sd/core/seg/setup.py b/hair_service_sd/core/seg/setup.py new file mode 100644 index 0000000..ca6bf5a --- /dev/null +++ b/hair_service_sd/core/seg/setup.py @@ -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 \ No newline at end of file diff --git a/hair_service_sd/core/simsun.ttc b/hair_service_sd/core/simsun.ttc new file mode 100644 index 0000000..cd356c7 --- /dev/null +++ b/hair_service_sd/core/simsun.ttc @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7a368113c36d516aac1a825acc4bbbc9906c4d197f7c649e7fe0f6f31e7475dd +size 10500792 diff --git a/hair_service_sd/core/utils/MomocvFaceAlignment1K.py b/hair_service_sd/core/utils/MomocvFaceAlignment1K.py new file mode 100644 index 0000000..1e62212 --- /dev/null +++ b/hair_service_sd/core/utils/MomocvFaceAlignment1K.py @@ -0,0 +1,462 @@ +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 algorithm_conf import ConfFactory +from core.utils.umeyama import umeyama + +import cv2 + +__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_dir = ConfFactory.getModelValue("model_dir") + weights = torch.load(os.path.join(self.model_dir, '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('MomocvFaceAlignment1K success') + + 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, crop_M): + dst_size = 256 + with torch.no_grad(): + + # add for change crop for 1024 + h, w, _ = img.shape + if h == 768: + crop_img = img[104:img.shape[0] - 104, 104:img.shape[1] - 104, :] + tmp = cv2.resize(crop_img, (dst_size, dst_size)) + + else: + tmp = cv2.warpAffine(img, crop_M, (dst_size, dst_size), flags=cv2.INTER_CUBIC) + + # cv2.imshow("img_paf_test_crop: ", tmp) + # cv2.waitKey() + # crop_img = img[220:img.shape[0] - 220, 266:img.shape[1] - 266, :] # 220 266 + # tmp = cv2.resize(crop_img, (dst_size, dst_size)) + + # cv2.imshow("detect face: h:{:d}".format(h), tmp) + # cv2.waitKey() + + 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)) + + if h == 768: + orig_pts[:, 0] = orig_pts[:, 0] * crop_img.shape[1] + orig_pts[:, 1] = orig_pts[:, 1] * crop_img.shape[0] + orig_pts[:, 0] += 104 + orig_pts[:, 1] += 104 + else: + orig_pts[:, 0] = orig_pts[:, 0] * dst_size + orig_pts[:, 1] = orig_pts[:, 1] * dst_size + orig_pts = landmark_processor.transform_points(orig_pts, crop_M, invert=True) + # orig_pts[:, 0] = orig_pts[:, 0] * crop_img.shape[1] + # orig_pts[:, 1] = orig_pts[:, 1] * crop_img.shape[0] + # orig_pts[:, 0] += 266 + # orig_pts[:, 1] += 220 + + 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.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # 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]]) + + 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 + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + + tmp = cv2.warpAffine(img, mat, (dst_size, dst_size)) + # 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 diff --git a/hair_service_sd/core/utils/__init__.py b/hair_service_sd/core/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hair_service_sd/core/utils/box_utils_Retina.py b/hair_service_sd/core/utils/box_utils_Retina.py new file mode 100644 index 0000000..c1d12bc --- /dev/null +++ b/hair_service_sd/core/utils/box_utils_Retina.py @@ -0,0 +1,330 @@ +import torch +import numpy as np + + +def point_form(boxes): + """ Convert prior_boxes to (xmin, ymin, xmax, ymax) + representation for comparison to point form ground truth data. + Args: + boxes: (tensor) center-size default boxes from priorbox layers. + Return: + boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes. + """ + return torch.cat((boxes[:, :2] - boxes[:, 2:]/2, # xmin, ymin + boxes[:, :2] + boxes[:, 2:]/2), 1) # xmax, ymax + + +def center_size(boxes): + """ Convert prior_boxes to (cx, cy, w, h) + representation for comparison to center-size form ground truth data. + Args: + boxes: (tensor) point_form boxes + Return: + boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes. + """ + return torch.cat((boxes[:, 2:] + boxes[:, :2])/2, # cx, cy + boxes[:, 2:] - boxes[:, :2], 1) # w, h + + +def intersect(box_a, box_b): + """ We resize both tensors to [A,B,2] without new malloc: + [A,2] -> [A,1,2] -> [A,B,2] + [B,2] -> [1,B,2] -> [A,B,2] + Then we compute the area of intersect between box_a and box_b. + Args: + box_a: (tensor) bounding boxes, Shape: [A,4]. + box_b: (tensor) bounding boxes, Shape: [B,4]. + Return: + (tensor) intersection area, Shape: [A,B]. + """ + A = box_a.size(0) + B = box_b.size(0) + max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2), + box_b[:, 2:].unsqueeze(0).expand(A, B, 2)) + min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2), + box_b[:, :2].unsqueeze(0).expand(A, B, 2)) + inter = torch.clamp((max_xy - min_xy), min=0) + return inter[:, :, 0] * inter[:, :, 1] + + +def jaccard(box_a, box_b): + """Compute the jaccard overlap of two sets of boxes. The jaccard overlap + is simply the intersection over union of two boxes. Here we operate on + ground truth boxes and default boxes. + E.g.: + A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B) + Args: + box_a: (tensor) Ground truth bounding boxes, Shape: [num_objects,4] + box_b: (tensor) Prior boxes from priorbox layers, Shape: [num_priors,4] + Return: + jaccard overlap: (tensor) Shape: [box_a.size(0), box_b.size(0)] + """ + inter = intersect(box_a, box_b) + area_a = ((box_a[:, 2]-box_a[:, 0]) * + (box_a[:, 3]-box_a[:, 1])).unsqueeze(1).expand_as(inter) # [A,B] + area_b = ((box_b[:, 2]-box_b[:, 0]) * + (box_b[:, 3]-box_b[:, 1])).unsqueeze(0).expand_as(inter) # [A,B] + union = area_a + area_b - inter + return inter / union # [A,B] + + +def matrix_iou(a, b): + """ + return iou of a and b, numpy version for data augenmentation + """ + lt = np.maximum(a[:, np.newaxis, :2], b[:, :2]) + rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:]) + + area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2) + area_a = np.prod(a[:, 2:] - a[:, :2], axis=1) + area_b = np.prod(b[:, 2:] - b[:, :2], axis=1) + return area_i / (area_a[:, np.newaxis] + area_b - area_i) + + +def matrix_iof(a, b): + """ + return iof of a and b, numpy version for data augenmentation + """ + lt = np.maximum(a[:, np.newaxis, :2], b[:, :2]) + rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:]) + + area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2) + area_a = np.prod(a[:, 2:] - a[:, :2], axis=1) + return area_i / np.maximum(area_a[:, np.newaxis], 1) + + +def match(threshold, truths, priors, variances, labels, landms, loc_t, conf_t, landm_t, idx): + """Match each prior box with the ground truth box of the highest jaccard + overlap, encode the bounding boxes, then return the matched indices + corresponding to both confidence and location preds. + Args: + threshold: (float) The overlap threshold used when mathing boxes. + truths: (tensor) Ground truth boxes, Shape: [num_obj, 4]. + priors: (tensor) Prior boxes from priorbox layers, Shape: [n_priors,4]. + variances: (tensor) Variances corresponding to each prior coord, + Shape: [num_priors, 4]. + labels: (tensor) All the class labels for the image, Shape: [num_obj]. + landms: (tensor) Ground truth landms, Shape [num_obj, 10]. + loc_t: (tensor) Tensor to be filled w/ endcoded location targets. + conf_t: (tensor) Tensor to be filled w/ matched indices for conf preds. + landm_t: (tensor) Tensor to be filled w/ endcoded landm targets. + idx: (int) current batch index + Return: + The matched indices corresponding to 1)location 2)confidence 3)landm preds. + """ + # jaccard index + overlaps = jaccard( + truths, + point_form(priors) + ) + # (Bipartite Matching) + # [1,num_objects] best prior for each ground truth + best_prior_overlap, best_prior_idx = overlaps.max(1, keepdim=True) + + # ignore hard gt + valid_gt_idx = best_prior_overlap[:, 0] >= 0.2 + best_prior_idx_filter = best_prior_idx[valid_gt_idx, :] + if best_prior_idx_filter.shape[0] <= 0: + loc_t[idx] = 0 + conf_t[idx] = 0 + return + + # [1,num_priors] best ground truth for each prior + best_truth_overlap, best_truth_idx = overlaps.max(0, keepdim=True) + best_truth_idx.squeeze_(0) + best_truth_overlap.squeeze_(0) + best_prior_idx.squeeze_(1) + best_prior_idx_filter.squeeze_(1) + best_prior_overlap.squeeze_(1) + best_truth_overlap.index_fill_(0, best_prior_idx_filter, 2) # ensure best prior + # TODO refactor: index best_prior_idx with long tensor + # ensure every gt matches with its prior of max overlap + for j in range(best_prior_idx.size(0)): # 判别此anchor是预测哪一个boxes + best_truth_idx[best_prior_idx[j]] = j + matches = truths[best_truth_idx] # Shape: [num_priors,4] 此处为每一个anchor对应的bbox取出来 + conf = labels[best_truth_idx] # Shape: [num_priors] 此处为每一个anchor对应的label取出来 + conf[best_truth_overlap < threshold] = 0 # label as background overlap<0.35的全部作为负样本 + loc = encode(matches, priors, variances) + + matches_landm = landms[best_truth_idx] + landm = encode_landm(matches_landm, priors, variances) + loc_t[idx] = loc # [num_priors,4] encoded offsets to learn + conf_t[idx] = conf # [num_priors] top class label for each prior + landm_t[idx] = landm + + +def encode(matched, priors, variances): + """Encode the variances from the priorbox layers into the ground truth boxes + we have matched (based on jaccard overlap) with the prior boxes. + Args: + matched: (tensor) Coords of ground truth for each prior in point-form + Shape: [num_priors, 4]. + priors: (tensor) Prior boxes in center-offset form + Shape: [num_priors,4]. + variances: (list[float]) Variances of priorboxes + Return: + encoded boxes (tensor), Shape: [num_priors, 4] + """ + + # dist b/t match center and prior's center + g_cxcy = (matched[:, :2] + matched[:, 2:])/2 - priors[:, :2] + # encode variance + g_cxcy /= (variances[0] * priors[:, 2:]) + # match wh / prior wh + g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:] + g_wh = torch.log(g_wh) / variances[1] + # return target for smooth_l1_loss + return torch.cat([g_cxcy, g_wh], 1) # [num_priors,4] + +def encode_landm(matched, priors, variances): + """Encode the variances from the priorbox layers into the ground truth boxes + we have matched (based on jaccard overlap) with the prior boxes. + Args: + matched: (tensor) Coords of ground truth for each prior in point-form + Shape: [num_priors, 10]. + priors: (tensor) Prior boxes in center-offset form + Shape: [num_priors,4]. + variances: (list[float]) Variances of priorboxes + Return: + encoded landm (tensor), Shape: [num_priors, 10] + """ + + # dist b/t match center and prior's center + matched = torch.reshape(matched, (matched.size(0), 5, 2)) + priors_cx = priors[:, 0].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2) + priors_cy = priors[:, 1].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2) + priors_w = priors[:, 2].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2) + priors_h = priors[:, 3].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2) + priors = torch.cat([priors_cx, priors_cy, priors_w, priors_h], dim=2) + g_cxcy = matched[:, :, :2] - priors[:, :, :2] + # encode variance + g_cxcy /= (variances[0] * priors[:, :, 2:]) + # g_cxcy /= priors[:, :, 2:] + g_cxcy = g_cxcy.reshape(g_cxcy.size(0), -1) + # return target for smooth_l1_loss + return g_cxcy + + +# Adapted from https://github.com/Hakuyume/chainer-ssd +def decode(loc, priors, variances): + """Decode locations from predictions using priors to undo + the encoding we did for offset regression at train time. + Args: + loc (tensor): location predictions for loc layers, + Shape: [num_priors,4] + priors (tensor): Prior boxes in center-offset form. + Shape: [num_priors,4]. + variances: (list[float]) Variances of priorboxes + Return: + decoded bounding box predictions + """ + + boxes = torch.cat(( + priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:], + priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1) + boxes[:, :2] -= boxes[:, 2:] / 2 + boxes[:, 2:] += boxes[:, :2] + return boxes + +def decode_landm(pre, priors, variances): + """Decode landm from predictions using priors to undo + the encoding we did for offset regression at train time. + Args: + pre (tensor): landm predictions for loc layers, + Shape: [num_priors,10] + priors (tensor): Prior boxes in center-offset form. + Shape: [num_priors,4]. + variances: (list[float]) Variances of priorboxes + Return: + decoded landm predictions + """ + landms = torch.cat((priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:], + priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:], + priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:], + priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:], + priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:], + ), dim=1) + return landms + + +def log_sum_exp(x): + """Utility function for computing log_sum_exp while determining + This will be used to determine unaveraged confidence loss across + all examples in a batch. + Args: + x (Variable(tensor)): conf_preds from conf layers + """ + x_max = x.data.max() + return torch.log(torch.sum(torch.exp(x-x_max), 1, keepdim=True)) + x_max + + +# Original author: Francisco Massa: +# https://github.com/fmassa/object-detection.torch +# Ported to PyTorch by Max deGroot (02/01/2017) +def nms(boxes, scores, overlap=0.5, top_k=200): + """Apply non-maximum suppression at test time to avoid detecting too many + overlapping bounding boxes for a given object. + Args: + boxes: (tensor) The location preds for the img, Shape: [num_priors,4]. + scores: (tensor) The class predscores for the img, Shape:[num_priors]. + overlap: (float) The overlap thresh for suppressing unnecessary boxes. + top_k: (int) The Maximum number of box preds to consider. + Return: + The indices of the kept boxes with respect to num_priors. + """ + + keep = torch.Tensor(scores.size(0)).fill_(0).long() + if boxes.numel() == 0: + return keep + x1 = boxes[:, 0] + y1 = boxes[:, 1] + x2 = boxes[:, 2] + y2 = boxes[:, 3] + area = torch.mul(x2 - x1, y2 - y1) + v, idx = scores.sort(0) # sort in ascending order + # I = I[v >= 0.01] + idx = idx[-top_k:] # indices of the top-k largest vals + xx1 = boxes.new() + yy1 = boxes.new() + xx2 = boxes.new() + yy2 = boxes.new() + w = boxes.new() + h = boxes.new() + + # keep = torch.Tensor() + count = 0 + while idx.numel() > 0: + i = idx[-1] # index of current largest val + # keep.append(i) + keep[count] = i + count += 1 + if idx.size(0) == 1: + break + idx = idx[:-1] # remove kept element from view + # load bboxes of next highest vals + torch.index_select(x1, 0, idx, out=xx1) + torch.index_select(y1, 0, idx, out=yy1) + torch.index_select(x2, 0, idx, out=xx2) + torch.index_select(y2, 0, idx, out=yy2) + # store element-wise max with next highest score + xx1 = torch.clamp(xx1, min=x1[i]) + yy1 = torch.clamp(yy1, min=y1[i]) + xx2 = torch.clamp(xx2, max=x2[i]) + yy2 = torch.clamp(yy2, max=y2[i]) + w.resize_as_(xx2) + h.resize_as_(yy2) + w = xx2 - xx1 + h = yy2 - yy1 + # check sizes of xx1 and xx2.. after each iteration + w = torch.clamp(w, min=0.0) + h = torch.clamp(h, min=0.0) + inter = w*h + # IoU = i / (area(a) + area(b) - i) + rem_areas = torch.index_select(area, 0, idx) # load remaining areas) + union = (rem_areas - inter) + area[i] + IoU = inter/union # store result in iou + # keep only elements with an IoU <= overlap + idx = idx[IoU.le(overlap)] + return keep, count + + diff --git a/hair_service_sd/core/utils/landmark_processor.py b/hair_service_sd/core/utils/landmark_processor.py new file mode 100644 index 0000000..2428a55 --- /dev/null +++ b/hair_service_sd/core/utils/landmark_processor.py @@ -0,0 +1,1566 @@ +import colorsys +import cv2 +import numpy as np +import random +import time +import math + +mean_face_x = np.array([ + 0.000213256, 0.0752622, 0.18113, 0.29077, 0.393397, 0.586856, 0.689483, 0.799124, + 0.904991, 0.98004, 0.490127, 0.490127, 0.490127, 0.490127, 0.36688, 0.426036, + 0.490127, 0.554217, 0.613373, 0.121737, 0.187122, 0.265825, 0.334606, 0.260918, + 0.182743, 0.645647, 0.714428, 0.793132, 0.858516, 0.79751, 0.719335, 0.254149, + 0.340985, 0.428858, 0.490127, 0.551395, 0.639268, 0.726104, 0.642159, 0.556721, + 0.490127, 0.423532, 0.338094, 0.290379, 0.428096, 0.490127, 0.552157, 0.689874, + 0.553364, 0.490127, 0.42689]) + +mean_face_y = np.array([ + 0.106454, 0.038915, 0.0187482, 0.0344891, 0.0773906, 0.0773906, 0.0344891, + 0.0187482, 0.038915, 0.106454, 0.203352, 0.307009, 0.409805, 0.515625, 0.587326, + 0.609345, 0.628106, 0.609345, 0.587326, 0.216423, 0.178758, 0.179852, 0.231733, + 0.245099, 0.244077, 0.231733, 0.179852, 0.178758, 0.216423, 0.244077, 0.245099, + 0.780233, 0.745405, 0.727388, 0.742578, 0.727388, 0.745405, 0.780233, 0.864805, + 0.902192, 0.909281, 0.902192, 0.864805, 0.784792, 0.778746, 0.785343, 0.778746, + 0.784792, 0.824182, 0.831803, 0.824182]) + +landmarks_2D = np.stack([mean_face_x, mean_face_y], axis=1) + +mean_face_x_1k = np.array([0.498047, 0.504671, 0.511286, 0.517984, 0.524451, 0.531086, 0.537574, 0.543902, 0.550305, 0.556545, 0.562891, 0.568903, 0.574994, 0.581092, 0.587026, 0.592814, 0.598603, 0.604259, 0.609802, 0.615311, 0.620821, 0.626135, 0.631542, 0.636597, 0.641621, 0.646680, 0.651656, 0.656505, 0.661374, 0.666025, 0.670571, 0.675225, 0.679616, 0.683987, 0.688181, 0.692460, 0.696634, 0.700552, 0.704500, 0.708402, 0.712261, 0.715842, 0.719373, 0.722930, 0.726486, 0.729785, 0.732875, 0.736135, 0.739115, 0.742179, 0.744925, 0.747688, 0.750314, 0.752887, 0.755412, 0.757737, 0.760038, 0.762291, 0.764312, 0.766333, 0.768259, 0.770363, 0.771888, 0.773763, 0.775485, 0.777069, 0.778656, 0.780126, 0.781483, 0.782879, 0.784085, 0.785261, 0.786402, 0.787338, 0.788398, 0.789310, 0.790038, 0.790796, 0.791281, 0.792001, 0.792397, 0.792903, 0.793221, 0.793554, 0.793613, 0.793848, 0.793856, 0.793925, 0.793911, 0.793825, 0.793601, 0.793426, 0.793199, 0.792833, 0.792528, 0.791962, 0.791501, 0.790897, 0.790382, 0.789660, 0.788884, 0.787702, 0.786300, 0.784951, 0.783363, 0.781861, 0.780103, 0.778348, 0.776405, 0.774489, 0.772390, 0.770098, 0.767540, 0.765030, 0.762359, 0.759349, 0.756513, 0.753197, 0.749784, 0.746353, 0.742572, 0.738633, 0.734433, 0.730200, 0.725658, 0.720657, 0.715711, 0.710406, 0.705108, 0.699482, 0.693569, 0.687476, 0.681209, 0.674681, 0.668094, 0.661239, 0.654234, 0.647131, 0.639989, 0.632582, 0.625134, 0.617439, 0.610009, 0.602125, 0.594318, 0.586546, 0.578466, 0.570582, 0.562638, 0.554622, 0.546602, 0.538546, 0.530536, 0.522423, 0.514258, 0.506205, 0.498047, 0.489889, 0.481836, 0.473671, 0.465558, 0.457548, 0.449491, 0.441472, 0.433455, 0.425511, 0.417628, 0.409547, 0.401776, 0.393969, 0.386085, 0.378655, 0.370960, 0.363512, 0.356105, 0.348963, 0.341860, 0.334855, 0.328000, 0.321412, 0.314885, 0.308618, 0.302525, 0.296612, 0.290986, 0.285687, 0.280382, 0.275437, 0.270436, 0.265894, 0.261660, 0.257461, 0.253522, 0.249740, 0.246310, 0.242897, 0.239580, 0.236744, 0.233735, 0.231064, 0.228553, 0.225996, 0.223704, 0.221605, 0.219689, 0.217746, 0.215990, 0.214232, 0.212730, 0.211143, 0.209794, 0.208392, 0.207210, 0.206433, 0.205712, 0.205197, 0.204593, 0.204132, 0.203565, 0.203261, 0.202895, 0.202667, 0.202492, 0.202268, 0.202183, 0.202169, 0.202238, 0.202246, 0.202480, 0.202540, 0.202873, 0.203191, 0.203696, 0.204093, 0.204813, 0.205297, 0.206055, 0.206783, 0.207696, 0.208755, 0.209692, 0.210833, 0.212009, 0.213215, 0.214611, 0.215967, 0.217438, 0.219024, 0.220609, 0.222330, 0.224206, 0.225731, 0.227835, 0.229761, 0.231781, 0.233803, 0.236055, 0.238356, 0.240681, 0.243207, 0.245780, 0.248405, 0.251168, 0.253915, 0.256979, 0.259958, 0.263219, 0.266309, 0.269608, 0.273163, 0.276721, 0.280252, 0.283833, 0.287692, 0.291593, 0.295541, 0.299459, 0.303633, 0.307913, 0.312107, 0.316478, 0.320869, 0.325523, 0.330069, 0.334719, 0.339589, 0.344438, 0.349414, 0.354473, 0.359497, 0.364552, 0.369959, 0.375273, 0.380783, 0.386292, 0.391835, 0.397491, 0.403280, 0.409068, 0.415002, 0.421100, 0.427190, 0.433203, 0.439548, 0.445789, 0.452191, 0.458520, 0.465008, 0.471643, 0.478110, 0.484807, 0.491422, 0.396255, 0.397888, 0.399588, 0.401344, 0.403150, 0.405001, 0.406895, 0.408831, 0.410807, 0.412822, 0.414875, 0.416966, 0.419095, 0.421262, 0.423465, 0.425705, 0.427981, 0.430294, 0.432644, 0.435030, 0.437451, 0.439908, 0.442400, 0.444928, 0.447489, 0.450084, 0.452712, 0.455372, 0.458062, 0.460782, 0.463531, 0.466307, 0.469108, 0.471933, 0.474781, 0.477648, 0.480532, 0.483432, 0.486345, 0.489267, 0.492194, 0.495123, 0.498047, 0.500971, 0.503899, 0.506827, 0.509749, 0.512662, 0.515562, 0.518446, 0.521313, 0.524160, 0.526986, 0.529787, 0.532563, 0.535312, 0.538032, 0.540722, 0.543382, 0.546009, 0.548605, 0.551166, 0.553693, 0.556186, 0.558642, 0.561064, 0.563450, 0.565799, 0.568113, 0.570389, 0.572629, 0.574832, 0.576998, 0.579127, 0.581219, 0.583272, 0.585287, 0.587263, 0.589198, 0.591093, 0.592944, 0.594749, 0.596506, 0.598206, 0.599839, 0.597799, 0.595660, 0.593444, 0.591167, 0.588838, 0.586464, 0.584049, 0.581599, 0.579114, 0.576598, 0.574052, 0.571476, 0.568873, 0.566242, 0.563585, 0.560901, 0.558193, 0.555461, 0.552704, 0.549922, 0.547116, 0.544286, 0.541432, 0.538555, 0.535655, 0.532732, 0.529786, 0.526819, 0.523832, 0.520830, 0.517033, 0.513236, 0.509439, 0.505641, 0.501844, 0.498047, 0.494250, 0.490452, 0.486655, 0.482858, 0.479061, 0.475264, 0.472262, 0.469275, 0.466308, 0.463362, 0.460439, 0.457538, 0.454662, 0.451808, 0.448978, 0.446172, 0.443390, 0.440633, 0.437900, 0.435192, 0.432509, 0.429852, 0.427221, 0.424617, 0.422042, 0.419496, 0.416979, 0.414495, 0.412044, 0.409630, 0.407256, 0.404927, 0.402650, 0.400434, 0.398294, 0.410228, 0.414399, 0.418630, 0.422899, 0.427196, 0.431515, 0.435855, 0.440212, 0.444587, 0.448978, 0.453385, 0.457806, 0.462242, 0.466691, 0.471152, 0.475622, 0.480101, 0.484586, 0.489074, 0.493564, 0.498047, 0.502530, 0.507020, 0.511508, 0.515993, 0.520472, 0.524942, 0.529403, 0.533852, 0.538288, 0.542709, 0.547116, 0.551507, 0.555882, 0.560238, 0.564578, 0.568898, 0.573194, 0.577464, 0.581695, 0.585866, 0.581642, 0.577369, 0.573066, 0.568743, 0.564401, 0.560044, 0.555673, 0.551287, 0.546890, 0.542478, 0.538055, 0.533622, 0.529181, 0.524733, 0.520280, 0.515825, 0.511370, 0.506919, 0.502475, 0.498047, 0.493618, 0.489175, 0.484723, 0.480269, 0.475814, 0.471361, 0.466913, 0.462472, 0.458039, 0.453616, 0.449204, 0.444806, 0.440421, 0.436050, 0.431693, 0.427351, 0.423027, 0.418725, 0.414452, 0.459656, 0.459227, 0.458691, 0.458058, 0.457338, 0.456539, 0.455664, 0.454719, 0.453708, 0.452633, 0.451497, 0.450302, 0.449049, 0.447741, 0.446378, 0.444962, 0.443494, 0.441973, 0.440402, 0.438780, 0.437107, 0.435384, 0.433612, 0.431790, 0.429919, 0.427998, 0.426028, 0.424009, 0.421938, 0.419814, 0.417632, 0.415384, 0.442766, 0.471804, 0.498047, 0.524290, 0.553327, 0.580709, 0.578462, 0.576279, 0.574156, 0.572085, 0.570065, 0.568095, 0.566175, 0.564304, 0.562481, 0.560709, 0.558987, 0.557314, 0.555692, 0.554120, 0.552600, 0.551131, 0.549715, 0.548353, 0.547044, 0.545792, 0.544597, 0.543461, 0.542386, 0.541375, 0.540430, 0.539555, 0.538755, 0.538036, 0.537403, 0.536867, 0.536438, 0.551823, 0.524746, 0.471347, 0.444271, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.370401, 0.394341, 0.393977, 0.392891, 0.391151, 0.388793, 0.385710, 0.382329, 0.378604, 0.374564, 0.370390, 0.366216, 0.362177, 0.358453, 0.355074, 0.351995, 0.349640, 0.347903, 0.346822, 0.346461, 0.346826, 0.347911, 0.349651, 0.352010, 0.355092, 0.358473, 0.362199, 0.366239, 0.370413, 0.374587, 0.378626, 0.382350, 0.385728, 0.388808, 0.391163, 0.392899, 0.393981, 0.423391, 0.421623, 0.419536, 0.417200, 0.414654, 0.411925, 0.409036, 0.406003, 0.402844, 0.399571, 0.396199, 0.392740, 0.389206, 0.385609, 0.381960, 0.378271, 0.374552, 0.370812, 0.367063, 0.363314, 0.359575, 0.355859, 0.352173, 0.348530, 0.344945, 0.341429, 0.337998, 0.334669, 0.331464, 0.328407, 0.325532, 0.322889, 0.320595, 0.322717, 0.325103, 0.327689, 0.330440, 0.333327, 0.336328, 0.339426, 0.342603, 0.345848, 0.349146, 0.352488, 0.355863, 0.359264, 0.362683, 0.366116, 0.369556, 0.372999, 0.376442, 0.379882, 0.383316, 0.386742, 0.390158, 0.393563, 0.396955, 0.400332, 0.403694, 0.407038, 0.410363, 0.413666, 0.416943, 0.420190, 0.625692, 0.649633, 0.649272, 0.648190, 0.646454, 0.644099, 0.641019, 0.637641, 0.633917, 0.629878, 0.625704, 0.621530, 0.617490, 0.613764, 0.610383, 0.607301, 0.604942, 0.603202, 0.602117, 0.601752, 0.602113, 0.603194, 0.604931, 0.607286, 0.610365, 0.613744, 0.617468, 0.621507, 0.625681, 0.629855, 0.633895, 0.637621, 0.641001, 0.644084, 0.646442, 0.648182, 0.649268, 0.572703, 0.574471, 0.576557, 0.578893, 0.581440, 0.584168, 0.587058, 0.590090, 0.593250, 0.596523, 0.599895, 0.603354, 0.606888, 0.610485, 0.614134, 0.617823, 0.621542, 0.625282, 0.629031, 0.632780, 0.636518, 0.640235, 0.643921, 0.647563, 0.651149, 0.654665, 0.658096, 0.661425, 0.664630, 0.667686, 0.670562, 0.673205, 0.675499, 0.673377, 0.670991, 0.668404, 0.665654, 0.662767, 0.659765, 0.656668, 0.653491, 0.650246, 0.646948, 0.643606, 0.640231, 0.636830, 0.633411, 0.629978, 0.626538, 0.623095, 0.619652, 0.616212, 0.612778, 0.609352, 0.605936, 0.602531, 0.599139, 0.595762, 0.592400, 0.589056, 0.585731, 0.582428, 0.579150, 0.575904, 0.552267, 0.555073, 0.558014, 0.561313, 0.565180, 0.569636, 0.574564, 0.579836, 0.585356, 0.591050, 0.596857, 0.602735, 0.608652, 0.614591, 0.620542, 0.626500, 0.632462, 0.638425, 0.644390, 0.649985, 0.655581, 0.661175, 0.666765, 0.672343, 0.677894, 0.683391, 0.688788, 0.694021, 0.699020, 0.703737, 0.708170, 0.712362, 0.716374, 0.720268, 0.724091, 0.727877, 0.731647, 0.726965, 0.722281, 0.717595, 0.712902, 0.708202, 0.703489, 0.698762, 0.694017, 0.689253, 0.684472, 0.679678, 0.674876, 0.670068, 0.665257, 0.660446, 0.655633, 0.650821, 0.646008, 0.640808, 0.635608, 0.630406, 0.625203, 0.619997, 0.614788, 0.609577, 0.604360, 0.599139, 0.593915, 0.588693, 0.583472, 0.578256, 0.573043, 0.567837, 0.562637, 0.557446, 0.264446, 0.268217, 0.272003, 0.275826, 0.279720, 0.283732, 0.287924, 0.292357, 0.297074, 0.302073, 0.307305, 0.312703, 0.318199, 0.323751, 0.329329, 0.334919, 0.340513, 0.346109, 0.351704, 0.357669, 0.363632, 0.369594, 0.375552, 0.381503, 0.387442, 0.393359, 0.399237, 0.405044, 0.410737, 0.416258, 0.421530, 0.426457, 0.430914, 0.434781, 0.438079, 0.441020, 0.443827, 0.438648, 0.433457, 0.428257, 0.423050, 0.417838, 0.412621, 0.407401, 0.402179, 0.396955, 0.391734, 0.386517, 0.381305, 0.376096, 0.370891, 0.365688, 0.360486, 0.355285, 0.350085, 0.345273, 0.340460, 0.335648, 0.330836, 0.326026, 0.321218, 0.316416, 0.311622, 0.306841, 0.302077, 0.297332, 0.292605, 0.287892, 0.283191, 0.278499, 0.273812, 0.269128 +]) +mean_face_y_1k = np.array([0.851392, 0.851204, 0.850802, 0.849951, 0.849050, 0.847896, 0.846372, 0.844708, 0.842887, 0.840805, 0.838470, 0.836093, 0.833361, 0.830589, 0.827781, 0.824650, 0.821361, 0.817917, 0.814463, 0.810803, 0.807127, 0.803118, 0.799272, 0.795088, 0.790897, 0.786585, 0.782195, 0.777726, 0.773128, 0.768407, 0.763595, 0.758801, 0.753939, 0.748915, 0.743745, 0.738664, 0.733510, 0.728055, 0.722725, 0.717442, 0.711952, 0.706343, 0.700725, 0.695036, 0.689315, 0.683557, 0.677857, 0.671952, 0.666097, 0.659596, 0.652993, 0.646386, 0.639687, 0.632843, 0.626218, 0.619403, 0.612554, 0.605551, 0.598705, 0.591735, 0.584939, 0.577854, 0.571094, 0.564026, 0.557009, 0.549978, 0.542887, 0.535838, 0.528892, 0.521764, 0.514596, 0.507444, 0.500318, 0.493164, 0.486134, 0.478851, 0.471725, 0.464586, 0.457279, 0.450071, 0.443049, 0.435691, 0.428445, 0.421397, 0.414262, 0.406894, 0.399680, 0.392547, 0.385399, 0.378182, 0.370759, 0.363690, 0.356390, 0.349346, 0.341952, 0.335000, 0.327902, 0.320707, 0.313459, 0.306325, 0.299213, 0.291181, 0.283309, 0.275338, 0.267185, 0.259441, 0.251545, 0.243652, 0.235765, 0.228054, 0.220194, 0.212369, 0.204820, 0.197130, 0.189558, 0.182061, 0.174602, 0.167242, 0.159851, 0.152736, 0.145606, 0.138595, 0.131706, 0.125067, 0.118485, 0.112101, 0.105897, 0.099812, 0.094088, 0.088477, 0.082962, 0.077854, 0.072924, 0.068273, 0.063748, 0.059715, 0.055721, 0.052042, 0.048762, 0.045388, 0.042616, 0.039888, 0.037449, 0.035191, 0.033201, 0.031339, 0.029751, 0.028369, 0.027112, 0.025987, 0.025155, 0.024428, 0.023762, 0.023446, 0.023057, 0.022894, 0.022946, 0.022894, 0.023057, 0.023446, 0.023762, 0.024428, 0.025155, 0.025987, 0.027112, 0.028369, 0.029751, 0.031339, 0.033201, 0.035191, 0.037449, 0.039888, 0.042616, 0.045388, 0.048762, 0.052042, 0.055721, 0.059715, 0.063748, 0.068273, 0.072924, 0.077854, 0.082962, 0.088477, 0.094088, 0.099812, 0.105897, 0.112101, 0.118485, 0.125067, 0.131706, 0.138595, 0.145606, 0.152736, 0.159851, 0.167242, 0.174602, 0.182061, 0.189558, 0.197130, 0.204820, 0.212369, 0.220194, 0.228054, 0.235765, 0.243652, 0.251545, 0.259441, 0.267185, 0.275338, 0.283309, 0.291181, 0.299213, 0.306325, 0.313459, 0.320707, 0.327902, 0.335000, 0.341952, 0.349346, 0.356390, 0.363690, 0.370759, 0.378182, 0.385399, 0.392547, 0.399680, 0.406894, 0.414262, 0.421397, 0.428445, 0.435691, 0.443049, 0.450071, 0.457279, 0.464586, 0.471725, 0.478851, 0.486134, 0.493164, 0.500318, 0.507444, 0.514596, 0.521764, 0.528892, 0.535838, 0.542887, 0.549978, 0.557009, 0.564026, 0.571094, 0.577854, 0.584939, 0.591735, 0.598705, 0.605551, 0.612554, 0.619403, 0.626218, 0.632843, 0.639687, 0.646386, 0.652993, 0.659596, 0.666097, 0.671952, 0.677857, 0.683557, 0.689315, 0.695036, 0.700725, 0.706343, 0.711952, 0.717442, 0.722725, 0.728055, 0.733510, 0.738664, 0.743745, 0.748915, 0.753939, 0.758801, 0.763595, 0.768407, 0.773128, 0.777726, 0.782195, 0.786585, 0.790897, 0.795088, 0.799272, 0.803118, 0.807127, 0.810803, 0.814463, 0.817917, 0.821361, 0.824650, 0.827781, 0.830589, 0.833361, 0.836093, 0.838470, 0.840805, 0.842887, 0.844708, 0.846372, 0.847896, 0.849050, 0.849951, 0.850802, 0.851204, 0.664237, 0.666525, 0.668791, 0.671031, 0.673242, 0.675423, 0.677571, 0.679687, 0.681769, 0.683818, 0.685829, 0.687805, 0.689742, 0.691641, 0.693499, 0.695315, 0.697088, 0.698816, 0.700499, 0.702133, 0.703718, 0.705251, 0.706730, 0.708155, 0.709522, 0.710829, 0.712076, 0.713258, 0.714375, 0.715424, 0.716404, 0.717311, 0.718144, 0.718900, 0.719577, 0.720171, 0.720681, 0.721102, 0.721432, 0.721666, 0.721798, 0.721824, 0.721732, 0.721824, 0.721798, 0.721666, 0.721432, 0.721102, 0.720681, 0.720171, 0.719577, 0.718900, 0.718144, 0.717311, 0.716404, 0.715424, 0.714375, 0.713258, 0.712076, 0.710829, 0.709522, 0.708155, 0.706730, 0.705251, 0.703718, 0.702133, 0.700499, 0.698816, 0.697088, 0.695315, 0.693499, 0.691641, 0.689742, 0.687805, 0.685829, 0.683818, 0.681769, 0.679687, 0.677571, 0.675423, 0.673242, 0.671031, 0.668791, 0.666525, 0.664237, 0.662342, 0.660514, 0.658743, 0.657022, 0.655347, 0.653715, 0.652122, 0.650567, 0.649048, 0.647565, 0.646117, 0.644703, 0.643326, 0.641984, 0.640678, 0.639409, 0.638178, 0.636987, 0.635837, 0.634730, 0.633668, 0.632655, 0.631695, 0.630791, 0.629948, 0.629175, 0.628479, 0.627873, 0.627375, 0.627012, 0.628162, 0.629312, 0.630462, 0.631611, 0.632761, 0.633911, 0.632761, 0.631611, 0.630462, 0.629312, 0.628162, 0.627012, 0.627375, 0.627873, 0.628479, 0.629175, 0.629948, 0.630791, 0.631695, 0.632655, 0.633668, 0.634730, 0.635837, 0.636987, 0.638178, 0.639409, 0.640678, 0.641984, 0.643326, 0.644703, 0.646117, 0.647565, 0.649048, 0.650567, 0.652122, 0.653715, 0.655347, 0.657022, 0.658743, 0.660514, 0.662342, 0.665075, 0.665788, 0.666477, 0.667149, 0.667807, 0.668452, 0.669084, 0.669701, 0.670303, 0.670890, 0.671461, 0.672014, 0.672548, 0.673062, 0.673552, 0.674017, 0.674454, 0.674860, 0.675234, 0.675572, 0.675871, 0.675572, 0.675234, 0.674860, 0.674454, 0.674017, 0.673552, 0.673062, 0.672548, 0.672014, 0.671461, 0.670890, 0.670303, 0.669701, 0.669084, 0.668452, 0.667807, 0.667149, 0.666477, 0.665788, 0.665075, 0.664381, 0.663748, 0.663164, 0.662625, 0.662130, 0.661679, 0.661275, 0.660920, 0.660618, 0.660372, 0.660186, 0.660063, 0.660004, 0.660012, 0.660086, 0.660229, 0.660443, 0.660728, 0.661089, 0.661532, 0.661089, 0.660728, 0.660443, 0.660229, 0.660086, 0.660012, 0.660004, 0.660063, 0.660186, 0.660372, 0.660618, 0.660920, 0.661275, 0.661679, 0.662130, 0.662625, 0.663164, 0.663748, 0.664381, 0.375535, 0.380755, 0.385970, 0.391178, 0.396378, 0.401568, 0.406749, 0.411919, 0.417077, 0.422223, 0.427357, 0.432478, 0.437586, 0.442679, 0.447758, 0.452822, 0.457871, 0.462903, 0.467918, 0.472916, 0.477896, 0.482858, 0.487799, 0.492722, 0.497623, 0.502504, 0.507362, 0.512198, 0.517010, 0.521792, 0.526540, 0.531243, 0.559981, 0.564298, 0.572818, 0.564298, 0.559981, 0.531243, 0.526540, 0.521792, 0.517010, 0.512198, 0.507362, 0.502504, 0.497623, 0.492722, 0.487799, 0.482858, 0.477896, 0.472916, 0.467918, 0.462903, 0.457871, 0.452822, 0.447758, 0.442679, 0.437586, 0.432478, 0.427357, 0.422223, 0.417077, 0.411919, 0.406749, 0.401568, 0.396378, 0.391178, 0.385970, 0.380755, 0.375535, 0.543598, 0.549005, 0.549005, 0.543598, 0.521853, 0.516702, 0.511551, 0.506400, 0.501249, 0.496098, 0.490947, 0.485796, 0.480645, 0.475494, 0.470343, 0.465192, 0.460041, 0.454890, 0.449739, 0.444588, 0.439437, 0.434286, 0.429135, 0.423984, 0.418833, 0.413682, 0.408531, 0.403380, 0.398229, 0.393078, 0.387927, 0.382776, 0.377625, 0.372474, 0.367323, 0.362172, 0.357021, 0.362533, 0.362584, 0.366719, 0.370758, 0.374482, 0.377860, 0.380939, 0.383295, 0.385031, 0.386113, 0.386473, 0.386109, 0.385023, 0.383283, 0.380925, 0.377842, 0.374461, 0.370736, 0.366696, 0.362560, 0.358348, 0.354309, 0.350585, 0.347206, 0.344127, 0.341772, 0.340035, 0.338954, 0.338593, 0.338958, 0.340043, 0.341783, 0.344142, 0.347224, 0.350605, 0.354331, 0.358371, 0.375192, 0.372191, 0.369350, 0.366674, 0.364163, 0.361818, 0.359640, 0.357632, 0.355795, 0.354130, 0.352640, 0.351325, 0.350186, 0.349223, 0.348437, 0.347827, 0.347393, 0.347135, 0.347052, 0.347144, 0.347412, 0.347854, 0.348473, 0.349267, 0.350238, 0.351387, 0.352715, 0.354223, 0.355912, 0.357785, 0.359846, 0.362104, 0.364585, 0.367095, 0.369405, 0.371520, 0.373447, 0.375190, 0.376755, 0.378148, 0.379374, 0.380440, 0.381352, 0.382117, 0.382743, 0.383236, 0.383605, 0.383856, 0.383996, 0.384031, 0.383968, 0.383811, 0.383568, 0.383241, 0.382836, 0.382357, 0.381806, 0.381188, 0.380505, 0.379760, 0.378955, 0.378093, 0.377176, 0.376206, 0.362533, 0.362560, 0.366696, 0.370736, 0.374461, 0.377842, 0.380925, 0.383283, 0.385023, 0.386109, 0.386473, 0.386113, 0.385031, 0.383295, 0.380939, 0.377860, 0.374482, 0.370758, 0.366719, 0.362584, 0.358371, 0.354331, 0.350605, 0.347224, 0.344142, 0.341783, 0.340043, 0.338958, 0.338593, 0.338954, 0.340035, 0.341772, 0.344127, 0.347206, 0.350585, 0.354309, 0.358348, 0.375192, 0.372191, 0.369350, 0.366674, 0.364163, 0.361818, 0.359640, 0.357632, 0.355795, 0.354130, 0.352640, 0.351325, 0.350186, 0.349223, 0.348437, 0.347827, 0.347393, 0.347135, 0.347052, 0.347144, 0.347412, 0.347854, 0.348473, 0.349267, 0.350238, 0.351387, 0.352715, 0.354223, 0.355912, 0.357785, 0.359846, 0.362104, 0.364585, 0.367095, 0.369405, 0.371520, 0.373447, 0.375190, 0.376755, 0.378148, 0.379374, 0.380440, 0.381352, 0.382117, 0.382743, 0.383236, 0.383605, 0.383856, 0.383996, 0.384031, 0.383968, 0.383811, 0.383568, 0.383241, 0.382836, 0.382357, 0.381806, 0.381188, 0.380505, 0.379760, 0.378955, 0.378093, 0.377176, 0.376206, 0.298620, 0.293794, 0.288991, 0.284301, 0.279916, 0.276030, 0.272714, 0.269943, 0.267662, 0.265804, 0.264290, 0.263036, 0.261967, 0.261022, 0.260156, 0.259339, 0.258551, 0.257778, 0.257011, 0.257406, 0.257815, 0.258259, 0.258769, 0.259392, 0.260193, 0.261257, 0.262688, 0.264582, 0.266996, 0.269911, 0.273238, 0.276859, 0.280671, 0.284596, 0.288585, 0.292606, 0.296640, 0.296323, 0.296006, 0.295688, 0.295372, 0.295058, 0.294752, 0.294461, 0.294192, 0.293955, 0.293756, 0.293596, 0.293472, 0.293376, 0.293299, 0.293234, 0.293177, 0.293123, 0.293071, 0.293667, 0.294260, 0.294850, 0.295430, 0.295996, 0.296538, 0.297043, 0.297496, 0.297883, 0.298195, 0.298431, 0.298597, 0.298704, 0.298760, 0.298773, 0.298749, 0.298694, 0.296640, 0.292606, 0.288585, 0.284596, 0.280671, 0.276859, 0.273238, 0.269911, 0.266996, 0.264582, 0.262688, 0.261257, 0.260193, 0.259392, 0.258769, 0.258259, 0.257815, 0.257406, 0.257011, 0.257778, 0.258551, 0.259339, 0.260156, 0.261022, 0.261967, 0.263036, 0.264290, 0.265804, 0.267662, 0.269943, 0.272714, 0.276030, 0.279916, 0.284301, 0.288991, 0.293794, 0.298620, 0.298694, 0.298749, 0.298773, 0.298760, 0.298704, 0.298597, 0.298431, 0.298195, 0.297883, 0.297496, 0.297043, 0.296538, 0.295996, 0.295430, 0.294850, 0.294260, 0.293667, 0.293071, 0.293123, 0.293177, 0.293234, 0.293299, 0.293376, 0.293472, 0.293596, 0.293756, 0.293955, 0.294192, 0.294461, 0.294752, 0.295058, 0.295372, 0.295688, 0.296006, 0.296323 +]) +landmarks_2D_1k = np.stack([mean_face_x_1k, mean_face_y_1k], axis=1) + +mean_face_x_137_22_client = np.array([0.4988282, 0.5449964, 0.5849726, 0.6179804, 0.643469, 0.6622178, 0.6730388, 0.676355, 0.6733304, + 0.6478118, + 0.5882786, + 0.4988282, + 0.4093778, + 0.349844, + 0.324326, + 0.32130139999999996, + 0.3246176, + 0.33543860000000003, + 0.35418740000000004, + 0.3796754, + 0.4126838, + 0.45266]) + +mean_face_y_137_22_client = np.array([0.7108352, 0.7000166, 0.6745382, 0.640106, 0.5996582, 0.5467124, 0.4916804, 0.4355282, 0.3795278, 0.2916416, + 0.23122520000000002, + 0.2137676, + 0.23122520000000002, + 0.2916416, + 0.3795278, + 0.4355282, + 0.4916804, + 0.5467124, + 0.5996582, + 0.640106, + 0.6745382, + 0.7000166]) + +landmarks_2D_137_22_clinet = np.stack([mean_face_x_137_22_client, mean_face_y_137_22_client], axis=1) + + +mat_face1024_256_full_face_client = np.array([[4.1666666e-01, 1.5257437e-17, -8.5333336e+01], + [-1.5257449e-17, 4.1666666e-01, -8.5333336e+01]]) + +mat_face1024_256_full_face_server = np.array([[4.1666666e-01, -1.5237085e-17, -8.5333336e+01], + [1.5237085e-17, 4.1666666e-01, -8.5333336e+01]]) + + +# 68 point landmark definitions +landmarks_68_pt = {"mouth": (48, 68), + "right_eyebrow": (17, 22), + "left_eyebrow": (22, 27), + "right_eye": (36, 42), + "left_eye": (42, 48), + "nose": (27, 36), # missed one point + "jaw": (0, 17)} + +def get_transform_mat_mmcv(landmark, output_size): + dst_size = output_size + + if len(landmark) == 68: + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[36] + landmark[39]) / 2 + right_eye = (landmark[42] + landmark[45]) / 2 + nose = (landmark[31] + landmark[35]) / 2 + left_mouth = (landmark[48] + landmark[60]) / 2 + right_mouth = (landmark[64] + landmark[54]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + elif len(landmark) == 87: + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[31] + landmark[35]) / 2 + right_eye = (landmark[39] + landmark[43]) / 2 + nose = landmark[62] + left_mouth = (landmark[66] + landmark[79]) / 2 + right_mouth = (landmark[72] + landmark[83]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + elif len(landmark) == 1000: + pt137 = pts_1k_to_137(landmark) + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (pt137[88] + pt137[96]) / 2 + right_eye = (pt137[105] + pt137[113]) / 2 + nose = pt137[83] + left_mouth = (pt137[22] + pt137[48]) / 2 + right_mouth = (pt137[56] + pt137[36]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + + +def umeyama(src, dst, estimate_scale): + """Estimate N-D similarity transformation with or without scaling. + Parameters + ---------- + src : (M, N) array + Source coordinates. + dst : (M, N) array + Destination coordinates. + estimate_scale : bool + Whether to estimate scaling factor. + Returns + ------- + T : (N + 1, N + 1) + The homogeneous similarity transformation matrix. The matrix contains + NaN values only if the problem is not well-conditioned. + References + ---------- + .. [1] "Least-squares estimation of transformation parameters between two + point patterns", Shinji Umeyama, PAMI 1991, DOI: 10.1109/34.88573 + """ + + num = src.shape[0] + dim = src.shape[1] + + # Compute mean of src and dst. + src_mean = src.mean(axis=0) + dst_mean = dst.mean(axis=0) + + # Subtract mean from src and dst. + src_demean = src - src_mean + dst_demean = dst - dst_mean + + # Eq. (38). + A = np.dot(dst_demean.T, src_demean) / num + + # Eq. (39). + d = np.ones((dim,), dtype=np.double) + if np.linalg.det(A) < 0: + d[dim - 1] = -1 + + T = np.eye(dim + 1, dtype=np.double) + + U, S, V = np.linalg.svd(A) + + # Eq. (40) and (43). + rank = np.linalg.matrix_rank(A) + if rank == 0: + return np.nan * T + elif rank == dim - 1: + if np.linalg.det(U) * np.linalg.det(V) > 0: + T[:dim, :dim] = np.dot(U, V) + else: + s = d[dim - 1] + d[dim - 1] = -1 + T[:dim, :dim] = np.dot(U, np.dot(np.diag(d), V)) + d[dim - 1] = s + else: + T[:dim, :dim] = np.dot(U, np.dot(np.diag(d), V.T)) + + if estimate_scale: + # Eq. (41) and (42). + scale = 1.0 / src_demean.var(axis=0).sum() * np.dot(S, d) + else: + scale = 1.0 + + T[:dim, dim] = dst_mean - scale * np.dot(T[:dim, :dim], src_mean.T) + T[:dim, :dim] *= scale + + return T + +def get_transform_mat_mmcv_bigger(landmark, output_size, forlabel=False): + dst_size = output_size + if len(landmark) == 87: + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[31] + landmark[35]) / 2 + right_eye = (landmark[39] + landmark[43]) / 2 + nose = landmark[62] + left_mouth = (landmark[66] + landmark[79]) / 2 + right_mouth = (landmark[72] + landmark[83]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + elif len(landmark) == 137: + if forlabel: + eye_dis = 0.4 + mouth_dis = 0.4 + g_Average_5point_180 = np.array([ + eye_dis, 0.4, + 1 - eye_dis, 0.4, + 0.5, 0.5, + mouth_dis, 0.6, + 1 - mouth_dis, 0.6 + ]) + else: + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[96] + landmark[88]) / 2 + right_eye = (landmark[105] + landmark[113]) / 2 + nose = landmark[83] + left_mouth = (landmark[22] + landmark[48]) / 2 + right_mouth = (landmark[56] + landmark[36]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + elif len(landmark) == 1000: + # left_eye = (landmarks_2D_1k[691] + landmarks_2D_1k[723]) / 2 + # right_eye = (landmarks_2D_1k[792] + landmarks_2D_1k[824]) / 2 + # nose = landmarks_2D_1k[621] + # left_mouth = (landmarks_2D_1k[467] + landmarks_2D_1k[468]) / 2 + # right_mouth = (landmarks_2D_1k[396] + landmarks_2D_1k[508]) / 2 + # pts5_dst = np.vstack((left_eye, right_eye, + # nose, + # left_mouth, right_mouth)) * dst_size + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + # print(g_Average_5point_180) + left_eye = (landmark[691] + landmark[723]) / 2 + right_eye = (landmark[792] + landmark[824]) / 2 + nose = landmark[621] + left_mouth = (landmark[467] + landmark[468]) / 2 + right_mouth = (landmark[396] + landmark[508]) / 2 + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + + image_to_face_mat = umeyama(pts5_src, pts5_dst, True)[:2] + + return image_to_face_mat + +def get_transform_mat_full_face(landmark, output_size): + dst_size = output_size + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k) + + # print("landmarks_2D_137 x: ", landmarks_2D_137[:22, 0]) + # print("landmarks_2D_137 y: ", landmarks_2D_137[:22, 1]) + + # exit() + + mat = umeyama(landmark[:22], landmarks_2D_137[:22] * dst_size, True)[0:2] + return mat + + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k[:312] * dst_size, True)[0:2] + return mat + + elif len(landmark) == 22: + landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k) + + mat = umeyama(landmark, landmarks_2D_137[:22] * dst_size, True)[0:2] + return mat + +def get_transform_mat_full_face_592(landmark, output_size, ratio): + dst_size = output_size + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + 0.5 + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.5) * ratio + 0.5 + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_137[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + elif len(landmark) == 22: + landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark, landmarks_2D_137[:22] * dst_size, True)[0:2] + return mat + +def get_transform_mat_full_face_592_v1(landmark, output_size, ratio=1.0, h_ratio=0.57): + dst_size = output_size + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + 0.5 + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.5) * ratio + h_ratio + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_137[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_full_face_592_client(landmark, output_size, ratio): + dst_size = output_size + landmarks_2D_137_22_clinet_tmp = landmarks_2D_137_22_clinet.copy() + landmarks_2D_137_22_clinet_tmp[:, 0] = (landmarks_2D_137_22_clinet_tmp[:, 0] - 0.5) * ratio + 0.5 + landmarks_2D_137_22_clinet_tmp[:, 1] = (landmarks_2D_137_22_clinet_tmp[:, 1] - 0.5) * ratio + 0.5 + + if len(landmark) == 22: + # landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark, landmarks_2D_137_22_clinet_tmp * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + landmark_137 = pts_1k_to_137(landmark) + + mat = umeyama(landmark_137[:22], landmarks_2D_137_22_clinet_tmp * dst_size, True)[0:2] + return mat + +def get_transform_mat_face592_full_face(img_size, detect_single_face_size, ratio): + + # face_landmark_client = landmarks_2D_137_22_clinet * img_size + # + # #TODO:客户端如何得到face 0.6 + # face_size2face_M_ratio = get_transform_mat_full_face_592_client(face_landmark_client, img_size, ratio) + + #服务端得到face 1 + face_landmark_server = landmarks_2D_1k * img_size + + face_size2face_M_ratio = get_transform_mat_full_face_592(face_landmark_server, img_size, ratio) + face_size2face_M_full = get_transform_mat_full_face(face_landmark_server, detect_single_face_size) + + + M_ori = np.zeros((3, 3), dtype=np.float32) + M_ori[:2, :] = cv2.invertAffineTransform(face_size2face_M_ratio) + M_ori[2:, :] = [0, 0, 1] + + matAffine_ori = np.zeros((3, 3), dtype=np.float32) + matAffine_ori[:2, :] = face_size2face_M_full + matAffine_ori[2:, :] = [0, 0, 1] + + new_mat = matAffine_ori.dot(M_ori) + + if False: + img = np.zeros((img_size, img_size, 3), dtype=np.uint8) + pred_label_int = face_landmark_server.copy().astype(np.int32) + img_client = img.copy() + for pt in pred_label_int: + cv2.circle(img_client, (pt[0], pt[1]), 1, (0, 0, 255), 1) + # cv2.imshow("img_client: ", img_client) + # + # img_server = np.zeros((img_size, img_size, 3), dtype=np.uint8) + # face_sever_22 = pts_1k_to_137(face_landmark_server)[:22].astype(np.int32) + # for pt in face_sever_22: + # cv2.circle(img_server, (pt[0], pt[1]), 1, (0, 255, 0), 1) + # cv2.imshow("img_server: ", img_server) + # cv2.waitKey() + # + # img_server_new = np.zeros((img_size, img_size, 3), dtype=np.uint8) + # face_sever_new = transform_points(face_landmark_server, face_size2face_M_full) + # face_sever_new_22 = pts_1k_to_137(face_sever_new)[:22].astype(np.int32) + # + # for pt in face_sever_new_22: + # cv2.circle(img_server_new, (pt[0], pt[1]), 1, (0, 255, 0), 1) + # cv2.imshow("img_server_new: ", img_server_new) + # cv2.waitKey() + face_landmark_server_new = transform_points(face_landmark_server, face_size2face_M_ratio) + img_server = np.zeros((img_size, img_size, 3), dtype=np.uint8) + pred_label_server_int = face_landmark_server_new.copy().astype(np.int32) + # img_client = img.copy() + for pt in pred_label_server_int: + cv2.circle(img_server, (pt[0], pt[1]), 1, (0, 0, 255), 1) + cv2.imshow("img_server: ", img_server) + + + img_new = np.zeros((detect_single_face_size, detect_single_face_size, 3), dtype=np.uint8) + pred_new_label_int = transform_points(face_landmark_server_new, new_mat[:2, :]).astype(np.int32) + for pt in pred_new_label_int: + cv2.circle(img_new, (pt[0], pt[1]), 1, (0, 255, 0), 1) + cv2.imshow("img_show new: ", img_new) + cv2.waitKey() + + return new_mat[:2, :] + +def get_transform_mat_for_eye(landmark, output_size): + dst_size = output_size + if len(landmark) == 137: + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[96] + landmark[88]) / 2 + right_eye = (landmark[105] + landmark[113]) / 2 + nose = landmark[83] + left_mouth = (landmark[22] + landmark[48]) / 2 + right_mouth = (landmark[56] + landmark[36]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + +def get_transform_mat_for_face_recognition(landmark, output_size): + g_Average_5point_180 = np.array([ + 57, 73, + 123, 73, + 90, 107, + 62, 134, + 118, 134 + ]) + dst_size = output_size + + if len(landmark) == 87: + left_eye = (landmark[17 + 19] + landmark[17 + 22]) / 2 + right_eye = (landmark[17 + 25] + landmark[17 + 28]) / 2 + nose = (landmark[17 + 14] + landmark[17 + 18]) / 2 + left_mouth = (landmark[17 + 31] + landmark[17 + 43]) / 2 + right_mouth = (landmark[17 + 47] + landmark[17 + 37]) / 2 + elif len(landmark) == 137: + left_eye = (landmark[88] + landmark[96]) / 2 + right_eye = (landmark[105] + landmark[113]) / 2 + nose = landmark[83] + left_mouth = (landmark[22] + landmark[48]) / 2 + right_mouth = (landmark[56] + landmark[36]) / 2 + elif len(landmark) == 1000: + left_eye = (landmark[691] + landmark[723]) / 2 + right_eye = (landmark[792] + landmark[824]) / 2 + nose = landmark[621] + left_mouth = (landmark[467] + landmark[468]) / 2 + right_mouth = (landmark[396] + landmark[508]) / 2 + else: + assert False + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) / 180 * dst_size + + mat = umeyama(pts5_src, pts5_dst, True) + + return mat + +def get_transform_mat_mmcv_hair(landmark, output_size, forlabel=False): + dst_size = output_size + + if len(landmark) == 1000: + eye_dis = 0.42 + mouth_dis = 0.42 + g_Average_5point_180 = np.array([ + eye_dis, 0.48, + 1 - eye_dis, 0.48, + 0.5, 0.53, + mouth_dis, 0.58, + 1 - mouth_dis, 0.58 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[691] + landmark[723]) / 2 + right_eye = (landmark[792] + landmark[824]) / 2 + nose = landmark[621] + left_mouth = (landmark[467] + landmark[468]) / 2 + right_mouth = (landmark[396] + landmark[508]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + else: + print("landmark < 1000 !!!") + assert False + return mat + + +def get_transform_mat_mmcv_seg(landmark, output_size, forlabel=False): + dst_size = output_size + eye_dis = 0.40 + mouth_dis = 0.40 + g_Average_5point_180 = np.array([ + eye_dis, 0.46, + 1 - eye_dis, 0.46, + 0.5, 0.55, + mouth_dis, 0.64, + 1 - mouth_dis, 0.64 + ]) + + if len(landmark) == 1000: + # print(g_Average_5point_180) + left_eye = (landmark[691] + landmark[723]) / 2 + right_eye = (landmark[792] + landmark[824]) / 2 + nose = landmark[621] + left_mouth = (landmark[467] + landmark[468]) / 2 + right_mouth = (landmark[396] + landmark[508]) / 2 + + elif len(landmark) == 137: + left_eye = (landmark[88] + landmark[96]) / 2 + right_eye = (landmark[105] + landmark[113]) / 2 + nose = landmark[83] + left_mouth = (landmark[22] + landmark[48]) / 2 + right_mouth = (landmark[56] + landmark[36]) / 2 + else: + print("landmark < 1000 !!!") + assert False + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + + +def get_transform_mat_two_sets(landmarks_src, landmarks_dst): + assert len(landmarks_src) == len(landmarks_dst) + assert len(landmarks_src) == 137 + mat = umeyama(landmarks_src[:22], landmarks_dst[:22], True)[0:2] + return mat + +def flip_points(landmark, width): + if len(landmark) == 137: + landmarks_order = np.array([1, 22, 21, 20, 19, 18, 17, 16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, # counter + 37, 36, 35, 34, 33, 32, 31, 30, 29, 28, 27, 26, 25, 24, 23, 48, 47, 46, 45, 44, 43, 42, 41, + 40, 39, 38, + 57, 56, 55, 54, 53, 52, 51, 50, 49, 64, 63, 62, 61, 60, 59, 58, # mouth + 79, 78, 77, 76, 75, 74, 73, 72, 71, 70, 69, 68, 67, 66, 65, 83, 82, 81, 80, 84, 85, 86, 87, + # nose + 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, # eye + 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, + 134, 133, 132, 131, 130, 137, 136, 135, # eyebrow + 126, 125, 124, 123, 122, 129, 128, 127], dtype=np.int32) - 1 + landmark_flip = landmark.copy() + landmark_flip[:, 0] = width - 1 - landmark_flip[:, 0] + landmark_flip = landmark_flip[landmarks_order, :] + return landmark_flip + elif len(landmark) == 1000: + dst = [0] + list(range(311, 156, -1)) + list(range(156, 0, -1)) + \ + list(range(396, 354, -1)) + list(range(354, 311, -1)) + list(range(467, 432, -1)) + list(range(432, 396, -1)) + \ + list(range(508, 488, -1)) + list(range(488, 467, -1)) + list(range(547, 528, -1)) + list(range(528, 508, -1)) + \ + list(range(616, 584, -1)) + [584, 583, 582, 581, 580] + list(range(579, 547, -1)) + [620, 619, 618, 617] + list(range(621, 654, 1)) + \ + [755] + list(range(774, 755, -1)) + list(range(791, 774, -1)) + list(range(792, 856, 1)) + \ + [654] + list(range(673, 654, -1)) + list(range(690, 673, -1)) + list(range(691, 755, 1)) + \ + list(range(964, 927, -1)) + list(range(999, 964, -1)) + \ + list(range(892, 855, -1)) + list(range(927, 892, -1)) + landmarks_order = np.array(dst, dtype=np.int32) + landmark_flip = landmark.copy() + landmark_flip[:, 0] = width - 1 - landmark_flip[:, 0] + landmark_flip = landmark_flip[landmarks_order, :] + return landmark_flip + else: + assert False + +def pts_1k_to_137(landmarks): + index_1k_to_137 = [0, 12, 24, 36, 48, 61, 74, 87, 100, 119, 137, 156, 175, 193, 212, 225, 238, 251, 264, 276, 288, + 300, 312, 318, 324, 330, 336, 342, 348, 354, 360, 366, 372, 378, 384, 390, 396, 402, 408, 414, + 420, 426, 432, 438, 444, 450, 456, 462, 468, 473, 478, 483, 488, 493, 498, 503, 508, 513, 518, + 523, 528, 533, 538, 543, 548, 556, 564, 571, 579, 580, 581, 582, 583, 584, 585, 593, 600, 608, + 616, 617, 618, 619, 620, 621, 632, 642, 653, 654, 691, 695, 699, 703, 707, 711, 715, 719, 723, + 727, 731, 735, 739, 743, 747, 751, 755, 792, 796, 800, 804, 808, 812, 816, 820, 824, 828, 832, + 836, 840, 844, 848, 852, 856, 865, 874, 883, 892, 901, 910, 919, 928, 937, 946, 955, 964, 973, + 982, 991] + landmarks_137 = landmarks[index_1k_to_137, :] + return landmarks_137 + +def transform_points(points, mat, invert=False): + if invert: + mat = cv2.invertAffineTransform(mat) + points = np.expand_dims(points, axis=1) + points = cv2.transform(points, mat, points.shape) + points = np.squeeze(points) + return points + +def calc_face_pitch(landmarks): + if not isinstance(landmarks, np.ndarray): + landmarks = np.array(landmarks) + t = ((landmarks[6][1] - landmarks[8][1]) + (landmarks[10][1] - landmarks[8][1])) / 2.0 + b = landmarks[8][1] + return float(b - t) + +def calc_face_yaw(landmarks): + if not isinstance(landmarks, np.ndarray): + landmarks = np.array(landmarks) + l = ((landmarks[27][0] - landmarks[0][0]) + (landmarks[28][0] - landmarks[1][0]) + ( + landmarks[29][0] - landmarks[2][0])) / 3.0 + r = ((landmarks[16][0] - landmarks[27][0]) + (landmarks[15][0] - landmarks[28][0]) + ( + landmarks[14][0] - landmarks[29][0])) / 3.0 + return float(r - l) + +# deprecated +def draw_pncc_features(fc_landmark, img_target, w=256, h=256, is_train=True): + assert False + return img_target + + +def draw_blur_no_mouth_img(img_target, fc_landmark, is_train=False): + h, w, c = img_target.shape + assert h == w + + if len(fc_landmark) == 1000: + fc_landmark = pts_1k_to_137(fc_landmark) + + + k_mid_size = int(w / 6.) + thresh = int(w / 15.) + if is_train: + k_size = random.choice(range(k_mid_size - thresh, k_mid_size + thresh, 2)) + else: + k_size = k_mid_size + if k_size % 2 == 0: + k_size += 1 + img_target_blur = cv2.blur(img_target, (k_size, k_size)) + + + hull_mask = np.zeros(img_target.shape, dtype=np.uint8) + if len(fc_landmark) == 1000: + fc_landmark = pts_1k_to_137(fc_landmark) + + cv2.fillPoly(hull_mask, fc_landmark[22:48][np.newaxis, :, :], (255, 255, 255)) + # cv2.imshow('hull_mask1', hull_mask) + # cv2.imshow('img_target_blur', img_target_blur) + # cv2.imshow('img_target1', img_target) + kernel_size = int(w / 35.) + if kernel_size % 2 == 0: + kernel_size += 1 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + hull_mask = cv2.dilate(hull_mask, kernel, iterations=1) .astype(np.uint8) + # cv2.imshow('hull_mask2', hull_mask) + img_target = img_target_blur * (hull_mask<125) + img_target * (hull_mask>125) + + # cv2.imshow('img_target2', img_target) + # cv2.waitKey() + + return img_target + + + +def draw_blur_img(img_target, fc_landmark, is_train=False): + h, w, c = img_target.shape + assert h == w + inpaint_mask = np.zeros((h, w), dtype=np.uint8) + + if len(fc_landmark) == 1000: + fc_landmark = pts_1k_to_137(fc_landmark) + + # cv2.imshow('img_target1', img_target) + + k_mid_size = int(w / 6.) + thresh = int(w / 15.) + if is_train: + k_size = random.choice(range(k_mid_size - thresh, k_mid_size + thresh, 2)) + else: + k_size = k_mid_size + if k_size % 2 == 0: + k_size += 1 + img_target = cv2.blur(img_target, (k_size, k_size)) + + # cv2.imshow('img_target2', img_target) + # cv2.waitKey() + + return img_target + +def draw_blurmore_img(img_target, fc_landmark, is_train=False): + h, w, c = img_target.shape + assert h == w + inpaint_mask = np.zeros((h, w), dtype=np.uint8) + + if len(fc_landmark) == 1000: + fc_landmark = pts_1k_to_137(fc_landmark) + + # cv2.imshow('img_target1', img_target) + + k_mid_size = int(w / 3.) + thresh = int(w / 15.) + if is_train: + k_size = random.choice(range(k_mid_size - thresh, k_mid_size + thresh, 2)) + else: + k_size = k_mid_size + if k_size % 2 == 0: + k_size += 1 + img_target = cv2.blur(img_target, (k_size, k_size)) + + # cv2.imshow('img_target2', img_target) + # cv2.waitKey() + + return img_target + + +def draw_comapre_mask(fc_landmark, w=256, h=256): + inpaint_mask = np.zeros((h, w), dtype=np.uint8) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(fc_landmark[129:137]), (255,)) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(fc_landmark[121:129]), (255,)) + + left_eye = np.around((fc_landmark[96] + fc_landmark[88]) / 2).astype(np.int32) + right_eye = np.around((fc_landmark[105] + fc_landmark[113]) / 2).astype(np.int32) + len = cv2.norm(fc_landmark[96].astype(np.int32) - fc_landmark[88].astype(np.int32)) + cv2.circle(inpaint_mask, (left_eye[0], left_eye[1]), 1, (255), int(len * 1.1)) + len = cv2.norm(fc_landmark[105].astype(np.int32) - fc_landmark[113].astype(np.int32)) + + cv2.circle(inpaint_mask, (right_eye[0], right_eye[1]), 1, (255), int(len * 1.1)) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(fc_landmark[64:87]), (255,)) + + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(fc_landmark[22:64]), (255,)) + + dilate_kernel_size = int(w / 23.) + if dilate_kernel_size % 2 == 0: + dilate_kernel_size += 1 + kernel = np.ones((dilate_kernel_size, dilate_kernel_size), np.uint8) + inpaint_mask = cv2.dilate(inpaint_mask, kernel) + + inpaint_mask = 1 - inpaint_mask.astype(np.float32) / 255 + + return inpaint_mask + +def draw_hull_mask(fc_landmark, w=256, h=256, is_gray=False): + if not is_gray: + hull_mask = np.zeros((h, w, 3), dtype=np.uint8) + + line_size1 = int(w / 85.) + line_size2 = int(w / 51.) + + if len(fc_landmark) == 137: + # cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark[0:22]), (64, 16, 32)) + + # left_brown = ((fc_landmark[129] + fc_landmark[133]) / 2).astype(np.int32) + # cv2.circle(hull_mask, (left_brown[0], left_brown[1]), 1, (0, 255, 0), line_size1) + # right_brwon = ((fc_landmark[121] + fc_landmark[125]) / 2).astype(np.int32) + # cv2.circle(hull_mask, (right_brwon[0], right_brwon[1]), 1, (0, 0, 255), line_size1) + + cv2.line(hull_mask, (int(fc_landmark[129, 0]), int(fc_landmark[129, 1])), + (int(fc_landmark[133, 0]), int(fc_landmark[133, 1])), (0, 255, 0), line_size1) + cv2.line(hull_mask, (int(fc_landmark[121, 0]), int(fc_landmark[121, 1])), + (int(fc_landmark[125, 0]), int(fc_landmark[125, 1])), (0, 0, 255), line_size1) + + cv2.fillPoly(hull_mask, fc_landmark[88:104][np.newaxis, :, :], (128, 128, 0)) + cv2.fillPoly(hull_mask, fc_landmark[105:121][np.newaxis, :, :], (128, 0, 128)) + + # cv2.line(hull_mask, (int(fc_landmark[86, 0]), int(fc_landmark[86, 1])), + # (int(fc_landmark[83, 0]), int(fc_landmark[83, 1])), (255, 0, 0), line_size2) + + cv2.fillPoly(hull_mask, fc_landmark[48:64][np.newaxis, :, :], (0, 128, 128)) + + # cv2.fillPoly(hull_mask, np.concatenate((fc_landmark[22:37], fc_landmark[56:47:-1]))[np.newaxis, :, :], + # (0, 128, 0)) + # + # cv2.fillPoly(hull_mask, + # np.concatenate((fc_landmark[47:35:-1], fc_landmark[56:64], [fc_landmark[48], fc_landmark[22]]))[ + # np.newaxis, :, :], (0, 0, 128)) + + eye = np.zeros((h, w, 3), dtype=np.uint8) + eye_mask1 = np.zeros((h, w, 1), dtype=np.uint8) + eye_mask2 = np.zeros((h, w, 1), dtype=np.uint8) + cv2.fillPoly(eye_mask1, fc_landmark[88:104][np.newaxis, :, :], (1,)) + cv2.fillPoly(eye_mask1, fc_landmark[105:121][np.newaxis, :, :], (1,)) + left_eye = fc_landmark[87] + right_eye = fc_landmark[104] + # left_eye = np.around((fc_landmark[96] + fc_landmark[88]) / 2).astype(np.int32) + # right_eye = np.around((fc_landmark[105] + fc_landmark[113]) / 2).astype(np.int32) + cv2.circle(eye_mask2, (left_eye[0], left_eye[1]), 1, (1,), line_size2) + cv2.circle(eye_mask2, (right_eye[0], right_eye[1]), 1, (1,), line_size2) + eye_mask = eye_mask1 & eye_mask2 + cv2.circle(eye, (left_eye[0], left_eye[1]), 1, (255, 255, 255), line_size2) + cv2.circle(eye, (right_eye[0], right_eye[1]), 1, (255, 255, 255), line_size2) + hull_mask = hull_mask * (1 - eye_mask[:, :, 0:1]) + eye * eye_mask[:, :, 0:1] + elif len(fc_landmark) == 1000: + # cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark[0:312]), (64, 16, 32)) + + # left_brown = ((fc_landmark[928] + fc_landmark[964]) / 2).astype(np.int32) + # cv2.circle(hull_mask, (left_brown[0], left_brown[1]), 1, (0, 255, 0), line_size1) + # right_brwon = ((fc_landmark[856] + fc_landmark[892]) / 2).astype(np.int32) + # cv2.circle(hull_mask, (right_brwon[0], right_brwon[1]), 1, (0, 0, 255), line_size1) + + cv2.line(hull_mask, (int(fc_landmark[928, 0]), int(fc_landmark[928, 1])), + (int(fc_landmark[964, 0]), int(fc_landmark[964, 1])), (0, 255, 0), line_size1) + cv2.line(hull_mask, (int(fc_landmark[856, 0]), int(fc_landmark[856, 1])), + (int(fc_landmark[892, 0]), int(fc_landmark[892, 1])), (0, 0, 255), line_size1) + + cv2.fillPoly(hull_mask, fc_landmark[691:755][np.newaxis, :, :], (128, 128, 0)) + cv2.fillPoly(hull_mask, fc_landmark[792:856][np.newaxis, :, :], (128, 0, 128)) + + # cv2.line(hull_mask, (int(fc_landmark[653, 0]), int(fc_landmark[653, 1])), + # (int(fc_landmark[621, 0]), int(fc_landmark[621, 1])), (255, 0, 0), line_size2) + + cv2.fillPoly(hull_mask, fc_landmark[468:548][np.newaxis, :, :], (0, 128, 128)) + + # cv2.fillPoly(hull_mask, np.concatenate((fc_landmark[312:397], fc_landmark[508:467:-1]))[np.newaxis, :, :], + # (0, 128, 0)) + # + # cv2.fillPoly(hull_mask, + # np.concatenate((fc_landmark[467:395:-1], fc_landmark[508:548], [fc_landmark[468], fc_landmark[312]]))[ + # np.newaxis, :, :], (0, 0, 128)) + + eye = np.zeros((h, w, 3), dtype=np.uint8) + eye_mask1 = np.zeros((h, w, 1), dtype=np.uint8) + eye_mask2 = np.zeros((h, w, 1), dtype=np.uint8) + cv2.fillPoly(eye_mask1, fc_landmark[691:755][np.newaxis, :, :], (1,)) + cv2.fillPoly(eye_mask1, fc_landmark[792:856][np.newaxis, :, :], (1,)) + left_eye = fc_landmark[654] + right_eye = fc_landmark[755] + # cv2.fillPoly(eye_mask2, fc_landmark[655:691][np.newaxis, :, :], (1,)) + cv2.circle(eye_mask2, (left_eye[0], left_eye[1]), 1, (1,), line_size2) + # cv2.fillPoly(eye_mask2, fc_landmark[756:792][np.newaxis, :, :], (1,)) + cv2.circle(eye_mask2, (right_eye[0], right_eye[1]), 1, (1,), line_size2) + eye_mask = eye_mask1 & eye_mask2 + # cv2.fillPoly(eye, fc_landmark[655:691][np.newaxis, :, :], (255, 255, 255)) + cv2.circle(eye, (left_eye[0], left_eye[1]), 1, (255, 255, 255), line_size2) + # cv2.fillPoly(eye, fc_landmark[756:792][np.newaxis, :, :], (255, 255, 255)) + cv2.circle(eye, (right_eye[0], right_eye[1]), 1, (255, 255, 255), line_size2) + hull_mask = hull_mask * (1 - eye_mask[:, :, 0:1]) + eye * eye_mask[:, :, 0:1] + else: + assert False + else: + hull_mask = np.zeros((h, w), dtype=np.uint8) + if len(fc_landmark) == 137: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (1)) + elif len(fc_landmark) == 1000: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (1)) + else: + assert False + + return hull_mask + + +def draw_makeup_mask(another_pts1k, another_mask, pts1k, hull_mask): + pts1kint = pts1k.astype(np.int32) + assert len(another_pts1k) == 1000 and len(pts1kint) == 1000 + cv2.fillPoly(another_mask, another_pts1k[691:755][np.newaxis, :, :], (4)) # eye + cv2.fillPoly(another_mask, another_pts1k[792:856][np.newaxis, :, :], (5)) # eye + cv2.fillPoly(another_mask, np.concatenate((another_pts1k[312:397], another_pts1k[508:467:-1]))[np.newaxis, :, :], (7)) # mouth + cv2.fillPoly(another_mask, np.concatenate((another_pts1k[467:395:-1], another_pts1k[508:548], [another_pts1k[468], another_pts1k[312]]))[np.newaxis, :, :], (9)) # mouth + cv2.fillPoly(another_mask, another_pts1k[928:1000][np.newaxis, :, :], (0)) # eyebrow + cv2.fillPoly(another_mask, another_pts1k[856:928][np.newaxis, :, :], (0)) # eyebrow + + + cv2.fillPoly(hull_mask, pts1kint[691:755][np.newaxis, :, :], (4)) # eye + cv2.fillPoly(hull_mask, pts1kint[792:856][np.newaxis, :, :], (5)) # eye + cv2.fillPoly(hull_mask, np.concatenate((pts1kint[312:397], pts1kint[508:467:-1]))[np.newaxis, :, :], (7)) # mouth + cv2.fillPoly(hull_mask, np.concatenate((pts1kint[467:395:-1], pts1kint[508:548], [pts1kint[468], pts1kint[312]]))[np.newaxis, :, :], (9)) # mouth + cv2.fillPoly(hull_mask, pts1kint[928:1000][np.newaxis, :, :], (0)) # eyebrow + cv2.fillPoly(hull_mask, pts1kint[856:928][np.newaxis, :, :], (0)) # eyebrow + return another_mask, hull_mask + +def draw_users_hull_mask(fc_landmark, w=256, h=256): + hull_mask = np.zeros((h, w), dtype=np.uint8) + if len(fc_landmark) == 87: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (1)) + elif len(fc_landmark) == 137: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (1)) + elif len(fc_landmark) == 1000: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (1)) + else: + assert False + return hull_mask + +def draw_bigger_hull_mask(fc_landmark, w=256, h=256): + + if len(fc_landmark) == 1000: + fc_landmark = pts_1k_to_137(fc_landmark) + + hull_mask = np.zeros((h, w), dtype=np.uint8) + if len(fc_landmark) == 137: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (255)) + elif len(fc_landmark) == 1000: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (255)) + else: + assert False + + kernel_size = int(w / 11.) + if kernel_size % 2 == 0: + kernel_size += 1 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + hull_mask = (cv2.dilate(hull_mask, kernel, iterations=1) / 255.).astype(np.uint8) + cv2.fillPoly(hull_mask, fc_landmark[22:48][np.newaxis, :, :], (10,)) + + + leye_mask = np.zeros((h, w), dtype=np.uint8) + cv2.fillPoly(leye_mask, fc_landmark[88:104][np.newaxis, :, :], (1,)) + cv2.fillPoly(leye_mask, fc_landmark[105:121][np.newaxis, :, :], (1,)) + kernel_size = int(w / 5.) + if kernel_size % 2 == 0: + kernel_size += 1 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + leye_mask = (cv2.dilate(leye_mask, kernel, iterations=1)).astype(np.uint8) + hull_mask[leye_mask > 0] = 10 + + # cv2.imshow('hull_mask', hull_mask * 25 ) + # cv2.waitKey() + + return hull_mask + + + +def get_kernel_size(div_num, output_img_size=256): + kernel_size = int(float(output_img_size) / div_num) + if kernel_size % 2 == 0: + kernel_size += 1 + return kernel_size + +class LKTracking(object): + def __init__(self, termcrit=cv2.TERM_CRITERIA_COUNT | cv2.TERM_CRITERIA_EPS, winSize=31, extSize=41, threshold=5): + self.termcrit = termcrit + self.winSize = winSize + self.extSize = extSize + self.threshold = threshold + + self.prepoints = [] + self.preimg_small_ = None + self.preimg_rect_ = np.array([0, 0, 0, 0]) + + def UpdatePoints(self, new_pts): + if len(new_pts) != len(self.prepoints): + self.prepoints = [] + return False + self.prepoints = new_pts + return True + + def TrackingPoints(self, curimg, curpoints): + need_init = False + + def Prepare(self, curimg, curpoints): + prepoints_int = curpoints.astype(np.int32) + pass + +def get_transform_mat_mouth(landmark, output_size): + # mean_mouth_x_4pts = np.array([0.15, 0.5, 0.75, 0.5]) + # mean_mouth_y_4pts = np.array([0.5, 0.15, 0.5, 0.75]) + mean_mouth_x_4pts = np.array([0.2, 0.5, 0.8, 0.5]) + mean_mouth_y_4pts = np.array([0.5, 0.2, 0.5, 0.8]) + landmarks_2D_mouth_4pts = np.stack([mean_mouth_x_4pts, mean_mouth_y_4pts], axis=1) + + dst_size = output_size + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k) + mat = umeyama(landmark[:22], landmarks_2D_137[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 236: + #mouth crop + landmark_mouth = [] + landmark_mouth.append(landmark[467-312]) + landmark_mouth.append(landmark[432-312]) + landmark_mouth.append(landmark[397-312]) + landmark_mouth.append(landmark[354-312]) + landmark_mouth = np.array(landmark_mouth) + mat = umeyama(landmark_mouth, landmarks_2D_mouth_4pts * dst_size, True)[0:2] + return mat + + +mean_mouth_x_4pts = np.array([0.2, 0.5, 0.8, 0.5]) +mean_mouth_y_4pts = np.array([0.5, 0.2, 0.5, 0.8]) +landmarks_2D_mouth_4pts = np.stack([mean_mouth_x_4pts, mean_mouth_y_4pts], axis=1) + +def get_transform_mat_for_mouth(landmark, output_size): + dst_size = output_size + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + # landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k) + # mat = umeyama(landmark[:22], landmarks_2D_137[:22] * dst_size, True)[0:2] + landmark_mouth = [] + landmark_mouth.append(landmark[22]) + landmark_mouth.append(landmark[42]) + landmark_mouth.append(landmark[36]) + landmark_mouth.append(landmark[29]) + landmark_mouth = np.array(landmark_mouth) + mat = umeyama(landmark_mouth, landmarks_2D_mouth_4pts * dst_size, True)[0:2] + + return mat + elif len(landmark) == 236: + # mouth crop + landmark_mouth = [] + landmark_mouth.append(landmark[467 - 312]) + landmark_mouth.append(landmark[432 - 312]) + landmark_mouth.append(landmark[397 - 312]) + landmark_mouth.append(landmark[354 - 312]) + landmark_mouth = np.array(landmark_mouth) + mat = umeyama(landmark_mouth, landmarks_2D_mouth_4pts * dst_size, True)[0:2] + # mat = umeyama(landmark[:312], landmarks_2D_1k[:312] * dst_size, True)[0:2] + return mat + + elif len(landmark) == 1000: + # mouth crop + landmark_mouth = [] + landmark_mouth.append(landmark[467]) + landmark_mouth.append(landmark[432]) + landmark_mouth.append(landmark[397]) + landmark_mouth.append(landmark[354]) + landmark_mouth = np.array(landmark_mouth) + mat = umeyama(landmark_mouth, landmarks_2D_mouth_4pts * dst_size, True)[0:2] + # mat = umeyama(landmark[:312], landmarks_2D_1k[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_full_face_ratio_skin(landmark, output_size, ratio=1.0, skin=0): + dst_size = output_size + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + # landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + 0.5 + # landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.5) * ratio + 0.4 + + if skin == 1: + ###################### 0.3 face rata 592 + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + 0.5 + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.6) * ratio + 0.3 + ###################### 0.3 face rata 592 + else: + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.4) * ratio + 0.45 + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.4) * ratio + 0.4 + + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_1k_tmp = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_1k_tmp[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_for_uv(landmark, output_size): + dst_size = output_size + if len(landmark) == 1000: + landmark = pts_1k_to_137(landmark) + + if len(landmark) == 137: + eye_dis = 0.35 + mouth_dis = 0.42 + g_Average_5point_180 = np.array([ + eye_dis, 0.24, + 1 - eye_dis, 0.24, + 0.5, 0.42, + mouth_dis, 0.55, + 1 - mouth_dis, 0.55 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[96] + landmark[88]) / 2 + right_eye = (landmark[105] + landmark[113]) / 2 + nose = landmark[83] + left_mouth = (landmark[22] + landmark[48]) / 2 + right_mouth = (landmark[56] + landmark[36]) / 2 + + pts5_dst = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_src = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + else: + assert False + +def align2stylegan(face_landmarks_1k, output_size=256): + face_landmarks_1k = np.float32(face_landmarks_1k) + x_scale = 1.0 + y_scale = 1.0 + em_scale = 0.1 + eye_left = (face_landmarks_1k[691] + face_landmarks_1k[723]) / 2 + eye_right = (face_landmarks_1k[792] + face_landmarks_1k[824]) / 2 + mouth_left = (face_landmarks_1k[467] + face_landmarks_1k[468]) / 2 + mouth_right = (face_landmarks_1k[396] + face_landmarks_1k[508]) / 2 + eye_avg = (eye_left + eye_right) * 0.5 + eye_to_eye = eye_right - eye_left + mouth_avg = (mouth_left + mouth_right) * 0.5 + eye_to_mouth = mouth_avg - eye_avg + x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1] + x /= np.hypot(*x) + x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8) + x *= x_scale + y = np.flipud(x) * [-y_scale, y_scale] + c = eye_avg + eye_to_mouth * em_scale + quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y]) + quad_ori = np.array(quad) + + rotate_radian = math.atan2((quad_ori[3][1] - quad_ori[0][1]), (quad_ori[3][0] - quad_ori[0][0])) + rotate_degree = rotate_radian / np.pi * 180 + scale = output_size / cv2.norm(quad_ori[3] - quad_ori[0]) + src_center = (quad_ori[0] + quad_ori[2]) * 0.5 + dst_center = np.float32([output_size / 2, output_size / 2]) + + M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), rotate_degree, scale) + M[:, 2] += dst_center - src_center + return M +def draw_crop_eye_bysize(img, pts1k, img_size, change_eyebrow): + pts137tmp = pts_1k_to_137(pts1k).astype(np.int32) + eye_mask = np.zeros((img_size, img_size, 3), dtype=np.uint8) + cv2.fillPoly(eye_mask, pts137tmp[88:104][np.newaxis, :, :], (1, 1, 1)) + cv2.fillPoly(eye_mask, pts137tmp[105:121][np.newaxis, :, :], (1, 1, 1)) + kernel_size = int(img_size / 198.) # 31 + if kernel_size % 2 == 0: + kernel_size += 1 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + eye_mask = (cv2.dilate(eye_mask, kernel, iterations=1)).astype(np.uint8) + + kernel_size = int(img_size / 7) # 8.2 + if kernel_size % 2 == 0: + kernel_size += 1 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size*2, kernel_size)) + eye_mask2 = (cv2.dilate(eye_mask, kernel, iterations=1)).astype(np.uint8) + eye_mask2 = eye_mask2 - eye_mask + + if change_eyebrow: + # eyebrow + eyebrow_mask = np.zeros((img_size, img_size, 3), dtype=np.uint8) + + x, y, w, h = cv2.boundingRect(pts137tmp[129:137]) + # cv2.rectangle(eyebrow_mask, (x, y), (x + w, y + h), (1, 1, 1), -1) + eyebrow_mask[y:y + h, x:x + w, :] = [1, 1, 1] + + x, y, w, h = cv2.boundingRect(pts137tmp[121:129]) + # cv2.rectangle(eyebrow_mask, (x, y), (x + w, y + h), (1, 1, 1), -1) + eyebrow_mask[y:y + h, x:x + w, :] = [1, 1, 1] + + kernel_size = int(592 / 21.) # 21 + if kernel_size % 2 == 0: + kernel_size += 1 + kernel_eyebrow = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (int(kernel_size * 5), int(kernel_size))) + eyebrow_mask = (cv2.dilate(eyebrow_mask, kernel_eyebrow, iterations=1)).astype(np.uint8) + eye_mask2 = (eye_mask2.astype(np.bool) | eyebrow_mask.astype(np.bool)).astype(np.float32) + else: + eyebrow_mask = np.zeros((img_size, img_size, 3), dtype=np.uint8) + cv2.fillPoly(eyebrow_mask, pts137tmp[129:137][np.newaxis, :, :], (1, 1, 1)) + cv2.fillPoly(eyebrow_mask, pts137tmp[121:129][np.newaxis, :, :], (1, 1, 1)) + eyebrow_mask = 1 - eyebrow_mask + eye_mask2 = (eye_mask2.astype(np.bool) & eyebrow_mask.astype(np.bool)).astype(np.float32) + + # nose + cv2.fillPoly(eye_mask2, pts137tmp[64:79][np.newaxis, :, :], (0, 0, 0)) + eye_mask2[(pts137tmp[133, 1] - img_size // 10):pts137tmp[65, 1], pts137tmp[64, 0]:pts137tmp[78, 0]] = 0 + + img[eye_mask2 > 0] = 0 + + cv2.circle(img, (pts137tmp[133, 0], pts137tmp[133, 1]), img_size // 50, (255, 0, 0), -1) + cv2.circle(img, (pts137tmp[121, 0], pts137tmp[121, 1]), img_size // 50, (255, 0, 0), -1) + + return img + + +# def draw_crop_eye_bysize_using_seg_mask(img, pts1k, mask, img_size, change_eyebrow): +# +# pts137tmp = pts_1k_to_137(pts1k).astype(np.int32) +# c0, c1, c2 = mask[:, :, 0], mask[:, :, 1], mask[:, :, 2] +# c0[c0 >= 100] = 255 +# c0[c0 < 100] = 0 +# c1[c1 >= 100] = 255 +# c1[c1 < 100] = 0 +# c2[c2 > 0] = 0 +# mask[:, :, 0] = c0 +# mask[:, :, 1] = c1 +# mask[:, :, 2] = c2 +# eye_index = (mask == [255, 0, 0]).all(axis=2) +# eyelids_index = (mask == [0, 255, 0]).all(axis=2) +# black_im = np.zeros((img_size, img_size, 3), dtype=np.uint8) +# black_im2 = np.zeros((img_size, img_size, 3), dtype=np.uint8) +# +# black_im[eye_index] = [255, 255, 255] +# black_im2[eye_index] = [255, 255, 255] +# kernel_size = int(img_size // 7) # 8.2 +# +# if kernel_size % 2 == 0: +# kernel_size += 1 +# kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size * 2, kernel_size)) +# black_im2 = cv2.dilate(black_im2, kernel).astype(np.uint8) +# +# endless_belt_im = black_im2 - black_im +# endless_belt_im = np.ones_like(endless_belt_im) * 255 - endless_belt_im +# +# if change_eyebrow: +# # eyebrow +# eyebrow_mask = np.zeros((img_size, img_size, 3), dtype=np.uint8) +# x, y, w, h = cv2.boundingRect(pts137tmp[129:137]) +# # cv2.rectangle(eyebrow_mask, (x, y), (x + w, y + h), (1, 1, 1), -1) +# eyebrow_mask[y:y + h, x:x + w, :] = [1, 1, 1] +# +# x, y, w, h = cv2.boundingRect(pts137tmp[121:129]) +# # cv2.rectangle(eyebrow_mask, (x, y), (x + w, y + h), (1, 1, 1), -1) +# eyebrow_mask[y:y + h, x:x + w, :] = [1, 1, 1] +# +# kernel_size = int(592 / 21.) # 21 +# if kernel_size % 2 == 0: +# kernel_size += 1 +# kernel_eyebrow = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (int(kernel_size * 5), int(kernel_size))) +# eyebrow_mask = (cv2.dilate(eyebrow_mask, kernel_eyebrow, iterations=1)).astype(np.uint8) +# eyebrow_mask = 1 - eyebrow_mask +# endless_belt_im = (endless_belt_im.astype(np.bool) & eyebrow_mask.astype(np.bool)).astype(np.float32) * 255 +# else: +# eyebrow_mask = np.zeros((img_size, img_size, 3), dtype=np.uint8) +# cv2.fillPoly(eyebrow_mask, pts137tmp[129:137][np.newaxis, :, :], (1, 1, 1)) +# cv2.fillPoly(eyebrow_mask, pts137tmp[121:129][np.newaxis, :, :], (1, 1, 1)) +# endless_belt_im = (endless_belt_im.astype(np.bool) | eyebrow_mask.astype(np.bool)).astype(np.float32) * 255 +# +# cv2.fillPoly(endless_belt_im, pts137tmp[64:79][np.newaxis, :, :], (255, 255, 255)) +# endless_belt_im[(pts137tmp[133, 1] - img_size // 10):pts137tmp[65, 1], pts137tmp[64, 0]:pts137tmp[78, 0], :] = 255 +# +# +# res = np.uint8(endless_belt_im / 255) +# res = np.uint8(res * img) +# +# # add eyelids semantic map +# eyelids_mask = np.zeros((img_size, img_size, 3)) +# eyelids_mask[eyelids_index] = [0, 255, 0] +# res = res + np.uint8(eyelids_mask) +# +# cv2.circle(res, (pts137tmp[133, 0], pts137tmp[133, 1]), img_size // 50, (255, 0, 0), -1) +# cv2.circle(res, (pts137tmp[121, 0], pts137tmp[121, 1]), img_size // 50, (255, 0, 0), -1) +# +# return res + + +def draw_crop_eye_bysize_using_seg_mask(img, pts1k, mask, img_size, change_eyebrow): + start = time.time() + mask = mask.copy() + pts137tmp = pts_1k_to_137(pts1k).astype(np.int32) + c0, c1, c2 = mask[:, :, 0], mask[:, :, 1], mask[:, :, 2] + c0[c0 >= 100] = 255 + c0[c0 < 100] = 0 + + c1[c1 >= 100] = 255 + c1[c1 < 100] = 0 + + # c2 也需要截断 + c2[c2 >= 100] = 255 + c2[c2 < 100] = 0 + + mask[:, :, 0] = c0 + mask[:, :, 1] = c1 + mask[:, :, 2] = c2 + eye_index = (mask == [255, 0, 0]).all(axis=2) | (mask == [0, 0, 255]).all(axis=2) + eyelids_index = (mask == [0, 255, 0]).all(axis=2) + black_im = np.zeros((img_size, img_size, 3), dtype=np.uint8) + black_im2 = np.zeros((img_size, img_size, 3), dtype=np.uint8) + + black_im[eye_index] = [255, 255, 255] + black_im2[eye_index] = [255, 255, 255] + + # cv2.imshow("black_im: ", black_im) + # cv2.waitKey() + + kernel_size = int(img_size // 7) # 8.2 + + if kernel_size % 2 == 0: + kernel_size += 1 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size * 2, kernel_size)) + black_im2 = cv2.dilate(black_im2, kernel).astype(np.uint8) + + endless_belt_im = black_im2 - black_im + endless_belt_im = np.ones_like(endless_belt_im) * 255 - endless_belt_im + + if change_eyebrow: + # eyebrow + eyebrow_mask = np.zeros((img_size, img_size, 3), dtype=np.uint8) + x, y, w, h = cv2.boundingRect(pts137tmp[129:137]) + cv2.rectangle(eyebrow_mask, (x, y), (x + w, y + h), (1, 1, 1), -1) + x, y, w, h = cv2.boundingRect(pts137tmp[121:129]) + cv2.rectangle(eyebrow_mask, (x, y), (x + w, y + h), (1, 1, 1), -1) + kernel_size = int(img_size / 21.) # 21 + if kernel_size % 2 == 0: + kernel_size += 1 + kernel_eyebrow = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (int(kernel_size * 5), int(kernel_size))) + eyebrow_mask = (cv2.dilate(eyebrow_mask, kernel_eyebrow, iterations=1)).astype(np.uint8) + eyebrow_mask = 1 - eyebrow_mask + endless_belt_im = (endless_belt_im.astype(np.bool) & eyebrow_mask.astype(np.bool)).astype(np.float32) * 255 + else: + eyebrow_mask = np.zeros((img_size, img_size, 3), dtype=np.uint8) + cv2.fillPoly(eyebrow_mask, pts137tmp[129:137][np.newaxis, :, :], (1, 1, 1)) + cv2.fillPoly(eyebrow_mask, pts137tmp[121:129][np.newaxis, :, :], (1, 1, 1)) + endless_belt_im = (endless_belt_im.astype(np.bool) | eyebrow_mask.astype(np.bool)).astype(np.float32) * 255 + + # nose + cv2.fillPoly(endless_belt_im, pts137tmp[64:79][np.newaxis, :, :], (255, 255, 255)) + endless_belt_im[(pts137tmp[133, 1] - img_size // 10):pts137tmp[65, 1], pts137tmp[64, 0]:pts137tmp[78, 0], :] = 255 + + res = np.uint8(endless_belt_im / 255) + res = np.uint8(res * img) + + # add eyelids semantic map + eyelids_mask = np.zeros((img_size, img_size, 3)) + eyelids_mask[eyelids_index] = [0, 255, 0] + res = res + np.uint8(eyelids_mask) + + cv2.circle(res, (pts137tmp[133, 0], pts137tmp[133, 1]), img_size // 50, (255, 0, 0), -1) + cv2.circle(res, (pts137tmp[121, 0], pts137tmp[121, 1]), img_size // 50, (255, 0, 0), -1) + + return res + +def get_transform_singleeye(landmark, output_size, forlabel=False): + dst_size = output_size + left_dis = 0.28 + g_Average_5point_180 = np.array([ + left_dis, 0.5, + 0.5, 0.499, + 1 - left_dis, 0.5, + ]) + pts3_dst = g_Average_5point_180.reshape((3, -1)) * dst_size + pts3_src = np.vstack((landmark[0], landmark[1], landmark[2])) + image_to_face_mat = umeyama(pts3_src, pts3_dst, True)[:2] + + return image_to_face_mat + +def get_transform_mat_full_face_ratio_stylegan(landmark, output_size, ratio=1.0): + dst_size = output_size + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + + #################### + # # ###################### train stylegan hair rate_0.4 size_512 + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + 0.5 + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.6) * ratio + 0.4 + + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_1k_tmp = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_1k_tmp[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_hair(landmark, output_size, ratio=0.5, w_ratio=0.5, h_ratio=0.40): + dst_size = output_size + + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + w_ratio + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.5) * ratio + h_ratio + + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_1k_tmp = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_1k_tmp[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_hair_ratio(landmark, output_size, ratio=1.0): + dst_size = output_size + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + 0.5 + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.5) * ratio + 0.45 + + # landmarks1k_contours = landmarks_2D_1k_tmp[:312] + # x_l, y_l = np.min(landmarks1k_contours, axis=0) + # x_h, y_h = np.max(landmarks1k_contours, axis=0) + # + # print("x_l: ", x_l, "x_h: ", x_h) + # print("y_l: ", y_l, "y_h: ", y_h) + # + # def draw_landmark(landmark_ori, img): + # landmark_full_int = (landmark_ori.copy() * 512).astype(np.int32) + # img_show = img.copy() + # for pt in landmark_full_int: + # cv2.circle(img_show, (pt[0], pt[1]), 2, (0, 0, 255), 1) + # + # return img_show.astype(np.uint8) + # + # img_temp = np.zeros((512, 512, 3), dtype=np.uint8) + # img_show = draw_landmark(landmarks_2D_1k_tmp, img_temp) + # cv2.imshow("img_show", img_show) + # cv2.waitKey() + + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_1k_tmp = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_1k_tmp[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_hair_ratio_v1(landmark, output_size, ratio=1.0, h_offset=0.5): + dst_size = output_size + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + 0.5 + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.5) * ratio + h_offset + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_1k_tmp = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_1k_tmp[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_sex(landmark, output_size, forlabel=False): + dst_size = output_size + if len(landmark) == 87: + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[31] + landmark[35]) / 2 + right_eye = (landmark[39] + landmark[43]) / 2 + nose = landmark[62] + left_mouth = (landmark[66] + landmark[79]) / 2 + right_mouth = (landmark[72] + landmark[83]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + elif len(landmark) == 137: + + eye_dis = 0.317 + mouth_dis = 0.345 + g_Average_5point_180 = np.array([ + eye_dis, 0.4, + 1 - eye_dis, 0.4, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[96] + landmark[88]) / 2 + right_eye = (landmark[105] + landmark[113]) / 2 + nose = landmark[83] + left_mouth = (landmark[22] + landmark[48]) / 2 + right_mouth = (landmark[56] + landmark[36]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + +def get_transform_mat_full_face_ratio_deeplab(landmark, output_size, ratio=0.3, w_ratio=0.5, h_ratio=0.45): + dst_size = output_size + + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + w_ratio + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.5) * ratio + h_ratio + + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_1k_tmp = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_1k_tmp[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_face_restore(landmark, output_size): + dst_size = output_size + if len(landmark) == 1000: + eye_dis = 0.4 + mouth_dis = 0.42 + g_Average_5point_180 = np.array([ + eye_dis, 0.47, + 1 - eye_dis, 0.47, + 0.5, 0.6, + mouth_dis, 0.71, + 1 - mouth_dis, 0.71 + ]) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + left_eye = (landmark[691] + landmark[723]) / 2 + right_eye = (landmark[792] + landmark[824]) / 2 + nose = landmark[621] + left_mouth = (landmark[467] + landmark[468]) / 2 + right_mouth = (landmark[396] + landmark[508]) / 2 + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + + image_to_face_mat = umeyama(pts5_src, pts5_dst, True)[:2] + + return image_to_face_mat diff --git a/hair_service_sd/core/utils/model_io.py b/hair_service_sd/core/utils/model_io.py new file mode 100644 index 0000000..6d3c2cd --- /dev/null +++ b/hair_service_sd/core/utils/model_io.py @@ -0,0 +1,25 @@ +import os +import torch + +def load_model_by_path(model_save_path, model, gpu_id = None): + if not os.path.exists(model_save_path): return + loc = 'cpu' if gpu_id is None else 'cuda:{}'.format(gpu_id) + pretrained_dict = torch.load(model_save_path, map_location=loc) + if 'state_dict' in pretrained_dict: pretrained_dict = pretrained_dict['state_dict'] + model_dict = model.state_dict() + # pretrained_dict.pop('netG.model.1.weight', '404') + pretrained_dict_new = {k: v for k, v in pretrained_dict.items() + if (k in model_dict and model_dict[k].data.shape == v.data.shape)} + # if 'netG.model.1.weight' not in pretrained_dict_new and 'netG.model.1.weight' in model_dict: + # pretrained_dict_new['netG.model.1.weight'] = model_dict['netG.model.1.weight'] + # pretrained_dict_new['netG.model.1.weight'][:,:12] = pretrained_dict['netG.model.1.weight'] + + model_dict.update(pretrained_dict_new) + model.load_state_dict(model_dict) + print("load model: ", model_save_path) + +def save_model_by_path(model_save_path, model): + save_dir, _ = os.path.split(model_save_path) + if not os.path.isdir(save_dir): os.makedirs(save_dir) + model_dic={k.replace('.module', ''):v for k,v in model.state_dict().items()} + torch.save(model_dic, model_save_path) \ No newline at end of file diff --git a/hair_service_sd/core/utils/params_3ddfa.py b/hair_service_sd/core/utils/params_3ddfa.py new file mode 100644 index 0000000..738e703 --- /dev/null +++ b/hair_service_sd/core/utils/params_3ddfa.py @@ -0,0 +1,8 @@ +import cv2 +import os + +SCALE_F = 1e4 +SCALE_ROTATE = 1e2 +SCALE_OFFSET = 1e-1 +SCALE_SHAPE = 1e-6 +SCALE_EXP = 1 \ No newline at end of file diff --git a/hair_service_sd/core/utils/umeyama.py b/hair_service_sd/core/utils/umeyama.py new file mode 100644 index 0000000..56c8b9e --- /dev/null +++ b/hair_service_sd/core/utils/umeyama.py @@ -0,0 +1,78 @@ +import numpy as np + +def umeyama(src, dst, estimate_scale): + """Estimate N-D similarity transformation with or without scaling. + Parameters + ---------- + src : (M, N) array + Source coordinates. + dst : (M, N) array + Destination coordinates. + estimate_scale : bool + Whether to estimate scaling factor. + Returns + ------- + T : (N + 1, N + 1) + The homogeneous similarity transformation matrix. The matrix contains + NaN values only if the problem is not well-conditioned. + References + ---------- + .. [1] "Least-squares estimation of transformation parameters between two + point patterns", Shinji Umeyama, PAMI 1991, DOI: 10.1109/34.88573 + """ + + num = src.shape[0] + dim = src.shape[1] + + # Compute mean of src and dst. + src_mean = src.mean(axis=0) + dst_mean = dst.mean(axis=0) + + # Subtract mean from src and dst. + src_demean = src - src_mean + dst_demean = dst - dst_mean + + # Eq. (38). + A = np.dot(dst_demean.T, src_demean) / num + + # Eq. (39). + d = np.ones((dim,), dtype=np.double) + if np.linalg.det(A) < 0: + d[dim - 1] = -1 + + T = np.eye(dim + 1, dtype=np.double) + U, S, V = np.linalg.svd(A) + # print('A',A) + # print('V',V) + # print('U', U) + # print('S', S) + # print('T', T) + + # Eq. (40) and (43). + rank = np.linalg.matrix_rank(A) + # print('rank', rank) + if rank == 0: + return np.nan * T + elif rank == dim - 1: + if np.linalg.det(U) * np.linalg.det(V) > 0: + T[:dim, :dim] = np.dot(U, V) + else: + s = d[dim - 1] + d[dim - 1] = -1 + T[:dim, :dim] = np.dot(U, np.dot(np.diag(d), V)) + d[dim - 1] = s + else: + T[:dim, :dim] = np.dot(U, np.dot(np.diag(d), V.T)) + + if estimate_scale: + # Eq. (41) and (42). + scale = 1.0 / src_demean.var(axis=0).sum() * np.dot(S, d) + else: + scale = 1.0 + # print('scale', scale) + # print('dst_maen',dst_mean) + # print('src_mean', src_mean) + T[:dim, dim] = dst_mean - scale * np.dot(T[:dim, :dim], src_mean.T) + T[:dim, :dim] *= scale + + return T \ No newline at end of file diff --git a/hair_service_sd/core/utils/util.py b/hair_service_sd/core/utils/util.py new file mode 100644 index 0000000..84f535f --- /dev/null +++ b/hair_service_sd/core/utils/util.py @@ -0,0 +1,287 @@ +import os +import cv2 +import torch +import logging +import numpy as np +# from utils.config import CONFIG +# import torch.distributed as dist + +def mkdirs(paths): + """create empty directories if they don't exist + Parameters: + paths (str list) -- a list of directory paths + """ + if isinstance(paths, list) and not isinstance(paths, str): + for path in paths: + os.makedirs(path) + else: + os.makedirs(paths) + +def make_dir(target_dir): + """ + Create dir if not exists + """ + if not os.path.exists(target_dir): + os.makedirs(target_dir) + + +def print_network(model, name): + """ + Print out the network information + """ + logger = logging.getLogger("Logger") + num_params = 0 + for p in model.parameters(): + num_params += p.numel() + + logger.info(model) + logger.info(name) + logger.info("Number of parameters: {}".format(num_params)) + + +def update_lr(lr, optimizer): + """ + update learning rates + """ + for param_group in optimizer.param_groups: + param_group['lr'] = lr + + +def warmup_lr(init_lr, step, iter_num): + """ + Warm up learning rate + """ + return step/iter_num*init_lr + + +def add_prefix_state_dict(state_dict, prefix="module"): + """ + add prefix from the key of pretrained state dict for Data-Parallel + """ + new_state_dict = {} + first_state_name = list(state_dict.keys())[0] + if not first_state_name.startswith(prefix): + for key, value in state_dict.items(): + new_state_dict[prefix+"."+key] = state_dict[key].float() + else: + for key, value in state_dict.items(): + new_state_dict[key] = state_dict[key].float() + return new_state_dict + + +def remove_prefix_state_dict(state_dict, prefix="module"): + """ + remove prefix from the key of pretrained state dict for Data-Parallel + """ + new_state_dict = {} + first_state_name = list(state_dict.keys())[0] + if not first_state_name.startswith(prefix): + for key, value in state_dict.items(): + new_state_dict[key] = state_dict[key].float() + else: + for key, value in state_dict.items(): + new_state_dict[key[len(prefix)+1:]] = state_dict[key].float() + return new_state_dict + +# +# def load_imagenet_pretrain(model, checkpoint_file): +# """ +# Load imagenet pretrained resnet +# Add zeros channel to the first convolution layer +# Since we have the spectral normalization, we need to do a little more +# """ +# checkpoint = torch.load(checkpoint_file, map_location = lambda storage, loc: storage.cuda(CONFIG.gpu)) +# state_dict = remove_prefix_state_dict(checkpoint['state_dict']) +# for key, value in state_dict.items(): +# state_dict[key] = state_dict[key].float() +# +# logger = logging.getLogger("Logger") +# logger.debug("Imagenet pretrained keys:") +# logger.debug(state_dict.keys()) +# logger.debug("Generator keys:") +# logger.debug(model.module.encoder.state_dict().keys()) +# logger.debug("Intersection keys:") +# logger.debug(set(model.module.encoder.state_dict().keys())&set(state_dict.keys())) +# +# weight_u = state_dict["conv1.module.weight_u"] +# weight_v = state_dict["conv1.module.weight_v"] +# weight_bar = state_dict["conv1.module.weight_bar"] +# +# logger.debug("weight_v: {}".format(weight_v)) +# logger.debug("weight_bar: {}".format(weight_bar.view(32, -1))) +# logger.debug("sigma: {}".format(weight_u.dot(weight_bar.view(32, -1).mv(weight_v)))) +# +# new_weight_v = torch.zeros(6, 3, 3).cuda() +# new_weight_bar = torch.zeros(32, 6, 3, 3).cuda() +# +# new_weight_v[:3, :, :].copy_(weight_v.view(3, 3, 3)) +# new_weight_bar[:, :3, :, :].copy_(weight_bar) +# +# logger.debug("new weight_v: {}".format(new_weight_v.view(-1))) +# logger.debug("new weight_bar: {}".format(new_weight_bar.view(32, -1))) +# logger.debug("new sigma: {}".format(weight_u.dot(new_weight_bar.view(32, -1).mv(new_weight_v.view(-1))))) +# +# state_dict["conv1.module.weight_v"] = new_weight_v.view(-1) +# state_dict["conv1.module.weight_bar"] = new_weight_bar +# +# model.module.encoder.load_state_dict(state_dict, strict=False) + + +def load_VGG_pretrain(model, checkpoint_file): + """ + Load imagenet pretrained resnet + Add zeros channel to the first convolution layer + Since we have the spectral normalization, we need to do a little more + """ + checkpoint = torch.load(checkpoint_file, map_location = lambda storage, loc: storage.cuda()) + backbone_state_dict = remove_prefix_state_dict(checkpoint['state_dict']) + + model.module.encoder.load_state_dict(backbone_state_dict, strict=False) + + +def get_unknown_tensor(trimap): + """ + get 1-channel unknown area tensor from the 3-channel/1-channel trimap tensor + """ + # if CONFIG.model.trimap_channel == 3: + weight = trimap[:, 1:2, :, :].float() + # else: + # weight = trimap.eq(1).float() + return weight + + +def get_gaborfilter(angles): + """ + generate gabor filter as the conv kernel + :param angles: number of different angles + """ + gabor_filter = [] + for angle in range(angles): + gabor_filter.append(cv2.getGaborKernel(ksize=(5,5), sigma=0.5, theta=angle*np.pi/8, lambd=5, gamma=0.5)) + gabor_filter = np.array(gabor_filter) + gabor_filter = np.expand_dims(gabor_filter, axis=1) + return gabor_filter.astype(np.float32) + + +def get_gradfilter(): + """ + generate gradient filter as the conv kernel + """ + grad_filter = [] + grad_filter.append([[-1, -2, -1], [0, 0, 0], [1, 2, 1]]) + grad_filter.append([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]) + grad_filter = np.array(grad_filter) + grad_filter = np.expand_dims(grad_filter, axis=1) + return grad_filter.astype(np.float32) + + +# def reduce_tensor_dict(tensor_dict, mode='mean'): +# """ +# average tensor dict over different GPUs +# """ +# for key, tensor in tensor_dict.items(): +# if tensor is not None: +# tensor_dict[key] = reduce_tensor(tensor, mode) +# return tensor_dict +# +# +# def reduce_tensor(tensor, mode='mean'): +# """ +# average tensor over different GPUs +# """ +# rt = tensor.clone() +# dist.all_reduce(rt, op=dist.ReduceOp.SUM) +# if mode == 'mean': +# rt /= CONFIG.world_size +# elif mode == 'sum': +# pass +# else: +# raise NotImplementedError("reduce mode can only be 'mean' or 'sum'") +# return rt + +def make_color_wheel(): + # from https://github.com/JiahuiYu/generative_inpainting/blob/master/inpaint_ops.py + RY, YG, GC, CB, BM, MR = (15, 6, 4, 11, 13, 6) + ncols = RY + YG + GC + CB + BM + MR + colorwheel = np.zeros([ncols, 3]) + col = 0 + # RY + colorwheel[0:RY, 0] = 255 + colorwheel[0:RY, 1] = np.transpose(np.floor(255*np.arange(0, RY) / RY)) + col += RY + # YG + colorwheel[col:col+YG, 0] = 255 - np.transpose(np.floor(255*np.arange(0, YG) / YG)) + colorwheel[col:col+YG, 1] = 255 + col += YG + # GC + colorwheel[col:col+GC, 1] = 255 + colorwheel[col:col+GC, 2] = np.transpose(np.floor(255*np.arange(0, GC) / GC)) + col += GC + # CB + colorwheel[col:col+CB, 1] = 255 - np.transpose(np.floor(255*np.arange(0, CB) / CB)) + colorwheel[col:col+CB, 2] = 255 + col += CB + # BM + colorwheel[col:col+BM, 2] = 255 + colorwheel[col:col+BM, 0] = np.transpose(np.floor(255*np.arange(0, BM) / BM)) + col += + BM + # MR + colorwheel[col:col+MR, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, MR) / MR)) + colorwheel[col:col+MR, 0] = 255 + return colorwheel + + +COLORWHEEL = make_color_wheel() + +def compute_color(u,v): + # from https://github.com/JiahuiYu/generative_inpainting/blob/master/inpaint_ops.py + h, w = u.shape + img = np.zeros([h, w, 3]) + nanIdx = np.isnan(u) | np.isnan(v) + u[nanIdx] = 0 + v[nanIdx] = 0 + colorwheel = COLORWHEEL + # colorwheel = make_color_wheel() + ncols = np.size(colorwheel, 0) + rad = np.sqrt(u**2+v**2) + a = np.arctan2(-v, -u) / np.pi + fk = (a+1) / 2 * (ncols - 1) + 1 + k0 = np.floor(fk).astype(int) + k1 = k0 + 1 + k1[k1 == ncols+1] = 1 + f = fk - k0 + for i in range(np.size(colorwheel,1)): + tmp = colorwheel[:, i] + col0 = tmp[k0-1] / 255 + col1 = tmp[k1-1] / 255 + col = (1-f) * col0 + f * col1 + idx = rad <= 1 + col[idx] = 1-rad[idx]*(1-col[idx]) + notidx = np.logical_not(idx) + col[notidx] *= 0.75 + img[:, :, i] = np.uint8(np.floor(255 * col*(1-nanIdx))) + return img + +def flow_to_image(flow): + # part from https://github.com/JiahuiYu/generative_inpainting/blob/master/inpaint_ops.py + maxrad = -1 + u = flow[0, :, :] + v = flow[1, :, :] + rad = np.sqrt(u ** 2 + v ** 2) + maxrad = max(maxrad, np.max(rad)) + u = u/(maxrad + np.finfo(float).eps) + v = v/(maxrad + np.finfo(float).eps) + img = compute_color(u, v) + + return img + + +if __name__ == "__main__": + import networks + logging.basicConfig(level=logging.DEBUG, format='[%(asctime)s] %(levelname)s: %(message)s', + datefmt='%m-%d %H:%M:%S') + G = networks.get_generator().cuda() + # load_imagenet_pretrain(G, CONFIG.model.imagenet_pretrain_path) + x = torch.randn(4,3,512,512).cuda() + y = torch.randn(4,3,512,512).cuda() + z = G(x, y) diff --git a/hair_service_sd/core/utils/utils_3ddfa.py b/hair_service_sd/core/utils/utils_3ddfa.py new file mode 100644 index 0000000..c9d8dcf --- /dev/null +++ b/hair_service_sd/core/utils/utils_3ddfa.py @@ -0,0 +1,461 @@ +import numpy as np +import math +import scipy.io as scio +import random +from core import face3d +from core.face3d import mesh +from core.utils import params_3ddfa +import cv2 + + +def RotationMatrix(angle_x, angle_y, angle_z): + phi = angle_x + gamma = angle_y + theta = angle_z + R_x = np.array([[1, 0, 0], [0, math.cos(phi), math.sin(phi)], [0, -math.sin(phi), math.cos(phi)]]) + R_y = np.array([[math.cos(gamma), 0, -math.sin(gamma)], [0, 1, 0], [math.sin(gamma), 0, math.cos(gamma)]]) + R_z = np.array([[math.cos(theta), math.sin(theta), 0], [-math.sin(theta), math.cos(theta), 0], [0, 0, 1]]) + R = R_z @ R_x @ R_y + return R + +def process_uv(uv_coords, uv_h=256, uv_w=256): + uv_coords[:, 0] = uv_coords[:, 0] * (uv_w - 1) + uv_coords[:, 1] = uv_coords[:, 1] * (uv_h - 1) + uv_coords[:, 1] = uv_h - uv_coords[:, 1] - 1 + uv_coords = np.hstack((uv_coords, np.zeros((uv_coords.shape[0], 1)))) # add z + return uv_coords + +def transform_params(params, tranform_mat, dst_img_height, src_img_height=256): + params_convert = np.zeros(76) + s_ori = np.sqrt(tranform_mat[:, 0:2].dot(np.transpose(tranform_mat[:, 0:2]))[0, 0]) + s = s_ori * params[0] + roll = -np.arcsin(tranform_mat[0, 1] / s_ori) + txy_ori = np.ones([3, 1]) + txy_ori[0] = params[4] + txy_ori[1] = src_img_height - 1 - params[5] + txy = tranform_mat.dot(txy_ori) + txy[1] = dst_img_height - 1 - txy[1] + params_convert[0] = s + params_convert[1:4] = params[1:4] + params_convert[3] += roll + params_convert[4] = txy[0] + params_convert[5] = txy[1] + params_convert[6:56] = params[6:56] + params_convert[56:] = params[56:] + return params_convert +class FaceModel(object): + + def __init__(self, img_size=384): + super(FaceModel, self).__init__() + self.init_status = False + from program_conf import ConfFactory + model_path = ConfFactory.getModelValue("model_dir") + # model_path = '/home/colomi/data/PycharmProjects/zao_service_test/models/model' + # model_path = '/home/colomi/Desktop/models' + face_model = scio.loadmat(model_path+'/face_3dfa/face_model.mat') + # face_model = scio.loadmat('data/face_model.mat') + self.mu_exp = face_model['mu_exp'].astype(np.float32) + self.mu_shape = face_model['mu_shape'].astype(np.float32) + self.w = face_model['w'].astype(np.float32) + self.w_exp = face_model['w_exp'].astype(np.float32) + self.sigma = face_model['sigma'].astype(np.float32).reshape((-1)) + self.sigma_exp = face_model['sigma_exp'].astype(np.float32).reshape((-1)) + self.tex = face_model['tex'].astype(np.float32).transpose((1, 0)) + self.tri = face_model['tri'].transpose((1, 0)).astype(np.int32) - 1 + + keypoint834 = scio.loadmat(model_path+'/face_3dfa/keypt834.mat') + all_parrale = keypoint834['parrale_834'] + all_iso_index = [] + for line in all_parrale: + ddd = (line[0] - 1).reshape((-1)) + all_iso_index.extend(((line[0] - 1).reshape((-1))).astype(np.int32).tolist()) + tmp_mu = self.mu_exp + self.mu_shape + tmp_mu = tmp_mu.reshape((-1, 3)) + # np.savetxt(r'H:\tmp\obj\tmp_mu.txt', tmp_mu, fmt='%f', delimiter=' ') + tmp_mu = tmp_mu[all_iso_index] + # np.savetxt(r'H:\tmp\obj\tmp_iso.txt', tmp_mu, fmt='%f', delimiter=' ') + + # --load uv coords + uv_coords = face3d.load.load_uv_coords(model_path+'/face_3dfa/BFM_UV.mat') + uv_h = uv_w = img_size + self.image_h = self.image_w = img_size + self.uv_coords = process_uv(uv_coords, uv_h, uv_w) + + trim_face = scio.loadmat(model_path+'/face_3dfa/trim_face.mat') + self.trim_idx = (trim_face['idx_face'].astype(np.int32) - 1).reshape((-1)) + self.trim_tri = trim_face['tri_face'].transpose((1, 0)).astype(np.int32) - 1 + self.trim_tri = self.trim_tri[:, ::-1] + keypoints834 = scio.loadmat(model_path+'/face_3dfa/keypt834.mat') + self.index_834 = (keypoints834['index_final'].astype(np.int32) - 1).reshape((-1)) + + self.mu_exp = self.mu_exp.reshape((-1, 3))[self.trim_idx, :].reshape((-1, 1)) + self.mu_shape = self.mu_shape.reshape((-1, 3))[self.trim_idx, :].reshape((-1, 1)) + self.w = self.w.reshape((-1, 3, 50))[self.trim_idx, :, :].reshape((-1, 50)) + self.w_exp = self.w_exp.reshape((-1, 3, 20))[self.trim_idx, :, :].reshape((-1, 20)) + self.tri = self.trim_tri + self.uv_coords = self.uv_coords[self.trim_idx] + + self.mu = self.mu_exp + self.mu_shape + + self.index30kTo137 = np.loadtxt(model_path+'/face_3dfa/index30kTo137.txt', dtype=np.int32) + self.init_status = True + + def running(self): + return self.init_status + + def draw_normal_depth_map_full_params(self, params, img_w, img_h, target_img=None, params_need_scale=False): + if params_need_scale: + f = params[0] / params_3ddfa.SCALE_F + phi = params[1] / params_3ddfa.SCALE_ROTATE + gamma = params[2] / params_3ddfa.SCALE_ROTATE + theta = params[3] / params_3ddfa.SCALE_ROTATE + t3d = np.array([params[4], params[5], 0]) / params_3ddfa.SCALE_OFFSET + alpha = (params[6:56] / params_3ddfa.SCALE_SHAPE)[:, np.newaxis] + alpha_exp = (params[56:] / params_3ddfa.SCALE_EXP)[:, np.newaxis] + else: + f = params[0] + phi = params[1] + gamma = params[2] + theta = params[3] + t3d = np.array([params[4], params[5], 0]) + alpha = params[6:56, np.newaxis] + alpha_exp = params[56:, np.newaxis] + + express3d = self.w_exp @ alpha_exp + express3d = np.reshape(express3d, (-1, 3)).transpose((1, 0)) + shape3d = self.mu_shape + self.mu_exp + self.w @ alpha + shape3d = np.reshape(shape3d, (-1, 3)).transpose((1, 0)) + vertex3d = shape3d + express3d + R = RotationMatrix(phi, gamma, theta) + project_vertex = f * (R @ vertex3d) + t3d.reshape((3, 1)) + project_vertex = project_vertex.transpose((1, 0)) + project_vertex[:, 1] = img_h - 1 - project_vertex[:, 1] + normal = face3d.mesh.light.get_normal(project_vertex, self.tri) + normal = (normal + 1) / 2 + + # np.savetxt('uv_coords.txt', self.uv_coords) + # np.savetxt('trim_tri.txt', self.trim_tri) + # np.savetxt('trim_tex.txt', trim_tex) + normal_map = mesh.render.render_colors(project_vertex, self.tri, normal, img_h, img_w, 3, BG=target_img) + # cv2.imwrite('uv_texture_map.png', uv_texture_map) + + z = project_vertex[:, 2:] + z = z - np.min(z) + z = z / np.max(z) + depth_map = mesh.render.render_colors(project_vertex, self.tri, z, img_h, img_w, 1) + + return normal_map, depth_map, project_vertex + + def draw_normal_depth_map_full_params_with_bellus(self, params, bellus_shape, pts137_2d, img_w, img_h, + target_img=None, params_need_scale=False): + if params_need_scale: + f = params[0] / params_3ddfa.SCALE_F + phi = params[1] / params_3ddfa.SCALE_ROTATE + gamma = params[2] / params_3ddfa.SCALE_ROTATE + theta = params[3] / params_3ddfa.SCALE_ROTATE + t3d = np.array([params[4], params[5], 0]) / params_3ddfa.SCALE_OFFSET + alpha = (params[6:56] / params_3ddfa.SCALE_SHAPE)[:, np.newaxis] + alpha_exp = (params[56:] / params_3ddfa.SCALE_EXP)[:, np.newaxis] + else: + f = params[0] + phi = params[1] + gamma = params[2] + theta = params[3] + t3d = np.array([params[4], params[5], 0]) + alpha = params[6:56, np.newaxis] + alpha_exp = params[56:, np.newaxis] + + express3d = self.w_exp @ alpha_exp + express3d = np.reshape(express3d, (-1, 3)).transpose((1, 0)) + + shape3d = bellus_shape + + # shape3d = self.mu_shape + self.mu_exp + self.w @ alpha + + shape3d = np.reshape(shape3d, (-1, 3)).transpose((1, 0)) + vertex3d = shape3d + express3d + R = RotationMatrix(phi, gamma, theta) + project_vertex = f * (R @ vertex3d) + t3d.reshape((3, 1)) + project_vertex = project_vertex.transpose((1, 0)) + project_vertex[:, 1] = img_h - 1 - project_vertex[:, 1] + + normal = face3d.mesh.light.get_normal(project_vertex, self.tri) + normal = (normal + 1) / 2 + + # np.savetxt('uv_coords.txt', self.uv_coords) + # np.savetxt('trim_tri.txt', self.trim_tri) + # np.savetxt('trim_tex.txt', trim_tex) + + valid_2d_index = [list(range(0, 9)), list(range(14, 22)), [22, 42, 36, 29], list(range(64, 87)), + list(range(88, 104)), list(range(105, 121))] + valid_2d_index = sum(valid_2d_index, []) + + tmp2_ = target_img.copy() + tmp_ = target_img.copy() + for pt in pts137_2d[valid_2d_index].astype(np.int32): + cv2.circle(tmp_, (pt[0], pt[1]), 1, (0, 1, 0), 2) + # cv2.imshow('tmp_', tmp_) + + index3d = self.index30kTo137[valid_2d_index] + # tmp_pts = project_vertex[index3d, :2] + # for pt in tmp_pts.astype(np.int32): + # cv2.circle(tmp_, (pt[0], pt[1]), 1, (0, 0, 1), 2) + # cv2.imshow('tmp_', tmp_) + + # Camera internals + + focal_length = 384 * 1.8 + center = (384 / 2, 384 / 2) + camera_matrix = np.array( + [[focal_length, 0, center[0]], + [0, focal_length, center[1]], + [0, 0, 1]], dtype="double" + ) + + express3d = self.w_exp @ alpha_exp + express3d = np.reshape(express3d, (-1, 3)).transpose((1, 0)) + shape3d = self.mu_shape + self.mu_exp + self.w @ alpha + shape3d = np.reshape(shape3d, (-1, 3)).transpose((1, 0)) + vertex3d = shape3d + express3d + R = RotationMatrix(phi, gamma, theta) + project_vertex_tmp = f * (R @ vertex3d) + t3d.reshape((3, 1)) + project_vertex_tmp = project_vertex_tmp.transpose((1, 0)) + project_vertex_tmp[:, 1] = img_h - 1 - project_vertex_tmp[:, 1] + + dist_coeffs = np.zeros((4, 1)) # Assuming no lens distortion + # (success, rotation_vector, translation_vector) = cv2.solvePnP(project_vertex[index3d], pts137_2d[valid_2d_index], camera_matrix, dist_coeffs, flags=cv2.SOLVEPNP_ITERATIVE) + (success, rotation_vector, translation_vector) = cv2.solvePnP(project_vertex[index3d], + project_vertex_tmp[index3d, :2], camera_matrix, + dist_coeffs, flags=cv2.SOLVEPNP_ITERATIVE) + + (reproject, jacobian) = cv2.projectPoints(project_vertex[index3d], rotation_vector, + translation_vector, camera_matrix, dist_coeffs) + reproject = reproject.squeeze() + for pt in reproject.astype(np.int32): + cv2.circle(tmp_, (pt[0], pt[1]), 1, (1, 0, 0), 2) + # cv2.imshow('tmp_', tmp_) + + (reproject, jacobian) = cv2.projectPoints(project_vertex, rotation_vector, + translation_vector, camera_matrix, dist_coeffs) + reproject = reproject.squeeze() + for pt in reproject.astype(np.int32): + cv2.circle(tmp2_, (pt[0], pt[1]), 1, (1, 0, 0), 1) + # cv2.imshow('tmp2_', tmp2_) + + project_vertex[:, :2] = reproject + + normal_map = mesh.render.render_colors(project_vertex, self.tri, normal, img_h, img_w, 3, BG=target_img) + # cv2.imwrite('uv_texture_map.png', uv_texture_map) + + z = project_vertex[:, 2:] + z = z - np.min(z) + z = z / np.max(z) + depth_map = mesh.render.render_colors(project_vertex, self.tri, z, img_h, img_w, 1) + + return normal_map, depth_map, project_vertex + + def draw_normal_depth_map_137(self, user_params, movie_params, img_w, img_h, params_need_scale=False, + target_img=None, alpha_value=1): + if params_need_scale: + f_movie = movie_params[0] / params_3ddfa.SCALE_F + phi_movie = movie_params[1] / params_3ddfa.SCALE_ROTATE + gamma_movie = movie_params[2] / params_3ddfa.SCALE_ROTATE + theta = movie_params[3] / params_3ddfa.SCALE_ROTATE + t3d_movie = np.array([movie_params[4], movie_params[5], 0]) / params_3ddfa.SCALE_OFFSET + alpha_movie = (movie_params[6:56] / params_3ddfa.SCALE_SHAPE)[:, np.newaxis] + alpha_exp_movie = (movie_params[56:] / params_3ddfa.SCALE_EXP)[:, np.newaxis] + alpha_user = movie_params[6:56, np.newaxis] / params_3ddfa.SCALE_SHAPE + \ + (user_params[6:56, np.newaxis] / params_3ddfa.SCALE_SHAPE - movie_params[6:56, + np.newaxis] / params_3ddfa.SCALE_SHAPE) * alpha_value + alpha_exp_user = (user_params[56:, np.newaxis] / params_3ddfa.SCALE_EXP)[:, np.newaxis] + else: + f_movie = movie_params[0] + phi_movie = movie_params[1] + gamma_movie = movie_params[2] + theta = movie_params[3] + t3d_movie = np.array([movie_params[4], movie_params[5], 0]) + alpha_movie = movie_params[6:56, np.newaxis] + alpha_exp_movie = movie_params[56:, np.newaxis] + alpha_user = movie_params[6:56, np.newaxis] + ( + user_params[6:56, np.newaxis] - movie_params[6:56, np.newaxis]) * alpha_value + alpha_exp_user = user_params[56:, np.newaxis] + + # np.savetxt(r'F:\workspace\faceswap_cpp\params.txt', params, fmt='%f', delimiter=' ') + + express3d_movie = self.mu_exp + self.w_exp @ alpha_exp_movie + express3d_movie = np.reshape(express3d_movie, (-1, 3)).transpose((1, 0)) + shape3d_user = self.mu_shape + self.w @ alpha_user + shape3d_user = np.reshape(shape3d_user, (-1, 3)).transpose((1, 0)) + vertex3d_mix = shape3d_user + express3d_movie + R_movie = RotationMatrix(phi_movie, gamma_movie, theta) + project_vertex_mix = f_movie * (R_movie @ vertex3d_mix) + t3d_movie.reshape((3, 1)) + project_vertex_mix = project_vertex_mix.transpose((1, 0)) + project_vertex_mix[:, 1] = img_h - 1 - project_vertex_mix[:, 1] + + project_vertex_normal = f_movie * (R_movie @ vertex3d_mix) + t3d_movie.reshape((3, 1)) + project_vertex_normal = project_vertex_normal.transpose((1, 0)) + project_vertex_normal[:, 1] = img_h - 1 - project_vertex_normal[:, 1] + + normal = face3d.mesh.light.get_normal(project_vertex_normal, self.tri) + normal = (normal + 1) / 2 + # normal_map = mesh.render.render_colors(project_vertex_normal, self.tri, normal, img_h, img_w, 3, BG=target_img) + # np.savetxt('uv_coords.txt', project_vertex_normal) + # np.savetxt('trim_tri.txt', self.tri) + # np.savetxt('trim_tex.txt', normal) + normal_map = mesh.render.render_colors(project_vertex_normal, self.tri, normal, img_h, img_w, 3, BG=target_img) + # cv2.imshow('normal_map', normal_map) + # cv2.waitKey() + # cv2.imwrite('normal_map.png', (normal_map * 255).astype(np.uint8)) + + z = project_vertex_normal[:, 2:] + z = z - np.min(z) + z = z / np.max(z) + gray_map = mesh.render.render_colors(project_vertex_normal, self.tri, z, img_h, img_w, 1) + + # movie + shape3d_movie = self.mu_shape + self.w @ alpha_movie + shape3d_movie = np.reshape(shape3d_movie, (-1, 3)).transpose((1, 0)) + vertex3d_movie = shape3d_movie + express3d_movie + project_vertex_movie = f_movie * (R_movie @ vertex3d_movie) + t3d_movie.reshape((3, 1)) + project_vertex_movie = project_vertex_movie.transpose((1, 0)) + project_vertex_movie[:, 1] = img_h - 1 - project_vertex_movie[:, 1] + point_usr_137_3DDFA = project_vertex_mix[self.index30kTo137, :2] + point_movie_137_3DDFA = project_vertex_movie[self.index30kTo137, :2] + + return normal_map, gray_map, point_usr_137_3DDFA, point_movie_137_3DDFA + + def get2Dpoint_137_orig(self, user_params, movie_params, img_w, img_h, params_need_scale=False, target_img=None, + alpha_value=1): + if params_need_scale: + f_movie = movie_params[0] / params_3ddfa.SCALE_F + phi_movie = movie_params[1] / params_3ddfa.SCALE_ROTATE + gamma_movie = movie_params[2] / params_3ddfa.SCALE_ROTATE + theta = movie_params[3] / params_3ddfa.SCALE_ROTATE + t3d_movie = np.array([movie_params[4], movie_params[5], 0]) / params_3ddfa.SCALE_OFFSET + alpha_movie = (movie_params[6:56] / params_3ddfa.SCALE_SHAPE)[:, np.newaxis] + alpha_exp_movie = (movie_params[56:] / params_3ddfa.SCALE_EXP)[:, np.newaxis] + alpha_user = movie_params[6:56, np.newaxis] / params_3ddfa.SCALE_SHAPE + \ + (user_params[6:56, np.newaxis] / params_3ddfa.SCALE_SHAPE - movie_params[6:56, + np.newaxis] / params_3ddfa.SCALE_SHAPE) * alpha_value + else: + f_movie = movie_params[0] + phi_movie = movie_params[1] + gamma_movie = movie_params[2] + theta = movie_params[3] + t3d_movie = np.array([movie_params[4], movie_params[5], 0]) + alpha_movie = movie_params[6:56, np.newaxis] + alpha_exp_movie = movie_params[56:, np.newaxis] + alpha_user = movie_params[6:56, np.newaxis] + ( + user_params[6:56, np.newaxis] - movie_params[6:56, np.newaxis]) * alpha_value + + # np.savetxt(r'F:\workspace\faceswap_cpp\params.txt', params, fmt='%f', delimiter=' ') + + shape3d_user = self.mu_shape + self.w @ alpha_user + shape3d_user = np.reshape(shape3d_user, (-1, 3)).transpose((1, 0)) + vertex3d_mix = shape3d_user + + R_movie = RotationMatrix(phi_movie, gamma_movie, theta) + # project_vertex_mix = f_movie * (R_movie @ vertex3d_mix) + t3d_movie.reshape((3, 1)) + project_vertex_mix = f_movie * (R_movie @ vertex3d_mix) + np.array([img_h / 2, img_w / 2, 0], + dtype=np.float32).reshape((3, 1)) + project_vertex_mix = project_vertex_mix.transpose((1, 0)) + project_vertex_mix[:, 1] = img_h - 1 - project_vertex_mix[:, 1] + + # movie + shape3d_movie = self.mu_shape + self.w @ alpha_movie + shape3d_movie = np.reshape(shape3d_movie, (-1, 3)).transpose((1, 0)) + vertex3d_movie = shape3d_movie + # project_vertex_movie = f_movie * (R_movie @ vertex3d_movie) + t3d_movie.reshape((3, 1)) + project_vertex_movie = f_movie * (R_movie @ vertex3d_movie) + np.array([img_h / 2, img_w / 2, 0], + dtype=np.float32).reshape((3, 1)) + project_vertex_movie = project_vertex_movie.transpose((1, 0)) + project_vertex_movie[:, 1] = img_h - 1 - project_vertex_movie[:, 1] + point_usr_137_3DDFA_orig = project_vertex_mix[self.index30kTo137, :2] + point_movie_137_3DDFA_orig = project_vertex_movie[self.index30kTo137, :2] + + return point_usr_137_3DDFA_orig, point_movie_137_3DDFA_orig + + def draw_uv_map(self, img_target, params, params_need_scale=False, img_w=256, img_h=256, is_train=False): + if params_need_scale: + f = params[0] / params_3ddfa.SCALE_F + phi = params[1] / params_3ddfa.SCALE_ROTATE + gamma = params[2] / params_3ddfa.SCALE_ROTATE + theta = params[3] / params_3ddfa.SCALE_ROTATE + t3d = np.array([params[4], params[5], 0]) / params_3ddfa.SCALE_OFFSET + alpha = (params[6:56] / params_3ddfa.SCALE_SHAPE)[:, np.newaxis] + alpha_exp = (params[56:] / params_3ddfa.SCALE_EXP)[:, np.newaxis] + else: + f = params[0] + phi = params[1] + gamma = params[2] + theta = params[3] + t3d = np.array([params[4], params[5], 0]) + alpha = params[6:56, np.newaxis] + alpha_exp = params[56:, np.newaxis] + + if is_train: + scale_ratio = 0.7 + random.random() * 0.3 + f = f * scale_ratio + else: + f = f * 0.85 + # np.savetxt(r'F:\workspace\faceswap_cpp\params.txt', params, fmt='%f', delimiter=' ') + + express3d = self.mu_exp + self.w_exp @ alpha_exp + express3d = np.reshape(express3d, (-1, 3)).transpose((1, 0)) + shape3d = self.mu_shape + self.w @ alpha + shape3d = np.reshape(shape3d, (-1, 3)).transpose((1, 0)) + vertex3d = shape3d + express3d + R = RotationMatrix(phi, gamma, theta) + project_vertex = f * (R @ vertex3d) + t3d.reshape((3, 1)) + project_vertex = project_vertex.transpose((1, 0)) + project_vertex[:, 1] = img_h - 1 - project_vertex[:, 1] + + trim_tex = np.zeros((project_vertex.shape[0], 3)) + project_vertex_int = project_vertex.astype(np.int32) + project_vertex_int = np.clip(project_vertex_int, 0, img_w - 1) + for j in range(project_vertex_int.shape[0]): + tmp_ = img_target[project_vertex_int[j, 1], project_vertex_int[j, 0]] + trim_tex[j] = tmp_ + + # np.savetxt('uv_coords.txt', self.uv_coords) + # np.savetxt('trim_tri.txt', self.trim_tri) + # np.savetxt('trim_tex.txt', trim_tex) + uv_texture_map = mesh.render.render_colors(self.uv_coords, self.trim_tri, trim_tex, img_h, img_w, c=3).astype( + np.uint8) + # cv2.imwrite('uv_texture_map.png', uv_texture_map) + # uv_texture_map = cv2.cvtColor(uv_texture_map, cv2.COLOR_RGB2BGR) + return uv_texture_map + + def compare_shape(self, pred_params, gt_params): + def parse_param_batch_noscale(param): + """Work for both numpy and tensor""" + N = param.shape[0] + f = param[:, 0] + R = np.zeros((N, 3, 3), dtype=np.float32) + for i in range(N): + R[i, :, :] = RotationMatrix(param[i, 1], param[i, 2], param[i, 3]) + f = f.reshape((N, 1, 1)) + p = f * R + offset = np.zeros((N, 3, 1), dtype=np.float32) + offset[:, :2, 0] = param[:, 4:6] + alpha_shp = param[:, 6:56].reshape((N, -1, 1)) + alpha_exp = param[:, 56:].reshape((N, -1, 1)) + return p, offset, alpha_shp, alpha_exp + + gt_pred_params = gt_params.copy() + gt_pred_params[:, 6:56] = pred_params[:, 6:56] + + pred_p, pred_offset, pred_alpha_shape, pred_alpha_exp = parse_param_batch_noscale(gt_pred_params) + gt_p, gt_offset, gt_alpha_shape, gt_alpha_exp = parse_param_batch_noscale(gt_params) + + N = pred_params.shape[0] + gt_vertex = gt_p @ (self.mu + self.w @ gt_alpha_shape + self.w_exp @ gt_alpha_exp) \ + .reshape((N, -1, 3)) \ + .transpose((0, 2, 1)) + gt_offset + pred_vertex = pred_p @ (self.mu + self.w @ pred_alpha_shape + self.w_exp @ pred_alpha_exp) \ + .reshape((N, -1, 3)) \ + .transpose((0, 2, 1)) + pred_offset + + diff = np.sqrt(np.sum((gt_vertex[:, :2, :] - pred_vertex[:, :2, :]) ** 2, axis=1)) + loss = np.mean(diff) + return loss \ No newline at end of file diff --git a/hair_service_sd/core/utils/weight_init.py b/hair_service_sd/core/utils/weight_init.py new file mode 100644 index 0000000..091fe13 --- /dev/null +++ b/hair_service_sd/core/utils/weight_init.py @@ -0,0 +1,32 @@ +#!/usr/bin/env python3 +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. + +import torch.nn as nn + + +def c2_xavier_fill(module: nn.Module): + """ + Initialize `module.weight` using the "XavierFill" implemented in Caffe2. + Also initializes `module.bias` to 0. + + Args: + module (torch.nn.Module): module to initialize. + """ + # Caffe2 implementation of XavierFill in fact + # corresponds to kaiming_uniform_ in PyTorch + nn.init.kaiming_uniform_(module.weight, a=1) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + +def c2_msra_fill(module: nn.Module): + """ + Initialize `module.weight` using the "MSRAFill" implemented in Caffe2. + Also initializes `module.bias` to 0. + + Args: + module (torch.nn.Module): module to initialize. + """ + nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu") + if module.bias is not None: + nn.init.constant_(module.bias, 0) diff --git a/hair_service_sd/data/front.jpg b/hair_service_sd/data/front.jpg new file mode 100644 index 0000000..cb77fe9 Binary files /dev/null and b/hair_service_sd/data/front.jpg differ diff --git a/hair_service_sd/data/hair_color_base.jpg b/hair_service_sd/data/hair_color_base.jpg new file mode 100644 index 0000000..7190905 Binary files /dev/null and b/hair_service_sd/data/hair_color_base.jpg differ diff --git a/hair_service_sd/data/mask.png b/hair_service_sd/data/mask.png new file mode 100644 index 0000000..f071476 Binary files /dev/null and b/hair_service_sd/data/mask.png differ diff --git a/hair_service_sd/data/template/config.json b/hair_service_sd/data/template/config.json new file mode 100644 index 0000000..5166992 --- /dev/null +++ b/hair_service_sd/data/template/config.json @@ -0,0 +1 @@ +{"gender": "girl", "version": "20210203", "ratio": "1"} \ No newline at end of file diff --git a/hair_service_sd/data/template/input_another_pose_hair_image.npy b/hair_service_sd/data/template/input_another_pose_hair_image.npy new file mode 100644 index 0000000..2b173c1 Binary files /dev/null and b/hair_service_sd/data/template/input_another_pose_hair_image.npy differ diff --git a/hair_service_sd/data/template/ref_baldseg_8uc3_768.png b/hair_service_sd/data/template/ref_baldseg_8uc3_768.png new file mode 100644 index 0000000..4169747 Binary files /dev/null and b/hair_service_sd/data/template/ref_baldseg_8uc3_768.png differ diff --git a/hair_service_sd/data/template/ref_landmark_f1k2_768.txt b/hair_service_sd/data/template/ref_landmark_f1k2_768.txt new file mode 100644 index 0000000..9848b43 --- /dev/null +++ b/hair_service_sd/data/template/ref_landmark_f1k2_768.txt @@ -0,0 +1,1000 @@ +3.910047607421875000e+02 4.782720031738281250e+02 +3.931997680664062500e+02 4.781407165527343750e+02 +3.951347656250000000e+02 4.777680664062500000e+02 +3.973148803710937500e+02 4.775318603515625000e+02 +3.994989013671875000e+02 4.773020629882812500e+02 +4.016657104492187500e+02 4.768167724609375000e+02 +4.035793457031250000e+02 4.760934448242187500e+02 +4.053948974609375000e+02 4.754965515136718750e+02 +4.075155639648437500e+02 4.748956909179687500e+02 +4.095679931640625000e+02 4.743517150878906250e+02 +4.113338012695312500e+02 4.732738952636718750e+02 +4.132869262695312500e+02 4.725390319824218750e+02 +4.151880493164062500e+02 4.716535644531250000e+02 +4.171122436523437500e+02 4.705672302246093750e+02 +4.189739990234375000e+02 4.694437561035156250e+02 +4.207155761718750000e+02 4.684048156738281250e+02 +4.224512329101562500e+02 4.672062988281250000e+02 +4.241761474609375000e+02 4.659199218750000000e+02 +4.258467407226562500e+02 4.646964111328125000e+02 +4.277152099609375000e+02 4.635527954101562500e+02 +4.293043823242187500e+02 4.623859252929687500e+02 +4.307888183593750000e+02 4.607523803710937500e+02 +4.325281982421875000e+02 4.594768981933593750e+02 +4.340862426757812500e+02 4.579970092773437500e+02 +4.356415405273437500e+02 4.565692138671875000e+02 +4.370926513671875000e+02 4.550765380859375000e+02 +4.385520629882812500e+02 4.536557312011718750e+02 +4.402352294921875000e+02 4.521771240234375000e+02 +4.416687011718750000e+02 4.507610778808593750e+02 +4.431079101562500000e+02 4.490620727539062500e+02 +4.442294311523437500e+02 4.473476562500000000e+02 +4.458046875000000000e+02 4.460588378906250000e+02 +4.470056762695312500e+02 4.443149108886718750e+02 +4.484832153320312500e+02 4.425657653808593750e+02 +4.498682861328125000e+02 4.411689453125000000e+02 +4.512639770507812500e+02 4.394785156250000000e+02 +4.524092407226562500e+02 4.377350158691406250e+02 +4.535640869140625000e+02 4.360216674804687500e+02 +4.550753173828125000e+02 4.344009399414062500e+02 +4.561343994140625000e+02 4.326965942382812500e+02 +4.573981323242187500e+02 4.310314025878906250e+02 +4.585409545898437500e+02 4.290661926269531250e+02 +4.594936523437500000e+02 4.273397521972656250e+02 +4.608695678710937500e+02 4.255390930175781250e+02 +4.620853881835937500e+02 4.238475646972656250e+02 +4.633486328125000000e+02 4.220666198730468750e+02 +4.640751953125000000e+02 4.202518615722656250e+02 +4.653312988281250000e+02 4.186289062500000000e+02 +4.663242187500000000e+02 4.166644897460937500e+02 +4.676533203125000000e+02 4.139621887207031250e+02 +4.687851562500000000e+02 4.111848144531250000e+02 +4.699882812500000000e+02 4.085221862792968750e+02 +4.710814819335937500e+02 4.058671569824218750e+02 +4.722793579101562500e+02 4.028602905273437500e+02 +4.734202880859375000e+02 4.001745605468750000e+02 +4.746752319335937500e+02 3.974074401855468750e+02 +4.755272216796875000e+02 3.946020812988281250e+02 +4.765712280273437500e+02 3.917247619628906250e+02 +4.775900268554687500e+02 3.888320007324218750e+02 +4.786235351562500000e+02 3.861441345214843750e+02 +4.793122558593750000e+02 3.831235961914062500e+02 +4.805646972656250000e+02 3.804020080566406250e+02 +4.811363525390625000e+02 3.776853637695312500e+02 +4.822098388671875000e+02 3.745228881835937500e+02 +4.829704589843750000e+02 3.717484130859375000e+02 +4.837393409179687500e+02 3.687555847167968750e+02 +4.847482299804687500e+02 3.660077514648437500e+02 +4.852825927734375000e+02 3.632245788574218750e+02 +4.860716552734375000e+02 3.601911926269531250e+02 +4.868696289062500000e+02 3.572536621093750000e+02 +4.875172729492187500e+02 3.541525878906250000e+02 +4.880067749023437500e+02 3.513582153320312500e+02 +4.884503173828125000e+02 3.484066467285156250e+02 +4.890531616210937500e+02 3.451851806640625000e+02 +4.897596435546875000e+02 3.425148925781250000e+02 +4.902989501953125000e+02 3.392861938476562500e+02 +4.905686035156250000e+02 3.364981689453125000e+02 +4.908748779296875000e+02 3.335456542968750000e+02 +4.912830810546875000e+02 3.305635375976562500e+02 +4.915813598632812500e+02 3.274760131835937500e+02 +4.918599853515625000e+02 3.243733215332031250e+02 +4.920621337890625000e+02 3.215319213867187500e+02 +4.922351074218750000e+02 3.185668640136718750e+02 +4.924135131835937500e+02 3.153148193359375000e+02 +4.924689331054687500e+02 3.127440490722656250e+02 +4.928894042968750000e+02 3.095771179199218750e+02 +4.928063354492187500e+02 3.063067321777343750e+02 +4.927054443359375000e+02 3.036191711425781250e+02 +4.926890869140625000e+02 3.005630493164062500e+02 +4.926867065429687500e+02 2.977387695312500000e+02 +4.926057128906250000e+02 2.944062500000000000e+02 +4.925239868164062500e+02 2.918076782226562500e+02 +4.924401855468750000e+02 2.885119323730468750e+02 +4.922729492187500000e+02 2.855444030761718750e+02 +4.920768432617187500e+02 2.823962707519531250e+02 +4.916710815429687500e+02 2.797327880859375000e+02 +4.914155883789062500e+02 2.766358337402343750e+02 +4.911994018554687500e+02 2.739381811523437500e+02 +4.908278198242187500e+02 2.704189453125000000e+02 +4.906377563476562500e+02 2.675762634277343750e+02 +4.902336425781250000e+02 2.647892150878906250e+02 +4.897935180664062500e+02 2.617091979980468750e+02 +4.894433593750000000e+02 2.587752685546875000e+02 +4.891079711914062500e+02 2.558663940429687500e+02 +4.884973144531250000e+02 2.528918762207031250e+02 +4.879754638671875000e+02 2.501762390136718750e+02 +4.873798828125000000e+02 2.474143371582031250e+02 +4.869479370117187500e+02 2.445428771972656250e+02 +4.861714477539062500e+02 2.416342163085937500e+02 +4.856568603515625000e+02 2.386917724609375000e+02 +4.848718872070312500e+02 2.361712646484375000e+02 +4.842299804687500000e+02 2.330022888183593750e+02 +4.833648681640625000e+02 2.303438720703125000e+02 +4.824683227539062500e+02 2.275167236328125000e+02 +4.815615234375000000e+02 2.246616821289062500e+02 +4.804127807617187500e+02 2.219893188476562500e+02 +4.794301147460937500e+02 2.192821350097656250e+02 +4.784436035156250000e+02 2.165496826171875000e+02 +4.770610961914062500e+02 2.140067749023437500e+02 +4.758833007812500000e+02 2.110261535644531250e+02 +4.748496704101562500e+02 2.087293701171875000e+02 +4.730812988281250000e+02 2.059562683105468750e+02 +4.718128662109375000e+02 2.035137023925781250e+02 +4.705297851562500000e+02 2.011148376464843750e+02 +4.690582885742187500e+02 1.985964660644531250e+02 +4.671060791015625000e+02 1.960516662597656250e+02 +4.652582397460937500e+02 1.937749328613281250e+02 +4.636505737304687500e+02 1.913393859863281250e+02 +4.619212036132812500e+02 1.892011108398437500e+02 +4.600159912109375000e+02 1.871256103515625000e+02 +4.576888427734375000e+02 1.848725280761718750e+02 +4.556460571289062500e+02 1.828733825683593750e+02 +4.535033569335937500e+02 1.809280395507812500e+02 +4.513851928710937500e+02 1.791069335937500000e+02 +4.492207031250000000e+02 1.773511352539062500e+02 +4.466425170898437500e+02 1.757018127441406250e+02 +4.442827148437500000e+02 1.739249877929687500e+02 +4.417765502929687500e+02 1.722103881835937500e+02 +4.393155517578125000e+02 1.709140319824218750e+02 +4.364564208984375000e+02 1.693395080566406250e+02 +4.341014404296875000e+02 1.683565368652343750e+02 +4.313510742187500000e+02 1.670603332519531250e+02 +4.288320312500000000e+02 1.657947082519531250e+02 +4.259161376953125000e+02 1.648233184814453125e+02 +4.233702392578125000e+02 1.638068847656250000e+02 +4.203672485351562500e+02 1.627771606445312500e+02 +4.175509643554687500e+02 1.619681701660156250e+02 +4.146489868164062500e+02 1.612653198242187500e+02 +4.117740478515625000e+02 1.606667480468750000e+02 +4.090285644531250000e+02 1.600175170898437500e+02 +4.060110473632812500e+02 1.596807556152343750e+02 +4.031548461914062500e+02 1.590272064208984375e+02 +4.003882446289062500e+02 1.586013488769531250e+02 +3.972160644531250000e+02 1.581306152343750000e+02 +3.945243530273437500e+02 1.581075744628906250e+02 +3.914724121093750000e+02 1.577026062011718750e+02 +3.885404357910156250e+02 1.574693298339843750e+02 +3.851850585937500000e+02 1.573929748535156250e+02 +3.822156066894531250e+02 1.578230895996093750e+02 +3.790032958984375000e+02 1.580787811279296875e+02 +3.759069824218750000e+02 1.583385620117187500e+02 +3.726146545410156250e+02 1.587290954589843750e+02 +3.698563842773437500e+02 1.589024963378906250e+02 +3.662501831054687500e+02 1.594544677734375000e+02 +3.632992858886718750e+02 1.599416809082031250e+02 +3.601953735351562500e+02 1.604930267333984375e+02 +3.573072204589843750e+02 1.612315216064453125e+02 +3.540190429687500000e+02 1.618346862792968750e+02 +3.511755065917968750e+02 1.626842346191406250e+02 +3.477412109375000000e+02 1.638114318847656250e+02 +3.447492065429687500e+02 1.645718994140625000e+02 +3.421724853515625000e+02 1.656198730468750000e+02 +3.388677978515625000e+02 1.666756286621093750e+02 +3.358128051757812500e+02 1.679354553222656250e+02 +3.329687500000000000e+02 1.695538024902343750e+02 +3.302599792480468750e+02 1.706882019042968750e+02 +3.273959960937500000e+02 1.722122497558593750e+02 +3.248004760742187500e+02 1.738005371093750000e+02 +3.220776977539062500e+02 1.755321350097656250e+02 +3.193132324218750000e+02 1.770840148925781250e+02 +3.168741455078125000e+02 1.790624084472656250e+02 +3.141356811523437500e+02 1.810252990722656250e+02 +3.115019226074218750e+02 1.829084167480468750e+02 +3.092257080078125000e+02 1.849979553222656250e+02 +3.069350280761718750e+02 1.871966857910156250e+02 +3.048293457031250000e+02 1.892181091308593750e+02 +3.025616149902343750e+02 1.916981811523437500e+02 +3.004593505859375000e+02 1.940289306640625000e+02 +2.985946655273437500e+02 1.965252075195312500e+02 +2.962413940429687500e+02 1.989359741210937500e+02 +2.944647216796875000e+02 2.012795104980468750e+02 +2.929493408203125000e+02 2.039413757324218750e+02 +2.909598999023437500e+02 2.066150512695312500e+02 +2.892709960937500000e+02 2.094328308105468750e+02 +2.880366516113281250e+02 2.120726013183593750e+02 +2.862335815429687500e+02 2.148235473632812500e+02 +2.846673583984375000e+02 2.176023254394531250e+02 +2.837269287109375000e+02 2.204185791015625000e+02 +2.821777954101562500e+02 2.232128601074218750e+02 +2.809744262695312500e+02 2.260358276367187500e+02 +2.797096557617187500e+02 2.289823913574218750e+02 +2.785889282226562500e+02 2.319153137207031250e+02 +2.775585937500000000e+02 2.347338867187500000e+02 +2.765081176757812500e+02 2.380631713867187500e+02 +2.756327514648437500e+02 2.410674438476562500e+02 +2.747382812500000000e+02 2.438032226562500000e+02 +2.737136840820312500e+02 2.468875122070312500e+02 +2.729382019042968750e+02 2.499927978515625000e+02 +2.722362670898437500e+02 2.527980346679687500e+02 +2.711337890625000000e+02 2.559432983398437500e+02 +2.708790588378906250e+02 2.592346191406250000e+02 +2.702111511230468750e+02 2.622360229492187500e+02 +2.696638183593750000e+02 2.653680114746093750e+02 +2.692619018554687500e+02 2.681720581054687500e+02 +2.687227172851562500e+02 2.712290954589843750e+02 +2.687125244140625000e+02 2.741842346191406250e+02 +2.684149780273437500e+02 2.769852294921875000e+02 +2.682730712890625000e+02 2.800444641113281250e+02 +2.681085815429687500e+02 2.829888000488281250e+02 +2.678695373535156250e+02 2.859908752441406250e+02 +2.676472167968750000e+02 2.890692443847656250e+02 +2.676504516601562500e+02 2.919165954589843750e+02 +2.677385864257812500e+02 2.947578735351562500e+02 +2.676535644531250000e+02 2.978811645507812500e+02 +2.677267456054687500e+02 3.008780212402343750e+02 +2.677018737792968750e+02 3.036448974609375000e+02 +2.680137329101562500e+02 3.067313232421875000e+02 +2.680429687500000000e+02 3.098110351562500000e+02 +2.681555175781250000e+02 3.127003173828125000e+02 +2.682051696777343750e+02 3.156906127929687500e+02 +2.686511230468750000e+02 3.186703796386718750e+02 +2.686446838378906250e+02 3.217354431152343750e+02 +2.691054992675781250e+02 3.246509399414062500e+02 +2.692631835937500000e+02 3.273279724121093750e+02 +2.699179077148437500e+02 3.305036010742187500e+02 +2.700928955078125000e+02 3.335144042968750000e+02 +2.703885803222656250e+02 3.365271606445312500e+02 +2.708822021484375000e+02 3.393918457031250000e+02 +2.713158264160156250e+02 3.423330383300781250e+02 +2.717955932617187500e+02 3.451097106933593750e+02 +2.723541564941406250e+02 3.481500244140625000e+02 +2.728264770507812500e+02 3.509885864257812500e+02 +2.734866027832031250e+02 3.536997375488281250e+02 +2.741121826171875000e+02 3.570650024414062500e+02 +2.748559570312500000e+02 3.598380126953125000e+02 +2.755387573242187500e+02 3.626068115234375000e+02 +2.760928955078125000e+02 3.656945800781250000e+02 +2.766440429687500000e+02 3.685693359375000000e+02 +2.774740600585937500e+02 3.713130798339843750e+02 +2.782913513183593750e+02 3.740830078125000000e+02 +2.792150268554687500e+02 3.772721557617187500e+02 +2.797438354492187500e+02 3.798110961914062500e+02 +2.807128906250000000e+02 3.827689208984375000e+02 +2.815528869628906250e+02 3.856153259277343750e+02 +2.824947509765625000e+02 3.882369995117187500e+02 +2.833540649414062500e+02 3.910834350585937500e+02 +2.843873291015625000e+02 3.940989990234375000e+02 +2.855980834960937500e+02 3.966461486816406250e+02 +2.866153564453125000e+02 3.994964294433593750e+02 +2.875024719238281250e+02 4.020555419921875000e+02 +2.887578735351562500e+02 4.048761596679687500e+02 +2.900878601074218750e+02 4.075897216796875000e+02 +2.914431152343750000e+02 4.103431091308593750e+02 +2.926695251464843750e+02 4.127597045898437500e+02 +2.939250488281250000e+02 4.156338195800781250e+02 +2.951120605468750000e+02 4.175483703613281250e+02 +2.964980773925781250e+02 4.196302185058593750e+02 +2.979119262695312500e+02 4.216356811523437500e+02 +2.993850708007812500e+02 4.237121887207031250e+02 +3.007662353515625000e+02 4.257619018554687500e+02 +3.023657226562500000e+02 4.278158569335937500e+02 +3.037879333496093750e+02 4.296451721191406250e+02 +3.053681335449218750e+02 4.316391601562500000e+02 +3.071166992187500000e+02 4.336812438964843750e+02 +3.086636352539062500e+02 4.355292053222656250e+02 +3.103882751464843750e+02 4.373770751953125000e+02 +3.121508789062500000e+02 4.391198425292968750e+02 +3.135588378906250000e+02 4.409849243164062500e+02 +3.157987670898437500e+02 4.427908630371093750e+02 +3.174847106933593750e+02 4.444118957519531250e+02 +3.192200317382812500e+02 4.461600036621093750e+02 +3.209035644531250000e+02 4.480081787109375000e+02 +3.229768981933593750e+02 4.494548645019531250e+02 +3.248511352539062500e+02 4.511229858398437500e+02 +3.266812133789062500e+02 4.526847229003906250e+02 +3.287023925781250000e+02 4.543664550781250000e+02 +3.308529052734375000e+02 4.558198547363281250e+02 +3.328441162109375000e+02 4.571927185058593750e+02 +3.347348632812500000e+02 4.586124267578125000e+02 +3.368096923828125000e+02 4.601754455566406250e+02 +3.388520202636718750e+02 4.614751586914062500e+02 +3.411101074218750000e+02 4.629147033691406250e+02 +3.433207397460937500e+02 4.641184387207031250e+02 +3.453942260742187500e+02 4.654356689453125000e+02 +3.475252685546875000e+02 4.664303588867187500e+02 +3.498818359375000000e+02 4.676504211425781250e+02 +3.519665527343750000e+02 4.689663696289062500e+02 +3.542409667968750000e+02 4.700110168457031250e+02 +3.565171508789062500e+02 4.709737548828125000e+02 +3.591160888671875000e+02 4.718554687500000000e+02 +3.612753906250000000e+02 4.728333129882812500e+02 +3.639052124023437500e+02 4.737481079101562500e+02 +3.663112792968750000e+02 4.742750854492187500e+02 +3.687968139648437500e+02 4.752687988281250000e+02 +3.709296875000000000e+02 4.758932189941406250e+02 +3.733759155273437500e+02 4.763841857910156250e+02 +3.758018188476562500e+02 4.770349426269531250e+02 +3.784192199707031250e+02 4.773708801269531250e+02 +3.808459472656250000e+02 4.777694702148437500e+02 +3.833079833984375000e+02 4.778888549804687500e+02 +3.861060180664062500e+02 4.780720214843750000e+02 +3.885011901855468750e+02 4.782825317382812500e+02 +3.521204833984375000e+02 4.167516784667968750e+02 +3.528087768554687500e+02 4.173226623535156250e+02 +3.534196777343750000e+02 4.183272705078125000e+02 +3.544356689453125000e+02 4.191568298339843750e+02 +3.551550903320312500e+02 4.197434692382812500e+02 +3.561969604492187500e+02 4.204717407226562500e+02 +3.567721252441406250e+02 4.212708740234375000e+02 +3.573645019531250000e+02 4.221080627441406250e+02 +3.584053344726562500e+02 4.227391967773437500e+02 +3.592808532714843750e+02 4.233082275390625000e+02 +3.600875854492187500e+02 4.241193542480468750e+02 +3.610968017578125000e+02 4.246537780761718750e+02 +3.620442199707031250e+02 4.253583984375000000e+02 +3.627587280273437500e+02 4.258468627929687500e+02 +3.638988037109375000e+02 4.265360107421875000e+02 +3.648369750976562500e+02 4.269735717773437500e+02 +3.658148803710937500e+02 4.275664062500000000e+02 +3.668965454101562500e+02 4.281504821777343750e+02 +3.677517700195312500e+02 4.287276611328125000e+02 +3.687322082519531250e+02 4.292100830078125000e+02 +3.699141235351562500e+02 4.295912170410156250e+02 +3.707911682128906250e+02 4.302386474609375000e+02 +3.719789428710937500e+02 4.305954589843750000e+02 +3.730458984375000000e+02 4.309553222656250000e+02 +3.739714355468750000e+02 4.313817138671875000e+02 +3.751072998046875000e+02 4.318345947265625000e+02 +3.761156005859375000e+02 4.321669006347656250e+02 +3.770516967773437500e+02 4.324555358886718750e+02 +3.781370849609375000e+02 4.326474609375000000e+02 +3.793493652343750000e+02 4.330824584960937500e+02 +3.802975769042968750e+02 4.332702331542968750e+02 +3.814768676757812500e+02 4.334466247558593750e+02 +3.826953735351562500e+02 4.336792297363281250e+02 +3.836239929199218750e+02 4.339237060546875000e+02 +3.848772583007812500e+02 4.342310485839843750e+02 +3.859263305664062500e+02 4.343762207031250000e+02 +3.869266357421875000e+02 4.345182189941406250e+02 +3.880844726562500000e+02 4.345042419433593750e+02 +3.892600708007812500e+02 4.345186462402343750e+02 +3.903241577148437500e+02 4.345789489746093750e+02 +3.912609252929687500e+02 4.345938110351562500e+02 +3.926121215820312500e+02 4.348665771484375000e+02 +3.936383666992187500e+02 4.347409667968750000e+02 +3.948381347656250000e+02 4.348464050292968750e+02 +3.955734252929687500e+02 4.349284973144531250e+02 +3.964852294921875000e+02 4.349472656250000000e+02 +3.977370605468750000e+02 4.348482360839843750e+02 +3.987133789062500000e+02 4.349392700195312500e+02 +3.996628417968750000e+02 4.348316040039062500e+02 +4.007484741210937500e+02 4.348317260742187500e+02 +4.017102661132812500e+02 4.347630920410156250e+02 +4.025545043945312500e+02 4.346860961914062500e+02 +4.037086791992187500e+02 4.345020446777343750e+02 +4.045956420898437500e+02 4.342686462402343750e+02 +4.055935058593750000e+02 4.341362609863281250e+02 +4.068196411132812500e+02 4.340357666015625000e+02 +4.076744384765625000e+02 4.337779541015625000e+02 +4.086466064453125000e+02 4.335263671875000000e+02 +4.096824340820312500e+02 4.331712646484375000e+02 +4.104812011718750000e+02 4.330651245117187500e+02 +4.114807128906250000e+02 4.326776428222656250e+02 +4.126278686523437500e+02 4.323018188476562500e+02 +4.134716186523437500e+02 4.318635559082031250e+02 +4.143526611328125000e+02 4.315397949218750000e+02 +4.154725952148437500e+02 4.311900939941406250e+02 +4.162621459960937500e+02 4.308366699218750000e+02 +4.171686401367187500e+02 4.303926696777343750e+02 +4.179757080078125000e+02 4.297747497558593750e+02 +4.188726196289062500e+02 4.294333190917968750e+02 +4.199852905273437500e+02 4.289374389648437500e+02 +4.206107788085937500e+02 4.284470825195312500e+02 +4.215678710937500000e+02 4.279853515625000000e+02 +4.223414306640625000e+02 4.273232727050781250e+02 +4.230730590820312500e+02 4.266233215332031250e+02 +4.238900756835937500e+02 4.262651062011718750e+02 +4.247690429687500000e+02 4.255344238281250000e+02 +4.255030517578125000e+02 4.249055786132812500e+02 +4.262783203125000000e+02 4.243013000488281250e+02 +4.269791870117187500e+02 4.236112976074218750e+02 +4.276716308593750000e+02 4.229107055664062500e+02 +4.282933349609375000e+02 4.223720397949218750e+02 +4.291115722656250000e+02 4.216362304687500000e+02 +4.297967529296875000e+02 4.209312133789062500e+02 +4.302330322265625000e+02 4.202318725585937500e+02 +4.309013061523437500e+02 4.194882507324218750e+02 +4.302960205078125000e+02 4.184688415527343750e+02 +4.296063842773437500e+02 4.176357421875000000e+02 +4.288562011718750000e+02 4.167015991210937500e+02 +4.281531982421875000e+02 4.158873901367187500e+02 +4.274471435546875000e+02 4.149989013671875000e+02 +4.266796264648437500e+02 4.141911315917968750e+02 +4.259470825195312500e+02 4.133565368652343750e+02 +4.251210937500000000e+02 4.124977416992187500e+02 +4.241779174804687500e+02 4.116471252441406250e+02 +4.234856567382812500e+02 4.108971862792968750e+02 +4.224335937500000000e+02 4.100395507812500000e+02 +4.215239868164062500e+02 4.093321838378906250e+02 +4.207324829101562500e+02 4.086655273437500000e+02 +4.199442138671875000e+02 4.080108032226562500e+02 +4.188093872070312500e+02 4.073389587402343750e+02 +4.179335937500000000e+02 4.065296936035156250e+02 +4.169863281250000000e+02 4.058353271484375000e+02 +4.158973388671875000e+02 4.053085021972656250e+02 +4.150784301757812500e+02 4.046409606933593750e+02 +4.140390625000000000e+02 4.039804382324218750e+02 +4.130779418945312500e+02 4.033994445800781250e+02 +4.119455566406250000e+02 4.027572021484375000e+02 +4.109661865234375000e+02 4.021667480468750000e+02 +4.099630737304687500e+02 4.018325500488281250e+02 +4.089788208007812500e+02 4.012012023925781250e+02 +4.078145141601562500e+02 4.007326354980468750e+02 +4.068986206054687500e+02 4.004776306152343750e+02 +4.056873168945312500e+02 3.999591064453125000e+02 +4.046847534179687500e+02 3.996046447753906250e+02 +4.035782470703125000e+02 3.994441528320312500e+02 +4.020933837890625000e+02 3.998591918945312500e+02 +4.005184326171875000e+02 4.002268371582031250e+02 +3.989674682617187500e+02 4.005736389160156250e+02 +3.975261840820312500e+02 4.010181274414062500e+02 +3.960938720703125000e+02 4.015845642089843750e+02 +3.945369873046875000e+02 4.018594055175781250e+02 +3.930253906250000000e+02 4.012137451171875000e+02 +3.913918457031250000e+02 4.009394226074218750e+02 +3.898701171875000000e+02 4.004072875976562500e+02 +3.882700805664062500e+02 3.998920898437500000e+02 +3.868731689453125000e+02 3.995216064453125000e+02 +3.854118347167968750e+02 3.987537841796875000e+02 +3.840134887695312500e+02 3.991007385253906250e+02 +3.826737670898437500e+02 3.993463745117187500e+02 +3.815607910156250000e+02 3.996596679687500000e+02 +3.804288330078125000e+02 4.000513305664062500e+02 +3.792362670898437500e+02 4.002554321289062500e+02 +3.779437255859375000e+02 4.008184814453125000e+02 +3.767947387695312500e+02 4.012247619628906250e+02 +3.755578002929687500e+02 4.017073364257812500e+02 +3.743472290039062500e+02 4.022062683105468750e+02 +3.732527465820312500e+02 4.026362609863281250e+02 +3.721012573242187500e+02 4.031929626464843750e+02 +3.709651489257812500e+02 4.038359680175781250e+02 +3.696183471679687500e+02 4.043455810546875000e+02 +3.685654296875000000e+02 4.048559265136718750e+02 +3.674141845703125000e+02 4.055009765625000000e+02 +3.663387451171875000e+02 4.062401428222656250e+02 +3.651137695312500000e+02 4.067765502929687500e+02 +3.641742248535156250e+02 4.074176025390625000e+02 +3.631356201171875000e+02 4.081200561523437500e+02 +3.620260009765625000e+02 4.088284912109375000e+02 +3.609375000000000000e+02 4.094896545410156250e+02 +3.601084289550781250e+02 4.102031250000000000e+02 +3.588945312500000000e+02 4.108911132812500000e+02 +3.578315429687500000e+02 4.116589660644531250e+02 +3.568225708007812500e+02 4.124814758300781250e+02 +3.559434814453125000e+02 4.133362426757812500e+02 +3.550498046875000000e+02 4.141184997558593750e+02 +3.540920104980468750e+02 4.149489746093750000e+02 +3.530513916015625000e+02 4.156572265625000000e+02 +3.575292053222656250e+02 4.165018615722656250e+02 +3.592604980468750000e+02 4.166200256347656250e+02 +3.610758666992187500e+02 4.165762329101562500e+02 +3.627486572265625000e+02 4.165775451660156250e+02 +3.646106567382812500e+02 4.167193908691406250e+02 +3.663956298828125000e+02 4.165784606933593750e+02 +3.682939758300781250e+02 4.164841613769531250e+02 +3.702015380859375000e+02 4.166382751464843750e+02 +3.718377685546875000e+02 4.166539916992187500e+02 +3.738204956054687500e+02 4.165650939941406250e+02 +3.756232299804687500e+02 4.167592468261718750e+02 +3.774773559570312500e+02 4.167646179199218750e+02 +3.793209228515625000e+02 4.168058471679687500e+02 +3.810430908203125000e+02 4.168282775878906250e+02 +3.828559570312500000e+02 4.168836669921875000e+02 +3.848069152832031250e+02 4.170431518554687500e+02 +3.865273437500000000e+02 4.170496215820312500e+02 +3.882499389648437500e+02 4.172210998535156250e+02 +3.900841674804687500e+02 4.173319702148437500e+02 +3.920105590820312500e+02 4.174277343750000000e+02 +3.936836547851562500e+02 4.174969177246093750e+02 +3.954889526367187500e+02 4.175051879882812500e+02 +3.969188842773437500e+02 4.175189208984375000e+02 +3.984749755859375000e+02 4.176730041503906250e+02 +3.999990234375000000e+02 4.175622863769531250e+02 +4.017514648437500000e+02 4.177708740234375000e+02 +4.032680053710937500e+02 4.178101806640625000e+02 +4.048945922851562500e+02 4.179051818847656250e+02 +4.064683227539062500e+02 4.178690490722656250e+02 +4.080950317382812500e+02 4.182566833496093750e+02 +4.096282348632812500e+02 4.181933898925781250e+02 +4.113403320312500000e+02 4.181448669433593750e+02 +4.128590087890625000e+02 4.184350585937500000e+02 +4.144869384765625000e+02 4.185231933593750000e+02 +4.160946044921875000e+02 4.186103820800781250e+02 +4.176349487304687500e+02 4.186224365234375000e+02 +4.194961547851562500e+02 4.189107360839843750e+02 +4.209221801757812500e+02 4.188243103027343750e+02 +4.226396484375000000e+02 4.189660034179687500e+02 +4.242702026367187500e+02 4.188227233886718750e+02 +4.256229248046875000e+02 4.190178222656250000e+02 +4.242955932617187500e+02 4.184355163574218750e+02 +4.227077026367187500e+02 4.177140502929687500e+02 +4.212763061523437500e+02 4.171837158203125000e+02 +4.197434082031250000e+02 4.165837402343750000e+02 +4.180620117187500000e+02 4.161244201660156250e+02 +4.166382446289062500e+02 4.156167907714843750e+02 +4.148741455078125000e+02 4.151088562011718750e+02 +4.134846801757812500e+02 4.143928527832031250e+02 +4.119786376953125000e+02 4.141283264160156250e+02 +4.102852172851562500e+02 4.137619934082031250e+02 +4.086368408203125000e+02 4.133719482421875000e+02 +4.069548339843750000e+02 4.130238342285156250e+02 +4.053729858398437500e+02 4.128310241699218750e+02 +4.037662353515625000e+02 4.125678405761718750e+02 +4.021102905273437500e+02 4.123402709960937500e+02 +4.005623779296875000e+02 4.123262634277343750e+02 +3.988540039062500000e+02 4.120371093750000000e+02 +3.971496582031250000e+02 4.121512145996093750e+02 +3.955463867187500000e+02 4.121418762207031250e+02 +3.940032348632812500e+02 4.122934265136718750e+02 +3.921443481445312500e+02 4.120115051269531250e+02 +3.902565307617187500e+02 4.118058471679687500e+02 +3.884299621582031250e+02 4.117985839843750000e+02 +3.865487060546875000e+02 4.116312561035156250e+02 +3.847177734375000000e+02 4.117279052734375000e+02 +3.829042968750000000e+02 4.116462707519531250e+02 +3.812183837890625000e+02 4.118165283203125000e+02 +3.791962585449218750e+02 4.116862792968750000e+02 +3.772154235839843750e+02 4.121430053710937500e+02 +3.755222778320312500e+02 4.121350097656250000e+02 +3.737372131347656250e+02 4.124214172363281250e+02 +3.718861694335937500e+02 4.127971496582031250e+02 +3.700177001953125000e+02 4.132359008789062500e+02 +3.681904907226562500e+02 4.134781188964843750e+02 +3.664437255859375000e+02 4.139185485839843750e+02 +3.645873413085937500e+02 4.143890380859375000e+02 +3.628308410644531250e+02 4.147691040039062500e+02 +3.610878906250000000e+02 4.152871093750000000e+02 +3.593110351562500000e+02 4.157409057617187500e+02 +3.770297851562500000e+02 3.023644409179687500e+02 +3.769004516601562500e+02 3.044742126464843750e+02 +3.769155883789062500e+02 3.066232604980468750e+02 +3.767077941894531250e+02 3.087592468261718750e+02 +3.764638366699218750e+02 3.108121337890625000e+02 +3.762814636230468750e+02 3.129788513183593750e+02 +3.759818420410156250e+02 3.149567260742187500e+02 +3.755031127929687500e+02 3.171455078125000000e+02 +3.752974853515625000e+02 3.190561218261718750e+02 +3.750824584960937500e+02 3.212796936035156250e+02 +3.746799621582031250e+02 3.233632507324218750e+02 +3.741690063476562500e+02 3.251835021972656250e+02 +3.738916320800781250e+02 3.274477844238281250e+02 +3.733205261230468750e+02 3.295006713867187500e+02 +3.728502807617187500e+02 3.314273986816406250e+02 +3.722926940917968750e+02 3.335574645996093750e+02 +3.717465820312500000e+02 3.356444091796875000e+02 +3.713123168945312500e+02 3.374712219238281250e+02 +3.706954040527343750e+02 3.396483764648437500e+02 +3.701931152343750000e+02 3.417878417968750000e+02 +3.695739746093750000e+02 3.436334228515625000e+02 +3.689635620117187500e+02 3.455536193847656250e+02 +3.683742980957031250e+02 3.476045227050781250e+02 +3.675649414062500000e+02 3.496122436523437500e+02 +3.670552368164062500e+02 3.517064208984375000e+02 +3.662591552734375000e+02 3.536615600585937500e+02 +3.655444946289062500e+02 3.554921875000000000e+02 +3.648411254882812500e+02 3.575023803710937500e+02 +3.641191406250000000e+02 3.594146423339843750e+02 +3.631390380859375000e+02 3.614341735839843750e+02 +3.623420104980468750e+02 3.633703002929687500e+02 +3.615927734375000000e+02 3.651572875976562500e+02 +3.738127746582031250e+02 3.747657775878906250e+02 +3.854759521484375000e+02 3.752520141601562500e+02 +3.957313232421875000e+02 3.789692077636718750e+02 +4.054473266601562500e+02 3.755355834960937500e+02 +4.153934326171875000e+02 3.751524963378906250e+02 +4.254553833007812500e+02 3.648309936523437500e+02 +4.246772460937500000e+02 3.629747314453125000e+02 +4.238004150390625000e+02 3.610342712402343750e+02 +4.230289916992187500e+02 3.592248840332031250e+02 +4.221752929687500000e+02 3.571557312011718750e+02 +4.215001831054687500e+02 3.551923828125000000e+02 +4.206057739257812500e+02 3.533284912109375000e+02 +4.200118408203125000e+02 3.513179626464843750e+02 +4.192122802734375000e+02 3.493679809570312500e+02 +4.186785888671875000e+02 3.474390563964843750e+02 +4.179393310546875000e+02 3.454389953613281250e+02 +4.171780395507812500e+02 3.435828552246093750e+02 +4.165917968750000000e+02 3.414419250488281250e+02 +4.159752807617187500e+02 3.394848022460937500e+02 +4.154730224609375000e+02 3.375299377441406250e+02 +4.147579956054687500e+02 3.355547485351562500e+02 +4.141737670898437500e+02 3.334330139160156250e+02 +4.136813354492187500e+02 3.314673767089843750e+02 +4.131109619140625000e+02 3.294137268066406250e+02 +4.125842285156250000e+02 3.274969482421875000e+02 +4.119983520507812500e+02 3.253722839355468750e+02 +4.116618652343750000e+02 3.234113159179687500e+02 +4.113226318359375000e+02 3.212356872558593750e+02 +4.106989135742187500e+02 3.192195129394531250e+02 +4.104084472656250000e+02 3.170971069335937500e+02 +4.100326538085937500e+02 3.152402038574218750e+02 +4.095477294921875000e+02 3.130978698730468750e+02 +4.092229003906250000e+02 3.110300903320312500e+02 +4.089182128906250000e+02 3.088378906250000000e+02 +4.087111816406250000e+02 3.067948608398437500e+02 +4.084342651367187500e+02 3.046477355957031250e+02 +4.081934814453125000e+02 3.026671752929687500e+02 +4.153786621093750000e+02 3.688175659179687500e+02 +4.063031005859375000e+02 3.699299621582031250e+02 +3.851812744140625000e+02 3.695863952636718750e+02 +3.745530090332031250e+02 3.684029541015625000e+02 +3.970175781250000000e+02 3.593204650878906250e+02 +3.968653564453125000e+02 3.572724304199218750e+02 +3.969441528320312500e+02 3.552400817871093750e+02 +3.966307983398437500e+02 3.531354675292968750e+02 +3.964943847656250000e+02 3.511116333007812500e+02 +3.963660888671875000e+02 3.489621887207031250e+02 +3.962761840820312500e+02 3.470785217285156250e+02 +3.962651367187500000e+02 3.449924316406250000e+02 +3.962250976562500000e+02 3.429673767089843750e+02 +3.960581054687500000e+02 3.408771362304687500e+02 +3.958803100585937500e+02 3.390036621093750000e+02 +3.958627929687500000e+02 3.368578491210937500e+02 +3.955939331054687500e+02 3.346062316894531250e+02 +3.955187377929687500e+02 3.327220458984375000e+02 +3.954448852539062500e+02 3.306238098144531250e+02 +3.953554687500000000e+02 3.286548156738281250e+02 +3.952377319335937500e+02 3.265642700195312500e+02 +3.953149414062500000e+02 3.243825378417968750e+02 +3.949652099609375000e+02 3.224709472656250000e+02 +3.948925170898437500e+02 3.203785095214843750e+02 +3.948510131835937500e+02 3.182610473632812500e+02 +3.947994995117187500e+02 3.163409729003906250e+02 +3.946992797851562500e+02 3.143557128906250000e+02 +3.946044921875000000e+02 3.122401123046875000e+02 +3.944739379882812500e+02 3.102077636718750000e+02 +3.943228149414062500e+02 3.081126403808593750e+02 +3.942737426757812500e+02 3.061366882324218750e+02 +3.941406860351562500e+02 3.039453430175781250e+02 +3.941163940429687500e+02 3.019885253906250000e+02 +3.939625854492187500e+02 2.998885498046875000e+02 +3.938740844726562500e+02 2.977395324707031250e+02 +3.936279907226562500e+02 2.958052368164062500e+02 +3.936812133789062500e+02 2.937925109863281250e+02 +3.370075683593750000e+02 2.964372558593750000e+02 +3.477243957519531250e+02 2.966025695800781250e+02 +3.474921264648437500e+02 2.983515319824218750e+02 +3.470780639648437500e+02 3.002081909179687500e+02 +3.463545532226562500e+02 3.018513793945312500e+02 +3.452401733398437500e+02 3.034513549804687500e+02 +3.438533935546875000e+02 3.046014709472656250e+02 +3.422747497558593750e+02 3.055956420898437500e+02 +3.406037597656250000e+02 3.065331115722656250e+02 +3.390270080566406250e+02 3.068665466308593750e+02 +3.369832763671875000e+02 3.070329895019531250e+02 +3.351445312500000000e+02 3.069108886718750000e+02 +3.332484436035156250e+02 3.063848876953125000e+02 +3.316392822265625000e+02 3.055252685546875000e+02 +3.301130371093750000e+02 3.044198608398437500e+02 +3.287681274414062500e+02 3.030424804687500000e+02 +3.278490600585937500e+02 3.016306152343750000e+02 +3.271745300292968750e+02 2.999621276855468750e+02 +3.266617431640625000e+02 2.980501403808593750e+02 +3.266002197265625000e+02 2.964011535644531250e+02 +3.265179443359375000e+02 2.945748291015625000e+02 +3.273939279785156250e+02 2.926670227050781250e+02 +3.282852783203125000e+02 2.910791015625000000e+02 +3.291219787597656250e+02 2.894396362304687500e+02 +3.304111938476562500e+02 2.879978942871093750e+02 +3.321921386718750000e+02 2.870217895507812500e+02 +3.336619567871093750e+02 2.864649047851562500e+02 +3.354833374023437500e+02 2.858363342285156250e+02 +3.373334350585937500e+02 2.858395996093750000e+02 +3.389926147460937500e+02 2.859585571289062500e+02 +3.408043823242187500e+02 2.863342590332031250e+02 +3.425104370117187500e+02 2.873048095703125000e+02 +3.441286010742187500e+02 2.883066101074218750e+02 +3.453389892578125000e+02 2.898455505371093750e+02 +3.462976379394531250e+02 2.913185119628906250e+02 +3.471510009765625000e+02 2.928807678222656250e+02 +3.476633911132812500e+02 2.946390380859375000e+02 +3.586881713867187500e+02 3.023515014648437500e+02 +3.579472045898437500e+02 3.007976684570312500e+02 +3.572587280273437500e+02 2.994201965332031250e+02 +3.564132080078125000e+02 2.980540466308593750e+02 +3.553127136230468750e+02 2.966205139160156250e+02 +3.541333618164062500e+02 2.954881896972656250e+02 +3.529988708496093750e+02 2.942803039550781250e+02 +3.517933349609375000e+02 2.930562744140625000e+02 +3.506250610351562500e+02 2.920354614257812500e+02 +3.490271911621093750e+02 2.910300292968750000e+02 +3.478405761718750000e+02 2.901397705078125000e+02 +3.463692016601562500e+02 2.895262145996093750e+02 +3.447445678710937500e+02 2.888222045898437500e+02 +3.432000732421875000e+02 2.882996826171875000e+02 +3.416818237304687500e+02 2.877415466308593750e+02 +3.400817260742187500e+02 2.873519897460937500e+02 +3.383666076660156250e+02 2.870903015136718750e+02 +3.367892456054687500e+02 2.868784179687500000e+02 +3.350701904296875000e+02 2.868544921875000000e+02 +3.334316711425781250e+02 2.868787841796875000e+02 +3.318397216796875000e+02 2.868224487304687500e+02 +3.302119750976562500e+02 2.870732421875000000e+02 +3.284316406250000000e+02 2.873132934570312500e+02 +3.270736083984375000e+02 2.877458801269531250e+02 +3.253529357910156250e+02 2.883095703125000000e+02 +3.238147277832031250e+02 2.886452026367187500e+02 +3.221575927734375000e+02 2.894009399414062500e+02 +3.207155761718750000e+02 2.900116577148437500e+02 +3.194310302734375000e+02 2.908625793457031250e+02 +3.181653442382812500e+02 2.915887756347656250e+02 +3.168468627929687500e+02 2.928825988769531250e+02 +3.155966186523437500e+02 2.938191528320312500e+02 +3.145454101562500000e+02 2.950351257324218750e+02 +3.154223632812500000e+02 2.962810363769531250e+02 +3.162520141601562500e+02 2.974008483886718750e+02 +3.172910766601562500e+02 2.986289367675781250e+02 +3.183339843750000000e+02 2.997000732421875000e+02 +3.195225830078125000e+02 3.006811828613281250e+02 +3.206997985839843750e+02 3.015661621093750000e+02 +3.220291137695312500e+02 3.023874206542968750e+02 +3.233753356933593750e+02 3.029280090332031250e+02 +3.245155029296875000e+02 3.036726379394531250e+02 +3.262146301269531250e+02 3.041983032226562500e+02 +3.276296997070312500e+02 3.046290588378906250e+02 +3.290662231445312500e+02 3.052132568359375000e+02 +3.305314331054687500e+02 3.056028442382812500e+02 +3.319151611328125000e+02 3.058771362304687500e+02 +3.334210815429687500e+02 3.057910461425781250e+02 +3.348599853515625000e+02 3.061490478515625000e+02 +3.366019287109375000e+02 3.061588439941406250e+02 +3.380304260253906250e+02 3.062843627929687500e+02 +3.396590576171875000e+02 3.061474914550781250e+02 +3.410475158691406250e+02 3.060791015625000000e+02 +3.425422363281250000e+02 3.060748291015625000e+02 +3.440933532714843750e+02 3.057759094238281250e+02 +3.456801147460937500e+02 3.056242675781250000e+02 +3.469487304687500000e+02 3.053547668457031250e+02 +3.485513305664062500e+02 3.051038818359375000e+02 +3.500129089355468750e+02 3.049170532226562500e+02 +3.514395141601562500e+02 3.045433654785156250e+02 +3.528832092285156250e+02 3.042445068359375000e+02 +3.543231811523437500e+02 3.037082214355468750e+02 +3.558008422851562500e+02 3.032126159667968750e+02 +3.572453613281250000e+02 3.027485351562500000e+02 +4.437272338867187500e+02 2.967629089355468750e+02 +4.542265014648437500e+02 2.969200439453125000e+02 +4.537081909179687500e+02 2.987091979980468750e+02 +4.533902587890625000e+02 3.005592956542968750e+02 +4.527069091796875000e+02 3.021545715332031250e+02 +4.515424194335937500e+02 3.035590820312500000e+02 +4.502153930664062500e+02 3.050177612304687500e+02 +4.486756591796875000e+02 3.059883117675781250e+02 +4.469882812500000000e+02 3.067131958007812500e+02 +4.453274536132812500e+02 3.071915588378906250e+02 +4.434951171875000000e+02 3.074196166992187500e+02 +4.417742919921875000e+02 3.071787109375000000e+02 +4.400050659179687500e+02 3.066697692871093750e+02 +4.382166137695312500e+02 3.058719177246093750e+02 +4.368788452148437500e+02 3.046940917968750000e+02 +4.355451049804687500e+02 3.034078979492187500e+02 +4.345660400390625000e+02 3.019886169433593750e+02 +4.337286987304687500e+02 3.004515991210937500e+02 +4.333356323242187500e+02 2.985209350585937500e+02 +4.330297241210937500e+02 2.968796691894531250e+02 +4.332910156250000000e+02 2.947609863281250000e+02 +4.339544067382812500e+02 2.930450744628906250e+02 +4.348093261718750000e+02 2.914680786132812500e+02 +4.357527465820312500e+02 2.900333862304687500e+02 +4.371106567382812500e+02 2.886365966796875000e+02 +4.385268554687500000e+02 2.876694335937500000e+02 +4.402557373046875000e+02 2.868555297851562500e+02 +4.419143066406250000e+02 2.864060668945312500e+02 +4.437724609375000000e+02 2.863373413085937500e+02 +4.456200561523437500e+02 2.864002685546875000e+02 +4.475073242187500000e+02 2.868964233398437500e+02 +4.489599609375000000e+02 2.876240539550781250e+02 +4.504296264648437500e+02 2.887193298339843750e+02 +4.517373046875000000e+02 2.901672058105468750e+02 +4.526318359375000000e+02 2.916263427734375000e+02 +4.534147338867187500e+02 2.933683471679687500e+02 +4.538630371093750000e+02 2.951434936523437500e+02 +4.233912353515625000e+02 3.028165588378906250e+02 +4.239356079101562500e+02 3.014843139648437500e+02 +4.245304565429687500e+02 3.001320800781250000e+02 +4.255767822265625000e+02 2.987652282714843750e+02 +4.263389892578125000e+02 2.975351867675781250e+02 +4.273240966796875000e+02 2.963959960937500000e+02 +4.286828002929687500e+02 2.954822082519531250e+02 +4.296702880859375000e+02 2.943658142089843750e+02 +4.309227905273437500e+02 2.935890808105468750e+02 +4.322281494140625000e+02 2.924817810058593750e+02 +4.336097412109375000e+02 2.917057495117187500e+02 +4.348453979492187500e+02 2.909993286132812500e+02 +4.361077270507812500e+02 2.903282775878906250e+02 +4.375888671875000000e+02 2.896724853515625000e+02 +4.389382934570312500e+02 2.891862182617187500e+02 +4.405220336914062500e+02 2.886862182617187500e+02 +4.420280761718750000e+02 2.882009582519531250e+02 +4.435343017578125000e+02 2.881647338867187500e+02 +4.451002197265625000e+02 2.877382812500000000e+02 +4.466730957031250000e+02 2.876398620605468750e+02 +4.481345825195312500e+02 2.875352478027343750e+02 +4.497673950195312500e+02 2.877914733886718750e+02 +4.512246093750000000e+02 2.878280334472656250e+02 +4.527464599609375000e+02 2.880759277343750000e+02 +4.542857666015625000e+02 2.884389648437500000e+02 +4.556397705078125000e+02 2.889220275878906250e+02 +4.570881347656250000e+02 2.893143615722656250e+02 +4.584274291992187500e+02 2.899386291503906250e+02 +4.598115844726562500e+02 2.907011718750000000e+02 +4.610299072265625000e+02 2.915794067382812500e+02 +4.621311645507812500e+02 2.925708007812500000e+02 +4.629083251953125000e+02 2.934942016601562500e+02 +4.639108886718750000e+02 2.947073974609375000e+02 +4.632938232421875000e+02 2.959713439941406250e+02 +4.625965576171875000e+02 2.971324157714843750e+02 +4.616710815429687500e+02 2.982861633300781250e+02 +4.607698364257812500e+02 2.992879333496093750e+02 +4.598544311523437500e+02 3.002314147949218750e+02 +4.586490478515625000e+02 3.012197875976562500e+02 +4.575346679687500000e+02 3.020882568359375000e+02 +4.564538574218750000e+02 3.028730773925781250e+02 +4.550679931640625000e+02 3.034374694824218750e+02 +4.538176879882812500e+02 3.040915222167968750e+02 +4.524442138671875000e+02 3.045444335937500000e+02 +4.510889282226562500e+02 3.049105834960937500e+02 +4.496782226562500000e+02 3.052831726074218750e+02 +4.481926269531250000e+02 3.054526367187500000e+02 +4.468398437500000000e+02 3.057025756835937500e+02 +4.454355468750000000e+02 3.059428405761718750e+02 +4.440140380859375000e+02 3.060188293457031250e+02 +4.424381713867187500e+02 3.061211853027343750e+02 +4.411302490234375000e+02 3.060521850585937500e+02 +4.397055053710937500e+02 3.059457092285156250e+02 +4.384063110351562500e+02 3.055949707031250000e+02 +4.370062866210937500e+02 3.058170776367187500e+02 +4.354328613281250000e+02 3.056553955078125000e+02 +4.341029052734375000e+02 3.054066162109375000e+02 +4.326022949218750000e+02 3.051792602539062500e+02 +4.313005981445312500e+02 3.048571472167968750e+02 +4.298206176757812500e+02 3.045515441894531250e+02 +4.286182861328125000e+02 3.041274719238281250e+02 +4.271894531250000000e+02 3.039939880371093750e+02 +4.258738403320312500e+02 3.035596923828125000e+02 +4.245727539062500000e+02 3.031623229980468750e+02 +4.190195922851562500e+02 2.631866149902343750e+02 +4.199384155273437500e+02 2.614685058593750000e+02 +4.210264892578125000e+02 2.597923889160156250e+02 +4.222739868164062500e+02 2.578305664062500000e+02 +4.236252441406250000e+02 2.561618041992187500e+02 +4.251543579101562500e+02 2.546264343261718750e+02 +4.268716430664062500e+02 2.530742187500000000e+02 +4.287089843750000000e+02 2.517674560546875000e+02 +4.307768554687500000e+02 2.510032043457031250e+02 +4.327308349609375000e+02 2.499283752441406250e+02 +4.347692260742187500e+02 2.493054809570312500e+02 +4.369383544921875000e+02 2.486110229492187500e+02 +4.391963500976562500e+02 2.480767822265625000e+02 +4.412813720703125000e+02 2.476284179687500000e+02 +4.435217285156250000e+02 2.471432495117187500e+02 +4.459555053710937500e+02 2.466680908203125000e+02 +4.478927001953125000e+02 2.463624267578125000e+02 +4.501729125976562500e+02 2.459187622070312500e+02 +4.522559204101562500e+02 2.454302673339843750e+02 +4.541750488281250000e+02 2.454819030761718750e+02 +4.563784790039062500e+02 2.457939453125000000e+02 +4.584494018554687500e+02 2.460440673828125000e+02 +4.604117431640625000e+02 2.462436828613281250e+02 +4.623783569335937500e+02 2.463335571289062500e+02 +4.644083251953125000e+02 2.467622070312500000e+02 +4.664375000000000000e+02 2.471948242187500000e+02 +4.684022216796875000e+02 2.478916320800781250e+02 +4.703129882812500000e+02 2.485737609863281250e+02 +4.720978393554687500e+02 2.495819702148437500e+02 +4.736296386718750000e+02 2.508729553222656250e+02 +4.750525512695312500e+02 2.520935058593750000e+02 +4.766403198242187500e+02 2.534428100585937500e+02 +4.781604614257812500e+02 2.550153198242187500e+02 +4.794153442382812500e+02 2.565595092773437500e+02 +4.807761840820312500e+02 2.580772399902343750e+02 +4.820009765625000000e+02 2.596038208007812500e+02 +4.833951416015625000e+02 2.613121337890625000e+02 +4.818289184570312500e+02 2.608347473144531250e+02 +4.802351684570312500e+02 2.603893127441406250e+02 +4.784777221679687500e+02 2.602135925292968750e+02 +4.766915893554687500e+02 2.598609924316406250e+02 +4.752946166992187500e+02 2.595364685058593750e+02 +4.734129028320312500e+02 2.592920532226562500e+02 +4.719865722656250000e+02 2.591613159179687500e+02 +4.704136962890625000e+02 2.587732849121093750e+02 +4.685772705078125000e+02 2.586391296386718750e+02 +4.668298339843750000e+02 2.584175109863281250e+02 +4.651864624023437500e+02 2.583602294921875000e+02 +4.635697021484375000e+02 2.581711425781250000e+02 +4.619950561523437500e+02 2.582466125488281250e+02 +4.601999511718750000e+02 2.584054870605468750e+02 +4.584484863281250000e+02 2.584765014648437500e+02 +4.568240966796875000e+02 2.583059997558593750e+02 +4.551835937500000000e+02 2.584429016113281250e+02 +4.534444580078125000e+02 2.584557495117187500e+02 +4.514877929687500000e+02 2.588838806152343750e+02 +4.496143798828125000e+02 2.590530395507812500e+02 +4.478853759765625000e+02 2.595275268554687500e+02 +4.458594970703125000e+02 2.599739685058593750e+02 +4.439983520507812500e+02 2.602551269531250000e+02 +4.421376953125000000e+02 2.606439514160156250e+02 +4.400956420898437500e+02 2.610155334472656250e+02 +4.381828002929687500e+02 2.613041992187500000e+02 +4.361448364257812500e+02 2.615097045898437500e+02 +4.344030151367187500e+02 2.618566284179687500e+02 +4.323035888671875000e+02 2.620546875000000000e+02 +4.307114257812500000e+02 2.622398681640625000e+02 +4.285688476562500000e+02 2.624992065429687500e+02 +4.265895385742187500e+02 2.625452880859375000e+02 +4.244920043945312500e+02 2.628674926757812500e+02 +4.227355957031250000e+02 2.628467102050781250e+02 +4.207977294921875000e+02 2.629879150390625000e+02 +2.928839721679687500e+02 2.635675048828125000e+02 +2.946077270507812500e+02 2.618333740234375000e+02 +2.964005737304687500e+02 2.602785034179687500e+02 +2.979085083007812500e+02 2.586958923339843750e+02 +2.996956787109375000e+02 2.569317626953125000e+02 +3.013282165527343750e+02 2.553381042480468750e+02 +3.032139282226562500e+02 2.538218383789062500e+02 +3.051825561523437500e+02 2.522536010742187500e+02 +3.070715332031250000e+02 2.509933776855468750e+02 +3.090723876953125000e+02 2.498932800292968750e+02 +3.111459045410156250e+02 2.490853271484375000e+02 +3.136820678710937500e+02 2.481297302246093750e+02 +3.156913757324218750e+02 2.476498718261718750e+02 +3.179194335937500000e+02 2.471979064941406250e+02 +3.205593261718750000e+02 2.468444824218750000e+02 +3.227147521972656250e+02 2.465094299316406250e+02 +3.251756591796875000e+02 2.462221984863281250e+02 +3.275546875000000000e+02 2.458421630859375000e+02 +3.299039916992187500e+02 2.457342834472656250e+02 +3.323659057617187500e+02 2.459216003417968750e+02 +3.346100769042968750e+02 2.462837524414062500e+02 +3.371808471679687500e+02 2.466419067382812500e+02 +3.394158325195312500e+02 2.469504699707031250e+02 +3.418619689941406250e+02 2.474119873046875000e+02 +3.441268005371093750e+02 2.477041625976562500e+02 +3.464942016601562500e+02 2.481552734375000000e+02 +3.487810058593750000e+02 2.486613769531250000e+02 +3.512681884765625000e+02 2.494647521972656250e+02 +3.533315429687500000e+02 2.503934936523437500e+02 +3.556939086914062500e+02 2.510322875976562500e+02 +3.576658325195312500e+02 2.523665161132812500e+02 +3.597265014648437500e+02 2.537709655761718750e+02 +3.612951049804687500e+02 2.553668823242187500e+02 +3.627548828125000000e+02 2.570787658691406250e+02 +3.642534790039062500e+02 2.587900390625000000e+02 +3.655438232421875000e+02 2.607092590332031250e+02 +3.667231140136718750e+02 2.626480407714843750e+02 +3.646480712890625000e+02 2.625995178222656250e+02 +3.624152526855468750e+02 2.625198364257812500e+02 +3.604819335937500000e+02 2.624482421875000000e+02 +3.584381408691406250e+02 2.624686279296875000e+02 +3.563471679687500000e+02 2.622581176757812500e+02 +3.540722656250000000e+02 2.621452331542968750e+02 +3.521099853515625000e+02 2.618910217285156250e+02 +3.500458374023437500e+02 2.617684631347656250e+02 +3.479780273437500000e+02 2.615444030761718750e+02 +3.457633666992187500e+02 2.611929931640625000e+02 +3.437684326171875000e+02 2.610061645507812500e+02 +3.417130737304687500e+02 2.606749267578125000e+02 +3.395646362304687500e+02 2.603321838378906250e+02 +3.375764770507812500e+02 2.599500427246093750e+02 +3.354345092773437500e+02 2.597213439941406250e+02 +3.333072509765625000e+02 2.593193664550781250e+02 +3.312770996093750000e+02 2.591120605468750000e+02 +3.293882446289062500e+02 2.587723083496093750e+02 +3.271489868164062500e+02 2.587781066894531250e+02 +3.251070556640625000e+02 2.588195495605468750e+02 +3.232276000976562500e+02 2.590333557128906250e+02 +3.210216674804687500e+02 2.590506591796875000e+02 +3.191218261718750000e+02 2.592118530273437500e+02 +3.169299926757812500e+02 2.592537841796875000e+02 +3.149517822265625000e+02 2.593973388671875000e+02 +3.130695190429687500e+02 2.596701049804687500e+02 +3.108565063476562500e+02 2.598765258789062500e+02 +3.089948425292968750e+02 2.601942443847656250e+02 +3.070269165039062500e+02 2.605869445800781250e+02 +3.049775695800781250e+02 2.609008483886718750e+02 +3.030289306640625000e+02 2.613616027832031250e+02 +3.009130554199218750e+02 2.616642456054687500e+02 +2.988822021484375000e+02 2.621275634765625000e+02 +2.968955688476562500e+02 2.626363830566406250e+02 +2.950447082519531250e+02 2.633565673828125000e+02 diff --git a/hair_service_sd/data/template/ref_matting_8uc3_768.png b/hair_service_sd/data/template/ref_matting_8uc3_768.png new file mode 100644 index 0000000..684e84c Binary files /dev/null and b/hair_service_sd/data/template/ref_matting_8uc3_768.png differ diff --git a/hair_service_sd/data/template/ref_matting_fg_8uc3_768.png b/hair_service_sd/data/template/ref_matting_fg_8uc3_768.png new file mode 100644 index 0000000..56702cf Binary files /dev/null and b/hair_service_sd/data/template/ref_matting_fg_8uc3_768.png differ diff --git a/hair_service_sd/data/template/ref_rgb_8uc3_768.png b/hair_service_sd/data/template/ref_rgb_8uc3_768.png new file mode 100644 index 0000000..1c663f9 Binary files /dev/null and b/hair_service_sd/data/template/ref_rgb_8uc3_768.png differ diff --git a/hair_service_sd/enhance_test.py b/hair_service_sd/enhance_test.py new file mode 100644 index 0000000..88ffd1b --- /dev/null +++ b/hair_service_sd/enhance_test.py @@ -0,0 +1,114 @@ +#coding:utf-8 +import torch +from uuid import uuid4 +import base64 +import os +import random +import shutil +import requests +import time +import json +import os.path as osp +from gen_super_image import webui_img2img +from gen_super_image import webui_img2img_diy +import urllib.request +import hashlib +import cv2 +from datetime import datetime +import numpy as np +from core.hairstyle_model import HairStyle_Model +import configparser +from common.logger import config + +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') + + + +def download_img(img_url, userId, isfix=False, ismask=False): + try: + img_name = osp.basename(img_url) + 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: + file_r = open(tmp_dir, 'rb') + md5_img = hashlib.md5(file_r.read()).hexdigest() + if isfix: + target_dir = osp.join(user_img_save_dir, userId, md5_img) + else: + target_dir = osp.join(user_img_save_dir, userId, md5_img) + os.makedirs(target_dir, exist_ok=True) + if ismask: + dst_file = osp.join(target_dir, md5_img + '.png') + else: + dst_file = osp.join(target_dir, md5_img + '.jpg') + shutil.move(tmp_dir, dst_file) + else: + return None, None + return dst_file, md5_img + except Exception as e: + print(e) + return None,None + +def download_img_new(img_url, userId, isfix=False, ismask=False): + try: + img_name = osp.basename(img_url) + tmp_dir = osp.join(user_img_tmp_dir, img_name) + # if osp.exists(tmp_dir): + # os.remove(tmp_dir) + print(img_url) + + r = requests.get(img_url) + # 写入图片 + with open(tmp_dir, "wb") as f: + f.write(r.content) + + + return tmp_dir, None + except Exception as e: + print(e) + return None,None + + +def hair_enhance(): + try: + # print('\n + hairEnhance input :', input) + img_path = "/home/data/hair/data/tmp/diy/1723090924587-1725452391390.jpg" + mask_path = "/home/data/hair/data/1/f1ae43b3-4788-45a6-8867-49f1458b3a08_mask.png" + # req_id = input['req_id'] + gender = "girl" + # user_id = input['user_id'] + + # 发型图 + final_img = cv2.imread(img_path) + + # mask图 + mask_dilate = cv2.imread(mask_path) + + # 获取性别 + in_gender = gender + sd_result = webui_img2img_diy(img=final_img, mask_img=mask_dilate, in_gender=in_gender, task_id="", tag="") + sd_save_path = "/home/data/hair/data/1/sd_res.png" + cv2.imwrite(sd_save_path, sd_result) + # cv2.imshow("sd_result", sd_result) + # cv2.waitKey(0) + + + except Exception as e: + print(e) + + + + +if __name__ == '__main__': + hair_diy() diff --git a/hair_service_sd/env.yaml b/hair_service_sd/env.yaml new file mode 100644 index 0000000..a1283a8 --- /dev/null +++ b/hair_service_sd/env.yaml @@ -0,0 +1,62 @@ +name: py37 +channels: + - defaults +dependencies: + - _libgcc_mutex=0.1=main + - _openmp_mutex=5.1=1_gnu + - ca-certificates=2023.01.10=h06a4308_0 + - certifi=2022.12.7=py37h06a4308_0 + - libedit=3.1.20221030=h5eee18b_0 + - libffi=3.2.1=hf484d3e_1007 + - libgcc-ng=11.2.0=h1234567_1 + - libgomp=11.2.0=h1234567_1 + - libstdcxx-ng=11.2.0=h1234567_1 + - ncurses=6.4=h6a678d5_0 + - openssl=1.0.2u=h7b6447c_0 + - pip=22.3.1=py37h06a4308_0 + - python=3.7.0=h6e4f718_3 + - readline=7.0=h7b6447c_5 + - setuptools=65.6.3=py37h06a4308_0 + - sqlite=3.33.0=h62c20be_0 + - tk=8.6.12=h1ccaba5_0 + - wheel=0.38.4=py37h06a4308_0 + - xz=5.2.10=h5eee18b_1 + - zlib=1.2.13=h5eee18b_0 + - pip: + - charset-normalizer==3.1.0 + - click==8.1.3 + - cos-python-sdk-v5==1.9.23 + - crcmod==1.7 + - cycler==0.11.0 + - flask==2.2.3 + - flask-cors==3.0.10 + - fonttools==4.38.0 + - gevent==22.10.2 + - greenlet==2.0.2 + - idna==3.4 + - importlib-metadata==6.0.0 + - itsdangerous==2.1.2 + - jinja2==3.1.2 + - kiwisolver==1.4.4 + - markupsafe==2.1.2 + - matplotlib==3.5.3 + - numpy==1.19.1 + - opencv-python==4.7.0.72 + - packaging==23.0 + - pillow==9.4.0 + - pycryptodome==3.17 + - pyparsing==3.0.9 + - python-dateutil==2.8.2 + - requests==2.28.2 + - scipy==1.7.3 + - six==1.16.0 + - torch==1.9.0+cu111 + - torchvision==0.10.0+cu111 + - typing-extensions==4.5.0 + - urllib3==1.26.15 + - werkzeug==2.2.3 + - xmltodict==0.13.0 + - zipp==3.15.0 + - zope-event==4.6 + - zope-interface==5.5.2 +prefix: /home/colo/anaconda3/envs/py37 diff --git a/hair_service_sd/face_enhance/face_enhancement.py b/hair_service_sd/face_enhance/face_enhancement.py new file mode 100644 index 0000000..55b852a --- /dev/null +++ b/hair_service_sd/face_enhance/face_enhancement.py @@ -0,0 +1,97 @@ +import os +import cv2 +import glob +import numpy as np +from utils import landmark_processor + +from face_enhance.face_gan_pt import FaceGAN +from time import time + +class FaceEnhancement(object): + def __init__(self, size=512, gpu_id=0): + 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 + ef = self.facegan.process(of) + 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) + + # cv2.imshow("img: ", img) + + # cv2.waitKey() + + 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) + + diff --git a/hair_service_sd/face_enhance/face_gan_pt.py b/hair_service_sd/face_enhance/face_gan_pt.py new file mode 100644 index 0000000..bd9d315 --- /dev/null +++ b/hair_service_sd/face_enhance/face_gan_pt.py @@ -0,0 +1,66 @@ +''' +@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 +# modelRoot = "/home/yangchaojie/Desktop/hairstyle/hairstyle_infer/weights" +modelRoot = "./weights" +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.model_dir = modelRoot + self.face_gan_model = os.path.join(self.model_dir, "face_enhance_0630.pt") + + self.model = torch.jit.load(self.face_gan_model).to(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) diff --git a/hair_service_sd/face_enhance/setup.py b/hair_service_sd/face_enhance/setup.py new file mode 100644 index 0000000..ca6bf5a --- /dev/null +++ b/hair_service_sd/face_enhance/setup.py @@ -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 \ No newline at end of file diff --git a/hair_service_sd/faceseg/face_seg.py b/hair_service_sd/faceseg/face_seg.py new file mode 100644 index 0000000..f795f50 --- /dev/null +++ b/hair_service_sd/faceseg/face_seg.py @@ -0,0 +1,75 @@ +import os +import pickle + +import torch +from faceseg.u2net import U2NET +from 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/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) + + 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 = "/mnt/DataDisk/my_projects/faceswap_hq/train_data/example" + 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["human_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() \ No newline at end of file diff --git a/hair_service_sd/faceseg/stm.py b/hair_service_sd/faceseg/stm.py new file mode 100644 index 0000000..c97834e --- /dev/null +++ b/hair_service_sd/faceseg/stm.py @@ -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) \ No newline at end of file diff --git a/hair_service_sd/faceseg/tma.py b/hair_service_sd/faceseg/tma.py new file mode 100644 index 0000000..db0b674 --- /dev/null +++ b/hair_service_sd/faceseg/tma.py @@ -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 diff --git a/hair_service_sd/faceseg/u2net.py b/hair_service_sd/faceseg/u2net.py new file mode 100644 index 0000000..91f2c9a --- /dev/null +++ b/hair_service_sd/faceseg/u2net.py @@ -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 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) diff --git a/hair_service_sd/feishu.py b/hair_service_sd/feishu.py new file mode 100644 index 0000000..090c082 --- /dev/null +++ b/hair_service_sd/feishu.py @@ -0,0 +1,41 @@ +import json + +import lark_oapi as lark +from lark_oapi.api.im.v1 import * + +def feishumessage(str): + content = f'{{"text":"{str}"}}' + + # 创建client + client = lark.Client.builder() \ + .app_id("cli_a71a255a053c500c") \ + .app_secret("3wZKMONhuGJuLHC4bzZelbb2H2ApvSLb") \ + .log_level(lark.LogLevel.DEBUG) \ + .build() + + # 构造请求对象 + request: CreateMessageRequest = CreateMessageRequest.builder() \ + .receive_id_type("chat_id") \ + .request_body(CreateMessageRequestBody.builder() + .receive_id("oc_cd69119457a806802670007c6358d621") + .msg_type("text") + .content(content) + .build()) \ + .build() + + # 发起请求 + response: CreateMessageResponse = client.im.v1.message.create(request) + + # 处理失败返回 + if not response.success(): + lark.logger.error( + f"client.im.v1.message.create failed, code: {response.code}, msg: {response.msg}, log_id: {response.get_log_id()}, resp: \n{json.dumps(json.loads(response.raw.content), indent=4, ensure_ascii=False)}") + return + + # 处理业务结果 + lark.logger.info(lark.JSON.marshal(response.data, indent=4)) + + +if __name__ == "__main__": + str = '1' + feishumessage(str) \ No newline at end of file diff --git a/hair_service_sd/gen_super_image.py b/hair_service_sd/gen_super_image.py new file mode 100644 index 0000000..1d86d9a --- /dev/null +++ b/hair_service_sd/gen_super_image.py @@ -0,0 +1,516 @@ +import io +import os.path +import time + +import cv2 +import base64 +import requests +from PIL import Image +import numpy as np +import json +from utils.call_hair_inter import call_hair_infer +from utils.call_hair_inter import call_hair_infer_diy +from common.logger import config +from uuid import uuid4 + +user_img_tmp_dir = config.get('default', 'tmp_dir') +version = config.get('default', 'version') +if version == "local": + current_webui_url = 'http://192.168.1.57:57860/' +else: + current_webui_url = 'http://0.0.0.0:57860/' + +class WebUISupersuperResolution: + def __init__(self): + self.url = f"{current_webui_url}sdapi/v1/extra-single-image" + self.body = None + + def encode_image_to_base64(self, img): + retval, bytes = cv2.imencode('.png', img) + encoded_image = base64.b64encode(bytes).decode('utf-8') + return encoded_image + + def send_request(self): + response = requests.post(url=self.url, json=self.body) + return response.json() + + # 图像初步超分 + def build_body(self, base_img): + self.body = { + # "show_extras_results": True, + # "gfpgan_visibility": 0, + "codeformer_visibility": 1, + "codeformer_weight": 1, + "upscaling_resize": 2, + # "upscaling_resize_w": 512, + # "upscaling_resize_h": 512, + "upscaling_crop": True, + "upscaler_1": "8x_NMKD-Superscale_150000_G", + "upscaler_2": "None", + "extras_upscaler_2_visibility": 0, + "image": self.encode_image_to_base64(base_img) + } + + +def interrogate(img): + url_interrogate = current_webui_url + 'sdapi/v1/interrogate' + + payload = json.dumps({ + # "model": "deepdanbooru", + "image": img + }) + headers = { + 'Content-Type': 'application/json' + } + response = requests.request("POST", url_interrogate, headers=headers, data=payload) + result = response.json()['caption'] + + return result + + +class ControlnetRequestImg2Img: + def __init__(self, prompt, net_prompt, mask_img): + self.url = f"{current_webui_url}sdapi/v1/img2img" + self.prompt = prompt + self.neg_prompt = net_prompt + self.body = None + self.mask = mask_img + + def read_mask(self): + img = self.mask + retval, bytes = cv2.imencode('.png', img) + encoded_image = base64.b64encode(bytes).decode('utf-8') + return encoded_image + + + def build_body_v2(self, dst_width, dst_height, cfg_scale, base_img, denoising_strength=0.7): + self.body = { + "prompt": self.prompt, + "negative_prompt": self.neg_prompt, + "sampler_name": "DPM++ 2M Karras", + "batch_size": 1, + "steps": 20, + "width": dst_width, + "height": dst_height, + "cfg_scale": cfg_scale, + "seed": 123456789, + "mask_blur": 11, + "init_images": [ + base_img + ], + "inpaint_full_res": False, + "inpainting_fill": 1, + "inpainting_mask_invert": 0, + "mask": self.read_mask(), + # "refiner_checkpoint": "majicmixRealistic_v7.safetensors", + # "refiner_switch_at": 0.5, + "denoising_strength": denoising_strength, + "alwayson_scripts": { + # "controlnet": { + # "args": [ + # { + # "enabled": True, + # "module": "openpose_full", + # "model": "openpose", + # "weight": 1.0, + # # "image": self.read_image(), + # "resize_mode": "Crop and Resize", + # "low_vram": False, + # "processor_res": 512, + # "guidance_start": 0.0, + # "guidance_end": 1.0, + # "control_mode": "Balanced", + # "pixel_perfect": True + # } + # ] + # } + # "controlnet": { + # "args": [ + # { + # "enabled": True, + # "module": "openpose_full", + # "model": "openpose", + # "weight": 1.0, + # "resize_mode": 1, + # "lowvram": False, + # # "processor_res": 512, + # # "guidance_start": 0.0, + # # "guidance_end": 1.0, + # # "control_mode": 0, + # # "pixel_perfect": True + # }, + # ] + # }, + } + } + # 打印去除掉图像的body + self.print_body_without_images() + + def print_body_without_images(self): + """打印body内容，但不包含init_images和mask字段""" + import copy + import json + + # 深拷贝body，避免修改原始数据 + print_body = copy.deepcopy(self.body) + + # 移除图像相关字段 + if 'init_images' in print_body: + print_body['init_images'] = [''] + if 'mask' in print_body: + print_body['mask'] = '' + + print("Request body (without images):") + print(json.dumps(print_body, indent=2, ensure_ascii=False)) + + def build_body_hr(self, dst_width, dst_height, cfg_scale, base_img, denoising_strength=0.7): + self.body = { + "prompt": self.prompt, + "negative_prompt": self.neg_prompt, + "sampler_name": "DPM++ 2M Karras", + "batch_size": 1, + "steps": 20, + "width": dst_width, + "height": dst_height, + "cfg_scale": cfg_scale, + "seed": 123456789, + "mask_blur": 11, + "init_images": [ + base_img + ], + "inpaint_full_res": False, + "inpainting_fill": 1, + "inpainting_mask_invert": 0, + "mask": self.read_mask(), + "refiner_checkpoint": "majicmixRealistic_v7.safetensors", + "refiner_switch_at": 0.5, + "denoising_strength": denoising_strength, + "alwayson_scripts": { + } + } + # 打印去除掉图像的body + self.print_body_without_images() + + def build_body_full_inpaint(self, dst_width, dst_height, cfg_scale, base_img): + self.body = { + "prompt": self.prompt, + "negative_prompt": self.neg_prompt, + "sampler_name": "DPM++ 2M Karras", + "batch_size": 1, + "steps": 20, + "width": dst_width, + "height": dst_height, + "cfg_scale": cfg_scale, + "seed": 123456789, + "mask_blur": 11, + "init_images": [ + base_img + ], + "inpaint_full_res": False, + "inpainting_fill": 1, + "inpainting_mask_invert": 0, + # "mask": self.read_mask(), + "refiner_checkpoint": "majicmixRealistic_v7.safetensors", + "refiner_switch_at": 0.5, + "denoising_strength": 0.7, + "alwayson_scripts": { + } + } + # 打印去除掉图像的body + self.print_body_without_images() + + + def build_body(self, dst_width, dst_height, cfg_scale, base_img): + self.body = { + "prompt": self.prompt, + "negative_prompt": self.neg_prompt, + "sampler_name": "DPM++ 2M Karras", + "batch_size": 1, + "steps": 30, + "width": dst_width, + "height": dst_height, + "cfg_scale": cfg_scale, + "seed": -1, + "mask_blur": 4, + "init_images": [ + base_img + ], + "inpaint_full_res": False, + "inpainting_fill": 1, + "inpainting_mask_invert": 1, + "mask": self.read_mask(), + "denoising_strength": 0.5, + "alwayson_scripts": { + "controlnet": { + "args": [ + { + "enabled": True, + "module": "openpose_full", + "model": "openpose", + "weight": 1.0, + # "image": self.read_image(), + "resize_mode": "Crop and Resize", + "low_vram": False, + "processor_res": 512, + "guidance_start": 0.0, + "guidance_end": 1.0, + "control_mode": "Balanced", + "pixel_perfect": True + } + ] + } + # "controlnet": { + # "args": [ + # { + # "enabled": True, + # "module": "openpose_full", + # "model": "openpose", + # "weight": 1.0, + # "resize_mode": 1, + # "lowvram": False, + # # "processor_res": 512, + # # "guidance_start": 0.0, + # # "guidance_end": 1.0, + # # "control_mode": 0, + # # "pixel_perfect": True + # }, + # ] + # }, + } + } + # 打印去除掉图像的body + self.print_body_without_images() + + def send_request(self): + response = requests.post(url=self.url, json=self.body) + return response.json() + + def encode_image_to_base64(self, img): + retval, bytes = cv2.imencode('.png', img) + encoded_image = base64.b64encode(bytes).decode('utf-8') + return encoded_image + + +def get_high_train_img(img_path, in_gender): + img = cv2.imread(img_path) + out_path = img_path + + # 如果图像长边尺寸小于1000，做超分 + if max(img.shape[1], img.shape[0]) < 1000: + # cv2.imshow("img orig", img) + # 发送超分请求 + img_super_res = WebUISupersuperResolution() + img_super_res.build_body(img) + print('sent hr request') + result = img_super_res.send_request()['image'] + print('Super resolution done!') + image_array = np.frombuffer(base64.b64decode(result.split(",", 1)[0]), np.uint8) + img = cv2.imdecode(image_array, cv2.IMREAD_COLOR) + print(img.shape[1], img.shape[0]) + + # cv2.imshow("img super", img) + # cv2.waitKey(0) + + # 做全图重绘到2000 + img_scale = 2000 / (max(img.shape[1], img.shape[0])) + if img_scale < 1.0: + img = cv2.resize(img, (0, 0), fx=img_scale, fy=img_scale, interpolation=cv2.INTER_LANCZOS4) + + print(img.shape[1], img.shape[0]) + + # cv2.imshow("orig", img) + + # 存储超分后的图片 + task_id = str(uuid4()) + super_image_save_path = os.path.join(user_img_tmp_dir, task_id + "_super.png") + cv2.imwrite(super_image_save_path, img) + out_path = super_image_save_path + + # todo can be del + # cv2.imshow("super", img) + # temp1 = os.path.join("/home/data/hair/data/test_tmp", str(uuid4()) + ".png") + # cv2.imwrite(temp1, img) + + # 直接请求增强接口 + # out = call_hair_enhance(super_image_save_path, "", task_id, in_gender) + # out_path = out["result"] + + print("out_path:", out_path) + return out_path + + +def super_process(in_img=None, in_mask_img=None, in_gender=None, material_save_path=None, train_lora_material_path=None, task_id=None, hair_id=None): + img = in_img + mask_img = in_mask_img + + # 如果图像长边尺寸大于1000，缩放到1000且不做超分 + if max(img.shape[1], img.shape[0]) > 1024: + scale = 1024 / max(img.shape[1], img.shape[0]) + img = cv2.resize(img, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_LANCZOS4) + + # 发送超分请求 + img_super_res = WebUISupersuperResolution() + img_super_res.build_body(img) + print('sent hr request') + result = img_super_res.send_request()['image'] + print('Super resolution done!') + image_array = np.frombuffer(base64.b64decode(result.split(",", 1)[0]), np.uint8) + img = cv2.imdecode(image_array, cv2.IMREAD_COLOR) + + print(img.shape[1], img.shape[0]) + + img_scale = 1500 / (max(img.shape[1], img.shape[0])) + if img_scale < 1.0: + img = cv2.resize(img, (0, 0), fx=img_scale, fy=img_scale, interpolation=cv2.INTER_LANCZOS4) + + print(img.shape[1], img.shape[0]) + mask_img = cv2.resize(mask_img, dsize=(img.shape[1], img.shape[0])) + + # cv2.imshow("super image", img) + super_image_save_path = os.path.join(material_save_path, "super.png") + cv2.imwrite(super_image_save_path, img) + + # cv2.imshow("mask image", mask_img) + gen_img_mask_save_path = os.path.join(material_save_path, "gen_img_mask.png") + cv2.imwrite(gen_img_mask_save_path, mask_img) + # cv2.waitKey(0) + + # 图像编码 + retval, bytes = cv2.imencode('.png', img) + encoded_image = base64.b64encode(bytes).decode('utf-8') + prompt = interrogate(encoded_image) + prompt = "" + # print("prompt: ", prompt) + + # prompt = ',easyphoto_face, easyphoto, 1person,face,suit' + + neg_prompt = '(nsfw:1.5), ng_deepnegative_v1_75t, (badhandv4:1.2), (worst quality:2), (low quality:2), (normal quality:2), lowres, bad anatomy, bad hands, ((monochrome)), ((grayscale)) watermark, moles, large breast, big breast, bad_pictures,easynegative' + if in_gender == "boy": + neg_prompt = '(nsfw:1.5),(worst quality:2),(low quality:2),(normal quality:2),lowers,normal quality,(monochrome:1.2),(grayscale:1.2),skin spots,acnes,skin blemishes,age spot,ugly face,glans,fat,missing fingers,extra fingers,extra arms,extra legs,watermark,text,error,blurry,jpeg artifacts,cropped,bad anatomy,double navel,muscle,nsfw,nude,no nipple,hair ornaments,bad_pictures,badhandv4,easynegative' + + control_net = ControlnetRequestImg2Img(prompt, neg_prompt, mask_img) + control_net.build_body(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image) + # 发送inpainting请求 + print('sent inpainting request') + + output = call_hair_infer(task_id, hair_id, train_lora_material_path, control_net.body) + print('Img2img done!') + # print(output) + result = output['images'][0] + res_img_encode = result.split(",", 1)[0] + + # image_array = np.frombuffer(base64.b64decode(res_img_encode), np.uint8) + # img_res = cv2.imdecode(image_array, cv2.IMREAD_COLOR) + # cv2.imwrite("/mnt/database2/online-server/hair-online/res_dir/90f21793-819f-46f6-91a9-d9a5259471101111.png", img_res) + # cv2.imshow("res_img:", img_res) + # cv2.waitKey(0) + + return res_img_encode + + +def encode_numpy_to_base64(img): + retval, bytes = cv2.imencode('.png', img) + encoded_image = base64.b64encode(bytes).decode('utf-8') + return encoded_image + +def webui_img2img(img=None, mask_img=None, in_gender=None, task_id=None, hair_id=None, lora_material_path=None, tag="", is_hr=False, denoising_strength=0.7, inference_port="57860"): + # url = "http://hairservice.tslead.net:57860/sdapi/v1/img2img" + + neg_prompt = '(nsfw:1.5), ng_deepnegative_v1_75t, (badhandv4:1.2), (worst quality:2), (low quality:2), (normal quality:2), lowres, bad anatomy, bad hands, ((monochrome)), ((grayscale)) watermark, moles, large breast, big breast, bad_pictures,easynegative' + if in_gender == "boy": + neg_prompt = '(nsfw:1.5),(worst quality:2),(low quality:2),(normal quality:2),lowers,normal quality,(monochrome:1.2),(grayscale:1.2),skin spots,acnes,skin blemishes,age spot,ugly face,glans,fat,missing fingers,extra fingers,extra arms,extra legs,watermark,text,error,blurry,jpeg artifacts,cropped,bad anatomy,double navel,muscle,nsfw,nude,no nipple,hair ornaments,bad_pictures,badhandv4,easynegative' + + # 图像编码 + retval, bytes = cv2.imencode('.png', img) + encoded_image = base64.b64encode(bytes).decode('utf-8') + prompt = tag + print("prompt:", prompt) + + control_net = ControlnetRequestImg2Img(prompt, neg_prompt, mask_img) + # control_net.build_body_hr(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength) + if not is_hr: + control_net.build_body_v2(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength) + else: + control_net.build_body_hr(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength) + + + # 发送inpainting请求 + print('sent inpainting request') + + start_inter = time.time() + output = call_hair_infer(task_id, hair_id, lora_material_path, control_net.body, is_hr, inference_port) + print('Img2img done!') + # print("--------------------- infer:", time.time() - start_inter) + + # print(output) + result = output['images'][0] + image = Image.open(io.BytesIO(base64.b64decode(result.split(",", 1)[0]))) + img_rgb = np.array(image) + img = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR) + return img + + +def webui_img2img_diy(img=None, mask_img=None, in_gender=None, task_id=None, tag="", inference_port="57860"): + # url = "http://hairservice.tslead.net:57860/sdapi/v1/img2img" + + neg_prompt = '(nsfw:1.5), ng_deepnegative_v1_75t, (badhandv4:1.2), (worst quality:2), (low quality:2), (normal quality:2), lowres, bad anatomy, bad hands, ((monochrome)), ((grayscale)) watermark, moles, large breast, big breast, bad_pictures,easynegative' + if in_gender == "boy": + neg_prompt = '(nsfw:1.5),(worst quality:2),(low quality:2),(normal quality:2),lowers,normal quality,(monochrome:1.2),(grayscale:1.2),skin spots,acnes,skin blemishes,age spot,ugly face,glans,fat,missing fingers,extra fingers,extra arms,extra legs,watermark,text,error,blurry,jpeg artifacts,cropped,bad anatomy,double navel,muscle,nsfw,nude,no nipple,hair ornaments,bad_pictures,badhandv4,easynegative' + + # 图像编码 + retval, bytes = cv2.imencode('.png', img) + encoded_image = base64.b64encode(bytes).decode('utf-8') + prompt = tag + print("prompt:", prompt) + + denoising_strength = 0.3 + print(f"diy strength:{denoising_strength}") + control_net = ControlnetRequestImg2Img(prompt, neg_prompt, mask_img) + control_net.build_body_v2(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength) + + # 发送inpainting请求 + print('sent inpainting request') + + start_inter = time.time() + output = call_hair_infer_diy(task_id, control_net.body, inference_port) + print('Img2img done!') + # print("--------------------- infer:", time.time() - start_inter) + + # print(output) + result = output['images'][0] + image = Image.open(io.BytesIO(base64.b64decode(result.split(",", 1)[0]))) + img_rgb = np.array(image) + img = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR) + return img + + +def webui_super_res_img(img, ratio): + url = f"{current_webui_url}sdapi/v1/extra-single-image" + request_dict = { + "resize_mode": 0, + "show_extras_results": False, + "gfpgan_visibility": 0, + "codeformer_visibility": 1, + "codeformer_weight": 1, + "upscaling_resize": ratio, + "upscaler_1": "8x_NMKD-Superscale_150000_G", + "upscale_first": False, + "image": encode_numpy_to_base64(img) + } + response = requests.post(url=url, json=request_dict) + ret_json = response.json() + result = ret_json['image'] + img = cv2.imdecode(np.frombuffer(base64.b64decode(result), np.uint8), cv2.IMREAD_COLOR) + return img + + + +if __name__ == '__main__': + in_img = cv2.imread("/mnt/database2/online-server/hair-online/res_dir/90f21793-819f-46f6-91a9-d9a525947110.png") + in_mask_img = cv2.imread("/mnt/database2/online-server/hair-online/ref_hairstyle/5cc660db-0970-4467-9ccb-8f895fcdf5be/hull_mask.png") + + # cv2.imshow("in_img", in_img) + # cv2.imshow("in_mask_img", in_mask_img) + # cv2.waitKey(0) + + super_process(in_img=in_img, in_mask_img=in_mask_img) diff --git a/hair_service_sd/gpt4v_caption.py b/hair_service_sd/gpt4v_caption.py new file mode 100644 index 0000000..1c81102 --- /dev/null +++ b/hair_service_sd/gpt4v_caption.py @@ -0,0 +1,81 @@ +import base64 +import time + +import requests +import cv2 +import json +import re +import os +import tqdm + +# OpenAI API Key +api_key = "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" + +# Function to encode the image +def encode_image(image_path): + img = cv2.imread(image_path) + scale = 500.0 / min(img.shape[:2]) + img = cv2.resize(img, (0, 0), fx=scale, fy=scale) + scaled_path = '/tmp/scaled_image.jpg' + cv2.imwrite(scaled_path, img) + with open(scaled_path, "rb") as image_file: + return base64.b64encode(image_file.read()).decode('utf-8') + + +def encode_image_v2(image_in): + img = image_in + scale = 500.0 / min(img.shape[:2]) + img = cv2.resize(img, (0, 0), fx=scale, fy=scale) + scaled_path = '/tmp/scaled_image.jpg' + cv2.imwrite(scaled_path, img) + with open(scaled_path, "rb") as image_file: + return base64.b64encode(image_file.read()).decode('utf-8') + +def caption_image(image_in): + # Getting the base64 string + base64_image = encode_image_v2(image_in) + + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {api_key}" + } + + payload = { + "model": "gpt-4-vision-preview", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "As an AI image tagging expert, please provide precise tags for the hairstyle in the image, To enhance CLIP model's understanding of the content. Please provide a detailed description of the hairstyle in the image, including but not limited to the color, style, length, curliness, hairline, highlights, gradients, etc. Your tags should be accurate, non-duplicative, and within a 10-20 word count range. Tags should be comma-separated. No need to provide any safety statements or precautions." + }, + { + "type": "image_url", + "image_url": { + "url": f"data:image/jpeg;base64,{base64_image}", + "detail": "low" + } + } + ] + } + ], + "max_tokens": 300 + } + + # try: + # response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload) + # + # tmp_json = response.json() + # return tmp_json['choices'][0]['message']['content'] + # except Exception as e: + # print(e) + # return "" + return "" + + + +if __name__ == '__main__': + prompt = caption_image('/home/chinatszrn/Downloads/abc/train_data/style1/07ebac82-4c0f-4dd2-84bd-bc34a059bd9b.png') + print(prompt) + diff --git a/hair_service_sd/gunicorn_config.py b/hair_service_sd/gunicorn_config.py new file mode 100644 index 0000000..4679150 --- /dev/null +++ b/hair_service_sd/gunicorn_config.py @@ -0,0 +1,50 @@ +#!/home/ubuntu/miniconda3/envs/yd/bin/python +# -*- coding: utf-8 -*- +import os +# 监听本机的端口 +bind = "0.0.0.0:7395" +# 未决连接的最大数量，即等待服务的客户的数量 +backlog = 2048 +# 进程数 +workers = 2 +# 线程数 +threads = 1 +# 工作模式为gevent +worker_class = 'gevent' +# 最大客户端并发数量，默认情况下这个值为1000。 +worker_connections = 1000 +# 超时 默认30秒 +timeout = 120 +# 连接上等待请求的秒数，默认情况下值为2 +keepalive = 50 +# 根目录,server.py所在目录 +# chdir = '' # 未使用，服务通过 gevent pywsgi 启动 +# 记录PID +pidfile = 'gunicorn.pid' + +def on_starting(server): + server._worker_id_overload = set() + +def nworkers_changed(server, new_value, old_value): + server._worker_id_current_workers = new_value + +def _next_worker_id(server): + if server._worker_id_overload: + return server._worker_id_overload.pop() + in_use = set(w._worker_id for w in server.WORKERS.values() if w.alive) + free = set(range(1, server._worker_id_current_workers + 1)) - in_use + return free.pop() + +def on_reload(server): + server._worker_id_overload = set(range(1, server.cfg.workers + 1)) + +def pre_fork(server, worker): + worker._worker_id = _next_worker_id(server) + +def post_fork(server, worker): + os.environ["APP_WORKER_ID"] = str(worker._worker_id) + + + + + diff --git a/hair_service_sd/hair_init.py b/hair_service_sd/hair_init.py new file mode 100644 index 0000000..64a2ec3 --- /dev/null +++ b/hair_service_sd/hair_init.py @@ -0,0 +1,296 @@ +#coding:utf-8 +import os + + + +import math +import shutil +import time + +from step05_detect_fa_hairmatting_inplace import pkl_process +from PIL import Image, ImageFont, ImageDraw +import cv2 +import numpy as np +import torch +import pickle +import json +import glob +import uuid +from gpt4v_caption import caption_image +from utils.call_hair_train import call_hair_train +import os +import os.path as osp +from common.logger import LogFactory +from core.process_modules import Process_Data, localtranslationwarpfastwithstrength,\ + Generator_Hair, chinClass, Generator_Fusion_Res, Change_Hair_Color, GenderClassifyProcessor, BodySeg, \ + localtranslationwarpfastwithstrength_v2, localtranslationwarpfastwithstrength_v2_soft, updateEndPosition +from process_modules import Get_Landmark as Get_Landmark_mtcnn +from core.face_enhance.face_enhancement import FaceEnhancement +from datetime import datetime +import random +from utils import landmark_processor +# from core.cos_module import COS_object as OSS_object +from core.oss_module import OSS_object +from core.faceseg.face_seg import FaceSeg +from common.logger import config +from process_modules import PersonProcessor_yolov5,KeypointsProcessor,Human_Keypoints,pt_conv_25_to_17 +from gen_super_image import get_high_train_img +from core.MMCVFaceRecognitionServer import MomocvFaceRecognitionServer +import process_modules + +from common.logger import config + +version = config.get('default', 'version') +if version == "local": + train_gpu_nums = 2 +else: + train_gpu_nums = 1 + +# webui_services = ["http://hairservice.tslead.net:32678/", "http://hairservice.tslead.net:32679/"] + +class HairInit(object): + __instance = None + __first_init = False + + def __new__(cls, gpu=True, use_enhance=False, infer_use_enhance=False, color_user_enhance=False): + if not cls.__instance: + cls.__instance = object.__new__(cls) + return cls.__instance + + def __init__(self, gpu=True, use_enhance=False, infer_use_enhance=False, color_user_enhance=False): + if not self.__first_init: + worker_id = int(os.environ.get('APP_WORKER_ID', 1)) + + rand_max = 9527 + self.gpu_index = (worker_id + rand_max) % train_gpu_nums + os.environ['CUDA_VISIBLE_DEVICES'] = str(self.gpu_index) + + print('current worker id {} set the gpu id :{}'.format(worker_id, self.gpu_index)) + device_id = self.gpu_index + + + self.get_landmark = Get_Landmark_mtcnn(gpu_id=device_id) + self.get_landmark_mtcnn = Get_Landmark_mtcnn(gpu_id=device_id) + self.face_recognition = MomocvFaceRecognitionServer(gpu_id=device_id) + + self.hair_size = 768 + self.color_output_size = 768 + + self.process_data = Process_Data(gpu, device_id) + self.process_data_infer = process_modules.Process_Data(gpu, device_id) + + self.generator_hair = Generator_Hair(gpu, device_id) + self.hair_fusion = Generator_Fusion_Res(gpu, device_id) + self.use_enhance = use_enhance + + self.face_enhance = FaceEnhancement(512, device_id) + self.change_haircolor = Change_Hair_Color(gpu, device_id) + self.logger_init = LogFactory.getLogger("init") + self.logger_process = LogFactory.getLogger("process") + self.logger_call = LogFactory.getLogger("call") + self.oss2 = OSS_object() + self.face_seg = FaceSeg(device_id) + self.chin_cls = chinClass(device_id) + model_path = "./weights/gender_models" + self.output_img_size = 128 + if not os.path.exists(model_path): + print("GenderClassifyProcessor don't have model!") + self.gender_model = GenderClassifyProcessor(gpu_id=device_id) + print("Load model finish ... ") + self.baseColor_dir = os.path.join(config.get('default', "haircolorDir"), + config.get('default', "baseColor_ID")) + + for i in range(2): + image = cv2.imread('data/front.jpg') + with torch.no_grad(): + tmp_dir = config.get('default', "tmp_dir") + if not os.path.exists(tmp_dir): + os.makedirs(tmp_dir) + self.infer_hairstyle(image, 'data/template', tmp_dir, 'test.jpg') + + self.effect_prepare_mask_fc32 = cv2.imread("./data/mask.png").astype(np.float32) / 255 + # body process + self.person_processor = PersonProcessor_yolov5(gpu_id=device_id) + self.keypoints_processor = KeypointsProcessor(gpu_id=device_id) + self.human_keypoint = Human_Keypoints(gpu=True, device_id=device_id) + print("Load model finish ... ") + + # infer init + # self.table_enlight = cv2.imread('./data/convert_enlight.png') + self.change_haircolor = Change_Hair_Color(gpu, device_id) + self.gender_classify = GenderClassifyProcessor(gpu_id=device_id) + + HairInit.__first_init = True + print("[HAIR_INIT] All init done. ") + + def infer_hairstyle(self, origin_img, hairstyle_dir, userinfo_dir, mask_newname): + config_path = os.path.join(hairstyle_dir, "config.json") + if not os.path.exists(config_path): + return None, None, 10001 + config_info = json.load(open(config_path, "rb")) + gender = config_info["gender"] + ratio = int(config_info["ratio"]) + # print("gender: ", gender, " ratio: ", ratio) + + user_rgb_8uc3_orisize = origin_img.copy() + tt = time.time() + landmark1k_dir = osp.join(userinfo_dir, 'kpt_1k.txt') + if not osp.exists(landmark1k_dir): + with torch.no_grad(): + landmarks_origin_img_1k = self.get_landmark.forward(origin_img) + if landmarks_origin_img_1k is None: + return None, None, 10001 + np.savetxt(landmark1k_dir, landmarks_origin_img_1k) + else: + landmarks_origin_img_1k = np.loadtxt(landmark1k_dir) + + user_bald_res_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_res_ori_raw.png') + # res_matting_8uc3_bald_orisize_dir = osp.join(userinfo_dir, 'res_matting_mask_ori.png') + user_baldseg_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_seg_ori.png') + user_baldseg_8uc3_768_dir = osp.join(userinfo_dir, 'user_baldseg_768.png') + user_bald_8uc3_768_dir = osp.join(userinfo_dir, 'bald_seg_768.png') + user_landmark_f1k2_768_dir = osp.join(userinfo_dir, 'landmark_f1k2_768.txt') + user_hairstyle_M_dir = osp.join(userinfo_dir, 'hairstyle_M.txt') + + pre_list = [user_bald_res_8uc3_orisize_dir, user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_768_dir, user_bald_8uc3_768_dir, + user_landmark_f1k2_768_dir, user_hairstyle_M_dir] + + condition_exist = True + for tmp_file in pre_list: + if not osp.exists(tmp_file): + condition_exist = False + if not condition_exist: + user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, user_baldseg_8uc3_768, user_bald_8uc3_768, \ + user_landmark_f1k2_768, user_hairstyle_M,_ = self.process_data.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio) + cv2.imwrite(user_bald_res_8uc3_orisize_dir, user_bald_res_8uc3_orisize) + cv2.imwrite(user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_orisize) + cv2.imwrite(user_baldseg_8uc3_768_dir, user_baldseg_8uc3_768) + cv2.imwrite(user_bald_8uc3_768_dir, user_bald_8uc3_768) + np.savetxt(user_landmark_f1k2_768_dir, user_landmark_f1k2_768) + np.savetxt(user_hairstyle_M_dir, user_hairstyle_M) + else: + user_bald_res_8uc3_orisize = cv2.imread(user_bald_res_8uc3_orisize_dir) + user_baldseg_8uc3_orisize = cv2.imread(user_baldseg_8uc3_orisize_dir) + # user_baldseg_8uc3_768 = cv2.imread(user_baldseg_8uc3_768_dir) + # user_bald_8uc3_768 = cv2.imread(user_bald_8uc3_768_dir) + # user_landmark_f1k2_768 = np.loadtxt(user_landmark_f1k2_768_dir) + # user_hairstyle_M = np.loadtxt(user_hairstyle_M_dir) + if ratio == 0: + user_hairstyle_M = self.process_data.get_hair_M_boy_v1(landmarks_origin_img_1k) + elif ratio == 1: + user_hairstyle_M = self.process_data.get_hair_M_girl_v1(landmarks_origin_img_1k) + elif ratio == 2: + user_hairstyle_M = self.process_data.get_hair_M_girl_v2(landmarks_origin_img_1k) + else: + user_hairstyle_M = self.process_data.get_hair_M_girl_v1(landmarks_origin_img_1k) + user_landmark_f1k2_768 = landmark_processor.transform_points(landmarks_origin_img_1k, user_hairstyle_M) + user_bald_8uc3_768 = cv2.warpAffine(user_bald_res_8uc3_orisize, user_hairstyle_M, + (self.hair_size, self.hair_size)) + user_baldseg_8uc3_768 = cv2.warpAffine(user_baldseg_8uc3_orisize, user_hairstyle_M, + (self.hair_size, self.hair_size)) + print('condition cosst:', time.time() - tt ) + # show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("show_concat", show_concat) + # cv2.waitKey() + + another_pose_hair_img_path = os.path.join(hairstyle_dir, "input_another_pose_hair_image.npy") + if not os.path.exists(another_pose_hair_img_path): + return None, None, 10002 + another_pose_hair_image = np.load(another_pose_hair_img_path) + # cv2.imshow('user_baldseg_8uc3_768', user_baldseg_8uc3_768) + # cv2.imshow('user_bald_8uc3_768', user_bald_8uc3_768) + + t0 = time.time() + # 换发型 + + hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image, + user_baldseg_8uc3_768, + user_bald_8uc3_768, + user_landmark_f1k2_768, gender) + self.logger_process.info('Generator_Hair_inference_use_pref costs:{}'.format(time.time() - t0)) + + t1 = time.time() + + hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste( + user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M) + + + res_matting_mask_ori = osp.join(userinfo_dir, mask_newname) + res_matting_mask_ori_raw = osp.join(userinfo_dir, 'res_matting_mask_ori_raw.png') + if osp.exists(res_matting_mask_ori): + os.remove(res_matting_mask_ori) + if osp.exists(res_matting_mask_ori_raw): + os.remove(res_matting_mask_ori_raw) + + self.logger_process.info('get_fusion_res_hairpaste costs:{}'.format(time.time() - t1)) + + user_res_8uc3_orisize = self.hair_fusion.inference(hair_gene_fusion_8uc3_orisize, + hair_gene_matte_8uc3_orisize, + landmarks_origin_img_1k, + user_baldseg_8uc3_orisize, + ratio) + hair_gene_matte_8uc3_orisize_cp = hair_gene_matte_8uc3_orisize.copy() + # show_concat = np.concatenate((hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize, user_baldseg_8uc3_orisize, user_res_8uc3_orisize), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_concat", show_concat) + # cv2.waitKey() + t2 = time.time() + + user_res_8uc3_orisize_for_haircolor = user_res_8uc3_orisize.copy() + + # if self.use_enhance: + user_res_8uc3_orisize_enhance = self.face_enhance.process(user_res_8uc3_orisize, landmarks_origin_img_1k) + t3 = time.time() + self.logger_process.info('use_enhance costs:{}'.format(t3 - t2)) + + user_bald_res_8uc3_orisize_enhance, user_baldseg_8uc3_orisize_enhance, user_baldseg_8uc3_768_enhance, user_bald_8uc3_768_enhance, \ + user_landmark_f1k2_768_enhance, user_hairstyle_M_enhance, user_matting_8uc3_bald_orisize = self.process_data.get_prepare_user_768_data( + user_res_8uc3_orisize_enhance, landmarks_origin_img_1k, ratio=ratio) + + t4 = time.time() + self.logger_process.info('gen bald costs:{}'.format(t4 - t3)) + + # 重新提取matting + _, hair_gene_matte_8uC0_orisize, _ = self.process_data.generator_matte.matte_inference(user_res_8uc3_orisize_enhance, landmarks_origin_img_1k) + t5 = time.time() + self.logger_process.info('gen matte_inference costs:{}'.format(t5 - t4)) + hair_gene_matte_8uc3_orisize = np.repeat(hair_gene_matte_8uC0_orisize[:, :, np.newaxis], 3, axis=2) + hair_gene_matte_8uc3_orisize = cv2.blur(hair_gene_matte_8uc3_orisize, (3, 3)) + + # user_res_8uc3_orisize = user_bald_res_8uc3_orisize_enhance.astype(np.float32)/255. * (1. - hair_gene_matte_8uc3_orisize.astype(np.float32) / 255.) + user_res_8uc3_orisize_enhance.astype(np.float32)/255. * hair_gene_matte_8uc3_orisize.astype(np.float32) / 255. + # user_res_8uc3_orisize = (user_res_8uc3_orisize*255).astype(np.uint8) + + user_res_fc32_orisize_enhance = user_res_8uc3_orisize_enhance.astype(np.float32)/255. + user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32)/255. + hair_gene_matte_fc32_orisize = hair_gene_matte_8uc3_orisize.astype(np.float32)/255. + hair_gene_matte_fc32_orisize = cv2.GaussianBlur(hair_gene_matte_fc32_orisize, (11, 11), 0, 0) + user_res_fc32_orisize_enhance = user_res_fc32_orisize_enhance * hair_gene_matte_fc32_orisize + user_res_fc32_orisize * (1 - hair_gene_matte_fc32_orisize) + user_res_8uc3_orisize_enhance = (user_res_fc32_orisize_enhance * 255).astype(np.uint8) + + # mid_show = np.concatenate((user_res_fc32_orisize_enhance, user_res_8uc3_orisize_enhance, hair_gene_matte_8uc3_orisize, user_res_8uc3_orisize), axis=1) + # ratio = 1536. / max(mid_show.shape[:2]) + # mid_show = cv2.resize(mid_show, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_show", mid_show) + # cv2.imshow("user_res_8uc3_orisize_enhance_fg", user_res_8uc3_orisize_enhance.astype(np.float32)/255. * hair_gene_matte_8uc3_orisize.astype(np.float32) / 255) + # cv2.waitKey() + + user_bald_res_8uc3_orisize_for_fusion_dir = osp.join(userinfo_dir, 'bald_res_ori.png') + # if osp.exists(user_bald_res_8uc3_orisize_for_fusion_dir): + # os.remove(user_bald_res_8uc3_orisize_for_fusion_dir) + cv2.imwrite(user_bald_res_8uc3_orisize_for_fusion_dir, user_bald_res_8uc3_orisize_enhance) + # if osp.exists(res_matting_mask_ori_raw): + # os.remove(res_matting_mask_ori_raw) + cv2.imwrite(res_matting_mask_ori_raw, hair_gene_matte_8uc3_orisize) + + # user_bald_res_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_res_ori.png') + # cv2.imwrite(user_bald_res_8uc3_orisize_dir, user_res_8uc3_orisize) + + + res_fix_img_mask_8uc4_orisize = np.concatenate((user_res_8uc3_orisize_enhance, hair_gene_matte_8uc3_orisize_cp[:, :, :1]), axis=2) + # cv2.imshow('') + cv2.imwrite(res_matting_mask_ori, res_fix_img_mask_8uc4_orisize) + print('costs:', time.time() - t5) + return user_res_8uc3_orisize_enhance, user_res_8uc3_orisize_for_haircolor, 0 diff --git a/hair_service_sd/hair_matting/Generator_Matte.py b/hair_service_sd/hair_matting/Generator_Matte.py new file mode 100644 index 0000000..79c2850 --- /dev/null +++ b/hair_service_sd/hair_matting/Generator_Matte.py @@ -0,0 +1,158 @@ +from hair_matting.matting.gca_matting_hair_single_fg import * +from hair_matting.seg.hairseg_single_model import Evaluator +from time import time +import cv2 +import numpy as np + + +class GenTrimap(object): + def __init__(self): + self.erosion_kernels = [None] + [cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size, size)) for size in range(1,30)] + + def __call__(self, alpha): + + fg_mask = np.zeros_like(alpha) + bg_mask = np.zeros_like(alpha) + fg_mask[alpha == 255] = 1 + bg_mask[alpha == 0] = 1 + + fg_mask = fg_mask.astype(np.int).astype(np.uint8) + bg_mask = bg_mask.astype(np.int).astype(np.uint8) + + fg_mask = cv2.erode(fg_mask, self.erosion_kernels[15]) + bg_mask = cv2.erode(bg_mask, self.erosion_kernels[29]) + + trimap = np.ones_like(alpha, dtype=np.uint8) * 128 + trimap[fg_mask == 1] = 255 + trimap[bg_mask == 1] = 0 + return trimap + +class Generator_Matte(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.debug_flag = False + self.model_dir = os.path.dirname(__file__) + self.output_img_size = 512 + + triseg_model_path = os.path.join(self.model_dir, 'deeplabv3_hair512_360_0520_wl.pth') # deeplabv3_hair512_520_0121 deeplabv3_hair512_360_0520_wl + self.triseg_model = Evaluator(gpu_id=device_id, output_img_size=self.output_img_size, nclass=3, seg_model_path=triseg_model_path) + + hair_matte_model_path = os.path.join(self.model_dir, 'gca-dist-fg-0430-latest_model.pth') # gca-dist-fg-0203-latest_model gca-dist-fg-0430-latest_model + self.matte_model = self.load_hair_matte_model(hair_matte_model_path) + + self.gen_trimap = GenTrimap() + + def remove_prefix_state_dict(self, state_dict, prefix="module"): + """ + remove prefix from the key of pretrained state dict for Data-Parallel + """ + new_state_dict = {} + first_state_name = list(state_dict.keys())[0] + if not first_state_name.startswith(prefix): + for key, value in state_dict.items(): + new_state_dict[key] = state_dict[key].float() + else: + for key, value in state_dict.items(): + new_state_dict[key[len(prefix) + 1:]] = state_dict[key].float() + return new_state_dict + + + def load_hair_matte_model(self, hair_matte_model_path): + # build model + model = networks.get_generator(encoder="resnet_gca_encoder_29", decoder="res_gca_decoder_22", num_class=4) + + # load checkpoint + checkpoint = torch.load(hair_matte_model_path, map_location=lambda storage, loc: storage) + model.load_state_dict(self.remove_prefix_state_dict(checkpoint['state_dict']), strict=True) + model.to(self.device) + # print("matte_model: ", model) + + # inference + model = model.eval() + return model + + def matte_inference(self, image, landmark1k): + + trimap = self.triseg_model.eval(image, landmark1k) + + ori_h, ori_w, _ = image.shape + + limit_size = 1600 + if ori_h > limit_size or ori_w > limit_size: + if ori_h > ori_w: + new_tri_h = limit_size + new_tri_w = int(ori_w * limit_size / ori_h) + else: + new_tri_w = limit_size + new_tri_h = int(ori_h * limit_size / ori_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 + + # cv2.imshow("image before: ", image) + # cv2.imshow("trimap before: ", trimap) + # cv2.waitKey() + trimap = self.gen_trimap(trimap_resize[:, :, 0]) + + # cv2.imwrite(os.path.join(args.output, image_name.replace(ext, "_trimap.png")), trimap) + + image_dict = generator_tensor_dict(image_resize, trimap) + + pred_fg, pred, offset = single_inference(self.matte_model, image_dict, device=self.device) + + pred_fg[trimap==1] = image_resize[trimap==1] + + if pred.shape[1] != image.shape[1] or pred.shape[0] != image.shape[0]: + pred = cv2.resize(pred, (image.shape[1], image.shape[0])) + return pred_fg, pred + +if __name__ == '__main__': + print('Torch Version: ', torch.__version__) + from utils.data_preprocess import * + + generator_matte = Generator_Matte(True, 0) + data_dir = "/media/DATA_4T/hair_data/check_fafeng_trimap/pick_test_fafeng/pick_liuhai" + for image_name in os.listdir(data_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(data_dir, image_name) + + start = time() + + image = cv2.imread(image_path) + + pt_path = image_path.replace(ext, '_landmark1k.txt') + + if not os.path.exists(pt_path): + continue + + landmark1k = np.loadtxt(pt_path) + + matting_start = time() + pred = generator_matte.matte_inference(image, landmark1k) + matting_end = time() + print("matting cost time :{:.4f}".format(matting_end - matting_start)) + + torch.cuda.empty_cache() + + pred = cv2.resize(pred, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_CUBIC) + # cv2.imshow("pred:", pred) + # cv2.waitKey() + + end = time() + print('end img_path: {}, Time cost : {:.4f}'.format(image_path, end-start)) + diff --git a/hair_service_sd/hair_matting/deeplabv3_hair512_360_0520_wl.pth b/hair_service_sd/hair_matting/deeplabv3_hair512_360_0520_wl.pth new file mode 100644 index 0000000..1d93ba2 --- /dev/null +++ b/hair_service_sd/hair_matting/deeplabv3_hair512_360_0520_wl.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9fd1c65a4c002dfb215771a3f2882740e08c49dc974e982a3e9a934e6f4355a5 +size 233017644 diff --git a/hair_service_sd/hair_matting/gca-dist-fg-0430-latest_model.pth b/hair_service_sd/hair_matting/gca-dist-fg-0430-latest_model.pth new file mode 100644 index 0000000..1531178 --- /dev/null +++ b/hair_service_sd/hair_matting/gca-dist-fg-0430-latest_model.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9003635a734d030978e8d35e13de3e226a06f56767c443c868562d30a51acdfb +size 302856998 diff --git a/hair_service_sd/hair_matting/matting/gca_matting_hair_single.py b/hair_service_sd/hair_matting/matting/gca_matting_hair_single.py new file mode 100644 index 0000000..2dbb013 --- /dev/null +++ b/hair_service_sd/hair_matting/matting/gca_matting_hair_single.py @@ -0,0 +1,209 @@ +import os +import cv2 +import argparse +import numpy as np + +import torch +from torch.nn import functional as F + +import utils +from matting import networks +from utils.data_preprocess import * +from utils import landmark_processor +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]) diff --git a/hair_service_sd/hair_matting/matting/gca_matting_hair_single_fg.py b/hair_service_sd/hair_matting/matting/gca_matting_hair_single_fg.py new file mode 100644 index 0000000..a4850d5 --- /dev/null +++ b/hair_service_sd/hair_matting/matting/gca_matting_hair_single_fg.py @@ -0,0 +1,197 @@ +import os +import cv2 +import argparse +import numpy as np + +import torch +from torch.nn import functional as F + +import utils +from hair_matting.matting import networks +from utils.data_preprocess import * +from utils import landmark_processor +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] + + 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) + diff --git a/hair_service_sd/hair_matting/matting/networks/__init__.py b/hair_service_sd/hair_matting/matting/networks/__init__.py new file mode 100644 index 0000000..aee2746 --- /dev/null +++ b/hair_service_sd/hair_matting/matting/networks/__init__.py @@ -0,0 +1 @@ +from .generators import * \ No newline at end of file diff --git a/hair_service_sd/hair_matting/matting/networks/decoders/__init__.py b/hair_service_sd/hair_matting/matting/networks/decoders/__init__.py new file mode 100644 index 0000000..a7eac0f --- /dev/null +++ b/hair_service_sd/hair_matting/matting/networks/decoders/__init__.py @@ -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) \ No newline at end of file diff --git a/hair_service_sd/hair_matting/matting/networks/decoders/res_gca_dec.py b/hair_service_sd/hair_matting/matting/networks/decoders/res_gca_dec.py new file mode 100644 index 0000000..befe48f --- /dev/null +++ b/hair_service_sd/hair_matting/matting/networks/decoders/res_gca_dec.py @@ -0,0 +1,28 @@ +from hair_matting.matting.networks.ops import GuidedCxtAtten, SpectralNorm +from hair_matting.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} + diff --git a/hair_service_sd/hair_matting/matting/networks/decoders/res_shortcut_dec.py b/hair_service_sd/hair_matting/matting/networks/decoders/res_shortcut_dec.py new file mode 100644 index 0000000..47d0a96 --- /dev/null +++ b/hair_service_sd/hair_matting/matting/networks/decoders/res_shortcut_dec.py @@ -0,0 +1,24 @@ +from hair_matting.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 + diff --git a/hair_service_sd/hair_matting/matting/networks/decoders/resnet_dec.py b/hair_service_sd/hair_matting/matting/networks/decoders/resnet_dec.py new file mode 100644 index 0000000..3c32ede --- /dev/null +++ b/hair_service_sd/hair_matting/matting/networks/decoders/resnet_dec.py @@ -0,0 +1,142 @@ +import logging +import torch.nn as nn +from hair_matting.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 diff --git a/hair_service_sd/hair_matting/matting/networks/encoders/__init__.py b/hair_service_sd/hair_matting/matting/networks/encoders/__init__.py new file mode 100644 index 0000000..71693a6 --- /dev/null +++ b/hair_service_sd/hair_matting/matting/networks/encoders/__init__.py @@ -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) diff --git a/hair_service_sd/hair_matting/matting/networks/encoders/res_gca_enc.py b/hair_service_sd/hair_matting/matting/networks/encoders/res_gca_enc.py new file mode 100644 index 0000000..c1cf25f --- /dev/null +++ b/hair_service_sd/hair_matting/matting/networks/encoders/res_gca_enc.py @@ -0,0 +1,97 @@ +import torch.nn as nn +import torch.nn.functional as F + +# from utils import CONFIG +from hair_matting.matting.networks.encoders.resnet_enc import ResNet_D +from hair_matting.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 matting.networks.encoders.resnet_enc import BasicBlock + m = ResGuidedCxtAtten(BasicBlock, [3, 4, 4, 2]) + for m in m.modules(): + print(m) diff --git a/hair_service_sd/hair_matting/matting/networks/encoders/res_shortcut_enc.py b/hair_service_sd/hair_matting/matting/networks/encoders/res_shortcut_enc.py new file mode 100644 index 0000000..0ac7bdd --- /dev/null +++ b/hair_service_sd/hair_matting/matting/networks/encoders/res_shortcut_enc.py @@ -0,0 +1,51 @@ +import torch.nn as nn +# from utils import CONFIG +from hair_matting.matting.networks.encoders.resnet_enc import ResNet_D +from hair_matting.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,...]} \ No newline at end of file diff --git a/hair_service_sd/hair_matting/matting/networks/encoders/resnet_enc.py b/hair_service_sd/hair_matting/matting/networks/encoders/resnet_enc.py new file mode 100644 index 0000000..2daf4ed --- /dev/null +++ b/hair_service_sd/hair_matting/matting/networks/encoders/resnet_enc.py @@ -0,0 +1,150 @@ +import logging +import torch.nn as nn +from hair_matting.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()) diff --git a/hair_service_sd/hair_matting/matting/networks/generators.py b/hair_service_sd/hair_matting/matting/networks/generators.py new file mode 100644 index 0000000..e1cf59b --- /dev/null +++ b/hair_service_sd/hair_matting/matting/networks/generators.py @@ -0,0 +1,60 @@ +import torch +import torch.nn as nn + +# from utils import CONFIG +from hair_matting.matting.networks import encoders, decoders + + +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) \ No newline at end of file diff --git a/hair_service_sd/hair_matting/matting/networks/ops.py b/hair_service_sd/hair_matting/matting/networks/ops.py new file mode 100644 index 0000000..a2df034 --- /dev/null +++ b/hair_service_sd/hair_matting/matting/networks/ops.py @@ -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) \ No newline at end of file diff --git a/hair_service_sd/hair_matting/matting/setup.py b/hair_service_sd/hair_matting/matting/setup.py new file mode 100644 index 0000000..ca6bf5a --- /dev/null +++ b/hair_service_sd/hair_matting/matting/setup.py @@ -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 \ No newline at end of file diff --git a/hair_service_sd/hair_matting/seg/hairseg_single_model.py b/hair_service_sd/hair_matting/seg/hairseg_single_model.py new file mode 100644 index 0000000..66fb717 --- /dev/null +++ b/hair_service_sd/hair_matting/seg/hairseg_single_model.py @@ -0,0 +1,95 @@ +import os +import torch + +from hair_matting.seg.networks.deeplabv3_plus import get_deeplabv3_plus +import numpy as np +import cv2 +import time +from utils import landmark_processor + +def label_to_mask(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 + +label_map = [ + [0, 0, 0], # + [128, 128, 128], + [255, 255, 255], +] + +class Evaluator(object): + def __init__(self, gpu_id, output_img_size, nclass, seg_model_path=None): + self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu') + # print("gpu_id: ", gpu_id) + + # create network + self.model = get_deeplabv3_plus(backbone='xception', nclass=nclass) + model_path = os.path.join(seg_model_path) + self.model.load_state_dict(torch.load(model_path, map_location=lambda storage, loc: storage)) + # print("seg device: ", self.model.device) + self.model.to(self.device) + self.model.eval() + + # images = torch.randn((1, 3, 512, 512)).to(self.device) + # torch.onnx.export(self.model, images, + # "deeplabv3_hair512_360_0520_wl.onnx", + # verbose=True, + # opset_version=11, + # input_names=['data'], + # do_constant_folding=True, + # output_names=['output']) + + # exit() + + self.output_img_size = output_img_size + self.nclass = nclass + + + def process_data(self, img): + img = (img.astype(np.float32) / 255).transpose((2, 0, 1)) + img = torch.from_numpy(img).unsqueeze(0) + + return img + + def eval(self, img, pts1k): + + orig_h, orig_w, _ = img.shape + + pre_start = time.time() + + M1 = landmark_processor.get_transform_mat_hair(pts1k, self.output_img_size, ratio=0.28, h_ratio=0.30) + img = cv2.warpAffine(img, M1, (self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4) + img = self.process_data(img) + img = img.to(self.device) + pre_end = time.time() + # print("seg pre cost time : {:.6f} s".format(pre_end - pre_start)) + + with torch.no_grad(): + # torch.cuda.synchronize() + start = time.time() + outputs = self.model(img) + end = time.time() + # print("seg forward cost time : {:.6f} s".format(end - start)) + pred = torch.argmax(outputs[0], 1) + # print("argmax cost time : {:.6f} s".format(time.time() - end)) + + pred = pred[0].detach().cpu().numpy() + predict = pred.astype(np.float32) + pred_detach = time.time() + # print("pred_detach cost time : {:.6f} s".format(pred_detach - end)) + + pred_mask = label_to_mask(predict) + + M1_invert = cv2.invertAffineTransform(M1) + img_pred = cv2.warpAffine(pred_mask, M1_invert, (orig_w, orig_h), flags=cv2.INTER_CUBIC) #flags=cv2.INTER_NEAREST + orig_mask = img_pred.copy() + # print("seg post cost time : {:.6f} s".format(time.time() - end)) + + return orig_mask + + + diff --git a/hair_service_sd/hair_matting/seg/networks/basic.py b/hair_service_sd/hair_matting/seg/networks/basic.py new file mode 100644 index 0000000..b241fad --- /dev/null +++ b/hair_service_sd/hair_matting/seg/networks/basic.py @@ -0,0 +1,462 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +__all__ = ['_ConvBNReLU', '_DWConvBNReLU', 'InvertedResidual', '_ASPP', '_FCNHead', + '_Hswish', '_ConvBNHswish', 'SEModule', 'Bottleneck', 'ShuffleNetUnit', + 'ShuffleNetV2Unit', 'InvertedIGCV3', 'MBConvBlock'] + + +class _ConvBNReLU(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, + dilation=1, groups=1, relu6=False, norm_layer=nn.BatchNorm2d, **kwargs): + super(_ConvBNReLU, self).__init__() + self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias=False) + self.bn = norm_layer(out_channels) + self.relu = nn.ReLU6(True) if relu6 else nn.ReLU(True) + + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + x = self.relu(x) + return x + + +class _FCNHead(nn.Module): + def __init__(self, in_channels, channels, norm_layer=nn.BatchNorm2d, norm_kwargs=None, **kwargs): + super(_FCNHead, self).__init__() + inter_channels = in_channels // 4 + self.block = nn.Sequential( + nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False), + norm_layer(inter_channels, **({} if norm_kwargs is None else norm_kwargs)), + nn.ReLU(True), + nn.Dropout(0.1), + nn.Conv2d(inter_channels, channels, 1) + ) + + def forward(self, x): + return self.block(x) + + +# ----------------------------------------------------------------- +# For MobileNet +# ----------------------------------------------------------------- +class _DWConvBNReLU(nn.Module): + """Depthwise Separable Convolution in MobileNet. + depthwise convolution + pointwise convolution + """ + + def __init__(self, in_channels, dw_channels, out_channels, stride, dilation=1, norm_layer=nn.BatchNorm2d, **kwargs): + super(_DWConvBNReLU, self).__init__() + self.conv = nn.Sequential( + _ConvBNReLU(in_channels, dw_channels, 3, stride, dilation, dilation, in_channels, norm_layer=norm_layer), + _ConvBNReLU(dw_channels, out_channels, 1, norm_layer=norm_layer)) + + def forward(self, x): + return self.conv(x) + + +# ----------------------------------------------------------------- +# For MobileNetV2 +# ----------------------------------------------------------------- +class InvertedResidual(nn.Module): + def __init__(self, in_channels, out_channels, stride, expand_ratio, + dilation=1, norm_layer=nn.BatchNorm2d, **kwargs): + super(InvertedResidual, self).__init__() + assert stride in [1, 2] + self.use_res_connect = stride == 1 and in_channels == out_channels + + layers = list() + inter_channels = int(round(in_channels * expand_ratio)) + if expand_ratio != 1: + # pw + layers.append(_ConvBNReLU(in_channels, inter_channels, 1, relu6=True, norm_layer=norm_layer)) + layers.extend([ + # dw + _ConvBNReLU(inter_channels, inter_channels, 3, stride, dilation, dilation, + groups=inter_channels, relu6=True, norm_layer=norm_layer), + # pw-linear + nn.Conv2d(inter_channels, out_channels, 1, bias=False), + norm_layer(out_channels)]) + self.conv = nn.Sequential(*layers) + + def forward(self, x): + if self.use_res_connect: + return x + self.conv(x) + else: + return self.conv(x) + + +# ----------------------------------------------------------------- +# ASPP: For MobileNetV2 +# ----------------------------------------------------------------- +class _AsppPooling(nn.Module): + def __init__(self, in_channels, out_channels, norm_layer, **kwargs): + super(_AsppPooling, self).__init__() + self.gap = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(in_channels, out_channels, 1, bias=False), + norm_layer(out_channels), + nn.ReLU(True) + ) + + def forward(self, x): +# size = x.size()[2:] + size = (48, 48) +# print("size: ", size) + pool = self.gap(x) +# out = F.interpolate(pool, size, mode='bilinear', align_corners=True) + out = F.interpolate(pool, size, mode='nearest') + return out + + +class _ASPP(nn.Module): + def __init__(self, in_channels, atrous_rates, norm_layer=nn.BatchNorm2d, **kwargs): + super(_ASPP, self).__init__() + out_channels = 256 + self.b0 = nn.Sequential( + nn.Conv2d(in_channels, out_channels, 1, bias=False), + norm_layer(out_channels), + nn.ReLU(True) + ) + + rate1, rate2, rate3 = tuple(atrous_rates) + self.b1 = _ConvBNReLU(in_channels, out_channels, 3, padding=rate1, dilation=rate1, norm_layer=norm_layer) + self.b2 = _ConvBNReLU(in_channels, out_channels, 3, padding=rate2, dilation=rate2, norm_layer=norm_layer) + self.b3 = _ConvBNReLU(in_channels, out_channels, 3, padding=rate3, dilation=rate3, norm_layer=norm_layer) + self.b4 = _AsppPooling(in_channels, out_channels, norm_layer=norm_layer) + + self.project = nn.Sequential( + nn.Conv2d(5 * out_channels, out_channels, 1, bias=False), + norm_layer(out_channels), + nn.ReLU(True), + nn.Dropout2d(0.5) + ) + + def forward(self, x): + feat1 = self.b0(x) + feat2 = self.b1(x) + feat3 = self.b2(x) + feat4 = self.b3(x) + feat5 = self.b4(x) + x = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1) + x = self.project(x) + return x + + +# ----------------------------------------------------------------- +# For MobileNetV3 +# ----------------------------------------------------------------- +class _Hswish(nn.Module): + def __init__(self, inplace=True): + super(_Hswish, self).__init__() + self.relu6 = nn.ReLU6(inplace) + + def forward(self, x): + return x * self.relu6(x + 3.) / 6. + + +class _Hsigmoid(nn.Module): + def __init__(self, inplace=True): + super(_Hsigmoid, self).__init__() + self.relu6 = nn.ReLU6(inplace) + + def forward(self, x): + return self.relu6(x + 3.) / 6. + + +class _ConvBNHswish(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, + dilation=1, groups=1, norm_layer=nn.BatchNorm2d, **kwargs): + super(_ConvBNHswish, self).__init__() + self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias=False) + self.bn = norm_layer(out_channels) + self.act = _Hswish(True) + + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + x = self.act(x) + return x + + +class SEModule(nn.Module): + def __init__(self, in_channels, reduction=4): + super(SEModule, self).__init__() + self.avg_pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(in_channels, in_channels // reduction, bias=False), + nn.ReLU(True), + nn.Linear(in_channels // reduction, in_channels, bias=False), + _Hsigmoid(True) + ) + + def forward(self, x): + n, c, _, _ = x.size() + out = self.avg_pool(x).view(n, c) + out = self.fc(out).view(n, c, 1, 1) + return x * out.expand_as(x) + + +class Identity(nn.Module): + def __init__(self, in_channels): + super(Identity, self).__init__() + + def forward(self, x): + return x + + +class Bottleneck(nn.Module): + def __init__(self, in_channels, out_channels, exp_size, kernel_size, stride, dilation=1, se=False, nl='RE', + norm_layer=nn.BatchNorm2d, **kwargs): + super(Bottleneck, self).__init__() + assert stride in [1, 2] + self.use_res_connect = stride == 1 and in_channels == out_channels + if nl == 'HS': + act = _Hswish + else: + act = nn.ReLU + if se: + SELayer = SEModule + else: + SELayer = Identity + + self.conv = nn.Sequential( + # pw + nn.Conv2d(in_channels, exp_size, 1, bias=False), + norm_layer(exp_size), + act(True), + # dw + nn.Conv2d(exp_size, exp_size, kernel_size, stride, (kernel_size - 1) // 2 * dilation, + dilation, groups=exp_size, bias=False), + norm_layer(exp_size), + SELayer(exp_size), + act(True), + # pw-linear + nn.Conv2d(exp_size, out_channels, 1, bias=False), + norm_layer(out_channels) + ) + + def forward(self, x): + if self.use_res_connect: + return x + self.conv(x) + else: + return self.conv(x) + + +# ----------------------------------------------------------------- +# For ShuffleNet +# ----------------------------------------------------------------- +def channel_shuffle(x, groups): + n, c, h, w = x.size() + + channels_per_group = c // groups + x = x.view(n, groups, channels_per_group, h, w) + x = torch.transpose(x, 1, 2).contiguous() + x = x.view(n, -1, h, w) + + return x + + +class ShuffleNetUnit(nn.Module): + def __init__(self, in_channels, out_channels, stride, groups, dilation=1, norm_layer=nn.BatchNorm2d, **kwargs): + super(ShuffleNetUnit, self).__init__() + self.stride = stride + self.groups = groups + self.dilation = dilation + assert stride in [1, 2, 3] + + inter_channels = out_channels // 4 + + if stride > 1: + self.shortcut = nn.AvgPool2d(3, stride, 1) + out_channels -= in_channels + elif dilation > 1: + out_channels -= in_channels + + g = 1 if in_channels == 24 else groups + self.conv1 = _ConvBNReLU(in_channels, inter_channels, 1, groups=g, norm_layer=norm_layer) + self.conv2 = _ConvBNReLU(inter_channels, inter_channels, 3, stride, dilation, + dilation, groups, norm_layer=norm_layer) + self.conv3 = nn.Sequential( + nn.Conv2d(inter_channels, out_channels, 1, groups=groups, bias=False), + norm_layer(out_channels)) + + def forward(self, x): + out = self.conv1(x) + out = channel_shuffle(out, self.groups) + out = self.conv2(out) + out = self.conv3(out) + if self.stride > 1: + x = self.shortcut(x) + out = torch.cat([out, x], dim=1) + elif self.dilation > 1: + out = torch.cat([out, x], dim=1) + else: + out = out + x + out = F.relu(out) + + return out + + +# ----------------------------------------------------------------- +# For ShuffleNetV2 +# ----------------------------------------------------------------- +class _DWConv(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, bias=False): + super(_DWConv, self).__init__() + self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, + padding, dilation, groups=in_channels, bias=bias) + + def forward(self, x): + return self.conv(x) + + +class ShuffleNetV2Unit(nn.Module): + def __init__(self, in_channels, out_channels, stride, dilation=1, norm_layer=nn.BatchNorm2d, **kwargs): + super(ShuffleNetV2Unit, self).__init__() + assert stride in [1, 2, 3] + self.stride = stride + self.dilation = dilation + + inter_channels = out_channels // 2 + + if (stride > 1) or (dilation > 1): + self.branch1 = nn.Sequential( + _DWConv(in_channels, in_channels, 3, stride, dilation, dilation), + norm_layer(in_channels), + _ConvBNReLU(in_channels, inter_channels, 1, norm_layer=norm_layer)) + self.branch2 = nn.Sequential( + _ConvBNReLU(in_channels if (stride > 1) else inter_channels, inter_channels, 1, norm_layer=norm_layer), + _DWConv(inter_channels, inter_channels, 3, stride, dilation, dilation), + norm_layer(inter_channels), + _ConvBNReLU(inter_channels, inter_channels, 1, norm_layer=norm_layer)) + + def forward(self, x): + if (self.stride == 1) and (self.dilation == 1): + x1, x2 = x.chunk(2, dim=1) + out = torch.cat((x1, self.branch2(x2)), dim=1) + else: + out = torch.cat((self.branch1(x), self.branch2(x)), dim=1) + out = channel_shuffle(out, 2) + + return out + + +# ----------------------------------------------------------------- +# For IGCV3 +# ----------------------------------------------------------------- +class PermutationBlock(nn.Module): + def __init__(self, groups): + super(PermutationBlock, self).__init__() + self.groups = groups + + def forward(self, x): + n, c, h, w = x.size() + x = x.view(n, self.groups, c // self.groups, h, w).permute(0, 2, 1, 3, 4).contiguous().view(n, c, h, w) + return x + + +class InvertedIGCV3(nn.Module): + def __init__(self, in_channels, out_channels, stride, expand_ratio, + dilation=1, norm_layer=nn.BatchNorm2d, **kwargs): + super(InvertedIGCV3, self).__init__() + assert stride in [1, 2] + self.use_res_connect = stride == 1 and in_channels == out_channels + + layers = list() + inter_channels = int(round(in_channels * expand_ratio)) + if expand_ratio != 1: + # pw + layers.append(_ConvBNReLU(in_channels, inter_channels, 1, + groups=2, relu6=True, norm_layer=norm_layer)) + # permutation + layers.append(PermutationBlock(groups=2)) + layers.extend([ + # dw + _ConvBNReLU(inter_channels, inter_channels, 3, stride, dilation, dilation, + groups=inter_channels, relu6=True, norm_layer=norm_layer), + # pw-linear + nn.Conv2d(inter_channels, out_channels, 1, groups=2, bias=False), + norm_layer(out_channels), + # permutation + PermutationBlock(groups=int(round(out_channels / 2))) + ]) + self.conv = nn.Sequential(*layers) + + def forward(self, x): + if self.use_res_connect: + return x + self.conv(x) + else: + return self.conv(x) + + +# ----------------------------------------------------------------- +# For EfficientNet +# ----------------------------------------------------------------- +class _Swish(nn.Module): + def __init__(self): + super(_Swish, self).__init__() + self.sigmoid = nn.Sigmoid() + + def forward(self, x): + return x * self.sigmoid(x) + + +class SEModuleV2(nn.Module): + def __init__(self, in_channels, se_ratio=0.25): + super(SEModuleV2, self).__init__() + self.avg_pool = nn.AdaptiveAvgPool2d(1) + se_channels = max(1, int(in_channels * se_ratio)) + self.fc = nn.Sequential( + nn.Conv2d(in_channels, se_channels, 1, bias=False), + _Swish(), + nn.Conv2d(se_channels, in_channels, 1, bias=False), + nn.Sigmoid() + ) + + def forward(self, x): + n, c, _, _ = x.size() + out = self.avg_pool(x) + out = self.fc(out) + return x * out.expand_as(x) + + +class MBConvBlock(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride, expand_ratio, + dilation=1, se_ratio=0.25, drop_connect_rate=0.2, norm_layer=nn.BatchNorm2d, **kwargs): + super(MBConvBlock, self).__init__() + assert stride in [1, 2] + self.use_res_connect = stride == 1 and in_channels == out_channels + self.drop_connect_rate = drop_connect_rate + use_se = (se_ratio is not None) and (0 < se_ratio <= 1.) + if use_se: + SELayer = SEModuleV2 + else: + SELayer = Identity + + layers = list() + inter_channels = int(round(in_channels * expand_ratio)) + if expand_ratio != 1: + layers.append(_ConvBNHswish(in_channels, inter_channels, 1, norm_layer=norm_layer)) + layers.extend([ + # dw + _ConvBNHswish(inter_channels, inter_channels, kernel_size, stride, kernel_size // 2 * dilation, dilation, + groups=inter_channels, norm_layer=norm_layer), # check act function + SELayer(inter_channels, se_ratio), + # pw-linear + nn.Conv2d(inter_channels, out_channels, 1, bias=False), + norm_layer(out_channels) + ]) + self.conv = nn.Sequential(*layers) + + if drop_connect_rate: + self.dropout = nn.Dropout2d(drop_connect_rate) + + def forward(self, x): + out = self.conv(x) + if self.use_res_connect: + if self.drop_connect_rate: + out = self.dropout(out) + out = x + out + return out diff --git a/hair_service_sd/hair_matting/seg/networks/deeplabv3.py b/hair_service_sd/hair_matting/seg/networks/deeplabv3.py new file mode 100644 index 0000000..a1c6c50 --- /dev/null +++ b/hair_service_sd/hair_matting/seg/networks/deeplabv3.py @@ -0,0 +1,187 @@ +"""Pyramid Scene Parsing Network""" +import torch +import torch.nn as nn +import torch.nn.functional as F + +from hair_matting.seg.networks.segbase import SegBaseModel +from hair_matting.seg.networks.fcn import _FCNHead + +__all__ = ['DeepLabV3', 'get_deeplabv3', 'get_deeplabv3_resnet50_voc', 'get_deeplabv3_resnet101_voc', + 'get_deeplabv3_resnet152_voc', 'get_deeplabv3_resnet50_ade', 'get_deeplabv3_resnet101_ade', + 'get_deeplabv3_resnet152_ade'] + + +class DeepLabV3(SegBaseModel): + r"""DeepLabV3 + + Parameters + ---------- + nclass : int + Number of categories for the training dataset. + backbone : string + Pre-trained dilated backbone network type (default:'resnet50'; 'resnet50', + 'resnet101' or 'resnet152'). + norm_layer : object + Normalization layer used in backbone network (default: :class:`nn.BatchNorm`; + for Synchronized Cross-GPU BachNormalization). + aux : bool + Auxiliary loss. + + Reference: + Chen, Liang-Chieh, et al. "Rethinking atrous convolution for semantic image segmentation." + arXiv preprint arXiv:1706.05587 (2017). + """ + + def __init__(self, nclass, backbone='resnet50', aux=False, pretrained_base=True, **kwargs): + super(DeepLabV3, self).__init__(nclass, aux, backbone, pretrained_base=pretrained_base, **kwargs) + self.head = _DeepLabHead(nclass, **kwargs) + if self.aux: + self.auxlayer = _FCNHead(1024, nclass, **kwargs) + + self.__setattr__('exclusive', ['head', 'auxlayer'] if aux else ['head']) + + def forward(self, x): + size = x.size()[2:] + _, _, c3, c4 = self.base_forward(x) + outputs = [] + x = self.head(c4) + x = F.interpolate(x, size, mode='bilinear', align_corners=True) + outputs.append(x) + + if self.aux: + auxout = self.auxlayer(c3) + auxout = F.interpolate(auxout, size, mode='bilinear', align_corners=True) + outputs.append(auxout) + return tuple(outputs) + + +class _DeepLabHead(nn.Module): + def __init__(self, nclass, norm_layer=nn.BatchNorm2d, norm_kwargs=None, **kwargs): + super(_DeepLabHead, self).__init__() + self.aspp = _ASPP(2048, [12, 24, 36], norm_layer=norm_layer, norm_kwargs=norm_kwargs, **kwargs) + self.block = nn.Sequential( + nn.Conv2d(256, 256, 3, padding=1, bias=False), + norm_layer(256, **({} if norm_kwargs is None else norm_kwargs)), + nn.ReLU(True), + nn.Dropout(0.1), + nn.Conv2d(256, nclass, 1) + ) + + def forward(self, x): + x = self.aspp(x) + return self.block(x) + + +class _ASPPConv(nn.Module): + def __init__(self, in_channels, out_channels, atrous_rate, norm_layer, norm_kwargs): + super(_ASPPConv, self).__init__() + self.block = nn.Sequential( + nn.Conv2d(in_channels, out_channels, 3, padding=atrous_rate, dilation=atrous_rate, bias=False), + norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)), + nn.ReLU(True) + ) + + def forward(self, x): + return self.block(x) + + +class _AsppPooling(nn.Module): + def __init__(self, in_channels, out_channels, norm_layer, norm_kwargs, **kwargs): + super(_AsppPooling, self).__init__() + self.gap = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(in_channels, out_channels, 1, bias=False), + norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)), + nn.ReLU(True) + ) + + def forward(self, x): + size = x.size()[2:] +# print("before gap: ", x.size()) + pool = self.gap(x) + out = F.interpolate(pool, size, mode='bilinear', align_corners=True) + return out + + +class _ASPP(nn.Module): + def __init__(self, in_channels, atrous_rates, norm_layer, norm_kwargs=None, **kwargs): + super(_ASPP, self).__init__() + out_channels = 256 + self.b0 = nn.Sequential( + nn.Conv2d(in_channels, out_channels, 1, bias=False), + norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)), + nn.ReLU(True) + ) + + rate1, rate2, rate3 = tuple(atrous_rates) + self.b1 = _ASPPConv(in_channels, out_channels, rate1, norm_layer, norm_kwargs) + self.b2 = _ASPPConv(in_channels, out_channels, rate2, norm_layer, norm_kwargs) + self.b3 = _ASPPConv(in_channels, out_channels, rate3, norm_layer, norm_kwargs) + self.b4 = _AsppPooling(in_channels, out_channels, norm_layer=norm_layer, norm_kwargs=norm_kwargs) + + self.project = nn.Sequential( + nn.Conv2d(5 * out_channels, out_channels, 1, bias=False), + norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)), + nn.ReLU(True), + nn.Dropout(0.5) + ) + + def forward(self, x): + feat1 = self.b0(x) + feat2 = self.b1(x) + feat3 = self.b2(x) + feat4 = self.b3(x) +# print("before b4: ", x.size()) + feat5 = self.b4(x) + x = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1) + x = self.project(x) + return x + + +def get_deeplabv3(dataset='pascal_voc', backbone='resnet50', pretrained=False, root='~/.torch/models', + pretrained_base=True, **kwargs): + acronyms = { + 'pascal_voc': 'pascal_voc', + 'pascal_aug': 'pascal_aug', + 'ade20k': 'ade', + 'coco': 'coco', + 'citys': 'citys', + } + from ..data.dataloader import datasets + model = DeepLabV3(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs) + if pretrained: + from .model_store import get_model_file + device = torch.device(kwargs['local_rank']) + model.load_state_dict(torch.load(get_model_file('deeplabv3_%s_%s' % (backbone, acronyms[dataset]), root=root), + map_location=device)) + return model + + +def get_deeplabv3_resnet50_voc(**kwargs): + return get_deeplabv3('pascal_voc', 'resnet50', **kwargs) + + +def get_deeplabv3_resnet101_voc(**kwargs): + return get_deeplabv3('pascal_voc', 'resnet101', **kwargs) + + +def get_deeplabv3_resnet152_voc(**kwargs): + return get_deeplabv3('pascal_voc', 'resnet152', **kwargs) + + +def get_deeplabv3_resnet50_ade(**kwargs): + return get_deeplabv3('ade20k', 'resnet50', **kwargs) + + +def get_deeplabv3_resnet101_ade(**kwargs): + return get_deeplabv3('ade20k', 'resnet101', **kwargs) + + +def get_deeplabv3_resnet152_ade(**kwargs): + return get_deeplabv3('ade20k', 'resnet152', **kwargs) + + +if __name__ == '__main__': + model = get_deeplabv3_resnet50_voc() + img = torch.randn(2, 3, 480, 480) + output = model(img) diff --git a/hair_service_sd/hair_matting/seg/networks/deeplabv3_plus.py b/hair_service_sd/hair_matting/seg/networks/deeplabv3_plus.py new file mode 100644 index 0000000..c843029 --- /dev/null +++ b/hair_service_sd/hair_matting/seg/networks/deeplabv3_plus.py @@ -0,0 +1,160 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +from hair_matting.seg.networks.xception import get_xception +from hair_matting.seg.networks.deeplabv3 import _ASPP +from hair_matting.seg.networks.fcn import _FCNHead +from hair_matting.seg.networks.basic import _ConvBNReLU + +__all__ = ['DeepLabV3Plus', 'get_deeplabv3_plus', 'get_deeplabv3_plus_xception_voc'] + + +class DeepLabV3Plus(nn.Module): + r"""DeepLabV3Plus + Parameters + ---------- + nclass : int + Number of categories for the training dataset. + backbone : string + Pre-trained dilated backbone network type (default:'xception'). + norm_layer : object + Normalization layer used in backbone network (default: :class:`nn.BatchNorm`; + for Synchronized Cross-GPU BachNormalization). + aux : bool + Auxiliary loss. + + Reference: + Chen, Liang-Chieh, et al. "Encoder-Decoder with Atrous Separable Convolution for Semantic + Image Segmentation." + """ + + def __init__(self, nclass, backbone='xception', aux=True, pretrained_base=True, dilated=True, **kwargs): + super(DeepLabV3Plus, self).__init__() + self.aux = aux + self.nclass = nclass + output_stride = 8 if dilated else 32 + + self.pretrained = get_xception(pretrained=pretrained_base, output_stride=output_stride, **kwargs) + + # deeplabv3 plus + self.head = _DeepLabHead(nclass, **kwargs) + if aux: + self.auxlayer = _FCNHead(728, nclass, **kwargs) + + def base_forward(self, x): + # Entry flow + x = self.pretrained.conv1(x) + x = self.pretrained.bn1(x) + x = self.pretrained.relu(x) + + x = self.pretrained.conv2(x) + x = self.pretrained.bn2(x) + x = self.pretrained.relu(x) + + x = self.pretrained.block1(x) + # add relu here + x = self.pretrained.relu(x) + low_level_feat = x + + x = self.pretrained.block2(x) + x = self.pretrained.block3(x) + + # Middle flow + x = self.pretrained.midflow(x) + mid_level_feat = x + + # Exit flow + x = self.pretrained.block20(x) + x = self.pretrained.relu(x) + x = self.pretrained.conv3(x) + x = self.pretrained.bn3(x) + x = self.pretrained.relu(x) + + x = self.pretrained.conv4(x) + x = self.pretrained.bn4(x) + x = self.pretrained.relu(x) + + x = self.pretrained.conv5(x) + x = self.pretrained.bn5(x) + x = self.pretrained.relu(x) + return low_level_feat, mid_level_feat, x + + def forward(self, x): +# print("x size: ", x.size()) + size = x.size()[2:] + c1, c3, c4 = self.base_forward(x) +# print("c1 size: ", c1.size()) +# print("c4 size: ", c4.size()) + outputs = list() + x = self.head(c4, c1) + x = F.interpolate(x, size, mode='bilinear', align_corners=True) + outputs.append(x) + if self.aux: + auxout = self.auxlayer(c3) + auxout = F.interpolate(auxout, size, mode='bilinear', align_corners=True) + outputs.append(auxout) + + # for save onnx + # y = torch.max(x, 1)[1].to(torch.float32) + # return y + + return tuple(outputs) + + +class _DeepLabHead(nn.Module): + def __init__(self, nclass, c1_channels=128, norm_layer=nn.BatchNorm2d, **kwargs): + super(_DeepLabHead, self).__init__() + self.aspp = _ASPP(2048, [12, 24, 36], norm_layer=norm_layer, **kwargs) + self.c1_block = _ConvBNReLU(c1_channels, 48, 3, padding=1, norm_layer=norm_layer) + self.block = nn.Sequential( + _ConvBNReLU(304, 256, 3, padding=1, norm_layer=norm_layer), + nn.Dropout(0.5), + _ConvBNReLU(256, 256, 3, padding=1, norm_layer=norm_layer), + nn.Dropout(0.1), + nn.Conv2d(256, nclass, 1)) + + def forward(self, x, c1): + size = c1.size()[2:] + c1 = self.c1_block(c1) +# print("c1", c1.size()) +# print("before aspp: ", x.size()) + x = self.aspp(x) +# print("after aspp: ", x.size()) + x = F.interpolate(x, size, mode='bilinear', align_corners=True) + return self.block(torch.cat([x, c1], dim=1)) + + +def get_deeplabv3_plus(dataset='pascal_voc', backbone='xception', pretrained=False, root='../ckpt', + pretrained_base=False, nclass=3, **kwargs): + acronyms = { + 'pascal_voc': 'pascal_voc', + 'pascal_aug': 'pascal_aug', + 'ade20k': 'ade', + 'coco': 'coco', + 'citys': 'citys', + } + #from light.data import datasets + + model = DeepLabV3Plus(nclass, backbone=backbone, pretrained_base=pretrained_base, **kwargs) + if pretrained: + pass + # if dataset not in acronyms.keys(): + # print("root:", root) + # model_path = os.path.join(root, "deeplabv3_plus_28.pth") + # model.load_state_dict(torch.load(model_path), strict=False) + # else: + # from .model_store import get_model_file + # device = torch.device(kwargs['local_rank']) + # model.load_state_dict( + # torch.load(get_model_file('deeplabv3_plus_%s_%s' % (backbone, acronyms[dataset]), root=root), + # map_location=device)) + return model + + +def get_deeplabv3_plus_xception_voc(**kwargs): + return get_deeplabv3_plus('pascal_voc', 'xception', **kwargs) + + +if __name__ == '__main__': + model = get_deeplabv3_plus_xception_voc() diff --git a/hair_service_sd/hair_matting/seg/networks/fcn.py b/hair_service_sd/hair_matting/seg/networks/fcn.py new file mode 100644 index 0000000..a1c75da --- /dev/null +++ b/hair_service_sd/hair_matting/seg/networks/fcn.py @@ -0,0 +1,222 @@ +import os +import torch +import torch.nn as nn +import torch.nn.functional as F + +from hair_matting.seg.networks.vgg import vgg16 + +__all__ = ['get_fcn32s', 'get_fcn16s', 'get_fcn8s', + 'get_fcn32s_vgg16_voc', 'get_fcn16s_vgg16_voc', 'get_fcn8s_vgg16_voc'] + + +class FCN32s(nn.Module): + """There are some difference from original fcn""" + + def __init__(self, nclass, backbone='vgg16', aux=False, pretrained_base=True, + norm_layer=nn.BatchNorm2d, **kwargs): + super(FCN32s, self).__init__() + self.aux = aux + if backbone == 'vgg16': + self.pretrained = vgg16(pretrained=pretrained_base).features + else: + raise RuntimeError('unknown backbone: {}'.format(backbone)) + self.head = _FCNHead(512, nclass, norm_layer) + if aux: + self.auxlayer = _FCNHead(512, nclass, norm_layer) + + self.__setattr__('exclusive', ['head', 'auxlayer'] if aux else ['head']) + + def forward(self, x): + size = x.size()[2:] + pool5 = self.pretrained(x) + + outputs = [] + out = self.head(pool5) + out = F.interpolate(out, size, mode='bilinear', align_corners=True) + outputs.append(out) + + if self.aux: + auxout = self.auxlayer(pool5) + auxout = F.interpolate(auxout, size, mode='bilinear', align_corners=True) + outputs.append(auxout) + + return tuple(outputs) + + +class FCN16s(nn.Module): + def __init__(self, nclass, backbone='vgg16', aux=False, pretrained_base=True, norm_layer=nn.BatchNorm2d, **kwargs): + super(FCN16s, self).__init__() + self.aux = aux + if backbone == 'vgg16': + self.pretrained = vgg16(pretrained=pretrained_base).features + else: + raise RuntimeError('unknown backbone: {}'.format(backbone)) + self.pool4 = nn.Sequential(*self.pretrained[:24]) + self.pool5 = nn.Sequential(*self.pretrained[24:]) + self.head = _FCNHead(512, nclass, norm_layer) + self.score_pool4 = nn.Conv2d(512, nclass, 1) + if aux: + self.auxlayer = _FCNHead(512, nclass, norm_layer) + + self.__setattr__('exclusive', ['head', 'score_pool4', 'auxlayer'] if aux else ['head', 'score_pool4']) + + def forward(self, x): + pool4 = self.pool4(x) + pool5 = self.pool5(pool4) + + outputs = [] + score_fr = self.head(pool5) + + score_pool4 = self.score_pool4(pool4) + + upscore2 = F.interpolate(score_fr, score_pool4.size()[2:], mode='bilinear', align_corners=True) + fuse_pool4 = upscore2 + score_pool4 + + out = F.interpolate(fuse_pool4, x.size()[2:], mode='bilinear', align_corners=True) + outputs.append(out) + + if self.aux: + auxout = self.auxlayer(pool5) + auxout = F.interpolate(auxout, x.size()[2:], mode='bilinear', align_corners=True) + outputs.append(auxout) + + return tuple(outputs) + + +class FCN8s(nn.Module): + def __init__(self, nclass, backbone='vgg16', aux=False, pretrained_base=True, norm_layer=nn.BatchNorm2d, **kwargs): + super(FCN8s, self).__init__() + self.aux = aux + if backbone == 'vgg16': + self.pretrained = vgg16(pretrained=pretrained_base).features + else: + raise RuntimeError('unknown backbone: {}'.format(backbone)) + self.pool3 = nn.Sequential(*self.pretrained[:17]) + self.pool4 = nn.Sequential(*self.pretrained[17:24]) + self.pool5 = nn.Sequential(*self.pretrained[24:]) + self.head = _FCNHead(512, nclass, norm_layer) + self.score_pool3 = nn.Conv2d(256, nclass, 1) + self.score_pool4 = nn.Conv2d(512, nclass, 1) + if aux: + self.auxlayer = _FCNHead(512, nclass, norm_layer) + + self.__setattr__('exclusive', + ['head', 'score_pool3', 'score_pool4', 'auxlayer'] if aux else ['head', 'score_pool3', + 'score_pool4']) + + def forward(self, x): + pool3 = self.pool3(x) + pool4 = self.pool4(pool3) + pool5 = self.pool5(pool4) + + outputs = [] + score_fr = self.head(pool5) + + score_pool4 = self.score_pool4(pool4) + score_pool3 = self.score_pool3(pool3) + + upscore2 = F.interpolate(score_fr, score_pool4.size()[2:], mode='bilinear', align_corners=True) + fuse_pool4 = upscore2 + score_pool4 + + upscore_pool4 = F.interpolate(fuse_pool4, score_pool3.size()[2:], mode='bilinear', align_corners=True) + fuse_pool3 = upscore_pool4 + score_pool3 + + out = F.interpolate(fuse_pool3, x.size()[2:], mode='bilinear', align_corners=True) + outputs.append(out) + + if self.aux: + auxout = self.auxlayer(pool5) + auxout = F.interpolate(auxout, x.size()[2:], mode='bilinear', align_corners=True) + outputs.append(auxout) + + return tuple(outputs) + + +class _FCNHead(nn.Module): + def __init__(self, in_channels, channels, norm_layer=nn.BatchNorm2d, **kwargs): + super(_FCNHead, self).__init__() + inter_channels = in_channels // 4 + self.block = nn.Sequential( + nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False), + norm_layer(inter_channels), + nn.ReLU(inplace=True), + nn.Dropout(0.1), + nn.Conv2d(inter_channels, channels, 1) + ) + + def forward(self, x): + return self.block(x) + + +def get_fcn32s(dataset='pascal_voc', backbone='vgg16', pretrained=False, root='~/.torch/models', + pretrained_base=True, **kwargs): + acronyms = { + 'pascal_voc': 'pascal_voc', + 'pascal_aug': 'pascal_aug', + 'ade20k': 'ade', + 'coco': 'coco', + 'citys': 'citys', + } + from ..data.dataloader import datasets + model = FCN32s(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs) + if pretrained: + from .model_store import get_model_file + device = torch.device(kwargs['local_rank']) + model.load_state_dict(torch.load(get_model_file('fcn32s_%s_%s' % (backbone, acronyms[dataset]), root=root), + map_location=device)) + return model + + +def get_fcn16s(dataset='pascal_voc', backbone='vgg16', pretrained=False, root='~/.torch/models', + pretrained_base=True, **kwargs): + acronyms = { + 'pascal_voc': 'pascal_voc', + 'pascal_aug': 'pascal_aug', + 'ade20k': 'ade', + 'coco': 'coco', + 'citys': 'citys', + } + from ..data.dataloader import datasets + model = FCN16s(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs) + if pretrained: + from .model_store import get_model_file + device = torch.device(kwargs['local_rank']) + model.load_state_dict(torch.load(get_model_file('fcn16s_%s_%s' % (backbone, acronyms[dataset]), root=root), + map_location=device)) + return model + + +def get_fcn8s(dataset='pascal_voc', backbone='vgg16', pretrained=False, root='~/.torch/models', + pretrained_base=True, **kwargs): + acronyms = { + 'pascal_voc': 'pascal_voc', + 'pascal_aug': 'pascal_aug', + 'ade20k': 'ade', + 'coco': 'coco', + 'citys': 'citys', + } + from ..data.dataloader import datasets + model = FCN8s(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs) + if pretrained: + from .model_store import get_model_file + device = torch.device(kwargs['local_rank']) + model.load_state_dict(torch.load(get_model_file('fcn8s_%s_%s' % (backbone, acronyms[dataset]), root=root), + map_location=device)) + return model + + +def get_fcn32s_vgg16_voc(**kwargs): + return get_fcn32s('pascal_voc', 'vgg16', **kwargs) + + +def get_fcn16s_vgg16_voc(**kwargs): + return get_fcn16s('pascal_voc', 'vgg16', **kwargs) + + +def get_fcn8s_vgg16_voc(**kwargs): + return get_fcn8s('pascal_voc', 'vgg16', **kwargs) + + +if __name__ == '__main__': + model = FCN16s(21) + print(model) diff --git a/hair_service_sd/hair_matting/seg/networks/jpu.py b/hair_service_sd/hair_matting/seg/networks/jpu.py new file mode 100644 index 0000000..db23bab --- /dev/null +++ b/hair_service_sd/hair_matting/seg/networks/jpu.py @@ -0,0 +1,68 @@ +"""Joint Pyramid Upsampling""" +import torch +import torch.nn as nn +import torch.nn.functional as F + +__all__ = ['JPU'] + + +class SeparableConv2d(nn.Module): + def __init__(self, inplanes, planes, kernel_size=3, stride=1, padding=1, + dilation=1, bias=False, norm_layer=nn.BatchNorm2d): + super(SeparableConv2d, self).__init__() + self.conv = nn.Conv2d(inplanes, inplanes, kernel_size, stride, padding, dilation, groups=inplanes, bias=bias) + self.bn = norm_layer(inplanes) + self.pointwise = nn.Conv2d(inplanes, planes, 1, bias=bias) + + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + x = self.pointwise(x) + return x + + +# copy from: https://github.com/wuhuikai/FastFCN/blob/master/encoding/nn/customize.py +class JPU(nn.Module): + def __init__(self, in_channels, width=512, norm_layer=nn.BatchNorm2d, **kwargs): + super(JPU, self).__init__() + + self.conv5 = nn.Sequential( + nn.Conv2d(in_channels[-1], width, 3, padding=1, bias=False), + norm_layer(width), + nn.ReLU(True)) + self.conv4 = nn.Sequential( + nn.Conv2d(in_channels[-2], width, 3, padding=1, bias=False), + norm_layer(width), + nn.ReLU(True)) + self.conv3 = nn.Sequential( + nn.Conv2d(in_channels[-3], width, 3, padding=1, bias=False), + norm_layer(width), + nn.ReLU(True)) + + self.dilation1 = nn.Sequential( + SeparableConv2d(3 * width, width, 3, padding=1, dilation=1, bias=False), + norm_layer(width), + nn.ReLU(True)) + self.dilation2 = nn.Sequential( + SeparableConv2d(3 * width, width, 3, padding=2, dilation=2, bias=False), + norm_layer(width), + nn.ReLU(True)) + self.dilation3 = nn.Sequential( + SeparableConv2d(3 * width, width, 3, padding=4, dilation=4, bias=False), + norm_layer(width), + nn.ReLU(True)) + self.dilation4 = nn.Sequential( + SeparableConv2d(3 * width, width, 3, padding=8, dilation=8, bias=False), + norm_layer(width), + nn.ReLU(True)) + + def forward(self, *inputs): + feats = [self.conv5(inputs[-1]), self.conv4(inputs[-2]), self.conv3(inputs[-3])] + size = feats[-1].size()[2:] + feats[-2] = F.interpolate(feats[-2], size, mode='bilinear', align_corners=True) + feats[-3] = F.interpolate(feats[-3], size, mode='bilinear', align_corners=True) + feat = torch.cat(feats, dim=1) + feat = torch.cat([self.dilation1(feat), self.dilation2(feat), self.dilation3(feat), self.dilation4(feat)], + dim=1) + + return inputs[0], inputs[1], inputs[2], feat diff --git a/hair_service_sd/hair_matting/seg/networks/resnetv1b.py b/hair_service_sd/hair_matting/seg/networks/resnetv1b.py new file mode 100644 index 0000000..21d67b7 --- /dev/null +++ b/hair_service_sd/hair_matting/seg/networks/resnetv1b.py @@ -0,0 +1,264 @@ +import torch +import torch.nn as nn +import torch.utils.model_zoo as model_zoo + +__all__ = ['ResNetV1b', 'resnet18_v1b', 'resnet34_v1b', 'resnet50_v1b', + 'resnet101_v1b', 'resnet152_v1b', 'resnet152_v1s', 'resnet101_v1s', 'resnet50_v1s'] + +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', +} + + +class BasicBlockV1b(nn.Module): + expansion = 1 + + def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None, + previous_dilation=1, norm_layer=nn.BatchNorm2d): + super(BasicBlockV1b, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, 3, stride, + dilation, dilation, bias=False) + self.bn1 = norm_layer(planes) + self.relu = nn.ReLU(True) + self.conv2 = nn.Conv2d(planes, planes, 3, 1, previous_dilation, + dilation=previous_dilation, bias=False) + 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 BottleneckV1b(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None, + previous_dilation=1, norm_layer=nn.BatchNorm2d): + super(BottleneckV1b, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False) + self.bn1 = norm_layer(planes) + self.conv2 = nn.Conv2d(planes, planes, 3, stride, + dilation, dilation, bias=False) + self.bn2 = norm_layer(planes) + self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False) + self.bn3 = norm_layer(planes * self.expansion) + self.relu = nn.ReLU(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 ResNetV1b(nn.Module): + + def __init__(self, block, layers, num_classes=1000, dilated=True, deep_stem=False, + zero_init_residual=False, norm_layer=nn.BatchNorm2d): + self.inplanes = 128 if deep_stem else 64 + super(ResNetV1b, self).__init__() + if deep_stem: + self.conv1 = nn.Sequential( + nn.Conv2d(3, 64, 3, 2, 1, bias=False), + norm_layer(64), + nn.ReLU(True), + nn.Conv2d(64, 64, 3, 1, 1, bias=False), + norm_layer(64), + nn.ReLU(True), + nn.Conv2d(64, 128, 3, 1, 1, bias=False) + ) + else: + self.conv1 = nn.Conv2d(3, 64, 7, 2, 3, bias=False) + self.bn1 = norm_layer(self.inplanes) + self.relu = nn.ReLU(True) + self.maxpool = nn.MaxPool2d(3, 2, 1) + self.layer1 = self._make_layer(block, 64, layers[0], norm_layer=norm_layer) + self.layer2 = self._make_layer(block, 128, layers[1], stride=2, norm_layer=norm_layer) + if dilated: + self.layer3 = self._make_layer(block, 256, layers[2], stride=1, dilation=2, norm_layer=norm_layer) + self.layer4 = self._make_layer(block, 512, layers[3], stride=1, dilation=4, norm_layer=norm_layer) + else: + self.layer3 = self._make_layer(block, 256, layers[2], stride=2, norm_layer=norm_layer) + self.layer4 = self._make_layer(block, 512, layers[3], stride=2, norm_layer=norm_layer) + 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.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + if zero_init_residual: + for m in self.modules(): + if isinstance(m, BottleneckV1b): + nn.init.constant_(m.bn3.weight, 0) + elif isinstance(m, BasicBlockV1b): + nn.init.constant_(m.bn2.weight, 0) + + def _make_layer(self, block, planes, blocks, stride=1, dilation=1, norm_layer=nn.BatchNorm2d): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.inplanes, planes * block.expansion, 1, stride, bias=False), + norm_layer(planes * block.expansion), + ) + + layers = [] + if dilation in (1, 2): + layers.append(block(self.inplanes, planes, stride, dilation=1, downsample=downsample, + previous_dilation=dilation, norm_layer=norm_layer)) + elif dilation == 4: + layers.append(block(self.inplanes, planes, stride, dilation=2, downsample=downsample, + previous_dilation=dilation, norm_layer=norm_layer)) + else: + raise RuntimeError("=> unknown dilation size: {}".format(dilation)) + self.inplanes = planes * block.expansion + for _ in range(1, blocks): + layers.append(block(self.inplanes, planes, dilation=dilation, + previous_dilation=dilation, norm_layer=norm_layer)) + + return nn.Sequential(*layers) + + def forward(self, x): + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + x = self.maxpool(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + + x = self.avgpool(x) + x = x.view(x.size(0), -1) + x = self.fc(x) + + return x + + +def resnet18_v1b(pretrained=False, **kwargs): + model = ResNetV1b(BasicBlockV1b, [2, 2, 2, 2], **kwargs) + if pretrained: + old_dict = model_zoo.load_url(model_urls['resnet18']) + model_dict = model.state_dict() + old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)} + model_dict.update(old_dict) + model.load_state_dict(model_dict) + return model + + +def resnet34_v1b(pretrained=False, **kwargs): + model = ResNetV1b(BasicBlockV1b, [3, 4, 6, 3], **kwargs) + if pretrained: + old_dict = model_zoo.load_url(model_urls['resnet34']) + model_dict = model.state_dict() + old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)} + model_dict.update(old_dict) + model.load_state_dict(model_dict) + return model + + +def resnet50_v1b(pretrained=False, **kwargs): + model = ResNetV1b(BottleneckV1b, [3, 4, 6, 3], **kwargs) + if pretrained: + old_dict = model_zoo.load_url(model_urls['resnet50']) + model_dict = model.state_dict() + old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)} + model_dict.update(old_dict) + model.load_state_dict(model_dict) + return model + + +def resnet101_v1b(pretrained=False, **kwargs): + model = ResNetV1b(BottleneckV1b, [3, 4, 23, 3], **kwargs) + if pretrained: + old_dict = model_zoo.load_url(model_urls['resnet101']) + model_dict = model.state_dict() + old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)} + model_dict.update(old_dict) + model.load_state_dict(model_dict) + return model + + +def resnet152_v1b(pretrained=False, **kwargs): + model = ResNetV1b(BottleneckV1b, [3, 8, 36, 3], **kwargs) + if pretrained: + old_dict = model_zoo.load_url(model_urls['resnet152']) + model_dict = model.state_dict() + old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)} + model_dict.update(old_dict) + model.load_state_dict(model_dict) + return model + + +def resnet50_v1s(pretrained=False, root='~/.torch/models', **kwargs): + model = ResNetV1b(BottleneckV1b, [3, 4, 6, 3], deep_stem=True, **kwargs) + if pretrained: + from ..model_store import get_resnet_file + model.load_state_dict(torch.load(get_resnet_file('resnet50', root=root)), strict=False) + return model + + +def resnet101_v1s(pretrained=False, root='~/.torch/models', **kwargs): + model = ResNetV1b(BottleneckV1b, [3, 4, 23, 3], deep_stem=True, **kwargs) + if pretrained: + from ..model_store import get_resnet_file + model.load_state_dict(torch.load(get_resnet_file('resnet101', root=root)), strict=False) + return model + + +def resnet152_v1s(pretrained=False, root='~/.torch/models', **kwargs): + model = ResNetV1b(BottleneckV1b, [3, 8, 36, 3], deep_stem=True, **kwargs) + if pretrained: + from ..model_store import get_resnet_file + model.load_state_dict(torch.load(get_resnet_file('resnet152', root=root)), strict=False) + return model + + +if __name__ == '__main__': + import torch + + img = torch.randn(4, 3, 224, 224) + model = resnet50_v1b(True) + output = model(img) diff --git a/hair_service_sd/hair_matting/seg/networks/segbase.py b/hair_service_sd/hair_matting/seg/networks/segbase.py new file mode 100644 index 0000000..3c1822d --- /dev/null +++ b/hair_service_sd/hair_matting/seg/networks/segbase.py @@ -0,0 +1,60 @@ +"""Base Model for Semantic Segmentation""" +import torch.nn as nn + +from hair_matting.seg.networks.jpu import JPU +from hair_matting.seg.networks.resnetv1b import resnet50_v1s, resnet101_v1s, resnet152_v1s + +__all__ = ['SegBaseModel'] + + +class SegBaseModel(nn.Module): + r"""Base Model for Semantic Segmentation + + Parameters + ---------- + backbone : string + Pre-trained dilated backbone network type (default:'resnet50'; 'resnet50', + 'resnet101' or 'resnet152'). + """ + + def __init__(self, nclass, aux, backbone='resnet50', jpu=False, pretrained_base=True, **kwargs): + super(SegBaseModel, self).__init__() + dilated = False if jpu else True + self.aux = aux + self.nclass = nclass + if backbone == 'resnet50': + self.pretrained = resnet50_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs) + elif backbone == 'resnet101': + self.pretrained = resnet101_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs) + elif backbone == 'resnet152': + self.pretrained = resnet152_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs) + else: + raise RuntimeError('unknown backbone: {}'.format(backbone)) + + self.jpu = JPU([512, 1024, 2048], width=512, **kwargs) if jpu else None + + def base_forward(self, x): + """forwarding pre-trained network""" + x = self.pretrained.conv1(x) + x = self.pretrained.bn1(x) + x = self.pretrained.relu(x) + x = self.pretrained.maxpool(x) + c1 = self.pretrained.layer1(x) + c2 = self.pretrained.layer2(c1) + c3 = self.pretrained.layer3(c2) + c4 = self.pretrained.layer4(c3) + + if self.jpu: + return self.jpu(c1, c2, c3, c4) + else: + return c1, c2, c3, c4 + + def evaluate(self, x): + """evaluating network with inputs and targets""" + return self.forward(x)[0] + + def demo(self, x): + pred = self.forward(x) + if self.aux: + pred = pred[0] + return pred diff --git a/hair_service_sd/hair_matting/seg/networks/vgg.py b/hair_service_sd/hair_matting/seg/networks/vgg.py new file mode 100644 index 0000000..fe5c163 --- /dev/null +++ b/hair_service_sd/hair_matting/seg/networks/vgg.py @@ -0,0 +1,191 @@ +import torch +import torch.nn as nn +import torch.utils.model_zoo as model_zoo + +__all__ = [ + 'VGG', 'vgg11', 'vgg11_bn', 'vgg13', 'vgg13_bn', 'vgg16', 'vgg16_bn', + 'vgg19_bn', 'vgg19', +] + +model_urls = { + 'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth', + 'vgg13': 'https://download.pytorch.org/models/vgg13-c768596a.pth', + 'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth', + 'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth', + 'vgg11_bn': 'https://download.pytorch.org/models/vgg11_bn-6002323d.pth', + 'vgg13_bn': 'https://download.pytorch.org/models/vgg13_bn-abd245e5.pth', + 'vgg16_bn': 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth', + 'vgg19_bn': 'https://download.pytorch.org/models/vgg19_bn-c79401a0.pth', +} + + +class VGG(nn.Module): + def __init__(self, features, num_classes=1000, init_weights=True): + super(VGG, self).__init__() + self.features = features + self.avgpool = nn.AdaptiveAvgPool2d((7, 7)) + self.classifier = nn.Sequential( + nn.Linear(512 * 7 * 7, 4096), + nn.ReLU(True), + nn.Dropout(), + nn.Linear(4096, 4096), + nn.ReLU(True), + nn.Dropout(), + nn.Linear(4096, num_classes) + ) + if init_weights: + self._initialize_weights() + + def forward(self, x): + x = self.features(x) + x = self.avgpool(x) + x = x.view(x.size(0), -1) + x = self.classifier(x) + return x + + def _initialize_weights(self): + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.BatchNorm2d): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.Linear): + nn.init.normal_(m.weight, 0, 0.01) + nn.init.constant_(m.bias, 0) + + +def make_layers(cfg, batch_norm=False): + layers = [] + in_channels = 3 + for v in cfg: + if v == 'M': + layers += [nn.MaxPool2d(kernel_size=2, stride=2)] + else: + conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1) + if batch_norm: + layers += (conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)) + else: + layers += [conv2d, nn.ReLU(inplace=True)] + in_channels = v + return nn.Sequential(*layers) + + +cfg = { + 'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], + 'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], + 'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'], + 'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'], +} + + +def vgg11(pretrained=False, **kwargs): + """VGG 11-layer model (configuration "A") + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['A']), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg11'])) + return model + + +def vgg11_bn(pretrained=False, **kwargs): + """VGG 11-layer model (configuration "A") with batch normalization + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['A'], batch_norm=True), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg11_bn'])) + return model + + +def vgg13(pretrained=False, **kwargs): + """VGG 13-layer model (configuration "B") + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['B']), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg13'])) + return model + + +def vgg13_bn(pretrained=False, **kwargs): + """VGG 13-layer model (configuration "B") with batch normalization + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['B'], batch_norm=True), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg13_bn'])) + return model + + +def vgg16(pretrained=False, **kwargs): + """VGG 16-layer model (configuration "D") + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['D']), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg16'])) + return model + + +def vgg16_bn(pretrained=False, **kwargs): + """VGG 16-layer model (configuration "D") with batch normalization + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['D'], batch_norm=True), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg16_bn'])) + return model + + +def vgg19(pretrained=False, **kwargs): + """VGG 19-layer model (configuration "E") + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['E']), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg19'])) + return model + + +def vgg19_bn(pretrained=False, **kwargs): + """VGG 19-layer model (configuration 'E') with batch normalization + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['E'], batch_norm=True), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg19_bn'])) + return model + + +if __name__ == '__main__': + img = torch.randn((4, 3, 480, 480)) + model = vgg16(pretrained=False) + out = model(img) diff --git a/hair_service_sd/hair_matting/seg/networks/xception.py b/hair_service_sd/hair_matting/seg/networks/xception.py new file mode 100644 index 0000000..52dc0b9 --- /dev/null +++ b/hair_service_sd/hair_matting/seg/networks/xception.py @@ -0,0 +1,411 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +__all__ = ['Enc', 'FCAttention', 'Xception65', 'Xception71', 'get_xception', 'get_xception_71', 'get_xception_a'] + + +class SeparableConv2d(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, dilation=1, bias=False, norm_layer=None): + super(SeparableConv2d, self).__init__() + self.kernel_size = kernel_size + self.dilation = dilation + + self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size, stride, 0, dilation, groups=in_channels, + bias=bias) + self.bn = norm_layer(in_channels) + self.pointwise = nn.Conv2d(in_channels, out_channels, 1, bias=bias) + + def forward(self, x): + x = self.fix_padding(x, self.kernel_size, self.dilation) + x = self.conv1(x) + x = self.bn(x) + x = self.pointwise(x) + + return x + + def fix_padding(self, x, 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(x, (pad_beg, pad_end, pad_beg, pad_end)) + return padded_inputs + + +class Block(nn.Module): + def __init__(self, in_channels, out_channels, reps, stride=1, dilation=1, norm_layer=None, + start_with_relu=True, grow_first=True, is_last=False): + super(Block, self).__init__() + if out_channels != in_channels or stride != 1: + self.skip = nn.Conv2d(in_channels, out_channels, 1, stride, bias=False) + self.skipbn = norm_layer(out_channels) + else: + self.skip = None + self.relu = nn.ReLU(True) + rep = list() + filters = in_channels + if grow_first: + if start_with_relu: + rep.append(self.relu) + rep.append(SeparableConv2d(in_channels, out_channels, 3, 1, dilation, norm_layer=norm_layer)) + rep.append(norm_layer(out_channels)) + filters = out_channels + for i in range(reps - 1): + if grow_first or start_with_relu: + rep.append(self.relu) + rep.append(SeparableConv2d(filters, filters, 3, 1, dilation, norm_layer=norm_layer)) + rep.append(norm_layer(filters)) + if not grow_first: + rep.append(self.relu) + rep.append(SeparableConv2d(in_channels, out_channels, 3, 1, dilation, norm_layer=norm_layer)) + if stride != 1: + rep.append(self.relu) + rep.append(SeparableConv2d(out_channels, out_channels, 3, stride, norm_layer=norm_layer)) + rep.append(norm_layer(out_channels)) + elif is_last: + rep.append(self.relu) + rep.append(SeparableConv2d(out_channels, out_channels, 3, 1, dilation, norm_layer=norm_layer)) + rep.append(norm_layer(out_channels)) + self.rep = nn.Sequential(*rep) + + def forward(self, x): + out = self.rep(x) + if self.skip is not None: + skip = self.skipbn(self.skip(x)) + else: + skip = x + out = out + skip + return out + + +class Xception65(nn.Module): + """Modified Aligned Xception + """ + + def __init__(self, num_classes=1000, output_stride=32, norm_layer=nn.BatchNorm2d): + super(Xception65, self).__init__() + if output_stride == 32: + entry_block3_stride = 2 + exit_block20_stride = 2 + middle_block_dilation = 1 + exit_block_dilations = (1, 1) + elif output_stride == 16: + entry_block3_stride = 2 + exit_block20_stride = 1 + middle_block_dilation = 1 + exit_block_dilations = (1, 2) + elif output_stride == 8: + entry_block3_stride = 1 + exit_block20_stride = 1 + middle_block_dilation = 2 + exit_block_dilations = (2, 4) + else: + raise NotImplementedError + # Entry flow + self.conv1 = nn.Conv2d(3, 32, 3, 2, 1, bias=False) + self.bn1 = norm_layer(32) + self.relu = nn.ReLU(True) + + self.conv2 = nn.Conv2d(32, 64, 3, 1, 1, bias=False) + self.bn2 = norm_layer(64) + + self.block1 = Block(64, 128, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False) + self.block2 = Block(128, 256, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False, grow_first=True) + self.block3 = Block(256, 728, reps=2, stride=entry_block3_stride, norm_layer=norm_layer, + start_with_relu=True, grow_first=True, is_last=True) + + # Middle flow + midflow = list() + for i in range(4, 20): + midflow.append(Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation, norm_layer=norm_layer, + start_with_relu=True, grow_first=True)) + self.midflow = nn.Sequential(*midflow) + + # Exit flow + self.block20 = Block(728, 1024, reps=2, stride=exit_block20_stride, dilation=exit_block_dilations[0], + norm_layer=norm_layer, start_with_relu=True, grow_first=False, is_last=True) + self.conv3 = SeparableConv2d(1024, 1536, 3, 1, dilation=exit_block_dilations[1], norm_layer=norm_layer) + self.bn3 = norm_layer(1536) + self.conv4 = SeparableConv2d(1536, 1536, 3, stride=1, dilation=exit_block_dilations[1], norm_layer=norm_layer) + self.bn4 = norm_layer(1536) + self.conv5 = SeparableConv2d(1536, 2048, 3, 1, dilation=exit_block_dilations[1], norm_layer=norm_layer) + self.bn5 = norm_layer(2048) + self.avgpool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Linear(2048, num_classes) + + 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) + x = self.relu(x) + # c1 = x + x = self.block2(x) + # c2 = x + x = self.block3(x) + + # Middle flow + x = self.midflow(x) + # c3 = 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) + + x = self.avgpool(x) + x = x.view(x.size(0), -1) + x = self.fc(x) + + return x + + +class Xception71(nn.Module): + """Modified Aligned Xception + """ + + def __init__(self, num_classes=1000, output_stride=32, norm_layer=nn.BatchNorm2d): + super(Xception71, self).__init__() + if output_stride == 32: + entry_block3_stride = 2 + exit_block20_stride = 2 + middle_block_dilation = 1 + exit_block_dilations = (1, 1) + elif output_stride == 16: + entry_block3_stride = 2 + exit_block20_stride = 1 + middle_block_dilation = 1 + exit_block_dilations = (1, 2) + elif output_stride == 8: + entry_block3_stride = 1 + exit_block20_stride = 1 + middle_block_dilation = 2 + exit_block_dilations = (2, 4) + else: + raise NotImplementedError + # Entry flow + self.conv1 = nn.Conv2d(3, 32, 3, 2, 1, bias=False) + self.bn1 = norm_layer(32) + self.relu = nn.ReLU(True) + + self.conv2 = nn.Conv2d(32, 64, 3, 1, 1, bias=False) + self.bn2 = norm_layer(64) + + self.block1 = Block(64, 128, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False) + self.block2 = nn.Sequential( + Block(128, 256, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False, grow_first=True), + Block(256, 728, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False, grow_first=True)) + self.block3 = Block(728, 728, reps=2, stride=entry_block3_stride, norm_layer=norm_layer, + start_with_relu=True, grow_first=True, is_last=True) + + # Middle flow + midflow = list() + for i in range(4, 20): + midflow.append(Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation, norm_layer=norm_layer, + start_with_relu=True, grow_first=True)) + self.midflow = nn.Sequential(*midflow) + + # Exit flow + self.block20 = Block(728, 1024, reps=2, stride=exit_block20_stride, dilation=exit_block_dilations[0], + norm_layer=norm_layer, start_with_relu=True, grow_first=False, is_last=True) + self.conv3 = SeparableConv2d(1024, 1536, 3, 1, dilation=exit_block_dilations[1], norm_layer=norm_layer) + self.bn3 = norm_layer(1536) + self.conv4 = SeparableConv2d(1536, 1536, 3, stride=1, dilation=exit_block_dilations[1], norm_layer=norm_layer) + self.bn4 = norm_layer(1536) + self.conv5 = SeparableConv2d(1536, 2048, 3, 1, dilation=exit_block_dilations[1], norm_layer=norm_layer) + self.bn5 = norm_layer(2048) + self.avgpool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Linear(2048, num_classes) + + 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) + x = self.relu(x) + # c1 = x + x = self.block2(x) + # c2 = x + x = self.block3(x) + + # Middle flow + x = self.midflow(x) + # c3 = 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) + + x = self.avgpool(x) + x = x.view(x.size(0), -1) + x = self.fc(x) + + return x + + +# ------------------------------------------------- +# For DFANet +# ------------------------------------------------- +class BlockA(nn.Module): + def __init__(self, in_channels, out_channels, stride=1, dilation=1, norm_layer=None, start_with_relu=True): + super(BlockA, self).__init__() + if out_channels != in_channels or stride != 1: + self.skip = nn.Conv2d(in_channels, out_channels, 1, stride, bias=False) + self.skipbn = norm_layer(out_channels) + else: + self.skip = None + self.relu = nn.ReLU(True) + rep = list() + inter_channels = out_channels // 4 + + if start_with_relu: + rep.append(self.relu) + rep.append(SeparableConv2d(in_channels, inter_channels, 3, 1, dilation, norm_layer=norm_layer)) + rep.append(norm_layer(inter_channels)) + + rep.append(self.relu) + rep.append(SeparableConv2d(inter_channels, inter_channels, 3, 1, dilation, norm_layer=norm_layer)) + rep.append(norm_layer(inter_channels)) + + if stride != 1: + rep.append(self.relu) + rep.append(SeparableConv2d(inter_channels, out_channels, 3, stride, norm_layer=norm_layer)) + rep.append(norm_layer(out_channels)) + else: + rep.append(self.relu) + rep.append(SeparableConv2d(inter_channels, out_channels, 3, 1, norm_layer=norm_layer)) + rep.append(norm_layer(out_channels)) + self.rep = nn.Sequential(*rep) + + def forward(self, x): + out = self.rep(x) + if self.skip is not None: + skip = self.skipbn(self.skip(x)) + else: + skip = x + out = out + skip + return out + + +class Enc(nn.Module): + def __init__(self, in_channels, out_channels, blocks, norm_layer=None): + super(Enc, self).__init__() + block = list() + block.append(BlockA(in_channels, out_channels, 2, norm_layer=norm_layer)) + for i in range(blocks - 1): + block.append(BlockA(out_channels, out_channels, 1, norm_layer=norm_layer)) + self.block = nn.Sequential(*block) + + def forward(self, x): + return self.block(x) + + +class FCAttention(nn.Module): + def __init__(self, in_channels, norm_layer=None): + super(FCAttention, self).__init__() + self.avgpool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Linear(in_channels, 1000) + self.conv = nn.Sequential( + nn.Conv2d(1000, in_channels, 1, bias=False), + norm_layer(in_channels), + nn.ReLU(True)) + + def forward(self, x): + n, c, _, _ = x.size() + att = self.avgpool(x).view(n, c) + att = self.fc(att).view(n, 1000, 1, 1) + att = self.conv(att) + return x * att.expand_as(x) + + +class XceptionA(nn.Module): + def __init__(self, num_classes=1000, norm_layer=nn.BatchNorm2d): + super(XceptionA, self).__init__() + self.conv1 = nn.Sequential(nn.Conv2d(3, 8, 3, 2, 1, bias=False), + norm_layer(8), + nn.ReLU(True)) + + self.enc2 = Enc(8, 48, 4, norm_layer=norm_layer) + self.enc3 = Enc(48, 96, 6, norm_layer=norm_layer) + self.enc4 = Enc(96, 192, 4, norm_layer=norm_layer) + + self.fca = FCAttention(192, norm_layer=norm_layer) + self.avgpool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Linear(192, num_classes) + + def forward(self, x): + x = self.conv1(x) + + x = self.enc2(x) + x = self.enc3(x) + x = self.enc4(x) + x = self.fca(x) + + x = self.avgpool(x) + x = x.view(x.size(0), -1) + x = self.fc(x) + + return x + + +# Constructor +def get_xception(pretrained=False, root='~/.torch/models', **kwargs): + model = Xception65(**kwargs) + if pretrained: + from ..model_store import get_model_file + model.load_state_dict(torch.load(get_model_file('xception', root=root))) + return model + + +def get_xception_71(pretrained=False, root='~/.torch/models', **kwargs): + model = Xception71(**kwargs) + if pretrained: + from ..model_store import get_model_file + model.load_state_dict(torch.load(get_model_file('xception71', root=root))) + return model + + +def get_xception_a(pretrained=False, root='~/.torch/models', **kwargs): + model = XceptionA(**kwargs) + if pretrained: + from ..model_store import get_model_file + model.load_state_dict(torch.load(get_model_file('xception_a', root=root))) + return model + + +if __name__ == '__main__': + model = get_xception_a() diff --git a/hair_service_sd/hair_matting/seg/setup.py b/hair_service_sd/hair_matting/seg/setup.py new file mode 100644 index 0000000..ca6bf5a --- /dev/null +++ b/hair_service_sd/hair_matting/seg/setup.py @@ -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 \ No newline at end of file diff --git a/hair_service_sd/hair_matting/seg/test_single.py b/hair_service_sd/hair_matting/seg/test_single.py new file mode 100644 index 0000000..0c9162d --- /dev/null +++ b/hair_service_sd/hair_matting/seg/test_single.py @@ -0,0 +1,24 @@ +from seg.hairseg_single_model import Evaluator +import os +import cv2 +import numpy as np +if __name__ == "__main__": + + data_path = "/home/liyang/project/matting/合格/origin" + dst_path = "/home/liyang/project/matting/合格/origin_seg_res1102" + if not os.path.exists(dst_path): + os.mkdir(dst_path) + seg_model = Evaluator(gpu_id=0, output_img_size=512, nclass=3) + for imgs in os.listdir(data_path): + if imgs.endswith(".txt"): + continue + + img_path = os.path.join(data_path, imgs) + img = cv2.imread(img_path) + if img_path.endswith('.png'): + kpts_1k = np.loadtxt(img_path.replace('.png', '_landmark1k.txt')) + elif img_path.endswith('.jpg'): + kpts_1k = np.loadtxt(img_path.replace('.jpg', '_landmark1k.txt')) + output_img_size = 512 + mask = seg_model.eval(img, kpts_1k) + cv2.imwrite(os.path.join(dst_path, imgs), mask) diff --git a/hair_service_sd/hairstyle_model_infer.py b/hair_service_sd/hairstyle_model_infer.py new file mode 100644 index 0000000..29a6636 --- /dev/null +++ b/hair_service_sd/hairstyle_model_infer.py @@ -0,0 +1,1672 @@ +import pickle +import time +import cv2 +import numpy as np +import torch +import json +import os +import configparser +from process_modules import PersonProcessor_yolov5,KeypointsProcessor,Human_Keypoints,pt_conv_25_to_17 +from process_modules import Get_Landmark, Process_Data, Generator_Hair, Generator_Fusion_Res, Change_Hair_Color, BodySeg +from prepare_ref_hairstyle_data import GenderClassifyProcessor +from face_enhance.face_enhancement import FaceEnhancement +from utils import landmark_processor +from common.logger import LogFactory +from faceseg.face_seg import FaceSeg +from hair_init import HairInit + +config = configparser.ConfigParser() # 创建对象 +config.read("config/configure.ini", encoding="utf-8") # 读取配置文件，如果配置文件不存在则创建 + +class HairStyle_Model_Infer(object): + def __init__(self, gpu=True, use_enhance=False): + hair_init = HairInit() + self.use_enhance = use_enhance + self.gpu_index = hair_init.gpu_index + self.get_landmark = hair_init.get_landmark + self.hair_size = hair_init.hair_size + self.process_data_infer = hair_init.process_data_infer + self.generator_hair = hair_init.generator_hair + self.hair_fusion = hair_init.hair_fusion + self.logger_call = hair_init.logger_call + self.face_seg = hair_init.face_seg + self.face_enhance = hair_init.face_enhance + self.change_haircolor = hair_init.change_haircolor + self.gender_classify = hair_init.gender_classify + + self.person_processor = hair_init.person_processor + self.keypoints_processor = hair_init.keypoints_processor + self.human_keypoint = hair_init.human_keypoint + + + + # 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.get_landmark = Get_Landmark(gpu_id=device_id) + # self.hair_size = 768 + # self.process_data = Process_Data(gpu, device_id) + # self.generator_hair = Generator_Hair(gpu, device_id) + # self.hair_fusion = Generator_Fusion_Res(gpu, device_id) + # self.logger_call = LogFactory.getLogger("call") + # self.face_seg = FaceSeg(device_id) + # self.use_enhance = use_enhance + # if self.use_enhance: + # self.face_enhance = FaceEnhancement(512, device_id) + # self.table_enlight = cv2.imread('./data/convert_enlight.png') + # self.change_haircolor = Change_Hair_Color(gpu, device_id) + # self.gender_classify = GenderClassifyProcessor(gpu_id=device_id) + # + # # body process + # self.person_processor = PersonProcessor_yolov5(gpu_id=device_id) + # self.keypoints_processor = KeypointsProcessor(gpu_id=device_id) + # self.human_keypoint = Human_Keypoints(gpu=True, device_id=device_id) + # print("Load model finish ... ") + + + def infer_haircolor_v4(self, user_rgb_8uc3_orisize, target_hair_color, haircolor_dir): + # haircolor_dir = '/home/data/hair/data/ref_color/3628746832766' + face_base, hair_matting, status = self.infer_haircolor_new(user_rgb_8uc3_orisize, haircolor_dir,target_hair_color, + return_matting=True) + user_rgb_8uc3_orisize = face_base + # 构建一个色板 + r, g, b = target_hair_color + tar_color = np.zeros_like(user_rgb_8uc3_orisize) + tar_color[:, :, 0] = b + tar_color[:, :, 1] = g + tar_color[:, :, 2] = r + + # 将色板和原图转换到LAB空间 + origin_img_LAB = cv2.cvtColor(user_rgb_8uc3_orisize, cv2.COLOR_BGR2LAB) + target_color_LAB = cv2.cvtColor(tar_color, cv2.COLOR_BGR2LAB) + + # img_res_LAB取色板的色度通道和原图的亮度通道 + img_res_LAB = np.concatenate([origin_img_LAB[:, :, 0:1], target_color_LAB[:, :, 1:3]], axis=2) + ret_color_img = cv2.cvtColor(img_res_LAB, cv2.COLOR_LAB2BGR) + + # 将头发的matting转换为三通道 + user_matting_mask_fc32_orisize = np.repeat(hair_matting[:, :, np.newaxis], 3, axis=2).astype(np.float32) / 255 + + ### change ycj + img_res_change_hsv_change_bgr_fc32 = ret_color_img.astype(np.float32) / 255 + ref_rgb_avg_std = ([b / 255.0, g / 255.0, r / 255.0], [0, 0, 0]) + face_base_new_fc32_orisize_new, src_avg_std = self.reinhard_rgb(img_res_change_hsv_change_bgr_fc32, + user_matting_mask_fc32_orisize, + ref_rgb_avg_std, + ratio=1.0) + + face_base_new_fc32_orisize_new = face_base_new_fc32_orisize_new * user_matting_mask_fc32_orisize + \ + (user_rgb_8uc3_orisize.astype(np.float32) / 255) * ( + 1 - user_matting_mask_fc32_orisize) + face_base_new_new = (face_base_new_fc32_orisize_new * 255).astype(np.uint8) + + return face_base_new_new, hair_matting, 0 + + def infer_hairstyle_random_ref(self, origin_img, ref_img): + ratio = 1 + if ratio == 0: + gender = "boy" + else: + gender = "girl" + + user_rgb_8uc3_orisize = origin_img + landmarks_origin_img_1k = self.get_landmark.forward(origin_img) + if landmarks_origin_img_1k is None: + return None, 10001 + + user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, \ + user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M = self.process_data.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio) + + # cv2.imshow("user_bald_res_8uc3_orisize", user_bald_res_8uc3_orisize) + # cv2.waitKey() + + # show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("show_concat", show_concat) + # cv2.waitKey() + + ref_landmark_1k2_f_orisize = self.get_landmark.forward(ref_img) + 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_img, ref_landmark_1k2_f_orisize, ratio=2) + another_pose_hair_image = self.process_data.Generator_reftensor(ref_rgb_8uc3_768, ref_matting_8uc3_768, + ref_baldseg_8uc3_768, + ref_landmark_f1k2_768) + if another_pose_hair_image is None: + return None, 10002 + # 换发型 + hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image, + user_baldseg_8uc3_768, + user_bald_8uc3_768, + user_landmark_f1k2_768, gender) + + hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste( + user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M) + + user_res_8uc3_orisize = self.hair_fusion.inference(hair_gene_fusion_8uc3_orisize, + hair_gene_matte_8uc3_orisize, + landmarks_origin_img_1k, + user_baldseg_8uc3_orisize, + ratio) + + # if self.use_enhance: + user_res_8uc3_orisize = self.face_enhance.process(user_res_8uc3_orisize, landmarks_origin_img_1k) + + # for pt in landmarks_origin_img_1k.astype(np.int32): + # cv2.circle(origin_img, (pt[0], pt[1]), 1, (0, 255, 0), -1) + # show_concat_orisize = np.concatenate((origin_img, user_res_8uc3_orisize), axis=1) + # ratio = 1024. / max(show_concat_orisize.shape[:2]) + # show_concat_orisize = cv2.resize(show_concat_orisize, (0, 0), fx=ratio, fy=ratio) + # ratio = 1024. / max(show_concat_768.shape[:2]) + # show_concat_768 = cv2.resize(show_concat_768, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("show_concat_orisize", show_concat_orisize) + # cv2.imshow("show_concat_768", show_concat_768) + # cv2.waitKey() + + return user_res_8uc3_orisize, 0 + + def draw_hull_mask(self, fc_landmark, w=256, h=256, is_gray=False): + hull_mask = np.zeros((h, w), dtype=np.float32) + if len(fc_landmark) == 1000: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (1)) + else: + raise Exception('landmark should be 1000') + return hull_mask + + + def get_1k(self, origin_img): + landmarks_origin_img_1k = self.get_landmark.forward(origin_img) + return landmarks_origin_img_1k + + def get_body_info(self, img_res): + with torch.no_grad(): + img_w, img_h = 768, 1024 + ori_h, ori_w, _ = img_res.shape + res = self.person_processor.forward(img_res) + if len(res['boxes']) == 0: + print('No person detected') + return [0,0,ori_w, ori_h] + filter_boxes = [] + + for ix, box_score in enumerate(res['scores']): + box_ = res['boxes'][ix] + box_w = box_[2] - box_[0] + box_h = box_[3] - box_[1] + max_box_len = max(box_w, box_h) + if box_score > 0.5 and max_box_len > 150: + filter_boxes.append([box_, box_w * box_h, box_h / img_res.shape[0]]) + + filter_boxes.sort(key=lambda x: x[1], reverse=True) + filter_boxes = list(filter(lambda x: x[2] > 0.2, filter_boxes)) + if len(filter_boxes) == 0: + return [0, 0, ori_w, ori_h] + filter_boxes = [item[0] for item in filter_boxes] + # return filter_boxes, res + # step2 person keypoints (version: hrnet 17pt) + keypoints = self.keypoints_processor.forward(img_res, filter_boxes) + + # step3 human keypoints: hands, face, body keypoints(25pt) + person_keypoints = keypoints[0] + human_box = filter_boxes[0].reshape((-1, 2)) + human_box[0][1] = max(0, human_box[0][1] - ori_h * 0.1) + # human_box[0][0] = max(0, human_box[0][0]-ori_w*0.1) + # human_box[1][0] = min(ori_w, human_box[1][0]+ori_w*0.1) + + knee_keypoints_left = person_keypoints[13] + knee_keypoints_right = person_keypoints[14] + if knee_keypoints_left[2] > 0.4 or knee_keypoints_right[2] > 0.4: + src_body_kpnts_select = person_keypoints[np.where(person_keypoints[:, 2] > 0.2)][:12][:, :2] + src_body_box = cv2.boundingRect(src_body_kpnts_select[np.newaxis, :, :]) + kpnts_bbox_tlx, kpnts_bbox_tly = src_body_box[:2] + kpnts_bbox_brx, kpnts_bbox_bry = kpnts_bbox_tlx + src_body_box[2] - 1, kpnts_bbox_tly + src_body_box[ + 3] - 1 + + kpnts_bbox_tlx, kpnts_bbox_tly = min(kpnts_bbox_tlx, human_box[0][0]), min(kpnts_bbox_tly, + human_box[0][1]) + kpnts_bbox_brx, kpnts_bbox_bry = min(kpnts_bbox_brx, human_box[1][0]), min(kpnts_bbox_bry, + human_box[1][1]) + box_center_x = int(kpnts_bbox_tlx / 2 + kpnts_bbox_brx / 2) + box_hight_now = kpnts_bbox_bry - kpnts_bbox_tly + kpnts_bbox_tlx = max(int(box_center_x - box_hight_now * 3 / 8), 0) + kpnts_bbox_brx = min(int(box_center_x + box_hight_now * 3 / 8), ori_w) + # cv2.rectangle(img_res, (kpnts_bbox_tlx, kpnts_bbox_tly), (kpnts_bbox_brx, kpnts_bbox_bry), (0, 255, 0), + # 2) + waist_keypoints_mean = (person_keypoints[11] + person_keypoints[12]) / 2 + + crop_img = img_res[kpnts_bbox_tly:kpnts_bbox_bry, kpnts_bbox_tlx:kpnts_bbox_brx] + return [kpnts_bbox_tlx, kpnts_bbox_tly,kpnts_bbox_brx, kpnts_bbox_bry] + return [0, 0, ori_w, ori_h] + + def infer_hairstyle(self, origin_img, hairstyle_dir, return_pt1k=False,use_enhance=False): + config_path = os.path.join(hairstyle_dir, "config.json") + if not os.path.exists(config_path): + return None, 10001 + config_info = json.load(open(config_path, "rb")) + gender = config_info["gender"] + ratio = int(config_info["ratio"]) + if ratio == 2: + ratio = 3 + print("gender: ", gender, " ratio: ", ratio) + + user_rgb_8uc3_orisize = origin_img + landmarks_origin_img_1k = self.get_landmark.forward_infer(origin_img) + + # 根据人脸关键点画出人脸区域的mask + hull_mask = self.draw_hull_mask(landmarks_origin_img_1k.astype(np.int32), w=origin_img.shape[1], h=origin_img.shape[0]) + user_mask_save_path = os.path.join(hairstyle_dir, "hull_mask.png") + cv2.imwrite(user_mask_save_path, hull_mask * 255) + + # cv2.imshow('hull_mask', hull_mask) + # cv2.waitKey(0) + + if landmarks_origin_img_1k is None: + return None, 10001 + + user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, \ + user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M, user_matting_8uc3_bald_orisize = self.process_data_infer.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio) + + # show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("show_concat", show_concat) + # cv2.waitKey() + + # cv2.imshow("user_matting_8uc3_bald_orisize", user_matting_8uc3_bald_orisize) + # cv2.waitKey(0) + + user_orig_mask_path = os.path.join(hairstyle_dir, "user_orig_mask.png") + cv2.imwrite(user_orig_mask_path, user_matting_8uc3_bald_orisize) + + another_pose_hair_img_path = os.path.join(hairstyle_dir, "input_another_pose_hair_image.npy") + if not os.path.exists(another_pose_hair_img_path): + return None, 10002 + another_pose_hair_image = np.load(another_pose_hair_img_path) + + # 换发型 + hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image, + user_baldseg_8uc3_768, + user_bald_8uc3_768, + user_landmark_f1k2_768, gender) + + hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data_infer.get_fusion_res_hairpaste( + user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M) + + # cv2.imshow("hair_gene_fusion_8uc3_orisize:", hair_gene_fusion_8uc3_orisize) + # cv2.imshow("hair_gene_matte_8uc3_orisize:", hair_gene_matte_8uc3_orisize) + # cv2.waitKey(0) + + gen_hair_mask_path = os.path.join(hairstyle_dir, "hair_mask.png") + cv2.imwrite(gen_hair_mask_path, hair_gene_matte_8uc3_orisize * 255) + + gen_hair_mask_path_2 = os.path.join(hairstyle_dir, "hair_mask_2.png") + cv2.imwrite(gen_hair_mask_path_2, hair_gene_matte_8uc3_orisize) + + if use_enhance: + user_res_8uc3_orisize = self.hair_fusion.inference(hair_gene_fusion_8uc3_orisize, + hair_gene_matte_8uc3_orisize, + landmarks_origin_img_1k, + user_baldseg_8uc3_orisize, + ratio) + hair_gene_fusion_8uc3_orisize_LAB = cv2.cvtColor(hair_gene_fusion_8uc3_orisize, cv2.COLOR_BGR2LAB) + user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) + mix_ratio = 0.2 + user_res_8uc3_orisize_LAB[:, :, 0] = hair_gene_fusion_8uc3_orisize_LAB[:, :, 0] * mix_ratio + user_res_8uc3_orisize_LAB[:, :, 0] * (1 - mix_ratio) + user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_LAB, cv2.COLOR_LAB2BGR) + + # show_concat = np.concatenate((hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize, + # user_baldseg_8uc3_orisize, user_res_8uc3_orisize), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_concat", show_concat) + # cv2.waitKey() + else: + user_res_8uc3_orisize = hair_gene_fusion_8uc3_orisize + + + # face_bbox = cv2.boundingRect(landmarks_origin_img_1k[np.newaxis, :, :]) + # face_max_len = max(face_bbox[2], face_bbox[3]) + # erode_kernel_size = int(face_max_len * 0.11) + # if erode_kernel_size % 2 == 0: + # erode_kernel_size += 1 + # blur_kernel_size = int(face_max_len * 0.05) + # if blur_kernel_size % 2 == 0: + # blur_kernel_size += 1 + # # print("blur_kernel_size: ", blur_kernels + # hair_gene_matte_8uc3_hard = hair_gene_matte_8uc3_orisize.copy() + # hair_gene_matte_8uc3_hard[hair_gene_matte_8uc3_orisize[:, :, 0] > 0] = 255 + # hair_gene_matte_8uc3_orisize_erode = cv2.erode(hair_gene_matte_8uc3_hard, np.ones((erode_kernel_size, erode_kernel_size), np.uint8), iterations=1) + # hair_gene_matte_fc32_orisize_erode = hair_gene_matte_8uc3_orisize_erode.astype(np.float32) / 255 + # hair_gene_matte_fc32_orisize = hair_gene_matte_8uc3_orisize.astype(np.float32) / 255 + # hair_gene_matte_fc32_orisize_circle = hair_gene_matte_fc32_orisize * (1 - hair_gene_matte_fc32_orisize_erode) + # hair_gene_matte_fc32_orisize_circle = cv2.GaussianBlur(hair_gene_matte_fc32_orisize_circle, (blur_kernel_size, blur_kernel_size), 0, 0) + # + # hair_gene_fusion_8uc3_orisize_LAB = cv2.cvtColor(hair_gene_fusion_8uc3_orisize, cv2.COLOR_BGR2LAB) + # user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) + # mix_ratio = 0.5 + # user_res_8uc3_orisize_LAB_new = user_res_8uc3_orisize_LAB.copy() + # user_res_8uc3_orisize_LAB_new[:, :, 0] = hair_gene_fusion_8uc3_orisize_LAB[:, :, 0] * mix_ratio + user_res_8uc3_orisize_LAB[:, :, 0] * (1 - mix_ratio) + # user_res_8uc3_orisize_LAB[:, :, 0] = user_res_8uc3_orisize_LAB_new[:, :, 0] * hair_gene_matte_fc32_orisize_circle[:, :, 0] + user_res_8uc3_orisize_LAB[:, :, 0] * (1 - hair_gene_matte_fc32_orisize_circle[:, :, 0]) + # user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_LAB, cv2.COLOR_LAB2BGR) + + # mid_show = np.concatenate((hair_gene_matte_fc32_orisize, hair_gene_matte_fc32_orisize_erode, hair_gene_matte_fc32_orisize_circle), axis=1) + # ratio = 1536. / max(mid_show.shape[:2]) + # mid_show = cv2.resize(mid_show, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("mid_show", mid_show) + # # cv2.imshow("user_res_8uc3_orisize", user_res_8uc3_orisize) + # cv2.waitKey() + + if use_enhance: + user_res_8uc3_orisize_enhance = self.face_enhance.process(user_res_8uc3_orisize, landmarks_origin_img_1k) + user_res_8uc3_orisize = user_bald_res_8uc3_orisize.astype(np.float32) / 255. * ( + 1. - hair_gene_matte_8uc3_orisize.astype( + np.float32) / 255.) + user_res_8uc3_orisize_enhance.astype( + np.float32) / 255. * hair_gene_matte_8uc3_orisize.astype(np.float32) / 255. + user_res_8uc3_orisize = (user_res_8uc3_orisize * 255).astype(np.uint8) + + # for pt in landmarks_origin_img_1k.astype(np.int32): + # cv2.circle(origin_img, (pt[0], pt[1]), 1, (0, 255, 0), -1) + # show_concat_orisize = np.concatenate((origin_img, user_res_8uc3_orisize), axis=1) + # ratio = 1024. / max(show_concat_orisize.shape[:2]) + # show_concat_orisize = cv2.resize(show_concat_orisize, (0, 0), fx=ratio, fy=ratio) + # ratio = 1024. / max(show_concat_768.shape[:2]) + # show_concat_768 = cv2.resize(show_concat_768, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("show_concat_orisize", show_concat_orisize) + # cv2.imshow("show_concat_768", show_concat_768) + # cv2.waitKey() + + if return_pt1k: + ret_dict = dict(user_res_8uc3_orisize=user_res_8uc3_orisize, + user_bald_res_8uc3_orisize=user_bald_res_8uc3_orisize, + landmarks_origin_img_1k=landmarks_origin_img_1k) + + return ret_dict, 0 + else: + return user_res_8uc3_orisize, 0 + + def infer_hairstyle_fix_face(self, origin_img, hairstyle_dir): + config_path = os.path.join(hairstyle_dir, "config.json") + if not os.path.exists(config_path): + return None, 10001 + config_info = json.load(open(config_path, "rb")) + gender = config_info["gender"] + ratio = int(config_info["ratio"]) + # print("gender: ", gender, " ratio: ", ratio) + + user_rgb_8uc3_orisize = origin_img + landmarks_origin_img_1k = self.get_landmark.forward(origin_img) + if landmarks_origin_img_1k is None: + return None, 10001 + + user_hair_matting_fg_orisize, user_hair_matting_8uc1_orisize = self.process_data.generator_matte.matte_inference(user_rgb_8uc3_orisize, landmarks_origin_img_1k) + user_hair_matting_fc32_orisize = user_hair_matting_8uc1_orisize[:, :, np.newaxis].astype(np.float32) / 255 + user_face_mask_fc32_orisize = self.face_seg.inference(origin_img, landmarks_origin_img_1k) + user_face_mask_fc32_orisize = np.repeat(user_face_mask_fc32_orisize, 3, axis=2) + user_face_mask_fc32_orisize = np.clip(user_face_mask_fc32_orisize, 0, 1) + + user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, \ + user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M,_ = self.process_data.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio) + + # show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("show_concat", show_concat) + # cv2.waitKey() + + another_pose_hair_img_path = os.path.join(hairstyle_dir, "input_another_pose_hair_image.npy") + if not os.path.exists(another_pose_hair_img_path): + return None, 10002 + another_pose_hair_image = np.load(another_pose_hair_img_path) + + # 换发型 + hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image, + user_baldseg_8uc3_768, + user_bald_8uc3_768, + user_landmark_f1k2_768, gender) + + hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste( + user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M) + + user_res_8uc3_orisize = self.hair_fusion.inference(hair_gene_fusion_8uc3_orisize, + hair_gene_matte_8uc3_orisize, + landmarks_origin_img_1k, + user_baldseg_8uc3_orisize, + ratio) + # for pt in landmarks_origin_img_1k.astype(np.int32): + # cv2.circle(origin_img, (pt[0], pt[1]), 1, (0, 255, 0), -1) + # show_concat_orisize = np.concatenate((origin_img, (user_face_mask_fc32_orisize*255).astype(np.uint8), hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize, user_baldseg_8uc3_orisize, user_res_8uc3_orisize), axis=1) + # ratio = 1536. / max(show_concat_orisize.shape[:2]) + # show_concat_orisize = cv2.resize(show_concat_orisize, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("show_concat_orisize", show_concat_orisize) + # cv2.waitKey() + + if self.use_enhance: + user_res_8uc3_orisize_enhance = self.face_enhance.process(user_res_8uc3_orisize, landmarks_origin_img_1k) + + # show_concat_orisize = np.concatenate((origin_img, user_res_8uc3_orisize, user_res_8uc3_orisize_enhance), axis=1) + # ratio = 1536. / max(show_concat_orisize.shape[:2]) + # show_concat_orisize = cv2.resize(show_concat_orisize, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("show_concat_orisize", show_concat_orisize) + # cv2.waitKey() + + user_res_8uc3_orisize = user_res_8uc3_orisize_enhance + + user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32) / 255 + hair_gene_matte_fc32_orisize = hair_gene_matte_8uc3_orisize.astype(np.float32) / 255 + user_rgb_fc32_orisize = user_rgb_8uc3_orisize.astype(np.float32) / 255 + user_hair_matting_fc32_orisize = cv2.dilate(user_hair_matting_fc32_orisize[:, :, 0], np.ones((35, 35), np.uint8))[:, :, np.newaxis] + user_face_mask_fc32_orisize = user_face_mask_fc32_orisize * (1 - hair_gene_matte_fc32_orisize) * (1 - user_hair_matting_fc32_orisize) + user_res_fc32_orisize = user_res_fc32_orisize * (1 - user_face_mask_fc32_orisize) + user_rgb_fc32_orisize * user_face_mask_fc32_orisize + + show_concat_orisize = np.concatenate((user_res_fc32_orisize, user_face_mask_fc32_orisize, user_rgb_fc32_orisize), axis=1) + ratio = 1536. / max(show_concat_orisize.shape[:2]) + show_concat_orisize = cv2.resize(show_concat_orisize, (0, 0), fx=ratio, fy=ratio) + cv2.imshow("show_concat_orisize", show_concat_orisize) + cv2.waitKey(10) + + return (user_res_fc32_orisize*255).astype(np.uint8), 0 + # return user_res_8uc3_orisize, 0 + + def infer_hairstyle_finetune(self, origin_img, hairstyle_dir): + config_path = os.path.join(hairstyle_dir, "config.json") + if not os.path.exists(config_path): + return None, 10001 + config_info = json.load(open(config_path, "rb")) + gender = config_info["gender"] + ratio = int(config_info["ratio"]) + # print("gender: ", gender, " ratio: ", ratio) + + user_rgb_8uc3_orisize = origin_img + landmarks_origin_img_1k = self.get_landmark.forward(origin_img) + if landmarks_origin_img_1k is None: + return None, 10001 + + user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, \ + user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M,_ = self.process_data.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio) + + # show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("show_concat", show_concat) + # cv2.waitKey() + + another_pose_hair_img_path = os.path.join(hairstyle_dir, "input_another_pose_hair_image.npy") + if not os.path.exists(another_pose_hair_img_path): + return None, 10002 + another_pose_hair_image = np.load(another_pose_hair_img_path) + + # 换发型 + hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image, + user_baldseg_8uc3_768, + user_bald_8uc3_768, + user_landmark_f1k2_768, gender) + + hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste( + user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M) + + # cv2.imshow("hair_gene_fusion_8uc3_orisize", hair_gene_fusion_8uc3_orisize) + # cv2.imshow("hair_gene_matte_8uc3_orisize", hair_gene_matte_8uc3_orisize) + # cv2.imshow("user_bald_res_8uc3_orisize", user_bald_res_8uc3_orisize) + # cv2.waitKey() + + user_res_8uc3_orisize = self.hair_fusion.inference(hair_gene_fusion_8uc3_orisize, + hair_gene_matte_8uc3_orisize, + landmarks_origin_img_1k, + user_baldseg_8uc3_orisize, + ratio) + # cv2.imshow("hair_gene_fusion_8uc3_orisize", hair_gene_fusion_8uc3_orisize) + # cv2.imshow("hair_gene_matte_8uc3_orisize", hair_gene_matte_8uc3_orisize) + # cv2.imshow("user_res_8uc3_orisize", user_res_8uc3_orisize) + # cv2.waitKey() + + if self.use_enhance: + user_res_8uc3_orisize = self.face_enhance.process(user_res_8uc3_orisize, landmarks_origin_img_1k) + + # for pt in landmarks_origin_img_1k.astype(np.int32): + # cv2.circle(origin_img, (pt[0], pt[1]), 1, (0, 255, 0), -1) + # show_concat_orisize = np.concatenate((origin_img, user_res_8uc3_orisize), axis=1) + # ratio = 1024. / max(show_concat_orisize.shape[:2]) + # show_concat_orisize = cv2.resize(show_concat_orisize, (0, 0), fx=ratio, fy=ratio) + # ratio = 1024. / max(show_concat_768.shape[:2]) + # show_concat_768 = cv2.resize(show_concat_768, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("show_concat_orisize", show_concat_orisize) + # cv2.imshow("show_concat_768", show_concat_768) + # cv2.waitKey() + + return user_res_8uc3_orisize, 0 + + def change_image_hair_color(self, image, color=[20, 20, 200]): + image = self.lookup(self.table_enlight, image) + b, g, r = color # [10, 50, 250] # [10, 250, 10] + tar_color = np.zeros_like(image) + tar_color[:, :, 0] = b + tar_color[:, :, 1] = g + tar_color[:, :, 2] = r + image_hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) + tar_hsv = cv2.cvtColor(tar_color, cv2.COLOR_BGR2HSV) + image_hsv[:, :, 0:2] = tar_hsv[:, :, 0:2] + changed = cv2.cvtColor(image_hsv.astype(np.uint8), cv2.COLOR_HSV2BGR) + # changed_sharpen = sharpen(changed) + return changed + + def lookup(self, luv, img): + table = luv + + b, g, r = img[:, :, 0] / 255.0, img[:, :, 1] / 255.0, img[:, :, 2] / 255.0 + b = b * 63 + + x1 = np.floor(np.floor(b) / 8.) + y1 = np.floor(b) - x1 * 8.0 + idx_x1 = x1 * 0.125 + 0.5 / 512.0 + (0.125 - 1.0 / 512.0) * g + idx_y1 = y1 * 0.125 + 0.5 / 512.0 + (0.125 - 1.0 / 512.0) * r + + x2 = np.floor(np.ceil(b) / 8.) + y2 = np.ceil(b) - x2 * 8.0 + idx_x2 = x2 * 0.125 + 0.5 / 512.0 + (0.125 - 1.0 / 512.0) * g + idx_y2 = y2 * 0.125 + 0.5 / 512.0 + (0.125 - 1.0 / 512.0) * r + + fract = np.mod(b * 63, 1.0)[:, :, np.newaxis] + + idx_x1 = (idx_x1 * 512).astype(np.int32).reshape(-1) + idx_y1 = (idx_y1 * 512).astype(np.int32).reshape(-1) + idx_x2 = (idx_x2 * 512).astype(np.int32).reshape(-1) + idx_y2 = (idx_y2 * 512).astype(np.int32).reshape(-1) + + color1 = table[idx_x1, idx_y1].reshape(img.shape[0], img.shape[1], 3) + color2 = table[idx_x2, idx_y2].reshape(img.shape[0], img.shape[1], 3) + final_img = np.clip(color1 * (1 - fract) + color2 * fract, 0, 255) + return final_img.astype(np.uint8) + def infer_hairtiaoran(self, origin_img, color, input_mask=None): + landmarks_origin_img_1k = self.get_landmark.forward(origin_img) + if landmarks_origin_img_1k is None: + return None, 10001 + hair_gene_matte_fg_8uc3_orisize, hair_gene_matte_8uc1_orisize = self.process_data.get_matte_img(origin_img, landmarks_origin_img_1k) + hair_matting_mask_8uc3_orisize = np.repeat(hair_gene_matte_8uc1_orisize[:, :, np.newaxis], 3, axis=2) + if input_mask is None: + input_mask = np.zeros_like(origin_img) + cv2.circle(input_mask, (int(landmarks_origin_img_1k[154, 0]), int(landmarks_origin_img_1k[154, 1])), int(origin_img.shape[0]*0.1), (255, 255, 255), -1) + origin_img_haircolor = self.change_image_hair_color(origin_img, color) + input_mask_blur = cv2.GaussianBlur(input_mask, (55, 55), 0, 0) + input_mask_blur_fc32 = input_mask_blur.astype(np.float32) / 255 + hair_matting_mask_fc32_orisize = hair_matting_mask_8uc3_orisize.astype(np.float32) / 255 + fusion_mask_fc32 = input_mask_blur_fc32 *hair_matting_mask_fc32_orisize + + origin_img_haircolor_fc32 = origin_img_haircolor.astype(np.float32) / 255 + origin_img_fc32 = origin_img.astype(np.float32) / 255 + + img_res_fc32 = origin_img_haircolor_fc32 * fusion_mask_fc32 + origin_img_fc32 * (1 - fusion_mask_fc32) + img_res = (img_res_fc32 * 255).astype(np.uint8) + + # for pt in landmarks_origin_img_1k.astype(np.int32): + # cv2.circle(origin_img, (pt[0], pt[1]), 2, (0, 255, 0), -1) + # show_concat = np.concatenate((origin_img, input_mask, hair_matting_mask_8uc3_orisize, img_res), axis=1) + # ratio = 1024. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("show_concat", show_concat) + # cv2.waitKey() + + return img_res, 0 + + 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 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 + def infer_hairstyle_diy(self, user_rgb_8uc3_orisize, ref_rgb_8uc3_orisize): + ref_landmark_1k2_f_orisize = self.get_landmark.forward(ref_rgb_8uc3_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 gender_res: + gender = "girl" + else: + gender = "boy" + 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) + another_pose_hair_image = self.Generator_reftensor(ref_rgb_8uc3_768, ref_matting_8uc3_768, + ref_baldseg_8uc3_768, + ref_landmark_f1k2_768) + + landmarks_origin_img_1k = self.get_landmark.forward(user_rgb_8uc3_orisize) + if landmarks_origin_img_1k is None: + return None, 10001 + + user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, \ + user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M,_ = self.process_data.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio) + + # show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1) + # ratio = 1536. / max(show_concat.shape[:2]) + # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("show_concat", show_concat) + # cv2.waitKey() + + # 换发型 + hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image, + user_baldseg_8uc3_768, + user_bald_8uc3_768, + user_landmark_f1k2_768, gender) + + hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste( + user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M) + + user_res_8uc3_orisize = self.hair_fusion.inference(hair_gene_fusion_8uc3_orisize, + hair_gene_matte_8uc3_orisize, + landmarks_origin_img_1k, + user_baldseg_8uc3_orisize, + ratio) + + if self.use_enhance: + user_res_8uc3_orisize = self.face_enhance.process(user_res_8uc3_orisize, landmarks_origin_img_1k) + + # for pt in landmarks_origin_img_1k.astype(np.int32): + # cv2.circle(origin_img, (pt[0], pt[1]), 1, (0, 255, 0), -1) + # show_concat_orisize = np.concatenate((origin_img, user_res_8uc3_orisize), axis=1) + # ratio = 1024. / max(show_concat_orisize.shape[:2]) + # show_concat_orisize = cv2.resize(show_concat_orisize, (0, 0), fx=ratio, fy=ratio) + # ratio = 1024. / max(show_concat_768.shape[:2]) + # show_concat_768 = cv2.resize(show_concat_768, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("show_concat_orisize", show_concat_orisize) + # cv2.imshow("show_concat_768", show_concat_768) + # cv2.waitKey() + + return user_res_8uc3_orisize, 0 + + def preprocess_haircolor(self, origin_img, origin_mask, ref_img, ref_mask): + ref_img_hsv = cv2.cvtColor(ref_img, cv2.COLOR_BGR2HSV) + loc_index = ref_mask[:, :, 0].nonzero() + color_val = ref_img_hsv[loc_index] + ref_img_mean_hsv = np.mean(color_val, axis=0) + + origin_img_hsv = cv2.cvtColor(origin_img, cv2.COLOR_BGR2HSV) + + loc_index = origin_mask[:, :, 0].nonzero() + color_val = origin_img_hsv[loc_index] + origin_img_mean_hsv = np.mean(color_val, axis=0) + + origin_img_hsv[:, :, :1] = ref_img_mean_hsv[:1] + # origin_img_hsv[:, :, 1:3] = ref_img_mean_hsv[1:3] + + res_img = cv2.cvtColor(origin_img_hsv.astype(np.uint8), cv2.COLOR_HSV2BGR) + origin_mask_fc32 = origin_mask.astype(np.float32) / 255 + res_img = (res_img.astype(np.float32) * origin_mask_fc32 + origin_img.astype(np.float32) * (1 - origin_mask_fc32)).astype(np.uint8) + + mid_show = np.concatenate((origin_img, origin_mask, res_img), axis=1) + ratio = 1536. / max(mid_show.shape[:2]) + mid_show = cv2.resize(mid_show, (0, 0), fx=ratio, fy=ratio) + cv2.imshow("mid_show", mid_show) + cv2.waitKey(0) + + return res_img + + def rein_hard_raw(self, user_res_8uc3_orisize, user_matting_8uc3_bald_orisize, + ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768): + ref_rgb_8uc3_change_color_768_hsv = cv2.cvtColor(ref_rgb_8uc3_change_color_768, cv2.COLOR_BGR2HSV) + loc_index = ref_matting_8uc3_change_color_768[:, :, 0].nonzero() + color_val = ref_rgb_8uc3_change_color_768_hsv[loc_index] + mean_ref_hair_hsv = np.mean(color_val, axis=0) + print("mean_ref_hair_hsv: ", mean_ref_hair_hsv) + user_res_8uc3_orisize_hsv = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2HSV) + loc_index = user_matting_8uc3_bald_orisize[:, :, 0].nonzero() + color_val = user_res_8uc3_orisize_hsv[loc_index] + mean_color_hair_hsv = np.mean(color_val, axis=0) + print("mean_color_hair_hsv: ", mean_color_hair_hsv) + user_res_8uc3_orisize_hsv[:, :, 1:3] = mean_ref_hair_hsv[1:3] + + img_res = cv2.cvtColor(user_res_8uc3_orisize_hsv, cv2.COLOR_HSV2BGR) + + user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) + img_res_LAB = cv2.cvtColor(img_res, cv2.COLOR_BGR2LAB) + img_res_LAB[:, :, 0] = user_res_8uc3_orisize_LAB[:, :, 0] + img_res = cv2.cvtColor(img_res_LAB, cv2.COLOR_LAB2BGR) + + return img_res + def rein_hard(self, user_res_8uc3_orisize, user_matting_8uc3_bald_orisize, + ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768): + ref_rgb_8uc3_change_color_768_hsv = cv2.cvtColor(ref_rgb_8uc3_change_color_768, cv2.COLOR_BGR2HSV) + loc_index = ref_matting_8uc3_change_color_768[:, :, 0].nonzero() + color_val = ref_rgb_8uc3_change_color_768_hsv[loc_index] + mean_ref_hair_hsv = np.mean(color_val, axis=0) + print("mean_ref_hair_hsv: ", mean_ref_hair_hsv) + user_res_8uc3_orisize_hsv = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2HSV) + loc_index = user_matting_8uc3_bald_orisize[:, :, 0].nonzero() + color_val = user_res_8uc3_orisize_hsv[loc_index] + mean_color_hair_hsv = np.mean(color_val, axis=0) + print("mean_color_hair_hsv: ", mean_color_hair_hsv) + # user_res_8uc3_orisize_hsv[:, :, 0:1] = user_res_8uc3_orisize_hsv[:, :, 0:1] + mean_ref_hair_hsv[0] - mean_color_hair_hsv[0] + user_res_8uc3_orisize_hsv[:, :, 0:1] = mean_ref_hair_hsv[0] + # user_res_8uc3_orisize_hsv[:, :, 1:2] = user_res_8uc3_orisize_hsv[:, :, 1:2] + (mean_ref_hair_hsv[1] - mean_color_hair_hsv[1]) * 0.1 + # user_res_8uc3_orisize_hsv[:, :, 0:1] = np.clip(user_res_8uc3_orisize_hsv[:, :, 0:1], 0, 180) + # user_res_8uc3_orisize_hsv[:, :, 1:2] = user_res_8uc3_orisize_hsv[:, :, 1:2] * mean_ref_hair_hsv[1] / mean_color_hair_hsv[1] + ratio_s = max(1.05, mean_ref_hair_hsv[1] * 0.85 / mean_color_hair_hsv[1]) + print("ratio_s: ", ratio_s) + user_res_8uc3_orisize_hsv[:, :, 1:2] = np.clip(user_res_8uc3_orisize_hsv[:, :, 1:2] * ratio_s, 0, 255) + + # user_res_8uc3_orisize_hsv[:, :, 2:3] = user_res_8uc3_orisize_hsv[:, :, 2:3] + (mean_ref_hair_hsv[2] - mean_color_hair_hsv[2]) * 0.4 + # user_res_8uc3_orisize_hsv[:, :, 2:3] = np.clip(user_res_8uc3_orisize_hsv[:, :, 2:3], 0, 255) + # user_res_8uc3_orisize_hsv[:, :, 2:3] = mean_ref_hair_hsv[2] + + ratio_v = 0.5 + user_res_8uc3_orisize_hsv[:, :, 2:3] = user_res_8uc3_orisize_hsv[:, :, 2:3] * ratio_v + mean_ref_hair_hsv[2] * (1 - ratio_v) + user_res_8uc3_orisize_hsv[:, :, 2:3] = np.clip(user_res_8uc3_orisize_hsv[:, :, 2:3], 0, 255) + + # ratio_v = mean_ref_hair_hsv[2] / mean_color_hair_hsv[2] + # ratio_v = ratio_v * 0.8 + 0.2 + # print("ratio_v: ", ratio_v) + # user_res_8uc3_orisize_hsv[:, :, 2:3] = np.clip(user_res_8uc3_orisize_hsv[:, :, 2:3] * ratio_v, 0, 255) + + img_res = cv2.cvtColor(user_res_8uc3_orisize_hsv, cv2.COLOR_HSV2BGR) + + user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) + img_res_LAB = cv2.cvtColor(img_res, cv2.COLOR_BGR2LAB) + # img_res_LAB[:, :, 0] = user_res_8uc3_orisize_LAB[:, :, 0] + mix_ratio = 0.5 + img_res_LAB[:, :, 0] = user_res_8uc3_orisize_LAB[:, :, 0] * mix_ratio + img_res_LAB[:, :, 0] * (1 - mix_ratio) + img_res = cv2.cvtColor(img_res_LAB, cv2.COLOR_LAB2BGR) + + # middle_show = np.concatenate((user_res_8uc3_orisize, img_res), axis=1) + # ratio = 1536. / max(middle_show.shape[:2]) + # middle_show = cv2.resize(middle_show, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow('middle_show', middle_show) + # cv2.waitKey(0) + + return img_res + + + def infer_haircolor(self, user_rgb_8uc3_orisize, haircolor_dir, return_matting=False): + landmarks_origin_img_1k = self.get_landmark.forward(user_rgb_8uc3_orisize) + if landmarks_origin_img_1k is None: + return None, 10001 + + _, user_matting_8uc1_bald_orisize = self.process_data.generator_matte.matte_inference(user_rgb_8uc3_orisize, + landmarks_origin_img_1k) + user_matting_8uc3_bald_orisize = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) + + ref_rgb_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_rgb_8uc3_color_768.png")) + ref_matting_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_matting_8uc3_color_768.png")) + + # user_rgb_8uc3_orisize = self.preprocess_haircolor(user_rgb_8uc3_orisize, user_matting_8uc3_bald_orisize, + # ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768) + + user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, user_hair_color_M = \ + self.process_data.get_prepare_hair_color_user_data(user_rgb_8uc3_orisize, + landmarks_origin_img_1k) + + # cv2.imshow("user_matting_8uc3_change_color_768", user_matting_8uc3_change_color_768) + # cv2.waitKey() + start_time = time.time() + hair_gene_color_8uc3_768, user_matting_mask_8uc3_768 = self.change_haircolor.Change_Hair_inference( + user_rgb_8uc3_change_color_768, + user_matting_8uc3_change_color_768, + ref_rgb_8uc3_change_color_768, + ref_matting_8uc3_change_color_768) + # print("Change_Hair_inference cost: ", time.time()-start_time) + # cv2.imshow("hair_gene_color_8uc3_768: ", hair_gene_color_8uc3_768) + # cv2.imshow("user_matting_mask_8uc3_768: ", user_matting_mask_8uc3_768) + # cv2.waitKey() + + user_matting_8uc3_orisize = cv2.warpAffine(user_matting_mask_8uc3_768, + cv2.invertAffineTransform(user_hair_color_M), + (user_rgb_8uc3_orisize.shape[1], + user_rgb_8uc3_orisize.shape[0]), flags=cv2.INTER_CUBIC) + + hair_gene_color_8uc3_orisize = cv2.warpAffine(hair_gene_color_8uc3_768, + cv2.invertAffineTransform(user_hair_color_M), + (user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), + dst=user_rgb_8uc3_orisize.copy(), + flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_TRANSPARENT) + + # ref_rgb_8uc3_change_color_orisize = cv2.resize(ref_rgb_8uc3_change_color_768, + # (user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), + # interpolation=cv2.INTER_CUBIC) + + # user_res_8uc3_orisize = user_rgb_8uc3_orisize * ( + # 1 - user_matting_8uc3_change_color_orisize / 255) + hair_gene_color_8uc3_orisize * ( + # user_matting_8uc3_change_color_orisize / 255) + # user_res_8uc3_orisize = (np.clip(user_res_8uc3_orisize, 0, 255)).astype(np.uint8) + + # cv2.imshow("user_rgb_8uc3_orisize", user_rgb_8uc3_orisize) + # cv2.imshow("hair_gene_color_8uc3_orisize", hair_gene_color_8uc3_orisize) + # cv2.imshow("user_matting_8uc3_change_color_orisize", user_matting_8uc3_change_color_orisize) + # cv2.waitKey() + + # if self.use_enhance: + user_res_8uc3_orisize = self.face_enhance.process(hair_gene_color_8uc3_orisize, landmarks_origin_img_1k) + + user_res_8uc3_orisize_hsv = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2HSV) + user_res_8uc3_orisize_hsv[:, :, 1:2] = np.clip(user_res_8uc3_orisize_hsv[:, :, 1:2] * 1.15, 0, 255) + user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_hsv, cv2.COLOR_HSV2BGR) + + user_res_8uc3_orisize2 = user_rgb_8uc3_orisize * ( + 1 - user_matting_8uc3_orisize / 255) + user_res_8uc3_orisize * ( + user_matting_8uc3_orisize / 255) + user_res_8uc3_orisize2 = (np.clip(user_res_8uc3_orisize2, 0, 255)).astype(np.uint8) + + # user_res_8uc3_orisize_hsv = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2HSV) + # user_res_8uc3_orisize_hsv[:, :, 1:2] = np.clip(user_res_8uc3_orisize_hsv[:, :, 1:2] * 1.1, 0, 255) + # user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_hsv, cv2.COLOR_HSV2BGR) + + # user_res_8uc3_orisize = self.rein_hard_raw(user_res_8uc3_orisize, user_matting_8uc3_bald_orisize, + # ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768) + # user_res_8uc3_orisize = self.rein_hard(user_res_8uc3_orisize, user_matting_8uc3_bald_orisize, + # ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768) + + # user_res_8uc3_orisize = self.rein_hard(user_res_8uc3_orisize, user_matting_8uc3_bald_orisize, + # ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768) + + # user_res_8uc3_orisize = user_rgb_8uc3_orisize * ( + # 1 - user_matting_8uc3_bald_orisize / 255) + user_res_8uc3_orisize * ( + # user_matting_8uc3_bald_orisize / 255) + # user_res_8uc3_orisize = (np.clip(user_res_8uc3_orisize, 0, 255)).astype(np.uint8) + + # middle_show = np.concatenate( + # (user_rgb_8uc3_orisize, user_matting_8uc3_bald_orisize, user_res_8uc3_orisize), axis=1) + # ratio = 1536. / max(middle_show.shape[:2]) + # middle_show = cv2.resize(middle_show, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow('middle_show', middle_show) + # cv2.waitKey(0) + + # user_res_8uc3_orisize = np.concatenate((user_res_8uc3_orisize, user_matting_8uc3_bald_orisize), axis=1) + if return_matting: + return user_res_8uc3_orisize2, user_matting_8uc1_bald_orisize, 0 + else: + return user_res_8uc3_orisize2, 0 + + + def get_hair_color(self, input_img, input_mask): + # 确保图像和掩膜具有相同的尺寸 + if input_img.shape[:2] != input_mask.shape[:2]: + return None # 如果没有找到头发区域 + + # 创建掩膜，提取头发区域 + hair_region = cv2.bitwise_and(input_img, input_img, mask=input_mask) + # 计算头发区域的颜色（去除黑色区域） + hair_pixels = hair_region[hair_region != 0] + + if len(hair_pixels) == 0: + return None # 如果没有找到头发区域 + + # 计算平均颜色 + average_color = np.mean(hair_pixels, axis=0) + is_black = (average_color[2] * 0.299 + average_color[1] * 0.587 + average_color[1] * 0.114) < 60 + return is_black + + + def infer_haircolor_tj(self, user_rgb_8uc3_orisize, haircolor_dir): + landmarks_origin_img_1k= self.get_landmark.forward(user_rgb_8uc3_orisize) + if landmarks_origin_img_1k is None: + return None, 10001 + t0 = time.time() + _, user_matting_8uc1_bald_orisize = self.process_data_infer.generator_matte.matte_inference(user_rgb_8uc3_orisize, + landmarks_origin_img_1k) + user_matting_mask_8uc3_orisize = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) + + ref_rgb_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_rgb_8uc3_color_768.png")) + ref_matting_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_matting_8uc3_color_768.png")) + + user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, user_hair_color_M = \ + self.process_data_infer.get_prepare_hair_color_user_data(user_rgb_8uc3_orisize,landmarks_origin_img_1k) + + # start_time = time.time() + hair_gene_color_8uc3_768, user_matting_mask_8uc3_768 = self.change_haircolor.Change_Hair_inference( + user_rgb_8uc3_change_color_768, + user_matting_8uc3_change_color_768, + ref_rgb_8uc3_change_color_768, + ref_matting_8uc3_change_color_768) + + hair_gene_color_8uc3_orisize = cv2.warpAffine(hair_gene_color_8uc3_768, + cv2.invertAffineTransform(user_hair_color_M), + (user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), + dst=user_rgb_8uc3_orisize.copy(), + flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_TRANSPARENT) + + user_res_8uc3_orisize = self.face_enhance.process(hair_gene_color_8uc3_orisize, landmarks_origin_img_1k) + user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32) / 255 + # user_res_fc32_orisize_hsv = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2HSV) + ref_rgb_fc32_change_color_768 = ref_rgb_8uc3_change_color_768.astype(np.float32) / 255 + ref_matting_fc32_change_color_768 = ref_matting_8uc3_change_color_768.astype(np.float32) / 255 + + ref_rgb_avg_std = self.getavgstd(ref_rgb_fc32_change_color_768, ref_matting_fc32_change_color_768) + user_matting_mask_fc32_orisize = user_matting_mask_8uc3_orisize.astype(np.float32) / 255 + user_res_fc32_orisize, src_avg_std = self.reinhard_rgb(user_res_fc32_orisize, user_matting_mask_fc32_orisize, + ref_rgb_avg_std, ratio=1.0) + + user_res_8uc3_orisize_reinhard = (user_res_fc32_orisize * 255).astype(np.uint8) + + user_res_8uc3_orisize_reinhard_LAB = cv2.cvtColor(user_res_8uc3_orisize_reinhard, cv2.COLOR_BGR2LAB) + user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) + mix_ratio = 0.1 + user_res_8uc3_orisize_reinhard_LAB[:, :, 0] = user_res_8uc3_orisize_LAB[:, :, + 0] * mix_ratio + user_res_8uc3_orisize_reinhard_LAB[:, :, 0] * ( + 1 - mix_ratio) + user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_reinhard_LAB, cv2.COLOR_LAB2BGR) + + user_res_8uc3_orisize2 = user_rgb_8uc3_orisize * ( + 1 - user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize * ( + user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize2 = (np.clip(user_res_8uc3_orisize2, 0, 255)).astype(np.uint8) + # print('change color, last proces,', time.time() - t3) + return user_res_8uc3_orisize2, 0 + + def infer_haircolor_new(self, user_rgb_8uc3_orisize, haircolor_dir, target_hair_color, return_matting=False): + landmarks_origin_img_1k= self.get_landmark.forward(user_rgb_8uc3_orisize) + if landmarks_origin_img_1k is None: + return None, 10001 + need_process = (target_hair_color[0] * 0.299 + target_hair_color[1] * 0.587 + target_hair_color[2] * 0.114) > 150 + _, user_matting_8uc1_bald_orisize = self.process_data_infer.generator_matte.matte_inference(user_rgb_8uc3_orisize, + landmarks_origin_img_1k) + + user_matting_mask_8uc3_orisize = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) + + ref_rgb_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_rgb_8uc3_color_768.png")) + ref_matting_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_matting_8uc3_color_768.png")) + + user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, user_hair_color_M = \ + self.process_data_infer.get_prepare_hair_color_user_data(user_rgb_8uc3_orisize, + landmarks_origin_img_1k) + + # cv2.imshow("user_matting_8uc3_change_color_768", user_matting_8uc3_change_color_768) + # cv2.waitKey() + start_time = time.time() + need_process=True + if need_process: + haircolor_dir_tmp = os.path.join(config.get('default', "haircolorDir"), config.get('default', "baseColor_ID")) + face_base, status1 = self.infer_haircolor_tj(user_rgb_8uc3_orisize, haircolor_dir_tmp) + # cv2.imwrite('/home/student/Desktop/tmp_color/need/face_base.png', face_base) + if status1 == 0: + return face_base, user_matting_8uc1_bald_orisize, 0 + else: + need_process = False + if not need_process: + hair_gene_color_8uc3_768, user_matting_mask_8uc3_768 = self.change_haircolor.Change_Hair_inference( + user_rgb_8uc3_change_color_768, + user_matting_8uc3_change_color_768, + ref_rgb_8uc3_change_color_768, + ref_matting_8uc3_change_color_768) + # cv2.imwrite("hair_gene_color_8uc3_768.png", hair_gene_color_8uc3_768) + # print("Change_Hair_inference cost: ", time.time()-start_time) + # cv2.imshow("hair_gene_color_8uc3_768: ", hair_gene_color_8uc3_768) + # cv2.imshow("user_matting_mask_8uc3_768: ", user_matting_mask_8uc3_768) + # cv2.waitKey() + + # user_matting_mask_8uc3_orisize = cv2.warpAffine(user_matting_mask_8uc3_768, + # cv2.invertAffineTransform(user_hair_color_M), + # (user_rgb_8uc3_orisize.shape[1], + # user_rgb_8uc3_orisize.shape[0]), flags=cv2.INTER_CUBIC) + + hair_gene_color_8uc3_orisize = cv2.warpAffine(hair_gene_color_8uc3_768, + cv2.invertAffineTransform(user_hair_color_M), + (user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), + dst=user_rgb_8uc3_orisize.copy(), + flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_TRANSPARENT) + + user_res_8uc3_orisize = self.face_enhance.process(hair_gene_color_8uc3_orisize, landmarks_origin_img_1k) + # cv2.imshow("user_res_8uc3_orisize", user_res_8uc3_orisize) + # cv2.waitKey() + + user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32) / 255 + # user_res_fc32_orisize_hsv = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2HSV) + ref_rgb_fc32_change_color_768 = ref_rgb_8uc3_change_color_768.astype(np.float32) / 255 + ref_matting_fc32_change_color_768 = ref_matting_8uc3_change_color_768.astype(np.float32) / 255 + + ref_rgb_avg_std = self.getavgstd(ref_rgb_fc32_change_color_768, ref_matting_fc32_change_color_768) + user_matting_mask_fc32_orisize = user_matting_mask_8uc3_orisize.astype(np.float32) / 255 + user_res_fc32_orisize, src_avg_std = self.reinhard_rgb(user_res_fc32_orisize, user_matting_mask_fc32_orisize, + ref_rgb_avg_std, ratio=1.0) + + user_res_8uc3_orisize_reinhard = (user_res_fc32_orisize * 255).astype(np.uint8) + + user_res_8uc3_orisize_reinhard_LAB = cv2.cvtColor(user_res_8uc3_orisize_reinhard, cv2.COLOR_BGR2LAB) + user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) + mix_ratio = 0.1 + user_res_8uc3_orisize_reinhard_LAB[:, :, 0] = user_res_8uc3_orisize_LAB[:, :, + 0] * mix_ratio + user_res_8uc3_orisize_reinhard_LAB[:, :, 0] * ( + 1 - mix_ratio) + user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_reinhard_LAB, cv2.COLOR_LAB2BGR) + + user_res_8uc3_orisize2 = user_rgb_8uc3_orisize * ( + 1 - user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize * ( + user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize2 = (np.clip(user_res_8uc3_orisize2, 0, 255)).astype(np.uint8) + # cv2.imwrite('/home/student/Desktop/tmp_color/need/face_base2.png', user_res_8uc3_orisize2) # 调试用，路径不存在会导致崩溃 + + if return_matting: + return user_res_8uc3_orisize2, user_matting_8uc1_bald_orisize, 0 + else: + return user_res_8uc3_orisize2, 0 + + + def infer_haircolor_0311(self, user_rgb_8uc3_orisize, haircolor_dir): + landmarks_origin_img_1k = self.get_landmark.forward(user_rgb_8uc3_orisize) + if landmarks_origin_img_1k is None: + return None, 10001 + + _, user_matting_8uc1_bald_orisize = self.process_data.generator_matte.matte_inference(user_rgb_8uc3_orisize, + landmarks_origin_img_1k) + user_matting_mask_8uc3_orisize = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) + + ref_rgb_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_rgb_8uc3_color_768.png")) + ref_matting_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_matting_8uc3_color_768.png")) + + # user_rgb_8uc3_orisize = self.preprocess_haircolor(user_rgb_8uc3_orisize, user_matting_8uc3_bald_orisize, + # ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768) + + user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, user_hair_color_M = \ + self.process_data.get_prepare_hair_color_user_data(user_rgb_8uc3_orisize, + landmarks_origin_img_1k) + + # cv2.imshow("user_matting_8uc3_change_color_768", user_matting_8uc3_change_color_768) + # cv2.waitKey() + # start_time = time.time() + hair_gene_color_8uc3_768, user_matting_mask_8uc3_768 = self.change_haircolor.Change_Hair_inference( + user_rgb_8uc3_change_color_768, + user_matting_8uc3_change_color_768, + ref_rgb_8uc3_change_color_768, + ref_matting_8uc3_change_color_768) + # print("Change_Hair_inference cost: ", time.time()-start_time) + # cv2.imshow("hair_gene_color_8uc3_768: ", hair_gene_color_8uc3_768) + # cv2.imshow("user_matting_mask_8uc3_768: ", user_matting_mask_8uc3_768) + # cv2.waitKey() + + # user_matting_8uc3_orisize = cv2.warpAffine(user_matting_mask_8uc3_768, + # cv2.invertAffineTransform(user_hair_color_M), + # (user_rgb_8uc3_orisize.shape[1], + # user_rgb_8uc3_orisize.shape[0]), flags=cv2.INTER_CUBIC) + + hair_gene_color_8uc3_orisize = cv2.warpAffine(hair_gene_color_8uc3_768, + cv2.invertAffineTransform(user_hair_color_M), + (user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), + dst=user_rgb_8uc3_orisize.copy(), + flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_TRANSPARENT) + + # ref_rgb_8uc3_change_color_orisize = cv2.resize(ref_rgb_8uc3_change_color_768, + # (user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), + # interpolation=cv2.INTER_CUBIC) + + # user_res_8uc3_orisize = user_rgb_8uc3_orisize * ( + # 1 - user_matting_8uc3_change_color_orisize / 255) + hair_gene_color_8uc3_orisize * ( + # user_matting_8uc3_change_color_orisize / 255) + # user_res_8uc3_orisize = (np.clip(user_res_8uc3_orisize, 0, 255)).astype(np.uint8) + + # cv2.imshow("user_rgb_8uc3_orisize", user_rgb_8uc3_orisize) + # cv2.imshow("hair_gene_color_8uc3_orisize", hair_gene_color_8uc3_orisize) + # cv2.imshow("user_matting_8uc3_change_color_orisize", user_matting_8uc3_change_color_orisize) + # cv2.waitKey() + + # if self.use_enhance: + user_res_8uc3_orisize = self.face_enhance.process(hair_gene_color_8uc3_orisize, landmarks_origin_img_1k) + + # user_res_8uc3_orisize_hsv = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2HSV) + # user_res_8uc3_orisize_hsv[:, :, 1:2] = np.clip(user_res_8uc3_orisize_hsv[:, :, 1:2] * 1.15, 0, 255) + # # user_res_8uc3_orisize_hsv[:, :, 2:3] = np.clip(user_res_8uc3_orisize_hsv[:, :, 2:3], 100, 255) + # user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_hsv, cv2.COLOR_HSV2BGR) + + user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32) / 255 + user_res_fc32_orisize_hsv = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2HSV) + ref_rgb_fc32_change_color_768 = ref_rgb_8uc3_change_color_768.astype(np.float32) / 255 + ref_rgb_fc32_change_color_768_hsv = cv2.cvtColor(ref_rgb_fc32_change_color_768, cv2.COLOR_BGR2HSV) + loc_index = ref_matting_8uc3_change_color_768[:, :, 0].nonzero() + color_val = ref_rgb_fc32_change_color_768_hsv[loc_index] + mean_ref_haircolor_hsv = np.mean(color_val, axis=0) + print("mean_ref_haircolor_hsv: ", mean_ref_haircolor_hsv) + loc_index = user_matting_mask_8uc3_orisize[:, :, 0].nonzero() + color_val = user_res_fc32_orisize_hsv[loc_index] + mean_user_haircolor_hsv = np.mean(color_val, axis=0) + print("mean_user_haircolor_hsv: ", mean_user_haircolor_hsv) + user_res_fc32_orisize_hsv[:, :, 0:1] = mean_ref_haircolor_hsv[0] + # ratio_h = mean_ref_haircolor_hsv[0] / mean_user_haircolor_hsv[0] + # user_res_fc32_orisize_hsv[:, :, 0:1] = np.clip(user_res_fc32_orisize_hsv[:, :, 0:1] + mean_ref_haircolor_hsv[0] - mean_user_haircolor_hsv[0], 0, 360.0) + user_res_fc32_orisize_hsv[:, :, 1:2] = np.clip(user_res_fc32_orisize_hsv[:, :, 1:2] * 1.15, 0, 1.0) + user_res_fc32_orisize_hsv[:, :, 2:3] = np.clip(user_res_fc32_orisize_hsv[:, :, 2:3], 0, 1.0) + user_res_fc32_orisize = cv2.cvtColor(user_res_fc32_orisize_hsv, cv2.COLOR_HSV2BGR) + user_res_8uc3_orisize_new = (user_res_fc32_orisize * 255).astype(np.uint8) + + user_res_8uc3_orisize_new_LAB = cv2.cvtColor(user_res_8uc3_orisize_new, cv2.COLOR_BGR2LAB) + user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) + user_res_8uc3_orisize_new_LAB[:, :, 0] = user_res_8uc3_orisize_LAB[:, :, 0] + user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_new_LAB, cv2.COLOR_LAB2BGR) + + user_res_8uc3_orisize2 = user_rgb_8uc3_orisize * (1 - user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize * (user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize2 = (np.clip(user_res_8uc3_orisize2, 0, 255)).astype(np.uint8) + + # middle_show = np.concatenate( + # (user_rgb_8uc3_orisize, user_matting_mask_8uc3_orisize, user_res_8uc3_orisize, user_res_8uc3_orisize2), axis=1) + # ratio = 1536. / max(middle_show.shape[:2]) + # middle_show = cv2.resize(middle_show, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow('middle_show', middle_show) + # cv2.waitKey(0) + + # user_res_8uc3_orisize_hsv = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2HSV) + # user_res_8uc3_orisize_hsv[:, :, 1:2] = np.clip(user_res_8uc3_orisize_hsv[:, :, 1:2] * 1.1, 0, 255) + # user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_hsv, cv2.COLOR_HSV2BGR) + + # user_res_8uc3_orisize = self.rein_hard_raw(user_res_8uc3_orisize, user_matting_8uc3_bald_orisize, + # ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768) + # user_res_8uc3_orisize = self.rein_hard(user_res_8uc3_orisize, user_matting_8uc3_bald_orisize, + # ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768) + + # user_res_8uc3_orisize = self.rein_hard(user_res_8uc3_orisize, user_matting_8uc3_bald_orisize, + # ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768) + + # user_res_8uc3_orisize = user_rgb_8uc3_orisize * ( + # 1 - user_matting_8uc3_bald_orisize / 255) + user_res_8uc3_orisize * ( + # user_matting_8uc3_bald_orisize / 255) + # user_res_8uc3_orisize = (np.clip(user_res_8uc3_orisize, 0, 255)).astype(np.uint8) + + # middle_show = np.concatenate( + # (user_rgb_8uc3_orisize, user_matting_8uc3_bald_orisize, user_res_8uc3_orisize), axis=1) + # ratio = 1536. / max(middle_show.shape[:2]) + # middle_show = cv2.resize(middle_show, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow('middle_show', middle_show) + # cv2.waitKey(0) + + # user_res_8uc3_orisize = np.concatenate((user_res_8uc3_orisize, user_matting_8uc3_bald_orisize), axis=1) + + return user_res_8uc3_orisize2, 0 + + def getavgstd(self, image, mask): + if mask.shape[2] == 1: + mask = np.repeat(mask, 3, axis=2) + mask_index = np.flatnonzero((mask > 0.1).any(axis=2)) + lab_layer = image.reshape(-1, 3)[mask_index] + if len(lab_layer) < 100: return None + lab_layer = np.float32(lab_layer) + avg = [] + std = [] + image_avg_l = np.mean(lab_layer[:, 0]) + image_std_l = np.std(lab_layer[:, 0]) + image_avg_a = np.mean(lab_layer[:, 1]) + image_std_a = np.std(lab_layer[:, 1]) + image_avg_b = np.mean(lab_layer[:, 2]) + image_std_b = np.std(lab_layer[:, 2]) + avg.append(image_avg_l) + avg.append(image_avg_a) + avg.append(image_avg_b) + std.append(image_std_l) + std.append(image_std_a) + std.append(image_std_b) + return avg, std + + def reinhard(self, origin_img, mask, tar_avg_std, ratio=0.5): + origin_img_lab = cv2.cvtColor(origin_img, cv2.COLOR_BGR2LAB) + src_avg_std = self.getavgstd(origin_img_lab, mask) + src_avg_std = np.float32(src_avg_std) + + # origin_img_lab[:, :, 0] = origin_img_lab[:, :, 0] + (tar_avg_std[0][0] - src_avg_std[0][0]) * ratio + origin_img_lab[:, :, 1] = origin_img_lab[:, :, 1] + (tar_avg_std[0][1] - src_avg_std[0][1]) * ratio + origin_img_lab[:, :, 2] = origin_img_lab[:, :, 2] + (tar_avg_std[0][2] - src_avg_std[0][2]) * ratio + img_ret = cv2.cvtColor(origin_img_lab, cv2.COLOR_LAB2BGR) + img_ret = np.clip(img_ret, 0.01, 0.99) + return img_ret, src_avg_std + def reinhard_rgb_old(self, origin_img, mask, tar_avg_std, ratio=0.5): + # origin_img_lab = cv2.cvtColor(origin_img, cv2.COLOR_BGR2LAB) + origin_img_lab = origin_img + src_avg_std = self.getavgstd(origin_img_lab, mask) + src_avg_std = np.float32(src_avg_std) + + origin_img_lab[:, :, 0] = origin_img_lab[:, :, 0] + (tar_avg_std[0][0] - src_avg_std[0][0]) * (tar_avg_std[1][0] / src_avg_std[1][0]) + origin_img_lab[:, :, 1] = origin_img_lab[:, :, 1] + (tar_avg_std[0][1] - src_avg_std[0][1]) * (tar_avg_std[1][1] / src_avg_std[1][1]) + origin_img_lab[:, :, 2] = origin_img_lab[:, :, 2] + (tar_avg_std[0][2] - src_avg_std[0][2]) * (tar_avg_std[1][2] / src_avg_std[1][2]) + # img_ret = cv2.cvtColor(origin_img_lab, cv2.COLOR_LAB2BGR) + img_ret = origin_img_lab + img_ret = np.clip(img_ret, 0.0, 1.0) + return img_ret, src_avg_std + def reinhard_rgb(self, origin_img, mask, tar_avg_std, ratio=0.5): + # origin_img_lab = cv2.cvtColor(origin_img, cv2.COLOR_BGR2LAB) + origin_img_lab = origin_img + src_avg_std = self.getavgstd(origin_img_lab, mask) + src_avg_std = np.float32(src_avg_std) + + origin_img_lab[:, :, 0] = origin_img_lab[:, :, 0] + (tar_avg_std[0][0] - src_avg_std[0][0]) * ratio + origin_img_lab[:, :, 1] = origin_img_lab[:, :, 1] + (tar_avg_std[0][1] - src_avg_std[0][1]) * ratio + origin_img_lab[:, :, 2] = origin_img_lab[:, :, 2] + (tar_avg_std[0][2] - src_avg_std[0][2]) * ratio + # img_ret = cv2.cvtColor(origin_img_lab, cv2.COLOR_LAB2BGR) + img_ret = origin_img_lab + img_ret = np.clip(img_ret, 0.0, 1.0) + return img_ret, src_avg_std + def infer_haircolor_0313(self, user_rgb_8uc3_orisize, haircolor_dir): + landmarks_origin_img_1k = self.get_landmark.forward(user_rgb_8uc3_orisize) + if landmarks_origin_img_1k is None: + return None, 10001 + + _, user_matting_8uc1_bald_orisize = self.process_data.generator_matte.matte_inference(user_rgb_8uc3_orisize, + landmarks_origin_img_1k) + user_matting_mask_8uc3_orisize = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) + + ref_rgb_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_rgb_8uc3_color_768.png")) + ref_matting_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_matting_8uc3_color_768.png")) + + user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, user_hair_color_M = \ + self.process_data.get_prepare_hair_color_user_data(user_rgb_8uc3_orisize, + landmarks_origin_img_1k) + + # start_time = time.time() + hair_gene_color_8uc3_768, user_matting_mask_8uc3_768 = self.change_haircolor.Change_Hair_inference( + user_rgb_8uc3_change_color_768, + user_matting_8uc3_change_color_768, + ref_rgb_8uc3_change_color_768, + ref_matting_8uc3_change_color_768) + + hair_gene_color_8uc3_orisize = cv2.warpAffine(hair_gene_color_8uc3_768, + cv2.invertAffineTransform(user_hair_color_M), + (user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), + dst=user_rgb_8uc3_orisize.copy(), + flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_TRANSPARENT) + + # if self.use_enhance: + # user_res_8uc3_orisize = self.face_enhance.process(hair_gene_color_8uc3_orisize, landmarks_origin_img_1k) + + user_res_8uc3_orisize = hair_gene_color_8uc3_orisize.copy() + + # user_rgb_8uc3_orisize_LAB = cv2.cvtColor(user_rgb_8uc3_orisize, cv2.COLOR_BGR2LAB) + # user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) + # user_res_8uc3_orisize_LAB[:, :, 0] = user_rgb_8uc3_orisize_LAB[:, :, 0] + # user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_LAB, cv2.COLOR_LAB2BGR) + + # middle_show = np.concatenate((user_rgb_8uc3_orisize, hair_gene_color_8uc3_orisize, user_res_8uc3_orisize, user_res_8uc3_orisize_new), axis=1) + # ratio = 1536. / max(middle_show.shape[:2]) + # middle_show = cv2.resize(middle_show, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow('middle_show', middle_show) + # cv2.waitKey(0) + + # #### version 1 + # user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32) / 255 + # user_res_fc32_orisize_hsv = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2HSV) + # ref_rgb_fc32_change_color_768 = ref_rgb_8uc3_change_color_768.astype(np.float32) / 255 + # ref_rgb_fc32_change_color_768_hsv = cv2.cvtColor(ref_rgb_fc32_change_color_768, cv2.COLOR_BGR2HSV) + # + # ref_hair_hsv_mean_std = self.getavgstd(ref_rgb_fc32_change_color_768_hsv, ref_matting_8uc3_change_color_768) + # user_hair_hsv_mean_std = self.getavgstd(user_res_fc32_orisize_hsv, user_matting_mask_8uc3_orisize) + # print("ref_hair_hsv_mean_std: ", ref_hair_hsv_mean_std) + # print("user_hair_hsv_mean_std: ", user_hair_hsv_mean_std) + # + # loc_index = ref_matting_8uc3_change_color_768[:, :, 0].nonzero() + # color_val = ref_rgb_fc32_change_color_768_hsv[loc_index] + # mean_ref_haircolor_hsv = np.mean(color_val, axis=0) + # print("mean_ref_haircolor_hsv: ", mean_ref_haircolor_hsv) + # loc_index = user_matting_mask_8uc3_orisize[:, :, 0].nonzero() + # color_val = user_res_fc32_orisize_hsv[loc_index] + # mean_user_haircolor_hsv = np.mean(color_val, axis=0) + # print("mean_user_haircolor_hsv: ", mean_user_haircolor_hsv) + # + # # user_res_fc32_orisize_hsv[:, :, 0:1] = user_res_fc32_orisize_hsv[:, :, 0:1] + mean_ref_haircolor_hsv[0] - mean_user_haircolor_hsv[0] + # user_res_fc32_orisize_hsv[:, :, 0:1] = user_res_fc32_orisize_hsv[:, :, 0:1] + ref_hair_hsv_mean_std[0][0] - user_hair_hsv_mean_std[0][0] + # # user_res_fc32_orisize_hsv[:, :, 0:1] = np.clip(user_res_fc32_orisize_hsv[:, :, 0:1], 0, 360) + # # user_res_fc32_orisize_hsv[:, :, 1:2] = user_res_fc32_orisize_hsv[:, :, 1:2] + mean_ref_haircolor_hsv[1] - mean_user_haircolor_hsv[1] + # user_res_fc32_orisize_hsv[:, :, 1:2] = user_res_fc32_orisize_hsv[:, :, 1:2] + ref_hair_hsv_mean_std[0][1] - user_hair_hsv_mean_std[0][1] + # user_res_fc32_orisize_hsv[:, :, 1:2] = np.clip(user_res_fc32_orisize_hsv[:, :, 1:2], 0, 1.0) + # # user_res_fc32_orisize_hsv[:, :, 2:3] = user_res_fc32_orisize_hsv[:, :, 2:3] + mean_ref_haircolor_hsv[2] - mean_user_haircolor_hsv[2] + # user_res_fc32_orisize_hsv[:, :, 2:3] = user_res_fc32_orisize_hsv[:, :, 2:3] + ref_hair_hsv_mean_std[0][2] - user_hair_hsv_mean_std[0][2] + # user_res_fc32_orisize_hsv[:, :, 2:3] = np.clip(user_res_fc32_orisize_hsv[:, :, 2:3], 0, 1.0) + # + # # user_res_fc32_orisize_hsv[:, :, 0:1] = mean_ref_haircolor_hsv[0] + # # user_res_fc32_orisize_hsv[:, :, 1:2] = np.clip(user_res_fc32_orisize_hsv[:, :, 1:2] * 1.15, 0, 1.0) + # # user_res_fc32_orisize_hsv[:, :, 2:3] = np.clip(user_res_fc32_orisize_hsv[:, :, 2:3], 0, 1.0) + # user_res_fc32_orisize = cv2.cvtColor(user_res_fc32_orisize_hsv, cv2.COLOR_HSV2BGR) + # user_res_8uc3_orisize_new = (user_res_fc32_orisize * 255).astype(np.uint8) + # + # user_res_8uc3_orisize_new_LAB = cv2.cvtColor(user_res_8uc3_orisize_new, cv2.COLOR_BGR2LAB) + # user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) + # user_res_8uc3_orisize_new_LAB[:, :, 0] = user_res_8uc3_orisize_LAB[:, :, 0] + # user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_new_LAB, cv2.COLOR_LAB2BGR) + # + # user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32) / 255 + # user_res_fc32_orisize_hsv = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2HSV) + # final_user_hair_hsv_mean_std = self.getavgstd(user_res_fc32_orisize_hsv, user_matting_mask_8uc3_orisize) + # print("final_user_hair_hsv_mean_std: ", final_user_hair_hsv_mean_std) + # #### version 1 + + #### version 2 + user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32) / 255 + # user_res_fc32_orisize_hsv = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2HSV) + ref_rgb_fc32_change_color_768 = ref_rgb_8uc3_change_color_768.astype(np.float32) / 255 + ref_matting_fc32_change_color_768 = ref_matting_8uc3_change_color_768.astype(np.float32) / 255 + # ref_rgb_fc32_change_color_768_hsv = cv2.cvtColor(ref_rgb_fc32_change_color_768, cv2.COLOR_BGR2HSV) + + # ref_rgb_fc32_change_color_768_lab = cv2.cvtColor(ref_rgb_fc32_change_color_768, cv2.COLOR_BGR2LAB) + # ref_avg_std = self.getavgstd(ref_rgb_fc32_change_color_768_lab, ref_matting_fc32_change_color_768) + # # user_res_fc32_orisize_lab = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2LAB) + # # user_avg_std = self.getavgstd(user_res_fc32_orisize_lab, user_matting_mask_8uc3_orisize) + # user_matting_mask_fc32_orisize = user_matting_mask_8uc3_orisize.astype(np.float32) / 255 + # user_res_fc32_orisize, src_avg_std = self.reinhard(user_res_fc32_orisize, user_matting_mask_fc32_orisize, ref_avg_std, ratio=0.99) + + ref_rgb_avg_std = self.getavgstd(ref_rgb_fc32_change_color_768, ref_matting_fc32_change_color_768) + # user_rgb_avg_std = self.getavgstd(user_res_fc32_orisize, user_matting_mask_8uc3_orisize) + user_matting_mask_fc32_orisize = user_matting_mask_8uc3_orisize.astype(np.float32) / 255 + user_res_fc32_orisize, src_avg_std = self.reinhard_rgb(user_res_fc32_orisize, user_matting_mask_fc32_orisize, ref_rgb_avg_std, ratio=0.65) + + user_res_8uc3_orisize_reinhard = (user_res_fc32_orisize * 255).astype(np.uint8) + #### version 2 + + # middle_show = np.concatenate((user_rgb_8uc3_orisize, user_res_8uc3_orisize, user_res_8uc3_orisize_reinhard), axis=1) + # ratio = 1536. / max(middle_show.shape[:2]) + # middle_show = cv2.resize(middle_show, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow('middle_show', middle_show) + # cv2.waitKey(0) + + user_res_8uc3_orisize_reinhard_LAB = cv2.cvtColor(user_res_8uc3_orisize_reinhard, cv2.COLOR_BGR2LAB) + user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) + mix_ratio = 0.1 + user_res_8uc3_orisize_reinhard_LAB[:, :, 0] = user_res_8uc3_orisize_LAB[:, :, 0] * mix_ratio + user_res_8uc3_orisize_reinhard_LAB[:, :, 0] * (1 - mix_ratio) + user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_reinhard_LAB, cv2.COLOR_LAB2BGR) + + user_res_8uc3_orisize2 = user_rgb_8uc3_orisize * (1 - user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize * (user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize2 = (np.clip(user_res_8uc3_orisize2, 0, 255)).astype(np.uint8) + + # middle_show = np.concatenate((user_rgb_8uc3_orisize, user_res_8uc3_orisize, user_matting_mask_8uc3_orisize, user_res_8uc3_orisize2), axis=1) + # ratio = 1536. / max(middle_show.shape[:2]) + # middle_show = cv2.resize(middle_show, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow('middle_show', middle_show) + # cv2.waitKey(0) + + # user_res_8uc3_orisize = np.concatenate((user_res_8uc3_orisize, user_matting_8uc3_bald_orisize), axis=1) + + return user_res_8uc3_orisize2, 0 + + def infer_haircolor_0313_basecolor(self, user_rgb_8uc3_orisize, haircolor_dir): + landmarks_origin_img_1k = self.get_landmark.forward(user_rgb_8uc3_orisize) + if landmarks_origin_img_1k is None: + return None, 10001 + + _, user_matting_8uc1_bald_orisize = self.process_data.generator_matte.matte_inference(user_rgb_8uc3_orisize, + landmarks_origin_img_1k) + user_matting_mask_8uc3_orisize = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) + + ref_rgb_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_rgb_8uc3_color_768.png")) + ref_matting_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_matting_8uc3_color_768.png")) + + user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, user_hair_color_M = \ + self.process_data.get_prepare_hair_color_user_data(user_rgb_8uc3_orisize, + landmarks_origin_img_1k) + + haircolor_dir_basecolor = "/home/yangchaojie/Desktop/hairstyle/hairstyle_infer/data/ref_haircolor/HDR10_443322" + ref_rgb_8uc3_change_color_768_basecolor = cv2.imread(os.path.join(haircolor_dir_basecolor, "ref_rgb_8uc3_color_768.png")) + ref_matting_8uc3_change_color_768_basecolor = cv2.imread(os.path.join(haircolor_dir_basecolor, "ref_matting_8uc3_color_768.png")) + + # start_time = time.time() + hair_gene_color_8uc3_768, user_matting_mask_8uc3_768 = self.change_haircolor.Change_Hair_inference( + user_rgb_8uc3_change_color_768, + user_matting_8uc3_change_color_768, + ref_rgb_8uc3_change_color_768_basecolor, + ref_matting_8uc3_change_color_768_basecolor) + + hair_gene_color_8uc3_orisize = cv2.warpAffine(hair_gene_color_8uc3_768, + cv2.invertAffineTransform(user_hair_color_M), + (user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), + dst=user_rgb_8uc3_orisize.copy(), + flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_TRANSPARENT) + + # if self.use_enhance: + user_res_8uc3_orisize = self.face_enhance.process(hair_gene_color_8uc3_orisize, landmarks_origin_img_1k) + + #### version 2 + user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32) / 255 + # user_res_fc32_orisize_hsv = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2HSV) + ref_rgb_fc32_change_color_768 = ref_rgb_8uc3_change_color_768.astype(np.float32) / 255 + ref_matting_fc32_change_color_768 = ref_matting_8uc3_change_color_768.astype(np.float32) / 255 + # ref_rgb_fc32_change_color_768_hsv = cv2.cvtColor(ref_rgb_fc32_change_color_768, cv2.COLOR_BGR2HSV) + + # ref_rgb_fc32_change_color_768_lab = cv2.cvtColor(ref_rgb_fc32_change_color_768, cv2.COLOR_BGR2LAB) + # ref_avg_std = self.getavgstd(ref_rgb_fc32_change_color_768_lab, ref_matting_fc32_change_color_768) + # # user_res_fc32_orisize_lab = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2LAB) + # # user_avg_std = self.getavgstd(user_res_fc32_orisize_lab, user_matting_mask_8uc3_orisize) + # user_matting_mask_fc32_orisize = user_matting_mask_8uc3_orisize.astype(np.float32) / 255 + # user_res_fc32_orisize, src_avg_std = self.reinhard(user_res_fc32_orisize, user_matting_mask_fc32_orisize, ref_avg_std, ratio=0.99) + + ref_rgb_avg_std = self.getavgstd(ref_rgb_fc32_change_color_768, ref_matting_fc32_change_color_768) + # user_rgb_avg_std = self.getavgstd(user_res_fc32_orisize, user_matting_mask_8uc3_orisize) + user_matting_mask_fc32_orisize = user_matting_mask_8uc3_orisize.astype(np.float32) / 255 + user_res_fc32_orisize, src_avg_std = self.reinhard_rgb(user_res_fc32_orisize, user_matting_mask_fc32_orisize, ref_rgb_avg_std, ratio=0.95) + + user_res_8uc3_orisize = (user_res_fc32_orisize * 255).astype(np.uint8) + #### version 2 + + user_res_8uc3_orisize2 = user_rgb_8uc3_orisize * (1 - user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize * (user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize2 = (np.clip(user_res_8uc3_orisize2, 0, 255)).astype(np.uint8) + + # middle_show = np.concatenate((user_rgb_8uc3_orisize, user_res_8uc3_orisize, user_matting_mask_8uc3_orisize, user_res_8uc3_orisize2), axis=1) + # ratio = 1536. / max(middle_show.shape[:2]) + # middle_show = cv2.resize(middle_show, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow('middle_show', middle_show) + # cv2.waitKey(0) + + # user_res_8uc3_orisize = np.concatenate((user_res_8uc3_orisize, user_matting_8uc3_bald_orisize), axis=1) + + return user_res_8uc3_orisize2, 0 + + def infer_haircolor_0313_rgb(self, user_rgb_8uc3_orisize, haircolor_dir): + landmarks_origin_img_1k = self.get_landmark.forward(user_rgb_8uc3_orisize) + if landmarks_origin_img_1k is None: + return None, 10001 + + _, user_matting_8uc1_bald_orisize = self.process_data.generator_matte.matte_inference(user_rgb_8uc3_orisize, + landmarks_origin_img_1k) + user_matting_mask_8uc3_orisize = np.repeat(user_matting_8uc1_bald_orisize[:, :, np.newaxis], 3, axis=2) + + ref_rgb_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_rgb_8uc3_color_768.png")) + ref_matting_8uc3_change_color_768 = cv2.imread(os.path.join(haircolor_dir, "ref_matting_8uc3_color_768.png")) + # if self.use_enhance: + user_res_8uc3_orisize = user_rgb_8uc3_orisize.copy() + + #### version 2 + user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32) / 255 + # user_res_fc32_orisize_hsv = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2HSV) + ref_rgb_fc32_change_color_768 = ref_rgb_8uc3_change_color_768.astype(np.float32) / 255 + ref_matting_fc32_change_color_768 = ref_matting_8uc3_change_color_768.astype(np.float32) / 255 + # ref_rgb_fc32_change_color_768_hsv = cv2.cvtColor(ref_rgb_fc32_change_color_768, cv2.COLOR_BGR2HSV) + + # ref_rgb_fc32_change_color_768_lab = cv2.cvtColor(ref_rgb_fc32_change_color_768, cv2.COLOR_BGR2LAB) + # ref_avg_std = self.getavgstd(ref_rgb_fc32_change_color_768_lab, ref_matting_fc32_change_color_768) + # # user_res_fc32_orisize_lab = cv2.cvtColor(user_res_fc32_orisize, cv2.COLOR_BGR2LAB) + # # user_avg_std = self.getavgstd(user_res_fc32_orisize_lab, user_matting_mask_8uc3_orisize) + # user_matting_mask_fc32_orisize = user_matting_mask_8uc3_orisize.astype(np.float32) / 255 + # user_res_fc32_orisize, src_avg_std = self.reinhard(user_res_fc32_orisize, user_matting_mask_fc32_orisize, ref_avg_std, ratio=0.99) + + ref_rgb_avg_std = self.getavgstd(ref_rgb_fc32_change_color_768, ref_matting_fc32_change_color_768) + # user_rgb_avg_std = self.getavgstd(user_res_fc32_orisize, user_matting_mask_8uc3_orisize) + user_matting_mask_fc32_orisize = user_matting_mask_8uc3_orisize.astype(np.float32) / 255 + user_res_fc32_orisize, src_avg_std = self.reinhard_rgb(user_res_fc32_orisize, user_matting_mask_fc32_orisize, ref_rgb_avg_std, ratio=1.0) + + user_res_8uc3_orisize = (user_res_fc32_orisize * 255).astype(np.uint8) + #### version 2 + + # user_res_8uc3_orisize_LAB = cv2.cvtColor(user_res_8uc3_orisize, cv2.COLOR_BGR2LAB) + # user_rgb_8uc3_orisize_LAB = cv2.cvtColor(user_rgb_8uc3_orisize, cv2.COLOR_BGR2LAB) + # mix_ratio = 0.5 + # user_res_8uc3_orisize_LAB[:, :, 0] = user_rgb_8uc3_orisize_LAB[:, :, 0] * mix_ratio + user_res_8uc3_orisize_LAB[:, :, 0] * (1 - mix_ratio) + # user_res_8uc3_orisize = cv2.cvtColor(user_res_8uc3_orisize_LAB, cv2.COLOR_LAB2BGR) + + user_res_8uc3_orisize2 = user_rgb_8uc3_orisize * (1 - user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize * (user_matting_mask_8uc3_orisize / 255) + user_res_8uc3_orisize2 = (np.clip(user_res_8uc3_orisize2, 0, 255)).astype(np.uint8) + + # middle_show = np.concatenate((user_rgb_8uc3_orisize, user_res_8uc3_orisize, user_matting_mask_8uc3_orisize, user_res_8uc3_orisize2), axis=1) + # ratio = 1536. / max(middle_show.shape[:2]) + # middle_show = cv2.resize(middle_show, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow('middle_show', middle_show) + # cv2.waitKey(0) + + # user_res_8uc3_orisize = np.concatenate((user_res_8uc3_orisize, user_matting_8uc3_bald_orisize), axis=1) + + return user_res_8uc3_orisize2, 0 + + def infer_bald(self, user_rgb_8uc3_orisize): + landmarks_origin_img_1k = self.get_landmark.forward(user_rgb_8uc3_orisize) + if landmarks_origin_img_1k is None: + return None, 10001 + # 返回光头 + user_res_8uc3_orisize, user_inter_res_8uc3_orisize, user_matting_8uc3_bald_768, user_baldseg_8uc3_bald_768 = self.process_data.get_user_blad( + user_rgb_8uc3_orisize, + landmarks_origin_img_1k) + if self.use_enhance: + user_res_8uc3_orisize = self.face_enhance.process(user_res_8uc3_orisize, landmarks_origin_img_1k) + return user_res_8uc3_orisize, 0 + + + def buff(self, img, img_skin, value1, value2): + img = img.astype(np.float32) + dx = value1 * 5 + fc = value1 * 12.5 + p = 80 + temp1 = cv2.bilateralFilter(img, dx, fc, fc) + temp2 = (temp1 - img + 128) + temp2 = np.clip(temp2, 0, 255) + temp3 = cv2.GaussianBlur(temp2, (2 * value2 - 1, 2 * value2 - 1), 0, 0) + temp4 = img + 2 * temp3 - 255 + temp4 = np.clip(temp4, 0, 255) + dst = img * ((100 - p) / 100) + temp4 * (p / 100) + img_skin_c = 1-img_skin + dst = dst * img_skin + img * img_skin_c + + # mid_show = np.concatenate((img, img_skin, dst), axis=1) + # ratio = 1224. / max(mid_show.shape[:2]) + # mid_show = cv2.resize(mid_show, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow('mid_show', mid_show) + # cv2.waitKey(0) + + return dst.astype(np.uint8) + def whitening(self, img, img_skin, value): + midtones_add = np.zeros(256) + for i in range(256): + midtones_add[i] = 0.667 * (1 - ((i - 127) / 127) * ((i - 127) / 127)) + lookup = np.zeros(256, dtype='uint8') + for i in range(256): + red = i + red += value * midtones_add[red] + red = max(0, red) + lookup[i] = np.uint(red) + w, h, c = img.shape + + img_skin = img_skin[:,:,-1] + index = np.where(img_skin == 1) + for i in range(index[0].shape[0]): + img[index[0][i], index[1][i], 0] = lookup[int(img[index[0][i], index[1][i], 0])] + img[index[0][i], index[1][i], 1] = lookup[int(img[index[0][i], index[1][i], 1])] + img[index[0][i], index[1][i], 2] = lookup[int(img[index[0][i], index[1][i], 2])] + + # for i in range(w): + # for j in range(h): + # if img_skin[i, j, 0] == 1: + # img[i, j, 0] = lookup[img[i, j, 0]] + # img[i, j, 1] = lookup[img[i, j, 1]] + # img[i, j, 2] = lookup[img[i, j, 2]] + return img + def infer_face(self, origin_img): + landmarks_origin_img_1k = self.get_landmark.forward(origin_img) + if landmarks_origin_img_1k is None: + return None + landmark_137 = landmark_processor.pts_1k_to_137(landmarks_origin_img_1k).astype(np.int32) + face_mask_ori = self.face_seg.inference(origin_img, landmarks_origin_img_1k) + + # if get_facecolor: + h, w, _ = origin_img.shape + inpaint_mask = np.zeros((h, w), dtype=np.uint8) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[129:137]), (255,)) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[121:129]), (255,)) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[22:48]), (255,)) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[88:104]), (255,)) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[105:121]), (255,)) + inpaint_mask = cv2.cvtColor(inpaint_mask, cv2.COLOR_GRAY2BGR) + face_mask_del = np.clip(face_mask_ori - (inpaint_mask / 255).astype(np.float32), 0, 1) + # face_image_del = (origin_img*face_mask_del).astype(np.uint8) + # face_mask_del = np.around(face_mask_del) + + dst_img_white = self.whitening(origin_img, face_mask_del, 10) + + # cv2.imwrite("/home/yangchaojie/Desktop/hairstyle/hairstyle_infer/test_data/test_face_white.jpg", dst_img_white) + # whiten_show = np.concatenate((origin_img, dst_img_white), axis=1) + # ratio = 1024. / max(whiten_show.shape[:2]) + # whiten_show = cv2.resize(whiten_show, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow('whiten_show', whiten_show) + # cv2.waitKey(0) + + # bounding_boxs = bounding_boxs.astype(np.int32) + face_left = landmark_137[16] + face_right = landmark_137[6] + face_top = landmark_137[11] + face_bottom = landmark_137[0] + x1, y1, x2, y2 = max(face_left[0] - 20, 0), max(face_top[1] - 20, 0), min(face_right[0] + 20, w), min( + face_bottom[1] + 20, h) + img_beauty = dst_img_white.copy() + dst_img_buff = self.buff(dst_img_white[y1:y2, x1:x2], face_mask_del[y1:y2, x1:x2], 4, 3) + img_beauty[y1:y2, x1:x2] = dst_img_buff + + # cv2.imwrite("/home/yangchaojie/Desktop/hairstyle/hairstyle_infer/test_data/test_face_buff.jpg", img_beauty) + # buff_show = np.concatenate((origin_img, dst_img_white, img_beauty), axis=1) + # ratio = 1224. / max(buff_show.shape[:2]) + # buff_show = cv2.resize(buff_show, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow('buff_show', buff_show) + # cv2.waitKey(0) + + return img_beauty + + +if __name__=='__main__': + hf = HairStyle_Model_Infer() + # user_img = cv2.imread() + # ref_color = + # hf.infer_haircolor_new() \ No newline at end of file diff --git a/hair_service_sd/infer_test.py b/hair_service_sd/infer_test.py new file mode 100644 index 0000000..4b87e51 --- /dev/null +++ b/hair_service_sd/infer_test.py @@ -0,0 +1,321 @@ +#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.callback import recall +from gen_super_image import webui_img2img, webui_img2img_diy, webui_super_res_img +from utils import enhance_hair +import configparser +from common.logger import config + +from queue import Queue +from concurrent.futures import ThreadPoolExecutor +from common.callback import * +from utils import call_hair_inter +from utils import landmark_processor +from change_color import process_infer, resize_pre_webui + + + +hairstyle_process = HairStyle_Model(gpu=True,use_enhance=True) +hairstyle_process_infer = HairStyle_Model_Infer(gpu=True, use_enhance=False) +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_hairstyle_v2(): + hairstyle_dir = config.get('default', 'hairstyleDir') + user_dir = config.get('default', 'userDir') + ref_img_dir = config.get('default', 'refImgDir') + res_dir = config.get('default', 'res_dir') + hair_template_material_dir = config.get('default', 'hair_template_material_dir') + train_dir = config.get('default', 'train_dir') + + start_time0 = time.time() + is_hr = False + + ret = { + "state": -1, + "msg": "fail", + "result": "", + "umd": "" + } + + # get req input + try: + hair_id = "042abc5d-262f-473e-bd62-79ab0a2fdce6" + task_id = hair_id + "_" + str(uuid4()) + hair_material_dir = os.path.join(train_dir, hair_id) + + user_img_url = "https://ydapp-1317132355.cos.ap-beijing.myqcloud.com/hair_mz/images/tmp/20240923166694465407_2.jpg" + userId = "47385" + user_img_path, _ = download_img(user_img_url, userId) + + print("--------------download:", time.time() - start_time0) + + except Exception as e: + print(e) + + + try: + # 获取用户图 user img + user_img_name = user_img_path[user_img_path.rfind("/") + 1:] + new_user_img_path = os.path.join(user_dir, user_img_name) + print("new_user_img_path: ", new_user_img_path) + + if os.path.exists(new_user_img_path): + os.remove(new_user_img_path) + + shutil.copy(user_img_path, new_user_img_path) + + # 获取发型图 ref hair img + template_ref_hair_name = "" + template_ref_hair_path = "" + hair_material_save_dir = os.path.join(hair_template_material_dir, hair_id) + + # upload train dir + train_img_save_dir = os.path.join(train_upload_dir, hair_id) + hair_name_lists = os.listdir(train_img_save_dir) + + for hair_name in hair_name_lists: + if "first##" in hair_name: + template_ref_hair_name = hair_name + template_ref_hair_path = os.path.join(train_img_save_dir, hair_name) + break + + + + new_hair_ref_img_path = os.path.join(ref_img_dir, template_ref_hair_name) + print("new_hair_ref_img_path: ", new_hair_ref_img_path) + + # hair_template_save_dir = os.path.join(hair_template_dir, hair_id) + # pkl_process(hair_template_save_dir) + + long_flag = False + long_txt_path = os.path.join(hair_material_dir, "long.txt") + if os.path.exists(long_txt_path): + long_flag = True + print("long_flag:", long_flag) + + # 1. gen hair res img material + start1 = time.time() + material_save_path = os.path.join(hairstyle_dir, hair_id) + print("first_hair_material_save_path :", material_save_path) + if not os.path.exists(material_save_path): + print("!!! gen material !!!!") + shutil.copy(template_ref_hair_path, new_hair_ref_img_path) + prepare_single(new_hair_ref_img_path, material_save_path, long_flag) + print("--------------prepare_single:", time.time() - start1) + + # 2. swap user hair to ref img hair + dst_path = os.path.join(res_dir, task_id + ".png") + if not os.path.exists(os.path.dirname(dst_path)): + os.makedirs(os.path.dirname(dst_path)) + + start2 = time.time() + origin_img = cv2.imread(new_user_img_path) + + # resize origin user img + user_scale = 1920 / max(origin_img.shape[0], origin_img.shape[1]) + if user_scale < 1.0: + origin_img = cv2.resize(origin_img, (0, 0), fx=user_scale, fy=user_scale, interpolation=cv2.INTER_LANCZOS4) + + ret_dict, status = hairstyle_process_infer.infer_hairstyle(origin_img, material_save_path, return_pt1k=True, + use_enhance=False) + img_res = ret_dict['user_res_8uc3_orisize'] + origin_pt1k = ret_dict['landmarks_origin_img_1k'] + user_bald_res_8uc3_orisize = ret_dict['user_bald_res_8uc3_orisize'] + + print("--------------hairstyle_process infer_hairstyle:", time.time() - start2) + + if status == 0: + # 读取新图抠图 + hair_matting_path = os.path.join(material_save_path, "hair_mask_2.png") + new_matting = cv2.imread(hair_matting_path, cv2.IMREAD_GRAYSCALE) + + # 读取结果图 + result_img = img_res + + user_orig_mask_path = os.path.join(material_save_path, "user_orig_mask.png") + origin_matting = cv2.imread(user_orig_mask_path, cv2.IMREAD_GRAYSCALE) + + # 获取头发处理的局部区域图像 + box_info = hairstyle_process.get_body_info(result_img) + dst_size = (576, 768) + + 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) + + crop_result = landmark_processor.high_quality_warpAffine(result_img, M, dst_size) + + + # matting_merge = np.concatenate([origin_matting[:, :, np.newaxis], new_matting[:, :, np.newaxis]], + # axis=2) + # matting_merge = np.max(matting_merge, axis=2) + + matting_merge = new_matting[:, :, np.newaxis] + crop_matting = cv2.warpAffine(matting_merge, M, dst_size) + crop_new_matting = cv2.warpAffine(new_matting, M, dst_size) + + mask = (crop_matting > 10).astype(np.float32) + if not is_hr: + mask_dilate = cv2.dilate(mask, np.ones((3, 9), np.uint8)) + else: + mask_dilate = cv2.dilate(mask, np.ones((6, 18), np.uint8)) + + pt1k = landmark_processor.transform_points(origin_pt1k, M) + # face_mask = landmark_processor.draw_hull_mask(pt1k.astype(np.int32), + # w=mask_dilate.shape[1], h=mask_dilate.shape[0], + # is_gray=True).astype(np.float32) + face_mask = landmark_processor.draw_half_mask(pt1k.astype(np.int32), + w=mask_dilate.shape[1], h=mask_dilate.shape[0], + is_gray=True).astype(np.float32) + + face_mask = cv2.erode(face_mask, np.ones((19, 19), np.uint8)) + face_mask = cv2.blur(face_mask, (11, 11)) + + crop_new_matting_f32 = crop_new_matting.astype(np.float32) / 255 + face_mask = np.clip(face_mask - crop_new_matting_f32, 0, 1) + mask_dilate = np.clip(mask_dilate - face_mask, 0, 1) + + final_img = crop_result + + mask_dilate = np.clip(mask_dilate * 255, 0, 255).astype(np.uint8) + + # get gender + config_json_path = os.path.join(material_save_path, "config.json") + with open(config_json_path, "r") as f: + config_json_content = json.load(f) + in_gender = config_json_content["gender"] + print("in_gender:", in_gender) + + images_dir = os.path.join(hair_material_dir, "images") + txt_dir = os.path.join(images_dir, os.listdir(images_dir)[0]) + txt_path = glob.glob(txt_dir + '/*.txt')[0] + with open(txt_path, 'r') as f: + p_tag = f.readline() + if "titor hairstyle, faceless, no human, gray background, simple background" in p_tag: + p_tag = p_tag[p_tag.find("simple background, ") + len("simple background, "):] + else: + p_tag = "" + + start4 = time.time() + # cv2.imshow("final_img", final_img) + # cv2.imshow("face_mask", face_mask) + # cv2.imshow("mask_dilate", mask_dilate) + + denoising_strength = 0.6 + + sd_result = webui_img2img(img=final_img, mask_img=mask_dilate, in_gender=in_gender, task_id=task_id, + hair_id=hair_id, lora_material_path=hair_material_dir, tag=p_tag, is_hr=is_hr, + denoising_strength=denoising_strength) + # cv2.imshow("final_img", cv2.resize(final_img, (768, 768), interpolation=cv2.INTER_AREA)) + # cv2.imshow("mask_dilate", cv2.resize(mask_dilate, (768, 768), interpolation=cv2.INTER_AREA)) + # cv2.waitKey(0) + print("--------------webui_img2img:", time.time() - start4) + + start5 = time.time() + origin_img_final = origin_img.copy() + + # restore the sd_result_small to origin_img + M_inv = cv2.invertAffineTransform(M) + cv2.warpAffine(sd_result, M_inv, (origin_img_final.shape[1], origin_img_final.shape[0]), + dst=origin_img_final, + borderMode=cv2.BORDER_TRANSPARENT, flags=cv2.INTER_LANCZOS4) + + # face_mask_origin = cv2.warpAffine(face_mask, M_inv, (origin_img_final.shape[1], origin_img_final.shape[0]))[ + # :, :, np.newaxis] + # origin_img_final = (origin_img_final * ( + # 1 - face_mask_origin) + user_bald_res_8uc3_orisize * face_mask_origin).astype(np.uint8) + + cv2.imwrite(dst_path, origin_img_final) + print("origin_img_final shape:", origin_img_final.shape) + + # cv2.imshow("origin_img_final", origin_img_final) + # cv2.waitKey(0) + + + # t_name_res = 'digital_cloth/' + task_id + ".png" + # res_url = oss_2.upload_file(dst_path, t_name_res) + # print('res url:', res_url) + # print("--------------write:", time.time() - start5) + ret_url = hairstyle_process.oss2.upload_file(dst_path, + "hair_mz/images/hairstyle/{}/{}".format(hair_id, + str(uuid4()) + '.jpg')) + + ret["msg"] = 'success' + ret['state'] = 0 + ret['result'] = ret_url + + print("---------------------------------- cost: ", time.time() - start_time0) + print("status :", status) + print("\n\n") + + print('\n\n\n\n') + + except Exception as e: + print(e) + ret["msg"] = str(traceback.print_exc()) + ret["result"] = "" + ret['state'] = -1 + +if __name__ == '__main__': + change_hairstyle_v2() \ No newline at end of file diff --git a/hair_service_sd/keypoints/experiments/coco/hrnet/coco25_384x288_adam_lr1e-3.yaml b/hair_service_sd/keypoints/experiments/coco/hrnet/coco25_384x288_adam_lr1e-3.yaml new file mode 100644 index 0000000..0a71dc8 --- /dev/null +++ b/hair_service_sd/keypoints/experiments/coco/hrnet/coco25_384x288_adam_lr1e-3.yaml @@ -0,0 +1,127 @@ +AUTO_RESUME: true +CUDNN: + BENCHMARK: true + DETERMINISTIC: false + ENABLED: true +DATA_DIR: '' +GPUS: (0,) +OUTPUT_DIR: 'output' +LOG_DIR: 'log' +WORKERS: 1 +PRINT_FREQ: 100 + +DATASET: + COLOR_RGB: true + DATASET: 'coco' + DATA_FORMAT: jpg + FLIP: true + NUM_JOINTS_HALF_BODY: 8 + PROB_HALF_BODY: 0.3 + ROOT: 'data/coco/' + ROT_FACTOR: 45 + SCALE_FACTOR: 0.35 + TEST_SET: 'val2017' + TRAIN_SET: 'train2017' +MODEL: + INIT_WEIGHTS: true + NAME: pose_hrnet + NUM_JOINTS: 25 + PRETRAINED: 'models/pytorch/imagenet/hrnet_w48-8ef0771d.pth' + TARGET_TYPE: gaussian + IMAGE_SIZE: + - 288 + - 384 + HEATMAP_SIZE: + - 72 + - 96 + SIGMA: 3 + EXTRA: + PRETRAINED_LAYERS: + - 'conv1' + - 'bn1' + - 'conv2' + - 'bn2' + - 'layer1' + - 'transition1' + - 'stage2' + - 'transition2' + - 'stage3' + - 'transition3' + - 'stage4' + FINAL_CONV_KERNEL: 1 + STAGE2: + NUM_MODULES: 1 + NUM_BRANCHES: 2 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + NUM_CHANNELS: + - 48 + - 96 + FUSE_METHOD: SUM + STAGE3: + NUM_MODULES: 4 + NUM_BRANCHES: 3 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + - 4 + NUM_CHANNELS: + - 48 + - 96 + - 192 + FUSE_METHOD: SUM + STAGE4: + NUM_MODULES: 3 + NUM_BRANCHES: 4 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + - 4 + - 4 + NUM_CHANNELS: + - 48 + - 96 + - 192 + - 384 + FUSE_METHOD: SUM +LOSS: + USE_TARGET_WEIGHT: true +TRAIN: + BATCH_SIZE_PER_GPU: 24 + SHUFFLE: true + BEGIN_EPOCH: 0 + END_EPOCH: 210 + OPTIMIZER: adam + LR: 0.001 + LR_FACTOR: 0.1 + LR_STEP: + - 170 + - 200 + WD: 0.0001 + GAMMA1: 0.99 + GAMMA2: 0.0 + MOMENTUM: 0.9 + NESTEROV: false +TEST: + BATCH_SIZE_PER_GPU: 1 + COCO_BBOX_FILE: 'data/coco/person_detection_results/COCO_val2017_detections_AP_H_56_person.json' + BBOX_THRE: 1.0 + IMAGE_THRE: 0.0 + IN_VIS_THRE: 0.2 + MODEL_FILE: '' + NMS_THRE: 1.0 + OKS_THRE: 0.9 + USE_GT_BBOX: true + FLIP_TEST: false + POST_PROCESS: true + SHIFT_HEATMAP: true +DEBUG: + DEBUG: true + SAVE_BATCH_IMAGES_GT: true + SAVE_BATCH_IMAGES_PRED: true + SAVE_HEATMAPS_GT: true + SAVE_HEATMAPS_PRED: true diff --git a/hair_service_sd/keypoints/experiments/coco/hrnet/w32_256x192_adam_lr1e-3.yaml b/hair_service_sd/keypoints/experiments/coco/hrnet/w32_256x192_adam_lr1e-3.yaml new file mode 100644 index 0000000..16854cf --- /dev/null +++ b/hair_service_sd/keypoints/experiments/coco/hrnet/w32_256x192_adam_lr1e-3.yaml @@ -0,0 +1,127 @@ +AUTO_RESUME: true +CUDNN: + BENCHMARK: true + DETERMINISTIC: false + ENABLED: true +DATA_DIR: '' +GPUS: (0,1,2,3) +OUTPUT_DIR: 'output' +LOG_DIR: 'log' +WORKERS: 24 +PRINT_FREQ: 100 + +DATASET: + COLOR_RGB: true + DATASET: 'coco' + DATA_FORMAT: jpg + FLIP: true + NUM_JOINTS_HALF_BODY: 8 + PROB_HALF_BODY: 0.3 + ROOT: 'data/coco/' + ROT_FACTOR: 45 + SCALE_FACTOR: 0.35 + TEST_SET: 'val2017' + TRAIN_SET: 'train2017' +MODEL: + INIT_WEIGHTS: true + NAME: pose_hrnet + NUM_JOINTS: 17 + PRETRAINED: 'models/pytorch/imagenet/hrnet_w32-36af842e.pth' + TARGET_TYPE: gaussian + IMAGE_SIZE: + - 192 + - 256 + HEATMAP_SIZE: + - 48 + - 64 + SIGMA: 2 + EXTRA: + PRETRAINED_LAYERS: + - 'conv1' + - 'bn1' + - 'conv2' + - 'bn2' + - 'layer1' + - 'transition1' + - 'stage2' + - 'transition2' + - 'stage3' + - 'transition3' + - 'stage4' + FINAL_CONV_KERNEL: 1 + STAGE2: + NUM_MODULES: 1 + NUM_BRANCHES: 2 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + NUM_CHANNELS: + - 32 + - 64 + FUSE_METHOD: SUM + STAGE3: + NUM_MODULES: 4 + NUM_BRANCHES: 3 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + - 4 + NUM_CHANNELS: + - 32 + - 64 + - 128 + FUSE_METHOD: SUM + STAGE4: + NUM_MODULES: 3 + NUM_BRANCHES: 4 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + - 4 + - 4 + NUM_CHANNELS: + - 32 + - 64 + - 128 + - 256 + FUSE_METHOD: SUM +LOSS: + USE_TARGET_WEIGHT: true +TRAIN: + BATCH_SIZE_PER_GPU: 32 + SHUFFLE: true + BEGIN_EPOCH: 0 + END_EPOCH: 210 + OPTIMIZER: adam + LR: 0.001 + LR_FACTOR: 0.1 + LR_STEP: + - 170 + - 200 + WD: 0.0001 + GAMMA1: 0.99 + GAMMA2: 0.0 + MOMENTUM: 0.9 + NESTEROV: false +TEST: + BATCH_SIZE_PER_GPU: 32 + COCO_BBOX_FILE: 'data/coco/person_detection_results/COCO_val2017_detections_AP_H_56_person.json' + BBOX_THRE: 1.0 + IMAGE_THRE: 0.0 + IN_VIS_THRE: 0.2 + MODEL_FILE: '' + NMS_THRE: 1.0 + OKS_THRE: 0.9 + USE_GT_BBOX: true + FLIP_TEST: true + POST_PROCESS: true + SHIFT_HEATMAP: true +DEBUG: + DEBUG: true + SAVE_BATCH_IMAGES_GT: true + SAVE_BATCH_IMAGES_PRED: true + SAVE_HEATMAPS_GT: true + SAVE_HEATMAPS_PRED: true diff --git a/hair_service_sd/keypoints/experiments/coco/hrnet/w32_384x288_adam_lr1e-3.yaml b/hair_service_sd/keypoints/experiments/coco/hrnet/w32_384x288_adam_lr1e-3.yaml new file mode 100644 index 0000000..57101e9 --- /dev/null +++ b/hair_service_sd/keypoints/experiments/coco/hrnet/w32_384x288_adam_lr1e-3.yaml @@ -0,0 +1,127 @@ +AUTO_RESUME: true +CUDNN: + BENCHMARK: true + DETERMINISTIC: false + ENABLED: true +DATA_DIR: '' +GPUS: (0,1,2,3) +OUTPUT_DIR: 'output' +LOG_DIR: 'log' +WORKERS: 24 +PRINT_FREQ: 100 + +DATASET: + COLOR_RGB: true + DATASET: 'coco' + DATA_FORMAT: jpg + FLIP: true + NUM_JOINTS_HALF_BODY: 8 + PROB_HALF_BODY: 0.3 + ROOT: 'data/coco/' + ROT_FACTOR: 45 + SCALE_FACTOR: 0.35 + TEST_SET: 'val2017' + TRAIN_SET: 'train2017' +MODEL: + INIT_WEIGHTS: true + NAME: pose_hrnet + NUM_JOINTS: 17 + PRETRAINED: 'models/pytorch/imagenet/hrnet_w32-36af842e.pth' + TARGET_TYPE: gaussian + IMAGE_SIZE: + - 288 + - 384 + HEATMAP_SIZE: + - 72 + - 96 + SIGMA: 3 + EXTRA: + PRETRAINED_LAYERS: + - 'conv1' + - 'bn1' + - 'conv2' + - 'bn2' + - 'layer1' + - 'transition1' + - 'stage2' + - 'transition2' + - 'stage3' + - 'transition3' + - 'stage4' + FINAL_CONV_KERNEL: 1 + STAGE2: + NUM_MODULES: 1 + NUM_BRANCHES: 2 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + NUM_CHANNELS: + - 32 + - 64 + FUSE_METHOD: SUM + STAGE3: + NUM_MODULES: 4 + NUM_BRANCHES: 3 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + - 4 + NUM_CHANNELS: + - 32 + - 64 + - 128 + FUSE_METHOD: SUM + STAGE4: + NUM_MODULES: 3 + NUM_BRANCHES: 4 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + - 4 + - 4 + NUM_CHANNELS: + - 32 + - 64 + - 128 + - 256 + FUSE_METHOD: SUM +LOSS: + USE_TARGET_WEIGHT: true +TRAIN: + BATCH_SIZE_PER_GPU: 32 + SHUFFLE: true + BEGIN_EPOCH: 0 + END_EPOCH: 210 + OPTIMIZER: adam + LR: 0.001 + LR_FACTOR: 0.1 + LR_STEP: + - 170 + - 200 + WD: 0.0001 + GAMMA1: 0.99 + GAMMA2: 0.0 + MOMENTUM: 0.9 + NESTEROV: false +TEST: + BATCH_SIZE_PER_GPU: 32 + COCO_BBOX_FILE: 'data/coco/person_detection_results/COCO_val2017_detections_AP_H_56_person.json' + BBOX_THRE: 1.0 + IMAGE_THRE: 0.0 + IN_VIS_THRE: 0.2 + MODEL_FILE: '' + NMS_THRE: 1.0 + OKS_THRE: 0.9 + USE_GT_BBOX: true + FLIP_TEST: true + POST_PROCESS: true + SHIFT_HEATMAP: true +DEBUG: + DEBUG: true + SAVE_BATCH_IMAGES_GT: true + SAVE_BATCH_IMAGES_PRED: true + SAVE_HEATMAPS_GT: true + SAVE_HEATMAPS_PRED: true diff --git a/hair_service_sd/keypoints/experiments/coco/hrnet/w48_256x192_adam_lr1e-3.yaml b/hair_service_sd/keypoints/experiments/coco/hrnet/w48_256x192_adam_lr1e-3.yaml new file mode 100644 index 0000000..45c7011 --- /dev/null +++ b/hair_service_sd/keypoints/experiments/coco/hrnet/w48_256x192_adam_lr1e-3.yaml @@ -0,0 +1,127 @@ +AUTO_RESUME: true +CUDNN: + BENCHMARK: true + DETERMINISTIC: false + ENABLED: true +DATA_DIR: '' +GPUS: (0,1,2,3) +OUTPUT_DIR: 'output' +LOG_DIR: 'log' +WORKERS: 24 +PRINT_FREQ: 100 + +DATASET: + COLOR_RGB: true + DATASET: 'coco' + DATA_FORMAT: jpg + FLIP: true + NUM_JOINTS_HALF_BODY: 8 + PROB_HALF_BODY: 0.3 + ROOT: 'data/coco/' + ROT_FACTOR: 45 + SCALE_FACTOR: 0.35 + TEST_SET: 'val2017' + TRAIN_SET: 'train2017' +MODEL: + INIT_WEIGHTS: true + NAME: pose_hrnet + NUM_JOINTS: 17 + PRETRAINED: 'models/pytorch/imagenet/hrnet_w48-8ef0771d.pth' + TARGET_TYPE: gaussian + IMAGE_SIZE: + - 192 + - 256 + HEATMAP_SIZE: + - 48 + - 64 + SIGMA: 2 + EXTRA: + PRETRAINED_LAYERS: + - 'conv1' + - 'bn1' + - 'conv2' + - 'bn2' + - 'layer1' + - 'transition1' + - 'stage2' + - 'transition2' + - 'stage3' + - 'transition3' + - 'stage4' + FINAL_CONV_KERNEL: 1 + STAGE2: + NUM_MODULES: 1 + NUM_BRANCHES: 2 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + NUM_CHANNELS: + - 48 + - 96 + FUSE_METHOD: SUM + STAGE3: + NUM_MODULES: 4 + NUM_BRANCHES: 3 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + - 4 + NUM_CHANNELS: + - 48 + - 96 + - 192 + FUSE_METHOD: SUM + STAGE4: + NUM_MODULES: 3 + NUM_BRANCHES: 4 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + - 4 + - 4 + NUM_CHANNELS: + - 48 + - 96 + - 192 + - 384 + FUSE_METHOD: SUM +LOSS: + USE_TARGET_WEIGHT: true +TRAIN: + BATCH_SIZE_PER_GPU: 32 + SHUFFLE: true + BEGIN_EPOCH: 0 + END_EPOCH: 210 + OPTIMIZER: adam + LR: 0.001 + LR_FACTOR: 0.1 + LR_STEP: + - 170 + - 200 + WD: 0.0001 + GAMMA1: 0.99 + GAMMA2: 0.0 + MOMENTUM: 0.9 + NESTEROV: false +TEST: + BATCH_SIZE_PER_GPU: 32 + COCO_BBOX_FILE: 'data/coco/person_detection_results/COCO_val2017_detections_AP_H_56_person.json' + BBOX_THRE: 1.0 + IMAGE_THRE: 0.0 + IN_VIS_THRE: 0.2 + MODEL_FILE: '' + NMS_THRE: 1.0 + OKS_THRE: 0.9 + USE_GT_BBOX: true + FLIP_TEST: true + POST_PROCESS: true + SHIFT_HEATMAP: true +DEBUG: + DEBUG: true + SAVE_BATCH_IMAGES_GT: true + SAVE_BATCH_IMAGES_PRED: true + SAVE_HEATMAPS_GT: true + SAVE_HEATMAPS_PRED: true diff --git a/hair_service_sd/keypoints/experiments/coco/hrnet/w48_384x288_adam_lr1e-3.yaml b/hair_service_sd/keypoints/experiments/coco/hrnet/w48_384x288_adam_lr1e-3.yaml new file mode 100644 index 0000000..d83b74b --- /dev/null +++ b/hair_service_sd/keypoints/experiments/coco/hrnet/w48_384x288_adam_lr1e-3.yaml @@ -0,0 +1,127 @@ +AUTO_RESUME: true +CUDNN: + BENCHMARK: true + DETERMINISTIC: false + ENABLED: true +DATA_DIR: '' +GPUS: (0,) +OUTPUT_DIR: 'output' +LOG_DIR: 'log' +WORKERS: 1 +PRINT_FREQ: 100 + +DATASET: + COLOR_RGB: true + DATASET: 'coco' + DATA_FORMAT: jpg + FLIP: true + NUM_JOINTS_HALF_BODY: 8 + PROB_HALF_BODY: 0.3 + ROOT: 'data/coco/' + ROT_FACTOR: 45 + SCALE_FACTOR: 0.35 + TEST_SET: 'val2017' + TRAIN_SET: 'train2017' +MODEL: + INIT_WEIGHTS: true + NAME: pose_hrnet + NUM_JOINTS: 17 + PRETRAINED: 'models/pytorch/imagenet/hrnet_w48-8ef0771d.pth' + TARGET_TYPE: gaussian + IMAGE_SIZE: + - 288 + - 384 + HEATMAP_SIZE: + - 72 + - 96 + SIGMA: 3 + EXTRA: + PRETRAINED_LAYERS: + - 'conv1' + - 'bn1' + - 'conv2' + - 'bn2' + - 'layer1' + - 'transition1' + - 'stage2' + - 'transition2' + - 'stage3' + - 'transition3' + - 'stage4' + FINAL_CONV_KERNEL: 1 + STAGE2: + NUM_MODULES: 1 + NUM_BRANCHES: 2 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + NUM_CHANNELS: + - 48 + - 96 + FUSE_METHOD: SUM + STAGE3: + NUM_MODULES: 4 + NUM_BRANCHES: 3 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + - 4 + NUM_CHANNELS: + - 48 + - 96 + - 192 + FUSE_METHOD: SUM + STAGE4: + NUM_MODULES: 3 + NUM_BRANCHES: 4 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + - 4 + - 4 + NUM_CHANNELS: + - 48 + - 96 + - 192 + - 384 + FUSE_METHOD: SUM +LOSS: + USE_TARGET_WEIGHT: true +TRAIN: + BATCH_SIZE_PER_GPU: 24 + SHUFFLE: true + BEGIN_EPOCH: 0 + END_EPOCH: 210 + OPTIMIZER: adam + LR: 0.001 + LR_FACTOR: 0.1 + LR_STEP: + - 170 + - 200 + WD: 0.0001 + GAMMA1: 0.99 + GAMMA2: 0.0 + MOMENTUM: 0.9 + NESTEROV: false +TEST: + BATCH_SIZE_PER_GPU: 1 + COCO_BBOX_FILE: 'data/coco/person_detection_results/COCO_val2017_detections_AP_H_56_person.json' + BBOX_THRE: 1.0 + IMAGE_THRE: 0.0 + IN_VIS_THRE: 0.2 + MODEL_FILE: '' + NMS_THRE: 1.0 + OKS_THRE: 0.9 + USE_GT_BBOX: true + FLIP_TEST: false + POST_PROCESS: true + SHIFT_HEATMAP: true +DEBUG: + DEBUG: true + SAVE_BATCH_IMAGES_GT: true + SAVE_BATCH_IMAGES_PRED: true + SAVE_HEATMAPS_GT: true + SAVE_HEATMAPS_PRED: true diff --git a/hair_service_sd/logging.conf b/hair_service_sd/logging.conf new file mode 100644 index 0000000..c50aa87 --- /dev/null +++ b/hair_service_sd/logging.conf @@ -0,0 +1,22 @@ +[loggers] +keys=root + +[handlers] +keys=console + +[formatters] +keys=fileFormatter + +[logger_root] +level=INFO +handlers=console + +[handler_console] +class=StreamHandler +args=(sys.stdout,) +level=INFO +formatter=fileFormatter + +[formatter_fileFormatter] +format=%(asctime)s -%(thread)d-%(filename)s- %(name)s - %(levelname)s - %(message)s +datefmt=%a, %d %b %Y %H:%M:%S \ No newline at end of file diff --git a/hair_service_sd/matting/gca_matting_hair_single.py b/hair_service_sd/matting/gca_matting_hair_single.py new file mode 100644 index 0000000..2dbb013 --- /dev/null +++ b/hair_service_sd/matting/gca_matting_hair_single.py @@ -0,0 +1,209 @@ +import os +import cv2 +import argparse +import numpy as np + +import torch +from torch.nn import functional as F + +import utils +from matting import networks +from utils.data_preprocess import * +from utils import landmark_processor +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]) diff --git a/hair_service_sd/matting/gca_matting_hair_single_fg.py b/hair_service_sd/matting/gca_matting_hair_single_fg.py new file mode 100644 index 0000000..dde45e0 --- /dev/null +++ b/hair_service_sd/matting/gca_matting_hair_single_fg.py @@ -0,0 +1,201 @@ +import os +import cv2 +import argparse +import numpy as np + +import torch +from torch.nn import functional as F + +import utils +from matting import networks +from utils.data_preprocess import * +from utils import landmark_processor +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) + diff --git a/hair_service_sd/matting/networks/__init__.py b/hair_service_sd/matting/networks/__init__.py new file mode 100644 index 0000000..aee2746 --- /dev/null +++ b/hair_service_sd/matting/networks/__init__.py @@ -0,0 +1 @@ +from .generators import * \ No newline at end of file diff --git a/hair_service_sd/matting/networks/decoders/__init__.py b/hair_service_sd/matting/networks/decoders/__init__.py new file mode 100644 index 0000000..a7eac0f --- /dev/null +++ b/hair_service_sd/matting/networks/decoders/__init__.py @@ -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) \ No newline at end of file diff --git a/hair_service_sd/matting/networks/decoders/res_gca_dec.py b/hair_service_sd/matting/networks/decoders/res_gca_dec.py new file mode 100644 index 0000000..7515bac --- /dev/null +++ b/hair_service_sd/matting/networks/decoders/res_gca_dec.py @@ -0,0 +1,28 @@ +from matting.networks.ops import GuidedCxtAtten, SpectralNorm +from 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} + diff --git a/hair_service_sd/matting/networks/decoders/res_shortcut_dec.py b/hair_service_sd/matting/networks/decoders/res_shortcut_dec.py new file mode 100644 index 0000000..d1c2d1e --- /dev/null +++ b/hair_service_sd/matting/networks/decoders/res_shortcut_dec.py @@ -0,0 +1,24 @@ +from 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 + diff --git a/hair_service_sd/matting/networks/decoders/resnet_dec.py b/hair_service_sd/matting/networks/decoders/resnet_dec.py new file mode 100644 index 0000000..ba36688 --- /dev/null +++ b/hair_service_sd/matting/networks/decoders/resnet_dec.py @@ -0,0 +1,142 @@ +import logging +import torch.nn as nn +from 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 diff --git a/hair_service_sd/matting/networks/encoders/__init__.py b/hair_service_sd/matting/networks/encoders/__init__.py new file mode 100644 index 0000000..71693a6 --- /dev/null +++ b/hair_service_sd/matting/networks/encoders/__init__.py @@ -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) diff --git a/hair_service_sd/matting/networks/encoders/res_gca_enc.py b/hair_service_sd/matting/networks/encoders/res_gca_enc.py new file mode 100644 index 0000000..25fe90d --- /dev/null +++ b/hair_service_sd/matting/networks/encoders/res_gca_enc.py @@ -0,0 +1,97 @@ +import torch.nn as nn +import torch.nn.functional as F + +# from utils import CONFIG +from matting.networks.encoders.resnet_enc import ResNet_D +from 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 matting.networks.encoders.resnet_enc import BasicBlock + m = ResGuidedCxtAtten(BasicBlock, [3, 4, 4, 2]) + for m in m.modules(): + print(m) diff --git a/hair_service_sd/matting/networks/encoders/res_shortcut_enc.py b/hair_service_sd/matting/networks/encoders/res_shortcut_enc.py new file mode 100644 index 0000000..99db4db --- /dev/null +++ b/hair_service_sd/matting/networks/encoders/res_shortcut_enc.py @@ -0,0 +1,51 @@ +import torch.nn as nn +# from utils import CONFIG +from matting.networks.encoders.resnet_enc import ResNet_D +from 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,...]} \ No newline at end of file diff --git a/hair_service_sd/matting/networks/encoders/resnet_enc.py b/hair_service_sd/matting/networks/encoders/resnet_enc.py new file mode 100644 index 0000000..729e173 --- /dev/null +++ b/hair_service_sd/matting/networks/encoders/resnet_enc.py @@ -0,0 +1,150 @@ +import logging +import torch.nn as nn +from 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()) diff --git a/hair_service_sd/matting/networks/generators.py b/hair_service_sd/matting/networks/generators.py new file mode 100644 index 0000000..9bde563 --- /dev/null +++ b/hair_service_sd/matting/networks/generators.py @@ -0,0 +1,60 @@ +import torch +import torch.nn as nn + +# from utils import CONFIG +from matting.networks import encoders, decoders + + +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) \ No newline at end of file diff --git a/hair_service_sd/matting/networks/ops.py b/hair_service_sd/matting/networks/ops.py new file mode 100644 index 0000000..a2df034 --- /dev/null +++ b/hair_service_sd/matting/networks/ops.py @@ -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) \ No newline at end of file diff --git a/hair_service_sd/matting/setup.py b/hair_service_sd/matting/setup.py new file mode 100644 index 0000000..ca6bf5a --- /dev/null +++ b/hair_service_sd/matting/setup.py @@ -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 \ No newline at end of file diff --git a/hair_service_sd/models/Generator_BaldSeg.py b/hair_service_sd/models/Generator_BaldSeg.py new file mode 100644 index 0000000..ca5625f --- /dev/null +++ b/hair_service_sd/models/Generator_BaldSeg.py @@ -0,0 +1,69 @@ +import os +import time +import torch +import torch.nn.parallel + +import numpy as np +from utils import landmark_processor +import cv2 + +# modelRoot = "/home/yangchaojie/Desktop/hairstyle/hairstyle_infer/weights" +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(self.model_dir, "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 \ No newline at end of file diff --git a/hair_service_sd/models/MomocvFaceAlignment1K.py b/hair_service_sd/models/MomocvFaceAlignment1K.py new file mode 100644 index 0000000..7fc7a46 --- /dev/null +++ b/hair_service_sd/models/MomocvFaceAlignment1K.py @@ -0,0 +1,429 @@ +import torch.nn as nn +import torch.utils.model_zoo as model_zoo +import torch +import numpy as np +import os +from utils import landmark_processor +from 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(): + input_numpy = np.zeros((1, 3, dst_size, dst_size), dtype=np.float32) + tmp = cv2.resize(img, (dst_size, dst_size)) + 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))) + orig_pts[:, 0] *= img.shape[1] + orig_pts[:, 1] *= img.shape[0] + 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.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # 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]]) + + 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 + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + + tmp = cv2.warpAffine(img, mat, (dst_size, dst_size)) + # 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 \ No newline at end of file diff --git a/hair_service_sd/models/box_utils.py b/hair_service_sd/models/box_utils.py new file mode 100644 index 0000000..1aa5fae --- /dev/null +++ b/hair_service_sd/models/box_utils.py @@ -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 == '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 diff --git a/hair_service_sd/models/common.py b/hair_service_sd/models/common.py new file mode 100644 index 0000000..3c4a0d7 --- /dev/null +++ b/hair_service_sd/models/common.py @@ -0,0 +1,97 @@ +# This file contains modules common to various models + + +from utils.utils import * + + +def DWConv(c1, c2, k=1, s=1, act=True): + # Depthwise convolution + return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act) + + +class Conv(nn.Module): + # Standard convolution + def __init__(self, c1, c2, k=1, s=1, g=1, act=True): # ch_in, ch_out, kernel, stride, groups + super(Conv, self).__init__() + p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # padding + self.conv = nn.Conv2d(c1, c2, k, s, p, groups=g, bias=False) + self.bn = nn.BatchNorm2d(c2) + self.act = nn.LeakyReLU(0.1, inplace=True) if act else nn.Identity() + + def forward(self, x): + return self.act(self.bn(self.conv(x))) + + def fuseforward(self, x): + return self.act(self.conv(x)) + + +class Bottleneck(nn.Module): + # Standard bottleneck + def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion + super(Bottleneck, self).__init__() + c_ = int(c2 * e) # hidden channels + self.cv1 = Conv(c1, c_, 1, 1) + self.cv2 = Conv(c_, c2, 3, 1, g=g) + self.add = shortcut and c1 == c2 + + def forward(self, x): + return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x)) + + +class BottleneckCSP(nn.Module): + # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks + def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion + super(BottleneckCSP, self).__init__() + c_ = int(c2 * e) # hidden channels + self.cv1 = Conv(c1, c_, 1, 1) + self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False) + self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False) + self.cv4 = Conv(c2, c2, 1, 1) + self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3) + self.act = nn.LeakyReLU(0.1, inplace=True) + self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)]) + + def forward(self, x): + y1 = self.cv3(self.m(self.cv1(x))) + y2 = self.cv2(x) + return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1)))) + + +class SPP(nn.Module): + # Spatial pyramid pooling layer used in YOLOv3-SPP + def __init__(self, c1, c2, k=(5, 9, 13)): + super(SPP, self).__init__() + c_ = c1 // 2 # hidden channels + self.cv1 = Conv(c1, c_, 1, 1) + self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1) + self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k]) + + def forward(self, x): + x = self.cv1(x) + return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1)) + + +class Flatten(nn.Module): + # Use after nn.AdaptiveAvgPool2d(1) to remove last 2 dimensions + def forward(self, x): + return x.view(x.size(0), -1) + + +class Focus(nn.Module): + # Focus wh information into c-space + def __init__(self, c1, c2, k=1): + super(Focus, self).__init__() + self.conv = Conv(c1 * 4, c2, k, 1) + + def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2) + return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)) + + +class Concat(nn.Module): + # Concatenate a list of tensors along dimension + def __init__(self, dimension=1): + super(Concat, self).__init__() + self.d = dimension + + def forward(self, x): + return torch.cat(x, self.d) diff --git a/hair_service_sd/models/config.py b/hair_service_sd/models/config.py new file mode 100644 index 0000000..591f349 --- /dev/null +++ b/hair_service_sd/models/config.py @@ -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 +} + diff --git a/hair_service_sd/models/detector.py b/hair_service_sd/models/detector.py new file mode 100644 index 0000000..0bba639 --- /dev/null +++ b/hair_service_sd/models/detector.py @@ -0,0 +1,353 @@ +import math +import numpy as np +from .model import PNet, RNet, ONet +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 utils import box_utils_Retina +from .layers.functions.prior_box import PriorBox +from .config import cfg_re50 +from .retinaface import RetinaFace + + +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 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 + + +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 diff --git a/hair_service_sd/models/experimental.py b/hair_service_sd/models/experimental.py new file mode 100644 index 0000000..60cb7aa --- /dev/null +++ b/hair_service_sd/models/experimental.py @@ -0,0 +1,85 @@ +from models.common import * + + +class Sum(nn.Module): + # Weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070 + def __init__(self, n, weight=False): # n: number of inputs + super(Sum, self).__init__() + self.weight = weight # apply weights boolean + self.iter = range(n - 1) # iter object + if weight: + self.w = nn.Parameter(-torch.arange(1., n) / 2, requires_grad=True) # layer weights + + def forward(self, x): + y = x[0] # no weight + if self.weight: + w = torch.sigmoid(self.w) * 2 + for i in self.iter: + y = y + x[i + 1] * w[i] + else: + for i in self.iter: + y = y + x[i + 1] + return y + + +class GhostConv(nn.Module): + # Ghost Convolution https://github.com/huawei-noah/ghostnet + def __init__(self, c1, c2, k=1, s=1, g=1, act=True): # ch_in, ch_out, kernel, stride, groups + super(GhostConv, self).__init__() + c_ = c2 // 2 # hidden channels + self.cv1 = Conv(c1, c_, k, s, g, act) + self.cv2 = Conv(c_, c_, 5, 1, c_, act) + + def forward(self, x): + y = self.cv1(x) + return torch.cat([y, self.cv2(y)], 1) + + +class GhostBottleneck(nn.Module): + # Ghost Bottleneck https://github.com/huawei-noah/ghostnet + def __init__(self, c1, c2, k, s): + super(GhostBottleneck, self).__init__() + c_ = c2 // 2 + self.conv = nn.Sequential(GhostConv(c1, c_, 1, 1), # pw + DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw + GhostConv(c_, c2, 1, 1, act=False)) # pw-linear + self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False), + Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity() + + def forward(self, x): + return self.conv(x) + self.shortcut(x) + + +class ConvPlus(nn.Module): + # Plus-shaped convolution + def __init__(self, c1, c2, k=3, s=1, g=1, bias=True): # ch_in, ch_out, kernel, stride, groups + super(ConvPlus, self).__init__() + self.cv1 = nn.Conv2d(c1, c2, (k, 1), s, (k // 2, 0), groups=g, bias=bias) + self.cv2 = nn.Conv2d(c1, c2, (1, k), s, (0, k // 2), groups=g, bias=bias) + + def forward(self, x): + return self.cv1(x) + self.cv2(x) + + +class MixConv2d(nn.Module): + # Mixed Depthwise Conv https://arxiv.org/abs/1907.09595 + def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True): + super(MixConv2d, self).__init__() + groups = len(k) + if equal_ch: # equal c_ per group + i = torch.linspace(0, groups - 1E-6, c2).floor() # c2 indices + c_ = [(i == g).sum() for g in range(groups)] # intermediate channels + else: # equal weight.numel() per group + b = [c2] + [0] * groups + a = np.eye(groups + 1, groups, k=-1) + a -= np.roll(a, 1, axis=1) + a *= np.array(k) ** 2 + a[0] = 1 + c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b + + self.m = nn.ModuleList([nn.Conv2d(c1, int(c_[g]), k[g], s, k[g] // 2, bias=False) for g in range(groups)]) + self.bn = nn.BatchNorm2d(c2) + self.act = nn.LeakyReLU(0.1, inplace=True) + + def forward(self, x): + return x + self.act(self.bn(torch.cat([m(x) for m in self.m], 1))) diff --git a/hair_service_sd/models/layers/__init__.py b/hair_service_sd/models/layers/__init__.py new file mode 100644 index 0000000..53a3f4b --- /dev/null +++ b/hair_service_sd/models/layers/__init__.py @@ -0,0 +1,2 @@ +from .functions import * +from .modules import * diff --git a/hair_service_sd/models/layers/data/FDDB/img_list.txt b/hair_service_sd/models/layers/data/FDDB/img_list.txt new file mode 100644 index 0000000..5cf3d31 --- /dev/null +++ b/hair_service_sd/models/layers/data/FDDB/img_list.txt @@ -0,0 +1,2845 @@ +2002/08/11/big/img_591 +2002/08/26/big/img_265 +2002/07/19/big/img_423 +2002/08/24/big/img_490 +2002/08/31/big/img_17676 +2002/07/31/big/img_228 +2002/07/24/big/img_402 +2002/08/04/big/img_769 +2002/07/19/big/img_581 +2002/08/13/big/img_723 +2002/08/12/big/img_821 +2003/01/17/big/img_610 +2002/08/13/big/img_1116 +2002/08/28/big/img_19238 +2002/08/21/big/img_660 +2002/08/14/big/img_607 +2002/08/05/big/img_3708 +2002/08/19/big/img_511 +2002/08/07/big/img_1316 +2002/07/25/big/img_1047 +2002/07/23/big/img_474 +2002/07/27/big/img_970 +2002/09/02/big/img_15752 +2002/09/01/big/img_16378 +2002/09/01/big/img_16189 +2002/08/26/big/img_276 +2002/07/24/big/img_518 +2002/08/14/big/img_1027 +2002/08/24/big/img_733 +2002/08/15/big/img_249 +2003/01/15/big/img_1371 +2002/08/07/big/img_1348 +2003/01/01/big/img_331 +2002/08/23/big/img_536 +2002/07/30/big/img_224 +2002/08/10/big/img_763 +2002/08/21/big/img_293 +2002/08/15/big/img_1211 +2002/08/15/big/img_1194 +2003/01/15/big/img_390 +2002/08/06/big/img_2893 +2002/08/17/big/img_691 +2002/08/07/big/img_1695 +2002/08/16/big/img_829 +2002/07/25/big/img_201 +2002/08/23/big/img_36 +2003/01/15/big/img_763 +2003/01/15/big/img_637 +2002/08/22/big/img_592 +2002/07/25/big/img_817 +2003/01/15/big/img_1219 +2002/08/05/big/img_3508 +2002/08/15/big/img_1108 +2002/07/19/big/img_488 +2003/01/16/big/img_704 +2003/01/13/big/img_1087 +2002/08/10/big/img_670 +2002/07/24/big/img_104 +2002/08/27/big/img_19823 +2002/09/01/big/img_16229 +2003/01/13/big/img_846 +2002/08/04/big/img_412 +2002/07/22/big/img_554 +2002/08/12/big/img_331 +2002/08/02/big/img_533 +2002/08/12/big/img_259 +2002/08/18/big/img_328 +2003/01/14/big/img_630 +2002/08/05/big/img_3541 +2002/08/06/big/img_2390 +2002/08/20/big/img_150 +2002/08/02/big/img_1231 +2002/08/16/big/img_710 +2002/08/19/big/img_591 +2002/07/22/big/img_725 +2002/07/24/big/img_820 +2003/01/13/big/img_568 +2002/08/22/big/img_853 +2002/08/09/big/img_648 +2002/08/23/big/img_528 +2003/01/14/big/img_888 +2002/08/30/big/img_18201 +2002/08/13/big/img_965 +2003/01/14/big/img_660 +2002/07/19/big/img_517 +2003/01/14/big/img_406 +2002/08/30/big/img_18433 +2002/08/07/big/img_1630 +2002/08/06/big/img_2717 +2002/08/21/big/img_470 +2002/07/23/big/img_633 +2002/08/20/big/img_915 +2002/08/16/big/img_893 +2002/07/29/big/img_644 +2002/08/15/big/img_529 +2002/08/16/big/img_668 +2002/08/07/big/img_1871 +2002/07/25/big/img_192 +2002/07/31/big/img_961 +2002/08/19/big/img_738 +2002/07/31/big/img_382 +2002/08/19/big/img_298 +2003/01/17/big/img_608 +2002/08/21/big/img_514 +2002/07/23/big/img_183 +2003/01/17/big/img_536 +2002/07/24/big/img_478 +2002/08/06/big/img_2997 +2002/09/02/big/img_15380 +2002/08/07/big/img_1153 +2002/07/31/big/img_967 +2002/07/31/big/img_711 +2002/08/26/big/img_664 +2003/01/01/big/img_326 +2002/08/24/big/img_775 +2002/08/08/big/img_961 +2002/08/16/big/img_77 +2002/08/12/big/img_296 +2002/07/22/big/img_905 +2003/01/13/big/img_284 +2002/08/13/big/img_887 +2002/08/24/big/img_849 +2002/07/30/big/img_345 +2002/08/18/big/img_419 +2002/08/01/big/img_1347 +2002/08/05/big/img_3670 +2002/07/21/big/img_479 +2002/08/08/big/img_913 +2002/09/02/big/img_15828 +2002/08/30/big/img_18194 +2002/08/08/big/img_471 +2002/08/22/big/img_734 +2002/08/09/big/img_586 +2002/08/09/big/img_454 +2002/07/29/big/img_47 +2002/07/19/big/img_381 +2002/07/29/big/img_733 +2002/08/20/big/img_327 +2002/07/21/big/img_96 +2002/08/06/big/img_2680 +2002/07/25/big/img_919 +2002/07/21/big/img_158 +2002/07/22/big/img_801 +2002/07/22/big/img_567 +2002/07/24/big/img_804 +2002/07/24/big/img_690 +2003/01/15/big/img_576 +2002/08/14/big/img_335 +2003/01/13/big/img_390 +2002/08/11/big/img_258 +2002/07/23/big/img_917 +2002/08/15/big/img_525 +2003/01/15/big/img_505 +2002/07/30/big/img_886 +2003/01/16/big/img_640 +2003/01/14/big/img_642 +2003/01/17/big/img_844 +2002/08/04/big/img_571 +2002/08/29/big/img_18702 +2003/01/15/big/img_240 +2002/07/29/big/img_553 +2002/08/10/big/img_354 +2002/08/18/big/img_17 +2003/01/15/big/img_782 +2002/07/27/big/img_382 +2002/08/14/big/img_970 +2003/01/16/big/img_70 +2003/01/16/big/img_625 +2002/08/18/big/img_341 +2002/08/26/big/img_188 +2002/08/09/big/img_405 +2002/08/02/big/img_37 +2002/08/13/big/img_748 +2002/07/22/big/img_399 +2002/07/25/big/img_844 +2002/08/12/big/img_340 +2003/01/13/big/img_815 +2002/08/26/big/img_5 +2002/08/10/big/img_158 +2002/08/18/big/img_95 +2002/07/29/big/img_1297 +2003/01/13/big/img_508 +2002/09/01/big/img_16680 +2003/01/16/big/img_338 +2002/08/13/big/img_517 +2002/07/22/big/img_626 +2002/08/06/big/img_3024 +2002/07/26/big/img_499 +2003/01/13/big/img_387 +2002/08/31/big/img_18025 +2002/08/13/big/img_520 +2003/01/16/big/img_576 +2002/07/26/big/img_121 +2002/08/25/big/img_703 +2002/08/26/big/img_615 +2002/08/17/big/img_434 +2002/08/02/big/img_677 +2002/08/18/big/img_276 +2002/08/05/big/img_3672 +2002/07/26/big/img_700 +2002/07/31/big/img_277 +2003/01/14/big/img_220 +2002/08/23/big/img_232 +2002/08/31/big/img_17422 +2002/07/22/big/img_508 +2002/08/13/big/img_681 +2003/01/15/big/img_638 +2002/08/30/big/img_18408 +2003/01/14/big/img_533 +2003/01/17/big/img_12 +2002/08/28/big/img_19388 +2002/08/08/big/img_133 +2002/07/26/big/img_885 +2002/08/19/big/img_387 +2002/08/27/big/img_19976 +2002/08/26/big/img_118 +2002/08/28/big/img_19146 +2002/08/05/big/img_3259 +2002/08/15/big/img_536 +2002/07/22/big/img_279 +2002/07/22/big/img_9 +2002/08/13/big/img_301 +2002/08/15/big/img_974 +2002/08/06/big/img_2355 +2002/08/01/big/img_1526 +2002/08/03/big/img_417 +2002/08/04/big/img_407 +2002/08/15/big/img_1029 +2002/07/29/big/img_700 +2002/08/01/big/img_1463 +2002/08/31/big/img_17365 +2002/07/28/big/img_223 +2002/07/19/big/img_827 +2002/07/27/big/img_531 +2002/07/19/big/img_845 +2002/08/20/big/img_382 +2002/07/31/big/img_268 +2002/08/27/big/img_19705 +2002/08/02/big/img_830 +2002/08/23/big/img_250 +2002/07/20/big/img_777 +2002/08/21/big/img_879 +2002/08/26/big/img_20146 +2002/08/23/big/img_789 +2002/08/06/big/img_2683 +2002/08/25/big/img_576 +2002/08/09/big/img_498 +2002/08/08/big/img_384 +2002/08/26/big/img_592 +2002/07/29/big/img_1470 +2002/08/21/big/img_452 +2002/08/30/big/img_18395 +2002/08/15/big/img_215 +2002/07/21/big/img_643 +2002/07/22/big/img_209 +2003/01/17/big/img_346 +2002/08/25/big/img_658 +2002/08/21/big/img_221 +2002/08/14/big/img_60 +2003/01/17/big/img_885 +2003/01/16/big/img_482 +2002/08/19/big/img_593 +2002/08/08/big/img_233 +2002/07/30/big/img_458 +2002/07/23/big/img_384 +2003/01/15/big/img_670 +2003/01/15/big/img_267 +2002/08/26/big/img_540 +2002/07/29/big/img_552 +2002/07/30/big/img_997 +2003/01/17/big/img_377 +2002/08/21/big/img_265 +2002/08/09/big/img_561 +2002/07/31/big/img_945 +2002/09/02/big/img_15252 +2002/08/11/big/img_276 +2002/07/22/big/img_491 +2002/07/26/big/img_517 +2002/08/14/big/img_726 +2002/08/08/big/img_46 +2002/08/28/big/img_19458 +2002/08/06/big/img_2935 +2002/07/29/big/img_1392 +2002/08/13/big/img_776 +2002/08/24/big/img_616 +2002/08/14/big/img_1065 +2002/07/29/big/img_889 +2002/08/18/big/img_188 +2002/08/07/big/img_1453 +2002/08/02/big/img_760 +2002/07/28/big/img_416 +2002/08/07/big/img_1393 +2002/08/26/big/img_292 +2002/08/26/big/img_301 +2003/01/13/big/img_195 +2002/07/26/big/img_532 +2002/08/20/big/img_550 +2002/08/05/big/img_3658 +2002/08/26/big/img_738 +2002/09/02/big/img_15750 +2003/01/17/big/img_451 +2002/07/23/big/img_339 +2002/08/16/big/img_637 +2002/08/14/big/img_748 +2002/08/06/big/img_2739 +2002/07/25/big/img_482 +2002/08/19/big/img_191 +2002/08/26/big/img_537 +2003/01/15/big/img_716 +2003/01/15/big/img_767 +2002/08/02/big/img_452 +2002/08/08/big/img_1011 +2002/08/10/big/img_144 +2003/01/14/big/img_122 +2002/07/24/big/img_586 +2002/07/24/big/img_762 +2002/08/20/big/img_369 +2002/07/30/big/img_146 +2002/08/23/big/img_396 +2003/01/15/big/img_200 +2002/08/15/big/img_1183 +2003/01/14/big/img_698 +2002/08/09/big/img_792 +2002/08/06/big/img_2347 +2002/07/31/big/img_911 +2002/08/26/big/img_722 +2002/08/23/big/img_621 +2002/08/05/big/img_3790 +2003/01/13/big/img_633 +2002/08/09/big/img_224 +2002/07/24/big/img_454 +2002/07/21/big/img_202 +2002/08/02/big/img_630 +2002/08/30/big/img_18315 +2002/07/19/big/img_491 +2002/09/01/big/img_16456 +2002/08/09/big/img_242 +2002/07/25/big/img_595 +2002/07/22/big/img_522 +2002/08/01/big/img_1593 +2002/07/29/big/img_336 +2002/08/15/big/img_448 +2002/08/28/big/img_19281 +2002/07/29/big/img_342 +2002/08/12/big/img_78 +2003/01/14/big/img_525 +2002/07/28/big/img_147 +2002/08/11/big/img_353 +2002/08/22/big/img_513 +2002/08/04/big/img_721 +2002/08/17/big/img_247 +2003/01/14/big/img_891 +2002/08/20/big/img_853 +2002/07/19/big/img_414 +2002/08/01/big/img_1530 +2003/01/14/big/img_924 +2002/08/22/big/img_468 +2002/08/18/big/img_354 +2002/08/30/big/img_18193 +2002/08/23/big/img_492 +2002/08/15/big/img_871 +2002/08/12/big/img_494 +2002/08/06/big/img_2470 +2002/07/23/big/img_923 +2002/08/26/big/img_155 +2002/08/08/big/img_669 +2002/07/23/big/img_404 +2002/08/28/big/img_19421 +2002/08/29/big/img_18993 +2002/08/25/big/img_416 +2003/01/17/big/img_434 +2002/07/29/big/img_1370 +2002/07/28/big/img_483 +2002/08/11/big/img_50 +2002/08/10/big/img_404 +2002/09/02/big/img_15057 +2003/01/14/big/img_911 +2002/09/01/big/img_16697 +2003/01/16/big/img_665 +2002/09/01/big/img_16708 +2002/08/22/big/img_612 +2002/08/28/big/img_19471 +2002/08/02/big/img_198 +2003/01/16/big/img_527 +2002/08/22/big/img_209 +2002/08/30/big/img_18205 +2003/01/14/big/img_114 +2003/01/14/big/img_1028 +2003/01/16/big/img_894 +2003/01/14/big/img_837 +2002/07/30/big/img_9 +2002/08/06/big/img_2821 +2002/08/04/big/img_85 +2003/01/13/big/img_884 +2002/07/22/big/img_570 +2002/08/07/big/img_1773 +2002/07/26/big/img_208 +2003/01/17/big/img_946 +2002/07/19/big/img_930 +2003/01/01/big/img_698 +2003/01/17/big/img_612 +2002/07/19/big/img_372 +2002/07/30/big/img_721 +2003/01/14/big/img_649 +2002/08/19/big/img_4 +2002/07/25/big/img_1024 +2003/01/15/big/img_601 +2002/08/30/big/img_18470 +2002/07/22/big/img_29 +2002/08/07/big/img_1686 +2002/07/20/big/img_294 +2002/08/14/big/img_800 +2002/08/19/big/img_353 +2002/08/19/big/img_350 +2002/08/05/big/img_3392 +2002/08/09/big/img_622 +2003/01/15/big/img_236 +2002/08/11/big/img_643 +2002/08/05/big/img_3458 +2002/08/12/big/img_413 +2002/08/22/big/img_415 +2002/08/13/big/img_635 +2002/08/07/big/img_1198 +2002/08/04/big/img_873 +2002/08/12/big/img_407 +2003/01/15/big/img_346 +2002/08/02/big/img_275 +2002/08/17/big/img_997 +2002/08/21/big/img_958 +2002/08/20/big/img_579 +2002/07/29/big/img_142 +2003/01/14/big/img_1115 +2002/08/16/big/img_365 +2002/07/29/big/img_1414 +2002/08/17/big/img_489 +2002/08/13/big/img_1010 +2002/07/31/big/img_276 +2002/07/25/big/img_1000 +2002/08/23/big/img_524 +2002/08/28/big/img_19147 +2003/01/13/big/img_433 +2002/08/20/big/img_205 +2003/01/01/big/img_458 +2002/07/29/big/img_1449 +2003/01/16/big/img_696 +2002/08/28/big/img_19296 +2002/08/29/big/img_18688 +2002/08/21/big/img_767 +2002/08/20/big/img_532 +2002/08/26/big/img_187 +2002/07/26/big/img_183 +2002/07/27/big/img_890 +2003/01/13/big/img_576 +2002/07/30/big/img_15 +2002/07/31/big/img_889 +2002/08/31/big/img_17759 +2003/01/14/big/img_1114 +2002/07/19/big/img_445 +2002/08/03/big/img_593 +2002/07/24/big/img_750 +2002/07/30/big/img_133 +2002/08/25/big/img_671 +2002/07/20/big/img_351 +2002/08/31/big/img_17276 +2002/08/05/big/img_3231 +2002/09/02/big/img_15882 +2002/08/14/big/img_115 +2002/08/02/big/img_1148 +2002/07/25/big/img_936 +2002/07/31/big/img_639 +2002/08/04/big/img_427 +2002/08/22/big/img_843 +2003/01/17/big/img_17 +2003/01/13/big/img_690 +2002/08/13/big/img_472 +2002/08/09/big/img_425 +2002/08/05/big/img_3450 +2003/01/17/big/img_439 +2002/08/13/big/img_539 +2002/07/28/big/img_35 +2002/08/16/big/img_241 +2002/08/06/big/img_2898 +2003/01/16/big/img_429 +2002/08/05/big/img_3817 +2002/08/27/big/img_19919 +2002/07/19/big/img_422 +2002/08/15/big/img_560 +2002/07/23/big/img_750 +2002/07/30/big/img_353 +2002/08/05/big/img_43 +2002/08/23/big/img_305 +2002/08/01/big/img_2137 +2002/08/30/big/img_18097 +2002/08/01/big/img_1389 +2002/08/02/big/img_308 +2003/01/14/big/img_652 +2002/08/01/big/img_1798 +2003/01/14/big/img_732 +2003/01/16/big/img_294 +2002/08/26/big/img_213 +2002/07/24/big/img_842 +2003/01/13/big/img_630 +2003/01/13/big/img_634 +2002/08/06/big/img_2285 +2002/08/01/big/img_2162 +2002/08/30/big/img_18134 +2002/08/02/big/img_1045 +2002/08/01/big/img_2143 +2002/07/25/big/img_135 +2002/07/20/big/img_645 +2002/08/05/big/img_3666 +2002/08/14/big/img_523 +2002/08/04/big/img_425 +2003/01/14/big/img_137 +2003/01/01/big/img_176 +2002/08/15/big/img_505 +2002/08/24/big/img_386 +2002/08/05/big/img_3187 +2002/08/15/big/img_419 +2003/01/13/big/img_520 +2002/08/04/big/img_444 +2002/08/26/big/img_483 +2002/08/05/big/img_3449 +2002/08/30/big/img_18409 +2002/08/28/big/img_19455 +2002/08/27/big/img_20090 +2002/07/23/big/img_625 +2002/08/24/big/img_205 +2002/08/08/big/img_938 +2003/01/13/big/img_527 +2002/08/07/big/img_1712 +2002/07/24/big/img_801 +2002/08/09/big/img_579 +2003/01/14/big/img_41 +2003/01/15/big/img_1130 +2002/07/21/big/img_672 +2002/08/07/big/img_1590 +2003/01/01/big/img_532 +2002/08/02/big/img_529 +2002/08/05/big/img_3591 +2002/08/23/big/img_5 +2003/01/14/big/img_882 +2002/08/28/big/img_19234 +2002/07/24/big/img_398 +2003/01/14/big/img_592 +2002/08/22/big/img_548 +2002/08/12/big/img_761 +2003/01/16/big/img_497 +2002/08/18/big/img_133 +2002/08/08/big/img_874 +2002/07/19/big/img_247 +2002/08/15/big/img_170 +2002/08/27/big/img_19679 +2002/08/20/big/img_246 +2002/08/24/big/img_358 +2002/07/29/big/img_599 +2002/08/01/big/img_1555 +2002/07/30/big/img_491 +2002/07/30/big/img_371 +2003/01/16/big/img_682 +2002/07/25/big/img_619 +2003/01/15/big/img_587 +2002/08/02/big/img_1212 +2002/08/01/big/img_2152 +2002/07/25/big/img_668 +2003/01/16/big/img_574 +2002/08/28/big/img_19464 +2002/08/11/big/img_536 +2002/07/24/big/img_201 +2002/08/05/big/img_3488 +2002/07/25/big/img_887 +2002/07/22/big/img_789 +2002/07/30/big/img_432 +2002/08/16/big/img_166 +2002/09/01/big/img_16333 +2002/07/26/big/img_1010 +2002/07/21/big/img_793 +2002/07/22/big/img_720 +2002/07/31/big/img_337 +2002/07/27/big/img_185 +2002/08/23/big/img_440 +2002/07/31/big/img_801 +2002/07/25/big/img_478 +2003/01/14/big/img_171 +2002/08/07/big/img_1054 +2002/09/02/big/img_15659 +2002/07/29/big/img_1348 +2002/08/09/big/img_337 +2002/08/26/big/img_684 +2002/07/31/big/img_537 +2002/08/15/big/img_808 +2003/01/13/big/img_740 +2002/08/07/big/img_1667 +2002/08/03/big/img_404 +2002/08/06/big/img_2520 +2002/07/19/big/img_230 +2002/07/19/big/img_356 +2003/01/16/big/img_627 +2002/08/04/big/img_474 +2002/07/29/big/img_833 +2002/07/25/big/img_176 +2002/08/01/big/img_1684 +2002/08/21/big/img_643 +2002/08/27/big/img_19673 +2002/08/02/big/img_838 +2002/08/06/big/img_2378 +2003/01/15/big/img_48 +2002/07/30/big/img_470 +2002/08/15/big/img_963 +2002/08/24/big/img_444 +2002/08/16/big/img_662 +2002/08/15/big/img_1209 +2002/07/24/big/img_25 +2002/08/06/big/img_2740 +2002/07/29/big/img_996 +2002/08/31/big/img_18074 +2002/08/04/big/img_343 +2003/01/17/big/img_509 +2003/01/13/big/img_726 +2002/08/07/big/img_1466 +2002/07/26/big/img_307 +2002/08/10/big/img_598 +2002/08/13/big/img_890 +2002/08/14/big/img_997 +2002/07/19/big/img_392 +2002/08/02/big/img_475 +2002/08/29/big/img_19038 +2002/07/29/big/img_538 +2002/07/29/big/img_502 +2002/08/02/big/img_364 +2002/08/31/big/img_17353 +2002/08/08/big/img_539 +2002/08/01/big/img_1449 +2002/07/22/big/img_363 +2002/08/02/big/img_90 +2002/09/01/big/img_16867 +2002/08/05/big/img_3371 +2002/07/30/big/img_342 +2002/08/07/big/img_1363 +2002/08/22/big/img_790 +2003/01/15/big/img_404 +2002/08/05/big/img_3447 +2002/09/01/big/img_16167 +2003/01/13/big/img_840 +2002/08/22/big/img_1001 +2002/08/09/big/img_431 +2002/07/27/big/img_618 +2002/07/31/big/img_741 +2002/07/30/big/img_964 +2002/07/25/big/img_86 +2002/07/29/big/img_275 +2002/08/21/big/img_921 +2002/07/26/big/img_892 +2002/08/21/big/img_663 +2003/01/13/big/img_567 +2003/01/14/big/img_719 +2002/07/28/big/img_251 +2003/01/15/big/img_1123 +2002/07/29/big/img_260 +2002/08/24/big/img_337 +2002/08/01/big/img_1914 +2002/08/13/big/img_373 +2003/01/15/big/img_589 +2002/08/13/big/img_906 +2002/07/26/big/img_270 +2002/08/26/big/img_313 +2002/08/25/big/img_694 +2003/01/01/big/img_327 +2002/07/23/big/img_261 +2002/08/26/big/img_642 +2002/07/29/big/img_918 +2002/07/23/big/img_455 +2002/07/24/big/img_612 +2002/07/23/big/img_534 +2002/07/19/big/img_534 +2002/07/19/big/img_726 +2002/08/01/big/img_2146 +2002/08/02/big/img_543 +2003/01/16/big/img_777 +2002/07/30/big/img_484 +2002/08/13/big/img_1161 +2002/07/21/big/img_390 +2002/08/06/big/img_2288 +2002/08/21/big/img_677 +2002/08/13/big/img_747 +2002/08/15/big/img_1248 +2002/07/31/big/img_416 +2002/09/02/big/img_15259 +2002/08/16/big/img_781 +2002/08/24/big/img_754 +2002/07/24/big/img_803 +2002/08/20/big/img_609 +2002/08/28/big/img_19571 +2002/09/01/big/img_16140 +2002/08/26/big/img_769 +2002/07/20/big/img_588 +2002/08/02/big/img_898 +2002/07/21/big/img_466 +2002/08/14/big/img_1046 +2002/07/25/big/img_212 +2002/08/26/big/img_353 +2002/08/19/big/img_810 +2002/08/31/big/img_17824 +2002/08/12/big/img_631 +2002/07/19/big/img_828 +2002/07/24/big/img_130 +2002/08/25/big/img_580 +2002/07/31/big/img_699 +2002/07/23/big/img_808 +2002/07/31/big/img_377 +2003/01/16/big/img_570 +2002/09/01/big/img_16254 +2002/07/21/big/img_471 +2002/08/01/big/img_1548 +2002/08/18/big/img_252 +2002/08/19/big/img_576 +2002/08/20/big/img_464 +2002/07/27/big/img_735 +2002/08/21/big/img_589 +2003/01/15/big/img_1192 +2002/08/09/big/img_302 +2002/07/31/big/img_594 +2002/08/23/big/img_19 +2002/08/29/big/img_18819 +2002/08/19/big/img_293 +2002/07/30/big/img_331 +2002/08/23/big/img_607 +2002/07/30/big/img_363 +2002/08/16/big/img_766 +2003/01/13/big/img_481 +2002/08/06/big/img_2515 +2002/09/02/big/img_15913 +2002/09/02/big/img_15827 +2002/09/02/big/img_15053 +2002/08/07/big/img_1576 +2002/07/23/big/img_268 +2002/08/21/big/img_152 +2003/01/15/big/img_578 +2002/07/21/big/img_589 +2002/07/20/big/img_548 +2002/08/27/big/img_19693 +2002/08/31/big/img_17252 +2002/07/31/big/img_138 +2002/07/23/big/img_372 +2002/08/16/big/img_695 +2002/07/27/big/img_287 +2002/08/15/big/img_315 +2002/08/10/big/img_361 +2002/07/29/big/img_899 +2002/08/13/big/img_771 +2002/08/21/big/img_92 +2003/01/15/big/img_425 +2003/01/16/big/img_450 +2002/09/01/big/img_16942 +2002/08/02/big/img_51 +2002/09/02/big/img_15379 +2002/08/24/big/img_147 +2002/08/30/big/img_18122 +2002/07/26/big/img_950 +2002/08/07/big/img_1400 +2002/08/17/big/img_468 +2002/08/15/big/img_470 +2002/07/30/big/img_318 +2002/07/22/big/img_644 +2002/08/27/big/img_19732 +2002/07/23/big/img_601 +2002/08/26/big/img_398 +2002/08/21/big/img_428 +2002/08/06/big/img_2119 +2002/08/29/big/img_19103 +2003/01/14/big/img_933 +2002/08/11/big/img_674 +2002/08/28/big/img_19420 +2002/08/03/big/img_418 +2002/08/17/big/img_312 +2002/07/25/big/img_1044 +2003/01/17/big/img_671 +2002/08/30/big/img_18297 +2002/07/25/big/img_755 +2002/07/23/big/img_471 +2002/08/21/big/img_39 +2002/07/26/big/img_699 +2003/01/14/big/img_33 +2002/07/31/big/img_411 +2002/08/16/big/img_645 +2003/01/17/big/img_116 +2002/09/02/big/img_15903 +2002/08/20/big/img_120 +2002/08/22/big/img_176 +2002/07/29/big/img_1316 +2002/08/27/big/img_19914 +2002/07/22/big/img_719 +2002/08/28/big/img_19239 +2003/01/13/big/img_385 +2002/08/08/big/img_525 +2002/07/19/big/img_782 +2002/08/13/big/img_843 +2002/07/30/big/img_107 +2002/08/11/big/img_752 +2002/07/29/big/img_383 +2002/08/26/big/img_249 +2002/08/29/big/img_18860 +2002/07/30/big/img_70 +2002/07/26/big/img_194 +2002/08/15/big/img_530 +2002/08/08/big/img_816 +2002/07/31/big/img_286 +2003/01/13/big/img_294 +2002/07/31/big/img_251 +2002/07/24/big/img_13 +2002/08/31/big/img_17938 +2002/07/22/big/img_642 +2003/01/14/big/img_728 +2002/08/18/big/img_47 +2002/08/22/big/img_306 +2002/08/20/big/img_348 +2002/08/15/big/img_764 +2002/08/08/big/img_163 +2002/07/23/big/img_531 +2002/07/23/big/img_467 +2003/01/16/big/img_743 +2003/01/13/big/img_535 +2002/08/02/big/img_523 +2002/08/22/big/img_120 +2002/08/11/big/img_496 +2002/08/29/big/img_19075 +2002/08/08/big/img_465 +2002/08/09/big/img_790 +2002/08/19/big/img_588 +2002/08/23/big/img_407 +2003/01/17/big/img_435 +2002/08/24/big/img_398 +2002/08/27/big/img_19899 +2003/01/15/big/img_335 +2002/08/13/big/img_493 +2002/09/02/big/img_15460 +2002/07/31/big/img_470 +2002/08/05/big/img_3550 +2002/07/28/big/img_123 +2002/08/01/big/img_1498 +2002/08/04/big/img_504 +2003/01/17/big/img_427 +2002/08/27/big/img_19708 +2002/07/27/big/img_861 +2002/07/25/big/img_685 +2002/07/31/big/img_207 +2003/01/14/big/img_745 +2002/08/31/big/img_17756 +2002/08/24/big/img_288 +2002/08/18/big/img_181 +2002/08/10/big/img_520 +2002/08/25/big/img_705 +2002/08/23/big/img_226 +2002/08/04/big/img_727 +2002/07/24/big/img_625 +2002/08/28/big/img_19157 +2002/08/23/big/img_586 +2002/07/31/big/img_232 +2003/01/13/big/img_240 +2003/01/14/big/img_321 +2003/01/15/big/img_533 +2002/07/23/big/img_480 +2002/07/24/big/img_371 +2002/08/21/big/img_702 +2002/08/31/big/img_17075 +2002/09/02/big/img_15278 +2002/07/29/big/img_246 +2003/01/15/big/img_829 +2003/01/15/big/img_1213 +2003/01/16/big/img_441 +2002/08/14/big/img_921 +2002/07/23/big/img_425 +2002/08/15/big/img_296 +2002/07/19/big/img_135 +2002/07/26/big/img_402 +2003/01/17/big/img_88 +2002/08/20/big/img_872 +2002/08/13/big/img_1110 +2003/01/16/big/img_1040 +2002/07/23/big/img_9 +2002/08/13/big/img_700 +2002/08/16/big/img_371 +2002/08/27/big/img_19966 +2003/01/17/big/img_391 +2002/08/18/big/img_426 +2002/08/01/big/img_1618 +2002/07/21/big/img_754 +2003/01/14/big/img_1101 +2003/01/16/big/img_1022 +2002/07/22/big/img_275 +2002/08/24/big/img_86 +2002/08/17/big/img_582 +2003/01/15/big/img_765 +2003/01/17/big/img_449 +2002/07/28/big/img_265 +2003/01/13/big/img_552 +2002/07/28/big/img_115 +2003/01/16/big/img_56 +2002/08/02/big/img_1232 +2003/01/17/big/img_925 +2002/07/22/big/img_445 +2002/07/25/big/img_957 +2002/07/20/big/img_589 +2002/08/31/big/img_17107 +2002/07/29/big/img_483 +2002/08/14/big/img_1063 +2002/08/07/big/img_1545 +2002/08/14/big/img_680 +2002/09/01/big/img_16694 +2002/08/14/big/img_257 +2002/08/11/big/img_726 +2002/07/26/big/img_681 +2002/07/25/big/img_481 +2003/01/14/big/img_737 +2002/08/28/big/img_19480 +2003/01/16/big/img_362 +2002/08/27/big/img_19865 +2003/01/01/big/img_547 +2002/09/02/big/img_15074 +2002/08/01/big/img_1453 +2002/08/22/big/img_594 +2002/08/28/big/img_19263 +2002/08/13/big/img_478 +2002/07/29/big/img_1358 +2003/01/14/big/img_1022 +2002/08/16/big/img_450 +2002/08/02/big/img_159 +2002/07/26/big/img_781 +2003/01/13/big/img_601 +2002/08/20/big/img_407 +2002/08/15/big/img_468 +2002/08/31/big/img_17902 +2002/08/16/big/img_81 +2002/07/25/big/img_987 +2002/07/25/big/img_500 +2002/08/02/big/img_31 +2002/08/18/big/img_538 +2002/08/08/big/img_54 +2002/07/23/big/img_686 +2002/07/24/big/img_836 +2003/01/17/big/img_734 +2002/08/16/big/img_1055 +2003/01/16/big/img_521 +2002/07/25/big/img_612 +2002/08/22/big/img_778 +2002/08/03/big/img_251 +2002/08/12/big/img_436 +2002/08/23/big/img_705 +2002/07/28/big/img_243 +2002/07/25/big/img_1029 +2002/08/20/big/img_287 +2002/08/29/big/img_18739 +2002/08/05/big/img_3272 +2002/07/27/big/img_214 +2003/01/14/big/img_5 +2002/08/01/big/img_1380 +2002/08/29/big/img_19097 +2002/07/30/big/img_486 +2002/08/29/big/img_18707 +2002/08/10/big/img_559 +2002/08/15/big/img_365 +2002/08/09/big/img_525 +2002/08/10/big/img_689 +2002/07/25/big/img_502 +2002/08/03/big/img_667 +2002/08/10/big/img_855 +2002/08/10/big/img_706 +2002/08/18/big/img_603 +2003/01/16/big/img_1055 +2002/08/31/big/img_17890 +2002/08/15/big/img_761 +2003/01/15/big/img_489 +2002/08/26/big/img_351 +2002/08/01/big/img_1772 +2002/08/31/big/img_17729 +2002/07/25/big/img_609 +2003/01/13/big/img_539 +2002/07/27/big/img_686 +2002/07/31/big/img_311 +2002/08/22/big/img_799 +2003/01/16/big/img_936 +2002/08/31/big/img_17813 +2002/08/04/big/img_862 +2002/08/09/big/img_332 +2002/07/20/big/img_148 +2002/08/12/big/img_426 +2002/07/24/big/img_69 +2002/07/27/big/img_685 +2002/08/02/big/img_480 +2002/08/26/big/img_154 +2002/07/24/big/img_598 +2002/08/01/big/img_1881 +2002/08/20/big/img_667 +2003/01/14/big/img_495 +2002/07/21/big/img_744 +2002/07/30/big/img_150 +2002/07/23/big/img_924 +2002/08/08/big/img_272 +2002/07/23/big/img_310 +2002/07/25/big/img_1011 +2002/09/02/big/img_15725 +2002/07/19/big/img_814 +2002/08/20/big/img_936 +2002/07/25/big/img_85 +2002/08/24/big/img_662 +2002/08/09/big/img_495 +2003/01/15/big/img_196 +2002/08/16/big/img_707 +2002/08/28/big/img_19370 +2002/08/06/big/img_2366 +2002/08/06/big/img_3012 +2002/08/01/big/img_1452 +2002/07/31/big/img_742 +2002/07/27/big/img_914 +2003/01/13/big/img_290 +2002/07/31/big/img_288 +2002/08/02/big/img_171 +2002/08/22/big/img_191 +2002/07/27/big/img_1066 +2002/08/12/big/img_383 +2003/01/17/big/img_1018 +2002/08/01/big/img_1785 +2002/08/11/big/img_390 +2002/08/27/big/img_20037 +2002/08/12/big/img_38 +2003/01/15/big/img_103 +2002/08/26/big/img_31 +2002/08/18/big/img_660 +2002/07/22/big/img_694 +2002/08/15/big/img_24 +2002/07/27/big/img_1077 +2002/08/01/big/img_1943 +2002/07/22/big/img_292 +2002/09/01/big/img_16857 +2002/07/22/big/img_892 +2003/01/14/big/img_46 +2002/08/09/big/img_469 +2002/08/09/big/img_414 +2003/01/16/big/img_40 +2002/08/28/big/img_19231 +2002/07/27/big/img_978 +2002/07/23/big/img_475 +2002/07/25/big/img_92 +2002/08/09/big/img_799 +2002/07/25/big/img_491 +2002/08/03/big/img_654 +2003/01/15/big/img_687 +2002/08/11/big/img_478 +2002/08/07/big/img_1664 +2002/08/20/big/img_362 +2002/08/01/big/img_1298 +2003/01/13/big/img_500 +2002/08/06/big/img_2896 +2002/08/30/big/img_18529 +2002/08/16/big/img_1020 +2002/07/29/big/img_892 +2002/08/29/big/img_18726 +2002/07/21/big/img_453 +2002/08/17/big/img_437 +2002/07/19/big/img_665 +2002/07/22/big/img_440 +2002/07/19/big/img_582 +2002/07/21/big/img_233 +2003/01/01/big/img_82 +2002/07/25/big/img_341 +2002/07/29/big/img_864 +2002/08/02/big/img_276 +2002/08/29/big/img_18654 +2002/07/27/big/img_1024 +2002/08/19/big/img_373 +2003/01/15/big/img_241 +2002/07/25/big/img_84 +2002/08/13/big/img_834 +2002/08/10/big/img_511 +2002/08/01/big/img_1627 +2002/08/08/big/img_607 +2002/08/06/big/img_2083 +2002/08/01/big/img_1486 +2002/08/08/big/img_700 +2002/08/01/big/img_1954 +2002/08/21/big/img_54 +2002/07/30/big/img_847 +2002/08/28/big/img_19169 +2002/07/21/big/img_549 +2002/08/03/big/img_693 +2002/07/31/big/img_1002 +2003/01/14/big/img_1035 +2003/01/16/big/img_622 +2002/07/30/big/img_1201 +2002/08/10/big/img_444 +2002/07/31/big/img_374 +2002/08/21/big/img_301 +2002/08/13/big/img_1095 +2003/01/13/big/img_288 +2002/07/25/big/img_232 +2003/01/13/big/img_967 +2002/08/26/big/img_360 +2002/08/05/big/img_67 +2002/08/29/big/img_18969 +2002/07/28/big/img_16 +2002/08/16/big/img_515 +2002/07/20/big/img_708 +2002/08/18/big/img_178 +2003/01/15/big/img_509 +2002/07/25/big/img_430 +2002/08/21/big/img_738 +2002/08/16/big/img_886 +2002/09/02/big/img_15605 +2002/09/01/big/img_16242 +2002/08/24/big/img_711 +2002/07/25/big/img_90 +2002/08/09/big/img_491 +2002/07/30/big/img_534 +2003/01/13/big/img_474 +2002/08/25/big/img_510 +2002/08/15/big/img_555 +2002/08/02/big/img_775 +2002/07/23/big/img_975 +2002/08/19/big/img_229 +2003/01/17/big/img_860 +2003/01/02/big/img_10 +2002/07/23/big/img_542 +2002/08/06/big/img_2535 +2002/07/22/big/img_37 +2002/08/06/big/img_2342 +2002/08/25/big/img_515 +2002/08/25/big/img_336 +2002/08/18/big/img_837 +2002/08/21/big/img_616 +2003/01/17/big/img_24 +2002/07/26/big/img_936 +2002/08/14/big/img_896 +2002/07/29/big/img_465 +2002/07/31/big/img_543 +2002/08/01/big/img_1411 +2002/08/02/big/img_423 +2002/08/21/big/img_44 +2002/07/31/big/img_11 +2003/01/15/big/img_628 +2003/01/15/big/img_605 +2002/07/30/big/img_571 +2002/07/23/big/img_428 +2002/08/15/big/img_942 +2002/07/26/big/img_531 +2003/01/16/big/img_59 +2002/08/02/big/img_410 +2002/07/31/big/img_230 +2002/08/19/big/img_806 +2003/01/14/big/img_462 +2002/08/16/big/img_370 +2002/08/13/big/img_380 +2002/08/16/big/img_932 +2002/07/19/big/img_393 +2002/08/20/big/img_764 +2002/08/15/big/img_616 +2002/07/26/big/img_267 +2002/07/27/big/img_1069 +2002/08/14/big/img_1041 +2003/01/13/big/img_594 +2002/09/01/big/img_16845 +2002/08/09/big/img_229 +2003/01/16/big/img_639 +2002/08/19/big/img_398 +2002/08/18/big/img_978 +2002/08/24/big/img_296 +2002/07/29/big/img_415 +2002/07/30/big/img_923 +2002/08/18/big/img_575 +2002/08/22/big/img_182 +2002/07/25/big/img_806 +2002/07/22/big/img_49 +2002/07/29/big/img_989 +2003/01/17/big/img_789 +2003/01/15/big/img_503 +2002/09/01/big/img_16062 +2003/01/17/big/img_794 +2002/08/15/big/img_564 +2003/01/15/big/img_222 +2002/08/01/big/img_1656 +2003/01/13/big/img_432 +2002/07/19/big/img_426 +2002/08/17/big/img_244 +2002/08/13/big/img_805 +2002/09/02/big/img_15067 +2002/08/11/big/img_58 +2002/08/22/big/img_636 +2002/07/22/big/img_416 +2002/08/13/big/img_836 +2002/08/26/big/img_363 +2002/07/30/big/img_917 +2003/01/14/big/img_206 +2002/08/12/big/img_311 +2002/08/31/big/img_17623 +2002/07/29/big/img_661 +2003/01/13/big/img_417 +2002/08/02/big/img_463 +2002/08/02/big/img_669 +2002/08/26/big/img_670 +2002/08/02/big/img_375 +2002/07/19/big/img_209 +2002/08/08/big/img_115 +2002/08/21/big/img_399 +2002/08/20/big/img_911 +2002/08/07/big/img_1212 +2002/08/20/big/img_578 +2002/08/22/big/img_554 +2002/08/21/big/img_484 +2002/07/25/big/img_450 +2002/08/03/big/img_542 +2002/08/15/big/img_561 +2002/07/23/big/img_360 +2002/08/30/big/img_18137 +2002/07/25/big/img_250 +2002/08/03/big/img_647 +2002/08/20/big/img_375 +2002/08/14/big/img_387 +2002/09/01/big/img_16990 +2002/08/28/big/img_19341 +2003/01/15/big/img_239 +2002/08/20/big/img_528 +2002/08/12/big/img_130 +2002/09/02/big/img_15108 +2003/01/15/big/img_372 +2002/08/16/big/img_678 +2002/08/04/big/img_623 +2002/07/23/big/img_477 +2002/08/28/big/img_19590 +2003/01/17/big/img_978 +2002/09/01/big/img_16692 +2002/07/20/big/img_109 +2002/08/06/big/img_2660 +2003/01/14/big/img_464 +2002/08/09/big/img_618 +2002/07/22/big/img_722 +2002/08/25/big/img_419 +2002/08/03/big/img_314 +2002/08/25/big/img_40 +2002/07/27/big/img_430 +2002/08/10/big/img_569 +2002/08/23/big/img_398 +2002/07/23/big/img_893 +2002/08/16/big/img_261 +2002/08/06/big/img_2668 +2002/07/22/big/img_835 +2002/09/02/big/img_15093 +2003/01/16/big/img_65 +2002/08/21/big/img_448 +2003/01/14/big/img_351 +2003/01/17/big/img_133 +2002/07/28/big/img_493 +2003/01/15/big/img_640 +2002/09/01/big/img_16880 +2002/08/15/big/img_350 +2002/08/20/big/img_624 +2002/08/25/big/img_604 +2002/08/06/big/img_2200 +2002/08/23/big/img_290 +2002/08/13/big/img_1152 +2003/01/14/big/img_251 +2002/08/02/big/img_538 +2002/08/22/big/img_613 +2003/01/13/big/img_351 +2002/08/18/big/img_368 +2002/07/23/big/img_392 +2002/07/25/big/img_198 +2002/07/25/big/img_418 +2002/08/26/big/img_614 +2002/07/23/big/img_405 +2003/01/14/big/img_445 +2002/07/25/big/img_326 +2002/08/10/big/img_734 +2003/01/14/big/img_530 +2002/08/08/big/img_561 +2002/08/29/big/img_18990 +2002/08/10/big/img_576 +2002/07/29/big/img_1494 +2002/07/19/big/img_198 +2002/08/10/big/img_562 +2002/07/22/big/img_901 +2003/01/14/big/img_37 +2002/09/02/big/img_15629 +2003/01/14/big/img_58 +2002/08/01/big/img_1364 +2002/07/27/big/img_636 +2003/01/13/big/img_241 +2002/09/01/big/img_16988 +2003/01/13/big/img_560 +2002/08/09/big/img_533 +2002/07/31/big/img_249 +2003/01/17/big/img_1007 +2002/07/21/big/img_64 +2003/01/13/big/img_537 +2003/01/15/big/img_606 +2002/08/18/big/img_651 +2002/08/24/big/img_405 +2002/07/26/big/img_837 +2002/08/09/big/img_562 +2002/08/01/big/img_1983 +2002/08/03/big/img_514 +2002/07/29/big/img_314 +2002/08/12/big/img_493 +2003/01/14/big/img_121 +2003/01/14/big/img_479 +2002/08/04/big/img_410 +2002/07/22/big/img_607 +2003/01/17/big/img_417 +2002/07/20/big/img_547 +2002/08/13/big/img_396 +2002/08/31/big/img_17538 +2002/08/13/big/img_187 +2002/08/12/big/img_328 +2003/01/14/big/img_569 +2002/07/27/big/img_1081 +2002/08/14/big/img_504 +2002/08/23/big/img_785 +2002/07/26/big/img_339 +2002/08/07/big/img_1156 +2002/08/07/big/img_1456 +2002/08/23/big/img_378 +2002/08/27/big/img_19719 +2002/07/31/big/img_39 +2002/07/31/big/img_883 +2003/01/14/big/img_676 +2002/07/29/big/img_214 +2002/07/26/big/img_669 +2002/07/25/big/img_202 +2002/08/08/big/img_259 +2003/01/17/big/img_943 +2003/01/15/big/img_512 +2002/08/05/big/img_3295 +2002/08/27/big/img_19685 +2002/08/08/big/img_277 +2002/08/30/big/img_18154 +2002/07/22/big/img_663 +2002/08/29/big/img_18914 +2002/07/31/big/img_908 +2002/08/27/big/img_19926 +2003/01/13/big/img_791 +2003/01/15/big/img_827 +2002/08/18/big/img_878 +2002/08/14/big/img_670 +2002/07/20/big/img_182 +2002/08/15/big/img_291 +2002/08/06/big/img_2600 +2002/07/23/big/img_587 +2002/08/14/big/img_577 +2003/01/15/big/img_585 +2002/07/30/big/img_310 +2002/08/03/big/img_658 +2002/08/10/big/img_157 +2002/08/19/big/img_811 +2002/07/29/big/img_1318 +2002/08/04/big/img_104 +2002/07/30/big/img_332 +2002/07/24/big/img_789 +2002/07/29/big/img_516 +2002/07/23/big/img_843 +2002/08/01/big/img_1528 +2002/08/13/big/img_798 +2002/08/07/big/img_1729 +2002/08/28/big/img_19448 +2003/01/16/big/img_95 +2002/08/12/big/img_473 +2002/07/27/big/img_269 +2003/01/16/big/img_621 +2002/07/29/big/img_772 +2002/07/24/big/img_171 +2002/07/19/big/img_429 +2002/08/07/big/img_1933 +2002/08/27/big/img_19629 +2002/08/05/big/img_3688 +2002/08/07/big/img_1691 +2002/07/23/big/img_600 +2002/07/29/big/img_666 +2002/08/25/big/img_566 +2002/08/06/big/img_2659 +2002/08/29/big/img_18929 +2002/08/16/big/img_407 +2002/08/18/big/img_774 +2002/08/19/big/img_249 +2002/08/06/big/img_2427 +2002/08/29/big/img_18899 +2002/08/01/big/img_1818 +2002/07/31/big/img_108 +2002/07/29/big/img_500 +2002/08/11/big/img_115 +2002/07/19/big/img_521 +2002/08/02/big/img_1163 +2002/07/22/big/img_62 +2002/08/13/big/img_466 +2002/08/21/big/img_956 +2002/08/23/big/img_602 +2002/08/20/big/img_858 +2002/07/25/big/img_690 +2002/07/19/big/img_130 +2002/08/04/big/img_874 +2002/07/26/big/img_489 +2002/07/22/big/img_548 +2002/08/10/big/img_191 +2002/07/25/big/img_1051 +2002/08/18/big/img_473 +2002/08/12/big/img_755 +2002/08/18/big/img_413 +2002/08/08/big/img_1044 +2002/08/17/big/img_680 +2002/08/26/big/img_235 +2002/08/20/big/img_330 +2002/08/22/big/img_344 +2002/08/09/big/img_593 +2002/07/31/big/img_1006 +2002/08/14/big/img_337 +2002/08/16/big/img_728 +2002/07/24/big/img_834 +2002/08/04/big/img_552 +2002/09/02/big/img_15213 +2002/07/25/big/img_725 +2002/08/30/big/img_18290 +2003/01/01/big/img_475 +2002/07/27/big/img_1083 +2002/08/29/big/img_18955 +2002/08/31/big/img_17232 +2002/08/08/big/img_480 +2002/08/01/big/img_1311 +2002/07/30/big/img_745 +2002/08/03/big/img_649 +2002/08/12/big/img_193 +2002/07/29/big/img_228 +2002/07/25/big/img_836 +2002/08/20/big/img_400 +2002/07/30/big/img_507 +2002/09/02/big/img_15072 +2002/07/26/big/img_658 +2002/07/28/big/img_503 +2002/08/05/big/img_3814 +2002/08/24/big/img_745 +2003/01/13/big/img_817 +2002/08/08/big/img_579 +2002/07/22/big/img_251 +2003/01/13/big/img_689 +2002/07/25/big/img_407 +2002/08/13/big/img_1050 +2002/08/14/big/img_733 +2002/07/24/big/img_82 +2003/01/17/big/img_288 +2003/01/15/big/img_475 +2002/08/14/big/img_620 +2002/08/21/big/img_167 +2002/07/19/big/img_300 +2002/07/26/big/img_219 +2002/08/01/big/img_1468 +2002/07/23/big/img_260 +2002/08/09/big/img_555 +2002/07/19/big/img_160 +2002/08/02/big/img_1060 +2003/01/14/big/img_149 +2002/08/15/big/img_346 +2002/08/24/big/img_597 +2002/08/22/big/img_502 +2002/08/30/big/img_18228 +2002/07/21/big/img_766 +2003/01/15/big/img_841 +2002/07/24/big/img_516 +2002/08/02/big/img_265 +2002/08/15/big/img_1243 +2003/01/15/big/img_223 +2002/08/04/big/img_236 +2002/07/22/big/img_309 +2002/07/20/big/img_656 +2002/07/31/big/img_412 +2002/09/01/big/img_16462 +2003/01/16/big/img_431 +2002/07/22/big/img_793 +2002/08/15/big/img_877 +2002/07/26/big/img_282 +2002/07/25/big/img_529 +2002/08/24/big/img_613 +2003/01/17/big/img_700 +2002/08/06/big/img_2526 +2002/08/24/big/img_394 +2002/08/21/big/img_521 +2002/08/25/big/img_560 +2002/07/29/big/img_966 +2002/07/25/big/img_448 +2003/01/13/big/img_782 +2002/08/21/big/img_296 +2002/09/01/big/img_16755 +2002/08/05/big/img_3552 +2002/09/02/big/img_15823 +2003/01/14/big/img_193 +2002/07/21/big/img_159 +2002/08/02/big/img_564 +2002/08/16/big/img_300 +2002/07/19/big/img_269 +2002/08/13/big/img_676 +2002/07/28/big/img_57 +2002/08/05/big/img_3318 +2002/07/31/big/img_218 +2002/08/21/big/img_898 +2002/07/29/big/img_109 +2002/07/19/big/img_854 +2002/08/23/big/img_311 +2002/08/14/big/img_318 +2002/07/25/big/img_523 +2002/07/21/big/img_678 +2003/01/17/big/img_690 +2002/08/28/big/img_19503 +2002/08/18/big/img_251 +2002/08/22/big/img_672 +2002/08/20/big/img_663 +2002/08/02/big/img_148 +2002/09/02/big/img_15580 +2002/07/25/big/img_778 +2002/08/14/big/img_565 +2002/08/12/big/img_374 +2002/08/13/big/img_1018 +2002/08/20/big/img_474 +2002/08/25/big/img_33 +2002/08/02/big/img_1190 +2002/08/08/big/img_864 +2002/08/14/big/img_1071 +2002/08/30/big/img_18103 +2002/08/18/big/img_533 +2003/01/16/big/img_650 +2002/07/25/big/img_108 +2002/07/26/big/img_81 +2002/07/27/big/img_543 +2002/07/29/big/img_521 +2003/01/13/big/img_434 +2002/08/26/big/img_674 +2002/08/06/big/img_2932 +2002/08/07/big/img_1262 +2003/01/15/big/img_201 +2003/01/16/big/img_673 +2002/09/02/big/img_15988 +2002/07/29/big/img_1306 +2003/01/14/big/img_1072 +2002/08/30/big/img_18232 +2002/08/05/big/img_3711 +2002/07/23/big/img_775 +2002/08/01/big/img_16 +2003/01/16/big/img_630 +2002/08/22/big/img_695 +2002/08/14/big/img_51 +2002/08/14/big/img_782 +2002/08/24/big/img_742 +2003/01/14/big/img_512 +2003/01/15/big/img_1183 +2003/01/15/big/img_714 +2002/08/01/big/img_2078 +2002/07/31/big/img_682 +2002/09/02/big/img_15687 +2002/07/26/big/img_518 +2002/08/27/big/img_19676 +2002/09/02/big/img_15969 +2002/08/02/big/img_931 +2002/08/25/big/img_508 +2002/08/29/big/img_18616 +2002/07/22/big/img_839 +2002/07/28/big/img_313 +2003/01/14/big/img_155 +2002/08/02/big/img_1105 +2002/08/09/big/img_53 +2002/08/16/big/img_469 +2002/08/15/big/img_502 +2002/08/20/big/img_575 +2002/07/25/big/img_138 +2003/01/16/big/img_579 +2002/07/19/big/img_352 +2003/01/14/big/img_762 +2003/01/01/big/img_588 +2002/08/02/big/img_981 +2002/08/21/big/img_447 +2002/09/01/big/img_16151 +2003/01/14/big/img_769 +2002/08/23/big/img_461 +2002/08/17/big/img_240 +2002/09/02/big/img_15220 +2002/07/19/big/img_408 +2002/09/02/big/img_15496 +2002/07/29/big/img_758 +2002/08/28/big/img_19392 +2002/08/06/big/img_2723 +2002/08/31/big/img_17752 +2002/08/23/big/img_469 +2002/08/13/big/img_515 +2002/09/02/big/img_15551 +2002/08/03/big/img_462 +2002/07/24/big/img_613 +2002/07/22/big/img_61 +2002/08/08/big/img_171 +2002/08/21/big/img_177 +2003/01/14/big/img_105 +2002/08/02/big/img_1017 +2002/08/22/big/img_106 +2002/07/27/big/img_542 +2002/07/21/big/img_665 +2002/07/23/big/img_595 +2002/08/04/big/img_657 +2002/08/29/big/img_19002 +2003/01/15/big/img_550 +2002/08/14/big/img_662 +2002/07/20/big/img_425 +2002/08/30/big/img_18528 +2002/07/26/big/img_611 +2002/07/22/big/img_849 +2002/08/07/big/img_1655 +2002/08/21/big/img_638 +2003/01/17/big/img_732 +2003/01/01/big/img_496 +2002/08/18/big/img_713 +2002/08/08/big/img_109 +2002/07/27/big/img_1008 +2002/07/20/big/img_559 +2002/08/16/big/img_699 +2002/08/31/big/img_17702 +2002/07/31/big/img_1013 +2002/08/01/big/img_2027 +2002/08/02/big/img_1001 +2002/08/03/big/img_210 +2002/08/01/big/img_2087 +2003/01/14/big/img_199 +2002/07/29/big/img_48 +2002/07/19/big/img_727 +2002/08/09/big/img_249 +2002/08/04/big/img_632 +2002/08/22/big/img_620 +2003/01/01/big/img_457 +2002/08/05/big/img_3223 +2002/07/27/big/img_240 +2002/07/25/big/img_797 +2002/08/13/big/img_430 +2002/07/25/big/img_615 +2002/08/12/big/img_28 +2002/07/30/big/img_220 +2002/07/24/big/img_89 +2002/08/21/big/img_357 +2002/08/09/big/img_590 +2003/01/13/big/img_525 +2002/08/17/big/img_818 +2003/01/02/big/img_7 +2002/07/26/big/img_636 +2003/01/13/big/img_1122 +2002/07/23/big/img_810 +2002/08/20/big/img_888 +2002/07/27/big/img_3 +2002/08/15/big/img_451 +2002/09/02/big/img_15787 +2002/07/31/big/img_281 +2002/08/05/big/img_3274 +2002/08/07/big/img_1254 +2002/07/31/big/img_27 +2002/08/01/big/img_1366 +2002/07/30/big/img_182 +2002/08/27/big/img_19690 +2002/07/29/big/img_68 +2002/08/23/big/img_754 +2002/07/30/big/img_540 +2002/08/27/big/img_20063 +2002/08/14/big/img_471 +2002/08/02/big/img_615 +2002/07/30/big/img_186 +2002/08/25/big/img_150 +2002/07/27/big/img_626 +2002/07/20/big/img_225 +2003/01/15/big/img_1252 +2002/07/19/big/img_367 +2003/01/15/big/img_582 +2002/08/09/big/img_572 +2002/08/08/big/img_428 +2003/01/15/big/img_639 +2002/08/28/big/img_19245 +2002/07/24/big/img_321 +2002/08/02/big/img_662 +2002/08/08/big/img_1033 +2003/01/17/big/img_867 +2002/07/22/big/img_652 +2003/01/14/big/img_224 +2002/08/18/big/img_49 +2002/07/26/big/img_46 +2002/08/31/big/img_18021 +2002/07/25/big/img_151 +2002/08/23/big/img_540 +2002/08/25/big/img_693 +2002/07/23/big/img_340 +2002/07/28/big/img_117 +2002/09/02/big/img_15768 +2002/08/26/big/img_562 +2002/07/24/big/img_480 +2003/01/15/big/img_341 +2002/08/10/big/img_783 +2002/08/20/big/img_132 +2003/01/14/big/img_370 +2002/07/20/big/img_720 +2002/08/03/big/img_144 +2002/08/20/big/img_538 +2002/08/01/big/img_1745 +2002/08/11/big/img_683 +2002/08/03/big/img_328 +2002/08/10/big/img_793 +2002/08/14/big/img_689 +2002/08/02/big/img_162 +2003/01/17/big/img_411 +2002/07/31/big/img_361 +2002/08/15/big/img_289 +2002/08/08/big/img_254 +2002/08/15/big/img_996 +2002/08/20/big/img_785 +2002/07/24/big/img_511 +2002/08/06/big/img_2614 +2002/08/29/big/img_18733 +2002/08/17/big/img_78 +2002/07/30/big/img_378 +2002/08/31/big/img_17947 +2002/08/26/big/img_88 +2002/07/30/big/img_558 +2002/08/02/big/img_67 +2003/01/14/big/img_325 +2002/07/29/big/img_1357 +2002/07/19/big/img_391 +2002/07/30/big/img_307 +2003/01/13/big/img_219 +2002/07/24/big/img_807 +2002/08/23/big/img_543 +2002/08/29/big/img_18620 +2002/07/22/big/img_769 +2002/08/26/big/img_503 +2002/07/30/big/img_78 +2002/08/14/big/img_1036 +2002/08/09/big/img_58 +2002/07/24/big/img_616 +2002/08/02/big/img_464 +2002/07/26/big/img_576 +2002/07/22/big/img_273 +2003/01/16/big/img_470 +2002/07/29/big/img_329 +2002/07/30/big/img_1086 +2002/07/31/big/img_353 +2002/09/02/big/img_15275 +2003/01/17/big/img_555 +2002/08/26/big/img_212 +2002/08/01/big/img_1692 +2003/01/15/big/img_600 +2002/07/29/big/img_825 +2002/08/08/big/img_68 +2002/08/10/big/img_719 +2002/07/31/big/img_636 +2002/07/29/big/img_325 +2002/07/21/big/img_515 +2002/07/22/big/img_705 +2003/01/13/big/img_818 +2002/08/09/big/img_486 +2002/08/22/big/img_141 +2002/07/22/big/img_303 +2002/08/09/big/img_393 +2002/07/29/big/img_963 +2002/08/02/big/img_1215 +2002/08/19/big/img_674 +2002/08/12/big/img_690 +2002/08/21/big/img_637 +2002/08/21/big/img_841 +2002/08/24/big/img_71 +2002/07/25/big/img_596 +2002/07/24/big/img_864 +2002/08/18/big/img_293 +2003/01/14/big/img_657 +2002/08/15/big/img_411 +2002/08/16/big/img_348 +2002/08/05/big/img_3157 +2002/07/20/big/img_663 +2003/01/13/big/img_654 +2003/01/16/big/img_433 +2002/08/30/big/img_18200 +2002/08/12/big/img_226 +2003/01/16/big/img_491 +2002/08/08/big/img_666 +2002/07/19/big/img_576 +2003/01/15/big/img_776 +2003/01/16/big/img_899 +2002/07/19/big/img_397 +2002/08/14/big/img_44 +2003/01/15/big/img_762 +2002/08/02/big/img_982 +2002/09/02/big/img_15234 +2002/08/17/big/img_556 +2002/08/21/big/img_410 +2002/08/21/big/img_386 +2002/07/19/big/img_690 +2002/08/05/big/img_3052 +2002/08/14/big/img_219 +2002/08/16/big/img_273 +2003/01/15/big/img_752 +2002/08/08/big/img_184 +2002/07/31/big/img_743 +2002/08/23/big/img_338 +2003/01/14/big/img_1055 +2002/08/05/big/img_3405 +2003/01/15/big/img_17 +2002/08/03/big/img_141 +2002/08/14/big/img_549 +2002/07/27/big/img_1034 +2002/07/31/big/img_932 +2002/08/30/big/img_18487 +2002/09/02/big/img_15814 +2002/08/01/big/img_2086 +2002/09/01/big/img_16535 +2002/07/22/big/img_500 +2003/01/13/big/img_400 +2002/08/25/big/img_607 +2002/08/30/big/img_18384 +2003/01/14/big/img_951 +2002/08/13/big/img_1150 +2002/08/08/big/img_1022 +2002/08/10/big/img_428 +2002/08/28/big/img_19242 +2002/08/05/big/img_3098 +2002/07/23/big/img_400 +2002/08/26/big/img_365 +2002/07/20/big/img_318 +2002/08/13/big/img_740 +2003/01/16/big/img_37 +2002/08/26/big/img_274 +2002/08/02/big/img_205 +2002/08/21/big/img_695 +2002/08/06/big/img_2289 +2002/08/20/big/img_794 +2002/08/18/big/img_438 +2002/08/07/big/img_1380 +2002/08/02/big/img_737 +2002/08/07/big/img_1651 +2002/08/15/big/img_1238 +2002/08/01/big/img_1681 +2002/08/06/big/img_3017 +2002/07/23/big/img_706 +2002/07/31/big/img_392 +2002/08/09/big/img_539 +2002/07/29/big/img_835 +2002/08/26/big/img_723 +2002/08/28/big/img_19235 +2003/01/16/big/img_353 +2002/08/10/big/img_150 +2002/08/29/big/img_19025 +2002/08/21/big/img_310 +2002/08/10/big/img_823 +2002/07/26/big/img_981 +2002/08/11/big/img_288 +2002/08/19/big/img_534 +2002/08/21/big/img_300 +2002/07/31/big/img_49 +2002/07/30/big/img_469 +2002/08/28/big/img_19197 +2002/08/25/big/img_205 +2002/08/10/big/img_390 +2002/08/23/big/img_291 +2002/08/26/big/img_230 +2002/08/18/big/img_76 +2002/07/23/big/img_409 +2002/08/14/big/img_1053 +2003/01/14/big/img_291 +2002/08/10/big/img_503 +2002/08/27/big/img_19928 +2002/08/03/big/img_563 +2002/08/17/big/img_250 +2002/08/06/big/img_2381 +2002/08/17/big/img_948 +2002/08/06/big/img_2710 +2002/07/22/big/img_696 +2002/07/31/big/img_670 +2002/08/12/big/img_594 +2002/07/29/big/img_624 +2003/01/17/big/img_934 +2002/08/03/big/img_584 +2002/08/22/big/img_1003 +2002/08/05/big/img_3396 +2003/01/13/big/img_570 +2002/08/02/big/img_219 +2002/09/02/big/img_15774 +2002/08/16/big/img_818 +2002/08/23/big/img_402 +2003/01/14/big/img_552 +2002/07/29/big/img_71 +2002/08/05/big/img_3592 +2002/08/16/big/img_80 +2002/07/27/big/img_672 +2003/01/13/big/img_470 +2003/01/16/big/img_702 +2002/09/01/big/img_16130 +2002/08/08/big/img_240 +2002/09/01/big/img_16338 +2002/07/26/big/img_312 +2003/01/14/big/img_538 +2002/07/20/big/img_695 +2002/08/30/big/img_18098 +2002/08/25/big/img_259 +2002/08/16/big/img_1042 +2002/08/09/big/img_837 +2002/08/31/big/img_17760 +2002/07/31/big/img_14 +2002/08/09/big/img_361 +2003/01/16/big/img_107 +2002/08/14/big/img_124 +2002/07/19/big/img_463 +2003/01/15/big/img_275 +2002/07/25/big/img_1151 +2002/07/29/big/img_1501 +2002/08/27/big/img_19889 +2002/08/29/big/img_18603 +2003/01/17/big/img_601 +2002/08/25/big/img_355 +2002/08/08/big/img_297 +2002/08/20/big/img_290 +2002/07/31/big/img_195 +2003/01/01/big/img_336 +2002/08/18/big/img_369 +2002/07/25/big/img_621 +2002/08/11/big/img_508 +2003/01/14/big/img_458 +2003/01/15/big/img_795 +2002/08/12/big/img_498 +2002/08/01/big/img_1734 +2002/08/02/big/img_246 +2002/08/16/big/img_565 +2002/08/11/big/img_475 +2002/08/22/big/img_408 +2002/07/28/big/img_78 +2002/07/21/big/img_81 +2003/01/14/big/img_697 +2002/08/14/big/img_661 +2002/08/15/big/img_507 +2002/08/19/big/img_55 +2002/07/22/big/img_152 +2003/01/14/big/img_470 +2002/08/03/big/img_379 +2002/08/22/big/img_506 +2003/01/16/big/img_966 +2002/08/18/big/img_698 +2002/08/24/big/img_528 +2002/08/23/big/img_10 +2002/08/01/big/img_1655 +2002/08/22/big/img_953 +2002/07/19/big/img_630 +2002/07/22/big/img_889 +2002/08/16/big/img_351 +2003/01/16/big/img_83 +2002/07/19/big/img_805 +2002/08/14/big/img_704 +2002/07/19/big/img_389 +2002/08/31/big/img_17765 +2002/07/29/big/img_606 +2003/01/17/big/img_939 +2002/09/02/big/img_15081 +2002/08/21/big/img_181 +2002/07/29/big/img_1321 +2002/07/21/big/img_497 +2002/07/20/big/img_539 +2002/08/24/big/img_119 +2002/08/01/big/img_1281 +2002/07/26/big/img_207 +2002/07/26/big/img_432 +2002/07/27/big/img_1006 +2002/08/05/big/img_3087 +2002/08/14/big/img_252 +2002/08/14/big/img_798 +2002/07/24/big/img_538 +2002/09/02/big/img_15507 +2002/08/08/big/img_901 +2003/01/14/big/img_557 +2002/08/07/big/img_1819 +2002/08/04/big/img_470 +2002/08/01/big/img_1504 +2002/08/16/big/img_1070 +2002/08/16/big/img_372 +2002/08/23/big/img_416 +2002/08/30/big/img_18208 +2002/08/01/big/img_2043 +2002/07/22/big/img_385 +2002/08/22/big/img_466 +2002/08/21/big/img_869 +2002/08/28/big/img_19429 +2002/08/02/big/img_770 +2002/07/23/big/img_433 +2003/01/14/big/img_13 +2002/07/27/big/img_953 +2002/09/02/big/img_15728 +2002/08/01/big/img_1361 +2002/08/29/big/img_18897 +2002/08/26/big/img_534 +2002/08/11/big/img_121 +2002/08/26/big/img_20130 +2002/07/31/big/img_363 +2002/08/13/big/img_978 +2002/07/25/big/img_835 +2002/08/02/big/img_906 +2003/01/14/big/img_548 +2002/07/30/big/img_80 +2002/07/26/big/img_982 +2003/01/16/big/img_99 +2002/08/19/big/img_362 +2002/08/24/big/img_376 +2002/08/07/big/img_1264 +2002/07/27/big/img_938 +2003/01/17/big/img_535 +2002/07/26/big/img_457 +2002/08/08/big/img_848 +2003/01/15/big/img_859 +2003/01/15/big/img_622 +2002/07/30/big/img_403 +2002/07/29/big/img_217 +2002/07/26/big/img_891 +2002/07/24/big/img_70 +2002/08/25/big/img_619 +2002/08/05/big/img_3375 +2002/08/01/big/img_2160 +2002/08/06/big/img_2227 +2003/01/14/big/img_117 +2002/08/14/big/img_227 +2002/08/13/big/img_565 +2002/08/19/big/img_625 +2002/08/03/big/img_812 +2002/07/24/big/img_41 +2002/08/16/big/img_235 +2002/07/29/big/img_759 +2002/07/21/big/img_433 +2002/07/29/big/img_190 +2003/01/16/big/img_435 +2003/01/13/big/img_708 +2002/07/30/big/img_57 +2002/08/22/big/img_162 +2003/01/01/big/img_558 +2003/01/15/big/img_604 +2002/08/16/big/img_935 +2002/08/20/big/img_394 +2002/07/28/big/img_465 +2002/09/02/big/img_15534 +2002/08/16/big/img_87 +2002/07/22/big/img_469 +2002/08/12/big/img_245 +2003/01/13/big/img_236 +2002/08/06/big/img_2736 +2002/08/03/big/img_348 +2003/01/14/big/img_218 +2002/07/26/big/img_232 +2003/01/15/big/img_244 +2002/07/25/big/img_1121 +2002/08/01/big/img_1484 +2002/07/26/big/img_541 +2002/08/07/big/img_1244 +2002/07/31/big/img_3 +2002/08/30/big/img_18437 +2002/08/29/big/img_19094 +2002/08/01/big/img_1355 +2002/08/19/big/img_338 +2002/07/19/big/img_255 +2002/07/21/big/img_76 +2002/08/25/big/img_199 +2002/08/12/big/img_740 +2002/07/30/big/img_852 +2002/08/15/big/img_599 +2002/08/23/big/img_254 +2002/08/19/big/img_125 +2002/07/24/big/img_2 +2002/08/04/big/img_145 +2002/08/05/big/img_3137 +2002/07/28/big/img_463 +2003/01/14/big/img_801 +2002/07/23/big/img_366 +2002/08/26/big/img_600 +2002/08/26/big/img_649 +2002/09/02/big/img_15849 +2002/07/26/big/img_248 +2003/01/13/big/img_200 +2002/08/07/big/img_1794 +2002/08/31/big/img_17270 +2002/08/23/big/img_608 +2003/01/13/big/img_837 +2002/08/23/big/img_581 +2002/08/20/big/img_754 +2002/08/18/big/img_183 +2002/08/20/big/img_328 +2002/07/22/big/img_494 +2002/07/29/big/img_399 +2002/08/28/big/img_19284 +2002/08/08/big/img_566 +2002/07/25/big/img_376 +2002/07/23/big/img_138 +2002/07/25/big/img_435 +2002/08/17/big/img_685 +2002/07/19/big/img_90 +2002/07/20/big/img_716 +2002/08/31/big/img_17458 +2002/08/26/big/img_461 +2002/07/25/big/img_355 +2002/08/06/big/img_2152 +2002/07/27/big/img_932 +2002/07/23/big/img_232 +2002/08/08/big/img_1020 +2002/07/31/big/img_366 +2002/08/06/big/img_2667 +2002/08/21/big/img_465 +2002/08/15/big/img_305 +2002/08/02/big/img_247 +2002/07/28/big/img_46 +2002/08/27/big/img_19922 +2002/08/23/big/img_643 +2003/01/13/big/img_624 +2002/08/23/big/img_625 +2002/08/05/big/img_3787 +2003/01/13/big/img_627 +2002/09/01/big/img_16381 +2002/08/05/big/img_3668 +2002/07/21/big/img_535 +2002/08/27/big/img_19680 +2002/07/22/big/img_413 +2002/07/29/big/img_481 +2003/01/15/big/img_496 +2002/07/23/big/img_701 +2002/08/29/big/img_18670 +2002/07/28/big/img_319 +2003/01/14/big/img_517 +2002/07/26/big/img_256 +2003/01/16/big/img_593 +2002/07/30/big/img_956 +2002/07/30/big/img_667 +2002/07/25/big/img_100 +2002/08/11/big/img_570 +2002/07/26/big/img_745 +2002/08/04/big/img_834 +2002/08/25/big/img_521 +2002/08/01/big/img_2148 +2002/09/02/big/img_15183 +2002/08/22/big/img_514 +2002/08/23/big/img_477 +2002/07/23/big/img_336 +2002/07/26/big/img_481 +2002/08/20/big/img_409 +2002/07/23/big/img_918 +2002/08/09/big/img_474 +2002/08/02/big/img_929 +2002/08/31/big/img_17932 +2002/08/19/big/img_161 +2002/08/09/big/img_667 +2002/07/31/big/img_805 +2002/09/02/big/img_15678 +2002/08/31/big/img_17509 +2002/08/29/big/img_18998 +2002/07/23/big/img_301 +2002/08/07/big/img_1612 +2002/08/06/big/img_2472 +2002/07/23/big/img_466 +2002/08/27/big/img_19634 +2003/01/16/big/img_16 +2002/08/14/big/img_193 +2002/08/21/big/img_340 +2002/08/27/big/img_19799 +2002/08/01/big/img_1345 +2002/08/07/big/img_1448 +2002/08/11/big/img_324 +2003/01/16/big/img_754 +2002/08/13/big/img_418 +2003/01/16/big/img_544 +2002/08/19/big/img_135 +2002/08/10/big/img_455 +2002/08/10/big/img_693 +2002/08/31/big/img_17967 +2002/08/28/big/img_19229 +2002/08/04/big/img_811 +2002/09/01/big/img_16225 +2003/01/16/big/img_428 +2002/09/02/big/img_15295 +2002/07/26/big/img_108 +2002/07/21/big/img_477 +2002/08/07/big/img_1354 +2002/08/23/big/img_246 +2002/08/16/big/img_652 +2002/07/27/big/img_553 +2002/07/31/big/img_346 +2002/08/04/big/img_537 +2002/08/08/big/img_498 +2002/08/29/big/img_18956 +2003/01/13/big/img_922 +2002/08/31/big/img_17425 +2002/07/26/big/img_438 +2002/08/19/big/img_185 +2003/01/16/big/img_33 +2002/08/10/big/img_252 +2002/07/29/big/img_598 +2002/08/27/big/img_19820 +2002/08/06/big/img_2664 +2002/08/20/big/img_705 +2003/01/14/big/img_816 +2002/08/03/big/img_552 +2002/07/25/big/img_561 +2002/07/25/big/img_934 +2002/08/01/big/img_1893 +2003/01/14/big/img_746 +2003/01/16/big/img_519 +2002/08/03/big/img_681 +2002/07/24/big/img_808 +2002/08/14/big/img_803 +2002/08/25/big/img_155 +2002/07/30/big/img_1107 +2002/08/29/big/img_18882 +2003/01/15/big/img_598 +2002/08/19/big/img_122 +2002/07/30/big/img_428 +2002/07/24/big/img_684 +2002/08/22/big/img_192 +2002/08/22/big/img_543 +2002/08/07/big/img_1318 +2002/08/18/big/img_25 +2002/07/26/big/img_583 +2002/07/20/big/img_464 +2002/08/19/big/img_664 +2002/08/24/big/img_861 +2002/09/01/big/img_16136 +2002/08/22/big/img_400 +2002/08/12/big/img_445 +2003/01/14/big/img_174 +2002/08/27/big/img_19677 +2002/08/31/big/img_17214 +2002/08/30/big/img_18175 +2003/01/17/big/img_402 +2002/08/06/big/img_2396 +2002/08/18/big/img_448 +2002/08/21/big/img_165 +2002/08/31/big/img_17609 +2003/01/01/big/img_151 +2002/08/26/big/img_372 +2002/09/02/big/img_15994 +2002/07/26/big/img_660 +2002/09/02/big/img_15197 +2002/07/29/big/img_258 +2002/08/30/big/img_18525 +2003/01/13/big/img_368 +2002/07/29/big/img_1538 +2002/07/21/big/img_787 +2002/08/18/big/img_152 +2002/08/06/big/img_2379 +2003/01/17/big/img_864 +2002/08/27/big/img_19998 +2002/08/01/big/img_1634 +2002/07/25/big/img_414 +2002/08/22/big/img_627 +2002/08/07/big/img_1669 +2002/08/16/big/img_1052 +2002/08/31/big/img_17796 +2002/08/18/big/img_199 +2002/09/02/big/img_15147 +2002/08/09/big/img_460 +2002/08/14/big/img_581 +2002/08/30/big/img_18286 +2002/07/26/big/img_337 +2002/08/18/big/img_589 +2003/01/14/big/img_866 +2002/07/20/big/img_624 +2002/08/01/big/img_1801 +2002/07/24/big/img_683 +2002/08/09/big/img_725 +2003/01/14/big/img_34 +2002/07/30/big/img_144 +2002/07/30/big/img_706 +2002/08/08/big/img_394 +2002/08/19/big/img_619 +2002/08/06/big/img_2703 +2002/08/29/big/img_19034 +2002/07/24/big/img_67 +2002/08/27/big/img_19841 +2002/08/19/big/img_427 +2003/01/14/big/img_333 +2002/09/01/big/img_16406 +2002/07/19/big/img_882 +2002/08/17/big/img_238 +2003/01/14/big/img_739 +2002/07/22/big/img_151 +2002/08/21/big/img_743 +2002/07/25/big/img_1048 +2002/07/30/big/img_395 +2003/01/13/big/img_584 +2002/08/13/big/img_742 +2002/08/13/big/img_1168 +2003/01/14/big/img_147 +2002/07/26/big/img_803 +2002/08/05/big/img_3298 +2002/08/07/big/img_1451 +2002/08/16/big/img_424 +2002/07/29/big/img_1069 +2002/09/01/big/img_16735 +2002/07/21/big/img_637 +2003/01/14/big/img_585 +2002/08/02/big/img_358 +2003/01/13/big/img_358 +2002/08/14/big/img_198 +2002/08/17/big/img_935 +2002/08/04/big/img_42 +2002/08/30/big/img_18245 +2002/07/25/big/img_158 +2002/08/22/big/img_744 +2002/08/06/big/img_2291 +2002/08/05/big/img_3044 +2002/07/30/big/img_272 +2002/08/23/big/img_641 +2002/07/24/big/img_797 +2002/07/30/big/img_392 +2003/01/14/big/img_447 +2002/07/31/big/img_898 +2002/08/06/big/img_2812 +2002/08/13/big/img_564 +2002/07/22/big/img_43 +2002/07/26/big/img_634 +2002/07/19/big/img_843 +2002/08/26/big/img_58 +2002/07/21/big/img_375 +2002/08/25/big/img_729 +2002/07/19/big/img_561 +2003/01/15/big/img_884 +2002/07/25/big/img_891 +2002/08/09/big/img_558 +2002/08/26/big/img_587 +2002/08/13/big/img_1146 +2002/09/02/big/img_15153 +2002/07/26/big/img_316 +2002/08/01/big/img_1940 +2002/08/26/big/img_90 +2003/01/13/big/img_347 +2002/07/25/big/img_520 +2002/08/29/big/img_18718 +2002/08/28/big/img_19219 +2002/08/13/big/img_375 +2002/07/20/big/img_719 +2002/08/31/big/img_17431 +2002/07/28/big/img_192 +2002/08/26/big/img_259 +2002/08/18/big/img_484 +2002/07/29/big/img_580 +2002/07/26/big/img_84 +2002/08/02/big/img_302 +2002/08/31/big/img_17007 +2003/01/15/big/img_543 +2002/09/01/big/img_16488 +2002/08/22/big/img_798 +2002/07/30/big/img_383 +2002/08/04/big/img_668 +2002/08/13/big/img_156 +2002/08/07/big/img_1353 +2002/07/25/big/img_281 +2003/01/14/big/img_587 +2003/01/15/big/img_524 +2002/08/19/big/img_726 +2002/08/21/big/img_709 +2002/08/26/big/img_465 +2002/07/31/big/img_658 +2002/08/28/big/img_19148 +2002/07/23/big/img_423 +2002/08/16/big/img_758 +2002/08/22/big/img_523 +2002/08/16/big/img_591 +2002/08/23/big/img_845 +2002/07/26/big/img_678 +2002/08/09/big/img_806 +2002/08/06/big/img_2369 +2002/07/29/big/img_457 +2002/07/19/big/img_278 +2002/08/30/big/img_18107 +2002/07/26/big/img_444 +2002/08/20/big/img_278 +2002/08/26/big/img_92 +2002/08/26/big/img_257 +2002/07/25/big/img_266 +2002/08/05/big/img_3829 +2002/07/26/big/img_757 +2002/07/29/big/img_1536 +2002/08/09/big/img_472 +2003/01/17/big/img_480 +2002/08/28/big/img_19355 +2002/07/26/big/img_97 +2002/08/06/big/img_2503 +2002/07/19/big/img_254 +2002/08/01/big/img_1470 +2002/08/21/big/img_42 +2002/08/20/big/img_217 +2002/08/06/big/img_2459 +2002/07/19/big/img_552 +2002/08/13/big/img_717 +2002/08/12/big/img_586 +2002/08/20/big/img_411 +2003/01/13/big/img_768 +2002/08/07/big/img_1747 +2002/08/15/big/img_385 +2002/08/01/big/img_1648 +2002/08/15/big/img_311 +2002/08/21/big/img_95 +2002/08/09/big/img_108 +2002/08/21/big/img_398 +2002/08/17/big/img_340 +2002/08/14/big/img_474 +2002/08/13/big/img_294 +2002/08/24/big/img_840 +2002/08/09/big/img_808 +2002/08/23/big/img_491 +2002/07/28/big/img_33 +2003/01/13/big/img_664 +2002/08/02/big/img_261 +2002/08/09/big/img_591 +2002/07/26/big/img_309 +2003/01/14/big/img_372 +2002/08/19/big/img_581 +2002/08/19/big/img_168 +2002/08/26/big/img_422 +2002/07/24/big/img_106 +2002/08/01/big/img_1936 +2002/08/05/big/img_3764 +2002/08/21/big/img_266 +2002/08/31/big/img_17968 +2002/08/01/big/img_1941 +2002/08/15/big/img_550 +2002/08/14/big/img_13 +2002/07/30/big/img_171 +2003/01/13/big/img_490 +2002/07/25/big/img_427 +2002/07/19/big/img_770 +2002/08/12/big/img_759 +2003/01/15/big/img_1360 +2002/08/05/big/img_3692 +2003/01/16/big/img_30 +2002/07/25/big/img_1026 +2002/07/22/big/img_288 +2002/08/29/big/img_18801 +2002/07/24/big/img_793 +2002/08/13/big/img_178 +2002/08/06/big/img_2322 +2003/01/14/big/img_560 +2002/08/18/big/img_408 +2003/01/16/big/img_915 +2003/01/16/big/img_679 +2002/08/07/big/img_1552 +2002/08/29/big/img_19050 +2002/08/01/big/img_2172 +2002/07/31/big/img_30 +2002/07/30/big/img_1019 +2002/07/30/big/img_587 +2003/01/13/big/img_773 +2002/07/30/big/img_410 +2002/07/28/big/img_65 +2002/08/05/big/img_3138 +2002/07/23/big/img_541 +2002/08/22/big/img_963 +2002/07/27/big/img_657 +2002/07/30/big/img_1051 +2003/01/16/big/img_150 +2002/07/31/big/img_519 +2002/08/01/big/img_1961 +2002/08/05/big/img_3752 +2002/07/23/big/img_631 +2003/01/14/big/img_237 +2002/07/28/big/img_21 +2002/07/22/big/img_813 +2002/08/05/big/img_3563 +2003/01/17/big/img_620 +2002/07/19/big/img_523 +2002/07/30/big/img_904 +2002/08/29/big/img_18642 +2002/08/11/big/img_492 +2002/08/01/big/img_2130 +2002/07/25/big/img_618 +2002/08/17/big/img_305 +2003/01/16/big/img_520 +2002/07/26/big/img_495 +2002/08/17/big/img_164 +2002/08/03/big/img_440 +2002/07/24/big/img_441 +2002/08/06/big/img_2146 +2002/08/11/big/img_558 +2002/08/02/big/img_545 +2002/08/31/big/img_18090 +2003/01/01/big/img_136 +2002/07/25/big/img_1099 +2003/01/13/big/img_728 +2003/01/16/big/img_197 +2002/07/26/big/img_651 +2002/08/11/big/img_676 +2003/01/15/big/img_10 +2002/08/21/big/img_250 +2002/08/14/big/img_325 +2002/08/04/big/img_390 +2002/07/24/big/img_554 +2003/01/16/big/img_333 +2002/07/31/big/img_922 +2002/09/02/big/img_15586 +2003/01/16/big/img_184 +2002/07/22/big/img_766 +2002/07/21/big/img_608 +2002/08/07/big/img_1578 +2002/08/17/big/img_961 +2002/07/27/big/img_324 +2002/08/05/big/img_3765 +2002/08/23/big/img_462 +2003/01/16/big/img_382 +2002/08/27/big/img_19838 +2002/08/01/big/img_1505 +2002/08/21/big/img_662 +2002/08/14/big/img_605 +2002/08/19/big/img_816 +2002/07/29/big/img_136 +2002/08/20/big/img_719 +2002/08/06/big/img_2826 +2002/08/10/big/img_630 +2003/01/17/big/img_973 +2002/08/14/big/img_116 +2002/08/02/big/img_666 +2002/08/21/big/img_710 +2002/08/05/big/img_55 +2002/07/31/big/img_229 +2002/08/01/big/img_1549 +2002/07/23/big/img_432 +2002/07/21/big/img_430 +2002/08/21/big/img_549 +2002/08/08/big/img_985 +2002/07/20/big/img_610 +2002/07/23/big/img_978 +2002/08/23/big/img_219 +2002/07/25/big/img_175 +2003/01/15/big/img_230 +2002/08/23/big/img_385 +2002/07/31/big/img_879 +2002/08/12/big/img_495 +2002/08/22/big/img_499 +2002/08/30/big/img_18322 +2002/08/15/big/img_795 +2002/08/13/big/img_835 +2003/01/17/big/img_930 +2002/07/30/big/img_873 +2002/08/11/big/img_257 +2002/07/31/big/img_593 +2002/08/21/big/img_916 +2003/01/13/big/img_814 +2002/07/25/big/img_722 +2002/08/16/big/img_379 +2002/07/31/big/img_497 +2002/07/22/big/img_602 +2002/08/21/big/img_642 +2002/08/21/big/img_614 +2002/08/23/big/img_482 +2002/07/29/big/img_603 +2002/08/13/big/img_705 +2002/07/23/big/img_833 +2003/01/14/big/img_511 +2002/07/24/big/img_376 +2002/08/17/big/img_1030 +2002/08/05/big/img_3576 +2002/08/16/big/img_540 +2002/07/22/big/img_630 +2002/08/10/big/img_180 +2002/08/14/big/img_905 +2002/08/29/big/img_18777 +2002/08/22/big/img_693 +2003/01/16/big/img_933 +2002/08/20/big/img_555 +2002/08/15/big/img_549 +2003/01/14/big/img_830 +2003/01/16/big/img_64 +2002/08/27/big/img_19670 +2002/08/22/big/img_729 +2002/07/27/big/img_981 +2002/08/09/big/img_458 +2003/01/17/big/img_884 +2002/07/25/big/img_639 +2002/08/31/big/img_18008 +2002/08/22/big/img_249 +2002/08/17/big/img_971 +2002/08/04/big/img_308 +2002/07/28/big/img_362 +2002/08/12/big/img_142 +2002/08/26/big/img_61 +2002/08/14/big/img_422 +2002/07/19/big/img_607 +2003/01/15/big/img_717 +2002/08/01/big/img_1475 +2002/08/29/big/img_19061 +2003/01/01/big/img_346 +2002/07/20/big/img_315 +2003/01/15/big/img_756 +2002/08/15/big/img_879 +2002/08/08/big/img_615 +2003/01/13/big/img_431 +2002/08/05/big/img_3233 +2002/08/24/big/img_526 +2003/01/13/big/img_717 +2002/09/01/big/img_16408 +2002/07/22/big/img_217 +2002/07/31/big/img_960 +2002/08/21/big/img_610 +2002/08/05/big/img_3753 +2002/08/03/big/img_151 +2002/08/21/big/img_267 +2002/08/01/big/img_2175 +2002/08/04/big/img_556 +2002/08/21/big/img_527 +2002/09/02/big/img_15800 +2002/07/27/big/img_156 +2002/07/20/big/img_590 +2002/08/15/big/img_700 +2002/08/08/big/img_444 +2002/07/25/big/img_94 +2002/07/24/big/img_778 +2002/08/14/big/img_694 +2002/07/20/big/img_666 +2002/08/02/big/img_200 +2002/08/02/big/img_578 +2003/01/17/big/img_332 +2002/09/01/big/img_16352 +2002/08/27/big/img_19668 +2002/07/23/big/img_823 +2002/08/13/big/img_431 +2003/01/16/big/img_463 +2002/08/27/big/img_19711 +2002/08/23/big/img_154 +2002/07/31/big/img_360 +2002/08/23/big/img_555 +2002/08/10/big/img_561 +2003/01/14/big/img_550 +2002/08/07/big/img_1370 +2002/07/30/big/img_1184 +2002/08/01/big/img_1445 +2002/08/23/big/img_22 +2002/07/30/big/img_606 +2003/01/17/big/img_271 +2002/08/31/big/img_17316 +2002/08/16/big/img_973 +2002/07/26/big/img_77 +2002/07/20/big/img_788 +2002/08/06/big/img_2426 +2002/08/07/big/img_1498 +2002/08/16/big/img_358 +2002/08/06/big/img_2851 +2002/08/12/big/img_359 +2002/08/01/big/img_1521 +2002/08/02/big/img_709 +2002/08/20/big/img_935 +2002/08/12/big/img_188 +2002/08/24/big/img_411 +2002/08/22/big/img_680 +2002/08/06/big/img_2480 +2002/07/20/big/img_627 +2002/07/30/big/img_214 +2002/07/25/big/img_354 +2002/08/02/big/img_636 +2003/01/15/big/img_661 +2002/08/07/big/img_1327 +2002/08/01/big/img_2108 +2002/08/31/big/img_17919 +2002/08/29/big/img_18768 +2002/08/05/big/img_3840 +2002/07/26/big/img_242 +2003/01/14/big/img_451 +2002/08/20/big/img_923 +2002/08/27/big/img_19908 +2002/08/16/big/img_282 +2002/08/19/big/img_440 +2003/01/01/big/img_230 +2002/08/08/big/img_212 +2002/07/20/big/img_443 +2002/08/25/big/img_635 +2003/01/13/big/img_1169 +2002/07/26/big/img_998 +2002/08/15/big/img_995 +2002/08/06/big/img_3002 +2002/07/29/big/img_460 +2003/01/14/big/img_925 +2002/07/23/big/img_539 +2002/08/16/big/img_694 +2003/01/13/big/img_459 +2002/07/23/big/img_249 +2002/08/20/big/img_539 +2002/08/04/big/img_186 +2002/08/26/big/img_264 +2002/07/22/big/img_704 +2002/08/25/big/img_277 +2002/08/22/big/img_988 +2002/07/29/big/img_504 +2002/08/05/big/img_3600 +2002/08/30/big/img_18380 +2003/01/14/big/img_937 +2002/08/21/big/img_254 +2002/08/10/big/img_130 +2002/08/20/big/img_339 +2003/01/14/big/img_428 +2002/08/20/big/img_889 +2002/08/31/big/img_17637 +2002/07/26/big/img_644 +2002/09/01/big/img_16776 +2002/08/06/big/img_2239 +2002/08/06/big/img_2646 +2003/01/13/big/img_491 +2002/08/10/big/img_579 +2002/08/21/big/img_713 +2002/08/22/big/img_482 +2002/07/22/big/img_167 +2002/07/24/big/img_539 +2002/08/14/big/img_721 +2002/07/25/big/img_389 +2002/09/01/big/img_16591 +2002/08/13/big/img_543 +2003/01/14/big/img_432 +2002/08/09/big/img_287 +2002/07/26/big/img_126 +2002/08/23/big/img_412 +2002/08/15/big/img_1034 +2002/08/28/big/img_19485 +2002/07/31/big/img_236 +2002/07/30/big/img_523 +2002/07/19/big/img_141 +2003/01/17/big/img_957 +2002/08/04/big/img_81 +2002/07/25/big/img_206 +2002/08/15/big/img_716 +2002/08/13/big/img_403 +2002/08/15/big/img_685 +2002/07/26/big/img_884 +2002/07/19/big/img_499 +2002/07/23/big/img_772 +2002/07/27/big/img_752 +2003/01/14/big/img_493 +2002/08/25/big/img_664 +2002/07/31/big/img_334 +2002/08/26/big/img_678 +2002/09/01/big/img_16541 +2003/01/14/big/img_347 +2002/07/23/big/img_187 +2002/07/30/big/img_1163 +2002/08/05/big/img_35 +2002/08/22/big/img_944 +2002/08/07/big/img_1239 +2002/07/29/big/img_1215 +2002/08/03/big/img_312 +2002/08/05/big/img_3523 +2002/07/29/big/img_218 +2002/08/13/big/img_672 +2002/08/16/big/img_205 +2002/08/17/big/img_594 +2002/07/29/big/img_1411 +2002/07/30/big/img_942 +2003/01/16/big/img_312 +2002/08/08/big/img_312 +2002/07/25/big/img_15 +2002/08/09/big/img_839 +2002/08/01/big/img_2069 +2002/08/31/big/img_17512 +2002/08/01/big/img_3 +2002/07/31/big/img_320 +2003/01/15/big/img_1265 +2002/08/14/big/img_563 +2002/07/31/big/img_167 +2002/08/20/big/img_374 +2002/08/13/big/img_406 +2002/08/08/big/img_625 +2002/08/02/big/img_314 +2002/08/27/big/img_19964 +2002/09/01/big/img_16670 +2002/07/31/big/img_599 +2002/08/29/big/img_18906 +2002/07/24/big/img_373 +2002/07/26/big/img_513 +2002/09/02/big/img_15497 +2002/08/19/big/img_117 +2003/01/01/big/img_158 +2002/08/24/big/img_178 +2003/01/13/big/img_935 +2002/08/13/big/img_609 +2002/08/30/big/img_18341 +2002/08/25/big/img_674 +2003/01/13/big/img_209 +2002/08/13/big/img_258 +2002/08/05/big/img_3543 +2002/08/07/big/img_1970 +2002/08/06/big/img_3004 +2003/01/17/big/img_487 +2002/08/24/big/img_873 +2002/08/29/big/img_18730 +2002/08/09/big/img_375 +2003/01/16/big/img_751 +2002/08/02/big/img_603 +2002/08/19/big/img_325 +2002/09/01/big/img_16420 +2002/08/05/big/img_3633 +2002/08/21/big/img_516 +2002/07/19/big/img_501 +2002/07/26/big/img_688 +2002/07/24/big/img_256 +2002/07/25/big/img_438 +2002/07/31/big/img_1017 +2002/08/22/big/img_512 +2002/07/21/big/img_543 +2002/08/08/big/img_223 +2002/08/19/big/img_189 +2002/08/12/big/img_630 +2002/07/30/big/img_958 +2002/07/28/big/img_208 +2002/08/31/big/img_17691 +2002/07/22/big/img_542 +2002/07/19/big/img_741 +2002/07/19/big/img_158 +2002/08/15/big/img_399 +2002/08/01/big/img_2159 +2002/08/14/big/img_455 +2002/08/17/big/img_1011 +2002/08/26/big/img_744 +2002/08/12/big/img_624 +2003/01/17/big/img_821 +2002/08/16/big/img_980 +2002/07/28/big/img_281 +2002/07/25/big/img_171 +2002/08/03/big/img_116 +2002/07/22/big/img_467 +2002/07/31/big/img_750 +2002/07/26/big/img_435 +2002/07/19/big/img_822 +2002/08/13/big/img_626 +2002/08/11/big/img_344 +2002/08/02/big/img_473 +2002/09/01/big/img_16817 +2002/08/01/big/img_1275 +2002/08/28/big/img_19270 +2002/07/23/big/img_607 +2002/08/09/big/img_316 +2002/07/29/big/img_626 +2002/07/24/big/img_824 +2002/07/22/big/img_342 +2002/08/08/big/img_794 +2002/08/07/big/img_1209 +2002/07/19/big/img_18 +2002/08/25/big/img_634 +2002/07/24/big/img_730 +2003/01/17/big/img_356 +2002/07/23/big/img_305 +2002/07/30/big/img_453 +2003/01/13/big/img_972 +2002/08/06/big/img_2610 +2002/08/29/big/img_18920 +2002/07/31/big/img_123 +2002/07/26/big/img_979 +2002/08/24/big/img_635 +2002/08/05/big/img_3704 +2002/08/07/big/img_1358 +2002/07/22/big/img_306 +2002/08/13/big/img_619 +2002/08/02/big/img_366 diff --git a/hair_service_sd/models/layers/data/__init__.py b/hair_service_sd/models/layers/data/__init__.py new file mode 100644 index 0000000..ea50eba --- /dev/null +++ b/hair_service_sd/models/layers/data/__init__.py @@ -0,0 +1,3 @@ +from .wider_face import WiderFaceDetection, detection_collate +from .data_augment import * +from .config import * diff --git a/hair_service_sd/models/layers/data/config.py b/hair_service_sd/models/layers/data/config.py new file mode 100644 index 0000000..591f349 --- /dev/null +++ b/hair_service_sd/models/layers/data/config.py @@ -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 +} + diff --git a/hair_service_sd/models/layers/data/data_augment.py b/hair_service_sd/models/layers/data/data_augment.py new file mode 100644 index 0000000..0c925b1 --- /dev/null +++ b/hair_service_sd/models/layers/data/data_augment.py @@ -0,0 +1,237 @@ +import cv2 +import numpy as np +import random +from utils.box_utils_Retina import matrix_iof + + +def _crop(image, boxes, labels, landm, img_dim): + height, width, _ = image.shape + pad_image_flag = True + + for _ in range(250): + """ + if random.uniform(0, 1) <= 0.2: + scale = 1.0 + else: + scale = random.uniform(0.3, 1.0) + """ + PRE_SCALES = [0.3, 0.45, 0.6, 0.8, 1.0] + scale = random.choice(PRE_SCALES) + short_side = min(width, height) + w = int(scale * short_side) + h = w + + if width == w: + l = 0 + else: + l = random.randrange(width - w) + if height == h: + t = 0 + else: + t = random.randrange(height - h) + roi = np.array((l, t, l + w, t + h)) + + value = matrix_iof(boxes, roi[np.newaxis]) + flag = (value >= 1) + if not flag.any(): + continue + + centers = (boxes[:, :2] + boxes[:, 2:]) / 2 + mask_a = np.logical_and(roi[:2] < centers, centers < roi[2:]).all(axis=1) + boxes_t = boxes[mask_a].copy() + labels_t = labels[mask_a].copy() + landms_t = landm[mask_a].copy() + landms_t = landms_t.reshape([-1, 5, 2]) + + if boxes_t.shape[0] == 0: + continue + + image_t = image[roi[1]:roi[3], roi[0]:roi[2]] + + boxes_t[:, :2] = np.maximum(boxes_t[:, :2], roi[:2]) + boxes_t[:, :2] -= roi[:2] + boxes_t[:, 2:] = np.minimum(boxes_t[:, 2:], roi[2:]) + boxes_t[:, 2:] -= roi[:2] + + # landm + landms_t[:, :, :2] = landms_t[:, :, :2] - roi[:2] + landms_t[:, :, :2] = np.maximum(landms_t[:, :, :2], np.array([0, 0])) + landms_t[:, :, :2] = np.minimum(landms_t[:, :, :2], roi[2:] - roi[:2]) + landms_t = landms_t.reshape([-1, 10]) + + + # make sure that the cropped image contains at least one face > 16 pixel at training image scale + b_w_t = (boxes_t[:, 2] - boxes_t[:, 0] + 1) / w * img_dim + b_h_t = (boxes_t[:, 3] - boxes_t[:, 1] + 1) / h * img_dim + mask_b = np.minimum(b_w_t, b_h_t) > 0.0 + boxes_t = boxes_t[mask_b] + labels_t = labels_t[mask_b] + landms_t = landms_t[mask_b] + + if boxes_t.shape[0] == 0: + continue + + pad_image_flag = False + + return image_t, boxes_t, labels_t, landms_t, pad_image_flag + return image, boxes, labels, landm, pad_image_flag + + +def _distort(image): + + def _convert(image, alpha=1, beta=0): + tmp = image.astype(float) * alpha + beta + tmp[tmp < 0] = 0 + tmp[tmp > 255] = 255 + image[:] = tmp + + image = image.copy() + + if random.randrange(2): + + #brightness distortion + if random.randrange(2): + _convert(image, beta=random.uniform(-32, 32)) + + #contrast distortion + if random.randrange(2): + _convert(image, alpha=random.uniform(0.5, 1.5)) + + image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) + + #saturation distortion + if random.randrange(2): + _convert(image[:, :, 1], alpha=random.uniform(0.5, 1.5)) + + #hue distortion + if random.randrange(2): + tmp = image[:, :, 0].astype(int) + random.randint(-18, 18) + tmp %= 180 + image[:, :, 0] = tmp + + image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR) + + else: + + #brightness distortion + if random.randrange(2): + _convert(image, beta=random.uniform(-32, 32)) + + image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) + + #saturation distortion + if random.randrange(2): + _convert(image[:, :, 1], alpha=random.uniform(0.5, 1.5)) + + #hue distortion + if random.randrange(2): + tmp = image[:, :, 0].astype(int) + random.randint(-18, 18) + tmp %= 180 + image[:, :, 0] = tmp + + image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR) + + #contrast distortion + if random.randrange(2): + _convert(image, alpha=random.uniform(0.5, 1.5)) + + return image + + +def _expand(image, boxes, fill, p): + if random.randrange(2): + return image, boxes + + height, width, depth = image.shape + + scale = random.uniform(1, p) + w = int(scale * width) + h = int(scale * height) + + left = random.randint(0, w - width) + top = random.randint(0, h - height) + + boxes_t = boxes.copy() + boxes_t[:, :2] += (left, top) + boxes_t[:, 2:] += (left, top) + expand_image = np.empty( + (h, w, depth), + dtype=image.dtype) + expand_image[:, :] = fill + expand_image[top:top + height, left:left + width] = image + image = expand_image + + return image, boxes_t + + +def _mirror(image, boxes, landms): + _, width, _ = image.shape + if random.randrange(2): + image = image[:, ::-1] + boxes = boxes.copy() + boxes[:, 0::2] = width - boxes[:, 2::-2] + + # landm + landms = landms.copy() + landms = landms.reshape([-1, 5, 2]) + landms[:, :, 0] = width - landms[:, :, 0] + tmp = landms[:, 1, :].copy() + landms[:, 1, :] = landms[:, 0, :] + landms[:, 0, :] = tmp + tmp1 = landms[:, 4, :].copy() + landms[:, 4, :] = landms[:, 3, :] + landms[:, 3, :] = tmp1 + landms = landms.reshape([-1, 10]) + + return image, boxes, landms + + +def _pad_to_square(image, rgb_mean, pad_image_flag): + if not pad_image_flag: + return image + height, width, _ = image.shape + long_side = max(width, height) + image_t = np.empty((long_side, long_side, 3), dtype=image.dtype) + image_t[:, :] = rgb_mean + image_t[0:0 + height, 0:0 + width] = image + return image_t + + +def _resize_subtract_mean(image, insize, rgb_mean): + interp_methods = [cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_NEAREST, cv2.INTER_LANCZOS4] + interp_method = interp_methods[random.randrange(5)] + image = cv2.resize(image, (insize, insize), interpolation=interp_method) + image = image.astype(np.float32) + image -= rgb_mean + return image.transpose(2, 0, 1) + + +class preproc(object): + + def __init__(self, img_dim, rgb_means): + self.img_dim = img_dim + self.rgb_means = rgb_means + + def __call__(self, image, targets): + assert targets.shape[0] > 0, "this image does not have gt" + + boxes = targets[:, :4].copy() + labels = targets[:, -1].copy() + landm = targets[:, 4:-1].copy() + + image_t, boxes_t, labels_t, landm_t, pad_image_flag = _crop(image, boxes, labels, landm, self.img_dim) + image_t = _distort(image_t) + image_t = _pad_to_square(image_t,self.rgb_means, pad_image_flag) + image_t, boxes_t, landm_t = _mirror(image_t, boxes_t, landm_t) + height, width, _ = image_t.shape + image_t = _resize_subtract_mean(image_t, self.img_dim, self.rgb_means) + boxes_t[:, 0::2] /= width + boxes_t[:, 1::2] /= height + + landm_t[:, 0::2] /= width + landm_t[:, 1::2] /= height + + labels_t = np.expand_dims(labels_t, 1) + targets_t = np.hstack((boxes_t, landm_t, labels_t)) + + return image_t, targets_t diff --git a/hair_service_sd/models/layers/data/wider_face.py b/hair_service_sd/models/layers/data/wider_face.py new file mode 100644 index 0000000..22f56ef --- /dev/null +++ b/hair_service_sd/models/layers/data/wider_face.py @@ -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) diff --git a/hair_service_sd/models/layers/functions/prior_box.py b/hair_service_sd/models/layers/functions/prior_box.py new file mode 100644 index 0000000..80c7f85 --- /dev/null +++ b/hair_service_sd/models/layers/functions/prior_box.py @@ -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 diff --git a/hair_service_sd/models/layers/modules/__init__.py b/hair_service_sd/models/layers/modules/__init__.py new file mode 100644 index 0000000..cf24bdd --- /dev/null +++ b/hair_service_sd/models/layers/modules/__init__.py @@ -0,0 +1,3 @@ +from .multibox_loss import MultiBoxLoss + +__all__ = ['MultiBoxLoss'] diff --git a/hair_service_sd/models/layers/modules/multibox_loss.py b/hair_service_sd/models/layers/modules/multibox_loss.py new file mode 100644 index 0000000..15c7c2f --- /dev/null +++ b/hair_service_sd/models/layers/modules/multibox_loss.py @@ -0,0 +1,125 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.autograd import Variable +from utils.box_utils_Retina import match, log_sum_exp +from 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 diff --git a/hair_service_sd/models/model.py b/hair_service_sd/models/model.py new file mode 100644 index 0000000..e71feba --- /dev/null +++ b/hair_service_sd/models/model.py @@ -0,0 +1,105 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from collections import OrderedDict +import numpy as np +import os + +class Flatten(nn.Module): + def __init__(self): + super(Flatten, self).__init__() + def forward(self, x): + x = x.transpose(3, 2).contiguous() + return x.view(x.size(0), -1) + +class PNet(nn.Module): + def __init__(self): + super(PNet, self).__init__() + self.model_path,_ = os.path.split(os.path.realpath(__file__)) + self.features = nn.Sequential(OrderedDict([ + ('conv1', nn.Conv2d(3, 10, 3, 1)), + ('prelu1', nn.PReLU(10)), + ('pool1', nn.MaxPool2d(2, 2, ceil_mode=True)), + ('conv2', nn.Conv2d(10, 16, 3, 1)), + ('prelu2', nn.PReLU(16)), + ('conv3', nn.Conv2d(16, 32, 3, 1)), + ('prelu3', nn.PReLU(32)) + ])) + self.conv4_1 = nn.Conv2d(32, 2, 1, 1) + self.conv4_2 = nn.Conv2d(32, 4, 1, 1) + weights = np.load('./weights/pnet.npy', allow_pickle=True)[()] + for n, p in self.named_parameters(): + p.data = torch.FloatTensor(weights[n]) + + def forward(self, x): + x = self.features(x) + a = self.conv4_1(x) + b = self.conv4_2(x) + a = F.softmax(a, dim=1) + return b, a + +class RNet(nn.Module): + def __init__(self): + super(RNet, self).__init__() + self.model_path,_ = os.path.split(os.path.realpath(__file__)) + self.features = nn.Sequential(OrderedDict([ + ('conv1', nn.Conv2d(3, 28, 3, 1)), + ('prelu1', nn.PReLU(28)), + ('pool1', nn.MaxPool2d(3, 2, ceil_mode=True)), + ('conv2', nn.Conv2d(28, 48, 3, 1)), + ('prelu2', nn.PReLU(48)), + ('pool2', nn.MaxPool2d(3, 2, ceil_mode=True)), + ('conv3', nn.Conv2d(48, 64, 2, 1)), + ('prelu3', nn.PReLU(64)), + ('flatten', Flatten()), + ('conv4', nn.Linear(576, 128)), + ('prelu4', nn.PReLU(128)) + ])) + self.conv5_1 = nn.Linear(128, 2) + self.conv5_2 = nn.Linear(128, 4) + weights = np.load('./weights/rnet.npy', allow_pickle=True)[()] + for n, p in self.named_parameters(): + p.data = torch.FloatTensor(weights[n]) + + def forward(self, x): + x = self.features(x) + a = self.conv5_1(x) + b = self.conv5_2(x) + a = F.softmax(a, dim=1) + return b, a + +class ONet(nn.Module): + def __init__(self): + super(ONet, self).__init__() + self.model_path,_ = os.path.split(os.path.realpath(__file__)) + self.features = nn.Sequential(OrderedDict([ + ('conv1', nn.Conv2d(3, 32, 3, 1)), + ('prelu1', nn.PReLU(32)), + ('pool1', nn.MaxPool2d(3, 2, ceil_mode=True)), + ('conv2', nn.Conv2d(32, 64, 3, 1)), + ('prelu2', nn.PReLU(64)), + ('pool2', nn.MaxPool2d(3, 2, ceil_mode=True)), + ('conv3', nn.Conv2d(64, 64, 3, 1)), + ('prelu3', nn.PReLU(64)), + ('pool3', nn.MaxPool2d(2, 2, ceil_mode=True)), + ('conv4', nn.Conv2d(64, 128, 2, 1)), + ('prelu4', nn.PReLU(128)), + ('flatten', Flatten()), + ('conv5', nn.Linear(1152, 256)), + ('drop5', nn.Dropout(0.25)), + ('prelu5', nn.PReLU(256)), + ])) + self.conv6_1 = nn.Linear(256, 2) + self.conv6_2 = nn.Linear(256, 4) + self.conv6_3 = nn.Linear(256, 10) + weights = np.load('./weights/onet.npy', allow_pickle=True)[()] + for n, p in self.named_parameters(): + p.data = torch.FloatTensor(weights[n]) + + def forward(self, x): + x = self.features(x) + a = self.conv6_1(x) + b = self.conv6_2(x) + c = self.conv6_3(x) + a = F.softmax(a, dim=1) + return c, b, a diff --git a/hair_service_sd/models/net.py b/hair_service_sd/models/net.py new file mode 100644 index 0000000..beb6040 --- /dev/null +++ b/hair_service_sd/models/net.py @@ -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 + diff --git a/hair_service_sd/models/nms/__init__.py b/hair_service_sd/models/nms/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hair_service_sd/models/nms/py_cpu_nms.py b/hair_service_sd/models/nms/py_cpu_nms.py new file mode 100644 index 0000000..260c5ba --- /dev/null +++ b/hair_service_sd/models/nms/py_cpu_nms.py @@ -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 diff --git a/hair_service_sd/models/resnet.py b/hair_service_sd/models/resnet.py new file mode 100644 index 0000000..28fb73e --- /dev/null +++ b/hair_service_sd/models/resnet.py @@ -0,0 +1,223 @@ +import torch.nn as nn +import torch.utils.model_zoo as model_zoo + + +__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, 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)) + 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.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) + + x = self.avgpool(x) + x = x.view(-1, 512) + + return self.fc(x) + + +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'])) + 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 diff --git a/hair_service_sd/models/retinaface.py b/hair_service_sd/models/retinaface.py new file mode 100644 index 0000000..d530bd8 --- /dev/null +++ b/hair_service_sd/models/retinaface.py @@ -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 models.net import MobileNetV1 as MobileNetV1 +from models.net import FPN as FPN +from 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 \ No newline at end of file diff --git a/hair_service_sd/models/seg_hrnet_ocr_to_onnx.py b/hair_service_sd/models/seg_hrnet_ocr_to_onnx.py new file mode 100644 index 0000000..8ec7104 --- /dev/null +++ b/hair_service_sd/models/seg_hrnet_ocr_to_onnx.py @@ -0,0 +1,688 @@ +# ------------------------------------------------------------------------------ +# Copyright (c) Microsoft +# Licensed under the MIT License. +# Written by Ke Sun (sunk@mail.ustc.edu.cn), Jingyi Xie (hsfzxjy@gmail.com) +# ------------------------------------------------------------------------------ + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import os +import logging +import functools + +import numpy as np + +import torch +import torch.nn as nn +import torch._utils +import torch.nn.functional as F + +ALIGN_CORNERS = True +BN_MOMENTUM = 0.1 +logger = logging.getLogger(__name__) +BatchNorm2d_class = BatchNorm2d = torch.nn.SyncBatchNorm +relu_inplace = True + + +class ModuleHelper: + + @staticmethod + def BNReLU(num_features, bn_type=None, **kwargs): + return nn.Sequential( + BatchNorm2d(num_features, **kwargs), + nn.ReLU() + ) + + @staticmethod + def BatchNorm2d(*args, **kwargs): + return BatchNorm2d + + +def conv3x3(in_planes, out_planes, stride=1): + """3x3 convolution with padding""" + return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, + padding=1, bias=False) + + +class SpatialGather_Module(nn.Module): + """ + Aggregate the context features according to the initial + predicted probability distribution. + Employ the soft-weighted method to aggregate the context. + """ + + def __init__(self, cls_num=0, scale=1): + super(SpatialGather_Module, self).__init__() + self.cls_num = cls_num + self.scale = scale + + def forward(self, feats, probs): + batch_size, c, h, w = probs.size(0), probs.size(1), probs.size(2), probs.size(3) + probs = probs.view(batch_size, c, -1) + feats = feats.view(batch_size, feats.size(1), -1) + feats = feats.permute(0, 2, 1) # batch x hw x c + probs = F.softmax(self.scale * probs, dim=2) # batch x k x hw + ocr_context = torch.matmul(probs, feats) \ + .permute(0, 2, 1).unsqueeze(3) # batch x k x c + return ocr_context + + +class _ObjectAttentionBlock(nn.Module): + ''' + The basic implementation for object context block + Input: + N X C X H X W + Parameters: + in_channels : the dimension of the input feature map + key_channels : the dimension after the key/query transform + scale : choose the scale to downsample the input feature maps (save memory cost) + bn_type : specify the bn type + Return: + N X C X H X W + ''' + + def __init__(self, + in_channels, + key_channels, + scale=1, + bn_type=None): + super(_ObjectAttentionBlock, self).__init__() + self.scale = scale + self.in_channels = in_channels + self.key_channels = key_channels + self.pool = nn.MaxPool2d(kernel_size=(scale, scale)) + self.f_pixel = nn.Sequential( + nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels, + kernel_size=1, stride=1, padding=0, bias=False), + ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type), + nn.Conv2d(in_channels=self.key_channels, out_channels=self.key_channels, + kernel_size=1, stride=1, padding=0, bias=False), + ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type), + ) + self.f_object = nn.Sequential( + nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels, + kernel_size=1, stride=1, padding=0, bias=False), + ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type), + nn.Conv2d(in_channels=self.key_channels, out_channels=self.key_channels, + kernel_size=1, stride=1, padding=0, bias=False), + ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type), + ) + self.f_down = nn.Sequential( + nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels, + kernel_size=1, stride=1, padding=0, bias=False), + ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type), + ) + self.f_up = nn.Sequential( + nn.Conv2d(in_channels=self.key_channels, out_channels=self.in_channels, + kernel_size=1, stride=1, padding=0, bias=False), + ModuleHelper.BNReLU(self.in_channels, bn_type=bn_type), + ) + + def forward(self, x, proxy): + batch_size, h, w = x.size(0), x.size(2), x.size(3) + if self.scale > 1: + x = self.pool(x) + + query = self.f_pixel(x).view(batch_size, self.key_channels, -1) + query = query.permute(0, 2, 1) + key = self.f_object(proxy).view(batch_size, self.key_channels, -1) + value = self.f_down(proxy).view(batch_size, self.key_channels, -1) + value = value.permute(0, 2, 1) + + sim_map = torch.matmul(query, key) + sim_map = (self.key_channels ** -.5) * sim_map + sim_map = F.softmax(sim_map, dim=-1) + + # add bg context ... + context = torch.matmul(sim_map, value) + context = context.permute(0, 2, 1).contiguous() + context = context.view(batch_size, self.key_channels, *x.size()[2:]) + context = self.f_up(context) + if self.scale > 1: + context = F.interpolate(input=context, size=(h, w), mode='bilinear', align_corners=ALIGN_CORNERS) + + return context + + +class ObjectAttentionBlock2D(_ObjectAttentionBlock): + def __init__(self, + in_channels, + key_channels, + scale=1, + bn_type=None): + super(ObjectAttentionBlock2D, self).__init__(in_channels, + key_channels, + scale, + bn_type=bn_type) + + +class SpatialOCR_Module(nn.Module): + """ + Implementation of the OCR module: + We aggregate the global object representation to update the representation for each pixel. + """ + + def __init__(self, + in_channels, + key_channels, + out_channels, + scale=1, + dropout=0.1, + bn_type=None): + super(SpatialOCR_Module, self).__init__() + self.object_context_block = ObjectAttentionBlock2D(in_channels, + key_channels, + scale, + bn_type) + _in_channels = 2 * in_channels + + self.conv_bn_dropout = nn.Sequential( + nn.Conv2d(_in_channels, out_channels, kernel_size=1, padding=0, bias=False), + ModuleHelper.BNReLU(out_channels, bn_type=bn_type), + nn.Dropout2d(dropout) + ) + + def forward(self, feats, proxy_feats): + context = self.object_context_block(feats, proxy_feats) + + output = self.conv_bn_dropout(torch.cat([context, feats], 1)) + + return output + + +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 = BatchNorm2d(planes, momentum=BN_MOMENTUM) + self.relu = nn.ReLU(inplace=relu_inplace) + self.conv2 = conv3x3(planes, planes) + self.bn2 = BatchNorm2d(planes, momentum=BN_MOMENTUM) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out = out + residual + out = self.relu(out) + + return out + + +class Bottleneck(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, stride=1, downsample=None): + super(Bottleneck, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) + self.bn1 = BatchNorm2d(planes, momentum=BN_MOMENTUM) + self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, + padding=1, bias=False) + self.bn2 = BatchNorm2d(planes, momentum=BN_MOMENTUM) + self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, + bias=False) + self.bn3 = BatchNorm2d(planes * self.expansion, + momentum=BN_MOMENTUM) + self.relu = nn.ReLU(inplace=relu_inplace) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out = out + residual + out = self.relu(out) + + return out + + +class HighResolutionModule(nn.Module): + def __init__(self, num_branches, blocks, num_blocks, num_inchannels, + num_channels, fuse_method, multi_scale_output=True): + super(HighResolutionModule, self).__init__() + self._check_branches( + num_branches, blocks, num_blocks, num_inchannels, num_channels) + + self.num_inchannels = num_inchannels + self.fuse_method = fuse_method + self.num_branches = num_branches + + self.multi_scale_output = multi_scale_output + + self.branches = self._make_branches( + num_branches, blocks, num_blocks, num_channels) + self.fuse_layers = self._make_fuse_layers() + self.relu = nn.ReLU(inplace=relu_inplace) + + def _check_branches(self, num_branches, blocks, num_blocks, + num_inchannels, num_channels): + if num_branches != len(num_blocks): + error_msg = 'NUM_BRANCHES({}) <> NUM_BLOCKS({})'.format( + num_branches, len(num_blocks)) + logger.error(error_msg) + raise ValueError(error_msg) + + if num_branches != len(num_channels): + error_msg = 'NUM_BRANCHES({}) <> NUM_CHANNELS({})'.format( + num_branches, len(num_channels)) + logger.error(error_msg) + raise ValueError(error_msg) + + if num_branches != len(num_inchannels): + error_msg = 'NUM_BRANCHES({}) <> NUM_INCHANNELS({})'.format( + num_branches, len(num_inchannels)) + logger.error(error_msg) + raise ValueError(error_msg) + + def _make_one_branch(self, branch_index, block, num_blocks, num_channels, + stride=1): + downsample = None + if stride != 1 or \ + self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.num_inchannels[branch_index], + num_channels[branch_index] * block.expansion, + kernel_size=1, stride=stride, bias=False), + BatchNorm2d(num_channels[branch_index] * block.expansion, + momentum=BN_MOMENTUM), + ) + + layers = [] + layers.append(block(self.num_inchannels[branch_index], + num_channels[branch_index], stride, downsample)) + self.num_inchannels[branch_index] = \ + num_channels[branch_index] * block.expansion + for i in range(1, num_blocks[branch_index]): + layers.append(block(self.num_inchannels[branch_index], + num_channels[branch_index])) + + return nn.Sequential(*layers) + + def _make_branches(self, num_branches, block, num_blocks, num_channels): + branches = [] + + for i in range(num_branches): + branches.append( + self._make_one_branch(i, block, num_blocks, num_channels)) + + return nn.ModuleList(branches) + + def _make_fuse_layers(self): + if self.num_branches == 1: + return None + + num_branches = self.num_branches + num_inchannels = self.num_inchannels + fuse_layers = [] + for i in range(num_branches if self.multi_scale_output else 1): + fuse_layer = [] + for j in range(num_branches): + if j > i: + fuse_layer.append(nn.Sequential( + nn.Conv2d(num_inchannels[j], + num_inchannels[i], + 1, + 1, + 0, + bias=False), + BatchNorm2d(num_inchannels[i], momentum=BN_MOMENTUM))) + elif j == i: + fuse_layer.append(None) + else: + conv3x3s = [] + for k in range(i - j): + if k == i - j - 1: + num_outchannels_conv3x3 = num_inchannels[i] + conv3x3s.append(nn.Sequential( + nn.Conv2d(num_inchannels[j], + num_outchannels_conv3x3, + 3, 2, 1, bias=False), + BatchNorm2d(num_outchannels_conv3x3, + momentum=BN_MOMENTUM))) + else: + num_outchannels_conv3x3 = num_inchannels[j] + conv3x3s.append(nn.Sequential( + nn.Conv2d(num_inchannels[j], + num_outchannels_conv3x3, + 3, 2, 1, bias=False), + BatchNorm2d(num_outchannels_conv3x3, + momentum=BN_MOMENTUM), + nn.ReLU(inplace=relu_inplace))) + fuse_layer.append(nn.Sequential(*conv3x3s)) + fuse_layers.append(nn.ModuleList(fuse_layer)) + + return nn.ModuleList(fuse_layers) + + def get_num_inchannels(self): + return self.num_inchannels + + def forward(self, x): + if self.num_branches == 1: + return [self.branches[0](x[0])] + + for i in range(self.num_branches): + x[i] = self.branches[i](x[i]) + + x_fuse = [] + for i in range(len(self.fuse_layers)): + y = x[0] if i == 0 else self.fuse_layers[i][0](x[0]) + for j in range(1, self.num_branches): + if i == j: + y = y + x[j] + elif j > i: + width_output = x[i].shape[-1] + height_output = x[i].shape[-2] + y = y + F.interpolate( + self.fuse_layers[i][j](x[j]), + size=[height_output, width_output], + mode='bilinear', align_corners=ALIGN_CORNERS) + else: + y = y + self.fuse_layers[i][j](x[j]) + x_fuse.append(self.relu(y)) + + return x_fuse + + +blocks_dict = { + 'BASIC': BasicBlock, + 'BOTTLENECK': Bottleneck +} + + +class HighResolutionNet(nn.Module): + + def __init__(self, config, **kwargs): + global ALIGN_CORNERS + extra = config.MODEL.EXTRA + super(HighResolutionNet, self).__init__() + ALIGN_CORNERS = config.MODEL.ALIGN_CORNERS + + # stem net + self.conv1 = nn.Conv2d(config.MODEL.INPUTC, 64, kernel_size=3, stride=2, padding=1, + bias=False) + self.bn1 = BatchNorm2d(64, momentum=BN_MOMENTUM) + self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1, + bias=False) + self.bn2 = BatchNorm2d(64, momentum=BN_MOMENTUM) + self.relu = nn.ReLU(inplace=relu_inplace) + + self.stage1_cfg = extra['STAGE1'] + num_channels = self.stage1_cfg['NUM_CHANNELS'][0] + block = blocks_dict[self.stage1_cfg['BLOCK']] + num_blocks = self.stage1_cfg['NUM_BLOCKS'][0] + self.layer1 = self._make_layer(block, 64, num_channels, num_blocks) + stage1_out_channel = block.expansion * num_channels + + self.stage2_cfg = extra['STAGE2'] + num_channels = self.stage2_cfg['NUM_CHANNELS'] + block = blocks_dict[self.stage2_cfg['BLOCK']] + num_channels = [ + num_channels[i] * block.expansion for i in range(len(num_channels))] + self.transition1 = self._make_transition_layer( + [stage1_out_channel], num_channels) + self.stage2, pre_stage_channels = self._make_stage( + self.stage2_cfg, num_channels) + + self.stage3_cfg = extra['STAGE3'] + num_channels = self.stage3_cfg['NUM_CHANNELS'] + block = blocks_dict[self.stage3_cfg['BLOCK']] + num_channels = [ + num_channels[i] * block.expansion for i in range(len(num_channels))] + self.transition2 = self._make_transition_layer( + pre_stage_channels, num_channels) + self.stage3, pre_stage_channels = self._make_stage( + self.stage3_cfg, num_channels) + + self.stage4_cfg = extra['STAGE4'] + num_channels = self.stage4_cfg['NUM_CHANNELS'] + block = blocks_dict[self.stage4_cfg['BLOCK']] + num_channels = [ + num_channels[i] * block.expansion for i in range(len(num_channels))] + self.transition3 = self._make_transition_layer( + pre_stage_channels, num_channels) + self.stage4, pre_stage_channels = self._make_stage( + self.stage4_cfg, num_channels, multi_scale_output=True) + + last_inp_channels = np.int(np.sum(pre_stage_channels)) + ocr_mid_channels = config.MODEL.OCR.MID_CHANNELS + ocr_key_channels = config.MODEL.OCR.KEY_CHANNELS + + self.conv3x3_ocr = nn.Sequential( + nn.Conv2d(last_inp_channels, ocr_mid_channels, + kernel_size=3, stride=1, padding=1), + BatchNorm2d(ocr_mid_channels), + nn.ReLU(inplace=relu_inplace), + ) + self.ocr_gather_head = SpatialGather_Module(config.DATASET.NUM_CLASSES) + + self.ocr_distri_head = SpatialOCR_Module(in_channels=ocr_mid_channels, + key_channels=ocr_key_channels, + out_channels=ocr_mid_channels, + scale=1, + dropout=0.05, + ) + self.cls_head = nn.Conv2d( + ocr_mid_channels, config.DATASET.NUM_CLASSES, kernel_size=1, stride=1, padding=0, bias=True) + + self.aux_head = nn.Sequential( + nn.Conv2d(last_inp_channels, last_inp_channels, + kernel_size=1, stride=1, padding=0), + BatchNorm2d(last_inp_channels), + nn.ReLU(inplace=relu_inplace), + nn.Conv2d(last_inp_channels, config.DATASET.NUM_CLASSES, + kernel_size=1, stride=1, padding=0, bias=True) + ) + + def _make_transition_layer( + self, num_channels_pre_layer, num_channels_cur_layer): + num_branches_cur = len(num_channels_cur_layer) + num_branches_pre = len(num_channels_pre_layer) + + transition_layers = [] + for i in range(num_branches_cur): + if i < num_branches_pre: + if num_channels_cur_layer[i] != num_channels_pre_layer[i]: + transition_layers.append(nn.Sequential( + nn.Conv2d(num_channels_pre_layer[i], + num_channels_cur_layer[i], + 3, + 1, + 1, + bias=False), + BatchNorm2d( + num_channels_cur_layer[i], momentum=BN_MOMENTUM), + nn.ReLU(inplace=relu_inplace))) + else: + transition_layers.append(None) + else: + conv3x3s = [] + for j in range(i + 1 - num_branches_pre): + inchannels = num_channels_pre_layer[-1] + outchannels = num_channels_cur_layer[i] \ + if j == i - num_branches_pre else inchannels + conv3x3s.append(nn.Sequential( + nn.Conv2d( + inchannels, outchannels, 3, 2, 1, bias=False), + BatchNorm2d(outchannels, momentum=BN_MOMENTUM), + nn.ReLU(inplace=relu_inplace))) + transition_layers.append(nn.Sequential(*conv3x3s)) + + return nn.ModuleList(transition_layers) + + def _make_layer(self, block, inplanes, planes, blocks, stride=1): + downsample = None + if stride != 1 or inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(inplanes, planes * block.expansion, + kernel_size=1, stride=stride, bias=False), + BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM), + ) + + layers = [] + layers.append(block(inplanes, planes, stride, downsample)) + inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(inplanes, planes)) + + return nn.Sequential(*layers) + + def _make_stage(self, layer_config, num_inchannels, + multi_scale_output=True): + num_modules = layer_config['NUM_MODULES'] + num_branches = layer_config['NUM_BRANCHES'] + num_blocks = layer_config['NUM_BLOCKS'] + num_channels = layer_config['NUM_CHANNELS'] + block = blocks_dict[layer_config['BLOCK']] + fuse_method = layer_config['FUSE_METHOD'] + + modules = [] + for i in range(num_modules): + # multi_scale_output is only used last module + if not multi_scale_output and i == num_modules - 1: + reset_multi_scale_output = False + else: + reset_multi_scale_output = True + modules.append( + HighResolutionModule(num_branches, + block, + num_blocks, + num_inchannels, + num_channels, + fuse_method, + reset_multi_scale_output) + ) + num_inchannels = modules[-1].get_num_inchannels() + + return nn.Sequential(*modules), num_inchannels + + def forward(self, x): + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + x = self.conv2(x) + x = self.bn2(x) + x = self.relu(x) + x = self.layer1(x) + + x_list = [] + for i in range(self.stage2_cfg['NUM_BRANCHES']): + if self.transition1[i] is not None: + x_list.append(self.transition1[i](x)) + else: + x_list.append(x) + y_list = self.stage2(x_list) + + x_list = [] + for i in range(self.stage3_cfg['NUM_BRANCHES']): + if self.transition2[i] is not None: + if i < self.stage2_cfg['NUM_BRANCHES']: + x_list.append(self.transition2[i](y_list[i])) + else: + x_list.append(self.transition2[i](y_list[-1])) + else: + x_list.append(y_list[i]) + y_list = self.stage3(x_list) + + x_list = [] + for i in range(self.stage4_cfg['NUM_BRANCHES']): + if self.transition3[i] is not None: + if i < self.stage3_cfg['NUM_BRANCHES']: + x_list.append(self.transition3[i](y_list[i])) + else: + x_list.append(self.transition3[i](y_list[-1])) + else: + x_list.append(y_list[i]) + x = self.stage4(x_list) + + # Upsampling + x0_h, x0_w = x[0].size(2), x[0].size(3) + x1 = F.interpolate(x[1], size=(x0_h, x0_w), + mode='bilinear', align_corners=ALIGN_CORNERS) + x2 = F.interpolate(x[2], size=(x0_h, x0_w), + mode='bilinear', align_corners=ALIGN_CORNERS) + x3 = F.interpolate(x[3], size=(x0_h, x0_w), + mode='bilinear', align_corners=ALIGN_CORNERS) + + feats = torch.cat([x[0], x1, x2, x3], 1) + + # out_aux_seg = [] + + # ocr + out_aux = self.aux_head(feats) + # compute contrast feature + feats = self.conv3x3_ocr(feats) + + context = self.ocr_gather_head(feats, out_aux) + feats = self.ocr_distri_head(feats, context) + + out = self.cls_head(feats) + + out = F.interpolate(out, size=(out.size(2) * 4, out.size(3) * 4), + mode='bilinear', align_corners=ALIGN_CORNERS) + + return out + + def init_weights(self, pretrained='', ): + logger.info('=> init weights from normal distribution') + for name, m in self.named_modules(): + if any(part in name for part in {'cls', 'aux', 'ocr'}): + # print('skipped', name) + continue + if isinstance(m, nn.Conv2d): + nn.init.normal_(m.weight, std=0.001) + elif isinstance(m, BatchNorm2d_class): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + if os.path.isfile(pretrained): + pretrained_dict = torch.load(pretrained, map_location={'cuda:0': 'cpu'}) + logger.info('=> loading pretrained model {}'.format(pretrained)) + model_dict = self.state_dict() + pretrained_dict = {k.replace('last_layer', 'aux_head').replace('model.', ''): v for k, v in + pretrained_dict.items()} + print(set(model_dict) - set(pretrained_dict)) + print(set(pretrained_dict) - set(model_dict)) + pretrained_dict = {k: v for k, v in pretrained_dict.items() + if k in model_dict.keys() and not k.startswith('conv1')} + # for k, _ in pretrained_dict.items(): + # logger.info( + # '=> loading {} pretrained model {}'.format(k, pretrained)) + model_dict.update(pretrained_dict) + self.load_state_dict(model_dict) + elif pretrained: + raise RuntimeError('No such file {}'.format(pretrained)) + + +def get_seg_model(cfg, **kwargs): + model = HighResolutionNet(cfg, **kwargs) + # model.init_weights(cfg.MODEL.PRETRAINED) + + return model diff --git a/hair_service_sd/models/yolo.py b/hair_service_sd/models/yolo.py new file mode 100644 index 0000000..c9e6c49 --- /dev/null +++ b/hair_service_sd/models/yolo.py @@ -0,0 +1,233 @@ +import argparse + +from models.experimental import * + + +class Detect(nn.Module): + def __init__(self, nc=80, anchors=()): # detection layer + super(Detect, self).__init__() + self.stride = None # strides computed during build + self.nc = nc # number of classes + self.no = nc + 5 # number of outputs per anchor + self.nl = len(anchors) # number of detection layers + self.na = len(anchors[0]) // 2 # number of anchors + self.grid = [torch.zeros(1)] * self.nl # init grid + a = torch.tensor(anchors).float().view(self.nl, -1, 2) + self.register_buffer('anchors', a) # shape(nl,na,2) + self.register_buffer('anchor_grid', a.clone().view(self.nl, 1, -1, 1, 1, 2)) # shape(nl,1,na,1,1,2) + self.export = False # onnx export + + def forward(self, x): + # x = x.copy() # for profiling + z = [] # inference output + self.training |= self.export + for i in range(self.nl): + bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85) + x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous() + + if not self.training: # inference + if self.grid[i].shape[2:4] != x[i].shape[2:4]: + self.grid[i] = self._make_grid(nx, ny).to(x[i].device) + + y = x[i].sigmoid() + y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i] # xy + y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh + z.append(y.view(bs, -1, self.no)) + + return x if self.training else (torch.cat(z, 1), x) + + @staticmethod + def _make_grid(nx=20, ny=20): + yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)]) + return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float() + + +class Model(nn.Module): + def __init__(self, model_cfg='yolov5s.yaml', ch=3, nc=None): # model, input channels, number of classes + super(Model, self).__init__() + if type(model_cfg) is dict: + self.md = model_cfg # model dict + else: # is *.yaml + with open(model_cfg) as f: + self.md = yaml.load(f, Loader=yaml.FullLoader) # model dict + + # Define model + if nc: + self.md['nc'] = nc # override yaml value + self.model, self.save = parse_model(self.md, ch=[ch]) # model, savelist, ch_out + # print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))]) + + # Build strides, anchors + m = self.model[-1] # Detect() + m.stride = torch.tensor([128 / x.shape[-2] for x in self.forward(torch.zeros(1, ch, 128, 128))]) # forward + m.anchors /= m.stride.view(-1, 1, 1) + check_anchor_order(m) + self.stride = m.stride + + # Init weights, biases + torch_utils.initialize_weights(self) + self._initialize_biases() # only run once + torch_utils.model_info(self) + print('') + + def forward(self, x, augment=False, profile=False): + if augment: + img_size = x.shape[-2:] # height, width + s = [0.83, 0.67] # scales + y = [] + for i, xi in enumerate((x, + torch_utils.scale_img(x.flip(3), s[0]), # flip-lr and scale + torch_utils.scale_img(x, s[1]), # scale + )): + # cv2.imwrite('img%g.jpg' % i, 255 * xi[0].numpy().transpose((1, 2, 0))[:, :, ::-1]) + y.append(self.forward_once(xi)[0]) + + y[1][..., :4] /= s[0] # scale + y[1][..., 0] = img_size[1] - y[1][..., 0] # flip lr + y[2][..., :4] /= s[1] # scale + return torch.cat(y, 1), None # augmented inference, train + else: + return self.forward_once(x, profile) # single-scale inference, train + + def forward_once(self, x, profile=False): + y, dt = [], [] # outputs + for m in self.model: + if m.f != -1: # if not from previous layer + x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers + + if profile: + try: + import thop + o = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 # FLOPS + except: + o = 0 + t = torch_utils.time_synchronized() + for _ in range(10): + _ = m(x) + dt.append((torch_utils.time_synchronized() - t) * 100) + print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type)) + + x = m(x) # run + y.append(x if m.i in self.save else None) # save output + + if profile: + print('%.1fms total' % sum(dt)) + return x + + def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency + # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1. + m = self.model[-1] # Detect() module + for f, s in zip(m.f, m.stride): #  from + mi = self.model[f % m.i] + b = mi.bias.view(m.na, -1) # conv.bias(255) to (3,85) + b[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image) + b[:, 5:] += math.log(0.6 / (m.nc - 0.99)) if cf is None else torch.log(cf / cf.sum()) # cls + mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True) + + def _print_biases(self): + m = self.model[-1] # Detect() module + for f in sorted([x % m.i for x in m.f]): #  from + b = self.model[f].bias.detach().view(m.na, -1).T # conv.bias(255) to (3,85) + print(('%g Conv2d.bias:' + '%10.3g' * 6) % (f, *b[:5].mean(1).tolist(), b[5:].mean())) + + # def _print_weights(self): + # for m in self.model.modules(): + # if type(m) is Bottleneck: + # print('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights + + def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers + print('Fusing layers...') + for m in self.model.modules(): + if type(m) is Conv: + m.conv = torch_utils.fuse_conv_and_bn(m.conv, m.bn) # update conv + m.bn = None # remove batchnorm + m.forward = m.fuseforward # update forward + torch_utils.model_info(self) + + +def parse_model(md, ch): # model_dict, input_channels(3) + print('\n%3s%15s%3s%10s %-40s%-30s' % ('', 'from', 'n', 'params', 'module', 'arguments')) + anchors, nc, gd, gw = md['anchors'], md['nc'], md['depth_multiple'], md['width_multiple'] + na = (len(anchors[0]) // 2) # number of anchors + no = na * (nc + 5) # number of outputs = anchors * (classes + 5) + + layers, save, c2 = [], [], ch[-1] # layers, savelist, ch out + for i, (f, n, m, args) in enumerate(md['backbone'] + md['head']): # from, number, module, args + m = eval(m) if isinstance(m, str) else m # eval strings + for j, a in enumerate(args): + try: + args[j] = eval(a) if isinstance(a, str) else a # eval strings + except: + pass + + n = max(round(n * gd), 1) if n > 1 else n # depth gain + if m in [nn.Conv2d, Conv, Bottleneck, SPP, DWConv, MixConv2d, Focus, ConvPlus, BottleneckCSP]: + c1, c2 = ch[f], args[0] + + # Normal + # if i > 0 and args[0] != no: # channel expansion factor + # ex = 1.75 # exponential (default 2.0) + # e = math.log(c2 / ch[1]) / math.log(2) + # c2 = int(ch[1] * ex ** e) + # if m != Focus: + c2 = make_divisible(c2 * gw, 8) if c2 != no else c2 + + # Experimental + # if i > 0 and args[0] != no: # channel expansion factor + # ex = 1 + gw # exponential (default 2.0) + # ch1 = 32 # ch[1] + # e = math.log(c2 / ch1) / math.log(2) # level 1-n + # c2 = int(ch1 * ex ** e) + # if m != Focus: + # c2 = make_divisible(c2, 8) if c2 != no else c2 + + args = [c1, c2, *args[1:]] + if m is BottleneckCSP: + args.insert(2, n) + n = 1 + elif m is nn.BatchNorm2d: + args = [ch[f]] + elif m is Concat: + c2 = sum([ch[-1 if x == -1 else x + 1] for x in f]) + elif m is Detect: + f = f or list(reversed([(-1 if j == i else j - 1) for j, x in enumerate(ch) if x == no])) + else: + c2 = ch[f] + + m_ = nn.Sequential(*[m(*args) for _ in range(n)]) if n > 1 else m(*args) # module + t = str(m)[8:-2].replace('__main__.', '') # module type + np = sum([x.numel() for x in m_.parameters()]) # number params + m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params + print('%3s%15s%3s%10.0f %-40s%-30s' % (i, f, n, np, t, args)) # print + save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist + layers.append(m_) + ch.append(c2) + return nn.Sequential(*layers), sorted(save) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml') + parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') + opt = parser.parse_args() + opt.cfg = check_file(opt.cfg) # check file + device = torch_utils.select_device(opt.device) + + # Create model + model = Model(opt.cfg).to(device) + model.train() + + # Profile + # img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device) + # y = model(img, profile=True) + + # ONNX export + # model.model[-1].export = True + # torch.onnx.export(model, img, opt.cfg.replace('.yaml', '.onnx'), verbose=True, opset_version=11) + + # Tensorboard + # from torch.utils.tensorboard import SummaryWriter + # tb_writer = SummaryWriter() + # print("Run 'tensorboard --logdir=models/runs' to view tensorboard at http://localhost:6006/") + # tb_writer.add_graph(model.model, img) # add model to tensorboard + # tb_writer.add_image('test', img[0], dataformats='CWH') # add model to tensorboard diff --git a/hair_service_sd/momocv/BigResNetStable.pth b/hair_service_sd/momocv/BigResNetStable.pth new file mode 100644 index 0000000..cb273a8 --- /dev/null +++ b/hair_service_sd/momocv/BigResNetStable.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ece946bd0ae8e4142eb8667f8b273b75a9564a93b08656ae50ff064c7cff08bf +size 7375394 diff --git a/hair_service_sd/momocv/BigResNetStable.py b/hair_service_sd/momocv/BigResNetStable.py new file mode 100644 index 0000000..b2af26c --- /dev/null +++ b/hair_service_sd/momocv/BigResNetStable.py @@ -0,0 +1,382 @@ +import torch +import torch.nn as nn +import math +import os +import cv2 +import numpy as np +# from utils.DATAIMG import DATAIMG +from utils import landmark_processor +import glob + +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 + '/relu', nn.ReLU()), + ) + +class BasicBlock(nn.Module): + def __init__(self, name, inplanes, planes, stride=2): + super(BasicBlock, self).__init__() + self.op_name = name + self.conv1 = conv_relu(name + '/conv1', inplanes, planes, kernel_size=3, stride=stride, padding=1) + self.conv2 = conv_relu(name + '/conv2', planes, planes, kernel_size=3, stride=1, padding=1) + self.downsample = conv_relu(name + '/sc_conv', inplanes, planes, kernel_size=1, stride=stride) + + def forward(self, x): + residual = x + out = self.conv1(x) + out = self.conv2(out) + if self.downsample is not None: + residual = self.downsample(x) + ret = out + residual + return ret + + +class BigResNetStable(nn.Module): + def __init__(self, name, in_channels, out_channels): + super(BigResNetStable, self).__init__() + self.op_name = name + + op_list = [] + + op_list += [conv_relu(name + '/first_conv', in_channels, 32, kernel_size=5, stride=2, padding=2)] + + ch_num = [32, 48, 64, 96, 128] + + op_list += [BasicBlock(name + '/stage%d' % (i + 1), ch_num[i], ch_num[i + 1]) for i in range(len(ch_num) - 1)] + + op_list += [flatten(name + '/flatten', 1)] + + op_list1 = [op_name(name + '/FC1/FC', nn.Linear(2048, 512)), + op_name(name + '/FC1/relu', nn.ReLU())] + + fullyconnected1 = [op_name(name + '/FC2', nn.Linear(512, out_channels))] + poselayer = [op_name('poselayer', nn.Linear(512, 3))] + tracking_probe = [op_name('tracking_probe', nn.Linear(2048, 1))] + occlusion_probe = [op_name('occlusion_probe', nn.Linear(2048, 87))] + + self.features = nn.Sequential(*op_list) + self.fc1 = nn.Sequential(*op_list1) + self.fc2 = nn.Sequential(*fullyconnected1) + self.poselayer = nn.Sequential(*poselayer) + self.tracking_probe = nn.Sequential(*tracking_probe) + self.occlusion_probe = nn.Sequential(*occlusion_probe) + + def forward(self, x): + features = self.features(x) + fc1 = self.fc1(features) + fullyconnected1 = self.fc2(fc1) + poselayer = self.poselayer(fc1) + tracking_probe = self.tracking_probe(features) + occlusion_probe = self.occlusion_probe(features) + return fullyconnected1.cpu().numpy(), poselayer.cpu().numpy(), tracking_probe.cpu().numpy(), occlusion_probe.cpu().numpy() + +class MomocvFaceAlignment(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 = BigResNetStable('BigResNetStable', 3, 174) + + self.model_path, _ = os.path.split(os.path.realpath(__file__)) + + weights = torch.load(os.path.join(self.model_path, 'BigResNetStable.pth'), + map_location=lambda storage, loc: storage) + self.face_alignment_net.load_state_dict(weights) + self.face_alignment_net.to(self.device) + self.face_alignment_net.eval() + + self.trackingFaceRects = [] + + + def crop_img(self, image, labels=None,img_size = 128): + ''' + :param imageimgs: Original graph + :param labels: The boxes of the original picture ; labels is a ndarray: [list,list] + :return: Coordinates and categories relative to the original + ''' + # Preprocessing + labels = np.array([labels]) + boxes_ret = np.zeros(labels.shape) + crop_imgs = [] + ret_M = [] + for i in range(labels.shape[0]): + box_orig = np.array(labels[i, :]) + box = box_orig.astype(np.int32).copy() + center = np.array([(box[0] + box[2]) / 2, (box[1] + box[3]) / 2]).astype(np.int32) + max_lenth = int(max(box[3] - box[1], box[2] - box[0]) / 2 * 1.0) + up = int(center[1] - max_lenth) + down = int(center[1] + max_lenth) + left = int(center[0] - max_lenth) + right = int(center[0] + max_lenth) + if up < 0: + up = 0 + down = max_lenth * 2 + if down > image.shape[0]: + down = image.shape[0] + if left < 0: + left = 0 + right = max_lenth * 2 + if right > image.shape[1]: + right = image.shape[1] + + crop_img = image[up:down, left:right, :].copy() + crop_img = cv2.resize(crop_img, (img_size, img_size)) + box_orig[0] -= left + box_orig[2] -= left + box_orig[1] -= up + box_orig[3] -= up + box_orig[0] *= img_size / (right - left) + box_orig[2] *= img_size / (right - left) + box_orig[1] *= img_size / (down - up) + box_orig[3] *= img_size / (down - up) + crop_imgs.append(crop_img) + ret_M.append(np.array([img_size / (right - left), img_size / (down - up), left, up])) + boxes_ret[i, :] = box_orig + return crop_imgs, boxes_ret, ret_M + + + def detect_from_bbox(self, img, bboxs): + dst_size = 128 + landmarks_res = [] + with torch.no_grad(): + input_numpy = np.zeros((len(bboxs), 3, dst_size, dst_size), dtype=np.float32) + for ix, bbox in enumerate(bboxs): + crop_imgs, boxes_ret, ret_M = self.crop_img(img, bbox, dst_size) + input_numpy[ix, :, :, :] = crop_imgs[0].transpose((2, 0, 1)).astype(np.float32) + + # cv2.imshow('tmp', tmp) + # cv2.waitKey() + + in_tensor = torch.from_numpy(input_numpy) + in_tensor = in_tensor.to(self.device) + fullyconnected1, poselayer, tracking_probe, occlusion_probe = self.face_alignment_net(in_tensor) + for ix, pts in enumerate(fullyconnected1): + orig_pts = ((np.reshape(pts, (2, 87)).transpose((1, 0))) * dst_size) + # print('pts',orig_pts) + # landmark2 = landmark2.cpu().detach().numpy()[0].reshape(2, -1).transpose((1, 0)).reshape(-1) + # orig_pts = ((orig_pts + 0.5) * 128) + orig_pts[:, 0] = orig_pts[:, 0] / ret_M[0][0] + ret_M[0][2] + orig_pts[:, 1] = orig_pts[:, 1] / ret_M[0][1] + ret_M[0][3] + + # orig_pts[0] = orig_pts[0] / ret_M[0][0] + ret_M[0][2] + # orig_pts[2] = orig_pts[2] / ret_M[0][0] + ret_M[0][2] + # orig_pts[1] = orig_pts[1] / ret_M[0][1] + ret_M[0][3] + # orig_pts[3] = orig_pts[3] / ret_M[0][1] + ret_M[0][3] + # cur_landmark2.append(orig_pts) + # orig_pts = landmark_processor.transform_points(orig_pts, all_mat[ix], invert=True) + landmarks_res.append(orig_pts) + tracking_probe = 1 / (1 + np.exp(-tracking_probe)) + occlusion_probe = 1 / (1 + np.exp(-occlusion_probe)) + return landmarks_res, poselayer, tracking_probe, occlusion_probe + + def detect(self, img, landmarks): + dst_size = 128 + 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_mmcv(landmark, dst_size) + + + all_mat.append(M) + + tmp = cv2.warpAffine(img, M, (dst_size, dst_size)) + + + input_numpy[ix, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) + # + # cv2.imshow('tmp', tmp) + # cv2.waitKey() + + in_tensor = torch.from_numpy(input_numpy) + in_tensor = in_tensor.to(self.device) + fullyconnected1, poselayer, tracking_probe, occlusion_probe = self.face_alignment_net(in_tensor) + for ix, pts in enumerate(fullyconnected1): + orig_pts = (np.reshape(pts, (2, 87)).transpose((1, 0)) * dst_size) + orig_pts = landmark_processor.transform_points(orig_pts, all_mat[ix], invert=True) + landmarks_res.append(orig_pts) + tracking_probe = 1 / (1 + np.exp(-tracking_probe)) + occlusion_probe = 1 / (1 + np.exp(-occlusion_probe)) + return landmarks_res, poselayer, tracking_probe, occlusion_probe + + def forward(self, images, landmarks): + dst_size = 128 + landmarks_res = [] + 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 = [] + for ix, img in enumerate(images): + landmark = landmarks[ix] + + M = landmark_processor.get_transform_mat_mmcv(landmark, dst_size) + + all_mat.append(M) + + tmp = cv2.warpAffine(img, M, (dst_size, dst_size)) + + input_numpy[ix, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) + # + # cv2.imshow('tmp', tmp) + # cv2.waitKey() + + in_tensor = torch.from_numpy(input_numpy) + in_tensor = in_tensor.to(self.device) + fullyconnected1, poselayer, tracking_probe, occlusion_probe = self.face_alignment_net(in_tensor) + for ix, pts in enumerate(fullyconnected1): + orig_pts = (np.reshape(pts, (2, 87)).transpose((1, 0)) * dst_size) + orig_pts = landmark_processor.transform_points(orig_pts, all_mat[ix], invert=True) + landmarks_res.append(orig_pts) + tracking_probe = 1 / (1 + np.exp(-tracking_probe)) + occlusion_probe = 1 / (1 + np.exp(-occlusion_probe)) + return landmarks_res, poselayer, tracking_probe, occlusion_probe + + def stable_forward(self, image, detected_faces): + for face_rect in detected_faces: + if len(self.trackingFaceRects) == 0: + new_tracking_rect = [face_rect, True, [0, 0], 0] + self.trackingFaceRects.append(new_tracking_rect) + + with torch.no_grad(): + landmarks = [] + for tracking_face_rect in 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 = 128 * 0.8 / min(d[2] - d[0], d[3] - d[1]) + dst_center = np.array([0.5, 0.5]) * 128 + offset = dst_center - src_center + print('hello') + else: + dst_left_anchor = np.array([0.395, 0.52]) * 128 + dst_right_anchor = np.array([1 - 0.395, 0.52]) * 128 + # use last anchors + src_center = (tracking_face_rect[2][0] + tracking_face_rect[2][1]) / 2 + rotate_radian = math.atan2(tracking_face_rect[2][1][1] - tracking_face_rect[2][0][1], tracking_face_rect[2][1][0] - tracking_face_rect[2][0][0]) + rotate_degree = rotate_radian / math.pi * 180 + print('degree', rotate_degree) + dst_anchor_len = cv2.norm(dst_left_anchor - dst_right_anchor) + src_anchor_len = cv2.norm(tracking_face_rect[2][0], tracking_face_rect[2][1]) + scale = dst_anchor_len / src_anchor_len + dst_center = (dst_left_anchor + dst_right_anchor) / 2 + offset = dst_center - src_center + + M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), rotate_degree, scale) + M[:, 2] += offset + + inp = cv2.warpAffine(image, M, (128, 128)) + + cv2.imshow('inp_stable', inp) + # cv2.waitKey() + + inp = inp.transpose((2, 0, 1)).astype(np.float32) + inp = inp[np.newaxis, :, :, :] + + in_tensor = torch.from_numpy(inp) + in_tensor = in_tensor.to(self.device) + fullyconnected1, poselayer, tracking_probe, occlusion_probe = self.face_alignment_net(in_tensor) + fullyconnected1 = fullyconnected1[0] + poselayer = poselayer[0] + tracking_probe = tracking_probe[0] + occlusion_probe = occlusion_probe[0] + orig_pts = (np.reshape(fullyconnected1, (2, 87)).transpose((1, 0)) * 128) + + orig_pts = landmark_processor.transform_points(orig_pts, M, invert=True) + # orig_pts = orig_pts.transpose((1, 0)) + fullyconnected1 = orig_pts + tracking_probe = 1 / (1 + np.exp(-tracking_probe)) + occlusion_probe = 1 / (1 + np.exp(-occlusion_probe)) + + # update tracking infos + tracking_face_rect[1] = False + tracking_face_rect[2] = [fullyconnected1[51], fullyconnected1[57]] + tracking_face_rect[3] = rotate_degree + + landmarks.append(fullyconnected1) + return landmarks + +if __name__ == '__main__': + all_jpegs = glob.glob(r'E:\deepfacelab_data\expression_dst\7201806132018061208311920180612083119\*.jpg') + for s_filename_path in all_jpegs: + img = cv2.imread(s_filename_path) + + dflpng = DATAIMG(str(s_filename_path), print_on_no_embedded_data=True) + if dflpng is None: + print('ERROR') + + landmarks = dflpng.get_landmarks() + + mmcv = MomocvFaceAlignment() + fullyconnected1, poselayer, tracking_probe, occlusion_probe = mmcv.forward([img], [landmarks]) + 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() + + + # anchor_dis = 0.445 + # dst_size = 128 + # anchors_template = np.array([[anchor_dis, 0.52], [1 - anchor_dis, 0.52]]) * dst_size + # + # src_center = (landmarks[31] + landmarks[35]) / 2 + # src_left_anchor = landmarks[31] + # src_right_anchor = landmarks[35] + # rotate_radian = math.atan2(src_right_anchor[1] - src_left_anchor[1], src_right_anchor[0] - src_left_anchor[0]) + # rotate_degree = rotate_radian / math.pi * 180 + # dst_eye_len = np.sqrt(np.sum((anchors_template[0] - anchors_template[1]) ** 2)) + # src_eye_len = np.sqrt(np.sum((src_left_anchor - src_right_anchor) ** 2)) + # scale = dst_eye_len / src_eye_len + # dst_center = (anchors_template[0] + anchors_template[1]) / 2 + # offset = dst_center - src_center + # + # M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), rotate_degree, scale) + # M[:, 2] += offset + # + # tmp = cv2.warpAffine(img, M, (dst_size, dst_size)) + # cv2.imshow('tmp', tmp) + # cv2.waitKey() + # + # mmcv = MomocvFaceAlignment() + # + # tmp_input = tmp.transpose((2, 0, 1))[np.newaxis, :, : :].astype(np.float32) + # + # # tmp_input = cv2.imread(r'E:\deepfacelab_data\workspace\input.png') + # # tmp_ori = tmp_input + # # tmp_input = tmp_input.transpose((2, 0, 1))[np.newaxis, :, :, :].astype(np.float32) + # fullyconnected1, poselayer, tracking_probe, occlusion_probe = mmcv.forward(torch.from_numpy(tmp_input)) + # fullyconnected1 = fullyconnected1[0] + # # fullyconnected1 = (np.reshape(fullyconnected1, (2, 87)).transpose((1, 0)) * dst_size).astype(np.int32) + # for i in range(87): + # cv2.circle(tmp, (int(fullyconnected1[i] * 128), int(fullyconnected1[i + 87] * 128)), 1, (255, 0, 0), 1) + # cv2.imshow('tmp_ori', tmp) + # cv2.waitKey() + # + # for ix, pt in enumerate(landmarks): + # cv2.circle(img, (int(pt[0]), int(pt[1])), 1, (255, 0, 0), 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() diff --git a/hair_service_sd/momocv/Face_ResNet.pth b/hair_service_sd/momocv/Face_ResNet.pth new file mode 100644 index 0000000..522a662 --- /dev/null +++ b/hair_service_sd/momocv/Face_ResNet.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a34ffdb26cd6780d41cb3ed313785c09bb7d60634a3f24237303fe7b2746d7e5 +size 12271767 diff --git a/hair_service_sd/momocv/FullFace.pth b/hair_service_sd/momocv/FullFace.pth new file mode 100644 index 0000000..867cc63 --- /dev/null +++ b/hair_service_sd/momocv/FullFace.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4236333ca60a7bffcb13e22e869f738593362e4d407e87e55b579401eeb05be8 +size 45329800 diff --git a/hair_service_sd/momocv/FullFace.py b/hair_service_sd/momocv/FullFace.py new file mode 100644 index 0000000..873298a --- /dev/null +++ b/hair_service_sd/momocv/FullFace.py @@ -0,0 +1,250 @@ +import torch +import torch.onnx +import torch.nn as nn +import sys +import os +from momocv.resnet import resnet18 +import re +import numpy as np +import cv2 +from sklearn import linear_model +from tqdm import tqdm +from momocv.LeftEye import get_left_eye_symbol +from utils import landmark_processor +from mtcnn.detector import MTCNNFaceDetector +import math + +class ResNet18(nn.Module): + def __init__(self, out_channels): + super(ResNet18, self).__init__() + self.op_name = 'ResNet18' + self.res18 = resnet18(pretrained=False) + self.res18.fc = nn.Linear(512, out_channels) + + def forward(self, x): + ret = self.res18(x) + return ret + +def get_full_face_symbol(output_nc=137 * 2): + return ResNet18(output_nc) + +class Model137(nn.Module): + def __init__(self, gpu_id=None): + super(Model137, self).__init__() + + self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu') + + self.face_alignment_net = ResNet18(137 * 2) + self.model_path, _ = os.path.split(os.path.realpath(__file__)) + weights = torch.load(os.path.join(self.model_path, 'FullFace.pth'), + map_location=lambda storage, loc: storage) + self.face_alignment_net.load_state_dict(weights) + self.to(self.device) + self.eval() + + def forward(self, imgs): + pred_key_pts = self.face_alignment_net(imgs) + return pred_key_pts + +class MomocvFaceAlignmentFinalV1(object): + def __init__(self, predict_eye=False, 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.predict_eye = predict_eye + + self.face_alignment_net = get_full_face_symbol() + self.model_path, _ = os.path.split(os.path.realpath(__file__)) + weights = torch.load(os.path.join(self.model_path, 'FullFace.pth'), + map_location=lambda storage, loc: storage) + self.face_alignment_net.load_state_dict(weights) + self.face_alignment_net.to(self.device) + self.face_alignment_net.eval() + + if predict_eye: + self.eye_alignment_net = get_left_eye_symbol() + self.model_path, _ = os.path.split(os.path.realpath(__file__)) + weights = torch.load(os.path.join(self.model_path, 'LeftEye.pth'), + map_location=lambda storage, loc: storage) + self.eye_alignment_net.load_state_dict(weights) + self.eye_alignment_net.to(self.device) + self.eye_alignment_net.eval() + + self.trackingFaceRects = [] + print('conansherry MomocvFaceAlignmentFinalV1') + + 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_mmcv_bigger(landmark, dst_size) + all_mat.append(M) + tmp = cv2.warpAffine(img, M, (dst_size, dst_size)) + 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, 137)).transpose((1, 0)) * dst_size) + orig_pts = landmark_processor.transform_points(orig_pts, all_mat[ix], invert=True) + if self.predict_eye: + eye_landmark = self.detect_eye(img, orig_pts) + orig_pts[87:104] = eye_landmark[0] + orig_pts[104:121] = eye_landmark[1] + landmarks_res.append(orig_pts) + return landmarks_res + + def detect_eye(self, img, landmarks): + dst_size = 96 + src_len = cv2.norm(landmarks[96] - landmarks[88]) + dst_len = 96 * 0.7 + degree = math.atan2(landmarks[88, 1] - landmarks[96, 1], landmarks[88, 0] - landmarks[96, 0]) + src_center = (landmarks[88] + landmarks[96]) / 2 + offset = np.array([0.5, 0.5]) * 96 - src_center + left_M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), math.degrees(degree), dst_len / src_len) + left_M[:, 2] += offset + left_eye_img = cv2.warpAffine(img, left_M, (dst_size, dst_size)) + + dst_size = 96 + src_len = cv2.norm(landmarks[105] - landmarks[113]) + dst_len = 96 * 0.7 + degree = math.atan2(landmarks[113, 1] - landmarks[105, 1], landmarks[113, 0] - landmarks[105, 0]) + src_center = (landmarks[105] + landmarks[113]) / 2 + offset = np.array([0.5, 0.5]) * 96 - src_center + right_M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), math.degrees(degree), dst_len / src_len) + right_M[:, 2] += offset + right_eye_img = cv2.warpAffine(img, right_M, (dst_size, dst_size)) + right_eye_img = cv2.flip(right_eye_img, 1) + + # cv2.imshow('left_eye_img', left_eye_img) + # cv2.imshow('right_eye_img', right_eye_img) + + with torch.no_grad(): + input_numpy = np.zeros((2, 3, dst_size, dst_size), dtype=np.float32) + input_numpy[0, :, :, :] = left_eye_img.transpose((2, 0, 1)).astype(np.float32) / 255 + input_numpy[1, :, :, :] = right_eye_img.transpose((2, 0, 1)).astype(np.float32) / 255 + in_tensor = torch.from_numpy(input_numpy) + in_tensor = in_tensor.to(self.device) + fullyconnected1 = self.eye_alignment_net(in_tensor).detach().cpu().numpy() + landmarks_res = [] + all_mat = [left_M, right_M] + for ix, pts in enumerate(fullyconnected1): + orig_pts = (np.reshape(pts, (2, 17)).transpose((1, 0)) * dst_size) + if ix == 1: + orig_pts[:, 0] = dst_size - orig_pts[:, 0] + orig_pts = landmark_processor.transform_points(orig_pts, all_mat[ix], invert=True) + landmarks_res.append(orig_pts) + return landmarks_res + + # cv2.waitKey() + + 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 = [] + eye_landmarks = [] + for tracking_face_rect in 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', inp) + orig_inp = inp + + inp = inp.transpose((2, 0, 1)).astype(np.float32) + inp = inp[np.newaxis, :, :, :] / 255 + + t0 = cv2.getTickCount() + + in_tensor = torch.from_numpy(inp) + in_tensor = in_tensor.to(self.device) + fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy() + fullyconnected1 = fullyconnected1[0] + orig_pts = (np.reshape(fullyconnected1, (2, 137)).transpose((1, 0)) * 256) + + t2 = cv2.getTickCount() + orig_pts = landmark_processor.transform_points(orig_pts, M, invert=True) + + if self.predict_eye: + eye_landmark = self.detect_eye(image, orig_pts) + orig_pts[87:104] = eye_landmark[0] + orig_pts[104:121] = eye_landmark[1] + eye_landmarks.append(eye_landmark) + # orig_pts = orig_pts.transpose((1, 0)) + fullyconnected1 = orig_pts + + # update tracking infos + tracking_face_rect[1] = False + tracking_face_rect[2] = [fullyconnected1[68], fullyconnected1[74], fullyconnected1[96], fullyconnected1[113]] + tracking_face_rect[3] = rotate_degree + tracking_face_rect[4] = fullyconnected1 + + landmarks.append(fullyconnected1) + return landmarks, eye_landmarks + +if __name__=='__main__': + gpu_id = 0 + net = MomocvFaceAlignmentFinalV1(predict_eye=True, gpu_id=gpu_id) + video_name = r'G:\all_online_videos\2019_04_09_21_57_25_10342e28-ed60-4568-8790-5a4431384031_Trim.mp4' + cap = cv2.VideoCapture(video_name) + face_detector = MTCNNFaceDetector(gpu_id=gpu_id) + + reset = False + bounding_boxes = [] + while True: + _, in_frame = cap.read() + if in_frame is None: + cap = cv2.VideoCapture(video_name) + _, in_frame = cap.read() + if len(bounding_boxes) == 0 or reset: + bounding_boxes, landmarks = face_detector.forward(in_frame, min_face_size=100, thresholds=[0.8, 0.9, 0.95]) + # for box_score in bounding_boxes: + # cv2.rectangle(in_frame, (int(box_score[0]), int(box_score[1])), + # (int(box_score[2]), int(box_score[3])), + # (0, 255, 0), + # 2) + # + # for pt in landmarks: + # for i in range(5): + # cv2.circle(in_frame, (int(pt[i]), int(pt[i + 5])), 1, (255, 0, 0), 2) + for _ in range(3): + landmark137, eye_landmark = net.stable_forward(in_frame, bounding_boxes, reset) + reset = False + + for pts in landmark137: + for ix, pt in enumerate(pts): + # cv2.putText(in_frame, str(ix), (int(pt[0]), int(pt[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (255, 0, 0)) + cv2.circle(in_frame, (pt[0], pt[1]), 1, (0, 255, 0), 1) + # for pts in eye_landmark: + # for eye in pts: + # for pt in eye: + # cv2.circle(in_frame, (pt[0], pt[1]), 1, (0, 0, 255), 1) + + cv2.imshow('in_frame', in_frame) + key = cv2.waitKey(10) + if key == ord('r'): + reset = True + bounding_boxes = [] diff --git a/hair_service_sd/momocv/LeftEye.pth b/hair_service_sd/momocv/LeftEye.pth new file mode 100644 index 0000000..99081f5 --- /dev/null +++ b/hair_service_sd/momocv/LeftEye.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:30238ab182eba7291094d6be816b7ca4062f9abf616a101eb0c91424cab03b59 +size 3216228 diff --git a/hair_service_sd/momocv/LeftEye.py b/hair_service_sd/momocv/LeftEye.py new file mode 100644 index 0000000..9694c48 --- /dev/null +++ b/hair_service_sd/momocv/LeftEye.py @@ -0,0 +1,94 @@ +import sys +import os +import cv2 +import torch +import torch.nn as nn + +def op_name(op_name, m): + m.op_name = op_name + return m + +def conv_bn_relu(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, momentum = 0.1, track_running_stats=True): + return nn.Sequential( + op_name(name, nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, False)), + op_name(name + '/bn', nn.BatchNorm2d(out_channels, momentum = momentum, track_running_stats = track_running_stats)), + op_name(name + '/relu', nn.ReLU(inplace=True)), + ) + + +# Define a Basic resnet block +class BasicResnetBlock(nn.Module): + def __init__(self, name, in_channels, out_channels, kernel_size=3, stride=2, padding=1, direct_plus=False): + super(BasicResnetBlock, self).__init__() + + self.op_name = name + + self.conv_block = nn.Sequential( + conv_bn_relu(name=name + '/conv1', in_channels=in_channels, out_channels=out_channels, + kernel_size=kernel_size, stride=stride, padding=padding), + + conv_bn_relu(name=name + '/conv2', in_channels=out_channels, out_channels=out_channels, + kernel_size=3, stride=1, padding=1) + ) + + if direct_plus and (in_channels == out_channels) and (stride == 1): + self.sc_block = None + else: + self.sc_block = conv_bn_relu(name=name + '/sc_conv', in_channels=in_channels, out_channels=out_channels, + kernel_size=1, stride=stride, padding=0) + + def forward(self, x): + if self.sc_block is None: + out = x + self.conv_block(x) + else: + out = self.conv_block(x) + self.sc_block(x) + return out + +class Flatten(nn.Module): + def __init__(self, axis): + super(Flatten, self).__init__() + self.axis = axis + + def forward(self, x): + assert self.axis == 1 + x = x.reshape(x.shape[0], -1) + # x = x.view(-1, 1152) + return x + +def flatten(name, axis): + return op_name(name, Flatten(axis)) + +def linear_bn_relu(name, in_features, out_features, momentum = 0.1, track_running_stats=True): + return nn.Sequential( + op_name(name + '/FC', nn.Linear(in_features, out_features)), + op_name(name + '/bn', nn.BatchNorm1d(out_features, momentum = momentum, track_running_stats = track_running_stats)), + op_name(name + '/relu', nn.ReLU(inplace=True)), + ) + +class LeftEye(nn.Module): + def __init__(self, name, in_channels, out_channels): + super(LeftEye, self).__init__() + self.op_name = name + + op_list = [] + + op_list += [conv_bn_relu(name + '/first_conv', in_channels, 24, kernel_size = 5, stride = 2, padding = 2) ] + + ch_num = [24, 32, 64, 96, 128] + + op_list += [BasicResnetBlock(name + '/stage%d'%(i + 1), ch_num[i], ch_num[i+1]) for i in range(len(ch_num) - 1) ] + + op_list += [flatten(name + '/flatten', 1), + linear_bn_relu(name + '/FC1', 1152, 256), + op_name(name + '/FC2', nn.Linear(256, out_channels))] + + self.conv_block = nn.Sequential(*op_list) + + def forward(self, x): + return self.conv_block(x) + +def get_left_eye_symbol(symbol_name='BigResNet', input_nc = 3, output_nc = 17 * 2): + return LeftEye(symbol_name, input_nc, output_nc) + +if __name__=='__main__': + pass diff --git a/hair_service_sd/momocv/MMCVFaceRecognition.py b/hair_service_sd/momocv/MMCVFaceRecognition.py new file mode 100644 index 0000000..b94bffe --- /dev/null +++ b/hair_service_sd/momocv/MMCVFaceRecognition.py @@ -0,0 +1,147 @@ +import torch +import torch.nn as nn +import math +import os +import cv2 +import numpy as np +from utils.DATAIMG import DATAIMG +from utils import landmark_processor +import glob +from mathlib.umeyama import umeyama + +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.PReLU(out_channels)), + ) + +class BasicBlock(nn.Module): + def __init__(self, i, j, inplanes): + super(BasicBlock, self).__init__() + + self.conv1 = conv_relu('conv{}_{}'.format(i, 2 * j), inplanes, inplanes, kernel_size=3, stride=1, padding=1) + self.conv2 = conv_relu('conv{}_{}'.format(i, 2 * j + 1), inplanes, inplanes, kernel_size=3, stride=1, padding=1) + + def forward(self, x): + residule = x + out = self.conv1(x) + out = self.conv2(out) + ret = out + residule + return ret + +class ResnetBlock(nn.Module): + def __init__(self, i, inplanes, outplanes, stride=2, n_blocks=1): + super(ResnetBlock, self).__init__() + + self.conv = [conv_relu('conv{}_{}'.format(i, 1), inplanes, outplanes, kernel_size=3, stride=stride, padding=1)] + for m in range(n_blocks): + self.conv.append(BasicBlock(i, m + 1, outplanes)) + self.conv = nn.Sequential(*self.conv) + + def forward(self, x): + ret = self.conv(x) + return ret + +class FaceRecognition(nn.Module): + def __init__(self): + super(FaceRecognition, self).__init__() + op_list = [] + ch_num = [3, 32, 64, 128, 128] + strides = [2, 2, 2, 2] + n_blocks = [1, 2, 4, 1] + op_list = [] + 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 += [flatten('flatten', 1)] + op_list += [op_name('fc5', nn.Linear(4608, 256))] + self.features = nn.Sequential(*op_list) + + model_path, _ = os.path.split(os.path.realpath(__file__)) + weights = torch.load(os.path.join(model_path, 'Face_ResNet.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 MomocvFaceRecognition(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 = FaceRecognition() + # self.model_path, _ = os.path.split(os.path.realpath(__file__)) + # weights = torch.load(os.path.join(self.model_path, 'Face_ResNet.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 = 90 + + 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] + + 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) - 128) / 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 + +if __name__ == '__main__': + all_jpegs = glob.glob(r'E:\deepfacelab_data\expression_dst\7201806132018061208311920180612083119\*.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 = MomocvFaceRecognition() + features = mmcv.forward([img, img], [landmarks, landmarks]) + # 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() diff --git a/hair_service_sd/momocv/MMCVFaceRecognitionServer.pth b/hair_service_sd/momocv/MMCVFaceRecognitionServer.pth new file mode 100644 index 0000000..dfde843 --- /dev/null +++ b/hair_service_sd/momocv/MMCVFaceRecognitionServer.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:752d0396a3a8ce943d7a915bb15063f27112ccb2dd8f4afa2bb9a42824eafbd9 +size 174375530 diff --git a/hair_service_sd/momocv/MMCVFaceRecognitionServer.py b/hair_service_sd/momocv/MMCVFaceRecognitionServer.py new file mode 100644 index 0000000..ba17aa4 --- /dev/null +++ b/hair_service_sd/momocv/MMCVFaceRecognitionServer.py @@ -0,0 +1,177 @@ +import torch +import torch.nn as nn +import math +import os +import cv2 +import numpy as np +from utils.DATAIMG import DATAIMG +from utils import landmark_processor +import glob +from mathlib.umeyama import umeyama + +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.split(os.path.realpath(__file__)) + weights = torch.load(os.path.join(model_path, '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.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] + + 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 + +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() diff --git a/hair_service_sd/momocv/Model_3ddfa.py b/hair_service_sd/momocv/Model_3ddfa.py new file mode 100644 index 0000000..cea54c0 --- /dev/null +++ b/hair_service_sd/momocv/Model_3ddfa.py @@ -0,0 +1,311 @@ +import torch.nn as nn +import torch.utils.model_zoo as model_zoo +import torch +import os +import numpy as np +import cv2 +from utils import landmark_processor, utils_3ddfa, params_3ddfa + +__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, 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(-1, 512) + 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.face_alignment_net = resnet18(pretrained=False, num_classes=76, no_branch=True, no_activate=True) + + model_path, _ = os.path.split(os.path.realpath(__file__)) + weights = torch.load(os.path.join(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) + + def forward(self, imgs): + pred_pose_shape_exp = self.face_alignment_net(imgs) + return pred_pose_shape_exp + + 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 diff --git a/hair_service_sd/momocv/MomocvFaceAlignment1K.py b/hair_service_sd/momocv/MomocvFaceAlignment1K.py new file mode 100644 index 0000000..a0cfb0d --- /dev/null +++ b/hair_service_sd/momocv/MomocvFaceAlignment1K.py @@ -0,0 +1,452 @@ +import torch.nn as nn +import torch.utils.model_zoo as model_zoo +import torch +import numpy as np +import os +from utils import landmark_processor, umeyama +import cv2 + +__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) + 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.eye_alignment_net = get_left_eye_symbol() + # self.model_path, _ = os.path.split(os.path.realpath(__file__)) + # weights = torch.load(os.path.join(self.model_path, 'LeftEye.pth'), map_location=lambda storage, loc: storage) + # self.eye_alignment_net.load_state_dict(weights) + # self.eye_alignment_net.to(self.device) + # self.eye_alignment_net.eval() + + 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_eye(self, img, landmarks): + # dst_size = 96 + # src_len = cv2.norm(landmarks[96] - landmarks[88]) + # dst_len = 96 * 0.7 + # degree = math.atan2(landmarks[88, 1] - landmarks[96, 1], landmarks[88, 0] - landmarks[96, 0]) + # src_center = (landmarks[88] + landmarks[96]) / 2 + # offset = np.array([0.5, 0.5]) * 96 - src_center + # left_M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), math.degrees(degree), dst_len / src_len) + # left_M[:, 2] += offset + # left_eye_img = cv2.warpAffine(img, left_M, (dst_size, dst_size)) + # + # dst_size = 96 + # src_len = cv2.norm(landmarks[105] - landmarks[113]) + # dst_len = 96 * 0.7 + # degree = math.atan2(landmarks[113, 1] - landmarks[105, 1], landmarks[113, 0] - landmarks[105, 0]) + # src_center = (landmarks[105] + landmarks[113]) / 2 + # offset = np.array([0.5, 0.5]) * 96 - src_center + # right_M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), math.degrees(degree), dst_len / src_len) + # right_M[:, 2] += offset + # right_eye_img = cv2.warpAffine(img, right_M, (dst_size, dst_size)) + # right_eye_img = cv2.flip(right_eye_img, 1) + # + # # cv2.imshow('left_eye_img', left_eye_img) + # # cv2.imshow('right_eye_img', right_eye_img) + # + # with torch.no_grad(): + # input_numpy = np.zeros((2, 3, dst_size, dst_size), dtype=np.float32) + # input_numpy[0, :, :, :] = left_eye_img.transpose((2, 0, 1)).astype(np.float32) / 255 + # input_numpy[1, :, :, :] = right_eye_img.transpose((2, 0, 1)).astype(np.float32) / 255 + # in_tensor = torch.from_numpy(input_numpy) + # in_tensor = in_tensor.to(self.device) + # fullyconnected1 = self.eye_alignment_net(in_tensor).detach().cpu().numpy() + # landmarks_res = [] + # all_mat = [left_M, right_M] + # for ix, pts in enumerate(fullyconnected1): + # orig_pts = (np.reshape(pts, (2, 17)).transpose((1, 0)) * dst_size) + # if ix == 1: + # orig_pts[:, 0] = dst_size - orig_pts[:, 0] + # orig_pts = landmark_processor.transform_points(orig_pts, all_mat[ix], invert=True) + # landmarks_res.append(orig_pts) + # return landmarks_res + + 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 + + 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.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # 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]]) + + 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 + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + + tmp = cv2.warpAffine(img, mat, (dst_size, dst_size)) + # 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 \ No newline at end of file diff --git a/hair_service_sd/momocv/SuperBigResNet.py b/hair_service_sd/momocv/SuperBigResNet.py new file mode 100644 index 0000000..4beb401 --- /dev/null +++ b/hair_service_sd/momocv/SuperBigResNet.py @@ -0,0 +1,229 @@ +import torch +import torch.nn as nn +import math +import sys +import os +from utils import landmark_processor +import cv2 +import numpy as np + +output_points = 137 * 2 +input_size = 192 + +def op_name(op_name, m): + m.op_name = op_name + return m + +class Flatten(nn.Module): + def __init__(self, axis): + 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): + return op_name(name, Flatten(axis)) + +def conv_bn_relu(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, momentum = 0.1, track_running_stats=True): + return nn.Sequential( + op_name(name, nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, False)), + op_name(name + '/bn', nn.BatchNorm2d(out_channels, momentum = momentum, track_running_stats = track_running_stats)), + op_name(name + '/relu', nn.ReLU(inplace=True)), + ) + +def linear_bn_relu(name, in_features, out_features, momentum = 0.1, track_running_stats=True): + return nn.Sequential( + op_name(name + '/FC', nn.Linear(in_features, out_features)), + op_name(name + '/bn', nn.BatchNorm1d(out_features, momentum = momentum, track_running_stats = track_running_stats)), + op_name(name + '/relu', nn.ReLU(inplace=True)), + ) + +# Define a Basic resnet block +class BasicResnetBlock(nn.Module): + def __init__(self, name, in_channels, out_channels, kernel_size=3, stride=2, padding=1, direct_plus=False): + super(BasicResnetBlock, self).__init__() + + self.op_name = name + + self.conv_block = nn.Sequential( + conv_bn_relu(name=name + '/conv1', in_channels=in_channels, out_channels=out_channels, + kernel_size=kernel_size, stride=stride, padding=padding), + + conv_bn_relu(name=name + '/conv2', in_channels=out_channels, out_channels=out_channels, + kernel_size=3, stride=1, padding=1) + ) + + if direct_plus and (in_channels == out_channels) and (stride == 1): + self.sc_block = None + else: + self.sc_block = conv_bn_relu(name=name + '/sc_conv', in_channels=in_channels, out_channels=out_channels, + kernel_size=1, stride=stride, padding=0) + + def forward(self, x): + if self.sc_block is None: + out = x + self.conv_block(x) + else: + out = self.conv_block(x) + self.sc_block(x) + return out + +#define SuperBigResNet Frame +class SuperBigResNet(nn.Module): + def __init__(self, name='SuperBigResNet', in_channels=3, out_channels=137 * 2): + super(SuperBigResNet, self).__init__() + self.op_name = name + + op_list = [] + + stride_cnt = 2 + + op_list += [conv_bn_relu(name + '/first_conv', in_channels, 48, kernel_size=5, stride=2, padding=2)] + + ch_num = [48, 64, 96, 128, 192] + for i in range(len(ch_num) - 1): + if ch_num[i] == ch_num[i+1]: + op_list += [BasicResnetBlock(name + '/stage%d'%(i + 1), ch_num[i], ch_num[i+1], kernel_size=3, stride=1, direct_plus=True)] + else: + stride_cnt = stride_cnt * 2 + op_list += [BasicResnetBlock(name + '/stage%d'%(i + 1), ch_num[i], ch_num[i+1], kernel_size=3, stride=2)] + + item_cnt = int((input_size / stride_cnt) * (input_size / stride_cnt) * ch_num[-1]) + op_list += [flatten(name + '/flatten', 1), + linear_bn_relu(name + '/FC1', item_cnt, 512), + op_name(name + '/FC2', nn.Linear(512, out_channels))] + + self.conv_block = nn.Sequential(*op_list) + + def forward(self, x): + res = self.conv_block(x) + return res.cpu().numpy() + +class MomocvFaceAlignmentBigger(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 = SuperBigResNet('SuperBigResNet', 3, 137 * 2) + + self.model_path, _ = os.path.split(os.path.realpath(__file__)) + + weights = torch.load(os.path.join(self.model_path, 'SuperBigResNet_latest.pth'), + map_location=lambda storage, loc: storage) + self.face_alignment_net.load_state_dict(weights) + self.face_alignment_net.to(self.device) + self.face_alignment_net.eval() + + self.trackingFaceRects = [] + + def detect(self, img, landmarks): + dst_size = 192 + 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_mmcv_bigger(landmark, dst_size) + + all_mat.append(M) + + tmp = cv2.warpAffine(img, M, (dst_size, dst_size)) + + input_numpy[ix, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) + + # cv2.imshow('tmp', tmp) + # cv2.waitKey() + + in_tensor = torch.from_numpy(input_numpy) + in_tensor = in_tensor.to(self.device) + fullyconnected1 = self.face_alignment_net(in_tensor) + for ix, pts in enumerate(fullyconnected1): + orig_pts = (np.reshape(pts, (2, 137)).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 stable_forward(self, image, detected_faces, reset=False): + if reset is True: + self.trackingFaceRects = [] + + for face_rect in detected_faces: + if len(self.trackingFaceRects) == 0: + new_tracking_rect = [face_rect, True, [0, 0], 0, None] + self.trackingFaceRects.append(new_tracking_rect) + + with torch.no_grad(): + landmarks = [] + for tracking_face_rect in 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 = 192 * 0.6 / min(d[2] - d[0], d[3] - d[1]) + dst_center = np.array([0.5, 0.5]) * 192 + offset = dst_center - src_center + print('hello') + M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), rotate_degree, scale) + M[:, 2] += offset + else: + # dst_left_anchor = np.array([0.403, 0.52]) * 192 + # dst_right_anchor = np.array([1 - 0.403, 0.52]) * 192 + # # use last anchors + # src_center = (tracking_face_rect[2][0] + tracking_face_rect[2][1]) / 2 + # rotate_radian = math.atan2(tracking_face_rect[2][1][1] - tracking_face_rect[2][0][1], tracking_face_rect[2][1][0] - tracking_face_rect[2][0][0]) + # rotate_degree = rotate_radian / math.pi * 180 + # print('degree', rotate_degree) + # dst_anchor_len = cv2.norm(dst_left_anchor - dst_right_anchor) + # src_anchor_len = cv2.norm(tracking_face_rect[2][2], tracking_face_rect[2][3]) + # scale = 71.0 / src_anchor_len + # + # # bounding_box = cv2.boundingRect(np.array(tracking_face_rect[2])) + # # a = bounding_box[2] * bounding_box[3] + # # print('a', a) + # # scale = float(65*65) / (bounding_box[2] * bounding_box[3]) + # # print('scale', scale) + # + # dst_center = (dst_left_anchor + dst_right_anchor) / 2 + # offset = dst_center - src_center + rotate_degree = 0 + M = landmark_processor.get_transform_mat_mmcv_bigger(tracking_face_rect[4], 192) + inp = cv2.warpAffine(image, M, (192, 192)) + + orig_inp = inp + + inp = inp.transpose((2, 0, 1)).astype(np.float32) + inp = inp[np.newaxis, :, :, :] + + t0 = cv2.getTickCount() + + in_tensor = torch.from_numpy(inp) + in_tensor = in_tensor.to(self.device) + fullyconnected1 = self.face_alignment_net(in_tensor) + fullyconnected1 = fullyconnected1[0] + orig_pts = (np.reshape(fullyconnected1, (2, 137)).transpose((1, 0)) * 192) + + t2 = cv2.getTickCount() + # print('cost', (t2 - t0) / cv2.getTickFrequency() * 1000) + + # for ix, pt in enumerate(orig_pts): + # cv2.putText(orig_inp, str(ix), (int(pt[0]), int(pt[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (255, 0, 0)) + # cv2.circle(orig_inp, (int(pt[0]), int(pt[1])), 1, (0, 255, 0), 1) + # cv2.imshow('inp_stable', orig_inp) + # cv2.waitKey() + + 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] = [fullyconnected1[68], fullyconnected1[74], fullyconnected1[96], fullyconnected1[113]] + tracking_face_rect[3] = rotate_degree + tracking_face_rect[4] = fullyconnected1 + + landmarks.append(fullyconnected1) + return landmarks + +if __name__=='__main__': + pass diff --git a/hair_service_sd/momocv/SuperBigResNetV2.py b/hair_service_sd/momocv/SuperBigResNetV2.py new file mode 100644 index 0000000..1c4c2a3 --- /dev/null +++ b/hair_service_sd/momocv/SuperBigResNetV2.py @@ -0,0 +1,102 @@ +import torch +import torch.nn as nn +import math +import torch.optim as optim +import sys +import os + +def op_name(op_name, m): + m.op_name = op_name + return m + +class Flatten(nn.Module): + def __init__(self, axis): + super(Flatten, self).__init__() + + def forward(self, x): + x = x.view(-1, 2560) + return x + +def flatten(name, axis): + return op_name(name, Flatten(axis)) + +def conv_bn_relu(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, momentum = 0.1, track_running_stats=True): + return nn.Sequential( + op_name(name, nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, False)), + op_name(name + '/bn', nn.BatchNorm2d(out_channels, momentum = momentum, track_running_stats = track_running_stats)), + op_name(name + '/relu', nn.ReLU(inplace=True)), + ) + +def linear_bn_relu(name, in_features, out_features, momentum = 0.1, track_running_stats=True): + return nn.Sequential( + op_name(name + '/FC', nn.Linear(in_features, out_features)), + op_name(name + '/bn', nn.BatchNorm1d(out_features, momentum = momentum, track_running_stats = track_running_stats)), + op_name(name + '/relu', nn.ReLU(inplace=True)), + ) + +# Define a Basic resnet block +class BasicResnetBlock(nn.Module): + def __init__(self, name, in_channels, out_channels, kernel_size=3, stride=2, padding=1, direct_plus=False): + super(BasicResnetBlock, self).__init__() + + self.op_name = name + + self.conv_block = nn.Sequential( + conv_bn_relu(name=name + '/conv1', in_channels=in_channels, out_channels=out_channels, + kernel_size=kernel_size, stride=stride, padding=padding), + + conv_bn_relu(name=name + '/conv2', in_channels=out_channels, out_channels=out_channels, + kernel_size=3, stride=1, padding=1) + ) + + if direct_plus and (in_channels == out_channels) and (stride == 1): + self.sc_block = None + else: + self.sc_block = conv_bn_relu(name=name + '/sc_conv', in_channels=in_channels, out_channels=out_channels, + kernel_size=1, stride=stride, padding=0) + + def forward(self, x): + if self.sc_block is None: + out = x + self.conv_block(x) + else: + out = self.conv_block(x) + self.sc_block(x) + return out + +class SuperBigResNetV2(nn.Module): + def __init__(self, name, in_channels, out_channels): + super(SuperBigResNetV2, self).__init__() + self.op_name = name + + input_size = 256 + + op_list = [] + + stride_cnt = 4 + + op_list += [conv_bn_relu(name + '/first_conv', in_channels, 48, kernel_size=5, stride=4, padding=2)] + + ch_num = [48, 64, 96, 128, 160] + for i in range(len(ch_num) - 1): + if ch_num[i] == ch_num[i+1]: + op_list += [BasicResnetBlock(name + '/stage%d'%(i + 1), ch_num[i], ch_num[i+1], kernel_size=3, stride=1, direct_plus=True)] + else: + stride_cnt = stride_cnt * 2 + op_list += [BasicResnetBlock(name + '/stage%d'%(i + 1), ch_num[i], ch_num[i+1], kernel_size=3, stride=2)] + + item_cnt = int((input_size / stride_cnt) * (input_size / stride_cnt) * ch_num[-1]) + op_list += [flatten(name + '/flatten', 1), + linear_bn_relu(name + '/FC1', item_cnt, 512), + op_name(name + '/FC2', nn.Linear(512, out_channels))] + + self.conv_block = nn.Sequential(*op_list) + + self.model_path, _ = os.path.split(os.path.realpath(__file__)) + + self.weights = torch.load(os.path.join(self.model_path, 'SuperBigResNetV2_latest.pth'), + map_location=lambda storage, loc: storage) + + def forward(self, x): + return self.conv_block(x) + +if __name__=='__main__': + pass diff --git a/hair_service_sd/momocv/SuperBigResNetV2_latest.pth b/hair_service_sd/momocv/SuperBigResNetV2_latest.pth new file mode 100644 index 0000000..a3085db --- /dev/null +++ b/hair_service_sd/momocv/SuperBigResNetV2_latest.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:82a5c042c29002a009f83345875a73d39209a435459a1fa46a95f30392d07e92 +size 9536531 diff --git a/hair_service_sd/momocv/SuperBigResNet_latest.pth b/hair_service_sd/momocv/SuperBigResNet_latest.pth new file mode 100644 index 0000000..00b61a2 --- /dev/null +++ b/hair_service_sd/momocv/SuperBigResNet_latest.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7206aaa727f403c6f8e660e9bf9d9967e6fdb3dced1f77232320e0ba0b5b2515 +size 19020307 diff --git a/hair_service_sd/momocv/face_3ddfa.pth b/hair_service_sd/momocv/face_3ddfa.pth new file mode 100644 index 0000000..9d2c828 --- /dev/null +++ b/hair_service_sd/momocv/face_3ddfa.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c0d5cc765269c4237fb29d2a825da9d0a39e59813fe2f681003cb1ac4b081f17 +size 44926130 diff --git a/hair_service_sd/momocv/face_alignment_1k.pth b/hair_service_sd/momocv/face_alignment_1k.pth new file mode 100644 index 0000000..dacb3ab --- /dev/null +++ b/hair_service_sd/momocv/face_alignment_1k.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:89042256066e0e4ff692f4cfa8a4727ee55c6ccf8dccb8d1a1907484da283582 +size 48873911 diff --git a/hair_service_sd/momocv/resnet.py b/hair_service_sd/momocv/resnet.py new file mode 100644 index 0000000..28fb73e --- /dev/null +++ b/hair_service_sd/momocv/resnet.py @@ -0,0 +1,223 @@ +import torch.nn as nn +import torch.utils.model_zoo as model_zoo + + +__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, 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)) + 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.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) + + x = self.avgpool(x) + x = x.view(-1, 512) + + return self.fc(x) + + +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'])) + 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 diff --git a/hair_service_sd/mtcnn/__init__.py b/hair_service_sd/mtcnn/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/hair_service_sd/mtcnn/__init__.py @@ -0,0 +1 @@ + diff --git a/hair_service_sd/mtcnn/box_utils.py b/hair_service_sd/mtcnn/box_utils.py new file mode 100644 index 0000000..d7a076f --- /dev/null +++ b/hair_service_sd/mtcnn/box_utils.py @@ -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 diff --git a/hair_service_sd/mtcnn/detector.py b/hair_service_sd/mtcnn/detector.py new file mode 100644 index 0000000..407c0b8 --- /dev/null +++ b/hair_service_sd/mtcnn/detector.py @@ -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 \ No newline at end of file diff --git a/hair_service_sd/mtcnn/detector_ly.py b/hair_service_sd/mtcnn/detector_ly.py new file mode 100644 index 0000000..407c0b8 --- /dev/null +++ b/hair_service_sd/mtcnn/detector_ly.py @@ -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 \ No newline at end of file diff --git a/hair_service_sd/mtcnn/model.py b/hair_service_sd/mtcnn/model.py new file mode 100644 index 0000000..a8fbc08 --- /dev/null +++ b/hair_service_sd/mtcnn/model.py @@ -0,0 +1,110 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from collections import OrderedDict +import numpy as np +import os +# from hairstyle_model import modelRoot +modelRoot = "./weights" +# modelRoot = "/home/yangchaojie/Desktop/hairstyle/hairstyle_infer/weights" +class Flatten(nn.Module): + def __init__(self): + super(Flatten, self).__init__() + def forward(self, x): + x = x.transpose(3, 2).contiguous() + return x.view(x.size(0), -1) + +class PNet(nn.Module): + def __init__(self): + super(PNet, self).__init__() + self.model_path = modelRoot + + self.features = nn.Sequential(OrderedDict([ + ('conv1', nn.Conv2d(3, 10, 3, 1)), + ('prelu1', nn.PReLU(10)), + ('pool1', nn.MaxPool2d(2, 2, ceil_mode=True)), + ('conv2', nn.Conv2d(10, 16, 3, 1)), + ('prelu2', nn.PReLU(16)), + ('conv3', nn.Conv2d(16, 32, 3, 1)), + ('prelu3', nn.PReLU(32)) + ])) + self.conv4_1 = nn.Conv2d(32, 2, 1, 1) + self.conv4_2 = nn.Conv2d(32, 4, 1, 1) + weights = np.load(os.path.join(self.model_path, 'pnet.npy'), allow_pickle=True)[()] + for n, p in self.named_parameters(): + p.data = torch.FloatTensor(weights[n]) + + def forward(self, x): + x = self.features(x) + a = self.conv4_1(x) + b = self.conv4_2(x) + a = F.softmax(a, dim=1) + return b, a + +class RNet(nn.Module): + def __init__(self): + super(RNet, self).__init__() + self.model_path = modelRoot + + self.features = nn.Sequential(OrderedDict([ + ('conv1', nn.Conv2d(3, 28, 3, 1)), + ('prelu1', nn.PReLU(28)), + ('pool1', nn.MaxPool2d(3, 2, ceil_mode=True)), + ('conv2', nn.Conv2d(28, 48, 3, 1)), + ('prelu2', nn.PReLU(48)), + ('pool2', nn.MaxPool2d(3, 2, ceil_mode=True)), + ('conv3', nn.Conv2d(48, 64, 2, 1)), + ('prelu3', nn.PReLU(64)), + ('flatten', Flatten()), + ('conv4', nn.Linear(576, 128)), + ('prelu4', nn.PReLU(128)) + ])) + self.conv5_1 = nn.Linear(128, 2) + self.conv5_2 = nn.Linear(128, 4) + weights = np.load(os.path.join(self.model_path, 'rnet.npy'), allow_pickle=True)[()] + for n, p in self.named_parameters(): + p.data = torch.FloatTensor(weights[n]) + + def forward(self, x): + x = self.features(x) + a = self.conv5_1(x) + b = self.conv5_2(x) + a = F.softmax(a, dim=1) + return b, a + +class ONet(nn.Module): + def __init__(self): + super(ONet, self).__init__() + self.model_path = modelRoot + + self.features = nn.Sequential(OrderedDict([ + ('conv1', nn.Conv2d(3, 32, 3, 1)), + ('prelu1', nn.PReLU(32)), + ('pool1', nn.MaxPool2d(3, 2, ceil_mode=True)), + ('conv2', nn.Conv2d(32, 64, 3, 1)), + ('prelu2', nn.PReLU(64)), + ('pool2', nn.MaxPool2d(3, 2, ceil_mode=True)), + ('conv3', nn.Conv2d(64, 64, 3, 1)), + ('prelu3', nn.PReLU(64)), + ('pool3', nn.MaxPool2d(2, 2, ceil_mode=True)), + ('conv4', nn.Conv2d(64, 128, 2, 1)), + ('prelu4', nn.PReLU(128)), + ('flatten', Flatten()), + ('conv5', nn.Linear(1152, 256)), + ('drop5', nn.Dropout(0.25)), + ('prelu5', nn.PReLU(256)), + ])) + self.conv6_1 = nn.Linear(256, 2) + self.conv6_2 = nn.Linear(256, 4) + self.conv6_3 = nn.Linear(256, 10) + weights = np.load(os.path.join(self.model_path, 'onet.npy'), allow_pickle=True)[()] + for n, p in self.named_parameters(): + p.data = torch.FloatTensor(weights[n]) + + def forward(self, x): + x = self.features(x) + a = self.conv6_1(x) + b = self.conv6_2(x) + c = self.conv6_3(x) + a = F.softmax(a, dim=1) + return c, b, a diff --git a/hair_service_sd/prepare_ref_hairstyle_data.py b/hair_service_sd/prepare_ref_hairstyle_data.py new file mode 100644 index 0000000..49338e0 --- /dev/null +++ b/hair_service_sd/prepare_ref_hairstyle_data.py @@ -0,0 +1,350 @@ +import cv2 +import numpy as np +import torch +import json +import time +import re +import os +from process_modules import Get_Landmark, Process_Data +from utils import landmark_processor + +class GenderClassifyProcessor(object): + def __init__(self, gpu_id=0): + model_path = "./weights/gender_models" + self.output_img_size = 128 + if not os.path.exists(model_path): + print("GenderClassifyProcessor don't have model!") + + if gpu_id == 'cpu': + self.device = torch.device(gpu_id) + else: + self.device = torch.device('cuda:{0}'.format(gpu_id)) + + self.gender_model = cv2.dnn.readNetFromCaffe(os.path.join(model_path, "gender.prototxt"), os.path.join(model_path, "gender.caffemodel")) + + def forward(self, img, landmark_137): + + image_to_face_mat = landmark_processor.get_transform_mat_sex(landmark_137, self.output_img_size) + gender_img = cv2.warpAffine(img, image_to_face_mat, (self.output_img_size, self.output_img_size), cv2.INTER_LANCZOS4) + inpBlob = cv2.dnn.blobFromImage(gender_img, 1.0, (self.output_img_size, self.output_img_size), (0, 0, 0), swapRB=False, + crop=False) + self.gender_model.setInput(inpBlob) + output = self.gender_model.forward() + + is_female = True + if output[0][0] > output[0][1]: + is_female = False + return is_female + + +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, long_flag=False): + + 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_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" + print("gender: ", gender, " ratio: ", ratio) + + # 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, long=long_flag) + + # 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 + +def prepare_single(img_path=None, dst_dir=None, long_flag=False): + # img_path = "/home/yangchaojie/Desktop/hairstyle/test_data/1123/pics/test3.jpg" + # img_path = "./test_data/ref_imgs/female1.jpg" + # dst_dir = "./data/ref_hairstyle/tmp" + + # img_path = "/home/szlc/Downloads/111/111/ddbc6cb6-766c-4f08-b1fd-22c6bab86432.jpg" + # dst_dir = "/home/szlc/Downloads/111/res" + if not os.path.exists(dst_dir): + os.makedirs(dst_dir) + prepare_data_process = Prepare_Ref_HairStyle_Data(True, device_id=0) + ref_rgb_8uc3_orisize = cv2.imread(img_path) + + # debug = True + debug = False + + version = "20221206" + start_time = time.time() + + ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768, gender, ratio = \ + prepare_data_process.get_prepare_ref_768_data(ref_rgb_8uc3_orisize, long_flag) + + + if not debug: + cv2.imwrite(os.path.join(dst_dir, "ref_rgb_8uc3_768.png"), ref_rgb_8uc3_768) + cv2.imwrite(os.path.join(dst_dir, "ref_matting_fg_8uc3_768.png"), ref_matting_fg_8uc3_768) + cv2.imwrite(os.path.join(dst_dir, "ref_matting_8uc3_768.png"), ref_matting_8uc3_768) + cv2.imwrite(os.path.join(dst_dir, "ref_baldseg_8uc3_768.png"), ref_baldseg_8uc3_768) + np.savetxt(os.path.join(dst_dir, "ref_landmark_f1k2_768.txt"), ref_landmark_f1k2_768) + + input_another_pose_hair_image = prepare_data_process.Generator_reftensor(ref_rgb_8uc3_768, ref_matting_8uc3_768, + ref_baldseg_8uc3_768, + ref_landmark_f1k2_768) + print("cost time: ", time.time() - start_time) + if debug: + print("gender: ", gender, " ratio: ", ratio) + show_concat = np.concatenate((ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768), axis=1) + # cv2.imshow("show_concat", cv2.resize(show_concat, (0, 0), fx=0.6, fy=0.6)) + # cv2.imshow("input_another_pose_hair_image", input_another_pose_hair_image) + # cv2.waitKey() + if not debug: + np.save(os.path.join(dst_dir, "input_another_pose_hair_image.npy"), input_another_pose_hair_image) + + config_dict = {} + + config_dict['gender'] = gender + config_dict['version'] = version + config_dict['ratio'] = str(ratio) + + with open(os.path.join(dst_dir, "config.json"), "w") as f: + json.dump(config_dict, f) + print("写入文件完成...", dst_dir) + + +def prepare_single_color(img_path=None, dst_dir=None, long_flag=False): + + if not os.path.exists(dst_dir): + os.makedirs(dst_dir) + prepare_data_process = Prepare_Ref_HairStyle_Data(True, device_id=0) + ref_rgb_8uc3_orisize = cv2.imread(img_path) + + # debug = True + debug = False + + version = "20221206" + start_time = time.time() + + ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768, gender, ratio = \ + prepare_data_process.get_prepare_ref_768_data(ref_rgb_8uc3_orisize, long_flag) + + + if not debug: + cv2.imwrite(os.path.join(dst_dir, "ref_rgb_8uc3_color_768.png"), ref_rgb_8uc3_768) + cv2.imwrite(os.path.join(dst_dir, "ref_matting_fg_8uc3_color_768.png"), ref_matting_fg_8uc3_768) + cv2.imwrite(os.path.join(dst_dir, "ref_matting_8uc3_color_768.png"), ref_matting_8uc3_768) + cv2.imwrite(os.path.join(dst_dir, "ref_baldseg_8uc3_color_768.png"), ref_baldseg_8uc3_768) + np.savetxt(os.path.join(dst_dir, "ref_landmark_f1k2_color_768.txt"), ref_landmark_f1k2_768) + + input_another_pose_hair_image = prepare_data_process.Generator_reftensor(ref_rgb_8uc3_768, ref_matting_8uc3_768, + ref_baldseg_8uc3_768, + ref_landmark_f1k2_768) + print("cost time: ", time.time() - start_time) + if debug: + print("gender: ", gender, " ratio: ", ratio) + show_concat = np.concatenate((ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768), axis=1) + # cv2.imshow("show_concat", cv2.resize(show_concat, (0, 0), fx=0.6, fy=0.6)) + # cv2.imshow("input_another_pose_hair_image", input_another_pose_hair_image) + # cv2.waitKey() + if not debug: + np.save(os.path.join(dst_dir, "input_another_pose_hair_image.npy"), input_another_pose_hair_image) + + config_dict = {} + + config_dict['gender'] = gender + config_dict['version'] = version + config_dict['ratio'] = str(ratio) + + with open(os.path.join(dst_dir, "config.json"), "w") as f: + json.dump(config_dict, f) + print("写入文件完成...", dst_dir) + +def prepare_multi(): + + img_dir = "/home/yangchaojie/Desktop/hairstyle/test_data/testdata_tj_0309/raw_pics" + save_dir = "/home/yangchaojie/Desktop/hairstyle/test_data/testdata_tj_0309/ref_haircolor" + # pattern = re.compile(r'^[^\.].*\.jpg$') + pattern = re.compile(r'^[^\.].*\.(jpg|JPG)$') + for dirpath, dirnames, filenames in os.walk(img_dir, followlinks=True): + for filename in filenames: + if not pattern.match(filename): continue + img_path = os.path.join(dirpath, filename) + dst_dir = os.path.join(save_dir, filename[:-4]) + + if not os.path.exists(dst_dir): + os.makedirs(dst_dir) + prepare_data_process = Prepare_Ref_HairStyle_Data(True, device_id=0) + ref_rgb_8uc3_orisize = cv2.imread(img_path) + + # debug = True + debug = False + + version = "20230309" + start_time = time.time() + + ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768, gender, ratio = \ + prepare_data_process.get_prepare_ref_768_data(ref_rgb_8uc3_orisize) + + + if not debug: + cv2.imwrite(os.path.join(dst_dir, "ref_rgb_8uc3_768.png"), ref_rgb_8uc3_768) + cv2.imwrite(os.path.join(dst_dir, "ref_matting_fg_8uc3_768.png"), ref_matting_fg_8uc3_768) + cv2.imwrite(os.path.join(dst_dir, "ref_matting_8uc3_768.png"), ref_matting_8uc3_768) + cv2.imwrite(os.path.join(dst_dir, "ref_baldseg_8uc3_768.png"), ref_baldseg_8uc3_768) + np.savetxt(os.path.join(dst_dir, "ref_landmark_f1k2_768.txt"), ref_landmark_f1k2_768) + + input_another_pose_hair_image = prepare_data_process.Generator_reftensor(ref_rgb_8uc3_768, ref_matting_8uc3_768, + ref_baldseg_8uc3_768, + ref_landmark_f1k2_768) + print("cost time: ", time.time() - start_time) + if debug: + print("gender: ", gender, " ratio: ", ratio) + show_concat = np.concatenate((ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768), axis=1) + cv2.imshow("show_concat", cv2.resize(show_concat, (0, 0), fx=0.6, fy=0.6)) + cv2.imshow("input_another_pose_hair_image", input_another_pose_hair_image) + cv2.waitKey() + if not debug: + np.save(os.path.join(dst_dir, "input_another_pose_hair_image.npy"), input_another_pose_hair_image) + + config_dict = {} + + config_dict['gender'] = gender + config_dict['version'] = version + config_dict['ratio'] = str(ratio) + + with open(os.path.join(dst_dir, "config.json"), "w") as f: + json.dump(config_dict, f) + print("写入文件完成...", dst_dir) + +if __name__ == '__main__': + prepare_single() + # prepare_multi() \ No newline at end of file diff --git a/hair_service_sd/process_modules.py b/hair_service_sd/process_modules.py new file mode 100644 index 0000000..1b79668 --- /dev/null +++ b/hair_service_sd/process_modules.py @@ -0,0 +1,3780 @@ +import os +import time +import torchvision +import pickle +import sys +import torch +from torch.nn import functional as F +from torch import nn +import numpy as np +import cv2 +import math +# import onnxruntime +from utils import landmark_processor +# import torchvision +from models.MomocvFaceAlignment1K import MomocvFaceAlignment1K +from mtcnn.detector import MTCNNFaceDetector +from utils import util +from matting.networks import generators +from seg.hairseg_single_model import Evaluator +from models.Generator_BaldSeg import Generator_BaldSeg_5c +from bodyseg.msc_distilling import DeepLab +from utils import model_io +modelRoot = "./weights" + + +class Pose_KeypointsProcessorV2(object): + def __init__(self, gpu_id=0): + import torch + sys.path.append("..") + from keypoints.lib.config import cfg, update_config + from keypoints.lib.models.pose_hrnet import get_pose_net + sys.path.remove("..") + + class Namespace: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + + args = Namespace(cfg=os.path.join(modelRoot, 'keypoints/experiments/coco/hrnet/coco25_384x288_adam_lr1e-3.yaml'), + opts=['TEST.MODEL_FILE', os.path.join(modelRoot, 'kpnts_detect_lyq_20200720.pth'), 'TEST.USE_GT_BBOX', 'False'], + dataDir='', + logDir='', + modelDir='', + prevModelDir='') + + update_config(cfg, args) + # print(cfg) + + self.model = get_pose_net(cfg, is_train=False) + self.model.eval() + + if cfg.TEST.MODEL_FILE: + if gpu_id == -1: + self.model.load_state_dict(torch.load(cfg.TEST.MODEL_FILE, map_location='cpu'), strict=False) + else: + # self.model.load_state_dict(cp['best_state_dict']) + self.model.load_state_dict(torch.load(cfg.TEST.MODEL_FILE, map_location=lambda storage, loc: storage), + strict=False) + print('load keypoints model weights') + + if gpu_id != -1: + self.model.cuda(gpu_id) + self.gpu_id = gpu_id + + self.image_size = [288, 384] + self.cfg = cfg + self.compare_table = [[0, 0], [1, 16], [2, 15], [3, 18], [4, 17], \ + [5, 5], [6, 2], [7, 6], [8, 3], [9, 7], \ + [10, 4], [11, 12], [12, 9], [13, 13], [14, 10], \ + [15, 14], [16, 11]] + + self.ref_lds = [[142.50674, 55.46675], [142.58651, 93.27755], [115.88326, 95.49268], + [107.85087, 135.16175], [103.58789, 169.31675], [169.52114, 89.90296], + [186.71601, 131.36675], [197.37347, 167.41925], [146.29747, 175.92391], + [125.37664, 174.99589], [141.95472, 241.42175], [150.48068, 311.62925], + [163.45916, 176.62695], [161.13813, 249.01175], [156.87515, 296.80147], + [135.56025, 49.77425], [150.48068, 47.87675], [124.90280, 55.46675], + [156.87515, 51.67175], [150.12675, 322.52566], [155.76106, 320.27593], + [158.85028, 295.46357], [154.74366, 336.29675], [146.21770, 336.29675], + [147.24197, 316.33891]] + + self.ref_lds = np.array(self.ref_lds) + self.compare_table = np.array(self.compare_table) + self.ref_lds = self.ref_lds[self.compare_table[:, 1]] + + # @staticmethod + def get_refimage(self, srcImg, landmark, refLandmark): + + def setInvM(M): + # return inv M + D = M[0, 0] * M[1, 1] - M[0, 1] * M[1, 0] + D = 1 / D if D != 0 else 0.0 + invM = np.zeros((2, 3)) + + invM[0, 0] = M[1, 1] * D + invM[0, 1] = M[0, 1] * D * (-1) + invM[1, 0] = M[1, 0] * D * (-1) + invM[1, 1] = M[0, 0] * D + + invM[0, 2] = -invM[0, 0] * M[0, 2] - invM[0, 1] * M[1, 2] + invM[1, 2] = -invM[1, 0] * M[0, 2] - invM[1, 1] * M[1, 2] + return invM + + def setTransMatrix(srcPts, dstPts): + ms_x = np.mean(srcPts[:, 0]) + ms_y = np.mean(srcPts[:, 1]) + srcPts -= np.array([ms_x, ms_y]) + + x = np.concatenate((srcPts[:, 0], srcPts[:, 1])) + y = np.concatenate((dstPts[:, 0], dstPts[:, 1])) + + a = y.dot(x) / x.dot(x) + offset = int(len(x) / 2) + b = 0 + for i in range(offset): + b += x[i] * y[offset + i] - x[offset + i] * y[i] + b /= x.dot(x) + + md_x = np.mean(dstPts[:, 0]) + md_y = np.mean(dstPts[:, 1]) + + M = np.zeros((2, 3)) + M[0, 0] = a + M[0, 1] = -b + M[1, 0] = b + M[1, 1] = a + M[0, 2] = md_x - M[0, 0] * ms_x - M[0, 1] * ms_y + M[1, 2] = md_y - M[1, 0] * ms_x - M[1, 1] * ms_y + + invM = setInvM(M) + return invM, M + + def transImage(srcImg, M, invM, back=False): + T = M if back else invM + dstImg = cv2.warpAffine(srcImg, T, (288, 384)) + return dstImg + + def transShape(srcShape, M, invM, forward=False): + T = invM if forward else M + x = T[0, 0] * srcShape[:, 0] + T[0, 1] * srcShape[:, 1] + T[0, 2] + y = T[1, 0] * srcShape[:, 1] + T[1, 1] * srcShape[:, 1] + T[1, 2] + # print(x, y) + tmp = zip(x, y) + return tmp + + landmark_tmp = landmark.copy() + M, invM = setTransMatrix(landmark_tmp, refLandmark) + dstImg = transImage(srcImg, M, invM) + return dstImg, invM + + # @staticmethod + def transform_points(self, points, mat, invert=False): + if invert: + mat = cv2.invertAffineTransform(mat) + points = np.expand_dims(points, axis=1) + points = cv2.transform(points, mat, points.shape) + points = np.squeeze(points) + return points + + # @staticmethod + def affine_trans(self, img, joints): + img1 = img.copy() + save_pts = joints.copy() + mask = [] + obj_lds = np.array(joints[:, :2]).copy() + tag = 0 + for i in joints[:, 2]: + if i > 0.6: + mask.append(True) + tag += 1 + else: + mask.append(False) + + if tag < 8: + mask = [False, True, True, False, False, + True, True, False, False, False, + False, False, False, False, False, + False, False] + + ref_lds = self.ref_lds.copy() + ref_lds = np.array(ref_lds) + mask = np.array(mask) + + ref_lds = ref_lds[mask] + obj_lds = obj_lds[mask] + + img, M = self.get_refimage(img, obj_lds, ref_lds) + return img, M + + def get_max_preds(self, batch_heatmaps): + ''' + get predictions from score maps + heatmaps: numpy.ndarray([batch_size, num_joints, height, width]) + ''' + assert isinstance(batch_heatmaps, np.ndarray), \ + 'batch_heatmaps should be numpy.ndarray' + assert batch_heatmaps.ndim == 4, 'batch_images should be 4-ndim' + + batch_size = batch_heatmaps.shape[0] + num_joints = batch_heatmaps.shape[1] + width = batch_heatmaps.shape[3] + heatmaps_reshaped = batch_heatmaps.reshape((batch_size, num_joints, -1)) + idx = np.argmax(heatmaps_reshaped, 2) + maxvals = np.amax(heatmaps_reshaped, 2) + + maxvals = maxvals.reshape((batch_size, num_joints, 1)) + idx = idx.reshape((batch_size, num_joints, 1)) + + preds = np.tile(idx, (1, 1, 2)).astype(np.float32) + + preds[:, :, 0] = (preds[:, :, 0]) % width + preds[:, :, 1] = np.floor((preds[:, :, 1]) / width) + + pred_mask = np.tile(np.greater(maxvals, 0.0), (1, 1, 2)) + pred_mask = pred_mask.astype(np.float32) + + preds *= pred_mask + return preds, maxvals + + def preprocess(self, img, pre_pts): + input, M = self.affine_trans(img, pre_pts) + # print(input1.shape) + + input = cv2.cvtColor(input, cv2.COLOR_BGR2RGB) + input = np.divide(np.subtract(input.astype(np.float32) / 255, [0.406, 0.456, 0.485]), + [0.225, 0.224, 0.229]) + + input = np.expand_dims(input.astype(np.float32).transpose((2, 0, 1)), axis=0) + input = torch.from_numpy(input) + return input, M + + def postprocess(self, output, M): + output = output.detach().cpu().numpy() + lds, maxvals = self.get_max_preds(output) + lds *= 4 + # print(lds, lds.shape) + lds = lds[0] + lds = self.transform_points(lds[:, :2], M, True) + res = np.concatenate((lds, maxvals[0]), axis=1) + + # print(res) + return res + + def forward(self, img, pre_lds): + ss = img.copy() + + img, M = self.preprocess(img, pre_lds) + + if self.gpu_id != -1: + img = img.cuda(self.gpu_id) + output = self.model(img) + + pts = self.postprocess(output, M) + return pts + +def pt_conv_25_to_17(pt25): + index_map = [ + 0, # 0 + 16, # 1 + 15, # 2 + 18, # 3 + 17, # 4 + 5, # 5 + 2, # 6 + 6, # 7 + 3, # 8 + 7, # 9 + 4, # 10 + 12, # 11 + 9, # 12 + 13, # 13 + 10, # 14 + 14, # 15 + 11, # 16 + ] + pt_17 = np.zeros((17, 3), dtype=np.float32) + for idx17, idx25 in enumerate(index_map): + pt_17[idx17, :] = pt25[idx25] + return pt_17 + +class Human_Keypoints: + def __init__(self, gpu=True, device_id=0): + if gpu: + self.gpu_id = device_id + else: + self.gpu_id = -1 + with torch.no_grad(): + # self.pose_keypoints_processor = Pose_KeypointsProcessor(gpu_id=self.gpu_id) + self.pose_keypoints_processor2 = Pose_KeypointsProcessorV2(gpu_id=self.gpu_id) + + def inference(self, input_img, human_rect, kpnts17, hand_seg=None): + # pose_keypoints = self.pose_keypoints_processor.forward(input_img, human_rect) + pose_keypoints = self.pose_keypoints_processor2.forward(input_img, kpnts17) # new version lyq + + # # check output 25 kpnts + # main_image_show = input_img.copy() + # for i in range(pose_keypoints.shape[0]): + # cv2.circle(main_image_show, (int(pose_keypoints[i][0]), int(pose_keypoints[i][1])), 5, (0, 255, 0), 5) # old green + # cv2.circle(main_image_show, (int(pose_keypoints2[i][0]), int(pose_keypoints2[i][1])), 5, (220, 0, 220), 5) # new red + # cv2.imshow("kpnts compare", main_image_show) + # cv2.waitKey() + + if len(pose_keypoints) == 0: + pose_keypoints = np.zeros((25, 3), dtype=np.float32) + return pose_keypoints + + def convert_pose_points(self, pose_keypoints, confidence): + if (pose_keypoints[11, 0] - pose_keypoints[10, 0]) ** 2 + (pose_keypoints[11, 1] - pose_keypoints[10, 1]) ** 2 < \ + 0.09 * ((pose_keypoints[10, 0] - pose_keypoints[9, 0]) ** 2 + (pose_keypoints[10, 1] - pose_keypoints[9, 1]) ** 2): + pose_keypoints[11, 2] = 0 + if (pose_keypoints[14, 0] - pose_keypoints[13, 0]) ** 2 + (pose_keypoints[14, 1] - pose_keypoints[13, 1]) ** 2 < \ + 0.09 * ((pose_keypoints[13, 0] - pose_keypoints[12, 0]) ** 2 + (pose_keypoints[13, 1] - pose_keypoints[12, 1]) ** 2): + pose_keypoints[14, 2] = 0 + + if pose_keypoints[10, 2] <= confidence: + pose_keypoints[11, 2] = 0 + if pose_keypoints[13, 2] <= confidence: + pose_keypoints[14, 2] = 0 + + ### heel + if (pose_keypoints[24, 0] - pose_keypoints[11, 0]) ** 2 + (pose_keypoints[24, 1] - pose_keypoints[11, 1]) ** 2 > \ + 0.2 * ((pose_keypoints[11, 0] - pose_keypoints[10, 0]) ** 2 + (pose_keypoints[11, 1] - pose_keypoints[10, 1]) ** 2): + pose_keypoints[24, 2] = 0 + if (pose_keypoints[21, 0] - pose_keypoints[14, 0]) ** 2 + (pose_keypoints[21, 1] - pose_keypoints[14, 1]) ** 2 > \ + 0.2 * ((pose_keypoints[14, 0] - pose_keypoints[13, 0]) ** 2 + (pose_keypoints[14, 1] - pose_keypoints[13, 1]) ** 2): + pose_keypoints[21, 2] = 0 + + ### foot keypoints + if pose_keypoints[11, 2] <= confidence: + pose_keypoints[22, 2] = 0 + pose_keypoints[23, 2] = 0 + pose_keypoints[24, 2] = 0 + if pose_keypoints[14, 2] <= confidence: + pose_keypoints[19, 2] = 0 + pose_keypoints[20, 2] = 0 + pose_keypoints[21, 2] = 0 + return pose_keypoints + + def draw_pose_keypoints(self, img, pose_keypoints, confidence): + if len(pose_keypoints) == 0: + return img + + # self.convert_pose_points(pose_keypoints, confidence) + + connect_body = [[0, 15], [0, 16], [15, 17], [16, 18], + [0, 1], [1, 2], [1, 5], [2, 3], + [3, 4], [5, 6], [6, 7], [1, 8], + [8, 9], [9, 10], [10, 11], [8, 12], + [12, 13], [13, 14], [14, 21], [14, 19], + [19, 20], [11, 24], [11, 22], [22, 23]] + color_body = [(128, 245, 223), (245, 223, 113), (35, 243, 46), (67, 187, 255), (222, 33, 245), (192, 133, 45)] + for i in range(len(pose_keypoints)): + if pose_keypoints[i, 2] > confidence: + cv2.circle(img, (int(pose_keypoints[i, 0]), int(pose_keypoints[i, 1])), 3, (0, 255, 0), 4) + cv2.putText(img, str(i), (int(pose_keypoints[i, 0]), int(pose_keypoints[i, 1])), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 0, 0), 1) + for i, con in enumerate(connect_body): + pt_0 = con[0] + pt_1 = con[1] + if pose_keypoints[pt_0, 2] > confidence and pose_keypoints[pt_1, 2] > confidence: + cv2.line(img, (int(pose_keypoints[pt_0, 0]), int(pose_keypoints[pt_0, 1])), + (int(pose_keypoints[pt_1, 0]), int(pose_keypoints[pt_1, 1])), + color_body[int(i / 4)], 3) + return img + def draw_single_hand_keypoints(self, img, pts, confidence): + connect_hand = [[0, 1], [1, 2], [2, 3], [3, 4], + [0, 5], [5, 6], [6, 7], [7, 8], + [0, 9], [9, 10], [10, 11], [11, 12], + [0, 13], [13, 14], [14, 15], [15, 16], + [0, 17], [17, 18], [18, 19], [19, 20]] + + color_hand = [(128, 245, 223), (245, 223, 113), (35, 243, 46), (67, 187, 255), (222, 33, 245)] + + for i in range(len(pts)): + if pts[i, 2] > confidence: + cv2.circle(img, (int(pts[i, 0]), int(pts[i, 1])), 1, (0, 255, 0), 2) + # cv2.putText(img, str(i), (int(pts[i, 0]), int(pts[i, 1])), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (255, 0, 0), 1) + for i, con in enumerate(connect_hand): + pt_0 = con[0] + pt_1 = con[1] + if pts[pt_0, 2] > confidence and pts[pt_1, 2] > confidence: + cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), + color_hand[int(i / 4)], 2) + + return img + + def draw_hand_keypoints(self, img, hands_list, hand_keypoints, confidence): + if len(hand_keypoints) == 0: + return img + connect_hand = [[0, 1], [1, 2], [2, 3], [3, 4], + [0, 5], [5, 6], [6, 7], [7, 8], + [0, 9], [9, 10], [10, 11], [11, 12], + [0, 13], [13, 14], [14, 15], [15, 16], + [0, 17], [17, 18], [18, 19], [19, 20]] + + color_hand = [(128, 245, 223), (245, 223, 113), (35, 243, 46), (67, 187, 255), (222, 33, 245)] + for single_hand_keypoints in hand_keypoints: + pts = single_hand_keypoints["hand_keypoints"] + for i in range(len(pts)): + if pts[i, 2] > confidence: + cv2.circle(img, (int(pts[i, 0]), int(pts[i, 1])), 1, (0, 255, 0), 2) + # cv2.putText(img, str(i), (int(pts[i, 0]), int(pts[i, 1])), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (255, 0, 0), 1) + for i, con in enumerate(connect_hand): + pt_0 = con[0] + pt_1 = con[1] + if pts[pt_0, 2] > confidence and pts[pt_1, 2] > confidence: + cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), + color_hand[int(i / 4)], 2) + for x, y, w, is_left in hands_list: + cv2.rectangle(img, (x, y), (x+w, y+w), (255, 0, 255), 1) + return img + + def draw_face_keypoints(self, img, face_keypoints): + if len(face_keypoints) == 0: return img + pts = face_keypoints + for i in range(len(pts)): + cv2.circle(img, (int(pts[i, 0]), int(pts[i, 1])), 1, (0, 255, 0), 1) + ### draw face_outline + for i in range(0, 16): + pt_0 = i + pt_1 = i + 1 + cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), + (0, 255, 0), 1) + ### draw left eye_brow + for i in range(17, 21): + pt_0 = i + pt_1 = i + 1 + cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), + (0, 255, 0), 1) + ### draw right eye_brow + for i in range(22, 26): + pt_0 = i + pt_1 = i + 1 + cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), + (0, 255, 0), 1) + ### draw right eye + for i in range(42, 47): + pt_0 = i + pt_1 = i + 1 + cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), + (0, 255, 0), 1) + ### draw left eye + for i in range(42, 47): + pt_0 = i + pt_1 = i + 1 + cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), + (0, 255, 0), 1) + ### draw left eye + for i in range(36, 41): + pt_0 = i + pt_1 = i + 1 + cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), + (0, 255, 0), 1) + ### draw nose line + for i in range(27, 30): + pt_0 = i + pt_1 = i + 1 + cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), + (0, 255, 0), 1) + + ### draw nose hole + for i in range(31, 35): + pt_0 = i + pt_1 = i + 1 + cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), + (0, 255, 0), 1) + ### draw big mouth + for i in range(48, 59): + pt_0 = i + pt_1 = i + 1 + cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), + (0, 255, 0), 1) + ### draw small mouth + for i in range(60, 69): + pt_0 = i + pt_1 = i + 1 + cv2.line(img, (int(pts[pt_0, 0]), int(pts[pt_0, 1])), (int(pts[pt_1, 0]), int(pts[pt_1, 1])), + (0, 255, 0), 1) + # cv2.rectangle(img, (face_rect[0], face_rect[1]), (face_rect[2], face_rect[3]), (0, 255, 255), 1) + return img + +class Get_Landmark(object): + def __init__(self, gpu_id=0): + self.img_size = 640 + self.face_alignmenter_1k = MomocvFaceAlignment1K(gpu_id=gpu_id) + self.face_detector = MTCNNFaceDetector(gpu_id=gpu_id) + def get_max_rect(self, bounding_boxes): + max_area = 0 + index = 0 + for i, box in enumerate(bounding_boxes): + width = box[2] - box[0] + height = box[3] - box[1] + if width * height > max_area: + index = i + max_area = width * height + return index + def forward(self, img): + # bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=50, + # thresholds=[0.6, 0.7, 0.9]) + bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=30, + thresholds=[0.6, 0.6, 0.6]) + if len(bounding_boxes) == 0: + return None + + if len(bounding_boxes) > 0: + box_index = self.get_max_rect(bounding_boxes) + pts5 = landmarks[box_index] + else: + return None + + landmarks1k = self.face_alignmenter_1k.detect_according_5pts(img, pts5) + + return landmarks1k + + def forward_infer(self, img): + bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=50) + + if len(bounding_boxes) == 0: + return None, None, None + + if len(bounding_boxes) >= 1: + box_index = self.get_max_rect(bounding_boxes) + pts5 = landmarks[box_index] + else: + return None, None, None + + landmarks1k = self.face_alignmenter_1k.detect_according_5pts(img, pts5) + return landmarks1k + def forward_color(self, img): + bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=50) + + if len(bounding_boxes) == 0: + return None + + if len(bounding_boxes) >= 1: + box_index = self.get_max_rect(bounding_boxes) + pts5 = landmarks[box_index] + else: + return None + + landmarks1k = self.face_alignmenter_1k.detect_according_5pts(img, pts5) + return landmarks1k + + + def forward_infer(self, img): + bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=50) + + if len(bounding_boxes) == 0: + return None, None, None + + if len(bounding_boxes) >= 1: + box_index = self.get_max_rect(bounding_boxes) + pts5 = landmarks[box_index] + else: + return None, None, None + + landmarks1k = self.face_alignmenter_1k.detect_according_5pts(img, pts5) + return landmarks1k + + def forward_diy(self, img): + # bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=100,thresholds=[0.6, 0.8, 0.8]) # mtcnn face detect input + bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=50) + + if len(bounding_boxes) == 0: + return None, None, None + + if len(bounding_boxes) >=1 : + box_index = self.get_max_rect(bounding_boxes) + pts5 = landmarks[box_index] + else: + return None, None, None + # for i in range(5): + # cv2.circle(img, (int(pts5[i*2]), int(pts5[i*2+1])), 1, (255, 0,0), 1) + # cv2.imshow('img', img) + # cv2.waitKey() + landmarks1k = self.face_alignmenter_1k.detect_according_5pts(img, pts5) + # pts137 = landmark_processor.pts_1k_to_137(landmarks1k) + # movie_params = self.model_3d.detect([img], [pts137])[0] + # pitch, yaw, roll = movie_params[1:4] + ret_info = [landmarks1k, bounding_boxes[box_index], None] + return ret_info + + def forward_v2(self, img): + # bounding_boxes, landmarks = self.face_detector.forward(img, min_face_size=50, + # thresholds=[0.6, 0.7, 0.9]) + bounding_boxes, landmarks = self.face_detector.forward_v2(img, min_face_size=30, + thresholds=[0.6, 0.6, 0.6]) + if len(bounding_boxes) == 0: + return None + + if len(bounding_boxes) > 0: + box_index = self.get_max_rect(bounding_boxes) + pts5 = landmarks[box_index] + else: + return None + + landmarks1k = self.face_alignmenter_1k.detect_according_5pts(img, pts5) + + return landmarks1k + + def get_face_shape(self, xiaba_type, landmark137): + + face_shape = ['chang', 'fang', 'yuan', 'tuoyuan', 'xin'] + forhead_idxs = [12, 13, 14] + mid_idxs = [15, 16, 17] + lowwer_idxs = [18, 19, 20] + face_width = max(landmark137[:,0]) - min(landmark137[:,0]) + face_widthest_idx = list(landmark137[:,0]).index(min(landmark137[:,0])) + face_scale_lw = 1.4 + face_scale_qx_top = 1.4 + face_scale_qx_bot = 1.2 + + face_height = landmark137[0][1] - landmark137[11][1] + face_width_qiane = landmark137[9][0] - landmark137[13][0] + face_width_xiahe = landmark137[3][0] - landmark137[19][0] + check_scale = face_width_qiane / face_width_xiahe + print('face_height', face_height) + print('face_width', face_width) + print('face_scale_lw', face_scale_lw) + print('face_width_qiane', face_width_qiane) + print('face_width_xiahe', face_width_xiahe) + print('face_scale_qx', check_scale) + if face_widthest_idx in forhead_idxs: + if face_height / face_width >= face_scale_lw: + return face_shape[0] + else: + if check_scale > face_scale_qx_bot and face_scale_qx_bot < face_scale_qx_top: + if xiaba_type == 'jian': + return face_shape[0] + elif xiaba_type == 'fang': + return face_shape[1] + else: + return face_shape[1] + elif check_scale > face_scale_qx_top: + if xiaba_type == 'jian': + return face_shape[4] + elif xiaba_type == 'fang': + return face_shape[4] + else: + return face_shape[1] + else: + return face_shape[1] + + elif face_widthest_idx in mid_idxs: + if face_height / face_width >= face_scale_lw: + return face_shape[0] + else: + if check_scale > face_scale_qx_bot and face_scale_qx_bot < face_scale_qx_top: + if xiaba_type == 'jian': + return face_shape[3] + elif xiaba_type == 'fang': + return face_shape[1] + else: + return face_shape[2] + elif check_scale > face_scale_qx_top: + if xiaba_type == 'jian': + return face_shape[3] + elif xiaba_type == 'fang': + return face_shape[3] + else: + return face_shape[2] + else: + return face_shape[1] + else: + if face_height / face_width >= face_scale_lw: + return face_shape[0] + else: + if check_scale > face_scale_qx_bot and face_scale_qx_bot < face_scale_qx_top: + if xiaba_type == 'jian': + return face_shape[0] + elif xiaba_type == 'fang': + return face_shape[1] + else: + return face_shape[1] + else: + return face_shape[1] +class GenTrimap(object): + def __init__(self): + self.erosion_kernels = [None] + [cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size, size)) for size in range(1,30)] + + def __call__(self, alpha): + + fg_mask = np.zeros_like(alpha) + bg_mask = np.zeros_like(alpha) + fg_mask[alpha == 255] = 1 + bg_mask[alpha == 0] = 1 + + fg_mask = fg_mask.astype(np.int32).astype(np.uint8) + bg_mask = bg_mask.astype(np.int32).astype(np.uint8) + + fg_mask = cv2.erode(fg_mask, self.erosion_kernels[15]) + bg_mask = cv2.erode(bg_mask, self.erosion_kernels[29]) + + trimap = np.ones_like(alpha, dtype=np.uint8) * 128 + trimap[fg_mask == 1] = 255 + trimap[bg_mask == 1] = 0 + return trimap +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) + alpha_pred, info_dict = model(image, trimap) + + fg_pred = alpha_pred[:, :-1, :, :] + alpha_pred = alpha_pred[:, -1, :, :].unsqueeze(1) + + 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 = util.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 = util.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 +class Generator_Matte(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.output_img_size = 512 + + triseg_model_path = os.path.join(modelRoot, 'deeplabv3_hair512_360_0520_wl.pth') + self.triseg_model = Evaluator(gpu_id=device_id, output_img_size=self.output_img_size, nclass=3, seg_model_path=triseg_model_path) + + hair_matte_model_path = os.path.join(modelRoot, 'gca-dist-fg-0430-latest_model.pth') # gca-dist-fg-0203-latest_model gca-dist-fg-0430-latest_model + self.matte_model = self.load_hair_matte_model(hair_matte_model_path) + + self.gen_trimap = GenTrimap() + + def load_hair_matte_model(self, hair_matte_model_path): + # build model + model = generators.get_generator(encoder="resnet_gca_encoder_29", decoder="res_gca_decoder_22", num_class=4) + + # load checkpoint + checkpoint = torch.load(hair_matte_model_path, map_location=lambda storage, loc: storage) + model.load_state_dict(util.remove_prefix_state_dict(checkpoint['state_dict']), strict=True) + model.to(self.device) + # print("matte_model: ", model) + + # inference + model = model.eval() + return model + + def matte_inference(self, image, landmark1k): + trimap = self.triseg_model.eval(image, landmark1k) + + ori_h, ori_w, _ = image.shape + + limit_size = 1600 + if ori_h > limit_size or ori_w > limit_size: + if ori_h > ori_w: + new_tri_h = limit_size + new_tri_w = int(ori_w * limit_size / ori_h) + else: + new_tri_w = limit_size + new_tri_h = int(ori_h * limit_size / ori_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 + + trimap = self.gen_trimap(trimap_resize[:, :, 0]) + + # cv2.imwrite(os.path.join(args.output, image_name.replace(ext, "_trimap.png")), trimap) + + image_dict = generator_tensor_dict(image_resize, trimap) + + pred_fg, pred, offset = single_inference(self.matte_model, image_dict, device=self.device) + + pred_fg[trimap == 1] = image_resize[trimap == 1] + + if pred.shape[1] != image.shape[1] or pred.shape[0] != image.shape[0]: + pred = cv2.resize(pred, (image.shape[1], image.shape[0])) + + return pred_fg, pred +class Generator_Bald(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.model_dir = modelRoot + generator_bald_model_path = os.path.join(self.model_dir, 'zxm_v7_gen_bald_add_changed_hair_pairdata_5seg_0611.pt') + print("load generator_bald_model_path", generator_bald_model_path) + self.load_generator_bald768_model(generator_bald_model_path) + self.output_size = 768 + + def load_generator_bald768_model(self, generator_bald_model_path): + self.generator_bald_model = torch.jit.load(generator_bald_model_path, map_location='cpu').to(self.device) + + def Geneator_Bald_inference(self, user_rgb_8uc3_bald_768, user_matting_8uc3_bald_768, user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768): + + """ + input： + 图像尺寸基于： 人脸 512， 图像均为3通道 + user_rgb_8uc3_bald_512: 输入图， size 512， uint8 （0-255） + user_matting_8uc3_bald_512： 输入图 matting alpha 值， uint8 （0-255） + user_baldseg_8uc3_bald_512： 输入图 光头分割， uint8 （0-255） + + user_landmark_f1k2_bald_512：输入图 关键点 1k*2 float32 + + output： + bald_gene_8uc3_bald_512: 生成 光头图， uint8 （0-255） + + """ + # cv2.imshow("user_rgb_8uc3_bald_768 ori", user_rgb_8uc3_bald_768) + # cv2.imshow("user_matting_8uc3_bald_768", user_matting_8uc3_bald_768) + # cv2.imshow("user_baldseg_8uc3_bald_768", user_baldseg_8uc3_bald_768) + + # user_baldseg_8uc3_bald_768_cp = user_baldseg_8uc3_bald_768.copy() + # cood_y = int((user_landmark_f1k2_bald_768[214, 1] + user_landmark_f1k2_bald_768[99, 1]) // 2) + # + # random_int = 30 # random.randint(0, 30) + # kernel = np.ones((random_int, 1), np.uint8) + # user_mask_erode = cv2.erode(user_baldseg_8uc3_bald_768, kernel, iterations=1) + # user_baldseg_8uc3_bald_768[:cood_y, :, :] = user_mask_erode[:cood_y, :, :] + + # bald_con = np.concatenate((user_baldseg_8uc3_bald_768_cp, user_baldseg_8uc3_bald_768), axis=1) + # cv2.imshow("user_rgb_8uc3_bald_768", user_rgb_8uc3_bald_768) + # cv2.imshow("user_matting_8uc3_bald_768", user_matting_8uc3_bald_768) + + # user_matting_8uc3_bald_768_nonzero = np.zeros_like(user_matting_8uc3_bald_768) + # user_matting_8uc3_bald_768_nonzero[(user_matting_8uc3_bald_768 > 50).all(axis=2)] = 255 + # + # cv2.imshow("user_matting_8uc3_bald_768 >0", user_matting_8uc3_bald_768_nonzero) + # cv2.imshow("bald_con", bald_con) + # cv2.waitKey() + + # cv2.imwrite("/media/DATA_4T/project/hairstyle/分割测试/debug_line_res/user_rgb_8uc3_bald_768.png", user_rgb_8uc3_bald_768) + # cv2.imwrite("/media/DATA_4T/project/hairstyle/分割测试/debug_line_res/user_matting_8uc3_bald_768.png", user_matting_8uc3_bald_768) + # cv2.imwrite("/media/DATA_4T/project/hairstyle/分割测试/debug_line_res/user_baldseg_8uc3_bald_768.png", user_baldseg_8uc3_bald_768) + # np.savetxt("/media/DATA_4T/project/hairstyle/分割测试/debug_line_res/user_landmark_f1k2_bald_768.txt", user_landmark_f1k2_bald_768) + + real_res = user_rgb_8uc3_bald_768.copy() + + user_landmark_f1k2_bald_768 = user_landmark_f1k2_bald_768.astype(np.int32) + # kernel = np.ones((8, 8), np.uint8) + kernel_1 = np.ones((20, 20), np.uint8) + kernel_2 = np.ones((40, 40), np.uint8) + user_matting_8uc3_bald_768_dilate_1 = cv2.dilate(user_matting_8uc3_bald_768[:, :, 0], kernel_1, iterations=1) + user_matting_8uc3_bald_768_dilate_2 = cv2.dilate(user_matting_8uc3_bald_768[:, :, 0], kernel_2, iterations=1) + + user_matting_8uc3_bald_768_dilate_2[user_matting_8uc3_bald_768_dilate_2 > 0] = 255 + + b = user_baldseg_8uc3_bald_768[:, :, 0] > 125 + g = user_baldseg_8uc3_bald_768[:, :, 1] > 0 + r = user_baldseg_8uc3_bald_768[:, :, 2] > 125 + + user_matting_8uc3_bald_768_blur = user_matting_8uc3_bald_768_dilate_2.copy() + # cv2.imshow("user_matting_8uc3_bald_768_blur noface", user_matting_8uc3_bald_768_blur) + + # TODO change noface mask + # user_matting_8uc3_bald_768_blur[g] = 255 + user_matting_8uc3_bald_768_blur = cv2.blur(user_matting_8uc3_bald_768_blur, (30, 30)) + # user_matting_8uc3_bald_768_blur_2 = cv2.blur(user_matting_8uc3_bald_768_blur, (20, 20)) + # user_matting_8uc3_bald_768_blur[(~b) & (~g) & (~r)] = user_matting_8uc3_bald_768_blur_2[(~b) & (~g) & (~r)] + + # cv2.imshow("user_matting_8uc3_bald_768_blur", user_matting_8uc3_bald_768_blur) + + # user_matting_8uc3_bald_768_dilate_1[(~b) * (~g)] = user_matting_8uc3_bald_768[:, :, 0][(~b) * (~g)] + user_matting_8uc3_bald_768_dilate_1[(~b) * (~g) * (~r)] = user_matting_8uc3_bald_768[:, :, 0][(~b) * (~g) * (~r)] + # user_matting_8uc3_bald_768_dilate_1[(~b) * g * r] = user_matting_8uc3_bald_768[:, :, 0][(~b) * g * r] + user_rgb_8uc3_bald_768[user_matting_8uc3_bald_768_dilate_1.astype(bool)] = 255 + + ######################################## + # cv2.fillPoly(user_baldseg_8uc3_bald_512, user_pts1k[:311][np.newaxis, :, :], (200, 175, 0)) # face + cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[928:999][np.newaxis, :, :], + (0, 125, 0)) # left eyebrow + cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[856:927][np.newaxis, :, :], + (125, 0, 0)) # right eyebrow + cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[691:754][np.newaxis, :, :], + (0, 0, 125)) # left eye + cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[792:855][np.newaxis, :, :], + (125, 125, 0)) # right eye + cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[548:616][np.newaxis, :, :], + (0, 125, 125)) # rose + cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[312:547][np.newaxis, :, :], + (125, 125, 125)) # mouth + cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[655:690][np.newaxis, :, :], + (0, 0, 175)) # left black_eyeball + cv2.fillPoly(user_baldseg_8uc3_bald_768, user_landmark_f1k2_bald_768[756:791][np.newaxis, :, :], + (125, 0, 125)) # right black_eyeball + + eyes_mask = np.zeros(user_baldseg_8uc3_bald_768[:, :, 0].shape) + cv2.fillPoly(eyes_mask, user_landmark_f1k2_bald_768[691:754][np.newaxis, :, :], 1) # left eye + cv2.fillPoly(eyes_mask, user_landmark_f1k2_bald_768[792:855][np.newaxis, :, :], 1) # right eye + kernel = np.ones((30, 30), np.uint8) + eyes_mask = cv2.dilate(eyes_mask, kernel, iterations=1) + cv2.fillPoly(eyes_mask, user_landmark_f1k2_bald_768[312:467][np.newaxis, :, :], 1) # mouth + user_rgb_8uc3_bald_768[eyes_mask.astype(bool)] = real_res[eyes_mask.astype(bool)] + + ######################### add by zxm + # face_erode_mask = np.zeros(user_baldseg_8uc3_bald_768[:, :, 0].shape) + # cv2.fillPoly(face_erode_mask, user_landmark_f1k2_bald_768[:311][np.newaxis, :, :], 1) # face + # kernel = np.ones((60, 60), np.uint8) + # face_erode_mask = cv2.erode(face_erode_mask, kernel, iterations=1) + # user_rgb_8uc3_bald_768[face_erode_mask.astype(bool)] = real_res[face_erode_mask.astype(bool)] + # user_rgb_8uc3_bald_768[user_matting_8uc3_bald_768_dilate_copy.astype(bool)] = 255 + ######################### add by zxm + + # repeat + # user_rgb_8uc3_bald_512[eyes_mask.astype(bool)] = real_res[eyes_mask.astype(bool)] + + # cv2.imshow("user_baldseg_8uc3_bald_768_condition", user_baldseg_8uc3_bald_768) + # cv2.imshow("user_rgb_8uc3_bald_768_input_paf", user_rgb_8uc3_bald_768) + # cv2.waitKey() + + # cv2.imwrite( + # "/media/liyang/DATA1/test_hair/换状后台数据_测试/badcase_0205/光头badcase/bald_input_0128/bald_condition.png", + # user_baldseg_8uc3_bald_768) + + # cv2.imwrite( + # "/media/liyang/DATA1/test_hair/换状后台数据_测试/badcase_0205/光头badcase/bald_input_0128/inpu_paf.png", + # user_rgb_8uc3_bald_768) + + + # cv2.imshow("user_baldseg_8uc3_bald_768", user_baldseg_8uc3_bald_768) + # cv2.imshow("user_rgb_8uc3_bald_768", user_rgb_8uc3_bald_768) + # cv2.waitKey() + + condition = user_baldseg_8uc3_bald_768 # np.concatenate((user_baldseg_8uc3_bald_512, user_hair_mask), axis=2) + condition = (condition.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] + input_paf = (user_rgb_8uc3_bald_768.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] + condition = torch.from_numpy(condition).to(self.device) + input_paf = torch.from_numpy(input_paf).to(self.device) + + # ******************* forward *******************$ + # test_fake, _, _ = model.preview(input_paf, condition) + with torch.no_grad(): + test_fake = self.generator_bald_model(input_paf, condition) + + test_res = (test_fake.detach()[0].to('cpu').numpy().transpose((1, 2, 0)) * 255).astype(np.uint8).copy() + + # show_concat = np.concatenate((user_rgb_8uc3_bald_768, user_baldseg_8uc3_bald_768, test_res), axis=1) + # cv2.imshow("bald_concat", show_concat) + # cv2.waitKey() + + # kernel = np.ones((25, 25), np.uint8) + # user_hair_mask_dilate = cv2.dilate(user_matting_8uc3_bald_768[:, :, 0], kernel) + # user_hair_mask_dilate[user_hair_mask_dilate > 0] = 255 + # user_hair_mask_blur = cv2.blur(user_hair_mask_dilate, (20, 20)) + test_res_clear = real_res * (1 - user_matting_8uc3_bald_768_blur[:, :, np.newaxis] / 255) + test_res * ( + user_matting_8uc3_bald_768_blur[:, :, np.newaxis] / 255) + bald_gene_8uc3_bald_768 = (np.clip(test_res_clear, 0, 255)).astype(np.uint8) + + # cv2.imshow("bald_gene_8uc3_bald_768", bald_gene_8uc3_bald_768) + # cv2.imshow("user_matting_8uc3_bald_768_blur", user_matting_8uc3_bald_768_blur) + # cv2.imshow("test_res", test_res) + # cv2.waitKey() + + # cv2.imwrite("/media/DATA_4T/project/hairstyle/分割测试/debug_line_res/test_res_768.png", test_res) + # cv2.imwrite("/media/DATA_4T/project/hairstyle/分割测试/debug_line_res/bald_gene_8uc3_bald_768.png", bald_gene_8uc3_bald_768) + # cv2.imwrite("/media/DATA_4T/project/hairstyle/分割测试/debug_line_res/user_matting_8uc3_bald_768_blur.png", user_matting_8uc3_bald_768_blur) + + return bald_gene_8uc3_bald_768, user_matting_8uc3_bald_768_blur +class Process_Data(object): + def __init__(self, gpu, device_id, save_name=None): + 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.generator_matte = Generator_Matte(gpu, device_id) + + self.generator_baldseg = Generator_BaldSeg_5c(gpu, device_id) + self.generator_bald = Generator_Bald(gpu, device_id) + + self.hair_size = 768 + self.bald_output_size = 768 + self.color_output_size = 768 + + def get_hair_M(self, landmark1k): + + hairstyle_M = landmark_processor.get_transform_mat_hair_ratio(landmark1k, 512, 0.5) + + return hairstyle_M + + def get_hair_M_girl_v1(self, landmark1k): + + hairstyle_M = landmark_processor.get_transform_mat_hair_ratio_v1(landmark1k, 768, ratio=0.5, + h_offset=0.45) # default 0.5 0.45 ratio=0.35, h_offset=0.32 + + return hairstyle_M + + def get_hair_M_girl_v2(self, landmark1k, long=False): + + hairstyle_M = landmark_processor.get_transform_mat_hair_ratio_v1(landmark1k, 768, ratio=0.35, + h_offset=0.32) # default 0.5 0.45 ratio=0.35, h_offset=0.32 + + if long: + print("------long process-----------") + hairstyle_M = landmark_processor.get_transform_mat_hair_ratio_v1(landmark1k, 768, ratio=0.27, + h_offset=0.28) # default 0.5 0.45 ratio=0.35, h_offset=0.32 + + return hairstyle_M + + + def get_hair_M_girl_v3(self, landmark1k): + print("get_hair_M_girl_v3 start......") + hairstyle_M = landmark_processor.get_transform_mat_hair_ratio_v1(landmark1k, 768, ratio=0.33, + h_offset=0.30) # default 0.5 0.45 ratio=0.35, h_offset=0.32 + print("get_hair_M_girl_v3 end......") + return hairstyle_M + + def get_hair_M_boy_v1(self, landmark1k): + + hairstyle_M = landmark_processor.get_transform_mat_hair_ratio_v1(landmark1k, 768, ratio=0.6, h_offset=0.65) + + return hairstyle_M + + def get_color_hair_M(self, landmark1k): + + # color_hair_M = landmark_processor.get_transform_mat_full_face_ratio_stylegan(landmark1k, self.color_output_size, 0.35) # 0.4 + color_hair_M = landmark_processor.get_transform_mat_hair_ratio_v1(landmark1k, self.color_output_size, + ratio=0.35, h_offset=0.35) + + return color_hair_M + + def get_prepare_data_bald(self, user_rgb_8uc3_orisize, user_baldseg_8uc3_orisize, + user_landmark_1k2_f_orisize): + user_pts1k = user_landmark_1k2_f_orisize.astype(np.int32) + user_bald_M = landmark_processor.get_transform_mat_full_face_ratio_stylegan(user_pts1k, self.bald_output_size, + 0.4) + + user_rgb_8uc3_bald_512 = cv2.warpAffine(user_rgb_8uc3_orisize, user_bald_M, + (self.bald_output_size, self.bald_output_size)) + + user_baldseg_8uc3_bald_512 = cv2.warpAffine(user_baldseg_8uc3_orisize, user_bald_M, + (self.bald_output_size, self.bald_output_size), + flags=cv2.INTER_NEAREST) + + user_landmark_f1k2_bald_512 = landmark_processor.transform_points(user_pts1k, user_bald_M) + + return user_rgb_8uc3_bald_512, user_baldseg_8uc3_bald_512, user_landmark_f1k2_bald_512, user_bald_M + + def get_prepare_data_bald_768(self, user_baldseg_8uc3_orisize, + user_landmark_1k2_f_orisize, user_color_M): + user_pts1k = user_landmark_1k2_f_orisize.astype(np.int32) + user_bald_M = landmark_processor.get_transform_mat_full_face_ratio_stylegan(user_pts1k, self.hair_size, 0.4) + + # user_rgb_8uc3_bald_768 = cv2.warpAffine(user_rgb_8uc3_orisize, user_color_M, (self.bald_output_size, self.bald_output_size)) + + user_baldseg_8uc3_bald_768 = cv2.warpAffine(user_baldseg_8uc3_orisize, user_color_M, + (self.bald_output_size, self.bald_output_size), + flags=cv2.INTER_NEAREST) + + user_baldseg_8uc3_bald_512 = cv2.warpAffine(user_baldseg_8uc3_orisize, user_bald_M, + (self.hair_size, self.hair_size), + flags=cv2.INTER_NEAREST) + + # user_baldseg_8uc3_bald_768_cp = user_baldseg_8uc3_bald_768.copy() + # + # user_pts1k_768 = landmark_processor.transform_points(user_pts1k, user_color_M) + # cood_y = int((user_pts1k_768[214, 1] + user_pts1k_768[99, 1]) // 2) + + # random_int = 30 # random.randint(0, 30) + # kernel = np.ones((random_int, 1), np.uint8) + # user_mask_erode = cv2.erode(user_baldseg_8uc3_bald_768, kernel, iterations=1) + # user_baldseg_8uc3_bald_768[:cood_y, :, :] = user_mask_erode[:cood_y, :, :] + # + # inter_res = np.concatenate((user_baldseg_8uc3_bald_768_cp, user_baldseg_8uc3_bald_768), axis=1) + + # cv2.imshow("inter_res", inter_res) + # cv2.waitKey() + + # user_baldseg_8uc3_orisize = cv2.warpAffine(user_baldseg_8uc3_bald_768, cv2.invertAffineTransform(user_color_M), + # dsize=(user_baldseg_8uc3_orisize.shape[1], user_baldseg_8uc3_orisize.shape[0]), flags=cv2.INTER_NEAREST) + + return user_baldseg_8uc3_bald_768, user_baldseg_8uc3_bald_512, user_bald_M # , user_baldseg_8uc3_orisize + + def get_user_blad(self, user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize): + """ + input： + + user_rgb_8uc3_orisize: 用户图 原始尺寸 + user_landmark_1k2_f_orisize: 用户图 1k 点 + + output: + + 图像尺寸基于： 人脸 512， 图像均为3通道 + + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + + """ + + # 得到512 小图 M 矩阵 + user_color_M = self.get_color_hair_M(user_landmark_1k2_f_orisize) + # 得到512 小图 M 矩阵 + # user_hairstyle_M = self.get_hair_M(user_landmark_1k2_f_orisize) + + # 1k 点 转换到 512 尺寸 + # user_landmark_f1k2_512 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_hairstyle_M) + user_landmark_f1k2_bald_768 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_color_M) + + # # 用户图 头发毛躁 warpaffine + usr_color_ratio = np.sqrt((user_color_M[0][0] * user_color_M[0][0]) + (user_color_M[1][0] * user_color_M[1][0])) + user_color_M_tmp = user_color_M / usr_color_ratio + user_rgb_8uc3_bald_768 = cv2.warpAffine(user_rgb_8uc3_orisize, user_color_M_tmp, ( + int(self.bald_output_size / usr_color_ratio), int(self.bald_output_size / usr_color_ratio)), + borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) + user_rgb_8uc3_bald_768 = cv2.resize(user_rgb_8uc3_bald_768, (self.bald_output_size, self.bald_output_size), + fx=usr_color_ratio, fy=usr_color_ratio, interpolation=cv2.INTER_AREA) + + user_rgb_8uc3_bald_768_bl_bg = cv2.warpAffine(user_rgb_8uc3_orisize, user_color_M_tmp, ( + int(self.bald_output_size / usr_color_ratio), int(self.bald_output_size / usr_color_ratio))) + user_rgb_8uc3_bald_768_bl_bg = cv2.resize(user_rgb_8uc3_bald_768_bl_bg, + (self.bald_output_size, self.bald_output_size), fx=usr_color_ratio, + fy=usr_color_ratio, interpolation=cv2.INTER_AREA) + + # 用户图得到 头发 matting 小图 + _, user_matting_8uc1_bald_768 = self.generator_matte.matte_inference(user_rgb_8uc3_bald_768, + user_landmark_f1k2_bald_768) + + # cv2.imshow('user_rgb_8uc3_bald_768', user_rgb_8uc3_bald_768) + # cv2.imshow('user_matting_8uc1_bald_768', user_matting_8uc1_bald_768) + # cv2.waitKey() + + if user_matting_8uc1_bald_768.shape[0] != 768 or user_matting_8uc1_bald_768.shape[1] != 768: + user_matting_8uc1_bald_768 = cv2.resize(user_matting_8uc1_bald_768, + (user_rgb_8uc3_bald_768.shape[1], user_rgb_8uc3_bald_768.shape[0]), + interpolation=cv2.INTER_CUBIC) + user_matting_8uc3_bald_768 = np.repeat(user_matting_8uc1_bald_768[:, :, np.newaxis], 3, axis=2) + + user_matting_8uc3_bald_orisize = cv2.warpAffine(user_matting_8uc1_bald_768, + cv2.invertAffineTransform(user_color_M), + ( + user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), + flags=cv2.INTER_CUBIC) + + # 光头分割 + user_baldseg_8uc3_orisize = self.generator_baldseg.forward(user_rgb_8uc3_orisize, + user_matting_8uc3_bald_orisize, + user_landmark_1k2_f_orisize) + + # 用户图 生成光头 + user_baldseg_8uc3_bald_768, user_baldseg_8uc3_bald_512, user_bald_M \ + = self.get_prepare_data_bald_768(user_baldseg_8uc3_orisize, + user_landmark_1k2_f_orisize, user_color_M) + + orig_h, orig_w, _ = user_rgb_8uc3_orisize.shape + + # TODO: 测试直接采用M 矩阵warp 到 user_hairstyle_M 后的512 尺寸 + # bald_gene_8uc3_bald_768, user_hair_mask_blur_768 = self.generator_bald.Geneator_Bald_inference(user_rgb_8uc3_bald_768.copy(), + # user_matting_8uc3_bald_768, + # user_baldseg_8uc3_bald_768, + # user_landmark_f1k2_bald_768) + + bald_gene_8uc3_bald_768, user_hair_mask_blur_768 = self.generator_bald.Geneator_Bald_inference( + user_rgb_8uc3_bald_768_bl_bg.copy(), + user_matting_8uc3_bald_768, + user_baldseg_8uc3_bald_768, + user_landmark_f1k2_bald_768) + + inv_M_user = cv2.invertAffineTransform(user_color_M) + bald_gene_8uc3_orisize = user_rgb_8uc3_orisize.copy() + + user_bald_gene_8uc3_orisize = cv2.warpAffine(bald_gene_8uc3_bald_768, inv_M_user, (orig_w, orig_h), + dst=bald_gene_8uc3_orisize.copy(), + borderMode=cv2.BORDER_TRANSPARENT) + + user_hair_mask_blur_orisize = cv2.warpAffine(user_hair_mask_blur_768, inv_M_user, (orig_w, orig_h), + flags=cv2.INTER_CUBIC) + + user_bald_res_8uc3_orisize = ( + (1 - user_hair_mask_blur_orisize[:, :, np.newaxis] / 255) * user_rgb_8uc3_orisize + \ + (user_hair_mask_blur_orisize[:, :, np.newaxis] / 255) * user_bald_gene_8uc3_orisize).astype( + np.uint8) + + return user_bald_res_8uc3_orisize, user_bald_gene_8uc3_orisize, user_matting_8uc3_bald_768, user_baldseg_8uc3_bald_768 + + def hair512_to_768(self, user_color_M, user_hairstyle_M): + + M_ori = np.zeros((3, 3), dtype=np.float32) + M_ori[:2, :] = cv2.invertAffineTransform(user_hairstyle_M) + M_ori[2:, :] = [0, 0, 1] + + matAffine_ori = np.zeros((3, 3), dtype=np.float32) + matAffine_ori[:2, :] = user_color_M + matAffine_ori[2:, :] = [0, 0, 1] + + new_mat = matAffine_ori.dot(M_ori) + return new_mat[:2, :] + + def get_prepare_data(self, user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize, ref_rgb_8uc3_orisize, + ref_landmark_1k2_f_orisize): + """ + input： + + user_rgb_8uc3_orisize: 用户图 原始尺寸 + user_landmark_1k2_f_orisize: 用户图 1k 点 + ref_rgb_8uc3_orisize: 参考图 原始尺寸 + ref_landmark_1k2_f_orisize: 参考图 1k 点 + + output: + + 图像尺寸基于： 人脸 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 + + user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 + user_hairstyle_M: 用户图 原图到512小图 M 矩阵 + """ + + # 得到512 小图 M 矩阵 + user_hairstyle_M = self.get_hair_M(user_landmark_1k2_f_orisize) + ref_hairstyle_M = self.get_hair_M(ref_landmark_1k2_f_orisize) + + # 1k 点 转换到 512 尺寸 + user_landmark_f1k2_512 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_hairstyle_M) + ref_landmark_f1k2_512 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, ref_hairstyle_M) + + # 用户图 头发毛躁 warpaffine + usr_ratio = np.sqrt( + (user_hairstyle_M[0][0] * user_hairstyle_M[0][0]) + (user_hairstyle_M[1][0] * user_hairstyle_M[1][0])) + user_hairstyle_M_tmp = user_hairstyle_M / usr_ratio + user_rgb_8uc3_512 = cv2.warpAffine(user_rgb_8uc3_orisize, user_hairstyle_M_tmp, + (int(self.hair_size / usr_ratio), int(self.hair_size / usr_ratio)), + borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) + user_rgb_8uc3_512 = cv2.resize(user_rgb_8uc3_512, (self.hair_size, self.hair_size), fx=usr_ratio, fy=usr_ratio, + interpolation=cv2.INTER_AREA) + + # 参考图 头发毛躁 warpaffine + ratio = np.sqrt( + (ref_hairstyle_M[0][0] * ref_hairstyle_M[0][0]) + (ref_hairstyle_M[1][0] * ref_hairstyle_M[1][0])) + ref_hairstyle_M_tmp = ref_hairstyle_M / ratio + ref_rgb_8uc3_512 = cv2.warpAffine(ref_rgb_8uc3_orisize, ref_hairstyle_M_tmp, + (int(self.hair_size / ratio), int(self.hair_size / ratio)), + borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) + ref_rgb_8uc3_512 = cv2.resize(ref_rgb_8uc3_512, (self.hair_size, self.hair_size), fx=ratio, fy=ratio, + interpolation=cv2.INTER_AREA) + + # ref_rgb_8uc3_512 = cv2.warpAffine(ref_rgb_8uc3_orisize, ref_hairstyle_M, (self.hair_size, self.hair_size), borderValue=[255, 255, 255]) + + # 参考图直接得到 头发 matting 小图 + _, ref_matte_pred_8uc1_512 = self.generator_matte.matte_inference(ref_rgb_8uc3_512, ref_landmark_f1k2_512) + + torch.cuda.empty_cache() + + ref_matting_8uc1_512 = cv2.resize(ref_matte_pred_8uc1_512, + (ref_rgb_8uc3_512.shape[1], ref_rgb_8uc3_512.shape[0]), + interpolation=cv2.INTER_CUBIC) + + ref_matting_8uc3_512 = np.repeat(ref_matting_8uc1_512[:, :, np.newaxis], 3, axis=2) + + # 光头分割 + user_baldseg_8uc3_orisize = self.generator_baldseg.forward(user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize) + ref_baldseg_8uc3_orisize = self.generator_baldseg.forward(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize) + + # 用户图 生成光头 + + user_rgb_8uc3_bald_512, user_baldseg_8uc3_bald_512, user_landmark_f1k2_bald_512, \ + user_bald_M = self.get_prepare_data_bald(user_rgb_8uc3_orisize, user_baldseg_8uc3_orisize, + user_landmark_1k2_f_orisize) + + # 用户图得到 头发 matting 小图 for blad + _, user_matting_8uc1_bald_512 = self.generator_matte.matte_inference(user_rgb_8uc3_bald_512, + user_landmark_f1k2_bald_512) + if user_matting_8uc1_bald_512.shape[0] != 512 or user_matting_8uc1_bald_512.shape[1] != 512: + user_matting_8uc1_bald_512 = cv2.resize(user_matting_8uc1_bald_512, + (user_rgb_8uc3_bald_512.shape[1], user_rgb_8uc3_bald_512.shape[0]), + interpolation=cv2.INTER_CUBIC) + user_matting_8uc3_bald_512 = np.repeat(user_matting_8uc1_bald_512[:, :, np.newaxis], 3, axis=2) + + orig_h, orig_w, _ = user_rgb_8uc3_orisize.shape + + # TODO: 测试直接采用M 矩阵warp 到 user_hairstyle_M 后的512 尺寸 + bald_gene_8uc3_bald_512 = self.generator_bald.Geneator_Bald_inference(user_rgb_8uc3_bald_512, + user_matting_8uc3_bald_512, + user_baldseg_8uc3_bald_512, + user_landmark_f1k2_bald_512) + + inv_M_user = cv2.invertAffineTransform(user_bald_M) + bald_gene_8uc3_orisize = user_rgb_8uc3_orisize.copy() + + user_bald_gene_8uc3_orisize = cv2.warpAffine(bald_gene_8uc3_bald_512, inv_M_user, (orig_w, orig_h), + dst=bald_gene_8uc3_orisize, borderMode=cv2.BORDER_TRANSPARENT) + + user_baldseg_8uc3_512 = cv2.warpAffine(user_baldseg_8uc3_orisize, user_hairstyle_M, + (self.hair_size, self.hair_size)) + user_bald_8uc3_512 = cv2.warpAffine(user_bald_gene_8uc3_orisize, user_hairstyle_M, + (self.hair_size, self.hair_size), borderValue=[255, 255, 255]) + + ref_baldseg_8uc3_512 = cv2.warpAffine(ref_baldseg_8uc3_orisize, ref_hairstyle_M, + (self.hair_size, self.hair_size)) + + return user_rgb_8uc3_512, user_baldseg_8uc3_512, user_bald_8uc3_512, user_landmark_f1k2_512, user_hairstyle_M, ref_rgb_8uc3_512, ref_matting_8uc3_512, ref_baldseg_8uc3_512, ref_landmark_f1k2_512 + + def get_prepare_user_data(self, user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize): + """ + input：F + + user_rgb_8uc3_orisize: 用户图 原始尺寸 + user_landmark_1k2_f_orisize: 用户图 1k 点 + + output: + + 图像尺寸基于： 人脸 512， 图像均为3通道 + + user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 + user_hairstyle_M: 用户图 原图到512小图 M 矩阵 + """ + + # 得到512 小图 M 矩阵 + user_hairstyle_M = self.get_hair_M(user_landmark_1k2_f_orisize) + + # 1k 点 转换到 512 尺寸 + user_landmark_f1k2_512 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_hairstyle_M) + + # 用户图 头发毛躁 warpaffine + usr_ratio = np.sqrt( + (user_hairstyle_M[0][0] * user_hairstyle_M[0][0]) + (user_hairstyle_M[1][0] * user_hairstyle_M[1][0])) + user_hairstyle_M_tmp = user_hairstyle_M / usr_ratio + user_rgb_8uc3_512 = cv2.warpAffine(user_rgb_8uc3_orisize, user_hairstyle_M_tmp, + (int(self.hair_size / usr_ratio), int(self.hair_size / usr_ratio)), + borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) + user_rgb_8uc3_512 = cv2.resize(user_rgb_8uc3_512, (self.hair_size, self.hair_size), fx=usr_ratio, fy=usr_ratio, + interpolation=cv2.INTER_AREA) + + # 光头分割 + user_baldseg_8uc3_orisize = self.generator_baldseg.forward(user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize) + + # 用户图 生成光头 + user_rgb_8uc3_bald_512, user_baldseg_8uc3_bald_512, user_landmark_f1k2_bald_512, \ + user_bald_M = self.get_prepare_data_bald(user_rgb_8uc3_orisize, user_baldseg_8uc3_orisize, + user_landmark_1k2_f_orisize) + + # 用户图得到 头发 matting 小图 for blad + _, user_matting_8uc1_bald_512 = self.generator_matte.matte_inference(user_rgb_8uc3_bald_512, + user_landmark_f1k2_bald_512) + if user_matting_8uc1_bald_512.shape[0] != 512 or user_matting_8uc1_bald_512.shape[1] != 512: + user_matting_8uc1_bald_512 = cv2.resize(user_matting_8uc1_bald_512, + (user_rgb_8uc3_bald_512.shape[1], user_rgb_8uc3_bald_512.shape[0]), + interpolation=cv2.INTER_CUBIC) + user_matting_8uc3_bald_512 = np.repeat(user_matting_8uc1_bald_512[:, :, np.newaxis], 3, axis=2) + + orig_h, orig_w, _ = user_rgb_8uc3_orisize.shape + + # TODO: 测试直接采用M 矩阵warp 到 user_hairstyle_M 后的512 尺寸 + bald_gene_8uc3_bald_512 = self.generator_bald.Geneator_Bald_inference( + user_rgb_8uc3_bald_512, + user_matting_8uc3_bald_512, + user_baldseg_8uc3_bald_512, + user_landmark_f1k2_bald_512) + + inv_M_user = cv2.invertAffineTransform(user_bald_M) + bald_gene_8uc3_orisize = user_rgb_8uc3_orisize.copy() + + user_bald_gene_8uc3_orisize = cv2.warpAffine(bald_gene_8uc3_bald_512, inv_M_user, (orig_w, orig_h), + dst=bald_gene_8uc3_orisize, borderMode=cv2.BORDER_TRANSPARENT) + + user_baldseg_8uc3_512 = cv2.warpAffine(user_baldseg_8uc3_orisize, user_hairstyle_M, + (self.hair_size, self.hair_size)) + user_bald_8uc3_512 = cv2.warpAffine(user_bald_gene_8uc3_orisize, user_hairstyle_M, + (self.hair_size, self.hair_size), borderValue=[255, 255, 255]) + + return user_rgb_8uc3_512, user_baldseg_8uc3_512, user_bald_8uc3_512, user_landmark_f1k2_512, user_hairstyle_M + + def get_prepare_user_768_data(self, user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize, ratio=1): + """ + input： + + user_rgb_8uc3_orisize: 用户图 原始尺寸 + user_landmark_1k2_f_orisize: 用户图 1k 点 + + output: + + 图像尺寸基于： 人脸 512， 图像均为3通道 + + user_baldseg_8uc3_768：用户图 光头分割 mask， uint8 （0-255） + user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + user_landmark_f1k2_768： 用户图 关键点 1k*2 float32 + user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 + user_color_M: 用户图 原图到768小图 M 矩阵 + user_hairstyle_M: 用户图 原图到512小图 M 矩阵 + """ + # 得到768 小图 M 矩阵 + + user_color_M = self.get_color_hair_M(user_landmark_1k2_f_orisize) + # 得到768 小图 M 矩阵 发型 + + if ratio == 0: + user_hairstyle_M = self.get_hair_M_boy_v1(user_landmark_1k2_f_orisize) + elif ratio == 1: + user_hairstyle_M = self.get_hair_M_girl_v1(user_landmark_1k2_f_orisize) + elif ratio == 2: + user_hairstyle_M = self.get_hair_M_girl_v2(user_landmark_1k2_f_orisize) + elif ratio == 3: + user_hairstyle_M = self.get_hair_M_girl_v3(user_landmark_1k2_f_orisize) + else: + user_hairstyle_M = self.get_hair_M_girl_v1(user_landmark_1k2_f_orisize) + + # np.save("/home/yangchaojie/Desktop/diffusion_model/lora/data_for_liveme/tmp/hair_swap/person/1706086560_1706086559583_ST0023_W0_045_0.npy", user_hairstyle_M) + + # user_color_M = user_hairstyle_M + + # 1k 点 转换到 768 尺寸 + user_landmark_f1k2_bald_768 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_color_M) + # 1k 点 转换到 768 尺寸 发型 + user_landmark_f1k2_768 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_hairstyle_M) + + # 用户图 头发毛躁 warpaffine 768 + usr_color_ratio = np.sqrt( + (user_color_M[0][0] * user_color_M[0][0]) + (user_color_M[1][0] * user_color_M[1][0])) + user_color_M_tmp = user_color_M / usr_color_ratio + user_rgb_8uc3_768 = cv2.warpAffine(user_rgb_8uc3_orisize, user_color_M_tmp, + (int(self.bald_output_size / usr_color_ratio), + int(self.bald_output_size / usr_color_ratio)), + borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) + user_rgb_8uc3_768 = cv2.resize(user_rgb_8uc3_768, (self.bald_output_size, self.bald_output_size), + fx=usr_color_ratio, fy=usr_color_ratio, + interpolation=cv2.INTER_AREA) + + user_rgb_8uc3_768_bald = cv2.warpAffine(user_rgb_8uc3_orisize, user_color_M_tmp, + (int(self.bald_output_size / usr_color_ratio), + int(self.bald_output_size / usr_color_ratio))) + + user_rgb_8uc3_768_bald = cv2.resize(user_rgb_8uc3_768_bald, (self.bald_output_size, self.bald_output_size), + fx=usr_color_ratio, fy=usr_color_ratio, + interpolation=cv2.INTER_AREA) + + # 用户图得到 头发 matting 小图 for blad + user_matting_fg_8uc3_bald_768, user_matting_8uc1_bald_768 = self.generator_matte.matte_inference( + user_rgb_8uc3_768, + user_landmark_f1k2_bald_768) + + if user_matting_8uc1_bald_768.shape[0] != 768 or user_matting_8uc1_bald_768.shape[1] != 768: + user_matting_8uc1_bald_768 = cv2.resize(user_matting_8uc1_bald_768, + (user_rgb_8uc3_768.shape[1], user_rgb_8uc3_768.shape[0]), + interpolation=cv2.INTER_CUBIC) + + # user_matting_fg_8uc3_bald_768 = cv2.resize(user_matting_fg_8uc3_bald_768, + # (user_rgb_8uc3_768.shape[1], user_rgb_8uc3_768.shape[0]), + # interpolation=cv2.INTER_CUBIC) + + user_matting_8uc3_bald_768 = np.repeat(user_matting_8uc1_bald_768[:, :, np.newaxis], 3, axis=2) + + user_matting_8uc3_bald_orisize = cv2.warpAffine(user_matting_8uc1_bald_768, + cv2.invertAffineTransform(user_color_M), + ( + user_rgb_8uc3_orisize.shape[1], user_rgb_8uc3_orisize.shape[0]), + flags=cv2.INTER_CUBIC) + + # 光头分割 + user_baldseg_8uc3_orisize = self.generator_baldseg.forward(user_rgb_8uc3_orisize, + user_matting_8uc3_bald_orisize, + user_landmark_1k2_f_orisize) + # 用户图 生成光头 需准备 768 尺寸； 512 尺寸各一 + + user_baldseg_8uc3_bald_768, user_baldseg_8uc3_bald_512, user_bald_M \ + = self.get_prepare_data_bald_768(user_baldseg_8uc3_orisize, + user_landmark_1k2_f_orisize, user_color_M) + + orig_h, orig_w, _ = user_rgb_8uc3_orisize.shape + + # cv2.imshow("user_matting_8uc3_bald_768",user_matting_8uc3_bald_768) + bald_gene_8uc3_bald_768, user_hair_mask_blur_768 = self.generator_bald.Geneator_Bald_inference( + user_rgb_8uc3_768_bald.copy(), + user_matting_8uc3_bald_768, + user_baldseg_8uc3_bald_768, + user_landmark_f1k2_bald_768) + + # cv2.imshow("user_hair_mask_blur_768", user_hair_mask_blur_768) + # cv2.waitKey() + + inv_M_user = cv2.invertAffineTransform(user_color_M) + bald_gene_8uc3_orisize = user_rgb_8uc3_orisize.copy() + + user_bald_gene_8uc3_orisize = cv2.warpAffine(bald_gene_8uc3_bald_768, inv_M_user, (orig_w, orig_h), + dst=bald_gene_8uc3_orisize.copy(), + borderMode=cv2.BORDER_TRANSPARENT) + + user_hair_mask_blur_orisize = cv2.warpAffine(user_hair_mask_blur_768, inv_M_user, (orig_w, orig_h), + flags=cv2.INTER_CUBIC) + + user_bald_res_8uc3_orisize = ( + (1 - user_hair_mask_blur_orisize[:, :, np.newaxis] / 255) * user_rgb_8uc3_orisize + \ + (user_hair_mask_blur_orisize[:, :, np.newaxis] / 255) * user_bald_gene_8uc3_orisize).astype( + np.uint8) + + user_baldseg_8uc3_768 = cv2.warpAffine(user_baldseg_8uc3_orisize, user_hairstyle_M, + (self.hair_size, self.hair_size)) # , borderValue=[255, 255, 255] + user_bald_8uc3_768 = cv2.warpAffine(user_bald_res_8uc3_orisize, user_hairstyle_M, + (self.hair_size, self.hair_size)) # , borderValue=[255, 255, 255] + + return user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, \ + user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M, user_matting_8uc3_bald_orisize + + def get_prepare_ref_data(self, ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize): + """ + input： + ref_rgb_8uc3_orisize: 参考图 原始尺寸 + ref_landmark_1k2_f_orisize: 参考图 1k 点 + + output: + + 图像尺寸基于： 人脸 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 + + """ + + # 得到512 小图 M 矩阵 + ref_hairstyle_M = self.get_hair_M(ref_landmark_1k2_f_orisize) + # ref_color_hair_M = self.get_color_hair_M(ref_landmark_1k2_f_orisize) + + # 1k 点 转换到 512 尺寸 + ref_landmark_f1k2_512 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, ref_hairstyle_M) + # ref_landmark_color_f1k2_512 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, ref_color_hair_M) + + # 参考图 头发毛躁 warpaffine + ratio = np.sqrt( + (ref_hairstyle_M[0][0] * ref_hairstyle_M[0][0]) + (ref_hairstyle_M[1][0] * ref_hairstyle_M[1][0])) + ref_hairstyle_M_tmp = ref_hairstyle_M / ratio + ref_rgb_8uc3_512 = cv2.warpAffine(ref_rgb_8uc3_orisize, ref_hairstyle_M_tmp, + (int(self.hair_size / ratio), int(self.hair_size / ratio)), + borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) + ref_rgb_8uc3_512 = cv2.resize(ref_rgb_8uc3_512, (self.hair_size, self.hair_size), fx=ratio, fy=ratio, + interpolation=cv2.INTER_AREA) + + # # 参考图 头发毛躁 warpaffine + # ratio_color = np.sqrt( + # (ref_color_hair_M[0][0] * ref_color_hair_M[0][0]) + (ref_color_hair_M[1][0] * ref_color_hair_M[1][0])) + # ref_color_hair_M_tmp = ref_color_hair_M / ratio_color + # ref_rgb_8uc3_768 = cv2.warpAffine(ref_rgb_8uc3_orisize, ref_color_hair_M_tmp, + # (int(self.bald_output_size / ratio_color), int(self.bald_output_size / ratio_color)), + # borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) + # ref_rgb_8uc3_768 = cv2.resize(ref_rgb_8uc3_768, (self.bald_output_size, self.bald_output_size), fx=ratio_color, fy=ratio_color, + # interpolation=cv2.INTER_AREA) + + # 参考图直接得到 头发 matting 小图 + _, ref_matte_pred_8uc1_512 = self.generator_matte.matte_inference(ref_rgb_8uc3_512, ref_landmark_f1k2_512) + torch.cuda.empty_cache() + + ref_matting_8uc1_512 = cv2.resize(ref_matte_pred_8uc1_512, + (ref_rgb_8uc3_512.shape[1], ref_rgb_8uc3_512.shape[0]), + interpolation=cv2.INTER_CUBIC) + + ref_matting_8uc3_512 = np.repeat(ref_matting_8uc1_512[:, :, np.newaxis], 3, axis=2) + + # 光头分割 + ref_baldseg_8uc3_orisize = self.generator_baldseg.forward(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize) + + + ref_baldseg_8uc3_512 = cv2.warpAffine(ref_baldseg_8uc3_orisize, ref_hairstyle_M, + (self.hair_size, self.hair_size), borderValue=[255, 255, 255]) + + return ref_rgb_8uc3_512, ref_matting_8uc3_512, ref_baldseg_8uc3_512, ref_landmark_f1k2_512 + + 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 + + def get_prepare_ref_768_data(self, ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize, ratio=1, long=False): + """ + input： + ref_rgb_8uc3_orisize: 参考图 原始尺寸 + ref_landmark_1k2_f_orisize: 参考图 1k 点 + + output: + + 图像尺寸基于： 人脸 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 + + """ + + # 得到512 小图 M 矩阵 + if ratio == 0: + ref_hairstyle_M = self.get_hair_M_boy_v1(ref_landmark_1k2_f_orisize) + elif ratio == 1: + ref_hairstyle_M = self.get_hair_M_girl_v1(ref_landmark_1k2_f_orisize) + elif ratio == 2: + ref_hairstyle_M = self.get_hair_M_girl_v2(ref_landmark_1k2_f_orisize, long=long) + else: + ref_hairstyle_M = self.get_hair_M_girl_v1(ref_landmark_1k2_f_orisize) + # ref_color_hair_M = self.get_color_hair_M(ref_landmark_1k2_f_orisize) + + # print("ref_hairstyle_M: ", ref_hairstyle_M) + + # 1k 点 转换到 512 尺寸 + ref_landmark_f1k2_768 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, ref_hairstyle_M) + + # 参考图 头发毛躁 warpaffine + ratio = np.sqrt( + (ref_hairstyle_M[0][0] * ref_hairstyle_M[0][0]) + (ref_hairstyle_M[1][0] * ref_hairstyle_M[1][0])) + ref_hairstyle_M_tmp = ref_hairstyle_M / ratio + ref_rgb_8uc3_768 = cv2.warpAffine(ref_rgb_8uc3_orisize, ref_hairstyle_M_tmp, + (int(self.bald_output_size / ratio), int(self.bald_output_size / ratio)), + borderMode=cv2.BORDER_CONSTANT) # , borderValue=[255, 255, 255] + ref_rgb_8uc3_768 = cv2.resize(ref_rgb_8uc3_768, (self.bald_output_size, self.bald_output_size), fx=ratio, + fy=ratio, + interpolation=cv2.INTER_AREA) + + # 参考图直接得到 头发 matting 小图 + ref_matte_fg_8uc3_768, ref_matte_pred_8uc1_768 = self.generator_matte.matte_inference(ref_rgb_8uc3_768, + ref_landmark_f1k2_768) + + torch.cuda.empty_cache() + + ref_matting_fg_8uc3_768 = cv2.resize(ref_matte_fg_8uc3_768, + (ref_rgb_8uc3_768.shape[1], ref_rgb_8uc3_768.shape[0]), + interpolation=cv2.INTER_CUBIC) + + ref_matting_8uc1_768 = cv2.resize(ref_matte_pred_8uc1_768, + (ref_rgb_8uc3_768.shape[1], ref_rgb_8uc3_768.shape[0]), + interpolation=cv2.INTER_CUBIC) + + ref_matting_8uc3_768 = np.repeat(ref_matting_8uc1_768[:, :, np.newaxis], 3, axis=2) + + ref_matting_8uc1_orisize = cv2.warpAffine(ref_matting_8uc1_768, cv2.invertAffineTransform(ref_hairstyle_M), + (ref_rgb_8uc3_orisize.shape[1], ref_rgb_8uc3_orisize.shape[0]), + flags=cv2.INTER_CUBIC) + + # 光头分割 + ref_baldseg_8uc3_orisize = self.generator_baldseg.forward(ref_rgb_8uc3_orisize, ref_matting_8uc1_orisize, + ref_landmark_1k2_f_orisize) + + ref_baldseg_8uc3_768 = cv2.warpAffine(ref_baldseg_8uc3_orisize, ref_hairstyle_M, + (self.bald_output_size, self.bald_output_size), + flags=cv2.INTER_NEAREST) # , borderValue=[255, 255, 255] + + return ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768 + + def get_prepare_ref_768_bald_data(self, ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize): + """ + input： + ref_rgb_8uc3_orisize: 参考图 原始尺寸 + ref_landmark_1k2_f_orisize: 参考图 1k 点 + + output: + + 图像尺寸基于： 人脸 512， 图像均为3通道 + ref_rgb_8uc3_768: 参考图， size 512， uint8 （0-255） + ref_matting_8uc3_768：参考图 matting alpha 值， uint8 （0-255） + ref_baldseg_8uc3_768: 参考图 光头分割， uint8 （0-255） + ref_landmark_f1k2_768： 参考图 关键点 1k*2 float32 + + """ + + # 得到512 小图 M 矩阵 + # ref_hairstyle_M = self.get_hair_M(ref_landmark_1k2_f_orisize) + ref_color_hair_M = self.get_color_hair_M(ref_landmark_1k2_f_orisize) + + # 1k 点 转换到 512 尺寸 + # ref_landmark_f1k2_512 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, ref_hairstyle_M) + ref_landmark_color_f1k2_768 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, ref_color_hair_M) + + # 参考图 头发毛躁 warpaffine + ratio = np.sqrt( + (ref_color_hair_M[0][0] * ref_color_hair_M[0][0]) + (ref_color_hair_M[1][0] * ref_color_hair_M[1][0])) + ref_color_hair_M_tmp = ref_color_hair_M / ratio + ref_rgb_8uc3_768 = cv2.warpAffine(ref_rgb_8uc3_orisize, ref_color_hair_M_tmp, + (int(self.color_output_size / ratio), int(self.color_output_size / ratio)), + borderMode=cv2.BORDER_CONSTANT) # , borderValue=[255, 255, 255] + ref_rgb_8uc3_768 = cv2.resize(ref_rgb_8uc3_768, (self.color_output_size, self.color_output_size), fx=ratio, + fy=ratio, + interpolation=cv2.INTER_AREA) + + # 参考图直接得到 头发 matting 小图 + _, ref_matte_pred_8uc1_768 = self.generator_matte.matte_inference(ref_rgb_8uc3_768, ref_landmark_color_f1k2_768) + torch.cuda.empty_cache() + + ref_matting_8uc1_768 = cv2.resize(ref_matte_pred_8uc1_768, + (ref_rgb_8uc3_768.shape[1], ref_rgb_8uc3_768.shape[0]), + interpolation=cv2.INTER_CUBIC) + + ref_matting_8uc3_768 = np.repeat(ref_matting_8uc1_768[:, :, np.newaxis], 3, axis=2) + + ref_matting_8uc1_orisize = cv2.warpAffine(ref_matting_8uc1_768, cv2.invertAffineTransform(ref_color_hair_M), + (ref_rgb_8uc3_orisize.shape[1], ref_rgb_8uc3_orisize.shape[0]), + flags=cv2.INTER_CUBIC) + + # 光头分割 + ref_baldseg_8uc3_orisize = self.generator_baldseg.forward(ref_rgb_8uc3_orisize, ref_matting_8uc1_orisize, + ref_landmark_1k2_f_orisize) + + ref_baldseg_8uc3_768 = cv2.warpAffine(ref_baldseg_8uc3_orisize, ref_color_hair_M, + (self.color_output_size, + self.color_output_size)) # , borderValue=[255, 255, 255] + + return ref_rgb_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_color_f1k2_768 + + def get_prepare_hair_color_data(self, user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize, ref_rgb_8uc3_orisize, + ref_landmark_1k2_f_orisize): + + """ + + input： + + user_rgb_8uc3_orisize: 用户图 原始尺寸 + user_landmark_1k2_f_orisize: 用户图 1k 点 + ref_rgb_8uc3_orisize: 参考图 原始尺寸 + ref_landmark_1k2_f_orisize: 参考图 1k 点 + + output: + + 图像尺寸基于： 人脸 768， 图像均为3通道 + user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768 + + user_rgb_8uc3_change_color_768: 用户图 align， size: 768 uint8 （0-255） + user_matting_8uc3_change_color_768: 用户图 matting， size: 768 uint8 （0-255） + + ref_rgb_8uc3_change_color_768: 参考图， align, size 512， uint8 （0-255） + ref_matting_8uc3_change_color_768: 参考图 matting alpha 值， uint8 （0-255） + + """ + + # 得到768 小图 M 矩阵 + user_hair_color_M = self.get_color_hair_M(user_landmark_1k2_f_orisize) + ref_hair_color_M = self.get_color_hair_M(ref_landmark_1k2_f_orisize) + + # 1k 点 转换到 768 尺寸 + user_landmark_f1k2_768 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_hair_color_M) + ref_landmark_f1k2_768 = landmark_processor.transform_points(ref_landmark_1k2_f_orisize, ref_hair_color_M) + + # 用户图 头发毛躁 warpaffine + usr_ratio = np.sqrt( + (user_hair_color_M[0][0] * user_hair_color_M[0][0]) + (user_hair_color_M[1][0] * user_hair_color_M[1][0])) + user_hair_color_M_tmp = user_hair_color_M / usr_ratio + user_rgb_8uc3_768 = cv2.warpAffine(user_rgb_8uc3_orisize, user_hair_color_M_tmp, + (int(self.color_output_size / usr_ratio), + int(self.color_output_size / usr_ratio))) + user_rgb_8uc3_768 = cv2.resize(user_rgb_8uc3_768, (self.color_output_size, self.color_output_size), + fx=usr_ratio, fy=usr_ratio, + interpolation=cv2.INTER_AREA) + + # 参考图 头发毛躁 warpaffine + ratio = np.sqrt( + (ref_hair_color_M[0][0] * ref_hair_color_M[0][0]) + (ref_hair_color_M[1][0] * ref_hair_color_M[1][0])) + ref_hair_color_M_tmp = ref_hair_color_M / ratio + ref_rgb_8uc3_768 = cv2.warpAffine(ref_rgb_8uc3_orisize, ref_hair_color_M_tmp, + (int(self.color_output_size / ratio), int(self.color_output_size / ratio))) + + ref_rgb_8uc3_768 = cv2.resize(ref_rgb_8uc3_768, (self.color_output_size, self.color_output_size), fx=ratio, + fy=ratio, + interpolation=cv2.INTER_AREA) + + # 参考图直接得到 头发 matting 小图 + _, ref_matte_pred_8uc1_768 = self.generator_matte.matte_inference(ref_rgb_8uc3_768, ref_landmark_f1k2_768) + + torch.cuda.empty_cache() + + ref_matting_8uc3_768 = np.repeat(ref_matte_pred_8uc1_768[:, :, np.newaxis], 3, axis=2) + + # 用户图得到 头发 matting 小图 + _, user_matting_8uc1_bald_768 = self.generator_matte.matte_inference(user_rgb_8uc3_768, user_landmark_f1k2_768) + + user_matting_8uc3_768 = np.repeat(user_matting_8uc1_bald_768[:, :, np.newaxis], 3, axis=2) + + return user_rgb_8uc3_768, user_matting_8uc3_768, ref_rgb_8uc3_768, ref_matting_8uc3_768 + + def get_prepare_hair_color_user_data(self, user_rgb_8uc3_orisize, user_landmark_1k2_f_orisize): + + """ + + input： + + user_rgb_8uc3_orisize: 用户图 原始尺寸 + user_landmark_1k2_f_orisize: 用户图 1k 点 + + output: + + 图像尺寸基于： 人脸 768， 图像均为3通道 + user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768 + + user_rgb_8uc3_change_color_768: 用户图 align， size: 768 uint8 （0-255） + user_matting_8uc3_change_color_768: 用户图 matting， size: 768 uint8 （0-255） + + """ + + # 得到768 小图 M 矩阵 + user_hair_color_M = self.get_color_hair_M(user_landmark_1k2_f_orisize) + + # 1k 点 转换到 768 尺寸 + user_landmark_f1k2_768 = landmark_processor.transform_points(user_landmark_1k2_f_orisize, user_hair_color_M) + + # # 用户图 头发毛躁 warpaffine + # usr_ratio = np.sqrt( + # (user_hair_color_M[0][0] * user_hair_color_M[0][0]) + (user_hair_color_M[1][0] * user_hair_color_M[1][0])) + # user_hair_color_M_tmp = user_hair_color_M / usr_ratio + # user_rgb_8uc3_768 = cv2.warpAffine(user_rgb_8uc3_orisize, user_hair_color_M_tmp, + # (int(self.color_output_size / usr_ratio), int(self.color_output_size / usr_ratio))) + # user_rgb_8uc3_768 = cv2.resize(user_rgb_8uc3_768, (self.color_output_size, self.color_output_size), fx=usr_ratio, fy=usr_ratio, + # interpolation=cv2.INTER_AREA) + + user_rgb_8uc3_768 = cv2.warpAffine(user_rgb_8uc3_orisize, user_hair_color_M, + (self.color_output_size, self.color_output_size)) + # + # usr_ratio = np.sqrt( + # (user_hair_color_M[0][0] * user_hair_color_M[0][0]) + (user_hair_color_M[1][0] * user_hair_color_M[1][0])) + # user_hair_color_M_tmp = user_hair_color_M / usr_ratio + # user_rgb_8uc3_bald_768 = cv2.warpAffine(user_rgb_8uc3_orisize, user_hair_color_M_tmp, + # (int(self.color_output_size / usr_ratio), int(self.color_output_size / usr_ratio)), + # borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255]) + # user_rgb_8uc3_bald_768 = cv2.resize(user_rgb_8uc3_bald_768, (self.color_output_size, self.color_output_size), fx=usr_ratio, fy=usr_ratio, + # interpolation=cv2.INTER_AREA) + + # 用户图得到 头发 matting 小图 + _, user_matting_8uc1_bald_768 = self.generator_matte.matte_inference(user_rgb_8uc3_768, user_landmark_f1k2_768) + + user_matting_8uc3_768 = np.repeat(user_matting_8uc1_bald_768[:, :, np.newaxis], 3, axis=2) + + return user_rgb_8uc3_768, user_matting_8uc3_768, user_hair_color_M + + def get_matte_img(self, img, landmark1k): + + matte_fg, matte_img = self.generator_matte.matte_inference(img, landmark1k) + + return matte_fg, matte_img + + def get_max_countour(self, contours): + index_contour = 0 + max_num = 0 + for i, contour in enumerate(contours): + if len(contour) > max_num: + max_num = len(contour) + index_contour = i + return index_contour + + def judge_hair_pos(self, user_matting_8uc3_bald_768, user_baldseg_8uc3_bald_768, retData): + + # cv2.imshow("user_matting_8uc3_bald_768", user_matting_8uc3_bald_768) + # cv2.imshow("user_baldseg_8uc3_bald_768", user_baldseg_8uc3_bald_768) + # cv2.waitKey() + + cloth_mask = np.zeros_like(user_baldseg_8uc3_bald_768, dtype=np.uint8) + cloth_pos = (user_baldseg_8uc3_bald_768[:, :, 0] < 10) & (user_baldseg_8uc3_bald_768[:, :, 1] < 10) & ( + user_baldseg_8uc3_bald_768[:, :, 2] > 250) + cloth_mask[cloth_pos] = 255 + + cloth_withhair_mask = (cloth_mask * user_matting_8uc3_bald_768.astype(np.float32) / 255).astype(np.uint8) + cloth_withhair_count = len(cloth_withhair_mask[cloth_withhair_mask[:, :, 0] > 0]) + # cv2.imshow("cloth_mask", cloth_mask) + # cv2.imshow("cloth_withhair_mask", cloth_withhair_mask) + + # print("cloth_withhair_count: ", cloth_withhair_count / (cloth_withhair_mask.shape[0] * cloth_withhair_mask.shape[1])) + + # cv2.waitKey() + + retData["cloth_withhair_count"] = str(cloth_withhair_count) + + def get_fusion_res(self, user_baldseg_8uc3_512, user_bald_8uc3_512, hair_gene_8uc3_512, user_landmark_f1k2_512): + """ + input： + + 图像尺寸基于： 人脸 512， 图像均为3通道 + + user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + hair_gene_8uc3_512: 用户图 换发型 贴合 光头， uint8 （0-255） + user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 + + output： + + hair_gene_fusion_8uc3_512: 生成 发型图， uint8 （0-255） + + """ + # "user_baldseg_8uc3_512, user_bald_8uc3_512, hair_gene_8uc3_512" + # cv2.imshow("user_baldseg_8uc3_512", user_baldseg_8uc3_512) + # cv2.imshow("user_bald_8uc3_512", user_bald_8uc3_512) + # cv2.imshow("hair_gene_8uc3_512", hair_gene_8uc3_512) + + # 换发型后的 matting 图 + hair_gene_matte_fg_8uc3_512, hair_gene_matte_8uc1_512 = self.get_matte_img(hair_gene_8uc3_512, + user_landmark_f1k2_512) + hair_gene_matte_32fc1_512 = hair_gene_matte_8uc1_512.astype(np.float32) / 255 + hair_gene_matte_32fc3_512 = np.repeat(hair_gene_matte_32fc1_512[:, :, np.newaxis], 3, axis=2) + + hair_bg = (user_bald_8uc3_512 * (1 - hair_gene_matte_32fc3_512)).astype(np.uint8) + + hair_bg_face = np.zeros_like(user_bald_8uc3_512) + + face_index = (user_baldseg_8uc3_512[:, :, 1] == 255) & (user_baldseg_8uc3_512[:, :, 0] == 0) & ( + user_baldseg_8uc3_512[:, :, 2] == 0) + + hair_bg_face[face_index] = user_bald_8uc3_512[face_index] + + hair_face_bg = np.zeros_like(user_bald_8uc3_512) + hair_face_bg[face_index] = hair_gene_8uc3_512[face_index] + + hair_face_mask = np.zeros_like(user_bald_8uc3_512, dtype=np.uint8) + hair_face_mask[face_index] = 255 + kernel_size = 7 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + hair_face_mask = cv2.erode(hair_face_mask, kernel, iterations=1).astype(np.uint8) + hair_face_mask_blur = cv2.blur(hair_face_mask, (5, 5)) + + # cv2.imshow("hair_face_bg ", hair_face_bg) + # cv2.imshow("hair_face_mask ", (hair_face_mask).astype(np.uint8)) + # cv2.imshow("hair_face_mask_blur ", hair_face_mask_blur) + # cv2.imshow("user_bald_8uc3_512 ", user_bald_8uc3_512) + # cv2.imshow("hair_gene_matte_fg_8uc3_512 ", hair_gene_matte_fg_8uc3_512) + + # paste_hair = np.zeros_like(user_bald_8uc3_512) + hair_gene_fusion_8uc3_512 = (user_bald_8uc3_512 * ( + 1 - hair_gene_matte_32fc3_512) + hair_gene_matte_fg_8uc3_512 * hair_gene_matte_32fc3_512).astype( + np.uint8) + # cv2.imshow("hair_gene_fusion_8uc3_512 nouse mask ", hair_gene_fusion_8uc3_512) + + hair_gene_fusion_8uc3_512 = (hair_gene_fusion_8uc3_512 * ( + 1 - hair_face_mask_blur.astype(np.float32) / 255) + hair_face_bg * hair_face_mask_blur.astype( + np.float32) / 255).astype(np.uint8) + # hair_gene_fusion_8uc3_512[face_index] = hair_face_bg[face_index] + + # cv2.waitKey() + + return hair_gene_fusion_8uc3_512 + + def get_fusion_res_hairpaste(self, user_bald_8uc3_orisize, hair_gene_8uc3_768, + user_landmark_f1k2_768, + user_hairstyle_M): + """ + input： + + 图像尺寸基于： 人脸 512， 图像均为3通道 + + user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + hair_gene_8uc3_512: 用户图 换发型 贴合 光头， uint8 （0-255） + user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 + + output： + + hair_gene_fusion_8uc3_512: 生成 发型图， uint8 （0-255） + + """ + hair_gene_matte_fg_8uc3_768, hair_gene_matte_8uc1_768 = self.get_matte_img(hair_gene_8uc3_768, + user_landmark_f1k2_768) + + # cv2.imshow("hair_gene_8uc3_768", hair_gene_8uc3_768) + # cv2.imshow("hair_gene_matte_8uc1_768", hair_gene_matte_8uc1_768) + # cv2.imshow("hair_gene_matte_fg_8uc3_768", hair_gene_matte_fg_8uc3_768) + # cv2.waitKey() + + # hair_gene_8uc3_orisize = cv2.warpAffine(hair_gene_8uc3_768, cv2.invertAffineTransform(user_hairstyle_M), + # (user_bald_8uc3_orisize.shape[1], + # user_bald_8uc3_orisize.shape[0]), + # dst=user_bald_8uc3_orisize.copy(), + # borderMode=cv2.BORDER_TRANSPARENT) + + + hair_gene_matte_8uc1_orisize = cv2.warpAffine(hair_gene_matte_8uc1_768, + cv2.invertAffineTransform(user_hairstyle_M), + (user_bald_8uc3_orisize.shape[1], + user_bald_8uc3_orisize.shape[0]), + flags=cv2.INTER_CUBIC) + + + hair_gene_matte_fg_8uc3_orisize = cv2.warpAffine(hair_gene_matte_fg_8uc3_768, + cv2.invertAffineTransform(user_hairstyle_M), + (user_bald_8uc3_orisize.shape[1], + user_bald_8uc3_orisize.shape[0]), + flags=cv2.INTER_CUBIC) + + hair_gene_matte_8uc3_orisize = np.repeat(hair_gene_matte_8uc1_orisize[:, :, np.newaxis], 3, axis=2) + + hair_gene_matte_32fc1_orisize = hair_gene_matte_8uc1_orisize.astype(np.float32) / 255 + hair_gene_matte_32fc3_orisize = np.repeat(hair_gene_matte_32fc1_orisize[:, :, np.newaxis], 3, axis=2) + + hair_gene_fusion_8uc3_orisize = (user_bald_8uc3_orisize * ( + 1 - hair_gene_matte_32fc3_orisize) + hair_gene_matte_fg_8uc3_orisize * hair_gene_matte_32fc3_orisize).astype( + np.uint8) + + return hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize + + def get_fusion_res_hairblur(self, user_baldseg_8uc3_orisize, user_bald_8uc3_orisize, hair_gene_8uc3_768, + user_landmark_f1k2_768, user_hairstyle_M, ref_hairstyle_dir): + """ + input： + + 图像尺寸基于： 人脸 512， 图像均为3通道 + + user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + hair_gene_8uc3_512: 用户图 换发型 贴合 光头， uint8 （0-255） + user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 + + output： + + hair_gene_fusion_8uc3_512: 生成 发型图， uint8 （0-255） + + """ + # cv2.imshow("user_baldseg_8uc3_512", user_baldseg_8uc3_512) + # cv2.imshow("user_bald_8uc3_512", user_bald_8uc3_512) + # cv2.imshow("hair_gene_8uc3_512", hair_gene_8uc3_512) + + # 换发型后的 matting 图 + + hair_gene_matte_fg_8uc3_768, hair_gene_matte_8uc1_768 = self.get_matte_img(hair_gene_8uc3_768, + user_landmark_f1k2_768) + # cv2.imshow("hair_gene_matte_8uc1_768 before", hair_gene_matte_8uc1_768) + # cv2.imshow("hair_gene_matte_fg_8uc3_768 before", hair_gene_matte_fg_8uc3_768) + + hair_gene_8uc3_orisize = cv2.warpAffine(hair_gene_8uc3_768, cv2.invertAffineTransform(user_hairstyle_M), + (user_bald_8uc3_orisize.shape[1], + user_bald_8uc3_orisize.shape[0]), + dst=user_bald_8uc3_orisize.copy(), + borderMode=cv2.BORDER_TRANSPARENT) + + default_hair_tail_mask = cv2.imread(os.path.join(ref_hairstyle_dir, "default_hair_tail_mask.png")) + + hair_gene_matte_8uc1_768 = ( + hair_gene_matte_8uc1_768 * default_hair_tail_mask[:, :, 0].astype(np.float32) / 255).astype( + np.uint8) + # cv2.imshow("hair_gene_matte_8uc1_768 after", hair_gene_matte_8uc1_768) + + hair_gene_matte_8uc1_orisize = cv2.warpAffine(hair_gene_matte_8uc1_768, + cv2.invertAffineTransform(user_hairstyle_M), + (user_bald_8uc3_orisize.shape[1], + user_bald_8uc3_orisize.shape[0]), + flags=cv2.INTER_CUBIC) + + hair_gene_matte_fg_8uc3_768 = ( + hair_gene_matte_fg_8uc3_768 * default_hair_tail_mask.astype(np.float32) / 255).astype(np.uint8) + # cv2.imshow("hair_gene_matte_fg_8uc3_768 after", hair_gene_matte_fg_8uc3_768) + # cv2.waitKey() + + hair_gene_matte_fg_8uc3_orisize = cv2.warpAffine(hair_gene_matte_fg_8uc3_768, + cv2.invertAffineTransform(user_hairstyle_M), + (user_bald_8uc3_orisize.shape[1], + user_bald_8uc3_orisize.shape[0]), + flags=cv2.INTER_CUBIC) + + hair_gene_matte_32fc1_orisize = hair_gene_matte_8uc1_orisize.astype(np.float32) / 255 + hair_gene_matte_32fc3_orisize = np.repeat(hair_gene_matte_32fc1_orisize[:, :, np.newaxis], 3, axis=2) + + # cv2.imshow("hair_gene_8uc3_512", hair_gene_8uc3_512) + # cv2.imshow("hair_gene_matte_8uc1_512", hair_gene_matte_8uc1_512) + # cv2.imshow("hair_gene_matte_32fc3_512", (hair_gene_matte_32fc3_512*255).astype(np.uint8)) + # cv2.waitKey() + + hair_bg_face = np.zeros_like(user_bald_8uc3_orisize) + + face_index = (user_baldseg_8uc3_orisize[:, :, 1] == 255) & (user_baldseg_8uc3_orisize[:, :, 0] == 0) & ( + user_baldseg_8uc3_orisize[:, :, 2] == 0) + + hair_bg_face[face_index] = user_bald_8uc3_orisize[face_index] + + hair_face_bg = np.zeros_like(user_bald_8uc3_orisize) + hair_face_bg[face_index] = hair_gene_8uc3_orisize[face_index] + + hair_face_mask = np.zeros_like(user_bald_8uc3_orisize, dtype=np.uint8) + hair_face_mask[face_index] = 255 + + hair_face_rec_mask = (((hair_gene_matte_8uc1_orisize > 0) & (hair_face_mask[:, :, 0] > 0)).astype( + np.float32) * 255).astype(np.uint8) + + # cv2.imshow("hair_face_rec_mask no resize", hair_face_rec_mask) + + hair_face_rec_mask = cv2.resize(hair_face_rec_mask, (768, 768), interpolation=cv2.INTER_CUBIC) + # cv2.imshow("hair_face_rec_mask ", np.repeat(hair_face_rec_mask[:, :, np.newaxis], 3, axis=2) ) + + kernel_size = 11 + kernel_dilate = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + hair_face_rec_mask_dilate = cv2.erode(hair_face_rec_mask, kernel_dilate, iterations=1) + + hair_face_rec_mask_blur = cv2.blur(hair_face_rec_mask_dilate, (31, 31)) + + hair_face_rec_mask_blur = cv2.resize(hair_face_rec_mask_blur, (user_bald_8uc3_orisize.shape[1], + user_bald_8uc3_orisize.shape[0]), + interpolation=cv2.INTER_CUBIC) + + # cv2.imshow("hair_face_rec_mask ", (hair_face_rec_mask * 255).astype(np.uint8)) + # cv2.imshow("hair_face_rec_mask_blur ", (np.repeat(hair_face_rec_mask_blur[:, :, np.newaxis], 3, axis=2)).astype(np.uint8)) + # cv2.waitKey() + + kernel_size_2 = 7 + kernel_erode_2 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size_2, kernel_size_2)) + + hair_face_mask_erode_2 = cv2.erode(hair_face_mask, kernel_erode_2, iterations=1).astype(np.uint8) + + hair_face_mask_blur = cv2.blur(hair_face_mask_erode_2, (5, 5)) + hair_face_mask_blur = cv2.resize(hair_face_mask_blur, (user_bald_8uc3_orisize.shape[1], + user_bald_8uc3_orisize.shape[0]), + interpolation=cv2.INTER_CUBIC) + + # cv2.imshow("hair_face_bg ", hair_face_bg) + # cv2.imshow("hair_face_mask ", (hair_face_mask).astype(np.uint8)) + # cv2.imshow("hair_face_mask_rolling ", (hair_face_mask_rolling).astype(np.uint8)) + # cv2.imshow("hair_face_mask_use_blur ", (hair_face_mask_use_blur).astype(np.uint8)) + # cv2.imshow("hair_face_mask_blur ", hair_face_mask_blur) + # cv2.imshow("hair_matte ", (hair_gene_matte_32fc3_512 * 255).astype(np.uint8)) + # cv2.imshow("user_bald_8uc3_512 ", user_bald_8uc3_512) + # cv2.imshow("hair_gene_matte_fg_8uc3_512 ", hair_gene_matte_fg_8uc3_512) + + # cv2.imshow("hair_matte blur", (hair_gene_matte_32fc3_512_blur * 255).astype(np.uint8)) + + # cv2.imshow("user_bald_8uc3_fg ", user_bald_8uc3_fg) + + # paste_hair = np.zeros_like(user_bald_8uc3_512) + hair_gene_fusion_8uc3_orisize = (user_bald_8uc3_orisize * ( + 1 - hair_gene_matte_32fc3_orisize) + hair_gene_matte_fg_8uc3_orisize * hair_gene_matte_32fc3_orisize).astype( + np.uint8) + # cv2.imshow("hair_gene_fusion_8uc3_512 nouse mask ", hair_gene_fusion_8uc3_512) + + hair_gene_fusion_8uc3_orisize_gene = hair_gene_fusion_8uc3_orisize.copy() + hair_gene_fusion_8uc3_orisize_gene = (hair_gene_fusion_8uc3_orisize_gene * ( + 1 - hair_face_mask_blur.astype(np.float32) / 255) + hair_face_bg * hair_face_mask_blur.astype( + np.float32) / 255).astype(np.uint8) + # cv2.imshow("hair_gene_fusion_8uc3_512_gene", hair_gene_fusion_8uc3_512_gene) + + hair_gene_fusion_8uc3_512 = (hair_gene_fusion_8uc3_orisize_gene * ( + hair_face_rec_mask_blur[:, :, np.newaxis].astype(np.float32) / 255) + + hair_gene_fusion_8uc3_orisize * ( + 1 - hair_face_rec_mask_blur[:, :, np.newaxis].astype( + np.float32) / 255)).astype(np.uint8) + # cv2.imshow("hair_gene_fusion_8uc3_512 use mask ", hair_gene_fusion_8uc3_512) + + # cv2.waitKey() + + return hair_gene_fusion_8uc3_512, hair_gene_8uc3_orisize + + def get_fusion_res_onlyhair(self, user_baldseg_8uc3_512, user_bald_8uc3_512, hair_gene_8uc3_512, + user_landmark_f1k2_512): + """ + input： + + 图像尺寸基于： 人脸 512， 图像均为3通道 + + user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + hair_gene_8uc3_512: 用户图 换发型 贴合 光头， uint8 （0-255） + user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 + + output： + + hair_gene_fusion_8uc3_512: 生成 发型图， uint8 （0-255） + + """ + # "user_baldseg_8uc3_512, user_bald_8uc3_512, hair_gene_8uc3_512" + cv2.imshow("user_baldseg_8uc3_512", user_baldseg_8uc3_512) + # cv2.imshow("user_bald_8uc3_512", user_bald_8uc3_512) + cv2.imshow("hair_gene_8uc3_512", hair_gene_8uc3_512) + + # 换发型后的 matting 图 + hair_gene_matte_fg_8uc3_512, hair_gene_matte_8uc1_512 = self.get_matte_img(hair_gene_8uc3_512, + user_landmark_f1k2_512) + hair_gene_matte_32fc1_512 = hair_gene_matte_8uc1_512.astype(np.float32) / 255 + hair_gene_matte_32fc3_512 = np.repeat(hair_gene_matte_32fc1_512[:, :, np.newaxis], 3, axis=2) + + hair_bg = (user_bald_8uc3_512 * (1 - hair_gene_matte_32fc3_512)).astype(np.uint8) + + hair_bg_face = np.zeros_like(user_bald_8uc3_512) + + face_index = (user_baldseg_8uc3_512[:, :, 1] == 255) & (user_baldseg_8uc3_512[:, :, 0] == 0) & ( + user_baldseg_8uc3_512[:, :, 2] == 0) + + hair_bg_face[face_index] = user_bald_8uc3_512[face_index] + + hair_face_bg = np.zeros_like(user_bald_8uc3_512) + hair_face_bg[face_index] = hair_gene_8uc3_512[face_index] + + hair_face_mask = np.zeros_like(user_bald_8uc3_512, dtype=np.uint8) + hair_face_mask[face_index] = 255 + kernel_size_1 = 15 + kernel_size_2 = 7 + kernel_erode_1 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size_1, kernel_size_1)) + kernel_erode_2 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size_2, kernel_size_2)) + + hair_face_mask_erode_1 = cv2.erode(hair_face_mask, kernel_erode_1, iterations=1).astype(np.uint8) + hair_face_mask_erode_2 = cv2.erode(hair_face_mask, kernel_erode_2, iterations=1).astype(np.uint8) + # hair_face_mask_dilate = cv2.dilate(hair_face_mask, kernel, iterations=1).astype(np.uint8) + + hair_face_mask_rolling = hair_face_mask_erode_2 - hair_face_mask_erode_1 + hair_face_mask_use = 255 - hair_face_mask_rolling + + hair_face_mask_blur = cv2.blur(hair_face_mask_erode_2, (5, 5)) + hair_face_mask_use_blur = cv2.blur(hair_face_mask_use, (11, 11)) + + cv2.imshow("hair_face_bg ", hair_face_bg) + cv2.imshow("hair_face_mask ", (hair_face_mask).astype(np.uint8)) + cv2.imshow("hair_face_mask_rolling ", (hair_face_mask_rolling).astype(np.uint8)) + cv2.imshow("hair_face_mask_use_blur ", (hair_face_mask_use_blur).astype(np.uint8)) + cv2.imshow("hair_face_mask_blur ", hair_face_mask_blur) + cv2.imshow("hair_matte ", (hair_gene_matte_32fc3_512 * 255).astype(np.uint8)) + cv2.imshow("user_bald_8uc3_512 ", user_bald_8uc3_512) + cv2.imshow("hair_gene_matte_fg_8uc3_512 ", hair_gene_matte_fg_8uc3_512) + + hair_gene_matte_32fc3_512_blur = cv2.blur(hair_gene_matte_32fc3_512, (3, 3)) + cv2.imshow("hair_matte blur", (hair_gene_matte_32fc3_512_blur * 255).astype(np.uint8)) + + user_bald_8uc3_white = np.full_like(user_bald_8uc3_512, (187, 202, 240), dtype=(np.uint8)) + user_bald_8uc3_fg = (user_bald_8uc3_white * ( + 1 - hair_gene_matte_32fc3_512) + hair_gene_matte_fg_8uc3_512 * hair_gene_matte_32fc3_512).astype( + np.uint8) + cv2.imshow("user_bald_8uc3_fg ", user_bald_8uc3_fg) + + # paste_hair = np.zeros_like(user_bald_8uc3_512) + hair_gene_fusion_8uc3_512 = (user_bald_8uc3_512 * ( + 1 - hair_gene_matte_32fc3_512) + hair_gene_matte_fg_8uc3_512 * hair_gene_matte_32fc3_512).astype( + np.uint8) + cv2.imshow("hair_gene_fusion_8uc3_512 nouse mask ", hair_gene_fusion_8uc3_512) + + hair_gene_fusion_8uc3_512_gene = hair_gene_fusion_8uc3_512.copy() + hair_gene_fusion_8uc3_512_gene = (hair_gene_fusion_8uc3_512_gene * ( + 1 - hair_face_mask_blur.astype(np.float32) / 255) + hair_face_bg * hair_face_mask_blur.astype( + np.float32) / 255).astype(np.uint8) + cv2.imshow("hair_gene_fusion_8uc3_512_gene", hair_gene_fusion_8uc3_512_gene) + + hair_gene_fusion_8uc3_512 = (hair_gene_fusion_8uc3_512_gene * (1 - hair_face_mask_use_blur.astype( + np.float32) / 255) + hair_gene_fusion_8uc3_512 * hair_face_mask_use_blur.astype(np.float32) / 255).astype( + np.uint8) + cv2.imshow("hair_gene_fusion_8uc3_512 use mask ", hair_gene_fusion_8uc3_512) + + cv2.waitKey() + + return hair_gene_fusion_8uc3_512 + + def get_fusion_res_forehead(self, user_baldseg_8uc3_512, user_bald_8uc3_512, hair_gene_8uc3_512, + user_landmark_f1k2_512): + """ + input： + + 图像尺寸基于： 人脸 512， 图像均为3通道 + + user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + hair_gene_8uc3_512: 用户图 换发型 贴合 光头， uint8 （0-255） + user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 + + output： + + hair_gene_fusion_8uc3_512: 生成 发型图， uint8 （0-255） + + """ + + # 换发型后的 matting 图 + hair_gene_matte_fg_8uc3_512, hair_gene_matte_8uc1_512 = self.get_matte_img(hair_gene_8uc3_512, + user_landmark_f1k2_512) + hair_gene_matte_32fc1_512 = hair_gene_matte_8uc1_512.astype(np.float32) / 255 + hair_gene_matte_32fc3_512 = np.repeat(hair_gene_matte_32fc1_512[:, :, np.newaxis], 3, axis=2) + + hair_bg = (user_bald_8uc3_512 * (1 - hair_gene_matte_32fc3_512)).astype(np.uint8) + + hair_bg_face = np.zeros_like(user_bald_8uc3_512) + + face_index = (user_baldseg_8uc3_512[:, :, 1] == 255) & (user_baldseg_8uc3_512[:, :, 0] == 0) & ( + user_baldseg_8uc3_512[:, :, 2] == 0) + + hair_bg_face[face_index] = user_bald_8uc3_512[face_index] + + hair_face_bg = np.zeros_like(user_bald_8uc3_512) + hair_face_bg[face_index] = hair_gene_8uc3_512[face_index] + + hair_face_mask = np.zeros_like(user_bald_8uc3_512, dtype=np.uint8) + hair_face_mask[face_index] = 255 + + user_landmark_f137_512 = landmark_processor.pts_1k_to_137(user_landmark_f1k2_512) + pts_leye_up = user_landmark_f137_512[89:96, :] + pts_reye_up = user_landmark_f137_512[106:113, :] + pts_leye_up_low = pts_leye_up.min(axis=0)[1] + pts_reye_up_low = pts_reye_up.min(axis=0)[1] + + pts_eye_up_low = min(pts_leye_up_low, pts_reye_up_low) + hair_face_mask[int(pts_eye_up_low):, :, :] = 0 + + kernel_size = 7 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + hair_face_mask = cv2.erode(hair_face_mask, kernel, iterations=1).astype(np.uint8) + hair_face_mask_blur = cv2.blur(hair_face_mask, (5, 5)) + + # cv2.imshow("hair_face_bg ", hair_face_bg) + # cv2.imshow("hair_face_mask ", (hair_face_mask).astype(np.uint8)) + # cv2.imshow("hair_face_mask_blur ", hair_face_mask_blur) + # cv2.imshow("user_bald_8uc3_512 ", user_bald_8uc3_512) + # cv2.imshow("hair_gene_matte_fg_8uc3_512 ", hair_gene_matte_fg_8uc3_512) + + # paste_hair = np.zeros_like(user_bald_8uc3_512) + hair_gene_fusion_8uc3_512 = (user_bald_8uc3_512 * ( + 1 - hair_gene_matte_32fc3_512) + hair_gene_matte_fg_8uc3_512 * hair_gene_matte_32fc3_512).astype( + np.uint8) + # cv2.imshow("hair_gene_fusion_8uc3_512 nouse mask ", hair_gene_fusion_8uc3_512) + + hair_gene_fusion_8uc3_512 = (hair_gene_fusion_8uc3_512 * ( + 1 - hair_face_mask_blur.astype(np.float32) / 255) + hair_face_bg * hair_face_mask_blur.astype( + np.float32) / 255).astype(np.uint8) + # hair_gene_fusion_8uc3_512[face_index] = hair_face_bg[face_index] + # cv2.imshow("hair_gene_fusion_8uc3_512 mask ", hair_gene_fusion_8uc3_512) + + cv2.waitKey() + + return hair_gene_fusion_8uc3_512 + +class PersonProcessor_yolov5(object): + def __init__(self, gpu_id=0): + self.img_size = 640 + self.gpu_id = gpu_id + self.device = torch.device("cuda:%d" % gpu_id) + self.confidence_threshold = 0.5 + self.nms_threshold = 0.8 + self.model = torch.load(os.path.join(modelRoot, "yolov5l.pt"), map_location=self.device)['model'].float() # load to FP32 + # torch.save(torch.load(weights, map_location=device), weights) # update model if SourceChangeWarning + # model.fuse() + + # 修复 Upsample 层缺失属性的问题 + for m in self.model.modules(): + if isinstance(m, torch.nn.Upsample): + if not hasattr(m, 'recompute_scale_factor'): + m.recompute_scale_factor = None # 或 False，根据需求 + + self.model.to(self.device).eval() + + print("init yolov5 model succeed!!!") + + + def forward(self, image): + # Padded resize + img = self.letterbox(image, new_shape=self.img_size)[0] + # Convert + img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416 + img = np.ascontiguousarray(img) + ### image is float32 + img = torch.from_numpy(img).to(self.device) + img = img.float() # uint8 to fp16/32 + img /= 255.0 + if img.ndimension() == 3: + img = img.unsqueeze(0) + + + # Inference + pred = self.model(img)[0] + + # Apply NMS + pred = self.non_max_suppression(pred, self.confidence_threshold, self.nms_threshold, classes=None, agnostic=False) + + boxes = [] + scores = [] + classes = [] + # Process detections + det = pred[0] + if det is not None and len(det): + # Rescale boxes from img_size to im0 size + det[:, :4] = self.scale_coords(img.shape[2:], det[:, :4], image.shape).round() + det = det.detach().cpu().numpy() + for i in range(len(det)): + if det[i, 5] == 0: + boxes.append([det[i, 0], det[i, 1], det[i, 2], det[i, 3]]) + scores.append(det[i, 4]) + classes.append(det[i, 5]) + # print("det: ", det) + + boxes = np.array(boxes).astype(np.int32) + return { + 'boxes': boxes, + 'scores': scores, + 'classes': classes, + } + + def process(self, in_frame): + """ + 返回的人体检测框的list，具体rect的格式是[x1, y1, x2, y2] + """ + res = self.forward(in_frame) + if len(res['boxes']) == 0: return None + filter_boxes = [] + for ix, box_score in enumerate(res['scores']): + box_ = res['boxes'][ix] + box_w = box_[2] - box_[0] + box_h = box_[3] - box_[1] + max_box_len = max(box_w, box_h) + if box_score > 0.5 and max_box_len > 150: + filter_boxes.append([box_, box_w * box_h, box_h / in_frame.shape[0]]) + + filter_boxes.sort(key=lambda x: x[1], reverse=True) + filter_boxes = list(filter(lambda x: x[2] > 0.2, filter_boxes)) + if len(filter_boxes) == 0: return None + filter_boxes = [item[0] for item in filter_boxes] + return filter_boxes + + def letterbox(self, img, new_shape=(416, 416), color=(114, 114, 114), auto=False, scaleFill=False, scaleup=True): + # Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232 + shape = img.shape[:2] # current shape [height, width] + if isinstance(new_shape, int): + new_shape = (new_shape, new_shape) + + # Scale ratio (new / old) + r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) + if not scaleup: # only scale down, do not scale up (for better test mAP) + r = min(r, 1.0) + + # Compute padding + ratio = r, r # width, height ratios + new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) + dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding + if auto: # minimum rectangle + dw, dh = np.mod(dw, 64), np.mod(dh, 64) # wh padding + elif scaleFill: # stretch + dw, dh = 0.0, 0.0 + new_unpad = new_shape + ratio = new_shape[0] / shape[1], new_shape[1] / shape[0] # width, height ratios + + dw /= 2 # divide padding into 2 sides + dh /= 2 + + if shape[::-1] != new_unpad: # resize + img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) + top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) + left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) + img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border + return img, ratio, (dw, dh) + + def scale_coords(self, img1_shape, coords, img0_shape, ratio_pad=None): + # Rescale coords (xyxy) from img1_shape to img0_shape + if ratio_pad is None: # calculate from img0_shape + gain = max(img1_shape) / max(img0_shape) # gain = old / new + pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding + else: + gain = ratio_pad[0][0] + pad = ratio_pad[1] + + coords[:, [0, 2]] -= pad[0] # x padding + coords[:, [1, 3]] -= pad[1] # y padding + coords[:, :4] /= gain + self.clip_coords(coords, img0_shape) + return coords + + def clip_coords(self, boxes, img_shape): + # Clip bounding xyxy bounding boxes to image shape (height, width) + boxes[:, 0].clamp_(0, img_shape[1]) # x1 + boxes[:, 1].clamp_(0, img_shape[0]) # y1 + boxes[:, 2].clamp_(0, img_shape[1]) # x2 + boxes[:, 3].clamp_(0, img_shape[0]) # y2 + + def non_max_suppression(self, prediction, conf_thres=0.1, iou_thres=0.6, merge=False, classes=None, agnostic=False): + """Performs Non-Maximum Suppression (NMS) on inference results + + Returns: + detections with shape: nx6 (x1, y1, x2, y2, conf, cls) + """ + if prediction.dtype is torch.float16: + prediction = prediction.float() # to FP32 + + nc = prediction[0].shape[1] - 5 # number of classes + xc = prediction[..., 4] > conf_thres # candidates + + # Settings + min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height + max_det = 300 # maximum number of detections per image + time_limit = 10.0 # seconds to quit after + redundant = True # require redundant detections + multi_label = nc > 1 # multiple labels per box (adds 0.5ms/img) + + t = time.time() + output = [None] * prediction.shape[0] + for xi, x in enumerate(prediction): # image index, image inference + # Apply constraints + # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height + x = x[xc[xi]] # confidence + + # If none remain process next image + if not x.shape[0]: + continue + + # Compute conf + x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf + + # Box (center x, center y, width, height) to (x1, y1, x2, y2) + box = self.xywh2xyxy(x[:, :4]) + + # Detections matrix nx6 (xyxy, conf, cls) + if multi_label: + i, j = (x[:, 5:] > conf_thres).nonzero().t() + x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1) + else: # best class only + conf, j = x[:, 5:].max(1, keepdim=True) + x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres] + + # Filter by class + if classes: + x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)] + + # Apply finite constraint + # if not torch.isfinite(x).all(): + # x = x[torch.isfinite(x).all(1)] + + # If none remain process next image + n = x.shape[0] # number of boxes + if not n: + continue + + # Sort by confidence + # x = x[x[:, 4].argsort(descending=True)] + + # Batched NMS + c = x[:, 5:6] * (0 if agnostic else max_wh) # classes + boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores + i = torchvision.ops.boxes.nms(boxes, scores, iou_thres) + if i.shape[0] > max_det: # limit detections + i = i[:max_det] + if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean) + try: # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4) + iou = self.box_iou(boxes[i], boxes) > iou_thres # iou matrix + weights = iou * scores[None] # box weights + x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes + if redundant: + i = i[iou.sum(1) > 1] # require redundancy + except: # possible CUDA error https://github.com/ultralytics/yolov3/issues/1139 + print(x, i, x.shape, i.shape) + pass + + output[xi] = x[i] + if (time.time() - t) > time_limit: + break # time limit exceeded + + return output + + def box_iou(self, box1, box2): + # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py + """ + Return intersection-over-union (Jaccard index) of boxes. + Both sets of boxes are expected to be in (x1, y1, x2, y2) format. + Arguments: + box1 (Tensor[N, 4]) + box2 (Tensor[M, 4]) + Returns: + iou (Tensor[N, M]): the NxM matrix containing the pairwise + IoU values for every element in boxes1 and boxes2 + """ + + def box_area(box): + # box = 4xn + return (box[2] - box[0]) * (box[3] - box[1]) + + area1 = box_area(box1.t()) + area2 = box_area(box2.t()) + + # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2) + inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2) + return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter) + + def xywh2xyxy(self, x): + # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right + y = torch.zeros_like(x) if isinstance(x, torch.Tensor) else np.zeros_like(x) + y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x + y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y + y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x + y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y + return y + +class KeypointsProcessor(object): + def __init__(self, gpu_id=0): + import torch + from keypoints.lib.config import cfg, update_config + from keypoints.lib.models.pose_hrnet import get_pose_net + + class Namespace: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + args = Namespace(cfg=os.path.join(modelRoot, 'keypoints/experiments/coco/hrnet/w48_384x288_adam_lr1e-3.yaml'), + opts=['TEST.MODEL_FILE', os.path.join(modelRoot, 'pose_hrnet_w48_384x288.pth'), 'TEST.USE_GT_BBOX', 'False'], + dataDir='', + logDir='', + modelDir='', + prevModelDir='') + + update_config(cfg, args) + # print(cfg) + + self.model = get_pose_net(cfg, is_train=False) + self.model.eval() + if gpu_id != -1: + self.model.cuda(gpu_id) + + if cfg.TEST.MODEL_FILE: + if gpu_id == -1: + self.model.load_state_dict(torch.load(cfg.TEST.MODEL_FILE, map_location='cpu'), strict=False) + else: + self.model.load_state_dict(torch.load(cfg.TEST.MODEL_FILE, map_location=torch.device(gpu_id)), strict=False) + # print('load keypoints model weights') + + + self.gpu_id = gpu_id + + self.image_size = [288, 384] + self.cfg = cfg + + @staticmethod + def _xywh2cs(x, y, w, h, aspect_ratio=288.0 / 384.0, pixel_std=200): + center = np.zeros((2), dtype=np.float32) + center[0] = x + w * 0.5 + center[1] = y + h * 0.5 + + if w > aspect_ratio * h: + h = w * 1.0 / aspect_ratio + elif w < aspect_ratio * h: + w = h * aspect_ratio + scale = np.array( + [w * 1.0 / pixel_std, h * 1.0 / pixel_std], + dtype=np.float32) + if center[0] != -1: + scale = scale * 1.25 + + return center, scale + + @staticmethod + def get_affine_transform( + center, scale, rot, output_size, + shift=np.array([0, 0], dtype=np.float32), inv=0 + ): + def get_dir(src_point, rot_rad): + sn, cs = np.sin(rot_rad), np.cos(rot_rad) + + src_result = [0, 0] + src_result[0] = src_point[0] * cs - src_point[1] * sn + src_result[1] = src_point[0] * sn + src_point[1] * cs + + return src_result + + def get_3rd_point(a, b): + direct = a - b + return b + np.array([-direct[1], direct[0]], dtype=np.float32) + + if not isinstance(scale, np.ndarray) and not isinstance(scale, list): + scale = np.array([scale, scale]) + + scale_tmp = scale * 200.0 + src_w = scale_tmp[0] + dst_w = output_size[0] + dst_h = output_size[1] + + rot_rad = np.pi * rot / 180 + src_dir = get_dir([0, src_w * -0.5], rot_rad) + dst_dir = np.array([0, dst_w * -0.5], np.float32) + + src = np.zeros((3, 2), dtype=np.float32) + dst = np.zeros((3, 2), dtype=np.float32) + src[0, :] = center + scale_tmp * shift + src[1, :] = center + src_dir + scale_tmp * shift + dst[0, :] = [dst_w * 0.5, dst_h * 0.5] + dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir + + src[2:, :] = get_3rd_point(src[0, :], src[1, :]) + dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :]) + + if inv: + trans = cv2.getAffineTransform(np.float32(dst), np.float32(src)) + else: + trans = cv2.getAffineTransform(np.float32(src), np.float32(dst)) + + return trans + + def forward(self, image, boxes): + from keypoints.lib.core.inference import get_final_preds + all_preds = [] + for ix, box in enumerate(boxes): + x1, y1, x2, y2 = box + x, y, w, h, = x1, y1, x2 - x1, y2 - y1 + c, s = self._xywh2cs(x, y, w, h) + r = 0 + trans = self.get_affine_transform(c, s, r, self.image_size) + input = cv2.warpAffine( + image, + trans, + (int(self.image_size[0]), int(self.image_size[1])), + flags=cv2.INTER_LINEAR) + + # cv2.imshow('keypoint_input', input) + + input_rgb = cv2.cvtColor(input, cv2.COLOR_BGR2RGB) + input_rgb = np.divide(np.subtract(input_rgb.astype(np.float32) / 255, [0.406, 0.456, 0.485]), [0.225, 0.224, 0.229]) + img = np.expand_dims(input_rgb.astype(np.float32).transpose((2, 0, 1)), axis=0) + img = torch.from_numpy(img) + if self.gpu_id != -1: + img = img.cuda(self.gpu_id) + output = self.model(img) + output_numpy = output.detach().cpu().numpy() + + c = np.array([c]) + s = np.array([s]) + preds, maxvals = get_final_preds(self.cfg, output_numpy, c, s) + all_preds.append(np.concatenate((preds, maxvals), axis=2)) + if len(all_preds) > 0: + res = np.concatenate(all_preds, axis=0) + else: + res = [] + return res + + def forward_keypoints(self, image, keypoints): + from keypoints.lib.core.inference import get_final_preds + all_preds = [] + for ix, keypoint in enumerate(keypoints): + + contours = [] + for pt in (keypoint): + x, y, prob = pt + if prob > 0.05: + contours.append([x, y]) + contours = np.array(contours) + box = cv2.boundingRect(contours) + + x, y, w, h, = box + h = h * 1.2 + w = w * 1.1 + c, s = self._xywh2cs(x, y, w, h) + r = 0 + trans = self.get_affine_transform(c, s, r, self.image_size) + input = cv2.warpAffine( + image, + trans, + (int(self.image_size[0]), int(self.image_size[1])), + flags=cv2.INTER_LINEAR) + + # cv2.imshow('keypoint_input', input) + + input_rgb = cv2.cvtColor(input, cv2.COLOR_BGR2RGB) + input_rgb = np.divide(np.subtract(input_rgb.astype(np.float32) / 255, [0.406, 0.456, 0.485]), [0.225, 0.224, 0.229]) + img = np.expand_dims(input_rgb.astype(np.float32).transpose((2, 0, 1)), axis=0) + img = torch.from_numpy(img) + if self.gpu_id != -1: + img = img.cuda(self.gpu_id) + output = self.model(img) + output_numpy = output.detach().cpu().numpy() + + c = np.array([c]) + s = np.array([s]) + preds, maxvals = get_final_preds(self.cfg, output_numpy, c, s) + all_preds.append(np.concatenate((preds, maxvals), axis=2)) + if len(all_preds) > 0: + res = np.concatenate(all_preds, axis=0) + else: + res = [] + return res + + def draw_connect_keypoints(self, keypoints, w, h, draw_line=True): + COCO_PERSON_KEYPOINT_NAMES = ( + "nose", + "left_eye", "right_eye", + "left_ear", "right_ear", + "left_shoulder", "right_shoulder", + "left_elbow", "right_elbow", + "left_wrist", "right_wrist", + "left_hip", "right_hip", + "left_knee", "right_knee", + "left_ankle", "right_ankle", + ) + COCO_PERSON_KEYPOINT_COLORS = ( + (102, 204, 255), + (51, 153, 255), + (102, 0, 204), + (51, 102, 255), + (153, 255, 204), + (128, 229, 255), + (153, 255, 153), + (102, 255, 224), + (255, 102, 0), + (255, 255, 77), + (153, 255, 204), + (191, 255, 128), + (255, 195, 77), + (77, 204, 22), + (22, 139, 77), + (0, 56, 138), + (138, 76, 23), + ) + KEYPOINT_CONNECTION_RULES = [ + # face + ("left_ear", "left_eye", (102, 204, 255)), + ("right_ear", "right_eye", (51, 153, 255)), + ("left_eye", "nose", (102, 0, 204)), + ("nose", "right_eye", (51, 102, 255)), + # upper-body + ("left_shoulder", "right_shoulder", (255, 128, 0)), + ("left_shoulder", "left_elbow", (153, 255, 204)), + ("right_shoulder", "right_elbow", (128, 229, 255)), + ("left_elbow", "left_wrist", (153, 255, 153)), + ("right_elbow", "right_wrist", (102, 255, 224)), + # lower-body + ("left_hip", "right_hip", (255, 102, 0)), + ("left_hip", "left_knee", (255, 255, 77)), + ("right_hip", "right_knee", (153, 255, 204)), + ("left_knee", "left_ankle", (191, 255, 128)), + ("right_knee", "right_ankle", (255, 195, 77)), + ] + image = np.zeros((h, w, 3), dtype=np.uint8) + for ix, keypoint in enumerate(keypoints): + visible = {} + for idx, pt in enumerate(keypoint): + x, y, prob = pt + keypoint_name = COCO_PERSON_KEYPOINT_NAMES[idx] + if x < 0 or x >= w: continue + if y < 0 or y >= h: continue + if prob < 0.2: continue + visible[keypoint_name] = (int(x), int(y), prob) + + for ix, (k, v) in enumerate(visible.items()): + if not draw_line and ( + k == 'nose' or k == 'left_eye' or k == 'right_eye' or k == 'left_ear' or k == 'right_ear'): + continue + cv2.circle(image, (v[0], v[1]), 10, color=COCO_PERSON_KEYPOINT_COLORS[ix], thickness=cv2.FILLED) + # cv2.putText(image, '{}%'.format(int(v[2] * 100)), (int(v[0] + 5), int(v[1])), cv2.FONT_HERSHEY_COMPLEX, 0.5, + # (0, 255, 0), 1) + + if not draw_line: + continue + + for kp0, kp1, color in KEYPOINT_CONNECTION_RULES: + if kp0 in visible and kp1 in visible: + x0, y0, _ = visible[kp0] + x1, y1, _ = visible[kp1] + color = (color[2], color[1], color[0]) + cv2.line(image, (x0, y0), (x1, y1), color=color, thickness=10, lineType=cv2.LINE_AA) + try: + ls_x, ls_y, _ = visible["left_shoulder"] + rs_x, rs_y, _ = visible["right_shoulder"] + mid_shoulder_x, mid_shoulder_y = int((ls_x + rs_x) / 2), int((ls_y + rs_y) / 2) + except KeyError: + pass + else: + # draw line from nose to mid-shoulder + nose_x, nose_y, _ = visible.get("nose", (None, None, None)) + if nose_x is not None: + cv2.line(image, (nose_x, nose_y), (mid_shoulder_x, mid_shoulder_y), color=(0, 0, 255), thickness=10, + lineType=cv2.LINE_AA) + + try: + # draw line from mid-shoulder to mid-hip + lh_x, lh_y, _ = visible["left_hip"] + rh_x, rh_y, _ = visible["right_hip"] + except KeyError: + pass + else: + mid_hip_x, mid_hip_y = int((lh_x + rh_x) / 2), int((lh_y + rh_y) / 2) + cv2.line(image, (mid_hip_x, mid_hip_y), (mid_shoulder_x, mid_shoulder_y), color=(0, 0, 255), + thickness=10, + lineType=cv2.LINE_AA) + return image + +class Generator_Hair(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.model_dir = modelRoot + + self.dst_height = 768 + self.dst_width = 768 + + generator_hair_path = os.path.join(self.model_dir, "master_hair_v8_onlyhair_nowarp_all_0709.pt") # default master_hair_v8_onlyhair_blur_aug_0413.pt master_hair_v8_onlyhair_min_0223 + # "master_hair_v8_onlyhair_nowarp_all_0709" # master_hair_v8_onlyhair_blur_aug_all_0703 + self.net = torch.jit.load(generator_hair_path, torch.device('cpu')).to(self.device) + # print("generate hair net: ", self.net) + self.net.eval() + + def Generator_Hair_inference(self, ref_rgb_8uc3_512, ref_matting_8uc3_512, ref_baldseg_8uc3_512, + ref_landmark_f1k2_512, + user_baldseg_8uc3_512, user_bald_8uc3_512, user_landmark_f1k2_512): + + """ + 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 + + user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 + + output： + + hair_gene_8uc3_512: 生成 发型图， uint8 （0-255） + + """ + nohair_mask = user_baldseg_8uc3_512.copy() + user_bald_mask = (user_baldseg_8uc3_512 == [0, 255, 0]).all(axis=2) + user_pts137 = landmark_processor.pts_1k_to_137(user_landmark_f1k2_512).astype(np.int32) + # 2 ********************** nohair_mask ********************** + # Label mouth + cv2.fillPoly(nohair_mask, + np.concatenate((user_pts137[47:35:-1], user_pts137[56:64], [user_pts137[48], user_pts137[22]]))[ + np.newaxis, :, :], (255, 255, 0)) + cv2.fillPoly(nohair_mask, np.concatenate((user_pts137[22:37], user_pts137[56:47:-1]))[np.newaxis, :, :], + (255, 0, 128)) + + cv2.fillPoly(nohair_mask, user_pts137[88:104][np.newaxis, :, :], (255, 0, 255)) # eye + cv2.fillPoly(nohair_mask, user_pts137[105:121][np.newaxis, :, :], (255, 0, 255)) # eye + + # Label nose + cv2.fillPoly(nohair_mask, user_pts137[64:79][np.newaxis, :, :], (255, 255, 255)) + + # Label eyebrow + cv2.fillPoly(nohair_mask, user_pts137[121:129][np.newaxis, :, :], (255, 128, 0)) + cv2.fillPoly(nohair_mask, user_pts137[129:137][np.newaxis, :, :], (255, 128, 0)) + + # 5 ********************** bald_img ********************** + bald_img = (user_bald_mask[:, :, np.newaxis] * user_bald_8uc3_512).astype(np.uint8) + + # 8 ********************** another pose hair_image ********************** + another_pose_image = ref_rgb_8uc3_512.copy() + + another_nohair_pose_mask = ref_baldseg_8uc3_512.copy() + 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_512).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_512.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_nohair_mask = torch.from_numpy( + (nohair_mask.astype(np.float32) / 255).transpose(2, 0, 1)[np.newaxis, :, :, :]).to(self.device) + # input_nohair_image = torch.from_numpy( (nohair_image.astype(np.float32) / 255).transpose(2, 0, 1)[np.newaxis, :, :, :]).to(self.device) + input_bald_img = torch.from_numpy( + (bald_img.astype(np.float32) / 255).transpose(2, 0, 1)[np.newaxis, :, :, :]).to(self.device) + input_another_pose_hair_image = torch.from_numpy( + (another_pose_hair_image.astype(np.float32) / 255).transpose(2, 0, 1)[np.newaxis, :, :, :]).to(self.device) + + zero_tensor = torch.from_numpy(np.zeros((1, 3, self.dst_height, self.dst_width), dtype=np.float32)).to( + self.device) + + with torch.no_grad(): + fake_image = self.net(input_nohair_mask, input_bald_img, input_another_pose_hair_image, zero_tensor) + fake_image_numpy = (fake_image[0].cpu().detach().numpy().transpose((1, 2, 0)) * 255).astype(np.uint8).copy() + + return fake_image_numpy + + def Generator_Hair_inference_use_pref(self, another_pose_hair_image, user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, gender="boy"): + + """ + input： + + 图像尺寸基于： 人脸 512， 图像均为3通道 + another_pose_hair_image: 参考图 处理好 numpy，float32 + user_baldseg_8uc3_512：用户图 光头分割 mask， uint8 （0-255） + user_bald_8uc3_512： 用户图 光头， uint8 （0-255） + user_landmark_f1k2_512： 用户图 关键点 1k*2 float32 + + output： + + hair_gene_8uc3_512: 生成 发型图， uint8 （0-255） + + """ + + nohair_mask = user_baldseg_8uc3_768.copy() + # user_bald_mask = (user_baldseg_8uc3_768 == [0, 255, 0]).all(axis=2) + user_pts137 = landmark_processor.pts_1k_to_137(user_landmark_f1k2_768).astype(np.int32) + # 2 ********************** nohair_mask ********************** + # Label mouth + cv2.fillPoly(nohair_mask, np.concatenate((user_pts137[47:35:-1], user_pts137[56:64], [user_pts137[48], user_pts137[22]]))[ + np.newaxis, :, :], (255, 255, 0)) + cv2.fillPoly(nohair_mask, np.concatenate((user_pts137[22:37], user_pts137[56:47:-1]))[np.newaxis, :, :], + (255, 0, 128)) + + cv2.fillPoly(nohair_mask, user_pts137[88:104][np.newaxis, :, :], (255, 0, 255)) # eye + cv2.fillPoly(nohair_mask, user_pts137[105:121][np.newaxis, :, :], (255, 0, 255)) # eye + + # Label nose + cv2.fillPoly(nohair_mask, user_pts137[64:79][np.newaxis, :, :], (255, 255, 255)) + + # Label eyebrow + cv2.fillPoly(nohair_mask, user_pts137[121:129][np.newaxis, :, :], (255, 128, 0)) + cv2.fillPoly(nohair_mask, user_pts137[129:137][np.newaxis, :, :], (255, 128, 0)) + + # # 4 ********************** nohair_image ********************** + nohair_image_orig_blur = cv2.blur(user_bald_8uc3_768, (self.dst_width // 5, self.dst_width // 5)) + # bald_img = nohair_image_orig_blur.copy() + clear_mask = np.zeros((self.dst_height, self.dst_width, 3), dtype=np.uint8) + cv2.fillPoly(clear_mask, np.concatenate((user_pts137[96:88:-1], user_pts137[105:114], user_pts137[2::-1], user_pts137[21:19:-1]))[np.newaxis, :, :], (1, 1, 1)) + + bald_img = nohair_image_orig_blur * (1 - clear_mask) + user_bald_8uc3_768 * clear_mask + + # inter_res = np.concatenate((user_baldseg_8uc3_768, user_bald_8uc3_768, bald_img), axis=1) + # cv2.imshow("gen hair inter_res: ", inter_res) + # cv2.waitKey() + + # 改动 for blur version + # nohair_image_orig_blur = cv2.blur(user_bald_8uc3_768, (self.dst_width // 5, self.dst_width // 5)) + # bald_img = nohair_image_orig_blur.copy() + + input_nohair_mask = torch.from_numpy((nohair_mask.astype(np.float32) / 255).transpose(2, 0, 1)[np.newaxis, :, :, :]).to(self.device) + input_bald_img = torch.from_numpy((bald_img.astype(np.float32) / 255).transpose(2, 0, 1)[np.newaxis, :, :, :]).to(self.device) + input_another_pose_hair_image = torch.from_numpy(another_pose_hair_image.transpose(2, 0, 1)[np.newaxis, :, :, :]).to(self.device) + + zero_tensor = torch.from_numpy(np.zeros((1, 3, self.dst_height, self.dst_width), dtype=np.float32)).to( + self.device) + + with torch.no_grad(): + fake_image = self.net(input_nohair_mask, input_bald_img, input_another_pose_hair_image, zero_tensor) + fake_image_numpy = (fake_image[0].cpu().detach().numpy().transpose((1, 2, 0)) * 255).astype(np.uint8).copy() + + # hair__align_concat_show = np.concatenate((nohair_mask, bald_img, (another_pose_hair_image*255).astype(np.uint8), fake_image_numpy), axis=1) + # hair__align_concat_show = cv2.resize(hair__align_concat_show, (0, 0), fx=0.5, fy=0.5) + # cv2.imshow("hair__align_concat_show", hair__align_concat_show) + # cv2.waitKey(0) + + return fake_image_numpy + +class Generator_Fusion_Res(Process_Data): + 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.model_dir = modelRoot + self.output_img_size = 768 + + fusion_model_path = os.path.join(self.model_dir, 'hair_fusion_0427.pt') # v1: 0412 v2:zxm_v2_use_paired_data_remap_add_encode_04_12_768 + self.load_fusion_model(fusion_model_path) + + + def load_fusion_model(self, hair_fusion_model_path): + self.fusion_model = torch.jit.load(hair_fusion_model_path, map_location='cpu').to(self.device) + + + def inference_girl(self, user_generator_hair_8uc3_orisize, user_generator_matte_8uc3_orisize, + user_generator_landmark_8uc3_orisize, user_bald_mask_8uc3_orisize): + """ + input: + user_generator_hair_8uc3_orisize: 换发型结果 原图尺寸 8uc3 + user_generator_matte_8uc3_orisize: 换发型结果matting alpha图 原图尺寸 8uc3 + + output: + fusion_res + + """ + ratio = 0.5 + h_offset = 0.45 + user_pts1k = user_generator_landmark_8uc3_orisize.astype(np.int32) + user_real = user_generator_hair_8uc3_orisize + user_hair_mask = user_generator_matte_8uc3_orisize + + # user_bald_mask_8uc3_orisize = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/5c_mask_xxq/00B4D4B8-8F26-7DF9-51D0-41492D30DDF720201230_boy_20_None_0.png") + + # cv2.imshow("user_bald_mask_8uc3_orisize", user_bald_mask_8uc3_orisize) + # cv2.imshow("user_generator_hair_8uc3_orisize", user_generator_hair_8uc3_orisize) + # cv2.imshow("user_generator_matte_8uc3_orisize", user_generator_matte_8uc3_orisize) + # cv2.waitKey() + + + user_bald_mask = user_bald_mask_8uc3_orisize + user_hair_mask_orig = user_hair_mask[:, :, 0].copy() + + user_real_save = user_real.copy() + orig_h, orig_w, _ = user_real_save.shape + + M_user = landmark_processor.get_transform_mat_hair_ratio_v1(user_pts1k, self.output_img_size, ratio, + h_offset) + inv_M_user = cv2.invertAffineTransform(M_user) + user_real = cv2.warpAffine(user_real, M_user, (self.output_img_size, self.output_img_size)) + user_hair_mask = cv2.warpAffine(user_hair_mask, M_user, (self.output_img_size, self.output_img_size), + flags=cv2.INTER_CUBIC) + user_bald_mask = cv2.warpAffine(user_bald_mask, M_user, (self.output_img_size, self.output_img_size), + flags=cv2.INTER_NEAREST) + + random_v = 15 + kernel_e = np.ones((random_v, random_v), np.uint8) + kernel_d = np.ones((random_v + 10, random_v + 10), np.uint8) + user_hair_mask_erode = cv2.erode(user_hair_mask[:, :, 0], kernel_e, iterations=1) + user_hair_mask = cv2.dilate(user_hair_mask[:, :, 0], kernel_d, iterations=1) + + blend = np.clip(user_real * (user_hair_mask_erode[:, :, np.newaxis] / 255), 0, 255) + lab_img_paf = cv2.cvtColor(user_real, cv2.COLOR_BGR2LAB) + lab_img_paf[:, :, 1] = 0 + lab_img_paf[:, :, 2] = 0 + # img_paf = np.clip(lab_img_paf * (user_hair_mask[:, :, np.newaxis] / 255) + user_real * ( + # 1 - user_hair_mask[:, :, np.newaxis] / 255), 0, 255) + img_paf = user_real.copy() + img_paf[user_hair_mask > 0] = lab_img_paf[user_hair_mask > 0] + + # cv2.imshow("blend", blend.astype(np.uint8)) + # cv2.imshow("user_hair_mask", user_hair_mask.astype(np.uint8)) + # cv2.imshow("img_paf", img_paf.astype(np.uint8)) + # cv2.waitKey() + + # img_paf = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/user_debug/fusion_test/img_paf.png") + # blend = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/user_debug/fusion_test/blend.png") + + condition = blend # np.concatenate((user_mask, user_hair_mask), axis=2) + condition = (condition.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] + input_paf = (img_paf.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] + condition = torch.from_numpy(condition).to(self.device) + input_paf = torch.from_numpy(input_paf).to(self.device) + + # ******************* forward ******************* + test_fake = self.fusion_model(input_paf, condition) + test_res = (test_fake.detach()[0].to('cpu').numpy().transpose((1, 2, 0)) * 255).astype(np.uint8).copy() + + input_out_show = np.concatenate((img_paf, user_bald_mask, (blend).astype(np.uint8), test_res), axis=1) + # cv2.imshow("input_out_show", input_out_show) + # cv2.waitKey() + + # cv2.imwrite("/media/DATA_4T/test_hair/debug_fusion/user_debug/test_new_res.png", test_res) + + test_res_orig = user_real_save.copy() + user_hair_mask_e_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) + user_hair_mask_d_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) + cv2.warpAffine(test_res, inv_M_user, (orig_w, orig_h), + dst=test_res_orig, borderMode=cv2.BORDER_TRANSPARENT) + cv2.warpAffine(user_hair_mask_erode, inv_M_user, (orig_w, orig_h), + dst=user_hair_mask_e_orig) + cv2.warpAffine(user_hair_mask, inv_M_user, (orig_w, orig_h), + dst=user_hair_mask_d_orig) + + # cv2.imshow("user_hair_mask", user_hair_mask) + # cv2.imshow("user_hair_mask_d_orig", user_hair_mask_d_orig) + # cv2.imshow("user_hair_mask_erode", user_hair_mask_erode) + # cv2.imshow("user_hair_mask_e_orig", user_hair_mask_e_orig) + # cv2.waitKey() + + user_bald_mask_b_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) + user_bald_mask_b = np.zeros(user_bald_mask[:, :, 0].shape, dtype=np.uint8) + user_bald_mask_b[user_bald_mask[:, :, 0] > 0] = 1 + user_bald_mask_b[user_bald_mask[:, :, 1] > 0] = 1 + user_bald_mask_b[user_bald_mask[:, :, 2] > 0] = 1 + user_bald_mask_b = cv2.dilate(user_bald_mask_b, np.ones((25, 25))) + cv2.warpAffine(user_bald_mask_b, inv_M_user, (orig_w, orig_h), + dst=user_bald_mask_b_orig) + + # user_hair_mask_d_orig_b = cv2.blur(user_hair_mask_d_orig, (30, 30)) + user_hair_mask_d_orig_b = user_hair_mask_d_orig + user_hair_mask_e_orig_b = cv2.blur(user_hair_mask_e_orig, (15, 15)) + + # user_hair_mask_d_orig[~(user_bald_mask_b_orig.astype(bool))] = user_hair_mask_e_orig[~(user_bald_mask_b_orig.astype(bool))] + # # user_hair_mask_d_orig_b = cv2.blur(user_hair_mask_d_orig, (15, 15)) + user_hair_mask_d_orig_b[~(user_bald_mask_b_orig.astype(bool))] = user_hair_mask_e_orig[ + ~(user_bald_mask_b_orig.astype(bool))] + user_hair_mask_d_orig_b = cv2.blur(user_hair_mask_d_orig_b, (15, 15)) + + test_res_clear_orig = user_real_save * ( + 1 - user_hair_mask_d_orig_b[:, :, np.newaxis] / 255) + test_res_orig * ( + user_hair_mask_d_orig_b[:, :, np.newaxis] / 255) + test_res_clear_orig = (np.clip(test_res_clear_orig, 0, 255)).astype(np.uint8) + + test_res_orig_lab = cv2.cvtColor(test_res_clear_orig, cv2.COLOR_BGR2LAB) + user_real_save_lab = cv2.cvtColor(user_real_save, cv2.COLOR_BGR2LAB) + test_res_orig_lab[:, :, 0] = user_real_save_lab[:, :, 0] + fusion_res = cv2.cvtColor(test_res_orig_lab, cv2.COLOR_LAB2BGR) + + user_hair_mask_e_orig_b[~(user_bald_mask_b_orig.astype(bool))] = user_hair_mask_d_orig_b[~(user_bald_mask_b_orig.astype(bool))] + fusion_res = test_res_clear_orig * ( + 1 - user_hair_mask_e_orig_b[:, :, np.newaxis] / 255) + fusion_res * ( + user_hair_mask_e_orig_b[:, :, np.newaxis] / 255) + fusion_res = (np.clip(fusion_res, 0, 255)).astype(np.uint8) + + # con = np.concatenate((user_hair_mask_orig, user_hair_mask_d_orig_b), axis=1) + # con = cv2.resize(con, (0, 0), fx=1 / 4, fy=1 / 4, interpolation=cv2.INTER_CUBIC) + # + con_res = np.concatenate(( user_real_save, fusion_res), axis=1) + # con_res = cv2.resize(con_res, (0, 0), fx=1 / 4, fy=1 / 4, interpolation=cv2.INTER_CUBIC) + # + # cv2.imshow("con", con) + # cv2.imshow("con_res", con_res) + # cv2.waitKey() + # cv2.imwrite("/media/DATA_4T/test_hair/debug_fusion/user_debug/con_new1_fuison.png", con_res) + + return fusion_res + + def inference_boy(self, user_generator_hair_8uc3_orisize, user_generator_matte_8uc3_orisize, + user_generator_landmark_8uc3_orisize, user_bald_mask_8uc3_orisize): + """ + input: + user_generator_hair_8uc3_orisize: 换发型结果 原图尺寸 8uc3 + user_generator_matte_8uc3_orisize: 换发型结果matting alpha图 原图尺寸 8uc3 + + output: + fusion_res + + """ + ratio = 0.6 + h_offset = 0.65 + user_pts1k = user_generator_landmark_8uc3_orisize.astype(np.int32) + user_real = user_generator_hair_8uc3_orisize + user_hair_mask = user_generator_matte_8uc3_orisize + + # user_bald_mask_8uc3_orisize = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/5c_mask_xxq/00B4D4B8-8F26-7DF9-51D0-41492D30DDF720201230_boy_20_None_0.png") + + # cv2.imshow("user_bald_mask_8uc3_orisize", user_bald_mask_8uc3_orisize) + # cv2.imshow("user_generator_hair_8uc3_orisize", user_generator_hair_8uc3_orisize) + # cv2.imshow("user_generator_matte_8uc3_orisize", user_generator_matte_8uc3_orisize) + # cv2.waitKey() + + user_bald_mask = user_bald_mask_8uc3_orisize + user_hair_mask_orig = user_hair_mask[:, :, 0].copy() + + user_real_save = user_real.copy() + orig_h, orig_w, _ = user_real_save.shape + + M_user = landmark_processor.get_transform_mat_hair_ratio_v1(user_pts1k, self.output_img_size, ratio, + h_offset) + inv_M_user = cv2.invertAffineTransform(M_user) + user_real = cv2.warpAffine(user_real, M_user, (self.output_img_size, self.output_img_size)) + user_hair_mask = cv2.warpAffine(user_hair_mask, M_user, (self.output_img_size, self.output_img_size), + flags=cv2.INTER_CUBIC) + user_bald_mask = cv2.warpAffine(user_bald_mask, M_user, (self.output_img_size, self.output_img_size), + flags=cv2.INTER_NEAREST) + + random_v = 15 + kernel_e = np.ones((random_v, random_v), np.uint8) + kernel_d = np.ones((random_v + 10, random_v + 10), np.uint8) + user_hair_mask_erode = cv2.erode(user_hair_mask[:, :, 0], kernel_e, iterations=1) + user_hair_mask = cv2.dilate(user_hair_mask[:, :, 0], kernel_d, iterations=1) + + blend = np.clip(user_real * (user_hair_mask_erode[:, :, np.newaxis] / 255), 0, 255) + lab_img_paf = cv2.cvtColor(user_real, cv2.COLOR_BGR2LAB) + lab_img_paf[:, :, 1] = 0 + lab_img_paf[:, :, 2] = 0 + # img_paf = np.clip(lab_img_paf * (user_hair_mask[:, :, np.newaxis] / 255) + user_real * ( + # 1 - user_hair_mask[:, :, np.newaxis] / 255), 0, 255) + img_paf = user_real.copy() + img_paf[user_hair_mask > 0] = lab_img_paf[user_hair_mask > 0] + + # cv2.imshow("blend", blend.astype(np.uint8)) + # cv2.imshow("user_hair_mask", user_hair_mask.astype(np.uint8)) + # cv2.imshow("img_paf", img_paf.astype(np.uint8)) + # cv2.waitKey() + + # img_paf = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/user_debug/fusion_test/img_paf.png") + # blend = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/user_debug/fusion_test/blend.png") + + condition = blend # np.concatenate((user_mask, user_hair_mask), axis=2) + condition = (condition.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] + input_paf = (img_paf.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] + condition = torch.from_numpy(condition).to(self.device) + input_paf = torch.from_numpy(input_paf).to(self.device) + + # ******************* forward ******************* + test_fake = self.fusion_model(input_paf, condition) + test_res = (test_fake.detach()[0].to('cpu').numpy().transpose((1, 2, 0)) * 255).astype(np.uint8).copy() + + input_out_show = np.concatenate((img_paf, user_bald_mask, (blend).astype(np.uint8), test_res), axis=1) + # cv2.imshow("input_out_show", input_out_show) + # cv2.waitKey() + + # cv2.imwrite("/media/DATA_4T/test_hair/debug_fusion/user_debug/test_new_res.png", test_res) + + test_res_orig = user_real_save.copy() + user_hair_mask_e_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) + user_hair_mask_d_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) + cv2.warpAffine(test_res, inv_M_user, (orig_w, orig_h), + dst=test_res_orig, borderMode=cv2.BORDER_TRANSPARENT) + cv2.warpAffine(user_hair_mask_erode, inv_M_user, (orig_w, orig_h), + dst=user_hair_mask_e_orig) + cv2.warpAffine(user_hair_mask, inv_M_user, (orig_w, orig_h), + dst=user_hair_mask_d_orig) + + # cv2.imshow("user_hair_mask", user_hair_mask) + # cv2.imshow("user_hair_mask_d_orig", user_hair_mask_d_orig) + # cv2.imshow("user_hair_mask_erode", user_hair_mask_erode) + # cv2.imshow("user_hair_mask_e_orig", user_hair_mask_e_orig) + # cv2.waitKey() + + user_bald_mask_b_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) + user_bald_mask_b = np.zeros(user_bald_mask[:, :, 0].shape, dtype=np.uint8) + user_bald_mask_b[user_bald_mask[:, :, 0] > 0] = 1 + user_bald_mask_b[user_bald_mask[:, :, 1] > 0] = 1 + user_bald_mask_b[user_bald_mask[:, :, 2] > 0] = 1 + user_bald_mask_b = cv2.dilate(user_bald_mask_b, np.ones((10, 10))) + cv2.warpAffine(user_bald_mask_b, inv_M_user, (orig_w, orig_h), + dst=user_bald_mask_b_orig) + + user_hair_mask_d_orig_b = user_hair_mask_d_orig + user_hair_mask_d_orig_b[~(user_bald_mask_b_orig.astype(bool))] = user_hair_mask_e_orig[ + ~(user_bald_mask_b_orig.astype(bool))] + user_hair_mask_d_orig_b = cv2.blur(user_hair_mask_d_orig_b, (15, 15)) + user_hair_mask_e_orig_b = cv2.blur(user_hair_mask_e_orig, (15, 15)) + + test_res_clear_orig = user_real_save * ( + 1 - user_hair_mask_d_orig_b[:, :, np.newaxis] / 255) + test_res_orig * ( + user_hair_mask_d_orig_b[:, :, np.newaxis] / 255) + test_res_clear_orig = (np.clip(test_res_clear_orig, 0, 255)).astype(np.uint8) + + test_res_orig_lab = cv2.cvtColor(test_res_clear_orig, cv2.COLOR_BGR2LAB) + user_real_save_lab = cv2.cvtColor(user_real_save, cv2.COLOR_BGR2LAB) + test_res_orig_lab[:, :, 0] = user_real_save_lab[:, :, 0] + fusion_res = cv2.cvtColor(test_res_orig_lab, cv2.COLOR_LAB2BGR) + + user_hair_mask_e_orig_b[~(user_bald_mask_b_orig.astype(bool))] = user_hair_mask_d_orig_b[ + ~(user_bald_mask_b_orig.astype(bool))] + fusion_res = test_res_clear_orig * ( + 1 - user_hair_mask_e_orig_b[:, :, np.newaxis] / 255) + fusion_res * ( + user_hair_mask_e_orig_b[:, :, np.newaxis] / 255) + fusion_res = (np.clip(fusion_res, 0, 255)).astype(np.uint8) + + # con = np.concatenate((user_hair_mask_orig, user_hair_mask_d_orig_b), axis=1) + # con = cv2.resize(con, (0, 0), fx=1 / 4, fy=1 / 4, interpolation=cv2.INTER_CUBIC) + # + # con_res = np.concatenate((user_real_save, fusion_res), axis=1) + # con_res = cv2.resize(con_res, (0, 0), fx=1 / 4, fy=1 / 4, interpolation=cv2.INTER_CUBIC) + # + # cv2.imshow("con", con) + # cv2.imshow("con_res", con_res) + # cv2.waitKey() + + return fusion_res + + def inference(self, user_generator_hair_8uc3_orisize, user_generator_matte_8uc3_orisize, + user_generator_landmark_8uc3_orisize, user_bald_mask_8uc3_orisize, ratio): + """ + input: + user_generator_hair_8uc3_orisize: 换发型结果 原图尺寸 8uc3 + user_generator_matte_8uc3_orisize: 换发型结果matting alpha图 原图尺寸 8uc3 + + output: + fusion_res + + """ + + user_pts1k = user_generator_landmark_8uc3_orisize.astype(np.int32) + user_real = user_generator_hair_8uc3_orisize + user_hair_mask = user_generator_matte_8uc3_orisize + + # user_bald_mask_8uc3_orisize = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/5c_mask_xxq/00B4D4B8-8F26-7DF9-51D0-41492D30DDF720201230_boy_20_None_0.png") + + # cv2.imshow("user_bald_mask_8uc3_orisize", user_bald_mask_8uc3_orisize) + # cv2.imshow("user_generator_hair_8uc3_orisize", user_generator_hair_8uc3_orisize) + # cv2.imshow("user_generator_matte_8uc3_orisize", user_generator_matte_8uc3_orisize) + # cv2.waitKey() + + + user_bald_mask = user_bald_mask_8uc3_orisize + user_hair_mask_orig = user_hair_mask[:, :, 0].copy() + + user_real_save = user_real.copy() + orig_h, orig_w, _ = user_real_save.shape + + if ratio == 0: + M_user = self.get_hair_M_boy_v1(user_pts1k) + elif ratio == 1: + M_user = self.get_hair_M_girl_v1(user_pts1k) + elif ratio == 2: + M_user = self.get_hair_M_girl_v2(user_pts1k) + else: + M_user = self.get_hair_M_girl_v1(user_pts1k) + + # M_user = landmark_processor.get_transform_mat_hair_ratio_v1(user_pts1k, self.output_img_size, ratio, + # h_offset) + inv_M_user = cv2.invertAffineTransform(M_user) + user_real = cv2.warpAffine(user_real, M_user, (self.output_img_size, self.output_img_size)) + user_hair_mask = cv2.warpAffine(user_hair_mask, M_user, (self.output_img_size, self.output_img_size), + flags=cv2.INTER_CUBIC) + user_bald_mask = cv2.warpAffine(user_bald_mask, M_user, (self.output_img_size, self.output_img_size), + flags=cv2.INTER_NEAREST) + + random_v = 15 + kernel_e = np.ones((random_v, random_v), np.uint8) + kernel_d = np.ones((random_v + 10, random_v + 10), np.uint8) + user_hair_mask_erode = cv2.erode(user_hair_mask[:, :, 0], kernel_e, iterations=1) + user_hair_mask = cv2.dilate(user_hair_mask[:, :, 0], kernel_d, iterations=1) + + blend = np.clip(user_real * (user_hair_mask_erode[:, :, np.newaxis] / 255), 0, 255) + lab_img_paf = cv2.cvtColor(user_real, cv2.COLOR_BGR2LAB) + lab_img_paf[:, :, 1] = 0 + lab_img_paf[:, :, 2] = 0 + # img_paf = np.clip(lab_img_paf * (user_hair_mask[:, :, np.newaxis] / 255) + user_real * ( + # 1 - user_hair_mask[:, :, np.newaxis] / 255), 0, 255) + img_paf = user_real.copy() + img_paf[user_hair_mask > 0] = lab_img_paf[user_hair_mask > 0] + + # cv2.imshow("blend", blend.astype(np.uint8)) + # cv2.imshow("user_hair_mask", user_hair_mask.astype(np.uint8)) + # cv2.imshow("img_paf", img_paf.astype(np.uint8)) + # cv2.waitKey() + + # img_paf = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/user_debug/fusion_test/img_paf.png") + # blend = cv2.imread("/media/DATA_4T/test_hair/debug_fusion/user_debug/fusion_test/blend.png") + + condition = blend # np.concatenate((user_mask, user_hair_mask), axis=2) + condition = (condition.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] + input_paf = (img_paf.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] + condition = torch.from_numpy(condition).to(self.device) + input_paf = torch.from_numpy(input_paf).to(self.device) + + # ******************* forward ******************* + test_fake = self.fusion_model(input_paf, condition) + test_res = (test_fake.detach()[0].to('cpu').numpy().transpose((1, 2, 0)) * 255).astype(np.uint8).copy() + + input_out_show = np.concatenate((img_paf, user_bald_mask, (blend).astype(np.uint8), test_res), axis=1) + ratio = 1536. / max(input_out_show.shape[:2]) + input_out_show = cv2.resize(input_out_show, (0, 0), fx=ratio, fy=ratio) + # cv2.imshow("fusion_input_out_show", input_out_show) + # cv2.waitKey() + + test_res_orig = user_real_save.copy() + user_hair_mask_e_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) + user_hair_mask_d_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) + cv2.warpAffine(test_res, inv_M_user, (orig_w, orig_h), + dst=test_res_orig, borderMode=cv2.BORDER_TRANSPARENT) + cv2.warpAffine(user_hair_mask_erode, inv_M_user, (orig_w, orig_h), + dst=user_hair_mask_e_orig) + cv2.warpAffine(user_hair_mask, inv_M_user, (orig_w, orig_h), + dst=user_hair_mask_d_orig) + + # cv2.imshow("user_hair_mask", user_hair_mask) + # cv2.imshow("user_hair_mask_d_orig", user_hair_mask_d_orig) + # cv2.imshow("user_hair_mask_erode", user_hair_mask_erode) + # cv2.imshow("user_hair_mask_e_orig", user_hair_mask_e_orig) + # cv2.waitKey() + + user_bald_mask_b_orig = np.zeros(user_real_save[:, :, 0].shape, dtype=np.uint8) + user_bald_mask_b = np.zeros(user_bald_mask[:, :, 0].shape, dtype=np.uint8) + user_bald_mask_b[user_bald_mask[:, :, 0] > 0] = 1 + user_bald_mask_b[user_bald_mask[:, :, 1] > 0] = 1 + user_bald_mask_b[user_bald_mask[:, :, 2] > 0] = 1 + + if ratio != 0: + user_bald_mask_b = cv2.dilate(user_bald_mask_b, np.ones((25, 25))) # default 25 + cv2.warpAffine(user_bald_mask_b, inv_M_user, (orig_w, orig_h), + dst=user_bald_mask_b_orig) + + user_hair_mask_d_orig_b = user_hair_mask_d_orig + user_hair_mask_e_orig_b = cv2.blur(user_hair_mask_e_orig, (15, 15)) + + user_hair_mask_d_orig_b[~(user_bald_mask_b_orig.astype(bool))] = user_hair_mask_e_orig[ + ~(user_bald_mask_b_orig.astype(bool))] + user_hair_mask_d_orig_b = cv2.blur(user_hair_mask_d_orig_b, (15, 15)) + else: + user_bald_mask_b = cv2.dilate(user_bald_mask_b, np.ones((10, 10))) + cv2.warpAffine(user_bald_mask_b, inv_M_user, (orig_w, orig_h), + dst=user_bald_mask_b_orig) + + user_hair_mask_d_orig_b = user_hair_mask_d_orig + user_hair_mask_d_orig_b[~(user_bald_mask_b_orig.astype(bool))] = user_hair_mask_e_orig[ + ~(user_bald_mask_b_orig.astype(bool))] + user_hair_mask_d_orig_b = cv2.blur(user_hair_mask_d_orig_b, (15, 15)) + user_hair_mask_e_orig_b = cv2.blur(user_hair_mask_e_orig, (15, 15)) + + test_res_clear_orig = user_real_save * ( + 1 - user_hair_mask_d_orig_b[:, :, np.newaxis] / 255) + test_res_orig * ( + user_hair_mask_d_orig_b[:, :, np.newaxis] / 255) + test_res_clear_orig = (np.clip(test_res_clear_orig, 0, 255)).astype(np.uint8) + + test_res_orig_lab = cv2.cvtColor(test_res_clear_orig, cv2.COLOR_BGR2LAB) + user_real_save_lab = cv2.cvtColor(user_real_save, cv2.COLOR_BGR2LAB) + test_res_orig_lab[:, :, 0] = user_real_save_lab[:, :, 0] + fusion_res = cv2.cvtColor(test_res_orig_lab, cv2.COLOR_LAB2BGR) + + user_hair_mask_e_orig_b[~(user_bald_mask_b_orig.astype(bool))] = user_hair_mask_d_orig_b[~(user_bald_mask_b_orig.astype(bool))] + fusion_res = test_res_clear_orig * ( + 1 - user_hair_mask_e_orig_b[:, :, np.newaxis] / 255) + fusion_res * ( + user_hair_mask_e_orig_b[:, :, np.newaxis] / 255) + fusion_res = (np.clip(fusion_res, 0, 255)).astype(np.uint8) + return fusion_res + +class Change_Hair_Color(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.model_dir = modelRoot + + self.dst_height = 768 + self.dst_width = 768 + + generator_hair_path = os.path.join(self.model_dir, "zxm_v1_encode_refer_add_body_0107_use_gendata_from_hisd_0721.pt") #last: master_gen_hair_zxm_change_hair_color_v1_encode_refer_add_body + self.net = torch.jit.load(generator_hair_path, torch.device('cpu')).to(self.device) + print("hair color net: ", self.net) + self.net.eval() + + def Change_Hair_inference(self, user_rgb_8uc3_change_color_768, user_matting_8uc3_change_color_768, ref_rgb_8uc3_change_color_768, ref_matting_8uc3_change_color_768): + + """ + input： + + 图像尺寸基于： 人脸 768， 图像均为3通道 + + user_rgb_8uc3_change_color_768: 用户图， size 768， uint8 （0-255） + user_matting_8uc3_change_color_768: 用户图 matting结果 size 768， uint8 （0-255） + user_landmark_f1k2_change_color_768: 用户图 关键点 1k*2 size 768， float32 + + ref_rgb_8uc3_change_color_768: 参考图， size 768， uint8 （0-255） + ref_matting_8uc3_change_color_768: 参考图， size 768， uint8 （0-255） + ref_landmark_f1k2_change_color_768: 参考图， 关键点 1k*2 size 768， float32 + + output： + + hair_gene_color_8uc3_768: 头发换颜色生成图， size 768， uint8 （0-255） + + """ + + # cv2.imshow("user_rgb_8uc3_change_color_768", user_rgb_8uc3_change_color_768) + # cv2.imshow("user_matting_8uc3_change_color_768", user_matting_8uc3_change_color_768) + # cv2.imshow("ref_rgb_8uc3_change_color_768", ref_rgb_8uc3_change_color_768) + # cv2.imshow("ref_matting_8uc3_change_color_768", ref_matting_8uc3_change_color_768) + + user_lab_user = cv2.cvtColor(user_rgb_8uc3_change_color_768, cv2.COLOR_BGR2LAB) + user_lab_user[:, :, 1] = 0 + user_lab_user[:, :, 2] = 0 + img_paf = user_rgb_8uc3_change_color_768.copy() + img_paf[user_matting_8uc3_change_color_768[:, :, 0].astype(bool)] = user_lab_user[user_matting_8uc3_change_color_768[:, :, 0].astype(bool)] + img_paf = (np.clip(img_paf, 0, 255)).astype(np.uint8) + + tmp = ref_matting_8uc3_change_color_768[:, :, 0] + tmp[tmp < 125] = 0 # tmp[tmp < 1] = 0 + + # mean_color = cv2.mean(refer_real, tmp) + condit_hair = (ref_rgb_8uc3_change_color_768 * (tmp[:, :, np.newaxis] / 255)).astype(np.uint8) + + # cv2.imshow("condit_hair", condit_hair) + # cv2.imshow("img_paf", img_paf) + # cv2.waitKey() + + condition = condit_hair # np.concatenate((user_mask, user_hair_mask), axis=2) + condition = (condition.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] + input_paf = (img_paf.astype(np.float32) / 255).transpose((2, 0, 1))[np.newaxis, :, :, :] + condition = torch.from_numpy(condition).to(self.device) + input_paf = torch.from_numpy(input_paf).to(self.device) + + # ******************* forward ******************* + with torch.no_grad(): + test_fake = self.net(input_paf, condition) + + hair_gene_color_8uc3_768 = (test_fake.detach()[0].to('cpu').numpy().transpose((1, 2, 0)) * 255).astype(np.uint8).copy() + + # hair_colr_mid_show = np.concatenate((condit_hair, img_paf, hair_gene_color_8uc3_768), axis=1) + # hair_colr_mid_show = cv2.resize(hair_colr_mid_show, (0, 0), fx=0.8, fy=0.8) + # cv2.imshow("hair_colr_mid_show", hair_colr_mid_show) + # cv2.waitKey(0) + + user_matting_mask_d = cv2.dilate(user_matting_8uc3_change_color_768, np.ones((30, 30), np.uint8)) + user_matting_mask_d_b = cv2.blur(user_matting_mask_d, (30, 30)) + + # hair_gene_color_8uc3_768 = user_rgb_8uc3_change_color_768 * (1 - user_matting_8uc3_change_color_768 / 255) + test_res * (user_matting_8uc3_change_color_768 / 255) + # hair_gene_color_8uc3_768 = (np.clip(hair_gene_color_8uc3_768, 0, 255)).astype(np.uint8) + + return hair_gene_color_8uc3_768, user_matting_mask_d_b + + +class BodySeg(): + def __init__(self, gpu_id=0): + self.gpu_id = gpu_id + self.device = torch.device(f"cuda:{gpu_id}") + + self.model = DeepLab() + model_io.load_model_by_path("./weights/human_seg_deeplabv3_288_384_sigmoid_msc.pth", self.model, gpu_id=gpu_id) + self.model.to(self.device) + self.model.eval() + self.output_img_size = [288, 384] + + self.resize_ratio = 2 + + def getM(self, center, angle, sx, sy): + angle = math.radians(angle) + alpha = math.cos(angle) + beta = math.sin(angle) + M = [[sx * alpha, sx * beta, (1 - sx * alpha) * center[0] - sx * beta * center[1]], + [-sy * beta, sy * alpha, sy * beta * center[0] + (1 - sy * alpha) * center[1]]] + return np.array(M) + def forward(self, frame): + # frame = cv2.imread(img_path).astype(np.float32) / 255 + + height, width = frame.shape[:2] + + x0, y0, x1, y1 = 0, 0, frame.shape[1], frame.shape[0] + + center_x, center_y = (x0 + x1) / 2, (y0 + y1) / 2 + random_scalex = min(self.output_img_size[0] * self.resize_ratio * 1.0 / (x1 - x0), + self.output_img_size[1] * self.resize_ratio * 1.0 / (y1 - y0)) + random_scaley = random_scalex + M = self.getM((center_x, center_y), 0, random_scalex, random_scaley) + M[:, 2] += [self.output_img_size[0] * self.resize_ratio / 2 - center_x, + self.output_img_size[1] * self.resize_ratio / 2 - center_y] + crop_img = cv2.warpAffine(frame, M, (self.output_img_size[0] * self.resize_ratio, self.output_img_size[1] * self.resize_ratio)) + crop_img_resize = cv2.resize(crop_img, (0, 0), fx=1 / self.resize_ratio, fy=1 / self.resize_ratio) + + input_tensor = torch.from_numpy(crop_img_resize.transpose([2, 0, 1])[np.newaxis]).to(self.device) + input_tensor = input_tensor.float() + output_tensor, _ = self.model(input_tensor) + output_tensor = F.interpolate(output_tensor, + size=[self.output_img_size[1] * self.resize_ratio, self.output_img_size[0] * self.resize_ratio], + mode='bilinear', align_corners=True) + # out_mask = torch.max(output_tensor[:1], 1)[1].detach().cpu().numpy().squeeze().astype(np.float32) + output_tensor = torch.sigmoid(output_tensor) + out_mask = output_tensor[0, 0].detach().cpu().numpy().squeeze().astype(np.float32) + input_output_show = np.concatenate((crop_img, np.repeat(out_mask[:, :, np.newaxis], 3, axis=2)), axis=1) + # cv2.imshow("input_output_show", input_output_show) + + M_inv = cv2.invertAffineTransform(M) + # mask_rawsize = cv2.warpAffine(out_mask, M_inv, (frame.shape[1], frame.shape[0])) + mask_rawsize = cv2.warpAffine(out_mask, M_inv, (frame.shape[1], frame.shape[0]), flags=cv2.INTER_NEAREST) + + mask_rawsize = np.repeat(mask_rawsize[:, :, np.newaxis], 3, axis=2) + return mask_rawsize + + def forward_kpts(self, frame, landmarks_1k): + + M = landmark_processor.get_transform_mat_bodyseg(landmarks_1k, output_size=self.output_img_size[0]*self.resize_ratio, ratio=0.7, offset=(0, 45)) + + crop_img = cv2.warpAffine(frame, M, (self.output_img_size[0] * self.resize_ratio, self.output_img_size[1] * self.resize_ratio)).astype(np.float32) / 255 + crop_img_resize = cv2.resize(crop_img, (0, 0), fx=1 / self.resize_ratio, fy=1 / self.resize_ratio) + + input_tensor = torch.from_numpy(crop_img_resize.transpose([2, 0, 1])[np.newaxis]).to(self.device) + input_tensor = input_tensor.float() + output_tensor, _ = self.model(input_tensor) + output_tensor = F.interpolate(output_tensor, + size=[self.output_img_size[1] * self.resize_ratio, self.output_img_size[0] * self.resize_ratio], + mode='bilinear', align_corners=True) + # out_mask = torch.max(output_tensor[:1], 1)[1].detach().cpu().numpy().squeeze().astype(np.float32) + output_tensor = torch.sigmoid(output_tensor) + out_mask = output_tensor[0, 0].detach().cpu().numpy().squeeze().astype(np.float32) + + # input_output_show = np.concatenate((crop_img, np.repeat(out_mask[:, :, np.newaxis], 3, axis=2)), axis=1) + # cv2.imshow("input_output_show", input_output_show) + # cv2.waitKey() + + M_inv = cv2.invertAffineTransform(M) + # mask_rawsize = cv2.warpAffine(out_mask, M_inv, (frame.shape[1], frame.shape[0])) + mask_rawsize = cv2.warpAffine(out_mask, M_inv, (frame.shape[1], frame.shape[0]), flags=cv2.INTER_NEAREST) + mask_rawsize = np.repeat(mask_rawsize[:, :, np.newaxis], 3, axis=2) + return mask_rawsize + diff --git a/hair_service_sd/requirements.txt b/hair_service_sd/requirements.txt new file mode 100644 index 0000000..1544336 --- /dev/null +++ b/hair_service_sd/requirements.txt @@ -0,0 +1,8 @@ +opencv-python +flask +flask_cors +gevent +pillow +cos-python-sdk-v5 +scipy +matplotlib \ No newline at end of file diff --git a/hair_service_sd/run_copy_cost_addcolor.py b/hair_service_sd/run_copy_cost_addcolor.py new file mode 100644 index 0000000..8623be9 --- /dev/null +++ b/hair_service_sd/run_copy_cost_addcolor.py @@ -0,0 +1,835 @@ +#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 +import pickle +from scipy.ndimage import label +from core.hairstyle_model import HairStyle_Model +from hairstyle_model_infer import HairStyle_Model_Infer +from flask import Flask, request, make_response, jsonify +from prepare_ref_hairstyle_data import prepare_single, prepare_single_color +from utils.callback import recall +from gen_super_image import webui_img2img, webui_img2img_diy, webui_super_res_img +from gevent import pywsgi +from flask_cors import CORS +from utils import enhance_hair +import configparser +from common.logger import config +from threading import Thread +# from meinheld import server +from queue import Queue +from concurrent.futures import ThreadPoolExecutor +from common.callback import * +from utils import call_hair_inter +from utils import landmark_processor +from change_color import process_infer, resize_pre_webui +from upload_oss import OSS_object +oss_2 = OSS_object() + +queue = Queue(1024) +app = Flask(__name__) +CORS(app, supports_credentials=True) +executor = ThreadPoolExecutor(8) +from multiprocessing import Process, Queue +from common.logger import config + +version = config.get('default', 'version') +if version == "local": + inference_services = { + 0: "57860" + } +else: + inference_services = { + 0: "57860", + 1: "57861", + 2: "57862", + } + + +hairstyle_process = HairStyle_Model(gpu=True,use_enhance=True) +hairstyle_process_infer = HairStyle_Model_Infer(gpu=True, use_enhance=False) +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 run_hairstyle(input_info): + ret_data = hairstyle_process.prepare_single_hairstyle(input_info) + hairId = osp.basename(input_info[0][-1]) + if ret_data['code'] == 200: + response = callback_hairstyle(hairId, ret_data['data'], True) + else: + response = callback_hairstyle(hairId, "", False) + print(response) + + +def run_hairstyle_v2(input_info): + ret_data = hairstyle_process.prepare_single_hairstyle_v2(input_info) + hairId = osp.basename(input_info[0][-1]) + print(ret_data) + + +def run_haircolor(sourceImage, dstdir,rgb): + ret_data = hairstyle_process.prepare_single_color_v2(sourceImage, dstdir, rgb) + if ret_data['code'] == 200: + response = callback_color(osp.basename(dstdir), ret_data['data'], True) + else: + response = callback_color(osp.basename(dstdir), "", False) + print(response) + # print('s') + + +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): + # 使用requests库下载图片 + response = requests.get(img_url) + if response.status_code == 200: + with open(tmp_dir, 'wb') as f: + f.write(response.content) + + 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 download_img_new(img_url, userId='ffff', 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:", img_url) + + r = requests.get(img_url) + # 写入图片 + with open(tmp_dir, "wb") as f: + f.write(r.content) + + img_type = imghdr.what(tmp_dir) + new_tmp_dir = tmp_dir.split(".")[0] + "." + img_type + shutil.move(tmp_dir, new_tmp_dir) + print("save path", new_tmp_dir) + + return new_tmp_dir, None + except Exception as e: + print(e) + return None,None + + +@app.route("/hairColor/v2", methods=['POST']) # 换发色新 +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() + input = request.json + + try: + print('\n + hairColor input :', input) + img_url = input['img'] + userId = input['userId'] + # color = input['colorId'] + rgb = input['rgb'] + try: + ratio = input['ratio'] + if ratio > 1: + ratio = ratio / 100 + except: + 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) + if len(color_id) != 3: + return make_response( + jsonify({'msg': 'color解析错误, must be [r,g,b]', 'result': '', 'umd': '', 'state': -1}), 400) + 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_id, + + req_id + '.jpg')) + print("---------------upload img:", time.time() - start_time3) + return make_response(jsonify( + {'msg': 'success', 'result': ret_url, 'umd': "", + 'state': 0}), 200) + else: + return make_response( + jsonify({'msg': '算法解析错误', 'result': '', 'umd': '', 'state': -1}), 400) + + except Exception as e: + print(e) + return make_response( + jsonify({'msg': '算法解析错误', 'result': '', 'umd': '', 'state': -1}), 400) + + +@app.route("/api/swapHair/v1", methods=['POST']) # 换发型新 +def change_hairstyle_v4(): + input_info = request.json + print(f"____swapHair____{input_info}") + hairstyle_dir = config.get('default', 'hairstyleDir') # 初始化配置路径 + user_dir = config.get('default', 'userDir') + ref_img_dir = config.get('default', 'refImgDir') + res_dir = config.get('default', 'res_dir') + hair_template_material_dir = config.get('default', 'hair_template_material_dir') + train_dir = config.get('default', 'train_dir') + userInfo_dir = config.get('default', 'userInfo_dir') + start_time_all = time.time() # 记录总耗时 + ret = { + "state": -1, + "msg": "fail", + "data": "", + "task_id": "" + } + # 功能1：获取请求参数 + try: + start_time = time.time() + hair_id = input_info['hair_id'] + task_id = input_info['task_id'] + is_hr_value = input_info['is_hr'] + is_hr = is_hr_value.lower() == "true" + #is_hr = False + hair_material_dir = os.path.join(train_dir, hair_id) + ret['task_id'] = task_id + user_img_url = input_info['user_img_path'] + user_img_path ,md5_img= download_img(user_img_url) + #shutil.copy(user_img_path,user_img_path.replace("/home/data/hair/data/tmp","/home/data/hair/data/userImage")) + #user_img_path = user_img_path.replace("/home/data/hair/data/tmp","/home/data/hair/data/userImage") + #user_img_path = user_img_url + # userId = input_info['userId'] + userId = 'test' + print(f"功能1：获取请求参数，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + print(e) + hairstyle_process.logger_call.error(f"时间:{datetime.now()}，参数错误：{e}") + ret['msg'] = '参数错误' + return make_response(jsonify(ret), 400) + # 功能2：加载用户图像 + try: + start_time = time.time() + user_img_name = os.path.basename(user_img_path) + new_user_img_path = os.path.join(user_dir, user_img_name) + if os.path.exists(new_user_img_path): + os.remove(new_user_img_path) + shutil.copy(user_img_path, new_user_img_path) + print(f"功能2：加载用户图像，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + print(e) + ret['msg'] = '用户图像加载失败' + return make_response(jsonify(ret), 400) + # 功能3：加载发型模板图像 + try: + start_time = time.time() + train_img_save_dir = os.path.join(train_upload_dir, hair_id) + hair_name_lists = os.listdir(train_img_save_dir) + template_ref_hair_name = next( + (name for name in hair_name_lists if "first##" in name), None) + if not template_ref_hair_name: + raise FileNotFoundError("未找到发型模板") + template_ref_hair_path = os.path.join(train_img_save_dir, template_ref_hair_name) + print(f"功能3：加载发型模板图像，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + print(e) + ret['msg'] = '发型模板加载失败' + return make_response(jsonify(ret), 400) + # 功能4：检查遮挡眼睛 + try: + start_time = time.time() + for root, dirs, files in os.walk(os.path.join(hair_template_material_dir, hair_id)): + hair_mask_img_path = next( + (os.path.join(root, file) for file in files if file.startswith("first##") and file.endswith("_matting.png")), + None + ) + if hair_mask_img_path: + break + if not hair_mask_img_path: + raise FileNotFoundError("未找到遮挡信息文件") + hair_img_pkl_path = hair_mask_img_path[:-12] + ".pkl" + with open(hair_img_pkl_path, 'rb') as fp: + data = pickle.load(fp) + hair_pt1k = data.get('human_pt1k') + print(f"功能4：检查遮挡眼睛，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + print(e) + ret['msg'] = '检查遮挡失败' + return make_response(jsonify(ret), 400) + # 功能5：生成发型材质文件 + try: + start_time = time.time() + long_flag = os.path.exists(os.path.join(hair_material_dir, "long.txt")) + material_save_path = os.path.join(hairstyle_dir, hair_id) + if not os.path.exists(material_save_path): + shutil.copy(template_ref_hair_path, os.path.join(ref_img_dir, template_ref_hair_name)) + prepare_single(os.path.join(ref_img_dir, template_ref_hair_name), material_save_path, long_flag) + print(f"功能5：生成发型材质文件，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + print(e) + ret['msg'] = '生成发型材质失败' + return make_response(jsonify(ret), 400) + # 功能6：推理发型 + try: + start_time = time.time() + origin_img = cv2.imread(new_user_img_path) + ref_img_path = os.path.join(material_save_path, "ref_rgb_8uc3_768.png") + if not os.path.exists(ref_img_path): + raise FileNotFoundError("推理目标图像不存在") + ref_img = cv2.imread(ref_img_path) + with torch.no_grad(): + img_res, status, _, landmarks_origin_img_1k, isEyeOccluded = hairstyle_process.infer_hairstyle_diy_jy( + origin_img, ref_img, os.path.join(userInfo_dir, task_id), f"{task_id}.png" + ) + if status != 0: + raise RuntimeError("发型推理失败") + print(f"功能6：推理发型，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + print(e) + ret['msg'] = '推理发型失败' + return make_response(jsonify(ret), 400) + + try: + # 功能7：判断头发遮挡眼睛 + start_time = time.time() + user_material_dir = os.path.join(userInfo_dir, task_id) + user_img_txt_path = os.path.join(user_material_dir, "kpt_1k.txt") + user_mask_img_path = os.path.join(user_material_dir, "user_orig_mask.png") + user_pt1k = np.loadtxt(user_img_txt_path) + result, user_landmarks_origin_img_137 = hairstyle_process.check_hair_covering_eyes(hair_pt1k,hair_mask_img_path,new_user_img_path, user_pt1k,user_mask_img_path) + if result == 1: print("result == 1: hair_covering_eyes") + else: print("result == 0") + print(f"功能7：判断头发遮挡眼睛，耗时：{time.time() - start_time:.3f}s") + result=0 + if result == 1: + print("lead to diy") + ret_url = change_hairstyle_use_diy(new_user_img_path, template_ref_hair_path, userId, task_id) + if ret_url == "": return make_response(jsonify(ret), 200) + else: + ret["msg"] = 'success' + ret['state'] = 0 + ret['data'] = ret_url + return make_response(jsonify(ret), 200) + else: + print("lead te change hairstyle") + + # 功能8：处理发型区域 + start_time = time.time() + hair_matting_path = os.path.join(user_material_dir, "hair_mask_2.png") + new_matting = cv2.imread(hair_matting_path, cv2.IMREAD_GRAYSCALE) + result_img = img_res + user_orig_mask_path = os.path.join(user_material_dir, "user_orig_mask.png") + origin_matting = cv2.imread(user_orig_mask_path, cv2.IMREAD_GRAYSCALE) + useTitorMask = True + if useTitorMask: + # if is_hr and max(origin_img.shape[:2]) < 2000: + # start_time = time.time() + # scale_ratio = 2000 / max(origin_img.shape[:2]) + # result_img = webui_super_res_img(result_img, scale_ratio) + # origin_img = cv2.resize(origin_img, (result_img.shape[1], result_img.shape[0]), + # interpolation=cv2.INTER_LANCZOS4) + # print(f"功能：resize 2000，耗时：{time.time() - start_time:.3f}s") + + dst_size = (1152, 1536) if is_hr else (576, 768) # 处理是否高清模式的分辨率 + print(dst_size) + box_info = hairstyle_process.get_body_info(img_res) # 获取头发处理的局部区域图像 + # cv2.imwrite(f"{task_id}_img_res", img_res) + box_w, box_h = box_info[2] - box_info[0], box_info[3] - box_info[1] + print(f"box_w, box_h : {box_w},{box_h}") + h = max(dst_size[1] , box_h) + w = max(dst_size[0] , box_w) + + 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) + crop_result = landmark_processor.high_quality_warpAffine(img_res, M, dst_size) + # crop_result = landmark_processor.high_quality_warpAffine(img_res, M, (w, h)) + + origin_matting = cv2.imread(user_orig_mask_path, cv2.IMREAD_GRAYSCALE) + new_matting = cv2.imread(hair_matting_path, cv2.IMREAD_GRAYSCALE) + left_point = user_landmarks_origin_img_137[14] # (x1, y1) 获取刘海区域的圆心和半径 + right_point = user_landmarks_origin_img_137[7] # (x2, y2) + center_x = int((left_point[0] + right_point[0]) // 2) + center_y = int((left_point[1] + right_point[1]) // 2) + radius = int(np.sqrt((right_point[0] - left_point[0]) ** 2 + (right_point[1] - left_point[1]) ** 2) / 2) + h, w = origin_matting.shape # 初始化刘海区域为空（全黑色） + part_bangs = np.zeros((h, w), dtype=np.uint8) + cv2.circle(part_bangs, (center_x, center_y), radius, 255, -1) # 在 part_bangs 上绘制圆作为刘海区域 圆形区域设为白色 + no_bang_result = cv2.subtract(origin_matting, part_bangs) # 去除刘海区域 + matting_merge = np.max( np.stack([no_bang_result, new_matting], axis=2), axis=2).astype(np.uint8)# 合并 origin_matting 和 new_matting + crop_matting = cv2.warpAffine(matting_merge, M, dst_size) + # crop_matting = cv2.warpAffine(matting_merge, M, (w, h)) + mask = (crop_matting > 10).astype(np.float32) + if not is_hr: + mask_dilate = cv2.dilate(mask, np.ones((3, 9), np.uint8)) + else: + mask_dilate = cv2.dilate(mask, np.ones((6, 18), np.uint8)) + final_img = crop_result + mask_dilate = np.clip(mask_dilate * 255, 0, 255).astype(np.uint8) + else: + ## jy_version + matting_merge = np.concatenate([origin_matting[:, :, np.newaxis], new_matting[:, :, np.newaxis]],axis=2) + matting_merge = np.max(matting_merge, axis=2) + mask = cv2.resize(matting_merge, (img_res.shape[1], img_res.shape[0])) + mask_dilate = cv2.dilate(mask, np.ones((3, 11), np.uint8)) + print(mask_dilate.shape) + save_mask_path = osp.join(user_material_dir, 'webui.png') + cv2.imwrite(save_mask_path, mask_dilate) + crop_img, crop_mask, M = resize_pre_webui(img_res, mask_dilate) + final_img = crop_img + mask_dilate = crop_mask + # cv2.imwrite(f"{task_id}_finalimg.jpg", final_img) + cv2.imwrite(f"{task_id}.jpg", mask_dilate) + config_json_path = os.path.join(material_save_path, "config.json") + with open(config_json_path, "r") as f: + config_json_content = json.load(f) + in_gender = config_json_content["gender"] + print("in_gender:", in_gender) + images_dir = os.path.join(hair_material_dir, "images") + txt_dir = os.path.join(images_dir, os.listdir(images_dir)[0]) + txt_path = glob.glob(txt_dir + '/*.txt')[0] + with open(txt_path, 'r') as f: + p_tag = f.readline() + if "titor hairstyle, faceless, no human, gray background, simple background" in p_tag: + p_tag = p_tag[p_tag.find("simple background, ") + len("simple background, "):] + else: + p_tag = "" + denoising_strength = 0.6 + print(f"功能8：处理发型区域，耗时：{time.time() - start_time:.3f}s") + + # 功能9：webui + start_time = time.time() + sd_result = webui_img2img(img=final_img, mask_img=mask_dilate, in_gender=in_gender, task_id=task_id, + hair_id=hair_id, lora_material_path=hair_material_dir, tag=p_tag, is_hr=is_hr, + denoising_strength=denoising_strength, inference_port="57860") + print(f"功能：webui，耗时：{time.time() - start_time:.3f}s") + + # 功能：后处理并上传结果 + start_time = time.time() + dst_path = os.path.join(res_dir, f"{task_id}.png") + if not os.path.exists(os.path.dirname(dst_path)): + os.makedirs(os.path.dirname(dst_path)) + origin_img_final = origin_img.copy() + M_inv = cv2.invertAffineTransform(M) + cv2.warpAffine(sd_result, M_inv, (origin_img.shape[1], origin_img.shape[0]), + dst=origin_img_final, borderMode=cv2.BORDER_TRANSPARENT, flags=cv2.INTER_LANCZOS4) + cv2.imwrite(dst_path, origin_img_final) + ret_url = oss_2.upload_file(dst_path, f"hair_mz/images/hairstyle/{hair_id}/{uuid4()}.jpg") + ret.update({"msg": "success", "state": 0, "data": ret_url}) + print(f"功能9：后处理并上传结果，耗时：{time.time() - start_time:.3f}s") + + except Exception as e: + print(e) + ret['msg'] = '推理发型失败' + return make_response(jsonify(ret), 400) + + # 打印总耗时 + print(f"all总耗时：{time.time() - start_time_all:.3f}s") + return make_response(jsonify(ret), 200) + +def change_hairstyle_use_diy(dir_user, dir_tar, userId, task_id): + print(f"____diy____") + userInfo_dir = config.get('default', 'userInfo_dir') + user_img_res_dir = config.get('default', 'res_dir') + + diy_start_time_all = time.time() + ret_url = "" + + # diy功能1：加载图像 + try: + start_time = time.time() + origin_img = cv2.imread(dir_user) + ref_img = cv2.imread(dir_tar) + + if origin_img is None or ref_img is None: + raise ValueError("图片可能损坏") + print(f"diy功能1：加载图像，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + hairstyle_process.logger_call.error(f"时间:{datetime.now()}，加载图像失败：{e}") + return make_response(jsonify({'msg': '图片可能损坏', 'data': "", 'umd': '', 'state': 1}), 200) + + # diy功能2：发型推理 + try: + start_time = time.time() + user_material_dir = os.path.join(userInfo_dir, task_id) + basefolder, imgname = osp.split(dir_user) + _, target_name = osp.split(dir_tar) + + with torch.no_grad(): + img_res, status, in_gender = hairstyle_process.infer_hairstyle_diy(origin_img, ref_img, user_material_dir, target_name) + + if status != 0: + raise RuntimeError("发型推理失败") + print(f"diy功能2：发型推理，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + hairstyle_process.logger_call.error(f"时间:{datetime.now()}，发型推理失败：{e}") + return "" + + + # diy功能3：保存发型结果图像 + try: + start_time = time.time() + img_res_dir = os.path.join(basefolder, 'diy') + os.makedirs(img_res_dir, exist_ok=True) + res_save_path = os.path.join(img_res_dir, f"{imgname[:-4]}-{target_name[:-4]}.jpg") + cv2.imwrite(res_save_path, img_res) + print(f"diy功能3：保存发型结果图像，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + hairstyle_process.logger_call.error(f"时间:{datetime.now()}，保存发型结果图像失败：{e}") + return "" + + # diy功能4：生成 mask + try: + start_time = time.time() + hair_mask_path = osp.join(user_material_dir, 'hair_mask_2.png') + orig_model_mask_path = osp.join(user_material_dir, 'user_orig_mask.png') + + new_matting = cv2.imread(hair_mask_path, cv2.IMREAD_GRAYSCALE) + origin_matting = cv2.imread(orig_model_mask_path, cv2.IMREAD_GRAYSCALE) + + matting_merge = np.max( + np.stack([origin_matting[:, :, np.newaxis], new_matting[:, :, np.newaxis]], axis=2), axis=2 + ).astype(np.uint8) + + mask_dilate = cv2.dilate(cv2.resize(matting_merge, (img_res.shape[1], img_res.shape[0])), np.ones((3, 11), np.uint8)) + save_mask_path = osp.join(user_material_dir, 'webui.png') + cv2.imwrite(save_mask_path, mask_dilate) + print(f"diy功能4：生成 mask，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + hairstyle_process.logger_call.error(f"时间:{datetime.now()}，生成 mask 失败：{e}") + return "" + + # diy功能5：增强图像 + try: + start_time = time.time() + crop_img, crop_mask, M = resize_pre_webui(img_res, mask_dilate) + crop_img_path = os.path.join(user_img_res_dir, f"{task_id}_crop.png") + crop_mask_path = os.path.join(user_img_res_dir, f"{task_id}_mask.png") + + cv2.imwrite(crop_img_path, crop_img) + cv2.imwrite(crop_mask_path, crop_mask) + + out = call_hair_inter.call_hair_enhance(crop_img_path, crop_mask_path, task_id, in_gender) + out_path = out["result"] + enhanced_img = cv2.imread(out_path) + print(f"diy功能5：增强图像，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + hairstyle_process.logger_call.error(f"时间:{datetime.now()}，增强图像失败：{e}") + return "" + + # diy功能6：恢复图像并保存 + try: + start_time = time.time() + origin_img_final = origin_img.copy() + 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) + + final_dst_path = os.path.join(user_img_res_dir, f"{task_id}_res.png") + cv2.imwrite(final_dst_path, origin_img_final) + print(f"diy功能6：恢复图像并保存，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + hairstyle_process.logger_call.error(f"时间:{datetime.now()}，恢复图像失败：{e}") + return "" + + # diy功能7：上传图像到 OSS + try: + start_time = time.time() + ret_url = oss_2.upload_file(final_dst_path, f"hair_mz/images/hairstyle/{userId}/{task_id}.jpg") + print(f"diy功能7：上传图像到 OSS，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + hairstyle_process.logger_call.error(f"时间:{datetime.now()}，上传图像到 OSS 失败：{e}") + return "" + + # 打印总耗时 + print(f"diy总耗时：{time.time() - diy_start_time_all:.3f}s") + return ret_url + + + + + +@app.route("/api/uploadHair/v1", methods=['POST']) # 新增发型模板新 +def add_new_hairstyle_v2(): + input_data = request.json + print(f"____uploadHair____{input_data}") + imgLists = input_data['img_lists'] + input_info = [] + hairID = input_data['hair_id'] + # cover_url = input_data['coverImg'] + userId = 'houtai' + hairstyle_process.logger_call.info('time:{}, fff, {}'.format(datetime.now().strftime("%Y-%m-%d, %H:%M:%S"), input_data)) + + if len(imgLists) == 0: + hairstyle_process.logger_call.error( + 'time:{}, 未选择图像传入, {}'.format(datetime.now().strftime("%Y-%m-%d, %H:%M:%S"), input_data)) + + return make_response(jsonify({'msg': '未选择图像传入', 'data': "", 'state': 1}), 200) + + for index,img_url in enumerate(imgLists): + + img_path ,md5_img= download_img(img_url, userId) + #img_path =img_url + if index == 0: + cover_path=img_path + if not os.path.exists(img_path): + hairstyle_process.logger_call.error('time:{}, 图像读取错误'.format(datetime.now().strftime("%Y-%m-%d, %H:%M:%S"))) + return make_response(jsonify({'msg': '图像读取错误', 'data': '', 'state': 1}), 200) + + # 存储原始图片 + dst_dir = osp.join(train_upload_dir, hairID) + # 存储用户参考图 + ref_user_save_dir = ref_user_dir + # 生成训练素材图 + train_material_save_dir = os.path.join(train_save_dir, hairID) + os.makedirs(dst_dir, exist_ok=True) + + # 拷贝训练数据 + file_name = img_path[img_path.rfind("/") + 1:] + new_file_path = os.path.join(dst_dir, file_name) + shutil.copy(img_path, new_file_path) + + input_info.append([new_file_path, cover_path, 0, dst_dir, ref_user_save_dir, train_material_save_dir]) + + + args = [input_info] + executor.submit(lambda p: run_hairstyle_v2(*p), args) + return make_response(jsonify({'msg': 'OK', 'data': "", 'state': 0}), 200) + + +@app.route("/api/hair/trainCallBack", methods=['POST']) # 发型训练回调 +def train_hair_callback(): + input_info = request.json + print(f"____trainCallBack____{input_info}") + ret_data = { + "task_id": "", + "state": -1, + "msg": "failed", + } + + hairstyle_dir = config.get('default', 'hairstyleDir') + ref_img_dir = config.get('default', 'refImgDir') + hair_template_material_dir = config.get('default', 'hair_template_material_dir') + + # get request params + try: + req_id = input_info['task_id'] + hair_id = input_info['hair_id'] + state = input_info['state'] + msg = input_info['msg'] + print("trainCallBack input_info: ", input_info) + except Exception as e: + print(e) + hairstyle_process.logger_call.error(f'params error, input is {input_info} error is {e}') + return make_response(jsonify(ret_data), 400) + + # gen hair material + hair_upload_save_dir = os.path.join(train_upload_dir, hair_id) + print("hair_upload_save_dir: ", hair_upload_save_dir) + + hair_name_lists = os.listdir(hair_upload_save_dir) + for hair_name in hair_name_lists: + if "first##" in hair_name and (hair_name.endswith(".png") or hair_name.endswith(".jpg") or hair_name.endswith(".jpeg")): + template_ref_hair_name = hair_name + template_ref_hair_path = os.path.join(hair_upload_save_dir, hair_name) + break + + new_hair_ref_img_path = os.path.join(ref_img_dir, template_ref_hair_name) + print("new_hair_ref_img_path: ", new_hair_ref_img_path) + shutil.copy(template_ref_hair_path, new_hair_ref_img_path) + + material_save_path = os.path.join(hairstyle_dir, hair_id) + print("first_hair_material_save_path :", material_save_path) + + long_flag = False + prepare_single(new_hair_ref_img_path, material_save_path, long_flag) + + # 训练完成，回调天津后端接口 + status = -1 + try: + if state == 0: + ret_data = {} + ret_data['code'] = 200 + ret_data['data'] = [] + ret_data['msg'] = 'success' + + cover_img_path = new_hair_ref_img_path + + ret_url = oss_2.upload_file(cover_img_path, "hair_mz/images/hairstyle/{}/{}".format(hair_id, + str(uuid4()) + '.jpg')) + print('recall success') + response, status = callback_hairstyle(req_id, state, msg, hair_id) + #response, status = callback_hairstyle(req_id, state, hair_id) + print(response) + else: + response, status = callback_hairstyle(req_id, state, msg, hair_id) + #response, status = callback_hairstyle(req_id, state, hair_id) + print(response) + + except Exception as e: + print(e) + print('recall failed, ConnectionError') + + if status != 0: + ret_data['state'] = -1 + ret_data['msg'] = 'failed' + ret_data['task_id'] = input_info['task_id'] + return make_response(jsonify(ret_data), 400) + + ret_data['state'] = 0 + ret_data['msg'] = 'success' + ret_data['task_id'] = input_info['task_id'] + return make_response(jsonify(ret_data), 200) + +@app.route("/hairEnhance/v1", methods=['POST']) # 发型增强 +def hair_enhance(): + input = request.json + print(f"____hairEnhance____{input}") + userInfo_dir = config.get('default', 'userInfo_dir') + try: + print('\n + hairEnhance input :', input) + img_path = input['img_path'] + mask_path = input['mask_path'] + req_id = input['req_id'] + in_gender = input['gender'] + + # 发型图 + final_img = cv2.imread(img_path) + print("final_img shape:", final_img.shape) + + # mask图 + if mask_path == "": + width, height = final_img.shape[1], final_img.shape[0] + mask_dilate = np.full((height, width, 3), 255, dtype=np.uint8) + print("mask_dilate shape:", mask_dilate.shape) + else: + mask_dilate = cv2.imread(mask_path) + + # cv2.imshow("final_img", final_img) + # cv2.imshow("mask_dilate", mask_dilate) + # cv2.waitKey(0) + + + sd_result = webui_img2img_diy(img=final_img, mask_img=mask_dilate, in_gender=in_gender, task_id="", tag="", inference_port="57860") + + material_dir = os.path.join(userInfo_dir, req_id) + if not os.path.exists(material_dir): + os.mkdir(material_dir) + sd_save_path = os.path.join(material_dir, str(uuid4()) + ".png") + cv2.imwrite(sd_save_path, sd_result) + + # cv2.imshow("inpaint", sd_result) + + + _, img_name = osp.split(img_path) + _, mask_name = osp.split(mask_path) + + # todo can be del + # temp2 = os.path.join("/home/data/hair/data/test_tmp", str(uuid4()) + ".png") + # cv2.imwrite(temp2, sd_result) + # + # cv2.waitKey(0) + + # ret_url = hairstyle_process.oss2.upload_file(sd_save_path, + # "hair_mz/images/{}/{}/{}".format(user_id, img_name[:-4], + # '{}-{}.jpg'.format( + # img_name[:-4], + # mask_name[:-4]))) + tmp_info = { + 'umd': "", + 'maskinfo': 'res_matting_mask_ori_fix.png', + 'hairId': '', + } + return make_response(jsonify({'msg': 'success', 'result': sd_save_path, 'umd': '', 'state': 0}), + 200) + + except Exception as e: + print(e) + + + +if __name__ == '__main__': + # from werkzeug.middleware.proxy_fix import ProxyFix + # server.listen(("0.0.0.0", 7393)) + # server.run(app) + # app.wsgi_app = ProxyFix(app.wsgi_app) + # app.run() + + server = pywsgi.WSGIServer(('0.0.0.0', 8801), app) + server.serve_forever() diff --git a/hair_service_sd/run_copy_cost_colorb64.py b/hair_service_sd/run_copy_cost_colorb64.py new file mode 100644 index 0000000..ebcf0a2 --- /dev/null +++ b/hair_service_sd/run_copy_cost_colorb64.py @@ -0,0 +1,830 @@ +#coding:utf-8 +import traceback +from uuid import uuid4 +import imghdr +import traceback +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 cv2 +from datetime import datetime +import glob +import numpy as np +import pickle +from core.hairstyle_model import HairStyle_Model +from hairstyle_model_infer import HairStyle_Model_Infer +from flask import Flask, request, make_response, jsonify +from prepare_ref_hairstyle_data import prepare_single, prepare_single_color +from gen_super_image import webui_img2img, webui_img2img_diy +from gevent import pywsgi +from flask_cors import CORS +from common.logger import config +from queue import Queue +from concurrent.futures import ThreadPoolExecutor +from common.callback import * +from utils import call_hair_inter +from utils import landmark_processor +from change_color import process_infer, resize_pre_webui +from upload_oss import OSS_object +oss_2 = OSS_object() + +queue = Queue(1024) +app = Flask(__name__) +CORS(app, supports_credentials=True) +executor = ThreadPoolExecutor(8) +from multiprocessing import Process, Queue +from common.logger import config + +version = config.get('default', 'version') +if version == "local": + inference_services = { + 0: "57860" + } +else: + inference_services = { + 0: "57860", + # 1: "57861", + # 2: "57862", + } + +user_img_save_dir = config.get('default', 'userDir') +if not os.path.exists(user_img_save_dir): + os.makedirs(user_img_save_dir) +user_img_tmp_dir = config.get('default', 'tmp_dir') +if not os.path.exists(user_img_tmp_dir): + os.makedirs(user_img_tmp_dir) +user_img_res_dir = config.get('default', 'res_dir') +if not os.path.exists(user_img_res_dir): + os.makedirs(user_img_res_dir) +ref_user_dir = config.get('default', 'ref_user_dir') +if not os.path.exists(ref_user_dir): + os.makedirs(ref_user_dir) +train_save_dir = config.get('default', 'train_dir') +if not os.path.exists(train_save_dir): + os.makedirs(train_save_dir) +hair_template_material_dir = config.get('default', 'hair_template_material_dir') +if not os.path.exists(hair_template_material_dir): + os.makedirs(hair_template_material_dir) +ref_color_dir = config.get('default', 'ref_color') +if not os.path.exists(ref_color_dir): + os.makedirs(ref_color_dir) +ref_color_imgs_dir = config.get('default', 'ref_color_img') +if not os.path.exists(ref_color_imgs_dir): + os.makedirs(ref_color_imgs_dir) +train_upload_dir = config.get('default', 'upload_train_dir') +if not os.path.exists(train_upload_dir): + os.makedirs(train_upload_dir) + +hairstyle_process = HairStyle_Model(gpu=True,use_enhance=True) +hairstyle_process_infer = HairStyle_Model_Infer(gpu=True, use_enhance=False) + + + +def run_hairstyle(input_info): + ret_data = hairstyle_process.prepare_single_hairstyle(input_info) + hairId = osp.basename(input_info[0][-1]) + if ret_data['code'] == 200: + response = callback_hairstyle(hairId, ret_data['data'], True) + else: + response = callback_hairstyle(hairId, "", False) + print(response) + + +def run_hairstyle_v2(input_info): + ret_data = hairstyle_process.prepare_single_hairstyle_v2(input_info) + hairId = osp.basename(input_info[0][-1]) + print(ret_data) + + +def run_haircolor(sourceImage, dstdir,rgb): + ret_data = hairstyle_process.prepare_single_color_v2(sourceImage, dstdir, rgb) + if ret_data['code'] == 200: + response = callback_color(osp.basename(dstdir), ret_data['data'], True) + else: + response = callback_color(osp.basename(dstdir), "", False) + print(response) + # print('s') + + +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) + + print(img_url) + + download_success = False + for i in range(3): + # 使用requests库下载图片 + response = requests.get(img_url) + if response.status_code == 200: + with open(tmp_dir, 'wb') as f: + f.write(response.content) + + if osp.exists(tmp_dir) and osp.getsize(tmp_dir) > 0: + download_success = True + break + + if download_success: + img_type = tmp_dir.split('.')[-1] + 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 + + +@app.route("/hairColor/v2", methods=['POST']) # 换发色新 +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() + input = request.json + + try: + img_b64 = input.get('img', '') + userId = input['userId'] + rgb = input['rgb'] + output_format = input.get('output_format', 'url') + + # 添加参数校验 + if not img_b64: + return make_response(jsonify({'msg': 'img参数不能为空', 'state': -1}), 400) + if not userId: + return make_response(jsonify({'msg': 'userId参数不能为空', 'state': -1}), 400) + if len(rgb) != 3: + return make_response(jsonify({'msg': 'rgb参数必须为长度为3的数组', 'state': -1}), 400) + + try: + ratio = input['ratio'] + if ratio > 1: + ratio = ratio / 100 + except: + ratio = 0.9 + + if 'data:image/' in img_b64: + img_b64 = img_b64.split(',')[1] + img_data = base64.b64decode(img_b64) + nparr = np.frombuffer(img_data, np.uint8) + img_np = cv2.imdecode(nparr, cv2.IMREAD_COLOR) + img_path = osp.join(user_img_tmp_dir, userId + '#' + str(datetime.now().strftime("%Y%m%d%H%M%S") +'.jpg')) + cv2.imwrite(img_path, img_np) + else: + img_path, _ = download_img(img_b64, userId) + if not img_path: + return make_response(jsonify({'msg': '图片下载失败', 'state': -1}), 400) + + 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) + if len(rgb) != 3: + return make_response( + jsonify({'msg': 'color解析错误, must be [r,g,b]', 'result': '', 'umd': '', 'state': -1}), 400) + 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 + ".jpg") + + 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): + if output_format == 'base64': + with open(res_path, 'rb') as f: + img_data = f.read() + base64_str = base64.b64encode(img_data).decode('utf-8') + return make_response(jsonify( + {'msg': 'success', 'result': base64_str, 'umd': "", + 'state': 0}), 200) + else: + ret_url = hairstyle_process.oss2.upload_file(res_path, + "hair_mz/images/hairstyle/{}/{}".format(color_id, + req_id + '.jpg')) + print("---------------upload img:", time.time() - start_time3) + return make_response(jsonify( + {'msg': 'success', 'result': ret_url, 'umd': "", + 'state': 0}), 200) + else: + return make_response( + jsonify({'msg': '算法解析错误', 'result': '', 'umd': '', 'state': -1}), 400) + + except Exception as e: + print(e) + return make_response( + jsonify({'msg': '算法解析错误', 'result': '', 'umd': '', 'state': -1}), 400) + + +@app.route("/api/swapHair/v1", methods=['POST']) # 换发型新 +def change_hairstyle_v4(): + input_info = request.json + print(f"____swapHair____{input_info}") + hairstyle_dir = config.get('default', 'hairstyleDir') # 初始化配置路径 + user_dir = config.get('default', 'userDir') + ref_img_dir = config.get('default', 'refImgDir') + res_dir = config.get('default', 'res_dir') + hair_template_material_dir = config.get('default', 'hair_template_material_dir') + train_dir = config.get('default', 'train_dir') + userInfo_dir = config.get('default', 'userInfo_dir') + start_time_all = time.time() # 记录总耗时 + ret = { + "state": -1, + "msg": "fail", + "data": "", + "task_id": "" + } + clean_temp_data = True + # 新增输出格式参数 + output_format = input_info.get('output_format', 'url') + + # 功能1：获取请求参数 + try: + start_time = time.time() + hair_id = input_info['hair_id'] + task_id = input_info['task_id'] + is_hr_value = input_info['is_hr'] + is_hr = is_hr_value.lower() == "true" + hair_material_dir = os.path.join(train_dir, hair_id) + ret['task_id'] = task_id + + # 处理图片输入（支持base64或URL） + user_img_input = input_info['user_img_path'] + if 'data:image/' in user_img_input: + img_b64 = user_img_input.split(',')[1] + img_data = base64.b64decode(img_b64) + nparr = np.frombuffer(img_data, np.uint8) + img_np = cv2.imdecode(nparr, cv2.IMREAD_COLOR) + user_img_path = osp.join(user_img_tmp_dir, task_id + '_' + str(datetime.now().strftime("%Y%m%d%H%M%S")) + '.jpg') + cv2.imwrite(user_img_path, img_np) + md5_img = None + else: + user_img_path, md5_img = download_img(user_img_input) + + userId = 'test' + print(f"功能1：获取请求参数，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + print(e) + hairstyle_process.logger_call.error(f"时间:{datetime.now()}，参数错误：{e}") + ret['msg'] = '参数错误' + return make_response(jsonify(ret), 400) + + # 功能2：加载用户图像 + try: + start_time = time.time() + user_img_name = os.path.basename(user_img_path) + new_user_img_path = os.path.join(user_dir, user_img_name) + if os.path.exists(new_user_img_path): + os.remove(new_user_img_path) + shutil.copy(user_img_path, new_user_img_path) + print(f"功能2：加载用户图像，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + print(e) + ret['msg'] = '用户图像加载失败' + return make_response(jsonify(ret), 400) + # 功能3：加载发型模板图像 + try: + start_time = time.time() + train_img_save_dir = os.path.join(train_upload_dir, hair_id) + hair_name_lists = os.listdir(train_img_save_dir) + template_ref_hair_name = next( + (name for name in hair_name_lists if "first##" in name), None) + if not template_ref_hair_name: + raise FileNotFoundError("未找到发型模板") + template_ref_hair_path = os.path.join(train_img_save_dir, template_ref_hair_name) + print(f"功能3：加载发型模板图像，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + print(e) + ret['msg'] = '发型模板加载失败' + return make_response(jsonify(ret), 400) + # 功能4：检查遮挡眼睛 + try: + start_time = time.time() + for root, dirs, files in os.walk(os.path.join(hair_template_material_dir, hair_id)): + hair_mask_img_path = next( + (os.path.join(root, file) for file in files if file.startswith("first##") and file.endswith("_matting.png")), + None + ) + if hair_mask_img_path: + break + if not hair_mask_img_path: + raise FileNotFoundError("未找到遮挡信息文件") + hair_img_pkl_path = hair_mask_img_path[:-12] + ".pkl" + with open(hair_img_pkl_path, 'rb') as fp: + data = pickle.load(fp) + hair_pt1k = data.get('human_pt1k') + print(f"功能4：检查遮挡眼睛，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + print(e) + ret['msg'] = '检查遮挡失败' + return make_response(jsonify(ret), 400) + # 功能5：生成发型材质文件 + try: + start_time = time.time() + long_flag = os.path.exists(os.path.join(hair_material_dir, "long.txt")) + material_save_path = os.path.join(hairstyle_dir, hair_id) + if not os.path.exists(material_save_path): + ret['msg'] = '生成发型材质失败' + return make_response(jsonify(ret), 400) + # shutil.copy(template_ref_hair_path, os.path.join(ref_img_dir, template_ref_hair_name)) + # prepare_single(os.path.join(ref_img_dir, template_ref_hair_name), material_save_path, long_flag) + print(f"功能5：生成发型材质文件，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + print(e) + ret['msg'] = '生成发型材质失败' + return make_response(jsonify(ret), 400) + # 功能6：推理发型 + try: + start_time = time.time() + origin_img = cv2.imread(new_user_img_path) + ref_img_path = os.path.join(material_save_path, "ref_rgb_8uc3_768.png") + if not os.path.exists(ref_img_path): + raise FileNotFoundError("推理目标图像不存在") + ref_img = cv2.imread(ref_img_path) + with torch.no_grad(): + img_res, status, _, landmarks_origin_img_1k, isEyeOccluded = hairstyle_process.infer_hairstyle_diy_jy( + origin_img, ref_img, os.path.join(userInfo_dir, task_id), f"{task_id}.png" + ) + if status != 0: + raise RuntimeError("发型推理失败") + print(f"功能6：推理发型，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + print(e) + ret['msg'] = '推理发型失败' + return make_response(jsonify(ret), 400) + + try: + # 功能7：判断头发遮挡眼睛 + start_time = time.time() + user_material_dir = os.path.join(userInfo_dir, task_id) + user_img_txt_path = os.path.join(user_material_dir, "kpt_1k.txt") + user_mask_img_path = os.path.join(user_material_dir, "user_orig_mask.png") + user_pt1k = np.loadtxt(user_img_txt_path) + result, user_landmarks_origin_img_137 = hairstyle_process.check_hair_covering_eyes(hair_pt1k,hair_mask_img_path,new_user_img_path, user_pt1k,user_mask_img_path) + if result == 1: print("result == 1: hair_covering_eyes") + else: print("result == 0") + print(f"功能7：判断头发遮挡眼睛，耗时：{time.time() - start_time:.3f}s") + result=0 + if result == 1: + print("lead to diy") + ret_url = change_hairstyle_use_diy(new_user_img_path, template_ref_hair_path, userId, task_id, output_format) + if ret_url == "": return make_response(jsonify(ret), 200) + else: + ret["msg"] = 'success' + ret['state'] = 0 + ret['data'] = ret_url + return make_response(jsonify(ret), 200) + else: + print("lead te change hairstyle") + + # 功能8：处理发型区域 + start_time = time.time() + hair_matting_path = os.path.join(user_material_dir, "hair_mask_2.png") + new_matting = cv2.imread(hair_matting_path, cv2.IMREAD_GRAYSCALE) + result_img = img_res + user_orig_mask_path = os.path.join(user_material_dir, "user_orig_mask.png") + origin_matting = cv2.imread(user_orig_mask_path, cv2.IMREAD_GRAYSCALE) + useTitorMask = True + if useTitorMask: + # if is_hr and max(origin_img.shape[:2]) < 2000: + # start_time = time.time() + # scale_ratio = 2000 / max(origin_img.shape[:2]) + # result_img = webui_super_res_img(result_img, scale_ratio) + # origin_img = cv2.resize(origin_img, (result_img.shape[1], result_img.shape[0]), + # interpolation=cv2.INTER_LANCZOS4) + # print(f"功能：resize 2000，耗时：{time.time() - start_time:.3f}s") + + dst_size = (1152, 1536) if is_hr else (576, 768) # 处理是否高清模式的分辨率 + print(dst_size) + box_info = hairstyle_process.get_body_info(img_res) # 获取头发处理的局部区域图像 + # cv2.imwrite(f"{task_id}_img_res", img_res) + box_w, box_h = box_info[2] - box_info[0], box_info[3] - box_info[1] + print(f"box_w, box_h : {box_w},{box_h}") + h = max(dst_size[1] , box_h) + w = max(dst_size[0] , box_w) + + 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) + crop_result = landmark_processor.high_quality_warpAffine(img_res, M, dst_size) + # crop_result = landmark_processor.high_quality_warpAffine(img_res, M, (w, h)) + + origin_matting = cv2.imread(user_orig_mask_path, cv2.IMREAD_GRAYSCALE) + new_matting = cv2.imread(hair_matting_path, cv2.IMREAD_GRAYSCALE) + left_point = user_landmarks_origin_img_137[14] # (x1, y1) 获取刘海区域的圆心和半径 + right_point = user_landmarks_origin_img_137[7] # (x2, y2) + center_x = int((left_point[0] + right_point[0]) // 2) + center_y = int((left_point[1] + right_point[1]) // 2) + radius = int(np.sqrt((right_point[0] - left_point[0]) ** 2 + (right_point[1] - left_point[1]) ** 2) / 2) + h, w = origin_matting.shape # 初始化刘海区域为空（全黑色） + part_bangs = np.zeros((h, w), dtype=np.uint8) + cv2.circle(part_bangs, (center_x, center_y), radius, 255, -1) # 在 part_bangs 上绘制圆作为刘海区域 圆形区域设为白色 + no_bang_result = cv2.subtract(origin_matting, part_bangs) # 去除刘海区域 + matting_merge = np.max( np.stack([no_bang_result, new_matting], axis=2), axis=2).astype(np.uint8)# 合并 origin_matting 和 new_matting + crop_matting = cv2.warpAffine(matting_merge, M, dst_size) + # crop_matting = cv2.warpAffine(matting_merge, M, (w, h)) + mask = (crop_matting > 10).astype(np.float32) + if not is_hr: + mask_dilate = cv2.dilate(mask, np.ones((3, 9), np.uint8)) + else: + mask_dilate = cv2.dilate(mask, np.ones((6, 18), np.uint8)) + final_img = crop_result + mask_dilate = np.clip(mask_dilate * 255, 0, 255).astype(np.uint8) + else: + ## jy_version + matting_merge = np.concatenate([origin_matting[:, :, np.newaxis], new_matting[:, :, np.newaxis]],axis=2) + matting_merge = np.max(matting_merge, axis=2) + mask = cv2.resize(matting_merge, (img_res.shape[1], img_res.shape[0])) + mask_dilate = cv2.dilate(mask, np.ones((3, 11), np.uint8)) + print(mask_dilate.shape) + save_mask_path = osp.join(user_material_dir, 'webui.png') + cv2.imwrite(save_mask_path, mask_dilate) + crop_img, crop_mask, M = resize_pre_webui(img_res, mask_dilate) + final_img = crop_img + mask_dilate = crop_mask + # cv2.imwrite(f"{task_id}_finalimg.jpg", final_img) + # cv2.imwrite(f"{task_id}.jpg", mask_dilate) + config_json_path = os.path.join(material_save_path, "config.json") + with open(config_json_path, "r") as f: + config_json_content = json.load(f) + in_gender = config_json_content["gender"] + print("in_gender:", in_gender) + images_dir = os.path.join(hair_material_dir, "images") + txt_dir = os.path.join(images_dir, os.listdir(images_dir)[0]) + txt_path = glob.glob(txt_dir + '/*.txt')[0] + with open(txt_path, 'r') as f: + p_tag = f.readline() + if "titor hairstyle, faceless, no human, gray background, simple background" in p_tag: + p_tag = p_tag[p_tag.find("simple background, ") + len("simple background, "):] + else: + p_tag = "" + denoising_strength = 0.6 + print(f"功能8：处理发型区域，耗时：{time.time() - start_time:.3f}s") + # cv2.imwrite(f"{task_id}_mask_dilate.jpg", mask_dilate) + # cv2.imwrite(f"{task_id}_final_img.jpg", final_img) + # 功能9：webui + start_time = time.time() + # cv2.imwrite(f"{task_id}_final_img.jpg", final_img) + # cv2.imwrite(f"{task_id}_mask_dilate.jpg", mask_dilate) + # print("final image shape:", final_img.shape) + # print("mask_dilate shape:", mask_dilate.shape) + # final_img = cv2.imread("/root/project/hair_service_sd/gt_final.jpg") + # mask_dilate = cv2.imread("/root/project/hair_service_sd/gt_mask_dilate.jpg") + sd_result = webui_img2img(img=final_img, mask_img=mask_dilate, in_gender=in_gender, task_id=task_id, + hair_id=hair_id, lora_material_path=hair_material_dir, tag=p_tag, is_hr=is_hr, + denoising_strength=denoising_strength, inference_port="57860") + print(f"功能：webui，耗时：{time.time() - start_time:.3f}s") + + # 功能：后处理并上传结果 + start_time = time.time() + dst_path = os.path.join(res_dir, f"{task_id}.jpg") + if not os.path.exists(os.path.dirname(dst_path)): + os.makedirs(os.path.dirname(dst_path)) + origin_img_final = origin_img.copy() + M_inv = cv2.invertAffineTransform(M) + cv2.warpAffine(sd_result, M_inv, (origin_img.shape[1], origin_img.shape[0]), + dst=origin_img_final, borderMode=cv2.BORDER_TRANSPARENT, flags=cv2.INTER_LANCZOS4) + cv2.imwrite(dst_path, origin_img_final) + + # colo CHANGE + if clean_temp_data: + try: + shutil.rmtree(user_material_dir) + os.remove(new_user_img_path) + except Exception as e: + print(f"删除临时数据失败: {e}") + + if os.path.exists(dst_path): + + start_time3 = time.time() # 添加start_time3定义 + if output_format == 'base64': + with open(dst_path, 'rb') as f: + img_data = f.read() + base64_str = base64.b64encode(img_data).decode('utf-8') + print(f"all总耗时：{time.time() - start_time_all:.3f}s") + os.remove(dst_path) + return make_response(jsonify( + {'msg': 'success', 'data': base64_str, + 'state': 0}), 200) + else: + ret_url = hairstyle_process.oss2.upload_file(dst_path, + f"hair_mz/images/hairstyle/{hair_id}/{uuid4()}.jpg") + print("---------------upload img:", time.time() - start_time3) + print(f"all总耗时：{time.time() - start_time_all:.3f}s") + os.remove(dst_path) + + return make_response(jsonify( + {'msg': 'success', 'data': ret_url, + 'state': 0}), 200) + + # ret_url = oss_2.upload_file(dst_path, f"hair_mz/images/hairstyle/{hair_id}/{uuid4()}.jpg") + # ret.update({"msg": "success", "state": 0, "data": ret_url}) + # print(f"功能9：后处理并上传结果，耗时：{time.time() - start_time:.3f}s") + + except Exception as e: + print(e) + traceback.print_exc() + ret['msg'] = '推理发型失败' + return make_response(jsonify(ret), 400) + + +def change_hairstyle_use_diy(dir_user, dir_tar, userId, task_id,output_format='url'): + print(f"____diy____") + userInfo_dir = config.get('default', 'userInfo_dir') + user_img_res_dir = config.get('default', 'res_dir') + + diy_start_time_all = time.time() + ret_url = "" + + # diy功能1：加载图像 + try: + start_time = time.time() + origin_img = cv2.imread(dir_user) + ref_img = cv2.imread(dir_tar) + + if origin_img is None or ref_img is None: + raise ValueError("图片可能损坏") + print(f"diy功能1：加载图像，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + hairstyle_process.logger_call.error(f"时间:{datetime.now()}，加载图像失败：{e}") + return make_response(jsonify({'msg': '图片可能损坏', 'data': "", 'umd': '', 'state': 1}), 200) + + # diy功能2：发型推理 + try: + start_time = time.time() + user_material_dir = os.path.join(userInfo_dir, task_id) + basefolder, imgname = osp.split(dir_user) + _, target_name = osp.split(dir_tar) + + with torch.no_grad(): + img_res, status, in_gender = hairstyle_process.infer_hairstyle_diy(origin_img, ref_img, user_material_dir, target_name) + + if status != 0: + raise RuntimeError("发型推理失败") + print(f"diy功能2：发型推理，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + hairstyle_process.logger_call.error(f"时间:{datetime.now()}，发型推理失败：{e}") + return "" + + + # diy功能3：保存发型结果图像 + try: + start_time = time.time() + img_res_dir = os.path.join(basefolder, 'diy') + os.makedirs(img_res_dir, exist_ok=True) + res_save_path = os.path.join(img_res_dir, f"{imgname[:-4]}-{target_name[:-4]}.jpg") + cv2.imwrite(res_save_path, img_res) + print(f"diy功能3：保存发型结果图像，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + hairstyle_process.logger_call.error(f"时间:{datetime.now()}，保存发型结果图像失败：{e}") + return "" + + # diy功能4：生成 mask + try: + start_time = time.time() + hair_mask_path = osp.join(user_material_dir, 'hair_mask_2.png') + orig_model_mask_path = osp.join(user_material_dir, 'user_orig_mask.png') + + new_matting = cv2.imread(hair_mask_path, cv2.IMREAD_GRAYSCALE) + origin_matting = cv2.imread(orig_model_mask_path, cv2.IMREAD_GRAYSCALE) + + matting_merge = np.max( + np.stack([origin_matting[:, :, np.newaxis], new_matting[:, :, np.newaxis]], axis=2), axis=2 + ).astype(np.uint8) + + mask_dilate = cv2.dilate(cv2.resize(matting_merge, (img_res.shape[1], img_res.shape[0])), np.ones((3, 11), np.uint8)) + save_mask_path = osp.join(user_material_dir, 'webui.png') + cv2.imwrite(save_mask_path, mask_dilate) + print(f"diy功能4：生成 mask，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + hairstyle_process.logger_call.error(f"时间:{datetime.now()}，生成 mask 失败：{e}") + return "" + + # diy功能5：增强图像 + try: + start_time = time.time() + crop_img, crop_mask, M = resize_pre_webui(img_res, mask_dilate) + crop_img_path = os.path.join(user_img_res_dir, f"{task_id}_crop.png") + crop_mask_path = os.path.join(user_img_res_dir, f"{task_id}_mask.png") + + cv2.imwrite(crop_img_path, crop_img) + cv2.imwrite(crop_mask_path, crop_mask) + + out = call_hair_inter.call_hair_enhance(crop_img_path, crop_mask_path, task_id, in_gender) + out_path = out["result"] + enhanced_img = cv2.imread(out_path) + print(f"diy功能5：增强图像，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + hairstyle_process.logger_call.error(f"时间:{datetime.now()}，增强图像失败：{e}") + return "" + + # diy功能6：恢复图像并保存 + try: + start_time = time.time() + origin_img_final = origin_img.copy() + 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) + + final_dst_path = os.path.join(user_img_res_dir, f"{task_id}_res.jpg") + cv2.imwrite(final_dst_path, origin_img_final) + print(f"diy功能6：恢复图像并保存，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + hairstyle_process.logger_call.error(f"时间:{datetime.now()}，恢复图像失败：{e}") + return "" + + # diy功能7：上传图像到 OSS + try: + start_time = time.time() + + if output_format == 'base64': + with open(final_dst_path, 'rb') as f: + img_data = f.read() + base64_str = base64.b64encode(img_data).decode('utf-8') + print(f"diy功能3：保存发型结果图像，耗时：{time.time() - start_time:.3f}s") + return base64_str + else: + ret_url = hairstyle_process.oss2.upload_file(final_dst_path, + "hair_mz/images/hairstyle/{userId}/{task_id}.jpg") + print("---------------upload img:", time.time() - start_time3) + print(f"diy功能3：保存发型结果图像，耗时：{time.time() - start_time:.3f}s") + return ret_url + # ret_url = oss_2.upload_file(final_dst_path, f"hair_mz/images/hairstyle/{userId}/{task_id}.jpg") + # print(f"diy功能7：上传图像到 OSS，耗时：{time.time() - start_time:.3f}s") + except Exception as e: + hairstyle_process.logger_call.error(f"时间:{datetime.now()}，上传图像到 OSS 失败：{e}") + return "" + + +@app.route("/api/uploadHair/v1", methods=['POST']) # 新增发型模板新 +def add_new_hairstyle_v2(): + input_data = request.json + print(f"____uploadHair____{input_data}") + imgLists = input_data['img_lists'] + input_info = [] + hairID = input_data['hair_id'] + # cover_url = input_data['coverImg'] + userId = 'houtai' + hairstyle_process.logger_call.info('time:{}, fff, {}'.format(datetime.now().strftime("%Y-%m-%d, %H:%M:%S"), input_data)) + + if len(imgLists) == 0: + hairstyle_process.logger_call.error( + 'time:{}, 未选择图像传入, {}'.format(datetime.now().strftime("%Y-%m-%d, %H:%M:%S"), input_data)) + + return make_response(jsonify({'msg': '未选择图像传入', 'data': "", 'state': 1}), 200) + + for index,img_url in enumerate(imgLists): + + img_path ,md5_img= download_img(img_url, userId) + #img_path =img_url + if index == 0: + cover_path=img_path + if not os.path.exists(img_path): + hairstyle_process.logger_call.error('time:{}, 图像读取错误'.format(datetime.now().strftime("%Y-%m-%d, %H:%M:%S"))) + return make_response(jsonify({'msg': '图像读取错误', 'data': '', 'state': 1}), 200) + + # 存储原始图片 + dst_dir = osp.join(train_upload_dir, hairID) + # 存储用户参考图 + ref_user_save_dir = ref_user_dir + # 生成训练素材图 + train_material_save_dir = os.path.join(train_save_dir, hairID) + os.makedirs(dst_dir, exist_ok=True) + + # 拷贝训练数据 + file_name = img_path[img_path.rfind("/") + 1:] + new_file_path = os.path.join(dst_dir, file_name) + shutil.copy(img_path, new_file_path) + + input_info.append([new_file_path, cover_path, 0, dst_dir, ref_user_save_dir, train_material_save_dir]) + + + args = [input_info] + executor.submit(lambda p: run_hairstyle_v2(*p), args) + return make_response(jsonify({'msg': 'OK', 'data': "", 'state': 0}), 200) + + +@app.route("/api/hair/trainCallBack", methods=['POST']) # 发型训练回调 +def train_hair_callback(): + input_info = request.json + print(f"____trainCallBack____{input_info}") + ret_data = { + "task_id": "", + "state": -1, + "msg": "failed", + } + + hairstyle_dir = config.get('default', 'hairstyleDir') + ref_img_dir = config.get('default', 'refImgDir') + hair_template_material_dir = config.get('default', 'hair_template_material_dir') + + # get request params + try: + req_id = input_info['task_id'] + hair_id = input_info['hair_id'] + state = input_info['state'] + msg = input_info['msg'] + print("trainCallBack input_info: ", input_info) + except Exception as e: + print(e) + hairstyle_process.logger_call.error(f'params error, input is {input_info} error is {e}') + return make_response(jsonify(ret_data), 400) + + # gen hair material + hair_upload_save_dir = os.path.join(train_upload_dir, hair_id) + print("hair_upload_save_dir: ", hair_upload_save_dir) + + hair_name_lists = os.listdir(hair_upload_save_dir) + for hair_name in hair_name_lists: + if "first##" in hair_name and (hair_name.endswith(".png") or hair_name.endswith(".jpg") or hair_name.endswith(".jpeg")): + template_ref_hair_name = hair_name + template_ref_hair_path = os.path.join(hair_upload_save_dir, hair_name) + break + + new_hair_ref_img_path = os.path.join(ref_img_dir, template_ref_hair_name) + print("new_hair_ref_img_path: ", new_hair_ref_img_path) + shutil.copy(template_ref_hair_path, new_hair_ref_img_path) + + material_save_path = os.path.join(hairstyle_dir, hair_id) + print("first_hair_material_save_path :", material_save_path) + + long_flag = False + prepare_single(new_hair_ref_img_path, material_save_path, long_flag) + + # 训练完成，回调天津后端接口 + status = -1 + try: + if state == 0: + ret_data = {} + ret_data['code'] = 200 + ret_data['data'] = [] + ret_data['msg'] = 'success' + + cover_img_path = new_hair_ref_img_path + + ret_url = oss_2.upload_file(cover_img_path, "hair_mz/images/hairstyle/{}/{}".format(hair_id, + str(uuid4()) + '.jpg')) + print('recall success') + response, status = callback_hairstyle(req_id, state, msg, hair_id) + print(response) + else: + response, status = callback_hairstyle(req_id, state, msg, hair_id) + print(response) + + except Exception as e: + print(e) + print('recall failed, ConnectionError') + + if status != 0: + ret_data['state'] = -1 + ret_data['msg'] = 'failed' + ret_data['task_id'] = input_info['task_id'] + return make_response(jsonify(ret_data), 400) + + ret_data['state'] = 0 + ret_data['msg'] = 'success' + ret_data['task_id'] = input_info['task_id'] + return make_response(jsonify(ret_data), 200) + + + +if __name__ == '__main__': + server = pywsgi.WSGIServer(('0.0.0.0', 8801), app) + server.serve_forever() diff --git a/hair_service_sd/scripts/__init__.py b/hair_service_sd/scripts/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hair_service_sd/scripts/tools.py b/hair_service_sd/scripts/tools.py new file mode 100644 index 0000000..504a15e --- /dev/null +++ b/hair_service_sd/scripts/tools.py @@ -0,0 +1,70 @@ +import time +import cv2 +import math +import numpy as np + + +def draw_protect_mask(sourceImage, landmark_137): + h, w, _ = sourceImage.shape + inpaint_mask = np.zeros((h, w), dtype=np.uint8) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[129:137]), (255,)) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[121:129]), (255,)) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[22:48]), (255,)) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[88:104]), (255,)) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[105:121]), (255,)) + return inpaint_mask + + +def localtranslationwarpfastwithstrength(srcimg, kpt137, startx, starty, endx, endy, radius, strength): + ddradius = float(radius * radius) + mask_keep = draw_protect_mask(srcimg, kpt137) + # mask_keep = cv2.imread('/home/colo/Pictures/test/102.png') + # mask_keep = cv2.cvtColor(mask_keep, cv2.COLOR_BGR2GRAY) + # copyimg = np.zeros(srcimg.shape, np.uint8) + # copyimg = srcimg.copy() + maskimg = np.zeros(srcimg.shape[:2], np.uint8) + cv2.circle(maskimg, (startx, starty), math.ceil(radius), (255, 255, 255), -1) + # cv2.imshow('maskimg_before', maskimg) + # cv2.imshow('maskimg', maskimg) + # cv2.imshow('mask_keep', mask_keep) + # cv2.waitKey() + maskimg = maskimg * (1-mask_keep/255).astype(np.uint8) + + k0 = 100 / strength # 计算公式中的|m-c|^2 + ddmc_x = (endx - startx) * (endx - startx) + ddmc_y = (endy - starty) * (endy - starty) + h, w, c = srcimg.shape + mapx = np.vstack([np.arange(w).astype(np.float32).reshape(1, -1)] * h) + mapy = np.hstack([np.arange(h).astype(np.float32).reshape(-1, 1)] * w) + distance_x = (mapx - startx) * (mapx - startx) + distance_y = (mapy - starty) * (mapy - starty) + distance = distance_x + distance_y + k1 = np.sqrt(distance) + ratio_x = (ddradius - distance_x) / (ddradius - distance_x + k0 * ddmc_x) + ratio_y = (ddradius - distance_y) / (ddradius - distance_y + k0 * ddmc_y) + ratio_x = ratio_x * ratio_x + ratio_y = ratio_y * ratio_y + ux = mapx - ratio_x * (endx - startx) * (1 - k1/radius) + uy = mapy - ratio_y * (endy - starty) * (1 - k1/radius) + np.copyto(ux, mapx, where=maskimg == 0) + np.copyto(uy, mapy, where=maskimg == 0) + ux = ux.astype(np.float32) + uy = uy.astype(np.float32) + copyimg = cv2.remap(srcimg, ux, uy, interpolation=cv2.INTER_LINEAR) + return copyimg + +image = cv2.imread("/home/colo/Pictures/for_hn/102.jpg") +processed_image = image.copy() +startx_left, starty_left, endx_left, endy_left = 170, 123, 190, 74 +# startx_right, starty_right, endx_right, endy_right = 287, 275, 192, 233 +radius = 60 +strength = 100 # 瘦左边脸 +t0 = time.time() +processed_image = localtranslationwarpfastwithstrength(processed_image, startx_left, starty_left, endx_left, endy_left, radius, strength) # 瘦右边脸 +# processed_image = localtranslationwarpfastwithstrength(processed_image, startx_right, starty_right, endx_right, endy_right, radius, strength) +# cv2.imwrite("thin.jpg", processed_image) +print('costs', time.time() - t0) +# cv2.imshow('image', image) + +# cv2.imshow('processed_image', processed_image) +# cv2.waitKey() diff --git a/hair_service_sd/seg/hairseg_single_model.py b/hair_service_sd/seg/hairseg_single_model.py new file mode 100644 index 0000000..faabac4 --- /dev/null +++ b/hair_service_sd/seg/hairseg_single_model.py @@ -0,0 +1,86 @@ +import os +import torch + +from seg.networks.deeplabv3_plus import get_deeplabv3_plus +import numpy as np +import cv2 +import time +from utils import landmark_processor + +def label_to_mask(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 + +label_map = [ + [0, 0, 0], # + [128, 128, 128], + [255, 255, 255], +] + +class Evaluator(object): + def __init__(self, gpu_id, output_img_size, nclass, seg_model_path=None): + self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu') + # print("gpu_id: ", gpu_id) + + # create network + self.model = get_deeplabv3_plus(backbone='xception', nclass=nclass) + model_path = os.path.join(seg_model_path) + self.model.load_state_dict(torch.load(model_path, map_location=lambda storage, loc: storage)) + # print("seg device: ", self.model.device) + self.model.to(self.device) + self.model.eval() + + # images = torch.randn((1, 3, 512, 512)).to(self.device) + # torch.onnx.export(self.model, images, + # "deeplabv3_hair512_360_0520_wl.onnx", + # verbose=True, + # opset_version=11, + # input_names=['data'], + # do_constant_folding=True, + # output_names=['output']) + + # exit() + + self.output_img_size = output_img_size + self.nclass = nclass + + + def process_data(self, img): + img = (img.astype(np.float32) / 255).transpose((2, 0, 1)) + img = torch.from_numpy(img).unsqueeze(0) + + return img + + def eval(self, img, pts1k): + orig_h, orig_w, _ = img.shape + + M1 = landmark_processor.get_transform_mat_hair(pts1k, self.output_img_size, ratio=0.28, h_ratio=0.3) + crop_img = cv2.warpAffine(img, M1, (self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4) + crop_img = self.process_data(crop_img) + crop_img = crop_img.to(self.device) + + with torch.no_grad(): + # torch.cuda.synchronize() + outputs = self.model(crop_img) + pred = torch.argmax(outputs[0], 1) + + pred = pred[0].detach().cpu().numpy() + predict = pred.astype(np.float32) + + pred_mask = label_to_mask(predict) + + M1_invert = cv2.invertAffineTransform(M1) + img_pred = cv2.warpAffine(pred_mask, M1_invert, (orig_w, orig_h), flags=cv2.INTER_CUBIC) #flags=cv2.INTER_NEAREST + orig_mask = img_pred.copy() + + # show_concat = np.concatenate((img, orig_mask), axis=1) + # cv2.imshow("show_concat", show_concat) + # cv2.waitKey() + return orig_mask + + + diff --git a/hair_service_sd/seg/networks/basic.py b/hair_service_sd/seg/networks/basic.py new file mode 100644 index 0000000..b241fad --- /dev/null +++ b/hair_service_sd/seg/networks/basic.py @@ -0,0 +1,462 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +__all__ = ['_ConvBNReLU', '_DWConvBNReLU', 'InvertedResidual', '_ASPP', '_FCNHead', + '_Hswish', '_ConvBNHswish', 'SEModule', 'Bottleneck', 'ShuffleNetUnit', + 'ShuffleNetV2Unit', 'InvertedIGCV3', 'MBConvBlock'] + + +class _ConvBNReLU(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, + dilation=1, groups=1, relu6=False, norm_layer=nn.BatchNorm2d, **kwargs): + super(_ConvBNReLU, self).__init__() + self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias=False) + self.bn = norm_layer(out_channels) + self.relu = nn.ReLU6(True) if relu6 else nn.ReLU(True) + + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + x = self.relu(x) + return x + + +class _FCNHead(nn.Module): + def __init__(self, in_channels, channels, norm_layer=nn.BatchNorm2d, norm_kwargs=None, **kwargs): + super(_FCNHead, self).__init__() + inter_channels = in_channels // 4 + self.block = nn.Sequential( + nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False), + norm_layer(inter_channels, **({} if norm_kwargs is None else norm_kwargs)), + nn.ReLU(True), + nn.Dropout(0.1), + nn.Conv2d(inter_channels, channels, 1) + ) + + def forward(self, x): + return self.block(x) + + +# ----------------------------------------------------------------- +# For MobileNet +# ----------------------------------------------------------------- +class _DWConvBNReLU(nn.Module): + """Depthwise Separable Convolution in MobileNet. + depthwise convolution + pointwise convolution + """ + + def __init__(self, in_channels, dw_channels, out_channels, stride, dilation=1, norm_layer=nn.BatchNorm2d, **kwargs): + super(_DWConvBNReLU, self).__init__() + self.conv = nn.Sequential( + _ConvBNReLU(in_channels, dw_channels, 3, stride, dilation, dilation, in_channels, norm_layer=norm_layer), + _ConvBNReLU(dw_channels, out_channels, 1, norm_layer=norm_layer)) + + def forward(self, x): + return self.conv(x) + + +# ----------------------------------------------------------------- +# For MobileNetV2 +# ----------------------------------------------------------------- +class InvertedResidual(nn.Module): + def __init__(self, in_channels, out_channels, stride, expand_ratio, + dilation=1, norm_layer=nn.BatchNorm2d, **kwargs): + super(InvertedResidual, self).__init__() + assert stride in [1, 2] + self.use_res_connect = stride == 1 and in_channels == out_channels + + layers = list() + inter_channels = int(round(in_channels * expand_ratio)) + if expand_ratio != 1: + # pw + layers.append(_ConvBNReLU(in_channels, inter_channels, 1, relu6=True, norm_layer=norm_layer)) + layers.extend([ + # dw + _ConvBNReLU(inter_channels, inter_channels, 3, stride, dilation, dilation, + groups=inter_channels, relu6=True, norm_layer=norm_layer), + # pw-linear + nn.Conv2d(inter_channels, out_channels, 1, bias=False), + norm_layer(out_channels)]) + self.conv = nn.Sequential(*layers) + + def forward(self, x): + if self.use_res_connect: + return x + self.conv(x) + else: + return self.conv(x) + + +# ----------------------------------------------------------------- +# ASPP: For MobileNetV2 +# ----------------------------------------------------------------- +class _AsppPooling(nn.Module): + def __init__(self, in_channels, out_channels, norm_layer, **kwargs): + super(_AsppPooling, self).__init__() + self.gap = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(in_channels, out_channels, 1, bias=False), + norm_layer(out_channels), + nn.ReLU(True) + ) + + def forward(self, x): +# size = x.size()[2:] + size = (48, 48) +# print("size: ", size) + pool = self.gap(x) +# out = F.interpolate(pool, size, mode='bilinear', align_corners=True) + out = F.interpolate(pool, size, mode='nearest') + return out + + +class _ASPP(nn.Module): + def __init__(self, in_channels, atrous_rates, norm_layer=nn.BatchNorm2d, **kwargs): + super(_ASPP, self).__init__() + out_channels = 256 + self.b0 = nn.Sequential( + nn.Conv2d(in_channels, out_channels, 1, bias=False), + norm_layer(out_channels), + nn.ReLU(True) + ) + + rate1, rate2, rate3 = tuple(atrous_rates) + self.b1 = _ConvBNReLU(in_channels, out_channels, 3, padding=rate1, dilation=rate1, norm_layer=norm_layer) + self.b2 = _ConvBNReLU(in_channels, out_channels, 3, padding=rate2, dilation=rate2, norm_layer=norm_layer) + self.b3 = _ConvBNReLU(in_channels, out_channels, 3, padding=rate3, dilation=rate3, norm_layer=norm_layer) + self.b4 = _AsppPooling(in_channels, out_channels, norm_layer=norm_layer) + + self.project = nn.Sequential( + nn.Conv2d(5 * out_channels, out_channels, 1, bias=False), + norm_layer(out_channels), + nn.ReLU(True), + nn.Dropout2d(0.5) + ) + + def forward(self, x): + feat1 = self.b0(x) + feat2 = self.b1(x) + feat3 = self.b2(x) + feat4 = self.b3(x) + feat5 = self.b4(x) + x = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1) + x = self.project(x) + return x + + +# ----------------------------------------------------------------- +# For MobileNetV3 +# ----------------------------------------------------------------- +class _Hswish(nn.Module): + def __init__(self, inplace=True): + super(_Hswish, self).__init__() + self.relu6 = nn.ReLU6(inplace) + + def forward(self, x): + return x * self.relu6(x + 3.) / 6. + + +class _Hsigmoid(nn.Module): + def __init__(self, inplace=True): + super(_Hsigmoid, self).__init__() + self.relu6 = nn.ReLU6(inplace) + + def forward(self, x): + return self.relu6(x + 3.) / 6. + + +class _ConvBNHswish(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, + dilation=1, groups=1, norm_layer=nn.BatchNorm2d, **kwargs): + super(_ConvBNHswish, self).__init__() + self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias=False) + self.bn = norm_layer(out_channels) + self.act = _Hswish(True) + + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + x = self.act(x) + return x + + +class SEModule(nn.Module): + def __init__(self, in_channels, reduction=4): + super(SEModule, self).__init__() + self.avg_pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(in_channels, in_channels // reduction, bias=False), + nn.ReLU(True), + nn.Linear(in_channels // reduction, in_channels, bias=False), + _Hsigmoid(True) + ) + + def forward(self, x): + n, c, _, _ = x.size() + out = self.avg_pool(x).view(n, c) + out = self.fc(out).view(n, c, 1, 1) + return x * out.expand_as(x) + + +class Identity(nn.Module): + def __init__(self, in_channels): + super(Identity, self).__init__() + + def forward(self, x): + return x + + +class Bottleneck(nn.Module): + def __init__(self, in_channels, out_channels, exp_size, kernel_size, stride, dilation=1, se=False, nl='RE', + norm_layer=nn.BatchNorm2d, **kwargs): + super(Bottleneck, self).__init__() + assert stride in [1, 2] + self.use_res_connect = stride == 1 and in_channels == out_channels + if nl == 'HS': + act = _Hswish + else: + act = nn.ReLU + if se: + SELayer = SEModule + else: + SELayer = Identity + + self.conv = nn.Sequential( + # pw + nn.Conv2d(in_channels, exp_size, 1, bias=False), + norm_layer(exp_size), + act(True), + # dw + nn.Conv2d(exp_size, exp_size, kernel_size, stride, (kernel_size - 1) // 2 * dilation, + dilation, groups=exp_size, bias=False), + norm_layer(exp_size), + SELayer(exp_size), + act(True), + # pw-linear + nn.Conv2d(exp_size, out_channels, 1, bias=False), + norm_layer(out_channels) + ) + + def forward(self, x): + if self.use_res_connect: + return x + self.conv(x) + else: + return self.conv(x) + + +# ----------------------------------------------------------------- +# For ShuffleNet +# ----------------------------------------------------------------- +def channel_shuffle(x, groups): + n, c, h, w = x.size() + + channels_per_group = c // groups + x = x.view(n, groups, channels_per_group, h, w) + x = torch.transpose(x, 1, 2).contiguous() + x = x.view(n, -1, h, w) + + return x + + +class ShuffleNetUnit(nn.Module): + def __init__(self, in_channels, out_channels, stride, groups, dilation=1, norm_layer=nn.BatchNorm2d, **kwargs): + super(ShuffleNetUnit, self).__init__() + self.stride = stride + self.groups = groups + self.dilation = dilation + assert stride in [1, 2, 3] + + inter_channels = out_channels // 4 + + if stride > 1: + self.shortcut = nn.AvgPool2d(3, stride, 1) + out_channels -= in_channels + elif dilation > 1: + out_channels -= in_channels + + g = 1 if in_channels == 24 else groups + self.conv1 = _ConvBNReLU(in_channels, inter_channels, 1, groups=g, norm_layer=norm_layer) + self.conv2 = _ConvBNReLU(inter_channels, inter_channels, 3, stride, dilation, + dilation, groups, norm_layer=norm_layer) + self.conv3 = nn.Sequential( + nn.Conv2d(inter_channels, out_channels, 1, groups=groups, bias=False), + norm_layer(out_channels)) + + def forward(self, x): + out = self.conv1(x) + out = channel_shuffle(out, self.groups) + out = self.conv2(out) + out = self.conv3(out) + if self.stride > 1: + x = self.shortcut(x) + out = torch.cat([out, x], dim=1) + elif self.dilation > 1: + out = torch.cat([out, x], dim=1) + else: + out = out + x + out = F.relu(out) + + return out + + +# ----------------------------------------------------------------- +# For ShuffleNetV2 +# ----------------------------------------------------------------- +class _DWConv(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, bias=False): + super(_DWConv, self).__init__() + self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, + padding, dilation, groups=in_channels, bias=bias) + + def forward(self, x): + return self.conv(x) + + +class ShuffleNetV2Unit(nn.Module): + def __init__(self, in_channels, out_channels, stride, dilation=1, norm_layer=nn.BatchNorm2d, **kwargs): + super(ShuffleNetV2Unit, self).__init__() + assert stride in [1, 2, 3] + self.stride = stride + self.dilation = dilation + + inter_channels = out_channels // 2 + + if (stride > 1) or (dilation > 1): + self.branch1 = nn.Sequential( + _DWConv(in_channels, in_channels, 3, stride, dilation, dilation), + norm_layer(in_channels), + _ConvBNReLU(in_channels, inter_channels, 1, norm_layer=norm_layer)) + self.branch2 = nn.Sequential( + _ConvBNReLU(in_channels if (stride > 1) else inter_channels, inter_channels, 1, norm_layer=norm_layer), + _DWConv(inter_channels, inter_channels, 3, stride, dilation, dilation), + norm_layer(inter_channels), + _ConvBNReLU(inter_channels, inter_channels, 1, norm_layer=norm_layer)) + + def forward(self, x): + if (self.stride == 1) and (self.dilation == 1): + x1, x2 = x.chunk(2, dim=1) + out = torch.cat((x1, self.branch2(x2)), dim=1) + else: + out = torch.cat((self.branch1(x), self.branch2(x)), dim=1) + out = channel_shuffle(out, 2) + + return out + + +# ----------------------------------------------------------------- +# For IGCV3 +# ----------------------------------------------------------------- +class PermutationBlock(nn.Module): + def __init__(self, groups): + super(PermutationBlock, self).__init__() + self.groups = groups + + def forward(self, x): + n, c, h, w = x.size() + x = x.view(n, self.groups, c // self.groups, h, w).permute(0, 2, 1, 3, 4).contiguous().view(n, c, h, w) + return x + + +class InvertedIGCV3(nn.Module): + def __init__(self, in_channels, out_channels, stride, expand_ratio, + dilation=1, norm_layer=nn.BatchNorm2d, **kwargs): + super(InvertedIGCV3, self).__init__() + assert stride in [1, 2] + self.use_res_connect = stride == 1 and in_channels == out_channels + + layers = list() + inter_channels = int(round(in_channels * expand_ratio)) + if expand_ratio != 1: + # pw + layers.append(_ConvBNReLU(in_channels, inter_channels, 1, + groups=2, relu6=True, norm_layer=norm_layer)) + # permutation + layers.append(PermutationBlock(groups=2)) + layers.extend([ + # dw + _ConvBNReLU(inter_channels, inter_channels, 3, stride, dilation, dilation, + groups=inter_channels, relu6=True, norm_layer=norm_layer), + # pw-linear + nn.Conv2d(inter_channels, out_channels, 1, groups=2, bias=False), + norm_layer(out_channels), + # permutation + PermutationBlock(groups=int(round(out_channels / 2))) + ]) + self.conv = nn.Sequential(*layers) + + def forward(self, x): + if self.use_res_connect: + return x + self.conv(x) + else: + return self.conv(x) + + +# ----------------------------------------------------------------- +# For EfficientNet +# ----------------------------------------------------------------- +class _Swish(nn.Module): + def __init__(self): + super(_Swish, self).__init__() + self.sigmoid = nn.Sigmoid() + + def forward(self, x): + return x * self.sigmoid(x) + + +class SEModuleV2(nn.Module): + def __init__(self, in_channels, se_ratio=0.25): + super(SEModuleV2, self).__init__() + self.avg_pool = nn.AdaptiveAvgPool2d(1) + se_channels = max(1, int(in_channels * se_ratio)) + self.fc = nn.Sequential( + nn.Conv2d(in_channels, se_channels, 1, bias=False), + _Swish(), + nn.Conv2d(se_channels, in_channels, 1, bias=False), + nn.Sigmoid() + ) + + def forward(self, x): + n, c, _, _ = x.size() + out = self.avg_pool(x) + out = self.fc(out) + return x * out.expand_as(x) + + +class MBConvBlock(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride, expand_ratio, + dilation=1, se_ratio=0.25, drop_connect_rate=0.2, norm_layer=nn.BatchNorm2d, **kwargs): + super(MBConvBlock, self).__init__() + assert stride in [1, 2] + self.use_res_connect = stride == 1 and in_channels == out_channels + self.drop_connect_rate = drop_connect_rate + use_se = (se_ratio is not None) and (0 < se_ratio <= 1.) + if use_se: + SELayer = SEModuleV2 + else: + SELayer = Identity + + layers = list() + inter_channels = int(round(in_channels * expand_ratio)) + if expand_ratio != 1: + layers.append(_ConvBNHswish(in_channels, inter_channels, 1, norm_layer=norm_layer)) + layers.extend([ + # dw + _ConvBNHswish(inter_channels, inter_channels, kernel_size, stride, kernel_size // 2 * dilation, dilation, + groups=inter_channels, norm_layer=norm_layer), # check act function + SELayer(inter_channels, se_ratio), + # pw-linear + nn.Conv2d(inter_channels, out_channels, 1, bias=False), + norm_layer(out_channels) + ]) + self.conv = nn.Sequential(*layers) + + if drop_connect_rate: + self.dropout = nn.Dropout2d(drop_connect_rate) + + def forward(self, x): + out = self.conv(x) + if self.use_res_connect: + if self.drop_connect_rate: + out = self.dropout(out) + out = x + out + return out diff --git a/hair_service_sd/seg/networks/deeplabv3.py b/hair_service_sd/seg/networks/deeplabv3.py new file mode 100644 index 0000000..d3fef95 --- /dev/null +++ b/hair_service_sd/seg/networks/deeplabv3.py @@ -0,0 +1,187 @@ +"""Pyramid Scene Parsing Network""" +import torch +import torch.nn as nn +import torch.nn.functional as F + +from seg.networks.segbase import SegBaseModel +from seg.networks.fcn import _FCNHead + +__all__ = ['DeepLabV3', 'get_deeplabv3', 'get_deeplabv3_resnet50_voc', 'get_deeplabv3_resnet101_voc', + 'get_deeplabv3_resnet152_voc', 'get_deeplabv3_resnet50_ade', 'get_deeplabv3_resnet101_ade', + 'get_deeplabv3_resnet152_ade'] + + +class DeepLabV3(SegBaseModel): + r"""DeepLabV3 + + Parameters + ---------- + nclass : int + Number of categories for the training dataset. + backbone : string + Pre-trained dilated backbone network type (default:'resnet50'; 'resnet50', + 'resnet101' or 'resnet152'). + norm_layer : object + Normalization layer used in backbone network (default: :class:`nn.BatchNorm`; + for Synchronized Cross-GPU BachNormalization). + aux : bool + Auxiliary loss. + + Reference: + Chen, Liang-Chieh, et al. "Rethinking atrous convolution for semantic image segmentation." + arXiv preprint arXiv:1706.05587 (2017). + """ + + def __init__(self, nclass, backbone='resnet50', aux=False, pretrained_base=True, **kwargs): + super(DeepLabV3, self).__init__(nclass, aux, backbone, pretrained_base=pretrained_base, **kwargs) + self.head = _DeepLabHead(nclass, **kwargs) + if self.aux: + self.auxlayer = _FCNHead(1024, nclass, **kwargs) + + self.__setattr__('exclusive', ['head', 'auxlayer'] if aux else ['head']) + + def forward(self, x): + size = x.size()[2:] + _, _, c3, c4 = self.base_forward(x) + outputs = [] + x = self.head(c4) + x = F.interpolate(x, size, mode='bilinear', align_corners=True) + outputs.append(x) + + if self.aux: + auxout = self.auxlayer(c3) + auxout = F.interpolate(auxout, size, mode='bilinear', align_corners=True) + outputs.append(auxout) + return tuple(outputs) + + +class _DeepLabHead(nn.Module): + def __init__(self, nclass, norm_layer=nn.BatchNorm2d, norm_kwargs=None, **kwargs): + super(_DeepLabHead, self).__init__() + self.aspp = _ASPP(2048, [12, 24, 36], norm_layer=norm_layer, norm_kwargs=norm_kwargs, **kwargs) + self.block = nn.Sequential( + nn.Conv2d(256, 256, 3, padding=1, bias=False), + norm_layer(256, **({} if norm_kwargs is None else norm_kwargs)), + nn.ReLU(True), + nn.Dropout(0.1), + nn.Conv2d(256, nclass, 1) + ) + + def forward(self, x): + x = self.aspp(x) + return self.block(x) + + +class _ASPPConv(nn.Module): + def __init__(self, in_channels, out_channels, atrous_rate, norm_layer, norm_kwargs): + super(_ASPPConv, self).__init__() + self.block = nn.Sequential( + nn.Conv2d(in_channels, out_channels, 3, padding=atrous_rate, dilation=atrous_rate, bias=False), + norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)), + nn.ReLU(True) + ) + + def forward(self, x): + return self.block(x) + + +class _AsppPooling(nn.Module): + def __init__(self, in_channels, out_channels, norm_layer, norm_kwargs, **kwargs): + super(_AsppPooling, self).__init__() + self.gap = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(in_channels, out_channels, 1, bias=False), + norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)), + nn.ReLU(True) + ) + + def forward(self, x): + size = x.size()[2:] +# print("before gap: ", x.size()) + pool = self.gap(x) + out = F.interpolate(pool, size, mode='bilinear', align_corners=True) + return out + + +class _ASPP(nn.Module): + def __init__(self, in_channels, atrous_rates, norm_layer, norm_kwargs=None, **kwargs): + super(_ASPP, self).__init__() + out_channels = 256 + self.b0 = nn.Sequential( + nn.Conv2d(in_channels, out_channels, 1, bias=False), + norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)), + nn.ReLU(True) + ) + + rate1, rate2, rate3 = tuple(atrous_rates) + self.b1 = _ASPPConv(in_channels, out_channels, rate1, norm_layer, norm_kwargs) + self.b2 = _ASPPConv(in_channels, out_channels, rate2, norm_layer, norm_kwargs) + self.b3 = _ASPPConv(in_channels, out_channels, rate3, norm_layer, norm_kwargs) + self.b4 = _AsppPooling(in_channels, out_channels, norm_layer=norm_layer, norm_kwargs=norm_kwargs) + + self.project = nn.Sequential( + nn.Conv2d(5 * out_channels, out_channels, 1, bias=False), + norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)), + nn.ReLU(True), + nn.Dropout(0.5) + ) + + def forward(self, x): + feat1 = self.b0(x) + feat2 = self.b1(x) + feat3 = self.b2(x) + feat4 = self.b3(x) +# print("before b4: ", x.size()) + feat5 = self.b4(x) + x = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1) + x = self.project(x) + return x + + +def get_deeplabv3(dataset='pascal_voc', backbone='resnet50', pretrained=False, root='~/.torch/models', + pretrained_base=True, **kwargs): + acronyms = { + 'pascal_voc': 'pascal_voc', + 'pascal_aug': 'pascal_aug', + 'ade20k': 'ade', + 'coco': 'coco', + 'citys': 'citys', + } + from ..data.dataloader import datasets + model = DeepLabV3(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs) + if pretrained: + from .model_store import get_model_file + device = torch.device(kwargs['local_rank']) + model.load_state_dict(torch.load(get_model_file('deeplabv3_%s_%s' % (backbone, acronyms[dataset]), root=root), + map_location=device)) + return model + + +def get_deeplabv3_resnet50_voc(**kwargs): + return get_deeplabv3('pascal_voc', 'resnet50', **kwargs) + + +def get_deeplabv3_resnet101_voc(**kwargs): + return get_deeplabv3('pascal_voc', 'resnet101', **kwargs) + + +def get_deeplabv3_resnet152_voc(**kwargs): + return get_deeplabv3('pascal_voc', 'resnet152', **kwargs) + + +def get_deeplabv3_resnet50_ade(**kwargs): + return get_deeplabv3('ade20k', 'resnet50', **kwargs) + + +def get_deeplabv3_resnet101_ade(**kwargs): + return get_deeplabv3('ade20k', 'resnet101', **kwargs) + + +def get_deeplabv3_resnet152_ade(**kwargs): + return get_deeplabv3('ade20k', 'resnet152', **kwargs) + + +if __name__ == '__main__': + model = get_deeplabv3_resnet50_voc() + img = torch.randn(2, 3, 480, 480) + output = model(img) diff --git a/hair_service_sd/seg/networks/deeplabv3_plus.py b/hair_service_sd/seg/networks/deeplabv3_plus.py new file mode 100644 index 0000000..d0636c6 --- /dev/null +++ b/hair_service_sd/seg/networks/deeplabv3_plus.py @@ -0,0 +1,160 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +from seg.networks.xception import get_xception +from seg.networks.deeplabv3 import _ASPP +from seg.networks.fcn import _FCNHead +from seg.networks.basic import _ConvBNReLU + +__all__ = ['DeepLabV3Plus', 'get_deeplabv3_plus', 'get_deeplabv3_plus_xception_voc'] + + +class DeepLabV3Plus(nn.Module): + r"""DeepLabV3Plus + Parameters + ---------- + nclass : int + Number of categories for the training dataset. + backbone : string + Pre-trained dilated backbone network type (default:'xception'). + norm_layer : object + Normalization layer used in backbone network (default: :class:`nn.BatchNorm`; + for Synchronized Cross-GPU BachNormalization). + aux : bool + Auxiliary loss. + + Reference: + Chen, Liang-Chieh, et al. "Encoder-Decoder with Atrous Separable Convolution for Semantic + Image Segmentation." + """ + + def __init__(self, nclass, backbone='xception', aux=True, pretrained_base=True, dilated=True, **kwargs): + super(DeepLabV3Plus, self).__init__() + self.aux = aux + self.nclass = nclass + output_stride = 8 if dilated else 32 + + self.pretrained = get_xception(pretrained=pretrained_base, output_stride=output_stride, **kwargs) + + # deeplabv3 plus + self.head = _DeepLabHead(nclass, **kwargs) + if aux: + self.auxlayer = _FCNHead(728, nclass, **kwargs) + + def base_forward(self, x): + # Entry flow + x = self.pretrained.conv1(x) + x = self.pretrained.bn1(x) + x = self.pretrained.relu(x) + + x = self.pretrained.conv2(x) + x = self.pretrained.bn2(x) + x = self.pretrained.relu(x) + + x = self.pretrained.block1(x) + # add relu here + x = self.pretrained.relu(x) + low_level_feat = x + + x = self.pretrained.block2(x) + x = self.pretrained.block3(x) + + # Middle flow + x = self.pretrained.midflow(x) + mid_level_feat = x + + # Exit flow + x = self.pretrained.block20(x) + x = self.pretrained.relu(x) + x = self.pretrained.conv3(x) + x = self.pretrained.bn3(x) + x = self.pretrained.relu(x) + + x = self.pretrained.conv4(x) + x = self.pretrained.bn4(x) + x = self.pretrained.relu(x) + + x = self.pretrained.conv5(x) + x = self.pretrained.bn5(x) + x = self.pretrained.relu(x) + return low_level_feat, mid_level_feat, x + + def forward(self, x): +# print("x size: ", x.size()) + size = x.size()[2:] + c1, c3, c4 = self.base_forward(x) +# print("c1 size: ", c1.size()) +# print("c4 size: ", c4.size()) + outputs = list() + x = self.head(c4, c1) + x = F.interpolate(x, size, mode='bilinear', align_corners=True) + outputs.append(x) + if self.aux: + auxout = self.auxlayer(c3) + auxout = F.interpolate(auxout, size, mode='bilinear', align_corners=True) + outputs.append(auxout) + + # for save onnx + # y = torch.max(x, 1)[1].to(torch.float32) + # return y + + return tuple(outputs) + + +class _DeepLabHead(nn.Module): + def __init__(self, nclass, c1_channels=128, norm_layer=nn.BatchNorm2d, **kwargs): + super(_DeepLabHead, self).__init__() + self.aspp = _ASPP(2048, [12, 24, 36], norm_layer=norm_layer, **kwargs) + self.c1_block = _ConvBNReLU(c1_channels, 48, 3, padding=1, norm_layer=norm_layer) + self.block = nn.Sequential( + _ConvBNReLU(304, 256, 3, padding=1, norm_layer=norm_layer), + nn.Dropout(0.5), + _ConvBNReLU(256, 256, 3, padding=1, norm_layer=norm_layer), + nn.Dropout(0.1), + nn.Conv2d(256, nclass, 1)) + + def forward(self, x, c1): + size = c1.size()[2:] + c1 = self.c1_block(c1) +# print("c1", c1.size()) +# print("before aspp: ", x.size()) + x = self.aspp(x) +# print("after aspp: ", x.size()) + x = F.interpolate(x, size, mode='bilinear', align_corners=True) + return self.block(torch.cat([x, c1], dim=1)) + + +def get_deeplabv3_plus(dataset='pascal_voc', backbone='xception', pretrained=False, root='../ckpt', + pretrained_base=False, nclass=3, **kwargs): + acronyms = { + 'pascal_voc': 'pascal_voc', + 'pascal_aug': 'pascal_aug', + 'ade20k': 'ade', + 'coco': 'coco', + 'citys': 'citys', + } + #from light.data import datasets + + model = DeepLabV3Plus(nclass, backbone=backbone, pretrained_base=pretrained_base, **kwargs) + if pretrained: + pass + # if dataset not in acronyms.keys(): + # print("root:", root) + # model_path = os.path.join(root, "deeplabv3_plus_28.pth") + # model.load_state_dict(torch.load(model_path), strict=False) + # else: + # from .model_store import get_model_file + # device = torch.device(kwargs['local_rank']) + # model.load_state_dict( + # torch.load(get_model_file('deeplabv3_plus_%s_%s' % (backbone, acronyms[dataset]), root=root), + # map_location=device)) + return model + + +def get_deeplabv3_plus_xception_voc(**kwargs): + return get_deeplabv3_plus('pascal_voc', 'xception', **kwargs) + + +if __name__ == '__main__': + model = get_deeplabv3_plus_xception_voc() diff --git a/hair_service_sd/seg/networks/fcn.py b/hair_service_sd/seg/networks/fcn.py new file mode 100644 index 0000000..fb1d981 --- /dev/null +++ b/hair_service_sd/seg/networks/fcn.py @@ -0,0 +1,222 @@ +import os +import torch +import torch.nn as nn +import torch.nn.functional as F + +from seg.networks.vgg import vgg16 + +__all__ = ['get_fcn32s', 'get_fcn16s', 'get_fcn8s', + 'get_fcn32s_vgg16_voc', 'get_fcn16s_vgg16_voc', 'get_fcn8s_vgg16_voc'] + + +class FCN32s(nn.Module): + """There are some difference from original fcn""" + + def __init__(self, nclass, backbone='vgg16', aux=False, pretrained_base=True, + norm_layer=nn.BatchNorm2d, **kwargs): + super(FCN32s, self).__init__() + self.aux = aux + if backbone == 'vgg16': + self.pretrained = vgg16(pretrained=pretrained_base).features + else: + raise RuntimeError('unknown backbone: {}'.format(backbone)) + self.head = _FCNHead(512, nclass, norm_layer) + if aux: + self.auxlayer = _FCNHead(512, nclass, norm_layer) + + self.__setattr__('exclusive', ['head', 'auxlayer'] if aux else ['head']) + + def forward(self, x): + size = x.size()[2:] + pool5 = self.pretrained(x) + + outputs = [] + out = self.head(pool5) + out = F.interpolate(out, size, mode='bilinear', align_corners=True) + outputs.append(out) + + if self.aux: + auxout = self.auxlayer(pool5) + auxout = F.interpolate(auxout, size, mode='bilinear', align_corners=True) + outputs.append(auxout) + + return tuple(outputs) + + +class FCN16s(nn.Module): + def __init__(self, nclass, backbone='vgg16', aux=False, pretrained_base=True, norm_layer=nn.BatchNorm2d, **kwargs): + super(FCN16s, self).__init__() + self.aux = aux + if backbone == 'vgg16': + self.pretrained = vgg16(pretrained=pretrained_base).features + else: + raise RuntimeError('unknown backbone: {}'.format(backbone)) + self.pool4 = nn.Sequential(*self.pretrained[:24]) + self.pool5 = nn.Sequential(*self.pretrained[24:]) + self.head = _FCNHead(512, nclass, norm_layer) + self.score_pool4 = nn.Conv2d(512, nclass, 1) + if aux: + self.auxlayer = _FCNHead(512, nclass, norm_layer) + + self.__setattr__('exclusive', ['head', 'score_pool4', 'auxlayer'] if aux else ['head', 'score_pool4']) + + def forward(self, x): + pool4 = self.pool4(x) + pool5 = self.pool5(pool4) + + outputs = [] + score_fr = self.head(pool5) + + score_pool4 = self.score_pool4(pool4) + + upscore2 = F.interpolate(score_fr, score_pool4.size()[2:], mode='bilinear', align_corners=True) + fuse_pool4 = upscore2 + score_pool4 + + out = F.interpolate(fuse_pool4, x.size()[2:], mode='bilinear', align_corners=True) + outputs.append(out) + + if self.aux: + auxout = self.auxlayer(pool5) + auxout = F.interpolate(auxout, x.size()[2:], mode='bilinear', align_corners=True) + outputs.append(auxout) + + return tuple(outputs) + + +class FCN8s(nn.Module): + def __init__(self, nclass, backbone='vgg16', aux=False, pretrained_base=True, norm_layer=nn.BatchNorm2d, **kwargs): + super(FCN8s, self).__init__() + self.aux = aux + if backbone == 'vgg16': + self.pretrained = vgg16(pretrained=pretrained_base).features + else: + raise RuntimeError('unknown backbone: {}'.format(backbone)) + self.pool3 = nn.Sequential(*self.pretrained[:17]) + self.pool4 = nn.Sequential(*self.pretrained[17:24]) + self.pool5 = nn.Sequential(*self.pretrained[24:]) + self.head = _FCNHead(512, nclass, norm_layer) + self.score_pool3 = nn.Conv2d(256, nclass, 1) + self.score_pool4 = nn.Conv2d(512, nclass, 1) + if aux: + self.auxlayer = _FCNHead(512, nclass, norm_layer) + + self.__setattr__('exclusive', + ['head', 'score_pool3', 'score_pool4', 'auxlayer'] if aux else ['head', 'score_pool3', + 'score_pool4']) + + def forward(self, x): + pool3 = self.pool3(x) + pool4 = self.pool4(pool3) + pool5 = self.pool5(pool4) + + outputs = [] + score_fr = self.head(pool5) + + score_pool4 = self.score_pool4(pool4) + score_pool3 = self.score_pool3(pool3) + + upscore2 = F.interpolate(score_fr, score_pool4.size()[2:], mode='bilinear', align_corners=True) + fuse_pool4 = upscore2 + score_pool4 + + upscore_pool4 = F.interpolate(fuse_pool4, score_pool3.size()[2:], mode='bilinear', align_corners=True) + fuse_pool3 = upscore_pool4 + score_pool3 + + out = F.interpolate(fuse_pool3, x.size()[2:], mode='bilinear', align_corners=True) + outputs.append(out) + + if self.aux: + auxout = self.auxlayer(pool5) + auxout = F.interpolate(auxout, x.size()[2:], mode='bilinear', align_corners=True) + outputs.append(auxout) + + return tuple(outputs) + + +class _FCNHead(nn.Module): + def __init__(self, in_channels, channels, norm_layer=nn.BatchNorm2d, **kwargs): + super(_FCNHead, self).__init__() + inter_channels = in_channels // 4 + self.block = nn.Sequential( + nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False), + norm_layer(inter_channels), + nn.ReLU(inplace=True), + nn.Dropout(0.1), + nn.Conv2d(inter_channels, channels, 1) + ) + + def forward(self, x): + return self.block(x) + + +def get_fcn32s(dataset='pascal_voc', backbone='vgg16', pretrained=False, root='~/.torch/models', + pretrained_base=True, **kwargs): + acronyms = { + 'pascal_voc': 'pascal_voc', + 'pascal_aug': 'pascal_aug', + 'ade20k': 'ade', + 'coco': 'coco', + 'citys': 'citys', + } + from ..data.dataloader import datasets + model = FCN32s(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs) + if pretrained: + from .model_store import get_model_file + device = torch.device(kwargs['local_rank']) + model.load_state_dict(torch.load(get_model_file('fcn32s_%s_%s' % (backbone, acronyms[dataset]), root=root), + map_location=device)) + return model + + +def get_fcn16s(dataset='pascal_voc', backbone='vgg16', pretrained=False, root='~/.torch/models', + pretrained_base=True, **kwargs): + acronyms = { + 'pascal_voc': 'pascal_voc', + 'pascal_aug': 'pascal_aug', + 'ade20k': 'ade', + 'coco': 'coco', + 'citys': 'citys', + } + from ..data.dataloader import datasets + model = FCN16s(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs) + if pretrained: + from .model_store import get_model_file + device = torch.device(kwargs['local_rank']) + model.load_state_dict(torch.load(get_model_file('fcn16s_%s_%s' % (backbone, acronyms[dataset]), root=root), + map_location=device)) + return model + + +def get_fcn8s(dataset='pascal_voc', backbone='vgg16', pretrained=False, root='~/.torch/models', + pretrained_base=True, **kwargs): + acronyms = { + 'pascal_voc': 'pascal_voc', + 'pascal_aug': 'pascal_aug', + 'ade20k': 'ade', + 'coco': 'coco', + 'citys': 'citys', + } + from ..data.dataloader import datasets + model = FCN8s(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs) + if pretrained: + from .model_store import get_model_file + device = torch.device(kwargs['local_rank']) + model.load_state_dict(torch.load(get_model_file('fcn8s_%s_%s' % (backbone, acronyms[dataset]), root=root), + map_location=device)) + return model + + +def get_fcn32s_vgg16_voc(**kwargs): + return get_fcn32s('pascal_voc', 'vgg16', **kwargs) + + +def get_fcn16s_vgg16_voc(**kwargs): + return get_fcn16s('pascal_voc', 'vgg16', **kwargs) + + +def get_fcn8s_vgg16_voc(**kwargs): + return get_fcn8s('pascal_voc', 'vgg16', **kwargs) + + +if __name__ == '__main__': + model = FCN16s(21) + print(model) diff --git a/hair_service_sd/seg/networks/jpu.py b/hair_service_sd/seg/networks/jpu.py new file mode 100644 index 0000000..db23bab --- /dev/null +++ b/hair_service_sd/seg/networks/jpu.py @@ -0,0 +1,68 @@ +"""Joint Pyramid Upsampling""" +import torch +import torch.nn as nn +import torch.nn.functional as F + +__all__ = ['JPU'] + + +class SeparableConv2d(nn.Module): + def __init__(self, inplanes, planes, kernel_size=3, stride=1, padding=1, + dilation=1, bias=False, norm_layer=nn.BatchNorm2d): + super(SeparableConv2d, self).__init__() + self.conv = nn.Conv2d(inplanes, inplanes, kernel_size, stride, padding, dilation, groups=inplanes, bias=bias) + self.bn = norm_layer(inplanes) + self.pointwise = nn.Conv2d(inplanes, planes, 1, bias=bias) + + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + x = self.pointwise(x) + return x + + +# copy from: https://github.com/wuhuikai/FastFCN/blob/master/encoding/nn/customize.py +class JPU(nn.Module): + def __init__(self, in_channels, width=512, norm_layer=nn.BatchNorm2d, **kwargs): + super(JPU, self).__init__() + + self.conv5 = nn.Sequential( + nn.Conv2d(in_channels[-1], width, 3, padding=1, bias=False), + norm_layer(width), + nn.ReLU(True)) + self.conv4 = nn.Sequential( + nn.Conv2d(in_channels[-2], width, 3, padding=1, bias=False), + norm_layer(width), + nn.ReLU(True)) + self.conv3 = nn.Sequential( + nn.Conv2d(in_channels[-3], width, 3, padding=1, bias=False), + norm_layer(width), + nn.ReLU(True)) + + self.dilation1 = nn.Sequential( + SeparableConv2d(3 * width, width, 3, padding=1, dilation=1, bias=False), + norm_layer(width), + nn.ReLU(True)) + self.dilation2 = nn.Sequential( + SeparableConv2d(3 * width, width, 3, padding=2, dilation=2, bias=False), + norm_layer(width), + nn.ReLU(True)) + self.dilation3 = nn.Sequential( + SeparableConv2d(3 * width, width, 3, padding=4, dilation=4, bias=False), + norm_layer(width), + nn.ReLU(True)) + self.dilation4 = nn.Sequential( + SeparableConv2d(3 * width, width, 3, padding=8, dilation=8, bias=False), + norm_layer(width), + nn.ReLU(True)) + + def forward(self, *inputs): + feats = [self.conv5(inputs[-1]), self.conv4(inputs[-2]), self.conv3(inputs[-3])] + size = feats[-1].size()[2:] + feats[-2] = F.interpolate(feats[-2], size, mode='bilinear', align_corners=True) + feats[-3] = F.interpolate(feats[-3], size, mode='bilinear', align_corners=True) + feat = torch.cat(feats, dim=1) + feat = torch.cat([self.dilation1(feat), self.dilation2(feat), self.dilation3(feat), self.dilation4(feat)], + dim=1) + + return inputs[0], inputs[1], inputs[2], feat diff --git a/hair_service_sd/seg/networks/resnetv1b.py b/hair_service_sd/seg/networks/resnetv1b.py new file mode 100644 index 0000000..21d67b7 --- /dev/null +++ b/hair_service_sd/seg/networks/resnetv1b.py @@ -0,0 +1,264 @@ +import torch +import torch.nn as nn +import torch.utils.model_zoo as model_zoo + +__all__ = ['ResNetV1b', 'resnet18_v1b', 'resnet34_v1b', 'resnet50_v1b', + 'resnet101_v1b', 'resnet152_v1b', 'resnet152_v1s', 'resnet101_v1s', 'resnet50_v1s'] + +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', +} + + +class BasicBlockV1b(nn.Module): + expansion = 1 + + def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None, + previous_dilation=1, norm_layer=nn.BatchNorm2d): + super(BasicBlockV1b, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, 3, stride, + dilation, dilation, bias=False) + self.bn1 = norm_layer(planes) + self.relu = nn.ReLU(True) + self.conv2 = nn.Conv2d(planes, planes, 3, 1, previous_dilation, + dilation=previous_dilation, bias=False) + 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 BottleneckV1b(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None, + previous_dilation=1, norm_layer=nn.BatchNorm2d): + super(BottleneckV1b, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False) + self.bn1 = norm_layer(planes) + self.conv2 = nn.Conv2d(planes, planes, 3, stride, + dilation, dilation, bias=False) + self.bn2 = norm_layer(planes) + self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False) + self.bn3 = norm_layer(planes * self.expansion) + self.relu = nn.ReLU(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 ResNetV1b(nn.Module): + + def __init__(self, block, layers, num_classes=1000, dilated=True, deep_stem=False, + zero_init_residual=False, norm_layer=nn.BatchNorm2d): + self.inplanes = 128 if deep_stem else 64 + super(ResNetV1b, self).__init__() + if deep_stem: + self.conv1 = nn.Sequential( + nn.Conv2d(3, 64, 3, 2, 1, bias=False), + norm_layer(64), + nn.ReLU(True), + nn.Conv2d(64, 64, 3, 1, 1, bias=False), + norm_layer(64), + nn.ReLU(True), + nn.Conv2d(64, 128, 3, 1, 1, bias=False) + ) + else: + self.conv1 = nn.Conv2d(3, 64, 7, 2, 3, bias=False) + self.bn1 = norm_layer(self.inplanes) + self.relu = nn.ReLU(True) + self.maxpool = nn.MaxPool2d(3, 2, 1) + self.layer1 = self._make_layer(block, 64, layers[0], norm_layer=norm_layer) + self.layer2 = self._make_layer(block, 128, layers[1], stride=2, norm_layer=norm_layer) + if dilated: + self.layer3 = self._make_layer(block, 256, layers[2], stride=1, dilation=2, norm_layer=norm_layer) + self.layer4 = self._make_layer(block, 512, layers[3], stride=1, dilation=4, norm_layer=norm_layer) + else: + self.layer3 = self._make_layer(block, 256, layers[2], stride=2, norm_layer=norm_layer) + self.layer4 = self._make_layer(block, 512, layers[3], stride=2, norm_layer=norm_layer) + 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.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + if zero_init_residual: + for m in self.modules(): + if isinstance(m, BottleneckV1b): + nn.init.constant_(m.bn3.weight, 0) + elif isinstance(m, BasicBlockV1b): + nn.init.constant_(m.bn2.weight, 0) + + def _make_layer(self, block, planes, blocks, stride=1, dilation=1, norm_layer=nn.BatchNorm2d): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.inplanes, planes * block.expansion, 1, stride, bias=False), + norm_layer(planes * block.expansion), + ) + + layers = [] + if dilation in (1, 2): + layers.append(block(self.inplanes, planes, stride, dilation=1, downsample=downsample, + previous_dilation=dilation, norm_layer=norm_layer)) + elif dilation == 4: + layers.append(block(self.inplanes, planes, stride, dilation=2, downsample=downsample, + previous_dilation=dilation, norm_layer=norm_layer)) + else: + raise RuntimeError("=> unknown dilation size: {}".format(dilation)) + self.inplanes = planes * block.expansion + for _ in range(1, blocks): + layers.append(block(self.inplanes, planes, dilation=dilation, + previous_dilation=dilation, norm_layer=norm_layer)) + + return nn.Sequential(*layers) + + def forward(self, x): + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + x = self.maxpool(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + + x = self.avgpool(x) + x = x.view(x.size(0), -1) + x = self.fc(x) + + return x + + +def resnet18_v1b(pretrained=False, **kwargs): + model = ResNetV1b(BasicBlockV1b, [2, 2, 2, 2], **kwargs) + if pretrained: + old_dict = model_zoo.load_url(model_urls['resnet18']) + model_dict = model.state_dict() + old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)} + model_dict.update(old_dict) + model.load_state_dict(model_dict) + return model + + +def resnet34_v1b(pretrained=False, **kwargs): + model = ResNetV1b(BasicBlockV1b, [3, 4, 6, 3], **kwargs) + if pretrained: + old_dict = model_zoo.load_url(model_urls['resnet34']) + model_dict = model.state_dict() + old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)} + model_dict.update(old_dict) + model.load_state_dict(model_dict) + return model + + +def resnet50_v1b(pretrained=False, **kwargs): + model = ResNetV1b(BottleneckV1b, [3, 4, 6, 3], **kwargs) + if pretrained: + old_dict = model_zoo.load_url(model_urls['resnet50']) + model_dict = model.state_dict() + old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)} + model_dict.update(old_dict) + model.load_state_dict(model_dict) + return model + + +def resnet101_v1b(pretrained=False, **kwargs): + model = ResNetV1b(BottleneckV1b, [3, 4, 23, 3], **kwargs) + if pretrained: + old_dict = model_zoo.load_url(model_urls['resnet101']) + model_dict = model.state_dict() + old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)} + model_dict.update(old_dict) + model.load_state_dict(model_dict) + return model + + +def resnet152_v1b(pretrained=False, **kwargs): + model = ResNetV1b(BottleneckV1b, [3, 8, 36, 3], **kwargs) + if pretrained: + old_dict = model_zoo.load_url(model_urls['resnet152']) + model_dict = model.state_dict() + old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)} + model_dict.update(old_dict) + model.load_state_dict(model_dict) + return model + + +def resnet50_v1s(pretrained=False, root='~/.torch/models', **kwargs): + model = ResNetV1b(BottleneckV1b, [3, 4, 6, 3], deep_stem=True, **kwargs) + if pretrained: + from ..model_store import get_resnet_file + model.load_state_dict(torch.load(get_resnet_file('resnet50', root=root)), strict=False) + return model + + +def resnet101_v1s(pretrained=False, root='~/.torch/models', **kwargs): + model = ResNetV1b(BottleneckV1b, [3, 4, 23, 3], deep_stem=True, **kwargs) + if pretrained: + from ..model_store import get_resnet_file + model.load_state_dict(torch.load(get_resnet_file('resnet101', root=root)), strict=False) + return model + + +def resnet152_v1s(pretrained=False, root='~/.torch/models', **kwargs): + model = ResNetV1b(BottleneckV1b, [3, 8, 36, 3], deep_stem=True, **kwargs) + if pretrained: + from ..model_store import get_resnet_file + model.load_state_dict(torch.load(get_resnet_file('resnet152', root=root)), strict=False) + return model + + +if __name__ == '__main__': + import torch + + img = torch.randn(4, 3, 224, 224) + model = resnet50_v1b(True) + output = model(img) diff --git a/hair_service_sd/seg/networks/segbase.py b/hair_service_sd/seg/networks/segbase.py new file mode 100644 index 0000000..75e8f64 --- /dev/null +++ b/hair_service_sd/seg/networks/segbase.py @@ -0,0 +1,60 @@ +"""Base Model for Semantic Segmentation""" +import torch.nn as nn + +from seg.networks.jpu import JPU +from seg.networks.resnetv1b import resnet50_v1s, resnet101_v1s, resnet152_v1s + +__all__ = ['SegBaseModel'] + + +class SegBaseModel(nn.Module): + r"""Base Model for Semantic Segmentation + + Parameters + ---------- + backbone : string + Pre-trained dilated backbone network type (default:'resnet50'; 'resnet50', + 'resnet101' or 'resnet152'). + """ + + def __init__(self, nclass, aux, backbone='resnet50', jpu=False, pretrained_base=True, **kwargs): + super(SegBaseModel, self).__init__() + dilated = False if jpu else True + self.aux = aux + self.nclass = nclass + if backbone == 'resnet50': + self.pretrained = resnet50_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs) + elif backbone == 'resnet101': + self.pretrained = resnet101_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs) + elif backbone == 'resnet152': + self.pretrained = resnet152_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs) + else: + raise RuntimeError('unknown backbone: {}'.format(backbone)) + + self.jpu = JPU([512, 1024, 2048], width=512, **kwargs) if jpu else None + + def base_forward(self, x): + """forwarding pre-trained network""" + x = self.pretrained.conv1(x) + x = self.pretrained.bn1(x) + x = self.pretrained.relu(x) + x = self.pretrained.maxpool(x) + c1 = self.pretrained.layer1(x) + c2 = self.pretrained.layer2(c1) + c3 = self.pretrained.layer3(c2) + c4 = self.pretrained.layer4(c3) + + if self.jpu: + return self.jpu(c1, c2, c3, c4) + else: + return c1, c2, c3, c4 + + def evaluate(self, x): + """evaluating network with inputs and targets""" + return self.forward(x)[0] + + def demo(self, x): + pred = self.forward(x) + if self.aux: + pred = pred[0] + return pred diff --git a/hair_service_sd/seg/networks/vgg.py b/hair_service_sd/seg/networks/vgg.py new file mode 100644 index 0000000..fe5c163 --- /dev/null +++ b/hair_service_sd/seg/networks/vgg.py @@ -0,0 +1,191 @@ +import torch +import torch.nn as nn +import torch.utils.model_zoo as model_zoo + +__all__ = [ + 'VGG', 'vgg11', 'vgg11_bn', 'vgg13', 'vgg13_bn', 'vgg16', 'vgg16_bn', + 'vgg19_bn', 'vgg19', +] + +model_urls = { + 'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth', + 'vgg13': 'https://download.pytorch.org/models/vgg13-c768596a.pth', + 'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth', + 'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth', + 'vgg11_bn': 'https://download.pytorch.org/models/vgg11_bn-6002323d.pth', + 'vgg13_bn': 'https://download.pytorch.org/models/vgg13_bn-abd245e5.pth', + 'vgg16_bn': 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth', + 'vgg19_bn': 'https://download.pytorch.org/models/vgg19_bn-c79401a0.pth', +} + + +class VGG(nn.Module): + def __init__(self, features, num_classes=1000, init_weights=True): + super(VGG, self).__init__() + self.features = features + self.avgpool = nn.AdaptiveAvgPool2d((7, 7)) + self.classifier = nn.Sequential( + nn.Linear(512 * 7 * 7, 4096), + nn.ReLU(True), + nn.Dropout(), + nn.Linear(4096, 4096), + nn.ReLU(True), + nn.Dropout(), + nn.Linear(4096, num_classes) + ) + if init_weights: + self._initialize_weights() + + def forward(self, x): + x = self.features(x) + x = self.avgpool(x) + x = x.view(x.size(0), -1) + x = self.classifier(x) + return x + + def _initialize_weights(self): + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.BatchNorm2d): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.Linear): + nn.init.normal_(m.weight, 0, 0.01) + nn.init.constant_(m.bias, 0) + + +def make_layers(cfg, batch_norm=False): + layers = [] + in_channels = 3 + for v in cfg: + if v == 'M': + layers += [nn.MaxPool2d(kernel_size=2, stride=2)] + else: + conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1) + if batch_norm: + layers += (conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)) + else: + layers += [conv2d, nn.ReLU(inplace=True)] + in_channels = v + return nn.Sequential(*layers) + + +cfg = { + 'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], + 'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], + 'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'], + 'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'], +} + + +def vgg11(pretrained=False, **kwargs): + """VGG 11-layer model (configuration "A") + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['A']), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg11'])) + return model + + +def vgg11_bn(pretrained=False, **kwargs): + """VGG 11-layer model (configuration "A") with batch normalization + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['A'], batch_norm=True), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg11_bn'])) + return model + + +def vgg13(pretrained=False, **kwargs): + """VGG 13-layer model (configuration "B") + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['B']), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg13'])) + return model + + +def vgg13_bn(pretrained=False, **kwargs): + """VGG 13-layer model (configuration "B") with batch normalization + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['B'], batch_norm=True), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg13_bn'])) + return model + + +def vgg16(pretrained=False, **kwargs): + """VGG 16-layer model (configuration "D") + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['D']), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg16'])) + return model + + +def vgg16_bn(pretrained=False, **kwargs): + """VGG 16-layer model (configuration "D") with batch normalization + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['D'], batch_norm=True), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg16_bn'])) + return model + + +def vgg19(pretrained=False, **kwargs): + """VGG 19-layer model (configuration "E") + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['E']), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg19'])) + return model + + +def vgg19_bn(pretrained=False, **kwargs): + """VGG 19-layer model (configuration 'E') with batch normalization + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + if pretrained: + kwargs['init_weights'] = False + model = VGG(make_layers(cfg['E'], batch_norm=True), **kwargs) + if pretrained: + model.load_state_dict(model_zoo.load_url(model_urls['vgg19_bn'])) + return model + + +if __name__ == '__main__': + img = torch.randn((4, 3, 480, 480)) + model = vgg16(pretrained=False) + out = model(img) diff --git a/hair_service_sd/seg/networks/xception.py b/hair_service_sd/seg/networks/xception.py new file mode 100644 index 0000000..52dc0b9 --- /dev/null +++ b/hair_service_sd/seg/networks/xception.py @@ -0,0 +1,411 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +__all__ = ['Enc', 'FCAttention', 'Xception65', 'Xception71', 'get_xception', 'get_xception_71', 'get_xception_a'] + + +class SeparableConv2d(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, dilation=1, bias=False, norm_layer=None): + super(SeparableConv2d, self).__init__() + self.kernel_size = kernel_size + self.dilation = dilation + + self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size, stride, 0, dilation, groups=in_channels, + bias=bias) + self.bn = norm_layer(in_channels) + self.pointwise = nn.Conv2d(in_channels, out_channels, 1, bias=bias) + + def forward(self, x): + x = self.fix_padding(x, self.kernel_size, self.dilation) + x = self.conv1(x) + x = self.bn(x) + x = self.pointwise(x) + + return x + + def fix_padding(self, x, 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(x, (pad_beg, pad_end, pad_beg, pad_end)) + return padded_inputs + + +class Block(nn.Module): + def __init__(self, in_channels, out_channels, reps, stride=1, dilation=1, norm_layer=None, + start_with_relu=True, grow_first=True, is_last=False): + super(Block, self).__init__() + if out_channels != in_channels or stride != 1: + self.skip = nn.Conv2d(in_channels, out_channels, 1, stride, bias=False) + self.skipbn = norm_layer(out_channels) + else: + self.skip = None + self.relu = nn.ReLU(True) + rep = list() + filters = in_channels + if grow_first: + if start_with_relu: + rep.append(self.relu) + rep.append(SeparableConv2d(in_channels, out_channels, 3, 1, dilation, norm_layer=norm_layer)) + rep.append(norm_layer(out_channels)) + filters = out_channels + for i in range(reps - 1): + if grow_first or start_with_relu: + rep.append(self.relu) + rep.append(SeparableConv2d(filters, filters, 3, 1, dilation, norm_layer=norm_layer)) + rep.append(norm_layer(filters)) + if not grow_first: + rep.append(self.relu) + rep.append(SeparableConv2d(in_channels, out_channels, 3, 1, dilation, norm_layer=norm_layer)) + if stride != 1: + rep.append(self.relu) + rep.append(SeparableConv2d(out_channels, out_channels, 3, stride, norm_layer=norm_layer)) + rep.append(norm_layer(out_channels)) + elif is_last: + rep.append(self.relu) + rep.append(SeparableConv2d(out_channels, out_channels, 3, 1, dilation, norm_layer=norm_layer)) + rep.append(norm_layer(out_channels)) + self.rep = nn.Sequential(*rep) + + def forward(self, x): + out = self.rep(x) + if self.skip is not None: + skip = self.skipbn(self.skip(x)) + else: + skip = x + out = out + skip + return out + + +class Xception65(nn.Module): + """Modified Aligned Xception + """ + + def __init__(self, num_classes=1000, output_stride=32, norm_layer=nn.BatchNorm2d): + super(Xception65, self).__init__() + if output_stride == 32: + entry_block3_stride = 2 + exit_block20_stride = 2 + middle_block_dilation = 1 + exit_block_dilations = (1, 1) + elif output_stride == 16: + entry_block3_stride = 2 + exit_block20_stride = 1 + middle_block_dilation = 1 + exit_block_dilations = (1, 2) + elif output_stride == 8: + entry_block3_stride = 1 + exit_block20_stride = 1 + middle_block_dilation = 2 + exit_block_dilations = (2, 4) + else: + raise NotImplementedError + # Entry flow + self.conv1 = nn.Conv2d(3, 32, 3, 2, 1, bias=False) + self.bn1 = norm_layer(32) + self.relu = nn.ReLU(True) + + self.conv2 = nn.Conv2d(32, 64, 3, 1, 1, bias=False) + self.bn2 = norm_layer(64) + + self.block1 = Block(64, 128, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False) + self.block2 = Block(128, 256, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False, grow_first=True) + self.block3 = Block(256, 728, reps=2, stride=entry_block3_stride, norm_layer=norm_layer, + start_with_relu=True, grow_first=True, is_last=True) + + # Middle flow + midflow = list() + for i in range(4, 20): + midflow.append(Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation, norm_layer=norm_layer, + start_with_relu=True, grow_first=True)) + self.midflow = nn.Sequential(*midflow) + + # Exit flow + self.block20 = Block(728, 1024, reps=2, stride=exit_block20_stride, dilation=exit_block_dilations[0], + norm_layer=norm_layer, start_with_relu=True, grow_first=False, is_last=True) + self.conv3 = SeparableConv2d(1024, 1536, 3, 1, dilation=exit_block_dilations[1], norm_layer=norm_layer) + self.bn3 = norm_layer(1536) + self.conv4 = SeparableConv2d(1536, 1536, 3, stride=1, dilation=exit_block_dilations[1], norm_layer=norm_layer) + self.bn4 = norm_layer(1536) + self.conv5 = SeparableConv2d(1536, 2048, 3, 1, dilation=exit_block_dilations[1], norm_layer=norm_layer) + self.bn5 = norm_layer(2048) + self.avgpool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Linear(2048, num_classes) + + 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) + x = self.relu(x) + # c1 = x + x = self.block2(x) + # c2 = x + x = self.block3(x) + + # Middle flow + x = self.midflow(x) + # c3 = 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) + + x = self.avgpool(x) + x = x.view(x.size(0), -1) + x = self.fc(x) + + return x + + +class Xception71(nn.Module): + """Modified Aligned Xception + """ + + def __init__(self, num_classes=1000, output_stride=32, norm_layer=nn.BatchNorm2d): + super(Xception71, self).__init__() + if output_stride == 32: + entry_block3_stride = 2 + exit_block20_stride = 2 + middle_block_dilation = 1 + exit_block_dilations = (1, 1) + elif output_stride == 16: + entry_block3_stride = 2 + exit_block20_stride = 1 + middle_block_dilation = 1 + exit_block_dilations = (1, 2) + elif output_stride == 8: + entry_block3_stride = 1 + exit_block20_stride = 1 + middle_block_dilation = 2 + exit_block_dilations = (2, 4) + else: + raise NotImplementedError + # Entry flow + self.conv1 = nn.Conv2d(3, 32, 3, 2, 1, bias=False) + self.bn1 = norm_layer(32) + self.relu = nn.ReLU(True) + + self.conv2 = nn.Conv2d(32, 64, 3, 1, 1, bias=False) + self.bn2 = norm_layer(64) + + self.block1 = Block(64, 128, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False) + self.block2 = nn.Sequential( + Block(128, 256, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False, grow_first=True), + Block(256, 728, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False, grow_first=True)) + self.block3 = Block(728, 728, reps=2, stride=entry_block3_stride, norm_layer=norm_layer, + start_with_relu=True, grow_first=True, is_last=True) + + # Middle flow + midflow = list() + for i in range(4, 20): + midflow.append(Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation, norm_layer=norm_layer, + start_with_relu=True, grow_first=True)) + self.midflow = nn.Sequential(*midflow) + + # Exit flow + self.block20 = Block(728, 1024, reps=2, stride=exit_block20_stride, dilation=exit_block_dilations[0], + norm_layer=norm_layer, start_with_relu=True, grow_first=False, is_last=True) + self.conv3 = SeparableConv2d(1024, 1536, 3, 1, dilation=exit_block_dilations[1], norm_layer=norm_layer) + self.bn3 = norm_layer(1536) + self.conv4 = SeparableConv2d(1536, 1536, 3, stride=1, dilation=exit_block_dilations[1], norm_layer=norm_layer) + self.bn4 = norm_layer(1536) + self.conv5 = SeparableConv2d(1536, 2048, 3, 1, dilation=exit_block_dilations[1], norm_layer=norm_layer) + self.bn5 = norm_layer(2048) + self.avgpool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Linear(2048, num_classes) + + 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) + x = self.relu(x) + # c1 = x + x = self.block2(x) + # c2 = x + x = self.block3(x) + + # Middle flow + x = self.midflow(x) + # c3 = 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) + + x = self.avgpool(x) + x = x.view(x.size(0), -1) + x = self.fc(x) + + return x + + +# ------------------------------------------------- +# For DFANet +# ------------------------------------------------- +class BlockA(nn.Module): + def __init__(self, in_channels, out_channels, stride=1, dilation=1, norm_layer=None, start_with_relu=True): + super(BlockA, self).__init__() + if out_channels != in_channels or stride != 1: + self.skip = nn.Conv2d(in_channels, out_channels, 1, stride, bias=False) + self.skipbn = norm_layer(out_channels) + else: + self.skip = None + self.relu = nn.ReLU(True) + rep = list() + inter_channels = out_channels // 4 + + if start_with_relu: + rep.append(self.relu) + rep.append(SeparableConv2d(in_channels, inter_channels, 3, 1, dilation, norm_layer=norm_layer)) + rep.append(norm_layer(inter_channels)) + + rep.append(self.relu) + rep.append(SeparableConv2d(inter_channels, inter_channels, 3, 1, dilation, norm_layer=norm_layer)) + rep.append(norm_layer(inter_channels)) + + if stride != 1: + rep.append(self.relu) + rep.append(SeparableConv2d(inter_channels, out_channels, 3, stride, norm_layer=norm_layer)) + rep.append(norm_layer(out_channels)) + else: + rep.append(self.relu) + rep.append(SeparableConv2d(inter_channels, out_channels, 3, 1, norm_layer=norm_layer)) + rep.append(norm_layer(out_channels)) + self.rep = nn.Sequential(*rep) + + def forward(self, x): + out = self.rep(x) + if self.skip is not None: + skip = self.skipbn(self.skip(x)) + else: + skip = x + out = out + skip + return out + + +class Enc(nn.Module): + def __init__(self, in_channels, out_channels, blocks, norm_layer=None): + super(Enc, self).__init__() + block = list() + block.append(BlockA(in_channels, out_channels, 2, norm_layer=norm_layer)) + for i in range(blocks - 1): + block.append(BlockA(out_channels, out_channels, 1, norm_layer=norm_layer)) + self.block = nn.Sequential(*block) + + def forward(self, x): + return self.block(x) + + +class FCAttention(nn.Module): + def __init__(self, in_channels, norm_layer=None): + super(FCAttention, self).__init__() + self.avgpool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Linear(in_channels, 1000) + self.conv = nn.Sequential( + nn.Conv2d(1000, in_channels, 1, bias=False), + norm_layer(in_channels), + nn.ReLU(True)) + + def forward(self, x): + n, c, _, _ = x.size() + att = self.avgpool(x).view(n, c) + att = self.fc(att).view(n, 1000, 1, 1) + att = self.conv(att) + return x * att.expand_as(x) + + +class XceptionA(nn.Module): + def __init__(self, num_classes=1000, norm_layer=nn.BatchNorm2d): + super(XceptionA, self).__init__() + self.conv1 = nn.Sequential(nn.Conv2d(3, 8, 3, 2, 1, bias=False), + norm_layer(8), + nn.ReLU(True)) + + self.enc2 = Enc(8, 48, 4, norm_layer=norm_layer) + self.enc3 = Enc(48, 96, 6, norm_layer=norm_layer) + self.enc4 = Enc(96, 192, 4, norm_layer=norm_layer) + + self.fca = FCAttention(192, norm_layer=norm_layer) + self.avgpool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Linear(192, num_classes) + + def forward(self, x): + x = self.conv1(x) + + x = self.enc2(x) + x = self.enc3(x) + x = self.enc4(x) + x = self.fca(x) + + x = self.avgpool(x) + x = x.view(x.size(0), -1) + x = self.fc(x) + + return x + + +# Constructor +def get_xception(pretrained=False, root='~/.torch/models', **kwargs): + model = Xception65(**kwargs) + if pretrained: + from ..model_store import get_model_file + model.load_state_dict(torch.load(get_model_file('xception', root=root))) + return model + + +def get_xception_71(pretrained=False, root='~/.torch/models', **kwargs): + model = Xception71(**kwargs) + if pretrained: + from ..model_store import get_model_file + model.load_state_dict(torch.load(get_model_file('xception71', root=root))) + return model + + +def get_xception_a(pretrained=False, root='~/.torch/models', **kwargs): + model = XceptionA(**kwargs) + if pretrained: + from ..model_store import get_model_file + model.load_state_dict(torch.load(get_model_file('xception_a', root=root))) + return model + + +if __name__ == '__main__': + model = get_xception_a() diff --git a/hair_service_sd/seg/setup.py b/hair_service_sd/seg/setup.py new file mode 100644 index 0000000..ca6bf5a --- /dev/null +++ b/hair_service_sd/seg/setup.py @@ -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 \ No newline at end of file diff --git a/hair_service_sd/seg/test_single.py b/hair_service_sd/seg/test_single.py new file mode 100644 index 0000000..0c9162d --- /dev/null +++ b/hair_service_sd/seg/test_single.py @@ -0,0 +1,24 @@ +from seg.hairseg_single_model import Evaluator +import os +import cv2 +import numpy as np +if __name__ == "__main__": + + data_path = "/home/liyang/project/matting/合格/origin" + dst_path = "/home/liyang/project/matting/合格/origin_seg_res1102" + if not os.path.exists(dst_path): + os.mkdir(dst_path) + seg_model = Evaluator(gpu_id=0, output_img_size=512, nclass=3) + for imgs in os.listdir(data_path): + if imgs.endswith(".txt"): + continue + + img_path = os.path.join(data_path, imgs) + img = cv2.imread(img_path) + if img_path.endswith('.png'): + kpts_1k = np.loadtxt(img_path.replace('.png', '_landmark1k.txt')) + elif img_path.endswith('.jpg'): + kpts_1k = np.loadtxt(img_path.replace('.jpg', '_landmark1k.txt')) + output_img_size = 512 + mask = seg_model.eval(img, kpts_1k) + cv2.imwrite(os.path.join(dst_path, imgs), mask) diff --git a/hair_service_sd/step05_detect_fa_hairmatting_inplace.py b/hair_service_sd/step05_detect_fa_hairmatting_inplace.py new file mode 100644 index 0000000..711bbda --- /dev/null +++ b/hair_service_sd/step05_detect_fa_hairmatting_inplace.py @@ -0,0 +1,101 @@ +import os +import pickle +import sys +import torch +import cv2 +import time +from multiprocessing import Process, Queue +from utils import landmark_processor +import numpy as np + + +def process_thread(sq, gpu_id): + from momocv.BigResNetStable import MomocvFaceAlignment + from models.detector import RetinaFaceDetector + from models.MomocvFaceAlignment1K import MomocvFaceAlignment1K + from hair_matting.Generator_Matte import Generator_Matte + + os.environ['CUDA_VISIBLE_DEVICES'] = f'{gpu_id}' + torch.set_grad_enabled(False) + face_detector = RetinaFaceDetector(gpu_id=0) + face_alignmenter_1k = MomocvFaceAlignment1K(gpu_id=0) + mmcv = MomocvFaceAlignment(gpu_id=gpu_id) + generator_matte = Generator_Matte(gpu=True, device_id=0) + + while True: + try: + img_file_path = sq.get(timeout=0.5) + except Exception as e: + break + try: + img = cv2.imread(img_file_path) + print("pkl img_file_path:", img_file_path) + bounding_boxes, landmarks = face_detector.forward(img, min_face_size=50) + + if len(bounding_boxes) == 0: continue + box_index = landmark_processor.get_max_rect(bounding_boxes) + max_face_rect = bounding_boxes[box_index].astype(np.int32) + pt1k = face_alignmenter_1k.stable_forward(img, [max_face_rect], reset=True)[0] + # for pt in pt1k.astype(np.int32): + # cv2.circle(img, (pt[0], pt[1]), 1, (0, 255, 0), -1) + # cv2.imshow("img", img) + # cv2.waitKey(0) + + for iter in range(3): + crop_size = 384 + image_to_face_mat = landmark_processor.get_transform_mat_full_face(pt1k, crop_size) + crop_face = cv2.warpAffine(img, image_to_face_mat, (crop_size, crop_size), flags=cv2.INTER_LANCZOS4, + borderMode=cv2.BORDER_REFLECT) + + # cv2.imshow('crop_face', crop_face) + # cv2.waitKey(0) + pt1k = face_alignmenter_1k.detect_single_face(crop_face) + pt1k = landmark_processor.transform_points(pt1k, image_to_face_mat, invert=True) + + # for pt in pt1k.astype(np.int32): + # cv2.circle(img, (pt[0], pt[1]), 1, (0, 255, 0), -1) + # cv2.imshow("img", img) + # cv2.waitKey(0) + + + landmarks87, poselayer, tracking_probe, occlusion_probe = mmcv.detect(img, [pt1k]) + # check valid + if tracking_probe[0] < 0.1: continue + + _, user_matting_8uc1 = generator_matte.matte_inference(img, pt1k) + + cv2.imwrite(os.path.splitext(img_file_path)[0] + '_matting.png', user_matting_8uc1) + with open(os.path.splitext(img_file_path)[0] + '.pkl', 'wb') as fp: + pickle.dump({'human_pt1k': pt1k, 'pt87': landmarks87[0]}, fp) + except Exception as e: + print(e) + + +def pkl_process(img_dir): + gpu_num = 1 + num_process = 1 + + import re, tqdm + + sq = Queue() + pattern = re.compile(r"^[^.].*\.((jpg)|(JPG)|(jpeg)|(JPEG)|(bmp)|(BMP)|(png)|(PNG))$") + for dirpath, dirnames, filenames in os.walk(img_dir): + for filename in filenames: + match = pattern.match(filename) + if not match: continue + if '_seg.png' in filename: continue + file_path = os.path.join(dirpath, filename) + # save_path = file_path.replace(img_dir, dst_dir) + sq.put(file_path) + + last_sq_size = sq.qsize() + if last_sq_size == 0: raise RuntimeError('sq is empty!') + pbar = tqdm.tqdm(total=last_sq_size) + + + process_thread(sq, 0) + + +if __name__ == "__main__": + img_dir = "" + pkl_process(img_dir) diff --git a/hair_service_sd/upload_oss.py b/hair_service_sd/upload_oss.py new file mode 100644 index 0000000..33292cf --- /dev/null +++ b/hair_service_sd/upload_oss.py @@ -0,0 +1,40 @@ +import time + +import oss2 +import os + +class OSS_object(): + def __init__(self): + access_key_id = os.getenv('OSS_TEST_ACCESS_KEY_ID', 'LTAI5tPZA6M67YRoxGPdJw1v') + access_key_secret = os.getenv('OSS_TEST_ACCESS_KEY_SECRET', 'BPydkvYFrXsbOYj3ix8UHQSLS4ivNP') + bucket_name = os.getenv('OSS_TEST_BUCKET', 'oss-aidigitalfield') + endpoint = os.getenv('OSS_TEST_ENDPOINT', 'oss-cn-beijing.aliyuncs.com') + # access_key_id = os.getenv('OSS_TEST_ACCESS_KEY_ID', 'LTAI5tMq9DivPYYkcpc6qhNP') + # access_key_secret = os.getenv('OSS_TEST_ACCESS_KEY_SECRET', 'XIkDAq7r4U9BVf7fkKECP46FXDLF9l') + # bucket_name = os.getenv('OSS_TEST_BUCKET', 'digit-person') + # endpoint = os.getenv('OSS_TEST_ENDPOINT', 'oss-cn-beijing.aliyuncs.com') + # access_key_id = os.getenv('OSS_TEST_ACCESS_KEY_ID', 'LTAI5tByNnrV4vRVioW66uq2') + # access_key_secret = os.getenv('OSS_TEST_ACCESS_KEY_SECRET', 'bLGYCKUNtQTDfzGfXqw06MWYomr5lw') + # bucket_name = os.getenv('OSS_TEST_BUCKET', 'mzyidong-tmp') + # endpoint = os.getenv('OSS_TEST_ENDPOINT', 'oss-cn-zhangjiakou-internal.aliyuncs.com') + for param in (access_key_id, access_key_secret, bucket_name, endpoint): + assert '<' not in param, '请设置参数：' + param + + self.bucket = oss2.Bucket(oss2.Auth(access_key_id, access_key_secret), endpoint, bucket_name) + + def upload_file(self, file, target_name): + t0 = time.time() + with open(oss2.to_unicode(file), 'rb') as f: + ret = self.bucket.put_object(target_name, f) + print(ret.headers['x-oss-request-id']) + + url = "https://oss-aidigitalfield.oss-cn-beijing.aliyuncs.com/{}".format(target_name) + print('耗时：{},签名url的地址为：{}'.format(time.time() - t0, url)) + return url + + +if __name__ == '__main__': + oss_2 = OSS_object() + t0 = time.time() + url = oss_2.upload_file('/home/szlc/Downloads/多中心录入模版.xlsx', 'tongji/多中心录入模版.xlsx') + print('url:', url) diff --git a/hair_service_sd/utils/MomocvFaceAlignment1K.py b/hair_service_sd/utils/MomocvFaceAlignment1K.py new file mode 100644 index 0000000..b0afc5c --- /dev/null +++ b/hair_service_sd/utils/MomocvFaceAlignment1K.py @@ -0,0 +1,462 @@ +import torch.nn as nn +import torch.utils.model_zoo as model_zoo +import torch +import numpy as np +import os +from utils import landmark_processor +from algorithm_conf import ConfFactory +from utils.umeyama import umeyama + +import cv2 + +__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_dir = ConfFactory.getModelValue("model_dir") + weights = torch.load(os.path.join(self.model_dir, '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('MomocvFaceAlignment1K success') + + 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, crop_M): + dst_size = 256 + with torch.no_grad(): + + # add for change crop for 1024 + h, w, _ = img.shape + if h == 768: + crop_img = img[104:img.shape[0] - 104, 104:img.shape[1] - 104, :] + tmp = cv2.resize(crop_img, (dst_size, dst_size)) + + else: + tmp = cv2.warpAffine(img, crop_M, (dst_size, dst_size), flags=cv2.INTER_CUBIC) + + # cv2.imshow("img_paf_test_crop: ", tmp) + # cv2.waitKey() + # crop_img = img[220:img.shape[0] - 220, 266:img.shape[1] - 266, :] # 220 266 + # tmp = cv2.resize(crop_img, (dst_size, dst_size)) + + # cv2.imshow("detect face: h:{:d}".format(h), tmp) + # cv2.waitKey() + + 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)) + + if h == 768: + orig_pts[:, 0] = orig_pts[:, 0] * crop_img.shape[1] + orig_pts[:, 1] = orig_pts[:, 1] * crop_img.shape[0] + orig_pts[:, 0] += 104 + orig_pts[:, 1] += 104 + else: + orig_pts[:, 0] = orig_pts[:, 0] * dst_size + orig_pts[:, 1] = orig_pts[:, 1] * dst_size + orig_pts = landmark_processor.transform_points(orig_pts, crop_M, invert=True) + # orig_pts[:, 0] = orig_pts[:, 0] * crop_img.shape[1] + # orig_pts[:, 1] = orig_pts[:, 1] * crop_img.shape[0] + # orig_pts[:, 0] += 266 + # orig_pts[:, 1] += 220 + + 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.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # 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]]) + + 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 + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + + tmp = cv2.warpAffine(img, mat, (dst_size, dst_size)) + # 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 diff --git a/hair_service_sd/utils/box_utils_Retina.py b/hair_service_sd/utils/box_utils_Retina.py new file mode 100644 index 0000000..c1d12bc --- /dev/null +++ b/hair_service_sd/utils/box_utils_Retina.py @@ -0,0 +1,330 @@ +import torch +import numpy as np + + +def point_form(boxes): + """ Convert prior_boxes to (xmin, ymin, xmax, ymax) + representation for comparison to point form ground truth data. + Args: + boxes: (tensor) center-size default boxes from priorbox layers. + Return: + boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes. + """ + return torch.cat((boxes[:, :2] - boxes[:, 2:]/2, # xmin, ymin + boxes[:, :2] + boxes[:, 2:]/2), 1) # xmax, ymax + + +def center_size(boxes): + """ Convert prior_boxes to (cx, cy, w, h) + representation for comparison to center-size form ground truth data. + Args: + boxes: (tensor) point_form boxes + Return: + boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes. + """ + return torch.cat((boxes[:, 2:] + boxes[:, :2])/2, # cx, cy + boxes[:, 2:] - boxes[:, :2], 1) # w, h + + +def intersect(box_a, box_b): + """ We resize both tensors to [A,B,2] without new malloc: + [A,2] -> [A,1,2] -> [A,B,2] + [B,2] -> [1,B,2] -> [A,B,2] + Then we compute the area of intersect between box_a and box_b. + Args: + box_a: (tensor) bounding boxes, Shape: [A,4]. + box_b: (tensor) bounding boxes, Shape: [B,4]. + Return: + (tensor) intersection area, Shape: [A,B]. + """ + A = box_a.size(0) + B = box_b.size(0) + max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2), + box_b[:, 2:].unsqueeze(0).expand(A, B, 2)) + min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2), + box_b[:, :2].unsqueeze(0).expand(A, B, 2)) + inter = torch.clamp((max_xy - min_xy), min=0) + return inter[:, :, 0] * inter[:, :, 1] + + +def jaccard(box_a, box_b): + """Compute the jaccard overlap of two sets of boxes. The jaccard overlap + is simply the intersection over union of two boxes. Here we operate on + ground truth boxes and default boxes. + E.g.: + A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B) + Args: + box_a: (tensor) Ground truth bounding boxes, Shape: [num_objects,4] + box_b: (tensor) Prior boxes from priorbox layers, Shape: [num_priors,4] + Return: + jaccard overlap: (tensor) Shape: [box_a.size(0), box_b.size(0)] + """ + inter = intersect(box_a, box_b) + area_a = ((box_a[:, 2]-box_a[:, 0]) * + (box_a[:, 3]-box_a[:, 1])).unsqueeze(1).expand_as(inter) # [A,B] + area_b = ((box_b[:, 2]-box_b[:, 0]) * + (box_b[:, 3]-box_b[:, 1])).unsqueeze(0).expand_as(inter) # [A,B] + union = area_a + area_b - inter + return inter / union # [A,B] + + +def matrix_iou(a, b): + """ + return iou of a and b, numpy version for data augenmentation + """ + lt = np.maximum(a[:, np.newaxis, :2], b[:, :2]) + rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:]) + + area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2) + area_a = np.prod(a[:, 2:] - a[:, :2], axis=1) + area_b = np.prod(b[:, 2:] - b[:, :2], axis=1) + return area_i / (area_a[:, np.newaxis] + area_b - area_i) + + +def matrix_iof(a, b): + """ + return iof of a and b, numpy version for data augenmentation + """ + lt = np.maximum(a[:, np.newaxis, :2], b[:, :2]) + rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:]) + + area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2) + area_a = np.prod(a[:, 2:] - a[:, :2], axis=1) + return area_i / np.maximum(area_a[:, np.newaxis], 1) + + +def match(threshold, truths, priors, variances, labels, landms, loc_t, conf_t, landm_t, idx): + """Match each prior box with the ground truth box of the highest jaccard + overlap, encode the bounding boxes, then return the matched indices + corresponding to both confidence and location preds. + Args: + threshold: (float) The overlap threshold used when mathing boxes. + truths: (tensor) Ground truth boxes, Shape: [num_obj, 4]. + priors: (tensor) Prior boxes from priorbox layers, Shape: [n_priors,4]. + variances: (tensor) Variances corresponding to each prior coord, + Shape: [num_priors, 4]. + labels: (tensor) All the class labels for the image, Shape: [num_obj]. + landms: (tensor) Ground truth landms, Shape [num_obj, 10]. + loc_t: (tensor) Tensor to be filled w/ endcoded location targets. + conf_t: (tensor) Tensor to be filled w/ matched indices for conf preds. + landm_t: (tensor) Tensor to be filled w/ endcoded landm targets. + idx: (int) current batch index + Return: + The matched indices corresponding to 1)location 2)confidence 3)landm preds. + """ + # jaccard index + overlaps = jaccard( + truths, + point_form(priors) + ) + # (Bipartite Matching) + # [1,num_objects] best prior for each ground truth + best_prior_overlap, best_prior_idx = overlaps.max(1, keepdim=True) + + # ignore hard gt + valid_gt_idx = best_prior_overlap[:, 0] >= 0.2 + best_prior_idx_filter = best_prior_idx[valid_gt_idx, :] + if best_prior_idx_filter.shape[0] <= 0: + loc_t[idx] = 0 + conf_t[idx] = 0 + return + + # [1,num_priors] best ground truth for each prior + best_truth_overlap, best_truth_idx = overlaps.max(0, keepdim=True) + best_truth_idx.squeeze_(0) + best_truth_overlap.squeeze_(0) + best_prior_idx.squeeze_(1) + best_prior_idx_filter.squeeze_(1) + best_prior_overlap.squeeze_(1) + best_truth_overlap.index_fill_(0, best_prior_idx_filter, 2) # ensure best prior + # TODO refactor: index best_prior_idx with long tensor + # ensure every gt matches with its prior of max overlap + for j in range(best_prior_idx.size(0)): # 判别此anchor是预测哪一个boxes + best_truth_idx[best_prior_idx[j]] = j + matches = truths[best_truth_idx] # Shape: [num_priors,4] 此处为每一个anchor对应的bbox取出来 + conf = labels[best_truth_idx] # Shape: [num_priors] 此处为每一个anchor对应的label取出来 + conf[best_truth_overlap < threshold] = 0 # label as background overlap<0.35的全部作为负样本 + loc = encode(matches, priors, variances) + + matches_landm = landms[best_truth_idx] + landm = encode_landm(matches_landm, priors, variances) + loc_t[idx] = loc # [num_priors,4] encoded offsets to learn + conf_t[idx] = conf # [num_priors] top class label for each prior + landm_t[idx] = landm + + +def encode(matched, priors, variances): + """Encode the variances from the priorbox layers into the ground truth boxes + we have matched (based on jaccard overlap) with the prior boxes. + Args: + matched: (tensor) Coords of ground truth for each prior in point-form + Shape: [num_priors, 4]. + priors: (tensor) Prior boxes in center-offset form + Shape: [num_priors,4]. + variances: (list[float]) Variances of priorboxes + Return: + encoded boxes (tensor), Shape: [num_priors, 4] + """ + + # dist b/t match center and prior's center + g_cxcy = (matched[:, :2] + matched[:, 2:])/2 - priors[:, :2] + # encode variance + g_cxcy /= (variances[0] * priors[:, 2:]) + # match wh / prior wh + g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:] + g_wh = torch.log(g_wh) / variances[1] + # return target for smooth_l1_loss + return torch.cat([g_cxcy, g_wh], 1) # [num_priors,4] + +def encode_landm(matched, priors, variances): + """Encode the variances from the priorbox layers into the ground truth boxes + we have matched (based on jaccard overlap) with the prior boxes. + Args: + matched: (tensor) Coords of ground truth for each prior in point-form + Shape: [num_priors, 10]. + priors: (tensor) Prior boxes in center-offset form + Shape: [num_priors,4]. + variances: (list[float]) Variances of priorboxes + Return: + encoded landm (tensor), Shape: [num_priors, 10] + """ + + # dist b/t match center and prior's center + matched = torch.reshape(matched, (matched.size(0), 5, 2)) + priors_cx = priors[:, 0].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2) + priors_cy = priors[:, 1].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2) + priors_w = priors[:, 2].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2) + priors_h = priors[:, 3].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2) + priors = torch.cat([priors_cx, priors_cy, priors_w, priors_h], dim=2) + g_cxcy = matched[:, :, :2] - priors[:, :, :2] + # encode variance + g_cxcy /= (variances[0] * priors[:, :, 2:]) + # g_cxcy /= priors[:, :, 2:] + g_cxcy = g_cxcy.reshape(g_cxcy.size(0), -1) + # return target for smooth_l1_loss + return g_cxcy + + +# Adapted from https://github.com/Hakuyume/chainer-ssd +def decode(loc, priors, variances): + """Decode locations from predictions using priors to undo + the encoding we did for offset regression at train time. + Args: + loc (tensor): location predictions for loc layers, + Shape: [num_priors,4] + priors (tensor): Prior boxes in center-offset form. + Shape: [num_priors,4]. + variances: (list[float]) Variances of priorboxes + Return: + decoded bounding box predictions + """ + + boxes = torch.cat(( + priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:], + priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1) + boxes[:, :2] -= boxes[:, 2:] / 2 + boxes[:, 2:] += boxes[:, :2] + return boxes + +def decode_landm(pre, priors, variances): + """Decode landm from predictions using priors to undo + the encoding we did for offset regression at train time. + Args: + pre (tensor): landm predictions for loc layers, + Shape: [num_priors,10] + priors (tensor): Prior boxes in center-offset form. + Shape: [num_priors,4]. + variances: (list[float]) Variances of priorboxes + Return: + decoded landm predictions + """ + landms = torch.cat((priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:], + priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:], + priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:], + priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:], + priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:], + ), dim=1) + return landms + + +def log_sum_exp(x): + """Utility function for computing log_sum_exp while determining + This will be used to determine unaveraged confidence loss across + all examples in a batch. + Args: + x (Variable(tensor)): conf_preds from conf layers + """ + x_max = x.data.max() + return torch.log(torch.sum(torch.exp(x-x_max), 1, keepdim=True)) + x_max + + +# Original author: Francisco Massa: +# https://github.com/fmassa/object-detection.torch +# Ported to PyTorch by Max deGroot (02/01/2017) +def nms(boxes, scores, overlap=0.5, top_k=200): + """Apply non-maximum suppression at test time to avoid detecting too many + overlapping bounding boxes for a given object. + Args: + boxes: (tensor) The location preds for the img, Shape: [num_priors,4]. + scores: (tensor) The class predscores for the img, Shape:[num_priors]. + overlap: (float) The overlap thresh for suppressing unnecessary boxes. + top_k: (int) The Maximum number of box preds to consider. + Return: + The indices of the kept boxes with respect to num_priors. + """ + + keep = torch.Tensor(scores.size(0)).fill_(0).long() + if boxes.numel() == 0: + return keep + x1 = boxes[:, 0] + y1 = boxes[:, 1] + x2 = boxes[:, 2] + y2 = boxes[:, 3] + area = torch.mul(x2 - x1, y2 - y1) + v, idx = scores.sort(0) # sort in ascending order + # I = I[v >= 0.01] + idx = idx[-top_k:] # indices of the top-k largest vals + xx1 = boxes.new() + yy1 = boxes.new() + xx2 = boxes.new() + yy2 = boxes.new() + w = boxes.new() + h = boxes.new() + + # keep = torch.Tensor() + count = 0 + while idx.numel() > 0: + i = idx[-1] # index of current largest val + # keep.append(i) + keep[count] = i + count += 1 + if idx.size(0) == 1: + break + idx = idx[:-1] # remove kept element from view + # load bboxes of next highest vals + torch.index_select(x1, 0, idx, out=xx1) + torch.index_select(y1, 0, idx, out=yy1) + torch.index_select(x2, 0, idx, out=xx2) + torch.index_select(y2, 0, idx, out=yy2) + # store element-wise max with next highest score + xx1 = torch.clamp(xx1, min=x1[i]) + yy1 = torch.clamp(yy1, min=y1[i]) + xx2 = torch.clamp(xx2, max=x2[i]) + yy2 = torch.clamp(yy2, max=y2[i]) + w.resize_as_(xx2) + h.resize_as_(yy2) + w = xx2 - xx1 + h = yy2 - yy1 + # check sizes of xx1 and xx2.. after each iteration + w = torch.clamp(w, min=0.0) + h = torch.clamp(h, min=0.0) + inter = w*h + # IoU = i / (area(a) + area(b) - i) + rem_areas = torch.index_select(area, 0, idx) # load remaining areas) + union = (rem_areas - inter) + area[i] + IoU = inter/union # store result in iou + # keep only elements with an IoU <= overlap + idx = idx[IoU.le(overlap)] + return keep, count + + diff --git a/hair_service_sd/utils/call_hair_inter.py b/hair_service_sd/utils/call_hair_inter.py new file mode 100644 index 0000000..2d8c9b0 --- /dev/null +++ b/hair_service_sd/utils/call_hair_inter.py @@ -0,0 +1,125 @@ +import requests +import json +from common.logger import config + +version = config.get('default', 'version') +if version == "local": + current_photo_service_url = 'http://192.168.1.57:32678/' +else: + current_photo_service_url = 'http://0.0.0.0:32678/' + +def call_hair_infer(task_id, hair_id, hair_material_dir, infer_req, is_hr, inference_port): + url = f"{current_photo_service_url}api/hair/inference" + payload = json.dumps({ + "task_id": task_id, + "hair_id": hair_id, + "hd_version_flag":is_hr, + "hair_material_dir": hair_material_dir, + "request_json": infer_req, + "inference_port": inference_port + }) + headers = { + 'Content-Type': 'application/json' + } + # print("---call infer payload:", payload) + + response = requests.request("POST", url, headers=headers, data=payload) + + # print(response.json()) + + return response.json() + +def call_hair_infer_diy(task_id, infer_req, inference_port): + url = f"{current_photo_service_url}api/hair/inference_diy" + payload = json.dumps({ + "task_id": task_id, + "request_json": infer_req, + "inference_port": inference_port + }) + headers = { + 'Content-Type': 'application/json' + } + # print("---call infer payload:", payload) + + response = requests.request("POST", url, headers=headers, data=payload) + + # print(response.json()) + + return response.json() + + + +def call_hair_enhance(img_path, mask_path, task_id, in_gender): + url = "http://127.0.0.1:7393/hairEnhance/v1" + payload = json.dumps({ + "img_path": img_path, + "mask_path": mask_path, + "req_id": task_id, + "gender": in_gender + }) + headers = { + 'Content-Type': 'application/json' + } + # print("---call infer payload:", payload) + + response = requests.request("POST", url, headers=headers, data=payload) + + # print(response.json()) + + return response.json() + +def face_info(img_path): + url = "http://127.0.0.1:7393/faceInfo/v1" + payload = json.dumps({ + "img": img_path, + "userId": 'fff', + "isLocal": True, + }) + headers = { + 'Content-Type': 'application/json' + } + # print("---call infer payload:", payload) + + response = requests.request("POST", url, headers=headers, data=payload) + + print(response.json()) + + return response.json() + +def is_same(img_url1, img_url2): + url = "http://127.0.0.1:7393/faceInfo/same" + payload = json.dumps({ + "img_url1": img_url1, + "img_url2": img_url2, + # "isLocal": True, + }) + headers = { + 'Content-Type': 'application/json' + } + # print("---call infer payload:", payload) + + response = requests.request("POST", url, headers=headers, data=payload) + + print(response.json()) + + return response.json() + +def hair_select(img_url): + url = "http://127.0.0.1:7393/hairStyle/hair_select" + payload = json.dumps({ + "img_url": img_url, + # "isLocal": True, + }) + headers = { + 'Content-Type': 'application/json' + } + # print("---call infer payload:", payload) + + response = requests.request("POST", url, headers=headers, data=payload) + + print(response.json()) + + return response.json() +# test_url = 'https://ydapp-1317132355.cos.ap-beijing.myqcloud.com/hair_mz/images/hairstyle/fb7d7a58-3228-40f2-84a6-337088ee31e2/2023031623425988.jpg' +# test_url = 'https://ydapp-1317132355.cos.ap-beijing.myqcloud.com/hair_mz/images/vaffflue12.png' +# hair_select(test_url) \ No newline at end of file diff --git a/hair_service_sd/utils/call_hair_train.py b/hair_service_sd/utils/call_hair_train.py new file mode 100644 index 0000000..68d2a19 --- /dev/null +++ b/hair_service_sd/utils/call_hair_train.py @@ -0,0 +1,30 @@ +import requests +import json +from common.logger import config + +version = config.get('default', 'version') +if version == "local": + current_photo_service_url = 'http://192.168.1.57:32678/' +else: + current_photo_service_url = 'http://0.0.0.0:32678/' + +def call_hair_train(task_id, hair_id, hair_material_dir, tag, is_tj="0", webui_addr=None, device_id=None): + url = f"{webui_addr}api/hair/train" + payload = json.dumps({ + "task_id": task_id, + "hair_id": hair_id, + "hair_material_dir": hair_material_dir, + "tag": tag, + "is_tj": is_tj, + "webui_addr": webui_addr, + "device_id":device_id + }) + headers = { + 'Content-Type': 'application/json' + } + print("---call train payload:", payload) + + response = requests.request("POST", url, headers=headers, data=payload) + + print(response.json()) + diff --git a/hair_service_sd/utils/callback.py b/hair_service_sd/utils/callback.py new file mode 100644 index 0000000..70dac81 --- /dev/null +++ b/hair_service_sd/utils/callback.py @@ -0,0 +1,20 @@ +import requests +import json + +def recall(req_id, state, message, clothId): + url = "http://192.168.11.220:9281/api/cloth/callBack" + payload = json.dumps({ + "taskId": req_id, + "status": state, + "clothId": clothId, + "msg": message + }) + headers = { + 'Content-Type': 'application/json' + } + print("---payload:", payload) + + response = requests.request("POST", url, headers=headers, data=payload) + + print(response.json()) + diff --git a/hair_service_sd/utils/data_preprocess.py b/hair_service_sd/utils/data_preprocess.py new file mode 100644 index 0000000..fd69e82 --- /dev/null +++ b/hair_service_sd/utils/data_preprocess.py @@ -0,0 +1,27 @@ +import os + +IMG_EXTENSIONS = [ + '.jpg', '.JPG', '.jpeg', '.JPEG', + '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', +] + +def is_image_file(filename): + return any(filename.endswith(extension) for extension in IMG_EXTENSIONS) + +def mkdirs(paths): + """create empty directories if they don't exist + Parameters: + paths (str list) -- a list of directory paths + """ + if isinstance(paths, list) and not isinstance(paths, str): + for path in paths: + make_dir(path) + else: + make_dir(paths) + +def make_dir(target_dir): + """ + Create dir if not exists + """ + if not os.path.exists(target_dir): + os.makedirs(target_dir) diff --git a/hair_service_sd/utils/enhance_hair.py b/hair_service_sd/utils/enhance_hair.py new file mode 100644 index 0000000..53798e2 --- /dev/null +++ b/hair_service_sd/utils/enhance_hair.py @@ -0,0 +1,80 @@ +import cv2 +import numpy as np + +import base64 +import requests +from common.logger import config + +version = config.get('default', 'version') +if version == "local": + webui_url = 'http://192.168.1.57:57860/' +else: + webui_url = 'http://0.0.0.0:57860/' + +def encode_numpy_to_base64(img): + retval, bytes = cv2.imencode('.png', img) + encoded_image = base64.b64encode(bytes).decode('utf-8') + return encoded_image + + +def webui_img2img(img, mask, prompt=''): + url = f"{webui_url}sdapi/v1/img2img" + request_dict = { + "prompt": prompt, + "negative_prompt": '(nsfw:1.5), ng_deepnegative_v1_75t, (badhandv4:1.2), (worst quality:2), (low quality:2), (normal quality:2), lowres, bad anatomy, bad hands, ((monochrome)), ((grayscale)) watermark, bad_pictures,easynegative', + "sampler_name": "DPM++ 2M Karras", + "batch_size": 1, + "steps": 30, + "width": img.shape[1], + "height": img.shape[0], + "cfg_scale": 7.0, + "seed": 123456789, + "mask_blur": 5, + "init_images": [ + encode_numpy_to_base64(img) + ], + "inpaint_full_res": False, + "inpainting_fill": 1, + "inpainting_mask_invert": 0, + "mask": encode_numpy_to_base64(mask), + # "refiner_checkpoint":"majicmixRealistic_v7.safetensors", + # "refiner_switch_at": 0.4, + "denoising_strength": 0.35, + "alwayson_scripts": { + } + } + response = requests.post(url=url, json=request_dict) + ret_json = response.json() + result = ret_json['images'][0] + img = cv2.imdecode(np.frombuffer(base64.b64decode(result.split(",", 1)[0]), np.uint8), cv2.IMREAD_COLOR) + return img + +def webui_super_res_img(img, ratio): + url = f"{webui_url}sdapi/v1/extra-single-image" + request_dict = { + "resize_mode": 0, + "show_extras_results": False, + "gfpgan_visibility": 0, + "codeformer_visibility": 1, + "codeformer_weight": 1, + "upscaling_resize": ratio, + "upscaler_1": "8x_NMKD-Superscale_150000_G", + "upscale_first": False, + "image": encode_numpy_to_base64(img) + } + response = requests.post(url=url, json=request_dict) + ret_json = response.json() + result = ret_json['image'] + img = cv2.imdecode(np.frombuffer(base64.b64decode(result), np.uint8), cv2.IMREAD_COLOR) + return img + + +def webui_tag_by_clip(img): + url = f"{webui_url}sdapi/v1/interrogate" + request_dict = { + "image": encode_numpy_to_base64(img), + "model": "clip" + } + response = requests.post(url=url, json=request_dict) + ret_json = response.json() + return ret_json['caption'] diff --git a/hair_service_sd/utils/landmark_processor.py b/hair_service_sd/utils/landmark_processor.py new file mode 100644 index 0000000..fe516f5 --- /dev/null +++ b/hair_service_sd/utils/landmark_processor.py @@ -0,0 +1,1670 @@ +import colorsys +import cv2 +import numpy as np +import random +import math +import time + +mean_face_x = np.array([ + 0.000213256, 0.0752622, 0.18113, 0.29077, 0.393397, 0.586856, 0.689483, 0.799124, + 0.904991, 0.98004, 0.490127, 0.490127, 0.490127, 0.490127, 0.36688, 0.426036, + 0.490127, 0.554217, 0.613373, 0.121737, 0.187122, 0.265825, 0.334606, 0.260918, + 0.182743, 0.645647, 0.714428, 0.793132, 0.858516, 0.79751, 0.719335, 0.254149, + 0.340985, 0.428858, 0.490127, 0.551395, 0.639268, 0.726104, 0.642159, 0.556721, + 0.490127, 0.423532, 0.338094, 0.290379, 0.428096, 0.490127, 0.552157, 0.689874, + 0.553364, 0.490127, 0.42689]) + +mean_face_y = np.array([ + 0.106454, 0.038915, 0.0187482, 0.0344891, 0.0773906, 0.0773906, 0.0344891, + 0.0187482, 0.038915, 0.106454, 0.203352, 0.307009, 0.409805, 0.515625, 0.587326, + 0.609345, 0.628106, 0.609345, 0.587326, 0.216423, 0.178758, 0.179852, 0.231733, + 0.245099, 0.244077, 0.231733, 0.179852, 0.178758, 0.216423, 0.244077, 0.245099, + 0.780233, 0.745405, 0.727388, 0.742578, 0.727388, 0.745405, 0.780233, 0.864805, + 0.902192, 0.909281, 0.902192, 0.864805, 0.784792, 0.778746, 0.785343, 0.778746, + 0.784792, 0.824182, 0.831803, 0.824182]) + +landmarks_2D = np.stack([mean_face_x, mean_face_y], axis=1) + +mean_face_x_1k = np.array([0.498047, 0.504671, 0.511286, 0.517984, 0.524451, 0.531086, 0.537574, 0.543902, 0.550305, 0.556545, 0.562891, 0.568903, 0.574994, 0.581092, 0.587026, 0.592814, 0.598603, 0.604259, 0.609802, 0.615311, 0.620821, 0.626135, 0.631542, 0.636597, 0.641621, 0.646680, 0.651656, 0.656505, 0.661374, 0.666025, 0.670571, 0.675225, 0.679616, 0.683987, 0.688181, 0.692460, 0.696634, 0.700552, 0.704500, 0.708402, 0.712261, 0.715842, 0.719373, 0.722930, 0.726486, 0.729785, 0.732875, 0.736135, 0.739115, 0.742179, 0.744925, 0.747688, 0.750314, 0.752887, 0.755412, 0.757737, 0.760038, 0.762291, 0.764312, 0.766333, 0.768259, 0.770363, 0.771888, 0.773763, 0.775485, 0.777069, 0.778656, 0.780126, 0.781483, 0.782879, 0.784085, 0.785261, 0.786402, 0.787338, 0.788398, 0.789310, 0.790038, 0.790796, 0.791281, 0.792001, 0.792397, 0.792903, 0.793221, 0.793554, 0.793613, 0.793848, 0.793856, 0.793925, 0.793911, 0.793825, 0.793601, 0.793426, 0.793199, 0.792833, 0.792528, 0.791962, 0.791501, 0.790897, 0.790382, 0.789660, 0.788884, 0.787702, 0.786300, 0.784951, 0.783363, 0.781861, 0.780103, 0.778348, 0.776405, 0.774489, 0.772390, 0.770098, 0.767540, 0.765030, 0.762359, 0.759349, 0.756513, 0.753197, 0.749784, 0.746353, 0.742572, 0.738633, 0.734433, 0.730200, 0.725658, 0.720657, 0.715711, 0.710406, 0.705108, 0.699482, 0.693569, 0.687476, 0.681209, 0.674681, 0.668094, 0.661239, 0.654234, 0.647131, 0.639989, 0.632582, 0.625134, 0.617439, 0.610009, 0.602125, 0.594318, 0.586546, 0.578466, 0.570582, 0.562638, 0.554622, 0.546602, 0.538546, 0.530536, 0.522423, 0.514258, 0.506205, 0.498047, 0.489889, 0.481836, 0.473671, 0.465558, 0.457548, 0.449491, 0.441472, 0.433455, 0.425511, 0.417628, 0.409547, 0.401776, 0.393969, 0.386085, 0.378655, 0.370960, 0.363512, 0.356105, 0.348963, 0.341860, 0.334855, 0.328000, 0.321412, 0.314885, 0.308618, 0.302525, 0.296612, 0.290986, 0.285687, 0.280382, 0.275437, 0.270436, 0.265894, 0.261660, 0.257461, 0.253522, 0.249740, 0.246310, 0.242897, 0.239580, 0.236744, 0.233735, 0.231064, 0.228553, 0.225996, 0.223704, 0.221605, 0.219689, 0.217746, 0.215990, 0.214232, 0.212730, 0.211143, 0.209794, 0.208392, 0.207210, 0.206433, 0.205712, 0.205197, 0.204593, 0.204132, 0.203565, 0.203261, 0.202895, 0.202667, 0.202492, 0.202268, 0.202183, 0.202169, 0.202238, 0.202246, 0.202480, 0.202540, 0.202873, 0.203191, 0.203696, 0.204093, 0.204813, 0.205297, 0.206055, 0.206783, 0.207696, 0.208755, 0.209692, 0.210833, 0.212009, 0.213215, 0.214611, 0.215967, 0.217438, 0.219024, 0.220609, 0.222330, 0.224206, 0.225731, 0.227835, 0.229761, 0.231781, 0.233803, 0.236055, 0.238356, 0.240681, 0.243207, 0.245780, 0.248405, 0.251168, 0.253915, 0.256979, 0.259958, 0.263219, 0.266309, 0.269608, 0.273163, 0.276721, 0.280252, 0.283833, 0.287692, 0.291593, 0.295541, 0.299459, 0.303633, 0.307913, 0.312107, 0.316478, 0.320869, 0.325523, 0.330069, 0.334719, 0.339589, 0.344438, 0.349414, 0.354473, 0.359497, 0.364552, 0.369959, 0.375273, 0.380783, 0.386292, 0.391835, 0.397491, 0.403280, 0.409068, 0.415002, 0.421100, 0.427190, 0.433203, 0.439548, 0.445789, 0.452191, 0.458520, 0.465008, 0.471643, 0.478110, 0.484807, 0.491422, 0.396255, 0.397888, 0.399588, 0.401344, 0.403150, 0.405001, 0.406895, 0.408831, 0.410807, 0.412822, 0.414875, 0.416966, 0.419095, 0.421262, 0.423465, 0.425705, 0.427981, 0.430294, 0.432644, 0.435030, 0.437451, 0.439908, 0.442400, 0.444928, 0.447489, 0.450084, 0.452712, 0.455372, 0.458062, 0.460782, 0.463531, 0.466307, 0.469108, 0.471933, 0.474781, 0.477648, 0.480532, 0.483432, 0.486345, 0.489267, 0.492194, 0.495123, 0.498047, 0.500971, 0.503899, 0.506827, 0.509749, 0.512662, 0.515562, 0.518446, 0.521313, 0.524160, 0.526986, 0.529787, 0.532563, 0.535312, 0.538032, 0.540722, 0.543382, 0.546009, 0.548605, 0.551166, 0.553693, 0.556186, 0.558642, 0.561064, 0.563450, 0.565799, 0.568113, 0.570389, 0.572629, 0.574832, 0.576998, 0.579127, 0.581219, 0.583272, 0.585287, 0.587263, 0.589198, 0.591093, 0.592944, 0.594749, 0.596506, 0.598206, 0.599839, 0.597799, 0.595660, 0.593444, 0.591167, 0.588838, 0.586464, 0.584049, 0.581599, 0.579114, 0.576598, 0.574052, 0.571476, 0.568873, 0.566242, 0.563585, 0.560901, 0.558193, 0.555461, 0.552704, 0.549922, 0.547116, 0.544286, 0.541432, 0.538555, 0.535655, 0.532732, 0.529786, 0.526819, 0.523832, 0.520830, 0.517033, 0.513236, 0.509439, 0.505641, 0.501844, 0.498047, 0.494250, 0.490452, 0.486655, 0.482858, 0.479061, 0.475264, 0.472262, 0.469275, 0.466308, 0.463362, 0.460439, 0.457538, 0.454662, 0.451808, 0.448978, 0.446172, 0.443390, 0.440633, 0.437900, 0.435192, 0.432509, 0.429852, 0.427221, 0.424617, 0.422042, 0.419496, 0.416979, 0.414495, 0.412044, 0.409630, 0.407256, 0.404927, 0.402650, 0.400434, 0.398294, 0.410228, 0.414399, 0.418630, 0.422899, 0.427196, 0.431515, 0.435855, 0.440212, 0.444587, 0.448978, 0.453385, 0.457806, 0.462242, 0.466691, 0.471152, 0.475622, 0.480101, 0.484586, 0.489074, 0.493564, 0.498047, 0.502530, 0.507020, 0.511508, 0.515993, 0.520472, 0.524942, 0.529403, 0.533852, 0.538288, 0.542709, 0.547116, 0.551507, 0.555882, 0.560238, 0.564578, 0.568898, 0.573194, 0.577464, 0.581695, 0.585866, 0.581642, 0.577369, 0.573066, 0.568743, 0.564401, 0.560044, 0.555673, 0.551287, 0.546890, 0.542478, 0.538055, 0.533622, 0.529181, 0.524733, 0.520280, 0.515825, 0.511370, 0.506919, 0.502475, 0.498047, 0.493618, 0.489175, 0.484723, 0.480269, 0.475814, 0.471361, 0.466913, 0.462472, 0.458039, 0.453616, 0.449204, 0.444806, 0.440421, 0.436050, 0.431693, 0.427351, 0.423027, 0.418725, 0.414452, 0.459656, 0.459227, 0.458691, 0.458058, 0.457338, 0.456539, 0.455664, 0.454719, 0.453708, 0.452633, 0.451497, 0.450302, 0.449049, 0.447741, 0.446378, 0.444962, 0.443494, 0.441973, 0.440402, 0.438780, 0.437107, 0.435384, 0.433612, 0.431790, 0.429919, 0.427998, 0.426028, 0.424009, 0.421938, 0.419814, 0.417632, 0.415384, 0.442766, 0.471804, 0.498047, 0.524290, 0.553327, 0.580709, 0.578462, 0.576279, 0.574156, 0.572085, 0.570065, 0.568095, 0.566175, 0.564304, 0.562481, 0.560709, 0.558987, 0.557314, 0.555692, 0.554120, 0.552600, 0.551131, 0.549715, 0.548353, 0.547044, 0.545792, 0.544597, 0.543461, 0.542386, 0.541375, 0.540430, 0.539555, 0.538755, 0.538036, 0.537403, 0.536867, 0.536438, 0.551823, 0.524746, 0.471347, 0.444271, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.370401, 0.394341, 0.393977, 0.392891, 0.391151, 0.388793, 0.385710, 0.382329, 0.378604, 0.374564, 0.370390, 0.366216, 0.362177, 0.358453, 0.355074, 0.351995, 0.349640, 0.347903, 0.346822, 0.346461, 0.346826, 0.347911, 0.349651, 0.352010, 0.355092, 0.358473, 0.362199, 0.366239, 0.370413, 0.374587, 0.378626, 0.382350, 0.385728, 0.388808, 0.391163, 0.392899, 0.393981, 0.423391, 0.421623, 0.419536, 0.417200, 0.414654, 0.411925, 0.409036, 0.406003, 0.402844, 0.399571, 0.396199, 0.392740, 0.389206, 0.385609, 0.381960, 0.378271, 0.374552, 0.370812, 0.367063, 0.363314, 0.359575, 0.355859, 0.352173, 0.348530, 0.344945, 0.341429, 0.337998, 0.334669, 0.331464, 0.328407, 0.325532, 0.322889, 0.320595, 0.322717, 0.325103, 0.327689, 0.330440, 0.333327, 0.336328, 0.339426, 0.342603, 0.345848, 0.349146, 0.352488, 0.355863, 0.359264, 0.362683, 0.366116, 0.369556, 0.372999, 0.376442, 0.379882, 0.383316, 0.386742, 0.390158, 0.393563, 0.396955, 0.400332, 0.403694, 0.407038, 0.410363, 0.413666, 0.416943, 0.420190, 0.625692, 0.649633, 0.649272, 0.648190, 0.646454, 0.644099, 0.641019, 0.637641, 0.633917, 0.629878, 0.625704, 0.621530, 0.617490, 0.613764, 0.610383, 0.607301, 0.604942, 0.603202, 0.602117, 0.601752, 0.602113, 0.603194, 0.604931, 0.607286, 0.610365, 0.613744, 0.617468, 0.621507, 0.625681, 0.629855, 0.633895, 0.637621, 0.641001, 0.644084, 0.646442, 0.648182, 0.649268, 0.572703, 0.574471, 0.576557, 0.578893, 0.581440, 0.584168, 0.587058, 0.590090, 0.593250, 0.596523, 0.599895, 0.603354, 0.606888, 0.610485, 0.614134, 0.617823, 0.621542, 0.625282, 0.629031, 0.632780, 0.636518, 0.640235, 0.643921, 0.647563, 0.651149, 0.654665, 0.658096, 0.661425, 0.664630, 0.667686, 0.670562, 0.673205, 0.675499, 0.673377, 0.670991, 0.668404, 0.665654, 0.662767, 0.659765, 0.656668, 0.653491, 0.650246, 0.646948, 0.643606, 0.640231, 0.636830, 0.633411, 0.629978, 0.626538, 0.623095, 0.619652, 0.616212, 0.612778, 0.609352, 0.605936, 0.602531, 0.599139, 0.595762, 0.592400, 0.589056, 0.585731, 0.582428, 0.579150, 0.575904, 0.552267, 0.555073, 0.558014, 0.561313, 0.565180, 0.569636, 0.574564, 0.579836, 0.585356, 0.591050, 0.596857, 0.602735, 0.608652, 0.614591, 0.620542, 0.626500, 0.632462, 0.638425, 0.644390, 0.649985, 0.655581, 0.661175, 0.666765, 0.672343, 0.677894, 0.683391, 0.688788, 0.694021, 0.699020, 0.703737, 0.708170, 0.712362, 0.716374, 0.720268, 0.724091, 0.727877, 0.731647, 0.726965, 0.722281, 0.717595, 0.712902, 0.708202, 0.703489, 0.698762, 0.694017, 0.689253, 0.684472, 0.679678, 0.674876, 0.670068, 0.665257, 0.660446, 0.655633, 0.650821, 0.646008, 0.640808, 0.635608, 0.630406, 0.625203, 0.619997, 0.614788, 0.609577, 0.604360, 0.599139, 0.593915, 0.588693, 0.583472, 0.578256, 0.573043, 0.567837, 0.562637, 0.557446, 0.264446, 0.268217, 0.272003, 0.275826, 0.279720, 0.283732, 0.287924, 0.292357, 0.297074, 0.302073, 0.307305, 0.312703, 0.318199, 0.323751, 0.329329, 0.334919, 0.340513, 0.346109, 0.351704, 0.357669, 0.363632, 0.369594, 0.375552, 0.381503, 0.387442, 0.393359, 0.399237, 0.405044, 0.410737, 0.416258, 0.421530, 0.426457, 0.430914, 0.434781, 0.438079, 0.441020, 0.443827, 0.438648, 0.433457, 0.428257, 0.423050, 0.417838, 0.412621, 0.407401, 0.402179, 0.396955, 0.391734, 0.386517, 0.381305, 0.376096, 0.370891, 0.365688, 0.360486, 0.355285, 0.350085, 0.345273, 0.340460, 0.335648, 0.330836, 0.326026, 0.321218, 0.316416, 0.311622, 0.306841, 0.302077, 0.297332, 0.292605, 0.287892, 0.283191, 0.278499, 0.273812, 0.269128 +]) +mean_face_y_1k = np.array([0.851392, 0.851204, 0.850802, 0.849951, 0.849050, 0.847896, 0.846372, 0.844708, 0.842887, 0.840805, 0.838470, 0.836093, 0.833361, 0.830589, 0.827781, 0.824650, 0.821361, 0.817917, 0.814463, 0.810803, 0.807127, 0.803118, 0.799272, 0.795088, 0.790897, 0.786585, 0.782195, 0.777726, 0.773128, 0.768407, 0.763595, 0.758801, 0.753939, 0.748915, 0.743745, 0.738664, 0.733510, 0.728055, 0.722725, 0.717442, 0.711952, 0.706343, 0.700725, 0.695036, 0.689315, 0.683557, 0.677857, 0.671952, 0.666097, 0.659596, 0.652993, 0.646386, 0.639687, 0.632843, 0.626218, 0.619403, 0.612554, 0.605551, 0.598705, 0.591735, 0.584939, 0.577854, 0.571094, 0.564026, 0.557009, 0.549978, 0.542887, 0.535838, 0.528892, 0.521764, 0.514596, 0.507444, 0.500318, 0.493164, 0.486134, 0.478851, 0.471725, 0.464586, 0.457279, 0.450071, 0.443049, 0.435691, 0.428445, 0.421397, 0.414262, 0.406894, 0.399680, 0.392547, 0.385399, 0.378182, 0.370759, 0.363690, 0.356390, 0.349346, 0.341952, 0.335000, 0.327902, 0.320707, 0.313459, 0.306325, 0.299213, 0.291181, 0.283309, 0.275338, 0.267185, 0.259441, 0.251545, 0.243652, 0.235765, 0.228054, 0.220194, 0.212369, 0.204820, 0.197130, 0.189558, 0.182061, 0.174602, 0.167242, 0.159851, 0.152736, 0.145606, 0.138595, 0.131706, 0.125067, 0.118485, 0.112101, 0.105897, 0.099812, 0.094088, 0.088477, 0.082962, 0.077854, 0.072924, 0.068273, 0.063748, 0.059715, 0.055721, 0.052042, 0.048762, 0.045388, 0.042616, 0.039888, 0.037449, 0.035191, 0.033201, 0.031339, 0.029751, 0.028369, 0.027112, 0.025987, 0.025155, 0.024428, 0.023762, 0.023446, 0.023057, 0.022894, 0.022946, 0.022894, 0.023057, 0.023446, 0.023762, 0.024428, 0.025155, 0.025987, 0.027112, 0.028369, 0.029751, 0.031339, 0.033201, 0.035191, 0.037449, 0.039888, 0.042616, 0.045388, 0.048762, 0.052042, 0.055721, 0.059715, 0.063748, 0.068273, 0.072924, 0.077854, 0.082962, 0.088477, 0.094088, 0.099812, 0.105897, 0.112101, 0.118485, 0.125067, 0.131706, 0.138595, 0.145606, 0.152736, 0.159851, 0.167242, 0.174602, 0.182061, 0.189558, 0.197130, 0.204820, 0.212369, 0.220194, 0.228054, 0.235765, 0.243652, 0.251545, 0.259441, 0.267185, 0.275338, 0.283309, 0.291181, 0.299213, 0.306325, 0.313459, 0.320707, 0.327902, 0.335000, 0.341952, 0.349346, 0.356390, 0.363690, 0.370759, 0.378182, 0.385399, 0.392547, 0.399680, 0.406894, 0.414262, 0.421397, 0.428445, 0.435691, 0.443049, 0.450071, 0.457279, 0.464586, 0.471725, 0.478851, 0.486134, 0.493164, 0.500318, 0.507444, 0.514596, 0.521764, 0.528892, 0.535838, 0.542887, 0.549978, 0.557009, 0.564026, 0.571094, 0.577854, 0.584939, 0.591735, 0.598705, 0.605551, 0.612554, 0.619403, 0.626218, 0.632843, 0.639687, 0.646386, 0.652993, 0.659596, 0.666097, 0.671952, 0.677857, 0.683557, 0.689315, 0.695036, 0.700725, 0.706343, 0.711952, 0.717442, 0.722725, 0.728055, 0.733510, 0.738664, 0.743745, 0.748915, 0.753939, 0.758801, 0.763595, 0.768407, 0.773128, 0.777726, 0.782195, 0.786585, 0.790897, 0.795088, 0.799272, 0.803118, 0.807127, 0.810803, 0.814463, 0.817917, 0.821361, 0.824650, 0.827781, 0.830589, 0.833361, 0.836093, 0.838470, 0.840805, 0.842887, 0.844708, 0.846372, 0.847896, 0.849050, 0.849951, 0.850802, 0.851204, 0.664237, 0.666525, 0.668791, 0.671031, 0.673242, 0.675423, 0.677571, 0.679687, 0.681769, 0.683818, 0.685829, 0.687805, 0.689742, 0.691641, 0.693499, 0.695315, 0.697088, 0.698816, 0.700499, 0.702133, 0.703718, 0.705251, 0.706730, 0.708155, 0.709522, 0.710829, 0.712076, 0.713258, 0.714375, 0.715424, 0.716404, 0.717311, 0.718144, 0.718900, 0.719577, 0.720171, 0.720681, 0.721102, 0.721432, 0.721666, 0.721798, 0.721824, 0.721732, 0.721824, 0.721798, 0.721666, 0.721432, 0.721102, 0.720681, 0.720171, 0.719577, 0.718900, 0.718144, 0.717311, 0.716404, 0.715424, 0.714375, 0.713258, 0.712076, 0.710829, 0.709522, 0.708155, 0.706730, 0.705251, 0.703718, 0.702133, 0.700499, 0.698816, 0.697088, 0.695315, 0.693499, 0.691641, 0.689742, 0.687805, 0.685829, 0.683818, 0.681769, 0.679687, 0.677571, 0.675423, 0.673242, 0.671031, 0.668791, 0.666525, 0.664237, 0.662342, 0.660514, 0.658743, 0.657022, 0.655347, 0.653715, 0.652122, 0.650567, 0.649048, 0.647565, 0.646117, 0.644703, 0.643326, 0.641984, 0.640678, 0.639409, 0.638178, 0.636987, 0.635837, 0.634730, 0.633668, 0.632655, 0.631695, 0.630791, 0.629948, 0.629175, 0.628479, 0.627873, 0.627375, 0.627012, 0.628162, 0.629312, 0.630462, 0.631611, 0.632761, 0.633911, 0.632761, 0.631611, 0.630462, 0.629312, 0.628162, 0.627012, 0.627375, 0.627873, 0.628479, 0.629175, 0.629948, 0.630791, 0.631695, 0.632655, 0.633668, 0.634730, 0.635837, 0.636987, 0.638178, 0.639409, 0.640678, 0.641984, 0.643326, 0.644703, 0.646117, 0.647565, 0.649048, 0.650567, 0.652122, 0.653715, 0.655347, 0.657022, 0.658743, 0.660514, 0.662342, 0.665075, 0.665788, 0.666477, 0.667149, 0.667807, 0.668452, 0.669084, 0.669701, 0.670303, 0.670890, 0.671461, 0.672014, 0.672548, 0.673062, 0.673552, 0.674017, 0.674454, 0.674860, 0.675234, 0.675572, 0.675871, 0.675572, 0.675234, 0.674860, 0.674454, 0.674017, 0.673552, 0.673062, 0.672548, 0.672014, 0.671461, 0.670890, 0.670303, 0.669701, 0.669084, 0.668452, 0.667807, 0.667149, 0.666477, 0.665788, 0.665075, 0.664381, 0.663748, 0.663164, 0.662625, 0.662130, 0.661679, 0.661275, 0.660920, 0.660618, 0.660372, 0.660186, 0.660063, 0.660004, 0.660012, 0.660086, 0.660229, 0.660443, 0.660728, 0.661089, 0.661532, 0.661089, 0.660728, 0.660443, 0.660229, 0.660086, 0.660012, 0.660004, 0.660063, 0.660186, 0.660372, 0.660618, 0.660920, 0.661275, 0.661679, 0.662130, 0.662625, 0.663164, 0.663748, 0.664381, 0.375535, 0.380755, 0.385970, 0.391178, 0.396378, 0.401568, 0.406749, 0.411919, 0.417077, 0.422223, 0.427357, 0.432478, 0.437586, 0.442679, 0.447758, 0.452822, 0.457871, 0.462903, 0.467918, 0.472916, 0.477896, 0.482858, 0.487799, 0.492722, 0.497623, 0.502504, 0.507362, 0.512198, 0.517010, 0.521792, 0.526540, 0.531243, 0.559981, 0.564298, 0.572818, 0.564298, 0.559981, 0.531243, 0.526540, 0.521792, 0.517010, 0.512198, 0.507362, 0.502504, 0.497623, 0.492722, 0.487799, 0.482858, 0.477896, 0.472916, 0.467918, 0.462903, 0.457871, 0.452822, 0.447758, 0.442679, 0.437586, 0.432478, 0.427357, 0.422223, 0.417077, 0.411919, 0.406749, 0.401568, 0.396378, 0.391178, 0.385970, 0.380755, 0.375535, 0.543598, 0.549005, 0.549005, 0.543598, 0.521853, 0.516702, 0.511551, 0.506400, 0.501249, 0.496098, 0.490947, 0.485796, 0.480645, 0.475494, 0.470343, 0.465192, 0.460041, 0.454890, 0.449739, 0.444588, 0.439437, 0.434286, 0.429135, 0.423984, 0.418833, 0.413682, 0.408531, 0.403380, 0.398229, 0.393078, 0.387927, 0.382776, 0.377625, 0.372474, 0.367323, 0.362172, 0.357021, 0.362533, 0.362584, 0.366719, 0.370758, 0.374482, 0.377860, 0.380939, 0.383295, 0.385031, 0.386113, 0.386473, 0.386109, 0.385023, 0.383283, 0.380925, 0.377842, 0.374461, 0.370736, 0.366696, 0.362560, 0.358348, 0.354309, 0.350585, 0.347206, 0.344127, 0.341772, 0.340035, 0.338954, 0.338593, 0.338958, 0.340043, 0.341783, 0.344142, 0.347224, 0.350605, 0.354331, 0.358371, 0.375192, 0.372191, 0.369350, 0.366674, 0.364163, 0.361818, 0.359640, 0.357632, 0.355795, 0.354130, 0.352640, 0.351325, 0.350186, 0.349223, 0.348437, 0.347827, 0.347393, 0.347135, 0.347052, 0.347144, 0.347412, 0.347854, 0.348473, 0.349267, 0.350238, 0.351387, 0.352715, 0.354223, 0.355912, 0.357785, 0.359846, 0.362104, 0.364585, 0.367095, 0.369405, 0.371520, 0.373447, 0.375190, 0.376755, 0.378148, 0.379374, 0.380440, 0.381352, 0.382117, 0.382743, 0.383236, 0.383605, 0.383856, 0.383996, 0.384031, 0.383968, 0.383811, 0.383568, 0.383241, 0.382836, 0.382357, 0.381806, 0.381188, 0.380505, 0.379760, 0.378955, 0.378093, 0.377176, 0.376206, 0.362533, 0.362560, 0.366696, 0.370736, 0.374461, 0.377842, 0.380925, 0.383283, 0.385023, 0.386109, 0.386473, 0.386113, 0.385031, 0.383295, 0.380939, 0.377860, 0.374482, 0.370758, 0.366719, 0.362584, 0.358371, 0.354331, 0.350605, 0.347224, 0.344142, 0.341783, 0.340043, 0.338958, 0.338593, 0.338954, 0.340035, 0.341772, 0.344127, 0.347206, 0.350585, 0.354309, 0.358348, 0.375192, 0.372191, 0.369350, 0.366674, 0.364163, 0.361818, 0.359640, 0.357632, 0.355795, 0.354130, 0.352640, 0.351325, 0.350186, 0.349223, 0.348437, 0.347827, 0.347393, 0.347135, 0.347052, 0.347144, 0.347412, 0.347854, 0.348473, 0.349267, 0.350238, 0.351387, 0.352715, 0.354223, 0.355912, 0.357785, 0.359846, 0.362104, 0.364585, 0.367095, 0.369405, 0.371520, 0.373447, 0.375190, 0.376755, 0.378148, 0.379374, 0.380440, 0.381352, 0.382117, 0.382743, 0.383236, 0.383605, 0.383856, 0.383996, 0.384031, 0.383968, 0.383811, 0.383568, 0.383241, 0.382836, 0.382357, 0.381806, 0.381188, 0.380505, 0.379760, 0.378955, 0.378093, 0.377176, 0.376206, 0.298620, 0.293794, 0.288991, 0.284301, 0.279916, 0.276030, 0.272714, 0.269943, 0.267662, 0.265804, 0.264290, 0.263036, 0.261967, 0.261022, 0.260156, 0.259339, 0.258551, 0.257778, 0.257011, 0.257406, 0.257815, 0.258259, 0.258769, 0.259392, 0.260193, 0.261257, 0.262688, 0.264582, 0.266996, 0.269911, 0.273238, 0.276859, 0.280671, 0.284596, 0.288585, 0.292606, 0.296640, 0.296323, 0.296006, 0.295688, 0.295372, 0.295058, 0.294752, 0.294461, 0.294192, 0.293955, 0.293756, 0.293596, 0.293472, 0.293376, 0.293299, 0.293234, 0.293177, 0.293123, 0.293071, 0.293667, 0.294260, 0.294850, 0.295430, 0.295996, 0.296538, 0.297043, 0.297496, 0.297883, 0.298195, 0.298431, 0.298597, 0.298704, 0.298760, 0.298773, 0.298749, 0.298694, 0.296640, 0.292606, 0.288585, 0.284596, 0.280671, 0.276859, 0.273238, 0.269911, 0.266996, 0.264582, 0.262688, 0.261257, 0.260193, 0.259392, 0.258769, 0.258259, 0.257815, 0.257406, 0.257011, 0.257778, 0.258551, 0.259339, 0.260156, 0.261022, 0.261967, 0.263036, 0.264290, 0.265804, 0.267662, 0.269943, 0.272714, 0.276030, 0.279916, 0.284301, 0.288991, 0.293794, 0.298620, 0.298694, 0.298749, 0.298773, 0.298760, 0.298704, 0.298597, 0.298431, 0.298195, 0.297883, 0.297496, 0.297043, 0.296538, 0.295996, 0.295430, 0.294850, 0.294260, 0.293667, 0.293071, 0.293123, 0.293177, 0.293234, 0.293299, 0.293376, 0.293472, 0.293596, 0.293756, 0.293955, 0.294192, 0.294461, 0.294752, 0.295058, 0.295372, 0.295688, 0.296006, 0.296323 +]) +landmarks_2D_1k = np.stack([mean_face_x_1k, mean_face_y_1k], axis=1) + +mean_face_x_137_22_client = np.array([0.4988282, 0.5449964, 0.5849726, 0.6179804, 0.643469, 0.6622178, 0.6730388, 0.676355, 0.6733304, + 0.6478118, + 0.5882786, + 0.4988282, + 0.4093778, + 0.349844, + 0.324326, + 0.32130139999999996, + 0.3246176, + 0.33543860000000003, + 0.35418740000000004, + 0.3796754, + 0.4126838, + 0.45266]) + +mean_face_y_137_22_client = np.array([0.7108352, 0.7000166, 0.6745382, 0.640106, 0.5996582, 0.5467124, 0.4916804, 0.4355282, 0.3795278, 0.2916416, + 0.23122520000000002, + 0.2137676, + 0.23122520000000002, + 0.2916416, + 0.3795278, + 0.4355282, + 0.4916804, + 0.5467124, + 0.5996582, + 0.640106, + 0.6745382, + 0.7000166]) + +landmarks_2D_137_22_clinet = np.stack([mean_face_x_137_22_client, mean_face_y_137_22_client], axis=1) + + +mat_face1024_256_full_face_client = np.array([[4.1666666e-01, 1.5257437e-17, -8.5333336e+01], + [-1.5257449e-17, 4.1666666e-01, -8.5333336e+01]]) + +mat_face1024_256_full_face_server = np.array([[4.1666666e-01, -1.5237085e-17, -8.5333336e+01], + [1.5237085e-17, 4.1666666e-01, -8.5333336e+01]]) + + +# 68 point landmark definitions +landmarks_68_pt = {"mouth": (48, 68), + "right_eyebrow": (17, 22), + "left_eyebrow": (22, 27), + "right_eye": (36, 42), + "left_eye": (42, 48), + "nose": (27, 36), # missed one point + "jaw": (0, 17)} + + +def get_max_rect(bounding_boxes): + max_area = 0 + index = 0 + for i, box in enumerate(bounding_boxes): + width = box[2] - box[0] + height = box[3] - box[1] + if width * height > max_area: + index = i + max_area = width * height + return index + + +def get_transform_mat_mmcv(landmark, output_size): + dst_size = output_size + + if len(landmark) == 68: + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[36] + landmark[39]) / 2 + right_eye = (landmark[42] + landmark[45]) / 2 + nose = (landmark[31] + landmark[35]) / 2 + left_mouth = (landmark[48] + landmark[60]) / 2 + right_mouth = (landmark[64] + landmark[54]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + elif len(landmark) == 87: + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[31] + landmark[35]) / 2 + right_eye = (landmark[39] + landmark[43]) / 2 + nose = landmark[62] + left_mouth = (landmark[66] + landmark[79]) / 2 + right_mouth = (landmark[72] + landmark[83]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + elif len(landmark) == 1000: + pt137 = pts_1k_to_137(landmark) + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (pt137[88] + pt137[96]) / 2 + right_eye = (pt137[105] + pt137[113]) / 2 + nose = pt137[83] + left_mouth = (pt137[22] + pt137[48]) / 2 + right_mouth = (pt137[56] + pt137[36]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + + +def umeyama(src, dst, estimate_scale): + """Estimate N-D similarity transformation with or without scaling. + Parameters + ---------- + src : (M, N) array + Source coordinates. + dst : (M, N) array + Destination coordinates. + estimate_scale : bool + Whether to estimate scaling factor. + Returns + ------- + T : (N + 1, N + 1) + The homogeneous similarity transformation matrix. The matrix contains + NaN values only if the problem is not well-conditioned. + References + ---------- + .. [1] "Least-squares estimation of transformation parameters between two + point patterns", Shinji Umeyama, PAMI 1991, DOI: 10.1109/34.88573 + """ + + num = src.shape[0] + dim = src.shape[1] + + # Compute mean of src and dst. + src_mean = src.mean(axis=0) + dst_mean = dst.mean(axis=0) + + # Subtract mean from src and dst. + src_demean = src - src_mean + dst_demean = dst - dst_mean + + # Eq. (38). + A = np.dot(dst_demean.T, src_demean) / num + + # Eq. (39). + d = np.ones((dim,), dtype=np.double) + if np.linalg.det(A) < 0: + d[dim - 1] = -1 + + T = np.eye(dim + 1, dtype=np.double) + + U, S, V = np.linalg.svd(A) + + # Eq. (40) and (43). + rank = np.linalg.matrix_rank(A) + if rank == 0: + return np.nan * T + elif rank == dim - 1: + if np.linalg.det(U) * np.linalg.det(V) > 0: + T[:dim, :dim] = np.dot(U, V) + else: + s = d[dim - 1] + d[dim - 1] = -1 + T[:dim, :dim] = np.dot(U, np.dot(np.diag(d), V)) + d[dim - 1] = s + else: + T[:dim, :dim] = np.dot(U, np.dot(np.diag(d), V.T)) + + if estimate_scale: + # Eq. (41) and (42). + scale = 1.0 / src_demean.var(axis=0).sum() * np.dot(S, d) + else: + scale = 1.0 + + T[:dim, dim] = dst_mean - scale * np.dot(T[:dim, :dim], src_mean.T) + T[:dim, :dim] *= scale + + return T + +def get_transform_mat_mmcv_bigger(landmark, output_size, forlabel=False): + dst_size = output_size + if len(landmark) == 87: + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[31] + landmark[35]) / 2 + right_eye = (landmark[39] + landmark[43]) / 2 + nose = landmark[62] + left_mouth = (landmark[66] + landmark[79]) / 2 + right_mouth = (landmark[72] + landmark[83]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + elif len(landmark) == 137: + if forlabel: + eye_dis = 0.4 + mouth_dis = 0.4 + g_Average_5point_180 = np.array([ + eye_dis, 0.4, + 1 - eye_dis, 0.4, + 0.5, 0.5, + mouth_dis, 0.6, + 1 - mouth_dis, 0.6 + ]) + else: + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[96] + landmark[88]) / 2 + right_eye = (landmark[105] + landmark[113]) / 2 + nose = landmark[83] + left_mouth = (landmark[22] + landmark[48]) / 2 + right_mouth = (landmark[56] + landmark[36]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + elif len(landmark) == 1000: + # left_eye = (landmarks_2D_1k[691] + landmarks_2D_1k[723]) / 2 + # right_eye = (landmarks_2D_1k[792] + landmarks_2D_1k[824]) / 2 + # nose = landmarks_2D_1k[621] + # left_mouth = (landmarks_2D_1k[467] + landmarks_2D_1k[468]) / 2 + # right_mouth = (landmarks_2D_1k[396] + landmarks_2D_1k[508]) / 2 + # pts5_dst = np.vstack((left_eye, right_eye, + # nose, + # left_mouth, right_mouth)) * dst_size + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + # print(g_Average_5point_180) + left_eye = (landmark[691] + landmark[723]) / 2 + right_eye = (landmark[792] + landmark[824]) / 2 + nose = landmark[621] + left_mouth = (landmark[467] + landmark[468]) / 2 + right_mouth = (landmark[396] + landmark[508]) / 2 + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + + image_to_face_mat = umeyama(pts5_src, pts5_dst, True)[:2] + + return image_to_face_mat + +def get_transform_mat_full_face(landmark, output_size): + dst_size = output_size + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k) + + # print("landmarks_2D_137 x: ", landmarks_2D_137[:22, 0]) + # print("landmarks_2D_137 y: ", landmarks_2D_137[:22, 1]) + + # exit() + + mat = umeyama(landmark[:22], landmarks_2D_137[:22] * dst_size, True)[0:2] + return mat + + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k[:312] * dst_size, True)[0:2] + return mat + + elif len(landmark) == 22: + landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k) + + mat = umeyama(landmark, landmarks_2D_137[:22] * dst_size, True)[0:2] + return mat + +def get_transform_mat_full_face_592(landmark, output_size, ratio): + dst_size = output_size + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + 0.5 + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.5) * ratio + 0.5 + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_137[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + elif len(landmark) == 22: + landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark, landmarks_2D_137[:22] * dst_size, True)[0:2] + return mat + +def get_transform_mat_full_face_592_v1(landmark, output_size, ratio=1.0, h_ratio=0.57): + dst_size = output_size + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + 0.5 + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.5) * ratio + h_ratio + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_137[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_full_face_592_client(landmark, output_size, ratio): + dst_size = output_size + landmarks_2D_137_22_clinet_tmp = landmarks_2D_137_22_clinet.copy() + landmarks_2D_137_22_clinet_tmp[:, 0] = (landmarks_2D_137_22_clinet_tmp[:, 0] - 0.5) * ratio + 0.5 + landmarks_2D_137_22_clinet_tmp[:, 1] = (landmarks_2D_137_22_clinet_tmp[:, 1] - 0.5) * ratio + 0.5 + + if len(landmark) == 22: + # landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark, landmarks_2D_137_22_clinet_tmp * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + landmark_137 = pts_1k_to_137(landmark) + + mat = umeyama(landmark_137[:22], landmarks_2D_137_22_clinet_tmp * dst_size, True)[0:2] + return mat + +def get_transform_mat_face592_full_face(img_size, detect_single_face_size, ratio): + + # face_landmark_client = landmarks_2D_137_22_clinet * img_size + # + # #TODO:客户端如何得到face 0.6 + # face_size2face_M_ratio = get_transform_mat_full_face_592_client(face_landmark_client, img_size, ratio) + + #服务端得到face 1 + face_landmark_server = landmarks_2D_1k * img_size + + face_size2face_M_ratio = get_transform_mat_full_face_592(face_landmark_server, img_size, ratio) + face_size2face_M_full = get_transform_mat_full_face(face_landmark_server, detect_single_face_size) + + + M_ori = np.zeros((3, 3), dtype=np.float32) + M_ori[:2, :] = cv2.invertAffineTransform(face_size2face_M_ratio) + M_ori[2:, :] = [0, 0, 1] + + matAffine_ori = np.zeros((3, 3), dtype=np.float32) + matAffine_ori[:2, :] = face_size2face_M_full + matAffine_ori[2:, :] = [0, 0, 1] + + new_mat = matAffine_ori.dot(M_ori) + + if False: + img = np.zeros((img_size, img_size, 3), dtype=np.uint8) + pred_label_int = face_landmark_server.copy().astype(np.int32) + img_client = img.copy() + for pt in pred_label_int: + cv2.circle(img_client, (pt[0], pt[1]), 1, (0, 0, 255), 1) + # cv2.imshow("img_client: ", img_client) + # + # img_server = np.zeros((img_size, img_size, 3), dtype=np.uint8) + # face_sever_22 = pts_1k_to_137(face_landmark_server)[:22].astype(np.int32) + # for pt in face_sever_22: + # cv2.circle(img_server, (pt[0], pt[1]), 1, (0, 255, 0), 1) + # cv2.imshow("img_server: ", img_server) + # cv2.waitKey() + # + # img_server_new = np.zeros((img_size, img_size, 3), dtype=np.uint8) + # face_sever_new = transform_points(face_landmark_server, face_size2face_M_full) + # face_sever_new_22 = pts_1k_to_137(face_sever_new)[:22].astype(np.int32) + # + # for pt in face_sever_new_22: + # cv2.circle(img_server_new, (pt[0], pt[1]), 1, (0, 255, 0), 1) + # cv2.imshow("img_server_new: ", img_server_new) + # cv2.waitKey() + face_landmark_server_new = transform_points(face_landmark_server, face_size2face_M_ratio) + img_server = np.zeros((img_size, img_size, 3), dtype=np.uint8) + pred_label_server_int = face_landmark_server_new.copy().astype(np.int32) + # img_client = img.copy() + for pt in pred_label_server_int: + cv2.circle(img_server, (pt[0], pt[1]), 1, (0, 0, 255), 1) + cv2.imshow("img_server: ", img_server) + + + img_new = np.zeros((detect_single_face_size, detect_single_face_size, 3), dtype=np.uint8) + pred_new_label_int = transform_points(face_landmark_server_new, new_mat[:2, :]).astype(np.int32) + for pt in pred_new_label_int: + cv2.circle(img_new, (pt[0], pt[1]), 1, (0, 255, 0), 1) + cv2.imshow("img_show new: ", img_new) + cv2.waitKey() + + return new_mat[:2, :] + +def get_transform_mat_for_eye(landmark, output_size): + dst_size = output_size + if len(landmark) == 137: + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[96] + landmark[88]) / 2 + right_eye = (landmark[105] + landmark[113]) / 2 + nose = landmark[83] + left_mouth = (landmark[22] + landmark[48]) / 2 + right_mouth = (landmark[56] + landmark[36]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + +def get_transform_mat_for_face_recognition(landmark, output_size): + g_Average_5point_180 = np.array([ + 57, 73, + 123, 73, + 90, 107, + 62, 134, + 118, 134 + ]) + dst_size = output_size + + if len(landmark) == 87: + left_eye = (landmark[17 + 19] + landmark[17 + 22]) / 2 + right_eye = (landmark[17 + 25] + landmark[17 + 28]) / 2 + nose = (landmark[17 + 14] + landmark[17 + 18]) / 2 + left_mouth = (landmark[17 + 31] + landmark[17 + 43]) / 2 + right_mouth = (landmark[17 + 47] + landmark[17 + 37]) / 2 + elif len(landmark) == 137: + left_eye = (landmark[88] + landmark[96]) / 2 + right_eye = (landmark[105] + landmark[113]) / 2 + nose = landmark[83] + left_mouth = (landmark[22] + landmark[48]) / 2 + right_mouth = (landmark[56] + landmark[36]) / 2 + elif len(landmark) == 1000: + left_eye = (landmark[691] + landmark[723]) / 2 + right_eye = (landmark[792] + landmark[824]) / 2 + nose = landmark[621] + left_mouth = (landmark[467] + landmark[468]) / 2 + right_mouth = (landmark[396] + landmark[508]) / 2 + else: + assert False + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) / 180 * dst_size + + mat = umeyama(pts5_src, pts5_dst, True) + + return mat + +def get_transform_mat_mmcv_hair(landmark, output_size, forlabel=False): + dst_size = output_size + + if len(landmark) == 1000: + eye_dis = 0.42 + mouth_dis = 0.42 + g_Average_5point_180 = np.array([ + eye_dis, 0.48, + 1 - eye_dis, 0.48, + 0.5, 0.53, + mouth_dis, 0.58, + 1 - mouth_dis, 0.58 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[691] + landmark[723]) / 2 + right_eye = (landmark[792] + landmark[824]) / 2 + nose = landmark[621] + left_mouth = (landmark[467] + landmark[468]) / 2 + right_mouth = (landmark[396] + landmark[508]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + else: + print("landmark < 1000 !!!") + assert False + return mat + + +def get_transform_mat_mmcv_seg(landmark, output_size, forlabel=False): + dst_size = output_size + eye_dis = 0.40 + mouth_dis = 0.40 + g_Average_5point_180 = np.array([ + eye_dis, 0.46, + 1 - eye_dis, 0.46, + 0.5, 0.55, + mouth_dis, 0.64, + 1 - mouth_dis, 0.64 + ]) + + if len(landmark) == 1000: + # print(g_Average_5point_180) + left_eye = (landmark[691] + landmark[723]) / 2 + right_eye = (landmark[792] + landmark[824]) / 2 + nose = landmark[621] + left_mouth = (landmark[467] + landmark[468]) / 2 + right_mouth = (landmark[396] + landmark[508]) / 2 + + elif len(landmark) == 137: + left_eye = (landmark[88] + landmark[96]) / 2 + right_eye = (landmark[105] + landmark[113]) / 2 + nose = landmark[83] + left_mouth = (landmark[22] + landmark[48]) / 2 + right_mouth = (landmark[56] + landmark[36]) / 2 + else: + print("landmark < 1000 !!!") + assert False + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + + +def get_transform_mat_two_sets(landmarks_src, landmarks_dst): + assert len(landmarks_src) == len(landmarks_dst) + assert len(landmarks_src) == 137 + mat = umeyama(landmarks_src[:22], landmarks_dst[:22], True)[0:2] + return mat + +def flip_points(landmark, width): + if len(landmark) == 137: + landmarks_order = np.array([1, 22, 21, 20, 19, 18, 17, 16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, # counter + 37, 36, 35, 34, 33, 32, 31, 30, 29, 28, 27, 26, 25, 24, 23, 48, 47, 46, 45, 44, 43, 42, 41, + 40, 39, 38, + 57, 56, 55, 54, 53, 52, 51, 50, 49, 64, 63, 62, 61, 60, 59, 58, # mouth + 79, 78, 77, 76, 75, 74, 73, 72, 71, 70, 69, 68, 67, 66, 65, 83, 82, 81, 80, 84, 85, 86, 87, + # nose + 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, # eye + 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, + 134, 133, 132, 131, 130, 137, 136, 135, # eyebrow + 126, 125, 124, 123, 122, 129, 128, 127], dtype=np.int32) - 1 + landmark_flip = landmark.copy() + landmark_flip[:, 0] = width - 1 - landmark_flip[:, 0] + landmark_flip = landmark_flip[landmarks_order, :] + return landmark_flip + elif len(landmark) == 1000: + dst = [0] + list(range(311, 156, -1)) + list(range(156, 0, -1)) + \ + list(range(396, 354, -1)) + list(range(354, 311, -1)) + list(range(467, 432, -1)) + list(range(432, 396, -1)) + \ + list(range(508, 488, -1)) + list(range(488, 467, -1)) + list(range(547, 528, -1)) + list(range(528, 508, -1)) + \ + list(range(616, 584, -1)) + [584, 583, 582, 581, 580] + list(range(579, 547, -1)) + [620, 619, 618, 617] + list(range(621, 654, 1)) + \ + [755] + list(range(774, 755, -1)) + list(range(791, 774, -1)) + list(range(792, 856, 1)) + \ + [654] + list(range(673, 654, -1)) + list(range(690, 673, -1)) + list(range(691, 755, 1)) + \ + list(range(964, 927, -1)) + list(range(999, 964, -1)) + \ + list(range(892, 855, -1)) + list(range(927, 892, -1)) + landmarks_order = np.array(dst, dtype=np.int32) + landmark_flip = landmark.copy() + landmark_flip[:, 0] = width - 1 - landmark_flip[:, 0] + landmark_flip = landmark_flip[landmarks_order, :] + return landmark_flip + else: + assert False + +def pts_1k_to_137(landmarks): + index_1k_to_137 = [0, 12, 24, 36, 48, 61, 74, 87, 100, 119, 137, 156, 175, 193, 212, 225, 238, 251, 264, 276, 288, + 300, 312, 318, 324, 330, 336, 342, 348, 354, 360, 366, 372, 378, 384, 390, 396, 402, 408, 414, + 420, 426, 432, 438, 444, 450, 456, 462, 468, 473, 478, 483, 488, 493, 498, 503, 508, 513, 518, + 523, 528, 533, 538, 543, 548, 556, 564, 571, 579, 580, 581, 582, 583, 584, 585, 593, 600, 608, + 616, 617, 618, 619, 620, 621, 632, 642, 653, 654, 691, 695, 699, 703, 707, 711, 715, 719, 723, + 727, 731, 735, 739, 743, 747, 751, 755, 792, 796, 800, 804, 808, 812, 816, 820, 824, 828, 832, + 836, 840, 844, 848, 852, 856, 865, 874, 883, 892, 901, 910, 919, 928, 937, 946, 955, 964, 973, + 982, 991] + landmarks_137 = landmarks[index_1k_to_137, :] + return landmarks_137 + +def get_transform_mat_full_face_to_target(landmark, dst_pt): + mat = umeyama(landmark[:312], dst_pt[:312], True)[0:2] + return mat + + +def decompose_affine_matrix(matrix): + # 确保输入矩阵是 2x3 的 + assert matrix.shape == (2, 3), "输入矩阵必须是 2x3 的仿射变换矩阵" + + # 提取平移分量 + tx = matrix[0, 2] + ty = matrix[1, 2] + + # 提取旋转、缩放分量 + a = matrix[0, 0] + b = matrix[0, 1] + c = matrix[1, 0] + d = matrix[1, 1] + + # 计算缩放因子 + scale_x = np.sqrt(a ** 2 + c ** 2) + scale_y = np.sqrt(b ** 2 + d ** 2) + + # 计算旋转角度 + theta = np.arctan2(c, a) + + return { + "translation": (tx, ty), + "scale": (scale_x, scale_y), + "rotation": np.degrees(theta) # 以度数表示 + } + + +# 定义一个高质量缩放图像的函数，主要是用于处理比如头发这种非常细微的图片 +def high_quality_warpAffine(origin_img, M, dst_size, const_value=(255, 255, 255)): + M_params = decompose_affine_matrix(M) + + scale_ratio = M_params['scale'][0] + M_big = M / scale_ratio + crop_face_big = cv2.warpAffine(origin_img, M_big, (int(dst_size[0] / scale_ratio), int(dst_size[1] / scale_ratio)), + cv2.BORDER_CONSTANT, borderValue=const_value) + + if scale_ratio < 1: + crop_face = cv2.resize(crop_face_big, (dst_size[0], dst_size[1]), interpolation=cv2.INTER_AREA) + else: + crop_face = cv2.resize(crop_face_big, (dst_size[0], dst_size[1]), interpolation=cv2.INTER_LANCZOS4) + + return crop_face + + +def transform_points(points, mat, invert=False): + if invert: + mat = cv2.invertAffineTransform(mat) + points = np.expand_dims(points, axis=1) + points = cv2.transform(points, mat, points.shape) + points = np.squeeze(points) + return points +def get_transform_mat_bodyseg(landmark, output_size, ratio=1.0, offset = (0, 0)): + dst_size = output_size + assert len(landmark) == 1000 + mat = umeyama(landmark[:312], landmarks_2D_1k[:312] * dst_size * ratio + dst_size * (1 - ratio) / 2 + offset, True)[0:2] + return mat +def align2stylegan(face_landmarks_1k, output_size=256): + face_landmarks_1k = np.float32(face_landmarks_1k) + x_scale = 1.0 + y_scale = 1.0 + em_scale = 0.1 + eye_left = (face_landmarks_1k[691] + face_landmarks_1k[723]) / 2 + eye_right = (face_landmarks_1k[792] + face_landmarks_1k[824]) / 2 + mouth_left = (face_landmarks_1k[467] + face_landmarks_1k[468]) / 2 + mouth_right = (face_landmarks_1k[396] + face_landmarks_1k[508]) / 2 + eye_avg = (eye_left + eye_right) * 0.5 + eye_to_eye = eye_right - eye_left + mouth_avg = (mouth_left + mouth_right) * 0.5 + eye_to_mouth = mouth_avg - eye_avg + x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1] + x /= np.hypot(*x) + x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8) + x *= x_scale + y = np.flipud(x) * [-y_scale, y_scale] + c = eye_avg + eye_to_mouth * em_scale + quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y]) + quad_ori = np.array(quad) + + rotate_radian = math.atan2((quad_ori[3][1] - quad_ori[0][1]), (quad_ori[3][0] - quad_ori[0][0])) + rotate_degree = rotate_radian / np.pi * 180 + scale = output_size / cv2.norm(quad_ori[3] - quad_ori[0]) + src_center = (quad_ori[0] + quad_ori[2]) * 0.5 + dst_center = np.float32([output_size / 2, output_size / 2]) + + M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), rotate_degree, scale) + M[:, 2] += dst_center - src_center + return M +def align2stylegan_ratio(face_landmarks_1k, ratio=1.0, output_size=256): + face_landmarks_1k = np.float32(face_landmarks_1k) + x_scale = 1.0 / ratio + y_scale = 1.0 / ratio + em_scale = 0.1 + eye_left = (face_landmarks_1k[691] + face_landmarks_1k[723]) / 2 + eye_right = (face_landmarks_1k[792] + face_landmarks_1k[824]) / 2 + mouth_left = (face_landmarks_1k[467] + face_landmarks_1k[468]) / 2 + mouth_right = (face_landmarks_1k[396] + face_landmarks_1k[508]) / 2 + eye_avg = (eye_left + eye_right) * 0.5 + eye_to_eye = eye_right - eye_left + mouth_avg = (mouth_left + mouth_right) * 0.5 + eye_to_mouth = mouth_avg - eye_avg + x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1] + x /= np.hypot(*x) + x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8) + x *= x_scale + y = np.flipud(x) * [-y_scale, y_scale] + c = eye_avg + eye_to_mouth * em_scale + quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y]) + quad_ori = np.array(quad) + + rotate_radian = math.atan2((quad_ori[3][1] - quad_ori[0][1]), (quad_ori[3][0] - quad_ori[0][0])) + rotate_degree = rotate_radian / np.pi * 180 + scale = output_size / cv2.norm(quad_ori[3] - quad_ori[0]) + src_center = (quad_ori[0] + quad_ori[2]) * 0.5 + dst_center = np.float32([output_size / 2, output_size / 2]) + + M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), rotate_degree, scale) + M[:, 2] += dst_center - src_center + return M +def calc_face_pitch(landmarks): + if not isinstance(landmarks, np.ndarray): + landmarks = np.array(landmarks) + t = ((landmarks[6][1] - landmarks[8][1]) + (landmarks[10][1] - landmarks[8][1])) / 2.0 + b = landmarks[8][1] + return float(b - t) + +def calc_face_yaw(landmarks): + if not isinstance(landmarks, np.ndarray): + landmarks = np.array(landmarks) + l = ((landmarks[27][0] - landmarks[0][0]) + (landmarks[28][0] - landmarks[1][0]) + ( + landmarks[29][0] - landmarks[2][0])) / 3.0 + r = ((landmarks[16][0] - landmarks[27][0]) + (landmarks[15][0] - landmarks[28][0]) + ( + landmarks[14][0] - landmarks[29][0])) / 3.0 + return float(r - l) + +# deprecated +def draw_pncc_features(fc_landmark, img_target, w=256, h=256, is_train=True): + assert False + return img_target + + +def draw_blur_no_mouth_img(img_target, fc_landmark, is_train=False): + h, w, c = img_target.shape + assert h == w + + if len(fc_landmark) == 1000: + fc_landmark = pts_1k_to_137(fc_landmark) + + + k_mid_size = int(w / 6.) + thresh = int(w / 15.) + if is_train: + k_size = random.choice(range(k_mid_size - thresh, k_mid_size + thresh, 2)) + else: + k_size = k_mid_size + if k_size % 2 == 0: + k_size += 1 + img_target_blur = cv2.blur(img_target, (k_size, k_size)) + + + hull_mask = np.zeros(img_target.shape, dtype=np.uint8) + if len(fc_landmark) == 1000: + fc_landmark = pts_1k_to_137(fc_landmark) + + cv2.fillPoly(hull_mask, fc_landmark[22:48][np.newaxis, :, :], (255, 255, 255)) + # cv2.imshow('hull_mask1', hull_mask) + # cv2.imshow('img_target_blur', img_target_blur) + # cv2.imshow('img_target1', img_target) + kernel_size = int(w / 35.) + if kernel_size % 2 == 0: + kernel_size += 1 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + hull_mask = cv2.dilate(hull_mask, kernel, iterations=1) .astype(np.uint8) + # cv2.imshow('hull_mask2', hull_mask) + img_target = img_target_blur * (hull_mask<125) + img_target * (hull_mask>125) + + # cv2.imshow('img_target2', img_target) + # cv2.waitKey() + + return img_target + + + +def draw_blur_img(img_target, fc_landmark, is_train=False): + h, w, c = img_target.shape + assert h == w + inpaint_mask = np.zeros((h, w), dtype=np.uint8) + + if len(fc_landmark) == 1000: + fc_landmark = pts_1k_to_137(fc_landmark) + + # cv2.imshow('img_target1', img_target) + + k_mid_size = int(w / 6.) + thresh = int(w / 15.) + if is_train: + k_size = random.choice(range(k_mid_size - thresh, k_mid_size + thresh, 2)) + else: + k_size = k_mid_size + if k_size % 2 == 0: + k_size += 1 + img_target = cv2.blur(img_target, (k_size, k_size)) + + # cv2.imshow('img_target2', img_target) + # cv2.waitKey() + + return img_target + +def draw_blurmore_img(img_target, fc_landmark, is_train=False): + h, w, c = img_target.shape + assert h == w + inpaint_mask = np.zeros((h, w), dtype=np.uint8) + + if len(fc_landmark) == 1000: + fc_landmark = pts_1k_to_137(fc_landmark) + + # cv2.imshow('img_target1', img_target) + + k_mid_size = int(w / 3.) + thresh = int(w / 15.) + if is_train: + k_size = random.choice(range(k_mid_size - thresh, k_mid_size + thresh, 2)) + else: + k_size = k_mid_size + if k_size % 2 == 0: + k_size += 1 + img_target = cv2.blur(img_target, (k_size, k_size)) + + # cv2.imshow('img_target2', img_target) + # cv2.waitKey() + + return img_target + + +def draw_comapre_mask(fc_landmark, w=256, h=256): + inpaint_mask = np.zeros((h, w), dtype=np.uint8) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(fc_landmark[129:137]), (255,)) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(fc_landmark[121:129]), (255,)) + + left_eye = np.around((fc_landmark[96] + fc_landmark[88]) / 2).astype(np.int32) + right_eye = np.around((fc_landmark[105] + fc_landmark[113]) / 2).astype(np.int32) + len = cv2.norm(fc_landmark[96].astype(np.int32) - fc_landmark[88].astype(np.int32)) + cv2.circle(inpaint_mask, (left_eye[0], left_eye[1]), 1, (255), int(len * 1.1)) + len = cv2.norm(fc_landmark[105].astype(np.int32) - fc_landmark[113].astype(np.int32)) + + cv2.circle(inpaint_mask, (right_eye[0], right_eye[1]), 1, (255), int(len * 1.1)) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(fc_landmark[64:87]), (255,)) + + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(fc_landmark[22:64]), (255,)) + + dilate_kernel_size = int(w / 23.) + if dilate_kernel_size % 2 == 0: + dilate_kernel_size += 1 + kernel = np.ones((dilate_kernel_size, dilate_kernel_size), np.uint8) + inpaint_mask = cv2.dilate(inpaint_mask, kernel) + + inpaint_mask = 1 - inpaint_mask.astype(np.float32) / 255 + + return inpaint_mask + +def draw_hull_mask(fc_landmark, w=256, h=256, is_gray=False): + if not is_gray: + hull_mask = np.zeros((h, w, 3), dtype=np.uint8) + + line_size1 = int(w / 85.) + line_size2 = int(w / 51.) + + if len(fc_landmark) == 137: + # cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark[0:22]), (64, 16, 32)) + + # left_brown = ((fc_landmark[129] + fc_landmark[133]) / 2).astype(np.int32) + # cv2.circle(hull_mask, (left_brown[0], left_brown[1]), 1, (0, 255, 0), line_size1) + # right_brwon = ((fc_landmark[121] + fc_landmark[125]) / 2).astype(np.int32) + # cv2.circle(hull_mask, (right_brwon[0], right_brwon[1]), 1, (0, 0, 255), line_size1) + + cv2.line(hull_mask, (int(fc_landmark[129, 0]), int(fc_landmark[129, 1])), + (int(fc_landmark[133, 0]), int(fc_landmark[133, 1])), (0, 255, 0), line_size1) + cv2.line(hull_mask, (int(fc_landmark[121, 0]), int(fc_landmark[121, 1])), + (int(fc_landmark[125, 0]), int(fc_landmark[125, 1])), (0, 0, 255), line_size1) + + cv2.fillPoly(hull_mask, fc_landmark[88:104][np.newaxis, :, :], (128, 128, 0)) + cv2.fillPoly(hull_mask, fc_landmark[105:121][np.newaxis, :, :], (128, 0, 128)) + + # cv2.line(hull_mask, (int(fc_landmark[86, 0]), int(fc_landmark[86, 1])), + # (int(fc_landmark[83, 0]), int(fc_landmark[83, 1])), (255, 0, 0), line_size2) + + cv2.fillPoly(hull_mask, fc_landmark[48:64][np.newaxis, :, :], (0, 128, 128)) + + # cv2.fillPoly(hull_mask, np.concatenate((fc_landmark[22:37], fc_landmark[56:47:-1]))[np.newaxis, :, :], + # (0, 128, 0)) + # + # cv2.fillPoly(hull_mask, + # np.concatenate((fc_landmark[47:35:-1], fc_landmark[56:64], [fc_landmark[48], fc_landmark[22]]))[ + # np.newaxis, :, :], (0, 0, 128)) + + eye = np.zeros((h, w, 3), dtype=np.uint8) + eye_mask1 = np.zeros((h, w, 1), dtype=np.uint8) + eye_mask2 = np.zeros((h, w, 1), dtype=np.uint8) + cv2.fillPoly(eye_mask1, fc_landmark[88:104][np.newaxis, :, :], (1,)) + cv2.fillPoly(eye_mask1, fc_landmark[105:121][np.newaxis, :, :], (1,)) + left_eye = fc_landmark[87] + right_eye = fc_landmark[104] + # left_eye = np.around((fc_landmark[96] + fc_landmark[88]) / 2).astype(np.int32) + # right_eye = np.around((fc_landmark[105] + fc_landmark[113]) / 2).astype(np.int32) + cv2.circle(eye_mask2, (left_eye[0], left_eye[1]), 1, (1,), line_size2) + cv2.circle(eye_mask2, (right_eye[0], right_eye[1]), 1, (1,), line_size2) + eye_mask = eye_mask1 & eye_mask2 + cv2.circle(eye, (left_eye[0], left_eye[1]), 1, (255, 255, 255), line_size2) + cv2.circle(eye, (right_eye[0], right_eye[1]), 1, (255, 255, 255), line_size2) + hull_mask = hull_mask * (1 - eye_mask[:, :, 0:1]) + eye * eye_mask[:, :, 0:1] + elif len(fc_landmark) == 1000: + # cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark[0:312]), (64, 16, 32)) + + # left_brown = ((fc_landmark[928] + fc_landmark[964]) / 2).astype(np.int32) + # cv2.circle(hull_mask, (left_brown[0], left_brown[1]), 1, (0, 255, 0), line_size1) + # right_brwon = ((fc_landmark[856] + fc_landmark[892]) / 2).astype(np.int32) + # cv2.circle(hull_mask, (right_brwon[0], right_brwon[1]), 1, (0, 0, 255), line_size1) + + cv2.line(hull_mask, (int(fc_landmark[928, 0]), int(fc_landmark[928, 1])), + (int(fc_landmark[964, 0]), int(fc_landmark[964, 1])), (0, 255, 0), line_size1) + cv2.line(hull_mask, (int(fc_landmark[856, 0]), int(fc_landmark[856, 1])), + (int(fc_landmark[892, 0]), int(fc_landmark[892, 1])), (0, 0, 255), line_size1) + + cv2.fillPoly(hull_mask, fc_landmark[691:755][np.newaxis, :, :], (128, 128, 0)) + cv2.fillPoly(hull_mask, fc_landmark[792:856][np.newaxis, :, :], (128, 0, 128)) + + # cv2.line(hull_mask, (int(fc_landmark[653, 0]), int(fc_landmark[653, 1])), + # (int(fc_landmark[621, 0]), int(fc_landmark[621, 1])), (255, 0, 0), line_size2) + + cv2.fillPoly(hull_mask, fc_landmark[468:548][np.newaxis, :, :], (0, 128, 128)) + + # cv2.fillPoly(hull_mask, np.concatenate((fc_landmark[312:397], fc_landmark[508:467:-1]))[np.newaxis, :, :], + # (0, 128, 0)) + # + # cv2.fillPoly(hull_mask, + # np.concatenate((fc_landmark[467:395:-1], fc_landmark[508:548], [fc_landmark[468], fc_landmark[312]]))[ + # np.newaxis, :, :], (0, 0, 128)) + + eye = np.zeros((h, w, 3), dtype=np.uint8) + eye_mask1 = np.zeros((h, w, 1), dtype=np.uint8) + eye_mask2 = np.zeros((h, w, 1), dtype=np.uint8) + cv2.fillPoly(eye_mask1, fc_landmark[691:755][np.newaxis, :, :], (1,)) + cv2.fillPoly(eye_mask1, fc_landmark[792:856][np.newaxis, :, :], (1,)) + left_eye = fc_landmark[654] + right_eye = fc_landmark[755] + # cv2.fillPoly(eye_mask2, fc_landmark[655:691][np.newaxis, :, :], (1,)) + cv2.circle(eye_mask2, (left_eye[0], left_eye[1]), 1, (1,), line_size2) + # cv2.fillPoly(eye_mask2, fc_landmark[756:792][np.newaxis, :, :], (1,)) + cv2.circle(eye_mask2, (right_eye[0], right_eye[1]), 1, (1,), line_size2) + eye_mask = eye_mask1 & eye_mask2 + # cv2.fillPoly(eye, fc_landmark[655:691][np.newaxis, :, :], (255, 255, 255)) + cv2.circle(eye, (left_eye[0], left_eye[1]), 1, (255, 255, 255), line_size2) + # cv2.fillPoly(eye, fc_landmark[756:792][np.newaxis, :, :], (255, 255, 255)) + cv2.circle(eye, (right_eye[0], right_eye[1]), 1, (255, 255, 255), line_size2) + hull_mask = hull_mask * (1 - eye_mask[:, :, 0:1]) + eye * eye_mask[:, :, 0:1] + else: + assert False + else: + hull_mask = np.zeros((h, w), dtype=np.uint8) + if len(fc_landmark) == 137: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (1)) + elif len(fc_landmark) == 1000: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (1)) + else: + assert False + + return hull_mask + +def draw_half_mask(fc_landmark, w=256, h=256, is_gray=False): + hull_mask = np.zeros((h, w), dtype=np.uint8) + pt2draws = np.concatenate([fc_landmark[:80], fc_landmark[236:312]], axis=0) + cv2.fillConvexPoly(hull_mask, cv2.convexHull(pt2draws), (1)) + return hull_mask + + + +def draw_makeup_mask(another_pts1k, another_mask, pts1k, hull_mask): + pts1kint = pts1k.astype(np.int32) + assert len(another_pts1k) == 1000 and len(pts1kint) == 1000 + cv2.fillPoly(another_mask, another_pts1k[691:755][np.newaxis, :, :], (4)) # eye + cv2.fillPoly(another_mask, another_pts1k[792:856][np.newaxis, :, :], (5)) # eye + cv2.fillPoly(another_mask, np.concatenate((another_pts1k[312:397], another_pts1k[508:467:-1]))[np.newaxis, :, :], (7)) # mouth + cv2.fillPoly(another_mask, np.concatenate((another_pts1k[467:395:-1], another_pts1k[508:548], [another_pts1k[468], another_pts1k[312]]))[np.newaxis, :, :], (9)) # mouth + cv2.fillPoly(another_mask, another_pts1k[928:1000][np.newaxis, :, :], (0)) # eyebrow + cv2.fillPoly(another_mask, another_pts1k[856:928][np.newaxis, :, :], (0)) # eyebrow + + + cv2.fillPoly(hull_mask, pts1kint[691:755][np.newaxis, :, :], (4)) # eye + cv2.fillPoly(hull_mask, pts1kint[792:856][np.newaxis, :, :], (5)) # eye + cv2.fillPoly(hull_mask, np.concatenate((pts1kint[312:397], pts1kint[508:467:-1]))[np.newaxis, :, :], (7)) # mouth + cv2.fillPoly(hull_mask, np.concatenate((pts1kint[467:395:-1], pts1kint[508:548], [pts1kint[468], pts1kint[312]]))[np.newaxis, :, :], (9)) # mouth + cv2.fillPoly(hull_mask, pts1kint[928:1000][np.newaxis, :, :], (0)) # eyebrow + cv2.fillPoly(hull_mask, pts1kint[856:928][np.newaxis, :, :], (0)) # eyebrow + return another_mask, hull_mask + +def draw_users_hull_mask(fc_landmark, w=256, h=256): + hull_mask = np.zeros((h, w), dtype=np.uint8) + if len(fc_landmark) == 87: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (1)) + elif len(fc_landmark) == 137: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (1)) + elif len(fc_landmark) == 1000: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (1)) + else: + assert False + return hull_mask + +def draw_bigger_hull_mask(fc_landmark, w=256, h=256): + + if len(fc_landmark) == 1000: + fc_landmark = pts_1k_to_137(fc_landmark) + + hull_mask = np.zeros((h, w), dtype=np.uint8) + if len(fc_landmark) == 137: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (255)) + elif len(fc_landmark) == 1000: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (255)) + else: + assert False + + kernel_size = int(w / 11.) + if kernel_size % 2 == 0: + kernel_size += 1 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + hull_mask = (cv2.dilate(hull_mask, kernel, iterations=1) / 255.).astype(np.uint8) + cv2.fillPoly(hull_mask, fc_landmark[22:48][np.newaxis, :, :], (10,)) + + + leye_mask = np.zeros((h, w), dtype=np.uint8) + cv2.fillPoly(leye_mask, fc_landmark[88:104][np.newaxis, :, :], (1,)) + cv2.fillPoly(leye_mask, fc_landmark[105:121][np.newaxis, :, :], (1,)) + kernel_size = int(w / 5.) + if kernel_size % 2 == 0: + kernel_size += 1 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + leye_mask = (cv2.dilate(leye_mask, kernel, iterations=1)).astype(np.uint8) + hull_mask[leye_mask > 0] = 10 + + # cv2.imshow('hull_mask', hull_mask * 25 ) + # cv2.waitKey() + + return hull_mask + + + +def get_kernel_size(div_num, output_img_size=256): + kernel_size = int(float(output_img_size) / div_num) + if kernel_size % 2 == 0: + kernel_size += 1 + return kernel_size + +class LKTracking(object): + def __init__(self, termcrit=cv2.TERM_CRITERIA_COUNT | cv2.TERM_CRITERIA_EPS, winSize=31, extSize=41, threshold=5): + self.termcrit = termcrit + self.winSize = winSize + self.extSize = extSize + self.threshold = threshold + + self.prepoints = [] + self.preimg_small_ = None + self.preimg_rect_ = np.array([0, 0, 0, 0]) + + def UpdatePoints(self, new_pts): + if len(new_pts) != len(self.prepoints): + self.prepoints = [] + return False + self.prepoints = new_pts + return True + + def TrackingPoints(self, curimg, curpoints): + need_init = False + + def Prepare(self, curimg, curpoints): + prepoints_int = curpoints.astype(np.int32) + pass + +def get_transform_mat_mouth(landmark, output_size): + # mean_mouth_x_4pts = np.array([0.15, 0.5, 0.75, 0.5]) + # mean_mouth_y_4pts = np.array([0.5, 0.15, 0.5, 0.75]) + mean_mouth_x_4pts = np.array([0.2, 0.5, 0.8, 0.5]) + mean_mouth_y_4pts = np.array([0.5, 0.2, 0.5, 0.8]) + landmarks_2D_mouth_4pts = np.stack([mean_mouth_x_4pts, mean_mouth_y_4pts], axis=1) + + dst_size = output_size + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k) + mat = umeyama(landmark[:22], landmarks_2D_137[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 236: + #mouth crop + landmark_mouth = [] + landmark_mouth.append(landmark[467-312]) + landmark_mouth.append(landmark[432-312]) + landmark_mouth.append(landmark[397-312]) + landmark_mouth.append(landmark[354-312]) + landmark_mouth = np.array(landmark_mouth) + mat = umeyama(landmark_mouth, landmarks_2D_mouth_4pts * dst_size, True)[0:2] + return mat + + +mean_mouth_x_4pts = np.array([0.2, 0.5, 0.8, 0.5]) +mean_mouth_y_4pts = np.array([0.5, 0.2, 0.5, 0.8]) +landmarks_2D_mouth_4pts = np.stack([mean_mouth_x_4pts, mean_mouth_y_4pts], axis=1) + +def get_transform_mat_for_mouth(landmark, output_size): + dst_size = output_size + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + # landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k) + # mat = umeyama(landmark[:22], landmarks_2D_137[:22] * dst_size, True)[0:2] + landmark_mouth = [] + landmark_mouth.append(landmark[22]) + landmark_mouth.append(landmark[42]) + landmark_mouth.append(landmark[36]) + landmark_mouth.append(landmark[29]) + landmark_mouth = np.array(landmark_mouth) + mat = umeyama(landmark_mouth, landmarks_2D_mouth_4pts * dst_size, True)[0:2] + + return mat + elif len(landmark) == 236: + # mouth crop + landmark_mouth = [] + landmark_mouth.append(landmark[467 - 312]) + landmark_mouth.append(landmark[432 - 312]) + landmark_mouth.append(landmark[397 - 312]) + landmark_mouth.append(landmark[354 - 312]) + landmark_mouth = np.array(landmark_mouth) + mat = umeyama(landmark_mouth, landmarks_2D_mouth_4pts * dst_size, True)[0:2] + # mat = umeyama(landmark[:312], landmarks_2D_1k[:312] * dst_size, True)[0:2] + return mat + + elif len(landmark) == 1000: + # mouth crop + landmark_mouth = [] + landmark_mouth.append(landmark[467]) + landmark_mouth.append(landmark[432]) + landmark_mouth.append(landmark[397]) + landmark_mouth.append(landmark[354]) + landmark_mouth = np.array(landmark_mouth) + mat = umeyama(landmark_mouth, landmarks_2D_mouth_4pts * dst_size, True)[0:2] + # mat = umeyama(landmark[:312], landmarks_2D_1k[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_full_face_ratio_skin(landmark, output_size, ratio=1.0, skin=0): + dst_size = output_size + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + # landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + 0.5 + # landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.5) * ratio + 0.4 + + if skin == 1: + ###################### 0.3 face rata 592 + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + 0.5 + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.6) * ratio + 0.3 + ###################### 0.3 face rata 592 + else: + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.4) * ratio + 0.45 + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.4) * ratio + 0.4 + + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_1k_tmp = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_1k_tmp[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_for_uv(landmark, output_size): + dst_size = output_size + if len(landmark) == 1000: + landmark = pts_1k_to_137(landmark) + + if len(landmark) == 137: + eye_dis = 0.35 + mouth_dis = 0.42 + g_Average_5point_180 = np.array([ + eye_dis, 0.24, + 1 - eye_dis, 0.24, + 0.5, 0.42, + mouth_dis, 0.55, + 1 - mouth_dis, 0.55 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[96] + landmark[88]) / 2 + right_eye = (landmark[105] + landmark[113]) / 2 + nose = landmark[83] + left_mouth = (landmark[22] + landmark[48]) / 2 + right_mouth = (landmark[56] + landmark[36]) / 2 + + pts5_dst = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_src = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + else: + assert False + + +def draw_crop_eye_bysize(img, pts1k, img_size, change_eyebrow): + pts137tmp = pts_1k_to_137(pts1k).astype(np.int32) + eye_mask = np.zeros((img_size, img_size, 3), dtype=np.uint8) + cv2.fillPoly(eye_mask, pts137tmp[88:104][np.newaxis, :, :], (1, 1, 1)) + cv2.fillPoly(eye_mask, pts137tmp[105:121][np.newaxis, :, :], (1, 1, 1)) + kernel_size = int(img_size / 198.) # 31 + if kernel_size % 2 == 0: + kernel_size += 1 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + eye_mask = (cv2.dilate(eye_mask, kernel, iterations=1)).astype(np.uint8) + + kernel_size = int(img_size / 7) # 8.2 + if kernel_size % 2 == 0: + kernel_size += 1 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size*2, kernel_size)) + eye_mask2 = (cv2.dilate(eye_mask, kernel, iterations=1)).astype(np.uint8) + eye_mask2 = eye_mask2 - eye_mask + + if change_eyebrow: + # eyebrow + eyebrow_mask = np.zeros((img_size, img_size, 3), dtype=np.uint8) + + x, y, w, h = cv2.boundingRect(pts137tmp[129:137]) + # cv2.rectangle(eyebrow_mask, (x, y), (x + w, y + h), (1, 1, 1), -1) + eyebrow_mask[y:y + h, x:x + w, :] = [1, 1, 1] + + x, y, w, h = cv2.boundingRect(pts137tmp[121:129]) + # cv2.rectangle(eyebrow_mask, (x, y), (x + w, y + h), (1, 1, 1), -1) + eyebrow_mask[y:y + h, x:x + w, :] = [1, 1, 1] + + kernel_size = int(592 / 21.) # 21 + if kernel_size % 2 == 0: + kernel_size += 1 + kernel_eyebrow = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (int(kernel_size * 5), int(kernel_size))) + eyebrow_mask = (cv2.dilate(eyebrow_mask, kernel_eyebrow, iterations=1)).astype(np.uint8) + eye_mask2 = (eye_mask2.astype(bool) | eyebrow_mask.astype(bool)).astype(np.float32) + else: + eyebrow_mask = np.zeros((img_size, img_size, 3), dtype=np.uint8) + cv2.fillPoly(eyebrow_mask, pts137tmp[129:137][np.newaxis, :, :], (1, 1, 1)) + cv2.fillPoly(eyebrow_mask, pts137tmp[121:129][np.newaxis, :, :], (1, 1, 1)) + eyebrow_mask = 1 - eyebrow_mask + eye_mask2 = (eye_mask2.astype(bool) & eyebrow_mask.astype(bool)).astype(np.float32) + + # nose + cv2.fillPoly(eye_mask2, pts137tmp[64:79][np.newaxis, :, :], (0, 0, 0)) + eye_mask2[(pts137tmp[133, 1] - img_size // 10):pts137tmp[65, 1], pts137tmp[64, 0]:pts137tmp[78, 0]] = 0 + + img[eye_mask2 > 0] = 0 + + cv2.circle(img, (pts137tmp[133, 0], pts137tmp[133, 1]), img_size // 50, (255, 0, 0), -1) + cv2.circle(img, (pts137tmp[121, 0], pts137tmp[121, 1]), img_size // 50, (255, 0, 0), -1) + + return img + + +# def draw_crop_eye_bysize_using_seg_mask(img, pts1k, mask, img_size, change_eyebrow): +# +# pts137tmp = pts_1k_to_137(pts1k).astype(np.int32) +# c0, c1, c2 = mask[:, :, 0], mask[:, :, 1], mask[:, :, 2] +# c0[c0 >= 100] = 255 +# c0[c0 < 100] = 0 +# c1[c1 >= 100] = 255 +# c1[c1 < 100] = 0 +# c2[c2 > 0] = 0 +# mask[:, :, 0] = c0 +# mask[:, :, 1] = c1 +# mask[:, :, 2] = c2 +# eye_index = (mask == [255, 0, 0]).all(axis=2) +# eyelids_index = (mask == [0, 255, 0]).all(axis=2) +# black_im = np.zeros((img_size, img_size, 3), dtype=np.uint8) +# black_im2 = np.zeros((img_size, img_size, 3), dtype=np.uint8) +# +# black_im[eye_index] = [255, 255, 255] +# black_im2[eye_index] = [255, 255, 255] +# kernel_size = int(img_size // 7) # 8.2 +# +# if kernel_size % 2 == 0: +# kernel_size += 1 +# kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size * 2, kernel_size)) +# black_im2 = cv2.dilate(black_im2, kernel).astype(np.uint8) +# +# endless_belt_im = black_im2 - black_im +# endless_belt_im = np.ones_like(endless_belt_im) * 255 - endless_belt_im +# +# if change_eyebrow: +# # eyebrow +# eyebrow_mask = np.zeros((img_size, img_size, 3), dtype=np.uint8) +# x, y, w, h = cv2.boundingRect(pts137tmp[129:137]) +# # cv2.rectangle(eyebrow_mask, (x, y), (x + w, y + h), (1, 1, 1), -1) +# eyebrow_mask[y:y + h, x:x + w, :] = [1, 1, 1] +# +# x, y, w, h = cv2.boundingRect(pts137tmp[121:129]) +# # cv2.rectangle(eyebrow_mask, (x, y), (x + w, y + h), (1, 1, 1), -1) +# eyebrow_mask[y:y + h, x:x + w, :] = [1, 1, 1] +# +# kernel_size = int(592 / 21.) # 21 +# if kernel_size % 2 == 0: +# kernel_size += 1 +# kernel_eyebrow = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (int(kernel_size * 5), int(kernel_size))) +# eyebrow_mask = (cv2.dilate(eyebrow_mask, kernel_eyebrow, iterations=1)).astype(np.uint8) +# eyebrow_mask = 1 - eyebrow_mask +# endless_belt_im = (endless_belt_im.astype(bool) & eyebrow_mask.astype(bool)).astype(np.float32) * 255 +# else: +# eyebrow_mask = np.zeros((img_size, img_size, 3), dtype=np.uint8) +# cv2.fillPoly(eyebrow_mask, pts137tmp[129:137][np.newaxis, :, :], (1, 1, 1)) +# cv2.fillPoly(eyebrow_mask, pts137tmp[121:129][np.newaxis, :, :], (1, 1, 1)) +# endless_belt_im = (endless_belt_im.astype(bool) | eyebrow_mask.astype(bool)).astype(np.float32) * 255 +# +# cv2.fillPoly(endless_belt_im, pts137tmp[64:79][np.newaxis, :, :], (255, 255, 255)) +# endless_belt_im[(pts137tmp[133, 1] - img_size // 10):pts137tmp[65, 1], pts137tmp[64, 0]:pts137tmp[78, 0], :] = 255 +# +# +# res = np.uint8(endless_belt_im / 255) +# res = np.uint8(res * img) +# +# # add eyelids semantic map +# eyelids_mask = np.zeros((img_size, img_size, 3)) +# eyelids_mask[eyelids_index] = [0, 255, 0] +# res = res + np.uint8(eyelids_mask) +# +# cv2.circle(res, (pts137tmp[133, 0], pts137tmp[133, 1]), img_size // 50, (255, 0, 0), -1) +# cv2.circle(res, (pts137tmp[121, 0], pts137tmp[121, 1]), img_size // 50, (255, 0, 0), -1) +# +# return res + + +def draw_crop_eye_bysize_using_seg_mask(img, pts1k, mask, img_size, change_eyebrow): + start = time.time() + mask = mask.copy() + pts137tmp = pts_1k_to_137(pts1k).astype(np.int32) + c0, c1, c2 = mask[:, :, 0], mask[:, :, 1], mask[:, :, 2] + c0[c0 >= 100] = 255 + c0[c0 < 100] = 0 + + c1[c1 >= 100] = 255 + c1[c1 < 100] = 0 + + # c2 也需要截断 + c2[c2 >= 100] = 255 + c2[c2 < 100] = 0 + + mask[:, :, 0] = c0 + mask[:, :, 1] = c1 + mask[:, :, 2] = c2 + eye_index = (mask == [255, 0, 0]).all(axis=2) | (mask == [0, 0, 255]).all(axis=2) + eyelids_index = (mask == [0, 255, 0]).all(axis=2) + black_im = np.zeros((img_size, img_size, 3), dtype=np.uint8) + black_im2 = np.zeros((img_size, img_size, 3), dtype=np.uint8) + + black_im[eye_index] = [255, 255, 255] + black_im2[eye_index] = [255, 255, 255] + + # cv2.imshow("black_im: ", black_im) + # cv2.waitKey() + + kernel_size = int(img_size // 7) # 8.2 + + if kernel_size % 2 == 0: + kernel_size += 1 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size * 2, kernel_size)) + black_im2 = cv2.dilate(black_im2, kernel).astype(np.uint8) + + endless_belt_im = black_im2 - black_im + endless_belt_im = np.ones_like(endless_belt_im) * 255 - endless_belt_im + + if change_eyebrow: + # eyebrow + eyebrow_mask = np.zeros((img_size, img_size, 3), dtype=np.uint8) + x, y, w, h = cv2.boundingRect(pts137tmp[129:137]) + cv2.rectangle(eyebrow_mask, (x, y), (x + w, y + h), (1, 1, 1), -1) + x, y, w, h = cv2.boundingRect(pts137tmp[121:129]) + cv2.rectangle(eyebrow_mask, (x, y), (x + w, y + h), (1, 1, 1), -1) + kernel_size = int(img_size / 21.) # 21 + if kernel_size % 2 == 0: + kernel_size += 1 + kernel_eyebrow = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (int(kernel_size * 5), int(kernel_size))) + eyebrow_mask = (cv2.dilate(eyebrow_mask, kernel_eyebrow, iterations=1)).astype(np.uint8) + eyebrow_mask = 1 - eyebrow_mask + endless_belt_im = (endless_belt_im.astype(bool) & eyebrow_mask.astype(bool)).astype(np.float32) * 255 + else: + eyebrow_mask = np.zeros((img_size, img_size, 3), dtype=np.uint8) + cv2.fillPoly(eyebrow_mask, pts137tmp[129:137][np.newaxis, :, :], (1, 1, 1)) + cv2.fillPoly(eyebrow_mask, pts137tmp[121:129][np.newaxis, :, :], (1, 1, 1)) + endless_belt_im = (endless_belt_im.astype(bool) | eyebrow_mask.astype(bool)).astype(np.float32) * 255 + + # nose + cv2.fillPoly(endless_belt_im, pts137tmp[64:79][np.newaxis, :, :], (255, 255, 255)) + endless_belt_im[(pts137tmp[133, 1] - img_size // 10):pts137tmp[65, 1], pts137tmp[64, 0]:pts137tmp[78, 0], :] = 255 + + res = np.uint8(endless_belt_im / 255) + res = np.uint8(res * img) + + # add eyelids semantic map + eyelids_mask = np.zeros((img_size, img_size, 3)) + eyelids_mask[eyelids_index] = [0, 255, 0] + res = res + np.uint8(eyelids_mask) + + cv2.circle(res, (pts137tmp[133, 0], pts137tmp[133, 1]), img_size // 50, (255, 0, 0), -1) + cv2.circle(res, (pts137tmp[121, 0], pts137tmp[121, 1]), img_size // 50, (255, 0, 0), -1) + + return res + +def get_transform_singleeye(landmark, output_size, forlabel=False): + dst_size = output_size + left_dis = 0.28 + g_Average_5point_180 = np.array([ + left_dis, 0.5, + 0.5, 0.499, + 1 - left_dis, 0.5, + ]) + pts3_dst = g_Average_5point_180.reshape((3, -1)) * dst_size + pts3_src = np.vstack((landmark[0], landmark[1], landmark[2])) + image_to_face_mat = umeyama(pts3_src, pts3_dst, True)[:2] + + return image_to_face_mat + +def get_transform_mat_full_face_ratio_stylegan(landmark, output_size, ratio=1.0): + dst_size = output_size + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + + #################### + # # ###################### train stylegan hair rate_0.4 size_512 + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + 0.5 + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.6) * ratio + 0.4 + + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_1k_tmp = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_1k_tmp[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_hair(landmark, output_size, ratio=0.5, w_ratio=0.5, h_ratio=0.40): + dst_size = output_size + + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + w_ratio + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.5) * ratio + h_ratio + + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_1k_tmp = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_1k_tmp[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat +def get_transform_mat_sex(landmark, output_size, forlabel=False): + dst_size = output_size + if len(landmark) == 87: + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[31] + landmark[35]) / 2 + right_eye = (landmark[39] + landmark[43]) / 2 + nose = landmark[62] + left_mouth = (landmark[66] + landmark[79]) / 2 + right_mouth = (landmark[72] + landmark[83]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + elif len(landmark) == 137: + + eye_dis = 0.317 + mouth_dis = 0.345 + g_Average_5point_180 = np.array([ + eye_dis, 0.4, + 1 - eye_dis, 0.4, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[96] + landmark[88]) / 2 + right_eye = (landmark[105] + landmark[113]) / 2 + nose = landmark[83] + left_mouth = (landmark[22] + landmark[48]) / 2 + right_mouth = (landmark[56] + landmark[36]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat +def get_transform_mat_hair_ratio(landmark, output_size, ratio=1.0): + dst_size = output_size + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + 0.5 + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.5) * ratio + 0.45 + + # landmarks1k_contours = landmarks_2D_1k_tmp[:312] + # x_l, y_l = np.min(landmarks1k_contours, axis=0) + # x_h, y_h = np.max(landmarks1k_contours, axis=0) + # + # print("x_l: ", x_l, "x_h: ", x_h) + # print("y_l: ", y_l, "y_h: ", y_h) + # + # def draw_landmark(landmark_ori, img): + # landmark_full_int = (landmark_ori.copy() * 512).astype(np.int32) + # img_show = img.copy() + # for pt in landmark_full_int: + # cv2.circle(img_show, (pt[0], pt[1]), 2, (0, 0, 255), 1) + # + # return img_show.astype(np.uint8) + # + # img_temp = np.zeros((512, 512, 3), dtype=np.uint8) + # img_show = draw_landmark(landmarks_2D_1k_tmp, img_temp) + # cv2.imshow("img_show", img_show) + # cv2.waitKey() + + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_1k_tmp = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_1k_tmp[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_hair_ratio_v1(landmark, output_size, ratio=1.0, h_offset=0.5): + dst_size = output_size + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + 0.5 + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.5) * ratio + h_offset + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_1k_tmp = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_1k_tmp[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_full_face_ratio_deeplab(landmark, output_size, ratio=0.3, w_ratio=0.5, h_ratio=0.45): + dst_size = output_size + + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + w_ratio + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.5) * ratio + h_ratio + + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_1k_tmp = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_1k_tmp[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_face_restore(landmark, output_size): + dst_size = output_size + if len(landmark) == 1000: + eye_dis = 0.4 + mouth_dis = 0.42 + g_Average_5point_180 = np.array([ + eye_dis, 0.47, + 1 - eye_dis, 0.47, + 0.5, 0.6, + mouth_dis, 0.71, + 1 - mouth_dis, 0.71 + ]) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + left_eye = (landmark[691] + landmark[723]) / 2 + right_eye = (landmark[792] + landmark[824]) / 2 + nose = landmark[621] + left_mouth = (landmark[467] + landmark[468]) / 2 + right_mouth = (landmark[396] + landmark[508]) / 2 + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + + image_to_face_mat = umeyama(pts5_src, pts5_dst, True)[:2] + + return image_to_face_mat diff --git a/hair_service_sd/utils/model_io.py b/hair_service_sd/utils/model_io.py new file mode 100644 index 0000000..6d3c2cd --- /dev/null +++ b/hair_service_sd/utils/model_io.py @@ -0,0 +1,25 @@ +import os +import torch + +def load_model_by_path(model_save_path, model, gpu_id = None): + if not os.path.exists(model_save_path): return + loc = 'cpu' if gpu_id is None else 'cuda:{}'.format(gpu_id) + pretrained_dict = torch.load(model_save_path, map_location=loc) + if 'state_dict' in pretrained_dict: pretrained_dict = pretrained_dict['state_dict'] + model_dict = model.state_dict() + # pretrained_dict.pop('netG.model.1.weight', '404') + pretrained_dict_new = {k: v for k, v in pretrained_dict.items() + if (k in model_dict and model_dict[k].data.shape == v.data.shape)} + # if 'netG.model.1.weight' not in pretrained_dict_new and 'netG.model.1.weight' in model_dict: + # pretrained_dict_new['netG.model.1.weight'] = model_dict['netG.model.1.weight'] + # pretrained_dict_new['netG.model.1.weight'][:,:12] = pretrained_dict['netG.model.1.weight'] + + model_dict.update(pretrained_dict_new) + model.load_state_dict(model_dict) + print("load model: ", model_save_path) + +def save_model_by_path(model_save_path, model): + save_dir, _ = os.path.split(model_save_path) + if not os.path.isdir(save_dir): os.makedirs(save_dir) + model_dic={k.replace('.module', ''):v for k,v in model.state_dict().items()} + torch.save(model_dic, model_save_path) \ No newline at end of file diff --git a/hair_service_sd/utils/render_offline.py b/hair_service_sd/utils/render_offline.py new file mode 100644 index 0000000..4772170 --- /dev/null +++ b/hair_service_sd/utils/render_offline.py @@ -0,0 +1,326 @@ +import pickle +import sys + +import cv2 +from PIL import Image +from OpenGL.GL import * +from OpenGL.GLU import * +from OpenGL.GLUT import * +import numpy as np + +# define shader code +vertex_code=''' +uniform float scale; +attribute vec2 position; +attribute vec4 color; +varying vec4 v_color; +attribute vec2 TexCoordIn; +varying vec2 TexCoordOut; + +void main() +{ + gl_Position = vec4(position*scale, 0.0, 1.0); + v_color = color; + TexCoordOut = TexCoordIn; + +}''' + +fragment_code=''' +varying vec4 v_color; +varying vec2 TexCoordOut; +uniform sampler2D Texture; +uniform vec2 originPosition; +uniform vec2 targetPosition; +vec2 curveWarp(vec2 textureCoord, vec2 originPosition, vec2 targetPosition, float radius) +{ + vec2 offset = vec2(0.0); + vec2 result = vec2(0.0); + + vec2 direction = targetPosition - originPosition; + + float infect = distance(textureCoord, originPosition)/radius; + + infect = 1.0 - infect; + infect = clamp(infect, 0.0, 1.0); + offset = direction * infect; + + result = textureCoord - offset; + return result; +} + +void main() +{ + vec2 coordinate = vec2(0.0); + float radius = 0.5; + coordinate = curveWarp(TexCoordOut,originPosition,targetPosition,radius); + gl_FragColor = v_color*0.000000000001 + texture2D(Texture, coordinate); + +}''' + + + +#define useful function +def display(): + glClear(GL_COLOR_BUFFER_BIT) + glDrawArrays(GL_TRIANGLE_STRIP, 0, 4) + + glutSwapBuffers() + +def reshape(width,height): + glViewport(0, 0, width, height) + +#step1 init the context +def init(img_path): + image = Image.open(img_path) + glutInit() + glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGB) + glutCreateWindow('Hello world!') + glutReshapeWindow(image.width,image.height) + glutReshapeFunc(reshape) + glutDisplayFunc(display) + return image + +#step 2 +def initShaderProgram(): + program = glCreateProgram() + vertex = glCreateShader(GL_VERTEX_SHADER) + fragment = glCreateShader(GL_FRAGMENT_SHADER) + + # Set shaders source + glShaderSource(vertex, vertex_code) + glShaderSource(fragment, fragment_code) + + # Compile shaders + glCompileShader(vertex) + glCompileShader(fragment) + + fragSuccess = glGetShaderiv(fragment, GL_COMPILE_STATUS) + vertSuccess = glGetShaderiv(vertex, GL_COMPILE_STATUS) + print("vertext shader compile success [%s]" % (vertSuccess,)) + print("fragment shader compile success [%s]" % (fragSuccess,)) + + if vertSuccess == 0: + print(glGetShaderInfoLog(vertex)) + sys.exit(0) + + if fragSuccess == 0: + print(glGetShaderInfoLog(fragment)) + sys.exit(0) + + + glAttachShader(program, vertex) + glAttachShader(program, fragment) + glLinkProgram(program) + linksucc=glGetProgramiv(program, GL_LINK_STATUS) + print("link program success [%s]" % (linksucc,)) + glUseProgram(program) + + return program + +#step 3.1 optional setup texture +def getTextureFromFile(image_file): + #convert file to bytes + image = image_file.transpose(Image.FLIP_TOP_BOTTOM) + image = image.convert("RGBA") + byteImage =np.array(list(image.getdata()), np.uint8) + + #setup texture + texIndex=glGenTextures(1) + glEnable( GL_TEXTURE_2D ) + glBindTexture(GL_TEXTURE_2D,texIndex) + + glPixelStorei(GL_UNPACK_ALIGNMENT,1) + glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP) + glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP) + glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR) + glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR) + #make the texture the default + glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, texIndex, 0) + glTexImage2D(GL_TEXTURE_2D,0,GL_RGBA,image.width,image.height,0,GL_RGBA,GL_UNSIGNED_BYTE,byteImage) + + return texIndex + + +#step 4.1 optional get the image from BufferFrame +def saveImageFromFBO(width, height, output_img_path): + glReadBuffer(GL_COLOR_ATTACHMENT0) + glPixelStorei(GL_PACK_ALIGNMENT, 1) + data = glReadPixels(0, 0, width, height, GL_RGB, GL_UNSIGNED_BYTE) + image = Image.new("RGB", (width, height), (0, 0, 0)) + image.frombytes(data) + image = image.transpose(Image.FLIP_TOP_BOTTOM) + image.save(output_img_path) + + +#step 4 optional if need to OSR generate a BufferFrame +def setupSelfDefineFBO(program, image, data, output_img_path): + fbWidth, fbHeight = image.width, image.height + + # Setup framebuffer + framebuffer = glGenFramebuffers(1) + glBindFramebuffer(GL_FRAMEBUFFER, framebuffer) + + # # Setup colorbuffer + # colorbuffer = glGenRenderbuffers(1) + # glBindRenderbuffer(GL_RENDERBUFFER, colorbuffer) + # glRenderbufferStorage(GL_RENDERBUFFER, GL_RGBA, fbWidth, fbHeight) + # glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_RENDERBUFFER, colorbuffer) + + # Setup depthbuffer + depthbuffer = glGenRenderbuffers (1) + glBindRenderbuffer (GL_RENDERBUFFER,depthbuffer) + glRenderbufferStorage (GL_RENDERBUFFER, GL_DEPTH_COMPONENT, image.width, image.height) + glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_DEPTH_ATTACHMENT, GL_RENDERBUFFER, depthbuffer) + + #first init VBO, then other parameters + buffer = glGenBuffers(1) # Request a buffer slot from GPU + glBindBuffer(GL_ARRAY_BUFFER, buffer) # Make this buffer the default one + glBufferData(GL_ARRAY_BUFFER, data.nbytes, data, GL_DYNAMIC_DRAW) # Upload data + + # Create texture to render to + # glBufferData(GL_FRAMEBUFFER, data.nbytes, data, GL_DYNAMIC_DRAW) + loc = glGetAttribLocation(program, "position") #get the index of the attribute in program + glEnableVertexAttribArray(loc) #allow this attribute decide by index can be use + stride = data.strides[0] #define how to read buffer + offset = ctypes.c_void_p(0) #define the offset where the data begin in buffer + glVertexAttribPointer(loc, 2, GL_FLOAT, False, stride, offset) + offset = ctypes.c_void_p(data.dtype["position"].itemsize) + loc = glGetAttribLocation(program, "color") + glEnableVertexAttribArray(loc) + glVertexAttribPointer(loc, 4, GL_FLOAT, False, stride, offset) + + #setup other parameters + loc = glGetUniformLocation(program, "scale") + glUniform1f(loc, 1.0) + + # originPosition = glGetUniformLocation(program, "originPosition") + # glUniform2f(originPosition, 0.5, 0.5) + # + # targetPosition = glGetUniformLocation(program, "targetPosition") + # glUniform2f(targetPosition, 0.47, 0.47) + # # glUniform2f(targetPosition, 0.5, 0.5) + + # following code to bind uniform texture if needed + aTexture = getTextureFromFile(image) + + glViewport(0, 0, fbWidth, fbHeight) + + glActiveTexture(GL_TEXTURE0) + glBindTexture(GL_TEXTURE_2D, aTexture) + loc = glGetUniformLocation(program, "Texture") + glUniform1i(loc, 0) + + loc = glGetAttribLocation(program, "TexCoordIn") + glEnableVertexAttribArray(loc) + offset=ctypes.c_void_p(data.dtype["color"].itemsize+8) + glVertexAttribPointer(loc, 2, GL_FLOAT, False, stride, offset) + + status = glCheckFramebufferStatus (GL_FRAMEBUFFER) + if status != GL_FRAMEBUFFER_COMPLETE: + print( "Error in framebuffer activation") + + glDrawArrays(GL_TRIANGLE_STRIP, 0, 4) + saveImageFromFBO(fbWidth, fbHeight, output_img_path) + + glBindFramebuffer(GL_FRAMEBUFFER, GL_NONE) + glDeleteTextures([aTexture]) + glDeleteFramebuffers(1, [framebuffer]) + + print('save image from FBO success') + +def pt_in_img(pt, img_w, img_h): + + if pt[0] < 0 or pt[1] < 0 or pt[0] >= img_w or pt[1] >= img_h: + return False + else: + return True + +def demo_test(): + input_img_path = "../test_data/pics/female003.jpg" + output_img_path = "../tmp.png" + + ###### define vertex and color array + + # data = np.zeros(4, dtype=[("position", np.float32, 2), ("color", np.float32, 4), ("textureCoord", np.float32, 2)]) + # data['color'] = [(1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1)] + # # data['position'] = [(-1, -1), (-1, 1), (1, -1), (1, 1)] + # data['position'] = [(-1, -1), (-1, 1), (0, -1), (0, 1)] + # # data['textureCoord'] = [(0, 0), (0, 1), (1, 0), (1, 1)] + # data['textureCoord'] = [(0, 0), (0, 1), (1, 0), (0.5, 1)] + + data = np.zeros(8, dtype=[("position", np.float32, 2), ("color", np.float32, 4), ("textureCoord", np.float32, 2)]) + data['color'] = [(1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1)] + data['position'] = [(-1, -1), (-1, +1), (0, -1), (0, +1), (0, -1), (0, +1), (+1, -1), (+1, +1)] + # data['textureCoord'] = [(0, 0), (0, 1), (0.5, 0), (0.5, 1), (0.5, 0), (0.5, 1), (1, 0), (1, 1)] + mid_x = 0.6 + data['textureCoord'] = [(0, 0), (0, 1), (mid_x, 0), (mid_x, 1), (mid_x, 0), (mid_x, 1), (1, 0), (1, 1)] + + # data = np.zeros(3, dtype=[("position", np.float32, 2), ("color", np.float32, 4), ("textureCoord", np.float32, 2)]) + # data['color'] = [(1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1)] + # data['position'] = [(-1, -1), (-1, +1), (1, -1)] + # # data['position'] = [(-1, -1), (1, 1), (1, -1)] + # data['textureCoord'] = [(0, 0), (0, 1), (0.5, 0)] + # # data['textureCoord'] = [(0, 0), (1, 1), (1, 0)] + + # data = np.zeros(6, dtype=[("position", np.float32, 2), ("color", np.float32, 4), ("textureCoord", np.float32, 2)]) + # data['color'] = [(1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1)] + # data['position'] = [(-1, -1), (-1, +1), (0, -1), (0, +1), (1, -1), (1, 1)] + # # data['textureCoord'] = [(0, 0), (0, 1), (0.5, 0), (0.5, 1), (1, 0), (1, 1)] + # data['textureCoord'] = [(0, 0), (0, 1), (0.5, 0), (0.5, 1), (1, 0), (1, 1)] + # mid_x = 0.4 + # data['textureCoord'] = [(0, 0), (0, 1), (mid_x, 0), (mid_x, 1), (1, 0), (1, 1)] + + + image = init(input_img_path) + program = initShaderProgram() + setupSelfDefineFBO(program, image, data, output_img_path) + +def demo_warp_pt137(): + input_img_path = "../test_data/pics/female003.jpg" + output_img_path = "../tmp.png" + + import pickle + with open("../test_render.pkl", "rb") as fp: + info = pickle.load(fp) + vertice = info["vertice"] + dst_vertice = info["dst_vertice"] + faces = info["faces"] + color_list = [] + position_list = [] + textureCoord_list = [] + img = cv2.imread(input_img_path) + img_h, img_w, _ = img.shape + for i in range(len(faces)): + # print("index: ", faces[i]) + pt1 = vertice[faces[i][0]] + pt2 = vertice[faces[i][1]] + pt3 = vertice[faces[i][2]] + dst_pt1 = dst_vertice[faces[i][0]] + dst_pt2 = dst_vertice[faces[i][1]] + dst_pt3 = dst_vertice[faces[i][2]] + # if pt1[0] != dst_pt1[0]: + # print("use warp!") + if pt_in_img(pt1, img_w, img_h) and pt_in_img(pt2, img_w, img_h) and pt_in_img(pt3, img_w, img_h): + color_list.append((1, 1, 0, 1)) + color_list.append((1, 1, 0, 1)) + color_list.append((1, 1, 0, 1)) + position_list.append((dst_pt1[0] * 2 / img_w - 1, dst_pt1[1] * 2 / img_h - 1)) + position_list.append((dst_pt2[0] * 2 / img_w - 1, dst_pt2[1] * 2 / img_h - 1)) + position_list.append((dst_pt3[0] * 2 / img_w - 1, dst_pt3[1] * 2 / img_h - 1)) + textureCoord_list.append((pt1[0] / img_w, pt1[1] / img_h)) + textureCoord_list.append((pt2[0] / img_w, pt2[1] / img_h)) + textureCoord_list.append((pt3[0] / img_w, pt3[1] / img_h)) + + data = np.zeros(len(color_list), + dtype=[("position", np.float32, 2), ("color", np.float32, 4), ("textureCoord", np.float32, 2)]) + data['color'] = color_list + data['position'] = position_list + data['textureCoord'] = textureCoord_list + + image = init(input_img_path) + program = initShaderProgram() + setupSelfDefineFBO(program, image, data, output_img_path) + +if __name__ == '__main__': + demo_test() + # demo_warp_pt137() \ No newline at end of file diff --git a/hair_service_sd/utils/torch_utils.py b/hair_service_sd/utils/torch_utils.py new file mode 100644 index 0000000..e069792 --- /dev/null +++ b/hair_service_sd/utils/torch_utils.py @@ -0,0 +1,202 @@ +import math +import os +import time +from copy import deepcopy + +import torch +import torch.backends.cudnn as cudnn +import torch.nn as nn +import torch.nn.functional as F +import torchvision.models as models + + +def init_seeds(seed=0): + torch.manual_seed(seed) + + # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html + if seed == 0: # slower, more reproducible + cudnn.deterministic = True + cudnn.benchmark = False + else: # faster, less reproducible + cudnn.deterministic = False + cudnn.benchmark = True + + +def select_device(device='', apex=False, batch_size=None): + # device = 'cpu' or '0' or '0,1,2,3' + cpu_request = device.lower() == 'cpu' + if device and not cpu_request: # if device requested other than 'cpu' + os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable + assert torch.cuda.is_available(), 'CUDA unavailable, invalid device %s requested' % device # check availablity + + cuda = False if cpu_request else torch.cuda.is_available() + if cuda: + c = 1024 ** 2 # bytes to MB + ng = torch.cuda.device_count() + if ng > 1 and batch_size: # check that batch_size is compatible with device_count + assert batch_size % ng == 0, 'batch-size %g not multiple of GPU count %g' % (batch_size, ng) + x = [torch.cuda.get_device_properties(i) for i in range(ng)] + s = 'Using CUDA ' + ('Apex ' if apex else '') # apex for mixed precision https://github.com/NVIDIA/apex + for i in range(0, ng): + if i == 1: + s = ' ' * len(s) + print("%sdevice%g _CudaDeviceProperties(name='%s', total_memory=%dMB)" % + (s, i, x[i].name, x[i].total_memory / c)) + else: + print('Using CPU') + + print('') # skip a line + return torch.device('cuda:0' if cuda else 'cpu') + + +def time_synchronized(): + torch.cuda.synchronize() if torch.cuda.is_available() else None + return time.time() + + +def initialize_weights(model): + for m in model.modules(): + t = type(m) + if t is nn.Conv2d: + pass # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + elif t is nn.BatchNorm2d: + m.eps = 1e-4 + m.momentum = 0.03 + elif t in [nn.LeakyReLU, nn.ReLU, nn.ReLU6]: + m.inplace = True + + +def find_modules(model, mclass=nn.Conv2d): + # finds layer indices matching module class 'mclass' + return [i for i, m in enumerate(model.module_list) if isinstance(m, mclass)] + + +def fuse_conv_and_bn(conv, bn): + # https://tehnokv.com/posts/fusing-batchnorm-and-conv/ + with torch.no_grad(): + # init + fusedconv = torch.nn.Conv2d(conv.in_channels, + conv.out_channels, + kernel_size=conv.kernel_size, + stride=conv.stride, + padding=conv.padding, + bias=True) + + # prepare filters + w_conv = conv.weight.clone().view(conv.out_channels, -1) + w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var))) + fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.size())) + + # prepare spatial bias + if conv.bias is not None: + b_conv = conv.bias + else: + b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) + b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps)) + fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn) + + return fusedconv + + +def model_info(model, verbose=False): + # Plots a line-by-line description of a PyTorch model + n_p = sum(x.numel() for x in model.parameters()) # number parameters + n_g = sum(x.numel() for x in model.parameters() if x.requires_grad) # number gradients + if verbose: + print('%5s %40s %9s %12s %20s %10s %10s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma')) + for i, (name, p) in enumerate(model.named_parameters()): + name = name.replace('module_list.', '') + print('%5g %40s %9s %12g %20s %10.3g %10.3g' % + (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std())) + + try: # FLOPS + from thop import profile + macs, _ = profile(model, inputs=(torch.zeros(1, 3, 480, 640),), verbose=False) + fs = ', %.1f GFLOPS' % (macs / 1E9 * 2) + except: + fs = '' + + print('Model Summary: %g layers, %g parameters, %g gradients%s' % (len(list(model.parameters())), n_p, n_g, fs)) + + +def load_classifier(name='resnet101', n=2): + # Loads a pretrained model reshaped to n-class output + model = models.__dict__[name](pretrained=True) + + # Display model properties + input_size = [3, 224, 224] + input_space = 'RGB' + input_range = [0, 1] + mean = [0.485, 0.456, 0.406] + std = [0.229, 0.224, 0.225] + for x in [input_size, input_space, input_range, mean, std]: + print(x + ' =', eval(x)) + + # Reshape output to n classes + filters = model.fc.weight.shape[1] + model.fc.bias = torch.nn.Parameter(torch.zeros(n), requires_grad=True) + model.fc.weight = torch.nn.Parameter(torch.zeros(n, filters), requires_grad=True) + model.fc.out_features = n + return model + + +def scale_img(img, ratio=1.0, same_shape=False): # img(16,3,256,416), r=ratio + # scales img(bs,3,y,x) by ratio + h, w = img.shape[2:] + s = (int(h * ratio), int(w * ratio)) # new size + img = F.interpolate(img, size=s, mode='bilinear', align_corners=False) # resize + if not same_shape: # pad/crop img + gs = 32 # (pixels) grid size + h, w = [math.ceil(x * ratio / gs) * gs for x in (h, w)] + return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447) # value = imagenet mean + + +class ModelEMA: + """ Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models + Keep a moving average of everything in the model state_dict (parameters and buffers). + This is intended to allow functionality like + https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage + A smoothed version of the weights is necessary for some training schemes to perform well. + E.g. Google's hyper-params for training MNASNet, MobileNet-V3, EfficientNet, etc that use + RMSprop with a short 2.4-3 epoch decay period and slow LR decay rate of .96-.99 requires EMA + smoothing of weights to match results. Pay attention to the decay constant you are using + relative to your update count per epoch. + To keep EMA from using GPU resources, set device='cpu'. This will save a bit of memory but + disable validation of the EMA weights. Validation will have to be done manually in a separate + process, or after the training stops converging. + This class is sensitive where it is initialized in the sequence of model init, + GPU assignment and distributed training wrappers. + I've tested with the sequence in my own train.py for torch.DataParallel, apex.DDP, and single-GPU. + """ + + def __init__(self, model, decay=0.9999, device=''): + # make a copy of the model for accumulating moving average of weights + self.ema = deepcopy(model) + self.ema.eval() + self.updates = 0 # number of EMA updates + self.decay = lambda x: decay * (1 - math.exp(-x / 2000)) # decay exponential ramp (to help early epochs) + self.device = device # perform ema on different device from model if set + if device: + self.ema.to(device=device) + for p in self.ema.parameters(): + p.requires_grad_(False) + + def update(self, model): + self.updates += 1 + d = self.decay(self.updates) + with torch.no_grad(): + if type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel): + msd, esd = model.module.state_dict(), self.ema.module.state_dict() + else: + msd, esd = model.state_dict(), self.ema.state_dict() + + for k, v in esd.items(): + if v.dtype.is_floating_point: + v *= d + v += (1. - d) * msd[k].detach() + + def update_attr(self, model): + # Assign attributes (which may change during training) + for k in model.__dict__.keys(): + if not k.startswith('_'): + setattr(self.ema, k, getattr(model, k)) diff --git a/hair_service_sd/utils/umeyama.py b/hair_service_sd/utils/umeyama.py new file mode 100644 index 0000000..aad6235 --- /dev/null +++ b/hair_service_sd/utils/umeyama.py @@ -0,0 +1,71 @@ +import numpy as np + +def umeyama(src, dst, estimate_scale): + """Estimate N-D similarity transformation with or without scaling. + Parameters + ---------- + src : (M, N) array + Source coordinates. + dst : (M, N) array + Destination coordinates. + estimate_scale : bool + Whether to estimate scaling factor. + Returns + ------- + T : (N + 1, N + 1) + The homogeneous similarity transformation matrix. The matrix contains + NaN values only if the problem is not well-conditioned. + References + ---------- + .. [1] "Least-squares estimation of transformation parameters between two + point patterns", Shinji Umeyama, PAMI 1991, DOI: 10.1109/34.88573 + """ + + num = src.shape[0] + dim = src.shape[1] + + # Compute mean of src and dst. + src_mean = src.mean(axis=0) + dst_mean = dst.mean(axis=0) + + # Subtract mean from src and dst. + src_demean = src - src_mean + dst_demean = dst - dst_mean + + # Eq. (38). + A = np.dot(dst_demean.T, src_demean) / num + + # Eq. (39). + d = np.ones((dim,), dtype=np.double) + if np.linalg.det(A) < 0: + d[dim - 1] = -1 + + T = np.eye(dim + 1, dtype=np.double) + + U, S, V = np.linalg.svd(A) + + # Eq. (40) and (43). + rank = np.linalg.matrix_rank(A) + if rank == 0: + return np.nan * T + elif rank == dim - 1: + if np.linalg.det(U) * np.linalg.det(V) > 0: + T[:dim, :dim] = np.dot(U, V) + else: + s = d[dim - 1] + d[dim - 1] = -1 + T[:dim, :dim] = np.dot(U, np.dot(np.diag(d), V)) + d[dim - 1] = s + else: + T[:dim, :dim] = np.dot(U, np.dot(np.diag(d), V.T)) + + if estimate_scale: + # Eq. (41) and (42). + scale = 1.0 / src_demean.var(axis=0).sum() * np.dot(S, d) + else: + scale = 1.0 + + T[:dim, dim] = dst_mean - scale * np.dot(T[:dim, :dim], src_mean.T) + T[:dim, :dim] *= scale + + return T \ No newline at end of file diff --git a/hair_service_sd/utils/util.py b/hair_service_sd/utils/util.py new file mode 100644 index 0000000..84f535f --- /dev/null +++ b/hair_service_sd/utils/util.py @@ -0,0 +1,287 @@ +import os +import cv2 +import torch +import logging +import numpy as np +# from utils.config import CONFIG +# import torch.distributed as dist + +def mkdirs(paths): + """create empty directories if they don't exist + Parameters: + paths (str list) -- a list of directory paths + """ + if isinstance(paths, list) and not isinstance(paths, str): + for path in paths: + os.makedirs(path) + else: + os.makedirs(paths) + +def make_dir(target_dir): + """ + Create dir if not exists + """ + if not os.path.exists(target_dir): + os.makedirs(target_dir) + + +def print_network(model, name): + """ + Print out the network information + """ + logger = logging.getLogger("Logger") + num_params = 0 + for p in model.parameters(): + num_params += p.numel() + + logger.info(model) + logger.info(name) + logger.info("Number of parameters: {}".format(num_params)) + + +def update_lr(lr, optimizer): + """ + update learning rates + """ + for param_group in optimizer.param_groups: + param_group['lr'] = lr + + +def warmup_lr(init_lr, step, iter_num): + """ + Warm up learning rate + """ + return step/iter_num*init_lr + + +def add_prefix_state_dict(state_dict, prefix="module"): + """ + add prefix from the key of pretrained state dict for Data-Parallel + """ + new_state_dict = {} + first_state_name = list(state_dict.keys())[0] + if not first_state_name.startswith(prefix): + for key, value in state_dict.items(): + new_state_dict[prefix+"."+key] = state_dict[key].float() + else: + for key, value in state_dict.items(): + new_state_dict[key] = state_dict[key].float() + return new_state_dict + + +def remove_prefix_state_dict(state_dict, prefix="module"): + """ + remove prefix from the key of pretrained state dict for Data-Parallel + """ + new_state_dict = {} + first_state_name = list(state_dict.keys())[0] + if not first_state_name.startswith(prefix): + for key, value in state_dict.items(): + new_state_dict[key] = state_dict[key].float() + else: + for key, value in state_dict.items(): + new_state_dict[key[len(prefix)+1:]] = state_dict[key].float() + return new_state_dict + +# +# def load_imagenet_pretrain(model, checkpoint_file): +# """ +# Load imagenet pretrained resnet +# Add zeros channel to the first convolution layer +# Since we have the spectral normalization, we need to do a little more +# """ +# checkpoint = torch.load(checkpoint_file, map_location = lambda storage, loc: storage.cuda(CONFIG.gpu)) +# state_dict = remove_prefix_state_dict(checkpoint['state_dict']) +# for key, value in state_dict.items(): +# state_dict[key] = state_dict[key].float() +# +# logger = logging.getLogger("Logger") +# logger.debug("Imagenet pretrained keys:") +# logger.debug(state_dict.keys()) +# logger.debug("Generator keys:") +# logger.debug(model.module.encoder.state_dict().keys()) +# logger.debug("Intersection keys:") +# logger.debug(set(model.module.encoder.state_dict().keys())&set(state_dict.keys())) +# +# weight_u = state_dict["conv1.module.weight_u"] +# weight_v = state_dict["conv1.module.weight_v"] +# weight_bar = state_dict["conv1.module.weight_bar"] +# +# logger.debug("weight_v: {}".format(weight_v)) +# logger.debug("weight_bar: {}".format(weight_bar.view(32, -1))) +# logger.debug("sigma: {}".format(weight_u.dot(weight_bar.view(32, -1).mv(weight_v)))) +# +# new_weight_v = torch.zeros(6, 3, 3).cuda() +# new_weight_bar = torch.zeros(32, 6, 3, 3).cuda() +# +# new_weight_v[:3, :, :].copy_(weight_v.view(3, 3, 3)) +# new_weight_bar[:, :3, :, :].copy_(weight_bar) +# +# logger.debug("new weight_v: {}".format(new_weight_v.view(-1))) +# logger.debug("new weight_bar: {}".format(new_weight_bar.view(32, -1))) +# logger.debug("new sigma: {}".format(weight_u.dot(new_weight_bar.view(32, -1).mv(new_weight_v.view(-1))))) +# +# state_dict["conv1.module.weight_v"] = new_weight_v.view(-1) +# state_dict["conv1.module.weight_bar"] = new_weight_bar +# +# model.module.encoder.load_state_dict(state_dict, strict=False) + + +def load_VGG_pretrain(model, checkpoint_file): + """ + Load imagenet pretrained resnet + Add zeros channel to the first convolution layer + Since we have the spectral normalization, we need to do a little more + """ + checkpoint = torch.load(checkpoint_file, map_location = lambda storage, loc: storage.cuda()) + backbone_state_dict = remove_prefix_state_dict(checkpoint['state_dict']) + + model.module.encoder.load_state_dict(backbone_state_dict, strict=False) + + +def get_unknown_tensor(trimap): + """ + get 1-channel unknown area tensor from the 3-channel/1-channel trimap tensor + """ + # if CONFIG.model.trimap_channel == 3: + weight = trimap[:, 1:2, :, :].float() + # else: + # weight = trimap.eq(1).float() + return weight + + +def get_gaborfilter(angles): + """ + generate gabor filter as the conv kernel + :param angles: number of different angles + """ + gabor_filter = [] + for angle in range(angles): + gabor_filter.append(cv2.getGaborKernel(ksize=(5,5), sigma=0.5, theta=angle*np.pi/8, lambd=5, gamma=0.5)) + gabor_filter = np.array(gabor_filter) + gabor_filter = np.expand_dims(gabor_filter, axis=1) + return gabor_filter.astype(np.float32) + + +def get_gradfilter(): + """ + generate gradient filter as the conv kernel + """ + grad_filter = [] + grad_filter.append([[-1, -2, -1], [0, 0, 0], [1, 2, 1]]) + grad_filter.append([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]) + grad_filter = np.array(grad_filter) + grad_filter = np.expand_dims(grad_filter, axis=1) + return grad_filter.astype(np.float32) + + +# def reduce_tensor_dict(tensor_dict, mode='mean'): +# """ +# average tensor dict over different GPUs +# """ +# for key, tensor in tensor_dict.items(): +# if tensor is not None: +# tensor_dict[key] = reduce_tensor(tensor, mode) +# return tensor_dict +# +# +# def reduce_tensor(tensor, mode='mean'): +# """ +# average tensor over different GPUs +# """ +# rt = tensor.clone() +# dist.all_reduce(rt, op=dist.ReduceOp.SUM) +# if mode == 'mean': +# rt /= CONFIG.world_size +# elif mode == 'sum': +# pass +# else: +# raise NotImplementedError("reduce mode can only be 'mean' or 'sum'") +# return rt + +def make_color_wheel(): + # from https://github.com/JiahuiYu/generative_inpainting/blob/master/inpaint_ops.py + RY, YG, GC, CB, BM, MR = (15, 6, 4, 11, 13, 6) + ncols = RY + YG + GC + CB + BM + MR + colorwheel = np.zeros([ncols, 3]) + col = 0 + # RY + colorwheel[0:RY, 0] = 255 + colorwheel[0:RY, 1] = np.transpose(np.floor(255*np.arange(0, RY) / RY)) + col += RY + # YG + colorwheel[col:col+YG, 0] = 255 - np.transpose(np.floor(255*np.arange(0, YG) / YG)) + colorwheel[col:col+YG, 1] = 255 + col += YG + # GC + colorwheel[col:col+GC, 1] = 255 + colorwheel[col:col+GC, 2] = np.transpose(np.floor(255*np.arange(0, GC) / GC)) + col += GC + # CB + colorwheel[col:col+CB, 1] = 255 - np.transpose(np.floor(255*np.arange(0, CB) / CB)) + colorwheel[col:col+CB, 2] = 255 + col += CB + # BM + colorwheel[col:col+BM, 2] = 255 + colorwheel[col:col+BM, 0] = np.transpose(np.floor(255*np.arange(0, BM) / BM)) + col += + BM + # MR + colorwheel[col:col+MR, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, MR) / MR)) + colorwheel[col:col+MR, 0] = 255 + return colorwheel + + +COLORWHEEL = make_color_wheel() + +def compute_color(u,v): + # from https://github.com/JiahuiYu/generative_inpainting/blob/master/inpaint_ops.py + h, w = u.shape + img = np.zeros([h, w, 3]) + nanIdx = np.isnan(u) | np.isnan(v) + u[nanIdx] = 0 + v[nanIdx] = 0 + colorwheel = COLORWHEEL + # colorwheel = make_color_wheel() + ncols = np.size(colorwheel, 0) + rad = np.sqrt(u**2+v**2) + a = np.arctan2(-v, -u) / np.pi + fk = (a+1) / 2 * (ncols - 1) + 1 + k0 = np.floor(fk).astype(int) + k1 = k0 + 1 + k1[k1 == ncols+1] = 1 + f = fk - k0 + for i in range(np.size(colorwheel,1)): + tmp = colorwheel[:, i] + col0 = tmp[k0-1] / 255 + col1 = tmp[k1-1] / 255 + col = (1-f) * col0 + f * col1 + idx = rad <= 1 + col[idx] = 1-rad[idx]*(1-col[idx]) + notidx = np.logical_not(idx) + col[notidx] *= 0.75 + img[:, :, i] = np.uint8(np.floor(255 * col*(1-nanIdx))) + return img + +def flow_to_image(flow): + # part from https://github.com/JiahuiYu/generative_inpainting/blob/master/inpaint_ops.py + maxrad = -1 + u = flow[0, :, :] + v = flow[1, :, :] + rad = np.sqrt(u ** 2 + v ** 2) + maxrad = max(maxrad, np.max(rad)) + u = u/(maxrad + np.finfo(float).eps) + v = v/(maxrad + np.finfo(float).eps) + img = compute_color(u, v) + + return img + + +if __name__ == "__main__": + import networks + logging.basicConfig(level=logging.DEBUG, format='[%(asctime)s] %(levelname)s: %(message)s', + datefmt='%m-%d %H:%M:%S') + G = networks.get_generator().cuda() + # load_imagenet_pretrain(G, CONFIG.model.imagenet_pretrain_path) + x = torch.randn(4,3,512,512).cuda() + y = torch.randn(4,3,512,512).cuda() + z = G(x, y) diff --git a/hair_service_sd/utils/utils.py b/hair_service_sd/utils/utils.py new file mode 100644 index 0000000..c33f41f --- /dev/null +++ b/hair_service_sd/utils/utils.py @@ -0,0 +1,1207 @@ +import glob +import math +import os +import random +import shutil +import subprocess +import time +from copy import copy +from pathlib import Path +from sys import platform + +import cv2 +import matplotlib +import matplotlib.pyplot as plt +import numpy as np +import torch +import torch.nn as nn +import torchvision +import yaml +from scipy.signal import butter, filtfilt +from tqdm import tqdm + +from . import torch_utils #  torch_utils, google_utils + +# Set printoptions +torch.set_printoptions(linewidth=320, precision=5, profile='long') +np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format}) # format short g, %precision=5 +matplotlib.rc('font', **{'size': 11}) + +# Prevent OpenCV from multithreading (to use PyTorch DataLoader) +cv2.setNumThreads(0) + + +def init_seeds(seed=0): + random.seed(seed) + np.random.seed(seed) + torch_utils.init_seeds(seed=seed) + + +def check_git_status(): + # Suggest 'git pull' if repo is out of date + if platform in ['linux', 'darwin']: + s = subprocess.check_output('if [ -d .git ]; then git fetch && git status -uno; fi', shell=True).decode('utf-8') + if 'Your branch is behind' in s: + print(s[s.find('Your branch is behind'):s.find('\n\n')] + '\n') + + +def check_img_size(img_size, s=32): + # Verify img_size is a multiple of stride s + new_size = make_divisible(img_size, s) # ceil gs-multiple + if new_size != img_size: + print('WARNING: --img-size %g must be multiple of max stride %g, updating to %g' % (img_size, s, new_size)) + return new_size + + +def check_anchors(dataset, model, thr=4.0, imgsz=640): + # Check anchor fit to data, recompute if necessary + print('\nAnalyzing anchors... ', end='') + m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1] # Detect() + shapes = imgsz * dataset.shapes / dataset.shapes.max(1, keepdims=True) + scale = np.random.uniform(0.9, 1.1, size=(shapes.shape[0], 1)) # augment scale + wh = torch.tensor(np.concatenate([l[:, 3:5] * s for s, l in zip(shapes * scale, dataset.labels)])).float() # wh + + def metric(k): # compute metric + r = wh[:, None] / k[None] + x = torch.min(r, 1. / r).min(2)[0] # ratio metric + best = x.max(1)[0] # best_x + return (best > 1. / thr).float().mean() #  best possible recall + + bpr = metric(m.anchor_grid.clone().cpu().view(-1, 2)) + print('Best Possible Recall (BPR) = %.4f' % bpr, end='') + if bpr < 0.99: # threshold to recompute + print('. Attempting to generate improved anchors, please wait...' % bpr) + na = m.anchor_grid.numel() // 2 # number of anchors + new_anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False) + new_bpr = metric(new_anchors.reshape(-1, 2)) + if new_bpr > bpr: # replace anchors + new_anchors = torch.tensor(new_anchors, device=m.anchors.device).type_as(m.anchors) + m.anchor_grid[:] = new_anchors.clone().view_as(m.anchor_grid) # for inference + m.anchors[:] = new_anchors.clone().view_as(m.anchors) / m.stride.to(m.anchors.device).view(-1, 1, 1) # loss + check_anchor_order(m) + print('New anchors saved to model. Update model *.yaml to use these anchors in the future.') + else: + print('Original anchors better than new anchors. Proceeding with original anchors.') + print('') # newline + + +def check_anchor_order(m): + # Check anchor order against stride order for YOLOv5 Detect() module m, and correct if necessary + a = m.anchor_grid.prod(-1).view(-1) # anchor area + da = a[-1] - a[0] # delta a + ds = m.stride[-1] - m.stride[0] # delta s + if da.sign() != ds.sign(): # same order + m.anchors[:] = m.anchors.flip(0) + m.anchor_grid[:] = m.anchor_grid.flip(0) + + +def check_file(file): + # Searches for file if not found locally + if os.path.isfile(file): + return file + else: + files = glob.glob('./**/' + file, recursive=True) # find file + assert len(files), 'File Not Found: %s' % file # assert file was found + return files[0] # return first file if multiple found + + +def make_divisible(x, divisor): + # Returns x evenly divisble by divisor + return math.ceil(x / divisor) * divisor + + +def labels_to_class_weights(labels, nc=80): + # Get class weights (inverse frequency) from training labels + if labels[0] is None: # no labels loaded + return torch.Tensor() + + labels = np.concatenate(labels, 0) # labels.shape = (866643, 5) for COCO + classes = labels[:, 0].astype(np.int) # labels = [class xywh] + weights = np.bincount(classes, minlength=nc) # occurences per class + + # Prepend gridpoint count (for uCE trianing) + # gpi = ((320 / 32 * np.array([1, 2, 4])) ** 2 * 3).sum() # gridpoints per image + # weights = np.hstack([gpi * len(labels) - weights.sum() * 9, weights * 9]) ** 0.5 # prepend gridpoints to start + + weights[weights == 0] = 1 # replace empty bins with 1 + weights = 1 / weights # number of targets per class + weights /= weights.sum() # normalize + return torch.from_numpy(weights) + + +def labels_to_image_weights(labels, nc=80, class_weights=np.ones(80)): + # Produces image weights based on class mAPs + n = len(labels) + class_counts = np.array([np.bincount(labels[i][:, 0].astype(np.int), minlength=nc) for i in range(n)]) + image_weights = (class_weights.reshape(1, nc) * class_counts).sum(1) + # index = random.choices(range(n), weights=image_weights, k=1) # weight image sample + return image_weights + + +def coco80_to_coco91_class(): # converts 80-index (val2014) to 91-index (paper) + # https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/ + # a = np.loadtxt('data/coco.names', dtype='str', delimiter='\n') + # b = np.loadtxt('data/coco_paper.names', dtype='str', delimiter='\n') + # x1 = [list(a[i] == b).index(True) + 1 for i in range(80)] # darknet to coco + # x2 = [list(b[i] == a).index(True) if any(b[i] == a) else None for i in range(91)] # coco to darknet + x = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34, + 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, + 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90] + return x + + +def xyxy2xywh(x): + # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right + y = torch.zeros_like(x) if isinstance(x, torch.Tensor) else np.zeros_like(x) + y[:, 0] = (x[:, 0] + x[:, 2]) / 2 # x center + y[:, 1] = (x[:, 1] + x[:, 3]) / 2 # y center + y[:, 2] = x[:, 2] - x[:, 0] # width + y[:, 3] = x[:, 3] - x[:, 1] # height + return y + + +def xywh2xyxy(x): + # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right + y = torch.zeros_like(x) if isinstance(x, torch.Tensor) else np.zeros_like(x) + y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x + y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y + y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x + y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y + return y + + +def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None): + # Rescale coords (xyxy) from img1_shape to img0_shape + if ratio_pad is None: # calculate from img0_shape + gain = max(img1_shape) / max(img0_shape) # gain = old / new + pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding + else: + gain = ratio_pad[0][0] + pad = ratio_pad[1] + + coords[:, [0, 2]] -= pad[0] # x padding + coords[:, [1, 3]] -= pad[1] # y padding + coords[:, :4] /= gain + clip_coords(coords, img0_shape) + return coords + + +def clip_coords(boxes, img_shape): + # Clip bounding xyxy bounding boxes to image shape (height, width) + boxes[:, 0].clamp_(0, img_shape[1]) # x1 + boxes[:, 1].clamp_(0, img_shape[0]) # y1 + boxes[:, 2].clamp_(0, img_shape[1]) # x2 + boxes[:, 3].clamp_(0, img_shape[0]) # y2 + + +def ap_per_class(tp, conf, pred_cls, target_cls): + """ Compute the average precision, given the recall and precision curves. + Source: https://github.com/rafaelpadilla/Object-Detection-Metrics. + # Arguments + tp: True positives (nparray, nx1 or nx10). + conf: Objectness value from 0-1 (nparray). + pred_cls: Predicted object classes (nparray). + target_cls: True object classes (nparray). + # Returns + The average precision as computed in py-faster-rcnn. + """ + + # Sort by objectness + i = np.argsort(-conf) + tp, conf, pred_cls = tp[i], conf[i], pred_cls[i] + + # Find unique classes + unique_classes = np.unique(target_cls) + + # Create Precision-Recall curve and compute AP for each class + pr_score = 0.1 # score to evaluate P and R https://github.com/ultralytics/yolov3/issues/898 + s = [unique_classes.shape[0], tp.shape[1]] # number class, number iou thresholds (i.e. 10 for mAP0.5...0.95) + ap, p, r = np.zeros(s), np.zeros(s), np.zeros(s) + for ci, c in enumerate(unique_classes): + i = pred_cls == c + n_gt = (target_cls == c).sum() # Number of ground truth objects + n_p = i.sum() # Number of predicted objects + + if n_p == 0 or n_gt == 0: + continue + else: + # Accumulate FPs and TPs + fpc = (1 - tp[i]).cumsum(0) + tpc = tp[i].cumsum(0) + + # Recall + recall = tpc / (n_gt + 1e-16) # recall curve + r[ci] = np.interp(-pr_score, -conf[i], recall[:, 0]) # r at pr_score, negative x, xp because xp decreases + + # Precision + precision = tpc / (tpc + fpc) # precision curve + p[ci] = np.interp(-pr_score, -conf[i], precision[:, 0]) # p at pr_score + + # AP from recall-precision curve + for j in range(tp.shape[1]): + ap[ci, j] = compute_ap(recall[:, j], precision[:, j]) + + # Plot + # fig, ax = plt.subplots(1, 1, figsize=(5, 5)) + # ax.plot(recall, precision) + # ax.set_xlabel('Recall') + # ax.set_ylabel('Precision') + # ax.set_xlim(0, 1.01) + # ax.set_ylim(0, 1.01) + # fig.tight_layout() + # fig.savefig('PR_curve.png', dpi=300) + + # Compute F1 score (harmonic mean of precision and recall) + f1 = 2 * p * r / (p + r + 1e-16) + + return p, r, ap, f1, unique_classes.astype('int32') + + +def compute_ap(recall, precision): + """ Compute the average precision, given the recall and precision curves. + Source: https://github.com/rbgirshick/py-faster-rcnn. + # Arguments + recall: The recall curve (list). + precision: The precision curve (list). + # Returns + The average precision as computed in py-faster-rcnn. + """ + + # Append sentinel values to beginning and end + mrec = np.concatenate(([0.], recall, [min(recall[-1] + 1E-3, 1.)])) + mpre = np.concatenate(([0.], precision, [0.])) + + # Compute the precision envelope + mpre = np.flip(np.maximum.accumulate(np.flip(mpre))) + + # Integrate area under curve + method = 'interp' # methods: 'continuous', 'interp' + if method == 'interp': + x = np.linspace(0, 1, 101) # 101-point interp (COCO) + ap = np.trapz(np.interp(x, mrec, mpre), x) # integrate + else: # 'continuous' + i = np.where(mrec[1:] != mrec[:-1])[0] # points where x axis (recall) changes + ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1]) # area under curve + + return ap + + +def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False): + # Returns the IoU of box1 to box2. box1 is 4, box2 is nx4 + box2 = box2.t() + + # Get the coordinates of bounding boxes + if x1y1x2y2: # x1, y1, x2, y2 = box1 + b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3] + b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3] + else: # transform from xywh to xyxy + b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2 + b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2 + b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2 + b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2 + + # Intersection area + inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \ + (torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0) + + # Union Area + w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + union = (w1 * h1 + 1e-16) + w2 * h2 - inter + + iou = inter / union # iou + if GIoU or DIoU or CIoU: + cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1) # convex (smallest enclosing box) width + ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1) # convex height + if GIoU: # Generalized IoU https://arxiv.org/pdf/1902.09630.pdf + c_area = cw * ch + 1e-16 # convex area + return iou - (c_area - union) / c_area # GIoU + if DIoU or CIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1 + # convex diagonal squared + c2 = cw ** 2 + ch ** 2 + 1e-16 + # centerpoint distance squared + rho2 = ((b2_x1 + b2_x2) - (b1_x1 + b1_x2)) ** 2 / 4 + ((b2_y1 + b2_y2) - (b1_y1 + b1_y2)) ** 2 / 4 + if DIoU: + return iou - rho2 / c2 # DIoU + elif CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47 + v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2) + with torch.no_grad(): + alpha = v / (1 - iou + v) + return iou - (rho2 / c2 + v * alpha) # CIoU + + return iou + + +def box_iou(box1, box2): + # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py + """ + Return intersection-over-union (Jaccard index) of boxes. + Both sets of boxes are expected to be in (x1, y1, x2, y2) format. + Arguments: + box1 (Tensor[N, 4]) + box2 (Tensor[M, 4]) + Returns: + iou (Tensor[N, M]): the NxM matrix containing the pairwise + IoU values for every element in boxes1 and boxes2 + """ + + def box_area(box): + # box = 4xn + return (box[2] - box[0]) * (box[3] - box[1]) + + area1 = box_area(box1.t()) + area2 = box_area(box2.t()) + + # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2) + inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2) + return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter) + + +def wh_iou(wh1, wh2): + # Returns the nxm IoU matrix. wh1 is nx2, wh2 is mx2 + wh1 = wh1[:, None] # [N,1,2] + wh2 = wh2[None] # [1,M,2] + inter = torch.min(wh1, wh2).prod(2) # [N,M] + return inter / (wh1.prod(2) + wh2.prod(2) - inter) # iou = inter / (area1 + area2 - inter) + + +class FocalLoss(nn.Module): + # Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5) + def __init__(self, loss_fcn, gamma=1.5, alpha=0.25): + super(FocalLoss, self).__init__() + self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss() + self.gamma = gamma + self.alpha = alpha + self.reduction = loss_fcn.reduction + self.loss_fcn.reduction = 'none' # required to apply FL to each element + + def forward(self, pred, true): + loss = self.loss_fcn(pred, true) + # p_t = torch.exp(-loss) + # loss *= self.alpha * (1.000001 - p_t) ** self.gamma # non-zero power for gradient stability + + # TF implementation https://github.com/tensorflow/addons/blob/v0.7.1/tensorflow_addons/losses/focal_loss.py + pred_prob = torch.sigmoid(pred) # prob from logits + p_t = true * pred_prob + (1 - true) * (1 - pred_prob) + alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha) + modulating_factor = (1.0 - p_t) ** self.gamma + loss *= alpha_factor * modulating_factor + + if self.reduction == 'mean': + return loss.mean() + elif self.reduction == 'sum': + return loss.sum() + else: # 'none' + return loss + + +def smooth_BCE(eps=0.1): # https://github.com/ultralytics/yolov3/issues/238#issuecomment-598028441 + # return positive, negative label smoothing BCE targets + return 1.0 - 0.5 * eps, 0.5 * eps + + +class BCEBlurWithLogitsLoss(nn.Module): + # BCEwithLogitLoss() with reduced missing label effects. + def __init__(self, alpha=0.05): + super(BCEBlurWithLogitsLoss, self).__init__() + self.loss_fcn = nn.BCEWithLogitsLoss(reduction='none') # must be nn.BCEWithLogitsLoss() + self.alpha = alpha + + def forward(self, pred, true): + loss = self.loss_fcn(pred, true) + pred = torch.sigmoid(pred) # prob from logits + dx = pred - true # reduce only missing label effects + # dx = (pred - true).abs() # reduce missing label and false label effects + alpha_factor = 1 - torch.exp((dx - 1) / (self.alpha + 1e-4)) + loss *= alpha_factor + return loss.mean() + + +def compute_loss(p, targets, model): # predictions, targets, model + ft = torch.cuda.FloatTensor if p[0].is_cuda else torch.Tensor + lcls, lbox, lobj = ft([0]), ft([0]), ft([0]) + tcls, tbox, indices, anchors = build_targets(p, targets, model) # targets + h = model.hyp # hyperparameters + red = 'mean' # Loss reduction (sum or mean) + + # Define criteria + BCEcls = nn.BCEWithLogitsLoss(pos_weight=ft([h['cls_pw']]), reduction=red) + BCEobj = nn.BCEWithLogitsLoss(pos_weight=ft([h['obj_pw']]), reduction=red) + + # class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3 + cp, cn = smooth_BCE(eps=0.0) + + # focal loss + g = h['fl_gamma'] # focal loss gamma + if g > 0: + BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g) + + # per output + nt = 0 # targets + for i, pi in enumerate(p): # layer index, layer predictions + b, a, gj, gi = indices[i] # image, anchor, gridy, gridx + tobj = torch.zeros_like(pi[..., 0]) # target obj + + nb = b.shape[0] # number of targets + if nb: + nt += nb # cumulative targets + ps = pi[b, a, gj, gi] # prediction subset corresponding to targets + + # GIoU + pxy = ps[:, :2].sigmoid() * 2. - 0.5 + pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i] + pbox = torch.cat((pxy, pwh), 1) # predicted box + giou = bbox_iou(pbox.t(), tbox[i], x1y1x2y2=False, GIoU=True) # giou(prediction, target) + lbox += (1.0 - giou).sum() if red == 'sum' else (1.0 - giou).mean() # giou loss + + # Obj + tobj[b, a, gj, gi] = (1.0 - model.gr) + model.gr * giou.detach().clamp(0).type(tobj.dtype) # giou ratio + + # Class + if model.nc > 1: # cls loss (only if multiple classes) + t = torch.full_like(ps[:, 5:], cn) # targets + t[range(nb), tcls[i]] = cp + lcls += BCEcls(ps[:, 5:], t) # BCE + + # Append targets to text file + # with open('targets.txt', 'a') as file: + # [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)] + + lobj += BCEobj(pi[..., 4], tobj) # obj loss + + lbox *= h['giou'] + lobj *= h['obj'] + lcls *= h['cls'] + bs = tobj.shape[0] # batch size + if red == 'sum': + g = 3.0 # loss gain + lobj *= g / bs + if nt: + lcls *= g / nt / model.nc + lbox *= g / nt + + loss = lbox + lobj + lcls + return loss * bs, torch.cat((lbox, lobj, lcls, loss)).detach() + + +def build_targets(p, targets, model): + # Build targets for compute_loss(), input targets(image,class,x,y,w,h) + det = model.module.model[-1] if type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel) \ + else model.model[-1] # Detect() module + na, nt = det.na, targets.shape[0] # number of anchors, targets + tcls, tbox, indices, anch = [], [], [], [] + gain = torch.ones(6, device=targets.device) # normalized to gridspace gain + off = torch.tensor([[1, 0], [0, 1], [-1, 0], [0, -1]], device=targets.device).float() # overlap offsets + at = torch.arange(na).view(na, 1).repeat(1, nt) # anchor tensor, same as .repeat_interleave(nt) + + style = 'rect4' + for i in range(det.nl): + anchors = det.anchors[i] + gain[2:] = torch.tensor(p[i].shape)[[3, 2, 3, 2]] # xyxy gain + + # Match targets to anchors + a, t, offsets = [], targets * gain, 0 + if nt: + r = t[None, :, 4:6] / anchors[:, None] # wh ratio + j = torch.max(r, 1. / r).max(2)[0] < model.hyp['anchor_t'] # compare + # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n) = wh_iou(anchors(3,2), gwh(n,2)) + a, t = at[j], t.repeat(na, 1, 1)[j] # filter + + # overlaps + gxy = t[:, 2:4] # grid xy + z = torch.zeros_like(gxy) + if style == 'rect2': + g = 0.2 # offset + j, k = ((gxy % 1. < g) & (gxy > 1.)).T + a, t = torch.cat((a, a[j], a[k]), 0), torch.cat((t, t[j], t[k]), 0) + offsets = torch.cat((z, z[j] + off[0], z[k] + off[1]), 0) * g + + elif style == 'rect4': + g = 0.5 # offset + j, k = ((gxy % 1. < g) & (gxy > 1.)).T + l, m = ((gxy % 1. > (1 - g)) & (gxy < (gain[[2, 3]] - 1.))).T + a, t = torch.cat((a, a[j], a[k], a[l], a[m]), 0), torch.cat((t, t[j], t[k], t[l], t[m]), 0) + offsets = torch.cat((z, z[j] + off[0], z[k] + off[1], z[l] + off[2], z[m] + off[3]), 0) * g + + # Define + b, c = t[:, :2].long().T # image, class + gxy = t[:, 2:4] # grid xy + gwh = t[:, 4:6] # grid wh + gij = (gxy - offsets).long() + gi, gj = gij.T # grid xy indices + + # Append + indices.append((b, a, gj, gi)) # image, anchor, grid indices + tbox.append(torch.cat((gxy - gij, gwh), 1)) # box + anch.append(anchors[a]) # anchors + tcls.append(c) # class + + return tcls, tbox, indices, anch + + +def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, merge=False, classes=None, agnostic=False): + """Performs Non-Maximum Suppression (NMS) on inference results + + Returns: + detections with shape: nx6 (x1, y1, x2, y2, conf, cls) + """ + if prediction.dtype is torch.float16: + prediction = prediction.float() # to FP32 + + nc = prediction[0].shape[1] - 5 # number of classes + xc = prediction[..., 4] > conf_thres # candidates + + # Settings + min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height + max_det = 300 # maximum number of detections per image + time_limit = 10.0 # seconds to quit after + redundant = True # require redundant detections + multi_label = nc > 1 # multiple labels per box (adds 0.5ms/img) + + t = time.time() + output = [None] * prediction.shape[0] + for xi, x in enumerate(prediction): # image index, image inference + # Apply constraints + # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height + x = x[xc[xi]] # confidence + + # If none remain process next image + if not x.shape[0]: + continue + + # Compute conf + x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf + + # Box (center x, center y, width, height) to (x1, y1, x2, y2) + box = xywh2xyxy(x[:, :4]) + + # Detections matrix nx6 (xyxy, conf, cls) + if multi_label: + i, j = (x[:, 5:] > conf_thres).nonzero().t() + x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1) + else: # best class only + conf, j = x[:, 5:].max(1, keepdim=True) + x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres] + + # Filter by class + if classes: + x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)] + + # Apply finite constraint + # if not torch.isfinite(x).all(): + # x = x[torch.isfinite(x).all(1)] + + # If none remain process next image + n = x.shape[0] # number of boxes + if not n: + continue + + # Sort by confidence + # x = x[x[:, 4].argsort(descending=True)] + + # Batched NMS + c = x[:, 5:6] * (0 if agnostic else max_wh) # classes + boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores + i = torchvision.ops.boxes.nms(boxes, scores, iou_thres) + if i.shape[0] > max_det: # limit detections + i = i[:max_det] + if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean) + try: # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4) + iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix + weights = iou * scores[None] # box weights + x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes + if redundant: + i = i[iou.sum(1) > 1] # require redundancy + except: # possible CUDA error https://github.com/ultralytics/yolov3/issues/1139 + print(x, i, x.shape, i.shape) + pass + + output[xi] = x[i] + if (time.time() - t) > time_limit: + break # time limit exceeded + + return output + + +def strip_optimizer(f='weights/best.pt'): # from utils.utils import *; strip_optimizer() + # Strip optimizer from *.pt files for lighter files (reduced by 1/2 size) + x = torch.load(f, map_location=torch.device('cpu')) + x['optimizer'] = None + x['model'].half() # to FP16 + torch.save(x, f) + print('Optimizer stripped from %s' % f) + + +def create_pretrained(f='weights/best.pt', s='weights/pretrained.pt'): # from utils.utils import *; create_pretrained() + # create pretrained checkpoint 's' from 'f' (create_pretrained(x, x) for x in glob.glob('./*.pt')) + device = torch.device('cpu') + x = torch.load(s, map_location=device) + + x['optimizer'] = None + x['training_results'] = None + x['epoch'] = -1 + x['model'].half() # to FP16 + for p in x['model'].parameters(): + p.requires_grad = True + torch.save(x, s) + print('%s saved as pretrained checkpoint %s' % (f, s)) + + +def coco_class_count(path='../coco/labels/train2014/'): + # Histogram of occurrences per class + nc = 80 # number classes + x = np.zeros(nc, dtype='int32') + files = sorted(glob.glob('%s/*.*' % path)) + for i, file in enumerate(files): + labels = np.loadtxt(file, dtype=np.float32).reshape(-1, 5) + x += np.bincount(labels[:, 0].astype('int32'), minlength=nc) + print(i, len(files)) + + +def coco_only_people(path='../coco/labels/train2017/'): # from utils.utils import *; coco_only_people() + # Find images with only people + files = sorted(glob.glob('%s/*.*' % path)) + for i, file in enumerate(files): + labels = np.loadtxt(file, dtype=np.float32).reshape(-1, 5) + if all(labels[:, 0] == 0): + print(labels.shape[0], file) + + +def crop_images_random(path='../images/', scale=0.50): # from utils.utils import *; crop_images_random() + # crops images into random squares up to scale fraction + # WARNING: overwrites images! + for file in tqdm(sorted(glob.glob('%s/*.*' % path))): + img = cv2.imread(file) # BGR + if img is not None: + h, w = img.shape[:2] + + # create random mask + a = 30 # minimum size (pixels) + mask_h = random.randint(a, int(max(a, h * scale))) # mask height + mask_w = mask_h # mask width + + # box + xmin = max(0, random.randint(0, w) - mask_w // 2) + ymin = max(0, random.randint(0, h) - mask_h // 2) + xmax = min(w, xmin + mask_w) + ymax = min(h, ymin + mask_h) + + # apply random color mask + cv2.imwrite(file, img[ymin:ymax, xmin:xmax]) + + +def coco_single_class_labels(path='../coco/labels/train2014/', label_class=43): + # Makes single-class coco datasets. from utils.utils import *; coco_single_class_labels() + if os.path.exists('new/'): + shutil.rmtree('new/') # delete output folder + os.makedirs('new/') # make new output folder + os.makedirs('new/labels/') + os.makedirs('new/images/') + for file in tqdm(sorted(glob.glob('%s/*.*' % path))): + with open(file, 'r') as f: + labels = np.array([x.split() for x in f.read().splitlines()], dtype=np.float32) + i = labels[:, 0] == label_class + if any(i): + img_file = file.replace('labels', 'images').replace('txt', 'jpg') + labels[:, 0] = 0 # reset class to 0 + with open('new/images.txt', 'a') as f: # add image to dataset list + f.write(img_file + '\n') + with open('new/labels/' + Path(file).name, 'a') as f: # write label + for l in labels[i]: + f.write('%g %.6f %.6f %.6f %.6f\n' % tuple(l)) + shutil.copyfile(src=img_file, dst='new/images/' + Path(file).name.replace('txt', 'jpg')) # copy images + + +def kmean_anchors(path='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=1000, verbose=True): + """ Creates kmeans-evolved anchors from training dataset + + Arguments: + path: path to dataset *.yaml, or a loaded dataset + n: number of anchors + img_size: image size used for training + thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0 + gen: generations to evolve anchors using genetic algorithm + + Return: + k: kmeans evolved anchors + + Usage: + from utils.utils import *; _ = kmean_anchors() + """ + thr = 1. / thr + + def metric(k, wh): # compute metrics + r = wh[:, None] / k[None] + x = torch.min(r, 1. / r).min(2)[0] # ratio metric + # x = wh_iou(wh, torch.tensor(k)) # iou metric + return x, x.max(1)[0] # x, best_x + + def fitness(k): # mutation fitness + _, best = metric(torch.tensor(k, dtype=torch.float32), wh) + return (best * (best > thr).float()).mean() # fitness + + def print_results(k): + k = k[np.argsort(k.prod(1))] # sort small to large + x, best = metric(k, wh0) + bpr, aat = (best > thr).float().mean(), (x > thr).float().mean() * n # best possible recall, anch > thr + print('thr=%.2f: %.4f best possible recall, %.2f anchors past thr' % (thr, bpr, aat)) + print('n=%g, img_size=%s, metric_all=%.3f/%.3f-mean/best, past_thr=%.3f-mean: ' % + (n, img_size, x.mean(), best.mean(), x[x > thr].mean()), end='') + for i, x in enumerate(k): + print('%i,%i' % (round(x[0]), round(x[1])), end=', ' if i < len(k) - 1 else '\n') # use in *.cfg + return k + + if isinstance(path, str): # *.yaml file + with open(path) as f: + data_dict = yaml.load(f, Loader=yaml.FullLoader) # model dict + from utils.datasets import LoadImagesAndLabels + dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True) + else: + dataset = path # dataset + + # Get label wh + shapes = img_size * dataset.shapes / dataset.shapes.max(1, keepdims=True) + wh0 = np.concatenate([l[:, 3:5] * s for s, l in zip(shapes, dataset.labels)]) # wh + + # Filter + i = (wh0 < 4.0).any(1).sum() + if i: + print('WARNING: Extremely small objects found. ' + '%g of %g labels are < 4 pixels in width or height.' % (i, len(wh0))) + wh = wh0[(wh0 >= 4.0).any(1)] # filter > 2 pixels + + # Kmeans calculation + from scipy.cluster.vq import kmeans + print('Running kmeans for %g anchors on %g points...' % (n, len(wh))) + s = wh.std(0) # sigmas for whitening + k, dist = kmeans(wh / s, n, iter=30) # points, mean distance + k *= s + wh = torch.tensor(wh, dtype=torch.float32) # filtered + wh0 = torch.tensor(wh0, dtype=torch.float32) # unflitered + k = print_results(k) + + # Plot + # k, d = [None] * 20, [None] * 20 + # for i in tqdm(range(1, 21)): + # k[i-1], d[i-1] = kmeans(wh / s, i) # points, mean distance + # fig, ax = plt.subplots(1, 2, figsize=(14, 7)) + # ax = ax.ravel() + # ax[0].plot(np.arange(1, 21), np.array(d) ** 2, marker='.') + # fig, ax = plt.subplots(1, 2, figsize=(14, 7)) # plot wh + # ax[0].hist(wh[wh[:, 0]<100, 0],400) + # ax[1].hist(wh[wh[:, 1]<100, 1],400) + # fig.tight_layout() + # fig.savefig('wh.png', dpi=200) + + # Evolve + npr = np.random + f, sh, mp, s = fitness(k), k.shape, 0.9, 0.1 # fitness, generations, mutation prob, sigma + pbar = tqdm(range(gen), desc='Evolving anchors with Genetic Algorithm') # progress bar + for _ in pbar: + v = np.ones(sh) + while (v == 1).all(): # mutate until a change occurs (prevent duplicates) + v = ((npr.random(sh) < mp) * npr.random() * npr.randn(*sh) * s + 1).clip(0.3, 3.0) + kg = (k.copy() * v).clip(min=2.0) + fg = fitness(kg) + if fg > f: + f, k = fg, kg.copy() + pbar.desc = 'Evolving anchors with Genetic Algorithm: fitness = %.4f' % f + if verbose: + print_results(k) + + return print_results(k) + + +def print_mutation(hyp, results, bucket=''): + # Print mutation results to evolve.txt (for use with train.py --evolve) + a = '%10s' * len(hyp) % tuple(hyp.keys()) # hyperparam keys + b = '%10.3g' * len(hyp) % tuple(hyp.values()) # hyperparam values + c = '%10.4g' * len(results) % results # results (P, R, mAP, F1, test_loss) + print('\n%s\n%s\nEvolved fitness: %s\n' % (a, b, c)) + + if bucket: + os.system('gsutil cp gs://%s/evolve.txt .' % bucket) # download evolve.txt + + with open('evolve.txt', 'a') as f: # append result + f.write(c + b + '\n') + x = np.unique(np.loadtxt('evolve.txt', ndmin=2), axis=0) # load unique rows + np.savetxt('evolve.txt', x[np.argsort(-fitness(x))], '%10.3g') # save sort by fitness + + if bucket: + os.system('gsutil cp evolve.txt gs://%s' % bucket) # upload evolve.txt + + +def apply_classifier(x, model, img, im0): + # applies a second stage classifier to yolo outputs + im0 = [im0] if isinstance(im0, np.ndarray) else im0 + for i, d in enumerate(x): # per image + if d is not None and len(d): + d = d.clone() + + # Reshape and pad cutouts + b = xyxy2xywh(d[:, :4]) # boxes + b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # rectangle to square + b[:, 2:] = b[:, 2:] * 1.3 + 30 # pad + d[:, :4] = xywh2xyxy(b).long() + + # Rescale boxes from img_size to im0 size + scale_coords(img.shape[2:], d[:, :4], im0[i].shape) + + # Classes + pred_cls1 = d[:, 5].long() + ims = [] + for j, a in enumerate(d): # per item + cutout = im0[i][int(a[1]):int(a[3]), int(a[0]):int(a[2])] + im = cv2.resize(cutout, (224, 224)) # BGR + # cv2.imwrite('test%i.jpg' % j, cutout) + + im = im[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416 + im = np.ascontiguousarray(im, dtype=np.float32) # uint8 to float32 + im /= 255.0 # 0 - 255 to 0.0 - 1.0 + ims.append(im) + + pred_cls2 = model(torch.Tensor(ims).to(d.device)).argmax(1) # classifier prediction + x[i] = x[i][pred_cls1 == pred_cls2] # retain matching class detections + + return x + + +def fitness(x): + # Returns fitness (for use with results.txt or evolve.txt) + w = [0.0, 0.0, 0.1, 0.9] # weights for [P, R, mAP@0.5, mAP@0.5:0.95] + return (x[:, :4] * w).sum(1) + + +def output_to_target(output, width, height): + """ + Convert a YOLO model output to target format + [batch_id, class_id, x, y, w, h, conf] + """ + if isinstance(output, torch.Tensor): + output = output.cpu().numpy() + + targets = [] + for i, o in enumerate(output): + if o is not None: + for pred in o: + box = pred[:4] + w = (box[2] - box[0]) / width + h = (box[3] - box[1]) / height + x = box[0] / width + w / 2 + y = box[1] / height + h / 2 + conf = pred[4] + cls = int(pred[5]) + + targets.append([i, cls, x, y, w, h, conf]) + + return np.array(targets) + + +# Plotting functions --------------------------------------------------------------------------------------------------- +def butter_lowpass_filtfilt(data, cutoff=1500, fs=50000, order=5): + # https://stackoverflow.com/questions/28536191/how-to-filter-smooth-with-scipy-numpy + def butter_lowpass(cutoff, fs, order): + nyq = 0.5 * fs + normal_cutoff = cutoff / nyq + b, a = butter(order, normal_cutoff, btype='low', analog=False) + return b, a + + b, a = butter_lowpass(cutoff, fs, order=order) + return filtfilt(b, a, data) # forward-backward filter + + +def plot_one_box(x, img, color=None, label=None, line_thickness=None): + # Plots one bounding box on image img + tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line/font thickness + color = color or [random.randint(0, 255) for _ in range(3)] + c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3])) + cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA) + if label: + tf = max(tl - 1, 1) # font thickness + t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0] + c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3 + cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled + cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA) + + +def plot_wh_methods(): # from utils.utils import *; plot_wh_methods() + # Compares the two methods for width-height anchor multiplication + # https://github.com/ultralytics/yolov3/issues/168 + x = np.arange(-4.0, 4.0, .1) + ya = np.exp(x) + yb = torch.sigmoid(torch.from_numpy(x)).numpy() * 2 + + fig = plt.figure(figsize=(6, 3), dpi=150) + plt.plot(x, ya, '.-', label='yolo method') + plt.plot(x, yb ** 2, '.-', label='^2 power method') + plt.plot(x, yb ** 2.5, '.-', label='^2.5 power method') + plt.xlim(left=-4, right=4) + plt.ylim(bottom=0, top=6) + plt.xlabel('input') + plt.ylabel('output') + plt.legend() + fig.tight_layout() + fig.savefig('comparison.png', dpi=200) + + +def plot_images(images, targets, paths=None, fname='images.jpg', names=None, max_size=640, max_subplots=16): + tl = 3 # line thickness + tf = max(tl - 1, 1) # font thickness + if os.path.isfile(fname): # do not overwrite + return None + + if isinstance(images, torch.Tensor): + images = images.cpu().float().numpy() + + if isinstance(targets, torch.Tensor): + targets = targets.cpu().numpy() + + # un-normalise + if np.max(images[0]) <= 1: + images *= 255 + + bs, _, h, w = images.shape # batch size, _, height, width + bs = min(bs, max_subplots) # limit plot images + ns = np.ceil(bs ** 0.5) # number of subplots (square) + + # Check if we should resize + scale_factor = max_size / max(h, w) + if scale_factor < 1: + h = math.ceil(scale_factor * h) + w = math.ceil(scale_factor * w) + + # Empty array for output + mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8) + + # Fix class - colour map + prop_cycle = plt.rcParams['axes.prop_cycle'] + # https://stackoverflow.com/questions/51350872/python-from-color-name-to-rgb + hex2rgb = lambda h: tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4)) + color_lut = [hex2rgb(h) for h in prop_cycle.by_key()['color']] + + for i, img in enumerate(images): + if i == max_subplots: # if last batch has fewer images than we expect + break + + block_x = int(w * (i // ns)) + block_y = int(h * (i % ns)) + + img = img.transpose(1, 2, 0) + if scale_factor < 1: + img = cv2.resize(img, (w, h)) + + mosaic[block_y:block_y + h, block_x:block_x + w, :] = img + if len(targets) > 0: + image_targets = targets[targets[:, 0] == i] + boxes = xywh2xyxy(image_targets[:, 2:6]).T + classes = image_targets[:, 1].astype('int') + gt = image_targets.shape[1] == 6 # ground truth if no conf column + conf = None if gt else image_targets[:, 6] # check for confidence presence (gt vs pred) + + boxes[[0, 2]] *= w + boxes[[0, 2]] += block_x + boxes[[1, 3]] *= h + boxes[[1, 3]] += block_y + for j, box in enumerate(boxes.T): + cls = int(classes[j]) + color = color_lut[cls % len(color_lut)] + cls = names[cls] if names else cls + if gt or conf[j] > 0.3: # 0.3 conf thresh + label = '%s' % cls if gt else '%s %.1f' % (cls, conf[j]) + plot_one_box(box, mosaic, label=label, color=color, line_thickness=tl) + + # Draw image filename labels + if paths is not None: + label = os.path.basename(paths[i])[:40] # trim to 40 char + t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0] + cv2.putText(mosaic, label, (block_x + 5, block_y + t_size[1] + 5), 0, tl / 3, [220, 220, 220], thickness=tf, + lineType=cv2.LINE_AA) + + # Image border + cv2.rectangle(mosaic, (block_x, block_y), (block_x + w, block_y + h), (255, 255, 255), thickness=3) + + if fname is not None: + mosaic = cv2.resize(mosaic, (int(ns * w * 0.5), int(ns * h * 0.5)), interpolation=cv2.INTER_AREA) + cv2.imwrite(fname, cv2.cvtColor(mosaic, cv2.COLOR_BGR2RGB)) + + return mosaic + + +def plot_lr_scheduler(optimizer, scheduler, epochs=300): + # Plot LR simulating training for full epochs + optimizer, scheduler = copy(optimizer), copy(scheduler) # do not modify originals + y = [] + for _ in range(epochs): + scheduler.step() + y.append(optimizer.param_groups[0]['lr']) + plt.plot(y, '.-', label='LR') + plt.xlabel('epoch') + plt.ylabel('LR') + plt.grid() + plt.xlim(0, epochs) + plt.ylim(0) + plt.tight_layout() + plt.savefig('LR.png', dpi=200) + + +def plot_test_txt(): # from utils.utils import *; plot_test() + # Plot test.txt histograms + x = np.loadtxt('test.txt', dtype=np.float32) + box = xyxy2xywh(x[:, :4]) + cx, cy = box[:, 0], box[:, 1] + + fig, ax = plt.subplots(1, 1, figsize=(6, 6), tight_layout=True) + ax.hist2d(cx, cy, bins=600, cmax=10, cmin=0) + ax.set_aspect('equal') + plt.savefig('hist2d.png', dpi=300) + + fig, ax = plt.subplots(1, 2, figsize=(12, 6), tight_layout=True) + ax[0].hist(cx, bins=600) + ax[1].hist(cy, bins=600) + plt.savefig('hist1d.png', dpi=200) + + +def plot_targets_txt(): # from utils.utils import *; plot_targets_txt() + # Plot targets.txt histograms + x = np.loadtxt('targets.txt', dtype=np.float32).T + s = ['x targets', 'y targets', 'width targets', 'height targets'] + fig, ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True) + ax = ax.ravel() + for i in range(4): + ax[i].hist(x[i], bins=100, label='%.3g +/- %.3g' % (x[i].mean(), x[i].std())) + ax[i].legend() + ax[i].set_title(s[i]) + plt.savefig('targets.jpg', dpi=200) + + +def plot_study_txt(f='study.txt', x=None): # from utils.utils import *; plot_study_txt() + # Plot study.txt generated by test.py + fig, ax = plt.subplots(2, 4, figsize=(10, 6), tight_layout=True) + ax = ax.ravel() + + fig2, ax2 = plt.subplots(1, 1, figsize=(8, 4), tight_layout=True) + for f in ['coco_study/study_coco_yolov5%s.txt' % x for x in ['s', 'm', 'l', 'x']]: + y = np.loadtxt(f, dtype=np.float32, usecols=[0, 1, 2, 3, 7, 8, 9], ndmin=2).T + x = np.arange(y.shape[1]) if x is None else np.array(x) + s = ['P', 'R', 'mAP@.5', 'mAP@.5:.95', 't_inference (ms/img)', 't_NMS (ms/img)', 't_total (ms/img)'] + for i in range(7): + ax[i].plot(x, y[i], '.-', linewidth=2, markersize=8) + ax[i].set_title(s[i]) + + j = y[3].argmax() + 1 + ax2.plot(y[6, :j], y[3, :j] * 1E2, '.-', linewidth=2, markersize=8, + label=Path(f).stem.replace('study_coco_', '').replace('yolo', 'YOLO')) + + ax2.plot(1E3 / np.array([209, 140, 97, 58, 35, 18]), [33.5, 39.1, 42.5, 45.9, 49., 50.5], + 'k.-', linewidth=2, markersize=8, alpha=.25, label='EfficientDet') + + ax2.grid() + ax2.set_xlim(0, 30) + ax2.set_ylim(28, 50) + ax2.set_yticks(np.arange(30, 55, 5)) + ax2.set_xlabel('GPU Speed (ms/img)') + ax2.set_ylabel('COCO AP val') + ax2.legend(loc='lower right') + plt.savefig('study_mAP_latency.png', dpi=300) + plt.savefig(f.replace('.txt', '.png'), dpi=200) + + +def plot_labels(labels): + # plot dataset labels + c, b = labels[:, 0], labels[:, 1:].transpose() # classees, boxes + + def hist2d(x, y, n=100): + xedges, yedges = np.linspace(x.min(), x.max(), n), np.linspace(y.min(), y.max(), n) + hist, xedges, yedges = np.histogram2d(x, y, (xedges, yedges)) + xidx = np.clip(np.digitize(x, xedges) - 1, 0, hist.shape[0] - 1) + yidx = np.clip(np.digitize(y, yedges) - 1, 0, hist.shape[1] - 1) + return np.log(hist[xidx, yidx]) + + fig, ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True) + ax = ax.ravel() + ax[0].hist(c, bins=int(c.max() + 1)) + ax[0].set_xlabel('classes') + ax[1].scatter(b[0], b[1], c=hist2d(b[0], b[1], 90), cmap='jet') + ax[1].set_xlabel('x') + ax[1].set_ylabel('y') + ax[2].scatter(b[2], b[3], c=hist2d(b[2], b[3], 90), cmap='jet') + ax[2].set_xlabel('width') + ax[2].set_ylabel('height') + plt.savefig('labels.png', dpi=200) + plt.close() + + +def plot_evolution_results(hyp): # from utils.utils import *; plot_evolution_results(hyp) + # Plot hyperparameter evolution results in evolve.txt + x = np.loadtxt('evolve.txt', ndmin=2) + f = fitness(x) + # weights = (f - f.min()) ** 2 # for weighted results + plt.figure(figsize=(12, 10), tight_layout=True) + matplotlib.rc('font', **{'size': 8}) + for i, (k, v) in enumerate(hyp.items()): + y = x[:, i + 7] + # mu = (y * weights).sum() / weights.sum() # best weighted result + mu = y[f.argmax()] # best single result + plt.subplot(4, 5, i + 1) + plt.plot(mu, f.max(), 'o', markersize=10) + plt.plot(y, f, '.') + plt.title('%s = %.3g' % (k, mu), fontdict={'size': 9}) # limit to 40 characters + print('%15s: %.3g' % (k, mu)) + plt.savefig('evolve.png', dpi=200) + + +def plot_results_overlay(start=0, stop=0): # from utils.utils import *; plot_results_overlay() + # Plot training 'results*.txt', overlaying train and val losses + s = ['train', 'train', 'train', 'Precision', 'mAP@0.5', 'val', 'val', 'val', 'Recall', 'mAP@0.5:0.95'] # legends + t = ['GIoU', 'Objectness', 'Classification', 'P-R', 'mAP-F1'] # titles + for f in sorted(glob.glob('results*.txt') + glob.glob('../../Downloads/results*.txt')): + results = np.loadtxt(f, usecols=[2, 3, 4, 8, 9, 12, 13, 14, 10, 11], ndmin=2).T + n = results.shape[1] # number of rows + x = range(start, min(stop, n) if stop else n) + fig, ax = plt.subplots(1, 5, figsize=(14, 3.5), tight_layout=True) + ax = ax.ravel() + for i in range(5): + for j in [i, i + 5]: + y = results[j, x] + ax[i].plot(x, y, marker='.', label=s[j]) + # y_smooth = butter_lowpass_filtfilt(y) + # ax[i].plot(x, np.gradient(y_smooth), marker='.', label=s[j]) + + ax[i].set_title(t[i]) + ax[i].legend() + ax[i].set_ylabel(f) if i == 0 else None # add filename + fig.savefig(f.replace('.txt', '.png'), dpi=200) + + +def plot_results(start=0, stop=0, bucket='', id=(), labels=()): # from utils.utils import *; plot_results() + # Plot training 'results*.txt' as seen in https://github.com/ultralytics/yolov5#reproduce-our-training + fig, ax = plt.subplots(2, 5, figsize=(12, 6)) + ax = ax.ravel() + s = ['GIoU', 'Objectness', 'Classification', 'Precision', 'Recall', + 'val GIoU', 'val Objectness', 'val Classification', 'mAP@0.5', 'mAP@0.5:0.95'] + if bucket: + os.system('rm -rf storage.googleapis.com') + files = ['https://storage.googleapis.com/%s/results%g.txt' % (bucket, x) for x in id] + else: + files = glob.glob('results*.txt') + glob.glob('../../Downloads/results*.txt') + for fi, f in enumerate(files): + try: + results = np.loadtxt(f, usecols=[2, 3, 4, 8, 9, 12, 13, 14, 10, 11], ndmin=2).T + n = results.shape[1] # number of rows + x = range(start, min(stop, n) if stop else n) + for i in range(10): + y = results[i, x] + if i in [0, 1, 2, 5, 6, 7]: + y[y == 0] = np.nan # dont show zero loss values + # y /= y[0] # normalize + label = labels[fi] if len(labels) else Path(f).stem + ax[i].plot(x, y, marker='.', label=label, linewidth=2, markersize=8) + ax[i].set_title(s[i]) + # if i in [5, 6, 7]: # share train and val loss y axes + # ax[i].get_shared_y_axes().join(ax[i], ax[i - 5]) + except: + print('Warning: Plotting error for %s, skipping file' % f) + + fig.tight_layout() + ax[1].legend() + fig.savefig('results.png', dpi=200) diff --git a/hair_service_sd/utils/weight_init.py b/hair_service_sd/utils/weight_init.py new file mode 100644 index 0000000..091fe13 --- /dev/null +++ b/hair_service_sd/utils/weight_init.py @@ -0,0 +1,32 @@ +#!/usr/bin/env python3 +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. + +import torch.nn as nn + + +def c2_xavier_fill(module: nn.Module): + """ + Initialize `module.weight` using the "XavierFill" implemented in Caffe2. + Also initializes `module.bias` to 0. + + Args: + module (torch.nn.Module): module to initialize. + """ + # Caffe2 implementation of XavierFill in fact + # corresponds to kaiming_uniform_ in PyTorch + nn.init.kaiming_uniform_(module.weight, a=1) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + +def c2_msra_fill(module: nn.Module): + """ + Initialize `module.weight` using the "MSRAFill" implemented in Caffe2. + Also initializes `module.bias` to 0. + + Args: + module (torch.nn.Module): module to initialize. + """ + nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu") + if module.bias is not None: + nn.init.constant_(module.bias, 0) diff --git a/hair_service_sd/watch_delete.py b/hair_service_sd/watch_delete.py new file mode 100644 index 0000000..c26277f --- /dev/null +++ b/hair_service_sd/watch_delete.py @@ -0,0 +1,21 @@ +import os +import sys +from apscheduler.schedulers.blocking import BlockingScheduler +import schedule +import time +from datetime import datetime +from common.logger import config + +user_dir = config.get('default', 'userDir') + +# 输出时间 +def my_job(): + os.system(f'rm -rf {user_dir}/*') + # print('fffffff') + print(datetime.now().strftime("%Y-%m-%d %H:%M:%S")) + + +schedule.every().day.at("17:13").do(my_job) +while True: + schedule.run_pending() + time.sleep(1) \ No newline at end of file diff --git a/kohya_ss_home/start_docker.sh b/kohya_ss_home/start_docker.sh new file mode 100755 index 0000000..670d29a --- /dev/null +++ b/kohya_ss_home/start_docker.sh @@ -0,0 +1 @@ +docker run --rm -it --gpus all -v $(pwd):/home/chinatszrn -v /mnt:/mnt $* --net=host chinatszrn/ubuntu:kohya_ss bash diff --git a/meidaojia/api_load_test.py b/meidaojia/api_load_test.py new file mode 100644 index 0000000..6127fd2 --- /dev/null +++ b/meidaojia/api_load_test.py @@ -0,0 +1,162 @@ +import requests +import time +import random +import string +from concurrent.futures import ThreadPoolExecutor, as_completed +import argparse +from threading import Lock + +# 全局计数器 +class RequestCounter: + def __init__(self): + self.sent = 0 + self.in_progress = 0 + self.completed = 0 + self.success = 0 + self.failed = 0 + self.lock = Lock() + + def increment_sent(self): + with self.lock: + self.sent += 1 + + def increment_in_progress(self): + with self.lock: + self.in_progress += 1 + + def decrement_in_progress(self): + with self.lock: + self.in_progress -= 1 + + def increment_completed(self): + with self.lock: + self.completed += 1 + + def increment_success(self): + with self.lock: + self.success += 1 + + def increment_failed(self): + with self.lock: + self.failed += 1 + + def get_stats(self): + with self.lock: + return { + 'sent': self.sent, + 'in_progress': self.in_progress, + 'completed': self.completed, + 'success': self.success, + 'failed': self.failed + } + +counter = RequestCounter() + + +def send_request(url, headers, data): + """发送单个请求""" + try: + # 更新计数器 + counter.increment_sent() + counter.increment_in_progress() + + response = requests.post(url, headers=headers, json=data) + + # 根据响应状态更新计数器 + if response.status_code == 200: + counter.increment_success() + else: + counter.increment_failed() + + return response.status_code, response.text + except Exception as e: + counter.increment_failed() + return None, str(e) + finally: + counter.decrement_in_progress() + counter.increment_completed() + +def print_stats(): + """打印实时统计信息""" + while True: + stats = counter.get_stats() + print(f"\rStats - Sent: {stats['sent']}, In Progress: {stats['in_progress']}, " + f"Completed: {stats['completed']} (Success: {stats['success']}, Failed: {stats['failed']})", + end="", flush=True) + time.sleep(10) + +def worker(url, headers, data, qps, duration): + """工作线程函数，控制QPS""" + requests_count = 0 + start_time = time.time() + end_time = start_time + duration + + while time.time() < end_time: + request_start = time.time() + status_code, response_text = send_request(url, headers, data) + requests_count += 1 + + # 控制QPS + elapsed = time.time() - request_start + sleep_time = max(0, (1.0 / qps) - elapsed) + if sleep_time > 0: + time.sleep(sleep_time) + + return requests_count + +def run_test(url, headers, data, qps, duration, concurrency): + """运行测试""" + total_requests = 0 + start_time = time.time() + + # 启动统计信息打印线程 + import threading + stats_thread = threading.Thread(target=print_stats, daemon=True) + stats_thread.start() + + with ThreadPoolExecutor(max_workers=concurrency) as executor: + futures = [executor.submit(worker, url, headers, data, qps, duration) + for _ in range(concurrency)] + + for future in as_completed(futures): + total_requests += future.result() + + end_time = time.time() + actual_duration = end_time - start_time + actual_qps = total_requests / actual_duration + + # 获取最终统计 + stats = counter.get_stats() + + print("\n\nTest Summary:") + print(f"Total requests: {total_requests}") + print(f" Success: {stats['success']}") + print(f" Failed: {stats['failed']}") + print(f"Success rate: {(stats['success']/total_requests)*100:.2f}%") + print(f"Test duration: {actual_duration:.2f} seconds") + print(f"Actual QPS: {actual_qps:.2f}") + print(f"Target QPS: {qps}") + print(f"Concurrency: {concurrency}") + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description='API Load Test Tool') + parser.add_argument('--qps', type=float, default=1, help='Requests per second (per thread)') + parser.add_argument('--duration', type=float, default=10, help='Test duration in seconds') + parser.add_argument('--concurrency', type=int, default=2, help='Number of concurrent threads') + + args = parser.parse_args() + + # 配置请求参数 + url = 'http://mq.aidigifi.meidaojia.com/hairColor/v2' + # url = 'https://692139771842565-http-8801.northwest1.gpugeek.com:8443/hairColor/v2' + headers = {'Content-Type': 'application/json'} + data = { + "img": "data:image/jpeg;base64,/9j/4AAQSkZJ....Euy+sn/oZrKn8CPlK2qR/9k=", + "userId": "18701620166", + "rgb":[77,99,0], + "ratio": 0.9, + "output_format": "base64" + } + + print(f"Starting test with QPS={args.qps}, duration={args.duration}s, concurrency={args.concurrency}") + run_test(url, headers, data, args.qps, args.duration, args.concurrency) \ No newline at end of file diff --git a/meidaojia/api_load_test_color.py b/meidaojia/api_load_test_color.py new file mode 100644 index 0000000..ee7f5ec --- /dev/null +++ b/meidaojia/api_load_test_color.py @@ -0,0 +1,203 @@ +import requests +import time +import base64 +import random +import string +from concurrent.futures import ThreadPoolExecutor, as_completed +import argparse +from threading import Lock + +# 全局计数器 +class RequestCounter: + def __init__(self): + self.sent = 0 + self.in_progress = 0 + self.completed = 0 + self.success = 0 + self.failed = 0 + self.lock = Lock() + + def increment_sent(self): + with self.lock: + self.sent += 1 + + def increment_in_progress(self): + with self.lock: + self.in_progress += 1 + + def decrement_in_progress(self): + with self.lock: + self.in_progress -= 1 + + def increment_completed(self): + with self.lock: + self.completed += 1 + + def increment_success(self): + with self.lock: + self.success += 1 + + def increment_failed(self): + with self.lock: + self.failed += 1 + + def get_stats(self): + with self.lock: + return { + 'sent': self.sent, + 'in_progress': self.in_progress, + 'completed': self.completed, + 'success': self.success, + 'failed': self.failed + } + +counter = RequestCounter() + +def generate_random_task_id(length=19): + """生成随机task_id""" + digits = string.digits + return ''.join(random.choice(digits) for _ in range(length)) + +def send_request(url, headers, data): + """发送单个请求""" + try: + # 更新计数器 + counter.increment_sent() + counter.increment_in_progress() + + # 每次请求生成新的随机task_id + data["task_id"] = generate_random_task_id() + response = requests.post(url, headers=headers, json=data) + + # 根据响应状态更新计数器 + if response.status_code == 200: + counter.increment_success() + else: + counter.increment_failed() + + return response.status_code, response.text + except Exception as e: + counter.increment_failed() + return None, str(e) + finally: + counter.decrement_in_progress() + counter.increment_completed() + +def print_stats(): + """打印实时统计信息""" + while True: + stats = counter.get_stats() + print(f"\rStats - Sent: {stats['sent']}, In Progress: {stats['in_progress']}, " + f"Completed: {stats['completed']} (Success: {stats['success']}, Failed: {stats['failed']})", + end="", flush=True) + time.sleep(10) + +def worker(url, headers, data, qps, duration): + """工作线程函数，控制QPS""" + requests_count = 0 + start_time = time.time() + end_time = start_time + duration + + while time.time() < end_time: + request_start = time.time() + status_code, response_text = send_request(url, headers, data) + requests_count += 1 + + # 控制QPS + elapsed = time.time() - request_start + sleep_time = max(0, (1.0 / qps) - elapsed) + if sleep_time > 0: + time.sleep(sleep_time) + + return requests_count + +def run_test(url, headers, data, qps, duration, concurrency): + """运行测试""" + total_requests = 0 + start_time = time.time() + + # 启动统计信息打印线程 + import threading + stats_thread = threading.Thread(target=print_stats, daemon=True) + stats_thread.start() + + with ThreadPoolExecutor(max_workers=concurrency) as executor: + futures = [executor.submit(worker, url, headers, data, qps, duration) + for _ in range(concurrency)] + + for future in as_completed(futures): + total_requests += future.result() + + end_time = time.time() + actual_duration = end_time - start_time + actual_qps = total_requests / actual_duration + + # 获取最终统计 + stats = counter.get_stats() + + print("\n\nTest Summary:") + print(f"Total requests: {total_requests}") + print(f" Success: {stats['success']}") + print(f" Failed: {stats['failed']}") + print(f"Success rate: {(stats['success']/total_requests)*100:.2f}%") + print(f"Test duration: {actual_duration:.2f} seconds") + print(f"Actual QPS: {actual_qps:.2f}") + print(f"Target QPS: {qps}") + print(f"Concurrency: {concurrency}") + + +def image_to_base64(file_path, mime_type=None): + """ + 将图片文件转换为带Base64前缀的Data URI字符串 + + 参数: + file_path (str): 图片文件路径 + mime_type (str): 可选，指定MIME类型。如果为None，则根据文件扩展名自动判断 + + 返回: + str: 带Base64前缀的Data URI字符串 + """ + # 如果没有指定MIME类型，根据文件扩展名推断 + if mime_type is None: + extension = file_path.split('.')[-1].lower() + mime_types = { + 'jpg': 'image/jpeg', + 'jpeg': 'image/jpeg', + 'png': 'image/png', + 'gif': 'image/gif', + 'webp': 'image/webp', + 'bmp': 'image/bmp' + } + mime_type = mime_types.get(extension, 'application/octet-stream') + + # 读取文件内容并编码为Base64 + with open(file_path, 'rb') as image_file: + encoded_string = base64.b64encode(image_file.read()).decode('utf-8') + + # 组合成Data URI格式 + return f"data:{mime_type};base64,{encoded_string}" + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description='API Load Test Tool') + parser.add_argument('--qps', type=float, default=1, help='Requests per second (per thread)') + parser.add_argument('--duration', type=float, default=10, help='Test duration in seconds') + parser.add_argument('--concurrency', type=int, default=2, help='Number of concurrent threads') + + args = parser.parse_args() + + # 配置请求参数 + url = 'http://mq.aidigifi.meidaojia.com/api/swapHair/v1' + headers = {'Content-Type': 'application/json'} + data = { + "hair_id": "1907651680352395265", + "task_id": "1907651680352395265", # 会被每次请求覆盖 + "user_img_path": "https://cdn.meidaojia.com/ZoeFiles/user2_1_%E5%89%AF%E6%9C%AC.JPG", + "is_hr": "false", + "output_format": "base64" + } + + img64_str = image_to_base64("aaa.jpg") + data['img'] = img64_str + + print(f"Starting test with QPS={args.qps}, duration={args.duration}s, concurrency={args.concurrency}") + run_test(url, headers, data, args.qps, args.duration, args.concurrency) \ No newline at end of file diff --git a/meidaojia/app.py b/meidaojia/app.py new file mode 100755 index 0000000..eb890f0 --- /dev/null +++ b/meidaojia/app.py @@ -0,0 +1,228 @@ +import time +import logging +import redis +import uuid +import json +from flask import Flask, request, jsonify +from datetime import datetime, timedelta +from config import QUEUE_NAME, DEFAULT_TIMEOUT, KEY_QUEUE_LOCK_NAME, acquire_lock + +# 创建 logger +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) # 设置 logger 的级别 + + +# 获取当前日期和时间 +now = datetime.now() + +# 格式化为字符串（例如：2023-10-25 14:30:45） +date_time_str = now.strftime("%Y-%m-%d_%H:%M:%S") + +# 创建文件 handler +file_handler = logging.FileHandler(f'/var/log/meidaojia/app_{date_time_str}.log') +file_handler.setLevel(logging.INFO) +file_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s') +file_handler.setFormatter(file_formatter) + +# 创建控制台 handler +console_handler = logging.StreamHandler() +console_handler.setLevel(logging.INFO) +console_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s') +console_handler.setFormatter(console_formatter) + +# 添加 handlers 到 logger +logger.addHandler(file_handler) +logger.addHandler(console_handler) + +app = Flask(__name__) + +# 配置 +from config import REDIS_HOST +from config import REDIS_PORT +from config import REDIS_DB + + +# 创建 Redis 连接池 +redis_pool = redis.ConnectionPool( + host=REDIS_HOST, + port=REDIS_PORT, + db=REDIS_DB, + max_connections=2000 # 根据实际情况调整 +) + +def get_redis_conn(): + """获取 Redis 连接""" + return redis.Redis(connection_pool=redis_pool) + +get_redis_conn().set(QUEUE_NAME, "[]") + +def pushStr2Queue(key, data_str): + redis_conn = get_redis_conn() + lock = acquire_lock(redis_conn, KEY_QUEUE_LOCK_NAME) + if lock: + try: + redis_conn.set(key, data_str, ex=60) + key_queue_str = redis_conn.get(QUEUE_NAME).decode("utf-8") + key_queue = json.loads(key_queue_str) + key_queue.append(key) + logger.info(f"pushStr2Queue queue len:{key_queue}") + key_queue_str = json.dumps(key_queue) + redis_conn.set(QUEUE_NAME, key_queue_str) + finally: + lock.release() + else: + logger.error(f"pushStr2Queue get lock error {data_str}") + + + + +@app.route('/hairColor/v2', methods=['POST']) +def api_hairColor_v2(): + # 获取请求参数 + data = request.get_json() + if not data or 'img' not in data: + return jsonify({'msg': 'Missing img parameter', "state":-1, "data":"" }), 400 + + if not data or 'rgb' not in data: + return jsonify({'msg': 'Missing "rgb": parameter', "state":-1, "data":""}), 400 + + if not data or 'ratio' not in data: + return jsonify({'msg': 'Missing "ratio": parameter', "state":-1, "data":""}), 400 + + if not data or 'output_format' not in data: + return jsonify({'msg': 'Missing "output_format": parameter', "state":-1, "data":""}), 400 + + + key = str(uuid.uuid4()) + redis_conn = get_redis_conn() + + # if len(data['img']) > 256: + # redis_conn.set(f"img_{key}", data['img'], ex=60) + # data['img'] = "base64" + + task_data = {} + task_data['key'] = key + task_data['api'] = "/hairColor/v2" + task_data['request'] = data + logger.info(f"request hairColor/v2 img:{data['img'][:64]}, rgb:{data['rgb']}, ratio:{data['ratio']} output_format:{data['output_format']}") + task_data_str = json.dumps(task_data) + pushStr2Queue(key, task_data_str) + timeout = DEFAULT_TIMEOUT + start_time = datetime.now() + result_key = f"result_{key}" + result_status_key = f"result_status_code_{key}" + while (datetime.now() - start_time).seconds < timeout: + if redis_conn.exists(result_key): + result_str = redis_conn.get(result_key) + return result_str, int(redis_conn.get(result_status_key)), {'Content-Type': 'application/json'} + time.sleep(0.1) + + logger.error(f"request hairColor time out ") + return jsonify({ + 'msg': f'Timeout after {timeout} seconds, http time out' + , "state":-1, "data":"" + }), 408 + + +@app.route('/api/swapHair/v1', methods=['POST']) +def api_swapHair_v1(): + if not request.is_json: + return jsonify({"error": "Request must be JSON"}), 400 + # 获取请求参数 + data = request.get_json() + if not data or 'hair_id' not in data: + return jsonify({'msg': 'Missing hair_id parameter', "state":-1, "data":"" }), 400 + if not data or 'task_id' not in data: + return jsonify({'msg': 'Missing task_id parameter', "state":-1, "data":""}), 400 + if not data or 'user_img_path' not in data: + return jsonify({'msg': 'Missing user_img_path parameter', "state":-1, "data":""}), 400 + if not data or 'is_hr' not in data: + return jsonify({'msg': 'Missing is_hr parameter', "state":-1, "data":""}), 400 + + if not data or 'output_format' not in data: + return jsonify({'msg': 'Missing output_format parameter', "state":-1, "data":""}), 400 + + redis_conn = get_redis_conn() + key = str(uuid.uuid4()) + + # if len(data['user_img_path']) > 256: + # redis_conn.set(f"img_{key}", data['user_img_path'], ex=60) + # data['user_img_path'] = "base64" + + task_data = {} + task_data['key'] = key + task_data['api'] = "/api/swapHair/v1" + task_data['request'] = data + + logger.info(f"request /api/swapHair/v1 user_img_path:{data['user_img_path'][:64]} hair_id:{data['hair_id']}") + + task_data_str = json.dumps(task_data) + pushStr2Queue(key, task_data_str) + timeout = DEFAULT_TIMEOUT + start_time = datetime.now() + result_key = f"result_{key}" + result_status_key = f"result_status_code_{key}" + while (datetime.now() - start_time).seconds < timeout: + if redis_conn.exists(result_key): + result_str = redis_conn.get(result_key) + return result_str, int(redis_conn.get(result_status_key)), {'Content-Type': 'application/json'} + time.sleep(0.1) + + logger.error(f"request /api/swapHair/v1 time out") + return jsonify({ + 'msg': f'Timeout after {timeout} seconds, http time out' + , "state":-1, "data":"" + }), 408 + + + +# @app.route('/api/uploadHair/v1', methods=['POST']) +# def api_uploadHair_v1(): +# # 获取请求参数 +# data = request.get_json() +# if not data or 'img_lists' not in data: +# return jsonify({'msg': 'Missing img_lists parameter', "state":-1, "data":"" }), 400 +# if not data or '"hair_id": ' not in data: +# return jsonify({'msg': 'Missing "hair_id": parameter', "state":-1, "data":""}), 400 + +# if not data or 'output_format' not in data: +# logger.warning(f"api_swapHair_v1 no output_format task_id:{data['task_id']}") + + +# key = str(uuid.uuid4()) +# redis_conn = get_redis_conn() +# task_data = {} +# task_data['key'] = key +# task_data['api'] = "/api/uploadHair/v1" +# task_data['request'] = data +# logger.info(f"request api_uploadHair_v1 task_data: {json.dumps(task_data)}") +# task_data_str = json.dumps(task_data) +# pushStr2Queue(task_data_str) +# timeout = DEFAULT_TIMEOUT +# start_time = datetime.now() +# result_key = f"result_{key}" +# while (datetime.now() - start_time).seconds < timeout: +# if redis_conn.exists(result_key): +# result_str = redis_conn.get(result_key) +# result = json.loads(result_str) +# if result['state'] == 0: +# return jsonify(result), 200 +# else: +# return jsonify(result), 500 +# time.sleep(0.1) + +# logger.error(f"request api_uploadHair_v1 time out") +# return jsonify({ +# 'msg': f'Timeout after {timeout} seconds, http time out' +# , "state":-1, "data":"" +# }), 408 + + +if __name__ == '__main__': + # 启动Flask应用，启用多线程处理 + app.run( + host='0.0.0.0', + port=80, + threaded=True, # 启用多线程处理并发请求 + debug=False # 生产环境应设置为False + ) \ No newline at end of file diff --git a/meidaojia/autossh.sh b/meidaojia/autossh.sh new file mode 100755 index 0000000..8a79473 --- /dev/null +++ b/meidaojia/autossh.sh @@ -0,0 +1,69 @@ +#!/bin/bash + +# 定义连接参数 +HOST="northwest1.gpugeek.com" +PORT="55139" +USER="root" +PASSWORD="uaUMSzmAHyka2ZcBvQHZ5My9PUNRP78T" +TARGET_DIR="/root/project/hair_service_sd" +COMMAND="if ! pgrep -f 'run_copy_cost_colorb64.py'; then bash start_services.sh; fi" + +# 使用expect工具自动化交互过程 +/usr/bin/expect < 0: + total_time = stats.end_time - stats.start_time + avg_response_time = sum(stats.response_times) / len(stats.response_times) + qps = stats.calculate_qps() + + print(f"测试总时间: {total_time:.3f} 秒") + print(f"平均响应时间: {avg_response_time:.3f} 秒") + print(f"最大响应时间: {max(stats.response_times):.3f} 秒") + print(f"最小响应时间: {min(stats.response_times):.3f} 秒") + print(f"QPS (每秒查询率): {qps:.2f}") + else: + print("所有请求都失败了，无法计算统计信息") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/meidaojia/request_rsp_time_test.py b/meidaojia/request_rsp_time_test.py new file mode 100644 index 0000000..c39df71 --- /dev/null +++ b/meidaojia/request_rsp_time_test.py @@ -0,0 +1,93 @@ +import requests +import random +import time +import base64 + +# url = 'https://693565319933957-http-8801.northwest1.gpugeek.com:8443/api/swapHair/v1' +url = 'http://172.17.110.92/api/swapHair/v1' + +headers = {'Content-Type': 'application/json'} +data = { + "hair_id": "1907651680352395265", + "task_id": "1907651680352395265", + "user_img_path": "https://cdn.meidaojia.com/ZoeFiles/user2_1_%E5%89%AF%E6%9C%AC.JPG", + "is_hr": "false", + "output_format": "base64" +} + +def image_to_base64(file_path, mime_type=None): + """ + 将图片文件转换为带Base64前缀的Data URI字符串 + + 参数: + file_path (str): 图片文件路径 + mime_type (str): 可选，指定MIME类型。如果为None，则根据文件扩展名自动判断 + + 返回: + str: 带Base64前缀的Data URI字符串 + """ + # 如果没有指定MIME类型，根据文件扩展名推断 + if mime_type is None: + extension = file_path.split('.')[-1].lower() + mime_types = { + 'jpg': 'image/jpeg', + 'jpeg': 'image/jpeg', + 'png': 'image/png', + 'gif': 'image/gif', + 'webp': 'image/webp', + 'bmp': 'image/bmp' + } + mime_type = mime_types.get(extension, 'application/octet-stream') + + # 读取文件内容并编码为Base64 + with open(file_path, 'rb') as image_file: + encoded_string = base64.b64encode(image_file.read()).decode('utf-8') + + # 组合成Data URI格式 + return f"data:{mime_type};base64,{encoded_string}" + +# img64_str = image_to_base64("aaa.jpg") +# data['user_img_path'] = img64_str + +def send_request(): + try: + data['task_id'] = str(int(random.uniform(0, 10000000))) + start_time = time.time() + response = requests.post(url, headers=headers, json=data) + end_time = time.time() + response_time = end_time - start_time + return response_time + except Exception as e: + print(f"请求发生错误: {e}") + return None + +def main(): + total_requests = 30 + response_times = [] + + for i in range(total_requests): + print(f"正在发送第 {i+1} 个请求...") + rt = send_request() + if rt is not None: + response_times.append(rt) + print(f"第 {i+1} 个请求完成，响应时间: {rt:.3f} 秒") + else: + print(f"第 {i+1} 个请求失败") + + # 单路串行，不需要延迟，但如果需要可以添加 + # time.sleep(1) + + if response_times: + avg_response_time = sum(response_times) / len(response_times) + print("\n统计结果:") + print(f"总请求数: {total_requests}") + print(f"成功请求数: {len(response_times)}") + print(f"失败请求数: {total_requests - len(response_times)}") + print(f"平均响应时间: {avg_response_time:.3f} 秒") + print(f"最大响应时间: {max(response_times):.3f} 秒") + print(f"最小响应时间: {min(response_times):.3f} 秒") + else: + print("所有请求都失败了，无法计算统计信息") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/meidaojia/servers.csv b/meidaojia/servers.csv new file mode 100644 index 0000000..811691e --- /dev/null +++ b/meidaojia/servers.csv @@ -0,0 +1,45 @@ +693573774479365,ssh -p 51884 root@northwest1.gpugeek.com,VZxzyqDNbucAEdm5nZQFFQ67v6X7vQcK,,,,,,,,,,,,,,,,, +693573706682373,ssh -p 46901 root@northwest1.gpugeek.com,ZDdguaaQsNRb8qmMaTu7q95eu2vG5Qc3,,,,,,,,,,,,,,,,, +693570752139269,ssh -p 44853 root@northwest1.gpugeek.com,Wavu5Cnewm65p5ABXqXuSMFfF2c77ZuG,,,,,,,,,,,,,,,,, +693570664882181,ssh -p 49563 root@northwest1.gpugeek.com,4mz3pKZdVff9X3hFzcfVY4h4VkzSTpcb,,,,,,,,,,,,,,,,, +693565557084165,ssh -p 57192 root@northwest1.gpugeek.com,5ztvzCwtEaBG8CYaCbyxZVXmzmUxuD7b,,,,,,,,,,,,,,,,, +693565452951557,ssh -p 40188 root@northwest1.gpugeek.com,wWAsrQ5yGvahpF2zHXVFrRterMeYWuN3,,,,,,,,,,,,,,,,, +693565319933957,ssh -p 55139 root@northwest1.gpugeek.com,uaUMSzmAHyka2ZcBvQHZ5My9PUNRP78T,,,,,,,,,,,,,,,,, +693597199511557,ssh -p 54874 root@northwest1.gpugeek.com,zN2SB3DZxNpnvpQzC7PBHun3CNh2ydGV,,,,,,,,,,,,,,,,, +693597082546181,ssh -p 40885 root@northwest1.gpugeek.com,HdYQ3QtzKdRgNQZaXWDKYSqWPkE2w53h,,,,,,,,,,,,,,,,, +693596619751429,ssh -p 44684 root@northwest1.gpugeek.com,tFF9WKSe5XRSDR8PWcNweF2cZ6ZdYpaW,,,,,,,,,,,,,,,,, +693596213850117,ssh -p 44307 root@northwest1.gpugeek.com,YSMNFffKTMY37Ze83UfmSr5XF65RP7dT,,,,,,,,,,,,,,,,, +693595248599045,ssh -p 54178 root@northwest1.gpugeek.com,akZ3qyRtqUum553sMNXEVNY3qzxMA3sV,,,,,,,,,,,,,,,,, +693595168616453,ssh -p 59719 root@northwest1.gpugeek.com,XqQzT8qtfUWFQP8E3PPMb7FMPRUYBehS,,,,,,,,,,,,,,,,, +693594319818757,ssh -p 56408 root@northwest1.gpugeek.com,46SMhqxfqkyvVu6qKyr9UW7GfTnK5GCp,,,,,,,,,,,,,,,,, +693594143621125,ssh -p 47103 root@northwest1.gpugeek.com,ZvxQPDvKQeVdRMnm3VTMD5vfdg8VpqCE,,,,,,,,,,,,,,,,, +693579634159621,ssh -p 45208 root@northwest1.gpugeek.com,467nsPAHNgS5fvzNNvzCaebzs3Z4DB8s,,,,,,,,,,,,,,,,, +693579093471237,ssh -p 40600 root@northwest1.gpugeek.com,mRabksFkwYV8NWx2NmAufTevNrFUeVBw,,,,,,,,,,,,,,,,, +693578025705477,ssh -p 58639 root@northwest1.gpugeek.com,WSXHUssQSXTCWSDuVCdcvG8DQyFPU2er,,,,,,,,,,,,,,,,, +693577880313861,ssh -p 59280 root@northwest1.gpugeek.com,CywKAemA3b2z6VFG5TzgYzAKy9bnGvNM,,,,,,,,,,,,,,,,, +693577753862149,ssh -p 58493 root@northwest1.gpugeek.com,yDxpPQbyxEuNbPQPapUV72vK3s9CRu3F,,,,,,,,,,,,,,,,, +693577626185733,ssh -p 59071 root@northwest1.gpugeek.com,UFNpw9B48ahAU8RrmcGVmSENazGCSSm5,,,,,,,,,,,,,,,,, +693577527365637,ssh -p 57796 root@northwest1.gpugeek.com,2sZhwEmS7RdgQeuCpSnQqgpfR4BQnhvd,,,,,,,,,,,,,,,,, +693575032029189,ssh -p 44408 root@northwest1.gpugeek.com,yM6wR5g8VGXbb4XfvA9ghYWVk9maTht8,,,,,,,,,,,,,,,,, +693574955724805,ssh -p 43820 root@northwest1.gpugeek.com,csQbqbCAkdwTZZSVdmxMq3svKhdpEAtF,,,,,,,,,,,,,,,,, +699531043577861,ssh -p 53934 root@northwest1.gpugeek.com,WD3WcCNTuz5VxmbAdmWFSYZVfhEXf8NK,,,,,,,,,,,,,,,,, +699531676672005,ssh -p 53976 root@northwest1.gpugeek.com,p37A8HRPcAU8SP4ssha4VBS9phgrykF9,,,,,,,,,,,,,,,,, +699552291840005,ssh -p 52102 root@northwest1.gpugeek.com,S5G4zWPQ8WMZ89Tk3afTgQAAfpqXSuM3,,,,,,,,,,,,,,,,, +699552392802309,ssh -p 48055 root@northwest1.gpugeek.com,afXHTcaPMqbWnqCfWv6tzSFkBKFFf7vS,,,,,,,,,,,,,,,,, +699552766849029,ssh -p 51140 root@northwest1.gpugeek.com,9YFyTAruNZwr8K44ZXp9ET5aG5c4DTa4,,,,,,,,,,,,,,,,, +699554474184709,ssh -p 40280 root@northwest1.gpugeek.com,sbSfgFa8V4bhBDpbqtccCvXkHF6gSfAn,,,,,,,,,,,,,,,,, +699554413039621,ssh -p 49817 root@northwest1.gpugeek.com,dRB7CRtcdgGwFe4nbxaEh8D7nXMmuYDX,,,,,,,,,,,,,,,,, +699554332229637,ssh -p 40625 root@northwest1.gpugeek.com,HQz2nC7tbTTxGZhC7YN4YFCg5WfcMDD2,,,,,,,,,,,,,,,,, +699554276352005,ssh -p 43763 root@northwest1.gpugeek.com,bYFCDzzN2MKtWHF4qvyK9PtG9xqurA5z,,,,,,,,,,,,,,,,, +699554118586373,ssh -p 49257 root@northwest1.gpugeek.com,GNMmwCpB3NcGryqznhB2acDmtgFPBYXs,,,,,,,,,,,,,,,,, +699554021834757,ssh -p 47192 root@northwest1.gpugeek.com,wFxwYDD6KapUSYgAanEcYzDRusQCvPBT,,,,,,,,,,,,,,,,, +699553939038213,ssh -p 47118 root@northwest1.gpugeek.com,fXU3wEvUE4ckCfSXwkVW7dW8xG3VqRh2,,,,,,,,,,,,,,,,, +699553843216389,ssh -p 40918 root@northwest1.gpugeek.com,3EYpyAx7s5gbHQEXWNZ6Td3QdMCZU8Mu,,,,,,,,,,,,,,,,, +699553742393349,ssh -p 57046 root@northwest1.gpugeek.com,a6KeZhBPMWecTtvsP3kbgcV7MZseyerY,,,,,,,,,,,,,,,,, +699553618411525,ssh -p 52041 root@northwest1.gpugeek.com,tev3z7knqaaSP2T74QtYqtfb9HdNy9WW,,,,,,,,,,,,,,,,, +699553554767877,ssh -p 42691 root@northwest1.gpugeek.com,r3bFQs4uNFyUCDN99f7deH2pXqdmRYEV,,,,,,,,,,,,,,,,, +699553450074117,ssh -p 42199 root@northwest1.gpugeek.com,hYgx3rbWpmYSFs6pEq8umZqcgAZyZa9p,,,,,,,,,,,,,,,,, +699553329291269,ssh -p 49696 root@northwest1.gpugeek.com,EVnYtmqBZTtdfzXPX39CukVugrX6HzEW,,,,,,,,,,,,,,,,, +699553242570757,ssh -p 57669 root@northwest1.gpugeek.com,FU9UmrBvMeCGvvxfDdX5fWCn4E9AtTBg,,,,,,,,,,,,,,,,, +699552977174533,ssh -p 41179 root@northwest1.gpugeek.com,NFx6ahHDQKpqFEBg4nrPMerureYX5dMV,,,,,,,,,,,,,,,,, +699552838074373,ssh -p 46220 root@northwest1.gpugeek.com,qZTCYqU4KpB9GbZfgxXAwcxwMcKM9eFz,,,,,,,,,,,,,,,,, diff --git a/meidaojia/servers_test.csv b/meidaojia/servers_test.csv new file mode 100644 index 0000000..914831e --- /dev/null +++ b/meidaojia/servers_test.csv @@ -0,0 +1,7 @@ +699553450074117,ssh -p 42199 root@northwest1.gpugeek.com,hYgx3rbWpmYSFs6pEq8umZqcgAZyZa9p,,,,,,,,,,,,,,,,, +699553742393349,ssh -p 57046 root@northwest1.gpugeek.com,a6KeZhBPMWecTtvsP3kbgcV7MZseyerY,,,,,,,,,,,,,,,,, +699553618411525,ssh -p 52041 root@northwest1.gpugeek.com,tev3z7knqaaSP2T74QtYqtfb9HdNy9WW,,,,,,,,,,,,,,,,, +699553242570757,ssh -p 57669 root@northwest1.gpugeek.com,FU9UmrBvMeCGvvxfDdX5fWCn4E9AtTBg,,,,,,,,,,,,,,,,, +699553554767877,ssh -p 42691 root@northwest1.gpugeek.com,r3bFQs4uNFyUCDN99f7deH2pXqdmRYEV,,,,,,,,,,,,,,,,, +699552977174533,ssh -p 41179 root@northwest1.gpugeek.com,NFx6ahHDQKpqFEBg4nrPMerureYX5dMV,,,,,,,,,,,,,,,,, +699552838074373,ssh -p 46220 root@northwest1.gpugeek.com,qZTCYqU4KpB9GbZfgxXAwcxwMcKM9eFz,,,,,,,,,,,,,,,,, diff --git a/meidaojia/static/index.html b/meidaojia/static/index.html new file mode 100644 index 0000000..256ca49 --- /dev/null +++ b/meidaojia/static/index.html @@ -0,0 +1,382 @@ + + + + + + 服务器状态监控 + + + +
+

服务器状态监控

+
+ +
+
+ 最后更新时间: - +
+
+ +
+
+ + + + \ No newline at end of file diff --git a/meidaojia/test_color.py b/meidaojia/test_color.py new file mode 100644 index 0000000..42cc941 --- /dev/null +++ b/meidaojia/test_color.py @@ -0,0 +1,146 @@ +import threading +import requests +import json +import base64 +import time +import random + + +# url = 'http://172.17.110.92/hairColor/v2' +url = 'https://779460252262853-http-8801.northwest1.gpugeek.com:8443/hairColor/v2' +# url = 'http://xiangsilian.com:18801/hairColor/v2' + + + +headers = {'Content-Type': 'application/json'} +data = { + "img": "data:image/jpeg;base64,/9j/4AAQSkZJ....Euy+sn/oZrKn8CPlK2qR/9k=", + "userId": "18701620166", + "rgb":[255,0,0], + "ratio": 0.9, + "output_format": "base64" +} + +def image_to_base64(file_path, mime_type=None): + """ + 将图片文件转换为带Base64前缀的Data URI字符串 + + 参数: + file_path (str): 图片文件路径 + mime_type (str): 可选，指定MIME类型。如果为None，则根据文件扩展名自动判断 + + 返回: + str: 带Base64前缀的Data URI字符串 + """ + # 如果没有指定MIME类型，根据文件扩展名推断 + if mime_type is None: + extension = file_path.split('.')[-1].lower() + mime_types = { + 'jpg': 'image/jpeg', + 'jpeg': 'image/jpeg', + 'png': 'image/png', + 'gif': 'image/gif', + 'webp': 'image/webp', + 'bmp': 'image/bmp' + } + mime_type = mime_types.get(extension, 'application/octet-stream') + + # 读取文件内容并编码为Base64 + with open(file_path, 'rb') as image_file: + encoded_string = base64.b64encode(image_file.read()).decode('utf-8') + + # 组合成Data URI格式 + return f"data:{mime_type};base64,{encoded_string}" + +img64_str = image_to_base64("qwerqwe.jpg") +data['img'] = img64_str + +# 存储所有请求的响应 +responses = [] +responses_time = [] + +# 线程锁，防止多线程写入冲突 +lock = threading.Lock() + + +def send_request(): + try: + start_time = time.time() + response = requests.post(url, headers=headers, json=data) + end_time = time.time() + response_time = end_time - start_time + print(f"相应时间:{response_time}") + result = response.json() + if 'state' in result: + if result['state'] != 0: + print(f"请求失败: {result}") + time.sleep(5) + with lock: # 加锁，防止多线程同时修改 responses + responses.append(response) + responses_time.append(response_time) + except Exception as e: + with lock: + error_msg = {"error": str(e)} + responses.append(error_msg) + print(f"请求失败: {error_msg}") # 打印错误信息 + time.sleep(5) # 延时5秒 + +# 创建并启动多个线程 +threads = [] +num_requests = 1 # 并发请求数量 +test_rand_num = 1 + +time_start = time.time() + +for _ in range(test_rand_num): + for _ in range(num_requests): + t = threading.Thread(target=send_request) + t.start() + threads.append(t) + + # 等待所有线程完成 + for t in threads: + t.join() + + +end_time = time.time() + +total_time = (end_time - time_start) +qps = (num_requests*test_rand_num)/total_time +print(f"num_requests:{num_requests} total_time:{total_time} qps:{qps}") + +totalReqTime = 0 +for t in responses_time: + totalReqTime += t +avgTime = totalReqTime/len(responses_time) +print(f"平均相应时间:{avgTime}") + + +def save_base64_image(base64_str, filename): + """ + 将 Base64 编码的图片字符串保存为本地 JPG 文件。 + + :param base64_str: Base64 编码的图片字符串（可能包含 data:image/jpeg;base64, 前缀） + :param filename: 保存的文件名（例如 'image.jpg'） + """ + # 如果包含 data URL 前缀，去掉它 + if base64_str.startswith("data:image"): + base64_str = base64_str.split(",", 1)[1] + + try: + image_data = base64.b64decode(base64_str) + with open(filename, "wb") as f: + f.write(image_data) + print(f"图片已保存为 {filename}") + except Exception as e: + print(f"保存图片失败: {e}") + +# 打印所有响应 +for i, resp in enumerate(responses, 1): + print(f"请求 {i} 结果:", {resp.status_code, resp.text[:128]}) + json_obj = json.loads(resp.text) + data_string = json_obj['result'] + save_base64_image(data_string, f"color_result{i}.jpg") + +print(f"num_requests:{num_requests} total_time:{total_time} qps:{qps}") +print(f"平均相应时间:{avgTime}") \ No newline at end of file diff --git a/meidaojia/test_hair.py b/meidaojia/test_hair.py new file mode 100644 index 0000000..fdd9a9b --- /dev/null +++ b/meidaojia/test_hair.py @@ -0,0 +1,128 @@ +import threading +import requests +import json +import base64 +import time +import random + +# url = 'http://xiangsilian.com:18801/api/swapHair/v1' +url = 'https://779460252262853-http-8801.northwest1.gpugeek.com:8443/api/swapHair/v1' + +headers = {'Content-Type': 'application/json'} +data = { + #"hair_id": "1907651680352395265", + "hair_id": "1954538119216013313", + "task_id": "3455643521235", + "user_img_path": "https://cdn.meidaojia.com/ZoeFiles/testuser_8.JPG", + "is_hr": "false", + "output_format": "base64" +} + + + +def image_to_base64(file_path, mime_type=None): + """ + 将图片文件转换为带Base64前缀的Data URI字符串 + + 参数: + file_path (str): 图片文件路径 + mime_type (str): 可选，指定MIME类型。如果为None，则根据文件扩展名自动判断 + + 返回: + str: 带Base64前缀的Data URI字符串 + """ + # 如果没有指定MIME类型，根据文件扩展名推断 + if mime_type is None: + extension = file_path.split('.')[-1].lower() + mime_types = { + 'jpg': 'image/jpeg', + 'jpeg': 'image/jpeg', + 'png': 'image/png', + 'gif': 'image/gif', + 'webp': 'image/webp', + 'bmp': 'image/bmp' + } + mime_type = mime_types.get(extension, 'application/octet-stream') + + # 读取文件内容并编码为Base64 + with open(file_path, 'rb') as image_file: + encoded_string = base64.b64encode(image_file.read()).decode('utf-8') + + # 组合成Data URI格式 + return f"data:{mime_type};base64,{encoded_string}" + +img64_str = image_to_base64("wertwert.jpg") +data['user_img_path'] = img64_str + +responses = [] + +lock = threading.Lock() + + +send_index = 0 + +def send_request(start_time): + try: + sleep_time = start_time + time.sleep(sleep_time) + global send_index + send_data = data + response = requests.post(url, headers=headers, json=send_data) + with lock: # 加锁，防止多线程同时修改 responses + responses.append(response) + except Exception as e: + with lock: + responses.append({"error": str(e)}) + +threads = [] +num_requests = 1 + +time_start = time.time() + +start_time = 0 +for _ in range(num_requests): + start_time = 0 + t = threading.Thread(target=send_request, args=(start_time,)) + t.start() + threads.append(t) + + +for t in threads: + t.join() + +end_time = time.time() + +total_time = (end_time - time_start) +qps = num_requests/total_time +print(f"num_requests:{num_requests} total_time:{total_time} qps:{qps}") + + + +import base64 +import os + +def save_base64_image(base64_str, filename): + """ + 将 Base64 编码的图片字符串保存为本地 JPG 文件。 + + :param base64_str: Base64 编码的图片字符串（可能包含 data:image/jpeg;base64, 前缀） + :param filename: 保存的文件名（例如 'image.jpg'） + """ + # 如果包含 data URL 前缀，去掉它 + if base64_str.startswith("data:image"): + base64_str = base64_str.split(",", 1)[1] + + try: + image_data = base64.b64decode(base64_str) + with open(filename, "wb") as f: + f.write(image_data) + print(f"图片已保存为 {filename}") + except Exception as e: + print(f"保存图片失败: {e}") + +# 打印所有响应 +for i, resp in enumerate(responses, 1): + print(f"请求 {i} 结果: {resp.status_code, resp.text[:128]}") + json_obj = json.loads(resp.text) + data_string = json_obj['data'] + save_base64_image(data_string, f"response_{i}.jpg") \ No newline at end of file diff --git a/meidaojia/worker.py b/meidaojia/worker.py new file mode 100755 index 0000000..27b835e --- /dev/null +++ b/meidaojia/worker.py @@ -0,0 +1,428 @@ +import time +import datetime +import uuid +import redis +import logging +import requests +import threading +import json +from datetime import datetime, timedelta + +from config import REDIS_HOST +from config import REDIS_PORT +from config import REDIS_DB +from config import QUEUE_NAME +from config import GPU_SERVER_LIST +from config import SERVER_LOG_EVENT_LEN +from config import SERVER_LOG_ +from config import GPU_SERVER_TIME_OUT +from config import KEY_QUEUE_LOCK_NAME, acquire_lock, SERVER_LIST_LOCK_NAME, SERVER_LIST_LOG_LOCK_NAME +import copy + +import logging +import csv + +# 创建 logger +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) # 设置 logger 的级别 + + +# 获取当前日期和时间 +now = datetime.now() + +# 格式化为字符串（例如：2023-10-25 14:30:45） +date_time_str = now.strftime("%Y-%m-%d_%H:%M:%S") + +# 创建文件 handler +file_handler = logging.FileHandler(f'/var/log/meidaojia/worker_{date_time_str}.log') +file_handler.setLevel(logging.INFO) +file_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s') +file_handler.setFormatter(file_formatter) + +# 创建控制台 handler +console_handler = logging.StreamHandler() +console_handler.setLevel(logging.INFO) +console_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s') +console_handler.setFormatter(console_formatter) + +# 添加 handlers 到 logger +logger.addHandler(file_handler) +logger.addHandler(console_handler) + +# 创建 Redis 连接池 +redis_pool = redis.ConnectionPool( + host=REDIS_HOST, + port=REDIS_PORT, + db=REDIS_DB, + max_connections=2000 # 根据实际情况调整 +) + +def get_redis_conn(): + """获取 Redis 连接""" + return redis.Redis(connection_pool=redis_pool) + +redis_conn = get_redis_conn() + +def get_time(timestamp): + # 转换为UTC时间 + utc_time = datetime.utcfromtimestamp(timestamp) + + # 转换为北京时间 (UTC+8) + beijing_time = utc_time + timedelta(hours=8) # 正确 + + # 格式化输出 + formatted_time = beijing_time.strftime("%Y-%m-%d %H:%M:%S") + return formatted_time + +def server_log_event(name, action, in_data, success=True): + + log_data = copy.deepcopy(in_data) + server_log_name = f"{SERVER_LOG_}{name}" + + lock = acquire_lock(redis_conn, f"server_log_lock_{name}") + if lock: + try: + logger.info(f"server_log_event name:{name} action:{action} success:{success}") + server_log_str = redis_conn.get(server_log_name).decode("utf-8") + server_log = json.loads(server_log_str) + event = {} + t = time.time() + event['time'] = t + event['time_str'] = get_time(t) + event['action'] = action + + server_log['last_call'] = "success" + if 'state' in log_data: + if log_data['state'] != 0: + server_log['last_call'] = "failure" + + if not success: + server_log['last_call'] = "failure" + + if 'data' in log_data: + imgstr = log_data['data'] + imglen = len(imgstr) + if imglen > 256: + log_data['data'] = imgstr[:64] + + if 'img' in log_data: + imgstr = log_data['img'] + imglen = len(imgstr) + if imglen > 256: + log_data['img'] = imgstr[:64] + + if 'result' in log_data: + imgstr = log_data['result'] + imglen = len(imgstr) + if imglen > 256: + log_data['result'] = imgstr[:64] + + if 'user_img_path' in log_data: + imgstr = log_data['user_img_path'] + imglen = len(imgstr) + if imglen > 256: + log_data['user_img_path'] = imgstr[:32] + + event['data'] = log_data + if(len(server_log['events']) > SERVER_LOG_EVENT_LEN): + server_log['events'].pop(0) + server_log['events'].append(event) + server_log['last_event'] = event + server_log_str = json.dumps(server_log) + redis_conn.set(f"{SERVER_LOG_}{name}", server_log_str) + + finally: + lock.release() + else: + logger.error(f"get lock none {server_log_name}") + +def registerGpuServer(name, url, can_use): + try: + server_list_str = redis_conn.get(GPU_SERVER_LIST).decode("utf-8") + server_list = json.loads(server_list_str) + if name not in server_list: + server_list.append(name) + server = {} + server['name'] = name + server['url'] = url + server['can_use'] = can_use + server_str = json.dumps(server) + logger.info(f"registerGpuServer, {server_str}") + redis_conn.set(name, server_str) + + server_log = {} + server_log['name'] = name + event = {} + t = time.time() + event['time'] = t + event['time_str'] = get_time(t) + event['action'] = "register" + event['data'] = f"url {url}" + server_log['last_event'] = event + server_log['events'] = [] + server_log['last_call'] = "success" + if(len(server_log['events']) > SERVER_LOG_EVENT_LEN): + server_log['events'].pop(0) + server_log['events'].append(event) + server_log_str = json.dumps(server_log) + redis_conn.set(f"{SERVER_LOG_}{name}", server_log_str) + + redis_conn.set(GPU_SERVER_LIST, json.dumps(server_list)) + except Exception as e: + logger.error(f"registerGpuServer{e.message}") + + +def get_remote_gpu_server(): + start_time = datetime.now() + while (datetime.now() - start_time).seconds < GPU_SERVER_TIME_OUT: + lock = acquire_lock(redis_conn, SERVER_LIST_LOCK_NAME) + try: + server_list_str = redis_conn.get(GPU_SERVER_LIST).decode("utf-8") + server_list = json.loads(server_list_str) + for name in server_list: + server_str = redis_conn.get(name).decode("utf-8") + server = json.loads(server_str) + if server['can_use']: + return server + finally: + lock.release() + time.sleep(0.2) + return None + +def call_remote_gpu_server(task_data_str, server=None): + task_data = json.loads(task_data_str) + key = task_data['key'] + if server == None: + server = get_remote_gpu_server() + + result_str = '' + result_status_code = '200' + if server: + headers = { + 'Content-Type': 'application/json' + } + + if 'hairColor' in task_data['api']: + data = { + "img": task_data['request']['img'], + "rgb": task_data['request']["rgb"], + "ratio": task_data['request']['ratio'], + "userId": task_data['request']['userId'], + "output_format":task_data['request']['output_format'] + } + else: + data = { + "hair_id": task_data['request']['hair_id'], + "task_id": task_data['request']["task_id"], + "user_img_path": task_data['request']['user_img_path'], + "is_hr": task_data['request']['is_hr'], + "output_format":task_data['request']['output_format'] + } + name = server['name'] + + # set server busy + server['can_use'] = False + server_str = json.dumps(server) + redis_conn.set(name, server_str) + + server_log_event(server['name'], "Request_GpuServer", data, True) + url = server['url'] + task_data['api'] + try: + logger.info(f"call get_remote_gpu_server {url}") + start_ms = int(time.time() * 1000) # 转换为毫秒级整数 + base64 = False + if data['output_format'] == 'base64': + base64 = True + + # if 'hairColor' in task_data['api']: + # if data['img'] == 'base64': + # imgstr = redis_conn.get(f"img_{key}").decode("utf-8") + # data['img'] = imgstr + # else: + # if data['user_img_path'] == 'base64': + # imgstr = redis_conn.get(f"img_{key}").decode("utf-8") + # data['user_img_path'] = imgstr + + response = requests.post(url, headers=headers, json=data, timeout=30) + + logger.info(f"call finshed {response.status_code} {response.headers} {response.text[:128]}") + if response.status_code == 200: + result = response.json() + end_ms = int(time.time() * 1000) + if 'hairColor' in task_data['api']: + if base64: + result['result'] = f"data:image/jpeg;base64,{result['result']}" + else: + if base64: + result['data'] = f"data:image/jpeg;base64,{result['data']}" + + result['process_time_ms'] = (end_ms - start_ms) + server_log_event(server['name'], "GpuServer_Response", result, True) + result_str = json.dumps(result) + result_status_code = '200' + if 'state' in result: + if result['state'] != 0: + result_status_code = '500' + else: + try: + result = response.json() + except json.JSONDecodeError: + result = { + "state":-1, + "data":"", + 'msg': f'status_code error{response.status_code} call_remote_gpu_server{url} {response.text}' + } + time.sleep(1) + result_str = json.dumps(result) + result_status_code = str(response.status_code) + server_log_event(server['name'], f"error GpuServer_Response ", result, False) + except Exception as e: + logger.info(f"End call Exception {url} {type(e).__name__}, 错误信息: {e}") + time.sleep(1) + result = { + "state":-1, + "data":"", + 'msg': f'Exception call call_remote_gpu_server{url}' + } + result_str = json.dumps(result) + result_status_code = '500' + server_log_event(server['name'], f"error GpuServer_Response ", result, False) + + logger.info(f"End call {url}") + # release server + server['can_use'] = True + server_str = json.dumps(server) + redis_conn.set(name, server_str) + + else: + result_str = json.dumps({ + "state":-1, + "data":"", + 'msg': f'Timeout after all gpu server is busy' + }) + result_status_code = '500' + + result_key = f"result_{key}" + redis_conn.set(result_key, result_str, ex=60) + result_status_key = f"result_status_code_{key}" + redis_conn.set(result_status_key, result_status_code, ex=60) + + +def check_server_work_call(server): + task_data = {} + request = {} + request['hair_id'] = "1907651680352395265" + request['task_id'] = "1907651680352395265" + request['user_img_path'] = "https://cdn.meidaojia.com/ZoeFiles/user2_1_%E5%89%AF%E6%9C%AC.JPG" + request['output_format'] = "url" + request['is_hr'] = "false" + task_data['request'] = request + key = str(uuid.uuid4()) + task_data['key'] = key + task_data['api'] = "/api/swapHair/v1" + task_data_str = json.dumps(task_data) + threading.Thread(target=call_remote_gpu_server, args=(task_data_str, server)).start() + +def check_server_work(): + server_list_str = redis_conn.get(GPU_SERVER_LIST).decode("utf-8") + server_list = json.loads(server_list_str) + for name in server_list: + server_str = redis_conn.get(name).decode("utf-8") + server = json.loads(server_str) + if not server['can_use']: + continue + + server_log_name = f"{SERVER_LOG_}{name}" + server_log_str = redis_conn.get(server_log_name).decode("utf-8") + server_log = json.loads(server_log_str) + if time.time() - server_log['last_event']['time'] > (5*60): + logger.info("check_server_work and call name") + check_server_work_call(server) + + +def get_task_queue_key(): + redis_conn = get_redis_conn() + lock = acquire_lock(redis_conn, KEY_QUEUE_LOCK_NAME) + if not lock: + logger.error("get_task_queue_key get lock error") + return + try: + key_queue_str = redis_conn.get(QUEUE_NAME).decode("utf-8") + key_queue = json.loads(key_queue_str) + queue_len = len(key_queue) + if(queue_len > 0): + logger.info(f"get_task_queue_key queue len:{queue_len}") + key = key_queue.pop(0) + key_queue_str = json.dumps(key_queue) + logger.info(f"get_task_queue_key queue len:{len(key_queue)}") + redis_conn.set(QUEUE_NAME, key_queue_str) + return key + else: + return None + finally: + lock.release() + +def main_worker(): + while True: + try: + key = get_task_queue_key() + if key: + task_data_str = redis_conn.get(key) + if task_data_str: + threading.Thread(target=call_remote_gpu_server, args=(task_data_str,)).start() + else: + logging.error(f"main_worker get task_data_str error with key{key}") + continue + except Exception as e: + logger.info(f"main_worker {e}") + # check_server_work() + time.sleep(0.1) + +serverindex = 0 +def regServer(instance): + global serverindex + registerGpuServer(f"s{serverindex}_1_{instance}", f"https://{instance}-http-8801.northwest1.gpugeek.com:8443", True) + # registerGpuServer(f"s{serverindex}_2_{instance}", f"https://{instance}-http-8801.northwest1.gpugeek.com:8443", True) + serverindex += 1 + +def get_server_names(csv_file='servers.csv'): + """ + 从servers.csv文件中提取所有主机名称 + + 参数: + csv_file (str): CSV文件路径，默认为'servers.csv' + + 返回: + list: 包含所有主机名称的列表 + """ + server_names = [] + + try: + with open(csv_file, mode='r', encoding='utf-8') as file: + csv_reader = csv.reader(file) + for row in csv_reader: + if row: # 确保不是空行 + # 取第一列并去除前后空格 + server_name = row[0].strip() + if server_name: # 确保名称不为空 + server_names.append(server_name) + except FileNotFoundError: + print(f"错误：文件 {csv_file} 未找到") + except Exception as e: + print(f"读取文件时出错: {e}") + + return server_names + +if __name__ == '__main__': + import os + print(f"[Worker] Starting with PID: {os.getpid()}") + redis_conn.set(GPU_SERVER_LIST, "[]") + + + # servers = get_server_names() + # for s in servers: + # regServer(s) + regServer("693570664882181") + regServer("693565557084165") + + main_worker() \ No newline at end of file diff --git a/photo_service/README.md b/photo_service/README.md new file mode 100644 index 0000000..513120e --- /dev/null +++ b/photo_service/README.md @@ -0,0 +1,3 @@ +# liveme_photo_service + +liveme数字写真线上服务代码。依赖webui…… \ No newline at end of file diff --git a/photo_service/api_service.py b/photo_service/api_service.py new file mode 100644 index 0000000..8b88934 --- /dev/null +++ b/photo_service/api_service.py @@ -0,0 +1,159 @@ +import os +import sys +from webui_im2im import ControlnetRequestImg2Img +import numpy as np +import base64 +import cv2 +import os,sys +from gevent import pywsgi, monkey + +monkey.patch_all() + +# 将当前工作目录切换到当前目录 +project_dir = os.path.dirname(os.path.abspath(__file__)) +os.chdir(project_dir) +sys.path.append(project_dir) + +from flask import Flask, request, jsonify +import global_variable as global_var +import json + +app = Flask(__name__) + + +@app.route('/template/list', methods=['POST']) +def get_template_list(): + import os + import json + try: + # 读取data/template.json文件，并返回 + with open('data/template.json', 'r') as f: + template_list = json.load(f) + ret_dict = dict(code=0, message='success', data=template_list) + # 返回结果作为 JSON 响应 + return jsonify(ret_dict) + + except Exception as e: + ret_dict = dict(code=-1, message=str(e), data=[]) + return jsonify(ret_dict) + + +@app.route('/user/list', methods=['POST']) +def get_user_list(): + try: + # 获取用户文件夹路径 + user_list_dir = os.path.join(global_var.service_data_dir, 'user_data') + + ret_user_list = [] + + # 遍历查找user_list_dir目录下所有的cfg.json文件 + for root, dirs, files in os.walk(user_list_dir): + for file in files: + if file != 'cfg.json': continue + json_path = os.path.join(root, file) + lora_path = os.path.join(root, 'lora.safetensors') + if not os.path.exists(lora_path): continue + # 读取cfg.json文件 + with open(json_path, 'r') as f: + user_info = json.load(f) + user_dict = dict(user_id=user_info['user_id'], face_img_url=user_info['face_img_url']) + ret_user_list.append(user_dict) + + ret_dict = dict(code=0, message='success', data=ret_user_list) + return jsonify(ret_dict) + + except Exception as e: + ret_dict = dict(code=-1, message='error', data=[]) + return jsonify(ret_dict) + + +@app.route('/user/generate', methods=['POST']) +def generate_photo(): + try: + # 获取请求参数 + request_data = request.get_json() + #判断'user_id'和'base_img'是否在请求参数中 + assert 'user_id' in request_data and 'base_img' in request_data, 'user_id and base_img is required' + user_id = request_data['user_id'] + base_img_b64 = request_data['base_img'] + + user_path = os.path.join(global_var.service_data_dir, 'user_data', user_id) + assert os.path.isdir(user_path), 'user_id not exist' + user_lora = os.path.join(user_path, 'lora.safetensors') + usr_config_path = os.path.join(user_path, 'cfg.json') + assert os.path.exists(user_lora) and os.path.exists(usr_config_path), 'lora file or config not exist' + + # 读取用户配置文件 + with open(usr_config_path, 'r') as f: + user_info = json.load(f) + lora_md5 = user_info['lora_md5'] + dst_lora_path = os.path.join(global_var.webui_lora_dir, lora_md5 + '.safetensors') + if not os.path.exists(dst_lora_path): + os.system(f'cp {user_lora} {dst_lora_path}') + + #将模板图像转化为numpy数组 + prompt = f',easyphoto_face, easyphoto, 1person,face,suit' + neg_prompt = '(worst quality:2),(low quality:2),(normal quality:2),lowres,watermark' + + image_array = np.frombuffer(base64.b64decode(base_img_b64), np.uint8) + base_img = cv2.imdecode(image_array, cv2.IMREAD_COLOR) + # base_img = cv2.resize(base_img, (512, 512)) + # cv2.imshow('image', base_img) + # cv2.waitKey() + + # 生成图片 + control_net = ControlnetRequestImg2Img(prompt, neg_prompt) + control_net.build_body(dst_width=base_img.shape[1], dst_height=base_img.shape[0], cfg_scale=3.5, base_img=base_img) + output = control_net.send_request() + generate_photo = output['images'][0] + + # 清理硬盘空间 + os.remove(dst_lora_path) + + # # 将生成的图片转化为base64编码 + # retval, bytes = cv2.imencode('.png', generate_photo) + # generate_photo = base64.b64encode(bytes).decode('utf-8') + return jsonify(dict(code=0, message='success', generate_photo_b64=generate_photo)) + + except Exception as e: + ret_dict = dict(code=-1, message=str(e), data=[]) + return jsonify(ret_dict) + + +def webd_service(): + # 用于启动webd的后台服务 + current_file_dir = os.path.dirname(os.path.abspath(__file__)) + webd_path = os.path.join(current_file_dir, 'webd', 'webd') + + print('webd server started...') + cmd = f"{webd_path} -w {global_var.service_data_dir} -g rlT -l 10219" + os.system(cmd) + + +if __name__ == '__main__': + # 服务启动的数据目录 + global_var.service_data_dir = sys.argv[1] + + # 本地webui的lora存储目录 + global_var.webui_lora_dir = sys.argv[2] + + # webui_server_port + global_var.webui_server_port = int(sys.argv[3]) + + # server_port + global_var.server_port = int(sys.argv[4]) + + # 检查webui_lora_dir目录是否存在 + assert os.path.isdir(global_var.webui_lora_dir), 'webui_lora_dir should be a directory' + + # 检查service_data_dir目录是否存在 + if not os.path.exists(global_var.service_data_dir): + os.makedirs(global_var.service_data_dir, exist_ok=True) + else: + assert os.path.isdir(global_var.service_data_dir), 'service_data_dir should be a directory' + + # 启动服务 + # app.run(debug=False, port=global_var.server_port, host='0.0.0.0') + + server = pywsgi.WSGIServer(('0.0.0.0', global_var.server_port), app) # test port + server.serve_forever() diff --git a/photo_service/create_dir.py b/photo_service/create_dir.py new file mode 100644 index 0000000..8ad4203 --- /dev/null +++ b/photo_service/create_dir.py @@ -0,0 +1,32 @@ +import os + +if __name__ == '__main__': + in_dir = "/mnt/database2/jiangqian/0808/online_orig_train_datas_2" + out_dir = "/mnt/database2/jiangqian/0808/online_train_datas_2" + + dirs = os.listdir(in_dir) + + for single_dir in dirs: + if "traindata" in single_dir: + in_single_dir = os.path.join(in_dir, single_dir) + out_single_dir = os.path.join(out_dir, single_dir) + + if not os.path.exists(out_single_dir): + os.makedirs(out_single_dir) + + images_dir = os.path.join(out_single_dir, "images") + if not os.path.exists(images_dir): + os.makedirs(images_dir) + + hairstyle_dir = os.path.join(images_dir, "1_hairstyle") + if not os.path.exists(hairstyle_dir): + os.makedirs(hairstyle_dir) + + for root, dirs, files in os.walk(in_single_dir): + for file in files: + in_file = os.path.join(root, file) + out_file = os.path.join(hairstyle_dir, file) + os.system("cp %s %s" % (in_file, out_file)) + print("copy %s to %s" % (in_file, out_file)) + + diff --git a/photo_service/data/template.json b/photo_service/data/template.json new file mode 100644 index 0000000..030bf66 --- /dev/null +++ b/photo_service/data/template.json @@ -0,0 +1,13 @@ +[ + "http://homenas.zhourunnan.cn:9212/zhourunnan/service_data/live_photo/template_data/template01.png", + "http://homenas.zhourunnan.cn:9212/zhourunnan/service_data/live_photo/template_data/template02.png", + "http://homenas.zhourunnan.cn:9212/zhourunnan/service_data/live_photo/template_data/template03.jpg", + "http://homenas.zhourunnan.cn:9212/zhourunnan/service_data/live_photo/template_data/template04.jpg", + "http://homenas.zhourunnan.cn:9212/zhourunnan/service_data/live_photo/template_data/template05.png", + "http://homenas.zhourunnan.cn:9212/zhourunnan/service_data/live_photo/template_data/template06.png", + "http://homenas.zhourunnan.cn:9212/zhourunnan/service_data/live_photo/template_data/template07.png", + "http://homenas.zhourunnan.cn:9212/zhourunnan/service_data/live_photo/template_data/template08.jpg", + "http://homenas.zhourunnan.cn:9212/zhourunnan/service_data/live_photo/template_data/template09.png", + "http://homenas.zhourunnan.cn:9212/zhourunnan/service_data/live_photo/template_data/template10.png", + "http://homenas.zhourunnan.cn:9212/zhourunnan/service_data/live_photo/template_data/template11.png" +] \ No newline at end of file diff --git a/photo_service/data/template01.png b/photo_service/data/template01.png new file mode 100644 index 0000000..cd27bf8 Binary files /dev/null and b/photo_service/data/template01.png differ diff --git a/photo_service/enhance_hair.py b/photo_service/enhance_hair.py new file mode 100755 index 0000000..ef8934c --- /dev/null +++ b/photo_service/enhance_hair.py @@ -0,0 +1,247 @@ +import cv2 +import os +import numpy as np +import tqdm + +from utils import landmark_processor +import base64 +import requests +from PIL import Image +import io + +def encode_numpy_to_base64(img): + retval, bytes = cv2.imencode('.png', img) + encoded_image = base64.b64encode(bytes).decode('utf-8') + return encoded_image + + +def webui_img2img(img, mask, prompt=''): + url = "http://127.0.0.1:57860/sdapi/v1/img2img" + request_dict = { + "prompt": prompt, + "negative_prompt": '(nsfw:1.5), ng_deepnegative_v1_75t, (badhandv4:1.2), (worst quality:2), (low quality:2), (normal quality:2), lowres, bad anatomy, ' + 'bad hands, ((monochrome)), ((grayscale)) watermark, large breast, big breast, bad_pictures,easynegative, faceless, no human, white background, simple background, ', + "sampler_name": "DPM++ 2M Karras", + "batch_size": 1, + "steps": 20, + "width": img.shape[1], + "height": img.shape[0], + "cfg_scale": 7.0, + "seed": 123456789, + "mask_blur": 11, + "init_images": [ + encode_numpy_to_base64(img) + ], + "inpaint_full_res": False, + "inpainting_fill": 1, + "inpainting_mask_invert": 0, + "mask": encode_numpy_to_base64(mask), + # "refiner_checkpoint":"majicmixRealistic_v7.safetensors", + # "refiner_switch_at": 0.4, + "denoising_strength": 0.7, + "alwayson_scripts": { + # "controlnet": { + # "args": [ + # { + # "enabled": True, + # "module": "openpose_full", + # "model": "openpose", + # "weight": 1.0, + # # "image": self.read_image(), + # "resize_mode": "Crop and Resize", + # "low_vram": False, + # "processor_res": 512, + # "guidance_start": 0.0, + # "guidance_end": 1.0, + # "control_mode": "Balanced", + # "pixel_perfect": False + # } + # ] + # } + } + } + response = requests.post(url=url, json=request_dict) + ret_json = response.json() + result = ret_json['images'][0] + img = cv2.imdecode(np.frombuffer(base64.b64decode(result.split(",", 1)[0]), np.uint8), cv2.IMREAD_COLOR) + return img + +def webui_super_res_img(img, ratio): + url = "http://127.0.0.1:57860/sdapi/v1/extra-single-image" + request_dict = { + "resize_mode": 0, + "show_extras_results": False, + "gfpgan_visibility": 0, + "codeformer_visibility": 1, + "codeformer_weight": 1, + "upscaling_resize": ratio, + "upscaler_1": "8x_NMKD-Superscale_150000_G", + "upscale_first": False, + "image": encode_numpy_to_base64(img) + } + response = requests.post(url=url, json=request_dict) + ret_json = response.json() + result = ret_json['image'] + img = cv2.imdecode(np.frombuffer(base64.b64decode(result), np.uint8), cv2.IMREAD_COLOR) + return img + + +def webui_tag_by_clip(img): + url = "http://127.0.0.1:57860/sdapi/v1/interrogate" + request_dict = { + "image": encode_numpy_to_base64(img), + "model": "clip" + } + response = requests.post(url=url, json=request_dict) + ret_json = response.json() + return ret_json['caption'] + + + + +if __name__ == '__main__': + hair_dir = '/mnt/database2/jiangqian/0808/exp2-data-zrn-0808/1816523294655647746' + + data2process_list = [] + + # 遍历查找hair_dir下的所有npy文件 + for root, dirs, files in os.walk(hair_dir): + for file in files: + if file.endswith('.npy'): + # 关键点路径 + pt1k_path = os.path.join(root, file) + origin_img_path = pt1k_path[:-4] + '.png' + origin_matting_path = pt1k_path[:-4] + '_origin_matting.png' + new_matting_path = pt1k_path[:-4] + '_new_matting.png' + result_img_path = pt1k_path[:-4] + '_res.png' + # 判断上面的文件是否存在 + if (not os.path.exists(origin_img_path) or not os.path.exists(origin_matting_path) + or not os.path.exists(new_matting_path) or not os.path.exists(result_img_path)): + continue + + ref_hair_path = pt1k_path[:-4] + '_orig_hair.png' + # if not os.path.exists(ref_hair_path): + # continue + lora_model_path = pt1k_path[:-4] + '_hairstyle_lora.safetensors' + # if not os.path.exists(lora_model_path): + # continue + data2process_list.append([pt1k_path, origin_img_path, origin_matting_path, + new_matting_path, result_img_path, ref_hair_path, lora_model_path]) + + crop_size = 768 + + webui_lora_dir = '/home/student/Documents/workspace_cxt_tianjing_hair/miaoya/webui_home/stable-diffusion-webui/models/Lora' + for pt1k_path, origin_img_path, origin_matting_path, new_matting_path, result_img_path, ref_hair_path, lora_model_path in tqdm.tqdm(data2process_list): + # if '508417f3-2c71-45cf-a75b-969b27ec7d8f' not in pt1k_path: continue + # 读取关键点 + pt1k = np.load(pt1k_path) + # 读取原图 + origin_img = cv2.imread(origin_img_path) + tmp_scale = 1920 / max(origin_img.shape[0], origin_img.shape[1]) + if tmp_scale < 1.0: + origin_img = cv2.resize(origin_img, (0, 0), fx=tmp_scale, fy=tmp_scale, interpolation=cv2.INTER_LANCZOS4) + + print("origin_img shape:", origin_img.shape) + # cv2.imshow("origin_img", origin_img) + + # 读取原图抠图 + origin_matting = cv2.imread(origin_matting_path, cv2.IMREAD_GRAYSCALE) + # 读取新图抠图 + new_matting = cv2.imread(new_matting_path, cv2.IMREAD_GRAYSCALE) + # 读取结果图 + result_img = cv2.imread(result_img_path) + print("result_img shape:", result_img.shape) + # cv2.imshow("result_img", result_img) + # # 读取参考头发 + # ref_hair = cv2.imread(ref_hair_path) + # if max(ref_hair.shape[:2]) < 300: continue + + + # 如何图像不清晰，进行超分辨率处理 + # if max(origin_img.shape[:2]) < 1500: + # scale_ratio = 2000 / max(origin_img.shape[:2]) + # result_img = webui_super_res_img(result_img, scale_ratio) + # origin_img = cv2.resize(origin_img, (result_img.shape[1], result_img.shape[0]), + # interpolation=cv2.INTER_LANCZOS4) + # origin_matting = cv2.resize(origin_matting, (result_img.shape[1], result_img.shape[0])) + # new_matting = cv2.resize(new_matting, (result_img.shape[1], result_img.shape[0])) + # pt1k = pt1k * scale_ratio + + # 获取头发处理的局部区域图像 + # M = landmark_processor.get_transform_mat_hair_ratio_v1(pt1k, crop_size, ratio=0.30, h_offset=0.32) + + scale = 768 / max(origin_img.shape[0], origin_img.shape[1]) + M = cv2.getRotationMatrix2D((0, 0), 0, scale) + + dst_size = (int(origin_img.shape[1] * scale), int(origin_img.shape[0] * scale)) + + # 高质量的从原图中截取头发区域 + crop_origin = landmark_processor.high_quality_warpAffine(origin_img, M, dst_size) + cv2.imwrite("./crop_origin.png", crop_origin) + + crop_result = landmark_processor.high_quality_warpAffine(result_img, M, dst_size) + cv2.imwrite("./crop_result.png", crop_result) + # tmp = cv2.warpAffine(origin_img, M, dst_size, flags=cv2.INTER_AREA) + + # 构造重绘的mask + matting_merge = np.concatenate([origin_matting[:,:, np.newaxis], new_matting[:,:, np.newaxis]], axis=2) + matting_merge = np.max(matting_merge, axis=2) + # matting_merge = new_matting + crop_matting = cv2.warpAffine(matting_merge, M, dst_size) + mask = (crop_matting > 10).astype(np.float32) + mask_dilate = cv2.dilate(mask, np.ones((3, 11), np.uint8)) + final_img = crop_result + + mask_dilate = np.clip(mask_dilate * 255, 0, 255).astype(np.uint8) + + # file_name = os.path.basename(pt1k_path)[:-4] + # save_dir = '/home/chinatszrn/Downloads/abc/ref_hair/dst_res/style3_tmp' + # cv2.imwrite(os.path.join(save_dir, file_name + '.png'), final_img) + # cv2.imwrite(os.path.join(save_dir, file_name + '_mask.png'), mask_dilate) + # # cv2.imwrite(os.path.join(save_dir, file_name + '_ref_hair.png'), ref_hair) + # continue + + # # 拷贝lora模型 + # os.system('cp {} {}'.format(lora_model_path, webui_lora_dir)) + # # 构建prompt提示词 + # lora_model_name = os.path.basename(pt1k_path)[:-4] + + # 对final_img进行打标 + # tag_result = webui_tag_by_clip(final_img) + tag_result = '' + + # 开始重绘 + prompt = f' titor hairstyle, easyphoto, ' + tag_result + # 对发型进行重绘 + + # cv2.imshow("final_img_0", final_img) + # cv2.imshow("mask_dilate", mask_dilate) + # cv2.waitKey(100) + + sd_result = webui_img2img(final_img, mask_dilate, prompt) + + final_img = origin_img.copy() + # 将重绘结果恢复到原图 + M_inv = cv2.invertAffineTransform(M) + cv2.warpAffine(sd_result, M_inv, (final_img.shape[1], final_img.shape[0]), dst=final_img, + borderMode=cv2.BORDER_TRANSPARENT, flags=cv2.INTER_LANCZOS4) + + # cv2.imshow("final_img", final_img) + # cv2.waitKey(0) + cv2.imwrite(result_img_path[:-4] + '_sd.png', final_img) + + # ref_hair_scale = final_img.shape[0] / ref_hair.shape[0] + # ref_hair = cv2.resize(ref_hair, (0, 0), fx=ref_hair_scale, fy=ref_hair_scale, interpolation=cv2.INTER_LANCZOS4) + # img2show = np.concatenate([origin_img, ref_hair, final_img], axis=1) + # + # # 显示结果 + # save_dir = '/home/chinatszrn/Downloads/exp' + # cv2.imwrite(os.path.join(save_dir, lora_model_name + '.png'), img2show) + # # cv2.imshow("origin_img", cv2.resize(img2show, (0, 0), fx=0.3, fy=0.3, interpolation=cv2.INTER_AREA)) + # cv2.imshow("sd_result", cv2.resize(sd_result, (0, 0), fx=0.3, fy=0.3, interpolation=cv2.INTER_AREA)) + # cv2.waitKey(1000) + + + + + diff --git a/photo_service/global_variable.py b/photo_service/global_variable.py new file mode 100644 index 0000000..b6d5fee --- /dev/null +++ b/photo_service/global_variable.py @@ -0,0 +1,34 @@ +# 服务数据目录 +service_data_dir = None + +# 生成视频的thread_num +video_generation_thread_num = None + +# avatar训练的thread_num +avatar_train_thread_num = None + +# 视频生成的任务队列 +video_generation_task_sq = None + +# avatar训练的任务队列 +avatar_train_task_sq = None + +# webui的lora存储目录 +webui_lora_dir = None + +# webui_server_port +webui_server_port = None + +# server ip address +server_ip = None + +# service port +server_port = 10239 + +# 是否启动视频的硬件编码 +video_hardware_encode = True + +#---------------------------------------------- + +# 发型lora训练队列 +hair_style_lora_train_task_sq = None diff --git a/photo_service/gpt4v_caption.py b/photo_service/gpt4v_caption.py new file mode 100755 index 0000000..b0d01ac --- /dev/null +++ b/photo_service/gpt4v_caption.py @@ -0,0 +1,68 @@ +import base64 +import time + +import requests +import cv2 +import json +import re +import os +import tqdm + +# OpenAI API Key +api_key = "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" + +# Function to encode the image +def encode_image(image_path): + img = cv2.imread(image_path) + scale = 500.0 / min(img.shape[:2]) + img = cv2.resize(img, (0, 0), fx=scale, fy=scale) + scaled_path = '/tmp/scaled_image.jpg' + cv2.imwrite(scaled_path, img) + with open(scaled_path, "rb") as image_file: + return base64.b64encode(image_file.read()).decode('utf-8') + +def caption_image(image_path): + # Getting the base64 string + base64_image = encode_image(image_path) + + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {api_key}" + } + + payload = { + "model": "gpt-4-vision-preview", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "As an AI image tagging expert, please provide precise tags for the hairstyle in the image, To enhance CLIP model's understanding of the content. Please provide a detailed description of the hairstyle in the image, including but not limited to the color, style, length, curliness, hairline, highlights, gradients, etc. Your tags should be accurate, non-duplicative, and within a 10-20 word count range. Tags should be comma-separated. No need to provide any safety statements or precautions." + }, + { + "type": "image_url", + "image_url": { + "url": f"data:image/jpeg;base64,{base64_image}", + "detail": "low" + } + } + ] + } + ], + "max_tokens": 300 + } + + try: + response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload) + tmp_json = response.json() + return tmp_json['choices'][0]['message']['content'] + except Exception as e: + print(e) + return "" + + +if __name__ == '__main__': + prompt = caption_image('/home/chinatszrn/Downloads/abc/train_data/style1/07ebac82-4c0f-4dd2-84bd-bc34a059bd9b.png') + print(prompt) + diff --git a/photo_service/gradio_demo.py b/photo_service/gradio_demo.py new file mode 100644 index 0000000..87a1787 --- /dev/null +++ b/photo_service/gradio_demo.py @@ -0,0 +1,169 @@ +import gradio as gr +import os,re +import numpy as np +import requests +import cv2 +import base64 +import json +from io import BytesIO +from PIL import Image + +api_service_url = 'http://127.0.0.1:1234' +# api_service_url = 'http://i-2.gpushare.com:53412' + +class ChangeFaceGui(): + def __init__(self): + self.template_list = self.get_template_list() + self.user_dict = self.get_user_dict() + # self.tmp_dir = './tmp' + # if not os.path.exists(self.tmp_dir): os.makedirs(self.tmp_dir) + self.user_img_list = [] + # 获取户图片 + for item in self.user_dict: + response = requests.get(item['face_img_url']) + image_data = BytesIO(response.content) + user_face_img = cv2.imdecode(np.frombuffer(image_data.read(), np.uint8), cv2.IMREAD_COLOR) + self.user_img_list.append(user_face_img) + item['face_img'] = user_face_img + + # img_path = f'{self.tmp_dir}/{item["user_id"]}.jpg' + # if not os.path.exists(img_path): + # cv2.imwrite(f'{self.tmp_dir}/{item["user_id"]}.jpg', user_face_img) + + + #请求得到模板图片的url + def get_template_list(self): + url = f"{api_service_url}/template/list" + payload = {} + headers = {} + response = requests.request("POST", url, headers=headers, data=payload) + print('请求模板照片成功！！') + return(response.json()['data']) + + + #请求得到用户图片及ID + def get_user_dict(self): + url = f"{api_service_url}/user/list" + payload = {} + headers = {} + response = requests.request("POST", url, headers=headers, data=payload) + user_dict = response.json()['data'] + print('请求用户照片成功！！') + return(user_dict) + + + #获取用户ID + def get_user_id(self, user_img): + user_small_img = cv2.resize(user_img, (100, 100)) + mse = [] + for item in self.user_dict: + user_face_img = item['face_img'] + # cv2.imwrite(f'{self.tmp_dir}/{item["user_id"]}_list.jpg', user_face_img) + user_face_img = cv2.resize(user_face_img, (100, 100)) + # 计算均方误差 + mse.append(np.mean((user_small_img - user_face_img) ** 2)) + + # 找到均方差最小值对应的ID + user_id = mse.index(min(mse)) + user_id = self.user_dict[user_id]['user_id'] + print('用户ID:', user_id) + return user_id + + + #虚拟试穿模块，输入是模特图和衣服图，输出是虚拟试穿的结果 + def take_photo(self, template_img, user_img): + # if not os.path.exists(self.tmp_dir):os.makedirs(self.tmp_dir) + + if template_img is None or user_img is None: + return None + + # 使用 Pillow 加载图像 + template_img = Image.open(template_img) + user_img = Image.open(user_img) + + # 将 Pillow 图像转换为 OpenCV 格式（BGR） + template_img = cv2.cvtColor(np.array(template_img), cv2.COLOR_RGB2BGR) + user_img = cv2.cvtColor(np.array(user_img), cv2.COLOR_RGB2BGR) + + # 将模板图片转换为base64格式 + retval, template_bytes = cv2.imencode('.jpg', template_img) + encoded_image = base64.b64encode(template_bytes).decode('utf-8') + user_id = self.get_user_id(user_img) + url = f"{api_service_url}/user/generate" + payload = json.dumps({ + "user_id": user_id, + "base_img": encoded_image + }) + headers = { + 'Content-Type': 'application/json' + } + print('请求api_service.py发送请求！！') + response = requests.request("POST", url, headers=headers, data=payload) + print('请求api_service.py发送请求成功！！') + if response.status_code != 200: + raise RuntimeError(f"Failed to send request to API service. Status code: {response.status_code}") + + ret_image_b64 = response.json().get('generate_photo_b64') + + if ret_image_b64 is None: + raise RuntimeError(f"ret image failed!") + + image_array = np.frombuffer(base64.b64decode(ret_image_b64), np.uint8) + result_image = cv2.imdecode(image_array, cv2.IMREAD_COLOR) + result_image = cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB) + return result_image + + + def start_gui(self): + user_img_list = [] + for item in self.user_img_list: + img = cv2.cvtColor(item, cv2.COLOR_BGR2RGB) + user_img_list.append(img) + + with gr.Blocks() as demo: + with gr.Row(): + gr.Markdown("# 数字力场效果展示") + + # 换脸 + with gr.Tab("数字写真"): + with gr.Row(): + gr.Markdown("# 数字写真") + with gr.Row(): + with gr.Column(): + #选择模板 + template_img = gr.Image(label="模版", sources='upload', min_width=384, width=384, height=384, type="filepath", value=self.template_list[0],interactive=True) + example_template = gr.Examples( + inputs=template_img, + examples_per_page=12, + examples= self.template_list) + with gr.Column(): + # 选择用户 + user_img = gr.Image(label="用户", sources='upload', min_width=384, width=384, height=384, type="filepath", value= user_img_list[0],interactive=False) + example_user = gr.Examples( + inputs=user_img, + examples_per_page=12, + examples=user_img_list) + + with gr.Column(): + output_img = gr.Image(label="结果展示", type="numpy", height=576, width=384) + with gr.Column(): + run_button = gr.Button(value="提交") + run_button.click(self.take_photo, inputs=[template_img, user_img], outputs=[output_img]) + + # #换衣服 + # with gr.Tab("换衣"): + # with gr.Row(): + # gr.Markdown("# 换衣") + # text_button = gr.Button("提交") + # # 换发型 + # with gr.Tab("换发型"): + # with gr.Row(): + # gr.Markdown("# 换发型") + # text_button = gr.Button("提交") + + demo.launch(server_name='0.0.0.0', server_port=8080) + + +if __name__ == '__main__': + demo = ChangeFaceGui() + demo.start_gui() diff --git a/photo_service/gradio_hair.py b/photo_service/gradio_hair.py new file mode 100644 index 0000000..1060376 --- /dev/null +++ b/photo_service/gradio_hair.py @@ -0,0 +1,102 @@ +import gradio as gr +import os,re +import numpy as np +import requests +import cv2 +import base64 +import json +from io import BytesIO +from PIL import Image + +# api_service_url = 'http://127.0.0.1:1234' +# api_service_url = 'http://i-2.gpushare.com:53412' +api_service_url = 'http://service.aicloud.fit:7393/api/hairStyle/v1' + +class ChangeHairGui(): + # def __init__(self): + # a = 1 + def change_hair(self, user_img, hair_img): + + if user_img is None or user_img is None: + return None + # + # if user_img.shape != hair_img.shape: + # hair_img = cv2.resize(hair_img, (user_img.shape[1], user_img.shape[0])) + # + # alpha = 0.5 # 图像1的权重 + # beta = 0.5 # 图像2的权重 + # gamma = 0 # 亮度调整常量（通常为0） + # + # result_image = cv2.addWeighted(user_img, alpha, hair_img, beta, gamma) + + + # 将 Pillow 图像转换为 OpenCV 格式（BGR） + user_img = cv2.cvtColor(np.array(user_img), cv2.COLOR_RGB2BGR) + hair_img = cv2.cvtColor(np.array(hair_img), cv2.COLOR_RGB2BGR) + + # 将模板图片转换为base64格式 + retval, user_bytes = cv2.imencode('.jpg', user_img) + encoded_user_image = base64.b64encode(user_bytes).decode('utf-8') + retval, hair_bytes = cv2.imencode('.jpg', hair_img) + encoded_hair_image = base64.b64encode(hair_bytes).decode('utf-8') + + url = api_service_url + + # 请求换发型接口 + payload = json.dumps({ + "user_img_base64": encoded_user_image, + "hair_ref_img_base64": encoded_hair_image + }) + headers = { + 'Content-Type': 'application/json' + } + print('请求api_service.py发送请求！！') + response = requests.request("POST", url, headers=headers, data=payload) + print('请求api_service.py发送请求成功！！') + if response.status_code != 200: + raise RuntimeError(f"Failed to send request to API service. Status code: {response.status_code}") + + ret_image_b64 = response.json().get('result') + + if ret_image_b64 is None: + raise RuntimeError(f"ret image failed!") + + image_array = np.frombuffer(base64.b64decode(ret_image_b64), np.uint8) + result_image = cv2.imdecode(image_array, cv2.IMREAD_COLOR) + result_image = cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB) + + dst_size = max(result_image.shape[0], result_image.shape[1]) + M = cv2.getRotationMatrix2D((result_image.shape[1] / 2, result_image.shape[0] / 2), 0, 1) + M[:, 2] += np.float32([dst_size / 2 - result_image.shape[1] / 2, dst_size / 2 - result_image.shape[0] / 2]) + result_image = cv2.warpAffine(result_image, M, (dst_size, dst_size), borderValue=(255, 255, 255)) + + return result_image + + + + def start_gui(self): + with gr.Blocks() as demo: + with gr.Row(): + gr.Markdown("# 数字力场换发型效果展示") + + # 换脸 + with gr.Tab("换发型"): + with gr.Row(): + gr.Markdown("#换发型") + with gr.Row(): + with gr.Column(): + user_img = gr.Image(label="请上传用户图片", type="numpy", height=384, width=384) + with gr.Column(): + hair_img = gr.Image(label="请上传发型图片", type="numpy", height=384, width=384) + with gr.Column(): + output_img = gr.Image(label="结果展示", type="numpy", height=384, width=384, format='png') + with gr.Column(): + run_button = gr.Button(value="提交") + run_button.click(self.change_hair, inputs=[user_img, hair_img], outputs=[output_img]) + + demo.launch(server_name='0.0.0.0', server_port=8080) + +if __name__ == '__main__': + + demo = ChangeHairGui() + demo.start_gui() diff --git a/photo_service/local_lora_train.py b/photo_service/local_lora_train.py new file mode 100644 index 0000000..083bcab --- /dev/null +++ b/photo_service/local_lora_train.py @@ -0,0 +1,120 @@ +import os +import sys +import uuid + +from webui_im2im import ControlnetRequestImg2Img +import numpy as np +import base64 +import cv2 +import os, sys +from gevent import pywsgi, monkey +from multiprocessing import Process, Queue +import glob + +# from gpt4v_caption import caption_image + +# monkey.patch_all() +# sys.setrecursionlimit(20000) + +# 将当前工作目录切换到当前目录 +project_dir = os.path.dirname(os.path.abspath(__file__)) +os.chdir(project_dir) +sys.path.append(project_dir) + + +import global_variable as global_var +import json + + + +def train_hair_lora(): + import os + import json + + try: + hair_train_dir = "/mnt/database2/jiangqian/0808/online_train_datas_2" + hair_material_dir_list = os.listdir(hair_train_dir) + + for single_hair_material in hair_material_dir_list: + # 获取请求参数 + task_id = str(uuid.uuid4()) + + request_data = {} + hair_material_dir = os.path.join(hair_train_dir, single_hair_material) + print("--------------------hair_material_dir:", hair_material_dir) + request_data['hair_material_dir'] = hair_material_dir + + img_dir = os.path.join(hair_material_dir, 'images') + + model_dir = os.path.join(hair_material_dir, 'model') + if not os.path.exists(model_dir): + os.makedirs(model_dir) + + #判断img_dir下面是否只有一个文件夹 + img_dir_list = os.listdir(img_dir) + train_image_dir = os.path.join(img_dir, img_dir_list[0]) + + #判断文件夹下面是否有图片 + request_data['train_image_dir'] = train_image_dir + + # 将请求数据放入队列 + train_thread(request_data) + + # 返回结果 + print("头发lora训练开始") + print("\n\n\n\n") + + except Exception as e: + print(e) + +def train_thread(task_dict): + try: + hair_material_dir = task_dict['hair_material_dir'] + images_dir = os.path.join(hair_material_dir, 'images') + model_dir = os.path.join(hair_material_dir, 'model') + train_image_dir = task_dict['train_image_dir'] + tag = "" + + sample_dir = os.path.join(model_dir, 'sample') + if not os.path.exists(sample_dir): + os.makedirs(sample_dir) + sample_txt = os.path.join(sample_dir, 'prompt.txt') + with open(sample_txt, 'w') as f_s: + f_s.write('titor hairstyle, easyphoto, faceless, no human, white background, simple background, ' + + tag + + ' --n low quality, worst quality, bad anatomy, bad composition, poor, low effort --h 512 ' + '--w 512 --s 30 --l 7') + + #训练头发lora + cmd_train = ( + 'docker run --rm --gpus all -v /home/student/Documents/workspace_cxt_tianjing_hair/miaoya/kohya_ss_home:/home/chinatszrn -v ' + '/mnt:/mnt -e PATH=/home/chinatszrn/.local/bin -w /home/chinatszrn/kohya_ss ' + '--net=host chinatszrn/ubuntu:kohya_ss accelerate launch --num_cpu_threads_per_process=2 "./train_network.py" --enable_bucket ' + '--min_bucket_reso=256 --max_bucket_reso=1800 --pretrained_model_name_or_path="/mnt/nas_hdd/米亚像馆/models/Stable-diffusion/majicmixRealistic_v7.safetensors" ' + f'--train_data_dir={images_dir} --resolution="768,768" ' + f'--output_dir={model_dir} ' + '--network_alpha="64" --save_model_as=safetensors --network_module=networks.lora --text_encoder_lr=5e-05 ' + '--unet_lr=0.0001 --network_dim=128 --output_name="hairstyle_lora" --lr_scheduler_num_cycles="20" ' + '--no_half_vae --learning_rate="0.0001" --lr_scheduler="cosine" --lr_warmup_steps="650" --train_batch_size="1" ' + '--max_train_steps="2000" --save_every_n_epochs="100" --mixed_precision="fp16" --save_precision="fp16" ' + '--caption_extension=".txt" --sample_sampler=ddim ' + f'--sample_prompts={sample_txt} --sample_every_n_epochs="1" ' + '--seed="1234" --cache_latents --optimizer_type="AdamW8bit" --max_data_loader_n_workers="0" --bucket_reso_steps=64 ' + '--xformers --bucket_no_upscale --noise_offset=0.0 --tokenizer_cache_dir="/home/chinatszrn/.cache/clip"') + print("cmd_train:", cmd_train) + os.system(cmd_train) + + lora_path = os.path.join(model_dir, 'hairstyle_lora.safetensors') + + except Exception as e: + print(e) + return + + +if __name__ == '__main__': + global_var.webui_lora_dir = '/home/student/Documents/workspace_cxt_tianjing_hair/miaoya/webui_home/stable-diffusion-webui/models/Lora' + + train_hair_lora() + + + diff --git a/photo_service/lora_train_service.py b/photo_service/lora_train_service.py new file mode 100644 index 0000000..90b87c4 --- /dev/null +++ b/photo_service/lora_train_service.py @@ -0,0 +1,315 @@ +import os +import sys +import uuid + +from webui_im2im import ControlnetRequestImg2Img +import numpy as np +import base64 +import cv2 +import os, sys +from gevent import pywsgi, monkey +from multiprocessing import Process, Queue +import glob + +# from gpt4v_caption import caption_image + +# monkey.patch_all() +# sys.setrecursionlimit(20000) + +# 将当前工作目录切换到当前目录 +project_dir = os.path.dirname(os.path.abspath(__file__)) +os.chdir(project_dir) +sys.path.append(project_dir) + +from flask import Flask, request, jsonify +import global_variable as global_var +import json + +app = Flask(__name__) + + +def send_request(url, state, task_dict): + import requests + import json + + msg = '头发lora训练失败' if state == -1 else '头发lora训练成功' + payload = json.dumps({ + "task_id": task_dict['task_id'] if task_dict is not None else '', + "hair_id": task_dict['hair_id'] if task_dict is not None else '', + "state": state, + "msg": msg, + "is_tj": task_dict['is_tj'] + }) + + print("payload:", payload) + + headers = { + 'Content-Type': 'application/json' + } + + requests.request("POST", url, headers=headers, data=payload) + + +@app.route('/api/hair/train', methods=['POST']) +def train_hair_lora(): + import os + import json + task_id = '' + try: + # 获取请求参数 + request_data = request.get_json() + assert 'task_id' in request_data, 'task_id is required' + task_id = request_data['task_id'] + assert 'hair_id' in request_data and 'hair_material_dir' in request_data, 'hair_id and hair_material_dir is required' + hair_material_dir = request_data['hair_material_dir'] + assert os.path.exists(hair_material_dir), f'{hair_material_dir} not exists' + img_dir = os.path.join(hair_material_dir, 'images') + assert os.path.exists(img_dir), f'{img_dir} not exists' + assert os.path.exists( + os.path.join(hair_material_dir, 'model')), f'{os.path.join(hair_material_dir, "model")} not exists' + + #判断img_dir下面是否只有一个文件夹 + img_dir_list = os.listdir(img_dir) + assert len(img_dir_list) == 1, f'{img_dir}下面只能有一个文件夹' + train_image_dir = os.path.join(img_dir, img_dir_list[0]) + assert os.path.isdir(train_image_dir), f'{train_image_dir}不是文件夹' + # 检查该文件夹是否以数字加下划线开头 + assert img_dir_list[0].split('_')[0].isdigit(), f'{img_dir_list[0]}不是数字开头' + + #判断文件夹下面是否有图片 + img_list = glob.glob(train_image_dir + '/*.png') + assert len(img_list) > 10, f'{os.path.join(img_dir, img_dir_list[0])}下面图片数量小于10张' + request_data['train_image_dir'] = train_image_dir + + # is tianjin company + try: + request_data['is_tj'] = request_data['is_tj'] + except Exception as e: + print(e) + request_data['is_tj'] = '0' + + # 将请求数据放入队列 + global_var.hair_style_lora_train_task_sq.put(request_data) + + # 返回结果 + ret_dict = dict(state=0, msg='头发lora训练开始', task_id=task_id) + return jsonify(ret_dict) + except Exception as e: + ret_dict = dict(state=-1, msg=str(e), task_id=task_id) + return jsonify(ret_dict) + + +@app.route('/api/hair/inference', methods=['POST']) +def inference_webui(): + import os + import requests + import time + url = "http://127.0.0.1:57860/sdapi/v1/img2img" + onediff_url = "http://127.0.0.1:9038/sdapi/v1/img2img" + + try: + t1 = time.time() + # 获取请求参数 + request_data = request.get_json() + # with open('request_data.json', 'w') as f: + # json.dump(request_data, f) + request_json = request_data.get('request_json') + # hd_version_flag = request_data.get('hd_version_flag', False) + hair_id = request_data['hair_id'] + selected_url = url + if 'refiner_checkpoint' not in request_json: + selected_url = onediff_url + request_json['script_name'] = 'onediff_diffusion_model' + print('user onediff') + + + + # if hd_version_flag: + # print('user hd version!!!!!') + # lora_file_name = 'hairstyle_hd_lora.safetensors' + # else: + # lora_file_name = 'hairstyle_lora.safetensors' + + hair_material_dir = request_data.get('hair_material_dir') + print(hair_material_dir) + assert os.path.exists(hair_material_dir), f'{hair_material_dir} not exists' + + hd_version_flag = True + lora_file_name = 'hairstyle_hd_lora.safetensors' + request_lora_path = os.path.join(hair_material_dir, 'model', lora_file_name) + if not os.path.exists(request_lora_path): + print('use low resolution lora') + hd_version_flag = False + lora_file_name = 'hairstyle_lora.safetensors' + request_lora_path = os.path.join(hair_material_dir, 'model', lora_file_name) + + assert os.path.exists(request_lora_path), f'{request_lora_path} not exists' + + if hd_version_flag: + tmp_lora_name = f'{hair_id}_hd' + else: + tmp_lora_name = f'{hair_id}' + lora_dst_path = os.path.join(global_var.webui_lora_dir, f'{tmp_lora_name}.safetensors') + if not os.path.exists(lora_dst_path): + os.system('cp {} {}'.format(request_lora_path, lora_dst_path)) + request_json['prompt'] = f', titor hairstyle, ' + request_json['prompt'] + request_json['negative_prompt'] = request_json['negative_prompt'] + ', faceless, no human' + print('pre_stage cost:', time.time() - t1) + + start_time = time.time() + response = requests.post(url=selected_url, json=request_json) + print('inference time:', time.time() - start_time) + + ret_json = response.json() + # if os.path.exists(lora_dst_path): + # os.remove(lora_dst_path) + # 返回结果 + return jsonify(ret_json) + except Exception as e: + ret_dict = dict(state=-1, msg=str(e)) + return jsonify(ret_dict) + + +@app.route('/api/hair/inference_diy', methods=['POST']) +def inference_diy_webui(): + import os + import requests + import time + url = "http://127.0.0.1:57860/sdapi/v1/img2img" + onediff_url = "http://127.0.0.1:9038/sdapi/v1/img2img" + + try: + t1 = time.time() + # 获取请求参数 + request_data = request.get_json() + # with open('request_data.json', 'w') as f: + # json.dump(request_data, f) + request_json = request_data.get('request_json') + + selected_url = url + if 'refiner_checkpoint' not in request_json: + selected_url = onediff_url + request_json['script_name'] = 'onediff_diffusion_model' + print('user onediff') + + request_json['prompt'] = f'titor hairstyle, ' + request_json['prompt'] + request_json['negative_prompt'] = request_json['negative_prompt'] + ', faceless, no human' + print('pre_stage cost:', time.time() - t1) + + start_time = time.time() + response = requests.post(url=selected_url, json=request_json) + print('inference time:', time.time() - start_time) + + ret_json = response.json() + # if os.path.exists(lora_dst_path): + # os.remove(lora_dst_path) + # 返回结果 + return jsonify(ret_json) + except Exception as e: + ret_dict = dict(state=-1, msg=str(e)) + return jsonify(ret_dict) + +def train_thread(sq, gpu_id): + while True: + url = 'http://service.aicloud.fit:7393/api/hair/trainCallBack' + task_dict = None + try: + task_dict = sq.get() + hair_material_dir = task_dict['hair_material_dir'] + images_dir = os.path.join(hair_material_dir, 'images') + model_dir = os.path.join(hair_material_dir, 'model') + train_image_dir = task_dict['train_image_dir'] + # tag = task_dict['tag'] + tag = "" + is_tj = task_dict['is_tj'] + + sample_dir = os.path.join(model_dir, 'sample') + if not os.path.exists(sample_dir): + os.makedirs(sample_dir) + sample_txt = os.path.join(sample_dir, 'prompt.txt') + with open(sample_txt, 'w') as f_s: + f_s.write('titor hairstyle, easyphoto, faceless, no human, white background, simple background, ' + + tag + + ' --n low quality, worst quality, bad anatomy, bad composition, poor, low effort --h 768 ' + '--w 768 --s 30 --l 7') + + # 给训练图片打标签 + # cmd_caption = ('docker run --rm --gpus all -v /home/chinatszrn/Documents/miaoya/kohya_ss_home:/home/chinatszrn ' + # '-v /mnt:/mnt -e PATH=/home/chinatszrn/.local/bin -w /home/chinatszrn/kohya_ss ' + # '--net=host chinatszrn/ubuntu:kohya_ss accelerate ' + # 'launch "./finetune/tag_images_by_wd14_tagger.py" --batch_size=2 ' + # '--general_threshold=0.5 --character_threshold=0.5 --caption_extension=".txt" ' + # '--model="SmilingWolf/wd-v1-4-convnextv2-tagger-v2" --max_data_loader_n_workers="2" ' + # '--debug --remove_underscore --frequency_tags --undesired_tags="1girl, 1boy" ' + # f'"{train_image_dir}"') + # os.system(cmd_caption) + + img_path_list = glob.glob(train_image_dir + '/*.png') + # 给训练图片打标签, gpt + for img_path in img_path_list: + # print(img_path) + txt_path = img_path[:-4] + '.txt' + + print("tag img_path: ", img_path) + # tags = caption_image(img_path) + tags = tag + with open(txt_path, 'w') as f: + f.write('titor hairstyle, easyphoto, faceless, no human, white background, simple background, ' + tags) + + #判断是否每个训练图片都有标签文件 + # for img_path in img_path_list: + # assert os.path.exists(img_path[:-4]+'.txt'), f'{img_path}没有对应的标签文件' + # with open(img_path[:-4]+'.txt', 'r') as f: + # tags = f.readline() + # with open(img_path[:-4]+'.txt', 'w') as f: + # f.write('titor hairstyle, faceless, no human, gray background, simple background, ' + tags) + + #训练头发lora + cmd_train = ( + 'docker run --rm --gpus "device=0" -v /home/student/Documents/workspace_cxt_tianjing_hair/miaoya/kohya_ss_home:/home/chinatszrn -v ' + '/mnt:/mnt -e PATH=/home/chinatszrn/.local/bin -w /home/chinatszrn/kohya_ss ' + '--net=host chinatszrn/ubuntu:kohya_ss accelerate launch --num_cpu_threads_per_process=2 "./train_network.py" --enable_bucket ' + '--min_bucket_reso=256 --max_bucket_reso=2048 --pretrained_model_name_or_path="/mnt/nas_hdd/米亚像馆/models/Stable-diffusion/majicmixRealistic_v7.safetensors" ' + f'--train_data_dir={images_dir} --resolution="2000,2000" ' + f'--output_dir={model_dir} ' + '--network_alpha="64" --save_model_as=safetensors --network_module=networks.lora --text_encoder_lr=5e-05 ' + '--unet_lr=0.0001 --network_dim=128 --output_name="hairstyle_hd_lora" --lr_scheduler_num_cycles="20" ' + '--no_half_vae --learning_rate="0.0001" --lr_scheduler="cosine" --lr_warmup_steps="650" --train_batch_size="1" ' + '--max_train_steps="1500" --save_every_n_epochs="100" --mixed_precision="fp16" --save_precision="fp16" ' + '--caption_extension=".txt" --sample_sampler=ddim ' + f'--sample_prompts={sample_txt} --sample_every_n_epochs="1000" ' + '--seed="1234" --cache_latents --optimizer_type="AdamW8bit" --max_data_loader_n_workers="0" --bucket_reso_steps=64 ' + '--xformers --bucket_no_upscale --noise_offset=0.0 --tokenizer_cache_dir="/home/chinatszrn/.cache/clip"') + print("cmd_train:", cmd_train) + os.system(cmd_train) + + lora_path = os.path.join(model_dir, 'hairstyle_hd_lora.safetensors') + + # tianjin callback + if is_tj == "1": + url = 'http://service.aicloud.fit:7395/api/hair/trainCallBack' + + if not os.path.exists(lora_path): + send_request(url, -1, task_dict) + else: + send_request(url, 0, task_dict) + + except Exception as e: + send_request(url, -1, None) + continue + + +if __name__ == '__main__': + global_var.webui_lora_dir = '/gz-fs/models/Lora' + + # avatar训练的任务队列 + global_var.hair_style_lora_train_task_sq = Queue() + p = Process(target=train_thread, args=( + global_var.hair_style_lora_train_task_sq, 0)) + p.start() + + # # 启动服务 + # app.run(debug=True, port=32678, host='0.0.0.0') + + server = pywsgi.WSGIServer(('0.0.0.0', 32678), app) # test port + server.serve_forever() diff --git a/photo_service/lora_train_service_1.py b/photo_service/lora_train_service_1.py new file mode 100644 index 0000000..3bec5a8 --- /dev/null +++ b/photo_service/lora_train_service_1.py @@ -0,0 +1,357 @@ +# -*- coding: utf-8 -*- + +import os +import sys +import uuid + +from webui_im2im import ControlnetRequestImg2Img +import numpy as np +import base64 +import cv2 +import os, sys +from gevent import pywsgi, monkey +from multiprocessing import Process, Queue +import glob +import datetime +# from gpt4v_caption import caption_image + +# monkey.patch_all() +# sys.setrecursionlimit(20000) + +# 将当前工作目录切换到当前目录 +project_dir = os.path.dirname(os.path.abspath(__file__)) +os.chdir(project_dir) +sys.path.append(project_dir) + +from flask import Flask, request, jsonify +import global_variable as global_var +import json +#from common.logger import config + +BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + +version = "online" +if version == "local": + current_url = 'http://service.aicloud.fit:7393/' + kohya_ss_home_dir = '/home/chinatszrn/Documents/miaoya/kohya_ss_home' + webui_lora_dir = '/home/chinatszrn/Documents/miaoya/webui_home/stable-diffusion-webui/models/Lora' + inference_use_onediff = False + callback_url = 'http://service.aicloud.fit:7395/api/hair/trainCallBack' +else: + current_url = 'http://0.0.0.0:7393/' + 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') + inference_use_onediff = False + callback_url = 'http://0.0.0.0:8801/api/hair/trainCallBack' +base_webui_port = '57860' +base_onediff_port = '9038' +base_webui_url = "http://127.0.0.1:57860/sdapi/v1/img2img" + + +app = Flask(__name__) + + +def send_request(url, state, task_dict): + import requests + import json + + msg = '头发lora训练失败' if state == -1 else '头发lora训练成功' + payload = json.dumps({ + "task_id": task_dict['task_id'] if task_dict is not None else '', + "hair_id": task_dict['hair_id'] if task_dict is not None else '', + "state": state, + "msg": msg, + "is_tj": task_dict['is_tj'] + }) + + print("payload:", payload) + + headers = { + 'Content-Type': 'application/json' + } + + requests.request("POST", url, headers=headers, data=payload) + + +@app.route('/api/hair/train', methods=['POST']) +def train_hair_lora(): + import os + import json + task_id = '' + try: + # 获取请求参数 + request_data = request.get_json() + assert 'task_id' in request_data, 'task_id is required' + task_id = request_data['task_id'] + assert 'hair_id' in request_data and 'hair_material_dir' in request_data, 'hair_id and hair_material_dir is required' + hair_material_dir = request_data['hair_material_dir'] + assert os.path.exists(hair_material_dir), f'{hair_material_dir} not exists' + img_dir = os.path.join(hair_material_dir, 'images') + assert os.path.exists(img_dir), f'{img_dir} not exists' + assert os.path.exists( + os.path.join(hair_material_dir, 'model')), f'{os.path.join(hair_material_dir, "model")} not exists' + + #判断img_dir下面是否只有一个文件夹 + img_dir_list = os.listdir(img_dir) + assert len(img_dir_list) == 1, f'{img_dir}下面只能有一个文件夹' + train_image_dir = os.path.join(img_dir, img_dir_list[0]) + assert os.path.isdir(train_image_dir), f'{train_image_dir}不是文件夹' + # 检查该文件夹是否以数字加下划线开头 + assert img_dir_list[0].split('_')[0].isdigit(), f'{img_dir_list[0]}不是数字开头' + + #判断文件夹下面是否有图片 + img_list = glob.glob(train_image_dir + '/*.png') + assert len(img_list) > 10, f'{os.path.join(img_dir, img_dir_list[0])}下面图片数量小于10张' + request_data['train_image_dir'] = train_image_dir + + # is tianjin company + try: + request_data['is_tj'] = request_data['is_tj'] + request_data['device_id'] = request_data['device_id'] + except Exception as e: + print(e) + request_data['is_tj'] = '0' + request_data['device_id'] = '1' + + # 将请求数据放入队列 + global_var.hair_style_lora_train_task_sq.put(request_data) + + # 返回结果 + ret_dict = dict(state=0, msg='头发lora训练开始', task_id=task_id) + return jsonify(ret_dict) + except Exception as e: + ret_dict = dict(state=-1, msg=str(e), task_id=task_id) + return jsonify(ret_dict) + + +@app.route('/api/hair/inference', methods=['POST']) +def inference_webui(): + import os + import requests + import time + + + try: + t1 = time.time() + # 获取请求参数 + request_data = request.get_json() + # with open('request_data.json', 'w') as f: + # json.dump(request_data, f) + request_json = request_data.get('request_json') + # hd_version_flag = request_data.get('hd_version_flag', False) + hair_id = request_data['hair_id'] + inference_port = request_data['inference_port'] + + if inference_use_onediff and 'refiner_checkpoint' not in request_json: + selected_url = base_webui_url.replace(base_webui_port, base_onediff_port) + request_json['script_name'] = 'onediff_diffusion_model' + print('user onediff') + else: + selected_url = base_webui_url.replace(base_webui_port, inference_port) + + + + # if hd_version_flag: + # print('user hd version!!!!!') + # lora_file_name = 'hairstyle_hd_lora.safetensors' + # else: + # lora_file_name = 'hairstyle_lora.safetensors' + + hair_material_dir = request_data.get('hair_material_dir') + print(hair_material_dir) + assert os.path.exists(hair_material_dir), f'{hair_material_dir} not exists' + + hd_version_flag = True + lora_file_name = 'hairstyle_hd_lora.safetensors' + request_lora_path = os.path.join(hair_material_dir, 'model', lora_file_name) + if not os.path.exists(request_lora_path): + print('use low resolution lora') + hd_version_flag = False + lora_file_name = 'hairstyle_lora.safetensors' + request_lora_path = os.path.join(hair_material_dir, 'model', lora_file_name) + + assert os.path.exists(request_lora_path), f'{request_lora_path} not exists' + + if hd_version_flag: + tmp_lora_name = f'{hair_id}_hd' + else: + tmp_lora_name = f'{hair_id}' + lora_dst_path = os.path.join(global_var.webui_lora_dir, f'{tmp_lora_name}.safetensors') + if not os.path.exists(lora_dst_path): + os.system('cp {} {}'.format(request_lora_path, lora_dst_path)) + request_json['prompt'] = f', titor hairstyle, ' + request_json['prompt'] + request_json['negative_prompt'] = request_json['negative_prompt'] + ', faceless, no human' + print('pre_stage cost:', time.time() - t1) + + start_time = time.time() + current_time = datetime.datetime.now() + print("*********cxt log1, infer***************, {}:{}:{}, {}".format(current_time.hour, current_time.minute, + current_time.second, selected_url)) + response = requests.post(url=selected_url, json=request_json) + print('inference time:', time.time() - start_time) + + ret_json = response.json() + # if os.path.exists(lora_dst_path): + # os.remove(lora_dst_path) + # 返回结果 + return jsonify(ret_json) + except Exception as e: + ret_dict = dict(state=-1, msg=str(e)) + return jsonify(ret_dict) + + +@app.route('/api/hair/inference_diy', methods=['POST']) +def inference_diy_webui(): + import os + import requests + import time + + try: + t1 = time.time() + # 获取请求参数 + request_data = request.get_json() + # with open('request_data.json', 'w') as f: + # json.dump(request_data, f) + request_json = request_data.get('request_json') + inference_port = request_data['inference_port'] + + if inference_use_onediff and 'refiner_checkpoint' not in request_json: + selected_url = base_webui_url.replace(base_webui_port, base_onediff_port) + request_json['script_name'] = 'onediff_diffusion_model' + print('user onediff') + else: + selected_url = base_webui_url.replace(base_webui_port, inference_port) + + request_json['prompt'] = f'titor hairstyle, ' + request_json['prompt'] + request_json['negative_prompt'] = request_json['negative_prompt'] + ', faceless, no human' + print('pre_stage cost:', time.time() - t1) + + start_time = time.time() + current_time = datetime.datetime.now() + print("*********cxt log1, diy***************, {}:{}:{}, {}".format(current_time.hour, current_time.minute, + current_time.second, selected_url)) + response = requests.post(url=selected_url, json=request_json) + print('inference time:', time.time() - start_time) + + ret_json = response.json() + # if os.path.exists(lora_dst_path): + # os.remove(lora_dst_path) + # 返回结果 + return jsonify(ret_json) + except Exception as e: + ret_dict = dict(state=-1, msg=str(e)) + return jsonify(ret_dict) + +def train_thread(sq, gpu_id): + while True: + url = f'{current_url}api/hair/trainCallBack' + task_dict = None + try: + task_dict = sq.get() + hair_material_dir = task_dict['hair_material_dir'] + images_dir = os.path.join(hair_material_dir, 'images') + model_dir = os.path.join(hair_material_dir, 'model') + train_image_dir = task_dict['train_image_dir'] + # tag = task_dict['tag'] + tag = "" + is_tj = task_dict['is_tj'] + device_id = task_dict['device_id'] + + sample_dir = os.path.join(model_dir, 'sample') + if not os.path.exists(sample_dir): + os.makedirs(sample_dir) + sample_txt = os.path.join(sample_dir, 'prompt.txt') + with open(sample_txt, 'w') as f_s: + f_s.write('titor hairstyle, easyphoto, faceless, no human, white background, simple background, ' + + tag + + ' --n low quality, worst quality, bad anatomy, bad composition, poor, low effort --h 768 ' + '--w 768 --s 30 --l 7') + + # 给训练图片打标签 + # cmd_caption = ('docker run --rm --gpus all -v /home/chinatszrn/Documents/miaoya/kohya_ss_home:/home/chinatszrn ' + # '-v /mnt:/mnt -e PATH=/home/chinatszrn/.local/bin -w /home/chinatszrn/kohya_ss ' + # '--net=host chinatszrn/ubuntu:kohya_ss accelerate ' + # 'launch "./finetune/tag_images_by_wd14_tagger.py" --batch_size=2 ' + # '--general_threshold=0.5 --character_threshold=0.5 --caption_extension=".txt" ' + # '--model="SmilingWolf/wd-v1-4-convnextv2-tagger-v2" --max_data_loader_n_workers="2" ' + # '--debug --remove_underscore --frequency_tags --undesired_tags="1girl, 1boy" ' + # f'"{train_image_dir}"') + # os.system(cmd_caption) + + img_path_list = glob.glob(train_image_dir + '/*.png') + # 给训练图片打标签, gpt + for img_path in img_path_list: + # print(img_path) + txt_path = img_path[:-4] + '.txt' + + print("tag img_path: ", img_path) + # tags = caption_image(img_path) + tags = tag + with open(txt_path, 'w') as f: + f.write('titor hairstyle, easyphoto, faceless, no human, white background, simple background, ' + tags) + + #判断是否每个训练图片都有标签文件 + # for img_path in img_path_list: + # assert os.path.exists(img_path[:-4]+'.txt'), f'{img_path}没有对应的标签文件' + # with open(img_path[:-4]+'.txt', 'r') as f: + # tags = f.readline() + # with open(img_path[:-4]+'.txt', 'w') as f: + # f.write('titor hairstyle, faceless, no human, gray background, simple background, ' + tags) + + os.system(f'chmod 777 {model_dir}') + #训练头发lora + 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') + cmd_train = ( + f'sudo docker run --rm --privileged=true --gpus "device=0" ' + f'-v {kohya_ss_home_dir}:/home/chinatszrn ' + f'-v {sd_models_dir}:/mnt/sd_models ' + '-e PATH=/home/chinatszrn/.local/bin -w /home/chinatszrn/kohya_ss ' + '--net=host chinatszrn/ubuntu:kohya_ss accelerate launch --num_cpu_threads_per_process=2 "./train_network.py" --enable_bucket ' + '--min_bucket_reso=256 --max_bucket_reso=2048 --pretrained_model_name_or_path="/mnt/sd_models/majicmixRealistic_v7.safetensors" ' + f'--train_data_dir={container_images_dir} --resolution="2000,2000" ' + f'--output_dir={container_model_dir} ' + '--network_alpha="64" --save_model_as=safetensors --network_module=networks.lora --text_encoder_lr=5e-05 ' + '--unet_lr=0.0001 --network_dim=128 --output_name="hairstyle_hd_lora" --lr_scheduler_num_cycles="20" ' + '--no_half_vae --learning_rate="0.0001" --lr_scheduler="cosine" --lr_warmup_steps="650" --train_batch_size="1" ' + '--max_train_steps="1500" --save_every_n_epochs="100" --mixed_precision="fp16" --save_precision="fp16" ' + '--caption_extension=".txt" --sample_sampler=ddim ' + f'--sample_prompts={sample_txt} --sample_every_n_epochs="1000" ' + '--seed="1234" --cache_latents --optimizer_type="AdamW8bit" --max_data_loader_n_workers="0" --bucket_reso_steps=64 ' + '--xformers --bucket_no_upscale --noise_offset=0.0 --tokenizer_cache_dir="/home/chinatszrn/.cache/clip"') + print("cmd_train:", cmd_train) + os.system(cmd_train) + + lora_path = os.path.join(model_dir, 'hairstyle_hd_lora.safetensors') + + # tianjin callback + if is_tj == "1": + url = f'{callback_url}' + + if not os.path.exists(lora_path): + send_request(url, -1, task_dict) + else: + send_request(url, 0, task_dict) + + except Exception as e: + send_request(url, -1, None) + continue + + +if __name__ == '__main__': + global_var.webui_lora_dir = webui_lora_dir + + # avatar训练的任务队列 + global_var.hair_style_lora_train_task_sq = Queue() + p = Process(target=train_thread, args=( + global_var.hair_style_lora_train_task_sq, 0)) + p.start() + + # # 启动服务 + # app.run(debug=True, port=32678, host='0.0.0.0') + + server = pywsgi.WSGIServer(('0.0.0.0', 32678), app) # test port + server.serve_forever() + + p.join() diff --git a/photo_service/lora_train_service_test.py b/photo_service/lora_train_service_test.py new file mode 100644 index 0000000..ce3d46a --- /dev/null +++ b/photo_service/lora_train_service_test.py @@ -0,0 +1,190 @@ +import os +import sys +import uuid + +from webui_im2im import ControlnetRequestImg2Img +import numpy as np +import base64 +import cv2 +import os,sys +from gevent import pywsgi, monkey +from multiprocessing import Process, Queue +import glob +from gpt4v_caption import caption_image + + +# 将当前工作目录切换到当前目录 +project_dir = os.path.dirname(os.path.abspath(__file__)) +os.chdir(project_dir) +sys.path.append(project_dir) + +import json + +def send_request(url, state, task_dict): + import requests + import json + + msg = '头发lora训练失败' if state == -1 else '头发lora训练成功' + payload = json.dumps({ + "task_id": task_dict['task_id'] if task_dict is not None else '', + "hair_id": task_dict['hair_id'] if task_dict is not None else '', + "state": state, + "msg": msg + }) + headers = { + 'Content-Type': 'application/json' + } + + requests.request("POST", url, headers=headers, data=payload) + + + +def train_hair_lora(): + import os + import json + task_id = '' + try: + # 获取请求参数 + request_data = "" + assert 'task_id' in request_data, 'task_id is required' + task_id = request_data['task_id'] + assert 'hair_id' in request_data and 'hair_material_dir' in request_data, 'hair_id and hair_material_dir is required' + hair_material_dir = request_data['hair_material_dir'] + assert os.path.exists(hair_material_dir), f'{hair_material_dir} not exists' + img_dir = os.path.join(hair_material_dir,'images') + assert os.path.exists(img_dir), f'{img_dir} not exists' + assert os.path.exists(os.path.join(hair_material_dir, 'model')), f'{os.path.join(hair_material_dir, "model")} not exists' + + #判断img_dir下面是否只有一个文件夹 + img_dir_list = os.listdir(img_dir) + assert len(img_dir_list) == 1, f'{img_dir}下面只能有一个文件夹' + train_image_dir = os.path.join(img_dir,img_dir_list[0]) + assert os.path.isdir(train_image_dir), f'{train_image_dir}不是文件夹' + # 检查该文件夹是否以数字加下划线开头 + assert img_dir_list[0].split('_')[0].isdigit(), f'{img_dir_list[0]}不是数字开头' + + #判断文件夹下面是否有图片 + img_list = glob.glob(train_image_dir + '/*.png') + assert len(img_list) > 10, f'{os.path.join(img_dir,img_dir_list[0])}下面图片数量小于10张' + request_data['train_image_dir'] = train_image_dir + except Exception as e: + print(e) + + +def inference_webui(): + import os + import requests + url = "http://127.0.0.1:57860/sdapi/v1/img2img" + try: + # 获取请求参数 + request_data = "" + # with open('request_data.json', 'w') as f: + # json.dump(request_data, f) + request_json = request_data.get('request_json') + hair_material_dir = request_data.get('hair_material_dir') + assert os.path.exists(hair_material_dir), f'{hair_material_dir} not exists' + request_lora_path = os.path.join(hair_material_dir, 'model', 'hairstyle_lora.safetensors') + assert os.path.exists(request_lora_path), f'{request_lora_path} not exists' + + tmp_lora_name = f'{uuid.uuid4()}' + lora_dst_path = "" + os.system('cp {} {}'.format(request_lora_path, lora_dst_path)) + request_json['prompt'] = f', titor hairstyle, ' + request_json['prompt'] + request_json['negative_prompt'] = request_json['negative_prompt'] + ', faceless, no human' + response = requests.post(url=url, json=request_json) + ret_json = response.json() + + if os.path.exists(lora_dst_path): + os.remove(lora_dst_path) + + except Exception as e: + print(e) + + +def tag(train_image_dir): + img_path_list = glob.glob(train_image_dir + '/*.png') + + # 给训练图片打标签, gpt + for img_path in img_path_list: + print(img_path) + txt_path = img_path[:-4] + '.txt' + # if not os.path.exists(txt_path): + tags = caption_image(img_path) + with open(txt_path, 'w') as f: + f.write('titor hairstyle, faceless, no human, gray background, simple background, ' + tags) + # else: + # continue + + +def train_thread(sq, gpu_id): + while True: + url = 'http://service.aicloud.fit:7393/api/hair/trainCallBack' + task_dict = None + try: + task_dict = sq.get() + hair_material_dir = task_dict['hair_material_dir'] + images_dir = os.path.join(hair_material_dir, 'images') + model_dir = os.path.join(hair_material_dir, 'model') + train_image_dir = task_dict['train_image_dir'] + # 给训练图片打标签 + # cmd_caption = ('docker run --rm --gpus all -v /home/chinatszrn/Documents/miaoya/kohya_ss_home:/home/chinatszrn ' + # '-v /mnt:/mnt -e PATH=/home/chinatszrn/.local/bin -w /home/chinatszrn/kohya_ss ' + # '--net=host chinatszrn/ubuntu:kohya_ss accelerate ' + # 'launch "./finetune/tag_images_by_wd14_tagger.py" --batch_size=2 ' + # '--general_threshold=0.5 --character_threshold=0.5 --caption_extension=".txt" ' + # '--model="SmilingWolf/wd-v1-4-convnextv2-tagger-v2" --max_data_loader_n_workers="2" ' + # '--debug --remove_underscore --frequency_tags --undesired_tags="1girl, 1boy" ' + # f'"{train_image_dir}"') + # os.system(cmd_caption) + + img_path_list = glob.glob(train_image_dir + '/*.png') + # 给训练图片打标签, gpt + for img_path in img_path_list: + tags = caption_image(img_path) + txt_path = img_path[:-4]+'.txt' + if not os.path.exists(txt_path): + with open(txt_path, 'w') as f: + f.write('titor hairstyle, faceless, no human, gray background, simple background, ' + tags) + else: + continue + + + #判断是否每个训练图片都有标签文件 + # for img_path in img_path_list: + # assert os.path.exists(img_path[:-4]+'.txt'), f'{img_path}没有对应的标签文件' + # with open(img_path[:-4]+'.txt', 'r') as f: + # tags = f.readline() + # with open(img_path[:-4]+'.txt', 'w') as f: + # f.write('titor hairstyle, faceless, no human, gray background, simple background, ' + tags) + + #训练头发lora + cmd_train = ('docker run --rm --gpus all -v /home/chinatszrn/Documents/miaoya/kohya_ss_home:/home/chinatszrn -v ' + '/mnt:/mnt -e PATH=/home/chinatszrn/.local/bin -w /home/chinatszrn/kohya_ss ' + '--net=host chinatszrn/ubuntu:kohya_ss accelerate launch --num_cpu_threads_per_process=2 "./train_network.py" --enable_bucket ' + '--min_bucket_reso=512 --max_bucket_reso=1800 --pretrained_model_name_or_path="/mnt/nas_hdd/米亚像馆/models/Stable-diffusion/majicmixRealistic_v7.safetensors" ' + f'--train_data_dir={images_dir} --resolution="1800,1800" ' + f'--output_dir={model_dir} ' + '--network_alpha="64" --save_model_as=safetensors --network_module=networks.lora --text_encoder_lr=5e-05 ' + '--unet_lr=0.0001 --network_dim=128 --output_name="hairstyle_lora" --lr_scheduler_num_cycles="12" ' + '--no_half_vae --learning_rate="0.0001" --lr_scheduler="cosine" --lr_warmup_steps="96" --train_batch_size="1" ' + '--max_train_steps="4000" --save_every_n_epochs="100" --mixed_precision="fp16" --save_precision="fp16" ' + '--seed="1234" --cache_latents --optimizer_type="AdamW8bit" --max_data_loader_n_workers="0" --bucket_reso_steps=64 ' + '--xformers --bucket_no_upscale --noise_offset=0.0 --tokenizer_cache_dir="/home/chinatszrn/.cache/clip"') + os.system(cmd_train) + + lora_path = os.path.join(model_dir, 'hairstyle_lora.safetensors') + + + if not os.path.exists(lora_path): + send_request(url, -1, task_dict) + else: + send_request(url, 0, task_dict) + + except Exception as e: + send_request(url, -1, None) + continue + + +if __name__ == '__main__': + + tag("/mnt/database2/online-server/hair-online/hair_lora_train_material/AAVWGW1NN-0KC-33B-24-/images/1_hairstyle") diff --git a/photo_service/mask.png b/photo_service/mask.png new file mode 100644 index 0000000..accc003 Binary files /dev/null and b/photo_service/mask.png differ diff --git a/photo_service/request_demo.py b/photo_service/request_demo.py new file mode 100644 index 0000000..9e8daab --- /dev/null +++ b/photo_service/request_demo.py @@ -0,0 +1,37 @@ +import requests +import cv2 +import numpy as np +import base64 +import sys +import json + +if __name__ == '__main__': + path = './data/template01.png' + img = cv2.imread(path) + retval, bytes = cv2.imencode('.jpg', img) + encoded_image = base64.b64encode(bytes).decode('utf-8') + user_id = '61AB9097' + + #给api_service.py发送请求 + url = "http://i-2.gpushare.com:53412/user/generate" + payload = json.dumps({ + "user_id": "61AB9097", + "base_img": encoded_image + }) + headers = { + 'Content-Type': 'application/json' + } + response = requests.request("POST", url, headers=headers, data=payload) + + if response.status_code != 200: + raise RuntimeError(f"Failed to send request to API service. Status code: {response.status_code}") + + ret_image_b64 = response.json().get('generate_photo_b64') + + if ret_image_b64 is None: + raise RuntimeError(f"ret image failed!") + + image_array = np.frombuffer(base64.b64decode(ret_image_b64), np.uint8) + image = cv2.imdecode(image_array, cv2.IMREAD_COLOR) + cv2.imshow('image', cv2.resize(image, (0, 0), fx=0.5, fy=0.5)) + cv2.waitKey() \ No newline at end of file diff --git a/photo_service/start_service.sh b/photo_service/start_service.sh new file mode 100644 index 0000000..8684706 --- /dev/null +++ b/photo_service/start_service.sh @@ -0,0 +1 @@ +python api_service.py /hy-tmp/photo_service/service_data /hy-tmp/stable-diffusion-webui/models/Lora 7860 1234 \ No newline at end of file diff --git a/photo_service/tmp.py b/photo_service/tmp.py new file mode 100644 index 0000000..b771275 --- /dev/null +++ b/photo_service/tmp.py @@ -0,0 +1,18 @@ +import cv2 +import os +import json +import base64 +from PIL import Image +import io + +if __name__ == '__main__': + with open('request_json.json', 'r') as f: + request_data = json.load(f) + request_json = request_data + img_base64 = request_json['init_images'][0] + mask_base64 = request_json['mask'] + image = Image.open(io.BytesIO(base64.b64decode(img_base64))) + image.save('image.png') + mask = Image.open(io.BytesIO(base64.b64decode(mask_base64))) + mask.save('mask.png') + diff --git a/photo_service/utils/landmark_processor.py b/photo_service/utils/landmark_processor.py new file mode 100755 index 0000000..8fa1c9f --- /dev/null +++ b/photo_service/utils/landmark_processor.py @@ -0,0 +1,1663 @@ +import colorsys +import cv2 +import numpy as np +import random +import math +import time + +mean_face_x = np.array([ + 0.000213256, 0.0752622, 0.18113, 0.29077, 0.393397, 0.586856, 0.689483, 0.799124, + 0.904991, 0.98004, 0.490127, 0.490127, 0.490127, 0.490127, 0.36688, 0.426036, + 0.490127, 0.554217, 0.613373, 0.121737, 0.187122, 0.265825, 0.334606, 0.260918, + 0.182743, 0.645647, 0.714428, 0.793132, 0.858516, 0.79751, 0.719335, 0.254149, + 0.340985, 0.428858, 0.490127, 0.551395, 0.639268, 0.726104, 0.642159, 0.556721, + 0.490127, 0.423532, 0.338094, 0.290379, 0.428096, 0.490127, 0.552157, 0.689874, + 0.553364, 0.490127, 0.42689]) + +mean_face_y = np.array([ + 0.106454, 0.038915, 0.0187482, 0.0344891, 0.0773906, 0.0773906, 0.0344891, + 0.0187482, 0.038915, 0.106454, 0.203352, 0.307009, 0.409805, 0.515625, 0.587326, + 0.609345, 0.628106, 0.609345, 0.587326, 0.216423, 0.178758, 0.179852, 0.231733, + 0.245099, 0.244077, 0.231733, 0.179852, 0.178758, 0.216423, 0.244077, 0.245099, + 0.780233, 0.745405, 0.727388, 0.742578, 0.727388, 0.745405, 0.780233, 0.864805, + 0.902192, 0.909281, 0.902192, 0.864805, 0.784792, 0.778746, 0.785343, 0.778746, + 0.784792, 0.824182, 0.831803, 0.824182]) + +landmarks_2D = np.stack([mean_face_x, mean_face_y], axis=1) + +mean_face_x_1k = np.array([0.498047, 0.504671, 0.511286, 0.517984, 0.524451, 0.531086, 0.537574, 0.543902, 0.550305, 0.556545, 0.562891, 0.568903, 0.574994, 0.581092, 0.587026, 0.592814, 0.598603, 0.604259, 0.609802, 0.615311, 0.620821, 0.626135, 0.631542, 0.636597, 0.641621, 0.646680, 0.651656, 0.656505, 0.661374, 0.666025, 0.670571, 0.675225, 0.679616, 0.683987, 0.688181, 0.692460, 0.696634, 0.700552, 0.704500, 0.708402, 0.712261, 0.715842, 0.719373, 0.722930, 0.726486, 0.729785, 0.732875, 0.736135, 0.739115, 0.742179, 0.744925, 0.747688, 0.750314, 0.752887, 0.755412, 0.757737, 0.760038, 0.762291, 0.764312, 0.766333, 0.768259, 0.770363, 0.771888, 0.773763, 0.775485, 0.777069, 0.778656, 0.780126, 0.781483, 0.782879, 0.784085, 0.785261, 0.786402, 0.787338, 0.788398, 0.789310, 0.790038, 0.790796, 0.791281, 0.792001, 0.792397, 0.792903, 0.793221, 0.793554, 0.793613, 0.793848, 0.793856, 0.793925, 0.793911, 0.793825, 0.793601, 0.793426, 0.793199, 0.792833, 0.792528, 0.791962, 0.791501, 0.790897, 0.790382, 0.789660, 0.788884, 0.787702, 0.786300, 0.784951, 0.783363, 0.781861, 0.780103, 0.778348, 0.776405, 0.774489, 0.772390, 0.770098, 0.767540, 0.765030, 0.762359, 0.759349, 0.756513, 0.753197, 0.749784, 0.746353, 0.742572, 0.738633, 0.734433, 0.730200, 0.725658, 0.720657, 0.715711, 0.710406, 0.705108, 0.699482, 0.693569, 0.687476, 0.681209, 0.674681, 0.668094, 0.661239, 0.654234, 0.647131, 0.639989, 0.632582, 0.625134, 0.617439, 0.610009, 0.602125, 0.594318, 0.586546, 0.578466, 0.570582, 0.562638, 0.554622, 0.546602, 0.538546, 0.530536, 0.522423, 0.514258, 0.506205, 0.498047, 0.489889, 0.481836, 0.473671, 0.465558, 0.457548, 0.449491, 0.441472, 0.433455, 0.425511, 0.417628, 0.409547, 0.401776, 0.393969, 0.386085, 0.378655, 0.370960, 0.363512, 0.356105, 0.348963, 0.341860, 0.334855, 0.328000, 0.321412, 0.314885, 0.308618, 0.302525, 0.296612, 0.290986, 0.285687, 0.280382, 0.275437, 0.270436, 0.265894, 0.261660, 0.257461, 0.253522, 0.249740, 0.246310, 0.242897, 0.239580, 0.236744, 0.233735, 0.231064, 0.228553, 0.225996, 0.223704, 0.221605, 0.219689, 0.217746, 0.215990, 0.214232, 0.212730, 0.211143, 0.209794, 0.208392, 0.207210, 0.206433, 0.205712, 0.205197, 0.204593, 0.204132, 0.203565, 0.203261, 0.202895, 0.202667, 0.202492, 0.202268, 0.202183, 0.202169, 0.202238, 0.202246, 0.202480, 0.202540, 0.202873, 0.203191, 0.203696, 0.204093, 0.204813, 0.205297, 0.206055, 0.206783, 0.207696, 0.208755, 0.209692, 0.210833, 0.212009, 0.213215, 0.214611, 0.215967, 0.217438, 0.219024, 0.220609, 0.222330, 0.224206, 0.225731, 0.227835, 0.229761, 0.231781, 0.233803, 0.236055, 0.238356, 0.240681, 0.243207, 0.245780, 0.248405, 0.251168, 0.253915, 0.256979, 0.259958, 0.263219, 0.266309, 0.269608, 0.273163, 0.276721, 0.280252, 0.283833, 0.287692, 0.291593, 0.295541, 0.299459, 0.303633, 0.307913, 0.312107, 0.316478, 0.320869, 0.325523, 0.330069, 0.334719, 0.339589, 0.344438, 0.349414, 0.354473, 0.359497, 0.364552, 0.369959, 0.375273, 0.380783, 0.386292, 0.391835, 0.397491, 0.403280, 0.409068, 0.415002, 0.421100, 0.427190, 0.433203, 0.439548, 0.445789, 0.452191, 0.458520, 0.465008, 0.471643, 0.478110, 0.484807, 0.491422, 0.396255, 0.397888, 0.399588, 0.401344, 0.403150, 0.405001, 0.406895, 0.408831, 0.410807, 0.412822, 0.414875, 0.416966, 0.419095, 0.421262, 0.423465, 0.425705, 0.427981, 0.430294, 0.432644, 0.435030, 0.437451, 0.439908, 0.442400, 0.444928, 0.447489, 0.450084, 0.452712, 0.455372, 0.458062, 0.460782, 0.463531, 0.466307, 0.469108, 0.471933, 0.474781, 0.477648, 0.480532, 0.483432, 0.486345, 0.489267, 0.492194, 0.495123, 0.498047, 0.500971, 0.503899, 0.506827, 0.509749, 0.512662, 0.515562, 0.518446, 0.521313, 0.524160, 0.526986, 0.529787, 0.532563, 0.535312, 0.538032, 0.540722, 0.543382, 0.546009, 0.548605, 0.551166, 0.553693, 0.556186, 0.558642, 0.561064, 0.563450, 0.565799, 0.568113, 0.570389, 0.572629, 0.574832, 0.576998, 0.579127, 0.581219, 0.583272, 0.585287, 0.587263, 0.589198, 0.591093, 0.592944, 0.594749, 0.596506, 0.598206, 0.599839, 0.597799, 0.595660, 0.593444, 0.591167, 0.588838, 0.586464, 0.584049, 0.581599, 0.579114, 0.576598, 0.574052, 0.571476, 0.568873, 0.566242, 0.563585, 0.560901, 0.558193, 0.555461, 0.552704, 0.549922, 0.547116, 0.544286, 0.541432, 0.538555, 0.535655, 0.532732, 0.529786, 0.526819, 0.523832, 0.520830, 0.517033, 0.513236, 0.509439, 0.505641, 0.501844, 0.498047, 0.494250, 0.490452, 0.486655, 0.482858, 0.479061, 0.475264, 0.472262, 0.469275, 0.466308, 0.463362, 0.460439, 0.457538, 0.454662, 0.451808, 0.448978, 0.446172, 0.443390, 0.440633, 0.437900, 0.435192, 0.432509, 0.429852, 0.427221, 0.424617, 0.422042, 0.419496, 0.416979, 0.414495, 0.412044, 0.409630, 0.407256, 0.404927, 0.402650, 0.400434, 0.398294, 0.410228, 0.414399, 0.418630, 0.422899, 0.427196, 0.431515, 0.435855, 0.440212, 0.444587, 0.448978, 0.453385, 0.457806, 0.462242, 0.466691, 0.471152, 0.475622, 0.480101, 0.484586, 0.489074, 0.493564, 0.498047, 0.502530, 0.507020, 0.511508, 0.515993, 0.520472, 0.524942, 0.529403, 0.533852, 0.538288, 0.542709, 0.547116, 0.551507, 0.555882, 0.560238, 0.564578, 0.568898, 0.573194, 0.577464, 0.581695, 0.585866, 0.581642, 0.577369, 0.573066, 0.568743, 0.564401, 0.560044, 0.555673, 0.551287, 0.546890, 0.542478, 0.538055, 0.533622, 0.529181, 0.524733, 0.520280, 0.515825, 0.511370, 0.506919, 0.502475, 0.498047, 0.493618, 0.489175, 0.484723, 0.480269, 0.475814, 0.471361, 0.466913, 0.462472, 0.458039, 0.453616, 0.449204, 0.444806, 0.440421, 0.436050, 0.431693, 0.427351, 0.423027, 0.418725, 0.414452, 0.459656, 0.459227, 0.458691, 0.458058, 0.457338, 0.456539, 0.455664, 0.454719, 0.453708, 0.452633, 0.451497, 0.450302, 0.449049, 0.447741, 0.446378, 0.444962, 0.443494, 0.441973, 0.440402, 0.438780, 0.437107, 0.435384, 0.433612, 0.431790, 0.429919, 0.427998, 0.426028, 0.424009, 0.421938, 0.419814, 0.417632, 0.415384, 0.442766, 0.471804, 0.498047, 0.524290, 0.553327, 0.580709, 0.578462, 0.576279, 0.574156, 0.572085, 0.570065, 0.568095, 0.566175, 0.564304, 0.562481, 0.560709, 0.558987, 0.557314, 0.555692, 0.554120, 0.552600, 0.551131, 0.549715, 0.548353, 0.547044, 0.545792, 0.544597, 0.543461, 0.542386, 0.541375, 0.540430, 0.539555, 0.538755, 0.538036, 0.537403, 0.536867, 0.536438, 0.551823, 0.524746, 0.471347, 0.444271, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.498047, 0.370401, 0.394341, 0.393977, 0.392891, 0.391151, 0.388793, 0.385710, 0.382329, 0.378604, 0.374564, 0.370390, 0.366216, 0.362177, 0.358453, 0.355074, 0.351995, 0.349640, 0.347903, 0.346822, 0.346461, 0.346826, 0.347911, 0.349651, 0.352010, 0.355092, 0.358473, 0.362199, 0.366239, 0.370413, 0.374587, 0.378626, 0.382350, 0.385728, 0.388808, 0.391163, 0.392899, 0.393981, 0.423391, 0.421623, 0.419536, 0.417200, 0.414654, 0.411925, 0.409036, 0.406003, 0.402844, 0.399571, 0.396199, 0.392740, 0.389206, 0.385609, 0.381960, 0.378271, 0.374552, 0.370812, 0.367063, 0.363314, 0.359575, 0.355859, 0.352173, 0.348530, 0.344945, 0.341429, 0.337998, 0.334669, 0.331464, 0.328407, 0.325532, 0.322889, 0.320595, 0.322717, 0.325103, 0.327689, 0.330440, 0.333327, 0.336328, 0.339426, 0.342603, 0.345848, 0.349146, 0.352488, 0.355863, 0.359264, 0.362683, 0.366116, 0.369556, 0.372999, 0.376442, 0.379882, 0.383316, 0.386742, 0.390158, 0.393563, 0.396955, 0.400332, 0.403694, 0.407038, 0.410363, 0.413666, 0.416943, 0.420190, 0.625692, 0.649633, 0.649272, 0.648190, 0.646454, 0.644099, 0.641019, 0.637641, 0.633917, 0.629878, 0.625704, 0.621530, 0.617490, 0.613764, 0.610383, 0.607301, 0.604942, 0.603202, 0.602117, 0.601752, 0.602113, 0.603194, 0.604931, 0.607286, 0.610365, 0.613744, 0.617468, 0.621507, 0.625681, 0.629855, 0.633895, 0.637621, 0.641001, 0.644084, 0.646442, 0.648182, 0.649268, 0.572703, 0.574471, 0.576557, 0.578893, 0.581440, 0.584168, 0.587058, 0.590090, 0.593250, 0.596523, 0.599895, 0.603354, 0.606888, 0.610485, 0.614134, 0.617823, 0.621542, 0.625282, 0.629031, 0.632780, 0.636518, 0.640235, 0.643921, 0.647563, 0.651149, 0.654665, 0.658096, 0.661425, 0.664630, 0.667686, 0.670562, 0.673205, 0.675499, 0.673377, 0.670991, 0.668404, 0.665654, 0.662767, 0.659765, 0.656668, 0.653491, 0.650246, 0.646948, 0.643606, 0.640231, 0.636830, 0.633411, 0.629978, 0.626538, 0.623095, 0.619652, 0.616212, 0.612778, 0.609352, 0.605936, 0.602531, 0.599139, 0.595762, 0.592400, 0.589056, 0.585731, 0.582428, 0.579150, 0.575904, 0.552267, 0.555073, 0.558014, 0.561313, 0.565180, 0.569636, 0.574564, 0.579836, 0.585356, 0.591050, 0.596857, 0.602735, 0.608652, 0.614591, 0.620542, 0.626500, 0.632462, 0.638425, 0.644390, 0.649985, 0.655581, 0.661175, 0.666765, 0.672343, 0.677894, 0.683391, 0.688788, 0.694021, 0.699020, 0.703737, 0.708170, 0.712362, 0.716374, 0.720268, 0.724091, 0.727877, 0.731647, 0.726965, 0.722281, 0.717595, 0.712902, 0.708202, 0.703489, 0.698762, 0.694017, 0.689253, 0.684472, 0.679678, 0.674876, 0.670068, 0.665257, 0.660446, 0.655633, 0.650821, 0.646008, 0.640808, 0.635608, 0.630406, 0.625203, 0.619997, 0.614788, 0.609577, 0.604360, 0.599139, 0.593915, 0.588693, 0.583472, 0.578256, 0.573043, 0.567837, 0.562637, 0.557446, 0.264446, 0.268217, 0.272003, 0.275826, 0.279720, 0.283732, 0.287924, 0.292357, 0.297074, 0.302073, 0.307305, 0.312703, 0.318199, 0.323751, 0.329329, 0.334919, 0.340513, 0.346109, 0.351704, 0.357669, 0.363632, 0.369594, 0.375552, 0.381503, 0.387442, 0.393359, 0.399237, 0.405044, 0.410737, 0.416258, 0.421530, 0.426457, 0.430914, 0.434781, 0.438079, 0.441020, 0.443827, 0.438648, 0.433457, 0.428257, 0.423050, 0.417838, 0.412621, 0.407401, 0.402179, 0.396955, 0.391734, 0.386517, 0.381305, 0.376096, 0.370891, 0.365688, 0.360486, 0.355285, 0.350085, 0.345273, 0.340460, 0.335648, 0.330836, 0.326026, 0.321218, 0.316416, 0.311622, 0.306841, 0.302077, 0.297332, 0.292605, 0.287892, 0.283191, 0.278499, 0.273812, 0.269128 +]) +mean_face_y_1k = np.array([0.851392, 0.851204, 0.850802, 0.849951, 0.849050, 0.847896, 0.846372, 0.844708, 0.842887, 0.840805, 0.838470, 0.836093, 0.833361, 0.830589, 0.827781, 0.824650, 0.821361, 0.817917, 0.814463, 0.810803, 0.807127, 0.803118, 0.799272, 0.795088, 0.790897, 0.786585, 0.782195, 0.777726, 0.773128, 0.768407, 0.763595, 0.758801, 0.753939, 0.748915, 0.743745, 0.738664, 0.733510, 0.728055, 0.722725, 0.717442, 0.711952, 0.706343, 0.700725, 0.695036, 0.689315, 0.683557, 0.677857, 0.671952, 0.666097, 0.659596, 0.652993, 0.646386, 0.639687, 0.632843, 0.626218, 0.619403, 0.612554, 0.605551, 0.598705, 0.591735, 0.584939, 0.577854, 0.571094, 0.564026, 0.557009, 0.549978, 0.542887, 0.535838, 0.528892, 0.521764, 0.514596, 0.507444, 0.500318, 0.493164, 0.486134, 0.478851, 0.471725, 0.464586, 0.457279, 0.450071, 0.443049, 0.435691, 0.428445, 0.421397, 0.414262, 0.406894, 0.399680, 0.392547, 0.385399, 0.378182, 0.370759, 0.363690, 0.356390, 0.349346, 0.341952, 0.335000, 0.327902, 0.320707, 0.313459, 0.306325, 0.299213, 0.291181, 0.283309, 0.275338, 0.267185, 0.259441, 0.251545, 0.243652, 0.235765, 0.228054, 0.220194, 0.212369, 0.204820, 0.197130, 0.189558, 0.182061, 0.174602, 0.167242, 0.159851, 0.152736, 0.145606, 0.138595, 0.131706, 0.125067, 0.118485, 0.112101, 0.105897, 0.099812, 0.094088, 0.088477, 0.082962, 0.077854, 0.072924, 0.068273, 0.063748, 0.059715, 0.055721, 0.052042, 0.048762, 0.045388, 0.042616, 0.039888, 0.037449, 0.035191, 0.033201, 0.031339, 0.029751, 0.028369, 0.027112, 0.025987, 0.025155, 0.024428, 0.023762, 0.023446, 0.023057, 0.022894, 0.022946, 0.022894, 0.023057, 0.023446, 0.023762, 0.024428, 0.025155, 0.025987, 0.027112, 0.028369, 0.029751, 0.031339, 0.033201, 0.035191, 0.037449, 0.039888, 0.042616, 0.045388, 0.048762, 0.052042, 0.055721, 0.059715, 0.063748, 0.068273, 0.072924, 0.077854, 0.082962, 0.088477, 0.094088, 0.099812, 0.105897, 0.112101, 0.118485, 0.125067, 0.131706, 0.138595, 0.145606, 0.152736, 0.159851, 0.167242, 0.174602, 0.182061, 0.189558, 0.197130, 0.204820, 0.212369, 0.220194, 0.228054, 0.235765, 0.243652, 0.251545, 0.259441, 0.267185, 0.275338, 0.283309, 0.291181, 0.299213, 0.306325, 0.313459, 0.320707, 0.327902, 0.335000, 0.341952, 0.349346, 0.356390, 0.363690, 0.370759, 0.378182, 0.385399, 0.392547, 0.399680, 0.406894, 0.414262, 0.421397, 0.428445, 0.435691, 0.443049, 0.450071, 0.457279, 0.464586, 0.471725, 0.478851, 0.486134, 0.493164, 0.500318, 0.507444, 0.514596, 0.521764, 0.528892, 0.535838, 0.542887, 0.549978, 0.557009, 0.564026, 0.571094, 0.577854, 0.584939, 0.591735, 0.598705, 0.605551, 0.612554, 0.619403, 0.626218, 0.632843, 0.639687, 0.646386, 0.652993, 0.659596, 0.666097, 0.671952, 0.677857, 0.683557, 0.689315, 0.695036, 0.700725, 0.706343, 0.711952, 0.717442, 0.722725, 0.728055, 0.733510, 0.738664, 0.743745, 0.748915, 0.753939, 0.758801, 0.763595, 0.768407, 0.773128, 0.777726, 0.782195, 0.786585, 0.790897, 0.795088, 0.799272, 0.803118, 0.807127, 0.810803, 0.814463, 0.817917, 0.821361, 0.824650, 0.827781, 0.830589, 0.833361, 0.836093, 0.838470, 0.840805, 0.842887, 0.844708, 0.846372, 0.847896, 0.849050, 0.849951, 0.850802, 0.851204, 0.664237, 0.666525, 0.668791, 0.671031, 0.673242, 0.675423, 0.677571, 0.679687, 0.681769, 0.683818, 0.685829, 0.687805, 0.689742, 0.691641, 0.693499, 0.695315, 0.697088, 0.698816, 0.700499, 0.702133, 0.703718, 0.705251, 0.706730, 0.708155, 0.709522, 0.710829, 0.712076, 0.713258, 0.714375, 0.715424, 0.716404, 0.717311, 0.718144, 0.718900, 0.719577, 0.720171, 0.720681, 0.721102, 0.721432, 0.721666, 0.721798, 0.721824, 0.721732, 0.721824, 0.721798, 0.721666, 0.721432, 0.721102, 0.720681, 0.720171, 0.719577, 0.718900, 0.718144, 0.717311, 0.716404, 0.715424, 0.714375, 0.713258, 0.712076, 0.710829, 0.709522, 0.708155, 0.706730, 0.705251, 0.703718, 0.702133, 0.700499, 0.698816, 0.697088, 0.695315, 0.693499, 0.691641, 0.689742, 0.687805, 0.685829, 0.683818, 0.681769, 0.679687, 0.677571, 0.675423, 0.673242, 0.671031, 0.668791, 0.666525, 0.664237, 0.662342, 0.660514, 0.658743, 0.657022, 0.655347, 0.653715, 0.652122, 0.650567, 0.649048, 0.647565, 0.646117, 0.644703, 0.643326, 0.641984, 0.640678, 0.639409, 0.638178, 0.636987, 0.635837, 0.634730, 0.633668, 0.632655, 0.631695, 0.630791, 0.629948, 0.629175, 0.628479, 0.627873, 0.627375, 0.627012, 0.628162, 0.629312, 0.630462, 0.631611, 0.632761, 0.633911, 0.632761, 0.631611, 0.630462, 0.629312, 0.628162, 0.627012, 0.627375, 0.627873, 0.628479, 0.629175, 0.629948, 0.630791, 0.631695, 0.632655, 0.633668, 0.634730, 0.635837, 0.636987, 0.638178, 0.639409, 0.640678, 0.641984, 0.643326, 0.644703, 0.646117, 0.647565, 0.649048, 0.650567, 0.652122, 0.653715, 0.655347, 0.657022, 0.658743, 0.660514, 0.662342, 0.665075, 0.665788, 0.666477, 0.667149, 0.667807, 0.668452, 0.669084, 0.669701, 0.670303, 0.670890, 0.671461, 0.672014, 0.672548, 0.673062, 0.673552, 0.674017, 0.674454, 0.674860, 0.675234, 0.675572, 0.675871, 0.675572, 0.675234, 0.674860, 0.674454, 0.674017, 0.673552, 0.673062, 0.672548, 0.672014, 0.671461, 0.670890, 0.670303, 0.669701, 0.669084, 0.668452, 0.667807, 0.667149, 0.666477, 0.665788, 0.665075, 0.664381, 0.663748, 0.663164, 0.662625, 0.662130, 0.661679, 0.661275, 0.660920, 0.660618, 0.660372, 0.660186, 0.660063, 0.660004, 0.660012, 0.660086, 0.660229, 0.660443, 0.660728, 0.661089, 0.661532, 0.661089, 0.660728, 0.660443, 0.660229, 0.660086, 0.660012, 0.660004, 0.660063, 0.660186, 0.660372, 0.660618, 0.660920, 0.661275, 0.661679, 0.662130, 0.662625, 0.663164, 0.663748, 0.664381, 0.375535, 0.380755, 0.385970, 0.391178, 0.396378, 0.401568, 0.406749, 0.411919, 0.417077, 0.422223, 0.427357, 0.432478, 0.437586, 0.442679, 0.447758, 0.452822, 0.457871, 0.462903, 0.467918, 0.472916, 0.477896, 0.482858, 0.487799, 0.492722, 0.497623, 0.502504, 0.507362, 0.512198, 0.517010, 0.521792, 0.526540, 0.531243, 0.559981, 0.564298, 0.572818, 0.564298, 0.559981, 0.531243, 0.526540, 0.521792, 0.517010, 0.512198, 0.507362, 0.502504, 0.497623, 0.492722, 0.487799, 0.482858, 0.477896, 0.472916, 0.467918, 0.462903, 0.457871, 0.452822, 0.447758, 0.442679, 0.437586, 0.432478, 0.427357, 0.422223, 0.417077, 0.411919, 0.406749, 0.401568, 0.396378, 0.391178, 0.385970, 0.380755, 0.375535, 0.543598, 0.549005, 0.549005, 0.543598, 0.521853, 0.516702, 0.511551, 0.506400, 0.501249, 0.496098, 0.490947, 0.485796, 0.480645, 0.475494, 0.470343, 0.465192, 0.460041, 0.454890, 0.449739, 0.444588, 0.439437, 0.434286, 0.429135, 0.423984, 0.418833, 0.413682, 0.408531, 0.403380, 0.398229, 0.393078, 0.387927, 0.382776, 0.377625, 0.372474, 0.367323, 0.362172, 0.357021, 0.362533, 0.362584, 0.366719, 0.370758, 0.374482, 0.377860, 0.380939, 0.383295, 0.385031, 0.386113, 0.386473, 0.386109, 0.385023, 0.383283, 0.380925, 0.377842, 0.374461, 0.370736, 0.366696, 0.362560, 0.358348, 0.354309, 0.350585, 0.347206, 0.344127, 0.341772, 0.340035, 0.338954, 0.338593, 0.338958, 0.340043, 0.341783, 0.344142, 0.347224, 0.350605, 0.354331, 0.358371, 0.375192, 0.372191, 0.369350, 0.366674, 0.364163, 0.361818, 0.359640, 0.357632, 0.355795, 0.354130, 0.352640, 0.351325, 0.350186, 0.349223, 0.348437, 0.347827, 0.347393, 0.347135, 0.347052, 0.347144, 0.347412, 0.347854, 0.348473, 0.349267, 0.350238, 0.351387, 0.352715, 0.354223, 0.355912, 0.357785, 0.359846, 0.362104, 0.364585, 0.367095, 0.369405, 0.371520, 0.373447, 0.375190, 0.376755, 0.378148, 0.379374, 0.380440, 0.381352, 0.382117, 0.382743, 0.383236, 0.383605, 0.383856, 0.383996, 0.384031, 0.383968, 0.383811, 0.383568, 0.383241, 0.382836, 0.382357, 0.381806, 0.381188, 0.380505, 0.379760, 0.378955, 0.378093, 0.377176, 0.376206, 0.362533, 0.362560, 0.366696, 0.370736, 0.374461, 0.377842, 0.380925, 0.383283, 0.385023, 0.386109, 0.386473, 0.386113, 0.385031, 0.383295, 0.380939, 0.377860, 0.374482, 0.370758, 0.366719, 0.362584, 0.358371, 0.354331, 0.350605, 0.347224, 0.344142, 0.341783, 0.340043, 0.338958, 0.338593, 0.338954, 0.340035, 0.341772, 0.344127, 0.347206, 0.350585, 0.354309, 0.358348, 0.375192, 0.372191, 0.369350, 0.366674, 0.364163, 0.361818, 0.359640, 0.357632, 0.355795, 0.354130, 0.352640, 0.351325, 0.350186, 0.349223, 0.348437, 0.347827, 0.347393, 0.347135, 0.347052, 0.347144, 0.347412, 0.347854, 0.348473, 0.349267, 0.350238, 0.351387, 0.352715, 0.354223, 0.355912, 0.357785, 0.359846, 0.362104, 0.364585, 0.367095, 0.369405, 0.371520, 0.373447, 0.375190, 0.376755, 0.378148, 0.379374, 0.380440, 0.381352, 0.382117, 0.382743, 0.383236, 0.383605, 0.383856, 0.383996, 0.384031, 0.383968, 0.383811, 0.383568, 0.383241, 0.382836, 0.382357, 0.381806, 0.381188, 0.380505, 0.379760, 0.378955, 0.378093, 0.377176, 0.376206, 0.298620, 0.293794, 0.288991, 0.284301, 0.279916, 0.276030, 0.272714, 0.269943, 0.267662, 0.265804, 0.264290, 0.263036, 0.261967, 0.261022, 0.260156, 0.259339, 0.258551, 0.257778, 0.257011, 0.257406, 0.257815, 0.258259, 0.258769, 0.259392, 0.260193, 0.261257, 0.262688, 0.264582, 0.266996, 0.269911, 0.273238, 0.276859, 0.280671, 0.284596, 0.288585, 0.292606, 0.296640, 0.296323, 0.296006, 0.295688, 0.295372, 0.295058, 0.294752, 0.294461, 0.294192, 0.293955, 0.293756, 0.293596, 0.293472, 0.293376, 0.293299, 0.293234, 0.293177, 0.293123, 0.293071, 0.293667, 0.294260, 0.294850, 0.295430, 0.295996, 0.296538, 0.297043, 0.297496, 0.297883, 0.298195, 0.298431, 0.298597, 0.298704, 0.298760, 0.298773, 0.298749, 0.298694, 0.296640, 0.292606, 0.288585, 0.284596, 0.280671, 0.276859, 0.273238, 0.269911, 0.266996, 0.264582, 0.262688, 0.261257, 0.260193, 0.259392, 0.258769, 0.258259, 0.257815, 0.257406, 0.257011, 0.257778, 0.258551, 0.259339, 0.260156, 0.261022, 0.261967, 0.263036, 0.264290, 0.265804, 0.267662, 0.269943, 0.272714, 0.276030, 0.279916, 0.284301, 0.288991, 0.293794, 0.298620, 0.298694, 0.298749, 0.298773, 0.298760, 0.298704, 0.298597, 0.298431, 0.298195, 0.297883, 0.297496, 0.297043, 0.296538, 0.295996, 0.295430, 0.294850, 0.294260, 0.293667, 0.293071, 0.293123, 0.293177, 0.293234, 0.293299, 0.293376, 0.293472, 0.293596, 0.293756, 0.293955, 0.294192, 0.294461, 0.294752, 0.295058, 0.295372, 0.295688, 0.296006, 0.296323 +]) +landmarks_2D_1k = np.stack([mean_face_x_1k, mean_face_y_1k], axis=1) + +mean_face_x_137_22_client = np.array([0.4988282, 0.5449964, 0.5849726, 0.6179804, 0.643469, 0.6622178, 0.6730388, 0.676355, 0.6733304, + 0.6478118, + 0.5882786, + 0.4988282, + 0.4093778, + 0.349844, + 0.324326, + 0.32130139999999996, + 0.3246176, + 0.33543860000000003, + 0.35418740000000004, + 0.3796754, + 0.4126838, + 0.45266]) + +mean_face_y_137_22_client = np.array([0.7108352, 0.7000166, 0.6745382, 0.640106, 0.5996582, 0.5467124, 0.4916804, 0.4355282, 0.3795278, 0.2916416, + 0.23122520000000002, + 0.2137676, + 0.23122520000000002, + 0.2916416, + 0.3795278, + 0.4355282, + 0.4916804, + 0.5467124, + 0.5996582, + 0.640106, + 0.6745382, + 0.7000166]) + +landmarks_2D_137_22_clinet = np.stack([mean_face_x_137_22_client, mean_face_y_137_22_client], axis=1) + + +mat_face1024_256_full_face_client = np.array([[4.1666666e-01, 1.5257437e-17, -8.5333336e+01], + [-1.5257449e-17, 4.1666666e-01, -8.5333336e+01]]) + +mat_face1024_256_full_face_server = np.array([[4.1666666e-01, -1.5237085e-17, -8.5333336e+01], + [1.5237085e-17, 4.1666666e-01, -8.5333336e+01]]) + + +# 68 point landmark definitions +landmarks_68_pt = {"mouth": (48, 68), + "right_eyebrow": (17, 22), + "left_eyebrow": (22, 27), + "right_eye": (36, 42), + "left_eye": (42, 48), + "nose": (27, 36), # missed one point + "jaw": (0, 17)} + + +def get_max_rect(bounding_boxes): + max_area = 0 + index = 0 + for i, box in enumerate(bounding_boxes): + width = box[2] - box[0] + height = box[3] - box[1] + if width * height > max_area: + index = i + max_area = width * height + return index + + +def get_transform_mat_mmcv(landmark, output_size): + dst_size = output_size + + if len(landmark) == 68: + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[36] + landmark[39]) / 2 + right_eye = (landmark[42] + landmark[45]) / 2 + nose = (landmark[31] + landmark[35]) / 2 + left_mouth = (landmark[48] + landmark[60]) / 2 + right_mouth = (landmark[64] + landmark[54]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + elif len(landmark) == 87: + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[31] + landmark[35]) / 2 + right_eye = (landmark[39] + landmark[43]) / 2 + nose = landmark[62] + left_mouth = (landmark[66] + landmark[79]) / 2 + right_mouth = (landmark[72] + landmark[83]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + elif len(landmark) == 1000: + pt137 = pts_1k_to_137(landmark) + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (pt137[88] + pt137[96]) / 2 + right_eye = (pt137[105] + pt137[113]) / 2 + nose = pt137[83] + left_mouth = (pt137[22] + pt137[48]) / 2 + right_mouth = (pt137[56] + pt137[36]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + + +def umeyama(src, dst, estimate_scale): + """Estimate N-D similarity transformation with or without scaling. + Parameters + ---------- + src : (M, N) array + Source coordinates. + dst : (M, N) array + Destination coordinates. + estimate_scale : bool + Whether to estimate scaling factor. + Returns + ------- + T : (N + 1, N + 1) + The homogeneous similarity transformation matrix. The matrix contains + NaN values only if the problem is not well-conditioned. + References + ---------- + .. [1] "Least-squares estimation of transformation parameters between two + point patterns", Shinji Umeyama, PAMI 1991, DOI: 10.1109/34.88573 + """ + + num = src.shape[0] + dim = src.shape[1] + + # Compute mean of src and dst. + src_mean = src.mean(axis=0) + dst_mean = dst.mean(axis=0) + + # Subtract mean from src and dst. + src_demean = src - src_mean + dst_demean = dst - dst_mean + + # Eq. (38). + A = np.dot(dst_demean.T, src_demean) / num + + # Eq. (39). + d = np.ones((dim,), dtype=np.double) + if np.linalg.det(A) < 0: + d[dim - 1] = -1 + + T = np.eye(dim + 1, dtype=np.double) + + U, S, V = np.linalg.svd(A) + + # Eq. (40) and (43). + rank = np.linalg.matrix_rank(A) + if rank == 0: + return np.nan * T + elif rank == dim - 1: + if np.linalg.det(U) * np.linalg.det(V) > 0: + T[:dim, :dim] = np.dot(U, V) + else: + s = d[dim - 1] + d[dim - 1] = -1 + T[:dim, :dim] = np.dot(U, np.dot(np.diag(d), V)) + d[dim - 1] = s + else: + T[:dim, :dim] = np.dot(U, np.dot(np.diag(d), V.T)) + + if estimate_scale: + # Eq. (41) and (42). + scale = 1.0 / src_demean.var(axis=0).sum() * np.dot(S, d) + else: + scale = 1.0 + + T[:dim, dim] = dst_mean - scale * np.dot(T[:dim, :dim], src_mean.T) + T[:dim, :dim] *= scale + + return T + +def get_transform_mat_mmcv_bigger(landmark, output_size, forlabel=False): + dst_size = output_size + if len(landmark) == 87: + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[31] + landmark[35]) / 2 + right_eye = (landmark[39] + landmark[43]) / 2 + nose = landmark[62] + left_mouth = (landmark[66] + landmark[79]) / 2 + right_mouth = (landmark[72] + landmark[83]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + elif len(landmark) == 137: + if forlabel: + eye_dis = 0.4 + mouth_dis = 0.4 + g_Average_5point_180 = np.array([ + eye_dis, 0.4, + 1 - eye_dis, 0.4, + 0.5, 0.5, + mouth_dis, 0.6, + 1 - mouth_dis, 0.6 + ]) + else: + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[96] + landmark[88]) / 2 + right_eye = (landmark[105] + landmark[113]) / 2 + nose = landmark[83] + left_mouth = (landmark[22] + landmark[48]) / 2 + right_mouth = (landmark[56] + landmark[36]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + elif len(landmark) == 1000: + # left_eye = (landmarks_2D_1k[691] + landmarks_2D_1k[723]) / 2 + # right_eye = (landmarks_2D_1k[792] + landmarks_2D_1k[824]) / 2 + # nose = landmarks_2D_1k[621] + # left_mouth = (landmarks_2D_1k[467] + landmarks_2D_1k[468]) / 2 + # right_mouth = (landmarks_2D_1k[396] + landmarks_2D_1k[508]) / 2 + # pts5_dst = np.vstack((left_eye, right_eye, + # nose, + # left_mouth, right_mouth)) * dst_size + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + # print(g_Average_5point_180) + left_eye = (landmark[691] + landmark[723]) / 2 + right_eye = (landmark[792] + landmark[824]) / 2 + nose = landmark[621] + left_mouth = (landmark[467] + landmark[468]) / 2 + right_mouth = (landmark[396] + landmark[508]) / 2 + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + + image_to_face_mat = umeyama(pts5_src, pts5_dst, True)[:2] + + return image_to_face_mat + +def get_transform_mat_full_face(landmark, output_size): + dst_size = output_size + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k) + + # print("landmarks_2D_137 x: ", landmarks_2D_137[:22, 0]) + # print("landmarks_2D_137 y: ", landmarks_2D_137[:22, 1]) + + # exit() + + mat = umeyama(landmark[:22], landmarks_2D_137[:22] * dst_size, True)[0:2] + return mat + + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k[:312] * dst_size, True)[0:2] + return mat + + elif len(landmark) == 22: + landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k) + + mat = umeyama(landmark, landmarks_2D_137[:22] * dst_size, True)[0:2] + return mat + +def get_transform_mat_full_face_592(landmark, output_size, ratio): + dst_size = output_size + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + 0.5 + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.5) * ratio + 0.5 + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_137[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + elif len(landmark) == 22: + landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark, landmarks_2D_137[:22] * dst_size, True)[0:2] + return mat + +def get_transform_mat_full_face_592_v1(landmark, output_size, ratio=1.0, h_ratio=0.57): + dst_size = output_size + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + 0.5 + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.5) * ratio + h_ratio + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_137[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_full_face_592_client(landmark, output_size, ratio): + dst_size = output_size + landmarks_2D_137_22_clinet_tmp = landmarks_2D_137_22_clinet.copy() + landmarks_2D_137_22_clinet_tmp[:, 0] = (landmarks_2D_137_22_clinet_tmp[:, 0] - 0.5) * ratio + 0.5 + landmarks_2D_137_22_clinet_tmp[:, 1] = (landmarks_2D_137_22_clinet_tmp[:, 1] - 0.5) * ratio + 0.5 + + if len(landmark) == 22: + # landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark, landmarks_2D_137_22_clinet_tmp * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + landmark_137 = pts_1k_to_137(landmark) + + mat = umeyama(landmark_137[:22], landmarks_2D_137_22_clinet_tmp * dst_size, True)[0:2] + return mat + +def get_transform_mat_face592_full_face(img_size, detect_single_face_size, ratio): + + # face_landmark_client = landmarks_2D_137_22_clinet * img_size + # + # #TODO:客户端如何得到face 0.6 + # face_size2face_M_ratio = get_transform_mat_full_face_592_client(face_landmark_client, img_size, ratio) + + #服务端得到face 1 + face_landmark_server = landmarks_2D_1k * img_size + + face_size2face_M_ratio = get_transform_mat_full_face_592(face_landmark_server, img_size, ratio) + face_size2face_M_full = get_transform_mat_full_face(face_landmark_server, detect_single_face_size) + + + M_ori = np.zeros((3, 3), dtype=np.float32) + M_ori[:2, :] = cv2.invertAffineTransform(face_size2face_M_ratio) + M_ori[2:, :] = [0, 0, 1] + + matAffine_ori = np.zeros((3, 3), dtype=np.float32) + matAffine_ori[:2, :] = face_size2face_M_full + matAffine_ori[2:, :] = [0, 0, 1] + + new_mat = matAffine_ori.dot(M_ori) + + if False: + img = np.zeros((img_size, img_size, 3), dtype=np.uint8) + pred_label_int = face_landmark_server.copy().astype(np.int32) + img_client = img.copy() + for pt in pred_label_int: + cv2.circle(img_client, (pt[0], pt[1]), 1, (0, 0, 255), 1) + cv2.imshow("img_client: ", img_client) + # + # img_server = np.zeros((img_size, img_size, 3), dtype=np.uint8) + # face_sever_22 = pts_1k_to_137(face_landmark_server)[:22].astype(np.int32) + # for pt in face_sever_22: + # cv2.circle(img_server, (pt[0], pt[1]), 1, (0, 255, 0), 1) + # cv2.imshow("img_server: ", img_server) + # cv2.waitKey() + # + # img_server_new = np.zeros((img_size, img_size, 3), dtype=np.uint8) + # face_sever_new = transform_points(face_landmark_server, face_size2face_M_full) + # face_sever_new_22 = pts_1k_to_137(face_sever_new)[:22].astype(np.int32) + # + # for pt in face_sever_new_22: + # cv2.circle(img_server_new, (pt[0], pt[1]), 1, (0, 255, 0), 1) + # cv2.imshow("img_server_new: ", img_server_new) + # cv2.waitKey() + face_landmark_server_new = transform_points(face_landmark_server, face_size2face_M_ratio) + img_server = np.zeros((img_size, img_size, 3), dtype=np.uint8) + pred_label_server_int = face_landmark_server_new.copy().astype(np.int32) + # img_client = img.copy() + for pt in pred_label_server_int: + cv2.circle(img_server, (pt[0], pt[1]), 1, (0, 0, 255), 1) + cv2.imshow("img_server: ", img_server) + + + img_new = np.zeros((detect_single_face_size, detect_single_face_size, 3), dtype=np.uint8) + pred_new_label_int = transform_points(face_landmark_server_new, new_mat[:2, :]).astype(np.int32) + for pt in pred_new_label_int: + cv2.circle(img_new, (pt[0], pt[1]), 1, (0, 255, 0), 1) + cv2.imshow("img_show new: ", img_new) + cv2.waitKey() + + return new_mat[:2, :] + +def get_transform_mat_for_eye(landmark, output_size): + dst_size = output_size + if len(landmark) == 137: + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[96] + landmark[88]) / 2 + right_eye = (landmark[105] + landmark[113]) / 2 + nose = landmark[83] + left_mouth = (landmark[22] + landmark[48]) / 2 + right_mouth = (landmark[56] + landmark[36]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + +def get_transform_mat_for_face_recognition(landmark, output_size): + g_Average_5point_180 = np.array([ + 57, 73, + 123, 73, + 90, 107, + 62, 134, + 118, 134 + ]) + dst_size = output_size + + if len(landmark) == 87: + left_eye = (landmark[17 + 19] + landmark[17 + 22]) / 2 + right_eye = (landmark[17 + 25] + landmark[17 + 28]) / 2 + nose = (landmark[17 + 14] + landmark[17 + 18]) / 2 + left_mouth = (landmark[17 + 31] + landmark[17 + 43]) / 2 + right_mouth = (landmark[17 + 47] + landmark[17 + 37]) / 2 + elif len(landmark) == 137: + left_eye = (landmark[88] + landmark[96]) / 2 + right_eye = (landmark[105] + landmark[113]) / 2 + nose = landmark[83] + left_mouth = (landmark[22] + landmark[48]) / 2 + right_mouth = (landmark[56] + landmark[36]) / 2 + elif len(landmark) == 1000: + left_eye = (landmark[691] + landmark[723]) / 2 + right_eye = (landmark[792] + landmark[824]) / 2 + nose = landmark[621] + left_mouth = (landmark[467] + landmark[468]) / 2 + right_mouth = (landmark[396] + landmark[508]) / 2 + else: + assert False + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) / 180 * dst_size + + mat = umeyama(pts5_src, pts5_dst, True) + + return mat + +def get_transform_mat_mmcv_hair(landmark, output_size, forlabel=False): + dst_size = output_size + + if len(landmark) == 1000: + eye_dis = 0.42 + mouth_dis = 0.42 + g_Average_5point_180 = np.array([ + eye_dis, 0.48, + 1 - eye_dis, 0.48, + 0.5, 0.53, + mouth_dis, 0.58, + 1 - mouth_dis, 0.58 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[691] + landmark[723]) / 2 + right_eye = (landmark[792] + landmark[824]) / 2 + nose = landmark[621] + left_mouth = (landmark[467] + landmark[468]) / 2 + right_mouth = (landmark[396] + landmark[508]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + else: + print("landmark < 1000 !!!") + assert False + return mat + + +def get_transform_mat_mmcv_seg(landmark, output_size, forlabel=False): + dst_size = output_size + eye_dis = 0.40 + mouth_dis = 0.40 + g_Average_5point_180 = np.array([ + eye_dis, 0.46, + 1 - eye_dis, 0.46, + 0.5, 0.55, + mouth_dis, 0.64, + 1 - mouth_dis, 0.64 + ]) + + if len(landmark) == 1000: + # print(g_Average_5point_180) + left_eye = (landmark[691] + landmark[723]) / 2 + right_eye = (landmark[792] + landmark[824]) / 2 + nose = landmark[621] + left_mouth = (landmark[467] + landmark[468]) / 2 + right_mouth = (landmark[396] + landmark[508]) / 2 + + elif len(landmark) == 137: + left_eye = (landmark[88] + landmark[96]) / 2 + right_eye = (landmark[105] + landmark[113]) / 2 + nose = landmark[83] + left_mouth = (landmark[22] + landmark[48]) / 2 + right_mouth = (landmark[56] + landmark[36]) / 2 + else: + print("landmark < 1000 !!!") + assert False + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + + +def get_transform_mat_two_sets(landmarks_src, landmarks_dst): + assert len(landmarks_src) == len(landmarks_dst) + assert len(landmarks_src) == 137 + mat = umeyama(landmarks_src[:22], landmarks_dst[:22], True)[0:2] + return mat + +def flip_points(landmark, width): + if len(landmark) == 137: + landmarks_order = np.array([1, 22, 21, 20, 19, 18, 17, 16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, # counter + 37, 36, 35, 34, 33, 32, 31, 30, 29, 28, 27, 26, 25, 24, 23, 48, 47, 46, 45, 44, 43, 42, 41, + 40, 39, 38, + 57, 56, 55, 54, 53, 52, 51, 50, 49, 64, 63, 62, 61, 60, 59, 58, # mouth + 79, 78, 77, 76, 75, 74, 73, 72, 71, 70, 69, 68, 67, 66, 65, 83, 82, 81, 80, 84, 85, 86, 87, + # nose + 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, # eye + 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, + 134, 133, 132, 131, 130, 137, 136, 135, # eyebrow + 126, 125, 124, 123, 122, 129, 128, 127], dtype=np.int32) - 1 + landmark_flip = landmark.copy() + landmark_flip[:, 0] = width - 1 - landmark_flip[:, 0] + landmark_flip = landmark_flip[landmarks_order, :] + return landmark_flip + elif len(landmark) == 1000: + dst = [0] + list(range(311, 156, -1)) + list(range(156, 0, -1)) + \ + list(range(396, 354, -1)) + list(range(354, 311, -1)) + list(range(467, 432, -1)) + list(range(432, 396, -1)) + \ + list(range(508, 488, -1)) + list(range(488, 467, -1)) + list(range(547, 528, -1)) + list(range(528, 508, -1)) + \ + list(range(616, 584, -1)) + [584, 583, 582, 581, 580] + list(range(579, 547, -1)) + [620, 619, 618, 617] + list(range(621, 654, 1)) + \ + [755] + list(range(774, 755, -1)) + list(range(791, 774, -1)) + list(range(792, 856, 1)) + \ + [654] + list(range(673, 654, -1)) + list(range(690, 673, -1)) + list(range(691, 755, 1)) + \ + list(range(964, 927, -1)) + list(range(999, 964, -1)) + \ + list(range(892, 855, -1)) + list(range(927, 892, -1)) + landmarks_order = np.array(dst, dtype=np.int32) + landmark_flip = landmark.copy() + landmark_flip[:, 0] = width - 1 - landmark_flip[:, 0] + landmark_flip = landmark_flip[landmarks_order, :] + return landmark_flip + else: + assert False + +def pts_1k_to_137(landmarks): + index_1k_to_137 = [0, 12, 24, 36, 48, 61, 74, 87, 100, 119, 137, 156, 175, 193, 212, 225, 238, 251, 264, 276, 288, + 300, 312, 318, 324, 330, 336, 342, 348, 354, 360, 366, 372, 378, 384, 390, 396, 402, 408, 414, + 420, 426, 432, 438, 444, 450, 456, 462, 468, 473, 478, 483, 488, 493, 498, 503, 508, 513, 518, + 523, 528, 533, 538, 543, 548, 556, 564, 571, 579, 580, 581, 582, 583, 584, 585, 593, 600, 608, + 616, 617, 618, 619, 620, 621, 632, 642, 653, 654, 691, 695, 699, 703, 707, 711, 715, 719, 723, + 727, 731, 735, 739, 743, 747, 751, 755, 792, 796, 800, 804, 808, 812, 816, 820, 824, 828, 832, + 836, 840, 844, 848, 852, 856, 865, 874, 883, 892, 901, 910, 919, 928, 937, 946, 955, 964, 973, + 982, 991] + landmarks_137 = landmarks[index_1k_to_137, :] + return landmarks_137 + +def get_transform_mat_full_face_to_target(landmark, dst_pt): + mat = umeyama(landmark[:312], dst_pt[:312], True)[0:2] + return mat + + +def decompose_affine_matrix(matrix): + # 确保输入矩阵是 2x3 的 + assert matrix.shape == (2, 3), "输入矩阵必须是 2x3 的仿射变换矩阵" + + # 提取平移分量 + tx = matrix[0, 2] + ty = matrix[1, 2] + + # 提取旋转、缩放分量 + a = matrix[0, 0] + b = matrix[0, 1] + c = matrix[1, 0] + d = matrix[1, 1] + + # 计算缩放因子 + scale_x = np.sqrt(a ** 2 + c ** 2) + scale_y = np.sqrt(b ** 2 + d ** 2) + + # 计算旋转角度 + theta = np.arctan2(c, a) + + return { + "translation": (tx, ty), + "scale": (scale_x, scale_y), + "rotation": np.degrees(theta) # 以度数表示 + } + + +# 定义一个高质量缩放图像的函数，主要是用于处理比如头发这种非常细微的图片 +def high_quality_warpAffine(origin_img, M, dst_size, const_value=(255, 255, 255)): + M_params = decompose_affine_matrix(M) + + scale_ratio = M_params['scale'][0] + M_big = M / scale_ratio + crop_face_big = cv2.warpAffine(origin_img, M_big, (int(dst_size[0] / scale_ratio), int(dst_size[1] / scale_ratio)), + cv2.BORDER_CONSTANT, borderValue=const_value) + + if scale_ratio < 1: + crop_face = cv2.resize(crop_face_big, (dst_size[0], dst_size[1]), interpolation=cv2.INTER_AREA) + else: + crop_face = cv2.resize(crop_face_big, (dst_size[0], dst_size[1]), interpolation=cv2.INTER_LANCZOS4) + + return crop_face + + +def transform_points(points, mat, invert=False): + if invert: + mat = cv2.invertAffineTransform(mat) + points = np.expand_dims(points, axis=1) + points = cv2.transform(points, mat, points.shape) + points = np.squeeze(points) + return points +def get_transform_mat_bodyseg(landmark, output_size, ratio=1.0, offset = (0, 0)): + dst_size = output_size + assert len(landmark) == 1000 + mat = umeyama(landmark[:312], landmarks_2D_1k[:312] * dst_size * ratio + dst_size * (1 - ratio) / 2 + offset, True)[0:2] + return mat +def align2stylegan(face_landmarks_1k, output_size=256): + face_landmarks_1k = np.float32(face_landmarks_1k) + x_scale = 1.0 + y_scale = 1.0 + em_scale = 0.1 + eye_left = (face_landmarks_1k[691] + face_landmarks_1k[723]) / 2 + eye_right = (face_landmarks_1k[792] + face_landmarks_1k[824]) / 2 + mouth_left = (face_landmarks_1k[467] + face_landmarks_1k[468]) / 2 + mouth_right = (face_landmarks_1k[396] + face_landmarks_1k[508]) / 2 + eye_avg = (eye_left + eye_right) * 0.5 + eye_to_eye = eye_right - eye_left + mouth_avg = (mouth_left + mouth_right) * 0.5 + eye_to_mouth = mouth_avg - eye_avg + x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1] + x /= np.hypot(*x) + x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8) + x *= x_scale + y = np.flipud(x) * [-y_scale, y_scale] + c = eye_avg + eye_to_mouth * em_scale + quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y]) + quad_ori = np.array(quad) + + rotate_radian = math.atan2((quad_ori[3][1] - quad_ori[0][1]), (quad_ori[3][0] - quad_ori[0][0])) + rotate_degree = rotate_radian / np.pi * 180 + scale = output_size / cv2.norm(quad_ori[3] - quad_ori[0]) + src_center = (quad_ori[0] + quad_ori[2]) * 0.5 + dst_center = np.float32([output_size / 2, output_size / 2]) + + M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), rotate_degree, scale) + M[:, 2] += dst_center - src_center + return M +def align2stylegan_ratio(face_landmarks_1k, ratio=1.0, output_size=256): + face_landmarks_1k = np.float32(face_landmarks_1k) + x_scale = 1.0 / ratio + y_scale = 1.0 / ratio + em_scale = 0.1 + eye_left = (face_landmarks_1k[691] + face_landmarks_1k[723]) / 2 + eye_right = (face_landmarks_1k[792] + face_landmarks_1k[824]) / 2 + mouth_left = (face_landmarks_1k[467] + face_landmarks_1k[468]) / 2 + mouth_right = (face_landmarks_1k[396] + face_landmarks_1k[508]) / 2 + eye_avg = (eye_left + eye_right) * 0.5 + eye_to_eye = eye_right - eye_left + mouth_avg = (mouth_left + mouth_right) * 0.5 + eye_to_mouth = mouth_avg - eye_avg + x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1] + x /= np.hypot(*x) + x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8) + x *= x_scale + y = np.flipud(x) * [-y_scale, y_scale] + c = eye_avg + eye_to_mouth * em_scale + quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y]) + quad_ori = np.array(quad) + + rotate_radian = math.atan2((quad_ori[3][1] - quad_ori[0][1]), (quad_ori[3][0] - quad_ori[0][0])) + rotate_degree = rotate_radian / np.pi * 180 + scale = output_size / cv2.norm(quad_ori[3] - quad_ori[0]) + src_center = (quad_ori[0] + quad_ori[2]) * 0.5 + dst_center = np.float32([output_size / 2, output_size / 2]) + + M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), rotate_degree, scale) + M[:, 2] += dst_center - src_center + return M +def calc_face_pitch(landmarks): + if not isinstance(landmarks, np.ndarray): + landmarks = np.array(landmarks) + t = ((landmarks[6][1] - landmarks[8][1]) + (landmarks[10][1] - landmarks[8][1])) / 2.0 + b = landmarks[8][1] + return float(b - t) + +def calc_face_yaw(landmarks): + if not isinstance(landmarks, np.ndarray): + landmarks = np.array(landmarks) + l = ((landmarks[27][0] - landmarks[0][0]) + (landmarks[28][0] - landmarks[1][0]) + ( + landmarks[29][0] - landmarks[2][0])) / 3.0 + r = ((landmarks[16][0] - landmarks[27][0]) + (landmarks[15][0] - landmarks[28][0]) + ( + landmarks[14][0] - landmarks[29][0])) / 3.0 + return float(r - l) + +# deprecated +def draw_pncc_features(fc_landmark, img_target, w=256, h=256, is_train=True): + assert False + return img_target + + +def draw_blur_no_mouth_img(img_target, fc_landmark, is_train=False): + h, w, c = img_target.shape + assert h == w + + if len(fc_landmark) == 1000: + fc_landmark = pts_1k_to_137(fc_landmark) + + + k_mid_size = int(w / 6.) + thresh = int(w / 15.) + if is_train: + k_size = random.choice(range(k_mid_size - thresh, k_mid_size + thresh, 2)) + else: + k_size = k_mid_size + if k_size % 2 == 0: + k_size += 1 + img_target_blur = cv2.blur(img_target, (k_size, k_size)) + + + hull_mask = np.zeros(img_target.shape, dtype=np.uint8) + if len(fc_landmark) == 1000: + fc_landmark = pts_1k_to_137(fc_landmark) + + cv2.fillPoly(hull_mask, fc_landmark[22:48][np.newaxis, :, :], (255, 255, 255)) + # cv2.imshow('hull_mask1', hull_mask) + # cv2.imshow('img_target_blur', img_target_blur) + # cv2.imshow('img_target1', img_target) + kernel_size = int(w / 35.) + if kernel_size % 2 == 0: + kernel_size += 1 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + hull_mask = cv2.dilate(hull_mask, kernel, iterations=1) .astype(np.uint8) + # cv2.imshow('hull_mask2', hull_mask) + img_target = img_target_blur * (hull_mask<125) + img_target * (hull_mask>125) + + # cv2.imshow('img_target2', img_target) + # cv2.waitKey() + + return img_target + + + +def draw_blur_img(img_target, fc_landmark, is_train=False): + h, w, c = img_target.shape + assert h == w + inpaint_mask = np.zeros((h, w), dtype=np.uint8) + + if len(fc_landmark) == 1000: + fc_landmark = pts_1k_to_137(fc_landmark) + + # cv2.imshow('img_target1', img_target) + + k_mid_size = int(w / 6.) + thresh = int(w / 15.) + if is_train: + k_size = random.choice(range(k_mid_size - thresh, k_mid_size + thresh, 2)) + else: + k_size = k_mid_size + if k_size % 2 == 0: + k_size += 1 + img_target = cv2.blur(img_target, (k_size, k_size)) + + # cv2.imshow('img_target2', img_target) + # cv2.waitKey() + + return img_target + +def draw_blurmore_img(img_target, fc_landmark, is_train=False): + h, w, c = img_target.shape + assert h == w + inpaint_mask = np.zeros((h, w), dtype=np.uint8) + + if len(fc_landmark) == 1000: + fc_landmark = pts_1k_to_137(fc_landmark) + + # cv2.imshow('img_target1', img_target) + + k_mid_size = int(w / 3.) + thresh = int(w / 15.) + if is_train: + k_size = random.choice(range(k_mid_size - thresh, k_mid_size + thresh, 2)) + else: + k_size = k_mid_size + if k_size % 2 == 0: + k_size += 1 + img_target = cv2.blur(img_target, (k_size, k_size)) + + # cv2.imshow('img_target2', img_target) + # cv2.waitKey() + + return img_target + + +def draw_comapre_mask(fc_landmark, w=256, h=256): + inpaint_mask = np.zeros((h, w), dtype=np.uint8) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(fc_landmark[129:137]), (255,)) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(fc_landmark[121:129]), (255,)) + + left_eye = np.around((fc_landmark[96] + fc_landmark[88]) / 2).astype(np.int32) + right_eye = np.around((fc_landmark[105] + fc_landmark[113]) / 2).astype(np.int32) + len = cv2.norm(fc_landmark[96].astype(np.int32) - fc_landmark[88].astype(np.int32)) + cv2.circle(inpaint_mask, (left_eye[0], left_eye[1]), 1, (255), int(len * 1.1)) + len = cv2.norm(fc_landmark[105].astype(np.int32) - fc_landmark[113].astype(np.int32)) + + cv2.circle(inpaint_mask, (right_eye[0], right_eye[1]), 1, (255), int(len * 1.1)) + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(fc_landmark[64:87]), (255,)) + + cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(fc_landmark[22:64]), (255,)) + + dilate_kernel_size = int(w / 23.) + if dilate_kernel_size % 2 == 0: + dilate_kernel_size += 1 + kernel = np.ones((dilate_kernel_size, dilate_kernel_size), np.uint8) + inpaint_mask = cv2.dilate(inpaint_mask, kernel) + + inpaint_mask = 1 - inpaint_mask.astype(np.float32) / 255 + + return inpaint_mask + +def draw_hull_mask(fc_landmark, w=256, h=256, is_gray=False): + if not is_gray: + hull_mask = np.zeros((h, w, 3), dtype=np.uint8) + + line_size1 = int(w / 85.) + line_size2 = int(w / 51.) + + if len(fc_landmark) == 137: + # cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark[0:22]), (64, 16, 32)) + + # left_brown = ((fc_landmark[129] + fc_landmark[133]) / 2).astype(np.int32) + # cv2.circle(hull_mask, (left_brown[0], left_brown[1]), 1, (0, 255, 0), line_size1) + # right_brwon = ((fc_landmark[121] + fc_landmark[125]) / 2).astype(np.int32) + # cv2.circle(hull_mask, (right_brwon[0], right_brwon[1]), 1, (0, 0, 255), line_size1) + + cv2.line(hull_mask, (int(fc_landmark[129, 0]), int(fc_landmark[129, 1])), + (int(fc_landmark[133, 0]), int(fc_landmark[133, 1])), (0, 255, 0), line_size1) + cv2.line(hull_mask, (int(fc_landmark[121, 0]), int(fc_landmark[121, 1])), + (int(fc_landmark[125, 0]), int(fc_landmark[125, 1])), (0, 0, 255), line_size1) + + cv2.fillPoly(hull_mask, fc_landmark[88:104][np.newaxis, :, :], (128, 128, 0)) + cv2.fillPoly(hull_mask, fc_landmark[105:121][np.newaxis, :, :], (128, 0, 128)) + + # cv2.line(hull_mask, (int(fc_landmark[86, 0]), int(fc_landmark[86, 1])), + # (int(fc_landmark[83, 0]), int(fc_landmark[83, 1])), (255, 0, 0), line_size2) + + cv2.fillPoly(hull_mask, fc_landmark[48:64][np.newaxis, :, :], (0, 128, 128)) + + # cv2.fillPoly(hull_mask, np.concatenate((fc_landmark[22:37], fc_landmark[56:47:-1]))[np.newaxis, :, :], + # (0, 128, 0)) + # + # cv2.fillPoly(hull_mask, + # np.concatenate((fc_landmark[47:35:-1], fc_landmark[56:64], [fc_landmark[48], fc_landmark[22]]))[ + # np.newaxis, :, :], (0, 0, 128)) + + eye = np.zeros((h, w, 3), dtype=np.uint8) + eye_mask1 = np.zeros((h, w, 1), dtype=np.uint8) + eye_mask2 = np.zeros((h, w, 1), dtype=np.uint8) + cv2.fillPoly(eye_mask1, fc_landmark[88:104][np.newaxis, :, :], (1,)) + cv2.fillPoly(eye_mask1, fc_landmark[105:121][np.newaxis, :, :], (1,)) + left_eye = fc_landmark[87] + right_eye = fc_landmark[104] + # left_eye = np.around((fc_landmark[96] + fc_landmark[88]) / 2).astype(np.int32) + # right_eye = np.around((fc_landmark[105] + fc_landmark[113]) / 2).astype(np.int32) + cv2.circle(eye_mask2, (left_eye[0], left_eye[1]), 1, (1,), line_size2) + cv2.circle(eye_mask2, (right_eye[0], right_eye[1]), 1, (1,), line_size2) + eye_mask = eye_mask1 & eye_mask2 + cv2.circle(eye, (left_eye[0], left_eye[1]), 1, (255, 255, 255), line_size2) + cv2.circle(eye, (right_eye[0], right_eye[1]), 1, (255, 255, 255), line_size2) + hull_mask = hull_mask * (1 - eye_mask[:, :, 0:1]) + eye * eye_mask[:, :, 0:1] + elif len(fc_landmark) == 1000: + # cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark[0:312]), (64, 16, 32)) + + # left_brown = ((fc_landmark[928] + fc_landmark[964]) / 2).astype(np.int32) + # cv2.circle(hull_mask, (left_brown[0], left_brown[1]), 1, (0, 255, 0), line_size1) + # right_brwon = ((fc_landmark[856] + fc_landmark[892]) / 2).astype(np.int32) + # cv2.circle(hull_mask, (right_brwon[0], right_brwon[1]), 1, (0, 0, 255), line_size1) + + cv2.line(hull_mask, (int(fc_landmark[928, 0]), int(fc_landmark[928, 1])), + (int(fc_landmark[964, 0]), int(fc_landmark[964, 1])), (0, 255, 0), line_size1) + cv2.line(hull_mask, (int(fc_landmark[856, 0]), int(fc_landmark[856, 1])), + (int(fc_landmark[892, 0]), int(fc_landmark[892, 1])), (0, 0, 255), line_size1) + + cv2.fillPoly(hull_mask, fc_landmark[691:755][np.newaxis, :, :], (128, 128, 0)) + cv2.fillPoly(hull_mask, fc_landmark[792:856][np.newaxis, :, :], (128, 0, 128)) + + # cv2.line(hull_mask, (int(fc_landmark[653, 0]), int(fc_landmark[653, 1])), + # (int(fc_landmark[621, 0]), int(fc_landmark[621, 1])), (255, 0, 0), line_size2) + + cv2.fillPoly(hull_mask, fc_landmark[468:548][np.newaxis, :, :], (0, 128, 128)) + + # cv2.fillPoly(hull_mask, np.concatenate((fc_landmark[312:397], fc_landmark[508:467:-1]))[np.newaxis, :, :], + # (0, 128, 0)) + # + # cv2.fillPoly(hull_mask, + # np.concatenate((fc_landmark[467:395:-1], fc_landmark[508:548], [fc_landmark[468], fc_landmark[312]]))[ + # np.newaxis, :, :], (0, 0, 128)) + + eye = np.zeros((h, w, 3), dtype=np.uint8) + eye_mask1 = np.zeros((h, w, 1), dtype=np.uint8) + eye_mask2 = np.zeros((h, w, 1), dtype=np.uint8) + cv2.fillPoly(eye_mask1, fc_landmark[691:755][np.newaxis, :, :], (1,)) + cv2.fillPoly(eye_mask1, fc_landmark[792:856][np.newaxis, :, :], (1,)) + left_eye = fc_landmark[654] + right_eye = fc_landmark[755] + # cv2.fillPoly(eye_mask2, fc_landmark[655:691][np.newaxis, :, :], (1,)) + cv2.circle(eye_mask2, (left_eye[0], left_eye[1]), 1, (1,), line_size2) + # cv2.fillPoly(eye_mask2, fc_landmark[756:792][np.newaxis, :, :], (1,)) + cv2.circle(eye_mask2, (right_eye[0], right_eye[1]), 1, (1,), line_size2) + eye_mask = eye_mask1 & eye_mask2 + # cv2.fillPoly(eye, fc_landmark[655:691][np.newaxis, :, :], (255, 255, 255)) + cv2.circle(eye, (left_eye[0], left_eye[1]), 1, (255, 255, 255), line_size2) + # cv2.fillPoly(eye, fc_landmark[756:792][np.newaxis, :, :], (255, 255, 255)) + cv2.circle(eye, (right_eye[0], right_eye[1]), 1, (255, 255, 255), line_size2) + hull_mask = hull_mask * (1 - eye_mask[:, :, 0:1]) + eye * eye_mask[:, :, 0:1] + else: + assert False + else: + hull_mask = np.zeros((h, w), dtype=np.uint8) + if len(fc_landmark) == 137: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (1)) + elif len(fc_landmark) == 1000: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (1)) + else: + assert False + + return hull_mask + + +def draw_makeup_mask(another_pts1k, another_mask, pts1k, hull_mask): + pts1kint = pts1k.astype(np.int32) + assert len(another_pts1k) == 1000 and len(pts1kint) == 1000 + cv2.fillPoly(another_mask, another_pts1k[691:755][np.newaxis, :, :], (4)) # eye + cv2.fillPoly(another_mask, another_pts1k[792:856][np.newaxis, :, :], (5)) # eye + cv2.fillPoly(another_mask, np.concatenate((another_pts1k[312:397], another_pts1k[508:467:-1]))[np.newaxis, :, :], (7)) # mouth + cv2.fillPoly(another_mask, np.concatenate((another_pts1k[467:395:-1], another_pts1k[508:548], [another_pts1k[468], another_pts1k[312]]))[np.newaxis, :, :], (9)) # mouth + cv2.fillPoly(another_mask, another_pts1k[928:1000][np.newaxis, :, :], (0)) # eyebrow + cv2.fillPoly(another_mask, another_pts1k[856:928][np.newaxis, :, :], (0)) # eyebrow + + + cv2.fillPoly(hull_mask, pts1kint[691:755][np.newaxis, :, :], (4)) # eye + cv2.fillPoly(hull_mask, pts1kint[792:856][np.newaxis, :, :], (5)) # eye + cv2.fillPoly(hull_mask, np.concatenate((pts1kint[312:397], pts1kint[508:467:-1]))[np.newaxis, :, :], (7)) # mouth + cv2.fillPoly(hull_mask, np.concatenate((pts1kint[467:395:-1], pts1kint[508:548], [pts1kint[468], pts1kint[312]]))[np.newaxis, :, :], (9)) # mouth + cv2.fillPoly(hull_mask, pts1kint[928:1000][np.newaxis, :, :], (0)) # eyebrow + cv2.fillPoly(hull_mask, pts1kint[856:928][np.newaxis, :, :], (0)) # eyebrow + return another_mask, hull_mask + +def draw_users_hull_mask(fc_landmark, w=256, h=256): + hull_mask = np.zeros((h, w), dtype=np.uint8) + if len(fc_landmark) == 87: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (1)) + elif len(fc_landmark) == 137: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (1)) + elif len(fc_landmark) == 1000: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (1)) + else: + assert False + return hull_mask + +def draw_bigger_hull_mask(fc_landmark, w=256, h=256): + + if len(fc_landmark) == 1000: + fc_landmark = pts_1k_to_137(fc_landmark) + + hull_mask = np.zeros((h, w), dtype=np.uint8) + if len(fc_landmark) == 137: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (255)) + elif len(fc_landmark) == 1000: + cv2.fillConvexPoly(hull_mask, cv2.convexHull(fc_landmark), (255)) + else: + assert False + + kernel_size = int(w / 11.) + if kernel_size % 2 == 0: + kernel_size += 1 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + hull_mask = (cv2.dilate(hull_mask, kernel, iterations=1) / 255.).astype(np.uint8) + cv2.fillPoly(hull_mask, fc_landmark[22:48][np.newaxis, :, :], (10,)) + + + leye_mask = np.zeros((h, w), dtype=np.uint8) + cv2.fillPoly(leye_mask, fc_landmark[88:104][np.newaxis, :, :], (1,)) + cv2.fillPoly(leye_mask, fc_landmark[105:121][np.newaxis, :, :], (1,)) + kernel_size = int(w / 5.) + if kernel_size % 2 == 0: + kernel_size += 1 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + leye_mask = (cv2.dilate(leye_mask, kernel, iterations=1)).astype(np.uint8) + hull_mask[leye_mask > 0] = 10 + + # cv2.imshow('hull_mask', hull_mask * 25 ) + # cv2.waitKey() + + return hull_mask + + + +def get_kernel_size(div_num, output_img_size=256): + kernel_size = int(float(output_img_size) / div_num) + if kernel_size % 2 == 0: + kernel_size += 1 + return kernel_size + +class LKTracking(object): + def __init__(self, termcrit=cv2.TERM_CRITERIA_COUNT | cv2.TERM_CRITERIA_EPS, winSize=31, extSize=41, threshold=5): + self.termcrit = termcrit + self.winSize = winSize + self.extSize = extSize + self.threshold = threshold + + self.prepoints = [] + self.preimg_small_ = None + self.preimg_rect_ = np.array([0, 0, 0, 0]) + + def UpdatePoints(self, new_pts): + if len(new_pts) != len(self.prepoints): + self.prepoints = [] + return False + self.prepoints = new_pts + return True + + def TrackingPoints(self, curimg, curpoints): + need_init = False + + def Prepare(self, curimg, curpoints): + prepoints_int = curpoints.astype(np.int32) + pass + +def get_transform_mat_mouth(landmark, output_size): + # mean_mouth_x_4pts = np.array([0.15, 0.5, 0.75, 0.5]) + # mean_mouth_y_4pts = np.array([0.5, 0.15, 0.5, 0.75]) + mean_mouth_x_4pts = np.array([0.2, 0.5, 0.8, 0.5]) + mean_mouth_y_4pts = np.array([0.5, 0.2, 0.5, 0.8]) + landmarks_2D_mouth_4pts = np.stack([mean_mouth_x_4pts, mean_mouth_y_4pts], axis=1) + + dst_size = output_size + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k) + mat = umeyama(landmark[:22], landmarks_2D_137[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 236: + #mouth crop + landmark_mouth = [] + landmark_mouth.append(landmark[467-312]) + landmark_mouth.append(landmark[432-312]) + landmark_mouth.append(landmark[397-312]) + landmark_mouth.append(landmark[354-312]) + landmark_mouth = np.array(landmark_mouth) + mat = umeyama(landmark_mouth, landmarks_2D_mouth_4pts * dst_size, True)[0:2] + return mat + + +mean_mouth_x_4pts = np.array([0.2, 0.5, 0.8, 0.5]) +mean_mouth_y_4pts = np.array([0.5, 0.2, 0.5, 0.8]) +landmarks_2D_mouth_4pts = np.stack([mean_mouth_x_4pts, mean_mouth_y_4pts], axis=1) + +def get_transform_mat_for_mouth(landmark, output_size): + dst_size = output_size + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + # landmarks_2D_137 = pts_1k_to_137(landmarks_2D_1k) + # mat = umeyama(landmark[:22], landmarks_2D_137[:22] * dst_size, True)[0:2] + landmark_mouth = [] + landmark_mouth.append(landmark[22]) + landmark_mouth.append(landmark[42]) + landmark_mouth.append(landmark[36]) + landmark_mouth.append(landmark[29]) + landmark_mouth = np.array(landmark_mouth) + mat = umeyama(landmark_mouth, landmarks_2D_mouth_4pts * dst_size, True)[0:2] + + return mat + elif len(landmark) == 236: + # mouth crop + landmark_mouth = [] + landmark_mouth.append(landmark[467 - 312]) + landmark_mouth.append(landmark[432 - 312]) + landmark_mouth.append(landmark[397 - 312]) + landmark_mouth.append(landmark[354 - 312]) + landmark_mouth = np.array(landmark_mouth) + mat = umeyama(landmark_mouth, landmarks_2D_mouth_4pts * dst_size, True)[0:2] + # mat = umeyama(landmark[:312], landmarks_2D_1k[:312] * dst_size, True)[0:2] + return mat + + elif len(landmark) == 1000: + # mouth crop + landmark_mouth = [] + landmark_mouth.append(landmark[467]) + landmark_mouth.append(landmark[432]) + landmark_mouth.append(landmark[397]) + landmark_mouth.append(landmark[354]) + landmark_mouth = np.array(landmark_mouth) + mat = umeyama(landmark_mouth, landmarks_2D_mouth_4pts * dst_size, True)[0:2] + # mat = umeyama(landmark[:312], landmarks_2D_1k[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_full_face_ratio_skin(landmark, output_size, ratio=1.0, skin=0): + dst_size = output_size + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + # landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + 0.5 + # landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.5) * ratio + 0.4 + + if skin == 1: + ###################### 0.3 face rata 592 + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + 0.5 + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.6) * ratio + 0.3 + ###################### 0.3 face rata 592 + else: + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.4) * ratio + 0.45 + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.4) * ratio + 0.4 + + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_1k_tmp = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_1k_tmp[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_for_uv(landmark, output_size): + dst_size = output_size + if len(landmark) == 1000: + landmark = pts_1k_to_137(landmark) + + if len(landmark) == 137: + eye_dis = 0.35 + mouth_dis = 0.42 + g_Average_5point_180 = np.array([ + eye_dis, 0.24, + 1 - eye_dis, 0.24, + 0.5, 0.42, + mouth_dis, 0.55, + 1 - mouth_dis, 0.55 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[96] + landmark[88]) / 2 + right_eye = (landmark[105] + landmark[113]) / 2 + nose = landmark[83] + left_mouth = (landmark[22] + landmark[48]) / 2 + right_mouth = (landmark[56] + landmark[36]) / 2 + + pts5_dst = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_src = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + else: + assert False + + +def draw_crop_eye_bysize(img, pts1k, img_size, change_eyebrow): + pts137tmp = pts_1k_to_137(pts1k).astype(np.int32) + eye_mask = np.zeros((img_size, img_size, 3), dtype=np.uint8) + cv2.fillPoly(eye_mask, pts137tmp[88:104][np.newaxis, :, :], (1, 1, 1)) + cv2.fillPoly(eye_mask, pts137tmp[105:121][np.newaxis, :, :], (1, 1, 1)) + kernel_size = int(img_size / 198.) # 31 + if kernel_size % 2 == 0: + kernel_size += 1 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + eye_mask = (cv2.dilate(eye_mask, kernel, iterations=1)).astype(np.uint8) + + kernel_size = int(img_size / 7) # 8.2 + if kernel_size % 2 == 0: + kernel_size += 1 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size*2, kernel_size)) + eye_mask2 = (cv2.dilate(eye_mask, kernel, iterations=1)).astype(np.uint8) + eye_mask2 = eye_mask2 - eye_mask + + if change_eyebrow: + # eyebrow + eyebrow_mask = np.zeros((img_size, img_size, 3), dtype=np.uint8) + + x, y, w, h = cv2.boundingRect(pts137tmp[129:137]) + # cv2.rectangle(eyebrow_mask, (x, y), (x + w, y + h), (1, 1, 1), -1) + eyebrow_mask[y:y + h, x:x + w, :] = [1, 1, 1] + + x, y, w, h = cv2.boundingRect(pts137tmp[121:129]) + # cv2.rectangle(eyebrow_mask, (x, y), (x + w, y + h), (1, 1, 1), -1) + eyebrow_mask[y:y + h, x:x + w, :] = [1, 1, 1] + + kernel_size = int(592 / 21.) # 21 + if kernel_size % 2 == 0: + kernel_size += 1 + kernel_eyebrow = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (int(kernel_size * 5), int(kernel_size))) + eyebrow_mask = (cv2.dilate(eyebrow_mask, kernel_eyebrow, iterations=1)).astype(np.uint8) + eye_mask2 = (eye_mask2.astype(bool) | eyebrow_mask.astype(bool)).astype(np.float32) + else: + eyebrow_mask = np.zeros((img_size, img_size, 3), dtype=np.uint8) + cv2.fillPoly(eyebrow_mask, pts137tmp[129:137][np.newaxis, :, :], (1, 1, 1)) + cv2.fillPoly(eyebrow_mask, pts137tmp[121:129][np.newaxis, :, :], (1, 1, 1)) + eyebrow_mask = 1 - eyebrow_mask + eye_mask2 = (eye_mask2.astype(bool) & eyebrow_mask.astype(bool)).astype(np.float32) + + # nose + cv2.fillPoly(eye_mask2, pts137tmp[64:79][np.newaxis, :, :], (0, 0, 0)) + eye_mask2[(pts137tmp[133, 1] - img_size // 10):pts137tmp[65, 1], pts137tmp[64, 0]:pts137tmp[78, 0]] = 0 + + img[eye_mask2 > 0] = 0 + + cv2.circle(img, (pts137tmp[133, 0], pts137tmp[133, 1]), img_size // 50, (255, 0, 0), -1) + cv2.circle(img, (pts137tmp[121, 0], pts137tmp[121, 1]), img_size // 50, (255, 0, 0), -1) + + return img + + +# def draw_crop_eye_bysize_using_seg_mask(img, pts1k, mask, img_size, change_eyebrow): +# +# pts137tmp = pts_1k_to_137(pts1k).astype(np.int32) +# c0, c1, c2 = mask[:, :, 0], mask[:, :, 1], mask[:, :, 2] +# c0[c0 >= 100] = 255 +# c0[c0 < 100] = 0 +# c1[c1 >= 100] = 255 +# c1[c1 < 100] = 0 +# c2[c2 > 0] = 0 +# mask[:, :, 0] = c0 +# mask[:, :, 1] = c1 +# mask[:, :, 2] = c2 +# eye_index = (mask == [255, 0, 0]).all(axis=2) +# eyelids_index = (mask == [0, 255, 0]).all(axis=2) +# black_im = np.zeros((img_size, img_size, 3), dtype=np.uint8) +# black_im2 = np.zeros((img_size, img_size, 3), dtype=np.uint8) +# +# black_im[eye_index] = [255, 255, 255] +# black_im2[eye_index] = [255, 255, 255] +# kernel_size = int(img_size // 7) # 8.2 +# +# if kernel_size % 2 == 0: +# kernel_size += 1 +# kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size * 2, kernel_size)) +# black_im2 = cv2.dilate(black_im2, kernel).astype(np.uint8) +# +# endless_belt_im = black_im2 - black_im +# endless_belt_im = np.ones_like(endless_belt_im) * 255 - endless_belt_im +# +# if change_eyebrow: +# # eyebrow +# eyebrow_mask = np.zeros((img_size, img_size, 3), dtype=np.uint8) +# x, y, w, h = cv2.boundingRect(pts137tmp[129:137]) +# # cv2.rectangle(eyebrow_mask, (x, y), (x + w, y + h), (1, 1, 1), -1) +# eyebrow_mask[y:y + h, x:x + w, :] = [1, 1, 1] +# +# x, y, w, h = cv2.boundingRect(pts137tmp[121:129]) +# # cv2.rectangle(eyebrow_mask, (x, y), (x + w, y + h), (1, 1, 1), -1) +# eyebrow_mask[y:y + h, x:x + w, :] = [1, 1, 1] +# +# kernel_size = int(592 / 21.) # 21 +# if kernel_size % 2 == 0: +# kernel_size += 1 +# kernel_eyebrow = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (int(kernel_size * 5), int(kernel_size))) +# eyebrow_mask = (cv2.dilate(eyebrow_mask, kernel_eyebrow, iterations=1)).astype(np.uint8) +# eyebrow_mask = 1 - eyebrow_mask +# endless_belt_im = (endless_belt_im.astype(bool) & eyebrow_mask.astype(bool)).astype(np.float32) * 255 +# else: +# eyebrow_mask = np.zeros((img_size, img_size, 3), dtype=np.uint8) +# cv2.fillPoly(eyebrow_mask, pts137tmp[129:137][np.newaxis, :, :], (1, 1, 1)) +# cv2.fillPoly(eyebrow_mask, pts137tmp[121:129][np.newaxis, :, :], (1, 1, 1)) +# endless_belt_im = (endless_belt_im.astype(bool) | eyebrow_mask.astype(bool)).astype(np.float32) * 255 +# +# cv2.fillPoly(endless_belt_im, pts137tmp[64:79][np.newaxis, :, :], (255, 255, 255)) +# endless_belt_im[(pts137tmp[133, 1] - img_size // 10):pts137tmp[65, 1], pts137tmp[64, 0]:pts137tmp[78, 0], :] = 255 +# +# +# res = np.uint8(endless_belt_im / 255) +# res = np.uint8(res * img) +# +# # add eyelids semantic map +# eyelids_mask = np.zeros((img_size, img_size, 3)) +# eyelids_mask[eyelids_index] = [0, 255, 0] +# res = res + np.uint8(eyelids_mask) +# +# cv2.circle(res, (pts137tmp[133, 0], pts137tmp[133, 1]), img_size // 50, (255, 0, 0), -1) +# cv2.circle(res, (pts137tmp[121, 0], pts137tmp[121, 1]), img_size // 50, (255, 0, 0), -1) +# +# return res + + +def draw_crop_eye_bysize_using_seg_mask(img, pts1k, mask, img_size, change_eyebrow): + start = time.time() + mask = mask.copy() + pts137tmp = pts_1k_to_137(pts1k).astype(np.int32) + c0, c1, c2 = mask[:, :, 0], mask[:, :, 1], mask[:, :, 2] + c0[c0 >= 100] = 255 + c0[c0 < 100] = 0 + + c1[c1 >= 100] = 255 + c1[c1 < 100] = 0 + + # c2 也需要截断 + c2[c2 >= 100] = 255 + c2[c2 < 100] = 0 + + mask[:, :, 0] = c0 + mask[:, :, 1] = c1 + mask[:, :, 2] = c2 + eye_index = (mask == [255, 0, 0]).all(axis=2) | (mask == [0, 0, 255]).all(axis=2) + eyelids_index = (mask == [0, 255, 0]).all(axis=2) + black_im = np.zeros((img_size, img_size, 3), dtype=np.uint8) + black_im2 = np.zeros((img_size, img_size, 3), dtype=np.uint8) + + black_im[eye_index] = [255, 255, 255] + black_im2[eye_index] = [255, 255, 255] + + # cv2.imshow("black_im: ", black_im) + # cv2.waitKey() + + kernel_size = int(img_size // 7) # 8.2 + + if kernel_size % 2 == 0: + kernel_size += 1 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size * 2, kernel_size)) + black_im2 = cv2.dilate(black_im2, kernel).astype(np.uint8) + + endless_belt_im = black_im2 - black_im + endless_belt_im = np.ones_like(endless_belt_im) * 255 - endless_belt_im + + if change_eyebrow: + # eyebrow + eyebrow_mask = np.zeros((img_size, img_size, 3), dtype=np.uint8) + x, y, w, h = cv2.boundingRect(pts137tmp[129:137]) + cv2.rectangle(eyebrow_mask, (x, y), (x + w, y + h), (1, 1, 1), -1) + x, y, w, h = cv2.boundingRect(pts137tmp[121:129]) + cv2.rectangle(eyebrow_mask, (x, y), (x + w, y + h), (1, 1, 1), -1) + kernel_size = int(img_size / 21.) # 21 + if kernel_size % 2 == 0: + kernel_size += 1 + kernel_eyebrow = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (int(kernel_size * 5), int(kernel_size))) + eyebrow_mask = (cv2.dilate(eyebrow_mask, kernel_eyebrow, iterations=1)).astype(np.uint8) + eyebrow_mask = 1 - eyebrow_mask + endless_belt_im = (endless_belt_im.astype(bool) & eyebrow_mask.astype(bool)).astype(np.float32) * 255 + else: + eyebrow_mask = np.zeros((img_size, img_size, 3), dtype=np.uint8) + cv2.fillPoly(eyebrow_mask, pts137tmp[129:137][np.newaxis, :, :], (1, 1, 1)) + cv2.fillPoly(eyebrow_mask, pts137tmp[121:129][np.newaxis, :, :], (1, 1, 1)) + endless_belt_im = (endless_belt_im.astype(bool) | eyebrow_mask.astype(bool)).astype(np.float32) * 255 + + # nose + cv2.fillPoly(endless_belt_im, pts137tmp[64:79][np.newaxis, :, :], (255, 255, 255)) + endless_belt_im[(pts137tmp[133, 1] - img_size // 10):pts137tmp[65, 1], pts137tmp[64, 0]:pts137tmp[78, 0], :] = 255 + + res = np.uint8(endless_belt_im / 255) + res = np.uint8(res * img) + + # add eyelids semantic map + eyelids_mask = np.zeros((img_size, img_size, 3)) + eyelids_mask[eyelids_index] = [0, 255, 0] + res = res + np.uint8(eyelids_mask) + + cv2.circle(res, (pts137tmp[133, 0], pts137tmp[133, 1]), img_size // 50, (255, 0, 0), -1) + cv2.circle(res, (pts137tmp[121, 0], pts137tmp[121, 1]), img_size // 50, (255, 0, 0), -1) + + return res + +def get_transform_singleeye(landmark, output_size, forlabel=False): + dst_size = output_size + left_dis = 0.28 + g_Average_5point_180 = np.array([ + left_dis, 0.5, + 0.5, 0.499, + 1 - left_dis, 0.5, + ]) + pts3_dst = g_Average_5point_180.reshape((3, -1)) * dst_size + pts3_src = np.vstack((landmark[0], landmark[1], landmark[2])) + image_to_face_mat = umeyama(pts3_src, pts3_dst, True)[:2] + + return image_to_face_mat + +def get_transform_mat_full_face_ratio_stylegan(landmark, output_size, ratio=1.0): + dst_size = output_size + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + + #################### + # # ###################### train stylegan hair rate_0.4 size_512 + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + 0.5 + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.6) * ratio + 0.4 + + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_1k_tmp = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_1k_tmp[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_hair(landmark, output_size, ratio=0.5, w_ratio=0.5, h_ratio=0.40): + dst_size = output_size + + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + w_ratio + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.5) * ratio + h_ratio + + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_1k_tmp = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_1k_tmp[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat +def get_transform_mat_sex(landmark, output_size, forlabel=False): + dst_size = output_size + if len(landmark) == 87: + eye_dis = 0.34 + mouth_dis = 0.34 + g_Average_5point_180 = np.array([ + eye_dis, 0.3, + 1 - eye_dis, 0.3, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[31] + landmark[35]) / 2 + right_eye = (landmark[39] + landmark[43]) / 2 + nose = landmark[62] + left_mouth = (landmark[66] + landmark[79]) / 2 + right_mouth = (landmark[72] + landmark[83]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat + elif len(landmark) == 137: + + eye_dis = 0.317 + mouth_dis = 0.345 + g_Average_5point_180 = np.array([ + eye_dis, 0.4, + 1 - eye_dis, 0.4, + 0.5, 0.6, + mouth_dis, 0.63, + 1 - mouth_dis, 0.63 + ]) + # print(g_Average_5point_180) + left_eye = (landmark[96] + landmark[88]) / 2 + right_eye = (landmark[105] + landmark[113]) / 2 + nose = landmark[83] + left_mouth = (landmark[22] + landmark[48]) / 2 + right_mouth = (landmark[56] + landmark[36]) / 2 + + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + + mat = umeyama(pts5_src, pts5_dst, True)[0:2] + return mat +def get_transform_mat_hair_ratio(landmark, output_size, ratio=1.0): + dst_size = output_size + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + 0.5 + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.5) * ratio + 0.45 + + # landmarks1k_contours = landmarks_2D_1k_tmp[:312] + # x_l, y_l = np.min(landmarks1k_contours, axis=0) + # x_h, y_h = np.max(landmarks1k_contours, axis=0) + # + # print("x_l: ", x_l, "x_h: ", x_h) + # print("y_l: ", y_l, "y_h: ", y_h) + # + # def draw_landmark(landmark_ori, img): + # landmark_full_int = (landmark_ori.copy() * 512).astype(np.int32) + # img_show = img.copy() + # for pt in landmark_full_int: + # cv2.circle(img_show, (pt[0], pt[1]), 2, (0, 0, 255), 1) + # + # return img_show.astype(np.uint8) + # + # img_temp = np.zeros((512, 512, 3), dtype=np.uint8) + # img_show = draw_landmark(landmarks_2D_1k_tmp, img_temp) + # cv2.imshow("img_show", img_show) + # cv2.waitKey() + + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_1k_tmp = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_1k_tmp[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_hair_ratio_v1(landmark, output_size, ratio=1.0, h_offset=0.5): + dst_size = output_size + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + 0.5 + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.5) * ratio + h_offset + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_1k_tmp = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_1k_tmp[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_full_face_ratio_deeplab(landmark, output_size, ratio=0.3, w_ratio=0.5, h_ratio=0.45): + dst_size = output_size + + landmarks_2D_1k_tmp = landmarks_2D_1k.copy() + landmarks_2D_1k_tmp[:, 0] = (landmarks_2D_1k_tmp[:, 0] - 0.5) * ratio + w_ratio + landmarks_2D_1k_tmp[:, 1] = (landmarks_2D_1k_tmp[:, 1] - 0.5) * ratio + h_ratio + + if len(landmark) == 87: + assert False + elif len(landmark) == 137: + landmarks_2D_1k_tmp = pts_1k_to_137(landmarks_2D_1k_tmp) + mat = umeyama(landmark[:22], landmarks_2D_1k_tmp[:22] * dst_size, True)[0:2] + return mat + elif len(landmark) == 1000: + mat = umeyama(landmark[:312], landmarks_2D_1k_tmp[:312] * dst_size, True)[0:2] + return mat + +def get_transform_mat_face_restore(landmark, output_size): + dst_size = output_size + if len(landmark) == 1000: + eye_dis = 0.4 + mouth_dis = 0.42 + g_Average_5point_180 = np.array([ + eye_dis, 0.47, + 1 - eye_dis, 0.47, + 0.5, 0.6, + mouth_dis, 0.71, + 1 - mouth_dis, 0.71 + ]) + pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size + left_eye = (landmark[691] + landmark[723]) / 2 + right_eye = (landmark[792] + landmark[824]) / 2 + nose = landmark[621] + left_mouth = (landmark[467] + landmark[468]) / 2 + right_mouth = (landmark[396] + landmark[508]) / 2 + pts5_src = np.vstack((left_eye, right_eye, + nose, + left_mouth, right_mouth)) + + image_to_face_mat = umeyama(pts5_src, pts5_dst, True)[:2] + + return image_to_face_mat diff --git a/photo_service/webd/changelog.txt b/photo_service/webd/changelog.txt new file mode 100644 index 0000000..9d22da2 --- /dev/null +++ b/photo_service/webd/changelog.txt @@ -0,0 +1,39 @@ +2024/02/23 +Fix popup 127.0.0.1 instead of real ip in some situations. + +2023/06/04 +Rename packages name with toolchains name. +Fix some toolchain's misconfiguration, which lead to unusable binary. +Add support for drag-drop upload and folder upload. + +2022/05/12 +Fix a issue when rename a file's name with question mark(?). + +2022/03/27 +Fix a potential security problem. +Add support for platform armv6. + +2022/01/27 +Fix unexpected quit when upload file to jffs2 filesystem. + +2022/01/24 +Fix unexpected quit on windows. +Add back web player, but not enabled by default. + +2022/01/10 +Fix UI mess up when multiple files selected to upload at once. + +2022/01/05 +Fix upload error when file > 4G, 64bits versions not affected. Thanks for Mover. + +2021/12/02 +Some minor UI correction. +Add installation script for Android. + +2021/11/30 +Fix for __libc_start_main@@GLIBC_2.34 +Fix parameter "-c". + +2021/11/29 +Try fix rename() error on Android. +Fix guest's permission settings. diff --git a/photo_service/webd/web/.player.htm b/photo_service/webd/web/.player.htm new file mode 100644 index 0000000..2cee5f8 --- /dev/null +++ b/photo_service/webd/web/.player.htm @@ -0,0 +1 @@ +
\ No newline at end of file diff --git a/photo_service/webd/webd b/photo_service/webd/webd new file mode 100755 index 0000000..c16424b Binary files /dev/null and b/photo_service/webd/webd differ diff --git a/photo_service/webd/webd.conf b/photo_service/webd/webd.conf new file mode 100644 index 0000000..447fdfe --- /dev/null +++ b/photo_service/webd/webd.conf @@ -0,0 +1,74 @@ +# NOTE: +# This file must be encoded in UTF-8. +# Directives and variable definition in this file are case-insensitive. +# Lines that begin with the hash character "#" are considered comments, and are ignored. + + +# Webd.Root: The directory that webd share on network. + +# Example for Linux: +# Webd.Root /mnt/sdb1 +# Example for Windows: +# Webd.Root "D:\my share" + + +# Webd.Listen: Bind webd to specific IP and/or port. +# Also, webd can bind to multiple addresses by use multiple "Webd.Listen" instructions. + +# Bind to port 9212 with IPv4: +# Webd.Listen 9212 +# Bind to port 9212 with both IPv4 and IPv6: +# Webd.Listen [::]:9212 + + +# User's permissions tag, can be set via one or more tag combinations: +# r: Access files. +# l: List directories. +# u: Upload file. +# m: Delete, move, or rename files. +# S: Show hidden files or directories. +# T: Use webpage to play media files. +# D: Add 'download' atrribute to file link. + +# For now, webd supports only two users. +# Each user may have it's own Username Password and Permissions. +# But they share the same web directory. + +# user1 has all permissions. +# Webd.User rlumS user1 pass1 + +# user2 can download and list files. +# Webd.User rl user2 pass2 + +# Guest can download and list files by default. +# Uncomment to disable all permissions for guest. +# Webd.Guest 0 + + +# Hide tray icon for Windows. +# Webd.Hide + + +# Specify the path of Browser for Windows if webd can not popup Browser by double clicking tray icon. +# Webd.Browser "C:\Program Files\Mozilla Firefox\firefox.exe" +# Or start Browser with extra paramters that set by a batch file. +# Webd.Browser "C:\Program Files\Mozilla Firefox\myFirefox.cmd" + + +# Envionment variables for webd. +# These should be set in the command line or system configration. +# +# Write log files to /var/log/webd-YYYY-MM-DD.log +# _LOG_DIR=/var/log/webd- +# +# Write log to syslog. +# _syslog=1 +# +# Set the maximum number of open file descriptors, linux only. +# _FD_LIMIT=10240 +# +# Switch to non-privileged user after startup, linux only. +# _RUNAS=nobody +# +# chroot after startup, linux only. +# _CHROOT_PATH=/mnt/sda1 diff --git a/photo_service/webd/webd_upload.sh b/photo_service/webd/webd_upload.sh new file mode 100755 index 0000000..8ae2a17 --- /dev/null +++ b/photo_service/webd/webd_upload.sh @@ -0,0 +1,33 @@ +#!/bin/sh +# 本脚本用于向webd的服务器上传文件，需要提供两个参数： 1、上传的文件路径。 2、上传的地址 +# bash ./webd_upload.sh abc.png http://nas.zhourunnan.cn:9212/zhourunnan + +b36enc() { + local b36=$(echo {0..9} {a..z}); b36=${b36// /} + awk \ + 'BEGIN { b = split(ARGV[1], D, ""); n = ARGV[2]; do { d = int(n / b); i = D[n - b * d + 1]; r = i r; n = d } while(n != 0); print r}' \ + "${b36}" "$1" +} + +upload(){ + local def_url="http://aaanas.zhourunnan.cn:9212/" + local def_cookie='AcMY8R290zo' + + local file="${1}" + local size=$(stat -c %s "${file}") + local time=$(stat -c %Y "${file}") + + local url="${2:-${def_url}}"; url="${url%/}/${file##*/}" + local cookie="${3:-${def_cookie}}" + + echo -e "\e[1;32mupload: ${url}\e[0m" + + wget -vdt1 -O- "${url//#/%23}?N$(b36enc ${time})" \ + --header='Content-Type: application/octet-stream' \ + --header="Cookie: u=${cookie}" \ + --header="RaOff: bytes=0/$(b36enc ${size})" \ + --post-file="${file}" + + echo -e "\n\e[1;32mfinish: ${url}\e[0m" + +}; upload "$@" diff --git a/photo_service/webui_im2im.py b/photo_service/webui_im2im.py new file mode 100755 index 0000000..3e67e92 --- /dev/null +++ b/photo_service/webui_im2im.py @@ -0,0 +1,106 @@ +import io +import cv2 +import base64 +import requests +from PIL import Image +import numpy as np + +""" + To use this example make sure you've done the following steps before executing: + 1. Ensure automatic1111 is running in api mode with the controlnet extension. + Use the following command in your terminal to activate: + ./webui.sh --no-half --api + 2. Validate python environment meet package dependencies. + If running in a local repo you'll likely need to pip install cv2, requests and PIL +""" + + +class ControlnetRequestImg2Img: + def __init__(self, prompt, net_prompt): + self.url = "http://127.0.0.1:7860/sdapi/v1/img2img" + self.prompt = prompt + self.neg_prompt = net_prompt + self.body = None + + def build_body(self, dst_width, dst_height, cfg_scale, base_img): + + self.body = { + "prompt": self.prompt, + "negative_prompt": self.neg_prompt, + "sampler_name": "Restart", + "batch_size": 1, + "steps": 30, + "width": dst_width, + "height": dst_height, + "cfg_scale": cfg_scale, + "seed": -1, + "init_images": [ + self.encode_image_to_base64(base_img) + ], + "denoising_strength": 0.4, + "alwayson_scripts": { + "controlnet": { + "args": [ + { + "enabled": True, + "module": "openpose_full", + "model": "openpose", + "weight": 1.0, + # "image": self.read_image(), + "resize_mode": "Crop and Resize", + "low_vram": False, + "processor_res": 512, + "guidance_start": 0.0, + "guidance_end": 1.0, + "control_mode": "Balanced", + "pixel_perfect": True + } + ] + } + } + } + + def send_request(self): + response = requests.post(url=self.url, json=self.body) + return response.json() + + def encode_image_to_base64(self, img): + retval, bytes = cv2.imencode('.png', img) + encoded_image = base64.b64encode(bytes).decode('utf-8') + return encoded_image + + def read_image(self): + img = cv2.imread(self.img_path) + retval, bytes = cv2.imencode('.png', img) + encoded_image = base64.b64encode(bytes).decode('utf-8') + return encoded_image + + def read_mask(self): + img = cv2.imread(self.mask) + retval, bytes = cv2.imencode('.png', img) + encoded_image = base64.b64encode(bytes).decode('utf-8') + return encoded_image + +def encode_image_to_base64(img): + retval, bytes = cv2.imencode('.jpg', img) + encoded_image = base64.b64encode(bytes).decode('utf-8') + return encoded_image + + + +if __name__ == '__main__': + path = '/home/chinatszrn/Downloads/photo_service/service_data/template_data/template01.png' + img = cv2.imread(path) + prompt = ',easyphoto_face, easyphoto, 1person,face,suit' + neg_prompt = '(worst quality:2),(low quality:2),(normal quality:2),lowres,watermark' + + + control_net = ControlnetRequestImg2Img(prompt, neg_prompt) + control_net.build_body(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=3.5, base_img=img) + output = control_net.send_request() + result = output['images'][0] + + image_array = np.frombuffer(base64.b64decode(result.split(",", 1)[0]), np.uint8) + image = cv2.imdecode(image_array, cv2.IMREAD_COLOR) + cv2.imshow('image', image) + cv2.waitKey() diff --git a/photo_service/webui_im2im2.py b/photo_service/webui_im2im2.py new file mode 100755 index 0000000..2bf0c41 --- /dev/null +++ b/photo_service/webui_im2im2.py @@ -0,0 +1,86 @@ +import io +import cv2 +import base64 +import requests +from PIL import Image + +class ControlnetRequestImg2Img: + def __init__(self, prompt, net_prompt, path, mask): + self.url = "http://127.0.0.1:57860/sdapi/v1/img2img" + self.prompt = prompt + self.neg_prompt = net_prompt + self.img_path = path + self.mask = mask + self.body = None + + def build_body(self): + img = cv2.imread(self.img_path) + self.body = { + "prompt": self.prompt, + "negative_prompt": self.neg_prompt, + "sampler_name": "DPM++ 2M Karras", + "batch_size": 1, + "steps": 30, + "width": img.shape[1], + "height": img.shape[0], + "cfg_scale": 7, + "seed": -1, + "mask_blur": 15, + "init_images": [ + self.read_image() + ], + "inpaint_full_res": True, + "inpainting_fill": 1, + "inpainting_mask_invert": 1, + "mask": self.read_mask(), + "denoising_strength": 0.4, + "alwayson_scripts": { + "controlnet": { + "args": [ + { + "enabled": True, + "module": "openpose_full", + "model": "openpose", + "weight": 1.0, + "resize_mode": 1, + "lowvram": False, + "processor_res": 512, + "guidance_start": 0.0, + "guidance_end": 1.0, + "control_mode": 0, + "pixel_perfect": True + }, + ] + }, + } + } + + def send_request(self): + response = requests.post(url=self.url, json=self.body) + return response.json() + + def read_image(self): + img = cv2.imread(self.img_path) + retval, bytes = cv2.imencode('.png', img) + encoded_image = base64.b64encode(bytes).decode('utf-8') + return encoded_image + + def read_mask(self): + img = cv2.imread(self.mask) + retval, bytes = cv2.imencode('.png', img) + encoded_image = base64.b64encode(bytes).decode('utf-8') + return encoded_image + + +if __name__ == '__main__': + path = '/home/chinatszrn/Downloads/user1_hr.png' + mask_path = '/home/chinatszrn/Downloads/user1_hr_mask.png' + prompt = 'a woman with long blonde hair and a blue shirt on a gray background with a gray background and a gray background, lyco art, An Gyeon, realistic face, a character portrait' + neg_prompt = '(nsfw:1.5), ng_deepnegative_v1_75t, (badhandv4:1.2), (worst quality:2), (low quality:2), (normal quality:2), lowres, bad anatomy, bad hands, ((monochrome)), ((grayscale)) watermark, moles, large breast, big breast, bad_pictures,easynegative' + + control_net = ControlnetRequestImg2Img(prompt, neg_prompt, path, mask_path) + control_net.build_body() + output = control_net.send_request() + result = output['images'][0] + image = Image.open(io.BytesIO(base64.b64decode(result.split(",", 1)[0]))) + image.save('save2.png') diff --git a/setup.sh b/setup.sh new file mode 100755 index 0000000..ec49ace --- /dev/null +++ b/setup.sh @@ -0,0 +1,80 @@ +#!/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" diff --git a/stable-diffusion-webui/.eslintignore b/stable-diffusion-webui/.eslintignore new file mode 100755 index 0000000..1cfd948 --- /dev/null +++ b/stable-diffusion-webui/.eslintignore @@ -0,0 +1,4 @@ +extensions +extensions-disabled +repositories +venv \ No newline at end of file diff --git a/stable-diffusion-webui/.eslintrc.js b/stable-diffusion-webui/.eslintrc.js new file mode 100755 index 0000000..2e7258f --- /dev/null +++ b/stable-diffusion-webui/.eslintrc.js @@ -0,0 +1,98 @@ +/* global module */ +module.exports = { + env: { + browser: true, + es2021: true, + }, + extends: "eslint:recommended", + parserOptions: { + ecmaVersion: "latest", + }, + rules: { + "arrow-spacing": "error", + "block-spacing": "error", + "brace-style": "error", + "comma-dangle": ["error", "only-multiline"], + "comma-spacing": "error", + "comma-style": ["error", "last"], + "curly": ["error", "multi-line", "consistent"], + "eol-last": "error", + "func-call-spacing": "error", + "function-call-argument-newline": ["error", "consistent"], + "function-paren-newline": ["error", "consistent"], + "indent": ["error", 4], + "key-spacing": "error", + "keyword-spacing": "error", + "linebreak-style": ["error", "unix"], + "no-extra-semi": "error", + "no-mixed-spaces-and-tabs": "error", + "no-multi-spaces": "error", + "no-redeclare": ["error", {builtinGlobals: false}], + "no-trailing-spaces": "error", + "no-unused-vars": "off", + "no-whitespace-before-property": "error", + "object-curly-newline": ["error", {consistent: true, multiline: true}], + "object-curly-spacing": ["error", "never"], + "operator-linebreak": ["error", "after"], + "quote-props": ["error", "consistent-as-needed"], + "semi": ["error", "always"], + "semi-spacing": "error", + "semi-style": ["error", "last"], + "space-before-blocks": "error", + "space-before-function-paren": ["error", "never"], + "space-in-parens": ["error", "never"], + "space-infix-ops": "error", + "space-unary-ops": "error", + "switch-colon-spacing": "error", + "template-curly-spacing": ["error", "never"], + "unicode-bom": "error", + }, + globals: { + //script.js + gradioApp: "readonly", + executeCallbacks: "readonly", + onAfterUiUpdate: "readonly", + onOptionsChanged: "readonly", + onUiLoaded: "readonly", + onUiUpdate: "readonly", + uiCurrentTab: "writable", + uiElementInSight: "readonly", + uiElementIsVisible: "readonly", + //ui.js + opts: "writable", + all_gallery_buttons: "readonly", + selected_gallery_button: "readonly", + selected_gallery_index: "readonly", + switch_to_txt2img: "readonly", + switch_to_img2img_tab: "readonly", + switch_to_img2img: "readonly", + switch_to_sketch: "readonly", + switch_to_inpaint: "readonly", + switch_to_inpaint_sketch: "readonly", + switch_to_extras: "readonly", + get_tab_index: "readonly", + create_submit_args: "readonly", + restart_reload: "readonly", + updateInput: "readonly", + onEdit: "readonly", + //extraNetworks.js + requestGet: "readonly", + popup: "readonly", + // profilerVisualization.js + createVisualizationTable: "readonly", + // from python + localization: "readonly", + // progrssbar.js + randomId: "readonly", + requestProgress: "readonly", + // imageviewer.js + modalPrevImage: "readonly", + modalNextImage: "readonly", + // localStorage.js + localSet: "readonly", + localGet: "readonly", + localRemove: "readonly", + // resizeHandle.js + setupResizeHandle: "writable" + } +}; diff --git a/stable-diffusion-webui/.git-blame-ignore-revs b/stable-diffusion-webui/.git-blame-ignore-revs new file mode 100755 index 0000000..4104da6 --- /dev/null +++ b/stable-diffusion-webui/.git-blame-ignore-revs @@ -0,0 +1,2 @@ +# Apply ESlint +9c54b78d9dde5601e916f308d9a9d6953ec39430 \ No newline at end of file diff --git a/stable-diffusion-webui/.github/ISSUE_TEMPLATE/bug_report.yml b/stable-diffusion-webui/.github/ISSUE_TEMPLATE/bug_report.yml new file mode 100755 index 0000000..c86bd8a --- /dev/null +++ b/stable-diffusion-webui/.github/ISSUE_TEMPLATE/bug_report.yml @@ -0,0 +1,105 @@ +name: Bug Report +description: You think something is broken in the UI +title: "[Bug]: " +labels: ["bug-report"] + +body: + - type: markdown + attributes: + value: | + > The title of the bug report should be short and descriptive. + > Use relevant keywords for searchability. + > Do not leave it blank, but also do not put an entire error log in it. + - type: checkboxes + attributes: + label: Checklist + description: | + Please perform basic debugging to see if extensions or configuration is the cause of the issue. + Basic debug procedure + 　1. Disable all third-party extensions - check if extension is the cause + 　2. Update extensions and webui - sometimes things just need to be updated + 　3. Backup and remove your config.json and ui-config.json - check if the issue is caused by bad configuration + 　4. Delete venv with third-party extensions disabled - sometimes extensions might cause wrong libraries to be installed + 　5. Try a fresh installation webui in a different directory - see if a clean installation solves the issue + Before making a issue report please, check that the issue hasn't been reported recently. + options: + - label: The issue exists after disabling all extensions + - label: The issue exists on a clean installation of webui + - label: The issue is caused by an extension, but I believe it is caused by a bug in the webui + - label: The issue exists in the current version of the webui + - label: The issue has not been reported before recently + - label: The issue has been reported before but has not been fixed yet + - type: markdown + attributes: + value: | + > Please fill this form with as much information as possible. Don't forget to "Upload Sysinfo" and "What browsers" and provide screenshots if possible + - type: textarea + id: what-did + attributes: + label: What happened? + description: Tell us what happened in a very clear and simple way + placeholder: | + txt2img is not working as intended. + validations: + required: true + - type: textarea + id: steps + attributes: + label: Steps to reproduce the problem + description: Please provide us with precise step by step instructions on how to reproduce the bug + placeholder: | + 1. Go to ... + 2. Press ... + 3. ... + validations: + required: true + - type: textarea + id: what-should + attributes: + label: What should have happened? + description: Tell us what you think the normal behavior should be + placeholder: | + WebUI should ... + validations: + required: true + - type: dropdown + id: browsers + attributes: + label: What browsers do you use to access the UI ? + multiple: true + options: + - Mozilla Firefox + - Google Chrome + - Brave + - Apple Safari + - Microsoft Edge + - Android + - iOS + - Other + - type: textarea + id: sysinfo + attributes: + label: Sysinfo + description: System info file, generated by WebUI. You can generate it in settings, on the Sysinfo page. Drag the file into the field to upload it. If you submit your report without including the sysinfo file, the report will be closed. If needed, review the report to make sure it includes no personal information you don't want to share. If you can't start WebUI, you can use --dump-sysinfo commandline argument to generate the file. + placeholder: | + 1. Go to WebUI Settings -> Sysinfo -> Download system info. + If WebUI fails to launch, use --dump-sysinfo commandline argument to generate the file + 2. Upload the Sysinfo as a attached file, Do NOT paste it in as plain text. + validations: + required: true + - type: textarea + id: logs + attributes: + label: Console logs + description: Please provide **full** cmd/terminal logs from the moment you started UI to the end of it, after the bug occurred. If it's very long, provide a link to pastebin or similar service. + render: Shell + validations: + required: true + - type: textarea + id: misc + attributes: + label: Additional information + description: | + Please provide us with any relevant additional info or context. + Examples: + 　I have updated my GPU driver recently. diff --git a/stable-diffusion-webui/.github/ISSUE_TEMPLATE/config.yml b/stable-diffusion-webui/.github/ISSUE_TEMPLATE/config.yml new file mode 100755 index 0000000..f58c94a --- /dev/null +++ b/stable-diffusion-webui/.github/ISSUE_TEMPLATE/config.yml @@ -0,0 +1,5 @@ +blank_issues_enabled: false +contact_links: + - name: WebUI Community Support + url: https://github.com/AUTOMATIC1111/stable-diffusion-webui/discussions + about: Please ask and answer questions here. diff --git a/stable-diffusion-webui/.github/ISSUE_TEMPLATE/feature_request.yml b/stable-diffusion-webui/.github/ISSUE_TEMPLATE/feature_request.yml new file mode 100755 index 0000000..35a8874 --- /dev/null +++ b/stable-diffusion-webui/.github/ISSUE_TEMPLATE/feature_request.yml @@ -0,0 +1,40 @@ +name: Feature request +description: Suggest an idea for this project +title: "[Feature Request]: " +labels: ["enhancement"] + +body: + - type: checkboxes + attributes: + label: Is there an existing issue for this? + description: Please search to see if an issue already exists for the feature you want, and that it's not implemented in a recent build/commit. + options: + - label: I have searched the existing issues and checked the recent builds/commits + required: true + - type: markdown + attributes: + value: | + *Please fill this form with as much information as possible, provide screenshots and/or illustrations of the feature if possible* + - type: textarea + id: feature + attributes: + label: What would your feature do ? + description: Tell us about your feature in a very clear and simple way, and what problem it would solve + validations: + required: true + - type: textarea + id: workflow + attributes: + label: Proposed workflow + description: Please provide us with step by step information on how you'd like the feature to be accessed and used + value: | + 1. Go to .... + 2. Press .... + 3. ... + validations: + required: true + - type: textarea + id: misc + attributes: + label: Additional information + description: Add any other context or screenshots about the feature request here. diff --git a/stable-diffusion-webui/.github/pull_request_template.md b/stable-diffusion-webui/.github/pull_request_template.md new file mode 100755 index 0000000..c9fcda2 --- /dev/null +++ b/stable-diffusion-webui/.github/pull_request_template.md @@ -0,0 +1,15 @@ +## Description + +* a simple description of what you're trying to accomplish +* a summary of changes in code +* which issues it fixes, if any + +## Screenshots/videos: + + +## Checklist: + +- [ ] I have read [contributing wiki page](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Contributing) +- [ ] I have performed a self-review of my own code +- [ ] My code follows the [style guidelines](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Contributing#code-style) +- [ ] My code passes [tests](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Tests) diff --git a/stable-diffusion-webui/.github/workflows/on_pull_request.yaml b/stable-diffusion-webui/.github/workflows/on_pull_request.yaml new file mode 100755 index 0000000..9326c6a --- /dev/null +++ b/stable-diffusion-webui/.github/workflows/on_pull_request.yaml @@ -0,0 +1,38 @@ +name: Linter + +on: + - push + - pull_request + +jobs: + lint-python: + name: ruff + runs-on: ubuntu-latest + if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name != github.event.pull_request.base.repo.full_name + steps: + - name: Checkout Code + uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: 3.11 + # NB: there's no cache: pip here since we're not installing anything + # from the requirements.txt file(s) in the repository; it's faster + # not to have GHA download an (at the time of writing) 4 GB cache + # of PyTorch and other dependencies. + - name: Install Ruff + run: pip install ruff==0.3.3 + - name: Run Ruff + run: ruff . + lint-js: + name: eslint + runs-on: ubuntu-latest + if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name != github.event.pull_request.base.repo.full_name + steps: + - name: Checkout Code + uses: actions/checkout@v4 + - name: Install Node.js + uses: actions/setup-node@v4 + with: + node-version: 18 + - run: npm i --ci + - run: npm run lint diff --git a/stable-diffusion-webui/.github/workflows/run_tests.yaml b/stable-diffusion-webui/.github/workflows/run_tests.yaml new file mode 100755 index 0000000..0610f4f --- /dev/null +++ b/stable-diffusion-webui/.github/workflows/run_tests.yaml @@ -0,0 +1,81 @@ +name: Tests + +on: + - push + - pull_request + +jobs: + test: + name: tests on CPU with empty model + runs-on: ubuntu-latest + if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name != github.event.pull_request.base.repo.full_name + steps: + - name: Checkout Code + uses: actions/checkout@v4 + - name: Set up Python 3.10 + uses: actions/setup-python@v5 + with: + python-version: 3.10.6 + cache: pip + cache-dependency-path: | + **/requirements*txt + launch.py + - name: Cache models + id: cache-models + uses: actions/cache@v4 + with: + path: models + key: "2023-12-30" + - name: Install test dependencies + run: pip install wait-for-it -r requirements-test.txt + env: + PIP_DISABLE_PIP_VERSION_CHECK: "1" + PIP_PROGRESS_BAR: "off" + - name: Setup environment + run: python launch.py --skip-torch-cuda-test --exit + env: + PIP_DISABLE_PIP_VERSION_CHECK: "1" + PIP_PROGRESS_BAR: "off" + TORCH_INDEX_URL: https://download.pytorch.org/whl/cpu + WEBUI_LAUNCH_LIVE_OUTPUT: "1" + PYTHONUNBUFFERED: "1" + - name: Print installed packages + run: pip freeze + - name: Start test server + run: > + python -m coverage run + --data-file=.coverage.server + launch.py + --skip-prepare-environment + --skip-torch-cuda-test + --test-server + --do-not-download-clip + --no-half + --disable-opt-split-attention + --use-cpu all + --api-server-stop + 2>&1 | tee output.txt & + - name: Run tests + run: | + wait-for-it --service 127.0.0.1:7860 -t 20 + python -m pytest -vv --junitxml=test/results.xml --cov . --cov-report=xml --verify-base-url test + - name: Kill test server + if: always() + run: curl -vv -XPOST http://127.0.0.1:7860/sdapi/v1/server-stop && sleep 10 + - name: Show coverage + run: | + python -m coverage combine .coverage* + python -m coverage report -i + python -m coverage html -i + - name: Upload main app output + uses: actions/upload-artifact@v4 + if: always() + with: + name: output + path: output.txt + - name: Upload coverage HTML + uses: actions/upload-artifact@v4 + if: always() + with: + name: htmlcov + path: htmlcov diff --git a/stable-diffusion-webui/.github/workflows/warns_merge_master.yml b/stable-diffusion-webui/.github/workflows/warns_merge_master.yml new file mode 100755 index 0000000..ae2aab6 --- /dev/null +++ b/stable-diffusion-webui/.github/workflows/warns_merge_master.yml @@ -0,0 +1,19 @@ +name: Pull requests can't target master branch + +"on": + pull_request: + types: + - opened + - synchronize + - reopened + branches: + - master + +jobs: + check: + runs-on: ubuntu-latest + steps: + - name: Warning marge into master + run: | + echo -e "::warning::This pull request directly merge into \"master\" branch, normally development happens on \"dev\" branch." + exit 1 diff --git a/stable-diffusion-webui/.gitignore b/stable-diffusion-webui/.gitignore new file mode 100755 index 0000000..e81ad31 --- /dev/null +++ b/stable-diffusion-webui/.gitignore @@ -0,0 +1,44 @@ +__pycache__ +*.ckpt +*.safetensors +*.pth +.DS_Store +/ESRGAN/* +/SwinIR/* +/repositories +/venv +/tmp +/model.ckpt +/models/**/* +/GFPGANv1.3.pth +/gfpgan/weights/*.pth +/ui-config.json +/outputs +/config.json +/log +/webui.settings.bat +/embeddings +/styles.csv +/params.txt +/styles.csv.bak +/webui-user.bat +/webui-user.sh +/interrogate +/user.css +/.idea +notification.mp3 +/SwinIR +/textual_inversion +.vscode +/extensions +/test/stdout.txt +/test/stderr.txt +/cache.json* +/config_states/ +/node_modules +/package-lock.json +/.coverage* +/test/test_outputs +/cache +trace.json +/sysinfo-????-??-??-??-??.json diff --git a/stable-diffusion-webui/.pylintrc b/stable-diffusion-webui/.pylintrc new file mode 100755 index 0000000..53254e5 --- /dev/null +++ b/stable-diffusion-webui/.pylintrc @@ -0,0 +1,3 @@ +# See https://pylint.pycqa.org/en/latest/user_guide/messages/message_control.html +[MESSAGES CONTROL] +disable=C,R,W,E,I diff --git a/stable-diffusion-webui/CHANGELOG.md b/stable-diffusion-webui/CHANGELOG.md new file mode 100755 index 0000000..1b4550b --- /dev/null +++ b/stable-diffusion-webui/CHANGELOG.md @@ -0,0 +1,1085 @@ +## 1.10.1 + +### Bug Fixes: +* fix image upscale on cpu ([#16275](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16275)) + + +## 1.10.0 + +### Features: +* A lot of performance improvements (see below in Performance section) +* Stable Diffusion 3 support ([#16030](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16030), [#16164](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16164), [#16212](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16212)) + * Recommended Euler sampler; DDIM and other timestamp samplers currently not supported + * T5 text model is disabled by default, enable it in settings +* New schedulers: + * Align Your Steps ([#15751](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15751)) + * KL Optimal ([#15608](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15608)) + * Normal ([#16149](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16149)) + * DDIM ([#16149](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16149)) + * Simple ([#16142](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16142)) + * Beta ([#16235](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16235)) +* New sampler: DDIM CFG++ ([#16035](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16035)) + +### Minor: +* Option to skip CFG on early steps ([#15607](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15607)) +* Add --models-dir option ([#15742](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15742)) +* Allow mobile users to open context menu by using two fingers press ([#15682](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15682)) +* Infotext: add Lora name as TI hashes for bundled Textual Inversion ([#15679](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15679)) +* Check model's hash after downloading it to prevent corruped downloads ([#15602](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15602)) +* More extension tag filtering options ([#15627](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15627)) +* When saving AVIF, use JPEG's quality setting ([#15610](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15610)) +* Add filename pattern: `[basename]` ([#15978](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15978)) +* Add option to enable clip skip for clip L on SDXL ([#15992](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15992)) +* Option to prevent screen sleep during generation ([#16001](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16001)) +* ToggleLivePriview button in image viewer ([#16065](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16065)) +* Remove ui flashing on reloading and fast scrollong ([#16153](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16153)) +* option to disable save button log.csv ([#16242](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16242)) + +### Extensions and API: +* Add process_before_every_sampling hook ([#15984](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15984)) +* Return HTTP 400 instead of 404 on invalid sampler error ([#16140](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16140)) + +### Performance: +* [Performance 1/6] use_checkpoint = False ([#15803](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15803)) +* [Performance 2/6] Replace einops.rearrange with torch native ops ([#15804](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15804)) +* [Performance 4/6] Precompute is_sdxl_inpaint flag ([#15806](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15806)) +* [Performance 5/6] Prevent unnecessary extra networks bias backup ([#15816](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15816)) +* [Performance 6/6] Add --precision half option to avoid casting during inference ([#15820](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15820)) +* [Performance] LDM optimization patches ([#15824](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15824)) +* [Performance] Keep sigmas on CPU ([#15823](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15823)) +* Check for nans in unet only once, after all steps have been completed +* Added pption to run torch profiler for image generation + +### Bug Fixes: +* Fix for grids without comprehensive infotexts ([#15958](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15958)) +* feat: lora partial update precede full update ([#15943](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15943)) +* Fix bug where file extension had an extra '.' under some circumstances ([#15893](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15893)) +* Fix corrupt model initial load loop ([#15600](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15600)) +* Allow old sampler names in API ([#15656](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15656)) +* more old sampler scheduler compatibility ([#15681](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15681)) +* Fix Hypertile xyz ([#15831](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15831)) +* XYZ CSV skipinitialspace ([#15832](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15832)) +* fix soft inpainting on mps and xpu, torch_utils.float64 ([#15815](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15815)) +* fix extention update when not on main branch ([#15797](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15797)) +* update pickle safe filenames +* use relative path for webui-assets css ([#15757](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15757)) +* When creating a virtual environment, upgrade pip in webui.bat/webui.sh ([#15750](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15750)) +* Fix AttributeError ([#15738](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15738)) +* use script_path for webui root in launch_utils ([#15705](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15705)) +* fix extra batch mode P Transparency ([#15664](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15664)) +* use gradio theme colors in css ([#15680](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15680)) +* Fix dragging text within prompt input ([#15657](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15657)) +* Add correct mimetype for .mjs files ([#15654](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15654)) +* QOL Items - handle metadata issues more cleanly for SD models, Loras and embeddings ([#15632](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15632)) +* replace wsl-open with wslpath and explorer.exe ([#15968](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15968)) +* Fix SDXL Inpaint ([#15976](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15976)) +* multi size grid ([#15988](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15988)) +* fix Replace preview ([#16118](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16118)) +* Possible fix of wrong scale in weight decomposition ([#16151](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16151)) +* Ensure use of python from venv on Mac and Linux ([#16116](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16116)) +* Prioritize python3.10 over python3 if both are available on Linux and Mac (with fallback) ([#16092](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16092)) +* stoping generation extras ([#16085](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16085)) +* Fix SD2 loading ([#16078](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16078), [#16079](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16079)) +* fix infotext Lora hashes for hires fix different lora ([#16062](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16062)) +* Fix sampler scheduler autocorrection warning ([#16054](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16054)) +* fix ui flashing on reloading and fast scrollong ([#16153](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16153)) +* fix upscale logic ([#16239](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16239)) +* [bug] do not break progressbar on non-job actions (add wrap_gradio_call_no_job) ([#16202](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16202)) +* fix OSError: cannot write mode P as JPEG ([#16194](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16194)) + +### Other: +* fix changelog #15883 -> #15882 ([#15907](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15907)) +* ReloadUI backgroundColor --background-fill-primary ([#15864](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15864)) +* Use different torch versions for Intel and ARM Macs ([#15851](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15851)) +* XYZ override rework ([#15836](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15836)) +* scroll extensions table on overflow ([#15830](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15830)) +* img2img batch upload method ([#15817](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15817)) +* chore: sync v1.8.0 packages according to changelog ([#15783](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15783)) +* Add AVIF MIME type support to mimetype definitions ([#15739](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15739)) +* Update imageviewer.js ([#15730](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15730)) +* no-referrer ([#15641](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15641)) +* .gitignore trace.json ([#15980](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15980)) +* Bump spandrel to 0.3.4 ([#16144](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16144)) +* Defunct --max-batch-count ([#16119](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16119)) +* docs: update bug_report.yml ([#16102](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16102)) +* Maintaining Project Compatibility for Python 3.9 Users Without Upgrade Requirements. ([#16088](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16088), [#16169](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16169), [#16192](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16192)) +* Update torch for ARM Macs to 2.3.1 ([#16059](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16059)) +* remove deprecated setting dont_fix_second_order_samplers_schedule ([#16061](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16061)) +* chore: fix typos ([#16060](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16060)) +* shlex.join launch args in console log ([#16170](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16170)) +* activate venv .bat ([#16231](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16231)) +* add ids to the resize tabs in img2img ([#16218](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16218)) +* update installation guide linux ([#16178](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16178)) +* Robust sysinfo ([#16173](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16173)) +* do not send image size on paste inpaint ([#16180](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16180)) +* Fix noisy DS_Store files for MacOS ([#16166](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/16166)) + + +## 1.9.4 + +### Bug Fixes: +* pin setuptools version to fix the startup error ([#15882](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15882)) + +## 1.9.3 + +### Bug Fixes: +* fix get_crop_region_v2 ([#15594](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15594)) + +## 1.9.2 + +### Extensions and API: +* restore 1.8.0-style naming of scripts + +## 1.9.1 + +### Minor: +* Add avif support ([#15582](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15582)) +* Add filename patterns: `[sampler_scheduler]` and `[scheduler]` ([#15581](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15581)) + +### Extensions and API: +* undo adding scripts to sys.modules +* Add schedulers API endpoint ([#15577](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15577)) +* Remove API upscaling factor limits ([#15560](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15560)) + +### Bug Fixes: +* Fix images do not match / Coordinate 'right' is less than 'left' ([#15534](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15534)) +* fix: remove_callbacks_for_function should also remove from the ordered map ([#15533](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15533)) +* fix x1 upscalers ([#15555](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15555)) +* Fix cls.__module__ value in extension script ([#15532](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15532)) +* fix typo in function call (eror -> error) ([#15531](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15531)) + +### Other: +* Hide 'No Image data blocks found.' message ([#15567](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15567)) +* Allow webui.sh to be runnable from arbitrary directories containing a .git file ([#15561](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15561)) +* Compatibility with Debian 11, Fedora 34+ and openSUSE 15.4+ ([#15544](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15544)) +* numpy DeprecationWarning product -> prod ([#15547](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15547)) +* get_crop_region_v2 ([#15583](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15583), [#15587](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15587)) + + +## 1.9.0 + +### Features: +* Make refiner switchover based on model timesteps instead of sampling steps ([#14978](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14978)) +* add an option to have old-style directory view instead of tree view; stylistic changes for extra network sorting/search controls +* add UI for reordering callbacks, support for specifying callback order in extension metadata ([#15205](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15205)) +* Sgm uniform scheduler for SDXL-Lightning models ([#15325](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15325)) +* Scheduler selection in main UI ([#15333](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15333), [#15361](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15361), [#15394](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15394)) + +### Minor: +* "open images directory" button now opens the actual dir ([#14947](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14947)) +* Support inference with LyCORIS BOFT networks ([#14871](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14871), [#14973](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14973)) +* make extra network card description plaintext by default, with an option to re-enable HTML as it was +* resize handle for extra networks ([#15041](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15041)) +* cmd args: `--unix-filenames-sanitization` and `--filenames-max-length` ([#15031](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15031)) +* show extra networks parameters in HTML table rather than raw JSON ([#15131](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15131)) +* Add DoRA (weight-decompose) support for LoRA/LoHa/LoKr ([#15160](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15160), [#15283](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15283)) +* Add '--no-prompt-history' cmd args for disable last generation prompt history ([#15189](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15189)) +* update preview on Replace Preview ([#15201](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15201)) +* only fetch updates for extensions' active git branches ([#15233](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15233)) +* put upscale postprocessing UI into an accordion ([#15223](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15223)) +* Support dragdrop for URLs to read infotext ([#15262](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15262)) +* use diskcache library for caching ([#15287](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15287), [#15299](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15299)) +* Allow PNG-RGBA for Extras Tab ([#15334](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15334)) +* Support cover images embedded in safetensors metadata ([#15319](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15319)) +* faster interrupt when using NN upscale ([#15380](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15380)) +* Extras upscaler: an input field to limit maximul side length for the output image ([#15293](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15293), [#15415](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15415), [#15417](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15417), [#15425](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15425)) +* add an option to hide postprocessing options in Extras tab + +### Extensions and API: +* ResizeHandleRow - allow overriden column scale parametr ([#15004](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15004)) +* call script_callbacks.ui_settings_callback earlier; fix extra-options-section built-in extension killing the ui if using a setting that doesn't exist +* make it possible to use zoom.js outside webui context ([#15286](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15286), [#15288](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15288)) +* allow variants for extension name in metadata.ini ([#15290](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15290)) +* make reloading UI scripts optional when doing Reload UI, and off by default +* put request: gr.Request at start of img2img function similar to txt2img +* open_folder as util ([#15442](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15442)) +* make it possible to import extensions' script files as `import scripts.` ([#15423](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15423)) + +### Performance: +* performance optimization for extra networks HTML pages +* optimization for extra networks filtering +* optimization for extra networks sorting + +### Bug Fixes: +* prevent escape button causing an interrupt when no generation has been made yet +* [bug] avoid doble upscaling in inpaint ([#14966](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14966)) +* possible fix for reload button not appearing in some cases for extra networks. +* fix: the `split_threshold` parameter does not work when running Split oversized images ([#15006](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15006)) +* Fix resize-handle visability for vertical layout (mobile) ([#15010](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15010)) +* register_tmp_file also for mtime ([#15012](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15012)) +* Protect alphas_cumprod during refiner switchover ([#14979](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14979)) +* Fix EXIF orientation in API image loading ([#15062](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15062)) +* Only override emphasis if actually used in prompt ([#15141](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15141)) +* Fix emphasis infotext missing from `params.txt` ([#15142](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15142)) +* fix extract_style_text_from_prompt #15132 ([#15135](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15135)) +* Fix Soft Inpaint for AnimateDiff ([#15148](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15148)) +* edit-attention: deselect surrounding whitespace ([#15178](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15178)) +* chore: fix font not loaded ([#15183](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15183)) +* use natural sort in extra networks when ordering by path +* Fix built-in lora system bugs caused by torch.nn.MultiheadAttention ([#15190](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15190)) +* Avoid error from None in get_learned_conditioning ([#15191](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15191)) +* Add entry to MassFileLister after writing metadata ([#15199](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15199)) +* fix issue with Styles when Hires prompt is used ([#15269](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15269), [#15276](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15276)) +* Strip comments from hires fix prompt ([#15263](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15263)) +* Make imageviewer event listeners browser consistent ([#15261](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15261)) +* Fix AttributeError in OFT when trying to get MultiheadAttention weight ([#15260](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15260)) +* Add missing .mean() back ([#15239](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15239)) +* fix "Restore progress" button ([#15221](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15221)) +* fix ui-config for InputAccordion [custom_script_source] ([#15231](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15231)) +* handle 0 wheel deltaY ([#15268](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15268)) +* prevent alt menu for firefox ([#15267](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15267)) +* fix: fix syntax errors ([#15179](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15179)) +* restore outputs path ([#15307](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15307)) +* Escape btn_copy_path filename ([#15316](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15316)) +* Fix extra networks buttons when filename contains an apostrophe ([#15331](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15331)) +* escape brackets in lora random prompt generator ([#15343](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15343)) +* fix: Python version check for PyTorch installation compatibility ([#15390](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15390)) +* fix typo in call_queue.py ([#15386](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15386)) +* fix: when find already_loaded model, remove loaded by array index ([#15382](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15382)) +* minor bug fix of sd model memory management ([#15350](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15350)) +* Fix CodeFormer weight ([#15414](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15414)) +* Fix: Remove script callbacks in ordered_callbacks_map ([#15428](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15428)) +* fix limited file write (thanks, Sylwia) +* Fix extra-single-image API not doing upscale failed ([#15465](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15465)) +* error handling paste_field callables ([#15470](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15470)) + +### Hardware: +* Add training support and change lspci for Ascend NPU ([#14981](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14981)) +* Update to ROCm5.7 and PyTorch ([#14820](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14820)) +* Better workaround for Navi1, removing --pre for Navi3 ([#15224](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15224)) +* Ascend NPU wiki page ([#15228](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15228)) + +### Other: +* Update comment for Pad prompt/negative prompt v0 to add a warning about truncation, make it override the v1 implementation +* support resizable columns for touch (tablets) ([#15002](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15002)) +* Fix #14591 using translated content to do categories mapping ([#14995](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14995)) +* Use `absolute` path for normalized filepath ([#15035](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15035)) +* resizeHandle handle double tap ([#15065](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15065)) +* --dat-models-path cmd flag ([#15039](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15039)) +* Add a direct link to the binary release ([#15059](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15059)) +* upscaler_utils: Reduce logging ([#15084](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15084)) +* Fix various typos with crate-ci/typos ([#15116](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15116)) +* fix_jpeg_live_preview ([#15102](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15102)) +* [alternative fix] can't load webui if selected wrong extra option in ui ([#15121](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15121)) +* Error handling for unsupported transparency ([#14958](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14958)) +* Add model description to searched terms ([#15198](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15198)) +* bump action version ([#15272](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15272)) +* PEP 604 annotations ([#15259](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15259)) +* Automatically Set the Scale by value when user selects an Upscale Model ([#15244](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15244)) +* move postprocessing-for-training into builtin extensions ([#15222](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15222)) +* type hinting in shared.py ([#15211](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15211)) +* update ruff to 0.3.3 +* Update pytorch lightning utilities ([#15310](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15310)) +* Add Size as an XYZ Grid option ([#15354](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15354)) +* Use HF_ENDPOINT variable for HuggingFace domain with default ([#15443](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15443)) +* re-add update_file_entry ([#15446](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15446)) +* create_infotext allow index and callable, re-work Hires prompt infotext ([#15460](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15460)) +* update restricted_opts to include more options for --hide-ui-dir-config ([#15492](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15492)) + + +## 1.8.0 + +### Features: +* Update torch to version 2.1.2 +* Soft Inpainting ([#14208](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14208)) +* FP8 support ([#14031](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14031), [#14327](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14327)) +* Support for SDXL-Inpaint Model ([#14390](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14390)) +* Use Spandrel for upscaling and face restoration architectures ([#14425](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14425), [#14467](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14467), [#14473](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14473), [#14474](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14474), [#14477](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14477), [#14476](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14476), [#14484](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14484), [#14500](https://github.com/AUTOMATIC1111/stable-difusion-webui/pull/14500), [#14501](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14501), [#14504](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14504), [#14524](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14524), [#14809](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14809)) +* Automatic backwards version compatibility (when loading infotexts from old images with program version specified, will add compatibility settings) +* Implement zero terminal SNR noise schedule option (**[SEED BREAKING CHANGE](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Seed-breaking-changes#180-dev-170-225-2024-01-01---zero-terminal-snr-noise-schedule-option)**, [#14145](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14145), [#14979](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14979)) +* Add a [✨] button to run hires fix on selected image in the gallery (with help from [#14598](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14598), [#14626](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14626), [#14728](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14728)) +* [Separate assets repository](https://github.com/AUTOMATIC1111/stable-diffusion-webui-assets); serve fonts locally rather than from google's servers +* Official LCM Sampler Support ([#14583](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14583)) +* Add support for DAT upscaler models ([#14690](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14690), [#15039](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15039)) +* Extra Networks Tree View ([#14588](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14588), [#14900](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14900)) +* NPU Support ([#14801](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14801)) +* Prompt comments support + +### Minor: +* Allow pasting in WIDTHxHEIGHT strings into the width/height fields ([#14296](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14296)) +* add option: Live preview in full page image viewer ([#14230](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14230), [#14307](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14307)) +* Add keyboard shortcuts for generate/skip/interrupt ([#14269](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14269)) +* Better TCMALLOC support on different platforms ([#14227](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14227), [#14883](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14883), [#14910](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14910)) +* Lora not found warning ([#14464](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14464)) +* Adding negative prompts to Loras in extra networks ([#14475](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14475)) +* xyz_grid: allow varying the seed along an axis separate from axis options ([#12180](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12180)) +* option to convert VAE to bfloat16 (implementation of [#9295](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/9295)) +* Better IPEX support ([#14229](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14229), [#14353](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14353), [#14559](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14559), [#14562](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14562), [#14597](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14597)) +* Option to interrupt after current generation rather than immediately ([#13653](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13653), [#14659](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14659)) +* Fullscreen Preview control fading/disable ([#14291](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14291)) +* Finer settings freezing control ([#13789](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13789)) +* Increase Upscaler Limits ([#14589](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14589)) +* Adjust brush size with hotkeys ([#14638](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14638)) +* Add checkpoint info to csv log file when saving images ([#14663](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14663)) +* Make more columns resizable ([#14740](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14740), [#14884](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14884)) +* Add an option to not overlay original image for inpainting for #14727 +* Add Pad conds v0 option to support same generation with DDIM as before 1.6.0 +* Add "Interrupting..." placeholder. +* Button for refresh extensions list ([#14857](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14857)) +* Add an option to disable normalization after calculating emphasis. ([#14874](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14874)) +* When counting tokens, also include enabled styles (can be disabled in settings to revert to previous behavior) +* Configuration for the [📂] button for image gallery ([#14947](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14947)) +* Support inference with LyCORIS BOFT networks ([#14871](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14871), [#14973](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14973)) +* support resizable columns for touch (tablets) ([#15002](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15002)) + +### Extensions and API: +* Removed packages from requirements: basicsr, gfpgan, realesrgan; as well as their dependencies: absl-py, addict, beautifulsoup4, future, gdown, grpcio, importlib-metadata, lmdb, lpips, Markdown, platformdirs, PySocks, soupsieve, tb-nightly, tensorboard-data-server, tomli, Werkzeug, yapf, zipp, soupsieve +* Enable task ids for API ([#14314](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14314)) +* add override_settings support for infotext API +* rename generation_parameters_copypaste module to infotext_utils +* prevent crash due to Script __init__ exception ([#14407](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14407)) +* Bump numpy to 1.26.2 ([#14471](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14471)) +* Add utility to inspect a model's dtype/device ([#14478](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14478)) +* Implement general forward method for all method in built-in lora ext ([#14547](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14547)) +* Execute model_loaded_callback after moving to target device ([#14563](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14563)) +* Add self to CFGDenoiserParams ([#14573](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14573)) +* Allow TLS with API only mode (--nowebui) ([#14593](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14593)) +* New callback: postprocess_image_after_composite ([#14657](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14657)) +* modules/api/api.py: add api endpoint to refresh embeddings list ([#14715](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14715)) +* set_named_arg ([#14773](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14773)) +* add before_token_counter callback and use it for prompt comments +* ResizeHandleRow - allow overridden column scale parameter ([#15004](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15004)) + +### Performance: +* Massive performance improvement for extra networks directories with a huge number of files in them in an attempt to tackle #14507 ([#14528](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14528)) +* Reduce unnecessary re-indexing extra networks directory ([#14512](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14512)) +* Avoid unnecessary `isfile`/`exists` calls ([#14527](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14527)) + +### Bug Fixes: +* fix multiple bugs related to styles multi-file support ([#14203](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14203), [#14276](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14276), [#14707](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14707)) +* Lora fixes ([#14300](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14300), [#14237](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14237), [#14546](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14546), [#14726](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14726)) +* Re-add setting lost as part of e294e46 ([#14266](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14266)) +* fix extras caption BLIP ([#14330](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14330)) +* include infotext into saved init image for img2img ([#14452](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14452)) +* xyz grid handle axis_type is None ([#14394](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14394)) +* Update Added (Fixed) IPV6 Functionality When there is No Webui Argument Passed webui.py ([#14354](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14354)) +* fix API thread safe issues of txt2img and img2img ([#14421](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14421)) +* handle selectable script_index is None ([#14487](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14487)) +* handle config.json failed to load ([#14525](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14525), [#14767](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14767)) +* paste infotext cast int as float ([#14523](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14523)) +* Ensure GRADIO_ANALYTICS_ENABLED is set early enough ([#14537](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14537)) +* Fix logging configuration again ([#14538](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14538)) +* Handle CondFunc exception when resolving attributes ([#14560](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14560)) +* Fix extras big batch crashes ([#14699](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14699)) +* Fix using wrong model caused by alias ([#14655](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14655)) +* Add # to the invalid_filename_chars list ([#14640](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14640)) +* Fix extension check for requirements ([#14639](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14639)) +* Fix tab indexes are reset after restart UI ([#14637](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14637)) +* Fix nested manual cast ([#14689](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14689)) +* Keep postprocessing upscale selected tab after restart ([#14702](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14702)) +* XYZ grid: filter out blank vals when axis is int or float type (like int axis seed) ([#14754](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14754)) +* fix CLIP Interrogator topN regex ([#14775](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14775)) +* Fix dtype error in MHA layer/change dtype checking mechanism for manual cast ([#14791](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14791)) +* catch load style.csv error ([#14814](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14814)) +* fix error when editing extra networks card +* fix extra networks metadata failing to work properly when you create the .json file with metadata for the first time. +* util.walk_files extensions case insensitive ([#14879](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14879)) +* if extensions page not loaded, prevent apply ([#14873](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14873)) +* call the right function for token counter in img2img +* Fix the bugs that search/reload will disappear when using other ExtraNetworks extensions ([#14939](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14939)) +* Gracefully handle mtime read exception from cache ([#14933](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14933)) +* Only trigger interrupt on `Esc` when interrupt button visible ([#14932](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14932)) +* Disable prompt token counters option actually disables token counting rather than just hiding results. +* avoid double upscaling in inpaint ([#14966](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14966)) +* Fix #14591 using translated content to do categories mapping ([#14995](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14995)) +* fix: the `split_threshold` parameter does not work when running Split oversized images ([#15006](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15006)) +* Fix resize-handle for mobile ([#15010](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15010), [#15065](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15065)) + +### Other: +* Assign id for "extra_options". Replace numeric field with slider. ([#14270](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14270)) +* change state dict comparison to ref compare ([#14216](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14216)) +* Bump torch-rocm to 5.6/5.7 ([#14293](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14293)) +* Base output path off data path ([#14446](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14446)) +* reorder training preprocessing modules in extras tab ([#14367](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14367)) +* Remove `cleanup_models` code ([#14472](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14472)) +* only rewrite ui-config when there is change ([#14352](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14352)) +* Fix lint issue from 501993eb ([#14495](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14495)) +* Update README.md ([#14548](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14548)) +* hires button, fix seeds () +* Logging: set formatter correctly for fallback logger too ([#14618](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14618)) +* Read generation info from infotexts rather than json for internal needs (save, extract seed from generated pic) ([#14645](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14645)) +* improve get_crop_region ([#14709](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14709)) +* Bump safetensors' version to 0.4.2 ([#14782](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14782)) +* add tooltip create_submit_box ([#14803](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14803)) +* extensions tab table row hover highlight ([#14885](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14885)) +* Always add timestamp to displayed image ([#14890](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14890)) +* Added core.filemode=false so doesn't track changes in file permission… ([#14930](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14930)) +* Normalize command-line argument paths ([#14934](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14934), [#15035](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15035)) +* Use original App Title in progress bar ([#14916](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14916)) +* register_tmp_file also for mtime ([#15012](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15012)) + +## 1.7.0 + +### Features: +* settings tab rework: add search field, add categories, split UI settings page into many +* add altdiffusion-m18 support ([#13364](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13364)) +* support inference with LyCORIS GLora networks ([#13610](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13610)) +* add lora-embedding bundle system ([#13568](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13568)) +* option to move prompt from top row into generation parameters +* add support for SSD-1B ([#13865](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13865)) +* support inference with OFT networks ([#13692](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13692)) +* script metadata and DAG sorting mechanism ([#13944](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13944)) +* support HyperTile optimization ([#13948](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13948)) +* add support for SD 2.1 Turbo ([#14170](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14170)) +* remove Train->Preprocessing tab and put all its functionality into Extras tab +* initial IPEX support for Intel Arc GPU ([#14171](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14171)) + +### Minor: +* allow reading model hash from images in img2img batch mode ([#12767](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12767)) +* add option to align with sgm repo's sampling implementation ([#12818](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12818)) +* extra field for lora metadata viewer: `ss_output_name` ([#12838](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12838)) +* add action in settings page to calculate all SD checkpoint hashes ([#12909](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12909)) +* add button to copy prompt to style editor ([#12975](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12975)) +* add --skip-load-model-at-start option ([#13253](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13253)) +* write infotext to gif images +* read infotext from gif images ([#13068](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13068)) +* allow configuring the initial state of InputAccordion in ui-config.json ([#13189](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13189)) +* allow editing whitespace delimiters for ctrl+up/ctrl+down prompt editing ([#13444](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13444)) +* prevent accidentally closing popup dialogs ([#13480](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13480)) +* added option to play notification sound or not ([#13631](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13631)) +* show the preview image in the full screen image viewer if available ([#13459](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13459)) +* support for webui.settings.bat ([#13638](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13638)) +* add an option to not print stack traces on ctrl+c +* start/restart generation by Ctrl (Alt) + Enter ([#13644](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13644)) +* update prompts_from_file script to allow concatenating entries with the general prompt ([#13733](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13733)) +* added a visible checkbox to input accordion +* added an option to hide all txt2img/img2img parameters in an accordion ([#13826](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13826)) +* added 'Path' sorting option for Extra network cards ([#13968](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13968)) +* enable prompt hotkeys in style editor ([#13931](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13931)) +* option to show batch img2img results in UI ([#14009](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14009)) +* infotext updates: add option to disregard certain infotext fields, add option to not include VAE in infotext, add explanation to infotext settings page, move some options to infotext settings page +* add FP32 fallback support on sd_vae_approx ([#14046](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14046)) +* support XYZ scripts / split hires path from unet ([#14126](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14126)) +* allow use of multiple styles csv files ([#14125](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14125)) +* make extra network card description plaintext by default, with an option (Treat card description as HTML) to re-enable HTML as it was (originally by [#13241](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13241)) + +### Extensions and API: +* update gradio to 3.41.2 +* support installed extensions list api ([#12774](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12774)) +* update pnginfo API to return dict with parsed values +* add noisy latent to `ExtraNoiseParams` for callback ([#12856](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12856)) +* show extension datetime in UTC ([#12864](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12864), [#12865](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12865), [#13281](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13281)) +* add an option to choose how to combine hires fix and refiner +* include program version in info response. ([#13135](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13135)) +* sd_unet support for SDXL +* patch DDPM.register_betas so that users can put given_betas in model yaml ([#13276](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13276)) +* xyz_grid: add prepare ([#13266](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13266)) +* allow multiple localization files with same language in extensions ([#13077](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13077)) +* add onEdit function for js and rework token-counter.js to use it +* fix the key error exception when processing override_settings keys ([#13567](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13567)) +* ability for extensions to return custom data via api in response.images ([#13463](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13463)) +* call state.jobnext() before postproces*() ([#13762](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13762)) +* add option to set notification sound volume ([#13884](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13884)) +* update Ruff to 0.1.6 ([#14059](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14059)) +* add Block component creation callback ([#14119](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14119)) +* catch uncaught exception with ui creation scripts ([#14120](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14120)) +* use extension name for determining an extension is installed in the index ([#14063](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14063)) +* update is_installed() from launch_utils.py to fix reinstalling already installed packages ([#14192](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14192)) + +### Bug Fixes: +* fix pix2pix producing bad results +* fix defaults settings page breaking when any of main UI tabs are hidden +* fix error that causes some extra networks to be disabled if both and are present in the prompt +* fix for Reload UI function: if you reload UI on one tab, other opened tabs will no longer stop working +* prevent duplicate resize handler ([#12795](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12795)) +* small typo: vae resolve bug ([#12797](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12797)) +* hide broken image crop tool ([#12792](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12792)) +* don't show hidden samplers in dropdown for XYZ script ([#12780](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12780)) +* fix style editing dialog breaking if it's opened in both img2img and txt2img tabs +* hide --gradio-auth and --api-auth values from /internal/sysinfo report +* add missing infotext for RNG in options ([#12819](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12819)) +* fix notification not playing when built-in webui tab is inactive ([#12834](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12834)) +* honor `--skip-install` for extension installers ([#12832](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12832)) +* don't print blank stdout in extension installers ([#12833](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12833), [#12855](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12855)) +* get progressbar to display correctly in extensions tab +* keep order in list of checkpoints when loading model that doesn't have a checksum +* fix inpainting models in txt2img creating black pictures +* fix generation params regex ([#12876](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12876)) +* fix batch img2img output dir with script ([#12926](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12926)) +* fix #13080 - Hypernetwork/TI preview generation ([#13084](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13084)) +* fix bug with sigma min/max overrides. ([#12995](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12995)) +* more accurate check for enabling cuDNN benchmark on 16XX cards ([#12924](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12924)) +* don't use multicond parser for negative prompt counter ([#13118](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13118)) +* fix data-sort-name containing spaces ([#13412](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13412)) +* update card on correct tab when editing metadata ([#13411](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13411)) +* fix viewing/editing metadata when filename contains an apostrophe ([#13395](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13395)) +* fix: --sd_model in "Prompts from file or textbox" script is not working ([#13302](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13302)) +* better Support for Portable Git ([#13231](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13231)) +* fix issues when webui_dir is not work_dir ([#13210](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13210)) +* fix: lora-bias-backup don't reset cache ([#13178](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13178)) +* account for customizable extra network separators whyen removing extra network text from the prompt ([#12877](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12877)) +* re fix batch img2img output dir with script ([#13170](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13170)) +* fix `--ckpt-dir` path separator and option use `short name` for checkpoint dropdown ([#13139](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13139)) +* consolidated allowed preview formats, Fix extra network `.gif` not woking as preview ([#13121](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13121)) +* fix venv_dir=- environment variable not working as expected on linux ([#13469](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13469)) +* repair unload sd checkpoint button +* edit-attention fixes ([#13533](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13533)) +* fix bug when using --gfpgan-models-path ([#13718](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13718)) +* properly apply sort order for extra network cards when selected from dropdown +* fixes generation restart not working for some users when 'Ctrl+Enter' is pressed ([#13962](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13962)) +* thread safe extra network list_items ([#13014](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13014)) +* fix not able to exit metadata popup when pop up is too big ([#14156](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14156)) +* fix auto focal point crop for opencv >= 4.8 ([#14121](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14121)) +* make 'use-cpu all' actually apply to 'all' ([#14131](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14131)) +* extras tab batch: actually use original filename +* make webui not crash when running with --disable-all-extensions option + +### Other: +* non-local condition ([#12814](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12814)) +* fix minor typos ([#12827](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12827)) +* remove xformers Python version check ([#12842](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12842)) +* style: file-metadata word-break ([#12837](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12837)) +* revert SGM noise multiplier change for img2img because it breaks hires fix +* do not change quicksettings dropdown option when value returned is `None` ([#12854](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12854)) +* [RC 1.6.0 - zoom is partly hidden] Update style.css ([#12839](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12839)) +* chore: change extension time format ([#12851](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12851)) +* WEBUI.SH - Use torch 2.1.0 release candidate for Navi 3 ([#12929](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12929)) +* add Fallback at images.read_info_from_image if exif data was invalid ([#13028](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13028)) +* update cmd arg description ([#12986](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12986)) +* fix: update shared.opts.data when add_option ([#12957](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12957), [#13213](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13213)) +* restore missing tooltips ([#12976](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12976)) +* use default dropdown padding on mobile ([#12880](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12880)) +* put enable console prompts option into settings from commandline args ([#13119](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13119)) +* fix some deprecated types ([#12846](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12846)) +* bump to torchsde==0.2.6 ([#13418](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13418)) +* update dragdrop.js ([#13372](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13372)) +* use orderdict as lru cache:opt/bug ([#13313](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13313)) +* XYZ if not include sub grids do not save sub grid ([#13282](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13282)) +* initialize state.time_start befroe state.job_count ([#13229](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13229)) +* fix fieldname regex ([#13458](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13458)) +* change denoising_strength default to None. ([#13466](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13466)) +* fix regression ([#13475](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13475)) +* fix IndexError ([#13630](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13630)) +* fix: checkpoints_loaded:{checkpoint:state_dict}, model.load_state_dict issue in dict value empty ([#13535](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13535)) +* update bug_report.yml ([#12991](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12991)) +* requirements_versions httpx==0.24.1 ([#13839](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13839)) +* fix parenthesis auto selection ([#13829](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13829)) +* fix #13796 ([#13797](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13797)) +* corrected a typo in `modules/cmd_args.py` ([#13855](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13855)) +* feat: fix randn found element of type float at pos 2 ([#14004](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14004)) +* adds tqdm handler to logging_config.py for progress bar integration ([#13996](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13996)) +* hotfix: call shared.state.end() after postprocessing done ([#13977](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13977)) +* fix dependency address patch 1 ([#13929](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13929)) +* save sysinfo as .json ([#14035](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14035)) +* move exception_records related methods to errors.py ([#14084](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14084)) +* compatibility ([#13936](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13936)) +* json.dump(ensure_ascii=False) ([#14108](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14108)) +* dir buttons start with / so only the correct dir will be shown and no… ([#13957](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13957)) +* alternate implementation for unet forward replacement that does not depend on hijack being applied +* re-add `keyedit_delimiters_whitespace` setting lost as part of commit e294e46 ([#14178](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14178)) +* fix `save_samples` being checked early when saving masked composite ([#14177](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14177)) +* slight optimization for mask and mask_composite ([#14181](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14181)) +* add import_hook hack to work around basicsr/torchvision incompatibility ([#14186](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14186)) + +## 1.6.1 + +### Bug Fixes: + * fix an error causing the webui to fail to start ([#13839](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13839)) + +## 1.6.0 + +### Features: + * refiner support [#12371](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12371) + * add NV option for Random number generator source setting, which allows to generate same pictures on CPU/AMD/Mac as on NVidia videocards + * add style editor dialog + * hires fix: add an option to use a different checkpoint for second pass ([#12181](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12181)) + * option to keep multiple loaded models in memory ([#12227](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12227)) + * new samplers: Restart, DPM++ 2M SDE Exponential, DPM++ 2M SDE Heun, DPM++ 2M SDE Heun Karras, DPM++ 2M SDE Heun Exponential, DPM++ 3M SDE, DPM++ 3M SDE Karras, DPM++ 3M SDE Exponential ([#12300](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12300), [#12519](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12519), [#12542](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12542)) + * rework DDIM, PLMS, UniPC to use CFG denoiser same as in k-diffusion samplers: + * makes all of them work with img2img + * makes prompt composition possible (AND) + * makes them available for SDXL + * always show extra networks tabs in the UI ([#11808](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/11808)) + * use less RAM when creating models ([#11958](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/11958), [#12599](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12599)) + * textual inversion inference support for SDXL + * extra networks UI: show metadata for SD checkpoints + * checkpoint merger: add metadata support + * prompt editing and attention: add support for whitespace after the number ([ red : green : 0.5 ]) (seed breaking change) ([#12177](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12177)) + * VAE: allow selecting own VAE for each checkpoint (in user metadata editor) + * VAE: add selected VAE to infotext + * options in main UI: add own separate setting for txt2img and img2img, correctly read values from pasted infotext, add setting for column count ([#12551](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12551)) + * add resize handle to txt2img and img2img tabs, allowing to change the amount of horizontable space given to generation parameters and resulting image gallery ([#12687](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12687), [#12723](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12723)) + * change default behavior for batching cond/uncond -- now it's on by default, and is disabled by an UI setting (Optimizatios -> Batch cond/uncond) - if you are on lowvram/medvram and are getting OOM exceptions, you will need to enable it + * show current position in queue and make it so that requests are processed in the order of arrival ([#12707](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12707)) + * add `--medvram-sdxl` flag that only enables `--medvram` for SDXL models + * prompt editing timeline has separate range for first pass and hires-fix pass (seed breaking change) ([#12457](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12457)) + +### Minor: + * img2img batch: RAM savings, VRAM savings, .tif, .tiff in img2img batch ([#12120](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12120), [#12514](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12514), [#12515](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12515)) + * postprocessing/extras: RAM savings ([#12479](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12479)) + * XYZ: in the axis labels, remove pathnames from model filenames + * XYZ: support hires sampler ([#12298](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12298)) + * XYZ: new option: use text inputs instead of dropdowns ([#12491](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12491)) + * add gradio version warning + * sort list of VAE checkpoints ([#12297](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12297)) + * use transparent white for mask in inpainting, along with an option to select the color ([#12326](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12326)) + * move some settings to their own section: img2img, VAE + * add checkbox to show/hide dirs for extra networks + * Add TAESD(or more) options for all the VAE encode/decode operation ([#12311](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12311)) + * gradio theme cache, new gradio themes, along with explanation that the user can input his own values ([#12346](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12346), [#12355](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12355)) + * sampler fixes/tweaks: s_tmax, s_churn, s_noise, s_tmax ([#12354](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12354), [#12356](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12356), [#12357](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12357), [#12358](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12358), [#12375](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12375), [#12521](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12521)) + * update README.md with correct instructions for Linux installation ([#12352](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12352)) + * option to not save incomplete images, on by default ([#12338](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12338)) + * enable cond cache by default + * git autofix for repos that are corrupted ([#12230](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12230)) + * allow to open images in new browser tab by middle mouse button ([#12379](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12379)) + * automatically open webui in browser when running "locally" ([#12254](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12254)) + * put commonly used samplers on top, make DPM++ 2M Karras the default choice + * zoom and pan: option to auto-expand a wide image, improved integration ([#12413](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12413), [#12727](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12727)) + * option to cache Lora networks in memory + * rework hires fix UI to use accordion + * face restoration and tiling moved to settings - use "Options in main UI" setting if you want them back + * change quicksettings items to have variable width + * Lora: add Norm module, add support for bias ([#12503](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12503)) + * Lora: output warnings in UI rather than fail for unfitting loras; switch to logging for error output in console + * support search and display of hashes for all extra network items ([#12510](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12510)) + * add extra noise param for img2img operations ([#12564](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12564)) + * support for Lora with bias ([#12584](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12584)) + * make interrupt quicker ([#12634](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12634)) + * configurable gallery height ([#12648](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12648)) + * make results column sticky ([#12645](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12645)) + * more hash filename patterns ([#12639](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12639)) + * make image viewer actually fit the whole page ([#12635](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12635)) + * make progress bar work independently from live preview display which results in it being updated a lot more often + * forbid Full live preview method for medvram and add a setting to undo the forbidding + * make it possible to localize tooltips and placeholders + * add option to align with sgm repo's sampling implementation ([#12818](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12818)) + * Restore faces and Tiling generation parameters have been moved to settings out of main UI + * if you want to put them back into main UI, use `Options in main UI` setting on the UI page. + +### Extensions and API: + * gradio 3.41.2 + * also bump versions for packages: transformers, GitPython, accelerate, scikit-image, timm, tomesd + * support tooltip kwarg for gradio elements: gr.Textbox(label='hello', tooltip='world') + * properly clear the total console progressbar when using txt2img and img2img from API + * add cmd_arg --disable-extra-extensions and --disable-all-extensions ([#12294](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12294)) + * shared.py and webui.py split into many files + * add --loglevel commandline argument for logging + * add a custom UI element that combines accordion and checkbox + * avoid importing gradio in tests because it spams warnings + * put infotext label for setting into OptionInfo definition rather than in a separate list + * make `StableDiffusionProcessingImg2Img.mask_blur` a property, make more inline with PIL `GaussianBlur` ([#12470](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12470)) + * option to make scripts UI without gr.Group + * add a way for scripts to register a callback for before/after just a single component's creation + * use dataclass for StableDiffusionProcessing + * store patches for Lora in a specialized module instead of inside torch + * support http/https URLs in API ([#12663](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12663), [#12698](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12698)) + * add extra noise callback ([#12616](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12616)) + * dump current stack traces when exiting with SIGINT + * add type annotations for extra fields of shared.sd_model + +### Bug Fixes: + * Don't crash if out of local storage quota for javascriot localStorage + * XYZ plot do not fail if an exception occurs + * fix missing TI hash in infotext if generation uses both negative and positive TI ([#12269](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12269)) + * localization fixes ([#12307](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12307)) + * fix sdxl model invalid configuration after the hijack + * correctly toggle extras checkbox for infotext paste ([#12304](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12304)) + * open raw sysinfo link in new page ([#12318](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12318)) + * prompt parser: Account for empty field in alternating words syntax ([#12319](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12319)) + * add tab and carriage return to invalid filename chars ([#12327](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12327)) + * fix api only Lora not working ([#12387](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12387)) + * fix options in main UI misbehaving when there's just one element + * make it possible to use a sampler from infotext even if it's hidden in the dropdown + * fix styles missing from the prompt in infotext when making a grid of batch of multiplie images + * prevent bogus progress output in console when calculating hires fix dimensions + * fix --use-textbox-seed + * fix broken `Lora/Networks: use old method` option ([#12466](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12466)) + * properly return `None` for VAE hash when using `--no-hashing` ([#12463](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12463)) + * MPS/macOS fixes and optimizations ([#12526](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12526)) + * add second_order to samplers that mistakenly didn't have it + * when refreshing cards in extra networks UI, do not discard user's custom resolution + * fix processing error that happens if batch_size is not a multiple of how many prompts/negative prompts there are ([#12509](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12509)) + * fix inpaint upload for alpha masks ([#12588](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12588)) + * fix exception when image sizes are not integers ([#12586](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12586)) + * fix incorrect TAESD Latent scale ([#12596](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12596)) + * auto add data-dir to gradio-allowed-path ([#12603](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12603)) + * fix exception if extensuions dir is missing ([#12607](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12607)) + * fix issues with api model-refresh and vae-refresh ([#12638](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12638)) + * fix img2img background color for transparent images option not being used ([#12633](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12633)) + * attempt to resolve NaN issue with unstable VAEs in fp32 mk2 ([#12630](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12630)) + * implement missing undo hijack for SDXL + * fix xyz swap axes ([#12684](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12684)) + * fix errors in backup/restore tab if any of config files are broken ([#12689](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12689)) + * fix SD VAE switch error after model reuse ([#12685](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12685)) + * fix trying to create images too large for the chosen format ([#12667](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12667)) + * create Gradio temp directory if necessary ([#12717](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12717)) + * prevent possible cache loss if exiting as it's being written by using an atomic operation to replace the cache with the new version + * set devices.dtype_unet correctly + * run RealESRGAN on GPU for non-CUDA devices ([#12737](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12737)) + * prevent extra network buttons being obscured by description for very small card sizes ([#12745](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12745)) + * fix error that causes some extra networks to be disabled if both and are present in the prompt + * fix defaults settings page breaking when any of main UI tabs are hidden + * fix incorrect save/display of new values in Defaults page in settings + * fix for Reload UI function: if you reload UI on one tab, other opened tabs will no longer stop working + * fix an error that prevents VAE being reloaded after an option change if a VAE near the checkpoint exists ([#12797](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12737)) + * hide broken image crop tool ([#12792](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12737)) + * don't show hidden samplers in dropdown for XYZ script ([#12780](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12737)) + * fix style editing dialog breaking if it's opened in both img2img and txt2img tabs + * fix a bug allowing users to bypass gradio and API authentication (reported by vysecurity) + * fix notification not playing when built-in webui tab is inactive ([#12834](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12834)) + * honor `--skip-install` for extension installers ([#12832](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12832)) + * don't print blank stdout in extension installers ([#12833](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12832), [#12855](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12855)) + * do not change quicksettings dropdown option when value returned is `None` ([#12854](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12854)) + * get progressbar to display correctly in extensions tab + + +## 1.5.2 + +### Bug Fixes: + * fix memory leak when generation fails + * update doggettx cross attention optimization to not use an unreasonable amount of memory in some edge cases -- suggestion by MorkTheOrk + + +## 1.5.1 + +### Minor: + * support parsing text encoder blocks in some new LoRAs + * delete scale checker script due to user demand + +### Extensions and API: + * add postprocess_batch_list script callback + +### Bug Fixes: + * fix TI training for SD1 + * fix reload altclip model error + * prepend the pythonpath instead of overriding it + * fix typo in SD_WEBUI_RESTARTING + * if txt2img/img2img raises an exception, finally call state.end() + * fix composable diffusion weight parsing + * restyle Startup profile for black users + * fix webui not launching with --nowebui + * catch exception for non git extensions + * fix some options missing from /sdapi/v1/options + * fix for extension update status always saying "unknown" + * fix display of extra network cards that have `<>` in the name + * update lora extension to work with python 3.8 + + +## 1.5.0 + +### Features: + * SD XL support + * user metadata system for custom networks + * extended Lora metadata editor: set activation text, default weight, view tags, training info + * Lora extension rework to include other types of networks (all that were previously handled by LyCORIS extension) + * show github stars for extensions + * img2img batch mode can read extra stuff from png info + * img2img batch works with subdirectories + * hotkeys to move prompt elements: alt+left/right + * restyle time taken/VRAM display + * add textual inversion hashes to infotext + * optimization: cache git extension repo information + * move generate button next to the generated picture for mobile clients + * hide cards for networks of incompatible Stable Diffusion version in Lora extra networks interface + * skip installing packages with pip if they all are already installed - startup speedup of about 2 seconds + +### Minor: + * checkbox to check/uncheck all extensions in the Installed tab + * add gradio user to infotext and to filename patterns + * allow gif for extra network previews + * add options to change colors in grid + * use natural sort for items in extra networks + * Mac: use empty_cache() from torch 2 to clear VRAM + * added automatic support for installing the right libraries for Navi3 (AMD) + * add option SWIN_torch_compile to accelerate SwinIR upscale + * suppress printing TI embedding info at start to console by default + * speedup extra networks listing + * added `[none]` filename token. + * removed thumbs extra networks view mode (use settings tab to change width/height/scale to get thumbs) + * add always_discard_next_to_last_sigma option to XYZ plot + * automatically switch to 32-bit float VAE if the generated picture has NaNs without the need for `--no-half-vae` commandline flag. + +### Extensions and API: + * api endpoints: /sdapi/v1/server-kill, /sdapi/v1/server-restart, /sdapi/v1/server-stop + * allow Script to have custom metaclass + * add model exists status check /sdapi/v1/options + * rename --add-stop-route to --api-server-stop + * add `before_hr` script callback + * add callback `after_extra_networks_activate` + * disable rich exception output in console for API by default, use WEBUI_RICH_EXCEPTIONS env var to enable + * return http 404 when thumb file not found + * allow replacing extensions index with environment variable + +### Bug Fixes: + * fix for catch errors when retrieving extension index #11290 + * fix very slow loading speed of .safetensors files when reading from network drives + * API cache cleanup + * fix UnicodeEncodeError when writing to file CLIP Interrogator batch mode + * fix warning of 'has_mps' deprecated from PyTorch + * fix problem with extra network saving images as previews losing generation info + * fix throwing exception when trying to resize image with I;16 mode + * fix for #11534: canvas zoom and pan extension hijacking shortcut keys + * fixed launch script to be runnable from any directory + * don't add "Seed Resize: -1x-1" to API image metadata + * correctly remove end parenthesis with ctrl+up/down + * fixing --subpath on newer gradio version + * fix: check fill size none zero when resize (fixes #11425) + * use submit and blur for quick settings textbox + * save img2img batch with images.save_image() + * prevent running preload.py for disabled extensions + * fix: previously, model name was added together with directory name to infotext and to [model_name] filename pattern; directory name is now not included + + +## 1.4.1 + +### Bug Fixes: + * add queue lock for refresh-checkpoints + +## 1.4.0 + +### Features: + * zoom controls for inpainting + * run basic torch calculation at startup in parallel to reduce the performance impact of first generation + * option to pad prompt/neg prompt to be same length + * remove taming_transformers dependency + * custom k-diffusion scheduler settings + * add an option to show selected settings in main txt2img/img2img UI + * sysinfo tab in settings + * infer styles from prompts when pasting params into the UI + * an option to control the behavior of the above + +### Minor: + * bump Gradio to 3.32.0 + * bump xformers to 0.0.20 + * Add option to disable token counters + * tooltip fixes & optimizations + * make it possible to configure filename for the zip download + * `[vae_filename]` pattern for filenames + * Revert discarding penultimate sigma for DPM-Solver++(2M) SDE + * change UI reorder setting to multiselect + * read version info form CHANGELOG.md if git version info is not available + * link footer API to Wiki when API is not active + * persistent conds cache (opt-in optimization) + +### Extensions: + * After installing extensions, webui properly restarts the process rather than reloads the UI + * Added VAE listing to web API. Via: /sdapi/v1/sd-vae + * custom unet support + * Add onAfterUiUpdate callback + * refactor EmbeddingDatabase.register_embedding() to allow unregistering + * add before_process callback for scripts + * add ability for alwayson scripts to specify section and let user reorder those sections + +### Bug Fixes: + * Fix dragging text to prompt + * fix incorrect quoting for infotext values with colon in them + * fix "hires. fix" prompt sharing same labels with txt2img_prompt + * Fix s_min_uncond default type int + * Fix for #10643 (Inpainting mask sometimes not working) + * fix bad styling for thumbs view in extra networks #10639 + * fix for empty list of optimizations #10605 + * small fixes to prepare_tcmalloc for Debian/Ubuntu compatibility + * fix --ui-debug-mode exit + * patch GitPython to not use leaky persistent processes + * fix duplicate Cross attention optimization after UI reload + * torch.cuda.is_available() check for SdOptimizationXformers + * fix hires fix using wrong conds in second pass if using Loras. + * handle exception when parsing generation parameters from png info + * fix upcast attention dtype error + * forcing Torch Version to 1.13.1 for RX 5000 series GPUs + * split mask blur into X and Y components, patch Outpainting MK2 accordingly + * don't die when a LoRA is a broken symlink + * allow activation of Generate Forever during generation + + +## 1.3.2 + +### Bug Fixes: + * fix files served out of tmp directory even if they are saved to disk + * fix postprocessing overwriting parameters + +## 1.3.1 + +### Features: + * revert default cross attention optimization to Doggettx + +### Bug Fixes: + * fix bug: LoRA don't apply on dropdown list sd_lora + * fix png info always added even if setting is not enabled + * fix some fields not applying in xyz plot + * fix "hires. fix" prompt sharing same labels with txt2img_prompt + * fix lora hashes not being added properly to infotex if there is only one lora + * fix --use-cpu failing to work properly at startup + * make --disable-opt-split-attention command line option work again + +## 1.3.0 + +### Features: + * add UI to edit defaults + * token merging (via dbolya/tomesd) + * settings tab rework: add a lot of additional explanations and links + * load extensions' Git metadata in parallel to loading the main program to save a ton of time during startup + * update extensions table: show branch, show date in separate column, and show version from tags if available + * TAESD - another option for cheap live previews + * allow choosing sampler and prompts for second pass of hires fix - hidden by default, enabled in settings + * calculate hashes for Lora + * add lora hashes to infotext + * when pasting infotext, use infotext's lora hashes to find local loras for `` entries whose hashes match loras the user has + * select cross attention optimization from UI + +### Minor: + * bump Gradio to 3.31.0 + * bump PyTorch to 2.0.1 for macOS and Linux AMD + * allow setting defaults for elements in extensions' tabs + * allow selecting file type for live previews + * show "Loading..." for extra networks when displaying for the first time + * suppress ENSD infotext for samplers that don't use it + * clientside optimizations + * add options to show/hide hidden files and dirs in extra networks, and to not list models/files in hidden directories + * allow whitespace in styles.csv + * add option to reorder tabs + * move some functionality (swap resolution and set seed to -1) to client + * option to specify editor height for img2img + * button to copy image resolution into img2img width/height sliders + * switch from pyngrok to ngrok-py + * lazy-load images in extra networks UI + * set "Navigate image viewer with gamepad" option to false by default, by request + * change upscalers to download models into user-specified directory (from commandline args) rather than the default models/<...> + * allow hiding buttons in ui-config.json + +### Extensions: + * add /sdapi/v1/script-info api + * use Ruff to lint Python code + * use ESlint to lint Javascript code + * add/modify CFG callbacks for Self-Attention Guidance extension + * add command and endpoint for graceful server stopping + * add some locals (prompts/seeds/etc) from processing function into the Processing class as fields + * rework quoting for infotext items that have commas in them to use JSON (should be backwards compatible except for cases where it didn't work previously) + * add /sdapi/v1/refresh-loras api checkpoint post request + * tests overhaul + +### Bug Fixes: + * fix an issue preventing the program from starting if the user specifies a bad Gradio theme + * fix broken prompts from file script + * fix symlink scanning for extra networks + * fix --data-dir ignored when launching via webui-user.bat COMMANDLINE_ARGS + * allow web UI to be ran fully offline + * fix inability to run with --freeze-settings + * fix inability to merge checkpoint without adding metadata + * fix extra networks' save preview image not adding infotext for jpeg/webm + * remove blinking effect from text in hires fix and scale resolution preview + * make links to `http://<...>.git` extensions work in the extension tab + * fix bug with webui hanging at startup due to hanging git process + + +## 1.2.1 + +### Features: + * add an option to always refer to LoRA by filenames + +### Bug Fixes: + * never refer to LoRA by an alias if multiple LoRAs have same alias or the alias is called none + * fix upscalers disappearing after the user reloads UI + * allow bf16 in safe unpickler (resolves problems with loading some LoRAs) + * allow web UI to be ran fully offline + * fix localizations not working + * fix error for LoRAs: `'LatentDiffusion' object has no attribute 'lora_layer_mapping'` + +## 1.2.0 + +### Features: + * do not wait for Stable Diffusion model to load at startup + * add filename patterns: `[denoising]` + * directory hiding for extra networks: dirs starting with `.` will hide their cards on extra network tabs unless specifically searched for + * LoRA: for the `<...>` text in prompt, use name of LoRA that is in the metadata of the file, if present, instead of filename (both can be used to activate LoRA) + * LoRA: read infotext params from kohya-ss's extension parameters if they are present and if his extension is not active + * LoRA: fix some LoRAs not working (ones that have 3x3 convolution layer) + * LoRA: add an option to use old method of applying LoRAs (producing same results as with kohya-ss) + * add version to infotext, footer and console output when starting + * add links to wiki for filename pattern settings + * add extended info for quicksettings setting and use multiselect input instead of a text field + +### Minor: + * bump Gradio to 3.29.0 + * bump PyTorch to 2.0.1 + * `--subpath` option for gradio for use with reverse proxy + * Linux/macOS: use existing virtualenv if already active (the VIRTUAL_ENV environment variable) + * do not apply localizations if there are none (possible frontend optimization) + * add extra `None` option for VAE in XYZ plot + * print error to console when batch processing in img2img fails + * create HTML for extra network pages only on demand + * allow directories starting with `.` to still list their models for LoRA, checkpoints, etc + * put infotext options into their own category in settings tab + * do not show licenses page when user selects Show all pages in settings + +### Extensions: + * tooltip localization support + * add API method to get LoRA models with prompt + +### Bug Fixes: + * re-add `/docs` endpoint + * fix gamepad navigation + * make the lightbox fullscreen image function properly + * fix squished thumbnails in extras tab + * keep "search" filter for extra networks when user refreshes the tab (previously it showed everything after you refreshed) + * fix webui showing the same image if you configure the generation to always save results into same file + * fix bug with upscalers not working properly + * fix MPS on PyTorch 2.0.1, Intel Macs + * make it so that custom context menu from contextMenu.js only disappears after user's click, ignoring non-user click events + * prevent Reload UI button/link from reloading the page when it's not yet ready + * fix prompts from file script failing to read contents from a drag/drop file + + +## 1.1.1 +### Bug Fixes: + * fix an error that prevents running webui on PyTorch<2.0 without --disable-safe-unpickle + +## 1.1.0 +### Features: + * switch to PyTorch 2.0.0 (except for AMD GPUs) + * visual improvements to custom code scripts + * add filename patterns: `[clip_skip]`, `[hasprompt<>]`, `[batch_number]`, `[generation_number]` + * add support for saving init images in img2img, and record their hashes in infotext for reproducibility + * automatically select current word when adjusting weight with ctrl+up/down + * add dropdowns for X/Y/Z plot + * add setting: Stable Diffusion/Random number generator source: makes it possible to make images generated from a given manual seed consistent across different GPUs + * support Gradio's theme API + * use TCMalloc on Linux by default; possible fix for memory leaks + * add optimization option to remove negative conditioning at low sigma values #9177 + * embed model merge metadata in .safetensors file + * extension settings backup/restore feature #9169 + * add "resize by" and "resize to" tabs to img2img + * add option "keep original size" to textual inversion images preprocess + * image viewer scrolling via analog stick + * button to restore the progress from session lost / tab reload + +### Minor: + * bump Gradio to 3.28.1 + * change "scale to" to sliders in Extras tab + * add labels to tool buttons to make it possible to hide them + * add tiled inference support for ScuNET + * add branch support for extension installation + * change Linux installation script to install into current directory rather than `/home/username` + * sort textual inversion embeddings by name (case-insensitive) + * allow styles.csv to be symlinked or mounted in docker + * remove the "do not add watermark to images" option + * make selected tab configurable with UI config + * make the extra networks UI fixed height and scrollable + * add `disable_tls_verify` arg for use with self-signed certs + +### Extensions: + * add reload callback + * add `is_hr_pass` field for processing + +### Bug Fixes: + * fix broken batch image processing on 'Extras/Batch Process' tab + * add "None" option to extra networks dropdowns + * fix FileExistsError for CLIP Interrogator + * fix /sdapi/v1/txt2img endpoint not working on Linux #9319 + * fix disappearing live previews and progressbar during slow tasks + * fix fullscreen image view not working properly in some cases + * prevent alwayson_scripts args param resizing script_arg list when they are inserted in it + * fix prompt schedule for second order samplers + * fix image mask/composite for weird resolutions #9628 + * use correct images for previews when using AND (see #9491) + * one broken image in img2img batch won't stop all processing + * fix image orientation bug in train/preprocess + * fix Ngrok recreating tunnels every reload + * fix `--realesrgan-models-path` and `--ldsr-models-path` not working + * fix `--skip-install` not working + * use SAMPLE file format in Outpainting Mk2 & Poorman + * do not fail all LoRAs if some have failed to load when making a picture + +## 1.0.0 + * everything diff --git a/stable-diffusion-webui/CITATION.cff b/stable-diffusion-webui/CITATION.cff new file mode 100755 index 0000000..2c781af --- /dev/null +++ b/stable-diffusion-webui/CITATION.cff @@ -0,0 +1,7 @@ +cff-version: 1.2.0 +message: "If you use this software, please cite it as below." +authors: + - given-names: AUTOMATIC1111 +title: "Stable Diffusion Web UI" +date-released: 2022-08-22 +url: "https://github.com/AUTOMATIC1111/stable-diffusion-webui" diff --git a/stable-diffusion-webui/CODEOWNERS b/stable-diffusion-webui/CODEOWNERS new file mode 100755 index 0000000..2c937f6 --- /dev/null +++ b/stable-diffusion-webui/CODEOWNERS @@ -0,0 +1,12 @@ +* @AUTOMATIC1111 + +# if you were managing a localization and were removed from this file, this is because +# the intended way to do localizations now is via extensions. See: +# https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Developing-extensions +# Make a repo with your localization and since you are still listed as a collaborator +# you can add it to the wiki page yourself. This change is because some people complained +# the git commit log is cluttered with things unrelated to almost everyone and +# because I believe this is the best overall for the project to handle localizations almost +# entirely without my oversight. + + diff --git a/stable-diffusion-webui/LICENSE.txt b/stable-diffusion-webui/LICENSE.txt new file mode 100755 index 0000000..211d32e --- /dev/null +++ b/stable-diffusion-webui/LICENSE.txt @@ -0,0 +1,663 @@ + GNU AFFERO GENERAL PUBLIC LICENSE + Version 3, 19 November 2007 + + Copyright (c) 2023 AUTOMATIC1111 + + Copyright (C) 2007 Free Software Foundation, Inc. + Everyone is permitted to copy and distribute verbatim copies + of this license document, but changing it is not allowed. + + Preamble + + The GNU Affero General Public License is a free, copyleft license for +software and other kinds of works, specifically designed to ensure +cooperation with the community in the case of network server software. + + The licenses for most software and other practical works are designed +to take away your freedom to share and change the works. By contrast, +our General Public Licenses are intended to guarantee your freedom to +share and change all versions of a program--to make sure it remains free +software for all its users. + + When we speak of free software, we are referring to freedom, not +price. Our General Public Licenses are designed to make sure that you +have the freedom to distribute copies of free software (and charge for +them if you wish), that you receive source code or can get it if you +want it, that you can change the software or use pieces of it in new +free programs, and that you know you can do these things. + + Developers that use our General Public Licenses protect your rights +with two steps: (1) assert copyright on the software, and (2) offer +you this License which gives you legal permission to copy, distribute +and/or modify the software. + + A secondary benefit of defending all users' freedom is that +improvements made in alternate versions of the program, if they +receive widespread use, become available for other developers to +incorporate. Many developers of free software are heartened and +encouraged by the resulting cooperation. However, in the case of +software used on network servers, this result may fail to come about. +The GNU General Public License permits making a modified version and +letting the public access it on a server without ever releasing its +source code to the public. + + The GNU Affero General Public License is designed specifically to +ensure that, in such cases, the modified source code becomes available +to the community. It requires the operator of a network server to +provide the source code of the modified version running there to the +users of that server. Therefore, public use of a modified version, on +a publicly accessible server, gives the public access to the source +code of the modified version. + + An older license, called the Affero General Public License and +published by Affero, was designed to accomplish similar goals. This is +a different license, not a version of the Affero GPL, but Affero has +released a new version of the Affero GPL which permits relicensing under +this license. + + The precise terms and conditions for copying, distribution and +modification follow. + + TERMS AND CONDITIONS + + 0. Definitions. + + "This License" refers to version 3 of the GNU Affero General Public License. + + "Copyright" also means copyright-like laws that apply to other kinds of +works, such as semiconductor masks. + + "The Program" refers to any copyrightable work licensed under this +License. Each licensee is addressed as "you". "Licensees" and +"recipients" may be individuals or organizations. + + To "modify" a work means to copy from or adapt all or part of the work +in a fashion requiring copyright permission, other than the making of an +exact copy. The resulting work is called a "modified version" of the +earlier work or a work "based on" the earlier work. + + A "covered work" means either the unmodified Program or a work based +on the Program. + + To "propagate" a work means to do anything with it that, without +permission, would make you directly or secondarily liable for +infringement under applicable copyright law, except executing it on a +computer or modifying a private copy. Propagation includes copying, +distribution (with or without modification), making available to the +public, and in some countries other activities as well. + + To "convey" a work means any kind of propagation that enables other +parties to make or receive copies. Mere interaction with a user through +a computer network, with no transfer of a copy, is not conveying. + + An interactive user interface displays "Appropriate Legal Notices" +to the extent that it includes a convenient and prominently visible +feature that (1) displays an appropriate copyright notice, and (2) +tells the user that there is no warranty for the work (except to the +extent that warranties are provided), that licensees may convey the +work under this License, and how to view a copy of this License. If +the interface presents a list of user commands or options, such as a +menu, a prominent item in the list meets this criterion. + + 1. Source Code. + + The "source code" for a work means the preferred form of the work +for making modifications to it. "Object code" means any non-source +form of a work. + + A "Standard Interface" means an interface that either is an official +standard defined by a recognized standards body, or, in the case of +interfaces specified for a particular programming language, one that +is widely used among developers working in that language. + + The "System Libraries" of an executable work include anything, other +than the work as a whole, that (a) is included in the normal form of +packaging a Major Component, but which is not part of that Major +Component, and (b) serves only to enable use of the work with that +Major Component, or to implement a Standard Interface for which an +implementation is available to the public in source code form. A +"Major Component", in this context, means a major essential component +(kernel, window system, and so on) of the specific operating system +(if any) on which the executable work runs, or a compiler used to +produce the work, or an object code interpreter used to run it. + + The "Corresponding Source" for a work in object code form means all +the source code needed to generate, install, and (for an executable +work) run the object code and to modify the work, including scripts to +control those activities. However, it does not include the work's +System Libraries, or general-purpose tools or generally available free +programs which are used unmodified in performing those activities but +which are not part of the work. For example, Corresponding Source +includes interface definition files associated with source files for +the work, and the source code for shared libraries and dynamically +linked subprograms that the work is specifically designed to require, +such as by intimate data communication or control flow between those +subprograms and other parts of the work. + + The Corresponding Source need not include anything that users +can regenerate automatically from other parts of the Corresponding +Source. + + The Corresponding Source for a work in source code form is that +same work. + + 2. Basic Permissions. + + All rights granted under this License are granted for the term of +copyright on the Program, and are irrevocable provided the stated +conditions are met. This License explicitly affirms your unlimited +permission to run the unmodified Program. The output from running a +covered work is covered by this License only if the output, given its +content, constitutes a covered work. This License acknowledges your +rights of fair use or other equivalent, as provided by copyright law. + + You may make, run and propagate covered works that you do not +convey, without conditions so long as your license otherwise remains +in force. You may convey covered works to others for the sole purpose +of having them make modifications exclusively for you, or provide you +with facilities for running those works, provided that you comply with +the terms of this License in conveying all material for which you do +not control copyright. Those thus making or running the covered works +for you must do so exclusively on your behalf, under your direction +and control, on terms that prohibit them from making any copies of +your copyrighted material outside their relationship with you. + + Conveying under any other circumstances is permitted solely under +the conditions stated below. Sublicensing is not allowed; section 10 +makes it unnecessary. + + 3. Protecting Users' Legal Rights From Anti-Circumvention Law. + + No covered work shall be deemed part of an effective technological +measure under any applicable law fulfilling obligations under article +11 of the WIPO copyright treaty adopted on 20 December 1996, or +similar laws prohibiting or restricting circumvention of such +measures. + + When you convey a covered work, you waive any legal power to forbid +circumvention of technological measures to the extent such circumvention +is effected by exercising rights under this License with respect to +the covered work, and you disclaim any intention to limit operation or +modification of the work as a means of enforcing, against the work's +users, your or third parties' legal rights to forbid circumvention of +technological measures. + + 4. Conveying Verbatim Copies. + + You may convey verbatim copies of the Program's source code as you +receive it, in any medium, provided that you conspicuously and +appropriately publish on each copy an appropriate copyright notice; +keep intact all notices stating that this License and any +non-permissive terms added in accord with section 7 apply to the code; +keep intact all notices of the absence of any warranty; and give all +recipients a copy of this License along with the Program. + + You may charge any price or no price for each copy that you convey, +and you may offer support or warranty protection for a fee. + + 5. Conveying Modified Source Versions. + + You may convey a work based on the Program, or the modifications to +produce it from the Program, in the form of source code under the +terms of section 4, provided that you also meet all of these conditions: + + a) The work must carry prominent notices stating that you modified + it, and giving a relevant date. + + b) The work must carry prominent notices stating that it is + released under this License and any conditions added under section + 7. This requirement modifies the requirement in section 4 to + "keep intact all notices". + + c) You must license the entire work, as a whole, under this + License to anyone who comes into possession of a copy. This + License will therefore apply, along with any applicable section 7 + additional terms, to the whole of the work, and all its parts, + regardless of how they are packaged. This License gives no + permission to license the work in any other way, but it does not + invalidate such permission if you have separately received it. + + d) If the work has interactive user interfaces, each must display + Appropriate Legal Notices; however, if the Program has interactive + interfaces that do not display Appropriate Legal Notices, your + work need not make them do so. + + A compilation of a covered work with other separate and independent +works, which are not by their nature extensions of the covered work, +and which are not combined with it such as to form a larger program, +in or on a volume of a storage or distribution medium, is called an +"aggregate" if the compilation and its resulting copyright are not +used to limit the access or legal rights of the compilation's users +beyond what the individual works permit. Inclusion of a covered work +in an aggregate does not cause this License to apply to the other +parts of the aggregate. + + 6. Conveying Non-Source Forms. + + You may convey a covered work in object code form under the terms +of sections 4 and 5, provided that you also convey the +machine-readable Corresponding Source under the terms of this License, +in one of these ways: + + a) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by the + Corresponding Source fixed on a durable physical medium + customarily used for software interchange. + + b) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by a + written offer, valid for at least three years and valid for as + long as you offer spare parts or customer support for that product + model, to give anyone who possesses the object code either (1) a + copy of the Corresponding Source for all the software in the + product that is covered by this License, on a durable physical + medium customarily used for software interchange, for a price no + more than your reasonable cost of physically performing this + conveying of source, or (2) access to copy the + Corresponding Source from a network server at no charge. + + c) Convey individual copies of the object code with a copy of the + written offer to provide the Corresponding Source. This + alternative is allowed only occasionally and noncommercially, and + only if you received the object code with such an offer, in accord + with subsection 6b. + + d) Convey the object code by offering access from a designated + place (gratis or for a charge), and offer equivalent access to the + Corresponding Source in the same way through the same place at no + further charge. You need not require recipients to copy the + Corresponding Source along with the object code. If the place to + copy the object code is a network server, the Corresponding Source + may be on a different server (operated by you or a third party) + that supports equivalent copying facilities, provided you maintain + clear directions next to the object code saying where to find the + Corresponding Source. Regardless of what server hosts the + Corresponding Source, you remain obligated to ensure that it is + available for as long as needed to satisfy these requirements. + + e) Convey the object code using peer-to-peer transmission, provided + you inform other peers where the object code and Corresponding + Source of the work are being offered to the general public at no + charge under subsection 6d. + + A separable portion of the object code, whose source code is excluded +from the Corresponding Source as a System Library, need not be +included in conveying the object code work. + + A "User Product" is either (1) a "consumer product", which means any +tangible personal property which is normally used for personal, family, +or household purposes, or (2) anything designed or sold for incorporation +into a dwelling. In determining whether a product is a consumer product, +doubtful cases shall be resolved in favor of coverage. For a particular +product received by a particular user, "normally used" refers to a +typical or common use of that class of product, regardless of the status +of the particular user or of the way in which the particular user +actually uses, or expects or is expected to use, the product. A product +is a consumer product regardless of whether the product has substantial +commercial, industrial or non-consumer uses, unless such uses represent +the only significant mode of use of the product. + + "Installation Information" for a User Product means any methods, +procedures, authorization keys, or other information required to install +and execute modified versions of a covered work in that User Product from +a modified version of its Corresponding Source. The information must +suffice to ensure that the continued functioning of the modified object +code is in no case prevented or interfered with solely because +modification has been made. + + If you convey an object code work under this section in, or with, or +specifically for use in, a User Product, and the conveying occurs as +part of a transaction in which the right of possession and use of the +User Product is transferred to the recipient in perpetuity or for a +fixed term (regardless of how the transaction is characterized), the +Corresponding Source conveyed under this section must be accompanied +by the Installation Information. But this requirement does not apply +if neither you nor any third party retains the ability to install +modified object code on the User Product (for example, the work has +been installed in ROM). + + The requirement to provide Installation Information does not include a +requirement to continue to provide support service, warranty, or updates +for a work that has been modified or installed by the recipient, or for +the User Product in which it has been modified or installed. Access to a +network may be denied when the modification itself materially and +adversely affects the operation of the network or violates the rules and +protocols for communication across the network. + + Corresponding Source conveyed, and Installation Information provided, +in accord with this section must be in a format that is publicly +documented (and with an implementation available to the public in +source code form), and must require no special password or key for +unpacking, reading or copying. + + 7. Additional Terms. + + "Additional permissions" are terms that supplement the terms of this +License by making exceptions from one or more of its conditions. +Additional permissions that are applicable to the entire Program shall +be treated as though they were included in this License, to the extent +that they are valid under applicable law. If additional permissions +apply only to part of the Program, that part may be used separately +under those permissions, but the entire Program remains governed by +this License without regard to the additional permissions. + + When you convey a copy of a covered work, you may at your option +remove any additional permissions from that copy, or from any part of +it. (Additional permissions may be written to require their own +removal in certain cases when you modify the work.) You may place +additional permissions on material, added by you to a covered work, +for which you have or can give appropriate copyright permission. + + Notwithstanding any other provision of this License, for material you +add to a covered work, you may (if authorized by the copyright holders of +that material) supplement the terms of this License with terms: + + a) Disclaiming warranty or limiting liability differently from the + terms of sections 15 and 16 of this License; or + + b) Requiring preservation of specified reasonable legal notices or + author attributions in that material or in the Appropriate Legal + Notices displayed by works containing it; or + + c) Prohibiting misrepresentation of the origin of that material, or + requiring that modified versions of such material be marked in + reasonable ways as different from the original version; or + + d) Limiting the use for publicity purposes of names of licensors or + authors of the material; or + + e) Declining to grant rights under trademark law for use of some + trade names, trademarks, or service marks; or + + f) Requiring indemnification of licensors and authors of that + material by anyone who conveys the material (or modified versions of + it) with contractual assumptions of liability to the recipient, for + any liability that these contractual assumptions directly impose on + those licensors and authors. + + All other non-permissive additional terms are considered "further +restrictions" within the meaning of section 10. If the Program as you +received it, or any part of it, contains a notice stating that it is +governed by this License along with a term that is a further +restriction, you may remove that term. If a license document contains +a further restriction but permits relicensing or conveying under this +License, you may add to a covered work material governed by the terms +of that license document, provided that the further restriction does +not survive such relicensing or conveying. + + If you add terms to a covered work in accord with this section, you +must place, in the relevant source files, a statement of the +additional terms that apply to those files, or a notice indicating +where to find the applicable terms. + + Additional terms, permissive or non-permissive, may be stated in the +form of a separately written license, or stated as exceptions; +the above requirements apply either way. + + 8. Termination. + + You may not propagate or modify a covered work except as expressly +provided under this License. Any attempt otherwise to propagate or +modify it is void, and will automatically terminate your rights under +this License (including any patent licenses granted under the third +paragraph of section 11). + + However, if you cease all violation of this License, then your +license from a particular copyright holder is reinstated (a) +provisionally, unless and until the copyright holder explicitly and +finally terminates your license, and (b) permanently, if the copyright +holder fails to notify you of the violation by some reasonable means +prior to 60 days after the cessation. + + Moreover, your license from a particular copyright holder is +reinstated permanently if the copyright holder notifies you of the +violation by some reasonable means, this is the first time you have +received notice of violation of this License (for any work) from that +copyright holder, and you cure the violation prior to 30 days after +your receipt of the notice. + + Termination of your rights under this section does not terminate the +licenses of parties who have received copies or rights from you under +this License. If your rights have been terminated and not permanently +reinstated, you do not qualify to receive new licenses for the same +material under section 10. + + 9. Acceptance Not Required for Having Copies. + + You are not required to accept this License in order to receive or +run a copy of the Program. Ancillary propagation of a covered work +occurring solely as a consequence of using peer-to-peer transmission +to receive a copy likewise does not require acceptance. However, +nothing other than this License grants you permission to propagate or +modify any covered work. These actions infringe copyright if you do +not accept this License. Therefore, by modifying or propagating a +covered work, you indicate your acceptance of this License to do so. + + 10. Automatic Licensing of Downstream Recipients. + + Each time you convey a covered work, the recipient automatically +receives a license from the original licensors, to run, modify and +propagate that work, subject to this License. You are not responsible +for enforcing compliance by third parties with this License. + + An "entity transaction" is a transaction transferring control of an +organization, or substantially all assets of one, or subdividing an +organization, or merging organizations. If propagation of a covered +work results from an entity transaction, each party to that +transaction who receives a copy of the work also receives whatever +licenses to the work the party's predecessor in interest had or could +give under the previous paragraph, plus a right to possession of the +Corresponding Source of the work from the predecessor in interest, if +the predecessor has it or can get it with reasonable efforts. + + You may not impose any further restrictions on the exercise of the +rights granted or affirmed under this License. For example, you may +not impose a license fee, royalty, or other charge for exercise of +rights granted under this License, and you may not initiate litigation +(including a cross-claim or counterclaim in a lawsuit) alleging that +any patent claim is infringed by making, using, selling, offering for +sale, or importing the Program or any portion of it. + + 11. Patents. + + A "contributor" is a copyright holder who authorizes use under this +License of the Program or a work on which the Program is based. The +work thus licensed is called the contributor's "contributor version". + + A contributor's "essential patent claims" are all patent claims +owned or controlled by the contributor, whether already acquired or +hereafter acquired, that would be infringed by some manner, permitted +by this License, of making, using, or selling its contributor version, +but do not include claims that would be infringed only as a +consequence of further modification of the contributor version. For +purposes of this definition, "control" includes the right to grant +patent sublicenses in a manner consistent with the requirements of +this License. + + Each contributor grants you a non-exclusive, worldwide, royalty-free +patent license under the contributor's essential patent claims, to +make, use, sell, offer for sale, import and otherwise run, modify and +propagate the contents of its contributor version. + + In the following three paragraphs, a "patent license" is any express +agreement or commitment, however denominated, not to enforce a patent +(such as an express permission to practice a patent or covenant not to +sue for patent infringement). To "grant" such a patent license to a +party means to make such an agreement or commitment not to enforce a +patent against the party. + + If you convey a covered work, knowingly relying on a patent license, +and the Corresponding Source of the work is not available for anyone +to copy, free of charge and under the terms of this License, through a +publicly available network server or other readily accessible means, +then you must either (1) cause the Corresponding Source to be so +available, or (2) arrange to deprive yourself of the benefit of the +patent license for this particular work, or (3) arrange, in a manner +consistent with the requirements of this License, to extend the patent +license to downstream recipients. "Knowingly relying" means you have +actual knowledge that, but for the patent license, your conveying the +covered work in a country, or your recipient's use of the covered work +in a country, would infringe one or more identifiable patents in that +country that you have reason to believe are valid. + + If, pursuant to or in connection with a single transaction or +arrangement, you convey, or propagate by procuring conveyance of, a +covered work, and grant a patent license to some of the parties +receiving the covered work authorizing them to use, propagate, modify +or convey a specific copy of the covered work, then the patent license +you grant is automatically extended to all recipients of the covered +work and works based on it. + + A patent license is "discriminatory" if it does not include within +the scope of its coverage, prohibits the exercise of, or is +conditioned on the non-exercise of one or more of the rights that are +specifically granted under this License. You may not convey a covered +work if you are a party to an arrangement with a third party that is +in the business of distributing software, under which you make payment +to the third party based on the extent of your activity of conveying +the work, and under which the third party grants, to any of the +parties who would receive the covered work from you, a discriminatory +patent license (a) in connection with copies of the covered work +conveyed by you (or copies made from those copies), or (b) primarily +for and in connection with specific products or compilations that +contain the covered work, unless you entered into that arrangement, +or that patent license was granted, prior to 28 March 2007. + + Nothing in this License shall be construed as excluding or limiting +any implied license or other defenses to infringement that may +otherwise be available to you under applicable patent law. + + 12. No Surrender of Others' Freedom. + + If conditions are imposed on you (whether by court order, agreement or +otherwise) that contradict the conditions of this License, they do not +excuse you from the conditions of this License. If you cannot convey a +covered work so as to satisfy simultaneously your obligations under this +License and any other pertinent obligations, then as a consequence you may +not convey it at all. For example, if you agree to terms that obligate you +to collect a royalty for further conveying from those to whom you convey +the Program, the only way you could satisfy both those terms and this +License would be to refrain entirely from conveying the Program. + + 13. Remote Network Interaction; Use with the GNU General Public License. + + Notwithstanding any other provision of this License, if you modify the +Program, your modified version must prominently offer all users +interacting with it remotely through a computer network (if your version +supports such interaction) an opportunity to receive the Corresponding +Source of your version by providing access to the Corresponding Source +from a network server at no charge, through some standard or customary +means of facilitating copying of software. This Corresponding Source +shall include the Corresponding Source for any work covered by version 3 +of the GNU General Public License that is incorporated pursuant to the +following paragraph. + + Notwithstanding any other provision of this License, you have +permission to link or combine any covered work with a work licensed +under version 3 of the GNU General Public License into a single +combined work, and to convey the resulting work. The terms of this +License will continue to apply to the part which is the covered work, +but the work with which it is combined will remain governed by version +3 of the GNU General Public License. + + 14. Revised Versions of this License. + + The Free Software Foundation may publish revised and/or new versions of +the GNU Affero General Public License from time to time. Such new versions +will be similar in spirit to the present version, but may differ in detail to +address new problems or concerns. + + Each version is given a distinguishing version number. If the +Program specifies that a certain numbered version of the GNU Affero General +Public License "or any later version" applies to it, you have the +option of following the terms and conditions either of that numbered +version or of any later version published by the Free Software +Foundation. If the Program does not specify a version number of the +GNU Affero General Public License, you may choose any version ever published +by the Free Software Foundation. + + If the Program specifies that a proxy can decide which future +versions of the GNU Affero General Public License can be used, that proxy's +public statement of acceptance of a version permanently authorizes you +to choose that version for the Program. + + Later license versions may give you additional or different +permissions. However, no additional obligations are imposed on any +author or copyright holder as a result of your choosing to follow a +later version. + + 15. Disclaimer of Warranty. + + THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY +APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT +HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY +OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, +THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM +IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF +ALL NECESSARY SERVICING, REPAIR OR CORRECTION. + + 16. Limitation of Liability. + + IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING +WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS +THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY +GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE +USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF +DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD +PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), +EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF +SUCH DAMAGES. + + 17. Interpretation of Sections 15 and 16. + + If the disclaimer of warranty and limitation of liability provided +above cannot be given local legal effect according to their terms, +reviewing courts shall apply local law that most closely approximates +an absolute waiver of all civil liability in connection with the +Program, unless a warranty or assumption of liability accompanies a +copy of the Program in return for a fee. + + END OF TERMS AND CONDITIONS + + How to Apply These Terms to Your New Programs + + If you develop a new program, and you want it to be of the greatest +possible use to the public, the best way to achieve this is to make it +free software which everyone can redistribute and change under these terms. + + To do so, attach the following notices to the program. It is safest +to attach them to the start of each source file to most effectively +state the exclusion of warranty; and each file should have at least +the "copyright" line and a pointer to where the full notice is found. + + + Copyright (C) + + This program is free software: you can redistribute it and/or modify + it under the terms of the GNU Affero General Public License as published by + the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU Affero General Public License for more details. + + You should have received a copy of the GNU Affero General Public License + along with this program. If not, see . + +Also add information on how to contact you by electronic and paper mail. + + If your software can interact with users remotely through a computer +network, you should also make sure that it provides a way for users to +get its source. For example, if your program is a web application, its +interface could display a "Source" link that leads users to an archive +of the code. There are many ways you could offer source, and different +solutions will be better for different programs; see section 13 for the +specific requirements. + + You should also get your employer (if you work as a programmer) or school, +if any, to sign a "copyright disclaimer" for the program, if necessary. +For more information on this, and how to apply and follow the GNU AGPL, see +. diff --git a/stable-diffusion-webui/README.md b/stable-diffusion-webui/README.md new file mode 100755 index 0000000..c630677 --- /dev/null +++ b/stable-diffusion-webui/README.md @@ -0,0 +1,205 @@ +# Stable Diffusion web UI +A web interface for Stable Diffusion, implemented using Gradio library. + +![](screenshot.png) + +## Features +[Detailed feature showcase with images](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features): +- Original txt2img and img2img modes +- One click install and run script (but you still must install python and git) +- Outpainting +- Inpainting +- Color Sketch +- Prompt Matrix +- Stable Diffusion Upscale +- Attention, specify parts of text that the model should pay more attention to + - a man in a `((tuxedo))` - will pay more attention to tuxedo + - a man in a `(tuxedo:1.21)` - alternative syntax + - select text and press `Ctrl+Up` or `Ctrl+Down` (or `Command+Up` or `Command+Down` if you're on a MacOS) to automatically adjust attention to selected text (code contributed by anonymous user) +- Loopback, run img2img processing multiple times +- X/Y/Z plot, a way to draw a 3 dimensional plot of images with different parameters +- Textual Inversion + - have as many embeddings as you want and use any names you like for them + - use multiple embeddings with different numbers of vectors per token + - works with half precision floating point numbers + - train embeddings on 8GB (also reports of 6GB working) +- Extras tab with: + - GFPGAN, neural network that fixes faces + - CodeFormer, face restoration tool as an alternative to GFPGAN + - RealESRGAN, neural network upscaler + - ESRGAN, neural network upscaler with a lot of third party models + - SwinIR and Swin2SR ([see here](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/2092)), neural network upscalers + - LDSR, Latent diffusion super resolution upscaling +- Resizing aspect ratio options +- Sampling method selection + - Adjust sampler eta values (noise multiplier) + - More advanced noise setting options +- Interrupt processing at any time +- 4GB video card support (also reports of 2GB working) +- Correct seeds for batches +- Live prompt token length validation +- Generation parameters + - parameters you used to generate images are saved with that image + - in PNG chunks for PNG, in EXIF for JPEG + - can drag the image to PNG info tab to restore generation parameters and automatically copy them into UI + - can be disabled in settings + - drag and drop an image/text-parameters to promptbox +- Read Generation Parameters Button, loads parameters in promptbox to UI +- Settings page +- Running arbitrary python code from UI (must run with `--allow-code` to enable) +- Mouseover hints for most UI elements +- Possible to change defaults/mix/max/step values for UI elements via text config +- Tiling support, a checkbox to create images that can be tiled like textures +- Progress bar and live image generation preview + - Can use a separate neural network to produce previews with almost none VRAM or compute requirement +- Negative prompt, an extra text field that allows you to list what you don't want to see in generated image +- Styles, a way to save part of prompt and easily apply them via dropdown later +- Variations, a way to generate same image but with tiny differences +- Seed resizing, a way to generate same image but at slightly different resolution +- CLIP interrogator, a button that tries to guess prompt from an image +- Prompt Editing, a way to change prompt mid-generation, say to start making a watermelon and switch to anime girl midway +- Batch Processing, process a group of files using img2img +- Img2img Alternative, reverse Euler method of cross attention control +- Highres Fix, a convenience option to produce high resolution pictures in one click without usual distortions +- Reloading checkpoints on the fly +- Checkpoint Merger, a tab that allows you to merge up to 3 checkpoints into one +- [Custom scripts](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Custom-Scripts) with many extensions from community +- [Composable-Diffusion](https://energy-based-model.github.io/Compositional-Visual-Generation-with-Composable-Diffusion-Models/), a way to use multiple prompts at once + - separate prompts using uppercase `AND` + - also supports weights for prompts: `a cat :1.2 AND a dog AND a penguin :2.2` +- No token limit for prompts (original stable diffusion lets you use up to 75 tokens) +- DeepDanbooru integration, creates danbooru style tags for anime prompts +- [xformers](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Xformers), major speed increase for select cards: (add `--xformers` to commandline args) +- via extension: [History tab](https://github.com/yfszzx/stable-diffusion-webui-images-browser): view, direct and delete images conveniently within the UI +- Generate forever option +- Training tab + - hypernetworks and embeddings options + - Preprocessing images: cropping, mirroring, autotagging using BLIP or deepdanbooru (for anime) +- Clip skip +- Hypernetworks +- Loras (same as Hypernetworks but more pretty) +- A separate UI where you can choose, with preview, which embeddings, hypernetworks or Loras to add to your prompt +- Can select to load a different VAE from settings screen +- Estimated completion time in progress bar +- API +- Support for dedicated [inpainting model](https://github.com/runwayml/stable-diffusion#inpainting-with-stable-diffusion) by RunwayML +- via extension: [Aesthetic Gradients](https://github.com/AUTOMATIC1111/stable-diffusion-webui-aesthetic-gradients), a way to generate images with a specific aesthetic by using clip images embeds (implementation of [https://github.com/vicgalle/stable-diffusion-aesthetic-gradients](https://github.com/vicgalle/stable-diffusion-aesthetic-gradients)) +- [Stable Diffusion 2.0](https://github.com/Stability-AI/stablediffusion) support - see [wiki](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#stable-diffusion-20) for instructions +- [Alt-Diffusion](https://arxiv.org/abs/2211.06679) support - see [wiki](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#alt-diffusion) for instructions +- Now without any bad letters! +- Load checkpoints in safetensors format +- Eased resolution restriction: generated image's dimensions must be a multiple of 8 rather than 64 +- Now with a license! +- Reorder elements in the UI from settings screen +- [Segmind Stable Diffusion](https://huggingface.co/segmind/SSD-1B) support + +## Installation and Running +Make sure the required [dependencies](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Dependencies) are met and follow the instructions available for: +- [NVidia](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Install-and-Run-on-NVidia-GPUs) (recommended) +- [AMD](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Install-and-Run-on-AMD-GPUs) GPUs. +- [Intel CPUs, Intel GPUs (both integrated and discrete)](https://github.com/openvinotoolkit/stable-diffusion-webui/wiki/Installation-on-Intel-Silicon) (external wiki page) +- [Ascend NPUs](https://github.com/wangshuai09/stable-diffusion-webui/wiki/Install-and-run-on-Ascend-NPUs) (external wiki page) + +Alternatively, use online services (like Google Colab): + +- [List of Online Services](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Online-Services) + +### Installation on Windows 10/11 with NVidia-GPUs using release package +1. Download `sd.webui.zip` from [v1.0.0-pre](https://github.com/AUTOMATIC1111/stable-diffusion-webui/releases/tag/v1.0.0-pre) and extract its contents. +2. Run `update.bat`. +3. Run `run.bat`. +> For more details see [Install-and-Run-on-NVidia-GPUs](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Install-and-Run-on-NVidia-GPUs) + +### Automatic Installation on Windows +1. Install [Python 3.10.6](https://www.python.org/downloads/release/python-3106/) (Newer version of Python does not support torch), checking "Add Python to PATH". +2. Install [git](https://git-scm.com/download/win). +3. Download the stable-diffusion-webui repository, for example by running `git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git`. +4. Run `webui-user.bat` from Windows Explorer as normal, non-administrator, user. + +### Automatic Installation on Linux +1. Install the dependencies: +```bash +# Debian-based: +sudo apt install wget git python3 python3-venv libgl1 libglib2.0-0 +# Red Hat-based: +sudo dnf install wget git python3 gperftools-libs libglvnd-glx +# openSUSE-based: +sudo zypper install wget git python3 libtcmalloc4 libglvnd +# Arch-based: +sudo pacman -S wget git python3 +``` +If your system is very new, you need to install python3.11 or python3.10: +```bash +# Ubuntu 24.04 +sudo add-apt-repository ppa:deadsnakes/ppa +sudo apt update +sudo apt install python3.11 + +# Manjaro/Arch +sudo pacman -S yay +yay -S python311 # do not confuse with python3.11 package + +# Only for 3.11 +# Then set up env variable in launch script +export python_cmd="python3.11" +# or in webui-user.sh +python_cmd="python3.11" +``` +2. Navigate to the directory you would like the webui to be installed and execute the following command: +```bash +wget -q https://raw.githubusercontent.com/AUTOMATIC1111/stable-diffusion-webui/master/webui.sh +``` +Or just clone the repo wherever you want: +```bash +git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui +``` + +3. Run `webui.sh`. +4. Check `webui-user.sh` for options. +### Installation on Apple Silicon + +Find the instructions [here](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Installation-on-Apple-Silicon). + +## Contributing +Here's how to add code to this repo: [Contributing](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Contributing) + +## Documentation + +The documentation was moved from this README over to the project's [wiki](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki). + +For the purposes of getting Google and other search engines to crawl the wiki, here's a link to the (not for humans) [crawlable wiki](https://github-wiki-see.page/m/AUTOMATIC1111/stable-diffusion-webui/wiki). + +## Credits +Licenses for borrowed code can be found in `Settings -> Licenses` screen, and also in `html/licenses.html` file. + +- Stable Diffusion - https://github.com/Stability-AI/stablediffusion, https://github.com/CompVis/taming-transformers, https://github.com/mcmonkey4eva/sd3-ref +- k-diffusion - https://github.com/crowsonkb/k-diffusion.git +- Spandrel - https://github.com/chaiNNer-org/spandrel implementing + - GFPGAN - https://github.com/TencentARC/GFPGAN.git + - CodeFormer - https://github.com/sczhou/CodeFormer + - ESRGAN - https://github.com/xinntao/ESRGAN + - SwinIR - https://github.com/JingyunLiang/SwinIR + - Swin2SR - https://github.com/mv-lab/swin2sr +- LDSR - https://github.com/Hafiidz/latent-diffusion +- MiDaS - https://github.com/isl-org/MiDaS +- Ideas for optimizations - https://github.com/basujindal/stable-diffusion +- Cross Attention layer optimization - Doggettx - https://github.com/Doggettx/stable-diffusion, original idea for prompt editing. +- Cross Attention layer optimization - InvokeAI, lstein - https://github.com/invoke-ai/InvokeAI (originally http://github.com/lstein/stable-diffusion) +- Sub-quadratic Cross Attention layer optimization - Alex Birch (https://github.com/Birch-san/diffusers/pull/1), Amin Rezaei (https://github.com/AminRezaei0x443/memory-efficient-attention) +- Textual Inversion - Rinon Gal - https://github.com/rinongal/textual_inversion (we're not using his code, but we are using his ideas). +- Idea for SD upscale - https://github.com/jquesnelle/txt2imghd +- Noise generation for outpainting mk2 - https://github.com/parlance-zz/g-diffuser-bot +- CLIP interrogator idea and borrowing some code - https://github.com/pharmapsychotic/clip-interrogator +- Idea for Composable Diffusion - https://github.com/energy-based-model/Compositional-Visual-Generation-with-Composable-Diffusion-Models-PyTorch +- xformers - https://github.com/facebookresearch/xformers +- DeepDanbooru - interrogator for anime diffusers https://github.com/KichangKim/DeepDanbooru +- Sampling in float32 precision from a float16 UNet - marunine for the idea, Birch-san for the example Diffusers implementation (https://github.com/Birch-san/diffusers-play/tree/92feee6) +- Instruct pix2pix - Tim Brooks (star), Aleksander Holynski (star), Alexei A. Efros (no star) - https://github.com/timothybrooks/instruct-pix2pix +- Security advice - RyotaK +- UniPC sampler - Wenliang Zhao - https://github.com/wl-zhao/UniPC +- TAESD - Ollin Boer Bohan - https://github.com/madebyollin/taesd +- LyCORIS - KohakuBlueleaf +- Restart sampling - lambertae - https://github.com/Newbeeer/diffusion_restart_sampling +- Hypertile - tfernd - https://github.com/tfernd/HyperTile +- Initial Gradio script - posted on 4chan by an Anonymous user. Thank you Anonymous user. +- (You) diff --git a/stable-diffusion-webui/_typos.toml b/stable-diffusion-webui/_typos.toml new file mode 100755 index 0000000..1c63fe7 --- /dev/null +++ b/stable-diffusion-webui/_typos.toml @@ -0,0 +1,5 @@ +[default.extend-words] +# Part of "RGBa" (Pillow's pre-multiplied alpha RGB mode) +Ba = "Ba" +# HSA is something AMD uses for their GPUs +HSA = "HSA" diff --git a/stable-diffusion-webui/configs/alt-diffusion-inference.yaml b/stable-diffusion-webui/configs/alt-diffusion-inference.yaml new file mode 100755 index 0000000..4944ab5 --- /dev/null +++ b/stable-diffusion-webui/configs/alt-diffusion-inference.yaml @@ -0,0 +1,72 @@ +model: + base_learning_rate: 1.0e-04 + target: ldm.models.diffusion.ddpm.LatentDiffusion + params: + linear_start: 0.00085 + linear_end: 0.0120 + num_timesteps_cond: 1 + log_every_t: 200 + timesteps: 1000 + first_stage_key: "jpg" + cond_stage_key: "txt" + image_size: 64 + channels: 4 + cond_stage_trainable: false # Note: different from the one we trained before + conditioning_key: crossattn + monitor: val/loss_simple_ema + scale_factor: 0.18215 + use_ema: False + + scheduler_config: # 10000 warmup steps + target: ldm.lr_scheduler.LambdaLinearScheduler + params: + warm_up_steps: [ 10000 ] + cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases + f_start: [ 1.e-6 ] + f_max: [ 1. ] + f_min: [ 1. ] + + unet_config: + target: ldm.modules.diffusionmodules.openaimodel.UNetModel + params: + image_size: 32 # unused + in_channels: 4 + out_channels: 4 + model_channels: 320 + attention_resolutions: [ 4, 2, 1 ] + num_res_blocks: 2 + channel_mult: [ 1, 2, 4, 4 ] + num_heads: 8 + use_spatial_transformer: True + transformer_depth: 1 + context_dim: 768 + use_checkpoint: False + legacy: False + + first_stage_config: + target: ldm.models.autoencoder.AutoencoderKL + params: + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + double_z: true + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: + - 1 + - 2 + - 4 + - 4 + num_res_blocks: 2 + attn_resolutions: [] + dropout: 0.0 + lossconfig: + target: torch.nn.Identity + + cond_stage_config: + target: modules.xlmr.BertSeriesModelWithTransformation + params: + name: "XLMR-Large" \ No newline at end of file diff --git a/stable-diffusion-webui/configs/alt-diffusion-m18-inference.yaml b/stable-diffusion-webui/configs/alt-diffusion-m18-inference.yaml new file mode 100755 index 0000000..c60dca8 --- /dev/null +++ b/stable-diffusion-webui/configs/alt-diffusion-m18-inference.yaml @@ -0,0 +1,73 @@ +model: + base_learning_rate: 1.0e-04 + target: ldm.models.diffusion.ddpm.LatentDiffusion + params: + linear_start: 0.00085 + linear_end: 0.0120 + num_timesteps_cond: 1 + log_every_t: 200 + timesteps: 1000 + first_stage_key: "jpg" + cond_stage_key: "txt" + image_size: 64 + channels: 4 + cond_stage_trainable: false # Note: different from the one we trained before + conditioning_key: crossattn + monitor: val/loss_simple_ema + scale_factor: 0.18215 + use_ema: False + + scheduler_config: # 10000 warmup steps + target: ldm.lr_scheduler.LambdaLinearScheduler + params: + warm_up_steps: [ 10000 ] + cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases + f_start: [ 1.e-6 ] + f_max: [ 1. ] + f_min: [ 1. ] + + unet_config: + target: ldm.modules.diffusionmodules.openaimodel.UNetModel + params: + image_size: 32 # unused + in_channels: 4 + out_channels: 4 + model_channels: 320 + attention_resolutions: [ 4, 2, 1 ] + num_res_blocks: 2 + channel_mult: [ 1, 2, 4, 4 ] + num_head_channels: 64 + use_spatial_transformer: True + use_linear_in_transformer: True + transformer_depth: 1 + context_dim: 1024 + use_checkpoint: False + legacy: False + + first_stage_config: + target: ldm.models.autoencoder.AutoencoderKL + params: + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + double_z: true + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: + - 1 + - 2 + - 4 + - 4 + num_res_blocks: 2 + attn_resolutions: [] + dropout: 0.0 + lossconfig: + target: torch.nn.Identity + + cond_stage_config: + target: modules.xlmr_m18.BertSeriesModelWithTransformation + params: + name: "XLMR-Large" diff --git a/stable-diffusion-webui/configs/instruct-pix2pix.yaml b/stable-diffusion-webui/configs/instruct-pix2pix.yaml new file mode 100755 index 0000000..564e50a --- /dev/null +++ b/stable-diffusion-webui/configs/instruct-pix2pix.yaml @@ -0,0 +1,98 @@ +# File modified by authors of InstructPix2Pix from original (https://github.com/CompVis/stable-diffusion). +# See more details in LICENSE. + +model: + base_learning_rate: 1.0e-04 + target: modules.models.diffusion.ddpm_edit.LatentDiffusion + params: + linear_start: 0.00085 + linear_end: 0.0120 + num_timesteps_cond: 1 + log_every_t: 200 + timesteps: 1000 + first_stage_key: edited + cond_stage_key: edit + # image_size: 64 + # image_size: 32 + image_size: 16 + channels: 4 + cond_stage_trainable: false # Note: different from the one we trained before + conditioning_key: hybrid + monitor: val/loss_simple_ema + scale_factor: 0.18215 + use_ema: false + + scheduler_config: # 10000 warmup steps + target: ldm.lr_scheduler.LambdaLinearScheduler + params: + warm_up_steps: [ 0 ] + cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases + f_start: [ 1.e-6 ] + f_max: [ 1. ] + f_min: [ 1. ] + + unet_config: + target: ldm.modules.diffusionmodules.openaimodel.UNetModel + params: + image_size: 32 # unused + in_channels: 8 + out_channels: 4 + model_channels: 320 + attention_resolutions: [ 4, 2, 1 ] + num_res_blocks: 2 + channel_mult: [ 1, 2, 4, 4 ] + num_heads: 8 + use_spatial_transformer: True + transformer_depth: 1 + context_dim: 768 + use_checkpoint: False + legacy: False + + first_stage_config: + target: ldm.models.autoencoder.AutoencoderKL + params: + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + double_z: true + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: + - 1 + - 2 + - 4 + - 4 + num_res_blocks: 2 + attn_resolutions: [] + dropout: 0.0 + lossconfig: + target: torch.nn.Identity + + cond_stage_config: + target: ldm.modules.encoders.modules.FrozenCLIPEmbedder + +data: + target: main.DataModuleFromConfig + params: + batch_size: 128 + num_workers: 1 + wrap: false + validation: + target: edit_dataset.EditDataset + params: + path: data/clip-filtered-dataset + cache_dir: data/ + cache_name: data_10k + split: val + min_text_sim: 0.2 + min_image_sim: 0.75 + min_direction_sim: 0.2 + max_samples_per_prompt: 1 + min_resize_res: 512 + max_resize_res: 512 + crop_res: 512 + output_as_edit: False + real_input: True diff --git a/stable-diffusion-webui/configs/sd3-inference.yaml b/stable-diffusion-webui/configs/sd3-inference.yaml new file mode 100755 index 0000000..bccb69d --- /dev/null +++ b/stable-diffusion-webui/configs/sd3-inference.yaml @@ -0,0 +1,5 @@ +model: + target: modules.models.sd3.sd3_model.SD3Inferencer + params: + shift: 3 + state_dict: null diff --git a/stable-diffusion-webui/configs/sd_xl_inpaint.yaml b/stable-diffusion-webui/configs/sd_xl_inpaint.yaml new file mode 100755 index 0000000..f40f45e --- /dev/null +++ b/stable-diffusion-webui/configs/sd_xl_inpaint.yaml @@ -0,0 +1,98 @@ +model: + target: sgm.models.diffusion.DiffusionEngine + params: + scale_factor: 0.13025 + disable_first_stage_autocast: True + + denoiser_config: + target: sgm.modules.diffusionmodules.denoiser.DiscreteDenoiser + params: + num_idx: 1000 + + weighting_config: + target: sgm.modules.diffusionmodules.denoiser_weighting.EpsWeighting + scaling_config: + target: sgm.modules.diffusionmodules.denoiser_scaling.EpsScaling + discretization_config: + target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization + + network_config: + target: sgm.modules.diffusionmodules.openaimodel.UNetModel + params: + adm_in_channels: 2816 + num_classes: sequential + use_checkpoint: False + in_channels: 9 + out_channels: 4 + model_channels: 320 + attention_resolutions: [4, 2] + num_res_blocks: 2 + channel_mult: [1, 2, 4] + num_head_channels: 64 + use_spatial_transformer: True + use_linear_in_transformer: True + transformer_depth: [1, 2, 10] # note: the first is unused (due to attn_res starting at 2) 32, 16, 8 --> 64, 32, 16 + context_dim: 2048 + spatial_transformer_attn_type: softmax-xformers + legacy: False + + conditioner_config: + target: sgm.modules.GeneralConditioner + params: + emb_models: + # crossattn cond + - is_trainable: False + input_key: txt + target: sgm.modules.encoders.modules.FrozenCLIPEmbedder + params: + layer: hidden + layer_idx: 11 + # crossattn and vector cond + - is_trainable: False + input_key: txt + target: sgm.modules.encoders.modules.FrozenOpenCLIPEmbedder2 + params: + arch: ViT-bigG-14 + version: laion2b_s39b_b160k + freeze: True + layer: penultimate + always_return_pooled: True + legacy: False + # vector cond + - is_trainable: False + input_key: original_size_as_tuple + target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND + params: + outdim: 256 # multiplied by two + # vector cond + - is_trainable: False + input_key: crop_coords_top_left + target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND + params: + outdim: 256 # multiplied by two + # vector cond + - is_trainable: False + input_key: target_size_as_tuple + target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND + params: + outdim: 256 # multiplied by two + + first_stage_config: + target: sgm.models.autoencoder.AutoencoderKLInferenceWrapper + params: + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + attn_type: vanilla-xformers + double_z: true + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: [1, 2, 4, 4] + num_res_blocks: 2 + attn_resolutions: [] + dropout: 0.0 + lossconfig: + target: torch.nn.Identity diff --git a/stable-diffusion-webui/configs/v1-inference.yaml b/stable-diffusion-webui/configs/v1-inference.yaml new file mode 100755 index 0000000..25c4d9e --- /dev/null +++ b/stable-diffusion-webui/configs/v1-inference.yaml @@ -0,0 +1,70 @@ +model: + base_learning_rate: 1.0e-04 + target: ldm.models.diffusion.ddpm.LatentDiffusion + params: + linear_start: 0.00085 + linear_end: 0.0120 + num_timesteps_cond: 1 + log_every_t: 200 + timesteps: 1000 + first_stage_key: "jpg" + cond_stage_key: "txt" + image_size: 64 + channels: 4 + cond_stage_trainable: false # Note: different from the one we trained before + conditioning_key: crossattn + monitor: val/loss_simple_ema + scale_factor: 0.18215 + use_ema: False + + scheduler_config: # 10000 warmup steps + target: ldm.lr_scheduler.LambdaLinearScheduler + params: + warm_up_steps: [ 10000 ] + cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases + f_start: [ 1.e-6 ] + f_max: [ 1. ] + f_min: [ 1. ] + + unet_config: + target: ldm.modules.diffusionmodules.openaimodel.UNetModel + params: + image_size: 32 # unused + in_channels: 4 + out_channels: 4 + model_channels: 320 + attention_resolutions: [ 4, 2, 1 ] + num_res_blocks: 2 + channel_mult: [ 1, 2, 4, 4 ] + num_heads: 8 + use_spatial_transformer: True + transformer_depth: 1 + context_dim: 768 + use_checkpoint: False + legacy: False + + first_stage_config: + target: ldm.models.autoencoder.AutoencoderKL + params: + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + double_z: true + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: + - 1 + - 2 + - 4 + - 4 + num_res_blocks: 2 + attn_resolutions: [] + dropout: 0.0 + lossconfig: + target: torch.nn.Identity + + cond_stage_config: + target: ldm.modules.encoders.modules.FrozenCLIPEmbedder diff --git a/stable-diffusion-webui/configs/v1-inpainting-inference.yaml b/stable-diffusion-webui/configs/v1-inpainting-inference.yaml new file mode 100755 index 0000000..68c199f --- /dev/null +++ b/stable-diffusion-webui/configs/v1-inpainting-inference.yaml @@ -0,0 +1,70 @@ +model: + base_learning_rate: 7.5e-05 + target: ldm.models.diffusion.ddpm.LatentInpaintDiffusion + params: + linear_start: 0.00085 + linear_end: 0.0120 + num_timesteps_cond: 1 + log_every_t: 200 + timesteps: 1000 + first_stage_key: "jpg" + cond_stage_key: "txt" + image_size: 64 + channels: 4 + cond_stage_trainable: false # Note: different from the one we trained before + conditioning_key: hybrid # important + monitor: val/loss_simple_ema + scale_factor: 0.18215 + finetune_keys: null + + scheduler_config: # 10000 warmup steps + target: ldm.lr_scheduler.LambdaLinearScheduler + params: + warm_up_steps: [ 2500 ] # NOTE for resuming. use 10000 if starting from scratch + cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases + f_start: [ 1.e-6 ] + f_max: [ 1. ] + f_min: [ 1. ] + + unet_config: + target: ldm.modules.diffusionmodules.openaimodel.UNetModel + params: + image_size: 32 # unused + in_channels: 9 # 4 data + 4 downscaled image + 1 mask + out_channels: 4 + model_channels: 320 + attention_resolutions: [ 4, 2, 1 ] + num_res_blocks: 2 + channel_mult: [ 1, 2, 4, 4 ] + num_heads: 8 + use_spatial_transformer: True + transformer_depth: 1 + context_dim: 768 + use_checkpoint: False + legacy: False + + first_stage_config: + target: ldm.models.autoencoder.AutoencoderKL + params: + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + double_z: true + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: + - 1 + - 2 + - 4 + - 4 + num_res_blocks: 2 + attn_resolutions: [] + dropout: 0.0 + lossconfig: + target: torch.nn.Identity + + cond_stage_config: + target: ldm.modules.encoders.modules.FrozenCLIPEmbedder diff --git a/stable-diffusion-webui/environment-wsl2.yaml b/stable-diffusion-webui/environment-wsl2.yaml new file mode 100755 index 0000000..0c4ae68 --- /dev/null +++ b/stable-diffusion-webui/environment-wsl2.yaml @@ -0,0 +1,11 @@ +name: automatic +channels: + - pytorch + - defaults +dependencies: + - python=3.10 + - pip=23.0 + - cudatoolkit=11.8 + - pytorch=2.0 + - torchvision=0.15 + - numpy=1.23 diff --git a/stable-diffusion-webui/extensions-builtin/LDSR/ldsr_model_arch.py b/stable-diffusion-webui/extensions-builtin/LDSR/ldsr_model_arch.py new file mode 100755 index 0000000..7cac36c --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/LDSR/ldsr_model_arch.py @@ -0,0 +1,250 @@ +import os +import gc +import time + +import numpy as np +import torch +import torchvision +from PIL import Image +from einops import rearrange, repeat +from omegaconf import OmegaConf +import safetensors.torch + +from ldm.models.diffusion.ddim import DDIMSampler +from ldm.util import instantiate_from_config, ismap +from modules import shared, sd_hijack, devices + +cached_ldsr_model: torch.nn.Module = None + + +# Create LDSR Class +class LDSR: + def load_model_from_config(self, half_attention): + global cached_ldsr_model + + if shared.opts.ldsr_cached and cached_ldsr_model is not None: + print("Loading model from cache") + model: torch.nn.Module = cached_ldsr_model + else: + print(f"Loading model from {self.modelPath}") + _, extension = os.path.splitext(self.modelPath) + if extension.lower() == ".safetensors": + pl_sd = safetensors.torch.load_file(self.modelPath, device="cpu") + else: + pl_sd = torch.load(self.modelPath, map_location="cpu") + sd = pl_sd["state_dict"] if "state_dict" in pl_sd else pl_sd + config = OmegaConf.load(self.yamlPath) + config.model.target = "ldm.models.diffusion.ddpm.LatentDiffusionV1" + model: torch.nn.Module = instantiate_from_config(config.model) + model.load_state_dict(sd, strict=False) + model = model.to(shared.device) + if half_attention: + model = model.half() + if shared.cmd_opts.opt_channelslast: + model = model.to(memory_format=torch.channels_last) + + sd_hijack.model_hijack.hijack(model) # apply optimization + model.eval() + + if shared.opts.ldsr_cached: + cached_ldsr_model = model + + return {"model": model} + + def __init__(self, model_path, yaml_path): + self.modelPath = model_path + self.yamlPath = yaml_path + + @staticmethod + def run(model, selected_path, custom_steps, eta): + example = get_cond(selected_path) + + n_runs = 1 + guider = None + ckwargs = None + ddim_use_x0_pred = False + temperature = 1. + eta = eta + custom_shape = None + + height, width = example["image"].shape[1:3] + split_input = height >= 128 and width >= 128 + + if split_input: + ks = 128 + stride = 64 + vqf = 4 # + model.split_input_params = {"ks": (ks, ks), "stride": (stride, stride), + "vqf": vqf, + "patch_distributed_vq": True, + "tie_braker": False, + "clip_max_weight": 0.5, + "clip_min_weight": 0.01, + "clip_max_tie_weight": 0.5, + "clip_min_tie_weight": 0.01} + else: + if hasattr(model, "split_input_params"): + delattr(model, "split_input_params") + + x_t = None + logs = None + for _ in range(n_runs): + if custom_shape is not None: + x_t = torch.randn(1, custom_shape[1], custom_shape[2], custom_shape[3]).to(model.device) + x_t = repeat(x_t, '1 c h w -> b c h w', b=custom_shape[0]) + + logs = make_convolutional_sample(example, model, + custom_steps=custom_steps, + eta=eta, quantize_x0=False, + custom_shape=custom_shape, + temperature=temperature, noise_dropout=0., + corrector=guider, corrector_kwargs=ckwargs, x_T=x_t, + ddim_use_x0_pred=ddim_use_x0_pred + ) + return logs + + def super_resolution(self, image, steps=100, target_scale=2, half_attention=False): + model = self.load_model_from_config(half_attention) + + # Run settings + diffusion_steps = int(steps) + eta = 1.0 + + + gc.collect() + devices.torch_gc() + + im_og = image + width_og, height_og = im_og.size + # If we can adjust the max upscale size, then the 4 below should be our variable + down_sample_rate = target_scale / 4 + wd = width_og * down_sample_rate + hd = height_og * down_sample_rate + width_downsampled_pre = int(np.ceil(wd)) + height_downsampled_pre = int(np.ceil(hd)) + + if down_sample_rate != 1: + print( + f'Downsampling from [{width_og}, {height_og}] to [{width_downsampled_pre}, {height_downsampled_pre}]') + im_og = im_og.resize((width_downsampled_pre, height_downsampled_pre), Image.LANCZOS) + else: + print(f"Down sample rate is 1 from {target_scale} / 4 (Not downsampling)") + + # pad width and height to multiples of 64, pads with the edge values of image to avoid artifacts + pad_w, pad_h = np.max(((2, 2), np.ceil(np.array(im_og.size) / 64).astype(int)), axis=0) * 64 - im_og.size + im_padded = Image.fromarray(np.pad(np.array(im_og), ((0, pad_h), (0, pad_w), (0, 0)), mode='edge')) + + logs = self.run(model["model"], im_padded, diffusion_steps, eta) + + sample = logs["sample"] + sample = sample.detach().cpu() + sample = torch.clamp(sample, -1., 1.) + sample = (sample + 1.) / 2. * 255 + sample = sample.numpy().astype(np.uint8) + sample = np.transpose(sample, (0, 2, 3, 1)) + a = Image.fromarray(sample[0]) + + # remove padding + a = a.crop((0, 0) + tuple(np.array(im_og.size) * 4)) + + del model + gc.collect() + devices.torch_gc() + + return a + + +def get_cond(selected_path): + example = {} + up_f = 4 + c = selected_path.convert('RGB') + c = torch.unsqueeze(torchvision.transforms.ToTensor()(c), 0) + c_up = torchvision.transforms.functional.resize(c, size=[up_f * c.shape[2], up_f * c.shape[3]], + antialias=True) + c_up = rearrange(c_up, '1 c h w -> 1 h w c') + c = rearrange(c, '1 c h w -> 1 h w c') + c = 2. * c - 1. + + c = c.to(shared.device) + example["LR_image"] = c + example["image"] = c_up + + return example + + +@torch.no_grad() +def convsample_ddim(model, cond, steps, shape, eta=1.0, callback=None, normals_sequence=None, + mask=None, x0=None, quantize_x0=False, temperature=1., score_corrector=None, + corrector_kwargs=None, x_t=None + ): + ddim = DDIMSampler(model) + bs = shape[0] + shape = shape[1:] + print(f"Sampling with eta = {eta}; steps: {steps}") + samples, intermediates = ddim.sample(steps, batch_size=bs, shape=shape, conditioning=cond, callback=callback, + normals_sequence=normals_sequence, quantize_x0=quantize_x0, eta=eta, + mask=mask, x0=x0, temperature=temperature, verbose=False, + score_corrector=score_corrector, + corrector_kwargs=corrector_kwargs, x_t=x_t) + + return samples, intermediates + + +@torch.no_grad() +def make_convolutional_sample(batch, model, custom_steps=None, eta=1.0, quantize_x0=False, custom_shape=None, temperature=1., noise_dropout=0., corrector=None, + corrector_kwargs=None, x_T=None, ddim_use_x0_pred=False): + log = {} + + z, c, x, xrec, xc = model.get_input(batch, model.first_stage_key, + return_first_stage_outputs=True, + force_c_encode=not (hasattr(model, 'split_input_params') + and model.cond_stage_key == 'coordinates_bbox'), + return_original_cond=True) + + if custom_shape is not None: + z = torch.randn(custom_shape) + print(f"Generating {custom_shape[0]} samples of shape {custom_shape[1:]}") + + z0 = None + + log["input"] = x + log["reconstruction"] = xrec + + if ismap(xc): + log["original_conditioning"] = model.to_rgb(xc) + if hasattr(model, 'cond_stage_key'): + log[model.cond_stage_key] = model.to_rgb(xc) + + else: + log["original_conditioning"] = xc if xc is not None else torch.zeros_like(x) + if model.cond_stage_model: + log[model.cond_stage_key] = xc if xc is not None else torch.zeros_like(x) + if model.cond_stage_key == 'class_label': + log[model.cond_stage_key] = xc[model.cond_stage_key] + + with model.ema_scope("Plotting"): + t0 = time.time() + + sample, intermediates = convsample_ddim(model, c, steps=custom_steps, shape=z.shape, + eta=eta, + quantize_x0=quantize_x0, mask=None, x0=z0, + temperature=temperature, score_corrector=corrector, corrector_kwargs=corrector_kwargs, + x_t=x_T) + t1 = time.time() + + if ddim_use_x0_pred: + sample = intermediates['pred_x0'][-1] + + x_sample = model.decode_first_stage(sample) + + try: + x_sample_noquant = model.decode_first_stage(sample, force_not_quantize=True) + log["sample_noquant"] = x_sample_noquant + log["sample_diff"] = torch.abs(x_sample_noquant - x_sample) + except Exception: + pass + + log["sample"] = x_sample + log["time"] = t1 - t0 + + return log diff --git a/stable-diffusion-webui/extensions-builtin/LDSR/preload.py b/stable-diffusion-webui/extensions-builtin/LDSR/preload.py new file mode 100755 index 0000000..cfd478d --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/LDSR/preload.py @@ -0,0 +1,6 @@ +import os +from modules import paths + + +def preload(parser): + parser.add_argument("--ldsr-models-path", type=str, help="Path to directory with LDSR model file(s).", default=os.path.join(paths.models_path, 'LDSR')) diff --git a/stable-diffusion-webui/extensions-builtin/LDSR/scripts/ldsr_model.py b/stable-diffusion-webui/extensions-builtin/LDSR/scripts/ldsr_model.py new file mode 100755 index 0000000..bd78dec --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/LDSR/scripts/ldsr_model.py @@ -0,0 +1,68 @@ +import os + +from modules.modelloader import load_file_from_url +from modules.upscaler import Upscaler, UpscalerData +from ldsr_model_arch import LDSR +from modules import shared, script_callbacks, errors +import sd_hijack_autoencoder # noqa: F401 +import sd_hijack_ddpm_v1 # noqa: F401 + + +class UpscalerLDSR(Upscaler): + def __init__(self, user_path): + self.name = "LDSR" + self.user_path = user_path + self.model_url = "https://heibox.uni-heidelberg.de/f/578df07c8fc04ffbadf3/?dl=1" + self.yaml_url = "https://heibox.uni-heidelberg.de/f/31a76b13ea27482981b4/?dl=1" + super().__init__() + scaler_data = UpscalerData("LDSR", None, self) + self.scalers = [scaler_data] + + def load_model(self, path: str): + # Remove incorrect project.yaml file if too big + yaml_path = os.path.join(self.model_path, "project.yaml") + old_model_path = os.path.join(self.model_path, "model.pth") + new_model_path = os.path.join(self.model_path, "model.ckpt") + + local_model_paths = self.find_models(ext_filter=[".ckpt", ".safetensors"]) + local_ckpt_path = next(iter([local_model for local_model in local_model_paths if local_model.endswith("model.ckpt")]), None) + local_safetensors_path = next(iter([local_model for local_model in local_model_paths if local_model.endswith("model.safetensors")]), None) + local_yaml_path = next(iter([local_model for local_model in local_model_paths if local_model.endswith("project.yaml")]), None) + + if os.path.exists(yaml_path): + statinfo = os.stat(yaml_path) + if statinfo.st_size >= 10485760: + print("Removing invalid LDSR YAML file.") + os.remove(yaml_path) + + if os.path.exists(old_model_path): + print("Renaming model from model.pth to model.ckpt") + os.rename(old_model_path, new_model_path) + + if local_safetensors_path is not None and os.path.exists(local_safetensors_path): + model = local_safetensors_path + else: + model = local_ckpt_path or load_file_from_url(self.model_url, model_dir=self.model_download_path, file_name="model.ckpt") + + yaml = local_yaml_path or load_file_from_url(self.yaml_url, model_dir=self.model_download_path, file_name="project.yaml") + + return LDSR(model, yaml) + + def do_upscale(self, img, path): + try: + ldsr = self.load_model(path) + except Exception: + errors.report(f"Failed loading LDSR model {path}", exc_info=True) + return img + ddim_steps = shared.opts.ldsr_steps + return ldsr.super_resolution(img, ddim_steps, self.scale) + + +def on_ui_settings(): + import gradio as gr + + shared.opts.add_option("ldsr_steps", shared.OptionInfo(100, "LDSR processing steps. Lower = faster", gr.Slider, {"minimum": 1, "maximum": 200, "step": 1}, section=('upscaling', "Upscaling"))) + shared.opts.add_option("ldsr_cached", shared.OptionInfo(False, "Cache LDSR model in memory", gr.Checkbox, {"interactive": True}, section=('upscaling', "Upscaling"))) + + +script_callbacks.on_ui_settings(on_ui_settings) diff --git a/stable-diffusion-webui/extensions-builtin/LDSR/sd_hijack_autoencoder.py b/stable-diffusion-webui/extensions-builtin/LDSR/sd_hijack_autoencoder.py new file mode 100755 index 0000000..c29d274 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/LDSR/sd_hijack_autoencoder.py @@ -0,0 +1,293 @@ +# The content of this file comes from the ldm/models/autoencoder.py file of the compvis/stable-diffusion repo +# The VQModel & VQModelInterface were subsequently removed from ldm/models/autoencoder.py when we moved to the stability-ai/stablediffusion repo +# As the LDSR upscaler relies on VQModel & VQModelInterface, the hijack aims to put them back into the ldm.models.autoencoder +import numpy as np +import torch +import pytorch_lightning as pl +import torch.nn.functional as F +from contextlib import contextmanager + +from torch.optim.lr_scheduler import LambdaLR + +from ldm.modules.ema import LitEma +from vqvae_quantize import VectorQuantizer2 as VectorQuantizer +from ldm.modules.diffusionmodules.model import Encoder, Decoder +from ldm.util import instantiate_from_config + +import ldm.models.autoencoder +from packaging import version + +class VQModel(pl.LightningModule): + def __init__(self, + ddconfig, + lossconfig, + n_embed, + embed_dim, + ckpt_path=None, + ignore_keys=None, + image_key="image", + colorize_nlabels=None, + monitor=None, + batch_resize_range=None, + scheduler_config=None, + lr_g_factor=1.0, + remap=None, + sane_index_shape=False, # tell vector quantizer to return indices as bhw + use_ema=False + ): + super().__init__() + self.embed_dim = embed_dim + self.n_embed = n_embed + self.image_key = image_key + self.encoder = Encoder(**ddconfig) + self.decoder = Decoder(**ddconfig) + self.loss = instantiate_from_config(lossconfig) + self.quantize = VectorQuantizer(n_embed, embed_dim, beta=0.25, + remap=remap, + sane_index_shape=sane_index_shape) + self.quant_conv = torch.nn.Conv2d(ddconfig["z_channels"], embed_dim, 1) + self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1) + if colorize_nlabels is not None: + assert type(colorize_nlabels)==int + self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1)) + if monitor is not None: + self.monitor = monitor + self.batch_resize_range = batch_resize_range + if self.batch_resize_range is not None: + print(f"{self.__class__.__name__}: Using per-batch resizing in range {batch_resize_range}.") + + self.use_ema = use_ema + if self.use_ema: + self.model_ema = LitEma(self) + print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.") + + if ckpt_path is not None: + self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys or []) + self.scheduler_config = scheduler_config + self.lr_g_factor = lr_g_factor + + @contextmanager + def ema_scope(self, context=None): + if self.use_ema: + self.model_ema.store(self.parameters()) + self.model_ema.copy_to(self) + if context is not None: + print(f"{context}: Switched to EMA weights") + try: + yield None + finally: + if self.use_ema: + self.model_ema.restore(self.parameters()) + if context is not None: + print(f"{context}: Restored training weights") + + def init_from_ckpt(self, path, ignore_keys=None): + sd = torch.load(path, map_location="cpu")["state_dict"] + keys = list(sd.keys()) + for k in keys: + for ik in ignore_keys or []: + if k.startswith(ik): + print("Deleting key {} from state_dict.".format(k)) + del sd[k] + missing, unexpected = self.load_state_dict(sd, strict=False) + print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys") + if missing: + print(f"Missing Keys: {missing}") + if unexpected: + print(f"Unexpected Keys: {unexpected}") + + def on_train_batch_end(self, *args, **kwargs): + if self.use_ema: + self.model_ema(self) + + def encode(self, x): + h = self.encoder(x) + h = self.quant_conv(h) + quant, emb_loss, info = self.quantize(h) + return quant, emb_loss, info + + def encode_to_prequant(self, x): + h = self.encoder(x) + h = self.quant_conv(h) + return h + + def decode(self, quant): + quant = self.post_quant_conv(quant) + dec = self.decoder(quant) + return dec + + def decode_code(self, code_b): + quant_b = self.quantize.embed_code(code_b) + dec = self.decode(quant_b) + return dec + + def forward(self, input, return_pred_indices=False): + quant, diff, (_,_,ind) = self.encode(input) + dec = self.decode(quant) + if return_pred_indices: + return dec, diff, ind + return dec, diff + + def get_input(self, batch, k): + x = batch[k] + if len(x.shape) == 3: + x = x[..., None] + x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float() + if self.batch_resize_range is not None: + lower_size = self.batch_resize_range[0] + upper_size = self.batch_resize_range[1] + if self.global_step <= 4: + # do the first few batches with max size to avoid later oom + new_resize = upper_size + else: + new_resize = np.random.choice(np.arange(lower_size, upper_size+16, 16)) + if new_resize != x.shape[2]: + x = F.interpolate(x, size=new_resize, mode="bicubic") + x = x.detach() + return x + + def training_step(self, batch, batch_idx, optimizer_idx): + # https://github.com/pytorch/pytorch/issues/37142 + # try not to fool the heuristics + x = self.get_input(batch, self.image_key) + xrec, qloss, ind = self(x, return_pred_indices=True) + + if optimizer_idx == 0: + # autoencode + aeloss, log_dict_ae = self.loss(qloss, x, xrec, optimizer_idx, self.global_step, + last_layer=self.get_last_layer(), split="train", + predicted_indices=ind) + + self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True) + return aeloss + + if optimizer_idx == 1: + # discriminator + discloss, log_dict_disc = self.loss(qloss, x, xrec, optimizer_idx, self.global_step, + last_layer=self.get_last_layer(), split="train") + self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True) + return discloss + + def validation_step(self, batch, batch_idx): + log_dict = self._validation_step(batch, batch_idx) + with self.ema_scope(): + self._validation_step(batch, batch_idx, suffix="_ema") + return log_dict + + def _validation_step(self, batch, batch_idx, suffix=""): + x = self.get_input(batch, self.image_key) + xrec, qloss, ind = self(x, return_pred_indices=True) + aeloss, log_dict_ae = self.loss(qloss, x, xrec, 0, + self.global_step, + last_layer=self.get_last_layer(), + split="val"+suffix, + predicted_indices=ind + ) + + discloss, log_dict_disc = self.loss(qloss, x, xrec, 1, + self.global_step, + last_layer=self.get_last_layer(), + split="val"+suffix, + predicted_indices=ind + ) + rec_loss = log_dict_ae[f"val{suffix}/rec_loss"] + self.log(f"val{suffix}/rec_loss", rec_loss, + prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True) + self.log(f"val{suffix}/aeloss", aeloss, + prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True) + if version.parse(pl.__version__) >= version.parse('1.4.0'): + del log_dict_ae[f"val{suffix}/rec_loss"] + self.log_dict(log_dict_ae) + self.log_dict(log_dict_disc) + return self.log_dict + + def configure_optimizers(self): + lr_d = self.learning_rate + lr_g = self.lr_g_factor*self.learning_rate + print("lr_d", lr_d) + print("lr_g", lr_g) + opt_ae = torch.optim.Adam(list(self.encoder.parameters())+ + list(self.decoder.parameters())+ + list(self.quantize.parameters())+ + list(self.quant_conv.parameters())+ + list(self.post_quant_conv.parameters()), + lr=lr_g, betas=(0.5, 0.9)) + opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(), + lr=lr_d, betas=(0.5, 0.9)) + + if self.scheduler_config is not None: + scheduler = instantiate_from_config(self.scheduler_config) + + print("Setting up LambdaLR scheduler...") + scheduler = [ + { + 'scheduler': LambdaLR(opt_ae, lr_lambda=scheduler.schedule), + 'interval': 'step', + 'frequency': 1 + }, + { + 'scheduler': LambdaLR(opt_disc, lr_lambda=scheduler.schedule), + 'interval': 'step', + 'frequency': 1 + }, + ] + return [opt_ae, opt_disc], scheduler + return [opt_ae, opt_disc], [] + + def get_last_layer(self): + return self.decoder.conv_out.weight + + def log_images(self, batch, only_inputs=False, plot_ema=False, **kwargs): + log = {} + x = self.get_input(batch, self.image_key) + x = x.to(self.device) + if only_inputs: + log["inputs"] = x + return log + xrec, _ = self(x) + if x.shape[1] > 3: + # colorize with random projection + assert xrec.shape[1] > 3 + x = self.to_rgb(x) + xrec = self.to_rgb(xrec) + log["inputs"] = x + log["reconstructions"] = xrec + if plot_ema: + with self.ema_scope(): + xrec_ema, _ = self(x) + if x.shape[1] > 3: + xrec_ema = self.to_rgb(xrec_ema) + log["reconstructions_ema"] = xrec_ema + return log + + def to_rgb(self, x): + assert self.image_key == "segmentation" + if not hasattr(self, "colorize"): + self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x)) + x = F.conv2d(x, weight=self.colorize) + x = 2.*(x-x.min())/(x.max()-x.min()) - 1. + return x + + +class VQModelInterface(VQModel): + def __init__(self, embed_dim, *args, **kwargs): + super().__init__(*args, embed_dim=embed_dim, **kwargs) + self.embed_dim = embed_dim + + def encode(self, x): + h = self.encoder(x) + h = self.quant_conv(h) + return h + + def decode(self, h, force_not_quantize=False): + # also go through quantization layer + if not force_not_quantize: + quant, emb_loss, info = self.quantize(h) + else: + quant = h + quant = self.post_quant_conv(quant) + dec = self.decoder(quant) + return dec + +ldm.models.autoencoder.VQModel = VQModel +ldm.models.autoencoder.VQModelInterface = VQModelInterface diff --git a/stable-diffusion-webui/extensions-builtin/LDSR/sd_hijack_ddpm_v1.py b/stable-diffusion-webui/extensions-builtin/LDSR/sd_hijack_ddpm_v1.py new file mode 100755 index 0000000..51ab182 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/LDSR/sd_hijack_ddpm_v1.py @@ -0,0 +1,1443 @@ +# This script is copied from the compvis/stable-diffusion repo (aka the SD V1 repo) +# Original filename: ldm/models/diffusion/ddpm.py +# The purpose to reinstate the old DDPM logic which works with VQ, whereas the V2 one doesn't +# Some models such as LDSR require VQ to work correctly +# The classes are suffixed with "V1" and added back to the "ldm.models.diffusion.ddpm" module + +import torch +import torch.nn as nn +import numpy as np +import pytorch_lightning as pl +from torch.optim.lr_scheduler import LambdaLR +from einops import rearrange, repeat +from contextlib import contextmanager +from functools import partial +from tqdm import tqdm +from torchvision.utils import make_grid +from pytorch_lightning.utilities.distributed import rank_zero_only + +from ldm.util import log_txt_as_img, exists, default, ismap, isimage, mean_flat, count_params, instantiate_from_config +from ldm.modules.ema import LitEma +from ldm.modules.distributions.distributions import normal_kl, DiagonalGaussianDistribution +from ldm.models.autoencoder import VQModelInterface, IdentityFirstStage, AutoencoderKL +from ldm.modules.diffusionmodules.util import make_beta_schedule, extract_into_tensor, noise_like +from ldm.models.diffusion.ddim import DDIMSampler + +import ldm.models.diffusion.ddpm + +__conditioning_keys__ = {'concat': 'c_concat', + 'crossattn': 'c_crossattn', + 'adm': 'y'} + + +def disabled_train(self, mode=True): + """Overwrite model.train with this function to make sure train/eval mode + does not change anymore.""" + return self + + +def uniform_on_device(r1, r2, shape, device): + return (r1 - r2) * torch.rand(*shape, device=device) + r2 + + +class DDPMV1(pl.LightningModule): + # classic DDPM with Gaussian diffusion, in image space + def __init__(self, + unet_config, + timesteps=1000, + beta_schedule="linear", + loss_type="l2", + ckpt_path=None, + ignore_keys=None, + load_only_unet=False, + monitor="val/loss", + use_ema=True, + first_stage_key="image", + image_size=256, + channels=3, + log_every_t=100, + clip_denoised=True, + linear_start=1e-4, + linear_end=2e-2, + cosine_s=8e-3, + given_betas=None, + original_elbo_weight=0., + v_posterior=0., # weight for choosing posterior variance as sigma = (1-v) * beta_tilde + v * beta + l_simple_weight=1., + conditioning_key=None, + parameterization="eps", # all assuming fixed variance schedules + scheduler_config=None, + use_positional_encodings=False, + learn_logvar=False, + logvar_init=0., + ): + super().__init__() + assert parameterization in ["eps", "x0"], 'currently only supporting "eps" and "x0"' + self.parameterization = parameterization + print(f"{self.__class__.__name__}: Running in {self.parameterization}-prediction mode") + self.cond_stage_model = None + self.clip_denoised = clip_denoised + self.log_every_t = log_every_t + self.first_stage_key = first_stage_key + self.image_size = image_size # try conv? + self.channels = channels + self.use_positional_encodings = use_positional_encodings + self.model = DiffusionWrapperV1(unet_config, conditioning_key) + count_params(self.model, verbose=True) + self.use_ema = use_ema + if self.use_ema: + self.model_ema = LitEma(self.model) + print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.") + + self.use_scheduler = scheduler_config is not None + if self.use_scheduler: + self.scheduler_config = scheduler_config + + self.v_posterior = v_posterior + self.original_elbo_weight = original_elbo_weight + self.l_simple_weight = l_simple_weight + + if monitor is not None: + self.monitor = monitor + if ckpt_path is not None: + self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys or [], only_model=load_only_unet) + + self.register_schedule(given_betas=given_betas, beta_schedule=beta_schedule, timesteps=timesteps, + linear_start=linear_start, linear_end=linear_end, cosine_s=cosine_s) + + self.loss_type = loss_type + + self.learn_logvar = learn_logvar + self.logvar = torch.full(fill_value=logvar_init, size=(self.num_timesteps,)) + if self.learn_logvar: + self.logvar = nn.Parameter(self.logvar, requires_grad=True) + + + def register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000, + linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): + if exists(given_betas): + betas = given_betas + else: + betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end, + cosine_s=cosine_s) + alphas = 1. - betas + alphas_cumprod = np.cumprod(alphas, axis=0) + alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1]) + + timesteps, = betas.shape + self.num_timesteps = int(timesteps) + self.linear_start = linear_start + self.linear_end = linear_end + assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep' + + to_torch = partial(torch.tensor, dtype=torch.float32) + + self.register_buffer('betas', to_torch(betas)) + self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod)) + self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev)) + + # calculations for diffusion q(x_t | x_{t-1}) and others + self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod))) + self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod))) + self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod))) + self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod))) + self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1))) + + # calculations for posterior q(x_{t-1} | x_t, x_0) + posterior_variance = (1 - self.v_posterior) * betas * (1. - alphas_cumprod_prev) / ( + 1. - alphas_cumprod) + self.v_posterior * betas + # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t) + self.register_buffer('posterior_variance', to_torch(posterior_variance)) + # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain + self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20)))) + self.register_buffer('posterior_mean_coef1', to_torch( + betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))) + self.register_buffer('posterior_mean_coef2', to_torch( + (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod))) + + if self.parameterization == "eps": + lvlb_weights = self.betas ** 2 / ( + 2 * self.posterior_variance * to_torch(alphas) * (1 - self.alphas_cumprod)) + elif self.parameterization == "x0": + lvlb_weights = 0.5 * np.sqrt(torch.Tensor(alphas_cumprod)) / (2. * 1 - torch.Tensor(alphas_cumprod)) + else: + raise NotImplementedError("mu not supported") + # TODO how to choose this term + lvlb_weights[0] = lvlb_weights[1] + self.register_buffer('lvlb_weights', lvlb_weights, persistent=False) + assert not torch.isnan(self.lvlb_weights).all() + + @contextmanager + def ema_scope(self, context=None): + if self.use_ema: + self.model_ema.store(self.model.parameters()) + self.model_ema.copy_to(self.model) + if context is not None: + print(f"{context}: Switched to EMA weights") + try: + yield None + finally: + if self.use_ema: + self.model_ema.restore(self.model.parameters()) + if context is not None: + print(f"{context}: Restored training weights") + + def init_from_ckpt(self, path, ignore_keys=None, only_model=False): + sd = torch.load(path, map_location="cpu") + if "state_dict" in list(sd.keys()): + sd = sd["state_dict"] + keys = list(sd.keys()) + for k in keys: + for ik in ignore_keys or []: + if k.startswith(ik): + print("Deleting key {} from state_dict.".format(k)) + del sd[k] + missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict( + sd, strict=False) + print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys") + if missing: + print(f"Missing Keys: {missing}") + if unexpected: + print(f"Unexpected Keys: {unexpected}") + + def q_mean_variance(self, x_start, t): + """ + Get the distribution q(x_t | x_0). + :param x_start: the [N x C x ...] tensor of noiseless inputs. + :param t: the number of diffusion steps (minus 1). Here, 0 means one step. + :return: A tuple (mean, variance, log_variance), all of x_start's shape. + """ + mean = (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start) + variance = extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape) + log_variance = extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape) + return mean, variance, log_variance + + def predict_start_from_noise(self, x_t, t, noise): + return ( + extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - + extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise + ) + + def q_posterior(self, x_start, x_t, t): + posterior_mean = ( + extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start + + extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t + ) + posterior_variance = extract_into_tensor(self.posterior_variance, t, x_t.shape) + posterior_log_variance_clipped = extract_into_tensor(self.posterior_log_variance_clipped, t, x_t.shape) + return posterior_mean, posterior_variance, posterior_log_variance_clipped + + def p_mean_variance(self, x, t, clip_denoised: bool): + model_out = self.model(x, t) + if self.parameterization == "eps": + x_recon = self.predict_start_from_noise(x, t=t, noise=model_out) + elif self.parameterization == "x0": + x_recon = model_out + if clip_denoised: + x_recon.clamp_(-1., 1.) + + model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t) + return model_mean, posterior_variance, posterior_log_variance + + @torch.no_grad() + def p_sample(self, x, t, clip_denoised=True, repeat_noise=False): + b, *_, device = *x.shape, x.device + model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, clip_denoised=clip_denoised) + noise = noise_like(x.shape, device, repeat_noise) + # no noise when t == 0 + nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1))) + return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise + + @torch.no_grad() + def p_sample_loop(self, shape, return_intermediates=False): + device = self.betas.device + b = shape[0] + img = torch.randn(shape, device=device) + intermediates = [img] + for i in tqdm(reversed(range(0, self.num_timesteps)), desc='Sampling t', total=self.num_timesteps): + img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long), + clip_denoised=self.clip_denoised) + if i % self.log_every_t == 0 or i == self.num_timesteps - 1: + intermediates.append(img) + if return_intermediates: + return img, intermediates + return img + + @torch.no_grad() + def sample(self, batch_size=16, return_intermediates=False): + image_size = self.image_size + channels = self.channels + return self.p_sample_loop((batch_size, channels, image_size, image_size), + return_intermediates=return_intermediates) + + def q_sample(self, x_start, t, noise=None): + noise = default(noise, lambda: torch.randn_like(x_start)) + return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + + extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise) + + def get_loss(self, pred, target, mean=True): + if self.loss_type == 'l1': + loss = (target - pred).abs() + if mean: + loss = loss.mean() + elif self.loss_type == 'l2': + if mean: + loss = torch.nn.functional.mse_loss(target, pred) + else: + loss = torch.nn.functional.mse_loss(target, pred, reduction='none') + else: + raise NotImplementedError("unknown loss type '{loss_type}'") + + return loss + + def p_losses(self, x_start, t, noise=None): + noise = default(noise, lambda: torch.randn_like(x_start)) + x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) + model_out = self.model(x_noisy, t) + + loss_dict = {} + if self.parameterization == "eps": + target = noise + elif self.parameterization == "x0": + target = x_start + else: + raise NotImplementedError(f"Parameterization {self.parameterization} not yet supported") + + loss = self.get_loss(model_out, target, mean=False).mean(dim=[1, 2, 3]) + + log_prefix = 'train' if self.training else 'val' + + loss_dict.update({f'{log_prefix}/loss_simple': loss.mean()}) + loss_simple = loss.mean() * self.l_simple_weight + + loss_vlb = (self.lvlb_weights[t] * loss).mean() + loss_dict.update({f'{log_prefix}/loss_vlb': loss_vlb}) + + loss = loss_simple + self.original_elbo_weight * loss_vlb + + loss_dict.update({f'{log_prefix}/loss': loss}) + + return loss, loss_dict + + def forward(self, x, *args, **kwargs): + # b, c, h, w, device, img_size, = *x.shape, x.device, self.image_size + # assert h == img_size and w == img_size, f'height and width of image must be {img_size}' + t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long() + return self.p_losses(x, t, *args, **kwargs) + + def get_input(self, batch, k): + x = batch[k] + if len(x.shape) == 3: + x = x[..., None] + x = rearrange(x, 'b h w c -> b c h w') + x = x.to(memory_format=torch.contiguous_format).float() + return x + + def shared_step(self, batch): + x = self.get_input(batch, self.first_stage_key) + loss, loss_dict = self(x) + return loss, loss_dict + + def training_step(self, batch, batch_idx): + loss, loss_dict = self.shared_step(batch) + + self.log_dict(loss_dict, prog_bar=True, + logger=True, on_step=True, on_epoch=True) + + self.log("global_step", self.global_step, + prog_bar=True, logger=True, on_step=True, on_epoch=False) + + if self.use_scheduler: + lr = self.optimizers().param_groups[0]['lr'] + self.log('lr_abs', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False) + + return loss + + @torch.no_grad() + def validation_step(self, batch, batch_idx): + _, loss_dict_no_ema = self.shared_step(batch) + with self.ema_scope(): + _, loss_dict_ema = self.shared_step(batch) + loss_dict_ema = {key + '_ema': loss_dict_ema[key] for key in loss_dict_ema} + self.log_dict(loss_dict_no_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True) + self.log_dict(loss_dict_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True) + + def on_train_batch_end(self, *args, **kwargs): + if self.use_ema: + self.model_ema(self.model) + + def _get_rows_from_list(self, samples): + n_imgs_per_row = len(samples) + denoise_grid = rearrange(samples, 'n b c h w -> b n c h w') + denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w') + denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row) + return denoise_grid + + @torch.no_grad() + def log_images(self, batch, N=8, n_row=2, sample=True, return_keys=None, **kwargs): + log = {} + x = self.get_input(batch, self.first_stage_key) + N = min(x.shape[0], N) + n_row = min(x.shape[0], n_row) + x = x.to(self.device)[:N] + log["inputs"] = x + + # get diffusion row + diffusion_row = [] + x_start = x[:n_row] + + for t in range(self.num_timesteps): + if t % self.log_every_t == 0 or t == self.num_timesteps - 1: + t = repeat(torch.tensor([t]), '1 -> b', b=n_row) + t = t.to(self.device).long() + noise = torch.randn_like(x_start) + x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) + diffusion_row.append(x_noisy) + + log["diffusion_row"] = self._get_rows_from_list(diffusion_row) + + if sample: + # get denoise row + with self.ema_scope("Plotting"): + samples, denoise_row = self.sample(batch_size=N, return_intermediates=True) + + log["samples"] = samples + log["denoise_row"] = self._get_rows_from_list(denoise_row) + + if return_keys: + if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0: + return log + else: + return {key: log[key] for key in return_keys} + return log + + def configure_optimizers(self): + lr = self.learning_rate + params = list(self.model.parameters()) + if self.learn_logvar: + params = params + [self.logvar] + opt = torch.optim.AdamW(params, lr=lr) + return opt + + +class LatentDiffusionV1(DDPMV1): + """main class""" + def __init__(self, + first_stage_config, + cond_stage_config, + num_timesteps_cond=None, + cond_stage_key="image", + cond_stage_trainable=False, + concat_mode=True, + cond_stage_forward=None, + conditioning_key=None, + scale_factor=1.0, + scale_by_std=False, + *args, **kwargs): + self.num_timesteps_cond = default(num_timesteps_cond, 1) + self.scale_by_std = scale_by_std + assert self.num_timesteps_cond <= kwargs['timesteps'] + # for backwards compatibility after implementation of DiffusionWrapper + if conditioning_key is None: + conditioning_key = 'concat' if concat_mode else 'crossattn' + if cond_stage_config == '__is_unconditional__': + conditioning_key = None + ckpt_path = kwargs.pop("ckpt_path", None) + ignore_keys = kwargs.pop("ignore_keys", []) + super().__init__(*args, conditioning_key=conditioning_key, **kwargs) + self.concat_mode = concat_mode + self.cond_stage_trainable = cond_stage_trainable + self.cond_stage_key = cond_stage_key + try: + self.num_downs = len(first_stage_config.params.ddconfig.ch_mult) - 1 + except Exception: + self.num_downs = 0 + if not scale_by_std: + self.scale_factor = scale_factor + else: + self.register_buffer('scale_factor', torch.tensor(scale_factor)) + self.instantiate_first_stage(first_stage_config) + self.instantiate_cond_stage(cond_stage_config) + self.cond_stage_forward = cond_stage_forward + self.clip_denoised = False + self.bbox_tokenizer = None + + self.restarted_from_ckpt = False + if ckpt_path is not None: + self.init_from_ckpt(ckpt_path, ignore_keys) + self.restarted_from_ckpt = True + + def make_cond_schedule(self, ): + self.cond_ids = torch.full(size=(self.num_timesteps,), fill_value=self.num_timesteps - 1, dtype=torch.long) + ids = torch.round(torch.linspace(0, self.num_timesteps - 1, self.num_timesteps_cond)).long() + self.cond_ids[:self.num_timesteps_cond] = ids + + @rank_zero_only + @torch.no_grad() + def on_train_batch_start(self, batch, batch_idx, dataloader_idx): + # only for very first batch + if self.scale_by_std and self.current_epoch == 0 and self.global_step == 0 and batch_idx == 0 and not self.restarted_from_ckpt: + assert self.scale_factor == 1., 'rather not use custom rescaling and std-rescaling simultaneously' + # set rescale weight to 1./std of encodings + print("### USING STD-RESCALING ###") + x = super().get_input(batch, self.first_stage_key) + x = x.to(self.device) + encoder_posterior = self.encode_first_stage(x) + z = self.get_first_stage_encoding(encoder_posterior).detach() + del self.scale_factor + self.register_buffer('scale_factor', 1. / z.flatten().std()) + print(f"setting self.scale_factor to {self.scale_factor}") + print("### USING STD-RESCALING ###") + + def register_schedule(self, + given_betas=None, beta_schedule="linear", timesteps=1000, + linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): + super().register_schedule(given_betas, beta_schedule, timesteps, linear_start, linear_end, cosine_s) + + self.shorten_cond_schedule = self.num_timesteps_cond > 1 + if self.shorten_cond_schedule: + self.make_cond_schedule() + + def instantiate_first_stage(self, config): + model = instantiate_from_config(config) + self.first_stage_model = model.eval() + self.first_stage_model.train = disabled_train + for param in self.first_stage_model.parameters(): + param.requires_grad = False + + def instantiate_cond_stage(self, config): + if not self.cond_stage_trainable: + if config == "__is_first_stage__": + print("Using first stage also as cond stage.") + self.cond_stage_model = self.first_stage_model + elif config == "__is_unconditional__": + print(f"Training {self.__class__.__name__} as an unconditional model.") + self.cond_stage_model = None + # self.be_unconditional = True + else: + model = instantiate_from_config(config) + self.cond_stage_model = model.eval() + self.cond_stage_model.train = disabled_train + for param in self.cond_stage_model.parameters(): + param.requires_grad = False + else: + assert config != '__is_first_stage__' + assert config != '__is_unconditional__' + model = instantiate_from_config(config) + self.cond_stage_model = model + + def _get_denoise_row_from_list(self, samples, desc='', force_no_decoder_quantization=False): + denoise_row = [] + for zd in tqdm(samples, desc=desc): + denoise_row.append(self.decode_first_stage(zd.to(self.device), + force_not_quantize=force_no_decoder_quantization)) + n_imgs_per_row = len(denoise_row) + denoise_row = torch.stack(denoise_row) # n_log_step, n_row, C, H, W + denoise_grid = rearrange(denoise_row, 'n b c h w -> b n c h w') + denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w') + denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row) + return denoise_grid + + def get_first_stage_encoding(self, encoder_posterior): + if isinstance(encoder_posterior, DiagonalGaussianDistribution): + z = encoder_posterior.sample() + elif isinstance(encoder_posterior, torch.Tensor): + z = encoder_posterior + else: + raise NotImplementedError(f"encoder_posterior of type '{type(encoder_posterior)}' not yet implemented") + return self.scale_factor * z + + def get_learned_conditioning(self, c): + if self.cond_stage_forward is None: + if hasattr(self.cond_stage_model, 'encode') and callable(self.cond_stage_model.encode): + c = self.cond_stage_model.encode(c) + if isinstance(c, DiagonalGaussianDistribution): + c = c.mode() + else: + c = self.cond_stage_model(c) + else: + assert hasattr(self.cond_stage_model, self.cond_stage_forward) + c = getattr(self.cond_stage_model, self.cond_stage_forward)(c) + return c + + def meshgrid(self, h, w): + y = torch.arange(0, h).view(h, 1, 1).repeat(1, w, 1) + x = torch.arange(0, w).view(1, w, 1).repeat(h, 1, 1) + + arr = torch.cat([y, x], dim=-1) + return arr + + def delta_border(self, h, w): + """ + :param h: height + :param w: width + :return: normalized distance to image border, + with min distance = 0 at border and max dist = 0.5 at image center + """ + lower_right_corner = torch.tensor([h - 1, w - 1]).view(1, 1, 2) + arr = self.meshgrid(h, w) / lower_right_corner + dist_left_up = torch.min(arr, dim=-1, keepdims=True)[0] + dist_right_down = torch.min(1 - arr, dim=-1, keepdims=True)[0] + edge_dist = torch.min(torch.cat([dist_left_up, dist_right_down], dim=-1), dim=-1)[0] + return edge_dist + + def get_weighting(self, h, w, Ly, Lx, device): + weighting = self.delta_border(h, w) + weighting = torch.clip(weighting, self.split_input_params["clip_min_weight"], + self.split_input_params["clip_max_weight"], ) + weighting = weighting.view(1, h * w, 1).repeat(1, 1, Ly * Lx).to(device) + + if self.split_input_params["tie_braker"]: + L_weighting = self.delta_border(Ly, Lx) + L_weighting = torch.clip(L_weighting, + self.split_input_params["clip_min_tie_weight"], + self.split_input_params["clip_max_tie_weight"]) + + L_weighting = L_weighting.view(1, 1, Ly * Lx).to(device) + weighting = weighting * L_weighting + return weighting + + def get_fold_unfold(self, x, kernel_size, stride, uf=1, df=1): # todo load once not every time, shorten code + """ + :param x: img of size (bs, c, h, w) + :return: n img crops of size (n, bs, c, kernel_size[0], kernel_size[1]) + """ + bs, nc, h, w = x.shape + + # number of crops in image + Ly = (h - kernel_size[0]) // stride[0] + 1 + Lx = (w - kernel_size[1]) // stride[1] + 1 + + if uf == 1 and df == 1: + fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride) + unfold = torch.nn.Unfold(**fold_params) + + fold = torch.nn.Fold(output_size=x.shape[2:], **fold_params) + + weighting = self.get_weighting(kernel_size[0], kernel_size[1], Ly, Lx, x.device).to(x.dtype) + normalization = fold(weighting).view(1, 1, h, w) # normalizes the overlap + weighting = weighting.view((1, 1, kernel_size[0], kernel_size[1], Ly * Lx)) + + elif uf > 1 and df == 1: + fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride) + unfold = torch.nn.Unfold(**fold_params) + + fold_params2 = dict(kernel_size=(kernel_size[0] * uf, kernel_size[0] * uf), + dilation=1, padding=0, + stride=(stride[0] * uf, stride[1] * uf)) + fold = torch.nn.Fold(output_size=(x.shape[2] * uf, x.shape[3] * uf), **fold_params2) + + weighting = self.get_weighting(kernel_size[0] * uf, kernel_size[1] * uf, Ly, Lx, x.device).to(x.dtype) + normalization = fold(weighting).view(1, 1, h * uf, w * uf) # normalizes the overlap + weighting = weighting.view((1, 1, kernel_size[0] * uf, kernel_size[1] * uf, Ly * Lx)) + + elif df > 1 and uf == 1: + fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride) + unfold = torch.nn.Unfold(**fold_params) + + fold_params2 = dict(kernel_size=(kernel_size[0] // df, kernel_size[0] // df), + dilation=1, padding=0, + stride=(stride[0] // df, stride[1] // df)) + fold = torch.nn.Fold(output_size=(x.shape[2] // df, x.shape[3] // df), **fold_params2) + + weighting = self.get_weighting(kernel_size[0] // df, kernel_size[1] // df, Ly, Lx, x.device).to(x.dtype) + normalization = fold(weighting).view(1, 1, h // df, w // df) # normalizes the overlap + weighting = weighting.view((1, 1, kernel_size[0] // df, kernel_size[1] // df, Ly * Lx)) + + else: + raise NotImplementedError + + return fold, unfold, normalization, weighting + + @torch.no_grad() + def get_input(self, batch, k, return_first_stage_outputs=False, force_c_encode=False, + cond_key=None, return_original_cond=False, bs=None): + x = super().get_input(batch, k) + if bs is not None: + x = x[:bs] + x = x.to(self.device) + encoder_posterior = self.encode_first_stage(x) + z = self.get_first_stage_encoding(encoder_posterior).detach() + + if self.model.conditioning_key is not None: + if cond_key is None: + cond_key = self.cond_stage_key + if cond_key != self.first_stage_key: + if cond_key in ['caption', 'coordinates_bbox']: + xc = batch[cond_key] + elif cond_key == 'class_label': + xc = batch + else: + xc = super().get_input(batch, cond_key).to(self.device) + else: + xc = x + if not self.cond_stage_trainable or force_c_encode: + if isinstance(xc, dict) or isinstance(xc, list): + # import pudb; pudb.set_trace() + c = self.get_learned_conditioning(xc) + else: + c = self.get_learned_conditioning(xc.to(self.device)) + else: + c = xc + if bs is not None: + c = c[:bs] + + if self.use_positional_encodings: + pos_x, pos_y = self.compute_latent_shifts(batch) + ckey = __conditioning_keys__[self.model.conditioning_key] + c = {ckey: c, 'pos_x': pos_x, 'pos_y': pos_y} + + else: + c = None + xc = None + if self.use_positional_encodings: + pos_x, pos_y = self.compute_latent_shifts(batch) + c = {'pos_x': pos_x, 'pos_y': pos_y} + out = [z, c] + if return_first_stage_outputs: + xrec = self.decode_first_stage(z) + out.extend([x, xrec]) + if return_original_cond: + out.append(xc) + return out + + @torch.no_grad() + def decode_first_stage(self, z, predict_cids=False, force_not_quantize=False): + if predict_cids: + if z.dim() == 4: + z = torch.argmax(z.exp(), dim=1).long() + z = self.first_stage_model.quantize.get_codebook_entry(z, shape=None) + z = rearrange(z, 'b h w c -> b c h w').contiguous() + + z = 1. / self.scale_factor * z + + if hasattr(self, "split_input_params"): + if self.split_input_params["patch_distributed_vq"]: + ks = self.split_input_params["ks"] # eg. (128, 128) + stride = self.split_input_params["stride"] # eg. (64, 64) + uf = self.split_input_params["vqf"] + bs, nc, h, w = z.shape + if ks[0] > h or ks[1] > w: + ks = (min(ks[0], h), min(ks[1], w)) + print("reducing Kernel") + + if stride[0] > h or stride[1] > w: + stride = (min(stride[0], h), min(stride[1], w)) + print("reducing stride") + + fold, unfold, normalization, weighting = self.get_fold_unfold(z, ks, stride, uf=uf) + + z = unfold(z) # (bn, nc * prod(**ks), L) + # 1. Reshape to img shape + z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L ) + + # 2. apply model loop over last dim + if isinstance(self.first_stage_model, VQModelInterface): + output_list = [self.first_stage_model.decode(z[:, :, :, :, i], + force_not_quantize=predict_cids or force_not_quantize) + for i in range(z.shape[-1])] + else: + + output_list = [self.first_stage_model.decode(z[:, :, :, :, i]) + for i in range(z.shape[-1])] + + o = torch.stack(output_list, axis=-1) # # (bn, nc, ks[0], ks[1], L) + o = o * weighting + # Reverse 1. reshape to img shape + o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L) + # stitch crops together + decoded = fold(o) + decoded = decoded / normalization # norm is shape (1, 1, h, w) + return decoded + else: + if isinstance(self.first_stage_model, VQModelInterface): + return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize) + else: + return self.first_stage_model.decode(z) + + else: + if isinstance(self.first_stage_model, VQModelInterface): + return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize) + else: + return self.first_stage_model.decode(z) + + # same as above but without decorator + def differentiable_decode_first_stage(self, z, predict_cids=False, force_not_quantize=False): + if predict_cids: + if z.dim() == 4: + z = torch.argmax(z.exp(), dim=1).long() + z = self.first_stage_model.quantize.get_codebook_entry(z, shape=None) + z = rearrange(z, 'b h w c -> b c h w').contiguous() + + z = 1. / self.scale_factor * z + + if hasattr(self, "split_input_params"): + if self.split_input_params["patch_distributed_vq"]: + ks = self.split_input_params["ks"] # eg. (128, 128) + stride = self.split_input_params["stride"] # eg. (64, 64) + uf = self.split_input_params["vqf"] + bs, nc, h, w = z.shape + if ks[0] > h or ks[1] > w: + ks = (min(ks[0], h), min(ks[1], w)) + print("reducing Kernel") + + if stride[0] > h or stride[1] > w: + stride = (min(stride[0], h), min(stride[1], w)) + print("reducing stride") + + fold, unfold, normalization, weighting = self.get_fold_unfold(z, ks, stride, uf=uf) + + z = unfold(z) # (bn, nc * prod(**ks), L) + # 1. Reshape to img shape + z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L ) + + # 2. apply model loop over last dim + if isinstance(self.first_stage_model, VQModelInterface): + output_list = [self.first_stage_model.decode(z[:, :, :, :, i], + force_not_quantize=predict_cids or force_not_quantize) + for i in range(z.shape[-1])] + else: + + output_list = [self.first_stage_model.decode(z[:, :, :, :, i]) + for i in range(z.shape[-1])] + + o = torch.stack(output_list, axis=-1) # # (bn, nc, ks[0], ks[1], L) + o = o * weighting + # Reverse 1. reshape to img shape + o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L) + # stitch crops together + decoded = fold(o) + decoded = decoded / normalization # norm is shape (1, 1, h, w) + return decoded + else: + if isinstance(self.first_stage_model, VQModelInterface): + return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize) + else: + return self.first_stage_model.decode(z) + + else: + if isinstance(self.first_stage_model, VQModelInterface): + return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize) + else: + return self.first_stage_model.decode(z) + + @torch.no_grad() + def encode_first_stage(self, x): + if hasattr(self, "split_input_params"): + if self.split_input_params["patch_distributed_vq"]: + ks = self.split_input_params["ks"] # eg. (128, 128) + stride = self.split_input_params["stride"] # eg. (64, 64) + df = self.split_input_params["vqf"] + self.split_input_params['original_image_size'] = x.shape[-2:] + bs, nc, h, w = x.shape + if ks[0] > h or ks[1] > w: + ks = (min(ks[0], h), min(ks[1], w)) + print("reducing Kernel") + + if stride[0] > h or stride[1] > w: + stride = (min(stride[0], h), min(stride[1], w)) + print("reducing stride") + + fold, unfold, normalization, weighting = self.get_fold_unfold(x, ks, stride, df=df) + z = unfold(x) # (bn, nc * prod(**ks), L) + # Reshape to img shape + z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L ) + + output_list = [self.first_stage_model.encode(z[:, :, :, :, i]) + for i in range(z.shape[-1])] + + o = torch.stack(output_list, axis=-1) + o = o * weighting + + # Reverse reshape to img shape + o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L) + # stitch crops together + decoded = fold(o) + decoded = decoded / normalization + return decoded + + else: + return self.first_stage_model.encode(x) + else: + return self.first_stage_model.encode(x) + + def shared_step(self, batch, **kwargs): + x, c = self.get_input(batch, self.first_stage_key) + loss = self(x, c) + return loss + + def forward(self, x, c, *args, **kwargs): + t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long() + if self.model.conditioning_key is not None: + assert c is not None + if self.cond_stage_trainable: + c = self.get_learned_conditioning(c) + if self.shorten_cond_schedule: # TODO: drop this option + tc = self.cond_ids[t].to(self.device) + c = self.q_sample(x_start=c, t=tc, noise=torch.randn_like(c.float())) + return self.p_losses(x, c, t, *args, **kwargs) + + def apply_model(self, x_noisy, t, cond, return_ids=False): + + if isinstance(cond, dict): + # hybrid case, cond is expected to be a dict + pass + else: + if not isinstance(cond, list): + cond = [cond] + key = 'c_concat' if self.model.conditioning_key == 'concat' else 'c_crossattn' + cond = {key: cond} + + if hasattr(self, "split_input_params"): + assert len(cond) == 1 # todo can only deal with one conditioning atm + assert not return_ids + ks = self.split_input_params["ks"] # eg. (128, 128) + stride = self.split_input_params["stride"] # eg. (64, 64) + + h, w = x_noisy.shape[-2:] + + fold, unfold, normalization, weighting = self.get_fold_unfold(x_noisy, ks, stride) + + z = unfold(x_noisy) # (bn, nc * prod(**ks), L) + # Reshape to img shape + z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L ) + z_list = [z[:, :, :, :, i] for i in range(z.shape[-1])] + + if self.cond_stage_key in ["image", "LR_image", "segmentation", + 'bbox_img'] and self.model.conditioning_key: # todo check for completeness + c_key = next(iter(cond.keys())) # get key + c = next(iter(cond.values())) # get value + assert (len(c) == 1) # todo extend to list with more than one elem + c = c[0] # get element + + c = unfold(c) + c = c.view((c.shape[0], -1, ks[0], ks[1], c.shape[-1])) # (bn, nc, ks[0], ks[1], L ) + + cond_list = [{c_key: [c[:, :, :, :, i]]} for i in range(c.shape[-1])] + + elif self.cond_stage_key == 'coordinates_bbox': + assert 'original_image_size' in self.split_input_params, 'BoundingBoxRescaling is missing original_image_size' + + # assuming padding of unfold is always 0 and its dilation is always 1 + n_patches_per_row = int((w - ks[0]) / stride[0] + 1) + full_img_h, full_img_w = self.split_input_params['original_image_size'] + # as we are operating on latents, we need the factor from the original image size to the + # spatial latent size to properly rescale the crops for regenerating the bbox annotations + num_downs = self.first_stage_model.encoder.num_resolutions - 1 + rescale_latent = 2 ** (num_downs) + + # get top left positions of patches as conforming for the bbbox tokenizer, therefore we + # need to rescale the tl patch coordinates to be in between (0,1) + tl_patch_coordinates = [(rescale_latent * stride[0] * (patch_nr % n_patches_per_row) / full_img_w, + rescale_latent * stride[1] * (patch_nr // n_patches_per_row) / full_img_h) + for patch_nr in range(z.shape[-1])] + + # patch_limits are tl_coord, width and height coordinates as (x_tl, y_tl, h, w) + patch_limits = [(x_tl, y_tl, + rescale_latent * ks[0] / full_img_w, + rescale_latent * ks[1] / full_img_h) for x_tl, y_tl in tl_patch_coordinates] + # patch_values = [(np.arange(x_tl,min(x_tl+ks, 1.)),np.arange(y_tl,min(y_tl+ks, 1.))) for x_tl, y_tl in tl_patch_coordinates] + + # tokenize crop coordinates for the bounding boxes of the respective patches + patch_limits_tknzd = [torch.LongTensor(self.bbox_tokenizer._crop_encoder(bbox))[None].to(self.device) + for bbox in patch_limits] # list of length l with tensors of shape (1, 2) + print(patch_limits_tknzd[0].shape) + # cut tknzd crop position from conditioning + assert isinstance(cond, dict), 'cond must be dict to be fed into model' + cut_cond = cond['c_crossattn'][0][..., :-2].to(self.device) + print(cut_cond.shape) + + adapted_cond = torch.stack([torch.cat([cut_cond, p], dim=1) for p in patch_limits_tknzd]) + adapted_cond = rearrange(adapted_cond, 'l b n -> (l b) n') + print(adapted_cond.shape) + adapted_cond = self.get_learned_conditioning(adapted_cond) + print(adapted_cond.shape) + adapted_cond = rearrange(adapted_cond, '(l b) n d -> l b n d', l=z.shape[-1]) + print(adapted_cond.shape) + + cond_list = [{'c_crossattn': [e]} for e in adapted_cond] + + else: + cond_list = [cond for i in range(z.shape[-1])] # Todo make this more efficient + + # apply model by loop over crops + output_list = [self.model(z_list[i], t, **cond_list[i]) for i in range(z.shape[-1])] + assert not isinstance(output_list[0], + tuple) # todo cant deal with multiple model outputs check this never happens + + o = torch.stack(output_list, axis=-1) + o = o * weighting + # Reverse reshape to img shape + o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L) + # stitch crops together + x_recon = fold(o) / normalization + + else: + x_recon = self.model(x_noisy, t, **cond) + + if isinstance(x_recon, tuple) and not return_ids: + return x_recon[0] + else: + return x_recon + + def _predict_eps_from_xstart(self, x_t, t, pred_xstart): + return (extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart) / \ + extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) + + def _prior_bpd(self, x_start): + """ + Get the prior KL term for the variational lower-bound, measured in + bits-per-dim. + This term can't be optimized, as it only depends on the encoder. + :param x_start: the [N x C x ...] tensor of inputs. + :return: a batch of [N] KL values (in bits), one per batch element. + """ + batch_size = x_start.shape[0] + t = torch.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device) + qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t) + kl_prior = normal_kl(mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0) + return mean_flat(kl_prior) / np.log(2.0) + + def p_losses(self, x_start, cond, t, noise=None): + noise = default(noise, lambda: torch.randn_like(x_start)) + x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) + model_output = self.apply_model(x_noisy, t, cond) + + loss_dict = {} + prefix = 'train' if self.training else 'val' + + if self.parameterization == "x0": + target = x_start + elif self.parameterization == "eps": + target = noise + else: + raise NotImplementedError() + + loss_simple = self.get_loss(model_output, target, mean=False).mean([1, 2, 3]) + loss_dict.update({f'{prefix}/loss_simple': loss_simple.mean()}) + + logvar_t = self.logvar[t].to(self.device) + loss = loss_simple / torch.exp(logvar_t) + logvar_t + # loss = loss_simple / torch.exp(self.logvar) + self.logvar + if self.learn_logvar: + loss_dict.update({f'{prefix}/loss_gamma': loss.mean()}) + loss_dict.update({'logvar': self.logvar.data.mean()}) + + loss = self.l_simple_weight * loss.mean() + + loss_vlb = self.get_loss(model_output, target, mean=False).mean(dim=(1, 2, 3)) + loss_vlb = (self.lvlb_weights[t] * loss_vlb).mean() + loss_dict.update({f'{prefix}/loss_vlb': loss_vlb}) + loss += (self.original_elbo_weight * loss_vlb) + loss_dict.update({f'{prefix}/loss': loss}) + + return loss, loss_dict + + def p_mean_variance(self, x, c, t, clip_denoised: bool, return_codebook_ids=False, quantize_denoised=False, + return_x0=False, score_corrector=None, corrector_kwargs=None): + t_in = t + model_out = self.apply_model(x, t_in, c, return_ids=return_codebook_ids) + + if score_corrector is not None: + assert self.parameterization == "eps" + model_out = score_corrector.modify_score(self, model_out, x, t, c, **corrector_kwargs) + + if return_codebook_ids: + model_out, logits = model_out + + if self.parameterization == "eps": + x_recon = self.predict_start_from_noise(x, t=t, noise=model_out) + elif self.parameterization == "x0": + x_recon = model_out + else: + raise NotImplementedError() + + if clip_denoised: + x_recon.clamp_(-1., 1.) + if quantize_denoised: + x_recon, _, [_, _, indices] = self.first_stage_model.quantize(x_recon) + model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t) + if return_codebook_ids: + return model_mean, posterior_variance, posterior_log_variance, logits + elif return_x0: + return model_mean, posterior_variance, posterior_log_variance, x_recon + else: + return model_mean, posterior_variance, posterior_log_variance + + @torch.no_grad() + def p_sample(self, x, c, t, clip_denoised=False, repeat_noise=False, + return_codebook_ids=False, quantize_denoised=False, return_x0=False, + temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None): + b, *_, device = *x.shape, x.device + outputs = self.p_mean_variance(x=x, c=c, t=t, clip_denoised=clip_denoised, + return_codebook_ids=return_codebook_ids, + quantize_denoised=quantize_denoised, + return_x0=return_x0, + score_corrector=score_corrector, corrector_kwargs=corrector_kwargs) + if return_codebook_ids: + raise DeprecationWarning("Support dropped.") + model_mean, _, model_log_variance, logits = outputs + elif return_x0: + model_mean, _, model_log_variance, x0 = outputs + else: + model_mean, _, model_log_variance = outputs + + noise = noise_like(x.shape, device, repeat_noise) * temperature + if noise_dropout > 0.: + noise = torch.nn.functional.dropout(noise, p=noise_dropout) + # no noise when t == 0 + nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1))) + + if return_codebook_ids: + return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, logits.argmax(dim=1) + if return_x0: + return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, x0 + else: + return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise + + @torch.no_grad() + def progressive_denoising(self, cond, shape, verbose=True, callback=None, quantize_denoised=False, + img_callback=None, mask=None, x0=None, temperature=1., noise_dropout=0., + score_corrector=None, corrector_kwargs=None, batch_size=None, x_T=None, start_T=None, + log_every_t=None): + if not log_every_t: + log_every_t = self.log_every_t + timesteps = self.num_timesteps + if batch_size is not None: + b = batch_size if batch_size is not None else shape[0] + shape = [batch_size] + list(shape) + else: + b = batch_size = shape[0] + if x_T is None: + img = torch.randn(shape, device=self.device) + else: + img = x_T + intermediates = [] + if cond is not None: + if isinstance(cond, dict): + cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else + [x[:batch_size] for x in cond[key]] for key in cond} + else: + cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size] + + if start_T is not None: + timesteps = min(timesteps, start_T) + iterator = tqdm(reversed(range(0, timesteps)), desc='Progressive Generation', + total=timesteps) if verbose else reversed( + range(0, timesteps)) + if type(temperature) == float: + temperature = [temperature] * timesteps + + for i in iterator: + ts = torch.full((b,), i, device=self.device, dtype=torch.long) + if self.shorten_cond_schedule: + assert self.model.conditioning_key != 'hybrid' + tc = self.cond_ids[ts].to(cond.device) + cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond)) + + img, x0_partial = self.p_sample(img, cond, ts, + clip_denoised=self.clip_denoised, + quantize_denoised=quantize_denoised, return_x0=True, + temperature=temperature[i], noise_dropout=noise_dropout, + score_corrector=score_corrector, corrector_kwargs=corrector_kwargs) + if mask is not None: + assert x0 is not None + img_orig = self.q_sample(x0, ts) + img = img_orig * mask + (1. - mask) * img + + if i % log_every_t == 0 or i == timesteps - 1: + intermediates.append(x0_partial) + if callback: + callback(i) + if img_callback: + img_callback(img, i) + return img, intermediates + + @torch.no_grad() + def p_sample_loop(self, cond, shape, return_intermediates=False, + x_T=None, verbose=True, callback=None, timesteps=None, quantize_denoised=False, + mask=None, x0=None, img_callback=None, start_T=None, + log_every_t=None): + + if not log_every_t: + log_every_t = self.log_every_t + device = self.betas.device + b = shape[0] + if x_T is None: + img = torch.randn(shape, device=device) + else: + img = x_T + + intermediates = [img] + if timesteps is None: + timesteps = self.num_timesteps + + if start_T is not None: + timesteps = min(timesteps, start_T) + iterator = tqdm(reversed(range(0, timesteps)), desc='Sampling t', total=timesteps) if verbose else reversed( + range(0, timesteps)) + + if mask is not None: + assert x0 is not None + assert x0.shape[2:3] == mask.shape[2:3] # spatial size has to match + + for i in iterator: + ts = torch.full((b,), i, device=device, dtype=torch.long) + if self.shorten_cond_schedule: + assert self.model.conditioning_key != 'hybrid' + tc = self.cond_ids[ts].to(cond.device) + cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond)) + + img = self.p_sample(img, cond, ts, + clip_denoised=self.clip_denoised, + quantize_denoised=quantize_denoised) + if mask is not None: + img_orig = self.q_sample(x0, ts) + img = img_orig * mask + (1. - mask) * img + + if i % log_every_t == 0 or i == timesteps - 1: + intermediates.append(img) + if callback: + callback(i) + if img_callback: + img_callback(img, i) + + if return_intermediates: + return img, intermediates + return img + + @torch.no_grad() + def sample(self, cond, batch_size=16, return_intermediates=False, x_T=None, + verbose=True, timesteps=None, quantize_denoised=False, + mask=None, x0=None, shape=None,**kwargs): + if shape is None: + shape = (batch_size, self.channels, self.image_size, self.image_size) + if cond is not None: + if isinstance(cond, dict): + cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else + [x[:batch_size] for x in cond[key]] for key in cond} + else: + cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size] + return self.p_sample_loop(cond, + shape, + return_intermediates=return_intermediates, x_T=x_T, + verbose=verbose, timesteps=timesteps, quantize_denoised=quantize_denoised, + mask=mask, x0=x0) + + @torch.no_grad() + def sample_log(self,cond,batch_size,ddim, ddim_steps,**kwargs): + + if ddim: + ddim_sampler = DDIMSampler(self) + shape = (self.channels, self.image_size, self.image_size) + samples, intermediates =ddim_sampler.sample(ddim_steps,batch_size, + shape,cond,verbose=False,**kwargs) + + else: + samples, intermediates = self.sample(cond=cond, batch_size=batch_size, + return_intermediates=True,**kwargs) + + return samples, intermediates + + + @torch.no_grad() + def log_images(self, batch, N=8, n_row=4, sample=True, ddim_steps=200, ddim_eta=1., return_keys=None, + quantize_denoised=True, inpaint=True, plot_denoise_rows=False, plot_progressive_rows=True, + plot_diffusion_rows=True, **kwargs): + + use_ddim = ddim_steps is not None + + log = {} + z, c, x, xrec, xc = self.get_input(batch, self.first_stage_key, + return_first_stage_outputs=True, + force_c_encode=True, + return_original_cond=True, + bs=N) + N = min(x.shape[0], N) + n_row = min(x.shape[0], n_row) + log["inputs"] = x + log["reconstruction"] = xrec + if self.model.conditioning_key is not None: + if hasattr(self.cond_stage_model, "decode"): + xc = self.cond_stage_model.decode(c) + log["conditioning"] = xc + elif self.cond_stage_key in ["caption"]: + xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["caption"]) + log["conditioning"] = xc + elif self.cond_stage_key == 'class_label': + xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["human_label"]) + log['conditioning'] = xc + elif isimage(xc): + log["conditioning"] = xc + if ismap(xc): + log["original_conditioning"] = self.to_rgb(xc) + + if plot_diffusion_rows: + # get diffusion row + diffusion_row = [] + z_start = z[:n_row] + for t in range(self.num_timesteps): + if t % self.log_every_t == 0 or t == self.num_timesteps - 1: + t = repeat(torch.tensor([t]), '1 -> b', b=n_row) + t = t.to(self.device).long() + noise = torch.randn_like(z_start) + z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise) + diffusion_row.append(self.decode_first_stage(z_noisy)) + + diffusion_row = torch.stack(diffusion_row) # n_log_step, n_row, C, H, W + diffusion_grid = rearrange(diffusion_row, 'n b c h w -> b n c h w') + diffusion_grid = rearrange(diffusion_grid, 'b n c h w -> (b n) c h w') + diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0]) + log["diffusion_row"] = diffusion_grid + + if sample: + # get denoise row + with self.ema_scope("Plotting"): + samples, z_denoise_row = self.sample_log(cond=c,batch_size=N,ddim=use_ddim, + ddim_steps=ddim_steps,eta=ddim_eta) + # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True) + x_samples = self.decode_first_stage(samples) + log["samples"] = x_samples + if plot_denoise_rows: + denoise_grid = self._get_denoise_row_from_list(z_denoise_row) + log["denoise_row"] = denoise_grid + + if quantize_denoised and not isinstance(self.first_stage_model, AutoencoderKL) and not isinstance( + self.first_stage_model, IdentityFirstStage): + # also display when quantizing x0 while sampling + with self.ema_scope("Plotting Quantized Denoised"): + samples, z_denoise_row = self.sample_log(cond=c,batch_size=N,ddim=use_ddim, + ddim_steps=ddim_steps,eta=ddim_eta, + quantize_denoised=True) + # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True, + # quantize_denoised=True) + x_samples = self.decode_first_stage(samples.to(self.device)) + log["samples_x0_quantized"] = x_samples + + if inpaint: + # make a simple center square + h, w = z.shape[2], z.shape[3] + mask = torch.ones(N, h, w).to(self.device) + # zeros will be filled in + mask[:, h // 4:3 * h // 4, w // 4:3 * w // 4] = 0. + mask = mask[:, None, ...] + with self.ema_scope("Plotting Inpaint"): + + samples, _ = self.sample_log(cond=c,batch_size=N,ddim=use_ddim, eta=ddim_eta, + ddim_steps=ddim_steps, x0=z[:N], mask=mask) + x_samples = self.decode_first_stage(samples.to(self.device)) + log["samples_inpainting"] = x_samples + log["mask"] = mask + + # outpaint + with self.ema_scope("Plotting Outpaint"): + samples, _ = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,eta=ddim_eta, + ddim_steps=ddim_steps, x0=z[:N], mask=mask) + x_samples = self.decode_first_stage(samples.to(self.device)) + log["samples_outpainting"] = x_samples + + if plot_progressive_rows: + with self.ema_scope("Plotting Progressives"): + img, progressives = self.progressive_denoising(c, + shape=(self.channels, self.image_size, self.image_size), + batch_size=N) + prog_row = self._get_denoise_row_from_list(progressives, desc="Progressive Generation") + log["progressive_row"] = prog_row + + if return_keys: + if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0: + return log + else: + return {key: log[key] for key in return_keys} + return log + + def configure_optimizers(self): + lr = self.learning_rate + params = list(self.model.parameters()) + if self.cond_stage_trainable: + print(f"{self.__class__.__name__}: Also optimizing conditioner params!") + params = params + list(self.cond_stage_model.parameters()) + if self.learn_logvar: + print('Diffusion model optimizing logvar') + params.append(self.logvar) + opt = torch.optim.AdamW(params, lr=lr) + if self.use_scheduler: + assert 'target' in self.scheduler_config + scheduler = instantiate_from_config(self.scheduler_config) + + print("Setting up LambdaLR scheduler...") + scheduler = [ + { + 'scheduler': LambdaLR(opt, lr_lambda=scheduler.schedule), + 'interval': 'step', + 'frequency': 1 + }] + return [opt], scheduler + return opt + + @torch.no_grad() + def to_rgb(self, x): + x = x.float() + if not hasattr(self, "colorize"): + self.colorize = torch.randn(3, x.shape[1], 1, 1).to(x) + x = nn.functional.conv2d(x, weight=self.colorize) + x = 2. * (x - x.min()) / (x.max() - x.min()) - 1. + return x + + +class DiffusionWrapperV1(pl.LightningModule): + def __init__(self, diff_model_config, conditioning_key): + super().__init__() + self.diffusion_model = instantiate_from_config(diff_model_config) + self.conditioning_key = conditioning_key + assert self.conditioning_key in [None, 'concat', 'crossattn', 'hybrid', 'adm'] + + def forward(self, x, t, c_concat: list = None, c_crossattn: list = None): + if self.conditioning_key is None: + out = self.diffusion_model(x, t) + elif self.conditioning_key == 'concat': + xc = torch.cat([x] + c_concat, dim=1) + out = self.diffusion_model(xc, t) + elif self.conditioning_key == 'crossattn': + cc = torch.cat(c_crossattn, 1) + out = self.diffusion_model(x, t, context=cc) + elif self.conditioning_key == 'hybrid': + xc = torch.cat([x] + c_concat, dim=1) + cc = torch.cat(c_crossattn, 1) + out = self.diffusion_model(xc, t, context=cc) + elif self.conditioning_key == 'adm': + cc = c_crossattn[0] + out = self.diffusion_model(x, t, y=cc) + else: + raise NotImplementedError() + + return out + + +class Layout2ImgDiffusionV1(LatentDiffusionV1): + # TODO: move all layout-specific hacks to this class + def __init__(self, cond_stage_key, *args, **kwargs): + assert cond_stage_key == 'coordinates_bbox', 'Layout2ImgDiffusion only for cond_stage_key="coordinates_bbox"' + super().__init__(*args, cond_stage_key=cond_stage_key, **kwargs) + + def log_images(self, batch, N=8, *args, **kwargs): + logs = super().log_images(*args, batch=batch, N=N, **kwargs) + + key = 'train' if self.training else 'validation' + dset = self.trainer.datamodule.datasets[key] + mapper = dset.conditional_builders[self.cond_stage_key] + + bbox_imgs = [] + map_fn = lambda catno: dset.get_textual_label(dset.get_category_id(catno)) + for tknzd_bbox in batch[self.cond_stage_key][:N]: + bboximg = mapper.plot(tknzd_bbox.detach().cpu(), map_fn, (256, 256)) + bbox_imgs.append(bboximg) + + cond_img = torch.stack(bbox_imgs, dim=0) + logs['bbox_image'] = cond_img + return logs + +ldm.models.diffusion.ddpm.DDPMV1 = DDPMV1 +ldm.models.diffusion.ddpm.LatentDiffusionV1 = LatentDiffusionV1 +ldm.models.diffusion.ddpm.DiffusionWrapperV1 = DiffusionWrapperV1 +ldm.models.diffusion.ddpm.Layout2ImgDiffusionV1 = Layout2ImgDiffusionV1 diff --git a/stable-diffusion-webui/extensions-builtin/LDSR/vqvae_quantize.py b/stable-diffusion-webui/extensions-builtin/LDSR/vqvae_quantize.py new file mode 100755 index 0000000..dd14b8f --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/LDSR/vqvae_quantize.py @@ -0,0 +1,147 @@ +# Vendored from https://raw.githubusercontent.com/CompVis/taming-transformers/24268930bf1dce879235a7fddd0b2355b84d7ea6/taming/modules/vqvae/quantize.py, +# where the license is as follows: +# +# Copyright (c) 2020 Patrick Esser and Robin Rombach and Björn Ommer +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR +# OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE +# OR OTHER DEALINGS IN THE SOFTWARE./ + +import torch +import torch.nn as nn +import numpy as np +from einops import rearrange + + +class VectorQuantizer2(nn.Module): + """ + Improved version over VectorQuantizer, can be used as a drop-in replacement. Mostly + avoids costly matrix multiplications and allows for post-hoc remapping of indices. + """ + + # NOTE: due to a bug the beta term was applied to the wrong term. for + # backwards compatibility we use the buggy version by default, but you can + # specify legacy=False to fix it. + def __init__(self, n_e, e_dim, beta, remap=None, unknown_index="random", + sane_index_shape=False, legacy=True): + super().__init__() + self.n_e = n_e + self.e_dim = e_dim + self.beta = beta + self.legacy = legacy + + self.embedding = nn.Embedding(self.n_e, self.e_dim) + self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e) + + self.remap = remap + if self.remap is not None: + self.register_buffer("used", torch.tensor(np.load(self.remap))) + self.re_embed = self.used.shape[0] + self.unknown_index = unknown_index # "random" or "extra" or integer + if self.unknown_index == "extra": + self.unknown_index = self.re_embed + self.re_embed = self.re_embed + 1 + print(f"Remapping {self.n_e} indices to {self.re_embed} indices. " + f"Using {self.unknown_index} for unknown indices.") + else: + self.re_embed = n_e + + self.sane_index_shape = sane_index_shape + + def remap_to_used(self, inds): + ishape = inds.shape + assert len(ishape) > 1 + inds = inds.reshape(ishape[0], -1) + used = self.used.to(inds) + match = (inds[:, :, None] == used[None, None, ...]).long() + new = match.argmax(-1) + unknown = match.sum(2) < 1 + if self.unknown_index == "random": + new[unknown] = torch.randint(0, self.re_embed, size=new[unknown].shape).to(device=new.device) + else: + new[unknown] = self.unknown_index + return new.reshape(ishape) + + def unmap_to_all(self, inds): + ishape = inds.shape + assert len(ishape) > 1 + inds = inds.reshape(ishape[0], -1) + used = self.used.to(inds) + if self.re_embed > self.used.shape[0]: # extra token + inds[inds >= self.used.shape[0]] = 0 # simply set to zero + back = torch.gather(used[None, :][inds.shape[0] * [0], :], 1, inds) + return back.reshape(ishape) + + def forward(self, z, temp=None, rescale_logits=False, return_logits=False): + assert temp is None or temp == 1.0, "Only for interface compatible with Gumbel" + assert rescale_logits is False, "Only for interface compatible with Gumbel" + assert return_logits is False, "Only for interface compatible with Gumbel" + # reshape z -> (batch, height, width, channel) and flatten + z = rearrange(z, 'b c h w -> b h w c').contiguous() + z_flattened = z.view(-1, self.e_dim) + # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z + + d = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \ + torch.sum(self.embedding.weight ** 2, dim=1) - 2 * \ + torch.einsum('bd,dn->bn', z_flattened, rearrange(self.embedding.weight, 'n d -> d n')) + + min_encoding_indices = torch.argmin(d, dim=1) + z_q = self.embedding(min_encoding_indices).view(z.shape) + perplexity = None + min_encodings = None + + # compute loss for embedding + if not self.legacy: + loss = self.beta * torch.mean((z_q.detach() - z) ** 2) + \ + torch.mean((z_q - z.detach()) ** 2) + else: + loss = torch.mean((z_q.detach() - z) ** 2) + self.beta * \ + torch.mean((z_q - z.detach()) ** 2) + + # preserve gradients + z_q = z + (z_q - z).detach() + + # reshape back to match original input shape + z_q = rearrange(z_q, 'b h w c -> b c h w').contiguous() + + if self.remap is not None: + min_encoding_indices = min_encoding_indices.reshape(z.shape[0], -1) # add batch axis + min_encoding_indices = self.remap_to_used(min_encoding_indices) + min_encoding_indices = min_encoding_indices.reshape(-1, 1) # flatten + + if self.sane_index_shape: + min_encoding_indices = min_encoding_indices.reshape( + z_q.shape[0], z_q.shape[2], z_q.shape[3]) + + return z_q, loss, (perplexity, min_encodings, min_encoding_indices) + + def get_codebook_entry(self, indices, shape): + # shape specifying (batch, height, width, channel) + if self.remap is not None: + indices = indices.reshape(shape[0], -1) # add batch axis + indices = self.unmap_to_all(indices) + indices = indices.reshape(-1) # flatten again + + # get quantized latent vectors + z_q = self.embedding(indices) + + if shape is not None: + z_q = z_q.view(shape) + # reshape back to match original input shape + z_q = z_q.permute(0, 3, 1, 2).contiguous() + + return z_q diff --git a/stable-diffusion-webui/extensions-builtin/Lora/extra_networks_lora.py b/stable-diffusion-webui/extensions-builtin/Lora/extra_networks_lora.py new file mode 100755 index 0000000..33edf46 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/Lora/extra_networks_lora.py @@ -0,0 +1,62 @@ +from modules import extra_networks, shared +import networks + + +class ExtraNetworkLora(extra_networks.ExtraNetwork): + def __init__(self): + super().__init__('lora') + + self.errors = {} + """mapping of network names to the number of errors the network had during operation""" + + remove_symbols = str.maketrans('', '', ":,") + + def activate(self, p, params_list): + additional = shared.opts.sd_lora + + self.errors.clear() + + if additional != "None" and additional in networks.available_networks and not any(x for x in params_list if x.items[0] == additional): + p.all_prompts = [x + f"" for x in p.all_prompts] + params_list.append(extra_networks.ExtraNetworkParams(items=[additional, shared.opts.extra_networks_default_multiplier])) + + names = [] + te_multipliers = [] + unet_multipliers = [] + dyn_dims = [] + for params in params_list: + assert params.items + + names.append(params.positional[0]) + + te_multiplier = float(params.positional[1]) if len(params.positional) > 1 else 1.0 + te_multiplier = float(params.named.get("te", te_multiplier)) + + unet_multiplier = float(params.positional[2]) if len(params.positional) > 2 else te_multiplier + unet_multiplier = float(params.named.get("unet", unet_multiplier)) + + dyn_dim = int(params.positional[3]) if len(params.positional) > 3 else None + dyn_dim = int(params.named["dyn"]) if "dyn" in params.named else dyn_dim + + te_multipliers.append(te_multiplier) + unet_multipliers.append(unet_multiplier) + dyn_dims.append(dyn_dim) + + networks.load_networks(names, te_multipliers, unet_multipliers, dyn_dims) + + if shared.opts.lora_add_hashes_to_infotext: + if not getattr(p, "is_hr_pass", False) or not hasattr(p, "lora_hashes"): + p.lora_hashes = {} + + for item in networks.loaded_networks: + if item.network_on_disk.shorthash and item.mentioned_name: + p.lora_hashes[item.mentioned_name.translate(self.remove_symbols)] = item.network_on_disk.shorthash + + if p.lora_hashes: + p.extra_generation_params["Lora hashes"] = ', '.join(f'{k}: {v}' for k, v in p.lora_hashes.items()) + + def deactivate(self, p): + if self.errors: + p.comment("Networks with errors: " + ", ".join(f"{k} ({v})" for k, v in self.errors.items())) + + self.errors.clear() diff --git a/stable-diffusion-webui/extensions-builtin/Lora/lora.py b/stable-diffusion-webui/extensions-builtin/Lora/lora.py new file mode 100755 index 0000000..6186538 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/Lora/lora.py @@ -0,0 +1,9 @@ +import networks + +list_available_loras = networks.list_available_networks + +available_loras = networks.available_networks +available_lora_aliases = networks.available_network_aliases +available_lora_hash_lookup = networks.available_network_hash_lookup +forbidden_lora_aliases = networks.forbidden_network_aliases +loaded_loras = networks.loaded_networks diff --git a/stable-diffusion-webui/extensions-builtin/Lora/lora_logger.py b/stable-diffusion-webui/extensions-builtin/Lora/lora_logger.py new file mode 100755 index 0000000..d51de29 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/Lora/lora_logger.py @@ -0,0 +1,33 @@ +import sys +import copy +import logging + + +class ColoredFormatter(logging.Formatter): + COLORS = { + "DEBUG": "\033[0;36m", # CYAN + "INFO": "\033[0;32m", # GREEN + "WARNING": "\033[0;33m", # YELLOW + "ERROR": "\033[0;31m", # RED + "CRITICAL": "\033[0;37;41m", # WHITE ON RED + "RESET": "\033[0m", # RESET COLOR + } + + def format(self, record): + colored_record = copy.copy(record) + levelname = colored_record.levelname + seq = self.COLORS.get(levelname, self.COLORS["RESET"]) + colored_record.levelname = f"{seq}{levelname}{self.COLORS['RESET']}" + return super().format(colored_record) + + +logger = logging.getLogger("lora") +logger.propagate = False + + +if not logger.handlers: + handler = logging.StreamHandler(sys.stdout) + handler.setFormatter( + ColoredFormatter("[%(name)s]-%(levelname)s: %(message)s") + ) + logger.addHandler(handler) diff --git a/stable-diffusion-webui/extensions-builtin/Lora/lora_patches.py b/stable-diffusion-webui/extensions-builtin/Lora/lora_patches.py new file mode 100755 index 0000000..59859e6 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/Lora/lora_patches.py @@ -0,0 +1,31 @@ +import torch + +import networks +from modules import patches + + +class LoraPatches: + def __init__(self): + self.Linear_forward = patches.patch(__name__, torch.nn.Linear, 'forward', networks.network_Linear_forward) + self.Linear_load_state_dict = patches.patch(__name__, torch.nn.Linear, '_load_from_state_dict', networks.network_Linear_load_state_dict) + self.Conv2d_forward = patches.patch(__name__, torch.nn.Conv2d, 'forward', networks.network_Conv2d_forward) + self.Conv2d_load_state_dict = patches.patch(__name__, torch.nn.Conv2d, '_load_from_state_dict', networks.network_Conv2d_load_state_dict) + self.GroupNorm_forward = patches.patch(__name__, torch.nn.GroupNorm, 'forward', networks.network_GroupNorm_forward) + self.GroupNorm_load_state_dict = patches.patch(__name__, torch.nn.GroupNorm, '_load_from_state_dict', networks.network_GroupNorm_load_state_dict) + self.LayerNorm_forward = patches.patch(__name__, torch.nn.LayerNorm, 'forward', networks.network_LayerNorm_forward) + self.LayerNorm_load_state_dict = patches.patch(__name__, torch.nn.LayerNorm, '_load_from_state_dict', networks.network_LayerNorm_load_state_dict) + self.MultiheadAttention_forward = patches.patch(__name__, torch.nn.MultiheadAttention, 'forward', networks.network_MultiheadAttention_forward) + self.MultiheadAttention_load_state_dict = patches.patch(__name__, torch.nn.MultiheadAttention, '_load_from_state_dict', networks.network_MultiheadAttention_load_state_dict) + + def undo(self): + self.Linear_forward = patches.undo(__name__, torch.nn.Linear, 'forward') + self.Linear_load_state_dict = patches.undo(__name__, torch.nn.Linear, '_load_from_state_dict') + self.Conv2d_forward = patches.undo(__name__, torch.nn.Conv2d, 'forward') + self.Conv2d_load_state_dict = patches.undo(__name__, torch.nn.Conv2d, '_load_from_state_dict') + self.GroupNorm_forward = patches.undo(__name__, torch.nn.GroupNorm, 'forward') + self.GroupNorm_load_state_dict = patches.undo(__name__, torch.nn.GroupNorm, '_load_from_state_dict') + self.LayerNorm_forward = patches.undo(__name__, torch.nn.LayerNorm, 'forward') + self.LayerNorm_load_state_dict = patches.undo(__name__, torch.nn.LayerNorm, '_load_from_state_dict') + self.MultiheadAttention_forward = patches.undo(__name__, torch.nn.MultiheadAttention, 'forward') + self.MultiheadAttention_load_state_dict = patches.undo(__name__, torch.nn.MultiheadAttention, '_load_from_state_dict') + diff --git a/stable-diffusion-webui/extensions-builtin/Lora/lyco_helpers.py b/stable-diffusion-webui/extensions-builtin/Lora/lyco_helpers.py new file mode 100755 index 0000000..3d4efd7 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/Lora/lyco_helpers.py @@ -0,0 +1,68 @@ +import torch + + +def make_weight_cp(t, wa, wb): + temp = torch.einsum('i j k l, j r -> i r k l', t, wb) + return torch.einsum('i j k l, i r -> r j k l', temp, wa) + + +def rebuild_conventional(up, down, shape, dyn_dim=None): + up = up.reshape(up.size(0), -1) + down = down.reshape(down.size(0), -1) + if dyn_dim is not None: + up = up[:, :dyn_dim] + down = down[:dyn_dim, :] + return (up @ down).reshape(shape) + + +def rebuild_cp_decomposition(up, down, mid): + up = up.reshape(up.size(0), -1) + down = down.reshape(down.size(0), -1) + return torch.einsum('n m k l, i n, m j -> i j k l', mid, up, down) + + +# copied from https://github.com/KohakuBlueleaf/LyCORIS/blob/dev/lycoris/modules/lokr.py +def factorization(dimension: int, factor:int=-1) -> tuple[int, int]: + ''' + return a tuple of two value of input dimension decomposed by the number closest to factor + second value is higher or equal than first value. + + In LoRA with Kroneckor Product, first value is a value for weight scale. + secon value is a value for weight. + + Because of non-commutative property, A⊗B ≠ B⊗A. Meaning of two matrices is slightly different. + + examples) + factor + -1 2 4 8 16 ... + 127 -> 1, 127 127 -> 1, 127 127 -> 1, 127 127 -> 1, 127 127 -> 1, 127 + 128 -> 8, 16 128 -> 2, 64 128 -> 4, 32 128 -> 8, 16 128 -> 8, 16 + 250 -> 10, 25 250 -> 2, 125 250 -> 2, 125 250 -> 5, 50 250 -> 10, 25 + 360 -> 8, 45 360 -> 2, 180 360 -> 4, 90 360 -> 8, 45 360 -> 12, 30 + 512 -> 16, 32 512 -> 2, 256 512 -> 4, 128 512 -> 8, 64 512 -> 16, 32 + 1024 -> 32, 32 1024 -> 2, 512 1024 -> 4, 256 1024 -> 8, 128 1024 -> 16, 64 + ''' + + if factor > 0 and (dimension % factor) == 0: + m = factor + n = dimension // factor + if m > n: + n, m = m, n + return m, n + if factor < 0: + factor = dimension + m, n = 1, dimension + length = m + n + while m length or new_m>factor: + break + else: + m, n = new_m, new_n + if m > n: + n, m = m, n + return m, n + diff --git a/stable-diffusion-webui/extensions-builtin/Lora/network.py b/stable-diffusion-webui/extensions-builtin/Lora/network.py new file mode 100755 index 0000000..8998743 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/Lora/network.py @@ -0,0 +1,228 @@ +from __future__ import annotations +import os +from collections import namedtuple +import enum + +import torch.nn as nn +import torch.nn.functional as F + +from modules import sd_models, cache, errors, hashes, shared +import modules.models.sd3.mmdit + +NetworkWeights = namedtuple('NetworkWeights', ['network_key', 'sd_key', 'w', 'sd_module']) + +metadata_tags_order = {"ss_sd_model_name": 1, "ss_resolution": 2, "ss_clip_skip": 3, "ss_num_train_images": 10, "ss_tag_frequency": 20} + + +class SdVersion(enum.Enum): + Unknown = 1 + SD1 = 2 + SD2 = 3 + SDXL = 4 + + +class NetworkOnDisk: + def __init__(self, name, filename): + self.name = name + self.filename = filename + self.metadata = {} + self.is_safetensors = os.path.splitext(filename)[1].lower() == ".safetensors" + + def read_metadata(): + metadata = sd_models.read_metadata_from_safetensors(filename) + + return metadata + + if self.is_safetensors: + try: + self.metadata = cache.cached_data_for_file('safetensors-metadata', "lora/" + self.name, filename, read_metadata) + except Exception as e: + errors.display(e, f"reading lora {filename}") + + if self.metadata: + m = {} + for k, v in sorted(self.metadata.items(), key=lambda x: metadata_tags_order.get(x[0], 999)): + m[k] = v + + self.metadata = m + + self.alias = self.metadata.get('ss_output_name', self.name) + + self.hash = None + self.shorthash = None + self.set_hash( + self.metadata.get('sshs_model_hash') or + hashes.sha256_from_cache(self.filename, "lora/" + self.name, use_addnet_hash=self.is_safetensors) or + '' + ) + + self.sd_version = self.detect_version() + + def detect_version(self): + if str(self.metadata.get('ss_base_model_version', "")).startswith("sdxl_"): + return SdVersion.SDXL + elif str(self.metadata.get('ss_v2', "")) == "True": + return SdVersion.SD2 + elif len(self.metadata): + return SdVersion.SD1 + + return SdVersion.Unknown + + def set_hash(self, v): + self.hash = v + self.shorthash = self.hash[0:12] + + if self.shorthash: + import networks + networks.available_network_hash_lookup[self.shorthash] = self + + def read_hash(self): + if not self.hash: + self.set_hash(hashes.sha256(self.filename, "lora/" + self.name, use_addnet_hash=self.is_safetensors) or '') + + def get_alias(self): + import networks + if shared.opts.lora_preferred_name == "Filename" or self.alias.lower() in networks.forbidden_network_aliases: + return self.name + else: + return self.alias + + +class Network: # LoraModule + def __init__(self, name, network_on_disk: NetworkOnDisk): + self.name = name + self.network_on_disk = network_on_disk + self.te_multiplier = 1.0 + self.unet_multiplier = 1.0 + self.dyn_dim = None + self.modules = {} + self.bundle_embeddings = {} + self.mtime = None + + self.mentioned_name = None + """the text that was used to add the network to prompt - can be either name or an alias""" + + +class ModuleType: + def create_module(self, net: Network, weights: NetworkWeights) -> Network | None: + return None + + +class NetworkModule: + def __init__(self, net: Network, weights: NetworkWeights): + self.network = net + self.network_key = weights.network_key + self.sd_key = weights.sd_key + self.sd_module = weights.sd_module + + if isinstance(self.sd_module, modules.models.sd3.mmdit.QkvLinear): + s = self.sd_module.weight.shape + self.shape = (s[0] // 3, s[1]) + elif hasattr(self.sd_module, 'weight'): + self.shape = self.sd_module.weight.shape + elif isinstance(self.sd_module, nn.MultiheadAttention): + # For now, only self-attn use Pytorch's MHA + # So assume all qkvo proj have same shape + self.shape = self.sd_module.out_proj.weight.shape + else: + self.shape = None + + self.ops = None + self.extra_kwargs = {} + if isinstance(self.sd_module, nn.Conv2d): + self.ops = F.conv2d + self.extra_kwargs = { + 'stride': self.sd_module.stride, + 'padding': self.sd_module.padding + } + elif isinstance(self.sd_module, nn.Linear): + self.ops = F.linear + elif isinstance(self.sd_module, nn.LayerNorm): + self.ops = F.layer_norm + self.extra_kwargs = { + 'normalized_shape': self.sd_module.normalized_shape, + 'eps': self.sd_module.eps + } + elif isinstance(self.sd_module, nn.GroupNorm): + self.ops = F.group_norm + self.extra_kwargs = { + 'num_groups': self.sd_module.num_groups, + 'eps': self.sd_module.eps + } + + self.dim = None + self.bias = weights.w.get("bias") + self.alpha = weights.w["alpha"].item() if "alpha" in weights.w else None + self.scale = weights.w["scale"].item() if "scale" in weights.w else None + + self.dora_scale = weights.w.get("dora_scale", None) + self.dora_norm_dims = len(self.shape) - 1 + + def multiplier(self): + if 'transformer' in self.sd_key[:20]: + return self.network.te_multiplier + else: + return self.network.unet_multiplier + + def calc_scale(self): + if self.scale is not None: + return self.scale + if self.dim is not None and self.alpha is not None: + return self.alpha / self.dim + + return 1.0 + + def apply_weight_decompose(self, updown, orig_weight): + # Match the device/dtype + orig_weight = orig_weight.to(updown.dtype) + dora_scale = self.dora_scale.to(device=orig_weight.device, dtype=updown.dtype) + updown = updown.to(orig_weight.device) + + merged_scale1 = updown + orig_weight + merged_scale1_norm = ( + merged_scale1.transpose(0, 1) + .reshape(merged_scale1.shape[1], -1) + .norm(dim=1, keepdim=True) + .reshape(merged_scale1.shape[1], *[1] * self.dora_norm_dims) + .transpose(0, 1) + ) + + dora_merged = ( + merged_scale1 * (dora_scale / merged_scale1_norm) + ) + final_updown = dora_merged - orig_weight + return final_updown + + def finalize_updown(self, updown, orig_weight, output_shape, ex_bias=None): + if self.bias is not None: + updown = updown.reshape(self.bias.shape) + updown += self.bias.to(orig_weight.device, dtype=updown.dtype) + updown = updown.reshape(output_shape) + + if len(output_shape) == 4: + updown = updown.reshape(output_shape) + + if orig_weight.size().numel() == updown.size().numel(): + updown = updown.reshape(orig_weight.shape) + + if ex_bias is not None: + ex_bias = ex_bias * self.multiplier() + + updown = updown * self.calc_scale() + + if self.dora_scale is not None: + updown = self.apply_weight_decompose(updown, orig_weight) + + return updown * self.multiplier(), ex_bias + + def calc_updown(self, target): + raise NotImplementedError() + + def forward(self, x, y): + """A general forward implementation for all modules""" + if self.ops is None: + raise NotImplementedError() + else: + updown, ex_bias = self.calc_updown(self.sd_module.weight) + return y + self.ops(x, weight=updown, bias=ex_bias, **self.extra_kwargs) + diff --git a/stable-diffusion-webui/extensions-builtin/Lora/network_full.py b/stable-diffusion-webui/extensions-builtin/Lora/network_full.py new file mode 100755 index 0000000..cf5fbbb --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/Lora/network_full.py @@ -0,0 +1,27 @@ +import network + + +class ModuleTypeFull(network.ModuleType): + def create_module(self, net: network.Network, weights: network.NetworkWeights): + if all(x in weights.w for x in ["diff"]): + return NetworkModuleFull(net, weights) + + return None + + +class NetworkModuleFull(network.NetworkModule): + def __init__(self, net: network.Network, weights: network.NetworkWeights): + super().__init__(net, weights) + + self.weight = weights.w.get("diff") + self.ex_bias = weights.w.get("diff_b") + + def calc_updown(self, orig_weight): + output_shape = self.weight.shape + updown = self.weight.to(orig_weight.device) + if self.ex_bias is not None: + ex_bias = self.ex_bias.to(orig_weight.device) + else: + ex_bias = None + + return self.finalize_updown(updown, orig_weight, output_shape, ex_bias) diff --git a/stable-diffusion-webui/extensions-builtin/Lora/network_glora.py b/stable-diffusion-webui/extensions-builtin/Lora/network_glora.py new file mode 100755 index 0000000..efe5c68 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/Lora/network_glora.py @@ -0,0 +1,33 @@ + +import network + +class ModuleTypeGLora(network.ModuleType): + def create_module(self, net: network.Network, weights: network.NetworkWeights): + if all(x in weights.w for x in ["a1.weight", "a2.weight", "alpha", "b1.weight", "b2.weight"]): + return NetworkModuleGLora(net, weights) + + return None + +# adapted from https://github.com/KohakuBlueleaf/LyCORIS +class NetworkModuleGLora(network.NetworkModule): + def __init__(self, net: network.Network, weights: network.NetworkWeights): + super().__init__(net, weights) + + if hasattr(self.sd_module, 'weight'): + self.shape = self.sd_module.weight.shape + + self.w1a = weights.w["a1.weight"] + self.w1b = weights.w["b1.weight"] + self.w2a = weights.w["a2.weight"] + self.w2b = weights.w["b2.weight"] + + def calc_updown(self, orig_weight): + w1a = self.w1a.to(orig_weight.device) + w1b = self.w1b.to(orig_weight.device) + w2a = self.w2a.to(orig_weight.device) + w2b = self.w2b.to(orig_weight.device) + + output_shape = [w1a.size(0), w1b.size(1)] + updown = ((w2b @ w1b) + ((orig_weight.to(dtype = w1a.dtype) @ w2a) @ w1a)) + + return self.finalize_updown(updown, orig_weight, output_shape) diff --git a/stable-diffusion-webui/extensions-builtin/Lora/network_hada.py b/stable-diffusion-webui/extensions-builtin/Lora/network_hada.py new file mode 100755 index 0000000..d179b29 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/Lora/network_hada.py @@ -0,0 +1,55 @@ +import lyco_helpers +import network + + +class ModuleTypeHada(network.ModuleType): + def create_module(self, net: network.Network, weights: network.NetworkWeights): + if all(x in weights.w for x in ["hada_w1_a", "hada_w1_b", "hada_w2_a", "hada_w2_b"]): + return NetworkModuleHada(net, weights) + + return None + + +class NetworkModuleHada(network.NetworkModule): + def __init__(self, net: network.Network, weights: network.NetworkWeights): + super().__init__(net, weights) + + if hasattr(self.sd_module, 'weight'): + self.shape = self.sd_module.weight.shape + + self.w1a = weights.w["hada_w1_a"] + self.w1b = weights.w["hada_w1_b"] + self.dim = self.w1b.shape[0] + self.w2a = weights.w["hada_w2_a"] + self.w2b = weights.w["hada_w2_b"] + + self.t1 = weights.w.get("hada_t1") + self.t2 = weights.w.get("hada_t2") + + def calc_updown(self, orig_weight): + w1a = self.w1a.to(orig_weight.device) + w1b = self.w1b.to(orig_weight.device) + w2a = self.w2a.to(orig_weight.device) + w2b = self.w2b.to(orig_weight.device) + + output_shape = [w1a.size(0), w1b.size(1)] + + if self.t1 is not None: + output_shape = [w1a.size(1), w1b.size(1)] + t1 = self.t1.to(orig_weight.device) + updown1 = lyco_helpers.make_weight_cp(t1, w1a, w1b) + output_shape += t1.shape[2:] + else: + if len(w1b.shape) == 4: + output_shape += w1b.shape[2:] + updown1 = lyco_helpers.rebuild_conventional(w1a, w1b, output_shape) + + if self.t2 is not None: + t2 = self.t2.to(orig_weight.device) + updown2 = lyco_helpers.make_weight_cp(t2, w2a, w2b) + else: + updown2 = lyco_helpers.rebuild_conventional(w2a, w2b, output_shape) + + updown = updown1 * updown2 + + return self.finalize_updown(updown, orig_weight, output_shape) diff --git a/stable-diffusion-webui/extensions-builtin/Lora/network_ia3.py b/stable-diffusion-webui/extensions-builtin/Lora/network_ia3.py new file mode 100755 index 0000000..549b9a7 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/Lora/network_ia3.py @@ -0,0 +1,30 @@ +import network + + +class ModuleTypeIa3(network.ModuleType): + def create_module(self, net: network.Network, weights: network.NetworkWeights): + if all(x in weights.w for x in ["weight"]): + return NetworkModuleIa3(net, weights) + + return None + + +class NetworkModuleIa3(network.NetworkModule): + def __init__(self, net: network.Network, weights: network.NetworkWeights): + super().__init__(net, weights) + + self.w = weights.w["weight"] + self.on_input = weights.w["on_input"].item() + + def calc_updown(self, orig_weight): + w = self.w.to(orig_weight.device) + + output_shape = [w.size(0), orig_weight.size(1)] + if self.on_input: + output_shape.reverse() + else: + w = w.reshape(-1, 1) + + updown = orig_weight * w + + return self.finalize_updown(updown, orig_weight, output_shape) diff --git a/stable-diffusion-webui/extensions-builtin/Lora/network_lokr.py b/stable-diffusion-webui/extensions-builtin/Lora/network_lokr.py new file mode 100755 index 0000000..4f3128e --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/Lora/network_lokr.py @@ -0,0 +1,64 @@ +import torch + +import lyco_helpers +import network + + +class ModuleTypeLokr(network.ModuleType): + def create_module(self, net: network.Network, weights: network.NetworkWeights): + has_1 = "lokr_w1" in weights.w or ("lokr_w1_a" in weights.w and "lokr_w1_b" in weights.w) + has_2 = "lokr_w2" in weights.w or ("lokr_w2_a" in weights.w and "lokr_w2_b" in weights.w) + if has_1 and has_2: + return NetworkModuleLokr(net, weights) + + return None + + +def make_kron(orig_shape, w1, w2): + if len(w2.shape) == 4: + w1 = w1.unsqueeze(2).unsqueeze(2) + w2 = w2.contiguous() + return torch.kron(w1, w2).reshape(orig_shape) + + +class NetworkModuleLokr(network.NetworkModule): + def __init__(self, net: network.Network, weights: network.NetworkWeights): + super().__init__(net, weights) + + self.w1 = weights.w.get("lokr_w1") + self.w1a = weights.w.get("lokr_w1_a") + self.w1b = weights.w.get("lokr_w1_b") + self.dim = self.w1b.shape[0] if self.w1b is not None else self.dim + self.w2 = weights.w.get("lokr_w2") + self.w2a = weights.w.get("lokr_w2_a") + self.w2b = weights.w.get("lokr_w2_b") + self.dim = self.w2b.shape[0] if self.w2b is not None else self.dim + self.t2 = weights.w.get("lokr_t2") + + def calc_updown(self, orig_weight): + if self.w1 is not None: + w1 = self.w1.to(orig_weight.device) + else: + w1a = self.w1a.to(orig_weight.device) + w1b = self.w1b.to(orig_weight.device) + w1 = w1a @ w1b + + if self.w2 is not None: + w2 = self.w2.to(orig_weight.device) + elif self.t2 is None: + w2a = self.w2a.to(orig_weight.device) + w2b = self.w2b.to(orig_weight.device) + w2 = w2a @ w2b + else: + t2 = self.t2.to(orig_weight.device) + w2a = self.w2a.to(orig_weight.device) + w2b = self.w2b.to(orig_weight.device) + w2 = lyco_helpers.make_weight_cp(t2, w2a, w2b) + + output_shape = [w1.size(0) * w2.size(0), w1.size(1) * w2.size(1)] + if len(orig_weight.shape) == 4: + output_shape = orig_weight.shape + + updown = make_kron(output_shape, w1, w2) + + return self.finalize_updown(updown, orig_weight, output_shape) diff --git a/stable-diffusion-webui/extensions-builtin/Lora/network_lora.py b/stable-diffusion-webui/extensions-builtin/Lora/network_lora.py new file mode 100755 index 0000000..8ee26c3 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/Lora/network_lora.py @@ -0,0 +1,94 @@ +import torch + +import lyco_helpers +import modules.models.sd3.mmdit +import network +from modules import devices + + +class ModuleTypeLora(network.ModuleType): + def create_module(self, net: network.Network, weights: network.NetworkWeights): + if all(x in weights.w for x in ["lora_up.weight", "lora_down.weight"]): + return NetworkModuleLora(net, weights) + + if all(x in weights.w for x in ["lora_A.weight", "lora_B.weight"]): + w = weights.w.copy() + weights.w.clear() + weights.w.update({"lora_up.weight": w["lora_B.weight"], "lora_down.weight": w["lora_A.weight"]}) + + return NetworkModuleLora(net, weights) + + return None + + +class NetworkModuleLora(network.NetworkModule): + def __init__(self, net: network.Network, weights: network.NetworkWeights): + super().__init__(net, weights) + + self.up_model = self.create_module(weights.w, "lora_up.weight") + self.down_model = self.create_module(weights.w, "lora_down.weight") + self.mid_model = self.create_module(weights.w, "lora_mid.weight", none_ok=True) + + self.dim = weights.w["lora_down.weight"].shape[0] + + def create_module(self, weights, key, none_ok=False): + weight = weights.get(key) + + if weight is None and none_ok: + return None + + is_linear = type(self.sd_module) in [torch.nn.Linear, torch.nn.modules.linear.NonDynamicallyQuantizableLinear, torch.nn.MultiheadAttention, modules.models.sd3.mmdit.QkvLinear] + is_conv = type(self.sd_module) in [torch.nn.Conv2d] + + if is_linear: + weight = weight.reshape(weight.shape[0], -1) + module = torch.nn.Linear(weight.shape[1], weight.shape[0], bias=False) + elif is_conv and key == "lora_down.weight" or key == "dyn_up": + if len(weight.shape) == 2: + weight = weight.reshape(weight.shape[0], -1, 1, 1) + + if weight.shape[2] != 1 or weight.shape[3] != 1: + module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], self.sd_module.kernel_size, self.sd_module.stride, self.sd_module.padding, bias=False) + else: + module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], (1, 1), bias=False) + elif is_conv and key == "lora_mid.weight": + module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], self.sd_module.kernel_size, self.sd_module.stride, self.sd_module.padding, bias=False) + elif is_conv and key == "lora_up.weight" or key == "dyn_down": + module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], (1, 1), bias=False) + else: + raise AssertionError(f'Lora layer {self.network_key} matched a layer with unsupported type: {type(self.sd_module).__name__}') + + with torch.no_grad(): + if weight.shape != module.weight.shape: + weight = weight.reshape(module.weight.shape) + module.weight.copy_(weight) + + module.to(device=devices.cpu, dtype=devices.dtype) + module.weight.requires_grad_(False) + + return module + + def calc_updown(self, orig_weight): + up = self.up_model.weight.to(orig_weight.device) + down = self.down_model.weight.to(orig_weight.device) + + output_shape = [up.size(0), down.size(1)] + if self.mid_model is not None: + # cp-decomposition + mid = self.mid_model.weight.to(orig_weight.device) + updown = lyco_helpers.rebuild_cp_decomposition(up, down, mid) + output_shape += mid.shape[2:] + else: + if len(down.shape) == 4: + output_shape += down.shape[2:] + updown = lyco_helpers.rebuild_conventional(up, down, output_shape, self.network.dyn_dim) + + return self.finalize_updown(updown, orig_weight, output_shape) + + def forward(self, x, y): + self.up_model.to(device=devices.device) + self.down_model.to(device=devices.device) + + return y + self.up_model(self.down_model(x)) * self.multiplier() * self.calc_scale() + + diff --git a/stable-diffusion-webui/extensions-builtin/Lora/network_norm.py b/stable-diffusion-webui/extensions-builtin/Lora/network_norm.py new file mode 100755 index 0000000..d25afcb --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/Lora/network_norm.py @@ -0,0 +1,28 @@ +import network + + +class ModuleTypeNorm(network.ModuleType): + def create_module(self, net: network.Network, weights: network.NetworkWeights): + if all(x in weights.w for x in ["w_norm", "b_norm"]): + return NetworkModuleNorm(net, weights) + + return None + + +class NetworkModuleNorm(network.NetworkModule): + def __init__(self, net: network.Network, weights: network.NetworkWeights): + super().__init__(net, weights) + + self.w_norm = weights.w.get("w_norm") + self.b_norm = weights.w.get("b_norm") + + def calc_updown(self, orig_weight): + output_shape = self.w_norm.shape + updown = self.w_norm.to(orig_weight.device) + + if self.b_norm is not None: + ex_bias = self.b_norm.to(orig_weight.device) + else: + ex_bias = None + + return self.finalize_updown(updown, orig_weight, output_shape, ex_bias) diff --git a/stable-diffusion-webui/extensions-builtin/Lora/network_oft.py b/stable-diffusion-webui/extensions-builtin/Lora/network_oft.py new file mode 100755 index 0000000..1c515eb --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/Lora/network_oft.py @@ -0,0 +1,118 @@ +import torch +import network +from einops import rearrange + + +class ModuleTypeOFT(network.ModuleType): + def create_module(self, net: network.Network, weights: network.NetworkWeights): + if all(x in weights.w for x in ["oft_blocks"]) or all(x in weights.w for x in ["oft_diag"]): + return NetworkModuleOFT(net, weights) + + return None + +# Supports both kohya-ss' implementation of COFT https://github.com/kohya-ss/sd-scripts/blob/main/networks/oft.py +# and KohakuBlueleaf's implementation of OFT/COFT https://github.com/KohakuBlueleaf/LyCORIS/blob/dev/lycoris/modules/diag_oft.py +class NetworkModuleOFT(network.NetworkModule): + def __init__(self, net: network.Network, weights: network.NetworkWeights): + + super().__init__(net, weights) + + self.lin_module = None + self.org_module: list[torch.Module] = [self.sd_module] + + self.scale = 1.0 + self.is_R = False + self.is_boft = False + + # kohya-ss/New LyCORIS OFT/BOFT + if "oft_blocks" in weights.w.keys(): + self.oft_blocks = weights.w["oft_blocks"] # (num_blocks, block_size, block_size) + self.alpha = weights.w.get("alpha", None) # alpha is constraint + self.dim = self.oft_blocks.shape[0] # lora dim + # Old LyCORIS OFT + elif "oft_diag" in weights.w.keys(): + self.is_R = True + self.oft_blocks = weights.w["oft_diag"] + # self.alpha is unused + self.dim = self.oft_blocks.shape[1] # (num_blocks, block_size, block_size) + + is_linear = type(self.sd_module) in [torch.nn.Linear, torch.nn.modules.linear.NonDynamicallyQuantizableLinear] + is_conv = type(self.sd_module) in [torch.nn.Conv2d] + is_other_linear = type(self.sd_module) in [torch.nn.MultiheadAttention] # unsupported + + if is_linear: + self.out_dim = self.sd_module.out_features + elif is_conv: + self.out_dim = self.sd_module.out_channels + elif is_other_linear: + self.out_dim = self.sd_module.embed_dim + + # LyCORIS BOFT + if self.oft_blocks.dim() == 4: + self.is_boft = True + self.rescale = weights.w.get('rescale', None) + if self.rescale is not None and not is_other_linear: + self.rescale = self.rescale.reshape(-1, *[1]*(self.org_module[0].weight.dim() - 1)) + + self.num_blocks = self.dim + self.block_size = self.out_dim // self.dim + self.constraint = (0 if self.alpha is None else self.alpha) * self.out_dim + if self.is_R: + self.constraint = None + self.block_size = self.dim + self.num_blocks = self.out_dim // self.dim + elif self.is_boft: + self.boft_m = self.oft_blocks.shape[0] + self.num_blocks = self.oft_blocks.shape[1] + self.block_size = self.oft_blocks.shape[2] + self.boft_b = self.block_size + + def calc_updown(self, orig_weight): + oft_blocks = self.oft_blocks.to(orig_weight.device) + eye = torch.eye(self.block_size, device=oft_blocks.device) + + if not self.is_R: + block_Q = oft_blocks - oft_blocks.transpose(-1, -2) # ensure skew-symmetric orthogonal matrix + if self.constraint != 0: + norm_Q = torch.norm(block_Q.flatten()) + new_norm_Q = torch.clamp(norm_Q, max=self.constraint.to(oft_blocks.device)) + block_Q = block_Q * ((new_norm_Q + 1e-8) / (norm_Q + 1e-8)) + oft_blocks = torch.matmul(eye + block_Q, (eye - block_Q).float().inverse()) + + R = oft_blocks.to(orig_weight.device) + + if not self.is_boft: + # This errors out for MultiheadAttention, might need to be handled up-stream + merged_weight = rearrange(orig_weight, '(k n) ... -> k n ...', k=self.num_blocks, n=self.block_size) + merged_weight = torch.einsum( + 'k n m, k n ... -> k m ...', + R, + merged_weight + ) + merged_weight = rearrange(merged_weight, 'k m ... -> (k m) ...') + else: + # TODO: determine correct value for scale + scale = 1.0 + m = self.boft_m + b = self.boft_b + r_b = b // 2 + inp = orig_weight + for i in range(m): + bi = R[i] # b_num, b_size, b_size + if i == 0: + # Apply multiplier/scale and rescale into first weight + bi = bi * scale + (1 - scale) * eye + inp = rearrange(inp, "(c g k) ... -> (c k g) ...", g=2, k=2**i * r_b) + inp = rearrange(inp, "(d b) ... -> d b ...", b=b) + inp = torch.einsum("b i j, b j ... -> b i ...", bi, inp) + inp = rearrange(inp, "d b ... -> (d b) ...") + inp = rearrange(inp, "(c k g) ... -> (c g k) ...", g=2, k=2**i * r_b) + merged_weight = inp + + # Rescale mechanism + if self.rescale is not None: + merged_weight = self.rescale.to(merged_weight) * merged_weight + + updown = merged_weight.to(orig_weight.device) - orig_weight.to(merged_weight.dtype) + output_shape = orig_weight.shape + return self.finalize_updown(updown, orig_weight, output_shape) diff --git a/stable-diffusion-webui/extensions-builtin/Lora/networks.py b/stable-diffusion-webui/extensions-builtin/Lora/networks.py new file mode 100755 index 0000000..72efdaa --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/Lora/networks.py @@ -0,0 +1,737 @@ +from __future__ import annotations +import gradio as gr +import logging +import os +import re + +import lora_patches +import network +import network_lora +import network_glora +import network_hada +import network_ia3 +import network_lokr +import network_full +import network_norm +import network_oft + +import torch +from typing import Union + +from modules import shared, devices, sd_models, errors, scripts, sd_hijack +import modules.textual_inversion.textual_inversion as textual_inversion +import modules.models.sd3.mmdit + +from lora_logger import logger + +module_types = [ + network_lora.ModuleTypeLora(), + network_hada.ModuleTypeHada(), + network_ia3.ModuleTypeIa3(), + network_lokr.ModuleTypeLokr(), + network_full.ModuleTypeFull(), + network_norm.ModuleTypeNorm(), + network_glora.ModuleTypeGLora(), + network_oft.ModuleTypeOFT(), +] + + +re_digits = re.compile(r"\d+") +re_x_proj = re.compile(r"(.*)_([qkv]_proj)$") +re_compiled = {} + +suffix_conversion = { + "attentions": {}, + "resnets": { + "conv1": "in_layers_2", + "conv2": "out_layers_3", + "norm1": "in_layers_0", + "norm2": "out_layers_0", + "time_emb_proj": "emb_layers_1", + "conv_shortcut": "skip_connection", + } +} + + +def convert_diffusers_name_to_compvis(key, is_sd2): + def match(match_list, regex_text): + regex = re_compiled.get(regex_text) + if regex is None: + regex = re.compile(regex_text) + re_compiled[regex_text] = regex + + r = re.match(regex, key) + if not r: + return False + + match_list.clear() + match_list.extend([int(x) if re.match(re_digits, x) else x for x in r.groups()]) + return True + + m = [] + + if match(m, r"lora_unet_conv_in(.*)"): + return f'diffusion_model_input_blocks_0_0{m[0]}' + + if match(m, r"lora_unet_conv_out(.*)"): + return f'diffusion_model_out_2{m[0]}' + + if match(m, r"lora_unet_time_embedding_linear_(\d+)(.*)"): + return f"diffusion_model_time_embed_{m[0] * 2 - 2}{m[1]}" + + if match(m, r"lora_unet_down_blocks_(\d+)_(attentions|resnets)_(\d+)_(.+)"): + suffix = suffix_conversion.get(m[1], {}).get(m[3], m[3]) + return f"diffusion_model_input_blocks_{1 + m[0] * 3 + m[2]}_{1 if m[1] == 'attentions' else 0}_{suffix}" + + if match(m, r"lora_unet_mid_block_(attentions|resnets)_(\d+)_(.+)"): + suffix = suffix_conversion.get(m[0], {}).get(m[2], m[2]) + return f"diffusion_model_middle_block_{1 if m[0] == 'attentions' else m[1] * 2}_{suffix}" + + if match(m, r"lora_unet_up_blocks_(\d+)_(attentions|resnets)_(\d+)_(.+)"): + suffix = suffix_conversion.get(m[1], {}).get(m[3], m[3]) + return f"diffusion_model_output_blocks_{m[0] * 3 + m[2]}_{1 if m[1] == 'attentions' else 0}_{suffix}" + + if match(m, r"lora_unet_down_blocks_(\d+)_downsamplers_0_conv"): + return f"diffusion_model_input_blocks_{3 + m[0] * 3}_0_op" + + if match(m, r"lora_unet_up_blocks_(\d+)_upsamplers_0_conv"): + return f"diffusion_model_output_blocks_{2 + m[0] * 3}_{2 if m[0]>0 else 1}_conv" + + if match(m, r"lora_te_text_model_encoder_layers_(\d+)_(.+)"): + if is_sd2: + if 'mlp_fc1' in m[1]: + return f"model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc1', 'mlp_c_fc')}" + elif 'mlp_fc2' in m[1]: + return f"model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc2', 'mlp_c_proj')}" + else: + return f"model_transformer_resblocks_{m[0]}_{m[1].replace('self_attn', 'attn')}" + + return f"transformer_text_model_encoder_layers_{m[0]}_{m[1]}" + + if match(m, r"lora_te2_text_model_encoder_layers_(\d+)_(.+)"): + if 'mlp_fc1' in m[1]: + return f"1_model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc1', 'mlp_c_fc')}" + elif 'mlp_fc2' in m[1]: + return f"1_model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc2', 'mlp_c_proj')}" + else: + return f"1_model_transformer_resblocks_{m[0]}_{m[1].replace('self_attn', 'attn')}" + + return key + + +def assign_network_names_to_compvis_modules(sd_model): + network_layer_mapping = {} + + if shared.sd_model.is_sdxl: + for i, embedder in enumerate(shared.sd_model.conditioner.embedders): + if not hasattr(embedder, 'wrapped'): + continue + + for name, module in embedder.wrapped.named_modules(): + network_name = f'{i}_{name.replace(".", "_")}' + network_layer_mapping[network_name] = module + module.network_layer_name = network_name + else: + cond_stage_model = getattr(shared.sd_model.cond_stage_model, 'wrapped', shared.sd_model.cond_stage_model) + + for name, module in cond_stage_model.named_modules(): + network_name = name.replace(".", "_") + network_layer_mapping[network_name] = module + module.network_layer_name = network_name + + for name, module in shared.sd_model.model.named_modules(): + network_name = name.replace(".", "_") + network_layer_mapping[network_name] = module + module.network_layer_name = network_name + + sd_model.network_layer_mapping = network_layer_mapping + + +class BundledTIHash(str): + def __init__(self, hash_str): + self.hash = hash_str + + def __str__(self): + return self.hash if shared.opts.lora_bundled_ti_to_infotext else '' + + +def load_network(name, network_on_disk): + net = network.Network(name, network_on_disk) + net.mtime = os.path.getmtime(network_on_disk.filename) + + sd = sd_models.read_state_dict(network_on_disk.filename) + + # this should not be needed but is here as an emergency fix for an unknown error people are experiencing in 1.2.0 + if not hasattr(shared.sd_model, 'network_layer_mapping'): + assign_network_names_to_compvis_modules(shared.sd_model) + + keys_failed_to_match = {} + is_sd2 = 'model_transformer_resblocks' in shared.sd_model.network_layer_mapping + if hasattr(shared.sd_model, 'diffusers_weight_map'): + diffusers_weight_map = shared.sd_model.diffusers_weight_map + elif hasattr(shared.sd_model, 'diffusers_weight_mapping'): + diffusers_weight_map = {} + for k, v in shared.sd_model.diffusers_weight_mapping(): + diffusers_weight_map[k] = v + shared.sd_model.diffusers_weight_map = diffusers_weight_map + else: + diffusers_weight_map = None + + matched_networks = {} + bundle_embeddings = {} + + for key_network, weight in sd.items(): + + if diffusers_weight_map: + key_network_without_network_parts, network_name, network_weight = key_network.rsplit(".", 2) + network_part = network_name + '.' + network_weight + else: + key_network_without_network_parts, _, network_part = key_network.partition(".") + + if key_network_without_network_parts == "bundle_emb": + emb_name, vec_name = network_part.split(".", 1) + emb_dict = bundle_embeddings.get(emb_name, {}) + if vec_name.split('.')[0] == 'string_to_param': + _, k2 = vec_name.split('.', 1) + emb_dict['string_to_param'] = {k2: weight} + else: + emb_dict[vec_name] = weight + bundle_embeddings[emb_name] = emb_dict + + if diffusers_weight_map: + key = diffusers_weight_map.get(key_network_without_network_parts, key_network_without_network_parts) + else: + key = convert_diffusers_name_to_compvis(key_network_without_network_parts, is_sd2) + + sd_module = shared.sd_model.network_layer_mapping.get(key, None) + + if sd_module is None: + m = re_x_proj.match(key) + if m: + sd_module = shared.sd_model.network_layer_mapping.get(m.group(1), None) + + # SDXL loras seem to already have correct compvis keys, so only need to replace "lora_unet" with "diffusion_model" + if sd_module is None and "lora_unet" in key_network_without_network_parts: + key = key_network_without_network_parts.replace("lora_unet", "diffusion_model") + sd_module = shared.sd_model.network_layer_mapping.get(key, None) + elif sd_module is None and "lora_te1_text_model" in key_network_without_network_parts: + key = key_network_without_network_parts.replace("lora_te1_text_model", "0_transformer_text_model") + sd_module = shared.sd_model.network_layer_mapping.get(key, None) + + # some SD1 Loras also have correct compvis keys + if sd_module is None: + key = key_network_without_network_parts.replace("lora_te1_text_model", "transformer_text_model") + sd_module = shared.sd_model.network_layer_mapping.get(key, None) + + # kohya_ss OFT module + elif sd_module is None and "oft_unet" in key_network_without_network_parts: + key = key_network_without_network_parts.replace("oft_unet", "diffusion_model") + sd_module = shared.sd_model.network_layer_mapping.get(key, None) + + # KohakuBlueLeaf OFT module + if sd_module is None and "oft_diag" in key: + key = key_network_without_network_parts.replace("lora_unet", "diffusion_model") + key = key_network_without_network_parts.replace("lora_te1_text_model", "0_transformer_text_model") + sd_module = shared.sd_model.network_layer_mapping.get(key, None) + + if sd_module is None: + keys_failed_to_match[key_network] = key + continue + + if key not in matched_networks: + matched_networks[key] = network.NetworkWeights(network_key=key_network, sd_key=key, w={}, sd_module=sd_module) + + matched_networks[key].w[network_part] = weight + + for key, weights in matched_networks.items(): + net_module = None + for nettype in module_types: + net_module = nettype.create_module(net, weights) + if net_module is not None: + break + + if net_module is None: + raise AssertionError(f"Could not find a module type (out of {', '.join([x.__class__.__name__ for x in module_types])}) that would accept those keys: {', '.join(weights.w)}") + + net.modules[key] = net_module + + embeddings = {} + for emb_name, data in bundle_embeddings.items(): + embedding = textual_inversion.create_embedding_from_data(data, emb_name, filename=network_on_disk.filename + "/" + emb_name) + embedding.loaded = None + embedding.shorthash = BundledTIHash(name) + embeddings[emb_name] = embedding + + net.bundle_embeddings = embeddings + + if keys_failed_to_match: + logging.debug(f"Network {network_on_disk.filename} didn't match keys: {keys_failed_to_match}") + + return net + + +def purge_networks_from_memory(): + while len(networks_in_memory) > shared.opts.lora_in_memory_limit and len(networks_in_memory) > 0: + name = next(iter(networks_in_memory)) + networks_in_memory.pop(name, None) + + devices.torch_gc() + + +def load_networks(names, te_multipliers=None, unet_multipliers=None, dyn_dims=None): + emb_db = sd_hijack.model_hijack.embedding_db + already_loaded = {} + + for net in loaded_networks: + if net.name in names: + already_loaded[net.name] = net + for emb_name, embedding in net.bundle_embeddings.items(): + if embedding.loaded: + emb_db.register_embedding_by_name(None, shared.sd_model, emb_name) + + loaded_networks.clear() + + unavailable_networks = [] + for name in names: + if name.lower() in forbidden_network_aliases and available_networks.get(name) is None: + unavailable_networks.append(name) + elif available_network_aliases.get(name) is None: + unavailable_networks.append(name) + + if unavailable_networks: + update_available_networks_by_names(unavailable_networks) + + networks_on_disk = [available_networks.get(name, None) if name.lower() in forbidden_network_aliases else available_network_aliases.get(name, None) for name in names] + if any(x is None for x in networks_on_disk): + list_available_networks() + + networks_on_disk = [available_networks.get(name, None) if name.lower() in forbidden_network_aliases else available_network_aliases.get(name, None) for name in names] + + failed_to_load_networks = [] + + for i, (network_on_disk, name) in enumerate(zip(networks_on_disk, names)): + net = already_loaded.get(name, None) + + if network_on_disk is not None: + if net is None: + net = networks_in_memory.get(name) + + if net is None or os.path.getmtime(network_on_disk.filename) > net.mtime: + try: + net = load_network(name, network_on_disk) + + networks_in_memory.pop(name, None) + networks_in_memory[name] = net + except Exception as e: + errors.display(e, f"loading network {network_on_disk.filename}") + continue + + net.mentioned_name = name + + network_on_disk.read_hash() + + if net is None: + failed_to_load_networks.append(name) + logging.info(f"Couldn't find network with name {name}") + continue + + net.te_multiplier = te_multipliers[i] if te_multipliers else 1.0 + net.unet_multiplier = unet_multipliers[i] if unet_multipliers else 1.0 + net.dyn_dim = dyn_dims[i] if dyn_dims else 1.0 + loaded_networks.append(net) + + for emb_name, embedding in net.bundle_embeddings.items(): + if embedding.loaded is None and emb_name in emb_db.word_embeddings: + logger.warning( + f'Skip bundle embedding: "{emb_name}"' + ' as it was already loaded from embeddings folder' + ) + continue + + embedding.loaded = False + if emb_db.expected_shape == -1 or emb_db.expected_shape == embedding.shape: + embedding.loaded = True + emb_db.register_embedding(embedding, shared.sd_model) + else: + emb_db.skipped_embeddings[name] = embedding + + if failed_to_load_networks: + lora_not_found_message = f'Lora not found: {", ".join(failed_to_load_networks)}' + sd_hijack.model_hijack.comments.append(lora_not_found_message) + if shared.opts.lora_not_found_warning_console: + print(f'\n{lora_not_found_message}\n') + if shared.opts.lora_not_found_gradio_warning: + gr.Warning(lora_not_found_message) + + purge_networks_from_memory() + + +def allowed_layer_without_weight(layer): + if isinstance(layer, torch.nn.LayerNorm) and not layer.elementwise_affine: + return True + + return False + + +def store_weights_backup(weight): + if weight is None: + return None + + return weight.to(devices.cpu, copy=True) + + +def restore_weights_backup(obj, field, weight): + if weight is None: + setattr(obj, field, None) + return + + getattr(obj, field).copy_(weight) + + +def network_restore_weights_from_backup(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, torch.nn.MultiheadAttention]): + weights_backup = getattr(self, "network_weights_backup", None) + bias_backup = getattr(self, "network_bias_backup", None) + + if weights_backup is None and bias_backup is None: + return + + if weights_backup is not None: + if isinstance(self, torch.nn.MultiheadAttention): + restore_weights_backup(self, 'in_proj_weight', weights_backup[0]) + restore_weights_backup(self.out_proj, 'weight', weights_backup[1]) + else: + restore_weights_backup(self, 'weight', weights_backup) + + if isinstance(self, torch.nn.MultiheadAttention): + restore_weights_backup(self.out_proj, 'bias', bias_backup) + else: + restore_weights_backup(self, 'bias', bias_backup) + + +def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, torch.nn.MultiheadAttention]): + """ + Applies the currently selected set of networks to the weights of torch layer self. + If weights already have this particular set of networks applied, does nothing. + If not, restores original weights from backup and alters weights according to networks. + """ + + network_layer_name = getattr(self, 'network_layer_name', None) + if network_layer_name is None: + return + + current_names = getattr(self, "network_current_names", ()) + wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in loaded_networks) + + weights_backup = getattr(self, "network_weights_backup", None) + if weights_backup is None and wanted_names != (): + if current_names != () and not allowed_layer_without_weight(self): + raise RuntimeError(f"{network_layer_name} - no backup weights found and current weights are not unchanged") + + if isinstance(self, torch.nn.MultiheadAttention): + weights_backup = (store_weights_backup(self.in_proj_weight), store_weights_backup(self.out_proj.weight)) + else: + weights_backup = store_weights_backup(self.weight) + + self.network_weights_backup = weights_backup + + bias_backup = getattr(self, "network_bias_backup", None) + if bias_backup is None and wanted_names != (): + if isinstance(self, torch.nn.MultiheadAttention) and self.out_proj.bias is not None: + bias_backup = store_weights_backup(self.out_proj.bias) + elif getattr(self, 'bias', None) is not None: + bias_backup = store_weights_backup(self.bias) + else: + bias_backup = None + + # Unlike weight which always has value, some modules don't have bias. + # Only report if bias is not None and current bias are not unchanged. + if bias_backup is not None and current_names != (): + raise RuntimeError("no backup bias found and current bias are not unchanged") + + self.network_bias_backup = bias_backup + + if current_names != wanted_names: + network_restore_weights_from_backup(self) + + for net in loaded_networks: + module = net.modules.get(network_layer_name, None) + if module is not None and hasattr(self, 'weight') and not isinstance(module, modules.models.sd3.mmdit.QkvLinear): + try: + with torch.no_grad(): + if getattr(self, 'fp16_weight', None) is None: + weight = self.weight + bias = self.bias + else: + weight = self.fp16_weight.clone().to(self.weight.device) + bias = getattr(self, 'fp16_bias', None) + if bias is not None: + bias = bias.clone().to(self.bias.device) + updown, ex_bias = module.calc_updown(weight) + + if len(weight.shape) == 4 and weight.shape[1] == 9: + # inpainting model. zero pad updown to make channel[1] 4 to 9 + updown = torch.nn.functional.pad(updown, (0, 0, 0, 0, 0, 5)) + + self.weight.copy_((weight.to(dtype=updown.dtype) + updown).to(dtype=self.weight.dtype)) + if ex_bias is not None and hasattr(self, 'bias'): + if self.bias is None: + self.bias = torch.nn.Parameter(ex_bias).to(self.weight.dtype) + else: + self.bias.copy_((bias + ex_bias).to(dtype=self.bias.dtype)) + except RuntimeError as e: + logging.debug(f"Network {net.name} layer {network_layer_name}: {e}") + extra_network_lora.errors[net.name] = extra_network_lora.errors.get(net.name, 0) + 1 + + continue + + module_q = net.modules.get(network_layer_name + "_q_proj", None) + module_k = net.modules.get(network_layer_name + "_k_proj", None) + module_v = net.modules.get(network_layer_name + "_v_proj", None) + module_out = net.modules.get(network_layer_name + "_out_proj", None) + + if isinstance(self, torch.nn.MultiheadAttention) and module_q and module_k and module_v and module_out: + try: + with torch.no_grad(): + # Send "real" orig_weight into MHA's lora module + qw, kw, vw = self.in_proj_weight.chunk(3, 0) + updown_q, _ = module_q.calc_updown(qw) + updown_k, _ = module_k.calc_updown(kw) + updown_v, _ = module_v.calc_updown(vw) + del qw, kw, vw + updown_qkv = torch.vstack([updown_q, updown_k, updown_v]) + updown_out, ex_bias = module_out.calc_updown(self.out_proj.weight) + + self.in_proj_weight += updown_qkv + self.out_proj.weight += updown_out + if ex_bias is not None: + if self.out_proj.bias is None: + self.out_proj.bias = torch.nn.Parameter(ex_bias) + else: + self.out_proj.bias += ex_bias + + except RuntimeError as e: + logging.debug(f"Network {net.name} layer {network_layer_name}: {e}") + extra_network_lora.errors[net.name] = extra_network_lora.errors.get(net.name, 0) + 1 + + continue + + if isinstance(self, modules.models.sd3.mmdit.QkvLinear) and module_q and module_k and module_v: + try: + with torch.no_grad(): + # Send "real" orig_weight into MHA's lora module + qw, kw, vw = self.weight.chunk(3, 0) + updown_q, _ = module_q.calc_updown(qw) + updown_k, _ = module_k.calc_updown(kw) + updown_v, _ = module_v.calc_updown(vw) + del qw, kw, vw + updown_qkv = torch.vstack([updown_q, updown_k, updown_v]) + self.weight += updown_qkv + + except RuntimeError as e: + logging.debug(f"Network {net.name} layer {network_layer_name}: {e}") + extra_network_lora.errors[net.name] = extra_network_lora.errors.get(net.name, 0) + 1 + + continue + + if module is None: + continue + + logging.debug(f"Network {net.name} layer {network_layer_name}: couldn't find supported operation") + extra_network_lora.errors[net.name] = extra_network_lora.errors.get(net.name, 0) + 1 + + self.network_current_names = wanted_names + + +def network_forward(org_module, input, original_forward): + """ + Old way of applying Lora by executing operations during layer's forward. + Stacking many loras this way results in big performance degradation. + """ + + if len(loaded_networks) == 0: + return original_forward(org_module, input) + + input = devices.cond_cast_unet(input) + + network_restore_weights_from_backup(org_module) + network_reset_cached_weight(org_module) + + y = original_forward(org_module, input) + + network_layer_name = getattr(org_module, 'network_layer_name', None) + for lora in loaded_networks: + module = lora.modules.get(network_layer_name, None) + if module is None: + continue + + y = module.forward(input, y) + + return y + + +def network_reset_cached_weight(self: Union[torch.nn.Conv2d, torch.nn.Linear]): + self.network_current_names = () + self.network_weights_backup = None + self.network_bias_backup = None + + +def network_Linear_forward(self, input): + if shared.opts.lora_functional: + return network_forward(self, input, originals.Linear_forward) + + network_apply_weights(self) + + return originals.Linear_forward(self, input) + + +def network_Linear_load_state_dict(self, *args, **kwargs): + network_reset_cached_weight(self) + + return originals.Linear_load_state_dict(self, *args, **kwargs) + + +def network_Conv2d_forward(self, input): + if shared.opts.lora_functional: + return network_forward(self, input, originals.Conv2d_forward) + + network_apply_weights(self) + + return originals.Conv2d_forward(self, input) + + +def network_Conv2d_load_state_dict(self, *args, **kwargs): + network_reset_cached_weight(self) + + return originals.Conv2d_load_state_dict(self, *args, **kwargs) + + +def network_GroupNorm_forward(self, input): + if shared.opts.lora_functional: + return network_forward(self, input, originals.GroupNorm_forward) + + network_apply_weights(self) + + return originals.GroupNorm_forward(self, input) + + +def network_GroupNorm_load_state_dict(self, *args, **kwargs): + network_reset_cached_weight(self) + + return originals.GroupNorm_load_state_dict(self, *args, **kwargs) + + +def network_LayerNorm_forward(self, input): + if shared.opts.lora_functional: + return network_forward(self, input, originals.LayerNorm_forward) + + network_apply_weights(self) + + return originals.LayerNorm_forward(self, input) + + +def network_LayerNorm_load_state_dict(self, *args, **kwargs): + network_reset_cached_weight(self) + + return originals.LayerNorm_load_state_dict(self, *args, **kwargs) + + +def network_MultiheadAttention_forward(self, *args, **kwargs): + network_apply_weights(self) + + return originals.MultiheadAttention_forward(self, *args, **kwargs) + + +def network_MultiheadAttention_load_state_dict(self, *args, **kwargs): + network_reset_cached_weight(self) + + return originals.MultiheadAttention_load_state_dict(self, *args, **kwargs) + + +def process_network_files(names: list[str] | None = None): + candidates = list(shared.walk_files(shared.cmd_opts.lora_dir, allowed_extensions=[".pt", ".ckpt", ".safetensors"])) + candidates += list(shared.walk_files(shared.cmd_opts.lyco_dir_backcompat, allowed_extensions=[".pt", ".ckpt", ".safetensors"])) + for filename in candidates: + if os.path.isdir(filename): + continue + name = os.path.splitext(os.path.basename(filename))[0] + # if names is provided, only load networks with names in the list + if names and name not in names: + continue + try: + entry = network.NetworkOnDisk(name, filename) + except OSError: # should catch FileNotFoundError and PermissionError etc. + errors.report(f"Failed to load network {name} from {filename}", exc_info=True) + continue + + available_networks[name] = entry + + if entry.alias in available_network_aliases: + forbidden_network_aliases[entry.alias.lower()] = 1 + + available_network_aliases[name] = entry + available_network_aliases[entry.alias] = entry + + +def update_available_networks_by_names(names: list[str]): + process_network_files(names) + + +def list_available_networks(): + available_networks.clear() + available_network_aliases.clear() + forbidden_network_aliases.clear() + available_network_hash_lookup.clear() + forbidden_network_aliases.update({"none": 1, "Addams": 1}) + + os.makedirs(shared.cmd_opts.lora_dir, exist_ok=True) + + process_network_files() + + +re_network_name = re.compile(r"(.*)\s*\([0-9a-fA-F]+\)") + + +def infotext_pasted(infotext, params): + if "AddNet Module 1" in [x[1] for x in scripts.scripts_txt2img.infotext_fields]: + return # if the other extension is active, it will handle those fields, no need to do anything + + added = [] + + for k in params: + if not k.startswith("AddNet Model "): + continue + + num = k[13:] + + if params.get("AddNet Module " + num) != "LoRA": + continue + + name = params.get("AddNet Model " + num) + if name is None: + continue + + m = re_network_name.match(name) + if m: + name = m.group(1) + + multiplier = params.get("AddNet Weight A " + num, "1.0") + + added.append(f"") + + if added: + params["Prompt"] += "\n" + "".join(added) + + +originals: lora_patches.LoraPatches = None + +extra_network_lora = None + +available_networks = {} +available_network_aliases = {} +loaded_networks = [] +loaded_bundle_embeddings = {} +networks_in_memory = {} +available_network_hash_lookup = {} +forbidden_network_aliases = {} + +list_available_networks() diff --git a/stable-diffusion-webui/extensions-builtin/Lora/preload.py b/stable-diffusion-webui/extensions-builtin/Lora/preload.py new file mode 100755 index 0000000..763f942 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/Lora/preload.py @@ -0,0 +1,8 @@ +import os +from modules import paths +from modules.paths_internal import normalized_filepath + + +def preload(parser): + parser.add_argument("--lora-dir", type=normalized_filepath, help="Path to directory with Lora networks.", default=os.path.join(paths.models_path, 'Lora')) + parser.add_argument("--lyco-dir-backcompat", type=normalized_filepath, help="Path to directory with LyCORIS networks (for backawards compatibility; can also use --lyco-dir).", default=os.path.join(paths.models_path, 'LyCORIS')) diff --git a/stable-diffusion-webui/extensions-builtin/Lora/scripts/lora_script.py b/stable-diffusion-webui/extensions-builtin/Lora/scripts/lora_script.py new file mode 100755 index 0000000..15be776 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/Lora/scripts/lora_script.py @@ -0,0 +1,102 @@ +import re + +import gradio as gr +from fastapi import FastAPI + +import network +import networks +import lora # noqa:F401 +import lora_patches +import extra_networks_lora +import ui_extra_networks_lora +from modules import script_callbacks, ui_extra_networks, extra_networks, shared + + +def unload(): + networks.originals.undo() + + +def before_ui(): + ui_extra_networks.register_page(ui_extra_networks_lora.ExtraNetworksPageLora()) + + networks.extra_network_lora = extra_networks_lora.ExtraNetworkLora() + extra_networks.register_extra_network(networks.extra_network_lora) + extra_networks.register_extra_network_alias(networks.extra_network_lora, "lyco") + + +networks.originals = lora_patches.LoraPatches() + +script_callbacks.on_model_loaded(networks.assign_network_names_to_compvis_modules) +script_callbacks.on_script_unloaded(unload) +script_callbacks.on_before_ui(before_ui) +script_callbacks.on_infotext_pasted(networks.infotext_pasted) + + +shared.options_templates.update(shared.options_section(('extra_networks', "Extra Networks"), { + "sd_lora": shared.OptionInfo("None", "Add network to prompt", gr.Dropdown, lambda: {"choices": ["None", *networks.available_networks]}, refresh=networks.list_available_networks), + "lora_preferred_name": shared.OptionInfo("Alias from file", "When adding to prompt, refer to Lora by", gr.Radio, {"choices": ["Alias from file", "Filename"]}), + "lora_add_hashes_to_infotext": shared.OptionInfo(True, "Add Lora hashes to infotext"), + "lora_bundled_ti_to_infotext": shared.OptionInfo(True, "Add Lora name as TI hashes for bundled Textual Inversion").info('"Add Textual Inversion hashes to infotext" needs to be enabled'), + "lora_show_all": shared.OptionInfo(False, "Always show all networks on the Lora page").info("otherwise, those detected as for incompatible version of Stable Diffusion will be hidden"), + "lora_hide_unknown_for_versions": shared.OptionInfo([], "Hide networks of unknown versions for model versions", gr.CheckboxGroup, {"choices": ["SD1", "SD2", "SDXL"]}), + "lora_in_memory_limit": shared.OptionInfo(0, "Number of Lora networks to keep cached in memory", gr.Number, {"precision": 0}), + "lora_not_found_warning_console": shared.OptionInfo(False, "Lora not found warning in console"), + "lora_not_found_gradio_warning": shared.OptionInfo(False, "Lora not found warning popup in webui"), +})) + + +shared.options_templates.update(shared.options_section(('compatibility', "Compatibility"), { + "lora_functional": shared.OptionInfo(False, "Lora/Networks: use old method that takes longer when you have multiple Loras active and produces same results as kohya-ss/sd-webui-additional-networks extension"), +})) + + +def create_lora_json(obj: network.NetworkOnDisk): + return { + "name": obj.name, + "alias": obj.alias, + "path": obj.filename, + "metadata": obj.metadata, + } + + +def api_networks(_: gr.Blocks, app: FastAPI): + @app.get("/sdapi/v1/loras") + async def get_loras(): + return [create_lora_json(obj) for obj in networks.available_networks.values()] + + @app.post("/sdapi/v1/refresh-loras") + async def refresh_loras(): + return networks.list_available_networks() + + +script_callbacks.on_app_started(api_networks) + +re_lora = re.compile("= 16 + + +re_word = re.compile(r"[-_\w']+") +re_comma = re.compile(r" *, *") + + +def build_tags(metadata): + tags = {} + + ss_tag_frequency = metadata.get("ss_tag_frequency", {}) + if ss_tag_frequency is not None and hasattr(ss_tag_frequency, 'items'): + for _, tags_dict in ss_tag_frequency.items(): + for tag, tag_count in tags_dict.items(): + tag = tag.strip() + tags[tag] = tags.get(tag, 0) + int(tag_count) + + if tags and is_non_comma_tagset(tags): + new_tags = {} + + for text, text_count in tags.items(): + for word in re.findall(re_word, text): + if len(word) < 3: + continue + + new_tags[word] = new_tags.get(word, 0) + text_count + + tags = new_tags + + ordered_tags = sorted(tags.keys(), key=tags.get, reverse=True) + + return [(tag, tags[tag]) for tag in ordered_tags] + + +class LoraUserMetadataEditor(ui_extra_networks_user_metadata.UserMetadataEditor): + def __init__(self, ui, tabname, page): + super().__init__(ui, tabname, page) + + self.select_sd_version = None + + self.taginfo = None + self.edit_activation_text = None + self.slider_preferred_weight = None + self.edit_notes = None + + def save_lora_user_metadata(self, name, desc, sd_version, activation_text, preferred_weight, negative_text, notes): + user_metadata = self.get_user_metadata(name) + user_metadata["description"] = desc + user_metadata["sd version"] = sd_version + user_metadata["activation text"] = activation_text + user_metadata["preferred weight"] = preferred_weight + user_metadata["negative text"] = negative_text + user_metadata["notes"] = notes + + self.write_user_metadata(name, user_metadata) + + def get_metadata_table(self, name): + table = super().get_metadata_table(name) + item = self.page.items.get(name, {}) + metadata = item.get("metadata") or {} + + keys = { + 'ss_output_name': "Output name:", + 'ss_sd_model_name': "Model:", + 'ss_clip_skip': "Clip skip:", + 'ss_network_module': "Kohya module:", + } + + for key, label in keys.items(): + value = metadata.get(key, None) + if value is not None and str(value) != "None": + table.append((label, html.escape(value))) + + ss_training_started_at = metadata.get('ss_training_started_at') + if ss_training_started_at: + table.append(("Date trained:", datetime.datetime.utcfromtimestamp(float(ss_training_started_at)).strftime('%Y-%m-%d %H:%M'))) + + ss_bucket_info = metadata.get("ss_bucket_info") + if ss_bucket_info and "buckets" in ss_bucket_info: + resolutions = {} + for _, bucket in ss_bucket_info["buckets"].items(): + resolution = bucket["resolution"] + resolution = f'{resolution[1]}x{resolution[0]}' + + resolutions[resolution] = resolutions.get(resolution, 0) + int(bucket["count"]) + + resolutions_list = sorted(resolutions.keys(), key=resolutions.get, reverse=True) + resolutions_text = html.escape(", ".join(resolutions_list[0:4])) + if len(resolutions) > 4: + resolutions_text += ", ..." + resolutions_text = f"{resolutions_text}" + + table.append(('Resolutions:' if len(resolutions_list) > 1 else 'Resolution:', resolutions_text)) + + image_count = 0 + for _, params in metadata.get("ss_dataset_dirs", {}).items(): + image_count += int(params.get("img_count", 0)) + + if image_count: + table.append(("Dataset size:", image_count)) + + return table + + def put_values_into_components(self, name): + user_metadata = self.get_user_metadata(name) + values = super().put_values_into_components(name) + + item = self.page.items.get(name, {}) + metadata = item.get("metadata") or {} + + tags = build_tags(metadata) + gradio_tags = [(tag, str(count)) for tag, count in tags[0:24]] + + return [ + *values[0:5], + item.get("sd_version", "Unknown"), + gr.HighlightedText.update(value=gradio_tags, visible=True if tags else False), + user_metadata.get('activation text', ''), + float(user_metadata.get('preferred weight', 0.0)), + user_metadata.get('negative text', ''), + gr.update(visible=True if tags else False), + gr.update(value=self.generate_random_prompt_from_tags(tags), visible=True if tags else False), + ] + + def generate_random_prompt(self, name): + item = self.page.items.get(name, {}) + metadata = item.get("metadata") or {} + tags = build_tags(metadata) + + return self.generate_random_prompt_from_tags(tags) + + def generate_random_prompt_from_tags(self, tags): + max_count = None + res = [] + for tag, count in tags: + if not max_count: + max_count = count + + v = random.random() * max_count + if count > v: + for x in "({[]})": + tag = tag.replace(x, '\\' + x) + res.append(tag) + + return ", ".join(sorted(res)) + + def create_extra_default_items_in_left_column(self): + + # this would be a lot better as gr.Radio but I can't make it work + self.select_sd_version = gr.Dropdown(['SD1', 'SD2', 'SDXL', 'Unknown'], value='Unknown', label='Stable Diffusion version', interactive=True) + + def create_editor(self): + self.create_default_editor_elems() + + self.taginfo = gr.HighlightedText(label="Training dataset tags") + self.edit_activation_text = gr.Text(label='Activation text', info="Will be added to prompt along with Lora") + self.slider_preferred_weight = gr.Slider(label='Preferred weight', info="Set to 0 to disable", minimum=0.0, maximum=2.0, step=0.01) + self.edit_negative_text = gr.Text(label='Negative prompt', info="Will be added to negative prompts") + with gr.Row() as row_random_prompt: + with gr.Column(scale=8): + random_prompt = gr.Textbox(label='Random prompt', lines=4, max_lines=4, interactive=False) + + with gr.Column(scale=1, min_width=120): + generate_random_prompt = gr.Button('Generate', size="lg", scale=1) + + self.edit_notes = gr.TextArea(label='Notes', lines=4) + + generate_random_prompt.click(fn=self.generate_random_prompt, inputs=[self.edit_name_input], outputs=[random_prompt], show_progress=False) + + def select_tag(activation_text, evt: gr.SelectData): + tag = evt.value[0] + + words = re.split(re_comma, activation_text) + if tag in words: + words = [x for x in words if x != tag and x.strip()] + return ", ".join(words) + + return activation_text + ", " + tag if activation_text else tag + + self.taginfo.select(fn=select_tag, inputs=[self.edit_activation_text], outputs=[self.edit_activation_text], show_progress=False) + + self.create_default_buttons() + + viewed_components = [ + self.edit_name, + self.edit_description, + self.html_filedata, + self.html_preview, + self.edit_notes, + self.select_sd_version, + self.taginfo, + self.edit_activation_text, + self.slider_preferred_weight, + self.edit_negative_text, + row_random_prompt, + random_prompt, + ] + + self.button_edit\ + .click(fn=self.put_values_into_components, inputs=[self.edit_name_input], outputs=viewed_components)\ + .then(fn=lambda: gr.update(visible=True), inputs=[], outputs=[self.box]) + + edited_components = [ + self.edit_description, + self.select_sd_version, + self.edit_activation_text, + self.slider_preferred_weight, + self.edit_negative_text, + self.edit_notes, + ] + + + self.setup_save_handler(self.button_save, self.save_lora_user_metadata, edited_components) diff --git a/stable-diffusion-webui/extensions-builtin/Lora/ui_extra_networks_lora.py b/stable-diffusion-webui/extensions-builtin/Lora/ui_extra_networks_lora.py new file mode 100755 index 0000000..42e0376 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/Lora/ui_extra_networks_lora.py @@ -0,0 +1,90 @@ +import os + +import network +import networks + +from modules import shared, ui_extra_networks +from modules.ui_extra_networks import quote_js +from ui_edit_user_metadata import LoraUserMetadataEditor + + +class ExtraNetworksPageLora(ui_extra_networks.ExtraNetworksPage): + def __init__(self): + super().__init__('Lora') + + def refresh(self): + networks.list_available_networks() + + def create_item(self, name, index=None, enable_filter=True): + lora_on_disk = networks.available_networks.get(name) + if lora_on_disk is None: + return + + path, ext = os.path.splitext(lora_on_disk.filename) + + alias = lora_on_disk.get_alias() + + search_terms = [self.search_terms_from_path(lora_on_disk.filename)] + if lora_on_disk.hash: + search_terms.append(lora_on_disk.hash) + item = { + "name": name, + "filename": lora_on_disk.filename, + "shorthash": lora_on_disk.shorthash, + "preview": self.find_preview(path) or self.find_embedded_preview(path, name, lora_on_disk.metadata), + "description": self.find_description(path), + "search_terms": search_terms, + "local_preview": f"{path}.{shared.opts.samples_format}", + "metadata": lora_on_disk.metadata, + "sort_keys": {'default': index, **self.get_sort_keys(lora_on_disk.filename)}, + "sd_version": lora_on_disk.sd_version.name, + } + + self.read_user_metadata(item) + activation_text = item["user_metadata"].get("activation text") + preferred_weight = item["user_metadata"].get("preferred weight", 0.0) + item["prompt"] = quote_js(f"") + + if activation_text: + item["prompt"] += " + " + quote_js(" " + activation_text) + + negative_prompt = item["user_metadata"].get("negative text") + item["negative_prompt"] = quote_js("") + if negative_prompt: + item["negative_prompt"] = quote_js('(' + negative_prompt + ':1)') + + sd_version = item["user_metadata"].get("sd version") + if sd_version in network.SdVersion.__members__: + item["sd_version"] = sd_version + sd_version = network.SdVersion[sd_version] + else: + sd_version = lora_on_disk.sd_version + + if shared.opts.lora_show_all or not enable_filter or not shared.sd_model: + pass + elif sd_version == network.SdVersion.Unknown: + model_version = network.SdVersion.SDXL if shared.sd_model.is_sdxl else network.SdVersion.SD2 if shared.sd_model.is_sd2 else network.SdVersion.SD1 + if model_version.name in shared.opts.lora_hide_unknown_for_versions: + return None + elif shared.sd_model.is_sdxl and sd_version != network.SdVersion.SDXL: + return None + elif shared.sd_model.is_sd2 and sd_version != network.SdVersion.SD2: + return None + elif shared.sd_model.is_sd1 and sd_version != network.SdVersion.SD1: + return None + + return item + + def list_items(self): + # instantiate a list to protect against concurrent modification + names = list(networks.available_networks) + for index, name in enumerate(names): + item = self.create_item(name, index) + if item is not None: + yield item + + def allowed_directories_for_previews(self): + return [shared.cmd_opts.lora_dir, shared.cmd_opts.lyco_dir_backcompat] + + def create_user_metadata_editor(self, ui, tabname): + return LoraUserMetadataEditor(ui, tabname, self) diff --git a/stable-diffusion-webui/extensions-builtin/ScuNET/preload.py b/stable-diffusion-webui/extensions-builtin/ScuNET/preload.py new file mode 100755 index 0000000..4ce82b1 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/ScuNET/preload.py @@ -0,0 +1,6 @@ +import os +from modules import paths + + +def preload(parser): + parser.add_argument("--scunet-models-path", type=str, help="Path to directory with ScuNET model file(s).", default=os.path.join(paths.models_path, 'ScuNET')) diff --git a/stable-diffusion-webui/extensions-builtin/ScuNET/scripts/scunet_model.py b/stable-diffusion-webui/extensions-builtin/ScuNET/scripts/scunet_model.py new file mode 100755 index 0000000..fe5e5a1 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/ScuNET/scripts/scunet_model.py @@ -0,0 +1,74 @@ +import sys + +import PIL.Image + +import modules.upscaler +from modules import devices, errors, modelloader, script_callbacks, shared, upscaler_utils + + +class UpscalerScuNET(modules.upscaler.Upscaler): + def __init__(self, dirname): + self.name = "ScuNET" + self.model_name = "ScuNET GAN" + self.model_name2 = "ScuNET PSNR" + self.model_url = "https://github.com/cszn/KAIR/releases/download/v1.0/scunet_color_real_gan.pth" + self.model_url2 = "https://github.com/cszn/KAIR/releases/download/v1.0/scunet_color_real_psnr.pth" + self.user_path = dirname + super().__init__() + model_paths = self.find_models(ext_filter=[".pth"]) + scalers = [] + add_model2 = True + for file in model_paths: + if file.startswith("http"): + name = self.model_name + else: + name = modelloader.friendly_name(file) + if name == self.model_name2 or file == self.model_url2: + add_model2 = False + try: + scaler_data = modules.upscaler.UpscalerData(name, file, self, 4) + scalers.append(scaler_data) + except Exception: + errors.report(f"Error loading ScuNET model: {file}", exc_info=True) + if add_model2: + scaler_data2 = modules.upscaler.UpscalerData(self.model_name2, self.model_url2, self) + scalers.append(scaler_data2) + self.scalers = scalers + + def do_upscale(self, img: PIL.Image.Image, selected_file): + devices.torch_gc() + try: + model = self.load_model(selected_file) + except Exception as e: + print(f"ScuNET: Unable to load model from {selected_file}: {e}", file=sys.stderr) + return img + + img = upscaler_utils.upscale_2( + img, + model, + tile_size=shared.opts.SCUNET_tile, + tile_overlap=shared.opts.SCUNET_tile_overlap, + scale=1, # ScuNET is a denoising model, not an upscaler + desc='ScuNET', + ) + devices.torch_gc() + return img + + def load_model(self, path: str): + device = devices.get_device_for('scunet') + if path.startswith("http"): + # TODO: this doesn't use `path` at all? + filename = modelloader.load_file_from_url(self.model_url, model_dir=self.model_download_path, file_name=f"{self.name}.pth") + else: + filename = path + return modelloader.load_spandrel_model(filename, device=device, expected_architecture='SCUNet') + + +def on_ui_settings(): + import gradio as gr + + shared.opts.add_option("SCUNET_tile", shared.OptionInfo(256, "Tile size for SCUNET upscalers.", gr.Slider, {"minimum": 0, "maximum": 512, "step": 16}, section=('upscaling', "Upscaling")).info("0 = no tiling")) + shared.opts.add_option("SCUNET_tile_overlap", shared.OptionInfo(8, "Tile overlap for SCUNET upscalers.", gr.Slider, {"minimum": 0, "maximum": 64, "step": 1}, section=('upscaling', "Upscaling")).info("Low values = visible seam")) + + +script_callbacks.on_ui_settings(on_ui_settings) diff --git a/stable-diffusion-webui/extensions-builtin/SwinIR/preload.py b/stable-diffusion-webui/extensions-builtin/SwinIR/preload.py new file mode 100755 index 0000000..e912c64 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/SwinIR/preload.py @@ -0,0 +1,6 @@ +import os +from modules import paths + + +def preload(parser): + parser.add_argument("--swinir-models-path", type=str, help="Path to directory with SwinIR model file(s).", default=os.path.join(paths.models_path, 'SwinIR')) diff --git a/stable-diffusion-webui/extensions-builtin/SwinIR/scripts/swinir_model.py b/stable-diffusion-webui/extensions-builtin/SwinIR/scripts/swinir_model.py new file mode 100755 index 0000000..16bf9b7 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/SwinIR/scripts/swinir_model.py @@ -0,0 +1,95 @@ +import logging +import sys + +import torch +from PIL import Image + +from modules import devices, modelloader, script_callbacks, shared, upscaler_utils +from modules.upscaler import Upscaler, UpscalerData + +SWINIR_MODEL_URL = "https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN.pth" + +logger = logging.getLogger(__name__) + + +class UpscalerSwinIR(Upscaler): + def __init__(self, dirname): + self._cached_model = None # keep the model when SWIN_torch_compile is on to prevent re-compile every runs + self._cached_model_config = None # to clear '_cached_model' when changing model (v1/v2) or settings + self.name = "SwinIR" + self.model_url = SWINIR_MODEL_URL + self.model_name = "SwinIR 4x" + self.user_path = dirname + super().__init__() + scalers = [] + model_files = self.find_models(ext_filter=[".pt", ".pth"]) + for model in model_files: + if model.startswith("http"): + name = self.model_name + else: + name = modelloader.friendly_name(model) + model_data = UpscalerData(name, model, self) + scalers.append(model_data) + self.scalers = scalers + + def do_upscale(self, img: Image.Image, model_file: str) -> Image.Image: + current_config = (model_file, shared.opts.SWIN_tile) + + if self._cached_model_config == current_config: + model = self._cached_model + else: + try: + model = self.load_model(model_file) + except Exception as e: + print(f"Failed loading SwinIR model {model_file}: {e}", file=sys.stderr) + return img + self._cached_model = model + self._cached_model_config = current_config + + img = upscaler_utils.upscale_2( + img, + model, + tile_size=shared.opts.SWIN_tile, + tile_overlap=shared.opts.SWIN_tile_overlap, + scale=model.scale, + desc="SwinIR", + ) + devices.torch_gc() + return img + + def load_model(self, path, scale=4): + if path.startswith("http"): + filename = modelloader.load_file_from_url( + url=path, + model_dir=self.model_download_path, + file_name=f"{self.model_name.replace(' ', '_')}.pth", + ) + else: + filename = path + + model_descriptor = modelloader.load_spandrel_model( + filename, + device=self._get_device(), + prefer_half=(devices.dtype == torch.float16), + expected_architecture="SwinIR", + ) + if getattr(shared.opts, 'SWIN_torch_compile', False): + try: + model_descriptor.model.compile() + except Exception: + logger.warning("Failed to compile SwinIR model, fallback to JIT", exc_info=True) + return model_descriptor + + def _get_device(self): + return devices.get_device_for('swinir') + + +def on_ui_settings(): + import gradio as gr + + shared.opts.add_option("SWIN_tile", shared.OptionInfo(192, "Tile size for all SwinIR.", gr.Slider, {"minimum": 16, "maximum": 512, "step": 16}, section=('upscaling', "Upscaling"))) + shared.opts.add_option("SWIN_tile_overlap", shared.OptionInfo(8, "Tile overlap, in pixels for SwinIR. Low values = visible seam.", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}, section=('upscaling', "Upscaling"))) + shared.opts.add_option("SWIN_torch_compile", shared.OptionInfo(False, "Use torch.compile to accelerate SwinIR.", gr.Checkbox, {"interactive": True}, section=('upscaling', "Upscaling")).info("Takes longer on first run")) + + +script_callbacks.on_ui_settings(on_ui_settings) diff --git a/stable-diffusion-webui/extensions-builtin/canvas-zoom-and-pan/javascript/zoom.js b/stable-diffusion-webui/extensions-builtin/canvas-zoom-and-pan/javascript/zoom.js new file mode 100755 index 0000000..7807f7f --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/canvas-zoom-and-pan/javascript/zoom.js @@ -0,0 +1,995 @@ +onUiLoaded(async() => { + const elementIDs = { + img2imgTabs: "#mode_img2img .tab-nav", + inpaint: "#img2maskimg", + inpaintSketch: "#inpaint_sketch", + rangeGroup: "#img2img_column_size", + sketch: "#img2img_sketch" + }; + const tabNameToElementId = { + "Inpaint sketch": elementIDs.inpaintSketch, + "Inpaint": elementIDs.inpaint, + "Sketch": elementIDs.sketch + }; + + + // Helper functions + // Get active tab + + /** + * Waits for an element to be present in the DOM. + */ + const waitForElement = (id) => new Promise(resolve => { + const checkForElement = () => { + const element = document.querySelector(id); + if (element) return resolve(element); + setTimeout(checkForElement, 100); + }; + checkForElement(); + }); + + function getActiveTab(elements, all = false) { + if (!elements.img2imgTabs) return null; + const tabs = elements.img2imgTabs.querySelectorAll("button"); + + if (all) return tabs; + + for (let tab of tabs) { + if (tab.classList.contains("selected")) { + return tab; + } + } + } + + // Get tab ID + function getTabId(elements) { + const activeTab = getActiveTab(elements); + if (!activeTab) return null; + return tabNameToElementId[activeTab.innerText]; + } + + // Wait until opts loaded + async function waitForOpts() { + for (; ;) { + if (window.opts && Object.keys(window.opts).length) { + return window.opts; + } + await new Promise(resolve => setTimeout(resolve, 100)); + } + } + + // Detect whether the element has a horizontal scroll bar + function hasHorizontalScrollbar(element) { + return element.scrollWidth > element.clientWidth; + } + + // Function for defining the "Ctrl", "Shift" and "Alt" keys + function isModifierKey(event, key) { + switch (key) { + case "Ctrl": + return event.ctrlKey; + case "Shift": + return event.shiftKey; + case "Alt": + return event.altKey; + default: + return false; + } + } + + // Check if hotkey is valid + function isValidHotkey(value) { + const specialKeys = ["Ctrl", "Alt", "Shift", "Disable"]; + return ( + (typeof value === "string" && + value.length === 1 && + /[a-z]/i.test(value)) || + specialKeys.includes(value) + ); + } + + // Normalize hotkey + function normalizeHotkey(hotkey) { + return hotkey.length === 1 ? "Key" + hotkey.toUpperCase() : hotkey; + } + + // Format hotkey for display + function formatHotkeyForDisplay(hotkey) { + return hotkey.startsWith("Key") ? hotkey.slice(3) : hotkey; + } + + // Create hotkey configuration with the provided options + function createHotkeyConfig(defaultHotkeysConfig, hotkeysConfigOpts) { + const result = {}; // Resulting hotkey configuration + const usedKeys = new Set(); // Set of used hotkeys + + // Iterate through defaultHotkeysConfig keys + for (const key in defaultHotkeysConfig) { + const userValue = hotkeysConfigOpts[key]; // User-provided hotkey value + const defaultValue = defaultHotkeysConfig[key]; // Default hotkey value + + // Apply appropriate value for undefined, boolean, or object userValue + if ( + userValue === undefined || + typeof userValue === "boolean" || + typeof userValue === "object" || + userValue === "disable" + ) { + result[key] = + userValue === undefined ? defaultValue : userValue; + } else if (isValidHotkey(userValue)) { + const normalizedUserValue = normalizeHotkey(userValue); + + // Check for conflicting hotkeys + if (!usedKeys.has(normalizedUserValue)) { + usedKeys.add(normalizedUserValue); + result[key] = normalizedUserValue; + } else { + console.error( + `Hotkey: ${formatHotkeyForDisplay( + userValue + )} for ${key} is repeated and conflicts with another hotkey. The default hotkey is used: ${formatHotkeyForDisplay( + defaultValue + )}` + ); + result[key] = defaultValue; + } + } else { + console.error( + `Hotkey: ${formatHotkeyForDisplay( + userValue + )} for ${key} is not valid. The default hotkey is used: ${formatHotkeyForDisplay( + defaultValue + )}` + ); + result[key] = defaultValue; + } + } + + return result; + } + + // Disables functions in the config object based on the provided list of function names + function disableFunctions(config, disabledFunctions) { + // Bind the hasOwnProperty method to the functionMap object to avoid errors + const hasOwnProperty = + Object.prototype.hasOwnProperty.bind(functionMap); + + // Loop through the disabledFunctions array and disable the corresponding functions in the config object + disabledFunctions.forEach(funcName => { + if (hasOwnProperty(funcName)) { + const key = functionMap[funcName]; + config[key] = "disable"; + } + }); + + // Return the updated config object + return config; + } + + /** + * The restoreImgRedMask function displays a red mask around an image to indicate the aspect ratio. + * If the image display property is set to 'none', the mask breaks. To fix this, the function + * temporarily sets the display property to 'block' and then hides the mask again after 300 milliseconds + * to avoid breaking the canvas. Additionally, the function adjusts the mask to work correctly on + * very long images. + */ + function restoreImgRedMask(elements) { + const mainTabId = getTabId(elements); + + if (!mainTabId) return; + + const mainTab = gradioApp().querySelector(mainTabId); + const img = mainTab.querySelector("img"); + const imageARPreview = gradioApp().querySelector("#imageARPreview"); + + if (!img || !imageARPreview) return; + + imageARPreview.style.transform = ""; + if (parseFloat(mainTab.style.width) > 865) { + const transformString = mainTab.style.transform; + const scaleMatch = transformString.match( + /scale\(([-+]?[0-9]*\.?[0-9]+)\)/ + ); + let zoom = 1; // default zoom + + if (scaleMatch && scaleMatch[1]) { + zoom = Number(scaleMatch[1]); + } + + imageARPreview.style.transformOrigin = "0 0"; + imageARPreview.style.transform = `scale(${zoom})`; + } + + if (img.style.display !== "none") return; + + img.style.display = "block"; + + setTimeout(() => { + img.style.display = "none"; + }, 400); + } + + const hotkeysConfigOpts = await waitForOpts(); + + // Default config + const defaultHotkeysConfig = { + canvas_hotkey_zoom: "Alt", + canvas_hotkey_adjust: "Ctrl", + canvas_hotkey_reset: "KeyR", + canvas_hotkey_fullscreen: "KeyS", + canvas_hotkey_move: "KeyF", + canvas_hotkey_overlap: "KeyO", + canvas_hotkey_shrink_brush: "KeyQ", + canvas_hotkey_grow_brush: "KeyW", + canvas_disabled_functions: [], + canvas_show_tooltip: true, + canvas_auto_expand: true, + canvas_blur_prompt: false, + }; + + const functionMap = { + "Zoom": "canvas_hotkey_zoom", + "Adjust brush size": "canvas_hotkey_adjust", + "Hotkey shrink brush": "canvas_hotkey_shrink_brush", + "Hotkey enlarge brush": "canvas_hotkey_grow_brush", + "Moving canvas": "canvas_hotkey_move", + "Fullscreen": "canvas_hotkey_fullscreen", + "Reset Zoom": "canvas_hotkey_reset", + "Overlap": "canvas_hotkey_overlap" + }; + + // Loading the configuration from opts + const preHotkeysConfig = createHotkeyConfig( + defaultHotkeysConfig, + hotkeysConfigOpts + ); + + // Disable functions that are not needed by the user + const hotkeysConfig = disableFunctions( + preHotkeysConfig, + preHotkeysConfig.canvas_disabled_functions + ); + + let isMoving = false; + let mouseX, mouseY; + let activeElement; + let interactedWithAltKey = false; + + const elements = Object.fromEntries( + Object.keys(elementIDs).map(id => [ + id, + gradioApp().querySelector(elementIDs[id]) + ]) + ); + const elemData = {}; + + // Apply functionality to the range inputs. Restore redmask and correct for long images. + const rangeInputs = elements.rangeGroup ? + Array.from(elements.rangeGroup.querySelectorAll("input")) : + [ + gradioApp().querySelector("#img2img_width input[type='range']"), + gradioApp().querySelector("#img2img_height input[type='range']") + ]; + + for (const input of rangeInputs) { + input?.addEventListener("input", () => restoreImgRedMask(elements)); + } + + function applyZoomAndPan(elemId, isExtension = true) { + const targetElement = gradioApp().querySelector(elemId); + + if (!targetElement) { + console.log("Element not found", elemId); + return; + } + + targetElement.style.transformOrigin = "0 0"; + + elemData[elemId] = { + zoom: 1, + panX: 0, + panY: 0 + }; + let fullScreenMode = false; + + // Create tooltip + function createTooltip() { + const toolTipElement = + targetElement.querySelector(".image-container"); + const tooltip = document.createElement("div"); + tooltip.className = "canvas-tooltip"; + + // Creating an item of information + const info = document.createElement("i"); + info.className = "canvas-tooltip-info"; + info.textContent = ""; + + // Create a container for the contents of the tooltip + const tooltipContent = document.createElement("div"); + tooltipContent.className = "canvas-tooltip-content"; + + // Define an array with hotkey information and their actions + const hotkeysInfo = [ + { + configKey: "canvas_hotkey_zoom", + action: "Zoom canvas", + keySuffix: " + wheel" + }, + { + configKey: "canvas_hotkey_adjust", + action: "Adjust brush size", + keySuffix: " + wheel" + }, + {configKey: "canvas_hotkey_reset", action: "Reset zoom"}, + { + configKey: "canvas_hotkey_fullscreen", + action: "Fullscreen mode" + }, + {configKey: "canvas_hotkey_move", action: "Move canvas"}, + {configKey: "canvas_hotkey_overlap", action: "Overlap"} + ]; + + // Create hotkeys array with disabled property based on the config values + const hotkeys = hotkeysInfo.map(info => { + const configValue = hotkeysConfig[info.configKey]; + const key = info.keySuffix ? + `${configValue}${info.keySuffix}` : + configValue.charAt(configValue.length - 1); + return { + key, + action: info.action, + disabled: configValue === "disable" + }; + }); + + for (const hotkey of hotkeys) { + if (hotkey.disabled) { + continue; + } + + const p = document.createElement("p"); + p.innerHTML = `${hotkey.key} - ${hotkey.action}`; + tooltipContent.appendChild(p); + } + + // Add information and content elements to the tooltip element + tooltip.appendChild(info); + tooltip.appendChild(tooltipContent); + + // Add a hint element to the target element + toolTipElement.appendChild(tooltip); + } + + //Show tool tip if setting enable + if (hotkeysConfig.canvas_show_tooltip) { + createTooltip(); + } + + // In the course of research, it was found that the tag img is very harmful when zooming and creates white canvases. This hack allows you to almost never think about this problem, it has no effect on webui. + function fixCanvas() { + const activeTab = getActiveTab(elements)?.textContent.trim(); + + if (activeTab && activeTab !== "img2img") { + const img = targetElement.querySelector(`${elemId} img`); + + if (img && img.style.display !== "none") { + img.style.display = "none"; + img.style.visibility = "hidden"; + } + } + } + + // Reset the zoom level and pan position of the target element to their initial values + function resetZoom() { + elemData[elemId] = { + zoomLevel: 1, + panX: 0, + panY: 0 + }; + + if (isExtension) { + targetElement.style.overflow = "hidden"; + } + + targetElement.isZoomed = false; + + fixCanvas(); + targetElement.style.transform = `scale(${elemData[elemId].zoomLevel}) translate(${elemData[elemId].panX}px, ${elemData[elemId].panY}px)`; + + const canvas = gradioApp().querySelector( + `${elemId} canvas[key="interface"]` + ); + + toggleOverlap("off"); + fullScreenMode = false; + + const closeBtn = targetElement.querySelector("button[aria-label='Remove Image']"); + if (closeBtn) { + closeBtn.addEventListener("click", resetZoom); + } + + if (canvas && isExtension) { + const parentElement = targetElement.closest('[id^="component-"]'); + if ( + canvas && + parseFloat(canvas.style.width) > parentElement.offsetWidth && + parseFloat(targetElement.style.width) > parentElement.offsetWidth + ) { + fitToElement(); + return; + } + + } + + if ( + canvas && + !isExtension && + parseFloat(canvas.style.width) > 865 && + parseFloat(targetElement.style.width) > 865 + ) { + fitToElement(); + return; + } + + targetElement.style.width = ""; + } + + // Toggle the zIndex of the target element between two values, allowing it to overlap or be overlapped by other elements + function toggleOverlap(forced = "") { + const zIndex1 = "0"; + const zIndex2 = "998"; + + targetElement.style.zIndex = + targetElement.style.zIndex !== zIndex2 ? zIndex2 : zIndex1; + + if (forced === "off") { + targetElement.style.zIndex = zIndex1; + } else if (forced === "on") { + targetElement.style.zIndex = zIndex2; + } + } + + // Adjust the brush size based on the deltaY value from a mouse wheel event + function adjustBrushSize( + elemId, + deltaY, + withoutValue = false, + percentage = 5 + ) { + const input = + gradioApp().querySelector( + `${elemId} input[aria-label='Brush radius']` + ) || + gradioApp().querySelector( + `${elemId} button[aria-label="Use brush"]` + ); + + if (input) { + input.click(); + if (!withoutValue) { + const maxValue = + parseFloat(input.getAttribute("max")) || 100; + const changeAmount = maxValue * (percentage / 100); + const newValue = + parseFloat(input.value) + + (deltaY > 0 ? -changeAmount : changeAmount); + input.value = Math.min(Math.max(newValue, 0), maxValue); + input.dispatchEvent(new Event("change")); + } + } + } + + // Reset zoom when uploading a new image + const fileInput = gradioApp().querySelector( + `${elemId} input[type="file"][accept="image/*"].svelte-116rqfv` + ); + fileInput.addEventListener("click", resetZoom); + + // Update the zoom level and pan position of the target element based on the values of the zoomLevel, panX and panY variables + function updateZoom(newZoomLevel, mouseX, mouseY) { + newZoomLevel = Math.max(0.1, Math.min(newZoomLevel, 15)); + + elemData[elemId].panX += + mouseX - (mouseX * newZoomLevel) / elemData[elemId].zoomLevel; + elemData[elemId].panY += + mouseY - (mouseY * newZoomLevel) / elemData[elemId].zoomLevel; + + targetElement.style.transformOrigin = "0 0"; + targetElement.style.transform = `translate(${elemData[elemId].panX}px, ${elemData[elemId].panY}px) scale(${newZoomLevel})`; + + toggleOverlap("on"); + if (isExtension) { + targetElement.style.overflow = "visible"; + } + + return newZoomLevel; + } + + // Change the zoom level based on user interaction + function changeZoomLevel(operation, e) { + if (isModifierKey(e, hotkeysConfig.canvas_hotkey_zoom)) { + e.preventDefault(); + + if (hotkeysConfig.canvas_hotkey_zoom === "Alt") { + interactedWithAltKey = true; + } + + let zoomPosX, zoomPosY; + let delta = 0.2; + if (elemData[elemId].zoomLevel > 7) { + delta = 0.9; + } else if (elemData[elemId].zoomLevel > 2) { + delta = 0.6; + } + + zoomPosX = e.clientX; + zoomPosY = e.clientY; + + fullScreenMode = false; + elemData[elemId].zoomLevel = updateZoom( + elemData[elemId].zoomLevel + + (operation === "+" ? delta : -delta), + zoomPosX - targetElement.getBoundingClientRect().left, + zoomPosY - targetElement.getBoundingClientRect().top + ); + + targetElement.isZoomed = true; + } + } + + /** + * This function fits the target element to the screen by calculating + * the required scale and offsets. It also updates the global variables + * zoomLevel, panX, and panY to reflect the new state. + */ + + function fitToElement() { + //Reset Zoom + targetElement.style.transform = `translate(${0}px, ${0}px) scale(${1})`; + + let parentElement; + + if (isExtension) { + parentElement = targetElement.closest('[id^="component-"]'); + } else { + parentElement = targetElement.parentElement; + } + + + // Get element and screen dimensions + const elementWidth = targetElement.offsetWidth; + const elementHeight = targetElement.offsetHeight; + + const screenWidth = parentElement.clientWidth; + const screenHeight = parentElement.clientHeight; + + // Get element's coordinates relative to the parent element + const elementRect = targetElement.getBoundingClientRect(); + const parentRect = parentElement.getBoundingClientRect(); + const elementX = elementRect.x - parentRect.x; + + // Calculate scale and offsets + const scaleX = screenWidth / elementWidth; + const scaleY = screenHeight / elementHeight; + const scale = Math.min(scaleX, scaleY); + + const transformOrigin = + window.getComputedStyle(targetElement).transformOrigin; + const [originX, originY] = transformOrigin.split(" "); + const originXValue = parseFloat(originX); + const originYValue = parseFloat(originY); + + const offsetX = + (screenWidth - elementWidth * scale) / 2 - + originXValue * (1 - scale); + const offsetY = + (screenHeight - elementHeight * scale) / 2.5 - + originYValue * (1 - scale); + + // Apply scale and offsets to the element + targetElement.style.transform = `translate(${offsetX}px, ${offsetY}px) scale(${scale})`; + + // Update global variables + elemData[elemId].zoomLevel = scale; + elemData[elemId].panX = offsetX; + elemData[elemId].panY = offsetY; + + fullScreenMode = false; + toggleOverlap("off"); + } + + /** + * This function fits the target element to the screen by calculating + * the required scale and offsets. It also updates the global variables + * zoomLevel, panX, and panY to reflect the new state. + */ + + // Fullscreen mode + function fitToScreen() { + const canvas = gradioApp().querySelector( + `${elemId} canvas[key="interface"]` + ); + + if (!canvas) return; + + if (canvas.offsetWidth > 862 || isExtension) { + targetElement.style.width = (canvas.offsetWidth + 2) + "px"; + } + + if (isExtension) { + targetElement.style.overflow = "visible"; + } + + if (fullScreenMode) { + resetZoom(); + fullScreenMode = false; + return; + } + + //Reset Zoom + targetElement.style.transform = `translate(${0}px, ${0}px) scale(${1})`; + + // Get scrollbar width to right-align the image + const scrollbarWidth = + window.innerWidth - document.documentElement.clientWidth; + + // Get element and screen dimensions + const elementWidth = targetElement.offsetWidth; + const elementHeight = targetElement.offsetHeight; + const screenWidth = window.innerWidth - scrollbarWidth; + const screenHeight = window.innerHeight; + + // Get element's coordinates relative to the page + const elementRect = targetElement.getBoundingClientRect(); + const elementY = elementRect.y; + const elementX = elementRect.x; + + // Calculate scale and offsets + const scaleX = screenWidth / elementWidth; + const scaleY = screenHeight / elementHeight; + const scale = Math.min(scaleX, scaleY); + + // Get the current transformOrigin + const computedStyle = window.getComputedStyle(targetElement); + const transformOrigin = computedStyle.transformOrigin; + const [originX, originY] = transformOrigin.split(" "); + const originXValue = parseFloat(originX); + const originYValue = parseFloat(originY); + + // Calculate offsets with respect to the transformOrigin + const offsetX = + (screenWidth - elementWidth * scale) / 2 - + elementX - + originXValue * (1 - scale); + const offsetY = + (screenHeight - elementHeight * scale) / 2 - + elementY - + originYValue * (1 - scale); + + // Apply scale and offsets to the element + targetElement.style.transform = `translate(${offsetX}px, ${offsetY}px) scale(${scale})`; + + // Update global variables + elemData[elemId].zoomLevel = scale; + elemData[elemId].panX = offsetX; + elemData[elemId].panY = offsetY; + + fullScreenMode = true; + toggleOverlap("on"); + } + + // Handle keydown events + function handleKeyDown(event) { + // Disable key locks to make pasting from the buffer work correctly + if ((event.ctrlKey && event.code === 'KeyV') || (event.ctrlKey && event.code === 'KeyC') || event.code === "F5") { + return; + } + + // before activating shortcut, ensure user is not actively typing in an input field + if (!hotkeysConfig.canvas_blur_prompt) { + if (event.target.nodeName === 'TEXTAREA' || event.target.nodeName === 'INPUT') { + return; + } + } + + + const hotkeyActions = { + [hotkeysConfig.canvas_hotkey_reset]: resetZoom, + [hotkeysConfig.canvas_hotkey_overlap]: toggleOverlap, + [hotkeysConfig.canvas_hotkey_fullscreen]: fitToScreen, + [hotkeysConfig.canvas_hotkey_shrink_brush]: () => adjustBrushSize(elemId, 10), + [hotkeysConfig.canvas_hotkey_grow_brush]: () => adjustBrushSize(elemId, -10) + }; + + const action = hotkeyActions[event.code]; + if (action) { + event.preventDefault(); + action(event); + } + + if ( + isModifierKey(event, hotkeysConfig.canvas_hotkey_zoom) || + isModifierKey(event, hotkeysConfig.canvas_hotkey_adjust) + ) { + event.preventDefault(); + } + } + + // Get Mouse position + function getMousePosition(e) { + mouseX = e.offsetX; + mouseY = e.offsetY; + } + + // Simulation of the function to put a long image into the screen. + // We detect if an image has a scroll bar or not, make a fullscreen to reveal the image, then reduce it to fit into the element. + // We hide the image and show it to the user when it is ready. + + targetElement.isExpanded = false; + function autoExpand() { + const canvas = document.querySelector(`${elemId} canvas[key="interface"]`); + if (canvas) { + if (hasHorizontalScrollbar(targetElement) && targetElement.isExpanded === false) { + targetElement.style.visibility = "hidden"; + setTimeout(() => { + fitToScreen(); + resetZoom(); + targetElement.style.visibility = "visible"; + targetElement.isExpanded = true; + }, 10); + } + } + } + + targetElement.addEventListener("mousemove", getMousePosition); + + //observers + // Creating an observer with a callback function to handle DOM changes + const observer = new MutationObserver((mutationsList, observer) => { + for (let mutation of mutationsList) { + // If the style attribute of the canvas has changed, by observation it happens only when the picture changes + if (mutation.type === 'attributes' && mutation.attributeName === 'style' && + mutation.target.tagName.toLowerCase() === 'canvas') { + targetElement.isExpanded = false; + setTimeout(resetZoom, 10); + } + } + }); + + // Apply auto expand if enabled + if (hotkeysConfig.canvas_auto_expand) { + targetElement.addEventListener("mousemove", autoExpand); + // Set up an observer to track attribute changes + observer.observe(targetElement, {attributes: true, childList: true, subtree: true}); + } + + // Handle events only inside the targetElement + let isKeyDownHandlerAttached = false; + + function handleMouseMove() { + if (!isKeyDownHandlerAttached) { + document.addEventListener("keydown", handleKeyDown); + isKeyDownHandlerAttached = true; + + activeElement = elemId; + } + } + + function handleMouseLeave() { + if (isKeyDownHandlerAttached) { + document.removeEventListener("keydown", handleKeyDown); + isKeyDownHandlerAttached = false; + + activeElement = null; + } + } + + // Add mouse event handlers + targetElement.addEventListener("mousemove", handleMouseMove); + targetElement.addEventListener("mouseleave", handleMouseLeave); + + // Reset zoom when click on another tab + if (elements.img2imgTabs) { + elements.img2imgTabs.addEventListener("click", resetZoom); + elements.img2imgTabs.addEventListener("click", () => { + // targetElement.style.width = ""; + if (parseInt(targetElement.style.width) > 865) { + setTimeout(fitToElement, 0); + } + }); + } + + targetElement.addEventListener("wheel", e => { + // change zoom level + const operation = (e.deltaY || -e.wheelDelta) > 0 ? "-" : "+"; + changeZoomLevel(operation, e); + + // Handle brush size adjustment with ctrl key pressed + if (isModifierKey(e, hotkeysConfig.canvas_hotkey_adjust)) { + e.preventDefault(); + + if (hotkeysConfig.canvas_hotkey_adjust === "Alt") { + interactedWithAltKey = true; + } + + // Increase or decrease brush size based on scroll direction + adjustBrushSize(elemId, e.deltaY); + } + }); + + // Handle the move event for pan functionality. Updates the panX and panY variables and applies the new transform to the target element. + function handleMoveKeyDown(e) { + + // Disable key locks to make pasting from the buffer work correctly + if ((e.ctrlKey && e.code === 'KeyV') || (e.ctrlKey && event.code === 'KeyC') || e.code === "F5") { + return; + } + + // before activating shortcut, ensure user is not actively typing in an input field + if (!hotkeysConfig.canvas_blur_prompt) { + if (e.target.nodeName === 'TEXTAREA' || e.target.nodeName === 'INPUT') { + return; + } + } + + + if (e.code === hotkeysConfig.canvas_hotkey_move) { + if (!e.ctrlKey && !e.metaKey && isKeyDownHandlerAttached) { + e.preventDefault(); + document.activeElement.blur(); + isMoving = true; + } + } + } + + function handleMoveKeyUp(e) { + if (e.code === hotkeysConfig.canvas_hotkey_move) { + isMoving = false; + } + } + + document.addEventListener("keydown", handleMoveKeyDown); + document.addEventListener("keyup", handleMoveKeyUp); + + + // Prevent firefox from opening main menu when alt is used as a hotkey for zoom or brush size + function handleAltKeyUp(e) { + if (e.key !== "Alt" || !interactedWithAltKey) { + return; + } + + e.preventDefault(); + interactedWithAltKey = false; + } + + document.addEventListener("keyup", handleAltKeyUp); + + + // Detect zoom level and update the pan speed. + function updatePanPosition(movementX, movementY) { + let panSpeed = 2; + + if (elemData[elemId].zoomLevel > 8) { + panSpeed = 3.5; + } + + elemData[elemId].panX += movementX * panSpeed; + elemData[elemId].panY += movementY * panSpeed; + + // Delayed redraw of an element + requestAnimationFrame(() => { + targetElement.style.transform = `translate(${elemData[elemId].panX}px, ${elemData[elemId].panY}px) scale(${elemData[elemId].zoomLevel})`; + toggleOverlap("on"); + }); + } + + function handleMoveByKey(e) { + if (isMoving && elemId === activeElement) { + updatePanPosition(e.movementX, e.movementY); + targetElement.style.pointerEvents = "none"; + + if (isExtension) { + targetElement.style.overflow = "visible"; + } + + } else { + targetElement.style.pointerEvents = "auto"; + } + } + + // Prevents sticking to the mouse + window.onblur = function() { + isMoving = false; + }; + + // Checks for extension + function checkForOutBox() { + const parentElement = targetElement.closest('[id^="component-"]'); + if (parentElement.offsetWidth < targetElement.offsetWidth && !targetElement.isExpanded) { + resetZoom(); + targetElement.isExpanded = true; + } + + if (parentElement.offsetWidth < targetElement.offsetWidth && elemData[elemId].zoomLevel == 1) { + resetZoom(); + } + + if (parentElement.offsetWidth < targetElement.offsetWidth && targetElement.offsetWidth * elemData[elemId].zoomLevel > parentElement.offsetWidth && elemData[elemId].zoomLevel < 1 && !targetElement.isZoomed) { + resetZoom(); + } + } + + if (isExtension) { + targetElement.addEventListener("mousemove", checkForOutBox); + } + + + window.addEventListener('resize', (e) => { + resetZoom(); + + if (isExtension) { + targetElement.isExpanded = false; + targetElement.isZoomed = false; + } + }); + + gradioApp().addEventListener("mousemove", handleMoveByKey); + + + } + + applyZoomAndPan(elementIDs.sketch, false); + applyZoomAndPan(elementIDs.inpaint, false); + applyZoomAndPan(elementIDs.inpaintSketch, false); + + // Make the function global so that other extensions can take advantage of this solution + const applyZoomAndPanIntegration = async(id, elementIDs) => { + const mainEl = document.querySelector(id); + if (id.toLocaleLowerCase() === "none") { + for (const elementID of elementIDs) { + const el = await waitForElement(elementID); + if (!el) break; + applyZoomAndPan(elementID); + } + return; + } + + if (!mainEl) return; + mainEl.addEventListener("click", async() => { + for (const elementID of elementIDs) { + const el = await waitForElement(elementID); + if (!el) break; + applyZoomAndPan(elementID); + } + }, {once: true}); + }; + + window.applyZoomAndPan = applyZoomAndPan; // Only 1 elements, argument elementID, for example applyZoomAndPan("#txt2img_controlnet_ControlNet_input_image") + + window.applyZoomAndPanIntegration = applyZoomAndPanIntegration; // for any extension + + /* + The function `applyZoomAndPanIntegration` takes two arguments: + + 1. `id`: A string identifier for the element to which zoom and pan functionality will be applied on click. + If the `id` value is "none", the functionality will be applied to all elements specified in the second argument without a click event. + + 2. `elementIDs`: An array of string identifiers for elements. Zoom and pan functionality will be applied to each of these elements on click of the element specified by the first argument. + If "none" is specified in the first argument, the functionality will be applied to each of these elements without a click event. + + Example usage: + applyZoomAndPanIntegration("#txt2img_controlnet", ["#txt2img_controlnet_ControlNet_input_image"]); + In this example, zoom and pan functionality will be applied to the element with the identifier "txt2img_controlnet_ControlNet_input_image" upon clicking the element with the identifier "txt2img_controlnet". + */ + + // More examples + // Add integration with ControlNet txt2img One TAB + // applyZoomAndPanIntegration("#txt2img_controlnet", ["#txt2img_controlnet_ControlNet_input_image"]); + + // Add integration with ControlNet txt2img Tabs + // applyZoomAndPanIntegration("#txt2img_controlnet",Array.from({ length: 10 }, (_, i) => `#txt2img_controlnet_ControlNet-${i}_input_image`)); + + // Add integration with Inpaint Anything + // applyZoomAndPanIntegration("None", ["#ia_sam_image", "#ia_sel_mask"]); +}); diff --git a/stable-diffusion-webui/extensions-builtin/canvas-zoom-and-pan/scripts/hotkey_config.py b/stable-diffusion-webui/extensions-builtin/canvas-zoom-and-pan/scripts/hotkey_config.py new file mode 100755 index 0000000..17b27b2 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/canvas-zoom-and-pan/scripts/hotkey_config.py @@ -0,0 +1,17 @@ +import gradio as gr +from modules import shared + +shared.options_templates.update(shared.options_section(('canvas_hotkey', "Canvas Hotkeys"), { + "canvas_hotkey_zoom": shared.OptionInfo("Alt", "Zoom canvas", gr.Radio, {"choices": ["Shift","Ctrl", "Alt"]}).info("If you choose 'Shift' you cannot scroll horizontally, 'Alt' can cause a little trouble in firefox"), + "canvas_hotkey_adjust": shared.OptionInfo("Ctrl", "Adjust brush size", gr.Radio, {"choices": ["Shift","Ctrl", "Alt"]}).info("If you choose 'Shift' you cannot scroll horizontally, 'Alt' can cause a little trouble in firefox"), + "canvas_hotkey_shrink_brush": shared.OptionInfo("Q", "Shrink the brush size"), + "canvas_hotkey_grow_brush": shared.OptionInfo("W", "Enlarge the brush size"), + "canvas_hotkey_move": shared.OptionInfo("F", "Moving the canvas").info("To work correctly in firefox, turn off 'Automatically search the page text when typing' in the browser settings"), + "canvas_hotkey_fullscreen": shared.OptionInfo("S", "Fullscreen Mode, maximizes the picture so that it fits into the screen and stretches it to its full width "), + "canvas_hotkey_reset": shared.OptionInfo("R", "Reset zoom and canvas position"), + "canvas_hotkey_overlap": shared.OptionInfo("O", "Toggle overlap").info("Technical button, needed for testing"), + "canvas_show_tooltip": shared.OptionInfo(True, "Enable tooltip on the canvas"), + "canvas_auto_expand": shared.OptionInfo(True, "Automatically expands an image that does not fit completely in the canvas area, similar to manually pressing the S and R buttons"), + "canvas_blur_prompt": shared.OptionInfo(False, "Take the focus off the prompt when working with a canvas"), + "canvas_disabled_functions": shared.OptionInfo(["Overlap"], "Disable function that you don't use", gr.CheckboxGroup, {"choices": ["Zoom","Adjust brush size","Hotkey enlarge brush","Hotkey shrink brush","Moving canvas","Fullscreen","Reset Zoom","Overlap"]}), +})) diff --git a/stable-diffusion-webui/extensions-builtin/canvas-zoom-and-pan/style.css b/stable-diffusion-webui/extensions-builtin/canvas-zoom-and-pan/style.css new file mode 100755 index 0000000..5d8054e --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/canvas-zoom-and-pan/style.css @@ -0,0 +1,66 @@ +.canvas-tooltip-info { + position: absolute; + top: 10px; + left: 10px; + cursor: help; + background-color: rgba(0, 0, 0, 0.3); + width: 20px; + height: 20px; + border-radius: 50%; + display: flex; + align-items: center; + justify-content: center; + flex-direction: column; + + z-index: 100; +} + +.canvas-tooltip-info::after { + content: ''; + display: block; + width: 2px; + height: 7px; + background-color: white; + margin-top: 2px; +} + +.canvas-tooltip-info::before { + content: ''; + display: block; + width: 2px; + height: 2px; + background-color: white; +} + +.canvas-tooltip-content { + display: none; + background-color: #f9f9f9; + color: #333; + border: 1px solid #ddd; + padding: 15px; + position: absolute; + top: 40px; + left: 10px; + width: 250px; + font-size: 16px; + opacity: 0; + border-radius: 8px; + box-shadow: 0px 8px 16px 0px rgba(0,0,0,0.2); + + z-index: 100; +} + +.canvas-tooltip:hover .canvas-tooltip-content { + display: block; + animation: fadeIn 0.5s; + opacity: 1; +} + +@keyframes fadeIn { + from {opacity: 0;} + to {opacity: 1;} +} + +.styler { + overflow:inherit !important; +} \ No newline at end of file diff --git a/stable-diffusion-webui/extensions-builtin/extra-options-section/scripts/extra_options_section.py b/stable-diffusion-webui/extensions-builtin/extra-options-section/scripts/extra_options_section.py new file mode 100755 index 0000000..67dc3e2 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/extra-options-section/scripts/extra_options_section.py @@ -0,0 +1,82 @@ +import math + +import gradio as gr +from modules import scripts, shared, ui_components, ui_settings, infotext_utils, errors +from modules.ui_components import FormColumn + + +class ExtraOptionsSection(scripts.Script): + section = "extra_options" + + def __init__(self): + self.comps = None + self.setting_names = None + + def title(self): + return "Extra options" + + def show(self, is_img2img): + return scripts.AlwaysVisible + + def ui(self, is_img2img): + self.comps = [] + self.setting_names = [] + self.infotext_fields = [] + extra_options = shared.opts.extra_options_img2img if is_img2img else shared.opts.extra_options_txt2img + elem_id_tabname = "extra_options_" + ("img2img" if is_img2img else "txt2img") + + mapping = {k: v for v, k in infotext_utils.infotext_to_setting_name_mapping} + + with gr.Blocks() as interface: + with gr.Accordion("Options", open=False, elem_id=elem_id_tabname) if shared.opts.extra_options_accordion and extra_options else gr.Group(elem_id=elem_id_tabname): + + row_count = math.ceil(len(extra_options) / shared.opts.extra_options_cols) + + for row in range(row_count): + with gr.Row(): + for col in range(shared.opts.extra_options_cols): + index = row * shared.opts.extra_options_cols + col + if index >= len(extra_options): + break + + setting_name = extra_options[index] + + with FormColumn(): + try: + comp = ui_settings.create_setting_component(setting_name) + except KeyError: + errors.report(f"Can't add extra options for {setting_name} in ui") + continue + + self.comps.append(comp) + self.setting_names.append(setting_name) + + setting_infotext_name = mapping.get(setting_name) + if setting_infotext_name is not None: + self.infotext_fields.append((comp, setting_infotext_name)) + + def get_settings_values(): + res = [ui_settings.get_value_for_setting(key) for key in self.setting_names] + return res[0] if len(res) == 1 else res + + interface.load(fn=get_settings_values, inputs=[], outputs=self.comps, queue=False, show_progress=False) + + return self.comps + + def before_process(self, p, *args): + for name, value in zip(self.setting_names, args): + if name not in p.override_settings: + p.override_settings[name] = value + + +shared.options_templates.update(shared.options_section(('settings_in_ui', "Settings in UI", "ui"), { + "settings_in_ui": shared.OptionHTML(""" +This page allows you to add some settings to the main interface of txt2img and img2img tabs. +"""), + "extra_options_txt2img": shared.OptionInfo([], "Settings for txt2img", ui_components.DropdownMulti, lambda: {"choices": list(shared.opts.data_labels.keys())}).js("info", "settingsHintsShowQuicksettings").info("setting entries that also appear in txt2img interfaces").needs_reload_ui(), + "extra_options_img2img": shared.OptionInfo([], "Settings for img2img", ui_components.DropdownMulti, lambda: {"choices": list(shared.opts.data_labels.keys())}).js("info", "settingsHintsShowQuicksettings").info("setting entries that also appear in img2img interfaces").needs_reload_ui(), + "extra_options_cols": shared.OptionInfo(1, "Number of columns for added settings", gr.Slider, {"step": 1, "minimum": 1, "maximum": 20}).info("displayed amount will depend on the actual browser window width").needs_reload_ui(), + "extra_options_accordion": shared.OptionInfo(False, "Place added settings into an accordion").needs_reload_ui() +})) + + diff --git a/stable-diffusion-webui/extensions-builtin/hypertile/hypertile.py b/stable-diffusion-webui/extensions-builtin/hypertile/hypertile.py new file mode 100755 index 0000000..0f40e2d --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/hypertile/hypertile.py @@ -0,0 +1,351 @@ +""" +Hypertile module for splitting attention layers in SD-1.5 U-Net and SD-1.5 VAE +Warn: The patch works well only if the input image has a width and height that are multiples of 128 +Original author: @tfernd Github: https://github.com/tfernd/HyperTile +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable + +from functools import wraps, cache + +import math +import torch.nn as nn +import random + +from einops import rearrange + + +@dataclass +class HypertileParams: + depth = 0 + layer_name = "" + tile_size: int = 0 + swap_size: int = 0 + aspect_ratio: float = 1.0 + forward = None + enabled = False + + + +# TODO add SD-XL layers +DEPTH_LAYERS = { + 0: [ + # SD 1.5 U-Net (diffusers) + "down_blocks.0.attentions.0.transformer_blocks.0.attn1", + "down_blocks.0.attentions.1.transformer_blocks.0.attn1", + "up_blocks.3.attentions.0.transformer_blocks.0.attn1", + "up_blocks.3.attentions.1.transformer_blocks.0.attn1", + "up_blocks.3.attentions.2.transformer_blocks.0.attn1", + # SD 1.5 U-Net (ldm) + "input_blocks.1.1.transformer_blocks.0.attn1", + "input_blocks.2.1.transformer_blocks.0.attn1", + "output_blocks.9.1.transformer_blocks.0.attn1", + "output_blocks.10.1.transformer_blocks.0.attn1", + "output_blocks.11.1.transformer_blocks.0.attn1", + # SD 1.5 VAE + "decoder.mid_block.attentions.0", + "decoder.mid.attn_1", + ], + 1: [ + # SD 1.5 U-Net (diffusers) + "down_blocks.1.attentions.0.transformer_blocks.0.attn1", + "down_blocks.1.attentions.1.transformer_blocks.0.attn1", + "up_blocks.2.attentions.0.transformer_blocks.0.attn1", + "up_blocks.2.attentions.1.transformer_blocks.0.attn1", + "up_blocks.2.attentions.2.transformer_blocks.0.attn1", + # SD 1.5 U-Net (ldm) + "input_blocks.4.1.transformer_blocks.0.attn1", + "input_blocks.5.1.transformer_blocks.0.attn1", + "output_blocks.6.1.transformer_blocks.0.attn1", + "output_blocks.7.1.transformer_blocks.0.attn1", + "output_blocks.8.1.transformer_blocks.0.attn1", + ], + 2: [ + # SD 1.5 U-Net (diffusers) + "down_blocks.2.attentions.0.transformer_blocks.0.attn1", + "down_blocks.2.attentions.1.transformer_blocks.0.attn1", + "up_blocks.1.attentions.0.transformer_blocks.0.attn1", + "up_blocks.1.attentions.1.transformer_blocks.0.attn1", + "up_blocks.1.attentions.2.transformer_blocks.0.attn1", + # SD 1.5 U-Net (ldm) + "input_blocks.7.1.transformer_blocks.0.attn1", + "input_blocks.8.1.transformer_blocks.0.attn1", + "output_blocks.3.1.transformer_blocks.0.attn1", + "output_blocks.4.1.transformer_blocks.0.attn1", + "output_blocks.5.1.transformer_blocks.0.attn1", + ], + 3: [ + # SD 1.5 U-Net (diffusers) + "mid_block.attentions.0.transformer_blocks.0.attn1", + # SD 1.5 U-Net (ldm) + "middle_block.1.transformer_blocks.0.attn1", + ], +} +# XL layers, thanks for GitHub@gel-crabs for the help +DEPTH_LAYERS_XL = { + 0: [ + # SD 1.5 U-Net (diffusers) + "down_blocks.0.attentions.0.transformer_blocks.0.attn1", + "down_blocks.0.attentions.1.transformer_blocks.0.attn1", + "up_blocks.3.attentions.0.transformer_blocks.0.attn1", + "up_blocks.3.attentions.1.transformer_blocks.0.attn1", + "up_blocks.3.attentions.2.transformer_blocks.0.attn1", + # SD 1.5 U-Net (ldm) + "input_blocks.4.1.transformer_blocks.0.attn1", + "input_blocks.5.1.transformer_blocks.0.attn1", + "output_blocks.3.1.transformer_blocks.0.attn1", + "output_blocks.4.1.transformer_blocks.0.attn1", + "output_blocks.5.1.transformer_blocks.0.attn1", + # SD 1.5 VAE + "decoder.mid_block.attentions.0", + "decoder.mid.attn_1", + ], + 1: [ + # SD 1.5 U-Net (diffusers) + #"down_blocks.1.attentions.0.transformer_blocks.0.attn1", + #"down_blocks.1.attentions.1.transformer_blocks.0.attn1", + #"up_blocks.2.attentions.0.transformer_blocks.0.attn1", + #"up_blocks.2.attentions.1.transformer_blocks.0.attn1", + #"up_blocks.2.attentions.2.transformer_blocks.0.attn1", + # SD 1.5 U-Net (ldm) + "input_blocks.4.1.transformer_blocks.1.attn1", + "input_blocks.5.1.transformer_blocks.1.attn1", + "output_blocks.3.1.transformer_blocks.1.attn1", + "output_blocks.4.1.transformer_blocks.1.attn1", + "output_blocks.5.1.transformer_blocks.1.attn1", + "input_blocks.7.1.transformer_blocks.0.attn1", + "input_blocks.8.1.transformer_blocks.0.attn1", + "output_blocks.0.1.transformer_blocks.0.attn1", + "output_blocks.1.1.transformer_blocks.0.attn1", + "output_blocks.2.1.transformer_blocks.0.attn1", + "input_blocks.7.1.transformer_blocks.1.attn1", + "input_blocks.8.1.transformer_blocks.1.attn1", + "output_blocks.0.1.transformer_blocks.1.attn1", + "output_blocks.1.1.transformer_blocks.1.attn1", + "output_blocks.2.1.transformer_blocks.1.attn1", + "input_blocks.7.1.transformer_blocks.2.attn1", + "input_blocks.8.1.transformer_blocks.2.attn1", + "output_blocks.0.1.transformer_blocks.2.attn1", + "output_blocks.1.1.transformer_blocks.2.attn1", + "output_blocks.2.1.transformer_blocks.2.attn1", + "input_blocks.7.1.transformer_blocks.3.attn1", + "input_blocks.8.1.transformer_blocks.3.attn1", + "output_blocks.0.1.transformer_blocks.3.attn1", + "output_blocks.1.1.transformer_blocks.3.attn1", + "output_blocks.2.1.transformer_blocks.3.attn1", + "input_blocks.7.1.transformer_blocks.4.attn1", + "input_blocks.8.1.transformer_blocks.4.attn1", + "output_blocks.0.1.transformer_blocks.4.attn1", + "output_blocks.1.1.transformer_blocks.4.attn1", + "output_blocks.2.1.transformer_blocks.4.attn1", + "input_blocks.7.1.transformer_blocks.5.attn1", + "input_blocks.8.1.transformer_blocks.5.attn1", + "output_blocks.0.1.transformer_blocks.5.attn1", + "output_blocks.1.1.transformer_blocks.5.attn1", + "output_blocks.2.1.transformer_blocks.5.attn1", + "input_blocks.7.1.transformer_blocks.6.attn1", + "input_blocks.8.1.transformer_blocks.6.attn1", + "output_blocks.0.1.transformer_blocks.6.attn1", + "output_blocks.1.1.transformer_blocks.6.attn1", + "output_blocks.2.1.transformer_blocks.6.attn1", + "input_blocks.7.1.transformer_blocks.7.attn1", + "input_blocks.8.1.transformer_blocks.7.attn1", + "output_blocks.0.1.transformer_blocks.7.attn1", + "output_blocks.1.1.transformer_blocks.7.attn1", + "output_blocks.2.1.transformer_blocks.7.attn1", + "input_blocks.7.1.transformer_blocks.8.attn1", + "input_blocks.8.1.transformer_blocks.8.attn1", + "output_blocks.0.1.transformer_blocks.8.attn1", + "output_blocks.1.1.transformer_blocks.8.attn1", + "output_blocks.2.1.transformer_blocks.8.attn1", + "input_blocks.7.1.transformer_blocks.9.attn1", + "input_blocks.8.1.transformer_blocks.9.attn1", + "output_blocks.0.1.transformer_blocks.9.attn1", + "output_blocks.1.1.transformer_blocks.9.attn1", + "output_blocks.2.1.transformer_blocks.9.attn1", + ], + 2: [ + # SD 1.5 U-Net (diffusers) + "mid_block.attentions.0.transformer_blocks.0.attn1", + # SD 1.5 U-Net (ldm) + "middle_block.1.transformer_blocks.0.attn1", + "middle_block.1.transformer_blocks.1.attn1", + "middle_block.1.transformer_blocks.2.attn1", + "middle_block.1.transformer_blocks.3.attn1", + "middle_block.1.transformer_blocks.4.attn1", + "middle_block.1.transformer_blocks.5.attn1", + "middle_block.1.transformer_blocks.6.attn1", + "middle_block.1.transformer_blocks.7.attn1", + "middle_block.1.transformer_blocks.8.attn1", + "middle_block.1.transformer_blocks.9.attn1", + ], + 3 : [] # TODO - separate layers for SD-XL +} + + +RNG_INSTANCE = random.Random() + +@cache +def get_divisors(value: int, min_value: int, /, max_options: int = 1) -> list[int]: + """ + Returns divisors of value that + x * min_value <= value + in big -> small order, amount of divisors is limited by max_options + """ + max_options = max(1, max_options) # at least 1 option should be returned + min_value = min(min_value, value) + divisors = [i for i in range(min_value, value + 1) if value % i == 0] # divisors in small -> big order + ns = [value // i for i in divisors[:max_options]] # has at least 1 element # big -> small order + return ns + + +def random_divisor(value: int, min_value: int, /, max_options: int = 1) -> int: + """ + Returns a random divisor of value that + x * min_value <= value + if max_options is 1, the behavior is deterministic + """ + ns = get_divisors(value, min_value, max_options=max_options) # get cached divisors + idx = RNG_INSTANCE.randint(0, len(ns) - 1) + + return ns[idx] + + +def set_hypertile_seed(seed: int) -> None: + RNG_INSTANCE.seed(seed) + + +@cache +def largest_tile_size_available(width: int, height: int) -> int: + """ + Calculates the largest tile size available for a given width and height + Tile size is always a power of 2 + """ + gcd = math.gcd(width, height) + largest_tile_size_available = 1 + while gcd % (largest_tile_size_available * 2) == 0: + largest_tile_size_available *= 2 + return largest_tile_size_available + + +def iterative_closest_divisors(hw:int, aspect_ratio:float) -> tuple[int, int]: + """ + Finds h and w such that h*w = hw and h/w = aspect_ratio + We check all possible divisors of hw and return the closest to the aspect ratio + """ + divisors = [i for i in range(2, hw + 1) if hw % i == 0] # all divisors of hw + pairs = [(i, hw // i) for i in divisors] # all pairs of divisors of hw + ratios = [w/h for h, w in pairs] # all ratios of pairs of divisors of hw + closest_ratio = min(ratios, key=lambda x: abs(x - aspect_ratio)) # closest ratio to aspect_ratio + closest_pair = pairs[ratios.index(closest_ratio)] # closest pair of divisors to aspect_ratio + return closest_pair + + +@cache +def find_hw_candidates(hw:int, aspect_ratio:float) -> tuple[int, int]: + """ + Finds h and w such that h*w = hw and h/w = aspect_ratio + """ + h, w = round(math.sqrt(hw * aspect_ratio)), round(math.sqrt(hw / aspect_ratio)) + # find h and w such that h*w = hw and h/w = aspect_ratio + if h * w != hw: + w_candidate = hw / h + # check if w is an integer + if not w_candidate.is_integer(): + h_candidate = hw / w + # check if h is an integer + if not h_candidate.is_integer(): + return iterative_closest_divisors(hw, aspect_ratio) + else: + h = int(h_candidate) + else: + w = int(w_candidate) + return h, w + + +def self_attn_forward(params: HypertileParams, scale_depth=True) -> Callable: + + @wraps(params.forward) + def wrapper(*args, **kwargs): + if not params.enabled: + return params.forward(*args, **kwargs) + + latent_tile_size = max(128, params.tile_size) // 8 + x = args[0] + + # VAE + if x.ndim == 4: + b, c, h, w = x.shape + + nh = random_divisor(h, latent_tile_size, params.swap_size) + nw = random_divisor(w, latent_tile_size, params.swap_size) + + if nh * nw > 1: + x = rearrange(x, "b c (nh h) (nw w) -> (b nh nw) c h w", nh=nh, nw=nw) # split into nh * nw tiles + + out = params.forward(x, *args[1:], **kwargs) + + if nh * nw > 1: + out = rearrange(out, "(b nh nw) c h w -> b c (nh h) (nw w)", nh=nh, nw=nw) + + # U-Net + else: + hw: int = x.size(1) + h, w = find_hw_candidates(hw, params.aspect_ratio) + assert h * w == hw, f"Invalid aspect ratio {params.aspect_ratio} for input of shape {x.shape}, hw={hw}, h={h}, w={w}" + + factor = 2 ** params.depth if scale_depth else 1 + nh = random_divisor(h, latent_tile_size * factor, params.swap_size) + nw = random_divisor(w, latent_tile_size * factor, params.swap_size) + + if nh * nw > 1: + x = rearrange(x, "b (nh h nw w) c -> (b nh nw) (h w) c", h=h // nh, w=w // nw, nh=nh, nw=nw) + + out = params.forward(x, *args[1:], **kwargs) + + if nh * nw > 1: + out = rearrange(out, "(b nh nw) hw c -> b nh nw hw c", nh=nh, nw=nw) + out = rearrange(out, "b nh nw (h w) c -> b (nh h nw w) c", h=h // nh, w=w // nw) + + return out + + return wrapper + + +def hypertile_hook_model(model: nn.Module, width, height, *, enable=False, tile_size_max=128, swap_size=1, max_depth=3, is_sdxl=False): + hypertile_layers = getattr(model, "__webui_hypertile_layers", None) + if hypertile_layers is None: + if not enable: + return + + hypertile_layers = {} + layers = DEPTH_LAYERS_XL if is_sdxl else DEPTH_LAYERS + + for depth in range(4): + for layer_name, module in model.named_modules(): + if any(layer_name.endswith(try_name) for try_name in layers[depth]): + params = HypertileParams() + module.__webui_hypertile_params = params + params.forward = module.forward + params.depth = depth + params.layer_name = layer_name + module.forward = self_attn_forward(params) + + hypertile_layers[layer_name] = 1 + + model.__webui_hypertile_layers = hypertile_layers + + aspect_ratio = width / height + tile_size = min(largest_tile_size_available(width, height), tile_size_max) + + for layer_name, module in model.named_modules(): + if layer_name in hypertile_layers: + params = module.__webui_hypertile_params + + params.tile_size = tile_size + params.swap_size = swap_size + params.aspect_ratio = aspect_ratio + params.enabled = enable and params.depth <= max_depth diff --git a/stable-diffusion-webui/extensions-builtin/hypertile/scripts/hypertile_script.py b/stable-diffusion-webui/extensions-builtin/hypertile/scripts/hypertile_script.py new file mode 100755 index 0000000..2c62267 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/hypertile/scripts/hypertile_script.py @@ -0,0 +1,122 @@ +import hypertile +from modules import scripts, script_callbacks, shared + + +class ScriptHypertile(scripts.Script): + name = "Hypertile" + + def title(self): + return self.name + + def show(self, is_img2img): + return scripts.AlwaysVisible + + def process(self, p, *args): + hypertile.set_hypertile_seed(p.all_seeds[0]) + + configure_hypertile(p.width, p.height, enable_unet=shared.opts.hypertile_enable_unet) + + self.add_infotext(p) + + def before_hr(self, p, *args): + + enable = shared.opts.hypertile_enable_unet_secondpass or shared.opts.hypertile_enable_unet + + # exclusive hypertile seed for the second pass + if enable: + hypertile.set_hypertile_seed(p.all_seeds[0]) + + configure_hypertile(p.hr_upscale_to_x, p.hr_upscale_to_y, enable_unet=enable) + + if enable and not shared.opts.hypertile_enable_unet: + p.extra_generation_params["Hypertile U-Net second pass"] = True + + self.add_infotext(p, add_unet_params=True) + + def add_infotext(self, p, add_unet_params=False): + def option(name): + value = getattr(shared.opts, name) + default_value = shared.opts.get_default(name) + return None if value == default_value else value + + if shared.opts.hypertile_enable_unet: + p.extra_generation_params["Hypertile U-Net"] = True + + if shared.opts.hypertile_enable_unet or add_unet_params: + p.extra_generation_params["Hypertile U-Net max depth"] = option('hypertile_max_depth_unet') + p.extra_generation_params["Hypertile U-Net max tile size"] = option('hypertile_max_tile_unet') + p.extra_generation_params["Hypertile U-Net swap size"] = option('hypertile_swap_size_unet') + + if shared.opts.hypertile_enable_vae: + p.extra_generation_params["Hypertile VAE"] = True + p.extra_generation_params["Hypertile VAE max depth"] = option('hypertile_max_depth_vae') + p.extra_generation_params["Hypertile VAE max tile size"] = option('hypertile_max_tile_vae') + p.extra_generation_params["Hypertile VAE swap size"] = option('hypertile_swap_size_vae') + + +def configure_hypertile(width, height, enable_unet=True): + hypertile.hypertile_hook_model( + shared.sd_model.first_stage_model, + width, + height, + swap_size=shared.opts.hypertile_swap_size_vae, + max_depth=shared.opts.hypertile_max_depth_vae, + tile_size_max=shared.opts.hypertile_max_tile_vae, + enable=shared.opts.hypertile_enable_vae, + ) + + hypertile.hypertile_hook_model( + shared.sd_model.model, + width, + height, + swap_size=shared.opts.hypertile_swap_size_unet, + max_depth=shared.opts.hypertile_max_depth_unet, + tile_size_max=shared.opts.hypertile_max_tile_unet, + enable=enable_unet, + is_sdxl=shared.sd_model.is_sdxl + ) + + +def on_ui_settings(): + import gradio as gr + + options = { + "hypertile_explanation": shared.OptionHTML(""" + Hypertile optimizes the self-attention layer within U-Net and VAE models, + resulting in a reduction in computation time ranging from 1 to 4 times. The larger the generated image is, the greater the + benefit. + """), + + "hypertile_enable_unet": shared.OptionInfo(False, "Enable Hypertile U-Net", infotext="Hypertile U-Net").info("enables hypertile for all modes, including hires fix second pass; noticeable change in details of the generated picture"), + "hypertile_enable_unet_secondpass": shared.OptionInfo(False, "Enable Hypertile U-Net for hires fix second pass", infotext="Hypertile U-Net second pass").info("enables hypertile just for hires fix second pass - regardless of whether the above setting is enabled"), + "hypertile_max_depth_unet": shared.OptionInfo(3, "Hypertile U-Net max depth", gr.Slider, {"minimum": 0, "maximum": 3, "step": 1}, infotext="Hypertile U-Net max depth").info("larger = more neural network layers affected; minor effect on performance"), + "hypertile_max_tile_unet": shared.OptionInfo(256, "Hypertile U-Net max tile size", gr.Slider, {"minimum": 0, "maximum": 512, "step": 16}, infotext="Hypertile U-Net max tile size").info("larger = worse performance"), + "hypertile_swap_size_unet": shared.OptionInfo(3, "Hypertile U-Net swap size", gr.Slider, {"minimum": 0, "maximum": 64, "step": 1}, infotext="Hypertile U-Net swap size"), + "hypertile_enable_vae": shared.OptionInfo(False, "Enable Hypertile VAE", infotext="Hypertile VAE").info("minimal change in the generated picture"), + "hypertile_max_depth_vae": shared.OptionInfo(3, "Hypertile VAE max depth", gr.Slider, {"minimum": 0, "maximum": 3, "step": 1}, infotext="Hypertile VAE max depth"), + "hypertile_max_tile_vae": shared.OptionInfo(128, "Hypertile VAE max tile size", gr.Slider, {"minimum": 0, "maximum": 512, "step": 16}, infotext="Hypertile VAE max tile size"), + "hypertile_swap_size_vae": shared.OptionInfo(3, "Hypertile VAE swap size ", gr.Slider, {"minimum": 0, "maximum": 64, "step": 1}, infotext="Hypertile VAE swap size"), + } + + for name, opt in options.items(): + opt.section = ('hypertile', "Hypertile") + shared.opts.add_option(name, opt) + + +def add_axis_options(): + xyz_grid = [x for x in scripts.scripts_data if x.script_class.__module__ == "xyz_grid.py"][0].module + xyz_grid.axis_options.extend([ + xyz_grid.AxisOption("[Hypertile] Unet First pass Enabled", str, xyz_grid.apply_override('hypertile_enable_unet', boolean=True), choices=xyz_grid.boolean_choice(reverse=True)), + xyz_grid.AxisOption("[Hypertile] Unet Second pass Enabled", str, xyz_grid.apply_override('hypertile_enable_unet_secondpass', boolean=True), choices=xyz_grid.boolean_choice(reverse=True)), + xyz_grid.AxisOption("[Hypertile] Unet Max Depth", int, xyz_grid.apply_override("hypertile_max_depth_unet"), confirm=xyz_grid.confirm_range(0, 3, '[Hypertile] Unet Max Depth'), choices=lambda: [str(x) for x in range(4)]), + xyz_grid.AxisOption("[Hypertile] Unet Max Tile Size", int, xyz_grid.apply_override("hypertile_max_tile_unet"), confirm=xyz_grid.confirm_range(0, 512, '[Hypertile] Unet Max Tile Size')), + xyz_grid.AxisOption("[Hypertile] Unet Swap Size", int, xyz_grid.apply_override("hypertile_swap_size_unet"), confirm=xyz_grid.confirm_range(0, 64, '[Hypertile] Unet Swap Size')), + xyz_grid.AxisOption("[Hypertile] VAE Enabled", str, xyz_grid.apply_override('hypertile_enable_vae', boolean=True), choices=xyz_grid.boolean_choice(reverse=True)), + xyz_grid.AxisOption("[Hypertile] VAE Max Depth", int, xyz_grid.apply_override("hypertile_max_depth_vae"), confirm=xyz_grid.confirm_range(0, 3, '[Hypertile] VAE Max Depth'), choices=lambda: [str(x) for x in range(4)]), + xyz_grid.AxisOption("[Hypertile] VAE Max Tile Size", int, xyz_grid.apply_override("hypertile_max_tile_vae"), confirm=xyz_grid.confirm_range(0, 512, '[Hypertile] VAE Max Tile Size')), + xyz_grid.AxisOption("[Hypertile] VAE Swap Size", int, xyz_grid.apply_override("hypertile_swap_size_vae"), confirm=xyz_grid.confirm_range(0, 64, '[Hypertile] VAE Swap Size')), + ]) + + +script_callbacks.on_ui_settings(on_ui_settings) +script_callbacks.on_before_ui(add_axis_options) diff --git a/stable-diffusion-webui/extensions-builtin/mobile/javascript/mobile.js b/stable-diffusion-webui/extensions-builtin/mobile/javascript/mobile.js new file mode 100755 index 0000000..bff1ace --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/mobile/javascript/mobile.js @@ -0,0 +1,34 @@ +var isSetupForMobile = false; + +function isMobile() { + for (var tab of ["txt2img", "img2img"]) { + var imageTab = gradioApp().getElementById(tab + '_results'); + if (imageTab && imageTab.offsetParent && imageTab.offsetLeft == 0) { + return true; + } + } + + return false; +} + +function reportWindowSize() { + if (gradioApp().querySelector('.toprow-compact-tools')) return; // not applicable for compact prompt layout + + var currentlyMobile = isMobile(); + if (currentlyMobile == isSetupForMobile) return; + isSetupForMobile = currentlyMobile; + + for (var tab of ["txt2img", "img2img"]) { + var button = gradioApp().getElementById(tab + '_generate_box'); + var target = gradioApp().getElementById(currentlyMobile ? tab + '_results' : tab + '_actions_column'); + target.insertBefore(button, target.firstElementChild); + + gradioApp().getElementById(tab + '_results').classList.toggle('mobile', currentlyMobile); + } +} + +window.addEventListener("resize", reportWindowSize); + +onUiLoaded(function() { + reportWindowSize(); +}); diff --git a/stable-diffusion-webui/extensions-builtin/postprocessing-for-training/scripts/postprocessing_autosized_crop.py b/stable-diffusion-webui/extensions-builtin/postprocessing-for-training/scripts/postprocessing_autosized_crop.py new file mode 100755 index 0000000..1e83de6 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/postprocessing-for-training/scripts/postprocessing_autosized_crop.py @@ -0,0 +1,64 @@ +from PIL import Image + +from modules import scripts_postprocessing, ui_components +import gradio as gr + + +def center_crop(image: Image, w: int, h: int): + iw, ih = image.size + if ih / h < iw / w: + sw = w * ih / h + box = (iw - sw) / 2, 0, iw - (iw - sw) / 2, ih + else: + sh = h * iw / w + box = 0, (ih - sh) / 2, iw, ih - (ih - sh) / 2 + return image.resize((w, h), Image.Resampling.LANCZOS, box) + + +def multicrop_pic(image: Image, mindim, maxdim, minarea, maxarea, objective, threshold): + iw, ih = image.size + err = lambda w, h: 1 - (lambda x: x if x < 1 else 1 / x)(iw / ih / (w / h)) + wh = max(((w, h) for w in range(mindim, maxdim + 1, 64) for h in range(mindim, maxdim + 1, 64) + if minarea <= w * h <= maxarea and err(w, h) <= threshold), + key=lambda wh: (wh[0] * wh[1], -err(*wh))[::1 if objective == 'Maximize area' else -1], + default=None + ) + return wh and center_crop(image, *wh) + + +class ScriptPostprocessingAutosizedCrop(scripts_postprocessing.ScriptPostprocessing): + name = "Auto-sized crop" + order = 4020 + + def ui(self): + with ui_components.InputAccordion(False, label="Auto-sized crop") as enable: + gr.Markdown('Each image is center-cropped with an automatically chosen width and height.') + with gr.Row(): + mindim = gr.Slider(minimum=64, maximum=2048, step=8, label="Dimension lower bound", value=384, elem_id="postprocess_multicrop_mindim") + maxdim = gr.Slider(minimum=64, maximum=2048, step=8, label="Dimension upper bound", value=768, elem_id="postprocess_multicrop_maxdim") + with gr.Row(): + minarea = gr.Slider(minimum=64 * 64, maximum=2048 * 2048, step=1, label="Area lower bound", value=64 * 64, elem_id="postprocess_multicrop_minarea") + maxarea = gr.Slider(minimum=64 * 64, maximum=2048 * 2048, step=1, label="Area upper bound", value=640 * 640, elem_id="postprocess_multicrop_maxarea") + with gr.Row(): + objective = gr.Radio(["Maximize area", "Minimize error"], value="Maximize area", label="Resizing objective", elem_id="postprocess_multicrop_objective") + threshold = gr.Slider(minimum=0, maximum=1, step=0.01, label="Error threshold", value=0.1, elem_id="postprocess_multicrop_threshold") + + return { + "enable": enable, + "mindim": mindim, + "maxdim": maxdim, + "minarea": minarea, + "maxarea": maxarea, + "objective": objective, + "threshold": threshold, + } + + def process(self, pp: scripts_postprocessing.PostprocessedImage, enable, mindim, maxdim, minarea, maxarea, objective, threshold): + if not enable: + return + + cropped = multicrop_pic(pp.image, mindim, maxdim, minarea, maxarea, objective, threshold) + if cropped is not None: + pp.image = cropped + else: + print(f"skipped {pp.image.width}x{pp.image.height} image (can't find suitable size within error threshold)") diff --git a/stable-diffusion-webui/extensions-builtin/postprocessing-for-training/scripts/postprocessing_caption.py b/stable-diffusion-webui/extensions-builtin/postprocessing-for-training/scripts/postprocessing_caption.py new file mode 100755 index 0000000..758222a --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/postprocessing-for-training/scripts/postprocessing_caption.py @@ -0,0 +1,30 @@ +from modules import scripts_postprocessing, ui_components, deepbooru, shared +import gradio as gr + + +class ScriptPostprocessingCeption(scripts_postprocessing.ScriptPostprocessing): + name = "Caption" + order = 4040 + + def ui(self): + with ui_components.InputAccordion(False, label="Caption") as enable: + option = gr.CheckboxGroup(value=["Deepbooru"], choices=["Deepbooru", "BLIP"], show_label=False) + + return { + "enable": enable, + "option": option, + } + + def process(self, pp: scripts_postprocessing.PostprocessedImage, enable, option): + if not enable: + return + + captions = [pp.caption] + + if "Deepbooru" in option: + captions.append(deepbooru.model.tag(pp.image)) + + if "BLIP" in option: + captions.append(shared.interrogator.interrogate(pp.image.convert("RGB"))) + + pp.caption = ", ".join([x for x in captions if x]) diff --git a/stable-diffusion-webui/extensions-builtin/postprocessing-for-training/scripts/postprocessing_create_flipped_copies.py b/stable-diffusion-webui/extensions-builtin/postprocessing-for-training/scripts/postprocessing_create_flipped_copies.py new file mode 100755 index 0000000..e7bd340 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/postprocessing-for-training/scripts/postprocessing_create_flipped_copies.py @@ -0,0 +1,32 @@ +from PIL import ImageOps, Image + +from modules import scripts_postprocessing, ui_components +import gradio as gr + + +class ScriptPostprocessingCreateFlippedCopies(scripts_postprocessing.ScriptPostprocessing): + name = "Create flipped copies" + order = 4030 + + def ui(self): + with ui_components.InputAccordion(False, label="Create flipped copies") as enable: + with gr.Row(): + option = gr.CheckboxGroup(value=["Horizontal"], choices=["Horizontal", "Vertical", "Both"], show_label=False) + + return { + "enable": enable, + "option": option, + } + + def process(self, pp: scripts_postprocessing.PostprocessedImage, enable, option): + if not enable: + return + + if "Horizontal" in option: + pp.extra_images.append(ImageOps.mirror(pp.image)) + + if "Vertical" in option: + pp.extra_images.append(pp.image.transpose(Image.Transpose.FLIP_TOP_BOTTOM)) + + if "Both" in option: + pp.extra_images.append(pp.image.transpose(Image.Transpose.FLIP_TOP_BOTTOM).transpose(Image.Transpose.FLIP_LEFT_RIGHT)) diff --git a/stable-diffusion-webui/extensions-builtin/postprocessing-for-training/scripts/postprocessing_focal_crop.py b/stable-diffusion-webui/extensions-builtin/postprocessing-for-training/scripts/postprocessing_focal_crop.py new file mode 100755 index 0000000..08fd2cc --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/postprocessing-for-training/scripts/postprocessing_focal_crop.py @@ -0,0 +1,54 @@ + +from modules import scripts_postprocessing, ui_components, errors +import gradio as gr + +from modules.textual_inversion import autocrop + + +class ScriptPostprocessingFocalCrop(scripts_postprocessing.ScriptPostprocessing): + name = "Auto focal point crop" + order = 4010 + + def ui(self): + with ui_components.InputAccordion(False, label="Auto focal point crop") as enable: + face_weight = gr.Slider(label='Focal point face weight', value=0.9, minimum=0.0, maximum=1.0, step=0.05, elem_id="postprocess_focal_crop_face_weight") + entropy_weight = gr.Slider(label='Focal point entropy weight', value=0.15, minimum=0.0, maximum=1.0, step=0.05, elem_id="postprocess_focal_crop_entropy_weight") + edges_weight = gr.Slider(label='Focal point edges weight', value=0.5, minimum=0.0, maximum=1.0, step=0.05, elem_id="postprocess_focal_crop_edges_weight") + debug = gr.Checkbox(label='Create debug image', elem_id="train_process_focal_crop_debug") + + return { + "enable": enable, + "face_weight": face_weight, + "entropy_weight": entropy_weight, + "edges_weight": edges_weight, + "debug": debug, + } + + def process(self, pp: scripts_postprocessing.PostprocessedImage, enable, face_weight, entropy_weight, edges_weight, debug): + if not enable: + return + + if not pp.shared.target_width or not pp.shared.target_height: + return + + dnn_model_path = None + try: + dnn_model_path = autocrop.download_and_cache_models() + except Exception: + errors.report("Unable to load face detection model for auto crop selection. Falling back to lower quality haar method.", exc_info=True) + + autocrop_settings = autocrop.Settings( + crop_width=pp.shared.target_width, + crop_height=pp.shared.target_height, + face_points_weight=face_weight, + entropy_points_weight=entropy_weight, + corner_points_weight=edges_weight, + annotate_image=debug, + dnn_model_path=dnn_model_path, + ) + + result, *others = autocrop.crop_image(pp.image, autocrop_settings) + + pp.image = result + pp.extra_images = [pp.create_copy(x, nametags=["focal-crop-debug"], disable_processing=True) for x in others] + diff --git a/stable-diffusion-webui/extensions-builtin/postprocessing-for-training/scripts/postprocessing_split_oversized.py b/stable-diffusion-webui/extensions-builtin/postprocessing-for-training/scripts/postprocessing_split_oversized.py new file mode 100755 index 0000000..888740e --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/postprocessing-for-training/scripts/postprocessing_split_oversized.py @@ -0,0 +1,71 @@ +import math + +from modules import scripts_postprocessing, ui_components +import gradio as gr + + +def split_pic(image, inverse_xy, width, height, overlap_ratio): + if inverse_xy: + from_w, from_h = image.height, image.width + to_w, to_h = height, width + else: + from_w, from_h = image.width, image.height + to_w, to_h = width, height + h = from_h * to_w // from_w + if inverse_xy: + image = image.resize((h, to_w)) + else: + image = image.resize((to_w, h)) + + split_count = math.ceil((h - to_h * overlap_ratio) / (to_h * (1.0 - overlap_ratio))) + y_step = (h - to_h) / (split_count - 1) + for i in range(split_count): + y = int(y_step * i) + if inverse_xy: + splitted = image.crop((y, 0, y + to_h, to_w)) + else: + splitted = image.crop((0, y, to_w, y + to_h)) + yield splitted + + +class ScriptPostprocessingSplitOversized(scripts_postprocessing.ScriptPostprocessing): + name = "Split oversized images" + order = 4000 + + def ui(self): + with ui_components.InputAccordion(False, label="Split oversized images") as enable: + with gr.Row(): + split_threshold = gr.Slider(label='Threshold', value=0.5, minimum=0.0, maximum=1.0, step=0.05, elem_id="postprocess_split_threshold") + overlap_ratio = gr.Slider(label='Overlap ratio', value=0.2, minimum=0.0, maximum=0.9, step=0.05, elem_id="postprocess_overlap_ratio") + + return { + "enable": enable, + "split_threshold": split_threshold, + "overlap_ratio": overlap_ratio, + } + + def process(self, pp: scripts_postprocessing.PostprocessedImage, enable, split_threshold, overlap_ratio): + if not enable: + return + + width = pp.shared.target_width + height = pp.shared.target_height + + if not width or not height: + return + + if pp.image.height > pp.image.width: + ratio = (pp.image.width * height) / (pp.image.height * width) + inverse_xy = False + else: + ratio = (pp.image.height * width) / (pp.image.width * height) + inverse_xy = True + + if ratio >= 1.0 or ratio > split_threshold: + return + + result, *others = split_pic(pp.image, inverse_xy, width, height, overlap_ratio) + + pp.image = result + pp.extra_images = [pp.create_copy(x) for x in others] + diff --git a/stable-diffusion-webui/extensions-builtin/prompt-bracket-checker/javascript/prompt-bracket-checker.js b/stable-diffusion-webui/extensions-builtin/prompt-bracket-checker/javascript/prompt-bracket-checker.js new file mode 100755 index 0000000..114cf94 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/prompt-bracket-checker/javascript/prompt-bracket-checker.js @@ -0,0 +1,42 @@ +// Stable Diffusion WebUI - Bracket checker +// By Hingashi no Florin/Bwin4L & @akx +// Counts open and closed brackets (round, square, curly) in the prompt and negative prompt text boxes in the txt2img and img2img tabs. +// If there's a mismatch, the keyword counter turns red and if you hover on it, a tooltip tells you what's wrong. + +function checkBrackets(textArea, counterElt) { + var counts = {}; + (textArea.value.match(/[(){}[\]]/g) || []).forEach(bracket => { + counts[bracket] = (counts[bracket] || 0) + 1; + }); + var errors = []; + + function checkPair(open, close, kind) { + if (counts[open] !== counts[close]) { + errors.push( + `${open}...${close} - Detected ${counts[open] || 0} opening and ${counts[close] || 0} closing ${kind}.` + ); + } + } + + checkPair('(', ')', 'round brackets'); + checkPair('[', ']', 'square brackets'); + checkPair('{', '}', 'curly brackets'); + counterElt.title = errors.join('\n'); + counterElt.classList.toggle('error', errors.length !== 0); +} + +function setupBracketChecking(id_prompt, id_counter) { + var textarea = gradioApp().querySelector("#" + id_prompt + " > label > textarea"); + var counter = gradioApp().getElementById(id_counter); + + if (textarea && counter) { + textarea.addEventListener("input", () => checkBrackets(textarea, counter)); + } +} + +onUiLoaded(function() { + setupBracketChecking('txt2img_prompt', 'txt2img_token_counter'); + setupBracketChecking('txt2img_neg_prompt', 'txt2img_negative_token_counter'); + setupBracketChecking('img2img_prompt', 'img2img_token_counter'); + setupBracketChecking('img2img_neg_prompt', 'img2img_negative_token_counter'); +}); diff --git a/stable-diffusion-webui/extensions-builtin/soft-inpainting/scripts/soft_inpainting.py b/stable-diffusion-webui/extensions-builtin/soft-inpainting/scripts/soft_inpainting.py new file mode 100755 index 0000000..0e62996 --- /dev/null +++ b/stable-diffusion-webui/extensions-builtin/soft-inpainting/scripts/soft_inpainting.py @@ -0,0 +1,760 @@ +import numpy as np +import gradio as gr +import math +from modules.ui_components import InputAccordion +import modules.scripts as scripts +from modules.torch_utils import float64 + + +class SoftInpaintingSettings: + def __init__(self, + mask_blend_power, + mask_blend_scale, + inpaint_detail_preservation, + composite_mask_influence, + composite_difference_threshold, + composite_difference_contrast): + self.mask_blend_power = mask_blend_power + self.mask_blend_scale = mask_blend_scale + self.inpaint_detail_preservation = inpaint_detail_preservation + self.composite_mask_influence = composite_mask_influence + self.composite_difference_threshold = composite_difference_threshold + self.composite_difference_contrast = composite_difference_contrast + + def add_generation_params(self, dest): + dest[enabled_gen_param_label] = True + dest[gen_param_labels.mask_blend_power] = self.mask_blend_power + dest[gen_param_labels.mask_blend_scale] = self.mask_blend_scale + dest[gen_param_labels.inpaint_detail_preservation] = self.inpaint_detail_preservation + dest[gen_param_labels.composite_mask_influence] = self.composite_mask_influence + dest[gen_param_labels.composite_difference_threshold] = self.composite_difference_threshold + dest[gen_param_labels.composite_difference_contrast] = self.composite_difference_contrast + + +# ------------------- Methods ------------------- + +def processing_uses_inpainting(p): + # TODO: Figure out a better way to determine if inpainting is being used by p + if getattr(p, "image_mask", None) is not None: + return True + + if getattr(p, "mask", None) is not None: + return True + + if getattr(p, "nmask", None) is not None: + return True + + return False + + +def latent_blend(settings, a, b, t): + """ + Interpolates two latent image representations according to the parameter t, + where the interpolated vectors' magnitudes are also interpolated separately. + The "detail_preservation" factor biases the magnitude interpolation towards + the larger of the two magnitudes. + """ + import torch + + # NOTE: We use inplace operations wherever possible. + + if len(t.shape) == 3: + # [4][w][h] to [1][4][w][h] + t2 = t.unsqueeze(0) + # [4][w][h] to [1][1][w][h] - the [4] seem redundant. + t3 = t[0].unsqueeze(0).unsqueeze(0) + else: + t2 = t + t3 = t[:, 0][:, None] + + one_minus_t2 = 1 - t2 + one_minus_t3 = 1 - t3 + + # Linearly interpolate the image vectors. + a_scaled = a * one_minus_t2 + b_scaled = b * t2 + image_interp = a_scaled + image_interp.add_(b_scaled) + result_type = image_interp.dtype + del a_scaled, b_scaled, t2, one_minus_t2 + + # Calculate the magnitude of the interpolated vectors. (We will remove this magnitude.) + # 64-bit operations are used here to allow large exponents. + current_magnitude = torch.norm(image_interp, p=2, dim=1, keepdim=True).to(float64(image_interp)).add_(0.00001) + + # Interpolate the powered magnitudes, then un-power them (bring them back to a power of 1). + a_magnitude = torch.norm(a, p=2, dim=1, keepdim=True).to(float64(a)).pow_(settings.inpaint_detail_preservation) * one_minus_t3 + b_magnitude = torch.norm(b, p=2, dim=1, keepdim=True).to(float64(b)).pow_(settings.inpaint_detail_preservation) * t3 + desired_magnitude = a_magnitude + desired_magnitude.add_(b_magnitude).pow_(1 / settings.inpaint_detail_preservation) + del a_magnitude, b_magnitude, t3, one_minus_t3 + + # Change the linearly interpolated image vectors' magnitudes to the value we want. + # This is the last 64-bit operation. + image_interp_scaling_factor = desired_magnitude + image_interp_scaling_factor.div_(current_magnitude) + image_interp_scaling_factor = image_interp_scaling_factor.to(result_type) + image_interp_scaled = image_interp + image_interp_scaled.mul_(image_interp_scaling_factor) + del current_magnitude + del desired_magnitude + del image_interp + del image_interp_scaling_factor + del result_type + + return image_interp_scaled + + +def get_modified_nmask(settings, nmask, sigma): + """ + Converts a negative mask representing the transparency of the original latent vectors being overlaid + to a mask that is scaled according to the denoising strength for this step. + + Where: + 0 = fully opaque, infinite density, fully masked + 1 = fully transparent, zero density, fully unmasked + + We bring this transparency to a power, as this allows one to simulate N number of blending operations + where N can be any positive real value. Using this one can control the balance of influence between + the denoiser and the original latents according to the sigma value. + + NOTE: "mask" is not used + """ + import torch + return torch.pow(nmask, (sigma ** settings.mask_blend_power) * settings.mask_blend_scale) + + +def apply_adaptive_masks( + settings: SoftInpaintingSettings, + nmask, + latent_orig, + latent_processed, + overlay_images, + width, height, + paste_to): + import torch + import modules.processing as proc + import modules.images as images + from PIL import Image, ImageOps, ImageFilter + + # TODO: Bias the blending according to the latent mask, add adjustable parameter for bias control. + if len(nmask.shape) == 3: + latent_mask = nmask[0].float() + else: + latent_mask = nmask[:, 0].float() + # convert the original mask into a form we use to scale distances for thresholding + mask_scalar = 1 - (torch.clamp(latent_mask, min=0, max=1) ** (settings.mask_blend_scale / 2)) + mask_scalar = (0.5 * (1 - settings.composite_mask_influence) + + mask_scalar * settings.composite_mask_influence) + mask_scalar = mask_scalar / (1.00001 - mask_scalar) + mask_scalar = mask_scalar.cpu().numpy() + + latent_distance = torch.norm(latent_processed - latent_orig, p=2, dim=1) + + kernel, kernel_center = get_gaussian_kernel(stddev_radius=1.5, max_radius=2) + + masks_for_overlay = [] + + for i, (distance_map, overlay_image) in enumerate(zip(latent_distance, overlay_images)): + converted_mask = distance_map.float().cpu().numpy() + converted_mask = weighted_histogram_filter(converted_mask, kernel, kernel_center, + percentile_min=0.9, percentile_max=1, min_width=1) + converted_mask = weighted_histogram_filter(converted_mask, kernel, kernel_center, + percentile_min=0.25, percentile_max=0.75, min_width=1) + + # The distance at which opacity of original decreases to 50% + if len(mask_scalar.shape) == 3: + if mask_scalar.shape[0] > i: + half_weighted_distance = settings.composite_difference_threshold * mask_scalar[i] + else: + half_weighted_distance = settings.composite_difference_threshold * mask_scalar[0] + else: + half_weighted_distance = settings.composite_difference_threshold * mask_scalar + + converted_mask = converted_mask / half_weighted_distance + + converted_mask = 1 / (1 + converted_mask ** settings.composite_difference_contrast) + converted_mask = smootherstep(converted_mask) + converted_mask = 1 - converted_mask + converted_mask = 255. * converted_mask + converted_mask = converted_mask.astype(np.uint8) + converted_mask = Image.fromarray(converted_mask) + converted_mask = images.resize_image(2, converted_mask, width, height) + converted_mask = proc.create_binary_mask(converted_mask, round=False) + + # Remove aliasing artifacts using a gaussian blur. + converted_mask = converted_mask.filter(ImageFilter.GaussianBlur(radius=4)) + + # Expand the mask to fit the whole image if needed. + if paste_to is not None: + converted_mask = proc.uncrop(converted_mask, + (overlay_image.width, overlay_image.height), + paste_to) + + masks_for_overlay.append(converted_mask) + + image_masked = Image.new('RGBa', (overlay_image.width, overlay_image.height)) + image_masked.paste(overlay_image.convert("RGBA").convert("RGBa"), + mask=ImageOps.invert(converted_mask.convert('L'))) + + overlay_images[i] = image_masked.convert('RGBA') + + return masks_for_overlay + + +def apply_masks( + settings, + nmask, + overlay_images, + width, height, + paste_to): + import torch + import modules.processing as proc + import modules.images as images + from PIL import Image, ImageOps, ImageFilter + + converted_mask = nmask[0].float() + converted_mask = torch.clamp(converted_mask, min=0, max=1).pow_(settings.mask_blend_scale / 2) + converted_mask = 255. * converted_mask + converted_mask = converted_mask.cpu().numpy().astype(np.uint8) + converted_mask = Image.fromarray(converted_mask) + converted_mask = images.resize_image(2, converted_mask, width, height) + converted_mask = proc.create_binary_mask(converted_mask, round=False) + + # Remove aliasing artifacts using a gaussian blur. + converted_mask = converted_mask.filter(ImageFilter.GaussianBlur(radius=4)) + + # Expand the mask to fit the whole image if needed. + if paste_to is not None: + converted_mask = proc.uncrop(converted_mask, + (width, height), + paste_to) + + masks_for_overlay = [] + + for i, overlay_image in enumerate(overlay_images): + masks_for_overlay[i] = converted_mask + + image_masked = Image.new('RGBa', (overlay_image.width, overlay_image.height)) + image_masked.paste(overlay_image.convert("RGBA").convert("RGBa"), + mask=ImageOps.invert(converted_mask.convert('L'))) + + overlay_images[i] = image_masked.convert('RGBA') + + return masks_for_overlay + + +def weighted_histogram_filter(img, kernel, kernel_center, percentile_min=0.0, percentile_max=1.0, min_width=1.0): + """ + Generalization convolution filter capable of applying + weighted mean, median, maximum, and minimum filters + parametrically using an arbitrary kernel. + + Args: + img (nparray): + The image, a 2-D array of floats, to which the filter is being applied. + kernel (nparray): + The kernel, a 2-D array of floats. + kernel_center (nparray): + The kernel center coordinate, a 1-D array with two elements. + percentile_min (float): + The lower bound of the histogram window used by the filter, + from 0 to 1. + percentile_max (float): + The upper bound of the histogram window used by the filter, + from 0 to 1. + min_width (float): + The minimum size of the histogram window bounds, in weight units. + Must be greater than 0. + + Returns: + (nparray): A filtered copy of the input image "img", a 2-D array of floats. + """ + + # Converts an index tuple into a vector. + def vec(x): + return np.array(x) + + kernel_min = -kernel_center + kernel_max = vec(kernel.shape) - kernel_center + + def weighted_histogram_filter_single(idx): + idx = vec(idx) + min_index = np.maximum(0, idx + kernel_min) + max_index = np.minimum(vec(img.shape), idx + kernel_max) + window_shape = max_index - min_index + + class WeightedElement: + """ + An element of the histogram, its weight + and bounds. + """ + + def __init__(self, value, weight): + self.value: float = value + self.weight: float = weight + self.window_min: float = 0.0 + self.window_max: float = 1.0 + + # Collect the values in the image as WeightedElements, + # weighted by their corresponding kernel values. + values = [] + for window_tup in np.ndindex(tuple(window_shape)): + window_index = vec(window_tup) + image_index = window_index + min_index + centered_kernel_index = image_index - idx + kernel_index = centered_kernel_index + kernel_center + element = WeightedElement(img[tuple(image_index)], kernel[tuple(kernel_index)]) + values.append(element) + + def sort_key(x: WeightedElement): + return x.value + + values.sort(key=sort_key) + + # Calculate the height of the stack (sum) + # and each sample's range they occupy in the stack + sum = 0 + for i in range(len(values)): + values[i].window_min = sum + sum += values[i].weight + values[i].window_max = sum + + # Calculate what range of this stack ("window") + # we want to get the weighted average across. + window_min = sum * percentile_min + window_max = sum * percentile_max + window_width = window_max - window_min + + # Ensure the window is within the stack and at least a certain size. + if window_width < min_width: + window_center = (window_min + window_max) / 2 + window_min = window_center - min_width / 2 + window_max = window_center + min_width / 2 + + if window_max > sum: + window_max = sum + window_min = sum - min_width + + if window_min < 0: + window_min = 0 + window_max = min_width + + value = 0 + value_weight = 0 + + # Get the weighted average of all the samples + # that overlap with the window, weighted + # by the size of their overlap. + for i in range(len(values)): + if window_min >= values[i].window_max: + continue + if window_max <= values[i].window_min: + break + + s = max(window_min, values[i].window_min) + e = min(window_max, values[i].window_max) + w = e - s + + value += values[i].value * w + value_weight += w + + return value / value_weight if value_weight != 0 else 0 + + img_out = img.copy() + + # Apply the kernel operation over each pixel. + for index in np.ndindex(img.shape): + img_out[index] = weighted_histogram_filter_single(index) + + return img_out + + +def smoothstep(x): + """ + The smoothstep function, input should be clamped to 0-1 range. + Turns a diagonal line (f(x) = x) into a sigmoid-like curve. + """ + return x * x * (3 - 2 * x) + + +def smootherstep(x): + """ + The smootherstep function, input should be clamped to 0-1 range. + Turns a diagonal line (f(x) = x) into a sigmoid-like curve. + """ + return x * x * x * (x * (6 * x - 15) + 10) + + +def get_gaussian_kernel(stddev_radius=1.0, max_radius=2): + """ + Creates a Gaussian kernel with thresholded edges. + + Args: + stddev_radius (float): + Standard deviation of the gaussian kernel, in pixels. + max_radius (int): + The size of the filter kernel. The number of pixels is (max_radius*2+1) ** 2. + The kernel is thresholded so that any values one pixel beyond this radius + is weighted at 0. + + Returns: + (nparray, nparray): A kernel array (shape: (N, N)), its center coordinate (shape: (2)) + """ + + # Evaluates a 0-1 normalized gaussian function for a given square distance from the mean. + def gaussian(sqr_mag): + return math.exp(-sqr_mag / (stddev_radius * stddev_radius)) + + # Helper function for converting a tuple to an array. + def vec(x): + return np.array(x) + + """ + Since a gaussian is unbounded, we need to limit ourselves + to a finite range. + We taper the ends off at the end of that range so they equal zero + while preserving the maximum value of 1 at the mean. + """ + zero_radius = max_radius + 1.0 + gauss_zero = gaussian(zero_radius * zero_radius) + gauss_kernel_scale = 1 / (1 - gauss_zero) + + def gaussian_kernel_func(coordinate): + x = coordinate[0] ** 2.0 + coordinate[1] ** 2.0 + x = gaussian(x) + x -= gauss_zero + x *= gauss_kernel_scale + x = max(0.0, x) + return x + + size = max_radius * 2 + 1 + kernel_center = max_radius + kernel = np.zeros((size, size)) + + for index in np.ndindex(kernel.shape): + kernel[index] = gaussian_kernel_func(vec(index) - kernel_center) + + return kernel, kernel_center + + +# ------------------- Constants ------------------- + + +default = SoftInpaintingSettings(1, 0.5, 4, 0, 0.5, 2) + +enabled_ui_label = "Soft inpainting" +enabled_gen_param_label = "Soft inpainting enabled" +enabled_el_id = "soft_inpainting_enabled" + +ui_labels = SoftInpaintingSettings( + "Schedule bias", + "Preservation strength", + "Transition contrast boost", + "Mask influence", + "Difference threshold", + "Difference contrast") + +ui_info = SoftInpaintingSettings( + "Shifts when preservation of original content occurs during denoising.", + "How strongly partially masked content should be preserved.", + "Amplifies the contrast that may be lost in partially masked regions.", + "How strongly the original mask should bias the difference threshold.", + "How much an image region can change before the original pixels are not blended in anymore.", + "How sharp the transition should be between blended and not blended.") + +gen_param_labels = SoftInpaintingSettings( + "Soft inpainting schedule bias", + "Soft inpainting preservation strength", + "Soft inpainting transition contrast boost", + "Soft inpainting mask influence", + "Soft inpainting difference threshold", + "Soft inpainting difference contrast") + +el_ids = SoftInpaintingSettings( + "mask_blend_power", + "mask_blend_scale", + "inpaint_detail_preservation", + "composite_mask_influence", + "composite_difference_threshold", + "composite_difference_contrast") + + +# ------------------- Script ------------------- + + +class Script(scripts.Script): + def __init__(self): + self.section = "inpaint" + self.masks_for_overlay = None + self.overlay_images = None + + def title(self): + return "Soft Inpainting" + + def show(self, is_img2img): + return scripts.AlwaysVisible if is_img2img else False + + def ui(self, is_img2img): + if not is_img2img: + return + + with InputAccordion(False, label=enabled_ui_label, elem_id=enabled_el_id) as soft_inpainting_enabled: + with gr.Group(): + gr.Markdown( + """ + Soft inpainting allows you to **seamlessly blend original content with inpainted content** according to the mask opacity. + **High _Mask blur_** values are recommended! + """) + + power = \ + gr.Slider(label=ui_labels.mask_blend_power, + info=ui_info.mask_blend_power, + minimum=0, + maximum=8, + step=0.1, + value=default.mask_blend_power, + elem_id=el_ids.mask_blend_power) + scale = \ + gr.Slider(label=ui_labels.mask_blend_scale, + info=ui_info.mask_blend_scale, + minimum=0, + maximum=8, + step=0.05, + value=default.mask_blend_scale, + elem_id=el_ids.mask_blend_scale) + detail = \ + gr.Slider(label=ui_labels.inpaint_detail_preservation, + info=ui_info.inpaint_detail_preservation, + minimum=1, + maximum=32, + step=0.5, + value=default.inpaint_detail_preservation, + elem_id=el_ids.inpaint_detail_preservation) + + gr.Markdown( + """ + ### Pixel Composite Settings + """) + + mask_inf = \ + gr.Slider(label=ui_labels.composite_mask_influence, + info=ui_info.composite_mask_influence, + minimum=0, + maximum=1, + step=0.05, + value=default.composite_mask_influence, + elem_id=el_ids.composite_mask_influence) + + dif_thresh = \ + gr.Slider(label=ui_labels.composite_difference_threshold, + info=ui_info.composite_difference_threshold, + minimum=0, + maximum=8, + step=0.25, + value=default.composite_difference_threshold, + elem_id=el_ids.composite_difference_threshold) + + dif_contr = \ + gr.Slider(label=ui_labels.composite_difference_contrast, + info=ui_info.composite_difference_contrast, + minimum=0, + maximum=8, + step=0.25, + value=default.composite_difference_contrast, + elem_id=el_ids.composite_difference_contrast) + + with gr.Accordion("Help", open=False): + gr.Markdown( + f""" + ### {ui_labels.mask_blend_power} + + The blending strength of original content is scaled proportionally with the decreasing noise level values at each step (sigmas). + This ensures that the influence of the denoiser and original content preservation is roughly balanced at each step. + This balance can be shifted using this parameter, controlling whether earlier or later steps have stronger preservation. + + - **Below 1**: Stronger preservation near the end (with low sigma) + - **1**: Balanced (proportional to sigma) + - **Above 1**: Stronger preservation in the beginning (with high sigma) + """) + gr.Markdown( + f""" + ### {ui_labels.mask_blend_scale} + + Skews whether partially masked image regions should be more likely to preserve the original content or favor inpainted content. + This may need to be adjusted depending on the {ui_labels.mask_blend_power}, CFG Scale, prompt and Denoising strength. + + - **Low values**: Favors generated content. + - **High values**: Favors original content. + """) + gr.Markdown( + f""" + ### {ui_labels.inpaint_detail_preservation} + + This parameter controls how the original latent vectors and denoised latent vectors are interpolated. + With higher values, the magnitude of the resulting blended vector will be closer to the maximum of the two interpolated vectors. + This can prevent the loss of contrast that occurs with linear interpolation. + + - **Low values**: Softer blending, details may fade. + - **High values**: Stronger contrast, may over-saturate colors. + """) + + gr.Markdown( + """ + ## Pixel Composite Settings + + Masks are generated based on how much a part of the image changed after denoising. + These masks are used to blend the original and final images together. + If the difference is low, the original pixels are used instead of the pixels returned by the inpainting process. + """) + + gr.Markdown( + f""" + ### {ui_labels.composite_mask_influence} + + This parameter controls how much the mask should bias this sensitivity to difference. + + - **0**: Ignore the mask, only consider differences in image content. + - **1**: Follow the mask closely despite image content changes. + """) + + gr.Markdown( + f""" + ### {ui_labels.composite_difference_threshold} + + This value represents the difference at which the original pixels will have less than 50% opacity. + + - **Low values**: Two images patches must be almost the same in order to retain original pixels. + - **High values**: Two images patches can be very different and still retain original pixels. + """) + + gr.Markdown( + f""" + ### {ui_labels.composite_difference_contrast} + + This value represents the contrast between the opacity of the original and inpainted content. + + - **Low values**: The blend will be more gradual and have longer transitions, but may cause ghosting. + - **High values**: Ghosting will be less common, but transitions may be very sudden. + """) + + self.infotext_fields = [(soft_inpainting_enabled, enabled_gen_param_label), + (power, gen_param_labels.mask_blend_power), + (scale, gen_param_labels.mask_blend_scale), + (detail, gen_param_labels.inpaint_detail_preservation), + (mask_inf, gen_param_labels.composite_mask_influence), + (dif_thresh, gen_param_labels.composite_difference_threshold), + (dif_contr, gen_param_labels.composite_difference_contrast)] + + self.paste_field_names = [] + for _, field_name in self.infotext_fields: + self.paste_field_names.append(field_name) + + return [soft_inpainting_enabled, + power, + scale, + detail, + mask_inf, + dif_thresh, + dif_contr] + + def process(self, p, enabled, power, scale, detail_preservation, mask_inf, dif_thresh, dif_contr): + if not enabled: + return + + if not processing_uses_inpainting(p): + return + + # Shut off the rounding it normally does. + p.mask_round = False + + settings = SoftInpaintingSettings(power, scale, detail_preservation, mask_inf, dif_thresh, dif_contr) + + # p.extra_generation_params["Mask rounding"] = False + settings.add_generation_params(p.extra_generation_params) + + def on_mask_blend(self, p, mba: scripts.MaskBlendArgs, enabled, power, scale, detail_preservation, mask_inf, + dif_thresh, dif_contr): + if not enabled: + return + + if not processing_uses_inpainting(p): + return + + if mba.is_final_blend: + mba.blended_latent = mba.current_latent + return + + settings = SoftInpaintingSettings(power, scale, detail_preservation, mask_inf, dif_thresh, dif_contr) + + # todo: Why is sigma 2D? Both values are the same. + mba.blended_latent = latent_blend(settings, + mba.init_latent, + mba.current_latent, + get_modified_nmask(settings, mba.nmask, mba.sigma[0])) + + def post_sample(self, p, ps: scripts.PostSampleArgs, enabled, power, scale, detail_preservation, mask_inf, + dif_thresh, dif_contr): + if not enabled: + return + + if not processing_uses_inpainting(p): + return + + nmask = getattr(p, "nmask", None) + if nmask is None: + return + + from modules import images + from modules.shared import opts + + settings = SoftInpaintingSettings(power, scale, detail_preservation, mask_inf, dif_thresh, dif_contr) + + # since the original code puts holes in the existing overlay images, + # we have to rebuild them. + self.overlay_images = [] + for img in p.init_images: + + image = images.flatten(img, opts.img2img_background_color) + + if p.paste_to is None and p.resize_mode != 3: + image = images.resize_image(p.resize_mode, image, p.width, p.height) + + self.overlay_images.append(image.convert('RGBA')) + + if len(p.init_images) == 1: + self.overlay_images = self.overlay_images * p.batch_size + + if getattr(ps.samples, 'already_decoded', False): + self.masks_for_overlay = apply_masks(settings=settings, + nmask=nmask, + overlay_images=self.overlay_images, + width=p.width, + height=p.height, + paste_to=p.paste_to) + else: + self.masks_for_overlay = apply_adaptive_masks(settings=settings, + nmask=nmask, + latent_orig=p.init_latent, + latent_processed=ps.samples, + overlay_images=self.overlay_images, + width=p.width, + height=p.height, + paste_to=p.paste_to) + + def postprocess_maskoverlay(self, p, ppmo: scripts.PostProcessMaskOverlayArgs, enabled, power, scale, + detail_preservation, mask_inf, dif_thresh, dif_contr): + if not enabled: + return + + if not processing_uses_inpainting(p): + return + + if self.masks_for_overlay is None: + return + + if self.overlay_images is None: + return + + ppmo.mask_for_overlay = self.masks_for_overlay[ppmo.index] + ppmo.overlay_image = self.overlay_images[ppmo.index] diff --git a/stable-diffusion-webui/html/card-no-preview.png b/stable-diffusion-webui/html/card-no-preview.png new file mode 100755 index 0000000..e2beb26 Binary files /dev/null and b/stable-diffusion-webui/html/card-no-preview.png differ diff --git a/stable-diffusion-webui/html/extra-networks-card.html b/stable-diffusion-webui/html/extra-networks-card.html new file mode 100755 index 0000000..f1d959a --- /dev/null +++ b/stable-diffusion-webui/html/extra-networks-card.html @@ -0,0 +1,9 @@ +
+ {background_image} +
{copy_path_button}{metadata_button}{edit_button}
+
+
{search_terms}
+ {name} + {description} +
+
diff --git a/stable-diffusion-webui/html/extra-networks-copy-path-button.html b/stable-diffusion-webui/html/extra-networks-copy-path-button.html new file mode 100755 index 0000000..50304b4 --- /dev/null +++ b/stable-diffusion-webui/html/extra-networks-copy-path-button.html @@ -0,0 +1,5 @@ +
+
\ No newline at end of file diff --git a/stable-diffusion-webui/html/extra-networks-edit-item-button.html b/stable-diffusion-webui/html/extra-networks-edit-item-button.html new file mode 100755 index 0000000..fd72860 --- /dev/null +++ b/stable-diffusion-webui/html/extra-networks-edit-item-button.html @@ -0,0 +1,4 @@ +
+
\ No newline at end of file diff --git a/stable-diffusion-webui/html/extra-networks-metadata-button.html b/stable-diffusion-webui/html/extra-networks-metadata-button.html new file mode 100755 index 0000000..4ef013b --- /dev/null +++ b/stable-diffusion-webui/html/extra-networks-metadata-button.html @@ -0,0 +1,4 @@ +
+
\ No newline at end of file diff --git a/stable-diffusion-webui/html/extra-networks-no-cards.html b/stable-diffusion-webui/html/extra-networks-no-cards.html new file mode 100755 index 0000000..389358d --- /dev/null +++ b/stable-diffusion-webui/html/extra-networks-no-cards.html @@ -0,0 +1,8 @@ +
+

Nothing here. Add some content to the following directories:

+ +
    +{dirs} +
+
+ diff --git a/stable-diffusion-webui/html/extra-networks-pane-dirs.html b/stable-diffusion-webui/html/extra-networks-pane-dirs.html new file mode 100755 index 0000000..d7c9661 --- /dev/null +++ b/stable-diffusion-webui/html/extra-networks-pane-dirs.html @@ -0,0 +1,8 @@ +
+
+ {dirs_html} +
+
+ {items_html} +
+
diff --git a/stable-diffusion-webui/html/extra-networks-pane-tree.html b/stable-diffusion-webui/html/extra-networks-pane-tree.html new file mode 100755 index 0000000..e4d92a3 --- /dev/null +++ b/stable-diffusion-webui/html/extra-networks-pane-tree.html @@ -0,0 +1,8 @@ +
+
+ {tree_html} +
+
+ {items_html} +
+
\ No newline at end of file diff --git a/stable-diffusion-webui/html/extra-networks-pane.html b/stable-diffusion-webui/html/extra-networks-pane.html new file mode 100755 index 0000000..9a67bae --- /dev/null +++ b/stable-diffusion-webui/html/extra-networks-pane.html @@ -0,0 +1,81 @@ +
+ + {pane_content} +
diff --git a/stable-diffusion-webui/html/extra-networks-tree-button.html b/stable-diffusion-webui/html/extra-networks-tree-button.html new file mode 100755 index 0000000..9dc2e2a --- /dev/null +++ b/stable-diffusion-webui/html/extra-networks-tree-button.html @@ -0,0 +1,23 @@ + +
+ + {action_list_item_action_leading} + + + {action_list_item_visual_leading} + + + {action_list_item_label} + + + {action_list_item_visual_trailing} + + + {action_list_item_action_trailing} + +
\ No newline at end of file diff --git a/stable-diffusion-webui/html/footer.html b/stable-diffusion-webui/html/footer.html new file mode 100755 index 0000000..8739a0f --- /dev/null +++ b/stable-diffusion-webui/html/footer.html @@ -0,0 +1,15 @@ +
+ API +  •  + Github +  •  + Gradio +  •  + Startup profile +  •  + Reload UI +
+
+
+{versions} +
diff --git a/stable-diffusion-webui/html/licenses.html b/stable-diffusion-webui/html/licenses.html new file mode 100755 index 0000000..e14bf3c --- /dev/null +++ b/stable-diffusion-webui/html/licenses.html @@ -0,0 +1,382 @@ + + +

InvokeAI

+Some code for compatibility with OSX is taken from lstein's repository. +
+MIT License
+
+Copyright (c) 2022 InvokeAI Team
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+ +

LDSR

+Code added by contirubtors, most likely copied from this repository. +
+MIT License
+
+Copyright (c) 2022 Machine Vision and Learning Group, LMU Munich
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+ +

CLIP Interrogator

+Some small amounts of code borrowed and reworked. +
+MIT License
+
+Copyright (c) 2022 pharmapsychotic
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+ +

Memory Efficient Attention

+The sub-quadratic cross attention optimization uses modified code from the Memory Efficient Attention package that Alex Birch optimized for 3D tensors. This license is updated to reflect that. +
+MIT License
+
+Copyright (c) 2023 Alex Birch
+Copyright (c) 2023 Amin Rezaei
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+ +

Scaled Dot Product Attention

+Some small amounts of code borrowed and reworked. +
+   Copyright 2023 The HuggingFace Team. All rights reserved.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+
+ +

Curated transformers

+The MPS workaround for nn.Linear on macOS 13.2.X is based on the MPS workaround for nn.Linear created by danieldk for Curated transformers +
+The MIT License (MIT)
+
+Copyright (C) 2021 ExplosionAI GmbH
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+
+ +

TAESD

+Tiny AutoEncoder for Stable Diffusion option for live previews +
+MIT License
+
+Copyright (c) 2023 Ollin Boer Bohan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
diff --git a/stable-diffusion-webui/javascript/aspectRatioOverlay.js b/stable-diffusion-webui/javascript/aspectRatioOverlay.js new file mode 100755 index 0000000..c8751fe --- /dev/null +++ b/stable-diffusion-webui/javascript/aspectRatioOverlay.js @@ -0,0 +1,113 @@ + +let currentWidth = null; +let currentHeight = null; +let arFrameTimeout = setTimeout(function() {}, 0); + +function dimensionChange(e, is_width, is_height) { + + if (is_width) { + currentWidth = e.target.value * 1.0; + } + if (is_height) { + currentHeight = e.target.value * 1.0; + } + + var inImg2img = gradioApp().querySelector("#tab_img2img").style.display == "block"; + + if (!inImg2img) { + return; + } + + var targetElement = null; + + var tabIndex = get_tab_index('mode_img2img'); + if (tabIndex == 0) { // img2img + targetElement = gradioApp().querySelector('#img2img_image div[data-testid=image] img'); + } else if (tabIndex == 1) { //Sketch + targetElement = gradioApp().querySelector('#img2img_sketch div[data-testid=image] img'); + } else if (tabIndex == 2) { // Inpaint + targetElement = gradioApp().querySelector('#img2maskimg div[data-testid=image] img'); + } else if (tabIndex == 3) { // Inpaint sketch + targetElement = gradioApp().querySelector('#inpaint_sketch div[data-testid=image] img'); + } + + + if (targetElement) { + + var arPreviewRect = gradioApp().querySelector('#imageARPreview'); + if (!arPreviewRect) { + arPreviewRect = document.createElement('div'); + arPreviewRect.id = "imageARPreview"; + gradioApp().appendChild(arPreviewRect); + } + + + + var viewportOffset = targetElement.getBoundingClientRect(); + + var viewportscale = Math.min(targetElement.clientWidth / targetElement.naturalWidth, targetElement.clientHeight / targetElement.naturalHeight); + + var scaledx = targetElement.naturalWidth * viewportscale; + var scaledy = targetElement.naturalHeight * viewportscale; + + var clientRectTop = (viewportOffset.top + window.scrollY); + var clientRectLeft = (viewportOffset.left + window.scrollX); + var clientRectCentreY = clientRectTop + (targetElement.clientHeight / 2); + var clientRectCentreX = clientRectLeft + (targetElement.clientWidth / 2); + + var arscale = Math.min(scaledx / currentWidth, scaledy / currentHeight); + var arscaledx = currentWidth * arscale; + var arscaledy = currentHeight * arscale; + + var arRectTop = clientRectCentreY - (arscaledy / 2); + var arRectLeft = clientRectCentreX - (arscaledx / 2); + var arRectWidth = arscaledx; + var arRectHeight = arscaledy; + + arPreviewRect.style.top = arRectTop + 'px'; + arPreviewRect.style.left = arRectLeft + 'px'; + arPreviewRect.style.width = arRectWidth + 'px'; + arPreviewRect.style.height = arRectHeight + 'px'; + + clearTimeout(arFrameTimeout); + arFrameTimeout = setTimeout(function() { + arPreviewRect.style.display = 'none'; + }, 2000); + + arPreviewRect.style.display = 'block'; + + } + +} + + +onAfterUiUpdate(function() { + var arPreviewRect = gradioApp().querySelector('#imageARPreview'); + if (arPreviewRect) { + arPreviewRect.style.display = 'none'; + } + var tabImg2img = gradioApp().querySelector("#tab_img2img"); + if (tabImg2img) { + var inImg2img = tabImg2img.style.display == "block"; + if (inImg2img) { + let inputs = gradioApp().querySelectorAll('input'); + inputs.forEach(function(e) { + var is_width = e.parentElement.id == "img2img_width"; + var is_height = e.parentElement.id == "img2img_height"; + + if ((is_width || is_height) && !e.classList.contains('scrollwatch')) { + e.addEventListener('input', function(e) { + dimensionChange(e, is_width, is_height); + }); + e.classList.add('scrollwatch'); + } + if (is_width) { + currentWidth = e.value * 1.0; + } + if (is_height) { + currentHeight = e.value * 1.0; + } + }); + } + } +}); diff --git a/stable-diffusion-webui/javascript/contextMenus.js b/stable-diffusion-webui/javascript/contextMenus.js new file mode 100755 index 0000000..e01fd67 --- /dev/null +++ b/stable-diffusion-webui/javascript/contextMenus.js @@ -0,0 +1,163 @@ + +var contextMenuInit = function() { + let eventListenerApplied = false; + let menuSpecs = new Map(); + + const uid = function() { + return Date.now().toString(36) + Math.random().toString(36).substring(2); + }; + + function showContextMenu(event, element, menuEntries) { + let oldMenu = gradioApp().querySelector('#context-menu'); + if (oldMenu) { + oldMenu.remove(); + } + + let baseStyle = window.getComputedStyle(uiCurrentTab); + + const contextMenu = document.createElement('nav'); + contextMenu.id = "context-menu"; + contextMenu.style.background = baseStyle.background; + contextMenu.style.color = baseStyle.color; + contextMenu.style.fontFamily = baseStyle.fontFamily; + contextMenu.style.top = event.pageY + 'px'; + contextMenu.style.left = event.pageX + 'px'; + + const contextMenuList = document.createElement('ul'); + contextMenuList.className = 'context-menu-items'; + contextMenu.append(contextMenuList); + + menuEntries.forEach(function(entry) { + let contextMenuEntry = document.createElement('a'); + contextMenuEntry.innerHTML = entry['name']; + contextMenuEntry.addEventListener("click", function() { + entry['func'](); + }); + contextMenuList.append(contextMenuEntry); + + }); + + gradioApp().appendChild(contextMenu); + } + + function appendContextMenuOption(targetElementSelector, entryName, entryFunction) { + + var currentItems = menuSpecs.get(targetElementSelector); + + if (!currentItems) { + currentItems = []; + menuSpecs.set(targetElementSelector, currentItems); + } + let newItem = { + id: targetElementSelector + '_' + uid(), + name: entryName, + func: entryFunction, + isNew: true + }; + + currentItems.push(newItem); + return newItem['id']; + } + + function removeContextMenuOption(uid) { + menuSpecs.forEach(function(v) { + let index = -1; + v.forEach(function(e, ei) { + if (e['id'] == uid) { + index = ei; + } + }); + if (index >= 0) { + v.splice(index, 1); + } + }); + } + + function addContextMenuEventListener() { + if (eventListenerApplied) { + return; + } + gradioApp().addEventListener("click", function(e) { + if (!e.isTrusted) { + return; + } + + let oldMenu = gradioApp().querySelector('#context-menu'); + if (oldMenu) { + oldMenu.remove(); + } + }); + ['contextmenu', 'touchstart'].forEach((eventType) => { + gradioApp().addEventListener(eventType, function(e) { + let ev = e; + if (eventType.startsWith('touch')) { + if (e.touches.length !== 2) return; + ev = e.touches[0]; + } + let oldMenu = gradioApp().querySelector('#context-menu'); + if (oldMenu) { + oldMenu.remove(); + } + menuSpecs.forEach(function(v, k) { + if (e.composedPath()[0].matches(k)) { + showContextMenu(ev, e.composedPath()[0], v); + e.preventDefault(); + } + }); + }); + }); + eventListenerApplied = true; + + } + + return [appendContextMenuOption, removeContextMenuOption, addContextMenuEventListener]; +}; + +var initResponse = contextMenuInit(); +var appendContextMenuOption = initResponse[0]; +var removeContextMenuOption = initResponse[1]; +var addContextMenuEventListener = initResponse[2]; + +(function() { + //Start example Context Menu Items + let generateOnRepeat = function(genbuttonid, interruptbuttonid) { + let genbutton = gradioApp().querySelector(genbuttonid); + let interruptbutton = gradioApp().querySelector(interruptbuttonid); + if (!interruptbutton.offsetParent) { + genbutton.click(); + } + clearInterval(window.generateOnRepeatInterval); + window.generateOnRepeatInterval = setInterval(function() { + if (!interruptbutton.offsetParent) { + genbutton.click(); + } + }, + 500); + }; + + let generateOnRepeat_txt2img = function() { + generateOnRepeat('#txt2img_generate', '#txt2img_interrupt'); + }; + + let generateOnRepeat_img2img = function() { + generateOnRepeat('#img2img_generate', '#img2img_interrupt'); + }; + + appendContextMenuOption('#txt2img_generate', 'Generate forever', generateOnRepeat_txt2img); + appendContextMenuOption('#txt2img_interrupt', 'Generate forever', generateOnRepeat_txt2img); + appendContextMenuOption('#img2img_generate', 'Generate forever', generateOnRepeat_img2img); + appendContextMenuOption('#img2img_interrupt', 'Generate forever', generateOnRepeat_img2img); + + let cancelGenerateForever = function() { + clearInterval(window.generateOnRepeatInterval); + }; + + appendContextMenuOption('#txt2img_interrupt', 'Cancel generate forever', cancelGenerateForever); + appendContextMenuOption('#txt2img_generate', 'Cancel generate forever', cancelGenerateForever); + appendContextMenuOption('#img2img_interrupt', 'Cancel generate forever', cancelGenerateForever); + appendContextMenuOption('#img2img_generate', 'Cancel generate forever', cancelGenerateForever); + +})(); +//End example Context Menu Items + +onAfterUiUpdate(addContextMenuEventListener); diff --git a/stable-diffusion-webui/javascript/dragdrop.js b/stable-diffusion-webui/javascript/dragdrop.js new file mode 100755 index 0000000..882562d --- /dev/null +++ b/stable-diffusion-webui/javascript/dragdrop.js @@ -0,0 +1,156 @@ +// allows drag-dropping files into gradio image elements, and also pasting images from clipboard + +function isValidImageList(files) { + return files && files?.length === 1 && ['image/png', 'image/gif', 'image/jpeg'].includes(files[0].type); +} + +function dropReplaceImage(imgWrap, files) { + if (!isValidImageList(files)) { + return; + } + + const tmpFile = files[0]; + + imgWrap.querySelector('.modify-upload button + button, .touch-none + div button + button')?.click(); + const callback = () => { + const fileInput = imgWrap.querySelector('input[type="file"]'); + if (fileInput) { + if (files.length === 0) { + files = new DataTransfer(); + files.items.add(tmpFile); + fileInput.files = files.files; + } else { + fileInput.files = files; + } + fileInput.dispatchEvent(new Event('change')); + } + }; + + if (imgWrap.closest('#pnginfo_image')) { + // special treatment for PNG Info tab, wait for fetch request to finish + const oldFetch = window.fetch; + window.fetch = async(input, options) => { + const response = await oldFetch(input, options); + if ('api/predict/' === input) { + const content = await response.text(); + window.fetch = oldFetch; + window.requestAnimationFrame(() => callback()); + return new Response(content, { + status: response.status, + statusText: response.statusText, + headers: response.headers + }); + } + return response; + }; + } else { + window.requestAnimationFrame(() => callback()); + } +} + +function eventHasFiles(e) { + if (!e.dataTransfer || !e.dataTransfer.files) return false; + if (e.dataTransfer.files.length > 0) return true; + if (e.dataTransfer.items.length > 0 && e.dataTransfer.items[0].kind == "file") return true; + + return false; +} + +function isURL(url) { + try { + const _ = new URL(url); + return true; + } catch { + return false; + } +} + +function dragDropTargetIsPrompt(target) { + if (target?.placeholder && target?.placeholder.indexOf("Prompt") >= 0) return true; + if (target?.parentNode?.parentNode?.className?.indexOf("prompt") > 0) return true; + return false; +} + +window.document.addEventListener('dragover', e => { + const target = e.composedPath()[0]; + if (!eventHasFiles(e)) return; + + var targetImage = target.closest('[data-testid="image"]'); + if (!dragDropTargetIsPrompt(target) && !targetImage) return; + + e.stopPropagation(); + e.preventDefault(); + e.dataTransfer.dropEffect = 'copy'; +}); + +window.document.addEventListener('drop', async e => { + const target = e.composedPath()[0]; + const url = e.dataTransfer.getData('text/uri-list') || e.dataTransfer.getData('text/plain'); + if (!eventHasFiles(e) && !isURL(url)) return; + + if (dragDropTargetIsPrompt(target)) { + e.stopPropagation(); + e.preventDefault(); + + const isImg2img = get_tab_index('tabs') == 1; + let prompt_image_target = isImg2img ? "img2img_prompt_image" : "txt2img_prompt_image"; + + const imgParent = gradioApp().getElementById(prompt_image_target); + const files = e.dataTransfer.files; + const fileInput = imgParent.querySelector('input[type="file"]'); + if (eventHasFiles(e) && fileInput) { + fileInput.files = files; + fileInput.dispatchEvent(new Event('change')); + } else if (url) { + try { + const request = await fetch(url); + if (!request.ok) { + console.error('Error fetching URL:', url, request.status); + return; + } + const data = new DataTransfer(); + data.items.add(new File([await request.blob()], 'image.png')); + fileInput.files = data.files; + fileInput.dispatchEvent(new Event('change')); + } catch (error) { + console.error('Error fetching URL:', url, error); + return; + } + } + } + + var targetImage = target.closest('[data-testid="image"]'); + if (targetImage) { + e.stopPropagation(); + e.preventDefault(); + const files = e.dataTransfer.files; + dropReplaceImage(targetImage, files); + return; + } +}); + +window.addEventListener('paste', e => { + const files = e.clipboardData.files; + if (!isValidImageList(files)) { + return; + } + + const visibleImageFields = [...gradioApp().querySelectorAll('[data-testid="image"]')] + .filter(el => uiElementIsVisible(el)) + .sort((a, b) => uiElementInSight(b) - uiElementInSight(a)); + + + if (!visibleImageFields.length) { + return; + } + + const firstFreeImageField = visibleImageFields + .filter(el => !el.querySelector('img'))?.[0]; + + dropReplaceImage( + firstFreeImageField ? + firstFreeImageField : + visibleImageFields[visibleImageFields.length - 1] + , files + ); +}); diff --git a/stable-diffusion-webui/javascript/edit-attention.js b/stable-diffusion-webui/javascript/edit-attention.js new file mode 100755 index 0000000..b07ba97 --- /dev/null +++ b/stable-diffusion-webui/javascript/edit-attention.js @@ -0,0 +1,156 @@ +function keyupEditAttention(event) { + let target = event.originalTarget || event.composedPath()[0]; + if (!target.matches("*:is([id*='_toprow'] [id*='_prompt'], .prompt) textarea")) return; + if (!(event.metaKey || event.ctrlKey)) return; + + let isPlus = event.key == "ArrowUp"; + let isMinus = event.key == "ArrowDown"; + if (!isPlus && !isMinus) return; + + let selectionStart = target.selectionStart; + let selectionEnd = target.selectionEnd; + let text = target.value; + + function selectCurrentParenthesisBlock(OPEN, CLOSE) { + if (selectionStart !== selectionEnd) return false; + + // Find opening parenthesis around current cursor + const before = text.substring(0, selectionStart); + let beforeParen = before.lastIndexOf(OPEN); + if (beforeParen == -1) return false; + + let beforeClosingParen = before.lastIndexOf(CLOSE); + if (beforeClosingParen != -1 && beforeClosingParen > beforeParen) return false; + + // Find closing parenthesis around current cursor + const after = text.substring(selectionStart); + let afterParen = after.indexOf(CLOSE); + if (afterParen == -1) return false; + + let afterOpeningParen = after.indexOf(OPEN); + if (afterOpeningParen != -1 && afterOpeningParen < afterParen) return false; + + // Set the selection to the text between the parenthesis + const parenContent = text.substring(beforeParen + 1, selectionStart + afterParen); + if (/.*:-?[\d.]+/s.test(parenContent)) { + const lastColon = parenContent.lastIndexOf(":"); + selectionStart = beforeParen + 1; + selectionEnd = selectionStart + lastColon; + } else { + selectionStart = beforeParen + 1; + selectionEnd = selectionStart + parenContent.length; + } + + target.setSelectionRange(selectionStart, selectionEnd); + return true; + } + + function selectCurrentWord() { + if (selectionStart !== selectionEnd) return false; + const whitespace_delimiters = {"Tab": "\t", "Carriage Return": "\r", "Line Feed": "\n"}; + let delimiters = opts.keyedit_delimiters; + + for (let i of opts.keyedit_delimiters_whitespace) { + delimiters += whitespace_delimiters[i]; + } + + // seek backward to find beginning + while (!delimiters.includes(text[selectionStart - 1]) && selectionStart > 0) { + selectionStart--; + } + + // seek forward to find end + while (!delimiters.includes(text[selectionEnd]) && selectionEnd < text.length) { + selectionEnd++; + } + + // deselect surrounding whitespace + while (text[selectionStart] == " " && selectionStart < selectionEnd) { + selectionStart++; + } + while (text[selectionEnd - 1] == " " && selectionEnd > selectionStart) { + selectionEnd--; + } + + target.setSelectionRange(selectionStart, selectionEnd); + return true; + } + + // If the user hasn't selected anything, let's select their current parenthesis block or word + if (!selectCurrentParenthesisBlock('<', '>') && !selectCurrentParenthesisBlock('(', ')') && !selectCurrentParenthesisBlock('[', ']')) { + selectCurrentWord(); + } + + event.preventDefault(); + + var closeCharacter = ')'; + var delta = opts.keyedit_precision_attention; + var start = selectionStart > 0 ? text[selectionStart - 1] : ""; + var end = text[selectionEnd]; + + if (start == '<') { + closeCharacter = '>'; + delta = opts.keyedit_precision_extra; + } else if (start == '(' && end == ')' || start == '[' && end == ']') { // convert old-style (((emphasis))) + let numParen = 0; + + while (text[selectionStart - numParen - 1] == start && text[selectionEnd + numParen] == end) { + numParen++; + } + + if (start == "[") { + weight = (1 / 1.1) ** numParen; + } else { + weight = 1.1 ** numParen; + } + + weight = Math.round(weight / opts.keyedit_precision_attention) * opts.keyedit_precision_attention; + + text = text.slice(0, selectionStart - numParen) + "(" + text.slice(selectionStart, selectionEnd) + ":" + weight + ")" + text.slice(selectionEnd + numParen); + selectionStart -= numParen - 1; + selectionEnd -= numParen - 1; + } else if (start != '(') { + // do not include spaces at the end + while (selectionEnd > selectionStart && text[selectionEnd - 1] == ' ') { + selectionEnd--; + } + + if (selectionStart == selectionEnd) { + return; + } + + text = text.slice(0, selectionStart) + "(" + text.slice(selectionStart, selectionEnd) + ":1.0)" + text.slice(selectionEnd); + + selectionStart++; + selectionEnd++; + } + + if (text[selectionEnd] != ':') return; + var weightLength = text.slice(selectionEnd + 1).indexOf(closeCharacter) + 1; + var weight = parseFloat(text.slice(selectionEnd + 1, selectionEnd + weightLength)); + if (isNaN(weight)) return; + + weight += isPlus ? delta : -delta; + weight = parseFloat(weight.toPrecision(12)); + if (Number.isInteger(weight)) weight += ".0"; + + if (closeCharacter == ')' && weight == 1) { + var endParenPos = text.substring(selectionEnd).indexOf(')'); + text = text.slice(0, selectionStart - 1) + text.slice(selectionStart, selectionEnd) + text.slice(selectionEnd + endParenPos + 1); + selectionStart--; + selectionEnd--; + } else { + text = text.slice(0, selectionEnd + 1) + weight + text.slice(selectionEnd + weightLength); + } + + target.focus(); + target.value = text; + target.selectionStart = selectionStart; + target.selectionEnd = selectionEnd; + + updateInput(target); +} + +addEventListener('keydown', (event) => { + keyupEditAttention(event); +}); diff --git a/stable-diffusion-webui/javascript/edit-order.js b/stable-diffusion-webui/javascript/edit-order.js new file mode 100755 index 0000000..ed4ef9a --- /dev/null +++ b/stable-diffusion-webui/javascript/edit-order.js @@ -0,0 +1,41 @@ +/* alt+left/right moves text in prompt */ + +function keyupEditOrder(event) { + if (!opts.keyedit_move) return; + + let target = event.originalTarget || event.composedPath()[0]; + if (!target.matches("*:is([id*='_toprow'] [id*='_prompt'], .prompt) textarea")) return; + if (!event.altKey) return; + + let isLeft = event.key == "ArrowLeft"; + let isRight = event.key == "ArrowRight"; + if (!isLeft && !isRight) return; + event.preventDefault(); + + let selectionStart = target.selectionStart; + let selectionEnd = target.selectionEnd; + let text = target.value; + let items = text.split(","); + let indexStart = (text.slice(0, selectionStart).match(/,/g) || []).length; + let indexEnd = (text.slice(0, selectionEnd).match(/,/g) || []).length; + let range = indexEnd - indexStart + 1; + + if (isLeft && indexStart > 0) { + items.splice(indexStart - 1, 0, ...items.splice(indexStart, range)); + target.value = items.join(); + target.selectionStart = items.slice(0, indexStart - 1).join().length + (indexStart == 1 ? 0 : 1); + target.selectionEnd = items.slice(0, indexEnd).join().length; + } else if (isRight && indexEnd < items.length - 1) { + items.splice(indexStart + 1, 0, ...items.splice(indexStart, range)); + target.value = items.join(); + target.selectionStart = items.slice(0, indexStart + 1).join().length + 1; + target.selectionEnd = items.slice(0, indexEnd + 2).join().length; + } + + event.preventDefault(); + updateInput(target); +} + +addEventListener('keydown', (event) => { + keyupEditOrder(event); +}); diff --git a/stable-diffusion-webui/javascript/extensions.js b/stable-diffusion-webui/javascript/extensions.js new file mode 100755 index 0000000..cc8ee22 --- /dev/null +++ b/stable-diffusion-webui/javascript/extensions.js @@ -0,0 +1,95 @@ + +function extensions_apply(_disabled_list, _update_list, disable_all) { + var disable = []; + var update = []; + const extensions_input = gradioApp().querySelectorAll('#extensions input[type="checkbox"]'); + if (extensions_input.length == 0) { + throw Error("Extensions page not yet loaded."); + } + extensions_input.forEach(function(x) { + if (x.name.startsWith("enable_") && !x.checked) { + disable.push(x.name.substring(7)); + } + + if (x.name.startsWith("update_") && x.checked) { + update.push(x.name.substring(7)); + } + }); + + restart_reload(); + + return [JSON.stringify(disable), JSON.stringify(update), disable_all]; +} + +function extensions_check() { + var disable = []; + + gradioApp().querySelectorAll('#extensions input[type="checkbox"]').forEach(function(x) { + if (x.name.startsWith("enable_") && !x.checked) { + disable.push(x.name.substring(7)); + } + }); + + gradioApp().querySelectorAll('#extensions .extension_status').forEach(function(x) { + x.innerHTML = "Loading..."; + }); + + + var id = randomId(); + requestProgress(id, gradioApp().getElementById('extensions_installed_html'), null, function() { + + }); + + return [id, JSON.stringify(disable)]; +} + +function install_extension_from_index(button, url) { + button.disabled = "disabled"; + button.value = "Installing..."; + + var textarea = gradioApp().querySelector('#extension_to_install textarea'); + textarea.value = url; + updateInput(textarea); + + gradioApp().querySelector('#install_extension_button').click(); +} + +function config_state_confirm_restore(_, config_state_name, config_restore_type) { + if (config_state_name == "Current") { + return [false, config_state_name, config_restore_type]; + } + let restored = ""; + if (config_restore_type == "extensions") { + restored = "all saved extension versions"; + } else if (config_restore_type == "webui") { + restored = "the webui version"; + } else { + restored = "the webui version and all saved extension versions"; + } + let confirmed = confirm("Are you sure you want to restore from this state?\nThis will reset " + restored + "."); + if (confirmed) { + restart_reload(); + gradioApp().querySelectorAll('#extensions .extension_status').forEach(function(x) { + x.innerHTML = "Loading..."; + }); + } + return [confirmed, config_state_name, config_restore_type]; +} + +function toggle_all_extensions(event) { + gradioApp().querySelectorAll('#extensions .extension_toggle').forEach(function(checkbox_el) { + checkbox_el.checked = event.target.checked; + }); +} + +function toggle_extension() { + let all_extensions_toggled = true; + for (const checkbox_el of gradioApp().querySelectorAll('#extensions .extension_toggle')) { + if (!checkbox_el.checked) { + all_extensions_toggled = false; + break; + } + } + + gradioApp().querySelector('#extensions .all_extensions_toggle').checked = all_extensions_toggled; +} diff --git a/stable-diffusion-webui/javascript/extraNetworks.js b/stable-diffusion-webui/javascript/extraNetworks.js new file mode 100755 index 0000000..c5cced9 --- /dev/null +++ b/stable-diffusion-webui/javascript/extraNetworks.js @@ -0,0 +1,712 @@ +function toggleCss(key, css, enable) { + var style = document.getElementById(key); + if (enable && !style) { + style = document.createElement('style'); + style.id = key; + style.type = 'text/css'; + document.head.appendChild(style); + } + if (style && !enable) { + document.head.removeChild(style); + } + if (style) { + style.innerHTML == ''; + style.appendChild(document.createTextNode(css)); + } +} + +function setupExtraNetworksForTab(tabname) { + function registerPrompt(tabname, id) { + var textarea = gradioApp().querySelector("#" + id + " > label > textarea"); + + if (!activePromptTextarea[tabname]) { + activePromptTextarea[tabname] = textarea; + } + + textarea.addEventListener("focus", function() { + activePromptTextarea[tabname] = textarea; + }); + } + + var tabnav = gradioApp().querySelector('#' + tabname + '_extra_tabs > div.tab-nav'); + var controlsDiv = document.createElement('DIV'); + controlsDiv.classList.add('extra-networks-controls-div'); + tabnav.appendChild(controlsDiv); + tabnav.insertBefore(controlsDiv, null); + + var this_tab = gradioApp().querySelector('#' + tabname + '_extra_tabs'); + this_tab.querySelectorAll(":scope > [id^='" + tabname + "_']").forEach(function(elem) { + // tabname_full = {tabname}_{extra_networks_tabname} + var tabname_full = elem.id; + var search = gradioApp().querySelector("#" + tabname_full + "_extra_search"); + var sort_dir = gradioApp().querySelector("#" + tabname_full + "_extra_sort_dir"); + var refresh = gradioApp().querySelector("#" + tabname_full + "_extra_refresh"); + var currentSort = ''; + + // If any of the buttons above don't exist, we want to skip this iteration of the loop. + if (!search || !sort_dir || !refresh) { + return; // `return` is equivalent of `continue` but for forEach loops. + } + + var applyFilter = function(force) { + var searchTerm = search.value.toLowerCase(); + gradioApp().querySelectorAll('#' + tabname + '_extra_tabs div.card').forEach(function(elem) { + var searchOnly = elem.querySelector('.search_only'); + var text = Array.prototype.map.call(elem.querySelectorAll('.search_terms, .description'), function(t) { + return t.textContent.toLowerCase(); + }).join(" "); + + var visible = text.indexOf(searchTerm) != -1; + if (searchOnly && searchTerm.length < 4) { + visible = false; + } + if (visible) { + elem.classList.remove("hidden"); + } else { + elem.classList.add("hidden"); + } + }); + + applySort(force); + }; + + var applySort = function(force) { + var cards = gradioApp().querySelectorAll('#' + tabname_full + ' div.card'); + var parent = gradioApp().querySelector('#' + tabname_full + "_cards"); + var reverse = sort_dir.dataset.sortdir == "Descending"; + var activeSearchElem = gradioApp().querySelector('#' + tabname_full + "_controls .extra-network-control--sort.extra-network-control--enabled"); + var sortKey = activeSearchElem ? activeSearchElem.dataset.sortkey : "default"; + var sortKeyDataField = "sort" + sortKey.charAt(0).toUpperCase() + sortKey.slice(1); + var sortKeyStore = sortKey + "-" + sort_dir.dataset.sortdir + "-" + cards.length; + + if (sortKeyStore == currentSort && !force) { + return; + } + currentSort = sortKeyStore; + + var sortedCards = Array.from(cards); + sortedCards.sort(function(cardA, cardB) { + var a = cardA.dataset[sortKeyDataField]; + var b = cardB.dataset[sortKeyDataField]; + if (!isNaN(a) && !isNaN(b)) { + return parseInt(a) - parseInt(b); + } + + return (a < b ? -1 : (a > b ? 1 : 0)); + }); + + if (reverse) { + sortedCards.reverse(); + } + + parent.innerHTML = ''; + + var frag = document.createDocumentFragment(); + sortedCards.forEach(function(card) { + frag.appendChild(card); + }); + parent.appendChild(frag); + }; + + search.addEventListener("input", function() { + applyFilter(); + }); + applySort(); + applyFilter(); + extraNetworksApplySort[tabname_full] = applySort; + extraNetworksApplyFilter[tabname_full] = applyFilter; + + var controls = gradioApp().querySelector("#" + tabname_full + "_controls"); + controlsDiv.insertBefore(controls, null); + + if (elem.style.display != "none") { + extraNetworksShowControlsForPage(tabname, tabname_full); + } + }); + + registerPrompt(tabname, tabname + "_prompt"); + registerPrompt(tabname, tabname + "_neg_prompt"); +} + +function extraNetworksMovePromptToTab(tabname, id, showPrompt, showNegativePrompt) { + if (!gradioApp().querySelector('.toprow-compact-tools')) return; // only applicable for compact prompt layout + + var promptContainer = gradioApp().getElementById(tabname + '_prompt_container'); + var prompt = gradioApp().getElementById(tabname + '_prompt_row'); + var negPrompt = gradioApp().getElementById(tabname + '_neg_prompt_row'); + var elem = id ? gradioApp().getElementById(id) : null; + + if (showNegativePrompt && elem) { + elem.insertBefore(negPrompt, elem.firstChild); + } else { + promptContainer.insertBefore(negPrompt, promptContainer.firstChild); + } + + if (showPrompt && elem) { + elem.insertBefore(prompt, elem.firstChild); + } else { + promptContainer.insertBefore(prompt, promptContainer.firstChild); + } + + if (elem) { + elem.classList.toggle('extra-page-prompts-active', showNegativePrompt || showPrompt); + } +} + + +function extraNetworksShowControlsForPage(tabname, tabname_full) { + gradioApp().querySelectorAll('#' + tabname + '_extra_tabs .extra-networks-controls-div > div').forEach(function(elem) { + var targetId = tabname_full + "_controls"; + elem.style.display = elem.id == targetId ? "" : "none"; + }); +} + + +function extraNetworksUnrelatedTabSelected(tabname) { // called from python when user selects an unrelated tab (generate) + extraNetworksMovePromptToTab(tabname, '', false, false); + + extraNetworksShowControlsForPage(tabname, null); +} + +function extraNetworksTabSelected(tabname, id, showPrompt, showNegativePrompt, tabname_full) { // called from python when user selects an extra networks tab + extraNetworksMovePromptToTab(tabname, id, showPrompt, showNegativePrompt); + + extraNetworksShowControlsForPage(tabname, tabname_full); +} + +function applyExtraNetworkFilter(tabname_full) { + var doFilter = function() { + var applyFunction = extraNetworksApplyFilter[tabname_full]; + + if (applyFunction) { + applyFunction(true); + } + }; + setTimeout(doFilter, 1); +} + +function applyExtraNetworkSort(tabname_full) { + var doSort = function() { + extraNetworksApplySort[tabname_full](true); + }; + setTimeout(doSort, 1); +} + +var extraNetworksApplyFilter = {}; +var extraNetworksApplySort = {}; +var activePromptTextarea = {}; + +function setupExtraNetworks() { + setupExtraNetworksForTab('txt2img'); + setupExtraNetworksForTab('img2img'); +} + +var re_extranet = /<([^:^>]+:[^:]+):[\d.]+>(.*)/; +var re_extranet_g = /<([^:^>]+:[^:]+):[\d.]+>/g; + +var re_extranet_neg = /\(([^:^>]+:[\d.]+)\)/; +var re_extranet_g_neg = /\(([^:^>]+:[\d.]+)\)/g; +function tryToRemoveExtraNetworkFromPrompt(textarea, text, isNeg) { + var m = text.match(isNeg ? re_extranet_neg : re_extranet); + var replaced = false; + var newTextareaText; + var extraTextBeforeNet = opts.extra_networks_add_text_separator; + if (m) { + var extraTextAfterNet = m[2]; + var partToSearch = m[1]; + var foundAtPosition = -1; + newTextareaText = textarea.value.replaceAll(isNeg ? re_extranet_g_neg : re_extranet_g, function(found, net, pos) { + m = found.match(isNeg ? re_extranet_neg : re_extranet); + if (m[1] == partToSearch) { + replaced = true; + foundAtPosition = pos; + return ""; + } + return found; + }); + if (foundAtPosition >= 0) { + if (extraTextAfterNet && newTextareaText.substr(foundAtPosition, extraTextAfterNet.length) == extraTextAfterNet) { + newTextareaText = newTextareaText.substr(0, foundAtPosition) + newTextareaText.substr(foundAtPosition + extraTextAfterNet.length); + } + if (newTextareaText.substr(foundAtPosition - extraTextBeforeNet.length, extraTextBeforeNet.length) == extraTextBeforeNet) { + newTextareaText = newTextareaText.substr(0, foundAtPosition - extraTextBeforeNet.length) + newTextareaText.substr(foundAtPosition); + } + } + } else { + newTextareaText = textarea.value.replaceAll(new RegExp(`((?:${extraTextBeforeNet})?${text})`, "g"), ""); + replaced = (newTextareaText != textarea.value); + } + + if (replaced) { + textarea.value = newTextareaText; + return true; + } + + return false; +} + +function updatePromptArea(text, textArea, isNeg) { + if (!tryToRemoveExtraNetworkFromPrompt(textArea, text, isNeg)) { + textArea.value = textArea.value + opts.extra_networks_add_text_separator + text; + } + + updateInput(textArea); +} + +function cardClicked(tabname, textToAdd, textToAddNegative, allowNegativePrompt) { + if (textToAddNegative.length > 0) { + updatePromptArea(textToAdd, gradioApp().querySelector("#" + tabname + "_prompt > label > textarea")); + updatePromptArea(textToAddNegative, gradioApp().querySelector("#" + tabname + "_neg_prompt > label > textarea"), true); + } else { + var textarea = allowNegativePrompt ? activePromptTextarea[tabname] : gradioApp().querySelector("#" + tabname + "_prompt > label > textarea"); + updatePromptArea(textToAdd, textarea); + } +} + +function saveCardPreview(event, tabname, filename) { + var textarea = gradioApp().querySelector("#" + tabname + '_preview_filename > label > textarea'); + var button = gradioApp().getElementById(tabname + '_save_preview'); + + textarea.value = filename; + updateInput(textarea); + + button.click(); + + event.stopPropagation(); + event.preventDefault(); +} + +function extraNetworksSearchButton(tabname, extra_networks_tabname, event) { + var searchTextarea = gradioApp().querySelector("#" + tabname + "_" + extra_networks_tabname + "_extra_search"); + var button = event.target; + var text = button.classList.contains("search-all") ? "" : button.textContent.trim(); + + searchTextarea.value = text; + updateInput(searchTextarea); +} + +function extraNetworksTreeProcessFileClick(event, btn, tabname, extra_networks_tabname) { + /** + * Processes `onclick` events when user clicks on files in tree. + * + * @param event The generated event. + * @param btn The clicked `tree-list-item` button. + * @param tabname The name of the active tab in the sd webui. Ex: txt2img, img2img, etc. + * @param extra_networks_tabname The id of the active extraNetworks tab. Ex: lora, checkpoints, etc. + */ + // NOTE: Currently unused. + return; +} + +function extraNetworksTreeProcessDirectoryClick(event, btn, tabname, extra_networks_tabname) { + /** + * Processes `onclick` events when user clicks on directories in tree. + * + * Here is how the tree reacts to clicks for various states: + * unselected unopened directory: Directory is selected and expanded. + * unselected opened directory: Directory is selected. + * selected opened directory: Directory is collapsed and deselected. + * chevron is clicked: Directory is expanded or collapsed. Selected state unchanged. + * + * @param event The generated event. + * @param btn The clicked `tree-list-item` button. + * @param tabname The name of the active tab in the sd webui. Ex: txt2img, img2img, etc. + * @param extra_networks_tabname The id of the active extraNetworks tab. Ex: lora, checkpoints, etc. + */ + var ul = btn.nextElementSibling; + // This is the actual target that the user clicked on within the target button. + // We use this to detect if the chevron was clicked. + var true_targ = event.target; + + function _expand_or_collapse(_ul, _btn) { + // Expands
    if it is collapsed, collapses otherwise. Updates button attributes. + if (_ul.hasAttribute("hidden")) { + _ul.removeAttribute("hidden"); + _btn.dataset.expanded = ""; + } else { + _ul.setAttribute("hidden", ""); + delete _btn.dataset.expanded; + } + } + + function _remove_selected_from_all() { + // Removes the `selected` attribute from all buttons. + var sels = document.querySelectorAll("div.tree-list-content"); + [...sels].forEach(el => { + delete el.dataset.selected; + }); + } + + function _select_button(_btn) { + // Removes `data-selected` attribute from all buttons then adds to passed button. + _remove_selected_from_all(); + _btn.dataset.selected = ""; + } + + function _update_search(_tabname, _extra_networks_tabname, _search_text) { + // Update search input with select button's path. + var search_input_elem = gradioApp().querySelector("#" + tabname + "_" + extra_networks_tabname + "_extra_search"); + search_input_elem.value = _search_text; + updateInput(search_input_elem); + } + + + // If user clicks on the chevron, then we do not select the folder. + if (true_targ.matches(".tree-list-item-action--leading, .tree-list-item-action-chevron")) { + _expand_or_collapse(ul, btn); + } else { + // User clicked anywhere else on the button. + if ("selected" in btn.dataset && !(ul.hasAttribute("hidden"))) { + // If folder is select and open, collapse and deselect button. + _expand_or_collapse(ul, btn); + delete btn.dataset.selected; + _update_search(tabname, extra_networks_tabname, ""); + } else if (!(!("selected" in btn.dataset) && !(ul.hasAttribute("hidden")))) { + // If folder is open and not selected, then we don't collapse; just select. + // NOTE: Double inversion sucks but it is the clearest way to show the branching here. + _expand_or_collapse(ul, btn); + _select_button(btn, tabname, extra_networks_tabname); + _update_search(tabname, extra_networks_tabname, btn.dataset.path); + } else { + // All other cases, just select the button. + _select_button(btn, tabname, extra_networks_tabname); + _update_search(tabname, extra_networks_tabname, btn.dataset.path); + } + } +} + +function extraNetworksTreeOnClick(event, tabname, extra_networks_tabname) { + /** + * Handles `onclick` events for buttons within an `extra-network-tree .tree-list--tree`. + * + * Determines whether the clicked button in the tree is for a file entry or a directory + * then calls the appropriate function. + * + * @param event The generated event. + * @param tabname The name of the active tab in the sd webui. Ex: txt2img, img2img, etc. + * @param extra_networks_tabname The id of the active extraNetworks tab. Ex: lora, checkpoints, etc. + */ + var btn = event.currentTarget; + var par = btn.parentElement; + if (par.dataset.treeEntryType === "file") { + extraNetworksTreeProcessFileClick(event, btn, tabname, extra_networks_tabname); + } else { + extraNetworksTreeProcessDirectoryClick(event, btn, tabname, extra_networks_tabname); + } +} + +function extraNetworksControlSortOnClick(event, tabname, extra_networks_tabname) { + /** Handles `onclick` events for Sort Mode buttons. */ + + var self = event.currentTarget; + var parent = event.currentTarget.parentElement; + + parent.querySelectorAll('.extra-network-control--sort').forEach(function(x) { + x.classList.remove('extra-network-control--enabled'); + }); + + self.classList.add('extra-network-control--enabled'); + + applyExtraNetworkSort(tabname + "_" + extra_networks_tabname); +} + +function extraNetworksControlSortDirOnClick(event, tabname, extra_networks_tabname) { + /** + * Handles `onclick` events for the Sort Direction button. + * + * Modifies the data attributes of the Sort Direction button to cycle between + * ascending and descending sort directions. + * + * @param event The generated event. + * @param tabname The name of the active tab in the sd webui. Ex: txt2img, img2img, etc. + * @param extra_networks_tabname The id of the active extraNetworks tab. Ex: lora, checkpoints, etc. + */ + if (event.currentTarget.dataset.sortdir == "Ascending") { + event.currentTarget.dataset.sortdir = "Descending"; + event.currentTarget.setAttribute("title", "Sort descending"); + } else { + event.currentTarget.dataset.sortdir = "Ascending"; + event.currentTarget.setAttribute("title", "Sort ascending"); + } + applyExtraNetworkSort(tabname + "_" + extra_networks_tabname); +} + +function extraNetworksControlTreeViewOnClick(event, tabname, extra_networks_tabname) { + /** + * Handles `onclick` events for the Tree View button. + * + * Toggles the tree view in the extra networks pane. + * + * @param event The generated event. + * @param tabname The name of the active tab in the sd webui. Ex: txt2img, img2img, etc. + * @param extra_networks_tabname The id of the active extraNetworks tab. Ex: lora, checkpoints, etc. + */ + var button = event.currentTarget; + button.classList.toggle("extra-network-control--enabled"); + var show = !button.classList.contains("extra-network-control--enabled"); + + var pane = gradioApp().getElementById(tabname + "_" + extra_networks_tabname + "_pane"); + pane.classList.toggle("extra-network-dirs-hidden", show); +} + +function extraNetworksControlRefreshOnClick(event, tabname, extra_networks_tabname) { + /** + * Handles `onclick` events for the Refresh Page button. + * + * In order to actually call the python functions in `ui_extra_networks.py` + * to refresh the page, we created an empty gradio button in that file with an + * event handler that refreshes the page. So what this function here does + * is it manually raises a `click` event on that button. + * + * @param event The generated event. + * @param tabname The name of the active tab in the sd webui. Ex: txt2img, img2img, etc. + * @param extra_networks_tabname The id of the active extraNetworks tab. Ex: lora, checkpoints, etc. + */ + var btn_refresh_internal = gradioApp().getElementById(tabname + "_" + extra_networks_tabname + "_extra_refresh_internal"); + btn_refresh_internal.dispatchEvent(new Event("click")); +} + +var globalPopup = null; +var globalPopupInner = null; + +function closePopup() { + if (!globalPopup) return; + globalPopup.style.display = "none"; +} + +function popup(contents) { + if (!globalPopup) { + globalPopup = document.createElement('div'); + globalPopup.classList.add('global-popup'); + + var close = document.createElement('div'); + close.classList.add('global-popup-close'); + close.addEventListener("click", closePopup); + close.title = "Close"; + globalPopup.appendChild(close); + + globalPopupInner = document.createElement('div'); + globalPopupInner.classList.add('global-popup-inner'); + globalPopup.appendChild(globalPopupInner); + + gradioApp().querySelector('.main').appendChild(globalPopup); + } + + globalPopupInner.innerHTML = ''; + globalPopupInner.appendChild(contents); + + globalPopup.style.display = "flex"; +} + +var storedPopupIds = {}; +function popupId(id) { + if (!storedPopupIds[id]) { + storedPopupIds[id] = gradioApp().getElementById(id); + } + + popup(storedPopupIds[id]); +} + +function extraNetworksFlattenMetadata(obj) { + const result = {}; + + // Convert any stringified JSON objects to actual objects + for (const key of Object.keys(obj)) { + if (typeof obj[key] === 'string') { + try { + const parsed = JSON.parse(obj[key]); + if (parsed && typeof parsed === 'object') { + obj[key] = parsed; + } + } catch (error) { + continue; + } + } + } + + // Flatten the object + for (const key of Object.keys(obj)) { + if (typeof obj[key] === 'object' && obj[key] !== null) { + const nested = extraNetworksFlattenMetadata(obj[key]); + for (const nestedKey of Object.keys(nested)) { + result[`${key}/${nestedKey}`] = nested[nestedKey]; + } + } else { + result[key] = obj[key]; + } + } + + // Special case for handling modelspec keys + for (const key of Object.keys(result)) { + if (key.startsWith("modelspec.")) { + result[key.replaceAll(".", "/")] = result[key]; + delete result[key]; + } + } + + // Add empty keys to designate hierarchy + for (const key of Object.keys(result)) { + const parts = key.split("/"); + for (let i = 1; i < parts.length; i++) { + const parent = parts.slice(0, i).join("/"); + if (!result[parent]) { + result[parent] = ""; + } + } + } + + return result; +} + +function extraNetworksShowMetadata(text) { + try { + let parsed = JSON.parse(text); + if (parsed && typeof parsed === 'object') { + parsed = extraNetworksFlattenMetadata(parsed); + const table = createVisualizationTable(parsed, 0); + popup(table); + return; + } + } catch (error) { + console.error(error); + } + + var elem = document.createElement('pre'); + elem.classList.add('popup-metadata'); + elem.textContent = text; + + popup(elem); + return; +} + +function requestGet(url, data, handler, errorHandler) { + var xhr = new XMLHttpRequest(); + var args = Object.keys(data).map(function(k) { + return encodeURIComponent(k) + '=' + encodeURIComponent(data[k]); + }).join('&'); + xhr.open("GET", url + "?" + args, true); + + xhr.onreadystatechange = function() { + if (xhr.readyState === 4) { + if (xhr.status === 200) { + try { + var js = JSON.parse(xhr.responseText); + handler(js); + } catch (error) { + console.error(error); + errorHandler(); + } + } else { + errorHandler(); + } + } + }; + var js = JSON.stringify(data); + xhr.send(js); +} + +function extraNetworksCopyCardPath(event) { + navigator.clipboard.writeText(event.target.getAttribute("data-clipboard-text")); + event.stopPropagation(); +} + +function extraNetworksRequestMetadata(event, extraPage) { + var showError = function() { + extraNetworksShowMetadata("there was an error getting metadata"); + }; + + var cardName = event.target.parentElement.parentElement.getAttribute("data-name"); + + requestGet("./sd_extra_networks/metadata", {page: extraPage, item: cardName}, function(data) { + if (data && data.metadata) { + extraNetworksShowMetadata(data.metadata); + } else { + showError(); + } + }, showError); + + event.stopPropagation(); +} + +var extraPageUserMetadataEditors = {}; + +function extraNetworksEditUserMetadata(event, tabname, extraPage) { + var id = tabname + '_' + extraPage + '_edit_user_metadata'; + + var editor = extraPageUserMetadataEditors[id]; + if (!editor) { + editor = {}; + editor.page = gradioApp().getElementById(id); + editor.nameTextarea = gradioApp().querySelector("#" + id + "_name" + ' textarea'); + editor.button = gradioApp().querySelector("#" + id + "_button"); + extraPageUserMetadataEditors[id] = editor; + } + + var cardName = event.target.parentElement.parentElement.getAttribute("data-name"); + editor.nameTextarea.value = cardName; + updateInput(editor.nameTextarea); + + editor.button.click(); + + popup(editor.page); + + event.stopPropagation(); +} + +function extraNetworksRefreshSingleCard(page, tabname, name) { + requestGet("./sd_extra_networks/get-single-card", {page: page, tabname: tabname, name: name}, function(data) { + if (data && data.html) { + var card = gradioApp().querySelector(`#${tabname}_${page.replace(" ", "_")}_cards > .card[data-name="${name}"]`); + + var newDiv = document.createElement('DIV'); + newDiv.innerHTML = data.html; + var newCard = newDiv.firstElementChild; + + newCard.style.display = ''; + card.parentElement.insertBefore(newCard, card); + card.parentElement.removeChild(card); + } + }); +} + +window.addEventListener("keydown", function(event) { + if (event.key == "Escape") { + closePopup(); + } +}); + +/** + * Setup custom loading for this script. + * We need to wait for all of our HTML to be generated in the extra networks tabs + * before we can actually run the `setupExtraNetworks` function. + * The `onUiLoaded` function actually runs before all of our extra network tabs are + * finished generating. Thus we needed this new method. + * + */ + +var uiAfterScriptsCallbacks = []; +var uiAfterScriptsTimeout = null; +var executedAfterScripts = false; + +function scheduleAfterScriptsCallbacks() { + clearTimeout(uiAfterScriptsTimeout); + uiAfterScriptsTimeout = setTimeout(function() { + executeCallbacks(uiAfterScriptsCallbacks); + }, 200); +} + +onUiLoaded(function() { + var mutationObserver = new MutationObserver(function(m) { + let existingSearchfields = gradioApp().querySelectorAll("[id$='_extra_search']").length; + let neededSearchfields = gradioApp().querySelectorAll("[id$='_extra_tabs'] > .tab-nav > button").length - 2; + + if (!executedAfterScripts && existingSearchfields >= neededSearchfields) { + mutationObserver.disconnect(); + executedAfterScripts = true; + scheduleAfterScriptsCallbacks(); + } + }); + mutationObserver.observe(gradioApp(), {childList: true, subtree: true}); +}); + +uiAfterScriptsCallbacks.push(setupExtraNetworks); diff --git a/stable-diffusion-webui/javascript/generationParams.js b/stable-diffusion-webui/javascript/generationParams.js new file mode 100755 index 0000000..7c0fd22 --- /dev/null +++ b/stable-diffusion-webui/javascript/generationParams.js @@ -0,0 +1,35 @@ +// attaches listeners to the txt2img and img2img galleries to update displayed generation param text when the image changes + +let txt2img_gallery, img2img_gallery, modal = undefined; +onAfterUiUpdate(function() { + if (!txt2img_gallery) { + txt2img_gallery = attachGalleryListeners("txt2img"); + } + if (!img2img_gallery) { + img2img_gallery = attachGalleryListeners("img2img"); + } + if (!modal) { + modal = gradioApp().getElementById('lightboxModal'); + modalObserver.observe(modal, {attributes: true, attributeFilter: ['style']}); + } +}); + +let modalObserver = new MutationObserver(function(mutations) { + mutations.forEach(function(mutationRecord) { + let selectedTab = gradioApp().querySelector('#tabs div button.selected')?.innerText; + if (mutationRecord.target.style.display === 'none' && (selectedTab === 'txt2img' || selectedTab === 'img2img')) { + gradioApp().getElementById(selectedTab + "_generation_info_button")?.click(); + } + }); +}); + +function attachGalleryListeners(tab_name) { + var gallery = gradioApp().querySelector('#' + tab_name + '_gallery'); + gallery?.addEventListener('click', () => gradioApp().getElementById(tab_name + "_generation_info_button").click()); + gallery?.addEventListener('keydown', (e) => { + if (e.keyCode == 37 || e.keyCode == 39) { // left or right arrow + gradioApp().getElementById(tab_name + "_generation_info_button").click(); + } + }); + return gallery; +} diff --git a/stable-diffusion-webui/javascript/hints.js b/stable-diffusion-webui/javascript/hints.js new file mode 100755 index 0000000..6de9372 --- /dev/null +++ b/stable-diffusion-webui/javascript/hints.js @@ -0,0 +1,203 @@ +// mouseover tooltips for various UI elements + +var titles = { + "Sampling steps": "How many times to improve the generated image iteratively; higher values take longer; very low values can produce bad results", + "Sampling method": "Which algorithm to use to produce the image", + "GFPGAN": "Restore low quality faces using GFPGAN neural network", + "Euler a": "Euler Ancestral - very creative, each can get a completely different picture depending on step count, setting steps higher than 30-40 does not help", + "DDIM": "Denoising Diffusion Implicit Models - best at inpainting", + "UniPC": "Unified Predictor-Corrector Framework for Fast Sampling of Diffusion Models", + "DPM adaptive": "Ignores step count - uses a number of steps determined by the CFG and resolution", + + "\u{1F4D0}": "Auto detect size from img2img", + "Batch count": "How many batches of images to create (has no impact on generation performance or VRAM usage)", + "Batch size": "How many image to create in a single batch (increases generation performance at cost of higher VRAM usage)", + "CFG Scale": "Classifier Free Guidance Scale - how strongly the image should conform to prompt - lower values produce more creative results", + "Seed": "A value that determines the output of random number generator - if you create an image with same parameters and seed as another image, you'll get the same result", + "\u{1f3b2}\ufe0f": "Set seed to -1, which will cause a new random number to be used every time", + "\u267b\ufe0f": "Reuse seed from last generation, mostly useful if it was randomized", + "\u2199\ufe0f": "Read generation parameters from prompt or last generation if prompt is empty into user interface.", + "\u{1f4c2}": "Open images output directory", + "\u{1f4be}": "Save style", + "\u{1f5d1}\ufe0f": "Clear prompt", + "\u{1f4cb}": "Apply selected styles to current prompt", + "\u{1f4d2}": "Paste available values into the field", + "\u{1f3b4}": "Show/hide extra networks", + "\u{1f300}": "Restore progress", + + "Inpaint a part of image": "Draw a mask over an image, and the script will regenerate the masked area with content according to prompt", + "SD upscale": "Upscale image normally, split result into tiles, improve each tile using img2img, merge whole image back", + + "Just resize": "Resize image to target resolution. Unless height and width match, you will get incorrect aspect ratio.", + "Crop and resize": "Resize the image so that entirety of target resolution is filled with the image. Crop parts that stick out.", + "Resize and fill": "Resize the image so that entirety of image is inside target resolution. Fill empty space with image's colors.", + + "Mask blur": "How much to blur the mask before processing, in pixels.", + "Masked content": "What to put inside the masked area before processing it with Stable Diffusion.", + "fill": "fill it with colors of the image", + "original": "keep whatever was there originally", + "latent noise": "fill it with latent space noise", + "latent nothing": "fill it with latent space zeroes", + "Inpaint at full resolution": "Upscale masked region to target resolution, do inpainting, downscale back and paste into original image", + + "Denoising strength": "Determines how little respect the algorithm should have for image's content. At 0, nothing will change, and at 1 you'll get an unrelated image. With values below 1.0, processing will take less steps than the Sampling Steps slider specifies.", + + "Skip": "Stop processing current image and continue processing.", + "Interrupt": "Stop processing images and return any results accumulated so far.", + "Save": "Write image to a directory (default - log/images) and generation parameters into csv file.", + + "X values": "Separate values for X axis using commas.", + "Y values": "Separate values for Y axis using commas.", + + "None": "Do not do anything special", + "Prompt matrix": "Separate prompts into parts using vertical pipe character (|) and the script will create a picture for every combination of them (except for the first part, which will be present in all combinations)", + "X/Y/Z plot": "Create grid(s) where images will have different parameters. Use inputs below to specify which parameters will be shared by columns and rows", + "Custom code": "Run Python code. Advanced user only. Must run program with --allow-code for this to work", + + "Prompt S/R": "Separate a list of words with commas, and the first word will be used as a keyword: script will search for this word in the prompt, and replace it with others", + "Prompt order": "Separate a list of words with commas, and the script will make a variation of prompt with those words for their every possible order", + + "Tiling": "Produce an image that can be tiled.", + "Tile overlap": "For SD upscale, how much overlap in pixels should there be between tiles. Tiles overlap so that when they are merged back into one picture, there is no clearly visible seam.", + + "Variation seed": "Seed of a different picture to be mixed into the generation.", + "Variation strength": "How strong of a variation to produce. At 0, there will be no effect. At 1, you will get the complete picture with variation seed (except for ancestral samplers, where you will just get something).", + "Resize seed from height": "Make an attempt to produce a picture similar to what would have been produced with same seed at specified resolution", + "Resize seed from width": "Make an attempt to produce a picture similar to what would have been produced with same seed at specified resolution", + + "Interrogate": "Reconstruct prompt from existing image and put it into the prompt field.", + + "Images filename pattern": "Use tags like [seed] and [date] to define how filenames for images are chosen. Leave empty for default.", + "Directory name pattern": "Use tags like [seed] and [date] to define how subdirectories for images and grids are chosen. Leave empty for default.", + "Max prompt words": "Set the maximum number of words to be used in the [prompt_words] option; ATTENTION: If the words are too long, they may exceed the maximum length of the file path that the system can handle", + + "Loopback": "Performs img2img processing multiple times. Output images are used as input for the next loop.", + "Loops": "How many times to process an image. Each output is used as the input of the next loop. If set to 1, behavior will be as if this script were not used.", + "Final denoising strength": "The denoising strength for the final loop of each image in the batch.", + "Denoising strength curve": "The denoising curve controls the rate of denoising strength change each loop. Aggressive: Most of the change will happen towards the start of the loops. Linear: Change will be constant through all loops. Lazy: Most of the change will happen towards the end of the loops.", + + "Style 1": "Style to apply; styles have components for both positive and negative prompts and apply to both", + "Style 2": "Style to apply; styles have components for both positive and negative prompts and apply to both", + "Apply style": "Insert selected styles into prompt fields", + "Create style": "Save current prompts as a style. If you add the token {prompt} to the text, the style uses that as a placeholder for your prompt when you use the style in the future.", + + "Checkpoint name": "Loads weights from checkpoint before making images. You can either use hash or a part of filename (as seen in settings) for checkpoint name. Recommended to use with Y axis for less switching.", + "Inpainting conditioning mask strength": "Only applies to inpainting models. Determines how strongly to mask off the original image for inpainting and img2img. 1.0 means fully masked, which is the default behaviour. 0.0 means a fully unmasked conditioning. Lower values will help preserve the overall composition of the image, but will struggle with large changes.", + + "Eta noise seed delta": "If this values is non-zero, it will be added to seed and used to initialize RNG for noises when using samplers with Eta. You can use this to produce even more variation of images, or you can use this to match images of other software if you know what you are doing.", + + "Filename word regex": "This regular expression will be used extract words from filename, and they will be joined using the option below into label text used for training. Leave empty to keep filename text as it is.", + "Filename join string": "This string will be used to join split words into a single line if the option above is enabled.", + + "Quicksettings list": "List of setting names, separated by commas, for settings that should go to the quick access bar at the top, rather than the usual setting tab. See modules/shared.py for setting names. Requires restarting to apply.", + + "Weighted sum": "Result = A * (1 - M) + B * M", + "Add difference": "Result = A + (B - C) * M", + "No interpolation": "Result = A", + + "Initialization text": "If the number of tokens is more than the number of vectors, some may be skipped.\nLeave the textbox empty to start with zeroed out vectors", + "Learning rate": "How fast should training go. Low values will take longer to train, high values may fail to converge (not generate accurate results) and/or may break the embedding (This has happened if you see Loss: nan in the training info textbox. If this happens, you need to manually restore your embedding from an older not-broken backup).\n\nYou can set a single numeric value, or multiple learning rates using the syntax:\n\n rate_1:max_steps_1, rate_2:max_steps_2, ...\n\nEG: 0.005:100, 1e-3:1000, 1e-5\n\nWill train with rate of 0.005 for first 100 steps, then 1e-3 until 1000 steps, then 1e-5 for all remaining steps.", + + "Clip skip": "Early stopping parameter for CLIP model; 1 is stop at last layer as usual, 2 is stop at penultimate layer, etc.", + + "Approx NN": "Cheap neural network approximation. Very fast compared to VAE, but produces pictures with 4 times smaller horizontal/vertical resolution and lower quality.", + "Approx cheap": "Very cheap approximation. Very fast compared to VAE, but produces pictures with 8 times smaller horizontal/vertical resolution and extremely low quality.", + + "Hires. fix": "Use a two step process to partially create an image at smaller resolution, upscale, and then improve details in it without changing composition", + "Hires steps": "Number of sampling steps for upscaled picture. If 0, uses same as for original.", + "Upscale by": "Adjusts the size of the image by multiplying the original width and height by the selected value. Ignored if either Resize width to or Resize height to are non-zero.", + "Resize width to": "Resizes image to this width. If 0, width is inferred from either of two nearby sliders.", + "Resize height to": "Resizes image to this height. If 0, height is inferred from either of two nearby sliders.", + "Discard weights with matching name": "Regular expression; if weights's name matches it, the weights is not written to the resulting checkpoint. Use ^model_ema to discard EMA weights.", + "Extra networks tab order": "Comma-separated list of tab names; tabs listed here will appear in the extra networks UI first and in order listed.", + "Negative Guidance minimum sigma": "Skip negative prompt for steps where image is already mostly denoised; the higher this value, the more skips there will be; provides increased performance in exchange for minor quality reduction." +}; + +function updateTooltip(element) { + if (element.title) return; // already has a title + + let text = element.textContent; + let tooltip = localization[titles[text]] || titles[text]; + + if (!tooltip) { + let value = element.value; + if (value) tooltip = localization[titles[value]] || titles[value]; + } + + if (!tooltip) { + // Gradio dropdown options have `data-value`. + let dataValue = element.dataset.value; + if (dataValue) tooltip = localization[titles[dataValue]] || titles[dataValue]; + } + + if (!tooltip) { + for (const c of element.classList) { + if (c in titles) { + tooltip = localization[titles[c]] || titles[c]; + break; + } + } + } + + if (tooltip) { + element.title = tooltip; + } +} + +// Nodes to check for adding tooltips. +const tooltipCheckNodes = new Set(); +// Timer for debouncing tooltip check. +let tooltipCheckTimer = null; + +function processTooltipCheckNodes() { + for (const node of tooltipCheckNodes) { + updateTooltip(node); + } + tooltipCheckNodes.clear(); +} + +onUiUpdate(function(mutationRecords) { + for (const record of mutationRecords) { + if (record.type === "childList" && record.target.classList.contains("options")) { + // This smells like a Gradio dropdown menu having changed, + // so let's enqueue an update for the input element that shows the current value. + let wrap = record.target.parentNode; + let input = wrap?.querySelector("input"); + if (input) { + input.title = ""; // So we'll even have a chance to update it. + tooltipCheckNodes.add(input); + } + } + for (const node of record.addedNodes) { + if (node.nodeType === Node.ELEMENT_NODE && !node.classList.contains("hide")) { + if (!node.title) { + if ( + node.tagName === "SPAN" || + node.tagName === "BUTTON" || + node.tagName === "P" || + node.tagName === "INPUT" || + (node.tagName === "LI" && node.classList.contains("item")) // Gradio dropdown item + ) { + tooltipCheckNodes.add(node); + } + } + node.querySelectorAll('span, button, p').forEach(n => tooltipCheckNodes.add(n)); + } + } + } + if (tooltipCheckNodes.size) { + clearTimeout(tooltipCheckTimer); + tooltipCheckTimer = setTimeout(processTooltipCheckNodes, 1000); + } +}); + +onUiLoaded(function() { + for (var comp of window.gradio_config.components) { + if (comp.props.webui_tooltip && comp.props.elem_id) { + var elem = gradioApp().getElementById(comp.props.elem_id); + if (elem) { + elem.title = comp.props.webui_tooltip; + } + } + } +}); diff --git a/stable-diffusion-webui/javascript/hires_fix.js b/stable-diffusion-webui/javascript/hires_fix.js new file mode 100755 index 0000000..0d04ab3 --- /dev/null +++ b/stable-diffusion-webui/javascript/hires_fix.js @@ -0,0 +1,18 @@ + +function onCalcResolutionHires(enable, width, height, hr_scale, hr_resize_x, hr_resize_y) { + function setInactive(elem, inactive) { + elem.classList.toggle('inactive', !!inactive); + } + + var hrUpscaleBy = gradioApp().getElementById('txt2img_hr_scale'); + var hrResizeX = gradioApp().getElementById('txt2img_hr_resize_x'); + var hrResizeY = gradioApp().getElementById('txt2img_hr_resize_y'); + + gradioApp().getElementById('txt2img_hires_fix_row2').style.display = opts.use_old_hires_fix_width_height ? "none" : ""; + + setInactive(hrUpscaleBy, opts.use_old_hires_fix_width_height || hr_resize_x > 0 || hr_resize_y > 0); + setInactive(hrResizeX, opts.use_old_hires_fix_width_height || hr_resize_x == 0); + setInactive(hrResizeY, opts.use_old_hires_fix_width_height || hr_resize_y == 0); + + return [enable, width, height, hr_scale, hr_resize_x, hr_resize_y]; +} diff --git a/stable-diffusion-webui/javascript/imageMaskFix.js b/stable-diffusion-webui/javascript/imageMaskFix.js new file mode 100755 index 0000000..900c56f --- /dev/null +++ b/stable-diffusion-webui/javascript/imageMaskFix.js @@ -0,0 +1,43 @@ +/** + * temporary fix for https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/668 + * @see https://github.com/gradio-app/gradio/issues/1721 + */ +function imageMaskResize() { + const canvases = gradioApp().querySelectorAll('#img2maskimg .touch-none canvas'); + if (!canvases.length) { + window.removeEventListener('resize', imageMaskResize); + return; + } + + const wrapper = canvases[0].closest('.touch-none'); + const previewImage = wrapper.previousElementSibling; + + if (!previewImage.complete) { + previewImage.addEventListener('load', imageMaskResize); + return; + } + + const w = previewImage.width; + const h = previewImage.height; + const nw = previewImage.naturalWidth; + const nh = previewImage.naturalHeight; + const portrait = nh > nw; + + const wW = Math.min(w, portrait ? h / nh * nw : w / nw * nw); + const wH = Math.min(h, portrait ? h / nh * nh : w / nw * nh); + + wrapper.style.width = `${wW}px`; + wrapper.style.height = `${wH}px`; + wrapper.style.left = `0px`; + wrapper.style.top = `0px`; + + canvases.forEach(c => { + c.style.width = c.style.height = ''; + c.style.maxWidth = '100%'; + c.style.maxHeight = '100%'; + c.style.objectFit = 'contain'; + }); +} + +onAfterUiUpdate(imageMaskResize); +window.addEventListener('resize', imageMaskResize); diff --git a/stable-diffusion-webui/javascript/imageviewer.js b/stable-diffusion-webui/javascript/imageviewer.js new file mode 100755 index 0000000..9b23f47 --- /dev/null +++ b/stable-diffusion-webui/javascript/imageviewer.js @@ -0,0 +1,268 @@ +// A full size 'lightbox' preview modal shown when left clicking on gallery previews +function closeModal() { + gradioApp().getElementById("lightboxModal").style.display = "none"; +} + +function showModal(event) { + const source = event.target || event.srcElement; + const modalImage = gradioApp().getElementById("modalImage"); + const modalToggleLivePreviewBtn = gradioApp().getElementById("modal_toggle_live_preview"); + modalToggleLivePreviewBtn.innerHTML = opts.js_live_preview_in_modal_lightbox ? "🗇" : "🗆"; + const lb = gradioApp().getElementById("lightboxModal"); + modalImage.src = source.src; + if (modalImage.style.display === 'none') { + lb.style.setProperty('background-image', 'url(' + source.src + ')'); + } + lb.style.display = "flex"; + lb.focus(); + + const tabTxt2Img = gradioApp().getElementById("tab_txt2img"); + const tabImg2Img = gradioApp().getElementById("tab_img2img"); + // show the save button in modal only on txt2img or img2img tabs + if (tabTxt2Img.style.display != "none" || tabImg2Img.style.display != "none") { + gradioApp().getElementById("modal_save").style.display = "inline"; + } else { + gradioApp().getElementById("modal_save").style.display = "none"; + } + event.stopPropagation(); +} + +function negmod(n, m) { + return ((n % m) + m) % m; +} + +function updateOnBackgroundChange() { + const modalImage = gradioApp().getElementById("modalImage"); + if (modalImage && modalImage.offsetParent) { + let currentButton = selected_gallery_button(); + let preview = gradioApp().querySelectorAll('.livePreview > img'); + if (opts.js_live_preview_in_modal_lightbox && preview.length > 0) { + // show preview image if available + modalImage.src = preview[preview.length - 1].src; + } else if (currentButton?.children?.length > 0 && modalImage.src != currentButton.children[0].src) { + modalImage.src = currentButton.children[0].src; + if (modalImage.style.display === 'none') { + const modal = gradioApp().getElementById("lightboxModal"); + modal.style.setProperty('background-image', `url(${modalImage.src})`); + } + } + } +} + +function modalImageSwitch(offset) { + var galleryButtons = all_gallery_buttons(); + + if (galleryButtons.length > 1) { + var result = selected_gallery_index(); + + if (result != -1) { + var nextButton = galleryButtons[negmod((result + offset), galleryButtons.length)]; + nextButton.click(); + const modalImage = gradioApp().getElementById("modalImage"); + const modal = gradioApp().getElementById("lightboxModal"); + modalImage.src = nextButton.children[0].src; + if (modalImage.style.display === 'none') { + modal.style.setProperty('background-image', `url(${modalImage.src})`); + } + setTimeout(function() { + modal.focus(); + }, 10); + } + } +} + +function saveImage() { + const tabTxt2Img = gradioApp().getElementById("tab_txt2img"); + const tabImg2Img = gradioApp().getElementById("tab_img2img"); + const saveTxt2Img = "save_txt2img"; + const saveImg2Img = "save_img2img"; + if (tabTxt2Img.style.display != "none") { + gradioApp().getElementById(saveTxt2Img).click(); + } else if (tabImg2Img.style.display != "none") { + gradioApp().getElementById(saveImg2Img).click(); + } else { + console.error("missing implementation for saving modal of this type"); + } +} + +function modalSaveImage(event) { + saveImage(); + event.stopPropagation(); +} + +function modalNextImage(event) { + modalImageSwitch(1); + event.stopPropagation(); +} + +function modalPrevImage(event) { + modalImageSwitch(-1); + event.stopPropagation(); +} + +function modalKeyHandler(event) { + switch (event.key) { + case "s": + saveImage(); + break; + case "ArrowLeft": + modalPrevImage(event); + break; + case "ArrowRight": + modalNextImage(event); + break; + case "Escape": + closeModal(); + break; + } +} + +function setupImageForLightbox(e) { + if (e.dataset.modded) { + return; + } + + e.dataset.modded = true; + e.style.cursor = 'pointer'; + e.style.userSelect = 'none'; + + e.addEventListener('mousedown', function(evt) { + if (evt.button == 1) { + open(evt.target.src); + evt.preventDefault(); + return; + } + }, true); + + e.addEventListener('click', function(evt) { + if (!opts.js_modal_lightbox || evt.button != 0) return; + + modalZoomSet(gradioApp().getElementById('modalImage'), opts.js_modal_lightbox_initially_zoomed); + evt.preventDefault(); + showModal(evt); + }, true); + +} + +function modalZoomSet(modalImage, enable) { + if (modalImage) modalImage.classList.toggle('modalImageFullscreen', !!enable); +} + +function modalZoomToggle(event) { + var modalImage = gradioApp().getElementById("modalImage"); + modalZoomSet(modalImage, !modalImage.classList.contains('modalImageFullscreen')); + event.stopPropagation(); +} + +function modalLivePreviewToggle(event) { + const modalToggleLivePreview = gradioApp().getElementById("modal_toggle_live_preview"); + opts.js_live_preview_in_modal_lightbox = !opts.js_live_preview_in_modal_lightbox; + modalToggleLivePreview.innerHTML = opts.js_live_preview_in_modal_lightbox ? "🗇" : "🗆"; + event.stopPropagation(); +} + +function modalTileImageToggle(event) { + const modalImage = gradioApp().getElementById("modalImage"); + const modal = gradioApp().getElementById("lightboxModal"); + const isTiling = modalImage.style.display === 'none'; + if (isTiling) { + modalImage.style.display = 'block'; + modal.style.setProperty('background-image', 'none'); + } else { + modalImage.style.display = 'none'; + modal.style.setProperty('background-image', `url(${modalImage.src})`); + } + + event.stopPropagation(); +} + +onAfterUiUpdate(function() { + var fullImg_preview = gradioApp().querySelectorAll('.gradio-gallery > div > img'); + if (fullImg_preview != null) { + fullImg_preview.forEach(setupImageForLightbox); + } + updateOnBackgroundChange(); +}); + +document.addEventListener("DOMContentLoaded", function() { + //const modalFragment = document.createDocumentFragment(); + const modal = document.createElement('div'); + modal.onclick = closeModal; + modal.id = "lightboxModal"; + modal.tabIndex = 0; + modal.addEventListener('keydown', modalKeyHandler, true); + + const modalControls = document.createElement('div'); + modalControls.className = 'modalControls gradio-container'; + modal.append(modalControls); + + const modalZoom = document.createElement('span'); + modalZoom.className = 'modalZoom cursor'; + modalZoom.innerHTML = '⤡'; + modalZoom.addEventListener('click', modalZoomToggle, true); + modalZoom.title = "Toggle zoomed view"; + modalControls.appendChild(modalZoom); + + const modalTileImage = document.createElement('span'); + modalTileImage.className = 'modalTileImage cursor'; + modalTileImage.innerHTML = '⊞'; + modalTileImage.addEventListener('click', modalTileImageToggle, true); + modalTileImage.title = "Preview tiling"; + modalControls.appendChild(modalTileImage); + + const modalSave = document.createElement("span"); + modalSave.className = "modalSave cursor"; + modalSave.id = "modal_save"; + modalSave.innerHTML = "🖫"; + modalSave.addEventListener("click", modalSaveImage, true); + modalSave.title = "Save Image(s)"; + modalControls.appendChild(modalSave); + + const modalToggleLivePreview = document.createElement('span'); + modalToggleLivePreview.className = 'modalToggleLivePreview cursor'; + modalToggleLivePreview.id = "modal_toggle_live_preview"; + modalToggleLivePreview.innerHTML = "🗆"; + modalToggleLivePreview.onclick = modalLivePreviewToggle; + modalToggleLivePreview.title = "Toggle live preview"; + modalControls.appendChild(modalToggleLivePreview); + + const modalClose = document.createElement('span'); + modalClose.className = 'modalClose cursor'; + modalClose.innerHTML = '×'; + modalClose.onclick = closeModal; + modalClose.title = "Close image viewer"; + modalControls.appendChild(modalClose); + + const modalImage = document.createElement('img'); + modalImage.id = 'modalImage'; + modalImage.onclick = closeModal; + modalImage.tabIndex = 0; + modalImage.addEventListener('keydown', modalKeyHandler, true); + modal.appendChild(modalImage); + + const modalPrev = document.createElement('a'); + modalPrev.className = 'modalPrev'; + modalPrev.innerHTML = '❮'; + modalPrev.tabIndex = 0; + modalPrev.addEventListener('click', modalPrevImage, true); + modalPrev.addEventListener('keydown', modalKeyHandler, true); + modal.appendChild(modalPrev); + + const modalNext = document.createElement('a'); + modalNext.className = 'modalNext'; + modalNext.innerHTML = '❯'; + modalNext.tabIndex = 0; + modalNext.addEventListener('click', modalNextImage, true); + modalNext.addEventListener('keydown', modalKeyHandler, true); + + modal.appendChild(modalNext); + + try { + gradioApp().appendChild(modal); + } catch (e) { + gradioApp().body.appendChild(modal); + } + + document.body.appendChild(modal); + +}); diff --git a/stable-diffusion-webui/javascript/imageviewerGamepad.js b/stable-diffusion-webui/javascript/imageviewerGamepad.js new file mode 100755 index 0000000..a22c7e6 --- /dev/null +++ b/stable-diffusion-webui/javascript/imageviewerGamepad.js @@ -0,0 +1,63 @@ +let gamepads = []; + +window.addEventListener('gamepadconnected', (e) => { + const index = e.gamepad.index; + let isWaiting = false; + gamepads[index] = setInterval(async() => { + if (!opts.js_modal_lightbox_gamepad || isWaiting) return; + const gamepad = navigator.getGamepads()[index]; + const xValue = gamepad.axes[0]; + if (xValue <= -0.3) { + modalPrevImage(e); + isWaiting = true; + } else if (xValue >= 0.3) { + modalNextImage(e); + isWaiting = true; + } + if (isWaiting) { + await sleepUntil(() => { + const xValue = navigator.getGamepads()[index].axes[0]; + if (xValue < 0.3 && xValue > -0.3) { + return true; + } + }, opts.js_modal_lightbox_gamepad_repeat); + isWaiting = false; + } + }, 10); +}); + +window.addEventListener('gamepaddisconnected', (e) => { + clearInterval(gamepads[e.gamepad.index]); +}); + +/* +Primarily for vr controller type pointer devices. +I use the wheel event because there's currently no way to do it properly with web xr. + */ +let isScrolling = false; +window.addEventListener('wheel', (e) => { + if (!opts.js_modal_lightbox_gamepad || isScrolling) return; + isScrolling = true; + + if (e.deltaX <= -0.6) { + modalPrevImage(e); + } else if (e.deltaX >= 0.6) { + modalNextImage(e); + } + + setTimeout(() => { + isScrolling = false; + }, opts.js_modal_lightbox_gamepad_repeat); +}); + +function sleepUntil(f, timeout) { + return new Promise((resolve) => { + const timeStart = new Date(); + const wait = setInterval(function() { + if (f() || new Date() - timeStart > timeout) { + clearInterval(wait); + resolve(); + } + }, 20); + }); +} diff --git a/stable-diffusion-webui/javascript/inputAccordion.js b/stable-diffusion-webui/javascript/inputAccordion.js new file mode 100755 index 0000000..7570309 --- /dev/null +++ b/stable-diffusion-webui/javascript/inputAccordion.js @@ -0,0 +1,68 @@ +function inputAccordionChecked(id, checked) { + var accordion = gradioApp().getElementById(id); + accordion.visibleCheckbox.checked = checked; + accordion.onVisibleCheckboxChange(); +} + +function setupAccordion(accordion) { + var labelWrap = accordion.querySelector('.label-wrap'); + var gradioCheckbox = gradioApp().querySelector('#' + accordion.id + "-checkbox input"); + var extra = gradioApp().querySelector('#' + accordion.id + "-extra"); + var span = labelWrap.querySelector('span'); + var linked = true; + + var isOpen = function() { + return labelWrap.classList.contains('open'); + }; + + var observerAccordionOpen = new MutationObserver(function(mutations) { + mutations.forEach(function(mutationRecord) { + accordion.classList.toggle('input-accordion-open', isOpen()); + + if (linked) { + accordion.visibleCheckbox.checked = isOpen(); + accordion.onVisibleCheckboxChange(); + } + }); + }); + observerAccordionOpen.observe(labelWrap, {attributes: true, attributeFilter: ['class']}); + + if (extra) { + labelWrap.insertBefore(extra, labelWrap.lastElementChild); + } + + accordion.onChecked = function(checked) { + if (isOpen() != checked) { + labelWrap.click(); + } + }; + + var visibleCheckbox = document.createElement('INPUT'); + visibleCheckbox.type = 'checkbox'; + visibleCheckbox.checked = isOpen(); + visibleCheckbox.id = accordion.id + "-visible-checkbox"; + visibleCheckbox.className = gradioCheckbox.className + " input-accordion-checkbox"; + span.insertBefore(visibleCheckbox, span.firstChild); + + accordion.visibleCheckbox = visibleCheckbox; + accordion.onVisibleCheckboxChange = function() { + if (linked && isOpen() != visibleCheckbox.checked) { + labelWrap.click(); + } + + gradioCheckbox.checked = visibleCheckbox.checked; + updateInput(gradioCheckbox); + }; + + visibleCheckbox.addEventListener('click', function(event) { + linked = false; + event.stopPropagation(); + }); + visibleCheckbox.addEventListener('input', accordion.onVisibleCheckboxChange); +} + +onUiLoaded(function() { + for (var accordion of gradioApp().querySelectorAll('.input-accordion')) { + setupAccordion(accordion); + } +}); diff --git a/stable-diffusion-webui/javascript/localStorage.js b/stable-diffusion-webui/javascript/localStorage.js new file mode 100755 index 0000000..dc1a36c --- /dev/null +++ b/stable-diffusion-webui/javascript/localStorage.js @@ -0,0 +1,26 @@ + +function localSet(k, v) { + try { + localStorage.setItem(k, v); + } catch (e) { + console.warn(`Failed to save ${k} to localStorage: ${e}`); + } +} + +function localGet(k, def) { + try { + return localStorage.getItem(k); + } catch (e) { + console.warn(`Failed to load ${k} from localStorage: ${e}`); + } + + return def; +} + +function localRemove(k) { + try { + return localStorage.removeItem(k); + } catch (e) { + console.warn(`Failed to remove ${k} from localStorage: ${e}`); + } +} diff --git a/stable-diffusion-webui/javascript/localization.js b/stable-diffusion-webui/javascript/localization.js new file mode 100755 index 0000000..8f00c18 --- /dev/null +++ b/stable-diffusion-webui/javascript/localization.js @@ -0,0 +1,205 @@ + +// localization = {} -- the dict with translations is created by the backend + +var ignore_ids_for_localization = { + setting_sd_hypernetwork: 'OPTION', + setting_sd_model_checkpoint: 'OPTION', + modelmerger_primary_model_name: 'OPTION', + modelmerger_secondary_model_name: 'OPTION', + modelmerger_tertiary_model_name: 'OPTION', + train_embedding: 'OPTION', + train_hypernetwork: 'OPTION', + txt2img_styles: 'OPTION', + img2img_styles: 'OPTION', + setting_random_artist_categories: 'OPTION', + setting_face_restoration_model: 'OPTION', + setting_realesrgan_enabled_models: 'OPTION', + extras_upscaler_1: 'OPTION', + extras_upscaler_2: 'OPTION', +}; + +var re_num = /^[.\d]+$/; +var re_emoji = /[\p{Extended_Pictographic}\u{1F3FB}-\u{1F3FF}\u{1F9B0}-\u{1F9B3}]/u; + +var original_lines = {}; +var translated_lines = {}; + +function hasLocalization() { + return window.localization && Object.keys(window.localization).length > 0; +} + +function textNodesUnder(el) { + var n, a = [], walk = document.createTreeWalker(el, NodeFilter.SHOW_TEXT, null, false); + while ((n = walk.nextNode())) a.push(n); + return a; +} + +function canBeTranslated(node, text) { + if (!text) return false; + if (!node.parentElement) return false; + + var parentType = node.parentElement.nodeName; + if (parentType == 'SCRIPT' || parentType == 'STYLE' || parentType == 'TEXTAREA') return false; + + if (parentType == 'OPTION' || parentType == 'SPAN') { + var pnode = node; + for (var level = 0; level < 4; level++) { + pnode = pnode.parentElement; + if (!pnode) break; + + if (ignore_ids_for_localization[pnode.id] == parentType) return false; + } + } + + if (re_num.test(text)) return false; + if (re_emoji.test(text)) return false; + return true; +} + +function getTranslation(text) { + if (!text) return undefined; + + if (translated_lines[text] === undefined) { + original_lines[text] = 1; + } + + var tl = localization[text]; + if (tl !== undefined) { + translated_lines[tl] = 1; + } + + return tl; +} + +function processTextNode(node) { + var text = node.textContent.trim(); + + if (!canBeTranslated(node, text)) return; + + var tl = getTranslation(text); + if (tl !== undefined) { + node.textContent = tl; + } +} + +function processNode(node) { + if (node.nodeType == 3) { + processTextNode(node); + return; + } + + if (node.title) { + let tl = getTranslation(node.title); + if (tl !== undefined) { + node.title = tl; + } + } + + if (node.placeholder) { + let tl = getTranslation(node.placeholder); + if (tl !== undefined) { + node.placeholder = tl; + } + } + + textNodesUnder(node).forEach(function(node) { + processTextNode(node); + }); +} + +function localizeWholePage() { + processNode(gradioApp()); + + function elem(comp) { + var elem_id = comp.props.elem_id ? comp.props.elem_id : "component-" + comp.id; + return gradioApp().getElementById(elem_id); + } + + for (var comp of window.gradio_config.components) { + if (comp.props.webui_tooltip) { + let e = elem(comp); + + let tl = e ? getTranslation(e.title) : undefined; + if (tl !== undefined) { + e.title = tl; + } + } + if (comp.props.placeholder) { + let e = elem(comp); + let textbox = e ? e.querySelector('[placeholder]') : null; + + let tl = textbox ? getTranslation(textbox.placeholder) : undefined; + if (tl !== undefined) { + textbox.placeholder = tl; + } + } + } +} + +function dumpTranslations() { + if (!hasLocalization()) { + // If we don't have any localization, + // we will not have traversed the app to find + // original_lines, so do that now. + localizeWholePage(); + } + var dumped = {}; + if (localization.rtl) { + dumped.rtl = true; + } + + for (const text in original_lines) { + if (dumped[text] !== undefined) continue; + dumped[text] = localization[text] || text; + } + + return dumped; +} + +function download_localization() { + var text = JSON.stringify(dumpTranslations(), null, 4); + + var element = document.createElement('a'); + element.setAttribute('href', 'data:text/plain;charset=utf-8,' + encodeURIComponent(text)); + element.setAttribute('download', "localization.json"); + element.style.display = 'none'; + document.body.appendChild(element); + + element.click(); + + document.body.removeChild(element); +} + +document.addEventListener("DOMContentLoaded", function() { + if (!hasLocalization()) { + return; + } + + onUiUpdate(function(m) { + m.forEach(function(mutation) { + mutation.addedNodes.forEach(function(node) { + processNode(node); + }); + }); + }); + + localizeWholePage(); + + if (localization.rtl) { // if the language is from right to left, + (new MutationObserver((mutations, observer) => { // wait for the style to load + mutations.forEach(mutation => { + mutation.addedNodes.forEach(node => { + if (node.tagName === 'STYLE') { + observer.disconnect(); + + for (const x of node.sheet.rules) { // find all rtl media rules + if (Array.from(x.media || []).includes('rtl')) { + x.media.appendMedium('all'); // enable them + } + } + } + }); + }); + })).observe(gradioApp(), {childList: true}); + } +}); diff --git a/stable-diffusion-webui/javascript/notification.js b/stable-diffusion-webui/javascript/notification.js new file mode 100755 index 0000000..3ee972a --- /dev/null +++ b/stable-diffusion-webui/javascript/notification.js @@ -0,0 +1,53 @@ +// Monitors the gallery and sends a browser notification when the leading image is new. + +let lastHeadImg = null; + +let notificationButton = null; + +onAfterUiUpdate(function() { + if (notificationButton == null) { + notificationButton = gradioApp().getElementById('request_notifications'); + + if (notificationButton != null) { + notificationButton.addEventListener('click', () => { + void Notification.requestPermission(); + }, true); + } + } + + const galleryPreviews = gradioApp().querySelectorAll('div[id^="tab_"] div[id$="_results"] .thumbnail-item > img'); + + if (galleryPreviews == null) return; + + const headImg = galleryPreviews[0]?.src; + + if (headImg == null || headImg == lastHeadImg) return; + + lastHeadImg = headImg; + + // play notification sound if available + const notificationAudio = gradioApp().querySelector('#audio_notification audio'); + if (notificationAudio) { + notificationAudio.volume = opts.notification_volume / 100.0 || 1.0; + notificationAudio.play(); + } + + if (document.hasFocus()) return; + + // Multiple copies of the images are in the DOM when one is selected. Dedup with a Set to get the real number generated. + const imgs = new Set(Array.from(galleryPreviews).map(img => img.src)); + + const notification = new Notification( + 'Stable Diffusion', + { + body: `Generated ${imgs.size > 1 ? imgs.size - opts.return_grid : 1} image${imgs.size > 1 ? 's' : ''}`, + icon: headImg, + image: headImg, + } + ); + + notification.onclick = function(_) { + parent.focus(); + this.close(); + }; +}); diff --git a/stable-diffusion-webui/javascript/profilerVisualization.js b/stable-diffusion-webui/javascript/profilerVisualization.js new file mode 100755 index 0000000..9822f4b --- /dev/null +++ b/stable-diffusion-webui/javascript/profilerVisualization.js @@ -0,0 +1,174 @@ + +function createRow(table, cellName, items) { + var tr = document.createElement('tr'); + var res = []; + + items.forEach(function(x, i) { + if (x === undefined) { + res.push(null); + return; + } + + var td = document.createElement(cellName); + td.textContent = x; + tr.appendChild(td); + res.push(td); + + var colspan = 1; + for (var n = i + 1; n < items.length; n++) { + if (items[n] !== undefined) { + break; + } + + colspan += 1; + } + + if (colspan > 1) { + td.colSpan = colspan; + } + }); + + table.appendChild(tr); + + return res; +} + +function createVisualizationTable(data, cutoff = 0, sort = "") { + var table = document.createElement('table'); + table.className = 'popup-table'; + + var keys = Object.keys(data); + if (sort === "number") { + keys = keys.sort(function(a, b) { + return data[b] - data[a]; + }); + } else { + keys = keys.sort(); + } + var items = keys.map(function(x) { + return {key: x, parts: x.split('/'), value: data[x]}; + }); + var maxLength = items.reduce(function(a, b) { + return Math.max(a, b.parts.length); + }, 0); + + var cols = createRow( + table, + 'th', + [ + cutoff === 0 ? 'key' : 'record', + cutoff === 0 ? 'value' : 'seconds' + ] + ); + cols[0].colSpan = maxLength; + + function arraysEqual(a, b) { + return !(a < b || b < a); + } + + var addLevel = function(level, parent, hide) { + var matching = items.filter(function(x) { + return x.parts[level] && !x.parts[level + 1] && arraysEqual(x.parts.slice(0, level), parent); + }); + if (sort === "number") { + matching = matching.sort(function(a, b) { + return b.value - a.value; + }); + } else { + matching = matching.sort(); + } + var othersTime = 0; + var othersList = []; + var othersRows = []; + var childrenRows = []; + matching.forEach(function(x) { + var visible = (cutoff === 0 && !hide) || (x.value >= cutoff && !hide); + + var cells = []; + for (var i = 0; i < maxLength; i++) { + cells.push(x.parts[i]); + } + cells.push(cutoff === 0 ? x.value : x.value.toFixed(3)); + var cols = createRow(table, 'td', cells); + for (i = 0; i < level; i++) { + cols[i].className = 'muted'; + } + + var tr = cols[0].parentNode; + if (!visible) { + tr.classList.add("hidden"); + } + + if (cutoff === 0 || x.value >= cutoff) { + childrenRows.push(tr); + } else { + othersTime += x.value; + othersList.push(x.parts[level]); + othersRows.push(tr); + } + + var children = addLevel(level + 1, parent.concat([x.parts[level]]), true); + if (children.length > 0) { + var cell = cols[level]; + var onclick = function() { + cell.classList.remove("link"); + cell.removeEventListener("click", onclick); + children.forEach(function(x) { + x.classList.remove("hidden"); + }); + }; + cell.classList.add("link"); + cell.addEventListener("click", onclick); + } + }); + + if (othersTime > 0) { + var cells = []; + for (var i = 0; i < maxLength; i++) { + cells.push(parent[i]); + } + cells.push(othersTime.toFixed(3)); + cells[level] = 'others'; + var cols = createRow(table, 'td', cells); + for (i = 0; i < level; i++) { + cols[i].className = 'muted'; + } + + var cell = cols[level]; + var tr = cell.parentNode; + var onclick = function() { + tr.classList.add("hidden"); + cell.classList.remove("link"); + cell.removeEventListener("click", onclick); + othersRows.forEach(function(x) { + x.classList.remove("hidden"); + }); + }; + + cell.title = othersList.join(", "); + cell.classList.add("link"); + cell.addEventListener("click", onclick); + + if (hide) { + tr.classList.add("hidden"); + } + + childrenRows.push(tr); + } + + return childrenRows; + }; + + addLevel(0, []); + + return table; +} + +function showProfile(path, cutoff = 0.05) { + requestGet(path, {}, function(data) { + data.records['total'] = data.total; + const table = createVisualizationTable(data.records, cutoff, "number"); + popup(table); + }); +} + diff --git a/stable-diffusion-webui/javascript/progressbar.js b/stable-diffusion-webui/javascript/progressbar.js new file mode 100755 index 0000000..23dea64 --- /dev/null +++ b/stable-diffusion-webui/javascript/progressbar.js @@ -0,0 +1,215 @@ +// code related to showing and updating progressbar shown as the image is being made + +function rememberGallerySelection() { + +} + +function getGallerySelectedIndex() { + +} + +function request(url, data, handler, errorHandler) { + var xhr = new XMLHttpRequest(); + xhr.open("POST", url, true); + xhr.setRequestHeader("Content-Type", "application/json"); + xhr.onreadystatechange = function() { + if (xhr.readyState === 4) { + if (xhr.status === 200) { + try { + var js = JSON.parse(xhr.responseText); + handler(js); + } catch (error) { + console.error(error); + errorHandler(); + } + } else { + errorHandler(); + } + } + }; + var js = JSON.stringify(data); + xhr.send(js); +} + +function pad2(x) { + return x < 10 ? '0' + x : x; +} + +function formatTime(secs) { + if (secs > 3600) { + return pad2(Math.floor(secs / 60 / 60)) + ":" + pad2(Math.floor(secs / 60) % 60) + ":" + pad2(Math.floor(secs) % 60); + } else if (secs > 60) { + return pad2(Math.floor(secs / 60)) + ":" + pad2(Math.floor(secs) % 60); + } else { + return Math.floor(secs) + "s"; + } +} + + +var originalAppTitle = undefined; + +onUiLoaded(function() { + originalAppTitle = document.title; +}); + +function setTitle(progress) { + var title = originalAppTitle; + + if (opts.show_progress_in_title && progress) { + title = '[' + progress.trim() + '] ' + title; + } + + if (document.title != title) { + document.title = title; + } +} + + +function randomId() { + return "task(" + Math.random().toString(36).slice(2, 7) + Math.random().toString(36).slice(2, 7) + Math.random().toString(36).slice(2, 7) + ")"; +} + +// starts sending progress requests to "/internal/progress" uri, creating progressbar above progressbarContainer element and +// preview inside gallery element. Cleans up all created stuff when the task is over and calls atEnd. +// calls onProgress every time there is a progress update +function requestProgress(id_task, progressbarContainer, gallery, atEnd, onProgress, inactivityTimeout = 40) { + var dateStart = new Date(); + var wasEverActive = false; + var parentProgressbar = progressbarContainer.parentNode; + var wakeLock = null; + + var requestWakeLock = async function() { + if (!opts.prevent_screen_sleep_during_generation || wakeLock) return; + try { + wakeLock = await navigator.wakeLock.request('screen'); + } catch (err) { + console.error('Wake Lock is not supported.'); + } + }; + + var releaseWakeLock = async function() { + if (!opts.prevent_screen_sleep_during_generation || !wakeLock) return; + try { + await wakeLock.release(); + wakeLock = null; + } catch (err) { + console.error('Wake Lock release failed', err); + } + }; + + var divProgress = document.createElement('div'); + divProgress.className = 'progressDiv'; + divProgress.style.display = opts.show_progressbar ? "block" : "none"; + var divInner = document.createElement('div'); + divInner.className = 'progress'; + + divProgress.appendChild(divInner); + parentProgressbar.insertBefore(divProgress, progressbarContainer); + + var livePreview = null; + + var removeProgressBar = function() { + releaseWakeLock(); + if (!divProgress) return; + + setTitle(""); + parentProgressbar.removeChild(divProgress); + if (gallery && livePreview) gallery.removeChild(livePreview); + atEnd(); + + divProgress = null; + }; + + var funProgress = function(id_task) { + requestWakeLock(); + request("./internal/progress", {id_task: id_task, live_preview: false}, function(res) { + if (res.completed) { + removeProgressBar(); + return; + } + + let progressText = ""; + + divInner.style.width = ((res.progress || 0) * 100.0) + '%'; + divInner.style.background = res.progress ? "" : "transparent"; + + if (res.progress > 0) { + progressText = ((res.progress || 0) * 100.0).toFixed(0) + '%'; + } + + if (res.eta) { + progressText += " ETA: " + formatTime(res.eta); + } + + setTitle(progressText); + + if (res.textinfo && res.textinfo.indexOf("\n") == -1) { + progressText = res.textinfo + " " + progressText; + } + + divInner.textContent = progressText; + + var elapsedFromStart = (new Date() - dateStart) / 1000; + + if (res.active) wasEverActive = true; + + if (!res.active && wasEverActive) { + removeProgressBar(); + return; + } + + if (elapsedFromStart > inactivityTimeout && !res.queued && !res.active) { + removeProgressBar(); + return; + } + + if (onProgress) { + onProgress(res); + } + + setTimeout(() => { + funProgress(id_task, res.id_live_preview); + }, opts.live_preview_refresh_period || 500); + }, function() { + removeProgressBar(); + }); + }; + + var funLivePreview = function(id_task, id_live_preview) { + request("./internal/progress", {id_task: id_task, id_live_preview: id_live_preview}, function(res) { + if (!divProgress) { + return; + } + + if (res.live_preview && gallery) { + var img = new Image(); + img.onload = function() { + if (!livePreview) { + livePreview = document.createElement('div'); + livePreview.className = 'livePreview'; + gallery.insertBefore(livePreview, gallery.firstElementChild); + } + + livePreview.appendChild(img); + if (livePreview.childElementCount > 2) { + livePreview.removeChild(livePreview.firstElementChild); + } + }; + img.src = res.live_preview; + } + + setTimeout(() => { + funLivePreview(id_task, res.id_live_preview); + }, opts.live_preview_refresh_period || 500); + }, function() { + removeProgressBar(); + }); + }; + + funProgress(id_task, 0); + + if (gallery) { + funLivePreview(id_task, 0); + } + +} diff --git a/stable-diffusion-webui/javascript/resizeHandle.js b/stable-diffusion-webui/javascript/resizeHandle.js new file mode 100755 index 0000000..4aeb14b --- /dev/null +++ b/stable-diffusion-webui/javascript/resizeHandle.js @@ -0,0 +1,205 @@ +(function() { + const GRADIO_MIN_WIDTH = 320; + const PAD = 16; + const DEBOUNCE_TIME = 100; + const DOUBLE_TAP_DELAY = 200; //ms + + const R = { + tracking: false, + parent: null, + parentWidth: null, + leftCol: null, + leftColStartWidth: null, + screenX: null, + lastTapTime: null, + }; + + let resizeTimer; + let parents = []; + + function setLeftColGridTemplate(el, width) { + el.style.gridTemplateColumns = `${width}px 16px 1fr`; + } + + function displayResizeHandle(parent) { + if (!parent.needHideOnMoblie) { + return true; + } + if (window.innerWidth < GRADIO_MIN_WIDTH * 2 + PAD * 4) { + parent.style.display = 'flex'; + parent.resizeHandle.style.display = "none"; + return false; + } else { + parent.style.display = 'grid'; + parent.resizeHandle.style.display = "block"; + return true; + } + } + + function afterResize(parent) { + if (displayResizeHandle(parent) && parent.style.gridTemplateColumns != parent.style.originalGridTemplateColumns) { + const oldParentWidth = R.parentWidth; + const newParentWidth = parent.offsetWidth; + const widthL = parseInt(parent.style.gridTemplateColumns.split(' ')[0]); + + const ratio = newParentWidth / oldParentWidth; + + const newWidthL = Math.max(Math.floor(ratio * widthL), parent.minLeftColWidth); + setLeftColGridTemplate(parent, newWidthL); + + R.parentWidth = newParentWidth; + } + } + + function setup(parent) { + + function onDoubleClick(evt) { + evt.preventDefault(); + evt.stopPropagation(); + + parent.style.gridTemplateColumns = parent.style.originalGridTemplateColumns; + } + + const leftCol = parent.firstElementChild; + const rightCol = parent.lastElementChild; + + parents.push(parent); + + parent.style.display = 'grid'; + parent.style.gap = '0'; + let leftColTemplate = ""; + if (parent.children[0].style.flexGrow) { + leftColTemplate = `${parent.children[0].style.flexGrow}fr`; + parent.minLeftColWidth = GRADIO_MIN_WIDTH; + parent.minRightColWidth = GRADIO_MIN_WIDTH; + parent.needHideOnMoblie = true; + } else { + leftColTemplate = parent.children[0].style.flexBasis; + parent.minLeftColWidth = parent.children[0].style.flexBasis.slice(0, -2) / 2; + parent.minRightColWidth = 0; + parent.needHideOnMoblie = false; + } + + if (!leftColTemplate) { + leftColTemplate = '1fr'; + } + + const gridTemplateColumns = `${leftColTemplate} ${PAD}px ${parent.children[1].style.flexGrow}fr`; + parent.style.gridTemplateColumns = gridTemplateColumns; + parent.style.originalGridTemplateColumns = gridTemplateColumns; + + const resizeHandle = document.createElement('div'); + resizeHandle.classList.add('resize-handle'); + parent.insertBefore(resizeHandle, rightCol); + parent.resizeHandle = resizeHandle; + + ['mousedown', 'touchstart'].forEach((eventType) => { + resizeHandle.addEventListener(eventType, (evt) => { + if (eventType.startsWith('mouse')) { + if (evt.button !== 0) return; + } else { + if (evt.changedTouches.length !== 1) return; + + const currentTime = new Date().getTime(); + if (R.lastTapTime && currentTime - R.lastTapTime <= DOUBLE_TAP_DELAY) { + onDoubleClick(evt); + return; + } + + R.lastTapTime = currentTime; + } + + evt.preventDefault(); + evt.stopPropagation(); + + document.body.classList.add('resizing'); + + R.tracking = true; + R.parent = parent; + R.parentWidth = parent.offsetWidth; + R.leftCol = leftCol; + R.leftColStartWidth = leftCol.offsetWidth; + if (eventType.startsWith('mouse')) { + R.screenX = evt.screenX; + } else { + R.screenX = evt.changedTouches[0].screenX; + } + }); + }); + + resizeHandle.addEventListener('dblclick', onDoubleClick); + + afterResize(parent); + } + + ['mousemove', 'touchmove'].forEach((eventType) => { + window.addEventListener(eventType, (evt) => { + if (eventType.startsWith('mouse')) { + if (evt.button !== 0) return; + } else { + if (evt.changedTouches.length !== 1) return; + } + + if (R.tracking) { + if (eventType.startsWith('mouse')) { + evt.preventDefault(); + } + evt.stopPropagation(); + + let delta = 0; + if (eventType.startsWith('mouse')) { + delta = R.screenX - evt.screenX; + } else { + delta = R.screenX - evt.changedTouches[0].screenX; + } + const leftColWidth = Math.max(Math.min(R.leftColStartWidth - delta, R.parent.offsetWidth - R.parent.minRightColWidth - PAD), R.parent.minLeftColWidth); + setLeftColGridTemplate(R.parent, leftColWidth); + } + }); + }); + + ['mouseup', 'touchend'].forEach((eventType) => { + window.addEventListener(eventType, (evt) => { + if (eventType.startsWith('mouse')) { + if (evt.button !== 0) return; + } else { + if (evt.changedTouches.length !== 1) return; + } + + if (R.tracking) { + evt.preventDefault(); + evt.stopPropagation(); + + R.tracking = false; + + document.body.classList.remove('resizing'); + } + }); + }); + + + window.addEventListener('resize', () => { + clearTimeout(resizeTimer); + + resizeTimer = setTimeout(function() { + for (const parent of parents) { + afterResize(parent); + } + }, DEBOUNCE_TIME); + }); + + setupResizeHandle = setup; +})(); + + +function setupAllResizeHandles() { + for (var elem of gradioApp().querySelectorAll('.resize-handle-row')) { + if (!elem.querySelector('.resize-handle') && !elem.children[0].classList.contains("hidden")) { + setupResizeHandle(elem); + } + } +} + + +onUiLoaded(setupAllResizeHandles); + diff --git a/stable-diffusion-webui/javascript/settings.js b/stable-diffusion-webui/javascript/settings.js new file mode 100755 index 0000000..b2d981c --- /dev/null +++ b/stable-diffusion-webui/javascript/settings.js @@ -0,0 +1,71 @@ +let settingsExcludeTabsFromShowAll = { + settings_tab_defaults: 1, + settings_tab_sysinfo: 1, + settings_tab_actions: 1, + settings_tab_licenses: 1, +}; + +function settingsShowAllTabs() { + gradioApp().querySelectorAll('#settings > div').forEach(function(elem) { + if (settingsExcludeTabsFromShowAll[elem.id]) return; + + elem.style.display = "block"; + }); +} + +function settingsShowOneTab() { + gradioApp().querySelector('#settings_show_one_page').click(); +} + +onUiLoaded(function() { + var edit = gradioApp().querySelector('#settings_search'); + var editTextarea = gradioApp().querySelector('#settings_search > label > input'); + var buttonShowAllPages = gradioApp().getElementById('settings_show_all_pages'); + var settings_tabs = gradioApp().querySelector('#settings div'); + + onEdit('settingsSearch', editTextarea, 250, function() { + var searchText = (editTextarea.value || "").trim().toLowerCase(); + + gradioApp().querySelectorAll('#settings > div[id^=settings_] div[id^=column_settings_] > *').forEach(function(elem) { + var visible = elem.textContent.trim().toLowerCase().indexOf(searchText) != -1; + elem.style.display = visible ? "" : "none"; + }); + + if (searchText != "") { + settingsShowAllTabs(); + } else { + settingsShowOneTab(); + } + }); + + settings_tabs.insertBefore(edit, settings_tabs.firstChild); + settings_tabs.appendChild(buttonShowAllPages); + + + buttonShowAllPages.addEventListener("click", settingsShowAllTabs); +}); + + +onOptionsChanged(function() { + if (gradioApp().querySelector('#settings .settings-category')) return; + + var sectionMap = {}; + gradioApp().querySelectorAll('#settings > div > button').forEach(function(x) { + sectionMap[x.textContent.trim()] = x; + }); + + opts._categories.forEach(function(x) { + var section = localization[x[0]] ?? x[0]; + var category = localization[x[1]] ?? x[1]; + + var span = document.createElement('SPAN'); + span.textContent = category; + span.className = 'settings-category'; + + var sectionElem = sectionMap[section]; + if (!sectionElem) return; + + sectionElem.parentElement.insertBefore(span, sectionElem); + }); +}); + diff --git a/stable-diffusion-webui/javascript/textualInversion.js b/stable-diffusion-webui/javascript/textualInversion.js new file mode 100755 index 0000000..20443fc --- /dev/null +++ b/stable-diffusion-webui/javascript/textualInversion.js @@ -0,0 +1,17 @@ + + + +function start_training_textual_inversion() { + gradioApp().querySelector('#ti_error').innerHTML = ''; + + var id = randomId(); + requestProgress(id, gradioApp().getElementById('ti_output'), gradioApp().getElementById('ti_gallery'), function() {}, function(progress) { + gradioApp().getElementById('ti_progress').innerHTML = progress.textinfo; + }); + + var res = Array.from(arguments); + + res[0] = id; + + return res; +} diff --git a/stable-diffusion-webui/javascript/token-counters.js b/stable-diffusion-webui/javascript/token-counters.js new file mode 100755 index 0000000..eeea7a5 --- /dev/null +++ b/stable-diffusion-webui/javascript/token-counters.js @@ -0,0 +1,87 @@ +let promptTokenCountUpdateFunctions = {}; + +function update_txt2img_tokens(...args) { + // Called from Gradio + update_token_counter("txt2img_token_button"); + update_token_counter("txt2img_negative_token_button"); + if (args.length == 2) { + return args[0]; + } + return args; +} + +function update_img2img_tokens(...args) { + // Called from Gradio + update_token_counter("img2img_token_button"); + update_token_counter("img2img_negative_token_button"); + if (args.length == 2) { + return args[0]; + } + return args; +} + +function update_token_counter(button_id) { + promptTokenCountUpdateFunctions[button_id]?.(); +} + + +function recalculatePromptTokens(name) { + promptTokenCountUpdateFunctions[name]?.(); +} + +function recalculate_prompts_txt2img() { + // Called from Gradio + recalculatePromptTokens('txt2img_prompt'); + recalculatePromptTokens('txt2img_neg_prompt'); + return Array.from(arguments); +} + +function recalculate_prompts_img2img() { + // Called from Gradio + recalculatePromptTokens('img2img_prompt'); + recalculatePromptTokens('img2img_neg_prompt'); + return Array.from(arguments); +} + +function setupTokenCounting(id, id_counter, id_button) { + var prompt = gradioApp().getElementById(id); + var counter = gradioApp().getElementById(id_counter); + var textarea = gradioApp().querySelector(`#${id} > label > textarea`); + + if (counter.parentElement == prompt.parentElement) { + return; + } + + prompt.parentElement.insertBefore(counter, prompt); + prompt.parentElement.style.position = "relative"; + + var func = onEdit(id, textarea, 800, function() { + if (counter.classList.contains("token-counter-visible")) { + gradioApp().getElementById(id_button)?.click(); + } + }); + promptTokenCountUpdateFunctions[id] = func; + promptTokenCountUpdateFunctions[id_button] = func; +} + +function toggleTokenCountingVisibility(id, id_counter, id_button) { + var counter = gradioApp().getElementById(id_counter); + + counter.style.display = opts.disable_token_counters ? "none" : "block"; + counter.classList.toggle("token-counter-visible", !opts.disable_token_counters); +} + +function runCodeForTokenCounters(fun) { + fun('txt2img_prompt', 'txt2img_token_counter', 'txt2img_token_button'); + fun('txt2img_neg_prompt', 'txt2img_negative_token_counter', 'txt2img_negative_token_button'); + fun('img2img_prompt', 'img2img_token_counter', 'img2img_token_button'); + fun('img2img_neg_prompt', 'img2img_negative_token_counter', 'img2img_negative_token_button'); +} + +onUiLoaded(function() { + runCodeForTokenCounters(setupTokenCounting); +}); + +onOptionsChanged(function() { + runCodeForTokenCounters(toggleTokenCountingVisibility); +}); diff --git a/stable-diffusion-webui/javascript/ui.js b/stable-diffusion-webui/javascript/ui.js new file mode 100755 index 0000000..2030963 --- /dev/null +++ b/stable-diffusion-webui/javascript/ui.js @@ -0,0 +1,436 @@ +// various functions for interaction with ui.py not large enough to warrant putting them in separate files + +function set_theme(theme) { + var gradioURL = window.location.href; + if (!gradioURL.includes('?__theme=')) { + window.location.replace(gradioURL + '?__theme=' + theme); + } +} + +function all_gallery_buttons() { + var allGalleryButtons = gradioApp().querySelectorAll('[style="display: block;"].tabitem div[id$=_gallery].gradio-gallery .thumbnails > .thumbnail-item.thumbnail-small'); + var visibleGalleryButtons = []; + allGalleryButtons.forEach(function(elem) { + if (elem.parentElement.offsetParent) { + visibleGalleryButtons.push(elem); + } + }); + return visibleGalleryButtons; +} + +function selected_gallery_button() { + return all_gallery_buttons().find(elem => elem.classList.contains('selected')) ?? null; +} + +function selected_gallery_index() { + return all_gallery_buttons().findIndex(elem => elem.classList.contains('selected')); +} + +function gallery_container_buttons(gallery_container) { + return gradioApp().querySelectorAll(`#${gallery_container} .thumbnail-item.thumbnail-small`); +} + +function selected_gallery_index_id(gallery_container) { + return Array.from(gallery_container_buttons(gallery_container)).findIndex(elem => elem.classList.contains('selected')); +} + +function extract_image_from_gallery(gallery) { + if (gallery.length == 0) { + return [null]; + } + if (gallery.length == 1) { + return [gallery[0]]; + } + + var index = selected_gallery_index(); + + if (index < 0 || index >= gallery.length) { + // Use the first image in the gallery as the default + index = 0; + } + + return [gallery[index]]; +} + +window.args_to_array = Array.from; // Compatibility with e.g. extensions that may expect this to be around + +function switch_to_txt2img() { + gradioApp().querySelector('#tabs').querySelectorAll('button')[0].click(); + + return Array.from(arguments); +} + +function switch_to_img2img_tab(no) { + gradioApp().querySelector('#tabs').querySelectorAll('button')[1].click(); + gradioApp().getElementById('mode_img2img').querySelectorAll('button')[no].click(); +} +function switch_to_img2img() { + switch_to_img2img_tab(0); + return Array.from(arguments); +} + +function switch_to_sketch() { + switch_to_img2img_tab(1); + return Array.from(arguments); +} + +function switch_to_inpaint() { + switch_to_img2img_tab(2); + return Array.from(arguments); +} + +function switch_to_inpaint_sketch() { + switch_to_img2img_tab(3); + return Array.from(arguments); +} + +function switch_to_extras() { + gradioApp().querySelector('#tabs').querySelectorAll('button')[2].click(); + + return Array.from(arguments); +} + +function get_tab_index(tabId) { + let buttons = gradioApp().getElementById(tabId).querySelector('div').querySelectorAll('button'); + for (let i = 0; i < buttons.length; i++) { + if (buttons[i].classList.contains('selected')) { + return i; + } + } + return 0; +} + +function create_tab_index_args(tabId, args) { + var res = Array.from(args); + res[0] = get_tab_index(tabId); + return res; +} + +function get_img2img_tab_index() { + let res = Array.from(arguments); + res.splice(-2); + res[0] = get_tab_index('mode_img2img'); + return res; +} + +function create_submit_args(args) { + var res = Array.from(args); + + // As it is currently, txt2img and img2img send back the previous output args (txt2img_gallery, generation_info, html_info) whenever you generate a new image. + // This can lead to uploading a huge gallery of previously generated images, which leads to an unnecessary delay between submitting and beginning to generate. + // I don't know why gradio is sending outputs along with inputs, but we can prevent sending the image gallery here, which seems to be an issue for some. + // If gradio at some point stops sending outputs, this may break something + if (Array.isArray(res[res.length - 3])) { + res[res.length - 3] = null; + } + + return res; +} + +function setSubmitButtonsVisibility(tabname, showInterrupt, showSkip, showInterrupting) { + gradioApp().getElementById(tabname + '_interrupt').style.display = showInterrupt ? "block" : "none"; + gradioApp().getElementById(tabname + '_skip').style.display = showSkip ? "block" : "none"; + gradioApp().getElementById(tabname + '_interrupting').style.display = showInterrupting ? "block" : "none"; +} + +function showSubmitButtons(tabname, show) { + setSubmitButtonsVisibility(tabname, !show, !show, false); +} + +function showSubmitInterruptingPlaceholder(tabname) { + setSubmitButtonsVisibility(tabname, false, true, true); +} + +function showRestoreProgressButton(tabname, show) { + var button = gradioApp().getElementById(tabname + "_restore_progress"); + if (!button) return; + button.style.setProperty('display', show ? 'flex' : 'none', 'important'); +} + +function submit() { + showSubmitButtons('txt2img', false); + + var id = randomId(); + localSet("txt2img_task_id", id); + + requestProgress(id, gradioApp().getElementById('txt2img_gallery_container'), gradioApp().getElementById('txt2img_gallery'), function() { + showSubmitButtons('txt2img', true); + localRemove("txt2img_task_id"); + showRestoreProgressButton('txt2img', false); + }); + + var res = create_submit_args(arguments); + + res[0] = id; + + return res; +} + +function submit_txt2img_upscale() { + var res = submit(...arguments); + + res[2] = selected_gallery_index(); + + return res; +} + +function submit_img2img() { + showSubmitButtons('img2img', false); + + var id = randomId(); + localSet("img2img_task_id", id); + + requestProgress(id, gradioApp().getElementById('img2img_gallery_container'), gradioApp().getElementById('img2img_gallery'), function() { + showSubmitButtons('img2img', true); + localRemove("img2img_task_id"); + showRestoreProgressButton('img2img', false); + }); + + var res = create_submit_args(arguments); + + res[0] = id; + res[1] = get_tab_index('mode_img2img'); + + return res; +} + +function submit_extras() { + showSubmitButtons('extras', false); + + var id = randomId(); + + requestProgress(id, gradioApp().getElementById('extras_gallery_container'), gradioApp().getElementById('extras_gallery'), function() { + showSubmitButtons('extras', true); + }); + + var res = create_submit_args(arguments); + + res[0] = id; + + console.log(res); + return res; +} + +function restoreProgressTxt2img() { + showRestoreProgressButton("txt2img", false); + var id = localGet("txt2img_task_id"); + + if (id) { + showSubmitInterruptingPlaceholder('txt2img'); + requestProgress(id, gradioApp().getElementById('txt2img_gallery_container'), gradioApp().getElementById('txt2img_gallery'), function() { + showSubmitButtons('txt2img', true); + }, null, 0); + } + + return id; +} + +function restoreProgressImg2img() { + showRestoreProgressButton("img2img", false); + + var id = localGet("img2img_task_id"); + + if (id) { + showSubmitInterruptingPlaceholder('img2img'); + requestProgress(id, gradioApp().getElementById('img2img_gallery_container'), gradioApp().getElementById('img2img_gallery'), function() { + showSubmitButtons('img2img', true); + }, null, 0); + } + + return id; +} + + +/** + * Configure the width and height elements on `tabname` to accept + * pasting of resolutions in the form of "width x height". + */ +function setupResolutionPasting(tabname) { + var width = gradioApp().querySelector(`#${tabname}_width input[type=number]`); + var height = gradioApp().querySelector(`#${tabname}_height input[type=number]`); + for (const el of [width, height]) { + el.addEventListener('paste', function(event) { + var pasteData = event.clipboardData.getData('text/plain'); + var parsed = pasteData.match(/^\s*(\d+)\D+(\d+)\s*$/); + if (parsed) { + width.value = parsed[1]; + height.value = parsed[2]; + updateInput(width); + updateInput(height); + event.preventDefault(); + } + }); + } +} + +onUiLoaded(function() { + showRestoreProgressButton('txt2img', localGet("txt2img_task_id")); + showRestoreProgressButton('img2img', localGet("img2img_task_id")); + setupResolutionPasting('txt2img'); + setupResolutionPasting('img2img'); +}); + + +function modelmerger() { + var id = randomId(); + requestProgress(id, gradioApp().getElementById('modelmerger_results_panel'), null, function() {}); + + var res = create_submit_args(arguments); + res[0] = id; + return res; +} + + +function ask_for_style_name(_, prompt_text, negative_prompt_text) { + var name_ = prompt('Style name:'); + return [name_, prompt_text, negative_prompt_text]; +} + +function confirm_clear_prompt(prompt, negative_prompt) { + if (confirm("Delete prompt?")) { + prompt = ""; + negative_prompt = ""; + } + + return [prompt, negative_prompt]; +} + + +var opts = {}; +onAfterUiUpdate(function() { + if (Object.keys(opts).length != 0) return; + + var json_elem = gradioApp().getElementById('settings_json'); + if (json_elem == null) return; + + var textarea = json_elem.querySelector('textarea'); + var jsdata = textarea.value; + opts = JSON.parse(jsdata); + + executeCallbacks(optionsAvailableCallbacks); /*global optionsAvailableCallbacks*/ + executeCallbacks(optionsChangedCallbacks); /*global optionsChangedCallbacks*/ + + Object.defineProperty(textarea, 'value', { + set: function(newValue) { + var valueProp = Object.getOwnPropertyDescriptor(HTMLTextAreaElement.prototype, 'value'); + var oldValue = valueProp.get.call(textarea); + valueProp.set.call(textarea, newValue); + + if (oldValue != newValue) { + opts = JSON.parse(textarea.value); + } + + executeCallbacks(optionsChangedCallbacks); + }, + get: function() { + var valueProp = Object.getOwnPropertyDescriptor(HTMLTextAreaElement.prototype, 'value'); + return valueProp.get.call(textarea); + } + }); + + json_elem.parentElement.style.display = "none"; +}); + +onOptionsChanged(function() { + var elem = gradioApp().getElementById('sd_checkpoint_hash'); + var sd_checkpoint_hash = opts.sd_checkpoint_hash || ""; + var shorthash = sd_checkpoint_hash.substring(0, 10); + + if (elem && elem.textContent != shorthash) { + elem.textContent = shorthash; + elem.title = sd_checkpoint_hash; + elem.href = "https://google.com/search?q=" + sd_checkpoint_hash; + } +}); + +let txt2img_textarea, img2img_textarea = undefined; + +function restart_reload() { + document.body.style.backgroundColor = "var(--background-fill-primary)"; + document.body.innerHTML = '

    Reloading...

    '; + var requestPing = function() { + requestGet("./internal/ping", {}, function(data) { + location.reload(); + }, function() { + setTimeout(requestPing, 500); + }); + }; + + setTimeout(requestPing, 2000); + + return []; +} + +// Simulate an `input` DOM event for Gradio Textbox component. Needed after you edit its contents in javascript, otherwise your edits +// will only visible on web page and not sent to python. +function updateInput(target) { + let e = new Event("input", {bubbles: true}); + Object.defineProperty(e, "target", {value: target}); + target.dispatchEvent(e); +} + + +var desiredCheckpointName = null; +function selectCheckpoint(name) { + desiredCheckpointName = name; + gradioApp().getElementById('change_checkpoint').click(); +} + +function currentImg2imgSourceResolution(w, h, scaleBy) { + var img = gradioApp().querySelector('#mode_img2img > div[style="display: block;"] img'); + return img ? [img.naturalWidth, img.naturalHeight, scaleBy] : [0, 0, scaleBy]; +} + +function updateImg2imgResizeToTextAfterChangingImage() { + // At the time this is called from gradio, the image has no yet been replaced. + // There may be a better solution, but this is simple and straightforward so I'm going with it. + + setTimeout(function() { + gradioApp().getElementById('img2img_update_resize_to').click(); + }, 500); + + return []; + +} + + + +function setRandomSeed(elem_id) { + var input = gradioApp().querySelector("#" + elem_id + " input"); + if (!input) return []; + + input.value = "-1"; + updateInput(input); + return []; +} + +function switchWidthHeight(tabname) { + var width = gradioApp().querySelector("#" + tabname + "_width input[type=number]"); + var height = gradioApp().querySelector("#" + tabname + "_height input[type=number]"); + if (!width || !height) return []; + + var tmp = width.value; + width.value = height.value; + height.value = tmp; + + updateInput(width); + updateInput(height); + return []; +} + + +var onEditTimers = {}; + +// calls func after afterMs milliseconds has passed since the input elem has been edited by user +function onEdit(editId, elem, afterMs, func) { + var edited = function() { + var existingTimer = onEditTimers[editId]; + if (existingTimer) clearTimeout(existingTimer); + + onEditTimers[editId] = setTimeout(func, afterMs); + }; + + elem.addEventListener("input", edited); + + return edited; +} diff --git a/stable-diffusion-webui/javascript/ui_settings_hints.js b/stable-diffusion-webui/javascript/ui_settings_hints.js new file mode 100755 index 0000000..d088f94 --- /dev/null +++ b/stable-diffusion-webui/javascript/ui_settings_hints.js @@ -0,0 +1,62 @@ +// various hints and extra info for the settings tab + +var settingsHintsSetup = false; + +onOptionsChanged(function() { + if (settingsHintsSetup) return; + settingsHintsSetup = true; + + gradioApp().querySelectorAll('#settings [id^=setting_]').forEach(function(div) { + var name = div.id.substr(8); + var commentBefore = opts._comments_before[name]; + var commentAfter = opts._comments_after[name]; + + if (!commentBefore && !commentAfter) return; + + var span = null; + if (div.classList.contains('gradio-checkbox')) span = div.querySelector('label span'); + else if (div.classList.contains('gradio-checkboxgroup')) span = div.querySelector('span').firstChild; + else if (div.classList.contains('gradio-radio')) span = div.querySelector('span').firstChild; + else span = div.querySelector('label span').firstChild; + + if (!span) return; + + if (commentBefore) { + var comment = document.createElement('DIV'); + comment.className = 'settings-comment'; + comment.innerHTML = commentBefore; + span.parentElement.insertBefore(document.createTextNode('\xa0'), span); + span.parentElement.insertBefore(comment, span); + span.parentElement.insertBefore(document.createTextNode('\xa0'), span); + } + if (commentAfter) { + comment = document.createElement('DIV'); + comment.className = 'settings-comment'; + comment.innerHTML = commentAfter; + span.parentElement.insertBefore(comment, span.nextSibling); + span.parentElement.insertBefore(document.createTextNode('\xa0'), span.nextSibling); + } + }); +}); + +function settingsHintsShowQuicksettings() { + requestGet("./internal/quicksettings-hint", {}, function(data) { + var table = document.createElement('table'); + table.className = 'popup-table'; + + data.forEach(function(obj) { + var tr = document.createElement('tr'); + var td = document.createElement('td'); + td.textContent = obj.name; + tr.appendChild(td); + + td = document.createElement('td'); + td.textContent = obj.label; + tr.appendChild(td); + + table.appendChild(tr); + }); + + popup(table); + }); +} diff --git a/stable-diffusion-webui/launch.py b/stable-diffusion-webui/launch.py new file mode 100755 index 0000000..cafab78 --- /dev/null +++ b/stable-diffusion-webui/launch.py @@ -0,0 +1,48 @@ +from modules import launch_utils + +args = launch_utils.args +python = launch_utils.python +git = launch_utils.git +index_url = launch_utils.index_url +dir_repos = launch_utils.dir_repos + +commit_hash = launch_utils.commit_hash +git_tag = launch_utils.git_tag + +run = launch_utils.run +is_installed = launch_utils.is_installed +repo_dir = launch_utils.repo_dir + +run_pip = launch_utils.run_pip +check_run_python = launch_utils.check_run_python +git_clone = launch_utils.git_clone +git_pull_recursive = launch_utils.git_pull_recursive +list_extensions = launch_utils.list_extensions +run_extension_installer = launch_utils.run_extension_installer +prepare_environment = launch_utils.prepare_environment +configure_for_tests = launch_utils.configure_for_tests +start = launch_utils.start + + +def main(): + if args.dump_sysinfo: + filename = launch_utils.dump_sysinfo() + + print(f"Sysinfo saved as {filename}. Exiting...") + + exit(0) + + launch_utils.startup_timer.record("initial startup") + + with launch_utils.startup_timer.subcategory("prepare environment"): + if not args.skip_prepare_environment: + prepare_environment() + + if args.test_server: + configure_for_tests() + + start() + + +if __name__ == "__main__": + main() diff --git a/stable-diffusion-webui/localizations/Put localization files here.txt b/stable-diffusion-webui/localizations/Put localization files here.txt new file mode 100755 index 0000000..e69de29 diff --git a/stable-diffusion-webui/modules/Roboto-Regular.ttf b/stable-diffusion-webui/modules/Roboto-Regular.ttf new file mode 100755 index 0000000..500b104 Binary files /dev/null and b/stable-diffusion-webui/modules/Roboto-Regular.ttf differ diff --git a/stable-diffusion-webui/modules/api/api.py b/stable-diffusion-webui/modules/api/api.py new file mode 100755 index 0000000..97ec751 --- /dev/null +++ b/stable-diffusion-webui/modules/api/api.py @@ -0,0 +1,928 @@ +import base64 +import io +import os +import time +import datetime +import uvicorn +import ipaddress +import requests +import gradio as gr +from threading import Lock +from io import BytesIO +from fastapi import APIRouter, Depends, FastAPI, Request, Response +from fastapi.security import HTTPBasic, HTTPBasicCredentials +from fastapi.exceptions import HTTPException +from fastapi.responses import JSONResponse +from fastapi.encoders import jsonable_encoder +from secrets import compare_digest + +import modules.shared as shared +from modules import sd_samplers, deepbooru, sd_hijack, images, scripts, ui, postprocessing, errors, restart, shared_items, script_callbacks, infotext_utils, sd_models, sd_schedulers +from modules.api import models +from modules.shared import opts +from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, process_images +from modules.textual_inversion.textual_inversion import create_embedding, train_embedding +from modules.hypernetworks.hypernetwork import create_hypernetwork, train_hypernetwork +from PIL import PngImagePlugin +from modules.sd_models_config import find_checkpoint_config_near_filename +from modules.realesrgan_model import get_realesrgan_models +from modules import devices +from typing import Any +import piexif +import piexif.helper +from contextlib import closing +from modules.progress import create_task_id, add_task_to_queue, start_task, finish_task, current_task + +def script_name_to_index(name, scripts): + try: + return [script.title().lower() for script in scripts].index(name.lower()) + except Exception as e: + raise HTTPException(status_code=422, detail=f"Script '{name}' not found") from e + + +def validate_sampler_name(name): + config = sd_samplers.all_samplers_map.get(name, None) + if config is None: + raise HTTPException(status_code=400, detail="Sampler not found") + + return name + + +def setUpscalers(req: dict): + reqDict = vars(req) + reqDict['extras_upscaler_1'] = reqDict.pop('upscaler_1', None) + reqDict['extras_upscaler_2'] = reqDict.pop('upscaler_2', None) + return reqDict + + +def verify_url(url): + """Returns True if the url refers to a global resource.""" + + import socket + from urllib.parse import urlparse + try: + parsed_url = urlparse(url) + domain_name = parsed_url.netloc + host = socket.gethostbyname_ex(domain_name) + for ip in host[2]: + ip_addr = ipaddress.ip_address(ip) + if not ip_addr.is_global: + return False + except Exception: + return False + + return True + + +def decode_base64_to_image(encoding): + if encoding.startswith("http://") or encoding.startswith("https://"): + if not opts.api_enable_requests: + raise HTTPException(status_code=500, detail="Requests not allowed") + + if opts.api_forbid_local_requests and not verify_url(encoding): + raise HTTPException(status_code=500, detail="Request to local resource not allowed") + + headers = {'user-agent': opts.api_useragent} if opts.api_useragent else {} + response = requests.get(encoding, timeout=30, headers=headers) + try: + image = images.read(BytesIO(response.content)) + return image + except Exception as e: + raise HTTPException(status_code=500, detail="Invalid image url") from e + + if encoding.startswith("data:image/"): + encoding = encoding.split(";")[1].split(",")[1] + try: + image = images.read(BytesIO(base64.b64decode(encoding))) + return image + except Exception as e: + raise HTTPException(status_code=500, detail="Invalid encoded image") from e + + +def encode_pil_to_base64(image): + with io.BytesIO() as output_bytes: + if isinstance(image, str): + return image + if opts.samples_format.lower() == 'png': + use_metadata = False + metadata = PngImagePlugin.PngInfo() + for key, value in image.info.items(): + if isinstance(key, str) and isinstance(value, str): + metadata.add_text(key, value) + use_metadata = True + image.save(output_bytes, format="PNG", pnginfo=(metadata if use_metadata else None), quality=opts.jpeg_quality) + + elif opts.samples_format.lower() in ("jpg", "jpeg", "webp"): + if image.mode in ("RGBA", "P"): + image = image.convert("RGB") + parameters = image.info.get('parameters', None) + exif_bytes = piexif.dump({ + "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(parameters or "", encoding="unicode") } + }) + if opts.samples_format.lower() in ("jpg", "jpeg"): + image.save(output_bytes, format="JPEG", exif = exif_bytes, quality=opts.jpeg_quality) + else: + image.save(output_bytes, format="WEBP", exif = exif_bytes, quality=opts.jpeg_quality) + + else: + raise HTTPException(status_code=500, detail="Invalid image format") + + bytes_data = output_bytes.getvalue() + + return base64.b64encode(bytes_data) + + +def api_middleware(app: FastAPI): + rich_available = False + try: + if os.environ.get('WEBUI_RICH_EXCEPTIONS', None) is not None: + import anyio # importing just so it can be placed on silent list + import starlette # importing just so it can be placed on silent list + from rich.console import Console + console = Console() + rich_available = True + except Exception: + pass + + @app.middleware("http") + async def log_and_time(req: Request, call_next): + ts = time.time() + res: Response = await call_next(req) + duration = str(round(time.time() - ts, 4)) + res.headers["X-Process-Time"] = duration + endpoint = req.scope.get('path', 'err') + if shared.cmd_opts.api_log and endpoint.startswith('/sdapi'): + print('API {t} {code} {prot}/{ver} {method} {endpoint} {cli} {duration}'.format( + t=datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"), + code=res.status_code, + ver=req.scope.get('http_version', '0.0'), + cli=req.scope.get('client', ('0:0.0.0', 0))[0], + prot=req.scope.get('scheme', 'err'), + method=req.scope.get('method', 'err'), + endpoint=endpoint, + duration=duration, + )) + return res + + def handle_exception(request: Request, e: Exception): + err = { + "error": type(e).__name__, + "detail": vars(e).get('detail', ''), + "body": vars(e).get('body', ''), + "errors": str(e), + } + if not isinstance(e, HTTPException): # do not print backtrace on known httpexceptions + message = f"API error: {request.method}: {request.url} {err}" + if rich_available: + print(message) + console.print_exception(show_locals=True, max_frames=2, extra_lines=1, suppress=[anyio, starlette], word_wrap=False, width=min([console.width, 200])) + else: + errors.report(message, exc_info=True) + return JSONResponse(status_code=vars(e).get('status_code', 500), content=jsonable_encoder(err)) + + @app.middleware("http") + async def exception_handling(request: Request, call_next): + try: + return await call_next(request) + except Exception as e: + return handle_exception(request, e) + + @app.exception_handler(Exception) + async def fastapi_exception_handler(request: Request, e: Exception): + return handle_exception(request, e) + + @app.exception_handler(HTTPException) + async def http_exception_handler(request: Request, e: HTTPException): + return handle_exception(request, e) + + +class Api: + def __init__(self, app: FastAPI, queue_lock: Lock): + if shared.cmd_opts.api_auth: + self.credentials = {} + for auth in shared.cmd_opts.api_auth.split(","): + user, password = auth.split(":") + self.credentials[user] = password + + self.router = APIRouter() + self.app = app + self.queue_lock = queue_lock + api_middleware(self.app) + self.add_api_route("/sdapi/v1/txt2img", self.text2imgapi, methods=["POST"], response_model=models.TextToImageResponse) + self.add_api_route("/sdapi/v1/img2img", self.img2imgapi, methods=["POST"], response_model=models.ImageToImageResponse) + self.add_api_route("/sdapi/v1/extra-single-image", self.extras_single_image_api, methods=["POST"], response_model=models.ExtrasSingleImageResponse) + self.add_api_route("/sdapi/v1/extra-batch-images", self.extras_batch_images_api, methods=["POST"], response_model=models.ExtrasBatchImagesResponse) + self.add_api_route("/sdapi/v1/png-info", self.pnginfoapi, methods=["POST"], response_model=models.PNGInfoResponse) + self.add_api_route("/sdapi/v1/progress", self.progressapi, methods=["GET"], response_model=models.ProgressResponse) + self.add_api_route("/sdapi/v1/interrogate", self.interrogateapi, methods=["POST"]) + self.add_api_route("/sdapi/v1/interrupt", self.interruptapi, methods=["POST"]) + self.add_api_route("/sdapi/v1/skip", self.skip, methods=["POST"]) + self.add_api_route("/sdapi/v1/options", self.get_config, methods=["GET"], response_model=models.OptionsModel) + self.add_api_route("/sdapi/v1/options", self.set_config, methods=["POST"]) + self.add_api_route("/sdapi/v1/cmd-flags", self.get_cmd_flags, methods=["GET"], response_model=models.FlagsModel) + self.add_api_route("/sdapi/v1/samplers", self.get_samplers, methods=["GET"], response_model=list[models.SamplerItem]) + self.add_api_route("/sdapi/v1/schedulers", self.get_schedulers, methods=["GET"], response_model=list[models.SchedulerItem]) + self.add_api_route("/sdapi/v1/upscalers", self.get_upscalers, methods=["GET"], response_model=list[models.UpscalerItem]) + self.add_api_route("/sdapi/v1/latent-upscale-modes", self.get_latent_upscale_modes, methods=["GET"], response_model=list[models.LatentUpscalerModeItem]) + self.add_api_route("/sdapi/v1/sd-models", self.get_sd_models, methods=["GET"], response_model=list[models.SDModelItem]) + self.add_api_route("/sdapi/v1/sd-vae", self.get_sd_vaes, methods=["GET"], response_model=list[models.SDVaeItem]) + self.add_api_route("/sdapi/v1/hypernetworks", self.get_hypernetworks, methods=["GET"], response_model=list[models.HypernetworkItem]) + self.add_api_route("/sdapi/v1/face-restorers", self.get_face_restorers, methods=["GET"], response_model=list[models.FaceRestorerItem]) + self.add_api_route("/sdapi/v1/realesrgan-models", self.get_realesrgan_models, methods=["GET"], response_model=list[models.RealesrganItem]) + self.add_api_route("/sdapi/v1/prompt-styles", self.get_prompt_styles, methods=["GET"], response_model=list[models.PromptStyleItem]) + self.add_api_route("/sdapi/v1/embeddings", self.get_embeddings, methods=["GET"], response_model=models.EmbeddingsResponse) + self.add_api_route("/sdapi/v1/refresh-embeddings", self.refresh_embeddings, methods=["POST"]) + self.add_api_route("/sdapi/v1/refresh-checkpoints", self.refresh_checkpoints, methods=["POST"]) + self.add_api_route("/sdapi/v1/refresh-vae", self.refresh_vae, methods=["POST"]) + self.add_api_route("/sdapi/v1/create/embedding", self.create_embedding, methods=["POST"], response_model=models.CreateResponse) + self.add_api_route("/sdapi/v1/create/hypernetwork", self.create_hypernetwork, methods=["POST"], response_model=models.CreateResponse) + self.add_api_route("/sdapi/v1/train/embedding", self.train_embedding, methods=["POST"], response_model=models.TrainResponse) + self.add_api_route("/sdapi/v1/train/hypernetwork", self.train_hypernetwork, methods=["POST"], response_model=models.TrainResponse) + self.add_api_route("/sdapi/v1/memory", self.get_memory, methods=["GET"], response_model=models.MemoryResponse) + self.add_api_route("/sdapi/v1/unload-checkpoint", self.unloadapi, methods=["POST"]) + self.add_api_route("/sdapi/v1/reload-checkpoint", self.reloadapi, methods=["POST"]) + self.add_api_route("/sdapi/v1/scripts", self.get_scripts_list, methods=["GET"], response_model=models.ScriptsList) + self.add_api_route("/sdapi/v1/script-info", self.get_script_info, methods=["GET"], response_model=list[models.ScriptInfo]) + self.add_api_route("/sdapi/v1/extensions", self.get_extensions_list, methods=["GET"], response_model=list[models.ExtensionItem]) + + if shared.cmd_opts.api_server_stop: + self.add_api_route("/sdapi/v1/server-kill", self.kill_webui, methods=["POST"]) + self.add_api_route("/sdapi/v1/server-restart", self.restart_webui, methods=["POST"]) + self.add_api_route("/sdapi/v1/server-stop", self.stop_webui, methods=["POST"]) + + self.default_script_arg_txt2img = [] + self.default_script_arg_img2img = [] + + txt2img_script_runner = scripts.scripts_txt2img + img2img_script_runner = scripts.scripts_img2img + + if not txt2img_script_runner.scripts or not img2img_script_runner.scripts: + ui.create_ui() + + if not txt2img_script_runner.scripts: + txt2img_script_runner.initialize_scripts(False) + if not self.default_script_arg_txt2img: + self.default_script_arg_txt2img = self.init_default_script_args(txt2img_script_runner) + + if not img2img_script_runner.scripts: + img2img_script_runner.initialize_scripts(True) + if not self.default_script_arg_img2img: + self.default_script_arg_img2img = self.init_default_script_args(img2img_script_runner) + + + + def add_api_route(self, path: str, endpoint, **kwargs): + if shared.cmd_opts.api_auth: + return self.app.add_api_route(path, endpoint, dependencies=[Depends(self.auth)], **kwargs) + return self.app.add_api_route(path, endpoint, **kwargs) + + def auth(self, credentials: HTTPBasicCredentials = Depends(HTTPBasic())): + if credentials.username in self.credentials: + if compare_digest(credentials.password, self.credentials[credentials.username]): + return True + + raise HTTPException(status_code=401, detail="Incorrect username or password", headers={"WWW-Authenticate": "Basic"}) + + def get_selectable_script(self, script_name, script_runner): + if script_name is None or script_name == "": + return None, None + + script_idx = script_name_to_index(script_name, script_runner.selectable_scripts) + script = script_runner.selectable_scripts[script_idx] + return script, script_idx + + def get_scripts_list(self): + t2ilist = [script.name for script in scripts.scripts_txt2img.scripts if script.name is not None] + i2ilist = [script.name for script in scripts.scripts_img2img.scripts if script.name is not None] + + return models.ScriptsList(txt2img=t2ilist, img2img=i2ilist) + + def get_script_info(self): + res = [] + + for script_list in [scripts.scripts_txt2img.scripts, scripts.scripts_img2img.scripts]: + res += [script.api_info for script in script_list if script.api_info is not None] + + return res + + def get_script(self, script_name, script_runner): + if script_name is None or script_name == "": + return None, None + + script_idx = script_name_to_index(script_name, script_runner.scripts) + return script_runner.scripts[script_idx] + + def init_default_script_args(self, script_runner): + #find max idx from the scripts in runner and generate a none array to init script_args + last_arg_index = 1 + for script in script_runner.scripts: + if last_arg_index < script.args_to: + last_arg_index = script.args_to + # None everywhere except position 0 to initialize script args + script_args = [None]*last_arg_index + script_args[0] = 0 + + # get default values + with gr.Blocks(): # will throw errors calling ui function without this + for script in script_runner.scripts: + if script.ui(script.is_img2img): + ui_default_values = [] + for elem in script.ui(script.is_img2img): + ui_default_values.append(elem.value) + script_args[script.args_from:script.args_to] = ui_default_values + return script_args + + def init_script_args(self, request, default_script_args, selectable_scripts, selectable_idx, script_runner, *, input_script_args=None): + script_args = default_script_args.copy() + + if input_script_args is not None: + for index, value in input_script_args.items(): + script_args[index] = value + + # position 0 in script_arg is the idx+1 of the selectable script that is going to be run when using scripts.scripts_*2img.run() + if selectable_scripts: + script_args[selectable_scripts.args_from:selectable_scripts.args_to] = request.script_args + script_args[0] = selectable_idx + 1 + + # Now check for always on scripts + if request.alwayson_scripts: + for alwayson_script_name in request.alwayson_scripts.keys(): + alwayson_script = self.get_script(alwayson_script_name, script_runner) + if alwayson_script is None: + raise HTTPException(status_code=422, detail=f"always on script {alwayson_script_name} not found") + # Selectable script in always on script param check + if alwayson_script.alwayson is False: + raise HTTPException(status_code=422, detail="Cannot have a selectable script in the always on scripts params") + # always on script with no arg should always run so you don't really need to add them to the requests + if "args" in request.alwayson_scripts[alwayson_script_name]: + # min between arg length in scriptrunner and arg length in the request + for idx in range(0, min((alwayson_script.args_to - alwayson_script.args_from), len(request.alwayson_scripts[alwayson_script_name]["args"]))): + script_args[alwayson_script.args_from + idx] = request.alwayson_scripts[alwayson_script_name]["args"][idx] + return script_args + + def apply_infotext(self, request, tabname, *, script_runner=None, mentioned_script_args=None): + """Processes `infotext` field from the `request`, and sets other fields of the `request` according to what's in infotext. + + If request already has a field set, and that field is encountered in infotext too, the value from infotext is ignored. + + Additionally, fills `mentioned_script_args` dict with index: value pairs for script arguments read from infotext. + """ + + if not request.infotext: + return {} + + possible_fields = infotext_utils.paste_fields[tabname]["fields"] + set_fields = request.model_dump(exclude_unset=True) if hasattr(request, "request") else request.dict(exclude_unset=True) # pydantic v1/v2 have different names for this + params = infotext_utils.parse_generation_parameters(request.infotext) + + def get_field_value(field, params): + value = field.function(params) if field.function else params.get(field.label) + if value is None: + return None + + if field.api in request.__fields__: + target_type = request.__fields__[field.api].type_ + else: + target_type = type(field.component.value) + + if target_type == type(None): + return None + + if isinstance(value, dict) and value.get('__type__') == 'generic_update': # this is a gradio.update rather than a value + value = value.get('value') + + if value is not None and not isinstance(value, target_type): + value = target_type(value) + + return value + + for field in possible_fields: + if not field.api: + continue + + if field.api in set_fields: + continue + + value = get_field_value(field, params) + if value is not None: + setattr(request, field.api, value) + + if request.override_settings is None: + request.override_settings = {} + + overridden_settings = infotext_utils.get_override_settings(params) + for _, setting_name, value in overridden_settings: + if setting_name not in request.override_settings: + request.override_settings[setting_name] = value + + if script_runner is not None and mentioned_script_args is not None: + indexes = {v: i for i, v in enumerate(script_runner.inputs)} + script_fields = ((field, indexes[field.component]) for field in possible_fields if field.component in indexes) + + for field, index in script_fields: + value = get_field_value(field, params) + + if value is None: + continue + + mentioned_script_args[index] = value + + return params + + def text2imgapi(self, txt2imgreq: models.StableDiffusionTxt2ImgProcessingAPI): + task_id = txt2imgreq.force_task_id or create_task_id("txt2img") + + script_runner = scripts.scripts_txt2img + + infotext_script_args = {} + self.apply_infotext(txt2imgreq, "txt2img", script_runner=script_runner, mentioned_script_args=infotext_script_args) + + selectable_scripts, selectable_script_idx = self.get_selectable_script(txt2imgreq.script_name, script_runner) + sampler, scheduler = sd_samplers.get_sampler_and_scheduler(txt2imgreq.sampler_name or txt2imgreq.sampler_index, txt2imgreq.scheduler) + + populate = txt2imgreq.copy(update={ # Override __init__ params + "sampler_name": validate_sampler_name(sampler), + "do_not_save_samples": not txt2imgreq.save_images, + "do_not_save_grid": not txt2imgreq.save_images, + }) + if populate.sampler_name: + populate.sampler_index = None # prevent a warning later on + + if not populate.scheduler and scheduler != "Automatic": + populate.scheduler = scheduler + + args = vars(populate) + args.pop('script_name', None) + args.pop('script_args', None) # will refeed them to the pipeline directly after initializing them + args.pop('alwayson_scripts', None) + args.pop('infotext', None) + + script_args = self.init_script_args(txt2imgreq, self.default_script_arg_txt2img, selectable_scripts, selectable_script_idx, script_runner, input_script_args=infotext_script_args) + + send_images = args.pop('send_images', True) + args.pop('save_images', None) + + add_task_to_queue(task_id) + + with self.queue_lock: + with closing(StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **args)) as p: + p.is_api = True + p.scripts = script_runner + p.outpath_grids = opts.outdir_txt2img_grids + p.outpath_samples = opts.outdir_txt2img_samples + + try: + shared.state.begin(job="scripts_txt2img") + start_task(task_id) + if selectable_scripts is not None: + p.script_args = script_args + processed = scripts.scripts_txt2img.run(p, *p.script_args) # Need to pass args as list here + else: + p.script_args = tuple(script_args) # Need to pass args as tuple here + processed = process_images(p) + finish_task(task_id) + finally: + shared.state.end() + shared.total_tqdm.clear() + + b64images = list(map(encode_pil_to_base64, processed.images)) if send_images else [] + + return models.TextToImageResponse(images=b64images, parameters=vars(txt2imgreq), info=processed.js()) + + def img2imgapi(self, img2imgreq: models.StableDiffusionImg2ImgProcessingAPI): + task_id = img2imgreq.force_task_id or create_task_id("img2img") + + init_images = img2imgreq.init_images + if init_images is None: + raise HTTPException(status_code=404, detail="Init image not found") + + mask = img2imgreq.mask + if mask: + mask = decode_base64_to_image(mask) + + script_runner = scripts.scripts_img2img + + infotext_script_args = {} + self.apply_infotext(img2imgreq, "img2img", script_runner=script_runner, mentioned_script_args=infotext_script_args) + + selectable_scripts, selectable_script_idx = self.get_selectable_script(img2imgreq.script_name, script_runner) + sampler, scheduler = sd_samplers.get_sampler_and_scheduler(img2imgreq.sampler_name or img2imgreq.sampler_index, img2imgreq.scheduler) + + populate = img2imgreq.copy(update={ # Override __init__ params + "sampler_name": validate_sampler_name(sampler), + "do_not_save_samples": not img2imgreq.save_images, + "do_not_save_grid": not img2imgreq.save_images, + "mask": mask, + }) + if populate.sampler_name: + populate.sampler_index = None # prevent a warning later on + + if not populate.scheduler and scheduler != "Automatic": + populate.scheduler = scheduler + + args = vars(populate) + args.pop('include_init_images', None) # this is meant to be done by "exclude": True in model, but it's for a reason that I cannot determine. + args.pop('script_name', None) + args.pop('script_args', None) # will refeed them to the pipeline directly after initializing them + args.pop('alwayson_scripts', None) + args.pop('infotext', None) + + script_args = self.init_script_args(img2imgreq, self.default_script_arg_img2img, selectable_scripts, selectable_script_idx, script_runner, input_script_args=infotext_script_args) + + send_images = args.pop('send_images', True) + args.pop('save_images', None) + + add_task_to_queue(task_id) + + with self.queue_lock: + with closing(StableDiffusionProcessingImg2Img(sd_model=shared.sd_model, **args)) as p: + p.init_images = [decode_base64_to_image(x) for x in init_images] + p.is_api = True + p.scripts = script_runner + p.outpath_grids = opts.outdir_img2img_grids + p.outpath_samples = opts.outdir_img2img_samples + + try: + shared.state.begin(job="scripts_img2img") + start_task(task_id) + if selectable_scripts is not None: + p.script_args = script_args + processed = scripts.scripts_img2img.run(p, *p.script_args) # Need to pass args as list here + else: + p.script_args = tuple(script_args) # Need to pass args as tuple here + processed = process_images(p) + finish_task(task_id) + finally: + shared.state.end() + shared.total_tqdm.clear() + + b64images = list(map(encode_pil_to_base64, processed.images)) if send_images else [] + + if not img2imgreq.include_init_images: + img2imgreq.init_images = None + img2imgreq.mask = None + + return models.ImageToImageResponse(images=b64images, parameters=vars(img2imgreq), info=processed.js()) + + def extras_single_image_api(self, req: models.ExtrasSingleImageRequest): + reqDict = setUpscalers(req) + + reqDict['image'] = decode_base64_to_image(reqDict['image']) + + with self.queue_lock: + result = postprocessing.run_extras(extras_mode=0, image_folder="", input_dir="", output_dir="", save_output=False, **reqDict) + + return models.ExtrasSingleImageResponse(image=encode_pil_to_base64(result[0][0]), html_info=result[1]) + + def extras_batch_images_api(self, req: models.ExtrasBatchImagesRequest): + reqDict = setUpscalers(req) + + image_list = reqDict.pop('imageList', []) + image_folder = [decode_base64_to_image(x.data) for x in image_list] + + with self.queue_lock: + result = postprocessing.run_extras(extras_mode=1, image_folder=image_folder, image="", input_dir="", output_dir="", save_output=False, **reqDict) + + return models.ExtrasBatchImagesResponse(images=list(map(encode_pil_to_base64, result[0])), html_info=result[1]) + + def pnginfoapi(self, req: models.PNGInfoRequest): + image = decode_base64_to_image(req.image.strip()) + if image is None: + return models.PNGInfoResponse(info="") + + geninfo, items = images.read_info_from_image(image) + if geninfo is None: + geninfo = "" + + params = infotext_utils.parse_generation_parameters(geninfo) + script_callbacks.infotext_pasted_callback(geninfo, params) + + return models.PNGInfoResponse(info=geninfo, items=items, parameters=params) + + def progressapi(self, req: models.ProgressRequest = Depends()): + # copy from check_progress_call of ui.py + + if shared.state.job_count == 0: + return models.ProgressResponse(progress=0, eta_relative=0, state=shared.state.dict(), textinfo=shared.state.textinfo) + + # avoid dividing zero + progress = 0.01 + + if shared.state.job_count > 0: + progress += shared.state.job_no / shared.state.job_count + if shared.state.sampling_steps > 0: + progress += 1 / shared.state.job_count * shared.state.sampling_step / shared.state.sampling_steps + + time_since_start = time.time() - shared.state.time_start + eta = (time_since_start/progress) + eta_relative = eta-time_since_start + + progress = min(progress, 1) + + shared.state.set_current_image() + + current_image = None + if shared.state.current_image and not req.skip_current_image: + current_image = encode_pil_to_base64(shared.state.current_image) + + return models.ProgressResponse(progress=progress, eta_relative=eta_relative, state=shared.state.dict(), current_image=current_image, textinfo=shared.state.textinfo, current_task=current_task) + + def interrogateapi(self, interrogatereq: models.InterrogateRequest): + image_b64 = interrogatereq.image + if image_b64 is None: + raise HTTPException(status_code=404, detail="Image not found") + + img = decode_base64_to_image(image_b64) + img = img.convert('RGB') + + # Override object param + with self.queue_lock: + if interrogatereq.model == "clip": + processed = shared.interrogator.interrogate(img) + elif interrogatereq.model == "deepdanbooru": + processed = deepbooru.model.tag(img) + else: + raise HTTPException(status_code=404, detail="Model not found") + + return models.InterrogateResponse(caption=processed) + + def interruptapi(self): + shared.state.interrupt() + + return {} + + def unloadapi(self): + sd_models.unload_model_weights() + + return {} + + def reloadapi(self): + sd_models.send_model_to_device(shared.sd_model) + + return {} + + def skip(self): + shared.state.skip() + + def get_config(self): + options = {} + for key in shared.opts.data.keys(): + metadata = shared.opts.data_labels.get(key) + if(metadata is not None): + options.update({key: shared.opts.data.get(key, shared.opts.data_labels.get(key).default)}) + else: + options.update({key: shared.opts.data.get(key, None)}) + + return options + + def set_config(self, req: dict[str, Any]): + checkpoint_name = req.get("sd_model_checkpoint", None) + if checkpoint_name is not None and checkpoint_name not in sd_models.checkpoint_aliases: + raise RuntimeError(f"model {checkpoint_name!r} not found") + + for k, v in req.items(): + shared.opts.set(k, v, is_api=True) + + shared.opts.save(shared.config_filename) + return + + def get_cmd_flags(self): + return vars(shared.cmd_opts) + + def get_samplers(self): + return [{"name": sampler[0], "aliases":sampler[2], "options":sampler[3]} for sampler in sd_samplers.all_samplers] + + def get_schedulers(self): + return [ + { + "name": scheduler.name, + "label": scheduler.label, + "aliases": scheduler.aliases, + "default_rho": scheduler.default_rho, + "need_inner_model": scheduler.need_inner_model, + } + for scheduler in sd_schedulers.schedulers] + + def get_upscalers(self): + return [ + { + "name": upscaler.name, + "model_name": upscaler.scaler.model_name, + "model_path": upscaler.data_path, + "model_url": None, + "scale": upscaler.scale, + } + for upscaler in shared.sd_upscalers + ] + + def get_latent_upscale_modes(self): + return [ + { + "name": upscale_mode, + } + for upscale_mode in [*(shared.latent_upscale_modes or {})] + ] + + def get_sd_models(self): + import modules.sd_models as sd_models + return [{"title": x.title, "model_name": x.model_name, "hash": x.shorthash, "sha256": x.sha256, "filename": x.filename, "config": find_checkpoint_config_near_filename(x)} for x in sd_models.checkpoints_list.values()] + + def get_sd_vaes(self): + import modules.sd_vae as sd_vae + return [{"model_name": x, "filename": sd_vae.vae_dict[x]} for x in sd_vae.vae_dict.keys()] + + def get_hypernetworks(self): + return [{"name": name, "path": shared.hypernetworks[name]} for name in shared.hypernetworks] + + def get_face_restorers(self): + return [{"name":x.name(), "cmd_dir": getattr(x, "cmd_dir", None)} for x in shared.face_restorers] + + def get_realesrgan_models(self): + return [{"name":x.name,"path":x.data_path, "scale":x.scale} for x in get_realesrgan_models(None)] + + def get_prompt_styles(self): + styleList = [] + for k in shared.prompt_styles.styles: + style = shared.prompt_styles.styles[k] + styleList.append({"name":style[0], "prompt": style[1], "negative_prompt": style[2]}) + + return styleList + + def get_embeddings(self): + db = sd_hijack.model_hijack.embedding_db + + def convert_embedding(embedding): + return { + "step": embedding.step, + "sd_checkpoint": embedding.sd_checkpoint, + "sd_checkpoint_name": embedding.sd_checkpoint_name, + "shape": embedding.shape, + "vectors": embedding.vectors, + } + + def convert_embeddings(embeddings): + return {embedding.name: convert_embedding(embedding) for embedding in embeddings.values()} + + return { + "loaded": convert_embeddings(db.word_embeddings), + "skipped": convert_embeddings(db.skipped_embeddings), + } + + def refresh_embeddings(self): + with self.queue_lock: + sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings(force_reload=True) + + def refresh_checkpoints(self): + with self.queue_lock: + shared.refresh_checkpoints() + + def refresh_vae(self): + with self.queue_lock: + shared_items.refresh_vae_list() + + def create_embedding(self, args: dict): + try: + shared.state.begin(job="create_embedding") + filename = create_embedding(**args) # create empty embedding + sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings() # reload embeddings so new one can be immediately used + return models.CreateResponse(info=f"create embedding filename: {filename}") + except AssertionError as e: + return models.TrainResponse(info=f"create embedding error: {e}") + finally: + shared.state.end() + + + def create_hypernetwork(self, args: dict): + try: + shared.state.begin(job="create_hypernetwork") + filename = create_hypernetwork(**args) # create empty embedding + return models.CreateResponse(info=f"create hypernetwork filename: {filename}") + except AssertionError as e: + return models.TrainResponse(info=f"create hypernetwork error: {e}") + finally: + shared.state.end() + + def train_embedding(self, args: dict): + try: + shared.state.begin(job="train_embedding") + apply_optimizations = shared.opts.training_xattention_optimizations + error = None + filename = '' + if not apply_optimizations: + sd_hijack.undo_optimizations() + try: + embedding, filename = train_embedding(**args) # can take a long time to complete + except Exception as e: + error = e + finally: + if not apply_optimizations: + sd_hijack.apply_optimizations() + return models.TrainResponse(info=f"train embedding complete: filename: {filename} error: {error}") + except Exception as msg: + return models.TrainResponse(info=f"train embedding error: {msg}") + finally: + shared.state.end() + + def train_hypernetwork(self, args: dict): + try: + shared.state.begin(job="train_hypernetwork") + shared.loaded_hypernetworks = [] + apply_optimizations = shared.opts.training_xattention_optimizations + error = None + filename = '' + if not apply_optimizations: + sd_hijack.undo_optimizations() + try: + hypernetwork, filename = train_hypernetwork(**args) + except Exception as e: + error = e + finally: + shared.sd_model.cond_stage_model.to(devices.device) + shared.sd_model.first_stage_model.to(devices.device) + if not apply_optimizations: + sd_hijack.apply_optimizations() + shared.state.end() + return models.TrainResponse(info=f"train embedding complete: filename: {filename} error: {error}") + except Exception as exc: + return models.TrainResponse(info=f"train embedding error: {exc}") + finally: + shared.state.end() + + def get_memory(self): + try: + import os + import psutil + process = psutil.Process(os.getpid()) + res = process.memory_info() # only rss is cross-platform guaranteed so we dont rely on other values + ram_total = 100 * res.rss / process.memory_percent() # and total memory is calculated as actual value is not cross-platform safe + ram = { 'free': ram_total - res.rss, 'used': res.rss, 'total': ram_total } + except Exception as err: + ram = { 'error': f'{err}' } + try: + import torch + if torch.cuda.is_available(): + s = torch.cuda.mem_get_info() + system = { 'free': s[0], 'used': s[1] - s[0], 'total': s[1] } + s = dict(torch.cuda.memory_stats(shared.device)) + allocated = { 'current': s['allocated_bytes.all.current'], 'peak': s['allocated_bytes.all.peak'] } + reserved = { 'current': s['reserved_bytes.all.current'], 'peak': s['reserved_bytes.all.peak'] } + active = { 'current': s['active_bytes.all.current'], 'peak': s['active_bytes.all.peak'] } + inactive = { 'current': s['inactive_split_bytes.all.current'], 'peak': s['inactive_split_bytes.all.peak'] } + warnings = { 'retries': s['num_alloc_retries'], 'oom': s['num_ooms'] } + cuda = { + 'system': system, + 'active': active, + 'allocated': allocated, + 'reserved': reserved, + 'inactive': inactive, + 'events': warnings, + } + else: + cuda = {'error': 'unavailable'} + except Exception as err: + cuda = {'error': f'{err}'} + return models.MemoryResponse(ram=ram, cuda=cuda) + + def get_extensions_list(self): + from modules import extensions + extensions.list_extensions() + ext_list = [] + for ext in extensions.extensions: + ext: extensions.Extension + ext.read_info_from_repo() + if ext.remote is not None: + ext_list.append({ + "name": ext.name, + "remote": ext.remote, + "branch": ext.branch, + "commit_hash":ext.commit_hash, + "commit_date":ext.commit_date, + "version":ext.version, + "enabled":ext.enabled + }) + return ext_list + + def launch(self, server_name, port, root_path): + self.app.include_router(self.router) + uvicorn.run( + self.app, + host=server_name, + port=port, + timeout_keep_alive=shared.cmd_opts.timeout_keep_alive, + root_path=root_path, + ssl_keyfile=shared.cmd_opts.tls_keyfile, + ssl_certfile=shared.cmd_opts.tls_certfile + ) + + def kill_webui(self): + restart.stop_program() + + def restart_webui(self): + if restart.is_restartable(): + restart.restart_program() + return Response(status_code=501) + + def stop_webui(request): + shared.state.server_command = "stop" + return Response("Stopping.") + diff --git a/stable-diffusion-webui/modules/api/models.py b/stable-diffusion-webui/modules/api/models.py new file mode 100755 index 0000000..79c4cbb --- /dev/null +++ b/stable-diffusion-webui/modules/api/models.py @@ -0,0 +1,329 @@ +import inspect + +from pydantic import BaseModel, Field, create_model +from typing import Any, Optional, Literal +from inflection import underscore +from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img +from modules.shared import sd_upscalers, opts, parser + +API_NOT_ALLOWED = [ + "self", + "kwargs", + "sd_model", + "outpath_samples", + "outpath_grids", + "sampler_index", + # "do_not_save_samples", + # "do_not_save_grid", + "extra_generation_params", + "overlay_images", + "do_not_reload_embeddings", + "seed_enable_extras", + "prompt_for_display", + "sampler_noise_scheduler_override", + "ddim_discretize" +] + +class ModelDef(BaseModel): + """Assistance Class for Pydantic Dynamic Model Generation""" + + field: str + field_alias: str + field_type: Any + field_value: Any + field_exclude: bool = False + + +class PydanticModelGenerator: + """ + Takes in created classes and stubs them out in a way FastAPI/Pydantic is happy about: + source_data is a snapshot of the default values produced by the class + params are the names of the actual keys required by __init__ + """ + + def __init__( + self, + model_name: str = None, + class_instance = None, + additional_fields = None, + ): + def field_type_generator(k, v): + field_type = v.annotation + + if field_type == 'Image': + # images are sent as base64 strings via API + field_type = 'str' + + return Optional[field_type] + + def merge_class_params(class_): + all_classes = list(filter(lambda x: x is not object, inspect.getmro(class_))) + parameters = {} + for classes in all_classes: + parameters = {**parameters, **inspect.signature(classes.__init__).parameters} + return parameters + + self._model_name = model_name + self._class_data = merge_class_params(class_instance) + + self._model_def = [ + ModelDef( + field=underscore(k), + field_alias=k, + field_type=field_type_generator(k, v), + field_value=None if isinstance(v.default, property) else v.default + ) + for (k,v) in self._class_data.items() if k not in API_NOT_ALLOWED + ] + + for fields in additional_fields: + self._model_def.append(ModelDef( + field=underscore(fields["key"]), + field_alias=fields["key"], + field_type=fields["type"], + field_value=fields["default"], + field_exclude=fields["exclude"] if "exclude" in fields else False)) + + def generate_model(self): + """ + Creates a pydantic BaseModel + from the json and overrides provided at initialization + """ + fields = { + d.field: (d.field_type, Field(default=d.field_value, alias=d.field_alias, exclude=d.field_exclude)) for d in self._model_def + } + DynamicModel = create_model(self._model_name, **fields) + DynamicModel.__config__.allow_population_by_field_name = True + DynamicModel.__config__.allow_mutation = True + return DynamicModel + +StableDiffusionTxt2ImgProcessingAPI = PydanticModelGenerator( + "StableDiffusionProcessingTxt2Img", + StableDiffusionProcessingTxt2Img, + [ + {"key": "sampler_index", "type": str, "default": "Euler"}, + {"key": "script_name", "type": str, "default": None}, + {"key": "script_args", "type": list, "default": []}, + {"key": "send_images", "type": bool, "default": True}, + {"key": "save_images", "type": bool, "default": False}, + {"key": "alwayson_scripts", "type": dict, "default": {}}, + {"key": "force_task_id", "type": str, "default": None}, + {"key": "infotext", "type": str, "default": None}, + ] +).generate_model() + +StableDiffusionImg2ImgProcessingAPI = PydanticModelGenerator( + "StableDiffusionProcessingImg2Img", + StableDiffusionProcessingImg2Img, + [ + {"key": "sampler_index", "type": str, "default": "Euler"}, + {"key": "init_images", "type": list, "default": None}, + {"key": "denoising_strength", "type": float, "default": 0.75}, + {"key": "mask", "type": str, "default": None}, + {"key": "include_init_images", "type": bool, "default": False, "exclude" : True}, + {"key": "script_name", "type": str, "default": None}, + {"key": "script_args", "type": list, "default": []}, + {"key": "send_images", "type": bool, "default": True}, + {"key": "save_images", "type": bool, "default": False}, + {"key": "alwayson_scripts", "type": dict, "default": {}}, + {"key": "force_task_id", "type": str, "default": None}, + {"key": "infotext", "type": str, "default": None}, + ] +).generate_model() + +class TextToImageResponse(BaseModel): + images: list[str] = Field(default=None, title="Image", description="The generated image in base64 format.") + parameters: dict + info: str + +class ImageToImageResponse(BaseModel): + images: list[str] = Field(default=None, title="Image", description="The generated image in base64 format.") + parameters: dict + info: str + +class ExtrasBaseRequest(BaseModel): + resize_mode: Literal[0, 1] = Field(default=0, title="Resize Mode", description="Sets the resize mode: 0 to upscale by upscaling_resize amount, 1 to upscale up to upscaling_resize_h x upscaling_resize_w.") + show_extras_results: bool = Field(default=True, title="Show results", description="Should the backend return the generated image?") + gfpgan_visibility: float = Field(default=0, title="GFPGAN Visibility", ge=0, le=1, allow_inf_nan=False, description="Sets the visibility of GFPGAN, values should be between 0 and 1.") + codeformer_visibility: float = Field(default=0, title="CodeFormer Visibility", ge=0, le=1, allow_inf_nan=False, description="Sets the visibility of CodeFormer, values should be between 0 and 1.") + codeformer_weight: float = Field(default=0, title="CodeFormer Weight", ge=0, le=1, allow_inf_nan=False, description="Sets the weight of CodeFormer, values should be between 0 and 1.") + upscaling_resize: float = Field(default=2, title="Upscaling Factor", gt=0, description="By how much to upscale the image, only used when resize_mode=0.") + upscaling_resize_w: int = Field(default=512, title="Target Width", ge=1, description="Target width for the upscaler to hit. Only used when resize_mode=1.") + upscaling_resize_h: int = Field(default=512, title="Target Height", ge=1, description="Target height for the upscaler to hit. Only used when resize_mode=1.") + upscaling_crop: bool = Field(default=True, title="Crop to fit", description="Should the upscaler crop the image to fit in the chosen size?") + upscaler_1: str = Field(default="None", title="Main upscaler", description=f"The name of the main upscaler to use, it has to be one of this list: {' , '.join([x.name for x in sd_upscalers])}") + upscaler_2: str = Field(default="None", title="Secondary upscaler", description=f"The name of the secondary upscaler to use, it has to be one of this list: {' , '.join([x.name for x in sd_upscalers])}") + extras_upscaler_2_visibility: float = Field(default=0, title="Secondary upscaler visibility", ge=0, le=1, allow_inf_nan=False, description="Sets the visibility of secondary upscaler, values should be between 0 and 1.") + upscale_first: bool = Field(default=False, title="Upscale first", description="Should the upscaler run before restoring faces?") + +class ExtraBaseResponse(BaseModel): + html_info: str = Field(title="HTML info", description="A series of HTML tags containing the process info.") + +class ExtrasSingleImageRequest(ExtrasBaseRequest): + image: str = Field(default="", title="Image", description="Image to work on, must be a Base64 string containing the image's data.") + +class ExtrasSingleImageResponse(ExtraBaseResponse): + image: str = Field(default=None, title="Image", description="The generated image in base64 format.") + +class FileData(BaseModel): + data: str = Field(title="File data", description="Base64 representation of the file") + name: str = Field(title="File name") + +class ExtrasBatchImagesRequest(ExtrasBaseRequest): + imageList: list[FileData] = Field(title="Images", description="List of images to work on. Must be Base64 strings") + +class ExtrasBatchImagesResponse(ExtraBaseResponse): + images: list[str] = Field(title="Images", description="The generated images in base64 format.") + +class PNGInfoRequest(BaseModel): + image: str = Field(title="Image", description="The base64 encoded PNG image") + +class PNGInfoResponse(BaseModel): + info: str = Field(title="Image info", description="A string with the parameters used to generate the image") + items: dict = Field(title="Items", description="A dictionary containing all the other fields the image had") + parameters: dict = Field(title="Parameters", description="A dictionary with parsed generation info fields") + +class ProgressRequest(BaseModel): + skip_current_image: bool = Field(default=False, title="Skip current image", description="Skip current image serialization") + +class ProgressResponse(BaseModel): + progress: float = Field(title="Progress", description="The progress with a range of 0 to 1") + eta_relative: float = Field(title="ETA in secs") + state: dict = Field(title="State", description="The current state snapshot") + current_image: str = Field(default=None, title="Current image", description="The current image in base64 format. opts.show_progress_every_n_steps is required for this to work.") + textinfo: str = Field(default=None, title="Info text", description="Info text used by WebUI.") + +class InterrogateRequest(BaseModel): + image: str = Field(default="", title="Image", description="Image to work on, must be a Base64 string containing the image's data.") + model: str = Field(default="clip", title="Model", description="The interrogate model used.") + +class InterrogateResponse(BaseModel): + caption: str = Field(default=None, title="Caption", description="The generated caption for the image.") + +class TrainResponse(BaseModel): + info: str = Field(title="Train info", description="Response string from train embedding or hypernetwork task.") + +class CreateResponse(BaseModel): + info: str = Field(title="Create info", description="Response string from create embedding or hypernetwork task.") + +fields = {} +for key, metadata in opts.data_labels.items(): + value = opts.data.get(key) + optType = opts.typemap.get(type(metadata.default), type(metadata.default)) if metadata.default else Any + + if metadata is not None: + fields.update({key: (Optional[optType], Field(default=metadata.default, description=metadata.label))}) + else: + fields.update({key: (Optional[optType], Field())}) + +OptionsModel = create_model("Options", **fields) + +flags = {} +_options = vars(parser)['_option_string_actions'] +for key in _options: + if(_options[key].dest != 'help'): + flag = _options[key] + _type = str + if _options[key].default is not None: + _type = type(_options[key].default) + flags.update({flag.dest: (_type, Field(default=flag.default, description=flag.help))}) + +FlagsModel = create_model("Flags", **flags) + +class SamplerItem(BaseModel): + name: str = Field(title="Name") + aliases: list[str] = Field(title="Aliases") + options: dict[str, str] = Field(title="Options") + +class SchedulerItem(BaseModel): + name: str = Field(title="Name") + label: str = Field(title="Label") + aliases: Optional[list[str]] = Field(title="Aliases") + default_rho: Optional[float] = Field(title="Default Rho") + need_inner_model: Optional[bool] = Field(title="Needs Inner Model") + +class UpscalerItem(BaseModel): + name: str = Field(title="Name") + model_name: Optional[str] = Field(title="Model Name") + model_path: Optional[str] = Field(title="Path") + model_url: Optional[str] = Field(title="URL") + scale: Optional[float] = Field(title="Scale") + +class LatentUpscalerModeItem(BaseModel): + name: str = Field(title="Name") + +class SDModelItem(BaseModel): + title: str = Field(title="Title") + model_name: str = Field(title="Model Name") + hash: Optional[str] = Field(title="Short hash") + sha256: Optional[str] = Field(title="sha256 hash") + filename: str = Field(title="Filename") + config: Optional[str] = Field(title="Config file") + +class SDVaeItem(BaseModel): + model_name: str = Field(title="Model Name") + filename: str = Field(title="Filename") + +class HypernetworkItem(BaseModel): + name: str = Field(title="Name") + path: Optional[str] = Field(title="Path") + +class FaceRestorerItem(BaseModel): + name: str = Field(title="Name") + cmd_dir: Optional[str] = Field(title="Path") + +class RealesrganItem(BaseModel): + name: str = Field(title="Name") + path: Optional[str] = Field(title="Path") + scale: Optional[int] = Field(title="Scale") + +class PromptStyleItem(BaseModel): + name: str = Field(title="Name") + prompt: Optional[str] = Field(title="Prompt") + negative_prompt: Optional[str] = Field(title="Negative Prompt") + + +class EmbeddingItem(BaseModel): + step: Optional[int] = Field(title="Step", description="The number of steps that were used to train this embedding, if available") + sd_checkpoint: Optional[str] = Field(title="SD Checkpoint", description="The hash of the checkpoint this embedding was trained on, if available") + sd_checkpoint_name: Optional[str] = Field(title="SD Checkpoint Name", description="The name of the checkpoint this embedding was trained on, if available. Note that this is the name that was used by the trainer; for a stable identifier, use `sd_checkpoint` instead") + shape: int = Field(title="Shape", description="The length of each individual vector in the embedding") + vectors: int = Field(title="Vectors", description="The number of vectors in the embedding") + +class EmbeddingsResponse(BaseModel): + loaded: dict[str, EmbeddingItem] = Field(title="Loaded", description="Embeddings loaded for the current model") + skipped: dict[str, EmbeddingItem] = Field(title="Skipped", description="Embeddings skipped for the current model (likely due to architecture incompatibility)") + +class MemoryResponse(BaseModel): + ram: dict = Field(title="RAM", description="System memory stats") + cuda: dict = Field(title="CUDA", description="nVidia CUDA memory stats") + + +class ScriptsList(BaseModel): + txt2img: list = Field(default=None, title="Txt2img", description="Titles of scripts (txt2img)") + img2img: list = Field(default=None, title="Img2img", description="Titles of scripts (img2img)") + + +class ScriptArg(BaseModel): + label: str = Field(default=None, title="Label", description="Name of the argument in UI") + value: Optional[Any] = Field(default=None, title="Value", description="Default value of the argument") + minimum: Optional[Any] = Field(default=None, title="Minimum", description="Minimum allowed value for the argumentin UI") + maximum: Optional[Any] = Field(default=None, title="Minimum", description="Maximum allowed value for the argumentin UI") + step: Optional[Any] = Field(default=None, title="Minimum", description="Step for changing value of the argumentin UI") + choices: Optional[list[str]] = Field(default=None, title="Choices", description="Possible values for the argument") + + +class ScriptInfo(BaseModel): + name: str = Field(default=None, title="Name", description="Script name") + is_alwayson: bool = Field(default=None, title="IsAlwayson", description="Flag specifying whether this script is an alwayson script") + is_img2img: bool = Field(default=None, title="IsImg2img", description="Flag specifying whether this script is an img2img script") + args: list[ScriptArg] = Field(title="Arguments", description="List of script's arguments") + +class ExtensionItem(BaseModel): + name: str = Field(title="Name", description="Extension name") + remote: str = Field(title="Remote", description="Extension Repository URL") + branch: str = Field(title="Branch", description="Extension Repository Branch") + commit_hash: str = Field(title="Commit Hash", description="Extension Repository Commit Hash") + version: str = Field(title="Version", description="Extension Version") + commit_date: str = Field(title="Commit Date", description="Extension Repository Commit Date") + enabled: bool = Field(title="Enabled", description="Flag specifying whether this extension is enabled") diff --git a/stable-diffusion-webui/modules/cache.py b/stable-diffusion-webui/modules/cache.py new file mode 100755 index 0000000..4790655 --- /dev/null +++ b/stable-diffusion-webui/modules/cache.py @@ -0,0 +1,123 @@ +import json +import os +import os.path +import threading + +import diskcache +import tqdm + +from modules.paths import data_path, script_path + +cache_filename = os.environ.get('SD_WEBUI_CACHE_FILE', os.path.join(data_path, "cache.json")) +cache_dir = os.environ.get('SD_WEBUI_CACHE_DIR', os.path.join(data_path, "cache")) +caches = {} +cache_lock = threading.Lock() + + +def dump_cache(): + """old function for dumping cache to disk; does nothing since diskcache.""" + + pass + + +def make_cache(subsection: str) -> diskcache.Cache: + return diskcache.Cache( + os.path.join(cache_dir, subsection), + size_limit=2**32, # 4 GB, culling oldest first + disk_min_file_size=2**18, # keep up to 256KB in Sqlite + ) + + +def convert_old_cached_data(): + try: + with open(cache_filename, "r", encoding="utf8") as file: + data = json.load(file) + except FileNotFoundError: + return + except Exception: + os.replace(cache_filename, os.path.join(script_path, "tmp", "cache.json")) + print('[ERROR] issue occurred while trying to read cache.json; old cache has been moved to tmp/cache.json') + return + + total_count = sum(len(keyvalues) for keyvalues in data.values()) + + with tqdm.tqdm(total=total_count, desc="converting cache") as progress: + for subsection, keyvalues in data.items(): + cache_obj = caches.get(subsection) + if cache_obj is None: + cache_obj = make_cache(subsection) + caches[subsection] = cache_obj + + for key, value in keyvalues.items(): + cache_obj[key] = value + progress.update(1) + + +def cache(subsection): + """ + Retrieves or initializes a cache for a specific subsection. + + Parameters: + subsection (str): The subsection identifier for the cache. + + Returns: + diskcache.Cache: The cache data for the specified subsection. + """ + + cache_obj = caches.get(subsection) + if not cache_obj: + with cache_lock: + if not os.path.exists(cache_dir) and os.path.isfile(cache_filename): + convert_old_cached_data() + + cache_obj = caches.get(subsection) + if not cache_obj: + cache_obj = make_cache(subsection) + caches[subsection] = cache_obj + + return cache_obj + + +def cached_data_for_file(subsection, title, filename, func): + """ + Retrieves or generates data for a specific file, using a caching mechanism. + + Parameters: + subsection (str): The subsection of the cache to use. + title (str): The title of the data entry in the subsection of the cache. + filename (str): The path to the file to be checked for modifications. + func (callable): A function that generates the data if it is not available in the cache. + + Returns: + dict or None: The cached or generated data, or None if data generation fails. + + The `cached_data_for_file` function implements a caching mechanism for data stored in files. + It checks if the data associated with the given `title` is present in the cache and compares the + modification time of the file with the cached modification time. If the file has been modified, + the cache is considered invalid and the data is regenerated using the provided `func`. + Otherwise, the cached data is returned. + + If the data generation fails, None is returned to indicate the failure. Otherwise, the generated + or cached data is returned as a dictionary. + """ + + existing_cache = cache(subsection) + ondisk_mtime = os.path.getmtime(filename) + + entry = existing_cache.get(title) + if entry: + cached_mtime = entry.get("mtime", 0) + if ondisk_mtime > cached_mtime: + entry = None + + if not entry or 'value' not in entry: + value = func() + if value is None: + return None + + entry = {'mtime': ondisk_mtime, 'value': value} + existing_cache[title] = entry + + dump_cache() + + return entry['value'] diff --git a/stable-diffusion-webui/modules/call_queue.py b/stable-diffusion-webui/modules/call_queue.py new file mode 100755 index 0000000..d469107 --- /dev/null +++ b/stable-diffusion-webui/modules/call_queue.py @@ -0,0 +1,134 @@ +import os.path +from functools import wraps +import html +import time + +from modules import shared, progress, errors, devices, fifo_lock, profiling + +queue_lock = fifo_lock.FIFOLock() + + +def wrap_queued_call(func): + def f(*args, **kwargs): + with queue_lock: + res = func(*args, **kwargs) + + return res + + return f + + +def wrap_gradio_gpu_call(func, extra_outputs=None): + @wraps(func) + def f(*args, **kwargs): + + # if the first argument is a string that says "task(...)", it is treated as a job id + if args and type(args[0]) == str and args[0].startswith("task(") and args[0].endswith(")"): + id_task = args[0] + progress.add_task_to_queue(id_task) + else: + id_task = None + + with queue_lock: + shared.state.begin(job=id_task) + progress.start_task(id_task) + + try: + res = func(*args, **kwargs) + progress.record_results(id_task, res) + finally: + progress.finish_task(id_task) + + shared.state.end() + + return res + + return wrap_gradio_call(f, extra_outputs=extra_outputs, add_stats=True) + + +def wrap_gradio_call(func, extra_outputs=None, add_stats=False): + @wraps(func) + def f(*args, **kwargs): + try: + res = func(*args, **kwargs) + finally: + shared.state.skipped = False + shared.state.interrupted = False + shared.state.stopping_generation = False + shared.state.job_count = 0 + shared.state.job = "" + return res + + return wrap_gradio_call_no_job(f, extra_outputs, add_stats) + + +def wrap_gradio_call_no_job(func, extra_outputs=None, add_stats=False): + @wraps(func) + def f(*args, extra_outputs_array=extra_outputs, **kwargs): + run_memmon = shared.opts.memmon_poll_rate > 0 and not shared.mem_mon.disabled and add_stats + if run_memmon: + shared.mem_mon.monitor() + t = time.perf_counter() + + try: + res = list(func(*args, **kwargs)) + except Exception as e: + # When printing out our debug argument list, + # do not print out more than a 100 KB of text + max_debug_str_len = 131072 + message = "Error completing request" + arg_str = f"Arguments: {args} {kwargs}"[:max_debug_str_len] + if len(arg_str) > max_debug_str_len: + arg_str += f" (Argument list truncated at {max_debug_str_len}/{len(arg_str)} characters)" + errors.report(f"{message}\n{arg_str}", exc_info=True) + + if extra_outputs_array is None: + extra_outputs_array = [None, ''] + + error_message = f'{type(e).__name__}: {e}' + res = extra_outputs_array + [f"
    {html.escape(error_message)}
    "] + + devices.torch_gc() + + if not add_stats: + return tuple(res) + + elapsed = time.perf_counter() - t + elapsed_m = int(elapsed // 60) + elapsed_s = elapsed % 60 + elapsed_text = f"{elapsed_s:.1f} sec." + if elapsed_m > 0: + elapsed_text = f"{elapsed_m} min. "+elapsed_text + + if run_memmon: + mem_stats = {k: -(v//-(1024*1024)) for k, v in shared.mem_mon.stop().items()} + active_peak = mem_stats['active_peak'] + reserved_peak = mem_stats['reserved_peak'] + sys_peak = mem_stats['system_peak'] + sys_total = mem_stats['total'] + sys_pct = sys_peak/max(sys_total, 1) * 100 + + toltip_a = "Active: peak amount of video memory used during generation (excluding cached data)" + toltip_r = "Reserved: total amount of video memory allocated by the Torch library " + toltip_sys = "System: peak amount of video memory allocated by all running programs, out of total capacity" + + text_a = f"A: {active_peak/1024:.2f} GB" + text_r = f"R: {reserved_peak/1024:.2f} GB" + text_sys = f"Sys: {sys_peak/1024:.1f}/{sys_total/1024:g} GB ({sys_pct:.1f}%)" + + vram_html = f"

    {text_a}, {text_r}, {text_sys}

    " + else: + vram_html = '' + + if shared.opts.profiling_enable and os.path.exists(shared.opts.profiling_filename): + profiling_html = f"

    [ Profile ]

    " + else: + profiling_html = '' + + # last item is always HTML + res[-1] += f"

    Time taken: {elapsed_text}

    {vram_html}{profiling_html}
    " + + return tuple(res) + + return f + diff --git a/stable-diffusion-webui/modules/cmd_args.py b/stable-diffusion-webui/modules/cmd_args.py new file mode 100755 index 0000000..e826798 --- /dev/null +++ b/stable-diffusion-webui/modules/cmd_args.py @@ -0,0 +1,128 @@ +import argparse +import json +import os +from modules.paths_internal import normalized_filepath, models_path, script_path, data_path, extensions_dir, extensions_builtin_dir, sd_default_config, sd_model_file # noqa: F401 + +parser = argparse.ArgumentParser() + +parser.add_argument("-f", action='store_true', help=argparse.SUPPRESS) # allows running as root; implemented outside of webui +parser.add_argument("--update-all-extensions", action='store_true', help="launch.py argument: download updates for all extensions when starting the program") +parser.add_argument("--skip-python-version-check", action='store_true', help="launch.py argument: do not check python version") +parser.add_argument("--skip-torch-cuda-test", action='store_true', help="launch.py argument: do not check if CUDA is able to work properly") +parser.add_argument("--reinstall-xformers", action='store_true', help="launch.py argument: install the appropriate version of xformers even if you have some version already installed") +parser.add_argument("--reinstall-torch", action='store_true', help="launch.py argument: install the appropriate version of torch even if you have some version already installed") +parser.add_argument("--update-check", action='store_true', help="launch.py argument: check for updates at startup") +parser.add_argument("--test-server", action='store_true', help="launch.py argument: configure server for testing") +parser.add_argument("--log-startup", action='store_true', help="launch.py argument: print a detailed log of what's happening at startup") +parser.add_argument("--skip-prepare-environment", action='store_true', help="launch.py argument: skip all environment preparation") +parser.add_argument("--skip-install", action='store_true', help="launch.py argument: skip installation of packages") +parser.add_argument("--dump-sysinfo", action='store_true', help="launch.py argument: dump limited sysinfo file (without information about extensions, options) to disk and quit") +parser.add_argument("--loglevel", type=str, help="log level; one of: CRITICAL, ERROR, WARNING, INFO, DEBUG", default=None) +parser.add_argument("--do-not-download-clip", action='store_true', help="do not download CLIP model even if it's not included in the checkpoint") +parser.add_argument("--data-dir", type=normalized_filepath, default=os.path.dirname(os.path.dirname(os.path.realpath(__file__))), help="base path where all user data is stored") +parser.add_argument("--models-dir", type=normalized_filepath, default=None, help="base path where models are stored; overrides --data-dir") +parser.add_argument("--config", type=normalized_filepath, default=sd_default_config, help="path to config which constructs model",) +parser.add_argument("--ckpt", type=normalized_filepath, default=sd_model_file, help="path to checkpoint of stable diffusion model; if specified, this checkpoint will be added to the list of checkpoints and loaded",) +parser.add_argument("--ckpt-dir", type=normalized_filepath, default=None, help="Path to directory with stable diffusion checkpoints") +parser.add_argument("--vae-dir", type=normalized_filepath, default=None, help="Path to directory with VAE files") +parser.add_argument("--gfpgan-dir", type=normalized_filepath, help="GFPGAN directory", default=('./src/gfpgan' if os.path.exists('./src/gfpgan') else './GFPGAN')) +parser.add_argument("--gfpgan-model", type=normalized_filepath, help="GFPGAN model file name", default=None) +parser.add_argument("--no-half", action='store_true', help="do not switch the model to 16-bit floats") +parser.add_argument("--no-half-vae", action='store_true', help="do not switch the VAE model to 16-bit floats") +parser.add_argument("--no-progressbar-hiding", action='store_true', help="do not hide progressbar in gradio UI (we hide it because it slows down ML if you have hardware acceleration in browser)") +parser.add_argument("--max-batch-count", type=int, default=16, help="does not do anything") +parser.add_argument("--embeddings-dir", type=normalized_filepath, default=os.path.join(data_path, 'embeddings'), help="embeddings directory for textual inversion (default: embeddings)") +parser.add_argument("--textual-inversion-templates-dir", type=normalized_filepath, default=os.path.join(script_path, 'textual_inversion_templates'), help="directory with textual inversion templates") +parser.add_argument("--hypernetwork-dir", type=normalized_filepath, default=os.path.join(models_path, 'hypernetworks'), help="hypernetwork directory") +parser.add_argument("--localizations-dir", type=normalized_filepath, default=os.path.join(script_path, 'localizations'), help="localizations directory") +parser.add_argument("--allow-code", action='store_true', help="allow custom script execution from webui") +parser.add_argument("--medvram", action='store_true', help="enable stable diffusion model optimizations for sacrificing a little speed for low VRM usage") +parser.add_argument("--medvram-sdxl", action='store_true', help="enable --medvram optimization just for SDXL models") +parser.add_argument("--lowvram", action='store_true', help="enable stable diffusion model optimizations for sacrificing a lot of speed for very low VRM usage") +parser.add_argument("--lowram", action='store_true', help="load stable diffusion checkpoint weights to VRAM instead of RAM") +parser.add_argument("--always-batch-cond-uncond", action='store_true', help="does not do anything") +parser.add_argument("--unload-gfpgan", action='store_true', help="does not do anything.") +parser.add_argument("--precision", type=str, help="evaluate at this precision", choices=["full", "half", "autocast"], default="autocast") +parser.add_argument("--upcast-sampling", action='store_true', help="upcast sampling. No effect with --no-half. Usually produces similar results to --no-half with better performance while using less memory.") +parser.add_argument("--share", action='store_true', help="use share=True for gradio and make the UI accessible through their site") +parser.add_argument("--ngrok", type=str, help="ngrok authtoken, alternative to gradio --share", default=None) +parser.add_argument("--ngrok-region", type=str, help="does not do anything.", default="") +parser.add_argument("--ngrok-options", type=json.loads, help='The options to pass to ngrok in JSON format, e.g.: \'{"authtoken_from_env":true, "basic_auth":"user:password", "oauth_provider":"google", "oauth_allow_emails":"user@asdf.com"}\'', default=dict()) +parser.add_argument("--enable-insecure-extension-access", action='store_true', help="enable extensions tab regardless of other options") +parser.add_argument("--codeformer-models-path", type=normalized_filepath, help="Path to directory with codeformer model file(s).", default=os.path.join(models_path, 'Codeformer')) +parser.add_argument("--gfpgan-models-path", type=normalized_filepath, help="Path to directory with GFPGAN model file(s).", default=os.path.join(models_path, 'GFPGAN')) +parser.add_argument("--esrgan-models-path", type=normalized_filepath, help="Path to directory with ESRGAN model file(s).", default=os.path.join(models_path, 'ESRGAN')) +parser.add_argument("--bsrgan-models-path", type=normalized_filepath, help="Path to directory with BSRGAN model file(s).", default=os.path.join(models_path, 'BSRGAN')) +parser.add_argument("--realesrgan-models-path", type=normalized_filepath, help="Path to directory with RealESRGAN model file(s).", default=os.path.join(models_path, 'RealESRGAN')) +parser.add_argument("--dat-models-path", type=normalized_filepath, help="Path to directory with DAT model file(s).", default=os.path.join(models_path, 'DAT')) +parser.add_argument("--clip-models-path", type=normalized_filepath, help="Path to directory with CLIP model file(s).", default=None) +parser.add_argument("--xformers", action='store_true', help="enable xformers for cross attention layers") +parser.add_argument("--force-enable-xformers", action='store_true', help="enable xformers for cross attention layers regardless of whether the checking code thinks you can run it; do not make bug reports if this fails to work") +parser.add_argument("--xformers-flash-attention", action='store_true', help="enable xformers with Flash Attention to improve reproducibility (supported for SD2.x or variant only)") +parser.add_argument("--deepdanbooru", action='store_true', help="does not do anything") +parser.add_argument("--opt-split-attention", action='store_true', help="prefer Doggettx's cross-attention layer optimization for automatic choice of optimization") +parser.add_argument("--opt-sub-quad-attention", action='store_true', help="prefer memory efficient sub-quadratic cross-attention layer optimization for automatic choice of optimization") +parser.add_argument("--sub-quad-q-chunk-size", type=int, help="query chunk size for the sub-quadratic cross-attention layer optimization to use", default=1024) +parser.add_argument("--sub-quad-kv-chunk-size", type=int, help="kv chunk size for the sub-quadratic cross-attention layer optimization to use", default=None) +parser.add_argument("--sub-quad-chunk-threshold", type=int, help="the percentage of VRAM threshold for the sub-quadratic cross-attention layer optimization to use chunking", default=None) +parser.add_argument("--opt-split-attention-invokeai", action='store_true', help="prefer InvokeAI's cross-attention layer optimization for automatic choice of optimization") +parser.add_argument("--opt-split-attention-v1", action='store_true', help="prefer older version of split attention optimization for automatic choice of optimization") +parser.add_argument("--opt-sdp-attention", action='store_true', help="prefer scaled dot product cross-attention layer optimization for automatic choice of optimization; requires PyTorch 2.*") +parser.add_argument("--opt-sdp-no-mem-attention", action='store_true', help="prefer scaled dot product cross-attention layer optimization without memory efficient attention for automatic choice of optimization, makes image generation deterministic; requires PyTorch 2.*") +parser.add_argument("--disable-opt-split-attention", action='store_true', help="prefer no cross-attention layer optimization for automatic choice of optimization") +parser.add_argument("--disable-nan-check", action='store_true', help="do not check if produced images/latent spaces have nans; useful for running without a checkpoint in CI") +parser.add_argument("--use-cpu", nargs='+', help="use CPU as torch device for specified modules", default=[], type=str.lower) +parser.add_argument("--use-ipex", action="store_true", help="use Intel XPU as torch device") +parser.add_argument("--disable-model-loading-ram-optimization", action='store_true', help="disable an optimization that reduces RAM use when loading a model") +parser.add_argument("--listen", action='store_true', help="launch gradio with 0.0.0.0 as server name, allowing to respond to network requests") +parser.add_argument("--port", type=int, help="launch gradio with given server port, you need root/admin rights for ports < 1024, defaults to 7860 if available", default=None) +parser.add_argument("--show-negative-prompt", action='store_true', help="does not do anything", default=False) +parser.add_argument("--ui-config-file", type=str, help="filename to use for ui configuration", default=os.path.join(data_path, 'ui-config.json')) +parser.add_argument("--hide-ui-dir-config", action='store_true', help="hide directory configuration from webui", default=False) +parser.add_argument("--freeze-settings", action='store_true', help="disable editing of all settings globally", default=False) +parser.add_argument("--freeze-settings-in-sections", type=str, help='disable editing settings in specific sections of the settings page by specifying a comma-delimited list such like "saving-images,upscaling". The list of setting names can be found in the modules/shared_options.py file', default=None) +parser.add_argument("--freeze-specific-settings", type=str, help='disable editing of individual settings by specifying a comma-delimited list like "samples_save,samples_format". The list of setting names can be found in the config.json file', default=None) +parser.add_argument("--ui-settings-file", type=str, help="filename to use for ui settings", default=os.path.join(data_path, 'config.json')) +parser.add_argument("--gradio-debug", action='store_true', help="launch gradio with --debug option") +parser.add_argument("--gradio-auth", type=str, help='set gradio authentication like "username:password"; or comma-delimit multiple like "u1:p1,u2:p2,u3:p3"', default=None) +parser.add_argument("--gradio-auth-path", type=normalized_filepath, help='set gradio authentication file path ex. "/path/to/auth/file" same auth format as --gradio-auth', default=None) +parser.add_argument("--gradio-img2img-tool", type=str, help='does not do anything') +parser.add_argument("--gradio-inpaint-tool", type=str, help="does not do anything") +parser.add_argument("--gradio-allowed-path", action='append', help="add path to gradio's allowed_paths, make it possible to serve files from it", default=[data_path]) +parser.add_argument("--opt-channelslast", action='store_true', help="change memory type for stable diffusion to channels last") +parser.add_argument("--styles-file", type=str, action='append', help="path or wildcard path of styles files, allow multiple entries.", default=[]) +parser.add_argument("--autolaunch", action='store_true', help="open the webui URL in the system's default browser upon launch", default=False) +parser.add_argument("--theme", type=str, help="launches the UI with light or dark theme", default=None) +parser.add_argument("--use-textbox-seed", action='store_true', help="use textbox for seeds in UI (no up/down, but possible to input long seeds)", default=False) +parser.add_argument("--disable-console-progressbars", action='store_true', help="do not output progressbars to console", default=False) +parser.add_argument("--enable-console-prompts", action='store_true', help="does not do anything", default=False) # Legacy compatibility, use as default value shared.opts.enable_console_prompts +parser.add_argument('--vae-path', type=normalized_filepath, help='Checkpoint to use as VAE; setting this argument disables all settings related to VAE', default=None) +parser.add_argument("--disable-safe-unpickle", action='store_true', help="disable checking pytorch models for malicious code", default=False) +parser.add_argument("--api", action='store_true', help="use api=True to launch the API together with the webui (use --nowebui instead for only the API)") +parser.add_argument("--api-auth", type=str, help='Set authentication for API like "username:password"; or comma-delimit multiple like "u1:p1,u2:p2,u3:p3"', default=None) +parser.add_argument("--api-log", action='store_true', help="use api-log=True to enable logging of all API requests") +parser.add_argument("--nowebui", action='store_true', help="use api=True to launch the API instead of the webui") +parser.add_argument("--ui-debug-mode", action='store_true', help="Don't load model to quickly launch UI") +parser.add_argument("--device-id", type=str, help="Select the default CUDA device to use (export CUDA_VISIBLE_DEVICES=0,1,etc might be needed before)", default=None) +parser.add_argument("--administrator", action='store_true', help="Administrator rights", default=False) +parser.add_argument("--cors-allow-origins", type=str, help="Allowed CORS origin(s) in the form of a comma-separated list (no spaces)", default=None) +parser.add_argument("--cors-allow-origins-regex", type=str, help="Allowed CORS origin(s) in the form of a single regular expression", default=None) +parser.add_argument("--tls-keyfile", type=str, help="Partially enables TLS, requires --tls-certfile to fully function", default=None) +parser.add_argument("--tls-certfile", type=str, help="Partially enables TLS, requires --tls-keyfile to fully function", default=None) +parser.add_argument("--disable-tls-verify", action="store_false", help="When passed, enables the use of self-signed certificates.", default=None) +parser.add_argument("--server-name", type=str, help="Sets hostname of server", default=None) +parser.add_argument("--gradio-queue", action='store_true', help="does not do anything", default=True) +parser.add_argument("--no-gradio-queue", action='store_true', help="Disables gradio queue; causes the webpage to use http requests instead of websockets; was the default in earlier versions") +parser.add_argument("--skip-version-check", action='store_true', help="Do not check versions of torch and xformers") +parser.add_argument("--no-hashing", action='store_true', help="disable sha256 hashing of checkpoints to help loading performance", default=False) +parser.add_argument("--no-download-sd-model", action='store_true', help="don't download SD1.5 model even if no model is found in --ckpt-dir", default=False) +parser.add_argument('--subpath', type=str, help='customize the subpath for gradio, use with reverse proxy') +parser.add_argument('--add-stop-route', action='store_true', help='does not do anything') +parser.add_argument('--api-server-stop', action='store_true', help='enable server stop/restart/kill via api') +parser.add_argument('--timeout-keep-alive', type=int, default=30, help='set timeout_keep_alive for uvicorn') +parser.add_argument("--disable-all-extensions", action='store_true', help="prevent all extensions from running regardless of any other settings", default=False) +parser.add_argument("--disable-extra-extensions", action='store_true', help="prevent all extensions except built-in from running regardless of any other settings", default=False) +parser.add_argument("--skip-load-model-at-start", action='store_true', help="if load a model at web start, only take effect when --nowebui") +parser.add_argument("--unix-filenames-sanitization", action='store_true', help="allow any symbols except '/' in filenames. May conflict with your browser and file system") +parser.add_argument("--filenames-max-length", type=int, default=128, help='maximal length of filenames of saved images. If you override it, it can conflict with your file system') +parser.add_argument("--no-prompt-history", action='store_true', help="disable read prompt from last generation feature; settings this argument will not create '--data_path/params.txt' file") diff --git a/stable-diffusion-webui/modules/codeformer_model.py b/stable-diffusion-webui/modules/codeformer_model.py new file mode 100755 index 0000000..c167d79 --- /dev/null +++ b/stable-diffusion-webui/modules/codeformer_model.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +import logging + +import torch + +from modules import ( + devices, + errors, + face_restoration, + face_restoration_utils, + modelloader, + shared, +) + +logger = logging.getLogger(__name__) + +model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth' +model_download_name = 'codeformer-v0.1.0.pth' + +# used by e.g. postprocessing_codeformer.py +codeformer: face_restoration.FaceRestoration | None = None + + +class FaceRestorerCodeFormer(face_restoration_utils.CommonFaceRestoration): + def name(self): + return "CodeFormer" + + def load_net(self) -> torch.Module: + for model_path in modelloader.load_models( + model_path=self.model_path, + model_url=model_url, + command_path=self.model_path, + download_name=model_download_name, + ext_filter=['.pth'], + ): + return modelloader.load_spandrel_model( + model_path, + device=devices.device_codeformer, + expected_architecture='CodeFormer', + ).model + raise ValueError("No codeformer model found") + + def get_device(self): + return devices.device_codeformer + + def restore(self, np_image, w: float | None = None): + if w is None: + w = getattr(shared.opts, "code_former_weight", 0.5) + + def restore_face(cropped_face_t): + assert self.net is not None + return self.net(cropped_face_t, weight=w, adain=True)[0] + + return self.restore_with_helper(np_image, restore_face) + + +def setup_model(dirname: str) -> None: + global codeformer + try: + codeformer = FaceRestorerCodeFormer(dirname) + shared.face_restorers.append(codeformer) + except Exception: + errors.report("Error setting up CodeFormer", exc_info=True) diff --git a/stable-diffusion-webui/modules/config_states.py b/stable-diffusion-webui/modules/config_states.py new file mode 100755 index 0000000..651793c --- /dev/null +++ b/stable-diffusion-webui/modules/config_states.py @@ -0,0 +1,198 @@ +""" +Supports saving and restoring webui and extensions from a known working set of commits +""" + +import os +import json +import tqdm + +from datetime import datetime +import git + +from modules import shared, extensions, errors +from modules.paths_internal import script_path, config_states_dir + +all_config_states = {} + + +def list_config_states(): + global all_config_states + + all_config_states.clear() + os.makedirs(config_states_dir, exist_ok=True) + + config_states = [] + for filename in os.listdir(config_states_dir): + if filename.endswith(".json"): + path = os.path.join(config_states_dir, filename) + try: + with open(path, "r", encoding="utf-8") as f: + j = json.load(f) + assert "created_at" in j, '"created_at" does not exist' + j["filepath"] = path + config_states.append(j) + except Exception as e: + print(f'[ERROR]: Config states {path}, {e}') + + config_states = sorted(config_states, key=lambda cs: cs["created_at"], reverse=True) + + for cs in config_states: + timestamp = datetime.fromtimestamp(cs["created_at"]).strftime('%Y-%m-%d %H:%M:%S') + name = cs.get("name", "Config") + full_name = f"{name}: {timestamp}" + all_config_states[full_name] = cs + + return all_config_states + + +def get_webui_config(): + webui_repo = None + + try: + if os.path.exists(os.path.join(script_path, ".git")): + webui_repo = git.Repo(script_path) + except Exception: + errors.report(f"Error reading webui git info from {script_path}", exc_info=True) + + webui_remote = None + webui_commit_hash = None + webui_commit_date = None + webui_branch = None + if webui_repo and not webui_repo.bare: + try: + webui_remote = next(webui_repo.remote().urls, None) + head = webui_repo.head.commit + webui_commit_date = webui_repo.head.commit.committed_date + webui_commit_hash = head.hexsha + webui_branch = webui_repo.active_branch.name + + except Exception: + webui_remote = None + + return { + "remote": webui_remote, + "commit_hash": webui_commit_hash, + "commit_date": webui_commit_date, + "branch": webui_branch, + } + + +def get_extension_config(): + ext_config = {} + + for ext in extensions.extensions: + ext.read_info_from_repo() + + entry = { + "name": ext.name, + "path": ext.path, + "enabled": ext.enabled, + "is_builtin": ext.is_builtin, + "remote": ext.remote, + "commit_hash": ext.commit_hash, + "commit_date": ext.commit_date, + "branch": ext.branch, + "have_info_from_repo": ext.have_info_from_repo + } + + ext_config[ext.name] = entry + + return ext_config + + +def get_config(): + creation_time = datetime.now().timestamp() + webui_config = get_webui_config() + ext_config = get_extension_config() + + return { + "created_at": creation_time, + "webui": webui_config, + "extensions": ext_config + } + + +def restore_webui_config(config): + print("* Restoring webui state...") + + if "webui" not in config: + print("Error: No webui data saved to config") + return + + webui_config = config["webui"] + + if "commit_hash" not in webui_config: + print("Error: No commit saved to webui config") + return + + webui_commit_hash = webui_config.get("commit_hash", None) + webui_repo = None + + try: + if os.path.exists(os.path.join(script_path, ".git")): + webui_repo = git.Repo(script_path) + except Exception: + errors.report(f"Error reading webui git info from {script_path}", exc_info=True) + return + + try: + webui_repo.git.fetch(all=True) + webui_repo.git.reset(webui_commit_hash, hard=True) + print(f"* Restored webui to commit {webui_commit_hash}.") + except Exception: + errors.report(f"Error restoring webui to commit{webui_commit_hash}") + + +def restore_extension_config(config): + print("* Restoring extension state...") + + if "extensions" not in config: + print("Error: No extension data saved to config") + return + + ext_config = config["extensions"] + + results = [] + disabled = [] + + for ext in tqdm.tqdm(extensions.extensions): + if ext.is_builtin: + continue + + ext.read_info_from_repo() + current_commit = ext.commit_hash + + if ext.name not in ext_config: + ext.disabled = True + disabled.append(ext.name) + results.append((ext, current_commit[:8], False, "Saved extension state not found in config, marking as disabled")) + continue + + entry = ext_config[ext.name] + + if "commit_hash" in entry and entry["commit_hash"]: + try: + ext.fetch_and_reset_hard(entry["commit_hash"]) + ext.read_info_from_repo() + if current_commit != entry["commit_hash"]: + results.append((ext, current_commit[:8], True, entry["commit_hash"][:8])) + except Exception as ex: + results.append((ext, current_commit[:8], False, ex)) + else: + results.append((ext, current_commit[:8], False, "No commit hash found in config")) + + if not entry.get("enabled", False): + ext.disabled = True + disabled.append(ext.name) + else: + ext.disabled = False + + shared.opts.disabled_extensions = disabled + shared.opts.save(shared.config_filename) + + print("* Finished restoring extensions. Results:") + for ext, prev_commit, success, result in results: + if success: + print(f" + {ext.name}: {prev_commit} -> {result}") + else: + print(f" ! {ext.name}: FAILURE ({result})") diff --git a/stable-diffusion-webui/modules/dat_model.py b/stable-diffusion-webui/modules/dat_model.py new file mode 100755 index 0000000..495d5f4 --- /dev/null +++ b/stable-diffusion-webui/modules/dat_model.py @@ -0,0 +1,79 @@ +import os + +from modules import modelloader, errors +from modules.shared import cmd_opts, opts +from modules.upscaler import Upscaler, UpscalerData +from modules.upscaler_utils import upscale_with_model + + +class UpscalerDAT(Upscaler): + def __init__(self, user_path): + self.name = "DAT" + self.user_path = user_path + self.scalers = [] + super().__init__() + + for file in self.find_models(ext_filter=[".pt", ".pth"]): + name = modelloader.friendly_name(file) + scaler_data = UpscalerData(name, file, upscaler=self, scale=None) + self.scalers.append(scaler_data) + + for model in get_dat_models(self): + if model.name in opts.dat_enabled_models: + self.scalers.append(model) + + def do_upscale(self, img, path): + try: + info = self.load_model(path) + except Exception: + errors.report(f"Unable to load DAT model {path}", exc_info=True) + return img + + model_descriptor = modelloader.load_spandrel_model( + info.local_data_path, + device=self.device, + prefer_half=(not cmd_opts.no_half and not cmd_opts.upcast_sampling), + expected_architecture="DAT", + ) + return upscale_with_model( + model_descriptor, + img, + tile_size=opts.DAT_tile, + tile_overlap=opts.DAT_tile_overlap, + ) + + def load_model(self, path): + for scaler in self.scalers: + if scaler.data_path == path: + if scaler.local_data_path.startswith("http"): + scaler.local_data_path = modelloader.load_file_from_url( + scaler.data_path, + model_dir=self.model_download_path, + ) + if not os.path.exists(scaler.local_data_path): + raise FileNotFoundError(f"DAT data missing: {scaler.local_data_path}") + return scaler + raise ValueError(f"Unable to find model info: {path}") + + +def get_dat_models(scaler): + return [ + UpscalerData( + name="DAT x2", + path="https://github.com/n0kovo/dat_upscaler_models/raw/main/DAT/DAT_x2.pth", + scale=2, + upscaler=scaler, + ), + UpscalerData( + name="DAT x3", + path="https://github.com/n0kovo/dat_upscaler_models/raw/main/DAT/DAT_x3.pth", + scale=3, + upscaler=scaler, + ), + UpscalerData( + name="DAT x4", + path="https://github.com/n0kovo/dat_upscaler_models/raw/main/DAT/DAT_x4.pth", + scale=4, + upscaler=scaler, + ), + ] diff --git a/stable-diffusion-webui/modules/deepbooru.py b/stable-diffusion-webui/modules/deepbooru.py new file mode 100755 index 0000000..fb043fe --- /dev/null +++ b/stable-diffusion-webui/modules/deepbooru.py @@ -0,0 +1,98 @@ +import os +import re + +import torch +import numpy as np + +from modules import modelloader, paths, deepbooru_model, devices, images, shared + +re_special = re.compile(r'([\\()])') + + +class DeepDanbooru: + def __init__(self): + self.model = None + + def load(self): + if self.model is not None: + return + + files = modelloader.load_models( + model_path=os.path.join(paths.models_path, "torch_deepdanbooru"), + model_url='https://github.com/AUTOMATIC1111/TorchDeepDanbooru/releases/download/v1/model-resnet_custom_v3.pt', + ext_filter=[".pt"], + download_name='model-resnet_custom_v3.pt', + ) + + self.model = deepbooru_model.DeepDanbooruModel() + self.model.load_state_dict(torch.load(files[0], map_location="cpu")) + + self.model.eval() + self.model.to(devices.cpu, devices.dtype) + + def start(self): + self.load() + self.model.to(devices.device) + + def stop(self): + if not shared.opts.interrogate_keep_models_in_memory: + self.model.to(devices.cpu) + devices.torch_gc() + + def tag(self, pil_image): + self.start() + res = self.tag_multi(pil_image) + self.stop() + + return res + + def tag_multi(self, pil_image, force_disable_ranks=False): + threshold = shared.opts.interrogate_deepbooru_score_threshold + use_spaces = shared.opts.deepbooru_use_spaces + use_escape = shared.opts.deepbooru_escape + alpha_sort = shared.opts.deepbooru_sort_alpha + include_ranks = shared.opts.interrogate_return_ranks and not force_disable_ranks + + pic = images.resize_image(2, pil_image.convert("RGB"), 512, 512) + a = np.expand_dims(np.array(pic, dtype=np.float32), 0) / 255 + + with torch.no_grad(), devices.autocast(): + x = torch.from_numpy(a).to(devices.device, devices.dtype) + y = self.model(x)[0].detach().cpu().numpy() + + probability_dict = {} + + for tag, probability in zip(self.model.tags, y): + if probability < threshold: + continue + + if tag.startswith("rating:"): + continue + + probability_dict[tag] = probability + + if alpha_sort: + tags = sorted(probability_dict) + else: + tags = [tag for tag, _ in sorted(probability_dict.items(), key=lambda x: -x[1])] + + res = [] + + filtertags = {x.strip().replace(' ', '_') for x in shared.opts.deepbooru_filter_tags.split(",")} + + for tag in [x for x in tags if x not in filtertags]: + probability = probability_dict[tag] + tag_outformat = tag + if use_spaces: + tag_outformat = tag_outformat.replace('_', ' ') + if use_escape: + tag_outformat = re.sub(re_special, r'\\\1', tag_outformat) + if include_ranks: + tag_outformat = f"({tag_outformat}:{probability:.3f})" + + res.append(tag_outformat) + + return ", ".join(res) + + +model = DeepDanbooru() diff --git a/stable-diffusion-webui/modules/deepbooru_model.py b/stable-diffusion-webui/modules/deepbooru_model.py new file mode 100755 index 0000000..7a53884 --- /dev/null +++ b/stable-diffusion-webui/modules/deepbooru_model.py @@ -0,0 +1,678 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +from modules import devices + +# see https://github.com/AUTOMATIC1111/TorchDeepDanbooru for more + + +class DeepDanbooruModel(nn.Module): + def __init__(self): + super(DeepDanbooruModel, self).__init__() + + self.tags = [] + + self.n_Conv_0 = nn.Conv2d(kernel_size=(7, 7), in_channels=3, out_channels=64, stride=(2, 2)) + self.n_MaxPool_0 = nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 2)) + self.n_Conv_1 = nn.Conv2d(kernel_size=(1, 1), in_channels=64, out_channels=256) + self.n_Conv_2 = nn.Conv2d(kernel_size=(1, 1), in_channels=64, out_channels=64) + self.n_Conv_3 = nn.Conv2d(kernel_size=(3, 3), in_channels=64, out_channels=64) + self.n_Conv_4 = nn.Conv2d(kernel_size=(1, 1), in_channels=64, out_channels=256) + self.n_Conv_5 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=64) + self.n_Conv_6 = nn.Conv2d(kernel_size=(3, 3), in_channels=64, out_channels=64) + self.n_Conv_7 = nn.Conv2d(kernel_size=(1, 1), in_channels=64, out_channels=256) + self.n_Conv_8 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=64) + self.n_Conv_9 = nn.Conv2d(kernel_size=(3, 3), in_channels=64, out_channels=64) + self.n_Conv_10 = nn.Conv2d(kernel_size=(1, 1), in_channels=64, out_channels=256) + self.n_Conv_11 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=512, stride=(2, 2)) + self.n_Conv_12 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=128) + self.n_Conv_13 = nn.Conv2d(kernel_size=(3, 3), in_channels=128, out_channels=128, stride=(2, 2)) + self.n_Conv_14 = nn.Conv2d(kernel_size=(1, 1), in_channels=128, out_channels=512) + self.n_Conv_15 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=128) + self.n_Conv_16 = nn.Conv2d(kernel_size=(3, 3), in_channels=128, out_channels=128) + self.n_Conv_17 = nn.Conv2d(kernel_size=(1, 1), in_channels=128, out_channels=512) + self.n_Conv_18 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=128) + self.n_Conv_19 = nn.Conv2d(kernel_size=(3, 3), in_channels=128, out_channels=128) + self.n_Conv_20 = nn.Conv2d(kernel_size=(1, 1), in_channels=128, out_channels=512) + self.n_Conv_21 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=128) + self.n_Conv_22 = nn.Conv2d(kernel_size=(3, 3), in_channels=128, out_channels=128) + self.n_Conv_23 = nn.Conv2d(kernel_size=(1, 1), in_channels=128, out_channels=512) + self.n_Conv_24 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=128) + self.n_Conv_25 = nn.Conv2d(kernel_size=(3, 3), in_channels=128, out_channels=128) + self.n_Conv_26 = nn.Conv2d(kernel_size=(1, 1), in_channels=128, out_channels=512) + self.n_Conv_27 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=128) + self.n_Conv_28 = nn.Conv2d(kernel_size=(3, 3), in_channels=128, out_channels=128) + self.n_Conv_29 = nn.Conv2d(kernel_size=(1, 1), in_channels=128, out_channels=512) + self.n_Conv_30 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=128) + self.n_Conv_31 = nn.Conv2d(kernel_size=(3, 3), in_channels=128, out_channels=128) + self.n_Conv_32 = nn.Conv2d(kernel_size=(1, 1), in_channels=128, out_channels=512) + self.n_Conv_33 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=128) + self.n_Conv_34 = nn.Conv2d(kernel_size=(3, 3), in_channels=128, out_channels=128) + self.n_Conv_35 = nn.Conv2d(kernel_size=(1, 1), in_channels=128, out_channels=512) + self.n_Conv_36 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=1024, stride=(2, 2)) + self.n_Conv_37 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=256) + self.n_Conv_38 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256, stride=(2, 2)) + self.n_Conv_39 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_40 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_41 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_42 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_43 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_44 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_45 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_46 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_47 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_48 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_49 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_50 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_51 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_52 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_53 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_54 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_55 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_56 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_57 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_58 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_59 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_60 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_61 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_62 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_63 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_64 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_65 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_66 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_67 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_68 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_69 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_70 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_71 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_72 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_73 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_74 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_75 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_76 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_77 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_78 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_79 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_80 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_81 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_82 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_83 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_84 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_85 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_86 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_87 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_88 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_89 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_90 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_91 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_92 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_93 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_94 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_95 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_96 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_97 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_98 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256, stride=(2, 2)) + self.n_Conv_99 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_100 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=1024, stride=(2, 2)) + self.n_Conv_101 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_102 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_103 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_104 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_105 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_106 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_107 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_108 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_109 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_110 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_111 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_112 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_113 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_114 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_115 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_116 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_117 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_118 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_119 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_120 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_121 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_122 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_123 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_124 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_125 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_126 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_127 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_128 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_129 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_130 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_131 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_132 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_133 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_134 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_135 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_136 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_137 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_138 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_139 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_140 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_141 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_142 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_143 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_144 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_145 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_146 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_147 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_148 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_149 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_150 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_151 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_152 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_153 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_154 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_155 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=256) + self.n_Conv_156 = nn.Conv2d(kernel_size=(3, 3), in_channels=256, out_channels=256) + self.n_Conv_157 = nn.Conv2d(kernel_size=(1, 1), in_channels=256, out_channels=1024) + self.n_Conv_158 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=2048, stride=(2, 2)) + self.n_Conv_159 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=512) + self.n_Conv_160 = nn.Conv2d(kernel_size=(3, 3), in_channels=512, out_channels=512, stride=(2, 2)) + self.n_Conv_161 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=2048) + self.n_Conv_162 = nn.Conv2d(kernel_size=(1, 1), in_channels=2048, out_channels=512) + self.n_Conv_163 = nn.Conv2d(kernel_size=(3, 3), in_channels=512, out_channels=512) + self.n_Conv_164 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=2048) + self.n_Conv_165 = nn.Conv2d(kernel_size=(1, 1), in_channels=2048, out_channels=512) + self.n_Conv_166 = nn.Conv2d(kernel_size=(3, 3), in_channels=512, out_channels=512) + self.n_Conv_167 = nn.Conv2d(kernel_size=(1, 1), in_channels=512, out_channels=2048) + self.n_Conv_168 = nn.Conv2d(kernel_size=(1, 1), in_channels=2048, out_channels=4096, stride=(2, 2)) + self.n_Conv_169 = nn.Conv2d(kernel_size=(1, 1), in_channels=2048, out_channels=1024) + self.n_Conv_170 = nn.Conv2d(kernel_size=(3, 3), in_channels=1024, out_channels=1024, stride=(2, 2)) + self.n_Conv_171 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=4096) + self.n_Conv_172 = nn.Conv2d(kernel_size=(1, 1), in_channels=4096, out_channels=1024) + self.n_Conv_173 = nn.Conv2d(kernel_size=(3, 3), in_channels=1024, out_channels=1024) + self.n_Conv_174 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=4096) + self.n_Conv_175 = nn.Conv2d(kernel_size=(1, 1), in_channels=4096, out_channels=1024) + self.n_Conv_176 = nn.Conv2d(kernel_size=(3, 3), in_channels=1024, out_channels=1024) + self.n_Conv_177 = nn.Conv2d(kernel_size=(1, 1), in_channels=1024, out_channels=4096) + self.n_Conv_178 = nn.Conv2d(kernel_size=(1, 1), in_channels=4096, out_channels=9176, bias=False) + + def forward(self, *inputs): + t_358, = inputs + t_359 = t_358.permute(*[0, 3, 1, 2]) + t_359_padded = F.pad(t_359, [2, 3, 2, 3], value=0) + t_360 = self.n_Conv_0(t_359_padded.to(self.n_Conv_0.bias.dtype) if devices.unet_needs_upcast else t_359_padded) + t_361 = F.relu(t_360) + t_361 = F.pad(t_361, [0, 1, 0, 1], value=float('-inf')) + t_362 = self.n_MaxPool_0(t_361) + t_363 = self.n_Conv_1(t_362) + t_364 = self.n_Conv_2(t_362) + t_365 = F.relu(t_364) + t_365_padded = F.pad(t_365, [1, 1, 1, 1], value=0) + t_366 = self.n_Conv_3(t_365_padded) + t_367 = F.relu(t_366) + t_368 = self.n_Conv_4(t_367) + t_369 = torch.add(t_368, t_363) + t_370 = F.relu(t_369) + t_371 = self.n_Conv_5(t_370) + t_372 = F.relu(t_371) + t_372_padded = F.pad(t_372, [1, 1, 1, 1], value=0) + t_373 = self.n_Conv_6(t_372_padded) + t_374 = F.relu(t_373) + t_375 = self.n_Conv_7(t_374) + t_376 = torch.add(t_375, t_370) + t_377 = F.relu(t_376) + t_378 = self.n_Conv_8(t_377) + t_379 = F.relu(t_378) + t_379_padded = F.pad(t_379, [1, 1, 1, 1], value=0) + t_380 = self.n_Conv_9(t_379_padded) + t_381 = F.relu(t_380) + t_382 = self.n_Conv_10(t_381) + t_383 = torch.add(t_382, t_377) + t_384 = F.relu(t_383) + t_385 = self.n_Conv_11(t_384) + t_386 = self.n_Conv_12(t_384) + t_387 = F.relu(t_386) + t_387_padded = F.pad(t_387, [0, 1, 0, 1], value=0) + t_388 = self.n_Conv_13(t_387_padded) + t_389 = F.relu(t_388) + t_390 = self.n_Conv_14(t_389) + t_391 = torch.add(t_390, t_385) + t_392 = F.relu(t_391) + t_393 = self.n_Conv_15(t_392) + t_394 = F.relu(t_393) + t_394_padded = F.pad(t_394, [1, 1, 1, 1], value=0) + t_395 = self.n_Conv_16(t_394_padded) + t_396 = F.relu(t_395) + t_397 = self.n_Conv_17(t_396) + t_398 = torch.add(t_397, t_392) + t_399 = F.relu(t_398) + t_400 = self.n_Conv_18(t_399) + t_401 = F.relu(t_400) + t_401_padded = F.pad(t_401, [1, 1, 1, 1], value=0) + t_402 = self.n_Conv_19(t_401_padded) + t_403 = F.relu(t_402) + t_404 = self.n_Conv_20(t_403) + t_405 = torch.add(t_404, t_399) + t_406 = F.relu(t_405) + t_407 = self.n_Conv_21(t_406) + t_408 = F.relu(t_407) + t_408_padded = F.pad(t_408, [1, 1, 1, 1], value=0) + t_409 = self.n_Conv_22(t_408_padded) + t_410 = F.relu(t_409) + t_411 = self.n_Conv_23(t_410) + t_412 = torch.add(t_411, t_406) + t_413 = F.relu(t_412) + t_414 = self.n_Conv_24(t_413) + t_415 = F.relu(t_414) + t_415_padded = F.pad(t_415, [1, 1, 1, 1], value=0) + t_416 = self.n_Conv_25(t_415_padded) + t_417 = F.relu(t_416) + t_418 = self.n_Conv_26(t_417) + t_419 = torch.add(t_418, t_413) + t_420 = F.relu(t_419) + t_421 = self.n_Conv_27(t_420) + t_422 = F.relu(t_421) + t_422_padded = F.pad(t_422, [1, 1, 1, 1], value=0) + t_423 = self.n_Conv_28(t_422_padded) + t_424 = F.relu(t_423) + t_425 = self.n_Conv_29(t_424) + t_426 = torch.add(t_425, t_420) + t_427 = F.relu(t_426) + t_428 = self.n_Conv_30(t_427) + t_429 = F.relu(t_428) + t_429_padded = F.pad(t_429, [1, 1, 1, 1], value=0) + t_430 = self.n_Conv_31(t_429_padded) + t_431 = F.relu(t_430) + t_432 = self.n_Conv_32(t_431) + t_433 = torch.add(t_432, t_427) + t_434 = F.relu(t_433) + t_435 = self.n_Conv_33(t_434) + t_436 = F.relu(t_435) + t_436_padded = F.pad(t_436, [1, 1, 1, 1], value=0) + t_437 = self.n_Conv_34(t_436_padded) + t_438 = F.relu(t_437) + t_439 = self.n_Conv_35(t_438) + t_440 = torch.add(t_439, t_434) + t_441 = F.relu(t_440) + t_442 = self.n_Conv_36(t_441) + t_443 = self.n_Conv_37(t_441) + t_444 = F.relu(t_443) + t_444_padded = F.pad(t_444, [0, 1, 0, 1], value=0) + t_445 = self.n_Conv_38(t_444_padded) + t_446 = F.relu(t_445) + t_447 = self.n_Conv_39(t_446) + t_448 = torch.add(t_447, t_442) + t_449 = F.relu(t_448) + t_450 = self.n_Conv_40(t_449) + t_451 = F.relu(t_450) + t_451_padded = F.pad(t_451, [1, 1, 1, 1], value=0) + t_452 = self.n_Conv_41(t_451_padded) + t_453 = F.relu(t_452) + t_454 = self.n_Conv_42(t_453) + t_455 = torch.add(t_454, t_449) + t_456 = F.relu(t_455) + t_457 = self.n_Conv_43(t_456) + t_458 = F.relu(t_457) + t_458_padded = F.pad(t_458, [1, 1, 1, 1], value=0) + t_459 = self.n_Conv_44(t_458_padded) + t_460 = F.relu(t_459) + t_461 = self.n_Conv_45(t_460) + t_462 = torch.add(t_461, t_456) + t_463 = F.relu(t_462) + t_464 = self.n_Conv_46(t_463) + t_465 = F.relu(t_464) + t_465_padded = F.pad(t_465, [1, 1, 1, 1], value=0) + t_466 = self.n_Conv_47(t_465_padded) + t_467 = F.relu(t_466) + t_468 = self.n_Conv_48(t_467) + t_469 = torch.add(t_468, t_463) + t_470 = F.relu(t_469) + t_471 = self.n_Conv_49(t_470) + t_472 = F.relu(t_471) + t_472_padded = F.pad(t_472, [1, 1, 1, 1], value=0) + t_473 = self.n_Conv_50(t_472_padded) + t_474 = F.relu(t_473) + t_475 = self.n_Conv_51(t_474) + t_476 = torch.add(t_475, t_470) + t_477 = F.relu(t_476) + t_478 = self.n_Conv_52(t_477) + t_479 = F.relu(t_478) + t_479_padded = F.pad(t_479, [1, 1, 1, 1], value=0) + t_480 = self.n_Conv_53(t_479_padded) + t_481 = F.relu(t_480) + t_482 = self.n_Conv_54(t_481) + t_483 = torch.add(t_482, t_477) + t_484 = F.relu(t_483) + t_485 = self.n_Conv_55(t_484) + t_486 = F.relu(t_485) + t_486_padded = F.pad(t_486, [1, 1, 1, 1], value=0) + t_487 = self.n_Conv_56(t_486_padded) + t_488 = F.relu(t_487) + t_489 = self.n_Conv_57(t_488) + t_490 = torch.add(t_489, t_484) + t_491 = F.relu(t_490) + t_492 = self.n_Conv_58(t_491) + t_493 = F.relu(t_492) + t_493_padded = F.pad(t_493, [1, 1, 1, 1], value=0) + t_494 = self.n_Conv_59(t_493_padded) + t_495 = F.relu(t_494) + t_496 = self.n_Conv_60(t_495) + t_497 = torch.add(t_496, t_491) + t_498 = F.relu(t_497) + t_499 = self.n_Conv_61(t_498) + t_500 = F.relu(t_499) + t_500_padded = F.pad(t_500, [1, 1, 1, 1], value=0) + t_501 = self.n_Conv_62(t_500_padded) + t_502 = F.relu(t_501) + t_503 = self.n_Conv_63(t_502) + t_504 = torch.add(t_503, t_498) + t_505 = F.relu(t_504) + t_506 = self.n_Conv_64(t_505) + t_507 = F.relu(t_506) + t_507_padded = F.pad(t_507, [1, 1, 1, 1], value=0) + t_508 = self.n_Conv_65(t_507_padded) + t_509 = F.relu(t_508) + t_510 = self.n_Conv_66(t_509) + t_511 = torch.add(t_510, t_505) + t_512 = F.relu(t_511) + t_513 = self.n_Conv_67(t_512) + t_514 = F.relu(t_513) + t_514_padded = F.pad(t_514, [1, 1, 1, 1], value=0) + t_515 = self.n_Conv_68(t_514_padded) + t_516 = F.relu(t_515) + t_517 = self.n_Conv_69(t_516) + t_518 = torch.add(t_517, t_512) + t_519 = F.relu(t_518) + t_520 = self.n_Conv_70(t_519) + t_521 = F.relu(t_520) + t_521_padded = F.pad(t_521, [1, 1, 1, 1], value=0) + t_522 = self.n_Conv_71(t_521_padded) + t_523 = F.relu(t_522) + t_524 = self.n_Conv_72(t_523) + t_525 = torch.add(t_524, t_519) + t_526 = F.relu(t_525) + t_527 = self.n_Conv_73(t_526) + t_528 = F.relu(t_527) + t_528_padded = F.pad(t_528, [1, 1, 1, 1], value=0) + t_529 = self.n_Conv_74(t_528_padded) + t_530 = F.relu(t_529) + t_531 = self.n_Conv_75(t_530) + t_532 = torch.add(t_531, t_526) + t_533 = F.relu(t_532) + t_534 = self.n_Conv_76(t_533) + t_535 = F.relu(t_534) + t_535_padded = F.pad(t_535, [1, 1, 1, 1], value=0) + t_536 = self.n_Conv_77(t_535_padded) + t_537 = F.relu(t_536) + t_538 = self.n_Conv_78(t_537) + t_539 = torch.add(t_538, t_533) + t_540 = F.relu(t_539) + t_541 = self.n_Conv_79(t_540) + t_542 = F.relu(t_541) + t_542_padded = F.pad(t_542, [1, 1, 1, 1], value=0) + t_543 = self.n_Conv_80(t_542_padded) + t_544 = F.relu(t_543) + t_545 = self.n_Conv_81(t_544) + t_546 = torch.add(t_545, t_540) + t_547 = F.relu(t_546) + t_548 = self.n_Conv_82(t_547) + t_549 = F.relu(t_548) + t_549_padded = F.pad(t_549, [1, 1, 1, 1], value=0) + t_550 = self.n_Conv_83(t_549_padded) + t_551 = F.relu(t_550) + t_552 = self.n_Conv_84(t_551) + t_553 = torch.add(t_552, t_547) + t_554 = F.relu(t_553) + t_555 = self.n_Conv_85(t_554) + t_556 = F.relu(t_555) + t_556_padded = F.pad(t_556, [1, 1, 1, 1], value=0) + t_557 = self.n_Conv_86(t_556_padded) + t_558 = F.relu(t_557) + t_559 = self.n_Conv_87(t_558) + t_560 = torch.add(t_559, t_554) + t_561 = F.relu(t_560) + t_562 = self.n_Conv_88(t_561) + t_563 = F.relu(t_562) + t_563_padded = F.pad(t_563, [1, 1, 1, 1], value=0) + t_564 = self.n_Conv_89(t_563_padded) + t_565 = F.relu(t_564) + t_566 = self.n_Conv_90(t_565) + t_567 = torch.add(t_566, t_561) + t_568 = F.relu(t_567) + t_569 = self.n_Conv_91(t_568) + t_570 = F.relu(t_569) + t_570_padded = F.pad(t_570, [1, 1, 1, 1], value=0) + t_571 = self.n_Conv_92(t_570_padded) + t_572 = F.relu(t_571) + t_573 = self.n_Conv_93(t_572) + t_574 = torch.add(t_573, t_568) + t_575 = F.relu(t_574) + t_576 = self.n_Conv_94(t_575) + t_577 = F.relu(t_576) + t_577_padded = F.pad(t_577, [1, 1, 1, 1], value=0) + t_578 = self.n_Conv_95(t_577_padded) + t_579 = F.relu(t_578) + t_580 = self.n_Conv_96(t_579) + t_581 = torch.add(t_580, t_575) + t_582 = F.relu(t_581) + t_583 = self.n_Conv_97(t_582) + t_584 = F.relu(t_583) + t_584_padded = F.pad(t_584, [0, 1, 0, 1], value=0) + t_585 = self.n_Conv_98(t_584_padded) + t_586 = F.relu(t_585) + t_587 = self.n_Conv_99(t_586) + t_588 = self.n_Conv_100(t_582) + t_589 = torch.add(t_587, t_588) + t_590 = F.relu(t_589) + t_591 = self.n_Conv_101(t_590) + t_592 = F.relu(t_591) + t_592_padded = F.pad(t_592, [1, 1, 1, 1], value=0) + t_593 = self.n_Conv_102(t_592_padded) + t_594 = F.relu(t_593) + t_595 = self.n_Conv_103(t_594) + t_596 = torch.add(t_595, t_590) + t_597 = F.relu(t_596) + t_598 = self.n_Conv_104(t_597) + t_599 = F.relu(t_598) + t_599_padded = F.pad(t_599, [1, 1, 1, 1], value=0) + t_600 = self.n_Conv_105(t_599_padded) + t_601 = F.relu(t_600) + t_602 = self.n_Conv_106(t_601) + t_603 = torch.add(t_602, t_597) + t_604 = F.relu(t_603) + t_605 = self.n_Conv_107(t_604) + t_606 = F.relu(t_605) + t_606_padded = F.pad(t_606, [1, 1, 1, 1], value=0) + t_607 = self.n_Conv_108(t_606_padded) + t_608 = F.relu(t_607) + t_609 = self.n_Conv_109(t_608) + t_610 = torch.add(t_609, t_604) + t_611 = F.relu(t_610) + t_612 = self.n_Conv_110(t_611) + t_613 = F.relu(t_612) + t_613_padded = F.pad(t_613, [1, 1, 1, 1], value=0) + t_614 = self.n_Conv_111(t_613_padded) + t_615 = F.relu(t_614) + t_616 = self.n_Conv_112(t_615) + t_617 = torch.add(t_616, t_611) + t_618 = F.relu(t_617) + t_619 = self.n_Conv_113(t_618) + t_620 = F.relu(t_619) + t_620_padded = F.pad(t_620, [1, 1, 1, 1], value=0) + t_621 = self.n_Conv_114(t_620_padded) + t_622 = F.relu(t_621) + t_623 = self.n_Conv_115(t_622) + t_624 = torch.add(t_623, t_618) + t_625 = F.relu(t_624) + t_626 = self.n_Conv_116(t_625) + t_627 = F.relu(t_626) + t_627_padded = F.pad(t_627, [1, 1, 1, 1], value=0) + t_628 = self.n_Conv_117(t_627_padded) + t_629 = F.relu(t_628) + t_630 = self.n_Conv_118(t_629) + t_631 = torch.add(t_630, t_625) + t_632 = F.relu(t_631) + t_633 = self.n_Conv_119(t_632) + t_634 = F.relu(t_633) + t_634_padded = F.pad(t_634, [1, 1, 1, 1], value=0) + t_635 = self.n_Conv_120(t_634_padded) + t_636 = F.relu(t_635) + t_637 = self.n_Conv_121(t_636) + t_638 = torch.add(t_637, t_632) + t_639 = F.relu(t_638) + t_640 = self.n_Conv_122(t_639) + t_641 = F.relu(t_640) + t_641_padded = F.pad(t_641, [1, 1, 1, 1], value=0) + t_642 = self.n_Conv_123(t_641_padded) + t_643 = F.relu(t_642) + t_644 = self.n_Conv_124(t_643) + t_645 = torch.add(t_644, t_639) + t_646 = F.relu(t_645) + t_647 = self.n_Conv_125(t_646) + t_648 = F.relu(t_647) + t_648_padded = F.pad(t_648, [1, 1, 1, 1], value=0) + t_649 = self.n_Conv_126(t_648_padded) + t_650 = F.relu(t_649) + t_651 = self.n_Conv_127(t_650) + t_652 = torch.add(t_651, t_646) + t_653 = F.relu(t_652) + t_654 = self.n_Conv_128(t_653) + t_655 = F.relu(t_654) + t_655_padded = F.pad(t_655, [1, 1, 1, 1], value=0) + t_656 = self.n_Conv_129(t_655_padded) + t_657 = F.relu(t_656) + t_658 = self.n_Conv_130(t_657) + t_659 = torch.add(t_658, t_653) + t_660 = F.relu(t_659) + t_661 = self.n_Conv_131(t_660) + t_662 = F.relu(t_661) + t_662_padded = F.pad(t_662, [1, 1, 1, 1], value=0) + t_663 = self.n_Conv_132(t_662_padded) + t_664 = F.relu(t_663) + t_665 = self.n_Conv_133(t_664) + t_666 = torch.add(t_665, t_660) + t_667 = F.relu(t_666) + t_668 = self.n_Conv_134(t_667) + t_669 = F.relu(t_668) + t_669_padded = F.pad(t_669, [1, 1, 1, 1], value=0) + t_670 = self.n_Conv_135(t_669_padded) + t_671 = F.relu(t_670) + t_672 = self.n_Conv_136(t_671) + t_673 = torch.add(t_672, t_667) + t_674 = F.relu(t_673) + t_675 = self.n_Conv_137(t_674) + t_676 = F.relu(t_675) + t_676_padded = F.pad(t_676, [1, 1, 1, 1], value=0) + t_677 = self.n_Conv_138(t_676_padded) + t_678 = F.relu(t_677) + t_679 = self.n_Conv_139(t_678) + t_680 = torch.add(t_679, t_674) + t_681 = F.relu(t_680) + t_682 = self.n_Conv_140(t_681) + t_683 = F.relu(t_682) + t_683_padded = F.pad(t_683, [1, 1, 1, 1], value=0) + t_684 = self.n_Conv_141(t_683_padded) + t_685 = F.relu(t_684) + t_686 = self.n_Conv_142(t_685) + t_687 = torch.add(t_686, t_681) + t_688 = F.relu(t_687) + t_689 = self.n_Conv_143(t_688) + t_690 = F.relu(t_689) + t_690_padded = F.pad(t_690, [1, 1, 1, 1], value=0) + t_691 = self.n_Conv_144(t_690_padded) + t_692 = F.relu(t_691) + t_693 = self.n_Conv_145(t_692) + t_694 = torch.add(t_693, t_688) + t_695 = F.relu(t_694) + t_696 = self.n_Conv_146(t_695) + t_697 = F.relu(t_696) + t_697_padded = F.pad(t_697, [1, 1, 1, 1], value=0) + t_698 = self.n_Conv_147(t_697_padded) + t_699 = F.relu(t_698) + t_700 = self.n_Conv_148(t_699) + t_701 = torch.add(t_700, t_695) + t_702 = F.relu(t_701) + t_703 = self.n_Conv_149(t_702) + t_704 = F.relu(t_703) + t_704_padded = F.pad(t_704, [1, 1, 1, 1], value=0) + t_705 = self.n_Conv_150(t_704_padded) + t_706 = F.relu(t_705) + t_707 = self.n_Conv_151(t_706) + t_708 = torch.add(t_707, t_702) + t_709 = F.relu(t_708) + t_710 = self.n_Conv_152(t_709) + t_711 = F.relu(t_710) + t_711_padded = F.pad(t_711, [1, 1, 1, 1], value=0) + t_712 = self.n_Conv_153(t_711_padded) + t_713 = F.relu(t_712) + t_714 = self.n_Conv_154(t_713) + t_715 = torch.add(t_714, t_709) + t_716 = F.relu(t_715) + t_717 = self.n_Conv_155(t_716) + t_718 = F.relu(t_717) + t_718_padded = F.pad(t_718, [1, 1, 1, 1], value=0) + t_719 = self.n_Conv_156(t_718_padded) + t_720 = F.relu(t_719) + t_721 = self.n_Conv_157(t_720) + t_722 = torch.add(t_721, t_716) + t_723 = F.relu(t_722) + t_724 = self.n_Conv_158(t_723) + t_725 = self.n_Conv_159(t_723) + t_726 = F.relu(t_725) + t_726_padded = F.pad(t_726, [0, 1, 0, 1], value=0) + t_727 = self.n_Conv_160(t_726_padded) + t_728 = F.relu(t_727) + t_729 = self.n_Conv_161(t_728) + t_730 = torch.add(t_729, t_724) + t_731 = F.relu(t_730) + t_732 = self.n_Conv_162(t_731) + t_733 = F.relu(t_732) + t_733_padded = F.pad(t_733, [1, 1, 1, 1], value=0) + t_734 = self.n_Conv_163(t_733_padded) + t_735 = F.relu(t_734) + t_736 = self.n_Conv_164(t_735) + t_737 = torch.add(t_736, t_731) + t_738 = F.relu(t_737) + t_739 = self.n_Conv_165(t_738) + t_740 = F.relu(t_739) + t_740_padded = F.pad(t_740, [1, 1, 1, 1], value=0) + t_741 = self.n_Conv_166(t_740_padded) + t_742 = F.relu(t_741) + t_743 = self.n_Conv_167(t_742) + t_744 = torch.add(t_743, t_738) + t_745 = F.relu(t_744) + t_746 = self.n_Conv_168(t_745) + t_747 = self.n_Conv_169(t_745) + t_748 = F.relu(t_747) + t_748_padded = F.pad(t_748, [0, 1, 0, 1], value=0) + t_749 = self.n_Conv_170(t_748_padded) + t_750 = F.relu(t_749) + t_751 = self.n_Conv_171(t_750) + t_752 = torch.add(t_751, t_746) + t_753 = F.relu(t_752) + t_754 = self.n_Conv_172(t_753) + t_755 = F.relu(t_754) + t_755_padded = F.pad(t_755, [1, 1, 1, 1], value=0) + t_756 = self.n_Conv_173(t_755_padded) + t_757 = F.relu(t_756) + t_758 = self.n_Conv_174(t_757) + t_759 = torch.add(t_758, t_753) + t_760 = F.relu(t_759) + t_761 = self.n_Conv_175(t_760) + t_762 = F.relu(t_761) + t_762_padded = F.pad(t_762, [1, 1, 1, 1], value=0) + t_763 = self.n_Conv_176(t_762_padded) + t_764 = F.relu(t_763) + t_765 = self.n_Conv_177(t_764) + t_766 = torch.add(t_765, t_760) + t_767 = F.relu(t_766) + t_768 = self.n_Conv_178(t_767) + t_769 = F.avg_pool2d(t_768, kernel_size=t_768.shape[-2:]) + t_770 = torch.squeeze(t_769, 3) + t_770 = torch.squeeze(t_770, 2) + t_771 = torch.sigmoid(t_770) + return t_771 + + def load_state_dict(self, state_dict, **kwargs): + self.tags = state_dict.get('tags', []) + + super(DeepDanbooruModel, self).load_state_dict({k: v for k, v in state_dict.items() if k != 'tags'}) + diff --git a/stable-diffusion-webui/modules/devices.py b/stable-diffusion-webui/modules/devices.py new file mode 100755 index 0000000..ee67914 --- /dev/null +++ b/stable-diffusion-webui/modules/devices.py @@ -0,0 +1,295 @@ +import sys +import contextlib +from functools import lru_cache + +import torch +from modules import errors, shared, npu_specific + +if sys.platform == "darwin": + from modules import mac_specific + +if shared.cmd_opts.use_ipex: + from modules import xpu_specific + + +def has_xpu() -> bool: + return shared.cmd_opts.use_ipex and xpu_specific.has_xpu + + +def has_mps() -> bool: + if sys.platform != "darwin": + return False + else: + return mac_specific.has_mps + + +def cuda_no_autocast(device_id=None) -> bool: + if device_id is None: + device_id = get_cuda_device_id() + return ( + torch.cuda.get_device_capability(device_id) == (7, 5) + and torch.cuda.get_device_name(device_id).startswith("NVIDIA GeForce GTX 16") + ) + + +def get_cuda_device_id(): + return ( + int(shared.cmd_opts.device_id) + if shared.cmd_opts.device_id is not None and shared.cmd_opts.device_id.isdigit() + else 0 + ) or torch.cuda.current_device() + + +def get_cuda_device_string(): + if shared.cmd_opts.device_id is not None: + return f"cuda:{shared.cmd_opts.device_id}" + + return "cuda" + + +def get_optimal_device_name(): + if torch.cuda.is_available(): + return get_cuda_device_string() + + if has_mps(): + return "mps" + + if has_xpu(): + return xpu_specific.get_xpu_device_string() + + if npu_specific.has_npu: + return npu_specific.get_npu_device_string() + + return "cpu" + + +def get_optimal_device(): + return torch.device(get_optimal_device_name()) + + +def get_device_for(task): + if task in shared.cmd_opts.use_cpu or "all" in shared.cmd_opts.use_cpu: + return cpu + + return get_optimal_device() + + +def torch_gc(): + + if torch.cuda.is_available(): + with torch.cuda.device(get_cuda_device_string()): + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + + if has_mps(): + mac_specific.torch_mps_gc() + + if has_xpu(): + xpu_specific.torch_xpu_gc() + + if npu_specific.has_npu: + torch_npu_set_device() + npu_specific.torch_npu_gc() + + +def torch_npu_set_device(): + # Work around due to bug in torch_npu, revert me after fixed, @see https://gitee.com/ascend/pytorch/issues/I8KECW?from=project-issue + if npu_specific.has_npu: + torch.npu.set_device(0) + + +def enable_tf32(): + if torch.cuda.is_available(): + + # enabling benchmark option seems to enable a range of cards to do fp16 when they otherwise can't + # see https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/4407 + if cuda_no_autocast(): + torch.backends.cudnn.benchmark = True + + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + + +errors.run(enable_tf32, "Enabling TF32") + +cpu: torch.device = torch.device("cpu") +fp8: bool = False +# Force fp16 for all models in inference. No casting during inference. +# This flag is controlled by "--precision half" command line arg. +force_fp16: bool = False +device: torch.device = None +device_interrogate: torch.device = None +device_gfpgan: torch.device = None +device_esrgan: torch.device = None +device_codeformer: torch.device = None +dtype: torch.dtype = torch.float16 +dtype_vae: torch.dtype = torch.float16 +dtype_unet: torch.dtype = torch.float16 +dtype_inference: torch.dtype = torch.float16 +unet_needs_upcast = False + + +def cond_cast_unet(input): + if force_fp16: + return input.to(torch.float16) + return input.to(dtype_unet) if unet_needs_upcast else input + + +def cond_cast_float(input): + return input.float() if unet_needs_upcast else input + + +nv_rng = None +patch_module_list = [ + torch.nn.Linear, + torch.nn.Conv2d, + torch.nn.MultiheadAttention, + torch.nn.GroupNorm, + torch.nn.LayerNorm, +] + + +def manual_cast_forward(target_dtype): + def forward_wrapper(self, *args, **kwargs): + if any( + isinstance(arg, torch.Tensor) and arg.dtype != target_dtype + for arg in args + ): + args = [arg.to(target_dtype) if isinstance(arg, torch.Tensor) else arg for arg in args] + kwargs = {k: v.to(target_dtype) if isinstance(v, torch.Tensor) else v for k, v in kwargs.items()} + + org_dtype = target_dtype + for param in self.parameters(): + if param.dtype != target_dtype: + org_dtype = param.dtype + break + + if org_dtype != target_dtype: + self.to(target_dtype) + result = self.org_forward(*args, **kwargs) + if org_dtype != target_dtype: + self.to(org_dtype) + + if target_dtype != dtype_inference: + if isinstance(result, tuple): + result = tuple( + i.to(dtype_inference) + if isinstance(i, torch.Tensor) + else i + for i in result + ) + elif isinstance(result, torch.Tensor): + result = result.to(dtype_inference) + return result + return forward_wrapper + + +@contextlib.contextmanager +def manual_cast(target_dtype): + applied = False + for module_type in patch_module_list: + if hasattr(module_type, "org_forward"): + continue + applied = True + org_forward = module_type.forward + if module_type == torch.nn.MultiheadAttention: + module_type.forward = manual_cast_forward(torch.float32) + else: + module_type.forward = manual_cast_forward(target_dtype) + module_type.org_forward = org_forward + try: + yield None + finally: + if applied: + for module_type in patch_module_list: + if hasattr(module_type, "org_forward"): + module_type.forward = module_type.org_forward + delattr(module_type, "org_forward") + + +def autocast(disable=False): + if disable: + return contextlib.nullcontext() + + if force_fp16: + # No casting during inference if force_fp16 is enabled. + # All tensor dtype conversion happens before inference. + return contextlib.nullcontext() + + if fp8 and device==cpu: + return torch.autocast("cpu", dtype=torch.bfloat16, enabled=True) + + if fp8 and dtype_inference == torch.float32: + return manual_cast(dtype) + + if dtype == torch.float32 or dtype_inference == torch.float32: + return contextlib.nullcontext() + + if has_xpu() or has_mps() or cuda_no_autocast(): + return manual_cast(dtype) + + return torch.autocast("cuda") + + +def without_autocast(disable=False): + return torch.autocast("cuda", enabled=False) if torch.is_autocast_enabled() and not disable else contextlib.nullcontext() + + +class NansException(Exception): + pass + + +def test_for_nans(x, where): + if shared.cmd_opts.disable_nan_check: + return + + if not torch.isnan(x[(0, ) * len(x.shape)]): + return + + if where == "unet": + message = "A tensor with NaNs was produced in Unet." + + if not shared.cmd_opts.no_half: + message += " This could be either because there's not enough precision to represent the picture, or because your video card does not support half type. Try setting the \"Upcast cross attention layer to float32\" option in Settings > Stable Diffusion or using the --no-half commandline argument to fix this." + + elif where == "vae": + message = "A tensor with NaNs was produced in VAE." + + if not shared.cmd_opts.no_half and not shared.cmd_opts.no_half_vae: + message += " This could be because there's not enough precision to represent the picture. Try adding --no-half-vae commandline argument to fix this." + else: + message = "A tensor with NaNs was produced." + + message += " Use --disable-nan-check commandline argument to disable this check." + + raise NansException(message) + + +@lru_cache +def first_time_calculation(): + """ + just do any calculation with pytorch layers - the first time this is done it allocates about 700MB of memory and + spends about 2.7 seconds doing that, at least with NVidia. + """ + + x = torch.zeros((1, 1)).to(device, dtype) + linear = torch.nn.Linear(1, 1).to(device, dtype) + linear(x) + + x = torch.zeros((1, 1, 3, 3)).to(device, dtype) + conv2d = torch.nn.Conv2d(1, 1, (3, 3)).to(device, dtype) + conv2d(x) + + +def force_model_fp16(): + """ + ldm and sgm has modules.diffusionmodules.util.GroupNorm32.forward, which + force conversion of input to float32. If force_fp16 is enabled, we need to + prevent this casting. + """ + assert force_fp16 + import sgm.modules.diffusionmodules.util as sgm_util + import ldm.modules.diffusionmodules.util as ldm_util + sgm_util.GroupNorm32 = torch.nn.GroupNorm + ldm_util.GroupNorm32 = torch.nn.GroupNorm + print("ldm/sgm GroupNorm32 replaced with normal torch.nn.GroupNorm due to `--precision half`.") diff --git a/stable-diffusion-webui/modules/errors.py b/stable-diffusion-webui/modules/errors.py new file mode 100755 index 0000000..3320f45 --- /dev/null +++ b/stable-diffusion-webui/modules/errors.py @@ -0,0 +1,150 @@ +import sys +import textwrap +import traceback + + +exception_records = [] + + +def format_traceback(tb): + return [[f"{x.filename}, line {x.lineno}, {x.name}", x.line] for x in traceback.extract_tb(tb)] + + +def format_exception(e, tb): + return {"exception": str(e), "traceback": format_traceback(tb)} + + +def get_exceptions(): + try: + return list(reversed(exception_records)) + except Exception as e: + return str(e) + + +def record_exception(): + _, e, tb = sys.exc_info() + if e is None: + return + + if exception_records and exception_records[-1] == e: + return + + exception_records.append(format_exception(e, tb)) + + if len(exception_records) > 5: + exception_records.pop(0) + + +def report(message: str, *, exc_info: bool = False) -> None: + """ + Print an error message to stderr, with optional traceback. + """ + + record_exception() + + for line in message.splitlines(): + print("***", line, file=sys.stderr) + if exc_info: + print(textwrap.indent(traceback.format_exc(), " "), file=sys.stderr) + print("---", file=sys.stderr) + + +def print_error_explanation(message): + record_exception() + + lines = message.strip().split("\n") + max_len = max([len(x) for x in lines]) + + print('=' * max_len, file=sys.stderr) + for line in lines: + print(line, file=sys.stderr) + print('=' * max_len, file=sys.stderr) + + +def display(e: Exception, task, *, full_traceback=False): + record_exception() + + print(f"{task or 'error'}: {type(e).__name__}", file=sys.stderr) + te = traceback.TracebackException.from_exception(e) + if full_traceback: + # include frames leading up to the try-catch block + te.stack = traceback.StackSummary(traceback.extract_stack()[:-2] + te.stack) + print(*te.format(), sep="", file=sys.stderr) + + message = str(e) + if "copying a param with shape torch.Size([640, 1024]) from checkpoint, the shape in current model is torch.Size([640, 768])" in message: + print_error_explanation(""" +The most likely cause of this is you are trying to load Stable Diffusion 2.0 model without specifying its config file. +See https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#stable-diffusion-20 for how to solve this. + """) + + +already_displayed = {} + + +def display_once(e: Exception, task): + record_exception() + + if task in already_displayed: + return + + display(e, task) + + already_displayed[task] = 1 + + +def run(code, task): + try: + code() + except Exception as e: + display(task, e) + + +def check_versions(): + from packaging import version + from modules import shared + + import torch + import gradio + + expected_torch_version = "2.1.2" + expected_xformers_version = "0.0.23.post1" + expected_gradio_version = "3.41.2" + + if version.parse(torch.__version__) < version.parse(expected_torch_version): + print_error_explanation(f""" +You are running torch {torch.__version__}. +The program is tested to work with torch {expected_torch_version}. +To reinstall the desired version, run with commandline flag --reinstall-torch. +Beware that this will cause a lot of large files to be downloaded, as well as +there are reports of issues with training tab on the latest version. + +Use --skip-version-check commandline argument to disable this check. + """.strip()) + + if shared.xformers_available: + import xformers + + if version.parse(xformers.__version__) < version.parse(expected_xformers_version): + print_error_explanation(f""" +You are running xformers {xformers.__version__}. +The program is tested to work with xformers {expected_xformers_version}. +To reinstall the desired version, run with commandline flag --reinstall-xformers. + +Use --skip-version-check commandline argument to disable this check. + """.strip()) + + if gradio.__version__ != expected_gradio_version: + print_error_explanation(f""" +You are running gradio {gradio.__version__}. +The program is designed to work with gradio {expected_gradio_version}. +Using a different version of gradio is extremely likely to break the program. + +Reasons why you have the mismatched gradio version can be: + - you use --skip-install flag. + - you use webui.py to start the program instead of launch.py. + - an extension installs the incompatible gradio version. + +Use --skip-version-check commandline argument to disable this check. + """.strip()) + diff --git a/stable-diffusion-webui/modules/esrgan_model.py b/stable-diffusion-webui/modules/esrgan_model.py new file mode 100755 index 0000000..09f09eb --- /dev/null +++ b/stable-diffusion-webui/modules/esrgan_model.py @@ -0,0 +1,62 @@ +from modules import modelloader, devices, errors +from modules.shared import opts +from modules.upscaler import Upscaler, UpscalerData +from modules.upscaler_utils import upscale_with_model + + +class UpscalerESRGAN(Upscaler): + def __init__(self, dirname): + self.name = "ESRGAN" + self.model_url = "https://github.com/cszn/KAIR/releases/download/v1.0/ESRGAN.pth" + self.model_name = "ESRGAN_4x" + self.scalers = [] + self.user_path = dirname + super().__init__() + model_paths = self.find_models(ext_filter=[".pt", ".pth"]) + scalers = [] + if len(model_paths) == 0: + scaler_data = UpscalerData(self.model_name, self.model_url, self, 4) + scalers.append(scaler_data) + for file in model_paths: + if file.startswith("http"): + name = self.model_name + else: + name = modelloader.friendly_name(file) + + scaler_data = UpscalerData(name, file, self, 4) + self.scalers.append(scaler_data) + + def do_upscale(self, img, selected_model): + try: + model = self.load_model(selected_model) + except Exception: + errors.report(f"Unable to load ESRGAN model {selected_model}", exc_info=True) + return img + model.to(devices.device_esrgan) + return esrgan_upscale(model, img) + + def load_model(self, path: str): + if path.startswith("http"): + # TODO: this doesn't use `path` at all? + filename = modelloader.load_file_from_url( + url=self.model_url, + model_dir=self.model_download_path, + file_name=f"{self.model_name}.pth", + ) + else: + filename = path + + return modelloader.load_spandrel_model( + filename, + device=('cpu' if devices.device_esrgan.type == 'mps' else None), + expected_architecture='ESRGAN', + ) + + +def esrgan_upscale(model, img): + return upscale_with_model( + model, + img, + tile_size=opts.ESRGAN_tile, + tile_overlap=opts.ESRGAN_tile_overlap, + ) diff --git a/stable-diffusion-webui/modules/extensions.py b/stable-diffusion-webui/modules/extensions.py new file mode 100755 index 0000000..40170dc --- /dev/null +++ b/stable-diffusion-webui/modules/extensions.py @@ -0,0 +1,299 @@ +from __future__ import annotations + +import configparser +import dataclasses +import os +import threading +import re + +from modules import shared, errors, cache, scripts +from modules.gitpython_hack import Repo +from modules.paths_internal import extensions_dir, extensions_builtin_dir, script_path # noqa: F401 + +extensions: list[Extension] = [] +extension_paths: dict[str, Extension] = {} +loaded_extensions: dict[str, Exception] = {} + + +os.makedirs(extensions_dir, exist_ok=True) + + +def active(): + if shared.cmd_opts.disable_all_extensions or shared.opts.disable_all_extensions == "all": + return [] + elif shared.cmd_opts.disable_extra_extensions or shared.opts.disable_all_extensions == "extra": + return [x for x in extensions if x.enabled and x.is_builtin] + else: + return [x for x in extensions if x.enabled] + + +@dataclasses.dataclass +class CallbackOrderInfo: + name: str + before: list + after: list + + +class ExtensionMetadata: + filename = "metadata.ini" + config: configparser.ConfigParser + canonical_name: str + requires: list + + def __init__(self, path, canonical_name): + self.config = configparser.ConfigParser() + + filepath = os.path.join(path, self.filename) + # `self.config.read()` will quietly swallow OSErrors (which FileNotFoundError is), + # so no need to check whether the file exists beforehand. + try: + self.config.read(filepath) + except Exception: + errors.report(f"Error reading {self.filename} for extension {canonical_name}.", exc_info=True) + + self.canonical_name = self.config.get("Extension", "Name", fallback=canonical_name) + self.canonical_name = canonical_name.lower().strip() + + self.requires = None + + def get_script_requirements(self, field, section, extra_section=None): + """reads a list of requirements from the config; field is the name of the field in the ini file, + like Requires or Before, and section is the name of the [section] in the ini file; additionally, + reads more requirements from [extra_section] if specified.""" + + x = self.config.get(section, field, fallback='') + + if extra_section: + x = x + ', ' + self.config.get(extra_section, field, fallback='') + + listed_requirements = self.parse_list(x.lower()) + res = [] + + for requirement in listed_requirements: + loaded_requirements = (x for x in requirement.split("|") if x in loaded_extensions) + relevant_requirement = next(loaded_requirements, requirement) + res.append(relevant_requirement) + + return res + + def parse_list(self, text): + """converts a line from config ("ext1 ext2, ext3 ") into a python list (["ext1", "ext2", "ext3"])""" + + if not text: + return [] + + # both "," and " " are accepted as separator + return [x for x in re.split(r"[,\s]+", text.strip()) if x] + + def list_callback_order_instructions(self): + for section in self.config.sections(): + if not section.startswith("callbacks/"): + continue + + callback_name = section[10:] + + if not callback_name.startswith(self.canonical_name): + errors.report(f"Callback order section for extension {self.canonical_name} is referencing the wrong extension: {section}") + continue + + before = self.parse_list(self.config.get(section, 'Before', fallback='')) + after = self.parse_list(self.config.get(section, 'After', fallback='')) + + yield CallbackOrderInfo(callback_name, before, after) + + +class Extension: + lock = threading.Lock() + cached_fields = ['remote', 'commit_date', 'branch', 'commit_hash', 'version'] + metadata: ExtensionMetadata + + def __init__(self, name, path, enabled=True, is_builtin=False, metadata=None): + self.name = name + self.path = path + self.enabled = enabled + self.status = '' + self.can_update = False + self.is_builtin = is_builtin + self.commit_hash = '' + self.commit_date = None + self.version = '' + self.branch = None + self.remote = None + self.have_info_from_repo = False + self.metadata = metadata if metadata else ExtensionMetadata(self.path, name.lower()) + self.canonical_name = metadata.canonical_name + + def to_dict(self): + return {x: getattr(self, x) for x in self.cached_fields} + + def from_dict(self, d): + for field in self.cached_fields: + setattr(self, field, d[field]) + + def read_info_from_repo(self): + if self.is_builtin or self.have_info_from_repo: + return + + def read_from_repo(): + with self.lock: + if self.have_info_from_repo: + return + + self.do_read_info_from_repo() + + return self.to_dict() + + try: + d = cache.cached_data_for_file('extensions-git', self.name, os.path.join(self.path, ".git"), read_from_repo) + self.from_dict(d) + except FileNotFoundError: + pass + self.status = 'unknown' if self.status == '' else self.status + + def do_read_info_from_repo(self): + repo = None + try: + if os.path.exists(os.path.join(self.path, ".git")): + repo = Repo(self.path) + except Exception: + errors.report(f"Error reading github repository info from {self.path}", exc_info=True) + + if repo is None or repo.bare: + self.remote = None + else: + try: + self.remote = next(repo.remote().urls, None) + commit = repo.head.commit + self.commit_date = commit.committed_date + if repo.active_branch: + self.branch = repo.active_branch.name + self.commit_hash = commit.hexsha + self.version = self.commit_hash[:8] + + except Exception: + errors.report(f"Failed reading extension data from Git repository ({self.name})", exc_info=True) + self.remote = None + + self.have_info_from_repo = True + + def list_files(self, subdir, extension): + dirpath = os.path.join(self.path, subdir) + if not os.path.isdir(dirpath): + return [] + + res = [] + for filename in sorted(os.listdir(dirpath)): + res.append(scripts.ScriptFile(self.path, filename, os.path.join(dirpath, filename))) + + res = [x for x in res if os.path.splitext(x.path)[1].lower() == extension and os.path.isfile(x.path)] + + return res + + def check_updates(self): + repo = Repo(self.path) + branch_name = f'{repo.remote().name}/{self.branch}' + for fetch in repo.remote().fetch(dry_run=True): + if self.branch and fetch.name != branch_name: + continue + if fetch.flags != fetch.HEAD_UPTODATE: + self.can_update = True + self.status = "new commits" + return + + try: + origin = repo.rev_parse(branch_name) + if repo.head.commit != origin: + self.can_update = True + self.status = "behind HEAD" + return + except Exception: + self.can_update = False + self.status = "unknown (remote error)" + return + + self.can_update = False + self.status = "latest" + + def fetch_and_reset_hard(self, commit=None): + repo = Repo(self.path) + if commit is None: + commit = f'{repo.remote().name}/{self.branch}' + # Fix: `error: Your local changes to the following files would be overwritten by merge`, + # because WSL2 Docker set 755 file permissions instead of 644, this results to the error. + repo.git.fetch(all=True) + repo.git.reset(commit, hard=True) + self.have_info_from_repo = False + + +def list_extensions(): + extensions.clear() + extension_paths.clear() + loaded_extensions.clear() + + if shared.cmd_opts.disable_all_extensions: + print("*** \"--disable-all-extensions\" arg was used, will not load any extensions ***") + elif shared.opts.disable_all_extensions == "all": + print("*** \"Disable all extensions\" option was set, will not load any extensions ***") + elif shared.cmd_opts.disable_extra_extensions: + print("*** \"--disable-extra-extensions\" arg was used, will only load built-in extensions ***") + elif shared.opts.disable_all_extensions == "extra": + print("*** \"Disable all extensions\" option was set, will only load built-in extensions ***") + + + # scan through extensions directory and load metadata + for dirname in [extensions_builtin_dir, extensions_dir]: + if not os.path.isdir(dirname): + continue + + for extension_dirname in sorted(os.listdir(dirname)): + path = os.path.join(dirname, extension_dirname) + if not os.path.isdir(path): + continue + + canonical_name = extension_dirname + metadata = ExtensionMetadata(path, canonical_name) + + # check for duplicated canonical names + already_loaded_extension = loaded_extensions.get(metadata.canonical_name) + if already_loaded_extension is not None: + errors.report(f'Duplicate canonical name "{canonical_name}" found in extensions "{extension_dirname}" and "{already_loaded_extension.name}". Former will be discarded.', exc_info=False) + continue + + is_builtin = dirname == extensions_builtin_dir + extension = Extension(name=extension_dirname, path=path, enabled=extension_dirname not in shared.opts.disabled_extensions, is_builtin=is_builtin, metadata=metadata) + extensions.append(extension) + extension_paths[extension.path] = extension + loaded_extensions[canonical_name] = extension + + for extension in extensions: + extension.metadata.requires = extension.metadata.get_script_requirements("Requires", "Extension") + + # check for requirements + for extension in extensions: + if not extension.enabled: + continue + + for req in extension.metadata.requires: + required_extension = loaded_extensions.get(req) + if required_extension is None: + errors.report(f'Extension "{extension.name}" requires "{req}" which is not installed.', exc_info=False) + continue + + if not required_extension.enabled: + errors.report(f'Extension "{extension.name}" requires "{required_extension.name}" which is disabled.', exc_info=False) + continue + + +def find_extension(filename): + parentdir = os.path.dirname(os.path.realpath(filename)) + + while parentdir != filename: + extension = extension_paths.get(parentdir) + if extension is not None: + return extension + + filename = parentdir + parentdir = os.path.dirname(filename) + + return None + diff --git a/stable-diffusion-webui/modules/extra_networks.py b/stable-diffusion-webui/modules/extra_networks.py new file mode 100755 index 0000000..893a928 --- /dev/null +++ b/stable-diffusion-webui/modules/extra_networks.py @@ -0,0 +1,225 @@ +import json +import os +import re +import logging +from collections import defaultdict + +from modules import errors + +extra_network_registry = {} +extra_network_aliases = {} + + +def initialize(): + extra_network_registry.clear() + extra_network_aliases.clear() + + +def register_extra_network(extra_network): + extra_network_registry[extra_network.name] = extra_network + + +def register_extra_network_alias(extra_network, alias): + extra_network_aliases[alias] = extra_network + + +def register_default_extra_networks(): + from modules.extra_networks_hypernet import ExtraNetworkHypernet + register_extra_network(ExtraNetworkHypernet()) + + +class ExtraNetworkParams: + def __init__(self, items=None): + self.items = items or [] + self.positional = [] + self.named = {} + + for item in self.items: + parts = item.split('=', 2) if isinstance(item, str) else [item] + if len(parts) == 2: + self.named[parts[0]] = parts[1] + else: + self.positional.append(item) + + def __eq__(self, other): + return self.items == other.items + + +class ExtraNetwork: + def __init__(self, name): + self.name = name + + def activate(self, p, params_list): + """ + Called by processing on every run. Whatever the extra network is meant to do should be activated here. + Passes arguments related to this extra network in params_list. + User passes arguments by specifying this in his prompt: + + + + Where name matches the name of this ExtraNetwork object, and arg1:arg2:arg3 are any natural number of text arguments + separated by colon. + + Even if the user does not mention this ExtraNetwork in his prompt, the call will still be made, with empty params_list - + in this case, all effects of this extra networks should be disabled. + + Can be called multiple times before deactivate() - each new call should override the previous call completely. + + For example, if this ExtraNetwork's name is 'hypernet' and user's prompt is: + + > "1girl, " + + params_list will be: + + [ + ExtraNetworkParams(items=["agm", "1.1"]), + ExtraNetworkParams(items=["ray"]) + ] + + """ + raise NotImplementedError + + def deactivate(self, p): + """ + Called at the end of processing for housekeeping. No need to do anything here. + """ + + raise NotImplementedError + + +def lookup_extra_networks(extra_network_data): + """returns a dict mapping ExtraNetwork objects to lists of arguments for those extra networks. + + Example input: + { + 'lora': [], + 'lyco': [], + 'hypernet': [] + } + + Example output: + + { + : [, ], + : [] + } + """ + + res = {} + + for extra_network_name, extra_network_args in list(extra_network_data.items()): + extra_network = extra_network_registry.get(extra_network_name, None) + alias = extra_network_aliases.get(extra_network_name, None) + + if alias is not None and extra_network is None: + extra_network = alias + + if extra_network is None: + logging.info(f"Skipping unknown extra network: {extra_network_name}") + continue + + res.setdefault(extra_network, []).extend(extra_network_args) + + return res + + +def activate(p, extra_network_data): + """call activate for extra networks in extra_network_data in specified order, then call + activate for all remaining registered networks with an empty argument list""" + + activated = [] + + for extra_network, extra_network_args in lookup_extra_networks(extra_network_data).items(): + + try: + extra_network.activate(p, extra_network_args) + activated.append(extra_network) + except Exception as e: + errors.display(e, f"activating extra network {extra_network.name} with arguments {extra_network_args}") + + for extra_network_name, extra_network in extra_network_registry.items(): + if extra_network in activated: + continue + + try: + extra_network.activate(p, []) + except Exception as e: + errors.display(e, f"activating extra network {extra_network_name}") + + if p.scripts is not None: + p.scripts.after_extra_networks_activate(p, batch_number=p.iteration, prompts=p.prompts, seeds=p.seeds, subseeds=p.subseeds, extra_network_data=extra_network_data) + + +def deactivate(p, extra_network_data): + """call deactivate for extra networks in extra_network_data in specified order, then call + deactivate for all remaining registered networks""" + + data = lookup_extra_networks(extra_network_data) + + for extra_network in data: + try: + extra_network.deactivate(p) + except Exception as e: + errors.display(e, f"deactivating extra network {extra_network.name}") + + for extra_network_name, extra_network in extra_network_registry.items(): + if extra_network in data: + continue + + try: + extra_network.deactivate(p) + except Exception as e: + errors.display(e, f"deactivating unmentioned extra network {extra_network_name}") + + +re_extra_net = re.compile(r"<(\w+):([^>]+)>") + + +def parse_prompt(prompt): + res = defaultdict(list) + + def found(m): + name = m.group(1) + args = m.group(2) + + res[name].append(ExtraNetworkParams(items=args.split(":"))) + + return "" + + prompt = re.sub(re_extra_net, found, prompt) + + return prompt, res + + +def parse_prompts(prompts): + res = [] + extra_data = None + + for prompt in prompts: + updated_prompt, parsed_extra_data = parse_prompt(prompt) + + if extra_data is None: + extra_data = parsed_extra_data + + res.append(updated_prompt) + + return res, extra_data + + +def get_user_metadata(filename, lister=None): + if filename is None: + return {} + + basename, ext = os.path.splitext(filename) + metadata_filename = basename + '.json' + + metadata = {} + try: + exists = lister.exists(metadata_filename) if lister else os.path.exists(metadata_filename) + if exists: + with open(metadata_filename, "r", encoding="utf8") as file: + metadata = json.load(file) + except Exception as e: + errors.display(e, f"reading extra network user metadata from {metadata_filename}") + + return metadata diff --git a/stable-diffusion-webui/modules/extra_networks_hypernet.py b/stable-diffusion-webui/modules/extra_networks_hypernet.py new file mode 100755 index 0000000..192f11b --- /dev/null +++ b/stable-diffusion-webui/modules/extra_networks_hypernet.py @@ -0,0 +1,28 @@ +from modules import extra_networks, shared +from modules.hypernetworks import hypernetwork + + +class ExtraNetworkHypernet(extra_networks.ExtraNetwork): + def __init__(self): + super().__init__('hypernet') + + def activate(self, p, params_list): + additional = shared.opts.sd_hypernetwork + + if additional != "None" and additional in shared.hypernetworks and not any(x for x in params_list if x.items[0] == additional): + hypernet_prompt_text = f"" + p.all_prompts = [f"{prompt}{hypernet_prompt_text}" for prompt in p.all_prompts] + params_list.append(extra_networks.ExtraNetworkParams(items=[additional, shared.opts.extra_networks_default_multiplier])) + + names = [] + multipliers = [] + for params in params_list: + assert params.items + + names.append(params.items[0]) + multipliers.append(float(params.items[1]) if len(params.items) > 1 else 1.0) + + hypernetwork.load_hypernetworks(names, multipliers) + + def deactivate(self, p): + pass diff --git a/stable-diffusion-webui/modules/extras.py b/stable-diffusion-webui/modules/extras.py new file mode 100755 index 0000000..4653c3f --- /dev/null +++ b/stable-diffusion-webui/modules/extras.py @@ -0,0 +1,330 @@ +import os +import re +import shutil +import json + + +import torch +import tqdm + +from modules import shared, images, sd_models, sd_vae, sd_models_config, errors +from modules.ui_common import plaintext_to_html +import gradio as gr +import safetensors.torch + + +def run_pnginfo(image): + if image is None: + return '', '', '' + + geninfo, items = images.read_info_from_image(image) + items = {**{'parameters': geninfo}, **items} + + info = '' + for key, text in items.items(): + info += f""" +
    +

    {plaintext_to_html(str(key))}

    +

    {plaintext_to_html(str(text))}

    +
    +""".strip()+"\n" + + if len(info) == 0: + message = "Nothing found in the image." + info = f"

    {message}

    " + + return '', geninfo, info + + +def create_config(ckpt_result, config_source, a, b, c): + def config(x): + res = sd_models_config.find_checkpoint_config_near_filename(x) if x else None + return res if res != shared.sd_default_config else None + + if config_source == 0: + cfg = config(a) or config(b) or config(c) + elif config_source == 1: + cfg = config(b) + elif config_source == 2: + cfg = config(c) + else: + cfg = None + + if cfg is None: + return + + filename, _ = os.path.splitext(ckpt_result) + checkpoint_filename = filename + ".yaml" + + print("Copying config:") + print(" from:", cfg) + print(" to:", checkpoint_filename) + shutil.copyfile(cfg, checkpoint_filename) + + +checkpoint_dict_skip_on_merge = ["cond_stage_model.transformer.text_model.embeddings.position_ids"] + + +def to_half(tensor, enable): + if enable and tensor.dtype == torch.float: + return tensor.half() + + return tensor + + +def read_metadata(primary_model_name, secondary_model_name, tertiary_model_name): + metadata = {} + + for checkpoint_name in [primary_model_name, secondary_model_name, tertiary_model_name]: + checkpoint_info = sd_models.checkpoints_list.get(checkpoint_name, None) + if checkpoint_info is None: + continue + + metadata.update(checkpoint_info.metadata) + + return json.dumps(metadata, indent=4, ensure_ascii=False) + + +def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format, config_source, bake_in_vae, discard_weights, save_metadata, add_merge_recipe, copy_metadata_fields, metadata_json): + shared.state.begin(job="model-merge") + + def fail(message): + shared.state.textinfo = message + shared.state.end() + return [*[gr.update() for _ in range(4)], message] + + def weighted_sum(theta0, theta1, alpha): + return ((1 - alpha) * theta0) + (alpha * theta1) + + def get_difference(theta1, theta2): + return theta1 - theta2 + + def add_difference(theta0, theta1_2_diff, alpha): + return theta0 + (alpha * theta1_2_diff) + + def filename_weighted_sum(): + a = primary_model_info.model_name + b = secondary_model_info.model_name + Ma = round(1 - multiplier, 2) + Mb = round(multiplier, 2) + + return f"{Ma}({a}) + {Mb}({b})" + + def filename_add_difference(): + a = primary_model_info.model_name + b = secondary_model_info.model_name + c = tertiary_model_info.model_name + M = round(multiplier, 2) + + return f"{a} + {M}({b} - {c})" + + def filename_nothing(): + return primary_model_info.model_name + + theta_funcs = { + "Weighted sum": (filename_weighted_sum, None, weighted_sum), + "Add difference": (filename_add_difference, get_difference, add_difference), + "No interpolation": (filename_nothing, None, None), + } + filename_generator, theta_func1, theta_func2 = theta_funcs[interp_method] + shared.state.job_count = (1 if theta_func1 else 0) + (1 if theta_func2 else 0) + + if not primary_model_name: + return fail("Failed: Merging requires a primary model.") + + primary_model_info = sd_models.checkpoints_list[primary_model_name] + + if theta_func2 and not secondary_model_name: + return fail("Failed: Merging requires a secondary model.") + + secondary_model_info = sd_models.checkpoints_list[secondary_model_name] if theta_func2 else None + + if theta_func1 and not tertiary_model_name: + return fail(f"Failed: Interpolation method ({interp_method}) requires a tertiary model.") + + tertiary_model_info = sd_models.checkpoints_list[tertiary_model_name] if theta_func1 else None + + result_is_inpainting_model = False + result_is_instruct_pix2pix_model = False + + if theta_func2: + shared.state.textinfo = "Loading B" + print(f"Loading {secondary_model_info.filename}...") + theta_1 = sd_models.read_state_dict(secondary_model_info.filename, map_location='cpu') + else: + theta_1 = None + + if theta_func1: + shared.state.textinfo = "Loading C" + print(f"Loading {tertiary_model_info.filename}...") + theta_2 = sd_models.read_state_dict(tertiary_model_info.filename, map_location='cpu') + + shared.state.textinfo = 'Merging B and C' + shared.state.sampling_steps = len(theta_1.keys()) + for key in tqdm.tqdm(theta_1.keys()): + if key in checkpoint_dict_skip_on_merge: + continue + + if 'model' in key: + if key in theta_2: + t2 = theta_2.get(key, torch.zeros_like(theta_1[key])) + theta_1[key] = theta_func1(theta_1[key], t2) + else: + theta_1[key] = torch.zeros_like(theta_1[key]) + + shared.state.sampling_step += 1 + del theta_2 + + shared.state.nextjob() + + shared.state.textinfo = f"Loading {primary_model_info.filename}..." + print(f"Loading {primary_model_info.filename}...") + theta_0 = sd_models.read_state_dict(primary_model_info.filename, map_location='cpu') + + print("Merging...") + shared.state.textinfo = 'Merging A and B' + shared.state.sampling_steps = len(theta_0.keys()) + for key in tqdm.tqdm(theta_0.keys()): + if theta_1 and 'model' in key and key in theta_1: + + if key in checkpoint_dict_skip_on_merge: + continue + + a = theta_0[key] + b = theta_1[key] + + # this enables merging an inpainting model (A) with another one (B); + # where normal model would have 4 channels, for latenst space, inpainting model would + # have another 4 channels for unmasked picture's latent space, plus one channel for mask, for a total of 9 + if a.shape != b.shape and a.shape[0:1] + a.shape[2:] == b.shape[0:1] + b.shape[2:]: + if a.shape[1] == 4 and b.shape[1] == 9: + raise RuntimeError("When merging inpainting model with a normal one, A must be the inpainting model.") + if a.shape[1] == 4 and b.shape[1] == 8: + raise RuntimeError("When merging instruct-pix2pix model with a normal one, A must be the instruct-pix2pix model.") + + if a.shape[1] == 8 and b.shape[1] == 4:#If we have an Instruct-Pix2Pix model... + theta_0[key][:, 0:4, :, :] = theta_func2(a[:, 0:4, :, :], b, multiplier)#Merge only the vectors the models have in common. Otherwise we get an error due to dimension mismatch. + result_is_instruct_pix2pix_model = True + else: + assert a.shape[1] == 9 and b.shape[1] == 4, f"Bad dimensions for merged layer {key}: A={a.shape}, B={b.shape}" + theta_0[key][:, 0:4, :, :] = theta_func2(a[:, 0:4, :, :], b, multiplier) + result_is_inpainting_model = True + else: + theta_0[key] = theta_func2(a, b, multiplier) + + theta_0[key] = to_half(theta_0[key], save_as_half) + + shared.state.sampling_step += 1 + + del theta_1 + + bake_in_vae_filename = sd_vae.vae_dict.get(bake_in_vae, None) + if bake_in_vae_filename is not None: + print(f"Baking in VAE from {bake_in_vae_filename}") + shared.state.textinfo = 'Baking in VAE' + vae_dict = sd_vae.load_vae_dict(bake_in_vae_filename, map_location='cpu') + + for key in vae_dict.keys(): + theta_0_key = 'first_stage_model.' + key + if theta_0_key in theta_0: + theta_0[theta_0_key] = to_half(vae_dict[key], save_as_half) + + del vae_dict + + if save_as_half and not theta_func2: + for key in theta_0.keys(): + theta_0[key] = to_half(theta_0[key], save_as_half) + + if discard_weights: + regex = re.compile(discard_weights) + for key in list(theta_0): + if re.search(regex, key): + theta_0.pop(key, None) + + ckpt_dir = shared.cmd_opts.ckpt_dir or sd_models.model_path + + filename = filename_generator() if custom_name == '' else custom_name + filename += ".inpainting" if result_is_inpainting_model else "" + filename += ".instruct-pix2pix" if result_is_instruct_pix2pix_model else "" + filename += "." + checkpoint_format + + output_modelname = os.path.join(ckpt_dir, filename) + + shared.state.nextjob() + shared.state.textinfo = "Saving" + print(f"Saving to {output_modelname}...") + + metadata = {} + + if save_metadata and copy_metadata_fields: + if primary_model_info: + metadata.update(primary_model_info.metadata) + if secondary_model_info: + metadata.update(secondary_model_info.metadata) + if tertiary_model_info: + metadata.update(tertiary_model_info.metadata) + + if save_metadata: + try: + metadata.update(json.loads(metadata_json)) + except Exception as e: + errors.display(e, "readin metadata from json") + + metadata["format"] = "pt" + + if save_metadata and add_merge_recipe: + merge_recipe = { + "type": "webui", # indicate this model was merged with webui's built-in merger + "primary_model_hash": primary_model_info.sha256, + "secondary_model_hash": secondary_model_info.sha256 if secondary_model_info else None, + "tertiary_model_hash": tertiary_model_info.sha256 if tertiary_model_info else None, + "interp_method": interp_method, + "multiplier": multiplier, + "save_as_half": save_as_half, + "custom_name": custom_name, + "config_source": config_source, + "bake_in_vae": bake_in_vae, + "discard_weights": discard_weights, + "is_inpainting": result_is_inpainting_model, + "is_instruct_pix2pix": result_is_instruct_pix2pix_model + } + + sd_merge_models = {} + + def add_model_metadata(checkpoint_info): + checkpoint_info.calculate_shorthash() + sd_merge_models[checkpoint_info.sha256] = { + "name": checkpoint_info.name, + "legacy_hash": checkpoint_info.hash, + "sd_merge_recipe": checkpoint_info.metadata.get("sd_merge_recipe", None) + } + + sd_merge_models.update(checkpoint_info.metadata.get("sd_merge_models", {})) + + add_model_metadata(primary_model_info) + if secondary_model_info: + add_model_metadata(secondary_model_info) + if tertiary_model_info: + add_model_metadata(tertiary_model_info) + + metadata["sd_merge_recipe"] = json.dumps(merge_recipe) + metadata["sd_merge_models"] = json.dumps(sd_merge_models) + + _, extension = os.path.splitext(output_modelname) + if extension.lower() == ".safetensors": + safetensors.torch.save_file(theta_0, output_modelname, metadata=metadata if len(metadata)>0 else None) + else: + torch.save(theta_0, output_modelname) + + sd_models.list_models() + created_model = next((ckpt for ckpt in sd_models.checkpoints_list.values() if ckpt.name == filename), None) + if created_model: + created_model.calculate_shorthash() + + create_config(output_modelname, config_source, primary_model_info, secondary_model_info, tertiary_model_info) + + print(f"Checkpoint saved to {output_modelname}.") + shared.state.textinfo = "Checkpoint saved" + shared.state.end() + + return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)], "Checkpoint saved to " + output_modelname] diff --git a/stable-diffusion-webui/modules/face_restoration.py b/stable-diffusion-webui/modules/face_restoration.py new file mode 100755 index 0000000..2c86c6c --- /dev/null +++ b/stable-diffusion-webui/modules/face_restoration.py @@ -0,0 +1,19 @@ +from modules import shared + + +class FaceRestoration: + def name(self): + return "None" + + def restore(self, np_image): + return np_image + + +def restore_faces(np_image): + face_restorers = [x for x in shared.face_restorers if x.name() == shared.opts.face_restoration_model or shared.opts.face_restoration_model is None] + if len(face_restorers) == 0: + return np_image + + face_restorer = face_restorers[0] + + return face_restorer.restore(np_image) diff --git a/stable-diffusion-webui/modules/face_restoration_utils.py b/stable-diffusion-webui/modules/face_restoration_utils.py new file mode 100755 index 0000000..1cbac23 --- /dev/null +++ b/stable-diffusion-webui/modules/face_restoration_utils.py @@ -0,0 +1,180 @@ +from __future__ import annotations + +import logging +import os +from functools import cached_property +from typing import TYPE_CHECKING, Callable + +import cv2 +import numpy as np +import torch + +from modules import devices, errors, face_restoration, shared + +if TYPE_CHECKING: + from facexlib.utils.face_restoration_helper import FaceRestoreHelper + +logger = logging.getLogger(__name__) + + +def bgr_image_to_rgb_tensor(img: np.ndarray) -> torch.Tensor: + """Convert a BGR NumPy image in [0..1] range to a PyTorch RGB float32 tensor.""" + assert img.shape[2] == 3, "image must be RGB" + if img.dtype == "float64": + img = img.astype("float32") + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + return torch.from_numpy(img.transpose(2, 0, 1)).float() + + +def rgb_tensor_to_bgr_image(tensor: torch.Tensor, *, min_max=(0.0, 1.0)) -> np.ndarray: + """ + Convert a PyTorch RGB tensor in range `min_max` to a BGR NumPy image in [0..1] range. + """ + tensor = tensor.squeeze(0).float().detach().cpu().clamp_(*min_max) + tensor = (tensor - min_max[0]) / (min_max[1] - min_max[0]) + assert tensor.dim() == 3, "tensor must be RGB" + img_np = tensor.numpy().transpose(1, 2, 0) + if img_np.shape[2] == 1: # gray image, no RGB/BGR required + return np.squeeze(img_np, axis=2) + return cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB) + + +def create_face_helper(device) -> FaceRestoreHelper: + from facexlib.detection import retinaface + from facexlib.utils.face_restoration_helper import FaceRestoreHelper + if hasattr(retinaface, 'device'): + retinaface.device = device + return FaceRestoreHelper( + upscale_factor=1, + face_size=512, + crop_ratio=(1, 1), + det_model='retinaface_resnet50', + save_ext='png', + use_parse=True, + device=device, + ) + + +def restore_with_face_helper( + np_image: np.ndarray, + face_helper: FaceRestoreHelper, + restore_face: Callable[[torch.Tensor], torch.Tensor], +) -> np.ndarray: + """ + Find faces in the image using face_helper, restore them using restore_face, and paste them back into the image. + + `restore_face` should take a cropped face image and return a restored face image. + """ + from torchvision.transforms.functional import normalize + np_image = np_image[:, :, ::-1] + original_resolution = np_image.shape[0:2] + + try: + logger.debug("Detecting faces...") + face_helper.clean_all() + face_helper.read_image(np_image) + face_helper.get_face_landmarks_5(only_center_face=False, resize=640, eye_dist_threshold=5) + face_helper.align_warp_face() + logger.debug("Found %d faces, restoring", len(face_helper.cropped_faces)) + for cropped_face in face_helper.cropped_faces: + cropped_face_t = bgr_image_to_rgb_tensor(cropped_face / 255.0) + normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True) + cropped_face_t = cropped_face_t.unsqueeze(0).to(devices.device_codeformer) + + try: + with torch.no_grad(): + cropped_face_t = restore_face(cropped_face_t) + devices.torch_gc() + except Exception: + errors.report('Failed face-restoration inference', exc_info=True) + + restored_face = rgb_tensor_to_bgr_image(cropped_face_t, min_max=(-1, 1)) + restored_face = (restored_face * 255.0).astype('uint8') + face_helper.add_restored_face(restored_face) + + logger.debug("Merging restored faces into image") + face_helper.get_inverse_affine(None) + img = face_helper.paste_faces_to_input_image() + img = img[:, :, ::-1] + if original_resolution != img.shape[0:2]: + img = cv2.resize( + img, + (0, 0), + fx=original_resolution[1] / img.shape[1], + fy=original_resolution[0] / img.shape[0], + interpolation=cv2.INTER_LINEAR, + ) + logger.debug("Face restoration complete") + finally: + face_helper.clean_all() + return img + + +class CommonFaceRestoration(face_restoration.FaceRestoration): + net: torch.Module | None + model_url: str + model_download_name: str + + def __init__(self, model_path: str): + super().__init__() + self.net = None + self.model_path = model_path + os.makedirs(model_path, exist_ok=True) + + @cached_property + def face_helper(self) -> FaceRestoreHelper: + return create_face_helper(self.get_device()) + + def send_model_to(self, device): + if self.net: + logger.debug("Sending %s to %s", self.net, device) + self.net.to(device) + if self.face_helper: + logger.debug("Sending face helper to %s", device) + self.face_helper.face_det.to(device) + self.face_helper.face_parse.to(device) + + def get_device(self): + raise NotImplementedError("get_device must be implemented by subclasses") + + def load_net(self) -> torch.Module: + raise NotImplementedError("load_net must be implemented by subclasses") + + def restore_with_helper( + self, + np_image: np.ndarray, + restore_face: Callable[[torch.Tensor], torch.Tensor], + ) -> np.ndarray: + try: + if self.net is None: + self.net = self.load_net() + except Exception: + logger.warning("Unable to load face-restoration model", exc_info=True) + return np_image + + try: + self.send_model_to(self.get_device()) + return restore_with_face_helper(np_image, self.face_helper, restore_face) + finally: + if shared.opts.face_restoration_unload: + self.send_model_to(devices.cpu) + + +def patch_facexlib(dirname: str) -> None: + import facexlib.detection + import facexlib.parsing + + det_facex_load_file_from_url = facexlib.detection.load_file_from_url + par_facex_load_file_from_url = facexlib.parsing.load_file_from_url + + def update_kwargs(kwargs): + return dict(kwargs, save_dir=dirname, model_dir=None) + + def facex_load_file_from_url(**kwargs): + return det_facex_load_file_from_url(**update_kwargs(kwargs)) + + def facex_load_file_from_url2(**kwargs): + return par_facex_load_file_from_url(**update_kwargs(kwargs)) + + facexlib.detection.load_file_from_url = facex_load_file_from_url + facexlib.parsing.load_file_from_url = facex_load_file_from_url2 diff --git a/stable-diffusion-webui/modules/fifo_lock.py b/stable-diffusion-webui/modules/fifo_lock.py new file mode 100755 index 0000000..c35b3ae --- /dev/null +++ b/stable-diffusion-webui/modules/fifo_lock.py @@ -0,0 +1,37 @@ +import threading +import collections + + +# reference: https://gist.github.com/vitaliyp/6d54dd76ca2c3cdfc1149d33007dc34a +class FIFOLock(object): + def __init__(self): + self._lock = threading.Lock() + self._inner_lock = threading.Lock() + self._pending_threads = collections.deque() + + def acquire(self, blocking=True): + with self._inner_lock: + lock_acquired = self._lock.acquire(False) + if lock_acquired: + return True + elif not blocking: + return False + + release_event = threading.Event() + self._pending_threads.append(release_event) + + release_event.wait() + return self._lock.acquire() + + def release(self): + with self._inner_lock: + if self._pending_threads: + release_event = self._pending_threads.popleft() + release_event.set() + + self._lock.release() + + __enter__ = acquire + + def __exit__(self, t, v, tb): + self.release() diff --git a/stable-diffusion-webui/modules/gfpgan_model.py b/stable-diffusion-webui/modules/gfpgan_model.py new file mode 100755 index 0000000..6def45c --- /dev/null +++ b/stable-diffusion-webui/modules/gfpgan_model.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +import logging +import os + +import torch + +from modules import ( + devices, + errors, + face_restoration, + face_restoration_utils, + modelloader, + shared, +) + +logger = logging.getLogger(__name__) +model_url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth" +model_download_name = "GFPGANv1.4.pth" +gfpgan_face_restorer: face_restoration.FaceRestoration | None = None + + +class FaceRestorerGFPGAN(face_restoration_utils.CommonFaceRestoration): + def name(self): + return "GFPGAN" + + def get_device(self): + return devices.device_gfpgan + + def load_net(self) -> torch.Module: + for model_path in modelloader.load_models( + model_path=self.model_path, + model_url=model_url, + command_path=self.model_path, + download_name=model_download_name, + ext_filter=['.pth'], + ): + if 'GFPGAN' in os.path.basename(model_path): + return modelloader.load_spandrel_model( + model_path, + device=self.get_device(), + expected_architecture='GFPGAN', + ).model + raise ValueError("No GFPGAN model found") + + def restore(self, np_image): + def restore_face(cropped_face_t): + assert self.net is not None + return self.net(cropped_face_t, return_rgb=False)[0] + + return self.restore_with_helper(np_image, restore_face) + + +def gfpgan_fix_faces(np_image): + if gfpgan_face_restorer: + return gfpgan_face_restorer.restore(np_image) + logger.warning("GFPGAN face restorer not set up") + return np_image + + +def setup_model(dirname: str) -> None: + global gfpgan_face_restorer + + try: + face_restoration_utils.patch_facexlib(dirname) + gfpgan_face_restorer = FaceRestorerGFPGAN(model_path=dirname) + shared.face_restorers.append(gfpgan_face_restorer) + except Exception: + errors.report("Error setting up GFPGAN", exc_info=True) diff --git a/stable-diffusion-webui/modules/gitpython_hack.py b/stable-diffusion-webui/modules/gitpython_hack.py new file mode 100755 index 0000000..b55f064 --- /dev/null +++ b/stable-diffusion-webui/modules/gitpython_hack.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +import io +import subprocess + +import git + + +class Git(git.Git): + """ + Git subclassed to never use persistent processes. + """ + + def _get_persistent_cmd(self, attr_name, cmd_name, *args, **kwargs): + raise NotImplementedError(f"Refusing to use persistent process: {attr_name} ({cmd_name} {args} {kwargs})") + + def get_object_header(self, ref: str | bytes) -> tuple[str, str, int]: + ret = subprocess.check_output( + [self.GIT_PYTHON_GIT_EXECUTABLE, "cat-file", "--batch-check"], + input=self._prepare_ref(ref), + cwd=self._working_dir, + timeout=2, + ) + return self._parse_object_header(ret) + + def stream_object_data(self, ref: str) -> tuple[str, str, int, Git.CatFileContentStream]: + # Not really streaming, per se; this buffers the entire object in memory. + # Shouldn't be a problem for our use case, since we're only using this for + # object headers (commit objects). + ret = subprocess.check_output( + [self.GIT_PYTHON_GIT_EXECUTABLE, "cat-file", "--batch"], + input=self._prepare_ref(ref), + cwd=self._working_dir, + timeout=30, + ) + bio = io.BytesIO(ret) + hexsha, typename, size = self._parse_object_header(bio.readline()) + return (hexsha, typename, size, self.CatFileContentStream(size, bio)) + + +class Repo(git.Repo): + GitCommandWrapperType = Git diff --git a/stable-diffusion-webui/modules/gradio_extensons.py b/stable-diffusion-webui/modules/gradio_extensons.py new file mode 100755 index 0000000..34e3870 --- /dev/null +++ b/stable-diffusion-webui/modules/gradio_extensons.py @@ -0,0 +1,83 @@ +import gradio as gr + +from modules import scripts, ui_tempdir, patches + + +def add_classes_to_gradio_component(comp): + """ + this adds gradio-* to the component for css styling (ie gradio-button to gr.Button), as well as some others + """ + + comp.elem_classes = [f"gradio-{comp.get_block_name()}", *(comp.elem_classes or [])] + + if getattr(comp, 'multiselect', False): + comp.elem_classes.append('multiselect') + + +def IOComponent_init(self, *args, **kwargs): + self.webui_tooltip = kwargs.pop('tooltip', None) + + if scripts.scripts_current is not None: + scripts.scripts_current.before_component(self, **kwargs) + + scripts.script_callbacks.before_component_callback(self, **kwargs) + + res = original_IOComponent_init(self, *args, **kwargs) + + add_classes_to_gradio_component(self) + + scripts.script_callbacks.after_component_callback(self, **kwargs) + + if scripts.scripts_current is not None: + scripts.scripts_current.after_component(self, **kwargs) + + return res + + +def Block_get_config(self): + config = original_Block_get_config(self) + + webui_tooltip = getattr(self, 'webui_tooltip', None) + if webui_tooltip: + config["webui_tooltip"] = webui_tooltip + + config.pop('example_inputs', None) + + return config + + +def BlockContext_init(self, *args, **kwargs): + if scripts.scripts_current is not None: + scripts.scripts_current.before_component(self, **kwargs) + + scripts.script_callbacks.before_component_callback(self, **kwargs) + + res = original_BlockContext_init(self, *args, **kwargs) + + add_classes_to_gradio_component(self) + + scripts.script_callbacks.after_component_callback(self, **kwargs) + + if scripts.scripts_current is not None: + scripts.scripts_current.after_component(self, **kwargs) + + return res + + +def Blocks_get_config_file(self, *args, **kwargs): + config = original_Blocks_get_config_file(self, *args, **kwargs) + + for comp_config in config["components"]: + if "example_inputs" in comp_config: + comp_config["example_inputs"] = {"serialized": []} + + return config + + +original_IOComponent_init = patches.patch(__name__, obj=gr.components.IOComponent, field="__init__", replacement=IOComponent_init) +original_Block_get_config = patches.patch(__name__, obj=gr.blocks.Block, field="get_config", replacement=Block_get_config) +original_BlockContext_init = patches.patch(__name__, obj=gr.blocks.BlockContext, field="__init__", replacement=BlockContext_init) +original_Blocks_get_config_file = patches.patch(__name__, obj=gr.blocks.Blocks, field="get_config_file", replacement=Blocks_get_config_file) + + +ui_tempdir.install_ui_tempdir_override() diff --git a/stable-diffusion-webui/modules/hashes.py b/stable-diffusion-webui/modules/hashes.py new file mode 100755 index 0000000..7c06246 --- /dev/null +++ b/stable-diffusion-webui/modules/hashes.py @@ -0,0 +1,84 @@ +import hashlib +import os.path + +from modules import shared +import modules.cache + +dump_cache = modules.cache.dump_cache +cache = modules.cache.cache + + +def calculate_sha256(filename): + hash_sha256 = hashlib.sha256() + blksize = 1024 * 1024 + + with open(filename, "rb") as f: + for chunk in iter(lambda: f.read(blksize), b""): + hash_sha256.update(chunk) + + return hash_sha256.hexdigest() + + +def sha256_from_cache(filename, title, use_addnet_hash=False): + hashes = cache("hashes-addnet") if use_addnet_hash else cache("hashes") + try: + ondisk_mtime = os.path.getmtime(filename) + except FileNotFoundError: + return None + + if title not in hashes: + return None + + cached_sha256 = hashes[title].get("sha256", None) + cached_mtime = hashes[title].get("mtime", 0) + + if ondisk_mtime > cached_mtime or cached_sha256 is None: + return None + + return cached_sha256 + + +def sha256(filename, title, use_addnet_hash=False): + hashes = cache("hashes-addnet") if use_addnet_hash else cache("hashes") + + sha256_value = sha256_from_cache(filename, title, use_addnet_hash) + if sha256_value is not None: + return sha256_value + + if shared.cmd_opts.no_hashing: + return None + + print(f"Calculating sha256 for {filename}: ", end='') + if use_addnet_hash: + with open(filename, "rb") as file: + sha256_value = addnet_hash_safetensors(file) + else: + sha256_value = calculate_sha256(filename) + print(f"{sha256_value}") + + hashes[title] = { + "mtime": os.path.getmtime(filename), + "sha256": sha256_value, + } + + dump_cache() + + return sha256_value + + +def addnet_hash_safetensors(b): + """kohya-ss hash for safetensors from https://github.com/kohya-ss/sd-scripts/blob/main/library/train_util.py""" + hash_sha256 = hashlib.sha256() + blksize = 1024 * 1024 + + b.seek(0) + header = b.read(8) + n = int.from_bytes(header, "little") + + offset = n + 8 + b.seek(offset) + for chunk in iter(lambda: b.read(blksize), b""): + hash_sha256.update(chunk) + + return hash_sha256.hexdigest() + diff --git a/stable-diffusion-webui/modules/hat_model.py b/stable-diffusion-webui/modules/hat_model.py new file mode 100755 index 0000000..a8f737b --- /dev/null +++ b/stable-diffusion-webui/modules/hat_model.py @@ -0,0 +1,43 @@ +import os +import sys + +from modules import modelloader, devices +from modules.shared import opts +from modules.upscaler import Upscaler, UpscalerData +from modules.upscaler_utils import upscale_with_model + + +class UpscalerHAT(Upscaler): + def __init__(self, dirname): + self.name = "HAT" + self.scalers = [] + self.user_path = dirname + super().__init__() + for file in self.find_models(ext_filter=[".pt", ".pth"]): + name = modelloader.friendly_name(file) + scale = 4 # TODO: scale might not be 4, but we can't know without loading the model + scaler_data = UpscalerData(name, file, upscaler=self, scale=scale) + self.scalers.append(scaler_data) + + def do_upscale(self, img, selected_model): + try: + model = self.load_model(selected_model) + except Exception as e: + print(f"Unable to load HAT model {selected_model}: {e}", file=sys.stderr) + return img + model.to(devices.device_esrgan) # TODO: should probably be device_hat + return upscale_with_model( + model, + img, + tile_size=opts.ESRGAN_tile, # TODO: should probably be HAT_tile + tile_overlap=opts.ESRGAN_tile_overlap, # TODO: should probably be HAT_tile_overlap + ) + + def load_model(self, path: str): + if not os.path.isfile(path): + raise FileNotFoundError(f"Model file {path} not found") + return modelloader.load_spandrel_model( + path, + device=devices.device_esrgan, # TODO: should probably be device_hat + expected_architecture='HAT', + ) diff --git a/stable-diffusion-webui/modules/hypernetworks/hypernetwork.py b/stable-diffusion-webui/modules/hypernetworks/hypernetwork.py new file mode 100755 index 0000000..75e1548 --- /dev/null +++ b/stable-diffusion-webui/modules/hypernetworks/hypernetwork.py @@ -0,0 +1,783 @@ +import datetime +import glob +import html +import os +import inspect +from contextlib import closing + +import modules.textual_inversion.dataset +import torch +import tqdm +from einops import rearrange, repeat +from ldm.util import default +from modules import devices, sd_models, shared, sd_samplers, hashes, sd_hijack_checkpoint, errors +from modules.textual_inversion import textual_inversion, saving_settings +from modules.textual_inversion.learn_schedule import LearnRateScheduler +from torch import einsum +from torch.nn.init import normal_, xavier_normal_, xavier_uniform_, kaiming_normal_, kaiming_uniform_, zeros_ + +from collections import deque +from statistics import stdev, mean + + +optimizer_dict = {optim_name : cls_obj for optim_name, cls_obj in inspect.getmembers(torch.optim, inspect.isclass) if optim_name != "Optimizer"} + +class HypernetworkModule(torch.nn.Module): + activation_dict = { + "linear": torch.nn.Identity, + "relu": torch.nn.ReLU, + "leakyrelu": torch.nn.LeakyReLU, + "elu": torch.nn.ELU, + "swish": torch.nn.Hardswish, + "tanh": torch.nn.Tanh, + "sigmoid": torch.nn.Sigmoid, + } + activation_dict.update({cls_name.lower(): cls_obj for cls_name, cls_obj in inspect.getmembers(torch.nn.modules.activation) if inspect.isclass(cls_obj) and cls_obj.__module__ == 'torch.nn.modules.activation'}) + + def __init__(self, dim, state_dict=None, layer_structure=None, activation_func=None, weight_init='Normal', + add_layer_norm=False, activate_output=False, dropout_structure=None): + super().__init__() + + self.multiplier = 1.0 + + assert layer_structure is not None, "layer_structure must not be None" + assert layer_structure[0] == 1, "Multiplier Sequence should start with size 1!" + assert layer_structure[-1] == 1, "Multiplier Sequence should end with size 1!" + + linears = [] + for i in range(len(layer_structure) - 1): + + # Add a fully-connected layer + linears.append(torch.nn.Linear(int(dim * layer_structure[i]), int(dim * layer_structure[i+1]))) + + # Add an activation func except last layer + if activation_func == "linear" or activation_func is None or (i >= len(layer_structure) - 2 and not activate_output): + pass + elif activation_func in self.activation_dict: + linears.append(self.activation_dict[activation_func]()) + else: + raise RuntimeError(f'hypernetwork uses an unsupported activation function: {activation_func}') + + # Add layer normalization + if add_layer_norm: + linears.append(torch.nn.LayerNorm(int(dim * layer_structure[i+1]))) + + # Everything should be now parsed into dropout structure, and applied here. + # Since we only have dropouts after layers, dropout structure should start with 0 and end with 0. + if dropout_structure is not None and dropout_structure[i+1] > 0: + assert 0 < dropout_structure[i+1] < 1, "Dropout probability should be 0 or float between 0 and 1!" + linears.append(torch.nn.Dropout(p=dropout_structure[i+1])) + # Code explanation : [1, 2, 1] -> dropout is missing when last_layer_dropout is false. [1, 2, 2, 1] -> [0, 0.3, 0, 0], when its True, [0, 0.3, 0.3, 0]. + + self.linear = torch.nn.Sequential(*linears) + + if state_dict is not None: + self.fix_old_state_dict(state_dict) + self.load_state_dict(state_dict) + else: + for layer in self.linear: + if type(layer) == torch.nn.Linear or type(layer) == torch.nn.LayerNorm: + w, b = layer.weight.data, layer.bias.data + if weight_init == "Normal" or type(layer) == torch.nn.LayerNorm: + normal_(w, mean=0.0, std=0.01) + normal_(b, mean=0.0, std=0) + elif weight_init == 'XavierUniform': + xavier_uniform_(w) + zeros_(b) + elif weight_init == 'XavierNormal': + xavier_normal_(w) + zeros_(b) + elif weight_init == 'KaimingUniform': + kaiming_uniform_(w, nonlinearity='leaky_relu' if 'leakyrelu' == activation_func else 'relu') + zeros_(b) + elif weight_init == 'KaimingNormal': + kaiming_normal_(w, nonlinearity='leaky_relu' if 'leakyrelu' == activation_func else 'relu') + zeros_(b) + else: + raise KeyError(f"Key {weight_init} is not defined as initialization!") + devices.torch_npu_set_device() + self.to(devices.device) + + def fix_old_state_dict(self, state_dict): + changes = { + 'linear1.bias': 'linear.0.bias', + 'linear1.weight': 'linear.0.weight', + 'linear2.bias': 'linear.1.bias', + 'linear2.weight': 'linear.1.weight', + } + + for fr, to in changes.items(): + x = state_dict.get(fr, None) + if x is None: + continue + + del state_dict[fr] + state_dict[to] = x + + def forward(self, x): + return x + self.linear(x) * (self.multiplier if not self.training else 1) + + def trainables(self): + layer_structure = [] + for layer in self.linear: + if type(layer) == torch.nn.Linear or type(layer) == torch.nn.LayerNorm: + layer_structure += [layer.weight, layer.bias] + return layer_structure + + +#param layer_structure : sequence used for length, use_dropout : controlling boolean, last_layer_dropout : for compatibility check. +def parse_dropout_structure(layer_structure, use_dropout, last_layer_dropout): + if layer_structure is None: + layer_structure = [1, 2, 1] + if not use_dropout: + return [0] * len(layer_structure) + dropout_values = [0] + dropout_values.extend([0.3] * (len(layer_structure) - 3)) + if last_layer_dropout: + dropout_values.append(0.3) + else: + dropout_values.append(0) + dropout_values.append(0) + return dropout_values + + +class Hypernetwork: + filename = None + name = None + + def __init__(self, name=None, enable_sizes=None, layer_structure=None, activation_func=None, weight_init=None, add_layer_norm=False, use_dropout=False, activate_output=False, **kwargs): + self.filename = None + self.name = name + self.layers = {} + self.step = 0 + self.sd_checkpoint = None + self.sd_checkpoint_name = None + self.layer_structure = layer_structure + self.activation_func = activation_func + self.weight_init = weight_init + self.add_layer_norm = add_layer_norm + self.use_dropout = use_dropout + self.activate_output = activate_output + self.last_layer_dropout = kwargs.get('last_layer_dropout', True) + self.dropout_structure = kwargs.get('dropout_structure', None) + if self.dropout_structure is None: + self.dropout_structure = parse_dropout_structure(self.layer_structure, self.use_dropout, self.last_layer_dropout) + self.optimizer_name = None + self.optimizer_state_dict = None + self.optional_info = None + + for size in enable_sizes or []: + self.layers[size] = ( + HypernetworkModule(size, None, self.layer_structure, self.activation_func, self.weight_init, + self.add_layer_norm, self.activate_output, dropout_structure=self.dropout_structure), + HypernetworkModule(size, None, self.layer_structure, self.activation_func, self.weight_init, + self.add_layer_norm, self.activate_output, dropout_structure=self.dropout_structure), + ) + self.eval() + + def weights(self): + res = [] + for layers in self.layers.values(): + for layer in layers: + res += layer.parameters() + return res + + def train(self, mode=True): + for layers in self.layers.values(): + for layer in layers: + layer.train(mode=mode) + for param in layer.parameters(): + param.requires_grad = mode + + def to(self, device): + for layers in self.layers.values(): + for layer in layers: + layer.to(device) + + return self + + def set_multiplier(self, multiplier): + for layers in self.layers.values(): + for layer in layers: + layer.multiplier = multiplier + + return self + + def eval(self): + for layers in self.layers.values(): + for layer in layers: + layer.eval() + for param in layer.parameters(): + param.requires_grad = False + + def save(self, filename): + state_dict = {} + optimizer_saved_dict = {} + + for k, v in self.layers.items(): + state_dict[k] = (v[0].state_dict(), v[1].state_dict()) + + state_dict['step'] = self.step + state_dict['name'] = self.name + state_dict['layer_structure'] = self.layer_structure + state_dict['activation_func'] = self.activation_func + state_dict['is_layer_norm'] = self.add_layer_norm + state_dict['weight_initialization'] = self.weight_init + state_dict['sd_checkpoint'] = self.sd_checkpoint + state_dict['sd_checkpoint_name'] = self.sd_checkpoint_name + state_dict['activate_output'] = self.activate_output + state_dict['use_dropout'] = self.use_dropout + state_dict['dropout_structure'] = self.dropout_structure + state_dict['last_layer_dropout'] = (self.dropout_structure[-2] != 0) if self.dropout_structure is not None else self.last_layer_dropout + state_dict['optional_info'] = self.optional_info if self.optional_info else None + + if self.optimizer_name is not None: + optimizer_saved_dict['optimizer_name'] = self.optimizer_name + + torch.save(state_dict, filename) + if shared.opts.save_optimizer_state and self.optimizer_state_dict: + optimizer_saved_dict['hash'] = self.shorthash() + optimizer_saved_dict['optimizer_state_dict'] = self.optimizer_state_dict + torch.save(optimizer_saved_dict, filename + '.optim') + + def load(self, filename): + self.filename = filename + if self.name is None: + self.name = os.path.splitext(os.path.basename(filename))[0] + + state_dict = torch.load(filename, map_location='cpu') + + self.layer_structure = state_dict.get('layer_structure', [1, 2, 1]) + self.optional_info = state_dict.get('optional_info', None) + self.activation_func = state_dict.get('activation_func', None) + self.weight_init = state_dict.get('weight_initialization', 'Normal') + self.add_layer_norm = state_dict.get('is_layer_norm', False) + self.dropout_structure = state_dict.get('dropout_structure', None) + self.use_dropout = True if self.dropout_structure is not None and any(self.dropout_structure) else state_dict.get('use_dropout', False) + self.activate_output = state_dict.get('activate_output', True) + self.last_layer_dropout = state_dict.get('last_layer_dropout', False) + # Dropout structure should have same length as layer structure, Every digits should be in [0,1), and last digit must be 0. + if self.dropout_structure is None: + self.dropout_structure = parse_dropout_structure(self.layer_structure, self.use_dropout, self.last_layer_dropout) + + if shared.opts.print_hypernet_extra: + if self.optional_info is not None: + print(f" INFO:\n {self.optional_info}\n") + + print(f" Layer structure: {self.layer_structure}") + print(f" Activation function: {self.activation_func}") + print(f" Weight initialization: {self.weight_init}") + print(f" Layer norm: {self.add_layer_norm}") + print(f" Dropout usage: {self.use_dropout}" ) + print(f" Activate last layer: {self.activate_output}") + print(f" Dropout structure: {self.dropout_structure}") + + optimizer_saved_dict = torch.load(self.filename + '.optim', map_location='cpu') if os.path.exists(self.filename + '.optim') else {} + + if self.shorthash() == optimizer_saved_dict.get('hash', None): + self.optimizer_state_dict = optimizer_saved_dict.get('optimizer_state_dict', None) + else: + self.optimizer_state_dict = None + if self.optimizer_state_dict: + self.optimizer_name = optimizer_saved_dict.get('optimizer_name', 'AdamW') + if shared.opts.print_hypernet_extra: + print("Loaded existing optimizer from checkpoint") + print(f"Optimizer name is {self.optimizer_name}") + else: + self.optimizer_name = "AdamW" + if shared.opts.print_hypernet_extra: + print("No saved optimizer exists in checkpoint") + + for size, sd in state_dict.items(): + if type(size) == int: + self.layers[size] = ( + HypernetworkModule(size, sd[0], self.layer_structure, self.activation_func, self.weight_init, + self.add_layer_norm, self.activate_output, self.dropout_structure), + HypernetworkModule(size, sd[1], self.layer_structure, self.activation_func, self.weight_init, + self.add_layer_norm, self.activate_output, self.dropout_structure), + ) + + self.name = state_dict.get('name', self.name) + self.step = state_dict.get('step', 0) + self.sd_checkpoint = state_dict.get('sd_checkpoint', None) + self.sd_checkpoint_name = state_dict.get('sd_checkpoint_name', None) + self.eval() + + def shorthash(self): + sha256 = hashes.sha256(self.filename, f'hypernet/{self.name}') + + return sha256[0:10] if sha256 else None + + +def list_hypernetworks(path): + res = {} + for filename in sorted(glob.iglob(os.path.join(path, '**/*.pt'), recursive=True), key=str.lower): + name = os.path.splitext(os.path.basename(filename))[0] + # Prevent a hypothetical "None.pt" from being listed. + if name != "None": + res[name] = filename + return res + + +def load_hypernetwork(name): + path = shared.hypernetworks.get(name, None) + + if path is None: + return None + + try: + hypernetwork = Hypernetwork() + hypernetwork.load(path) + return hypernetwork + except Exception: + errors.report(f"Error loading hypernetwork {path}", exc_info=True) + return None + + +def load_hypernetworks(names, multipliers=None): + already_loaded = {} + + for hypernetwork in shared.loaded_hypernetworks: + if hypernetwork.name in names: + already_loaded[hypernetwork.name] = hypernetwork + + shared.loaded_hypernetworks.clear() + + for i, name in enumerate(names): + hypernetwork = already_loaded.get(name, None) + if hypernetwork is None: + hypernetwork = load_hypernetwork(name) + + if hypernetwork is None: + continue + + hypernetwork.set_multiplier(multipliers[i] if multipliers else 1.0) + shared.loaded_hypernetworks.append(hypernetwork) + + +def apply_single_hypernetwork(hypernetwork, context_k, context_v, layer=None): + hypernetwork_layers = (hypernetwork.layers if hypernetwork is not None else {}).get(context_k.shape[2], None) + + if hypernetwork_layers is None: + return context_k, context_v + + if layer is not None: + layer.hyper_k = hypernetwork_layers[0] + layer.hyper_v = hypernetwork_layers[1] + + context_k = devices.cond_cast_unet(hypernetwork_layers[0](devices.cond_cast_float(context_k))) + context_v = devices.cond_cast_unet(hypernetwork_layers[1](devices.cond_cast_float(context_v))) + return context_k, context_v + + +def apply_hypernetworks(hypernetworks, context, layer=None): + context_k = context + context_v = context + for hypernetwork in hypernetworks: + context_k, context_v = apply_single_hypernetwork(hypernetwork, context_k, context_v, layer) + + return context_k, context_v + + +def attention_CrossAttention_forward(self, x, context=None, mask=None, **kwargs): + h = self.heads + + q = self.to_q(x) + context = default(context, x) + + context_k, context_v = apply_hypernetworks(shared.loaded_hypernetworks, context, self) + k = self.to_k(context_k) + v = self.to_v(context_v) + + q, k, v = (rearrange(t, 'b n (h d) -> (b h) n d', h=h) for t in (q, k, v)) + + sim = einsum('b i d, b j d -> b i j', q, k) * self.scale + + if mask is not None: + mask = rearrange(mask, 'b ... -> b (...)') + max_neg_value = -torch.finfo(sim.dtype).max + mask = repeat(mask, 'b j -> (b h) () j', h=h) + sim.masked_fill_(~mask, max_neg_value) + + # attention, what we cannot get enough of + attn = sim.softmax(dim=-1) + + out = einsum('b i j, b j d -> b i d', attn, v) + out = rearrange(out, '(b h) n d -> b n (h d)', h=h) + return self.to_out(out) + + +def stack_conds(conds): + if len(conds) == 1: + return torch.stack(conds) + + # same as in reconstruct_multicond_batch + token_count = max([x.shape[0] for x in conds]) + for i in range(len(conds)): + if conds[i].shape[0] != token_count: + last_vector = conds[i][-1:] + last_vector_repeated = last_vector.repeat([token_count - conds[i].shape[0], 1]) + conds[i] = torch.vstack([conds[i], last_vector_repeated]) + + return torch.stack(conds) + + +def statistics(data): + if len(data) < 2: + std = 0 + else: + std = stdev(data) + total_information = f"loss:{mean(data):.3f}" + u"\u00B1" + f"({std/ (len(data) ** 0.5):.3f})" + recent_data = data[-32:] + if len(recent_data) < 2: + std = 0 + else: + std = stdev(recent_data) + recent_information = f"recent 32 loss:{mean(recent_data):.3f}" + u"\u00B1" + f"({std / (len(recent_data) ** 0.5):.3f})" + return total_information, recent_information + + +def create_hypernetwork(name, enable_sizes, overwrite_old, layer_structure=None, activation_func=None, weight_init=None, add_layer_norm=False, use_dropout=False, dropout_structure=None): + # Remove illegal characters from name. + name = "".join( x for x in name if (x.isalnum() or x in "._- ")) + assert name, "Name cannot be empty!" + + fn = os.path.join(shared.cmd_opts.hypernetwork_dir, f"{name}.pt") + if not overwrite_old: + assert not os.path.exists(fn), f"file {fn} already exists" + + if type(layer_structure) == str: + layer_structure = [float(x.strip()) for x in layer_structure.split(",")] + + if use_dropout and dropout_structure and type(dropout_structure) == str: + dropout_structure = [float(x.strip()) for x in dropout_structure.split(",")] + else: + dropout_structure = [0] * len(layer_structure) + + hypernet = modules.hypernetworks.hypernetwork.Hypernetwork( + name=name, + enable_sizes=[int(x) for x in enable_sizes], + layer_structure=layer_structure, + activation_func=activation_func, + weight_init=weight_init, + add_layer_norm=add_layer_norm, + use_dropout=use_dropout, + dropout_structure=dropout_structure + ) + hypernet.save(fn) + + shared.reload_hypernetworks() + + +def train_hypernetwork(id_task, hypernetwork_name: str, learn_rate: float, batch_size: int, gradient_step: int, data_root: str, log_directory: str, training_width: int, training_height: int, varsize: bool, steps: int, clip_grad_mode: str, clip_grad_value: float, shuffle_tags: bool, tag_drop_out: bool, latent_sampling_method: str, use_weight: bool, create_image_every: int, save_hypernetwork_every: int, template_filename: str, preview_from_txt2img: bool, preview_prompt: str, preview_negative_prompt: str, preview_steps: int, preview_sampler_name: str, preview_cfg_scale: float, preview_seed: int, preview_width: int, preview_height: int): + from modules import images, processing + + save_hypernetwork_every = save_hypernetwork_every or 0 + create_image_every = create_image_every or 0 + template_file = textual_inversion.textual_inversion_templates.get(template_filename, None) + textual_inversion.validate_train_inputs(hypernetwork_name, learn_rate, batch_size, gradient_step, data_root, template_file, template_filename, steps, save_hypernetwork_every, create_image_every, log_directory, name="hypernetwork") + template_file = template_file.path + + path = shared.hypernetworks.get(hypernetwork_name, None) + hypernetwork = Hypernetwork() + hypernetwork.load(path) + shared.loaded_hypernetworks = [hypernetwork] + + shared.state.job = "train-hypernetwork" + shared.state.textinfo = "Initializing hypernetwork training..." + shared.state.job_count = steps + + hypernetwork_name = hypernetwork_name.rsplit('(', 1)[0] + filename = os.path.join(shared.cmd_opts.hypernetwork_dir, f'{hypernetwork_name}.pt') + + log_directory = os.path.join(log_directory, datetime.datetime.now().strftime("%Y-%m-%d"), hypernetwork_name) + unload = shared.opts.unload_models_when_training + + if save_hypernetwork_every > 0: + hypernetwork_dir = os.path.join(log_directory, "hypernetworks") + os.makedirs(hypernetwork_dir, exist_ok=True) + else: + hypernetwork_dir = None + + if create_image_every > 0: + images_dir = os.path.join(log_directory, "images") + os.makedirs(images_dir, exist_ok=True) + else: + images_dir = None + + checkpoint = sd_models.select_checkpoint() + + initial_step = hypernetwork.step or 0 + if initial_step >= steps: + shared.state.textinfo = "Model has already been trained beyond specified max steps" + return hypernetwork, filename + + scheduler = LearnRateScheduler(learn_rate, steps, initial_step) + + clip_grad = torch.nn.utils.clip_grad_value_ if clip_grad_mode == "value" else torch.nn.utils.clip_grad_norm_ if clip_grad_mode == "norm" else None + if clip_grad: + clip_grad_sched = LearnRateScheduler(clip_grad_value, steps, initial_step, verbose=False) + + if shared.opts.training_enable_tensorboard: + tensorboard_writer = textual_inversion.tensorboard_setup(log_directory) + + # dataset loading may take a while, so input validations and early returns should be done before this + shared.state.textinfo = f"Preparing dataset from {html.escape(data_root)}..." + + pin_memory = shared.opts.pin_memory + + ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=training_width, height=training_height, repeats=shared.opts.training_image_repeats_per_epoch, placeholder_token=hypernetwork_name, model=shared.sd_model, cond_model=shared.sd_model.cond_stage_model, device=devices.device, template_file=template_file, include_cond=True, batch_size=batch_size, gradient_step=gradient_step, shuffle_tags=shuffle_tags, tag_drop_out=tag_drop_out, latent_sampling_method=latent_sampling_method, varsize=varsize, use_weight=use_weight) + + if shared.opts.save_training_settings_to_txt: + saved_params = dict( + model_name=checkpoint.model_name, model_hash=checkpoint.shorthash, num_of_dataset_images=len(ds), + **{field: getattr(hypernetwork, field) for field in ['layer_structure', 'activation_func', 'weight_init', 'add_layer_norm', 'use_dropout', ]} + ) + saving_settings.save_settings_to_file(log_directory, {**saved_params, **locals()}) + + latent_sampling_method = ds.latent_sampling_method + + dl = modules.textual_inversion.dataset.PersonalizedDataLoader(ds, latent_sampling_method=latent_sampling_method, batch_size=ds.batch_size, pin_memory=pin_memory) + + old_parallel_processing_allowed = shared.parallel_processing_allowed + + if unload: + shared.parallel_processing_allowed = False + shared.sd_model.cond_stage_model.to(devices.cpu) + shared.sd_model.first_stage_model.to(devices.cpu) + + weights = hypernetwork.weights() + hypernetwork.train() + + # Here we use optimizer from saved HN, or we can specify as UI option. + if hypernetwork.optimizer_name in optimizer_dict: + optimizer = optimizer_dict[hypernetwork.optimizer_name](params=weights, lr=scheduler.learn_rate) + optimizer_name = hypernetwork.optimizer_name + else: + print(f"Optimizer type {hypernetwork.optimizer_name} is not defined!") + optimizer = torch.optim.AdamW(params=weights, lr=scheduler.learn_rate) + optimizer_name = 'AdamW' + + if hypernetwork.optimizer_state_dict: # This line must be changed if Optimizer type can be different from saved optimizer. + try: + optimizer.load_state_dict(hypernetwork.optimizer_state_dict) + except RuntimeError as e: + print("Cannot resume from saved optimizer!") + print(e) + + scaler = torch.cuda.amp.GradScaler() + + batch_size = ds.batch_size + gradient_step = ds.gradient_step + # n steps = batch_size * gradient_step * n image processed + steps_per_epoch = len(ds) // batch_size // gradient_step + max_steps_per_epoch = len(ds) // batch_size - (len(ds) // batch_size) % gradient_step + loss_step = 0 + _loss_step = 0 #internal + # size = len(ds.indexes) + # loss_dict = defaultdict(lambda : deque(maxlen = 1024)) + loss_logging = deque(maxlen=len(ds) * 3) # this should be configurable parameter, this is 3 * epoch(dataset size) + # losses = torch.zeros((size,)) + # previous_mean_losses = [0] + # previous_mean_loss = 0 + # print("Mean loss of {} elements".format(size)) + + steps_without_grad = 0 + + last_saved_file = "" + last_saved_image = "" + forced_filename = "" + + pbar = tqdm.tqdm(total=steps - initial_step) + try: + sd_hijack_checkpoint.add() + + for _ in range((steps-initial_step) * gradient_step): + if scheduler.finished: + break + if shared.state.interrupted: + break + for j, batch in enumerate(dl): + # works as a drop_last=True for gradient accumulation + if j == max_steps_per_epoch: + break + scheduler.apply(optimizer, hypernetwork.step) + if scheduler.finished: + break + if shared.state.interrupted: + break + + if clip_grad: + clip_grad_sched.step(hypernetwork.step) + + with devices.autocast(): + x = batch.latent_sample.to(devices.device, non_blocking=pin_memory) + if use_weight: + w = batch.weight.to(devices.device, non_blocking=pin_memory) + if tag_drop_out != 0 or shuffle_tags: + shared.sd_model.cond_stage_model.to(devices.device) + c = shared.sd_model.cond_stage_model(batch.cond_text).to(devices.device, non_blocking=pin_memory) + shared.sd_model.cond_stage_model.to(devices.cpu) + else: + c = stack_conds(batch.cond).to(devices.device, non_blocking=pin_memory) + if use_weight: + loss = shared.sd_model.weighted_forward(x, c, w)[0] / gradient_step + del w + else: + loss = shared.sd_model.forward(x, c)[0] / gradient_step + del x + del c + + _loss_step += loss.item() + scaler.scale(loss).backward() + + # go back until we reach gradient accumulation steps + if (j + 1) % gradient_step != 0: + continue + loss_logging.append(_loss_step) + if clip_grad: + clip_grad(weights, clip_grad_sched.learn_rate) + + scaler.step(optimizer) + scaler.update() + hypernetwork.step += 1 + pbar.update() + optimizer.zero_grad(set_to_none=True) + loss_step = _loss_step + _loss_step = 0 + + steps_done = hypernetwork.step + 1 + + epoch_num = hypernetwork.step // steps_per_epoch + epoch_step = hypernetwork.step % steps_per_epoch + + description = f"Training hypernetwork [Epoch {epoch_num}: {epoch_step+1}/{steps_per_epoch}]loss: {loss_step:.7f}" + pbar.set_description(description) + if hypernetwork_dir is not None and steps_done % save_hypernetwork_every == 0: + # Before saving, change name to match current checkpoint. + hypernetwork_name_every = f'{hypernetwork_name}-{steps_done}' + last_saved_file = os.path.join(hypernetwork_dir, f'{hypernetwork_name_every}.pt') + hypernetwork.optimizer_name = optimizer_name + if shared.opts.save_optimizer_state: + hypernetwork.optimizer_state_dict = optimizer.state_dict() + save_hypernetwork(hypernetwork, checkpoint, hypernetwork_name, last_saved_file) + hypernetwork.optimizer_state_dict = None # dereference it after saving, to save memory. + + + + if shared.opts.training_enable_tensorboard: + epoch_num = hypernetwork.step // len(ds) + epoch_step = hypernetwork.step - (epoch_num * len(ds)) + 1 + mean_loss = sum(loss_logging) / len(loss_logging) + textual_inversion.tensorboard_add(tensorboard_writer, loss=mean_loss, global_step=hypernetwork.step, step=epoch_step, learn_rate=scheduler.learn_rate, epoch_num=epoch_num) + + textual_inversion.write_loss(log_directory, "hypernetwork_loss.csv", hypernetwork.step, steps_per_epoch, { + "loss": f"{loss_step:.7f}", + "learn_rate": scheduler.learn_rate + }) + + if images_dir is not None and steps_done % create_image_every == 0: + forced_filename = f'{hypernetwork_name}-{steps_done}' + last_saved_image = os.path.join(images_dir, forced_filename) + hypernetwork.eval() + rng_state = torch.get_rng_state() + cuda_rng_state = None + if torch.cuda.is_available(): + cuda_rng_state = torch.cuda.get_rng_state_all() + shared.sd_model.cond_stage_model.to(devices.device) + shared.sd_model.first_stage_model.to(devices.device) + + p = processing.StableDiffusionProcessingTxt2Img( + sd_model=shared.sd_model, + do_not_save_grid=True, + do_not_save_samples=True, + ) + + p.disable_extra_networks = True + + if preview_from_txt2img: + p.prompt = preview_prompt + p.negative_prompt = preview_negative_prompt + p.steps = preview_steps + p.sampler_name = sd_samplers.samplers_map[preview_sampler_name.lower()] + p.cfg_scale = preview_cfg_scale + p.seed = preview_seed + p.width = preview_width + p.height = preview_height + else: + p.prompt = batch.cond_text[0] + p.steps = 20 + p.width = training_width + p.height = training_height + + preview_text = p.prompt + + with closing(p): + processed = processing.process_images(p) + image = processed.images[0] if len(processed.images) > 0 else None + + if unload: + shared.sd_model.cond_stage_model.to(devices.cpu) + shared.sd_model.first_stage_model.to(devices.cpu) + torch.set_rng_state(rng_state) + if torch.cuda.is_available(): + torch.cuda.set_rng_state_all(cuda_rng_state) + hypernetwork.train() + if image is not None: + shared.state.assign_current_image(image) + if shared.opts.training_enable_tensorboard and shared.opts.training_tensorboard_save_images: + textual_inversion.tensorboard_add_image(tensorboard_writer, + f"Validation at epoch {epoch_num}", image, + hypernetwork.step) + last_saved_image, last_text_info = images.save_image(image, images_dir, "", p.seed, p.prompt, shared.opts.samples_format, processed.infotexts[0], p=p, forced_filename=forced_filename, save_to_dirs=False) + last_saved_image += f", prompt: {preview_text}" + + shared.state.job_no = hypernetwork.step + + shared.state.textinfo = f""" +

    +Loss: {loss_step:.7f}
    +Step: {steps_done}
    +Last prompt: {html.escape(batch.cond_text[0])}
    +Last saved hypernetwork: {html.escape(last_saved_file)}
    +Last saved image: {html.escape(last_saved_image)}
    +

    +""" + except Exception: + errors.report("Exception in training hypernetwork", exc_info=True) + finally: + pbar.leave = False + pbar.close() + hypernetwork.eval() + sd_hijack_checkpoint.remove() + + + + filename = os.path.join(shared.cmd_opts.hypernetwork_dir, f'{hypernetwork_name}.pt') + hypernetwork.optimizer_name = optimizer_name + if shared.opts.save_optimizer_state: + hypernetwork.optimizer_state_dict = optimizer.state_dict() + save_hypernetwork(hypernetwork, checkpoint, hypernetwork_name, filename) + + del optimizer + hypernetwork.optimizer_state_dict = None # dereference it after saving, to save memory. + shared.sd_model.cond_stage_model.to(devices.device) + shared.sd_model.first_stage_model.to(devices.device) + shared.parallel_processing_allowed = old_parallel_processing_allowed + + return hypernetwork, filename + +def save_hypernetwork(hypernetwork, checkpoint, hypernetwork_name, filename): + old_hypernetwork_name = hypernetwork.name + old_sd_checkpoint = hypernetwork.sd_checkpoint if hasattr(hypernetwork, "sd_checkpoint") else None + old_sd_checkpoint_name = hypernetwork.sd_checkpoint_name if hasattr(hypernetwork, "sd_checkpoint_name") else None + try: + hypernetwork.sd_checkpoint = checkpoint.shorthash + hypernetwork.sd_checkpoint_name = checkpoint.model_name + hypernetwork.name = hypernetwork_name + hypernetwork.save(filename) + except: + hypernetwork.sd_checkpoint = old_sd_checkpoint + hypernetwork.sd_checkpoint_name = old_sd_checkpoint_name + hypernetwork.name = old_hypernetwork_name + raise diff --git a/stable-diffusion-webui/modules/hypernetworks/ui.py b/stable-diffusion-webui/modules/hypernetworks/ui.py new file mode 100755 index 0000000..3519104 --- /dev/null +++ b/stable-diffusion-webui/modules/hypernetworks/ui.py @@ -0,0 +1,38 @@ +import html + +import gradio as gr +import modules.hypernetworks.hypernetwork +from modules import devices, sd_hijack, shared + +not_available = ["hardswish", "multiheadattention"] +keys = [x for x in modules.hypernetworks.hypernetwork.HypernetworkModule.activation_dict if x not in not_available] + + +def create_hypernetwork(name, enable_sizes, overwrite_old, layer_structure=None, activation_func=None, weight_init=None, add_layer_norm=False, use_dropout=False, dropout_structure=None): + filename = modules.hypernetworks.hypernetwork.create_hypernetwork(name, enable_sizes, overwrite_old, layer_structure, activation_func, weight_init, add_layer_norm, use_dropout, dropout_structure) + + return gr.Dropdown.update(choices=sorted(shared.hypernetworks)), f"Created: {filename}", "" + + +def train_hypernetwork(*args): + shared.loaded_hypernetworks = [] + + assert not shared.cmd_opts.lowvram, 'Training models with lowvram is not possible' + + try: + sd_hijack.undo_optimizations() + + hypernetwork, filename = modules.hypernetworks.hypernetwork.train_hypernetwork(*args) + + res = f""" +Training {'interrupted' if shared.state.interrupted else 'finished'} at {hypernetwork.step} steps. +Hypernetwork saved to {html.escape(filename)} +""" + return res, "" + except Exception: + raise + finally: + shared.sd_model.cond_stage_model.to(devices.device) + shared.sd_model.first_stage_model.to(devices.device) + sd_hijack.apply_optimizations() + diff --git a/stable-diffusion-webui/modules/images.py b/stable-diffusion-webui/modules/images.py new file mode 100755 index 0000000..74b4356 --- /dev/null +++ b/stable-diffusion-webui/modules/images.py @@ -0,0 +1,877 @@ +from __future__ import annotations + +import datetime +import functools +import pytz +import io +import math +import os +from collections import namedtuple +import re + +import numpy as np +import piexif +import piexif.helper +from PIL import Image, ImageFont, ImageDraw, ImageColor, PngImagePlugin, ImageOps +# pillow_avif needs to be imported somewhere in code for it to work +import pillow_avif # noqa: F401 +import string +import json +import hashlib + +from modules import sd_samplers, shared, script_callbacks, errors +from modules.paths_internal import roboto_ttf_file +from modules.shared import opts + +LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS) + + +def get_font(fontsize: int): + try: + return ImageFont.truetype(opts.font or roboto_ttf_file, fontsize) + except Exception: + return ImageFont.truetype(roboto_ttf_file, fontsize) + + +def image_grid(imgs, batch_size=1, rows=None): + if rows is None: + if opts.n_rows > 0: + rows = opts.n_rows + elif opts.n_rows == 0: + rows = batch_size + elif opts.grid_prevent_empty_spots: + rows = math.floor(math.sqrt(len(imgs))) + while len(imgs) % rows != 0: + rows -= 1 + else: + rows = math.sqrt(len(imgs)) + rows = round(rows) + if rows > len(imgs): + rows = len(imgs) + + cols = math.ceil(len(imgs) / rows) + + params = script_callbacks.ImageGridLoopParams(imgs, cols, rows) + script_callbacks.image_grid_callback(params) + + w, h = map(max, zip(*(img.size for img in imgs))) + grid_background_color = ImageColor.getcolor(opts.grid_background_color, 'RGB') + grid = Image.new('RGB', size=(params.cols * w, params.rows * h), color=grid_background_color) + + for i, img in enumerate(params.imgs): + img_w, img_h = img.size + w_offset, h_offset = 0 if img_w == w else (w - img_w) // 2, 0 if img_h == h else (h - img_h) // 2 + grid.paste(img, box=(i % params.cols * w + w_offset, i // params.cols * h + h_offset)) + + return grid + + +class Grid(namedtuple("_Grid", ["tiles", "tile_w", "tile_h", "image_w", "image_h", "overlap"])): + @property + def tile_count(self) -> int: + """ + The total number of tiles in the grid. + """ + return sum(len(row[2]) for row in self.tiles) + + +def split_grid(image: Image.Image, tile_w: int = 512, tile_h: int = 512, overlap: int = 64) -> Grid: + w, h = image.size + + non_overlap_width = tile_w - overlap + non_overlap_height = tile_h - overlap + + cols = math.ceil((w - overlap) / non_overlap_width) + rows = math.ceil((h - overlap) / non_overlap_height) + + dx = (w - tile_w) / (cols - 1) if cols > 1 else 0 + dy = (h - tile_h) / (rows - 1) if rows > 1 else 0 + + grid = Grid([], tile_w, tile_h, w, h, overlap) + for row in range(rows): + row_images = [] + + y = int(row * dy) + + if y + tile_h >= h: + y = h - tile_h + + for col in range(cols): + x = int(col * dx) + + if x + tile_w >= w: + x = w - tile_w + + tile = image.crop((x, y, x + tile_w, y + tile_h)) + + row_images.append([x, tile_w, tile]) + + grid.tiles.append([y, tile_h, row_images]) + + return grid + + +def combine_grid(grid): + def make_mask_image(r): + r = r * 255 / grid.overlap + r = r.astype(np.uint8) + return Image.fromarray(r, 'L') + + mask_w = make_mask_image(np.arange(grid.overlap, dtype=np.float32).reshape((1, grid.overlap)).repeat(grid.tile_h, axis=0)) + mask_h = make_mask_image(np.arange(grid.overlap, dtype=np.float32).reshape((grid.overlap, 1)).repeat(grid.image_w, axis=1)) + + combined_image = Image.new("RGB", (grid.image_w, grid.image_h)) + for y, h, row in grid.tiles: + combined_row = Image.new("RGB", (grid.image_w, h)) + for x, w, tile in row: + if x == 0: + combined_row.paste(tile, (0, 0)) + continue + + combined_row.paste(tile.crop((0, 0, grid.overlap, h)), (x, 0), mask=mask_w) + combined_row.paste(tile.crop((grid.overlap, 0, w, h)), (x + grid.overlap, 0)) + + if y == 0: + combined_image.paste(combined_row, (0, 0)) + continue + + combined_image.paste(combined_row.crop((0, 0, combined_row.width, grid.overlap)), (0, y), mask=mask_h) + combined_image.paste(combined_row.crop((0, grid.overlap, combined_row.width, h)), (0, y + grid.overlap)) + + return combined_image + + +class GridAnnotation: + def __init__(self, text='', is_active=True): + self.text = text + self.is_active = is_active + self.size = None + + +def draw_grid_annotations(im, width, height, hor_texts, ver_texts, margin=0): + + color_active = ImageColor.getcolor(opts.grid_text_active_color, 'RGB') + color_inactive = ImageColor.getcolor(opts.grid_text_inactive_color, 'RGB') + color_background = ImageColor.getcolor(opts.grid_background_color, 'RGB') + + def wrap(drawing, text, font, line_length): + lines = [''] + for word in text.split(): + line = f'{lines[-1]} {word}'.strip() + if drawing.textlength(line, font=font) <= line_length: + lines[-1] = line + else: + lines.append(word) + return lines + + def draw_texts(drawing, draw_x, draw_y, lines, initial_fnt, initial_fontsize): + for line in lines: + fnt = initial_fnt + fontsize = initial_fontsize + while drawing.multiline_textsize(line.text, font=fnt)[0] > line.allowed_width and fontsize > 0: + fontsize -= 1 + fnt = get_font(fontsize) + drawing.multiline_text((draw_x, draw_y + line.size[1] / 2), line.text, font=fnt, fill=color_active if line.is_active else color_inactive, anchor="mm", align="center") + + if not line.is_active: + drawing.line((draw_x - line.size[0] // 2, draw_y + line.size[1] // 2, draw_x + line.size[0] // 2, draw_y + line.size[1] // 2), fill=color_inactive, width=4) + + draw_y += line.size[1] + line_spacing + + fontsize = (width + height) // 25 + line_spacing = fontsize // 2 + + fnt = get_font(fontsize) + + pad_left = 0 if sum([sum([len(line.text) for line in lines]) for lines in ver_texts]) == 0 else width * 3 // 4 + + cols = im.width // width + rows = im.height // height + + assert cols == len(hor_texts), f'bad number of horizontal texts: {len(hor_texts)}; must be {cols}' + assert rows == len(ver_texts), f'bad number of vertical texts: {len(ver_texts)}; must be {rows}' + + calc_img = Image.new("RGB", (1, 1), color_background) + calc_d = ImageDraw.Draw(calc_img) + + for texts, allowed_width in zip(hor_texts + ver_texts, [width] * len(hor_texts) + [pad_left] * len(ver_texts)): + items = [] + texts + texts.clear() + + for line in items: + wrapped = wrap(calc_d, line.text, fnt, allowed_width) + texts += [GridAnnotation(x, line.is_active) for x in wrapped] + + for line in texts: + bbox = calc_d.multiline_textbbox((0, 0), line.text, font=fnt) + line.size = (bbox[2] - bbox[0], bbox[3] - bbox[1]) + line.allowed_width = allowed_width + + hor_text_heights = [sum([line.size[1] + line_spacing for line in lines]) - line_spacing for lines in hor_texts] + ver_text_heights = [sum([line.size[1] + line_spacing for line in lines]) - line_spacing * len(lines) for lines in ver_texts] + + pad_top = 0 if sum(hor_text_heights) == 0 else max(hor_text_heights) + line_spacing * 2 + + result = Image.new("RGB", (im.width + pad_left + margin * (cols-1), im.height + pad_top + margin * (rows-1)), color_background) + + for row in range(rows): + for col in range(cols): + cell = im.crop((width * col, height * row, width * (col+1), height * (row+1))) + result.paste(cell, (pad_left + (width + margin) * col, pad_top + (height + margin) * row)) + + d = ImageDraw.Draw(result) + + for col in range(cols): + x = pad_left + (width + margin) * col + width / 2 + y = pad_top / 2 - hor_text_heights[col] / 2 + + draw_texts(d, x, y, hor_texts[col], fnt, fontsize) + + for row in range(rows): + x = pad_left / 2 + y = pad_top + (height + margin) * row + height / 2 - ver_text_heights[row] / 2 + + draw_texts(d, x, y, ver_texts[row], fnt, fontsize) + + return result + + +def draw_prompt_matrix(im, width, height, all_prompts, margin=0): + prompts = all_prompts[1:] + boundary = math.ceil(len(prompts) / 2) + + prompts_horiz = prompts[:boundary] + prompts_vert = prompts[boundary:] + + hor_texts = [[GridAnnotation(x, is_active=pos & (1 << i) != 0) for i, x in enumerate(prompts_horiz)] for pos in range(1 << len(prompts_horiz))] + ver_texts = [[GridAnnotation(x, is_active=pos & (1 << i) != 0) for i, x in enumerate(prompts_vert)] for pos in range(1 << len(prompts_vert))] + + return draw_grid_annotations(im, width, height, hor_texts, ver_texts, margin) + + +def resize_image(resize_mode, im, width, height, upscaler_name=None): + """ + Resizes an image with the specified resize_mode, width, and height. + + Args: + resize_mode: The mode to use when resizing the image. + 0: Resize the image to the specified width and height. + 1: Resize the image to fill the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, cropping the excess. + 2: Resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, filling empty with data from image. + im: The image to resize. + width: The width to resize the image to. + height: The height to resize the image to. + upscaler_name: The name of the upscaler to use. If not provided, defaults to opts.upscaler_for_img2img. + """ + + upscaler_name = upscaler_name or opts.upscaler_for_img2img + + def resize(im, w, h): + if upscaler_name is None or upscaler_name == "None" or im.mode == 'L': + return im.resize((w, h), resample=LANCZOS) + + scale = max(w / im.width, h / im.height) + + if scale > 1.0: + upscalers = [x for x in shared.sd_upscalers if x.name == upscaler_name] + if len(upscalers) == 0: + upscaler = shared.sd_upscalers[0] + print(f"could not find upscaler named {upscaler_name or ''}, using {upscaler.name} as a fallback") + else: + upscaler = upscalers[0] + + im = upscaler.scaler.upscale(im, scale, upscaler.data_path) + + if im.width != w or im.height != h: + im = im.resize((w, h), resample=LANCZOS) + + return im + + if resize_mode == 0: + res = resize(im, width, height) + + elif resize_mode == 1: + ratio = width / height + src_ratio = im.width / im.height + + src_w = width if ratio > src_ratio else im.width * height // im.height + src_h = height if ratio <= src_ratio else im.height * width // im.width + + resized = resize(im, src_w, src_h) + res = Image.new("RGB", (width, height)) + res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2)) + + else: + ratio = width / height + src_ratio = im.width / im.height + + src_w = width if ratio < src_ratio else im.width * height // im.height + src_h = height if ratio >= src_ratio else im.height * width // im.width + + resized = resize(im, src_w, src_h) + res = Image.new("RGB", (width, height)) + res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2)) + + if ratio < src_ratio: + fill_height = height // 2 - src_h // 2 + if fill_height > 0: + res.paste(resized.resize((width, fill_height), box=(0, 0, width, 0)), box=(0, 0)) + res.paste(resized.resize((width, fill_height), box=(0, resized.height, width, resized.height)), box=(0, fill_height + src_h)) + elif ratio > src_ratio: + fill_width = width // 2 - src_w // 2 + if fill_width > 0: + res.paste(resized.resize((fill_width, height), box=(0, 0, 0, height)), box=(0, 0)) + res.paste(resized.resize((fill_width, height), box=(resized.width, 0, resized.width, height)), box=(fill_width + src_w, 0)) + + return res + + +if not shared.cmd_opts.unix_filenames_sanitization: + invalid_filename_chars = '#<>:"/\\|?*\n\r\t' +else: + invalid_filename_chars = '/' +invalid_filename_prefix = ' ' +invalid_filename_postfix = ' .' +re_nonletters = re.compile(r'[\s' + string.punctuation + ']+') +re_pattern = re.compile(r"(.*?)(?:\[([^\[\]]+)\]|$)") +re_pattern_arg = re.compile(r"(.*)<([^>]*)>$") +max_filename_part_length = shared.cmd_opts.filenames_max_length +NOTHING_AND_SKIP_PREVIOUS_TEXT = object() + + +def sanitize_filename_part(text, replace_spaces=True): + if text is None: + return None + + if replace_spaces: + text = text.replace(' ', '_') + + text = text.translate({ord(x): '_' for x in invalid_filename_chars}) + text = text.lstrip(invalid_filename_prefix)[:max_filename_part_length] + text = text.rstrip(invalid_filename_postfix) + return text + + +@functools.cache +def get_scheduler_str(sampler_name, scheduler_name): + """Returns {Scheduler} if the scheduler is applicable to the sampler""" + if scheduler_name == 'Automatic': + config = sd_samplers.find_sampler_config(sampler_name) + scheduler_name = config.options.get('scheduler', 'Automatic') + return scheduler_name.capitalize() + + +@functools.cache +def get_sampler_scheduler_str(sampler_name, scheduler_name): + """Returns the '{Sampler} {Scheduler}' if the scheduler is applicable to the sampler""" + return f'{sampler_name} {get_scheduler_str(sampler_name, scheduler_name)}' + + +def get_sampler_scheduler(p, sampler): + """Returns '{Sampler} {Scheduler}' / '{Scheduler}' / 'NOTHING_AND_SKIP_PREVIOUS_TEXT'""" + if hasattr(p, 'scheduler') and hasattr(p, 'sampler_name'): + if sampler: + sampler_scheduler = get_sampler_scheduler_str(p.sampler_name, p.scheduler) + else: + sampler_scheduler = get_scheduler_str(p.sampler_name, p.scheduler) + return sanitize_filename_part(sampler_scheduler, replace_spaces=False) + return NOTHING_AND_SKIP_PREVIOUS_TEXT + + +class FilenameGenerator: + replacements = { + 'basename': lambda self: self.basename or 'img', + 'seed': lambda self: self.seed if self.seed is not None else '', + 'seed_first': lambda self: self.seed if self.p.batch_size == 1 else self.p.all_seeds[0], + 'seed_last': lambda self: NOTHING_AND_SKIP_PREVIOUS_TEXT if self.p.batch_size == 1 else self.p.all_seeds[-1], + 'steps': lambda self: self.p and self.p.steps, + 'cfg': lambda self: self.p and self.p.cfg_scale, + 'width': lambda self: self.image.width, + 'height': lambda self: self.image.height, + 'styles': lambda self: self.p and sanitize_filename_part(", ".join([style for style in self.p.styles if not style == "None"]) or "None", replace_spaces=False), + 'sampler': lambda self: self.p and sanitize_filename_part(self.p.sampler_name, replace_spaces=False), + 'sampler_scheduler': lambda self: self.p and get_sampler_scheduler(self.p, True), + 'scheduler': lambda self: self.p and get_sampler_scheduler(self.p, False), + 'model_hash': lambda self: getattr(self.p, "sd_model_hash", shared.sd_model.sd_model_hash), + 'model_name': lambda self: sanitize_filename_part(shared.sd_model.sd_checkpoint_info.name_for_extra, replace_spaces=False), + 'date': lambda self: datetime.datetime.now().strftime('%Y-%m-%d'), + 'datetime': lambda self, *args: self.datetime(*args), # accepts formats: [datetime], [datetime], [datetime