| """Predict with an exported ONNX model. Needs only onnxruntime, numpy, and Pillow (no PyTorch). |
| |
| python predict_onnx.py --model grid9.int8.onnx --image challenge.png --target bicycles |
| """ |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import onnxruntime as ort |
| from PIL import Image |
|
|
| from targets import target_name |
|
|
| MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) |
| STD = np.array([0.229, 0.224, 0.225], dtype=np.float32) |
| IMAGE_NAMES = ("image.jpg", "image.jpeg", "image.png", "image.webp") |
|
|
|
|
| def preprocess(image, size): |
| """Matches torchvision Resize (bilinear, PIL), ToTensor and Normalize.""" |
| image = image.convert("RGB").resize((size, size), Image.Resampling.BILINEAR) |
| return ((np.asarray(image, dtype=np.float32) / 255 - MEAN) / STD).transpose(2, 0, 1) |
|
|
|
|
| class GridPredictor: |
| def __init__(self, model, threads=None): |
| if threads is not None and threads < 1: |
| raise ValueError("CPU threads must be positive") |
| options = ort.SessionOptions() |
| if threads: |
| options.intra_op_num_threads = threads |
| self.session = ort.InferenceSession(str(model), options, providers=["CPUExecutionProvider"]) |
| self.info = json.loads(self.session.get_modelmeta().custom_metadata_map["grid_vision"]) |
|
|
| def inputs(self, source): |
| """Model input for a challenge image, or a dataset item folder holding image.jpg.""" |
| source, size, tiles = Path(source), self.info["size"], self.info["tiles"] |
| if source.is_dir(): |
| found = next((source / name for name in IMAGE_NAMES if (source / name).is_file()), None) |
| if found is None: |
| raise ValueError(f"No image in {source} (expected one of {IMAGE_NAMES})") |
| source = found |
| with Image.open(source) as opened: |
| image = opened.convert("RGB") |
| if tiles == 9: |
| |
| w, h = image.size |
| xs, ys = [round(w * i / 3) for i in range(4)], [round(h * i / 3) for i in range(4)] |
| return np.stack( |
| [ |
| preprocess(image.crop((xs[c], ys[r], xs[c + 1], ys[r + 1])), size) |
| for r in range(3) |
| for c in range(3) |
| ] |
| ) |
| return preprocess(image, size)[None] |
|
|
| def probabilities(self, inputs): |
| """Per-class probabilities, shaped (classes, cells).""" |
| output = self.session.run(None, {self.session.get_inputs()[0].name: inputs})[0] |
| return output.T if self.info["tiles"] == 9 else output[0] |
|
|
| def predict(self, source, target): |
| name = target_name(target) |
| if name not in self.info["classes"]: |
| raise ValueError(f"Unknown target {name!r}; supported: {self.info['classes']}") |
| probabilities = self.probabilities(self.inputs(source))[self.info["classes"].index(name)] |
| cutoff = self.info["thresholds"].get(name, self.info["threshold"]) |
| return { |
| "target": name, |
| "threshold": cutoff, |
| "selected_cells": [int(i) for i in np.flatnonzero(probabilities >= cutoff)], |
| "probabilities": probabilities.tolist(), |
| } |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser( |
| description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter |
| ) |
| parser.add_argument("--model", required=True) |
| parser.add_argument("--image", required=True, help="Challenge image, or a dataset item folder") |
| parser.add_argument("--target", required=True) |
| parser.add_argument("--threads", type=int, help="CPU threads (default: onnxruntime's choice)") |
| args = parser.parse_args() |
| print( |
| json.dumps( |
| GridPredictor(args.model, args.threads).predict(args.image, args.target), indent=2 |
| ) |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
