File size: 5,750 Bytes
10880af e6863b5 10880af e6863b5 10880af e6863b5 10880af e6863b5 10880af e6863b5 10880af e6863b5 10880af e6863b5 943fb39 10880af e6863b5 10880af e6863b5 943fb39 10880af 943fb39 10880af 943fb39 10880af e6863b5 | 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 | """Turn ink_9um outputs into the page's ZeroGPU layer.
Input: gpu/<winding>/receipt.json plus the forward TIFF it names (uint8, one per winding, the 9.362 um
surface-volume grid). Those TIFFs (about 25 MB each) are not shipped with the Space; anyone can regenerate
them with the model card's command on the public surface volumes listed in data/gpu/receipts/.
Writes data/gpu/<w>.jpg (3D texture, display-filtered), data/hi/<w>_gpu.jpg (close-up view, the model card's
display rescale only), data/gpu/receipts/<w>.json (the run records) and data/gpu.json (the page's summary).
UNROLL_GPU and UNROLL_OUT override the input and output folders.
"""
import glob
import json
import os
import sys
import numpy as np
import tifffile
from PIL import Image
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, f"{ROOT}/tools")
from build import CELL_MM, HI_MAX, PPMM, local_contrast # noqa: E402
OUT = os.environ.get("UNROLL_OUT", f"{ROOT}/data")
GPU_IN = os.environ.get("UNROLL_GPU", f"{ROOT}/gpu")
# run-record fields that are local plumbing, not provenance
DROP = {"space", "status"}
def say(items):
"""'a', 'a and b', 'a, b and c'."""
return items[0] if len(items) == 1 else f"{', '.join(items[:-1])} and {items[-1]}"
def main():
os.makedirs(f"{OUT}/gpu/receipts", exist_ok=True)
manifest = json.load(open(f"{OUT}/scroll.json"))
by_id = {w["id"]: w for w in manifest["windings"]}
wins, tot_patches, tot_measured, tot_calls = {}, 0, 0.0, 0
device = None
dirs = sorted(glob.glob(f"{GPU_IN}/w*/"))
if not dirs:
sys.exit(f"no ZeroGPU outputs under {GPU_IN}/w*/ (see the docstring)")
for d in dirs:
w = os.path.basename(d.rstrip("/"))
receipt = json.load(open(f"{d}receipt.json"))
record = {k: v for k, v in receipt.items() if k not in DROP}
record["gpu_seconds"] = {k: v for k, v in receipt["gpu_seconds"].items() if k != "note"}
record["inputs"] = {"direction": receipt["inputs"]["direction"], "geometry": receipt["inputs"]["geometry"],
"source": receipt["inputs"]["source"]}
json.dump(record, open(f"{OUT}/gpu/receipts/{w}.json", "w"), indent=1, sort_keys=True)
tif = f"{d}{receipt['outputs']['forward']['file']}"
a = tifffile.imread(tif).astype(np.float32) / 255
win = by_id[w]
assert a.shape == (win["gridH"] * 20, win["gridW"] * 20), (w, a.shape)
disp = np.clip((a - 0.25) / 0.5, 0, 1) # model card: display rescale (p - 0.25) / 0.5
disp[a == 0] = 0
im = Image.fromarray((disp * 255).round().astype(np.uint8))
hi = im.copy()
hi.thumbnail((HI_MAX, HI_MAX), Image.LANCZOS)
hi.save(f"{OUT}/hi/{w}_gpu.jpg", quality=84, optimize=True, progressive=True)
tw = int(round(win["gridW"] * CELL_MM * PPMM * 1.5))
th = int(round(win["gridH"] * CELL_MM * PPMM * 1.5))
tex = local_contrast(im.resize((tw, th), Image.LANCZOS), 2.2 * PPMM * 1.5)
tex.save(f"{OUT}/gpu/{w}.jpg", quality=86, optimize=True)
g = receipt["gpu_seconds"]
p = receipt["patches"]["forward"]["nonempty"]
device = (receipt.get("device") or {}).get("device_name")
tot_patches += p
tot_measured += g["measured"]
tot_calls += g["calls"]
args = receipt["args"]
wins[w] = {
"tex": f"data/gpu/{w}.jpg",
"receipt": f"data/gpu/receipts/{w}.json",
"hi": f"hi/{w}_gpu.jpg",
"patches": p,
"gpuSeconds": round(g["measured"], 1),
"calls": g["calls"],
"overlap": args["overlap"],
"started": receipt["calls"][0]["started"],
"meta": (f"ink_9um (9 碌m): scrollprize/ink_9um (seed 42, step 75000) on the 9.4 碌m scan (9.362 碌m voxels) 路 "
f"{p:,} patches of 128脳128 px, overlap {args['overlap']}"
f"{' (patch seams visible)' if args['overlap'] < 0.5 else ''} 路 {g['measured']:.1f} s of GPU in "
f"{g['calls']} call{'s' if g['calls'] > 1 else ''} 路 {receipt['calls'][0]['started'][:10]} 路 "
f"drag to pan, scroll or pinch to zoom, double-click to fit"),
}
print(w, p, g["measured"], tex.size)
ids = sorted(wins)
full = [w for w in ids if wins[w]["overlap"] >= 0.5]
low = [w for w in ids if wins[w]["overlap"] < 0.5]
out = {
"windings": wins,
"totals": {"windings": len(wins), "patches": tot_patches, "gpuSeconds": round(tot_measured, 1),
"calls": tot_calls, "device": device},
"about": (f"<b>ink_9um (9 碌m)</b> is the cyan band. Chase Hendrick ran the public "
f"<a href=\"https://huggingface.co/scrollprize/ink_9um\" target=\"_blank\" rel=\"noopener\">"
f"scrollprize/ink_9um</a> model, which the Vesuvius Challenge trained, over windings {ids[0]} to {ids[-1]} "
f"and processed the output for this viewer: {tot_patches:,} patches in "
f"{tot_measured:.0f} s of measured GPU time on half an NVIDIA RTX PRO 6000 Blackwell. It reads the "
f"9.4 碌m scan (9.362 碌m voxels), about four times coarser than the 2.4 碌m scan behind the amber layer, so its ink map "
f"is noisier. {say(full)} ran with patch overlap 0.5; {say(low)} ran with less overlap, so "
f"their 128 px patch grid shows as faint square tiles. That grid is a stitching artifact, not ink. "
f"Click a cyan winding to compare the two maps."),
}
json.dump(out, open(f"{OUT}/gpu.json", "w"), indent=1)
print(json.dumps(out["totals"]))
if __name__ == "__main__":
main()
|