export async function getFeatureVectors()

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;
}