534 lines
9.4 KiB
JavaScript
534 lines
9.4 KiB
JavaScript
const fs = require('fs');
|
|
const os = require('os');
|
|
const path = require('path');
|
|
const { spawnSync } = require('child_process');
|
|
|
|
const db = require('../db/knex');
|
|
|
|
const PROJECT_ROOT = path.resolve(__dirname, '..');
|
|
|
|
const CONFIG_PATH = path.join(
|
|
PROJECT_ROOT,
|
|
'native',
|
|
'face-knn',
|
|
'config.json',
|
|
);
|
|
|
|
function loadConfig() {
|
|
const data = JSON.parse(
|
|
fs.readFileSync(CONFIG_PATH, 'utf8'),
|
|
);
|
|
|
|
return {
|
|
enabled:
|
|
data.enabled !== false,
|
|
|
|
mode:
|
|
String(data.mode || 'shadow'),
|
|
|
|
backend:
|
|
String(data.backend || 'opencl'),
|
|
|
|
fallbackBackend:
|
|
String(data.fallbackBackend || 'cpu'),
|
|
|
|
dimension:
|
|
Number(data.dimension || 128),
|
|
|
|
topK:
|
|
Number(data.topK || 5),
|
|
|
|
recognitionThreshold:
|
|
Number(data.recognitionThreshold ?? 0.50),
|
|
|
|
scoreTolerance:
|
|
Number(data.scoreTolerance ?? 0.00001),
|
|
|
|
launcher:
|
|
path.resolve(
|
|
PROJECT_ROOT,
|
|
String(
|
|
data.launcher
|
|
|| 'native/face-knn/run-face-knn.sh'
|
|
),
|
|
),
|
|
|
|
runtimeDirectory:
|
|
path.resolve(
|
|
PROJECT_ROOT,
|
|
String(
|
|
data.runtimeDirectory
|
|
|| 'var/face-knn'
|
|
),
|
|
),
|
|
};
|
|
}
|
|
|
|
function validateEmbedding(
|
|
embedding,
|
|
faceId,
|
|
dimension,
|
|
) {
|
|
if (!Buffer.isBuffer(embedding)) {
|
|
throw new Error(
|
|
`Embedding non Buffer per ${faceId}`,
|
|
);
|
|
}
|
|
|
|
const expectedBytes =
|
|
dimension * Float32Array.BYTES_PER_ELEMENT;
|
|
|
|
if (embedding.length !== expectedBytes) {
|
|
throw new Error(
|
|
`Embedding non valido per ${faceId}: `
|
|
+ `${embedding.length} byte, `
|
|
+ `attesi ${expectedBytes}`,
|
|
);
|
|
}
|
|
}
|
|
|
|
function normalizeFaceIds(faceIds, label) {
|
|
if (!Array.isArray(faceIds) || faceIds.length === 0) {
|
|
throw new Error(
|
|
`${label} mancante o vuoto`,
|
|
);
|
|
}
|
|
|
|
const normalized = faceIds.map(
|
|
(faceId) => String(faceId || '').trim(),
|
|
);
|
|
|
|
if (normalized.some((faceId) => !faceId)) {
|
|
throw new Error(
|
|
`${label} contiene face_id non valido`,
|
|
);
|
|
}
|
|
|
|
if (new Set(normalized).size !== normalized.length) {
|
|
throw new Error(
|
|
`${label} contiene face_id duplicati`,
|
|
);
|
|
}
|
|
|
|
return normalized;
|
|
}
|
|
|
|
async function loadFacesByIds(
|
|
connection,
|
|
faceIds,
|
|
dimension,
|
|
) {
|
|
const normalizedIds = normalizeFaceIds(
|
|
faceIds,
|
|
'faceIds',
|
|
);
|
|
|
|
const rows = await connection('faces as face')
|
|
.join(
|
|
'photos as photo',
|
|
'photo.id',
|
|
'face.photo_id',
|
|
)
|
|
.whereIn('face.id', normalizedIds)
|
|
.whereNotNull('face.embedding')
|
|
.select(
|
|
'face.id as face_id',
|
|
'face.embedding',
|
|
'face.face_index',
|
|
'photo.name as photo_name',
|
|
);
|
|
|
|
const byId = new Map(
|
|
rows.map((row) => [
|
|
row.face_id,
|
|
row,
|
|
]),
|
|
);
|
|
|
|
return normalizedIds.map((faceId) => {
|
|
const row = byId.get(faceId);
|
|
|
|
if (!row) {
|
|
throw new Error(
|
|
`Volto o embedding mancante: ${faceId}`,
|
|
);
|
|
}
|
|
|
|
validateEmbedding(
|
|
row.embedding,
|
|
faceId,
|
|
dimension,
|
|
);
|
|
|
|
return row;
|
|
});
|
|
}
|
|
|
|
async function loadActivePrototypes(
|
|
connection,
|
|
personId,
|
|
user,
|
|
dimension,
|
|
) {
|
|
const rows = await connection(
|
|
'person_face_prototypes as prototype',
|
|
)
|
|
.join(
|
|
'faces as face',
|
|
'face.id',
|
|
'prototype.face_id',
|
|
)
|
|
.join(
|
|
'photos as photo',
|
|
'photo.id',
|
|
'face.photo_id',
|
|
)
|
|
.where({
|
|
'prototype.person_id': personId,
|
|
'prototype.user': user,
|
|
'prototype.active': 1,
|
|
})
|
|
.whereNotNull('face.embedding')
|
|
.orderBy(
|
|
'prototype.selection_order',
|
|
'asc',
|
|
)
|
|
.select(
|
|
'prototype.selection_order',
|
|
'prototype.selection_method',
|
|
'prototype.face_id',
|
|
'face.embedding',
|
|
'face.face_index',
|
|
'photo.name as photo_name',
|
|
);
|
|
|
|
if (rows.length === 0) {
|
|
throw new Error(
|
|
'Nessun prototype attivo disponibile',
|
|
);
|
|
}
|
|
|
|
for (let index = 0; index < rows.length; index += 1) {
|
|
const row = rows[index];
|
|
|
|
validateEmbedding(
|
|
row.embedding,
|
|
row.face_id,
|
|
dimension,
|
|
);
|
|
|
|
const expectedOrder = index + 1;
|
|
|
|
if (
|
|
Number(row.selection_order)
|
|
!== expectedOrder
|
|
) {
|
|
throw new Error(
|
|
'selection_order non contiguo: '
|
|
+ `atteso ${expectedOrder}, `
|
|
+ `trovato ${row.selection_order}`,
|
|
);
|
|
}
|
|
}
|
|
|
|
return rows;
|
|
}
|
|
|
|
function writeBinaryFile(filePath, rows) {
|
|
fs.writeFileSync(
|
|
filePath,
|
|
Buffer.concat(
|
|
rows.map((row) => row.embedding),
|
|
),
|
|
);
|
|
}
|
|
|
|
function buildSelfMap(
|
|
queryRows,
|
|
prototypeRows,
|
|
) {
|
|
const prototypeIndexByFaceId =
|
|
new Map(
|
|
prototypeRows.map(
|
|
(prototype, index) => [
|
|
prototype.face_id,
|
|
index,
|
|
],
|
|
),
|
|
);
|
|
|
|
return queryRows.map((query) => {
|
|
const index =
|
|
prototypeIndexByFaceId.get(
|
|
query.face_id,
|
|
);
|
|
|
|
return index == null
|
|
? -1
|
|
: index;
|
|
});
|
|
}
|
|
|
|
function writeSelfMap(
|
|
filePath,
|
|
selfMap,
|
|
) {
|
|
fs.writeFileSync(
|
|
filePath,
|
|
`${selfMap.join('\n')}\n`,
|
|
'utf8',
|
|
);
|
|
}
|
|
|
|
function runNativeKnn({
|
|
config,
|
|
queryPath,
|
|
databasePath,
|
|
selfMapPath,
|
|
queryCount,
|
|
databaseCount,
|
|
backend,
|
|
}) {
|
|
const args = [
|
|
'--queries',
|
|
queryPath,
|
|
'--database',
|
|
databasePath,
|
|
'--query-count',
|
|
String(queryCount),
|
|
'--database-count',
|
|
String(databaseCount),
|
|
'--dimension',
|
|
String(config.dimension),
|
|
'--top-k',
|
|
String(
|
|
Math.min(
|
|
config.topK,
|
|
databaseCount,
|
|
),
|
|
),
|
|
'--threshold',
|
|
String(config.recognitionThreshold),
|
|
'--self-map',
|
|
selfMapPath,
|
|
'--backend',
|
|
backend,
|
|
];
|
|
|
|
const result = spawnSync(
|
|
config.launcher,
|
|
args,
|
|
{
|
|
cwd: PROJECT_ROOT,
|
|
encoding: 'utf8',
|
|
stdio: [
|
|
'ignore',
|
|
'pipe',
|
|
'pipe',
|
|
],
|
|
maxBuffer: 16 * 1024 * 1024,
|
|
},
|
|
);
|
|
|
|
if (result.error) {
|
|
throw result.error;
|
|
}
|
|
|
|
if (result.status !== 0) {
|
|
throw new Error(
|
|
`face_knn ${backend} fallito `
|
|
+ `(status=${result.status}):\n`
|
|
+ String(result.stderr || '')
|
|
+ String(result.stdout || ''),
|
|
);
|
|
}
|
|
|
|
let output;
|
|
|
|
try {
|
|
output = JSON.parse(
|
|
String(result.stdout || ''),
|
|
);
|
|
} catch (error) {
|
|
throw new Error(
|
|
'Output JSON face_knn non valido: '
|
|
+ error.message
|
|
+ '\nstdout:\n'
|
|
+ String(result.stdout || '').slice(0, 2000)
|
|
+ '\nstderr:\n'
|
|
+ String(result.stderr || '').slice(0, 2000),
|
|
);
|
|
}
|
|
|
|
return {
|
|
output,
|
|
stderr:
|
|
String(result.stderr || ''),
|
|
command: {
|
|
launcher: config.launcher,
|
|
args,
|
|
},
|
|
};
|
|
}
|
|
|
|
async function runFaceKnn({
|
|
connection = db,
|
|
personId,
|
|
user,
|
|
queryFaceIds,
|
|
backend,
|
|
keepTemporaryFiles = false,
|
|
}) {
|
|
const config = loadConfig();
|
|
|
|
if (!config.enabled) {
|
|
throw new Error(
|
|
'face_knn disabilitato dalla configurazione',
|
|
);
|
|
}
|
|
|
|
const normalizedQueryIds =
|
|
normalizeFaceIds(
|
|
queryFaceIds,
|
|
'queryFaceIds',
|
|
);
|
|
|
|
const [prototypeRows, queryRows] =
|
|
await Promise.all([
|
|
loadActivePrototypes(
|
|
connection,
|
|
personId,
|
|
user,
|
|
config.dimension,
|
|
),
|
|
loadFacesByIds(
|
|
connection,
|
|
normalizedQueryIds,
|
|
config.dimension,
|
|
),
|
|
]);
|
|
|
|
fs.mkdirSync(
|
|
config.runtimeDirectory,
|
|
{
|
|
recursive: true,
|
|
},
|
|
);
|
|
|
|
const temporaryDirectory =
|
|
fs.mkdtempSync(
|
|
path.join(
|
|
config.runtimeDirectory,
|
|
'run-',
|
|
),
|
|
);
|
|
|
|
const queryPath =
|
|
path.join(
|
|
temporaryDirectory,
|
|
'queries.bin',
|
|
);
|
|
|
|
const databasePath =
|
|
path.join(
|
|
temporaryDirectory,
|
|
'prototypes.bin',
|
|
);
|
|
|
|
const selfMapPath =
|
|
path.join(
|
|
temporaryDirectory,
|
|
'self-map.txt',
|
|
);
|
|
|
|
const selfMap =
|
|
buildSelfMap(
|
|
queryRows,
|
|
prototypeRows,
|
|
);
|
|
|
|
try {
|
|
writeBinaryFile(
|
|
queryPath,
|
|
queryRows,
|
|
);
|
|
|
|
writeBinaryFile(
|
|
databasePath,
|
|
prototypeRows,
|
|
);
|
|
|
|
writeSelfMap(
|
|
selfMapPath,
|
|
selfMap,
|
|
);
|
|
|
|
const nativeResult =
|
|
runNativeKnn({
|
|
config,
|
|
queryPath,
|
|
databasePath,
|
|
selfMapPath,
|
|
queryCount:
|
|
queryRows.length,
|
|
databaseCount:
|
|
prototypeRows.length,
|
|
backend:
|
|
backend || config.backend,
|
|
});
|
|
|
|
return {
|
|
config,
|
|
output:
|
|
nativeResult.output,
|
|
stderr:
|
|
nativeResult.stderr,
|
|
command:
|
|
nativeResult.command,
|
|
queryRows:
|
|
queryRows.map((row, index) => ({
|
|
index,
|
|
face_id: row.face_id,
|
|
face_index: row.face_index,
|
|
photo_name: row.photo_name,
|
|
excluded_index:
|
|
selfMap[index],
|
|
})),
|
|
prototypeRows:
|
|
prototypeRows.map(
|
|
(row, index) => ({
|
|
index,
|
|
face_id: row.face_id,
|
|
face_index: row.face_index,
|
|
photo_name: row.photo_name,
|
|
selection_order:
|
|
Number(row.selection_order),
|
|
selection_method:
|
|
row.selection_method,
|
|
}),
|
|
),
|
|
temporaryDirectory:
|
|
keepTemporaryFiles
|
|
? temporaryDirectory
|
|
: null,
|
|
};
|
|
} finally {
|
|
if (!keepTemporaryFiles) {
|
|
fs.rmSync(
|
|
temporaryDirectory,
|
|
{
|
|
recursive: true,
|
|
force: true,
|
|
},
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
module.exports = {
|
|
loadConfig,
|
|
loadFacesByIds,
|
|
loadActivePrototypes,
|
|
buildSelfMap,
|
|
runFaceKnn,
|
|
};
|