Upload phone screen classifier

README.md

@@ -110,6 +110,7 @@ The exported model accepts dynamic batch, height, and width.
 The helpers intentionally return only `screen` and `safety`.
 
 Keep `onnx/model.onnx.data` beside `onnx/model.onnx`; ONNX Runtime loads the external tensor data when it opens the graph.
+By default the helpers load the FP32 model at `onnx/model.onnx`. Pass `fp16` to load `onnx/model.fp16.onnx`.
 
 ### Download
 
@@ -135,15 +136,19 @@ Import the helper from the exported model folder:
 from inference.python import Classifier, classify
 
 print(classify("example.png"))
+print(classify("example.png", model_format="fp16"))
 
 classifier = Classifier(model_dir)
 print(classifier.classify_batch(["one.png", "two.png"]))
+
+fp16_classifier = Classifier(model_dir, model_format="fp16")
+print(fp16_classifier.classify_batch(["one.png", "two.png"]))
 ```
 
 Or run it directly:
 
 ```bash
-python inference/python.py example.png another.png
+python inference/python.py example.png another.png --model-format fp16
 ```
 
 ### TypeScript
@@ -160,15 +165,19 @@ Import the helper from the exported model folder:
 import { Classifier, classify } from './inference/typescript.ts';
 
 console.log(await classify('example.png'));
+console.log(await classify('example.png', modelDir, 'fp16'));
 
 const classifier = await Classifier.create(modelDir);
 console.log(await classifier.classifyBatch(['one.png', 'two.png']));
+
+const fp16Classifier = await Classifier.create(modelDir, 'fp16');
+console.log(await fp16Classifier.classifyBatch(['one.png', 'two.png']));
 ```
 
 Or run it directly:
 
 ```bash
-bun inference/typescript.ts example.png another.png
+bun inference/typescript.ts example.png another.png --model-format fp16
 ```
 
 ## Training Data

inference/python.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import json
 from pathlib import Path
-from typing import Sequence, TypedDict, cast
+from typing import Literal, Sequence, TypedDict, cast
 
 import numpy as np
 import onnxruntime as ort
@@ -11,6 +11,11 @@ from PIL import Image, ImageOps
 
 MODEL_DIR = Path(__file__).resolve().parent.parent
 PADDING_MULTIPLE = 32
+ModelFormat = Literal["fp32", "fp16"]
+MODEL_FILENAMES: dict[ModelFormat, str] = {
+    "fp32": "model.onnx",
+    "fp16": "model.fp16.onnx",
+}
 
 
 class Preprocess(TypedDict):
@@ -38,9 +43,10 @@ class Classifier:
         self,
         model_dir: str | Path = MODEL_DIR,
         providers: Sequence[str] = ("CPUExecutionProvider",),
+        model_format: ModelFormat = "fp32",
     ) -> None:
         self.directory = Path(model_dir)
-        self.session = ort.InferenceSession(str(self.directory / "onnx" / "model.onnx"), providers=list(providers))
+        self.session = ort.InferenceSession(str(model_path(self.directory, model_format)), providers=list(providers))
         self.preprocess = load_preprocess(self.directory / "preprocess.json")
         self.labels = load_labels(self.directory / "inference" / "labels.json")
 
@@ -55,12 +61,27 @@ class Classifier:
         return decode_predictions(screen_logits, safety_logits, self.labels)
 
 
-def classify(image_path: str | Path, model_dir: str | Path = MODEL_DIR) -> Prediction:
-    return Classifier(model_dir).classify(image_path)
+def classify(
+    image_path: str | Path,
+    model_dir: str | Path = MODEL_DIR,
+    model_format: ModelFormat = "fp32",
+) -> Prediction:
+    return Classifier(model_dir, model_format=model_format).classify(image_path)
 
 
-def classify_batch(image_paths: Sequence[str | Path], model_dir: str | Path = MODEL_DIR) -> list[Prediction]:
-    return Classifier(model_dir).classify_batch(image_paths)
+def classify_batch(
+    image_paths: Sequence[str | Path],
+    model_dir: str | Path = MODEL_DIR,
+    model_format: ModelFormat = "fp32",
+) -> list[Prediction]:
+    return Classifier(model_dir, model_format=model_format).classify_batch(image_paths)
+
+
+def model_path(directory: Path, model_format: ModelFormat) -> Path:
+    path = directory / "onnx" / MODEL_FILENAMES[model_format]
+    if not path.is_file():
+        raise FileNotFoundError(f"Missing {model_format} ONNX model: {path}")
+    return path
 
 
 def preprocess_image(image_path: Path, preprocess: Preprocess) -> np.ndarray:
