File size: 5,817 Bytes
5bcd1fb ecf5ff4 5bcd1fb ecf5ff4 5bcd1fb ecf5ff4 5bcd1fb ecf5ff4 5bcd1fb ecf5ff4 5bcd1fb ecf5ff4 5bcd1fb ecf5ff4 5bcd1fb ecf5ff4 5bcd1fb | 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 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | from __future__ import annotations
import json
from pathlib import Path
from typing import Literal, 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
ModelFormat = Literal["fp32", "fp16"]
MODEL_FILENAMES: dict[ModelFormat, str] = {
"fp32": "model.onnx",
"fp16": "model.fp16.onnx",
}
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",),
model_format: ModelFormat = "fp32",
) -> None:
self.directory = Path(model_dir)
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")
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,
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,
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:
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))
parser.add_argument("--model-format", choices=tuple(MODEL_FILENAMES), default="fp32")
args = parser.parse_args()
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))
|