Spaces:
Running on Zero
Running on Zero
| """Lasagna Fields: make villa's lasagna field (cos, grad_mag, nx, ny) for one box of a public scan. | |
| Everything that is not the web page: the scan list (the open data's metadata.json), reading a box of CT from the | |
| open-data bucket, running a model on it, reading the field back, the agreement check against a published field, | |
| pictures and the download. app.py is the page and the API. | |
| The one model today is villa's own lasagna: its `predict3d` (lasagna/preprocess_cos_omezarr.py at a pinned commit) | |
| with the `scrollprize/lasagna` checkpoint, called in this process exactly as its command line would be, so the field | |
| is villa's field, not a copy of it. MODELS is the slot for other models (a faster student, a better net): each entry | |
| names its runner, and a model is offered only when its entry says it is released. | |
| python3 lasfield.py selftest the plumbing on synthetic CT (no network, no model) | |
| python3 lasfield.py run VOLUME Z Y X [SIZE] one field on this computer (CUDA if present, else CPU) | |
| """ | |
| import concurrent.futures | |
| import gzip | |
| import hashlib | |
| import io | |
| import json | |
| import math | |
| import os | |
| import re | |
| import shutil | |
| import subprocess | |
| import sys | |
| import tarfile | |
| import tempfile | |
| import threading | |
| import time | |
| import urllib.error | |
| import urllib.request | |
| import zipfile | |
| from pathlib import Path | |
| import numpy as np | |
| BUCKET = "https://vesuvius-challenge-open-data.s3.us-east-1.amazonaws.com" | |
| METADATA = BUCKET + "/metadata.json" | |
| USER_AGENT = "lasagna-fields-space (huggingface.co/spaces/SharpMeow/lasagna-fields)" | |
| # Scans this Space never reads: PHercParis4 is kept out of every public tool of ours (Chase's rule). | |
| REFUSED_SAMPLES = {"phercparis4"} | |
| CACHE = Path(os.environ.get("LASFIELD_CACHE") or Path(tempfile.gettempdir()) / "lasagna-fields") | |
| # villa's lasagna at the commit the ten whole-scan fields of this Space were made with (docs/tools/lasagna-fields.md); | |
| # lasagna/ imports vesuvius from vesuvius/src (train_unet_3d builds the net with vesuvius' NetworkFromConfig). | |
| VILLA = dict(repo="https://github.com/ScrollPrize/villa.git", | |
| tarball="https://codeload.github.com/ScrollPrize/villa/tar.gz/{commit}", | |
| commit="db3dbea26185a9e09ab997bfa21a2a00477f0d56", keep=("lasagna/", "vesuvius/src/"), | |
| check="lasagna/preprocess_cos_omezarr.py") | |
| WEIGHTS = dict(repo="scrollprize/lasagna", revision="de45f01147adea9f4ad0f621b1ccd885a1f03760", file="model_current.pt", | |
| sha256="4cc0c40a846b296769b9f7982470559078662a1a0b0e1f74ff41ed9341c0d600") | |
| TILE = 192 # predict3d's tile at this checkpoint (its patch size), stride TILE - overlap 64 | |
| MARGIN = TILE + 16 + 4 # CT read around the box: a tile, predict3d's border and the scaledown rounding | |
| # predict3d widens a crop to whole output chunks (64 voxels of the quarter grid: 256 input voxels), so boxes are whole | |
| # 256-voxel blocks of the input level: anything smaller costs the same, and the field matches villa's whole-scan run. | |
| ALIGN = 256 | |
| SIZES = (256, 512) | |
| DEFAULT_SIZE = 256 | |
| TARGET_UM = 9.4 # the model's training resolution: 1.2 m scans at level 0, 0.2 m (2.4 um) scans at level 2 | |
| CHANNELS = ("cos", "grad_mag", "nx", "ny") | |
| class FieldError(Exception): | |
| """A refusal or failure shown to the visitor as is.""" | |
| def log(message): | |
| print(f"[lasfield {time.strftime('%H:%M:%S')}] {message}", flush=True) | |
| # ---------------------------------------------------------------- http | |
| def http_get(url, timeout=60, tries=4): | |
| """Bytes at `url`, or None for a missing object (an unwritten chunk is a 403 or 404 on the bucket).""" | |
| err = None | |
| for i in range(tries): | |
| try: | |
| req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT}) | |
| with urllib.request.urlopen(req, timeout=timeout) as r: | |
| return r.read() | |
| except urllib.error.HTTPError as exc: | |
| if exc.code in (403, 404): | |
| return None | |
| err = exc | |
| except (urllib.error.URLError, OSError) as exc: | |
| err = exc | |
| time.sleep(1.5 * (i + 1)) | |
| raise FieldError(f"could not read {url}: {err}") | |
| def get_json(url): | |
| raw = http_get(url) | |
| if raw is None: | |
| raise FieldError(f"{url} does not exist") | |
| if raw[:2] == b"\x1f\x8b": | |
| raw = gzip.decompress(raw) | |
| return json.loads(raw) | |
| # ---------------------------------------------------------------- the scans | |
| _CATALOG = {"at": 0.0, "volumes": None} | |
| _LOCK = threading.Lock() | |
| def level_for(um): | |
| """The CT level whose voxels are nearest the model's 9.4 um (0 for 8.6 to 9.4 um scans, 2 for 2.4 um).""" | |
| return max(0, int(round(math.log2(TARGET_UM / float(um))))) if um else 0 | |
| def catalog(refresh_s=6 * 3600): | |
| """Every volume of the open data's metadata.json this Space may read, by volume id.""" | |
| with _LOCK: | |
| if _CATALOG["volumes"] is not None and time.time() - _CATALOG["at"] < refresh_s: | |
| return _CATALOG["volumes"] | |
| meta = get_json(METADATA) | |
| vols = {} | |
| for sid, s in sorted((meta.get("samples") or {}).items()): | |
| if sid.lower() in REFUSED_SAMPLES: | |
| continue | |
| for vid, v in (s.get("volumes") or {}).items(): | |
| pr = v.get("properties") or {} | |
| um, shape = pr.get("pixel_size_um"), pr.get("shape") | |
| if not um or not v.get("long_id"): | |
| continue | |
| vols[vid] = {"volume": vid, "sample": sid, "long_id": v["long_id"], "voxel_um": float(um), | |
| "shape_zyx": shape, "level": level_for(um), | |
| "superseded_by": pr.get("superseded_by"), | |
| "url": f"{BUCKET}/{sid}/volumes/{v['long_id']}/"} | |
| _CATALOG.update(at=time.time(), volumes=vols) | |
| return vols | |
| def volume(vid): | |
| """The catalog row for a volume id (the 14-digit id or its long name), or a refusal.""" | |
| vid = str(vid or "").strip().rstrip("/").split("/")[-1] | |
| vols = catalog() | |
| if vid in vols: | |
| return vols[vid] | |
| for v in vols.values(): | |
| if v["long_id"] == vid or v["long_id"].startswith(vid + "-"): | |
| return v | |
| if any(vid.lower().startswith(s) for s in REFUSED_SAMPLES): | |
| raise FieldError("this Space does not read PHercParis4") | |
| raise FieldError(f"{vid!r} is not a volume in the Vesuvius Challenge open data (metadata.json); " | |
| "give its 14-digit volume id, such as 20260102150214") | |
| def zarray(v, level): | |
| meta = get_json(f"{v['url']}{level}/.zarray") | |
| if len(meta.get("shape", ())) != 3: | |
| raise FieldError(f"{v['long_id']} level {level} is not a 3D array") | |
| return meta | |
| _PUBLISHED = {} | |
| # the ten whole-scan fields this Space hosts, made with the same villa commit and checkpoint as Make a field: volume id | |
| # to (sample, run folder) in the public bucket; the check uses them where the open data has no field | |
| OWN_BUCKET = "https://huggingface.co/buckets/SharpMeow/lasagna-fields/resolve" | |
| OWN_FIELDS = { | |
| "20250521115057": ("PHerc0175A", "20250521115057-lasagna-de45f011"), | |
| "20250521125822": ("PHerc0175B", "20250521125822-lasagna-de45f011"), | |
| "20250521133212": ("PHerc0306B", "20250521133212-lasagna-de45f011"), | |
| "20250521140913": ("PHerc0483A", "20250521140913-lasagna-de45f011"), | |
| "20251124083638": ("PHerc0483B", "20251124083638-lasagna-de45f011"), | |
| "20250521151210": ("PHerc0490A", "20250521151210-lasagna-de45f011"), | |
| "20250521151215": ("PHerc0490B", "20250521151215-lasagna-de45f011"), | |
| "20250821151531": ("PHerc0841", "20250821151531-lasagna-de45f011"), | |
| "20250728152254": ("PHerc0846A", "20250728152254-lasagna-de45f011"), | |
| "20250804142305": ("PHerc0846B", "20250804142305-lasagna-de45f011"), | |
| } | |
| def published_field(v): | |
| """The open data's own lasagna run for this volume at its default level, if there is one: its manifest's URL | |
| and groups. The run folder names start with the volume id; an -L2 suffix says it read level 2.""" | |
| key = v["volume"] | |
| if key in _PUBLISHED: | |
| return _PUBLISHED[key] | |
| found = _published_run(v) | |
| if found is None: | |
| # a remasked volume replaces an older one in the same voxel frame (metadata.json's superseded_by); the field | |
| # published for the older one is the one to check against, with the mask edge the only difference | |
| for old in catalog().values(): | |
| if (old.get("superseded_by") or {}).get("volume_id") == key and old["level"] == v["level"]: | |
| found = _published_run(old) | |
| if found: | |
| found = dict(found, note=f"published for {old['volume']}, which this volume replaces in the " | |
| "same voxel frame") | |
| break | |
| if found is None and key in OWN_FIELDS and v["level"] == 0: | |
| found = _own_run(*OWN_FIELDS[key]) | |
| _PUBLISHED[key] = found | |
| return found | |
| def _own_run(sample, run): | |
| """One of this Space's ten whole-scan fields, described like an open-data run.""" | |
| url = f"{OWN_BUCKET}/{sample}/representations/predictions/lasagna/{run}/" | |
| try: | |
| manifest = get_json(f"{url}{sample}.lasagna.json") | |
| except FieldError as exc: | |
| log(f"own field lookup for {sample} failed: {exc}") | |
| return None | |
| return {"run": run, "url": url, "manifest": f"{url}{sample}.lasagna.json", "groups": manifest.get("groups"), | |
| "base_shape_zyx": manifest.get("base_shape_zyx"), "source": "lasagna-fields", | |
| "note": "this Space's own whole-scan field, made with the same villa commit and checkpoint"} | |
| def _published_run(v): | |
| found = None | |
| try: | |
| prefix = f"{v['sample']}/representations/predictions/lasagna/" | |
| raw = http_get(f"{BUCKET}/?list-type=2&delimiter=/&prefix={prefix}") or b"" | |
| runs = re.findall(rb"<Prefix>([^<]+)</Prefix>", raw) | |
| for run in sorted(r.decode() for r in runs if r.decode() != prefix): | |
| name = run[len(prefix):].rstrip("/") | |
| tail = name.rsplit("-", 1)[-1] | |
| lvl = int(tail[1:]) if re.fullmatch(r"L\d", tail) else 0 | |
| if not name.startswith(v["volume"]) or lvl != v["level"]: | |
| continue | |
| listing = http_get(f"{BUCKET}/?list-type=2&delimiter=/&prefix={run}") or b"" | |
| js = [k.decode() for k in re.findall(rb"<Key>([^<]+\.lasagna\.json)</Key>", listing)] | |
| if js: | |
| manifest = get_json(f"{BUCKET}/{js[0]}") | |
| found = {"run": name, "url": f"{BUCKET}/{run}", "manifest": f"{BUCKET}/{js[0]}", | |
| "groups": manifest.get("groups"), "base_shape_zyx": manifest.get("base_shape_zyx"), | |
| "source": "open-data"} | |
| except FieldError as exc: | |
| log(f"published field lookup for {v['volume']} failed: {exc}") | |
| return found | |
| # ---------------------------------------------------------------- a box | |
| def plan(vid, z, y, x, size=DEFAULT_SIZE, level=None, model="lasagna"): | |
| """Where the field will be made: the `size`-voxel box of whole ALIGN blocks at the input level that holds | |
| level-0 voxel (z, y, x) (as near its middle as the grid allows), and the CT read around it. Refuses a point | |
| outside the scan or a size not offered.""" | |
| v = volume(vid) | |
| spec = MODELS.get(model) | |
| if not spec or not spec["offered"]: | |
| raise FieldError(f"model {model!r} is not offered; choose from {', '.join(k for k, m in MODELS.items() if m['offered'])}") | |
| size = int(size) | |
| if size not in SIZES: | |
| raise FieldError(f"size {size} is not offered; choose from {', '.join(map(str, SIZES))}") | |
| level = v["level"] if level in (None, "", "auto") else int(level) | |
| if not 0 <= level <= 5: | |
| raise FieldError("level must be 0 to 5") | |
| base = zarray(v, 0) | |
| meta = zarray(v, level) | |
| shape = [int(s) for s in meta["shape"]] | |
| f = 2 ** level | |
| center = [int(round(float(c))) for c in (z, y, x)] | |
| for c, n in zip(center, base["shape"]): | |
| if not 0 <= c < n: | |
| raise FieldError(f"the point {center} is outside the scan (level-0 shape {base['shape']})") | |
| origin = [max(0, min((c // f - size // 2 + ALIGN // 2) // ALIGN * ALIGN, (n - 1) // ALIGN * ALIGN)) | |
| for c, n in zip(center, shape)] | |
| box_size = [min(size, n - o) for o, n in zip(origin, shape)] | |
| read = [[max(0, o - MARGIN), min(n, o + s + MARGIN)] for o, s, n in zip(origin, box_size, shape)] | |
| ch = meta["chunks"] | |
| nchunks = math.prod((b - 1) // c - a // c + 1 for (a, b), c in zip(read, ch)) | |
| out = {"volume": v["volume"], "sample": v["sample"], "long_id": v["long_id"], "voxel_um": v["voxel_um"], | |
| "url": v["url"], "level": level, "level_shape_zyx": shape, "base_shape_zyx": base["shape"], | |
| "center_level0_zyx": center, "origin_zyx": origin, "size_zyx": box_size, | |
| "origin_level0_zyx": [o * f for o in origin], "size_level0_zyx": [s * f for s in box_size], | |
| "read_zyx": read, "chunks": nchunks, "chunk_bytes": math.prod(ch) * np.dtype(meta["dtype"]).itemsize, | |
| "model": model} | |
| out["tiles"] = tiles(out) | |
| out["gpu_s"] = gpu_seconds(out, model) | |
| return out | |
| def tiles(p_or_size, tile=TILE, overlap=64, border=16, sd_max=4): | |
| """Tiles predict3d runs for a box. predict3d puts tiles on one lattice over the whole scan (so a box's field is | |
| the whole-scan field there) and runs every lattice tile that touches the box plus its border; with villa loaded | |
| this asks villa's own _canonical_local_tile_positions, else it counts the same lattice here.""" | |
| if isinstance(p_or_size, dict): | |
| f = 2 ** p_or_size["level"] | |
| starts = p_or_size["origin_zyx"] | |
| sizes = p_or_size["size_zyx"] | |
| shape = p_or_size["level_shape_zyx"] | |
| else: | |
| sizes, starts, shape = list(p_or_size), [10 * 1024] * 3, [100 * 1024] * 3 | |
| stride = tile - overlap | |
| try: | |
| from tiled_predict3d import _canonical_local_tile_positions as positions | |
| except ImportError: | |
| positions = None | |
| n = 1 | |
| for s0, sz, vol in zip(starts, sizes, shape): | |
| padded = -(-(sz + 2 * border) // sd_max) * sd_max | |
| if positions is not None: | |
| n *= len(positions(volume_size=vol, crop_start=s0, crop_padded_size=padded, tile_size=tile, | |
| stride=stride, border=border, scaledown_multiple=sd_max)) | |
| else: | |
| first, last = max(0, (s0 - tile) // stride), (s0 + padded) // stride | |
| n *= sum(1 for k in range(first, last + 1) if k * stride - s0 < padded and k * stride - s0 + tile > 0) | |
| return n | |
| def gpu_seconds(p_or_size, model="lasagna"): | |
| """The GPU time to ask ZeroGPU for: the model's start-up plus its per-tile time, with a margin of 1.5.""" | |
| spec = MODELS[model] | |
| return int(math.ceil(1.5 * (spec["startup_s"] + spec["tile_s"] * tiles(p_or_size)))) | |
| # ---------------------------------------------------------------- the CT | |
| def stage(p, root): | |
| """The CT around the box on local disk, laid out like the open data: an OME-Zarr group whose level 0 is only its | |
| metadata (predict3d reads the scan's base shape from it) and whose input level holds the chunks it reads.""" | |
| root = Path(root) | |
| shutil.rmtree(root, ignore_errors=True) | |
| meta0 = get_json(f"{p['url']}0/.zarray") | |
| meta = get_json(f"{p['url']}{p['level']}/.zarray") | |
| zattrs = get_json(f"{p['url']}.zattrs") | |
| sep = meta.get("dimension_separator", ".") | |
| (root / "0").mkdir(parents=True) | |
| (root / "0" / ".zarray").write_text(json.dumps(meta0)) | |
| lvl = root / str(p["level"]) | |
| lvl.mkdir(parents=True, exist_ok=True) | |
| (lvl / ".zarray").write_text(json.dumps(meta)) | |
| (root / ".zgroup").write_text(json.dumps({"zarr_format": 2})) | |
| (root / ".zattrs").write_text(json.dumps(zattrs)) | |
| ch = meta["chunks"] | |
| keys = [(i, j, k) for i in range(p["read_zyx"][0][0] // ch[0], (p["read_zyx"][0][1] - 1) // ch[0] + 1) | |
| for j in range(p["read_zyx"][1][0] // ch[1], (p["read_zyx"][1][1] - 1) // ch[1] + 1) | |
| for k in range(p["read_zyx"][2][0] // ch[2], (p["read_zyx"][2][1] - 1) // ch[2] + 1)] | |
| def one(key): | |
| name = sep.join(map(str, key)) | |
| raw = http_get(f"{p['url']}{p['level']}/{name}") | |
| if raw is None: | |
| return 0 | |
| dest = lvl.joinpath(*name.split("/")) | |
| dest.parent.mkdir(parents=True, exist_ok=True) | |
| dest.write_bytes(raw) | |
| return len(raw) | |
| t = time.time() | |
| with concurrent.futures.ThreadPoolExecutor(32) as ex: | |
| sizes = list(ex.map(one, keys)) | |
| return {"chunks": len(keys), "written": sum(1 for s in sizes if s), "bytes": sum(sizes), | |
| "seconds": round(time.time() - t, 2), "input": str(lvl)} | |
| def read_box(arr_path, box): | |
| """One zarr array's values in a box [[z0, z1], [y0, y1], [x0, x1]] (zarr 3 reads villa's v2 output).""" | |
| import zarr | |
| a = zarr.open_array(str(arr_path), mode="r") | |
| return np.asarray(a[box[0][0]:box[0][1], box[1][0]:box[1][1], box[2][0]:box[2][1]]) | |
| def ct_box(p, root): | |
| lvl = Path(root) / str(p["level"]) | |
| import zarr | |
| a = zarr.open_array(str(lvl), mode="r") | |
| o, s = p["origin_zyx"], p["size_zyx"] | |
| return np.asarray(a[o[0]:o[0] + s[0], o[1]:o[1] + s[1], o[2]:o[2] + s[2]]) | |
| # ---------------------------------------------------------------- code and weights | |
| def sha256_file(path, block=1 << 22): | |
| h = hashlib.sha256() | |
| with open(path, "rb") as fh: | |
| for part in iter(lambda: fh.read(block), b""): | |
| h.update(part) | |
| return h.hexdigest() | |
| def _fetch_villa(dest): | |
| """villa's lasagna/ and vesuvius/src/ at the pinned commit: a sparse, blobless git fetch, else GitHub's tarball.""" | |
| commit = VILLA["commit"] | |
| 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", *[k.rstrip("/") for k in VILLA["keep"]]], | |
| ["fetch", "-q", "--depth", "1", "--filter=blob:none", "origin", commit], | |
| ["checkout", "-q", "FETCH_HEAD"]): | |
| subprocess.run(git + cmd, capture_output=True, text=True, check=True, timeout=900) | |
| head = subprocess.run(git + ["rev-parse", "HEAD"], capture_output=True, text=True, check=True).stdout.strip() | |
| if head != commit: | |
| raise RuntimeError(f"git fetched {head}, not {commit}") | |
| shutil.rmtree(tmp / ".git", ignore_errors=True) | |
| except (OSError, RuntimeError, subprocess.SubprocessError) as exc: | |
| log(f"git fetch of villa failed ({getattr(exc, 'stderr', None) or exc}); trying the tarball") | |
| shutil.rmtree(tmp, ignore_errors=True) | |
| tmp.mkdir(parents=True) | |
| root = tmp.resolve() | |
| req = urllib.request.Request(VILLA["tarball"].format(commit=commit), headers={"User-Agent": USER_AGENT}) | |
| with urllib.request.urlopen(req, timeout=900) as resp, tarfile.open(fileobj=resp, mode="r|gz") as tf: | |
| for m in tf: | |
| _, _, rel = m.name.partition("/") | |
| if not rel.startswith(VILLA["keep"]) or m.issym() or m.islnk(): | |
| continue | |
| target = (root / rel).resolve() | |
| if root not in target.parents: | |
| 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) | |
| if not (tmp / VILLA["check"]).is_file(): | |
| raise RuntimeError(f"villa at {commit} has no {VILLA['check']}") | |
| (tmp / "COMMIT").write_text(commit + "\n") | |
| shutil.rmtree(dest, ignore_errors=True) | |
| os.replace(tmp, dest) | |
| def ensure_villa(cache=CACHE): | |
| """villa's lasagna on sys.path (as top-level modules and as the `lasagna` package, the way its own install lays | |
| them out) with vesuvius/src beside it. LASFIELD_VILLA names a local villa checkout instead.""" | |
| if os.environ.get("LASFIELD_VILLA"): | |
| root = Path(os.environ["LASFIELD_VILLA"]).expanduser().resolve() | |
| source = "local checkout (LASFIELD_VILLA)" | |
| else: | |
| root = Path(cache) / f"villa-{VILLA['commit'][:12]}" | |
| if not (root / "COMMIT").is_file(): | |
| log(f"fetching villa {VILLA['commit'][:12]}") | |
| _fetch_villa(root) | |
| source = "pinned download" | |
| shim = Path(cache) / "villa-shim" | |
| shim.mkdir(parents=True, exist_ok=True) | |
| link = shim / "lasagna" | |
| if not link.exists(): | |
| link.symlink_to(root / "lasagna", target_is_directory=True) | |
| for path in (root / "vesuvius" / "src", shim, root / "lasagna"): | |
| if str(path) not in sys.path: | |
| sys.path.insert(0, str(path)) | |
| os.environ["PYTHONPATH"] = os.pathsep.join([str(root / "lasagna"), str(shim), str(root / "vesuvius" / "src"), | |
| os.environ.get("PYTHONPATH", "")]).rstrip(os.pathsep) | |
| return {"source": source, "commit": VILLA["commit"], "path": str(root)} | |
| def ensure_weights(cache=CACHE): | |
| """The checkpoint from the Hub at its pinned revision, refused unless its SHA-256 matches. LASFIELD_WEIGHTS names | |
| a local copy of model_current.pt instead (its SHA-256 is still checked).""" | |
| if os.environ.get("LASFIELD_WEIGHTS"): | |
| path = Path(os.environ["LASFIELD_WEIGHTS"]).expanduser() | |
| else: | |
| from huggingface_hub import hf_hub_download | |
| path = Path(hf_hub_download(WEIGHTS["repo"], WEIGHTS["file"], revision=WEIGHTS["revision"], | |
| cache_dir=str(Path(cache) / "hub"))) | |
| digest = sha256_file(path) | |
| if digest != WEIGHTS["sha256"]: | |
| raise RuntimeError(f"{WEIGHTS['repo']}/{WEIGHTS['file']} has SHA-256 {digest}, not {WEIGHTS['sha256']}") | |
| return path | |
| # ---------------------------------------------------------------- the models | |
| def predict3d_argv(p, ct_input, out_json, checkpoint, device): | |
| """villa predict3d's command line for the box, as its CLI takes it: --crop is x y z w h d in the scan's base | |
| (level 0) voxels, which predict3d divides by the input level's factor (its log: base_crop, input_crop). | |
| Synchronous flush and accumulation (no worker processes): the stored field is the same.""" | |
| o, s = p["origin_level0_zyx"], p["size_level0_zyx"] | |
| return ["predict3d", "--input", str(ct_input), "--output", str(out_json), "--unet-checkpoint", str(checkpoint), | |
| "--tile-size", str(TILE), "--crop", str(o[2]), str(o[1]), str(o[0]), str(s[2]), str(s[1]), str(s[0]), | |
| "--device", device, "--no-download", "--flush-workers", "0", "--accumulator-workers", "0"] | |
| POOL_CPUS = 8 # pyramid processes predict3d may start (its default is every CPU of the host) | |
| def _pool_friendly(): | |
| """Let predict3d's OME-Zarr pyramid step start its process pool when it runs inside a daemon process (ZeroGPU | |
| runs a GPU call in one, and Python refuses children there), and cap that pool at POOL_CPUS. The pool's processes | |
| only downsample chunks on the CPU. Restored on exit.""" | |
| import multiprocessing as mp | |
| proc = mp.current_process() | |
| was_daemon = proc._config.get("daemon") | |
| real = mp.cpu_count | |
| mp.cpu_count = lambda: max(1, min(POOL_CPUS, real())) | |
| if was_daemon: | |
| proc._config["daemon"] = False | |
| try: | |
| yield | |
| finally: | |
| mp.cpu_count = real | |
| if was_daemon: | |
| proc._config["daemon"] = was_daemon | |
| def run_villa_lasagna(p, ct_input, out_dir, device, checkpoint): | |
| """villa's predict3d in this process. Returns its exit code, wall seconds and the tail of its log.""" | |
| import contextlib | |
| import preprocess_cos_omezarr as pco | |
| out_dir = Path(out_dir) | |
| shutil.rmtree(out_dir, ignore_errors=True) | |
| out_dir.mkdir(parents=True) | |
| out_json = out_dir / f"{p['sample']}-{p['volume']}.lasagna.json" | |
| argv = predict3d_argv(p, ct_input, out_json, checkpoint, device) | |
| buf = io.StringIO() | |
| t = time.time() | |
| class Tee(io.TextIOBase): | |
| def write(self, s): | |
| buf.write(s) | |
| sys.__stdout__.write(s) | |
| return len(s) | |
| try: | |
| with contextlib.redirect_stdout(Tee()), contextlib.redirect_stderr(Tee()), _pool_friendly(): | |
| rc = pco.cli_main(argv) | |
| except SystemExit as exc: | |
| rc = exc.code if isinstance(exc.code, int) else 1 | |
| return {"exit": int(rc or 0), "wall_s": round(time.time() - t, 2), "argv": argv, | |
| "log_tail": buf.getvalue()[-6000:], "manifest": str(out_json)} | |
| # The model slot. `offered` False keeps a model off the page and the API until it passes its own held-out test; | |
| # Lasagnetta (the lasagna student) and any better net come in here with their own runner, checkpoint and costs. | |
| MODELS = { | |
| "lasagna": {"title": "villa lasagna (exact)", "offered": True, "run": run_villa_lasagna, | |
| "about": "villa's own predict3d with the scrollprize/lasagna checkpoint (MIT): the same field villa " | |
| "makes, at the same grids (cos at half, the rest at a quarter of the input level)", | |
| # 25 Oct 2026, L4: 310 ms a 192^3 tile end to end; the GPU here is faster, so these are ceilings | |
| "startup_s": 25, "tile_s": 0.35}, | |
| } | |
| # ---------------------------------------------------------------- the field | |
| def field_boxes(p, manifest): | |
| """Each channel's array path and the box covering the plan's box in that channel's grid.""" | |
| man = json.loads(Path(manifest).read_text()) | |
| root = Path(manifest).parent | |
| out = {} | |
| for name in CHANNELS: | |
| g = man["groups"][name] | |
| sd = 2 ** int(g["scaledown"]) | |
| box = [[o // sd, -(-(o + s) // sd)] for o, s in zip(p["origin_level0_zyx"], p["size_level0_zyx"])] | |
| out[name] = {"path": root / g["zarr"], "scaledown": int(g["scaledown"]), "box": box} | |
| return out, man | |
| def read_field(p, manifest): | |
| boxes, man = field_boxes(p, manifest) | |
| arrays = {n: read_box(b["path"], b["box"]) for n, b in boxes.items()} | |
| return arrays, boxes, man | |
| def decode_normal(nx_u8, ny_u8): | |
| """The unit normal from stored nx, ny (villa's encode_normal_nxny_u8 flips n so that nz >= 0).""" | |
| nx = (nx_u8.astype(np.float32) - 128.0) / 127.0 | |
| ny = (ny_u8.astype(np.float32) - 128.0) / 127.0 | |
| nz = np.sqrt(np.clip(1.0 - nx * nx - ny * ny, 0.0, 1.0)) | |
| n = np.stack([nx, ny, nz]) | |
| return n / (np.linalg.norm(n, axis=0, keepdims=True) + 1e-8) | |
| def _corr(a, b, mask): | |
| a, b = a[mask].astype(np.float64), b[mask].astype(np.float64) | |
| if a.size < 100 or a.std() == 0 or b.std() == 0: | |
| return None | |
| return round(float(np.corrcoef(a, b)[0, 1]), 4) | |
| def compare(mine, theirs): | |
| """Agreement of two fields on the same grids: Pearson r of cos and grad_mag where either is non-zero, and the | |
| median angle between normals (sign ignored) where both are, over all those voxels and over the strong half (our | |
| grad_mag at or above its median there). Where the layer signal is weak the normal is barely defined and tiny | |
| numerical differences turn it, so the strong half is the fair test.""" | |
| m_cos = (mine["cos"] > 0) | (theirs["cos"] > 0) | |
| m_oth = (mine["grad_mag"] > 0) | (theirs["grad_mag"] > 0) | |
| both = (mine["nx"] > 0) & (theirs["nx"] > 0) | |
| out = {"cos_r": _corr(mine["cos"], theirs["cos"], m_cos), | |
| "grad_mag_r": _corr(mine["grad_mag"], theirs["grad_mag"], m_oth), "normal_median_deg": None, | |
| "voxels": int(m_cos.sum())} | |
| if both.sum() >= 100: | |
| a, b = decode_normal(mine["nx"], mine["ny"]), decode_normal(theirs["nx"], theirs["ny"]) | |
| dot = np.abs((a * b).sum(axis=0)) | |
| out["normal_median_deg"] = round(float(np.degrees(np.arccos(np.clip(np.median(dot[both]), 0, 1)))), 2) | |
| g = mine["grad_mag"] | |
| strong = both & (g >= np.median(g[both])) | |
| if strong.sum() >= 100: | |
| out["normal_median_deg_strong"] = round( | |
| float(np.degrees(np.arccos(np.clip(np.median(dot[strong]), 0, 1)))), 2) | |
| return out | |
| def read_published(pub, boxes, shift=(0, 0, 0)): | |
| """The published field's values in the same boxes, optionally shifted (in each channel's voxels).""" | |
| import zarr | |
| out = {} | |
| for name, b in boxes.items(): | |
| g = (pub["groups"] or {}).get(name) | |
| if not g or int(g["scaledown"]) != b["scaledown"]: | |
| return None | |
| a = zarr.open_array(pub["url"] + g["zarr"], mode="r", zarr_format=2) | |
| box = [[lo + d, hi + d] for (lo, hi), d in zip(b["box"], shift)] | |
| if any(lo < 0 or hi > n for (lo, hi), n in zip(box, a.shape)): | |
| return None | |
| out[name] = np.asarray(a[box[0][0]:box[0][1], box[1][0]:box[1][1], box[2][0]:box[2][1]]) | |
| return out | |
| CONTROL_SHIFT = 48 # input-level voxels: the control reads the published field this far away in y and x | |
| OPEN_DATA_NORMALS = ("cos and grad_mag should match the open data's field exactly. Its normals can differ in " | |
| "parts of a box: villa changed how lasagna estimates normals after the open data's fields " | |
| "were made (April 2026, by their run names), and Make a field runs the later code, the same " | |
| "as this Space's ten whole-scan fields, which it matches exactly") | |
| def agreement(p, arrays, boxes): | |
| """How close this field is to the open data's own lasagna field for the scan, where one exists, with a control | |
| that must score lower: the published field read CONTROL_SHIFT input voxels away along y and x (a few papyrus | |
| layers), the same physical offset on every channel's grid.""" | |
| pub = published_field(volume(p["volume"])) | |
| if not pub: | |
| return {"published": None, "note": "the open data has no lasagna field for this volume at this level"} | |
| shift0 = CONTROL_SHIFT * 2 ** p["level"] | |
| try: | |
| same = read_published(pub, boxes) | |
| shifted = {} | |
| for name, b in boxes.items(): | |
| d = shift0 // 2 ** b["scaledown"] | |
| r = read_published(pub, {name: b}, (0, d, d)) | |
| if r is None: | |
| shifted = None | |
| break | |
| shifted[name] = r[name] | |
| except Exception as exc: # noqa: BLE001 (reported, never fatal) | |
| return {"published": pub["run"], "error": f"{type(exc).__name__}: {exc}"[:300]} | |
| if same is None: | |
| return {"published": pub["run"], "note": "the published field uses other grids; not compared"} | |
| out = {"published": pub["run"], "published_url": pub["url"], "published_source": pub.get("source"), | |
| "same_box": compare(arrays, same)} | |
| if pub.get("note"): | |
| out["published_note"] = pub["note"] | |
| if pub.get("source") == "open-data": | |
| out["normals_note"] = OPEN_DATA_NORMALS | |
| if shifted is not None: | |
| out["control_shifted"] = compare(arrays, shifted) | |
| out["control_note"] = (f"the published field {CONTROL_SHIFT} input voxels away in y and x: " | |
| "it must agree less than the same box") | |
| return out | |
| # ---------------------------------------------------------------- pictures and the download | |
| def _u8(img, lo=None, hi=None): | |
| img = img.astype(np.float32) | |
| lo = np.percentile(img, 1) if lo is None else lo | |
| hi = np.percentile(img, 99) if hi is None else hi | |
| return np.clip((img - lo) / max(hi - lo, 1e-6) * 255, 0, 255).astype(np.uint8) | |
| def _up(img, n): | |
| return np.repeat(np.repeat(img, n, axis=0), n, axis=1) if n > 1 else img | |
| def _resize(img, shape): | |
| """An image (H, W) or (H, W, 3) uint8 resized to `shape` (H, W) with bilinear smoothing, for display only.""" | |
| from PIL import Image | |
| return np.asarray(Image.fromarray(img).resize((int(shape[1]), int(shape[0])), Image.BILINEAR)) | |
| def picture(ct, arrays, boxes, axis=0): | |
| """One middle slice through the box along `axis` (0 z, 1 y, 2 x): CT, cos, grad_mag and the normal as color | |
| (|x| red, |y| green, |z| blue), side by side at the CT's grid. Returns an RGB uint8 array.""" | |
| def mid(a): | |
| return np.take(a, a.shape[axis] // 2, axis=axis) | |
| ct_img = _u8(mid(ct)) | |
| shape = ct_img.shape | |
| panels = [np.stack([ct_img] * 3, -1)] | |
| for name in ("cos", "grad_mag"): | |
| img = _u8(mid(arrays[name]), 0, 255 if name == "cos" else None) | |
| panels.append(np.stack([_resize(img, shape)] * 3, -1)) | |
| n = decode_normal(mid(arrays["nx"]), mid(arrays["ny"])) | |
| rgb = (np.abs(np.moveaxis(n, 0, -1)) * 255).astype(np.uint8) | |
| rgb[mid(arrays["nx"]) == 0] = 0 | |
| panels.append(_resize(rgb, shape)) | |
| gap = np.full((shape[0], 6, 3), 255, np.uint8) | |
| row = [] | |
| for i, t in enumerate(panels): | |
| row += [t] + ([gap] if i < len(panels) - 1 else []) | |
| return np.concatenate(row, axis=1) | |
| PANELS = ("CT", "cos", "grad_mag", "normal") | |
| LEGEND = "normal as color: x red, y green, z blue · villa lasagna, model output" | |
| def compose(pictures, title="", note=LEGEND): | |
| """The three rows (z, y, x middle slices) under one header naming the panels, each row named on its left, on a | |
| white card. Text is sized for the card to be shown about 700 px wide. Returns an RGB uint8 array.""" | |
| from PIL import Image, ImageDraw, ImageFont | |
| rows = [_fit(r, pictures[0].shape[:2]) if r.shape != pictures[0].shape else r for r in pictures] | |
| h, w = rows[0].shape[:2] | |
| k = max(1.0, w / 640) | |
| try: | |
| font = ImageFont.load_default(size=round(14 * k)) | |
| bold = ImageFont.load_default(size=round(16 * k)) | |
| except TypeError: # Pillow before 10.1 | |
| font = bold = ImageFont.load_default() | |
| pw = (w - 18) // 4 | |
| left, gap = round(66 * k), round(8 * k) | |
| top = round(26 * k) + (round(30 * k) if title else 0) | |
| foot = round(26 * k) if note else 0 | |
| canvas = Image.new("RGB", (left + w + gap, top + 3 * h + 3 * gap + foot), (255, 255, 255)) | |
| d = ImageDraw.Draw(canvas) | |
| if title: | |
| d.text((left, round(6 * k)), title, fill=(35, 32, 28), font=bold) | |
| for i, name in enumerate(PANELS): | |
| d.text((left + i * (pw + 6) + 2, top - round(21 * k)), name, fill=(93, 85, 75), font=font) | |
| for i, (r, name) in enumerate(zip(rows, ("z slice", "y slice", "x slice"))): | |
| y0 = top + i * (h + gap) | |
| canvas.paste(Image.fromarray(r), (left, y0)) | |
| d.text((round(6 * k), y0 + h // 2 - round(8 * k)), name, fill=(93, 85, 75), font=font) | |
| if note: | |
| d.text((left, top + 3 * h + 2 * gap + round(4 * k)), note, fill=(93, 85, 75), font=font) | |
| return np.asarray(canvas) | |
| def png(rgb): | |
| from PIL import Image | |
| buf = io.BytesIO() | |
| Image.fromarray(rgb).save(buf, format="PNG", optimize=True) | |
| return buf.getvalue() | |
| def package(out_dir, p, arrays, receipt, dest): | |
| """The download: villa's own output folder (its .lasagna.json and four OME-Zarr channels, holding only this | |
| box, in the scan's base coordinates, so lasagna's fit opens it as is), field.npz with the box's arrays, and the | |
| receipt.""" | |
| dest = Path(dest) | |
| with zipfile.ZipFile(dest, "w", zipfile.ZIP_DEFLATED, compresslevel=6) as zf: | |
| for f in sorted(Path(out_dir).rglob("*")): | |
| if f.is_file(): | |
| zf.write(f, f"villa/{f.relative_to(out_dir)}") | |
| buf = io.BytesIO() | |
| np.savez_compressed(buf, **arrays) | |
| zf.writestr("field.npz", buf.getvalue()) | |
| zf.writestr("receipt.json", json.dumps(receipt, indent=1)) | |
| zf.writestr("README.txt", README_TXT) | |
| return dest | |
| README_TXT = """Lasagna field for one box, made by the Lasagna Fields Space (huggingface.co/spaces/SharpMeow/lasagna-fields). | |
| villa/ villa lasagna predict3d's own output: SCAN.lasagna.json and SCAN_{cos,grad_mag,nx,ny}.ome.zarr. Only the | |
| chunks of this box are written; the arrays have the whole scan's shape, in its base (level 0) frame, so | |
| villa's lasagna fit (lasagna/fit.py) opens the manifest as is. | |
| field.npz the box as numpy arrays: cos and grad_mag, nx and ny (uint8, as villa stores them). Grids: see receipt.json | |
| (boxes: each channel's voxel box and scaledown from level 0). | |
| receipt.json what ran: scan, box, model, villa commit, checkpoint SHA-256, timings, and the agreement check. | |
| Encodings (villa): cos = clip(cos * 255); grad_mag = clip(grad_mag * 1000 * grad_mag_factor); nx, ny = the unit | |
| normal's x and y times 127 plus 128, with the normal flipped so its z is not negative. | |
| Scan data: Vesuvius Challenge open data, CC BY-NC 4.0. Model and code: ScrollPrize/villa lasagna, MIT. | |
| This is model output: geometry inputs, not a reading of any text. | |
| """ | |
| # ---------------------------------------------------------------- one run, end to end (no GPU decorator here) | |
| def prepare(p, work): | |
| """Everything before the GPU: code, weights and the CT on disk.""" | |
| t = time.time() | |
| code = ensure_villa() | |
| ck = ensure_weights() | |
| st = stage(p, Path(work) / "ct") | |
| return {"code": code, "checkpoint": str(ck), "stage": st, "prepare_s": round(time.time() - t, 2)} | |
| def finish(p, prep, ran, work, out_root): | |
| """Everything after the GPU: read the field, compare, draw, pack.""" | |
| if ran["exit"] != 0: | |
| raise FieldError(f"villa predict3d exited {ran['exit']}: {ran['log_tail'][-800:]}") | |
| arrays, boxes, man = read_field(p, ran["manifest"]) | |
| if not any(int(a.max()) for a in arrays.values()): | |
| raise FieldError("the field is empty: the box is probably outside the scroll (masked CT is all zero there)") | |
| ct = ct_box(p, Path(work) / "ct") | |
| pictures = [picture(ct, arrays, boxes, axis) for axis in (0, 1, 2)] | |
| agree = agreement(p, arrays, boxes) | |
| receipt = {"made_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), "plan": p, | |
| "model": {"id": p["model"], "title": MODELS[p["model"]]["title"], | |
| "villa_commit": VILLA["commit"], "checkpoint": f"{WEIGHTS['repo']}@{WEIGHTS['revision']}", | |
| "checkpoint_file": WEIGHTS["file"], "checkpoint_sha256": WEIGHTS["sha256"]}, | |
| "predict3d_argv": [a if not a.startswith("/") else Path(a).name for a in ran["argv"]], | |
| "manifest": {k: man[k] for k in man if k != "groups"}, | |
| "boxes": {n: {"scaledown": b["scaledown"], "box_zyx": b["box"]} for n, b in boxes.items()}, | |
| "timings": {"stage_s": prep["stage"]["seconds"], "stage_bytes": prep["stage"]["bytes"], | |
| "model_s": ran["wall_s"]}, | |
| "agreement": agree, | |
| "claim": "Model output: a lasagna field is geometry fitted by a model, not a reading of any text."} | |
| out_root = Path(out_root) | |
| out_root.mkdir(parents=True, exist_ok=True) | |
| stem = f"{p['sample']}-{p['volume']}-z{p['center_level0_zyx'][0]}-y{p['center_level0_zyx'][1]}-x{p['center_level0_zyx'][2]}-s{p['size_zyx'][0]}" | |
| zpath = package(Path(ran["manifest"]).parent, p, arrays, receipt, out_root / f"{stem}.zip") | |
| return {"receipt": receipt, "zip": str(zpath), "pictures": pictures, "arrays": arrays} | |
| def make(p, device=None, work=None, out_root=None, runner=None): | |
| """One field on this computer: stage, run, finish. `runner(p, ct_input, out_dir, device, checkpoint)` replaces | |
| the model call (the Space passes its ZeroGPU-decorated one).""" | |
| import torch | |
| device = device or ("cuda" if torch.cuda.is_available() else "cpu") | |
| work = Path(work or tempfile.mkdtemp(prefix="lasfield-")) | |
| prep = prepare(p, work) | |
| run = runner or MODELS[p["model"]]["run"] | |
| ran = run(p, prep["stage"]["input"], work / "out", device, prep["checkpoint"]) | |
| return finish(p, prep, ran, work, out_root or work / "dl") | |
| # ---------------------------------------------------------------- self-test (no network, no model) | |
| def selftest(): | |
| """The parts that need neither network nor model: grids, the tile count, the GPU ask, comparison and the | |
| control, decoding and pictures. Each check can fail.""" | |
| ok = [] | |
| assert level_for(9.362) == 0 and level_for(8.64) == 0 and level_for(2.399) == 2 and level_for(1.129) == 3 | |
| ok.append("levels") | |
| assert tiles([256] * 3) < tiles([512] * 3) | |
| assert gpu_seconds([256] * 3) < gpu_seconds([512] * 3) | |
| ok.append("tiles") | |
| rng = np.random.default_rng(0) | |
| z = np.arange(64)[:, None, None] * np.ones((1, 64, 64)) | |
| cos = (np.cos(z / 3.0) * 127 + 128).astype(np.uint8) | |
| f = {"cos": cos, "grad_mag": (cos // 2 + 1).astype(np.uint8), "nx": np.full_like(cos, 128), | |
| "ny": np.full_like(cos, 128)} | |
| g = {k: np.clip(v.astype(int) + rng.integers(-3, 4, v.shape), 1, 255).astype(np.uint8) for k, v in f.items()} | |
| c_same = compare(f, g) | |
| c_shift = compare(f, {k: np.roll(v, 5, axis=0) for k, v in g.items()}) | |
| assert c_same["cos_r"] > 0.99 and c_shift["cos_r"] < 0.5, (c_same, c_shift) | |
| assert c_same["normal_median_deg"] < 3, c_same | |
| ok.append("compare and control") | |
| n = decode_normal(np.array([255], np.uint8), np.array([128], np.uint8)) | |
| assert abs(n[0, 0] - 1) < 0.01 and abs(n[2, 0]) < 0.1 | |
| ok.append("normals") | |
| boxes = {"cos": {"scaledown": 1}, "grad_mag": {"scaledown": 2}, "nx": {"scaledown": 2}, "ny": {"scaledown": 2}} | |
| ct = rng.integers(0, 255, (128, 128, 128)).astype(np.uint8) | |
| arr = {"cos": cos, "grad_mag": f["grad_mag"][::2, ::2, ::2], "nx": f["nx"][::2, ::2, ::2], | |
| "ny": f["ny"][::2, ::2, ::2]} | |
| pic = picture(ct, arr, boxes, 0) | |
| assert pic.shape == (128, 4 * 128 + 18, 3), pic.shape | |
| card = compose([pic, pic, pic], "test") | |
| assert card.shape[0] > 3 * 128 and card.shape[1] > pic.shape[1], card.shape | |
| ok.append("pictures") | |
| argv = predict3d_argv({"origin_level0_zyx": [1, 2, 3], "size_level0_zyx": [4, 5, 6]}, "in", "out.json", "ck.pt", | |
| "cuda") | |
| assert argv[argv.index("--crop") + 1:argv.index("--crop") + 7] == ["3", "2", "1", "6", "5", "4"] | |
| ok.append("predict3d crop order") | |
| return {"passed": True, "checks": ok} | |
| def main(argv=None): | |
| argv = list(sys.argv[1:] if argv is None else argv) | |
| if not argv or argv[0] == "selftest": | |
| print(json.dumps(selftest(), indent=1)) | |
| return 0 | |
| if argv[0] == "run": | |
| vid, z, y, x = argv[1:5] | |
| size = int(argv[5]) if len(argv) > 5 else DEFAULT_SIZE | |
| p = plan(vid, z, y, x, size) | |
| log(json.dumps(p)) | |
| r = make(p, out_root=os.environ.get("LASFIELD_OUT") or ".") | |
| print(json.dumps(r["receipt"], indent=1)) | |
| print(r["zip"]) | |
| return 0 | |
| print(__doc__) | |
| return 2 | |
| if __name__ == "__main__": | |
| sys.exit(main()) | |