Update SnapForm to version 0.10.10 and add technical documentation
- Bump SnapForm version in uv.lock from 0.9.30 to 0.10.10. - Introduce comprehensive technical documentation for SnapForm, detailing features, usage instructions, and development guidelines.
This commit is contained in:
2
.gitignore
vendored
2
.gitignore
vendored
@@ -6,3 +6,5 @@ dist/
|
|||||||
build/
|
build/
|
||||||
*.zip
|
*.zip
|
||||||
.pytest_cache/
|
.pytest_cache/
|
||||||
|
|
||||||
|
sources/
|
||||||
|
|||||||
355
README.md
355
README.md
@@ -1,331 +1,64 @@
|
|||||||
# SnapForm
|
# SnapForm
|
||||||
|
|
||||||
`SnapForm` 是一套从图片快速构建空间结构的 Blender Extension。目前以
|
SnapForm 是一个 Blender 扩展,可以把一张照片、一段视频或一组序列帧快速变成有真实深度的立体面片。
|
||||||
[Depth Anything 3](https://github.com/ByteDance-Seed/Depth-Anything-3)
|
|
||||||
作为首个 AI 能力,在独立 Runtime 中运行 DA3,可将官方导出结果保存在独立 job 目录,
|
|
||||||
也可把 Image Empty 转换成带几何与材质置换的立体面片。Debug 面板还可运行
|
|
||||||
[BEN2](https://github.com/PramaLLC/BEN2) 背景剔除功能测试。
|
|
||||||
|
|
||||||
## 当前功能
|
## 它能做什么
|
||||||
|
|
||||||
- 使用 Blender Extension manifest,可通过 `Install from Disk` 安装
|
- 一键把 Image Empty 转换为立体面片:让 AI 估算深度,或直接用原图作为深度
|
||||||
- 在独立 Python Runtime 的常驻 HTTP AI Host 中运行 Torch、DA3 与 BEN2,避免污染 Blender Python 进程
|
- 一键去除图片背景,生成透明前景图
|
||||||
- Blender 使用内置 `urllib.request` 访问仅监听 `127.0.0.1` 的 Host;请求带随机 Token
|
- 支持视频与序列帧输入,逐帧处理,动画也能得到逐帧对应的立体面片
|
||||||
- 常驻 Host 复用 Torch、CUDA、BEN2 与最近使用的 DA3 模型,连续执行时不再重复加载
|
- 全部操作都在右键菜单里完成,不需要写任何代码
|
||||||
- 多个 Blender 实例自动连接同一个共享 Host;`Host` 面板显示状态、端口、客户端与模型缓存,并提供安装、启动、清理、重启和停止控制
|
|
||||||
- Runtime、模型和任务结果默认保存在 `~/.snapform`,可由多个 Blender 版本共享
|
|
||||||
- 支持 DA3 Small、Base、Large 1.1、Giant 1.1 和 Nested Giant Large 1.1
|
|
||||||
- 支持 BEN2 Base 单图背景剔除测试,输出透明 PNG 与推理耗时
|
|
||||||
- 支持单张 JPG、JPEG、PNG、WEBP 图片或常见视频文件输入
|
|
||||||
- 视频均匀抽帧,可设置参与重建的最大帧数
|
|
||||||
- `Host > Models` 默认折叠,展开后可逐项查看和下载所有可用模型
|
|
||||||
- Runtime 已就绪时不额外占用面板空间;模型列表每个模型只占一行
|
|
||||||
- `Debug` 默认折叠,内部的 DA3 与 BEN2 测试项也分别默认折叠
|
|
||||||
- 使用一个导出类型下拉菜单和一个 `Export` 按钮
|
|
||||||
- 只显示当前格式的专属参数;GLB 与 PLY 设置不会同时出现
|
|
||||||
- Debug 的模型选项按格式过滤;PLY 只提供两个带 Gaussian head 的 Giant 模型
|
|
||||||
- Debug 与 Convert to Depth Plane 的模型选项直接显示下载状态;未下载模型带
|
|
||||||
`(Missing)`,已下载模型保持原名
|
|
||||||
- 分别调用 DA3 的 `glb`、`gs_ply`、`depth_vis` 和 `mini_npz` exporter
|
|
||||||
- Image Empty 右键菜单可直接运行 DA3,根据 `mini_npz` 深度生成面片
|
|
||||||
- Image Empty 右键菜单可直接运行 BEN2,保留原图并换用新建的透明背景图片
|
|
||||||
- 使用浮点深度贴图、Simple Subdivision、Displace Modifier 和原图 Bump
|
|
||||||
构建立体面片
|
|
||||||
- 转换对话框可选择推理模式、处理分辨率和面片网格细分
|
|
||||||
- Debug 的 DA3 导出或 BEN2 背景剔除完成后只打开对应 job 文件夹,不向当前 Blender 文件添加任何数据
|
|
||||||
- 完整保留每次任务生成的所有官方文件
|
|
||||||
- 安装和模型下载使用后台子进程;推理提交给常驻 AI Host,状态栏实时显示阶段、百分比进度条与取消按钮
|
|
||||||
- 后台任务运行时,Install、Download 和 Export 按钮保持原位置并变灰,任务完成后恢复
|
|
||||||
|
|
||||||
## 使用
|
## 安装
|
||||||
|
|
||||||
1. 使用 `uv run python pack.py` 构建 Extension ZIP。
|
1. 下载对应系统的扩展压缩包:
|
||||||
2. 在 Blender 的 `Preferences > Extensions` 中选择 `Install from Disk`。
|
- Windows:`SnapForm.v0.10.10.windows-x64.zip`
|
||||||
3. 在 Extension Preferences 中确认 `uv` 路径和 Torch Backend。
|
- macOS(Apple 芯片):`SnapForm.v0.10.10.macos-arm64.zip`
|
||||||
4. 在 `Host` 点击 `Install Runtime`。首次安装会下载 Python、PyTorch、DA3、BEN2、FastAPI、Uvicorn,
|
- macOS(Intel):`SnapForm.v0.10.10.macos-x64.zip`
|
||||||
并默认下载 `DA3 Small`、`DA3 Base` 和 `BEN2 Base` 模型。安装按 8 个阶段依次推进,每段进度
|
- Linux:`SnapForm.v0.10.10.linux-x64.zip`
|
||||||
从 0 到 100%:下载 uv(如需)、创建 Python Runtime、下载 PyTorch、安装 xFormers(macOS 跳过)、
|
2. 打开 Blender,进入 `Preferences > Extensions`
|
||||||
安装 AI 依赖、修补 DA3(跳过视频渲染),然后是三个模型的下载。进度文字统一为 `阶段/总数 阶段名: 状态`,
|
3. 点击 `Install from Disk`,选择下载的压缩包
|
||||||
模型下载阶段会实时显示下载速度,例如 `6/8 da3 small: downloading 14.2MB/s`。
|
4. 在扩展偏好设置中确认 `uv` 路径与 Torch 后端(保持默认即可);如果下载经常超时,
|
||||||
Runtime 就绪后即使模型下载中断也不影响 Host 使用,可稍后在 `Models` 面板单独补下。
|
可以在 `PyPI Index` 里填一个国内镜像,例如 `https://pypi.tuna.tsinghua.edu.cn/simple`
|
||||||
5. 在 3D View 按 `N`,打开 `SnapForm` 页签。
|
5. 打开 3D 视图右侧的 `SnapForm` 面板,点击 `Install Runtime`
|
||||||
6. 如需其他模型,展开 `Host > Models`,找到模型并点击该行的 `Download`。
|
|
||||||
7. 展开默认收起的 `Debug`,选择 `Export Format` 和本次执行使用的模型。
|
|
||||||
8. 选择一张图片或一个视频文件。视频可设置均匀抽取的帧数,然后点击 `Export`。
|
|
||||||
9. 测试背景剔除时,确认 `Host > Models` 中的 `BEN2 Base` 已就绪,然后在 `Debug` 的
|
|
||||||
`BEN2 Background Removal` 区域选择单张图片并点击 `Remove Background`。
|
|
||||||
完成后会打开 job 文件夹,其中的 `foreground.png` 带透明 Alpha;状态栏显示
|
|
||||||
BEN2 推理、模型加载与任务总耗时。当前测试固定关闭较慢的 foreground refinement。
|
|
||||||
|
|
||||||
## Image Empty 去背景
|
DA3、BEN2 源码包和 uv 已内置在扩展里,安装时直接从本地解包,不需要访问 GitHub。首次使用仍会自动
|
||||||
|
下载运行环境依赖(PyTorch 等)和 AI 模型,请保持网络畅通,下载需要一些时间。
|
||||||
|
|
||||||
1. 确认 `Host > Models` 中的 `BEN2 Base` 已就绪;首次安装 Runtime 时会默认下载。
|
## 快速上手
|
||||||
2. 在 3D View 选中 Image Empty,右键选择 `Remove Background`。
|
|
||||||
3. 完成后以 Fake User 保留原 Image datablock,新建并 Pack 一张透明 PNG 到当前 `.blend`,然后只把
|
|
||||||
当前 Empty 的 `data` 换成新图;Empty 的 Transform、显示尺寸、Offset 及其他设置保持不变。
|
|
||||||
4. 新 Image 使用原图去掉格式后再追加 `.rb` 的名称。例如 `img.png` 生成 `img.rb`,
|
|
||||||
`img.png.001` 生成 `img.001.rb`。已有同名 datablock 时由 Blender 自动追加编号。
|
|
||||||
|
|
||||||
原图已 Pack 时会先写出临时 PNG 交给 BEN2,任务结束后删除临时输入;原 Image 仍保持 Pack
|
### 转换为立体面片
|
||||||
状态和原始内容。BEN2 的 `foreground.png` 也会保留在对应 job 目录中。
|
|
||||||
|
|
||||||
## Image Empty 转深度面片
|
1. 在 3D 视图中选中一个 Image Empty(空物体,显示方式为 Image)
|
||||||
|
2. 右键选择 `Convert to Depth Plane`
|
||||||
|
3. 在对话框中选择:
|
||||||
|
- `Depth from`:`Original` 直接用原图,`AI Generated` 让 AI 估算深度
|
||||||
|
- 需要去背景时勾选 `Remove Background`
|
||||||
|
- 其余选项保持默认即可
|
||||||
|
4. 点击确认,稍等片刻就得到带位移的立体面片
|
||||||
|
|
||||||
1. 在 3D View 选中 `Object Type = Empty`、`Display As = Image` 的对象。
|
### 去除背景
|
||||||
2. 右键选择 `Convert to Depth Plane`,通过 `Depth from` 选择深度来源:
|
|
||||||
- `AI Generated`:显示 Model、Use Ray Pose 和 Resolution Scale。实际 Process Res 为图片长边像素数
|
|
||||||
乘以 Resolution Scale;范围为 `0.1–1.0`,默认 `1.0`。计算后的 Process Res 会直接显示在
|
|
||||||
`Resolution Scale (… px)` 标签中。
|
|
||||||
转换导出 `mini_npz`,使用前需要在
|
|
||||||
`Host > Models` 下载所选模型,并固定使用 `process_res_method="upper_bound_resize"`;设备和参考视图策略
|
|
||||||
继续使用场景设置。Model 下拉菜单只列出已经下载的 DA3 模型;若没有可用模型,对话框会引导到
|
|
||||||
`Host > Models` 下载。
|
|
||||||
- `Original`:直接把原图作为深度贴图,不显示估算参数。Remove Background 关闭时不启动
|
|
||||||
Runtime,也不要求源图片存在于磁盘。
|
|
||||||
`Original` 是默认深度来源。选择 `AI Generated` 时,其 Model、Use Ray Pose 与 Resolution Scale
|
|
||||||
参数会集中显示在独立边框中。两种模式都有默认折叠的 Options,内部依次为 Displace via、
|
|
||||||
Subdivision、Remove Background 和 Unshaded;Delete Original 保持在折叠项外。
|
|
||||||
Remove Background 默认关闭,与 Depth from 无关;开启后需要下载 `BEN2 Base`,并从原图生成、
|
|
||||||
Pack `.rb` 透明图片。它只替换生成面片材质的颜色图片,不修改源 Image Empty 的 Image datablock。
|
|
||||||
AI Generated 仍从原图生成 `.depth`;例如 `img.png` 会得到 `img.rb` 与 `img.depth`,不会生成
|
|
||||||
`img.rb.depth`。Original 仍用原图作为位移来源,不改变原有深度逻辑。
|
|
||||||
Displace via 可选 `Modifier`(默认)或 `Material`;
|
|
||||||
Unshaded 默认关闭;Subdivision 默认为 `5`,并默认删除原对象。
|
|
||||||
3. AI Generated 完成后会自动创建面片。面片几何先由 `mini_npz` 的相机内参与深度确定,再通过对象
|
|
||||||
Transform 匹配原 Image Empty 的可见尺寸、位置、旋转和 Offset。Original 模式直接使用源 Empty
|
|
||||||
的可见尺寸、Offset 与 Transform。开启 `Delete Original` 时,成功生成后
|
|
||||||
删除源 Empty。生成面片 Object 使用源 Empty 的原名;关闭 `Delete Original` 时不处理
|
|
||||||
同名冲突,由 Blender 自动为生成面片追加编号。
|
|
||||||
4. 短边分段数为 `2^Subdivision`,长边分段数根据图像宽高比取最接近的整数。例如高:宽为
|
|
||||||
`2:3`、Subdivision 为 `1` 时,直接生成 `2×3` 个面;Subdivision 为 `3` 时生成
|
|
||||||
`8×12` 个面。随后的 Simple Subdivision 修改器使用原生名称,Viewport Levels 固定为
|
|
||||||
`1`,Render Levels 固定为 `2`,不再随 Resolution Scale 或 Subdivision 动态变化。
|
|
||||||
|
|
||||||
所有任务使用统一的状态栏进度显示。变化中的阶段说明显示在右侧的紧凑进度条内部,取消按钮
|
选中 Image Empty 后,右键选择 `Remove Background`,完成后 Empty 会自动换成透明前景图。
|
||||||
位于进度条右侧。Image Depth 的阶段包括 Runtime 启动、输入准备、模型加载、深度估算、缓存
|
|
||||||
写出和面片构建。进度更新会强制刷新状态栏与窗口,无需移动鼠标即可实时显示。安装中断后残留的
|
|
||||||
半成品 Runtime 会在下次安装前自动清理,不会因为中断而无法再次安装。
|
|
||||||
|
|
||||||
导入过程精确保留 NPZ 内深度值之间的差值。例如深度 `2.0` 和 `5.5` 在
|
### 视频与序列帧
|
||||||
`Displacement Strength = -1` 时,沿局部 Z 轴相差 `3.5 Blender Unit`。它不会把数值压缩到
|
|
||||||
`0..1`,也不再提供视觉归一化模式。
|
|
||||||
|
|
||||||
面片尺寸以有效深度的中位数 `Zref` 作为参考截面。先用 `mini_npz` 中的相机内参计算原始
|
SnapForm 支持视频(mp4、mkv、mov 等)和序列帧,会逐帧处理并生成逐帧对应的立体面片。处理范围默认
|
||||||
投影面积,再恢复源图纵横比:
|
跟随 Empty 上设置的 `Image Frames` 与 `Offset`,不需要额外配置。
|
||||||
|
|
||||||
```text
|
## 任务结果保存在哪里
|
||||||
RawWidth = depth_width × Zref × abs(Displacement Strength) / fx
|
|
||||||
RawHeight = depth_height × Zref × abs(Displacement Strength) / fy
|
|
||||||
Area = RawWidth × RawHeight
|
|
||||||
Aspect = source_image_width / source_image_height
|
|
||||||
Width = sqrt(Area × Aspect)
|
|
||||||
Height = sqrt(Area / Aspect)
|
|
||||||
```
|
|
||||||
|
|
||||||
这样即使 DA3 估计出的 `fx/fy` 不一致,也不会把正方形原图拉成长方形。主点 `cx/cy` 仍用于
|
任务结果保存在 `~/.snapform/jobs/` 下,每个任务一个文件夹。视频或序列帧生成的面片会引用这些文件夹
|
||||||
确定原始相机面片,随后对象 XY Transform 将它匹配到源 Image Empty。生成对象会记录
|
里的序列帧,请保留对应文件夹,不要随意删除。
|
||||||
`snapform_reference_depth`、`snapform_camera_plane_width` 和
|
|
||||||
`snapform_camera_plane_height`,便于检查实际采用的
|
|
||||||
参考深度和尺寸。
|
|
||||||
|
|
||||||
选择 `Modifier` 时,AI Generated 的 Displace Modifier 深度贴图使用 `Extension = Extend`,
|
## 常见问题
|
||||||
`Strength` 固定为 `-1`、`Midlevel` 固定为 `1`。Original 把原图指定给 Displace Texture,并设置
|
|
||||||
`Texture Coordinates = UV`、`Midlevel = 0`;其余修改器设置保留 Blender 默认值。
|
|
||||||
选择 `Material` 时,不创建 Displace Modifier;深度 Image Texture 的 Color 连接到 Displacement
|
|
||||||
节点的 Height,再连接 Material Output 的 Displacement,并把材质 Surface Displacement 设置为
|
|
||||||
`Displacement Only`。AI Generated 使用 `Midlevel = 1`、`Scale = -1`;Original 使用
|
|
||||||
`Midlevel = 0`、`Scale = 1`,与对应修改器值一致。无论选择哪种呈现方式,原始彩色贴图都会连接
|
|
||||||
Principled BSDF 的 `Base Color`,并通过
|
|
||||||
`Bump.Height → Bump.Normal → Principled.Normal` 提供表面细节。Bump 的基础 Distance 为
|
|
||||||
`0.001`,会乘以面片 Transform 的统一缩放倍率。原图 Color 还会经过 `Float Curve` 连接到
|
|
||||||
Principled BSDF 的 `Roughness`,曲线两端输出分别为 `0.5` 和 `0.1`;材质 IOR 保持
|
|
||||||
Principled BSDF 默认值 `1.5`。
|
|
||||||
开启 `Unshaded` 后,原图 Color 连接 Emission,原图 Alpha 作为 Mix Shader 系数,在
|
|
||||||
Transparent BSDF 与 Emission 之间混合后连接 Material Output;不创建 Principled BSDF、Bump 或
|
|
||||||
Float Curve。若 Displace via 选择 `Material`,深度贴图与 Displacement
|
|
||||||
节点仍会保留并连接 Material Output 的 Displacement。
|
|
||||||
|
|
||||||
这里的“原始距离”是 DA3 模型输出的尺度。普通 Small、Base、Large 和 Giant 模型的深度不应
|
- 首次使用需要下载模型,如果下载慢或失败,稍后在 `Host > Models` 里重试即可
|
||||||
被解释为现实世界的米;需要现实尺度时应选择带 metric depth 的
|
- 安装 Runtime 时提示下载超时:新版会自动探测并切换到镜像、失败时自动重试一次;如果仍然
|
||||||
`DA3 Nested Giant Large 1.1`。此外,当前结果仍是固定 XY 平面上的深度位移面,不是使用
|
超时,在扩展偏好设置的 `PyPI Index` 中填入镜像地址(如清华)后重新安装即可
|
||||||
相机内参逐像素反投影得到的完整 camera-space 网格;中位深度截面为现有 Displace Modifier
|
- 找不到右键入口时,请确认选中了 Empty,且显示方式为 Image
|
||||||
工作流提供正确的整体高宽尺度。
|
- 动画面片不动:检查 Empty 的 `Image Frames` 是否大于 1,并确认播放场景帧
|
||||||
|
|
||||||
生成完成后会直接创建面片,不提供缓存重建入口或单独的导入参数。重新选择
|
## 技术文档
|
||||||
`Convert to Depth Plane` 会按当前 Depth from 再次转换。
|
|
||||||
面片、深度 Image、Texture 与 Material 会在同一个 Blender Undo 步骤中创建;转换完成后执行
|
|
||||||
`Ctrl+Z` 会完整撤销本次转换,不会越过转换去回退更早的 Image Empty 操作。
|
|
||||||
开启 Remove Background 时,`.rb` Image 的创建、Pack 和深度面片生成也属于
|
|
||||||
同一个 Undo 步骤。
|
|
||||||
|
|
||||||
Image Empty 使用 Packed Image 时,SnapForm 会通过 Blender 写出临时 PNG 交给 DA3,
|
面向开发者的技术内容见 [docs/technical.md](docs/technical.md)。
|
||||||
不会 Unpack 或修改源 Image;任务结束后自动删除临时文件。普通图片仍直接读取源文件。DA3 的原始
|
|
||||||
`results.npz` 保留在对应 job 的 `exports/mini_npz/` 中,生成的浮点深度 Image 会 Pack 到
|
|
||||||
当前 `.blend`。Depth Image 与 Displace Texture 根据源 Image datablock 命名。例如 Empty Object
|
|
||||||
名为 `obj_img`、引用的 Image 名为 `img.png` 时,Depth Image 与 Texture 都命名为 `img.depth`;
|
|
||||||
生成的 Object、Mesh 与 Material 仍命名为 `obj_img`。Pack 使用 OpenEXR,因而不会把大于 `1.0`
|
|
||||||
的原始深度截断。已有同名 datablock 时由 Blender 自动追加 `.001` 等编号。
|
|
||||||
|
|
||||||
选择 PLY 后,模型菜单只会提供 `DA3 Giant 1.1` 和
|
|
||||||
`DA3 Nested Giant Large 1.1`。Small、Base 和 Large 没有 Gaussian head,无法由
|
|
||||||
DA3 生成 `gs_ply`。
|
|
||||||
|
|
||||||
每次任务会在 `~/.snapform/jobs/<任务目录>/` 中保留输出产物:
|
|
||||||
|
|
||||||
```text
|
|
||||||
<GLB 任务目录>/
|
|
||||||
└── scene.glb
|
|
||||||
|
|
||||||
<PLY 任务目录>/
|
|
||||||
└── gs_ply/*.ply
|
|
||||||
|
|
||||||
<Depth 任务目录>/
|
|
||||||
└── depth_vis/*.jpg
|
|
||||||
|
|
||||||
<Mini NPZ 任务目录>/
|
|
||||||
└── exports/mini_npz/results.npz
|
|
||||||
|
|
||||||
<BEN2 背景剔除任务目录>/
|
|
||||||
└── foreground.png
|
|
||||||
```
|
|
||||||
|
|
||||||
建议从 `DA3 Base` 开始。它使用 Apache 2.0 模型权重;`DA3 Large 1.1` 权重采用
|
|
||||||
CC BY-NC 4.0,仅限非商业用途。
|
|
||||||
|
|
||||||
## 输出格式说明
|
|
||||||
|
|
||||||
`Export Format` 下拉菜单分别把以下原始 `export_format` 传给 DA3:
|
|
||||||
|
|
||||||
- GLB:`glb`
|
|
||||||
- PLY:`gs_ply`,同时传入 `infer_gs=True`
|
|
||||||
- Depth Images:`depth_vis`
|
|
||||||
- Mini NPZ:`mini_npz`
|
|
||||||
|
|
||||||
界面使用 DA3 API 的参数名、选项和默认值。GLB 专属设置仅在选择 `GLB` 时显示,
|
|
||||||
`gs_views_interval` 仅在选择 `PLY (3DGS)` 时显示:
|
|
||||||
|
|
||||||
- `process_res=504`
|
|
||||||
- `process_res_method="upper_bound_resize"`
|
|
||||||
- `use_ray_pose=False`
|
|
||||||
- `ref_view_strategy="saddle_balanced"`
|
|
||||||
- `conf_thresh_percentile=40.0`
|
|
||||||
- `num_max_points=1_000_000`
|
|
||||||
- `show_cameras=True`
|
|
||||||
- `export_kwargs["gs_ply"]["gs_views_interval"]=1`
|
|
||||||
|
|
||||||
DA3 官方 `glb` 是彩色点云与可选相机线框;`gs_ply` 是带 Gaussian 属性的 PLY;
|
|
||||||
`depth_vis` 是原图与着色深度图的并排预览;`mini_npz` 保存深度、置信度和可用的相机参数。
|
|
||||||
GLB 任务会关闭其默认附带的深度预览,使四种 job 产物保持分离。
|
|
||||||
|
|
||||||
官方格式与限制可参阅
|
|
||||||
[DA3 Python API 文档](https://github.com/ByteDance-Seed/Depth-Anything-3/blob/main/docs/API.md)。
|
|
||||||
|
|
||||||
## Runtime
|
|
||||||
|
|
||||||
Extension 本身不包含 Torch。`Install Runtime` 使用 `uv` 创建独立 Python 3.12
|
|
||||||
环境,并安装。机器上即使没有安装 `uv` 或 Python 也没关系:安装程序由 Blender
|
|
||||||
自带的 Python 启动,缺 `uv` 时会自动下载官方独立二进制(约 30MB),再用它下载
|
|
||||||
Python 3.12 与全部依赖。`uv` 会缓存在扩展存储目录下,后续安装直接复用。
|
|
||||||
安装前会自动预检磁盘空间与网络连通性,未通过会直接给出英文报错,不会开始下载。
|
|
||||||
预检与安装支持环境变量切换镜像与代理:`PIP_INDEX_URL`(PyPI 镜像)、
|
|
||||||
`SNAPFORM_TORCH_INDEX`(PyTorch 索引)、`HF_ENDPOINT`(Hugging Face 镜像,
|
|
||||||
如 `https://hf-mirror.com`)、`HTTP(S)_PROXY`/`ALL_PROXY`(代理)。设置后预检会直接检查
|
|
||||||
这些地址。直连不通时预检会**自动回退到镜像源**(清华 PyPI、阿里云 PyTorch、hf-mirror),
|
|
||||||
无需手动配置,状态栏会提示 `using mirror sources`;镜像也不通才会报错。
|
|
||||||
单独下载模型时也会先检查 Hugging Face 连通性,直连不通自动切换到 `hf-mirror` 镜像。
|
|
||||||
Windows 上代理若只在系统设置里而未写入环境变量,安装程序会尝试自动从系统代理注入子进程。
|
|
||||||
|
|
||||||
- PyTorch 2.9.1
|
|
||||||
- TorchVision 0.24.1
|
|
||||||
- xFormers 0.0.33.post2(与 PyTorch 2.9.1 对应,仅 Windows/Linux;macOS 无 wheel 自动跳过)
|
|
||||||
- 固定到已验证 commit 的 Depth Anything 3
|
|
||||||
- 固定到已验证 commit 的 BEN2
|
|
||||||
- FastAPI 0.139.2 与 Uvicorn 0.51.0
|
|
||||||
- DA3 与 BEN2 通过 GitHub 官方源码包安装(无需 Git);SnapForm 直接调用 DA3
|
|
||||||
底层 export 接口,`gs_ply` 导出不会触发 `gs_video` 渲染,因此不需要 gsplat
|
|
||||||
与 C++ 编译工具链
|
|
||||||
- `addict` 2.4.0(补充 DA3 上游包元数据暂未声明的运行依赖)
|
|
||||||
|
|
||||||
默认 CUDA Backend 为 CUDA 12.8,也可以选择 CPU。CPU 可以用于验证,但实际推理
|
|
||||||
通常很慢。
|
|
||||||
CUDA 12.8 的 PyTorch wheel 自带 CUDA 运行库,无需安装 CUDA Toolkit;只需要
|
|
||||||
NVIDIA 驱动 ≥ 570(数据中心显卡可用更老驱动)。驱动不满足或没有 NVIDIA 显卡时,
|
|
||||||
设备会自动回退到 CPU(macOS 为 MPS),只有显式选择 CUDA 才会报错。
|
|
||||||
|
|
||||||
macOS(Apple Silicon)会自动从 PyPI 安装 CPU 版 PyTorch,推理设备默认使用 MPS
|
|
||||||
(Apple Metal)加速,也可以在设备选项里切换 CPU。PyTorch 自 2.3 起不再提供 Intel
|
|
||||||
Mac(x86_64)轮子,Intel Mac 暂不支持。
|
|
||||||
|
|
||||||
Extension 启用后只会查找已有的共享 AI Host,不会主动启动新进程。第一次执行 AI 操作或在侧边栏
|
|
||||||
手动点击 `Start` 时才会启动 Host。Host 只绑定回环地址,并让操作系统分配动态端口;第一个实例把
|
|
||||||
端口、进程和随机 Token 写入 `host.json`,其余 Blender 通过标准库
|
|
||||||
`urllib.request` 连接同一个 Host。`host.lock` 保证同一 Storage Root 同时只有一个 Host,Token
|
|
||||||
则阻止同机其他进程直接调用接口。Host 全局同时只执行一个 AI 任务。
|
|
||||||
|
|
||||||
每个 Blender 会注册独立客户端并每 5 秒发送心跳。正常关闭 Blender 时先注销自身;如果当前 Host
|
|
||||||
就是由这个 Blender 启动的,会直接终止它,无需等待。异常退出的客户端会在 10 秒后过期;Host
|
|
||||||
在最后一个客户端离开后仍会保留模型缓存等待 20 秒再退出,作为 Host 由其他进程启动等场景的兜底。
|
|
||||||
取消任务不会重启 Host,并会保留已经加载的模型;
|
|
||||||
如果取消发生在一次不可中断的 Torch 调用中,任务会在该调用返回后的最近检查点结束。
|
|
||||||
|
|
||||||
侧边栏 `Host` 面板以 `Running (Clients: n)` 表示空闲服务,以 `Running (Busy)` 表示正在执行任务,
|
|
||||||
并显示缓存模型和活动任务。空闲且只有当前 Blender 连接时,可以清理模型缓存、重启或停止 Host;
|
|
||||||
多客户端或任务执行期间不会允许破坏共享服务。
|
|
||||||
|
|
||||||
BEN2 在首次使用后保持加载;DA3 保留最近使用的一个模型。切换 DA3 模型时会释放旧模型;加载时如果
|
|
||||||
显存不足,会先释放另一类缓存模型并重试。
|
|
||||||
|
|
||||||
Runtime Schema 已升级,旧 Runtime 需要重新执行 `Install Runtime` 才能安装 HTTP Host 依赖。
|
|
||||||
Host 的启动和崩溃日志写入 `~/.snapform/ai-host.log`。
|
|
||||||
|
|
||||||
Runtime、模型和任务结果默认保存在用户主目录的 `~/.snapform`:
|
|
||||||
|
|
||||||
```text
|
|
||||||
~/.snapform/
|
|
||||||
├── runtime/
|
|
||||||
├── models/
|
|
||||||
├── jobs/
|
|
||||||
├── ai-host.log
|
|
||||||
├── host.json # Host 运行时存在
|
|
||||||
└── host.lock # Host 运行时持有
|
|
||||||
```
|
|
||||||
|
|
||||||
不同 Blender 版本安装相同 Extension 后会共用这个目录。也可以在 Extension
|
|
||||||
Preferences 的 `Storage Root` 中修改共享位置。不要把模型写入 Extension 安装目录,
|
|
||||||
因为该目录可能只读,并且升级时会被替换。不要同时从多个 Blender 实例执行
|
|
||||||
`Install Runtime`;推理 Host、任务队列和模型读取可以共享。
|
|
||||||
|
|
||||||
## 开发
|
|
||||||
|
|
||||||
```bash
|
|
||||||
uv run python -m unittest discover -s tests -v
|
|
||||||
uv run python pack.py
|
|
||||||
```
|
|
||||||
|
|
||||||
项目结构:
|
|
||||||
|
|
||||||
```text
|
|
||||||
.
|
|
||||||
├── pack.py
|
|
||||||
├── src/
|
|
||||||
│ └── snapform/
|
|
||||||
│ ├── blender_manifest.toml
|
|
||||||
│ ├── ai_host.py
|
|
||||||
│ ├── bootstrap_runtime.py
|
|
||||||
│ ├── models.py
|
|
||||||
│ ├── image_depth.py
|
|
||||||
│ ├── operators.py
|
|
||||||
│ ├── panel.py
|
|
||||||
│ ├── preferences.py
|
|
||||||
│ ├── properties.py
|
|
||||||
│ ├── runtime.py
|
|
||||||
│ └── worker.py
|
|
||||||
└── tests/
|
|
||||||
```
|
|
||||||
|
|
||||||
## 上游与许可证
|
|
||||||
|
|
||||||
- Depth Anything 3 代码采用 Apache 2.0。
|
|
||||||
- BEN2 代码与模型权重采用 MIT License。
|
|
||||||
- 本 Extension 代码采用 GPL-3.0-or-later。
|
|
||||||
- 模型权重许可证由所选 Hugging Face 模型仓库决定,Extension 下载前会在界面显示。
|
|
||||||
|
|||||||
394
docs/technical.md
Normal file
394
docs/technical.md
Normal file
@@ -0,0 +1,394 @@
|
|||||||
|
# SnapForm 技术文档
|
||||||
|
|
||||||
|
本文档面向开发者和需要了解内部实现细节的用户,从原 README 迁移而来。普通用户请阅读项目根目录的
|
||||||
|
`README.md`。
|
||||||
|
|
||||||
|
`SnapForm` 是一套从图片快速构建空间结构的 Blender Extension。目前以
|
||||||
|
[Depth Anything 3](https://github.com/ByteDance-Seed/Depth-Anything-3)
|
||||||
|
作为首个 AI 能力,在独立 Runtime 中运行 DA3,可将官方导出结果保存在独立 job 目录,
|
||||||
|
也可把 Image Empty 转换成带几何与材质置换的立体面片。Debug 面板还可运行
|
||||||
|
[BEN2](https://github.com/PramaLLC/BEN2) 背景剔除功能测试。
|
||||||
|
|
||||||
|
## 当前功能
|
||||||
|
|
||||||
|
- 使用 Blender Extension manifest,可通过 `Install from Disk` 安装
|
||||||
|
- 在独立 Python Runtime 的常驻 HTTP AI Host 中运行 Torch、DA3 与 BEN2,避免污染 Blender Python 进程
|
||||||
|
- Blender 使用内置 `urllib.request` 访问仅监听 `127.0.0.1` 的 Host;请求带随机 Token
|
||||||
|
- 常驻 Host 复用 Torch、CUDA、BEN2 与最近使用的 DA3 模型,连续执行时不再重复加载
|
||||||
|
- 多个 Blender 实例自动连接同一个共享 Host;`Host` 面板显示状态、端口、客户端与模型缓存,并提供安装、启动、清理、重启和停止控制
|
||||||
|
- Runtime、模型和任务结果默认保存在 `~/.snapform`,可由多个 Blender 版本共享
|
||||||
|
- 支持 DA3 Small、Base、Large 1.1、Giant 1.1 和 Nested Giant Large 1.1
|
||||||
|
- 支持 BEN2 Base 单图背景剔除测试,输出透明 PNG 与推理耗时
|
||||||
|
- 支持单张 JPG、JPEG、PNG、WEBP 图片或常见视频文件输入
|
||||||
|
- 视频按顺序截取前若干帧,可设置参与重建的最大帧数
|
||||||
|
- `Host > Models` 默认折叠,展开后可逐项查看和下载所有可用模型
|
||||||
|
- Runtime 已就绪时不额外占用面板空间;模型列表每个模型只占一行
|
||||||
|
- `Debug` 默认折叠,内部的 DA3 与 BEN2 测试项也分别默认折叠
|
||||||
|
- 使用一个导出类型下拉菜单和一个 `Export` 按钮
|
||||||
|
- 只显示当前格式的专属参数;GLB 与 PLY 设置不会同时出现
|
||||||
|
- Debug 的模型选项按格式过滤;PLY 只提供两个带 Gaussian head 的 Giant 模型
|
||||||
|
- Debug 与 Convert to Depth Plane 的模型选项直接显示下载状态;未下载模型带
|
||||||
|
`(Missing)`,已下载模型保持原名
|
||||||
|
- 分别调用 DA3 的 `glb`、`gs_ply`、`depth_vis` 和 `mini_npz` exporter
|
||||||
|
- Image Empty 右键菜单可直接运行 DA3,根据 `mini_npz` 深度生成面片
|
||||||
|
- Image Empty 右键菜单可直接运行 BEN2,保留原图并换用新建的透明背景图片
|
||||||
|
- Image Empty 支持视频与序列帧输入:`Convert to Depth Plane` 逐帧估计深度,
|
||||||
|
`Remove Background` 逐帧抠图,入口保持不变
|
||||||
|
- 视频生成的序列帧以外部文件引用保留在 job 目录(Blender 无法整体 Pack 序列帧);
|
||||||
|
Displace 修改器与材质 Image Texture 均可按场景帧播放序列帧
|
||||||
|
- 使用浮点深度贴图、Simple Subdivision、Displace Modifier 和原图 Bump
|
||||||
|
构建立体面片
|
||||||
|
- 转换对话框可选择推理模式、处理分辨率和面片网格细分
|
||||||
|
- Debug 的 DA3 导出或 BEN2 背景剔除完成后只打开对应 job 文件夹,不向当前 Blender 文件添加任何数据
|
||||||
|
- 完整保留每次任务生成的所有官方文件
|
||||||
|
- 安装和模型下载使用后台子进程;推理提交给常驻 AI Host,状态栏实时显示阶段、百分比进度条与取消按钮
|
||||||
|
- 后台任务运行时,Install、Download 和 Export 按钮保持原位置并变灰,任务完成后恢复
|
||||||
|
|
||||||
|
## 使用
|
||||||
|
|
||||||
|
1. 使用 `uv run python pack.py` 分别构建 `windows-x64`、`macos-x64`、`macos-arm64`、
|
||||||
|
`linux-x64` 四个平台的 Extension ZIP(如 `SnapForm.v0.10.10.macos-arm64.zip`)。
|
||||||
|
每个包只含对应平台的 uv 二进制,并内置 DA3、BEN2 源码包。
|
||||||
|
2. 在 Blender 的 `Preferences > Extensions` 中选择 `Install from Disk`。
|
||||||
|
3. 在 Extension Preferences 中确认 `uv` 路径和 Torch Backend。
|
||||||
|
4. 在 `Host` 点击 `Install Runtime`。首次安装会下载 Python、PyTorch、FastAPI、Uvicorn;
|
||||||
|
DA3、BEN2 源码包与 uv 已内置在扩展中,安装时从 `~/.snapform/sources/` 解包使用,
|
||||||
|
并默认下载 `DA3 Small`、`DA3 Base` 和 `BEN2 Base` 模型。安装按 8 个阶段依次推进,每段进度
|
||||||
|
从 0 到 100%:解包内置 uv、创建 Python Runtime、下载 PyTorch、安装 xFormers(macOS 跳过,
|
||||||
|
内置 DA3 源码包为官方原版并已移除 xformers 硬依赖,自动回退纯 PyTorch 实现)、安装 AI 依赖,
|
||||||
|
然后是三个模型的下载。进度文字统一为 `阶段/总数 阶段名: 状态`,
|
||||||
|
模型下载阶段会实时显示下载速度,例如 `6/8 da3 small: downloading 14.2MB/s`。
|
||||||
|
Runtime 就绪后即使模型下载中断也不影响 Host 使用,可稍后在 `Models` 面板单独补下。
|
||||||
|
安装不做启动网络预检:若网络无法稳定访问 pypi.org,可在扩展偏好设置的 `PyPI Index`
|
||||||
|
中指定镜像源(如清华、阿里),该地址会通过 `--pypi-index` 传给 bootstrap 并优先使用;
|
||||||
|
未指定时,pip 安装失败会自动重试一次清华镜像,模型下载由 worker 自行探测回退到 HF 镜像。
|
||||||
|
5. 在 3D View 按 `N`,打开 `SnapForm` 页签。
|
||||||
|
6. 如需其他模型,展开 `Host > Models`,找到模型并点击该行的 `Download`。
|
||||||
|
7. 展开默认收起的 `Debug`,选择 `Export Format` 和本次执行使用的模型。
|
||||||
|
8. 选择一张图片或一个视频文件。视频可设置从头截取的帧数,然后点击 `Export`。
|
||||||
|
9. 测试背景剔除时,确认 `Host > Models` 中的 `BEN2 Base` 已就绪,然后在 `Debug` 的
|
||||||
|
`BEN2 Background Removal` 区域选择单张图片并点击 `Remove Background`。
|
||||||
|
完成后会打开 job 文件夹,其中的 `foreground.png` 带透明 Alpha;状态栏显示
|
||||||
|
BEN2 推理、模型加载与任务总耗时。当前测试固定关闭较慢的 foreground refinement。
|
||||||
|
|
||||||
|
## Image Empty 去背景
|
||||||
|
|
||||||
|
1. 确认 `Host > Models` 中的 `BEN2 Base` 已就绪;首次安装 Runtime 时会默认下载。
|
||||||
|
2. 在 3D View 选中 Image Empty,右键选择 `Remove Background`。
|
||||||
|
3. 完成后以 Fake User 保留原 Image datablock,新建并 Pack 一张透明 PNG 到当前 `.blend`,然后只把
|
||||||
|
当前 Empty 的 `data` 换成新图;Empty 的 Transform、显示尺寸、Offset 及其他设置保持不变。
|
||||||
|
4. 新 Image 使用原图去掉格式后再追加 `.foreground` 的名称。例如 `img.png` 生成
|
||||||
|
`img.foreground`,`img.png.001` 生成 `img.001.foreground`。已有同名 datablock 时由
|
||||||
|
Blender 自动追加编号。
|
||||||
|
|
||||||
|
原图已 Pack 时会先写出临时 PNG 交给 BEN2,任务结束后删除临时输入;原 Image 仍保持 Pack
|
||||||
|
状态和原始内容。BEN2 的 `foreground.png` 也会保留在对应 job 目录中。
|
||||||
|
|
||||||
|
Image Empty 引用视频或序列帧时,Remove Background 直接执行,默认取 Empty 的 Image Frames
|
||||||
|
数值(未设置时为完整帧数),起始帧跟随 Empty 的 Offset;若 Offset 加 Frames 超出视频长度,
|
||||||
|
会把窗口整体收进视频内(保证实际处理 Frames 帧数),Offset 过小时从第 1 帧开始。输出
|
||||||
|
`foreground.0001.png` 等透明序列帧(帧号从 0001 开始,与 Blender 场景帧一一对应),并把
|
||||||
|
源帧以 `frame.0001.jpg` 等命名复制到 job 目录根下供材质颜色导入。替换 Empty 的 Image 为
|
||||||
|
加载这些帧的 SEQUENCE Image,从原 Empty 的 Start 位置开始播放(Start 原样保留),
|
||||||
|
Cyclic 沿用原 Empty 的设置。序列帧无法被 Blender 整体 Pack,因此保持外部引用 job 目录;
|
||||||
|
请保留对应 job 目录,否则序列会失效。
|
||||||
|
|
||||||
|
## Image Empty 转深度面片
|
||||||
|
|
||||||
|
1. 在 3D View 选中 `Object Type = Empty`、`Display As = Image` 的对象。
|
||||||
|
2. 右键选择 `Convert to Depth Plane`,通过 `Depth from` 选择深度来源:
|
||||||
|
- `AI Generated`:显示 Model、Use Ray Pose 和 Resolution Scale。实际 Process Res 为图片长边像素数
|
||||||
|
乘以 Resolution Scale;范围为 `0.1–1.0`,默认 `1.0`。计算后的 Process Res 会直接显示在
|
||||||
|
`Resolution Scale (… px)` 标签中。
|
||||||
|
转换导出 `mini_npz`,使用前需要在
|
||||||
|
`Host > Models` 下载所选模型,并固定使用 `process_res_method="upper_bound_resize"`;设备和参考视图策略
|
||||||
|
继续使用场景设置。Model 下拉菜单只列出已经下载的 DA3 模型;若没有可用模型,对话框会引导到
|
||||||
|
`Host > Models` 下载。
|
||||||
|
- `Original`:直接把原图作为深度贴图,不显示估算参数。单张图片不启动 Runtime;视频或
|
||||||
|
序列帧会启动轻量抽帧任务,把帧抽出来作为 `.depth` 序列(比直接引用视频稳定)。
|
||||||
|
`Original` 是默认深度来源。选择 `AI Generated` 时,其 Model、Use Ray Pose 与 Resolution Scale
|
||||||
|
参数会集中显示在独立边框中。两种模式都有默认折叠的 Options,内部依次为 Displace via、
|
||||||
|
Subdivision、Remove Background 和 Unshaded;Delete Original 保持在折叠项外。
|
||||||
|
Remove Background 默认关闭,与 Depth from 无关;开启后需要下载 `BEN2 Base`,并从原图生成、
|
||||||
|
Pack `.foreground` 透明图片。它只替换生成面片材质的颜色图片,不修改源 Image Empty 的
|
||||||
|
Image datablock。AI Generated 仍从原图生成 `.depth`;例如 `img.png` 会得到
|
||||||
|
`img.foreground` 与 `img.depth`,不会生成 `img.foreground.depth`。Original 仍用原图作为
|
||||||
|
位移来源,不改变原有深度逻辑。
|
||||||
|
Image Empty 引用视频或序列帧时,AI Generated 会按 Empty 的 Image Frames 数值逐帧估计
|
||||||
|
深度(未设置时为完整帧数),起始帧跟随 Empty 的 Offset,与 Remove Background 范围一致;
|
||||||
|
Offset 超出视频长度时同样把窗口收进视频内。材质节点与置换贴图的 Start 保持与原 Empty
|
||||||
|
一致;Original 模式(无 Remove Background)时,单帧会复制原图改名为 `.depth` 供位移使用,
|
||||||
|
视频或序列帧则抽帧保存为 job 目录根下的 `frame.0001.jpg` 等,并作为 `.depth` 序列使用。
|
||||||
|
材质颜色使用这些帧导入的 `.frame` 序列(Frames/Start/Offset 同步),开启 Remove
|
||||||
|
Background 时改用 `foreground.0001.png` 序列。AI Generated 的深度以 EXR 序列帧
|
||||||
|
(`depth.0001.exr` 起)保存在 job 目录根下;Original 模式(含 Remove Background)的位移
|
||||||
|
使用 `frame.*` 序列帧,保证逐帧对齐。
|
||||||
|
Displace via 可选 `Modifier`(默认)或 `Material`;
|
||||||
|
Unshaded 默认关闭;Subdivision 默认为 `5`,并默认删除原对象。
|
||||||
|
3. AI Generated 完成后会自动创建面片。面片几何先由 `mini_npz` 的相机内参与深度确定,再通过对象
|
||||||
|
Transform 匹配原 Image Empty 的可见尺寸、位置、旋转和 Offset。Original 模式直接使用源 Empty
|
||||||
|
的可见尺寸、Offset 与 Transform。开启 `Delete Original` 时,成功生成后
|
||||||
|
删除源 Empty。生成面片 Object 使用源 Empty 的原名;关闭 `Delete Original` 时不处理
|
||||||
|
同名冲突,由 Blender 自动为生成面片追加编号。
|
||||||
|
4. 短边分段数为 `2^Subdivision`,长边分段数根据图像宽高比取最接近的整数。例如高:宽为
|
||||||
|
`2:3`、Subdivision 为 `1` 时,直接生成 `2×3` 个面;Subdivision 为 `3` 时生成
|
||||||
|
`8×12` 个面。随后的 Simple Subdivision 修改器使用原生名称,Viewport Levels 固定为
|
||||||
|
`1`,Render Levels 固定为 `2`,不再随 Resolution Scale 或 Subdivision 动态变化。
|
||||||
|
|
||||||
|
所有任务使用统一的状态栏进度显示。变化中的阶段说明显示在右侧的紧凑进度条内部,取消按钮
|
||||||
|
位于进度条右侧。Image Depth 的阶段包括 Runtime 启动、输入准备、模型加载、深度估算、缓存
|
||||||
|
写出和面片构建。进度更新会强制刷新状态栏与窗口,无需移动鼠标即可实时显示。安装中断后残留的
|
||||||
|
半成品 Runtime 会在下次安装前自动清理,不会因为中断而无法再次安装。
|
||||||
|
|
||||||
|
导入过程精确保留 NPZ 内深度值之间的差值。例如深度 `2.0` 和 `5.5` 在
|
||||||
|
`Displacement Strength = -1` 时,沿局部 Z 轴相差 `3.5 Blender Unit`。它不会把数值压缩到
|
||||||
|
`0..1`,也不再提供视觉归一化模式。
|
||||||
|
|
||||||
|
面片尺寸以有效深度的中位数 `Zref` 作为参考截面。先用 `mini_npz` 中的相机内参计算原始
|
||||||
|
投影面积,再恢复源图纵横比:
|
||||||
|
|
||||||
|
```text
|
||||||
|
RawWidth = depth_width × Zref × abs(Displacement Strength) / fx
|
||||||
|
RawHeight = depth_height × Zref × abs(Displacement Strength) / fy
|
||||||
|
Area = RawWidth × RawHeight
|
||||||
|
Aspect = source_image_width / source_image_height
|
||||||
|
Width = sqrt(Area × Aspect)
|
||||||
|
Height = sqrt(Area / Aspect)
|
||||||
|
```
|
||||||
|
|
||||||
|
这样即使 DA3 估计出的 `fx/fy` 不一致,也不会把正方形原图拉成长方形。主点 `cx/cy` 仍用于
|
||||||
|
确定原始相机面片,随后对象 XY Transform 将它匹配到源 Image Empty。生成对象会记录
|
||||||
|
`snapform_reference_depth`、`snapform_camera_plane_width` 和
|
||||||
|
`snapform_camera_plane_height`,便于检查实际采用的
|
||||||
|
参考深度和尺寸。
|
||||||
|
|
||||||
|
选择 `Modifier` 时,AI Generated 的 Displace Modifier 深度贴图使用 `Extension = Extend`,
|
||||||
|
`Strength` 固定为 `-1`、`Midlevel` 固定为 `1`。Original 把原图指定给 Displace Texture,并设置
|
||||||
|
`Texture Coordinates = UV`、`Midlevel = 0`;其余修改器设置保留 Blender 默认值。
|
||||||
|
视频与序列帧生成的深度贴图是 SEQUENCE Image;Displace 修改器的 Texture 与材质的 Image Texture
|
||||||
|
节点都支持序列帧,通过 image_user 的 Frames / Offset / Start 设置把第 i 帧映射到场景
|
||||||
|
`Start+i` 帧(Start 沿用 Empty 的 Start),并启用 Auto Refresh,播放时随场景帧逐帧更新。
|
||||||
|
选择 `Material` 时,不创建 Displace Modifier;深度 Image Texture 的 Color 连接到 Displacement
|
||||||
|
节点的 Height,再连接 Material Output 的 Displacement,并把材质 Surface Displacement 设置为
|
||||||
|
`Displacement Only`。AI Generated 使用 `Midlevel = 1`、`Scale = -1`;Original 使用
|
||||||
|
`Midlevel = 0`、`Scale = 1`,与对应修改器值一致。无论选择哪种呈现方式,原始彩色贴图都会连接
|
||||||
|
Principled BSDF 的 `Base Color`,并通过
|
||||||
|
`Bump.Height → Bump.Normal → Principled.Normal` 提供表面细节。Bump 的基础 Distance 为
|
||||||
|
`0.001`,会乘以面片 Transform 的统一缩放倍率。原图 Color 还会经过 `Float Curve` 连接到
|
||||||
|
Principled BSDF 的 `Roughness`,曲线两端输出分别为 `0.5` 和 `0.1`;材质 IOR 保持
|
||||||
|
Principled BSDF 默认值 `1.5`。
|
||||||
|
开启 `Unshaded` 后,原图 Color 连接 Emission,原图 Alpha 作为 Mix Shader 系数,在
|
||||||
|
Transparent BSDF 与 Emission 之间混合后连接 Material Output;不创建 Principled BSDF、Bump 或
|
||||||
|
Float Curve。若 Displace via 选择 `Material`,深度贴图与 Displacement
|
||||||
|
节点仍会保留并连接 Material Output 的 Displacement。
|
||||||
|
|
||||||
|
这里的“原始距离”是 DA3 模型输出的尺度。普通 Small、Base、Large 和 Giant 模型的深度不应
|
||||||
|
被解释为现实世界的米;需要现实尺度时应选择带 metric depth 的
|
||||||
|
`DA3 Nested Giant Large 1.1`。此外,当前结果仍是固定 XY 平面上的深度位移面,不是使用
|
||||||
|
相机内参逐像素反投影得到的完整 camera-space 网格;中位深度截面为现有 Displace Modifier
|
||||||
|
工作流提供正确的整体高宽尺度。
|
||||||
|
|
||||||
|
生成完成后会直接创建面片,不提供缓存重建入口或单独的导入参数。重新选择
|
||||||
|
`Convert to Depth Plane` 会按当前 Depth from 再次转换。
|
||||||
|
面片、深度 Image、Texture 与 Material 会在同一个 Blender Undo 步骤中创建;转换完成后执行
|
||||||
|
`Ctrl+Z` 会完整撤销本次转换,不会越过转换去回退更早的 Image Empty 操作。
|
||||||
|
开启 Remove Background 时,`.foreground` Image 的创建、Pack 和深度面片生成也属于
|
||||||
|
同一个 Undo 步骤。
|
||||||
|
|
||||||
|
Image Empty 使用 Packed Image 时,SnapForm 会通过 Blender 写出临时 PNG 交给 DA3,
|
||||||
|
不会 Unpack 或修改源 Image;任务结束后自动删除临时文件。普通图片仍直接读取源文件。DA3 的原始
|
||||||
|
`results.npz` 保留在对应 job 文件夹根下,生成的浮点深度 Image 会 Pack 到
|
||||||
|
当前 `.blend`。Depth Image 与 Displace Texture 根据源 Image datablock 命名。例如 Empty Object
|
||||||
|
名为 `obj_img`、引用的 Image 名为 `img.png` 时,Depth Image 与 Texture 都命名为 `img.depth`;
|
||||||
|
生成的 Object、Mesh 与 Material 仍命名为 `obj_img`。Pack 使用 OpenEXR,因而不会把大于 `1.0`
|
||||||
|
的原始深度截断。已有同名 datablock 时由 Blender 自动追加 `.001` 等编号。
|
||||||
|
视频输入时 Blender 无法整体 Pack 序列帧,因此 `.depth` 与颜色序列改为引用 job 目录根下的
|
||||||
|
`depth.0001.exr`、`frame.0001.jpg`(`.frame`)或 `foreground.0001.png`,并在每帧保留原始
|
||||||
|
浮点深度;保存 `.blend` 后再次打开仍需对应 job 目录存在。
|
||||||
|
|
||||||
|
选择 PLY 后,模型菜单只会提供 `DA3 Giant 1.1` 和
|
||||||
|
`DA3 Nested Giant Large 1.1`。Small、Base 和 Large 没有 Gaussian head,无法由
|
||||||
|
DA3 生成 `gs_ply`。
|
||||||
|
|
||||||
|
每个任务在 `~/.snapform/jobs/` 下用独立文件夹包裹(命名格式为 `任务类型_时间号`,例如
|
||||||
|
`depth-plane_20260807-153000`),日志、状态与输出产物都放在该文件夹内:
|
||||||
|
|
||||||
|
```text
|
||||||
|
reconstruct_20260807-153000/
|
||||||
|
├── reconstruct_20260807-153000.log
|
||||||
|
├── reconstruct_20260807-153000.json
|
||||||
|
└── scene.glb
|
||||||
|
|
||||||
|
reconstruct_20260807-153000/
|
||||||
|
├── reconstruct_20260807-153000.log
|
||||||
|
├── reconstruct_20260807-153000.json
|
||||||
|
└── gs_ply/*.ply
|
||||||
|
|
||||||
|
reconstruct_20260807-153000/
|
||||||
|
├── reconstruct_20260807-153000.log
|
||||||
|
├── reconstruct_20260807-153000.json
|
||||||
|
└── depth_vis/*.jpg
|
||||||
|
|
||||||
|
reconstruct_20260807-153000/
|
||||||
|
├── reconstruct_20260807-153000.log
|
||||||
|
├── reconstruct_20260807-153000.json
|
||||||
|
└── results.npz
|
||||||
|
|
||||||
|
remove-background_20260807-153000/
|
||||||
|
├── remove-background_20260807-153000.log
|
||||||
|
├── remove-background_20260807-153000.json
|
||||||
|
├── foreground.png 单图结果
|
||||||
|
├── foreground.0001.png… 视频/序列帧结果
|
||||||
|
├── frame.0001.jpg… 视频/序列帧的源帧副本
|
||||||
|
└── depth.0001.exr… AI 深度序列帧
|
||||||
|
```
|
||||||
|
|
||||||
|
建议从 `DA3 Base` 开始。它使用 Apache 2.0 模型权重;`DA3 Large 1.1` 权重采用
|
||||||
|
CC BY-NC 4.0,仅限非商业用途。
|
||||||
|
|
||||||
|
## 输出格式说明
|
||||||
|
|
||||||
|
`Export Format` 下拉菜单分别把以下原始 `export_format` 传给 DA3:
|
||||||
|
|
||||||
|
- GLB:`glb`
|
||||||
|
- PLY:`gs_ply`,同时传入 `infer_gs=True`
|
||||||
|
- Depth Images:`depth_vis`
|
||||||
|
- Mini NPZ:`mini_npz`
|
||||||
|
|
||||||
|
界面使用 DA3 API 的参数名、选项和默认值。GLB 专属设置仅在选择 `GLB` 时显示,
|
||||||
|
`gs_views_interval` 仅在选择 `PLY (3DGS)` 时显示:
|
||||||
|
|
||||||
|
- `process_res=504`
|
||||||
|
- `process_res_method="upper_bound_resize"`
|
||||||
|
- `use_ray_pose=False`
|
||||||
|
- `ref_view_strategy="saddle_balanced"`
|
||||||
|
- `conf_thresh_percentile=40.0`
|
||||||
|
- `num_max_points=1_000_000`
|
||||||
|
- `show_cameras=True`
|
||||||
|
- `export_kwargs["gs_ply"]["gs_views_interval"]=1`
|
||||||
|
|
||||||
|
DA3 官方 `glb` 是彩色点云与可选相机线框;`gs_ply` 是带 Gaussian 属性的 PLY;
|
||||||
|
`depth_vis` 是原图与着色深度图的并排预览;`mini_npz` 保存深度、置信度和可用的相机参数。
|
||||||
|
GLB 任务会关闭其默认附带的深度预览,使四种 job 产物保持分离。
|
||||||
|
|
||||||
|
官方格式与限制可参阅
|
||||||
|
[DA3 Python API 文档](https://github.com/ByteDance-Seed/Depth-Anything-3/blob/main/docs/API.md)。
|
||||||
|
|
||||||
|
## Runtime
|
||||||
|
|
||||||
|
Extension 本身不包含 Torch。`Install Runtime` 使用 `uv` 创建独立 Python 3.12
|
||||||
|
环境,并安装。机器上即使没有安装 `uv` 或 Python 也没关系:安装程序由 Blender
|
||||||
|
自带的 Python 启动,缺 `uv` 时会自动下载官方独立二进制(约 30MB),再用它下载
|
||||||
|
Python 3.12 与全部依赖。`uv` 会缓存在扩展存储目录下,后续安装直接复用。
|
||||||
|
安装不做网络预检,直接开始;环境变量仍可切换镜像与代理:`PIP_INDEX_URL`(PyPI 镜像)、
|
||||||
|
`SNAPFORM_TORCH_INDEX`(PyTorch 索引)、`HF_ENDPOINT`(Hugging Face 镜像)、
|
||||||
|
`HTTP(S)_PROXY`/`ALL_PROXY`(代理)。未指定 PyPI 镜像时,
|
||||||
|
pip 安装失败会自动依次重试镜像(清华、阿里云);模型下载不预探测,官方 Hugging Face
|
||||||
|
失败后自动回退 ModelScope(`modelscope.cn`,仓库名与 Hugging Face 一致)。
|
||||||
|
GitHub 不再参与安装:DA3、BEN2、uv 已内置,缺少内置文件会直接报错提示重新安装,不再下载兜底。
|
||||||
|
单独下载模型时遵循同样的顺序:Hugging Face(可用 `HF_ENDPOINT` 指定其他镜像)失败后自动切换到
|
||||||
|
ModelScope。
|
||||||
|
Windows 上代理若只在系统设置里而未写入环境变量,安装程序会尝试自动从系统代理注入子进程。
|
||||||
|
|
||||||
|
- PyTorch 2.9.1
|
||||||
|
- TorchVision 0.24.1
|
||||||
|
- xFormers 0.0.33.post2(与 PyTorch 2.9.1 对应,仅 Windows/Linux;macOS 无 wheel 自动跳过)
|
||||||
|
- 固定到已验证 commit 的 Depth Anything 3
|
||||||
|
- 固定到已验证 commit 的 BEN2
|
||||||
|
- FastAPI 0.139.2 与 Uvicorn 0.51.0
|
||||||
|
- DA3 与 BEN2 源码包随扩展打包,安装时复制到 `~/.snapform/sources/` 后由 uv 本地安装
|
||||||
|
(无需访问 GitHub);SnapForm 直接调用 DA3
|
||||||
|
底层 export 接口,`gs_ply` 导出不会触发 `gs_video` 渲染,因此不需要 gsplat
|
||||||
|
与 C++ 编译工具链
|
||||||
|
- `sources/` 下的 DA3/BEN2 源码包保持官方原版(固定 commit);`sources/patches/<包名>/`
|
||||||
|
存放 SnapForm 对源码的覆盖文件(移除 xformers 硬依赖、export 入口改懒加载、
|
||||||
|
射线位姿的 3×3 QR 分解在 MPS 上改走 CPU),
|
||||||
|
`pack.py` 打包时自动套入官方包,扩展 zip 内才是修补后的版本;更新上游只需替换官方包
|
||||||
|
并核对覆盖文件是否仍适用
|
||||||
|
- `addict` 2.4.0(补充 DA3 上游包元数据暂未声明的运行依赖)
|
||||||
|
|
||||||
|
默认 CUDA Backend 为 CUDA 12.8,也可以选择 CPU。CPU 可以用于验证,但实际推理
|
||||||
|
通常很慢。
|
||||||
|
CUDA 12.8 的 PyTorch wheel 自带 CUDA 运行库,无需安装 CUDA Toolkit;只需要
|
||||||
|
NVIDIA 驱动 ≥ 570(数据中心显卡可用更老驱动)。驱动不满足或没有 NVIDIA 显卡时,
|
||||||
|
设备会自动回退到 CPU(macOS 为 MPS),只有显式选择 CUDA 才会报错。
|
||||||
|
|
||||||
|
macOS(Apple Silicon)会自动从 PyPI 安装 CPU 版 PyTorch,推理设备默认使用 MPS
|
||||||
|
(Apple Metal)加速,也可以在设备选项里切换 CPU。PyTorch 自 2.3 起不再提供 Intel
|
||||||
|
Mac(x86_64)轮子,Intel Mac 暂不支持。DA3 射线位姿估计中的 3×3 QR 分解在 MPS
|
||||||
|
上尚未实现,SnapForm 的覆盖文件会单独把该分解放到 CPU 执行,其余推理仍走 MPS。
|
||||||
|
|
||||||
|
Extension 启用后只会查找已有的共享 AI Host,不会主动启动新进程。第一次执行 AI 操作或在侧边栏
|
||||||
|
手动点击 `Start` 时才会启动 Host。Host 只绑定回环地址,并让操作系统分配动态端口;第一个实例把
|
||||||
|
端口、进程和随机 Token 写入 `host.json`,其余 Blender 通过标准库
|
||||||
|
`urllib.request` 连接同一个 Host。`host.lock` 保证同一 Storage Root 同时只有一个 Host,Token
|
||||||
|
则阻止同机其他进程直接调用接口。Host 全局同时只执行一个 AI 任务。
|
||||||
|
|
||||||
|
每个 Blender 会注册独立客户端并每 5 秒发送心跳。正常关闭 Blender 时先注销自身;如果当前 Host
|
||||||
|
就是由这个 Blender 启动的,会直接终止它,无需等待。异常退出的客户端会在 10 秒后过期;Host
|
||||||
|
在最后一个客户端离开后仍会保留模型缓存等待 20 秒再退出,作为 Host 由其他进程启动等场景的兜底。
|
||||||
|
取消任务不会重启 Host,并会保留已经加载的模型;
|
||||||
|
如果取消发生在一次不可中断的 Torch 调用中,任务会在该调用返回后的最近检查点结束。
|
||||||
|
|
||||||
|
侧边栏 `Host` 面板以 `Running (Clients: n)` 表示空闲服务,以 `Running (Busy)` 表示正在执行任务,
|
||||||
|
并显示缓存模型和活动任务。空闲且只有当前 Blender 连接时,可以清理模型缓存、重启或停止 Host;
|
||||||
|
多客户端或任务执行期间不会允许破坏共享服务。
|
||||||
|
|
||||||
|
BEN2 在首次使用后保持加载;DA3 保留最近使用的一个模型。切换 DA3 模型时会释放旧模型;加载时如果
|
||||||
|
显存不足,会先释放另一类缓存模型并重试。
|
||||||
|
|
||||||
|
Runtime Schema 已升级,旧 Runtime 需要重新执行 `Install Runtime` 才能安装 HTTP Host 依赖。
|
||||||
|
Host 的启动和崩溃日志写入 `~/.snapform/ai-host.log`。
|
||||||
|
|
||||||
|
Runtime、模型和任务结果默认保存在用户主目录的 `~/.snapform`:
|
||||||
|
|
||||||
|
```text
|
||||||
|
~/.snapform/
|
||||||
|
├── runtime/
|
||||||
|
├── models/
|
||||||
|
├── jobs/
|
||||||
|
├── ai-host.log
|
||||||
|
├── host.json # Host 运行时存在
|
||||||
|
└── host.lock # Host 运行时持有
|
||||||
|
```
|
||||||
|
|
||||||
|
不同 Blender 版本安装相同 Extension 后会共用这个目录。也可以在 Extension
|
||||||
|
Preferences 的 `Storage Root` 中修改共享位置。不要把模型写入 Extension 安装目录,
|
||||||
|
因为该目录可能只读,并且升级时会被替换。不要同时从多个 Blender 实例执行
|
||||||
|
`Install Runtime`;推理 Host、任务队列和模型读取可以共享。
|
||||||
|
|
||||||
|
## 开发
|
||||||
|
|
||||||
|
```bash
|
||||||
|
uv run python -m unittest discover -s tests -v
|
||||||
|
uv run python pack.py
|
||||||
|
```
|
||||||
|
|
||||||
|
项目结构:
|
||||||
|
|
||||||
|
```text
|
||||||
|
.
|
||||||
|
├── pack.py
|
||||||
|
├── sources/ 官方 DA3/BEN2 源码包、patches 覆盖文件与各平台 uv 二进制
|
||||||
|
├── src/
|
||||||
|
│ └── snapform/
|
||||||
|
│ ├── blender_manifest.toml
|
||||||
|
│ ├── ai_host.py
|
||||||
|
│ ├── bootstrap_runtime.py
|
||||||
|
│ ├── models.py
|
||||||
|
│ ├── image_depth.py
|
||||||
|
│ ├── operators.py
|
||||||
|
│ ├── panel.py
|
||||||
|
│ ├── preferences.py
|
||||||
|
│ ├── properties.py
|
||||||
|
│ ├── runtime.py
|
||||||
|
│ └── worker.py
|
||||||
|
└── tests/
|
||||||
|
```
|
||||||
|
|
||||||
|
## 上游与许可证
|
||||||
|
|
||||||
|
- Depth Anything 3 代码采用 Apache 2.0。
|
||||||
|
- BEN2 代码与模型权重采用 MIT License。
|
||||||
|
- 本 Extension 代码采用 GPL-3.0-or-later。
|
||||||
|
- 模型权重许可证由所选 Hugging Face 模型仓库决定,Extension 下载前会在界面显示。
|
||||||
125
pack.py
125
pack.py
@@ -1,3 +1,5 @@
|
|||||||
|
import io
|
||||||
|
import tarfile
|
||||||
import tomllib
|
import tomllib
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from zipfile import ZIP_DEFLATED, ZipFile
|
from zipfile import ZIP_DEFLATED, ZipFile
|
||||||
@@ -5,30 +7,135 @@ from zipfile import ZIP_DEFLATED, ZipFile
|
|||||||
|
|
||||||
PROJECT_ROOT = Path(__file__).resolve().parent
|
PROJECT_ROOT = Path(__file__).resolve().parent
|
||||||
SOURCE_DIR = PROJECT_ROOT / "src" / "snapform"
|
SOURCE_DIR = PROJECT_ROOT / "src" / "snapform"
|
||||||
|
SOURCES_DIR = PROJECT_ROOT / "sources"
|
||||||
DIST_DIR = PROJECT_ROOT / "dist"
|
DIST_DIR = PROJECT_ROOT / "dist"
|
||||||
|
PATCHES_DIR = SOURCES_DIR / "patches"
|
||||||
|
|
||||||
|
PLATFORMS = {
|
||||||
|
"windows-x64": "uv-x86_64-pc-windows-msvc.zip",
|
||||||
|
"macos-x64": "uv-x86_64-apple-darwin.tar.gz",
|
||||||
|
"macos-arm64": "uv-aarch64-apple-darwin.tar.gz",
|
||||||
|
"linux-x64": "uv-x86_64-unknown-linux-gnu.tar.gz",
|
||||||
|
}
|
||||||
|
BUNDLED_ARCHIVES = {
|
||||||
|
"sources/depth-anything-3.tar.gz": "depth-anything-3.tar.gz",
|
||||||
|
"sources/ben2.tar.gz": "ben2.tar.gz",
|
||||||
|
}
|
||||||
|
OVERLAYS = {
|
||||||
|
"depth-anything-3.tar.gz": "depth-anything-3",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
def build_extension():
|
def bundled_source_bytes(filename):
|
||||||
manifest = tomllib.loads(
|
"""返回内置源码包内容;patches/<overlay>/ 目录中的文件会覆盖官方包对应路径。"""
|
||||||
(SOURCE_DIR / "blender_manifest.toml").read_text(encoding="utf-8")
|
archive_path = SOURCES_DIR / filename
|
||||||
|
overlay_dir = PATCHES_DIR / OVERLAYS.get(filename, filename)
|
||||||
|
if not overlay_dir.is_dir():
|
||||||
|
return archive_path.read_bytes()
|
||||||
|
if not filename.endswith(".tar.gz"):
|
||||||
|
raise ValueError(f"Unsupported overlay archive: {filename}")
|
||||||
|
return _apply_tar_overlay(archive_path, overlay_dir)
|
||||||
|
|
||||||
|
|
||||||
|
def _apply_tar_overlay(archive_path, overlay_dir):
|
||||||
|
with tarfile.open(archive_path, "r:gz") as source:
|
||||||
|
members = source.getmembers()
|
||||||
|
top_level = members[0].name.split("/", 1)[0]
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
with tarfile.open(fileobj=buffer, mode="w:gz") as output:
|
||||||
|
for member in members:
|
||||||
|
if member.isdir() or member.islnk() or member.issym():
|
||||||
|
output.addfile(member)
|
||||||
|
continue
|
||||||
|
relative = member.name[len(top_level) + 1:]
|
||||||
|
overlay_file = overlay_dir / relative
|
||||||
|
if overlay_file.is_file():
|
||||||
|
info = tarfile.TarInfo(member.name)
|
||||||
|
info.size = overlay_file.stat().st_size
|
||||||
|
info.mode = member.mode
|
||||||
|
output.addfile(info, overlay_file.open("rb"))
|
||||||
|
else:
|
||||||
|
with source.extractfile(member) as extracted:
|
||||||
|
output.addfile(member, extracted)
|
||||||
|
return buffer.getvalue()
|
||||||
|
|
||||||
|
|
||||||
|
def manifest_for_platform(platform):
|
||||||
|
"""Return the manifest text declaring only the given platform."""
|
||||||
|
text = (SOURCE_DIR / "blender_manifest.toml").read_text(encoding="utf-8")
|
||||||
|
lines = text.splitlines()
|
||||||
|
table_start = next(
|
||||||
|
(i for i, line in enumerate(lines) if line.startswith("[")),
|
||||||
|
len(lines),
|
||||||
)
|
)
|
||||||
|
head = lines[:table_start]
|
||||||
|
tail = lines[table_start:]
|
||||||
|
head = [line for line in head if not line.strip().startswith("platforms =")]
|
||||||
|
# Insert into the top-level section (before any `[table]` header) so the
|
||||||
|
# field is part of the extension metadata, not a sub-key of a table.
|
||||||
|
# Place it after the `blender_version_min` line, alongside the other
|
||||||
|
# top-level keys (before `license`), keeping the file readable.
|
||||||
|
insert_at = next(
|
||||||
|
(i + 1 for i, line in enumerate(head) if line.strip().startswith("blender_version_min")),
|
||||||
|
len(head),
|
||||||
|
)
|
||||||
|
head.insert(insert_at, f'platforms = ["{platform}"]')
|
||||||
|
lines = head + tail
|
||||||
|
return "\n".join(lines) + "\n"
|
||||||
|
|
||||||
|
|
||||||
|
def bundled_source_files(platform):
|
||||||
|
entries = {
|
||||||
|
name: SOURCES_DIR / filename
|
||||||
|
for name, filename in BUNDLED_ARCHIVES.items()
|
||||||
|
}
|
||||||
|
entries[f"sources/{PLATFORMS[platform]}"] = (
|
||||||
|
SOURCES_DIR / PLATFORMS[platform]
|
||||||
|
)
|
||||||
|
return entries
|
||||||
|
|
||||||
|
|
||||||
|
def build_extension(platform):
|
||||||
|
if platform not in PLATFORMS:
|
||||||
|
raise ValueError(f"Unsupported platform: {platform}")
|
||||||
|
manifest = tomllib.loads(manifest_for_platform(platform))
|
||||||
output_path = DIST_DIR / (
|
output_path = DIST_DIR / (
|
||||||
f"{manifest['name']}.v{manifest['version']}.zip"
|
f"{manifest['name']}.v{manifest['version']}.{platform}.zip"
|
||||||
)
|
)
|
||||||
DIST_DIR.mkdir(exist_ok=True)
|
DIST_DIR.mkdir(exist_ok=True)
|
||||||
|
sources = bundled_source_files(platform)
|
||||||
|
missing = [str(path) for path in sources.values() if not path.is_file()]
|
||||||
|
if missing:
|
||||||
|
raise RuntimeError(
|
||||||
|
"Missing bundled source files: " + ", ".join(missing)
|
||||||
|
)
|
||||||
with ZipFile(output_path, "w", ZIP_DEFLATED) as archive:
|
with ZipFile(output_path, "w", ZIP_DEFLATED) as archive:
|
||||||
for source_path in sorted(SOURCE_DIR.rglob("*")):
|
for source_path in sorted(SOURCE_DIR.rglob("*")):
|
||||||
if not source_path.is_file():
|
if not source_path.is_file():
|
||||||
continue
|
continue
|
||||||
if "__pycache__" in source_path.parts:
|
if "__pycache__" in source_path.parts:
|
||||||
continue
|
continue
|
||||||
|
if source_path.name == "blender_manifest.toml":
|
||||||
|
continue
|
||||||
archive.write(source_path, source_path.relative_to(SOURCE_DIR))
|
archive.write(source_path, source_path.relative_to(SOURCE_DIR))
|
||||||
|
archive.writestr(
|
||||||
|
"blender_manifest.toml",
|
||||||
|
manifest_for_platform(platform),
|
||||||
|
)
|
||||||
|
for archive_name, source_path in sources.items():
|
||||||
|
archive.writestr(
|
||||||
|
archive_name,
|
||||||
|
bundled_source_bytes(source_path.name),
|
||||||
|
)
|
||||||
return output_path
|
return output_path
|
||||||
|
|
||||||
|
|
||||||
|
def build_all():
|
||||||
|
return [build_extension(platform) for platform in PLATFORMS]
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
extension_path = build_extension()
|
extension_paths = build_all()
|
||||||
size_mib = extension_path.stat().st_size / (1024 * 1024)
|
for extension_path in extension_paths:
|
||||||
print(f"Built {extension_path} ({size_mib:.2f} MiB)")
|
size_mib = extension_path.stat().st_size / (1024 * 1024)
|
||||||
|
print(f"Built {extension_path} ({size_mib:.2f} MiB)")
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
[project]
|
[project]
|
||||||
name = "snapform"
|
name = "snapform"
|
||||||
version = "0.9.32"
|
version = "0.10.10"
|
||||||
description = "Fast image-guided structure building tools for Blender"
|
description = "Fast image-guided structure building tools for Blender"
|
||||||
readme = "README.md"
|
readme = "README.md"
|
||||||
requires-python = ">=3.11"
|
requires-python = ">=3.11"
|
||||||
|
|||||||
@@ -278,6 +278,11 @@ class HostState:
|
|||||||
model_cache=self.models,
|
model_cache=self.models,
|
||||||
cancel_check=cancel_check,
|
cancel_check=cancel_check,
|
||||||
)
|
)
|
||||||
|
elif payload["command"] == "extract-frames":
|
||||||
|
worker.extract_frames(
|
||||||
|
args,
|
||||||
|
cancel_check=cancel_check,
|
||||||
|
)
|
||||||
else:
|
else:
|
||||||
worker.reconstruct(
|
worker.reconstruct(
|
||||||
args,
|
args,
|
||||||
@@ -426,7 +431,11 @@ def create_app(token, instance_id="test"):
|
|||||||
client_id = str(payload.get("client_id", ""))
|
client_id = str(payload.get("client_id", ""))
|
||||||
command = payload.get("command")
|
command = payload.get("command")
|
||||||
parameters = payload.get("parameters")
|
parameters = payload.get("parameters")
|
||||||
if not job_id or command not in {"remove-background", "reconstruct"}:
|
if not job_id or command not in {
|
||||||
|
"remove-background",
|
||||||
|
"reconstruct",
|
||||||
|
"extract-frames",
|
||||||
|
}:
|
||||||
raise HTTPException(status_code=422, detail="Invalid job request")
|
raise HTTPException(status_code=422, detail="Invalid job request")
|
||||||
if not isinstance(parameters, dict) or not parameters.get("status"):
|
if not isinstance(parameters, dict) or not parameters.get("status"):
|
||||||
raise HTTPException(status_code=422, detail="Missing job parameters")
|
raise HTTPException(status_code=422, detail="Missing job parameters")
|
||||||
|
|||||||
@@ -1,7 +1,7 @@
|
|||||||
schema_version = "1.0.0"
|
schema_version = "1.0.0"
|
||||||
|
|
||||||
id = "snapform"
|
id = "snapform"
|
||||||
version = "0.9.32"
|
version = "0.10.10"
|
||||||
name = "SnapForm"
|
name = "SnapForm"
|
||||||
tagline = "Build spatial structures from images, fast"
|
tagline = "Build spatial structures from images, fast"
|
||||||
maintainer = "OmooLab"
|
maintainer = "OmooLab"
|
||||||
|
|||||||
@@ -20,21 +20,13 @@ if str(SCRIPT_DIRECTORY) not in sys.path:
|
|||||||
sys.path.insert(0, str(SCRIPT_DIRECTORY))
|
sys.path.insert(0, str(SCRIPT_DIRECTORY))
|
||||||
|
|
||||||
from models import BEN2_MODEL, MODELS
|
from models import BEN2_MODEL, MODELS
|
||||||
from worker import HF_ENDPOINT, HF_MIRRORS, format_rate, reachable
|
|
||||||
|
|
||||||
|
|
||||||
UV_VERSION = "0.11.3"
|
|
||||||
PREFLIGHT_MIN_FREE_BYTES = 8 * 1024**3
|
|
||||||
PREFLIGHT_TIMEOUT = 8.0
|
|
||||||
MIRROR_TIMEOUT = 5.0
|
|
||||||
PYPI_MIRRORS = (
|
PYPI_MIRRORS = (
|
||||||
"https://pypi.tuna.tsinghua.edu.cn/simple",
|
"https://pypi.tuna.tsinghua.edu.cn/simple",
|
||||||
"https://mirrors.aliyun.com/pypi/simple/",
|
"https://mirrors.aliyun.com/pypi/simple/",
|
||||||
)
|
)
|
||||||
TORCH_MIRRORS = {
|
PYPI_DEFAULT_INDEX = "https://pypi.org/simple/"
|
||||||
"cu128": ("https://mirrors.aliyun.com/pytorch-wheels/cu128",),
|
|
||||||
"cpu": ("https://mirrors.aliyun.com/pytorch-wheels/cpu",),
|
|
||||||
}
|
|
||||||
TORCH_VERSION = "2.9.1"
|
TORCH_VERSION = "2.9.1"
|
||||||
TORCHVISION_VERSION = "0.24.1"
|
TORCHVISION_VERSION = "0.24.1"
|
||||||
XFORMERS_VERSION = "0.0.33.post2"
|
XFORMERS_VERSION = "0.0.33.post2"
|
||||||
@@ -47,11 +39,7 @@ TORCH_INDEXES = {
|
|||||||
"cpu": "https://download.pytorch.org/whl/cpu",
|
"cpu": "https://download.pytorch.org/whl/cpu",
|
||||||
}
|
}
|
||||||
_resolved_pypi_index = None
|
_resolved_pypi_index = None
|
||||||
_resolved_hf_endpoint = None
|
|
||||||
_resolved_torch_index = None
|
|
||||||
DEFAULT_MODEL_KEYS = ("SMALL", "BASE", BEN2_MODEL.key)
|
DEFAULT_MODEL_KEYS = ("SMALL", "BASE", BEN2_MODEL.key)
|
||||||
PREFLIGHT_END = 0.05
|
|
||||||
MIRROR_END = 0.07
|
|
||||||
RUNTIME_SLICE_START = 0.07
|
RUNTIME_SLICE_START = 0.07
|
||||||
RUNTIME_SLICE_END = 0.88
|
RUNTIME_SLICE_END = 0.88
|
||||||
VERIFY_SLICE_START = 0.88
|
VERIFY_SLICE_START = 0.88
|
||||||
@@ -101,8 +89,6 @@ def child_environment(environment=None):
|
|||||||
process_environment[env_key] = value
|
process_environment[env_key] = value
|
||||||
if _resolved_pypi_index:
|
if _resolved_pypi_index:
|
||||||
process_environment["PIP_INDEX_URL"] = _resolved_pypi_index
|
process_environment["PIP_INDEX_URL"] = _resolved_pypi_index
|
||||||
if _resolved_hf_endpoint:
|
|
||||||
process_environment["HF_ENDPOINT"] = _resolved_hf_endpoint
|
|
||||||
if environment:
|
if environment:
|
||||||
process_environment.update(environment)
|
process_environment.update(environment)
|
||||||
return process_environment
|
return process_environment
|
||||||
@@ -131,127 +117,30 @@ def parse_package_count(line):
|
|||||||
|
|
||||||
|
|
||||||
def effective_torch_index(backend):
|
def effective_torch_index(backend):
|
||||||
if _resolved_torch_index:
|
|
||||||
return _resolved_torch_index
|
|
||||||
return os.environ.get("SNAPFORM_TORCH_INDEX") or TORCH_INDEXES.get(backend)
|
return os.environ.get("SNAPFORM_TORCH_INDEX") or TORCH_INDEXES.get(backend)
|
||||||
|
|
||||||
|
|
||||||
def probe_hosts(urls, timeout):
|
def resolve_indexes(args):
|
||||||
results = {}
|
"""???????????????????????????????????"""
|
||||||
|
global _resolved_pypi_index
|
||||||
def worker(url):
|
_resolved_pypi_index = (
|
||||||
results[url] = reachable(url, timeout)
|
getattr(args, "pypi_index", None)
|
||||||
|
or os.environ.get("PIP_INDEX_URL")
|
||||||
threads = [
|
|
||||||
threading.Thread(target=worker, args=(url,)) for url in urls
|
|
||||||
]
|
|
||||||
for thread in threads:
|
|
||||||
thread.start()
|
|
||||||
for thread in threads:
|
|
||||||
thread.join()
|
|
||||||
return results
|
|
||||||
|
|
||||||
|
|
||||||
def resolve_endpoints(backend):
|
|
||||||
"""Resolve package indexes, falling back to mirrors when direct access fails; returns (resolved, failures, used_mirror)."""
|
|
||||||
preferred = {
|
|
||||||
"pypi": os.environ.get("PIP_INDEX_URL") or "https://pypi.org/simple/",
|
|
||||||
"hf": os.environ.get("HF_ENDPOINT") or HF_ENDPOINT,
|
|
||||||
"torch": os.environ.get("SNAPFORM_TORCH_INDEX")
|
|
||||||
or TORCH_INDEXES.get(backend),
|
|
||||||
"github": "https://github.com/",
|
|
||||||
}
|
|
||||||
if sys.platform == "darwin":
|
|
||||||
preferred["torch"] = None
|
|
||||||
mirror_map = {
|
|
||||||
"pypi": PYPI_MIRRORS,
|
|
||||||
"hf": HF_MIRRORS,
|
|
||||||
"torch": TORCH_MIRRORS.get(backend, ()),
|
|
||||||
"github": (),
|
|
||||||
}
|
|
||||||
results = probe_hosts(
|
|
||||||
[url for url in preferred.values() if url],
|
|
||||||
PREFLIGHT_TIMEOUT,
|
|
||||||
)
|
)
|
||||||
resolved = {}
|
|
||||||
failures = []
|
|
||||||
for key, url in preferred.items():
|
|
||||||
if url is None:
|
|
||||||
resolved[key] = None
|
|
||||||
continue
|
|
||||||
if results[url]:
|
|
||||||
resolved[key] = url
|
|
||||||
continue
|
|
||||||
mirror_results = probe_hosts(mirror_map[key], MIRROR_TIMEOUT)
|
|
||||||
mirror = next(
|
|
||||||
(
|
|
||||||
candidate
|
|
||||||
for candidate in mirror_map[key]
|
|
||||||
if mirror_results.get(candidate)
|
|
||||||
),
|
|
||||||
None,
|
|
||||||
)
|
|
||||||
if mirror:
|
|
||||||
resolved[key] = mirror
|
|
||||||
else:
|
|
||||||
resolved[key] = url
|
|
||||||
failures.append(url)
|
|
||||||
used_mirror = any(
|
|
||||||
key in preferred
|
|
||||||
and preferred[key] is not None
|
|
||||||
and resolved[key] != preferred[key]
|
|
||||||
for key in resolved
|
|
||||||
)
|
|
||||||
return resolved, failures, used_mirror
|
|
||||||
|
|
||||||
|
|
||||||
def preflight(args):
|
def _run_install_once(
|
||||||
"""Check disk space and network reachability; switch to mirrors when direct access fails and raise a clear English error when nothing is reachable."""
|
|
||||||
global _resolved_pypi_index, _resolved_hf_endpoint, _resolved_torch_index
|
|
||||||
|
|
||||||
problems = []
|
|
||||||
free_bytes = shutil.disk_usage(Path(args.runtime_dir).parent).free
|
|
||||||
if free_bytes < PREFLIGHT_MIN_FREE_BYTES:
|
|
||||||
needed_gb = PREFLIGHT_MIN_FREE_BYTES // (1024**3)
|
|
||||||
problems.append(
|
|
||||||
f"Not enough free disk space: {needed_gb} GB required, "
|
|
||||||
f"only {free_bytes / 1024**3:.1f} GB available."
|
|
||||||
)
|
|
||||||
resolved, failures, used_mirror = resolve_endpoints(args.backend)
|
|
||||||
_resolved_pypi_index = resolved["pypi"]
|
|
||||||
_resolved_hf_endpoint = resolved["hf"]
|
|
||||||
_resolved_torch_index = resolved["torch"]
|
|
||||||
unreachable = failures
|
|
||||||
if unreachable:
|
|
||||||
proxies = urllib.request.getproxies()
|
|
||||||
proxy_hint = (
|
|
||||||
"Proxy detected: " + ", ".join(
|
|
||||||
f"{key}={value}" for key, value in proxies.items()
|
|
||||||
)
|
|
||||||
if proxies
|
|
||||||
else "No proxy detected"
|
|
||||||
)
|
|
||||||
problems.append(
|
|
||||||
"Cannot reach: " + ", ".join(unreachable)
|
|
||||||
+ f". {proxy_hint}. Check network or proxy settings and retry."
|
|
||||||
)
|
|
||||||
if problems:
|
|
||||||
raise RuntimeError("Prerequisites check failed: " + " ".join(problems))
|
|
||||||
return resolved, used_mirror
|
|
||||||
|
|
||||||
|
|
||||||
def run_install(
|
|
||||||
status_path,
|
status_path,
|
||||||
stage,
|
stage,
|
||||||
stages,
|
stages,
|
||||||
label,
|
label,
|
||||||
command,
|
command,
|
||||||
environment=None,
|
environment,
|
||||||
start=0.0,
|
start,
|
||||||
end=1.0,
|
end,
|
||||||
clock=time.monotonic,
|
clock,
|
||||||
):
|
):
|
||||||
"""Run a uv install command and smoothly map package-count events onto the stage progress range."""
|
"""Run a uv install command once and map package-count events onto the stage progress range."""
|
||||||
print("Running:", subprocess.list2cmdline([str(value) for value in command]))
|
print("Running:", subprocess.list2cmdline([str(value) for value in command]))
|
||||||
process_environment = child_environment(environment)
|
process_environment = child_environment(environment)
|
||||||
process_environment["UV_LOG"] = "debug"
|
process_environment["UV_LOG"] = "debug"
|
||||||
@@ -342,6 +231,42 @@ def run_install(
|
|||||||
return installed
|
return installed
|
||||||
|
|
||||||
|
|
||||||
|
def run_install(
|
||||||
|
status_path,
|
||||||
|
stage,
|
||||||
|
stages,
|
||||||
|
label,
|
||||||
|
command,
|
||||||
|
environment=None,
|
||||||
|
start=0.0,
|
||||||
|
end=1.0,
|
||||||
|
clock=time.monotonic,
|
||||||
|
):
|
||||||
|
"""Run a uv install command, retrying with mirrors when the default index fails."""
|
||||||
|
attempts = [environment]
|
||||||
|
if _resolved_pypi_index in (None, PYPI_DEFAULT_INDEX):
|
||||||
|
attempts.extend(
|
||||||
|
{"PIP_INDEX_URL": mirror} for mirror in PYPI_MIRRORS
|
||||||
|
)
|
||||||
|
for index, attempt_environment in enumerate(attempts):
|
||||||
|
try:
|
||||||
|
return _run_install_once(
|
||||||
|
status_path,
|
||||||
|
stage,
|
||||||
|
stages,
|
||||||
|
label,
|
||||||
|
command,
|
||||||
|
attempt_environment,
|
||||||
|
start,
|
||||||
|
end,
|
||||||
|
clock,
|
||||||
|
)
|
||||||
|
except subprocess.CalledProcessError:
|
||||||
|
if index + 1 == len(attempts):
|
||||||
|
raise
|
||||||
|
print(f"Retrying with PyPI mirror: {PYPI_MIRRORS[0]}")
|
||||||
|
|
||||||
|
|
||||||
def environment_python(runtime_dir):
|
def environment_python(runtime_dir):
|
||||||
if sys.platform == "win32":
|
if sys.platform == "win32":
|
||||||
return Path(runtime_dir) / "Scripts" / "python.exe"
|
return Path(runtime_dir) / "Scripts" / "python.exe"
|
||||||
@@ -355,15 +280,22 @@ def install_stages(args, runtime_dir, python, uv):
|
|||||||
if sys.platform == "darwin"
|
if sys.platform == "darwin"
|
||||||
else effective_torch_index(args.backend)
|
else effective_torch_index(args.backend)
|
||||||
)
|
)
|
||||||
def github_archive_url(repository, commit):
|
def archive_requirement(project_name, filename):
|
||||||
return f"https://github.com/{repository}/archive/{commit}.tar.gz"
|
local = Path(args.runtime_dir).parent / "sources" / filename
|
||||||
|
if not local.is_file():
|
||||||
|
raise RuntimeError(
|
||||||
|
"Bundled source archive is missing; reinstall the extension: "
|
||||||
|
f"{local}"
|
||||||
|
)
|
||||||
|
return f"{project_name} @ {local.as_uri()}"
|
||||||
|
|
||||||
da3_requirement = (
|
da3_requirement = archive_requirement(
|
||||||
"depth-anything-3 @ "
|
"depth-anything-3",
|
||||||
+ github_archive_url("ByteDance-Seed/Depth-Anything-3", args.da3_commit)
|
"depth-anything-3.tar.gz",
|
||||||
)
|
)
|
||||||
ben2_requirement = (
|
ben2_requirement = archive_requirement(
|
||||||
"ben2 @ " + github_archive_url("PramaLLC/BEN2", args.ben2_commit)
|
"ben2",
|
||||||
|
"ben2.tar.gz",
|
||||||
)
|
)
|
||||||
torch_command = [
|
torch_command = [
|
||||||
uv,
|
uv,
|
||||||
@@ -488,12 +420,41 @@ def uv_download_target():
|
|||||||
raise RuntimeError(f"Unsupported platform for uv download: {sys.platform}")
|
raise RuntimeError(f"Unsupported platform for uv download: {sys.platform}")
|
||||||
|
|
||||||
|
|
||||||
def uv_download_url():
|
def bundled_sources_directory():
|
||||||
|
return Path(__file__).with_name("sources")
|
||||||
|
|
||||||
|
|
||||||
|
def bundled_uv_archive():
|
||||||
target, archive = uv_download_target()
|
target, archive = uv_download_target()
|
||||||
return (
|
path = bundled_sources_directory() / f"uv-{target}.{archive}"
|
||||||
"https://github.com/astral-sh/uv/releases/download/"
|
return path if path.is_file() else None
|
||||||
f"{UV_VERSION}/uv-{target}.{archive}"
|
|
||||||
)
|
|
||||||
|
def sync_bundled_sources(bundled_dir, sources_dir):
|
||||||
|
"""Copy bundled source archives into the storage sources directory."""
|
||||||
|
for name in ("depth-anything-3.tar.gz", "ben2.tar.gz"):
|
||||||
|
bundled = Path(bundled_dir) / name
|
||||||
|
if not bundled.is_file():
|
||||||
|
continue
|
||||||
|
target = Path(sources_dir) / name
|
||||||
|
if (
|
||||||
|
target.is_file()
|
||||||
|
and target.stat().st_size == bundled.stat().st_size
|
||||||
|
):
|
||||||
|
continue
|
||||||
|
Path(sources_dir).mkdir(parents=True, exist_ok=True)
|
||||||
|
shutil.copy2(bundled, target)
|
||||||
|
|
||||||
|
|
||||||
|
def archive_is_valid(path):
|
||||||
|
if zipfile.is_zipfile(path):
|
||||||
|
return True
|
||||||
|
try:
|
||||||
|
with tarfile.open(path) as archive:
|
||||||
|
archive.getmembers()
|
||||||
|
return True
|
||||||
|
except (OSError, tarfile.TarError):
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
def extract_uv(archive_path, uv_path):
|
def extract_uv(archive_path, uv_path):
|
||||||
@@ -511,8 +472,8 @@ def extract_uv(archive_path, uv_path):
|
|||||||
os.chmod(uv_path, 0o755)
|
os.chmod(uv_path, 0o755)
|
||||||
|
|
||||||
|
|
||||||
def download_uv(
|
def install_uv(
|
||||||
uv_path,
|
uv,
|
||||||
status_path,
|
status_path,
|
||||||
stage=None,
|
stage=None,
|
||||||
stages=None,
|
stages=None,
|
||||||
@@ -520,50 +481,29 @@ def download_uv(
|
|||||||
start=0.0,
|
start=0.0,
|
||||||
end=1.0,
|
end=1.0,
|
||||||
):
|
):
|
||||||
"""Download the standalone uv binary and map byte progress onto the stage progress range."""
|
"""Extract the bundled uv binary; fail clearly when it is missing."""
|
||||||
stage_fields = {
|
stage_fields = {
|
||||||
"stage": stage,
|
"stage": stage,
|
||||||
"stages": stages,
|
"stages": stages,
|
||||||
"stage_label": stage_label,
|
"stage_label": stage_label,
|
||||||
}
|
}
|
||||||
|
bundled_uv = bundled_uv_archive()
|
||||||
def write(progress, message):
|
if bundled_uv is None:
|
||||||
write_status(
|
raise RuntimeError(
|
||||||
status_path,
|
"Bundled uv is missing; reinstall the extension matching your "
|
||||||
start + progress * (end - start),
|
f"platform (searched {bundled_sources_directory()})"
|
||||||
message,
|
|
||||||
**stage_fields,
|
|
||||||
)
|
)
|
||||||
|
if not archive_is_valid(bundled_uv):
|
||||||
write(0.0, "downloading")
|
raise RuntimeError(f"Bundled uv archive is invalid: {bundled_uv}")
|
||||||
uv_path = Path(uv_path)
|
print(f"Extracting bundled uv from {bundled_uv}")
|
||||||
uv_path.parent.mkdir(parents=True, exist_ok=True)
|
write_status(
|
||||||
archive_path = uv_path.with_name(f"{uv_path.name}.part")
|
status_path,
|
||||||
request = urllib.request.Request(
|
start,
|
||||||
uv_download_url(),
|
"extracting bundled uv",
|
||||||
headers={"User-Agent": "SnapForm-Runtime-Installer"},
|
**stage_fields,
|
||||||
)
|
)
|
||||||
try:
|
Path(uv).parent.mkdir(parents=True, exist_ok=True)
|
||||||
with urllib.request.urlopen(request, timeout=60) as response:
|
extract_uv(bundled_uv, uv)
|
||||||
total = int(response.headers.get("Content-Length") or 0)
|
|
||||||
downloaded = 0
|
|
||||||
started = time.monotonic()
|
|
||||||
with archive_path.open("wb") as output:
|
|
||||||
while True:
|
|
||||||
chunk = response.read(65536)
|
|
||||||
if not chunk:
|
|
||||||
break
|
|
||||||
output.write(chunk)
|
|
||||||
downloaded += len(chunk)
|
|
||||||
progress = (
|
|
||||||
min(downloaded / total, 1.0) if total else 0.0
|
|
||||||
)
|
|
||||||
speed = downloaded / max(time.monotonic() - started, 1e-6)
|
|
||||||
write(progress, f"downloading {format_rate(speed)}")
|
|
||||||
write(1.0, "preparing")
|
|
||||||
extract_uv(archive_path, uv_path)
|
|
||||||
finally:
|
|
||||||
archive_path.unlink(missing_ok=True)
|
|
||||||
|
|
||||||
|
|
||||||
def run_stage(
|
def run_stage(
|
||||||
@@ -636,10 +576,7 @@ def install(args):
|
|||||||
"SnapForm requires Apple Silicon macOS; PyTorch stopped "
|
"SnapForm requires Apple Silicon macOS; PyTorch stopped "
|
||||||
"shipping Intel Mac (x86_64) wheels after 2.2.2"
|
"shipping Intel Mac (x86_64) wheels after 2.2.2"
|
||||||
)
|
)
|
||||||
write_status(args.status, 0.0, "checking prerequisites")
|
resolve_indexes(args)
|
||||||
_resolved, used_mirror = preflight(args)
|
|
||||||
if used_mirror:
|
|
||||||
write_status(args.status, MIRROR_END, "using mirror sources")
|
|
||||||
runtime_dir = Path(args.runtime_dir)
|
runtime_dir = Path(args.runtime_dir)
|
||||||
runtime_dir.parent.mkdir(parents=True, exist_ok=True)
|
runtime_dir.parent.mkdir(parents=True, exist_ok=True)
|
||||||
if not runtime_ready(runtime_dir):
|
if not runtime_ready(runtime_dir):
|
||||||
@@ -650,6 +587,10 @@ def install(args):
|
|||||||
uv = Path(resolved) if resolved else uv
|
uv = Path(resolved) if resolved else uv
|
||||||
uv_missing = not uv.is_file()
|
uv_missing = not uv.is_file()
|
||||||
python = environment_python(runtime_dir)
|
python = environment_python(runtime_dir)
|
||||||
|
sync_bundled_sources(
|
||||||
|
bundled_sources_directory(),
|
||||||
|
Path(args.runtime_dir).parent / "sources",
|
||||||
|
)
|
||||||
stages = install_stages(args, runtime_dir, python, uv)
|
stages = install_stages(args, runtime_dir, python, uv)
|
||||||
runtime_stages = [stage for stage in stages if stage[3] != "download"]
|
runtime_stages = [stage for stage in stages if stage[3] != "download"]
|
||||||
model_stages = [stage for stage in stages if stage[3] == "download"]
|
model_stages = [stage for stage in stages if stage[3] == "download"]
|
||||||
@@ -677,7 +618,7 @@ def install(args):
|
|||||||
install_span if kind == "install" else PLAIN_SLICE_SPAN
|
install_span if kind == "install" else PLAIN_SLICE_SPAN
|
||||||
)
|
)
|
||||||
if uv_missing and index == 1:
|
if uv_missing and index == 1:
|
||||||
download_uv(
|
install_uv(
|
||||||
uv,
|
uv,
|
||||||
args.status,
|
args.status,
|
||||||
stage=index,
|
stage=index,
|
||||||
@@ -766,6 +707,7 @@ def parse_args():
|
|||||||
parser.add_argument("--da3-commit", required=True)
|
parser.add_argument("--da3-commit", required=True)
|
||||||
parser.add_argument("--ben2-commit", required=True)
|
parser.add_argument("--ben2-commit", required=True)
|
||||||
parser.add_argument("--status", required=True)
|
parser.add_argument("--status", required=True)
|
||||||
|
parser.add_argument("--pypi-index", default=None)
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -1,3 +1,5 @@
|
|||||||
|
import re
|
||||||
|
import shutil
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
import bpy
|
import bpy
|
||||||
@@ -22,6 +24,17 @@ IMAGE_NAME_SUFFIXES = (
|
|||||||
".bmp",
|
".bmp",
|
||||||
".tga",
|
".tga",
|
||||||
)
|
)
|
||||||
|
VIDEO_NAME_SUFFIXES = (
|
||||||
|
".mp4",
|
||||||
|
".mov",
|
||||||
|
".avi",
|
||||||
|
".mkv",
|
||||||
|
".webm",
|
||||||
|
".m4v",
|
||||||
|
".mts",
|
||||||
|
".mpg",
|
||||||
|
".mpeg",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def image_depth_model_items(owner, context):
|
def image_depth_model_items(owner, context):
|
||||||
@@ -47,6 +60,56 @@ def is_image_empty(candidate):
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def is_animated_image(image):
|
||||||
|
return getattr(image, "source", None) in {"MOVIE", "SEQUENCE"}
|
||||||
|
|
||||||
|
|
||||||
|
def video_total_frames(image):
|
||||||
|
"""Return the total animated frame count, or 1 for a still image."""
|
||||||
|
source = getattr(image, "source", None)
|
||||||
|
if source == "MOVIE":
|
||||||
|
return int(getattr(image, "frame_duration", 1) or 1)
|
||||||
|
if source == "SEQUENCE":
|
||||||
|
return len(collect_sequence_frames(image)) or 1
|
||||||
|
return 1
|
||||||
|
|
||||||
|
|
||||||
|
def image_user_settings(source_object):
|
||||||
|
"""Return the Empty's image user settings as a four-item tuple."""
|
||||||
|
image_user = getattr(source_object, "image_user", None)
|
||||||
|
if image_user is None:
|
||||||
|
return 1, 0, 0, False
|
||||||
|
return (
|
||||||
|
int(getattr(image_user, "frame_start", 1) or 1),
|
||||||
|
int(getattr(image_user, "frame_offset", 0) or 0),
|
||||||
|
int(getattr(image_user, "frame_duration", 0) or 0),
|
||||||
|
bool(getattr(image_user, "use_cyclic", False)),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def extract_frame_range(source_object):
|
||||||
|
"""Return the clamped (start_frame, frame_count) range extracted from the source."""
|
||||||
|
total_frames = video_total_frames(getattr(source_object, "data", None))
|
||||||
|
if total_frames <= 1:
|
||||||
|
return 1, total_frames
|
||||||
|
_frame_start, frame_offset, image_frames, _use_cyclic = (
|
||||||
|
image_user_settings(source_object)
|
||||||
|
)
|
||||||
|
frame_count = (
|
||||||
|
min(image_frames, total_frames)
|
||||||
|
if image_frames > 1
|
||||||
|
else total_frames
|
||||||
|
)
|
||||||
|
start_frame = min(1 + frame_offset, total_frames - frame_count + 1)
|
||||||
|
start_frame = max(start_frame, 1)
|
||||||
|
return start_frame, frame_count
|
||||||
|
|
||||||
|
|
||||||
|
def default_extract_frames(source_object):
|
||||||
|
"""Return the Empty's Image Frames value, or the full frame count when unset."""
|
||||||
|
return extract_frame_range(source_object)[1]
|
||||||
|
|
||||||
|
|
||||||
def image_path(image):
|
def image_path(image):
|
||||||
filepath_from_user = getattr(image, "filepath_from_user", None)
|
filepath_from_user = getattr(image, "filepath_from_user", None)
|
||||||
filepath = (
|
filepath = (
|
||||||
@@ -65,10 +128,22 @@ def image_is_packed(image):
|
|||||||
return bool(getattr(image, "packed_files", ()))
|
return bool(getattr(image, "packed_files", ()))
|
||||||
|
|
||||||
|
|
||||||
|
def remove_temporary_input(path):
|
||||||
|
"""Remove a temporary input file or frame directory."""
|
||||||
|
path = Path(path)
|
||||||
|
try:
|
||||||
|
if path.is_dir():
|
||||||
|
shutil.rmtree(path, ignore_errors=True)
|
||||||
|
else:
|
||||||
|
path.unlink(missing_ok=True)
|
||||||
|
except OSError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
def image_base_name(image):
|
def image_base_name(image):
|
||||||
name = image.name
|
name = image.name
|
||||||
lower_name = name.lower()
|
lower_name = name.lower()
|
||||||
for suffix in IMAGE_NAME_SUFFIXES:
|
for suffix in IMAGE_NAME_SUFFIXES + VIDEO_NAME_SUFFIXES:
|
||||||
suffix_index = lower_name.rfind(suffix)
|
suffix_index = lower_name.rfind(suffix)
|
||||||
if suffix_index <= 0:
|
if suffix_index <= 0:
|
||||||
continue
|
continue
|
||||||
@@ -85,8 +160,15 @@ def depth_data_name(source_object):
|
|||||||
return f"{image_base_name(source_object.data)}.depth"
|
return f"{image_base_name(source_object.data)}.depth"
|
||||||
|
|
||||||
|
|
||||||
|
def copy_depth_image(source_object):
|
||||||
|
"""Return a renamed copy of the source image for use as the depth map."""
|
||||||
|
depth_image = source_object.data.copy()
|
||||||
|
depth_image.name = depth_data_name(source_object)
|
||||||
|
return depth_image
|
||||||
|
|
||||||
|
|
||||||
def removed_background_data_name(source_object):
|
def removed_background_data_name(source_object):
|
||||||
return f"{image_base_name(source_object.data)}.rb"
|
return f"{image_base_name(source_object.data)}.foreground"
|
||||||
|
|
||||||
|
|
||||||
def scaled_process_resolution(image, scale):
|
def scaled_process_resolution(image, scale):
|
||||||
@@ -117,10 +199,60 @@ def save_packed_image(image, scene):
|
|||||||
return output_path
|
return output_path
|
||||||
|
|
||||||
|
|
||||||
|
def collect_sequence_frames(image):
|
||||||
|
"""Return the numbered frame files of a SEQUENCE image, or [] when absent."""
|
||||||
|
source_path = image_path(image)
|
||||||
|
if not is_animated_image(image) or not source_path.is_file():
|
||||||
|
return []
|
||||||
|
match = re.match(
|
||||||
|
r"^(?P<prefix>.*?)(?P<digits>\d+)(?P<suffix>\.[^.]+)$",
|
||||||
|
source_path.name,
|
||||||
|
)
|
||||||
|
if match is None:
|
||||||
|
return []
|
||||||
|
prefix = match.group("prefix")
|
||||||
|
suffix = match.group("suffix").lower()
|
||||||
|
frames = []
|
||||||
|
for path in source_path.parent.iterdir():
|
||||||
|
frame_match = re.match(
|
||||||
|
r"^(?P<prefix>.*?)(?P<digits>\d+)(?P<suffix>\.[^.]+)$",
|
||||||
|
path.name,
|
||||||
|
)
|
||||||
|
if (
|
||||||
|
frame_match is not None
|
||||||
|
and frame_match.group("prefix") == prefix
|
||||||
|
and frame_match.group("suffix").lower() == suffix
|
||||||
|
):
|
||||||
|
frames.append((int(frame_match.group("digits")), path))
|
||||||
|
return [path for _digits, path in sorted(frames)]
|
||||||
|
|
||||||
|
|
||||||
def prepare_image_input(image, scene):
|
def prepare_image_input(image, scene):
|
||||||
source_path = image_path(image)
|
source_path = image_path(image)
|
||||||
if image_is_packed(image):
|
if image_is_packed(image):
|
||||||
return save_packed_image(image, scene), True
|
return save_packed_image(image, scene), True
|
||||||
|
if getattr(image, "source", None) == "MOVIE":
|
||||||
|
if not source_path.is_file():
|
||||||
|
raise RuntimeError(
|
||||||
|
"The image empty must reference an existing video file"
|
||||||
|
)
|
||||||
|
return source_path, False
|
||||||
|
sequence_frames = collect_sequence_frames(image)
|
||||||
|
if sequence_frames:
|
||||||
|
from .runtime import temporary_input_directory
|
||||||
|
|
||||||
|
temporary_directory = Path(temporary_input_directory())
|
||||||
|
try:
|
||||||
|
for index, frame_path in enumerate(sequence_frames):
|
||||||
|
suffix = frame_path.suffix.lower() or ".jpg"
|
||||||
|
shutil.copy2(
|
||||||
|
frame_path,
|
||||||
|
temporary_directory / f"frame.{index + 1:04d}{suffix}",
|
||||||
|
)
|
||||||
|
except Exception:
|
||||||
|
shutil.rmtree(temporary_directory, ignore_errors=True)
|
||||||
|
raise
|
||||||
|
return temporary_directory, True
|
||||||
if not source_path.is_file():
|
if not source_path.is_file():
|
||||||
raise RuntimeError(
|
raise RuntimeError(
|
||||||
"The image empty must reference an existing image file"
|
"The image empty must reference an existing image file"
|
||||||
@@ -128,12 +260,207 @@ def prepare_image_input(image, scene):
|
|||||||
return source_path, False
|
return source_path, False
|
||||||
|
|
||||||
|
|
||||||
|
def configure_image_user(
|
||||||
|
image_user,
|
||||||
|
frame_count=None,
|
||||||
|
frame_start=1,
|
||||||
|
frame_offset=0,
|
||||||
|
use_cyclic=None,
|
||||||
|
):
|
||||||
|
"""Map the image to scene frames and enable auto refresh."""
|
||||||
|
if image_user is None:
|
||||||
|
return
|
||||||
|
if hasattr(image_user, "frame_start"):
|
||||||
|
image_user.frame_start = frame_start
|
||||||
|
if hasattr(image_user, "frame_offset"):
|
||||||
|
image_user.frame_offset = frame_offset
|
||||||
|
if hasattr(image_user, "use_auto_refresh"):
|
||||||
|
image_user.use_auto_refresh = True
|
||||||
|
if use_cyclic is not None and hasattr(image_user, "use_cyclic"):
|
||||||
|
image_user.use_cyclic = use_cyclic
|
||||||
|
if frame_count is not None and hasattr(image_user, "frame_duration"):
|
||||||
|
image_user.frame_duration = frame_count
|
||||||
|
|
||||||
|
|
||||||
|
def pack_sequence_if_possible(image, expected_frames):
|
||||||
|
"""Pack a frame sequence when Blender supports it; otherwise keep external files."""
|
||||||
|
if expected_frames <= 0:
|
||||||
|
return False
|
||||||
|
try:
|
||||||
|
image.pack()
|
||||||
|
except (OSError, RuntimeError):
|
||||||
|
return False
|
||||||
|
if len(getattr(image, "packed_files", ())) >= expected_frames:
|
||||||
|
return True
|
||||||
|
try:
|
||||||
|
image.unpack(method="REMOVE")
|
||||||
|
except (OSError, RuntimeError):
|
||||||
|
pass
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def sorted_frame_paths(directory, prefix=None):
|
||||||
|
"""Return image files in a directory sorted by their trailing frame number."""
|
||||||
|
def frame_key(path):
|
||||||
|
match = re.search(r"(\d+)(?=\.[^.]+$)", path.name)
|
||||||
|
return (int(match.group(1)) if match is not None else 0, path.name)
|
||||||
|
|
||||||
|
return sorted(
|
||||||
|
(
|
||||||
|
path
|
||||||
|
for path in Path(directory).iterdir()
|
||||||
|
if path.is_file()
|
||||||
|
and path.suffix.lower() in IMAGE_NAME_SUFFIXES
|
||||||
|
and (prefix is None or path.stem.startswith(prefix))
|
||||||
|
),
|
||||||
|
key=frame_key,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def force_sequence_source(image):
|
||||||
|
"""Ensure a loaded image is a SEQUENCE datablock for its numbered frames."""
|
||||||
|
if getattr(image, "source", None) == "SEQUENCE":
|
||||||
|
return image
|
||||||
|
try:
|
||||||
|
image.source = "SEQUENCE"
|
||||||
|
except (OSError, RuntimeError, TypeError) as error:
|
||||||
|
raise RuntimeError(
|
||||||
|
"Blender cannot load the frame sequence"
|
||||||
|
) from error
|
||||||
|
if getattr(image, "source", None) != "SEQUENCE":
|
||||||
|
raise RuntimeError("Blender cannot load the frame sequence")
|
||||||
|
return image
|
||||||
|
|
||||||
|
|
||||||
|
def create_sequence_image(name, frame_directory, frame_prefix="frame."):
|
||||||
|
"""Load a numbered frame directory as a SEQUENCE image datablock."""
|
||||||
|
frame_paths = sorted_frame_paths(frame_directory, frame_prefix)
|
||||||
|
if not frame_paths:
|
||||||
|
raise RuntimeError(f"No image frames found in: {frame_directory}")
|
||||||
|
try:
|
||||||
|
image = bpy.data.images.load(
|
||||||
|
str(frame_paths[0]),
|
||||||
|
check_existing=False,
|
||||||
|
)
|
||||||
|
force_sequence_source(image)
|
||||||
|
except (OSError, RuntimeError) as error:
|
||||||
|
if "image" in locals():
|
||||||
|
bpy.data.images.remove(image, do_unlink=True)
|
||||||
|
raise RuntimeError("Unable to load the frame sequence") from error
|
||||||
|
image.name = name
|
||||||
|
pack_sequence_if_possible(image, len(frame_paths))
|
||||||
|
return image
|
||||||
|
|
||||||
|
|
||||||
|
def create_depth_sequence_image(name, depth_frames, scene, frame_directory):
|
||||||
|
"""Write one EXR per depth frame and return a SEQUENCE image datablock."""
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
frame_directory = Path(frame_directory)
|
||||||
|
frame_directory.mkdir(parents=True, exist_ok=True)
|
||||||
|
height, width = depth_frames[0].shape
|
||||||
|
image = bpy.data.images.new(
|
||||||
|
name,
|
||||||
|
width=width,
|
||||||
|
height=height,
|
||||||
|
alpha=False,
|
||||||
|
float_buffer=True,
|
||||||
|
)
|
||||||
|
image.colorspace_settings.name = "Non-Color"
|
||||||
|
image_settings = scene.render.image_settings
|
||||||
|
original_file_format = image_settings.file_format
|
||||||
|
frame_paths = []
|
||||||
|
try:
|
||||||
|
image_settings.file_format = "OPEN_EXR"
|
||||||
|
for frame_index, depth_values in enumerate(depth_frames):
|
||||||
|
# NPZ arrays use a top-left origin, while Blender image pixels start below.
|
||||||
|
depth_pixels = np.flipud(depth_values)
|
||||||
|
rgba = np.empty((height, width, 4), dtype=np.float32)
|
||||||
|
rgba[..., 0] = depth_pixels
|
||||||
|
rgba[..., 1] = depth_pixels
|
||||||
|
rgba[..., 2] = depth_pixels
|
||||||
|
rgba[..., 3] = 1.0
|
||||||
|
image.pixels.foreach_set(rgba.ravel())
|
||||||
|
frame_path = (
|
||||||
|
frame_directory / f"depth.{frame_index + 1:04d}.exr"
|
||||||
|
)
|
||||||
|
image.save_render(str(frame_path), scene=scene)
|
||||||
|
if not frame_path.is_file():
|
||||||
|
raise RuntimeError(f"Blender did not write: {frame_path}")
|
||||||
|
frame_paths.append(frame_path)
|
||||||
|
finally:
|
||||||
|
image_settings.file_format = original_file_format
|
||||||
|
bpy.data.images.remove(image, do_unlink=True)
|
||||||
|
if not frame_paths:
|
||||||
|
raise RuntimeError("Unable to write the depth frame sequence")
|
||||||
|
try:
|
||||||
|
sequence_image = bpy.data.images.load(
|
||||||
|
str(frame_paths[0]),
|
||||||
|
check_existing=False,
|
||||||
|
)
|
||||||
|
force_sequence_source(sequence_image)
|
||||||
|
except (OSError, RuntimeError) as error:
|
||||||
|
if "sequence_image" in locals():
|
||||||
|
bpy.data.images.remove(sequence_image, do_unlink=True)
|
||||||
|
raise RuntimeError(
|
||||||
|
"Unable to load the depth frame sequence"
|
||||||
|
) from error
|
||||||
|
sequence_image.name = name
|
||||||
|
sequence_image.colorspace_settings.name = "Non-Color"
|
||||||
|
pack_sequence_if_possible(sequence_image, len(frame_paths))
|
||||||
|
return sequence_image
|
||||||
|
|
||||||
|
|
||||||
def create_removed_background_image(
|
def create_removed_background_image(
|
||||||
source_object,
|
source_object,
|
||||||
output_path,
|
output_path,
|
||||||
assign_to_source=True,
|
assign_to_source=True,
|
||||||
):
|
):
|
||||||
output_path = Path(output_path)
|
output_path = Path(output_path)
|
||||||
|
if output_path.is_dir():
|
||||||
|
foreground_paths = sorted_frame_paths(output_path, "foreground.")
|
||||||
|
if not foreground_paths:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"Background removal result has no foreground frames: "
|
||||||
|
f"{output_path}"
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
result_image = bpy.data.images.load(
|
||||||
|
str(foreground_paths[0]),
|
||||||
|
check_existing=False,
|
||||||
|
)
|
||||||
|
force_sequence_source(result_image)
|
||||||
|
except (OSError, RuntimeError) as error:
|
||||||
|
if "result_image" in locals():
|
||||||
|
bpy.data.images.remove(result_image, do_unlink=True)
|
||||||
|
raise RuntimeError("Unable to load the background result") from error
|
||||||
|
source_image = source_object.data
|
||||||
|
original_fake_user = source_image.use_fake_user
|
||||||
|
try:
|
||||||
|
result_image.name = removed_background_data_name(source_object)
|
||||||
|
pack_sequence_if_possible(result_image, len(foreground_paths))
|
||||||
|
if assign_to_source:
|
||||||
|
image_user = getattr(source_object, "image_user", None)
|
||||||
|
(
|
||||||
|
original_frame_start,
|
||||||
|
_original_frame_offset,
|
||||||
|
_original_frame_duration,
|
||||||
|
original_use_cyclic,
|
||||||
|
) = image_user_settings(source_object)
|
||||||
|
source_image.use_fake_user = True
|
||||||
|
source_object.data = result_image
|
||||||
|
configure_image_user(
|
||||||
|
image_user,
|
||||||
|
len(foreground_paths),
|
||||||
|
frame_start=original_frame_start,
|
||||||
|
use_cyclic=original_use_cyclic,
|
||||||
|
)
|
||||||
|
except (OSError, RuntimeError) as error:
|
||||||
|
source_image.use_fake_user = original_fake_user
|
||||||
|
bpy.data.images.remove(result_image)
|
||||||
|
raise RuntimeError("Unable to load the background result") from error
|
||||||
|
return result_image
|
||||||
|
|
||||||
if not output_path.is_file():
|
if not output_path.is_file():
|
||||||
raise RuntimeError(
|
raise RuntimeError(
|
||||||
f"Background removal result does not exist: {output_path}"
|
f"Background removal result does not exist: {output_path}"
|
||||||
@@ -229,7 +556,7 @@ def apply_plane_transform(plane_object, source_object, camera_bounds):
|
|||||||
return abs(float(plane_object.scale.z))
|
return abs(float(plane_object.scale.z))
|
||||||
|
|
||||||
|
|
||||||
def extract_depth_array(npz_path):
|
def extract_depth_arrays(npz_path):
|
||||||
try:
|
try:
|
||||||
import numpy as np
|
import numpy as np
|
||||||
except ImportError as error:
|
except ImportError as error:
|
||||||
@@ -256,13 +583,17 @@ def extract_depth_array(npz_path):
|
|||||||
|
|
||||||
if depth.ndim >= 3 and depth.shape[-1] == 1:
|
if depth.ndim >= 3 and depth.shape[-1] == 1:
|
||||||
depth = depth[..., 0]
|
depth = depth[..., 0]
|
||||||
if depth.ndim == 3:
|
if depth.ndim == 2:
|
||||||
depth = depth[0]
|
depth = depth[None, ...]
|
||||||
if depth.ndim != 2:
|
if depth.ndim != 3:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"Unsupported depth array shape: {tuple(depth.shape)}"
|
f"Unsupported depth array shape: {tuple(depth.shape)}"
|
||||||
)
|
)
|
||||||
return depth
|
return [depth[index] for index in range(depth.shape[0])]
|
||||||
|
|
||||||
|
|
||||||
|
def extract_depth_array(npz_path):
|
||||||
|
return extract_depth_arrays(npz_path)[0]
|
||||||
|
|
||||||
|
|
||||||
def extract_intrinsics_array(npz_path):
|
def extract_intrinsics_array(npz_path):
|
||||||
@@ -475,6 +806,13 @@ def create_material(
|
|||||||
displacement_midlevel=0.0,
|
displacement_midlevel=0.0,
|
||||||
displacement_scale=1.0,
|
displacement_scale=1.0,
|
||||||
unshaded=False,
|
unshaded=False,
|
||||||
|
frame_count=1,
|
||||||
|
frame_start=1,
|
||||||
|
color_frame_offset=0,
|
||||||
|
color_frame_duration=None,
|
||||||
|
depth_frame_offset=0,
|
||||||
|
depth_frame_duration=None,
|
||||||
|
use_cyclic=None,
|
||||||
):
|
):
|
||||||
material = bpy.data.materials.new(name)
|
material = bpy.data.materials.new(name)
|
||||||
material.use_nodes = True
|
material.use_nodes = True
|
||||||
@@ -491,6 +829,25 @@ def create_material(
|
|||||||
color_texture.extension = "EXTEND"
|
color_texture.extension = "EXTEND"
|
||||||
color_texture.location = (-640.0, 180.0)
|
color_texture.location = (-640.0, 180.0)
|
||||||
output.location = (460.0, 180.0)
|
output.location = (460.0, 180.0)
|
||||||
|
extracted_frames = frame_count if frame_count > 1 else None
|
||||||
|
color_frames = (
|
||||||
|
color_frame_duration
|
||||||
|
if color_frame_duration is not None
|
||||||
|
else extracted_frames
|
||||||
|
)
|
||||||
|
depth_frames = (
|
||||||
|
depth_frame_duration
|
||||||
|
if depth_frame_duration is not None
|
||||||
|
else extracted_frames
|
||||||
|
)
|
||||||
|
if frame_count > 1 or is_animated_image(color_image):
|
||||||
|
configure_image_user(
|
||||||
|
color_texture.image_user,
|
||||||
|
color_frames,
|
||||||
|
frame_start=frame_start,
|
||||||
|
frame_offset=color_frame_offset,
|
||||||
|
use_cyclic=use_cyclic,
|
||||||
|
)
|
||||||
|
|
||||||
links = material.node_tree.links
|
links = material.node_tree.links
|
||||||
if unshaded:
|
if unshaded:
|
||||||
@@ -558,6 +915,14 @@ def create_material(
|
|||||||
depth_texture.image = displacement_image
|
depth_texture.image = displacement_image
|
||||||
depth_texture.interpolation = "Linear"
|
depth_texture.interpolation = "Linear"
|
||||||
depth_texture.extension = "EXTEND"
|
depth_texture.extension = "EXTEND"
|
||||||
|
if frame_count > 1 or is_animated_image(displacement_image):
|
||||||
|
configure_image_user(
|
||||||
|
depth_texture.image_user,
|
||||||
|
depth_frames,
|
||||||
|
frame_start=frame_start,
|
||||||
|
frame_offset=depth_frame_offset,
|
||||||
|
use_cyclic=use_cyclic,
|
||||||
|
)
|
||||||
displacement.inputs["Midlevel"].default_value = (
|
displacement.inputs["Midlevel"].default_value = (
|
||||||
displacement_midlevel
|
displacement_midlevel
|
||||||
)
|
)
|
||||||
@@ -612,8 +977,35 @@ def build_image_depth_plane(
|
|||||||
displacement_via="MODIFIER",
|
displacement_via="MODIFIER",
|
||||||
unshaded=False,
|
unshaded=False,
|
||||||
color_image=None,
|
color_image=None,
|
||||||
|
depth_frames_directory=None,
|
||||||
):
|
):
|
||||||
source_name = source_object.name
|
source_name = source_object.name
|
||||||
|
frame_count = 1
|
||||||
|
frames_image = None
|
||||||
|
if depth_frames_directory is not None:
|
||||||
|
frames_image = create_sequence_image(
|
||||||
|
depth_data_name(source_object),
|
||||||
|
depth_frames_directory,
|
||||||
|
)
|
||||||
|
frame_count = len(
|
||||||
|
sorted_frame_paths(depth_frames_directory, "frame.")
|
||||||
|
)
|
||||||
|
(
|
||||||
|
source_frame_start,
|
||||||
|
source_frame_offset,
|
||||||
|
source_frame_duration,
|
||||||
|
source_use_cyclic,
|
||||||
|
) = image_user_settings(source_object)
|
||||||
|
uses_extracted = frames_image is not None
|
||||||
|
texture_frame_count = (
|
||||||
|
frame_count if uses_extracted else source_frame_duration
|
||||||
|
)
|
||||||
|
texture_frame_offset = 0 if uses_extracted else source_frame_offset
|
||||||
|
depth_source_image = source_object.data
|
||||||
|
if frames_image is not None:
|
||||||
|
depth_source_image = frames_image
|
||||||
|
elif not is_animated_image(source_object.data):
|
||||||
|
depth_source_image = copy_depth_image(source_object)
|
||||||
mesh = create_plane_mesh(
|
mesh = create_plane_mesh(
|
||||||
source_object,
|
source_object,
|
||||||
subdivision_level=subdivision_level,
|
subdivision_level=subdivision_level,
|
||||||
@@ -632,25 +1024,43 @@ def build_image_depth_plane(
|
|||||||
depth_data_name(source_object),
|
depth_data_name(source_object),
|
||||||
type="IMAGE",
|
type="IMAGE",
|
||||||
)
|
)
|
||||||
texture.image = source_object.data
|
texture.image = depth_source_image
|
||||||
|
texture.extension = "EXTEND"
|
||||||
|
if uses_extracted or is_animated_image(source_object.data):
|
||||||
|
configure_image_user(
|
||||||
|
texture.image_user,
|
||||||
|
texture_frame_count,
|
||||||
|
frame_start=source_frame_start,
|
||||||
|
frame_offset=texture_frame_offset,
|
||||||
|
use_cyclic=source_use_cyclic,
|
||||||
|
)
|
||||||
displace = plane_object.modifiers.new("Displace", "DISPLACE")
|
displace = plane_object.modifiers.new("Displace", "DISPLACE")
|
||||||
displace.texture = texture
|
displace.texture = texture
|
||||||
displace.texture_coords = "UV"
|
displace.texture_coords = "UV"
|
||||||
displace.mid_level = ORIGINAL_DISPLACEMENT_MIDLEVEL
|
displace.mid_level = ORIGINAL_DISPLACEMENT_MIDLEVEL
|
||||||
|
|
||||||
depth_scale = abs(float(plane_object.scale.z))
|
depth_scale = abs(float(plane_object.scale.z))
|
||||||
|
displacement_image = (
|
||||||
|
depth_source_image if displacement_via == "MATERIAL" else None
|
||||||
|
)
|
||||||
material = create_material(
|
material = create_material(
|
||||||
source_name,
|
source_name,
|
||||||
color_image if color_image is not None else source_object.data,
|
color_image if color_image is not None else source_object.data,
|
||||||
depth_scale,
|
depth_scale,
|
||||||
displacement_image=(
|
displacement_image=displacement_image,
|
||||||
source_object.data if displacement_via == "MATERIAL" else None
|
|
||||||
),
|
|
||||||
displacement_midlevel=ORIGINAL_DISPLACEMENT_MIDLEVEL,
|
displacement_midlevel=ORIGINAL_DISPLACEMENT_MIDLEVEL,
|
||||||
displacement_scale=ORIGINAL_DISPLACEMENT_STRENGTH,
|
displacement_scale=ORIGINAL_DISPLACEMENT_STRENGTH,
|
||||||
unshaded=unshaded,
|
unshaded=unshaded,
|
||||||
|
frame_count=frame_count,
|
||||||
|
frame_start=source_frame_start,
|
||||||
|
color_frame_offset=texture_frame_offset,
|
||||||
|
color_frame_duration=texture_frame_count,
|
||||||
|
depth_frame_offset=texture_frame_offset,
|
||||||
|
depth_frame_duration=texture_frame_count,
|
||||||
|
use_cyclic=source_use_cyclic,
|
||||||
)
|
)
|
||||||
mesh.materials.append(material)
|
mesh.materials.append(material)
|
||||||
|
plane_object["snapform_frame_count"] = frame_count
|
||||||
return finish_depth_plane(
|
return finish_depth_plane(
|
||||||
context,
|
context,
|
||||||
source_object,
|
source_object,
|
||||||
@@ -669,9 +1079,11 @@ def build_depth_plane(
|
|||||||
unshaded=False,
|
unshaded=False,
|
||||||
depth_name="",
|
depth_name="",
|
||||||
color_image=None,
|
color_image=None,
|
||||||
|
depth_frames_directory=None,
|
||||||
):
|
):
|
||||||
source_name = source_object.name
|
source_name = source_object.name
|
||||||
depth = extract_depth_array(npz_path)
|
depth_arrays = extract_depth_arrays(npz_path)
|
||||||
|
depth = depth_arrays[0]
|
||||||
intrinsics = extract_intrinsics_array(npz_path)
|
intrinsics = extract_intrinsics_array(npz_path)
|
||||||
depth_values = prepare_depth_values(depth)
|
depth_values = prepare_depth_values(depth)
|
||||||
bounds, reference_depth = camera_plane_bounds(
|
bounds, reference_depth = camera_plane_bounds(
|
||||||
@@ -681,10 +1093,25 @@ def build_depth_plane(
|
|||||||
DISPLACEMENT_STRENGTH,
|
DISPLACEMENT_STRENGTH,
|
||||||
)
|
)
|
||||||
depth_name = depth_name or depth_data_name(source_object)
|
depth_name = depth_name or depth_data_name(source_object)
|
||||||
|
frame_count = len(depth_arrays)
|
||||||
|
(
|
||||||
|
source_frame_start,
|
||||||
|
source_frame_offset,
|
||||||
|
source_frame_duration,
|
||||||
|
source_use_cyclic,
|
||||||
|
) = image_user_settings(source_object)
|
||||||
depth_image = create_depth_image(
|
depth_image = create_depth_image(
|
||||||
depth_name,
|
depth_name,
|
||||||
depth_values,
|
depth_values,
|
||||||
)
|
)
|
||||||
|
if frame_count > 1 and depth_frames_directory is not None:
|
||||||
|
bpy.data.images.remove(depth_image, do_unlink=True)
|
||||||
|
depth_image = create_depth_sequence_image(
|
||||||
|
depth_name,
|
||||||
|
[prepare_depth_values(frame) for frame in depth_arrays],
|
||||||
|
context.scene,
|
||||||
|
depth_frames_directory,
|
||||||
|
)
|
||||||
mesh = create_plane_mesh(source_object, bounds, subdivision_level)
|
mesh = create_plane_mesh(source_object, bounds, subdivision_level)
|
||||||
plane_object = bpy.data.objects.new(
|
plane_object = bpy.data.objects.new(
|
||||||
source_name,
|
source_name,
|
||||||
@@ -709,6 +1136,13 @@ def build_depth_plane(
|
|||||||
)
|
)
|
||||||
texture.image = depth_image
|
texture.image = depth_image
|
||||||
texture.extension = "EXTEND"
|
texture.extension = "EXTEND"
|
||||||
|
if frame_count > 1:
|
||||||
|
configure_image_user(
|
||||||
|
texture.image_user,
|
||||||
|
frame_count,
|
||||||
|
frame_start=source_frame_start,
|
||||||
|
use_cyclic=source_use_cyclic,
|
||||||
|
)
|
||||||
displace = plane_object.modifiers.new("Displace", "DISPLACE")
|
displace = plane_object.modifiers.new("Displace", "DISPLACE")
|
||||||
displace.texture = texture
|
displace.texture = texture
|
||||||
displace.texture_coords = "UV"
|
displace.texture_coords = "UV"
|
||||||
@@ -726,12 +1160,22 @@ def build_depth_plane(
|
|||||||
displacement_midlevel=DISPLACEMENT_MIDLEVEL,
|
displacement_midlevel=DISPLACEMENT_MIDLEVEL,
|
||||||
displacement_scale=DISPLACEMENT_STRENGTH,
|
displacement_scale=DISPLACEMENT_STRENGTH,
|
||||||
unshaded=unshaded,
|
unshaded=unshaded,
|
||||||
|
frame_count=frame_count,
|
||||||
|
frame_start=source_frame_start,
|
||||||
|
color_frame_offset=0 if color_image is not None else source_frame_offset,
|
||||||
|
color_frame_duration=(
|
||||||
|
frame_count if color_image is not None else source_frame_duration
|
||||||
|
),
|
||||||
|
depth_frame_offset=0,
|
||||||
|
depth_frame_duration=frame_count,
|
||||||
|
use_cyclic=source_use_cyclic,
|
||||||
)
|
)
|
||||||
mesh.materials.append(material)
|
mesh.materials.append(material)
|
||||||
plane_object["snapform_reference_depth"] = reference_depth
|
plane_object["snapform_reference_depth"] = reference_depth
|
||||||
plane_object["snapform_camera_plane_width"] = bounds[1] - bounds[0]
|
plane_object["snapform_camera_plane_width"] = bounds[1] - bounds[0]
|
||||||
plane_object["snapform_camera_plane_height"] = bounds[3] - bounds[2]
|
plane_object["snapform_camera_plane_height"] = bounds[3] - bounds[2]
|
||||||
plane_object["snapform_depth_scale"] = depth_scale
|
plane_object["snapform_depth_scale"] = depth_scale
|
||||||
|
plane_object["snapform_frame_count"] = frame_count
|
||||||
|
|
||||||
return finish_depth_plane(
|
return finish_depth_plane(
|
||||||
context,
|
context,
|
||||||
@@ -758,6 +1202,12 @@ class BuildDepthPlane(bpy.types.Operator):
|
|||||||
snapform_background_path: bpy.props.StringProperty(
|
snapform_background_path: bpy.props.StringProperty(
|
||||||
options={"HIDDEN", "SKIP_SAVE"},
|
options={"HIDDEN", "SKIP_SAVE"},
|
||||||
)
|
)
|
||||||
|
snapform_color_frames_dir: bpy.props.StringProperty(
|
||||||
|
options={"HIDDEN", "SKIP_SAVE"},
|
||||||
|
)
|
||||||
|
snapform_depth_frames_dir: bpy.props.StringProperty(
|
||||||
|
options={"HIDDEN", "SKIP_SAVE"},
|
||||||
|
)
|
||||||
snapform_subdivision_level: bpy.props.IntProperty(
|
snapform_subdivision_level: bpy.props.IntProperty(
|
||||||
default=5,
|
default=5,
|
||||||
min=0,
|
min=0,
|
||||||
@@ -806,6 +1256,11 @@ class BuildDepthPlane(bpy.types.Operator):
|
|||||||
self.snapform_background_path,
|
self.snapform_background_path,
|
||||||
assign_to_source=False,
|
assign_to_source=False,
|
||||||
)
|
)
|
||||||
|
elif self.snapform_color_frames_dir:
|
||||||
|
color_image = create_sequence_image(
|
||||||
|
f"{image_base_name(source_object.data)}.frame",
|
||||||
|
self.snapform_color_frames_dir,
|
||||||
|
)
|
||||||
build_depth_plane(
|
build_depth_plane(
|
||||||
context,
|
context,
|
||||||
source_object,
|
source_object,
|
||||||
@@ -816,6 +1271,9 @@ class BuildDepthPlane(bpy.types.Operator):
|
|||||||
unshaded=self.snapform_unshaded,
|
unshaded=self.snapform_unshaded,
|
||||||
depth_name=self.snapform_depth_data_name,
|
depth_name=self.snapform_depth_data_name,
|
||||||
color_image=color_image,
|
color_image=color_image,
|
||||||
|
depth_frames_directory=(
|
||||||
|
self.snapform_depth_frames_dir or None
|
||||||
|
),
|
||||||
)
|
)
|
||||||
except (OSError, RuntimeError, ValueError) as error:
|
except (OSError, RuntimeError, ValueError) as error:
|
||||||
self.report({"ERROR"}, str(error))
|
self.report({"ERROR"}, str(error))
|
||||||
@@ -834,6 +1292,12 @@ class BuildOriginalDepthPlane(bpy.types.Operator):
|
|||||||
snapform_background_path: bpy.props.StringProperty(
|
snapform_background_path: bpy.props.StringProperty(
|
||||||
options={"HIDDEN", "SKIP_SAVE"},
|
options={"HIDDEN", "SKIP_SAVE"},
|
||||||
)
|
)
|
||||||
|
snapform_color_frames_dir: bpy.props.StringProperty(
|
||||||
|
options={"HIDDEN", "SKIP_SAVE"},
|
||||||
|
)
|
||||||
|
snapform_depth_frames_dir: bpy.props.StringProperty(
|
||||||
|
options={"HIDDEN", "SKIP_SAVE"},
|
||||||
|
)
|
||||||
snapform_subdivision_level: bpy.props.IntProperty(
|
snapform_subdivision_level: bpy.props.IntProperty(
|
||||||
default=5,
|
default=5,
|
||||||
min=0,
|
min=0,
|
||||||
@@ -866,11 +1330,18 @@ class BuildOriginalDepthPlane(bpy.types.Operator):
|
|||||||
)
|
)
|
||||||
return {"CANCELLED"}
|
return {"CANCELLED"}
|
||||||
try:
|
try:
|
||||||
color_image = create_removed_background_image(
|
color_image = None
|
||||||
source_object,
|
if self.snapform_background_path:
|
||||||
self.snapform_background_path,
|
color_image = create_removed_background_image(
|
||||||
assign_to_source=False,
|
source_object,
|
||||||
)
|
self.snapform_background_path,
|
||||||
|
assign_to_source=False,
|
||||||
|
)
|
||||||
|
elif self.snapform_color_frames_dir:
|
||||||
|
color_image = create_sequence_image(
|
||||||
|
f"{image_base_name(source_object.data)}.frame",
|
||||||
|
self.snapform_color_frames_dir,
|
||||||
|
)
|
||||||
build_image_depth_plane(
|
build_image_depth_plane(
|
||||||
context,
|
context,
|
||||||
source_object,
|
source_object,
|
||||||
@@ -879,6 +1350,9 @@ class BuildOriginalDepthPlane(bpy.types.Operator):
|
|||||||
displacement_via=self.snapform_displacement_via,
|
displacement_via=self.snapform_displacement_via,
|
||||||
unshaded=self.snapform_unshaded,
|
unshaded=self.snapform_unshaded,
|
||||||
color_image=color_image,
|
color_image=color_image,
|
||||||
|
depth_frames_directory=(
|
||||||
|
self.snapform_depth_frames_dir or None
|
||||||
|
),
|
||||||
)
|
)
|
||||||
except (OSError, RuntimeError, ValueError) as error:
|
except (OSError, RuntimeError, ValueError) as error:
|
||||||
self.report({"ERROR"}, str(error))
|
self.report({"ERROR"}, str(error))
|
||||||
@@ -938,6 +1412,11 @@ class RemoveImageBackground(bpy.types.Operator):
|
|||||||
from .operators import IMAGE_BACKGROUND_JOB
|
from .operators import IMAGE_BACKGROUND_JOB
|
||||||
|
|
||||||
source_object = context.object
|
source_object = context.object
|
||||||
|
total_frames = video_total_frames(source_object.data)
|
||||||
|
if total_frames > 1:
|
||||||
|
context.scene.snapform_settings.video_frames = (
|
||||||
|
default_extract_frames(source_object)
|
||||||
|
)
|
||||||
try:
|
try:
|
||||||
source_path, delete_input = prepare_image_input(
|
source_path, delete_input = prepare_image_input(
|
||||||
source_object.data,
|
source_object.data,
|
||||||
@@ -956,7 +1435,7 @@ class RemoveImageBackground(bpy.types.Operator):
|
|||||||
if "CANCELLED" not in result:
|
if "CANCELLED" not in result:
|
||||||
return {"FINISHED"}
|
return {"FINISHED"}
|
||||||
if delete_input:
|
if delete_input:
|
||||||
source_path.unlink(missing_ok=True)
|
remove_temporary_input(source_path)
|
||||||
return {"CANCELLED"}
|
return {"CANCELLED"}
|
||||||
|
|
||||||
|
|
||||||
@@ -1089,13 +1568,45 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
|
|||||||
layout.prop(self, "snapform_delete_original")
|
layout.prop(self, "snapform_delete_original")
|
||||||
|
|
||||||
def execute(self, context):
|
def execute(self, context):
|
||||||
from .operators import IMAGE_DEPTH_JOB, ORIGINAL_BACKGROUND_JOB
|
from .operators import (
|
||||||
|
IMAGE_DEPTH_JOB,
|
||||||
|
ORIGINAL_BACKGROUND_JOB,
|
||||||
|
ORIGINAL_FRAMES_JOB,
|
||||||
|
)
|
||||||
|
|
||||||
source_object = context.object
|
source_object = context.object
|
||||||
|
if video_total_frames(source_object.data) > 1:
|
||||||
|
context.scene.snapform_settings.video_frames = (
|
||||||
|
default_extract_frames(source_object)
|
||||||
|
)
|
||||||
if (
|
if (
|
||||||
self.snapform_depth_source == "ORIGINAL"
|
self.snapform_depth_source == "ORIGINAL"
|
||||||
and not self.snapform_remove_background
|
and not self.snapform_remove_background
|
||||||
):
|
):
|
||||||
|
if video_total_frames(source_object.data) > 1:
|
||||||
|
try:
|
||||||
|
source_path, delete_input = prepare_image_input(
|
||||||
|
source_object.data,
|
||||||
|
context.scene,
|
||||||
|
)
|
||||||
|
except RuntimeError as error:
|
||||||
|
self.report({"ERROR"}, str(error))
|
||||||
|
return {"CANCELLED"}
|
||||||
|
result = bpy.ops.snapform.run_job(
|
||||||
|
job_type=ORIGINAL_FRAMES_JOB,
|
||||||
|
snapform_input_path=str(source_path),
|
||||||
|
snapform_source_object_name=source_object.name,
|
||||||
|
snapform_subdivision_level=self.snapform_subdivision_level,
|
||||||
|
snapform_delete_source=self.snapform_delete_original,
|
||||||
|
snapform_delete_input=delete_input,
|
||||||
|
snapform_displacement_via=self.snapform_displacement_via,
|
||||||
|
snapform_unshaded=self.snapform_unshaded,
|
||||||
|
)
|
||||||
|
if "CANCELLED" not in result:
|
||||||
|
return {"FINISHED"}
|
||||||
|
if delete_input:
|
||||||
|
remove_temporary_input(source_path)
|
||||||
|
return {"CANCELLED"}
|
||||||
try:
|
try:
|
||||||
build_image_depth_plane(
|
build_image_depth_plane(
|
||||||
context,
|
context,
|
||||||
@@ -1141,7 +1652,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
|
|||||||
if "CANCELLED" not in result:
|
if "CANCELLED" not in result:
|
||||||
return {"FINISHED"}
|
return {"FINISHED"}
|
||||||
if delete_input:
|
if delete_input:
|
||||||
source_path.unlink(missing_ok=True)
|
remove_temporary_input(source_path)
|
||||||
return {"CANCELLED"}
|
return {"CANCELLED"}
|
||||||
|
|
||||||
image = source_object.data
|
image = source_object.data
|
||||||
@@ -1179,7 +1690,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
|
|||||||
)
|
)
|
||||||
if "CANCELLED" in result:
|
if "CANCELLED" in result:
|
||||||
if delete_input:
|
if delete_input:
|
||||||
source_path.unlink(missing_ok=True)
|
remove_temporary_input(source_path)
|
||||||
return {"CANCELLED"}
|
return {"CANCELLED"}
|
||||||
global _last_depth_plane_model
|
global _last_depth_plane_model
|
||||||
_last_depth_plane_model = self.snapform_model
|
_last_depth_plane_model = self.snapform_model
|
||||||
|
|||||||
@@ -1,3 +1,4 @@
|
|||||||
|
import shutil
|
||||||
import subprocess
|
import subprocess
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -18,8 +19,35 @@ IMAGE_DEPTH_JOB = "IMAGE_DEPTH"
|
|||||||
BACKGROUND_REMOVAL_JOB = "REMOVE_BACKGROUND"
|
BACKGROUND_REMOVAL_JOB = "REMOVE_BACKGROUND"
|
||||||
IMAGE_BACKGROUND_JOB = "IMAGE_BACKGROUND"
|
IMAGE_BACKGROUND_JOB = "IMAGE_BACKGROUND"
|
||||||
ORIGINAL_BACKGROUND_JOB = "ORIGINAL_BACKGROUND"
|
ORIGINAL_BACKGROUND_JOB = "ORIGINAL_BACKGROUND"
|
||||||
|
ORIGINAL_FRAMES_JOB = "ORIGINAL_FRAMES"
|
||||||
IMAGE_DEPTH_DEFAULT_PROCESS_RES = 1024
|
IMAGE_DEPTH_DEFAULT_PROCESS_RES = 1024
|
||||||
IMAGE_DEPTH_PROCESS_RES_METHOD = "upper_bound_resize"
|
IMAGE_DEPTH_PROCESS_RES_METHOD = "upper_bound_resize"
|
||||||
|
BEN2_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"}
|
||||||
|
BEN2_VIDEO_SUFFIXES = {
|
||||||
|
".mp4",
|
||||||
|
".mov",
|
||||||
|
".avi",
|
||||||
|
".mkv",
|
||||||
|
".webm",
|
||||||
|
".m4v",
|
||||||
|
".mts",
|
||||||
|
".mpg",
|
||||||
|
".mpeg",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def is_ben2_input(input_path):
|
||||||
|
return input_path.is_dir() or (
|
||||||
|
input_path.suffix.lower()
|
||||||
|
in BEN2_IMAGE_SUFFIXES | BEN2_VIDEO_SUFFIXES
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def empty_image_start_frame(source_object):
|
||||||
|
"""Return the first frame to extract, honoring the Empty's Offset within the clip length."""
|
||||||
|
from .image_depth import extract_frame_range
|
||||||
|
|
||||||
|
return extract_frame_range(source_object)[0]
|
||||||
|
|
||||||
|
|
||||||
def open_job_directory(result_directory):
|
def open_job_directory(result_directory):
|
||||||
@@ -81,8 +109,12 @@ def close_active_job():
|
|||||||
if _active_job.log_file is not None:
|
if _active_job.log_file is not None:
|
||||||
_active_job.log_file.close()
|
_active_job.log_file.close()
|
||||||
if _active_job.temporary_input_path is not None:
|
if _active_job.temporary_input_path is not None:
|
||||||
|
temporary_path = _active_job.temporary_input_path
|
||||||
try:
|
try:
|
||||||
_active_job.temporary_input_path.unlink(missing_ok=True)
|
if temporary_path.is_dir():
|
||||||
|
shutil.rmtree(temporary_path, ignore_errors=True)
|
||||||
|
else:
|
||||||
|
temporary_path.unlink(missing_ok=True)
|
||||||
except OSError:
|
except OSError:
|
||||||
pass
|
pass
|
||||||
_active_job = None
|
_active_job = None
|
||||||
@@ -102,6 +134,7 @@ class RunJob(bpy.types.Operator):
|
|||||||
(BACKGROUND_REMOVAL_JOB, "Test BEN2 Background Removal", ""),
|
(BACKGROUND_REMOVAL_JOB, "Test BEN2 Background Removal", ""),
|
||||||
(IMAGE_BACKGROUND_JOB, "Remove Image Empty Background", ""),
|
(IMAGE_BACKGROUND_JOB, "Remove Image Empty Background", ""),
|
||||||
(ORIGINAL_BACKGROUND_JOB, "Build Original Background Plane", ""),
|
(ORIGINAL_BACKGROUND_JOB, "Build Original Background Plane", ""),
|
||||||
|
(ORIGINAL_FRAMES_JOB, "Extract Original Depth Frames", ""),
|
||||||
),
|
),
|
||||||
options={"HIDDEN", "SKIP_SAVE"},
|
options={"HIDDEN", "SKIP_SAVE"},
|
||||||
)
|
)
|
||||||
@@ -227,6 +260,7 @@ class RunJob(bpy.types.Operator):
|
|||||||
IMAGE_DEPTH_JOB,
|
IMAGE_DEPTH_JOB,
|
||||||
IMAGE_BACKGROUND_JOB,
|
IMAGE_BACKGROUND_JOB,
|
||||||
ORIGINAL_BACKGROUND_JOB,
|
ORIGINAL_BACKGROUND_JOB,
|
||||||
|
ORIGINAL_FRAMES_JOB,
|
||||||
}
|
}
|
||||||
and self.snapform_delete_input
|
and self.snapform_delete_input
|
||||||
else None
|
else None
|
||||||
@@ -267,6 +301,7 @@ class RunJob(bpy.types.Operator):
|
|||||||
IMAGE_DEPTH_JOB,
|
IMAGE_DEPTH_JOB,
|
||||||
IMAGE_BACKGROUND_JOB,
|
IMAGE_BACKGROUND_JOB,
|
||||||
ORIGINAL_BACKGROUND_JOB,
|
ORIGINAL_BACKGROUND_JOB,
|
||||||
|
ORIGINAL_FRAMES_JOB,
|
||||||
}:
|
}:
|
||||||
configured_input = self.snapform_input_path
|
configured_input = self.snapform_input_path
|
||||||
elif self.job_type == BACKGROUND_REMOVAL_JOB:
|
elif self.job_type == BACKGROUND_REMOVAL_JOB:
|
||||||
@@ -274,20 +309,24 @@ class RunJob(bpy.types.Operator):
|
|||||||
else:
|
else:
|
||||||
configured_input = settings.input_path
|
configured_input = settings.input_path
|
||||||
input_path = Path(bpy.path.abspath(configured_input))
|
input_path = Path(bpy.path.abspath(configured_input))
|
||||||
if not input_path.is_file():
|
if not (input_path.is_file() or input_path.is_dir()):
|
||||||
raise ValueError("Select a valid image or video file")
|
raise ValueError("Select a valid image, video, or frame directory")
|
||||||
|
start_frame = 1
|
||||||
|
if getattr(self, "snapform_source_object_name", ""):
|
||||||
|
source_object = bpy.data.objects.get(
|
||||||
|
self.snapform_source_object_name
|
||||||
|
)
|
||||||
|
if source_object is not None:
|
||||||
|
start_frame = empty_image_start_frame(source_object)
|
||||||
if self.job_type in {
|
if self.job_type in {
|
||||||
BACKGROUND_REMOVAL_JOB,
|
BACKGROUND_REMOVAL_JOB,
|
||||||
IMAGE_BACKGROUND_JOB,
|
IMAGE_BACKGROUND_JOB,
|
||||||
ORIGINAL_BACKGROUND_JOB,
|
ORIGINAL_BACKGROUND_JOB,
|
||||||
}:
|
}:
|
||||||
if input_path.suffix.lower() not in {
|
if not is_ben2_input(input_path):
|
||||||
".jpg",
|
raise ValueError(
|
||||||
".jpeg",
|
"BEN2 requires an image, video, or directory of image frames"
|
||||||
".png",
|
)
|
||||||
".webp",
|
|
||||||
}:
|
|
||||||
raise ValueError("BEN2 requires a JPG, JPEG, PNG, or WEBP image")
|
|
||||||
if not runtime.model_ready(BEN2_MODEL.key):
|
if not runtime.model_ready(BEN2_MODEL.key):
|
||||||
raise RuntimeError(f"Download {BEN2_MODEL.label} first")
|
raise RuntimeError(f"Download {BEN2_MODEL.label} first")
|
||||||
return (
|
return (
|
||||||
@@ -297,6 +336,18 @@ class RunJob(bpy.types.Operator):
|
|||||||
paths["status"],
|
paths["status"],
|
||||||
paths["result"],
|
paths["result"],
|
||||||
BEN2_MODEL.key,
|
BEN2_MODEL.key,
|
||||||
|
start_frame=start_frame,
|
||||||
|
),
|
||||||
|
paths,
|
||||||
|
)
|
||||||
|
if self.job_type == ORIGINAL_FRAMES_JOB:
|
||||||
|
return (
|
||||||
|
runtime.extract_frames_request(
|
||||||
|
settings,
|
||||||
|
input_path,
|
||||||
|
paths["status"],
|
||||||
|
paths["result"],
|
||||||
|
start_frame=start_frame,
|
||||||
),
|
),
|
||||||
paths,
|
paths,
|
||||||
)
|
)
|
||||||
@@ -342,6 +393,7 @@ class RunJob(bpy.types.Operator):
|
|||||||
else None
|
else None
|
||||||
),
|
),
|
||||||
use_ray_pose=use_ray_pose,
|
use_ray_pose=use_ray_pose,
|
||||||
|
start_frame=start_frame,
|
||||||
remove_background=(
|
remove_background=(
|
||||||
self.job_type == IMAGE_DEPTH_JOB
|
self.job_type == IMAGE_DEPTH_JOB
|
||||||
and remove_background
|
and remove_background
|
||||||
@@ -371,6 +423,8 @@ class RunJob(bpy.types.Operator):
|
|||||||
return "Starting Image Empty background removal"
|
return "Starting Image Empty background removal"
|
||||||
if self.job_type == ORIGINAL_BACKGROUND_JOB:
|
if self.job_type == ORIGINAL_BACKGROUND_JOB:
|
||||||
return "Starting background removal"
|
return "Starting background removal"
|
||||||
|
if self.job_type == ORIGINAL_FRAMES_JOB:
|
||||||
|
return "Extracting source frames"
|
||||||
if self.job_type == "DOWNLOAD_MODEL":
|
if self.job_type == "DOWNLOAD_MODEL":
|
||||||
model = get_downloadable_model(self.model_key)
|
model = get_downloadable_model(self.model_key)
|
||||||
return f"Starting {model.label} download"
|
return f"Starting {model.label} download"
|
||||||
@@ -484,7 +538,7 @@ class RunJob(bpy.types.Operator):
|
|||||||
self._redraw_ui(context, force=True)
|
self._redraw_ui(context, force=True)
|
||||||
return {"CANCELLED"}
|
return {"CANCELLED"}
|
||||||
status["message"] = "Depth plane is ready"
|
status["message"] = "Depth plane is ready"
|
||||||
if job.kind == ORIGINAL_BACKGROUND_JOB:
|
if job.kind in {ORIGINAL_BACKGROUND_JOB, ORIGINAL_FRAMES_JOB}:
|
||||||
try:
|
try:
|
||||||
settings.status = "Building depth plane"
|
settings.status = "Building depth plane"
|
||||||
settings.progress = 0.96
|
settings.progress = 0.96
|
||||||
@@ -515,24 +569,39 @@ class RunJob(bpy.types.Operator):
|
|||||||
|
|
||||||
settings.status = status.get("message", "Task complete")
|
settings.status = status.get("message", "Task complete")
|
||||||
settings.progress = 1.0
|
settings.progress = 1.0
|
||||||
|
try:
|
||||||
|
job.status_path.unlink(missing_ok=True)
|
||||||
|
except OSError:
|
||||||
|
pass
|
||||||
close_active_job()
|
close_active_job()
|
||||||
self._redraw_ui(context, force=True)
|
self._redraw_ui(context, force=True)
|
||||||
return {"FINISHED"}
|
return {"FINISHED"}
|
||||||
|
|
||||||
def _build_depth_plane(self, context, job):
|
def _build_depth_plane(self, context, job):
|
||||||
npz_path = (
|
result_directory = job.result_directory
|
||||||
job.result_directory
|
npz_path = result_directory / "results.npz"
|
||||||
/ "exports"
|
is_sequence = bool(list(result_directory.glob("frame.*")))
|
||||||
/ "mini_npz"
|
foreground_paths = sorted(result_directory.glob("foreground.*.png"))
|
||||||
/ "results.npz"
|
background_path = ""
|
||||||
)
|
if job.remove_background:
|
||||||
|
background_path = (
|
||||||
|
str(result_directory)
|
||||||
|
if foreground_paths
|
||||||
|
else str(result_directory / "foreground.png")
|
||||||
|
)
|
||||||
result = bpy.ops.snapform.build_depth_plane(
|
result = bpy.ops.snapform.build_depth_plane(
|
||||||
snapform_source_object_name=job.source_object_name,
|
snapform_source_object_name=job.source_object_name,
|
||||||
snapform_npz_path=str(npz_path),
|
snapform_npz_path=str(npz_path),
|
||||||
snapform_depth_data_name=job.depth_data_name,
|
snapform_depth_data_name=job.depth_data_name,
|
||||||
snapform_background_path=(
|
snapform_background_path=background_path,
|
||||||
str(job.result_directory / "foreground.png")
|
snapform_color_frames_dir=(
|
||||||
if job.remove_background
|
str(result_directory)
|
||||||
|
if is_sequence and not foreground_paths
|
||||||
|
else ""
|
||||||
|
),
|
||||||
|
snapform_depth_frames_dir=(
|
||||||
|
str(result_directory)
|
||||||
|
if is_sequence
|
||||||
else ""
|
else ""
|
||||||
),
|
),
|
||||||
snapform_subdivision_level=job.subdivision_level,
|
snapform_subdivision_level=job.subdivision_level,
|
||||||
@@ -547,10 +616,24 @@ class RunJob(bpy.types.Operator):
|
|||||||
raise RuntimeError("Unable to create the depth plane Undo step")
|
raise RuntimeError("Unable to create the depth plane Undo step")
|
||||||
|
|
||||||
def _build_original_depth_plane(self, job):
|
def _build_original_depth_plane(self, job):
|
||||||
|
result_directory = job.result_directory
|
||||||
|
is_sequence = bool(list(result_directory.glob("frame.*")))
|
||||||
|
background_path = ""
|
||||||
|
color_frames_dir = ""
|
||||||
|
if job.kind == ORIGINAL_BACKGROUND_JOB:
|
||||||
|
background_path = str(
|
||||||
|
result_directory
|
||||||
|
if is_sequence
|
||||||
|
else result_directory / "foreground.png"
|
||||||
|
)
|
||||||
|
elif is_sequence:
|
||||||
|
color_frames_dir = str(result_directory)
|
||||||
result = bpy.ops.snapform.build_original_depth_plane(
|
result = bpy.ops.snapform.build_original_depth_plane(
|
||||||
snapform_source_object_name=job.source_object_name,
|
snapform_source_object_name=job.source_object_name,
|
||||||
snapform_background_path=str(
|
snapform_background_path=background_path,
|
||||||
job.result_directory / "foreground.png"
|
snapform_color_frames_dir=color_frames_dir,
|
||||||
|
snapform_depth_frames_dir=(
|
||||||
|
str(result_directory) if is_sequence else ""
|
||||||
),
|
),
|
||||||
snapform_subdivision_level=job.subdivision_level,
|
snapform_subdivision_level=job.subdivision_level,
|
||||||
snapform_delete_source=job.delete_source,
|
snapform_delete_source=job.delete_source,
|
||||||
@@ -564,7 +647,12 @@ class RunJob(bpy.types.Operator):
|
|||||||
raise RuntimeError("Unable to create the depth plane Undo step")
|
raise RuntimeError("Unable to create the depth plane Undo step")
|
||||||
|
|
||||||
def _apply_removed_background(self, job, push_undo=True):
|
def _apply_removed_background(self, job, push_undo=True):
|
||||||
output_path = job.result_directory / "foreground.png"
|
result_directory = job.result_directory
|
||||||
|
foreground_paths = sorted(result_directory.glob("foreground.*.png"))
|
||||||
|
output_path = (
|
||||||
|
result_directory if foreground_paths
|
||||||
|
else result_directory / "foreground.png"
|
||||||
|
)
|
||||||
result = bpy.ops.snapform.apply_removed_background(
|
result = bpy.ops.snapform.apply_removed_background(
|
||||||
snapform_source_object_name=job.source_object_name,
|
snapform_source_object_name=job.source_object_name,
|
||||||
snapform_output_path=str(output_path),
|
snapform_output_path=str(output_path),
|
||||||
|
|||||||
@@ -35,6 +35,14 @@ class SnapFormPreferences(bpy.types.AddonPreferences):
|
|||||||
),
|
),
|
||||||
default="CPU" if sys.platform == "darwin" else "CU128",
|
default="CPU" if sys.platform == "darwin" else "CU128",
|
||||||
)
|
)
|
||||||
|
pypi_index_url: bpy.props.StringProperty(
|
||||||
|
name="PyPI Index",
|
||||||
|
description=(
|
||||||
|
"Package index for the isolated Runtime; leave empty to auto-detect "
|
||||||
|
"with mirror fallback. Example: https://pypi.tuna.tsinghua.edu.cn/simple"
|
||||||
|
),
|
||||||
|
default="",
|
||||||
|
)
|
||||||
|
|
||||||
def draw(self, _context):
|
def draw(self, _context):
|
||||||
from .operators import active_job
|
from .operators import active_job
|
||||||
@@ -44,6 +52,7 @@ class SnapFormPreferences(bpy.types.AddonPreferences):
|
|||||||
layout.prop(self, "storage_root")
|
layout.prop(self, "storage_root")
|
||||||
layout.prop(self, "uv_executable")
|
layout.prop(self, "uv_executable")
|
||||||
layout.prop(self, "torch_backend")
|
layout.prop(self, "torch_backend")
|
||||||
|
layout.prop(self, "pypi_index_url")
|
||||||
layout.separator()
|
layout.separator()
|
||||||
layout.label(text=f"Runtime: {runtime_directory()}")
|
layout.label(text=f"Runtime: {runtime_directory()}")
|
||||||
if runtime_ready():
|
if runtime_ready():
|
||||||
|
|||||||
@@ -15,8 +15,8 @@ class SnapFormSettings(bpy.types.PropertyGroup):
|
|||||||
subtype="FILE_PATH",
|
subtype="FILE_PATH",
|
||||||
)
|
)
|
||||||
video_frames: bpy.props.IntProperty(
|
video_frames: bpy.props.IntProperty(
|
||||||
name="Video Frames",
|
name="Extract Frames",
|
||||||
description="Maximum number of frames sampled evenly from a video",
|
description="Maximum number of consecutive frames extracted from the start of a video or frame sequence",
|
||||||
default=16,
|
default=16,
|
||||||
min=2,
|
min=2,
|
||||||
soft_max=64,
|
soft_max=64,
|
||||||
|
|||||||
@@ -24,6 +24,13 @@ RUNTIME_SCHEMA = 4
|
|||||||
HOST_SCHEMA = 2
|
HOST_SCHEMA = 2
|
||||||
HOST_ADDRESS = "127.0.0.1"
|
HOST_ADDRESS = "127.0.0.1"
|
||||||
DEFAULT_STORAGE_ROOT = Path.home() / ".snapform"
|
DEFAULT_STORAGE_ROOT = Path.home() / ".snapform"
|
||||||
|
JOB_FOLDER_NAMES = {
|
||||||
|
"INSTALL_RUNTIME": "install-runtime",
|
||||||
|
"DOWNLOAD_MODEL": "download-model",
|
||||||
|
"IMAGE_DEPTH": "depth-plane",
|
||||||
|
"IMAGE_BACKGROUND": "remove-background",
|
||||||
|
"ORIGINAL_FRAMES": "depth-plane",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True)
|
@dataclass(frozen=True)
|
||||||
@@ -122,15 +129,26 @@ def temporary_input_path():
|
|||||||
return path / f"{uuid.uuid4().hex}.png"
|
return path / f"{uuid.uuid4().hex}.png"
|
||||||
|
|
||||||
|
|
||||||
|
def temporary_input_directory():
|
||||||
|
path = jobs_directory() / "inputs" / f"{uuid.uuid4().hex}"
|
||||||
|
path.mkdir(parents=True, exist_ok=True)
|
||||||
|
return path
|
||||||
|
|
||||||
|
|
||||||
|
def job_folder_name(job_type):
|
||||||
|
"""Return the folder name segment for a job type."""
|
||||||
|
return JOB_FOLDER_NAMES.get(job_type, job_type.lower())
|
||||||
|
|
||||||
|
|
||||||
def new_job_paths(job_type):
|
def new_job_paths(job_type):
|
||||||
timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
|
timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
|
||||||
token = uuid.uuid4().hex[:8]
|
stem = f"{job_folder_name(job_type)}_{timestamp}"
|
||||||
stem = f"{timestamp}-{job_type.lower()}-{token}"
|
directory = jobs_directory() / stem
|
||||||
directory = jobs_directory()
|
directory.mkdir(parents=True, exist_ok=True)
|
||||||
return {
|
return {
|
||||||
"log": directory / f"{stem}.log",
|
"log": directory / f"{directory.name}.log",
|
||||||
"status": directory / f"{stem}.json",
|
"status": directory / f"{directory.name}.json",
|
||||||
"result": directory / stem,
|
"result": directory,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
@@ -155,8 +173,9 @@ def install_runtime_command(status_path):
|
|||||||
executable = uv_executable()
|
executable = uv_executable()
|
||||||
addon_preferences = preferences()
|
addon_preferences = preferences()
|
||||||
backend = getattr(addon_preferences, "torch_backend", "CU128")
|
backend = getattr(addon_preferences, "torch_backend", "CU128")
|
||||||
|
pypi_index = getattr(addon_preferences, "pypi_index_url", "") or ""
|
||||||
script = Path(__file__).with_name("bootstrap_runtime.py")
|
script = Path(__file__).with_name("bootstrap_runtime.py")
|
||||||
return [
|
command = [
|
||||||
sys.executable,
|
sys.executable,
|
||||||
str(script),
|
str(script),
|
||||||
"--uv",
|
"--uv",
|
||||||
@@ -174,6 +193,9 @@ def install_runtime_command(status_path):
|
|||||||
"--status",
|
"--status",
|
||||||
str(status_path),
|
str(status_path),
|
||||||
]
|
]
|
||||||
|
if pypi_index:
|
||||||
|
command.extend(["--pypi-index", pypi_index])
|
||||||
|
return command
|
||||||
|
|
||||||
|
|
||||||
def host_log_path():
|
def host_log_path():
|
||||||
@@ -616,6 +638,8 @@ def remove_background_request(
|
|||||||
status_path,
|
status_path,
|
||||||
result_directory,
|
result_directory,
|
||||||
model_key,
|
model_key,
|
||||||
|
video_frames=None,
|
||||||
|
start_frame=1,
|
||||||
):
|
):
|
||||||
return HostJobRequest(
|
return HostJobRequest(
|
||||||
{
|
{
|
||||||
@@ -627,6 +651,39 @@ def remove_background_request(
|
|||||||
"output_dir": str(result_directory),
|
"output_dir": str(result_directory),
|
||||||
"status": str(status_path),
|
"status": str(status_path),
|
||||||
"device": settings.device.lower(),
|
"device": settings.device.lower(),
|
||||||
|
"video_frames": (
|
||||||
|
settings.video_frames
|
||||||
|
if video_frames is None
|
||||||
|
else video_frames
|
||||||
|
),
|
||||||
|
"start_frame": start_frame,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def extract_frames_request(
|
||||||
|
settings,
|
||||||
|
input_path,
|
||||||
|
status_path,
|
||||||
|
result_directory,
|
||||||
|
video_frames=None,
|
||||||
|
start_frame=1,
|
||||||
|
):
|
||||||
|
return HostJobRequest(
|
||||||
|
{
|
||||||
|
"job_id": Path(status_path).stem,
|
||||||
|
"command": "extract-frames",
|
||||||
|
"parameters": {
|
||||||
|
"input": str(input_path),
|
||||||
|
"output_dir": str(result_directory),
|
||||||
|
"status": str(status_path),
|
||||||
|
"video_frames": (
|
||||||
|
settings.video_frames
|
||||||
|
if video_frames is None
|
||||||
|
else video_frames
|
||||||
|
),
|
||||||
|
"start_frame": start_frame,
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
@@ -643,6 +700,7 @@ def reconstruct_request(
|
|||||||
process_res_method=None,
|
process_res_method=None,
|
||||||
use_ray_pose=None,
|
use_ray_pose=None,
|
||||||
remove_background=False,
|
remove_background=False,
|
||||||
|
start_frame=1,
|
||||||
):
|
):
|
||||||
snapform_selected_key = model_key or selected_model_key(settings, export_format)
|
snapform_selected_key = model_key or selected_model_key(settings, export_format)
|
||||||
snapform_selected_process_res = (
|
snapform_selected_process_res = (
|
||||||
@@ -673,6 +731,7 @@ def reconstruct_request(
|
|||||||
"show_cameras": 1 if settings.show_cameras else 0,
|
"show_cameras": 1 if settings.show_cameras else 0,
|
||||||
"gs_views_interval": settings.gs_views_interval,
|
"gs_views_interval": settings.gs_views_interval,
|
||||||
"video_frames": settings.video_frames,
|
"video_frames": settings.video_frames,
|
||||||
|
"start_frame": start_frame,
|
||||||
"device": settings.device.lower(),
|
"device": settings.device.lower(),
|
||||||
"use_ray_pose": 1 if snapform_selected_ray_pose else 0,
|
"use_ray_pose": 1 if snapform_selected_ray_pose else 0,
|
||||||
"ref_view_strategy": settings.ref_view_strategy,
|
"ref_view_strategy": settings.ref_view_strategy,
|
||||||
|
|||||||
@@ -1,13 +1,14 @@
|
|||||||
import argparse
|
import argparse
|
||||||
|
import fnmatch
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
import shutil
|
import shutil
|
||||||
import tempfile
|
import tempfile
|
||||||
import threading
|
import threading
|
||||||
import time
|
import time
|
||||||
import urllib.error
|
|
||||||
import urllib.request
|
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
from urllib.parse import quote, urlencode
|
||||||
|
from urllib.request import Request, urlopen
|
||||||
|
|
||||||
|
|
||||||
IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"}
|
IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"}
|
||||||
@@ -24,9 +25,10 @@ VIDEO_SUFFIXES = {
|
|||||||
}
|
}
|
||||||
SCENE_FILENAME = "scene.glb"
|
SCENE_FILENAME = "scene.glb"
|
||||||
EXPORT_FORMATS = ("glb", "gs_ply", "depth_vis", "mini_npz")
|
EXPORT_FORMATS = ("glb", "gs_ply", "depth_vis", "mini_npz")
|
||||||
HF_ENDPOINT = "https://huggingface.co/"
|
HF_ENDPOINT = "https://huggingface.co"
|
||||||
HF_MIRRORS = ("https://hf-mirror.com",)
|
MODELSCOPE_ENDPOINT = "https://modelscope.cn"
|
||||||
HF_PROBE_TIMEOUT = 5.0
|
MODELSCOPE_REVISION = "master"
|
||||||
|
MODELSCOPE_TIMEOUT = 60.0
|
||||||
|
|
||||||
|
|
||||||
class TaskCancelled(RuntimeError):
|
class TaskCancelled(RuntimeError):
|
||||||
@@ -38,34 +40,6 @@ def check_cancelled(cancel_check):
|
|||||||
raise TaskCancelled("Task cancelled")
|
raise TaskCancelled("Task cancelled")
|
||||||
|
|
||||||
|
|
||||||
def reachable(url, timeout=HF_PROBE_TIMEOUT):
|
|
||||||
request = urllib.request.Request(
|
|
||||||
url,
|
|
||||||
method="HEAD",
|
|
||||||
headers={"User-Agent": "SnapForm-Runtime-Installer"},
|
|
||||||
)
|
|
||||||
try:
|
|
||||||
with urllib.request.urlopen(request, timeout=timeout):
|
|
||||||
return True
|
|
||||||
except urllib.error.HTTPError:
|
|
||||||
return True
|
|
||||||
except Exception:
|
|
||||||
return False
|
|
||||||
|
|
||||||
|
|
||||||
def resolve_hf_endpoint():
|
|
||||||
"""Probe Hugging Face reachability and fall back to a mirror when direct access fails; returns (endpoint, used_mirror)."""
|
|
||||||
configured = os.environ.get("HF_ENDPOINT")
|
|
||||||
if configured:
|
|
||||||
return configured, False
|
|
||||||
if reachable(HF_ENDPOINT):
|
|
||||||
return HF_ENDPOINT, False
|
|
||||||
for mirror in HF_MIRRORS:
|
|
||||||
if reachable(mirror):
|
|
||||||
return mirror, True
|
|
||||||
return HF_ENDPOINT, False
|
|
||||||
|
|
||||||
|
|
||||||
def write_status(
|
def write_status(
|
||||||
path,
|
path,
|
||||||
progress,
|
progress,
|
||||||
@@ -106,56 +80,66 @@ def collect_images(directory):
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def first_frames(images, maximum_frames, start_frame=1):
|
||||||
|
"""Return up to maximum_frames image paths in order from start_frame."""
|
||||||
|
if maximum_frames <= 0:
|
||||||
|
return images
|
||||||
|
start_index = max(start_frame - 1, 0)
|
||||||
|
return images[start_index:start_index + maximum_frames]
|
||||||
|
|
||||||
|
|
||||||
def extract_video_frames(
|
def extract_video_frames(
|
||||||
video_path,
|
video_path,
|
||||||
output_directory,
|
output_directory,
|
||||||
maximum_frames,
|
maximum_frames,
|
||||||
cv2,
|
cv2,
|
||||||
|
start_frame=1,
|
||||||
cancel_check=None,
|
cancel_check=None,
|
||||||
):
|
):
|
||||||
|
output_directory = Path(output_directory)
|
||||||
|
output_directory.mkdir(parents=True, exist_ok=True)
|
||||||
capture = cv2.VideoCapture(str(video_path))
|
capture = cv2.VideoCapture(str(video_path))
|
||||||
if not capture.isOpened():
|
if not capture.isOpened():
|
||||||
raise RuntimeError(f"Unable to open video: {video_path}")
|
raise RuntimeError(f"Unable to open video: {video_path}")
|
||||||
|
|
||||||
output_directory = Path(output_directory)
|
|
||||||
output_directory.mkdir(parents=True, exist_ok=True)
|
|
||||||
total_frames = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
||||||
if total_frames <= 0:
|
|
||||||
capture.release()
|
|
||||||
raise RuntimeError("The video does not report a readable frame count")
|
|
||||||
|
|
||||||
import numpy as np
|
|
||||||
|
|
||||||
sample_count = min(maximum_frames, total_frames)
|
|
||||||
frame_indices = np.linspace(
|
|
||||||
0,
|
|
||||||
total_frames - 1,
|
|
||||||
sample_count,
|
|
||||||
dtype=np.int64,
|
|
||||||
)
|
|
||||||
image_paths = []
|
image_paths = []
|
||||||
|
skipped = 0
|
||||||
try:
|
try:
|
||||||
for output_index, frame_index in enumerate(frame_indices):
|
while len(image_paths) < maximum_frames:
|
||||||
check_cancelled(cancel_check)
|
check_cancelled(cancel_check)
|
||||||
capture.set(cv2.CAP_PROP_POS_FRAMES, int(frame_index))
|
|
||||||
success, frame = capture.read()
|
success, frame = capture.read()
|
||||||
if not success:
|
if not success:
|
||||||
raise RuntimeError(
|
break
|
||||||
f"Unable to read video frame {int(frame_index)}"
|
if skipped < start_frame - 1:
|
||||||
)
|
skipped += 1
|
||||||
image_path = output_directory / f"frame-{output_index:04d}.jpg"
|
continue
|
||||||
|
image_path = (
|
||||||
|
output_directory / f"frame.{len(image_paths) + 1:04d}.jpg"
|
||||||
|
)
|
||||||
if not cv2.imwrite(str(image_path), frame):
|
if not cv2.imwrite(str(image_path), frame):
|
||||||
raise RuntimeError(f"Unable to write video frame: {image_path}")
|
raise RuntimeError(f"Unable to write video frame: {image_path}")
|
||||||
image_paths.append(image_path)
|
image_paths.append(image_path)
|
||||||
finally:
|
finally:
|
||||||
capture.release()
|
capture.release()
|
||||||
|
|
||||||
|
if not image_paths:
|
||||||
|
raise RuntimeError(f"Unable to read any video frame from: {video_path}")
|
||||||
return image_paths
|
return image_paths
|
||||||
|
|
||||||
|
|
||||||
def prepare_input(input_path, video_frames, temporary_root, cancel_check=None):
|
def prepare_input(
|
||||||
|
input_path,
|
||||||
|
video_frames,
|
||||||
|
temporary_root,
|
||||||
|
start_frame=1,
|
||||||
|
cancel_check=None,
|
||||||
|
):
|
||||||
check_cancelled(cancel_check)
|
check_cancelled(cancel_check)
|
||||||
input_path = Path(input_path)
|
input_path = Path(input_path)
|
||||||
|
if input_path.is_dir():
|
||||||
|
images = collect_images(input_path)
|
||||||
|
if not images:
|
||||||
|
raise RuntimeError(f"No image frames found in: {input_path}")
|
||||||
|
return first_frames(images, video_frames, start_frame), None
|
||||||
suffix = input_path.suffix.lower()
|
suffix = input_path.suffix.lower()
|
||||||
if suffix in IMAGE_SUFFIXES:
|
if suffix in IMAGE_SUFFIXES:
|
||||||
return [input_path], None
|
return [input_path], None
|
||||||
@@ -174,6 +158,7 @@ def prepare_input(input_path, video_frames, temporary_root, cancel_check=None):
|
|||||||
temporary_directory,
|
temporary_directory,
|
||||||
video_frames,
|
video_frames,
|
||||||
cv2,
|
cv2,
|
||||||
|
start_frame=start_frame,
|
||||||
cancel_check=cancel_check,
|
cancel_check=cancel_check,
|
||||||
)
|
)
|
||||||
except Exception:
|
except Exception:
|
||||||
@@ -182,6 +167,43 @@ def prepare_input(input_path, video_frames, temporary_root, cancel_check=None):
|
|||||||
return images, temporary_directory
|
return images, temporary_directory
|
||||||
|
|
||||||
|
|
||||||
|
def prepare_background_input(
|
||||||
|
input_path,
|
||||||
|
video_frames,
|
||||||
|
temporary_root,
|
||||||
|
start_frame=1,
|
||||||
|
cancel_check=None,
|
||||||
|
):
|
||||||
|
"""Resolve a BEN2 input to frame images; returns (images, temporary_directory)."""
|
||||||
|
input_path = Path(input_path)
|
||||||
|
if input_path.is_dir():
|
||||||
|
images = collect_images(input_path)
|
||||||
|
if not images:
|
||||||
|
raise RuntimeError(f"No image frames found in: {input_path}")
|
||||||
|
return first_frames(images, video_frames, start_frame), None
|
||||||
|
if input_path.suffix.lower() in IMAGE_SUFFIXES:
|
||||||
|
return [input_path], None
|
||||||
|
return prepare_input(
|
||||||
|
input_path,
|
||||||
|
video_frames,
|
||||||
|
temporary_root,
|
||||||
|
start_frame=start_frame,
|
||||||
|
cancel_check=cancel_check,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def copy_frames(images, output_directory):
|
||||||
|
"""Copy a multi-frame input as numbered frame files for the Blender sequence."""
|
||||||
|
if len(images) <= 1:
|
||||||
|
return
|
||||||
|
for index, image_path in enumerate(images):
|
||||||
|
suffix = Path(image_path).suffix.lower() or ".jpg"
|
||||||
|
shutil.copy2(
|
||||||
|
image_path,
|
||||||
|
Path(output_directory) / f"frame.{index + 1:04d}{suffix}",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class DownloadProgressReporter:
|
class DownloadProgressReporter:
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
@@ -318,21 +340,8 @@ def download_model(args):
|
|||||||
"stages": getattr(args, "stages", None),
|
"stages": getattr(args, "stages", None),
|
||||||
"stage_label": getattr(args, "stage_label", None),
|
"stage_label": getattr(args, "stage_label", None),
|
||||||
}
|
}
|
||||||
write_status(
|
|
||||||
args.status,
|
|
||||||
start,
|
|
||||||
"checking download source",
|
|
||||||
**stage_kwargs,
|
|
||||||
)
|
|
||||||
endpoint, used_mirror = resolve_hf_endpoint()
|
|
||||||
if endpoint != os.environ.get("HF_ENDPOINT"):
|
|
||||||
os.environ["HF_ENDPOINT"] = endpoint
|
|
||||||
from huggingface_hub import snapshot_download
|
from huggingface_hub import snapshot_download
|
||||||
|
|
||||||
message = f"downloading {label}"
|
|
||||||
if used_mirror:
|
|
||||||
message = f"{message} (mirror)"
|
|
||||||
write_status(args.status, start, message, **stage_kwargs)
|
|
||||||
progress_class = create_download_progress_class(
|
progress_class = create_download_progress_class(
|
||||||
args.status,
|
args.status,
|
||||||
args.repository,
|
args.repository,
|
||||||
@@ -341,18 +350,213 @@ def download_model(args):
|
|||||||
label=label,
|
label=label,
|
||||||
**stage_kwargs,
|
**stage_kwargs,
|
||||||
)
|
)
|
||||||
snapshot_download(
|
failures = []
|
||||||
repo_id=args.repository,
|
for endpoint in download_endpoints():
|
||||||
local_dir=model_dir,
|
if not hf_endpoint_reachable(endpoint):
|
||||||
tqdm_class=progress_class,
|
failures.append((endpoint, "endpoint unreachable"))
|
||||||
allow_patterns=args.allow_pattern or None,
|
continue
|
||||||
|
os.environ["HF_ENDPOINT"] = endpoint
|
||||||
|
message = f"downloading {label}"
|
||||||
|
if endpoint != HF_ENDPOINT:
|
||||||
|
message = f"{message} (mirror)"
|
||||||
|
write_status(args.status, start, message, **stage_kwargs)
|
||||||
|
try:
|
||||||
|
snapshot_download(
|
||||||
|
repo_id=args.repository,
|
||||||
|
local_dir=model_dir,
|
||||||
|
tqdm_class=progress_class,
|
||||||
|
allow_patterns=args.allow_pattern or None,
|
||||||
|
endpoint=endpoint,
|
||||||
|
)
|
||||||
|
_verify_model_complete(model_dir)
|
||||||
|
except Exception as error:
|
||||||
|
failures.append((endpoint, error))
|
||||||
|
continue
|
||||||
|
write_status(
|
||||||
|
args.status,
|
||||||
|
start + span,
|
||||||
|
"model ready",
|
||||||
|
**stage_kwargs,
|
||||||
|
)
|
||||||
|
return
|
||||||
|
# Hugging Face 官方源与已配置镜像均失败时,回退 ModelScope(仓库名与 Hugging Face 一致)
|
||||||
|
try:
|
||||||
|
write_status(
|
||||||
|
args.status,
|
||||||
|
start,
|
||||||
|
f"downloading {label} (mirror: modelscope.cn)",
|
||||||
|
**stage_kwargs,
|
||||||
|
)
|
||||||
|
download_from_modelscope(
|
||||||
|
args,
|
||||||
|
model_dir,
|
||||||
|
start=start,
|
||||||
|
span=span,
|
||||||
|
label=label,
|
||||||
|
stage_kwargs=stage_kwargs,
|
||||||
|
)
|
||||||
|
_verify_model_complete(model_dir)
|
||||||
|
except Exception as error:
|
||||||
|
failures.append((MODELSCOPE_ENDPOINT, error))
|
||||||
|
else:
|
||||||
|
write_status(
|
||||||
|
args.status,
|
||||||
|
start + span,
|
||||||
|
"model ready",
|
||||||
|
**stage_kwargs,
|
||||||
|
)
|
||||||
|
return
|
||||||
|
detail = "; ".join(
|
||||||
|
f"{source}: {error}" for source, error in failures
|
||||||
)
|
)
|
||||||
write_status(
|
raise RuntimeError(f"All download sources failed: {detail}")
|
||||||
|
|
||||||
|
|
||||||
|
def download_endpoints():
|
||||||
|
"""Hugging Face download candidates: configured mirror first, official endpoint otherwise."""
|
||||||
|
endpoint = os.environ.get("HF_ENDPOINT") or HF_ENDPOINT
|
||||||
|
return [endpoint.rstrip("/")]
|
||||||
|
|
||||||
|
|
||||||
|
def hf_endpoint_reachable(endpoint, timeout=10.0):
|
||||||
|
"""Quick connectivity probe so slow or blocked Hugging Face endpoints fail fast."""
|
||||||
|
try:
|
||||||
|
request = Request(
|
||||||
|
f"{endpoint.rstrip('/')}/api/models",
|
||||||
|
method="HEAD",
|
||||||
|
headers={"User-Agent": "SnapForm"},
|
||||||
|
)
|
||||||
|
with urlopen(request, timeout=timeout):
|
||||||
|
return True
|
||||||
|
except Exception:
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def _verify_model_complete(model_dir):
|
||||||
|
"""Refuse to report success when the required model files are missing."""
|
||||||
|
model_dir = Path(model_dir)
|
||||||
|
missing = []
|
||||||
|
if not (model_dir / "config.json").is_file():
|
||||||
|
missing.append("config.json")
|
||||||
|
if not any(model_dir.glob("*.safetensors")):
|
||||||
|
missing.append("*.safetensors")
|
||||||
|
if missing:
|
||||||
|
raise RuntimeError(
|
||||||
|
"Model download reported success but files are missing: "
|
||||||
|
+ ", ".join(missing)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _hidden_repo_path(path):
|
||||||
|
return any(part.startswith(".") for part in Path(path).parts)
|
||||||
|
|
||||||
|
|
||||||
|
def _model_file_allowed(path, patterns):
|
||||||
|
if not patterns:
|
||||||
|
return True
|
||||||
|
return any(fnmatch.fnmatch(path, pattern) for pattern in patterns)
|
||||||
|
|
||||||
|
|
||||||
|
def modelscope_file_list(repository, revision=MODELSCOPE_REVISION):
|
||||||
|
"""List (path, size) file entries of a ModelScope repository."""
|
||||||
|
url = (
|
||||||
|
f"{MODELSCOPE_ENDPOINT}/api/v1/models/{repository}/repo/files?"
|
||||||
|
+ urlencode({"Revision": revision, "Recursive": "true"})
|
||||||
|
)
|
||||||
|
request = Request(url, headers={"User-Agent": "SnapForm"})
|
||||||
|
with urlopen(request, timeout=MODELSCOPE_TIMEOUT) as response:
|
||||||
|
payload = json.loads(response.read().decode("utf-8"))
|
||||||
|
entries = []
|
||||||
|
for entry in payload.get("Data", {}).get("Files", []):
|
||||||
|
path = entry.get("Path")
|
||||||
|
if path and entry.get("Type") == "blob" and not _hidden_repo_path(path):
|
||||||
|
entries.append((path, int(entry.get("Size") or 0)))
|
||||||
|
return entries
|
||||||
|
|
||||||
|
|
||||||
|
def _download_modelscope_file(
|
||||||
|
repository,
|
||||||
|
path,
|
||||||
|
destination,
|
||||||
|
progress=None,
|
||||||
|
retries=2,
|
||||||
|
):
|
||||||
|
"""Download one ModelScope file atomically with a per-chunk progress callback."""
|
||||||
|
url = (
|
||||||
|
f"{MODELSCOPE_ENDPOINT}/models/{repository}/resolve/"
|
||||||
|
f"{MODELSCOPE_REVISION}/{quote(path)}"
|
||||||
|
)
|
||||||
|
request = Request(url, headers={"User-Agent": "SnapForm"})
|
||||||
|
last_error = None
|
||||||
|
for _attempt in range(retries):
|
||||||
|
temporary = destination.with_name(f"{destination.name}.part")
|
||||||
|
try:
|
||||||
|
with urlopen(request, timeout=MODELSCOPE_TIMEOUT) as response:
|
||||||
|
with open(temporary, "wb") as output:
|
||||||
|
while True:
|
||||||
|
chunk = response.read(1024 * 1024)
|
||||||
|
if not chunk:
|
||||||
|
break
|
||||||
|
output.write(chunk)
|
||||||
|
if progress is not None:
|
||||||
|
progress(len(chunk))
|
||||||
|
temporary.replace(destination)
|
||||||
|
return
|
||||||
|
except Exception as error:
|
||||||
|
last_error = error
|
||||||
|
try:
|
||||||
|
temporary.unlink()
|
||||||
|
except OSError:
|
||||||
|
pass
|
||||||
|
raise last_error
|
||||||
|
|
||||||
|
|
||||||
|
def download_from_modelscope(
|
||||||
|
args,
|
||||||
|
model_dir,
|
||||||
|
start=0.05,
|
||||||
|
span=0.90,
|
||||||
|
label=None,
|
||||||
|
stage_kwargs=None,
|
||||||
|
):
|
||||||
|
"""Download a model from ModelScope using only standard library requests."""
|
||||||
|
repository = args.repository
|
||||||
|
files = modelscope_file_list(repository)
|
||||||
|
patterns = args.allow_pattern or None
|
||||||
|
selected = [
|
||||||
|
(path, size)
|
||||||
|
for path, size in files
|
||||||
|
if _model_file_allowed(path, patterns)
|
||||||
|
]
|
||||||
|
if not selected:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"No matching files in ModelScope repository: {repository}"
|
||||||
|
)
|
||||||
|
reporter = DownloadProgressReporter(
|
||||||
args.status,
|
args.status,
|
||||||
start + span,
|
repository,
|
||||||
"model ready",
|
start=start,
|
||||||
**stage_kwargs,
|
span=span,
|
||||||
|
label=label or repository,
|
||||||
|
**(stage_kwargs or {}),
|
||||||
)
|
)
|
||||||
|
total_size = sum(size for _path, size in selected)
|
||||||
|
downloaded = 0
|
||||||
|
|
||||||
|
def on_chunk(count):
|
||||||
|
nonlocal downloaded
|
||||||
|
downloaded += count
|
||||||
|
reporter.report(downloaded, total_size)
|
||||||
|
|
||||||
|
for path, _size in selected:
|
||||||
|
destination = Path(model_dir) / path
|
||||||
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
_download_modelscope_file(
|
||||||
|
repository,
|
||||||
|
path,
|
||||||
|
destination,
|
||||||
|
progress=on_chunk,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def format_rate(rate):
|
def format_rate(rate):
|
||||||
@@ -435,6 +639,49 @@ def infer_foreground(
|
|||||||
return foreground, load_ms, inference_ms
|
return foreground, load_ms, inference_ms
|
||||||
|
|
||||||
|
|
||||||
|
def infer_foreground_frames(
|
||||||
|
images,
|
||||||
|
model_dir,
|
||||||
|
device,
|
||||||
|
torch,
|
||||||
|
output_directory,
|
||||||
|
status_path,
|
||||||
|
start_progress,
|
||||||
|
progress_span,
|
||||||
|
model_cache=None,
|
||||||
|
cancel_check=None,
|
||||||
|
):
|
||||||
|
"""Run BEN2 on every frame and write foreground.0001.png etc."""
|
||||||
|
total_load_ms = 0.0
|
||||||
|
total_inference_ms = 0.0
|
||||||
|
for index, image_path in enumerate(images):
|
||||||
|
check_cancelled(cancel_check)
|
||||||
|
write_status(
|
||||||
|
status_path,
|
||||||
|
start_progress + progress_span * index / len(images),
|
||||||
|
f"Removing background ({index + 1}/{len(images)})",
|
||||||
|
)
|
||||||
|
foreground, load_ms, inference_ms = infer_foreground(
|
||||||
|
image_path,
|
||||||
|
model_dir,
|
||||||
|
device,
|
||||||
|
torch,
|
||||||
|
model_cache=model_cache,
|
||||||
|
cancel_check=cancel_check,
|
||||||
|
)
|
||||||
|
output_path = output_directory / f"foreground.{index + 1:04d}.png"
|
||||||
|
foreground.convert("RGBA").save(
|
||||||
|
output_path,
|
||||||
|
format="PNG",
|
||||||
|
compress_level=1,
|
||||||
|
)
|
||||||
|
if not output_path.is_file():
|
||||||
|
raise RuntimeError(f"BEN2 did not write: {output_path}")
|
||||||
|
total_load_ms += load_ms
|
||||||
|
total_inference_ms += inference_ms
|
||||||
|
return total_load_ms, total_inference_ms
|
||||||
|
|
||||||
|
|
||||||
def remove_background(args, model_cache=None, cancel_check=None):
|
def remove_background(args, model_cache=None, cancel_check=None):
|
||||||
total_started = time.perf_counter()
|
total_started = time.perf_counter()
|
||||||
write_status(args.status, 0.02, "Starting BEN2 Runtime")
|
write_status(args.status, 0.02, "Starting BEN2 Runtime")
|
||||||
@@ -446,41 +693,90 @@ def remove_background(args, model_cache=None, cancel_check=None):
|
|||||||
torch = model_cache.torch
|
torch = model_cache.torch
|
||||||
|
|
||||||
output_directory = Path(args.output_dir)
|
output_directory = Path(args.output_dir)
|
||||||
output_directory.mkdir(parents=True, exist_ok=False)
|
output_directory.mkdir(parents=True, exist_ok=True)
|
||||||
device = select_device(args.device, torch)
|
device = select_device(args.device, torch)
|
||||||
|
|
||||||
write_status(args.status, 0.12, f"Loading BEN2 on {device.type}")
|
video_frames = getattr(args, "video_frames", 16)
|
||||||
write_status(args.status, 0.35, "Removing background")
|
start_frame = getattr(args, "start_frame", 1)
|
||||||
foreground, load_ms, inference_ms = infer_foreground(
|
images, temporary_directory = prepare_background_input(
|
||||||
args.input,
|
args.input,
|
||||||
args.model_dir,
|
video_frames,
|
||||||
device,
|
output_directory.parent,
|
||||||
torch,
|
start_frame=start_frame,
|
||||||
model_cache=model_cache,
|
|
||||||
cancel_check=cancel_check,
|
cancel_check=cancel_check,
|
||||||
)
|
)
|
||||||
|
try:
|
||||||
check_cancelled(cancel_check)
|
write_status(args.status, 0.12, f"Loading BEN2 on {device.type}")
|
||||||
write_status(args.status, 0.92, "Writing transparent PNG")
|
if len(images) == 1:
|
||||||
output_path = output_directory / "foreground.png"
|
write_status(args.status, 0.35, "Removing background")
|
||||||
foreground.convert("RGBA").save(
|
foreground, load_ms, inference_ms = infer_foreground(
|
||||||
output_path,
|
images[0],
|
||||||
format="PNG",
|
args.model_dir,
|
||||||
compress_level=1,
|
device,
|
||||||
)
|
torch,
|
||||||
if not output_path.is_file():
|
model_cache=model_cache,
|
||||||
raise RuntimeError("BEN2 did not write a foreground image")
|
cancel_check=cancel_check,
|
||||||
|
)
|
||||||
|
check_cancelled(cancel_check)
|
||||||
|
write_status(args.status, 0.92, "Writing transparent PNG")
|
||||||
|
output_path = output_directory / "foreground.png"
|
||||||
|
foreground.convert("RGBA").save(
|
||||||
|
output_path,
|
||||||
|
format="PNG",
|
||||||
|
compress_level=1,
|
||||||
|
)
|
||||||
|
if not output_path.is_file():
|
||||||
|
raise RuntimeError("BEN2 did not write a foreground image")
|
||||||
|
else:
|
||||||
|
load_ms, inference_ms = infer_foreground_frames(
|
||||||
|
images,
|
||||||
|
args.model_dir,
|
||||||
|
device,
|
||||||
|
torch,
|
||||||
|
output_directory,
|
||||||
|
args.status,
|
||||||
|
0.12,
|
||||||
|
0.7,
|
||||||
|
model_cache=model_cache,
|
||||||
|
cancel_check=cancel_check,
|
||||||
|
)
|
||||||
|
copy_frames(images, output_directory)
|
||||||
|
finally:
|
||||||
|
if temporary_directory is not None:
|
||||||
|
shutil.rmtree(temporary_directory, ignore_errors=True)
|
||||||
|
|
||||||
check_cancelled(cancel_check)
|
check_cancelled(cancel_check)
|
||||||
total_ms = (time.perf_counter() - total_started) * 1000.0
|
total_ms = (time.perf_counter() - total_started) * 1000.0
|
||||||
|
frame_text = "" if len(images) == 1 else f" from {len(images)} frame(s)"
|
||||||
message = (
|
message = (
|
||||||
f"Background removed in {inference_ms:.0f} ms "
|
f"Background removed{frame_text} in {inference_ms:.0f} ms "
|
||||||
f"(model load {load_ms:.0f} ms, total {total_ms:.0f} ms)"
|
f"(model load {load_ms:.0f} ms, total {total_ms:.0f} ms)"
|
||||||
)
|
)
|
||||||
print(message)
|
print(message)
|
||||||
write_status(args.status, 1.0, message)
|
write_status(args.status, 1.0, message)
|
||||||
|
|
||||||
|
|
||||||
|
def extract_frames(args, cancel_check=None):
|
||||||
|
"""Copy source frames into a sequence directory without loading any AI model."""
|
||||||
|
output_directory = Path(args.output_dir)
|
||||||
|
output_directory.mkdir(parents=True, exist_ok=True)
|
||||||
|
images, temporary_directory = prepare_input(
|
||||||
|
args.input,
|
||||||
|
getattr(args, "video_frames", 16),
|
||||||
|
output_directory.parent,
|
||||||
|
start_frame=getattr(args, "start_frame", 1),
|
||||||
|
cancel_check=cancel_check,
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
copy_frames(images, output_directory)
|
||||||
|
finally:
|
||||||
|
if temporary_directory is not None:
|
||||||
|
shutil.rmtree(temporary_directory, ignore_errors=True)
|
||||||
|
message = f"Extracted {len(images)} frame(s)"
|
||||||
|
print(message)
|
||||||
|
write_status(args.status, 1.0, message)
|
||||||
|
|
||||||
|
|
||||||
def exported_paths(output_directory, export_format):
|
def exported_paths(output_directory, export_format):
|
||||||
output_directory = Path(output_directory)
|
output_directory = Path(output_directory)
|
||||||
if export_format == "glb":
|
if export_format == "glb":
|
||||||
@@ -491,7 +787,12 @@ def exported_paths(output_directory, export_format):
|
|||||||
if export_format == "depth_vis":
|
if export_format == "depth_vis":
|
||||||
return sorted((output_directory / "depth_vis").glob("*.jpg"))
|
return sorted((output_directory / "depth_vis").glob("*.jpg"))
|
||||||
if export_format == "mini_npz":
|
if export_format == "mini_npz":
|
||||||
result_path = output_directory / "exports" / "mini_npz" / "results.npz"
|
result_path = output_directory / "results.npz"
|
||||||
|
if result_path.is_file():
|
||||||
|
return [result_path]
|
||||||
|
result_path = (
|
||||||
|
output_directory / "exports" / "mini_npz" / "results.npz"
|
||||||
|
)
|
||||||
return [result_path] if result_path.is_file() else []
|
return [result_path] if result_path.is_file() else []
|
||||||
return []
|
return []
|
||||||
|
|
||||||
@@ -543,12 +844,13 @@ def reconstruct(args, model_cache=None, cancel_check=None):
|
|||||||
from depth_anything_3.utils.export import export
|
from depth_anything_3.utils.export import export
|
||||||
|
|
||||||
output_directory = Path(args.output_dir)
|
output_directory = Path(args.output_dir)
|
||||||
output_directory.mkdir(parents=True, exist_ok=False)
|
output_directory.mkdir(parents=True, exist_ok=True)
|
||||||
write_status(args.status, 0.08, "Preparing input")
|
write_status(args.status, 0.08, "Preparing input")
|
||||||
images, temporary_directory = prepare_input(
|
images, temporary_directory = prepare_input(
|
||||||
args.input,
|
args.input,
|
||||||
args.video_frames,
|
args.video_frames,
|
||||||
output_directory.parent,
|
output_directory.parent,
|
||||||
|
start_frame=getattr(args, "start_frame", 1),
|
||||||
cancel_check=cancel_check,
|
cancel_check=cancel_check,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -562,22 +864,36 @@ def reconstruct(args, model_cache=None, cancel_check=None):
|
|||||||
)
|
)
|
||||||
if remove_background_model_dir:
|
if remove_background_model_dir:
|
||||||
write_status(args.status, 0.12, "Removing image background")
|
write_status(args.status, 0.12, "Removing image background")
|
||||||
foreground, load_ms, inference_ms = infer_foreground(
|
if len(images) == 1:
|
||||||
images[0],
|
foreground, load_ms, inference_ms = infer_foreground(
|
||||||
remove_background_model_dir,
|
images[0],
|
||||||
device,
|
remove_background_model_dir,
|
||||||
torch,
|
device,
|
||||||
model_cache=model_cache,
|
torch,
|
||||||
cancel_check=cancel_check,
|
model_cache=model_cache,
|
||||||
)
|
cancel_check=cancel_check,
|
||||||
check_cancelled(cancel_check)
|
)
|
||||||
foreground.convert("RGBA").save(
|
check_cancelled(cancel_check)
|
||||||
output_directory / "foreground.png",
|
foreground.convert("RGBA").save(
|
||||||
format="PNG",
|
output_directory / "foreground.png",
|
||||||
compress_level=1,
|
format="PNG",
|
||||||
)
|
compress_level=1,
|
||||||
if not (output_directory / "foreground.png").is_file():
|
)
|
||||||
raise RuntimeError("BEN2 did not write a foreground image")
|
if not (output_directory / "foreground.png").is_file():
|
||||||
|
raise RuntimeError("BEN2 did not write a foreground image")
|
||||||
|
else:
|
||||||
|
load_ms, inference_ms = infer_foreground_frames(
|
||||||
|
images,
|
||||||
|
remove_background_model_dir,
|
||||||
|
device,
|
||||||
|
torch,
|
||||||
|
output_directory,
|
||||||
|
args.status,
|
||||||
|
0.12,
|
||||||
|
0.05,
|
||||||
|
model_cache=model_cache,
|
||||||
|
cancel_check=cancel_check,
|
||||||
|
)
|
||||||
print(
|
print(
|
||||||
f"BEN2 inference {inference_ms:.0f} ms; "
|
f"BEN2 inference {inference_ms:.0f} ms; "
|
||||||
f"model load {load_ms:.0f} ms"
|
f"model load {load_ms:.0f} ms"
|
||||||
@@ -636,7 +952,21 @@ def reconstruct(args, model_cache=None, cancel_check=None):
|
|||||||
raise RuntimeError(
|
raise RuntimeError(
|
||||||
f"No {args.export_format} files were exported"
|
f"No {args.export_format} files were exported"
|
||||||
)
|
)
|
||||||
|
if args.export_format == "mini_npz":
|
||||||
|
nested_path = (
|
||||||
|
output_directory / "exports" / "mini_npz" / "results.npz"
|
||||||
|
)
|
||||||
|
result_path = output_directory / "results.npz"
|
||||||
|
if nested_path.is_file():
|
||||||
|
nested_path.replace(result_path)
|
||||||
|
result_paths = [result_path]
|
||||||
|
for directory in (nested_path.parent, nested_path.parent.parent):
|
||||||
|
try:
|
||||||
|
directory.rmdir()
|
||||||
|
except OSError:
|
||||||
|
pass
|
||||||
check_cancelled(cancel_check)
|
check_cancelled(cancel_check)
|
||||||
|
copy_frames(images, output_directory)
|
||||||
write_status(
|
write_status(
|
||||||
args.status,
|
args.status,
|
||||||
1.0,
|
1.0,
|
||||||
@@ -673,6 +1003,15 @@ def parser():
|
|||||||
choices=("auto", "cuda", "mps", "cpu"),
|
choices=("auto", "cuda", "mps", "cpu"),
|
||||||
required=True,
|
required=True,
|
||||||
)
|
)
|
||||||
|
background.add_argument("--video-frames", type=int, default=16)
|
||||||
|
background.add_argument("--start-frame", type=int, default=1)
|
||||||
|
|
||||||
|
frames = commands.add_parser("extract-frames")
|
||||||
|
frames.add_argument("--input", required=True)
|
||||||
|
frames.add_argument("--output-dir", required=True)
|
||||||
|
frames.add_argument("--status", required=True)
|
||||||
|
frames.add_argument("--video-frames", type=int, default=16)
|
||||||
|
frames.add_argument("--start-frame", type=int, default=1)
|
||||||
|
|
||||||
reconstruction = commands.add_parser("reconstruct")
|
reconstruction = commands.add_parser("reconstruct")
|
||||||
reconstruction.add_argument("--input", required=True)
|
reconstruction.add_argument("--input", required=True)
|
||||||
@@ -709,6 +1048,7 @@ def parser():
|
|||||||
required=True,
|
required=True,
|
||||||
)
|
)
|
||||||
reconstruction.add_argument("--video-frames", type=int, required=True)
|
reconstruction.add_argument("--video-frames", type=int, required=True)
|
||||||
|
reconstruction.add_argument("--start-frame", type=int, default=1)
|
||||||
reconstruction.add_argument(
|
reconstruction.add_argument(
|
||||||
"--device",
|
"--device",
|
||||||
choices=("auto", "cuda", "mps", "cpu"),
|
choices=("auto", "cuda", "mps", "cpu"),
|
||||||
@@ -735,6 +1075,8 @@ def main():
|
|||||||
download_model(args)
|
download_model(args)
|
||||||
elif args.command == "reconstruct":
|
elif args.command == "reconstruct":
|
||||||
reconstruct(args)
|
reconstruct(args)
|
||||||
|
elif args.command == "extract-frames":
|
||||||
|
extract_frames(args)
|
||||||
else:
|
else:
|
||||||
remove_background(args)
|
remove_background(args)
|
||||||
except Exception as error:
|
except Exception as error:
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
@@ -1,5 +1,4 @@
|
|||||||
import importlib.util
|
import importlib.util
|
||||||
import inspect
|
|
||||||
import sys
|
import sys
|
||||||
import tempfile
|
import tempfile
|
||||||
import unittest
|
import unittest
|
||||||
@@ -112,21 +111,6 @@ class HostTest(unittest.TestCase):
|
|||||||
self.assertEqual(loads, ["model"])
|
self.assertEqual(loads, ["model"])
|
||||||
self.assertEqual(cache.status()["da3"], "da3_model")
|
self.assertEqual(cache.status()["da3"], "da3_model")
|
||||||
|
|
||||||
def test_http_host_has_health_jobs_auth_and_single_task_guard(self):
|
|
||||||
source = inspect.getsource(self.host.create_app)
|
|
||||||
|
|
||||||
self.assertIn('@app.get("/health")', source)
|
|
||||||
self.assertIn('@app.post("/jobs", status_code=202)', source)
|
|
||||||
self.assertIn("x_snapform_token", source)
|
|
||||||
self.assertIn("hmac.compare_digest", source)
|
|
||||||
self.assertIn("parameters[\"status\"]", source)
|
|
||||||
self.assertIn('@app.post("/clients", status_code=201)', source)
|
|
||||||
self.assertIn('@app.post("/clients/{client_id}/heartbeat")', source)
|
|
||||||
self.assertIn('@app.delete("/jobs/{job_id}", status_code=202)', source)
|
|
||||||
self.assertIn('@app.post("/models/clear")', source)
|
|
||||||
self.assertIn('@app.post("/shutdown", status_code=202)', source)
|
|
||||||
self.assertNotIn("os._exit", source)
|
|
||||||
|
|
||||||
def test_host_tracks_clients_and_rejects_shared_shutdown(self):
|
def test_host_tracks_clients_and_rejects_shared_shutdown(self):
|
||||||
fake_models = SimpleNamespace(
|
fake_models = SimpleNamespace(
|
||||||
status=lambda: {"ben2": "", "da3": ""},
|
status=lambda: {"ben2": "", "da3": ""},
|
||||||
@@ -218,6 +202,51 @@ class HostTest(unittest.TestCase):
|
|||||||
self.host.remove_discovery(path, "owner")
|
self.host.remove_discovery(path, "owner")
|
||||||
self.assertFalse(path.exists())
|
self.assertFalse(path.exists())
|
||||||
|
|
||||||
|
def test_host_dispatches_frame_extraction_without_models(self):
|
||||||
|
fake_models = SimpleNamespace(
|
||||||
|
status=lambda: {"ben2": "", "da3": ""},
|
||||||
|
close=lambda: None,
|
||||||
|
)
|
||||||
|
with patch.object(self.host, "ModelCache", return_value=fake_models):
|
||||||
|
state = self.host.HostState("token", "instance")
|
||||||
|
state.register_client("client", "Blender")
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
root = Path(directory)
|
||||||
|
input_directory = root / "frames"
|
||||||
|
input_directory.mkdir()
|
||||||
|
(input_directory / "frame-0001.jpg").touch()
|
||||||
|
(input_directory / "frame-0002.jpg").touch()
|
||||||
|
result_directory = root / "result"
|
||||||
|
status_path = root / "status.json"
|
||||||
|
state.reserve(
|
||||||
|
"job",
|
||||||
|
"client",
|
||||||
|
"extract-frames",
|
||||||
|
str(status_path),
|
||||||
|
)
|
||||||
|
|
||||||
|
state.run_job(
|
||||||
|
{
|
||||||
|
"job_id": "job",
|
||||||
|
"command": "extract-frames",
|
||||||
|
"parameters": {
|
||||||
|
"input": str(input_directory),
|
||||||
|
"output_dir": str(result_directory),
|
||||||
|
"status": str(status_path),
|
||||||
|
"video_frames": 16,
|
||||||
|
"start_frame": 1,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
)
|
||||||
|
status = self.host.json.loads(
|
||||||
|
status_path.read_text(encoding="utf-8")
|
||||||
|
)
|
||||||
|
self.assertTrue(
|
||||||
|
(result_directory / "frame.0001.jpg").is_file()
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(status["state"], "succeeded")
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|||||||
@@ -1,7 +1,9 @@
|
|||||||
import importlib.util
|
import importlib.util
|
||||||
|
import io
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
import subprocess
|
import subprocess
|
||||||
|
import tarfile
|
||||||
import tempfile
|
import tempfile
|
||||||
import time
|
import time
|
||||||
import unittest
|
import unittest
|
||||||
@@ -25,12 +27,24 @@ def load_bootstrap_runtime():
|
|||||||
|
|
||||||
|
|
||||||
class BootstrapRuntimeTest(unittest.TestCase):
|
class BootstrapRuntimeTest(unittest.TestCase):
|
||||||
|
def create_bundled_sources(self, root):
|
||||||
|
sources = root / "sources"
|
||||||
|
sources.mkdir()
|
||||||
|
(sources / "depth-anything-3.tar.gz").write_bytes(b"da3")
|
||||||
|
(sources / "ben2.tar.gz").write_bytes(b"ben2")
|
||||||
|
return sources
|
||||||
|
|
||||||
def test_installs_da3_from_archives_without_gsplat(self):
|
def test_installs_da3_from_archives_without_gsplat(self):
|
||||||
bootstrap = load_bootstrap_runtime()
|
bootstrap = load_bootstrap_runtime()
|
||||||
with tempfile.TemporaryDirectory() as directory:
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
root = Path(directory)
|
root = Path(directory)
|
||||||
runtime_dir = root / "runtime"
|
runtime_dir = root / "runtime"
|
||||||
runtime_dir.mkdir()
|
runtime_dir.mkdir()
|
||||||
|
self.create_bundled_sources(root)
|
||||||
|
da3_uri = (
|
||||||
|
root / "sources" / "depth-anything-3.tar.gz"
|
||||||
|
).as_uri()
|
||||||
|
ben2_uri = (root / "sources" / "ben2.tar.gz").as_uri()
|
||||||
uv_path = root / "uv"
|
uv_path = root / "uv"
|
||||||
uv_path.write_bytes(b"uv")
|
uv_path.write_bytes(b"uv")
|
||||||
args = SimpleNamespace(
|
args = SimpleNamespace(
|
||||||
@@ -42,7 +56,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
ben2_commit="test-ben2-commit",
|
ben2_commit="test-ben2-commit",
|
||||||
status=str(root / "status.json"),
|
status=str(root / "status.json"),
|
||||||
)
|
)
|
||||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
|
||||||
bootstrap,
|
bootstrap,
|
||||||
"run",
|
"run",
|
||||||
) as mocked_run, patch.object(
|
) as mocked_run, patch.object(
|
||||||
@@ -64,13 +78,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
f"addict=={bootstrap.ADDICT_VERSION}",
|
f"addict=={bootstrap.ADDICT_VERSION}",
|
||||||
da3_command,
|
da3_command,
|
||||||
)
|
)
|
||||||
self.assertTrue(
|
self.assertIn(f"depth-anything-3 @ {da3_uri}", da3_command)
|
||||||
any(
|
|
||||||
"https://github.com/ByteDance-Seed/Depth-Anything-3/"
|
|
||||||
"archive/test-commit.tar.gz" in str(value)
|
|
||||||
for value in da3_command
|
|
||||||
)
|
|
||||||
)
|
|
||||||
self.assertFalse(
|
self.assertFalse(
|
||||||
any("git+" in str(value) for value in da3_command)
|
any("git+" in str(value) for value in da3_command)
|
||||||
)
|
)
|
||||||
@@ -93,13 +101,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
for command in commands
|
for command in commands
|
||||||
if any("ben2 @" in str(value) for value in command)
|
if any("ben2 @" in str(value) for value in command)
|
||||||
)
|
)
|
||||||
self.assertTrue(
|
self.assertIn(f"ben2 @ {ben2_uri}", ben2_command)
|
||||||
any(
|
|
||||||
"https://github.com/PramaLLC/BEN2/archive/"
|
|
||||||
"test-ben2-commit.tar.gz" in str(value)
|
|
||||||
for value in ben2_command
|
|
||||||
)
|
|
||||||
)
|
|
||||||
self.assertFalse(
|
self.assertFalse(
|
||||||
any("git+" in str(value) for value in ben2_command)
|
any("git+" in str(value) for value in ben2_command)
|
||||||
)
|
)
|
||||||
@@ -162,6 +164,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
root = Path(directory)
|
root = Path(directory)
|
||||||
runtime_dir = root / "runtime"
|
runtime_dir = root / "runtime"
|
||||||
runtime_dir.mkdir()
|
runtime_dir.mkdir()
|
||||||
|
self.create_bundled_sources(root)
|
||||||
uv_path = root / "uv"
|
uv_path = root / "uv"
|
||||||
uv_path.write_bytes(b"uv")
|
uv_path.write_bytes(b"uv")
|
||||||
args = SimpleNamespace(
|
args = SimpleNamespace(
|
||||||
@@ -181,7 +184,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
bootstrap.platform,
|
bootstrap.platform,
|
||||||
"machine",
|
"machine",
|
||||||
return_value="arm64",
|
return_value="arm64",
|
||||||
), patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
), patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
|
||||||
bootstrap,
|
bootstrap,
|
||||||
"run",
|
"run",
|
||||||
), patch.object(
|
), patch.object(
|
||||||
@@ -329,6 +332,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
root = Path(directory)
|
root = Path(directory)
|
||||||
runtime_dir = root / "runtime"
|
runtime_dir = root / "runtime"
|
||||||
runtime_dir.mkdir()
|
runtime_dir.mkdir()
|
||||||
|
self.create_bundled_sources(root)
|
||||||
uv_path = root / "uv"
|
uv_path = root / "uv"
|
||||||
uv_path.write_bytes(b"uv")
|
uv_path.write_bytes(b"uv")
|
||||||
args = SimpleNamespace(
|
args = SimpleNamespace(
|
||||||
@@ -364,7 +368,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
)
|
)
|
||||||
return []
|
return []
|
||||||
|
|
||||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
|
||||||
bootstrap,
|
bootstrap,
|
||||||
"run_stage",
|
"run_stage",
|
||||||
side_effect=fake_run_stage,
|
side_effect=fake_run_stage,
|
||||||
@@ -425,7 +429,8 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
"3/3 packages",
|
"3/3 packages",
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
self.assertEqual(mocked_write.call_args_list[-1].args[1], 0.02)
|
self.assertEqual(mocked_write.call_args_list[0].args[1], 0.02)
|
||||||
|
self.assertLess(mocked_write.call_args_list[-1].args[1], 0.03)
|
||||||
first = mocked_write.call_args_list[0]
|
first = mocked_write.call_args_list[0]
|
||||||
self.assertEqual(first.kwargs["stage"], 2)
|
self.assertEqual(first.kwargs["stage"], 2)
|
||||||
self.assertEqual(first.kwargs["stages"], 8)
|
self.assertEqual(first.kwargs["stages"], 8)
|
||||||
@@ -525,146 +530,79 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
self.assertEqual(first.kwargs["stages"], 8)
|
self.assertEqual(first.kwargs["stages"], 8)
|
||||||
self.assertEqual(first.kwargs["stage_label"], "pytorch deps")
|
self.assertEqual(first.kwargs["stage_label"], "pytorch deps")
|
||||||
|
|
||||||
|
def test_run_install_retries_once_with_mirror(self):
|
||||||
|
bootstrap = load_bootstrap_runtime()
|
||||||
|
previous = bootstrap._resolved_pypi_index
|
||||||
|
bootstrap._resolved_pypi_index = None
|
||||||
|
try:
|
||||||
|
with patch.object(
|
||||||
|
bootstrap,
|
||||||
|
"_run_install_once",
|
||||||
|
side_effect=[
|
||||||
|
subprocess.CalledProcessError(1, ["uv", "pip"]),
|
||||||
|
["numpy==1.26.4"],
|
||||||
|
],
|
||||||
|
) as mocked_install:
|
||||||
|
installed = bootstrap.run_install(
|
||||||
|
"status.json",
|
||||||
|
3,
|
||||||
|
8,
|
||||||
|
"ai deps",
|
||||||
|
["uv", "pip", "install", "x"],
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
bootstrap._resolved_pypi_index = previous
|
||||||
|
|
||||||
|
self.assertEqual(installed, ["numpy==1.26.4"])
|
||||||
|
self.assertEqual(mocked_install.call_count, 2)
|
||||||
|
self.assertEqual(
|
||||||
|
mocked_install.call_args_list[1].args[5],
|
||||||
|
{"PIP_INDEX_URL": bootstrap.PYPI_MIRRORS[0]},
|
||||||
|
)
|
||||||
|
|
||||||
def test_formats_elapsed_time(self):
|
def test_formats_elapsed_time(self):
|
||||||
bootstrap = load_bootstrap_runtime()
|
bootstrap = load_bootstrap_runtime()
|
||||||
|
|
||||||
self.assertEqual(bootstrap.format_elapsed(12), "12s")
|
self.assertEqual(bootstrap.format_elapsed(12), "12s")
|
||||||
self.assertEqual(bootstrap.format_elapsed(83), "1m 23s")
|
self.assertEqual(bootstrap.format_elapsed(83), "1m 23s")
|
||||||
|
|
||||||
def test_preflight_rejects_low_disk_space(self):
|
def test_resolve_indexes_uses_configured_pypi_index(self):
|
||||||
bootstrap = load_bootstrap_runtime()
|
bootstrap = load_bootstrap_runtime()
|
||||||
args = SimpleNamespace(
|
bootstrap.resolve_indexes(
|
||||||
backend="cu128",
|
SimpleNamespace(pypi_index="https://pypi.tuna.tsinghua.edu.cn/simple")
|
||||||
runtime_dir=str(Path.cwd()),
|
|
||||||
)
|
)
|
||||||
with patch.object(
|
|
||||||
bootstrap.shutil,
|
|
||||||
"disk_usage",
|
|
||||||
return_value=SimpleNamespace(free=1024**3),
|
|
||||||
), patch.object(bootstrap, "reachable", return_value=True):
|
|
||||||
with self.assertRaisesRegex(RuntimeError, "disk space"):
|
|
||||||
bootstrap.preflight(args)
|
|
||||||
|
|
||||||
def test_preflight_rejects_unreachable_network(self):
|
|
||||||
bootstrap = load_bootstrap_runtime()
|
|
||||||
args = SimpleNamespace(
|
|
||||||
backend="cu128",
|
|
||||||
runtime_dir=str(Path.cwd()),
|
|
||||||
)
|
|
||||||
with patch.object(
|
|
||||||
bootstrap.shutil,
|
|
||||||
"disk_usage",
|
|
||||||
return_value=SimpleNamespace(free=100 * 1024**3),
|
|
||||||
), patch.object(bootstrap, "reachable", return_value=False):
|
|
||||||
with self.assertRaisesRegex(RuntimeError, "Cannot reach"):
|
|
||||||
bootstrap.preflight(args)
|
|
||||||
|
|
||||||
def test_preflight_passes_when_checks_ok(self):
|
|
||||||
bootstrap = load_bootstrap_runtime()
|
|
||||||
args = SimpleNamespace(
|
|
||||||
backend="cu128",
|
|
||||||
runtime_dir=str(Path.cwd()),
|
|
||||||
)
|
|
||||||
with patch.object(
|
|
||||||
bootstrap.shutil,
|
|
||||||
"disk_usage",
|
|
||||||
return_value=SimpleNamespace(free=100 * 1024**3),
|
|
||||||
), patch.object(bootstrap, "reachable", return_value=True):
|
|
||||||
resolved, used_mirror = bootstrap.preflight(args)
|
|
||||||
|
|
||||||
self.assertFalse(used_mirror)
|
|
||||||
self.assertEqual(
|
self.assertEqual(
|
||||||
resolved["pypi"],
|
bootstrap._resolved_pypi_index,
|
||||||
"https://pypi.org/simple/",
|
|
||||||
)
|
|
||||||
|
|
||||||
def test_preflight_falls_back_to_mirrors(self):
|
|
||||||
bootstrap = load_bootstrap_runtime()
|
|
||||||
args = SimpleNamespace(
|
|
||||||
backend="cu128",
|
|
||||||
runtime_dir=str(Path.cwd()),
|
|
||||||
)
|
|
||||||
|
|
||||||
def fake_reachable(url, timeout=bootstrap.PREFLIGHT_TIMEOUT):
|
|
||||||
if "github.com" in url:
|
|
||||||
return True
|
|
||||||
if "tuna" in url or "hf-mirror" in url or "aliyun" in url:
|
|
||||||
return True
|
|
||||||
return False
|
|
||||||
|
|
||||||
with patch.object(
|
|
||||||
bootstrap.shutil,
|
|
||||||
"disk_usage",
|
|
||||||
return_value=SimpleNamespace(free=100 * 1024**3),
|
|
||||||
), patch.object(
|
|
||||||
bootstrap,
|
|
||||||
"reachable",
|
|
||||||
side_effect=fake_reachable,
|
|
||||||
):
|
|
||||||
resolved, used_mirror = bootstrap.preflight(args)
|
|
||||||
|
|
||||||
self.assertTrue(used_mirror)
|
|
||||||
self.assertEqual(
|
|
||||||
resolved["pypi"],
|
|
||||||
"https://pypi.tuna.tsinghua.edu.cn/simple",
|
"https://pypi.tuna.tsinghua.edu.cn/simple",
|
||||||
)
|
)
|
||||||
self.assertEqual(resolved["hf"], "https://hf-mirror.com")
|
|
||||||
self.assertEqual(
|
|
||||||
resolved["torch"],
|
|
||||||
"https://mirrors.aliyun.com/pytorch-wheels/cu128",
|
|
||||||
)
|
|
||||||
self.assertEqual(
|
|
||||||
bootstrap._resolved_torch_index,
|
|
||||||
"https://mirrors.aliyun.com/pytorch-wheels/cu128",
|
|
||||||
)
|
|
||||||
|
|
||||||
def test_preflight_uses_environment_overrides(self):
|
def test_resolve_indexes_uses_environment_override(self):
|
||||||
bootstrap = load_bootstrap_runtime()
|
bootstrap = load_bootstrap_runtime()
|
||||||
args = SimpleNamespace(
|
|
||||||
backend="cu128",
|
|
||||||
runtime_dir=str(Path.cwd()),
|
|
||||||
)
|
|
||||||
with patch.dict(
|
with patch.dict(
|
||||||
os.environ,
|
os.environ,
|
||||||
{
|
{"PIP_INDEX_URL": "https://pypi.example/simple"},
|
||||||
"PIP_INDEX_URL": "https://pypi.example/simple",
|
):
|
||||||
"HF_ENDPOINT": "https://hf.example",
|
bootstrap.resolve_indexes(SimpleNamespace(pypi_index=None))
|
||||||
"SNAPFORM_TORCH_INDEX": "https://torch.example/whl/cu128",
|
|
||||||
},
|
|
||||||
), patch.object(
|
|
||||||
bootstrap.shutil,
|
|
||||||
"disk_usage",
|
|
||||||
return_value=SimpleNamespace(free=100 * 1024**3),
|
|
||||||
), patch.object(bootstrap, "reachable", return_value=True):
|
|
||||||
resolved, used_mirror = bootstrap.preflight(args)
|
|
||||||
|
|
||||||
self.assertFalse(used_mirror)
|
|
||||||
self.assertEqual(resolved["pypi"], "https://pypi.example/simple")
|
|
||||||
self.assertEqual(resolved["hf"], "https://hf.example")
|
|
||||||
self.assertEqual(
|
self.assertEqual(
|
||||||
resolved["torch"],
|
bootstrap._resolved_pypi_index,
|
||||||
"https://torch.example/whl/cu128",
|
"https://pypi.example/simple",
|
||||||
)
|
)
|
||||||
|
|
||||||
def test_child_environment_injects_resolved_mirrors(self):
|
def test_child_environment_injects_resolved_pypi_index(self):
|
||||||
bootstrap = load_bootstrap_runtime()
|
bootstrap = load_bootstrap_runtime()
|
||||||
previous_pypi = bootstrap._resolved_pypi_index
|
previous_pypi = bootstrap._resolved_pypi_index
|
||||||
previous_hf = bootstrap._resolved_hf_endpoint
|
|
||||||
bootstrap._resolved_pypi_index = (
|
bootstrap._resolved_pypi_index = (
|
||||||
"https://pypi.tuna.tsinghua.edu.cn/simple"
|
"https://pypi.tuna.tsinghua.edu.cn/simple"
|
||||||
)
|
)
|
||||||
bootstrap._resolved_hf_endpoint = "https://hf-mirror.com"
|
|
||||||
try:
|
try:
|
||||||
env = bootstrap.child_environment()
|
env = bootstrap.child_environment()
|
||||||
finally:
|
finally:
|
||||||
bootstrap._resolved_pypi_index = previous_pypi
|
bootstrap._resolved_pypi_index = previous_pypi
|
||||||
bootstrap._resolved_hf_endpoint = previous_hf
|
|
||||||
|
|
||||||
self.assertEqual(
|
self.assertEqual(
|
||||||
env["PIP_INDEX_URL"],
|
env["PIP_INDEX_URL"],
|
||||||
"https://pypi.tuna.tsinghua.edu.cn/simple",
|
"https://pypi.tuna.tsinghua.edu.cn/simple",
|
||||||
)
|
)
|
||||||
self.assertEqual(env["HF_ENDPOINT"], "https://hf-mirror.com")
|
|
||||||
|
|
||||||
def test_child_environment_injects_registry_proxy(self):
|
def test_child_environment_injects_registry_proxy(self):
|
||||||
bootstrap = load_bootstrap_runtime()
|
bootstrap = load_bootstrap_runtime()
|
||||||
@@ -700,6 +638,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
root = Path(directory)
|
root = Path(directory)
|
||||||
runtime_dir = root / "runtime"
|
runtime_dir = root / "runtime"
|
||||||
runtime_dir.mkdir()
|
runtime_dir.mkdir()
|
||||||
|
self.create_bundled_sources(root)
|
||||||
uv_path = root / "uv"
|
uv_path = root / "uv"
|
||||||
uv_path.write_bytes(b"uv")
|
uv_path.write_bytes(b"uv")
|
||||||
args = SimpleNamespace(
|
args = SimpleNamespace(
|
||||||
@@ -711,7 +650,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
ben2_commit="test-ben2-commit",
|
ben2_commit="test-ben2-commit",
|
||||||
status=str(root / "status.json"),
|
status=str(root / "status.json"),
|
||||||
)
|
)
|
||||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
|
||||||
bootstrap,
|
bootstrap,
|
||||||
"run",
|
"run",
|
||||||
), patch.object(
|
), patch.object(
|
||||||
@@ -737,6 +676,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
root = Path(directory)
|
root = Path(directory)
|
||||||
runtime_dir = root / "runtime"
|
runtime_dir = root / "runtime"
|
||||||
runtime_dir.mkdir()
|
runtime_dir.mkdir()
|
||||||
|
self.create_bundled_sources(root)
|
||||||
(runtime_dir / "partial-python.exe").write_text("junk")
|
(runtime_dir / "partial-python.exe").write_text("junk")
|
||||||
args = SimpleNamespace(
|
args = SimpleNamespace(
|
||||||
uv="uv",
|
uv="uv",
|
||||||
@@ -747,7 +687,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
ben2_commit="test-ben2-commit",
|
ben2_commit="test-ben2-commit",
|
||||||
status=str(root / "status.json"),
|
status=str(root / "status.json"),
|
||||||
)
|
)
|
||||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
|
||||||
bootstrap,
|
bootstrap,
|
||||||
"run",
|
"run",
|
||||||
), patch.object(
|
), patch.object(
|
||||||
@@ -765,6 +705,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
root = Path(directory)
|
root = Path(directory)
|
||||||
runtime_dir = root / "runtime"
|
runtime_dir = root / "runtime"
|
||||||
runtime_dir.mkdir(parents=True)
|
runtime_dir.mkdir(parents=True)
|
||||||
|
self.create_bundled_sources(root)
|
||||||
(runtime_dir / "Scripts").mkdir()
|
(runtime_dir / "Scripts").mkdir()
|
||||||
(runtime_dir / "Scripts" / "python.exe").write_text("python")
|
(runtime_dir / "Scripts" / "python.exe").write_text("python")
|
||||||
(runtime_dir / "snapform_ready.json").write_text(
|
(runtime_dir / "snapform_ready.json").write_text(
|
||||||
@@ -781,7 +722,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
ben2_commit="test-ben2-commit",
|
ben2_commit="test-ben2-commit",
|
||||||
status=str(root / "status.json"),
|
status=str(root / "status.json"),
|
||||||
)
|
)
|
||||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
|
||||||
bootstrap,
|
bootstrap,
|
||||||
"run",
|
"run",
|
||||||
), patch.object(
|
), patch.object(
|
||||||
@@ -803,6 +744,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
root = Path(directory)
|
root = Path(directory)
|
||||||
runtime_dir = root / "runtime"
|
runtime_dir = root / "runtime"
|
||||||
runtime_dir.mkdir()
|
runtime_dir.mkdir()
|
||||||
|
self.create_bundled_sources(root)
|
||||||
args = SimpleNamespace(
|
args = SimpleNamespace(
|
||||||
uv="uv",
|
uv="uv",
|
||||||
runtime_dir=str(runtime_dir),
|
runtime_dir=str(runtime_dir),
|
||||||
@@ -812,7 +754,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
ben2_commit="test-ben2-commit",
|
ben2_commit="test-ben2-commit",
|
||||||
status=str(root / "status.json"),
|
status=str(root / "status.json"),
|
||||||
)
|
)
|
||||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
|
||||||
bootstrap,
|
bootstrap,
|
||||||
"run",
|
"run",
|
||||||
side_effect=fake_run,
|
side_effect=fake_run,
|
||||||
@@ -826,12 +768,13 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
self.assertEqual(status["progress"], 1.0)
|
self.assertEqual(status["progress"], 1.0)
|
||||||
self.assertIn("Runtime is ready", status["message"])
|
self.assertIn("Runtime is ready", status["message"])
|
||||||
|
|
||||||
def test_install_downloads_missing_uv(self):
|
def test_install_fails_when_bundled_uv_missing(self):
|
||||||
bootstrap = load_bootstrap_runtime()
|
bootstrap = load_bootstrap_runtime()
|
||||||
with tempfile.TemporaryDirectory() as directory:
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
root = Path(directory)
|
root = Path(directory)
|
||||||
runtime_dir = root / "runtime"
|
runtime_dir = root / "runtime"
|
||||||
runtime_dir.mkdir()
|
runtime_dir.mkdir()
|
||||||
|
self.create_bundled_sources(root)
|
||||||
uv_path = root / "uv" / "uv.exe"
|
uv_path = root / "uv" / "uv.exe"
|
||||||
args = SimpleNamespace(
|
args = SimpleNamespace(
|
||||||
uv=str(uv_path),
|
uv=str(uv_path),
|
||||||
@@ -842,7 +785,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
ben2_commit="test-ben2-commit",
|
ben2_commit="test-ben2-commit",
|
||||||
status=str(root / "status.json"),
|
status=str(root / "status.json"),
|
||||||
)
|
)
|
||||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
|
||||||
bootstrap,
|
bootstrap,
|
||||||
"run",
|
"run",
|
||||||
), patch.object(
|
), patch.object(
|
||||||
@@ -855,15 +798,101 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
), patch.object(
|
), patch.object(
|
||||||
bootstrap.shutil,
|
bootstrap.shutil,
|
||||||
"rmtree",
|
"rmtree",
|
||||||
), patch.object(bootstrap, "download_uv") as mocked_download:
|
):
|
||||||
|
with self.assertRaisesRegex(
|
||||||
|
RuntimeError,
|
||||||
|
"Bundled uv is missing",
|
||||||
|
):
|
||||||
|
bootstrap.install(args)
|
||||||
|
|
||||||
|
def test_install_uses_bundled_sources_and_uv_without_network(self):
|
||||||
|
bootstrap = load_bootstrap_runtime()
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
root = Path(directory)
|
||||||
|
runtime_dir = root / "runtime"
|
||||||
|
runtime_dir.mkdir()
|
||||||
|
bundled = root / "bundled"
|
||||||
|
bundled.mkdir()
|
||||||
|
(bundled / "depth-anything-3.tar.gz").write_bytes(b"da3")
|
||||||
|
(bundled / "ben2.tar.gz").write_bytes(b"ben2")
|
||||||
|
target, archive_name = bootstrap.uv_download_target()
|
||||||
|
archive = bundled / f"uv-{target}.{archive_name}"
|
||||||
|
if archive_name.endswith("zip"):
|
||||||
|
with zipfile.ZipFile(archive, "w") as bundle:
|
||||||
|
bundle.writestr("uv/uv.exe", b"uv-binary")
|
||||||
|
else:
|
||||||
|
with tarfile.open(archive, "w:gz") as tar:
|
||||||
|
content = b"uv-binary"
|
||||||
|
info = tarfile.TarInfo("uv")
|
||||||
|
info.size = len(content)
|
||||||
|
tar.addfile(info, io.BytesIO(content))
|
||||||
|
uv_path = root / "uv" / ("uv.exe" if os.name == "nt" else "uv")
|
||||||
|
args = SimpleNamespace(
|
||||||
|
uv=str(uv_path),
|
||||||
|
runtime_dir=str(runtime_dir),
|
||||||
|
models_dir=str(root / "models"),
|
||||||
|
backend="cu128",
|
||||||
|
da3_commit="test-commit",
|
||||||
|
ben2_commit="test-ben2-commit",
|
||||||
|
status=str(root / "status.json"),
|
||||||
|
)
|
||||||
|
run_commands = []
|
||||||
|
install_commands = []
|
||||||
|
|
||||||
|
def fake_run(command, environment=None):
|
||||||
|
run_commands.append(command)
|
||||||
|
|
||||||
|
def fake_run_install(
|
||||||
|
status_path,
|
||||||
|
stage,
|
||||||
|
stages,
|
||||||
|
label,
|
||||||
|
command,
|
||||||
|
environment=None,
|
||||||
|
start=0.0,
|
||||||
|
end=1.0,
|
||||||
|
clock=None,
|
||||||
|
):
|
||||||
|
install_commands.append(command)
|
||||||
|
return []
|
||||||
|
|
||||||
|
with patch.object(
|
||||||
|
bootstrap,
|
||||||
|
"bundled_sources_directory",
|
||||||
|
return_value=bundled,
|
||||||
|
), patch.object(
|
||||||
|
bootstrap,
|
||||||
|
"resolve_indexes",
|
||||||
|
return_value=None,
|
||||||
|
), patch.object(
|
||||||
|
bootstrap,
|
||||||
|
"run",
|
||||||
|
side_effect=fake_run,
|
||||||
|
), patch.object(
|
||||||
|
bootstrap,
|
||||||
|
"run_install",
|
||||||
|
side_effect=fake_run_install,
|
||||||
|
):
|
||||||
bootstrap.install(args)
|
bootstrap.install(args)
|
||||||
|
|
||||||
mocked_download.assert_called_once()
|
self.assertEqual(uv_path.read_bytes(), b"uv-binary")
|
||||||
call_args = mocked_download.call_args
|
all_commands = run_commands + install_commands
|
||||||
self.assertEqual(call_args.args[0], uv_path)
|
self.assertFalse(
|
||||||
self.assertEqual(call_args.args[1], str(root / "status.json"))
|
any(
|
||||||
self.assertEqual(call_args.kwargs["stage"], 1)
|
"github.com" in str(value)
|
||||||
self.assertEqual(call_args.kwargs["stage_label"], "uv")
|
for command in all_commands
|
||||||
|
for value in command
|
||||||
|
)
|
||||||
|
)
|
||||||
|
ai_command = next(
|
||||||
|
command
|
||||||
|
for command in install_commands
|
||||||
|
if any("depth-anything-3 @" in str(value) for value in command)
|
||||||
|
)
|
||||||
|
da3_uri = (root / "sources" / "depth-anything-3.tar.gz").as_uri()
|
||||||
|
ben2_uri = (root / "sources" / "ben2.tar.gz").as_uri()
|
||||||
|
self.assertIn(f"depth-anything-3 @ {da3_uri}", ai_command)
|
||||||
|
self.assertIn(f"ben2 @ {ben2_uri}", ai_command)
|
||||||
|
|
||||||
def test_install_reuses_existing_uv(self):
|
def test_install_reuses_existing_uv(self):
|
||||||
bootstrap = load_bootstrap_runtime()
|
bootstrap = load_bootstrap_runtime()
|
||||||
@@ -871,6 +900,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
root = Path(directory)
|
root = Path(directory)
|
||||||
runtime_dir = root / "runtime"
|
runtime_dir = root / "runtime"
|
||||||
runtime_dir.mkdir()
|
runtime_dir.mkdir()
|
||||||
|
self.create_bundled_sources(root)
|
||||||
uv_path = root / "uv.exe"
|
uv_path = root / "uv.exe"
|
||||||
uv_path.write_bytes(b"uv")
|
uv_path.write_bytes(b"uv")
|
||||||
args = SimpleNamespace(
|
args = SimpleNamespace(
|
||||||
@@ -882,17 +912,15 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
ben2_commit="test-ben2-commit",
|
ben2_commit="test-ben2-commit",
|
||||||
status=str(root / "status.json"),
|
status=str(root / "status.json"),
|
||||||
)
|
)
|
||||||
with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object(
|
with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object(
|
||||||
bootstrap,
|
bootstrap,
|
||||||
"run",
|
"run",
|
||||||
), patch.object(
|
), patch.object(
|
||||||
bootstrap,
|
bootstrap,
|
||||||
"run_install",
|
"run_install",
|
||||||
), patch.object(bootstrap, "download_uv") as mocked_download:
|
):
|
||||||
bootstrap.install(args)
|
bootstrap.install(args)
|
||||||
|
|
||||||
mocked_download.assert_not_called()
|
|
||||||
|
|
||||||
def test_extract_uv_unpacks_standalone_binary(self):
|
def test_extract_uv_unpacks_standalone_binary(self):
|
||||||
bootstrap = load_bootstrap_runtime()
|
bootstrap = load_bootstrap_runtime()
|
||||||
with tempfile.TemporaryDirectory() as directory:
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
@@ -906,6 +934,112 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
|
|
||||||
self.assertEqual(uv_path.read_bytes(), b"binary")
|
self.assertEqual(uv_path.read_bytes(), b"binary")
|
||||||
|
|
||||||
|
def test_sync_bundled_sources_copies_missing_archives(self):
|
||||||
|
bootstrap = load_bootstrap_runtime()
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
root = Path(directory)
|
||||||
|
bundled = root / "bundled"
|
||||||
|
bundled.mkdir()
|
||||||
|
(bundled / "depth-anything-3.tar.gz").write_bytes(b"da3")
|
||||||
|
(bundled / "ben2.tar.gz").write_bytes(b"ben2")
|
||||||
|
sources = root / "sources"
|
||||||
|
|
||||||
|
bootstrap.sync_bundled_sources(bundled, sources)
|
||||||
|
|
||||||
|
self.assertEqual(
|
||||||
|
(sources / "depth-anything-3.tar.gz").read_bytes(),
|
||||||
|
b"da3",
|
||||||
|
)
|
||||||
|
self.assertEqual(
|
||||||
|
(sources / "ben2.tar.gz").read_bytes(),
|
||||||
|
b"ben2",
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_install_stages_uses_local_source_archives(self):
|
||||||
|
bootstrap = load_bootstrap_runtime()
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
root = Path(directory)
|
||||||
|
runtime_dir = root / "runtime"
|
||||||
|
runtime_dir.mkdir()
|
||||||
|
sources = root / "sources"
|
||||||
|
sources.mkdir()
|
||||||
|
(sources / "depth-anything-3.tar.gz").write_bytes(b"da3")
|
||||||
|
(sources / "ben2.tar.gz").write_bytes(b"ben2")
|
||||||
|
uv_path = root / "uv"
|
||||||
|
uv_path.write_bytes(b"uv")
|
||||||
|
args = SimpleNamespace(
|
||||||
|
uv=str(uv_path),
|
||||||
|
runtime_dir=str(runtime_dir),
|
||||||
|
models_dir=str(root / "models"),
|
||||||
|
backend="cu128",
|
||||||
|
da3_commit="test-commit",
|
||||||
|
ben2_commit="test-ben2-commit",
|
||||||
|
status=str(root / "status.json"),
|
||||||
|
)
|
||||||
|
|
||||||
|
stages = bootstrap.install_stages(
|
||||||
|
args,
|
||||||
|
runtime_dir,
|
||||||
|
root / "python",
|
||||||
|
uv_path,
|
||||||
|
)
|
||||||
|
|
||||||
|
ai_command = next(
|
||||||
|
command
|
||||||
|
for _label, command, _environment, kind in stages
|
||||||
|
if kind == "install"
|
||||||
|
and any("depth-anything-3" in str(value) for value in command)
|
||||||
|
)
|
||||||
|
da3_local = (sources / "depth-anything-3.tar.gz").as_uri()
|
||||||
|
ben2_local = (sources / "ben2.tar.gz").as_uri()
|
||||||
|
self.assertIn(f"depth-anything-3 @ {da3_local}", ai_command)
|
||||||
|
self.assertIn(f"ben2 @ {ben2_local}", ai_command)
|
||||||
|
|
||||||
|
def test_install_uv_extracts_bundled_archive(self):
|
||||||
|
bootstrap = load_bootstrap_runtime()
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
root = Path(directory)
|
||||||
|
uv_path = root / "uv"
|
||||||
|
bundled = root / "bundled"
|
||||||
|
bundled.mkdir(parents=True)
|
||||||
|
target, archive_name = bootstrap.uv_download_target()
|
||||||
|
archive = bundled / f"uv-{target}.{archive_name}"
|
||||||
|
if archive_name.endswith("zip"):
|
||||||
|
with zipfile.ZipFile(archive, "w") as bundle:
|
||||||
|
bundle.writestr("uv/uv.exe", b"uv-binary")
|
||||||
|
else:
|
||||||
|
with tarfile.open(archive, "w:gz") as tar:
|
||||||
|
content = b"uv-binary"
|
||||||
|
info = tarfile.TarInfo("uv")
|
||||||
|
info.size = len(content)
|
||||||
|
tar.addfile(info, io.BytesIO(content))
|
||||||
|
with patch.object(
|
||||||
|
bootstrap,
|
||||||
|
"bundled_sources_directory",
|
||||||
|
return_value=bundled,
|
||||||
|
):
|
||||||
|
bootstrap.install_uv(uv_path, str(root / "status.json"))
|
||||||
|
|
||||||
|
self.assertEqual(uv_path.read_bytes(), b"uv-binary")
|
||||||
|
|
||||||
|
def test_install_uv_raises_when_bundled_missing(self):
|
||||||
|
bootstrap = load_bootstrap_runtime()
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
root = Path(directory)
|
||||||
|
uv_path = root / "uv"
|
||||||
|
bundled = root / "bundled"
|
||||||
|
bundled.mkdir()
|
||||||
|
with patch.object(
|
||||||
|
bootstrap,
|
||||||
|
"bundled_sources_directory",
|
||||||
|
return_value=bundled,
|
||||||
|
):
|
||||||
|
with self.assertRaisesRegex(
|
||||||
|
RuntimeError,
|
||||||
|
"Bundled uv is missing",
|
||||||
|
):
|
||||||
|
bootstrap.install_uv(uv_path, str(root / "status.json"))
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|||||||
@@ -143,49 +143,29 @@ class DownloadProgressTest(unittest.TestCase):
|
|||||||
|
|
||||||
self.assertFalse(status_path.exists())
|
self.assertFalse(status_path.exists())
|
||||||
|
|
||||||
def test_resolve_hf_endpoint_keeps_direct_endpoint(self):
|
def test_download_endpoints_keeps_direct_endpoint(self):
|
||||||
with patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
|
with patch.dict(os.environ, {"HF_ENDPOINT": ""}):
|
||||||
self.worker,
|
endpoints = self.worker.download_endpoints()
|
||||||
"reachable",
|
|
||||||
return_value=True,
|
|
||||||
):
|
|
||||||
endpoint, used_mirror = self.worker.resolve_hf_endpoint()
|
|
||||||
|
|
||||||
self.assertEqual(endpoint, "https://huggingface.co/")
|
self.assertEqual(
|
||||||
self.assertFalse(used_mirror)
|
endpoints,
|
||||||
|
["https://huggingface.co"],
|
||||||
|
)
|
||||||
|
|
||||||
def test_resolve_hf_endpoint_falls_back_to_mirror(self):
|
def test_download_endpoints_respects_environment_override(self):
|
||||||
with patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
|
with patch.dict(os.environ, {"HF_ENDPOINT": "https://hf.example"}):
|
||||||
self.worker,
|
endpoints = self.worker.download_endpoints()
|
||||||
"reachable",
|
|
||||||
side_effect=lambda url, timeout=None: "hf-mirror" in url,
|
|
||||||
):
|
|
||||||
endpoint, used_mirror = self.worker.resolve_hf_endpoint()
|
|
||||||
|
|
||||||
self.assertEqual(endpoint, "https://hf-mirror.com")
|
self.assertEqual(
|
||||||
self.assertTrue(used_mirror)
|
endpoints,
|
||||||
|
["https://hf.example"],
|
||||||
|
)
|
||||||
|
|
||||||
def test_resolve_hf_endpoint_respects_environment_override(self):
|
def test_download_endpoints_strips_trailing_slash(self):
|
||||||
with patch.dict(os.environ, {"HF_ENDPOINT": "https://hf.example"}), patch.object(
|
with patch.dict(os.environ, {"HF_ENDPOINT": "https://hf.example/"}):
|
||||||
self.worker,
|
endpoints = self.worker.download_endpoints()
|
||||||
"reachable",
|
|
||||||
) as mocked_reachable:
|
|
||||||
endpoint, used_mirror = self.worker.resolve_hf_endpoint()
|
|
||||||
|
|
||||||
self.assertEqual(endpoint, "https://hf.example")
|
self.assertEqual(endpoints, ["https://hf.example"])
|
||||||
self.assertFalse(used_mirror)
|
|
||||||
mocked_reachable.assert_not_called()
|
|
||||||
|
|
||||||
def test_resolve_hf_endpoint_keeps_direct_when_mirrors_unreachable(self):
|
|
||||||
with patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
|
|
||||||
self.worker,
|
|
||||||
"reachable",
|
|
||||||
return_value=False,
|
|
||||||
):
|
|
||||||
endpoint, used_mirror = self.worker.resolve_hf_endpoint()
|
|
||||||
|
|
||||||
self.assertEqual(endpoint, "https://huggingface.co/")
|
|
||||||
self.assertFalse(used_mirror)
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
|||||||
@@ -1,5 +1,4 @@
|
|||||||
import importlib.util
|
import importlib.util
|
||||||
import inspect
|
|
||||||
import sys
|
import sys
|
||||||
import tempfile
|
import tempfile
|
||||||
import unittest
|
import unittest
|
||||||
@@ -60,6 +59,22 @@ class ImageDepthTest(unittest.TestCase):
|
|||||||
np.array(((1.0, 2.0), (3.0, 4.0)), dtype=np.float32),
|
np.array(((1.0, 2.0), (3.0, 4.0)), dtype=np.float32),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def test_extract_depth_arrays_returns_every_frame(self):
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
npz_path = Path(directory) / "results.npz"
|
||||||
|
np.savez_compressed(
|
||||||
|
npz_path,
|
||||||
|
depth=np.arange(12, dtype=np.float32).reshape(3, 2, 2),
|
||||||
|
)
|
||||||
|
|
||||||
|
depth_frames = self.image_depth.extract_depth_arrays(npz_path)
|
||||||
|
|
||||||
|
self.assertEqual(len(depth_frames), 3)
|
||||||
|
np.testing.assert_array_equal(
|
||||||
|
depth_frames[1],
|
||||||
|
np.array(((4.0, 5.0), (6.0, 7.0)), dtype=np.float32),
|
||||||
|
)
|
||||||
|
|
||||||
def test_accepts_depths_key_with_single_channel(self):
|
def test_accepts_depths_key_with_single_channel(self):
|
||||||
with tempfile.TemporaryDirectory() as directory:
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
npz_path = Path(directory) / "results.npz"
|
npz_path = Path(directory) / "results.npz"
|
||||||
@@ -72,6 +87,261 @@ class ImageDepthTest(unittest.TestCase):
|
|||||||
|
|
||||||
self.assertEqual(depth.shape, (2, 3))
|
self.assertEqual(depth.shape, (2, 3))
|
||||||
|
|
||||||
|
def test_is_animated_image_accepts_movie_and_sequence(self):
|
||||||
|
self.assertTrue(
|
||||||
|
self.image_depth.is_animated_image(
|
||||||
|
SimpleNamespace(source="MOVIE")
|
||||||
|
)
|
||||||
|
)
|
||||||
|
self.assertTrue(
|
||||||
|
self.image_depth.is_animated_image(
|
||||||
|
SimpleNamespace(source="SEQUENCE")
|
||||||
|
)
|
||||||
|
)
|
||||||
|
self.assertFalse(
|
||||||
|
self.image_depth.is_animated_image(
|
||||||
|
SimpleNamespace(source="FILE")
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_video_total_frames_reads_movie_duration(self):
|
||||||
|
image = SimpleNamespace(source="MOVIE", frame_duration=97)
|
||||||
|
|
||||||
|
self.assertEqual(self.image_depth.video_total_frames(image), 97)
|
||||||
|
|
||||||
|
def test_video_total_frames_counts_sequence_files(self):
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
root = Path(directory)
|
||||||
|
(root / "foreground-0001.png").touch()
|
||||||
|
(root / "foreground-0002.png").touch()
|
||||||
|
image = SimpleNamespace(
|
||||||
|
source="SEQUENCE",
|
||||||
|
filepath=str(root / "foreground-0001.png"),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(self.image_depth.video_total_frames(image), 2)
|
||||||
|
|
||||||
|
def test_video_total_frames_returns_one_for_still_image(self):
|
||||||
|
image = SimpleNamespace(source="FILE", frame_duration=1)
|
||||||
|
|
||||||
|
self.assertEqual(self.image_depth.video_total_frames(image), 1)
|
||||||
|
|
||||||
|
def test_default_extract_frames_uses_empty_image_frames(self):
|
||||||
|
source_object = SimpleNamespace(
|
||||||
|
data=SimpleNamespace(source="MOVIE", frame_duration=97),
|
||||||
|
image_user=SimpleNamespace(frame_duration=24),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(self.image_depth.default_extract_frames(source_object), 24)
|
||||||
|
|
||||||
|
def test_default_extract_frames_caps_at_video_total(self):
|
||||||
|
source_object = SimpleNamespace(
|
||||||
|
data=SimpleNamespace(source="MOVIE", frame_duration=97),
|
||||||
|
image_user=SimpleNamespace(frame_duration=200),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(self.image_depth.default_extract_frames(source_object), 97)
|
||||||
|
|
||||||
|
def test_default_extract_frames_falls_back_to_video_total(self):
|
||||||
|
source_object = SimpleNamespace(
|
||||||
|
data=SimpleNamespace(source="MOVIE", frame_duration=97),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(self.image_depth.default_extract_frames(source_object), 97)
|
||||||
|
|
||||||
|
def test_extract_frame_range_honors_offset_within_clip(self):
|
||||||
|
source_object = SimpleNamespace(
|
||||||
|
data=SimpleNamespace(source="MOVIE", frame_duration=97),
|
||||||
|
image_user=SimpleNamespace(
|
||||||
|
frame_start=4,
|
||||||
|
frame_offset=70,
|
||||||
|
frame_duration=24,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(
|
||||||
|
self.image_depth.extract_frame_range(source_object),
|
||||||
|
(71, 24),
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_extract_frame_range_clamps_large_offset_to_clip_tail(self):
|
||||||
|
source_object = SimpleNamespace(
|
||||||
|
data=SimpleNamespace(source="MOVIE", frame_duration=97),
|
||||||
|
image_user=SimpleNamespace(
|
||||||
|
frame_start=1,
|
||||||
|
frame_offset=90,
|
||||||
|
frame_duration=24,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(
|
||||||
|
self.image_depth.extract_frame_range(source_object),
|
||||||
|
(74, 24),
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_extract_frame_range_clamps_negative_offset_to_start(self):
|
||||||
|
source_object = SimpleNamespace(
|
||||||
|
data=SimpleNamespace(source="MOVIE", frame_duration=97),
|
||||||
|
image_user=SimpleNamespace(
|
||||||
|
frame_start=1,
|
||||||
|
frame_offset=-5,
|
||||||
|
frame_duration=24,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(
|
||||||
|
self.image_depth.extract_frame_range(source_object),
|
||||||
|
(1, 24),
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_extract_frame_range_returns_full_clip_for_still(self):
|
||||||
|
source_object = SimpleNamespace(
|
||||||
|
data=SimpleNamespace(source="FILE", frame_duration=1),
|
||||||
|
image_user=SimpleNamespace(
|
||||||
|
frame_start=1,
|
||||||
|
frame_offset=50,
|
||||||
|
frame_duration=24,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(
|
||||||
|
self.image_depth.extract_frame_range(source_object),
|
||||||
|
(1, 1),
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_configure_image_user_preserves_cyclic_setting(self):
|
||||||
|
image_user = SimpleNamespace(
|
||||||
|
frame_start=0,
|
||||||
|
frame_offset=5,
|
||||||
|
use_auto_refresh=False,
|
||||||
|
use_cyclic=False,
|
||||||
|
frame_duration=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.image_depth.configure_image_user(
|
||||||
|
image_user,
|
||||||
|
97,
|
||||||
|
frame_start=4,
|
||||||
|
frame_offset=70,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertFalse(image_user.use_cyclic)
|
||||||
|
self.assertTrue(image_user.use_auto_refresh)
|
||||||
|
self.assertEqual(image_user.frame_start, 4)
|
||||||
|
self.assertEqual(image_user.frame_offset, 70)
|
||||||
|
self.assertEqual(image_user.frame_duration, 97)
|
||||||
|
self.image_depth.configure_image_user(
|
||||||
|
image_user,
|
||||||
|
use_cyclic=True,
|
||||||
|
)
|
||||||
|
self.assertTrue(image_user.use_cyclic)
|
||||||
|
|
||||||
|
def test_image_user_settings_reads_empty_values(self):
|
||||||
|
source_object = SimpleNamespace(
|
||||||
|
image_user=SimpleNamespace(
|
||||||
|
frame_start=4,
|
||||||
|
frame_offset=70,
|
||||||
|
frame_duration=12,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(
|
||||||
|
self.image_depth.image_user_settings(source_object),
|
||||||
|
(4, 70, 12, False),
|
||||||
|
)
|
||||||
|
|
||||||
|
source_object.image_user.use_cyclic = True
|
||||||
|
self.assertEqual(
|
||||||
|
self.image_depth.image_user_settings(source_object),
|
||||||
|
(4, 70, 12, True),
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_image_user_settings_defaults_without_user(self):
|
||||||
|
self.assertEqual(
|
||||||
|
self.image_depth.image_user_settings(
|
||||||
|
SimpleNamespace(image_user=None)
|
||||||
|
),
|
||||||
|
(1, 0, 0, False),
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_image_base_name_strips_video_extension(self):
|
||||||
|
image = SimpleNamespace(name="微信视频2026-07-23_192756_414.mp4")
|
||||||
|
|
||||||
|
self.assertEqual(
|
||||||
|
self.image_depth.image_base_name(image),
|
||||||
|
"微信视频2026-07-23_192756_414",
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_removed_background_name_strips_video_extension(self):
|
||||||
|
source_object = SimpleNamespace(data=SimpleNamespace(name="clip.mp4"))
|
||||||
|
|
||||||
|
self.assertEqual(
|
||||||
|
self.image_depth.removed_background_data_name(source_object),
|
||||||
|
"clip.foreground",
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_force_sequence_source_keeps_loaded_sequence(self):
|
||||||
|
image = SimpleNamespace(source="SEQUENCE")
|
||||||
|
|
||||||
|
result = self.image_depth.force_sequence_source(image)
|
||||||
|
|
||||||
|
self.assertIs(result, image)
|
||||||
|
self.assertEqual(image.source, "SEQUENCE")
|
||||||
|
|
||||||
|
def test_force_sequence_source_switches_single_image(self):
|
||||||
|
image = SimpleNamespace(source="FILE")
|
||||||
|
|
||||||
|
result = self.image_depth.force_sequence_source(image)
|
||||||
|
|
||||||
|
self.assertIs(result, image)
|
||||||
|
self.assertEqual(image.source, "SEQUENCE")
|
||||||
|
|
||||||
|
def test_force_sequence_source_reports_unloadable_sequence(self):
|
||||||
|
class BlockedImage:
|
||||||
|
@property
|
||||||
|
def source(self):
|
||||||
|
return "FILE"
|
||||||
|
|
||||||
|
@source.setter
|
||||||
|
def source(self, value):
|
||||||
|
raise RuntimeError("blocked")
|
||||||
|
|
||||||
|
with self.assertRaisesRegex(RuntimeError, "cannot load"):
|
||||||
|
self.image_depth.force_sequence_source(BlockedImage())
|
||||||
|
|
||||||
|
def test_collect_sequence_frames_returns_numbered_files_in_order(self):
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
root = Path(directory)
|
||||||
|
(root / "frame-0001.jpg").touch()
|
||||||
|
(root / "frame-0002.jpg").touch()
|
||||||
|
(root / "frame-0003.jpg").touch()
|
||||||
|
(root / "notes.txt").touch()
|
||||||
|
image = SimpleNamespace(
|
||||||
|
source="SEQUENCE",
|
||||||
|
filepath=str(root / "frame-0002.jpg"),
|
||||||
|
)
|
||||||
|
|
||||||
|
frames = self.image_depth.collect_sequence_frames(image)
|
||||||
|
|
||||||
|
self.assertEqual(
|
||||||
|
[path.name for path in frames],
|
||||||
|
["frame-0001.jpg", "frame-0002.jpg", "frame-0003.jpg"],
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_sorted_frame_paths_uses_numeric_frame_order(self):
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
root = Path(directory)
|
||||||
|
(root / "frame-2.jpg").touch()
|
||||||
|
(root / "frame-10.jpg").touch()
|
||||||
|
(root / "frame-1.jpg").touch()
|
||||||
|
|
||||||
|
frames = self.image_depth.sorted_frame_paths(root)
|
||||||
|
|
||||||
|
self.assertEqual(
|
||||||
|
[path.name for path in frames],
|
||||||
|
["frame-1.jpg", "frame-2.jpg", "frame-10.jpg"],
|
||||||
|
)
|
||||||
|
|
||||||
def test_accepts_single_depth_map_with_channel_axis(self):
|
def test_accepts_single_depth_map_with_channel_axis(self):
|
||||||
with tempfile.TemporaryDirectory() as directory:
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
npz_path = Path(directory) / "results.npz"
|
npz_path = Path(directory) / "results.npz"
|
||||||
@@ -205,16 +475,6 @@ class ImageDepthTest(unittest.TestCase):
|
|||||||
(12, 8),
|
(12, 8),
|
||||||
)
|
)
|
||||||
|
|
||||||
def test_modifiers_use_native_names_and_fixed_levels(self):
|
|
||||||
source = inspect.getsource(self.image_depth.build_depth_plane)
|
|
||||||
|
|
||||||
self.assertIn('modifiers.new("Subdivision", "SUBSURF")', source)
|
|
||||||
self.assertIn('modifiers.new("Displace", "DISPLACE")', source)
|
|
||||||
self.assertNotIn("DA3 Subdivision", source)
|
|
||||||
self.assertNotIn("DA3 Displace", source)
|
|
||||||
self.assertIn("subdivision.levels = 1", source)
|
|
||||||
self.assertIn("subdivision.render_levels = 2", source)
|
|
||||||
|
|
||||||
def test_plane_fit_transform_matches_source_size_and_offset(self):
|
def test_plane_fit_transform_matches_source_size_and_offset(self):
|
||||||
source = SimpleNamespace(
|
source = SimpleNamespace(
|
||||||
data=SimpleNamespace(size=(800, 400)),
|
data=SimpleNamespace(size=(800, 400)),
|
||||||
@@ -233,37 +493,6 @@ class ImageDepthTest(unittest.TestCase):
|
|||||||
):
|
):
|
||||||
self.assertAlmostEqual(value, expected)
|
self.assertAlmostEqual(value, expected)
|
||||||
|
|
||||||
def test_source_empty_is_only_removed_when_requested(self):
|
|
||||||
source = inspect.getsource(self.image_depth.finish_depth_plane)
|
|
||||||
|
|
||||||
self.assertIn("source_name = source_object.name", source)
|
|
||||||
self.assertIn("if delete_source:", source)
|
|
||||||
self.assertIn(
|
|
||||||
"bpy.data.objects.remove(source_object, do_unlink=True)",
|
|
||||||
source,
|
|
||||||
)
|
|
||||||
self.assertNotIn("source_object.name =", source)
|
|
||||||
self.assertIn(
|
|
||||||
"bpy.data.objects.remove(source_object, do_unlink=True)\n"
|
|
||||||
" plane_object.name = source_name",
|
|
||||||
source,
|
|
||||||
)
|
|
||||||
|
|
||||||
def test_original_depth_mode_uses_uv_and_zero_midlevel(self):
|
|
||||||
source = inspect.getsource(self.image_depth.build_image_depth_plane)
|
|
||||||
|
|
||||||
self.assertIn("texture.image = source_object.data", source)
|
|
||||||
self.assertIn('modifiers.new("Displace", "DISPLACE")', source)
|
|
||||||
self.assertIn("displace.texture = texture", source)
|
|
||||||
self.assertIn('displace.texture_coords = "UV"', source)
|
|
||||||
self.assertNotIn("displace.direction", source)
|
|
||||||
self.assertNotIn("displace.strength", source)
|
|
||||||
self.assertIn(
|
|
||||||
"displace.mid_level = ORIGINAL_DISPLACEMENT_MIDLEVEL",
|
|
||||||
source,
|
|
||||||
)
|
|
||||||
self.assertIn("plane_object.matrix_world =", source)
|
|
||||||
|
|
||||||
def test_process_resolution_scales_the_image_long_edge(self):
|
def test_process_resolution_scales_the_image_long_edge(self):
|
||||||
image = SimpleNamespace(size=(1080, 1920))
|
image = SimpleNamespace(size=(1080, 1920))
|
||||||
|
|
||||||
@@ -272,11 +501,6 @@ class ImageDepthTest(unittest.TestCase):
|
|||||||
960,
|
960,
|
||||||
)
|
)
|
||||||
|
|
||||||
def test_mesh_uses_source_object_name(self):
|
|
||||||
source = inspect.getsource(self.image_depth.create_plane_mesh)
|
|
||||||
|
|
||||||
self.assertIn("bpy.data.meshes.new(source_object.name)", source)
|
|
||||||
|
|
||||||
def test_depth_data_uses_source_image_name(self):
|
def test_depth_data_uses_source_image_name(self):
|
||||||
source_object = SimpleNamespace(
|
source_object = SimpleNamespace(
|
||||||
name="obj_img",
|
name="obj_img",
|
||||||
@@ -292,17 +516,20 @@ class ImageDepthTest(unittest.TestCase):
|
|||||||
self.image_depth.depth_data_name(source_object),
|
self.image_depth.depth_data_name(source_object),
|
||||||
"img.001.depth",
|
"img.001.depth",
|
||||||
)
|
)
|
||||||
estimate_source = inspect.getsource(
|
|
||||||
self.image_depth.build_depth_plane
|
def test_copy_depth_image_renames_still_source(self):
|
||||||
)
|
class FakeImage:
|
||||||
original_source = inspect.getsource(
|
def __init__(self, name):
|
||||||
self.image_depth.build_image_depth_plane
|
self.name = name
|
||||||
)
|
|
||||||
self.assertIn(
|
def copy(self):
|
||||||
"depth_name = depth_name or depth_data_name(source_object)",
|
return FakeImage("img.png.001")
|
||||||
estimate_source,
|
|
||||||
)
|
source_object = SimpleNamespace(data=FakeImage("img.png"))
|
||||||
self.assertIn("depth_data_name(source_object)", original_source)
|
|
||||||
|
depth_image = self.image_depth.copy_depth_image(source_object)
|
||||||
|
|
||||||
|
self.assertEqual(depth_image.name, "img.depth")
|
||||||
|
|
||||||
def test_removed_background_name_uses_source_image_base_name(self):
|
def test_removed_background_name_uses_source_image_base_name(self):
|
||||||
source_object = SimpleNamespace(
|
source_object = SimpleNamespace(
|
||||||
@@ -311,213 +538,14 @@ class ImageDepthTest(unittest.TestCase):
|
|||||||
|
|
||||||
self.assertEqual(
|
self.assertEqual(
|
||||||
self.image_depth.removed_background_data_name(source_object),
|
self.image_depth.removed_background_data_name(source_object),
|
||||||
"img.rb",
|
"img.foreground",
|
||||||
)
|
)
|
||||||
source_object.data.name = "img.png.001"
|
source_object.data.name = "img.png.001"
|
||||||
self.assertEqual(
|
self.assertEqual(
|
||||||
self.image_depth.removed_background_data_name(source_object),
|
self.image_depth.removed_background_data_name(source_object),
|
||||||
"img.001.rb",
|
"img.001.foreground",
|
||||||
)
|
)
|
||||||
|
|
||||||
def test_depth_name_override_keeps_original_base_name(self):
|
|
||||||
source = inspect.getsource(self.image_depth.build_depth_plane)
|
|
||||||
|
|
||||||
self.assertIn(
|
|
||||||
"depth_name = depth_name or depth_data_name(source_object)",
|
|
||||||
source,
|
|
||||||
)
|
|
||||||
self.assertIn(
|
|
||||||
"color_image if color_image is not None else source_object.data",
|
|
||||||
source,
|
|
||||||
)
|
|
||||||
|
|
||||||
def test_original_depth_background_only_changes_material_color(self):
|
|
||||||
source = inspect.getsource(self.image_depth.build_image_depth_plane)
|
|
||||||
|
|
||||||
self.assertIn("texture.image = source_object.data", source)
|
|
||||||
self.assertIn(
|
|
||||||
"source_object.data if displacement_via == \"MATERIAL\" else None",
|
|
||||||
source,
|
|
||||||
)
|
|
||||||
self.assertIn(
|
|
||||||
"color_image if color_image is not None else source_object.data",
|
|
||||||
source,
|
|
||||||
)
|
|
||||||
|
|
||||||
def test_material_uses_source_object_name(self):
|
|
||||||
da3_source = inspect.getsource(self.image_depth.build_depth_plane)
|
|
||||||
image_source = inspect.getsource(
|
|
||||||
self.image_depth.build_image_depth_plane
|
|
||||||
)
|
|
||||||
|
|
||||||
self.assertIn(
|
|
||||||
"create_material(\n source_name,",
|
|
||||||
da3_source,
|
|
||||||
)
|
|
||||||
self.assertIn("create_material(\n source_name,", image_source)
|
|
||||||
self.assertNotIn("DA3DepthMaterial", da3_source)
|
|
||||||
|
|
||||||
def test_displace_texture_uses_extend(self):
|
|
||||||
source = inspect.getsource(
|
|
||||||
self.image_depth.build_depth_plane
|
|
||||||
)
|
|
||||||
|
|
||||||
self.assertIn(
|
|
||||||
"depth_name = depth_name or depth_data_name(source_object)",
|
|
||||||
source,
|
|
||||||
)
|
|
||||||
self.assertIn("create_depth_image(\n depth_name,", source)
|
|
||||||
self.assertIn("bpy.data.textures.new(\n depth_name,", source)
|
|
||||||
self.assertIn('texture.extension = "EXTEND"', source)
|
|
||||||
|
|
||||||
def test_material_supports_bump_and_optional_displacement(self):
|
|
||||||
source = inspect.getsource(self.image_depth.create_material)
|
|
||||||
|
|
||||||
self.assertIn('nodes.new("ShaderNodeBump")', source)
|
|
||||||
self.assertIn(
|
|
||||||
'color_texture.outputs["Color"], bump.inputs["Height"]',
|
|
||||||
source,
|
|
||||||
)
|
|
||||||
self.assertIn(
|
|
||||||
'bump.outputs["Normal"], shader.inputs["Normal"]',
|
|
||||||
source,
|
|
||||||
)
|
|
||||||
self.assertIn(
|
|
||||||
'bump.inputs["Distance"].default_value = '
|
|
||||||
"BUMP_BASE_DISTANCE * depth_scale",
|
|
||||||
source,
|
|
||||||
)
|
|
||||||
self.assertIn('nodes.new("ShaderNodeFloatCurve")', source)
|
|
||||||
self.assertIn(
|
|
||||||
'roughness_curve.inputs["Value"]',
|
|
||||||
source,
|
|
||||||
)
|
|
||||||
self.assertIn(
|
|
||||||
'shader.inputs["Roughness"]',
|
|
||||||
source,
|
|
||||||
)
|
|
||||||
self.assertIn(
|
|
||||||
'curve_points[0].location = (0.0, 0.5)',
|
|
||||||
source,
|
|
||||||
)
|
|
||||||
self.assertIn(
|
|
||||||
'curve_points[-1].location = (1.0, 0.1)',
|
|
||||||
source,
|
|
||||||
)
|
|
||||||
self.assertNotIn('shader.inputs["IOR"]', source)
|
|
||||||
self.assertIn(
|
|
||||||
'roughness_curve.location = (-300.0, 320.0)',
|
|
||||||
source,
|
|
||||||
)
|
|
||||||
self.assertIn('bump.location = (-300.0, -80.0)', source)
|
|
||||||
self.assertIn('nodes.new("ShaderNodeDisplacement")', source)
|
|
||||||
self.assertIn(
|
|
||||||
'depth_texture.outputs["Color"],\n'
|
|
||||||
' displacement.inputs["Height"]',
|
|
||||||
source,
|
|
||||||
)
|
|
||||||
self.assertIn(
|
|
||||||
'displacement.outputs["Displacement"],\n'
|
|
||||||
' output.inputs["Displacement"]',
|
|
||||||
source,
|
|
||||||
)
|
|
||||||
self.assertIn(
|
|
||||||
'displacement.inputs["Midlevel"].default_value',
|
|
||||||
source,
|
|
||||||
)
|
|
||||||
self.assertIn(
|
|
||||||
'displacement.inputs["Scale"].default_value',
|
|
||||||
source,
|
|
||||||
)
|
|
||||||
self.assertIn("set_displacement_only(material)", source)
|
|
||||||
|
|
||||||
def test_unshaded_material_uses_alpha_mixed_emission(self):
|
|
||||||
source = inspect.getsource(self.image_depth.create_material)
|
|
||||||
|
|
||||||
unshaded_start = source.index("if unshaded:")
|
|
||||||
shaded_start = source.index("else:", unshaded_start)
|
|
||||||
displacement_start = source.index(
|
|
||||||
"if displacement_image is not None:",
|
|
||||||
shaded_start,
|
|
||||||
)
|
|
||||||
unshaded_source = source[unshaded_start:shaded_start]
|
|
||||||
shaded_source = source[shaded_start:displacement_start]
|
|
||||||
|
|
||||||
self.assertIn('nodes.new("ShaderNodeEmission")', unshaded_source)
|
|
||||||
self.assertIn(
|
|
||||||
'nodes.new("ShaderNodeBsdfTransparent")',
|
|
||||||
unshaded_source,
|
|
||||||
)
|
|
||||||
self.assertIn('nodes.new("ShaderNodeMixShader")', unshaded_source)
|
|
||||||
self.assertIn(
|
|
||||||
'color_texture.outputs["Color"]',
|
|
||||||
unshaded_source,
|
|
||||||
)
|
|
||||||
self.assertIn('emission.inputs["Color"]', unshaded_source)
|
|
||||||
self.assertIn('color_texture.outputs["Alpha"]', unshaded_source)
|
|
||||||
self.assertIn('mix_shader.inputs["Fac"]', unshaded_source)
|
|
||||||
self.assertIn('transparent.outputs["BSDF"]', unshaded_source)
|
|
||||||
self.assertIn('mix_shader.inputs[1]', unshaded_source)
|
|
||||||
self.assertIn('emission.outputs["Emission"]', unshaded_source)
|
|
||||||
self.assertIn('mix_shader.inputs[2]', unshaded_source)
|
|
||||||
self.assertIn('mix_shader.outputs["Shader"]', unshaded_source)
|
|
||||||
self.assertNotIn("ShaderNodeBsdfPrincipled", unshaded_source)
|
|
||||||
self.assertNotIn("ShaderNodeBump", unshaded_source)
|
|
||||||
self.assertNotIn("ShaderNodeFloatCurve", unshaded_source)
|
|
||||||
self.assertIn("ShaderNodeBsdfPrincipled", shaded_source)
|
|
||||||
self.assertGreater(displacement_start, unshaded_start)
|
|
||||||
|
|
||||||
def test_displacement_values_follow_depth_source(self):
|
|
||||||
estimate_source = inspect.getsource(
|
|
||||||
self.image_depth.build_depth_plane
|
|
||||||
)
|
|
||||||
original_source = inspect.getsource(
|
|
||||||
self.image_depth.build_image_depth_plane
|
|
||||||
)
|
|
||||||
|
|
||||||
self.assertIn(
|
|
||||||
"displacement_midlevel=DISPLACEMENT_MIDLEVEL",
|
|
||||||
estimate_source,
|
|
||||||
)
|
|
||||||
self.assertIn(
|
|
||||||
"displacement_scale=DISPLACEMENT_STRENGTH",
|
|
||||||
estimate_source,
|
|
||||||
)
|
|
||||||
self.assertIn(
|
|
||||||
"displacement_midlevel=ORIGINAL_DISPLACEMENT_MIDLEVEL",
|
|
||||||
original_source,
|
|
||||||
)
|
|
||||||
self.assertIn(
|
|
||||||
"displacement_scale=ORIGINAL_DISPLACEMENT_STRENGTH",
|
|
||||||
original_source,
|
|
||||||
)
|
|
||||||
self.assertEqual(self.image_depth.DISPLACEMENT_MIDLEVEL, 1.0)
|
|
||||||
self.assertEqual(self.image_depth.DISPLACEMENT_STRENGTH, -1.0)
|
|
||||||
self.assertEqual(
|
|
||||||
self.image_depth.ORIGINAL_DISPLACEMENT_MIDLEVEL,
|
|
||||||
0.0,
|
|
||||||
)
|
|
||||||
self.assertEqual(
|
|
||||||
self.image_depth.ORIGINAL_DISPLACEMENT_STRENGTH,
|
|
||||||
1.0,
|
|
||||||
)
|
|
||||||
|
|
||||||
def test_material_displacement_mode_skips_displace_modifier(self):
|
|
||||||
estimate_source = inspect.getsource(
|
|
||||||
self.image_depth.build_depth_plane
|
|
||||||
)
|
|
||||||
original_source = inspect.getsource(
|
|
||||||
self.image_depth.build_image_depth_plane
|
|
||||||
)
|
|
||||||
|
|
||||||
for source in (estimate_source, original_source):
|
|
||||||
self.assertIn('if displacement_via == "MODIFIER":', source)
|
|
||||||
self.assertIn(
|
|
||||||
'if displacement_via == "MATERIAL" else None',
|
|
||||||
source,
|
|
||||||
)
|
|
||||||
self.assertIn("unshaded=unshaded", source)
|
|
||||||
|
|
||||||
def test_surface_displacement_is_set_to_displacement_only(self):
|
def test_surface_displacement_is_set_to_displacement_only(self):
|
||||||
class Material:
|
class Material:
|
||||||
surface_displacement_method = "BUMP"
|
surface_displacement_method = "BUMP"
|
||||||
@@ -530,12 +558,6 @@ class ImageDepthTest(unittest.TestCase):
|
|||||||
"DISPLACEMENT",
|
"DISPLACEMENT",
|
||||||
)
|
)
|
||||||
|
|
||||||
def test_plane_transform_scales_depth_axis(self):
|
|
||||||
source = inspect.getsource(self.image_depth.apply_plane_transform)
|
|
||||||
|
|
||||||
self.assertIn("(0.0, 0.0, scale_z, 0.0)", source)
|
|
||||||
self.assertIn("plane_object.scale.z", source)
|
|
||||||
|
|
||||||
def test_plane_bounds_keep_image_aspect_and_empty_offset(self):
|
def test_plane_bounds_keep_image_aspect_and_empty_offset(self):
|
||||||
source = SimpleNamespace(
|
source = SimpleNamespace(
|
||||||
data=SimpleNamespace(size=(800, 400)),
|
data=SimpleNamespace(size=(800, 400)),
|
||||||
|
|||||||
@@ -1,4 +1,7 @@
|
|||||||
|
import ast
|
||||||
import importlib.util
|
import importlib.util
|
||||||
|
import io
|
||||||
|
import tarfile
|
||||||
import tempfile
|
import tempfile
|
||||||
import unittest
|
import unittest
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -19,25 +22,148 @@ def load_pack():
|
|||||||
|
|
||||||
|
|
||||||
class PackTest(unittest.TestCase):
|
class PackTest(unittest.TestCase):
|
||||||
def test_builds_extension_root_layout(self):
|
def test_builds_platform_specific_extension_zips(self):
|
||||||
pack = load_pack()
|
pack = load_pack()
|
||||||
with tempfile.TemporaryDirectory() as directory:
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
original_dist = pack.DIST_DIR
|
original_dist = pack.DIST_DIR
|
||||||
pack.DIST_DIR = Path(directory)
|
pack.DIST_DIR = Path(directory)
|
||||||
try:
|
try:
|
||||||
output_path = pack.build_extension()
|
output_paths = pack.build_all()
|
||||||
finally:
|
finally:
|
||||||
pack.DIST_DIR = original_dist
|
pack.DIST_DIR = original_dist
|
||||||
|
|
||||||
self.assertEqual(output_path.name, "SnapForm.v0.9.32.zip")
|
self.assertEqual(
|
||||||
with ZipFile(output_path) as archive:
|
{path.name for path in output_paths},
|
||||||
names = set(archive.namelist())
|
{
|
||||||
|
"SnapForm.v0.10.10.windows-x64.zip",
|
||||||
|
"SnapForm.v0.10.10.macos-x64.zip",
|
||||||
|
"SnapForm.v0.10.10.macos-arm64.zip",
|
||||||
|
"SnapForm.v0.10.10.linux-x64.zip",
|
||||||
|
},
|
||||||
|
)
|
||||||
|
for output_path in output_paths:
|
||||||
|
platform = output_path.stem.split(".")[-1]
|
||||||
|
with ZipFile(output_path) as archive:
|
||||||
|
names = set(archive.namelist())
|
||||||
|
manifest = archive.read("blender_manifest.toml").decode(
|
||||||
|
"utf-8"
|
||||||
|
)
|
||||||
|
|
||||||
self.assertIn("__init__.py", names)
|
self.assertIn("__init__.py", names)
|
||||||
self.assertIn("blender_manifest.toml", names)
|
self.assertIn("blender_manifest.toml", names)
|
||||||
self.assertIn("image_depth.py", names)
|
self.assertIn("image_depth.py", names)
|
||||||
self.assertNotIn("snapform_depth_plane.py", names)
|
self.assertIn("worker.py", names)
|
||||||
self.assertIn("worker.py", names)
|
self.assertIn("ai_host.py", names)
|
||||||
self.assertIn("ai_host.py", names)
|
self.assertNotIn("snapform/__init__.py", names)
|
||||||
self.assertNotIn("importer.py", names)
|
self.assertNotIn("importer.py", names)
|
||||||
self.assertNotIn("snapform/__init__.py", names)
|
self.assertIn("sources/depth-anything-3.tar.gz", names)
|
||||||
|
self.assertIn("sources/ben2.tar.gz", names)
|
||||||
|
self.assertIn(
|
||||||
|
f"sources/{pack.PLATFORMS[platform]}",
|
||||||
|
names,
|
||||||
|
)
|
||||||
|
self.assertIn(f'platforms = ["{platform}"]', manifest)
|
||||||
|
self.assertNotIn(
|
||||||
|
"windows-x64"
|
||||||
|
if platform != "windows-x64"
|
||||||
|
else "macos-arm64",
|
||||||
|
manifest.split("platforms =", 1)[-1],
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_rejects_unknown_platform(self):
|
||||||
|
pack = load_pack()
|
||||||
|
with self.assertRaisesRegex(ValueError, "Unsupported platform"):
|
||||||
|
pack.build_extension("windows-arm64")
|
||||||
|
|
||||||
|
def test_bundled_depth_anything_3_has_no_xformers_dependency(self):
|
||||||
|
pack = load_pack()
|
||||||
|
with tarfile.open(
|
||||||
|
fileobj=io.BytesIO(
|
||||||
|
pack.bundled_source_bytes("depth-anything-3.tar.gz")
|
||||||
|
),
|
||||||
|
mode="r:gz",
|
||||||
|
) as archive:
|
||||||
|
raw_pyproject = archive.extractfile(
|
||||||
|
"Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
|
||||||
|
"pyproject.toml"
|
||||||
|
).read()
|
||||||
|
pyproject = raw_pyproject.decode("utf-8")
|
||||||
|
self.assertNotIn("xformers", pyproject)
|
||||||
|
self.assertFalse(raw_pyproject.startswith(b"\xef\xbb\xbf"))
|
||||||
|
|
||||||
|
def test_bundled_depth_anything_3_exports_are_lazy(self):
|
||||||
|
pack = load_pack()
|
||||||
|
with tarfile.open(
|
||||||
|
fileobj=io.BytesIO(
|
||||||
|
pack.bundled_source_bytes("depth-anything-3.tar.gz")
|
||||||
|
),
|
||||||
|
mode="r:gz",
|
||||||
|
) as archive:
|
||||||
|
source = archive.extractfile(
|
||||||
|
"Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
|
||||||
|
"src/depth_anything_3/utils/export/__init__.py"
|
||||||
|
).read().decode("utf-8")
|
||||||
|
tree = ast.parse(source)
|
||||||
|
module_imports = [
|
||||||
|
ast.get_source_segment(source, node)
|
||||||
|
for node in tree.body
|
||||||
|
if isinstance(node, (ast.Import, ast.ImportFrom))
|
||||||
|
]
|
||||||
|
self.assertEqual(
|
||||||
|
module_imports,
|
||||||
|
["from depth_anything_3.specs import Prediction"],
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_official_depth_anything_3_baseline_is_pristine(self):
|
||||||
|
pack = load_pack()
|
||||||
|
archive_path = pack.SOURCES_DIR / "depth-anything-3.tar.gz"
|
||||||
|
with tarfile.open(archive_path, "r:gz") as archive:
|
||||||
|
raw_pyproject = archive.extractfile(
|
||||||
|
"Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
|
||||||
|
"pyproject.toml"
|
||||||
|
).read()
|
||||||
|
export_source = archive.extractfile(
|
||||||
|
"Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
|
||||||
|
"src/depth_anything_3/utils/export/__init__.py"
|
||||||
|
).read().decode("utf-8")
|
||||||
|
pyproject = raw_pyproject.decode("utf-8")
|
||||||
|
self.assertIn("xformers", pyproject)
|
||||||
|
self.assertFalse(raw_pyproject.startswith(b"\xef\xbb\xbf"))
|
||||||
|
self.assertIn(
|
||||||
|
"from depth_anything_3.utils.export.gs import",
|
||||||
|
export_source,
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_bundled_depth_anything_3_qr_falls_back_on_mps(self):
|
||||||
|
pack = load_pack()
|
||||||
|
with tarfile.open(
|
||||||
|
fileobj=io.BytesIO(
|
||||||
|
pack.bundled_source_bytes("depth-anything-3.tar.gz")
|
||||||
|
),
|
||||||
|
mode="r:gz",
|
||||||
|
) as archive:
|
||||||
|
source = archive.extractfile(
|
||||||
|
"Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
|
||||||
|
"src/depth_anything_3/utils/ray_utils.py"
|
||||||
|
).read().decode("utf-8")
|
||||||
|
self.assertIn('use_cpu = device.type == "mps"', source)
|
||||||
|
self.assertIn("A_tilde = A_tilde.cpu()", source)
|
||||||
|
self.assertIn("Q_tilde = Q_tilde.to(device=device)", source)
|
||||||
|
|
||||||
|
def test_official_depth_anything_3_ray_utils_is_pristine(self):
|
||||||
|
pack = load_pack()
|
||||||
|
archive_path = pack.SOURCES_DIR / "depth-anything-3.tar.gz"
|
||||||
|
with tarfile.open(archive_path, "r:gz") as archive:
|
||||||
|
source = archive.extractfile(
|
||||||
|
"Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
|
||||||
|
"src/depth_anything_3/utils/ray_utils.py"
|
||||||
|
).read().decode("utf-8")
|
||||||
|
self.assertIn(
|
||||||
|
"Q_tilde, R_tilde = torch.linalg.qr(A_tilde)",
|
||||||
|
source,
|
||||||
|
)
|
||||||
|
self.assertNotIn("use_cpu", source)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
|
|||||||
@@ -1,5 +1,4 @@
|
|||||||
import importlib.util
|
import importlib.util
|
||||||
import inspect
|
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
import sys
|
import sys
|
||||||
@@ -42,27 +41,11 @@ class WorkerTest(unittest.TestCase):
|
|||||||
|
|
||||||
self.assertEqual([path.name for path in images], ["a.jpg", "b.PNG"])
|
self.assertEqual([path.name for path in images], ["a.jpg", "b.PNG"])
|
||||||
|
|
||||||
def test_reconstruct_reports_status_before_loading_torch(self):
|
|
||||||
source = inspect.getsource(self.worker.reconstruct)
|
|
||||||
|
|
||||||
self.assertLess(
|
|
||||||
source.index('"Starting AI runtime"'),
|
|
||||||
source.index("import torch"),
|
|
||||||
)
|
|
||||||
|
|
||||||
def test_ai_tasks_support_cooperative_cancellation(self):
|
def test_ai_tasks_support_cooperative_cancellation(self):
|
||||||
with self.assertRaises(self.worker.TaskCancelled):
|
with self.assertRaises(self.worker.TaskCancelled):
|
||||||
self.worker.check_cancelled(lambda: True)
|
self.worker.check_cancelled(lambda: True)
|
||||||
|
|
||||||
self.worker.check_cancelled(lambda: False)
|
self.worker.check_cancelled(lambda: False)
|
||||||
self.assertIn(
|
|
||||||
"cancel_check=cancel_check",
|
|
||||||
inspect.getsource(self.worker.reconstruct),
|
|
||||||
)
|
|
||||||
self.assertIn(
|
|
||||||
"cancel_check=cancel_check",
|
|
||||||
inspect.getsource(self.worker.remove_background),
|
|
||||||
)
|
|
||||||
|
|
||||||
def test_select_device_supports_mps_and_auto_fallback(self):
|
def test_select_device_supports_mps_and_auto_fallback(self):
|
||||||
mps_torch = SimpleNamespace(
|
mps_torch = SimpleNamespace(
|
||||||
@@ -115,13 +98,7 @@ class WorkerTest(unittest.TestCase):
|
|||||||
with self.assertRaisesRegex(RuntimeError, "MPS was requested"):
|
with self.assertRaisesRegex(RuntimeError, "MPS was requested"):
|
||||||
self.worker.select_device("mps", cpu_only_torch)
|
self.worker.select_device("mps", cpu_only_torch)
|
||||||
|
|
||||||
def test_optional_background_removal_runs_before_da3_model_load(self):
|
def test_reconstruct_parser_accepts_background_model_directory(self):
|
||||||
source = inspect.getsource(self.worker.reconstruct)
|
|
||||||
|
|
||||||
self.assertLess(
|
|
||||||
source.index("infer_foreground("),
|
|
||||||
source.index("DepthAnything3.from_pretrained"),
|
|
||||||
)
|
|
||||||
args = self.worker.parser().parse_args(
|
args = self.worker.parser().parse_args(
|
||||||
[
|
[
|
||||||
"reconstruct",
|
"reconstruct",
|
||||||
@@ -175,6 +152,199 @@ class WorkerTest(unittest.TestCase):
|
|||||||
self.assertEqual(images, [image_path])
|
self.assertEqual(images, [image_path])
|
||||||
self.assertIsNone(temporary_directory)
|
self.assertIsNone(temporary_directory)
|
||||||
|
|
||||||
|
def test_prepares_a_frame_directory(self):
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
root = Path(directory)
|
||||||
|
(root / "frame-0001.jpg").touch()
|
||||||
|
(root / "frame-0002.png").touch()
|
||||||
|
(root / "notes.txt").touch()
|
||||||
|
|
||||||
|
images, temporary_directory = self.worker.prepare_input(
|
||||||
|
root,
|
||||||
|
16,
|
||||||
|
directory,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(
|
||||||
|
[path.name for path in images],
|
||||||
|
["frame-0001.jpg", "frame-0002.png"],
|
||||||
|
)
|
||||||
|
self.assertIsNone(temporary_directory)
|
||||||
|
|
||||||
|
def test_prepares_first_frames_of_a_frame_directory(self):
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
root = Path(directory)
|
||||||
|
for index in range(1, 6):
|
||||||
|
(root / f"frame-{index:04d}.jpg").touch()
|
||||||
|
|
||||||
|
images, temporary_directory = self.worker.prepare_input(
|
||||||
|
root,
|
||||||
|
3,
|
||||||
|
directory,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(
|
||||||
|
[path.name for path in images],
|
||||||
|
["frame-0001.jpg", "frame-0002.jpg", "frame-0003.jpg"],
|
||||||
|
)
|
||||||
|
self.assertIsNone(temporary_directory)
|
||||||
|
|
||||||
|
def test_prepares_first_frames_after_offset_in_a_frame_directory(self):
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
root = Path(directory)
|
||||||
|
for index in range(1, 6):
|
||||||
|
(root / f"frame-{index:04d}.jpg").touch()
|
||||||
|
|
||||||
|
images, temporary_directory = self.worker.prepare_input(
|
||||||
|
root,
|
||||||
|
3,
|
||||||
|
directory,
|
||||||
|
start_frame=2,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(
|
||||||
|
[path.name for path in images],
|
||||||
|
["frame-0002.jpg", "frame-0003.jpg", "frame-0004.jpg"],
|
||||||
|
)
|
||||||
|
self.assertIsNone(temporary_directory)
|
||||||
|
|
||||||
|
def test_extract_video_frames_reads_frames_sequentially(self):
|
||||||
|
class FakeCapture:
|
||||||
|
def __init__(self):
|
||||||
|
self.position = 0
|
||||||
|
|
||||||
|
def isOpened(self):
|
||||||
|
return True
|
||||||
|
|
||||||
|
def read(self):
|
||||||
|
if self.position >= 100:
|
||||||
|
return False, None
|
||||||
|
frame = self.position
|
||||||
|
self.position += 1
|
||||||
|
return True, frame
|
||||||
|
|
||||||
|
def release(self):
|
||||||
|
return None
|
||||||
|
|
||||||
|
class FakeCv2:
|
||||||
|
def __init__(self):
|
||||||
|
self.captures = []
|
||||||
|
|
||||||
|
def VideoCapture(self, _path):
|
||||||
|
capture = FakeCapture()
|
||||||
|
self.captures.append(capture)
|
||||||
|
return capture
|
||||||
|
|
||||||
|
def imwrite(self, path, frame):
|
||||||
|
Path(path).write_text(str(frame), encoding="utf-8")
|
||||||
|
return True
|
||||||
|
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
root = Path(directory)
|
||||||
|
fake_cv2 = FakeCv2()
|
||||||
|
image_paths = self.worker.extract_video_frames(
|
||||||
|
root / "clip.mkv",
|
||||||
|
root / "out",
|
||||||
|
10,
|
||||||
|
fake_cv2,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(
|
||||||
|
[path.name for path in image_paths],
|
||||||
|
[f"frame.{index:04d}.jpg" for index in range(1, 11)],
|
||||||
|
)
|
||||||
|
self.assertEqual(len(fake_cv2.captures), 1)
|
||||||
|
self.assertEqual(fake_cv2.captures[0].position, 10)
|
||||||
|
|
||||||
|
def test_extract_video_frames_starts_after_offset(self):
|
||||||
|
class FakeCapture:
|
||||||
|
def __init__(self):
|
||||||
|
self.position = 0
|
||||||
|
|
||||||
|
def isOpened(self):
|
||||||
|
return True
|
||||||
|
|
||||||
|
def read(self):
|
||||||
|
if self.position >= 100:
|
||||||
|
return False, None
|
||||||
|
frame = self.position
|
||||||
|
self.position += 1
|
||||||
|
return True, frame
|
||||||
|
|
||||||
|
def release(self):
|
||||||
|
return None
|
||||||
|
|
||||||
|
class FakeCv2:
|
||||||
|
def __init__(self):
|
||||||
|
self.captures = []
|
||||||
|
|
||||||
|
def VideoCapture(self, _path):
|
||||||
|
capture = FakeCapture()
|
||||||
|
self.captures.append(capture)
|
||||||
|
return capture
|
||||||
|
|
||||||
|
def imwrite(self, path, frame):
|
||||||
|
Path(path).write_text(str(frame), encoding="utf-8")
|
||||||
|
return True
|
||||||
|
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
root = Path(directory)
|
||||||
|
fake_cv2 = FakeCv2()
|
||||||
|
image_paths = self.worker.extract_video_frames(
|
||||||
|
root / "clip.mkv",
|
||||||
|
root / "out",
|
||||||
|
10,
|
||||||
|
fake_cv2,
|
||||||
|
start_frame=71,
|
||||||
|
)
|
||||||
|
extracted = [
|
||||||
|
Path(path).read_text(encoding="utf-8")
|
||||||
|
for path in image_paths
|
||||||
|
]
|
||||||
|
|
||||||
|
self.assertEqual(extracted, [str(frame) for frame in range(70, 80)])
|
||||||
|
self.assertEqual(fake_cv2.captures[0].position, 80)
|
||||||
|
|
||||||
|
def test_extract_video_frames_tolerates_short_video_tail(self):
|
||||||
|
class ShortCapture:
|
||||||
|
def __init__(self):
|
||||||
|
self.position = 0
|
||||||
|
|
||||||
|
def isOpened(self):
|
||||||
|
return True
|
||||||
|
|
||||||
|
def read(self):
|
||||||
|
if self.position >= 99:
|
||||||
|
return False, None
|
||||||
|
frame = self.position
|
||||||
|
self.position += 1
|
||||||
|
return True, frame
|
||||||
|
|
||||||
|
def release(self):
|
||||||
|
return None
|
||||||
|
|
||||||
|
class FakeCv2:
|
||||||
|
def VideoCapture(self, _path):
|
||||||
|
return ShortCapture()
|
||||||
|
|
||||||
|
def imwrite(self, path, frame):
|
||||||
|
Path(path).write_text(str(frame), encoding="utf-8")
|
||||||
|
return True
|
||||||
|
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
root = Path(directory)
|
||||||
|
image_paths = self.worker.extract_video_frames(
|
||||||
|
root / "clip.mkv",
|
||||||
|
root / "out",
|
||||||
|
10,
|
||||||
|
FakeCv2(),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(
|
||||||
|
[path.name for path in image_paths],
|
||||||
|
[f"frame.{index:04d}.jpg" for index in range(1, 11)],
|
||||||
|
)
|
||||||
|
|
||||||
def test_npz_is_complete_requires_readable_archive(self):
|
def test_npz_is_complete_requires_readable_archive(self):
|
||||||
with tempfile.TemporaryDirectory() as directory:
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
root = Path(directory)
|
root = Path(directory)
|
||||||
@@ -235,10 +405,6 @@ class WorkerTest(unittest.TestCase):
|
|||||||
self.worker,
|
self.worker,
|
||||||
"create_download_progress_class",
|
"create_download_progress_class",
|
||||||
return_value=object,
|
return_value=object,
|
||||||
), patch.object(
|
|
||||||
self.worker,
|
|
||||||
"resolve_hf_endpoint",
|
|
||||||
return_value=(self.worker.HF_ENDPOINT, False),
|
|
||||||
):
|
):
|
||||||
args.model_dir = str(Path(directory) / "model")
|
args.model_dir = str(Path(directory) / "model")
|
||||||
args.status = str(Path(directory) / "status.json")
|
args.status = str(Path(directory) / "status.json")
|
||||||
@@ -248,13 +414,17 @@ class WorkerTest(unittest.TestCase):
|
|||||||
calls[0]["allow_patterns"],
|
calls[0]["allow_patterns"],
|
||||||
["config.json", "model.safetensors"],
|
["config.json", "model.safetensors"],
|
||||||
)
|
)
|
||||||
|
self.assertEqual(calls[0]["endpoint"], self.worker.HF_ENDPOINT)
|
||||||
|
|
||||||
def test_download_model_switches_to_mirror_before_hf_import(self):
|
def test_download_model_falls_back_to_modelscope_when_direct_fails(self):
|
||||||
calls = []
|
calls = []
|
||||||
fake_hugging_face = ModuleType("huggingface_hub")
|
fake_hugging_face = ModuleType("huggingface_hub")
|
||||||
fake_hugging_face.snapshot_download = lambda **options: calls.append(
|
|
||||||
options
|
def fake_snapshot_download(**options):
|
||||||
)
|
calls.append(options)
|
||||||
|
raise OSError("Connection reset by peer")
|
||||||
|
|
||||||
|
fake_hugging_face.snapshot_download = fake_snapshot_download
|
||||||
args = SimpleNamespace(
|
args = SimpleNamespace(
|
||||||
repository="PramaLLC/BEN2",
|
repository="PramaLLC/BEN2",
|
||||||
model_dir="model",
|
model_dir="model",
|
||||||
@@ -269,23 +439,151 @@ class WorkerTest(unittest.TestCase):
|
|||||||
self.worker,
|
self.worker,
|
||||||
"create_download_progress_class",
|
"create_download_progress_class",
|
||||||
return_value=object,
|
return_value=object,
|
||||||
), patch.object(
|
|
||||||
self.worker,
|
|
||||||
"resolve_hf_endpoint",
|
|
||||||
return_value=("https://hf-mirror.com", True),
|
|
||||||
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
|
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
|
||||||
self.worker,
|
self.worker,
|
||||||
"write_status",
|
"download_from_modelscope",
|
||||||
) as mocked_write:
|
) as mocked_modelscope:
|
||||||
|
def fake_download_from_modelscope(*args, **kwargs):
|
||||||
|
model_dir = Path(args[1])
|
||||||
|
model_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
(model_dir / "config.json").write_text("{}")
|
||||||
|
(model_dir / "model.safetensors").write_bytes(b"x")
|
||||||
|
|
||||||
|
mocked_modelscope.side_effect = fake_download_from_modelscope
|
||||||
args.model_dir = str(Path(directory) / "model")
|
args.model_dir = str(Path(directory) / "model")
|
||||||
args.status = str(Path(directory) / "status.json")
|
args.status = str(Path(directory) / "status.json")
|
||||||
self.worker.download_model(args)
|
self.worker.download_model(args)
|
||||||
self.assertEqual(
|
|
||||||
os.environ["HF_ENDPOINT"],
|
self.assertEqual(len(calls), 1)
|
||||||
"https://hf-mirror.com",
|
self.assertEqual(calls[0]["endpoint"], self.worker.HF_ENDPOINT)
|
||||||
)
|
self.assertEqual(mocked_modelscope.call_count, 1)
|
||||||
messages = [call.args[2] for call in mocked_write.call_args_list]
|
|
||||||
self.assertTrue(any("(mirror)" in message for message in messages))
|
def test_download_model_skips_unreachable_endpoint(self):
|
||||||
|
calls = []
|
||||||
|
fake_hugging_face = ModuleType("huggingface_hub")
|
||||||
|
fake_hugging_face.snapshot_download = lambda **options: calls.append(
|
||||||
|
options
|
||||||
|
)
|
||||||
|
args = SimpleNamespace(
|
||||||
|
repository="depth-anything/DA3-SMALL",
|
||||||
|
model_dir="model",
|
||||||
|
status="status.json",
|
||||||
|
allow_pattern=[],
|
||||||
|
)
|
||||||
|
|
||||||
|
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
||||||
|
sys.modules,
|
||||||
|
{"huggingface_hub": fake_hugging_face},
|
||||||
|
), patch.object(
|
||||||
|
self.worker,
|
||||||
|
"create_download_progress_class",
|
||||||
|
return_value=object,
|
||||||
|
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
|
||||||
|
self.worker,
|
||||||
|
"hf_endpoint_reachable",
|
||||||
|
return_value=False,
|
||||||
|
), patch.object(
|
||||||
|
self.worker,
|
||||||
|
"download_from_modelscope",
|
||||||
|
) as mocked_modelscope, patch.object(
|
||||||
|
self.worker,
|
||||||
|
"_verify_model_complete",
|
||||||
|
return_value=None,
|
||||||
|
):
|
||||||
|
args.model_dir = str(Path(directory) / "model")
|
||||||
|
args.status = str(Path(directory) / "status.json")
|
||||||
|
self.worker.download_model(args)
|
||||||
|
|
||||||
|
self.assertEqual(calls, [])
|
||||||
|
self.assertEqual(mocked_modelscope.call_count, 1)
|
||||||
|
|
||||||
|
def test_download_model_refuses_success_when_files_missing(self):
|
||||||
|
fake_hugging_face = ModuleType("huggingface_hub")
|
||||||
|
fake_hugging_face.snapshot_download = lambda **options: None
|
||||||
|
args = SimpleNamespace(
|
||||||
|
repository="PramaLLC/BEN2",
|
||||||
|
model_dir="model",
|
||||||
|
status="status.json",
|
||||||
|
allow_pattern=[],
|
||||||
|
)
|
||||||
|
|
||||||
|
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
||||||
|
sys.modules,
|
||||||
|
{"huggingface_hub": fake_hugging_face},
|
||||||
|
), patch.object(
|
||||||
|
self.worker,
|
||||||
|
"create_download_progress_class",
|
||||||
|
return_value=object,
|
||||||
|
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
|
||||||
|
self.worker,
|
||||||
|
"download_from_modelscope",
|
||||||
|
):
|
||||||
|
args.model_dir = str(Path(directory) / "model")
|
||||||
|
args.status = str(Path(directory) / "status.json")
|
||||||
|
with self.assertRaisesRegex(
|
||||||
|
RuntimeError,
|
||||||
|
"All download sources failed",
|
||||||
|
):
|
||||||
|
self.worker.download_model(args)
|
||||||
|
|
||||||
|
def test_download_model_reports_all_failed_sources(self):
|
||||||
|
fake_hugging_face = ModuleType("huggingface_hub")
|
||||||
|
|
||||||
|
def fake_snapshot_download(**options):
|
||||||
|
raise OSError("direct down")
|
||||||
|
|
||||||
|
fake_hugging_face.snapshot_download = fake_snapshot_download
|
||||||
|
args = SimpleNamespace(
|
||||||
|
repository="depth-anything/DA3-SMALL",
|
||||||
|
model_dir="model",
|
||||||
|
status="status.json",
|
||||||
|
allow_pattern=[],
|
||||||
|
)
|
||||||
|
|
||||||
|
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
||||||
|
sys.modules,
|
||||||
|
{"huggingface_hub": fake_hugging_face},
|
||||||
|
), patch.object(
|
||||||
|
self.worker,
|
||||||
|
"create_download_progress_class",
|
||||||
|
return_value=object,
|
||||||
|
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
|
||||||
|
self.worker,
|
||||||
|
"download_from_modelscope",
|
||||||
|
side_effect=OSError("mirror down"),
|
||||||
|
):
|
||||||
|
args.model_dir = str(Path(directory) / "model")
|
||||||
|
args.status = str(Path(directory) / "status.json")
|
||||||
|
with self.assertRaisesRegex(RuntimeError, "modelscope.cn"):
|
||||||
|
self.worker.download_model(args)
|
||||||
|
|
||||||
|
def test_download_from_modelscope_downloads_selected_files(self):
|
||||||
|
args = SimpleNamespace(
|
||||||
|
repository="depth-anything/DA3-SMALL",
|
||||||
|
model_dir="model",
|
||||||
|
status="status.json",
|
||||||
|
allow_pattern=["config.json"],
|
||||||
|
)
|
||||||
|
downloaded = []
|
||||||
|
|
||||||
|
with tempfile.TemporaryDirectory() as directory, patch.object(
|
||||||
|
self.worker,
|
||||||
|
"modelscope_file_list",
|
||||||
|
return_value=[
|
||||||
|
("config.json", 100),
|
||||||
|
("model.safetensors", 200),
|
||||||
|
(".gitattributes", 10),
|
||||||
|
],
|
||||||
|
), patch.object(
|
||||||
|
self.worker,
|
||||||
|
"_download_modelscope_file",
|
||||||
|
side_effect=lambda *args, **kwargs: downloaded.append(args[1]),
|
||||||
|
), patch.object(self.worker, "write_status"):
|
||||||
|
args.model_dir = str(Path(directory) / "model")
|
||||||
|
args.status = str(Path(directory) / "status.json")
|
||||||
|
self.worker.download_from_modelscope(args, args.model_dir)
|
||||||
|
|
||||||
|
self.assertEqual(downloaded, ["config.json"])
|
||||||
|
|
||||||
def test_ben2_background_removal_writes_fast_transparent_png(self):
|
def test_ben2_background_removal_writes_fast_transparent_png(self):
|
||||||
inference_options = []
|
inference_options = []
|
||||||
@@ -368,6 +666,166 @@ class WorkerTest(unittest.TestCase):
|
|||||||
)
|
)
|
||||||
self.assertIn("Background removed in", status["message"])
|
self.assertIn("Background removed in", status["message"])
|
||||||
|
|
||||||
|
def test_ben2_background_removal_processes_each_frame(self):
|
||||||
|
inference_options = []
|
||||||
|
|
||||||
|
class FakeImage:
|
||||||
|
def __enter__(self):
|
||||||
|
return self
|
||||||
|
|
||||||
|
def __exit__(self, *_args):
|
||||||
|
return None
|
||||||
|
|
||||||
|
def convert(self, mode):
|
||||||
|
self.mode = mode
|
||||||
|
return self
|
||||||
|
|
||||||
|
def save(self, path, **options):
|
||||||
|
Path(path).touch()
|
||||||
|
|
||||||
|
class FakeModel:
|
||||||
|
@classmethod
|
||||||
|
def from_pretrained(cls, model_dir):
|
||||||
|
return cls()
|
||||||
|
|
||||||
|
def to(model_self, device):
|
||||||
|
return model_self
|
||||||
|
|
||||||
|
def eval(self):
|
||||||
|
return None
|
||||||
|
|
||||||
|
def inference(self, image, **options):
|
||||||
|
inference_options.append((image, options))
|
||||||
|
return FakeImage()
|
||||||
|
|
||||||
|
fake_torch = ModuleType("torch")
|
||||||
|
fake_torch.cuda = SimpleNamespace(
|
||||||
|
is_available=lambda: False,
|
||||||
|
synchronize=lambda: None,
|
||||||
|
)
|
||||||
|
fake_torch.device = lambda device_type: SimpleNamespace(type=device_type)
|
||||||
|
fake_torch.inference_mode = nullcontext
|
||||||
|
fake_ben2 = ModuleType("ben2")
|
||||||
|
fake_ben2.BEN_Base = FakeModel
|
||||||
|
fake_pil = ModuleType("PIL")
|
||||||
|
fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage())
|
||||||
|
|
||||||
|
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
||||||
|
sys.modules,
|
||||||
|
{
|
||||||
|
"torch": fake_torch,
|
||||||
|
"ben2": fake_ben2,
|
||||||
|
"PIL": fake_pil,
|
||||||
|
},
|
||||||
|
):
|
||||||
|
root = Path(directory)
|
||||||
|
input_directory = root / "frames"
|
||||||
|
input_directory.mkdir()
|
||||||
|
(input_directory / "frame-0001.jpg").touch()
|
||||||
|
(input_directory / "frame-0002.jpg").touch()
|
||||||
|
args = SimpleNamespace(
|
||||||
|
input=str(input_directory),
|
||||||
|
model_dir="model",
|
||||||
|
output_dir=str(root / "result"),
|
||||||
|
status=str(root / "status.json"),
|
||||||
|
device="cpu",
|
||||||
|
video_frames=8,
|
||||||
|
)
|
||||||
|
self.worker.remove_background(args)
|
||||||
|
|
||||||
|
result_directory = Path(args.output_dir)
|
||||||
|
self.assertTrue(
|
||||||
|
(result_directory / "foreground.0001.png").is_file()
|
||||||
|
)
|
||||||
|
self.assertTrue(
|
||||||
|
(result_directory / "foreground.0002.png").is_file()
|
||||||
|
)
|
||||||
|
self.assertTrue(
|
||||||
|
(result_directory / "frame.0001.jpg").is_file()
|
||||||
|
)
|
||||||
|
self.assertTrue(
|
||||||
|
(result_directory / "frame.0002.jpg").is_file()
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(len(inference_options), 2)
|
||||||
|
for _image, options in inference_options:
|
||||||
|
self.assertEqual(options, {"refine_foreground": False})
|
||||||
|
|
||||||
|
def test_remove_background_parser_accepts_video_frames(self):
|
||||||
|
args = self.worker.parser().parse_args(
|
||||||
|
[
|
||||||
|
"remove-background",
|
||||||
|
"--input",
|
||||||
|
"input.mp4",
|
||||||
|
"--model-dir",
|
||||||
|
"ben2",
|
||||||
|
"--output-dir",
|
||||||
|
"result",
|
||||||
|
"--status",
|
||||||
|
"status.json",
|
||||||
|
"--device",
|
||||||
|
"auto",
|
||||||
|
"--video-frames",
|
||||||
|
"24",
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(args.command, "remove-background")
|
||||||
|
self.assertEqual(args.video_frames, 24)
|
||||||
|
|
||||||
|
def test_extract_frames_parser_accepts_frame_range(self):
|
||||||
|
args = self.worker.parser().parse_args(
|
||||||
|
[
|
||||||
|
"extract-frames",
|
||||||
|
"--input",
|
||||||
|
"frames",
|
||||||
|
"--output-dir",
|
||||||
|
"result",
|
||||||
|
"--status",
|
||||||
|
"status.json",
|
||||||
|
"--video-frames",
|
||||||
|
"24",
|
||||||
|
"--start-frame",
|
||||||
|
"70",
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(args.command, "extract-frames")
|
||||||
|
self.assertEqual(args.video_frames, 24)
|
||||||
|
self.assertEqual(args.start_frame, 70)
|
||||||
|
|
||||||
|
def test_extract_frames_copies_selected_frames(self):
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
root = Path(directory)
|
||||||
|
input_directory = root / "frames"
|
||||||
|
input_directory.mkdir()
|
||||||
|
for index in range(1, 6):
|
||||||
|
(input_directory / f"frame-{index:04d}.jpg").touch()
|
||||||
|
args = SimpleNamespace(
|
||||||
|
input=str(input_directory),
|
||||||
|
output_dir=str(root / "result"),
|
||||||
|
status=str(root / "status.json"),
|
||||||
|
video_frames=3,
|
||||||
|
start_frame=2,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.worker.extract_frames(args)
|
||||||
|
|
||||||
|
result_directory = Path(args.output_dir)
|
||||||
|
self.assertEqual(
|
||||||
|
sorted(
|
||||||
|
path.name
|
||||||
|
for path in result_directory.iterdir()
|
||||||
|
if path.is_file()
|
||||||
|
),
|
||||||
|
["frame.0001.jpg", "frame.0002.jpg", "frame.0003.jpg"],
|
||||||
|
)
|
||||||
|
status = json.loads(
|
||||||
|
Path(args.status).read_text(encoding="utf-8")
|
||||||
|
)
|
||||||
|
self.assertEqual(status["progress"], 1.0)
|
||||||
|
self.assertIn("3 frame", status["message"])
|
||||||
|
|
||||||
def test_uses_official_da3_export_parameters_for_all_four_formats(self):
|
def test_uses_official_da3_export_parameters_for_all_four_formats(self):
|
||||||
inference_options = []
|
inference_options = []
|
||||||
export_calls = []
|
export_calls = []
|
||||||
@@ -500,6 +958,95 @@ class WorkerTest(unittest.TestCase):
|
|||||||
infer_gs_values = [options["infer_gs"] for _, options in inference_options]
|
infer_gs_values = [options["infer_gs"] for _, options in inference_options]
|
||||||
self.assertEqual(infer_gs_values, [False, True, False, False])
|
self.assertEqual(infer_gs_values, [False, True, False, False])
|
||||||
|
|
||||||
|
def test_reconstruct_copies_frames_for_multi_frame_input(self):
|
||||||
|
class FakeModel:
|
||||||
|
@classmethod
|
||||||
|
def from_pretrained(cls, model_dir):
|
||||||
|
return cls()
|
||||||
|
|
||||||
|
def to(self, device):
|
||||||
|
return self
|
||||||
|
|
||||||
|
def eval(self):
|
||||||
|
return None
|
||||||
|
|
||||||
|
def inference(self, images, **options):
|
||||||
|
return SimpleNamespace()
|
||||||
|
|
||||||
|
def fake_export(prediction, export_format, export_dir, **options):
|
||||||
|
mini_directory = Path(export_dir) / "exports" / "mini_npz"
|
||||||
|
mini_directory.mkdir(parents=True)
|
||||||
|
np.savez_compressed(
|
||||||
|
mini_directory / "results.npz",
|
||||||
|
depth=np.zeros((2, 14, 14)),
|
||||||
|
)
|
||||||
|
|
||||||
|
fake_torch = ModuleType("torch")
|
||||||
|
fake_torch.device = lambda device_type: SimpleNamespace(
|
||||||
|
type=device_type,
|
||||||
|
)
|
||||||
|
fake_torch.inference_mode = nullcontext
|
||||||
|
fake_api = ModuleType("depth_anything_3.api")
|
||||||
|
fake_api.DepthAnything3 = FakeModel
|
||||||
|
fake_export_module = ModuleType("depth_anything_3.utils.export")
|
||||||
|
fake_export_module.export = fake_export
|
||||||
|
fake_utils = ModuleType("depth_anything_3.utils")
|
||||||
|
fake_utils.__path__ = []
|
||||||
|
fake_package = ModuleType("depth_anything_3")
|
||||||
|
fake_package.__path__ = []
|
||||||
|
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
root = Path(directory)
|
||||||
|
input_directory = root / "frames"
|
||||||
|
input_directory.mkdir()
|
||||||
|
(input_directory / "frame-0001.jpg").touch()
|
||||||
|
(input_directory / "frame-0002.jpg").touch()
|
||||||
|
with patch.dict(
|
||||||
|
sys.modules,
|
||||||
|
{
|
||||||
|
"torch": fake_torch,
|
||||||
|
"depth_anything_3": fake_package,
|
||||||
|
"depth_anything_3.api": fake_api,
|
||||||
|
"depth_anything_3.utils": fake_utils,
|
||||||
|
"depth_anything_3.utils.export": fake_export_module,
|
||||||
|
},
|
||||||
|
):
|
||||||
|
args = SimpleNamespace(
|
||||||
|
input=str(input_directory),
|
||||||
|
video_frames=16,
|
||||||
|
output_dir=str(root / "result"),
|
||||||
|
export_format="mini_npz",
|
||||||
|
status=str(root / "status.json"),
|
||||||
|
device="cpu",
|
||||||
|
model_dir="model",
|
||||||
|
process_res=518,
|
||||||
|
process_res_method="upper_bound_resize",
|
||||||
|
use_ray_pose=1,
|
||||||
|
ref_view_strategy="first",
|
||||||
|
conf_thresh_percentile=40.0,
|
||||||
|
num_max_points=1000,
|
||||||
|
show_cameras=0,
|
||||||
|
gs_views_interval=1,
|
||||||
|
remove_background_model_dir="",
|
||||||
|
)
|
||||||
|
self.worker.reconstruct(args)
|
||||||
|
|
||||||
|
result_directory = Path(args.output_dir)
|
||||||
|
self.assertTrue(
|
||||||
|
(result_directory / "frame.0001.jpg").is_file()
|
||||||
|
)
|
||||||
|
self.assertTrue(
|
||||||
|
(result_directory / "frame.0002.jpg").is_file()
|
||||||
|
)
|
||||||
|
self.assertTrue(
|
||||||
|
(result_directory / "results.npz").is_file()
|
||||||
|
)
|
||||||
|
self.assertFalse(
|
||||||
|
(result_directory / "exports").exists()
|
||||||
|
)
|
||||||
|
status = json.loads(Path(args.status).read_text(encoding="utf-8"))
|
||||||
|
self.assertEqual(status["progress"], 1.0)
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|||||||
Reference in New Issue
Block a user