intermediateView Python source

Classify images with an ONNX model

This example labels eight photos with ResNet-18, an image classifier trained on the 1,000 ImageNet categories. It runs the model inside the engine with ONNX Runtime, returns the top three labels per photo, and lists the less confident predictions with the runner-up label, for a person to check.

Run it

Run all commands below from the repository root (sail-examples/):

uv run multimodal/classify-images/main.py

The first run downloads the model (about 47 MB) from Hugging Face into the local Hugging Face cache. No account or key is needed. The model revision is pinned, so every run uses the same weights.

For these eight bundled photos, every top label matches the photo's subject. Five of them score 0.6 or more. The golden retriever, the Siamese cat and the tabby cat score less, so they land in the review table. For the Siamese cat, the runner-up (Persian cat) is close behind.

Output

The newest capture appears first. Older output stays here when a result changes.

Recorded output · Sail 0.7.3

Captured on Sail 0.7.3. Latest capture: 2026-11-05, Python 3.12.14.

The example has changed since this capture. Rerun it to refresh the output.

+---------------------+----------------+-----+
|file                 |label           |score|
+---------------------+----------------+-----+
|espresso.jpg         |espresso        |0.779|
|golden-retriever.jpg |golden retriever|0.493|
|monarch-butterfly.jpg|monarch         |0.994|
|red-fox.jpg          |red fox         |0.954|
|siamese-cat.jpg      |Siamese cat     |0.413|
|strawberries.jpg     |strawberry      |0.995|
|tabby-cat.jpg        |tabby           |0.533|
|zebra.jpg            |zebra           |0.999|
+---------------------+----------------+-----+

Needs review:
+--------------------+----------------+-----+-----------+---------------+
|file                |label           |score|runner_up  |runner_up_score|
+--------------------+----------------+-----+-----------+---------------+
|golden-retriever.jpg|golden retriever|0.493|kuvasz     |0.231          |
|siamese-cat.jpg     |Siamese cat     |0.413|Persian cat|0.379          |
|tabby-cat.jpg       |tabby           |0.533|hyena      |0.155          |
+--------------------+----------------+-----+-----------+---------------+

Python code

main.py148 lines

View main.py on GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148# /// script
# requires-python = ">=3.12"
# dependencies = [
#   "certifi==2026.10.19",
#   "pysail==0.7.3",
#   "pyspark-client==4.1.3",
#   "pyarrow==26.0.0",
#   "pandas==3.0.6",
#   "numpy==2.5.4",
#   "onnxruntime==1.31.1",
#   "huggingface-hub==0.36.0",
#   "pillow==13.0.0",
# ]
# ///
"""Classify a folder of photos with an ImageNet model inside the engine.

The pattern: a table of image paths goes through `mapInArrow`, which loads the ONNX model once
per partition, decodes each image, runs the model on the whole batch, and returns the top three
labels per image. Only labels and scores leave the function, never pixels.
"""

import atexit
import shutil
import tempfile
from pathlib import Path

import pyarrow as pa
from pyspark.sql import DataFrame, SparkSession
from pyspark.sql import functions as F

MODEL_REPO = "Xenova/resnet-18"  # ResNet-18 trained on ImageNet-1k, exported to ONNX
MODEL_REVISION = "549a3007ab1a7a2a4d3af94c287ebbb27f8f65c5"  # pinned: every run uses the same weights
IMAGES = Path(__file__).resolve().parent / "images"
TOP_K = 3
INFERENCE_BATCH_SIZE = 16  # bound model tensors independently of Arrow input batch size
REVIEW_BELOW = 0.6  # predictions less confident than this go to a review table

PREDICTION = pa.struct([("label", pa.string()), ("score", pa.float32())])
ARROW_SCHEMA = pa.schema(
    [("file", pa.string()), ("label", pa.string()), ("score", pa.float32()), ("top3", pa.list_(PREDICTION))]
)
SCHEMA = "file string, label string, score float, top3 array<struct<label: string, score: float>>"


def temp_dir(prefix: str) -> str:
    """A temporary directory that is removed when the process exits."""
    path = tempfile.mkdtemp(prefix=prefix)
    atexit.register(shutil.rmtree, path, ignore_errors=True)
    return path


def preprocess(path: str):
    """The model's own preprocessing: shortest side to 256, center crop 224, ImageNet normalization."""
    import numpy as np
    from PIL import Image, ImageOps

    with Image.open(path) as image:
        image = ImageOps.exif_transpose(image).convert("RGB")
    # The shorter side becomes exactly 256; the longer side is scaled and truncated, as transformers does.
    short, long = sorted(image.size)
    resized = (256, int(256 * long / short))
    image = image.resize(resized if image.width <= image.height else resized[::-1], Image.BICUBIC)
    left, top = (image.width - 224) // 2, (image.height - 224) // 2
    image = image.crop((left, top, left + 224, top + 224))
    pixels = np.asarray(image, dtype=np.float32) / 255.0
    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
    std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
    return ((pixels - mean) / std).transpose(2, 0, 1)


