File size: 3,038 Bytes
d667566 | 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 | """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: # tile-offset test-time augmentation
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()
|