第二步SAM模型接入

2026-09-18 20:29:29 +08:00
parent 1a7b2f14a9
commit f095c28858
16 changed files with 788 additions and 26 deletions
@@ -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;
+  }
+}