Fit a z-stack over the bridge

README.md

@@ -25,6 +25,39 @@ The server writes its address and a generated token to `~/.flimkit/bridge.json`,
 
 It writes `~/.flimkit/qupath-bridge.json` alongside it, carrying the same address under the older protocol name. QuPath extensions built before the server moved out check for that name and stop pairing without it. The second file goes when those versions are retired.
 
+## Z-stacks
+
+`POST /v1/zstack` runs FLIMKit's z-stack fit over a folder of slices, one file per slice named `region_z1.ptu`, `region_z2.ptu` and so on. A folder holding several regions is fitted as several stacks in one run.
+
+```
+POST /v1/zstack/scan   {"ptu_dir": "/path/slices"}
+POST /v1/zstack        {"ptu_dir": "/path/slices", "params": {"n_exp": 2, "z_step_um": 2.0}}
+GET  /v1/zstack/defaults
+POST /v1/zstack/export {"group_dir": "/path/out/RegionA", "format": "ome-tiff"}
+```
+
+The scan is there so a client can say what it found before offering any settings, and refuse an empty folder without starting a job. The fit itself runs as a job, so `GET /v1/jobs/{id}` reports progress by slice and `DELETE` cancels at the next one.
+
+Each stack is fitted as one FOV: the decay is pooled over every slice, the lifetimes are fitted once from the pooled decay and locked, and each slice then gets a per-pixel fit with only the amplitudes free. A single slice rarely carries the photons to identify two lifetimes; the stack usually does.
+
+Each stack comes back as an OME-Zarr store, a `(C, Z, Y, X)` volume with a channel per map: intensity in photons, `tau_mean_int` and `tau_mean_amp` in ns, an `alpha_N` per component, and the bound fraction and chi-squared maps when they were computed. `zarr` and `ome-zarr` are dependencies of this package, not an extra, so the store is always what a client gets unless it asks for something else. `ome-tiff` is that something else, same channels and same axes, for a viewer that will not read the store; `/v1/zstack/export` rewrites a finished run in either format without fitting it again, which is also how to correct a z step entered wrong.
+
+The store carries one resolution level on purpose. A pyramid is built by interpolating between neighbours, and a lifetime map is mostly NaN where nothing was fitted, so the coarse levels a viewer shows when zoomed out would smear that NaN over the pixels that did fit. Chunking, one chunk per slice per channel, is what keeps a large volume readable.
+
+FLIMKit's own per-slice `.npy` maps, `_zseries.csv` and `_zseries.json` are left where it writes them, and the reply names them. Its per-slice PNGs and detail plots are off by default here, since each one costs a second summed fit per slice; `save_plots` turns them back on.
+
+### When the client cannot see the disk
+
+The reply names a path, so a client on the same machine opens the store itself and nothing large goes over HTTP. Over a forwarded port that path means nothing locally, and for that there is:
+
+```
+GET /v1/zstack/volume.ome.tif?group_dir=/path/out/RegionA&z_step_um=2.0
+```
+
+which builds the volume from the same slice maps and streams it as one OME-TIFF, in blocks out of a temporary file rather than through one bytes object, since a fitted stack is routinely larger than memory. The temporary file goes whether the transfer finished or not. Shape, channels, units and voxel size ride on `X-FLIMKit-Volume-*` headers, and the store on the server is left alone.
+
+The slices themselves are still a server-side path: `ptu_dir` is opened by FLIMKit, not uploaded. A forwarded client picking a local folder gets a 404 naming the folder, before any job starts.
+
 ## Versions
 
 `GET /v1/status` reports `protocol_version`, `bridge_version` and `flimkit_version`. A client checks `protocol_version`, which is what governs whether the two can talk. `bridge_version` is for display: the server and each client version independently now, so a difference between them is ordinary rather than a problem.

flimkit_bridge/dataset_routes.py

@@ -510,6 +510,121 @@ def work(progress, cancel):
             'basename': basename, 'n_tiles': n_tiles, 'params_used': params}
 
 
