| """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 |
|
|
| OUT = os.environ.get("UNROLL_OUT", f"{ROOT}/data") |
| GPU_IN = os.environ.get("UNROLL_GPU", f"{ROOT}/gpu") |
| |
| 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) |
| 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() |
|
|
