75 lines
2.3 KiB
TypeScript
75 lines
2.3 KiB
TypeScript
import * as ort from 'onnxruntime-node';
|
|
import { existsSync } from 'node:fs';
|
|
|
|
const MODEL_PATH = process.env.INPAINT_MODEL_PATH || '/models/migan_pipeline_v2.onnx';
|
|
|
|
let sessionPromise: Promise<ort.InferenceSession | null> | null = null;
|
|
|
|
export async function getInpaintSession(): Promise<ort.InferenceSession | null> {
|
|
if (!sessionPromise) {
|
|
sessionPromise = (async () => {
|
|
const candidates = [
|
|
MODEL_PATH,
|
|
'/home/locpham/RecipesCam/docker/frontend/public/models/migan_pipeline_v2.onnx',
|
|
'../frontend/public/models/migan_pipeline_v2.onnx',
|
|
'./models/migan_pipeline_v2.onnx',
|
|
];
|
|
let foundPath: string | null = null;
|
|
for (const p of candidates) {
|
|
if (existsSync(p)) {
|
|
foundPath = p;
|
|
break;
|
|
}
|
|
}
|
|
if (!foundPath) {
|
|
console.warn('[inpaintService] Model file not found in candidates:', candidates);
|
|
return null;
|
|
}
|
|
|
|
console.log('[inpaintService] Loading ONNX model from:', foundPath);
|
|
const t0 = Date.now();
|
|
const session = await ort.InferenceSession.create(foundPath, {
|
|
executionProviders: ['cpu'],
|
|
graphOptimizationLevel: 'all',
|
|
});
|
|
console.log('[inpaintService] Model loaded in ' + (Date.now() - t0) + 'ms');
|
|
return session;
|
|
})().catch((err) => {
|
|
console.error('[inpaintService] Failed to load ONNX model:', err);
|
|
sessionPromise = null;
|
|
return null;
|
|
});
|
|
}
|
|
return sessionPromise;
|
|
}
|
|
|
|
export async function runServerInpaint(
|
|
srcBytes: Uint8Array,
|
|
maskBytes: Uint8Array,
|
|
w: number,
|
|
h: number
|
|
): Promise<Uint8Array | null> {
|
|
const session = await getInpaintSession();
|
|
if (!session) return null;
|
|
|
|
const in0 = new ort.Tensor('uint8', srcBytes, [1, 3, h, w]);
|
|
const in1 = new ort.Tensor('uint8', maskBytes, [1, 1, h, w]);
|
|
|
|
const out = await session.run({
|
|
[session.inputNames[0]]: in0,
|
|
[session.inputNames[1]]: in1,
|
|
});
|
|
|
|
const tensor = out[session.outputNames[0]];
|
|
const data = tensor.data as Uint8Array;
|
|
const patch = new Uint8Array(4 * w * h);
|
|
const plane = w * h;
|
|
for (let i = 0; i < plane; i++) {
|
|
patch[i * 4] = data[i];
|
|
patch[i * 4 + 1] = data[plane + i];
|
|
patch[i * 4 + 2] = data[2 * plane + i];
|
|
patch[i * 4 + 3] = 255;
|
|
}
|
|
return patch;
|
|
}
|