Spaces:
Running on Zero
Running on Zero
File size: 24,944 Bytes
99b82f1 faa2ae9 99b82f1 faa2ae9 99b82f1 faa2ae9 99b82f1 faa2ae9 99b82f1 faa2ae9 99b82f1 faa2ae9 99b82f1 | 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 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 | """Ink Lens engine: an openly licensed 9 um ink model on one region, three depth windows, both depth orders.
The model is built by villa's own flat ink inference (vesuvius.ink_detection.inference.infer.configure_model, at the
commit pinned in VILLA_COMMIT) from a checkpoint downloaded by URL and refused unless its SHA-256 matches. The patch
schedule (128 px patches, stride 64, the last flush with the end), the layer choice (center crop, reversed for the
control), villa's per-patch robust normalization and the Hann blend are villa's; two things differ from villa's
command line, both to fit one ZeroGPU call:
1. the robust normalization runs batched on the device (robust_torch), checked against villa's CPU function
(robust_np is a line-for-line copy of villa's normalize_robust for a patch with no mask; tests/test_inklens.py
compares them, and compares robust_np with villa's own function when villa is importable);
2. the maps stay float (0 to 1) instead of villa's uint8 encoding, so the control-set threshold is not quantized.
run_reference() is the same algorithm in numpy with any predict function: the tests run it with a stand-in model,
and the torch path is checked against it.
python3 engine.py estimate 1024 1024 --windows 3 patches and the GPU seconds one call asks for
python3 engine.py fetch d9v2 download and check a checkpoint into the cache
"""
from __future__ import annotations
import hashlib
import math
import os
import platform
import shutil
import subprocess
import sys
import tarfile
import tempfile
import time
import urllib.request
from pathlib import Path
VERSION = "0.1.0"
VILLA_REPO = "https://github.com/ScrollPrize/villa.git"
VILLA_TARBALL = "https://codeload.github.com/ScrollPrize/villa/tar.gz/{commit}"
VILLA_COMMIT = "8aa507c1215f1fa8758f13d0538f7af9710e02d7" # ScrollPrize/villa main, 18 Oct 2026
CACHE = Path(os.environ.get("INKLENS_CACHE") or Path(tempfile.gettempdir()) / "ink-lens")
USER_AGENT = f"ink-lens/{VERSION}"
# Openly licensed 9 um readers that drop into villa's flat inference. Public held-out scores (PHerc0841, crops with a
# 64 px inner edge, raw / reversed AUC) are in docs/results.json; d9v2 is the default.
READERS = {
"d9v2": {
"title": "d9v2 (TAUIL-Abd-Elilah, MIT)",
"url": "https://github.com/TAUIL-Abd-Elilah/pherc0826-first-letters-search/releases/download/v1.0/"
"d9v2_ft-012000.pth",
"file": "d9v2_ft-012000.pth",
"sha256": "50d2ad0ef7690a6a422640804d38f01409da8bc96cfc1ab72f45324a4a18f966",
"license": "MIT",
"source": "https://github.com/TAUIL-Abd-Elilah/pherc0826-first-letters-search",
"held_out": "PHerc0841 w00, ag896, ag405 (its reader/test.txt)",
},
"ink_9um": {
"title": "ink_9um seed 42 step 75,000 (Vesuvius Challenge, MIT)",
"url": "https://huggingface.co/scrollprize/ink_9um/resolve/main/hybrid_3d2d-seed42/step-075000.pth",
"file": "ink_9um-seed42-step-075000.pth",
"sha256": "e635558ae6a1a807a7e5ec1e83adfd45bc3c0ac53883ea43f1d4e085d62a9cab",
"license": "MIT",
"source": "https://huggingface.co/scrollprize/ink_9um",
"held_out": "PHerc0841 (its label dataset has no PHerc0841 segment)",
},
}
READER_CHOICES = ("d9v2", "ink_9um", "both")
PATCH, STRIDE, DEPTH = 128, 64, 17
STACK = 21 # layers kept around the middle: the three windows below fit in it
WINDOW_OFFSETS = {1: (2,), 3: (0, 2, 4)} # window starts inside the 21-layer stack (2 is villa's own center)
BATCH = 64
MIN_SIDE, MAX_SIDE = PATCH, 2048
MIN_LAYERS, MAX_LAYERS = DEPTH, 65
# One call's duration: a fixed start plus a per-patch allowance, capped. The per-patch figure is an estimate (3 to
# 6 ms per patch with batched fp16 and normalization on the GPU) with a margin; every receipt records the measured
# seconds so the first calls calibrate it.
GPU_START_S, GPU_PER_PATCH_S, GPU_MAX_S = 4.0, 0.012, 40
class EngineError(ValueError):
"""A request the engine refuses; the message says why."""
# ---------------------------------------------------------------- sizing
def positions(length, patch=PATCH, stride=STRIDE):
"""Patch starts along one axis, the last flush with the end (villa's _sliding_positions_1d)."""
if length <= patch:
return [0]
out = list(range(0, length - patch + 1, max(1, stride)))
if out[-1] != length - patch:
out.append(length - patch)
return out
def readers_of(choice):
if choice not in READER_CHOICES:
raise EngineError(f"reader {choice!r}: choose one of {', '.join(READER_CHOICES)}")
return ("d9v2", "ink_9um") if choice == "both" else (choice,)
def patches(height, width, windows=3, reader="d9v2"):
"""Patches one call runs: per pass the patch grid, times windows, times two depth orders, times readers."""
if windows not in WINDOW_OFFSETS:
raise EngineError(f"windows {windows}: use 1 or 3")
return len(positions(height)) * len(positions(width)) * windows * 2 * len(readers_of(reader))
def gpu_seconds(n_patches):
"""The duration a call asks ZeroGPU for; refuses a call that would need more than GPU_MAX_S."""
need = math.ceil(GPU_START_S + n_patches * GPU_PER_PATCH_S)
if need > GPU_MAX_S:
raise EngineError(f"{n_patches} patches need about {need} GPU seconds, more than the {GPU_MAX_S} s one call "
"may ask for; use a smaller region, one depth window, or local mode")
return need
def estimate(height, width, windows=3, reader="d9v2"):
check_side(height, width)
n = patches(height, width, windows, reader)
out = {"height": int(height), "width": int(width), "windows": int(windows), "reader": reader, "patches": n}
try:
out["gpu_seconds_requested"] = gpu_seconds(n)
except EngineError as exc:
out["gpu_seconds_requested"] = None
out["refused"] = str(exc)
return out
def check_side(height, width):
for name, side in (("height", height), ("width", width)):
if not MIN_SIDE <= int(side) <= MAX_SIDE:
raise EngineError(f"{name} {side} px: a region is {MIN_SIDE} to {MAX_SIDE} px a side at 9.6 um")
def check_volume(volume, windows=3):
"""Refuse anything but a uint8 (layers, rows, columns) region of the right size and depth."""
import numpy as np
a = np.asarray(volume)
if a.ndim != 3:
raise EngineError(f"a surface volume has 3 axes (layers, rows, columns), this one has {a.ndim}")
if a.dtype != np.uint8:
raise EngineError(f"type {a.dtype}: the engine takes uint8 (volio.to_uint8 converts uint16)")
L, H, W = a.shape
need = STACK if windows == 3 else DEPTH
if not need <= L <= MAX_LAYERS:
raise EngineError(f"{L} layers: {windows} depth window{'s' if windows > 1 else ''} need{'' if windows > 1 else 's'} "
f"{need} to {MAX_LAYERS} layers, layers first")
check_side(H, W)
if not a.any():
raise EngineError("the region is all zero (off the segment, or a hole in it)")
return L, H, W
def stack(volume, windows=3):
"""The middle STACK layers (or all, if fewer), as villa's center crop chooses them (upper center when uneven)."""
L = volume.shape[0]
n = min(L, STACK)
start = L // 2 - n // 2
return volume[start:start + n]
def window_indices(n_layers, offset, direction):
"""Layer indices of one 17-layer window inside the stack; reversed for the control."""
if n_layers < STACK: # one window, centered (villa's own choice)
start = n_layers // 2 - DEPTH // 2
else:
start = offset
idx = list(range(start, start + DEPTH))
return idx[::-1] if direction == "reverse" else idx
def passes(windows):
"""(offset, direction) for every pass of one reader: each window in both orders."""
return [(o, d) for o in WINDOW_OFFSETS[windows] for d in ("forward", "reverse")]
# ---------------------------------------------------------------- villa's normalization, numpy and torch
def robust_np(patch):
"""villa's normalize_robust (vesuvius/image_proc/intensity/normalization.py at VILLA_COMMIT) with no mask:
clip to the 1st and 99th percentiles, subtract the median, divide by 1.4826 x MAD, with villa's fallbacks."""
import numpy as np
arr = np.asarray(patch).astype(np.float32, copy=True)
valid = arr.reshape(-1)
if valid.size == 0:
return arr
std_preclip = float(np.std(valid))
lower = float(np.percentile(valid, 1.0))
upper = float(np.percentile(valid, 99.0))
np.clip(arr, lower, upper, out=arr)
valid = arr.reshape(-1)
median = float(np.median(valid))
mad = float(np.median(np.abs(valid - median)))
scaled = 1.4826 * mad
if not np.isfinite(scaled) or scaled < 1e-6:
span = upper - lower
cands = []
if np.isfinite(std_preclip):
cands.append(abs(std_preclip))
if np.isfinite(span):
cands.append(abs(span) / 2.0)
scaled = next((c for c in cands if c >= 1e-6), None)
if scaled is None or not np.isfinite(scaled):
scaled = 1.0
arr -= median
arr /= scaled
np.nan_to_num(arr, copy=False, nan=0.0, posinf=0.0, neginf=0.0)
return arr
def _lerp_sorted(torch, s, q):
"""numpy's linear percentile of each row of an already sorted (B, N) tensor."""
n = s.shape[1]
pos = q * (n - 1)
lo = int(math.floor(pos))
hi = min(lo + 1, n - 1)
frac = pos - lo
return s[:, lo] + (s[:, hi] - s[:, lo]) * frac
def robust_torch(x):
"""robust_np on every row of a (B, N) float32 tensor at once, on its device."""
import torch
x = x.float()
s, _ = torch.sort(x, dim=1)
std_pre = x.std(dim=1, unbiased=False)
lower, upper = _lerp_sorted(torch, s, 0.01), _lerp_sorted(torch, s, 0.99)
x = torch.maximum(torch.minimum(x, upper[:, None]), lower[:, None])
s = torch.maximum(torch.minimum(s, upper[:, None]), lower[:, None])
median = _lerp_sorted(torch, s, 0.5)
dev, _ = torch.sort((x - median[:, None]).abs(), dim=1)
scaled = 1.4826 * _lerp_sorted(torch, dev, 0.5)
half_span = (upper - lower).abs() / 2.0
ok = torch.isfinite(scaled) & (scaled >= 1e-6)
fb = torch.where(torch.isfinite(std_pre) & (std_pre.abs() >= 1e-6), std_pre.abs(),
torch.where(torch.isfinite(half_span) & (half_span >= 1e-6), half_span,
torch.ones_like(half_span)))
scaled = torch.where(ok, scaled, fb)
out = (x - median[:, None]) / scaled[:, None]
return torch.nan_to_num(out, nan=0.0, posinf=0.0, neginf=0.0)
def hann2d():
"""villa's floored Hann blend map (compute_importance_map_2d, mode hann), in numpy."""
import numpy as np
w1 = np.hanning(PATCH).astype(np.float32) # numpy's hanning is the symmetric (periodic=False) window
w = np.outer(w1, w1)
w /= max(float(w.max()), float(np.finfo(np.float32).eps))
return np.maximum(w, 0.001).astype(np.float32)
# ---------------------------------------------------------------- the numpy reference
def run_reference(volume, predict, windows=3):
"""The whole algorithm in numpy: predict(batch (B, 1, 17, 128, 128) float32) -> (B, 128, 128) probabilities.
Returns {"forward": F, "reverse": R} (float32, 0 to 1), each the mean over the windows."""
import numpy as np
vol = stack(np.asarray(volume), windows)
L, H, W = vol.shape
ys, xs = positions(H), positions(W)
wmap = hann2d()
maps = {"forward": [], "reverse": []}
for off, direction in passes(windows):
idx = window_indices(L, off, direction)
psum = np.zeros((H, W), np.float32)
wsum = np.zeros((H, W), np.float32)
for y0 in ys:
for x0 in xs:
raw = vol[idx, y0:y0 + PATCH, x0:x0 + PATCH]
if not raw.any():
continue
p = np.asarray(predict(robust_np(raw)[None, None]), np.float32)[0]
psum[y0:y0 + PATCH, x0:x0 + PATCH] += p * wmap
wsum[y0:y0 + PATCH, x0:x0 + PATCH] += wmap
prob = np.divide(psum, wsum, out=np.zeros_like(psum), where=wsum > 1e-6)
maps[direction].append(np.clip(prob, 0, 1))
return {d: np.mean(m, axis=0).astype(np.float32) for d, m in maps.items()}
# ---------------------------------------------------------------- downloads
def sha256_file(path, block=1 << 22):
h = hashlib.sha256()
with open(path, "rb") as fh:
for chunk in iter(lambda: fh.read(block), b""):
h.update(chunk)
return h.hexdigest()
def _download(url, dest):
dest.parent.mkdir(parents=True, exist_ok=True)
part = dest.with_name(dest.name + ".part")
req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
with urllib.request.urlopen(req, timeout=600) as resp, open(part, "wb") as out:
shutil.copyfileobj(resp, out, 1 << 20)
os.replace(part, dest)
def ensure_checkpoint(name, cache=CACHE):
"""A reader's checkpoint, downloaded once and refused unless its SHA-256 matches READERS. INKLENS_<NAME> (for
example INKLENS_D9V2) may name a local copy; its SHA-256 is still checked."""
spec = READERS[name]
env = os.environ.get("INKLENS_" + name.upper())
path = Path(env).expanduser() if env else Path(cache) / "checkpoints" / spec["file"]
if not path.is_file():
if env:
raise EngineError(f"INKLENS_{name.upper()}={env} does not exist")
_download(spec["url"], path)
digest = sha256_file(path)
if digest != spec["sha256"]:
raise EngineError(f"{path} has SHA-256 {digest}, expected {spec['sha256']} for {name}; refusing it")
return path
def _git_commit(path):
try:
return subprocess.run(["git", "-C", str(path), "rev-parse", "HEAD"], capture_output=True, text=True,
check=True).stdout.strip()
except (OSError, subprocess.CalledProcessError):
return "unknown"
def _villa_tarball(tmp, commit):
root = tmp.resolve()
kept = 0
req = urllib.request.Request(VILLA_TARBALL.format(commit=commit), headers={"User-Agent": USER_AGENT})
with urllib.request.urlopen(req, timeout=600) as resp, tarfile.open(fileobj=resp, mode="r|gz") as tf:
for m in tf:
_, _, rel = m.name.partition("/")
if not rel.startswith("vesuvius/src/"):
continue
target = (root / rel).resolve()
if root not in target.parents or m.issym() or m.islnk():
continue
if m.isdir():
target.mkdir(parents=True, exist_ok=True)
elif m.isfile():
target.parent.mkdir(parents=True, exist_ok=True)
with tf.extractfile(m) as src, open(target, "wb") as out:
shutil.copyfileobj(src, out)
kept += 1
if not kept:
raise EngineError(f"the villa tarball for {commit} held no vesuvius/src files")
def ensure_villa(cache=CACHE):
"""villa's vesuvius package on sys.path: VILLA_SRC (a local vesuvius/src) when set, else the pinned commit,
downloaded once (sparse git fetch, else GitHub's tarball)."""
if os.environ.get("VILLA_SRC"):
src = Path(os.environ["VILLA_SRC"]).expanduser().resolve()
info = {"source": "VILLA_SRC", "commit": _git_commit(src)}
else:
dest = Path(cache) / f"villa-{VILLA_COMMIT[:12]}"
if not (dest / "COMMIT").is_file():
tmp = dest.with_name(dest.name + ".part")
shutil.rmtree(tmp, ignore_errors=True)
tmp.mkdir(parents=True)
try:
git = ["git", "-C", str(tmp)]
for cmd in (["init", "-q"], ["remote", "add", "origin", VILLA_REPO],
["sparse-checkout", "set", "vesuvius/src"],
["fetch", "-q", "--depth", "1", "--filter=blob:none", "origin", VILLA_COMMIT],
["checkout", "-q", "FETCH_HEAD"]):
subprocess.run(git + cmd, capture_output=True, text=True, check=True, timeout=600)
except (OSError, subprocess.SubprocessError):
shutil.rmtree(tmp, ignore_errors=True)
tmp.mkdir(parents=True)
_villa_tarball(tmp, VILLA_COMMIT)
(tmp / "COMMIT").write_text(VILLA_COMMIT + "\n")
shutil.rmtree(dest, ignore_errors=True)
os.replace(tmp, dest)
src = dest / "vesuvius" / "src"
info = {"source": "pinned download", "commit": (dest / "COMMIT").read_text().strip()}
if not (src / "vesuvius" / "ink_detection" / "inference" / "infer.py").is_file():
raise EngineError(f"{src} holds no villa vesuvius package")
if str(src) not in sys.path:
sys.path.insert(0, str(src))
info["pinned"] = VILLA_COMMIT
return info
def load_villa_model(name, cache=CACHE):
"""(model, preprocessing, amp dtype, info) for one reader, built by villa's configure_model. The checkpoint is
opened with torch.load(weights_only=True): villa's loader allows any pickle, so it is replaced for this call."""
import argparse
import torch
villa = ensure_villa(cache)
path = ensure_checkpoint(name, cache)
from vesuvius.ink_detection.inference import infer
original = infer.load_checkpoint
infer.load_checkpoint = lambda p: torch.load(str(p), map_location="cpu", weights_only=True)
try:
cm = infer.configure_model(argparse.Namespace(checkpoint=path, amp_dtype="auto"))
finally:
infer.load_checkpoint = original
got = (int(cm.patch_size), int(cm.input_depth))
if got != (PATCH, DEPTH):
raise EngineError(f"{name} has patch and depth {got}; Ink Lens is built for {(PATCH, DEPTH)}")
if cm.preprocessing not in ("tifxyz_robust", "divide_255"):
raise EngineError(f"{name} asks for preprocessing {cm.preprocessing!r}, which Ink Lens does not run")
info = {"name": name, "sha256": READERS[name]["sha256"], "url": READERS[name]["url"],
"license": READERS[name]["license"], "preprocessing": cm.preprocessing,
"amp_on_cuda": str(cm.amp_dtype).replace("torch.", ""), "villa": villa, "loader": "weights_only=True"}
return cm.model.eval(), cm.preprocessing, cm.amp_dtype, info
# ---------------------------------------------------------------- the torch engine
def pick_device(prefer=None):
import torch
if prefer:
return torch.device(prefer)
if torch.cuda.is_available():
return torch.device("cuda")
if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
return torch.device("mps")
return torch.device("cpu")
class Engine:
"""Readers loaded once; run() does one region on one device.
models: {name: (module, preprocessing, amp_dtype, info)} to use instead of downloading (the tests pass a
stand-in); otherwise every name in `readers` is loaded through villa."""
def __init__(self, readers=("d9v2",), cache=CACHE, models=None, place_on_cuda=True):
t0 = time.time()
import torch
self.torch = torch
self.models = dict(models) if models else {n: load_villa_model(n, cache) for n in readers}
if place_on_cuda and torch.cuda.is_available():
for m in self.models.values():
m[0].to("cuda")
self.startup_s = round(time.time() - t0, 1)
def describe(self):
return {"engine": f"ink-lens {VERSION}", "readers": {n: m[3] for n, m in self.models.items()},
"patch": PATCH, "stride": STRIDE, "depth": DEPTH, "stack": STACK, "blend": "hann (villa, floored)",
"batch": BATCH, "torch": self.torch.__version__}
def run(self, volume, reader="d9v2", windows=3, device=None, batch=BATCH):
"""{"forward": F, "reverse": R, "per_reader": {...}, ...} for one uint8 region, float32 maps in 0 to 1."""
import numpy as np
torch = self.torch
names = readers_of(reader)
missing = [n for n in names if n not in self.models]
if missing:
raise EngineError(f"reader {', '.join(missing)} is not loaded here")
check_volume(volume, windows)
dev = pick_device(device)
t0 = time.perf_counter()
vol = stack(np.asarray(volume), windows)
L, H, W = vol.shape
stack_start = int(np.asarray(volume).shape[0] // 2 - L // 2)
vt = torch.from_numpy(np.ascontiguousarray(vol)).to(dev)
coords = [(y, x) for y in positions(H) for x in positions(W)]
wmap = torch.from_numpy(hann2d()).to(dev)
per_reader, n_patches = {}, 0
for name in names:
model, prep, amp, _ = self.models[name]
model = model.to(dev)
maps = {"forward": [], "reverse": []}
for off, direction in passes(windows):
idx = torch.tensor(window_indices(L, off, direction), device=dev)
sub = vt.index_select(0, idx)
psum = torch.zeros((H, W), dtype=torch.float32, device=dev)
wsum = torch.zeros((H, W), dtype=torch.float32, device=dev)
for b0 in range(0, len(coords), batch):
chunk = coords[b0:b0 + batch]
raw = torch.stack([sub[:, y:y + PATCH, x:x + PATCH] for y, x in chunk]) # (B, 17, 128, 128)
keep = raw.reshape(len(chunk), -1).any(dim=1) # villa skips all-zero patches, before normalizing
if not bool(keep.any()):
continue
if not bool(keep.all()):
chunk = [c for c, k in zip(chunk, keep.tolist()) if k]
raw = raw[keep]
flat = raw.reshape(len(chunk), -1).float()
norm = robust_torch(flat) if prep == "tifxyz_robust" else flat / 255.0
images = norm.reshape(len(chunk), 1, DEPTH, PATCH, PATCH)
ctx = (torch.autocast("cuda", dtype=amp) if dev.type == "cuda" and amp is not None
else _null())
with torch.inference_mode(), ctx:
logits = model(images)
prob = logits.float().sigmoid()
if prob.shape[-2:] != (PATCH, PATCH):
prob = torch.nn.functional.interpolate(prob, size=(PATCH, PATCH), mode="bilinear",
align_corners=False)
prob = prob[:, 0]
for (y, x), p in zip(chunk, prob):
psum[y:y + PATCH, x:x + PATCH] += p * wmap
wsum[y:y + PATCH, x:x + PATCH] += wmap
n_patches += len(chunk)
out = torch.where(wsum > 1e-6, psum / wsum.clamp_min(1e-6), torch.zeros_like(psum)).clamp(0, 1)
maps[direction].append(out)
per_reader[name] = {d: torch.stack(m).mean(0).cpu().numpy().astype(np.float32) for d, m in maps.items()}
if dev.type == "cuda":
torch.cuda.synchronize(dev)
F = np.mean([per_reader[n]["forward"] for n in names], axis=0).astype(np.float32)
R = np.mean([per_reader[n]["reverse"] for n in names], axis=0).astype(np.float32)
name = torch.cuda.get_device_name(dev) if dev.type == "cuda" else (platform.processor() or platform.machine())
return {"forward": F, "reverse": R, "per_reader": per_reader, "readers": list(names), "windows": windows,
"window_layers": {f"{o}:{d}": window_indices(L, o, d) for o, d in passes(windows)},
"stack_layers": int(L), "stack_start": stack_start, "patches": n_patches, "device": str(dev), "device_name": name,
"seconds": round(time.perf_counter() - t0, 3)}
class _null:
def __enter__(self):
return None
def __exit__(self, *exc):
return False
def main(argv=None):
import argparse
import json
ap = argparse.ArgumentParser(description="Ink Lens engine helpers")
sub = ap.add_subparsers(dest="cmd", required=True)
e = sub.add_parser("estimate")
e.add_argument("height", type=int)
e.add_argument("width", type=int)
e.add_argument("--windows", type=int, default=3, choices=(1, 3))
e.add_argument("--reader", default="d9v2", choices=READER_CHOICES)
f = sub.add_parser("fetch")
f.add_argument("reader", choices=sorted(READERS))
a = ap.parse_args(argv)
try:
if a.cmd == "estimate":
print(json.dumps(estimate(a.height, a.width, a.windows, a.reader), indent=1))
else:
print(json.dumps({"reader": a.reader, "path": str(ensure_checkpoint(a.reader))}))
except EngineError as exc:
print(f"refused: {exc}", file=sys.stderr)
return 2
return 0
if __name__ == "__main__":
sys.exit(main())
|