+def zstack_defaults(state):
+    from flimkit_bridge import zstack
+    return zstack.defaults()
+
+
+def scan_zstack(state, payload):
+    from flimkit_bridge import zstack
+    directory = (payload or {}).get('ptu_dir') or (payload or {}).get('path')
+    if not directory:
+        raise RouteError(400, 'a ptu_dir is required, the folder holding the slices')
+    try:
+        groups = zstack.scan(directory)
+    except FileNotFoundError as exc:
+        raise RouteError(404, str(exc))
+    except Exception as exc:
+        raise RouteError(400, str(exc))
+    from flimkit_bridge import volumes
+    return {'ptu_dir': str(directory), 'stacks': groups,
+            'n_stacks': len(groups),
+            'n_slices': sum(group['n_slices'] for group in groups),
+            'volume_formats': list(volumes.FORMATS)}
+
+
+def run_zstack(state, payload):
+    from flimkit_bridge import zstack
+    jobs = _jobs(state)
+    payload = payload or {}
+    directory = payload.get('ptu_dir') or payload.get('path')
+    if not directory:
+        raise RouteError(400, 'a ptu_dir is required, the folder holding the slices')
+    try:
+        params = zstack.merge_params(payload.get('params'))
+    except ValueError as exc:
+        raise RouteError(400, str(exc))
+    try:
+        ptu_dir, output_dir, groups = zstack.resolve(
+            directory, payload.get('output_dir'))
+    except FileNotFoundError as exc:
+        raise RouteError(404, str(exc))
+    except ValueError as exc:
+        raise RouteError(400, str(exc))
+    args = zstack.build_args(ptu_dir, output_dir, params)
+    n_slices = sum(group['n_slices'] for group in groups)
+
+    def work(progress, cancel):
+        progress(0, n_slices,
+                 f'{len(groups)} z-stack(s) over {n_slices} slices')
+        result = zstack.run(args, progress, cancel)
+        if cancel.is_set():
+            return None
+        return zstack.summarise(result, output_dir, groups, params)
+
+    job_id = jobs.submit('zstack_fit', work)
+    return {'job': job_id, 'ptu_dir': str(ptu_dir),
+            'output_dir': str(output_dir),
+            'stacks': [{k: v for k, v in group.items() if k != 'files'}
+                       for group in groups],
+            'n_stacks': len(groups), 'n_slices': n_slices,
+            'params_used': params}
+
+
+ZSTACK_VOLUME_PATH = '/v1/zstack/volume.ome.tif'
+
+
+def zstack_volume(state, query):
+    from flimkit_bridge import volumes, zstack
+    options = {key: values[-1] for key, values in parse_qs(query).items()}
+    group_dir = options.get('group_dir')
+    if not group_dir:
+        raise RouteError(400, 'group_dir is required, one stack of fitted slices')
+    try:
+        z_step = float(options.get('z_step_um', 1.0))
+        side = float(options['pixel_size_um']) if 'pixel_size_um' in options else None
+    except ValueError:
+        raise RouteError(400, 'z_step_um and pixel_size_um must be numbers')
+    try:
+        return zstack.stream_volume(group_dir, label=options.get('label'),
+                                    z_step_um=z_step, pixel_size_um=side)
+    except FileNotFoundError as exc:
+        raise RouteError(404, str(exc))
+    except volumes.NothingToStack as exc:
+        raise RouteError(409, str(exc))
+    except Exception as exc:
+        raise RouteError(500, str(exc))
+
+
+def export_zstack(state, payload):
+    from flimkit_bridge import zstack
+    jobs = _jobs(state)
+    payload = payload or {}
+    group_dir = payload.get('group_dir')
+    if not group_dir:
+        raise RouteError(400, 'a group_dir is required, one stack of fitted slices')
+    from flimkit_bridge import volumes
+    wanted = payload.get('format', 'ome-zarr')
+    if wanted not in volumes.FORMATS:
+        raise RouteError(400, f'format must be one of {list(volumes.FORMATS)}')
+    from pathlib import Path
+    if not Path(group_dir).expanduser().is_dir():
+        raise RouteError(404, f'no such folder: {group_dir}')
+
+    def work(progress, cancel):
+        progress(0, 1, f'writing {wanted}')
+        found = zstack.export(
+            group_dir, label=payload.get('label'),
+            output_dir=payload.get('output_dir'), volume_format=wanted,
+            z_step_um=payload.get('z_step_um', 1.0),
+            pixel_size_um=payload.get('pixel_size_um'))
+        progress(1, 1, 'written')
+        return {'products': [found]}
+
+    job_id = jobs.submit('zstack_export', work)
+    return {'job': job_id, 'group_dir': str(group_dir), 'format': wanted}
+
+
 def fit_pixels(state, ident, payload):
     from flimkit_bridge import fitting
     registry = _registry(state)

