| """Tiled full-image inference.""" |
| from __future__ import annotations |
|
|
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| from .config import PREP, STRIDE, TILE |
| from .decode import decode_heatmap, postprocess, wbf |
| from .preprocess import load_prepared, stream_tensors |
| from .tiles import pad_to, tile_origins |
|
|
|
|
| @torch.no_grad() |
| def predict_pair(model, split: str, idx: int, device: str = "cuda", |
| prep_root: Path | None = None, tile: int = TILE, stride: int = STRIDE, |
| batch: int = 4, score_thr: float = 0.05, offsets: tuple[int, ...] = (0,), |
| use_snap: bool = True, use_polarity: bool = True): |
| """Returns (boxes (N,4) in template pixels, scores (N,)).""" |
| t, p, tn, pn = load_prepared(split, idx, prep_root or PREP) |
| h, w = t.shape[:2] |
| t, p = pad_to(t, tile), pad_to(p, tile) |
| tnp, pnp = pad_to(tn, tile), pad_to(pn, tile) |
| a_full, b_full = stream_tensors(t, p, tnp, pnp) |
|
|
| all_boxes, all_scores = [], [] |
| for shift in offsets: |
| origins = tile_origins(t.shape[0], t.shape[1], tile, stride) |
| if shift: |
| origins = [(max(0, min(t.shape[1] - tile, x + shift)), |
| max(0, min(t.shape[0] - tile, y + shift))) for x, y in origins] |
| for i in range(0, len(origins), batch): |
| chunk = origins[i:i + batch] |
| a = np.stack([a_full[:, y:y + tile, x:x + tile] for x, y in chunk]) |
| b = np.stack([b_full[:, y:y + tile, x:x + tile] for x, y in chunk]) |
| ta = torch.from_numpy(a).to(device) |
| tb = torch.from_numpy(b).to(device) |
| with torch.autocast(device_type="cuda" if "cuda" in device else "cpu", |
| dtype=torch.bfloat16, enabled="cuda" in device): |
| out = model(ta, tb) |
| for k, (x0, y0) in enumerate(chunk): |
| boxes, scores = decode_heatmap( |
| out["hm"][k].float(), out["wh"][k].float(), out["off"][k].float(), |
| score_thr=score_thr, |
| ) |
| if len(boxes): |
| boxes[:, [0, 2]] += x0 |
| boxes[:, [1, 3]] += y0 |
| all_boxes.append(boxes) |
| all_scores.append(scores) |
|
|
| if not all_boxes: |
| return np.zeros((0, 4), np.float32), np.zeros(0, np.float32) |
|
|
| boxes = np.concatenate(all_boxes, 0) |
| scores = np.concatenate(all_scores, 0) |
| boxes, scores = wbf(boxes, scores) |
| boxes, scores = postprocess(boxes, scores, tn, pn, use_snap, use_polarity) |
| keep = (boxes[:, 0] < w) & (boxes[:, 1] < h) |
| return boxes[keep], scores[keep] |
|
|
|
|
| def load_model(ckpt_path: Path, device: str = "cuda", backbone: str = "convnext_tiny"): |
| from .model import SiamCenterNet |
|
|
| model = SiamCenterNet(backbone=backbone, pretrained=False) |
| state = torch.load(ckpt_path, map_location="cpu") |
| model.load_state_dict(state["model"] if "model" in state else state) |
| return model.to(device).eval() |
|
|
