File size: 3,870 Bytes
db208b4 | 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 | """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:
# Same rounding as training's grid_boxes.
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()
|