in storage-reverse-image-search/functions/src/common/feature_vectors.ts [88:128]
export async function getFeatureVectors(
images: string[],
batchSize = 32
): Promise<number[][]> {
console.log(`→ getFeatureVectors called with ${images.length} images`);
if (images.length === 0) {
return [];
}
// Lazy-load the model
if (!model) {
const fromTFHub = config.modelFromTFHub;
console.log(
`Loading model from ${config.modelUrl} (fromTFHub=${fromTFHub})`
);
model = await tf.loadGraphModel(config.modelUrl, {fromTFHub});
}
const preppedTensors: tf.Tensor4D[] = await Promise.all(
images.map(img => _loadImageAsTensor(img))
);
const allEmbeddings: number[][] = [];
for (let i = 0; i < preppedTensors.length; i += batchSize) {
const slice = preppedTensors.slice(i, i + batchSize);
const batchEmbeddings: number[][] = tf.tidy(() => {
const batch = tf.concat(slice, 0) as tf.Tensor4D;
const rawOut = model.predict(batch) as tf.Tensor;
return rawOut.arraySync() as number[][];
});
slice.forEach(t => t.dispose());
allEmbeddings.push(...batchEmbeddings);
}
return allEmbeddings;
}