From 1613bd3194da3fdb431134d45e32beff13aaae4b Mon Sep 17 00:00:00 2001 From: nan Date: Sat, 22 Aug 2026 22:17:10 +0800 Subject: [PATCH] Add image perspective editing and related tools - Implemented the Image Perspective tool for editing image corners in Blender. - Created a new operator `EditImagePerspective` to handle the perspective transformation. - Added a preview feature for the perspective adjustment. - Introduced the `ImagePolylineTool` for editing image masks with polylines. - Developed backend support for Real-ESRGAN model inference, including ONNX integration. - Added functionality for image masking and background removal. - Created unit tests for the Real-ESRGAN ONNX model to ensure correct preprocessing and inference. --- README.md | 20 +- docs/development.md | 3 +- docs/features.md | 12 +- docs/image-tools-technical.md | 93 +-- docs/image-tools.md | 91 ++- docs/index.md | 2 +- docs/models-and-storage.md | 2 +- docs/outputs.md | 5 +- docs/remove-background.md | 2 +- docs/server.md | 10 +- docs/workflows.md | 6 +- mkdocs.yml | 2 +- src/anyimage/__init__.py | 23 +- src/anyimage/model_catalog.py | 35 +- src/anyimage/operators/__init__.py | 12 +- src/anyimage/operators/debug.py | 29 +- src/anyimage/operators/image_edit/__init__.py | 23 + .../{image_extract => image_edit}/box.py | 14 +- .../{image_extract => image_edit}/common.py | 237 +++---- .../{image_extract => image_edit}/lasso.py | 14 +- src/anyimage/operators/image_edit/line.py | 20 + src/anyimage/operators/image_edit/mask.py | 583 ++++++++++++++++++ .../perspective.py | 219 +------ src/anyimage/operators/image_edit/polyline.py | 20 + .../operators/image_extract/__init__.py | 33 - .../operators/image_extract/region.py | 378 ------------ src/anyimage/panel.py | 48 +- src/anyimage/preferences.py | 9 +- src/anyimage/properties.py | 34 +- src/anyimage/server/__init__.py | 16 +- src/anyimage/server/jobs/debug_realesrgan.py | 26 + .../jobs/{image_extract.py => image_mask.py} | 28 +- src/anyimage/server/model_catalog.py | 35 +- src/anyimage/server/model_manager.py | 29 + src/anyimage/server/models/__init__.py | 10 +- src/anyimage/server/models/ben2.py | 15 +- src/anyimage/server/models/onnx_realesrgan.py | 51 ++ src/anyimage/server/models/realesrgan.py | 12 + sync_models.py | 21 +- tests/test_addon.py | 485 ++++++++++----- tests/test_models.py | 29 + tests/test_pack.py | 22 +- tests/test_realesrgan_onnx.py | 75 +++ tests/test_server.py | 54 +- tests/test_worker.py | 48 ++ 45 files changed, 1802 insertions(+), 1133 deletions(-) create mode 100644 src/anyimage/operators/image_edit/__init__.py rename src/anyimage/operators/{image_extract => image_edit}/box.py (51%) rename src/anyimage/operators/{image_extract => image_edit}/common.py (54%) rename src/anyimage/operators/{image_extract => image_edit}/lasso.py (51%) create mode 100644 src/anyimage/operators/image_edit/line.py create mode 100644 src/anyimage/operators/image_edit/mask.py rename src/anyimage/operators/{image_extract => image_edit}/perspective.py (80%) create mode 100644 src/anyimage/operators/image_edit/polyline.py delete mode 100644 src/anyimage/operators/image_extract/__init__.py delete mode 100644 src/anyimage/operators/image_extract/region.py create mode 100644 src/anyimage/server/jobs/debug_realesrgan.py rename src/anyimage/server/jobs/{image_extract.py => image_mask.py} (84%) create mode 100644 src/anyimage/server/models/onnx_realesrgan.py create mode 100644 src/anyimage/server/models/realesrgan.py create mode 100644 tests/test_realesrgan_onnx.py diff --git a/README.md b/README.md index 854b1c4..b06e515 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ AnyImage 是一个 Blender 扩展，可以把一张照片、一段视频或一 - 在 3D View、Node Editor 与绘画模式中按 `Ctrl+V`（macOS 为 `Cmd+V`）直接粘贴并 Pack 剪贴板图片 - 不运行 AI，把 Image Empty 转换为可实时调整 Thickness 与 Subdivide 的 Mesh Plane - 从 Image Empty 套索提取不规则区域，通过 Poisson 或 Teddy / Chordal Axis 生成封闭 Balloon，并可用 BEN2 分析轮廓、MoGe-2 添加可调深度；源 Empty 保持不变 -- 用 Image Box / Lasso 原位提取 Image Empty 内容，按矩形或手绘范围裁切并 Pack 新图，也可用 BEN2 进一步收紧 Alpha +- 用 Image Box / Lasso / Line / Polyline 直接清理 Image Empty；按 `F` 反选，也可用 BEN2 细化选区 - 用 Image Perspective 勾勒四个角点，把透视四边形矫正为矩形裁切图，并沿用源 Image Empty 的空间朝向 - 一键把 Image Empty 转换为立体面片：让 AI 估算深度，或直接用原图作为深度 - 把单张 Image Empty 转换为透视纹理网格，并自动匹配原对象的 Transform 与 Offset @@ -58,27 +58,27 @@ Environment Bootstrap 使用连续的 0–100%：uv 为 1–20%、Python 为 20 - `Keep Original` 默认关闭，位于分割线下方；开启后保留源 Image Empty 4. 点击确认，稍等片刻就得到带位移的立体面片。`O DA3 Depth Plane` 与 `O MoGe2 Depth Plane` 都使用 `.color` 的 `Alpha > 0.5` 划分前景范围，范围外由 `Background Depth` 控制。 -### 套索提取图片 +### 编辑图片选区 -在 3D Viewport 左侧工具栏选择 Image Box / Lasso 工具组。默认的 Image Box 拖出矩形区域；长按工具按钮可切换 Image Lasso 并绘制自由轮廓。点击另一张 Image Empty 时只切换选择，不会把这次点击当作取区起点；点击当前 Active 图片、空白处或非图片物体时，才以当前 Active 图片开始取区，非图片物体不会被移动或抢走 Active。当前没有 Active 图片时会提示先点击 Image Empty。两者使用默认鼠标光标和类似 Select Tool 的白色虚线，分别采用 Sculpt Mode 的 Box Mask 与 Lasso Mask 图标，共用相同设置。该工具组与 Balloon Lasso 相邻，并通过一个间隔与前方的 Add Cube 分开。Image Box / Lasso 与 Balloon Lasso 都只放在工具栏，不出现在右键 AnyImage 子菜单中。 +在 3D Viewport 左侧工具栏选择 Image Mask 工具组。Box 拖出矩形，Lasso 手绘轮廓，Line 清理有向直线的一侧，Polyline 逐点构建多边形。点击另一张 Image Empty 时只切换选择；点击当前 Active 图片、空白处或非图片物体时，才开始编辑当前图片。没有 Active 图片时会提示先选择 Image Empty，切换其他工具会取消当前绘制。 -区域外的 RGBA 会统一设为透明黑，避免后续只读取 RGB 的图像处理重新识别透明区内容。输出画布取最终 Alpha 大于 0 的最小 bounding box，再按 Padding 向四周增加指定的透明像素，默认 0 px。对象原点会移动到新画布中心，显示尺寸与图片比例同步更新，裁切后的像素仍与原图在世界空间重合。Continuous 默认关闭并原位更新；开启后源 Empty 保持 Active，每次操作持续创建一个向当前 View 稍微靠近的新结果，避免与源图重叠闪烁。Remove Background 默认关闭；开启后复用标准 Remove Background 的 BEN2 单图处理，在选区内删除背景，再计算最终 bounding box。该选项仅支持静态图片，并需要 AI Environment 与 BEN2。 +绘制中按 `F` 切换反选：默认保留选区内，反选后保留选区外。Line 按住 `Ctrl` 以 5° 吸附；Polyline 可单击起点圆环或按 `Enter` 完成。`Refine Selection` 只用于 Box / Lasso：先用 BEN2 将几何选区收紧到其中的前景，再应用反选。结果直接替换当前 Empty 的图片，按最终非零 Alpha 紧裁切，并更新显示尺寸与 Image Offset；对象 Transform 不变。透明区保留源 RGB，避免纹理过滤产生黑边。 ### 四点透视矫正 -在 Image Box / Lasso 工具组中切换到 `Image Perspective`。工具使用 Sculpt Polyline Mask 图标；四个角可以按任意顺序点击，完成后会在源图片坐标中自动求出不交错的凸四边形，并以原图 X/Y 方向确定输出旋转。绘制时显示白色虚线和半透明选区面，`Backspace` 删除最近一个角点，右键或 `Esc` 取消；切换其他工具也会立即退出。首次点击遵循同一选择规则：点另一张图片只切换 Active，下一次点击才开始四点取区。Perspective 不依赖当前 View 的相机内参，因此透视和正交 View 都可以使用。 +`Image Perspective` 是 Image Mask 工具组之后的独立工具，位于该组与 Balloon Lasso 之间。四个角可以按任意顺序点击，完成后会在源图片坐标中自动求出不交错的凸四边形，并以原图 X/Y 方向确定输出旋转。绘制时显示白色虚线和半透明选区面，`Backspace` 删除最近一个角点，右键或 `Esc` 取消；切换其他工具也会立即退出。首次点击遵循同一选择规则。Perspective 不依赖当前 View 的相机内参，因此透视和正交 View 都可以使用。 第四个角点按下时会在 Viewport 中央显示带蓝色边框的矫正图片预览，中央显示当前比例，下方显示鼠标操作提示；鼠标同时移到 Viewport 中心，为左右两个方向保留相同的调节空间。水平移动连续调整宽高比，按住 `Ctrl` 吸附到 `1:1`、`4:3`、`3:4`、`16:9`、`9:16`，下一次左键确认并创建；`Backspace` 返回第三点，右键或 `Esc` 取消。初始比例取四边形两组对边的平均长度比。预览使用一次生成的 512 × 512 临时图片，移动鼠标时只改变 UI 显示比例。 确认后，工具把选区透视矫正为矩形图片，按最终非零 Alpha 范围裁切并 Pack 到 `.blend`。结果 Empty 保持源图的朝向和深度，以四边形在源图上的 bounding box 近似确定中心与显示范围，不再尝试从当前 View 反解真实空间矩形。四边形完全位于源图范围外，或矫正结果没有可见 Alpha 像素时，只显示警告且不创建结果。输出最长边不超过 8192 px，总像素不超过约 16 MP。 -Image Perspective 与 Image Box / Lasso 共用 Tool Settings。`Padding` 在矫正后的 Alpha 裁切外增加透明像素；`Continuous` 控制原位更新或连续创建分层副本；`Remove Background` 默认关闭，开启后先复用标准 BEN2 去背景处理，再执行相同的透视矫正、裁切和 Pack 流程。 +Image Perspective 是 deform 工具，不提供 Refine 或反选。它与其他 Image 工具共享选图、绘制样式、紧裁切和 Empty 写回逻辑。 -完整操作说明见 [Image Box、Lasso 与 Perspective](docs/image-tools.md)。 +完整操作说明见 [Image 编辑工具](docs/image-tools.md)。 ### 套索生成 Balloon -选中单张图片的 Image Empty，右键打开 `AnyImage` 子菜单并选择 `Balloon Lasso Tool`。Blender 会激活 3D Viewport 左侧工具栏底部、紧邻 Image Box / Lasso 工具组的独立 Balloon Lasso；该工具组与 Add Cube 之间保留一个间隔。Balloon 使用 Poly Build 的绿色图标、Annotate 同款铅笔光标，Tool Header 与状态栏共同提示当前工具状态，不再用黄色外框或强制修改 Shading、Overlay。Blender 的公开 Workspace Tool API 不支持按 Active Object 动态隐藏工具，因此图标在 Object Mode 中保持可见，但只有 Active Object 是 Image Empty 时才能启动 Stroke。 +选中单张图片的 Image Empty，右键打开 `AnyImage` 子菜单并选择 `Balloon Lasso Tool`。Blender 会激活 3D Viewport 左侧工具栏底部、位于独立 Image Perspective 之后的 Balloon Lasso；Image Mask 工具组与 Add Cube 之间保留一个间隔。Balloon 使用 Poly Build 的绿色图标、Annotate 同款铅笔光标，Tool Header 与状态栏共同提示当前工具状态，不再用黄色外框或强制修改 Shading、Overlay。Blender 的公开 Workspace Tool API 不支持按 Active Object 动态隐藏工具，因此图标在 Object Mode 中保持可见，但只有 Active Object 是 Image Empty 时才能启动 Stroke。 按住左键沿目标轮廓绘制，松开后生成一个带原图 UV 的封闭网格。每次拖动都是一个独立完成的顶层 `Extract Balloon` Operator，工具本身继续保持激活，下一次拖动会创建另一个 Balloon；因此 `Ctrl+Z` 只撤销最近一次生成，`Ctrl+Shift+Z` 重做。每次生成后源 Image Empty 都保持选中和 Active。未绘制时没有持续 Modal 拦截事件，视图导航和其他原生快捷键由 Blender 正常处理；绘制中可用右键或 `Esc` 取消当前轨迹。选择左侧工具栏中的其他工具即可离开 Balloon Lasso。 @@ -111,7 +111,7 @@ MoGe-2 的像素级 normal 会作为 Object Space normal 贴图接入 `O MoGe2 S Bump 合并；网格本身不写入 custom normal，因此重拓扑降低顶点数后仍能保留细腻的法线细节。 当前此功能仅支持单张图片。 -MoGe-2 ViT-S、ViT-B、ViT-L、DA3 Small、DA3 Base 与 BEN2 全部使用 ONNX。模型执行、相机与深度恢复、 +MoGe-2 ViT-S、ViT-B、ViT-L、DA3 Small、DA3 Base、BEN2 与 Debug 使用的 Real-ESRGAN General x4v3 全部使用 ONNX。模型执行、相机与深度恢复、 点云重投影和网格生成均不依赖 PyTorch。Windows 使用 DirectML 加速；Apple M4 或更新芯片可使用 CoreML 的 `MLProgram` 模型格式，较早或无法识别的 Mac 只提供 CPU；provider 不可用时会回退到 ONNX Runtime CPU， 并在任务状态中显示实际 provider。 @@ -131,7 +131,7 @@ Layer Color 同时驱动 Principled BSDF 的 Base Color、Subsurface Radius 和 AnyImage 支持视频（mp4、mkv、mov 等）和序列帧，会逐帧处理并生成逐帧对应的立体面片。`.color` 与 `.depth` 按源 Empty Image 的类型载入：普通单图为 `IMAGE`，Movie/Sequence 即使只处理一帧也仍为 `SEQUENCE`。处理范围默认 跟随 Empty 上设置的 `Image Frames` 与 `Offset`，不需要额外配置。DA3 的转换任务复用同一个 ONNX Session 逐帧推理，显存峰值不随视频帧数平方增长。Debug 只保存模型 API 的直接返回：DA3 多帧输入 -保留联合推理，MoGe-2 与 BEN2 按各自的单图 API 逐帧保存。 +保留联合推理，MoGe-2、BEN2 与 Real-ESRGAN General x4v3 按各自的单图 API 逐帧保存。General x4v3 直接输出 4× RGB `prediction*.png`，不做切片、强度混合或图片替换。 ## 任务结果保存在哪里 diff --git a/docs/development.md b/docs/development.md index 6b7bd61..404b98f 100644 --- a/docs/development.md +++ b/docs/development.md @@ -37,7 +37,7 @@ Mesh Plane、Depth Plane、Surface、Background 和 Balloon 按功能组织。 Server 的模型模块返回 Prediction，几何模块处理数组与网格，输出模块写文件。Blender 数据只在 Operator Response 的主线程阶段创建。 -Image Box、Lasso 与 Perspective 位于 `operators/image_extract/` 功能包。模块职责、交互状态、像素管线、背景清理和 Empty 定位见 [Image 工具技术实现](image-tools-technical.md)。 +Image Box、Lasso、Line、Polyline 与 Perspective 位于 `operators/image_edit/` 功能包。模块职责、交互状态、像素管线、Refine 和 Empty 定位见 [Image 工具技术实现](image-tools-technical.md)。 ## 项目结构 @@ -122,5 +122,6 @@ BLENDER_BIN="/c/Program Files/Blender Foundation/Blender 4.5/blender.exe" - Depth Anything 3 代码采用 Apache 2.0 - BEN2 ONNX 模型采用其仓库声明的 Apache 2.0 License +- Real-ESRGAN General x4v3 ONNX 权重采用其模型仓库声明的 BSD 3-Clause License - AnyImage Extension 代码采用 GPL-3.0-or-later - 模型权重许可证由所选 Hugging Face 模型仓库决定 diff --git a/docs/features.md b/docs/features.md index 74c158f..8efc02b 100644 --- a/docs/features.md +++ b/docs/features.md @@ -11,7 +11,7 @@ | Depth Plane | 快速制作可调节的立体图片 | 单图、视频、序列帧 | 带颜色、深度与 Geometry Nodes 的面片 | | Surface | 从单张图片恢复透视空间 | 单图 | 相机空间网格、材质与可选匹配相机 | | Remove Background | 提取透明前景 | 单图、视频、序列帧 | 透明图片或透明序列 | -| Image Lasso | 手绘提取并原位更新图片 | 静态 Image Empty | Pack 的透明裁切图 | +| Image Mask | 用 Box、Lasso、Line 或 Polyline 直接清理图片 | 静态 Image Empty | 原位更新并 Pack 的透明裁切图 | | Image Perspective | 四点透视矫正裁切 | 静态 Image Empty | Pack 的矩形裁切图 | | Balloon | 沿轮廓绘制鼓起造型 | 静态 Image Empty | 带原图 UV 的封闭网格 | | Debug | 检查模型原始输出 | 模型支持的图片或帧序列 | 保存在 Job 目录中的原始 Prediction | @@ -30,9 +30,9 @@ Balloon 作为不依赖 AI 的基础工具单独排列；三个 AI 功能在缺 转换任务完成后，AnyImage 在一个 Undo 步骤中创建网格、图片、材质和修改器。源 Image Empty 是否保留由转换选项决定。 -Image Box、Image Lasso 与 Image Perspective 组成一个工具组，共用源图选择、`Remove Background`、`Continuous` 和 `Padding`。Box 提取矩形，Lasso 提取自由轮廓，Perspective 将任意顺序输入的四个角矫正为指定比例的矩形。结果都是 Pack 的 Image Empty，并沿用源图朝向与深度。 +Image Box、Lasso、Line 与 Polyline 组成 Mask 工具组，并共用 `EditImageMask`。四者用 `F` 切换反选；`Refine Selection` 只用于 Box / Lasso。Perspective 是位于 Mask 组与 Balloon Lasso 之间的独立工具，将任意顺序输入的四个角矫正为指定比例的矩形。结果直接更新当前 Image Empty，并 Pack 到 `.blend`。 -操作步骤见 [Image Box、Lasso 与 Perspective](image-tools.md)；处理管线与模块职责见 [Image 工具技术实现](image-tools-technical.md)。 +操作步骤见 [Image 编辑工具](image-tools.md)；处理管线与模块职责见 [Image 工具技术实现](image-tools-technical.md)。 Balloon 默认在 Blender 内部生成，无需 Environment 或模型。Tool Settings 提供模式、厚度、采样距离、材质复用、`Contour`、`Depth` 与 MoGe-2 `Depth Model`；AI 未就绪时只禁用这三个控件，并在 Tool Settings 中显示安装按钮。Depth 会把物理距离相对化并稳健拟合初始基面，保留未截断的正负起伏。Base Plane 子 Empty 的旋转定义基面朝向，沿自身局部 Z 移动定义 middle level。 @@ -50,7 +50,7 @@ Blender 无法整体 Pack 图片序列，因此动画结果会继续引用 Job AI 任务在独立 Python Environment 中通过 ONNX Runtime 执行。Blender 负责界面、任务提交和结果导入，Job Server 负责模型与文件处理。 -默认 Storage Root 为 `~/.anyimage`，其中保存 Environment、模型、日志和 Job 结果。连续任务会复用已加载的模型 Session；DA3 与 MoGe-2 共享一个深度模型缓存位置，BEN2 可同时保留。 +默认 Storage Root 为 `~/.anyimage`，其中保存 Environment、模型、日志和 Job 结果。连续任务会复用已加载的模型 Session；DA3 与 MoGe-2 共享一个深度模型缓存位置，BEN2 与 Real-ESRGAN 可同时保留。 安装、下载和推理都在后台运行，并通过状态栏提供进度与取消操作。详细设计见[Job Server](server.md)。 @@ -62,6 +62,8 @@ AI 任务在独立 Python Environment 中通过 ONNX Runtime 执行。Blender - DA3 Small、DA3 Base - BEN2 Base +Debug 另外提供 Real-ESRGAN General x4v3，用于直接检查 4× RGB 放大能力；它当前不属于正式图片工具。 + 首次安装会准备 MoGe-2 ViT-S Normal 与 BEN2 Base，其他模型可从 AnyImage 扩展偏好设置按需下载。 -Debug 工作流保存模型 API 的直接返回，适合验证模型、设备与输入。Depth Plane 和 Surface 则在模型结果之上继续生成 Blender 可用的贴图与几何。文件字段见[输出格式](outputs.md)。 +Debug 工作流保存模型 API 的直接返回，适合验证模型、设备与输入。General x4v3 对每张输入整图推理并保存 `prediction*.png`，不加入切片、降噪强度、Alpha 合成或 Blender 图片写回。Depth Plane 和 Surface 则在模型结果之上继续生成 Blender 可用的贴图与几何。文件字段见[输出格式](outputs.md)。 diff --git a/docs/image-tools-technical.md b/docs/image-tools-technical.md index a66a1cc..45fb199 100644 --- a/docs/image-tools-technical.md +++ b/docs/image-tools-technical.md @@ -1,77 +1,84 @@ # Image 工具技术实现 -Image Box、Lasso 与 Perspective 位于 `operators/image_extract/`。三者共享源图选择、设置、图片写回、空间定位、Continuous 分层和可选背景清理。 +Image 工具位于 `operators/image_edit/`。Box、Lasso、Line 与 Polyline 共享直接编辑管线；Perspective 在工具栏中独立显示并处理 Homography，但复用源图选择、绘制和 Empty 写回。 ## 模块职责 ```mermaid flowchart LR - Box["box.py\nWorkspaceTool"] --> Region["region.py\nBox / Lasso Operator"] - Lasso["lasso.py\nWorkspaceTool"] --> Region - Perspective["perspective.py\n四点与透视矫正"] - Region --> Common["common.py\n共享选择、设置与结果写回"] + Tools["box / lasso / line / polyline\nWorkspaceTool"] --> Mask["mask.py\n手势、Refine、Invert"] + Perspective["perspective.py\n四点、比例、Homography"] + Mask --> Common["common.py\n选择、绘制、裁切、Empty 写回"] Perspective --> Common - Common --> Job["extract-image-background\n可选 BEN2 Job"] + Mask --> Job["refine-image-mask\nBEN2 内容 mask"] ``` -`box.py` 与 `lasso.py` 只声明 WorkspaceTool。两者通过 `selection_shape` 复用 `ExtractImageRegion`；Perspective 保留独立 Operator，因为它还有四点输入、比例预览和 Homography。 +- `box.py`、`lasso.py`、`line.py`、`polyline.py` 只声明 WorkspaceTool。 +- `mask.py` 用 `EditImageMask` 统一四种手势；`RefineImageMask` 负责异步 BEN2 响应。 +- `perspective.py` 维护四点输入、比例预览和全分辨率重采样。 +- `common.py` 只放跨两条管线实际共用的逻辑。 ## 首次点击 -`resolve_image_extract_click()` 将首次点击归为两种结果： +`resolve_image_edit_click()` 临时调用 Viewport Select 判断点击目标： -| 点击目标 | 结果 | +| 点击目标 | 行为 | | --- | --- | -| 另一张 Image Empty | 只选择并设为 Active | -| 当前 Active 图片、空白或非图片对象 | 使用当前 Active 图片开始取区 | +| 另一张 Image Empty | 选择并设为 Active，本次不编辑 | +| 当前 Active 图片、空白或非图片对象 | 恢复原选择并开始编辑当前图片 | -这段逻辑由三个工具共用。点击非图片对象后会恢复原选择，避免触发对象移动或更换源图。 +这避免工具在无有效源图时退化为对象移动，也让用户无需切回 Select Tool。 -## 处理管线 +## Mask 管线 ```mermaid flowchart TD - Input["屏幕选区"] --> Pixels["转换到源图片像素坐标"] - Pixels --> Background{"Remove Background?"} - Background -- 否 --> Local["Blender 内处理 RGBA"] - Background -- 是 --> Server["选区画布 → BEN2 → 局部图片"] - Local --> Tool{"工具类型"} - Server --> Tool - Tool -- Box / Lasso --> Crop["Alpha bounds + Padding"] - Tool -- Perspective --> Warp["预乘 Alpha Homography"] - Warp --> Crop - Crop --> Result["创建并 Pack Image\n定位 Empty"] + Gesture["屏幕手势"] --> Polygon["源图片像素 polygon / mask"] + Polygon --> Refine{"Refine Selection?"} + Refine -- 否 --> Content["几何 mask"] + Refine -- 是 --> Server["选区外 RGBA 清零 → BEN2"] + Server --> Content2["BEN2 Alpha × 几何 mask"] + Content --> Invert{"Invert?"} + Content2 --> Invert + Invert -- 否 --> Final["final mask"] + Invert -- 是 --> Reverse["1 - final mask"] + Final --> Apply["源 Alpha × mask"] + Reverse --> Apply + Apply --> Crop["按非零 Alpha 紧裁切\n保留源 RGB"] + Crop --> Write["替换 Image 并更新 Empty 显示"] ``` -Box 与 Lasso 将选区外 RGBA 清为透明黑，再按非零 Alpha 的最小范围裁切。Perspective 先将四点规范为沿源图 X/Y 方向的凸四边形，再按确认的宽高比重采样；点击顺序不参与输出旋转。 +Box / Lasso 的顺序固定为 **Refine → Invert**。BEN2 只提供内容 Alpha；最终 RGB 始终来自源图。Blender 侧只修改 Alpha，不清理透明区 RGB，以避免 Viewport 纹理过滤产生黑边。Line / Polyline 强制关闭 Refine，只执行几何 mask 与 Invert。 -透视预览固定生成一次 512 × 512 临时 Image。鼠标移动只改变屏幕上的预览比例；确认后才执行全分辨率重采样。 +Line 的屏幕坐标 Y 轴与图片坐标 Y 轴方向相反，因此像素阶段会反转线段端点，使半透明预览侧与实际保留侧一致；`Ctrl` 只影响屏幕端点的 5° 吸附。零长度点击在投影前直接取消，投影结果不足两个点时也会安全终止。Polyline 接近起点时绘制闭合圆环，并支持单击起点或 `Enter` 完成；完成时显式闭合，不把当前游标位置作为额外顶点。 -## 背景清理 +## Perspective 管线 -`extract-image-background` 与标准 `remove-background` Job 共用 `remove_image_background()`。图片工具的 Job 只增加三项选区语义： +四点先按源图片 X/Y 方向规范为凸四边形，输入顺序不参与输出方向。第四点后只生成一次 512 × 512 临时纹理；移动鼠标改变 UI 比例，确认时才进行全分辨率 Homography。 -- polygon mask； -- 局部图片在源图中的 bounds； -- Box / Lasso 使用的透明 Padding。 +透视重采样使用预乘 Alpha 插值，随后恢复 RGB、清理透明像素并紧裁切。输出最长边限制为 8192 px，总像素限制约 16 MP。 -Perspective 请求的 Job Padding 固定为 0。局部结果回到 Blender 后先换算为源图坐标，再执行透视矫正和最终 Padding。 +Perspective 不调用 BEN2，也不支持 Invert。它是 deform 工具，不属于 mask 清理语义。 -## Empty 定位与 Continuous +## Empty 写回 -Box 与 Lasso 根据裁切前后的像素 bounds 换算 Empty 的中心和显示尺寸。Perspective 使用源四边形的轴对齐 bounds 近似放置结果。三者都沿用源 Empty 的矩阵，因此朝向与深度不变。 +`replace_empty_image_pixels()` 创建并 Pack 新 Image datablock，然后替换当前 Empty 的 `data`。它不复制对象，也不修改 `matrix_world`。 -Continuous 关闭时直接更新源 Empty；开启时复制对象，并沿操作开始时捕获的 View 方向偏移： +`align_empty_display_to_bounds()` 根据源画布、旧 `plane_bounds`、结果尺寸和放置 bounds 更新： -```text -offset = display_size × world_scale × 0.001 × layer -``` +- `empty_display_size`； +- `empty_image_offset`。 -`layer` 记录在源对象的 `_anyimage_extract_layer` Custom Property 中，使连续结果依次靠近当前 View。 +Box、Lasso、Line 与 Polyline 的放置 bounds 等于最终裁切范围，因此像素与原图片平面严格对齐。Perspective 的结果比例可能不同于四边形轴对齐范围，因此只保持中心和覆盖尺度近似一致。 -## 维护约束 +## Server 边界 -- WorkspaceTool 的显示设置放在 `box.py`、`lasso.py`、`perspective.py`。 -- Box / Lasso 的共享交互放在 `region.py`，跨三个工具的行为放在 `common.py`。 -- Blender 数据只在主线程 Operator 中读写；Server 参数保持为可序列化字典。 -- 普通裁切不启动 Job Server；只有 Remove Background 需要 Environment 与 BEN2。 +普通 Mask 和 Perspective 完全在 Blender 侧执行。只有 Box / Lasso 的 **Refine Selection** 提交 `refine-image-mask`： + +1. 根据 polygon 创建局部透明画布； +2. 清空选区外 RGBA；BEN2 公共入口再以输入 Alpha 遮住透明 RGB； +3. 使用共享的 `remove_image_background()` 运行 BEN2； +4. 再乘几何 selection mask； +5. 返回紧裁切的 `image.png` 与源图坐标 `bounds.json`。 + +Blender 只读取返回图片的 Alpha，并把它应用到原图 RGB。相同的 BEN2 输入清理规则也用于 Remove Background、Depth Plane、Surface 和 Balloon；正式输出使用源 RGB 与 BEN2 Alpha。Blender 数据只在主线程 Operator 中读写，Job 参数保持为可序列化字典。 diff --git a/docs/image-tools.md b/docs/image-tools.md index 6f1ade9..e03105a 100644 --- a/docs/image-tools.md +++ b/docs/image-tools.md @@ -1,85 +1,82 @@ -# Image Box、Lasso 与 Perspective +# Image 编辑工具 -这三个工具用于从 Image Empty 中提取局部图片。结果仍是 Image Empty，并自动 Pack 到 `.blend`。 - -它们只处理静态图片。视频和序列帧请使用其他转换工具。 +Image Box、Lasso、Line、Polyline 与 Perspective 直接编辑静态 Image Empty，并把结果 Pack 到当前 `.blend`。前四个工具组成 Mask 工具组；Perspective 是紧随其后的独立工具，Balloon Lasso 位于 Perspective 之后。 ## 选择工具 -| 工具 | 适合 | +| 工具 | 用途 | | --- | --- | -| **Image Box** | 提取矩形区域，操作最快 | -| **Image Lasso** | 沿物体或不规则边界自由绘制 | -| **Image Perspective** | 将海报、屏幕、门窗等四边形内容矫正为矩形 | +| **Image Box** | 用矩形保留或清理图片区域 | +| **Image Lasso** | 手绘自由选区 | +| **Image Line** | 用一条直线清理一侧 | +| **Image Polyline** | 逐点构建精确多边形 | +| **Image Perspective** | 将任意顺序输入的四个角矫正为矩形 | -在 3D Viewport 左侧工具栏找到 Image 工具组。长按当前图标可以切换三个工具；设置位于 Tool Settings。 +在 3D Viewport 左侧工具栏长按 Image Mask 图标，可以在 Box、Lasso、Line 与 Polyline 之间切换。**Refine Selection** 只显示在 Box / Lasso；Line / Polyline 保持几何直线边界。Perspective 使用独立图标。
-**截图占位：工具组与 Tool Settings** +**截图占位：Image 工具组与 Tool Settings** -建议同时显示三个工具入口，以及 Remove Background、Continuous、Padding。 +显示包含四个工具的 Image Mask 组、独立 Perspective、Balloon Lasso，以及 Box / Lasso 的 Refine Selection。 目标文件：`images/image-tools/tool-group.png`
-## 选择源图片 +## 选择图片 -如果当前 Active 已是目标图片，可以直接开始取区。 +直接在当前 Active 图片上开始绘制。如果第一次点击的是另一张 Image Empty，这次点击只会选中它；下一次操作才开始编辑。点击空白处或非图片对象时，继续编辑当前 Active 图片，不会移动对象。 -如果要处理另一张图片，先点击它。这一次点击只负责选择；下一次点击或拖动才开始取区。点击空白处或其他类型的对象时，工具继续使用当前 Active 图片。 +没有 Active Image Empty 时，工具会提示先选择图片。切换到其他工具会立即取消正在进行的绘制。 -!!! tip - 最简单的习惯是：先点一下目标图片确认 Active，再开始绘制。 +## 清理选区 -## 使用 Box 或 Lasso +### Box 与 Lasso -=== "Image Box" +- Box：按住左键拖出矩形，松开后执行。 +- Lasso：按住左键绘制轮廓，松开后执行。 - 1. 按住左键，从区域一角拖到另一角。 - 2. 松开左键，立即生成结果。 +### Line -=== "Image Lasso" +从起点拖到终点，松开后保留直线左侧的内容。按住 `Ctrl` 以 5° 为单位吸附方向；需要另一侧时按 `F` 翻转。 - 1. 按住左键，沿需要保留的区域绘制。 - 2. 松开左键，立即生成结果。 +### Polyline -选区可以越过图片边缘。只要范围内仍有可见像素，工具就会按实际内容裁切。 +逐点单击构建多边形。鼠标接近起点时会显示闭合圆环；单击圆环或按 `Enter` 完成，`Backspace` 删除最近一点。 + +绘制过程中按 `F` 切换反选。反选是一个持续到本次操作完成的状态，不需要按住： + +- 关闭：保留选区内； +- 开启：保留选区外。 + +Box / Lasso 启用 **Refine Selection** 后，BEN2 会先把几何选区收紧到其中的前景，再应用反选。例如框住一个苹果并开启反选，最终会清理苹果本身、保留其余内容。Line / Polyline 不提供 Refine，以保持直线边界和明确的裁切方向。 + +结果会直接替换当前 Empty 的图片，并把画布紧裁切到最终可见 Alpha。对象 Transform 不变，显示尺寸与 Image Offset 会更新，使剩余像素保持在原图片平面上的位置。透明区只修改 Alpha，保留源 RGB 供纹理过滤使用；完全透明的结果不会写入。
-**截图占位：Box 与 Lasso 取区** +**截图占位：Polyline 与反选** -建议用一张并排图展示矩形拖拽和自由轮廓，不需要额外结果图。 +显示半透明选区面、虚线边界，以及状态栏中的 `F toggles Invert`。 -目标文件：`images/image-tools/box-lasso.png` +目标文件：`images/image-tools/polyline-invert.png`
-## 使用 Perspective +## 四点透视 -1. 依次点击内容的四个角，顺序不限。 -2. 第四点完成后，Viewport 中央会显示矫正预览。 -3. 左右移动鼠标调整宽高比。 -4. 单击左键确认并生成结果。 +1. 任意顺序单击四个角点。 +2. 第四点后，在 Viewport 中央检查矫正预览。 +3. 左右移动鼠标调整宽高比；按住 `Ctrl` 可吸附到 `1:1`、`4:3`、`3:4`、`16:9`、`9:16`。 +4. 左键确认。 -调整比例时按住 `Ctrl`，可以吸附到 `1:1`、`4:3`、`3:4`、`16:9` 和 `9:16`。`Backspace` 返回上一个角点；右键或 `Esc` 取消。 +`Backspace` 返回上一步，右键或 `Esc` 取消。四点会自动排序，点击顺序不会改变输出旋转。 + +Perspective 只执行图片 Homography 与紧裁切，不做 Refine，也不尝试从当前 View 推导真实相机。结果保持源 Empty 的 Transform，并按原四边形覆盖范围近似更新显示位置与尺寸。选区完全位于图片外或没有可见像素时不会执行。
**截图占位：Perspective 比例确认** -建议截取中央图片预览、蓝色边框、比例文字和底部操作提示。 +显示中央图片预览、蓝色边框、当前比例与底部操作提示。 目标文件：`images/image-tools/perspective-aspect.png`
-## 结果设置 - -- **Remove Background**：使用 BEN2 删除选区内的背景，再按可见内容裁切。首次使用需要先完成 AI Environment 与 BEN2 配置。 -- **Continuous**：保留源图片，连续创建新结果。新图片会向当前 View 稍微靠近，减少重叠闪烁。 -- **Padding**：在可见内容四周增加透明像素，单位为 px。 - -Continuous 关闭时，结果会原位替换当前 Image Empty；开启时，源图片保持 Active，便于继续提取其他区域。 - -## Perspective 的结果位置 - -Perspective 会矫正图片内容，但不会尝试还原真实相机或空间平面。结果沿用源图片的朝向和深度，并大致放在原四边形覆盖的区域。 - -需要了解交互状态、像素处理和模块职责时，请阅读 [Image 工具技术实现](image-tools-technical.md)。 +实现结构与像素规则见 [Image 工具技术实现](image-tools-technical.md)。 diff --git a/docs/index.md b/docs/index.md index 4e9fce5..0e1db14 100644 --- a/docs/index.md +++ b/docs/index.md @@ -9,7 +9,7 @@ AnyImage 是一款 Blender 扩展，可以把照片、视频或序列帧转换 - 用 **Depth Plane** 快速得到可调节的浮雕式面片 - 用 **Surface** 从单张图片生成透视纹理网格 - 用 **Remove Background** 一键获得透明前景 -- 用 **Image Perspective** 将四点选区矫正为沿用源 Empty 朝向的矩形裁切图 +- 用 **Image Box / Lasso / Line / Polyline** 直接清理图片，或用 **Image Perspective** 做四点透视矫正 - 处理视频和序列帧，生成可随时间播放的 Depth Plane - 使用 Geometry Nodes 继续调整深度、细分、重拓扑与平滑效果 diff --git a/docs/models-and-storage.md b/docs/models-and-storage.md index 47e006a..cca8320 100644 --- a/docs/models-and-storage.md +++ b/docs/models-and-storage.md @@ -9,7 +9,7 @@ AnyImage 把运行环境、模型和任务结果放在 Blender 之外，多个 B - **Ready** 表示模型文件已经就绪。 - **Download** 用于按需下载模型。 - 默认安装会准备 MoGe-2 ViT-S Normal 和 BEN2 Base。 -- DA3 与更大的 MoGe-2 模型需要时再下载即可。 +- DA3、更大的 MoGe-2，以及只供 Debug 使用的 Real-ESRGAN General x4v3 需要时再下载即可。 较大的模型不一定适合所有图片。建议先用默认模型完成构图，再针对细节与性能需求比较其他模型。 diff --git a/docs/outputs.md b/docs/outputs.md index 0f6f1e4..bfc6ae1 100644 --- a/docs/outputs.md +++ b/docs/outputs.md @@ -9,11 +9,12 @@ | `generate-depth` | `results.npz`、深度与颜色序列 | Depth Plane 所需深度、相机和贴图 | | `generate-surface` | `mesh.npz`、`normal.png`、颜色层 | Surface 网格、相机与材质输入 | | `remove-background` | `color.0001.png` 起 | 带 Alpha 的颜色层 | -| `extract-image-background` | `selection-input.png`、`image.png`、`bounds.json` | 图片工具选区内去背景后的局部图片与原图坐标 bounds | +| `refine-image-mask` | `selection-input.png`、`image.png`、`bounds.json` | BEN2 细化后的局部内容 mask 与原图坐标 bounds | | `analyze-balloon` | `lasso-input.png`，按选项输出 `foreground.png`、`contour.json`、`depth.npz`、`normal.png` | 套索约束输入、BEN2 轮廓与 MoGe-2 深度、法线 | | `debug-da3` | `prediction.npz` | DA3 API 返回 | | `debug-moge2` | `prediction*.npz` | MoGe-2 API 返回 | | `debug-ben2` | `foreground*.png` | BEN2 API 返回 | +| `debug-realesrgan-general-x4v3` | `prediction*.png` | Real-ESRGAN General x4v3 的 4× RGB API 返回 | ## NPZ 字段 @@ -23,7 +24,7 @@ Balloon Depth 的 `depth.npz` 包含相对 `depth_signal`、归一化拟合平 Surface 的 `mesh.npz` 包含 `vertices`、`faces`、`uv`、`intrinsics`、`image_size` 与 `reference_depth`，以及 Geometry Nodes 使用的网格属性。 -Debug 文件保持模型 API 的数据形态，用于比较推理结果；正式 Job 则提供 Blender 导入所需的派生数据。 +Debug 文件保持模型 API 的数据形态，用于比较推理结果；正式 Job 则提供 Blender 导入所需的派生数据。General x4v3 的浮点 RGB 返回只做 `[0, 1]` 限制和 8-bit PNG 序列化，不做额外画质处理。 ## 图片序列 diff --git a/docs/remove-background.md b/docs/remove-background.md index fe6bf76..6047378 100644 --- a/docs/remove-background.md +++ b/docs/remove-background.md @@ -35,4 +35,4 @@ Depth Plane 和 Surface 的转换对话框也提供 **Remove Background**： 当前正式入口不提供 BEN2 refinement 参数。可以在结果材质中继续处理 Alpha，或在外部图像软件中修整生成的颜色图。 -输入图片已有 Alpha 时，AnyImage 会先用该 Alpha 遮住 BEN2 推理输入中的透明 RGB，再将结果 Alpha 计算为“输入 Alpha × BEN2 Alpha”。因此 Image Box / Lasso 或外部工具创建的透明范围不会被去背景重新扩张；软边缘也会继续保留。 +所有 BEN2 任务都会先用输入 Alpha 遮住推理输入中的透明 RGB。正式输出只采用 BEN2 的 Alpha，颜色继续使用源 RGB；已有 Alpha 时，结果 Alpha 为“输入 Alpha × BEN2 Alpha”。因此 Image Box / Lasso 或外部工具创建的透明范围不会被去背景重新扩张，软边缘也会继续保留，同时避免透明黑参与 Viewport 纹理过滤。 diff --git a/docs/server.md b/docs/server.md index 70f39bd..3b5c088 100644 --- a/docs/server.md +++ b/docs/server.md @@ -59,7 +59,9 @@ Environment 使用 Python 3.12、ONNX Runtime、NumPy、Pillow、PyAV、trimesh Balloon 的 `analyze-balloon` Job 会裁切套索输入并清零选区外 RGBA。`Contour` 用 BEN2 回传最大连通轮廓；`Depth` 用 MoGe-2 对有效相机距离稳健拟合二维基面，并把物理距离换算为不受绝对拍摄距离影响、未截断的相对深度信号。Job 同时把 MoGe normal 转换为 tangent-space 纹理，铺回原图尺寸并以中性法线填充无效区域。Blender 主线程负责采样深度属性、创建 Base Plane Empty、组装 `O Balloon Shape` 修改器与 `O Balloon Layer` 材质。 -图片 Box、Lasso 与 Perspective 的可选 `Remove Background` 使用 `extract-image-background` Job。它与标准 `remove-background` Job 共用单图 BEN2 去背景函数，再叠加图片工具专属的 polygon mask、局部 bounds 与透明 Padding。 +Image Box 与 Lasso 的 `Refine Selection` 使用 `refine-image-mask` Job。它与标准 `remove-background` Job 共用单图 BEN2 函数；Server 在推理前清空 polygon 外的 RGBA，推理后再次应用几何 mask，并返回局部 Alpha 与源图 bounds。Line、Polyline 与 Perspective 不提交 Job。 + +所有 BEN2 调用都在公共模型入口用输入 Alpha 遮住透明 RGB。Remove Background、Depth Plane、Surface 与 Image Refine 的正式结果保留源 RGB，只采用 BEN2 Alpha；Balloon 同样只消费 BEN2 Alpha。该边界既防止模型识别透明区隐藏内容，也避免把透明黑写回 Blender 后产生过滤黑边。 正式功能支持： @@ -67,10 +69,14 @@ Balloon 的 `analyze-balloon` Job 会裁切套索输入并清零选区外 RGBA - DA3 Small、DA3 Base - BEN2 Base -首次安装会下载 MoGe-2 ViT-S Normal 与 BEN2 Base。模型保存在 `models/`，可由多个 Blender 版本共享。Server 缓存 BEN2 和最近使用的一个深度模型，以控制内存占用。 +Debug 任务另外支持 Real-ESRGAN General x4v3；它只提供模型直通测试，不属于正式功能。 + +首次安装会下载 MoGe-2 ViT-S Normal 与 BEN2 Base。Real-ESRGAN General x4v3 由扩展偏好设置按需下载。模型保存在 `models/`，可由多个 Blender 版本共享。Server 缓存 BEN2、Real-ESRGAN 和最近使用的一个深度模型，以控制内存占用。 模型下载优先使用 Hugging Face，并以 OmooLab 镜像作为回退。AnyImage 用同一个准备 Operator 封装 BlendJob 安装任务，并按状态显示 `Set Up AI Environment…` 或 `Download Required Models…`；`Use Mirror` 属性只影响 Environment Bootstrap，不改变模型下载策略。Blender UI 通过同一个状态函数检查 Environment 指纹与默认模型文件：缺 Environment 时执行完整安装，只有默认模型缺失时仅提交补齐模型任务。 +`debug-realesrgan-general-x4v3` 将每张输入的完整 RGB 动态尺寸张量直接送入 ONNX，并把 4× 返回保存为 `prediction*.png`。该 Job 不实现 tile、强度混合、Alpha 合成或 Blender 结果导入。 + ## 任务生命周期 Server 在第一次提交 Job 或手动 Start 时启动。任务进入单 Worker FIFO Queue，同一时间执行一个推理、下载或媒体任务。 diff --git a/docs/workflows.md b/docs/workflows.md index 07d61cd..74ba53e 100644 --- a/docs/workflows.md +++ b/docs/workflows.md @@ -8,8 +8,8 @@ AnyImage 提供多种日常工作流。选择结果时，先考虑你需要的 | Depth Plane | 人物、产品、插画、动态背景、快速场景搭建 | 图片、视频、序列帧 | 带深度位移的平面 | 快速、可动画、Geometry Nodes 参数直观 | | Surface | 建筑、室内、街景、需要匹配视角的单图 | 单张图片 | 透视纹理网格 | 空间感更强，可创建匹配相机 | | Remove Background | 抠出人物或物体，继续作为 Image Empty 使用 | 图片、视频、序列帧 | 带透明 Alpha 的颜色图 | 不创建网格，保留 Empty 的位置和显示设置 | -| Image Box / Lasso | 矩形或手绘提取局部内容并按 Alpha 范围裁切 | 静态图片 | Pack 的透明裁切图 | 支持 Padding 与 Continuous，不创建网格 | -| Image Perspective | 海报、屏幕、门窗等四边形内容的透视矫正裁切 | 静态图片 | Pack 的矩形图片 | 四点任意顺序，共用 Remove Background、Padding 与 Continuous | +| Image Box / Lasso / Line / Polyline | 用几何选区直接清理图片 | 静态图片 | 原位更新的透明裁切图 | `F` 反选，可用 BEN2 细化选区 | +| Image Perspective | 海报、屏幕、门窗等四边形内容的透视矫正裁切 | 静态图片 | 原位更新的矩形图片 | 四点任意顺序，交互调整比例 | ## 什么时候选择 Mesh Plane @@ -37,4 +37,4 @@ Surface 会根据模型估算的相机和深度构建透视网格。相较 Depth ## 什么时候使用 Image Perspective -当图片中的矩形物体因为透视变成四边形，并且希望得到一张正视矩形裁切图时，使用 3D Viewport 工具栏中的 **Image Perspective**。它只把四点当作图片裁切与 Homography 输入，不假设当前 View 的焦距与原图拍摄相机一致。四点顺序不影响输出方向；结果 Empty 保持源图朝向和深度，并按源图中的选区范围近似放置。需要持续提取多个区域时，可开启与 Image Box / Lasso 共用的 Continuous。 +当图片中的矩形物体因为透视变成四边形，并且希望得到一张正视矩形裁切图时，使用 3D Viewport 工具栏中的 **Image Perspective**。它只把四点当作图片裁切与 Homography 输入，不假设当前 View 的焦距与原图拍摄相机一致。四点顺序不影响输出方向；结果 Empty 保持源图 Transform，并按源图中的选区范围近似更新显示位置与尺寸。 diff --git a/mkdocs.yml b/mkdocs.yml index 06ca7cc..f513eb5 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -15,7 +15,7 @@ nav: - 粘贴剪贴板图片: clipboard-image.md - 使用指南: - 选择转换方式: workflows.md - - Image Box、Lasso 与 Perspective: image-tools.md + - Image 编辑工具: image-tools.md - Mesh Plane: mesh-plane.md - Depth Plane: guide-depth-plane.md - Surface: guide-surface.md diff --git a/src/anyimage/__init__.py b/src/anyimage/__init__.py index e601577..212e44c 100644 --- a/src/anyimage/__init__.py +++ b/src/anyimage/__init__.py @@ -6,7 +6,13 @@ from .clipboard_image import CLASSES as CLIPBOARD_CLASSES from .runtime import runtime from .operators import CLASSES as OPERATOR_CLASSES from .operators.balloon import BalloonTool -from .operators.image_extract import ImageBoxTool, ImageLassoTool, ImagePerspectiveTool +from .operators.image_edit import ( + ImageBoxTool, + ImageLassoTool, + ImageLineTool, + ImagePerspectiveTool, + ImagePolylineTool, +) from .preferences import AnyImagePreferences from .properties import AnyImageSettings from .menu import AnyImageImageMenu, draw_image_context_menu @@ -15,6 +21,7 @@ from .panel import ( Da3OutputDebugPanel, DebugPanel, Moge2ApiDebugPanel, + RealesrganGeneralX4v3DebugPanel, ServerPanel, ) @@ -30,6 +37,7 @@ CLASSES = ( Moge2ApiDebugPanel, Da3OutputDebugPanel, Ben2RemovalDebugPanel, + RealesrganGeneralX4v3DebugPanel, ) @@ -47,10 +55,19 @@ def register(): separator=False, ) bpy.utils.register_tool( - ImagePerspectiveTool, + ImageLineTool, after={ImageLassoTool.bl_idname}, separator=False, ) + bpy.utils.register_tool( + ImagePolylineTool, + after={ImageLineTool.bl_idname}, + separator=False, + ) + bpy.utils.register_tool( + ImagePerspectiveTool, + separator=False, + ) bpy.utils.register_tool(BalloonTool, separator=False) bpy.types.Scene.anyimage_settings = bpy.props.PointerProperty( type=AnyImageSettings, @@ -66,6 +83,8 @@ def unregister(): clipboard_image.unregister_keymaps() bpy.utils.unregister_tool(BalloonTool) bpy.utils.unregister_tool(ImagePerspectiveTool) + bpy.utils.unregister_tool(ImagePolylineTool) + bpy.utils.unregister_tool(ImageLineTool) bpy.utils.unregister_tool(ImageLassoTool) bpy.utils.unregister_tool(ImageBoxTool) bpy.types.VIEW3D_MT_object_context_menu.remove(draw_image_context_menu) diff --git a/src/anyimage/model_catalog.py b/src/anyimage/model_catalog.py index 4e5260b..f63f039 100644 --- a/src/anyimage/model_catalog.py +++ b/src/anyimage/model_catalog.py @@ -89,6 +89,30 @@ BEN2_MODEL = ModelSpec( ) +REALESRGAN_MODEL = ModelSpec( + key="REALESRGAN_GENERAL_X4V3", + label="Real-ESRGAN General x4v3", + directory_name="realesr-general-x4v3", + repository="jonathanst29/tinier-upscale-models", + license_name="BSD 3-Clause", + description="Fast general-purpose 4x image upscaler using ONNX Runtime", + supports_gs=False, + download_patterns=("realesr-general-x4v3.onnx",), + ready_patterns=("realesr-general-x4v3.onnx",), + r2_directory="realesr-general-x4v3", + r2_files=(( + "realesr-general-x4v3.onnx", + 4866396, + "924ebad6532777303582d4ce7811849b88869231a3ae7093e0f21200249df8d5", + ),), +) + + +UPSCALE_MODELS = { + REALESRGAN_MODEL.key: REALESRGAN_MODEL, +} + + MOGE2_MODELS = { "MOGE2_VITS_NORMAL": ModelSpec( key="MOGE2_VITS_NORMAL", @@ -146,6 +170,7 @@ MOGE2_MODEL = MOGE2_MODELS["MOGE2_VITS_NORMAL"] DOWNLOADABLE_MODELS = { **MODELS, BEN2_MODEL.key: BEN2_MODEL, + **UPSCALE_MODELS, **MOGE2_MODELS, } MODEL_NUMBERS = { @@ -181,8 +206,14 @@ def model_record(model, ready=False): "license": model.license_name, "supports_gs": model.supports_gs, "enabled": model.enabled, - "family": "moge2" if model.key.startswith("MOGE2_") else ( - "ben2" if model.key == BEN2_MODEL.key else "da3" + "family": ( + "moge2" + if model.key.startswith("MOGE2_") + else "ben2" + if model.key == BEN2_MODEL.key + else "upscale" + if model.key in UPSCALE_MODELS + else "da3" ), "ready": bool(ready), } diff --git a/src/anyimage/operators/__init__.py b/src/anyimage/operators/__init__.py index 4bb6373..274e9a5 100644 --- a/src/anyimage/operators/__init__.py +++ b/src/anyimage/operators/__init__.py @@ -4,8 +4,13 @@ from .balloon import ( ExtractBalloon, AnalyzeBalloon, ) -from .debug import DebugBen2, DebugDa3, DebugMoge2 -from .image_extract import CLASSES as IMAGE_EXTRACT_CLASSES +from .debug import ( + DebugBen2, + DebugDa3, + DebugMoge2, + DebugRealesrganGeneralX4v3, +) +from .image_edit import CLASSES as IMAGE_EDIT_CLASSES from .depth_plane import ( ConvertToDepthPlane, ExtractOriginalFrames, @@ -32,6 +37,7 @@ CLASSES = ( DebugDa3, DebugMoge2, DebugBen2, + DebugRealesrganGeneralX4v3, GenerateDa3Depth, GenerateMoge2Depth, GenerateDa3Surface, @@ -40,7 +46,7 @@ CLASSES = ( GenerateOriginalBackground, ExtractOriginalFrames, ClearModels, - *IMAGE_EXTRACT_CLASSES, + *IMAGE_EDIT_CLASSES, ExtractBalloon, AnalyzeBalloon, ActivateBalloonTool, diff --git a/src/anyimage/operators/debug.py b/src/anyimage/operators/debug.py index c943f6e..238e290 100644 --- a/src/anyimage/operators/debug.py +++ b/src/anyimage/operators/debug.py @@ -1,7 +1,10 @@ import bpy from ..runtime import JobOperatorBase -from ..properties import BEN2_MODEL_KEY +from ..properties import ( + BEN2_MODEL_KEY, + REALESRGAN_MODEL_KEY, +) from .common import ( background_input_path, da3_parameters, @@ -89,3 +92,27 @@ class DebugBen2(JobOperatorBase, bpy.types.Operator): def response(self, _context, result): open_job_directory(result.directory) return f"Removed background to {result.directory}" + + +class DebugRealesrganGeneralX4v3(JobOperatorBase, bpy.types.Operator): + bl_idname = "anyimage.debug_realesrgan_general_x4v3" + bl_label = "Debug Real-ESRGAN General x4v3" + job_type = "debug-realesrgan-general-x4v3" + starting_message = "Starting Real-ESRGAN General x4v3 inference" + + def request(self, context): + require_environment() + settings = context.scene.anyimage_settings + require_model(REALESRGAN_MODEL_KEY) + return { + "input": str(input_path(settings.input_path)), + "model": REALESRGAN_MODEL_KEY, + "device": settings.device.lower(), + "video_frames": settings.video_frames, + "start_frame": 1, + } + + def response(self, _context, result): + open_job_directory(result.directory) + return f"Saved Real-ESRGAN API prediction to {result.directory}" + diff --git a/src/anyimage/operators/image_edit/__init__.py b/src/anyimage/operators/image_edit/__init__.py new file mode 100644 index 0000000..706ffd8 --- /dev/null +++ b/src/anyimage/operators/image_edit/__init__.py @@ -0,0 +1,23 @@ +from .box import ImageBoxTool +from .lasso import ImageLassoTool +from .line import ImageLineTool +from .mask import ActivateImageEditTool, EditImageMask, RefineImageMask +from .perspective import EditImagePerspective, ImagePerspectiveTool +from .polyline import ImagePolylineTool + + +CLASSES = ( + EditImageMask, + RefineImageMask, + ActivateImageEditTool, + EditImagePerspective, +) + +__all__ = ( + "CLASSES", + "ImageBoxTool", + "ImageLassoTool", + "ImageLineTool", + "ImagePolylineTool", + "ImagePerspectiveTool", +) diff --git a/src/anyimage/operators/image_extract/box.py b/src/anyimage/operators/image_edit/box.py similarity index 51% rename from src/anyimage/operators/image_extract/box.py rename to src/anyimage/operators/image_edit/box.py index 98ed8b1..3bef843 100644 --- a/src/anyimage/operators/image_extract/box.py +++ b/src/anyimage/operators/image_edit/box.py @@ -1,7 +1,7 @@ import bpy -from .common import draw_image_extract_settings -from .region import ExtractImageRegion +from .common import draw_image_mask_settings +from .mask import EditImageMask class ImageBoxTool(bpy.types.WorkSpaceTool): @@ -9,16 +9,16 @@ class ImageBoxTool(bpy.types.WorkSpaceTool): bl_context_mode = "OBJECT" bl_idname = "anyimage.image_box" bl_label = "Image Box" - bl_description = "Extract a boxed region from an Image Empty" + bl_description = "Edit an Image Empty with a rectangular mask" bl_icon = "ops.sculpt.border_mask" - bl_operator = ExtractImageRegion.bl_idname + bl_operator = EditImageMask.bl_idname bl_keymap = ( ( - ExtractImageRegion.bl_idname, + EditImageMask.bl_idname, {"type": "LEFTMOUSE", "value": "PRESS"}, - {"properties": [("selection_shape", "BOX")]}, + {"properties": [("gesture", "BOX")]}, ), ) def draw_settings(context, layout, _tool): - draw_image_extract_settings(context, layout) + draw_image_mask_settings(context, layout) diff --git a/src/anyimage/operators/image_extract/common.py b/src/anyimage/operators/image_edit/common.py similarity index 54% rename from src/anyimage/operators/image_extract/common.py rename to src/anyimage/operators/image_edit/common.py index 6077753..40aa73b 100644 --- a/src/anyimage/operators/image_extract/common.py +++ b/src/anyimage/operators/image_edit/common.py @@ -8,19 +8,12 @@ from ...properties import ( ai_ready, ai_setup_label, ai_status, - image_extract_continuous_property as continuous_property, - image_extract_remove_background_property as remove_background_property, - image_extract_padding_property as padding_property, ) from ..common import production_device, require_environment, require_model from ..image import image_base_name, is_image_empty, plane_bounds -EXTRACT_LAYER_PROPERTY = "_anyimage_extract_layer" -EXTRACT_LAYER_STEP = 0.001 - - -def image_extract_poll(context): +def image_edit_poll(context): return context.area is not None and context.area.type == "VIEW_3D" @@ -29,12 +22,15 @@ def _restore_selection(context, active_object, selected_objects): for selected in selected_objects: if bpy.data.objects.get(selected.name) is not None: selected.select_set(True) - if active_object is not None and bpy.data.objects.get(active_object.name) is not None: + if ( + active_object is not None + and bpy.data.objects.get(active_object.name) is not None + ): context.view_layer.objects.active = active_object -def resolve_image_extract_click(context, event, report): - """Resolve the first click into source selection or extraction start.""" +def resolve_image_edit_click(context, event, report): + """Resolve the first click into source selection or image editing.""" active_object = context.active_object selected_objects = tuple(context.selected_objects) try: @@ -57,23 +53,16 @@ def resolve_image_extract_click(context, event, report): return None, False -def capture_view_direction(context): - from mathutils import Vector - - direction = ( - context.region_data.view_matrix.inverted().to_3x3() - @ Vector((0.0, 0.0, 1.0)) - ).normalized() - return json.dumps(tuple(direction), separators=(",", ":")) - - -def decode_view_direction(data): - from mathutils import Vector - - direction = Vector(json.loads(data)) - if direction.length <= 1e-8: - raise ValueError("The stored View direction is invalid") - return direction.normalized() +def active_view3d_tool_id(context): + tools = getattr(getattr(context, "workspace", None), "tools", None) + from_mode = getattr(tools, "from_space_view3d_mode", None) + if not callable(from_mode): + return None + try: + tool = from_mode(getattr(context, "mode", "OBJECT"), create=False) + except (RuntimeError, TypeError): + return None + return getattr(tool, "idname", None) def polygon_bounds(polygon): @@ -93,12 +82,10 @@ def polygon_bounds(polygon): return left, top, right, bottom -def crop_pixels_to_alpha(pixels, bounds, padding=0): - """Crop Blender-order RGBA pixels to nonzero Alpha plus pixel padding.""" +def crop_pixels_to_alpha(pixels, bounds): + """Crop Blender-order RGBA pixels to their nonzero Alpha bounds.""" import numpy as np - if padding < 0: - raise ValueError("Padding must not be negative") left, top, right, bottom = bounds width = right - left height = bottom - top @@ -107,65 +94,61 @@ def crop_pixels_to_alpha(pixels, bounds, padding=0): ) visible_y, visible_x = np.nonzero(rgba[:, :, 3] > 0.0) if not len(visible_x): - raise ValueError("The selection contains no visible image pixels") + raise ValueError("The edit contains no visible image pixels") content_left = int(visible_x.min()) content_top = int(visible_y.min()) content_right = int(visible_x.max()) + 1 content_bottom = int(visible_y.max()) + 1 - output_width = content_right - content_left + padding * 2 - output_height = content_bottom - content_top + padding * 2 - output = np.zeros((output_height, output_width, 4), dtype=np.float32) - output[ - padding : padding + content_bottom - content_top, - padding : padding + content_right - content_left, - ] = rgba[content_top:content_bottom, content_left:content_right] - output[output[:, :, 3] <= 0.0] = 0.0 + output = rgba[content_top:content_bottom, content_left:content_right].copy() return np.flipud(output).ravel(), ( - left + content_left - padding, - top + content_top - padding, - left + content_right + padding, - top + content_bottom + padding, + left + content_left, + top + content_top, + left + content_right, + top + content_bottom, ) -def align_empty_to_crop(source_object, matrix_world, source_size, crop_bounds): - from mathutils import Matrix - - source_width, source_height = source_size - left, top, right, bottom = crop_bounds +def align_empty_display_to_bounds( + source_object, source_size, output_size, placement_bounds +): + """Place a new image canvas without changing the Empty transform.""" + source_width, source_height = (float(value) for value in source_size) + output_width, output_height = (float(value) for value in output_size) + left, top, right, bottom = (float(value) for value in placement_bounds) x_min, x_max, y_min, y_max = plane_bounds(source_object) local_width = x_max - x_min local_height = y_max - y_min - crop_x_min = x_min + left * local_width / source_width - crop_x_max = x_min + right * local_width / source_width - crop_y_max = y_max - top * local_height / source_height - crop_y_min = y_max - bottom * local_height / source_height - center = ( - (crop_x_min + crop_x_max) * 0.5, - (crop_y_min + crop_y_max) * 0.5, - 0.0, + desired_x_min = x_min + left * local_width / source_width + desired_x_max = x_min + right * local_width / source_width + desired_y_max = y_max - top * local_height / source_height + desired_y_min = y_max - bottom * local_height / source_height + desired_width = desired_x_max - desired_x_min + desired_height = desired_y_max - desired_y_min + display_size = max(desired_width, desired_height) + aspect = output_width / output_height + if aspect >= 1.0: + display_width = display_size + display_height = display_size / aspect + else: + display_width = display_size * aspect + display_height = display_size + center_x = (desired_x_min + desired_x_max) * 0.5 + center_y = (desired_y_min + desired_y_max) * 0.5 + source_object.empty_display_size = display_size + source_object.empty_image_offset = ( + center_x / display_width - 0.5, + center_y / display_height - 0.5, ) - source_object.matrix_world = matrix_world @ Matrix.Translation(center) - source_object.empty_display_size = max( - crop_x_max - crop_x_min, - crop_y_max - crop_y_min, - ) - source_object.empty_image_offset = (-0.5, -0.5) def replace_empty_image_pixels( - source_object, - matrix_world, - pixels, - output_size, - placement_bounds, - suffix=".extract", + source_object, pixels, output_size, placement_bounds, suffix=".edit" ): source_image = source_object.data source_size = tuple(source_image.size) output_width, output_height = (int(value) for value in output_size) if output_width <= 0 or output_height <= 0: - raise ValueError("The extracted image dimensions must be positive") + raise ValueError("The edited image dimensions must be positive") image = bpy.data.images.new( f"{image_base_name(source_image)}{suffix}", width=output_width, @@ -178,8 +161,11 @@ def replace_empty_image_pixels( image.pixels.foreach_set(pixels) image.pack() image.update() - align_empty_to_crop( - source_object, matrix_world, source_size, placement_bounds + align_empty_display_to_bounds( + source_object, + source_size, + (output_width, output_height), + placement_bounds, ) source_object.data = image except Exception: @@ -188,59 +174,7 @@ def replace_empty_image_pixels( return image -def replace_empty_image_from_path(source_object, matrix_world, path, crop_bounds): - source_size = tuple(source_object.data.size) - colorspace = source_object.data.colorspace_settings.name - try: - image = bpy.data.images.load(str(path), check_existing=False) - except (OSError, RuntimeError) as error: - raise RuntimeError("Unable to load the extracted image") from error - try: - image.name = f"{image_base_name(source_object.data)}.extract" - image.colorspace_settings.name = colorspace - image.alpha_mode = "STRAIGHT" - image.pack() - align_empty_to_crop(source_object, matrix_world, source_size, crop_bounds) - source_object.data = image - except Exception: - bpy.data.images.remove(image, do_unlink=True) - raise - return image - - -def image_extract_target(context, source_object, continuous): - if not continuous: - return source_object, 0 - layer = int(source_object.get(EXTRACT_LAYER_PROPERTY, 0)) + 1 - source_object[EXTRACT_LAYER_PROPERTY] = layer - target = source_object.copy() - target[EXTRACT_LAYER_PROPERTY] = layer - collection = ( - source_object.users_collection[0] - if source_object.users_collection - else context.collection - ) - collection.objects.link(target) - return target, layer - - -def offset_continuous_result(target, view_direction_data, layer): - if layer <= 0: - return - direction = decode_view_direction(view_direction_data) - world_scale = max(abs(value) for value in target.matrix_world.to_scale()) - distance = target.empty_display_size * world_scale * EXTRACT_LAYER_STEP * layer - target.matrix_world.translation += direction * distance - - -def select_image_extract_result(context, source_object, target, continuous): - bpy.ops.object.select_all(action="DESELECT") - selected = source_object if continuous else target - selected.select_set(True) - context.view_layer.objects.active = selected - - -def image_extract_background_request(input_path, selection_data, padding): +def image_mask_refine_request(input_path, selection_data): require_environment() require_model(BEN2_MODEL_KEY) polygon = json.loads(selection_data) @@ -248,28 +182,29 @@ def image_extract_background_request(input_path, selection_data, padding): return { "input": str(input_path), "polygon": polygon, - "padding": padding, "model": BEN2_MODEL_KEY, "device": production_device().lower(), } -def apply_image_extract_settings(operator, context): +def apply_image_mask_settings(operator, context): settings = context.scene.anyimage_settings - operator.remove_background = ( - ai_ready() and settings.image_extract_remove_background + operator.refine_selection = ( + operator.gesture in {"BOX", "LASSO"} + and ai_ready() + and settings.image_edit_refine_selection ) - operator.continuous = settings.image_extract_continuous - operator.padding = settings.image_extract_padding -def draw_image_extract_operator_settings(layout, operator): +def draw_image_mask_operator_settings(layout, operator): + if operator.gesture not in {"BOX", "LASSO"}: + layout.prop(operator, "invert") + return status = ai_status() row = layout.row() row.enabled = status["ready"] - row.prop(operator, "remove_background") - layout.prop(operator, "continuous") - layout.prop(operator, "padding") + row.prop(operator, "refine_selection") + layout.prop(operator, "invert") if not status["ready"]: layout.operator( "anyimage.setup_ai_environment", @@ -278,14 +213,12 @@ def draw_image_extract_operator_settings(layout, operator): ) -def draw_image_extract_settings(context, layout): +def draw_image_mask_settings(context, layout): settings = context.scene.anyimage_settings status = ai_status() row = layout.row() row.enabled = status["ready"] - row.prop(settings, "image_extract_remove_background") - layout.prop(settings, "image_extract_continuous") - layout.prop(settings, "image_extract_padding") + row.prop(settings, "image_edit_refine_selection") if not status["ready"]: layout.operator( "anyimage.setup_ai_environment", @@ -347,3 +280,27 @@ def draw_dashed_line(vertices, color=(1.0, 1.0, 1.0, 0.9)): shader.uniform_float("color", color) batch.draw(shader) gpu.state.blend_set("NONE") + + +def triangle_fan_vertices(vertices): + vertices = tuple(tuple(float(value) for value in vertex) for vertex in vertices) + triangles = [] + for index in range(1, len(vertices) - 1): + triangles.extend((vertices[0], vertices[index], vertices[index + 1])) + return tuple(triangles) + + +def draw_polygon_fill(vertices, color=(1.0, 1.0, 1.0, 0.16)): + import gpu + from gpu_extras.batch import batch_for_shader + + triangles = triangle_fan_vertices(vertices) + if not triangles: + return + shader = gpu.shader.from_builtin("UNIFORM_COLOR") + batch = batch_for_shader(shader, "TRIS", {"pos": triangles}) + gpu.state.blend_set("ALPHA") + shader.bind() + shader.uniform_float("color", color) + batch.draw(shader) + gpu.state.blend_set("NONE") diff --git a/src/anyimage/operators/image_extract/lasso.py b/src/anyimage/operators/image_edit/lasso.py similarity index 51% rename from src/anyimage/operators/image_extract/lasso.py rename to src/anyimage/operators/image_edit/lasso.py index 35c7961..39b1986 100644 --- a/src/anyimage/operators/image_extract/lasso.py +++ b/src/anyimage/operators/image_edit/lasso.py @@ -1,7 +1,7 @@ import bpy -from .common import draw_image_extract_settings -from .region import ExtractImageRegion +from .common import draw_image_mask_settings +from .mask import EditImageMask class ImageLassoTool(bpy.types.WorkSpaceTool): @@ -9,16 +9,16 @@ class ImageLassoTool(bpy.types.WorkSpaceTool): bl_context_mode = "OBJECT" bl_idname = "anyimage.image_lasso" bl_label = "Image Lasso" - bl_description = "Extract a lassoed region from an Image Empty" + bl_description = "Edit an Image Empty with a freehand mask" bl_icon = "ops.sculpt.lasso_mask" - bl_operator = ExtractImageRegion.bl_idname + bl_operator = EditImageMask.bl_idname bl_keymap = ( ( - ExtractImageRegion.bl_idname, + EditImageMask.bl_idname, {"type": "LEFTMOUSE", "value": "PRESS"}, - {"properties": [("selection_shape", "LASSO")]}, + {"properties": [("gesture", "LASSO")]}, ), ) def draw_settings(context, layout, _tool): - draw_image_extract_settings(context, layout) + draw_image_mask_settings(context, layout) diff --git a/src/anyimage/operators/image_edit/line.py b/src/anyimage/operators/image_edit/line.py new file mode 100644 index 0000000..c111d84 --- /dev/null +++ b/src/anyimage/operators/image_edit/line.py @@ -0,0 +1,20 @@ +import bpy + +from .mask import EditImageMask + + +class ImageLineTool(bpy.types.WorkSpaceTool): + bl_space_type = "VIEW_3D" + bl_context_mode = "OBJECT" + bl_idname = "anyimage.image_line" + bl_label = "Image Line" + bl_description = "Edit one side of an Image Empty with a line" + bl_icon = "ops.sculpt.line_mask" + bl_operator = EditImageMask.bl_idname + bl_keymap = ( + ( + EditImageMask.bl_idname, + {"type": "LEFTMOUSE", "value": "PRESS"}, + {"properties": [("gesture", "LINE")]}, + ), + ) diff --git a/src/anyimage/operators/image_edit/mask.py b/src/anyimage/operators/image_edit/mask.py new file mode 100644 index 0000000..1b98a51 --- /dev/null +++ b/src/anyimage/operators/image_edit/mask.py @@ -0,0 +1,583 @@ +import json +import math + +import bpy + +from ...properties import ( + image_edit_refine_selection_property as refine_selection_property, +) +from ...runtime import JobOperatorBase +from ..balloon.operator import ( + deserialize_matrix, + drawing_in_region, + hidden_string_property, + screen_path_to_image_pixels, + serialize_matrix, + simplify_screen_path, +) +from ..common import HIDDEN, cleanup_temporary_input, push_undo +from ..image import ( + is_image_empty, + plane_bounds, + prepare_image_input, + remove_temporary_input, + video_total_frames, +) +from .common import ( + active_view3d_tool_id, + apply_image_mask_settings, + crop_pixels_to_alpha, + draw_dashed_line, + draw_image_mask_operator_settings, + draw_polygon_fill, + image_edit_poll, + image_mask_refine_request, + polygon_bounds, + replace_empty_image_pixels, + resolve_image_edit_click, +) + + +GESTURE_TOOL_IDS = { + "BOX": "anyimage.image_box", + "LASSO": "anyimage.image_lasso", + "LINE": "anyimage.image_line", + "POLYLINE": "anyimage.image_polyline", +} +LINE_SNAP_RADIANS = math.radians(5.0) +POLYLINE_CLOSE_RADIUS = 8.0 + + +def points_inside_polygon(points, polygon): + import numpy as np + + points = np.asarray(points, dtype=np.float64) + polygon = np.asarray(polygon, dtype=np.float64) + inside = np.zeros(len(points), dtype=bool) + previous = polygon[-1] + for current in polygon: + crosses = (current[1] > points[:, 1]) != (previous[1] > points[:, 1]) + edge_x = ( + (previous[0] - current[0]) + * (points[:, 1] - current[1]) + / (previous[1] - current[1] + 1e-300) + + current[0] + ) + inside ^= crosses & (points[:, 0] < edge_x) + previous = current + return inside + + +def polygon_mask(image_size, polygon): + import numpy as np + + width, height = (int(value) for value in image_size) + if width <= 0 or height <= 0: + raise ValueError("The source image has no readable dimensions") + polygon_bounds(polygon) + y, x = np.indices((height, width), dtype=np.float64) + centers = np.column_stack((x.ravel() + 0.5, y.ravel() + 0.5)) + return points_inside_polygon(centers, polygon).reshape((height, width)).astype( + np.float32 + ) + + +def refined_mask_from_pixels(mask_pixels, mask_size, bounds, source_size): + import numpy as np + + mask_width, mask_height = (int(value) for value in mask_size) + source_width, source_height = (int(value) for value in source_size) + local = np.flipud( + np.asarray(mask_pixels, dtype=np.float32).reshape( + (mask_height, mask_width, 4) + ) + )[:, :, 3] + result = np.zeros((source_height, source_width), dtype=np.float32) + left, top, right, bottom = (int(value) for value in bounds) + source_left = max(left, 0) + source_top = max(top, 0) + source_right = min(right, source_width) + source_bottom = min(bottom, source_height) + if source_right <= source_left or source_bottom <= source_top: + return result + result[source_top:source_bottom, source_left:source_right] = local[ + source_top - top : source_bottom - top, + source_left - left : source_right - left, + ] + return result + + +def edit_image_mask_pixels(source_pixels, source_size, selection_mask, invert=False): + """Apply a mask while preserving source RGB, then tightly crop the canvas.""" + import numpy as np + + width, height = (int(value) for value in source_size) + rgba = np.flipud( + np.asarray(source_pixels, dtype=np.float32).reshape((height, width, 4)) + ).copy() + mask = np.asarray(selection_mask, dtype=np.float32) + if mask.shape != (height, width): + raise ValueError("The refined selection does not match the source image") + mask = np.clip(mask, 0.0, 1.0) + if invert: + mask = 1.0 - mask + rgba[:, :, 3] *= mask + return crop_pixels_to_alpha(np.flipud(rgba).ravel(), (0, 0, width, height)) + + +def edit_image_object(source_object, selection_mask, invert=False): + import numpy as np + + image = source_object.data + source_size = tuple(image.size) + source_pixels = np.empty(len(image.pixels), dtype=np.float32) + image.pixels.foreach_get(source_pixels) + pixels, crop_bounds = edit_image_mask_pixels( + source_pixels, source_size, selection_mask, invert + ) + left, top, right, bottom = crop_bounds + return replace_empty_image_pixels( + source_object, + pixels, + (right - left, bottom - top), + crop_bounds, + suffix=".mask", + ) + + +def gesture_property(): + return bpy.props.EnumProperty( + items=( + ("BOX", "Box", "Edit with a rectangular gesture"), + ("LASSO", "Lasso", "Edit with a freehand gesture"), + ("LINE", "Line", "Edit one side of a line"), + ("POLYLINE", "Polyline", "Edit with a clicked polygon"), + ), + default="BOX", + options={"HIDDEN", "SKIP_SAVE"}, + ) + + +def box_screen_path(start, end): + x0, y0 = start + x1, y1 = end + if abs(x1 - x0) < 2.0 or abs(y1 - y0) < 2.0: + return () + return ((x0, y0), (x1, y0), (x1, y1), (x0, y1)) + + +def snapped_line_endpoint(start, end, snap): + if not snap: + return tuple(float(value) for value in end) + delta_x = float(end[0]) - float(start[0]) + delta_y = float(end[1]) - float(start[1]) + length = math.hypot(delta_x, delta_y) + if length <= 0.0: + return tuple(float(value) for value in end) + angle = round(math.atan2(delta_y, delta_x) / LINE_SNAP_RADIANS) + angle *= LINE_SNAP_RADIANS + return ( + float(start[0]) + math.cos(angle) * length, + float(start[1]) + math.sin(angle) * length, + ) + + +def line_half_plane_polygon(start, end, rectangle_size): + """Clip a rectangle to the left side of one directed line.""" + width, height = (float(value) for value in rectangle_size) + if math.dist(start, end) < 2.0: + return () + vertices = [(0.0, 0.0), (width, 0.0), (width, height), (0.0, height)] + + def side(point): + return (end[0] - start[0]) * (point[1] - start[1]) - ( + end[1] - start[1] + ) * (point[0] - start[0]) + + result = [] + previous = vertices[-1] + previous_side = side(previous) + for current in vertices: + current_side = side(current) + if (current_side >= 0.0) != (previous_side >= 0.0): + factor = previous_side / (previous_side - current_side) + result.append( + ( + previous[0] + (current[0] - previous[0]) * factor, + previous[1] + (current[1] - previous[1]) * factor, + ) + ) + if current_side >= 0.0: + result.append(current) + previous = current + previous_side = current_side + return tuple(result) + + +def image_line_half_plane_polygon(line, image_size): + """Keep the image-space half-plane shown by the screen-space preview.""" + if len(line) != 2: + raise ValueError("Drag a longer line across the image plane") + return line_half_plane_polygon(line[1], line[0], image_size) + + +class EditImageMask(bpy.types.Operator): + bl_idname = "anyimage.edit_image_mask" + bl_label = "Edit Image Mask" + bl_description = "Edit and tightly crop an Image Empty with a gesture" + bl_options = {"REGISTER", "UNDO"} + + source_object_name: hidden_string_property() + source_matrix_data: hidden_string_property() + selection_data: hidden_string_property() + refine_selection: refine_selection_property() + invert: bpy.props.BoolProperty( + name="Invert", + description="Keep the opposite side of the refined gesture mask", + default=False, + options={"SKIP_SAVE"}, + ) + gesture: gesture_property() + + _path = None + _cursor = None + _handle = None + _area_pointer = 0 + _region_pointer = 0 + _region_size = (0, 0) + + def draw(self, _context): + draw_image_mask_operator_settings(self.layout, self) + + @classmethod + def poll(cls, context): + return image_edit_poll(context) + + def _screen_polygon(self): + if self.gesture == "BOX": + return box_screen_path(self._path[0], self._cursor) + if self.gesture == "LASSO": + return tuple(simplify_screen_path(self._path)) + if self.gesture == "LINE": + return line_half_plane_polygon( + self._path[0], self._cursor, self._region_size + ) + vertices = list(self._path) + if self._cursor is not None: + vertices.append(self._cursor) + return tuple(vertices) + + def _draw_selection(self): + if not self._path or not drawing_in_region( + self._area_pointer, self._region_pointer + ): + return + polygon = self._screen_polygon() + if len(polygon) >= 3: + draw_polygon_fill( + polygon, + (1.0, 0.25, 0.2, 0.16) + if self.invert + else (1.0, 1.0, 1.0, 0.16), + ) + if self.gesture == "LINE": + draw_dashed_line((self._path[0], self._cursor)) + return + if len(polygon) > 1: + draw_dashed_line((*polygon, polygon[0])) + if ( + self.gesture == "POLYLINE" + and len(self._path) >= 3 + and self._cursor is not None + and math.dist(self._cursor, self._path[0]) <= POLYLINE_CLOSE_RADIUS + ): + self._draw_polyline_close_indicator() + + def _draw_polyline_close_indicator(self): + import gpu + from gpu_extras.batch import batch_for_shader + + center_x, center_y = self._path[0] + vertices = tuple( + ( + center_x + math.cos(index * math.tau / 24.0) * POLYLINE_CLOSE_RADIUS, + center_y + math.sin(index * math.tau / 24.0) * POLYLINE_CLOSE_RADIUS, + ) + for index in range(25) + ) + shader = gpu.shader.from_builtin("UNIFORM_COLOR") + batch = batch_for_shader(shader, "LINE_STRIP", {"pos": vertices}) + gpu.state.blend_set("ALPHA") + gpu.state.line_width_set(2.0) + shader.bind() + shader.uniform_float("color", (1.0, 1.0, 1.0, 0.95)) + batch.draw(shader) + gpu.state.line_width_set(1.0) + gpu.state.blend_set("NONE") + + def _set_status(self, context): + invert = "On" if self.invert else "Off" + if self.gesture == "POLYLINE": + action = "click points; click start or Enter to finish; Backspace removes" + elif self.gesture == "LINE": + action = "drag line; Ctrl snaps 5°; release to edit" + else: + action = "drag and release to edit" + context.workspace.status_text_set( + f"Image {self.gesture.title()}: {action}; F toggles Invert ({invert}); " + "RMB or Esc cancels" + ) + + def _finish(self, context): + if self._handle is not None: + bpy.types.SpaceView3D.draw_handler_remove(self._handle, "WINDOW") + self._handle = None + context.workspace.status_text_set(None) + if context.area is not None: + context.area.tag_redraw() + + def cancel(self, context): + self._finish(context) + + def invoke(self, context, event): + source_object, should_edit = resolve_image_edit_click( + context, event, self.report + ) + if not should_edit: + return {"FINISHED"} if source_object is not None else {"CANCELLED"} + if source_object is None: + return {"CANCELLED"} + if video_total_frames(source_object.data) > 1: + self.report({"WARNING"}, "Image tools only support a still image") + return {"CANCELLED"} + apply_image_mask_settings(self, context) + self.invert = False + self.source_object_name = source_object.name + self.source_matrix_data = serialize_matrix(source_object.matrix_world) + self.selection_data = "" + point = (float(event.mouse_region_x), float(event.mouse_region_y)) + self._path = [point] + self._cursor = point + self._area_pointer = context.area.as_pointer() + self._region_pointer = context.region.as_pointer() + self._region_size = (context.region.width, context.region.height) + self._handle = bpy.types.SpaceView3D.draw_handler_add( + self._draw_selection, (), "WINDOW", "POST_PIXEL" + ) + self._set_status(context) + context.window_manager.modal_handler_add(self) + context.area.tag_redraw() + return {"RUNNING_MODAL"} + + def _complete(self, context): + source_object = bpy.data.objects.get(self.source_object_name) + if not is_image_empty(source_object): + raise RuntimeError("The source Image Empty is no longer available") + matrix = deserialize_matrix(self.source_matrix_data) + if self.gesture == "LINE": + line = screen_path_to_image_pixels( + (self._path[0], self._cursor), + context.region, + context.region_data, + matrix, + plane_bounds(source_object), + tuple(source_object.data.size), + ) + polygon = image_line_half_plane_polygon( + line, tuple(source_object.data.size) + ) + else: + path = self._screen_polygon() + if len(path) < 3: + raise ValueError("Draw a larger image region") + polygon = screen_path_to_image_pixels( + path, + context.region, + context.region_data, + matrix, + plane_bounds(source_object), + tuple(source_object.data.size), + ) + polygon_bounds(polygon) + self.selection_data = json.dumps(polygon, separators=(",", ":")) + self._finish(context) + return self.execute(context) + + def modal(self, context, event): + active_tool = active_view3d_tool_id(context) + if active_tool is not None and active_tool != GESTURE_TOOL_IDS[self.gesture]: + self._finish(context) + return {"CANCELLED"} + if event.type in {"ESC", "RIGHTMOUSE"} and event.value == "PRESS": + self._finish(context) + return {"CANCELLED"} + if event.type == "F" and event.value == "PRESS": + self.invert = not self.invert + self._set_status(context) + context.area.tag_redraw() + return {"RUNNING_MODAL"} + if self.gesture == "POLYLINE": + return self._modal_polyline(context, event) + if event.type == "MOUSEMOVE": + point = (float(event.mouse_region_x), float(event.mouse_region_y)) + if self.gesture == "LASSO": + self._path.append(point) + self._cursor = point + elif self.gesture == "LINE": + self._cursor = snapped_line_endpoint(self._path[0], point, event.ctrl) + else: + self._cursor = point + context.area.tag_redraw() + return {"RUNNING_MODAL"} + if event.type == "LEFTMOUSE" and event.value == "RELEASE": + point = (float(event.mouse_region_x), float(event.mouse_region_y)) + if self.gesture == "LASSO": + self._path.append(point) + self._cursor = point + elif self.gesture == "LINE": + self._cursor = snapped_line_endpoint(self._path[0], point, event.ctrl) + if math.dist(self._path[0], self._cursor) < 2.0: + self._finish(context) + return {"CANCELLED"} + else: + self._cursor = point + try: + return self._complete(context) + except (RuntimeError, TypeError, ValueError) as error: + self.report({"ERROR"}, str(error)) + self._finish(context) + return {"CANCELLED"} + return {"PASS_THROUGH"} + + def _modal_polyline(self, context, event): + if event.type == "BACK_SPACE" and event.value == "PRESS": + if len(self._path) > 1: + self._path.pop() + context.area.tag_redraw() + return {"RUNNING_MODAL"} + if event.type in {"RET", "NUMPAD_ENTER"} and event.value == "PRESS": + return self._confirm_polyline(context) + if event.type == "MOUSEMOVE": + self._cursor = ( + float(event.mouse_region_x), + float(event.mouse_region_y), + ) + context.area.tag_redraw() + return {"RUNNING_MODAL"} + if event.type == "LEFTMOUSE" and event.value == "PRESS": + point = (float(event.mouse_region_x), float(event.mouse_region_y)) + if ( + len(self._path) >= 3 + and math.dist(point, self._path[0]) <= POLYLINE_CLOSE_RADIUS + ): + return self._confirm_polyline(context) + if math.dist(point, self._path[-1]) >= 2.0: + self._path.append(point) + context.area.tag_redraw() + return {"RUNNING_MODAL"} + return {"PASS_THROUGH"} + + def _confirm_polyline(self, context): + self._cursor = self._path[0] + try: + return self._complete(context) + except (RuntimeError, TypeError, ValueError) as error: + self.report({"ERROR"}, str(error)) + self._finish(context) + return {"CANCELLED"} + + def execute(self, context): + source_object = bpy.data.objects.get(self.source_object_name) + if not is_image_empty(source_object): + self.report({"ERROR"}, "The source Image Empty is no longer available") + return {"CANCELLED"} + try: + polygon = json.loads(self.selection_data) + polygon_bounds(polygon) + if video_total_frames(source_object.data) > 1: + raise RuntimeError("Image tools only support a still image") + if self.refine_selection: + source_path, delete_input = prepare_image_input( + source_object.data, context.scene + ) + result = bpy.ops.anyimage.refine_image_mask( + input_path=str(source_path), + delete_input=delete_input, + source_object_name=self.source_object_name, + selection_data=self.selection_data, + invert=self.invert, + ) + if "CANCELLED" in result: + if delete_input: + remove_temporary_input(source_path) + return {"CANCELLED"} + return {"FINISHED"} + edit_image_object( + source_object, + polygon_mask(tuple(source_object.data.size), polygon), + self.invert, + ) + except (json.JSONDecodeError, RuntimeError, TypeError, ValueError) as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + return {"FINISHED"} + + +class RefineImageMask(JobOperatorBase, bpy.types.Operator): + bl_idname = "anyimage.refine_image_mask" + bl_label = "Refine Image Mask" + job_type = "refine-image-mask" + starting_message = "Refining image selection" + + input_path: bpy.props.StringProperty(options=HIDDEN) + delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) + source_object_name: hidden_string_property() + selection_data: hidden_string_property() + invert: bpy.props.BoolProperty(default=False, options=HIDDEN) + + def request(self, _context): + return image_mask_refine_request(self.input_path, self.selection_data) + + def response(self, _context, result): + source_object = bpy.data.objects.get(self.source_object_name) + if not is_image_empty(source_object): + raise RuntimeError("The source Image Empty is no longer available") + import numpy as np + + bounds = tuple(json.loads(result.file("bounds").read_text(encoding="utf-8"))) + refined_image = bpy.data.images.load( + str(result.file("image")), check_existing=False + ) + try: + refined_pixels = np.empty(len(refined_image.pixels), dtype=np.float32) + refined_image.pixels.foreach_get(refined_pixels) + mask = refined_mask_from_pixels( + refined_pixels, + tuple(refined_image.size), + bounds, + tuple(source_object.data.size), + ) + finally: + bpy.data.images.remove(refined_image, do_unlink=True) + edit_image_object(source_object, mask, self.invert) + push_undo("Edit Image Mask") + return "Image selection refined" + + def cleanup(self): + cleanup_temporary_input(self.input_path, self.delete_input) + + +class ActivateImageEditTool(bpy.types.Operator): + bl_idname = "anyimage.activate_image_edit_tool" + bl_label = "Image Edit Tool" + bl_description = "Activate the Image Box tool in the 3D Viewport" + bl_options = {"INTERNAL"} + + @classmethod + def poll(cls, context): + return image_edit_poll(context) + + def execute(self, _context): + return bpy.ops.wm.tool_set_by_id( + name="anyimage.image_box", space_type="VIEW_3D" + ) diff --git a/src/anyimage/operators/image_extract/perspective.py b/src/anyimage/operators/image_edit/perspective.py similarity index 80% rename from src/anyimage/operators/image_extract/perspective.py rename to src/anyimage/operators/image_edit/perspective.py index a620072..9e40b91 100644 --- a/src/anyimage/operators/image_extract/perspective.py +++ b/src/anyimage/operators/image_edit/perspective.py @@ -3,7 +3,6 @@ import math import bpy -from ...runtime import JobOperatorBase from ..balloon.operator import ( deserialize_matrix, drawing_in_region, @@ -11,33 +10,21 @@ from ..balloon.operator import ( screen_path_to_image_pixels, serialize_matrix, ) -from ..common import cleanup_temporary_input, push_undo from ..image import ( image_base_name, is_image_empty, plane_bounds, - prepare_image_input, - remove_temporary_input, video_total_frames, ) from .common import ( - apply_image_extract_settings, - capture_view_direction, - continuous_property, - image_extract_background_request, + active_view3d_tool_id, crop_pixels_to_alpha, draw_dashed_line, - draw_image_extract_operator_settings, - draw_image_extract_settings, - remove_background_property, - resolve_image_extract_click, - image_extract_poll, - image_extract_target, - offset_continuous_result, - padding_property, + draw_polygon_fill, + image_edit_poll, polygon_bounds, replace_empty_image_pixels, - select_image_extract_result, + resolve_image_edit_click, ) @@ -189,42 +176,6 @@ def perspective_quad_overlaps_image(points, image_size): return area > 1e-6 -def triangle_fan_vertices(vertices): - vertices = tuple(tuple(float(value) for value in vertex) for vertex in vertices) - triangles = [] - for index in range(1, len(vertices) - 1): - triangles.extend((vertices[0], vertices[index], vertices[index + 1])) - return tuple(triangles) - - -def draw_polygon_fill(vertices): - import gpu - from gpu_extras.batch import batch_for_shader - - triangles = triangle_fan_vertices(vertices) - if not triangles: - return - shader = gpu.shader.from_builtin("UNIFORM_COLOR") - batch = batch_for_shader(shader, "TRIS", {"pos": triangles}) - gpu.state.blend_set("ALPHA") - shader.bind() - shader.uniform_float("color", (1.0, 1.0, 1.0, 0.16)) - batch.draw(shader) - gpu.state.blend_set("NONE") - - -def active_view3d_tool_id(context): - tools = getattr(getattr(context, "workspace", None), "tools", None) - from_mode = getattr(tools, "from_space_view3d_mode", None) - if not callable(from_mode): - return None - try: - tool = from_mode(getattr(context, "mode", "OBJECT"), create=False) - except (RuntimeError, TypeError): - return None - return getattr(tool, "idname", None) - - def aspect_ratio_from_mouse(initial_ratio, horizontal_delta): ratio = initial_ratio * math.pow( 2.0, float(horizontal_delta) / ASPECT_PIXELS_PER_DOUBLE @@ -506,70 +457,46 @@ def extract_perspective_pixels( source_size, polygon, aspect_ratio, - padding, - placement_polygon=None, ): output_size = perspective_output_size(polygon, aspect_ratio) pixels = warp_perspective_pixels( source_pixels, source_size, polygon, output_size ) - pixels, crop_bounds = crop_pixels_to_alpha( - pixels, (0, 0, *output_size), padding - ) + pixels, crop_bounds = crop_pixels_to_alpha(pixels, (0, 0, *output_size)) cropped_size = ( crop_bounds[2] - crop_bounds[0], crop_bounds[3] - crop_bounds[1], ) placement_bounds = perspective_placement_bounds( - polygon if placement_polygon is None else placement_polygon, - output_size, - crop_bounds, + polygon, output_size, crop_bounds ) return pixels, cropped_size, placement_bounds def create_perspective_result( - context, source_object, - source_matrix, pixels, output_size, placement_bounds, - continuous, - view_direction_data, ): - target, layer = image_extract_target(context, source_object, continuous) - try: - replace_empty_image_pixels( - target, - source_matrix, - pixels, - output_size, - placement_bounds, - suffix=".perspective", - ) - offset_continuous_result(target, view_direction_data, layer) - except Exception: - if target is not source_object: - bpy.data.objects.remove(target, do_unlink=True) - raise - select_image_extract_result(context, source_object, target, continuous) - return target + return replace_empty_image_pixels( + source_object, + pixels, + output_size, + placement_bounds, + suffix=".perspective", + ) -class ExtractImagePerspective(bpy.types.Operator): - bl_idname = "anyimage.extract_image_perspective" - bl_label = "Extract Perspective Image" - bl_description = "Rectify four image corners while staying on the source plane" +class EditImagePerspective(bpy.types.Operator): + bl_idname = "anyimage.edit_image_perspective" + bl_label = "Edit Image Perspective" + bl_description = "Deform four image corners into a rectangular image" bl_options = {"REGISTER", "UNDO"} source_object_name: hidden_string_property() source_matrix_data: hidden_string_property() - view_direction_data: hidden_string_property() selection_data: hidden_string_property() - remove_background: remove_background_property() - continuous: continuous_property() - padding: padding_property() aspect_ratio: bpy.props.FloatProperty( default=1.0, min=MIN_INTERACTIVE_ASPECT, @@ -591,12 +518,9 @@ class ExtractImagePerspective(bpy.types.Operator): _area_pointer = 0 _region_pointer = 0 - def draw(self, _context): - draw_image_extract_operator_settings(self.layout, self) - @classmethod def poll(cls, context): - return image_extract_poll(context) + return image_edit_poll(context) def _draw_preview(self): if self._preview_image is None: @@ -688,10 +612,10 @@ class ExtractImagePerspective(bpy.types.Operator): self._finish(context) def invoke(self, context, event): - source_object, should_extract = resolve_image_extract_click( + source_object, should_edit = resolve_image_edit_click( context, event, self.report ) - if not should_extract: + if not should_edit: return {"FINISHED"} if source_object is not None else {"CANCELLED"} if source_object is None: return {"CANCELLED"} @@ -701,10 +625,8 @@ class ExtractImagePerspective(bpy.types.Operator): "Image Perspective only supports a still image", ) return {"CANCELLED"} - apply_image_extract_settings(self, context) self.source_object_name = source_object.name self.source_matrix_data = serialize_matrix(source_object.matrix_world) - self.view_direction_data = capture_view_direction(context) self.selection_data = "" self._phase = "POINTS" self._preview_image = None @@ -874,26 +796,6 @@ class ExtractImagePerspective(bpy.types.Operator): "Perspective region does not overlap the source image", ) return {"CANCELLED"} - if self.remove_background: - source_path, delete_input = prepare_image_input( - source_object.data, context.scene - ) - result = bpy.ops.anyimage.remove_image_perspective_background( - input_path=str(source_path), - delete_input=delete_input, - source_object_name=self.source_object_name, - source_matrix_data=self.source_matrix_data, - view_direction_data=self.view_direction_data, - selection_data=self.selection_data, - aspect_ratio=self.aspect_ratio, - continuous=self.continuous, - padding=self.padding, - ) - if "CANCELLED" in result: - if delete_input: - remove_temporary_input(source_path) - return {"CANCELLED"} - return {"FINISHED"} import numpy as np source_pixels = self._source_pixels @@ -909,17 +811,12 @@ class ExtractImagePerspective(bpy.types.Operator): tuple(source_object.data.size), polygon, self.aspect_ratio, - self.padding, ) create_perspective_result( - context, source_object, - deserialize_matrix(self.source_matrix_data), pixels, output_size, placement_bounds, - self.continuous, - self.view_direction_data, ) except ( json.JSONDecodeError, @@ -933,88 +830,18 @@ class ExtractImagePerspective(bpy.types.Operator): return {"FINISHED"} -class RemoveImagePerspectiveBackground(JobOperatorBase, bpy.types.Operator): - bl_idname = "anyimage.remove_image_perspective_background" - bl_label = "Remove Perspective Image Background" - job_type = "extract-image-background" - starting_message = "Removing perspective image background" - - input_path: bpy.props.StringProperty(options={"HIDDEN"}) - delete_input: bpy.props.BoolProperty(default=False, options={"HIDDEN"}) - source_object_name: hidden_string_property() - source_matrix_data: hidden_string_property() - view_direction_data: hidden_string_property() - selection_data: hidden_string_property() - aspect_ratio: bpy.props.FloatProperty(options={"HIDDEN"}) - continuous: continuous_property() - padding: padding_property() - - def request(self, _context): - return image_extract_background_request( - self.input_path, self.selection_data, 0 - ) - - def response(self, context, result): - source_object = bpy.data.objects.get(self.source_object_name) - if not is_image_empty(source_object): - raise RuntimeError("The source Image Empty is no longer available") - import numpy as np - - polygon = validate_perspective_quad(json.loads(self.selection_data)) - background_bounds = tuple( - json.loads(result.file("bounds").read_text(encoding="utf-8")) - ) - background_image = bpy.data.images.load( - str(result.file("image")), check_existing=False - ) - try: - background_pixels = np.empty( - len(background_image.pixels), dtype=np.float32 - ) - background_image.pixels.foreach_get(background_pixels) - local_polygon = polygon - np.asarray(background_bounds[:2]) - pixels, output_size, placement_bounds = extract_perspective_pixels( - background_pixels, - tuple(background_image.size), - local_polygon, - self.aspect_ratio, - self.padding, - placement_polygon=polygon, - ) - finally: - bpy.data.images.remove(background_image, do_unlink=True) - create_perspective_result( - context, - source_object, - deserialize_matrix(self.source_matrix_data), - pixels, - output_size, - placement_bounds, - self.continuous, - self.view_direction_data, - ) - push_undo("Extract Perspective Image") - return "Perspective extraction with background removal is ready" - - def cleanup(self): - cleanup_temporary_input(self.input_path, self.delete_input) - - class ImagePerspectiveTool(bpy.types.WorkSpaceTool): bl_space_type = "VIEW_3D" bl_context_mode = "OBJECT" bl_idname = IMAGE_PERSPECTIVE_TOOL_ID bl_label = "Image Perspective" bl_description = "Rectify four image corners on the source Image Empty plane" - bl_icon = "ops.sculpt.polyline_mask" - bl_operator = ExtractImagePerspective.bl_idname + bl_icon = "ops.sculpt.line_project" + bl_operator = EditImagePerspective.bl_idname bl_keymap = ( ( - ExtractImagePerspective.bl_idname, + EditImagePerspective.bl_idname, {"type": "LEFTMOUSE", "value": "PRESS"}, None, ), ) - - def draw_settings(context, layout, _tool): - draw_image_extract_settings(context, layout) diff --git a/src/anyimage/operators/image_edit/polyline.py b/src/anyimage/operators/image_edit/polyline.py new file mode 100644 index 0000000..01bafdc --- /dev/null +++ b/src/anyimage/operators/image_edit/polyline.py @@ -0,0 +1,20 @@ +import bpy + +from .mask import EditImageMask + + +class ImagePolylineTool(bpy.types.WorkSpaceTool): + bl_space_type = "VIEW_3D" + bl_context_mode = "OBJECT" + bl_idname = "anyimage.image_polyline" + bl_label = "Image Polyline" + bl_description = "Edit an Image Empty with a clicked polygon mask" + bl_icon = "ops.sculpt.polyline_mask" + bl_operator = EditImageMask.bl_idname + bl_keymap = ( + ( + EditImageMask.bl_idname, + {"type": "LEFTMOUSE", "value": "PRESS"}, + {"properties": [("gesture", "POLYLINE")]}, + ), + ) diff --git a/src/anyimage/operators/image_extract/__init__.py b/src/anyimage/operators/image_extract/__init__.py deleted file mode 100644 index a555df5..0000000 --- a/src/anyimage/operators/image_extract/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -from .box import ImageBoxTool -from .lasso import ImageLassoTool -from .perspective import ( - ExtractImagePerspective, - ImagePerspectiveTool, - RemoveImagePerspectiveBackground, -) -from .region import ( - ActivateImageExtractTool, - ExtractImageRegion, - RemoveImageRegionBackground, -) - - -CLASSES = ( - ExtractImageRegion, - RemoveImageRegionBackground, - ActivateImageExtractTool, - ExtractImagePerspective, - RemoveImagePerspectiveBackground, -) - - -TOOLS = (ImageBoxTool, ImageLassoTool, ImagePerspectiveTool) - - -__all__ = ( - "CLASSES", - "TOOLS", - "ImageBoxTool", - "ImageLassoTool", - "ImagePerspectiveTool", -) diff --git a/src/anyimage/operators/image_extract/region.py b/src/anyimage/operators/image_extract/region.py deleted file mode 100644 index 3e1ff89..0000000 --- a/src/anyimage/operators/image_extract/region.py +++ /dev/null @@ -1,378 +0,0 @@ -import json - -import bpy - -from ...runtime import JobOperatorBase -from ..balloon.operator import ( - deserialize_matrix, - drawing_in_region, - hidden_string_property, - screen_path_to_image_pixels, - serialize_matrix, - simplify_screen_path, -) -from ..common import HIDDEN, cleanup_temporary_input, push_undo -from ..image import ( - is_image_empty, - plane_bounds, - prepare_image_input, - remove_temporary_input, - video_total_frames, -) -from .common import ( - apply_image_extract_settings, - capture_view_direction, - continuous_property, - image_extract_background_request, - crop_pixels_to_alpha, - draw_dashed_line, - draw_image_extract_operator_settings, - remove_background_property, - resolve_image_extract_click, - image_extract_poll, - image_extract_target, - offset_continuous_result, - padding_property, - polygon_bounds, - replace_empty_image_from_path, - replace_empty_image_pixels, - select_image_extract_result, -) - - -def points_inside_polygon(points, polygon): - import numpy as np - - points = np.asarray(points, dtype=np.float64) - polygon = np.asarray(polygon, dtype=np.float64) - inside = np.zeros(len(points), dtype=bool) - previous = polygon[-1] - for current in polygon: - crosses = (current[1] > points[:, 1]) != (previous[1] > points[:, 1]) - edge_x = ( - (previous[0] - current[0]) - * (points[:, 1] - current[1]) - / (previous[1] - current[1] + 1e-300) - + current[0] - ) - inside ^= crosses & (points[:, 0] < edge_x) - previous = current - return inside - - -def extract_region_pixels(source_pixels, source_size, polygon): - """Return Blender-order RGBA pixels and the unbounded selection rectangle.""" - import numpy as np - - source_width, source_height = (int(value) for value in source_size) - if source_width <= 0 or source_height <= 0: - raise ValueError("The source image has no readable dimensions") - source = np.asarray(source_pixels, dtype=np.float32) - if source.size != source_width * source_height * 4: - raise ValueError("The source image pixels do not match its dimensions") - source = np.flipud(source.reshape((source_height, source_width, 4))) - left, top, right, bottom = polygon_bounds(polygon) - width = right - left - height = bottom - top - result = np.zeros((height, width, 4), dtype=np.float32) - - source_left = max(left, 0) - source_top = max(top, 0) - source_right = min(right, source_width) - source_bottom = min(bottom, source_height) - if source_right <= source_left or source_bottom <= source_top: - raise ValueError("The selection does not overlap the source image") - result[ - source_top - top : source_bottom - top, - source_left - left : source_right - left, - ] = source[source_top:source_bottom, source_left:source_right] - - y, x = np.indices((height, width), dtype=np.float64) - centers = np.column_stack( - ((x.ravel() + left + 0.5), (y.ravel() + top + 0.5)) - ) - selected = points_inside_polygon(centers, polygon).reshape((height, width)) - result[~selected] = 0.0 - result[result[:, :, 3] <= 0.0] = 0.0 - return np.flipud(result).ravel(), (left, top, right, bottom) - - -def selection_shape_property(): - return bpy.props.EnumProperty( - items=( - ("LASSO", "Lasso", "Draw a freehand image extraction region"), - ("BOX", "Box", "Draw a rectangular image extraction region"), - ), - default="BOX", - options={"HIDDEN", "SKIP_SAVE"}, - ) - - -def box_screen_path(path): - if len(path) < 2: - return () - x0, y0 = path[0] - x1, y1 = path[-1] - if abs(x1 - x0) < 2.0 or abs(y1 - y0) < 2.0: - return () - return ((x0, y0), (x1, y0), (x1, y1), (x0, y1)) - - -class ExtractImageRegion(bpy.types.Operator): - bl_idname = "anyimage.extract_image_region" - bl_label = "Extract Image Region" - bl_description = "Extract a boxed or lassoed region from an Image Empty" - bl_options = {"REGISTER", "UNDO"} - - source_object_name: hidden_string_property() - source_matrix_data: hidden_string_property() - view_direction_data: hidden_string_property() - selection_data: hidden_string_property() - remove_background: remove_background_property() - continuous: continuous_property() - padding: padding_property() - selection_shape: selection_shape_property() - - _path = None - _handle = None - _area_pointer = 0 - _region_pointer = 0 - - def draw(self, _context): - draw_image_extract_operator_settings(self.layout, self) - - @classmethod - def poll(cls, context): - return image_extract_poll(context) - - def _draw_selection(self): - if not self._path or not drawing_in_region( - self._area_pointer, self._region_pointer - ): - return - vertices = ( - list(box_screen_path(self._path)) - if self.selection_shape == "BOX" - else list(self._path) - ) - if len(vertices) > 2: - vertices.append(vertices[0]) - draw_dashed_line(vertices) - - def _finish(self, context): - if self._handle is not None: - bpy.types.SpaceView3D.draw_handler_remove(self._handle, "WINDOW") - self._handle = None - context.workspace.status_text_set(None) - if context.area is not None: - context.area.tag_redraw() - - def cancel(self, context): - self._finish(context) - - def invoke(self, context, event): - source_object, should_extract = resolve_image_extract_click( - context, event, self.report - ) - if not should_extract: - return {"FINISHED"} if source_object is not None else {"CANCELLED"} - if source_object is None: - return {"CANCELLED"} - if video_total_frames(source_object.data) > 1: - self.report({"WARNING"}, "Image tools only support a still image") - return {"CANCELLED"} - apply_image_extract_settings(self, context) - self.source_object_name = source_object.name - self.source_matrix_data = serialize_matrix(source_object.matrix_world) - self.view_direction_data = capture_view_direction(context) - self.selection_data = "" - self._path = [(event.mouse_region_x, event.mouse_region_y)] - self._area_pointer = context.area.as_pointer() - self._region_pointer = context.region.as_pointer() - self._handle = bpy.types.SpaceView3D.draw_handler_add( - self._draw_selection, (), "WINDOW", "POST_PIXEL" - ) - context.workspace.status_text_set( - f"Image {self.selection_shape.title()}: release LMB to extract, " - "RMB or Esc to cancel" - ) - context.window_manager.modal_handler_add(self) - context.area.tag_redraw() - return {"RUNNING_MODAL"} - - def modal(self, context, event): - if event.type in {"ESC", "RIGHTMOUSE"} and event.value == "PRESS": - self._finish(context) - return {"CANCELLED"} - if event.type == "MOUSEMOVE": - self._path.append((event.mouse_region_x, event.mouse_region_y)) - context.area.tag_redraw() - return {"RUNNING_MODAL"} - if event.type == "LEFTMOUSE" and event.value == "RELEASE": - self._path.append((event.mouse_region_x, event.mouse_region_y)) - path = ( - list(box_screen_path(self._path)) - if self.selection_shape == "BOX" - else simplify_screen_path(self._path) - ) - if len(path) < 3: - self.report({"WARNING"}, "Draw a larger image region") - self._finish(context) - return {"CANCELLED"} - source_object = bpy.data.objects.get(self.source_object_name) - if not is_image_empty(source_object): - self.report({"ERROR"}, "The source Image Empty is no longer available") - self._finish(context) - return {"CANCELLED"} - try: - polygon = screen_path_to_image_pixels( - path, - context.region, - context.region_data, - deserialize_matrix(self.source_matrix_data), - plane_bounds(source_object), - tuple(source_object.data.size), - ) - self.selection_data = json.dumps(polygon, separators=(",", ":")) - except (RuntimeError, TypeError, ValueError) as error: - self.report({"ERROR"}, str(error)) - self._finish(context) - return {"CANCELLED"} - self._finish(context) - return self.execute(context) - return {"PASS_THROUGH"} - - def execute(self, context): - source_object = bpy.data.objects.get(self.source_object_name) - if not is_image_empty(source_object): - self.report({"ERROR"}, "The source Image Empty is no longer available") - return {"CANCELLED"} - try: - polygon = json.loads(self.selection_data) - polygon_bounds(polygon) - if video_total_frames(source_object.data) > 1: - raise RuntimeError("Image tools only support a still image") - if self.remove_background: - source_path, delete_input = prepare_image_input( - source_object.data, context.scene - ) - result = bpy.ops.anyimage.remove_image_region_background( - input_path=str(source_path), - delete_input=delete_input, - source_object_name=self.source_object_name, - source_matrix_data=self.source_matrix_data, - view_direction_data=self.view_direction_data, - selection_data=self.selection_data, - continuous=self.continuous, - padding=self.padding, - ) - if "CANCELLED" in result: - if delete_input: - remove_temporary_input(source_path) - return {"CANCELLED"} - return {"FINISHED"} - import numpy as np - - image = source_object.data - source_pixels = np.empty(len(image.pixels), dtype=np.float32) - image.pixels.foreach_get(source_pixels) - pixels, crop_bounds = extract_region_pixels( - source_pixels, tuple(image.size), polygon - ) - pixels, crop_bounds = crop_pixels_to_alpha( - pixels, crop_bounds, self.padding - ) - target, layer = image_extract_target( - context, source_object, self.continuous - ) - try: - left, top, right, bottom = crop_bounds - replace_empty_image_pixels( - target, - deserialize_matrix(self.source_matrix_data), - pixels, - (right - left, bottom - top), - crop_bounds, - ) - offset_continuous_result( - target, self.view_direction_data, layer - ) - except Exception: - if target is not source_object: - bpy.data.objects.remove(target, do_unlink=True) - raise - select_image_extract_result( - context, source_object, target, self.continuous - ) - except (json.JSONDecodeError, RuntimeError, TypeError, ValueError) as error: - self.report({"ERROR"}, str(error)) - return {"CANCELLED"} - return {"FINISHED"} - - -class RemoveImageRegionBackground(JobOperatorBase, bpy.types.Operator): - bl_idname = "anyimage.remove_image_region_background" - bl_label = "Remove Image Region Background" - job_type = "extract-image-background" - starting_message = "Removing image region background" - - input_path: bpy.props.StringProperty(options=HIDDEN) - delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) - source_object_name: hidden_string_property() - source_matrix_data: hidden_string_property() - view_direction_data: hidden_string_property() - selection_data: hidden_string_property() - continuous: continuous_property() - padding: padding_property() - - def request(self, _context): - return image_extract_background_request( - self.input_path, self.selection_data, self.padding - ) - - def response(self, context, result): - source_object = bpy.data.objects.get(self.source_object_name) - if not is_image_empty(source_object): - raise RuntimeError("The source Image Empty is no longer available") - crop_bounds = tuple( - json.loads(result.file("bounds").read_text(encoding="utf-8")) - ) - target, layer = image_extract_target( - context, source_object, self.continuous - ) - try: - replace_empty_image_from_path( - target, - deserialize_matrix(self.source_matrix_data), - result.file("image"), - crop_bounds, - ) - offset_continuous_result(target, self.view_direction_data, layer) - except Exception: - if target is not source_object: - bpy.data.objects.remove(target, do_unlink=True) - raise - select_image_extract_result( - context, source_object, target, self.continuous - ) - push_undo("Extract Image Region") - return "Image extraction with background removal is ready" - - def cleanup(self): - cleanup_temporary_input(self.input_path, self.delete_input) - - -class ActivateImageExtractTool(bpy.types.Operator): - bl_idname = "anyimage.activate_image_extract_tool" - bl_label = "Image Extract Tool" - bl_description = "Activate the Image Box tool in the 3D Viewport" - bl_options = {"INTERNAL"} - - @classmethod - def poll(cls, context): - return image_extract_poll(context) - - def execute(self, _context): - return bpy.ops.wm.tool_set_by_id( - name="anyimage.image_box", space_type="VIEW_3D" - ) diff --git a/src/anyimage/panel.py b/src/anyimage/panel.py index 1cbb3cb..1d3ba2c 100644 --- a/src/anyimage/panel.py +++ b/src/anyimage/panel.py @@ -1,7 +1,11 @@ import bpy from .runtime import runtime -from .properties import BEN2_MODEL_KEY, model_ready +from .properties import ( + BEN2_MODEL_KEY, + REALESRGAN_MODEL_KEY, + model_ready, +) class ServerPanel(bpy.types.Panel): bl_label = "Job Server" @@ -221,3 +225,45 @@ class Ben2RemovalDebugPanel(bpy.types.Panel): text="Remove Background", icon="IMAGE_ALPHA", ) + + +class RealesrganGeneralX4v3DebugPanel(bpy.types.Panel): + bl_label = "Real-ESRGAN General x4v3" + bl_idname = "ANYIMAGE_PT_debug_realesrgan_general_x4v3" + bl_parent_id = DebugPanel.bl_idname + bl_space_type = "VIEW_3D" + bl_region_type = "UI" + bl_category = "AnyImage" + bl_options = {"DEFAULT_CLOSED"} + bl_order = 3 + + def draw(self, context): + layout = self.layout + settings = context.scene.anyimage_settings + layout.prop(settings, "input_path") + input_path = bpy.path.abspath(settings.input_path) + if input_path.lower().endswith( + ( + ".mp4", + ".mov", + ".avi", + ".mkv", + ".webm", + ".m4v", + ".mts", + ".mpg", + ".mpeg", + ) + ): + layout.prop(settings, "video_frames") + run_row = layout.row() + run_row.enabled = ( + not runtime.server_busy() + and model_ready(REALESRGAN_MODEL_KEY) + ) + run_row.operator( + "anyimage.debug_realesrgan_general_x4v3", + text="Run General x4v3 API", + icon="EXPORT", + ) + diff --git a/src/anyimage/preferences.py b/src/anyimage/preferences.py index 31dfc66..476f8e7 100644 --- a/src/anyimage/preferences.py +++ b/src/anyimage/preferences.py @@ -2,7 +2,6 @@ import bpy from .runtime import DEFAULT_STORAGE_ROOT, runtime from .properties import ( - BEN2_MODEL_KEY, ai_setup_label, ai_status, device_enum_items, @@ -61,18 +60,12 @@ class AnyImagePreferences(bpy.types.AddonPreferences): layout.separator(type="LINE") layout.label(text="Models") models = model_catalog() - for family in ("moge2", "da3"): + for family in ("moge2", "da3", "ben2", "upscale"): for model in models: if model.get("family") == family and ( family != "da3" or model.get("enabled") ): self._draw_model(layout, model) - ben2 = next( - (model for model in models if model.get("key") == BEN2_MODEL_KEY), - None, - ) - if ben2: - self._draw_model(layout, ben2) layout.separator() layout.operator( "anyimage.open_server_log", diff --git a/src/anyimage/properties.py b/src/anyimage/properties.py index fcad599..82bcb27 100644 --- a/src/anyimage/properties.py +++ b/src/anyimage/properties.py @@ -11,6 +11,7 @@ from . import model_catalog as shared_model_catalog BEN2_MODEL_KEY = "BEN2_BASE" +REALESRGAN_MODEL_KEY = "REALESRGAN_GENERAL_X4V3" DEFAULT_DA3_MODEL_KEY = "BASE" DEFAULT_MOGE2_MODEL_KEY = "MOGE2_VITS_NORMAL" REQUIRED_MODEL_KEYS = (DEFAULT_MOGE2_MODEL_KEY, BEN2_MODEL_KEY) @@ -280,37 +281,14 @@ def balloon_contour_analysis_property(): ) -def image_extract_remove_background_property(): +def image_edit_refine_selection_property(): return bpy.props.BoolProperty( - name="Remove Background", - description=( - "Use BEN2 to remove the background inside the selected image region" - ), + name="Refine Selection", + description="Use BEN2 to fit the gesture mask to foreground content", default=False, ) -def image_extract_continuous_property(): - return bpy.props.BoolProperty( - name="Continuous", - description=( - "Keep the source Image Empty active and continuously create extracted copies" - ), - default=False, - ) - - -def image_extract_padding_property(): - return bpy.props.IntProperty( - name="Padding", - description="Transparent padding in pixels around the visible Alpha bounds", - default=0, - min=0, - soft_max=256, - subtype="PIXEL", - ) - - def balloon_depth_analysis_property(): return bpy.props.BoolProperty( name="Depth", @@ -337,9 +315,7 @@ def balloon_radial_zoom_property(scale): class AnyImageSettings(bpy.types.PropertyGroup): - image_extract_remove_background: image_extract_remove_background_property() - image_extract_continuous: image_extract_continuous_property() - image_extract_padding: image_extract_padding_property() + image_edit_refine_selection: image_edit_refine_selection_property() balloon_mode: balloon_mode_property() balloon_thickness: balloon_thickness_property() balloon_resolution: balloon_resolution_property() diff --git a/src/anyimage/server/__init__.py b/src/anyimage/server/__init__.py index 268e044..3b933f5 100644 --- a/src/anyimage/server/__init__.py +++ b/src/anyimage/server/__init__.py @@ -9,8 +9,9 @@ from .jobs.balloon import run as _analyze_balloon from .jobs.debug_ben2 import run as _debug_ben2 from .jobs.debug_da3 import run as _debug_da3 from .jobs.debug_moge2 import run as _debug_moge2 +from .jobs.debug_realesrgan import run as _debug_realesrgan from .jobs.frames import run as _extract_frames -from .jobs.image_extract import run as _extract_image_background +from .jobs.image_mask import run as _refine_image_mask from .model_manager import ModelManager, ModelSessionCache @@ -39,9 +40,9 @@ def analyze_balloon(context, parameters): return _analyze_balloon(context, parameters) -@server.job("extract-image-background") -def extract_image_background(context, parameters): - return _extract_image_background(context, parameters) +@server.job("refine-image-mask") +def refine_image_mask(context, parameters): + return _refine_image_mask(context, parameters) @server.job("remove-background") @@ -79,4 +80,11 @@ def debug_ben2(context, parameters): return _debug_ben2(context, parameters) +@server.job("debug-realesrgan-general-x4v3") +def debug_realesrgan(context, parameters): + return _debug_realesrgan(context, parameters) + + + + __all__ = ("ModelManager", "ModelSessionCache", "server") diff --git a/src/anyimage/server/jobs/debug_realesrgan.py b/src/anyimage/server/jobs/debug_realesrgan.py new file mode 100644 index 0000000..f7111c6 --- /dev/null +++ b/src/anyimage/server/jobs/debug_realesrgan.py @@ -0,0 +1,26 @@ +from ..media import copy_frames +from ..models.realesrgan import infer_one +from .common import input_frames + + +def run(context, parameters): + context.progress(0.02, "Preparing Real-ESRGAN API input") + context.directory.mkdir(parents=True, exist_ok=True) + with input_frames(context, parameters) as images: + for index, image in enumerate(images): + context.check_cancelled() + prediction = infer_one( + context.resource("model_manager"), + parameters, + image, + ) + suffix = "" if len(images) == 1 else f".{index + 1:04d}" + prediction.save( + context.directory / f"prediction{suffix}.png", + format="PNG", + compress_level=1, + ) + copy_frames(images, context.directory) + context.check_cancelled() + context.progress(1.0, "Saved raw Real-ESRGAN API prediction") + return {"directory": "."} diff --git a/src/anyimage/server/jobs/image_extract.py b/src/anyimage/server/jobs/image_mask.py similarity index 84% rename from src/anyimage/server/jobs/image_extract.py rename to src/anyimage/server/jobs/image_mask.py index 466366d..946c5e0 100644 --- a/src/anyimage/server/jobs/image_extract.py +++ b/src/anyimage/server/jobs/image_mask.py @@ -67,13 +67,11 @@ def apply_selection_mask(foreground, selection_mask): return Image.fromarray(pixels, mode="RGBA") -def crop_to_alpha(image, bounds, padding=0): - """Crop an RGBA image to nonzero Alpha and expand it by pixel padding.""" +def crop_to_alpha(image, bounds): + """Crop an RGBA mask image to its nonzero Alpha bounds.""" import numpy as np from PIL import Image - if padding < 0: - raise ValueError("Padding must not be negative") pixels = np.asarray(image.convert("RGBA"), dtype=np.uint8) visible_y, visible_x = np.nonzero(pixels[:, :, 3] > 0) if not len(visible_x): @@ -85,23 +83,23 @@ def crop_to_alpha(image, bounds, padding=0): output = Image.new( "RGBA", ( - content_right - content_left + padding * 2, - content_bottom - content_top + padding * 2, + content_right - content_left, + content_bottom - content_top, ), (0, 0, 0, 0), ) - output.paste(image, (padding - content_left, padding - content_top)) + output.paste(image, (-content_left, -content_top)) left, top, _right, _bottom = bounds return output, ( - left + content_left - padding, - top + content_top - padding, - left + content_right + padding, - top + content_bottom + padding, + left + content_left, + top + content_top, + left + content_right, + top + content_bottom, ) def run(context, parameters): - """Remove the background from an image selection and return source bounds.""" + """Refine a gesture mask with BEN2 and return its source-image bounds.""" started = time.perf_counter() context.progress(0.05, "Preparing image selection") context.check_cancelled() @@ -124,15 +122,13 @@ def run(context, parameters): context.check_cancelled() context.progress(0.9, "Applying selection mask") output, bounds = crop_to_alpha( - apply_selection_mask(foreground, selection_mask), - bounds, - int(parameters.get("padding", 0)), + apply_selection_mask(foreground, selection_mask), bounds ) output_path = output_directory / "image.png" output.save(output_path, format="PNG", compress_level=1) bounds_path = output_directory / "bounds.json" bounds_path.write_text(json.dumps(bounds, separators=(",", ":")), encoding="utf-8") elapsed_ms = (time.perf_counter() - started) * 1000.0 - message = f"Image background removal finished in {elapsed_ms:.0f} ms" + message = f"Image selection refined in {elapsed_ms:.0f} ms" context.progress(1.0, message) return {"image": output_path.name, "bounds": bounds_path.name} diff --git a/src/anyimage/server/model_catalog.py b/src/anyimage/server/model_catalog.py index 4e5260b..f63f039 100644 --- a/src/anyimage/server/model_catalog.py +++ b/src/anyimage/server/model_catalog.py @@ -89,6 +89,30 @@ BEN2_MODEL = ModelSpec( ) +REALESRGAN_MODEL = ModelSpec( + key="REALESRGAN_GENERAL_X4V3", + label="Real-ESRGAN General x4v3", + directory_name="realesr-general-x4v3", + repository="jonathanst29/tinier-upscale-models", + license_name="BSD 3-Clause", + description="Fast general-purpose 4x image upscaler using ONNX Runtime", + supports_gs=False, + download_patterns=("realesr-general-x4v3.onnx",), + ready_patterns=("realesr-general-x4v3.onnx",), + r2_directory="realesr-general-x4v3", + r2_files=(( + "realesr-general-x4v3.onnx", + 4866396, + "924ebad6532777303582d4ce7811849b88869231a3ae7093e0f21200249df8d5", + ),), +) + + +UPSCALE_MODELS = { + REALESRGAN_MODEL.key: REALESRGAN_MODEL, +} + + MOGE2_MODELS = { "MOGE2_VITS_NORMAL": ModelSpec( key="MOGE2_VITS_NORMAL", @@ -146,6 +170,7 @@ MOGE2_MODEL = MOGE2_MODELS["MOGE2_VITS_NORMAL"] DOWNLOADABLE_MODELS = { **MODELS, BEN2_MODEL.key: BEN2_MODEL, + **UPSCALE_MODELS, **MOGE2_MODELS, } MODEL_NUMBERS = { @@ -181,8 +206,14 @@ def model_record(model, ready=False): "license": model.license_name, "supports_gs": model.supports_gs, "enabled": model.enabled, - "family": "moge2" if model.key.startswith("MOGE2_") else ( - "ben2" if model.key == BEN2_MODEL.key else "da3" + "family": ( + "moge2" + if model.key.startswith("MOGE2_") + else "ben2" + if model.key == BEN2_MODEL.key + else "upscale" + if model.key in UPSCALE_MODELS + else "da3" ), "ready": bool(ready), } diff --git a/src/anyimage/server/model_manager.py b/src/anyimage/server/model_manager.py index 53aa533..20a081a 100644 --- a/src/anyimage/server/model_manager.py +++ b/src/anyimage/server/model_manager.py @@ -17,6 +17,8 @@ class ModelSessionCache: self.da3_key = None self.moge_model = None self.moge_key = None + self.realesrgan_model = None + self.realesrgan_key = None def get_ben2(self, model_dir, device): requested_device = getattr(device, "type", device) @@ -68,6 +70,18 @@ class ModelSessionCache: self.moge_key = key return model + def get_realesrgan(self, model_dir, requested_device): + from .models.onnx_realesrgan import create_session + + key = (str(Path(model_dir).resolve()), requested_device, "onnx") + if self.realesrgan_model is not None and self.realesrgan_key == key: + return self.realesrgan_model + self.release_realesrgan() + model = create_session(model_dir, requested_device) + self.realesrgan_model = model + self.realesrgan_key = key + return model + def release_ben2(self): self.ben2_model = None self.ben2_key = None @@ -83,13 +97,20 @@ class ModelSessionCache: self.moge_key = None self._release_runtime_cache() + def release_realesrgan(self): + self.realesrgan_model = None + self.realesrgan_key = None + self._release_runtime_cache() + def close(self): self.ben2_model = None self.da3_model = None self.moge_model = None + self.realesrgan_model = None self.ben2_key = None self.da3_key = None self.moge_key = None + self.realesrgan_key = None self._release_runtime_cache() def snapshot(self): @@ -97,6 +118,11 @@ class ModelSessionCache: "ben2": Path(self.ben2_key[0]).name if self.ben2_key else "", "da3": Path(self.da3_key[0]).name if self.da3_key else "", "moge2": Path(self.moge_key[0]).name if self.moge_key else "", + "realesrgan": ( + Path(self.realesrgan_key[0]).name + if self.realesrgan_key + else "" + ), } def _release_runtime_cache(self): @@ -143,6 +169,9 @@ class ModelManager: def get_moge2(self, model_dir, device): return self.sessions.get_moge2(model_dir, device) + def get_realesrgan(self, model_dir, device): + return self.sessions.get_realesrgan(model_dir, device) + def snapshot(self): return {"loaded": self.sessions.snapshot()} diff --git a/src/anyimage/server/models/__init__.py b/src/anyimage/server/models/__init__.py index 0cc87e6..3309119 100644 --- a/src/anyimage/server/models/__init__.py +++ b/src/anyimage/server/models/__init__.py @@ -1,8 +1,14 @@ -from . import ben2, da3, moge2 +from . import ben2, da3, moge2, realesrgan def geometry_adapter(model_key): return moge2 if model_key.startswith("MOGE2_") else da3 -__all__ = ("ben2", "da3", "moge2", "geometry_adapter") +__all__ = ( + "ben2", + "da3", + "moge2", + "realesrgan", + "geometry_adapter", +) diff --git a/src/anyimage/server/models/ben2.py b/src/anyimage/server/models/ben2.py index 8449cf5..52a409a 100644 --- a/src/anyimage/server/models/ben2.py +++ b/src/anyimage/server/models/ben2.py @@ -32,15 +32,12 @@ def infer_foreground( check_cancelled(cancel_check) with Image.open(input_path) as opened_image: source_rgba = opened_image.convert("RGBA") - if preserve_input_alpha: - image = Image.new("RGB", source_rgba.size, (0, 0, 0)) - image.paste( - source_rgba.convert("RGB"), - (0, 0), - source_rgba.getchannel("A"), - ) - else: - image = source_rgba.convert("RGB") + image = Image.new("RGB", source_rgba.size, (0, 0, 0)) + image.paste( + source_rgba.convert("RGB"), + (0, 0), + source_rgba.getchannel("A"), + ) inference_started = time.perf_counter() foreground = infer(model, image) inference_ms = (time.perf_counter() - inference_started) * 1000.0 diff --git a/src/anyimage/server/models/onnx_realesrgan.py b/src/anyimage/server/models/onnx_realesrgan.py new file mode 100644 index 0000000..b10cce3 --- /dev/null +++ b/src/anyimage/server/models/onnx_realesrgan.py @@ -0,0 +1,51 @@ +"""Real-ESRGAN General x4v3 ONNX inference.""" + +from pathlib import Path + +try: + from .onnx_runtime import create_session as create_runtime_session +except ImportError: + from onnx_runtime import create_session as create_runtime_session + + +ONNX_FILENAME = "realesr-general-x4v3.onnx" +SCALE = 4 + + +def create_session(model_directory, requested_device="auto"): + path = Path(model_directory) / ONNX_FILENAME + return create_runtime_session(path, requested_device) + + +def preprocess(image): + import numpy as np + + image_array = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0 + return np.ascontiguousarray(image_array.transpose(2, 0, 1)[None]) + + +def infer(session, image): + import numpy as np + from PIL import Image + + input_tensor = preprocess(image) + input_name = session.get_inputs()[0].name + output_name = session.get_outputs()[0].name + prediction = np.asarray( + session.run([output_name], {input_name: input_tensor})[0], + dtype=np.float32, + ) + expected_shape = ( + 1, + 3, + input_tensor.shape[2] * SCALE, + input_tensor.shape[3] * SCALE, + ) + if prediction.shape != expected_shape: + raise RuntimeError( + "Real-ESRGAN General x4v3 returned an unexpected output shape: " + f"{prediction.shape}" + ) + pixels = np.clip(prediction[0].transpose(1, 2, 0), 0.0, 1.0) + pixels = np.rint(pixels * 255.0).astype(np.uint8) + return Image.fromarray(pixels, mode="RGB") diff --git a/src/anyimage/server/models/realesrgan.py b/src/anyimage/server/models/realesrgan.py new file mode 100644 index 0000000..34feb53 --- /dev/null +++ b/src/anyimage/server/models/realesrgan.py @@ -0,0 +1,12 @@ +def infer_one(model_manager, parameters, image_path): + from PIL import Image + + from .onnx_realesrgan import infer + + model_directory = model_manager.directory(parameters["model"]) + model = model_manager.get_realesrgan( + model_directory, + parameters["device"], + ) + with Image.open(image_path) as opened: + return infer(model, opened) diff --git a/sync_models.py b/sync_models.py index d9d1aff..1a6b3e9 100644 --- a/sync_models.py +++ b/sync_models.py @@ -22,11 +22,11 @@ class ModelFile: size: int sha256: str sources: tuple[str, ...] + filename: str = "model.onnx" @property def destination(self): - filename = "model_fp16.onnx" if self.folder.endswith("/onnx") else "model.onnx" - return MODELS_DIR / self.folder / filename + return MODELS_DIR / self.folder / self.filename def hf_url(repository, filename): @@ -58,6 +58,23 @@ MODEL_FILES = ( hf_url("onnx-community/BEN2-ONNX", "onnx/model_fp16.onnx"), r2_url("BEN2-ONNX", "onnx/model_fp16.onnx"), ), + "model_fp16.onnx", + ), + ModelFile( + "realesr-general-x4v3", + 4866396, + "924ebad6532777303582d4ce7811849b88869231a3ae7093e0f21200249df8d5", + ( + hf_url( + "jonathanst29/tinier-upscale-models", + "realesr-general-x4v3.onnx", + ), + r2_url( + "realesr-general-x4v3", + "realesr-general-x4v3.onnx", + ), + ), + "realesr-general-x4v3.onnx", ), ModelFile( "moge-2-vits-normal-onnx", diff --git a/tests/test_addon.py b/tests/test_addon.py index f8790a3..11150ba 100644 --- a/tests/test_addon.py +++ b/tests/test_addon.py @@ -87,6 +87,7 @@ class AddonTest(unittest.TestCase): cls.depth_operator = importlib.import_module( "anyimage.operators.depth_plane.operator" ) + cls.debug = importlib.import_module("anyimage.operators.debug") cls.depth_build = importlib.import_module( "anyimage.operators.depth_plane.build" ) @@ -95,20 +96,26 @@ class AddonTest(unittest.TestCase): ) cls.image_background = importlib.import_module("anyimage.operators.background") cls.image_mesh = importlib.import_module("anyimage.operators.surface") - cls.image_extract_common = importlib.import_module( - "anyimage.operators.image_extract.common" + cls.image_edit_common = importlib.import_module( + "anyimage.operators.image_edit.common" ) - cls.image_region = importlib.import_module( - "anyimage.operators.image_extract.region" + cls.image_mask = importlib.import_module( + "anyimage.operators.image_edit.mask" ) cls.image_box = importlib.import_module( - "anyimage.operators.image_extract.box" + "anyimage.operators.image_edit.box" ) cls.image_lasso = importlib.import_module( - "anyimage.operators.image_extract.lasso" + "anyimage.operators.image_edit.lasso" + ) + cls.image_line = importlib.import_module( + "anyimage.operators.image_edit.line" + ) + cls.image_polyline = importlib.import_module( + "anyimage.operators.image_edit.polyline" ) cls.image_perspective = importlib.import_module( - "anyimage.operators.image_extract.perspective" + "anyimage.operators.image_edit.perspective" ) cls.balloon = importlib.import_module("anyimage.operators.balloon.operator") cls.balloon_geometry = importlib.import_module( @@ -253,6 +260,57 @@ class AddonTest(unittest.TestCase): self.assertIn(self.job_runtime.JobOperatorBase, operator.__mro__) self.assertIn(self.fake_bpy.types.Operator, operator.__mro__) + def test_realesrgan_debug_operator_is_a_direct_model_request(self): + operator_type = self.debug.DebugRealesrganGeneralX4v3 + self.assertIn(operator_type, self.addon.CLASSES) + self.assertEqual( + operator_type.bl_idname, + "anyimage.debug_realesrgan_general_x4v3", + ) + self.assertEqual( + operator_type.job_type, + "debug-realesrgan-general-x4v3", + ) + + with tempfile.TemporaryDirectory() as directory: + image_path = Path(directory) / "input.png" + image_path.touch() + settings = SimpleNamespace( + input_path=str(image_path), + device="CPU", + video_frames=3, + ) + context = SimpleNamespace( + scene=SimpleNamespace(anyimage_settings=settings) + ) + with ( + patch.object(self.debug, "require_environment"), + patch.object(self.debug, "require_model") as require_model, + ): + request = operator_type().request(context) + + require_model.assert_called_once_with( + self.addon.properties.REALESRGAN_MODEL_KEY + ) + self.assertEqual( + request, + { + "input": str(image_path), + "model": "REALESRGAN_GENERAL_X4V3", + "device": "cpu", + "video_frames": 3, + "start_frame": 1, + }, + ) + self.assertNotIn("tile", request) + self.assertNotIn("strength", request) + + panel = importlib.import_module("anyimage.panel") + self.assertIn( + panel.RealesrganGeneralX4v3DebugPanel, + self.addon.CLASSES, + ) + def test_default_request_sends_every_operator_property(self): class Example(self.job_runtime.JobOperatorBase): value: object() @@ -614,6 +672,62 @@ class AddonTest(unittest.TestCase): panel = importlib.import_module("anyimage.panel") self.assertFalse(hasattr(panel, "ModelsPanel")) + def test_preferences_draws_upscale_model_download(self): + preference_module = importlib.import_module("anyimage.preferences") + downloads = [] + + class Layout: + enabled = True + + def prop(self, *_args, **_options): + return None + + def separator(self, **_options): + return None + + def label(self, **_options): + return None + + def row(self, **_options): + return self + + def operator(self, identifier, **_options): + operator = SimpleNamespace(model="") + if identifier == "anyimage.download_model": + downloads.append(operator) + return operator + + preferences = preference_module.AnyImagePreferences() + preferences.layout = Layout() + catalog = ({ + "key": "REALESRGAN_GENERAL_X4V3", + "label": "Real-ESRGAN General x4v3", + "family": "upscale", + "enabled": True, + "ready": False, + },) + with ( + patch.object( + preference_module, + "ai_status", + return_value={ + "environment_ready": True, + "missing_models": (), + "ready": True, + }, + ), + patch.object( + preference_module, + "model_catalog", + return_value=catalog, + ), + patch.object(preference_module.runtime, "server_busy", return_value=False), + ): + preferences.draw(None) + + self.assertEqual(len(downloads), 1) + self.assertEqual(downloads[0].model, "REALESRGAN_GENERAL_X4V3") + def test_mesh_plane_operator_is_registered_with_material_options(self): operator = self.mesh_plane.ConvertToMeshPlane self.assertIn(operator, self.addon.CLASSES) @@ -793,141 +907,213 @@ class AddonTest(unittest.TestCase): self.assertGreater(self.balloon_geometry.polygon_area(contour), 0.0) self.assertLess(float(lengths.std() / lengths.mean()), 0.15) - def test_image_region_clears_rgba_outside_the_selection(self): + def test_image_mask_preserves_source_rgb_outside_the_selection(self): import numpy as np top_down = np.full((3, 3, 4), (0.2, 0.4, 0.6, 1.0), dtype=np.float32) source = np.flipud(top_down).ravel() - pixels, bounds = self.image_region.extract_region_pixels( + selection = self.image_mask.polygon_mask( + (3, 3), ((0, 0), (3, 0), (0, 3)) + ) + pixels, bounds = self.image_mask.edit_image_mask_pixels( source, (3, 3), - ((0, 0), (3, 0), (0, 3)), + selection, ) - result = np.flipud(np.asarray(pixels).reshape((3, 3, 4))) + width = bounds[2] - bounds[0] + height = bounds[3] - bounds[1] + result = np.flipud(np.asarray(pixels).reshape((height, width, 4))) - self.assertEqual(bounds, (0, 0, 3, 3)) - np.testing.assert_allclose(result[2, 2], (0.0, 0.0, 0.0, 0.0)) + self.assertEqual(bounds, (0, 0, 2, 2)) + np.testing.assert_allclose(result[1, 1], (0.2, 0.4, 0.6, 0.0)) self.assertEqual(float(result[0, 0, 3]), 1.0) - def test_image_region_clears_hidden_rgb_from_source_transparency(self): + def test_image_mask_preserves_hidden_rgb_from_source_transparency(self): import numpy as np top_down = np.full((2, 2, 4), (0.2, 0.4, 0.6, 1.0), dtype=np.float32) top_down[0, 0, 3] = 0.0 - pixels, _bounds = self.image_region.extract_region_pixels( + pixels, _bounds = self.image_mask.edit_image_mask_pixels( np.flipud(top_down).ravel(), (2, 2), - ((0, 0), (2, 0), (2, 2), (0, 2)), + np.ones((2, 2), dtype=np.float32), ) result = np.flipud(np.asarray(pixels).reshape((2, 2, 4))) - np.testing.assert_allclose(result[0, 0], (0.0, 0.0, 0.0, 0.0)) + np.testing.assert_allclose(result[0, 0], (0.2, 0.4, 0.6, 0.0)) - def test_image_region_keeps_out_of_image_canvas_transparent(self): + def test_image_mask_invert_runs_after_refinement(self): import numpy as np - top_down = np.ones((2, 2, 4), dtype=np.float32) - pixels, bounds = self.image_region.extract_region_pixels( + top_down = np.full((2, 2, 4), (0.2, 0.4, 0.6, 1.0), dtype=np.float32) + refined = np.asarray(((1.0, 0.25), (0.0, 1.0)), dtype=np.float32) + pixels, bounds = self.image_mask.edit_image_mask_pixels( np.flipud(top_down).ravel(), (2, 2), - ((-1, -1), (2, -1), (2, 2), (-1, 2)), + refined, + invert=True, ) - result = np.flipud(np.asarray(pixels).reshape((3, 3, 4))) + result = np.flipud(np.asarray(pixels).reshape((2, 2, 4))) - self.assertEqual(bounds, (-1, -1, 2, 2)) + self.assertEqual(bounds, (0, 0, 2, 2)) self.assertEqual(float(result[0, 0, 3]), 0.0) - self.assertEqual(float(result[1, 1, 3]), 1.0) + self.assertAlmostEqual(float(result[0, 1, 3]), 0.75) + np.testing.assert_allclose(result[0, 0], (0.2, 0.4, 0.6, 0.0)) - def test_image_region_rejects_a_fully_out_of_image_region(self): + def test_image_mask_rejects_a_fully_transparent_result(self): import numpy as np - with self.assertRaisesRegex(ValueError, "does not overlap"): - self.image_region.extract_region_pixels( + with self.assertRaisesRegex(ValueError, "no visible"): + self.image_mask.edit_image_mask_pixels( np.ones(2 * 2 * 4, dtype=np.float32), (2, 2), - ((3, 3), (4, 3), (4, 4), (3, 4)), + np.zeros((2, 2), dtype=np.float32), ) - def test_image_region_crops_to_visible_alpha_with_pixel_padding(self): + def test_image_mask_crops_to_visible_alpha_without_padding(self): import numpy as np - rgba = np.zeros((4, 5, 4), dtype=np.float32) - rgba[1:3, 2:4] = (0.2, 0.4, 0.6, 1.0) - pixels, bounds = self.image_extract_common.crop_pixels_to_alpha( + rgba = np.full((4, 5, 4), (0.8, 0.7, 0.6, 0.0), dtype=np.float32) + rgba[1, 2] = (0.2, 0.4, 0.6, 1.0) + rgba[2, 3] = (0.2, 0.4, 0.6, 1.0) + pixels, bounds = self.image_edit_common.crop_pixels_to_alpha( np.flipud(rgba).ravel(), (10, 20, 15, 24), - padding=1, ) - result = np.flipud(np.asarray(pixels).reshape((4, 4, 4))) + result = np.flipud(np.asarray(pixels).reshape((2, 2, 4))) - self.assertEqual(bounds, (11, 20, 15, 24)) - self.assertEqual(result.shape, (4, 4, 4)) - self.assertTrue(np.all(result[[0, -1], :, 3] == 0.0)) - self.assertTrue(np.all(result[:, [0, -1], 3] == 0.0)) - self.assertTrue(np.all(result[1:3, 1:3, 3] == 1.0)) + self.assertEqual(bounds, (12, 21, 14, 23)) + self.assertEqual(result.shape, (2, 2, 4)) + np.testing.assert_allclose(result[0, 1], (0.8, 0.7, 0.6, 0.0)) - def test_image_extract_tools_and_operators_are_registered(self): + def test_image_mask_restores_refined_alpha_to_source_coordinates(self): + import numpy as np + + top_down = np.zeros((2, 2, 4), dtype=np.float32) + top_down[:, :, 3] = ((1.0, 0.5), (0.0, 1.0)) + mask = self.image_mask.refined_mask_from_pixels( + np.flipud(top_down).ravel(), + (2, 2), + (1, 1, 3, 3), + (4, 4), + ) + + expected = np.zeros((4, 4), dtype=np.float32) + expected[1:3, 1:3] = top_down[:, :, 3] + np.testing.assert_allclose(mask, expected) + + def test_image_edit_aligns_canvas_without_changing_object_transform(self): + matrix = object() + image_object = SimpleNamespace( + matrix_world=matrix, + empty_display_size=4.0, + empty_image_offset=(0.0, 0.0), + ) + with patch.object( + self.image_edit_common, + "plane_bounds", + return_value=(-2.0, 2.0, -1.0, 1.0), + ): + self.image_edit_common.align_empty_display_to_bounds( + image_object, + (100, 50), + (50, 25), + (50, 0, 100, 25), + ) + + self.assertIs(image_object.matrix_world, matrix) + self.assertEqual(image_object.empty_display_size, 2.0) + self.assertEqual(image_object.empty_image_offset, (0.0, 0.0)) + + def test_image_edit_tools_and_operators_are_registered(self): import inspect - self.assertIn(self.image_region.ExtractImageRegion, self.addon.CLASSES) - self.assertIn( - self.image_region.RemoveImageRegionBackground, self.addon.CLASSES + operator = self.image_mask.EditImageMask + self.assertIn(operator, self.addon.CLASSES) + self.assertIn(self.image_mask.RefineImageMask, self.addon.CLASSES) + self.assertIn(self.image_mask.ActivateImageEditTool, self.addon.CLASSES) + tool_cases = ( + (self.image_box.ImageBoxTool, "BOX", "ops.sculpt.border_mask"), + (self.image_lasso.ImageLassoTool, "LASSO", "ops.sculpt.lasso_mask"), + (self.image_line.ImageLineTool, "LINE", "ops.sculpt.line_mask"), + ( + self.image_polyline.ImagePolylineTool, + "POLYLINE", + "ops.sculpt.polyline_mask", + ), ) - self.assertIn(self.image_region.ActivateImageExtractTool, self.addon.CLASSES) - self.assertNotIn(self.image_lasso.ImageLassoTool, self.addon.CLASSES) - self.assertNotIn(self.image_box.ImageBoxTool, self.addon.CLASSES) - self.assertEqual( - self.image_lasso.ImageLassoTool.bl_operator, - self.image_region.ExtractImageRegion.bl_idname, + for tool, gesture, icon in tool_cases: + self.assertNotIn(tool, self.addon.CLASSES) + self.assertEqual(tool.bl_operator, operator.bl_idname) + self.assertEqual(tool.bl_icon, icon) + self.assertEqual( + tool.bl_keymap[0][2]["properties"], [("gesture", gesture)] + ) + self.assertNotIn("bl_cursor", tool.__dict__) + self.assertIn("draw_settings", self.image_box.ImageBoxTool.__dict__) + self.assertIn("draw_settings", self.image_lasso.ImageLassoTool.__dict__) + self.assertNotIn("draw_settings", self.image_line.ImageLineTool.__dict__) + self.assertNotIn( + "draw_settings", self.image_polyline.ImagePolylineTool.__dict__ ) - self.assertEqual( - self.image_lasso.ImageLassoTool.bl_icon, - "ops.sculpt.lasso_mask", - ) - self.assertEqual( - self.image_box.ImageBoxTool.bl_icon, - "ops.sculpt.border_mask", - ) - self.assertEqual( - self.image_box.ImageBoxTool.bl_keymap[0][2]["properties"], - [("selection_shape", "BOX")], - ) - self.assertEqual( - self.image_lasso.ImageLassoTool.bl_keymap[0][2]["properties"], - [("selection_shape", "LASSO")], - ) - self.assertNotIn("bl_cursor", self.image_box.ImageBoxTool.__dict__) - self.assertNotIn("bl_cursor", self.image_lasso.ImageLassoTool.__dict__) - self.assertEqual( - self.image_region.ExtractImageRegion.__annotations__["selection_shape"][ - "default" - ], - "BOX", - ) - operator = self.image_region.ExtractImageRegion - self.assertIn("continuous", operator.__annotations__) - self.assertFalse(operator.__annotations__["continuous"]["default"]) - self.assertIn("padding", operator.__annotations__) - self.assertEqual(operator.__annotations__["padding"]["default"], 0) + self.assertEqual(operator.__annotations__["gesture"]["default"], "BOX") + self.assertIn("invert", operator.__annotations__) + self.assertNotIn("continuous", operator.__annotations__) + self.assertNotIn("padding", operator.__annotations__) settings = self.addon.properties.AnyImageSettings.__annotations__ - self.assertFalse(settings["image_extract_continuous"]["default"]) - self.assertEqual(settings["image_extract_padding"]["default"], 0) - self.assertEqual( - settings["image_extract_remove_background"]["name"], - "Remove Background", - ) - source_picker = inspect.getsource( - self.image_extract_common.resolve_image_extract_click - ) - continuous_offset = inspect.getsource( - self.image_extract_common.offset_continuous_result - ) + self.assertFalse(settings["image_edit_refine_selection"]["default"]) + self.assertNotIn("image_extract_continuous", settings) + self.assertNotIn("image_extract_padding", settings) + self.assertNotIn("image_extract_remove_background", settings) + source_picker = inspect.getsource(self.image_edit_common.resolve_image_edit_click) self.assertIn("bpy.ops.view3d.select(", source_picker) self.assertIn("Click an Image Empty", source_picker) self.assertIn("_restore_selection(", source_picker) - self.assertIn("EXTRACT_LAYER_STEP * layer", continuous_offset) + modal_source = inspect.getsource(operator.modal) + self.assertIn('event.type == "F"', modal_source) + self.assertIn("self.invert = not self.invert", modal_source) - def test_image_extract_first_click_selects_another_image_without_extracting(self): + def test_image_line_and_polyline_disable_refine_selection(self): + settings = SimpleNamespace(image_edit_refine_selection=True) + context = SimpleNamespace(scene=SimpleNamespace(anyimage_settings=settings)) + with patch.object(self.image_edit_common, "ai_ready", return_value=True): + for gesture in ("LINE", "POLYLINE"): + operator = SimpleNamespace( + gesture=gesture, + refine_selection=True, + ) + self.image_edit_common.apply_image_mask_settings(operator, context) + self.assertFalse(operator.refine_selection) + + operator = SimpleNamespace(gesture="BOX", refine_selection=False) + self.image_edit_common.apply_image_mask_settings(operator, context) + self.assertTrue(operator.refine_selection) + + def test_image_line_click_is_cancelled_before_projection(self): + import inspect + + modal_source = inspect.getsource(self.image_mask.EditImageMask.modal) + projection_source = inspect.getsource( + self.image_mask.image_line_half_plane_polygon + ) + self.assertIn("math.dist(self._path[0], self._cursor) < 2.0", modal_source) + self.assertIn("if len(line) != 2", projection_source) + + def test_image_polyline_closes_from_start_or_enter(self): + import inspect + + modal_source = inspect.getsource( + self.image_mask.EditImageMask._modal_polyline + ) + draw_source = inspect.getsource( + self.image_mask.EditImageMask._draw_selection + ) + self.assertIn("return self._confirm_polyline(context)", modal_source) + self.assertNotIn("DOUBLE_CLICK", modal_source) + self.assertIn('event.type in {"RET", "NUMPAD_ENTER"}', modal_source) + self.assertIn("_draw_polyline_close_indicator()", draw_source) + + def test_image_edit_first_click_selects_another_image_without_editing(self): selected = [] def image(name): @@ -962,8 +1148,8 @@ class AddonTest(unittest.TestCase): clicked.name: clicked, } try: - source, should_extract = ( - self.image_extract_common.resolve_image_extract_click( + source, should_edit = ( + self.image_edit_common.resolve_image_edit_click( context, event, lambda *_args: None ) ) @@ -975,10 +1161,10 @@ class AddonTest(unittest.TestCase): self.fake_bpy.ops = previous_ops self.assertIs(source, clicked) - self.assertFalse(should_extract) + self.assertFalse(should_edit) self.assertEqual(selected, []) - def test_image_extract_first_click_on_non_image_keeps_active_image(self): + def test_image_edit_first_click_on_non_image_keeps_active_image(self): selection_calls = [] active = SimpleNamespace( name="Active", @@ -1009,8 +1195,8 @@ class AddonTest(unittest.TestCase): ) self.fake_bpy.data.objects = {active.name: active, clicked.name: clicked} try: - source, should_extract = ( - self.image_extract_common.resolve_image_extract_click( + source, should_edit = ( + self.image_edit_common.resolve_image_edit_click( context, event, lambda *_args: None ) ) @@ -1022,7 +1208,7 @@ class AddonTest(unittest.TestCase): self.fake_bpy.ops = previous_ops self.assertIs(source, active) - self.assertTrue(should_extract) + self.assertTrue(should_edit) self.assertIs(view_objects.active, active) self.assertEqual(selection_calls, [False, True]) @@ -1045,7 +1231,7 @@ class AddonTest(unittest.TestCase): atol=1e-6, ) - def test_image_perspective_crops_alpha_and_applies_padding(self): + def test_image_perspective_tightly_crops_alpha(self): import numpy as np top_down = np.zeros((4, 4, 4), dtype=np.float32) @@ -1056,16 +1242,13 @@ class AddonTest(unittest.TestCase): (4, 4), ((0, 0), (4, 0), (4, 4), (0, 4)), 1.0, - 1, ) ) - result = np.flipud(np.asarray(pixels).reshape((4, 4, 4))) + result = np.flipud(np.asarray(pixels).reshape((2, 2, 4))) - self.assertEqual(output_size, (4, 4)) - self.assertEqual(placement_bounds, (0.0, 0.0, 4.0, 4.0)) - self.assertTrue(np.all(result[[0, -1], :, 3] == 0.0)) - self.assertTrue(np.all(result[:, [0, -1], 3] == 0.0)) - self.assertTrue(np.all(result[1:3, 1:3, 3] == 1.0)) + self.assertEqual(output_size, (2, 2)) + self.assertEqual(placement_bounds, (1.0, 1.0, 3.0, 3.0)) + self.assertTrue(np.all(result[:, :, 3] == 1.0)) def test_image_perspective_rejects_non_convex_corners(self): with self.assertRaisesRegex(ValueError, "convex"): @@ -1100,7 +1283,7 @@ class AddonTest(unittest.TestCase): def test_image_perspective_builds_polyline_mask_fill_triangles(self): self.assertEqual( - self.image_perspective.triangle_fan_vertices( + self.image_edit_common.triangle_fan_vertices( ((0, 0), (2, 0), (2, 2), (0, 2)) ), ( @@ -1153,7 +1336,7 @@ class AddonTest(unittest.TestCase): self.assertAlmostEqual(ratio, 2.0) - def test_image_perspective_maps_padding_to_source_plane_bounds(self): + def test_image_perspective_maps_crop_to_source_plane_bounds(self): bounds = self.image_perspective.perspective_placement_bounds( ((10, 20), (30, 20), (30, 40), (10, 40)), (100, 50), @@ -1162,39 +1345,29 @@ class AddonTest(unittest.TestCase): self.assertEqual(bounds, (8.0, 18.0, 32.0, 42.0)) - def test_image_perspective_tool_uses_four_clicks_and_shared_style(self): + def test_image_perspective_tool_uses_four_clicks_and_deform_behavior(self): import inspect - operator = self.image_perspective.ExtractImagePerspective + operator = self.image_perspective.EditImagePerspective tool = self.image_perspective.ImagePerspectiveTool settings = self.addon.properties.AnyImageSettings.__annotations__ self.assertIn(operator, self.addon.CLASSES) - self.assertIn( - self.image_perspective.RemoveImagePerspectiveBackground, - self.addon.CLASSES, - ) self.assertNotIn(tool, self.addon.CLASSES) self.assertEqual(tool.bl_operator, operator.bl_idname) - self.assertEqual(tool.bl_icon, "ops.sculpt.polyline_mask") + self.assertEqual(tool.bl_icon, "ops.sculpt.line_project") self.assertNotIn("bl_cursor", tool.__dict__) - self.assertIn("draw_settings", tool.__dict__) + self.assertNotIn("draw_settings", tool.__dict__) self.assertNotIn("image_perspective_aspect_mode", settings) self.assertNotIn("image_perspective_aspect_ratio", settings) self.assertNotIn("aspect_mode", operator.__annotations__) self.assertIn("aspect_ratio", operator.__annotations__) - self.assertIn("remove_background", operator.__annotations__) - self.assertIn("continuous", operator.__annotations__) - self.assertIn("padding", operator.__annotations__) + self.assertNotIn("refine_selection", operator.__annotations__) + self.assertNotIn("continuous", operator.__annotations__) + self.assertNotIn("padding", operator.__annotations__) modal_source = inspect.getsource(operator.modal) invoke_source = inspect.getsource(operator.invoke) execute_source = inspect.getsource(operator.execute) - background_request_source = inspect.getsource( - self.image_perspective.RemoveImagePerspectiveBackground.request - ) - background_response_source = inspect.getsource( - self.image_perspective.RemoveImagePerspectiveBackground.response - ) draw_source = inspect.getsource(operator._draw_overlay) preview_source = inspect.getsource(operator._draw_preview) finish_source = inspect.getsource(operator._finish) @@ -1208,17 +1381,9 @@ class AddonTest(unittest.TestCase): self.assertIn("active_view3d_tool_id(context)", modal_source) self.assertIn("drawing_in_region(", modal_source) self.assertNotIn("context.region_data.is_perspective", invoke_source) - self.assertIn("resolve_image_extract_click(", invoke_source) - self.assertIn("apply_image_extract_settings(self, context)", invoke_source) - self.assertIn( - "bpy.ops.anyimage.remove_image_perspective_background(", - execute_source, - ) + self.assertIn("resolve_image_edit_click(", invoke_source) self.assertIn("extract_perspective_pixels(", execute_source) self.assertIn("create_perspective_result(", execute_source) - self.assertIn("image_extract_background_request(", background_request_source) - self.assertIn("placement_polygon=polygon", background_response_source) - self.assertIn("self.padding", background_response_source) self.assertIn("draw_polygon_fill(vertices)", draw_source) self.assertIn("self._draw_preview()", draw_source) self.assertIn("draw_texture_2d(", preview_source) @@ -1230,16 +1395,16 @@ class AddonTest(unittest.TestCase): def test_image_box_builds_a_rectangular_screen_path(self): self.assertEqual( - self.image_region.box_screen_path(((2, 3), (8, 9))), + self.image_mask.box_screen_path((2, 3), (8, 9)), ((2, 3), (8, 3), (8, 9), (2, 9)), ) self.assertEqual( - self.image_region.box_screen_path(((2, 3), (3, 9))), + self.image_mask.box_screen_path((2, 3), (3, 9)), (), ) - def test_image_extract_selection_uses_white_dashed_segments(self): - segments = self.image_extract_common.dashed_line_segments( + def test_image_edit_selection_uses_white_dashed_segments(self): + segments = self.image_edit_common.dashed_line_segments( ((0, 0), (10, 0)), dash_length=3.0, gap_length=2.0 ) @@ -1248,22 +1413,29 @@ class AddonTest(unittest.TestCase): ((0.0, 0.0), (3.0, 0.0), (5.0, 0.0), (8.0, 0.0)), ) - def test_image_extract_activation_defaults_to_box(self): + def test_image_edit_activation_defaults_to_box(self): import inspect - source = inspect.getsource(self.image_region.ActivateImageExtractTool.execute) + source = inspect.getsource(self.image_mask.ActivateImageEditTool.execute) self.assertIn('name="anyimage.image_box"', source) - def test_continuous_keeps_the_source_empty_selected(self): - import inspect + def test_image_line_snaps_and_keeps_the_previewed_half_plane(self): + import math - execute_source = inspect.getsource(self.image_region.ExtractImageRegion.execute) - response_source = inspect.getsource( - self.image_region.RemoveImageRegionBackground.response + endpoint = self.image_mask.snapped_line_endpoint((0, 0), (10, 1), True) + self.assertAlmostEqual(math.degrees(math.atan2(endpoint[1], endpoint[0])), 5.0) + preview = self.image_mask.line_half_plane_polygon( + (0, 5), (10, 5), (10, 10) + ) + image = self.image_mask.image_line_half_plane_polygon( + ((0, 5), (10, 5)), (10, 10) + ) + self.assertEqual( + self.image_edit_common.polygon_bounds(preview), (0, 5, 10, 10) + ) + self.assertEqual( + self.image_edit_common.polygon_bounds(image), (0, 0, 10, 5) ) - expected = "select_image_extract_result(" - self.assertIn(expected, execute_source) - self.assertIn(expected, response_source) def test_balloon_interior_points_have_uniform_spacing(self): import numpy as np @@ -2160,12 +2332,25 @@ class AddonTest(unittest.TestCase): }, ), ( - self.image_perspective.ImagePerspectiveTool, + self.image_line.ImageLineTool, { "after": {self.image_lasso.ImageLassoTool.bl_idname}, "separator": False, }, ), + ( + self.image_polyline.ImagePolylineTool, + { + "after": {self.image_line.ImageLineTool.bl_idname}, + "separator": False, + }, + ), + ( + self.image_perspective.ImagePerspectiveTool, + { + "separator": False, + }, + ), ( self.balloon.BalloonTool, { diff --git a/tests/test_models.py b/tests/test_models.py index a2b07a7..b692c51 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -89,6 +89,35 @@ class ModelTest(unittest.TestCase): "BEN2-ONNX", ) + def test_realesrgan_general_x4v3_uses_verified_onnx_artifact(self): + model = self.models.get_downloadable_model( + "REALESRGAN_GENERAL_X4V3" + ) + + self.assertEqual( + model.repository, + "jonathanst29/tinier-upscale-models", + ) + self.assertEqual(model.license_name, "BSD 3-Clause") + self.assertEqual( + model.download_patterns, + ("realesr-general-x4v3.onnx",), + ) + self.assertEqual(model.ready_patterns, model.download_patterns) + self.assertEqual(model.r2_directory, "realesr-general-x4v3") + self.assertEqual( + model.r2_files, + (( + "realesr-general-x4v3.onnx", + 4866396, + "924ebad6532777303582d4ce7811849b88869231a3ae7093e0f21200249df8d5", + ),), + ) + self.assertEqual( + self.models.model_record(model)["family"], + "upscale", + ) + def test_enabled_da3_models_use_onnx_and_r2_checksums(self): for model in self.models.MODELS.values(): if not model.enabled: diff --git a/tests/test_pack.py b/tests/test_pack.py index 3b0f41f..1d922cc 100644 --- a/tests/test_pack.py +++ b/tests/test_pack.py @@ -68,14 +68,14 @@ class PackTest(unittest.TestCase): self.assertIn("operators/mesh_plane.py", names) self.assertIn("operators/image.py", names) self.assertIn("operators/balloon/operator.py", names) - self.assertIn("operators/image_extract/__init__.py", names) - self.assertIn("operators/image_extract/common.py", names) - self.assertIn("operators/image_extract/region.py", names) - self.assertIn("operators/image_extract/box.py", names) - self.assertIn("operators/image_extract/lasso.py", names) - self.assertIn("operators/image_extract/perspective.py", names) - self.assertNotIn("operators/image_lasso.py", names) - self.assertNotIn("operators/image_perspective.py", names) + self.assertIn("operators/image_edit/__init__.py", names) + self.assertIn("operators/image_edit/common.py", names) + self.assertIn("operators/image_edit/mask.py", names) + self.assertIn("operators/image_edit/box.py", names) + self.assertIn("operators/image_edit/lasso.py", names) + self.assertIn("operators/image_edit/line.py", names) + self.assertIn("operators/image_edit/polyline.py", names) + self.assertIn("operators/image_edit/perspective.py", names) self.assertIn("operators/surface.py", names) self.assertIn( "assets/O_AnyImage.blend", @@ -86,8 +86,7 @@ class PackTest(unittest.TestCase): self.assertIn("server/geometry/grid_mesh.py", names) self.assertIn("server/geometry/surface.py", names) self.assertIn("server/jobs/background.py", names) - self.assertIn("server/jobs/image_extract.py", names) - self.assertNotIn("server/jobs/image_lasso.py", names) + self.assertIn("server/jobs/image_mask.py", names) self.assertIn("server/models/ben2.py", names) self.assertIn("server/models/da3.py", names) self.assertIn("server/models/moge2.py", names) @@ -96,6 +95,7 @@ class PackTest(unittest.TestCase): self.assertIn("server/jobs/debug_da3.py", names) self.assertIn("server/jobs/debug_moge2.py", names) self.assertIn("server/jobs/debug_ben2.py", names) + self.assertIn("server/jobs/debug_realesrgan.py", names) self.assertNotIn("server/processing.py", names) self.assertNotIn("server/progress.py", names) self.assertIn("model_catalog.py", names) @@ -103,6 +103,8 @@ class PackTest(unittest.TestCase): self.assertIn("server/models/onnx_ben2.py", names) self.assertIn("server/models/onnx_da3.py", names) self.assertIn("server/models/onnx_moge2.py", names) + self.assertIn("server/models/onnx_realesrgan.py", names) + self.assertIn("server/models/realesrgan.py", names) self.assertNotIn("server/da3.py", names) self.assertNotIn("server/moge2.py", names) self.assertIn("server/__init__.py", names) diff --git a/tests/test_realesrgan_onnx.py b/tests/test_realesrgan_onnx.py new file mode 100644 index 0000000..262dbbf --- /dev/null +++ b/tests/test_realesrgan_onnx.py @@ -0,0 +1,75 @@ +import sys +import unittest +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +from PIL import Image + + +SOURCE_DIR = ( + Path(__file__).resolve().parents[1] + / "src" + / "anyimage" + / "server" + / "models" +) +sys.path.insert(0, str(SOURCE_DIR)) + +import onnx_realesrgan + + +class FakeSession: + def get_inputs(self): + return [SimpleNamespace(name="input")] + + def get_outputs(self): + return [SimpleNamespace(name="output")] + + def run(self, output_names, feeds): + self.output_names = output_names + self.feeds = feeds + source = feeds["input"] + return [np.repeat(np.repeat(source, 4, axis=2), 4, axis=3)] + + +class RealesrganOnnxTest(unittest.TestCase): + def test_preprocess_matches_dynamic_rgb_layout(self): + image = Image.new("RGB", (7, 5), (64, 128, 192)) + + result = onnx_realesrgan.preprocess(image) + + self.assertEqual(result.shape, (1, 3, 5, 7)) + self.assertEqual(result.dtype, np.float32) + self.assertTrue(result.flags.c_contiguous) + np.testing.assert_allclose( + result[0, :, 0, 0], + np.array((64, 128, 192), dtype=np.float32) / 255.0, + ) + + def test_infer_returns_direct_four_x_rgb_prediction(self): + session = FakeSession() + image = Image.new("RGB", (7, 5), (10, 20, 30)) + + result = onnx_realesrgan.infer(session, image) + + self.assertEqual(result.mode, "RGB") + self.assertEqual(result.size, (28, 20)) + self.assertEqual(session.output_names, ["output"]) + self.assertEqual(session.feeds["input"].shape, (1, 3, 5, 7)) + self.assertEqual(result.getpixel((0, 0)), (10, 20, 30)) + + def test_infer_rejects_non_four_x_output(self): + class InvalidSession(FakeSession): + def run(self, _output_names, feeds): + return [feeds["input"]] + + with self.assertRaisesRegex(RuntimeError, "unexpected output shape"): + onnx_realesrgan.infer( + InvalidSession(), + Image.new("RGB", (7, 5)), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_server.py b/tests/test_server.py index d0a738a..771525f 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -79,7 +79,14 @@ class ServerTest(unittest.TestCase): snapshot = manager.snapshot() self.assertEqual( snapshot, - {"loaded": {"ben2": "", "da3": "", "moge2": ""}}, + { + "loaded": { + "ben2": "", + "da3": "", + "moge2": "", + "realesrgan": "", + } + }, ) self.assertIn("server.model_catalog", sys.modules) self.assertNotIn("model_catalog", sys.modules) @@ -94,6 +101,8 @@ class ServerTest(unittest.TestCase): cache.da3_key = None cache.moge_model = None cache.moge_key = None + cache.realesrgan_model = None + cache.realesrgan_key = None return cache def create_server(self): @@ -184,6 +193,25 @@ class ServerTest(unittest.TestCase): self.assertIsNone(cache.da3_model) self.assertIsNotNone(cache.ben2_model) + def test_realesrgan_model_is_reused_for_matching_runtime_key(self): + loads = [] + realesrgan = ModuleType("server.models.onnx_realesrgan") + realesrgan.create_session = lambda model_dir, device: ( + loads.append((model_dir, device)) or object() + ) + cache = self.create_cache() + + with patch.dict( + sys.modules, + {"server.models.onnx_realesrgan": realesrgan}, + ): + first = cache.get_realesrgan("model", "cpu") + second = cache.get_realesrgan("model", "cpu") + + self.assertIs(first, second) + self.assertEqual(loads, [("model", "cpu")]) + self.assertEqual(cache.snapshot()["realesrgan"], "model") + def test_dedicated_server_runs_only_one_job_at_a_time(self): server = self.server_module.JobServer("Test", storage_root=self.storage_root) first_started = threading.Event() @@ -220,7 +248,7 @@ class ServerTest(unittest.TestCase): "download-model", "download-required-models", "analyze-balloon", - "extract-image-background", + "refine-image-mask", "remove-background", "extract-frames", "generate-depth", @@ -228,6 +256,7 @@ class ServerTest(unittest.TestCase): "debug-da3", "debug-moge2", "debug-ben2", + "debug-realesrgan-general-x4v3", }, ) @@ -401,16 +430,16 @@ class ServerTest(unittest.TestCase): self.assertEqual(tuple(result[0, 0]), (20, 40, 60, 64)) self.assertEqual(tuple(result[3, 3]), (0, 0, 0, 0)) - def test_image_extract_canvas_clears_rgba_outside_the_selection(self): + def test_image_mask_canvas_clears_rgba_outside_the_selection(self): import numpy as np from PIL import Image - image_extract = sys.modules["server.jobs.image_extract"] + image_mask = sys.modules["server.jobs.image_mask"] source_path = self.storage_root / "source.png" pixels = np.full((2, 2, 4), (20, 40, 60, 255), dtype=np.uint8) Image.fromarray(pixels, mode="RGBA").save(source_path) - canvas, _selection, bounds = image_extract.create_selection_canvas( + canvas, _selection, bounds = image_mask.create_selection_canvas( source_path, ((-1, -1), (2, -1), (-1, 2)), ) @@ -421,25 +450,22 @@ class ServerTest(unittest.TestCase): self.assertEqual(tuple(result[2, 2]), (0, 0, 0, 0)) self.assertEqual(tuple(result[1, 1]), (20, 40, 60, 255)) - def test_image_extract_crops_to_alpha_with_pixel_padding(self): + def test_image_mask_crops_tightly_to_alpha(self): import numpy as np from PIL import Image - image_extract = sys.modules["server.jobs.image_extract"] + image_mask = sys.modules["server.jobs.image_mask"] pixels = np.zeros((4, 5, 4), dtype=np.uint8) pixels[1:3, 2:4] = (20, 40, 60, 255) - result, bounds = image_extract.crop_to_alpha( + result, bounds = image_mask.crop_to_alpha( Image.fromarray(pixels, mode="RGBA"), (10, 20, 15, 24), - padding=1, ) result = np.asarray(result) - self.assertEqual(bounds, (11, 20, 15, 24)) - self.assertEqual(result.shape, (4, 4, 4)) - self.assertTrue(np.all(result[[0, -1], :, 3] == 0)) - self.assertTrue(np.all(result[:, [0, -1], 3] == 0)) - self.assertTrue(np.all(result[1:3, 1:3, 3] == 255)) + self.assertEqual(bounds, (12, 21, 14, 23)) + self.assertEqual(result.shape, (2, 2, 4)) + self.assertTrue(np.all(result[:, :, 3] == 255)) def test_balloon_job_sends_the_masked_crop_to_ben2(self): import numpy as np diff --git a/tests/test_worker.py b/tests/test_worker.py index 57ade11..f8d9dfc 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -115,6 +115,7 @@ class WorkerTest(unittest.TestCase): cls.ben2 = load_server_module("models.ben2") cls.da3_model = load_server_module("models.da3") cls.moge2_model = load_server_module("models.moge2") + cls.realesrgan_model = load_server_module("models.realesrgan") cls.surface_outputs = load_server_module("outputs.surface") cls.debug_da3_outputs = load_server_module("outputs.debug_da3") cls.debug_moge2_outputs = load_server_module("outputs.debug_moge2") @@ -124,6 +125,9 @@ class WorkerTest(unittest.TestCase): cls.debug_da3_job = load_server_module("jobs.debug_da3") cls.debug_moge2_job = load_server_module("jobs.debug_moge2") cls.debug_ben2_job = load_server_module("jobs.debug_ben2") + cls.debug_realesrgan_job = load_server_module( + "jobs.debug_realesrgan" + ) cls.frames_job = load_server_module("jobs.frames") cls.foreground_job = load_server_module("jobs.foreground") cls.download = load_server_module("model_download") @@ -608,13 +612,24 @@ class WorkerTest(unittest.TestCase): "cpu", preserve_input_alpha=True, ) + raw_result, _load_ms, _inference_ms = self.ben2.infer_foreground( + input_path, + "model", + "cpu", + ) inference_pixels = received[0] result_pixels = np.asarray(result.convert("RGBA"), dtype=np.uint8) self.assertEqual(tuple(inference_pixels[0, 0]), (0, 0, 0)) self.assertEqual(tuple(inference_pixels[0, 1]), (0, 128, 0)) + self.assertEqual(tuple(received[1][0, 0]), (0, 0, 0)) + self.assertEqual(tuple(received[1][0, 1]), (0, 128, 0)) self.assertEqual(tuple(result_pixels[0, 0]), (255, 0, 0, 0)) self.assertEqual(tuple(result_pixels[0, 1]), (0, 255, 0, 64)) + self.assertEqual( + tuple(np.asarray(raw_result.convert("RGBA"))[0, 0]), + (10, 20, 30, 128), + ) def test_ben2_background_removal_processes_each_frame(self): inference_options = [] @@ -874,6 +889,39 @@ class WorkerTest(unittest.TestCase): (context.directory / "foreground.png").read_bytes(), b"foreground", ) + + def test_realesrgan_debug_saves_each_direct_api_prediction(self): + from PIL import Image + + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + frames = root / "frames" + frames.mkdir() + (frames / "a.png").touch() + (frames / "b.png").touch() + context = FakeJobContext(root / "result", "debug-realesrgan") + with patch.object( + self.debug_realesrgan_job, + "infer_one", + return_value=Image.new("RGB", (8, 8), (10, 20, 30)), + ) as infer_one: + self.debug_realesrgan_job.run( + context, + { + "input": str(frames), + "model": "REALESRGAN_GENERAL_X4V3", + "video_frames": 2, + "device": "cpu", + }, + ) + + self.assertEqual(infer_one.call_count, 2) + predictions = sorted(context.directory.glob("prediction.*.png")) + self.assertEqual(len(predictions), 2) + with Image.open(predictions[0]) as prediction: + self.assertEqual(prediction.size, (8, 8)) + self.assertEqual(prediction.mode, "RGB") + def test_da3_removes_background_after_depth_using_original_input(self): order = [] inference_inputs = []