Reuse shared features for paired 2D and 3D output

hecate.py

@@ -423,7 +423,50 @@ def main():
     if abs(args.spacing_um/model.sampling_um-1) > .02:
         parser.error(f'Resample the render to {model.sampling_um} um first (no automatic resampling)')
     volume = open_volume(args.input,args.array)
-    if output is not None:
+    if output is not None and output3d is not None:
+        from hecate_shared_output import predict_both
+        import zarr
+        from numcodecs import Blosc
+        output.parent.mkdir(parents=True,exist_ok=True)
+        output3d.parent.mkdir(parents=True,exist_ok=True)
+        partial = output3d.with_name(output3d.name+'.partial')
+        if partial.exists():
+            parser.error(f'Incomplete output exists: {partial}; remove it or choose another output')
+        options = {'zarr_format':2} if int(zarr.__version__.split('.')[0])>=3 else {}
+        result3 = zarr.open_array(str(partial),mode='w',shape=volume.shape,
+            chunks=(1,256,256),dtype='uint8',fill_value=0,
+            compressor=Blosc(cname='zstd',clevel=3,shuffle=Blosc.BITSHUFFLE),**options)
+        metadata = {
+            'checkpoint':Path(args.checkpoint).name,'sampling_um':args.spacing_um,
+            'shape_yx':list(volume.shape[-2:]),'reverse_z':args.reverse,
+            'central_depth':model.patch_size[0],'precision':args.precision,
+            'stride':args.stride or model.patch_size[1]//2,
+            'probabilities':'uint8 = round(255 * probability); no min-max scaling'
+        }
+        with tempfile.TemporaryDirectory(prefix='hecate-') as tmp:
+            result2 = np.memmap(Path(tmp)/'prediction',mode='w+',dtype='uint8',
+                                shape=volume.shape[-2:])
+            predict_both(model,volume,result2,result3,reverse=args.reverse,
+                         stride=args.stride,batch_size=args.batch_size,
+                         precision=args.precision)
+            result2.flush()
+            partial_png = output.with_suffix('.partial')
+            write_png(result2,partial_png)
+            partial_png.replace(output)
+            del result2
+        depth = volume.shape[0]
+        d = model.patch_size[0]
+        start = depth//2-d//2
+        if args.reverse:
+            start = depth-start-d
+        result3.attrs.update({**metadata,'axes':['z','y','x'],
+            'spacing_um':[args.spacing_um]*3,'z_order':'same as input',
+            'evaluated_z_interval':[start+model.margin,start+d-model.margin],
+            'outside_evaluated_z':'zero; not evaluated',
+            'background':'empty CT columns zeroed; no intensity threshold'})
+        partial.rename(output3d)
+        print(output3d)
+    elif output is not None:
         output.parent.mkdir(parents=True,exist_ok=True)
         with tempfile.TemporaryDirectory(prefix='hecate-') as tmp:
             result = np.memmap(Path(tmp)/'prediction',mode='w+',dtype='uint8',shape=volume.shape[-2:])
@@ -442,7 +485,7 @@ def main():
     if output is not None:
         output.with_suffix('.json').write_text(json.dumps(metadata,indent=2)+'\n')
         print(output)
-    if output3d is not None:
+    if output3d is not None and output is None:
         import zarr
         from numcodecs import Blosc
         output3d.parent.mkdir(parents=True,exist_ok=True)

hecate_shared_output.py

