#!/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()} }")