classifier.js

import { pipeline } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@4.2.0';

export const MODEL_OPTIONS = [
  {
    key: 'resnet18',
    label: 'ResNet-18',
    modelId: 'Xenova/resnet-18',
    task: 'image-classification',
    note: 'Recommended classroom default: a compact CNN with 1,000 ImageNet labels.',
    dtype: 'q8',
  },
  {
    key: 'vit',
    label: 'ViT-B/16',
    modelId: 'Xenova/vit-base-patch16-224',
    task: 'image-classification',
    note: 'A Vision Transformer that compares patterns across image patches.',
    dtype: 'q8',
  },
  {
    key: 'clip',
    label: 'CLIP ViT-B/32',
    modelId: 'Xenova/clip-vit-base-patch32',
    task: 'zero-shot-image-classification',
    note: 'A vision-language model that ranks your own candidate descriptions.',
    dtype: 'q8',
  },
];

const classifierCache = new Map();

export function getModelOption(key) {
  return MODEL_OPTIONS.find((model) => model.key === key) || MODEL_OPTIONS[0];
}

export async function loadClassifier(modelKey, progressCallback) {
  const option = getModelOption(modelKey);
  if (classifierCache.has(option.key)) return classifierCache.get(option.key);

  const classifier = await pipeline(option.task, option.modelId, {
    dtype: option.dtype,
    progress_callback: progressCallback,
  });

  classifierCache.set(option.key, classifier);
  return classifier;
}

export async function classifyImage(imageDataUrl, options = {}) {
  const modelKey = options.modelKey || MODEL_OPTIONS[0].key;
  const topK = Number.isFinite(options.topK) ? options.topK : 5;
  const option = getModelOption(modelKey);
  const classifier = await loadClassifier(modelKey, options.progressCallback);

  let predictions;
  if (option.task === 'zero-shot-image-classification') {
    predictions = await classifier(imageDataUrl, options.candidateLabels, {
      hypothesis_template: 'a photo of {}',
    });
  } else {
    predictions = await classifier(imageDataUrl, { top_k: topK });
  }

  return [...predictions]
    .sort((a, b) => Number(b.score) - Number(a.score))
    .slice(0, topK);
}