@@ -0,0 +1,203 @@
+"""Shared-feature tiled Hecate prediction for paired 2D and 3D outputs.
+
+This keeps the pinned Hecate provider and checkpoint untouched.  It computes the
+existing 2D attention head and 3D head from one ``features_logits`` call per tile,
+then blends both outputs with the same Hann weights.  It is an implementation
+helper, not a provider qualification or authorization to run on target material.
+"""
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import numpy as np
+
+
+def shared_logits(model, image, valid):
+    """Return the pinned Hecate 2D and 3D logits using one shared feature pass.
+
+    ``image`` is normalized B1ZYX float input; ``valid`` is the corresponding
+    boolean support mask.  The projections mirror ``Hecate.forward`` and
+    ``Hecate.forward_3d`` in the pinned provider, including 2.4 um depth margins.
+    """
+    import torch
+    import torch.nn.functional as F
+
+    if image.ndim != 5 or image.shape[1] != 1:
+        raise ValueError("image must have B1ZYX shape")
+    if valid is not None and tuple(valid.shape) != tuple(image.shape):
+        raise ValueError("valid support must match image shape")
+
+    features, logits = model.features_logits(image)
+    depth = image.shape[2]
+    margin = int(model.margin)
+    inner_depth = depth - 2 * margin
+    if inner_depth <= 0:
+        raise ValueError("model margin leaves no evaluated depth")
+
+    scores = model.canonical.decoder.depth_collapse.attn_conv(features).float()
+    half = scores.shape[2] // 2
+    z = (torch.arange(scores.shape[2], device=scores.device, dtype=torch.float32) - half) / half
+    scores = scores + model.depth_coordinate_scale * z.view(1, 1, -1, 1, 1)
+    if valid is None:
+        support = torch.ones_like(logits, dtype=torch.bool)
+    else:
+        support = F.max_pool3d(
+            valid[:, :, margin:depth - margin].float(),
+            (1, model.xy_stride, model.xy_stride),
+        ).bool()
+    weights = scores.masked_fill(~support, -1e4).softmax(2) * support
+    weights = weights / weights.sum(2, keepdim=True).clamp_min(1e-8)
+    probability_2d = F.interpolate(
+        (weights * logits).sum(2).sigmoid(), size=image.shape[-2:],
+        mode="bilinear", align_corners=False,
+    )
+    eps = torch.finfo(torch.float32).eps
+    logits_2d = torch.logit(probability_2d.clamp(eps, 1 - eps))
+
+    logits_3d = F.interpolate(
+        logits, size=(inner_depth, *image.shape[-2:]),
+        mode="trilinear", align_corners=False,
+    )
+    if margin:
+        logits_3d = F.pad(logits_3d, (0, 0, 0, 0, margin, margin), value=-20.0)
+    return logits_2d.float(), logits_3d.float()
+
+
+def _axis_starts(length: int, patch: int, stride: int) -> list[int]:
+    # Kept byte-for-byte equivalent in behavior to the pinned helper.
+    return list(range(0, max(0, length - patch), stride)) + [max(0, length - patch)]
+
+
+def predict_both(
+    model,
+    volume,
+    output_2d: np.ndarray,
+    output_3d: np.ndarray,
+    *,
+    reverse: bool = False,
+    stride: int | None = None,
+    batch_size: int = 1,
+    precision: str = "fp32",
+    before_batch: Callable[[], None] | None = None,
+    progress: Callable[[int, int], None] | None = None,
+) -> dict:
+    """Stream one input into both uint8 outputs, sharing Hecate's feature pass.
+
+    Accumulators are disk-backed temporary arrays.  Inputs and outputs follow
+    the pinned provider's ZYX/HW layouts, central-depth selection, reverse order,
+    Hann overlap weighting, and uint8 probability conversion.
+    """
+    import tempfile
+    from contextlib import ExitStack
+
+    import torch
+
+    if len(volume.shape) != 3 or np.dtype(volume.dtype) != np.dtype("uint8"):
+        raise ValueError("volume must be an unnormalized uint8 Z,Y,X array")
+    depth, height, width = map(int, volume.shape)
+    model_depth, patch, _ = model.patch_size
+    stride = patch // 2 if stride is None else int(stride)
+    batch_size = int(batch_size)
+    if depth < model_depth or not 0 < stride <= patch or batch_size < 1:
+        raise ValueError("insufficient depth, invalid stride, or invalid batch size")
+    if tuple(output_2d.shape) != (height, width) or output_2d.dtype != np.uint8:
+        raise ValueError("output_2d must be uint8 with shape (Y,X)")
+    if tuple(output_3d.shape) != (depth, height, width) or output_3d.dtype != np.uint8:
+        raise ValueError("output_3d must be uint8 with shape (Z,Y,X)")
+
+    device = next(model.parameters()).device
+    if precision not in ("fp32", "bf16") or (precision == "bf16" and device.type != "cuda"):
+        raise ValueError("bf16 requires CUDA; otherwise use fp32")
+
+    z0 = depth // 2 - model_depth // 2
+    selected_z0 = depth - z0 - model_depth if reverse else z0
+    zs = slice(selected_z0, selected_z0 + model_depth)
+    window_t = torch.hann_window(patch, periodic=False)
+    window_t = window_t[:, None] * window_t[None, :]
+    window_t = (window_t / window_t.max().clamp_min(torch.finfo(torch.float32).eps)).clamp_min(0.001)
+    window_np = window_t.numpy()
+    window_t = window_t.to(device)
+    origins = [(y, x) for y in _axis_starts(height, patch, stride)
+               for x in _axis_starts(width, patch, stride)]
+    output_3d[:] = 0
+    blank_tiles = 0
+
+    with tempfile.TemporaryDirectory(prefix="argus-hecate-shared-") as tmp, ExitStack() as stack:
+        def accumulator(name: str, shape: tuple[int, ...]):
+            array = np.memmap(f"{tmp}/{name}", mode="w+", dtype="float32", shape=shape)
+            stack.callback(array._mmap.close)
+            return array
+
+        numerator_2d = accumulator("sum2d", (height, width))
+        numerator_3d = accumulator("sum3d", (model_depth, height, width))
+        denominator = accumulator("weight", (height, width))
+        numerator_2d[:] = 0
+        numerator_3d[:] = 0
+        denominator[:] = 0
+
+        for start in range(0, len(origins), batch_size):
+            if before_batch is not None:
+                before_batch()
+            active, patches = [], []
+            for y, x in origins[start:start + batch_size]:
+                ph, pw = min(patch, height - y), min(patch, width - x)
+                denominator[y:y + ph, x:x + pw] += window_np[:ph, :pw]
+                raw = np.asarray(volume[zs, y:y + ph, x:x + pw])
+                if reverse:
+                    raw = raw[::-1]
+                if not raw.any():
+                    blank_tiles += 1
+                    continue
+                padded = np.zeros((model_depth, patch, patch), dtype=np.uint8)
+                padded[:, :ph, :pw] = raw
+                patches.append(padded)
+                active.append((y, x, ph, pw))
+
+            if patches:
+                raw_batch = np.stack(patches)
+                support = torch.from_numpy(raw_batch).to(device).any(1)
+                normalized = raw_batch.astype(np.float32)
+                normalized /= model.divisor
+                image = torch.from_numpy(normalized[:, None]).to(device)
+                valid = support[:, None, None].expand(-1, 1, model_depth, -1, -1)
+                with torch.inference_mode(), torch.autocast(
+                    device_type=device.type, dtype=torch.bfloat16, enabled=precision == "bf16"
+                ):
+                    logits_2d, logits_3d = shared_logits(model, image, valid)
+                    probability_2d = logits_2d.float().sigmoid()[:, 0]
+                    probability_3d = logits_3d.float().sigmoid()[:, 0]
+                    weighted_2d = (probability_2d * support * window_t).cpu().numpy()
+                    weighted_3d = (probability_3d * support[:, None] * window_t).cpu().numpy()
+                for pred2, pred3, (y, x, ph, pw) in zip(weighted_2d, weighted_3d, active):
+                    numerator_2d[y:y + ph, x:x + pw] += pred2[:ph, :pw]
+                    numerator_3d[:, y:y + ph, x:x + pw] += pred3[:, :ph, :pw]
+            if progress is not None:
+                progress(min(start + batch_size, len(origins)), len(origins))
+
+        if model.margin:
+            numerator_3d[:model.margin] = 0
+            numerator_3d[-model.margin:] = 0
+        for y in range(0, height, 128):
+            for x in range(0, width, 256):
+                y1, x1 = min(height, y + 128), min(width, x + 256)
+                den = denominator[y:y1, x:x1]
+                if not np.isfinite(den).all() or np.any(den <= 0):
+                    raise FloatingPointError("non-finite or uncovered output pixels")
+                prob2 = numerator_2d[y:y1, x:x1] / den
+                prob3 = numerator_3d[:, y:y1, x:x1] / den[None]
+                if not np.isfinite(prob2).all() or not np.isfinite(prob3).all():
+                    raise FloatingPointError("non-finite inference output")
+                values2 = np.rint(prob2.clip(0, 1) * 255).astype(np.uint8)
+                values3 = np.rint(prob3.clip(0, 1) * 255).astype(np.uint8)
+                output_2d[y:y1, x:x1] = values2
+                output_3d[zs, y:y1, x:x1] = values3[::-1] if reverse else values3
+
+        numerator_2d.flush()
+        numerator_3d.flush()
+        denominator.flush()
+        del numerator_2d, numerator_3d, denominator
+
+    return {"tile_count": len(origins), "blank_tiles_skipped": blank_tiles,
+            "evaluated_z_interval": [selected_z0 + model.margin,
+                                     selected_z0 + model_depth - model.margin],
+            "reverse": bool(reverse), "stride": stride, "precision": precision}

