initial commit

src/lib/classifier.ts

@@ -0,0 +1,197 @@
+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);
+}