Upload phone screen classifier

README.md

@@ -0,0 +1,201 @@
+---
+model_name: Yap Phone Screenshot Classifier
+library_name: 'timm'
+base_model: 'timm/mobilenetv4_conv_medium.e250_r384_in12k'
+tags:
+    - image-classification
+    - mobile-screenshots
+    - phone-screenshots
+    - screenshot-analysis
+    - content-safety
+    - timm
+datasets:
+    - yapwithai/phone-screenshots
+---
+
+# Yap Phone Screenshot Classifier
+
+Yap Phone Screenshot Classifier predicts two labels for one image:
+
+1. `screen`: the phone-screen category.
+2. `safety`: the content-safety category.
+
+The model returns exactly two outputs: `screen` and `safety`.
+
+The model is fine-tuned from `timm/mobilenetv4_conv_medium.e250_r384_in12k` on the Yap phone screenshot dataset: [`yapwithai/phone-screenshots`](https://huggingface.co/datasets/yapwithai/phone-screenshots).
+
+## What It Is For
+
+This model is built for routing and filtering mobile screenshot workflows. It separates UI classification from safety classification so an application can answer both of these questions before running more expensive analysis:
+
+- What kind of screen is this?
+- Which safety category does this image belong to?
+
+`generic` and `other` are real screen labels when included in the exported label map. `generic` means phone UI that is not one of the more specific trained screen classes. `other` means non-phone UI or images outside the phone-screenshot distribution.
+
+## Files
+
+| File                      | Purpose                                                        |
+| :------------------------ | :------------------------------------------------------------- |
+| `onnx/model.onnx`         | ONNX model for CPU/server inference.                           |
+| `onnx/model.fp16.onnx`    | Optional FP16 ONNX candidate.                                  |
+| `onnx/model.int8.onnx`    | Optional INT8 ONNX candidate.                                  |
+| `onnx/model.onnx.data`    | External ONNX weight data loaded beside `model.onnx`.          |
+| `model.safetensors`       | PyTorch state dict for reproducibility and continued training. |
+| `config.json`             | Model identity, base model, output names, and label arrays.    |
+| `preprocess.json`         | Resize and normalization contract used during training/export. |
+| `train.json`              | Sanitized training recipe for post-training or reproduction.   |
+| `inference/python.py`     | Python helper for ONNX inference from image paths.             |
+| `inference/typescript.ts` | TypeScript helper for ONNX inference from image paths.         |
+| `inference/labels.json`   | Screen and safety labels used to decode model logits.          |
+| `README.md`               | This model card and runnable inference reference.              |
+
+## Outputs
+
+The ONNX graph has two outputs in this exact order:
+
+| Output   | Shape                         | Meaning                 |
+| :------- | :---------------------------- | :---------------------- |
+| `screen` | `[batch, screen_class_count]` | Screen-category logits. |
+| `safety` | `[batch, safety_class_count]` | Content-safety logits.  |
+
+Inference uses `argmax` for both heads in this version.
+
+## Test Results
+
+| Model | Test | Output | Accuracy | Balanced Acc | Macro F1 | Top-2 | Images |
+| :-- | --: | --: | --: | --: | --: | --: | --: |
+| fp32 | full test | screen | 0.9154 | 0.8385 | 0.7265 | 0.9719 | 23615 |
+| fp32 | full test | safety | 0.9575 | 0.8941 | 0.8600 | 0.9921 | 23615 |
+| fp32 | screen-balanced test | screen | 0.8377 | 0.8421 | 0.8133 | 0.9400 | 6246 |
+| fp32 | screen-balanced test | safety | 0.9776 | 0.9285 | 0.9110 | 0.9978 | 6246 |
+| fp32 | safety-balanced test | screen | 0.9580 | 0.7752 | 0.6547 | 0.9870 | 3000 |
+| fp32 | safety-balanced test | safety | 0.8957 | 0.8957 | 0.8947 | 0.9847 | 3000 |
+| fp16 | full test | screen | 0.9153 | 0.8384 | 0.7264 | 0.9719 | 23615 |
+| fp16 | full test | safety | 0.9575 | 0.8941 | 0.8600 | 0.9921 | 23615 |
+| fp16 | screen-balanced test | screen | 0.8373 | 0.8419 | 0.8130 | 0.9400 | 6246 |
+| fp16 | screen-balanced test | safety | 0.9776 | 0.9285 | 0.9110 | 0.9978 | 6246 |
+| fp16 | safety-balanced test | screen | 0.9580 | 0.7752 | 0.6547 | 0.9870 | 3000 |
+| fp16 | safety-balanced test | safety | 0.8957 | 0.8957 | 0.8947 | 0.9847 | 3000 |
+| int8 | full test | screen | 0.0018 | 0.0424 | 0.0024 | 0.0278 | 23615 |
+| int8 | full test | safety | 0.1013 | 0.3502 | 0.0696 | 0.9431 | 23615 |
+| int8 | screen-balanced test | screen | 0.0062 | 0.0424 | 0.0031 | 0.1039 | 6246 |
+| int8 | screen-balanced test | safety | 0.1703 | 0.3686 | 0.1153 | 0.9241 | 6246 |
+| int8 | safety-balanced test | screen | 0.0003 | 0.0417 | 0.0000 | 0.0140 | 3000 |
+| int8 | safety-balanced test | safety | 0.3493 | 0.3493 | 0.1992 | 0.6667 | 3000 |
+
+## CPU Timing
+
+| Model | Images/s | Load Mean | Model Mean | Total Mean | Total Median | Total P95 | Provider |
+| :-- | --: | --: | --: | --: | --: | --: | --: |
+| fp32 | 20.0538 | 0.12 ms | 49.52 ms | 49.87 ms | 57.84 ms | 85.82 ms | onnxruntime:CPUExecutionProvider |
+| fp16 | 20.8351 | 0.11 ms | 47.69 ms | 48.00 ms | 48.60 ms | 63.65 ms | onnxruntime:CPUExecutionProvider |
+| int8 | 72.6504 | 0.14 ms | 13.36 ms | 13.76 ms | 14.07 ms | 18.01 ms | onnxruntime:CPUExecutionProvider |
+
+Timing is measured with ONNX Runtime CPU execution on `Apple M4 Max (16 logical cores)`. Total latency includes image load/preprocess, model inference, and label decoding.
+
+## Classification Labels
+
+For the screen and safety label lists, see [`config.json`](config.json).
+
+## Preprocessing
+
+Use the preprocessing contract in `preprocess.json`.
+
+1. Read the image.
+2. Apply EXIF orientation.
+3. Convert to RGB.
+4. Resize so the longest side is `1024` pixels and keep the original aspect ratio.
+5. Do not crop, stretch, square-pad before preprocessing, or horizontally flip.
+6. Convert to `float32` in `[0, 1]`.
+7. Normalize with the mean and standard deviation from `preprocess.json`.
+8. Pad each normalized tensor to the batch maximum height and width, rounded up to a multiple of 32.
+
+The exported model accepts dynamic batch, height, and width.
+
+## ONNX Inference
+
+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.
+
+### Download
+
+Download the exported model folder from Hugging Face:
+
+```python
+from huggingface_hub import snapshot_download
+
+model_dir = snapshot_download("yapwithai/phone-screen-classifier")
+```
+
+### Python
+
+Install the runtime dependencies:
+
+```bash
+python -m pip install numpy pillow onnxruntime
+```
+
+Import the helper from the exported model folder:
+
+```python
+from inference.python import Classifier, classify
+
+print(classify("example.png"))
+
+classifier = Classifier(model_dir)
+print(classifier.classify_batch(["one.png", "two.png"]))
+```
+
+Or run it directly:
+
+```bash
+python inference/python.py example.png another.png
+```
+
+### TypeScript
+
+Install the runtime dependencies:
+
+```bash
+bun add sharp onnxruntime-node
+```
+
+Import the helper from the exported model folder:
+
+```ts
+import { Classifier, classify } from './inference/typescript.ts';
+
+console.log(await classify('example.png'));
+
+const classifier = await Classifier.create(modelDir);
+console.log(await classifier.classifyBatch(['one.png', 'two.png']));
+```
+
+Or run it directly:
+
+```bash
+bun inference/typescript.ts example.png another.png
+```
+
+## Training Data
+
+The model is trained from local shards of [`yapwithai/phone-screenshots`](https://huggingface.co/datasets/yapwithai/phone-screenshots). Screen labels come from the dataset `label` field. Safety labels come from the dataset `safety` field. `fine_label` is not used.
+
+Training uses one minimum train-count floor for both heads. The default is `--min-train-count 300`. Labels with fewer than `--min-train-count` train samples are not trained directly. Screen labels below the floor are folded into `generic`; safety labels below the floor are omitted from the safety head. The exported `config.json` and `inference/labels.json` files are the source of truth for the active screen and safety labels.
+
+Training uses adaptive label sampling. `--screen-target-ratio` and `--safety-target-ratio` set the desired effective per-epoch class ratio; `--screen-max-repeat` and `--safety-max-repeat` cap how often smaller classes can repeat in one epoch.
+
+## Citation
+
+If you use this model, please cite:
+
+```bibtex
+@misc{phone-screen-classifier,
+  title={Yap Phone Screenshot Classifier},
+  year={2026},
+  publisher={Yap With AI},
+  url={https://huggingface.co/yapwithai/phone-screen-classifier}
+}
+```

checkpoints/checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7e22dfcf02e0993c23ea02f729d62190e85ff9e8ae5cf01046c0b78942c94e09
+size 104626873

config.json

@@ -0,0 +1,41 @@
+{
+  "architecture": "timm-spatial-flat-multitask",
+  "model": "timm/mobilenetv4_conv_medium.e250_r384_in12k",
+  "name": "Yap Phone Screenshot Classifier",
+  "outputs": [
+    "screen",
+    "safety"
+  ],
+  "resize_longest_side_px": 1024,
+  "safety_labels": [
+    "hot",
+    "nsfw",
+    "safe"
+  ],
+  "screen_labels": [
+    "audio",
+    "calendar",
+    "camera",
+    "comments",
+    "contacts",
+    "email",
+    "excel",
+    "feed",
+    "finances",
+    "food",
+    "gallery",
+    "game",
+    "generic",
+    "maps",
+    "message-chat",
+    "message-list",
+    "notes",
+    "other",
+    "post",
+    "profile",
+    "reading",
+    "shopping",
+    "streaming",
+    "weather"
+  ]
+}

inference/labels.json

@@ -0,0 +1,68 @@
+{
+  "safety": {
+    "label_to_id": {
+      "hot": 0,
+      "nsfw": 1,
+      "safe": 2
+    },
+    "labels": [
+      "hot",
+      "nsfw",
+      "safe"
+    ]
+  },
+  "screen": {
+    "label_to_id": {
+      "audio": 0,
+      "calendar": 1,
+      "camera": 2,
+      "comments": 3,
+      "contacts": 4,
+      "email": 5,
+      "excel": 6,
+      "feed": 7,
+      "finances": 8,
+      "food": 9,
+      "gallery": 10,
+      "game": 11,
+      "generic": 12,
+      "maps": 13,
+      "message-chat": 14,
+      "message-list": 15,
+      "notes": 16,
+      "other": 17,
+      "post": 18,
+      "profile": 19,
+      "reading": 20,
+      "shopping": 21,
+      "streaming": 22,
+      "weather": 23
+    },
+    "labels": [
+      "audio",
+      "calendar",
+      "camera",
+      "comments",
+      "contacts",
+      "email",
+      "excel",
+      "feed",
+      "finances",
+      "food",
+      "gallery",
+      "game",
+      "generic",
+      "maps",
+      "message-chat",
+      "message-list",
+      "notes",
+      "other",
+      "post",
+      "profile",
+      "reading",
+      "shopping",
+      "streaming",
+      "weather"
+    ]
+  }
+}

inference/python.py

@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Sequence, TypedDict, cast
+
+import numpy as np
+import onnxruntime as ort
+from PIL import Image, ImageOps
+
+
+MODEL_DIR = Path(__file__).resolve().parent.parent
+PADDING_MULTIPLE = 32
+
+
+class Preprocess(TypedDict):
+    resize_longest_side_px: int
+    mean: list[float]
+    std: list[float]
+
+
+class LabelSet(TypedDict):
+    labels: list[str]
+
+
+class Labels(TypedDict):
+    screen: LabelSet
+    safety: LabelSet
+
+
+class Prediction(TypedDict):
+    screen: str
+    safety: str
+
+
+class Classifier:
+    def __init__(
+        self,
+        model_dir: str | Path = MODEL_DIR,
+        providers: Sequence[str] = ("CPUExecutionProvider",),
+    ) -> None:
+        self.directory = Path(model_dir)
+        self.session = ort.InferenceSession(str(self.directory / "onnx" / "model.onnx"), providers=list(providers))
+        self.preprocess = load_preprocess(self.directory / "preprocess.json")
+        self.labels = load_labels(self.directory / "inference" / "labels.json")
+
+    def classify(self, image_path: str | Path) -> Prediction:
+        return self.classify_batch([image_path])[0]
+
+    def classify_batch(self, image_paths: Sequence[str | Path]) -> list[Prediction]:
+        if not image_paths:
+            return []
+        images = [preprocess_image(Path(image_path), self.preprocess) for image_path in image_paths]
+        screen_logits, safety_logits = self.session.run(None, {"image": collate_images(images)})
+        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_batch(image_paths: Sequence[str | Path], model_dir: str | Path = MODEL_DIR) -> list[Prediction]:
+    return Classifier(model_dir).classify_batch(image_paths)
+
+
+def preprocess_image(image_path: Path, preprocess: Preprocess) -> np.ndarray:
+    with Image.open(image_path) as opened:
+        image = to_training_rgb(opened)
+    resized = resize_image(image, preprocess["resize_longest_side_px"])
+    array = np.asarray(resized).astype("float32") / 255.0
+    mean = np.asarray(preprocess["mean"], dtype="float32")
+    std = np.asarray(preprocess["std"], dtype="float32")
+    array = (array - mean) / std
+    return np.transpose(array, (2, 0, 1))
+
+
+def to_training_rgb(image: Image.Image) -> Image.Image:
+    image = ImageOps.exif_transpose(image)
+    if image.mode == "P" and isinstance(image.info.get("transparency"), bytes):
+        image = image.convert("RGBA")
+    if image.mode in ("RGBA", "LA", "PA"):
+        rgba = image.convert("RGBA")
+        background = Image.new("RGBA", rgba.size, (255, 255, 255, 255))
+        image = Image.alpha_composite(background, rgba)
+    return image.convert("RGB")
+
+
+def collate_images(images: Sequence[np.ndarray]) -> np.ndarray:
+    height = round_up(max(int(image.shape[1]) for image in images))
+    width = round_up(max(int(image.shape[2]) for image in images))
+    batch = np.zeros((len(images), 3, height, width), dtype="float32")
+    for index, image in enumerate(images):
+        image_height = int(image.shape[1])
+        image_width = int(image.shape[2])
+        batch[index, :, :image_height, :image_width] = image
+    return batch
+
+
+def resize_image(image: Image.Image, image_size: int) -> Image.Image:
+    scale = image_size / max(image.width, image.height)
+    width = max(1, round(image.width * scale))
+    height = max(1, round(image.height * scale))
+    return image.resize((width, height), Image.Resampling.BICUBIC)
+
+
+def round_up(value: int, multiple: int = PADDING_MULTIPLE) -> int:
+    return ((value + multiple - 1) // multiple) * multiple
+
+
+def decode_predictions(screen_logits: np.ndarray, safety_logits: np.ndarray, labels: Labels) -> list[Prediction]:
+    # ONNX emits flat screen and safety logits.
+    screen_indices = top_indices(screen_logits)
+    safety_indices = top_indices(safety_logits)
+    return [
+        {
+            "screen": labels["screen"]["labels"][screen_index],
+            "safety": labels["safety"]["labels"][safety_index],
+        }
+        for screen_index, safety_index in zip(screen_indices, safety_indices, strict=True)
+    ]
+
+
+def load_labels(path: Path) -> Labels:
+    return cast(Labels, json.loads(path.read_text(encoding="utf-8")))
+
+
+def load_preprocess(path: Path) -> Preprocess:
+    return cast(Preprocess, json.loads(path.read_text(encoding="utf-8")))
+
+
+def top_indices(logits: np.ndarray) -> list[int]:
+    return [int(index) for index in np.argmax(logits, axis=1)]
+
+
+if __name__ == "__main__":
+    import argparse
+
+    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))
+    args = parser.parse_args()
+    predictions = classify_batch(args.images, args.model_dir)
+    value: Prediction | list[Prediction] = predictions[0] if len(predictions) == 1 else predictions
+    print(json.dumps(value, indent=2))

inference/typescript.ts

@@ -0,0 +1,283 @@
+import sharp from 'sharp';
+import * as ort from 'onnxruntime-node';
+import { readFile } from 'node:fs/promises';
+import { dirname, resolve } from 'node:path';
+import { fileURLToPath, pathToFileURL } from 'node:url';
+
+const MODEL_DIR = resolve(dirname(fileURLToPath(import.meta.url)), '..');
+const PADDING_MULTIPLE = 32;
+
+export class Classifier {
+    private constructor(
+        private readonly session: ort.InferenceSession,
+        private readonly preprocess: {
+            resize_longest_side_px: number;
+            mean: [number, number, number];
+            std: [number, number, number];
+        },
+        private readonly labels: {
+            screen: { labels: string[] };
+            safety: { labels: string[] };
+        },
+    ) {}
+
+    /**
+     * Load the exported ONNX classifier and sidecar metadata.
+     *
+     * @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'), {
+            executionProviders: ['cpu'],
+        });
+        const preprocess = JSON.parse(await readFile(resolve(modelDir, 'preprocess.json'), 'utf8')) as {
+            resize_longest_side_px: number;
+            mean: [number, number, number];
+            std: [number, number, number];
+        };
+        const labels = JSON.parse(await readFile(resolve(modelDir, 'inference/labels.json'), 'utf8')) as {
+            screen: { labels: string[] };
+            safety: { labels: string[] };
+        };
+        return new Classifier(session, preprocess, labels);
+    }
+
+    /**
+     * Classify one image.
+     *
+     * @param imagePath Path to an image file.
+     * @returns Predicted public screen and safety labels.
+     */
+    async classify(imagePath: string): Promise<{ screen: string; safety: string }> {
+        const [prediction] = await this.classifyBatch([imagePath]);
+        if (prediction === undefined) {
+            throw new Error('No prediction was produced.');
+        }
+        return prediction;
+    }
+
+    /**
+     * Classify a batch of images.
+     *
+     * @param imagePaths Paths to image files.
+     * @returns Predicted public screen and safety labels for each image.
+     */
+    async classifyBatch(imagePaths: readonly string[]): Promise<Array<{ screen: string; safety: string }>> {
+        if (imagePaths.length === 0) {
+            return [];
+        }
+        const images = await Promise.all(imagePaths.map((imagePath) => preprocessImage(imagePath, this.preprocess)));
+        const output = await this.session.run({ image: collateImages(images) });
+        return decodePredictions(outputTensor(output, 'screen'), outputTensor(output, 'safety'), this.labels);
+    }
+}
+
+/**
+ * Classify one image with a classifier loaded from disk.
+ *
+ * @param imagePath Path to an image file.
+ * @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);
+    if (imagePath.length === 0) {
+        throw new Error('Image path must not be empty.');
+    }
+    return classifier.classify(imagePath);
+}
+
+/**
+ * Classify a batch of images with a classifier loaded from disk.
+ *
+ * @param imagePaths Paths to image files.
+ * @param modelDir Directory containing the exported model artifacts.
+ * @returns Predicted public screen and safety labels for each image.
+ */
+export async function classifyBatch(
+    imagePaths: readonly string[],
+    modelDir = MODEL_DIR,
+): Promise<Array<{ screen: string; safety: string }>> {
+    const classifier = await Classifier.create(modelDir);
+    if (imagePaths.length === 0) {
+        return [];
+    }
+    return classifier.classifyBatch(imagePaths);
+}
+
+async function preprocessImage(
+    imagePath: string,
+    preprocess: {
+        resize_longest_side_px: number;
+        mean: [number, number, number];
+        std: [number, number, number];
+    },
+): Promise<{ data: Float32Array; height: number; width: number }> {
+    const { data, info: metadata } = await sharp(imagePath)
+        .rotate()
+        .resize({
+            width: preprocess.resize_longest_side_px,
+            height: preprocess.resize_longest_side_px,
+            fit: 'inside',
+            kernel: 'cubic',
+        })
+        .removeAlpha()
+        .toColourspace('srgb')
+        .raw()
+        .toBuffer({ resolveWithObject: true });
+
+    if (metadata.channels !== 3) {
+        throw new Error(`Expected RGB image data, got ${String(metadata.channels)} channels.`);
+    }
+
+    const pixels = metadata.width * metadata.height;
+    const tensor = new Float32Array(3 * pixels);
+
+    for (let pixel = 0; pixel < pixels; pixel += 1) {
+        const offset = pixel * 3;
+        tensor[pixel] = (channel(data, offset) / 255 - preprocess.mean[0]) / preprocess.std[0];
+        tensor[pixels + pixel] = (channel(data, offset + 1) / 255 - preprocess.mean[1]) / preprocess.std[1];
+        tensor[2 * pixels + pixel] = (channel(data, offset + 2) / 255 - preprocess.mean[2]) / preprocess.std[2];
+    }
+
+    return {
+        data: tensor,
+        height: metadata.height,
+        width: metadata.width,
+    };
+}
+
+function collateImages(
+    images: ReadonlyArray<{ data: Float32Array; height: number; width: number }>,
+): ort.TypedTensor<'float32'> {
+    const height = Math.ceil(Math.max(...images.map((image) => image.height)) / PADDING_MULTIPLE) * PADDING_MULTIPLE;
+    const width = Math.ceil(Math.max(...images.map((image) => image.width)) / PADDING_MULTIPLE) * PADDING_MULTIPLE;
+    const imagePixels = height * width;
+    const tensor = new Float32Array(images.length * 3 * imagePixels);
+
+    for (let imageIndex = 0; imageIndex < images.length; imageIndex += 1) {
+        const image = imageAt(images, imageIndex);
+        const sourcePixels = image.height * image.width;
+        for (let channelIndex = 0; channelIndex < 3; channelIndex += 1) {
+            for (let row = 0; row < image.height; row += 1) {
+                const sourceStart = channelIndex * sourcePixels + row * image.width;
+                const targetStart = imageIndex * 3 * imagePixels + channelIndex * imagePixels + row * width;
+                tensor.set(image.data.subarray(sourceStart, sourceStart + image.width), targetStart);
+            }
+        }
+    }
+
+    return new ort.Tensor('float32', tensor, [images.length, 3, height, width]);
+}
+
+function decodePredictions(
+    screen: ort.TypedTensor<'float32'>,
+    safety: ort.TypedTensor<'float32'>,
+    labels: {
+        screen: { labels: string[] };
+        safety: { labels: string[] };
+    },
+): Array<{ screen: string; safety: string }> {
+    // ONNX emits flat screen and safety logits.
+    const batchSize = tensorDim(screen, 0);
+    const screenClassCount = tensorDim(screen, 1);
+    const safetyBatchSize = tensorDim(safety, 0);
+    const safetyClassCount = tensorDim(safety, 1);
+    if (safetyBatchSize !== batchSize) {
+        throw new Error(
+            `Safety batch size ${String(safetyBatchSize)} does not match screen batch size ${String(batchSize)}.`,
+        );
+    }
+
+    const predictions: Array<{ screen: string; safety: string }> = [];
+    for (let row = 0; row < batchSize; row += 1) {
+        predictions.push({
+            screen: labelAt(labels.screen.labels, topIndex(screen.data, row * screenClassCount, screenClassCount)),
+            safety: labelAt(labels.safety.labels, topIndex(safety.data, row * safetyClassCount, safetyClassCount)),
+        });
+    }
+    return predictions;
+}
+
+function outputTensor(output: ort.InferenceSession.ReturnType, name: string): ort.TypedTensor<'float32'> {
+    const tensor = output[name];
+    if (tensor === undefined) {
+        throw new Error(`Missing ONNX output "${name}".`);
+    }
+    if (tensor.type !== 'float32') {
+        throw new Error(`Expected "${name}" to be float32, got ${tensor.type}.`);
+    }
+    return tensor as ort.TypedTensor<'float32'>;
+}
+
+function tensorDim(tensor: ort.TypedTensor<'float32'>, index: number): number {
+    const dimension = tensor.dims[index];
+    if (typeof dimension !== 'number') {
+        throw new Error(`Missing tensor dimension ${String(index)}.`);
+    }
+    return dimension;
+}
+
+function labelAt(labels: readonly string[], index: number): string {
+    const label = labels[index];
+    if (label === undefined) {
+        throw new Error(`Missing label for class ${String(index)}.`);
+    }
+    return label;
+}
+
+function topIndex(data: Float32Array, offset: number, count: number): number {
+    let top = 0;
+    let score = valueAt(data, offset);
+    for (let index = 1; index < count; index += 1) {
+        const candidate = valueAt(data, offset + index);
+        if (candidate > score) {
+            top = index;
+            score = candidate;
+        }
+    }
+    return top;
+}
+
+function imageAt(
+    images: ReadonlyArray<{ data: Float32Array; height: number; width: number }>,
+    index: number,
+): { data: Float32Array; height: number; width: number } {
+    const image = images[index];
+    if (image === undefined) {
+        throw new Error(`Missing image at index ${String(index)}.`);
+    }
+    return image;
+}
+
+function valueAt(data: Float32Array, index: number): number {
+    const value = data[index];
+    if (value === undefined) {
+        throw new Error(`Missing tensor value at offset ${String(index)}.`);
+    }
+    return value;
+}
+
+function channel(data: Buffer, index: number): number {
+    const channelByte = data[index];
+    if (channelByte === undefined) {
+        throw new Error(`Missing image channel at offset ${String(index)}.`);
+    }
+    return channelByte;
+}
+
+const script = process.argv[1];
+if (script !== undefined && import.meta.url === pathToFileURL(resolve(script)).href) {
+    const args = process.argv.slice(2);
+    const modelFlag = args.indexOf('--model-dir');
+    const modelDir = modelFlag >= 0 ? args[modelFlag + 1] : MODEL_DIR;
+    if (modelFlag >= 0) {
+        args.splice(modelFlag, 2);
+    }
+    if (args.length === 0 || modelDir === undefined) {
+        throw new Error('Usage: bun inference/typescript.ts <image...> [--model-dir export_dir]');
+    }
+    const predictions = await classifyBatch(args, modelDir);
+    process.stdout.write(`${JSON.stringify(predictions.length === 1 ? predictions[0] : predictions, null, 2)}\n`);
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8d335efa9d3683c28fa2936ef9fdd75aa493482bb38cf624f39efcc81126ed8e
+size 34197844

onnx/model.fp16.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:93e24db05bb256043c9b9cb64b177f0e1b4514340d8e92dfb24fa739040918bc
+size 17432378

onnx/model.int8.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f1acefd4c5a1204c6c7b1db7bd8e31b59117982136379a60c4dae05d5798e155
+size 9542131

onnx/model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a9a7af13190c85f9ba3794de60cf00986942ecee92fd4ab879532b137c3de6dc
+size 569419

onnx/model.onnx.data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5577fd3b7f9e020044ef91d8ca710310304db065974783a1359be1458f4901c7
+size 33751040

preprocess.json

@@ -0,0 +1,17 @@
+{
+  "crop": false,
+  "horizontal_flip": false,
+  "mean": [
+    0.485,
+    0.456,
+    0.406
+  ],
+  "preserve_aspect_ratio": true,
+  "resize_longest_side_px": 1024,
+  "std": [
+    0.229,
+    0.224,
+    0.225
+  ],
+  "stretch": false
+}

train.json

@@ -0,0 +1,203 @@
+{
+  "augmentation": {
+    "artifact_codecs": [
+      "JPEG",
+      "WEBP"
+    ],
+    "artifact_probability": 0.15,
+    "artifact_quality_max": 98,
+    "artifact_quality_min": 88,
+    "brightness_max": 1.1,
+    "brightness_min": 0.9,
+    "brightness_probability": 0.3,
+    "contrast_max": 1.1,
+    "contrast_min": 0.9,
+    "contrast_probability": 0.3
+  },
+  "counts": {
+    "test_safety": {
+      "hot": 1344,
+      "nsfw": 1559,
+      "safe": 20712
+    },
+    "test_screen": {
+      "audio": 111,
+      "calendar": 101,
+      "camera": 41,
+      "comments": 87,
+      "contacts": 33,
+      "email": 68,
+      "excel": 40,
+      "feed": 381,
+      "finances": 840,
+      "food": 285,
+      "gallery": 66,
+      "game": 330,
+      "generic": 3820,
+      "maps": 129,
+      "message-chat": 207,
+      "message-list": 38,
+      "notes": 59,
+      "other": 15549,
+      "post": 685,
+      "profile": 174,
+      "reading": 109,
+      "shopping": 381,
+      "streaming": 51,
+      "weather": 30
+    },
+    "train_safety": {
+      "hot": 14997,
+      "nsfw": 18086,
+      "safe": 234685
+    },
+    "train_screen": {
+      "audio": 1262,
+      "calendar": 1152,
+      "camera": 477,
+      "comments": 993,
+      "contacts": 374,
+      "email": 781,
+      "excel": 455,
+      "feed": 4323,
+      "finances": 9533,
+      "food": 3238,
+      "gallery": 753,
+      "game": 3748,
+      "generic": 43298,
+      "maps": 1468,
+      "message-chat": 2346,
+      "message-list": 442,
+      "notes": 678,
+      "other": 176223,
+      "post": 7773,
+      "profile": 1975,
+      "reading": 1246,
+      "shopping": 4324,
+      "streaming": 579,
+      "weather": 344
+    },
+    "val_safety": {
+      "hot": 1334,
+      "nsfw": 1696,
+      "safe": 20585
+    },
+    "val_screen": {
+      "audio": 111,
+      "calendar": 101,
+      "camera": 41,
+      "comments": 87,
+      "contacts": 33,
+      "email": 68,
+      "excel": 40,
+      "feed": 381,
+      "finances": 840,
+      "food": 285,
+      "gallery": 66,
+      "game": 330,
+      "generic": 3820,
+      "maps": 129,
+      "message-chat": 207,
+      "message-list": 38,
+      "notes": 59,
+      "other": 15549,
+      "post": 685,
+      "profile": 174,
+      "reading": 109,
+      "shopping": 381,
+      "streaming": 51,
+      "weather": 30
+    }
+  },
+  "dataset": {
+    "path": "output/dataset/phone-screenshots",
+    "repo": "yapwithai/phone-screenshots",
+    "safety_field": "safety",
+    "screen_field": "label"
+  },
+  "evaluation": {
+    "balanced_eval_max_per_class": 1000
+  },
+  "labels": {
+    "safety": [
+      "hot",
+      "nsfw",
+      "safe"
+    ],
+    "screen": [
+      "audio",
+      "calendar",
+      "camera",
+      "comments",
+      "contacts",
+      "email",
+      "excel",
+      "feed",
+      "finances",
+      "food",
+      "gallery",
+      "game",
+      "generic",
+      "maps",
+      "message-chat",
+      "message-list",
+      "notes",
+      "other",
+      "post",
+      "profile",
+      "reading",
+      "shopping",
+      "streaming",
+      "weather"
+    ]
+  },
+  "loss": {
+    "safety": "cross_entropy",
+    "safety_batch_probability": 0.35,
+    "safety_loss_weight": 0.5,
+    "screen": "cross_entropy"
+  },
+  "model": "timm/mobilenetv4_conv_medium.e250_r384_in12k",
+  "optimizer": {
+    "backbone_lr": 5.9999999999999995e-05,
+    "head_lr": 0.0003,
+    "lr": 0.0003,
+    "weight_decay": 0.05
+  },
+  "preprocess": {
+    "crop": false,
+    "horizontal_flip": false,
+    "mean": [
+      0.485,
+      0.456,
+      0.406
+    ],
+    "preserve_aspect_ratio": true,
+    "resize_longest_side_px": 1024,
+    "std": [
+      0.229,
+      0.224,
+      0.225
+    ],
+    "stretch": false
+  },
+  "sampling": {
+    "folded_screen_labels": {},
+    "min_train_count": 300,
+    "mode": "adaptive_label_aspect_bucketed",
+    "safety_max_repeat": 8,
+    "safety_target_ratio": 2.0,
+    "screen_max_repeat": 6,
+    "screen_target_ratio": 5.0
+  },
+  "schema": 3,
+  "training": {
+    "batch_size": 8,
+    "effective_batch_size": 64,
+    "epochs": 4,
+    "grad_accum_steps": 8,
+    "seed": 1337,
+    "target_effective_batch_size": 64,
+    "workers": 4
+  }
+}