musclemap-mlx 0.1.0__tar.gz

musclemap_mlx-0.1.0/.gitignore

@@ -0,0 +1,38 @@
+# MedMLX standard
+__pycache__/
+*.py[cod]
+*.egg-info/
+.venv
+build/
+dist/
+.pytest_cache/
+.ruff_cache/
+.mypy_cache/
+.hypothesis/
+.coverage
+htmlcov/
+.DS_Store
+# Never commit secrets (Hugging Face tokens)
+.env
+# Never commit data or weights
+/data/
+/work/
+/weights/
+/models/
+/checkpoints/
+/outputs/
+*.nii
+*.nii.gz
+*.dcm
+*.nrrd
+*.mha
+*.mhd
+*.safetensors
+*.pt
+*.pth
+*.ckpt
+*.bin
+*.h5
+*.npy
+*.npz
+!tests/fixtures/**/*.npz

musclemap_mlx-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 MuscleMap
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

musclemap_mlx-0.1.0/PKG-INFO

@@ -0,0 +1,155 @@
+Metadata-Version: 2.5
+Name: musclemap-mlx
+Version: 0.1.0
+Summary: Segment 113 whole-body muscle and bone structures with MuscleMap v1.4 and native MLX
+Project-URL: Homepage, https://medmlx.org
+Project-URL: Repository, https://github.com/MedMLX/musclemap-mlx
+Project-URL: Issues, https://github.com/MedMLX/musclemap-mlx/issues
+Project-URL: Weights, https://huggingface.co/MedMLX/musclemap-mlx
+Project-URL: Upstream, https://github.com/MuscleMap/MuscleMap
+Project-URL: MedMLX, https://github.com/MedMLX/medmlx
+Author: Joseph Sandoval, Alfredo Sandoval Jr.
+License-Expression: MIT
+License-File: LICENSE
+License-File: THIRD_PARTY_NOTICES.md
+License-File: UPSTREAM_LICENSE
+Keywords: apple-silicon,ct,medical-imaging,metal,mlx,mri,musclemap,segmentation,whole-body
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: MacOS :: MacOS X
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Image Processing
+Classifier: Topic :: Scientific/Engineering :: Medical Science Apps.
+Classifier: Typing :: Typed
+Requires-Python: <3.14,>=3.13
+Requires-Dist: medmlx-core<0.3,>=0.2.7
+Requires-Dist: mlx==0.32.4; platform_system == 'Darwin' and platform_machine == 'arm64'
+Requires-Dist: nibabel<6,>=5.3
+Requires-Dist: numpy<3,>=2.0
+Requires-Dist: scipy<2,>=1.14
+Description-Content-Type: text/markdown
+
+# musclemap-mlx
+
+Segment 113 muscle and bone structures in MRI or CT NIfTI images with
+MuscleMap v1.4, running natively through MLX and Metal on Apple Silicon.
+Weights download automatically on first use.
+
+## Install and segment
+
+Requires an Apple Silicon Mac, Metal and Python 3.13. With
+[uv](https://docs.astral.sh/uv/):
+
+```sh
+uv venv --python 3.13
+source .venv/bin/activate
+uv pip install musclemap-mlx
+
+# Real MRI example from TotalSegmentator's official test data.
+curl -L https://raw.githubusercontent.com/wasserth/TotalSegmentator/d89e47f73d24f5fba732a715aeb6fdc0ae34526b/tests/reference_files/example_mr_sm.nii.gz -o mri.nii.gz
+musclemap-mlx segment --input mri.nii.gz --output-dir out
+```
+
+The command writes `out/mask.nii.gz` and prints JSON containing the output path,
+label table, device, memory use and chunking settings. To use the corresponding
+CT example, replace `example_mr_sm.nii.gz` with `example_ct_sm.nii.gz`.
+
+### Python
+
+```python
+from musclemap_mlx.medmlx import run
+
+result = run(image_path="mri.nii.gz", output_dir="out-python")
+print(result["outputs"]["mask"])
+```
+
+### MedMLX
+
+```sh
+uv pip install 'medmlx[segment]>=0.9.1,<0.10'
+medmlx describe musclemap
+medmlx setup musclemap
+medmlx run musclemap --arg image_path=mri.nii.gz output_dir=out-medmlx
+```
+
+Or use its Python interface:
+
+```python
+import medmlx
+
+medmlx.setup_model("musclemap")
+job = medmlx.run("musclemap", arguments={
+    "image_path": "mri.nii.gz",
+    "output_dir": "out-medmlx-python",
+})
+print(job["result"])
+```
+
+The MCP server exposes the same model through `list_models`, `describe_model`,
+`model_readiness`, `prepare_model` and `run_model`; poll `get_job` for setup and
+segmentation results. Start it with `medmlx serve --root .`.
+
+## Inputs and outputs
+
+Input is one finite 3D `.nii` or `.nii.gz` MRI or CT volume. Millimetre spatial
+units are supported; an unspecified unit is interpreted as millimetres, matching
+upstream. Other declared spatial units and 4D images are rejected.
+
+The output is an `int16` mask with background 0 and the original 113 anatomical
+IDs, including 1101 for left levator scapulae and 8222 for right plantaris. Its
+shape, affine, qform, sform and spatial-unit metadata follow the input. The
+largest 6-connected component of each label is retained after chunk assembly.
+The JSON result lists every anatomical ID and name.
+
+| Option | Meaning | Default |
+| --- | --- | --- |
+| `--input` | 3D NIfTI image | Required |
+| `--output-dir` | Destination for `mask.nii.gz` | Required |
+| `--model-dir` | Local folder containing `model.safetensors` | Automatic download/cache |
+| `--chunk-size` | Source storage-axis slices per independent pass | Eight |
+| `--overlap` | Gaussian window overlap, integer 0–99 percent | 90 |
+| `--overwrite` | Replace an existing output | Off |
+
+Each chunk is independently normalized. For integer source images, chunking
+also reproduces upstream's source-header NIfTI round trip before normalization.
+Changing chunk size can change the mask. A chunk size at least the source depth
+uses one whole-volume pass. Inference uses MLX on Metal; there is no CPU or
+PyTorch network fallback.
+
+## Weights and offline use
+
+The approximately 105 MB float32 bundle is hosted at
+[MedMLX/musclemap-mlx](https://huggingface.co/MedMLX/musclemap-mlx/tree/a798b3073e3b2c449987c13aa193626aaacc3092),
+pinned to revision `a798b3073e3b2c449987c13aa193626aaacc3092`. It includes the
+publisher configuration, original label metadata and MIT notices.
+
+```sh
+# Pre-download to the shared cache; subsequent runs can work offline.
+musclemap-mlx download
+
+# Export a local bundle, then supply it explicitly.
+musclemap-mlx download --model-dir musclemap-weights
+musclemap-mlx segment --input mri.nii.gz --model-dir musclemap-weights --output-dir out-local
+```
+
+MedMLX and this package share the `MEDMLX_ASSET_CACHE` cache setting. Python
+callers supply a local bundle with `weights_path="musclemap-weights"`.
+
+## Qualification and license
+
+The native implementation was compared with the pinned MuscleMap v1.4 reference
+on real CT and MRI examples. Reproduction commands, acceptance criteria and
+measurements are in
+[docs/technical.md](https://github.com/MedMLX/musclemap-mlx/blob/dev/docs/technical.md).
+These two compact examples establish implementation parity; anatomical accuracy
+against manual segmentations and coverage of all 113 structures were not measured.
+
+Code and publisher weights are MIT licensed. See
+[UPSTREAM_LICENSE](https://github.com/MedMLX/musclemap-mlx/blob/dev/UPSTREAM_LICENSE)
+and [THIRD_PARTY_NOTICES.md](https://github.com/MedMLX/musclemap-mlx/blob/dev/THIRD_PARTY_NOTICES.md).
+The original checkpoint and configuration are published in
+[Zenodo record 21929873](https://zenodo.org/records/21929873).

musclemap_mlx-0.1.0/README.md

@@ -0,0 +1,120 @@
+# musclemap-mlx
+
+Segment 113 muscle and bone structures in MRI or CT NIfTI images with
+MuscleMap v1.4, running natively through MLX and Metal on Apple Silicon.
+Weights download automatically on first use.
+
+## Install and segment
+
+Requires an Apple Silicon Mac, Metal and Python 3.13. With
+[uv](https://docs.astral.sh/uv/):
+
+```sh
+uv venv --python 3.13
+source .venv/bin/activate
+uv pip install musclemap-mlx
+
+# Real MRI example from TotalSegmentator's official test data.
+curl -L https://raw.githubusercontent.com/wasserth/TotalSegmentator/d89e47f73d24f5fba732a715aeb6fdc0ae34526b/tests/reference_files/example_mr_sm.nii.gz -o mri.nii.gz
+musclemap-mlx segment --input mri.nii.gz --output-dir out
+```
+
+The command writes `out/mask.nii.gz` and prints JSON containing the output path,
+label table, device, memory use and chunking settings. To use the corresponding
+CT example, replace `example_mr_sm.nii.gz` with `example_ct_sm.nii.gz`.
+
+### Python
+
+```python
+from musclemap_mlx.medmlx import run
+
+result = run(image_path="mri.nii.gz", output_dir="out-python")
+print(result["outputs"]["mask"])
+```
+
+### MedMLX
+
+```sh
+uv pip install 'medmlx[segment]>=0.9.1,<0.10'
+medmlx describe musclemap
+medmlx setup musclemap
+medmlx run musclemap --arg image_path=mri.nii.gz output_dir=out-medmlx
+```
+
+Or use its Python interface:
+
+```python
+import medmlx
+
+medmlx.setup_model("musclemap")
+job = medmlx.run("musclemap", arguments={
+    "image_path": "mri.nii.gz",
+    "output_dir": "out-medmlx-python",
+})
+print(job["result"])
+```
+
+The MCP server exposes the same model through `list_models`, `describe_model`,
+`model_readiness`, `prepare_model` and `run_model`; poll `get_job` for setup and
+segmentation results. Start it with `medmlx serve --root .`.
+
+## Inputs and outputs
+
+Input is one finite 3D `.nii` or `.nii.gz` MRI or CT volume. Millimetre spatial
+units are supported; an unspecified unit is interpreted as millimetres, matching
+upstream. Other declared spatial units and 4D images are rejected.
+
+The output is an `int16` mask with background 0 and the original 113 anatomical
+IDs, including 1101 for left levator scapulae and 8222 for right plantaris. Its
+shape, affine, qform, sform and spatial-unit metadata follow the input. The
+largest 6-connected component of each label is retained after chunk assembly.
+The JSON result lists every anatomical ID and name.
+
+| Option | Meaning | Default |
+| --- | --- | --- |
+| `--input` | 3D NIfTI image | Required |
+| `--output-dir` | Destination for `mask.nii.gz` | Required |
+| `--model-dir` | Local folder containing `model.safetensors` | Automatic download/cache |
+| `--chunk-size` | Source storage-axis slices per independent pass | Eight |
+| `--overlap` | Gaussian window overlap, integer 0–99 percent | 90 |
+| `--overwrite` | Replace an existing output | Off |
+
+Each chunk is independently normalized. For integer source images, chunking
+also reproduces upstream's source-header NIfTI round trip before normalization.
+Changing chunk size can change the mask. A chunk size at least the source depth
+uses one whole-volume pass. Inference uses MLX on Metal; there is no CPU or
+PyTorch network fallback.
+
+## Weights and offline use
+
+The approximately 105 MB float32 bundle is hosted at
+[MedMLX/musclemap-mlx](https://huggingface.co/MedMLX/musclemap-mlx/tree/a798b3073e3b2c449987c13aa193626aaacc3092),
+pinned to revision `a798b3073e3b2c449987c13aa193626aaacc3092`. It includes the
+publisher configuration, original label metadata and MIT notices.
+
+```sh
+# Pre-download to the shared cache; subsequent runs can work offline.
+musclemap-mlx download
+
+# Export a local bundle, then supply it explicitly.
+musclemap-mlx download --model-dir musclemap-weights
+musclemap-mlx segment --input mri.nii.gz --model-dir musclemap-weights --output-dir out-local
+```
+
+MedMLX and this package share the `MEDMLX_ASSET_CACHE` cache setting. Python
+callers supply a local bundle with `weights_path="musclemap-weights"`.
+
+## Qualification and license
+
+The native implementation was compared with the pinned MuscleMap v1.4 reference
+on real CT and MRI examples. Reproduction commands, acceptance criteria and
+measurements are in
+[docs/technical.md](https://github.com/MedMLX/musclemap-mlx/blob/dev/docs/technical.md).
+These two compact examples establish implementation parity; anatomical accuracy
+against manual segmentations and coverage of all 113 structures were not measured.
+
+Code and publisher weights are MIT licensed. See
+[UPSTREAM_LICENSE](https://github.com/MedMLX/musclemap-mlx/blob/dev/UPSTREAM_LICENSE)
+and [THIRD_PARTY_NOTICES.md](https://github.com/MedMLX/musclemap-mlx/blob/dev/THIRD_PARTY_NOTICES.md).
+The original checkpoint and configuration are published in
+[Zenodo record 21929873](https://zenodo.org/records/21929873).

musclemap_mlx-0.1.0/THIRD_PARTY_NOTICES.md

@@ -0,0 +1,35 @@
+# Third-party notices
+
+## MuscleMap
+
+- Source: https://github.com/MuscleMap/MuscleMap
+- Revision: `35a2138a4a3a63be713bae5584632ffc6433b960`
+- Copyright: Copyright (c) 2023 MuscleMap
+- License: MIT; original text retained in [UPSTREAM_LICENSE](UPSTREAM_LICENSE).
+  The package [LICENSE](LICENSE) retains the same MIT terms.
+- Adapted: native MLX implementation of the published 2D MONAI UNet and its
+  inference preprocessing/postprocessing. The original anatomical label IDs,
+  configuration semantics and largest-connected-component helper are preserved.
+  Upstream source remains unchanged and read-only in the development-only
+  `reference/` submodule.
+
+## Dependencies, weights and data
+
+Dependencies declared in [pyproject.toml](pyproject.toml) are installed
+separately and keep their own licenses. MLX is the runtime backend. MONAI 1.3.2
+and PyTorch 2.4.1 are used only in the development reference environment; the
+converter uses PyTorch's restricted `weights_only=True` deserializer and the
+runtime never imports Torch.
+
+Zenodo record [21929873](https://zenodo.org/records/21929873) supplies the
+contrast-agnostic wholebody v1.4 checkpoint, original JSON configuration and
+publisher `LICENSE` under MIT. The license grant matches
+[UPSTREAM_LICENSE](UPSTREAM_LICENSE). Converted weights retain all original
+float32 parameters and are hosted at
+[MedMLX/musclemap-mlx](https://huggingface.co/MedMLX/musclemap-mlx/tree/a798b3073e3b2c449987c13aa193626aaacc3092),
+with configuration, labels, license and attribution. No weights or medical images
+are included in package distributions.
+
+## Showcase image
+
+No showcase image or patient data is included.

musclemap_mlx-0.1.0/UPSTREAM_LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 MuscleMap
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

musclemap_mlx-0.1.0/pyproject.toml

@@ -0,0 +1,116 @@
+[build-system]
+requires = ["hatchling>=1.27"]
+build-backend = "hatchling.build"
+
+[project]
+name = "musclemap-mlx"
+version = "0.1.0"
+description = "Segment 113 whole-body muscle and bone structures with MuscleMap v1.4 and native MLX"
+readme = "README.md"
+requires-python = ">=3.13,<3.14"
+license = "MIT"
+license-files = ["LICENSE", "UPSTREAM_LICENSE", "THIRD_PARTY_NOTICES.md"]
+authors = [
+  { name = "Joseph Sandoval" },
+  { name = "Alfredo Sandoval Jr." },
+]
+keywords = [
+  "mlx",
+  "apple-silicon",
+  "metal",
+  "medical-imaging",
+  "musclemap",
+  "whole-body",
+  "segmentation",
+  "mri",
+  "ct",
+]
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "Intended Audience :: Developers",
+  "Intended Audience :: Science/Research",
+  "Operating System :: MacOS :: MacOS X",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3 :: Only",
+  "Programming Language :: Python :: 3.13",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+  "Topic :: Scientific/Engineering :: Image Processing",
+  "Topic :: Scientific/Engineering :: Medical Science Apps.",
+  "Typing :: Typed",
+]
+dependencies = [
+  "medmlx-core>=0.2.7,<0.3",
+  "mlx==0.32.4; platform_system == 'Darwin' and platform_machine == 'arm64'",
+  "nibabel>=5.3,<6",
+  "numpy>=2.0,<3",
+  "scipy>=1.14,<2",
+]
+
+[project.urls]
+Homepage = "https://medmlx.org"
+Repository = "https://github.com/MedMLX/musclemap-mlx"
+Issues = "https://github.com/MedMLX/musclemap-mlx/issues"
+Weights = "https://huggingface.co/MedMLX/musclemap-mlx"
+Upstream = "https://github.com/MuscleMap/MuscleMap"
+MedMLX = "https://github.com/MedMLX/medmlx"
+[project.scripts]
+"musclemap-mlx" = "musclemap_mlx.cli:main"
+
+[project.entry-points."medmlx.models"]
+musclemap = "musclemap_mlx.medmlx:MODEL"
+[dependency-groups]
+dev = [
+  "monai==1.6.1",
+  "pyright==1.1.414",
+  "pytest>=8.4,<10",
+  "ruff>=0.13,<1",
+  "safetensors>=0.4,<1",
+  "scipy-stubs>=1.18.1,<1.19",
+  "torch==2.14.1",
+]
+
+[tool.uv]
+environments = ["sys_platform == 'darwin' and platform_machine == 'arm64'"]
+required-environments = ["sys_platform == 'darwin' and platform_machine == 'arm64'"]
+constraint-dependencies = ["setuptools>=84.1.0", "filelock>=4.2.1"]
+[tool.hatch.build.targets.wheel]
+packages = ["src/musclemap_mlx"]
+[tool.hatch.build.targets.sdist]
+only-include = [
+  "src/musclemap_mlx",
+  "pyproject.toml",
+  "README.md",
+  "LICENSE",
+  "THIRD_PARTY_NOTICES.md",
+  "UPSTREAM_LICENSE",
+]
+
+[tool.ruff]
+line-length = 100
+target-version = "py313"
+extend-exclude = ["reference"]
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "UP", "B", "SIM", "RUF", "TID251", "TID252", "ANN401"]
+
+[tool.ruff.lint.flake8-tidy-imports]
+ban-relative-imports = "all"
+
+[tool.ruff.lint.flake8-tidy-imports.banned-api]
+"typing.Any" = { msg = "Use a precise type and parse untyped values at the boundary." }
+"typing_extensions.Any" = { msg = "Use a precise type and parse untyped values at the boundary." }
+[tool.pyright]
+include = ["src", "tests"]
+extraPaths = ["src"]
+pythonVersion = "3.13"
+typeCheckingMode = "strict"
+enableTypeIgnoreComments = false
+reportUnnecessaryTypeIgnoreComment = "error"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "--strict-config --strict-markers"
+markers = [
+  "metal: needs the MLX Metal GPU and the shared Perfkit GPU lock",
+  "oracle: compares against a recorded MONAI 1.3.2 reference; opt-in",
+]

musclemap_mlx-0.1.0/src/musclemap_mlx/__init__.py

@@ -0,0 +1,11 @@
+"""MuscleMap v1.4 whole-body segmentation with native MLX on Apple Silicon."""
+
+from __future__ import annotations
+
+from importlib.metadata import version
+
+
+def __getattr__(name: str) -> str:
+    if name == "__version__":
+        return version("musclemap-mlx")
+    raise AttributeError(name)

musclemap_mlx-0.1.0/src/musclemap_mlx/assets.py

@@ -0,0 +1,45 @@
+"""Check publisher asset completeness without deserializing checkpoint data."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import cast
+
+from musclemap_mlx.config import CONFIG_NAME, MODEL_ID
+
+ASSETS = {
+    "contrast_agnostic_wholebody_model.pth": 104884148,
+    "contrast_agnostic_wholebody_model.json": 10555,
+    "LICENSE": 1087,
+}
+
+
+def verify_assets(folder: Path) -> dict[str, str]:
+    """Check publisher file sizes and return each asset's absolute path.
+
+    Raises:
+        ValueError: An asset is missing or incomplete for Zenodo record 21929873.
+    """
+    paths: dict[str, str] = {}
+    for name, size in ASSETS.items():
+        path = folder / name
+        if not path.is_file():
+            raise ValueError(f"Missing asset: {path}; obtain it from Zenodo record 21929873")
+        if path.stat().st_size != size:
+            raise ValueError(f"Unexpected size for {name}; expected {size} bytes")
+        paths[name] = path.absolute().as_posix()
+    return paths
+
+
+def inspect_assets(folder: Path) -> dict[str, object]:
+    """Check asset completeness and return the original configuration as JSON data."""
+    verify_assets(folder)
+    config = cast(object, json.loads((folder / CONFIG_NAME).read_text(encoding="utf-8")))
+    return {
+        "status": "assets-complete",
+        "model_id": MODEL_ID,
+        "record": "https://zenodo.org/records/21929873",
+        "configuration": config,
+        "checkpoint_deserialized": False,
+    }