Streamline the project, then update the model.
This commit is contained in:
flybirdxx
@@ -63,6 +63,7 @@ ComfyUI will automatically install the dependencies in `requirements.txt` on sta
|
||||
- timm
|
||||
- einops
|
||||
- lazyconfig
|
||||
- safetensors
|
||||
|
||||
### 3. Automatic Model Download
|
||||
|
||||
@@ -71,7 +72,7 @@ ComfyUI will automatically install the dependencies in `requirements.txt` on sta
|
||||
The first time you use the `Apply SDMatte` node, it will automatically check for and download the necessary model weights from Hugging Face. The models will be stored in:
|
||||
`ComfyUI/models/SDMatte/`
|
||||
|
||||
You can select between the standard (`SDMatte.pth`) and enhanced (`SDMatte_plus.pth`) versions directly within the node.
|
||||
You can select between the standard (`SDMatte.safetensors`) and enhanced (`SDMatte_plus.safetensors`) versions directly within the node.
|
||||
|
||||
### 4. Restart ComfyUI
|
||||
|
||||
@@ -85,7 +86,7 @@ Restart ComfyUI to load the new custom nodes.
|
||||
|
||||
- **Function**: Loads the model and applies it for matting in a single node.
|
||||
- **Input**:
|
||||
- `ckpt_name`: Select the model to use (`SDMatte.pth` or `SDMatte_plus.pth`). It will be downloaded automatically if not found.
|
||||
- `ckpt_name`: Select the model to use (`SDMatte.safetensors` or `SDMatte_plus.safetensors`). It will be downloaded automatically if not found.
|
||||
- `image`: Input image (ComfyUI IMAGE format)
|
||||
- `trimap`: Trimap mask (ComfyUI MASK format)
|
||||
- `inference_size`: Inference resolution (512/640/768/896/1024)
|
||||
@@ -144,7 +145,7 @@ The plugin has built-in memory optimization strategies (automatically enabled):
|
||||
A: Ensure the plugin directory structure is correct, restart ComfyUI, check console for error messages.
|
||||
|
||||
### Q: Model loading failed?
|
||||
A: Check SDMatte.pth file path, ensure base model directory structure is complete, view console for detailed error messages.
|
||||
A: Check SDMatte.safetensors file path, ensure base model directory structure is complete, view console for detailed error messages.
|
||||
|
||||
### Q: Insufficient VRAM during inference?
|
||||
A: Try reducing inference resolution, enable `force_cpu` option, or close other VRAM-consuming programs.
|
||||
@@ -157,8 +158,8 @@ A: First run needs to compile CUDA kernels, subsequent inference will be signifi
|
||||
|
||||
### Q: Which model version should I choose?
|
||||
A:
|
||||
- **SDMatte.pth (Standard)**: Smaller file (~11GB), faster inference, suitable for most scenarios
|
||||
- **SDMatte_plus.pth (Enhanced)**: Larger file, higher accuracy, suitable for professional use with extremely high quality requirements
|
||||
- **SDMatte.safetensors (Standard)**: Smaller file (~11GB), faster inference, suitable for most scenarios
|
||||
- **SDMatte_plus.safetensors (Enhanced)**: Larger file, higher accuracy, suitable for professional use with extremely high quality requirements
|
||||
- Recommend testing with standard version first, upgrade to enhanced version if higher quality is needed
|
||||
|
||||
## 📋 System Requirements
|
||||
@@ -167,10 +168,23 @@ A:
|
||||
- **Python**: 3.8+
|
||||
- **PyTorch**: 1.12+ (CUDA support recommended)
|
||||
- **VRAM**: 8GB+ recommended (CPU inference supported)
|
||||
- **Dependencies**: diffusers, timm, einops, lazyconfig
|
||||
- **Dependencies**: diffusers, timm, einops, lazyconfig, safetensors
|
||||
|
||||
## 📝 Changelog
|
||||
|
||||
### v1.5.0 (2025-01-XX)
|
||||
- 🔄 **Model Format Update**:
|
||||
- Migrated from `.pth` to `.safetensors` format for better security and performance
|
||||
- Updated model download URLs to use Hugging Face repository (1038lab/SDMatte)
|
||||
- Improved model loading with SafeTensors library for safer weight handling
|
||||
- 🔧 **Technical Improvements**:
|
||||
- Enhanced model loading stability with better error handling
|
||||
- Optimized memory usage during model loading process
|
||||
- Improved compatibility with latest ComfyUI versions
|
||||
- 📚 **Documentation Updates**:
|
||||
- Updated installation instructions to reflect new model format
|
||||
- Added information about SafeTensors format benefits
|
||||
|
||||
### v1.3.0 (2025-08-17)
|
||||
- ✨ **New Features**:
|
||||
- Implemented automatic model downloading and checking. Models are now stored in `ComfyUI/models/SDMatte/`.
|
||||
|
||||
+20
-6
@@ -63,6 +63,7 @@ ComfyUI 会在启动时自动安装 `requirements.txt` 中的依赖包：
|
||||
- timm
|
||||
- einops
|
||||
- lazyconfig
|
||||
- safetensors
|
||||
|
||||
### 3. 自动模型下载
|
||||
|
||||
@@ -71,7 +72,7 @@ ComfyUI 会在启动时自动安装 `requirements.txt` 中的依赖包：
|
||||
首次使用 `Apply SDMatte` 节点时，它会自动检查并从 Hugging Face 下载所需的模型权重。模型将被存放在：
|
||||
`ComfyUI/models/SDMatte/`
|
||||
|
||||
您可以直接在节点内选择使用标准版 (`SDMatte.pth`) 或增强版 (`SDMatte_plus.pth`)。
|
||||
您可以直接在节点内选择使用标准版 (`SDMatte.safetensors`) 或增强版 (`SDMatte_plus.safetensors`)。
|
||||
|
||||
### 4. 重启 ComfyUI
|
||||
|
||||
@@ -85,7 +86,7 @@ ComfyUI 会在启动时自动安装 `requirements.txt` 中的依赖包：
|
||||
|
||||
- **功能**：在一个节点内完成模型加载和抠图应用。
|
||||
- **输入**：
|
||||
- `ckpt_name`：选择要使用的模型（`SDMatte.pth` 或 `SDMatte_plus.pth`）。如果本地不存在，将自动下载。
|
||||
- `ckpt_name`：选择要使用的模型（`SDMatte.safetensors` 或 `SDMatte_plus.safetensors`）。如果本地不存在，将自动下载。
|
||||
- `image`：输入图像（ComfyUI IMAGE 格式）
|
||||
- `trimap`：三值图掩码（ComfyUI MASK 格式）
|
||||
- `inference_size`：推理分辨率（512/640/768/896/1024）
|
||||
@@ -144,7 +145,7 @@ ComfyUI 会在启动时自动安装 `requirements.txt` 中的依赖包：
|
||||
A: 确保插件目录结构正确，重启 ComfyUI，检查控制台是否有错误信息。
|
||||
|
||||
### Q: 模型加载失败？
|
||||
A: 检查 SDMatte.pth 文件路径，确保基础模型目录结构完整，查看控制台详细错误信息。
|
||||
A: 检查 SDMatte.safetensors 文件路径，确保基础模型目录结构完整，查看控制台详细错误信息。
|
||||
|
||||
### Q: 推理时显存不足？
|
||||
A: 尝试降低推理分辨率，启用 `force_cpu` 选项，或关闭其他占用显存的程序。
|
||||
@@ -157,8 +158,8 @@ A: 首次运行需要编译 CUDA 内核，后续推理会显著加速。
|
||||
|
||||
### Q: 应该选择哪个模型版本？
|
||||
A:
|
||||
- **SDMatte.pth（标准版）**：文件较小（~11GB），推理速度快，适合大多数场景
|
||||
- **SDMatte_plus.pth（增强版）**：文件较大，精度更高，适合对质量要求极高的专业用途
|
||||
- **SDMatte.safetensors（标准版）**：文件较小（~11GB），推理速度快，适合大多数场景
|
||||
- **SDMatte_plus.safetensors（增强版）**：文件较大，精度更高，适合对质量要求极高的专业用途
|
||||
- 建议先使用标准版测试，如需更高质量再升级到增强版
|
||||
|
||||
## 📋 系统要求
|
||||
@@ -167,10 +168,23 @@ A:
|
||||
- **Python**：3.8+
|
||||
- **PyTorch**：1.12+ （支持 CUDA 推荐）
|
||||
- **显存**：8GB+ 推荐（支持 CPU 推理）
|
||||
- **依赖**：diffusers, timm, einops, lazyconfig
|
||||
- **依赖**：diffusers, timm, einops, lazyconfig, safetensors
|
||||
|
||||
## 📝 版本更新日志
|
||||
|
||||
### v1.5.0 (2025-01-XX)
|
||||
- 🔄 **模型格式更新**：
|
||||
- 从 `.pth` 格式迁移到 `.safetensors` 格式，提供更好的安全性和性能
|
||||
- 更新模型下载链接，使用 Hugging Face 仓库 (1038lab/SDMatte)
|
||||
- 使用 SafeTensors 库改进模型加载，提供更安全的权重处理
|
||||
- 🔧 **技术改进**：
|
||||
- 增强模型加载稳定性，提供更好的错误处理
|
||||
- 优化模型加载过程中的内存使用
|
||||
- 改进与最新 ComfyUI 版本的兼容性
|
||||
- 📚 **文档更新**：
|
||||
- 更新安装说明以反映新的模型格式
|
||||
- 添加 SafeTensors 格式优势的相关信息
|
||||
|
||||
### v1.3.0 (2025-08-17)
|
||||
- ✨ **新增功能**：
|
||||
- 实现模型自动下载与检查功能，模型现在存放于 `ComfyUI/models/SDMatte/` 目录。
|
||||
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
Before Width: | Height: | Size: 256 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 310 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 168 KiB |
@@ -7,4 +7,5 @@ omegaconf>=2.3.0
|
||||
fvcore>=0.1.5.post20221221
|
||||
torchvision>=0.16.0
|
||||
opencv-python>=4.8.0
|
||||
safetensors>=0.4.0
|
||||
|
||||
|
||||
+25
-92
@@ -1,21 +1,16 @@
|
||||
import os
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
from torchvision import transforms
|
||||
|
||||
import folder_paths
|
||||
import comfy
|
||||
from comfy.utils import load_torch_file
|
||||
|
||||
|
||||
# SDMatte Model Downloader
|
||||
# Correctly get the 'models' directory by finding the parent of the checkpoints directory
|
||||
models_dir = os.path.dirname(folder_paths.get_folder_paths("checkpoints")[0])
|
||||
# Get the models directory from ComfyUI
|
||||
models_dir = folder_paths.models_dir
|
||||
MODEL_DIR = os.path.join(models_dir, "SDMatte")
|
||||
MODEL_URLS = {
|
||||
"SDMatte.pth": "https://huggingface.co/LongfeiHuang/SDMatte/resolve/main/SDMatte.pth",
|
||||
"SDMatte_plus.pth": "https://huggingface.co/LongfeiHuang/SDMatte/resolve/main/SDMatte_plus.pth"
|
||||
"SDMatte.safetensors": "https://huggingface.co/1038lab/SDMatte/resolve/main/SDMatte.safetensors",
|
||||
"SDMatte_plus.safetensors": "https://huggingface.co/1038lab/SDMatte/resolve/main/SDMatte_plus.safetensors"
|
||||
}
|
||||
os.makedirs(MODEL_DIR, exist_ok=True)
|
||||
|
||||
@@ -58,7 +53,6 @@ def download_model(model_name, models_dir=MODEL_DIR, model_urls=MODEL_URLS):
|
||||
return target_path
|
||||
|
||||
except (ImportError, ModuleNotFoundError):
|
||||
print("[SDMatte] Warning: 'requests' and 'tqdm' not found. Downloading without progress bar.")
|
||||
import urllib.request
|
||||
try:
|
||||
urllib.request.urlretrieve(url, tmp_path)
|
||||
@@ -66,22 +60,18 @@ def download_model(model_name, models_dir=MODEL_DIR, model_urls=MODEL_URLS):
|
||||
print(f"[SDMatte] Download complete: {target_path}")
|
||||
return target_path
|
||||
except Exception as e_url:
|
||||
print(f"[SDMatte] Error downloading with urllib: {e_url}")
|
||||
if os.path.exists(tmp_path):
|
||||
os.remove(tmp_path)
|
||||
raise
|
||||
except Exception as e:
|
||||
print(f"[SDMatte] Error downloading model: {e}")
|
||||
if os.path.exists(tmp_path):
|
||||
os.remove(tmp_path)
|
||||
raise
|
||||
|
||||
# 延迟导入核心模型，避免在依赖未安装时阻断节点注册
|
||||
SDMatteCore = None
|
||||
|
||||
|
||||
def _resize_norm_image_bchw(image_bchw: torch.Tensor, size_hw=(1024, 1024)) -> torch.Tensor:
|
||||
# image_bchw in [0,1], float32
|
||||
resize = transforms.Resize(size_hw, antialias=True)
|
||||
norm = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
|
||||
x = resize(image_bchw)
|
||||
@@ -125,10 +115,7 @@ class SDMatteApply:
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"force_cpu": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "强制使用CPU推理。仅在GPU显存不足时启用，会显著降低速度"
|
||||
}),
|
||||
"force_cpu": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
}
|
||||
|
||||
@@ -142,7 +129,6 @@ class SDMatteApply:
|
||||
if force_cpu:
|
||||
device = torch.device('cpu')
|
||||
|
||||
# Load model
|
||||
global SDMatteCore
|
||||
if SDMatteCore is None:
|
||||
from .src.modeling.SDMatte.meta_arch import SDMatte as SDMatteCore
|
||||
@@ -152,9 +138,7 @@ class SDMatteApply:
|
||||
required_subdirs = ["text_encoder", "vae", "unet", "scheduler", "tokenizer"]
|
||||
missing = [d for d in required_subdirs if not os.path.isdir(os.path.join(pretrained_repo, d))]
|
||||
if missing:
|
||||
raise FileNotFoundError(
|
||||
f"本地基底缺少目录: {missing}. 期望路径: {pretrained_repo}，需包含 {required_subdirs} 子目录。"
|
||||
)
|
||||
raise FileNotFoundError(f"Missing directories: {missing}. Expected path: {pretrained_repo}")
|
||||
|
||||
sdmatte_model = SDMatteCore(
|
||||
pretrained_model_name_or_path=pretrained_repo,
|
||||
@@ -170,30 +154,12 @@ class SDMatteApply:
|
||||
|
||||
ckpt_path = download_model(ckpt_name)
|
||||
|
||||
# 更健壮的加载与 state_dict 提取
|
||||
state_root = None
|
||||
try:
|
||||
from torch.serialization import add_safe_globals
|
||||
try:
|
||||
from omegaconf.listconfig import ListConfig
|
||||
add_safe_globals([ListConfig])
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
from omegaconf.base import ContainerMetadata
|
||||
add_safe_globals([ContainerMetadata])
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
state_root = torch.load(ckpt_path, map_location='cpu', weights_only=True)
|
||||
except Exception:
|
||||
try:
|
||||
state_root = torch.load(ckpt_path, map_location='cpu', weights_only=False)
|
||||
except TypeError:
|
||||
state_root = torch.load(ckpt_path, map_location='cpu')
|
||||
except Exception:
|
||||
state_root = load_torch_file(ckpt_path)
|
||||
from safetensors import safe_open
|
||||
state_dict = {}
|
||||
with safe_open(ckpt_path, framework="pt", device="cpu") as f:
|
||||
for key in f.keys():
|
||||
state_dict[key] = f.get_tensor(key)
|
||||
state_root = state_dict
|
||||
|
||||
candidate_keys = [
|
||||
'state_dict','model_state_dict','params','weights',
|
||||
@@ -213,40 +179,32 @@ class SDMatteApply:
|
||||
sdmatte_model.eval()
|
||||
sdmatte_model.to(device)
|
||||
|
||||
# 自动显存优化（内置默认设置）
|
||||
if device.type == 'cuda':
|
||||
try:
|
||||
torch.cuda.empty_cache() # 预清理显存
|
||||
torch.cuda.empty_cache()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 启用注意力切片以节省显存
|
||||
try:
|
||||
unet = getattr(sdmatte_model, 'unet', None)
|
||||
if unet is not None and hasattr(unet, 'set_attn_processor'):
|
||||
from diffusers.models.attention_processor import SlicedAttnProcessor
|
||||
# 使用适中的切片大小平衡性能和显存
|
||||
unet.set_attn_processor(SlicedAttnProcessor(slice_size=1))
|
||||
except Exception as e:
|
||||
print(f'[SDMatte] 注意力优化跳过: {e}')
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
B, H, W, C = image.shape
|
||||
orig_h, orig_w = H, W
|
||||
|
||||
# IMAGE: (B,H,W,C)->(B,C,H,W) in [0,1]
|
||||
img_bchw = image.permute(0, 3, 1, 2).contiguous().to(device)
|
||||
img_in = _resize_norm_image_bchw(img_bchw, (int(inference_size), int(inference_size)))
|
||||
|
||||
# 构造必要的数据字典字段
|
||||
is_trans = torch.tensor([1 if is_transparent else 0] * B, device=device)
|
||||
|
||||
data = {"image": img_in, "is_trans": is_trans, "caption": [""] * B}
|
||||
|
||||
def to_b1hw(x):
|
||||
return _resize_mask_b1hw(x.unsqueeze(1).contiguous().to(device), (int(inference_size), int(inference_size)))
|
||||
|
||||
# 处理 trimap 输入
|
||||
# 将[0,1]范围的trimap转换为[-1,1]范围，与训练时保持一致
|
||||
tri = to_b1hw(trimap) * 2 - 1
|
||||
data["trimap"] = tri
|
||||
data["trimap_coords"] = torch.tensor([[0,0,1,1]]*B, dtype=tri.dtype, device=device)
|
||||
@@ -258,68 +216,43 @@ class SDMatteApply:
|
||||
else:
|
||||
pred_alpha = sdmatte_model(data)
|
||||
|
||||
# 还原到原尺寸
|
||||
out = transforms.Resize((orig_h, orig_w))(pred_alpha)
|
||||
out = out.squeeze(1).clamp(0, 1).detach().cpu()
|
||||
|
||||
# 遮罩优化：使用trimap约束来过滤不需要的区域
|
||||
if mask_refine:
|
||||
trimap_cpu = trimap.cpu() # (B, H, W)
|
||||
trimap_cpu = trimap.cpu()
|
||||
|
||||
# 创建严格的前景约束
|
||||
# 只在trimap明确标记为前景(白色区域)的地方保留高alpha值
|
||||
# 在未知区域(灰色)允许模型输出，但在背景区域(黑色)强制为0
|
||||
foreground_regions = trimap_cpu > trimap_constraint # 前景区域
|
||||
background_regions = trimap_cpu < (1.0 - trimap_constraint) # 背景区域
|
||||
unknown_regions = ~(foreground_regions | background_regions) # 未知区域
|
||||
foreground_regions = trimap_cpu > trimap_constraint
|
||||
background_regions = trimap_cpu < (1.0 - trimap_constraint)
|
||||
unknown_regions = ~(foreground_regions | background_regions)
|
||||
|
||||
# 优化alpha遮罩
|
||||
refined_alpha = out.clone()
|
||||
|
||||
# 背景区域强制为0
|
||||
refined_alpha[background_regions] = 0.0
|
||||
|
||||
# 前景区域保持原值或增强
|
||||
refined_alpha[foreground_regions] = torch.clamp(refined_alpha[foreground_regions] * 1.2, 0, 1)
|
||||
|
||||
# 未知区域保持模型输出，但应用阈值过滤
|
||||
alpha_threshold = 0.3 # 过滤掉低置信度的区域
|
||||
alpha_threshold = 0.3
|
||||
low_confidence = (refined_alpha < alpha_threshold) & unknown_regions
|
||||
refined_alpha[low_confidence] = 0.0
|
||||
|
||||
out = refined_alpha
|
||||
|
||||
# 根据输出模式创建不同的抠图结果
|
||||
# out: (B, H, W), image: (B, H, W, C)
|
||||
alpha_expanded = out.unsqueeze(-1) # (B, H, W, 1)
|
||||
alpha_expanded = out.unsqueeze(-1)
|
||||
|
||||
if output_mode == "alpha_only":
|
||||
# 只输出alpha遮罩，图像输出为黑色
|
||||
matted_image = torch.zeros_like(image.cpu())
|
||||
elif output_mode == "matted_rgba":
|
||||
# 创建RGBA图像：RGB通道保持原图，A通道为alpha遮罩
|
||||
# 背景变为透明，只保留前景对象
|
||||
matted_image = torch.cat([
|
||||
image.cpu(), # RGB通道保持原图
|
||||
alpha_expanded.expand(-1, -1, -1, 1) # A通道为遮罩本身
|
||||
image.cpu(),
|
||||
alpha_expanded.expand(-1, -1, -1, 1)
|
||||
], dim=-1)
|
||||
# RGBA模式输出4通道图像，保持透明背景
|
||||
# 注意：ComfyUI可能不支持4通道显示，但数据格式正确
|
||||
elif output_mode == "matted_rgb":
|
||||
# 只保留前景对象，背景变为黑色
|
||||
# 使用trimap的前景信息来更精确地提取对象
|
||||
trimap_cpu = trimap.cpu() # (B, H, W)
|
||||
trimap_expanded = trimap_cpu.unsqueeze(-1) # (B, H, W, 1)
|
||||
|
||||
# 只在trimap标记为前景(>0.8)或未知区域(0.2-0.8)的地方保留原图
|
||||
# 结合alpha遮罩进一步细化
|
||||
trimap_cpu = trimap.cpu()
|
||||
trimap_expanded = trimap_cpu.unsqueeze(-1)
|
||||
foreground_mask = (trimap_expanded > 0.2) & (alpha_expanded > 0.1)
|
||||
matted_image = image.cpu() * foreground_mask.float()
|
||||
else:
|
||||
# 默认：直接应用alpha遮罩
|
||||
matted_image = image.cpu() * alpha_expanded
|
||||
|
||||
# 推理后清理显存
|
||||
if device.type == 'cuda':
|
||||
try:
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
-21
@@ -1,21 +0,0 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2023 SHI Labs
|
||||
|
||||
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.
|
||||
-120
@@ -1,120 +0,0 @@
|
||||
<div align="center">
|
||||
<p align="center">
|
||||
<img src="./assets/logo.png" alt="SDMatte Logo" height="52">
|
||||
</p>
|
||||
|
||||
<h1 align="center">Grafting Diffusion Models for Interactive Matting</h1>
|
||||
|
||||
<p align="center">
|
||||
<b>
|
||||
Longfei Huang<sup>1,2*</sup>, Yu Liang<sup>2*</sup>, Hao Zhang<sup>2</sup>,
|
||||
Jinwei Chen<sup>2</sup>, Wei Dong<sup>2</sup>,<br>
|
||||
Lunde Chen<sup>1</sup>, Wanyu Liu<sup>1</sup>, Bo Li<sup>2</sup>,
|
||||
Peng-Tao Jiang<sup>2†</sup>
|
||||
</b><br>
|
||||
<sup>1</sup> Shanghai University
|
||||
<sup>2</sup> vivo Mobile Communication Co., Ltd.<br>
|
||||
<sup>*</sup> Equal contribution
|
||||
<sup>†</sup> Corresponding author
|
||||
</p>
|
||||
|
||||
[](https://github.com/vivoCameraResearch/SDMatte)
|
||||
[](https://arxiv.org/abs/2508.00443)
|
||||
[](https://opensource.org/license/MIT)
|
||||
[](https://huggingface.co/LongfeiHuang/SDMatte)
|
||||
[](https://github.com/vivoCameraResearch)
|
||||
|
||||
<strong><i>SDMatte</i> is an interactive image matting method based on stable diffusion, which supports three types of visual prompts (points, boxes, and masks) for accurately extracting target objects from natural images.</strong>
|
||||
|
||||
<div style="width: 100%; text-align: center; margin:auto;">
|
||||
<img style="width:100%" src="assets/result_show.png">
|
||||
</div>
|
||||
</div>
|
||||
|
||||
## 📢 News
|
||||
- [2025.08.04] Released evaluation codes.
|
||||
- [2025.08.01] Published the [arXiv preprint](https://arxiv.org/abs/2508.00443).
|
||||
- [2025.07.31] This repo is created.
|
||||
- [2025.06.26] Paper accepted by ICCV 2025.
|
||||
|
||||
## 📖 Overview
|
||||
### Abstract
|
||||
|
||||
Recent interactive matting methods have demonstrated satisfactory performance in capturing the primary regions of objects, but they fall short in extracting fine-grained details in edge regions. Diffusion models trained on billions of image-text pairs, demonstrate exceptional capability in modeling highly complex data distributions and synthesizing realistic texture details, while exhibiting robust text-driven interaction capabilities, making them an attractive solution for interactive matting.
|
||||
To this end, we propose SDMatte, a diffusion-driven interactive matting model, with three key contributions.
|
||||
First, we exploit the powerful priors of the pre-trained U-Net within diffusion models and transform the text-driven interaction mechanism into a visual prompt-driven interaction mechanism to enable interactive matting.
|
||||
Second, we integrate coordinate embeddings of visual prompts and opacity embeddings of objects into U-Net, enhancing SDMatte's sensitivity to spatial position information and opacity information.
|
||||
Third, we propose a masked self-attention mechanism and a visual prompt-driven interaction mechanism that enable the model to focus on areas specified by visual prompts, leading to better performance.
|
||||
Extensive experiments on multiple datasets demonstrate the superior performance of our method, validating its effectiveness in interactive matting.
|
||||
|
||||
### Architecture
|
||||
|
||||
<div align="center">
|
||||
<img src="assets/pipeline.png" width="100%" height="100%"/>
|
||||
</div><br/>
|
||||
|
||||
### Visualization
|
||||
|
||||
<div align="center">
|
||||
<img src="assets/compare.png" width="100%" height="100%"/>
|
||||
</div>
|
||||
|
||||
## 🛠️ Installation
|
||||
|
||||
1. Create a conda virtual env and activate it.
|
||||
|
||||
```
|
||||
conda create -n SDMatte python==3.10
|
||||
conda activate SDMatte
|
||||
```
|
||||
2. Install packages.
|
||||
```
|
||||
cd path/to/SDMatte
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
3. Install [detectron2](https://github.com/facebookresearch/detectron2) , follow its [documentation](https://detectron2.readthedocs.io/en/latest/).
|
||||
For SDMatte, we recommend to build it from latest source code.
|
||||
```
|
||||
python -m pip install 'git+https://github.com/facebookresearch/detectron2.git'
|
||||
```
|
||||
|
||||
## 📂 Dataset
|
||||
|
||||
* To train **SDMatte** and **LiteSDMatte**, please prepare the following datasets: [Composition-1K](https://sites.google.com/view/deepimagematting), [DIS-646](https://github.com/yuhaoliu7456/CVPR2020-HAttMatting), [AM-2K](https://github.com/JizhiziLi/GFM), [UHRSD](https://github.com/iCVTEAM/PGNet), [RefMatte](https://github.com/jizhiziLi/rim), and [BG-20K](https://github.com/JizhiziLi/GFM).
|
||||
|
||||
* To train **SDMatte<sup>*</sup>**, please prepare the following datasets: [Composition-1K](https://sites.google.com/view/deepimagematting), [DIS-646](https://github.com/yuhaoliu7456/CVPR2020-HAttMatting), [AM-2K](https://github.com/JizhiziLi/GFM), [COCO-Matte](https://github.com/XiaRho/SEMat), and [BG-20K](https://github.com/JizhiziLi/GFM).
|
||||
|
||||
* To evaluate **SDMatte**, please prepare the following test datasets: [AIM-500](https://github.com/JizhiziLi/AIM), [AM-2K](https://github.com/JizhiziLi/GFM), [P3M-500](https://github.com/JizhiziLi/P3M), and [RefMatte-RW-100](https://github.com/jizhiziLi/rim).
|
||||
|
||||
* Check lines 15–56 and 529–530 in `SDMatte/data/dataset.py` to modify the data path to your data path.
|
||||
|
||||
## 🔍 Test
|
||||
|
||||
* Download the model weights from the Hugging Face repository [LongfeiHuang/SDMatte](https://huggingface.co/LongfeiHuang/SDMatte)
|
||||
|
||||
* Modify the `pretrained_model_name_or_path` field in `configs/SDMatte.py` to the directory path containing the configuration files, so that the model can be properly initialized.
|
||||
|
||||
* Modify the `CKPT_DIR` parameter in `script/test.sh` to the specific path of the downloaded weight file, and then run the script using `bash script/test.sh`.
|
||||
|
||||
## 📋 TodoList
|
||||
|
||||
- [x] Publish paper on arXiv
|
||||
- [x] Release source code for SDMatte
|
||||
- [x] Release evaluation codes
|
||||
- [ ] Release pretrained weights for SDMatte and SDMatte*
|
||||
- [ ] Release source code for LiteSDMatte
|
||||
- [ ] Release pretrained weights for LiteSDMatte
|
||||
- [ ] Release training code
|
||||
- [ ] Deploy interactive demo using Gradio or on Hugging Face Spaces
|
||||
|
||||
## 📝 License
|
||||
|
||||
This project is licensed under <a rel="license" href="./LICENSE">MIT</a>. Redistribution and use should follow this license.
|
||||
|
||||
## 👏 Acknowledgement
|
||||
|
||||
Our repo is built upon [Stable Diffusion 2](https://github.com/Stability-AI/stablediffusion), [TAESD](https://github.com/madebyollin/taesd), and [BK-SDM](https://github.com/Nota-NetsPresso/BK-SDM). We sincerely thank the authors for their contributions to the community.
|
||||
|
||||
## 📧 Contact
|
||||
|
||||
If you have any questions, please feel free to reach us at `2946399650fly@shu.edu.cn` or ` pt.jiang@vivo.com`.
|
||||
@@ -1,3 +0,0 @@
|
||||
---
|
||||
license: mit
|
||||
---
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 894 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 17 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 315 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 261 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 1.3 MiB |
@@ -1,114 +0,0 @@
|
||||
from easydict import EasyDict
|
||||
from detectron2.config import LazyCall as L
|
||||
from detectron2 import model_zoo
|
||||
from detectron2.solver import WarmupParamScheduler
|
||||
from fvcore.common.param_scheduler import ExponentialParamScheduler
|
||||
from omegaconf import OmegaConf
|
||||
from torch.utils.data import DataLoader
|
||||
from torch.utils.data.distributed import DistributedSampler
|
||||
from data.dataset import DataGenerator
|
||||
from evaluation import MattingEvaluator
|
||||
from modeling.SDMatte import SDMatte
|
||||
|
||||
hy_dict = EasyDict(
|
||||
bsz=1,
|
||||
num_workers=8,
|
||||
num_GPU=8,
|
||||
crop_size=512,
|
||||
epoch_num=10,
|
||||
data_num=24401,
|
||||
losses=["unknown_l1_loss", "known_l1_loss", "known_lap_loss", "unknown_lap_loss"],
|
||||
output_dir="output/",
|
||||
init_checkpoint=None,
|
||||
lr=5e-5,
|
||||
resume=False,
|
||||
psm="gauss",
|
||||
radius=25,
|
||||
distillation=False,
|
||||
distill_losses=["mse_loss"],
|
||||
model_kwargs=EasyDict(
|
||||
pretrained_model_name_or_path="LongfeiHuang/SDMatte",
|
||||
load_weight=False,
|
||||
conv_scale=3,
|
||||
num_inference_steps=1,
|
||||
aux_input="bbox_mask",
|
||||
add_noise=False,
|
||||
use_dis_loss=True,
|
||||
use_aux_input=True,
|
||||
use_coor_input=True,
|
||||
use_attention_mask=True,
|
||||
residual_connection=False,
|
||||
use_encoder_hidden_states=True,
|
||||
use_attention_mask_list=[True, True, True],
|
||||
use_encoder_hidden_states_list=[False, True, False],
|
||||
),
|
||||
)
|
||||
|
||||
max_iter = int(hy_dict.data_num / hy_dict.bsz / hy_dict.num_GPU * hy_dict.epoch_num)
|
||||
val_step = int(max_iter / 10)
|
||||
|
||||
model = L(SDMatte)(**hy_dict.model_kwargs)
|
||||
|
||||
dataloader = OmegaConf.create()
|
||||
train_dataset = L(DataGenerator)(
|
||||
set_list=["Composition-1K", "DIS-646", "AM-2K", "RefMatte", "UHRSD"],
|
||||
phase="train",
|
||||
crop_size=hy_dict.crop_size,
|
||||
psm=hy_dict.psm,
|
||||
radius=hy_dict.radius,
|
||||
)
|
||||
|
||||
aim500_test_dataset = L(DataGenerator)(set_list="AIM-500", phase="test", psm=hy_dict.psm, radius=hy_dict.radius)
|
||||
am2k_test_dataset = L(DataGenerator)(set_list="AM-2K", phase="test", psm=hy_dict.psm, radius=hy_dict.radius)
|
||||
p3m500_test_dataset = L(DataGenerator)(set_list="P3M-500-NP", phase="test", psm=hy_dict.psm, radius=hy_dict.radius)
|
||||
rw100_test_dataset = L(DataGenerator)(set_list="RefMatte_RW_100", phase="test", psm=hy_dict.psm, radius=hy_dict.radius)
|
||||
|
||||
dataloader.train = L(DataLoader)(
|
||||
dataset=train_dataset,
|
||||
batch_size=hy_dict.bsz,
|
||||
shuffle=False,
|
||||
num_workers=hy_dict.num_workers,
|
||||
pin_memory=True,
|
||||
sampler=L(DistributedSampler)(
|
||||
dataset=train_dataset,
|
||||
drop_last=True,
|
||||
),
|
||||
drop_last=True,
|
||||
)
|
||||
|
||||
dataloader.aim500_test = L(DataLoader)(dataset=aim500_test_dataset, shuffle=False, batch_size=1, num_workers=8, pin_memory=True)
|
||||
dataloader.am2k_test = L(DataLoader)(dataset=am2k_test_dataset, shuffle=False, batch_size=1, num_workers=8, pin_memory=True)
|
||||
dataloader.p3m500_test = L(DataLoader)(dataset=p3m500_test_dataset, shuffle=False, batch_size=1, num_workers=8, pin_memory=True)
|
||||
dataloader.rw100_test = L(DataLoader)(dataset=rw100_test_dataset, shuffle=False, batch_size=1, num_workers=8, pin_memory=True)
|
||||
|
||||
dataloader.evaluator = L(MattingEvaluator)()
|
||||
|
||||
train = EasyDict(
|
||||
output_dir=hy_dict.output_dir,
|
||||
init_checkpoint=hy_dict.init_checkpoint,
|
||||
max_iter=max_iter,
|
||||
amp=EasyDict(enabled=False),
|
||||
ddp=EasyDict(
|
||||
broadcast_buffers=True,
|
||||
find_unused_parameters=True,
|
||||
fp16_compression=True,
|
||||
),
|
||||
checkpointer=EasyDict(period=val_step, max_to_keep=100),
|
||||
eval_period=val_step,
|
||||
log_period=10,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
|
||||
optimizer = model_zoo.get_config("common/optim.py").AdamW
|
||||
optimizer.lr = hy_dict.lr
|
||||
|
||||
|
||||
lr_multiplier = L(WarmupParamScheduler)(
|
||||
scheduler=L(ExponentialParamScheduler)(
|
||||
start_value=1.0,
|
||||
decay=5e-3,
|
||||
),
|
||||
warmup_length=200 / train.max_iter,
|
||||
warmup_factor=0.001,
|
||||
)
|
||||
@@ -1,234 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class BinaryCrossEntropyLoss(nn.Module):
|
||||
def __init__(self):
|
||||
super(BinaryCrossEntropyLoss, self).__init__()
|
||||
self.bce_loss = nn.BCEWithLogitsLoss()
|
||||
|
||||
def forward(self, logits, targets):
|
||||
loss = self.bce_loss(logits, targets)
|
||||
return loss
|
||||
|
||||
|
||||
class DiceLoss(nn.Module):
|
||||
def __init__(self, smooth=1):
|
||||
super(DiceLoss, self).__init__()
|
||||
self.smooth = smooth
|
||||
|
||||
def forward(self, logits, targets):
|
||||
logits = torch.sigmoid(logits)
|
||||
intersection = (logits * targets).sum()
|
||||
dice = (2.0 * intersection + self.smooth) / (logits.sum() + targets.sum() + self.smooth)
|
||||
return 1 - dice
|
||||
|
||||
|
||||
class DistillationCriterion(nn.Module):
|
||||
def __init__(self, losses=None):
|
||||
super(DistillationCriterion, self).__init__()
|
||||
self.losses = losses
|
||||
|
||||
def forward(self, s_fp=None, t_fp=None):
|
||||
if "mse_loss" in self.losses:
|
||||
losses = sum(F.mse_loss(a, b, reduction="mean") for a, b in zip(s_fp, t_fp))
|
||||
if "l1_loss" in self.losses:
|
||||
losses += sum(F.l1_loss(a, b) for a, b in zip(s_fp, t_fp))
|
||||
return losses * 0.01
|
||||
|
||||
|
||||
class MattingCriterion(nn.Module):
|
||||
def __init__(self, losses=None, size=512):
|
||||
super(MattingCriterion, self).__init__()
|
||||
self.losses = losses
|
||||
self.size = size
|
||||
|
||||
def mean_flat(self, tensor):
|
||||
"""
|
||||
Take the mean over all non-batch dimensions.
|
||||
"""
|
||||
return tensor.mean(dim=list(range(1, len(tensor.shape))))
|
||||
|
||||
def loss_gradient_penalty(self, sample_map, preds, targets):
|
||||
preds = preds
|
||||
targets = targets
|
||||
|
||||
# sample_map for unknown area
|
||||
scale = sample_map.shape[0] * self.size * self.size / torch.sum(sample_map)
|
||||
|
||||
# gradient in x
|
||||
sobel_x_kernel = torch.tensor([[[[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]]]).type(dtype=preds.type())
|
||||
delta_pred_x = F.conv2d(preds, weight=sobel_x_kernel, padding=1)
|
||||
delta_gt_x = F.conv2d(targets, weight=sobel_x_kernel, padding=1)
|
||||
|
||||
# gradient in y
|
||||
sobel_y_kernel = torch.tensor([[[[-1, -2, -1], [0, 0, 0], [1, 2, 1]]]]).type(dtype=preds.type())
|
||||
delta_pred_y = F.conv2d(preds, weight=sobel_y_kernel, padding=1)
|
||||
delta_gt_y = F.conv2d(targets, weight=sobel_y_kernel, padding=1)
|
||||
|
||||
# loss
|
||||
loss = (
|
||||
F.l1_loss(delta_pred_x * sample_map, delta_gt_x * sample_map) * scale
|
||||
+ F.l1_loss(delta_pred_y * sample_map, delta_gt_y * sample_map) * scale
|
||||
+ 0.01 * torch.mean(torch.abs(delta_pred_x * sample_map)) * scale
|
||||
+ 0.01 * torch.mean(torch.abs(delta_pred_y * sample_map)) * scale
|
||||
)
|
||||
|
||||
return dict(loss_gradient_penalty=loss)
|
||||
|
||||
def lap_loss(self, preds, targets):
|
||||
loss = laplacian_loss(preds, targets)
|
||||
|
||||
return dict(lap_loss=loss)
|
||||
|
||||
def unknown_lap_loss(self, sample_map, preds, targets):
|
||||
if torch.sum(sample_map) == 0:
|
||||
scale = 0
|
||||
else:
|
||||
scale = sample_map.shape[0] * self.size * self.size / torch.sum(sample_map)
|
||||
|
||||
loss = laplacian_loss(preds * sample_map, targets * sample_map) * scale
|
||||
return dict(unknown_lap_loss=loss)
|
||||
|
||||
def known_lap_loss(self, sample_map, preds, targets):
|
||||
new_sample_map = torch.zeros_like(sample_map)
|
||||
new_sample_map[sample_map == 0] = 1
|
||||
|
||||
if torch.sum(new_sample_map) == 0:
|
||||
scale = 0
|
||||
else:
|
||||
scale = new_sample_map.shape[0] * self.size * self.size / torch.sum(new_sample_map)
|
||||
|
||||
loss = laplacian_loss(preds * new_sample_map, targets * new_sample_map) * scale
|
||||
return dict(known_lap_loss=loss)
|
||||
|
||||
def unknown_l1_loss(self, sample_map, preds, targets):
|
||||
if torch.sum(sample_map) == 0:
|
||||
scale = 0
|
||||
else:
|
||||
scale = sample_map.shape[0] * self.size * self.size / torch.sum(sample_map)
|
||||
|
||||
loss = F.l1_loss(preds * sample_map, targets * sample_map) * scale
|
||||
return dict(unknown_l1_loss=loss)
|
||||
|
||||
def known_l1_loss(self, sample_map, preds, targets):
|
||||
new_sample_map = torch.zeros_like(sample_map)
|
||||
new_sample_map[sample_map == 0] = 1
|
||||
|
||||
if torch.sum(new_sample_map) == 0:
|
||||
scale = 0
|
||||
else:
|
||||
scale = new_sample_map.shape[0] * self.size * self.size / torch.sum(new_sample_map)
|
||||
|
||||
loss = F.l1_loss(preds * new_sample_map, targets * new_sample_map) * scale
|
||||
|
||||
if torch.isnan(loss):
|
||||
raise ValueError("The computed loss is NaN. Check the input values or computation.")
|
||||
|
||||
return dict(known_l1_loss=loss)
|
||||
|
||||
def l1_loss(self, preds, targets):
|
||||
|
||||
loss = F.l1_loss(preds, targets)
|
||||
return dict(l1_loss=loss)
|
||||
|
||||
def mse_loss(self, preds, targets):
|
||||
loss = F.mse_loss(preds.float(), targets.float(), reduction="mean")
|
||||
return dict(mse_loss=loss)
|
||||
|
||||
def dice_loss(self, pred, target, smooth=1e-5):
|
||||
target = (target > 0.5).to(torch.uint8)
|
||||
pred_flat = pred.view(-1)
|
||||
target_flat = target.view(-1)
|
||||
|
||||
# Compute the intersection and union
|
||||
intersection = (pred_flat * target_flat).sum() # sum of element-wise multiplication
|
||||
union = pred_flat.sum() + target_flat.sum() # sum of all elements in pred and target
|
||||
|
||||
# Compute the Dice coefficient
|
||||
dice = (2.0 * intersection + smooth) / (union + smooth)
|
||||
|
||||
loss = 1 - dice
|
||||
return dict(dice_loss=loss)
|
||||
|
||||
def bce_loss(self, probs, target):
|
||||
loss_fn = nn.BCELoss()
|
||||
target = (target > 0.5).to(torch.uint8)
|
||||
loss = loss_fn(probs, target)
|
||||
return dict(bce_loss=loss)
|
||||
|
||||
def forward(self, pred=None, label=None, trimap=None):
|
||||
if trimap is not None:
|
||||
sample_map = torch.zeros_like(trimap)
|
||||
sample_map[trimap == 0.0] = 1
|
||||
losses = dict()
|
||||
for k in self.losses:
|
||||
if k in ["unknown_l1_loss", "known_l1_loss", "unknown_lap_loss", "known_lap_loss", "loss_gradient_penalty"]:
|
||||
losses.update(getattr(self, k)(sample_map, pred, label))
|
||||
else:
|
||||
losses.update(getattr(self, k)(pred, label))
|
||||
return losses
|
||||
|
||||
|
||||
# -----------------Laplacian Loss-------------------------#
|
||||
def laplacian_loss(pred, true, max_levels=5):
|
||||
kernel = gauss_kernel(device=pred.device, dtype=pred.dtype)
|
||||
pred_pyramid = laplacian_pyramid(pred, kernel, max_levels)
|
||||
true_pyramid = laplacian_pyramid(true, kernel, max_levels)
|
||||
loss = 0
|
||||
for level in range(max_levels):
|
||||
loss += (2**level) * F.l1_loss(pred_pyramid[level], true_pyramid[level])
|
||||
return loss / max_levels
|
||||
|
||||
|
||||
def laplacian_pyramid(img, kernel, max_levels):
|
||||
current = img
|
||||
pyramid = []
|
||||
for _ in range(max_levels):
|
||||
current = crop_to_even_size(current)
|
||||
down = downsample(current, kernel)
|
||||
up = upsample(down, kernel)
|
||||
diff = current - up
|
||||
pyramid.append(diff)
|
||||
current = down
|
||||
return pyramid
|
||||
|
||||
|
||||
def gauss_kernel(device="cpu", dtype=torch.float32):
|
||||
kernel = torch.tensor(
|
||||
[[1, 4, 6, 4, 1], [4, 16, 24, 16, 4], [6, 24, 36, 24, 6], [4, 16, 24, 16, 4], [1, 4, 6, 4, 1]], device=device, dtype=dtype
|
||||
)
|
||||
kernel /= 256
|
||||
kernel = kernel[None, None, :, :]
|
||||
return kernel
|
||||
|
||||
|
||||
def gauss_convolution(img, kernel):
|
||||
B, C, H, W = img.shape
|
||||
img = img.reshape(B * C, 1, H, W)
|
||||
img = F.pad(img, (2, 2, 2, 2), mode="reflect")
|
||||
img = F.conv2d(img, kernel)
|
||||
img = img.reshape(B, C, H, W)
|
||||
return img
|
||||
|
||||
|
||||
def downsample(img, kernel):
|
||||
img = gauss_convolution(img, kernel)
|
||||
img = img[:, :, ::2, ::2]
|
||||
return img
|
||||
|
||||
|
||||
def upsample(img, kernel):
|
||||
B, C, H, W = img.shape
|
||||
out = torch.zeros((B, C, H * 2, W * 2), device=img.device, dtype=img.dtype)
|
||||
out[:, :, ::2, ::2] = img * 4
|
||||
out = gauss_convolution(out, kernel)
|
||||
return out
|
||||
|
||||
|
||||
def crop_to_even_size(img):
|
||||
H, W = img.shape[2:]
|
||||
H = H - H % 2
|
||||
W = W - W % 2
|
||||
return img[:, :, :H, :W]
|
||||
@@ -1,2 +0,0 @@
|
||||
from .dataset import *
|
||||
from .prefetcher import *
|
||||
-1011
File diff suppressed because it is too large
Load Diff
@@ -1,43 +0,0 @@
|
||||
import torch
|
||||
import cv2
|
||||
|
||||
|
||||
class Prefetcher:
|
||||
"""
|
||||
Modified from the data_prefetcher in https://github.com/NVIDIA/apex/blob/master/examples/imagenet/main_amp.py
|
||||
"""
|
||||
|
||||
def __init__(self, loader):
|
||||
self.orig_loader = loader
|
||||
self.stream = torch.cuda.Stream()
|
||||
self.next_sample = None
|
||||
|
||||
def preload(self):
|
||||
try:
|
||||
self.next_sample = next(self.loader)
|
||||
except (StopIteration, IndexError):
|
||||
self.next_sample = None
|
||||
return
|
||||
|
||||
with torch.cuda.stream(self.stream):
|
||||
for key, value in self.next_sample.items():
|
||||
if isinstance(value, torch.Tensor):
|
||||
self.next_sample[key] = value.cuda(non_blocking=True)
|
||||
|
||||
def __next__(self):
|
||||
torch.cuda.current_stream().wait_stream(self.stream)
|
||||
sample = self.next_sample
|
||||
if sample is not None:
|
||||
for key, value in sample.items():
|
||||
if isinstance(value, torch.Tensor):
|
||||
sample[key].record_stream(torch.cuda.current_stream())
|
||||
self.preload()
|
||||
else:
|
||||
# throw stop exception if there is no more data to perform as a default dataloader
|
||||
raise StopIteration("No samples in loader. example: `iterator = iter(Prefetcher(loader)); " "data = next(iterator)`")
|
||||
return sample
|
||||
|
||||
def __iter__(self):
|
||||
self.loader = iter(self.orig_loader)
|
||||
self.preload()
|
||||
return self
|
||||
@@ -1,40 +0,0 @@
|
||||
o_099cc48f.jpg
|
||||
o_0bcaa573.jpg
|
||||
o_0c8f03f0.jpg
|
||||
o_12c828bc.jpg
|
||||
o_1462ee19.jpg
|
||||
o_1c321c56.jpg
|
||||
o_1ea2b894.jpg
|
||||
o_1f05f13c.jpg
|
||||
o_247d445d.jpg
|
||||
o_2749e288.jpg
|
||||
o_2972dbc2.jpg
|
||||
o_2b9227eb.jpg
|
||||
o_3c0b2617.jpg
|
||||
o_3c65e352.jpg
|
||||
o_428276f8.jpg
|
||||
o_542332e9.jpg
|
||||
o_5a248199.jpg
|
||||
o_5d53cc84.jpg
|
||||
o_5dad52d4.jpg
|
||||
o_660c8ca2.jpg
|
||||
o_6cd9e070.jpg
|
||||
o_7000c826.jpg
|
||||
o_740b4c5d.jpg
|
||||
o_741da4bb.jpg
|
||||
o_87e77b32.jpg
|
||||
o_87ed5d9e.jpg
|
||||
o_8861e554.jpg
|
||||
o_912d43be.jpg
|
||||
o_aaa92d0a.jpg
|
||||
o_b71d875b.jpg
|
||||
o_bbcef72a.jpg
|
||||
o_bd571707.jpg
|
||||
o_bde965af.jpg
|
||||
o_c665e90e.jpg
|
||||
o_cc63fd4a.jpg
|
||||
o_dc288b1a.jpg
|
||||
o_ddae6d29.jpg
|
||||
o_e8ce31a6.jpg
|
||||
o_fabaf6e3.jpg
|
||||
o_fae80c63.jpg
|
||||
@@ -1,127 +0,0 @@
|
||||
035A4301.jpg
|
||||
035A4302.jpg
|
||||
035A4304.jpg
|
||||
035A4305.jpg
|
||||
035A4308.jpg
|
||||
035A4310.jpg
|
||||
035A4311.jpg
|
||||
035A4313.jpg
|
||||
035A4316.jpg
|
||||
035A4320.jpg
|
||||
035A4321.jpg
|
||||
035A4323.jpg
|
||||
035A4326.jpg
|
||||
035A4327.jpg
|
||||
035A4329.jpg
|
||||
035A4330.jpg
|
||||
035A4331.jpg
|
||||
035A4332.jpg
|
||||
035A4334.jpg
|
||||
035A4335.jpg
|
||||
035A4336.jpg
|
||||
035A4338.jpg
|
||||
035A4340.jpg
|
||||
035A4430.jpg
|
||||
035A4431.jpg
|
||||
035A4432.jpg
|
||||
035A4435.jpg
|
||||
035A4436.jpg
|
||||
035A4438.jpg
|
||||
035A4440.jpg
|
||||
035A4441.jpg
|
||||
035A4442.jpg
|
||||
035A4447.jpg
|
||||
035A4448.jpg
|
||||
035A4449.jpg
|
||||
035A4450.jpg
|
||||
035A4452.jpg
|
||||
035A4453.jpg
|
||||
035A4454.jpg
|
||||
035A4455.jpg
|
||||
035A4456.jpg
|
||||
035A4457.jpg
|
||||
035A4458.jpg
|
||||
035A4461.jpg
|
||||
035A4463.jpg
|
||||
035A4476.jpg
|
||||
035A4477.jpg
|
||||
035A4478.jpg
|
||||
035A4508.jpg
|
||||
035A4509.jpg
|
||||
035A4510.jpg
|
||||
035A4512.jpg
|
||||
035A4513.jpg
|
||||
035A4514.jpg
|
||||
035A4515.jpg
|
||||
035A4516.jpg
|
||||
035A4517.jpg
|
||||
035A4523.jpg
|
||||
035A4525.jpg
|
||||
035A4526.jpg
|
||||
035A4527.jpg
|
||||
035A4528.jpg
|
||||
035A4529.jpg
|
||||
035A4530.jpg
|
||||
035A4531.jpg
|
||||
035A4532.jpg
|
||||
035A4533.jpg
|
||||
035A4534.jpg
|
||||
035A4535.jpg
|
||||
035A4536.jpg
|
||||
035A4538.jpg
|
||||
035A4539.jpg
|
||||
035A4540.jpg
|
||||
035A4541.jpg
|
||||
035A4542.jpg
|
||||
035A4543.jpg
|
||||
035A4545.jpg
|
||||
035A4546.jpg
|
||||
035A4547.jpg
|
||||
035A4548.jpg
|
||||
035A4551.jpg
|
||||
035A4552.jpg
|
||||
035A4553.jpg
|
||||
035A4554.jpg
|
||||
035A4555.jpg
|
||||
035A4556.jpg
|
||||
cobweb-1160456_1920.jpg
|
||||
eve-drink-crystal-festive.jpg
|
||||
glass-100967_1920.jpg
|
||||
glass-565914_1920.jpg
|
||||
glass-602889_1920.jpg
|
||||
glass-containers-1205652_1920.jpg
|
||||
ice-winter-weather-cold-42266.jpg
|
||||
income-996554_1920.jpg
|
||||
lace-876415.jpg
|
||||
lace-876415_1280.jpg
|
||||
milk-518067_1920.jpg
|
||||
milk.jpg
|
||||
milkweed-summer-nature-plant-42309.jpg
|
||||
nature-1281639_1920.jpg
|
||||
pexels-photo-106368.jpg
|
||||
pexels-photo-87383.jpg
|
||||
pexels-photo-97364.jpg
|
||||
pexels-photo.jpg
|
||||
service-359183_1920.jpg
|
||||
smoke_00005.jpg
|
||||
smoke_00010.jpg
|
||||
smoke_00015.jpg
|
||||
smoke_00020.jpg
|
||||
smoke_00025.jpg
|
||||
smoke_00030.jpg
|
||||
smoke_00035.jpg
|
||||
smoke_00040.jpg
|
||||
smoke_00045.jpg
|
||||
smoke_00050.jpg
|
||||
smoke_00055.jpg
|
||||
smoke_00060.jpg
|
||||
smoke_00065.jpg
|
||||
smoke_00070.jpg
|
||||
smoke_00075.jpg
|
||||
smoke_00080.jpg
|
||||
smoke_00085.jpg
|
||||
smoke_00090.jpg
|
||||
spider-web-617769_1920.jpg
|
||||
wedding-846926_1920.jpg
|
||||
wedding-dress-1174168.jpg
|
||||
wine-440575_1920.jpg
|
||||
@@ -1,127 +0,0 @@
|
||||
004c8d27c4063952a98616dd3c8ab316.png
|
||||
13(2).png
|
||||
131686738165901828.png
|
||||
3(2).png
|
||||
44c2a0d472b713ff81f33106b920b12a.png
|
||||
538724499685900405.png
|
||||
539641011387760661.png
|
||||
55b3cab8N50e79e9f.png
|
||||
604673748289192179.png
|
||||
606189768665464996.png
|
||||
613fc8224aac4e71620b753d61ad6e01.png
|
||||
624753897218113578.png
|
||||
657454154710122500.png
|
||||
664308724952072193.png
|
||||
6c1029de9bd68913b3b5eaa63716a86c.png
|
||||
79294f2d8502c34d841cb416a15319d4.png
|
||||
GT25.png
|
||||
a6f749195231b7564f2ed5662c8cb9dd.png
|
||||
agata-create-2122jrt79G8-unsplash.png
|
||||
alcohol-beverage-celebration-cocktail-342971.png
|
||||
bride-2819673_1920.png
|
||||
bride-442894_1920.png
|
||||
c736dff95796e1e99b8d5419a60e6078.png
|
||||
c9b7350aaa4faa4c194a6817f428f744.png
|
||||
champagne-2635041_1920.png
|
||||
champagne-glasses-1940275_1920.png
|
||||
clear-stemmed-glass-almost-full-170366.png
|
||||
clear-wine-glass-with-red-wine-3620231.png
|
||||
da4e9e7c748f8acc039921c468f55e82.png
|
||||
glass-100967_1920.png
|
||||
glass-783678_1920.png
|
||||
glass-783680_1920.png
|
||||
glass-of-wine-1973136_1920.png
|
||||
h_44.png
|
||||
long-stem-wine-glass-filled-with-red-liquid-1682888.png
|
||||
still-2607661_1920.png
|
||||
t_0.png
|
||||
t_10.png
|
||||
t_101.png
|
||||
t_11.png
|
||||
t_12.png
|
||||
t_13.png
|
||||
t_14.png
|
||||
t_15.png
|
||||
t_16.png
|
||||
t_17.png
|
||||
t_27.png
|
||||
t_29.png
|
||||
t_30.png
|
||||
t_31.png
|
||||
t_32.png
|
||||
t_33.png
|
||||
t_35.png
|
||||
t_36.png
|
||||
t_38.png
|
||||
t_39.png
|
||||
t_4.png
|
||||
t_40.png
|
||||
t_41.png
|
||||
t_44.png
|
||||
t_5.png
|
||||
t_52.png
|
||||
t_53.png
|
||||
t_54.png
|
||||
t_55.png
|
||||
t_56.png
|
||||
t_57.png
|
||||
t_58.png
|
||||
t_59.png
|
||||
t_60.png
|
||||
t_61.png
|
||||
t_62.png
|
||||
t_63.png
|
||||
t_64.png
|
||||
t_65.png
|
||||
t_66.png
|
||||
t_67.png
|
||||
t_68.png
|
||||
t_69.png
|
||||
t_7.png
|
||||
t_70.png
|
||||
t_72.png
|
||||
t_73.png
|
||||
t_74.png
|
||||
t_75.png
|
||||
t_76.png
|
||||
t_77.png
|
||||
t_78.png
|
||||
t_79.png
|
||||
t_8.png
|
||||
t_80.png
|
||||
t_81.png
|
||||
t_82.png
|
||||
t_83.png
|
||||
t_84.png
|
||||
t_85.png
|
||||
t_86.png
|
||||
t_87.png
|
||||
t_88.png
|
||||
t_89.png
|
||||
t_90.png
|
||||
t_91.png
|
||||
test_3.png
|
||||
test_33.png
|
||||
test_39.png
|
||||
test_44.png
|
||||
test_48.png
|
||||
train_145.png
|
||||
train_480.png
|
||||
train_488.png
|
||||
train_539.png
|
||||
train_557.png
|
||||
train_568.png
|
||||
train_570.png
|
||||
water-1061021_1920.png
|
||||
water-1821660_1920.png
|
||||
wedding-dress-366543_1920.png
|
||||
wine-1543170_1280.png
|
||||
wine-2160516_1920.png
|
||||
wine-224039_1920.png
|
||||
wine-3900655_1920.png
|
||||
wine-4778099_1920.png
|
||||
wine-4799372_1920.png
|
||||
wine-glass-filled-with-water-1088542.png
|
||||
wine-glasses-2300713_1280_1.png
|
||||
wine-glasses-2300713_1280_2.png
|
||||
wine-glasses-2300713_1280_3.png
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -1,148 +0,0 @@
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
import argparse
|
||||
from tqdm import tqdm
|
||||
from detectron2.evaluation import DatasetEvaluator
|
||||
import numpy as np
|
||||
import torch
|
||||
from utils import compute_mse_loss, compute_sad_loss, compute_mad_loss, compute_gradient_loss, compute_connectivity_error
|
||||
from data import DATA_TEST_ARGS, DATA_TEST_PATH
|
||||
|
||||
|
||||
class MattingEvaluator(DatasetEvaluator):
|
||||
def __init__(self):
|
||||
self.pred = []
|
||||
self.gt = []
|
||||
|
||||
def reset(self):
|
||||
self.pred = []
|
||||
self.gt = []
|
||||
|
||||
def process(self, input, output):
|
||||
"""
|
||||
collect model_output and gt
|
||||
"""
|
||||
H, W = input["hw"][0].item(), input["hw"][1].item()
|
||||
gt = (input["alpha"] + 1.0) / 2.0
|
||||
gt = gt.flatten(0, 2).detach().cpu().numpy() * 255
|
||||
pred = output.flatten(0, 2).detach().cpu().numpy() * 255
|
||||
gt = cv2.resize(gt, (W, H), interpolation=cv2.INTER_LINEAR).astype(np.uint8)
|
||||
pred = cv2.resize(pred, (W, H), interpolation=cv2.INTER_LINEAR).astype(np.uint8)
|
||||
self.pred.append(pred)
|
||||
self.gt.append(gt)
|
||||
|
||||
def evaluate(self):
|
||||
mse_loss = []
|
||||
mad_loss = []
|
||||
sad_loss = []
|
||||
# grad_loss = []
|
||||
# conn_loss = []
|
||||
for idx in tqdm(range(len(self.gt))):
|
||||
label = self.gt[idx].astype(np.float32)
|
||||
pred = self.pred[idx].astype(np.float32)
|
||||
|
||||
if pred.shape != label.shape:
|
||||
pred = cv2.resize(pred, (label.shape[1], label.shape[0]), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
# calculate loss
|
||||
mse_loss_ = compute_mse_loss(pred, label)
|
||||
mad_loss_ = compute_mad_loss(pred, label)
|
||||
sad_loss_ = compute_sad_loss(pred, label)
|
||||
# grad_loss_ = compute_gradient_loss(pred, label)
|
||||
# conn_loss_ = compute_connectivity_error(pred, label)
|
||||
|
||||
# save for average
|
||||
mse_loss.append(mse_loss_)
|
||||
mad_loss.append(mad_loss_)
|
||||
sad_loss.append(sad_loss_)
|
||||
# grad_loss.append(grad_loss_)
|
||||
# conn_loss.append(conn_loss_)
|
||||
|
||||
mse = np.array(mse_loss).mean() * 1000
|
||||
mad = np.array(mad_loss).mean() * 1000
|
||||
sad = np.array(sad_loss).mean()
|
||||
# grad = np.array(grad_loss).mean()
|
||||
# conn = np.array(conn_loss).mean()
|
||||
results = {
|
||||
"MSE": mse,
|
||||
"MAD": mad,
|
||||
"SAD": sad,
|
||||
# "Grad": grad,
|
||||
# "Conn": conn,
|
||||
}
|
||||
return results
|
||||
|
||||
|
||||
def evaluate(args):
|
||||
mse_loss = []
|
||||
mad_loss = []
|
||||
sad_loss = []
|
||||
grad_loss = []
|
||||
conn_loss = []
|
||||
|
||||
label_dir = os.path.join(DATA_TEST_PATH[args.setname], DATA_TEST_ARGS[args.setname][1])
|
||||
name_list = os.listdir(label_dir)
|
||||
if os.path.exists(args.result_path):
|
||||
with open(args.result_path, "r+") as f:
|
||||
if f.read().strip(): # 检查文件是否非空
|
||||
f.seek(0) # 回到文件开头
|
||||
f.truncate()
|
||||
for name in tqdm(name_list):
|
||||
alpha_path = os.path.join(label_dir, name)
|
||||
pred_path = os.path.join(args.pred_dir, name)
|
||||
label = cv2.imread(alpha_path, 0).astype(np.float32)
|
||||
pred = cv2.imread(pred_path, 0).astype(np.float32)
|
||||
if pred.shape != label.shape:
|
||||
pred = cv2.resize(pred, (label.shape[1], label.shape[0]), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
# calculate loss
|
||||
mse_loss_ = compute_mse_loss(pred, label) * 1000
|
||||
mad_loss_ = compute_mad_loss(pred, label) * 1000
|
||||
sad_loss_ = compute_sad_loss(pred, label)
|
||||
grad_loss_ = compute_gradient_loss(pred, label)
|
||||
conn_loss_ = compute_connectivity_error(pred, label)
|
||||
|
||||
# save for average
|
||||
mse_loss.append(mse_loss_) # mean l2 loss per unknown pixel
|
||||
mad_loss.append(mad_loss_) # l1 loss on unknown area
|
||||
sad_loss.append(sad_loss_)
|
||||
grad_loss.append(grad_loss_)
|
||||
conn_loss.append(conn_loss_)
|
||||
results_line = (
|
||||
f"{name}: MSE: {mse_loss_:.2f} MAD: {mad_loss_:.2f} SAD: {sad_loss_:.2f} Grad: {grad_loss_:.2f} Conn: {conn_loss_:.2f}"
|
||||
)
|
||||
with open(args.result_path, "a") as f:
|
||||
f.write(f"\n{results_line}")
|
||||
|
||||
mse = np.array(mse_loss).mean()
|
||||
mad = np.array(mad_loss).mean()
|
||||
sad = np.array(sad_loss).mean()
|
||||
grad = np.array(grad_loss).mean()
|
||||
conn = np.array(conn_loss).mean()
|
||||
|
||||
print(f"MSE: {mse:.2f} MAD: {mad:.2f} SAD: {sad:.2f} Grad: {grad:.2f} Conn: {conn:.2f}")
|
||||
# print(f"MSE: {mse:.2f} MAD: {mad:.2f} SAD: {sad:.2f}")
|
||||
with open(args.result_path, "a") as f:
|
||||
f.write(f"\n{args.setname}: MSE: {mse:.2f} MAD: {mad:.2f} SAD: {sad:.2f} Grad: {grad:.2f} Conn: {conn:.2f}")
|
||||
with open(args.result_path, "r") as f:
|
||||
lines = f.readlines()
|
||||
sorted_lines = sorted(
|
||||
lines,
|
||||
key=lambda line: (
|
||||
0 if args.setname in line else 1,
|
||||
-float(line.split("MSE:")[1].split()[0]) if "MSE:" in line else float("-inf"),
|
||||
),
|
||||
)
|
||||
with open(args.result_path, "w") as f:
|
||||
f.writelines(sorted_lines)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--pred-dir", type=str, default="infer_output/", required=True, help="output dir")
|
||||
parser.add_argument("--setname", type=str, default="dataset/test/alphas/", help="testset name")
|
||||
parser.add_argument("--result-path", type=str, help="save the evl result")
|
||||
args = parser.parse_args()
|
||||
|
||||
evaluate(args)
|
||||
@@ -1,122 +0,0 @@
|
||||
"""
|
||||
Inference for Composition-1k Dataset.
|
||||
|
||||
Run:
|
||||
python inference.py \
|
||||
--config-dir path/to/config
|
||||
--checkpoint-dir path/to/checkpoint
|
||||
--inference-dir path/to/inference
|
||||
--data-dir path/to/data
|
||||
"""
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
from torch.utils.data import Dataset, DataLoader
|
||||
from torchvision.transforms import functional as F
|
||||
from os.path import join
|
||||
from detectron2.checkpoint import DetectionCheckpointer
|
||||
from detectron2.config import LazyConfig, instantiate
|
||||
from detectron2.engine import default_argument_parser
|
||||
import warnings
|
||||
from data.dataset import DataGenerator
|
||||
import time
|
||||
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
|
||||
# model and output
|
||||
def matting_inference(
|
||||
config_dir="",
|
||||
checkpoint_dir="",
|
||||
inference_dir="",
|
||||
setname="",
|
||||
):
|
||||
# initializing model
|
||||
torch.set_grad_enabled(False)
|
||||
cfg = LazyConfig.load(config_dir)
|
||||
model = instantiate(cfg.model)
|
||||
model.to(cfg.train.device)
|
||||
DetectionCheckpointer(model).load(checkpoint_dir)
|
||||
model.eval()
|
||||
|
||||
# initializing dataset
|
||||
test_dataset = DataGenerator(set_list=setname, phase="test", psm=cfg.hy_dict.psm, radius=cfg.hy_dict.radius)
|
||||
test_loader = DataLoader(dataset=test_dataset, shuffle=False, batch_size=1, num_workers=8, pin_memory=True)
|
||||
|
||||
# inferencing
|
||||
os.makedirs(inference_dir, exist_ok=True)
|
||||
# image_dir = inference_dir.replace(args.setname, args.setname + "_Image")
|
||||
# os.makedirs(image_dir, exist_ok=True)
|
||||
|
||||
start_time = time.time()
|
||||
for data in tqdm(test_loader):
|
||||
image_name = data["image_name"][0]
|
||||
H, W = data["hw"][0].item(), data["hw"][1].item()
|
||||
|
||||
# # image
|
||||
# image = data["image"]
|
||||
# image = (image + 1) / 2
|
||||
# image = image.squeeze(0).permute(1, 2, 0) * 255
|
||||
# image = cv2.resize(image.detach().cpu().numpy(), (W, H)).astype(np.uint8)
|
||||
# red = np.array([[255.0, 0.0, 0.0]], dtype=image.dtype)
|
||||
# blue = np.array([[0, 0.0, 255.0]], dtype=image.dtype)
|
||||
# if cfg.hy_dict.model_kwargs.aux_input == "bbox_mask":
|
||||
# coords = data["bbox_coords"][0]
|
||||
# x_min, y_min, x_max, y_max = (
|
||||
# int(coords[0].item() * W),
|
||||
# int(coords[1].item() * H),
|
||||
# int(coords[2].item() * W),
|
||||
# int(coords[3].item() * H),
|
||||
# )
|
||||
# image[y_min:y_max, x_min, :] = red
|
||||
# image[y_min:y_max, x_max, :] = red
|
||||
# image[y_min, x_min:x_max, :] = red
|
||||
# image[y_max, x_min:x_max, :] = red
|
||||
# elif cfg.hy_dict.model_kwargs.aux_input == "point_mask":
|
||||
# point_mask = data["point_mask"]
|
||||
# point_mask = (point_mask + 1) / 2
|
||||
# # point_mask = point_mask * 255.0
|
||||
# point_mask = cv2.resize(point_mask.flatten(0, 2).detach().cpu().numpy(), (W, H), interpolation=cv2.INTER_NEAREST)
|
||||
# if len(np.unique(point_mask)) == 3:
|
||||
# image[point_mask == 0] = red
|
||||
# image[point_mask == 1] = blue
|
||||
# elif len(np.unique(point_mask)) == 2:
|
||||
# image[point_mask == 1] = red
|
||||
# image = F.to_pil_image(image).convert("RGB")
|
||||
# image.save(join(image_dir, image_name[:-4] + ".jpg"))
|
||||
|
||||
with torch.no_grad():
|
||||
pred = model(data)
|
||||
output = pred.flatten(0, 2) * 255
|
||||
output = cv2.resize(output.detach().cpu().numpy(), (W, H)).astype(np.uint8)
|
||||
output = F.to_pil_image(output).convert("RGB")
|
||||
output.save(join(inference_dir, image_name))
|
||||
torch.cuda.empty_cache()
|
||||
end_time = time.time()
|
||||
total_time = end_time - start_time
|
||||
avg_time = total_time / len(test_loader)
|
||||
|
||||
print(f"Start time: {start_time:.4f} seconds")
|
||||
print(f"End time: {end_time:.4f} seconds")
|
||||
print(f"Total time: {total_time:.4f} seconds")
|
||||
print(f"Average time per iteration: {avg_time:.4f} seconds")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# add argument we need:
|
||||
parser = default_argument_parser()
|
||||
parser.add_argument("--config-dir", type=str, required=True)
|
||||
parser.add_argument("--checkpoint-dir", type=str, required=True)
|
||||
parser.add_argument("--inference-dir", type=str, required=True)
|
||||
parser.add_argument("--setname", type=str, required=True)
|
||||
|
||||
args = parser.parse_args()
|
||||
matting_inference(
|
||||
config_dir=args.config_dir,
|
||||
checkpoint_dir=args.checkpoint_dir,
|
||||
inference_dir=args.inference_dir,
|
||||
setname=args.setname,
|
||||
)
|
||||
@@ -1,44 +0,0 @@
|
||||
from safetensors.torch import load_file, save_file
|
||||
import torch
|
||||
|
||||
# 加载 .safetensors 权重
|
||||
weights_path = (
|
||||
"stabilityai-stable-diffusion-2/unet/diffusion_pytorch_model.safetensors"
|
||||
)
|
||||
state_dict = load_file(weights_path)
|
||||
|
||||
# 定义缺失键和对应维度
|
||||
missing_keys = [
|
||||
"bbox_embedding.linear_1.weight",
|
||||
"bbox_embedding.linear_1.bias",
|
||||
"bbox_embedding.linear_2.weight",
|
||||
"bbox_embedding.linear_2.bias",
|
||||
"point_embedding.linear_2.weight",
|
||||
"point_embedding.linear_1.weight",
|
||||
"point_embedding.linear_1.bias",
|
||||
"point_embedding.linear_2.bias",
|
||||
]
|
||||
point_dim = 1680
|
||||
bbox_dim = 1280 # 示例，实际需要根据模型架构确定
|
||||
output_dim = 1280 # 示例，实际需要根据模型架构确定
|
||||
|
||||
# 将新增权重初始化为零
|
||||
for key in missing_keys:
|
||||
if "_1.weight" in key and "bbox" in key:
|
||||
state_dict[key] = torch.zeros((output_dim, bbox_dim)) # 初始化权重为零
|
||||
elif "_2.weight" in key and "bbox" in key:
|
||||
state_dict[key] = torch.zeros((output_dim, output_dim)) # 初始化权重为零
|
||||
elif "_1.weight" in key and "point" in key:
|
||||
state_dict[key] = torch.zeros((output_dim, point_dim))
|
||||
elif "_2.weight" in key and "point" in key:
|
||||
state_dict[key] = torch.zeros((output_dim, output_dim)) # 初始化权重为零
|
||||
elif "bias" in key:
|
||||
state_dict[key] = torch.zeros(output_dim) # 初始化偏置为零
|
||||
|
||||
# 保存更新后的权重文件
|
||||
new_weights_path = (
|
||||
"stabilityai-stable-diffusion-2/unet/diffusion_pytorch_model.safetensors"
|
||||
)
|
||||
save_file(state_dict, new_weights_path)
|
||||
|
||||
print(f"Updated weights saved to {new_weights_path}")
|
||||
@@ -1,16 +0,0 @@
|
||||
torch==2.2.0
|
||||
torchvision==0.17.0
|
||||
tensorboard
|
||||
timm==0.5.4
|
||||
opencv-python==4.5.4.60
|
||||
setuptools==58.2.0
|
||||
easydict
|
||||
scikit-image==0.24.0
|
||||
fairscale==0.4.13
|
||||
diffusers==0.30.3
|
||||
transformers==4.45.1
|
||||
numpy==1.24.1
|
||||
Pillow==9.2.0
|
||||
huggingface-hub==0.25.1
|
||||
accelerate==0.34.2
|
||||
cached_property==2.0.1
|
||||
@@ -1,19 +0,0 @@
|
||||
# AIM-500 AM-2K P3M-500-NP RefMatte_RW_100
|
||||
|
||||
SETNAME="AM-2K"
|
||||
CONFIG_DIR="configs/SDMatte.py"
|
||||
CKPT_DIR="SDMatte/SDMatte.pth"
|
||||
INFER_DIR="infer_output/SDMatte/bbox/${SETNAME}"
|
||||
RESULT_PATH="${INFER_DIR}_result.txt"
|
||||
|
||||
# # 使用变量执行命令
|
||||
CUDA_VISIBLE_DEVICES=0 python inference.py \
|
||||
--config-dir "$CONFIG_DIR" \
|
||||
--checkpoint-dir "$CKPT_DIR" \
|
||||
--inference-dir "$INFER_DIR" \
|
||||
--setname "$SETNAME"
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0 python evaluation.py \
|
||||
--pred-dir "$INFER_DIR" \
|
||||
--setname "$SETNAME" \
|
||||
--result-path "$RESULT_PATH"
|
||||
@@ -1,126 +0,0 @@
|
||||
"""
|
||||
Reimplement evaluation.mat provided by Adobe in python
|
||||
Output of `compute_gradient_loss` is sightly different from the MATLAB version provided by Adobe (less than 0.1%)
|
||||
Output of `compute_connectivity_error` is smaller than the MATLAB version (~5%, maybe MATLAB has a different algorithm)
|
||||
So do not report results calculated by these functions in your paper.
|
||||
Evaluate your inference with the MATLAB file `DIM_evaluation_code/evaluate.m`.
|
||||
|
||||
by Yaoyi Li
|
||||
"""
|
||||
|
||||
import scipy.ndimage
|
||||
import numpy as np
|
||||
from skimage.measure import label
|
||||
import scipy.ndimage.morphology
|
||||
|
||||
|
||||
def gauss(x, sigma):
|
||||
y = np.exp(-x ** 2 / (2 * sigma ** 2)) / (sigma * np.sqrt(2 * np.pi))
|
||||
return y
|
||||
|
||||
|
||||
def dgauss(x, sigma):
|
||||
y = -x * gauss(x, sigma) / (sigma ** 2)
|
||||
return y
|
||||
|
||||
|
||||
def gaussgradient(im, sigma):
|
||||
epsilon = 1e-2
|
||||
halfsize = np.ceil(sigma * np.sqrt(-2 * np.log(np.sqrt(2 * np.pi) * sigma * epsilon))).astype(np.int_)
|
||||
size = 2 * halfsize + 1
|
||||
hx = np.zeros((size, size))
|
||||
for i in range(0, size):
|
||||
for j in range(0, size):
|
||||
u = [i - halfsize, j - halfsize]
|
||||
hx[i, j] = gauss(u[0], sigma) * dgauss(u[1], sigma)
|
||||
|
||||
hx = hx / np.sqrt(np.sum(np.abs(hx) * np.abs(hx)))
|
||||
hy = hx.transpose()
|
||||
|
||||
gx = scipy.ndimage.convolve(im, hx, mode='nearest')
|
||||
gy = scipy.ndimage.convolve(im, hy, mode='nearest')
|
||||
|
||||
return gx, gy
|
||||
|
||||
|
||||
def compute_gradient_loss(pred, target, trimap=None):
|
||||
|
||||
pred = pred / 255.0
|
||||
target = target / 255.0
|
||||
|
||||
pred_x, pred_y = gaussgradient(pred, 1.4)
|
||||
target_x, target_y = gaussgradient(target, 1.4)
|
||||
|
||||
pred_amp = np.sqrt(pred_x ** 2 + pred_y ** 2)
|
||||
target_amp = np.sqrt(target_x ** 2 + target_y ** 2)
|
||||
|
||||
error_map = (pred_amp - target_amp) ** 2
|
||||
if trimap is not None:
|
||||
loss = np.sum(error_map[trimap == 128])
|
||||
else:
|
||||
loss = np.sum(error_map)
|
||||
return loss / 1000.
|
||||
|
||||
|
||||
def getLargestCC(segmentation):
|
||||
labels = label(segmentation, connectivity=1)
|
||||
largestCC = labels == np.argmax(np.bincount(labels.flat))
|
||||
return largestCC
|
||||
|
||||
|
||||
def compute_connectivity_error(pred, target, trimap=None, step=0.1):
|
||||
pred = pred / 255.0
|
||||
target = target / 255.0
|
||||
h, w = pred.shape
|
||||
|
||||
thresh_steps = list(np.arange(0, 1 + step, step))
|
||||
l_map = np.ones_like(pred, dtype=np.float32) * -1
|
||||
for i in range(1, len(thresh_steps)):
|
||||
pred_alpha_thresh = (pred >= thresh_steps[i]).astype(np.int_)
|
||||
target_alpha_thresh = (target >= thresh_steps[i]).astype(np.int_)
|
||||
|
||||
omega = getLargestCC(pred_alpha_thresh * target_alpha_thresh).astype(np.int_)
|
||||
flag = ((l_map == -1) & (omega == 0)).astype(np.int_)
|
||||
l_map[flag == 1] = thresh_steps[i - 1]
|
||||
|
||||
l_map[l_map == -1] = 1
|
||||
|
||||
pred_d = pred - l_map
|
||||
target_d = target - l_map
|
||||
pred_phi = 1 - pred_d * (pred_d >= 0.15).astype(np.int_)
|
||||
target_phi = 1 - target_d * (target_d >= 0.15).astype(np.int_)
|
||||
if trimap is not None:
|
||||
loss = np.sum(np.abs(pred_phi - target_phi)[trimap == 128])
|
||||
else:
|
||||
loss = np.sum(np.abs(pred_phi - target_phi))
|
||||
|
||||
return loss / 1000.
|
||||
|
||||
def compute_mse_loss(pred, target, trimap=None):
|
||||
error_map = (pred - target) / 255.0
|
||||
if trimap is not None:
|
||||
loss = np.sum((error_map ** 2) * (trimap == 128)) / (np.sum(trimap == 128) + 1e-8)
|
||||
else:
|
||||
loss = np.sum(error_map**2) / (pred.shape[0] * pred.shape[1])
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
def compute_sad_loss(pred, target, trimap=None):
|
||||
error_map = np.abs((pred - target) / 255.0)
|
||||
if trimap is not None:
|
||||
loss = np.sum(error_map * (trimap == 128))
|
||||
return loss / 1000, np.sum(trimap == 128) / 1000
|
||||
else:
|
||||
loss = np.sum(error_map)
|
||||
return loss / 1000
|
||||
|
||||
|
||||
def compute_mad_loss(pred, target, trimap=None):
|
||||
error_map = np.abs((pred - target) / 255.0)
|
||||
if trimap is not None:
|
||||
loss = np.sum(error_map * (trimap == 128)) / (np.sum(trimap == 128) + 1e-8)
|
||||
else:
|
||||
loss = np.sum(error_map) / (pred.shape[0] * pred.shape[1])
|
||||
|
||||
return loss
|
||||
@@ -1,79 +0,0 @@
|
||||
|
||||
from detectron2.engine.train_loop import HookBase
|
||||
from detectron2.utils.events import EventWriter, get_event_storage
|
||||
from cached_property import cached_property
|
||||
from typing import Optional
|
||||
from detectron2.utils.file_io import PathManager
|
||||
from detectron2.engine.defaults import CommonMetricPrinter, TensorboardXWriter, JSONWriter
|
||||
|
||||
import os
|
||||
import torch
|
||||
|
||||
def power_default_writers(output_dir: str, max_iter: Optional[int] = None):
|
||||
"""
|
||||
Build a list of :class:`EventWriter` to be used.
|
||||
It now consists of a :class:`CommonMetricPrinter`,
|
||||
:class:`TensorboardXWriter` and :class:`JSONWriter`.
|
||||
|
||||
Args:
|
||||
output_dir: directory to store JSON metrics and tensorboard events
|
||||
max_iter: the total number of iterations
|
||||
|
||||
Returns:
|
||||
list[EventWriter]: a list of :class:`EventWriter` objects.
|
||||
"""
|
||||
PathManager.mkdirs(output_dir)
|
||||
return [
|
||||
# It may not always print what you want to see, since it prints "common" metrics only.
|
||||
CommonMetricPrinter(max_iter),
|
||||
JSONWriter(os.path.join(output_dir, "metrics.json")),
|
||||
PowerTensorboardXWriter(output_dir),
|
||||
]
|
||||
|
||||
|
||||
class PowerTensorboardXWriter(TensorboardXWriter):
|
||||
"""
|
||||
Write all scalars to a tensorboard file.
|
||||
"""
|
||||
|
||||
def __init__(self, log_dir: str, window_size: int = 20, **kwargs):
|
||||
"""
|
||||
Args:
|
||||
log_dir (str): the directory to save the output events
|
||||
window_size (int): the scalars will be median-smoothed by this window size
|
||||
|
||||
kwargs: other arguments passed to `torch.utils.tensorboard.SummaryWriter(...)`
|
||||
"""
|
||||
super().__init__(log_dir, window_size, **kwargs)
|
||||
|
||||
def write(self):
|
||||
storage = get_event_storage()
|
||||
new_last_write = self._last_write
|
||||
for k, (v, iter) in storage.latest_with_smoothing_hint(self._window_size).items():
|
||||
if iter > self._last_write:
|
||||
self._writer.add_scalar(k, v, iter)
|
||||
new_last_write = max(new_last_write, iter)
|
||||
self._last_write = new_last_write
|
||||
|
||||
# storage.put_{image,histogram} is only meant to be used by
|
||||
# tensorboard writer. So we access its internal fields directly from here.
|
||||
if len(storage._vis_data) >= 1:
|
||||
for img_name, img, step_num in storage._vis_data:
|
||||
self._writer.add_images(img_name, img, step_num)
|
||||
# Storage stores all image data and rely on this writer to clear them.
|
||||
# As a result it assumes only one writer will use its image data.
|
||||
# An alternative design is to let storage store limited recent
|
||||
# data (e.g. only the most recent image) that all writers can access.
|
||||
# In that case a writer may not see all image data if its period is long.
|
||||
storage.clear_images()
|
||||
|
||||
if len(storage._t) == 2:
|
||||
t, step_num = storage._t
|
||||
self._writer.add_text("t", t, step_num)
|
||||
storage._t = ()
|
||||
|
||||
|
||||
if len(storage._histograms) >= 1:
|
||||
for params in storage._histograms:
|
||||
self._writer.add_histogram_raw(**params)
|
||||
storage.clear_histograms()
|
||||
Reference in New Issue
Block a user