第二步SAM模型接入
This commit is contained in:
Ukawka
+3
-2
@@ -9,5 +9,6 @@ VECTOR_ENGINE_IMAGE_REQUEST_TIMEOUT_MS=1200000
|
||||
SPINE_UNDERSTANDING_MODEL=gpt-6-astra
|
||||
SPINE_COMPLETION_MODEL=qwen-image-3.0
|
||||
|
||||
# Checkpoint location only; the SAM runtime is integrated in milestone 2.
|
||||
SAM_MODEL_DIR=F:/Study/sam3.1
|
||||
# Local source + checkpoint root. SAM runs in a temporary ai-spine GPU process.
|
||||
SAM_MODEL_DIR=F:/Study/SAM3
|
||||
SAM_CODE_DIR=F:/Study/SAM3/sam3
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
2. 使用 SAM 根据语义、粗框和程序候选生成 mask，经语义/几何检查后从原图提取可见像素。
|
||||
3. 使用图片编辑模型补全被遮挡的部件，再进行几何验收和 Spine 导出。
|
||||
|
||||
目前第一阶段已经可以完整运行。SAM 3.1 checkpoint 路径和图片编辑适配器已经接入配置，但 SAM 推理、补全结果合并及新管线的 Spine 导出尚未完成。
|
||||
目前第一阶段和第二阶段均可运行。第二阶段在 `ai-spine` Conda 环境的临时子进程中加载本地 SAM 3.1 权重，使用 GPU 分割并生成半透明标注图、白底拆解图和透明部件。分割结果仍需语义复核；补全结果合并及新管线的 Spine 导出尚未完成。
|
||||
|
||||
仓库仍保留旧版“AI 生成白底拆件图 → OpenCV 拆分 → Genielabs 定位 → Spine 导出”流程，用于对照和回归测试。
|
||||
|
||||
@@ -16,6 +16,7 @@
|
||||
|
||||
- [项目修改约定](AGENTS.md)：每次修改都必须遵守，Python 固定使用 `ai-spine` Conda 环境。
|
||||
- [第一阶段：视觉 LLM 角色结构理解](docs/1-understanding-stage.md)：执行流程、代码对应关系、数据格式和输出文件。
|
||||
- [第二阶段：本地 SAM 分割](docs/2-segmentation-stage.md)：环境、运行命令、权重适配、输出和质量边界。
|
||||
- [模型阶段配置与实验说明](docs/model-stages.md)：模型切换、协议适配、图片编辑探针、缓存和错误处理。
|
||||
|
||||
## 环境准备
|
||||
@@ -47,8 +48,9 @@ VECTOR_ENGINE_API_KEY=你的令牌
|
||||
SPINE_UNDERSTANDING_MODEL=deepseek-flash
|
||||
SPINE_COMPLETION_MODEL=qwen-image-3.0
|
||||
|
||||
# 当前只记录 checkpoint；SAM 推理尚未接入
|
||||
SAM_MODEL_DIR=F:/Study/sam3.1
|
||||
# 本地 SAM 根目录和源码仓库；分割结束后子进程退出
|
||||
SAM_MODEL_DIR=F:/Study/SAM3
|
||||
SAM_CODE_DIR=F:/Study/SAM3/sam3
|
||||
```
|
||||
|
||||
模型 ID 会原样发送给中转站，不进行隐式别名转换。不要把 API Key 写进 JSON 配置或提交到版本库。
|
||||
@@ -65,11 +67,13 @@ SAM_MODEL_DIR=F:/Study/sam3.1
|
||||
- `Pillow`、`numpy`、`opencv-python`；
|
||||
- 相邻目录中的 `spine-animation-ai` 项目。
|
||||
|
||||
在 `AI_spine` 目录安装缺失的 Python 依赖（已满足版本要求的包会保留）：
|
||||
当前本机验证组合为 Python 3.11、PyTorch 2.7.1 + CUDA 12.6、torchvision 0.22.1。在 `AI_spine` 目录补全依赖；已有此 CUDA 版本时跳过第一条安装命令：
|
||||
|
||||
```powershell
|
||||
conda activate ai-spine
|
||||
python -m pip install torch==2.7.1 torchvision==0.22.1 --index-url https://download.pytorch.org/whl/cu126
|
||||
python -m pip install -r .\requirements.txt
|
||||
python -m pip install --no-deps --no-build-isolation -e F:\Study\SAM3\sam3
|
||||
```
|
||||
|
||||
不激活终端也可明确指定安装目标：
|
||||
@@ -168,6 +172,22 @@ node .\pipeline.mjs --mode annotated-parts `
|
||||
|
||||
每组模型实验应使用独立的输出目录，避免覆盖对照结果。只有输入、配置、代码和产物哈希全部一致时，才可以通过 `--resume` 复用结果。
|
||||
|
||||
### 使用本地 SAM 生成标注和拆解图
|
||||
|
||||
复用已有第一阶段结果，不重新请求 DeepSeek，也不需要中转站密钥：
|
||||
|
||||
```powershell
|
||||
conda run --no-capture-output -n ai-spine node .\tools\model-check.mjs `
|
||||
--stage segmentation `
|
||||
--input-dir .\output\sombrero-grounding-v2 `
|
||||
--out-dir .\output\sombrero-sam-local `
|
||||
--sam-dir F:/Study/SAM3
|
||||
```
|
||||
|
||||
打开输出目录的 `report.html`：`annotated.png` 是原图上的不同半透明颜色标注，`parts-sheet.png` 是白底拆解图，`parts/` 保存原像素透明部件。运行期间加载 GPU 权重，完成后子进程退出，不启动常驻服务。
|
||||
|
||||
加入 `--dry-run` 可检查环境与本地路径，不加载权重推理。输入必须含 `source.png`、v2 `semantic-parts.json`、`segmentation-prompts.json`。完整参数与质量检查见[第二阶段文档](docs/2-segmentation-stage.md)。
|
||||
|
||||
### 测试图片补全接口
|
||||
|
||||
此命令只验证图片编辑适配器并保存原始候选，不代表补全部件已经验收通过：
|
||||
|
||||
@@ -39,7 +39,7 @@
|
||||
"segmentation": {
|
||||
"provider": "sam-local",
|
||||
"model_dir_env": "SAM_MODEL_DIR",
|
||||
"default_model_dir": "F:/Study/sam3.1"
|
||||
"default_model_dir": "F:/Study/SAM3"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -39,7 +39,7 @@
|
||||
"segmentation": {
|
||||
"provider": "sam-local",
|
||||
"model_dir_env": "SAM_MODEL_DIR",
|
||||
"default_model_dir": "F:/Study/sam3.1"
|
||||
"default_model_dir": "F:/Study/SAM3"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -363,8 +363,8 @@ POST /v1/responses
|
||||
= 可信的像素 mask + 透明拆件 + 隐藏区域补全 + 最终骨骼
|
||||
~~~
|
||||
|
||||
下一阶段应读取 <code>semantic-parts.json</code> 和 <code>segmentation-prompts.json</code>，仅使用 status=candidate 的语义、粗框和程序候选点。模型原始正负点不传给 SAM。候选始终 semantic_verified=false：框内可能混有多个部件，未来仍需要 mask 语义选择、越出角色 Alpha 比例等 QA。
|
||||
已实现的[第二阶段](2-segmentation-stage.md)读取 <code>source.png</code>、<code>semantic-parts.json</code> 和 <code>segmentation-prompts.json</code>，仅处理 status=candidate 的可见部件。SAM 图像概念接口接收文本和框；程序候选点参与返回 mask 的排序，不作为点提示传入此接口。模型原始正负点不传给 SAM。候选始终 semantic_verified=false：框内可能混有多个部件，Alpha、重叠和连通域检查不能替代语义复核。
|
||||
|
||||
已有 v1 结果可以通过 <code>--parts-file 旧计划.json --reference 原图.png --out-dir 新目录</code> 离线迁移与校验，不请求模型。迁移按 ID 推断骨骼/附件类型，袖子优先绑定同侧上臂，面部绑定头部；部分可见部件开启遮挡补全，骨骼连接处保守建议延展。这些推断在 notes 标注为待复核，不替代模型或人工判断。输入计划哈希和所有新产物参与 resume 校验。
|
||||
|
||||
<code>manifest.json</code> 当前只记录 SAM checkpoint 是否存在，<code>runtime_status</code> 仍为 <code>not_integrated</code>，不会加载权重或占用 GPU。
|
||||
第一阶段 <code>manifest.json</code> 只记录 SAM 路径信息，<code>runtime_status=local_subprocess_not_probed</code>，第一阶段不加载 SAM。第二阶段单独探测 Conda/CUDA、严格加载本地权重，并在自己的 manifest 中记录运行环境、权重哈希及子进程退出状态。
|
||||
|
||||
@@ -0,0 +1,108 @@
|
||||
# 第二阶段：本地 SAM 分割
|
||||
|
||||
本阶段读取第一阶段的语义计划，使用本地 SAM 3.1 图像检测分支生成部件 mask，从原图提取可见像素，并输出两张预览图。Python 固定使用 `ai-spine` Conda 环境，遵守 [AGENTS.md](../AGENTS.md)。
|
||||
|
||||
## 运行方式
|
||||
|
||||
在 `AI_spine` 目录执行：
|
||||
|
||||
```powershell
|
||||
conda run --no-capture-output -n ai-spine node .\tools\model-check.mjs `
|
||||
--stage segmentation `
|
||||
--input-dir .\output\sombrero-grounding-v2 `
|
||||
--out-dir .\output\sombrero-sam-local `
|
||||
--sam-dir F:/Study/SAM3
|
||||
```
|
||||
|
||||
输出目录必须是新目录。相同输入、代码、权重、参数及完整产物可加 `--resume` 复用；失败结果保留，另建目录重跑。独立分割不调用视觉 LLM 或图片编辑接口，不需要 API Key。
|
||||
|
||||
输入目录必须包含：
|
||||
|
||||
| 文件 | 要求 |
|
||||
| --- | --- |
|
||||
| `source.png` | 第一阶段预处理后的原尺寸 RGBA |
|
||||
| `semantic-parts.json` | 通过结构校验的 v2 语义计划 |
|
||||
| `segmentation-prompts.json` | 与计划 ID、画布匹配的程序候选提示 |
|
||||
|
||||
仅 `status=candidate` 的可见部件进入推理。hidden/blocked 记入跳过列表；无可用 SAM mask 的部件记录失败原因，不用矩形或候选点伪造分割。
|
||||
|
||||
可以在命令末尾加 `--dry-run` 检查代码路径、权重存在性、Conda、依赖导入和 CUDA；它不加载权重，因此不验证权重内容或完整推理显存需求。
|
||||
|
||||
如需从第一阶段接续运行，在原结构理解命令中使用 `--stop-after segment --sam-dir F:/Study/SAM3`。第二阶段输出到理解目录下的 `segmentation/`。原理解命令仍可能请求付费模型；已有计划优先使用上面的独立命令。
|
||||
|
||||
## 环境与权重
|
||||
|
||||
本机已经验证：Python 3.11、PyTorch `2.7.1+cu126`、torchvision `0.22.1+cu126`、NumPy `1.26.4`、OpenCV `4.11.0.86`、Windows Triton `3.3.1.post21`，GPU 为 RTX 3060 Laptop 6 GB。
|
||||
|
||||
环境补全命令见 [README](../README.md#python-与外部项目)，依赖记录在 [requirements.txt](../requirements.txt)。程序不会自动安装依赖、创建环境或切换到 CPU。
|
||||
|
||||
默认布局：
|
||||
|
||||
```text
|
||||
F:/Study/SAM3/
|
||||
sam3/ 源码仓库
|
||||
sam3/model_builder.py
|
||||
sam3.1checkpoint/
|
||||
sam3.1_multiplex.pt
|
||||
```
|
||||
|
||||
路径优先级为 CLI 覆盖、环境变量、配置默认值。`SAM_MODEL_DIR` 指定根目录，`SAM_CODE_DIR` 可单独指定源码仓库。若 `.env` 中仍是旧路径，修改为上述路径或显式传入 CLI 参数。
|
||||
|
||||
| 参数 | 含义 |
|
||||
| --- | --- |
|
||||
| `--sam-dir` | SAM 根目录 |
|
||||
| `--sam-code-dir` | 源码仓库目录，内部应包含 `sam3/model_builder.py` |
|
||||
| `--sam-checkpoint` | 本地权重文件 |
|
||||
| `--sam-confidence` | SAM 候选分数阈值，默认 0.2，不是语义正确率 |
|
||||
| `--sam-precision` | 默认 bf16 自动混合精度；可设 fp32，可能需要更多显存 |
|
||||
|
||||
Node 先启动环境探测子进程，再启动一次推理子进程。模型只加载到后者；没有监听端口或常驻服务，正常结束、报错、超时或收到取消信号后退出。推理超时为 20 分钟。禁用 Hugging Face 在线下载，权重只能来自本地。
|
||||
|
||||
`sam_runtime.py` 从 multiplex checkpoint 提取 `detector.*`。SAM 3.1 的检测分支有三个特征尺度；图像 builder 的第四尺度原本会被 scalp 丢弃，因此适配器移除该未使用尺度并取消 scalp。忽略视频交互/传播 neck，随后 `strict=True` 加载所有实际使用参数；结构不匹配会报错。未修改外部 SAM 源码。
|
||||
|
||||
保留 FP32 参数和复数位置编码缓冲区，仅在推理上下文开启 BF16 autocast。不能直接把整个模型转换为 BF16，这会破坏复数缓冲区及上游要求 FP32 的解码器路径。
|
||||
|
||||
## 提示、选择与质量检查
|
||||
|
||||
1. 在白色背景上合成参考图，保持原尺寸；GPU 只编码一次图像。
|
||||
2. 每个部件先用语义文本推理，再加入归一化粗框推理，保留两组候选。
|
||||
3. 根据模型分数、粗框内比例、越出角色 Alpha 的比例及候选点命中情况排序。拒绝前景不足 4 像素或框内比例低于 10% 的结果。
|
||||
4. 最终 mask 与原图 Alpha 前景取交集，从原图提取 RGB 和 Alpha，不重绘、不缩放。原始选中 mask 另存用于检查越界。
|
||||
5. 记录漏分、重叠、碎片化、候选点未覆盖，以及不同部件 mask IoU ≥ 0.8 的疑似重复。
|
||||
|
||||
当前使用的 SAM 图像概念接口接受文本与框，不接受点。第一阶段的 Alpha 候选点仅用于 mask 排序；模型原始点仍只用于第一阶段诊断。这次 mask 来自真实 GPU 模型推理。
|
||||
|
||||
粗框、候选点及模型分数都不能保证部件语义正确。因此所有结果保持 `quality_status=needs_review`、`semantic_verified=false`。重叠 mask 保留，未强行分配成互斥部件，也未补全遮挡纹理。人工复核通过后才能继续补全或导出。
|
||||
|
||||
## 输出
|
||||
|
||||
| 文件 | 内容 |
|
||||
| --- | --- |
|
||||
| `annotated.png` | 原尺寸角色图，各候选部件用不同半透明颜色叠加，透明背景显示为白色 |
|
||||
| `parts-sheet.png` | 白底拆解图，部件保持原尺寸、分开排布，并标注 ID 和颜色 |
|
||||
| `parts/<id>.png` | 紧裁剪的 RGBA 部件，保留原图像素及 Alpha |
|
||||
| `masks/<id>.png` | 原图尺寸黑白 mask，已裁掉 Alpha 背景 |
|
||||
| `candidates/<id>-raw.png` | 选中的 SAM 原始 mask，裁掉背景前 |
|
||||
| `unassigned-mask.png` | 白色表示没有被任何部件 mask 覆盖的角色像素 |
|
||||
| `overlap-mask.png` | 白色表示被多个部件 mask 覆盖的像素 |
|
||||
| `segmentation.json` | 裁剪坐标、拆解图坐标、颜色、质量指标、运行环境和模型加载信息 |
|
||||
| `candidate-diagnostics.json` | 两种提示的候选分数、排序、拒绝原因及最终选择 |
|
||||
| `report.html` | 两张预览、部件链接、告警、漏分及重复检查 |
|
||||
| `worker.log` / `worker-error.json` | 子进程日志 / 失败诊断（发生错误时） |
|
||||
| `manifest.json` | 状态、输入/代码/权重/产物哈希，完成时记录子进程已退出 |
|
||||
|
||||
`crop_xyxy` 使用原图坐标，右下角为开区间；`sheet_xyxy` 使用白底拆解图坐标。覆盖率是像素覆盖比例，不是准确率。
|
||||
|
||||
## 代码与验证记录
|
||||
|
||||
- [stages/segmentation-stage.mjs](../stages/segmentation-stage.mjs)：输入校验、Conda 子进程、哈希、缓存、日志和失败处理。
|
||||
- [sam_runtime.py](../sam_runtime.py)：环境探测和本地权重严格加载。
|
||||
- [sam_segment.py](../sam_segment.py)：GPU 推理、两组候选选择和诊断。
|
||||
- [segmentation_outputs.py](../segmentation_outputs.py)：mask QA、原像素拆件、两张预览和报告。
|
||||
- [tests/test_segmentation.py](../tests/test_segmentation.py)：Alpha 越界、候选选择、提示完整性、像素保真和重叠/重复检查。
|
||||
|
||||
2026-09-18 实测 `output/sombrero-grounding-v2` → `output/sombrero-sam31-v3`：28 个部件返回 mask，严格加载 1130 个检测参数/缓冲张量；推理及渲染约 11.64 秒（不含环境探测和加载权重），PyTorch 峰值已分配显存约 4.97 GiB，结束后无 SAM 子进程。
|
||||
|
||||
这组结果覆盖原图前景约 67.18%，仍有 64286 像素未覆盖、30569 像素重叠。`shin_left` 和 `foot_left` 的 mask IoU 约 0.843，实际都分到了鞋子；躯干及其他肢体也存在碎片或错分。文本候选对照未解决这些问题。这说明推理和输出链路已跑通，但当前旧语义计划与提示策略尚不能产出合格的最终拆件；需要修正部件粗框并继续验证提示及 mask 语义。
|
||||
|
||||
同日 `conda run --no-capture-output -n ai-spine npm test` 的 17 项测试全部通过（包括 Python 分割回归用例）；真实产物 `--resume` 哈希校验复用通过。测试通过仅验证程序行为，不表示上述分割质量合格。
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
本轮实现结构理解、图片编辑适配器、结构 JSON 校验、Alpha 几何校验、候选种子、诊断预览和实验缓存。几何层需要 Python 与 requirements.txt，不运行旧的“拆件图 → SIFT 定位”流程。
|
||||
|
||||
结构理解当前停在 `understand`；SAM 分割、掩码修正、隐藏部件合并、Spine 骨骼导出是后续阶段。图片编辑测试保存的是原始候选，不是验收通过的最终资产。
|
||||
结构理解支持停在 `understand`，也可通过 `--stop-after segment` 接续本地 SAM 分割，或使用 `--stage segmentation --input-dir` 独立执行第二阶段。掩码语义修正、隐藏部件合并、Spine 骨骼导出尚未完成。图片编辑测试保存的是原始候选，不是验收通过的最终资产。
|
||||
|
||||
## 环境与模型
|
||||
|
||||
@@ -13,7 +13,7 @@
|
||||
```dotenv
|
||||
SPINE_UNDERSTANDING_MODEL=gpt-6-astra
|
||||
SPINE_COMPLETION_MODEL=qwen-image-3.0
|
||||
SAM_MODEL_DIR=F:/Study/sam3.1
|
||||
SAM_MODEL_DIR=F:/Study/SAM3
|
||||
```
|
||||
|
||||
不配置这三项时，`configs/astra-qwen.json` 已提供相同默认值。不会自动改写用户的 `.env`。
|
||||
@@ -153,13 +153,13 @@ node .\tools\model-check.mjs --stage completion --image .\local-part.png --promp
|
||||
|
||||
同一命令增加 `--resume`，仅当上次完整结束，且模型、协议、参数、提示词、Schema、源图、适配器代码和所有产物校验一致时复用。模型变化时要求新的输出目录，防止覆盖对照组。失败任务也保留原目录，使用新目录重跑。
|
||||
|
||||
本轮实现单阶段缓存；跨 SAM/补全/导出的依赖缓存留到后续阶段接入。对比补全模型时，固定相同局部画布、任务文本与掩码。
|
||||
理解、SAM 和补全探针各自记录独立缓存。SAM 额外校验输入计划、候选提示、本地源码和权重哈希。对比补全模型时，固定相同局部画布、任务文本与掩码。
|
||||
|
||||
显式 429/5xx 最多按配置重试，超时或网络断开不自动重试，因为服务端可能已经执行任务。完整请求不保证模型内容可复现；seed 只在提供方确认支持时放入 params，usage 不存在时记录 null，不估造价格或 token 数。
|
||||
|
||||
## SAM 路径
|
||||
|
||||
默认 `F:/Study/sam3.1`，检查 `sam3.1_multiplex.pt` 是否存在。该目录 README 明确说明它仅提供 checkpoint，不是可直接调用的推理服务。本轮仅记录路径和 checkpoint 存在性，状态为 `not_integrated`，不加载权重、不占用 GPU。
|
||||
默认根目录 `F:/Study/SAM3`：源码位于 `sam3/`，权重位于 `sam3.1checkpoint/sam3.1_multiplex.pt`。可通过 `--sam-dir`、`--sam-code-dir`、`--sam-checkpoint` 覆盖。独立分割无需中转站凭据，在 `ai-spine` Conda 临时子进程加载 GPU 模型，结束即退出。命令、权重适配及当前质量限制见[第二阶段文档](2-segmentation-stage.md)。
|
||||
|
||||
## 测试
|
||||
|
||||
|
||||
+8
-4
@@ -84,9 +84,13 @@ export function publicStage(stage) {
|
||||
};
|
||||
}
|
||||
|
||||
export function samInfo(loaded) {
|
||||
export function samInfo(loaded, overrides = {}) {
|
||||
const stage = loaded.config.stages.segmentation || {};
|
||||
const modelDir = loaded.env[stage.model_dir_env] || stage.default_model_dir || '';
|
||||
const checkpoint = modelDir ? path.join(modelDir, 'sam3.1_multiplex.pt') : '';
|
||||
return {model_dir: modelDir, checkpoint, checkpoint_exists: Boolean(checkpoint && existsSync(checkpoint)), runtime_status: 'not_integrated'};
|
||||
const modelDir = overrides.samDir || loaded.env[stage.model_dir_env] || stage.default_model_dir || '';
|
||||
const checkpoint = overrides.samCheckpoint || (modelDir ? [path.join(modelDir, 'sam3.1checkpoint', 'sam3.1_multiplex.pt'), path.join(modelDir, 'sam3.1_multiplex.pt')].find(existsSync) || path.join(modelDir, 'sam3.1checkpoint', 'sam3.1_multiplex.pt') : '');
|
||||
const codeDir = overrides.samCodeDir || loaded.env.SAM_CODE_DIR || (modelDir ? path.join(modelDir, 'sam3') : '');
|
||||
return {model_dir: modelDir, code_dir: codeDir, checkpoint,
|
||||
checkpoint_exists: Boolean(checkpoint && existsSync(checkpoint)),
|
||||
code_exists: Boolean(codeDir && existsSync(path.join(codeDir, 'sam3', 'model_builder.py'))),
|
||||
runtime_status: 'local_subprocess_not_probed'};
|
||||
}
|
||||
|
||||
+22
-3
@@ -1,3 +1,22 @@
|
||||
Pillow>=10.0
|
||||
numpy>=1.26
|
||||
opencv-python>=4.9
|
||||
Pillow>=10.1
|
||||
numpy==1.26.4
|
||||
opencv-python==4.11.0.86
|
||||
# SAM runtime. Install the CUDA torch/torchvision wheels first as documented.
|
||||
torch==2.7.1
|
||||
torchvision==0.22.1
|
||||
setuptools==80.9.0
|
||||
wheel
|
||||
timm>=1.0.17
|
||||
tqdm
|
||||
ftfy==6.1.1
|
||||
regex
|
||||
iopath>=0.1.10
|
||||
typing_extensions
|
||||
huggingface_hub
|
||||
einops
|
||||
psutil
|
||||
scipy
|
||||
pycocotools
|
||||
triton-windows>=3.3,<3.4; platform_system == "Windows"
|
||||
triton>=3.3,<3.4; platform_system == "Linux"
|
||||
# SAM source is provided separately through --sam-code-dir / SAM_CODE_DIR.
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
"""Local SAM image detector from a SAM 3.1 multiplex checkpoint.
|
||||
|
||||
No downloads, service, video tracker, or silently uninitialized parameters.
|
||||
"""
|
||||
import gc
|
||||
import hashlib
|
||||
import importlib.metadata
|
||||
import os
|
||||
from pathlib import Path
|
||||
import sys
|
||||
|
||||
|
||||
def check_environment(code_dir):
|
||||
if Path(sys.prefix).name != 'ai-spine' or not (Path(sys.prefix) / 'conda-meta').is_dir():
|
||||
raise RuntimeError('Run inside Conda ai-spine: conda run -n ai-spine ...')
|
||||
code = Path(code_dir).resolve()
|
||||
if not (code / 'sam3' / 'model_builder.py').is_file():
|
||||
raise FileNotFoundError(f'SAM source checkout missing: {code}')
|
||||
sys.path.insert(0, str(code))
|
||||
os.environ['HF_HUB_OFFLINE'] = '1'
|
||||
os.environ['TRANSFORMERS_OFFLINE'] = '1'
|
||||
import torch
|
||||
import sam3
|
||||
from sam3.model_builder import build_sam3_image_model # noqa: F401
|
||||
from sam3.model.sam3_image_processor import Sam3Processor # noqa: F401
|
||||
if Path(sam3.__file__).resolve().parent != code / 'sam3':
|
||||
raise RuntimeError('Imported SAM package does not match --sam-code-dir')
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError('CUDA GPU unavailable; this stage requires GPU inference')
|
||||
source_hash = hashlib.sha256()
|
||||
for file in sorted((code / 'sam3').rglob('*')):
|
||||
if file.is_file() and (file.suffix == '.py' or file.name.endswith('.txt.gz')):
|
||||
source_hash.update(file.relative_to(code).as_posix().encode())
|
||||
source_hash.update(file.read_bytes())
|
||||
versions = {}
|
||||
for name in ['torch', 'torchvision', 'numpy', 'Pillow', 'opencv-python', 'timm', 'triton-windows']:
|
||||
try:
|
||||
versions[name] = importlib.metadata.version(name)
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
versions[name] = None
|
||||
return {'python': sys.executable, 'sam_source': str(code), 'sam_source_sha256': source_hash.hexdigest(),
|
||||
'versions': versions, 'cuda': torch.version.cuda, 'gpu': torch.cuda.get_device_name(0),
|
||||
'gpu_total_bytes': torch.cuda.get_device_properties(0).total_memory,
|
||||
'bf16_supported': torch.cuda.is_bf16_supported(), 'network_requested': False}
|
||||
|
||||
|
||||
def load_detector(checkpoint, precision='bf16'):
|
||||
import torch
|
||||
from sam3.model_builder import build_sam3_image_model
|
||||
model = build_sam3_image_model(device='cpu', load_from_HF=False,
|
||||
enable_inst_interactivity=False, compile=False)
|
||||
weights = torch.load(checkpoint, map_location='cpu', weights_only=True, mmap=True)
|
||||
weights = weights.get('model', weights)
|
||||
detector = {key.removeprefix('detector.'): value for key, value in weights.items() if key.startswith('detector.')}
|
||||
if not detector:
|
||||
raise RuntimeError('Checkpoint has no detector.* weights; refusing an uninitialized model')
|
||||
# The image builder uses four scales then scalp=1 discards the last one.
|
||||
# SAM 3.1 stores only the three used detection scales. Remove the unused
|
||||
# fourth scale and scalp, preserving exactly the same three feature maps.
|
||||
multiplex = any(k.startswith('backbone.vision_backbone.interactive_convs.') for k in detector)
|
||||
if multiplex:
|
||||
neck = model.backbone.vision_backbone
|
||||
if len(neck.convs) != 4 or model.backbone.scalp != 1:
|
||||
raise RuntimeError('Unsupported SAM image backbone layout; review adapter for this source version')
|
||||
neck.convs = torch.nn.ModuleList(list(neck.convs)[:3])
|
||||
neck.scale_factors = list(neck.scale_factors)[:3]
|
||||
model.backbone.scalp = 0
|
||||
ignored_prefixes = ('backbone.vision_backbone.interactive_convs.',
|
||||
'backbone.vision_backbone.propagation_convs.') if multiplex else ()
|
||||
selected = {k: v for k, v in detector.items() if not k.startswith(ignored_prefixes)}
|
||||
# strict=True checks every active parameter and every tensor shape.
|
||||
model.load_state_dict(selected, strict=True)
|
||||
info = {'checkpoint_kind': 'sam3.1_multiplex_detector' if multiplex else 'sam3_detector',
|
||||
'strict_load': True, 'loaded_tensors': len(selected),
|
||||
'ignored_video_neck_tensors': len(detector) - len(selected),
|
||||
'instance_interactivity': False, 'compile': False, 'precision': precision}
|
||||
del selected, detector, weights
|
||||
gc.collect()
|
||||
if precision == 'bf16' and not torch.cuda.is_bf16_supported():
|
||||
raise RuntimeError('BF16 unsupported on this GPU; use --sam-precision fp32')
|
||||
# Keep FP32 parameters and complex RoPE buffers intact. The upstream
|
||||
# decoder explicitly disables autocast for its FP32 FFN; casting the whole
|
||||
# model to BF16 would both break it and discard complex imaginary values.
|
||||
model = model.to(device='cuda').eval()
|
||||
return model, info
|
||||
+177
@@ -0,0 +1,177 @@
|
||||
#!/usr/bin/env python3
|
||||
"""One-shot GPU worker: semantic ROI prompts -> real SAM masks -> previews."""
|
||||
import argparse
|
||||
from contextlib import nullcontext
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import time
|
||||
import traceback
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from sam_runtime import check_environment, load_detector
|
||||
from segmentation_outputs import assess_candidate, render_outputs, write_report
|
||||
|
||||
|
||||
def choose_candidate(masks, scores, alpha, box, points):
|
||||
"""Rank real model candidates; ROI and seeds are hints, not ground truth."""
|
||||
ranked = []
|
||||
for index, (mask, score) in enumerate(zip(masks, scores)):
|
||||
qa = assess_candidate(mask, alpha, box, points)
|
||||
reasons = []
|
||||
if qa['area_px'] < 4:
|
||||
reasons.append('EMPTY_OR_TINY_FOREGROUND')
|
||||
if qa['inside_roi_fraction'] < .1:
|
||||
reasons.append('OUTSIDE_TARGET_ROI')
|
||||
rank = float(score) * (.25 + .75 * qa['inside_roi_fraction']) * (1 - qa['outside_alpha_ratio'])
|
||||
rank *= .8 + .2 * qa['seed_hit_fraction']
|
||||
ranked.append({'index': index, 'model_score': float(score), 'rank_score': rank,
|
||||
'qa': qa, 'rejected': reasons})
|
||||
eligible = [c for c in ranked if not c['rejected']]
|
||||
return max(eligible, key=lambda c: c['rank_score']) if eligible else None, ranked
|
||||
|
||||
|
||||
def validate_prompts(plan, prompts, image):
|
||||
w, h = image.size
|
||||
if (prompts.get('width'), prompts.get('height')) != (w, h):
|
||||
raise ValueError('Prompt canvas does not match source image')
|
||||
parts = {p['id']: p for p in plan['parts']}
|
||||
entries = prompts['parts']
|
||||
if len(entries) != len(parts) or {p['id'] for p in entries} != set(parts):
|
||||
raise ValueError('Prompt IDs must match the semantic plan exactly')
|
||||
alpha = np.array(image)[:, :, 3]
|
||||
for entry in entries:
|
||||
part = parts[entry['id']]
|
||||
if entry['status'] not in ['candidate', 'hidden', 'blocked']:
|
||||
raise ValueError('Unknown prompt status')
|
||||
if entry['status'] != 'candidate':
|
||||
continue
|
||||
if part['visibility'] == 'hidden' or entry.get('seed_source') != 'alpha_roi_distance_transform':
|
||||
raise ValueError('Only visible, program-generated candidate prompts are supported')
|
||||
box = entry['bbox_px_xyxy']
|
||||
if len(box) != 4 or any(type(v) is not int for v in box) or not (0 <= box[0] < box[2] <= w and 0 <= box[1] < box[3] <= h):
|
||||
raise ValueError('Invalid pixel ROI')
|
||||
if not entry['positive_points_px']:
|
||||
raise ValueError('Candidate prompt has no foreground seeds')
|
||||
for point in entry['positive_points_px']:
|
||||
if len(point) != 2 or any(type(v) is not int for v in point):
|
||||
raise ValueError('Invalid pixel seed')
|
||||
x, y = point
|
||||
if not (box[0] <= x < box[2] and box[1] <= y < box[3]) or alpha[y, x] == 0:
|
||||
raise ValueError('Candidate seed is outside foreground/ROI')
|
||||
return parts
|
||||
|
||||
|
||||
def run(args, runtime):
|
||||
import torch
|
||||
from sam3.model.sam3_image_processor import Sam3Processor
|
||||
out = Path(args.output)
|
||||
plan = json.loads((out / 'semantic-parts.json').read_text(encoding='utf-8'))
|
||||
prompts = json.loads((out / 'segmentation-prompts.json').read_text(encoding='utf-8'))
|
||||
image = Image.open(out / 'source.png').convert('RGBA')
|
||||
parts_by_id = validate_prompts(plan, prompts, image)
|
||||
alpha = np.array(image)[:, :, 3]
|
||||
print('Loading local SAM detector weights...', flush=True)
|
||||
model, loading = load_detector(args.checkpoint, args.precision)
|
||||
processor = Sam3Processor(model, device='cuda', confidence_threshold=args.confidence)
|
||||
rgb = Image.alpha_composite(Image.new('RGBA', image.size, 'white'), image).convert('RGB')
|
||||
candidates_dir = out / 'candidates'
|
||||
candidates_dir.mkdir(exist_ok=True)
|
||||
selected_parts, skipped, diagnostics = [], [], []
|
||||
started = time.monotonic()
|
||||
context = torch.autocast('cuda', dtype=torch.bfloat16) if args.precision == 'bf16' else nullcontext()
|
||||
with torch.inference_mode(), context:
|
||||
print('Encoding reference image on GPU...', flush=True)
|
||||
state = processor.set_image(rgb)
|
||||
for number, entry in enumerate(prompts['parts'], 1):
|
||||
name = entry['id']
|
||||
if entry['status'] != 'candidate':
|
||||
skipped.append({'id': name, 'reason': entry['status']})
|
||||
continue
|
||||
part = parts_by_id[name]
|
||||
print(f'[{number}/{len(prompts["parts"])}] {name}: text / text + coarse ROI', flush=True)
|
||||
processor.reset_all_prompts(state)
|
||||
# The official image concept API accepts text and boxes, not points.
|
||||
# Alpha seeds are used only to rank the returned mask candidates.
|
||||
processor.set_text_prompt(part['semantic'], state)
|
||||
x0, y0, x1, y1 = entry['bbox_px_xyxy']
|
||||
w, h = image.size
|
||||
# Preserve text-only candidates: an inaccurate coarse box can
|
||||
# steer the concept detector toward an unrelated fragment.
|
||||
text_masks = state['masks'].detach().cpu().numpy().reshape(-1, h, w).copy()
|
||||
text_scores = state['scores'].detach().float().cpu().numpy().reshape(-1).copy()
|
||||
processor.add_geometric_prompt([(x0+x1)/(2*w), (y0+y1)/(2*h), (x1-x0)/w, (y1-y0)/h], True, state)
|
||||
box_masks = state['masks'].detach().cpu().numpy().reshape(-1, h, w)
|
||||
box_scores = state['scores'].detach().float().cpu().numpy().reshape(-1)
|
||||
masks = np.concatenate([text_masks, box_masks])
|
||||
scores = np.concatenate([text_scores, box_scores])
|
||||
best, ranked = choose_candidate(masks, scores, alpha, entry['bbox_px_xyxy'], entry['positive_points_px'])
|
||||
for candidate in ranked:
|
||||
candidate['prompt_mode'] = 'text_only' if candidate['index'] < len(text_masks) else 'text_and_box'
|
||||
diagnostic = {'id': name, 'text_prompt': part['semantic'], 'candidates': ranked,
|
||||
'selected_index': best['index'] if best else None}
|
||||
diagnostics.append(diagnostic)
|
||||
if best is None:
|
||||
skipped.append({'id': name, 'reason': 'NO_ACCEPTABLE_SAM_MASK', 'candidate_count': len(ranked)})
|
||||
continue
|
||||
mask = masks[best['index']]
|
||||
Image.fromarray(mask.astype(np.uint8)*255).save(candidates_dir / f'{name}-raw.png')
|
||||
warnings = ['SEMANTIC_REVIEW_REQUIRED']
|
||||
if best['qa']['outside_alpha_ratio'] > .1:
|
||||
warnings.append('MASK_OUTSIDE_CHARACTER')
|
||||
if best['qa']['inside_roi_fraction'] < .5:
|
||||
warnings.append('MASK_EXTENDS_BEYOND_COARSE_ROI')
|
||||
if best['qa']['seed_hit_fraction'] < 1:
|
||||
warnings.append('CANDIDATE_SEED_NOT_COVERED')
|
||||
if best['qa']['components_at_least_4px'] > 10:
|
||||
warnings.append('FRAGMENTED_MASK')
|
||||
selected_parts.append({'id': name, 'semantic': part['semantic'], 'rig_role': part['rig_role'],
|
||||
'mask': mask, 'model_score': best['model_score'], 'qa': best['qa'],
|
||||
'prompt_mode': best['prompt_mode'],
|
||||
'warnings': warnings, 'semantic_verified': False})
|
||||
if not selected_parts:
|
||||
(out / 'candidate-diagnostics.json').write_text(json.dumps(diagnostics, indent=2), encoding='utf-8')
|
||||
raise RuntimeError('SAM produced no acceptable masks; see candidate-diagnostics.json')
|
||||
result = render_outputs(image, selected_parts, out)
|
||||
result.update(schema_version=1, quality_status='needs_review', skipped=skipped, runtime=runtime,
|
||||
model_loading=loading, elapsed_seconds=time.monotonic()-started,
|
||||
peak_cuda_memory_bytes=torch.cuda.max_memory_allocated(),
|
||||
prompt_strategy='rank text_only and text_and_box candidates; alpha seeds rank masks only')
|
||||
(out / 'segmentation.json').write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding='utf-8')
|
||||
(out / 'candidate-diagnostics.json').write_text(json.dumps(diagnostics, ensure_ascii=False, indent=2), encoding='utf-8')
|
||||
write_report(out, result)
|
||||
print(f'SAM finished: {len(selected_parts)} masks, {len(skipped)} skipped. GPU worker exiting.', flush=True)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--code-dir', required=True)
|
||||
parser.add_argument('--checkpoint', required=True)
|
||||
parser.add_argument('--output')
|
||||
parser.add_argument('--probe', action='store_true')
|
||||
parser.add_argument('--confidence', type=float, default=.2)
|
||||
parser.add_argument('--precision', choices=['bf16', 'fp32'], default='bf16')
|
||||
args = parser.parse_args()
|
||||
if not Path(args.checkpoint).is_file():
|
||||
parser.error('Local checkpoint does not exist')
|
||||
if not 0 < args.confidence < 1:
|
||||
parser.error('--confidence must be between 0 and 1')
|
||||
try:
|
||||
runtime = check_environment(args.code_dir)
|
||||
if args.probe:
|
||||
print(json.dumps(runtime))
|
||||
return
|
||||
if not args.output:
|
||||
parser.error('--output is required')
|
||||
run(args, runtime)
|
||||
except Exception as error:
|
||||
if args.output and Path(args.output).is_dir():
|
||||
(Path(args.output) / 'worker-error.json').write_text(json.dumps({'type': type(error).__name__, 'message': str(error), 'traceback': traceback.format_exc()}, indent=2), encoding='utf-8')
|
||||
traceback.print_exc()
|
||||
raise SystemExit(1)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,130 @@
|
||||
"""Measured mask QA and previews made only from the original character pixels."""
|
||||
import colorsys
|
||||
import html
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import cv2
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
|
||||
|
||||
def assess_candidate(mask, alpha, bbox, points):
|
||||
mask = np.asarray(mask, dtype=bool)
|
||||
if mask.shape != alpha.shape:
|
||||
raise ValueError('Mask dimensions do not match the source')
|
||||
clipped = mask & (alpha > 0)
|
||||
_, _, stats, _ = cv2.connectedComponentsWithStats(clipped.astype(np.uint8), 8)
|
||||
area, foreground_area = int(mask.sum()), int(clipped.sum())
|
||||
x0, y0, x1, y1 = bbox
|
||||
hit = sum(bool(clipped[y, x]) for x, y in points)
|
||||
return {
|
||||
'raw_area_px': area, 'area_px': foreground_area,
|
||||
'outside_alpha_ratio': (area - foreground_area) / max(1, area),
|
||||
'inside_roi_fraction': int(clipped[y0:y1, x0:x1].sum()) / max(1, foreground_area),
|
||||
'seed_hit_fraction': hit / max(1, len(points)),
|
||||
'components_at_least_4px': int((stats[1:, cv2.CC_STAT_AREA] >= 4).sum()),
|
||||
}
|
||||
|
||||
|
||||
def render_outputs(image, parts, out_dir):
|
||||
out = Path(out_dir)
|
||||
(out / 'parts').mkdir(parents=True, exist_ok=True)
|
||||
(out / 'masks').mkdir(parents=True, exist_ok=True)
|
||||
rgba = np.array(image.convert('RGBA'))
|
||||
h, w = rgba.shape[:2]
|
||||
coverage = np.zeros((h, w), np.uint16)
|
||||
white = Image.new('RGBA', (w, h), 'white')
|
||||
overlay = Image.alpha_composite(white, image.convert('RGBA'))
|
||||
records, crops, clipped_masks = [], [], []
|
||||
seen = set()
|
||||
for index, part in enumerate(parts):
|
||||
name = part['id']
|
||||
if not re.fullmatch('[a-z][a-z0-9_-]{0,63}', name) or name in seen:
|
||||
raise ValueError(f'Invalid or duplicate part ID: {name}')
|
||||
seen.add(name)
|
||||
mask = np.asarray(part['mask'], dtype=bool)
|
||||
if mask.shape != (h, w):
|
||||
raise ValueError(f'{name}: mask dimensions mismatch')
|
||||
mask = mask & (rgba[:, :, 3] > 0)
|
||||
clipped_masks.append(mask)
|
||||
ys, xs = np.nonzero(mask)
|
||||
if not len(xs):
|
||||
raise ValueError(f'{name}: empty foreground mask')
|
||||
coverage += mask
|
||||
bbox = [int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1]
|
||||
color = [round(v * 255) for v in colorsys.hsv_to_rgb((index * .61803398875) % 1, .75, .95)]
|
||||
cut = rgba.copy()
|
||||
cut[~mask] = 0
|
||||
crop = Image.fromarray(cut).crop(bbox)
|
||||
crop.save(out / 'parts' / f'{name}.png')
|
||||
Image.fromarray(mask.astype(np.uint8) * 255).save(out / 'masks' / f'{name}.png')
|
||||
tint = np.zeros_like(rgba)
|
||||
tint[mask, :3] = color
|
||||
tint[mask, 3] = 100
|
||||
overlay = Image.alpha_composite(overlay, Image.fromarray(tint))
|
||||
records.append({k: v for k, v in part.items() if k != 'mask'} | {
|
||||
'color_rgb': color, 'crop_xyxy': bbox, 'area_px': int(mask.sum()),
|
||||
'image': f'parts/{name}.png', 'mask': f'masks/{name}.png',
|
||||
})
|
||||
crops.append(crop)
|
||||
near_duplicates = []
|
||||
for i, a in enumerate(clipped_masks):
|
||||
for j in range(i + 1, len(clipped_masks)):
|
||||
b = clipped_masks[j]
|
||||
iou = int((a & b).sum()) / max(1, int((a | b).sum()))
|
||||
if iou >= .8:
|
||||
near_duplicates.append({'part_a': records[i]['id'], 'part_b': records[j]['id'], 'iou': iou})
|
||||
for index in [i, j]:
|
||||
warnings = records[index].setdefault('warnings', []).copy()
|
||||
if 'NEAR_DUPLICATE_MASK' not in warnings:
|
||||
warnings.append('NEAR_DUPLICATE_MASK')
|
||||
records[index]['warnings'] = warnings
|
||||
# Keep native scale, separated by padding and a label band. No rotation or
|
||||
# resampling of source pixels is used in the exploded sheet.
|
||||
pad, label = 24, 26
|
||||
cell_widths = [max(c.width, len(r['id']) * 8 + 8) for c, r in zip(crops, records)]
|
||||
sheet_width = max(640, max(cell_widths, default=1) + pad * 2,
|
||||
int(math.sqrt(sum((cw + pad) * (c.height + label + pad) for cw, c in zip(cell_widths, crops)) * 1.5)))
|
||||
x, y, row_height = pad, pad, 0
|
||||
for record, crop, cw in zip(records, crops, cell_widths):
|
||||
if x + cw + pad > sheet_width and x != pad:
|
||||
x, y, row_height = pad, y + row_height + pad, 0
|
||||
record['sheet_xyxy'] = [x, y + label, x + crop.width, y + label + crop.height]
|
||||
x += cw + pad
|
||||
row_height = max(row_height, label + crop.height)
|
||||
sheet = Image.new('RGB', (sheet_width, max(100, y + row_height + pad)), 'white')
|
||||
draw = ImageDraw.Draw(sheet)
|
||||
font = ImageFont.load_default(size=14)
|
||||
for record, crop in zip(records, crops):
|
||||
x0, y0, _, _ = record['sheet_xyxy']
|
||||
draw.rectangle((x0, y0 - label + 4, x0 + 10, y0 - label + 14), fill=tuple(record['color_rgb']))
|
||||
draw.text((x0 + 16, y0 - label + 1), record['id'], fill='black', font=font)
|
||||
sheet.paste(crop, (x0, y0), crop)
|
||||
overlay.convert('RGB').save(out / 'annotated.png')
|
||||
sheet.save(out / 'parts-sheet.png')
|
||||
Image.fromarray(((rgba[:, :, 3] > 0) & (coverage == 0)).astype(np.uint8) * 255).save(out / 'unassigned-mask.png')
|
||||
Image.fromarray((coverage > 1).astype(np.uint8) * 255).save(out / 'overlap-mask.png')
|
||||
total = int((rgba[:, :, 3] > 0).sum())
|
||||
return {'parts': records, 'near_duplicate_masks': near_duplicates, 'coverage': {
|
||||
'foreground_pixels': total, 'covered_pixels': int((coverage > 0).sum()),
|
||||
'coverage_ratio': int((coverage > 0).sum()) / max(1, total),
|
||||
'overlap_pixels': int((coverage > 1).sum()),
|
||||
'unassigned_pixels': int(((rgba[:, :, 3] > 0) & (coverage == 0)).sum()),
|
||||
}, 'outputs': {'annotated': 'annotated.png', 'parts_sheet': 'parts-sheet.png'}}
|
||||
|
||||
|
||||
def write_report(out_dir, result):
|
||||
out = Path(out_dir)
|
||||
rows = ''.join(f'<tr><td><span style="background:rgb({",".join(map(str, p["color_rgb"]))})"> </span> <a href="{p["image"]}">{html.escape(p["id"])}</a></td><td>{html.escape(p["semantic"])}</td><td>{p["model_score"]:.3f}</td><td>{html.escape(", ".join(p.get("warnings", [])))}</td></tr>' for p in result['parts'])
|
||||
body = f'''<!doctype html><html lang="zh-CN"><meta charset="utf-8"><title>SAM 部件分割</title>
|
||||
<style>body{{font:16px system-ui;max-width:1400px;margin:24px auto;padding:16px}}img{{max-width:100%;height:auto}}td,th{{border:1px solid #ccc;padding:8px}}table{{border-collapse:collapse}}pre{{white-space:pre-wrap}}</style>
|
||||
<h1>SAM 本地分割结果</h1><p>颜色表示候选部件 mask。白底拆解图保留原始可见像素，尚未补全遮挡部分。重叠 mask 保留并计入诊断，不代表已经获得互斥、语义正确的最终拆件。</p>
|
||||
<h2>原图半透明标注</h2><img src="annotated.png"><h2>白底拆解图</h2><img src="parts-sheet.png">
|
||||
<h2>部件</h2><table><tr><th>ID / 颜色</th><th>语义</th><th>模型分数</th><th>需复核</th></tr>{rows}</table>
|
||||
<h2>未生成的部件</h2><pre>{html.escape(json.dumps(result['skipped'], ensure_ascii=False, indent=2))}</pre>
|
||||
<h2>覆盖与重叠</h2><p>覆盖率只表示有 mask 的像素比例，不是分割准确率。模型分数也不代表部件语义已验证。</p><p><a href="unassigned-mask.png">漏分区域（白色）</a> · <a href="overlap-mask.png">重叠区域（白色）</a></p><pre>{html.escape(json.dumps(result['coverage'], indent=2))}</pre>
|
||||
<h2>疑似重复分割</h2><p>以下不同部件的 mask IoU ≥ 0.8，需要检查是否分到了同一物体。</p><pre>{html.escape(json.dumps(result['near_duplicate_masks'], indent=2))}</pre></html>'''
|
||||
(out / 'report.html').write_text(body, encoding='utf-8')
|
||||
@@ -0,0 +1,101 @@
|
||||
import {spawn} from 'node:child_process';
|
||||
import {createHash} from 'node:crypto';
|
||||
import {createReadStream, createWriteStream, existsSync, readFileSync, readdirSync, writeFileSync} from 'node:fs';
|
||||
import path from 'node:path';
|
||||
import {projectRoot, samInfo} from '../lib/config.mjs';
|
||||
import {selectPython} from '../lib/python-runtime.mjs';
|
||||
import {readImage, sha256} from '../lib/images.mjs';
|
||||
import {validateParts} from '../lib/validate-parts.mjs';
|
||||
import {prepareRun, jsonFile, finishRun} from '../lib/run-artifacts.mjs';
|
||||
|
||||
const worker = path.join(projectRoot, 'sam_segment.py');
|
||||
|
||||
export function runWorker(python, args, {logFile, timeoutMs = 1200000} = {}) {
|
||||
return new Promise((resolve, reject) => {
|
||||
const child = spawn(python, ['-u', worker, ...args], {cwd: projectRoot, windowsHide: true,
|
||||
env: {...process.env, PYTHONUTF8: '1', HF_HUB_OFFLINE: '1', TRANSFORMERS_OFFLINE: '1'},
|
||||
stdio: ['ignore', 'pipe', 'pipe']});
|
||||
const log = logFile ? createWriteStream(logFile) : null;
|
||||
let stdout = '', stderr = '', timedOut = false, interrupted = false;
|
||||
const timer = setTimeout(() => { timedOut = true; child.kill(); }, timeoutMs);
|
||||
const stop = () => { interrupted = true; child.kill(); };
|
||||
process.once('SIGINT', stop);
|
||||
process.once('SIGTERM', stop);
|
||||
child.stdout.on('data', (chunk) => { stdout += chunk.toString(); log?.write(chunk); if (log) process.stdout.write(chunk); });
|
||||
child.stderr.on('data', (chunk) => { stderr = (stderr + chunk.toString()).slice(-20000); log?.write(chunk); if (log) process.stderr.write(chunk); });
|
||||
const cleanup = () => { clearTimeout(timer); process.removeListener('SIGINT', stop); process.removeListener('SIGTERM', stop); };
|
||||
child.once('error', (error) => { cleanup(); log?.end(); reject(error); });
|
||||
child.once('close', (code) => {
|
||||
cleanup();
|
||||
const done = () => {
|
||||
if (code !== 0 || timedOut || interrupted) reject(new Error(timedOut ? 'SAM 子进程超时，已终止' : interrupted ? 'SAM 子进程已取消' : `SAM 子进程失败 (${code})：${stderr.slice(-3500)}`));
|
||||
else resolve(stdout);
|
||||
};
|
||||
if (log) log.end(done); else done();
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
async function fileHash(file) {
|
||||
const hash = createHash('sha256');
|
||||
for await (const chunk of createReadStream(file)) hash.update(chunk);
|
||||
return hash.digest('hex');
|
||||
}
|
||||
|
||||
function filesIn(dir, prefix = '') {
|
||||
return readdirSync(dir, {withFileTypes: true}).flatMap((entry) => entry.isDirectory()
|
||||
? filesIn(path.join(dir, entry.name), `${prefix}${entry.name}/`) : [`${prefix}${entry.name}`]);
|
||||
}
|
||||
|
||||
export async function runSegmentation(options) {
|
||||
const {loaded, inputDir, outDir, resume = false, dryRun = false} = options;
|
||||
if (!inputDir) throw new Error('SAM 分割需要 --input-dir 指向已完成的结构理解目录');
|
||||
if (!outDir && !dryRun) throw new Error('SAM 分割需要独立的 --out-dir');
|
||||
const sourceDir = path.resolve(inputDir);
|
||||
const image = readImage(path.join(sourceDir, 'source.png'), {pngOnly: true});
|
||||
const planBytes = readFileSync(path.join(sourceDir, 'semantic-parts.json'));
|
||||
const promptBytes = readFileSync(path.join(sourceDir, 'segmentation-prompts.json'));
|
||||
const plan = JSON.parse(planBytes);
|
||||
const validation = validateParts(plan);
|
||||
if (!validation.valid || plan.schema_version !== 2) throw new Error(`需要有效的 v2 计划；请先离线迁移/校验：${validation.errors.join(', ')}`);
|
||||
const prompts = JSON.parse(promptBytes);
|
||||
if (prompts.width !== image.width || prompts.height !== image.height) throw new Error('候选点画布尺寸与 source.png 不一致');
|
||||
const sam = samInfo(loaded, options);
|
||||
if (!sam.checkpoint_exists || !sam.code_exists) throw new Error(`SAM 代码或权重不存在：${JSON.stringify(sam)}`);
|
||||
const python = selectPython(options.python);
|
||||
const confidence = Number(options.samConfidence ?? .2);
|
||||
const precision = options.samPrecision || 'bf16';
|
||||
if (!(confidence > 0 && confidence < 1) || !['bf16', 'fp32'].includes(precision)) throw new Error('sam-confidence 必须在 0～1 之间，sam-precision 仅支持 bf16/fp32');
|
||||
const baseArgs = ['--code-dir', sam.code_dir, '--checkpoint', sam.checkpoint, '--confidence', String(confidence), '--precision', precision];
|
||||
const runtime = JSON.parse((await runWorker(python, [...baseArgs, '--probe'], {timeoutMs: 120000})).trim());
|
||||
if (dryRun) { console.log(JSON.stringify({sam, runtime, confidence, precision, network_requested: false}, null, 2)); return {dryRun: true}; }
|
||||
const codeFiles = ['sam_segment.py', 'sam_runtime.py', 'segmentation_outputs.py', 'stages/segmentation-stage.mjs', 'lib/python-runtime.mjs', 'lib/validate-parts.mjs', 'schemas/semantic-parts.schema.json'];
|
||||
const identity = {schema_version: 1, stage: 'segmentation', runtime, sam, confidence, precision,
|
||||
inputs: {source_sha256: image.sha256, plan_sha256: sha256(planBytes), prompts_sha256: sha256(promptBytes)},
|
||||
checkpoint_sha256: await fileHash(sam.checkpoint),
|
||||
code_sha256: sha256(codeFiles.map((f) => `${f}\n${readFileSync(path.join(projectRoot, f), 'utf8')}`).join('\n'))};
|
||||
const fingerprint = sha256(JSON.stringify(identity));
|
||||
const run = prepareRun(outDir, fingerprint, resume);
|
||||
if (run.reused) return {out: run.out, reused: true};
|
||||
const manifest = {...identity, fingerprint, status: 'running', quality_status: 'pending', started_at: new Date().toISOString()};
|
||||
jsonFile(run.manifestPath, manifest);
|
||||
writeFileSync(path.join(run.out, 'source.png'), image.bytes);
|
||||
writeFileSync(path.join(run.out, 'semantic-parts.json'), planBytes);
|
||||
writeFileSync(path.join(run.out, 'segmentation-prompts.json'), promptBytes);
|
||||
try {
|
||||
await runWorker(python, [...baseArgs, '--output', run.out], {logFile: path.join(run.out, 'worker.log')});
|
||||
const result = JSON.parse(readFileSync(path.join(run.out, 'segmentation.json'), 'utf8'));
|
||||
manifest.quality_status = 'needs_review';
|
||||
manifest.completed_at = new Date().toISOString();
|
||||
manifest.part_count = result.parts.length;
|
||||
manifest.skipped_count = result.skipped.length;
|
||||
manifest.coverage = result.coverage;
|
||||
manifest.gpu_worker_exited = true;
|
||||
finishRun(run, manifest, filesIn(run.out).filter((name) => name !== 'manifest.json'));
|
||||
return {out: run.out, reused: false, parts: result.parts.length};
|
||||
} catch (error) {
|
||||
manifest.status = 'failed'; manifest.quality_status = 'failed'; manifest.error = error.message;
|
||||
jsonFile(run.manifestPath, manifest);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
@@ -267,7 +267,7 @@ test('pipeline new mode reaches a mock relay with geometry; model listing and un
|
||||
const list = await exec(process.execPath, ['tools/model-check.mjs', '--list-models'], {cwd: projectRoot, env});
|
||||
assert.match(list.stdout, /relay-vision-test/u);
|
||||
assert.equal(requests, 2);
|
||||
await assert.rejects(exec(process.execPath, ['pipeline.mjs', '--mode', 'annotated-parts', '--stop-after', 'segment'], {cwd: projectRoot, env}), (error) => /仅支持 understand/u.test(error.stderr));
|
||||
await assert.rejects(exec(process.execPath, ['pipeline.mjs', '--mode', 'annotated-parts', '--stop-after', 'export'], {cwd: projectRoot, env}), (error) => /仅支持 understand/u.test(error.stderr));
|
||||
assert.equal(requests, 2);
|
||||
});
|
||||
|
||||
@@ -290,3 +290,15 @@ test('Python geometry regression suite', async () => {
|
||||
const result = await exec(geometryPython(), ['-m', 'unittest', 'discover', '-s', 'tests', '-p', 'test_geometry.py'], {cwd: projectRoot});
|
||||
assert.match(result.stderr, /OK/u);
|
||||
});
|
||||
|
||||
test('Python segmentation QA, prompt integrity and rendering regressions (no model/GPU needed)', async () => {
|
||||
const result = await exec(geometryPython(), ['-m', 'unittest', 'discover', '-s', 'tests', '-p', 'test_segmentation.py'], {cwd: projectRoot});
|
||||
assert.match(result.stderr, /OK/u);
|
||||
});
|
||||
|
||||
test('SAM stage validates local input without requiring relay credentials', async (t) => {
|
||||
const dir = temporary(t);
|
||||
const envFile = path.join(dir, 'empty.env'); writeFileSync(envFile, '');
|
||||
const env = {...process.env, VECTOR_ENGINE_BASE_URL: '', VECTOR_ENGINE_API_KEY: ''};
|
||||
await assert.rejects(exec(process.execPath, ['tools/model-check.mjs', '--env-file', envFile, '--stage', 'segmentation', '--out-dir', path.join(dir, 'result')], {cwd: projectRoot, env}), (error) => /--input-dir/u.test(error.stderr));
|
||||
});
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
import copy
|
||||
import json
|
||||
from pathlib import Path
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from sam_segment import choose_candidate, validate_prompts
|
||||
from segmentation_outputs import assess_candidate, render_outputs, write_report
|
||||
|
||||
|
||||
class SegmentationTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.rgba = np.zeros((24, 32, 4), dtype=np.uint8)
|
||||
self.rgba[4:20, 5:27] = [120, 50, 200, 255]
|
||||
self.rgba[5, 6, 3] = 80
|
||||
self.image = Image.fromarray(self.rgba)
|
||||
self.mask = self.rgba[:, :, 3] > 0
|
||||
|
||||
def test_background_leak_measured_before_clipping(self):
|
||||
raw = np.ones((24, 32), dtype=bool)
|
||||
qa = assess_candidate(raw, self.rgba[:, :, 3], [5, 4, 27, 20], [[6, 5]])
|
||||
self.assertGreater(qa['outside_alpha_ratio'], .5)
|
||||
self.assertEqual(qa['seed_hit_fraction'], 1)
|
||||
self.assertEqual(qa['area_px'], 16 * 22)
|
||||
|
||||
def test_rank_prefers_local_mask_and_never_invents_empty_mask(self):
|
||||
wrong = self.mask.copy()
|
||||
local = np.zeros_like(wrong)
|
||||
local[8:13, 8:13] = True
|
||||
best, _ = choose_candidate([wrong, local], [.95, .7], self.rgba[:, :, 3], [8, 8, 13, 13], [[10, 10]])
|
||||
self.assertEqual(best['index'], 1)
|
||||
best, _ = choose_candidate([np.zeros_like(local)], [.99], self.rgba[:, :, 3], [8, 8, 13, 13], [[10, 10]])
|
||||
self.assertIsNone(best)
|
||||
|
||||
def test_prompts_reject_transparent_points_mismatched_canvas_and_duplicate_ids(self):
|
||||
plan = {'parts': [{'id': 'hand', 'visibility': 'visible'}]}
|
||||
prompts = {'width': 32, 'height': 24, 'parts': [{'id': 'hand', 'status': 'candidate',
|
||||
'seed_source': 'alpha_roi_distance_transform', 'bbox_px_xyxy': [0, 0, 32, 24],
|
||||
'positive_points_px': [[10, 10]]}]}
|
||||
validate_prompts(plan, prompts, self.image)
|
||||
for variant in range(3):
|
||||
bad = copy.deepcopy(prompts)
|
||||
if variant == 0:
|
||||
bad['parts'][0]['positive_points_px'] = [[0, 0]]
|
||||
elif variant == 1:
|
||||
bad['width'] = 99
|
||||
else:
|
||||
bad['parts'].append(copy.deepcopy(bad['parts'][0]))
|
||||
with self.assertRaises(ValueError):
|
||||
validate_prompts(plan, bad, self.image)
|
||||
|
||||
def test_render_preserves_cutout_pixels_alpha_and_measures_overlap(self):
|
||||
parts = [{'id': 'hand', 'semantic': 'hand', 'mask': self.mask, 'model_score': .9},
|
||||
{'id': 'forearm', 'semantic': 'forearm', 'mask': self.mask, 'model_score': .8}]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
result = render_outputs(self.image, parts, directory)
|
||||
self.assertEqual(result['coverage']['coverage_ratio'], 1)
|
||||
self.assertEqual(result['coverage']['overlap_pixels'], int(self.mask.sum()))
|
||||
self.assertEqual(result['near_duplicate_masks'][0]['iou'], 1)
|
||||
self.assertIn('NEAR_DUPLICATE_MASK', result['parts'][0]['warnings'])
|
||||
cut = np.array(Image.open(Path(directory) / 'parts/hand.png'))
|
||||
np.testing.assert_array_equal(cut, self.rgba[4:20, 5:27])
|
||||
self.assertEqual(Image.open(Path(directory) / 'annotated.png').size, self.image.size)
|
||||
sheet = np.array(Image.open(Path(directory) / 'parts-sheet.png'))
|
||||
self.assertTrue(np.all(sheet[0, 0] == 255))
|
||||
a, b = [p['sheet_xyxy'] for p in result['parts']]
|
||||
self.assertTrue(a[2] <= b[0] or b[2] <= a[0] or a[3] <= b[1] or b[3] <= a[1])
|
||||
result['skipped'] = []
|
||||
write_report(directory, result)
|
||||
self.assertTrue((Path(directory) / 'report.html').exists())
|
||||
json.dumps(result)
|
||||
|
||||
def test_renderer_rejects_empty_masks_and_unsafe_ids(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
for name, mask in [('hand', np.zeros_like(self.mask)), ('../bad', self.mask)]:
|
||||
with self.assertRaises(ValueError):
|
||||
render_outputs(self.image, [{'id': name, 'mask': mask}], directory)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
+25
-4
@@ -5,10 +5,12 @@ import {loadConfig, resolveStage, publicStage, samInfo} from '../lib/config.mjs'
|
||||
import {requestJson} from '../providers/http.mjs';
|
||||
import {runUnderstanding, runCompletionProbe} from '../stages/model-stage.mjs';
|
||||
import {readImage} from '../lib/images.mjs';
|
||||
import {runSegmentation} from '../stages/segmentation-stage.mjs';
|
||||
|
||||
const help = `模型接入与结构理解（新增模式，不调用旧拆件流程）
|
||||
node pipeline.mjs --mode annotated-parts --reference image.png --out-dir output/run --stop-after understand
|
||||
node tools/model-check.mjs --stage completion --image local.png --prompt-file task.txt --out-dir output/edit
|
||||
node tools/model-check.mjs --stage segmentation --input-dir output/understand --out-dir output/segment --sam-dir F:/Study/SAM3
|
||||
node tools/model-check.mjs --list-models
|
||||
node tools/model-check.mjs --dry-run
|
||||
|
||||
@@ -23,11 +25,17 @@ const help = `模型接入与结构理解（新增模式，不调用旧拆件流
|
||||
--resume 输入/配置/代码和产物哈希相同时复用
|
||||
--parts-file FILE 离线读取已有 v1/v2 计划，做 Alpha 校验与种子生成，不请求模型
|
||||
--python FILE ai-spine Conda 环境解释器（默认使用已激活环境）
|
||||
--input-dir DIR 已完成的结构理解目录（SAM 阶段输入）
|
||||
--sam-dir DIR 本地 SAM 根目录，包含 sam3 与 sam3.1checkpoint
|
||||
--sam-code-dir DIR SAM 源码仓库目录覆盖
|
||||
--sam-checkpoint FILE 本地权重路径覆盖
|
||||
--sam-confidence FLOAT 模型候选分数阈值，默认 0.2
|
||||
--sam-precision NAME bf16（默认）或 fp32
|
||||
--image FILE 图片编辑输入，可重复；多图需配置 capabilities
|
||||
--mask FILE 透明处可编辑的 RGBA PNG；需配置 capabilities
|
||||
--dry-run 检查配置与本地输入，不请求模型、不生成产物
|
||||
|
||||
本轮 --stop-after 仅支持 understand。SAM 路径已配置，但分割/合并/导出尚未接入。`;
|
||||
--stop-after 支持 understand / segment。SAM 在临时 GPU 子进程运行，完成后退出；补全合并/导出尚未接入。`;
|
||||
|
||||
export function parseOptions(argv) {
|
||||
const options = {stage: 'understanding', imagePaths: []};
|
||||
@@ -37,6 +45,8 @@ export function parseOptions(argv) {
|
||||
['--understanding-protocol', 'understandingProtocol'], ['--output-mode', 'outputMode'],
|
||||
['--stage', 'stage'], ['--mask', 'maskPath'], ['--prompt-file', 'promptFile'], ['--stop-after', 'stopAfter'],
|
||||
['--mode', 'mode'], ['--python', 'python'], ['--stream', 'streamRaw'], ['--parts-file', 'partsFile'],
|
||||
['--input-dir', 'inputDir'], ['--sam-dir', 'samDir'], ['--sam-code-dir', 'samCodeDir'],
|
||||
['--sam-checkpoint', 'samCheckpoint'], ['--sam-confidence', 'samConfidence'], ['--sam-precision', 'samPrecision'],
|
||||
]);
|
||||
for (let i = 0; i < argv.length; i += 1) {
|
||||
const flag = argv[i];
|
||||
@@ -57,15 +67,22 @@ export function parseOptions(argv) {
|
||||
export async function main(argv = process.argv.slice(2)) {
|
||||
const options = parseOptions(argv);
|
||||
if (options.help) { console.log(help); return; }
|
||||
if (!['understanding', 'completion'].includes(options.stage)) throw new Error('--stage 仅支持 understanding 或 completion');
|
||||
if (!['understanding', 'completion', 'segmentation'].includes(options.stage)) throw new Error('--stage 仅支持 understanding / completion / segmentation');
|
||||
if (options.partsFile && (options.stage !== 'understanding' || options.listModels)) throw new Error('--parts-file 仅用于离线结构理解');
|
||||
if (options.streamRaw !== undefined && !['true', 'false'].includes(options.streamRaw)) throw new Error('--stream 仅支持 true 或 false');
|
||||
options.stream = options.streamRaw === undefined ? undefined : options.streamRaw === 'true';
|
||||
if (options.mode && options.mode !== 'annotated-parts') throw new Error('新入口仅支持 --mode annotated-parts');
|
||||
if (options.stopAfter && (options.stopAfter !== 'understand' || options.stage !== 'understanding')) {
|
||||
throw new Error('本轮仅落地模型接入和结构理解；--stop-after 目前仅支持 understand');
|
||||
if (options.stopAfter && (!['understand', 'segment'].includes(options.stopAfter) || options.stage === 'completion' || (options.stage === 'segmentation' && options.stopAfter !== 'segment'))) {
|
||||
throw new Error('--stop-after 仅支持 understand / segment，且必须匹配阶段');
|
||||
}
|
||||
const loaded = loadConfig(options);
|
||||
if (options.stage === 'segmentation' || (options.stopAfter === 'segment' && options.inputDir)) {
|
||||
if (options.reference || options.partsFile || options.imagePaths.length || options.maskPath || options.promptFile || options.listModels) throw new Error('独立 SAM 阶段使用 --input-dir，不接受模型请求或编辑参数');
|
||||
const result = await runSegmentation({...options, loaded});
|
||||
if (!result.dryRun) console.log(`SAM 分割${result.reused ? '复用完成' : '完成'}：${path.join(result.out, 'report.html')}`);
|
||||
return result;
|
||||
}
|
||||
if (options.inputDir) throw new Error('--input-dir 仅用于 SAM 分割阶段');
|
||||
const overrides = options.stage === 'understanding'
|
||||
? {model: options.understandingModel, protocol: options.understandingProtocol, outputMode: options.outputMode, stream: options.stream}
|
||||
: {model: options.completionModel};
|
||||
@@ -93,6 +110,10 @@ export async function main(argv = process.argv.slice(2)) {
|
||||
console.log(options.partsFile ? '角色结构理解：离线计划几何校验' : `角色结构理解：${stage.model} (${stage.protocol})`);
|
||||
result = await runUnderstanding({...options, loaded, stage});
|
||||
console.log(`${result.reused ? '复用已校验产物' : '结构理解完成，语义待复核'}：${path.join(result.out, 'report.html')}`);
|
||||
if (options.stopAfter === 'segment') {
|
||||
result = await runSegmentation({...options, loaded, inputDir: result.out, outDir: path.join(result.out, 'segmentation')});
|
||||
console.log(`SAM 分割完成：${path.join(result.out, 'report.html')}`);
|
||||
}
|
||||
} else {
|
||||
if (options.reference) throw new Error('图片编辑使用 --image 指定局部画布；额外参考图也通过重复 --image 指定');
|
||||
console.log(`图片编辑接入测试：${stage.model}`);
|
||||
|
||||
Reference in New Issue
Block a user