File size: 3,361 Bytes
aec3f16 | 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 | #!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Quickstart: the model card's Python snippet on the shipped sample, on one Blackhole p150.
pip install -e . # once, from the repo root, on top of an environment that has ttnn (tt-metal)
python examples/quickstart.py [input] [--variant bgr|rgb] [--out-dir examples/output]
Writes <out-dir>/quickstart.json (the same JSON as POST /predict: the published mono8 mask and the 3-class map, both
as PNG, plus the pixel counts per class), <out-dir>/quickstart_mask.png (the mask itself: 255 = foreground, the
input's size) and <out-dir>/quickstart.png (the input with foreground tinted magenta and road violet). The default
input is found relative to this file (runs from any directory); an input given on the command line is relative to
the current directory. The weights come from this repo's weights/ at the pinned commit (your HF cache).
"""
import argparse
import json
from pathlib import Path
REPO = Path(__file__).resolve().parents[1]
ap = argparse.ArgumentParser()
ap.add_argument("input", nargs="?", default=str(REPO / "code" / "tt_sceneseg" / "samples" / "highway_normal_1.png"))
ap.add_argument("--variant", choices=("bgr", "rgb"), default=None,
help="input channel order: bgr = the deployed VisionPilot order (default), rgb = the training order")
ap.add_argument("--out-dir", default=str(REPO / "examples" / "output"))
ap.add_argument("--device-id", type=int, default=0)
args = ap.parse_args()
out_dir = Path(args.out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
# --- the model card snippet --------------------------------------------------------------------------------------
from tt_sceneseg import SceneSeg
with SceneSeg.from_pretrained(device_id=args.device_id, variant=args.variant) as model: # weights -> HF cache, trace captured
out = model(args.input) # path, bytes, PIL image or RGB uint8 HxWx3 array of any size
mask = out.mask # uint8 (H, W) at the input size: 255 = foreground (the VisionPilot topic)
print(out.class_counts()) # pixels per class of out.class_map (background / foreground / road)
# ------------------------------------------------------------------------------------------------------------------
(out_dir / "quickstart.json").write_text(json.dumps(out.to_dict(), indent=1))
# pictures of the result: the published mask, and the class map over the input
import numpy as np # noqa: E402
from PIL import Image # noqa: E402
Image.fromarray(mask).save(out_dir / "quickstart_mask.png")
image = np.asarray(Image.open(args.input).convert("RGB"), np.float32)
tint = image.copy()
tint[out.class_map == 2] = 0.65 * tint[out.class_map == 2] + 0.35 * np.array([74, 58, 167], np.float32) # road
tint[out.class_map == 1] = 0.45 * tint[out.class_map == 1] + 0.55 * np.array([232, 123, 164], np.float32) # foreground
Image.fromarray(tint.round().clip(0, 255).astype(np.uint8)).save(out_dir / "quickstart.png")
fg = float((mask == 255).mean())
print(f"variant {out.meta['channel_order']}: mask {mask.shape}, foreground {100 * fg:.2f} % of pixels -> "
f"{out_dir / 'quickstart.json'}, {out_dir / 'quickstart_mask.png'}, {out_dir / 'quickstart.png'} "
f"timing_ms={ {k: round(v, 2) for k, v in out.timing_ms.items()} }")
|