122 lines
3.5 KiB
TypeScript
122 lines
3.5 KiB
TypeScript
import { spawnSync } from 'node:child_process';
|
||
import * as fs from 'node:fs';
|
||
import * as os from 'node:os';
|
||
import * as path from 'node:path';
|
||
import { fileURLToPath } from 'node:url';
|
||
import type { SampledPoint } from './pointcloudSample.js';
|
||
|
||
function scriptPath(): string {
|
||
const here = path.dirname(fileURLToPath(import.meta.url));
|
||
return path.join(here, '..', 'scripts', 'deviation_heatmap_gpu.py');
|
||
}
|
||
|
||
function resolvePython3(): string {
|
||
const fromEnv = process.env.PYTHON3_PATH?.trim();
|
||
if (fromEnv) return fromEnv;
|
||
for (const p of ['/usr/bin/python3', '/usr/local/bin/python3']) {
|
||
try {
|
||
if (fs.existsSync(p)) return p;
|
||
} catch {
|
||
/* ignore */
|
||
}
|
||
}
|
||
return 'python3';
|
||
}
|
||
|
||
export function deviationGpuEnabled(): boolean {
|
||
const v = process.env.ASSEMBLY_USE_GPU?.trim().toLowerCase();
|
||
return v === '1' || v === 'true' || v === 'yes';
|
||
}
|
||
|
||
export type GpuDeviationResult = {
|
||
colors: Float32Array;
|
||
vertexCount: number;
|
||
dispMin: number;
|
||
dispMax: number;
|
||
};
|
||
|
||
/**
|
||
* 调用 scripts/deviation_heatmap_gpu.py（PyTorch CUDA KNN）。
|
||
* 返回 null 时由调用方回退 CPU。
|
||
*/
|
||
export function tryRunDeviationGpu(
|
||
cadFilePath: string,
|
||
sampled: SampledPoint[],
|
||
): GpuDeviationResult | null {
|
||
if (!deviationGpuEnabled()) return null;
|
||
const script = scriptPath();
|
||
if (!fs.existsSync(script)) return null;
|
||
if (sampled.length === 0) return null;
|
||
|
||
const py = resolvePython3();
|
||
const tmp = fs.mkdtempSync(path.join(os.tmpdir(), 'dev-gpu-'));
|
||
const ptsFile = path.join(tmp, 'points.bin');
|
||
const buf = Buffer.allocUnsafe(sampled.length * 12);
|
||
for (let i = 0; i < sampled.length; i++) {
|
||
const o = i * 12;
|
||
buf.writeFloatLE(sampled[i]!.x, o);
|
||
buf.writeFloatLE(sampled[i]!.y, o + 4);
|
||
buf.writeFloatLE(sampled[i]!.z, o + 8);
|
||
}
|
||
fs.writeFileSync(ptsFile, buf);
|
||
|
||
try {
|
||
const r = spawnSync(
|
||
py,
|
||
[script, '--cad', cadFilePath, '--points', ptsFile, '--out-dir', tmp],
|
||
{
|
||
encoding: 'utf-8',
|
||
maxBuffer: 64 * 1024 * 1024,
|
||
timeout: 30 * 60 * 1000,
|
||
env: { ...process.env, PYTHONUNBUFFERED: '1' },
|
||
},
|
||
);
|
||
|
||
if (r.error) {
|
||
console.warn('[DeviationGPU]', r.error.message);
|
||
return null;
|
||
}
|
||
if (r.status === 2) {
|
||
console.warn('[DeviationGPU] 回退 CPU：', (r.stdout || '').trim().slice(-500) || 'exit 2');
|
||
return null;
|
||
}
|
||
if (r.status !== 0) {
|
||
console.warn('[DeviationGPU] 失败 status=', r.status, (r.stderr || r.stdout || '').slice(0, 800));
|
||
return null;
|
||
}
|
||
|
||
const resultJsonPath = path.join(tmp, 'result.json');
|
||
if (!fs.existsSync(resultJsonPath)) return null;
|
||
const meta = JSON.parse(fs.readFileSync(resultJsonPath, 'utf-8')) as {
|
||
ok?: boolean;
|
||
vertexCount?: number;
|
||
dispMin?: number;
|
||
dispMax?: number;
|
||
};
|
||
if (!meta.ok || !meta.vertexCount) return null;
|
||
|
||
const colorsPath = path.join(tmp, 'colors.bin');
|
||
if (!fs.existsSync(colorsPath)) return null;
|
||
const cBuf = fs.readFileSync(colorsPath);
|
||
const expectedBytes = meta.vertexCount * 3 * 4;
|
||
if (cBuf.length !== expectedBytes) {
|
||
console.warn('[DeviationGPU] colors.bin 长度异常', cBuf.length, '!=', expectedBytes);
|
||
return null;
|
||
}
|
||
const colors = new Float32Array(cBuf.buffer, cBuf.byteOffset, meta.vertexCount * 3);
|
||
|
||
return {
|
||
colors,
|
||
vertexCount: meta.vertexCount,
|
||
dispMin: meta.dispMin ?? 0,
|
||
dispMax: meta.dispMax ?? 1,
|
||
};
|
||
} finally {
|
||
try {
|
||
fs.rmSync(tmp, { recursive: true, force: true });
|
||
} catch {
|
||
/* ignore */
|
||
}
|
||
}
|
||
}
|