File size: 2,196 Bytes
041c6b2 | 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 | """Runs each model on the sample images on the CPU and compares its output with the first model's."""
import os
import sys
import time
import numpy as np
import onnxruntime as ort
from PIL import Image
samples_dir, out_dir, *models = sys.argv[1:]
os.makedirs(out_dir, exist_ok=True)
MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
def pixels(path):
image = Image.open(path).convert("RGB")
resized = np.asarray(image.resize((1024, 1024), Image.BILINEAR), dtype=np.float32) / 255
return image, ((resized - MEAN) / STD).transpose(2, 0, 1)[None]
def sigmoid(x):
return 1 / (1 + np.exp(-x))
samples = [(name, *pixels(os.path.join(samples_dir, name))) for name in sorted(os.listdir(samples_dir))]
reference = {}
for index, path in enumerate(models):
label = os.path.splitext(os.path.basename(path))[0]
options = ort.SessionOptions()
options.log_severity_level = 3
started = time.time()
session = ort.InferenceSession(path, options, providers=["CPUExecutionProvider"])
print(f"{label}: loaded in {time.time() - started:.1f} s, {os.path.getsize(path) / 1e6:.0f} MB")
input_name = session.get_inputs()[0].name
for name, image, x in samples:
started = time.time()
logits = session.run(None, {input_name: x})[0].astype(np.float32)[0, 0]
took = time.time() - started
alpha = sigmoid(logits)
Image.fromarray((alpha * 255).round().astype(np.uint8)).resize(image.size, Image.BILINEAR).save(os.path.join(out_dir, f"{os.path.splitext(name)[0]}-{label}.png"))
if index == 0:
reference[name] = (logits, alpha)
print(f" {name}: {took:.1f} s, logits {logits.min():.2f}..{logits.max():.2f}, {100 * (alpha > 0.5).mean():.1f}% foreground")
continue
ref_logits, ref_alpha = reference[name]
diff = np.abs(alpha - ref_alpha)
print(
f" {name}: {took:.1f} s, max logit diff {np.abs(logits - ref_logits).max():.2e}, "
f"alpha diff max {diff.max():.2e} mean {diff.mean():.2e}, "
f"pixels off by more than 1/255: {100 * (diff > 1 / 255).mean():.3f}%"
)
|