Added biometrics for avatars.

src/biometrics.js

@@ -0,0 +1,204 @@
+'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;
+	}
+
+	const tensor = human.tf.node.decodeImage(buffer, 3);
+
+	try {
+		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,
+};