@@ -138,7 +159,8 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Classify images with the exported ONNX model.")
     parser.add_argument("images", nargs="+")
     parser.add_argument("--model-dir", default=str(MODEL_DIR))
+    parser.add_argument("--model-format", choices=tuple(MODEL_FILENAMES), default="fp32")
     args = parser.parse_args()
-    predictions = classify_batch(args.images, args.model_dir)
+    predictions = classify_batch(args.images, args.model_dir, cast(ModelFormat, args.model_format))
     value: Prediction | list[Prediction] = predictions[0] if len(predictions) == 1 else predictions
     print(json.dumps(value, indent=2))

inference/typescript.ts

@@ -6,6 +6,11 @@ import { fileURLToPath, pathToFileURL } from 'node:url';
 
 const MODEL_DIR = resolve(dirname(fileURLToPath(import.meta.url)), '..');
 const PADDING_MULTIPLE = 32;
+export type ModelFormat = 'fp32' | 'fp16';
+const MODEL_FILES: Record<ModelFormat, string> = {
+    fp32: 'model.onnx',
+    fp16: 'model.fp16.onnx',
+};
 
 export class Classifier {
     private constructor(
@@ -27,8 +32,8 @@ export class Classifier {
      * @param modelDir Directory containing the exported model artifacts.
      * @returns Ready-to-use classifier instance.
      */
-    static async create(modelDir = MODEL_DIR): Promise<Classifier> {
-        const session = await ort.InferenceSession.create(resolve(modelDir, 'onnx/model.onnx'), {
+    static async create(modelDir = MODEL_DIR, modelFormat: ModelFormat = 'fp32'): Promise<Classifier> {
+        const session = await ort.InferenceSession.create(resolve(modelDir, 'onnx', MODEL_FILES[modelFormat]), {
             executionProviders: ['cpu'],
         });
         const preprocess = JSON.parse(await readFile(resolve(modelDir, 'preprocess.json'), 'utf8')) as {
@@ -80,8 +85,12 @@ export class Classifier {
  * @param modelDir Directory containing the exported model artifacts.
  * @returns Predicted public screen and safety labels.
  */
-export async function classify(imagePath: string, modelDir = MODEL_DIR): Promise<{ screen: string; safety: string }> {
-    const classifier = await Classifier.create(modelDir);
+export async function classify(
+    imagePath: string,
+    modelDir = MODEL_DIR,
+    modelFormat: ModelFormat = 'fp32',
+): Promise<{ screen: string; safety: string }> {
+    const classifier = await Classifier.create(modelDir, modelFormat);
     if (imagePath.length === 0) {
         throw new Error('Image path must not be empty.');
     }
@@ -98,8 +107,9 @@ export async function classify(imagePath: string, modelDir = MODEL_DIR): Promise
 export async function classifyBatch(
     imagePaths: readonly string[],
     modelDir = MODEL_DIR,
+    modelFormat: ModelFormat = 'fp32',
 ): Promise<Array<{ screen: string; safety: string }>> {
-    const classifier = await Classifier.create(modelDir);
+    const classifier = await Classifier.create(modelDir, modelFormat);
     if (imagePaths.length === 0) {
         return [];
     }
@@ -275,9 +285,21 @@ if (script !== undefined && import.meta.url === pathToFileURL(resolve(script)).h
     if (modelFlag >= 0) {
         args.splice(modelFlag, 2);
     }
+    const formatFlag = args.indexOf('--model-format');
+    const modelFormat = formatFlag >= 0 ? parseModelFormat(args[formatFlag + 1]) : 'fp32';
+    if (formatFlag >= 0) {
+        args.splice(formatFlag, 2);
+    }
     if (args.length === 0 || modelDir === undefined) {
-        throw new Error('Usage: bun inference/typescript.ts <image...> [--model-dir export_dir]');
+        throw new Error('Usage: bun inference/typescript.ts <image...> [--model-dir export_dir] [--model-format fp32|fp16]');
     }
-    const predictions = await classifyBatch(args, modelDir);
+    const predictions = await classifyBatch(args, modelDir, modelFormat);
     process.stdout.write(`${JSON.stringify(predictions.length === 1 ? predictions[0] : predictions, null, 2)}\n`);
 }
+
+function parseModelFormat(value: string | undefined): ModelFormat {
+    if (value === 'fp32' || value === 'fp16') {
+        return value;
+    }
+    throw new Error('Model format must be fp32 or fp16.');
+}