spine-segment 0.1.0__py3-none-any.whl

spine_segment/__init__.py

@@ -0,0 +1,16 @@
+from spine_segment.api import segment_file, segment_files
+from spine_segment.backend import (
+    SegmentationResult,
+    SpineSegmentBackend,
+    SpineSegmentBackendError,
+    load_backend,
+)
+
+__all__ = [
+    "SegmentationResult",
+    "SpineSegmentBackend",
+    "SpineSegmentBackendError",
+    "load_backend",
+    "segment_file",
+    "segment_files",
+]

spine_segment/api.py

@@ -0,0 +1,214 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+import json
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import SimpleITK as sitk
+
+from spine_segment.backend import SpineSegmentBackend, SpineSegmentBackendError
+from spine_segment.compartments import CortTrabConfig, derive_cort_trab_labels
+from spine_segment.device import resolve_device
+from spine_segment.io import read_image, write_image
+from spine_segment.outputs import OutputPaths, build_output_paths
+
+
+@dataclass(frozen=True, slots=True)
+class SegmentRunResult:
+    input_path: Path
+    output_paths: OutputPaths
+    device: str
+    metadata: dict[str, Any]
+
+
+def segment_file(
+    input_path: str | Path,
+    output_dir: str | Path,
+    *,
+    backend: SpineSegmentBackend,
+    device: str = "auto",
+    overwrite: bool = False,
+    cort_trab_config: CortTrabConfig | None = None,
+    level_only: bool = False,
+    localization_only: bool = False,
+) -> SegmentRunResult:
+    source_path = Path(input_path)
+    image = read_image(source_path, pixel_type=sitk.sitkFloat32)
+    selected_device = resolve_device(device)
+    output_paths = build_output_paths(source_path, output_dir)
+    if localization_only:
+        localization = backend.localize(
+            image=image,
+            source_path=source_path,
+            device=selected_device,
+        )
+        centroids = localization.centroids
+        metadata = dict(localization.metadata or {})
+        metadata["device"] = selected_device
+        metadata["localization_only"] = True
+        write_json(
+            {
+                "input": str(source_path),
+                "coordinate_system": "voxel_xyz",
+                "centroids": centroids,
+                "metadata": metadata,
+            },
+            output_paths.centroids,
+            overwrite=overwrite,
+        )
+        return SegmentRunResult(
+            input_path=source_path,
+            output_paths=output_paths,
+            device=selected_device,
+            metadata=metadata,
+        )
+
+    result = backend.segment(
+        image=image,
+        source_path=source_path,
+        device=selected_device,
+        level_only=level_only,
+    )
+
+    write_image(result.vertebral_level, output_paths.vertebral_level, overwrite=overwrite)
+    centroids = centroids_from_labelmap(result.vertebral_level)
+    if level_only:
+        metadata = dict(result.metadata or {})
+        metadata["device"] = selected_device
+        metadata["level_only"] = True
+        write_centroids_json(
+            input_path=source_path,
+            centroids=centroids,
+            metadata=metadata,
+            path=output_paths.centroids,
+            overwrite=overwrite,
+        )
+        return SegmentRunResult(
+            input_path=source_path,
+            output_paths=output_paths,
+            device=selected_device,
+            metadata=metadata,
+        )
+
+    if result.process_body is None:
+        raise SpineSegmentBackendError(
+            "Standalone spine backend did not return a process/body labelmap."
+        )
+
+    cort_trab = result.cort_trab or derive_cort_trab_labels(
+        image=image,
+        process_body=result.process_body,
+        vertebral_level=result.vertebral_level,
+        config=cort_trab_config,
+    )
+
+    write_image(result.process_body, output_paths.process_body, overwrite=overwrite)
+    write_image(cort_trab, output_paths.cort_trab, overwrite=overwrite)
+
+    metadata = dict(result.metadata or {})
+    metadata["device"] = selected_device
+    write_centroids_json(
+        input_path=source_path,
+        centroids=centroids,
+        metadata=metadata,
+        path=output_paths.centroids,
+        overwrite=overwrite,
+    )
+    return SegmentRunResult(
+        input_path=source_path,
+        output_paths=output_paths,
+        device=selected_device,
+        metadata=metadata,
+    )
+
+
+def segment_files(
+    input_paths: list[str | Path],
+    output_dir: str | Path,
+    *,
+    backend: SpineSegmentBackend,
+    device: str = "auto",
+    overwrite: bool = False,
+    cort_trab_config: CortTrabConfig | None = None,
+    level_only: bool = False,
+    localization_only: bool = False,
+) -> list[SegmentRunResult]:
+    return [
+        segment_file(
+            path,
+            output_dir,
+            backend=backend,
+            device=device,
+            overwrite=overwrite,
+            cort_trab_config=cort_trab_config,
+            level_only=level_only,
+            localization_only=localization_only,
+        )
+        for path in input_paths
+    ]
+
+
+def write_json(payload: dict[str, Any], path: str | Path, *, overwrite: bool = False) -> Path:
+    output_path = Path(path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    if output_path.exists() and not overwrite:
+        raise FileExistsError(
+            f"Refusing to overwrite existing output without --overwrite: {output_path}"
+        )
+    output_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
+    return output_path
+
+
+def write_centroids_json(
+    *,
+    input_path: Path,
+    centroids: dict[str, dict[str, Any]],
+    metadata: dict[str, Any],
+    path: str | Path,
+    overwrite: bool,
+) -> Path:
+    return write_json(
+        {
+            "input": str(input_path),
+            "coordinate_system": "voxel_xyz",
+            "centroids": centroids,
+            "metadata": metadata,
+        },
+        path,
+        overwrite=overwrite,
+    )
+
+
+def centroids_from_labelmap(labelmap: sitk.Image) -> dict[str, dict[str, Any]]:
+    labels = sitk.GetArrayFromImage(labelmap)
+    centroids: dict[str, dict[str, Any]] = {}
+    for raw_label in sorted(int(value) for value in np.unique(labels) if int(value) != 0):
+        coords_zyx = np.argwhere(labels == raw_label)
+        if coords_zyx.size == 0:
+            continue
+        centroid_zyx = coords_zyx.mean(axis=0)
+        voxel_xyz = (
+            float(centroid_zyx[2]),
+            float(centroid_zyx[1]),
+            float(centroid_zyx[0]),
+        )
+        physical_xyz = labelmap.TransformContinuousIndexToPhysicalPoint(voxel_xyz)
+        centroids[str(raw_label)] = {
+            "label": raw_label,
+            "index": _verse_index_from_label(raw_label),
+            "voxel_xyz": [float(value) for value in voxel_xyz],
+            "physical_xyz": [float(value) for value in physical_xyz],
+            "voxel_count": int(coords_zyx.shape[0]),
+            "source": "segmentation",
+        }
+    return centroids
+
+
+def _verse_index_from_label(label: int) -> int | None:
+    if 1 <= int(label) <= 25:
+        return int(label) - 1
+    if int(label) == 28:
+        return 25
+    return None

spine_segment/backend.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from importlib import import_module
+from pathlib import Path
+from typing import Any, Protocol, runtime_checkable
+
+import SimpleITK as sitk
+
+
+class SpineSegmentBackendError(RuntimeError):
+    """Raised when the standalone spine backend cannot run."""
+
+
+@dataclass(slots=True)
+class SegmentationResult:
+    vertebral_level: sitk.Image
+    process_body: sitk.Image | None = None
+    cort_trab: sitk.Image | None = None
+    metadata: dict[str, Any] | None = None
+
+
+@dataclass(slots=True)
+class LocalizationResult:
+    centroids: dict[str, dict[str, Any]]
+    metadata: dict[str, Any] | None = None
+
+
+@runtime_checkable
+class SpineSegmentBackend(Protocol):
+    def segment(
+        self,
+        *,
+        image: sitk.Image,
+        source_path: Path,
+        device: str,
+        level_only: bool = False,
+    ) -> SegmentationResult: ...
+
+    def localize(
+        self,
+        *,
+        image: sitk.Image,
+        source_path: Path,
+        device: str,
+    ) -> LocalizationResult: ...
+
+
+class UnavailableBackend:
+    def segment(
+        self,
+        *,
+        image: sitk.Image,
+        source_path: Path,
+        device: str,
+        level_only: bool = False,
+    ) -> SegmentationResult:
+        raise SpineSegmentBackendError(
+            "No standalone spine backend is wired yet. "
+            "Provide --backend module_path:factory or set SPINE_SEGMENT_BACKEND "
+            "to the PyTorch runtime entrypoint for the vertebral-level and "
+            "process/body models."
+        )
+
+    def localize(
+        self,
+        *,
+        image: sitk.Image,
+        source_path: Path,
+        device: str,
+    ) -> LocalizationResult:
+        raise SpineSegmentBackendError(
+            "No standalone spine backend is wired yet. "
+            "Provide --backend module_path:factory or set SPINE_SEGMENT_BACKEND."
+        )
+
+
+def load_backend(spec: str | None) -> SpineSegmentBackend:
+    if spec is None or not str(spec).strip():
+        return UnavailableBackend()
+
+    module_name, separator, attr_name = str(spec).partition(":")
+    module = import_module(module_name)
+    target_name = attr_name if separator else "create_backend"
+    target = getattr(module, target_name)
+    instance = target() if callable(target) else target
+
+    if not isinstance(instance, SpineSegmentBackend) and not hasattr(instance, "segment"):
+        raise SpineSegmentBackendError(
+            f"Backend '{spec}' does not expose a segment(...) method."
+        )
+    return instance

spine_segment/backend_template.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import SimpleITK as sitk
+
+from spine_segment.backend import SegmentationResult, SpineSegmentBackend
+
+
+class ExampleBackend(SpineSegmentBackend):
+    """
+    Replace this class with the real PyTorch runtime.
+
+    The backend must return:
+    - vertebral_level: integer vertebral labelmap in scan space
+    - process_body: integer process/body labelmap in scan space
+    - cort_trab: optional integer cort/trab labelmap in scan space
+    """
+
+    def segment(
+        self,
+        *,
+        image: sitk.Image,
+        source_path: Path,
+        device: str,
+    ) -> SegmentationResult:
+        raise NotImplementedError(
+            "Replace spine_segment.backend_template.ExampleBackend with the "
+            "PyTorch vertebral segmentation runtime."
+        )
+
+
+def create_backend() -> ExampleBackend:
+    return ExampleBackend()

spine_segment/checkpoint_mapping.py

@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+import re
+
+
+@dataclass(frozen=True, slots=True)
+class MappingRule:
+    pattern: re.Pattern[str]
+    replacement: str
+    tensor_kind: str
+
+
+@dataclass(frozen=True, slots=True)
+class ConversionCoverage:
+    converted_count: int
+    ignored_source_count: int
+    missing_target_keys: tuple[str, ...]
+    unmapped_source_names: tuple[str, ...]
+    shape_mismatches: tuple[dict[str, object], ...]
+
+
+_VARIABLE_VALUE_SUFFIX = "/.ATTRIBUTES/VARIABLE_VALUE"
+_IGNORED_EXACT_NAMES = {
+    "_CHECKPOINTABLE_OBJECT_GRAPH",
+    "save_counter",
+    "sigmas",
+}
+_IGNORED_PREFIXES = (
+    "optimizer/",
+    "save_counter/",
+)
+_IGNORED_SUBSTRINGS = (
+    "/.OPTIMIZER_SLOT/",
+)
+
+_UNET_BLOCK_RULE = re.compile(
+    r"^model/(?P<prefix>unet|scnet_local|scnet_spatial)/"
+    r"(?P<block>contracting_layers|expanding_layers)/level(?P<level>\d+)/"
+    r"layers/(?P<layer_index>\d+)/(?P<var>kernel|bias)$"
+)
+_PREDICTION_RULE = re.compile(
+    r"^model/prediction/layers/0/(?P<var>kernel|bias)$"
+)
+_LOCAL_HEATMAP_RULE = re.compile(
+    r"^model/local_heatmaps/layers/0/(?P<var>kernel|bias)$"
+)
+_SPATIAL_HEATMAP_RULE = re.compile(
+    r"^model/spatial_heatmaps/(?P<var>kernel|bias)$"
+)
+
+
+def normalize_tf_variable_name(variable_name: str) -> str:
+    normalized = str(variable_name).strip()
+    if normalized.endswith(_VARIABLE_VALUE_SUFFIX):
+        normalized = normalized[: -len(_VARIABLE_VALUE_SUFFIX)]
+    return normalized
+
+
+def should_ignore_tf_variable(variable_name: str) -> bool:
+    normalized = normalize_tf_variable_name(variable_name)
+    if normalized in _IGNORED_EXACT_NAMES:
+        return True
+    if any(token in normalized for token in _IGNORED_SUBSTRINGS):
+        return True
+    return normalized.startswith(_IGNORED_PREFIXES)
+
+
+def map_tf_variable_to_torch_key(variable_name: str) -> str | None:
+    if should_ignore_tf_variable(variable_name):
+        return None
+
+    normalized = normalize_tf_variable_name(variable_name)
+    match = _UNET_BLOCK_RULE.match(normalized)
+    if match:
+        prefix = match.group("prefix")
+        block = match.group("block")
+        level = int(match.group("level"))
+        layer_index = int(match.group("layer_index"))
+        variable = "weight" if match.group("var") == "kernel" else "bias"
+        conv_index = layer_index // 2
+        pytorch_conv_index = conv_index * 2
+
+        if prefix == "unet":
+            base = "unet"
+        elif prefix == "scnet_local":
+            base = "scnet_local"
+        elif prefix == "scnet_spatial":
+            base = "scnet_spatial"
+        else:
+            return None
+
+        block_name = "contracting" if block == "contracting_layers" else "expanding"
+        return f"{base}.{block_name}.{level}.{pytorch_conv_index}.{variable}"
+
+    match = _PREDICTION_RULE.match(normalized)
+    if match:
+        variable = "weight" if match.group("var") == "kernel" else "bias"
+        return f"prediction.{variable}"
+
+    match = _LOCAL_HEATMAP_RULE.match(normalized)
+    if match:
+        variable = "weight" if match.group("var") == "kernel" else "bias"
+        return f"local_heatmaps.0.{variable}"
+
+    match = _SPATIAL_HEATMAP_RULE.match(normalized)
+    if match:
+        variable = "weight" if match.group("var") == "kernel" else "bias"
+        return f"spatial_heatmaps.{variable}"
+
+    return None
+
+
+def needs_conv_kernel_transpose(variable_name: str) -> bool:
+    return normalize_tf_variable_name(variable_name).endswith("/kernel")
+
+
+def check_conversion_coverage(
+    *,
+    converted_keys: set[str],
+    target_keys: set[str],
+    ignored_source_names: list[str],
+    unmapped_source_names: list[str],
+    shape_mismatches: list[dict[str, object]],
+) -> ConversionCoverage:
+    return ConversionCoverage(
+        converted_count=len(converted_keys),
+        ignored_source_count=len(ignored_source_names),
+        missing_target_keys=tuple(sorted(target_keys - converted_keys)),
+        unmapped_source_names=tuple(sorted(unmapped_source_names)),
+        shape_mismatches=tuple(shape_mismatches),
+    )

spine_segment/cli.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+
+from spine_segment.api import segment_files
+from spine_segment.backend import SpineSegmentBackendError, load_backend
+from spine_segment.compartments import CortTrabConfig
+from spine_segment.io import expand_input_paths
+from spine_segment.native_torch_backend import create_backend as create_native_backend
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="spine-segment",
+        description="Standalone vertebral segmentation CLI.",
+    )
+    parser.add_argument("inputs", nargs="+", help="Input CT volumes (.nii or .nii.gz).")
+    parser.add_argument(
+        "--output",
+        required=True,
+        help="Directory that will receive outputs.",
+    )
+    parser.add_argument(
+        "--device",
+        default="auto",
+        choices=("auto", "cuda", "mps", "cpu"),
+        help="Runtime device preference.",
+    )
+    parser.add_argument(
+        "--backend",
+        default=os.environ.get("SPINE_SEGMENT_BACKEND", ""),
+        help="Backend factory in the form module_path:factory. "
+        "Defaults to the native PyTorch backend.",
+    )
+    parser.add_argument(
+        "--model-bundle",
+        default=os.environ.get("SPINE_SEGMENT_MODEL_BUNDLE", ""),
+        help="Path to a spine-segment model bundle. Defaults to "
+        "SPINE_SEGMENT_MODEL_BUNDLE, ./model-bundle, ./build/model-bundle-pytorch, "
+        "./build/model-bundle, or an automatically downloaded cached bundle.",
+    )
+    parser.add_argument(
+        "--no-model-download",
+        action="store_true",
+        default=os.environ.get("SPINE_SEGMENT_NO_DOWNLOAD", "").strip().lower()
+        not in ("", "0", "false", "no", "off"),
+        help="Disable automatic first-run model bundle download.",
+    )
+    parser.add_argument(
+        "--overwrite",
+        action="store_true",
+        help="Overwrite existing outputs.",
+    )
+    parser.add_argument(
+        "--level-only",
+        action="store_true",
+        help="Write *_vertebral-level.nii.gz and *_centroids.json; skip process/body plus cort/trab stages.",
+    )
+    parser.add_argument(
+        "--localization-only",
+        action="store_true",
+        help="Write only *_centroids.json with detected vertebral centroids in original scan voxel coordinates.",
+    )
+    parser.add_argument(
+        "--vertebra-batch-size",
+        type=int,
+        default=1,
+        help="Number of MDAT vertebra segmentation patches to run per PyTorch batch.",
+    )
+    parser.add_argument(
+        "--process-body-batch-size",
+        type=int,
+        default=4,
+        help="Number of process/body relabel crops to run per PyTorch batch.",
+    )
+    parser.add_argument(
+        "--cortical-threshold-hu",
+        type=float,
+        default=500.0,
+        help="HU threshold used for connected cortical assignment within the outer shell.",
+    )
+    parser.add_argument(
+        "--cort-trab-isotropic-spacing-mm",
+        type=float,
+        default=1.0,
+        help="Isotropic spacing used internally for cort/trab shell generation.",
+    )
+    parser.add_argument(
+        "--cortical-max-thickness-mm",
+        type=float,
+        default=6.0,
+        help="Maximum distance from the vertebral surface considered for threshold-based cortical assignment.",
+    )
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    inputs = expand_input_paths(list(args.inputs))
+    if not inputs:
+        parser.error("No input images were resolved.")
+
+    missing = [path for path in inputs if not path.exists()]
+    if missing:
+        parser.error(f"Input does not exist: {missing[0]}")
+
+    if args.backend:
+        backend = load_backend(args.backend)
+    else:
+        backend = create_native_backend(
+            args.model_bundle or "./build/model-bundle",
+            allow_download=not bool(args.no_model_download),
+        )
+    if hasattr(backend, "vertebra_segmentation_batch_size"):
+        backend.vertebra_segmentation_batch_size = max(1, int(args.vertebra_batch_size))
+    if hasattr(backend, "process_body_batch_size"):
+        backend.process_body_batch_size = max(1, int(args.process_body_batch_size))
+    cort_trab_config = CortTrabConfig(
+        cortical_threshold_hu=float(args.cortical_threshold_hu),
+        isotropic_spacing_mm=float(args.cort_trab_isotropic_spacing_mm),
+        cortical_max_thickness_mm=float(args.cortical_max_thickness_mm),
+    )
+    if args.localization_only and args.level_only:
+        parser.error("--localization-only and --level-only are mutually exclusive.")
+
+    try:
+        results = segment_files(
+            input_paths=inputs,
+            output_dir=Path(args.output),
+            backend=backend,
+            device=args.device,
+            overwrite=bool(args.overwrite),
+            cort_trab_config=cort_trab_config,
+            level_only=bool(args.level_only),
+            localization_only=bool(args.localization_only),
+        )
+    except SpineSegmentBackendError as exc:
+        print(f"[spine-segment] backend error: {exc}")
+        return 2
+    except FileExistsError as exc:
+        print(f"[spine-segment] output exists: {exc}")
+        return 2
+
+    for result in results:
+        fields = [
+            f"input={result.input_path}",
+            f"device={result.device}",
+        ]
+        if args.localization_only:
+            fields.append(f"centroids={result.output_paths.centroids}")
+        else:
+            fields.append(f"vertebral_level={result.output_paths.vertebral_level}")
+            fields.append(f"centroids={result.output_paths.centroids}")
+            if not args.level_only:
+                fields.append(f"process_body={result.output_paths.process_body}")
+                fields.append(f"cort_trab={result.output_paths.cort_trab}")
+        print("[spine-segment] " + " ".join(fields))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())