first commit
This commit is contained in:
fengwk
@@ -0,0 +1,3 @@
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
.pytest_cache/
|
||||
@@ -0,0 +1,288 @@
|
||||
# ComfyUI-UniversalRemove
|
||||
|
||||
一个自包含的 ComfyUI remove 核心扩展。
|
||||
|
||||
## 设计目标
|
||||
|
||||
- 不硬编码任何本地路径
|
||||
- 不依赖其他自定义节点
|
||||
- 只使用 ComfyUI 常见核心类型：`IMAGE` / `MASK` / `STRING` / `INT` / `FLOAT`
|
||||
- 面向 **图像批 / 视频帧批** 工作，适合作为通用 remove 流程的中间层
|
||||
|
||||
## 当前提供的节点
|
||||
|
||||
### 0. `UR Demo Image`
|
||||
|
||||
生成一个无需外部图片文件的测试图像，主要用于示例工作流和安装验证。
|
||||
|
||||
### 1. `UR Remove Prompt Builder`
|
||||
|
||||
输入目标描述，输出三路提示词：
|
||||
|
||||
- `segmentation_prompt`
|
||||
- `remove_positive_prompt`
|
||||
- `remove_negative_prompt`
|
||||
|
||||
用途：
|
||||
|
||||
- `segmentation_prompt` 给任意文本选目标模块
|
||||
- `remove_positive_prompt` 给关键帧修复模型
|
||||
- `remove_negative_prompt` 给关键帧修复模型的负向提示
|
||||
|
||||
### 2. `UR Region Editor`
|
||||
|
||||
交互式前端 region 编辑器，输出标准 `regions_json`。
|
||||
|
||||
支持：
|
||||
|
||||
- 加载本地图片作为编辑背景
|
||||
- 加载本地视频作为编辑背景并拖动时间轴
|
||||
- 使用当前视频时间写入关键帧编号
|
||||
- 绘制矩形区域
|
||||
- 绘制圆形区域
|
||||
- 绘制多边形区域
|
||||
- 删除选中区域
|
||||
- 清空全部区域
|
||||
|
||||
注意：背景图片 / 视频只用于前端交互预览，不会写进 workflow。真正可移植、可复用的是输出的 `regions_json`。
|
||||
|
||||
### 3. `UR Boxes To Mask Batch`
|
||||
|
||||
把 JSON 形式的矩形框转换成 `MASK` 批次。
|
||||
|
||||
支持：
|
||||
|
||||
- `xyxy`
|
||||
- `xywh`
|
||||
- 归一化坐标
|
||||
- 三种广播模式：
|
||||
- `repeat_to_batch`
|
||||
- `single_frame_only`
|
||||
- `hold_after_frame`
|
||||
|
||||
用途：
|
||||
|
||||
- 手工框选移除区域
|
||||
- 从外部工具导入 bbox
|
||||
|
||||
### 4. `UR Regions To Mask Batch`
|
||||
|
||||
把关键帧 region JSON 转换成整段图像批 / 视频帧批的 `MASK`。
|
||||
|
||||
支持：
|
||||
|
||||
- `rect`
|
||||
- `circle`
|
||||
- `polygon`
|
||||
- 单关键帧
|
||||
- 多关键帧
|
||||
- 四种广播模式：
|
||||
- `repeat_to_batch`
|
||||
- `single_frame_only`
|
||||
- `hold_after_frame`
|
||||
- `hold_until_next_keyframe`
|
||||
|
||||
用途：
|
||||
|
||||
- 用 JSON 描述“关键帧上的圈选区域”
|
||||
- 为视频帧批生成对应 mask
|
||||
- 后续接 `UR Mask Processor` / `UR Keyframe Prep` / `UR Soft Composite`
|
||||
|
||||
### 5. `UR Mask Processor`
|
||||
|
||||
对 mask 做 remove 常用处理：
|
||||
|
||||
- 扩张
|
||||
- 高斯羽化
|
||||
- 预览叠加
|
||||
|
||||
### 6. `UR Keyframe Prep`
|
||||
|
||||
从图像批中选出关键帧，并准备 remove 输入：
|
||||
|
||||
- 选关键帧
|
||||
- 对应 mask 对齐
|
||||
- 按最长边缩放
|
||||
- pad 到指定倍数
|
||||
- 用填充色生成 `filled_frame`
|
||||
|
||||
其中 `filled_frame` 可以直接送进任意关键帧修复 / inpaint / edit 模型。
|
||||
|
||||
### 7. `UR Soft Composite`
|
||||
|
||||
一个不依赖 BrushNet 的软融合节点，用于把修复结果融合回原图 / 原视频帧批。
|
||||
|
||||
## 推荐工作流
|
||||
|
||||
> 这个扩展刻意只负责 remove 核心流程，不强绑某个检测器、分割器、视频模型或视频 IO 节点。
|
||||
|
||||
### A. 手工框选 remove
|
||||
|
||||
```text
|
||||
Frames/Image Batch
|
||||
-> UR Region Editor
|
||||
-> UR Regions To Mask Batch
|
||||
-> UR Mask Processor
|
||||
-> UR Keyframe Prep
|
||||
-> 关键帧修复模型
|
||||
-> (可选) 视频传播模型
|
||||
-> UR Soft Composite
|
||||
-> Output
|
||||
```
|
||||
|
||||
### B. 提示词驱动 remove
|
||||
|
||||
```text
|
||||
Frames/Image Batch
|
||||
-> UR Remove Prompt Builder
|
||||
-> 任意文本选目标模块（输出 MASK）
|
||||
-> UR Mask Processor
|
||||
-> UR Keyframe Prep
|
||||
-> 关键帧修复模型
|
||||
-> (可选) 视频传播模型
|
||||
-> UR Soft Composite
|
||||
-> Output
|
||||
```
|
||||
|
||||
## 为什么没有内置“文本直接找目标”
|
||||
|
||||
这是刻意的设计边界：
|
||||
|
||||
- remove 核心流程应该通用、可复用
|
||||
- 文本找目标属于检测 / 分割 / 跟踪能力
|
||||
- 这部分通常意味着额外模型、额外权重、额外依赖
|
||||
|
||||
因此当前版本先把 remove 主流程做成稳定、自包含、无外部节点硬依赖的核心层。
|
||||
|
||||
如果后续要继续扩展，最自然的下一步是在这个扩展内部新增：
|
||||
|
||||
- `UR Prompt Mask Locator`
|
||||
- `UR Video Track Mask`
|
||||
- `UR Model Bundle`
|
||||
|
||||
但这些应当作为后续迭代，而不是在 V1 就把依赖链做死。
|
||||
|
||||
## 安装方式
|
||||
|
||||
把整个目录放进：
|
||||
|
||||
- `ComfyUI/custom_nodes/ComfyUI-UniversalRemove`
|
||||
|
||||
然后重启 ComfyUI。
|
||||
|
||||
## 示例工作流
|
||||
|
||||
当前已提供一个可直接导入的最小工作流：
|
||||
|
||||
- `example_workflows/Universal Remove - Interactive Region Editor.json`
|
||||
|
||||
这个 workflow 的特点：
|
||||
|
||||
- 只依赖 ComfyUI core + 本扩展
|
||||
- 可直接导入并运行，不需要额外图片文件
|
||||
- 可用 `UR Region Editor` 交互绘制 `regions_json`
|
||||
- 也可以手工编辑 `regions_json`
|
||||
- 可预览：
|
||||
- mask preview
|
||||
- filled keyframe
|
||||
- soft composite result
|
||||
|
||||
这个示例是 **remove 核心流程示例**，不内置外部检测模型或修复模型，因此更适合作为：
|
||||
|
||||
- remove 管线起点
|
||||
- 接入你自己的 inpaint / edit / video propagation 模型前的可运行 scaffold
|
||||
|
||||
## 测试
|
||||
|
||||
仓库内已提供基础单元测试：
|
||||
|
||||
- `tests/test_universal_remove.py`
|
||||
|
||||
运行方式：
|
||||
|
||||
```bash
|
||||
uv pip install --python /path/to/ComfyUI/.venv/bin/python -r requirements-dev.txt
|
||||
/path/to/ComfyUI/.venv/bin/python -m pytest tests/test_universal_remove.py
|
||||
```
|
||||
|
||||
## 视频输入与关键帧圈选
|
||||
|
||||
当前扩展不内置视频读取器，但已经内置 `UR Region Editor` 前端交互式圈选 UI。
|
||||
|
||||
推荐接法是：
|
||||
|
||||
```text
|
||||
任意视频读取节点 / 你的脚本
|
||||
-> IMAGE frame batch
|
||||
-> UR Region Editor 输出 regions_json
|
||||
-> UR Regions To Mask Batch
|
||||
```
|
||||
|
||||
也就是说，本扩展支持的是：
|
||||
|
||||
- 输入视频拆帧后的 `IMAGE` 批次
|
||||
- 通过 `frame_index` 指定关键帧
|
||||
- 通过 `UR Region Editor` 鼠标绘制或手写 `regions_json` 描述关键帧上的圈选区域
|
||||
- 输出整段帧批可用的 `MASK`
|
||||
|
||||
`UR Region Editor` 的本地视频背景只用于交互定位，不替代实际视频读取节点；真正进入后端计算的是上游 `IMAGE` frame batch 和编辑器输出的 `regions_json`。
|
||||
|
||||
## region JSON 示例
|
||||
|
||||
### 单关键帧矩形 + 圆形
|
||||
|
||||
```json
|
||||
[
|
||||
{"type": "rect", "frame_index": 0, "x": 160, "y": 120, "w": 260, "h": 140},
|
||||
{"type": "circle", "frame_index": 0, "cx": 520, "cy": 250, "r": 70}
|
||||
]
|
||||
```
|
||||
|
||||
### 多关键帧多边形
|
||||
|
||||
```json
|
||||
[
|
||||
{
|
||||
"type": "polygon",
|
||||
"frame_index": 0,
|
||||
"points": [[120, 160], [240, 150], [280, 260], [150, 290]]
|
||||
},
|
||||
{
|
||||
"type": "polygon",
|
||||
"frame_index": 48,
|
||||
"points": [[180, 170], [310, 160], [350, 280], [210, 310]]
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
配合 `broadcast_mode=hold_until_next_keyframe` 时，第一个关键帧 mask 会使用到下一个关键帧之前，第二个关键帧 mask 会使用到结尾。
|
||||
|
||||
## bbox JSON 示例
|
||||
|
||||
### `xyxy`
|
||||
|
||||
```json
|
||||
[
|
||||
[100, 120, 420, 220],
|
||||
[520, 50, 700, 160]
|
||||
]
|
||||
```
|
||||
|
||||
### `xywh`
|
||||
|
||||
```json
|
||||
[
|
||||
{"x": 100, "y": 120, "w": 320, "h": 100},
|
||||
{"x": 520, "y": 50, "w": 180, "h": 110}
|
||||
]
|
||||
```
|
||||
|
||||
## 备注
|
||||
|
||||
这个扩展当前最适合做：
|
||||
|
||||
- 通用 remove 流程骨架
|
||||
- 去水印 / 去字幕 / 去 logo
|
||||
- 有上游 mask 时的角色移除 / 物体移除
|
||||
|
||||
如果你下一步要做“文本直接选目标并跟踪整段视频”，建议继续在这个目录里加新节点，而不是让核心 remove 节点反向依赖别的 custom nodes。
|
||||
@@ -0,0 +1,8 @@
|
||||
try:
|
||||
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
except ImportError:
|
||||
from nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
WEB_DIRECTORY = "./web"
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|
||||
@@ -0,0 +1,723 @@
|
||||
{
|
||||
"revision": 0,
|
||||
"last_node_id": 9,
|
||||
"last_link_id": 12,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 1,
|
||||
"type": "UR Demo Image",
|
||||
"pos": [
|
||||
-980,
|
||||
260
|
||||
],
|
||||
"size": [
|
||||
400,
|
||||
480
|
||||
],
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "width",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "width"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "height",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "height"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "batch_size",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "batch_size"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
1,
|
||||
4
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "UR Demo Image"
|
||||
},
|
||||
"widgets_values": [
|
||||
768,
|
||||
512,
|
||||
1
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"type": "UR Regions To Mask Batch",
|
||||
"pos": [
|
||||
-460,
|
||||
230
|
||||
],
|
||||
"size": [
|
||||
360,
|
||||
320
|
||||
],
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 1
|
||||
},
|
||||
{
|
||||
"name": "regions_json",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "regions_json"
|
||||
},
|
||||
"link": 12
|
||||
},
|
||||
{
|
||||
"name": "default_frame_index",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "default_frame_index"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "broadcast_mode",
|
||||
"type": "COMBO",
|
||||
"widget": {
|
||||
"name": "broadcast_mode"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "normalized",
|
||||
"type": "BOOLEAN",
|
||||
"widget": {
|
||||
"name": "normalized"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "preview_frame",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "preview_frame"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "preview_color",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "preview_color"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "preview_opacity",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "preview_opacity"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "mask",
|
||||
"type": "MASK",
|
||||
"links": [
|
||||
2
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "selected_frame",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
3
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "preview",
|
||||
"type": "IMAGE",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "UR Regions To Mask Batch"
|
||||
},
|
||||
"widgets_values": [
|
||||
"[{\"type\":\"rect\",\"frame_index\":0,\"x\":160,\"y\":120,\"w\":260,\"h\":140},{\"type\":\"circle\",\"frame_index\":0,\"cx\":520,\"cy\":250,\"r\":70}]",
|
||||
0,
|
||||
"repeat_to_batch",
|
||||
false,
|
||||
0,
|
||||
"#FF3B30",
|
||||
0.35
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 3,
|
||||
"type": "UR Mask Processor",
|
||||
"pos": [
|
||||
-10,
|
||||
230
|
||||
],
|
||||
"size": [
|
||||
330,
|
||||
260
|
||||
],
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mask",
|
||||
"type": "MASK",
|
||||
"link": 2
|
||||
},
|
||||
{
|
||||
"name": "expand_pixels",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "expand_pixels"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "blur_sigma",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "blur_sigma"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "blur_kernel",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "blur_kernel"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "preview_color",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "preview_color"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "preview_opacity",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "preview_opacity"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"link": 3
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "mask",
|
||||
"type": "MASK",
|
||||
"links": [
|
||||
5,
|
||||
7
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "mask_image",
|
||||
"type": "IMAGE",
|
||||
"links": null
|
||||
},
|
||||
{
|
||||
"name": "preview",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
11
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "UR Mask Processor"
|
||||
},
|
||||
"widgets_values": [
|
||||
16,
|
||||
4.0,
|
||||
0,
|
||||
"#00FF00",
|
||||
0.35
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 4,
|
||||
"type": "UR Keyframe Prep",
|
||||
"pos": [
|
||||
400,
|
||||
230
|
||||
],
|
||||
"size": [
|
||||
340,
|
||||
250
|
||||
],
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 4
|
||||
},
|
||||
{
|
||||
"name": "mask",
|
||||
"type": "MASK",
|
||||
"link": 5
|
||||
},
|
||||
{
|
||||
"name": "frame_index",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "frame_index"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "fill_color",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "fill_color"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "resize_longest_side",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "resize_longest_side"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "pad_to_multiple",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "pad_to_multiple"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "original_frame",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
6
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "prepared_frame",
|
||||
"type": "IMAGE",
|
||||
"links": null
|
||||
},
|
||||
{
|
||||
"name": "prepared_mask",
|
||||
"type": "MASK",
|
||||
"links": null
|
||||
},
|
||||
{
|
||||
"name": "filled_frame",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
8,
|
||||
9
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "prepared_width",
|
||||
"type": "INT",
|
||||
"links": null
|
||||
},
|
||||
{
|
||||
"name": "prepared_height",
|
||||
"type": "INT",
|
||||
"links": null
|
||||
},
|
||||
{
|
||||
"name": "frame_count",
|
||||
"type": "INT",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "UR Keyframe Prep"
|
||||
},
|
||||
"widgets_values": [
|
||||
0,
|
||||
"#00FF00",
|
||||
0,
|
||||
16
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 5,
|
||||
"type": "UR Soft Composite",
|
||||
"pos": [
|
||||
840,
|
||||
230
|
||||
],
|
||||
"size": [
|
||||
300,
|
||||
160
|
||||
],
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "inpaint",
|
||||
"type": "IMAGE",
|
||||
"link": 8
|
||||
},
|
||||
{
|
||||
"name": "original",
|
||||
"type": "IMAGE",
|
||||
"link": 6
|
||||
},
|
||||
{
|
||||
"name": "mask",
|
||||
"type": "MASK",
|
||||
"link": 7
|
||||
},
|
||||
{
|
||||
"name": "kernel",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "kernel"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "sigma",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "sigma"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
10
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "soft_mask",
|
||||
"type": "MASK",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "UR Soft Composite"
|
||||
},
|
||||
"widgets_values": [
|
||||
25,
|
||||
10.0
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 6,
|
||||
"type": "PreviewImage",
|
||||
"pos": [
|
||||
300,
|
||||
620
|
||||
],
|
||||
"size": [
|
||||
320,
|
||||
260
|
||||
],
|
||||
"flags": {},
|
||||
"order": 6,
|
||||
"mode": 4,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "images",
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 9
|
||||
}
|
||||
],
|
||||
"outputs": [],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"Node name for S&R": "PreviewImage"
|
||||
},
|
||||
"widgets_values": []
|
||||
},
|
||||
{
|
||||
"id": 7,
|
||||
"type": "PreviewImage",
|
||||
"pos": [
|
||||
700,
|
||||
620
|
||||
],
|
||||
"size": [
|
||||
320,
|
||||
260
|
||||
],
|
||||
"flags": {},
|
||||
"order": 7,
|
||||
"mode": 4,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "images",
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 10
|
||||
}
|
||||
],
|
||||
"outputs": [],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"Node name for S&R": "PreviewImage"
|
||||
},
|
||||
"widgets_values": []
|
||||
},
|
||||
{
|
||||
"id": 8,
|
||||
"type": "PreviewImage",
|
||||
"pos": [
|
||||
-40,
|
||||
620
|
||||
],
|
||||
"size": [
|
||||
320,
|
||||
260
|
||||
],
|
||||
"flags": {},
|
||||
"order": 8,
|
||||
"mode": 4,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "images",
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 11
|
||||
}
|
||||
],
|
||||
"outputs": [],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"Node name for S&R": "PreviewImage"
|
||||
},
|
||||
"widgets_values": []
|
||||
},
|
||||
{
|
||||
"id": 9,
|
||||
"type": "UR Region Editor",
|
||||
"pos": [
|
||||
-900,
|
||||
-250
|
||||
],
|
||||
"size": [
|
||||
680,
|
||||
680
|
||||
],
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "regions_json",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "regions_json"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "canvas_width",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "canvas_width"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "canvas_height",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "canvas_height"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "default_frame_index",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "default_frame_index"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "regions_json",
|
||||
"type": "STRING",
|
||||
"links": [
|
||||
12
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "canvas_width",
|
||||
"type": "INT",
|
||||
"links": null
|
||||
},
|
||||
{
|
||||
"name": "canvas_height",
|
||||
"type": "INT",
|
||||
"links": null
|
||||
},
|
||||
{
|
||||
"name": "region_count",
|
||||
"type": "INT",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "UR Region Editor"
|
||||
},
|
||||
"widgets_values": [
|
||||
"[{\"type\":\"rect\",\"frame_index\":0,\"x\":160,\"y\":120,\"w\":260,\"h\":140},{\"type\":\"circle\",\"frame_index\":0,\"cx\":520,\"cy\":250,\"r\":70}]",
|
||||
768,
|
||||
512,
|
||||
0
|
||||
]
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
1,
|
||||
1,
|
||||
0,
|
||||
2,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
2,
|
||||
2,
|
||||
0,
|
||||
3,
|
||||
0,
|
||||
"MASK"
|
||||
],
|
||||
[
|
||||
3,
|
||||
2,
|
||||
1,
|
||||
3,
|
||||
6,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
4,
|
||||
1,
|
||||
0,
|
||||
4,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
5,
|
||||
3,
|
||||
0,
|
||||
4,
|
||||
1,
|
||||
"MASK"
|
||||
],
|
||||
[
|
||||
6,
|
||||
4,
|
||||
0,
|
||||
5,
|
||||
1,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
7,
|
||||
3,
|
||||
0,
|
||||
5,
|
||||
2,
|
||||
"MASK"
|
||||
],
|
||||
[
|
||||
8,
|
||||
4,
|
||||
3,
|
||||
5,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
9,
|
||||
4,
|
||||
3,
|
||||
6,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
10,
|
||||
5,
|
||||
0,
|
||||
7,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
11,
|
||||
3,
|
||||
2,
|
||||
8,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
12,
|
||||
9,
|
||||
0,
|
||||
2,
|
||||
1,
|
||||
"STRING"
|
||||
]
|
||||
],
|
||||
"version": 0.4
|
||||
}
|
||||
@@ -0,0 +1,843 @@
|
||||
import json
|
||||
import math
|
||||
from typing import List, Sequence, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
CATEGORY = "Universal Remove"
|
||||
|
||||
|
||||
def _ensure_image_batch(images: torch.Tensor) -> torch.Tensor:
|
||||
if not isinstance(images, torch.Tensor):
|
||||
raise TypeError("Expected IMAGE input as torch.Tensor")
|
||||
if images.ndim == 3:
|
||||
images = images.unsqueeze(0)
|
||||
if images.ndim != 4:
|
||||
raise ValueError(f"Expected IMAGE tensor with 4 dims, got {images.ndim}")
|
||||
return images.clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def _ensure_mask_batch(mask: torch.Tensor) -> torch.Tensor:
|
||||
if not isinstance(mask, torch.Tensor):
|
||||
raise TypeError("Expected MASK input as torch.Tensor")
|
||||
if mask.ndim == 2:
|
||||
mask = mask.unsqueeze(0)
|
||||
elif mask.ndim == 4 and mask.shape[1] == 1:
|
||||
mask = mask[:, 0]
|
||||
elif mask.ndim == 4 and mask.shape[-1] == 1:
|
||||
mask = mask[..., 0]
|
||||
if mask.ndim != 3:
|
||||
raise ValueError(f"Expected MASK tensor with 3 dims, got {mask.ndim}")
|
||||
return mask.float().clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def _tile_to_batch(tensor: torch.Tensor, target_batch: int) -> torch.Tensor:
|
||||
if tensor.shape[0] == target_batch:
|
||||
return tensor
|
||||
if tensor.shape[0] <= 0:
|
||||
raise ValueError("Cannot tile an empty batch")
|
||||
repeats = math.ceil(target_batch / tensor.shape[0])
|
||||
repeat_shape = [repeats] + [1] * (tensor.ndim - 1)
|
||||
return tensor.repeat(*repeat_shape)[:target_batch]
|
||||
|
||||
|
||||
def _match_mask_to_images(mask: torch.Tensor, images: torch.Tensor) -> torch.Tensor:
|
||||
mask = _ensure_mask_batch(mask)
|
||||
images = _ensure_image_batch(images)
|
||||
if mask.shape[0] != images.shape[0]:
|
||||
mask = _tile_to_batch(mask, images.shape[0])
|
||||
if mask.shape[1] != images.shape[1] or mask.shape[2] != images.shape[2]:
|
||||
mask = F.interpolate(
|
||||
mask.unsqueeze(1),
|
||||
size=(images.shape[1], images.shape[2]),
|
||||
mode="nearest",
|
||||
).squeeze(1)
|
||||
return mask.clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def _to_bchw(images: torch.Tensor) -> torch.Tensor:
|
||||
return images.permute(0, 3, 1, 2)
|
||||
|
||||
|
||||
def _to_bhwc(images: torch.Tensor) -> torch.Tensor:
|
||||
return images.permute(0, 2, 3, 1)
|
||||
|
||||
|
||||
def _resize_images(images: torch.Tensor, size: Tuple[int, int]) -> torch.Tensor:
|
||||
bchw = _to_bchw(images)
|
||||
resized = F.interpolate(bchw, size=size, mode="bilinear", align_corners=False)
|
||||
return _to_bhwc(resized).clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def _resize_masks(mask: torch.Tensor, size: Tuple[int, int]) -> torch.Tensor:
|
||||
resized = F.interpolate(mask.unsqueeze(1), size=size, mode="nearest")
|
||||
return resized.squeeze(1).clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def _pad_images(images: torch.Tensor, target_h: int, target_w: int) -> torch.Tensor:
|
||||
pad_h = max(0, target_h - images.shape[1])
|
||||
pad_w = max(0, target_w - images.shape[2])
|
||||
if pad_h == 0 and pad_w == 0:
|
||||
return images
|
||||
bchw = _to_bchw(images)
|
||||
padded = F.pad(bchw, (0, pad_w, 0, pad_h), mode="constant", value=0.0)
|
||||
return _to_bhwc(padded).clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def _pad_masks(mask: torch.Tensor, target_h: int, target_w: int) -> torch.Tensor:
|
||||
pad_h = max(0, target_h - mask.shape[1])
|
||||
pad_w = max(0, target_w - mask.shape[2])
|
||||
if pad_h == 0 and pad_w == 0:
|
||||
return mask
|
||||
padded = F.pad(mask.unsqueeze(1), (0, pad_w, 0, pad_h), mode="constant", value=0.0)
|
||||
return padded.squeeze(1).clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def _parse_hex_color(color: str) -> torch.Tensor:
|
||||
text = color.strip()
|
||||
if text.startswith("#"):
|
||||
text = text[1:]
|
||||
if len(text) != 6:
|
||||
raise ValueError("Color must be a 6-digit hex string, e.g. #00FF00")
|
||||
rgb = [int(text[i : i + 2], 16) / 255.0 for i in range(0, 6, 2)]
|
||||
return torch.tensor(rgb, dtype=torch.float32)
|
||||
|
||||
|
||||
def _overlay_mask(images: torch.Tensor, mask: torch.Tensor, color: torch.Tensor, opacity: float) -> torch.Tensor:
|
||||
images = _ensure_image_batch(images)
|
||||
mask = _match_mask_to_images(mask, images)
|
||||
color = color.to(device=images.device, dtype=images.dtype).view(1, 1, 1, 3)
|
||||
alpha = mask.unsqueeze(-1).clamp(0.0, 1.0) * float(opacity)
|
||||
return (images * (1.0 - alpha) + color * alpha).clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def _gaussian_kernel_1d(kernel_size: int, sigma: float, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
|
||||
if kernel_size % 2 == 0:
|
||||
kernel_size += 1
|
||||
radius = kernel_size // 2
|
||||
if sigma <= 0:
|
||||
sigma = max(float(kernel_size) / 6.0, 1e-6)
|
||||
coords = torch.arange(-radius, radius + 1, device=device, dtype=dtype)
|
||||
kernel = torch.exp(-0.5 * (coords / sigma) ** 2)
|
||||
kernel = kernel / kernel.sum()
|
||||
return kernel
|
||||
|
||||
|
||||
def _gaussian_blur_mask(mask: torch.Tensor, kernel_size: int, sigma: float) -> torch.Tensor:
|
||||
mask = _ensure_mask_batch(mask)
|
||||
if kernel_size <= 1 and sigma <= 0:
|
||||
return mask
|
||||
if kernel_size <= 1:
|
||||
kernel_size = int(max(3, math.ceil(sigma * 6)))
|
||||
if kernel_size % 2 == 0:
|
||||
kernel_size += 1
|
||||
|
||||
kernel_1d = _gaussian_kernel_1d(kernel_size, sigma, mask.device, mask.dtype)
|
||||
kernel_x = kernel_1d.view(1, 1, 1, -1)
|
||||
kernel_y = kernel_1d.view(1, 1, -1, 1)
|
||||
|
||||
blurred = mask.unsqueeze(1)
|
||||
pad = kernel_size // 2
|
||||
blurred = F.pad(blurred, (pad, pad, 0, 0), mode="replicate")
|
||||
blurred = F.conv2d(blurred, kernel_x)
|
||||
blurred = F.pad(blurred, (0, 0, pad, pad), mode="replicate")
|
||||
blurred = F.conv2d(blurred, kernel_y)
|
||||
return blurred.squeeze(1).clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def _expand_mask(mask: torch.Tensor, expand_pixels: int) -> torch.Tensor:
|
||||
mask = _ensure_mask_batch(mask)
|
||||
if expand_pixels <= 0:
|
||||
return mask
|
||||
kernel = expand_pixels * 2 + 1
|
||||
expanded = F.max_pool2d(mask.unsqueeze(1), kernel_size=kernel, stride=1, padding=expand_pixels)
|
||||
return expanded.squeeze(1).clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def _fill_masked_region(images: torch.Tensor, mask: torch.Tensor, color: torch.Tensor) -> torch.Tensor:
|
||||
images = _ensure_image_batch(images)
|
||||
mask = _match_mask_to_images(mask, images)
|
||||
color = color.to(device=images.device, dtype=images.dtype).view(1, 1, 1, 3)
|
||||
return (images * (1.0 - mask.unsqueeze(-1)) + color * mask.unsqueeze(-1)).clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def _normalize_box_values(values: Sequence[float], width: int, height: int, normalized: bool) -> Tuple[float, float, float, float]:
|
||||
x1, y1, x2, y2 = values
|
||||
if normalized:
|
||||
x1 *= width
|
||||
x2 *= width
|
||||
y1 *= height
|
||||
y2 *= height
|
||||
return x1, y1, x2, y2
|
||||
|
||||
|
||||
def _parse_single_box(item, coordinate_format: str, width: int, height: int, normalized: bool) -> Tuple[int, int, int, int]:
|
||||
if isinstance(item, dict):
|
||||
if {"x", "y", "w", "h"}.issubset(item):
|
||||
x1 = float(item["x"])
|
||||
y1 = float(item["y"])
|
||||
w = float(item["w"])
|
||||
h = float(item["h"])
|
||||
if normalized:
|
||||
x1 *= width
|
||||
y1 *= height
|
||||
w *= width
|
||||
h *= height
|
||||
x2 = x1 + w
|
||||
y2 = y1 + h
|
||||
elif {"x1", "y1", "x2", "y2"}.issubset(item):
|
||||
x1, y1, x2, y2 = _normalize_box_values(
|
||||
[float(item["x1"]), float(item["y1"]), float(item["x2"]), float(item["y2"])],
|
||||
width,
|
||||
height,
|
||||
normalized,
|
||||
)
|
||||
elif {"left", "top", "right", "bottom"}.issubset(item):
|
||||
x1, y1, x2, y2 = _normalize_box_values(
|
||||
[float(item["left"]), float(item["top"]), float(item["right"]), float(item["bottom"])],
|
||||
width,
|
||||
height,
|
||||
normalized,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported box object: {item}")
|
||||
elif isinstance(item, (list, tuple)) and len(item) == 4:
|
||||
values = [float(v) for v in item]
|
||||
if coordinate_format == "xywh":
|
||||
x1, y1, w, h = values
|
||||
if normalized:
|
||||
x1 *= width
|
||||
y1 *= height
|
||||
w *= width
|
||||
h *= height
|
||||
x2 = x1 + w
|
||||
y2 = y1 + h
|
||||
else:
|
||||
x1, y1, x2, y2 = _normalize_box_values(values, width, height, normalized)
|
||||
else:
|
||||
raise ValueError(f"Unsupported box entry: {item}")
|
||||
|
||||
left = int(math.floor(min(x1, x2)))
|
||||
right = int(math.ceil(max(x1, x2)))
|
||||
top = int(math.floor(min(y1, y2)))
|
||||
bottom = int(math.ceil(max(y1, y2)))
|
||||
|
||||
left = max(0, min(width, left))
|
||||
right = max(0, min(width, right))
|
||||
top = max(0, min(height, top))
|
||||
bottom = max(0, min(height, bottom))
|
||||
return left, top, right, bottom
|
||||
|
||||
|
||||
def _parse_boxes_json(boxes_json: str, coordinate_format: str, width: int, height: int, normalized: bool) -> List[Tuple[int, int, int, int]]:
|
||||
text = boxes_json.strip()
|
||||
if not text:
|
||||
return []
|
||||
data = json.loads(text)
|
||||
if isinstance(data, dict):
|
||||
if "boxes" in data:
|
||||
data = data["boxes"]
|
||||
else:
|
||||
data = [data]
|
||||
if not isinstance(data, list):
|
||||
raise ValueError("boxes_json must be a JSON array or object")
|
||||
boxes = []
|
||||
for item in data:
|
||||
box = _parse_single_box(item, coordinate_format, width, height, normalized)
|
||||
if box[2] > box[0] and box[3] > box[1]:
|
||||
boxes.append(box)
|
||||
return boxes
|
||||
|
||||
|
||||
def _scale_point(x: float, y: float, width: int, height: int, normalized: bool) -> Tuple[float, float]:
|
||||
if normalized:
|
||||
x *= width
|
||||
y *= height
|
||||
return x, y
|
||||
|
||||
|
||||
def _region_frame_index(region: dict, default_frame_index: int, batch_size: int) -> int:
|
||||
index = int(region.get("frame_index", default_frame_index))
|
||||
return max(0, min(index, batch_size - 1))
|
||||
|
||||
|
||||
def _region_type(region: dict) -> str:
|
||||
value = str(region.get("type", region.get("shape", "rect"))).strip().lower()
|
||||
aliases = {
|
||||
"rectangle": "rect",
|
||||
"box": "rect",
|
||||
"bbox": "rect",
|
||||
"ellipse": "circle",
|
||||
"poly": "polygon",
|
||||
}
|
||||
return aliases.get(value, value)
|
||||
|
||||
|
||||
def _polygon_to_mask(points: Sequence[Sequence[float]], width: int, height: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
|
||||
if len(points) < 3:
|
||||
return torch.zeros((height, width), dtype=dtype, device=device)
|
||||
|
||||
px = torch.arange(width, dtype=dtype, device=device).view(1, width) + 0.5
|
||||
py = torch.arange(height, dtype=dtype, device=device).view(height, 1) + 0.5
|
||||
inside = torch.zeros((height, width), dtype=torch.bool, device=device)
|
||||
|
||||
wrapped = list(points) + [points[0]]
|
||||
eps = torch.tensor(1e-6, dtype=dtype, device=device)
|
||||
for current, nxt in zip(wrapped[:-1], wrapped[1:]):
|
||||
x1 = torch.tensor(float(current[0]), dtype=dtype, device=device)
|
||||
y1 = torch.tensor(float(current[1]), dtype=dtype, device=device)
|
||||
x2 = torch.tensor(float(nxt[0]), dtype=dtype, device=device)
|
||||
y2 = torch.tensor(float(nxt[1]), dtype=dtype, device=device)
|
||||
crosses_y = torch.logical_xor(y1 > py, y2 > py)
|
||||
x_intersection = (x2 - x1) * (py - y1) / (y2 - y1 + eps) + x1
|
||||
intersects = crosses_y & (px < x_intersection)
|
||||
inside = torch.logical_xor(inside, intersects)
|
||||
|
||||
return inside.to(dtype=dtype)
|
||||
|
||||
|
||||
def _draw_region(mask: torch.Tensor, region: dict, width: int, height: int, normalized: bool) -> None:
|
||||
shape = _region_type(region)
|
||||
if shape == "rect":
|
||||
box = _parse_single_box(region, "xyxy", width, height, normalized)
|
||||
x1, y1, x2, y2 = box
|
||||
if x2 > x1 and y2 > y1:
|
||||
mask[y1:y2, x1:x2] = 1.0
|
||||
return
|
||||
|
||||
if shape == "circle":
|
||||
if "center" in region:
|
||||
cx, cy = region["center"]
|
||||
else:
|
||||
cx = region.get("cx", region.get("x", 0.0))
|
||||
cy = region.get("cy", region.get("y", 0.0))
|
||||
cx, cy = _scale_point(float(cx), float(cy), width, height, normalized)
|
||||
radius = float(region.get("r", region.get("radius", 0.0)))
|
||||
if normalized:
|
||||
radius *= min(width, height)
|
||||
if radius <= 0:
|
||||
return
|
||||
yy, xx = torch.meshgrid(
|
||||
torch.arange(height, dtype=mask.dtype, device=mask.device),
|
||||
torch.arange(width, dtype=mask.dtype, device=mask.device),
|
||||
indexing="ij",
|
||||
)
|
||||
circle = (xx + 0.5 - cx) ** 2 + (yy + 0.5 - cy) ** 2 <= radius ** 2
|
||||
mask[circle] = 1.0
|
||||
return
|
||||
|
||||
if shape == "polygon":
|
||||
raw_points = region.get("points", [])
|
||||
points = []
|
||||
for point in raw_points:
|
||||
if isinstance(point, dict):
|
||||
x = float(point.get("x", 0.0))
|
||||
y = float(point.get("y", 0.0))
|
||||
else:
|
||||
x = float(point[0])
|
||||
y = float(point[1])
|
||||
points.append(_scale_point(x, y, width, height, normalized))
|
||||
mask.copy_(torch.maximum(mask, _polygon_to_mask(points, width, height, mask.device, mask.dtype)))
|
||||
return
|
||||
|
||||
raise ValueError(f"Unsupported region type: {shape}")
|
||||
|
||||
|
||||
def _parse_regions_json(regions_json: str) -> List[dict]:
|
||||
text = regions_json.strip()
|
||||
if not text:
|
||||
return []
|
||||
data = json.loads(text)
|
||||
if isinstance(data, dict):
|
||||
if "regions" in data:
|
||||
data = data["regions"]
|
||||
else:
|
||||
data = [data]
|
||||
if not isinstance(data, list):
|
||||
raise ValueError("regions_json must be a JSON array or object")
|
||||
for region in data:
|
||||
if not isinstance(region, dict):
|
||||
raise ValueError(f"Region entries must be objects, got: {region}")
|
||||
return data
|
||||
|
||||
|
||||
def _apply_keyframe_masks(keyframe_masks: dict, batch: int, height: int, width: int, dtype: torch.dtype, device: torch.device, broadcast_mode: str) -> torch.Tensor:
|
||||
mask = torch.zeros((batch, height, width), dtype=dtype, device=device)
|
||||
if not keyframe_masks:
|
||||
return mask
|
||||
|
||||
sorted_items = sorted(keyframe_masks.items())
|
||||
if broadcast_mode == "repeat_to_batch":
|
||||
combined = torch.zeros((height, width), dtype=dtype, device=device)
|
||||
for _, key_mask in sorted_items:
|
||||
combined = torch.maximum(combined, key_mask)
|
||||
return combined.unsqueeze(0).repeat(batch, 1, 1)
|
||||
|
||||
if broadcast_mode == "hold_until_next_keyframe":
|
||||
for idx, (start_frame, key_mask) in enumerate(sorted_items):
|
||||
end_frame = sorted_items[idx + 1][0] if idx + 1 < len(sorted_items) else batch
|
||||
if end_frame > start_frame:
|
||||
mask[start_frame:end_frame] = key_mask
|
||||
return mask
|
||||
|
||||
for frame, key_mask in sorted_items:
|
||||
if broadcast_mode == "hold_after_frame":
|
||||
mask[frame:] = key_mask
|
||||
else:
|
||||
mask[frame] = key_mask
|
||||
return mask
|
||||
|
||||
|
||||
class URDemoImage:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"width": ("INT", {"default": 768, "min": 64, "max": 8192, "step": 8}),
|
||||
"height": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 8}),
|
||||
"batch_size": ("INT", {"default": 1, "min": 1, "max": 256}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = CATEGORY
|
||||
|
||||
def generate(self, width: int, height: int, batch_size: int):
|
||||
x = torch.linspace(0.0, 1.0, int(width), dtype=torch.float32).view(1, 1, width)
|
||||
y = torch.linspace(0.0, 1.0, int(height), dtype=torch.float32).view(1, height, 1)
|
||||
red = 0.15 + 0.35 * x.expand(1, height, width)
|
||||
green = 0.18 + 0.45 * y.expand(1, height, width)
|
||||
blue = 0.25 + 0.25 * (1.0 - x.expand(1, height, width))
|
||||
image = torch.stack([red, green, blue], dim=-1).repeat(int(batch_size), 1, 1, 1)
|
||||
|
||||
x1 = max(0, min(int(width), int(width * 0.22)))
|
||||
x2 = max(0, min(int(width), int(width * 0.55)))
|
||||
y1 = max(0, min(int(height), int(height * 0.24)))
|
||||
y2 = max(0, min(int(height), int(height * 0.50)))
|
||||
image[:, y1:y2, x1:x2, 0] = 0.92
|
||||
image[:, y1:y2, x1:x2, 1] = 0.16
|
||||
image[:, y1:y2, x1:x2, 2] = 0.12
|
||||
return (image.clamp(0.0, 1.0),)
|
||||
|
||||
|
||||
class URRemovePromptBuilder:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"target_prompt": ("STRING", {"multiline": False, "default": "object to remove"}),
|
||||
"scene_hint": ("STRING", {"multiline": True, "default": ""}),
|
||||
"preserve_prompt": ("STRING", {"multiline": True, "default": ""}),
|
||||
"extra_negative_prompt": ("STRING", {"multiline": True, "default": ""}),
|
||||
"fill_mode": (["background_reconstruct", "clean_plate", "direct_replace"],),
|
||||
"masked_edit_only": ("BOOLEAN", {"default": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "STRING", "STRING")
|
||||
RETURN_NAMES = ("segmentation_prompt", "remove_positive_prompt", "remove_negative_prompt")
|
||||
FUNCTION = "build"
|
||||
CATEGORY = CATEGORY
|
||||
|
||||
def build(
|
||||
self,
|
||||
target_prompt: str,
|
||||
scene_hint: str,
|
||||
preserve_prompt: str,
|
||||
extra_negative_prompt: str,
|
||||
fill_mode: str,
|
||||
masked_edit_only: bool,
|
||||
):
|
||||
target = target_prompt.strip() or "target object"
|
||||
segmentation_prompt = target
|
||||
|
||||
positive_parts = [f"Remove {target} from the masked region."]
|
||||
if fill_mode == "background_reconstruct":
|
||||
positive_parts.append("Reconstruct a plausible background behind the removed content.")
|
||||
elif fill_mode == "clean_plate":
|
||||
positive_parts.append("Create a clean plate result that matches the surrounding scene.")
|
||||
else:
|
||||
positive_parts.append("Replace the masked region with visually coherent scene content.")
|
||||
|
||||
if masked_edit_only:
|
||||
positive_parts.append("Only modify the masked region and keep unmasked content unchanged.")
|
||||
if scene_hint.strip():
|
||||
positive_parts.append(f"Scene context: {scene_hint.strip()}.")
|
||||
if preserve_prompt.strip():
|
||||
positive_parts.append(f"Preserve: {preserve_prompt.strip()}.")
|
||||
|
||||
negative_parts = [
|
||||
"artifacts",
|
||||
"ghosting",
|
||||
"duplicate subject",
|
||||
"warped anatomy",
|
||||
"broken perspective",
|
||||
"extra text",
|
||||
"watermark",
|
||||
"subtitle",
|
||||
"logo",
|
||||
"blur",
|
||||
"low quality",
|
||||
]
|
||||
if extra_negative_prompt.strip():
|
||||
negative_parts.append(extra_negative_prompt.strip())
|
||||
|
||||
return (
|
||||
segmentation_prompt,
|
||||
" ".join(positive_parts),
|
||||
", ".join(negative_parts),
|
||||
)
|
||||
|
||||
|
||||
class URRegionEditor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"regions_json": (
|
||||
"STRING",
|
||||
{
|
||||
"multiline": True,
|
||||
"default": "[]",
|
||||
},
|
||||
),
|
||||
"canvas_width": ("INT", {"default": 768, "min": 64, "max": 8192, "step": 8}),
|
||||
"canvas_height": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 8}),
|
||||
"default_frame_index": ("INT", {"default": 0, "min": 0, "max": 1000000}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "INT", "INT", "INT")
|
||||
RETURN_NAMES = ("regions_json", "canvas_width", "canvas_height", "region_count")
|
||||
FUNCTION = "emit"
|
||||
CATEGORY = CATEGORY
|
||||
|
||||
def emit(self, regions_json: str, canvas_width: int, canvas_height: int, default_frame_index: int):
|
||||
regions = _parse_regions_json(regions_json)
|
||||
normalized_regions = []
|
||||
for region in regions:
|
||||
copied = dict(region)
|
||||
copied["frame_index"] = int(copied.get("frame_index", default_frame_index))
|
||||
normalized_regions.append(copied)
|
||||
text = json.dumps(normalized_regions, ensure_ascii=False)
|
||||
return (text, int(canvas_width), int(canvas_height), len(normalized_regions))
|
||||
|
||||
|
||||
class URBoxesToMaskBatch:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"boxes_json": ("STRING", {"multiline": True, "default": "[]"}),
|
||||
"frame_index": ("INT", {"default": 0, "min": 0, "max": 1000000}),
|
||||
"coordinate_format": (["xyxy", "xywh"],),
|
||||
"broadcast_mode": (["repeat_to_batch", "single_frame_only", "hold_after_frame"],),
|
||||
"normalized": ("BOOLEAN", {"default": False}),
|
||||
"preview_color": ("STRING", {"multiline": False, "default": "#FF3B30"}),
|
||||
"preview_opacity": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK", "IMAGE", "IMAGE")
|
||||
RETURN_NAMES = ("mask", "selected_frame", "preview")
|
||||
FUNCTION = "build_masks"
|
||||
CATEGORY = CATEGORY
|
||||
|
||||
def build_masks(
|
||||
self,
|
||||
images: torch.Tensor,
|
||||
boxes_json: str,
|
||||
frame_index: int,
|
||||
coordinate_format: str,
|
||||
broadcast_mode: str,
|
||||
normalized: bool,
|
||||
preview_color: str,
|
||||
preview_opacity: float,
|
||||
):
|
||||
images = _ensure_image_batch(images)
|
||||
batch, height, width, _ = images.shape
|
||||
index = max(0, min(int(frame_index), batch - 1))
|
||||
|
||||
boxes = _parse_boxes_json(boxes_json, coordinate_format, width, height, normalized)
|
||||
selected_mask = torch.zeros((1, height, width), dtype=images.dtype, device=images.device)
|
||||
for x1, y1, x2, y2 in boxes:
|
||||
selected_mask[:, y1:y2, x1:x2] = 1.0
|
||||
|
||||
if broadcast_mode == "repeat_to_batch":
|
||||
mask = selected_mask.repeat(batch, 1, 1)
|
||||
elif broadcast_mode == "hold_after_frame":
|
||||
mask = torch.zeros((batch, height, width), dtype=images.dtype, device=images.device)
|
||||
mask[index:] = selected_mask
|
||||
else:
|
||||
mask = torch.zeros((batch, height, width), dtype=images.dtype, device=images.device)
|
||||
mask[index] = selected_mask[0]
|
||||
|
||||
preview = _overlay_mask(
|
||||
images[index : index + 1],
|
||||
selected_mask,
|
||||
_parse_hex_color(preview_color),
|
||||
preview_opacity,
|
||||
)
|
||||
return (mask, images[index : index + 1], preview)
|
||||
|
||||
|
||||
class URRegionsToMaskBatch:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"regions_json": (
|
||||
"STRING",
|
||||
{
|
||||
"multiline": True,
|
||||
"default": '[{"type":"rect","frame_index":0,"x":160,"y":120,"w":260,"h":140}]',
|
||||
},
|
||||
),
|
||||
"default_frame_index": ("INT", {"default": 0, "min": 0, "max": 1000000}),
|
||||
"broadcast_mode": (["repeat_to_batch", "single_frame_only", "hold_after_frame", "hold_until_next_keyframe"],),
|
||||
"normalized": ("BOOLEAN", {"default": False}),
|
||||
"preview_frame": ("INT", {"default": 0, "min": 0, "max": 1000000}),
|
||||
"preview_color": ("STRING", {"multiline": False, "default": "#FF3B30"}),
|
||||
"preview_opacity": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK", "IMAGE", "IMAGE")
|
||||
RETURN_NAMES = ("mask", "selected_frame", "preview")
|
||||
FUNCTION = "build_masks"
|
||||
CATEGORY = CATEGORY
|
||||
|
||||
def build_masks(
|
||||
self,
|
||||
images: torch.Tensor,
|
||||
regions_json: str,
|
||||
default_frame_index: int,
|
||||
broadcast_mode: str,
|
||||
normalized: bool,
|
||||
preview_frame: int,
|
||||
preview_color: str,
|
||||
preview_opacity: float,
|
||||
):
|
||||
images = _ensure_image_batch(images)
|
||||
batch, height, width, _ = images.shape
|
||||
default_index = max(0, min(int(default_frame_index), batch - 1))
|
||||
preview_index = max(0, min(int(preview_frame), batch - 1))
|
||||
|
||||
keyframe_masks = {}
|
||||
for region in _parse_regions_json(regions_json):
|
||||
frame = _region_frame_index(region, default_index, batch)
|
||||
if frame not in keyframe_masks:
|
||||
keyframe_masks[frame] = torch.zeros((height, width), dtype=images.dtype, device=images.device)
|
||||
_draw_region(keyframe_masks[frame], region, width, height, normalized)
|
||||
|
||||
mask = _apply_keyframe_masks(
|
||||
keyframe_masks,
|
||||
batch,
|
||||
height,
|
||||
width,
|
||||
images.dtype,
|
||||
images.device,
|
||||
broadcast_mode,
|
||||
)
|
||||
preview = _overlay_mask(
|
||||
images[preview_index : preview_index + 1],
|
||||
mask[preview_index : preview_index + 1],
|
||||
_parse_hex_color(preview_color),
|
||||
preview_opacity,
|
||||
)
|
||||
return (mask, images[preview_index : preview_index + 1], preview)
|
||||
|
||||
|
||||
class URMaskProcessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"mask": ("MASK",),
|
||||
"expand_pixels": ("INT", {"default": 0, "min": 0, "max": 2048}),
|
||||
"blur_sigma": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 512.0, "step": 0.01}),
|
||||
"blur_kernel": ("INT", {"default": 0, "min": 0, "max": 2049}),
|
||||
"preview_color": ("STRING", {"multiline": False, "default": "#00FF00"}),
|
||||
"preview_opacity": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
},
|
||||
"optional": {
|
||||
"image": ("IMAGE",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK", "IMAGE", "IMAGE")
|
||||
RETURN_NAMES = ("mask", "mask_image", "preview")
|
||||
FUNCTION = "process"
|
||||
CATEGORY = CATEGORY
|
||||
|
||||
def process(
|
||||
self,
|
||||
mask: torch.Tensor,
|
||||
expand_pixels: int,
|
||||
blur_sigma: float,
|
||||
blur_kernel: int,
|
||||
preview_color: str,
|
||||
preview_opacity: float,
|
||||
image: torch.Tensor = None,
|
||||
):
|
||||
mask = _ensure_mask_batch(mask)
|
||||
if expand_pixels > 0:
|
||||
mask = _expand_mask(mask, expand_pixels)
|
||||
if blur_sigma > 0.0 or blur_kernel > 1:
|
||||
mask = _gaussian_blur_mask(mask, blur_kernel, blur_sigma)
|
||||
mask = mask.clamp(0.0, 1.0)
|
||||
|
||||
mask_image = mask.unsqueeze(-1).repeat(1, 1, 1, 3)
|
||||
if image is not None:
|
||||
image = _ensure_image_batch(image)
|
||||
preview = _overlay_mask(image, mask, _parse_hex_color(preview_color), preview_opacity)
|
||||
else:
|
||||
preview = _overlay_mask(mask_image, mask, _parse_hex_color(preview_color), preview_opacity)
|
||||
return (mask, mask_image, preview)
|
||||
|
||||
|
||||
class URKeyframePrep:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"mask": ("MASK",),
|
||||
"frame_index": ("INT", {"default": 0, "min": 0, "max": 1000000}),
|
||||
"fill_color": ("STRING", {"multiline": False, "default": "#00FF00"}),
|
||||
"resize_longest_side": ("INT", {"default": 0, "min": 0, "max": 8192, "step": 8}),
|
||||
"pad_to_multiple": ("INT", {"default": 16, "min": 1, "max": 256}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "IMAGE", "MASK", "IMAGE", "INT", "INT", "INT")
|
||||
RETURN_NAMES = (
|
||||
"original_frame",
|
||||
"prepared_frame",
|
||||
"prepared_mask",
|
||||
"filled_frame",
|
||||
"prepared_width",
|
||||
"prepared_height",
|
||||
"frame_count",
|
||||
)
|
||||
FUNCTION = "prepare"
|
||||
CATEGORY = CATEGORY
|
||||
|
||||
def prepare(
|
||||
self,
|
||||
images: torch.Tensor,
|
||||
mask: torch.Tensor,
|
||||
frame_index: int,
|
||||
fill_color: str,
|
||||
resize_longest_side: int,
|
||||
pad_to_multiple: int,
|
||||
):
|
||||
images = _ensure_image_batch(images)
|
||||
mask = _match_mask_to_images(mask, images)
|
||||
batch, height, width, _ = images.shape
|
||||
index = max(0, min(int(frame_index), batch - 1))
|
||||
|
||||
original_frame = images[index : index + 1]
|
||||
prepared_frame = original_frame.clone()
|
||||
prepared_mask = mask[index : index + 1].clone()
|
||||
|
||||
if resize_longest_side > 0:
|
||||
longest = max(height, width)
|
||||
scale = float(resize_longest_side) / float(longest)
|
||||
target_h = max(1, int(round(height * scale)))
|
||||
target_w = max(1, int(round(width * scale)))
|
||||
prepared_frame = _resize_images(prepared_frame, (target_h, target_w))
|
||||
prepared_mask = _resize_masks(prepared_mask, (target_h, target_w))
|
||||
|
||||
prepared_h = prepared_frame.shape[1]
|
||||
prepared_w = prepared_frame.shape[2]
|
||||
if pad_to_multiple > 1:
|
||||
target_h = int(math.ceil(prepared_h / pad_to_multiple) * pad_to_multiple)
|
||||
target_w = int(math.ceil(prepared_w / pad_to_multiple) * pad_to_multiple)
|
||||
prepared_frame = _pad_images(prepared_frame, target_h, target_w)
|
||||
prepared_mask = _pad_masks(prepared_mask, target_h, target_w)
|
||||
|
||||
color = _parse_hex_color(fill_color)
|
||||
filled_frame = _fill_masked_region(prepared_frame, prepared_mask, color)
|
||||
return (
|
||||
original_frame,
|
||||
prepared_frame,
|
||||
prepared_mask,
|
||||
filled_frame,
|
||||
prepared_frame.shape[2],
|
||||
prepared_frame.shape[1],
|
||||
batch,
|
||||
)
|
||||
|
||||
|
||||
class URSoftComposite:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"inpaint": ("IMAGE",),
|
||||
"original": ("IMAGE",),
|
||||
"mask": ("MASK",),
|
||||
"kernel": ("INT", {"default": 15, "min": 1, "max": 2049}),
|
||||
"sigma": ("FLOAT", {"default": 7.5, "min": 0.0, "max": 512.0, "step": 0.01}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK")
|
||||
RETURN_NAMES = ("image", "soft_mask")
|
||||
FUNCTION = "composite"
|
||||
CATEGORY = CATEGORY
|
||||
|
||||
def composite(
|
||||
self,
|
||||
inpaint: torch.Tensor,
|
||||
original: torch.Tensor,
|
||||
mask: torch.Tensor,
|
||||
kernel: int,
|
||||
sigma: float,
|
||||
):
|
||||
original = _ensure_image_batch(original)
|
||||
inpaint = _ensure_image_batch(inpaint)
|
||||
target_batch = max(original.shape[0], inpaint.shape[0], _ensure_mask_batch(mask).shape[0])
|
||||
|
||||
original = _tile_to_batch(original, target_batch)
|
||||
inpaint = _tile_to_batch(inpaint, target_batch)
|
||||
mask = _tile_to_batch(_ensure_mask_batch(mask), target_batch)
|
||||
mask = _match_mask_to_images(mask, original)
|
||||
|
||||
if inpaint.shape[1] != original.shape[1] or inpaint.shape[2] != original.shape[2]:
|
||||
inpaint = _resize_images(inpaint, (original.shape[1], original.shape[2]))
|
||||
|
||||
soft_mask = _gaussian_blur_mask(mask, kernel, sigma).clamp(0.0, 1.0)
|
||||
composited = original * (1.0 - soft_mask.unsqueeze(-1)) + inpaint * soft_mask.unsqueeze(-1)
|
||||
return (composited.clamp(0.0, 1.0), soft_mask)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"UR Demo Image": URDemoImage,
|
||||
"UR Remove Prompt Builder": URRemovePromptBuilder,
|
||||
"UR Region Editor": URRegionEditor,
|
||||
"UR Boxes To Mask Batch": URBoxesToMaskBatch,
|
||||
"UR Regions To Mask Batch": URRegionsToMaskBatch,
|
||||
"UR Mask Processor": URMaskProcessor,
|
||||
"UR Keyframe Prep": URKeyframePrep,
|
||||
"UR Soft Composite": URSoftComposite,
|
||||
}
|
||||
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"UR Demo Image": "UR Demo Image",
|
||||
"UR Remove Prompt Builder": "UR Remove Prompt Builder",
|
||||
"UR Region Editor": "UR Region Editor",
|
||||
"UR Boxes To Mask Batch": "UR Boxes To Mask Batch",
|
||||
"UR Regions To Mask Batch": "UR Regions To Mask Batch",
|
||||
"UR Mask Processor": "UR Mask Processor",
|
||||
"UR Keyframe Prep": "UR Keyframe Prep",
|
||||
"UR Soft Composite": "UR Soft Composite",
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
[pytest]
|
||||
testpaths = tests
|
||||
@@ -0,0 +1 @@
|
||||
pytest>=7.8.0
|
||||
@@ -0,0 +1,244 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
MODULE_PATH = Path(__file__).resolve().parents[1] / "nodes.py"
|
||||
SPEC = importlib.util.spec_from_file_location("universal_remove_nodes", MODULE_PATH)
|
||||
MODULE = importlib.util.module_from_spec(SPEC)
|
||||
assert SPEC is not None and SPEC.loader is not None
|
||||
SPEC.loader.exec_module(MODULE)
|
||||
|
||||
|
||||
def test_extension_entrypoint_exports_nodes_and_web_directory():
|
||||
package_path = MODULE_PATH.parent / "__init__.py"
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"ComfyUI_UniversalRemove_test",
|
||||
package_path,
|
||||
submodule_search_locations=[str(MODULE_PATH.parent)],
|
||||
)
|
||||
package = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = package
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(package)
|
||||
|
||||
assert package.WEB_DIRECTORY == "./web"
|
||||
assert "UR Region Editor" in package.NODE_CLASS_MAPPINGS
|
||||
assert "UR Regions To Mask Batch" in package.NODE_CLASS_MAPPINGS
|
||||
|
||||
|
||||
def test_prompt_builder_outputs_expected_strings():
|
||||
node = MODULE.URRemovePromptBuilder()
|
||||
segmentation_prompt, positive_prompt, negative_prompt = node.build(
|
||||
target_prompt="watermark",
|
||||
scene_hint="night city street",
|
||||
preserve_prompt="keep the lighting and camera framing",
|
||||
extra_negative_prompt="text remnants",
|
||||
fill_mode="background_reconstruct",
|
||||
masked_edit_only=True,
|
||||
)
|
||||
|
||||
assert segmentation_prompt == "watermark"
|
||||
assert "Remove watermark" in positive_prompt
|
||||
assert "night city street" in positive_prompt
|
||||
assert "keep the lighting and camera framing" in positive_prompt
|
||||
assert "text remnants" in negative_prompt
|
||||
|
||||
|
||||
def test_demo_image_is_self_contained_and_batched():
|
||||
node = MODULE.URDemoImage()
|
||||
(image,) = node.generate(width=64, height=32, batch_size=2)
|
||||
|
||||
assert image.shape == (2, 32, 64, 3)
|
||||
assert image.min().item() >= 0.0
|
||||
assert image.max().item() <= 1.0
|
||||
|
||||
|
||||
def test_region_editor_normalizes_regions_and_counts_them():
|
||||
node = MODULE.URRegionEditor()
|
||||
regions_json, canvas_width, canvas_height, region_count = node.emit(
|
||||
regions_json='[{"type":"rect","x":1,"y":2,"w":3,"h":4}]',
|
||||
canvas_width=64,
|
||||
canvas_height=32,
|
||||
default_frame_index=5,
|
||||
)
|
||||
|
||||
assert '"frame_index": 5' in regions_json
|
||||
assert canvas_width == 64
|
||||
assert canvas_height == 32
|
||||
assert region_count == 1
|
||||
|
||||
|
||||
def test_boxes_to_mask_batch_repeats_mask_to_batch():
|
||||
node = MODULE.URBoxesToMaskBatch()
|
||||
images = torch.zeros((2, 4, 4, 3), dtype=torch.float32)
|
||||
|
||||
mask, selected_frame, preview = node.build_masks(
|
||||
images=images,
|
||||
boxes_json="[[1, 1, 3, 3]]",
|
||||
frame_index=0,
|
||||
coordinate_format="xyxy",
|
||||
broadcast_mode="repeat_to_batch",
|
||||
normalized=False,
|
||||
preview_color="#FF0000",
|
||||
preview_opacity=0.5,
|
||||
)
|
||||
|
||||
expected = torch.tensor(
|
||||
[
|
||||
[0, 0, 0, 0],
|
||||
[0, 1, 1, 0],
|
||||
[0, 1, 1, 0],
|
||||
[0, 0, 0, 0],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
assert mask.shape == (2, 4, 4)
|
||||
assert torch.equal(mask[0], expected)
|
||||
assert torch.equal(mask[1], expected)
|
||||
assert selected_frame.shape == (1, 4, 4, 3)
|
||||
assert preview.shape == (1, 4, 4, 3)
|
||||
|
||||
|
||||
def test_boxes_to_mask_batch_handles_reversed_out_of_bounds_box():
|
||||
node = MODULE.URBoxesToMaskBatch()
|
||||
images = torch.zeros((1, 4, 4, 3), dtype=torch.float32)
|
||||
|
||||
mask, _, _ = node.build_masks(
|
||||
images=images,
|
||||
boxes_json="[[5, 5, -1, -1]]",
|
||||
frame_index=0,
|
||||
coordinate_format="xyxy",
|
||||
broadcast_mode="single_frame_only",
|
||||
normalized=False,
|
||||
preview_color="#FF0000",
|
||||
preview_opacity=0.5,
|
||||
)
|
||||
|
||||
assert mask.shape == (1, 4, 4)
|
||||
assert mask.sum().item() == 16.0
|
||||
|
||||
|
||||
def test_regions_to_mask_batch_supports_rect_circle_and_polygon():
|
||||
node = MODULE.URRegionsToMaskBatch()
|
||||
images = torch.zeros((1, 10, 10, 3), dtype=torch.float32)
|
||||
regions_json = """
|
||||
[
|
||||
{"type": "rect", "frame_index": 0, "x": 1, "y": 1, "w": 2, "h": 2},
|
||||
{"type": "circle", "frame_index": 0, "cx": 7, "cy": 7, "r": 1.5},
|
||||
{"type": "polygon", "frame_index": 0, "points": [[5, 1], [8, 1], [8, 3], [5, 3]]}
|
||||
]
|
||||
"""
|
||||
|
||||
mask, selected_frame, preview = node.build_masks(
|
||||
images=images,
|
||||
regions_json=regions_json,
|
||||
default_frame_index=0,
|
||||
broadcast_mode="single_frame_only",
|
||||
normalized=False,
|
||||
preview_frame=0,
|
||||
preview_color="#FF0000",
|
||||
preview_opacity=0.5,
|
||||
)
|
||||
|
||||
assert mask.shape == (1, 10, 10)
|
||||
assert mask[0, 1:3, 1:3].sum().item() == 4.0
|
||||
assert mask[0, 7, 7].item() == 1.0
|
||||
assert mask[0, 1:3, 5:8].sum().item() > 0.0
|
||||
assert selected_frame.shape == (1, 10, 10, 3)
|
||||
assert preview.shape == (1, 10, 10, 3)
|
||||
|
||||
|
||||
def test_regions_to_mask_batch_holds_between_keyframes():
|
||||
node = MODULE.URRegionsToMaskBatch()
|
||||
images = torch.zeros((4, 6, 6, 3), dtype=torch.float32)
|
||||
regions_json = """
|
||||
[
|
||||
{"type": "rect", "frame_index": 0, "x": 0, "y": 0, "w": 2, "h": 2},
|
||||
{"type": "rect", "frame_index": 2, "x": 4, "y": 4, "w": 2, "h": 2}
|
||||
]
|
||||
"""
|
||||
|
||||
mask, _, _ = node.build_masks(
|
||||
images=images,
|
||||
regions_json=regions_json,
|
||||
default_frame_index=0,
|
||||
broadcast_mode="hold_until_next_keyframe",
|
||||
normalized=False,
|
||||
preview_frame=0,
|
||||
preview_color="#FF0000",
|
||||
preview_opacity=0.5,
|
||||
)
|
||||
|
||||
assert mask.shape == (4, 6, 6)
|
||||
assert mask[0, 0:2, 0:2].sum().item() == 4.0
|
||||
assert mask[1, 0:2, 0:2].sum().item() == 4.0
|
||||
assert mask[2, 4:6, 4:6].sum().item() == 4.0
|
||||
assert mask[3, 4:6, 4:6].sum().item() == 4.0
|
||||
assert mask[1, 4:6, 4:6].sum().item() == 0.0
|
||||
|
||||
|
||||
def test_mask_processor_can_expand_mask_without_image():
|
||||
node = MODULE.URMaskProcessor()
|
||||
mask = torch.zeros((1, 5, 5), dtype=torch.float32)
|
||||
mask[:, 2, 2] = 1.0
|
||||
|
||||
processed_mask, mask_image, preview = node.process(
|
||||
mask=mask,
|
||||
expand_pixels=1,
|
||||
blur_sigma=0.0,
|
||||
blur_kernel=0,
|
||||
preview_color="#00FF00",
|
||||
preview_opacity=0.5,
|
||||
image=None,
|
||||
)
|
||||
|
||||
assert processed_mask.sum().item() == 9.0
|
||||
assert mask_image.shape == (1, 5, 5, 3)
|
||||
assert preview.shape == (1, 5, 5, 3)
|
||||
|
||||
|
||||
def test_keyframe_prep_fills_masked_region_and_pads():
|
||||
node = MODULE.URKeyframePrep()
|
||||
images = torch.zeros((1, 3, 5, 3), dtype=torch.float32)
|
||||
mask = torch.zeros((1, 3, 5), dtype=torch.float32)
|
||||
mask[:, 1, 2] = 1.0
|
||||
|
||||
original_frame, prepared_frame, prepared_mask, filled_frame, prepared_width, prepared_height, frame_count = node.prepare(
|
||||
images=images,
|
||||
mask=mask,
|
||||
frame_index=0,
|
||||
fill_color="#00FF00",
|
||||
resize_longest_side=0,
|
||||
pad_to_multiple=4,
|
||||
)
|
||||
|
||||
assert original_frame.shape == (1, 3, 5, 3)
|
||||
assert prepared_frame.shape == (1, 4, 8, 3)
|
||||
assert prepared_mask.shape == (1, 4, 8)
|
||||
assert prepared_width == 8
|
||||
assert prepared_height == 4
|
||||
assert frame_count == 1
|
||||
assert torch.allclose(filled_frame[0, 1, 2], torch.tensor([0.0, 1.0, 0.0]))
|
||||
|
||||
|
||||
def test_soft_composite_uses_mask_for_direct_blend_when_no_blur():
|
||||
node = MODULE.URSoftComposite()
|
||||
original = torch.zeros((1, 4, 4, 3), dtype=torch.float32)
|
||||
inpaint = torch.ones((1, 4, 4, 3), dtype=torch.float32)
|
||||
mask = torch.zeros((1, 4, 4), dtype=torch.float32)
|
||||
mask[:, 1:3, 1:3] = 1.0
|
||||
|
||||
composited, soft_mask = node.composite(
|
||||
inpaint=inpaint,
|
||||
original=original,
|
||||
mask=mask,
|
||||
kernel=1,
|
||||
sigma=0.0,
|
||||
)
|
||||
|
||||
assert torch.equal(soft_mask, mask)
|
||||
assert torch.allclose(composited[0, 0, 0], torch.tensor([0.0, 0.0, 0.0]))
|
||||
assert torch.allclose(composited[0, 1, 1], torch.tensor([1.0, 1.0, 1.0]))
|
||||
@@ -0,0 +1,535 @@
|
||||
import { app } from "../../scripts/app.js";
|
||||
|
||||
function chainCallback(object, property, callback) {
|
||||
const original = object[property];
|
||||
if (original) {
|
||||
object[property] = function (...args) {
|
||||
const result = original.apply(this, args);
|
||||
callback.apply(this, args);
|
||||
return result;
|
||||
};
|
||||
} else {
|
||||
object[property] = callback;
|
||||
}
|
||||
}
|
||||
|
||||
function findWidget(node, name) {
|
||||
return node.widgets?.find((widget) => widget.name === name);
|
||||
}
|
||||
|
||||
function clamp(value, min, max) {
|
||||
return Math.max(min, Math.min(max, value));
|
||||
}
|
||||
|
||||
function parseRegions(value) {
|
||||
try {
|
||||
const parsed = JSON.parse(value || "[]");
|
||||
if (Array.isArray(parsed)) return parsed;
|
||||
if (parsed && Array.isArray(parsed.regions)) return parsed.regions;
|
||||
if (parsed && typeof parsed === "object") return [parsed];
|
||||
} catch (error) {
|
||||
console.warn("UR Region Editor: failed to parse regions_json", error);
|
||||
}
|
||||
return [];
|
||||
}
|
||||
|
||||
function normalizeRect(region) {
|
||||
const x = Number(region.x ?? region.x1 ?? 0);
|
||||
const y = Number(region.y ?? region.y1 ?? 0);
|
||||
if (region.w != null || region.h != null) {
|
||||
return {
|
||||
x,
|
||||
y,
|
||||
w: Number(region.w ?? 0),
|
||||
h: Number(region.h ?? 0),
|
||||
};
|
||||
}
|
||||
const x2 = Number(region.x2 ?? region.right ?? x);
|
||||
const y2 = Number(region.y2 ?? region.bottom ?? y);
|
||||
return {
|
||||
x: Math.min(x, x2),
|
||||
y: Math.min(y, y2),
|
||||
w: Math.abs(x2 - x),
|
||||
h: Math.abs(y2 - y),
|
||||
};
|
||||
}
|
||||
|
||||
function createButton(label, title) {
|
||||
const button = document.createElement("button");
|
||||
button.textContent = label;
|
||||
button.title = title || label;
|
||||
button.style.cssText = "height:24px;padding:0 8px;border:1px solid #555;background:#2b2b2b;color:#ddd;border-radius:4px;cursor:pointer;font-size:12px;";
|
||||
button.addEventListener("mouseenter", () => {
|
||||
button.style.background = "#3a3a3a";
|
||||
});
|
||||
button.addEventListener("mouseleave", () => {
|
||||
button.style.background = button.dataset.active === "true" ? "#4b5f7a" : "#2b2b2b";
|
||||
});
|
||||
return button;
|
||||
}
|
||||
|
||||
function createNumberInput(value, min, title) {
|
||||
const input = document.createElement("input");
|
||||
input.type = "number";
|
||||
input.value = String(value);
|
||||
input.min = String(min ?? 0);
|
||||
input.title = title || "";
|
||||
input.style.cssText = "height:22px;width:76px;background:#111;color:#ddd;border:1px solid #555;border-radius:4px;padding:0 4px;font-size:12px;";
|
||||
return input;
|
||||
}
|
||||
|
||||
class RegionEditorWidget {
|
||||
constructor(node) {
|
||||
this.node = node;
|
||||
this.regionsWidget = findWidget(node, "regions_json");
|
||||
this.widthWidget = findWidget(node, "canvas_width");
|
||||
this.heightWidget = findWidget(node, "canvas_height");
|
||||
this.defaultFrameWidget = findWidget(node, "default_frame_index");
|
||||
|
||||
this.regions = parseRegions(this.regionsWidget?.value);
|
||||
this.mode = "rect";
|
||||
this.activeFrame = Number(this.defaultFrameWidget?.value ?? 0);
|
||||
this.selectedIndex = -1;
|
||||
this.dragStart = null;
|
||||
this.currentDraft = null;
|
||||
this.polygonDraft = [];
|
||||
this.background = null;
|
||||
this.video = null;
|
||||
this.videoFps = 24;
|
||||
|
||||
this.root = document.createElement("div");
|
||||
this.root.style.cssText = "display:flex;flex-direction:column;gap:6px;width:100%;height:100%;box-sizing:border-box;background:#121212;border:1px solid #303030;border-radius:6px;padding:6px;overflow:hidden;";
|
||||
|
||||
this.toolbar = document.createElement("div");
|
||||
this.toolbar.style.cssText = "display:flex;align-items:center;gap:6px;flex-wrap:wrap;";
|
||||
|
||||
this.rectBtn = createButton("Rect", "Draw rectangle by dragging");
|
||||
this.circleBtn = createButton("Circle", "Draw circle by dragging");
|
||||
this.polyBtn = createButton("Poly", "Click points, then Finish Poly");
|
||||
this.finishPolyBtn = createButton("Finish Poly", "Finish current polygon");
|
||||
this.undoPointBtn = createButton("Undo Point", "Remove last polygon point");
|
||||
this.deleteBtn = createButton("Delete", "Delete selected region");
|
||||
this.clearBtn = createButton("Clear", "Clear all regions");
|
||||
this.loadBgBtn = createButton("Load BG", "Load local image/video as editor background only");
|
||||
|
||||
this.frameInput = createNumberInput(this.activeFrame, 0, "Current keyframe index");
|
||||
this.frameLabel = document.createElement("span");
|
||||
this.frameLabel.textContent = "Frame";
|
||||
this.frameLabel.style.cssText = "color:#aaa;font-size:12px;";
|
||||
|
||||
this.fileInput = document.createElement("input");
|
||||
this.fileInput.type = "file";
|
||||
this.fileInput.accept = "image/*,video/*";
|
||||
this.fileInput.style.display = "none";
|
||||
|
||||
this.videoControls = document.createElement("div");
|
||||
this.videoControls.style.cssText = "display:none;align-items:center;gap:6px;width:100%;";
|
||||
this.playBtn = createButton("Play", "Play/pause background video");
|
||||
this.videoSlider = document.createElement("input");
|
||||
this.videoSlider.type = "range";
|
||||
this.videoSlider.min = "0";
|
||||
this.videoSlider.max = "0";
|
||||
this.videoSlider.step = "0.001";
|
||||
this.videoSlider.value = "0";
|
||||
this.videoSlider.style.cssText = "flex:1;";
|
||||
this.fpsInput = createNumberInput(this.videoFps, 1, "FPS used to convert video time to frame index");
|
||||
this.useVideoFrameBtn = createButton("Use Time", "Set keyframe from current video time");
|
||||
|
||||
this.canvas = document.createElement("canvas");
|
||||
this.canvas.width = Number(this.widthWidget?.value ?? 768);
|
||||
this.canvas.height = Number(this.heightWidget?.value ?? 512);
|
||||
this.canvas.style.cssText = "width:100%;height:auto;max-height:520px;object-fit:contain;background:#181818;border:1px solid #333;border-radius:4px;cursor:crosshair;";
|
||||
this.ctx = this.canvas.getContext("2d");
|
||||
|
||||
this.status = document.createElement("div");
|
||||
this.status.style.cssText = "color:#aaa;font-size:12px;white-space:pre-wrap;line-height:1.3;";
|
||||
|
||||
this.toolbar.append(
|
||||
this.rectBtn,
|
||||
this.circleBtn,
|
||||
this.polyBtn,
|
||||
this.finishPolyBtn,
|
||||
this.undoPointBtn,
|
||||
this.deleteBtn,
|
||||
this.clearBtn,
|
||||
this.loadBgBtn,
|
||||
this.frameLabel,
|
||||
this.frameInput,
|
||||
this.fileInput,
|
||||
);
|
||||
this.videoControls.append(this.playBtn, this.videoSlider, document.createTextNode("FPS"), this.fpsInput, this.useVideoFrameBtn);
|
||||
this.root.append(this.toolbar, this.videoControls, this.canvas, this.status);
|
||||
|
||||
this.bindEvents();
|
||||
this.setMode("rect");
|
||||
this.syncWidgets(false);
|
||||
this.render();
|
||||
}
|
||||
|
||||
bindEvents() {
|
||||
this.rectBtn.onclick = () => this.setMode("rect");
|
||||
this.circleBtn.onclick = () => this.setMode("circle");
|
||||
this.polyBtn.onclick = () => this.setMode("polygon");
|
||||
this.finishPolyBtn.onclick = () => this.finishPolygon();
|
||||
this.undoPointBtn.onclick = () => {
|
||||
this.polygonDraft.pop();
|
||||
this.render();
|
||||
};
|
||||
this.deleteBtn.onclick = () => this.deleteSelected();
|
||||
this.clearBtn.onclick = () => {
|
||||
this.regions = [];
|
||||
this.selectedIndex = -1;
|
||||
this.polygonDraft = [];
|
||||
this.syncWidgets(true);
|
||||
this.render();
|
||||
};
|
||||
this.loadBgBtn.onclick = () => this.fileInput.click();
|
||||
|
||||
this.frameInput.addEventListener("change", () => {
|
||||
this.activeFrame = Math.max(0, Number(this.frameInput.value || 0));
|
||||
if (this.defaultFrameWidget) this.defaultFrameWidget.value = this.activeFrame;
|
||||
this.render();
|
||||
});
|
||||
this.fpsInput.addEventListener("change", () => {
|
||||
this.videoFps = Math.max(1, Number(this.fpsInput.value || 24));
|
||||
});
|
||||
|
||||
this.fileInput.addEventListener("change", () => this.loadBackgroundFile());
|
||||
this.playBtn.onclick = () => this.toggleVideo();
|
||||
this.videoSlider.addEventListener("input", () => {
|
||||
if (!this.video) return;
|
||||
this.video.currentTime = Number(this.videoSlider.value || 0);
|
||||
this.drawVideoFrame();
|
||||
});
|
||||
this.useVideoFrameBtn.onclick = () => {
|
||||
if (!this.video) return;
|
||||
this.activeFrame = Math.max(0, Math.round(this.video.currentTime * this.videoFps));
|
||||
this.frameInput.value = String(this.activeFrame);
|
||||
if (this.defaultFrameWidget) this.defaultFrameWidget.value = this.activeFrame;
|
||||
this.render();
|
||||
};
|
||||
|
||||
this.canvas.addEventListener("mousedown", (event) => this.onMouseDown(event));
|
||||
this.canvas.addEventListener("mousemove", (event) => this.onMouseMove(event));
|
||||
this.canvas.addEventListener("mouseup", (event) => this.onMouseUp(event));
|
||||
this.canvas.addEventListener("dblclick", () => this.finishPolygon());
|
||||
this.canvas.addEventListener("contextmenu", (event) => {
|
||||
event.preventDefault();
|
||||
const point = this.toCanvasPoint(event);
|
||||
const index = this.hitTest(point.x, point.y);
|
||||
if (index >= 0) {
|
||||
this.selectedIndex = index;
|
||||
this.deleteSelected();
|
||||
}
|
||||
});
|
||||
document.addEventListener("keydown", (event) => {
|
||||
if (!this.root.isConnected) return;
|
||||
if (event.key === "Delete" || event.key === "Backspace") this.deleteSelected();
|
||||
if (event.key === "Enter" && this.mode === "polygon") this.finishPolygon();
|
||||
if (event.key === "Escape") {
|
||||
this.currentDraft = null;
|
||||
this.polygonDraft = [];
|
||||
this.render();
|
||||
}
|
||||
});
|
||||
|
||||
if (this.regionsWidget) {
|
||||
const previous = this.regionsWidget.callback;
|
||||
this.regionsWidget.callback = (...args) => {
|
||||
if (previous) previous.apply(this.regionsWidget, args);
|
||||
this.regions = parseRegions(this.regionsWidget.value);
|
||||
this.render();
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
setMode(mode) {
|
||||
this.mode = mode;
|
||||
for (const [button, name] of [[this.rectBtn, "rect"], [this.circleBtn, "circle"], [this.polyBtn, "polygon"]]) {
|
||||
button.dataset.active = name === mode ? "true" : "false";
|
||||
button.style.background = name === mode ? "#4b5f7a" : "#2b2b2b";
|
||||
}
|
||||
this.currentDraft = null;
|
||||
this.render();
|
||||
}
|
||||
|
||||
syncWidgets(markDirty) {
|
||||
const text = JSON.stringify(this.regions);
|
||||
if (this.regionsWidget) this.regionsWidget.value = text;
|
||||
if (this.widthWidget) this.widthWidget.value = this.canvas.width;
|
||||
if (this.heightWidget) this.heightWidget.value = this.canvas.height;
|
||||
if (this.defaultFrameWidget) this.defaultFrameWidget.value = this.activeFrame;
|
||||
if (markDirty) app.graph?.setDirtyCanvas(true, true);
|
||||
}
|
||||
|
||||
loadBackgroundFile() {
|
||||
const file = this.fileInput.files?.[0];
|
||||
if (!file) return;
|
||||
const url = URL.createObjectURL(file);
|
||||
this.background = null;
|
||||
this.video = null;
|
||||
this.videoControls.style.display = "none";
|
||||
if (file.type.startsWith("video/")) {
|
||||
const video = document.createElement("video");
|
||||
video.src = url;
|
||||
video.muted = true;
|
||||
video.playsInline = true;
|
||||
video.addEventListener("loadedmetadata", () => {
|
||||
this.video = video;
|
||||
this.canvas.width = video.videoWidth || this.canvas.width;
|
||||
this.canvas.height = video.videoHeight || this.canvas.height;
|
||||
this.videoSlider.max = String(video.duration || 0);
|
||||
this.videoControls.style.display = "flex";
|
||||
this.drawVideoFrame();
|
||||
this.syncWidgets(true);
|
||||
});
|
||||
video.addEventListener("timeupdate", () => {
|
||||
this.videoSlider.value = String(video.currentTime);
|
||||
this.drawVideoFrame();
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
const image = new Image();
|
||||
image.onload = () => {
|
||||
this.background = image;
|
||||
this.canvas.width = image.width;
|
||||
this.canvas.height = image.height;
|
||||
this.syncWidgets(true);
|
||||
this.render();
|
||||
};
|
||||
image.src = url;
|
||||
}
|
||||
|
||||
toggleVideo() {
|
||||
if (!this.video) return;
|
||||
if (this.video.paused) {
|
||||
this.video.play();
|
||||
this.playBtn.textContent = "Pause";
|
||||
} else {
|
||||
this.video.pause();
|
||||
this.playBtn.textContent = "Play";
|
||||
}
|
||||
}
|
||||
|
||||
drawVideoFrame() {
|
||||
this.render();
|
||||
}
|
||||
|
||||
toCanvasPoint(event) {
|
||||
const rect = this.canvas.getBoundingClientRect();
|
||||
return {
|
||||
x: clamp((event.clientX - rect.left) * (this.canvas.width / rect.width), 0, this.canvas.width),
|
||||
y: clamp((event.clientY - rect.top) * (this.canvas.height / rect.height), 0, this.canvas.height),
|
||||
};
|
||||
}
|
||||
|
||||
onMouseDown(event) {
|
||||
event.preventDefault();
|
||||
const point = this.toCanvasPoint(event);
|
||||
const hit = this.hitTest(point.x, point.y);
|
||||
if (hit >= 0 && event.shiftKey) {
|
||||
this.selectedIndex = hit;
|
||||
this.render();
|
||||
return;
|
||||
}
|
||||
if (this.mode === "polygon") {
|
||||
this.polygonDraft.push(point);
|
||||
this.render();
|
||||
return;
|
||||
}
|
||||
this.dragStart = point;
|
||||
this.currentDraft = { type: this.mode, frame_index: this.activeFrame, start: point, end: point };
|
||||
}
|
||||
|
||||
onMouseMove(event) {
|
||||
if (!this.currentDraft) return;
|
||||
this.currentDraft.end = this.toCanvasPoint(event);
|
||||
this.render();
|
||||
}
|
||||
|
||||
onMouseUp(event) {
|
||||
if (!this.currentDraft) return;
|
||||
this.currentDraft.end = this.toCanvasPoint(event);
|
||||
const region = this.draftToRegion(this.currentDraft);
|
||||
this.currentDraft = null;
|
||||
if (region) {
|
||||
this.regions.push(region);
|
||||
this.selectedIndex = this.regions.length - 1;
|
||||
this.syncWidgets(true);
|
||||
}
|
||||
this.render();
|
||||
}
|
||||
|
||||
draftToRegion(draft) {
|
||||
const x1 = draft.start.x;
|
||||
const y1 = draft.start.y;
|
||||
const x2 = draft.end.x;
|
||||
const y2 = draft.end.y;
|
||||
if (draft.type === "rect") {
|
||||
const x = Math.min(x1, x2);
|
||||
const y = Math.min(y1, y2);
|
||||
const w = Math.abs(x2 - x1);
|
||||
const h = Math.abs(y2 - y1);
|
||||
if (w < 3 || h < 3) return null;
|
||||
return { type: "rect", frame_index: this.activeFrame, x: Math.round(x), y: Math.round(y), w: Math.round(w), h: Math.round(h) };
|
||||
}
|
||||
if (draft.type === "circle") {
|
||||
const r = Math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2);
|
||||
if (r < 3) return null;
|
||||
return { type: "circle", frame_index: this.activeFrame, cx: Math.round(x1), cy: Math.round(y1), r: Math.round(r) };
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
finishPolygon() {
|
||||
if (this.polygonDraft.length < 3) return;
|
||||
this.regions.push({
|
||||
type: "polygon",
|
||||
frame_index: this.activeFrame,
|
||||
points: this.polygonDraft.map((point) => [Math.round(point.x), Math.round(point.y)]),
|
||||
});
|
||||
this.selectedIndex = this.regions.length - 1;
|
||||
this.polygonDraft = [];
|
||||
this.syncWidgets(true);
|
||||
this.render();
|
||||
}
|
||||
|
||||
deleteSelected() {
|
||||
if (this.selectedIndex < 0 || this.selectedIndex >= this.regions.length) return;
|
||||
this.regions.splice(this.selectedIndex, 1);
|
||||
this.selectedIndex = -1;
|
||||
this.syncWidgets(true);
|
||||
this.render();
|
||||
}
|
||||
|
||||
hitTest(x, y) {
|
||||
for (let i = this.regions.length - 1; i >= 0; i--) {
|
||||
const region = this.regions[i];
|
||||
if (Number(region.frame_index ?? 0) !== this.activeFrame) continue;
|
||||
const type = String(region.type || "rect").toLowerCase();
|
||||
if (type === "rect") {
|
||||
const rect = normalizeRect(region);
|
||||
if (x >= rect.x && x <= rect.x + rect.w && y >= rect.y && y <= rect.y + rect.h) return i;
|
||||
} else if (type === "circle") {
|
||||
const dx = x - Number(region.cx ?? region.x ?? 0);
|
||||
const dy = y - Number(region.cy ?? region.y ?? 0);
|
||||
if (dx * dx + dy * dy <= Number(region.r ?? region.radius ?? 0) ** 2) return i;
|
||||
}
|
||||
}
|
||||
return -1;
|
||||
}
|
||||
|
||||
render() {
|
||||
const ctx = this.ctx;
|
||||
ctx.clearRect(0, 0, this.canvas.width, this.canvas.height);
|
||||
ctx.fillStyle = "#181818";
|
||||
ctx.fillRect(0, 0, this.canvas.width, this.canvas.height);
|
||||
if (this.video) ctx.drawImage(this.video, 0, 0, this.canvas.width, this.canvas.height);
|
||||
else if (this.background) ctx.drawImage(this.background, 0, 0, this.canvas.width, this.canvas.height);
|
||||
else this.drawGrid();
|
||||
|
||||
this.regions.forEach((region, index) => this.drawRegion(region, index === this.selectedIndex));
|
||||
if (this.currentDraft) this.drawDraft(this.currentDraft);
|
||||
if (this.polygonDraft.length) this.drawPolygonDraft();
|
||||
this.status.textContent = `Mode: ${this.mode} | Frame: ${this.activeFrame} | Regions: ${this.regions.length}\nLoad BG is editor-only; regions_json is portable and saved in the workflow.`;
|
||||
}
|
||||
|
||||
drawGrid() {
|
||||
const ctx = this.ctx;
|
||||
ctx.strokeStyle = "#2a2a2a";
|
||||
ctx.lineWidth = 1;
|
||||
for (let x = 0; x <= this.canvas.width; x += 64) {
|
||||
ctx.beginPath();
|
||||
ctx.moveTo(x, 0);
|
||||
ctx.lineTo(x, this.canvas.height);
|
||||
ctx.stroke();
|
||||
}
|
||||
for (let y = 0; y <= this.canvas.height; y += 64) {
|
||||
ctx.beginPath();
|
||||
ctx.moveTo(0, y);
|
||||
ctx.lineTo(this.canvas.width, y);
|
||||
ctx.stroke();
|
||||
}
|
||||
ctx.fillStyle = "#777";
|
||||
ctx.font = "16px sans-serif";
|
||||
ctx.textAlign = "center";
|
||||
ctx.fillText("Load image/video background, then draw regions", this.canvas.width / 2, this.canvas.height / 2);
|
||||
}
|
||||
|
||||
drawRegion(region, selected) {
|
||||
const ctx = this.ctx;
|
||||
const type = String(region.type || "rect").toLowerCase();
|
||||
ctx.save();
|
||||
ctx.strokeStyle = selected ? "#ffcc00" : "#35c759";
|
||||
ctx.fillStyle = selected ? "rgba(255,204,0,0.20)" : "rgba(53,199,89,0.20)";
|
||||
ctx.lineWidth = selected ? 3 : 2;
|
||||
if (type === "rect") {
|
||||
const rect = normalizeRect(region);
|
||||
ctx.fillRect(rect.x, rect.y, rect.w, rect.h);
|
||||
ctx.strokeRect(rect.x, rect.y, rect.w, rect.h);
|
||||
} else if (type === "circle") {
|
||||
ctx.beginPath();
|
||||
ctx.arc(Number(region.cx ?? region.x ?? 0), Number(region.cy ?? region.y ?? 0), Number(region.r ?? region.radius ?? 0), 0, Math.PI * 2);
|
||||
ctx.fill();
|
||||
ctx.stroke();
|
||||
} else if (type === "polygon") {
|
||||
this.drawPolygonPath(region.points || []);
|
||||
ctx.fill();
|
||||
ctx.stroke();
|
||||
}
|
||||
ctx.restore();
|
||||
}
|
||||
|
||||
drawPolygonPath(points) {
|
||||
const ctx = this.ctx;
|
||||
ctx.beginPath();
|
||||
points.forEach((point, index) => {
|
||||
const x = Array.isArray(point) ? point[0] : point.x;
|
||||
const y = Array.isArray(point) ? point[1] : point.y;
|
||||
if (index === 0) ctx.moveTo(x, y);
|
||||
else ctx.lineTo(x, y);
|
||||
});
|
||||
ctx.closePath();
|
||||
}
|
||||
|
||||
drawDraft(draft) {
|
||||
const region = this.draftToRegion(draft);
|
||||
if (!region) return;
|
||||
this.drawRegion(region, true);
|
||||
}
|
||||
|
||||
drawPolygonDraft() {
|
||||
const ctx = this.ctx;
|
||||
ctx.save();
|
||||
ctx.strokeStyle = "#ffcc00";
|
||||
ctx.fillStyle = "#ffcc00";
|
||||
ctx.lineWidth = 2;
|
||||
ctx.beginPath();
|
||||
this.polygonDraft.forEach((point, index) => {
|
||||
if (index === 0) ctx.moveTo(point.x, point.y);
|
||||
else ctx.lineTo(point.x, point.y);
|
||||
ctx.fillRect(point.x - 3, point.y - 3, 6, 6);
|
||||
});
|
||||
ctx.stroke();
|
||||
ctx.restore();
|
||||
}
|
||||
}
|
||||
|
||||
app.registerExtension({
|
||||
name: "ComfyUI.UniversalRemove.RegionEditor",
|
||||
async beforeRegisterNodeDef(nodeType, nodeData) {
|
||||
if (nodeData?.name !== "UR Region Editor") return;
|
||||
chainCallback(nodeType.prototype, "onNodeCreated", function () {
|
||||
this.serialize_widgets = true;
|
||||
const editor = new RegionEditorWidget(this);
|
||||
this.urRegionEditor = editor;
|
||||
const widget = this.addDOMWidget("region_editor", "region_editor", editor.root);
|
||||
widget.computeSize = (width) => [width, Math.max(420, this.size?.[1] ? this.size[1] - 160 : 420)];
|
||||
this.setSize([Math.max(this.size?.[0] || 680, 680), Math.max(this.size?.[1] || 680, 680)]);
|
||||
chainCallback(this, "onResize", function () {
|
||||
editor.render();
|
||||
});
|
||||
});
|
||||
},
|
||||
});
|
||||
Reference in New Issue
Block a user