traxxx/src/biometrics.js

'use strict';

const config = require('config');
const path = require('path');
const fsPromises = require('fs').promises;
const Human = require('@vladmandic/human').default;

const knex = require('./knex');
const logger = require('./logger')(__filename);

const humanConfig = {
	modelBasePath: `file://${path.join(__dirname, '../node_modules/@vladmandic/human/models')}/`,
	backend: 'tensorflow', // uses tfjs-node under the hood
	face: {
		enabled: true,
		detector: { rotation: false },
		mesh: { enabled: false },
		iris: { enabled: false },
		description: { enabled: true },
		emotion: { enabled: true },
	},
	body: { enabled: false },
	hand: { enabled: false },
	gesture: { enabled: false },
};

let humanInstance = null;
const nsPerSec = 1e9;

async function getHumanInstance() {
	if (!humanInstance) {
		humanInstance = new Human(humanConfig);

		await humanInstance.load();
		await humanInstance.warmup();
	}

	return humanInstance;
}

async function getImageBuffer(mediaEntry) {
	if (mediaEntry.is_s3) {
		if (!config.s3.enabled) {
			logger.verbose(`Skipping biometrics for ${mediaEntry.media_id}, S3 not enabled`);
			return null;
		}

		const res = await fetch(`https://${config.s3.bucket}/${mediaEntry.path}`);

		if (res.ok) {
			return Buffer.from(await res.arrayBuffer());
		}

		return null;
	}

	try {
		const buffer = await fsPromises.readFile(path.join(config.media.path, mediaEntry.path));

		return buffer;
	} catch (_error) {
		return null;
	}
}

async function getBiometrics(avatarEntry) {
	if (!avatarEntry) {
		return null;
	}

	const human = await getHumanInstance();
	const buffer = await getImageBuffer(avatarEntry);

	if (!buffer) {
		return null;
	}

	let tensor;

	try {
		tensor = human.tf.node.decodeImage(buffer, 3);
		const result = await human.detect(tensor);

		return result;
	} catch (error) {
		logger.error(error);

		return null;
	} finally {
		tensor.dispose();
	}
}

function getAvatarSource(actorIds, includePhotos) {
	if (includePhotos) {
		return knex('actors_avatars')
			.select('actor_id', 'media_id')
			.modify((builder) => {
				if (actorIds) {
					builder.whereIn('actor_id', actorIds);
				}
			});
	}

	return knex('actors')
		.select('id as actor_id', 'avatar_media_id as media_id')
		.whereNotNull('avatar_media_id')
		.modify((builder) => {
			if (actorIds) {
				builder.whereIn('id', actorIds);
			}
		});
}

async function setBiometrics(actorIds, shouldUpdate = false, includePhotos = false) {
	if (!config.media.biometrics.enabled) {
		return;
	}

	const startTime = process.hrtime();

	const avatarSource = getAvatarSource(actorIds, includePhotos);

	const avatarEntries = await knex
		.select(
			'avatar_source.actor_id',
			'media.id as media_id',
			'media.path',
			'media.is_s3',
			knex.raw('coalesce(json_agg(media_biometrics.biometrics) filter (where media_biometrics.id is not null), \'[]\') as biometrics'),
		)
		.from(avatarSource.as('avatar_source'))
		.leftJoin('media', 'media.id', 'avatar_source.media_id')
		.leftJoin('media_biometrics', 'media_biometrics.media_id', 'media.id')
		.groupBy('avatar_source.actor_id', 'media.id', 'media.path', 'media.is_s3');

	await avatarEntries.reduce(async (chain, avatarEntry, avatarIndex) => {
		await chain;

		if (avatarEntry.biometrics.length > 0 && !shouldUpdate) {
			logger.verbose(`Skipping biometrics for ${avatarEntry.media_id}, already present`);
			return;
		}

		const biometrics = await getBiometrics(avatarEntry);

		if (!biometrics) {
			return;
		}

		const curatedBiometrics = biometrics.face.map((face) => ({
			biometrics: {
				age: face.age,
				gender: face.gender,
				genderScore: face.genderScore,
				emotion: face.emotion,
				box: face.box,
				score: face.score,
				...Object.fromEntries(Object.entries(face.annotations).map(([key, annotation]) => [key, annotation.filter(Boolean)[0] || null])),
			},
			embedding: face.embedding,
		}));

		if (curatedBiometrics.length === 0) {
			return;
		}

		if (shouldUpdate) {
			await knex('media_biometrics')
				.where('media_id', avatarEntry.media_id)
				.delete();
		}

		await knex('media_biometrics')
			.insert(curatedBiometrics.map((face, index) => ({
				media_id: avatarEntry.media_id,
				face_index: index,
				width: biometrics.width,
				height: biometrics.height,
				biometrics: JSON.stringify(face.biometrics),
				embedding: knex.raw('?::vector', [JSON.stringify(face.embedding)]),
				updated_at: knex.raw('now()'),
			})))
			.onConflict(['media_id', 'face_index'])
			.ignore();

		const processed = avatarIndex + 1;
		const [elapsedSec, elapsedNs] = process.hrtime(startTime);
		const elapsedMs = (elapsedSec * nsPerSec + elapsedNs) / 1e6;
		const avgMsPerEntry = elapsedMs / processed;
		const remainingMs = avgMsPerEntry * (avatarEntries.length - processed);
		const remainingMinutes = (remainingMs / 1000 / 60).toFixed(1);

		logger.info(`Set biometrics for ${avatarEntry.media_id} (${processed}/${avatarEntries.length}, ~${remainingMinutes}m remaining)`);
	}, Promise.resolve());

	const diffTime = process.hrtime(startTime);
	const totalMs = (diffTime[0] * nsPerSec + diffTime[1]) / 1e6;

	logger.info(`Biometrics done in ${(totalMs / 1000).toFixed(1)}s for ${actorIds?.length || 'all'} actors and ${avatarEntries.length} avatars`);
}

module.exports = {
	setBiometrics,
};