feat: initial commit - massive-cloud point cloud processing system
This commit is contained in:
115
server/assemblyGpuNearest.ts
Normal file
115
server/assemblyGpuNearest.ts
Normal file
@@ -0,0 +1,115 @@
|
||||
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', 'assembly_nearest_gpu.py');
|
||||
}
|
||||
|
||||
function resolvePython3(): string {
|
||||
const fromEnv = process.env.PYTHON3_PATH?.trim();
|
||||
if (fromEnv) return fromEnv;
|
||||
if (process.platform === 'win32') return 'python3';
|
||||
for (const p of ['/usr/bin/python3', '/usr/local/bin/python3']) {
|
||||
try {
|
||||
if (fs.existsSync(p)) return p;
|
||||
} catch {
|
||||
/* ignore */
|
||||
}
|
||||
}
|
||||
return 'python3';
|
||||
}
|
||||
|
||||
/** 环境变量 ASSEMBLY_USE_GPU=1 / true 时尝试 GPU(失败则由调用方回退 CPU) */
|
||||
export function assemblyGpuNearestEnabled(): boolean {
|
||||
const v = process.env.ASSEMBLY_USE_GPU?.trim().toLowerCase();
|
||||
return v === '1' || v === 'true' || v === 'yes';
|
||||
}
|
||||
|
||||
export type GpuNearestNumbers = {
|
||||
assignments: number[];
|
||||
pointDistances: number[];
|
||||
};
|
||||
|
||||
/**
|
||||
* 调用 scripts/assembly_nearest_gpu.py(PyTorch3D CUDA)。
|
||||
* 返回 null:未启用、脚本缺失、无 CUDA、导入失败(exit 2)或其它错误。
|
||||
*/
|
||||
export function tryRunAssemblyNearestGpu(manifestDir: string, sampled: SampledPoint[]): GpuNearestNumbers | null {
|
||||
if (!assemblyGpuNearestEnabled()) 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(), 'asm-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, '--points', ptsFile, '--manifest-dir', manifestDir, '--out-dir', tmp],
|
||||
{
|
||||
encoding: 'utf-8',
|
||||
maxBuffer: 64 * 1024 * 1024,
|
||||
env: { ...process.env, PYTHONUNBUFFERED: '1' },
|
||||
},
|
||||
);
|
||||
|
||||
if (r.error) {
|
||||
console.warn('[AssemblyGPU]', r.error.message);
|
||||
return null;
|
||||
}
|
||||
if (r.status === 2) {
|
||||
console.warn('[AssemblyGPU] 回退 CPU:', (r.stdout || '').trim().slice(-500) || 'exit 2');
|
||||
return null;
|
||||
}
|
||||
if (r.status !== 0) {
|
||||
console.warn('[AssemblyGPU] 失败 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;
|
||||
n?: number;
|
||||
assignmentsFile?: string;
|
||||
distancesFile?: string;
|
||||
};
|
||||
if (!meta.ok || !meta.assignmentsFile || !meta.distancesFile) return null;
|
||||
if (meta.n !== sampled.length) {
|
||||
console.warn('[AssemblyGPU] 点数不一致', meta.n, sampled.length);
|
||||
return null;
|
||||
}
|
||||
|
||||
const aBuf = fs.readFileSync(meta.assignmentsFile);
|
||||
const dBuf = fs.readFileSync(meta.distancesFile);
|
||||
const n = sampled.length;
|
||||
if (aBuf.length !== n * 4 || dBuf.length !== n * 4) {
|
||||
console.warn('[AssemblyGPU] 二进制长度异常');
|
||||
return null;
|
||||
}
|
||||
const assignments = Array.from(new Int32Array(aBuf.buffer, aBuf.byteOffset, n));
|
||||
const pointDistances = Array.from(new Float32Array(dBuf.buffer, dBuf.byteOffset, n));
|
||||
return { assignments, pointDistances };
|
||||
} finally {
|
||||
try {
|
||||
fs.rmSync(tmp, { recursive: true, force: true });
|
||||
} catch {
|
||||
/* ignore */
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user