File size: 3,311 Bytes
f64e572 | 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 | # SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
"""Render SuperPoint keypoints produced by the tt-nn model onto the sample image.
Usage (run from repo root):
TT_METAL_DIR=/path/to/tt-metal DEVICE_ID=3 \
PYTHONPATH=.:$TT_METAL_DIR:$TT_METAL_DIR/ttnn \
TT_METAL_HOME=$TT_METAL_DIR ARCH_NAME=blackhole \
python models/visualize.py
Writes media/sample.png: the resized input with the top-500 keypoints overlaid
as cyan circles sized by relative score.
"""
from __future__ import annotations
import argparse
import os
from pathlib import Path
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image, ImageDraw
import ttnn
from models.reference.superpoint_reference import (
DEFAULT_NATURAL_IMAGE,
get_natural_input,
load_reference_model,
)
from models.tt.superpoint_ttnn import TtSuperPoint
HEIGHT, WIDTH = 480, 640
TOP_K = 500
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--image", type=Path, default=DEFAULT_NATURAL_IMAGE)
parser.add_argument("--out", type=Path, default=Path("media/sample.png"))
parser.add_argument("--device-id", type=int, default=int(os.environ.get("DEVICE_ID", "0")))
parser.add_argument("--top-k", type=int, default=TOP_K)
args = parser.parse_args()
device = ttnn.CreateDevice(device_id=args.device_id, l1_small_size=32 * 1024)
torch_model = load_reference_model()
pixel_values = get_natural_input(path=args.image, batch_size=1, height=HEIGHT, width=WIDTH)
tt_model = TtSuperPoint(torch_model, device, input_height=HEIGHT, input_width=WIDTH)
tt_in = tt_model.allocate_input(batch_size=1)
tt_model.load_input(tt_in, pixel_values)
s_sm, d_norm = tt_model.run_device_compute(tt_in, b=1)
ttnn.synchronize_device(device)
# Device softmax output + host fold + NMS.
enc_h, enc_w = HEIGHT // 8, WIDTH // 8
scores_nhwc = ttnn.to_torch(s_sm).reshape(1, enc_h, enc_w, 65)
scores_nchw = scores_nhwc.permute(0, 3, 1, 2).contiguous().float()
scores_pre = tt_model._decode_keypoints(scores_nchw, apply_nms=False)
scores_nms = tt_model._simple_nms(scores_pre, tt_model.nms_radius)
flat = scores_nms[0].flatten()
k_eff = min(args.top_k, flat.numel())
score_values, idx = torch.topk(flat, k_eff)
ys = (idx // WIDTH).tolist()
xs = (idx % WIDTH).tolist()
scores = score_values.tolist()
# Draw on the original pixel values (resized 480×640 RGB in [0, 1]).
img_t = pixel_values[0].permute(1, 2, 0).clamp(0, 1)
img_np = (img_t.numpy() * 255.0).astype(np.uint8)
pil = Image.fromarray(img_np)
draw = ImageDraw.Draw(pil, "RGBA")
if scores:
s_min, s_max = min(scores), max(scores)
span = max(s_max - s_min, 1e-9)
for x, y, s in zip(xs, ys, scores):
r = 2 + int(3 * (s - s_min) / span) # 2..5 px radius
draw.ellipse((x - r, y - r, x + r, y + r), outline=(0, 255, 255, 230), width=1)
args.out.parent.mkdir(parents=True, exist_ok=True)
pil.save(args.out)
print(f"wrote {args.out} ({len(xs)} keypoints, top-K={args.top_k})")
ttnn.deallocate(s_sm)
ttnn.deallocate(d_norm)
ttnn.deallocate(tt_in)
ttnn.close_device(device)
if __name__ == "__main__":
main()
|