Spaces:
Running on Zero
Running on Zero
| """Read what someone uploads into a uint8 array: one 2D layer (H, W) or a layer stack (L, H, W). | |
| Accepted: a single image (PNG, JPEG, WebP, BMP, one-page TIFF), a multi-page TIFF of surface layers, a ZIP of layer | |
| images (ordered by the numbers in their names, as villa and VC3D name them: 00.tif, 01.tif, ...), or an .npz with an | |
| array "volume" (Ink Lens's region format). Color is turned to gray, 16-bit is scaled to 8-bit. Nothing is kept. | |
| """ | |
| from __future__ import annotations | |
| import io | |
| import re | |
| import zipfile | |
| from pathlib import Path | |
| import numpy as np | |
| MIN_SIDE, MAX_SIDE = 128, 2048 | |
| MAX_LAYERS = 65 | |
| MAX_BYTES = 400 * 1024 * 1024 | |
| IMAGE_EXT = (".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff") | |
| class UploadError(ValueError): | |
| """Shown to the person as it is.""" | |
| def to_u8(a): | |
| """Gray uint8 from a 2D or (H, W, C) array of any integer or float type.""" | |
| a = np.asarray(a) | |
| if a.ndim == 3 and a.shape[-1] in (3, 4): | |
| rgb_u8 = a.dtype == np.uint8 | |
| a = a[..., :3].astype(np.float32) @ np.array([0.299, 0.587, 0.114], np.float32) | |
| if rgb_u8: | |
| return np.clip(a + 0.5, 0, 255).astype(np.uint8) | |
| if a.ndim != 2: | |
| raise UploadError(f"expected a 2D image, got shape {tuple(a.shape)}") | |
| if a.dtype == np.uint8: | |
| return a | |
| if a.dtype == np.uint16: | |
| return (a // 257).astype(np.uint8) | |
| is_float = a.dtype.kind == "f" | |
| a = np.nan_to_num(a.astype(np.float32)) | |
| lo, hi = float(a.min()), float(a.max()) | |
| if is_float and lo >= 0 and hi <= 1.0: | |
| return np.clip(a * 255 + 0.5, 0, 255).astype(np.uint8) | |
| if lo >= 0 and hi <= 255: | |
| return np.clip(a + 0.5, 0, 255).astype(np.uint8) | |
| if hi <= lo: | |
| return np.zeros(a.shape, np.uint8) | |
| return np.clip((a - lo) / (hi - lo) * 255 + 0.5, 0, 255).astype(np.uint8) | |
| def _image(data, name): | |
| if name.lower().endswith((".tif", ".tiff")): | |
| import tifffile | |
| a = tifffile.imread(io.BytesIO(data)) | |
| if a.ndim == 2 or (a.ndim == 3 and a.shape[-1] in (3, 4)): | |
| return [to_u8(a)] | |
| if a.ndim in (3, 4): | |
| return [to_u8(l) for l in a] | |
| raise UploadError(f"cannot read a TIFF of shape {tuple(a.shape)} as layers") | |
| from PIL import Image | |
| with Image.open(io.BytesIO(data)) as im: | |
| return [to_u8(np.asarray(im.convert("L") if im.mode not in ("L", "I;16", "I", "F") else im))] | |
| def _natural(name): | |
| return [int(t) if t.isdigit() else t.lower() for t in re.split(r"(\d+)", Path(name).name)] | |
| def read(path): | |
| """(array, kind, note): kind is "image" for one layer, "stack" for several.""" | |
| path = Path(path) | |
| size = path.stat().st_size | |
| if size > MAX_BYTES: | |
| raise UploadError(f"the file is {size / 2**20:.0f} MB; the limit is {MAX_BYTES // 2**20} MB") | |
| name = path.name.lower() | |
| if name.endswith(".npz"): | |
| with np.load(path, allow_pickle=False) as z: | |
| if "volume" not in z: | |
| raise UploadError('an .npz needs an array named "volume" (layers, height, width)') | |
| layers = [to_u8(v) for v in np.asarray(z["volume"])] | |
| note = f"{len(layers)} layers from the .npz" | |
| elif name.endswith(".zip"): | |
| with zipfile.ZipFile(path) as zf: | |
| names = sorted((n for n in zf.namelist() if n.lower().endswith(IMAGE_EXT) | |
| and not Path(n).name.startswith(".") and "__MACOSX" not in n), key=_natural) | |
| if not names: | |
| raise UploadError("the ZIP holds no layer images (.tif, .png or .jpg)") | |
| if len(names) > MAX_LAYERS: | |
| raise UploadError(f"the ZIP holds {len(names)} images; send at most {MAX_LAYERS} layers") | |
| layers = [] | |
| for n in names: | |
| layers.extend(_image(zf.read(n), n)) | |
| note = f"{len(layers)} layers from the ZIP, in name order ({Path(names[0]).name} first)" | |
| elif name.endswith(IMAGE_EXT): | |
| layers = _image(path.read_bytes(), name) | |
| note = "one image" if len(layers) == 1 else f"{len(layers)} pages of the TIFF as layers" | |
| else: | |
| raise UploadError("send an image (PNG, JPEG, TIFF), a multi-page TIFF or ZIP of layers, or an .npz") | |
| shapes = {l.shape for l in layers} | |
| if len(shapes) != 1: | |
| raise UploadError(f"the layers differ in size: {sorted(shapes)}") | |
| h, w = layers[0].shape | |
| if min(h, w) < MIN_SIDE: | |
| raise UploadError(f"the image is {w} x {h} px; it needs at least {MIN_SIDE} px on each side") | |
| if len(layers) > MAX_LAYERS: | |
| raise UploadError(f"{len(layers)} layers; send at most {MAX_LAYERS}") | |
| if max(h, w) > MAX_SIDE: | |
| raise UploadError(f"the image is {w} x {h} px; crop it to at most {MAX_SIDE} px on each side") | |
| if len(layers) == 1: | |
| return layers[0], "image", note | |
| return np.stack(layers), "stack", note | |