hotdog-app/src/lib/classifier.ts

import { classifyImageWithVision } from './vision';
import * as tf from '@tensorflow/tfjs';
import '@tensorflow/tfjs-react-native';
import * as mobilenet from '@tensorflow-models/mobilenet';
import { decodeJpeg, fetch } from '@tensorflow/tfjs-react-native';
import * as ImageManipulator from 'expo-image-manipulator';
import { Platform } from 'react-native';

const HOT_DOG_IMAGENET_LABEL = 'hotdog, hot dog, red hot';
const MODEL_INPUT_SIZE = 224;

const HOT_DOG_LABEL_PATTERNS = [
  /hot[\s_-]?dog/,
  /frankfurter/,
  /\bwiener\b/,
  /red[\s_-]?hot/,
];

const CONFUSABLE_TOP_LABELS = [
  /pizza/,
  /burger/,
  /sandwich/,
  /person/,
  /dog\b(?!.*hot)/,
  /cat\b/,
  /car\b/,
  /phone/,
];

let mobileNetModel: mobilenet.MobileNet | null = null;
let initPromise: Promise<void> | null = null;
let useVision = Platform.OS === 'ios';

export type ClassificationResult = {
  isHotDog: boolean;
  confidence: number;
  topLabel: string;
};

function normalizeLabel(label: string): string {
  return label.toLowerCase().replace(/_/g, ' ').trim();
}

function isHotDogLabel(label: string): boolean {
  const normalized = normalizeLabel(label);
  return HOT_DOG_LABEL_PATTERNS.some((pattern) => pattern.test(normalized));
}

function isConfusableLabel(label: string): boolean {
  const normalized = normalizeLabel(label);
  return CONFUSABLE_TOP_LABELS.some((pattern) => pattern.test(normalized));
}

function decideFromHotDogScore(
  hotDogScore: number,
  topLabel: string,
  runnerUpScore: number,
): boolean {
  if (hotDogScore < 0.18) return false;

  const margin = hotDogScore - runnerUpScore;
  const topIsConfusable = isConfusableLabel(topLabel) && !isHotDogLabel(topLabel);

  if (hotDogScore >= 0.45) return true;
  if (hotDogScore >= 0.28 && margin >= 0.05) return true;
  if (hotDogScore >= 0.22 && !topIsConfusable) return true;

  return false;
}

async function prepareImageUri(uri: string, quality: number): Promise<string> {
  const resized = await ImageManipulator.manipulateAsync(
    uri,
    [{ resize: { width: MODEL_INPUT_SIZE } }],
    {
      compress: quality,
      format: ImageManipulator.SaveFormat.JPEG,
    },
  );
  return resized.uri;
}

async function initMobileNet(): Promise<void> {
  await tf.ready();
  await tf.setBackend('rn-webgl');
  mobileNetModel = await mobilenet.load({ version: 2, alpha: 1.0 });

  const warmupTensor = tf.zeros<tf.Rank.R3>([
    MODEL_INPUT_SIZE,
    MODEL_INPUT_SIZE,
    3,
  ]);
  try {
    await mobileNetModel.classify(warmupTensor, 1);
  } finally {
    warmupTensor.dispose();
  }
}

export async function initClassifier(): Promise<void> {
  if (initPromise) return initPromise;

  initPromise = (async () => {
    if (useVision) {
      return;
    }
    await initMobileNet();
  })();

  return initPromise;
}

export function isClassifierReady(): boolean {
  return useVision || mobileNetModel !== null;
}

async function classifyWithAppleVision(
  uri: string,
): Promise<ClassificationResult> {
  const result = await classifyImageWithVision(uri, {
    minimumConfidence: 0.12,
    maxResults: 15,
    iosUseMlKit: true,
  });

  if (!result.success || result.labels.length === 0) {
    throw new Error(result.error ?? 'Vision classification failed');
  }

  const sorted = [...result.labels].sort((a, b) => b.confidence - a.confidence);
  const hotDogHits = sorted.filter((label) => isHotDogLabel(label.identifier));
  const hotDogScore = hotDogHits[0]?.confidence ?? 0;
  const top = sorted[0];
  const runnerUp = sorted.find((label) => !isHotDogLabel(label.identifier));

  return {
    isHotDog: decideFromHotDogScore(
      hotDogScore,
      top.identifier,
      runnerUp?.confidence ?? 0,
    ),
    confidence: hotDogScore,
    topLabel: top.identifier,
  };
}

async function classifyWithMobileNet(
  uri: string,
): Promise<ClassificationResult> {
  if (!mobileNetModel) {
    throw new Error('Classifier not initialized');
  }

  const preparedUri = await prepareImageUri(uri, 0.85);
  const response = await fetch(preparedUri, {}, { isBinary: true });
  const rawImageData = await response.arrayBuffer();
  const imageTensor = decodeJpeg(new Uint8Array(rawImageData));

  try {
    const predictions = await mobileNetModel.classify(imageTensor, 20);
    const hotDogPrediction = predictions.find(
      (p) => p.className === HOT_DOG_IMAGENET_LABEL,
    );
    const hotDogScore = hotDogPrediction?.probability ?? 0;
    const top = predictions[0];
    const runnerUp = predictions.find(
      (p) => p.className !== HOT_DOG_IMAGENET_LABEL,
    );

    return {
      isHotDog: decideFromHotDogScore(
        hotDogScore,
        top?.className ?? 'unknown',
        runnerUp?.probability ?? 0,
      ),
      confidence: hotDogScore,
      topLabel: top?.className ?? 'unknown',
    };
  } finally {
    imageTensor.dispose();
  }
}

export async function classifyImageUri(
  uri: string,
): Promise<ClassificationResult> {
  if (useVision) {
    try {
      return await classifyWithAppleVision(uri);
    } catch {
      useVision = false;
      await initMobileNet();
    }
  }

  return classifyWithMobileNet(uri);
}