flimkit_bridge/server.py

@@ -108,6 +108,12 @@ def do_GET(self):
                     return
                 self._route(lambda r: r.pipeline_defaults(state))
                 return
+            if self.path == '/v1/zstack/defaults':
+                if not self._authorized():
+                    self.send_error(401)
+                    return
+                self._route(lambda r: r.zstack_defaults(state))
+                return
             if self.path == '/v1/phasor/settings':
                 if not self._authorized():
                     self.send_error(401)
@@ -199,6 +205,24 @@ def do_POST(self):
                     return
                 self._route(lambda routes: routes.run_pipeline(state, self._read_json()))
                 return
+            if self.path == '/v1/zstack/scan':
+                if not self._authorized():
+                    self.send_error(401)
+                    return
+                self._route(lambda routes: routes.scan_zstack(state, self._read_json()))
+                return
+            if self.path == '/v1/zstack/export':
+                if not self._authorized():
+                    self.send_error(401)
+                    return
+                self._route(lambda routes: routes.export_zstack(state, self._read_json()))
+                return
+            if self.path == '/v1/zstack':
+                if not self._authorized():
+                    self.send_error(401)
+                    return
+                self._route(lambda routes: routes.run_zstack(state, self._read_json()))
+                return
             from flimkit_bridge import dataset_routes as _routes
             stats_match = _routes.PLANE_STATS_RE.match(self.path)
             if stats_match:
@@ -281,6 +305,12 @@ def do_DELETE(self):
         def _dataset_get(self):
             from flimkit_bridge import dataset_routes as routes
             parsed = urlparse(self.path)
+            if parsed.path == routes.ZSTACK_VOLUME_PATH:
+                if not self._authorized():
+                    self.send_error(401)
+                    return True
+                self._send_volume(parsed.query)
+                return True
             if parsed.path == '/v1/datasets':
                 if not self._authorized():
                     self.send_error(401)
@@ -335,6 +365,41 @@ def _dataset_get(self):
                 return True
             return False
 
+        def _send_volume(self, query):
+            import shutil
+
+            from flimkit_bridge import dataset_routes as routes
+            try:
+                built = routes.zstack_volume(state, query)
+            except routes.RouteError as problem:
+                self.send_error(problem.status, str(problem))
+                return
+            except Exception as exc:
+                self.send_error(500, str(exc))
+                return
+            # A fitted stack can be gigabytes, so it goes out of a file in
+            # blocks rather than through one bytes object in memory.
+            try:
+                size = os.path.getsize(built['file'])
+                self.send_response(200)
+                self.send_header('Content-Type', 'image/tiff')
+                self.send_header('Content-Length', str(size))
+                self.send_header('X-FLIMKit-Volume-Label', built['label'])
+                self.send_header('X-FLIMKit-Volume-Axes', 'ZCYX')
+                self.send_header('X-FLIMKit-Volume-Channels',
+                                 ','.join(built['channels']))
+                self.send_header('X-FLIMKit-Volume-Units',
+                                 ','.join(built['units']))
+                self.send_header('X-FLIMKit-Volume-Shape',
+                                 ','.join(str(n) for n in built['shape']))
+                self.send_header('X-FLIMKit-Voxel-Size-Um',
+                                 ','.join(str(v) for v in built['voxel_size_um']))
+                self.end_headers()
+                with open(built['file'], 'rb') as source:
+                    shutil.copyfileobj(source, self.wfile, 1024 * 1024)
+            finally:
+                shutil.rmtree(built['holding'], ignore_errors=True)
+
         def _send_plane(self, ident, name, query):
             from flimkit_bridge import dataset_routes as routes
             try: