File size: 2,573 Bytes
028c4b5 | 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 | """Run a bundled Pasma ONNX graph on one local image (no image upload)."""
from __future__ import annotations
import argparse
import json
from pathlib import Path
import numpy as np
import onnxruntime as ort
from PIL import Image, ImageOps
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('image', type=Path)
parser.add_argument('--model', type=Path, default=Path(__file__).resolve().parents[1] / 'models/pasma_small_320_fp32.onnx')
args = parser.parse_args()
metadata = args.model.parent
labels = json.loads((metadata / 'labels_order.json').read_text())
session = ort.InferenceSession(str(args.model), providers=['CPUExecutionProvider'])
size = session.get_inputs()[0].shape[-1]
with Image.open(args.image) as source:
image = ImageOps.exif_transpose(source).convert('RGB')
scale = size / max(image.size)
image = image.resize(tuple(max(1, round(v * scale)) for v in image.size), Image.Resampling.BILINEAR)
canvas = Image.new('RGB', (size, size), (128, 128, 128))
canvas.paste(image, ((size - image.width) // 2, (size - image.height) // 2))
x = np.asarray(canvas, dtype=np.float32).transpose(2, 0, 1)[None] / 255.0
coarse, types, materials = session.run(['prob_coarse', 'prob_type', 'prob_material'], {'image': x})
top3 = np.argsort(-types[0])[:3]
chosen = labels['garment_types'][int(top3[0])]
# Marginal argmax may differ; use this mapping for the displayed pair.
category = labels['type_to_coarse'][chosen]
prefix_to_run = {'pasma_small_320': 'r4_edgenext_320', 'pasma_nano_256': 'r3_edgenext_s',
'pasma_pico_256': 'r3_convnext_pico'}
run = next(run for prefix, run in prefix_to_run.items() if args.model.stem.startswith(prefix))
thresholds = json.loads((metadata / 'material_thresholds.json').read_text())['thresholds'][run]
suggestions = [{'label': name, 'score': float(materials[0, i])}
for i, name in enumerate(labels['materials'])
if thresholds[name]['fitted'] and materials[0, i] >= thresholds[name]['threshold']]
print(json.dumps({'top3': [{'type': labels['garment_types'][int(i)], 'score': float(types[0, i])}
for i in top3], 'category_for_selected_type': category,
'coarse_marginal_argmax': labels['coarse_categories'][int(coarse[0].argmax())],
'visual_material_suggestions': sorted(suggestions, key=lambda v: -v['score'])}, indent=2))
if __name__ == '__main__':
main()
|