test_hecate_shared_output.py

@@ -0,0 +1,208 @@
+"""Parity tests for Hecate's shared 2D/3D feature projection."""
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+import numpy as np
+from torch import nn
+
+from hecate_shared_output import predict_both, shared_logits
+
+
+class _Collapse(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn_conv = nn.Conv3d(1, 1, 1, bias=True)
+
+
+class _Decoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.depth_collapse = _Collapse()
+
+
+class _Canonical(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.decoder = _Decoder()
+
+
+class _FakeHecate(nn.Module):
+    def __init__(self, margin=0, xy_stride=1):
+        super().__init__()
+        self.anchor = nn.Parameter(torch.zeros(()))
+        self.canonical = _Canonical()
+        self.depth_coordinate_scale = nn.Parameter(torch.tensor(0.17))
+        self.margin = margin
+        self.xy_stride = xy_stride
+        self.divisor = 255.0
+        self.patch_size = (6, 8, 8)
+        self.feature_calls = 0
+
+    def features_logits(self, image):
+        self.feature_calls += 1
+        features = image[:, :, self.margin:image.shape[2] - self.margin]
+        if self.xy_stride > 1:
+            features = F.avg_pool3d(
+                features, (1, self.xy_stride, self.xy_stride),
+                stride=(1, self.xy_stride, self.xy_stride),
+            )
+        logits = features * 0.73 + 0.11
+        return features, logits
+
+    def forward(self, image, valid=None):
+        features, logits = self.features_logits(image)
+        scores = self.canonical.decoder.depth_collapse.attn_conv(features).float()
+        half = scores.shape[2] // 2
+        z = (torch.arange(scores.shape[2], device=scores.device, dtype=torch.float32) - half) / half
+        scores = scores + self.depth_coordinate_scale * z.view(1, 1, -1, 1, 1)
+        if valid is None:
+            support = torch.ones_like(logits, dtype=torch.bool)
+        else:
+            support = F.max_pool3d(
+                valid[:, :, self.margin:image.shape[2] - self.margin].float(),
+                (1, self.xy_stride, self.xy_stride),
+            ).bool()
+        weights = scores.masked_fill(~support, -1e4).softmax(2) * support
+        weights = weights / weights.sum(2, keepdim=True).clamp_min(1e-8)
+        probability = F.interpolate(
+            (weights * logits).sum(2).sigmoid(), size=image.shape[-2:],
+            mode="bilinear", align_corners=False,
+        )
+        eps = torch.finfo(torch.float32).eps
+        return torch.logit(probability.clamp(eps, 1 - eps))
+
+    def forward_3d(self, image):
+        _, logits = self.features_logits(image)
+        full = F.interpolate(
+            logits, size=(image.shape[2] - 2 * self.margin, *image.shape[-2:]),
+            mode="trilinear", align_corners=False,
+        )
+        return F.pad(full, (0, 0, 0, 0, self.margin, self.margin), value=-20.0)
+
+
+def test_shared_heads_match_separate_pinned_head_math_without_duplicate_features():
+    torch.manual_seed(7)
+    model = _FakeHecate().eval()
+    image = torch.rand(2, 1, 6, 12, 10)
+    valid = torch.ones_like(image, dtype=torch.bool)
+    valid[:, :, :, :2, :3] = False
+
+    expected_2d = model(image, valid)
+    expected_3d = model.forward_3d(image)
+    separate_calls = model.feature_calls
+    assert separate_calls == 2
+
+    actual_2d, actual_3d = shared_logits(model, image, valid)
+    assert model.feature_calls - separate_calls == 1
+    assert torch.equal(actual_2d, expected_2d)
+    assert torch.equal(actual_3d, expected_3d)
+
+
+def test_shared_heads_preserve_two_point_four_micron_margin_padding():
+    torch.manual_seed(11)
+    model = _FakeHecate(margin=1, xy_stride=2).eval()
+    image = torch.rand(1, 1, 6, 12, 10)
+    valid = torch.ones_like(image, dtype=torch.bool)
+    valid[:, :, :, :2, :3] = False
+
+    expected_2d = model(image, valid)
+    expected_3d = model.forward_3d(image)
+    actual_2d, actual_3d = shared_logits(model, image, valid)
+    assert torch.equal(actual_2d, expected_2d)
+    assert torch.equal(actual_3d, expected_3d)
+    assert torch.all(actual_3d[:, :, 0] == -20.0)
+    assert torch.all(actual_3d[:, :, -1] == -20.0)
+
+
+@torch.inference_mode()
+def _stock_predict(model, volume, output, *, reverse, stride, three_d):
+    """Small CPU reference matching the pinned provider's separate-head tiler."""
+    d, patch, _ = model.patch_size
+    depth, height, width = volume.shape
+    z0 = depth // 2 - d // 2
+    zs = slice(depth - z0 - d, depth - z0) if reverse else slice(z0, z0 + d)
+    window = torch.hann_window(patch, periodic=False)
+    window = window[:, None] * window[None, :]
+    window = (window / window.max().clamp_min(torch.finfo(torch.float32).eps)).clamp_min(0.001)
+    numerator = torch.zeros((d, height, width) if three_d else (height, width), dtype=torch.float32)
+    denominator = torch.zeros((height, width), dtype=torch.float32)
+    starts_y = list(range(0, max(0, height - patch), stride)) + [max(0, height - patch)]
+    starts_x = list(range(0, max(0, width - patch), stride)) + [max(0, width - patch)]
+    for y in starts_y:
+        for x in starts_x:
+            ph, pw = min(patch, height - y), min(patch, width - x)
+            denominator[y:y + ph, x:x + pw] += window[:ph, :pw]
+            raw = np.asarray(volume[zs, y:y + ph, x:x + pw])
+            if reverse:
+                raw = raw[::-1]
+            if not raw.any():
+                continue
+            padded = np.zeros((d, patch, patch), dtype=np.uint8)
+            padded[:, :ph, :pw] = raw
+            support = torch.from_numpy(padded[None]).any(1)
+            image = torch.from_numpy(padded.astype(np.float32)[None, None] / model.divisor)
+            valid = support[:, None, None].expand(-1, 1, d, -1, -1)
+            logits = model.forward_3d(image) if three_d else model(image, valid)
+            probability = logits.float().sigmoid()[:, 0]
+            if three_d:
+                weighted = probability * support[:, None] * window
+                numerator[:, y:y + ph, x:x + pw] += weighted[0, :, :ph, :pw]
+            else:
+                weighted = probability * support * window
+                numerator[y:y + ph, x:x + pw] += weighted[0, :ph, :pw]
+    if three_d and model.margin:
+        numerator[:model.margin] = 0
+        numerator[-model.margin:] = 0
+    probability = numerator / (denominator[None] if three_d else denominator)
+    values = torch.round(probability.clamp(0, 1) * 255).to(torch.uint8).numpy()
+    if three_d:
+        output[zs] = values[::-1] if reverse else values
+    else:
+        output[:] = values
+
+
+def test_tiled_outputs_match_two_stock_passes_and_skip_blank_tiles():
+    torch.manual_seed(19)
+    model = _FakeHecate().eval()
+    volume = np.zeros((6, 13, 15), dtype=np.uint8)
+    volume[:, :7, :7] = np.random.default_rng(19).integers(1, 255, size=(6, 7, 7), dtype=np.uint8)
+    expected_2d = np.zeros((13, 15), dtype=np.uint8)
+    expected_3d = np.zeros((6, 13, 15), dtype=np.uint8)
+    _stock_predict(model, volume, expected_2d, reverse=False, stride=4, three_d=False)
+    _stock_predict(model, volume, expected_3d, reverse=False, stride=4, three_d=True)
+
+    actual_2d = np.zeros_like(expected_2d)
+    actual_3d = np.zeros_like(expected_3d)
+    result = predict_both(model, volume, actual_2d, actual_3d, stride=4, batch_size=1)
+    assert result["tile_count"] == 9
+    assert result["blank_tiles_skipped"] > 0
+    assert np.array_equal(actual_2d, expected_2d)
+    assert np.array_equal(actual_3d, expected_3d)
+
+
+def test_center_depth_crop_and_reverse_order_match_two_stock_passes():
+    """Longer inputs select the centered model-depth window in either direction."""
+    for reverse in (False, True):
+        torch.manual_seed(31)
+        model = _FakeHecate().eval()
+        rng = np.random.default_rng(31)
+        volume = rng.integers(0, 256, size=(9, 11, 13), dtype=np.uint8)
+        expected_2d = np.zeros((11, 13), dtype=np.uint8)
+        expected_3d = np.zeros((9, 11, 13), dtype=np.uint8)
+        _stock_predict(model, volume, expected_2d, reverse=reverse, stride=4, three_d=False)
+        _stock_predict(model, volume, expected_3d, reverse=reverse, stride=4, three_d=True)
+
+        actual_2d = np.zeros_like(expected_2d)
+        actual_3d = np.zeros_like(expected_3d)
+        result = predict_both(
+            model, volume, actual_2d, actual_3d, reverse=reverse, stride=4, batch_size=1,
+        )
+        z0 = volume.shape[0] // 2 - model.patch_size[0] // 2
+        selected_z0 = volume.shape[0] - z0 - model.patch_size[0] if reverse else z0
+        assert result["tile_count"] == 6
+        assert np.array_equal(actual_2d, expected_2d)
+        assert np.array_equal(actual_3d, expected_3d)
+        assert np.all(actual_3d[:selected_z0] == 0)
+        assert np.all(actual_3d[selected_z0 + model.patch_size[0]:] == 0)
+        assert np.any(actual_3d[selected_z0:selected_z0 + model.patch_size[0]] > 0)