diff --git a/client/src/views/ehtool/DetectionWorkflow.js b/client/src/views/ehtool/DetectionWorkflow.js index f9eb8e41..cadbe7dc 100644 --- a/client/src/views/ehtool/DetectionWorkflow.js +++ b/client/src/views/ehtool/DetectionWorkflow.js @@ -39,6 +39,11 @@ import { getProofreadingMaskPath, getTrainingReadyCorrectedMask, } from "./proofreadingPaths"; +import { + isAuthoritativeMaskForPlane, + parsePyramidHeaders, + resolvePyramidLevelForQuality, +} from "./pyramidMetadata"; const { Sider, Content } = Layout; const { Title, Text } = Typography; @@ -48,6 +53,12 @@ const parsePositiveInt = (value, fallback) => { return Number.isFinite(parsed) && parsed > 0 ? parsed : fallback; }; +const pyramidCacheIdentity = (pyramid) => { + if (!pyramid) return "unresolved"; + const revision = pyramid.revision || pyramid.source || "unknown-source"; + return `${encodeURIComponent(revision)}@${Number(pyramid.level) || 0}`; +}; + const PREVIEW_MAX_DIM = parsePositiveInt( process.env.REACT_APP_EH_PREVIEW_MAX_DIM, 384, @@ -319,6 +330,7 @@ function DetectionWorkflow({ maskAllBase64: null, maskActiveBase64: null, maskRawBase64: null, + maskRawAuthoritative: false, zIndex: 0, axis: "xy", total: 0, @@ -337,6 +349,8 @@ function DetectionWorkflow({ const filmstripFrameCache = useRef(new Map()); const filmstripFrameCacheOrder = useRef([]); const filmstripFrameCacheLimit = 240; + const pyramidRevisionRef = useRef(null); + const pyramidLevelByQualityRef = useRef(new Map()); const isScrubbingRef = useRef(false); const scrubTrackRef = useRef({ zIndex: 0, at: 0 }); const isFastScrubRef = useRef(false); @@ -860,8 +874,13 @@ function DetectionWorkflow({ includeActive, quality, includeRaw = false, - }) => - [ + }) => { + const revision = pyramidRevisionRef.current || "unresolved"; + const level = resolvePyramidLevelForQuality( + pyramidLevelByQualityRef.current, + quality, + ); + return `${[ instanceId, axis, zIndex, @@ -869,7 +888,8 @@ function DetectionWorkflow({ includeActive ? "active" : "noactive", includeRaw ? "raw" : "noraw", quality, - ].join(":"); + ].join(":")}|pyr=${encodeURIComponent(revision)}@${level}`; + }; const buildFilmstripBatchKey = ({ sessionId: sid, @@ -880,8 +900,13 @@ function DetectionWorkflow({ kind, maxDim, quality, - }) => - [ + }) => { + const revision = pyramidRevisionRef.current || "unresolved"; + const level = + pyramidLevelByQualityRef.current.get(`${quality}:${kind}`) ?? + pyramidLevelByQualityRef.current.get(quality) ?? + "auto"; + return `${[ sid, instanceId, axis, @@ -890,12 +915,49 @@ function DetectionWorkflow({ kind, maxDim || "full", quality, - ].join(":"); + ].join(":")}|pyr=${encodeURIComponent(revision)}@${level}`; + }; + + const resolvePyramidCacheKey = (key, payload) => { + if (!payload?.pyramid) return key; + const baseKey = String(key).split("|pyr=")[0]; + const identity = pyramidCacheIdentity(payload?.pyramid); + return `${baseKey}|pyr=${identity}`; + }; + + const observePyramidMetadata = (pyramid, quality) => { + if (!pyramid) return; + const revision = pyramid.revision || pyramid.source || null; + if ( + revision && + pyramidRevisionRef.current && + revision !== pyramidRevisionRef.current + ) { + // A changed source revision invalidates all rendered frame identities; + // authoritative edit coordinates and payloads remain untouched. + clearFrameCaches(); + pyramidLevelByQualityRef.current.clear(); + } + if (revision) pyramidRevisionRef.current = revision; + if (quality) pyramidLevelByQualityRef.current.set(quality, pyramid.level); + }; const cacheFilmstripBatch = (key, batch) => { - const existing = filmstripBatchCache.current.get(key); + const baseKey = String(key).split("|pyr=")[0]; + const pyramid = batch?.meta?.pyramid; + const revision = + pyramidRevisionRef.current || + pyramid?.revision || + pyramid?.source || + "unresolved"; + const resolvedKey = pyramid + ? `${baseKey}|pyr=${encodeURIComponent(revision)}@${ + Number(pyramid.level) || 0 + }` + : key; + const existing = filmstripBatchCache.current.get(resolvedKey); if (existing && existing !== batch) { - const prefix = `${key}:`; + const prefix = `${resolvedKey}:`; filmstripFrameCache.current.forEach((_, frameKey) => { if (!frameKey.startsWith(prefix)) return; filmstripFrameCache.current.delete(frameKey); @@ -905,8 +967,8 @@ function DetectionWorkflow({ (frameKey) => !frameKey.startsWith(prefix), ); } - filmstripBatchCache.current.set(key, batch); - touchCacheOrder(filmstripBatchCacheOrder.current, key); + filmstripBatchCache.current.set(resolvedKey, batch); + touchCacheOrder(filmstripBatchCacheOrder.current, resolvedKey); while (filmstripBatchCacheOrder.current.length > filmstripBatchCacheLimit) { const oldest = filmstripBatchCacheOrder.current.shift(); filmstripBatchCache.current.delete(oldest); @@ -920,6 +982,7 @@ function DetectionWorkflow({ (frameKey) => !frameKey.startsWith(prefix), ); } + return resolvedKey; }; const cacheFilmstripFrame = (key, frameUrl) => { @@ -932,12 +995,13 @@ function DetectionWorkflow({ }; const cachePreview = (key, payload) => { - const existing = previewCache.current.get(key); + const resolvedKey = resolvePyramidCacheKey(key, payload); + const existing = previewCache.current.get(resolvedKey); if (existing && existing !== payload) { revokePayloadUrls(existing); } - previewCache.current.set(key, payload); - touchCacheOrder(previewCacheOrder.current, key); + previewCache.current.set(resolvedKey, payload); + touchCacheOrder(previewCacheOrder.current, resolvedKey); while (previewCacheOrder.current.length > previewCacheLimit) { const oldest = previewCacheOrder.current.shift(); const cached = previewCache.current.get(oldest); @@ -947,12 +1011,13 @@ function DetectionWorkflow({ }; const cacheView = (key, payload) => { - const existing = viewCache.current.get(key); + const resolvedKey = resolvePyramidCacheKey(key, payload); + const existing = viewCache.current.get(resolvedKey); if (existing && existing !== payload) { revokePayloadUrls(existing); } - viewCache.current.set(key, payload); - touchCacheOrder(viewCacheOrder.current, key); + viewCache.current.set(resolvedKey, payload); + touchCacheOrder(viewCacheOrder.current, resolvedKey); while (viewCacheOrder.current.length > viewCacheLimit) { const oldest = viewCacheOrder.current.shift(); const cached = viewCache.current.get(oldest); @@ -971,16 +1036,23 @@ function DetectionWorkflow({ maskAllBase64: payload.maskAllBase64 ?? prev.maskAllBase64, maskActiveBase64: payload.maskActiveBase64 ?? prev.maskActiveBase64, maskRawBase64: - payload.maskRawBase64 ?? - (payload.zIndex === prev.zIndex && + payload.zIndex === prev.zIndex && payload.axis === prev.axis && - payload.instanceId === prev.instanceId + payload.instanceId === prev.instanceId && + prev.maskRawAuthoritative ? prev.maskRawBase64 - : null), + : null, + maskRawAuthoritative: + payload.zIndex === prev.zIndex && + payload.axis === prev.axis && + payload.instanceId === prev.instanceId && + prev.maskRawAuthoritative, zIndex: payload.zIndex, axis: payload.axis, total: payload.total, instanceId: payload.instanceId ?? prev.instanceId, + pyramid: payload.pyramid ?? prev.pyramid ?? null, + pyramidByKind: payload.pyramidByKind ?? prev.pyramidByKind ?? {}, })); setAxisTotal(payload.total); }; @@ -1153,6 +1225,15 @@ function DetectionWorkflow({ metaResponse?.headers?.["x-total-layers"] ?? totalLayers ?? 0, ); const resolvedAxis = metaResponse?.headers?.["x-axis"] ?? axis; + const pyramidByKind = {}; + responses.forEach((response, idx) => { + const pyramid = parsePyramidHeaders(response?.headers); + if (pyramid) pyramidByKind[kinds[idx]] = pyramid; + }); + const pyramid = pyramidByKind.image || pyramidByKind[kinds[0]] || null; + if (pyramidByKind.image) { + observePyramidMetadata(pyramidByKind.image, quality); + } const payload = { imageBase64: null, @@ -1167,6 +1248,10 @@ function DetectionWorkflow({ batchCount: 1, quality, kindSet: kinds, + pyramid, + pyramidByKind, + maskRawAuthoritative: + quality === "full" && kinds.includes("mask_active_binary"), }; responses.forEach((response, idx) => { @@ -1240,6 +1325,14 @@ function DetectionWorkflow({ responses.forEach((response, idx) => { const { kind, key } = missingKinds[idx]; + const pyramid = parsePyramidHeaders(response?.headers); + if (pyramid) { + pyramidLevelByQualityRef.current.set( + `${quality}:${kind}`, + pyramid.level, + ); + if (kind === "image") observePyramidMetadata(pyramid, quality); + } const entry = { blob: response.data, meta: { @@ -1253,10 +1346,11 @@ function DetectionWorkflow({ axis: response?.headers?.["x-axis"] ?? axis, frameHeight: Number(response?.headers?.["x-frame-height"] ?? 0) || null, + pyramid, }, }; - cacheFilmstripBatch(key, entry); - entriesByKind.set(kind, { key, ...entry }); + const resolvedKey = cacheFilmstripBatch(key, entry); + entriesByKind.set(kind, { key: resolvedKey, ...entry }); }); } @@ -1284,6 +1378,14 @@ function DetectionWorkflow({ batchCount: resolvedCount, quality, kindSet: kinds, + pyramid: meta?.pyramid || null, + pyramidByKind: Object.fromEntries( + Array.from(entriesByKind.entries()).map(([kind, entry]) => [ + kind, + entry?.meta?.pyramid || null, + ]), + ), + maskRawAuthoritative: false, }; for (let idx = 0; idx < kinds.length; idx += 1) { @@ -1318,7 +1420,7 @@ function DetectionWorkflow({ preferFilmstrip = ENABLE_FILMSTRIP_PREVIEW, signal, }) => { - const kinds = ["image", "mask_active_binary"]; + const kinds = ["image"]; if (!imageOnly) { if (includeActive) kinds.push("mask_active"); if (includeAll) kinds.push("mask_all"); @@ -1457,6 +1559,7 @@ function DetectionWorkflow({ maskAllBase64: null, maskActiveBase64: null, maskRawBase64: null, + maskRawAuthoritative: false, })); setSliderZ(resolvedIndex); setCommittedZ(resolvedIndex); @@ -1793,7 +1896,11 @@ function DetectionWorkflow({ maxDim: null, quality: "full", }); - const merged = { ...cached, maskRawBase64: rawPayload.maskRawBase64 }; + const merged = { + ...cached, + maskRawBase64: rawPayload.maskRawBase64, + maskRawAuthoritative: true, + }; cacheView(cacheKey, merged); lastFullRequestKeyRef.current = requestIdentity; setViewState(merged); @@ -1936,6 +2043,7 @@ function DetectionWorkflow({ maskAllBase64: null, maskActiveBase64: null, maskRawBase64: null, + maskRawAuthoritative: false, })); setSliderZ(axisIndex); setCommittedZ(axisIndex); @@ -2158,12 +2266,13 @@ function DetectionWorkflow({ ) => { if (!sessionId || !activeInstanceId) return; const targetIndex = clampSliceIndex(zIndex, axisTotal || totalLayers); - const rawMaskMatchesCurrentSlice = - planeState?.axis === axis && - planeState?.zIndex === targetIndex && - planeState?.instanceId === activeInstanceId && - Boolean(planeState?.maskRawBase64); - if (!rawMaskMatchesCurrentSlice) { + const authoritativeMaskMatchesCurrentSlice = isAuthoritativeMaskForPlane({ + planeState, + axis, + targetIndex, + instanceId: activeInstanceId, + }); + if (!authoritativeMaskMatchesCurrentSlice) { logProofreadingEvent( "proofreading_mask_save_blocked_stale_mask", { @@ -2172,6 +2281,7 @@ function DetectionWorkflow({ viewStateAxis: planeState?.axis, viewStateInstanceId: planeState?.instanceId, hasRawMask: Boolean(planeState?.maskRawBase64), + maskRawAuthoritative: Boolean(planeState?.maskRawAuthoritative), }, { level: "WARNING" }, ); diff --git a/client/src/views/ehtool/DetectionWorkflow.pyramid.test.js b/client/src/views/ehtool/DetectionWorkflow.pyramid.test.js new file mode 100644 index 00000000..8cb5f22d --- /dev/null +++ b/client/src/views/ehtool/DetectionWorkflow.pyramid.test.js @@ -0,0 +1,71 @@ +import { + isAuthoritativeMaskForPlane, + parsePyramidHeaders, + resolvePyramidLevelForQuality, +} from "./pyramidMetadata"; + +describe("DetectionWorkflow pyramid metadata", () => { + test("parses encoded coordinate metadata", () => { + expect( + parsePyramidHeaders({ + "x-pyramid-level": "2", + "x-pyramid-authoritative-level": "0", + "x-pyramid-scale": "[4,2,2]", + "x-pyramid-translation": "[0,0,0]", + "x-pyramid-base-shape": "[64,512,512]", + "x-pyramid-dataset-key": "pyramid%2F2", + "x-pyramid-source": "%2Fdata%2Fimage.zarr", + "x-pyramid-revision": "revision%3A7", + }), + ).toEqual({ + level: 2, + authoritativeLevel: 0, + scale: [4, 2, 2], + translation: [0, 0, 0], + baseShape: [64, 512, 512], + datasetKey: "pyramid/2", + source: "/data/image.zarr", + revision: "revision:7", + }); + }); + + test("rejects malformed level metadata and ignores invalid vectors", () => { + expect(parsePyramidHeaders({})).toBeNull(); + expect(parsePyramidHeaders({ "x-pyramid-level": "coarse" })).toBeNull(); + expect( + parsePyramidHeaders({ + "x-pyramid-level": "1", + "x-pyramid-scale": '[2,"bad",2]', + }).scale, + ).toBeNull(); + }); + + test("preview labels share the observed server preview level", () => { + const levels = new Map([["preview", 2]]); + expect(resolvePyramidLevelForQuality(levels, "thumb-192")).toBe(2); + expect(resolvePyramidLevelForQuality(levels, "preview-384")).toBe(2); + expect(resolvePyramidLevelForQuality(levels, "full")).toBe("auto"); + }); + + test("only full-resolution masks are eligible for save", () => { + const base = { + axis: "xy", + zIndex: 6, + instanceId: 12, + maskRawBase64: "blob:mask", + }; + const matches = (planeState) => + isAuthoritativeMaskForPlane({ + planeState, + axis: "xy", + targetIndex: 6, + instanceId: 12, + }); + + expect(matches({ ...base, maskRawAuthoritative: false })).toBe(false); + expect(matches({ ...base, maskRawAuthoritative: true })).toBe(true); + expect(matches({ ...base, zIndex: 5, maskRawAuthoritative: true })).toBe( + false, + ); + }); +}); diff --git a/client/src/views/ehtool/pyramidMetadata.js b/client/src/views/ehtool/pyramidMetadata.js new file mode 100644 index 00000000..64c985b1 --- /dev/null +++ b/client/src/views/ehtool/pyramidMetadata.js @@ -0,0 +1,66 @@ +const parsePyramidVectorHeader = (value) => { + if (!value) return null; + try { + const parsed = JSON.parse(value); + if (!Array.isArray(parsed)) return null; + const vector = parsed.map(Number); + return vector.every(Number.isFinite) ? vector : null; + } catch (_error) { + return null; + } +}; + +const decodePyramidHeader = (value) => { + if (!value) return null; + try { + return decodeURIComponent(value); + } catch (_error) { + return value; + } +}; + +export const parsePyramidHeaders = (headers = {}) => { + const rawLevel = headers["x-pyramid-level"]; + if (rawLevel === undefined || rawLevel === null || rawLevel === "") { + return null; + } + const level = Number(rawLevel); + if (!Number.isFinite(level)) return null; + const authoritativeLevel = Number( + headers["x-pyramid-authoritative-level"] ?? 0, + ); + const datasetKey = decodePyramidHeader(headers["x-pyramid-dataset-key"]); + const source = decodePyramidHeader(headers["x-pyramid-source"]); + const explicitRevision = decodePyramidHeader(headers["x-pyramid-revision"]); + const revision = + explicitRevision || [source, datasetKey].filter(Boolean).join("#") || null; + return { + level, + authoritativeLevel: Number.isFinite(authoritativeLevel) + ? authoritativeLevel + : 0, + scale: parsePyramidVectorHeader(headers["x-pyramid-scale"]), + translation: parsePyramidVectorHeader(headers["x-pyramid-translation"]), + baseShape: parsePyramidVectorHeader(headers["x-pyramid-base-shape"]), + datasetKey, + source, + revision, + }; +}; + +export const resolvePyramidLevelForQuality = (levels, quality) => + levels.get(quality) ?? + (quality !== "full" ? levels.get("preview") : undefined) ?? + "auto"; + +export const isAuthoritativeMaskForPlane = ({ + planeState, + axis, + targetIndex, + instanceId, +}) => + planeState?.axis === axis && + planeState?.zIndex === targetIndex && + planeState?.instanceId === instanceId && + Boolean(planeState?.maskRawBase64) && + planeState?.maskRawAuthoritative === true; diff --git a/docs/decisions/dbos-durable-operations-spike.md b/docs/decisions/dbos-durable-operations-spike.md index a7227fab..2caca5ef 100644 --- a/docs/decisions/dbos-durable-operations-spike.md +++ b/docs/decisions/dbos-durable-operations-spike.md @@ -50,18 +50,18 @@ system database. ## Gates -| Gate | Required evidence | Result | -| --- | --- | --- | -| Runtime compatibility | Installs on the repository's Python 3.10-3.11 range | **Pass.** DBOS 2.28.0 declares Python >=3.10; the spike ran on 3.11. | -| Idempotent submission | Submitting the same workflow ID twice executes one workflow and one set of external markers | **Pass.** Both handles have the same ID; every marker is written once. | -| Durable progress | Progress is queryable outside the worker and remains available after completion or process death | **Pass.** `DBOS.set_event` progress is read through `DBOSClient`. | -| Queued cancellation | Cancelling enqueued work removes it before any external effect | **Pass.** Status becomes `CANCELLED`; no marker directory is created. | -| Running cancellation | Cancellation stops work at a documented durable boundary | **Pass with constraint.** Cancellation preempts at the next step boundary; it does not interrupt an ordinary blocking synchronous step. | -| Single-server restart | A killed process resumes from its last completed step without repeating that step's external effect | **Pass.** A replacement with the same executor identity recovers the `PENDING` workflow and completes the remaining markers. | -| Mid-step crash safety | Killing a process during a non-transactional external side effect cannot duplicate or corrupt that effect | **Not proven.** The test kills after the step and progress event are durable. Production steps still require idempotent outputs or transactional integration. | -| Postgres and multiple executors | Recovery, queue concurrency, and cancellation work with the intended production topology | **Not run; production gate fails.** SQLite is explicitly a development/test backend. | -| PyTC subprocess control | Training/inference subprocesses are killed, reaped, and reconciled correctly on cancel/restart | **Not run; production gate fails.** A blocking `Popen` step is not sufficient. | -| Product-state projection | DBOS state and `WorkflowOperation` cannot diverge under crashes | **Not designed; production gate fails.** A single source of truth and projection strategy is required. | +| Gate | Required evidence | Result | +| ------------------------------- | --------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| Runtime compatibility | Installs on the repository's Python 3.10-3.11 range | **Pass.** DBOS 2.28.0 declares Python >=3.10; the spike ran on 3.11. | +| Idempotent submission | Submitting the same workflow ID twice executes one workflow and one set of external markers | **Pass.** Both handles have the same ID; every marker is written once. | +| Durable progress | Progress is queryable outside the worker and remains available after completion or process death | **Pass.** `DBOS.set_event` progress is read through `DBOSClient`. | +| Queued cancellation | Cancelling enqueued work removes it before any external effect | **Pass.** Status becomes `CANCELLED`; no marker directory is created. | +| Running cancellation | Cancellation stops work at a documented durable boundary | **Pass with constraint.** Cancellation preempts at the next step boundary; it does not interrupt an ordinary blocking synchronous step. | +| Single-server restart | A killed process resumes from its last completed step without repeating that step's external effect | **Pass.** A replacement with the same executor identity recovers the `PENDING` workflow and completes the remaining markers. | +| Mid-step crash safety | Killing a process during a non-transactional external side effect cannot duplicate or corrupt that effect | **Not proven.** The test kills after the step and progress event are durable. Production steps still require idempotent outputs or transactional integration. | +| Postgres and multiple executors | Recovery, queue concurrency, and cancellation work with the intended production topology | **Not run; production gate fails.** SQLite is explicitly a development/test backend. | +| PyTC subprocess control | Training/inference subprocesses are killed, reaped, and reconciled correctly on cancel/restart | **Not run; production gate fails.** A blocking `Popen` step is not sufficient. | +| Product-state projection | DBOS state and `WorkflowOperation` cannot diverge under crashes | **Not designed; production gate fails.** A single source of truth and projection strategy is required. | ## Findings diff --git a/docs/synthetic-core-project.md b/docs/synthetic-core-project.md index a137f308..ae1eb78f 100644 --- a/docs/synthetic-core-project.md +++ b/docs/synthetic-core-project.md @@ -6,12 +6,12 @@ The local app now starts with a deterministic synthetic segmentation project by The generated project lives at `.pytc/synthetic-core-project` and contains four compressed, chunked HDF5 image volumes: -| Volume | Initial state | Intended workflow role | -| --- | --- | --- | -| `train-01` | Ground truth | Training source | -| `train-02` | Ground truth | Training source | +| Volume | Initial state | Intended workflow role | +| ----------- | --------------- | --------------------------- | +| `train-01` | Ground truth | Training source | +| `train-02` | Ground truth | Training source | | `review-01` | Imperfect draft | Proofreading and correction | -| `target-01` | Image only | Inference target | +| `target-01` | Image only | Inference target | The expected progress state is always **4 total / 2 ground truth / 1 needs proofreading / 1 missing segmentation**. Baseline and corrected candidate predictions are prepopulated for comparison. The data are only for interaction and systems testing; they are not scientific evidence or a model-quality benchmark. diff --git a/server_api/auth/database.py b/server_api/auth/database.py index 80137595..2720f091 100644 --- a/server_api/auth/database.py +++ b/server_api/auth/database.py @@ -14,9 +14,9 @@ engine = create_engine( DATABASE_URL, - connect_args={"check_same_thread": False} - if DATABASE_URL.startswith("sqlite:") - else {}, + connect_args=( + {"check_same_thread": False} if DATABASE_URL.startswith("sqlite:") else {} + ), ) SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) Base = declarative_base() diff --git a/server_api/ehtool/data_manager.py b/server_api/ehtool/data_manager.py index 3e4d42fe..1bae8e6d 100644 --- a/server_api/ehtool/data_manager.py +++ b/server_api/ehtool/data_manager.py @@ -9,6 +9,7 @@ import shutil import tempfile import logging +import threading from datetime import datetime, timezone import glob import numpy as np @@ -19,7 +20,12 @@ from scipy import ndimage from collections import OrderedDict from app_event_logger import append_app_event -from server_api.workflows.volume_io import load_volume, split_dataset_ref +from server_api.workflows.volume_io import ( + VolumeStore, + load_volume, + open_volume_store, + split_dataset_ref, +) from .utils import ( to_uint8, @@ -45,7 +51,14 @@ class DataManager: """ def __init__(self): - self.image_volume: Optional[np.ndarray] = None + # Single-file image artifacts stay storage-backed. ``image_volume`` + # intentionally remains the compatibility-facing array-like object, + # while masks are still materialized at authoritative level 0. + self.image_volume: Optional[Any] = None + self._image_store: Optional[VolumeStore] = None + self._image_level_stores: Dict[int, VolumeStore] = {} + self._image_source_revision: Optional[str] = None + self._image_store_lock = threading.RLock() self.mask_volume: Optional[np.ndarray] = None self.mask_path: Optional[str] = None self.dataset_path: Optional[str] = None @@ -102,16 +115,26 @@ def load_dataset( self, dataset_path: str, mask_path: Optional[str] = None ) -> Dict[str, Any]: """Load image dataset and optional mask dataset""" + self._close_image_stores() # Discover and load images - image_data = self._load_volume(dataset_path) + image_data = self._load_image_volume(dataset_path) # Load masks if provided mask_data = None if mask_path: - mask_data = self._load_volume(mask_path) + try: + mask_data = self._load_volume(mask_path) + except Exception: + image_store = image_data.get("store") + if image_store is not None: + image_store.close() + raise # Validate mask dimensions match image if image_data["num_slices"] != mask_data["num_slices"]: + image_store = image_data.get("store") + if image_store is not None: + image_store.close() raise ValueError( f"Mask layer count ({mask_data['num_slices']}) does not match " f"image layer count ({image_data['num_slices']})" @@ -121,6 +144,9 @@ def load_dataset( img_shape = image_data["shape"] mask_shape = mask_data["shape"] if img_shape[-2:] != mask_shape[-2:]: + image_store = image_data.get("store") + if image_store is not None: + image_store.close() raise ValueError( f"Mask dimensions {mask_shape[-2:]} do not match " f"image dimensions {img_shape[-2:]}" @@ -128,9 +154,14 @@ def load_dataset( # Store volume data self.image_volume = image_data["volume"] + self._image_store = image_data.get("store") + self._image_level_stores = ( + {0: self._image_store} if self._image_store is not None else {} + ) self.mask_volume = mask_data["volume"] if mask_data else None self.mask_path = mask_path self.dataset_path = dataset_path + self._image_source_revision = self._build_image_source_revision(dataset_path) self.is_3d = image_data["is_3d"] self.total_layers = image_data["num_slices"] self.image_shape = image_data["shape"] @@ -178,6 +209,38 @@ def load_dataset( "has_masks": mask_data is not None, } + def close(self) -> None: + """Release storage-backed image resources owned by this manager.""" + self._close_image_stores() + + def __del__(self) -> None: + try: + self._close_image_stores() + except Exception: + pass + + def _close_image_stores(self) -> None: + with self._image_store_lock: + stores = list(self._image_level_stores.values()) + if self._image_store is not None: + stores.append(self._image_store) + seen: set[int] = set() + for store in stores: + identity = id(store) + if identity in seen: + continue + seen.add(identity) + try: + store.close() + except Exception: + logger.debug( + "Failed to close proofreading image store", exc_info=True + ) + self._image_level_stores = {} + self._image_store = None + self._image_source_revision = None + self.image_volume = None + def save_mask(self, layer_index: int, mask_base64: str) -> None: """Update mask for a specific layer and save to disk""" import base64 @@ -615,6 +678,197 @@ def _write_volume_to_files(self, files: List[str], volume: np.ndarray) -> None: else: self._atomic_write_image(target, slice_data) + def _image_levels(self) -> Tuple[Any, ...]: + if self._image_store is None: + return () + return tuple(getattr(self._image_store.metadata, "levels", ()) or ()) + + def _image_store_for_level(self, level: int) -> Optional[VolumeStore]: + with self._image_store_lock: + if self._image_store is None: + return None + cached = self._image_level_stores.get(level) + if cached is not None: + return cached + store = open_volume_store(str(self.dataset_path), level=level) + self._image_level_stores[level] = store + return store + + def _pyramid_level_metadata(self, level: int) -> Optional[Any]: + levels = self._image_levels() + if levels and 0 <= level < len(levels): + return levels[level] + return None + + def _select_image_pyramid_level( + self, axis: str, *, quality: str, max_dim: Optional[int] + ) -> int: + """Choose a native image mip before reading any pixels.""" + if quality != "preview" or not max_dim or max_dim <= 0: + return 0 + levels = self._image_levels() + if len(levels) <= 1: + return 0 + + axis = axis.lower() + plane_axes = { + "xy": (-2, -1), + "zx": (0, -1), + "zy": (0, 1), + }.get(axis) + if plane_axes is None: + return 0 + + selected = 0 + base = levels[0] + for candidate in levels: + shape = tuple(int(value) for value in candidate.shape) + if len(shape) == 2: + largest = max(shape) + elif len(shape) == 3: + largest = max(shape[plane_axes[0]], shape[plane_axes[1]]) + else: + continue + candidate_scale = tuple(float(value) for value in candidate.scale) + if candidate_scale and any(value <= 0 for value in candidate_scale): + continue + # The base-resolution label overlay is rendered independently. Until + # the client applies full affine transforms, only use native mips + # whose in-plane origin is identical to the authoritative level. + base_translation = tuple(float(value) for value in base.translation) + candidate_translation = tuple( + float(value) for value in candidate.translation + ) + if base_translation and candidate_translation: + if any( + abs(candidate_translation[index] - base_translation[index]) > 1e-9 + for index in plane_axes + ): + continue + base_scale = tuple(float(value) for value in base.scale) + if base_scale and candidate_scale: + compatible_extent = all( + np.isclose( + float(base.shape[index]) * base_scale[index], + float(candidate.shape[index]) * candidate_scale[index], + rtol=0.02, + atol=max(base_scale[index], candidate_scale[index]), + ) + for index in plane_axes + ) + if not compatible_extent: + continue + # Do not choose a level that has already fallen below the requested + # output size; that would enlarge a low-resolution native mip. + if largest >= int(max_dim): + selected = int(candidate.index) + + return selected + + def _map_base_coordinate(self, base_index: int, axis_index: int, level: int) -> int: + levels = self._image_levels() + if not levels or level == 0: + shape = tuple(int(value) for value in self.image_shape or ()) + return max(0, min(int(base_index), shape[axis_index] - 1)) + + base = levels[0] + target = levels[level] + ndim = len(target.shape) + base_scale = tuple(base.scale) or tuple(1.0 for _ in range(ndim)) + base_translation = tuple(base.translation) or tuple(0.0 for _ in range(ndim)) + target_scale = tuple(target.scale) or tuple(1.0 for _ in range(ndim)) + target_translation = tuple(target.translation) or tuple( + 0.0 for _ in range(ndim) + ) + world_coordinate = float(base_index) * float(base_scale[axis_index]) + float( + base_translation[axis_index] + ) + target_axis_scale = float(target_scale[axis_index]) + if not np.isfinite(target_axis_scale) or target_axis_scale <= 0: + raise ValueError("Pyramid coordinate scale must be finite and positive") + target_coordinate = ( + world_coordinate - float(target_translation[axis_index]) + ) / target_axis_scale + mapped = int(np.floor(target_coordinate + 0.5)) + return max(0, min(mapped, int(target.shape[axis_index]) - 1)) + + def _read_image_slice_axis( + self, axis: str, index: int, *, level: int = 0, enhance: bool = True + ) -> np.ndarray: + """Read one bounded image plane, mapping a base coordinate to ``level``.""" + axis = axis.lower() + store = self._image_store_for_level(level) + source = store if store is not None else self.image_volume + if source is None: + raise ValueError("Image volume is not available") + shape = tuple(int(value) for value in source.shape) + + if len(shape) == 2: + image = ( + store.read(label="proofreading image") if store else np.asarray(source) + ) + elif len(shape) == 3: + coordinate_axis = {"xy": 0, "zx": 1, "zy": 2}.get(axis) + if coordinate_axis is None: + raise ValueError(f"Unsupported axis: {axis}") + mapped_index = self._map_base_coordinate(index, coordinate_axis, level) + crop = [slice(None), slice(None), slice(None)] + crop[coordinate_axis] = slice(mapped_index, mapped_index + 1) + if store is not None: + image = store.read(tuple(crop), label="proofreading image") + else: + image = np.asarray(source[tuple(crop)]) + image = np.squeeze(image, axis=coordinate_axis) + else: + raise ValueError(f"Unsupported image volume dimensions: {len(shape)}") + + image = ensure_grayscale_2d(np.asarray(image)) + return enhance_contrast(image) if enhance else to_uint8(image) + + def _pyramid_perf_meta(self, level: int) -> Dict[str, Any]: + level_meta = self._pyramid_level_metadata(level) + base_meta = self._pyramid_level_metadata(0) + metadata = self._image_store.metadata if self._image_store is not None else None + shape = tuple(int(value) for value in (self.image_shape or ())) + return { + "pyramid_level": int(level), + "pyramid_scale": ( + list(level_meta.scale) if level_meta else [1.0] * len(shape) + ), + "pyramid_translation": ( + list(level_meta.translation) if level_meta else [0.0] * len(shape) + ), + "pyramid_base_shape": list(base_meta.shape) if base_meta else list(shape), + "pyramid_dataset_key": ( + level_meta.dataset_key + if level_meta is not None + else getattr(metadata, "dataset_key", None) + ), + "pyramid_source": ( + str(metadata.path) + if metadata is not None + else str(self.dataset_path or "") + ), + "pyramid_revision": self._image_source_revision, + "pyramid_authoritative_level": 0, + } + + def _resolve_base_axis_index( + self, axis: str, index: Optional[int] + ) -> Tuple[str, int, int]: + if self.instance_volume is None: + raise ValueError("Instance volume is not available") + axis = axis.lower() + if self.instance_volume.ndim == 2: + return "xy", 0, 1 + axis_dimension = {"xy": 0, "zx": 1, "zy": 2}.get(axis) + if axis_dimension is None: + raise ValueError(f"Unsupported axis: {axis}") + total = int(self.instance_volume.shape[axis_dimension]) + default = 0 if axis == "xy" else total // 2 + resolved = default if index is None else int(index) + return axis, max(0, min(resolved, total - 1)), total + def get_layer( self, layer_index: int, enhance: bool = True ) -> Tuple[np.ndarray, Optional[np.ndarray]]: @@ -625,17 +879,7 @@ def get_layer( ) # Get image slice - if self.image_volume.ndim == 3: - image = self.image_volume[layer_index] - else: - image = self.image_volume - - image = ensure_grayscale_2d(image) - - if enhance: - image = enhance_contrast(image) - else: - image = to_uint8(image) + image = self._read_image_slice_axis("xy", layer_index, level=0, enhance=enhance) # Get mask slice if exists mask = None @@ -1153,8 +1397,7 @@ def get_instance_slice( if self.instance_volume.ndim == 2: z_index = 0 - image = ensure_grayscale_2d(self.image_volume) - image = enhance_contrast(image) + image = self._read_image_slice_axis("xy", 0, level=0, enhance=True) label_slice = self.instance_volume else: if z_index is None: @@ -1171,13 +1414,24 @@ def get_instance_slice_axis( self, instance_id: int, axis: str, index: Optional[int] = None ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, int, int]: """Return image/mask slice for a given axis (xy, zx, zy).""" + label_slice, active_mask, axis, index, total = ( + self._get_instance_label_slice_axis(instance_id, axis, index) + ) + image_slice = self._read_image_slice_axis(axis, index, level=0, enhance=True) + return image_slice, label_slice, active_mask, index, total + + def _get_instance_label_slice_axis( + self, instance_id: int, axis: str, index: Optional[int] = None + ) -> Tuple[np.ndarray, np.ndarray, str, int, int]: + """Return label planes without touching the image pyramid.""" if self.instance_volume is None: raise ValueError("Instance volume is not available") axis = axis.lower() if self.instance_volume.ndim == 2: - axis = "xy" - index = 0 + label_slice = ensure_grayscale_2d(self.instance_volume) + active_mask = (label_slice == instance_id).astype(np.uint8) + return label_slice, active_mask, "xy", 0, 1 if axis not in {"xy", "zx", "zy"}: raise ValueError(f"Unsupported axis: {axis}") @@ -1185,35 +1439,27 @@ def get_instance_slice_axis( if axis == "xy": z_index = 0 if index is None else int(index) z_index = max(0, min(z_index, self.total_layers - 1)) - image, _ = self.get_layer(z_index, enhance=True) label_slice = self.instance_volume[z_index] active_mask = (label_slice == instance_id).astype(np.uint8) - return image, label_slice, active_mask, z_index, self.total_layers + return label_slice, active_mask, axis, z_index, self.total_layers - # For ZX and ZY, we swap axes for a view slice - volume = self.image_volume label_volume = self.instance_volume if axis == "zx": - # Slice across Y dimension - max_index = volume.shape[1] - 1 + max_index = label_volume.shape[1] - 1 index = max_index // 2 if index is None else int(index) index = max(0, min(index, max_index)) - image_slice = volume[:, index, :] label_slice = label_volume[:, index, :] - total = volume.shape[1] + total = label_volume.shape[1] else: # zy - max_index = volume.shape[2] - 1 + max_index = label_volume.shape[2] - 1 index = max_index // 2 if index is None else int(index) index = max(0, min(index, max_index)) - image_slice = volume[:, :, index] label_slice = label_volume[:, :, index] - total = volume.shape[2] + total = label_volume.shape[2] - image_slice = ensure_grayscale_2d(image_slice) - image_slice = enhance_contrast(image_slice) label_slice = ensure_grayscale_2d(label_slice) active_mask = (label_slice == instance_id).astype(np.uint8) - return image_slice, label_slice, active_mask, index, total + return label_slice, active_mask, axis, index, total def _resize_to_max_dim( self, array: np.ndarray, max_dim: Optional[int], is_mask: bool = False @@ -1343,27 +1589,36 @@ def get_instance_image_bytes( quality = (quality or "full").lower() if quality not in {"full", "preview"}: raise ValueError(f"Unsupported quality: {quality}") - ( - image, - label_slice, - active_mask, - resolved_index, - total, - ) = self.get_instance_slice_axis( - instance_id=instance_id, axis=axis, index=z_index - ) - kind = kind.lower() output_format, media_type = self._normalize_output_format(kind, format) resize_ms = 0.0 max_dim_value = int(max_dim) if max_dim is not None else None if quality == "preview" and (not max_dim_value or max_dim_value <= 0): max_dim_value = 384 + pyramid_level = ( + self._select_image_pyramid_level( + axis, quality=quality, max_dim=max_dim_value + ) + if kind == "image" + else 0 + ) + if kind == "image": + axis, resolved_index, total = self._resolve_base_axis_index(axis, z_index) + image = None + label_slice = active_mask = None + else: + label_slice, active_mask, axis, resolved_index, total = ( + self._get_instance_label_slice_axis( + instance_id=instance_id, axis=axis, index=z_index + ) + ) + image = None if max_dim_value and max_dim_value > 0: cache_key = ( instance_id, axis, resolved_index, + pyramid_level, kind, max_dim_value, quality, @@ -1371,14 +1626,22 @@ def get_instance_image_bytes( ) cached = self._cache_get(self._resized_cache, cache_key) if cached: + perf_meta = self._pyramid_perf_meta(pyramid_level) + perf_meta.update( + {"cache_hit": True, "decode_ms": 0.0, "resize_ms": 0.0} + ) return ( cached, resolved_index, total, axis, media_type, - {"cache_hit": True, "decode_ms": 0.0, "resize_ms": 0.0}, + perf_meta, ) + if kind == "image" and image is None: + image = self._read_image_slice_axis( + axis, resolved_index, level=pyramid_level, enhance=True + ) resize_started = time.perf_counter() if kind == "image": array = image @@ -1396,11 +1659,12 @@ def get_instance_image_bytes( elif kind == "mask_raw": if axis != "xy": raise ValueError("Raw mask only supported for XY view") - _, mask_raw = self.get_layer(resolved_index, enhance=False) - if mask_raw is None: - array = np.zeros_like(image) + if self.mask_volume is None: + array = np.zeros_like(label_slice) + elif self.mask_volume.ndim == 3: + array = ensure_grayscale_2d(self.mask_volume[resolved_index]) else: - array = ensure_grayscale_2d(mask_raw) + array = ensure_grayscale_2d(self.mask_volume) array = self._resize_to_max_dim(array, max_dim_value, is_mask=True) else: raise ValueError(f"Unsupported image kind: {kind}") @@ -1415,6 +1679,7 @@ def get_instance_image_bytes( instance_id, axis, resolved_index, + pyramid_level, kind, max_dim_value, quality, @@ -1423,13 +1688,15 @@ def get_instance_image_bytes( encoded_bytes, self._resized_cache_limit, ) + perf_meta = self._pyramid_perf_meta(pyramid_level) + perf_meta.update({"cache_hit": False, "decode_ms": 0.0, "resize_ms": resize_ms}) return ( encoded_bytes, resolved_index, total, axis, media_type, - {"cache_hit": False, "decode_ms": 0.0, "resize_ms": resize_ms}, + perf_meta, ) def get_instance_filmstrip_bytes( @@ -1463,9 +1730,9 @@ def get_instance_filmstrip_bytes( elif axis == "xy": total = self.total_layers elif axis == "zx": - total = int(self.image_volume.shape[1]) + total = int(self.instance_volume.shape[1]) else: - total = int(self.image_volume.shape[2]) + total = int(self.instance_volume.shape[2]) if z_count is None: z_count = 1 @@ -1478,11 +1745,19 @@ def get_instance_filmstrip_bytes( max_dim_value = int(max_dim) if max_dim is not None else None if quality == "preview" and (not max_dim_value or max_dim_value <= 0): max_dim_value = 384 + pyramid_level = ( + self._select_image_pyramid_level( + axis, quality=quality, max_dim=max_dim_value + ) + if kind == "image" + else 0 + ) cache_key = ( instance_id, axis, z_start, z_count, + pyramid_level, kind, max_dim_value, quality, @@ -1494,6 +1769,8 @@ def get_instance_filmstrip_bytes( cached_bytes, cached_height = cached else: cached_bytes, cached_height = cached, int(max_dim_value or 0) + perf_meta = self._pyramid_perf_meta(pyramid_level) + perf_meta.update({"cache_hit": True, "decode_ms": 0.0, "resize_ms": 0.0}) return ( cached_bytes, z_start, @@ -1502,7 +1779,7 @@ def get_instance_filmstrip_bytes( axis, cached_height, media_type, - {"cache_hit": True, "decode_ms": 0.0, "resize_ms": 0.0}, + perf_meta, ) resize_started = time.perf_counter() @@ -1513,6 +1790,7 @@ def get_instance_filmstrip_bytes( instance_id, axis, z_index, + pyramid_level, kind, max_dim_value or 0, quality, @@ -1523,9 +1801,18 @@ def get_instance_filmstrip_bytes( else None ) if frame is None: - image, label_slice, active_mask, _, _ = self.get_instance_slice_axis( - instance_id=instance_id, axis=axis, index=z_index - ) + if kind == "image": + image = self._read_image_slice_axis( + axis, z_index, level=pyramid_level, enhance=True + ) + label_slice = active_mask = None + else: + label_slice, active_mask, _, _, _ = ( + self._get_instance_label_slice_axis( + instance_id=instance_id, axis=axis, index=z_index + ) + ) + image = None if kind == "image": frame = image @@ -1543,12 +1830,12 @@ def get_instance_filmstrip_bytes( elif kind == "mask_raw": if axis != "xy": raise ValueError("Raw mask only supported for XY view") - _, raw_mask = self.get_layer(z_index, enhance=False) - frame = ( - np.zeros_like(image) - if raw_mask is None - else ensure_grayscale_2d(raw_mask) - ) + if self.mask_volume is None: + frame = np.zeros_like(label_slice) + elif self.mask_volume.ndim == 3: + frame = ensure_grayscale_2d(self.mask_volume[z_index]) + else: + frame = ensure_grayscale_2d(self.mask_volume) frame = self._resize_to_max_dim(frame, max_dim_value, is_mask=True) else: raise ValueError(f"Unsupported image kind: {kind}") @@ -1576,6 +1863,8 @@ def get_instance_filmstrip_bytes( (encoded_bytes, frame_height), self._filmstrip_cache_limit, ) + perf_meta = self._pyramid_perf_meta(pyramid_level) + perf_meta.update({"cache_hit": False, "decode_ms": 0.0, "resize_ms": resize_ms}) return ( encoded_bytes, z_start, @@ -1584,7 +1873,7 @@ def get_instance_filmstrip_bytes( axis, frame_height, media_type, - {"cache_hit": False, "decode_ms": 0.0, "resize_ms": resize_ms}, + perf_meta, ) def get_sparse_active_mask( @@ -1624,6 +1913,70 @@ def get_sparse_active_mask( "axis": axis, } + def _load_image_volume(self, path: str) -> Dict[str, Any]: + """Open random-access image artifacts without materializing the volume.""" + file_path, _dataset_key = split_dataset_ref(path) + path_obj = Path(file_path) + storage_backed = path_obj.is_file() or ( + path_obj.is_dir() and path_obj.name.lower().endswith((".zarr", ".n5")) + ) + if not storage_backed: + return self._load_volume(path) + + store = open_volume_store(path, level=0) + try: + shape = tuple(int(value) for value in store.shape) + axes = tuple(axis.name.lower() for axis in store.metadata.axes) + expected_axes = ("y", "x") if store.ndim == 2 else ("z", "y", "x") + if ( + store.metadata.format in {"zarr", "n5"} + and axes + and axes != expected_axes + ): + raise ValueError( + "Proofreading image axes must be " + f"{''.join(expected_axes).upper()}, got {axes!r}" + ) + if store.ndim == 2: + return { + "volume": store, + "store": store, + "shape": shape, + "num_slices": 1, + "is_3d": False, + } + if store.ndim == 3: + return { + "volume": store, + "store": store, + "shape": shape, + "num_slices": shape[0], + "is_3d": True, + } + raise ValueError(f"Unsupported volume dimensions: {store.ndim}") + except Exception: + store.close() + raise + + @staticmethod + def _build_image_source_revision(path: str) -> str: + """Return a cheap cache revision for an immutable image artifact.""" + file_path, dataset_key = split_dataset_ref(path) + source = Path(file_path).resolve() + candidates = [source] + if source.is_dir(): + candidates.extend( + source / name for name in (".zattrs", ".zgroup", "zarr.json") + ) + stats = [] + for candidate in candidates: + try: + stat = candidate.stat() + except OSError: + continue + stats.append(f"{candidate.name}:{stat.st_mtime_ns}:{stat.st_size}") + return "|".join((str(source), dataset_key or "", *stats)) + def _load_volume(self, path: str) -> Dict[str, Any]: """Load volume data from a path""" file_path, _dataset_key = split_dataset_ref(path) diff --git a/server_api/ehtool/router.py b/server_api/ehtool/router.py index ff9714c0..cd1988ea 100644 --- a/server_api/ehtool/router.py +++ b/server_api/ehtool/router.py @@ -6,9 +6,11 @@ from fastapi import APIRouter, Depends, HTTPException, status, Response from sqlalchemy.orm import Session from typing import List, Optional +import json import math import logging import time +from urllib.parse import quote logger = logging.getLogger(__name__) @@ -92,6 +94,63 @@ def _append_ehtool_event(event: str, level: str = "INFO", **fields): logger.debug("Failed to append EHTool app event", exc_info=True) +def _pyramid_response_metadata(perf_meta): + """Normalize optional pyramid selection metadata from ``DataManager``. + + Proofreading edits continue to use authoritative, full-resolution voxel + coordinates. This metadata only describes the image representation used + to render the response. + """ + perf_meta = perf_meta or {} + vector_fields = { + "scale": "pyramid_scale", + "translation": "pyramid_translation", + "base_shape": "pyramid_base_shape", + } + metadata = { + "level": perf_meta.get("pyramid_level"), + "dataset_key": perf_meta.get("pyramid_dataset_key"), + "source": perf_meta.get("pyramid_source"), + "revision": perf_meta.get("pyramid_revision"), + "authoritative_level": perf_meta.get("pyramid_authoritative_level"), + } + for field, perf_key in vector_fields.items(): + value = perf_meta.get(perf_key) + metadata[field] = list(value) if value is not None else None + return {key: value for key, value in metadata.items() if value is not None} + + +def _pyramid_response_headers(metadata): + """Encode normalized pyramid metadata into compact ASCII-safe headers.""" + if not metadata: + return {} + headers = {} + scalar_headers = { + "level": "X-Pyramid-Level", + "authoritative_level": "X-Pyramid-Authoritative-Level", + } + vector_headers = { + "scale": "X-Pyramid-Scale", + "translation": "X-Pyramid-Translation", + "base_shape": "X-Pyramid-Base-Shape", + } + encoded_headers = { + "dataset_key": "X-Pyramid-Dataset-Key", + "source": "X-Pyramid-Source", + "revision": "X-Pyramid-Revision", + } + for field, header in scalar_headers.items(): + if field in metadata: + headers[header] = str(metadata[field]) + for field, header in vector_headers.items(): + if field in metadata: + headers[header] = json.dumps(metadata[field], separators=(",", ":")) + for field, header in encoded_headers.items(): + if field in metadata: + headers[header] = quote(str(metadata[field]), safe="") + return headers + + def get_data_manager(session_id: int, db: Session) -> DataManager: """Get or create DataManager for a session""" if session_id not in _data_managers: @@ -129,6 +188,8 @@ async def load_detection_dataset( current_user: User = Depends(get_current_user), db: Session = Depends(get_db), ): + data_manager = None + manager_cached = False try: workflow = None if request.workflow_id: @@ -180,6 +241,7 @@ async def load_detection_dataset( # Cache DataManager _data_managers[db_session.id] = data_manager + manager_cached = True if workflow: workflow_patch = { @@ -240,12 +302,18 @@ async def load_detection_dataset( ) except FileNotFoundError as e: + if data_manager is not None and not manager_cached: + data_manager.close() _append_ehtool_event("proofreading_load_failed", level="ERROR", error=str(e)) raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e)) except ValueError as e: + if data_manager is not None and not manager_cached: + data_manager.close() _append_ehtool_event("proofreading_load_failed", level="ERROR", error=str(e)) raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)) except Exception as e: + if data_manager is not None and not manager_cached: + data_manager.close() _append_ehtool_event("proofreading_load_failed", level="ERROR", error=str(e)) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, @@ -563,6 +631,7 @@ async def get_instance_image( except ValueError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc elapsed_ms = (time.perf_counter() - started_at) * 1000.0 + pyramid_meta = _pyramid_response_metadata(perf_meta) _append_ehtool_event( "proofreading_instance_image_served", session_id=session_id, @@ -579,6 +648,8 @@ async def get_instance_image( cache_hit=bool(perf_meta.get("cache_hit")), decode_ms=round(float(perf_meta.get("decode_ms", 0.0)), 2), resize_ms=round(float(perf_meta.get("resize_ms", 0.0)), 2), + performance=perf_meta, + pyramid=pyramid_meta or None, ) headers = { @@ -588,6 +659,7 @@ async def get_instance_image( "X-Cache-Hit": "1" if perf_meta.get("cache_hit") else "0", "X-Decode-MS": f"{float(perf_meta.get('decode_ms', 0.0)):.2f}", "X-Resize-MS": f"{float(perf_meta.get('resize_ms', 0.0)):.2f}", + **_pyramid_response_headers(pyramid_meta), } return Response(content=image_bytes, media_type=media_type, headers=headers) @@ -857,6 +929,7 @@ async def get_instance_filmstrip( except ValueError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc elapsed_ms = (time.perf_counter() - started_at) * 1000.0 + pyramid_meta = _pyramid_response_metadata(perf_meta) _append_ehtool_event( "proofreading_instance_filmstrip_served", session_id=session_id, @@ -875,6 +948,8 @@ async def get_instance_filmstrip( cache_hit=bool(perf_meta.get("cache_hit")), decode_ms=round(float(perf_meta.get("decode_ms", 0.0)), 2), resize_ms=round(float(perf_meta.get("resize_ms", 0.0)), 2), + performance=perf_meta, + pyramid=pyramid_meta or None, ) headers = { @@ -886,6 +961,7 @@ async def get_instance_filmstrip( "X-Cache-Hit": "1" if perf_meta.get("cache_hit") else "0", "X-Decode-MS": f"{float(perf_meta.get('decode_ms', 0.0)):.2f}", "X-Resize-MS": f"{float(perf_meta.get('resize_ms', 0.0)):.2f}", + **_pyramid_response_headers(pyramid_meta), } return Response(content=image_bytes, media_type=media_type, headers=headers) @@ -1057,7 +1133,8 @@ async def delete_detection_session( ) if session_id in _data_managers: - del _data_managers[session_id] + data_manager = _data_managers.pop(session_id) + data_manager.close() db.delete(db_session) db.commit() diff --git a/server_api/main.py b/server_api/main.py index c97f8c41..589a4115 100644 --- a/server_api/main.py +++ b/server_api/main.py @@ -292,7 +292,26 @@ def _ensure_sqlite_column(table_name: str, column_name: str, ddl: str) -> None: allow_credentials=True, allow_methods=["*"], allow_headers=["*"], - expose_headers=["x-request-id"], + expose_headers=[ + "x-request-id", + "x-axis", + "x-cache-hit", + "x-decode-ms", + "x-frame-height", + "x-pyramid-authoritative-level", + "x-pyramid-base-shape", + "x-pyramid-dataset-key", + "x-pyramid-level", + "x-pyramid-revision", + "x-pyramid-scale", + "x-pyramid-source", + "x-pyramid-translation", + "x-resize-ms", + "x-total-layers", + "x-z-count", + "x-z-index", + "x-z-start", + ], ) logger = logging.getLogger(__name__) diff --git a/tests/test_ehtool_data_manager.py b/tests/test_ehtool_data_manager.py index b289e30f..0a563475 100644 --- a/tests/test_ehtool_data_manager.py +++ b/tests/test_ehtool_data_manager.py @@ -2,10 +2,94 @@ import pytest import tifffile +import server_api.ehtool.data_manager as data_manager_module from server_api.ehtool.data_manager import DataManager from server_api.ehtool.utils import array_to_base64, glasbey_color, labels_to_rgba h5py = pytest.importorskip("h5py") +zarr = pytest.importorskip("zarr") + + +def _write_proofreading_pyramid(path, *, coarse_translation=(0.0, 0.0, 0.0)): + fine = np.arange(8 * 12 * 16, dtype=np.uint16).reshape(8, 12, 16) + coarse = fine[::2, ::2, ::2] + root = zarr.open_group(str(path), mode="w") + create_array = getattr(root, "create_array", root.create_dataset) + create_array("0", data=fine, chunks=(2, 4, 4)) + create_array("1", data=coarse, chunks=(1, 3, 4)) + root.attrs["multiscales"] = [ + { + "version": "0.4", + "axes": [ + {"name": "z", "type": "space"}, + {"name": "y", "type": "space"}, + {"name": "x", "type": "space"}, + ], + "datasets": [ + { + "path": "0", + "coordinateTransformations": [ + {"type": "scale", "scale": [1.0, 1.0, 1.0]}, + {"type": "translation", "translation": [0.0, 0.0, 0.0]}, + ], + }, + { + "path": "1", + "coordinateTransformations": [ + {"type": "scale", "scale": [2.0, 2.0, 2.0]}, + { + "type": "translation", + "translation": list(coarse_translation), + }, + ], + }, + ], + } + ] + return fine, coarse + + +class _RecordingStore: + def __init__(self, store, level): + self._store = store + self.level = level + self.reads = [] + self.close_calls = 0 + + @property + def metadata(self): + return self._store.metadata + + @property + def shape(self): + return self._store.shape + + @property + def ndim(self): + return self._store.ndim + + def read(self, crop=None, **kwargs): + self.reads.append(crop) + return self._store.read(crop, **kwargs) + + def close(self): + self.close_calls += 1 + self._store.close() + + +def _record_pyramid_reads(monkeypatch): + real_open = data_manager_module.open_volume_store + opened = [] + + def recording_open(path, *, level=None, **kwargs): + store = _RecordingStore( + real_open(path, level=level, **kwargs), 0 if level is None else level + ) + opened.append(store) + return store + + monkeypatch.setattr(data_manager_module, "open_volume_store", recording_open) + return opened def test_semantic_mask_edits_persist_instance_artifact_and_binary_mask(tmp_path): @@ -117,3 +201,185 @@ def test_hdf5_project_load_accepts_main_and_data_dataset_names(tmp_path): assert manager.image_volume.shape == (3, 8, 8) assert manager.mask_volume.shape == (3, 8, 8) assert manager.mask_volume.dtype == np.uint16 + + +def test_preview_reads_native_mip_before_io_and_overlays_read_no_image( + tmp_path, monkeypatch +): + image_path = tmp_path / "image.zarr" + mask_path = tmp_path / "mask.tif" + _write_proofreading_pyramid(image_path) + mask = np.zeros((8, 12, 16), dtype=np.uint8) + mask[6, 2:6, 3:8] = 255 + tifffile.imwrite(mask_path, mask) + opened = _record_pyramid_reads(monkeypatch) + + manager = DataManager() + manager.load_dataset(str(image_path), str(mask_path)) + manager.ensure_instances() + instance_id = manager.instances[0]["id"] + assert opened[0].reads == [] + + preview = manager.get_instance_image_bytes( + instance_id=instance_id, + z_index=6, + axis="xy", + kind="image", + max_dim=6, + quality="preview", + ) + + level_one = next(store for store in opened if store.level == 1) + assert opened[0].reads == [] + assert level_one.reads == [(slice(3, 4), slice(None), slice(None))] + assert preview[1:4] == (6, 8, "xy") + assert preview[-1]["pyramid_level"] == 1 + assert preview[-1]["pyramid_scale"] == [2.0, 2.0, 2.0] + assert preview[-1]["pyramid_authoritative_level"] == 0 + assert preview[-1]["pyramid_revision"] + + read_count = sum(len(store.reads) for store in opened) + manager.get_instance_image_bytes( + instance_id=instance_id, + z_index=6, + axis="xy", + kind="mask_active", + max_dim=6, + quality="preview", + ) + manager.get_instance_filmstrip_bytes( + instance_id=instance_id, + axis="xy", + z_start=5, + z_count=2, + kind="mask_all", + max_dim=6, + quality="preview", + ) + assert sum(len(store.reads) for store in opened) == read_count + + full = manager.get_instance_image_bytes( + instance_id=instance_id, + z_index=6, + axis="xy", + kind="image", + max_dim=None, + quality="full", + ) + assert opened[0].reads == [(slice(6, 7), slice(None), slice(None))] + assert full[-1]["pyramid_level"] == 0 + + +def test_coarse_preview_then_edit_persists_authoritative_voxels(tmp_path): + image_path = tmp_path / "image.zarr" + mask_path = tmp_path / "mask.tif" + _write_proofreading_pyramid(image_path) + mask = np.zeros((8, 12, 16), dtype=np.uint8) + mask[6, 2:6, 3:8] = 255 + tifffile.imwrite(mask_path, mask) + + manager = DataManager() + manager.load_dataset(str(image_path), str(mask_path)) + manager.ensure_instances() + instance_id = manager.instances[0]["id"] + preview = manager.get_instance_image_bytes( + instance_id=instance_id, + z_index=6, + axis="xy", + kind="image", + max_dim=6, + quality="preview", + ) + assert preview[-1]["pyramid_level"] == 1 + + edited = np.zeros((12, 16), dtype=np.uint8) + edited[1:5, 1:5] = 255 + result = manager.save_instance_mask_slice( + instance_id=instance_id, + axis="xy", + index=6, + mask_base64=array_to_base64(edited, format="PNG"), + ) + + assert result["z_index"] == 6 + assert manager.mask_volume[6, 1, 1] == 255 + assert manager.mask_volume[3, 1, 1] == 0 + manager.close() + + reloaded = DataManager() + reloaded.load_dataset(str(image_path), str(mask_path)) + reloaded.ensure_instances() + assert reloaded.mask_volume[6, 1, 1] == 255 + assert reloaded.mask_volume[3, 1, 1] == 0 + + +def test_translated_mip_is_not_used_for_level_zero_overlay(tmp_path): + image_path = tmp_path / "translated.zarr" + mask_path = tmp_path / "mask.tif" + _write_proofreading_pyramid(image_path, coarse_translation=(0.0, 1.0, 0.0)) + mask = np.zeros((8, 12, 16), dtype=np.uint8) + mask[:, 2:4, 2:4] = 7 + tifffile.imwrite(mask_path, mask) + + manager = DataManager() + manager.load_dataset(str(image_path), str(mask_path)) + manager.ensure_instances() + response = manager.get_instance_image_bytes( + instance_id=7, + z_index=2, + axis="xy", + kind="image", + max_dim=6, + quality="preview", + ) + + assert response[-1]["pyramid_level"] == 0 + + +def test_reload_and_close_release_each_pyramid_store_once(tmp_path, monkeypatch): + image_path = tmp_path / "image.zarr" + mask_path = tmp_path / "mask.tif" + _write_proofreading_pyramid(image_path) + mask = np.zeros((8, 12, 16), dtype=np.uint8) + mask[:, 2:4, 2:4] = 1 + tifffile.imwrite(mask_path, mask) + opened = _record_pyramid_reads(monkeypatch) + + manager = DataManager() + manager.load_dataset(str(image_path), str(mask_path)) + manager.ensure_instances() + manager.get_instance_image_bytes( + instance_id=1, + z_index=2, + axis="xy", + kind="image", + max_dim=6, + quality="preview", + ) + original_stores = list(opened) + + manager.load_dataset(str(image_path), str(mask_path)) + assert [store.close_calls for store in original_stores] == [1, 1] + replacement = opened[-1] + manager.close() + manager.close() + assert replacement.close_calls == 1 + assert manager._image_store is None + assert manager._image_level_stores == {} + + +def test_proofreading_rejects_non_zyx_ngff_axes(tmp_path): + image_path = tmp_path / "permuted.zarr" + root = zarr.open_group(str(image_path), mode="w") + create_array = getattr(root, "create_array", root.create_dataset) + create_array("0", data=np.zeros((16, 12, 8), dtype=np.uint8)) + root.attrs["multiscales"] = [ + { + "axes": ["x", "y", "z"], + "datasets": [{"path": "0"}], + } + ] + + manager = DataManager() + with pytest.raises(ValueError, match="must be ZYX"): + manager.load_dataset(str(image_path))