def classify(batches):
    """mapInArrow function: everything before the loop runs once per partition."""
    import json

    import numpy as np
    import onnxruntime as ort
    from huggingface_hub import hf_hub_download

    # Download where the task runs (cached per machine), so this also works on cluster workers.
    model_path = hf_hub_download(MODEL_REPO, "onnx/model.onnx", revision=MODEL_REVISION)
    config_path = hf_hub_download(MODEL_REPO, "config.json", revision=MODEL_REVISION)
    # ImageNet labels list synonyms ("tabby, tabby cat"); keep the first name.
    id2label = json.loads(Path(config_path).read_text())["id2label"]
    labels = {int(i): name.split(",")[0] for i, name in id2label.items()}
    options = ort.SessionOptions()
    options.intra_op_num_threads = 1  # the engine runs partitions in parallel already
    session = ort.InferenceSession(model_path, sess_options=options, providers=["CPUExecutionProvider"])

    for input_batch in batches:
        for start in range(0, input_batch.num_rows, INFERENCE_BATCH_SIZE):
            batch = input_batch.slice(start, INFERENCE_BATCH_SIZE)
            paths = batch.column("path").to_pylist()
            if not paths:
                continue
            pixels = np.stack([preprocess(path) for path in paths])
            (logits,) = session.run(["logits"], {"pixel_values": pixels})
            exp = np.exp(logits - logits.max(axis=1, keepdims=True))
            probabilities = exp / exp.sum(axis=1, keepdims=True)
            rows = {name: [] for name in ARROW_SCHEMA.names}
            for path, p in zip(paths, probabilities):
                best = np.argsort(-p)[:TOP_K]
                rows["file"].append(Path(path).name)
                rows["label"].append(labels[int(best[0])])
                rows["score"].append(float(p[best[0]]))
                rows["top3"].append([{"label": labels[int(i)], "score": float(p[i])} for i in best])
            yield pa.RecordBatch.from_pydict(rows, schema=ARROW_SCHEMA)


def run(spark: SparkSession) -> dict[str, DataFrame]:
    # A table of paths, one row per image. Each task opens only the files in its own rows.
    files = sorted(str(p) for p in IMAGES.iterdir() if p.suffix.lower() in {".jpg", ".jpeg", ".png"})
    paths = spark.createDataFrame([(p,) for p in files], "path string").repartition(2)

    # Classify, then save once: the queries below must not run the model again.
    predictions_path = temp_dir("predictions-")
    paths.mapInArrow(classify, SCHEMA).write.mode("overwrite").parquet(predictions_path)
    predictions = spark.read.parquet(predictions_path)

    # Low-confidence predictions, with the runner-up label, for a person to check.
    review = predictions.filter(F.col("score") < REVIEW_BELOW).select(
        "file",
        "label",
        "score",
        F.col("top3")[1]["label"].alias("runner_up"),
        F.col("top3")[1]["score"].alias("runner_up_score"),
    )
    return {"predictions": predictions.orderBy("file"), "review": review.orderBy("file")}


if __name__ == "__main__":
    from pysail.spark import SparkConnectServer

    server = SparkConnectServer()
    server.start()
    _, port = server.listening_address
    spark = SparkSession.builder.remote(f"sc://localhost:{port}").getOrCreate()
    results = run(spark)
    results["predictions"].select("file", "label", F.round("score", 3).alias("score")).show(truncate=False)
    print("Needs review:")
    results["review"].select(
        "file",
        "label",
        F.round("score", 3).alias("score"),
        "runner_up",
        F.round("runner_up_score", 3).alias("runner_up_score"),
    ).show(truncate=False)
    spark.stop()
    server.stop()

How it works

• A table of paths. The image files become a DataFrame with one path row per file, split into two partitions. Each task opens only the files in its own rows.
• Load once per partition. mapInArrow downloads the model where the task runs (cached per machine), then opens the ONNX Runtime session before its loop over batches.
• Preprocess like the model was trained. Each image is resized so its shorter side is 256 pixels, center-cropped to 224 by 224, and normalized with the ImageNet mean and standard deviation. This follows the model's preprocessor_config.json (checked to give the same pixels as the transformers 4.57.1 image processor on these photos).
• Batch inference. Each Arrow batch is split into groups of at most 16 images per model call. A softmax turns the scores into probabilities, and the function returns the best label, its probability, and the top three as an ARRAY<STRUCT<label, score>>.
• Save once. The predictions are written to Parquet and read back, so the review query never runs the model again.
• Review in SQL. Predictions below 0.6 are selected with their runner-up label, taken from the top3 array.

Use your own images

Point IMAGES at your own folder, or build the paths table from any list of files. The paths must be readable from wherever the UDF runs: local paths on one machine, object storage URLs (read with a client such as fsspec) on a cluster. A file that cannot be decoded raises an error and fails the job. To record such files and keep going, catch the error per file and return a status column, as the OCR example does.

Evaluate the 0.6 cutoff on labeled images for your task before using it to route predictions for review.

To use another classifier, change MODEL_REPO and MODEL_REVISION, and match preprocess to the model's preprocessing settings (preprocessor_config.json in its repository).

Photos

All eight photos are from Wikimedia Commons and dedicated to the public domain under CC0. They were resized to at most 640 pixels on the longer side.

File	Source
espresso.jpg	Antoccino Espresso
golden-retriever.jpg	Golden retriever 2
monarch-butterfly.jpg	Danaus plexippus on a Zinnia
red-fox.jpg	Vulpes vulpes ssp fulvus
siamese-cat.jpg	Siamese cat HD
strawberries.jpg	Strawberries (249191619)
tabby-cat.jpg	Tabby cat laying on the ground
zebra.jpg	Zebra no Zoológico de Brasília março de 2025

Related Sail behavior

• Process media through a table of file paths
• Load a model once per partition
• Set Python batch sizes through the environment

Keep exploring

AI data preparationSearch photos by description →Find a photo by describing what you want to see. AI data preparationEmbed text and search it with vector functions →Find related documents using text embeddings and SQL.