File size: 3,823 Bytes
58ba6df | 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 | # SPDX-License-Identifier: Apache-2.0
"""Round 9: robustness sweep -- device vs CPU fp32 reference on out-of-distribution versions of the 8 annotated
hands (dark, bright, low contrast, gaussian noise, salt-and-pepper, blocky) plus pure random uint8 crops (what the
server warm-up gate uses). Same served TtHamer path as acc.py; model knobs from the environment.
usage: robust.py [--ref-cache PATH] prints per-group mean / min PCC and ``ROBUST {json}``"""
from __future__ import annotations
import argparse, json, os, sys
from pathlib import Path
_REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_REPO))
def inputs():
import torch
from hamer.preprocess import preprocess_pil, normalize_raw
from PIL import Image
media = _REPO.parent / "media"
boxes = json.loads((media / "bboxes.json").read_text())
g = torch.Generator().manual_seed(7)
out = []
for name in sorted(boxes):
img = Image.open(media / name).convert("RGB")
for hi, e in enumerate(boxes[name]):
u, _ = preprocess_pil(img, tuple(e["bbox"]), is_right=bool(e.get("is_right", 1)), raw=True)
f = u.float()
var = {
"dark": f * 0.25,
"bright": f * 1.8,
"lowcon": 128 + (f - 128) * 0.3,
"noise25": f + 25 * torch.randn(f.shape, generator=g),
"saltpep": torch.where(torch.rand(f.shape, generator=g) < 0.05,
255 * (torch.rand(f.shape, generator=g) < 0.5).float(), f),
"blocky": torch.nn.functional.interpolate(torch.nn.functional.avg_pool2d(f, 16), scale_factor=16),
}
for k, v in var.items():
out.append((f"{k}:{name}#{hi}", v.round().clamp(0, 255).to(torch.uint8)))
for i in range(8):
out.append((f"random:{i}", torch.randint(0, 256, (1, 3, 256, 192), generator=g, dtype=torch.uint8)))
return [(n, u, normalize_raw(u)) for n, u in out]
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--ref-cache", default=str(Path(os.environ.get("TMPDIR", "/tmp")) / "hamer_robust_ref.pt"))
a = ap.parse_args()
os.environ.setdefault("TT_FUSED_HASH", "0")
import torch
from hamer import Hamer, build_from_reference, load_hamer_checkpoint
hf = Path.home() / ".cache/huggingface/hub/models--changh95--hamer-weights/snapshots/e8217b46d54b4e8afaaf0dd6efb6b62341c97d32"
ref = Hamer().eval()
load_hamer_checkpoint(ref, hf / "hamer.ckpt", hf / "mano_mean_params.npz")
xs = inputs()
cache = Path(a.ref_cache)
refs = None
if cache.exists():
d = torch.load(cache)
if d.get("names") == [n for n, _, _ in xs]:
refs = d["out"]
if refs is None:
with torch.no_grad():
refs = [ref(x) for _, _, x in xs]
torch.save({"names": [n for n, _, _ in xs], "out": refs}, cache)
def pcc(u, v):
u = u.reshape(-1).double(); v = v.reshape(-1).double()
u = u - u.mean(); v = v - v.mean()
return float((u @ v) / (u.norm() * v.norm()).clamp_min(1e-12))
tt = build_from_reference(ref, device_id=0, strict=True)
raw = getattr(tt.fused_opts, "input", "norm") == "u8"
try:
with torch.no_grad():
tt(xs[0][1] if raw else xs[0][2]); tt(xs[0][1] if raw else xs[0][2])
res = {}
for (n, u, x), r in zip(xs, refs):
o = tt(u if raw else x)
res.setdefault(n.split(":")[0], []).append(pcc(o, r))
finally:
tt.close()
summ = {k: {"mean": sum(v) / len(v), "min": min(v)} for k, v in res.items()}
for k, v in summ.items():
print(f"[robust] {k:8s} mean {v['mean']:.6f} min {v['min']:.6f}")
print("ROBUST " + json.dumps(summ))
if __name__ == "__main__":
main()
|