[ { "path": ".autorc", "content": "{\n \"onlyPublishWithReleaseLabel\": false,\n \"baseBranch\": \"master\",\n \"prereleaseBranches\": [\"alpha\"],\n \"author\": \"Nobrainer Bot \",\n \"noVersionPrefix\": true,\n \"plugins\": [\"git-tag\"]\n}\n" }, { "path": ".dockerignore", "content": ".git/\ndocker/\n.idea/\n" }, { "path": ".flake8", "content": "[flake8]\nmax-line-length = 100\nexclude =\n .git/\n __pycache__/\n build/\n dist/\n _version.py\n versioneer.py\nignore =\n E203\n W503\n" }, { "path": ".gitattributes", "content": "nobrainer/_version.py export-subst\n" }, { "path": ".github/EC2_GPU_RUNNER.md", "content": "# EC2 GPU Runner Setup\n\nThis document describes how to configure the AWS EC2 instance used as a\nself-hosted GitHub Actions runner for GPU integration tests.\n\nThe workflow (`guide-notebooks-ec2.yml`) uses\n[machulav/ec2-github-runner](https://github.com/machulav/ec2-github-runner) to\nstart an ephemeral EC2 instance, run GPU tests, and terminate the instance\nautomatically.\n\n## AMI preparation\n\nStart from the **AWS Deep Learning Base AMI (Amazon Linux 2023)** or any\nAmazon Linux 2023 AMI with NVIDIA drivers and CUDA pre-installed. The AMI must\nbe in the same region as the `AWS_REGION` variable configured in GitHub.\n\n### 1. Launch an instance to build the AMI\n\n```bash\naws ec2 run-instances \\\n --image-id ami-XXXXXXXX \\\n --instance-type g4dn.xlarge \\\n --key-name your-key-pair \\\n --security-group-ids sg-XXXXXXXX \\\n --subnet-id subnet-XXXXXXXX \\\n --block-device-mappings '[{\"DeviceName\":\"/dev/xvda\",\"Ebs\":{\"VolumeSize\":100}}]' \\\n --tag-specifications 'ResourceType=instance,Tags=[{Key=Name,Value=nobrainer-ami-builder}]'\n```\n\n### 2. SSH in as ec2-user and configure\n\n```bash\nssh -i your-key.pem ec2-user@\n```\n\nAll commands below run as `ec2-user`.\n\n#### Install system dependencies\n\n```bash\nsudo dnf install -y jq git\n```\n\n#### Install uv\n\n```bash\ncurl -LsSf https://astral.sh/uv/install.sh | sh\nsource $HOME/.local/bin/env\n```\n\n#### Create the pre-installed nobrainer venv\n\nThe CI workflow expects a venv at `~/nobrainer-env` with heavy dependencies\n(torch, monai, pyro-ppl) already installed. This avoids re-downloading ~2 GB\nof packages on every CI run.\n\n```bash\nuv venv --python 3.14 ~/nobrainer-env\n\n# Install the heavy GPU dependencies into the base venv\nuv pip install \\\n torch \\\n monai \\\n pyro-ppl \\\n pytorch-lightning \\\n pytest\n```\n\n#### Verify GPU access\n\n```bash\nsource ~/nobrainer-env/bin/activate\npython -c \"\nimport torch\nassert torch.cuda.is_available(), 'CUDA not available'\nprint(f'GPU: {torch.cuda.get_device_name(0)}')\nprint(f'CUDA: {torch.version.cuda}')\nprint(f'PyTorch: {torch.__version__}')\n\"\ndeactivate\n```\n\n### 3. Create the AMI\n\nStop the instance (or use `--no-reboot`), then:\n\n```bash\naws ec2 create-image \\\n --instance-id i-XXXXXXXXX \\\n --name \"nobrainer-pytorch-gpu-$(date +%Y%m%d)\" \\\n --description \"Amazon Linux 2023 + CUDA + PyTorch + uv for nobrainer GPU CI\" \\\n --no-reboot\n```\n\nNote the resulting AMI ID — this goes into the `AWS_IMAGE_ID` variable.\n\n### 4. Terminate the builder instance\n\n```bash\naws ec2 terminate-instances --instance-id i-XXXXXXXXX\n```\n\n## GitHub configuration\n\n### Secrets (Settings → Secrets → Actions)\n\n| Name | Description |\n|------|-------------|\n| `AWS_KEY_ID` | IAM access key with EC2 RunInstances/TerminateInstances/DescribeInstances permissions |\n| `AWS_KEY_SECRET` | Corresponding secret access key |\n| `GH_TOKEN` | GitHub PAT with `repo` scope (used by machulav/ec2-github-runner to register the runner) |\n\n### Variables (Settings → Variables → Actions)\n\n| Name | Example | Description |\n|------|---------|-------------|\n| `AWS_REGION` | `us-east-1` | Region where the AMI lives |\n| `AWS_IMAGE_ID` | `ami-0abc123def456` | The AMI created above |\n| `AWS_INSTANCE_TYPE` | `g4dn.xlarge` | 1x T4 GPU (~$0.53/hr); `p3.2xlarge` for V100 |\n| `AWS_SUBNET` | `subnet-0abc123` | Must have internet access for runner registration |\n| `AWS_SECURITY_GROUP` | `sg-0abc123` | Allow outbound HTTPS (port 443) |\n\n## IAM policy (minimum permissions)\n\n```json\n{\n \"Version\": \"2012-10-17\",\n \"Statement\": [\n {\n \"Effect\": \"Allow\",\n \"Action\": [\n \"ec2:RunInstances\",\n \"ec2:TerminateInstances\",\n \"ec2:DescribeInstances\",\n \"ec2:DescribeInstanceStatus\",\n \"ec2:CreateTags\"\n ],\n \"Resource\": \"*\"\n },\n {\n \"Effect\": \"Allow\",\n \"Action\": \"iam:PassRole\",\n \"Resource\": \"*\"\n }\n ]\n}\n```\n\n## Updating the base venv\n\nWhen upgrading PyTorch or other dependencies, SSH into a running instance (or\nlaunch the AMI), update `~/nobrainer-env`, and create a new AMI snapshot:\n\n```bash\nssh -i your-key.pem ec2-user@\ncd ~/nobrainer\nuv pip install --upgrade torch monai pyro-ppl pytorch-lightning\n# Then create a new AMI and update AWS_IMAGE_ID in GitHub variables\n```\n" }, { "path": ".github/ISSUE_TEMPLATE/bug_report.md", "content": "---\nname: Bug report\nabout: Report a bug (e.g., something not working as described, missing/incorrect documentation).\ntitle: ''\nlabels: 'bug'\nassignees: ''\n\n---\n\n" }, { "path": ".github/ISSUE_TEMPLATE/documentation.md", "content": "---\nname: Documentation improvement\nabout: Request improvements to the documentation and tutorials.\ntitle: ''\nlabels: 'documentation'\nassignees: ''\n\n---\n\n" }, { "path": ".github/ISSUE_TEMPLATE/feature_request.md", "content": "---\nname: Feature request\nabout: Propose a new feature or a change to an existing feature.\ntitle: ''\nlabels: 'feature'\nassignees: ''\n\n---\n\n" }, { "path": ".github/ISSUE_TEMPLATE/maintenance.md", "content": "---\nname: Maintenance and delivery\nabout: Suggestions and requests regarding the infrastructure for development, testing, and delivery.\ntitle: ''\nlabels: 'maintenance'\nassignees: ''\n\n---\n" }, { "path": ".github/ISSUE_TEMPLATE/question.md", "content": "---\nname: Question\nabout: Not sure if you are using Nobrainer correctly, or other questions? This is the place.\ntitle: ''\nlabels: 'question'\nassignees: ''\n\n---\n\n" }, { "path": ".github/PULL_REQUEST_TEMPLATE.md", "content": "## Types of changes\n\n- [ ] Bug fix (non-breaking change which fixes an issue)\n- [ ] New feature (non-breaking change which adds functionality)\n- [ ] Breaking change (fix or feature that would cause existing functionality to change)\n\n## Summary\n\n\n## Checklist\n\n- [ ] I have added tests to cover my changes\n- [ ] I have updated documentation (if necessary)\n\n## Acknowledgment\n- [ ] I acknowledge that this contribution will be available under the Apache 2 license.\n" }, { "path": ".github/workflows/ci.yml", "content": "name: CI\n\non:\n push:\n branches: [main, master]\n pull_request:\n branches: [main, master, alpha]\n\njobs:\n unit-tests:\n runs-on: ${{ matrix.os }}\n strategy:\n fail-fast: false\n matrix:\n os: [ubuntu-latest]\n python-version: [\"3.12\", \"3.13\", \"3.14\"]\n\n steps:\n - uses: actions/checkout@v4\n with:\n fetch-depth: 0\n fetch-tags: true\n\n - name: Install uv\n uses: astral-sh/setup-uv@v4\n\n - name: Cache sample brain data\n uses: actions/cache@v4\n with:\n path: /tmp/nobrainer-data\n key: nobrainer-sample-data-v1\n\n - name: Set up Python ${{ matrix.python-version }}\n run: uv venv --python ${{ matrix.python-version }}\n\n - name: Install dependencies\n run: |\n uv pip install \\\n \".[bayesian,generative,zarr,dev]\" \\\n monai \\\n pyro-ppl\n\n - name: Test with pytest (CPU, skip GPU)\n run: |\n uv run pytest nobrainer/tests/unit/ -v \\\n -m \"not gpu\" \\\n --no-header\n\n - name: Run sr-tests (somewhat realistic tests)\n run: |\n uv run pytest nobrainer/sr-tests/ -v \\\n -m \"not gpu\" \\\n --no-header \\\n --tb=short\n\n - name: Research loop smoke test (5 min budget, no API key)\n run: |\n mkdir -p /tmp/research-smoke\n cp nobrainer/research/templates/train_bayesian_vnet.py /tmp/research-smoke/train.py\n cp nobrainer/research/templates/prepare.py /tmp/research-smoke/prepare.py\n uv run nobrainer research run \\\n --working-dir /tmp/research-smoke \\\n --model-family meshnet \\\n --max-experiments 2 \\\n --budget-minutes 5 || true\n test -f /tmp/research-smoke/run_summary.md && echo \"run_summary.md exists\" || echo \"WARN: no run_summary.md\"\n\n image-build:\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v4\n with:\n fetch-depth: 0\n fetch-tags: true\n\n - name: Test CPU Docker image build\n run: |\n docker build -t neuronets/nobrainer:ci-cpu -f docker/cpu.Dockerfile .\n" }, { "path": ".github/workflows/guide-notebooks-ec2.yml", "content": "name: GPU Tests - EC2\nrun-name: ${{ github.ref_name }} - GPU Tests - EC2\non:\n push:\n branches: [main, master]\n # PRs require approval label from a repo admin before this workflow runs.\n # This prevents untrusted PR code from executing on the self-hosted GPU runner.\n pull_request:\n branches: [main, master, alpha]\n types: [labeled, synchronize]\n\njobs:\n start-runner:\n name: Start self-hosted EC2 runner\n runs-on: ubuntu-latest\n # Only run on PRs if an admin added the 'gpu-test-approved' label\n if: >-\n github.event_name == 'push' ||\n (github.event_name == 'pull_request' &&\n contains(github.event.pull_request.labels.*.name, 'gpu-test-approved'))\n outputs:\n label: ${{ steps.start-ec2-runner.outputs.label }}\n ec2-instance-id: ${{ steps.start-ec2-runner.outputs.ec2-instance-id }}\n multi_gpu: ${{ steps.gpu-config.outputs.multi_gpu }}\n steps:\n - name: Parse GPU config labels\n id: gpu-config\n if: github.event_name == 'pull_request'\n run: |\n LABELS='${{ toJSON(github.event.pull_request.labels.*.name) }}'\n\n # Priority: gpu-instance: (exact) > gpu-multi (multi-GPU) > gpu-family: (family default)\n INSTANCE=$(echo \"$LABELS\" | jq -r '[.[] | select(startswith(\"gpu-instance:\"))] | if length > 0 then .[0] | split(\":\")[1] else \"\" end')\n\n # gpu-multi label → pick a multi-GPU instance for DDP/model-parallel tests\n MULTI_GPU=$(echo \"$LABELS\" | jq -r '[.[] | select(. == \"gpu-multi\")] | if length > 0 then \"true\" else \"\" end')\n if [ -z \"$INSTANCE\" ] && [ -n \"$MULTI_GPU\" ]; then\n INSTANCE=\"g5.12xlarge\" # 4x A10G GPUs\n echo \"gpu-multi label → selecting $INSTANCE (4 GPUs)\"\n fi\n\n # gpu-family: label → pick default instance from that family\n # Supported families: g4dn, g5, g6, p3, p4d, p5\n if [ -z \"$INSTANCE\" ]; then\n FAMILY=$(echo \"$LABELS\" | jq -r '[.[] | select(startswith(\"gpu-family:\"))] | if length > 0 then .[0] | split(\":\")[1] else \"\" end')\n if [ -n \"$FAMILY\" ]; then\n case \"$FAMILY\" in\n g4dn) INSTANCE=\"g4dn.xlarge\" ;; # 1x T4\n g5) INSTANCE=\"g5.xlarge\" ;; # 1x A10G\n g6) INSTANCE=\"g6.xlarge\" ;; # 1x L4\n p3) INSTANCE=\"p3.2xlarge\" ;; # 1x V100\n p4d) INSTANCE=\"p4d.24xlarge\" ;; # 8x A100\n p5) INSTANCE=\"p5.48xlarge\" ;; # 8x H100\n *) echo \"Unknown GPU family: $FAMILY\"; INSTANCE=\"\" ;;\n esac\n if [ -n \"$INSTANCE\" ]; then\n echo \"gpu-family:${FAMILY} → selecting $INSTANCE\"\n fi\n fi\n fi\n\n # Default to spot pricing; gpu-ondemand:true overrides to on-demand\n ONDEMAND=$(echo \"$LABELS\" | jq -r '[.[] | select(. == \"gpu-ondemand:true\")] | if length > 0 then \"true\" else \"\" end')\n if [ -n \"$ONDEMAND\" ]; then\n MARKET=\"\"\n else\n MARKET=\"spot\"\n fi\n\n echo \"instance=${INSTANCE}\" >> $GITHUB_OUTPUT\n echo \"market_type=${MARKET}\" >> $GITHUB_OUTPUT\n echo \"multi_gpu=${MULTI_GPU}\" >> $GITHUB_OUTPUT\n echo \"Parsed labels: instance=${INSTANCE:-default}, market=${MARKET:-spot}, multi_gpu=${MULTI_GPU:-false}\"\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@v6\n with:\n aws-access-key-id: ${{ secrets.AWS_KEY_ID }}\n aws-secret-access-key: ${{ secrets.AWS_KEY_SECRET }}\n aws-region: ${{ vars.AWS_REGION }}\n\n - name: Start EC2 runner\n id: start-ec2-runner\n uses: machulav/ec2-github-runner@v2.5.2\n with:\n mode: start\n github-token: ${{ secrets.GH_TOKEN }}\n ec2-image-id: ${{ vars.AWS_IMAGE_ID }}\n ec2-instance-type: ${{ steps.gpu-config.outputs.instance || vars.AWS_INSTANCE_TYPE }}\n subnet-id: ${{ vars.AWS_SUBNET }}\n security-group-id: ${{ vars.AWS_SECURITY_GROUP }}\n market-type: ${{ steps.gpu-config.outputs.market_type || 'spot' }}\n\n gpu-tests:\n needs: start-runner\n runs-on: ${{ needs.start-runner.outputs.label }}\n env:\n # The GitHub Actions runner runs as root on the EC2 instance, but\n # the AMI was set up as ec2-user. Use absolute paths to ec2-user's\n # home directory for the pre-installed venv and uv binary.\n EC2_USER_HOME: /home/ec2-user\n steps:\n - name: Checkout\n uses: actions/checkout@v4\n\n - name: Log GPU runner config\n run: |\n echo \"Instance type: $(curl -s http://169.254.169.254/latest/meta-data/instance-type 2>/dev/null || echo 'unknown')\"\n echo \"Availability zone: $(curl -s http://169.254.169.254/latest/meta-data/placement/availability-zone 2>/dev/null || echo 'unknown')\"\n echo \"Market type: $(curl -s http://169.254.169.254/latest/meta-data/instance-life-cycle 2>/dev/null || echo 'unknown')\"\n\n - name: Set up venv from pre-installed base\n run: |\n set -ex\n BASE_VENV=\"${EC2_USER_HOME}/nobrainer-env\"\n export PATH=\"${EC2_USER_HOME}/.local/bin:$PATH\"\n\n if [ -d \"$BASE_VENV\" ]; then\n echo \"Found pre-installed base venv at $BASE_VENV\"\n # Copy the base venv so the AMI stays clean for next run\n cp -a \"$BASE_VENV\" .venv\n else\n echo \"No base venv found — creating from scratch\"\n uv venv --python 3.14\n fi\n\n # Install nobrainer from checkout on top of the base layer\n uv pip install \\\n \".[bayesian,generative,zarr,dev]\" \\\n monai \\\n pyro-ppl \\\n matplotlib\n\n - name: Verify GPU access\n run: |\n export PATH=\"${EC2_USER_HOME}/.local/bin:$PATH\"\n uv run python -c \"\n import torch\n assert torch.cuda.is_available(), 'CUDA not available'\n n = torch.cuda.device_count()\n print(f'GPUs: {n}')\n for i in range(n):\n print(f' [{i}] {torch.cuda.get_device_name(i)}')\n print(f'CUDA: {torch.version.cuda}')\n print(f'PyTorch: {torch.__version__}')\n \"\n\n - name: Run full test suite (including GPU)\n run: |\n export PATH=\"${EC2_USER_HOME}/.local/bin:$PATH\"\n uv run pytest nobrainer/tests/ nobrainer/sr-tests/ -v \\\n --no-header \\\n --tb=short\n\n - name: Run multi-GPU tests (DDP + model parallel)\n if: needs.start-runner.outputs.multi_gpu == 'true'\n run: |\n export PATH=\"${EC2_USER_HOME}/.local/bin:$PATH\"\n echo \"=== Multi-GPU DDP and model-parallel tests ===\"\n uv run pytest nobrainer/tests/gpu/ -v \\\n --no-header \\\n --tb=short \\\n -k \"multi_gpu or ddp or model_parallel\"\n\n stop-runner:\n name: Stop self-hosted EC2 runner\n needs:\n - start-runner\n - gpu-tests\n runs-on: ubuntu-latest\n if: ${{ always() && needs.start-runner.result == 'success' }}\n steps:\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@v6\n with:\n aws-access-key-id: ${{ secrets.AWS_KEY_ID }}\n aws-secret-access-key: ${{ secrets.AWS_KEY_SECRET }}\n aws-region: ${{ vars.AWS_REGION }}\n - name: Stop EC2 runner\n uses: machulav/ec2-github-runner@v2.5.2\n with:\n mode: stop\n github-token: ${{ secrets.GH_TOKEN }}\n label: ${{ needs.start-runner.outputs.label }}\n ec2-instance-id: ${{ needs.start-runner.outputs.ec2-instance-id }}\n" }, { "path": ".github/workflows/kwyk-reproduction-ec2.yml", "content": "name: KWYK Reproduction - EC2 GPU\nrun-name: ${{ github.ref_name }} - KWYK Reproduction\n\non:\n workflow_dispatch:\n inputs:\n mode:\n description: \"Run mode\"\n required: true\n default: \"smoke-test\"\n type: choice\n options:\n - smoke-test\n - small-train\n - full\n instance_type:\n description: \"EC2 instance type\"\n required: false\n default: \"\"\n type: string\n on_demand:\n description: \"Use on-demand pricing (not spot)\"\n required: false\n default: false\n type: boolean\n pull_request:\n branches: [main, master, alpha]\n types: [labeled, synchronize]\n\njobs:\n start-runner:\n name: Start self-hosted EC2 runner\n runs-on: ubuntu-latest\n if: >-\n github.event_name == 'workflow_dispatch' ||\n (github.event_name == 'pull_request' &&\n contains(github.event.pull_request.labels.*.name, 'kwyk-gpu-test'))\n outputs:\n label: ${{ steps.start-ec2-runner.outputs.label }}\n ec2-instance-id: ${{ steps.start-ec2-runner.outputs.ec2-instance-id }}\n steps:\n - name: Determine instance config\n id: gpu-config\n run: |\n if [ \"${{ github.event_name }}\" = \"workflow_dispatch\" ]; then\n INSTANCE=\"${{ inputs.instance_type }}\"\n if [ \"${{ inputs.on_demand }}\" = \"true\" ]; then\n MARKET=\"\"\n else\n MARKET=\"spot\"\n fi\n else\n # PR: parse labels\n LABELS='${{ toJSON(github.event.pull_request.labels.*.name) }}'\n INSTANCE=$(echo \"$LABELS\" | jq -r '[.[] | select(startswith(\"gpu-instance:\"))] | if length > 0 then .[0] | split(\":\")[1] else \"\" end')\n ONDEMAND=$(echo \"$LABELS\" | jq -r '[.[] | select(. == \"gpu-ondemand:true\")] | if length > 0 then \"true\" else \"\" end')\n if [ -n \"$ONDEMAND\" ]; then MARKET=\"\"; else MARKET=\"spot\"; fi\n fi\n echo \"instance=${INSTANCE}\" >> $GITHUB_OUTPUT\n echo \"market_type=${MARKET}\" >> $GITHUB_OUTPUT\n echo \"Config: instance=${INSTANCE:-default}, market=${MARKET:-spot}\"\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@v6\n with:\n aws-access-key-id: ${{ secrets.AWS_KEY_ID }}\n aws-secret-access-key: ${{ secrets.AWS_KEY_SECRET }}\n aws-region: ${{ vars.AWS_REGION }}\n\n - name: Start EC2 runner\n id: start-ec2-runner\n uses: machulav/ec2-github-runner@v2.5.2\n with:\n mode: start\n github-token: ${{ secrets.GH_TOKEN }}\n ec2-image-id: ${{ vars.AWS_IMAGE_ID }}\n ec2-instance-type: ${{ steps.gpu-config.outputs.instance || vars.AWS_INSTANCE_TYPE }}\n subnet-id: ${{ vars.AWS_SUBNET }}\n security-group-id: ${{ vars.AWS_SECURITY_GROUP }}\n market-type: ${{ steps.gpu-config.outputs.market_type || 'spot' }}\n\n kwyk-reproduction:\n needs: start-runner\n runs-on: ${{ needs.start-runner.outputs.label }}\n timeout-minutes: 120\n env:\n EC2_USER_HOME: /home/ec2-user\n steps:\n - name: Checkout\n uses: actions/checkout@v4\n with:\n fetch-depth: 0\n fetch-tags: true\n\n - name: Log GPU runner config\n run: |\n echo \"Instance type: $(curl -s http://169.254.169.254/latest/meta-data/instance-type 2>/dev/null || echo 'unknown')\"\n echo \"Market type: $(curl -s http://169.254.169.254/latest/meta-data/instance-life-cycle 2>/dev/null || echo 'unknown')\"\n\n - name: Install git-annex\n run: |\n # Runner runs as root but EC2_USER_HOME points to ec2-user.\n # Add both possible bin dirs to PATH.\n export PATH=\"/root/.local/bin:${EC2_USER_HOME}/.local/bin:$PATH\"\n if ! command -v git-annex &>/dev/null; then\n echo \"Installing git-annex via uv...\"\n uv tool install git-annex\n fi\n git-annex version\n # Persist PATH for subsequent steps\n echo \"/root/.local/bin\" >> $GITHUB_PATH\n\n - name: Set up venv from pre-installed base\n run: |\n set -ex\n BASE_VENV=\"${EC2_USER_HOME}/nobrainer-env\"\n export PATH=\"${EC2_USER_HOME}/.local/bin:$PATH\"\n\n if [ -d \"$BASE_VENV\" ]; then\n echo \"Found pre-installed base venv at $BASE_VENV\"\n cp -a \"$BASE_VENV\" .venv\n else\n echo \"No base venv found — creating from scratch\"\n uv venv --python 3.14\n fi\n\n uv pip install \\\n \".[bayesian,generative,zarr,versioning,dev]\" \\\n monai pyro-ppl datalad matplotlib pyyaml scipy nibabel\n\n - name: Verify GPU access\n run: |\n export PATH=\"${EC2_USER_HOME}/.local/bin:$PATH\"\n uv run python -c \"\n import torch\n assert torch.cuda.is_available(), 'CUDA not available'\n print(f'GPU: {torch.cuda.get_device_name(0)}')\n print(f'CUDA: {torch.version.cuda}')\n print(f'PyTorch: {torch.__version__}')\n print(f'Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB')\n \"\n\n - name: Cache sample brain data\n uses: actions/cache@v4\n with:\n path: /tmp/nobrainer-data\n key: nobrainer-sample-data-v1\n\n - name: Run kwyk sr-tests smoke test\n run: |\n export PATH=\"${EC2_USER_HOME}/.local/bin:$PATH\"\n uv run pytest nobrainer/sr-tests/test_kwyk_smoke.py -v \\\n --no-header \\\n --tb=short\n\n - name: Determine run mode\n id: mode\n run: |\n if [ \"${{ github.event_name }}\" = \"workflow_dispatch\" ]; then\n echo \"mode=${{ inputs.mode }}\" >> $GITHUB_OUTPUT\n else\n echo \"mode=smoke-test\" >> $GITHUB_OUTPUT\n fi\n\n - name: Run kwyk reproduction scripts (smoke test)\n if: steps.mode.outputs.mode == 'smoke-test'\n run: |\n set -ex\n export PATH=\"${EC2_USER_HOME}/.local/bin:$PATH\"\n cd scripts/kwyk_reproduction\n\n # Smoke test: skip DataLad (requires git-annex on AMI).\n # The sr-tests already validated the pipeline with get_data().\n # Here we verify the reproduction scripts parse configs correctly\n # and that the training loop works end-to-end with sample data.\n # Try DataLad first; fall back to sample brain data if it fails\n uv run python 01_assemble_dataset.py \\\n --datasets ds000114 \\\n --output-csv manifest.csv \\\n --output-dir data \\\n --label-mapping binary \\\n || {\n echo \"DataLad assembly failed, falling back to sample brain data\"\n uv run python -c \"\n import csv; from nobrainer.utils import get_data\n src = get_data(); pairs = []\n with open(src) as f:\n r = csv.reader(f); next(r)\n pairs = list(r)[:5]\n splits = ['train','train','train','val','test']\n with open('manifest.csv', 'w', newline='') as f:\n w = csv.DictWriter(f, ['t1w_path','label_path','split']); w.writeheader()\n for i,(t1,lbl) in enumerate(pairs):\n w.writerow(dict(t1w_path=t1, label_path=lbl, split=splits[i]))\n print('Manifest created with', len(pairs), 'volumes')\n \"\n }\n\n echo \"=== Step 2: Train deterministic MeshNet (2 epochs) ===\"\n uv run python 02_train_meshnet.py \\\n --manifest manifest.csv \\\n --config config.yaml \\\n --output-dir checkpoints/meshnet \\\n --epochs 2\n\n echo \"=== Step 3a: MC dropout variant ===\"\n uv run python 03_train_bayesian.py \\\n --manifest manifest.csv \\\n --config config.yaml \\\n --variant bwn_multi \\\n --warmstart checkpoints/meshnet \\\n --output-dir checkpoints/bwn_multi \\\n --epochs 2\n\n echo \"=== Step 3b: Spike-and-slab variant (2 epochs) ===\"\n uv run python 03_train_bayesian.py \\\n --manifest manifest.csv \\\n --config config.yaml \\\n --variant bvwn_multi_prior \\\n --warmstart checkpoints/meshnet \\\n --output-dir checkpoints/bvwn_multi_prior \\\n --epochs 2\n\n echo \"=== Checking outputs ===\"\n ls -la checkpoints/meshnet/ checkpoints/bwn_multi/ checkpoints/bvwn_multi_prior/ 2>/dev/null || true\n\n - name: Run kwyk reproduction scripts (small training)\n if: steps.mode.outputs.mode == 'small-train'\n run: |\n set -ex\n export PATH=\"${EC2_USER_HOME}/.local/bin:$PATH\"\n cd scripts/kwyk_reproduction\n\n # Try DataLad first; fall back to sample brain data if it fails\n uv run python 01_assemble_dataset.py \\\n --datasets ds000114 \\\n --output-csv manifest.csv \\\n --output-dir data \\\n --label-mapping binary \\\n || {\n echo \"DataLad assembly failed, falling back to sample brain data\"\n uv run python -c \"\n import csv; from nobrainer.utils import get_data\n src = get_data(); pairs = []\n with open(src) as f:\n r = csv.reader(f); next(r)\n pairs = list(r)[:5]\n splits = ['train','train','train','val','test']\n with open('manifest.csv', 'w', newline='') as f:\n w = csv.DictWriter(f, ['t1w_path','label_path','split']); w.writeheader()\n for i,(t1,lbl) in enumerate(pairs):\n w.writerow(dict(t1w_path=t1, label_path=lbl, split=splits[i]))\n print('Manifest created with', len(pairs), 'volumes')\n \"\n }\n\n echo \"=== Step 2: Train deterministic MeshNet (20 epochs) ===\"\n uv run python 02_train_meshnet.py \\\n --manifest manifest.csv \\\n --config config.yaml \\\n --output-dir checkpoints/meshnet \\\n --epochs 20\n\n echo \"=== Step 3a: MC dropout variant ===\"\n uv run python 03_train_bayesian.py \\\n --manifest manifest.csv \\\n --config config.yaml \\\n --variant bwn_multi \\\n --warmstart checkpoints/meshnet \\\n --output-dir checkpoints/bwn_multi \\\n --epochs 20\n\n echo \"=== Step 3b: Spike-and-slab variant (20 epochs) ===\"\n uv run python 03_train_bayesian.py \\\n --manifest manifest.csv \\\n --config config.yaml \\\n --variant bvwn_multi_prior \\\n --warmstart checkpoints/meshnet \\\n --output-dir checkpoints/bvwn_multi_prior \\\n --epochs 20\n\n echo \"=== Step 3c: Gaussian Bayesian variant (20 epochs) ===\"\n uv run python 03_train_bayesian.py \\\n --manifest manifest.csv \\\n --config config.yaml \\\n --variant bayesian_gaussian \\\n --warmstart checkpoints/meshnet \\\n --output-dir checkpoints/bayesian_gaussian \\\n --epochs 20\n\n echo \"=== Checking outputs ===\"\n ls -la checkpoints/*/ 2>/dev/null || true\n\n - name: Run kwyk reproduction scripts (full)\n if: steps.mode.outputs.mode == 'full'\n timeout-minutes: 1440\n run: |\n set -ex\n export PATH=\"${EC2_USER_HOME}/.local/bin:$PATH\"\n cd scripts/kwyk_reproduction\n\n uv run python 01_assemble_dataset.py \\\n --datasets ds000114 ds000228 ds002609 ds001021 ds002105 \\\n --output-csv manifest.csv \\\n --output-dir data \\\n --label-mapping binary --conform\n\n uv run python 02_train_meshnet.py \\\n --manifest manifest.csv \\\n --config config.yaml \\\n --output-dir checkpoints/meshnet \\\n --epochs 50\n\n for variant in bwn_multi bvwn_multi_prior bayesian_gaussian; do\n echo \"=== Training $variant (50 epochs) ===\"\n uv run python 03_train_bayesian.py \\\n --manifest manifest.csv \\\n --config config.yaml \\\n --variant $variant \\\n --warmstart checkpoints/meshnet \\\n --output-dir checkpoints/$variant \\\n --epochs 50\n done\n\n for variant in meshnet bwn_multi bvwn_multi_prior bayesian_gaussian; do\n echo \"=== Evaluating $variant ===\"\n uv run python 04_evaluate.py \\\n --model checkpoints/$variant/model.pth \\\n --manifest manifest.csv \\\n --split test \\\n --n-samples 10 \\\n --output-dir results/$variant\n done\n\n uv run python 05_compare_kwyk.py \\\n --new-model checkpoints/bvwn_multi_prior/model.pth \\\n --kwyk-dir ../../kwyk \\\n --manifest manifest.csv \\\n --output-dir results/comparison || echo \"WARN: kwyk comparison failed (container may not be available)\"\n\n uv run python 06_block_size_sweep.py \\\n --manifest manifest.csv \\\n --block-sizes 32 64 128 \\\n --output-dir results/sweep\n\n - name: Upload artifacts\n if: always()\n uses: actions/upload-artifact@v4\n with:\n name: kwyk-reproduction-${{ steps.mode.outputs.mode }}\n path: |\n scripts/kwyk_reproduction/figures/\n scripts/kwyk_reproduction/results/\n scripts/kwyk_reproduction/checkpoints/*/croissant.json\n retention-days: 30\n if-no-files-found: warn\n\n stop-runner:\n name: Stop self-hosted EC2 runner\n needs:\n - start-runner\n - kwyk-reproduction\n runs-on: ubuntu-latest\n if: ${{ always() && needs.start-runner.result == 'success' }}\n steps:\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@v6\n with:\n aws-access-key-id: ${{ secrets.AWS_KEY_ID }}\n aws-secret-access-key: ${{ secrets.AWS_KEY_SECRET }}\n aws-region: ${{ vars.AWS_REGION }}\n - name: Stop EC2 runner\n uses: machulav/ec2-github-runner@v2.5.2\n with:\n mode: stop\n github-token: ${{ secrets.GH_TOKEN }}\n label: ${{ needs.start-runner.outputs.label }}\n ec2-instance-id: ${{ needs.start-runner.outputs.ec2-instance-id }}\n" }, { "path": ".github/workflows/publish.yml", "content": "name: Publish to PyPI on GitHub release\n\non:\n release:\n types: [published]\n\njobs:\n pypi-release:\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v4\n\n - name: Install uv\n uses: astral-sh/setup-uv@v4\n\n - name: Build and publish\n run: |\n uv build\n uv publish\n env:\n UV_PUBLISH_TOKEN: ${{ secrets.PYPI_TOKEN }}\n" }, { "path": ".github/workflows/release.yml", "content": "name: Auto-release on PR merge\n\non:\n pull_request:\n branches: [master, alpha]\n types: [closed]\n\nenv:\n AUTO_VERSION: v11.2.1\n\njobs:\n auto-release:\n name: Create release\n runs-on: ubuntu-latest\n # Stable release: merged PR to master with 'release' label\n # Alpha pre-release: merged PR to alpha (book validation runs as\n # a separate check via validate-book.yml on every PR push)\n if: >-\n github.event.pull_request.merged == true &&\n (\n contains(github.event.pull_request.labels.*.name, 'release') ||\n github.event.pull_request.base.ref == 'alpha'\n )\n steps:\n - uses: actions/checkout@v4\n with:\n fetch-depth: 0\n fetch-tags: true\n\n - name: Unset header\n run: git config --local --unset http.https://github.com/.extraheader\n\n - name: Download auto\n run: |\n auto_download_url=\"$(curl -fsSL https://api.github.com/repos/intuit/auto/releases/tags/$AUTO_VERSION | jq -r '.assets[] | select(.name == \"auto-linux.gz\") | .browser_download_url')\"\n wget -O- \"$auto_download_url\" | gunzip > ~/auto\n chmod a+x ~/auto\n\n - name: Create release\n run: ~/auto shipit -vv\n env:\n GH_TOKEN: ${{ secrets.AUTO_USER_TOKEN }}\n" }, { "path": ".github/workflows/validate-book.yml", "content": "name: Validate nobrainer-book tutorials\n\non:\n workflow_dispatch: # Manual trigger only — not part of CI checks\n\njobs:\n validate-book:\n name: Run nobrainer-book tutorials\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v4\n\n - name: Install uv\n uses: astral-sh/setup-uv@v4\n\n - name: Set up Python\n run: uv venv --python 3.14\n\n - name: Install nobrainer and tutorial deps from PR branch\n run: |\n uv pip install \\\n \".[bayesian,generative,zarr,dev]\" \\\n monai \\\n pyro-ppl \\\n matplotlib \\\n nilearn\n\n - name: Clone nobrainer-book (matching branch or alpha)\n run: |\n PR_BRANCH=\"${{ github.head_ref }}\"\n BOOK_REPO=\"https://github.com/neuronets/nobrainer-book.git\"\n\n # Try the PR's branch name first (for lockstep development),\n # fall back to alpha\n if git ls-remote --heads \"$BOOK_REPO\" \"$PR_BRANCH\" | grep -q .; then\n echo \"Using matching book branch: $PR_BRANCH\"\n git clone --branch \"$PR_BRANCH\" --depth 1 \"$BOOK_REPO\" /tmp/nobrainer-book\n else\n echo \"No matching branch '$PR_BRANCH' on nobrainer-book, using alpha\"\n git clone --branch alpha --depth 1 \"$BOOK_REPO\" /tmp/nobrainer-book\n fi\n\n - name: Run book tutorials\n run: |\n for script in /tmp/nobrainer-book/docs/nobrainer-guides/scripts/0*.py /tmp/nobrainer-book/docs/nobrainer-guides/scripts/1[01]*.py; do\n echo \"=== Running $(basename $script) ===\"\n uv run python \"$script\" || {\n echo \"FAILED: $script\"\n exit 1\n }\n done\n echo \"All book tutorials passed\"\n" }, { "path": ".gitignore", "content": "# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.py,cover\n.hypothesis/\n.pytest_cache/\ncover/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\ndb.sqlite3-journal\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\n.pybuilder/\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pyenv\n# For a library or package, you might want to ignore these files since the code is\n# intended to run in multiple environments; otherwise, check them in:\n# .python-version\n\n# pipenv\n# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n# However, in case of collaboration, if having platform-specific dependencies or dependencies\n# having no cross-platform support, pipenv may install dependencies that don't work, or not\n# install all needed dependencies.\n#Pipfile.lock\n\n# PEP 582; used by e.g. github.com/David-OConnor/pyflow\n__pypackages__/\n\n# Celery stuff\ncelerybeat-schedule\ncelerybeat.pid\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n\n# pytype static type analyzer\n.pytype/\n\n# Cython debug symbols\ncython_debug/\n\n# pycharm\n.idea/\n\n# guide data\nguide/data/\n\n# Model artifacts\n*.pth\nbrain_mask_extraction_model/\ndata/\n\n# Model artifacts\n*.pth\nbrain_mask_extraction_model/\n" }, { "path": ".pre-commit-config.yaml", "content": "# See https://pre-commit.com for more information\n# See https://pre-commit.com/hooks.html for more hooks\nci:\n skip: [codespell]\n\nrepos:\n- repo: https://github.com/pre-commit/pre-commit-hooks\n rev: v4.5.0\n hooks:\n - id: check-added-large-files\n - id: check-yaml\n - id: end-of-file-fixer\n - id: trailing-whitespace\n- repo: https://github.com/psf/black\n rev: 24.3.0\n hooks:\n - id: black\n- repo: https://github.com/PyCQA/flake8\n rev: 7.0.0\n hooks:\n - id: flake8\n- repo: https://github.com/PyCQA/isort\n rev: 5.13.2\n hooks:\n - id: isort\n exclude: ^(nobrainer/_version\\.py|versioneer\\.py)$\n- repo: https://github.com/codespell-project/codespell\n rev: v2.2.6\n hooks:\n - id: codespell\n exclude: ^(nobrainer/_version\\.py|versioneer\\.py|pyproject\\.toml|CHANGELOG\\.md)$\n" }, { "path": ".zenodo.json", "content": "{\n \"creators\": [\n {\n \"affiliation\": \"Stony Brook University\",\n \"name\": \"Kaczmarzyk, Jakub\",\n \"orcid\": \"0000-0002-5544-7577\"\n },\n {\n \"affiliation\": \"NIMH\",\n \"name\": \"McClure, Patrick\"\n },\n {\n \"affiliation\": \"MIT\",\n \"name\": \"Zulfikar, Wazeer\"\n },\n {\n \"affiliation\": \"MIT\",\n \"name\": \"Rana, Aakanksha\",\n \"orcid\": \"0000-0002-8350-7602\"\n },\n {\n \"affiliation\": \"MIT\",\n \"name\": \"Rajaei, Hoda\",\n \"orcid\": \"0000-0002-0754-5586\"\n },\n {\n \"affiliation\": \"University of Washington\",\n \"name\": \"Richie-Halford, Adam\",\n \"orcid\": \"0000-0001-9276-9084\"\n },\n {\n \"affiliation\": \"Department of Psychology, Stanford University\",\n \"name\": \"Bansal, Shashank\",\n \"orcid\": \"0000-0002-1252-8772\"\n },\n {\n \"affiliation\": \"MIT\",\n \"name\": \"Jarecka, Dorota\",\n \"orcid\": \"0000-0001-8282-2988\"\n },\n {\n \"affiliation\": \"NIMH\",\n \"name\": \"Lee, John\"\n },\n {\n \"affiliation\": \"MIT, HMS\",\n \"name\": \"Ghosh, Satrajit\",\n \"orcid\": \"0000-0002-5312-6729\"\n }\n ],\n \"keywords\": [\n \"neuroimaging\",\n \"deep learning\",\n \"bayesian neural network\"\n ],\n \"license\": \"Apache-2.0\",\n \"upload_type\": \"software\"\n}\n" }, { "path": "CHANGELOG.md", "content": "# 1.2.1 (Thu Apr 04 2024)\n\n#### 🐛 Bug Fix\n\n- Fix PGAN notebook [#319](https://github.com/neuronets/nobrainer/pull/319) ([@satra](https://github.com/satra) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- fix dependencies [#318](https://github.com/neuronets/nobrainer/pull/318) ([@satra](https://github.com/satra))\n- Update setup.cfg to add cuda option [#309](https://github.com/neuronets/nobrainer/pull/309) ([@satra](https://github.com/satra))\n- [pre-commit.ci] pre-commit autoupdate [#294](https://github.com/neuronets/nobrainer/pull/294) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra) [@hvgazula](https://github.com/hvgazula))\n\n#### Authors: 3\n\n- [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot])\n- H Gazula ([@hvgazula](https://github.com/hvgazula))\n- Satrajit Ghosh ([@satra](https://github.com/satra))\n\n---\n\n# 1.2.0 (Fri Mar 22 2024)\n\n#### 🚀 Enhancement\n\n- Dev [#295](https://github.com/neuronets/nobrainer/pull/295) ([@hvgazula](https://github.com/hvgazula) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))\n- Update setup.cfg [#299](https://github.com/neuronets/nobrainer/pull/299) ([@satra](https://github.com/satra))\n\n#### 🐛 Bug Fix\n\n- Update release.yml ([@satra](https://github.com/satra))\n- change hyphenation [#275](https://github.com/neuronets/nobrainer/pull/275) ([@satra](https://github.com/satra))\n- update precommit checks [#275](https://github.com/neuronets/nobrainer/pull/275) ([@satra](https://github.com/satra))\n- fix docker syntax [#275](https://github.com/neuronets/nobrainer/pull/275) ([@satra](https://github.com/satra))\n- remove trained models [#275](https://github.com/neuronets/nobrainer/pull/275) ([@satra](https://github.com/satra))\n- [pre-commit.ci] pre-commit autoupdate [#269](https://github.com/neuronets/nobrainer/pull/269) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n\n#### Authors: 3\n\n- [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot])\n- H Gazula ([@hvgazula](https://github.com/hvgazula))\n- Satrajit Ghosh ([@satra](https://github.com/satra))\n\n---\n\n# 1.1.1 (Sat Oct 07 2023)\n\n#### 🐛 Bug Fix\n\n- Small changes to support long, preemptable training runs [#267](https://github.com/neuronets/nobrainer/pull/267) ([@ohinds](https://github.com/ohinds) [@satra](https://github.com/satra))\n- [pre-commit.ci] pre-commit autoupdate [#268](https://github.com/neuronets/nobrainer/pull/268) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n\n#### Authors: 3\n\n- [@ohinds](https://github.com/ohinds)\n- [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot])\n- Satrajit Ghosh ([@satra](https://github.com/satra))\n\n---\n\n# 1.1.0 (Tue Sep 19 2023)\n\n#### 🚀 Enhancement\n\n- Changes required to support the warmstart guide notebook [#266](https://github.com/neuronets/nobrainer/pull/266) ([@ohinds](https://github.com/ohinds) [@satra](https://github.com/satra) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n\n#### 🐛 Bug Fix\n\n- Fix some typos using codespell [#262](https://github.com/neuronets/nobrainer/pull/262) ([@yarikoptic](https://github.com/yarikoptic) [@satra](https://github.com/satra))\n- [pre-commit.ci] pre-commit autoupdate [#263](https://github.com/neuronets/nobrainer/pull/263) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))\n- Remove unnecessary keepalive runner [#265](https://github.com/neuronets/nobrainer/pull/265) ([@ohinds](https://github.com/ohinds))\n- Dynamically provision self-hosted runner [#264](https://github.com/neuronets/nobrainer/pull/264) ([@ohinds](https://github.com/ohinds))\n\n#### Authors: 4\n\n- [@ohinds](https://github.com/ohinds)\n- [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot])\n- Satrajit Ghosh ([@satra](https://github.com/satra))\n- Yaroslav Halchenko ([@yarikoptic](https://github.com/yarikoptic))\n\n---\n\n# 1.0.0 (Thu Aug 31 2023)\n\n#### 💥 Breaking Change\n\n- `nobrainer` dataset API rework [#261](https://github.com/neuronets/nobrainer/pull/261) ([@ohinds](https://github.com/ohinds) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))\n\n#### 🐛 Bug Fix\n\n- Fix Dockerfiles and add GitHub Actions job [#252](https://github.com/neuronets/nobrainer/pull/252) ([@kabilar](https://github.com/kabilar) [@satra](https://github.com/satra))\n- Self-hosted runner weekly keepalive [#260](https://github.com/neuronets/nobrainer/pull/260) ([@ohinds](https://github.com/ohinds))\n- [pre-commit.ci] pre-commit autoupdate [#253](https://github.com/neuronets/nobrainer/pull/253) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))\n- Processing model checkpointing [#256](https://github.com/neuronets/nobrainer/pull/256) ([@ohinds](https://github.com/ohinds) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))\n- Verify that the ec2 instance we started is where we are running. [#257](https://github.com/neuronets/nobrainer/pull/257) ([@ohinds](https://github.com/ohinds))\n- Training from warm start with multiple GPUs [#251](https://github.com/neuronets/nobrainer/pull/251) ([@ohinds](https://github.com/ohinds) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n\n#### Authors: 4\n\n- [@ohinds](https://github.com/ohinds)\n- [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot])\n- Kabilar Gunalan ([@kabilar](https://github.com/kabilar))\n- Satrajit Ghosh ([@satra](https://github.com/satra))\n\n---\n\n# 0.5.0 (Wed Jul 19 2023)\n\n#### 🚀 Enhancement\n\n- Remove guide [#243](https://github.com/neuronets/nobrainer/pull/243) ([@ohinds](https://github.com/ohinds) [@satra](https://github.com/satra))\n\n#### 🐛 Bug Fix\n\n- [pre-commit.ci] pre-commit autoupdate [#247](https://github.com/neuronets/nobrainer/pull/247) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))\n- Fix #246: Transforms now return labels if passed [#250](https://github.com/neuronets/nobrainer/pull/250) ([@ohinds](https://github.com/ohinds) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- Check out the master branch of the book repo on CI [#249](https://github.com/neuronets/nobrainer/pull/249) ([@ohinds](https://github.com/ohinds))\n- Nobrainer book guide examples run on AWS EC2 as a form of regression testing [#248](https://github.com/neuronets/nobrainer/pull/248) ([@ohinds](https://github.com/ohinds))\n- [pre-commit.ci] pre-commit autoupdate [#236](https://github.com/neuronets/nobrainer/pull/236) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))\n- Multi-GPU support [#242](https://github.com/neuronets/nobrainer/pull/242) ([@ohinds](https://github.com/ohinds) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- [pre-commit.ci] pre-commit autoupdate [#234](https://github.com/neuronets/nobrainer/pull/234) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- [pre-commit.ci] pre-commit autoupdate [#232](https://github.com/neuronets/nobrainer/pull/232) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- [pre-commit.ci] pre-commit autoupdate [#231](https://github.com/neuronets/nobrainer/pull/231) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n\n#### ⚠️ Pushed to `master`\n\n- Update setup.cfg ([@satra](https://github.com/satra))\n- Update ci.yml ([@satra](https://github.com/satra))\n- update python and tensorflow versions ([@satra](https://github.com/satra))\n- [CI] update python and auto versions ([@satra](https://github.com/satra))\n- replace special branch ([@satra](https://github.com/satra))\n\n#### Authors: 3\n\n- [@ohinds](https://github.com/ohinds)\n- [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot])\n- Satrajit Ghosh ([@satra](https://github.com/satra))\n\n---\n\n# 0.4.0 (Tue Oct 18 2022)\n\n#### 🚀 Enhancement\n\n- update actions [#230](https://github.com/neuronets/nobrainer/pull/230) ([@satra](https://github.com/satra))\n- [pre-commit.ci] pre-commit autoupdate [#229](https://github.com/neuronets/nobrainer/pull/229) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n\n#### 🐛 Bug Fix\n\n- Enh/api [#228](https://github.com/neuronets/nobrainer/pull/228) ([@satra](https://github.com/satra) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- [pre-commit.ci] pre-commit autoupdate [#227](https://github.com/neuronets/nobrainer/pull/227) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- Create unet_lstm.py [#207](https://github.com/neuronets/nobrainer/pull/207) ([@Aakanksha-Rana](https://github.com/Aakanksha-Rana) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra) [@Hoda1394](https://github.com/Hoda1394))\n- [pre-commit.ci] pre-commit autoupdate [#226](https://github.com/neuronets/nobrainer/pull/226) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- [pre-commit.ci] pre-commit autoupdate [#225](https://github.com/neuronets/nobrainer/pull/225) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- [pre-commit.ci] pre-commit autoupdate [#224](https://github.com/neuronets/nobrainer/pull/224) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- [pre-commit.ci] pre-commit autoupdate [#223](https://github.com/neuronets/nobrainer/pull/223) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- [WIP] Updating and simplifying the end-user API of Nobrainer [#215](https://github.com/neuronets/nobrainer/pull/215) ([@Aakanksha-Rana](https://github.com/Aakanksha-Rana) [@Hoda1394](https://github.com/Hoda1394) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))\n- [pre-commit.ci] pre-commit autoupdate [#220](https://github.com/neuronets/nobrainer/pull/220) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- [pre-commit.ci] pre-commit autoupdate [#218](https://github.com/neuronets/nobrainer/pull/218) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- Hard-Coded Scaling is removed for ProgressiveGANs, GanTrainer and ProgressiveAE Training and Model files [#209](https://github.com/neuronets/nobrainer/pull/209) ([@Aakanksha-Rana](https://github.com/Aakanksha-Rana) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))\n- generation api [#216](https://github.com/neuronets/nobrainer/pull/216) ([@Hoda1394](https://github.com/Hoda1394) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- [WIP] segmentation api [#213](https://github.com/neuronets/nobrainer/pull/213) ([@Aakanksha-Rana](https://github.com/Aakanksha-Rana) [@Hoda1394](https://github.com/Hoda1394) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- file name corrected for bayesian_meshnet.py [#211](https://github.com/neuronets/nobrainer/pull/211) ([@Hoda1394](https://github.com/Hoda1394) [@satra](https://github.com/satra))\n- [pre-commit.ci] pre-commit autoupdate [#212](https://github.com/neuronets/nobrainer/pull/212) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))\n- Adding vox2vox model [#205](https://github.com/neuronets/nobrainer/pull/205) ([@Aakanksha-Rana](https://github.com/Aakanksha-Rana) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@Hoda1394](https://github.com/Hoda1394))\n- added tensorflow-addons as install dependency [#206](https://github.com/neuronets/nobrainer/pull/206) ([@Hoda1394](https://github.com/Hoda1394))\n- Composable Data Augmentations [#189](https://github.com/neuronets/nobrainer/pull/189) ([@Aakanksha-Rana](https://github.com/Aakanksha-Rana) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))\n\n#### ⚠️ Pushed to `master`\n\n- update dockerfiles ([@satra](https://github.com/satra))\n\n#### Authors: 4\n\n- [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot])\n- Aakanksha Rana ([@Aakanksha-Rana](https://github.com/Aakanksha-Rana))\n- Hoda Rajaei ([@Hoda1394](https://github.com/Hoda1394))\n- Satrajit Ghosh ([@satra](https://github.com/satra))\n\n---\n\n# 0.3.0 (Tue Jan 11 2022)\n\n#### 🚀 Enhancement\n\n- Update README.md [#203](https://github.com/neuronets/nobrainer/pull/203) ([@satra](https://github.com/satra))\n\n#### 🐛 Bug Fix\n\n- Progressive auto encoder [#196](https://github.com/neuronets/nobrainer/pull/196) (alice.bizeul@inf.ethz.ch [@satra](https://github.com/satra) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- [pre-commit.ci] pre-commit autoupdate [#197](https://github.com/neuronets/nobrainer/pull/197) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n\n#### ⚠️ Pushed to `master`\n\n- Update README.md ([@satra](https://github.com/satra))\n\n#### Authors: 3\n\n- [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot])\n- Alice Elsa Marie Bizeul (alice.bizeul@inf.ethz.ch)\n- Satrajit Ghosh ([@satra](https://github.com/satra))\n\n---\n\n# 0.2.1 (Fri Dec 24 2021)\n\n#### 🐛 Bug Fix\n\n- fix: update docker files and citation pointer [#194](https://github.com/neuronets/nobrainer/pull/194) ([@satra](https://github.com/satra) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n\n#### Authors: 2\n\n- [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot])\n- Satrajit Ghosh ([@satra](https://github.com/satra))\n\n---\n\n# 0.2.0 (Fri Dec 24 2021)\n\n#### 🚀 Enhancement\n\n- Update README.md and move transforms [#187](https://github.com/neuronets/nobrainer/pull/187) ([@Hoda1394](https://github.com/Hoda1394) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))\n\n#### 🐛 Bug Fix\n\n- Documentation update [#184](https://github.com/neuronets/nobrainer/pull/184) ([@Aakanksha-Rana](https://github.com/Aakanksha-Rana) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))\n- Correction of G_dice and multi_class dice_calculation [#164](https://github.com/neuronets/nobrainer/pull/164) ([@kaczmarj](https://github.com/kaczmarj) [@Aakanksha-Rana](https://github.com/Aakanksha-Rana) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))\n- DOC: Update link to CITATION file [#190](https://github.com/neuronets/nobrainer/pull/190) ([@arokem](https://github.com/arokem))\n- Bayesian: 3D bayes_by_backprop_layer + Distributed Weight Consolidation [#185](https://github.com/neuronets/nobrainer/pull/185) ([@Aakanksha-Rana](https://github.com/Aakanksha-Rana) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))\n- [pre-commit.ci] pre-commit autoupdate [#188](https://github.com/neuronets/nobrainer/pull/188) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- Fix tf & tfp compatibility with python version [#186](https://github.com/neuronets/nobrainer/pull/186) ([@Hoda1394](https://github.com/Hoda1394) [@satra](https://github.com/satra))\n- [pre-commit.ci] pre-commit autoupdate [#183](https://github.com/neuronets/nobrainer/pull/183) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- Added Clipping constraint and KLD loss with CD [#180](https://github.com/neuronets/nobrainer/pull/180) ([@Aakanksha-Rana](https://github.com/Aakanksha-Rana) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))\n- ENH: Allow multichannel input [#177](https://github.com/neuronets/nobrainer/pull/177) ([@richford](https://github.com/richford) [@satra](https://github.com/satra))\n- Add citation [#181](https://github.com/neuronets/nobrainer/pull/181) ([@Aakanksha-Rana](https://github.com/Aakanksha-Rana) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))\n- [pre-commit.ci] pre-commit autoupdate [#182](https://github.com/neuronets/nobrainer/pull/182) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- adding brain siamese network models and test cases [#172](https://github.com/neuronets/nobrainer/pull/172) ([@dhritimandas](https://github.com/dhritimandas) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))\n- Added dcgan architecture with tests [#66](https://github.com/neuronets/nobrainer/pull/66) ([@wazeerzulfikar](https://github.com/wazeerzulfikar) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))\n- BF: Add block_length param to dataset.interleave [#175](https://github.com/neuronets/nobrainer/pull/175) ([@richford](https://github.com/richford) [@satra](https://github.com/satra))\n- Weightnorm feature for Variational layer [#166](https://github.com/neuronets/nobrainer/pull/166) ([@Aakanksha-Rana](https://github.com/Aakanksha-Rana) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))\n- [pre-commit.ci] pre-commit autoupdate [#179](https://github.com/neuronets/nobrainer/pull/179) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- [pre-commit.ci] pre-commit autoupdate [#178](https://github.com/neuronets/nobrainer/pull/178) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- [pre-commit.ci] pre-commit autoupdate [#173](https://github.com/neuronets/nobrainer/pull/173) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- [pre-commit.ci] pre-commit autoupdate [#170](https://github.com/neuronets/nobrainer/pull/170) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- [pre-commit.ci] pre-commit autoupdate [#165](https://github.com/neuronets/nobrainer/pull/165) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- [pre-commit.ci] pre-commit autoupdate [#163](https://github.com/neuronets/nobrainer/pull/163) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- fix: readme to remove qualified install [#162](https://github.com/neuronets/nobrainer/pull/162) ([@satra](https://github.com/satra) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- fix: remove unnecessary step in release action [#158](https://github.com/neuronets/nobrainer/pull/158) ([@satra](https://github.com/satra))\n\n#### ⚠️ Pushed to `master`\n\n- Update .zenodo.json ([@satra](https://github.com/satra))\n\n#### Authors: 9\n\n- [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot])\n- Aakanksha Rana ([@Aakanksha-Rana](https://github.com/Aakanksha-Rana))\n- Adam Richie-Halford ([@richford](https://github.com/richford))\n- Ariel Rokem ([@arokem](https://github.com/arokem))\n- Dhritiman Das ([@dhritimandas](https://github.com/dhritimandas))\n- Hoda Rajaei ([@Hoda1394](https://github.com/Hoda1394))\n- Jakub Kaczmarzyk ([@kaczmarj](https://github.com/kaczmarj))\n- Satrajit Ghosh ([@satra](https://github.com/satra))\n- Wazeer Zulfikar ([@wazeerzulfikar](https://github.com/wazeerzulfikar))\n\n---\n\n# 0.1.1 (Tue Jun 22 2021)\n\n#### 🐛 Bug Fix\n\n- fix: replace key retrieval and normalizers [#157](https://github.com/neuronets/nobrainer/pull/157) ([@satra](https://github.com/satra))\n- fix: separate auto release and publish to pypi [#156](https://github.com/neuronets/nobrainer/pull/156) ([@satra](https://github.com/satra))\n- [pre-commit.ci] pre-commit autoupdate [#154](https://github.com/neuronets/nobrainer/pull/154) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n\n#### ⚠️ Pushed to `master`\n\n- fix: add twine upload to release ([@satra](https://github.com/satra))\n\n#### Authors: 2\n\n- [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot])\n- Satrajit Ghosh ([@satra](https://github.com/satra))\n\n---\n\n# 0.1.0 (Sat Jun 19 2021)\n\n#### 🚀 Enhancement\n\n- fix: change nobrainer models repo, readme, and guide notebooks [#150](https://github.com/neuronets/nobrainer/pull/150) ([@satra](https://github.com/satra) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- Enh/docker [#148](https://github.com/neuronets/nobrainer/pull/148) ([@satra](https://github.com/satra))\n- fix: Update release.yml to handle branch protection [#147](https://github.com/neuronets/nobrainer/pull/147) ([@satra](https://github.com/satra))\n\n#### 🐛 Bug Fix\n\n- add zenodo file [#153](https://github.com/neuronets/nobrainer/pull/153) ([@satra](https://github.com/satra) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- enh: change standardize in dataset creation to a callable [#152](https://github.com/neuronets/nobrainer/pull/152) ([@satra](https://github.com/satra))\n- ENH: Bayesian neural network architectures with training requisites [#126](https://github.com/neuronets/nobrainer/pull/126) ([@Aakanksha-Rana](https://github.com/Aakanksha-Rana) [@satra](https://github.com/satra) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- fix: standardize not imported for dataset creation [#151](https://github.com/neuronets/nobrainer/pull/151) ([@satra](https://github.com/satra))\n- add estimator prediction function for kwyk model [#149](https://github.com/neuronets/nobrainer/pull/149) ([@Hoda1394](https://github.com/Hoda1394) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))\n- Fix for the generate cli pytest and generation guide notebook [#146](https://github.com/neuronets/nobrainer/pull/146) ([@wazeerzulfikar](https://github.com/wazeerzulfikar))\n- Created using Colaboratory [#144](https://github.com/neuronets/nobrainer/pull/144) ([@satra](https://github.com/satra) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- [pre-commit.ci] pre-commit autoupdate [#142](https://github.com/neuronets/nobrainer/pull/142) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- FIX: keeping compatibility with official tensorflow docker images [#141](https://github.com/neuronets/nobrainer/pull/141) ([@satra](https://github.com/satra))\n- Add bayesian prediction [#128](https://github.com/neuronets/nobrainer/pull/128) ([@Hoda1394](https://github.com/Hoda1394) [@satra](https://github.com/satra) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- ENH: Transform composition [#113](https://github.com/neuronets/nobrainer/pull/113) ([@Aakanksha-Rana](https://github.com/Aakanksha-Rana) [@satra](https://github.com/satra) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))\n- fixes https://github.com/neuronets/nobrainer/issues/124 [#125](https://github.com/neuronets/nobrainer/pull/125) ([@shashankbansal6](https://github.com/shashankbansal6) [@satra](https://github.com/satra))\n- fix: update dockerfiles to tensorflow 2.5.0 [#140](https://github.com/neuronets/nobrainer/pull/140) ([@satra](https://github.com/satra))\n- ENH: Add progressive gan to nobrainer [#138](https://github.com/neuronets/nobrainer/pull/138) ([@wazeerzulfikar](https://github.com/wazeerzulfikar) [@satra](https://github.com/satra) [@djarecka](https://github.com/djarecka))\n- add: release mechanism [#139](https://github.com/neuronets/nobrainer/pull/139) ([@satra](https://github.com/satra))\n- Add progressiveGAN for 3D brain MR images [#114](https://github.com/neuronets/nobrainer/pull/114) ([@wazeerzulfikar](https://github.com/wazeerzulfikar))\n- Enh/notebooks [#137](https://github.com/neuronets/nobrainer/pull/137) ([@satra](https://github.com/satra))\n- add CI workflow badge [#133](https://github.com/neuronets/nobrainer/pull/133) ([@kaczmarj](https://github.com/kaczmarj))\n- move CI to github actions [#131](https://github.com/neuronets/nobrainer/pull/131) ([@Hoda1394](https://github.com/Hoda1394))\n- Enh/minor updates [#123](https://github.com/neuronets/nobrainer/pull/123) ([@kaczmarj](https://github.com/kaczmarj))\n- export `LC_ALL` and `LANG` + use tensorflow 2.3.1 [#118](https://github.com/neuronets/nobrainer/pull/118) ([@kaczmarj](https://github.com/kaczmarj))\n- use specific version of cloudpickle to fix import error [#108](https://github.com/neuronets/nobrainer/pull/108) ([@kaczmarj](https://github.com/kaczmarj))\n- force reinstall of python dependencies in travis ci [#102](https://github.com/neuronets/nobrainer/pull/102) ([@kaczmarj](https://github.com/kaczmarj))\n\n#### ⚠️ Pushed to `master`\n\n- fix: use token for auto with repo access ([@satra](https://github.com/satra))\n- Created using Colaboratory ([@satra](https://github.com/satra))\n- Update release.yml ([@satra](https://github.com/satra))\n\n#### Authors: 8\n\n- [@Hoda1394](https://github.com/Hoda1394)\n- [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot])\n- Aakanksha Rana ([@Aakanksha-Rana](https://github.com/Aakanksha-Rana))\n- Dorota Jarecka ([@djarecka](https://github.com/djarecka))\n- Jakub Kaczmarzyk ([@kaczmarj](https://github.com/kaczmarj))\n- Satrajit Ghosh ([@satra](https://github.com/satra))\n- Shashank Bansal ([@shashankbansal6](https://github.com/shashankbansal6))\n- Wazeer Zulfikar ([@wazeerzulfikar](https://github.com/wazeerzulfikar))\n" }, { "path": "CITATION", "content": "Please follow this DOI (https://doi.org/10.5281/zenodo.4995077) to find\nthe latest citation on Zenodo. The different citation formats are available\nin the Share and Export sections of the page. On a desktop browser these\nare on the bottom right of the page.\n" }, { "path": "CLAUDE.md", "content": "# Nobrainer Development Guidelines\n\n## Project Overview\n\nNobrainer is a PyTorch-based deep learning library for 3D brain MRI segmentation. It provides scikit-learn-style estimators (`Segmentation`, `Generation`), Bayesian models (VWN/FFG, Pyro-based), and a comprehensive data pipeline (MONAI transforms, Zarr3 stores, SynthSeg generation).\n\n## Technology Stack\n\n- **Python**: 3.12+; CI matrix 3.12/3.13/3.14\n- **Package management**: `uv` throughout (never pip/conda/poetry)\n- **ML framework**: PyTorch >= 2.0\n- **Medical imaging**: MONAI >= 1.3 (transforms, losses, metrics, model wrappers)\n- **Bayesian**: Pyro-ppl >= 1.9 (optional `[bayesian]` extra)\n- **Data**: Zarr >= 3.0 (optional `[zarr]` extra), NIfTI via nibabel\n- **Testing**: pytest; pre-commit (black, flake8, isort, codespell)\n- **CI**: GitHub Actions; EC2 GPU runner for GPU tests\n\n## Commands\n\n```bash\n# Install\nuv pip install -e \".[all]\"\n\n# Test (CPU)\nuv run pytest nobrainer/tests/unit/ -m \"not gpu\" --tb=short\n\n# SR-tests (somewhat realistic, need sample brain data)\nuv run pytest nobrainer/sr-tests/ -m \"not gpu\"\n\n# Lint\nuv run pre-commit run --all-files\n```\n\n## Code Conventions\n\n- All models: `(B, C, D, H, W)` input → `(B, n_classes, D, H, W)` output\n- Factory functions: `model_name(n_classes=1, in_channels=1, **kwargs) -> nn.Module`\n- Bayesian models: `supports_mc = True` class attribute; `forward(x, **kwargs)` accepts `mc=True/False`\n- Prediction: use `model_supports_mc(model)` to check, never `try/except TypeError`\n- Labels: always squeeze channel dim + cast to `long` before `CrossEntropyLoss`\n- Device selection: `nobrainer.gpu.get_device()` (CUDA > MPS > CPU)\n- Data augmentation: `TrainableCompose` wraps MONAI Compose; `Augmentation()` wrapper auto-skips during predict\n\n## Key Modules\n\n| Module | Purpose |\n|--------|---------|\n| `models/` | MeshNet, SegFormer3D, UNet, SwinUNETR, SegResNet, Bayesian variants |\n| `processing/` | Segmentation/Generation estimators, Dataset builder |\n| `augmentation/` | SynthSeg generator, TrainableCompose, profiles |\n| `datasets/` | OpenNeuro fetching, Zarr3 store management |\n| `training.py` | `fit()` with DDP, AMP, validation, callbacks |\n| `prediction.py` | Block-based predict, strided reassembly, MC uncertainty |\n| `losses.py` | Dice, FocalLoss, DiceCE, ELBO, class weights |\n| `gpu.py` | Device detection, auto batch size, multi-GPU scaling |\n| `slurm.py` | SLURM preemption handler, checkpoint/resume |\n| `experiment.py` | Local JSONL/CSV + optional W&B tracking |\n\n## Development Workflow (Speckit Constitution)\n\nWhen working on new features or significant changes, follow these principles:\n\n### I. Specification-First\n\nEvery feature MUST begin with a written specification before implementation:\n- Prioritized user stories with independently testable acceptance scenarios\n- Functional requirements written as verifiable constraints (MUST/SHOULD)\n- Measurable success criteria that are technology-agnostic\n\n### II. Incremental Planning\n\nPlans are built in ordered phases — no phase may be skipped:\n- **Phase 0 — Research**: Resolve all unknowns before design\n- **Phase 1 — Design**: Data model, interface contracts, quickstart documented\n- **Phase 2 — Tasks**: Actionable task list organized by user story priority\n\nImplementation MUST NOT begin until tasks exist.\n\n### III. Independent User-Story Delivery\n\n- Each P1 story MUST produce a viable MVP with standalone value\n- Stories MUST NOT have hard runtime dependencies on lower-priority stories\n- Tasks MUST be labeled with their owning story (`[US1]`, `[US2]`, etc.)\n\n### IV. Constitution Compliance Gate\n\nEvery plan MUST include a Constitution Check evaluated before research and after design. Violations MUST be justified with a simpler alternative explicitly rejected.\n\n### V. Simplicity & YAGNI\n\n- Prefer the simplest architecture that satisfies current user stories\n- Do not introduce abstractions for hypothetical future requirements\n- Complexity MUST be justified against a concrete, present need\n\n### VI. Git Commit Discipline\n\n- Feature work on dedicated branches (`###-feature-name`)\n- Planning artifacts committed after each speckit command\n- Each completed task results in at least one commit\n- Prefer new commits over amending\n\n### VII. Technology Stack Standards\n\n- Python: `uv` for all environment and package management\n- Containers: Docker only\n- No substitutions without justified amendment\n\n## Quality Gates\n\n| Gate | Condition |\n|------|-----------|\n| G1 | spec.md has ≥1 user story with acceptance scenarios |\n| G2 | All NEEDS CLARIFICATION resolved before design |\n| G3 | Constitution Check passes (or violations justified) |\n| G4 | tasks.md exists and all tasks reference a user story |\n| G5 | P1 story independently verified before P2 work |\n| G6 | All planning artifacts committed to feature branch |\n\n## Speckit Commands (if available)\n\n```\n/speckit.specify → spec.md\n/speckit.clarify → spec.md (revised)\n/speckit.plan → plan.md, research.md, data-model.md, quickstart.md\n/speckit.tasks → tasks.md\n/speckit.implement → code\n/speckit.analyze → consistency report\n```\n\nIf speckit is not installed, follow the principles above manually.\n" }, { "path": "LICENSE", "content": "Copyright 2021 The Nobrainer Authors.\n\nLicensed under the Apache License, Version 2.0 (the \"License\");\nyou may not use this file except in compliance with the License.\nYou may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\nUnless required by applicable law or agreed to in writing, software\ndistributed under the License is distributed on an \"AS IS\" BASIS,\nWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\nSee the License for the specific language governing permissions and\nlimitations under the License.\n" }, { "path": "MANIFEST.in", "content": "# This line includes versioneer.py in sdists, which is necessary for wheels\n# built from sdists to have the version set in their metadata.\ninclude versioneer.py\ninclude CHANGELOG.md tox.ini\n\ngraft nobrainer\n\nglobal-exclude *.py[cod]\ninclude nobrainer/_version.py\n" }, { "path": "README.md", "content": "# Nobrainer\n\n![Build status](https://github.com/neuronets/nobrainer/actions/workflows/ci.yml/badge.svg)\n\n_Nobrainer_ is a deep learning framework for 3D brain image processing built on\n**PyTorch** and **MONAI**. It provides segmentation models (deterministic and\nBayesian), generative models, a MONAI-native data pipeline, block-based\nprediction with uncertainty quantification, and CLI tools for inference and\nautomated hyperparameter search.\n\nPre-trained models for brain extraction, segmentation, and generation are\navailable in the [trained-models](https://github.com/neuronets/trained-models)\nrepository.\n\nThe _Nobrainer_ project is supported by NIH RF1MH121885 and is distributed\nunder the Apache 2.0 license.\n\n## Models\n\n### Segmentation\n\n| Model | Backend | Application |\n|-------|---------|-------------|\n| [UNet](nobrainer/models/segmentation.py) | MONAI | segmentation |\n| [VNet](nobrainer/models/segmentation.py) | MONAI | segmentation |\n| [Attention U-Net](nobrainer/models/segmentation.py) | MONAI | segmentation |\n| [UNETR](nobrainer/models/segmentation.py) | MONAI | segmentation |\n| [MeshNet](nobrainer/models/meshnet.py) | PyTorch | segmentation |\n| [HighResNet](nobrainer/models/highresnet.py) | PyTorch | segmentation |\n\n### Bayesian (uncertainty quantification)\n\n| Model | Backend | Application |\n|-------|---------|-------------|\n| [Bayesian VNet](nobrainer/models/bayesian/bayesian_vnet.py) | Pyro | segmentation + uncertainty |\n| [Bayesian MeshNet](nobrainer/models/bayesian/bayesian_meshnet.py) | Pyro | segmentation + uncertainty |\n\n### Generative\n\n| Model | Backend | Application |\n|-------|---------|-------------|\n| [Progressive GAN](nobrainer/models/generative/progressivegan.py) | PyTorch Lightning | brain generation |\n| [DCGAN](nobrainer/models/generative/dcgan.py) | PyTorch Lightning | brain generation |\n\n### Other\n\n| Model | Application |\n|-------|-------------|\n| [Autoencoder](nobrainer/models/autoencoder.py) | representation learning |\n| [SimSiam](nobrainer/models/simsiam.py) | self-supervised learning |\n\n### Custom layers\n\n- `BernoulliDropout`, `ConcreteDropout`, `GaussianDropout` — stochastic regularization\n- `BayesianConv3d`, `BayesianLinear` — Pyro-based weight uncertainty layers\n- `MaxPool4D` — 4D max pooling via reshape\n\n### Losses and metrics\n\n**Losses**: Dice, Generalized Dice, Jaccard, Tversky, ELBO (Bayesian), Wasserstein, Gradient Penalty\n\n**Metrics**: Dice, Jaccard, Hausdorff distance (all via MONAI)\n\n## Installation\n\n### pip / uv\n\n```bash\nuv venv --python 3.14\nsource .venv/bin/activate\nuv pip install nobrainer\n```\n\nFor Bayesian and generative model support:\n\n```bash\nuv pip install \"nobrainer[bayesian,generative]\" monai pyro-ppl\n```\n\n### Docker\n\nGPU image (requires NVIDIA driver on host):\n\n```bash\ndocker pull neuronets/nobrainer:latest-gpu-pt\ndocker run --gpus all --rm neuronets/nobrainer:latest-gpu-pt predict --help\n```\n\nCPU-only image:\n\n```bash\ndocker pull neuronets/nobrainer:latest-cpu-pt\ndocker run --rm neuronets/nobrainer:latest-cpu-pt predict --help\n```\n\n## Quick start\n\n### Tutorials\n\nSee the [Nobrainer Book](https://neuronets.dev/nobrainer-book/) for 11\nprogressive tutorials — from installation to contributing.\n\n### sr-tests (somewhat realistic tests)\n\n`nobrainer/sr-tests/` contains pytest integration tests that exercise the\nreal API with real brain data. They run in CI on every push:\n\n```bash\npytest nobrainer/sr-tests/ -v -m \"not gpu\" --tb=short\n```\n\n### Simple API (3 lines)\n\n```python\nfrom nobrainer.processing import Segmentation, Dataset\n\nds = Dataset.from_files(filepaths, block_shape=(128, 128, 128), n_classes=2).batch(2)\nresult = Segmentation(\"unet\").fit(ds, epochs=5).predict(\"brain.nii.gz\")\n```\n\nModels are saved with [Croissant-ML](https://mlcommons.org/croissant/) metadata\nfor reproducibility:\n\n```python\nseg.save(\"my_model\") # Creates model.pth + croissant.json\nseg = Segmentation.load(\"my_model\")\n```\n\n### Brain segmentation (CLI)\n\n```bash\nnobrainer predict \\\n --model unet_brainmask.pth \\\n --model-type unet \\\n --n-classes 2 \\\n input_T1w.nii.gz output_mask.nii.gz\n```\n\n### Brain segmentation (Python)\n\n```python\nimport torch\nimport nobrainer\nfrom nobrainer.prediction import predict\n\nmodel = nobrainer.models.unet(n_classes=2)\nmodel.load_state_dict(torch.load(\"unet_brainmask.pth\"))\nmodel.eval()\n\nresult = predict(\n inputs=\"input_T1w.nii.gz\",\n model=model,\n block_shape=(128, 128, 128),\n device=\"cuda\",\n)\nresult.to_filename(\"output_mask.nii.gz\")\n```\n\n### Bayesian inference with uncertainty maps\n\n```python\nfrom nobrainer.prediction import predict_with_uncertainty\n\nmodel = nobrainer.models.bayesian_vnet(n_classes=2)\nmodel.load_state_dict(torch.load(\"bayesian_vnet.pth\"))\n\nlabel, variance, entropy = predict_with_uncertainty(\n inputs=\"input_T1w.nii.gz\",\n model=model,\n n_samples=10,\n block_shape=(128, 128, 128),\n device=\"cuda\",\n)\nlabel.to_filename(\"label.nii.gz\")\nvariance.to_filename(\"variance.nii.gz\")\nentropy.to_filename(\"entropy.nii.gz\")\n```\n\n### Brain generation\n\n```bash\nnobrainer generate \\\n --model progressivegan.ckpt \\\n --model-type progressivegan \\\n output_synthetic.nii.gz\n```\n\n### Zarr v3 data pipeline\n\n```python\nfrom nobrainer.io import nifti_to_zarr, zarr_to_nifti\n\n# Convert NIfTI to sharded Zarr v3 with multi-resolution pyramid\nnifti_to_zarr(\"brain_T1w.nii.gz\", \"brain.zarr\", chunk_shape=(64, 64, 64), levels=3)\n\n# Load Zarr stores directly in the training pipeline\nfrom nobrainer.dataset import get_dataset\n\nloader = get_dataset(\n data=[{\"image\": \"brain.zarr\", \"label\": \"label.zarr\"}],\n batch_size=2,\n)\n\n# Round-trip back to NIfTI\nzarr_to_nifti(\"brain.zarr\", \"brain_roundtrip.nii.gz\")\n```\n\n### Training a model\n\n```python\nimport torch\nfrom nobrainer.dataset import get_dataset\nfrom nobrainer.losses import dice\n\ndata_files = [\n {\"image\": f\"sub-{i:03d}_T1w.nii.gz\", \"label\": f\"sub-{i:03d}_label.nii.gz\"}\n for i in range(1, 101)\n]\nloader = get_dataset(data=data_files, batch_size=2, augment=True, cache=True)\n\nmodel = nobrainer.models.unet(n_classes=2).cuda()\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-4)\ncriterion = dice()\n\nfor epoch in range(50):\n model.train()\n for batch in loader:\n images, labels = batch[\"image\"].cuda(), batch[\"label\"].cuda()\n optimizer.zero_grad()\n loss = criterion(model(images), labels)\n loss.backward()\n optimizer.step()\n\ntorch.save(model.state_dict(), \"unet_trained.pth\")\n```\n\n## Automated research (autoresearch)\n\nNobrainer includes an automated hyperparameter search loop that uses an LLM\nto propose training modifications overnight:\n\n```bash\nnobrainer research run \\\n --working-dir ./research/bayesian_vnet \\\n --model-family bayesian_vnet \\\n --max-experiments 15 \\\n --budget-hours 8\n```\n\nImproved models are versioned via DataLad:\n\n```bash\nnobrainer research commit \\\n --run-dir ./research/bayesian_vnet \\\n --trained-models-path ~/trained-models \\\n --model-family bayesian_vnet\n```\n\n## GPU test dispatch (nobrainer-runner)\n\n[nobrainer-runner](https://github.com/neuronets/nobrainer-runner) submits GPU\ntest suites to Slurm clusters or cloud instances (AWS Batch, GCP Batch):\n\n```bash\nnobrainer-runner submit --profile mycluster --gpus 1 \"pytest tests/ -m gpu\"\nnobrainer-runner status $JOB_ID\nnobrainer-runner results --format json $JOB_ID\n```\n\n## Package layout\n\n- `nobrainer.models` — segmentation, Bayesian, and generative `torch.nn.Module` models\n- `nobrainer.losses` — Dice, Jaccard, Tversky, ELBO, Wasserstein (MONAI-backed)\n- `nobrainer.metrics` — Dice, Jaccard, Hausdorff (MONAI-backed)\n- `nobrainer.dataset` — MONAI `CacheDataset` + `DataLoader` pipeline\n- `nobrainer.prediction` — block-based `predict()` and `predict_with_uncertainty()`\n- `nobrainer.io` — `convert_tfrecords()`, `convert_weights()` (TF → PyTorch migration)\n- `nobrainer.layers` — dropout layers, Bayesian layers, MaxPool4D\n- `nobrainer.research` — autoresearch loop and DataLad model versioning\n- `nobrainer.cli` — Click CLI (`predict`, `generate`, `research`, `commit`, `info`)\n\n## Development and releases\n\nNobrainer uses a two-branch release workflow:\n\n| Branch | Purpose | PyPI version |\n|--------|---------|--------------|\n| `master` | Stable releases | `uv pip install nobrainer` |\n| `alpha` | Pre-releases for testing | `uv pip install --pre nobrainer` |\n\n**Alpha workflow**: Feature branches merge to `alpha`. Each merge triggers\nbook tutorial validation (using a matching branch on\n[nobrainer-book](https://github.com/neuronets/nobrainer-book) if available,\notherwise the book's `alpha` branch) followed by an automatic pre-release\ntag (e.g., `0.5.0-alpha.0`).\n\n**Stable workflow**: When `alpha` is merged to `master` with the `release`\nlabel, a stable version is tagged and published to PyPI.\n\n**GPU CI**: PRs to `master` can request GPU testing on EC2 by adding the\n`gpu-test-approved` label. Instance type and spot pricing are configurable\nvia `gpu-instance:` and `gpu-spot:true` labels.\n\n## Citation\n\nIf you use this package, please [cite](https://github.com/neuronets/nobrainer/blob/master/CITATION) it.\n\n## Questions or issues\n\nPlease [submit a GitHub issue](https://github.com/neuronets/helpdesk/issues/new/choose).\n" }, { "path": "conftest.py", "content": "\"\"\"Root conftest.py — auto-skip GPU tests when CUDA is unavailable.\"\"\"\n\nfrom __future__ import annotations\n\nimport pytest\nimport torch\n\n\ndef pytest_collection_modifyitems(config, items):\n \"\"\"Skip tests marked with @pytest.mark.gpu when CUDA is not available.\"\"\"\n if torch.cuda.is_available():\n return\n skip_gpu = pytest.mark.skip(reason=\"CUDA not available — skipping GPU test\")\n for item in items:\n if item.get_closest_marker(\"gpu\"):\n item.add_marker(skip_gpu)\n" }, { "path": "docker/README.md", "content": "# Nobrainer in a container\n\nThe Dockerfiles in this directory can be used to create Docker images to use _Nobrainer_ on CPU or GPU.\n\n## Build images\n\n```bash\ncd /code/nobrainer # Top-level nobrainer directory\ndocker build -t neuronets/nobrainer:master-cpu -f docker/cpu.Dockerfile .\ndocker build -t neuronets/nobrainer:master-gpu -f docker/gpu.Dockerfile .\n```\n\n# Convert Docker images to Singularity containers\n\nUsing Singularity version 3.x, Docker images can be converted to Singularity containers using the `singularity` command-line tool.\n\n## Pulling from DockerHub\n\nIn most cases (e.g., working on a HPC cluster), the _Nobrainer_ singularity container can be created with:\n\n```bash\nsingularity pull docker://neuronets/nobrainer:master-gpu\n```\n\n## Building from local Docker cache\n\nIf you built a _Nobrainer_ Docker images locally and would like to convert it to a Singularity container, you can do so with:\n\n```bash\nsudo singularity pull docker-daemon://neuronets/nobrainer:master-gpu\n```\n\nPlease note the use of `sudo` here. This is necessary for interacting with the Docker daemon.\n" }, { "path": "docker/cpu.Dockerfile", "content": "FROM python:3.14-slim\nENV DEBIAN_FRONTEND=noninteractive\nRUN apt-get update && apt-get install -y --no-install-recommends \\\n git \\\n && rm -rf /var/lib/apt/lists/*\nCOPY [\".\", \"/opt/nobrainer\"]\nRUN pip install --no-cache-dir uv \\\n && uv pip install --system \\\n \"torch\" \\\n \"/opt/nobrainer[bayesian,generative]\" \\\n monai \\\n pyro-ppl \\\n --index-url https://download.pytorch.org/whl/cpu \\\n --extra-index-url https://pypi.org/simple \\\n && rm -rf /root/.cache/uv\nENV LC_ALL=C.UTF-8 \\\n LANG=C.UTF-8\nWORKDIR \"/work\"\nLABEL maintainer=\"Satrajit Ghosh \"\nLABEL org.opencontainers.image.title=\"nobrainer-cpu-pytorch\"\nLABEL org.opencontainers.image.description=\"nobrainer with PyTorch CPU-only support\"\nENTRYPOINT [\"nobrainer\"]\n" }, { "path": "docker/gpu.Dockerfile", "content": "FROM python:3.14-slim\nENV DEBIAN_FRONTEND=noninteractive\nRUN apt-get update && apt-get install -y --no-install-recommends \\\n git \\\n && rm -rf /var/lib/apt/lists/*\nCOPY [\".\", \"/opt/nobrainer\"]\nRUN pip install --no-cache-dir uv \\\n && uv pip install --system \\\n torch \\\n \"/opt/nobrainer[bayesian,generative,versioning]\" \\\n monai \\\n pyro-ppl \\\n && rm -rf /root/.cache/uv\nENV LC_ALL=C.UTF-8 \\\n LANG=C.UTF-8\nWORKDIR \"/work\"\nLABEL maintainer=\"Satrajit Ghosh \"\nLABEL org.opencontainers.image.title=\"nobrainer-gpu-pytorch\"\nLABEL org.opencontainers.image.description=\"nobrainer with PyTorch GPU support (CUDA via host driver)\"\nENTRYPOINT [\"nobrainer\"]\n" }, { "path": "nobrainer/__init__.py", "content": "try:\n from ._version import __version__ # noqa: F401\nexcept (ImportError, ModuleNotFoundError):\n try:\n from . import _version # noqa: F401\n\n __version__ = _version.get_versions()[\"version\"]\n except (ImportError, AttributeError):\n __version__ = \"0.0.0.dev0\"\n\n# Lazy imports: submodules are available via nobrainer.io, nobrainer.models, etc.\n# but are not eagerly loaded to avoid requiring optional dependencies (monai,\n# pyro-ppl, pytorch-lightning) at import time.\n" }, { "path": "nobrainer/_version.py", "content": "# This file helps to compute a version number in source trees obtained from\n# git-archive tarball (such as those provided by githubs download-from-tag\n# feature). Distribution tarballs (built by setup.py sdist) and build\n# directories (produced by setup.py build) will contain a much shorter file\n# that just contains the computed version number.\n\n# This file is released into the public domain. Generated by\n# versioneer-0.21 (https://github.com/python-versioneer/python-versioneer)\n\n\"\"\"Git implementation of _version.py.\"\"\"\n\nimport errno\nimport os\nimport re\nimport subprocess\nimport sys\nfrom typing import Callable, Dict\n\n\ndef get_keywords():\n \"\"\"Get the keywords needed to look up the version information.\"\"\"\n # these strings will be replaced by git during git-archive.\n # setup.py/versioneer.py will grep for the variable names, so they must\n # each be defined on a line of their own. _version.py will just call\n # get_keywords().\n git_refnames = \"$Format:%d$\"\n git_full = \"$Format:%H$\"\n git_date = \"$Format:%ci$\"\n keywords = {\"refnames\": git_refnames, \"full\": git_full, \"date\": git_date}\n return keywords\n\n\nclass VersioneerConfig:\n \"\"\"Container for Versioneer configuration parameters.\"\"\"\n\n\ndef get_config():\n \"\"\"Create, populate and return the VersioneerConfig() object.\"\"\"\n # these strings are filled in when 'setup.py versioneer' creates\n # _version.py\n cfg = VersioneerConfig()\n cfg.VCS = \"git\"\n cfg.style = \"pep440\"\n cfg.tag_prefix = \"\"\n cfg.parentdir_prefix = \"\"\n cfg.versionfile_source = \"nobrainer/_version.py\"\n cfg.verbose = False\n return cfg\n\n\nclass NotThisMethod(Exception):\n \"\"\"Exception raised if a method is not valid for the current scenario.\"\"\"\n\n\nLONG_VERSION_PY: Dict[str, str] = {}\nHANDLERS: Dict[str, Dict[str, Callable]] = {}\n\n\ndef register_vcs_handler(vcs, method): # decorator\n \"\"\"Create decorator to mark a method as the handler of a VCS.\"\"\"\n\n def decorate(f):\n \"\"\"Store f in HANDLERS[vcs][method].\"\"\"\n if vcs not in HANDLERS:\n HANDLERS[vcs] = {}\n HANDLERS[vcs][method] = f\n return f\n\n return decorate\n\n\ndef run_command(commands, args, cwd=None, verbose=False, hide_stderr=False, env=None):\n \"\"\"Call the given command(s).\"\"\"\n assert isinstance(commands, list)\n process = None\n for command in commands:\n try:\n dispcmd = str([command] + args)\n # remember shell=False, so use git.cmd on windows, not just git\n process = subprocess.Popen(\n [command] + args,\n cwd=cwd,\n env=env,\n stdout=subprocess.PIPE,\n stderr=(subprocess.PIPE if hide_stderr else None),\n )\n break\n except OSError:\n e = sys.exc_info()[1]\n if e.errno == errno.ENOENT:\n continue\n if verbose:\n print(\"unable to run %s\" % dispcmd)\n print(e)\n return None, None\n else:\n if verbose:\n print(\"unable to find command, tried %s\" % (commands,))\n return None, None\n stdout = process.communicate()[0].strip().decode()\n if process.returncode != 0:\n if verbose:\n print(\"unable to run %s (error)\" % dispcmd)\n print(\"stdout was %s\" % stdout)\n return None, process.returncode\n return stdout, process.returncode\n\n\ndef versions_from_parentdir(parentdir_prefix, root, verbose):\n \"\"\"Try to determine the version from the parent directory name.\n\n Source tarballs conventionally unpack into a directory that includes both\n the project name and a version string. We will also support searching up\n two directory levels for an appropriately named parent directory\n \"\"\"\n rootdirs = []\n\n for _ in range(3):\n dirname = os.path.basename(root)\n if dirname.startswith(parentdir_prefix):\n return {\n \"version\": dirname[len(parentdir_prefix) :],\n \"full-revisionid\": None,\n \"dirty\": False,\n \"error\": None,\n \"date\": None,\n }\n rootdirs.append(root)\n root = os.path.dirname(root) # up a level\n\n if verbose:\n print(\n \"Tried directories %s but none started with prefix %s\"\n % (str(rootdirs), parentdir_prefix)\n )\n raise NotThisMethod(\"rootdir doesn't start with parentdir_prefix\")\n\n\n@register_vcs_handler(\"git\", \"get_keywords\")\ndef git_get_keywords(versionfile_abs):\n \"\"\"Extract version information from the given file.\"\"\"\n # the code embedded in _version.py can just fetch the value of these\n # keywords. When used from setup.py, we don't want to import _version.py,\n # so we do it with a regexp instead. This function is not used from\n # _version.py.\n keywords = {}\n try:\n with open(versionfile_abs, \"r\") as fobj:\n for line in fobj:\n if line.strip().startswith(\"git_refnames =\"):\n mo = re.search(r'=\\s*\"(.*)\"', line)\n if mo:\n keywords[\"refnames\"] = mo.group(1)\n if line.strip().startswith(\"git_full =\"):\n mo = re.search(r'=\\s*\"(.*)\"', line)\n if mo:\n keywords[\"full\"] = mo.group(1)\n if line.strip().startswith(\"git_date =\"):\n mo = re.search(r'=\\s*\"(.*)\"', line)\n if mo:\n keywords[\"date\"] = mo.group(1)\n except OSError:\n pass\n return keywords\n\n\n@register_vcs_handler(\"git\", \"keywords\")\ndef git_versions_from_keywords(keywords, tag_prefix, verbose):\n \"\"\"Get version information from git keywords.\"\"\"\n if \"refnames\" not in keywords:\n raise NotThisMethod(\"Short version file found\")\n date = keywords.get(\"date\")\n if date is not None:\n # Use only the last line. Previous lines may contain GPG signature\n # information.\n date = date.splitlines()[-1]\n\n # git-2.2.0 added \"%cI\", which expands to an ISO-8601 -compliant\n # datestamp. However we prefer \"%ci\" (which expands to an \"ISO-8601\n # -like\" string, which we must then edit to make compliant), because\n # it's been around since git-1.5.3, and it's too difficult to\n # discover which version we're using, or to work around using an\n # older one.\n date = date.strip().replace(\" \", \"T\", 1).replace(\" \", \"\", 1)\n refnames = keywords[\"refnames\"].strip()\n if refnames.startswith(\"$Format\"):\n if verbose:\n print(\"keywords are unexpanded, not using\")\n raise NotThisMethod(\"unexpanded keywords, not a git-archive tarball\")\n refs = {r.strip() for r in refnames.strip(\"()\").split(\",\")}\n # starting in git-1.8.3, tags are listed as \"tag: foo-1.0\" instead of\n # just \"foo-1.0\". If we see a \"tag: \" prefix, prefer those.\n TAG = \"tag: \"\n tags = {r[len(TAG) :] for r in refs if r.startswith(TAG)}\n if not tags:\n # Either we're using git < 1.8.3, or there really are no tags. We use\n # a heuristic: assume all version tags have a digit. The old git %d\n # expansion behaves like git log --decorate=short and strips out the\n # refs/heads/ and refs/tags/ prefixes that would let us distinguish\n # between branches and tags. By ignoring refnames without digits, we\n # filter out many common branch names like \"release\" and\n # \"stabilization\", as well as \"HEAD\" and \"master\".\n tags = {r for r in refs if re.search(r\"\\d\", r)}\n if verbose:\n print(\"discarding '%s', no digits\" % \",\".join(refs - tags))\n if verbose:\n print(\"likely tags: %s\" % \",\".join(sorted(tags)))\n for ref in sorted(tags):\n # sorting will prefer e.g. \"2.0\" over \"2.0rc1\"\n if ref.startswith(tag_prefix):\n r = ref[len(tag_prefix) :]\n # Filter out refs that exactly match prefix or that don't start\n # with a number once the prefix is stripped (mostly a concern\n # when prefix is '')\n if not re.match(r\"\\d\", r):\n continue\n if verbose:\n print(\"picking %s\" % r)\n return {\n \"version\": r,\n \"full-revisionid\": keywords[\"full\"].strip(),\n \"dirty\": False,\n \"error\": None,\n \"date\": date,\n }\n # no suitable tags, so version is \"0+unknown\", but full hex is still there\n if verbose:\n print(\"no suitable tags, using unknown + full revision id\")\n return {\n \"version\": \"0+unknown\",\n \"full-revisionid\": keywords[\"full\"].strip(),\n \"dirty\": False,\n \"error\": \"no suitable tags\",\n \"date\": None,\n }\n\n\n@register_vcs_handler(\"git\", \"pieces_from_vcs\")\ndef git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command):\n \"\"\"Get version from 'git describe' in the root of the source tree.\n\n This only gets called if the git-archive 'subst' keywords were *not*\n expanded, and _version.py hasn't already been rewritten with a short\n version string, meaning we're inside a checked out source tree.\n \"\"\"\n GITS = [\"git\"]\n TAG_PREFIX_REGEX = \"*\"\n if sys.platform == \"win32\":\n GITS = [\"git.cmd\", \"git.exe\"]\n TAG_PREFIX_REGEX = r\"\\*\"\n\n _, rc = runner(GITS, [\"rev-parse\", \"--git-dir\"], cwd=root, hide_stderr=True)\n if rc != 0:\n if verbose:\n print(\"Directory %s not under git control\" % root)\n raise NotThisMethod(\"'git rev-parse --git-dir' returned error\")\n\n # if there is a tag matching tag_prefix, this yields TAG-NUM-gHEX[-dirty]\n # if there isn't one, this yields HEX[-dirty] (no NUM)\n describe_out, rc = runner(\n GITS,\n [\n \"describe\",\n \"--tags\",\n \"--dirty\",\n \"--always\",\n \"--long\",\n \"--match\",\n \"%s%s\" % (tag_prefix, TAG_PREFIX_REGEX),\n ],\n cwd=root,\n )\n # --long was added in git-1.5.5\n if describe_out is None:\n raise NotThisMethod(\"'git describe' failed\")\n describe_out = describe_out.strip()\n full_out, rc = runner(GITS, [\"rev-parse\", \"HEAD\"], cwd=root)\n if full_out is None:\n raise NotThisMethod(\"'git rev-parse' failed\")\n full_out = full_out.strip()\n\n pieces = {}\n pieces[\"long\"] = full_out\n pieces[\"short\"] = full_out[:7] # maybe improved later\n pieces[\"error\"] = None\n\n branch_name, rc = runner(GITS, [\"rev-parse\", \"--abbrev-ref\", \"HEAD\"], cwd=root)\n # --abbrev-ref was added in git-1.6.3\n if rc != 0 or branch_name is None:\n raise NotThisMethod(\"'git rev-parse --abbrev-ref' returned error\")\n branch_name = branch_name.strip()\n\n if branch_name == \"HEAD\":\n # If we aren't exactly on a branch, pick a branch which represents\n # the current commit. If all else fails, we are on a branchless\n # commit.\n branches, rc = runner(GITS, [\"branch\", \"--contains\"], cwd=root)\n # --contains was added in git-1.5.4\n if rc != 0 or branches is None:\n raise NotThisMethod(\"'git branch --contains' returned error\")\n branches = branches.split(\"\\n\")\n\n # Remove the first line if we're running detached\n if \"(\" in branches[0]:\n branches.pop(0)\n\n # Strip off the leading \"* \" from the list of branches.\n branches = [branch[2:] for branch in branches]\n if \"master\" in branches:\n branch_name = \"master\"\n elif not branches:\n branch_name = None\n else:\n # Pick the first branch that is returned. Good or bad.\n branch_name = branches[0]\n\n pieces[\"branch\"] = branch_name\n\n # parse describe_out. It will be like TAG-NUM-gHEX[-dirty] or HEX[-dirty]\n # TAG might have hyphens.\n git_describe = describe_out\n\n # look for -dirty suffix\n dirty = git_describe.endswith(\"-dirty\")\n pieces[\"dirty\"] = dirty\n if dirty:\n git_describe = git_describe[: git_describe.rindex(\"-dirty\")]\n\n # now we have TAG-NUM-gHEX or HEX\n\n if \"-\" in git_describe:\n # TAG-NUM-gHEX\n mo = re.search(r\"^(.+)-(\\d+)-g([0-9a-f]+)$\", git_describe)\n if not mo:\n # unparsable. Maybe git-describe is misbehaving?\n pieces[\"error\"] = \"unable to parse git-describe output: '%s'\" % describe_out\n return pieces\n\n # tag\n full_tag = mo.group(1)\n if not full_tag.startswith(tag_prefix):\n if verbose:\n fmt = \"tag '%s' doesn't start with prefix '%s'\"\n print(fmt % (full_tag, tag_prefix))\n pieces[\"error\"] = \"tag '%s' doesn't start with prefix '%s'\" % (\n full_tag,\n tag_prefix,\n )\n return pieces\n pieces[\"closest-tag\"] = full_tag[len(tag_prefix) :]\n\n # distance: number of commits since tag\n pieces[\"distance\"] = int(mo.group(2))\n\n # commit: short hex revision ID\n pieces[\"short\"] = mo.group(3)\n\n else:\n # HEX: no tags\n pieces[\"closest-tag\"] = None\n count_out, rc = runner(GITS, [\"rev-list\", \"HEAD\", \"--count\"], cwd=root)\n pieces[\"distance\"] = int(count_out) # total number of commits\n\n # commit date: see ISO-8601 comment in git_versions_from_keywords()\n date = runner(GITS, [\"show\", \"-s\", \"--format=%ci\", \"HEAD\"], cwd=root)[0].strip()\n # Use only the last line. Previous lines may contain GPG signature\n # information.\n date = date.splitlines()[-1]\n pieces[\"date\"] = date.strip().replace(\" \", \"T\", 1).replace(\" \", \"\", 1)\n\n return pieces\n\n\ndef plus_or_dot(pieces):\n \"\"\"Return a + if we don't already have one, else return a .\"\"\"\n if \"+\" in pieces.get(\"closest-tag\", \"\"):\n return \".\"\n return \"+\"\n\n\ndef render_pep440(pieces):\n \"\"\"Build up version string, with post-release \"local version identifier\".\n\n Our goal: TAG[+DISTANCE.gHEX[.dirty]] . Note that if you\n get a tagged build and then dirty it, you'll get TAG+0.gHEX.dirty\n\n Exceptions:\n 1: no tags. git_describe was just HEX. 0+untagged.DISTANCE.gHEX[.dirty]\n \"\"\"\n if pieces[\"closest-tag\"]:\n rendered = pieces[\"closest-tag\"]\n if pieces[\"distance\"] or pieces[\"dirty\"]:\n rendered += plus_or_dot(pieces)\n rendered += \"%d.g%s\" % (pieces[\"distance\"], pieces[\"short\"])\n if pieces[\"dirty\"]:\n rendered += \".dirty\"\n else:\n # exception #1\n rendered = \"0+untagged.%d.g%s\" % (pieces[\"distance\"], pieces[\"short\"])\n if pieces[\"dirty\"]:\n rendered += \".dirty\"\n return rendered\n\n\ndef render_pep440_branch(pieces):\n \"\"\"TAG[[.dev0]+DISTANCE.gHEX[.dirty]] .\n\n The \".dev0\" means not master branch. Note that .dev0 sorts backwards\n (a feature branch will appear \"older\" than the master branch).\n\n Exceptions:\n 1: no tags. 0[.dev0]+untagged.DISTANCE.gHEX[.dirty]\n \"\"\"\n if pieces[\"closest-tag\"]:\n rendered = pieces[\"closest-tag\"]\n if pieces[\"distance\"] or pieces[\"dirty\"]:\n if pieces[\"branch\"] != \"master\":\n rendered += \".dev0\"\n rendered += plus_or_dot(pieces)\n rendered += \"%d.g%s\" % (pieces[\"distance\"], pieces[\"short\"])\n if pieces[\"dirty\"]:\n rendered += \".dirty\"\n else:\n # exception #1\n rendered = \"0\"\n if pieces[\"branch\"] != \"master\":\n rendered += \".dev0\"\n rendered += \"+untagged.%d.g%s\" % (pieces[\"distance\"], pieces[\"short\"])\n if pieces[\"dirty\"]:\n rendered += \".dirty\"\n return rendered\n\n\ndef pep440_split_post(ver):\n \"\"\"Split pep440 version string at the post-release segment.\n\n Returns the release segments before the post-release and the\n post-release version number (or -1 if no post-release segment is present).\n \"\"\"\n vc = str.split(ver, \".post\")\n return vc[0], int(vc[1] or 0) if len(vc) == 2 else None\n\n\ndef render_pep440_pre(pieces):\n \"\"\"TAG[.postN.devDISTANCE] -- No -dirty.\n\n Exceptions:\n 1: no tags. 0.post0.devDISTANCE\n \"\"\"\n if pieces[\"closest-tag\"]:\n if pieces[\"distance\"]:\n # update the post release segment\n tag_version, post_version = pep440_split_post(pieces[\"closest-tag\"])\n rendered = tag_version\n if post_version is not None:\n rendered += \".post%d.dev%d\" % (post_version + 1, pieces[\"distance\"])\n else:\n rendered += \".post0.dev%d\" % (pieces[\"distance\"])\n else:\n # no commits, use the tag as the version\n rendered = pieces[\"closest-tag\"]\n else:\n # exception #1\n rendered = \"0.post0.dev%d\" % pieces[\"distance\"]\n return rendered\n\n\ndef render_pep440_post(pieces):\n \"\"\"TAG[.postDISTANCE[.dev0]+gHEX] .\n\n The \".dev0\" means dirty. Note that .dev0 sorts backwards\n (a dirty tree will appear \"older\" than the corresponding clean one),\n but you shouldn't be releasing software with -dirty anyways.\n\n Exceptions:\n 1: no tags. 0.postDISTANCE[.dev0]\n \"\"\"\n if pieces[\"closest-tag\"]:\n rendered = pieces[\"closest-tag\"]\n if pieces[\"distance\"] or pieces[\"dirty\"]:\n rendered += \".post%d\" % pieces[\"distance\"]\n if pieces[\"dirty\"]:\n rendered += \".dev0\"\n rendered += plus_or_dot(pieces)\n rendered += \"g%s\" % pieces[\"short\"]\n else:\n # exception #1\n rendered = \"0.post%d\" % pieces[\"distance\"]\n if pieces[\"dirty\"]:\n rendered += \".dev0\"\n rendered += \"+g%s\" % pieces[\"short\"]\n return rendered\n\n\ndef render_pep440_post_branch(pieces):\n \"\"\"TAG[.postDISTANCE[.dev0]+gHEX[.dirty]] .\n\n The \".dev0\" means not master branch.\n\n Exceptions:\n 1: no tags. 0.postDISTANCE[.dev0]+gHEX[.dirty]\n \"\"\"\n if pieces[\"closest-tag\"]:\n rendered = pieces[\"closest-tag\"]\n if pieces[\"distance\"] or pieces[\"dirty\"]:\n rendered += \".post%d\" % pieces[\"distance\"]\n if pieces[\"branch\"] != \"master\":\n rendered += \".dev0\"\n rendered += plus_or_dot(pieces)\n rendered += \"g%s\" % pieces[\"short\"]\n if pieces[\"dirty\"]:\n rendered += \".dirty\"\n else:\n # exception #1\n rendered = \"0.post%d\" % pieces[\"distance\"]\n if pieces[\"branch\"] != \"master\":\n rendered += \".dev0\"\n rendered += \"+g%s\" % pieces[\"short\"]\n if pieces[\"dirty\"]:\n rendered += \".dirty\"\n return rendered\n\n\ndef render_pep440_old(pieces):\n \"\"\"TAG[.postDISTANCE[.dev0]] .\n\n The \".dev0\" means dirty.\n\n Exceptions:\n 1: no tags. 0.postDISTANCE[.dev0]\n \"\"\"\n if pieces[\"closest-tag\"]:\n rendered = pieces[\"closest-tag\"]\n if pieces[\"distance\"] or pieces[\"dirty\"]:\n rendered += \".post%d\" % pieces[\"distance\"]\n if pieces[\"dirty\"]:\n rendered += \".dev0\"\n else:\n # exception #1\n rendered = \"0.post%d\" % pieces[\"distance\"]\n if pieces[\"dirty\"]:\n rendered += \".dev0\"\n return rendered\n\n\ndef render_git_describe(pieces):\n \"\"\"TAG[-DISTANCE-gHEX][-dirty].\n\n Like 'git describe --tags --dirty --always'.\n\n Exceptions:\n 1: no tags. HEX[-dirty] (note: no 'g' prefix)\n \"\"\"\n if pieces[\"closest-tag\"]:\n rendered = pieces[\"closest-tag\"]\n if pieces[\"distance\"]:\n rendered += \"-%d-g%s\" % (pieces[\"distance\"], pieces[\"short\"])\n else:\n # exception #1\n rendered = pieces[\"short\"]\n if pieces[\"dirty\"]:\n rendered += \"-dirty\"\n return rendered\n\n\ndef render_git_describe_long(pieces):\n \"\"\"TAG-DISTANCE-gHEX[-dirty].\n\n Like 'git describe --tags --dirty --always -long'.\n The distance/hash is unconditional.\n\n Exceptions:\n 1: no tags. HEX[-dirty] (note: no 'g' prefix)\n \"\"\"\n if pieces[\"closest-tag\"]:\n rendered = pieces[\"closest-tag\"]\n rendered += \"-%d-g%s\" % (pieces[\"distance\"], pieces[\"short\"])\n else:\n # exception #1\n rendered = pieces[\"short\"]\n if pieces[\"dirty\"]:\n rendered += \"-dirty\"\n return rendered\n\n\ndef render(pieces, style):\n \"\"\"Render the given version pieces into the requested style.\"\"\"\n if pieces[\"error\"]:\n return {\n \"version\": \"unknown\",\n \"full-revisionid\": pieces.get(\"long\"),\n \"dirty\": None,\n \"error\": pieces[\"error\"],\n \"date\": None,\n }\n\n if not style or style == \"default\":\n style = \"pep440\" # the default\n\n if style == \"pep440\":\n rendered = render_pep440(pieces)\n elif style == \"pep440-branch\":\n rendered = render_pep440_branch(pieces)\n elif style == \"pep440-pre\":\n rendered = render_pep440_pre(pieces)\n elif style == \"pep440-post\":\n rendered = render_pep440_post(pieces)\n elif style == \"pep440-post-branch\":\n rendered = render_pep440_post_branch(pieces)\n elif style == \"pep440-old\":\n rendered = render_pep440_old(pieces)\n elif style == \"git-describe\":\n rendered = render_git_describe(pieces)\n elif style == \"git-describe-long\":\n rendered = render_git_describe_long(pieces)\n else:\n raise ValueError(\"unknown style '%s'\" % style)\n\n return {\n \"version\": rendered,\n \"full-revisionid\": pieces[\"long\"],\n \"dirty\": pieces[\"dirty\"],\n \"error\": None,\n \"date\": pieces.get(\"date\"),\n }\n\n\ndef get_versions():\n \"\"\"Get version information or return default if unable to do so.\"\"\"\n # I am in _version.py, which lives at ROOT/VERSIONFILE_SOURCE. If we have\n # __file__, we can work backwards from there to the root. Some\n # py2exe/bbfreeze/non-CPython implementations don't do __file__, in which\n # case we can only use expanded keywords.\n\n cfg = get_config()\n verbose = cfg.verbose\n\n try:\n return git_versions_from_keywords(get_keywords(), cfg.tag_prefix, verbose)\n except NotThisMethod:\n pass\n\n try:\n root = os.path.realpath(__file__)\n # versionfile_source is the relative path from the top of the source\n # tree (where the .git directory might live) to this file. Invert\n # this to find the root from __file__.\n for _ in cfg.versionfile_source.split(\"/\"):\n root = os.path.dirname(root)\n except NameError:\n return {\n \"version\": \"0+unknown\",\n \"full-revisionid\": None,\n \"dirty\": None,\n \"error\": \"unable to find root of source tree\",\n \"date\": None,\n }\n\n try:\n pieces = git_pieces_from_vcs(cfg.tag_prefix, root, verbose)\n return render(pieces, cfg.style)\n except NotThisMethod:\n pass\n\n try:\n if cfg.parentdir_prefix:\n return versions_from_parentdir(cfg.parentdir_prefix, root, verbose)\n except NotThisMethod:\n pass\n\n return {\n \"version\": \"0+unknown\",\n \"full-revisionid\": None,\n \"dirty\": None,\n \"error\": \"unable to compute version\",\n \"date\": None,\n }\n" }, { "path": "nobrainer/augmentation/__init__.py", "content": "\"\"\"Data augmentation: transform tagging, profiles, and SynthSeg generation.\"\"\"\n\nfrom .profiles import get_augmentation_profile\nfrom .synthseg import SynthSegGenerator\nfrom .transforms import Augmentation, TrainableCompose\n\n__all__ = [\n \"Augmentation\",\n \"SynthSegGenerator\",\n \"TrainableCompose\",\n \"get_augmentation_profile\",\n]\n" }, { "path": "nobrainer/augmentation/profiles.py", "content": "\"\"\"Predefined augmentation profiles for brain imaging.\n\nEach profile returns a list of MONAI dictionary transforms wrapped with\n:class:`~nobrainer.augmentation.transforms.Augmentation` so they are\nautomatically skipped during inference.\n\nProfiles: ``\"none\"``, ``\"light\"``, ``\"standard\"``, ``\"heavy\"``.\n\nUsage::\n\n from nobrainer.augmentation.profiles import get_augmentation_profile\n\n transforms = get_augmentation_profile(\"standard\", keys=[\"image\", \"label\"])\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom .transforms import Augmentation\n\n\ndef get_augmentation_profile(\n name: str,\n keys: list[str] | None = None,\n) -> list:\n \"\"\"Return a list of augmentation transforms for the given profile.\n\n All returned transforms are wrapped with :class:`Augmentation` so\n :class:`TrainableCompose` will skip them during inference.\n\n Parameters\n ----------\n name : str\n Profile name: ``\"none\"``, ``\"light\"``, ``\"standard\"``, ``\"heavy\"``.\n keys : list of str or None\n MONAI dictionary keys (default ``[\"image\", \"label\"]``).\n\n Returns\n -------\n list\n List of ``Augmentation``-wrapped MONAI transforms.\n \"\"\"\n from monai.transforms import RandAffined, RandFlipd, RandGaussianNoised\n\n if keys is None:\n keys = [\"image\", \"label\"]\n img_keys = [k for k in keys if k == \"image\"]\n has_label = \"label\" in keys\n modes = [\"bilinear\", \"nearest\"] if has_label else [\"bilinear\"]\n\n if name == \"none\":\n return []\n\n if name == \"light\":\n return [\n Augmentation(RandFlipd(keys=keys, prob=0.5, spatial_axis=0)),\n Augmentation(RandFlipd(keys=keys, prob=0.5, spatial_axis=1)),\n Augmentation(RandFlipd(keys=keys, prob=0.5, spatial_axis=2)),\n ]\n\n if name == \"standard\":\n return [\n Augmentation(\n RandAffined(\n keys=keys,\n prob=0.5,\n rotate_range=(0.15, 0.15, 0.15),\n scale_range=(0.1, 0.1, 0.1),\n mode=modes,\n padding_mode=\"border\",\n )\n ),\n Augmentation(RandFlipd(keys=keys, prob=0.5, spatial_axis=0)),\n Augmentation(RandFlipd(keys=keys, prob=0.5, spatial_axis=1)),\n Augmentation(RandFlipd(keys=keys, prob=0.5, spatial_axis=2)),\n Augmentation(\n RandGaussianNoised(keys=img_keys, prob=0.2, mean=0.0, std=0.1)\n ),\n ]\n\n if name == \"heavy\":\n return [\n Augmentation(\n RandAffined(\n keys=keys,\n prob=0.8,\n rotate_range=(0.3, 0.3, 0.3),\n scale_range=(0.2, 0.2, 0.2),\n mode=modes,\n padding_mode=\"border\",\n )\n ),\n Augmentation(RandFlipd(keys=keys, prob=0.5, spatial_axis=0)),\n Augmentation(RandFlipd(keys=keys, prob=0.5, spatial_axis=1)),\n Augmentation(RandFlipd(keys=keys, prob=0.5, spatial_axis=2)),\n Augmentation(\n RandGaussianNoised(keys=img_keys, prob=0.5, mean=0.0, std=0.15)\n ),\n ]\n\n available = \"none, light, standard, heavy\"\n raise ValueError(f\"Unknown augmentation profile '{name}'. Available: {available}\")\n" }, { "path": "nobrainer/augmentation/synthseg.py", "content": "\"\"\"SynthSeg-style synthetic brain data generator.\n\nEnhanced implementation following Billot et al. (2023) with:\n- GMM tissue class grouping (labels grouped by tissue type)\n- Spatial augmentation (elastic deformation, rotation, scaling, flipping)\n- Resolution randomization (downsample + upsample)\n- Configurable intensity priors\n\nReference: Billot et al., \"SynthSeg: Segmentation of brain MRI scans\nof any contrast and resolution without retraining\", Medical Image Analysis, 2023.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom pathlib import Path\n\nimport nibabel as nib\nimport numpy as np\nimport torch\nimport torch.utils.data\n\n\nclass SynthSegGenerator(torch.utils.data.Dataset):\n \"\"\"SynthSeg-style synthetic brain data generator.\n\n Generates synthetic brain images from label maps with domain\n randomization for contrast-agnostic training.\n\n Parameters\n ----------\n label_maps : list of str or Path\n Paths to NIfTI label-map files (e.g., FreeSurfer aparc+aseg).\n n_samples_per_map : int\n Number of synthetic samples per label map.\n generation_classes : dict or None\n Tissue class grouping: ``{\"WM\": [2, 41], ...}``.\n Labels in the same class share one intensity distribution.\n None = use default FreeSurfer tissue classes.\n intensity_prior : tuple of float\n ``(min, max)`` bounds for sampling per-class mean intensities.\n std_prior : tuple of float\n ``(min, max)`` bounds for sampling per-class std.\n noise_std : float\n Additive Gaussian noise std.\n bias_field_std : float\n Bias field magnitude (std of polynomial coefficients).\n elastic_std : float\n Elastic deformation magnitude (0 = disabled).\n rotation_range : float\n Max rotation in degrees per axis (0 = disabled).\n scaling_bounds : float\n Max scaling fraction (e.g., 0.2 = ±20%).\n flipping : bool\n Enable random left-right flipping with label remapping.\n randomize_resolution : bool\n Simulate variable acquisition resolution.\n resolution_range : tuple of float\n ``(min_mm, max_mm)`` per-axis resolution range.\n \"\"\"\n\n def __init__(\n self,\n label_maps: list[str | Path],\n n_samples_per_map: int = 10,\n generation_classes: dict[str, list[int]] | None = None,\n intensity_prior: tuple[float, float] = (0.0, 250.0),\n std_prior: tuple[float, float] = (0.0, 35.0),\n noise_std: float = 0.1,\n bias_field_std: float = 0.7,\n elastic_std: float = 4.0,\n rotation_range: float = 15.0,\n scaling_bounds: float = 0.2,\n flipping: bool = True,\n randomize_resolution: bool = True,\n resolution_range: tuple[float, float] = (1.0, 3.0),\n seed: int | None = None,\n ) -> None:\n self.label_maps = [Path(p) for p in label_maps]\n self._seed = seed\n self.n_samples_per_map = n_samples_per_map\n self.intensity_prior = intensity_prior\n self.std_prior = std_prior\n self.noise_std = noise_std\n self.bias_field_std = bias_field_std\n self.elastic_std = elastic_std\n self.rotation_range = rotation_range\n self.scaling_bounds = scaling_bounds\n self.flipping = flipping\n self.randomize_resolution = randomize_resolution\n self.resolution_range = resolution_range\n\n # Load tissue class mapping\n if generation_classes is None:\n from nobrainer.data.tissue_classes import FREESURFER_TISSUE_CLASSES\n\n self.generation_classes = FREESURFER_TISSUE_CLASSES\n else:\n self.generation_classes = generation_classes\n\n # Build reverse lookup: label_id → class_name\n self._label_to_class: dict[int, str] = {}\n for cls_name, label_ids in self.generation_classes.items():\n for lid in label_ids:\n self._label_to_class[lid] = cls_name\n\n def __len__(self) -> int:\n return len(self.label_maps) * self.n_samples_per_map\n\n def _get_rng(self, idx: int) -> np.random.Generator:\n \"\"\"Get a seeded RNG for reproducibility, or unseeded if no seed.\"\"\"\n if self._seed is not None:\n return np.random.default_rng(self._seed + idx)\n return np.random.default_rng()\n\n def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:\n map_idx = idx // self.n_samples_per_map\n label_path = self.label_maps[map_idx]\n\n # Load label map\n label_data = np.asarray(nib.load(label_path).dataobj, dtype=np.int32)\n\n # 1. GMM intensity generation (per tissue class)\n image = self._generate_intensities(label_data)\n\n # 2. Spatial augmentation (elastic + affine + flip)\n if self.elastic_std > 0 or self.rotation_range > 0 or self.flipping:\n image, label_data = self._spatial_augmentation(image, label_data)\n\n # 3. Resolution randomization\n if self.randomize_resolution:\n image = self._randomize_resolution(image)\n\n # 4. Bias field\n if self.bias_field_std > 0:\n image = self._add_bias_field(image)\n\n # 5. Gaussian noise\n if self.noise_std > 0:\n image = image + np.random.normal(0, self.noise_std, image.shape).astype(\n np.float32\n )\n\n # Convert to tensors with channel dim [1, D, H, W]\n image_t = torch.from_numpy(image).float().unsqueeze(0)\n label_t = torch.from_numpy(label_data).long().unsqueeze(0)\n\n return {\"image\": image_t, \"label\": label_t}\n\n # ------------------------------------------------------------------\n # GMM intensity generation\n # ------------------------------------------------------------------\n\n def _generate_intensities(self, label_data: np.ndarray) -> np.ndarray:\n \"\"\"Generate image by sampling GMM intensities per tissue class.\"\"\"\n rng = np.random.default_rng()\n unique_labels = np.unique(label_data)\n\n # Sample one (mean, std) per tissue class\n class_params: dict[str, tuple[float, float]] = {}\n for cls_name in self.generation_classes:\n mean = rng.uniform(*self.intensity_prior)\n std = rng.uniform(*self.std_prior)\n class_params[cls_name] = (mean, std)\n\n # Fill each label region from its class distribution\n image = np.zeros_like(label_data, dtype=np.float32)\n for lab in unique_labels:\n mask = label_data == lab\n n_vox = int(mask.sum())\n if n_vox == 0:\n continue\n\n cls_name = self._label_to_class.get(lab)\n if cls_name is not None and cls_name in class_params:\n mean, std = class_params[cls_name]\n else:\n # Unknown label: sample fresh random params\n mean = rng.uniform(*self.intensity_prior)\n std = rng.uniform(*self.std_prior)\n\n image[mask] = rng.normal(mean, max(std, 1e-6), size=n_vox).astype(\n np.float32\n )\n\n return image\n\n # ------------------------------------------------------------------\n # Spatial augmentation\n # ------------------------------------------------------------------\n\n def _spatial_augmentation(\n self, image: np.ndarray, label: np.ndarray\n ) -> tuple[np.ndarray, np.ndarray]:\n \"\"\"Apply elastic deformation, affine transform, and flipping.\"\"\"\n from scipy.ndimage import map_coordinates\n\n D, H, W = image.shape\n\n # Build coordinate grid\n coords = np.mgrid[:D, :H, :W].astype(np.float32) # (3, D, H, W)\n\n # Elastic deformation: smooth random displacement field\n if self.elastic_std > 0:\n # Sample on coarse grid, smooth, then resize\n coarse_shape = (max(4, D // 8), max(4, H // 8), max(4, W // 8))\n rng = np.random.default_rng()\n for axis in range(3):\n displacement = rng.normal(0, self.elastic_std, coarse_shape).astype(\n np.float32\n )\n # Smooth\n from scipy.ndimage import gaussian_filter, zoom\n\n displacement = gaussian_filter(displacement, sigma=2.0)\n # Resize to full volume\n zoom_factors = (\n D / coarse_shape[0],\n H / coarse_shape[1],\n W / coarse_shape[2],\n )\n displacement = zoom(displacement, zoom_factors, order=1)\n # Crop/pad to exact shape if needed\n displacement = displacement[:D, :H, :W]\n coords[axis] += displacement\n\n # Affine: rotation + scaling\n if self.rotation_range > 0 or self.scaling_bounds > 0:\n center = np.array([D / 2, H / 2, W / 2])\n coords_centered = coords.reshape(3, -1) - center[:, None]\n\n # Build rotation matrix (Euler angles)\n rng = np.random.default_rng()\n angles = rng.uniform(-self.rotation_range, self.rotation_range, size=3)\n angles_rad = np.deg2rad(angles)\n Rx = _rot_x(angles_rad[0])\n Ry = _rot_y(angles_rad[1])\n Rz = _rot_z(angles_rad[2])\n R = Rz @ Ry @ Rx\n\n # Scaling\n if self.scaling_bounds > 0:\n scale = rng.uniform(\n 1 - self.scaling_bounds, 1 + self.scaling_bounds, size=3\n )\n S = np.diag(scale)\n R = R @ S\n\n coords_centered = R @ coords_centered\n coords = (coords_centered + center[:, None]).reshape(3, D, H, W)\n\n # Apply spatial transform\n image_out = map_coordinates(image, coords, order=3, mode=\"nearest\")\n label_out = map_coordinates(\n label.astype(np.float32), coords, order=0, mode=\"nearest\"\n ).astype(np.int32)\n\n # Flipping\n if self.flipping and np.random.random() > 0.5:\n image_out = np.flip(image_out, axis=2).copy() # flip W axis (L/R)\n label_out = np.flip(label_out, axis=2).copy()\n label_out = self._remap_lr_labels(label_out)\n\n return image_out.astype(np.float32), label_out\n\n @staticmethod\n def _remap_lr_labels(label: np.ndarray) -> np.ndarray:\n \"\"\"Swap left/right FreeSurfer labels after L/R flip.\"\"\"\n from nobrainer.data.tissue_classes import FREESURFER_LR_PAIRS\n\n result = label.copy()\n for left, right in FREESURFER_LR_PAIRS:\n left_mask = label == left\n right_mask = label == right\n result[left_mask] = right\n result[right_mask] = left\n return result\n\n # ------------------------------------------------------------------\n # Resolution randomization\n # ------------------------------------------------------------------\n\n def _randomize_resolution(self, image: np.ndarray) -> np.ndarray:\n \"\"\"Simulate variable MRI acquisition resolution.\"\"\"\n from scipy.ndimage import gaussian_filter, zoom\n\n rng = np.random.default_rng()\n target_res = rng.uniform(*self.resolution_range, size=3)\n\n # Downsample with anti-aliasing\n sigmas = [max(0, (r - 1) / 2) for r in target_res]\n blurred = gaussian_filter(image, sigma=sigmas)\n\n # Downsample then upsample\n down_factors = [1.0 / r for r in target_res]\n downsampled = zoom(blurred, down_factors, order=1)\n up_factors = [image.shape[i] / downsampled.shape[i] for i in range(3)]\n upsampled = zoom(downsampled, up_factors, order=1)\n\n # Ensure exact shape match\n D, H, W = image.shape\n return upsampled[:D, :H, :W].astype(np.float32)\n\n # ------------------------------------------------------------------\n # Bias field\n # ------------------------------------------------------------------\n\n def _add_bias_field(self, image: np.ndarray) -> np.ndarray:\n \"\"\"Apply smooth multiplicative bias field.\"\"\"\n D, H, W = image.shape\n order = 3\n\n coords_d = np.linspace(-1, 1, D)\n coords_h = np.linspace(-1, 1, H)\n coords_w = np.linspace(-1, 1, W)\n\n rng = np.random.default_rng()\n coeffs = rng.normal(0, self.bias_field_std, (order + 1, order + 1, order + 1))\n\n bias = np.zeros_like(image)\n for i in range(order + 1):\n for j in range(order + 1):\n for k in range(order + 1):\n term = coeffs[i, j, k]\n term = term * np.power(coords_d, i)[:, None, None]\n term = term * np.power(coords_h, j)[None, :, None]\n term = term * np.power(coords_w, k)[None, None, :]\n bias += term\n\n bias = np.exp(bias)\n return (image * bias).astype(np.float32)\n\n\n# ------------------------------------------------------------------\n# Rotation matrix helpers\n# ------------------------------------------------------------------\n\n\ndef _rot_x(angle: float) -> np.ndarray:\n c, s = np.cos(angle), np.sin(angle)\n return np.array([[1, 0, 0], [0, c, -s], [0, s, c]])\n\n\ndef _rot_y(angle: float) -> np.ndarray:\n c, s = np.cos(angle), np.sin(angle)\n return np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]])\n\n\ndef _rot_z(angle: float) -> np.ndarray:\n c, s = np.cos(angle), np.sin(angle)\n return np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])\n" }, { "path": "nobrainer/augmentation/transforms.py", "content": "\"\"\"Augmentation tagging for MONAI transform pipelines.\n\nExtends MONAI's ``Compose`` so individual transforms can be tagged as\n**augmentation** (train-only) or **preprocessing** (always runs). During\ninference/prediction, augmentation-tagged transforms are automatically\nskipped.\n\nUsage::\n\n from nobrainer.augmentation.transforms import Augmentation, TrainableCompose\n from monai.transforms import RandAffined, RandGaussianNoised, LoadImaged\n\n pipeline = TrainableCompose([\n LoadImaged(keys=[\"image\", \"label\"]), # preprocessing\n Augmentation(RandAffined(keys=[\"image\", \"label\"], ...)), # train-only\n Augmentation(RandGaussianNoised(keys=[\"image\"], ...)), # train-only\n ])\n\n # Training: all transforms run\n result = pipeline(data, mode=\"train\")\n\n # Predict: augmentation transforms are skipped\n result = pipeline(data, mode=\"predict\")\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom typing import Any\n\nfrom monai.transforms import Compose\n\n\nclass Augmentation:\n \"\"\"Wrapper that tags a MONAI transform as train-only (augmentation).\n\n When used inside a :class:`TrainableCompose`, this transform is\n automatically skipped when ``mode=\"predict\"``.\n\n Parameters\n ----------\n transform : callable\n Any MONAI dictionary transform.\n \"\"\"\n\n is_augmentation = True\n\n def __init__(self, transform: Any) -> None:\n self.transform = transform\n\n def __call__(self, data: Any) -> Any:\n return self.transform(data)\n\n def __repr__(self) -> str:\n return f\"Augmentation({self.transform!r})\"\n\n\nclass TrainableCompose(Compose):\n \"\"\"MONAI Compose that skips augmentation-tagged transforms in predict mode.\n\n Behaves identically to ``monai.transforms.Compose`` in train mode.\n In predict mode, any transform wrapped with :class:`Augmentation`\n (or having ``is_augmentation = True``) is skipped.\n\n Parameters\n ----------\n transforms : list\n List of MONAI transforms, optionally wrapped with :class:`Augmentation`.\n mode : str\n Default mode: ``\"train\"`` or ``\"predict\"``. Can be overridden\n per-call via ``__call__(data, mode=...)``.\n \"\"\"\n\n def __init__(self, transforms: list, mode: str = \"train\") -> None:\n super().__init__(transforms)\n self._mode = mode\n\n @property\n def mode(self) -> str:\n return self._mode\n\n @mode.setter\n def mode(self, value: str) -> None:\n if value not in (\"train\", \"predict\"):\n raise ValueError(f\"mode must be 'train' or 'predict', got '{value}'\")\n self._mode = value\n\n def __call__(self, data: Any, mode: str | None = None, **kwargs) -> Any:\n \"\"\"Apply transforms, skipping augmentation in predict mode.\n\n Extra keyword arguments (e.g., ``end``, ``threading``) are passed\n through to MONAI's ``Compose.__call__`` for CacheDataset compat.\n \"\"\"\n active_mode = mode or self._mode\n\n if active_mode == \"train\":\n # All transforms run — pass through MONAI kwargs\n return super().__call__(data, **kwargs)\n\n # Predict mode: skip augmentation transforms\n result = data\n for t in self.transforms:\n if getattr(t, \"is_augmentation\", False):\n continue\n result = t(result)\n return result\n" }, { "path": "nobrainer/cli/__init__.py", "content": "" }, { "path": "nobrainer/cli/main.py", "content": "\"\"\"Main command-line interface for nobrainer.\"\"\"\n\nfrom __future__ import annotations\n\nimport datetime\nimport os\nimport platform\nimport sys\n\nimport click\nimport nibabel as nib\nimport numpy as np\nimport torch\n\nfrom .. import __version__\nfrom ..prediction import predict as _predict\nfrom ..training import get_device\n\n_option_kwds = {\"show_default\": True}\n\n\nclass JSONParamType(click.ParamType):\n name = \"json\"\n\n def convert(self, value, param, ctx):\n try:\n import json\n\n return json.loads(value)\n except Exception:\n self.fail(f\"{value} is not valid JSON\", param, ctx)\n\n\n@click.group()\n@click.version_option(__version__, message=\"%(prog)s version %(version)s\")\ndef cli():\n \"\"\"A framework for developing neural network models for 3D image processing.\"\"\"\n return\n\n\n@cli.command()\n@click.argument(\"infile\")\n@click.argument(\"outfile\")\n@click.option(\n \"-m\",\n \"--model\",\n type=click.Path(exists=True),\n required=True,\n help=\"Path to PyTorch model file (.pth) or model name.\",\n **_option_kwds,\n)\n@click.option(\n \"--model-type\",\n default=\"unet\",\n help=(\n \"Model architecture: unet, vnet, attention_unet, unetr, meshnet, \"\n \"highresnet, bayesian_vnet, bayesian_meshnet.\"\n ),\n **_option_kwds,\n)\n@click.option(\n \"--n-classes\",\n type=int,\n default=1,\n help=\"Number of output classes.\",\n **_option_kwds,\n)\n@click.option(\n \"--in-channels\",\n type=int,\n default=1,\n help=\"Number of input channels.\",\n **_option_kwds,\n)\n@click.option(\n \"-b\",\n \"--block-shape\",\n default=(128, 128, 128),\n type=int,\n nargs=3,\n help=\"Shape of sub-volumes on which to predict.\",\n **_option_kwds,\n)\n@click.option(\n \"--batch-size\",\n type=int,\n default=4,\n help=\"Number of blocks to process per forward pass.\",\n **_option_kwds,\n)\n@click.option(\n \"--n-samples\",\n type=int,\n default=1,\n help=\"Monte-Carlo samples for Bayesian uncertainty estimation (>1 enables MC-Dropout).\",\n **_option_kwds,\n)\n@click.option(\n \"--device\",\n default=\"auto\",\n help='Compute device: \"auto\", \"cpu\", \"cuda\", \"cuda:0\", …',\n **_option_kwds,\n)\n@click.option(\n \"-v\", \"--verbose\", is_flag=True, help=\"Print progress messages.\", **_option_kwds\n)\ndef predict(\n *,\n infile,\n outfile,\n model,\n model_type,\n n_classes,\n in_channels,\n block_shape,\n batch_size,\n n_samples,\n device,\n verbose,\n):\n \"\"\"Predict labels from a NIfTI volume using a trained PyTorch model.\n\n The predictions are saved to OUTFILE.\n \"\"\"\n if os.path.exists(outfile):\n raise FileExistsError(f\"Output file already exists: {outfile}\")\n\n # Resolve device\n if device == \"auto\":\n _device = get_device()\n else:\n _device = torch.device(device)\n\n if verbose:\n click.echo(f\"Using device: {_device}\")\n\n # Load model architecture + weights\n from ..models import get as _get_model\n\n try:\n factory = _get_model(model_type)\n pt_model = factory(n_classes=n_classes, in_channels=in_channels)\n state = torch.load(model, map_location=_device, weights_only=True)\n pt_model.load_state_dict(state, strict=False)\n except Exception as exc:\n click.echo(click.style(f\"ERROR: could not load model: {exc}\", fg=\"red\"))\n raise SystemExit(1) from exc\n\n if verbose:\n click.echo(\"Running prediction ...\")\n\n if n_samples > 1:\n from ..prediction import predict_with_uncertainty\n\n try:\n label_img, var_img, entropy_img = predict_with_uncertainty(\n infile,\n pt_model,\n n_samples=n_samples,\n block_shape=block_shape,\n batch_size=batch_size,\n device=_device,\n )\n nib.save(label_img, outfile)\n nib.save(var_img, outfile.replace(\".nii\", \"_var.nii\"))\n nib.save(entropy_img, outfile.replace(\".nii\", \"_entropy.nii\"))\n except NotImplementedError:\n click.echo(\n click.style(\n \"predict_with_uncertainty not yet implemented; \"\n \"falling back to deterministic predict()\",\n fg=\"yellow\",\n )\n )\n out_img = _predict(\n infile,\n pt_model,\n block_shape=block_shape,\n batch_size=batch_size,\n device=_device,\n )\n nib.save(out_img, outfile)\n else:\n out_img = _predict(\n infile,\n pt_model,\n block_shape=block_shape,\n batch_size=batch_size,\n device=_device,\n )\n nib.save(out_img, outfile)\n\n if verbose:\n click.echo(click.style(f\"Output saved to {outfile}\", fg=\"green\"))\n\n\n@cli.command()\n@click.option(\n \"-i\",\n \"--input\",\n \"input_paths\",\n multiple=True,\n type=click.Path(exists=True),\n required=True,\n help=\"TFRecord file(s) to convert.\",\n **_option_kwds,\n)\n@click.option(\n \"-o\",\n \"--output-dir\",\n required=True,\n type=click.Path(),\n help=\"Output directory for NIfTI or HDF5 files.\",\n **_option_kwds,\n)\n@click.option(\n \"--format\",\n \"output_format\",\n default=\"nifti\",\n type=click.Choice([\"nifti\", \"hdf5\"]),\n help=\"Output format.\",\n **_option_kwds,\n)\n@click.option(\n \"-v\", \"--verbose\", is_flag=True, help=\"Print progress messages.\", **_option_kwds\n)\ndef convert_tfrecords(*, input_paths, output_dir, output_format, verbose):\n \"\"\"Convert TFRecord files to NIfTI or HDF5 (no TensorFlow required).\"\"\"\n from ..io import convert_tfrecords as _convert\n\n if verbose:\n click.echo(f\"Converting {len(input_paths)} TFRecord file(s) …\")\n\n out_paths = _convert(\n tfrecord_paths=list(input_paths),\n output_dir=output_dir,\n output_format=output_format,\n )\n\n if verbose:\n for p in out_paths:\n click.echo(f\" → {p}\")\n click.echo(click.style(f\"Done. {len(out_paths)} files written.\", fg=\"green\"))\n\n\n@cli.command()\n@click.argument(\"output\", type=click.Path())\n@click.option(\n \"-i\",\n \"--images\",\n multiple=True,\n type=click.Path(exists=True),\n required=True,\n help=\"Image NIfTI files.\",\n **_option_kwds,\n)\n@click.option(\n \"-l\",\n \"--labels\",\n multiple=True,\n type=click.Path(exists=True),\n required=True,\n help=\"Label NIfTI files (same order as --images).\",\n **_option_kwds,\n)\n@click.option(\n \"--chunk-shape\",\n default=\"32,32,32\",\n help=\"Chunk shape (comma-separated).\",\n **_option_kwds,\n)\n@click.option(\"--no-conform\", is_flag=True, help=\"Disable auto-conforming.\")\n@click.option(\"-v\", \"--verbose\", is_flag=True, help=\"Print progress.\")\ndef convert_to_zarr(*, output, images, labels, chunk_shape, no_conform, verbose):\n \"\"\"Convert NIfTI image+label pairs to a sharded Zarr3 store.\"\"\"\n from ..datasets.zarr_store import create_zarr_store\n\n if len(images) != len(labels):\n click.echo(\n click.style(\n f\"Error: {len(images)} images but {len(labels)} labels.\", fg=\"red\"\n )\n )\n sys.exit(1)\n\n pairs = list(zip(images, labels))\n chunks = tuple(int(x) for x in chunk_shape.split(\",\"))\n\n if verbose:\n click.echo(f\"Converting {len(pairs)} pairs → {output}\")\n\n store_path = create_zarr_store(\n pairs,\n output,\n chunk_shape=chunks,\n conform=not no_conform,\n )\n click.echo(click.style(f\"Zarr store created: {store_path}\", fg=\"green\"))\n\n\n@cli.command()\ndef merge():\n \"\"\"Merge multiple models trained with variational weights.\"\"\"\n click.echo(\"Not implemented yet.\")\n sys.exit(-2)\n\n\n@cli.command()\n@click.argument(\"outfile\")\n@click.option(\n \"-m\",\n \"--model\",\n type=click.Path(exists=True),\n required=True,\n help=\"Path to model checkpoint (.ckpt) or weights (.pth).\",\n **_option_kwds,\n)\n@click.option(\n \"--model-type\",\n default=\"progressivegan\",\n type=click.Choice([\"progressivegan\", \"dcgan\"]),\n help=\"Generative model architecture.\",\n **_option_kwds,\n)\n@click.option(\n \"--latent-size\",\n type=int,\n default=512,\n help=\"Latent vector dimension.\",\n **_option_kwds,\n)\n@click.option(\n \"--n-samples\",\n type=int,\n default=1,\n help=\"Number of images to generate.\",\n **_option_kwds,\n)\n@click.option(\n \"--device\",\n default=\"auto\",\n help='Compute device: \"auto\", \"cpu\", \"cuda\", …',\n **_option_kwds,\n)\n@click.option(\n \"-v\", \"--verbose\", is_flag=True, help=\"Print progress messages.\", **_option_kwds\n)\ndef generate(\n *,\n outfile,\n model,\n model_type,\n latent_size,\n n_samples,\n device,\n verbose,\n):\n \"\"\"Generate brain volumes from a trained GAN model.\n\n Saves OUTFILE (NIfTI) for each generated sample. When ``--n-samples > 1``\n the file stem is suffixed with ``_0``, ``_1``, … before the extension.\n \"\"\"\n import os\n\n if device == \"auto\":\n _device = get_device()\n else:\n _device = torch.device(device)\n\n if verbose:\n click.echo(f\"Using device: {_device}\")\n\n from ..models import get as _get_model\n\n try:\n factory = _get_model(model_type)\n pt_model = factory(latent_size=latent_size)\n # Support both .ckpt (Lightning) and .pth (state dict)\n if model.endswith(\".ckpt\"):\n model_cls = type(pt_model)\n pt_model = model_cls.load_from_checkpoint(model, map_location=_device)\n else:\n state = torch.load(model, map_location=_device, weights_only=True)\n pt_model.load_state_dict(state, strict=False)\n except Exception as exc:\n click.echo(click.style(f\"ERROR: could not load model: {exc}\", fg=\"red\"))\n raise SystemExit(1) from exc\n\n pt_model = pt_model.to(_device)\n pt_model.eval()\n\n if verbose:\n click.echo(f\"Generating {n_samples} sample(s) …\")\n\n stem, ext = os.path.splitext(outfile)\n if ext == \".gz\":\n stem, ext2 = os.path.splitext(stem)\n ext = ext2 + ext\n\n with torch.no_grad():\n for i in range(n_samples):\n z = torch.randn(1, latent_size, device=_device)\n out = pt_model.generator(z) # (1, 1, D, H, W)\n arr = out.squeeze().cpu().numpy()\n img = nib.Nifti1Image(arr.astype(np.float32), np.eye(4))\n path = f\"{stem}_{i}{ext}\" if n_samples > 1 else outfile\n nib.save(img, path)\n if verbose:\n click.echo(f\" Saved {path}\")\n\n if verbose:\n click.echo(click.style(\"Done.\", fg=\"green\"))\n\n\n@cli.command()\n@click.option(\n \"--working-dir\",\n required=True,\n type=click.Path(),\n help=\"Directory with train script and data_manifest.json.\",\n **_option_kwds,\n)\n@click.option(\n \"--model-family\",\n default=\"bayesian_vnet\",\n help=\"Model family to use for training.\",\n **_option_kwds,\n)\n@click.option(\n \"--max-experiments\",\n type=int,\n default=10,\n help=\"Maximum number of experiments.\",\n **_option_kwds,\n)\n@click.option(\n \"--budget-hours\",\n type=float,\n default=8.0,\n help=\"Wall-clock budget in hours.\",\n **_option_kwds,\n)\n@click.option(\n \"--budget-minutes\",\n type=float,\n default=None,\n help=\"Wall-clock budget in minutes (overrides --budget-hours).\",\n **_option_kwds,\n)\n@click.option(\n \"-v\",\n \"--verbose\",\n is_flag=True,\n help=\"Print per-experiment progress.\",\n **_option_kwds,\n)\ndef research(\n *,\n working_dir,\n model_family,\n max_experiments,\n budget_hours,\n budget_minutes,\n verbose,\n):\n \"\"\"Run the autoresearch experiment loop.\n\n Proposes hyperparameter configs (via Anthropic API or random grid),\n runs training experiments, and keeps improvements.\n Writes ``run_summary.md`` in WORKING_DIR on completion.\n \"\"\"\n from ..research.loop import run_loop\n\n if verbose:\n import logging\n\n logging.basicConfig(level=logging.INFO, format=\"%(levelname)s %(message)s\")\n\n budget_seconds = None\n if budget_minutes is not None:\n budget_seconds = budget_minutes * 60\n results = run_loop(\n working_dir=working_dir,\n model_family=model_family,\n max_experiments=max_experiments,\n budget_hours=budget_hours,\n budget_seconds=budget_seconds,\n )\n\n # Progress table\n click.echo(\n f\"\\n{'run_id':>6} {'val_dice':>10} {'outcome':<12} {'failure_reason'}\"\n )\n click.echo(\"-\" * 55)\n for r in results:\n dice_str = f\"{r.val_dice:.4f}\" if r.val_dice is not None else \"—\"\n click.echo(\n f\"{r.run_id:>6} {dice_str:>10} {r.outcome:<12} {r.failure_reason or '—'}\"\n )\n\n summary_path = click.format_filename(f\"{working_dir}/run_summary.md\")\n click.echo(click.style(f\"\\nSummary written to {summary_path}\", fg=\"green\"))\n\n\n@cli.command()\n@click.option(\n \"--model-path\",\n required=True,\n type=click.Path(exists=True),\n help=\"Path to best_model.pth file.\",\n **_option_kwds,\n)\n@click.option(\n \"--config-path\",\n required=True,\n type=click.Path(exists=True),\n help=\"Path to best_config.json file.\",\n **_option_kwds,\n)\n@click.option(\n \"--trained-models-path\",\n required=True,\n type=click.Path(),\n help=\"Root of the DataLad-managed trained_models dataset.\",\n **_option_kwds,\n)\n@click.option(\n \"--model-family\",\n default=\"bayesian_vnet\",\n help=\"Model family name (used as subdirectory).\",\n **_option_kwds,\n)\n@click.option(\n \"--val-dice\",\n type=float,\n required=True,\n help=\"Validation Dice score of the best model.\",\n **_option_kwds,\n)\n@click.option(\n \"--source-run-id\",\n default=\"\",\n help=\"Run ID string for traceability.\",\n **_option_kwds,\n)\ndef commit(\n *,\n model_path,\n config_path,\n trained_models_path,\n model_family,\n val_dice,\n source_run_id,\n):\n \"\"\"Version the best model with DataLad and push to OSF.\n\n Copies model weights and config into the trained_models DataLad dataset,\n generates a model card, saves with DataLad, and pushes to OSF.\n \"\"\"\n from ..research.loop import commit_best_model\n\n try:\n result = commit_best_model(\n best_model_path=model_path,\n best_config_path=config_path,\n trained_models_path=trained_models_path,\n model_family=model_family,\n val_dice=val_dice,\n source_run_id=source_run_id,\n )\n except ImportError as exc:\n click.echo(click.style(f\"ERROR: {exc}\", fg=\"red\"))\n raise SystemExit(1) from exc\n\n click.echo(f\"Model versioned at: {result['path']}\")\n click.echo(f\"DataLad commit: {result['datalad_commit']}\")\n if result.get(\"osf_url\"):\n click.echo(click.style(f\"OSF URL: {result['osf_url']}\", fg=\"green\"))\n else:\n click.echo(click.style(\"OSF push skipped (no remote configured)\", fg=\"yellow\"))\n\n\n@cli.command()\ndef save():\n \"\"\"Save a model to PyTorch format.\"\"\"\n click.echo(\"Not implemented yet.\")\n sys.exit(-2)\n\n\n@cli.command()\ndef evaluate():\n \"\"\"Evaluate a model's predictions against known labels.\"\"\"\n click.echo(\"Not implemented yet.\")\n sys.exit(-2)\n\n\n@cli.command()\ndef info():\n \"\"\"Return information about this system.\"\"\"\n uname = platform.uname()\n cuda_available = torch.cuda.is_available()\n cuda_devices = torch.cuda.device_count() if cuda_available else 0\n s = f\"\"\"\\\nPython:\n Version: {platform.python_version()}\n Implementation: {platform.python_implementation()}\n 64-bit: {sys.maxsize > 2**32}\n Packages:\n Nobrainer: {__version__}\n Nibabel: {nib.__version__}\n Numpy: {np.__version__}\n PyTorch: {torch.__version__}\n CUDA available: {cuda_available}\n CUDA devices: {cuda_devices}\n\nSystem:\n OSType: {uname.system}\n Release: {uname.release}\n Version: {uname.version}\n Architecture: {uname.machine}\n\nTimestamp: {datetime.datetime.utcnow().strftime('%Y/%m/%d %T')}\"\"\"\n click.echo(s)\n\n\n# For debugging only.\nif __name__ == \"__main__\":\n cli()\n" }, { "path": "nobrainer/cli/tests/__init__.py", "content": "" }, { "path": "nobrainer/cli/tests/main_test.py", "content": "\"\"\"Tests for `nobrainer.cli.main`.\"\"\"\n\nimport csv\nfrom pathlib import Path\n\nfrom click.testing import CliRunner\nimport nibabel as nib\nimport numpy as np\nimport pytest\n\nfrom .. import main as climain\nfrom ...io import read_csv\nfrom ...models.meshnet import meshnet\nfrom ...models.progressivegan import progressivegan\nfrom ...utils import get_data\n\n\ndef test_convert_nonscalar_labels(tmp_path):\n runner = CliRunner()\n with runner.isolated_filesystem():\n csvpath = get_data(tmp_path)\n tfrecords_template = Path(\"data/shard-{shard:03d}.tfrecords\")\n tfrecords_template.parent.mkdir(exist_ok=True)\n args = \"\"\"\\\n convert --csv={} --tfrecords-template={} --volume-shape 256 256 256\n --examples-per-shard=2 --to-ras --no-verify-volumes\n \"\"\".format(\n csvpath, tfrecords_template\n )\n result = runner.invoke(climain.cli, args.split())\n assert result.exit_code == 0\n assert Path(\"data/shard-000.tfrecords\").is_file()\n assert Path(\"data/shard-001.tfrecords\").is_file()\n assert Path(\"data/shard-002.tfrecords\").is_file()\n assert Path(\"data/shard-003.tfrecords\").is_file()\n assert Path(\"data/shard-004.tfrecords\").is_file()\n assert not Path(\"data/shard-005.tfrecords\").is_file()\n\n\ndef test_convert_scalar_int_labels(tmp_path):\n runner = CliRunner()\n with runner.isolated_filesystem():\n csvpath = get_data(str(tmp_path))\n # Make labels scalars.\n data = [(x, 0) for (x, _) in read_csv(csvpath)]\n csvpath = tmp_path.with_suffix(\".new.csv\")\n with open(csvpath, \"w\", newline=\"\") as myfile:\n wr = csv.writer(myfile, quoting=csv.QUOTE_ALL)\n wr.writerows(data)\n tfrecords_template = Path(\"data/shard-{shard:03d}.tfrecords\")\n tfrecords_template.parent.mkdir(exist_ok=True)\n args = \"\"\"\\\n convert --csv={} --tfrecords-template={} --volume-shape 256 256 256\n --examples-per-shard=2 --to-ras --no-verify-volumes\n \"\"\".format(\n csvpath, tfrecords_template\n )\n result = runner.invoke(climain.cli, args.split())\n assert result.exit_code == 0\n assert Path(\"data/shard-000.tfrecords\").is_file()\n assert Path(\"data/shard-001.tfrecords\").is_file()\n assert Path(\"data/shard-002.tfrecords\").is_file()\n assert Path(\"data/shard-003.tfrecords\").is_file()\n assert Path(\"data/shard-004.tfrecords\").is_file()\n assert not Path(\"data/shard-005.tfrecords\").is_file()\n\n\ndef test_convert_scalar_float_labels(tmp_path):\n runner = CliRunner()\n with runner.isolated_filesystem():\n csvpath = get_data(str(tmp_path))\n # Make labels scalars.\n data = [(x, 1.0) for (x, _) in read_csv(csvpath)]\n csvpath = tmp_path.with_suffix(\".new.csv\")\n with open(csvpath, \"w\", newline=\"\") as myfile:\n wr = csv.writer(myfile, quoting=csv.QUOTE_ALL)\n wr.writerows(data)\n tfrecords_template = Path(\"data/shard-{shard:03d}.tfrecords\")\n tfrecords_template.parent.mkdir(exist_ok=True)\n args = \"\"\"\\\n convert --csv={} --tfrecords-template={} --volume-shape 256 256 256\n --examples-per-shard=2 --to-ras --no-verify-volumes\n \"\"\".format(\n csvpath, tfrecords_template\n )\n result = runner.invoke(climain.cli, args.split())\n assert result.exit_code == 0\n assert Path(\"data/shard-000.tfrecords\").is_file()\n assert Path(\"data/shard-001.tfrecords\").is_file()\n assert Path(\"data/shard-002.tfrecords\").is_file()\n assert Path(\"data/shard-003.tfrecords\").is_file()\n assert Path(\"data/shard-004.tfrecords\").is_file()\n assert not Path(\"data/shard-005.tfrecords\").is_file()\n\n\ndef test_convert_multi_resolution(tmp_path):\n runner = CliRunner()\n with runner.isolated_filesystem():\n csvpath = get_data(str(tmp_path))\n # Make labels scalars.\n data = [(x, 1.0) for (x, _) in read_csv(csvpath)]\n csvpath = tmp_path.with_suffix(\".new.csv\")\n with open(csvpath, \"w\", newline=\"\") as myfile:\n wr = csv.writer(myfile, quoting=csv.QUOTE_ALL)\n wr.writerows(data)\n tfrecords_template = Path(\"data/shard-{shard:03d}.tfrecords\")\n tfrecords_template.parent.mkdir(exist_ok=True)\n args = \"\"\"\\\n convert --csv={} --tfrecords-template={} --volume-shape 256 256 256 --start-resolution 64\n --examples-per-shard=2 --no-verify-volumes --multi-resolution\n \"\"\".format(\n csvpath, tfrecords_template\n )\n result = runner.invoke(climain.cli, args.split())\n assert result.exit_code == 0\n\n resolutions = [64, 128, 256]\n for res in resolutions:\n assert Path(\"data/shard-000-res-{:03d}.tfrecords\".format(res)).is_file()\n assert Path(\"data/shard-001-res-{:03d}.tfrecords\".format(res)).is_file()\n assert Path(\"data/shard-002-res-{:03d}.tfrecords\".format(res)).is_file()\n assert Path(\"data/shard-003-res-{:03d}.tfrecords\".format(res)).is_file()\n assert Path(\"data/shard-004-res-{:03d}.tfrecords\".format(res)).is_file()\n assert not Path(\"data/shard-005-res-{:03d}.tfrecords\".format(res)).is_file()\n\n\n@pytest.mark.xfail\ndef test_merge():\n assert False\n\n\ndef test_predict():\n runner = CliRunner()\n with runner.isolated_filesystem():\n model = meshnet(1, (10, 10, 10, 1))\n model_path = \"model.h5\"\n model.save(model_path)\n\n img_path = \"features.nii.gz\"\n nib.Nifti1Image(np.random.randn(20, 20, 20), np.eye(4)).to_filename(img_path)\n out_path = \"predictions.nii.gz\"\n\n args = \"\"\"\\\n predict --model={} --block-shape 10 10 10 --resize-features-to 20 20 20\n --largest-label --rotate-and-predict {} {}\n \"\"\".format(\n model_path, img_path, out_path\n )\n\n result = runner.invoke(climain.cli, args.split())\n assert result.exit_code == 0\n assert Path(\"predictions.nii.gz\").is_file()\n assert nib.load(out_path).shape == (20, 20, 20)\n\n\ndef test_generate():\n runner = CliRunner()\n with runner.isolated_filesystem():\n generator, _ = progressivegan(\n latent_size=256, g_fmap_base=1024, d_fmap_base=1024\n )\n resolutions = [8, 16]\n Path(\"models\").mkdir(exist_ok=True)\n for res in resolutions:\n generator.add_resolution()\n generator([np.random.random((1, 256)), 1.0]) # to build the model by a call\n model_path = \"models/generator_res_{}\".format(res)\n generator.save(model_path)\n assert Path(model_path).is_dir()\n\n out_path = \"generated.nii.gz\"\n\n args = \"\"\"\\\n generate --model {} --multi-resolution --latent-size 256 {}\n \"\"\".format(\n \"models\", out_path\n )\n result = runner.invoke(climain.cli, args.split())\n assert result.exit_code == 0\n for res in resolutions:\n assert Path(\"generated_res_{}.nii.gz\".format(res)).is_file()\n assert nib.load(\"generated_res_{}.nii.gz\".format(res)).shape == (\n res,\n res,\n res,\n )\n\n\n@pytest.mark.xfail\ndef test_save():\n assert False\n\n\n@pytest.mark.xfail\ndef test_evaluate():\n assert False\n\n\ndef test_info():\n runner = CliRunner()\n result = runner.invoke(climain.cli, [\"info\"])\n assert result.exit_code == 0\n assert \"Python\" in result.output\n assert \"System\" in result.output\n assert \"Timestamp\" in result.output\n" }, { "path": "nobrainer/dataset.py", "content": "\"\"\"PyTorch dataset utilities backed by MONAI.\"\"\"\n\nfrom __future__ import annotations\n\nfrom pathlib import Path\nfrom typing import Any, Callable\n\nfrom monai.data import CacheDataset, DataLoader\nfrom monai.transforms import (\n EnsureChannelFirstd,\n LoadImaged,\n NormalizeIntensityd,\n Orientationd,\n Spacingd,\n)\nimport numpy as np\nimport torch\n\n\ndef get_dataset(\n image_paths: list[str | Path],\n label_paths: list[str | Path] | None = None,\n block_shape: tuple[int, int, int] | None = None,\n batch_size: int = 1,\n num_workers: int = 0,\n augment: bool = False,\n binarize_labels: bool | set | Callable = False,\n target_spacing: tuple[float, float, float] = (1.0, 1.0, 1.0),\n cache_rate: float = 1.0,\n **kwargs: Any,\n) -> DataLoader:\n \"\"\"Build a MONAI-backed :class:`torch.utils.data.DataLoader`.\n\n Applies the following transform chain:\n\n ``LoadImaged → EnsureChannelFirstd → Orientationd(\"RAS\")\n → Spacingd(*target_spacing) → NormalizeIntensityd``\n → (if augment) ``RandAffined, RandFlipd, RandGaussianNoised``\n\n Parameters\n ----------\n image_paths : list\n Paths to input NIfTI volumes.\n label_paths : list or None\n Paths to corresponding label NIfTI volumes. ``None`` for\n inference-only datasets.\n block_shape : tuple or None\n If provided, spatial patch size ``(D, H, W)`` extracted by MONAI's\n ``RandSpatialCropd``. ``None`` loads full volumes.\n batch_size : int\n Number of samples per mini-batch.\n num_workers : int\n Number of DataLoader worker processes.\n augment : bool\n Whether to apply random spatial and intensity augmentations.\n target_spacing : tuple of float\n Voxel spacing (mm) to resample volumes to.\n cache_rate : float\n Fraction of dataset to cache in memory (1.0 = all).\n **kwargs\n Additional keyword arguments forwarded to :class:`DataLoader`.\n\n Returns\n -------\n DataLoader\n PyTorch DataLoader that yields batches of ``{\"image\": tensor}``\n (or ``{\"image\": tensor, \"label\": tensor}`` when labels are given).\n \"\"\"\n if label_paths is not None and len(image_paths) != len(label_paths):\n raise ValueError(\n f\"len(image_paths)={len(image_paths)} != len(label_paths)={len(label_paths)}\"\n )\n\n has_labels = label_paths is not None\n\n # Build data dicts\n if has_labels:\n data = [\n {\"image\": str(img), \"label\": str(lbl)}\n for img, lbl in zip(image_paths, label_paths)\n ]\n keys = [\"image\", \"label\"]\n else:\n data = [{\"image\": str(img)} for img in image_paths]\n keys = [\"image\"]\n\n # Core transforms — use NibabelReader to support .mgz and other formats\n transforms: list[Any] = [\n LoadImaged(keys=keys, image_only=False, reader=\"NibabelReader\"),\n EnsureChannelFirstd(keys=keys),\n Orientationd(keys=keys, axcodes=\"RAS\"),\n Spacingd(\n keys=keys,\n pixdim=target_spacing,\n mode=[\"bilinear\", \"nearest\"] if has_labels else [\"bilinear\"],\n ),\n NormalizeIntensityd(keys=[\"image\"], nonzero=True, channel_wise=True),\n ]\n\n # Optional label binarization (e.g., FreeSurfer parcellation → brain mask)\n if binarize_labels and has_labels:\n from monai.transforms import Lambdad\n\n if callable(binarize_labels) and binarize_labels is not True:\n transforms.append(Lambdad(keys=[\"label\"], func=binarize_labels))\n elif isinstance(binarize_labels, set):\n label_set = binarize_labels\n\n def _remap(x):\n import torch\n\n mask = torch.zeros_like(x)\n for val in label_set:\n mask = mask | (x == val)\n return mask.float()\n\n transforms.append(Lambdad(keys=[\"label\"], func=_remap))\n else:\n transforms.append(Lambdad(keys=[\"label\"], func=lambda x: (x > 0).float()))\n\n # Optional augmentation — supports bool or profile name\n if augment:\n from nobrainer.augmentation.profiles import get_augmentation_profile\n\n profile_name = augment if isinstance(augment, str) else \"standard\"\n aug_transforms = get_augmentation_profile(profile_name, keys=keys)\n transforms += aug_transforms\n\n if block_shape is not None:\n from monai.transforms import RandSpatialCropd\n\n transforms.append(\n RandSpatialCropd(keys=keys, roi_size=block_shape, random_size=False)\n )\n\n # Use TrainableCompose so augmentation can be skipped during predict\n from nobrainer.augmentation.transforms import TrainableCompose\n\n compose = TrainableCompose(transforms)\n\n dataset = CacheDataset(\n data=data,\n transform=compose,\n cache_rate=cache_rate,\n num_workers=max(0, num_workers),\n )\n\n return DataLoader(\n dataset,\n batch_size=batch_size,\n shuffle=True,\n num_workers=num_workers,\n pin_memory=torch.cuda.is_available(),\n **kwargs,\n )\n\n\n# ---------------------------------------------------------------------------\n# Zarr v3 dataset (requires [zarr] extras)\n# ---------------------------------------------------------------------------\n\n\nclass ZarrDataset(torch.utils.data.Dataset):\n \"\"\"PyTorch Dataset backed by Zarr v3 stores.\n\n Each item in *data_list* is a dict with ``\"image\"`` (and optionally\n ``\"label\"``) keys pointing to ``.zarr`` store paths.\n \"\"\"\n\n def __init__(\n self,\n data_list: list[dict[str, str]],\n transform: Any | None = None,\n zarr_level: int = 0,\n ):\n self.data = data_list\n self.transform = transform\n self.level = zarr_level\n\n def __len__(self) -> int:\n return len(self.data)\n\n def __getitem__(self, idx: int) -> dict:\n import zarr\n\n item = self.data[idx]\n store = zarr.open_group(str(item[\"image\"]), mode=\"r\")\n img_arr = np.asarray(store[str(self.level)]).astype(np.float32)\n\n result: dict[str, Any] = {\"image\": img_arr[None]} # add channel dim\n\n if \"label\" in item:\n lbl_store = zarr.open_group(str(item[\"label\"]), mode=\"r\")\n lbl_arr = np.asarray(lbl_store[str(self.level)]).astype(np.float32)\n result[\"label\"] = lbl_arr[None]\n\n if self.transform is not None:\n result = self.transform(result)\n\n # Convert to tensors if still numpy\n for k, v in result.items():\n if isinstance(v, np.ndarray):\n result[k] = torch.from_numpy(v)\n\n return result\n\n\ndef _is_zarr_path(path: str | Path) -> bool:\n \"\"\"Check if a path looks like a Zarr store.\"\"\"\n return str(path).rstrip(\"/\").endswith(\".zarr\")\n\n\ndef _get_zarr_dataset(\n data: list[dict[str, str]],\n batch_size: int,\n num_workers: int,\n augment: bool,\n zarr_level: int,\n **kwargs: Any,\n) -> DataLoader:\n \"\"\"Build a DataLoader from Zarr v3 stores.\"\"\"\n transform = None\n if augment:\n import monai.transforms as mt\n\n transform = mt.Compose(\n [\n mt.RandAffined(\n keys=list(data[0].keys()),\n prob=0.5,\n rotate_range=(0.1, 0.1, 0.1),\n ),\n mt.RandFlipd(keys=list(data[0].keys()), prob=0.5),\n ]\n )\n dataset = ZarrDataset(data, transform=transform, zarr_level=zarr_level)\n return DataLoader(\n dataset,\n batch_size=batch_size,\n shuffle=True,\n num_workers=num_workers,\n pin_memory=torch.cuda.is_available(),\n **kwargs,\n )\n\n\n__all__ = [\"get_dataset\", \"ZarrDataset\"]\n" }, { "path": "nobrainer/datasets/__init__.py", "content": "\"\"\"Dataset fetching utilities for various neuroimaging sources.\n\nEach submodule provides functions to install and fetch data from a\nspecific source. All require the ``[versioning]`` optional extra\n(``datalad``, ``git-annex``) unless noted otherwise.\n\nAvailable sources\n-----------------\n- :mod:`nobrainer.datasets.openneuro` — OpenNeuro raw + derivatives\n\"\"\"\n\nfrom __future__ import annotations\n\n\ndef _check_datalad():\n \"\"\"Import datalad.api, raising a clear error if not available.\"\"\"\n try:\n import datalad.api as dl\n\n return dl\n except ImportError:\n raise ImportError(\n \"DataLad is required for dataset fetching. \"\n \"Install with: pip install 'nobrainer[versioning]'\\n\"\n \"Also install git-annex: uv tool install git-annex\"\n ) from None\n\n\n__all__ = [\"openneuro\"]\n" }, { "path": "nobrainer/datasets/openneuro.py", "content": "\"\"\"Fetch datasets from OpenNeuro and OpenNeuro Derivatives via DataLad.\n\nRequires the ``[versioning]`` extra (``datalad >= 0.19``) and the\n``git-annex`` PyPI package (``uv tool install git-annex`` or\n``pip install git-annex``).\n\nExamples\n--------\nFetch fmriprep derivatives and get T1w + aparc+aseg pairs::\n\n from nobrainer.datasets.openneuro import (\n install_derivatives,\n find_subject_pairs,\n write_manifest,\n )\n\n ds_path = install_derivatives(\"ds000114\", \"/tmp/data\")\n pairs = find_subject_pairs(ds_path)\n write_manifest(pairs, \"manifest.csv\")\n\nFetch a raw OpenNeuro dataset::\n\n from nobrainer.datasets.openneuro import install_dataset\n ds_path = install_dataset(\"ds000114\", \"/tmp/data\")\n\nFetch specific files without auto-discovery::\n\n from nobrainer.datasets.openneuro import (\n install_derivatives,\n glob_dataset,\n fetch_files,\n )\n\n ds_path = install_derivatives(\"ds000114\", \"/tmp/data\")\n bold_files = glob_dataset(ds_path, \"sub-*/func/*_bold.nii.gz\")\n fetched = fetch_files(ds_path, bold_files[:5])\n\"\"\"\n\nfrom __future__ import annotations\n\nimport logging\nfrom pathlib import Path\n\nlogger = logging.getLogger(__name__)\n\n_OPENNEURO_GH = \"https://github.com/OpenNeuroDatasets\"\n_OPENNEURO_DERIV_GH = \"https://github.com/OpenNeuroDerivatives\"\n\n\ndef _dl():\n \"\"\"Lazy import of datalad.api.\"\"\"\n from nobrainer.datasets import _check_datalad\n\n return _check_datalad()\n\n\n# ---------------------------------------------------------------------------\n# Install (lightweight clone, no bulk download)\n# ---------------------------------------------------------------------------\n\n\ndef install_dataset(\n dataset_id: str,\n path: str | Path,\n) -> Path:\n \"\"\"Clone an OpenNeuro dataset (metadata only, no file content).\n\n Parameters\n ----------\n dataset_id : str\n OpenNeuro accession (e.g. ``\"ds000114\"``).\n path : str or Path\n Base directory. The dataset is cloned into\n ``/``.\n\n Returns\n -------\n Path\n Absolute path to the installed dataset directory.\n \"\"\"\n dl = _dl()\n dest = Path(path) / dataset_id\n if dest.exists():\n logger.info(\"Dataset %s already at %s\", dataset_id, dest)\n return dest.resolve()\n\n source = f\"{_OPENNEURO_GH}/{dataset_id}.git\"\n logger.info(\"Installing %s from %s\", dataset_id, source)\n dl.install(source=source, path=str(dest))\n return dest.resolve()\n\n\ndef install_derivatives(\n dataset_id: str,\n path: str | Path,\n derivative: str = \"fmriprep\",\n) -> Path:\n \"\"\"Clone an OpenNeuro Derivatives dataset (metadata only).\n\n Parameters\n ----------\n dataset_id : str\n OpenNeuro accession (e.g. ``\"ds000114\"``).\n path : str or Path\n Base directory. Cloned into ``/-``.\n derivative : str\n Pipeline name (default ``\"fmriprep\"``). Common values:\n ``\"fmriprep\"``, ``\"mriqc\"``, ``\"freesurfer\"``.\n\n Returns\n -------\n Path\n Absolute path to the installed derivative directory.\n \"\"\"\n dl = _dl()\n dest = Path(path) / f\"{dataset_id}-{derivative}\"\n if dest.exists():\n logger.info(\"Derivative %s-%s already at %s\", dataset_id, derivative, dest)\n return dest.resolve()\n\n source = f\"{_OPENNEURO_DERIV_GH}/{dataset_id}-{derivative}.git\"\n logger.info(\"Installing %s-%s from %s\", dataset_id, derivative, source)\n dl.install(source=source, path=str(dest))\n return dest.resolve()\n\n\n# ---------------------------------------------------------------------------\n# File discovery and download\n# ---------------------------------------------------------------------------\n\n\ndef glob_dataset(\n dataset_dir: str | Path,\n pattern: str,\n) -> list[Path]:\n \"\"\"Glob a DataLad dataset directory (metadata only, no download).\n\n Works on the git tree — returned paths may be git-annex symlinks\n whose content hasn't been fetched yet.\n\n Parameters\n ----------\n dataset_dir : str or Path\n Root of the DataLad dataset.\n pattern : str\n Glob pattern (e.g. ``\"sub-*/anat/*_T1w.nii.gz\"``).\n\n Returns\n -------\n list of Path\n Sorted matching paths.\n \"\"\"\n return sorted(Path(dataset_dir).glob(pattern))\n\n\ndef fetch_files(\n dataset_dir: str | Path,\n paths: list[str | Path],\n) -> list[Path]:\n \"\"\"Download specific files from a DataLad dataset.\n\n Parameters\n ----------\n dataset_dir : str or Path\n Root of the DataLad dataset.\n paths : list of str or Path\n Files to download (absolute or relative to *dataset_dir*).\n\n Returns\n -------\n list of Path\n Paths whose content was successfully downloaded.\n \"\"\"\n dl = _dl()\n dataset_dir = Path(dataset_dir)\n\n try:\n dl.get([str(p) for p in paths], dataset=str(dataset_dir))\n except Exception as exc:\n logger.warning(\"datalad get failed: %s\", exc)\n\n return [p for p in (Path(x) for x in paths) if _file_ok(p)]\n\n\n# ---------------------------------------------------------------------------\n# Paired file discovery (structural MRI)\n# ---------------------------------------------------------------------------\n\n\ndef _extract_subject_id(path: Path) -> str:\n \"\"\"Extract ``sub-XX`` from a BIDS-style path.\n\n Checks directory components first (``sub-01/anat/...``), then\n parses the filename (``sub-01_desc-preproc_T1w.nii.gz``).\n \"\"\"\n # Check directory parts (e.g. .../sub-01/anat/...)\n for part in path.parts[:-1]: # skip filename\n if part.startswith(\"sub-\"):\n return part\n # Parse from filename\n name = path.name\n if name.startswith(\"sub-\"):\n return name.split(\"_\")[0]\n return name\n\n\ndef _file_ok(p: Path) -> bool:\n \"\"\"True if *p* is a real file with nonzero size.\"\"\"\n try:\n return p.stat().st_size > 0\n except OSError:\n return False\n\n\ndef find_subject_pairs(\n dataset_dir: str | Path,\n feature_pattern: str | None = None,\n label_pattern: str | None = None,\n native_space: bool = True,\n download: bool = True,\n) -> list[dict[str, str]]:\n \"\"\"Discover and optionally download paired (feature, label) files.\n\n The default patterns find native-space preprocessed T1w images and\n aparc+aseg parcellations from fmriprep derivatives.\n\n Strategy:\n\n 1. Glob the dataset tree (git metadata only) to find label files.\n 2. For each label, find the matching feature file for the same\n subject.\n 3. Download each pair via ``datalad get``.\n 4. Verify both files are accessible before including them.\n\n Parameters\n ----------\n dataset_dir : str or Path\n Root of a DataLad dataset (typically an fmriprep derivative).\n feature_pattern : str or None\n Glob for feature files. When *None*, discovers the best\n native-space T1w pattern automatically.\n label_pattern : str or None\n Glob for label files. When *None*, tries\n ``*desc-aparcaseg_dseg.nii.gz`` then ``*desc-aseg_dseg.nii.gz``.\n native_space : bool\n Prefer native-space files (no ``space-`` token). Default True.\n download : bool\n If True (default), download each pair via ``datalad get``.\n\n Returns\n -------\n list of dict\n Each dict: ``{\"subject_id\", \"t1w_path\", \"label_path\"}``.\n \"\"\"\n dataset_dir = Path(dataset_dir)\n pairs: list[dict[str, str]] = []\n\n # --- Discover label files ---\n if label_pattern is not None:\n label_files = glob_dataset(dataset_dir, label_pattern)\n else:\n label_files = []\n for pat in [\n \"sub-*/anat/*desc-aparcaseg_dseg.nii.gz\",\n \"sub-*/anat/*desc-aseg_dseg.nii.gz\",\n ]:\n label_files = glob_dataset(dataset_dir, pat)\n if label_files:\n logger.info(\"Found %d labels matching %s\", len(label_files), pat)\n break\n\n if not label_files:\n logger.warning(\"No label files found in %s\", dataset_dir)\n return pairs\n\n # --- Match each label to a feature file ---\n for label_path in label_files:\n sub_id = _extract_subject_id(label_path)\n anat_dir = label_path.parent\n\n if feature_pattern is not None:\n feat_candidates = sorted(anat_dir.glob(feature_pattern))\n else:\n feat_candidates = [\n p\n for p in anat_dir.glob(f\"{sub_id}*desc-preproc_T1w.nii.gz\")\n if (not native_space) or (\"space-\" not in p.name)\n ]\n if not feat_candidates:\n feat_candidates = sorted(anat_dir.glob(f\"{sub_id}*_T1w.nii.gz\"))[:1]\n\n if not feat_candidates:\n logger.warning(\"No feature file for %s\", sub_id)\n continue\n\n feat_path = feat_candidates[0]\n\n if download:\n logger.info(\"Downloading pair for %s\", sub_id)\n fetch_files(dataset_dir, [feat_path, label_path])\n\n feat_ok = _file_ok(feat_path) if download else True\n label_ok = _file_ok(label_path) if download else True\n\n if feat_ok and label_ok:\n pairs.append(\n {\n \"subject_id\": sub_id,\n \"t1w_path\": str(feat_path),\n \"label_path\": str(label_path),\n }\n )\n else:\n logger.warning(\"Skipping %s: files not accessible\", sub_id)\n\n logger.info(\"Found %d paired subjects in %s\", len(pairs), dataset_dir.name)\n return pairs\n\n\n# ---------------------------------------------------------------------------\n# Manifest writing\n# ---------------------------------------------------------------------------\n\n\ndef write_manifest(\n pairs: list[dict[str, str]],\n output_path: str | Path,\n split_ratios: tuple[int, int, int] = (80, 10, 10),\n seed: int = 42,\n) -> Path:\n \"\"\"Write a manifest CSV with train/val/test split.\n\n Parameters\n ----------\n pairs : list of dict\n Each dict has ``\"subject_id\"``, ``\"t1w_path\"``, ``\"label_path\"``.\n Optionally ``\"dataset_id\"``.\n output_path : str or Path\n Destination CSV.\n split_ratios : tuple of int\n (train, val, test) percentages.\n seed : int\n Random seed for reproducible splits.\n\n Returns\n -------\n Path\n Written CSV path.\n \"\"\"\n import csv\n\n import numpy as np\n\n output_path = Path(output_path)\n output_path.parent.mkdir(parents=True, exist_ok=True)\n\n rng = np.random.default_rng(seed)\n indices = rng.permutation(len(pairs))\n total = sum(split_ratios)\n n_train = int(len(pairs) * split_ratios[0] / total)\n n_val = int(len(pairs) * split_ratios[1] / total)\n\n for i, idx in enumerate(indices):\n if i < n_train:\n pairs[idx][\"split\"] = \"train\"\n elif i < n_train + n_val:\n pairs[idx][\"split\"] = \"val\"\n else:\n pairs[idx][\"split\"] = \"test\"\n\n fieldnames = [\"subject_id\", \"dataset_id\", \"t1w_path\", \"label_path\", \"split\"]\n if not any(\"dataset_id\" in p for p in pairs):\n fieldnames = [f for f in fieldnames if f != \"dataset_id\"]\n\n with open(output_path, \"w\", newline=\"\") as f:\n writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction=\"ignore\")\n writer.writeheader()\n writer.writerows(pairs)\n\n counts = {\n s: sum(1 for p in pairs if p.get(\"split\") == s)\n for s in (\"train\", \"val\", \"test\")\n }\n logger.info(\n \"Manifest: %s — %d subjects (train=%d, val=%d, test=%d)\",\n output_path,\n len(pairs),\n counts[\"train\"],\n counts[\"val\"],\n counts[\"test\"],\n )\n return output_path\n" }, { "path": "nobrainer/datasets/zarr_store.py", "content": "\"\"\"Multi-subject Zarr3 dataset store with sharding.\n\nConverts NIfTI collections into a single sharded Zarr3 store where\nsubjects are stacked along a 4th dimension: ``images[N, D, H, W]``\nand ``labels[N, D, H, W]``. This layout enables efficient partial I/O\nfor training: reading one subject's patch is a single seek into one\nshard file.\n\nRequires the ``[zarr]`` optional extra (``zarr >= 3.0``).\n\"\"\"\n\nfrom __future__ import annotations\n\nimport json\nimport logging\nfrom pathlib import Path\nfrom typing import Any\n\nimport numpy as np\n\nlogger = logging.getLogger(__name__)\n\n\ndef _conform_volume(img, target_shape, target_voxel_size=(1.0, 1.0, 1.0)):\n \"\"\"Conform a nibabel image to target shape and voxel size.\"\"\"\n from nibabel.processing import conform\n\n return conform(img, out_shape=target_shape, voxel_size=target_voxel_size)\n\n\ndef _infer_target_shape(\n image_paths: list[str | Path],\n max_scan: int = 50,\n) -> tuple[tuple[int, int, int], tuple[float, float, float]]:\n \"\"\"Infer target shape and voxel size from input volumes.\n\n Uses the median shape and modal voxel size across a sample of volumes.\n \"\"\"\n import nibabel as nib\n\n shapes = []\n voxel_sizes = []\n for p in image_paths[:max_scan]:\n img = nib.load(p)\n shapes.append(img.shape[:3])\n voxel_sizes.append(tuple(np.abs(img.header.get_zooms()[:3])))\n\n # Median shape (rounded to nearest integer)\n median_shape = tuple(int(np.median([s[i] for s in shapes])) for i in range(3))\n\n # Modal voxel size (most common, or median if all different)\n from collections import Counter\n\n vox_counts = Counter(voxel_sizes)\n if vox_counts:\n modal_voxel = vox_counts.most_common(1)[0][0]\n else:\n modal_voxel = (1.0, 1.0, 1.0)\n\n return median_shape, modal_voxel\n\n\ndef create_zarr_store(\n image_label_pairs: list[tuple[str, str]],\n output_path: str | Path,\n subject_ids: list[str] | None = None,\n chunk_shape: tuple[int, int, int] = (32, 32, 32),\n shard_shape: tuple[int, int, int] | None = None,\n compressor: str = \"blosc\",\n conform: bool = True,\n target_shape: tuple[int, int, int] | None = None,\n target_voxel_size: tuple[float, float, float] | None = None,\n) -> Path:\n \"\"\"Convert NIfTI pairs into a single sharded Zarr3 store.\n\n When ``conform=True`` (default), volumes are conformed to a uniform\n shape so they can be stacked into 4D arrays ``images[N, D, H, W]``\n and ``labels[N, D, H, W]``. The target shape is inferred from the\n data (median shape) unless explicitly provided.\n\n Parameters\n ----------\n image_label_pairs : list of (str, str)\n List of ``(image_path, label_path)`` tuples.\n output_path : str or Path\n Output Zarr store directory.\n subject_ids : list of str or None\n Subject identifiers. If None, auto-generated as ``sub-000``, etc.\n chunk_shape : tuple of int\n Spatial chunk dimensions (default 32³).\n shard_shape : tuple of int or None\n Shard dimensions. None = auto (full array or large multiple).\n compressor : str\n Compression codec name (default ``\"blosc\"``).\n conform : bool\n Auto-conform volumes to uniform shape (default True).\n target_shape : tuple of int or None\n Target spatial shape. None = infer from data.\n target_voxel_size : tuple of float or None\n Target voxel size. None = infer from data.\n\n Returns\n -------\n Path\n Path to the created Zarr store.\n \"\"\"\n import nibabel as nib\n import zarr\n\n output_path = Path(output_path)\n n_subjects = len(image_label_pairs)\n\n if subject_ids is None:\n subject_ids = [f\"sub-{i:03d}\" for i in range(n_subjects)]\n\n if len(subject_ids) != n_subjects:\n raise ValueError(\n f\"subject_ids length ({len(subject_ids)}) != pairs ({n_subjects})\"\n )\n\n image_paths = [p[0] for p in image_label_pairs]\n\n # Infer or validate target shape\n if conform:\n if target_shape is None or target_voxel_size is None:\n inferred_shape, inferred_voxel = _infer_target_shape(image_paths)\n if target_shape is None:\n target_shape = inferred_shape\n if target_voxel_size is None:\n target_voxel_size = inferred_voxel\n logger.info(\n \"Inferred target: shape=%s, voxel_size=%s\",\n target_shape,\n target_voxel_size,\n )\n else:\n # Check all shapes are the same\n first_img = nib.load(image_paths[0])\n target_shape = first_img.shape[:3]\n for p in image_paths[1:]:\n img = nib.load(p)\n if img.shape[:3] != target_shape:\n raise ValueError(\n f\"Non-uniform shapes detected ({img.shape[:3]} vs {target_shape}). \"\n \"Use conform=True to auto-conform, or ensure all volumes match.\"\n )\n\n D, H, W = target_shape\n full_chunk = (1, *chunk_shape) # one subject per chunk along axis 0\n\n # Shard shape: group subjects into shards for balanced write parallelism\n # and read efficiency. Default: ~50 subjects per shard → manageable\n # file count while allowing parallel writes across shards.\n subjects_per_shard = 50\n if shard_shape is not None:\n full_shard = shard_shape\n else:\n full_shard = (min(subjects_per_shard, n_subjects), D, H, W)\n\n # Create store\n store = zarr.open_group(str(output_path), mode=\"w\")\n\n # Create sharded 4D arrays\n n_shards = int(np.ceil(n_subjects / full_shard[0]))\n images_arr = store.create_array(\n \"images\",\n shape=(n_subjects, D, H, W),\n chunks=full_chunk,\n shards=full_shard,\n dtype=np.float32,\n )\n labels_arr = store.create_array(\n \"labels\",\n shape=(n_subjects, D, H, W),\n chunks=full_chunk,\n shards=full_shard,\n dtype=np.int32,\n )\n logger.info(\n \"Created sharded Zarr3: shape=%s, chunks=%s, shards=%s (%d shard files)\",\n (n_subjects, D, H, W),\n full_chunk,\n full_shard,\n n_shards,\n )\n\n # Write volumes — parallel across shards.\n # Each shard is independent, so we can write to different shards\n # concurrently. Within a shard, writes are sequential.\n import concurrent.futures\n import os\n\n n_workers = min(os.cpu_count() or 1, n_shards, 8)\n\n def _write_shard_group(shard_idx):\n \"\"\"Load and write all subjects belonging to one shard.\"\"\"\n start = shard_idx * full_shard[0]\n end = min(start + full_shard[0], n_subjects)\n for i in range(start, end):\n img_path, lbl_path = image_label_pairs[i]\n img = nib.load(img_path)\n lbl = nib.load(lbl_path)\n if conform:\n img = _conform_volume(img, target_shape, target_voxel_size)\n lbl = _conform_volume(lbl, target_shape, target_voxel_size)\n images_arr[i] = np.asarray(img.dataobj, dtype=np.float32)[:D, :H, :W]\n labels_arr[i] = np.asarray(lbl.dataobj, dtype=np.int32)[:D, :H, :W]\n return end - start\n\n logger.info(\n \"Writing %d volumes across %d shards with %d workers...\",\n n_subjects,\n n_shards,\n n_workers,\n )\n with concurrent.futures.ThreadPoolExecutor(max_workers=n_workers) as pool:\n futures = [pool.submit(_write_shard_group, s) for s in range(n_shards)]\n done = 0\n for future in concurrent.futures.as_completed(futures):\n done += future.result()\n logger.info(\"Stored %d/%d volumes\", done, n_subjects)\n\n # Store metadata\n store.attrs[\"n_subjects\"] = n_subjects\n store.attrs[\"subject_ids\"] = subject_ids\n store.attrs[\"volume_shape\"] = list(target_shape)\n store.attrs[\"chunk_shape\"] = list(chunk_shape)\n store.attrs[\"layout\"] = \"stacked\"\n store.attrs[\"image_dtype\"] = \"float32\"\n store.attrs[\"label_dtype\"] = \"int32\"\n if conform:\n store.attrs[\"conformed\"] = True\n store.attrs[\"target_shape\"] = [int(x) for x in target_shape]\n store.attrs[\"target_voxel_size\"] = [float(x) for x in target_voxel_size]\n else:\n store.attrs[\"conformed\"] = False\n\n logger.info(\n \"Zarr store created: %s (%d subjects, shape=%s)\",\n output_path,\n n_subjects,\n target_shape,\n )\n return output_path.resolve()\n\n\ndef store_info(store_path: str | Path) -> dict[str, Any]:\n \"\"\"Return store metadata without reading voxel data.\n\n Parameters\n ----------\n store_path : str or Path\n Path to a Zarr store.\n\n Returns\n -------\n dict\n Store metadata including n_subjects, volume_shape, subject_ids, etc.\n \"\"\"\n import zarr\n\n store = zarr.open_group(str(store_path), mode=\"r\")\n return dict(store.attrs)\n\n\ndef create_partition(\n store_path: str | Path,\n ratios: tuple[int, int, int] = (80, 10, 10),\n seed: int = 42,\n output_path: str | Path | None = None,\n) -> Path:\n \"\"\"Generate a partition index JSON file.\n\n Parameters\n ----------\n store_path : str or Path\n Path to the Zarr store.\n ratios : tuple of int\n (train, val, test) percentages.\n seed : int\n Random seed for reproducibility.\n output_path : str or Path or None\n Output JSON path. None = ``_partition.json``.\n\n Returns\n -------\n Path\n Path to the written partition JSON file.\n \"\"\"\n info = store_info(store_path)\n subject_ids = info[\"subject_ids\"]\n n = len(subject_ids)\n\n rng = np.random.default_rng(seed)\n indices = rng.permutation(n)\n\n total = sum(ratios)\n n_train = int(n * ratios[0] / total)\n n_val = int(n * ratios[1] / total)\n\n train_ids = [subject_ids[i] for i in indices[:n_train]]\n val_ids = [subject_ids[i] for i in indices[n_train : n_train + n_val]]\n test_ids = [subject_ids[i] for i in indices[n_train + n_val :]]\n\n partition = {\n \"seed\": seed,\n \"ratios\": list(ratios),\n \"n_subjects\": n,\n \"store_path\": str(store_path),\n \"partitions\": {\n \"train\": train_ids,\n \"val\": val_ids,\n \"test\": test_ids,\n },\n }\n\n if output_path is None:\n output_path = Path(str(store_path) + \"_partition.json\")\n output_path = Path(output_path)\n\n with open(output_path, \"w\") as f:\n json.dump(partition, f, indent=2)\n\n logger.info(\n \"Partition created: %s (train=%d, val=%d, test=%d)\",\n output_path,\n len(train_ids),\n len(val_ids),\n len(test_ids),\n )\n return output_path\n\n\ndef load_partition(partition_path: str | Path) -> dict[str, list[str]]:\n \"\"\"Load a partition index and return ``{split: [subject_ids]}``.\n\n Parameters\n ----------\n partition_path : str or Path\n Path to a partition JSON file.\n\n Returns\n -------\n dict\n ``{\"train\": [...], \"val\": [...], \"test\": [...]}``.\n \"\"\"\n with open(partition_path) as f:\n data = json.load(f)\n return data[\"partitions\"]\n" }, { "path": "nobrainer/distributed_learning/dwc.py", "content": "import numpy as np\n\n# Distributed weight consolidation for Bayesian Deep Neural Networks\n# Implemented according to the:\n# McClure, Patrick, et al. Distributed weight consolidation: a brain segmentation case study.\n# Advances in neural information processing systems 31 (2018): 4093.\n\n\ndef distributed_weight_consolidation(model_weights, model_priors):\n # model_weights is a list of weights of client-models; models = [model1, model2, model3...]\n # model_priors is a list of priors of client models sames as models\n num_layers = int(len(model_weights[0]) / 2.0)\n num_datasets = np.shape(model_weights)[0]\n consolidated_model = model_weights[0]\n mean_idx = [i for i in range(0, len(model_weights[0])) if i % 2 == 0]\n std_idx = [i for i in range(0, len(model_weights[0])) if i % 2 != 0]\n ep = 1e-5\n for i in range(num_layers):\n num_1 = 0\n num_2 = 0\n den_1 = 0\n den_2 = 0\n for m in range(num_datasets):\n model = model_weights[m]\n prior = model_priors[m]\n mu_s = model[mean_idx[i]]\n mu_o = prior[mean_idx[i]]\n sig_s = model[std_idx[i]]\n sig_o = prior[std_idx[i]]\n d1 = np.power(sig_s, 2) + ep\n d2 = np.power(sig_o, 2) + ep\n num_1 = num_1 + (mu_s / d1)\n num_2 = num_2 + (mu_o / d2)\n den_1 = den_1 + (1.0 / d1)\n den_2 = den_2 + (1.0 / d2)\n consolidated_model[mean_idx[i]] = (num_1 - num_2) / (den_1 - den_2)\n consolidated_model[std_idx[i]] = 1 / (den_1 - den_2)\n return consolidated_model\n" }, { "path": "nobrainer/experiment.py", "content": "\"\"\"Experiment tracking: local file logger + optional Weights & Biases.\n\nProvides a unified interface for logging training metrics. The local\nlogger always works (writes JSON lines + CSV to the output directory).\nW&B integration is optional and auto-detected.\n\nUsage::\n\n from nobrainer.experiment import ExperimentTracker\n\n # Local-only (writes to output_dir/metrics.jsonl + metrics.csv)\n tracker = ExperimentTracker(output_dir=\"checkpoints/bvwn\", config={...})\n\n # With W&B (if wandb is installed and WANDB_API_KEY is set)\n tracker = ExperimentTracker(\n output_dir=\"checkpoints/bvwn\",\n config={\"lr\": 1e-4, \"filters\": 96},\n project=\"kwyk-reproduction\",\n tags=[\"bvwn_multi_prior\", \"50-class\"],\n )\n\n for epoch in range(epochs):\n tracker.log({\"epoch\": epoch, \"train_loss\": loss, \"val_dice\": dice})\n\n tracker.finish()\n\"\"\"\n\nfrom __future__ import annotations\n\nimport csv\nimport json\nimport logging\nimport os\nfrom pathlib import Path\nfrom typing import Any\n\nlogger = logging.getLogger(__name__)\n\n\nclass _LocalLogger:\n \"\"\"Write metrics to JSON lines + CSV in the output directory.\"\"\"\n\n def __init__(self, output_dir: Path) -> None:\n self.output_dir = Path(output_dir)\n self.output_dir.mkdir(parents=True, exist_ok=True)\n self.jsonl_path = self.output_dir / \"metrics.jsonl\"\n self.csv_path = self.output_dir / \"metrics.csv\"\n self._csv_writer = None\n self._csv_file = None\n self._fieldnames: list[str] | None = None\n\n def log(self, metrics: dict[str, Any]) -> None:\n # JSON lines (append)\n with open(self.jsonl_path, \"a\") as f:\n f.write(json.dumps(metrics, default=str) + \"\\n\")\n\n # CSV (create header on first call, append rows)\n if self._csv_writer is None:\n self._fieldnames = list(metrics.keys())\n self._csv_file = open(self.csv_path, \"w\", newline=\"\")\n self._csv_writer = csv.DictWriter(\n self._csv_file, fieldnames=self._fieldnames, extrasaction=\"ignore\"\n )\n self._csv_writer.writeheader()\n self._csv_writer.writerow(metrics)\n self._csv_file.flush()\n\n def log_config(self, config: dict[str, Any]) -> None:\n with open(self.output_dir / \"config.json\", \"w\") as f:\n json.dump(config, f, indent=2, default=str)\n\n def finish(self) -> None:\n if self._csv_file is not None:\n self._csv_file.close()\n self._csv_file = None\n self._csv_writer = None\n\n\nclass _WandbLogger:\n \"\"\"Log metrics to Weights & Biases.\"\"\"\n\n def __init__(\n self,\n config: dict[str, Any],\n project: str | None,\n name: str | None,\n tags: list[str] | None,\n ) -> None:\n import wandb\n\n self._wandb = wandb\n self._run = wandb.init(\n project=project or \"nobrainer\",\n name=name,\n config=config,\n tags=tags,\n reinit=True,\n )\n\n def log(self, metrics: dict[str, Any]) -> None:\n self._wandb.log(metrics)\n\n def log_config(self, config: dict[str, Any]) -> None:\n self._run.config.update(config, allow_val_change=True)\n\n def finish(self) -> None:\n self._wandb.finish()\n\n\nclass ExperimentTracker:\n \"\"\"Unified experiment tracker with local + optional W&B backends.\n\n The local backend always runs, writing ``metrics.jsonl``,\n ``metrics.csv``, and ``config.json`` to *output_dir*. W&B is\n activated when:\n\n 1. ``wandb`` is installed, AND\n 2. ``WANDB_API_KEY`` is set or ``use_wandb=True`` is passed.\n\n Parameters\n ----------\n output_dir : str or Path\n Directory for local metric files.\n config : dict, optional\n Hyperparameters / configuration to log.\n project : str, optional\n W&B project name (default ``\"nobrainer\"``).\n name : str, optional\n W&B run name.\n tags : list of str, optional\n W&B run tags.\n use_wandb : bool or None\n Force W&B on/off. None = auto-detect (use if installed + key set).\n \"\"\"\n\n def __init__(\n self,\n output_dir: str | Path,\n config: dict[str, Any] | None = None,\n project: str | None = None,\n name: str | None = None,\n tags: list[str] | None = None,\n use_wandb: bool | None = None,\n ) -> None:\n self._backends: list[Any] = []\n\n # Local logger (always active)\n local = _LocalLogger(Path(output_dir))\n self._backends.append(local)\n\n # Save config locally\n if config:\n local.log_config(config)\n\n # W&B (optional)\n if use_wandb is None:\n use_wandb = (\n os.environ.get(\"WANDB_API_KEY\") is not None\n or os.environ.get(\"WANDB_MODE\") == \"offline\"\n )\n if use_wandb:\n try:\n wb = _WandbLogger(\n config=config or {},\n project=project,\n name=name,\n tags=tags,\n )\n self._backends.append(wb)\n logger.info(\"W&B tracking enabled (project=%s)\", project)\n except Exception as exc:\n logger.warning(\"W&B init failed: %s — using local only\", exc)\n\n backend_names = [type(b).__name__ for b in self._backends]\n logger.info(\"Experiment tracking: %s\", \", \".join(backend_names))\n\n def log(self, metrics: dict[str, Any]) -> None:\n \"\"\"Log a dict of metrics to all backends.\"\"\"\n for backend in self._backends:\n backend.log(metrics)\n\n def log_config(self, config: dict[str, Any]) -> None:\n \"\"\"Log/update configuration to all backends.\"\"\"\n for backend in self._backends:\n backend.log_config(config)\n\n def finish(self) -> None:\n \"\"\"Finalize all backends (flush files, end W&B run).\"\"\"\n for backend in self._backends:\n backend.finish()\n\n def callback(self, **extra_fields) -> callable:\n \"\"\"Return a training callback that logs epoch metrics.\n\n The returned callable has signature ``(epoch, logs, model)`` —\n matching the callback protocol in :func:`nobrainer.training.fit`\n and :class:`Segmentation.fit`.\n\n Parameters\n ----------\n **extra_fields\n Extra key-value pairs included in every log entry (e.g.,\n ``variant=\"bvwn_multi_prior\"``).\n\n Example::\n\n tracker = ExperimentTracker(\"checkpoints/bvwn\", config={...})\n seg.fit(ds, epochs=50, callbacks=[tracker.callback(variant=\"ssd\")])\n tracker.finish()\n \"\"\"\n\n def _cb(epoch: int, logs: dict, model: Any) -> None:\n self.log({\"epoch\": epoch, **logs, **extra_fields})\n\n return _cb\n" }, { "path": "nobrainer/gpu.py", "content": "\"\"\"GPU utilities: device detection, memory profiling, batch size optimization.\n\nExamples\n--------\nAuto-select the best batch size for a model and block shape::\n\n from nobrainer.gpu import auto_batch_size, gpu_info\n\n info = gpu_info()\n print(info)\n # [{'name': 'Tesla T4', 'memory_gb': 15.1, 'id': 0}, ...]\n\n batch_size = auto_batch_size(\n model=my_model,\n block_shape=(32, 32, 32),\n n_classes=2,\n target_memory_fraction=0.85,\n )\n print(f\"Optimal batch size: {batch_size}\")\n\nScale batch size for multi-GPU::\n\n from nobrainer.gpu import scale_for_multi_gpu\n\n effective_batch, per_gpu_batch, n_gpus = scale_for_multi_gpu(\n base_batch_size=32,\n block_shape=(32, 32, 32),\n )\n # On 4x T4: effective=128, per_gpu=32, n_gpus=4\n\"\"\"\n\nfrom __future__ import annotations\n\nimport logging\nfrom typing import Any\n\nimport torch\nimport torch.nn as nn\n\nlogger = logging.getLogger(__name__)\n\n\ndef get_device() -> torch.device:\n \"\"\"Select the best available device: CUDA > MPS > CPU.\"\"\"\n if torch.cuda.is_available():\n return torch.device(\"cuda\")\n if hasattr(torch.backends, \"mps\") and torch.backends.mps.is_available():\n return torch.device(\"mps\")\n return torch.device(\"cpu\")\n\n\ndef gpu_count() -> int:\n \"\"\"Return the number of CUDA GPUs available (0 if none).\"\"\"\n if torch.cuda.is_available():\n return torch.cuda.device_count()\n return 0\n\n\ndef gpu_info() -> list[dict[str, Any]]:\n \"\"\"Return a list of dicts with GPU name, memory, and id.\n\n Returns an empty list if no CUDA GPUs are available.\n \"\"\"\n info = []\n if not torch.cuda.is_available():\n return info\n for i in range(torch.cuda.device_count()):\n props = torch.cuda.get_device_properties(i)\n info.append(\n {\n \"id\": i,\n \"name\": props.name,\n \"memory_gb\": round(props.total_memory / 1e9, 1),\n \"compute_capability\": f\"{props.major}.{props.minor}\",\n }\n )\n return info\n\n\ndef _estimate_memory_per_sample(\n model: nn.Module,\n block_shape: tuple[int, int, int],\n n_classes: int = 2,\n in_channels: int = 1,\n dtype: torch.dtype = torch.float32,\n forward_kwargs: dict | None = None,\n) -> float:\n \"\"\"Estimate GPU memory (bytes) for one training sample.\n\n Runs a forward + backward pass with batch_size=1 and measures the\n peak allocated memory. The model is moved to GPU temporarily.\n\n Parameters\n ----------\n model : nn.Module\n Model to profile.\n block_shape : tuple of int\n Spatial dimensions of one input patch.\n n_classes : int\n Number of output classes.\n in_channels : int\n Number of input channels.\n dtype : torch.dtype\n Input data type.\n\n Returns\n -------\n float\n Estimated bytes per sample (forward + backward + optimizer overhead).\n \"\"\"\n if forward_kwargs is None:\n forward_kwargs = {}\n\n if not torch.cuda.is_available():\n raise RuntimeError(\"CUDA required for memory estimation\")\n\n device = torch.device(\"cuda\")\n model = model.to(device)\n model.train()\n\n torch.cuda.reset_peak_memory_stats(device)\n torch.cuda.empty_cache()\n\n baseline = torch.cuda.memory_allocated(device)\n\n x = torch.randn(1, in_channels, *block_shape, device=device, dtype=dtype)\n labels = torch.randint(0, n_classes, (1, *block_shape), device=device)\n\n # Pass forward_kwargs if model accepts them (e.g. mc_vwn, mc_dropout)\n try:\n out = model(x, **forward_kwargs)\n except TypeError:\n out = model(x)\n loss = nn.CrossEntropyLoss()(out, labels)\n loss.backward()\n\n peak = torch.cuda.max_memory_allocated(device) - baseline\n\n # Clean up\n model.zero_grad(set_to_none=True)\n del x, labels, out, loss\n torch.cuda.empty_cache()\n model.cpu()\n\n return float(peak)\n\n\ndef auto_batch_size(\n model: nn.Module,\n block_shape: tuple[int, int, int],\n n_classes: int = 2,\n in_channels: int = 1,\n target_memory_fraction: float = 0.85,\n gpu_id: int = 0,\n min_batch: int = 1,\n max_batch: int = 512,\n forward_kwargs: dict | None = None,\n) -> int:\n \"\"\"Estimate the largest batch size that fits in GPU memory.\n\n Profiles one sample, then scales to fill ``target_memory_fraction``\n of the GPU.\n\n Parameters\n ----------\n model : nn.Module\n Model to profile (will be temporarily moved to GPU).\n block_shape : tuple of int\n Spatial dimensions ``(D, H, W)`` of one input patch.\n n_classes : int\n Number of output classes.\n in_channels : int\n Number of input channels.\n target_memory_fraction : float\n Fraction of total GPU memory to target (default 0.85).\n gpu_id : int\n Which GPU to profile.\n min_batch : int\n Minimum batch size to return.\n max_batch : int\n Maximum batch size to return.\n\n Returns\n -------\n int\n Recommended batch size for one GPU.\n \"\"\"\n if not torch.cuda.is_available():\n logger.warning(\"No CUDA — returning min_batch=%d\", min_batch)\n return min_batch\n\n total_mem = torch.cuda.get_device_properties(gpu_id).total_memory\n target_mem = total_mem * target_memory_fraction\n\n try:\n mem_per_sample = _estimate_memory_per_sample(\n model,\n block_shape,\n n_classes,\n in_channels,\n forward_kwargs=forward_kwargs,\n )\n except RuntimeError as e:\n logger.warning(\"Memory estimation failed: %s — returning min_batch\", e)\n return min_batch\n\n # Account for ~20% overhead (optimizer state, fragmentation)\n effective_per_sample = mem_per_sample * 1.2\n batch = int(target_mem / effective_per_sample)\n batch = max(min_batch, min(batch, max_batch))\n\n logger.info(\n \"auto_batch_size: %.1f GB total, %.1f MB/sample, \"\n \"target %.0f%% → batch_size=%d\",\n total_mem / 1e9,\n mem_per_sample / 1e6,\n target_memory_fraction * 100,\n batch,\n )\n return batch\n\n\ndef scale_for_multi_gpu(\n base_batch_size: int,\n block_shape: tuple[int, int, int] | None = None,\n model: nn.Module | None = None,\n n_classes: int = 2,\n target_memory_fraction: float = 0.85,\n) -> tuple[int, int, int]:\n \"\"\"Scale batch size for multi-GPU training.\n\n If ``model`` is provided, uses :func:`auto_batch_size` to determine\n the per-GPU batch size based on actual memory profiling. Otherwise,\n divides ``base_batch_size`` evenly across available GPUs.\n\n Parameters\n ----------\n base_batch_size : int\n Desired effective (global) batch size.\n block_shape : tuple of int, optional\n Spatial dimensions for memory profiling.\n model : nn.Module, optional\n Model for memory profiling. If None, uses simple division.\n n_classes : int\n Number of output classes (for profiling).\n target_memory_fraction : float\n Target GPU memory fraction (for profiling).\n\n Returns\n -------\n effective_batch : int\n Total batch size across all GPUs.\n per_gpu_batch : int\n Batch size per GPU.\n n_gpus : int\n Number of GPUs to use.\n \"\"\"\n n_gpus = gpu_count()\n if n_gpus == 0:\n return base_batch_size, base_batch_size, 0\n\n if model is not None and block_shape is not None:\n per_gpu = auto_batch_size(\n model,\n block_shape,\n n_classes=n_classes,\n target_memory_fraction=target_memory_fraction,\n )\n else:\n per_gpu = max(1, base_batch_size // n_gpus)\n\n effective = per_gpu * n_gpus\n\n logger.info(\n \"Multi-GPU scaling: %d GPUs × %d per-GPU = %d effective batch\",\n n_gpus,\n per_gpu,\n effective,\n )\n return effective, per_gpu, n_gpus\n" }, { "path": "nobrainer/io.py", "content": "\"\"\"Input/output utilities for nobrainer (PyTorch, no TensorFlow).\"\"\"\n\nfrom __future__ import annotations\n\nimport csv\nimport hashlib\nfrom pathlib import Path\nimport struct\nfrom typing import Any\n\nimport h5py\nimport nibabel as nib\nimport numpy as np\nimport torch\nimport torch.nn as nn\n\n# ---------------------------------------------------------------------------\n# CSV helpers (no TF dependency)\n# ---------------------------------------------------------------------------\n\n\ndef read_csv(\n filepath: str | Path, skip_header: bool = True, delimiter: str = \",\"\n) -> list:\n \"\"\"Return list of tuples from a CSV file.\"\"\"\n with open(filepath, newline=\"\") as f:\n reader = csv.reader(f, delimiter=delimiter)\n if skip_header:\n next(reader)\n return [tuple(row) for row in reader]\n\n\ndef read_mapping(\n filepath: str | Path, skip_header: bool = True, delimiter: str = \",\"\n) -> dict[str, str]:\n \"\"\"Read CSV as dict; first column → keys, second → values.\"\"\"\n rows = read_csv(filepath, skip_header=skip_header, delimiter=delimiter)\n return {row[0]: row[1] for row in rows}\n\n\n# ---------------------------------------------------------------------------\n# TFRecord conversion (T022)\n# ---------------------------------------------------------------------------\n\n\ndef _compute_sha256(path: str | Path) -> str:\n h = hashlib.sha256()\n with open(path, \"rb\") as f:\n for chunk in iter(lambda: f.read(65536), b\"\"):\n h.update(chunk)\n return h.hexdigest()\n\n\ndef _parse_tfrecord_file(path: str | Path):\n \"\"\"Yield raw TFRecord byte strings from a .tfrecord file.\n\n TFRecord format: [length:uint64][masked_crc32:uint32][data][masked_crc32:uint32]\n \"\"\"\n with open(path, \"rb\") as f:\n while True:\n header = f.read(12)\n if not header:\n break\n (length,) = struct.unpack_from(\" list[str]:\n \"\"\"Convert TFRecord files to NIfTI or HDF5.\n\n Uses the ``tfrecord`` PyPI package — no TensorFlow required.\n\n Parameters\n ----------\n tfrecord_paths : list\n Paths to ``.tfrecord`` files.\n output_dir : str or Path\n Directory where converted files are written.\n volume_shape : tuple or None\n Expected shape ``(D, H, W, C)`` of the stored arrays. Used\n to validate/reshape the parsed tensors.\n output_format : str\n ``\"nifti\"`` (writes ``.nii.gz``) or ``\"hdf5\"`` (writes ``.h5``).\n affine : ndarray or None\n 4×4 affine matrix for NIfTI files. Defaults to identity.\n verify_checksum : bool\n Compute SHA-256 of each output file after writing.\n\n Returns\n -------\n list of str\n Paths to converted output files.\n \"\"\"\n import tfrecord # noqa: F401 (optional dep)\n from tfrecord.reader import tfrecord_loader\n\n output_dir = Path(output_dir)\n output_dir.mkdir(parents=True, exist_ok=True)\n if affine is None:\n affine = np.eye(4)\n\n out_paths: list[str] = []\n for rec_path in tfrecord_paths:\n rec_path = Path(rec_path)\n loader = tfrecord_loader(\n str(rec_path),\n index_path=None,\n description={\"volume\": \"byte\", \"label\": \"byte\"},\n )\n for i, record in enumerate(loader):\n volume_bytes = record.get(\"volume\") or record.get(\"image\")\n label_bytes = record.get(\"label\")\n\n vol_arr = np.frombuffer(volume_bytes, dtype=np.float32)\n if volume_shape is not None:\n vol_arr = vol_arr.reshape(volume_shape)\n\n # TF stores (D,H,W,C), PyTorch wants (C,D,H,W)\n if vol_arr.ndim == 4:\n vol_arr = np.transpose(vol_arr, (3, 0, 1, 2))\n\n stem = rec_path.stem\n if output_format == \"hdf5\":\n out_path = output_dir / f\"{stem}_{i:04d}.h5\"\n with h5py.File(out_path, \"w\") as hf:\n hf.create_dataset(\"volume\", data=vol_arr, compression=\"gzip\")\n if label_bytes is not None:\n lbl_arr = np.frombuffer(label_bytes, dtype=np.float32)\n if volume_shape is not None:\n lbl_arr = lbl_arr.reshape(volume_shape)\n if lbl_arr.ndim == 4:\n lbl_arr = np.transpose(lbl_arr, (3, 0, 1, 2))\n hf.create_dataset(\"label\", data=lbl_arr, compression=\"gzip\")\n else:\n # NIfTI: use first channel as spatial volume\n spatial = vol_arr[0] if vol_arr.ndim == 4 else vol_arr\n img = nib.Nifti1Image(spatial.astype(np.float32), affine)\n out_path = output_dir / f\"{stem}_{i:04d}.nii.gz\"\n nib.save(img, str(out_path))\n\n out_paths.append(str(out_path))\n if verify_checksum:\n _compute_sha256(out_path) # validates file integrity\n\n return out_paths\n\n\n# ---------------------------------------------------------------------------\n# Weight conversion: TF Keras H5 → PyTorch (T024)\n# ---------------------------------------------------------------------------\n\n# Mapping patterns: (keras_layer_keyword, param_suffix) → pytorch_param_name\n_CONV_MAPPING = {\n \"kernel\": \"weight\",\n \"bias\": \"bias\",\n \"gamma\": \"weight\", # BatchNorm scale\n \"beta\": \"bias\", # BatchNorm shift\n \"moving_mean\": \"running_mean\",\n \"moving_variance\": \"running_var\",\n}\n\n\ndef _keras_conv3d_to_pytorch(w: np.ndarray) -> np.ndarray:\n \"\"\"Transpose Conv3D weights from Keras (D,H,W,Cin,Cout) → PyTorch (Cout,Cin,D,H,W).\"\"\"\n if w.ndim == 5:\n return np.transpose(w, (4, 3, 0, 1, 2))\n return w\n\n\ndef convert_weights(\n h5_path: str | Path,\n pt_model: nn.Module,\n layer_mapping: dict[str, str] | None = None,\n output_path: str | Path | None = None,\n verify: bool = False,\n) -> dict[str, torch.Tensor]:\n \"\"\"Load Keras ``.h5`` weights and map them to a PyTorch model.\n\n No TensorFlow is required; weights are read directly with ``h5py``.\n\n Parameters\n ----------\n h5_path : str or Path\n Path to the Keras ``.h5`` weight file.\n pt_model : nn.Module\n Target PyTorch model whose ``state_dict`` will receive the weights.\n layer_mapping : dict or None\n ``{keras_layer_name: pytorch_submodule_name}`` mapping. When\n ``None``, an automatic heuristic attempts to match by index.\n output_path : str or Path or None\n If provided, save the converted state dict to ``.pth``.\n verify : bool\n Run a brief forward-pass verification after loading (raises if\n shapes mismatch).\n\n Returns\n -------\n dict\n The loaded (possibly partial) state dict.\n \"\"\"\n h5_path = Path(h5_path)\n state = pt_model.state_dict()\n new_state: dict[str, torch.Tensor] = {}\n\n with h5py.File(h5_path, \"r\") as hf:\n # Traverse all datasets in the H5 file\n def _collect(name: str, obj: Any) -> None:\n if not isinstance(obj, h5py.Dataset):\n return\n w = obj[()] # numpy array\n # Apply weight transposition for Conv3D kernels\n if \"kernel\" in name and w.ndim == 5:\n w = _keras_conv3d_to_pytorch(w)\n # Determine target PyTorch parameter name\n pt_name = _map_name(name, layer_mapping, state)\n if pt_name is not None and pt_name in state:\n tensor = torch.from_numpy(w.copy())\n if tensor.shape == state[pt_name].shape:\n new_state[pt_name] = tensor\n\n hf.visititems(_collect)\n\n # Load matched weights; keep existing for unmatched\n combined = {**state, **new_state}\n pt_model.load_state_dict(combined, strict=False)\n\n if output_path is not None:\n torch.save(combined, str(output_path))\n\n if verify:\n pt_model.eval()\n dummy = torch.zeros(1, 1, 32, 32, 32)\n with torch.no_grad():\n _ = pt_model(dummy)\n\n return new_state\n\n\ndef _map_name(\n h5_name: str,\n mapping: dict[str, str] | None,\n state: dict[str, torch.Tensor],\n) -> str | None:\n \"\"\"Attempt to resolve an H5 dataset path to a PyTorch state-dict key.\"\"\"\n # Simple heuristic: look for a state-dict key that contains the leaf name\n parts = h5_name.replace(\"/\", \".\").split(\".\")\n leaf = parts[-1]\n pt_leaf = _CONV_MAPPING.get(leaf, leaf)\n if mapping:\n for k, v in mapping.items():\n if k in h5_name:\n candidate = f\"{v}.{pt_leaf}\"\n if candidate in state:\n return candidate\n # Fallback: direct match\n candidate = \".\".join(parts[:-1] + [pt_leaf])\n return candidate if candidate in state else None\n\n\n# ---------------------------------------------------------------------------\n# Zarr v3 conversion (requires [zarr] extras)\n# ---------------------------------------------------------------------------\n\n\ndef nifti_to_zarr(\n input_path: str | Path,\n output_path: str | Path,\n chunk_shape: tuple[int, int, int] = (64, 64, 64),\n shard_shape: tuple[int, int, int] | None = None,\n compressor: str = \"blosc\",\n levels: int = 1,\n) -> Path:\n \"\"\"Convert a NIfTI file to a sharded Zarr v3 store.\n\n Uses ``niizarr.nii2zarr`` for NIfTI-Zarr specification compliance\n (NIfTI header, OME multiscale metadata) and adds nobrainer provenance.\n\n Parameters\n ----------\n input_path : path to .nii or .nii.gz file\n output_path : path for the output .zarr directory\n chunk_shape : inner chunk dimensions\n shard_shape : outer shard dimensions; None lets niizarr choose\n compressor : compression codec (\"blosc\" or \"zlib\")\n levels : number of resolution levels (1 = single level, -1 = auto)\n\n Returns\n -------\n Path to the created .zarr store\n \"\"\"\n import datetime\n\n import zarr\n\n import nobrainer\n\n img = nib.load(str(input_path))\n output_path = Path(output_path)\n\n try:\n import niizarr\n\n niizarr.nii2zarr(\n str(input_path),\n str(output_path),\n chunk=chunk_shape,\n shard=shard_shape,\n nb_levels=levels,\n compressor=compressor,\n zarr_version=3,\n )\n except ImportError:\n # Fallback: manual Zarr v3 creation without niizarr\n arr = np.asarray(img.dataobj, dtype=np.float32)\n clamped_chunk = tuple(min(c, s) for c, s in zip(chunk_shape, arr.shape))\n if shard_shape is None:\n eff_shard = tuple(min(c * 2, s) for c, s in zip(clamped_chunk, arr.shape))\n else:\n eff_shard = tuple(max(c, s) for c, s in zip(clamped_chunk, shard_shape))\n store = zarr.open_group(str(output_path), mode=\"w\", zarr_format=3)\n store.create_array(\"0\", data=arr, chunks=clamped_chunk, shards=eff_shard)\n if levels > 1:\n from scipy.ndimage import zoom\n\n for lvl in range(1, levels):\n factor = 1 / 2**lvl\n down = zoom(arr, factor, order=1).astype(arr.dtype)\n lc = tuple(min(c, s) for c, s in zip(clamped_chunk, down.shape))\n ls = tuple(min(c * 2, s) for c, s in zip(lc, down.shape))\n store.create_array(str(lvl), data=down, chunks=lc, shards=ls)\n store.attrs[\"nifti_affine\"] = img.affine.tolist()\n\n # Store provenance in group attrs\n store = zarr.open_group(str(output_path), mode=\"r+\")\n store.attrs[\"nobrainer_provenance\"] = {\n \"source_file\": str(Path(input_path).name),\n \"created_at\": datetime.datetime.now(datetime.timezone.utc).isoformat(),\n \"tool\": \"nobrainer.io.nifti_to_zarr\",\n \"nobrainer_version\": nobrainer.__version__,\n \"chunk_shape\": list(chunk_shape),\n \"levels\": levels,\n }\n\n return output_path\n\n\ndef zarr_to_nifti(\n input_path: str | Path,\n output_path: str | Path,\n level: int = 0,\n) -> Path:\n \"\"\"Convert a Zarr v3 store back to NIfTI.\n\n Tries ``niizarr.zarr2nii`` first for NIfTI-Zarr spec compliance.\n Falls back to reading the array + stored affine if niizarr is\n unavailable or fails.\n\n Parameters\n ----------\n input_path : path to .zarr directory\n output_path : path for the output .nii.gz file\n level : resolution level to export (0 = full resolution)\n\n Returns\n -------\n Path to the created NIfTI file\n \"\"\"\n import zarr\n\n output_path = Path(output_path)\n\n # Try niizarr first\n try:\n import niizarr\n\n img = niizarr.zarr2nii(str(input_path), level=level)\n if img.affine is not None:\n nib.save(img, str(output_path))\n return output_path\n except Exception:\n pass\n\n # Fallback: manual read\n store = zarr.open_group(str(input_path), mode=\"r\")\n arr = np.asarray(store[str(level)])\n affine = np.array(store.attrs.get(\"nifti_affine\", np.eye(4).tolist()))\n img = nib.Nifti1Image(arr.astype(np.float32), affine)\n nib.save(img, str(output_path))\n return output_path\n\n\n__all__ = [\n \"read_csv\",\n \"read_mapping\",\n \"convert_tfrecords\",\n \"convert_weights\",\n \"nifti_to_zarr\",\n \"zarr_to_nifti\",\n]\n" }, { "path": "nobrainer/layers/InstanceNorm.py", "content": "import logging\n\nfrom ..layers.groupnorm import GroupNormalization\n\n\nclass InstanceNormalization(GroupNormalization):\n \"\"\"Instance normalization layer.\n Instance Normalization is an specific case of ```GroupNormalization```since\n it normalizes all features of one channel. The Groupsize is equal to the\n channel size. Empirically, its accuracy is more stable than batch norm in a\n wide range of small batch sizes, if learning rate is adjusted linearly\n with batch sizes.\n Arguments\n axis: Integer, the axis that should be normalized.\n epsilon: Small float added to variance to avoid dividing by zero.\n center: If True, add offset of `beta` to normalized tensor.\n If False, `beta` is ignored.\n scale: If True, multiply by `gamma`.\n If False, `gamma` is not used.\n beta_initializer: Initializer for the beta weight.\n gamma_initializer: Initializer for the gamma weight.\n beta_regularizer: Optional regularizer for the beta weight.\n gamma_regularizer: Optional regularizer for the gamma weight.\n beta_constraint: Optional constraint for the beta weight.\n gamma_constraint: Optional constraint for the gamma weight.\n Input shape\n Arbitrary. Use the keyword argument `input_shape`\n (tuple of integers, does not include the samples axis)\n when using this layer as the first layer in a model.\n Output shape\n Same shape as input.\n References\n - [Instance Normalization: The Missing Ingredient for Fast Stylization]\n (https://arxiv.org/abs/1607.08022)\n \"\"\"\n\n def __init__(self, **kwargs):\n if \"groups\" in kwargs:\n logging.warning(\"The given value for groups will be overwritten.\")\n\n kwargs[\"groups\"] = -1\n super().__init__(**kwargs)\n" }, { "path": "nobrainer/layers/__init__.py", "content": "from .bernoulli_dropout import BernoulliDropout\nfrom .concrete_dropout import ConcreteDropout\nfrom .gaussian_dropout import GaussianDropout\nfrom .maxpool4d import MaxPool4D\nfrom .padding import ZeroPadding3DChannels\n\n__all__ = [\n \"BernoulliDropout\",\n \"ConcreteDropout\",\n \"GaussianDropout\",\n \"MaxPool4D\",\n \"ZeroPadding3DChannels\",\n]\n" }, { "path": "nobrainer/layers/bernoulli_dropout.py", "content": "\"\"\"Bernoulli dropout layer for PyTorch.\"\"\"\n\nimport torch\nimport torch.nn as nn\n\n\nclass BernoulliDropout(nn.Module):\n \"\"\"Bernoulli dropout layer.\n\n Multiplies input by a Bernoulli mask sampled with keep probability\n ``1 - rate``. When ``scale_during_training`` is ``True`` the output\n is rescaled by ``1 / keep_prob`` so that the expected value is\n preserved (inverted dropout). When it is ``False`` the raw Bernoulli\n mask is applied and the output is scaled by ``keep_prob`` at test\n time.\n\n Parameters\n ----------\n rate : float\n Drop probability (0 ≤ rate < 1).\n is_monte_carlo : bool\n When ``True`` the stochastic mask is applied regardless of\n ``training`` mode (enables MC-Dropout inference).\n scale_during_training : bool\n When ``True`` uses inverted dropout (scale at train time).\n When ``False`` scales at test time instead.\n seed : int or None\n Optional RNG seed (used to create a per-layer Generator).\n\n References\n ----------\n Dropout: A Simple Way to Prevent Neural Networks from Overfitting.\n N. Srivastava et al., JMLR 2014.\n \"\"\"\n\n def __init__(\n self,\n rate: float,\n is_monte_carlo: bool,\n scale_during_training: bool = True,\n seed: int | None = None,\n ) -> None:\n super().__init__()\n if not 0.0 <= rate < 1.0:\n raise ValueError(f\"rate must be in [0, 1), got {rate}\")\n self.rate = rate\n self.is_monte_carlo = is_monte_carlo\n self.scale_during_training = scale_during_training\n self.keep_prob = 1.0 - rate\n self._generator: torch.Generator | None = None\n if seed is not None:\n self._generator = torch.Generator()\n self._generator.manual_seed(seed)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n apply_mask = self.is_monte_carlo or self.training\n if apply_mask:\n mask = torch.bernoulli(\n torch.full_like(x, self.keep_prob), generator=self._generator\n )\n out = x * mask\n return out / self.keep_prob if self.scale_during_training else out\n # deterministic path\n return x if self.scale_during_training else self.keep_prob * x\n\n def extra_repr(self) -> str:\n return (\n f\"rate={self.rate}, is_monte_carlo={self.is_monte_carlo}, \"\n f\"scale_during_training={self.scale_during_training}\"\n )\n" }, { "path": "nobrainer/layers/concrete_dropout.py", "content": "\"\"\"Concrete Dropout layer for PyTorch.\"\"\"\n\nimport math\n\nimport torch\nimport torch.nn as nn\n\n\nclass ConcreteDropout(nn.Module):\n \"\"\"Concrete (relaxed Bernoulli) dropout layer.\n\n Learns a per-channel drop probability ``p_post`` end-to-end via a\n differentiable relaxation of the Bernoulli mask. A KL-divergence\n regulariser between ``p_post`` and a fixed prior ``p_prior = 0.5``\n is accumulated in ``self.kl_loss`` after each forward call.\n\n Parameters\n ----------\n in_channels : int\n Number of input channels (last dimension of the input tensor).\n is_monte_carlo : bool\n When ``True`` the stochastic concrete mask is applied regardless\n of ``training`` mode.\n temperature : float\n Temperature of the concrete distribution (lower → more binary).\n use_expectation : bool\n At test time, use ``x * p_post`` instead of the identity.\n scale_factor : float\n Normalisation factor for the KL regulariser.\n seed : int or None\n Optional RNG seed.\n\n References\n ----------\n Concrete Dropout. Y. Gal, J. Hron & A. Kendall, NeurIPS 2017.\n \"\"\"\n\n def __init__(\n self,\n in_channels: int,\n is_monte_carlo: bool = False,\n temperature: float = 0.02,\n use_expectation: bool = False,\n scale_factor: float = 1.0,\n seed: int | None = None,\n ) -> None:\n super().__init__()\n self.is_monte_carlo = is_monte_carlo\n self.temperature = temperature\n self.use_expectation = use_expectation\n self.scale_factor = scale_factor\n self._generator: torch.Generator | None = None\n if seed is not None:\n self._generator = torch.Generator()\n self._generator.manual_seed(seed)\n\n # Learnable drop probability (per channel), initialised near 0.9\n self.p_logit = nn.Parameter(torch.full((in_channels,), math.log(0.9 / 0.1)))\n # Fixed prior p = 0.5 → logit = 0\n self.register_buffer(\"p_prior\", torch.full((in_channels,), 0.5))\n self.kl_loss: torch.Tensor = torch.tensor(0.0)\n\n @property\n def p_post(self) -> torch.Tensor:\n \"\"\"Dropout probability clipped to (0.05, 0.95).\"\"\"\n return torch.sigmoid(self.p_logit).clamp(0.05, 0.95)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n apply_mask = self.is_monte_carlo or self.training\n if apply_mask:\n out = self._apply_concrete(x)\n else:\n out = x * self.p_post if self.use_expectation else x\n self.kl_loss = self._kl_divergence()\n return out\n\n def _apply_concrete(self, x: torch.Tensor) -> torch.Tensor:\n eps = torch.finfo(x.dtype).eps\n p = self.p_post # (C,)\n noise = torch.rand(\n x.shape, dtype=x.dtype, device=x.device, generator=self._generator\n ).clamp(eps, 1.0 - eps)\n z = torch.sigmoid(\n (\n torch.log(p + eps)\n - torch.log(1.0 - p + eps)\n + torch.log(noise)\n - torch.log(1.0 - noise)\n )\n / self.temperature\n )\n return x * z\n\n def _kl_divergence(self) -> torch.Tensor:\n eps = 1e-7\n p = self.p_post\n pr = self.p_prior\n kl = p * (torch.log(p + eps) - torch.log(pr + eps)) + (1 - p) * (\n torch.log(1 - p + eps) - torch.log(1 - pr + eps)\n )\n return kl.sum() / self.scale_factor\n\n def extra_repr(self) -> str:\n return (\n f\"is_monte_carlo={self.is_monte_carlo}, temperature={self.temperature}, \"\n f\"scale_factor={self.scale_factor}\"\n )\n" }, { "path": "nobrainer/layers/gaussian_dropout.py", "content": "\"\"\"Gaussian dropout layer for PyTorch.\"\"\"\n\nimport math\n\nimport torch\nimport torch.nn as nn\n\n\nclass GaussianDropout(nn.Module):\n \"\"\"Gaussian (multiplicative) dropout layer.\n\n Multiplies the input by noise sampled from ``Normal(1, σ²)`` where\n σ is derived from ``rate``. When ``scale_during_training`` is\n ``True``, σ = sqrt(rate / (1 - rate)) (variance-preserving during\n training); otherwise σ = sqrt(rate * (1 - rate)).\n\n Parameters\n ----------\n rate : float\n Drop probability (0 ≤ rate < 1).\n is_monte_carlo : bool\n When ``True``, noise is applied regardless of ``training`` mode.\n scale_during_training : bool\n Selects which σ formula is used (see above).\n seed : int or None\n Optional RNG seed.\n\n References\n ----------\n Dropout: A Simple Way to Prevent Neural Networks from Overfitting.\n N. Srivastava et al., JMLR 2014.\n \"\"\"\n\n def __init__(\n self,\n rate: float,\n is_monte_carlo: bool,\n scale_during_training: bool = True,\n seed: int | None = None,\n ) -> None:\n super().__init__()\n if not 0.0 <= rate < 1.0:\n raise ValueError(f\"rate must be in [0, 1), got {rate}\")\n self.rate = rate\n self.is_monte_carlo = is_monte_carlo\n self.scale_during_training = scale_during_training\n self._generator: torch.Generator | None = None\n if seed is not None:\n self._generator = torch.Generator()\n self._generator.manual_seed(seed)\n\n if scale_during_training:\n self._stddev = math.sqrt(rate / (1.0 - rate))\n else:\n self._stddev = math.sqrt(rate * (1.0 - rate))\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n if self.is_monte_carlo or self.training:\n noise = torch.randn_like(x, generator=self._generator) * self._stddev + 1.0\n return x * noise\n return x\n\n def extra_repr(self) -> str:\n return (\n f\"rate={self.rate}, is_monte_carlo={self.is_monte_carlo}, \"\n f\"scale_during_training={self.scale_during_training}\"\n )\n" }, { "path": "nobrainer/layers/maxpool4d.py", "content": "\"\"\"MaxPool4D layer for PyTorch.\n\nImplements 4-D max-pooling (N, C, V, D, H, W) by treating the volume\ndimension V as a batch dimension and applying ``nn.MaxPool3d`` over\n(D, H, W). This avoids the need for a custom CUDA kernel.\n\"\"\"\n\nimport torch\nimport torch.nn as nn\n\n\nclass MaxPool4D(nn.Module):\n \"\"\"Max-pooling over 4 spatial dimensions.\n\n Expects input of shape ``(N, C, V, D, H, W)`` and applies\n ``kernel_size`` / ``stride`` / ``padding`` along the last 3\n dimensions (D, H, W). The volume dimension V is reduced with\n ``pool_v`` if ``> 1``.\n\n Parameters\n ----------\n kernel_size : int or tuple\n Kernel size for the (D, H, W) axes.\n stride : int or tuple or None\n Stride; defaults to ``kernel_size``.\n padding : int or tuple\n Zero-padding added to all spatial sides.\n pool_v : int\n Max-pool kernel size along the volume (V) axis. ``1`` leaves\n V unchanged.\n \"\"\"\n\n def __init__(\n self,\n kernel_size: int | tuple[int, int, int],\n stride: int | tuple[int, int, int] | None = None,\n padding: int | tuple[int, int, int] = 0,\n pool_v: int = 1,\n ) -> None:\n super().__init__()\n self.pool3d = nn.MaxPool3d(kernel_size, stride=stride, padding=padding)\n self.pool_v = pool_v\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n if x.dim() != 6:\n raise ValueError(\n f\"MaxPool4D expects 6-D input (N, C, V, D, H, W), got {x.dim()}-D\"\n )\n N, C, V, D, H, W = x.shape\n # Merge batch and volume dims so MaxPool3d sees (N*V, C, D, H, W)\n out = self.pool3d(x.view(N * V, C, D, H, W))\n _, _, D2, H2, W2 = out.shape\n out = out.view(N, C, V, D2, H2, W2)\n if self.pool_v > 1:\n # Max over V with stride pool_v (using unfold for non-overlapping)\n out = out.unfold(2, self.pool_v, self.pool_v).amax(dim=-1)\n return out\n\n def extra_repr(self) -> str:\n return f\"pool3d={self.pool3d}, pool_v={self.pool_v}\"\n" }, { "path": "nobrainer/layers/padding.py", "content": "\"\"\"Custom padding layers for nobrainer (PyTorch).\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\nclass ZeroPadding3DChannels(nn.Module):\n \"\"\"Pad the channel dimension of a 5-D tensor symmetrically with zeros.\n\n Expects input of shape ``(N, C, D, H, W)`` and pads ``C`` by\n ``padding`` on each side, yielding ``(N, C + 2*padding, D, H, W)``.\n\n Parameters\n ----------\n padding : int\n Number of zero channels to prepend and append.\n \"\"\"\n\n def __init__(self, padding: int) -> None:\n super().__init__()\n self.padding = padding\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n # F.pad pads in reverse dim order; last two entries pad dim 1 (C)\n return F.pad(x, (0, 0, 0, 0, 0, 0, self.padding, self.padding))\n\n def extra_repr(self) -> str:\n return f\"padding={self.padding}\"\n" }, { "path": "nobrainer/layers/tests/__init__.py", "content": "" }, { "path": "nobrainer/losses.py", "content": "\"\"\"Loss functions for 3-D semantic segmentation (PyTorch / MONAI).\"\"\"\n\nfrom __future__ import annotations\n\nfrom monai.losses import DiceLoss, GeneralizedDiceLoss, TverskyLoss\nimport torch\n\n# ---------------------------------------------------------------------------\n# Convenience factory functions\n# ---------------------------------------------------------------------------\n\n\ndef dice(\n sigmoid: bool = False,\n softmax: bool = False,\n squared_pred: bool = False,\n smooth_nr: float = 1e-5,\n smooth_dr: float = 1e-5,\n **kwargs,\n) -> DiceLoss:\n \"\"\"Return a MONAI ``DiceLoss`` instance.\n\n Parameters\n ----------\n sigmoid : bool\n Apply sigmoid to predictions before computing Dice.\n softmax : bool\n Apply softmax to predictions before computing Dice.\n squared_pred : bool\n Use squared predictions in the denominator.\n smooth_nr, smooth_dr : float\n Numerator/denominator smoothing to avoid division by zero.\n **kwargs\n Extra keyword arguments forwarded to ``monai.losses.DiceLoss``.\n \"\"\"\n return DiceLoss(\n sigmoid=sigmoid,\n softmax=softmax,\n squared_pred=squared_pred,\n smooth_nr=smooth_nr,\n smooth_dr=smooth_dr,\n **kwargs,\n )\n\n\ndef generalized_dice(\n sigmoid: bool = False,\n softmax: bool = False,\n smooth_nr: float = 1e-5,\n smooth_dr: float = 1e-5,\n **kwargs,\n) -> GeneralizedDiceLoss:\n \"\"\"Return a MONAI ``GeneralizedDiceLoss`` instance.\"\"\"\n return GeneralizedDiceLoss(\n sigmoid=sigmoid,\n softmax=softmax,\n smooth_nr=smooth_nr,\n smooth_dr=smooth_dr,\n **kwargs,\n )\n\n\ndef jaccard(\n sigmoid: bool = False,\n softmax: bool = False,\n smooth_nr: float = 1e-5,\n smooth_dr: float = 1e-5,\n **kwargs,\n) -> DiceLoss:\n \"\"\"Return a Dice loss configured for Jaccard (IoU) computation.\n\n The Jaccard index equals ``intersection / union``; setting\n ``jaccard=True`` in MONAI's ``DiceLoss`` switches the denominator\n accordingly.\n \"\"\"\n return DiceLoss(\n sigmoid=sigmoid,\n softmax=softmax,\n jaccard=True,\n smooth_nr=smooth_nr,\n smooth_dr=smooth_dr,\n **kwargs,\n )\n\n\ndef tversky(\n alpha: float = 0.3,\n beta: float = 0.7,\n sigmoid: bool = False,\n softmax: bool = False,\n smooth_nr: float = 1e-5,\n smooth_dr: float = 1e-5,\n **kwargs,\n) -> TverskyLoss:\n \"\"\"Return a MONAI ``TverskyLoss`` instance.\n\n Parameters\n ----------\n alpha : float\n Weight of false positives.\n beta : float\n Weight of false negatives.\n \"\"\"\n return TverskyLoss(\n alpha=alpha,\n beta=beta,\n sigmoid=sigmoid,\n softmax=softmax,\n smooth_nr=smooth_nr,\n smooth_dr=smooth_dr,\n **kwargs,\n )\n\n\n# ---------------------------------------------------------------------------\n# Stubs — implemented in US2 (elbo) and US3 (wasserstein)\n# ---------------------------------------------------------------------------\n\n\ndef elbo(\n model: torch.nn.Module,\n kl_weight: float,\n reconstruction_loss: torch.Tensor,\n) -> torch.Tensor:\n \"\"\"Compute ELBO = reconstruction_loss + kl_weight * KL.\n\n The KL term is accumulated by Pyro sampling during the forward pass\n of Bayesian modules (:class:`~nobrainer.models.bayesian.layers.BayesianConv3d`\n and :class:`~nobrainer.models.bayesian.layers.BayesianLinear`).\n\n Parameters\n ----------\n model : nn.Module\n A model with one or more Bayesian layers whose ``.kl`` attributes\n have been populated by a recent forward pass.\n kl_weight : float\n Scalar multiplier for the KL divergence term (often ``1 / N_data``\n or ``1 / N_batches``).\n reconstruction_loss : torch.Tensor\n Scalar reconstruction loss (e.g., Dice or cross-entropy) already\n computed for the current batch.\n\n Returns\n -------\n torch.Tensor\n Scalar ELBO = reconstruction_loss + kl_weight * KL.\n \"\"\"\n from .models.bayesian.utils import accumulate_kl\n\n kl = accumulate_kl(model)\n return reconstruction_loss + kl_weight * kl\n\n\ndef wasserstein(y_true: torch.Tensor, y_pred: torch.Tensor) -> torch.Tensor:\n \"\"\"Wasserstein critic loss: ``E[D(fake)] - E[D(real)]``.\n\n Parameters\n ----------\n y_true : torch.Tensor\n Critic scores for real samples, shape ``(N,)`` or ``(N, 1)``.\n y_pred : torch.Tensor\n Critic scores for fake samples, shape ``(N,)`` or ``(N, 1)``.\n\n Returns\n -------\n torch.Tensor\n Scalar Wasserstein critic loss (minimised by the discriminator).\n \"\"\"\n return y_pred.mean() - y_true.mean()\n\n\ndef gradient_penalty(\n discriminator: torch.nn.Module,\n real: torch.Tensor,\n fake: torch.Tensor,\n lambda_gp: float = 10.0,\n) -> torch.Tensor:\n \"\"\"WGAN-GP gradient penalty.\n\n Interpolates between ``real`` and ``fake`` samples and penalises the\n discriminator gradient norm for deviating from 1.\n\n Parameters\n ----------\n discriminator : nn.Module\n The discriminator / critic network.\n real : torch.Tensor\n Real samples, shape ``(N, C, D, H, W)``.\n fake : torch.Tensor\n Generated samples, same shape as ``real``.\n lambda_gp : float\n Penalty weight (default 10, standard WGAN-GP value).\n\n Returns\n -------\n torch.Tensor\n Scalar gradient penalty term.\n \"\"\"\n b = real.size(0)\n eps = torch.rand(b, *([1] * (real.dim() - 1)), device=real.device)\n interp = (eps * real + (1.0 - eps) * fake.detach()).requires_grad_(True)\n d_interp = discriminator(interp)\n grads = torch.autograd.grad(\n outputs=d_interp,\n inputs=interp,\n grad_outputs=torch.ones_like(d_interp),\n create_graph=True,\n retain_graph=True,\n )[0]\n gp = ((grads.norm(2, dim=list(range(1, real.dim()))) - 1) ** 2).mean()\n return lambda_gp * gp\n\n\n# ---------------------------------------------------------------------------\n# Class weights and weighted losses\n# ---------------------------------------------------------------------------\n\n\ndef compute_class_weights(\n label_paths: list[str],\n n_classes: int,\n label_mapping: str | None = None,\n method: str = \"inverse_frequency\",\n max_samples: int | None = None,\n) -> torch.Tensor:\n \"\"\"Compute per-class weights from label volumes.\n\n Scans label files to count voxel frequencies per class, then converts\n to weights. Useful for imbalanced segmentation (e.g., 50-class brain\n parcellation where small structures are underrepresented).\n\n Parameters\n ----------\n label_paths : list of str\n Paths to label NIfTI/MGZ files.\n n_classes : int\n Number of target classes.\n label_mapping : str or None\n Label mapping name (e.g., ``\"50-class\"``) or CSV path.\n If None, labels are used as-is.\n method : str\n ``\"inverse_frequency\"`` (1/freq, normalized) or\n ``\"median_frequency\"`` (median_freq/freq, as in SegNet).\n max_samples : int or None\n Limit scanning to this many files (for speed).\n\n Returns\n -------\n torch.Tensor\n Shape ``(n_classes,)`` float tensor of weights.\n \"\"\"\n import nibabel as nib\n import numpy as np\n\n counts = np.zeros(n_classes, dtype=np.float64)\n paths = label_paths[:max_samples] if max_samples else label_paths\n\n remap_fn = None\n if label_mapping is not None:\n from nobrainer.processing.dataset import _load_label_mapping\n\n remap_fn = _load_label_mapping(label_mapping)\n\n for path in paths:\n arr = np.asarray(nib.load(path).dataobj, dtype=np.int32)\n if remap_fn is not None:\n arr = remap_fn(arr)\n for c in range(n_classes):\n counts[c] += (arr == c).sum()\n\n # Avoid division by zero\n counts = np.maximum(counts, 1.0)\n total = counts.sum()\n\n if method == \"median_frequency\":\n freqs = counts / total\n median_freq = np.median(freqs[freqs > 0])\n weights = median_freq / freqs\n else:\n # inverse_frequency: weight = total / (n_classes * count)\n weights = total / (n_classes * counts)\n\n # Normalize so mean weight = 1\n weights = weights / weights.mean()\n\n return torch.tensor(weights, dtype=torch.float32)\n\n\ndef weighted_cross_entropy(\n weight: torch.Tensor | None = None,\n label_smoothing: float = 0.0,\n) -> torch.nn.CrossEntropyLoss:\n \"\"\"Return a ``CrossEntropyLoss`` with optional per-class weights.\n\n Parameters\n ----------\n weight : torch.Tensor or None\n Per-class weights, shape ``(n_classes,)``.\n label_smoothing : float\n Label smoothing factor (default 0).\n \"\"\"\n return torch.nn.CrossEntropyLoss(\n weight=weight,\n label_smoothing=label_smoothing,\n )\n\n\nclass HammingLoss(torch.nn.Module):\n \"\"\"Hamming loss: fraction of misclassified voxels.\n\n A differentiable approximation of Hamming distance using soft\n predictions: ``g·(1-p) + (1-g)·p`` averaged over spatial dims.\n\n For use as a loss function with logits, set ``from_logits=True``\n to apply softmax first.\n\n Parameters\n ----------\n from_logits : bool\n Apply softmax to predictions (default True).\n \"\"\"\n\n def __init__(self, from_logits: bool = True) -> None:\n super().__init__()\n self.from_logits = from_logits\n\n def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:\n if self.from_logits:\n pred = torch.softmax(pred, dim=1)\n\n # One-hot encode target if needed\n if target.ndim == pred.ndim - 1:\n n_classes = pred.shape[1]\n target_oh = (\n torch.nn.functional.one_hot(target.long(), n_classes)\n .permute(0, 4, 1, 2, 3)\n .float()\n )\n else:\n target_oh = target.float()\n\n # Hamming: g*(1-p) + (1-g)*p = fraction of disagreement\n loss = target_oh * (1 - pred) + (1 - target_oh) * pred\n return loss.mean()\n\n\ndef hamming(from_logits: bool = True) -> HammingLoss:\n \"\"\"Return a :class:`HammingLoss` instance.\"\"\"\n return HammingLoss(from_logits=from_logits)\n\n\nclass DiceCELoss(torch.nn.Module):\n \"\"\"Combined Dice + weighted CrossEntropy loss.\n\n Commonly used for imbalanced segmentation tasks. The Dice component\n is inherently class-balanced; the CE component can use per-class\n weights.\n\n Parameters\n ----------\n weight : torch.Tensor or None\n Per-class weights for the CE term.\n dice_weight : float\n Relative weight of the Dice term (default 1.0).\n ce_weight : float\n Relative weight of the CE term (default 1.0).\n softmax : bool\n Apply softmax to predictions for the Dice term.\n label_smoothing : float\n Label smoothing for the CE term.\n \"\"\"\n\n def __init__(\n self,\n weight: torch.Tensor | None = None,\n dice_weight: float = 1.0,\n ce_weight: float = 1.0,\n softmax: bool = True,\n label_smoothing: float = 0.0,\n ) -> None:\n super().__init__()\n self.dice_loss = DiceLoss(softmax=softmax, to_onehot_y=True)\n self.ce_loss = torch.nn.CrossEntropyLoss(\n weight=weight, label_smoothing=label_smoothing\n )\n self.dice_weight = dice_weight\n self.ce_weight = ce_weight\n\n def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:\n # Dice expects target with channel dim\n if target.ndim == pred.ndim - 1:\n target_dice = target.unsqueeze(1)\n else:\n target_dice = target\n d = self.dice_loss(pred, target_dice)\n\n # CE expects target without channel dim, as long\n if target.ndim == pred.ndim:\n target_ce = target.squeeze(1)\n else:\n target_ce = target\n if target_ce.dtype != torch.long:\n target_ce = target_ce.long()\n ce = self.ce_loss(pred, target_ce)\n\n return self.dice_weight * d + self.ce_weight * ce\n\n\nclass FocalLoss(torch.nn.Module):\n \"\"\"Focal Loss for imbalanced multi-class segmentation.\n\n Down-weights well-classified examples and focuses on hard ones.\n ``FL(p) = -α · (1 - p)^γ · log(p)``\n\n Parameters\n ----------\n gamma : float\n Focusing parameter (default 2.0). Higher = more focus on hard examples.\n alpha : torch.Tensor or None\n Per-class weights. None = uniform.\n \"\"\"\n\n def __init__(\n self,\n gamma: float = 2.0,\n alpha: torch.Tensor | None = None,\n ) -> None:\n super().__init__()\n self.gamma = gamma\n if alpha is not None:\n self.register_buffer(\"alpha\", alpha)\n else:\n self.alpha = None\n\n def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:\n if target.ndim == pred.ndim - 1:\n pass # expected: target is (B, D, H, W), pred is (B, C, D, H, W)\n elif target.ndim == pred.ndim and target.shape[1] == 1:\n target = target.squeeze(1)\n target = target.long()\n\n ce = torch.nn.functional.cross_entropy(pred, target, reduction=\"none\")\n p = torch.exp(-ce) # probability of correct class\n focal_weight = (1 - p) ** self.gamma\n\n if self.alpha is not None:\n alpha_t = self.alpha[target]\n focal_weight = focal_weight * alpha_t\n\n return (focal_weight * ce).mean()\n\n\ndef focal(gamma: float = 2.0, alpha: torch.Tensor | None = None) -> FocalLoss:\n \"\"\"Return a :class:`FocalLoss` instance.\"\"\"\n return FocalLoss(gamma=gamma, alpha=alpha)\n\n\n# ---------------------------------------------------------------------------\n# Registry\n# ---------------------------------------------------------------------------\n\n_losses = {\n \"dice\": dice,\n \"generalized_dice\": generalized_dice,\n \"jaccard\": jaccard,\n \"tversky\": tversky,\n \"elbo\": elbo,\n \"wasserstein\": wasserstein,\n \"gradient_penalty\": gradient_penalty,\n \"hamming\": hamming,\n \"focal\": focal,\n \"weighted_cross_entropy\": weighted_cross_entropy,\n \"dice_ce\": DiceCELoss,\n}\n\n\ndef get(name: str):\n \"\"\"Return loss factory by name (case-insensitive).\"\"\"\n try:\n return _losses[name.lower()]\n except KeyError:\n avail = \", \".join(_losses)\n raise ValueError(f\"Unknown loss '{name}'. Available: {avail}\") from None\n" }, { "path": "nobrainer/metrics.py", "content": "\"\"\"Evaluation metrics for 3-D semantic segmentation (PyTorch / MONAI).\"\"\"\n\nfrom __future__ import annotations\n\nfrom monai.metrics import DiceMetric, HausdorffDistanceMetric, MeanIoU\nimport torch\n\n# ---------------------------------------------------------------------------\n# Factory functions returning configured MONAI metric objects\n# ---------------------------------------------------------------------------\n\n\ndef dice_metric(\n include_background: bool = True,\n reduction: str = \"mean\",\n **kwargs,\n) -> DiceMetric:\n \"\"\"Return a MONAI ``DiceMetric`` instance.\n\n Parameters\n ----------\n include_background : bool\n Include the background class in the Dice computation.\n reduction : str\n Reduction applied over the batch (``\"mean\"``, ``\"sum\"``, ``\"none\"``).\n \"\"\"\n return DiceMetric(\n include_background=include_background,\n reduction=reduction,\n **kwargs,\n )\n\n\ndef generalized_dice_metric(\n include_background: bool = True,\n reduction: str = \"mean\",\n **kwargs,\n) -> DiceMetric:\n \"\"\"Return a ``DiceMetric`` configured for multi-class (generalised) Dice.\n\n MONAI's ``DiceMetric`` computes per-class Dice and averages over\n classes, which is equivalent to Generalized Dice when class weights\n are uniform.\n \"\"\"\n return DiceMetric(\n include_background=include_background,\n reduction=reduction,\n **kwargs,\n )\n\n\ndef jaccard_metric(\n include_background: bool = True,\n reduction: str = \"mean\",\n **kwargs,\n) -> MeanIoU:\n \"\"\"Return a MONAI ``MeanIoU`` (Jaccard) metric instance.\"\"\"\n return MeanIoU(\n include_background=include_background,\n reduction=reduction,\n **kwargs,\n )\n\n\ndef tversky_metric(\n include_background: bool = True,\n reduction: str = \"mean\",\n **kwargs,\n) -> DiceMetric:\n \"\"\"Return a ``DiceMetric`` used as a Tversky metric proxy.\n\n Tversky with alpha=beta=0.5 equals Dice. For asymmetric Tversky,\n compute the Tversky index manually and wrap it in a custom metric.\n \"\"\"\n return DiceMetric(\n include_background=include_background,\n reduction=reduction,\n **kwargs,\n )\n\n\ndef hausdorff_metric(\n include_background: bool = False,\n distance_metric: str = \"euclidean\",\n percentile: float | None = 95.0,\n directed: bool = False,\n **kwargs,\n) -> HausdorffDistanceMetric:\n \"\"\"Return a MONAI ``HausdorffDistanceMetric`` instance.\n\n Parameters\n ----------\n include_background : bool\n Include background class in distance computation.\n distance_metric : str\n ``\"euclidean\"``, ``\"chessboard\"``, or ``\"taxicab\"``.\n percentile : float or None\n If set, computes the *n*-th percentile Hausdorff distance (e.g.\n 95 for HD95). ``None`` returns the maximum (HD100).\n directed : bool\n Compute directed (asymmetric) Hausdorff distance.\n \"\"\"\n return HausdorffDistanceMetric(\n include_background=include_background,\n distance_metric=distance_metric,\n percentile=percentile,\n directed=directed,\n **kwargs,\n )\n\n\ndef hamming_metric(reduction: str = \"mean\") -> \"HammingMetric\":\n \"\"\"Return a Hamming distance metric (fraction of misclassified voxels).\n\n Unlike MONAI metrics, this is a simple callable that takes\n ``(y_pred, y_true)`` integer label tensors and returns the mean\n fraction of disagreeing voxels.\n \"\"\"\n return HammingMetric(reduction=reduction)\n\n\nclass HammingMetric:\n \"\"\"Hamming distance metric: fraction of voxels where prediction != label.\"\"\"\n\n def __init__(self, reduction: str = \"mean\") -> None:\n self.reduction = reduction\n\n def __call__(\n self,\n y_pred: torch.Tensor,\n y_true: torch.Tensor,\n ) -> torch.Tensor:\n ne = (y_pred != y_true).float()\n # Average over spatial dims per sample\n spatial = list(range(1, ne.ndim))\n per_sample = ne.mean(dim=spatial)\n if self.reduction == \"mean\":\n return per_sample.mean()\n if self.reduction == \"sum\":\n return per_sample.sum()\n return per_sample # \"none\"\n\n\n# ---------------------------------------------------------------------------\n# Registry\n# ---------------------------------------------------------------------------\n\n_metrics = {\n \"dice\": dice_metric,\n \"generalized_dice\": generalized_dice_metric,\n \"jaccard\": jaccard_metric,\n \"tversky\": tversky_metric,\n \"hausdorff\": hausdorff_metric,\n \"hamming\": hamming_metric,\n}\n\n\ndef get(name: str):\n \"\"\"Return metric factory by name (case-insensitive).\"\"\"\n try:\n return _metrics[name.lower()]\n except KeyError:\n avail = \", \".join(_metrics)\n raise ValueError(f\"Unknown metric '{name}'. Available: {avail}\") from None\n" }, { "path": "nobrainer/models/__init__.py", "content": "\"\"\"Nobrainer model registry (PyTorch).\"\"\"\n\nfrom pprint import pprint\n\nfrom .autoencoder import autoencoder\nfrom .highresnet import highresnet\nfrom .meshnet import meshnet\nfrom .segformer3d import segformer3d\nfrom .segmentation import attention_unet, segresnet, swin_unetr, unet, unetr, vnet\nfrom .simsiam import simsiam\n\n__all__ = [\"get\", \"list_available_models\"]\n\n# Core models (always available)\n_models = {\n \"unet\": unet,\n \"vnet\": vnet,\n \"attention_unet\": attention_unet,\n \"unetr\": unetr,\n \"meshnet\": meshnet,\n \"highresnet\": highresnet,\n \"autoencoder\": autoencoder,\n \"simsiam\": simsiam,\n \"swin_unetr\": swin_unetr,\n \"segresnet\": segresnet,\n \"segformer3d\": segformer3d,\n}\n\n# Optional: Bayesian models (require pyro-ppl)\ntry:\n from .bayesian import bayesian_meshnet, bayesian_vnet\n\n _models[\"bayesian_vnet\"] = bayesian_vnet\n _models[\"bayesian_meshnet\"] = bayesian_meshnet\nexcept ImportError:\n pass\n\n# KWYK MeshNet (VWN-based, no Pyro dependency)\nfrom .bayesian.kwyk_meshnet import kwyk_meshnet # noqa: E402\n\n_models[\"kwyk_meshnet\"] = kwyk_meshnet\n\n# Optional: Generative models (require pytorch-lightning)\ntry:\n from .generative import dcgan, progressivegan\n\n _models[\"progressivegan\"] = progressivegan\n _models[\"dcgan\"] = dcgan\nexcept ImportError:\n pass\n\n\ndef get(name: str):\n \"\"\"Return factory callable for a model by name (case-insensitive).\n\n Parameters\n ----------\n name : str\n Model name.\n\n Returns\n -------\n Callable that constructs a ``torch.nn.Module``.\n \"\"\"\n if not isinstance(name, str):\n raise ValueError(\"Model name must be a string.\")\n key = name.lower()\n if key in _models:\n return _models[key]\n # Check if it's an optional model that wasn't loaded\n optional = {\n \"bayesian_vnet\": \"pyro-ppl\",\n \"bayesian_meshnet\": \"pyro-ppl\",\n \"progressivegan\": \"pytorch-lightning\",\n \"dcgan\": \"pytorch-lightning\",\n }\n if key in optional:\n raise ImportError(\n f\"Model '{name}' requires '{optional[key]}'. \"\n f\"Install with: uv pip install {optional[key]}\"\n )\n avail = \", \".join(_models)\n raise ValueError(f\"Unknown model '{name}'. Available: {avail}.\")\n\n\ndef available_models() -> list[str]:\n return list(_models)\n\n\ndef list_available_models() -> None:\n pprint(available_models())\n" }, { "path": "nobrainer/models/_constants.py", "content": "\"\"\"Shared constants for nobrainer models.\"\"\"\n\nfrom __future__ import annotations\n\n# Dilation schedules indexed by receptive field size.\n# Used by MeshNet, BayesianMeshNet, and KWYKMeshNet.\nDILATION_SCHEDULES: dict[int, list[int]] = {\n 37: [1, 1, 1, 2, 4, 8, 1],\n 67: [1, 1, 2, 4, 8, 16, 1],\n 129: [1, 2, 4, 8, 16, 32, 1],\n}\n" }, { "path": "nobrainer/models/_utils.py", "content": "\"\"\"Shared utilities for nobrainer models and training.\"\"\"\n\nfrom __future__ import annotations\n\nfrom pathlib import Path\n\nimport nibabel as nib\nimport numpy as np\nimport torch\n\n\ndef unpack_batch(\n batch: dict | list | tuple,\n device: torch.device,\n) -> tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Extract image and label tensors from a batch, move to device.\n\n Handles both dict-style (MONAI) and tuple-style (TensorDataset) batches.\n Squeezes label channel dim and casts to long for CrossEntropyLoss.\n\n Parameters\n ----------\n batch : dict, list, or tuple\n A batch from a DataLoader.\n device : torch.device\n Target device.\n\n Returns\n -------\n images : torch.Tensor\n Shape ``(B, C, D, H, W)`` on *device*.\n labels : torch.Tensor\n Shape ``(B, D, H, W)`` long dtype on *device*.\n \"\"\"\n if isinstance(batch, dict):\n images = batch[\"image\"].to(device)\n labels = batch[\"label\"].to(device)\n elif isinstance(batch, (list, tuple)):\n images = batch[0].to(device)\n labels = batch[1].to(device)\n else:\n raise TypeError(f\"Unsupported batch type: {type(batch)}\")\n\n # Squeeze channel dim from labels if present\n if labels.ndim == images.ndim and labels.shape[1] == 1:\n labels = labels.squeeze(1)\n # Cast float labels to long for CrossEntropyLoss\n if labels.dtype in (torch.float32, torch.float64):\n labels = labels.long()\n\n return images, labels\n\n\ndef load_input(\n inputs: str | Path | np.ndarray | nib.Nifti1Image,\n) -> tuple[np.ndarray, np.ndarray | None]:\n \"\"\"Load a 3D volume from various input types.\n\n Parameters\n ----------\n inputs : str, Path, ndarray, or Nifti1Image\n Input volume.\n\n Returns\n -------\n arr : np.ndarray\n 3D array, shape ``(D, H, W)``.\n affine : np.ndarray or None\n 4x4 affine matrix (None if input is raw array).\n \"\"\"\n if isinstance(inputs, (str, Path)):\n img = nib.load(inputs)\n return np.asarray(img.dataobj, dtype=np.float32), img.affine\n elif isinstance(inputs, nib.Nifti1Image):\n return np.asarray(inputs.dataobj, dtype=np.float32), inputs.affine\n elif isinstance(inputs, np.ndarray):\n return inputs.astype(np.float32), None\n else:\n raise TypeError(f\"Unsupported input type: {type(inputs)}\")\n\n\ndef model_supports_mc(model: torch.nn.Module) -> bool:\n \"\"\"Check if a model supports the ``mc`` keyword argument in forward().\n\n Returns True if the model has a ``supports_mc`` class attribute set\n to True, or if its forward method accepts an ``mc`` parameter.\n \"\"\"\n if getattr(model, \"supports_mc\", False):\n return True\n # Check the forward signature\n import inspect\n\n sig = inspect.signature(model.forward)\n return \"mc\" in sig.parameters\n" }, { "path": "nobrainer/models/autoencoder.py", "content": "\"\"\"Symmetric 3-D autoencoder (PyTorch).\n\nEncodes a 3-D volume into a flat latent vector and reconstructs it via\ntransposed convolutions.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport math\n\nimport torch\nimport torch.nn as nn\n\n\nclass Autoencoder(nn.Module):\n \"\"\"Symmetric 3-D convolutional autoencoder.\n\n Dynamically builds encoder depth from the spatial size of the input.\n\n Parameters\n ----------\n input_shape : tuple of int\n Volume shape ``(D, H, W)`` (spatial dims only).\n in_channels : int\n Number of input channels (1 for single-modality MRI).\n encoding_dim : int\n Size of the flat latent code.\n n_base_filters : int\n Base filter count; doubled each encoder level.\n batchnorm : bool\n Whether to apply Batch Normalisation in conv blocks.\n \"\"\"\n\n def __init__(\n self,\n input_shape: tuple[int, int, int] = (64, 64, 64),\n in_channels: int = 1,\n encoding_dim: int = 512,\n n_base_filters: int = 16,\n batchnorm: bool = True,\n ) -> None:\n super().__init__()\n D = input_shape[0]\n n_levels = int(math.log2(D))\n\n # Build encoder\n enc_layers: list[nn.Module] = []\n ch_in = in_channels\n self._enc_channels: list[int] = []\n for i in range(n_levels):\n ch_out = min(n_base_filters * (2**i), encoding_dim)\n self._enc_channels.append(ch_out)\n block: list[nn.Module] = [\n nn.Conv3d(\n ch_in,\n ch_out,\n kernel_size=4,\n stride=2,\n padding=1,\n bias=not batchnorm,\n ),\n ]\n if batchnorm:\n block.append(nn.BatchNorm3d(ch_out))\n block.append(nn.ReLU(inplace=True))\n enc_layers.extend(block)\n ch_in = ch_out\n\n self.encoder_conv = nn.Sequential(*enc_layers)\n self.encoder_fc = nn.Linear(ch_in, encoding_dim)\n\n # Build decoder (mirror of encoder)\n dec_ch = list(reversed(self._enc_channels))\n self.decoder_fc = nn.Linear(encoding_dim, dec_ch[0])\n\n dec_layers: list[nn.Module] = []\n all_out = dec_ch[1:] + [in_channels]\n for i, ch_out in enumerate(all_out):\n ch_in_d = dec_ch[i]\n is_last = i == len(all_out) - 1\n act_d: nn.Module = (\n nn.Sigmoid() if is_last else nn.LeakyReLU(0.2, inplace=True)\n )\n use_bn = batchnorm and not is_last\n block_d: list[nn.Module] = [\n nn.ConvTranspose3d(\n ch_in_d, ch_out, kernel_size=4, stride=2, padding=1, bias=not use_bn\n ),\n ]\n if use_bn:\n block_d.append(nn.BatchNorm3d(ch_out))\n block_d.append(act_d)\n dec_layers.extend(block_d)\n\n self.decoder_conv = nn.Sequential(*dec_layers)\n\n def encode(self, x: torch.Tensor) -> torch.Tensor:\n h = self.encoder_conv(x) # (N, C, 1, 1, 1)\n return self.encoder_fc(h.flatten(1))\n\n def decode(self, z: torch.Tensor) -> torch.Tensor:\n h = self.decoder_fc(z).view(z.size(0), -1, 1, 1, 1)\n return self.decoder_conv(h)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return self.decode(self.encode(x))\n\n\ndef autoencoder(\n input_shape: tuple[int, int, int] = (64, 64, 64),\n in_channels: int = 1,\n encoding_dim: int = 512,\n n_base_filters: int = 16,\n batchnorm: bool = True,\n **kwargs,\n) -> Autoencoder:\n \"\"\"Factory function for :class:`Autoencoder`.\"\"\"\n return Autoencoder(\n input_shape=input_shape,\n in_channels=in_channels,\n encoding_dim=encoding_dim,\n n_base_filters=n_base_filters,\n batchnorm=batchnorm,\n )\n\n\n__all__ = [\"Autoencoder\", \"autoencoder\"]\n" }, { "path": "nobrainer/models/bayesian/__init__.py", "content": "\"\"\"Bayesian model sub-package.\n\nTwo flavours of Bayesian convolution are provided:\n\n* **Bayes-by-backprop** (``BayesianConv3d``, ``BayesianMeshNet``) — Pyro-based,\n weight uncertainty via learned mu/sigma, supports standard_normal/laplace priors.\n* **Variational Weight Normalization** (``VWNConv3d``, ``KWYKMeshNet``) — matches\n the original kwyk architecture (McClure et al. 2019) with weight normalization,\n local reparameterization, and Bernoulli or Concrete dropout.\n\"\"\"\n\nfrom .bayesian_meshnet import BayesianMeshNet, bayesian_meshnet\nfrom .bayesian_vnet import BayesianVNet, bayesian_vnet\nfrom .kwyk_meshnet import KWYKMeshNet, kwyk_meshnet\nfrom .layers import BayesianConv3d, BayesianLinear\nfrom .utils import accumulate_kl\nfrom .vwn_layers import ConcreteDropout3d, FFGConv3d, VWNConv3d\n\n__all__ = [\n \"BayesianConv3d\",\n \"BayesianLinear\",\n \"BayesianMeshNet\",\n \"BayesianVNet\",\n \"ConcreteDropout3d\",\n \"FFGConv3d\",\n \"KWYKMeshNet\",\n \"VWNConv3d\",\n \"accumulate_kl\",\n \"bayesian_meshnet\",\n \"bayesian_vnet\",\n \"kwyk_meshnet\",\n]\n" }, { "path": "nobrainer/models/bayesian/bayesian_meshnet.py", "content": "\"\"\"Bayesian MeshNet: dilated-convolution segmentation with weight uncertainty.\n\nReplaces every ``nn.Conv3d`` in the 7-layer dilated architecture with\n:class:`~nobrainer.models.bayesian.layers.BayesianConv3d`.\n\nReference\n---------\nFedorov A. et al., \"End-to-end learning of brain tissue segmentation\nfrom imperfect labeling\", IJCNN 2017. arXiv:1612.00940.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom pyro.nn import PyroModule\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom nobrainer.models._constants import ( # noqa: E501\n DILATION_SCHEDULES as _DILATION_SCHEDULES,\n)\n\nfrom .layers import BayesianConv3d\n\n\nclass _BayesConvBNActDrop(PyroModule):\n \"\"\"Single dilated Bayesian conv layer with BN + ELU/ReLU + spatial dropout.\"\"\"\n\n def __init__(\n self,\n in_ch: int,\n out_ch: int,\n dilation: int,\n activation: str,\n dropout_rate: float,\n prior_type: str,\n **sas_kwargs,\n ) -> None:\n super().__init__()\n padding = dilation # same-size output for 3×3×3 kernel\n self.conv = BayesianConv3d(\n in_ch,\n out_ch,\n kernel_size=3,\n padding=padding,\n dilation=dilation,\n bias=False,\n prior_type=prior_type,\n **sas_kwargs,\n )\n self.bn = nn.BatchNorm3d(out_ch)\n self.act_fn = {\"relu\": F.relu, \"elu\": F.elu}[activation.lower()]\n self.dropout = nn.Dropout3d(p=dropout_rate)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return self.dropout(self.act_fn(self.bn(self.conv(x))))\n\n\nclass BayesianMeshNet(PyroModule):\n \"\"\"3-D MeshNet with Bayesian convolutional layers.\n\n Identical dilated-convolution schedule as :class:`~nobrainer.models.meshnet.MeshNet`\n but all ``nn.Conv3d`` layers are replaced with :class:`BayesianConv3d`.\n\n Parameters\n ----------\n n_classes : int\n Number of output segmentation classes.\n in_channels : int\n Number of input image channels.\n filters : int\n Feature-map count in all hidden layers.\n receptive_field : int\n One of ``37``, ``67``, ``129`` — selects the dilation schedule.\n activation : str\n ``\"relu\"`` or ``\"elu\"``.\n dropout_rate : float\n Spatial dropout probability (0 = disabled).\n prior_type : str\n ``\"standard_normal\"``, ``\"laplace\"``, or ``\"spike_and_slab\"``.\n kl_weight : float\n Scalar applied to the summed KL when computing the ELBO.\n Stored as an attribute; not used internally during forward.\n spike_sigma : float\n Spike component σ for spike-and-slab prior (default 0.001).\n slab_sigma : float\n Slab component σ for spike-and-slab prior (default 1.0).\n prior_pi : float\n Prior probability of the spike component (default 0.5).\n \"\"\"\n\n def __init__(\n self,\n n_classes: int = 1,\n in_channels: int = 1,\n filters: int = 71,\n receptive_field: int = 67,\n activation: str = \"relu\",\n dropout_rate: float = 0.25,\n prior_type: str = \"standard_normal\",\n kl_weight: float = 1.0,\n spike_sigma: float = 0.001,\n slab_sigma: float = 1.0,\n prior_pi: float = 0.5,\n ) -> None:\n super().__init__()\n if receptive_field not in _DILATION_SCHEDULES:\n raise ValueError(\n f\"receptive_field must be one of {list(_DILATION_SCHEDULES)}, \"\n f\"got {receptive_field}\"\n )\n self.kl_weight = kl_weight\n self.prior_type = prior_type\n dilations = _DILATION_SCHEDULES[receptive_field]\n self._n_layers = len(dilations)\n\n # Extra kwargs for spike-and-slab layers\n sas_kwargs = {}\n if prior_type == \"spike_and_slab\":\n sas_kwargs = {\n \"spike_sigma\": spike_sigma,\n \"slab_sigma\": slab_sigma,\n \"prior_pi\": prior_pi,\n }\n\n # Register each Bayesian layer as a named attribute so Pyro assigns\n # unique sample site names (nn.ModuleList does not propagate names).\n for i, dil in enumerate(dilations):\n in_ch = in_channels if i == 0 else filters\n layer = _BayesConvBNActDrop(\n in_ch,\n filters,\n dil,\n activation,\n dropout_rate,\n prior_type,\n **sas_kwargs,\n )\n setattr(self, f\"layer_{i}\", layer)\n\n # Final 1×1×1 classifier — deterministic\n self.classifier = nn.Conv3d(filters, n_classes, kernel_size=1)\n\n supports_mc = True\n\n def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:\n h = x\n for i in range(self._n_layers):\n h = getattr(self, f\"layer_{i}\")(h)\n return self.classifier(h)\n\n\ndef bayesian_meshnet(\n n_classes: int = 1,\n in_channels: int = 1,\n filters: int = 71,\n receptive_field: int = 67,\n activation: str = \"relu\",\n dropout_rate: float = 0.25,\n prior_type: str = \"standard_normal\",\n kl_weight: float = 1.0,\n spike_sigma: float = 0.001,\n slab_sigma: float = 1.0,\n prior_pi: float = 0.5,\n **kwargs,\n) -> BayesianMeshNet:\n \"\"\"Factory function for :class:`BayesianMeshNet`.\"\"\"\n return BayesianMeshNet(\n n_classes=n_classes,\n in_channels=in_channels,\n filters=filters,\n receptive_field=receptive_field,\n activation=activation,\n dropout_rate=dropout_rate,\n prior_type=prior_type,\n kl_weight=kl_weight,\n spike_sigma=spike_sigma,\n slab_sigma=slab_sigma,\n prior_pi=prior_pi,\n )\n\n\n__all__ = [\"BayesianMeshNet\", \"bayesian_meshnet\"]\n" }, { "path": "nobrainer/models/bayesian/bayesian_vnet.py", "content": "\"\"\"Bayesian V-Net: encoder-decoder segmentation with weight uncertainty.\n\nReplaces the standard ``nn.Conv3d`` convolutions with\n:class:`~nobrainer.models.bayesian.layers.BayesianConv3d` (mean-field\nvariational inference via Pyro), preserving the residual encoder-decoder\narchitecture of V-Net.\n\nReference\n---------\nMilletari F. et al., \"V-Net: Fully Convolutional Neural Networks for\nVolumetric Medical Image Segmentation\", 3DV 2016. arXiv:1606.04797.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom pyro.nn import PyroModule\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom .layers import BayesianConv3d\n\n\nclass _BayesResBlock(PyroModule):\n \"\"\"Two-layer residual block with BayesianConv3d and a skip connection.\"\"\"\n\n def __init__(self, channels: int, prior_type: str = \"standard_normal\") -> None:\n super().__init__()\n self.conv1 = BayesianConv3d(\n channels, channels, kernel_size=3, padding=1, prior_type=prior_type\n )\n self.conv2 = BayesianConv3d(\n channels, channels, kernel_size=3, padding=1, prior_type=prior_type\n )\n self.bn1 = nn.BatchNorm3d(channels)\n self.bn2 = nn.BatchNorm3d(channels)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n h = F.elu(self.bn1(self.conv1(x)))\n h = self.bn2(self.conv2(h))\n return F.elu(h + x)\n\n\nclass _EncoderBlock(PyroModule):\n \"\"\"One encoder level: project channels → residual block → max-pool.\"\"\"\n\n def __init__(\n self,\n in_ch: int,\n out_ch: int,\n prior_type: str = \"standard_normal\",\n ) -> None:\n super().__init__()\n self.proj = BayesianConv3d(in_ch, out_ch, kernel_size=1, prior_type=prior_type)\n self.bn_proj = nn.BatchNorm3d(out_ch)\n self.res = _BayesResBlock(out_ch, prior_type=prior_type)\n self.pool = nn.MaxPool3d(kernel_size=2)\n\n def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:\n h = F.elu(self.bn_proj(self.proj(x)))\n h = self.res(h)\n return self.pool(h), h # (down-sampled, skip)\n\n\nclass _DecoderBlock(PyroModule):\n \"\"\"One decoder level: up-sample → concat skip → project → residual block.\"\"\"\n\n def __init__(\n self,\n in_ch: int,\n skip_ch: int,\n out_ch: int,\n prior_type: str = \"standard_normal\",\n ) -> None:\n super().__init__()\n self.upsample = nn.ConvTranspose3d(in_ch, out_ch, kernel_size=2, stride=2)\n self.proj = BayesianConv3d(\n out_ch + skip_ch, out_ch, kernel_size=1, prior_type=prior_type\n )\n self.bn_proj = nn.BatchNorm3d(out_ch)\n self.res = _BayesResBlock(out_ch, prior_type=prior_type)\n\n def forward(self, x: torch.Tensor, skip: torch.Tensor) -> torch.Tensor:\n h = self.upsample(x)\n if h.shape != skip.shape:\n h = F.interpolate(\n h, size=skip.shape[2:], mode=\"trilinear\", align_corners=False\n )\n h = torch.cat([h, skip], dim=1)\n h = F.elu(self.bn_proj(self.proj(h)))\n return self.res(h)\n\n\nclass BayesianVNet(PyroModule):\n \"\"\"3-D V-Net with Bayesian convolutional layers.\n\n All ``nn.Conv3d`` layers in the encoder and decoder are replaced with\n :class:`BayesianConv3d`. Upsampling transposed convolutions remain\n deterministic.\n\n Parameters\n ----------\n n_classes : int\n Number of output segmentation classes.\n in_channels : int\n Number of input image channels.\n base_filters : int\n Feature-map count at the first encoder level (doubles each level).\n levels : int\n Number of encoder/decoder levels (default 4).\n prior_type : str\n ``\"standard_normal\"`` or ``\"laplace\"`` — forwarded to Bayesian layers.\n kl_weight : float\n Scalar applied to the summed KL divergence when computing the ELBO.\n Stored as an attribute; not used internally during forward.\n \"\"\"\n\n def __init__(\n self,\n n_classes: int = 1,\n in_channels: int = 1,\n base_filters: int = 16,\n levels: int = 4,\n prior_type: str = \"standard_normal\",\n kl_weight: float = 1.0,\n ) -> None:\n super().__init__()\n self.kl_weight = kl_weight\n self._levels = levels\n\n ch = [base_filters * (2**i) for i in range(levels)]\n\n # Input projection\n self.input_proj = BayesianConv3d(\n in_channels, ch[0], kernel_size=3, padding=1, prior_type=prior_type\n )\n self.input_bn = nn.BatchNorm3d(ch[0])\n\n # Encoder — registered as individually named attributes so Pyro can\n # assign unique site names (nn.ModuleList does not propagate names).\n # encoder_i: ch[i] → ch[i+1]; skip tensor has ch[i+1] channels.\n for i in range(levels - 1):\n enc = _EncoderBlock(ch[i], ch[i + 1], prior_type)\n setattr(self, f\"encoder_{i}\", enc)\n\n # Bottom residual block (no pooling)\n self.bottom_res = _BayesResBlock(ch[-1], prior_type=prior_type)\n\n # Decoder — decoder_i processes the stage closest to the bottom first.\n # decoder_i: in_ch = ch[L-1-i], skip_ch = ch[L-1-i], out_ch = ch[L-2-i]\n # (upsampled out_ch channels are cat'd with skip_ch to give\n # out_ch + skip_ch channels before the projection layer)\n L = levels\n for i in range(L - 1):\n in_ch = ch[L - 1 - i]\n skip_ch = ch[L - 1 - i] # skip from encoder_{L-2-i} has ch[L-1-i] chans\n out_ch = ch[L - 2 - i]\n dec = _DecoderBlock(in_ch, skip_ch, out_ch, prior_type)\n setattr(self, f\"decoder_{i}\", dec)\n\n # Final 1×1×1 classifier — deterministic\n self.classifier = nn.Conv3d(ch[0], n_classes, kernel_size=1)\n\n supports_mc = True\n\n def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:\n h = F.elu(self.input_bn(self.input_proj(x)))\n\n skips: list[torch.Tensor] = []\n for i in range(self._levels - 1):\n enc = getattr(self, f\"encoder_{i}\")\n h, skip = enc(h)\n skips.append(skip)\n\n h = self.bottom_res(h)\n\n for i in range(self._levels - 1):\n dec = getattr(self, f\"decoder_{i}\")\n skip = skips[self._levels - 2 - i]\n h = dec(h, skip)\n\n return self.classifier(h)\n\n\ndef bayesian_vnet(\n n_classes: int = 1,\n in_channels: int = 1,\n base_filters: int = 16,\n levels: int = 4,\n prior_type: str = \"standard_normal\",\n kl_weight: float = 1.0,\n **kwargs,\n) -> BayesianVNet:\n \"\"\"Factory function for :class:`BayesianVNet`.\"\"\"\n return BayesianVNet(\n n_classes=n_classes,\n in_channels=in_channels,\n base_filters=base_filters,\n levels=levels,\n prior_type=prior_type,\n kl_weight=kl_weight,\n )\n\n\n__all__ = [\"BayesianVNet\", \"bayesian_vnet\"]\n" }, { "path": "nobrainer/models/bayesian/kwyk_meshnet.py", "content": "\"\"\"KWYK MeshNet variants — matching McClure et al. (2019) architecture.\n\nAll three kwyk models use Fully Factorized Gaussian (FFG) convolutions\nwith learned per-weight μ and σ, and the local reparameterization trick\n(Kingma et al. 2015). They differ in the dropout layer:\n\n* **bwn** / **bwn_multi**: FFG conv + Bernoulli dropout\n (``bwn`` disables dropout at inference; ``bwn_multi`` keeps it on)\n* **bvwn_multi_prior**: FFG conv + Concrete dropout (learned per-filter rate)\n This is the \"spike-and-slab dropout\" (SSD) model from the paper.\n\nReference\n---------\nMcClure P. et al., \"Knowing What You Know in Brain Segmentation Using\nBayesian Deep Neural Networks\", Front. Neuroinform. 2019.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom nobrainer.models._constants import ( # noqa: E501\n DILATION_SCHEDULES as _DILATION_SCHEDULES,\n)\n\nfrom .vwn_layers import ConcreteDropout3d, FFGConv3d\n\n\nclass _VWNLayerBernoulli(nn.Module):\n \"\"\"VWN conv + ReLU + Bernoulli dropout (bwn / bwn_multi).\"\"\"\n\n def __init__(\n self,\n in_ch: int,\n out_ch: int,\n dilation: int,\n dropout_rate: float,\n sigma_init: float,\n ) -> None:\n super().__init__()\n self.conv = FFGConv3d(\n in_ch,\n out_ch,\n kernel_size=3,\n padding=dilation,\n dilation=dilation,\n bias=False,\n sigma_init=sigma_init,\n )\n self.dropout = nn.Dropout3d(p=dropout_rate)\n\n def forward(\n self,\n x: torch.Tensor,\n mc_vwn: bool = True,\n mc_dropout: bool = True,\n ) -> torch.Tensor:\n # Original TF order: conv -> dropout -> relu (meshnetbwn.py:59-61)\n h = self.conv(x, mc=mc_vwn)\n if mc_dropout:\n h = self.dropout(h)\n return F.relu(h)\n\n\nclass _VWNLayerConcrete(nn.Module):\n \"\"\"VWN conv + ReLU + Concrete dropout (bvwn_multi_prior).\"\"\"\n\n def __init__(\n self,\n in_ch: int,\n out_ch: int,\n dilation: int,\n sigma_init: float,\n concrete_temperature: float = 0.02,\n concrete_init_p: float = 0.9,\n ) -> None:\n super().__init__()\n self.conv = FFGConv3d(\n in_ch,\n out_ch,\n kernel_size=3,\n padding=dilation,\n dilation=dilation,\n bias=False,\n sigma_init=sigma_init,\n )\n self.dropout = ConcreteDropout3d(\n out_ch,\n temperature=concrete_temperature,\n init_p=concrete_init_p,\n )\n\n def forward(\n self,\n x: torch.Tensor,\n mc_vwn: bool = True,\n mc_dropout: bool = True,\n ) -> torch.Tensor:\n # Original TF order: conv -> dropout -> relu (meshnetbvwn.py:54-55)\n h = self.conv(x, mc=mc_vwn)\n h = self.dropout(h, mc=mc_dropout)\n return F.relu(h)\n\n\nclass KWYKMeshNet(nn.Module):\n \"\"\"KWYK MeshNet with variational weight normalization.\n\n This is the architecture used in McClure et al. (2019). All layers\n use VWN convolutions; the ``dropout_type`` parameter selects between\n Bernoulli (``\"bernoulli\"``) and Concrete (``\"concrete\"``) dropout.\n\n Parameters\n ----------\n n_classes : int\n Number of output segmentation classes.\n in_channels : int\n Number of input image channels.\n filters : int\n Feature-map count in all hidden layers.\n receptive_field : int\n One of ``37``, ``67``, ``129`` — selects the dilation schedule.\n dropout_type : str\n ``\"bernoulli\"`` for bwn/bwn_multi, ``\"concrete\"`` for bvwn_multi_prior.\n dropout_rate : float\n For Bernoulli dropout (ignored for concrete).\n sigma_init : float\n Initial value for weight sigma (default 1e-4, matching kwyk).\n concrete_temperature : float\n Temperature for concrete dropout (default 0.02).\n concrete_init_p : float\n Initial dropout probability for concrete dropout (default 0.9).\n \"\"\"\n\n def __init__(\n self,\n n_classes: int = 1,\n in_channels: int = 1,\n filters: int = 71,\n receptive_field: int = 67,\n dropout_type: str = \"bernoulli\",\n dropout_rate: float = 0.25,\n sigma_init: float = 1e-4,\n concrete_temperature: float = 0.02,\n concrete_init_p: float = 0.9,\n ) -> None:\n super().__init__()\n if receptive_field not in _DILATION_SCHEDULES:\n raise ValueError(\n f\"receptive_field must be one of {list(_DILATION_SCHEDULES)}, \"\n f\"got {receptive_field}\"\n )\n self.dropout_type = dropout_type\n dilations = _DILATION_SCHEDULES[receptive_field]\n self._n_layers = len(dilations)\n\n for i, dil in enumerate(dilations):\n in_ch = in_channels if i == 0 else filters\n if dropout_type == \"concrete\":\n layer = _VWNLayerConcrete(\n in_ch,\n filters,\n dil,\n sigma_init,\n concrete_temperature,\n concrete_init_p,\n )\n else:\n layer = _VWNLayerBernoulli(\n in_ch,\n filters,\n dil,\n dropout_rate,\n sigma_init,\n )\n setattr(self, f\"layer_{i}\", layer)\n\n self.classifier = nn.Conv3d(filters, n_classes, kernel_size=1)\n\n def forward(\n self,\n x: torch.Tensor,\n mc: bool | None = None,\n mc_vwn: bool = True,\n mc_dropout: bool = True,\n ) -> torch.Tensor:\n \"\"\"Forward pass.\n\n Parameters\n ----------\n x : Tensor\n Input ``(B, 1, D, H, W)``.\n mc : bool or None\n Legacy convenience flag. If provided, sets both ``mc_vwn``\n and ``mc_dropout`` to the same value (backward compat).\n mc_vwn : bool\n If True, use stochastic VWN reparameterization.\n If False, use deterministic mean weights only.\n mc_dropout : bool\n If True, apply stochastic dropout.\n If False, skip dropout (Bernoulli) or use expectation (Concrete).\n\n Note\n ----\n The original TF bwn model trains with ``mc_vwn=False, mc_dropout=True``\n (deterministic weights + stochastic dropout).\n \"\"\"\n if mc is not None:\n mc_vwn = mc\n mc_dropout = mc\n\n h = x\n for i in range(self._n_layers):\n h = getattr(self, f\"layer_{i}\")(h, mc_vwn=mc_vwn, mc_dropout=mc_dropout)\n return self.classifier(h)\n\n def kl_divergence(self) -> torch.Tensor:\n \"\"\"Sum KL divergence from all VWN conv layers.\"\"\"\n kl = torch.tensor(0.0, device=next(self.parameters()).device)\n for m in self.modules():\n if isinstance(m, FFGConv3d):\n kl = kl + m.kl\n return kl\n\n def concrete_regularization(self) -> torch.Tensor:\n \"\"\"Sum concrete dropout regularization (0 for bernoulli models).\"\"\"\n reg = torch.tensor(0.0, device=next(self.parameters()).device)\n for m in self.modules():\n if isinstance(m, ConcreteDropout3d):\n reg = reg + m.regularization()\n return reg\n\n\ndef kwyk_meshnet(\n n_classes: int = 1,\n in_channels: int = 1,\n filters: int = 71,\n receptive_field: int = 67,\n dropout_type: str = \"bernoulli\",\n dropout_rate: float = 0.25,\n sigma_init: float = 1e-4,\n concrete_temperature: float = 0.02,\n concrete_init_p: float = 0.9,\n **kwargs,\n) -> KWYKMeshNet:\n \"\"\"Factory function for :class:`KWYKMeshNet`.\"\"\"\n return KWYKMeshNet(\n n_classes=n_classes,\n in_channels=in_channels,\n filters=filters,\n receptive_field=receptive_field,\n dropout_type=dropout_type,\n dropout_rate=dropout_rate,\n sigma_init=sigma_init,\n concrete_temperature=concrete_temperature,\n concrete_init_p=concrete_init_p,\n )\n\n\n__all__ = [\"KWYKMeshNet\", \"kwyk_meshnet\"]\n" }, { "path": "nobrainer/models/bayesian/layers.py", "content": "\"\"\"Bayesian convolutional and linear layers as Pyro modules.\n\nBoth ``BayesianConv3d`` and ``BayesianLinear`` implement weight\nuncertainty by maintaining learnable ``weight_mu`` and ``weight_sigma``\nparameters. During each stochastic forward pass they sample a weight\nmatrix from ``Normal(weight_mu, softplus(weight_sigma))`` and accumulate\nthe KL divergence against the prior into ``self.kl``.\n\nThree prior types are supported (matching the kwyk study variants):\n\n* ``\"standard_normal\"`` — N(0, 1) prior, standard Bayes-by-backprop.\n* ``\"laplace\"`` — tight Normal N(0, 0.1) approximation of a Laplace prior.\n* ``\"spike_and_slab\"`` — mixture prior ``π·N(0, σ₁) + (1-π)·N(0, σ₂)``\n where σ₁ (spike) is small and σ₂ (slab) is large. Each weight also\n learns a log-odds ``z_logit`` controlling how much mass is on the spike\n vs slab, implementing variational spike-and-slab dropout (SSD) as in\n McClure et al. (2019).\n\"\"\"\n\nfrom __future__ import annotations\n\nimport math\n\nimport pyro\nimport pyro.distributions as dist\nfrom pyro.nn import PyroModule, PyroParam\nimport torch\nfrom torch.distributions import constraints\nimport torch.nn.functional as F\n\n# ---------------------------------------------------------------------------\n# KL helpers\n# ---------------------------------------------------------------------------\n\n\ndef _kl_normal_normal(\n mu: torch.Tensor,\n sigma: torch.Tensor,\n prior_mu: float,\n prior_sigma: float,\n) -> torch.Tensor:\n \"\"\"Analytic KL(N(mu, sigma) || N(prior_mu, prior_sigma)).\"\"\"\n return (\n torch.log(prior_sigma / (sigma + 1e-8))\n + (sigma**2 + (mu - prior_mu) ** 2) / (2 * prior_sigma**2)\n - 0.5\n ).sum()\n\n\ndef _kl_spike_and_slab(\n mu: torch.Tensor,\n sigma: torch.Tensor,\n z_logit: torch.Tensor,\n spike_sigma: float,\n slab_sigma: float,\n prior_pi: float,\n) -> torch.Tensor:\n \"\"\"KL divergence for spike-and-slab variational posterior.\n\n The variational posterior is:\n q(w, z) = Bernoulli(z; sigmoid(z_logit)) · N(w; mu, sigma)\n\n The prior is:\n p(w, z) = (pi·N(0, spike_sigma) + (1-pi)·N(0, slab_sigma))\n\n We use the closed-form approximation from Louizos et al. (2017) and\n the practical version used in the kwyk spike-and-slab dropout.\n \"\"\"\n z = torch.sigmoid(z_logit)\n\n # Log-likelihood under spike and slab components\n log_spike = -0.5 * math.log(2 * math.pi * spike_sigma**2) - (\n mu**2 + sigma**2\n ) / (2 * spike_sigma**2)\n log_slab = -0.5 * math.log(2 * math.pi * slab_sigma**2) - (\n mu**2 + sigma**2\n ) / (2 * slab_sigma**2)\n\n # Entropy of the Bernoulli gate\n entropy_z = -(z * torch.log(z + 1e-8) + (1 - z) * torch.log(1 - z + 1e-8))\n\n # KL = E_q[log q - log p]\n # log q(w|z=slab) - log p(w) where p is the mixture\n kl_per_weight = (\n z * (-0.5 * torch.log(2 * math.pi * sigma**2 + 1e-8) - 0.5 - log_slab)\n + (1 - z) * (-log_spike)\n - entropy_z\n + z * math.log(1 - prior_pi + 1e-8)\n + (1 - z) * math.log(prior_pi + 1e-8)\n )\n\n return kl_per_weight.sum()\n\n\n# ---------------------------------------------------------------------------\n# Bayesian layers\n# ---------------------------------------------------------------------------\n\n\nclass BayesianConv3d(PyroModule):\n \"\"\"3-D convolution with learnable weight distribution (Pyro).\n\n Parameters\n ----------\n in_channels, out_channels : int\n Standard convolution channel counts.\n kernel_size : int\n Cubic kernel side length.\n stride, padding, dilation : int\n Standard ``nn.Conv3d`` arguments.\n bias : bool\n Whether to include a deterministic bias term.\n prior_type : str\n ``\"standard_normal\"`` (σ=1), ``\"laplace\"`` (tight Normal σ=0.1),\n or ``\"spike_and_slab\"`` (mixture prior with learnable gates).\n spike_sigma : float\n Spike component σ for spike-and-slab prior (default 0.001).\n slab_sigma : float\n Slab component σ for spike-and-slab prior (default 1.0).\n prior_pi : float\n Prior probability of the spike component (default 0.5).\n \"\"\"\n\n def __init__(\n self,\n in_channels: int,\n out_channels: int,\n kernel_size: int = 3,\n stride: int = 1,\n padding: int = 0,\n dilation: int = 1,\n bias: bool = True,\n prior_type: str = \"standard_normal\",\n spike_sigma: float = 0.001,\n slab_sigma: float = 1.0,\n prior_pi: float = 0.5,\n ) -> None:\n super().__init__()\n self.in_channels = in_channels\n self.out_channels = out_channels\n self.kernel_size = kernel_size\n self.stride = stride\n self.padding = padding\n self.dilation = dilation\n self.prior_type = prior_type\n self.spike_sigma = spike_sigma\n self.slab_sigma = slab_sigma\n self.prior_pi = prior_pi\n\n weight_shape = (\n out_channels,\n in_channels,\n kernel_size,\n kernel_size,\n kernel_size,\n )\n # Kaiming init for mu\n fan_in = in_channels * kernel_size**3\n std_init = math.sqrt(2.0 / fan_in)\n self.weight_mu = PyroParam(\n torch.zeros(weight_shape).normal_(0, std_init),\n constraint=constraints.real,\n )\n self.weight_rho = PyroParam(\n torch.full(weight_shape, -3.0), # softplus(-3) ≈ 0.05\n constraint=constraints.real,\n )\n\n # Spike-and-slab gate logits (one per weight)\n if prior_type == \"spike_and_slab\":\n self.z_logit = PyroParam(\n torch.full(weight_shape, 2.0), # sigmoid(2) ≈ 0.88 → mostly slab\n constraint=constraints.real,\n )\n\n if bias:\n self.bias_mu = PyroParam(\n torch.zeros(out_channels), constraint=constraints.real\n )\n self.bias_rho = PyroParam(\n torch.full((out_channels,), -3.0), constraint=constraints.real\n )\n else:\n self.bias_mu = None\n self.bias_rho = None\n\n if prior_type == \"standard_normal\":\n self.prior_sigma = 1.0\n elif prior_type == \"laplace\":\n self.prior_sigma = 0.1\n else:\n self.prior_sigma = slab_sigma # used as fallback only\n self.kl: torch.Tensor = torch.tensor(0.0)\n\n @property\n def weight_sigma(self) -> torch.Tensor:\n return F.softplus(self.weight_rho)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n weight = pyro.sample(\n f\"{self._pyro_name}.weight\",\n dist.Normal(self.weight_mu, self.weight_sigma + 1e-8).to_event(\n self.weight_mu.dim()\n ),\n )\n\n if self.prior_type == \"spike_and_slab\":\n # Apply spike-and-slab mask: sample Bernoulli gate, mask weights\n z_prob = torch.sigmoid(self.z_logit)\n z_mask = torch.bernoulli(z_prob)\n weight = weight * z_mask\n self.kl = _kl_spike_and_slab(\n self.weight_mu,\n self.weight_sigma,\n self.z_logit,\n self.spike_sigma,\n self.slab_sigma,\n self.prior_pi,\n )\n else:\n self.kl = _kl_normal_normal(\n self.weight_mu, self.weight_sigma, 0.0, self.prior_sigma\n )\n\n bias = None\n if self.bias_mu is not None:\n bias_sigma = F.softplus(self.bias_rho)\n bias = pyro.sample(\n f\"{self._pyro_name}.bias\",\n dist.Normal(self.bias_mu, bias_sigma + 1e-8).to_event(1),\n )\n self.kl = self.kl + _kl_normal_normal(\n self.bias_mu, bias_sigma, 0.0, self.prior_sigma\n )\n\n return F.conv3d(x, weight, bias, self.stride, self.padding, self.dilation)\n\n\nclass BayesianLinear(PyroModule):\n \"\"\"Fully-connected layer with learnable weight distribution (Pyro).\n\n Parameters\n ----------\n in_features, out_features : int\n Standard ``nn.Linear`` dimensions.\n bias : bool\n Whether to include a deterministic bias term.\n prior_type : str\n ``\"standard_normal\"``, ``\"laplace\"``, or ``\"spike_and_slab\"``.\n spike_sigma : float\n Spike component σ for spike-and-slab prior (default 0.001).\n slab_sigma : float\n Slab component σ for spike-and-slab prior (default 1.0).\n prior_pi : float\n Prior probability of the spike component (default 0.5).\n \"\"\"\n\n def __init__(\n self,\n in_features: int,\n out_features: int,\n bias: bool = True,\n prior_type: str = \"standard_normal\",\n spike_sigma: float = 0.001,\n slab_sigma: float = 1.0,\n prior_pi: float = 0.5,\n ) -> None:\n super().__init__()\n self.in_features = in_features\n self.out_features = out_features\n self.prior_type = prior_type\n self.spike_sigma = spike_sigma\n self.slab_sigma = slab_sigma\n self.prior_pi = prior_pi\n\n std_init = math.sqrt(2.0 / in_features)\n self.weight_mu = PyroParam(\n torch.zeros(out_features, in_features).normal_(0, std_init),\n constraint=constraints.real,\n )\n self.weight_rho = PyroParam(\n torch.full((out_features, in_features), -3.0),\n constraint=constraints.real,\n )\n\n if prior_type == \"spike_and_slab\":\n self.z_logit = PyroParam(\n torch.full((out_features, in_features), 2.0),\n constraint=constraints.real,\n )\n\n if bias:\n self.bias_mu = PyroParam(\n torch.zeros(out_features), constraint=constraints.real\n )\n self.bias_rho = PyroParam(\n torch.full((out_features,), -3.0), constraint=constraints.real\n )\n else:\n self.bias_mu = None\n self.bias_rho = None\n\n if prior_type == \"standard_normal\":\n self.prior_sigma = 1.0\n elif prior_type == \"laplace\":\n self.prior_sigma = 0.1\n else:\n self.prior_sigma = slab_sigma\n self.kl: torch.Tensor = torch.tensor(0.0)\n\n @property\n def weight_sigma(self) -> torch.Tensor:\n return F.softplus(self.weight_rho)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n weight = pyro.sample(\n f\"{self._pyro_name}.weight\",\n dist.Normal(self.weight_mu, self.weight_sigma + 1e-8).to_event(2),\n )\n\n if self.prior_type == \"spike_and_slab\":\n z_prob = torch.sigmoid(self.z_logit)\n z_mask = torch.bernoulli(z_prob)\n weight = weight * z_mask\n self.kl = _kl_spike_and_slab(\n self.weight_mu,\n self.weight_sigma,\n self.z_logit,\n self.spike_sigma,\n self.slab_sigma,\n self.prior_pi,\n )\n else:\n self.kl = _kl_normal_normal(\n self.weight_mu, self.weight_sigma, 0.0, self.prior_sigma\n )\n\n bias = None\n if self.bias_mu is not None:\n bias_sigma = F.softplus(self.bias_rho)\n bias = pyro.sample(\n f\"{self._pyro_name}.bias\",\n dist.Normal(self.bias_mu, bias_sigma + 1e-8).to_event(1),\n )\n self.kl = self.kl + _kl_normal_normal(\n self.bias_mu, bias_sigma, 0.0, self.prior_sigma\n )\n\n return F.linear(x, weight, bias)\n" }, { "path": "nobrainer/models/bayesian/utils.py", "content": "\"\"\"Utility functions for Bayesian models.\"\"\"\n\nfrom __future__ import annotations\n\nimport torch\n\nfrom .layers import BayesianConv3d, BayesianLinear\nfrom .vwn_layers import ConcreteDropout3d, FFGConv3d\n\n\ndef accumulate_kl(model: torch.nn.Module) -> torch.Tensor:\n \"\"\"Sum KL divergence from all Bayesian layers in ``model``.\n\n Works with both Pyro-based models (BayesianConv3d, BayesianLinear)\n and VWN/FFG models (FFGConv3d, ConcreteDropout3d).\n\n Parameters\n ----------\n model : nn.Module\n A model containing one or more Bayesian layers.\n\n Returns\n -------\n torch.Tensor\n Scalar KL sum.\n \"\"\"\n kl = torch.tensor(0.0)\n for m in model.modules():\n # Pyro-based layers\n if isinstance(m, (BayesianConv3d, BayesianLinear)):\n kl = kl + m.kl\n # VWN/FFG layers\n elif isinstance(m, FFGConv3d):\n kl = kl + m.kl\n # Concrete dropout regularization\n elif isinstance(m, ConcreteDropout3d):\n kl = kl + m.kl_divergence()\n return kl\n" }, { "path": "nobrainer/models/bayesian/vwn_layers.py", "content": "\"\"\"Fully Factorized Gaussian (FFG) layers with local reparameterization.\n\nThese layers implement the convolution used in McClure et al. (2019),\nSection 2.2.3.2 (\"Spike-and-Slab Dropout with Learned Model Uncertainty\"):\n\n* Each weight has learnable mean ``μ_{f,t}`` and std ``σ_{f,t}``\n* **Local reparameterization trick** (Kingma et al. 2015): instead of\n sampling weights, the output distribution is computed directly:\n ``output ~ N(conv(x, μ), conv(x², σ²))`` (Eqs. 12-14)\n* The **spike-and-slab dropout (SSD)** model combines this with\n **concrete dropout** (Gal et al. 2017): ``output_v = b_f · (g_f * h)_v``\n where ``b_f`` is a per-filter concrete dropout mask (Eq. 11).\n\nThe KL divergence has two terms (Eq. 16 in paper):\n 1. Bernoulli KL for concrete dropout gates (Eq. 17)\n 2. Gaussian KL for each weight: ``KL(N(μ,σ) || N(μ_prior, σ_prior))`` (Eq. 18)\n\nPrior parameters from the paper: ``p_prior=0.5, μ_prior=0, σ_prior=0.1``\n\nTwo dropout variants:\n* **Bernoulli dropout** — standard ``nn.Dropout3d``, fixed rate (BD model)\n* **Concrete dropout** — per-filter learnable drop rate (SSD model)\n\"\"\"\n\nfrom __future__ import annotations\n\nimport math\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\nclass FFGConv3d(nn.Module):\n \"\"\"3-D convolution with Variational Weight Normalization + learned sigma.\n\n Verified against the actual kwyk trained model (``neuronets/kwyk:latest``).\n Each layer stores: ``v``, ``g``, ``kernel_a``, ``bias_m``, ``bias_a``.\n\n **Mean weights** use weight normalization (Salimans & Kingma 2016):\n ``kernel_m = g · v / ||v||``\n\n **Sigma** is learned per weight:\n ``kernel_sigma = |kernel_a|``\n\n During stochastic forward passes (``mc=True``), the **local\n reparameterization trick** (Kingma et al. 2015, Eqs. 12-14 in\n McClure et al. 2019) computes the output distribution directly:\n\n ``μ* = conv(x, kernel_m)``\n ``σ*² = conv(x², kernel_sigma²)``\n ``output = μ* + σ* · ε, ε ~ N(0, 1)``\n\n In deterministic mode (``mc=False``) only the mean path is used.\n\n Parameters\n ----------\n in_channels, out_channels : int\n Standard convolution channel counts.\n kernel_size : int\n Cubic kernel side length.\n stride, padding, dilation : int\n Standard ``nn.Conv3d`` arguments.\n bias : bool\n Whether to include a bias term (with its own sigma).\n sigma_init : float\n Initial value for ``|kernel_a|`` (default 1e-4, matching kwyk).\n prior_mu : float\n Prior mean for KL (default 0.0).\n prior_sigma : float\n Prior std for KL (default 0.1, matching paper Eq. 18).\n \"\"\"\n\n def __init__(\n self,\n in_channels: int,\n out_channels: int,\n kernel_size: int = 3,\n stride: int = 1,\n padding: int = 0,\n dilation: int = 1,\n bias: bool = True,\n sigma_init: float = 1e-4,\n prior_mu: float = 0.0,\n prior_sigma: float = 0.1,\n ) -> None:\n super().__init__()\n self.in_channels = in_channels\n self.out_channels = out_channels\n self.kernel_size = kernel_size\n self.stride = stride\n self.padding = padding\n self.dilation = dilation\n self.prior_mu = prior_mu\n self.prior_sigma = prior_sigma\n\n k = kernel_size\n weight_shape = (out_channels, in_channels, k, k, k)\n g_shape = (out_channels, 1, 1, 1, 1)\n\n # Weight normalization: kernel_m = g * v / ||v||\n self.v = nn.Parameter(torch.empty(weight_shape))\n nn.init.kaiming_normal_(self.v, mode=\"fan_in\", nonlinearity=\"relu\")\n self.g = nn.Parameter(torch.full(g_shape, math.sqrt(2.0)))\n\n # Learned sigma: kernel_sigma = |kernel_a|\n self.kernel_a = nn.Parameter(torch.full(weight_shape, sigma_init))\n\n if bias:\n self.bias_m = nn.Parameter(torch.zeros(out_channels))\n self.bias_a = nn.Parameter(torch.full((out_channels,), sigma_init))\n else:\n self.register_parameter(\"bias_m\", None)\n self.register_parameter(\"bias_a\", None)\n\n # Accumulated KL (updated each forward pass)\n self.kl: torch.Tensor = torch.tensor(0.0)\n\n @property\n def kernel_m(self) -> torch.Tensor:\n \"\"\"Mean weight: ``g · v / ||v||``.\"\"\"\n v_norm = F.normalize(self.v.flatten(1), dim=1).view_as(self.v)\n return self.g * v_norm\n\n @property\n def weight_sigma(self) -> torch.Tensor:\n \"\"\"Weight std: ``|kernel_a|``.\"\"\"\n return torch.abs(self.kernel_a)\n\n def forward(self, x: torch.Tensor, mc: bool = True) -> torch.Tensor:\n \"\"\"Forward pass with optional stochastic sampling.\"\"\"\n km = self.kernel_m\n out_mean = F.conv3d(\n x, km, self.bias_m, self.stride, self.padding, self.dilation\n )\n\n if not mc:\n return out_mean\n\n # Local reparameterization trick (Eqs. 12-14)\n sigma = self.weight_sigma\n out_var = F.conv3d(\n x.pow(2), sigma.pow(2), None, self.stride, self.padding, self.dilation\n )\n if self.bias_a is not None:\n bias_sigma = torch.abs(self.bias_a)\n out_var = out_var + bias_sigma.pow(2).view(1, -1, 1, 1, 1)\n\n noise = torch.randn_like(out_mean)\n out = out_mean + torch.sqrt(out_var + 1e-8) * noise\n\n # KL(N(kernel_m, sigma) || N(prior_mu, prior_sigma)) — Eq. 18\n self.kl = (\n torch.log(self.prior_sigma / (sigma + 1e-8))\n + (sigma.pow(2) + (km - self.prior_mu).pow(2)) / (2 * self.prior_sigma**2)\n - 0.5\n ).sum()\n\n return out\n\n\n# Backward-compatible alias\nVWNConv3d = FFGConv3d\n\n\nclass ConcreteDropout3d(nn.Module):\n \"\"\"Concrete dropout (Gal et al. 2017) with per-filter learnable rate.\n\n Instead of a fixed dropout probability, each output filter learns its\n own drop rate ``p`` via a continuous relaxation of Bernoulli sampling\n (Eq. 10 in McClure et al. 2019).\n\n Parameters\n ----------\n n_filters : int\n Number of filters (one ``p`` per filter).\n temperature : float\n Concrete distribution temperature (default 0.02, matching paper).\n init_p : float\n Initial dropout probability (default 0.9, matching kwyk code).\n prior_p : float\n Prior dropout probability for KL (default 0.5, matching paper).\n \"\"\"\n\n def __init__(\n self,\n n_filters: int,\n temperature: float = 0.02,\n init_p: float = 0.9,\n prior_p: float = 0.5,\n ) -> None:\n super().__init__()\n self.temperature = temperature\n self.prior_p = prior_p\n # Store as raw logit; p = sigmoid(p_logit) to keep in (0, 1)\n init_logit = math.log(init_p / (1 - init_p + 1e-8))\n self.p_logit = nn.Parameter(torch.full((n_filters,), init_logit))\n\n @property\n def p(self) -> torch.Tensor:\n \"\"\"Per-filter dropout probabilities, clamped to [0.05, 0.95].\"\"\"\n return torch.sigmoid(self.p_logit).clamp(0.05, 0.95)\n\n def forward(self, x: torch.Tensor, mc: bool = True) -> torch.Tensor:\n \"\"\"Apply concrete dropout (Eq. 10).\n\n Parameters\n ----------\n x : Tensor\n Input ``(B, C, D, H, W)``.\n mc : bool\n If True, sample from concrete distribution.\n If False, scale by ``p`` (expectation).\n \"\"\"\n p = self.p.view(1, -1, 1, 1, 1)\n\n if not mc:\n return x * p\n\n # Concrete relaxation of Bernoulli (Eq. 10)\n eps = 1e-8\n noise = torch.rand_like(x[:1]) # (1, C, D, H, W)\n z = torch.sigmoid(\n (\n torch.log(p + eps)\n - torch.log(1 - p + eps)\n + torch.log(noise + eps)\n - torch.log(1 - noise + eps)\n )\n / self.temperature\n )\n return x * z\n\n def kl_divergence(self) -> torch.Tensor:\n \"\"\"KL(q_p || p_prior) for Bernoulli distributions (Eq. 17).\"\"\"\n p = self.p\n pp = self.prior_p\n eps = 1e-8\n return (\n p * torch.log(p / (pp + eps) + eps)\n + (1 - p) * torch.log((1 - p) / (1 - pp + eps) + eps)\n ).sum()\n\n def regularization(self) -> torch.Tensor:\n \"\"\"Alias for kl_divergence (backward compat).\"\"\"\n return self.kl_divergence()\n" }, { "path": "nobrainer/models/bayesian/warmstart.py", "content": "\"\"\"Warm-start a Bayesian model from a trained deterministic model.\"\"\"\n\nfrom __future__ import annotations\n\nimport logging\nfrom pathlib import Path\n\nimport torch\nimport torch.nn as nn\n\nfrom nobrainer.models.bayesian.layers import BayesianConv3d\nfrom nobrainer.models.bayesian.vwn_layers import FFGConv3d\n\nlogger = logging.getLogger(__name__)\n\n\ndef warmstart_bayesian_from_deterministic(\n bayesian_model: nn.Module,\n deterministic_model: nn.Module,\n initial_rho: float = -3.0,\n) -> int:\n \"\"\"Transfer deterministic Conv3d weights to BayesianConv3d weight_mu.\n\n Matches layers by position (not name) since the deterministic MeshNet\n uses ``nn.Sequential`` (``encoder.N.block.0``) while the Bayesian\n MeshNet uses named attributes (``layer_N.conv``).\n\n For each matching pair:\n\n * **Conv3d -> BayesianConv3d**: copies ``weight`` to ``weight_mu``,\n fills ``weight_rho`` with *initial_rho*, and handles bias if\n present.\n * **BatchNorm3d -> BatchNorm3d**: copies ``weight``, ``bias``,\n ``running_mean``, and ``running_var``.\n\n Parameters\n ----------\n bayesian_model : nn.Module\n Target Bayesian model whose parameters will be overwritten.\n deterministic_model : nn.Module\n Source deterministic model with trained weights.\n initial_rho : float, optional\n Value to fill ``weight_rho`` (and ``bias_rho``) with.\n ``softplus(-3.0) ≈ 0.05``. Default is ``-3.0``.\n\n Returns\n -------\n int\n Number of layers whose weights were transferred.\n \"\"\"\n # First try name-based matching (works if architectures share naming)\n transferred = _transfer_by_name(bayesian_model, deterministic_model, initial_rho)\n if transferred > 0:\n return transferred\n\n # Fall back to positional matching (different naming conventions)\n return _transfer_by_position(bayesian_model, deterministic_model, initial_rho)\n\n\ndef _transfer_by_name(\n bayesian_model: nn.Module,\n deterministic_model: nn.Module,\n initial_rho: float,\n) -> int:\n \"\"\"Match layers by module name.\"\"\"\n det_modules = dict(deterministic_model.named_modules())\n bayes_modules = dict(bayesian_model.named_modules())\n transferred = 0\n\n for name, bayes_mod in bayes_modules.items():\n if name not in det_modules:\n continue\n det_mod = det_modules[name]\n transferred += _transfer_pair(det_mod, bayes_mod, name, initial_rho)\n\n if transferred > 0:\n logger.info(\n \"Warm-started %d layers (name-matched) from deterministic model.\",\n transferred,\n )\n return transferred\n\n\ndef _transfer_by_position(\n bayesian_model: nn.Module,\n deterministic_model: nn.Module,\n initial_rho: float,\n) -> int:\n \"\"\"Match Conv3d/BN layers by position (order of appearance).\"\"\"\n # Collect Conv3d layers from deterministic model\n det_convs = [\n (n, m)\n for n, m in deterministic_model.named_modules()\n if isinstance(m, nn.Conv3d)\n ]\n det_bns = [\n (n, m)\n for n, m in deterministic_model.named_modules()\n if isinstance(m, nn.BatchNorm3d)\n ]\n\n # Collect BayesianConv3d layers from Bayesian model\n bayes_convs = [\n (n, m)\n for n, m in bayesian_model.named_modules()\n if isinstance(m, BayesianConv3d)\n ]\n bayes_bns = [\n (n, m)\n for n, m in bayesian_model.named_modules()\n if isinstance(m, nn.BatchNorm3d)\n ]\n\n transferred = 0\n\n # Transfer Conv3d -> BayesianConv3d by position\n for i, ((det_name, det_conv), (bay_name, bay_conv)) in enumerate(\n zip(det_convs, bayes_convs)\n ):\n if det_conv.weight.shape != bay_conv.weight_mu.shape:\n logger.warning(\n \"Shape mismatch at position %d: det %s %s vs bay %s %s\",\n i,\n det_name,\n det_conv.weight.shape,\n bay_name,\n bay_conv.weight_mu.shape,\n )\n continue\n\n bay_conv.weight_mu.data.copy_(det_conv.weight.data)\n bay_conv.weight_rho.data.fill_(initial_rho)\n\n if det_conv.bias is not None and bay_conv.bias_mu is not None:\n bay_conv.bias_mu.data.copy_(det_conv.bias.data)\n bay_conv.bias_rho.data.fill_(initial_rho)\n\n transferred += 1\n logger.debug(\"Transferred Conv3d[%d] %s -> %s\", i, det_name, bay_name)\n\n # Transfer BatchNorm3d by position\n for i, ((det_name, det_bn), (bay_name, bay_bn)) in enumerate(\n zip(det_bns, bayes_bns)\n ):\n if det_bn.weight is not None and bay_bn.weight is not None:\n bay_bn.weight.data.copy_(det_bn.weight.data)\n if det_bn.bias is not None and bay_bn.bias is not None:\n bay_bn.bias.data.copy_(det_bn.bias.data)\n if det_bn.running_mean is not None:\n bay_bn.running_mean.copy_(det_bn.running_mean)\n if det_bn.running_var is not None:\n bay_bn.running_var.copy_(det_bn.running_var)\n\n transferred += 1\n logger.debug(\"Transferred BatchNorm3d[%d] %s -> %s\", i, det_name, bay_name)\n\n logger.info(\n \"Warm-started %d layers (position-matched) from deterministic model.\",\n transferred,\n )\n return transferred\n\n\ndef _transfer_pair(\n det_mod: nn.Module,\n bayes_mod: nn.Module,\n name: str,\n initial_rho: float,\n) -> int:\n \"\"\"Transfer weights for a single matching pair. Returns 1 if transferred.\"\"\"\n is_conv = isinstance(det_mod, nn.Conv3d)\n is_bayes_conv = isinstance(bayes_mod, BayesianConv3d)\n\n if is_conv and is_bayes_conv:\n bayes_mod.weight_mu.data.copy_(det_mod.weight.data)\n bayes_mod.weight_rho.data.fill_(initial_rho)\n\n if det_mod.bias is not None and bayes_mod.bias_mu is not None:\n bayes_mod.bias_mu.data.copy_(det_mod.bias.data)\n bayes_mod.bias_rho.data.fill_(initial_rho)\n\n logger.debug(\"Transferred Conv3d weights: %s\", name)\n return 1\n\n if isinstance(det_mod, nn.BatchNorm3d) and isinstance(bayes_mod, nn.BatchNorm3d):\n if det_mod.weight is not None and bayes_mod.weight is not None:\n bayes_mod.weight.data.copy_(det_mod.weight.data)\n if det_mod.bias is not None and bayes_mod.bias is not None:\n bayes_mod.bias.data.copy_(det_mod.bias.data)\n if det_mod.running_mean is not None:\n bayes_mod.running_mean.copy_(det_mod.running_mean)\n if det_mod.running_var is not None:\n bayes_mod.running_var.copy_(det_mod.running_var)\n\n logger.debug(\"Transferred BatchNorm3d params: %s\", name)\n return 1\n\n return 0\n\n\n# ---------------------------------------------------------------------------\n# KWYK MeshNet warm-start (VWN-based, no Pyro)\n# ---------------------------------------------------------------------------\n\n\ndef warmstart_kwyk_from_deterministic(\n kwyk_model: nn.Module,\n det_weights_path: str | Path,\n get_model_fn=None,\n) -> int:\n \"\"\"Transfer deterministic MeshNet weights to a KWYKMeshNet.\n\n For each VWN conv layer, the deterministic weight ``w`` is decomposed\n into weight normalization form: ``v = w``, ``g = ||w||`` per filter.\n The sigma parameters (``kernel_a``) are left at their initial values.\n\n Parameters\n ----------\n kwyk_model : nn.Module\n Target KWYKMeshNet.\n det_weights_path : str or Path\n Path to a deterministic MeshNet ``model.pth``.\n get_model_fn : callable, optional\n Model factory (``nobrainer.models.get``). If None, imported lazily.\n\n Returns\n -------\n int\n Number of layers transferred.\n \"\"\"\n if get_model_fn is None:\n from nobrainer.models import get as get_model_fn\n\n det_weights_path = Path(det_weights_path)\n state = torch.load(det_weights_path, weights_only=True)\n\n # Separate encoder conv weights from classifier — sorted() puts\n # \"classifier\" before \"encoder\" alphabetically, so we must filter\n # to avoid misaligning the layer pairing.\n encoder_convs = []\n classifier_w = None\n classifier_b = None\n for k in sorted(state.keys()):\n v = state[k]\n if k == \"classifier.weight\" and v.ndim == 5:\n classifier_w = v\n elif k == \"classifier.bias\":\n classifier_b = v\n elif \"weight\" in k and v.ndim == 5:\n encoder_convs.append((k, v))\n\n # Collect FFGConv3d layers from the kwyk model\n kwyk_convs = [\n (n, m) for n, m in kwyk_model.named_modules() if isinstance(m, FFGConv3d)\n ]\n\n transferred = 0\n for (det_name, det_w), (kwyk_name, kwyk_conv) in zip(encoder_convs, kwyk_convs):\n if det_w.shape != kwyk_conv.v.shape:\n logger.warning(\n \"Shape mismatch: %s %s vs %s %s\",\n det_name,\n det_w.shape,\n kwyk_name,\n kwyk_conv.v.shape,\n )\n continue\n\n # Decompose w into weight-norm form: v = w, g = ||w|| per filter\n kwyk_conv.v.data.copy_(det_w)\n # g = ||v|| per output filter (over in_channels * k * k * k)\n norms = det_w.flatten(1).norm(dim=1).view_as(kwyk_conv.g)\n kwyk_conv.g.data.copy_(norms)\n\n transferred += 1\n logger.debug(\"Transferred Conv3d %s -> %s\", det_name, kwyk_name)\n\n # Transfer classifier separately (regular Conv3d, not FFGConv3d)\n if classifier_w is not None and hasattr(kwyk_model, \"classifier\"):\n kwyk_model.classifier.weight.data.copy_(classifier_w)\n if classifier_b is not None and kwyk_model.classifier.bias is not None:\n kwyk_model.classifier.bias.data.copy_(classifier_b)\n transferred += 1\n logger.debug(\"Transferred classifier\")\n\n logger.info(\"Warm-started %d layers from deterministic model.\", transferred)\n return transferred\n" }, { "path": "nobrainer/models/generative/__init__.py", "content": "\"\"\"Generative model sub-package (Phase 5 — US3).\"\"\"\n\nfrom .dcgan import DCGAN, dcgan\nfrom .progressivegan import ProgressiveGAN, progressivegan\n\n__all__ = [\n \"DCGAN\",\n \"ProgressiveGAN\",\n \"dcgan\",\n \"progressivegan\",\n]\n" }, { "path": "nobrainer/models/generative/dcgan.py", "content": "\"\"\"DCGAN implemented as a PyTorch Lightning module.\n\nStandard alternating generator/discriminator training using BCE loss.\n\nReference\n---------\nRadford A. et al., \"Unsupervised Representation Learning with Deep\nConvolutional Generative Adversarial Networks\", ICLR 2016. arXiv:1511.06434.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom typing import Any\n\nimport pytorch_lightning as pl\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n# ---------------------------------------------------------------------------\n# Building blocks\n# ---------------------------------------------------------------------------\n\n\nclass _GenBlock(nn.Module):\n \"\"\"Transposed-conv + BN + ReLU.\"\"\"\n\n def __init__(\n self,\n in_ch: int,\n out_ch: int,\n kernel_size: int = 4,\n stride: int = 2,\n padding: int = 1,\n ) -> None:\n super().__init__()\n self.block = nn.Sequential(\n nn.ConvTranspose3d(\n in_ch, out_ch, kernel_size, stride=stride, padding=padding\n ),\n nn.BatchNorm3d(out_ch),\n nn.ReLU(inplace=True),\n )\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return self.block(x)\n\n\nclass _DiscBlock(nn.Module):\n \"\"\"Conv + (optional BN) + LeakyReLU.\"\"\"\n\n def __init__(\n self,\n in_ch: int,\n out_ch: int,\n kernel_size: int = 4,\n stride: int = 2,\n padding: int = 1,\n use_bn: bool = True,\n ) -> None:\n super().__init__()\n layers: list[nn.Module] = [\n nn.Conv3d(in_ch, out_ch, kernel_size, stride=stride, padding=padding),\n ]\n if use_bn:\n layers.append(nn.BatchNorm3d(out_ch))\n layers.append(nn.LeakyReLU(0.2, inplace=True))\n self.block = nn.Sequential(*layers)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return self.block(x)\n\n\n# ---------------------------------------------------------------------------\n# Generator\n# ---------------------------------------------------------------------------\n\n\nclass _DCGenerator(nn.Module):\n \"\"\"4-level transposed-conv generator; outputs (N, 1, 32, 32, 32).\"\"\"\n\n def __init__(self, latent_size: int = 128, n_filters: int = 64) -> None:\n super().__init__()\n nf = n_filters\n self.net = nn.Sequential(\n # latent (N, Z, 1, 1, 1) → (N, nf*8, 4, 4, 4)\n nn.ConvTranspose3d(latent_size, nf * 8, kernel_size=4, stride=1, padding=0),\n nn.BatchNorm3d(nf * 8),\n nn.ReLU(inplace=True),\n # (N, nf*8, 4, 4, 4) → (N, nf*4, 8, 8, 8)\n _GenBlock(nf * 8, nf * 4),\n # (N, nf*4, 8, 8, 8) → (N, nf*2, 16, 16, 16)\n _GenBlock(nf * 4, nf * 2),\n # (N, nf*2, 16, 16, 16) → (N, nf, 32, 32, 32)\n _GenBlock(nf * 2, nf),\n # (N, nf, 32, 32, 32) → (N, 1, 32, 32, 32)\n nn.ConvTranspose3d(nf, 1, kernel_size=4, stride=2, padding=1),\n nn.Tanh(),\n )\n\n def forward(self, z: torch.Tensor) -> torch.Tensor:\n return self.net(z.view(*z.shape, 1, 1, 1))\n\n\n# ---------------------------------------------------------------------------\n# Discriminator\n# ---------------------------------------------------------------------------\n\n\nclass _DCDiscriminator(nn.Module):\n \"\"\"4-level conv discriminator; expects (N, 1, 64, 64, 64).\"\"\"\n\n def __init__(self, n_filters: int = 64) -> None:\n super().__init__()\n nf = n_filters\n self.net = nn.Sequential(\n # (N, 1, 64, 64, 64) → (N, nf, 32, 32, 32); no BN on first layer\n _DiscBlock(1, nf, use_bn=False),\n # → (N, nf*2, 16, 16, 16)\n _DiscBlock(nf, nf * 2),\n # → (N, nf*4, 8, 8, 8)\n _DiscBlock(nf * 2, nf * 4),\n # → (N, nf*8, 4, 4, 4)\n _DiscBlock(nf * 4, nf * 8),\n # → (N, 1, 1, 1, 1)\n nn.Conv3d(nf * 8, 1, kernel_size=4, stride=1, padding=0),\n nn.Flatten(),\n )\n\n def forward(self, img: torch.Tensor) -> torch.Tensor:\n return self.net(img)\n\n\n# ---------------------------------------------------------------------------\n# Lightning module\n# ---------------------------------------------------------------------------\n\n\nclass DCGAN(pl.LightningModule):\n \"\"\"DCGAN as a PyTorch Lightning module.\n\n Uses binary cross-entropy (non-saturating G loss) with standard\n alternating G/D updates.\n\n Parameters\n ----------\n latent_size : int\n Dimension of the latent noise vector.\n n_filters : int\n Base channel count for generator and discriminator.\n lr : float\n Learning rate for Adam.\n beta1 : float\n Adam beta1.\n \"\"\"\n\n def __init__(\n self,\n latent_size: int = 128,\n n_filters: int = 64,\n lr: float = 2e-4,\n beta1: float = 0.5,\n ) -> None:\n super().__init__()\n self.save_hyperparameters()\n self.latent_size = latent_size\n self.lr = lr\n self.beta1 = beta1\n\n self.generator = _DCGenerator(latent_size, n_filters)\n self.discriminator = _DCDiscriminator(n_filters)\n self.automatic_optimization = False\n\n # Fixed noise for visualisation\n self._fixed_z = None\n\n def _sample_z(self, n: int) -> torch.Tensor:\n return torch.randn(n, self.latent_size, device=self.device)\n\n def training_step(self, batch: Any, batch_idx: int) -> None:\n opt_g, opt_d = self.optimizers()\n\n real = batch[\"image\"] if isinstance(batch, dict) else batch[0]\n b = real.size(0)\n\n real_label = torch.ones(b, 1, device=self.device)\n fake_label = torch.zeros(b, 1, device=self.device)\n\n # --- Discriminator step ---\n opt_d.zero_grad()\n z = self._sample_z(b)\n fake = self.generator(z).detach()\n # Resize real to discriminator input size if necessary\n if real.shape[-1] != 64:\n real_in = F.interpolate(\n real, size=(64, 64, 64), mode=\"trilinear\", align_corners=False\n )\n else:\n real_in = real\n if fake.shape[-1] != 64:\n fake_in = F.interpolate(\n fake, size=(64, 64, 64), mode=\"trilinear\", align_corners=False\n )\n else:\n fake_in = fake\n d_real = F.binary_cross_entropy_with_logits(\n self.discriminator(real_in), real_label\n )\n d_fake = F.binary_cross_entropy_with_logits(\n self.discriminator(fake_in), fake_label\n )\n d_loss = (d_real + d_fake) * 0.5\n self.manual_backward(d_loss)\n opt_d.step()\n\n # --- Generator step ---\n opt_g.zero_grad()\n z = self._sample_z(b)\n fake = self.generator(z)\n if fake.shape[-1] != 64:\n fake_in = F.interpolate(\n fake, size=(64, 64, 64), mode=\"trilinear\", align_corners=False\n )\n else:\n fake_in = fake\n g_loss = F.binary_cross_entropy_with_logits(\n self.discriminator(fake_in), real_label\n )\n self.manual_backward(g_loss)\n opt_g.step()\n\n self.log_dict({\"g_loss\": g_loss, \"d_loss\": d_loss}, prog_bar=True)\n\n def configure_optimizers(self):\n opt_g = torch.optim.Adam(\n self.generator.parameters(), lr=self.lr, betas=(self.beta1, 0.999)\n )\n opt_d = torch.optim.Adam(\n self.discriminator.parameters(), lr=self.lr, betas=(self.beta1, 0.999)\n )\n return [opt_g, opt_d]\n\n\ndef dcgan(\n latent_size: int = 128,\n n_filters: int = 64,\n **kwargs,\n) -> DCGAN:\n \"\"\"Factory function for :class:`DCGAN`.\"\"\"\n return DCGAN(latent_size=latent_size, n_filters=n_filters, **kwargs)\n\n\n__all__ = [\"DCGAN\", \"dcgan\"]\n" }, { "path": "nobrainer/models/generative/progressivegan.py", "content": "\"\"\"ProgressiveGAN implemented as a PyTorch Lightning module.\n\nGrows the generator and discriminator from 4³ to the target resolution in\nstages. Each stage fades in a new layer using a learnable ``alpha``\nparameter that rises from 0 to 1 during the fade-in phase.\n\nReference\n---------\nKarras T. et al., \"Progressive Growing of GANs for Improved Quality,\nStability, and Variation\", ICLR 2018. arXiv:1710.10196.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom typing import Any\n\nimport pytorch_lightning as pl\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n# ---------------------------------------------------------------------------\n# Building blocks\n# ---------------------------------------------------------------------------\n\n\ndef _pixel_norm(x: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:\n \"\"\"Pixel-wise feature vector normalisation (ProGAN style).\"\"\"\n return x / (x.pow(2).mean(dim=1, keepdim=True) + eps).sqrt()\n\n\nclass _ConvBlock(nn.Module):\n def __init__(self, in_ch: int, out_ch: int, use_pixel_norm: bool = True) -> None:\n super().__init__()\n self.conv = nn.Conv3d(in_ch, out_ch, kernel_size=3, padding=1)\n self.use_pixel_norm = use_pixel_norm\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n x = F.leaky_relu(self.conv(x), 0.2)\n if self.use_pixel_norm:\n x = _pixel_norm(x)\n return x\n\n\nclass _ToRGB(nn.Module):\n def __init__(self, in_ch: int) -> None:\n super().__init__()\n self.conv = nn.Conv3d(in_ch, 1, kernel_size=1)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return self.conv(x)\n\n\nclass _FromRGB(nn.Module):\n def __init__(self, out_ch: int) -> None:\n super().__init__()\n self.conv = nn.Conv3d(1, out_ch, kernel_size=1)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return F.leaky_relu(self.conv(x), 0.2)\n\n\n# ---------------------------------------------------------------------------\n# Generator\n# ---------------------------------------------------------------------------\n\n\nclass _Generator(nn.Module):\n \"\"\"Progressive generator. Each stage doubles the spatial resolution.\"\"\"\n\n def __init__(\n self,\n latent_size: int,\n fmap_base: int,\n fmap_max: int,\n resolution_schedule: list[int],\n ) -> None:\n super().__init__()\n self.resolution_schedule = resolution_schedule\n self.current_level = 0\n\n def nf(level: int) -> int:\n return min(int(fmap_base / (2**level)), fmap_max)\n\n # Level 0: latent → 4³ feature map\n self.init_block = nn.Sequential(\n nn.ConvTranspose3d(latent_size, nf(0), kernel_size=4, stride=1, padding=0),\n nn.LeakyReLU(0.2),\n _ConvBlock(nf(0), nf(0)),\n )\n self.to_rgb_blocks = nn.ModuleList([_ToRGB(nf(0))])\n self.upsample_blocks = nn.ModuleList()\n\n for level in range(1, len(resolution_schedule)):\n block = nn.Sequential(\n _ConvBlock(nf(level - 1), nf(level)),\n _ConvBlock(nf(level), nf(level)),\n )\n self.upsample_blocks.append(block)\n self.to_rgb_blocks.append(_ToRGB(nf(level)))\n\n self.alpha: float = 1.0\n\n def forward(self, z: torch.Tensor) -> torch.Tensor:\n x = self.init_block(z.view(*z.shape, 1, 1, 1))\n\n if self.current_level == 0:\n return torch.tanh(self.to_rgb_blocks[0](x))\n\n # Grow through levels up to current_level - 1, then fade in last level\n for i in range(self.current_level - 1):\n x = F.interpolate(x, scale_factor=2, mode=\"trilinear\", align_corners=False)\n x = self.upsample_blocks[i](x)\n\n # Fade-in: blend previous RGB with new upsampled RGB\n prev_rgb = self.to_rgb_blocks[self.current_level - 1](x)\n prev_rgb = F.interpolate(\n prev_rgb, scale_factor=2, mode=\"trilinear\", align_corners=False\n )\n\n x = F.interpolate(x, scale_factor=2, mode=\"trilinear\", align_corners=False)\n x = self.upsample_blocks[self.current_level - 1](x)\n new_rgb = self.to_rgb_blocks[self.current_level](x)\n\n out = self.alpha * new_rgb + (1.0 - self.alpha) * prev_rgb\n return torch.tanh(out)\n\n\n# ---------------------------------------------------------------------------\n# Discriminator\n# ---------------------------------------------------------------------------\n\n\nclass _Discriminator(nn.Module):\n \"\"\"Progressive discriminator. Mirror of the generator.\"\"\"\n\n def __init__(\n self,\n fmap_base: int,\n fmap_max: int,\n resolution_schedule: list[int],\n ) -> None:\n super().__init__()\n self.resolution_schedule = resolution_schedule\n self.current_level = 0\n\n def nf(level: int) -> int:\n return min(int(fmap_base / (2**level)), fmap_max)\n\n # Level 0 (4³): feature → 1 (real/fake)\n self.final_block = nn.Sequential(\n _ConvBlock(nf(0), nf(0), use_pixel_norm=False),\n nn.AdaptiveAvgPool3d(1),\n nn.Flatten(),\n nn.Linear(nf(0), 1),\n )\n self.from_rgb_blocks = nn.ModuleList([_FromRGB(nf(0))])\n self.downsample_blocks = nn.ModuleList()\n\n for level in range(1, len(resolution_schedule)):\n block = nn.Sequential(\n _ConvBlock(nf(level), nf(level), use_pixel_norm=False),\n _ConvBlock(nf(level), nf(level - 1), use_pixel_norm=False),\n )\n self.downsample_blocks.append(block)\n self.from_rgb_blocks.append(_FromRGB(nf(level)))\n\n self.alpha: float = 1.0\n\n def forward(self, img: torch.Tensor) -> torch.Tensor:\n if self.current_level == 0:\n x = self.from_rgb_blocks[0](img)\n return self.final_block(x)\n\n # Fade-in: blend downsampled previous level with new level\n prev_img = F.avg_pool3d(img, kernel_size=2, stride=2)\n prev_x = self.from_rgb_blocks[self.current_level - 1](prev_img)\n\n x = self.from_rgb_blocks[self.current_level](img)\n x = self.downsample_blocks[self.current_level - 1](x)\n x = F.avg_pool3d(x, kernel_size=2, stride=2)\n\n x = self.alpha * x + (1.0 - self.alpha) * prev_x\n\n for i in range(self.current_level - 2, -1, -1):\n x = self.downsample_blocks[i](x)\n x = F.avg_pool3d(x, kernel_size=2, stride=2)\n\n return self.final_block(x)\n\n\n# ---------------------------------------------------------------------------\n# Lightning module\n# ---------------------------------------------------------------------------\n\n\nclass ProgressiveGAN(pl.LightningModule):\n \"\"\"ProgressiveGAN as a PyTorch Lightning module.\n\n Parameters\n ----------\n latent_size : int\n Dimension of the latent noise vector.\n label_size : int\n Conditioning label dimension (0 = unconditional).\n fmap_base : int\n Base feature-map count used to compute per-level channels.\n fmap_max : int\n Maximum feature-map count at any level.\n resolution_schedule : list[int]\n Spatial resolutions to train (e.g. ``[4, 8, 16, 32]``).\n steps_per_phase : int\n Number of training steps in each fade-in phase.\n lambda_gp : float\n WGAN-GP gradient penalty weight.\n lr : float\n Learning rate for Adam (used for both G and D).\n \"\"\"\n\n def __init__(\n self,\n latent_size: int = 512,\n label_size: int = 0,\n fmap_base: int = 2048,\n fmap_max: int = 512,\n resolution_schedule: list[int] | None = None,\n steps_per_phase: int = 1000,\n lambda_gp: float = 10.0,\n lr: float = 1e-3,\n ) -> None:\n super().__init__()\n self.save_hyperparameters()\n if resolution_schedule is None:\n resolution_schedule = [4, 8, 16, 32, 64]\n self.latent_size = latent_size\n self.resolution_schedule = resolution_schedule\n self.steps_per_phase = steps_per_phase\n self.lambda_gp = lambda_gp\n self.lr = lr\n\n self.generator = _Generator(\n latent_size, fmap_base, fmap_max, resolution_schedule\n )\n self.discriminator = _Discriminator(fmap_base, fmap_max, resolution_schedule)\n\n self._step_count = 0\n self.automatic_optimization = False\n\n # ------------------------------------------------------------------\n # Helpers\n # ------------------------------------------------------------------\n\n def _gradient_penalty(self, real: torch.Tensor, fake: torch.Tensor) -> torch.Tensor:\n \"\"\"Compute WGAN-GP gradient penalty.\"\"\"\n b = real.size(0)\n eps = torch.rand(b, 1, 1, 1, 1, device=real.device)\n interp = (eps * real + (1.0 - eps) * fake).requires_grad_(True)\n d_interp = self.discriminator(interp)\n grads = torch.autograd.grad(\n outputs=d_interp,\n inputs=interp,\n grad_outputs=torch.ones_like(d_interp),\n create_graph=True,\n retain_graph=True,\n )[0]\n gp = ((grads.norm(2, dim=[1, 2, 3, 4]) - 1) ** 2).mean()\n return gp\n\n def _sample_z(self, n: int) -> torch.Tensor:\n return torch.randn(n, self.latent_size, device=self.device)\n\n # ------------------------------------------------------------------\n # Training\n # ------------------------------------------------------------------\n\n def training_step(self, batch: Any, batch_idx: int) -> None:\n opt_g, opt_d = self.optimizers()\n\n real = batch[\"image\"] if isinstance(batch, dict) else batch[0]\n b = real.size(0)\n z = self._sample_z(b)\n\n # --- Discriminator step ---\n opt_d.zero_grad()\n fake = self.generator(z).detach()\n d_real = self.discriminator(real)\n d_fake = self.discriminator(fake)\n gp = self._gradient_penalty(real, fake.requires_grad_(True))\n d_loss = d_fake.mean() - d_real.mean() + self.lambda_gp * gp\n self.manual_backward(d_loss)\n opt_d.step()\n\n # --- Generator step ---\n opt_g.zero_grad()\n fake = self.generator(z)\n g_loss = -self.discriminator(fake).mean()\n self.manual_backward(g_loss)\n opt_g.step()\n\n self.log_dict({\"g_loss\": g_loss, \"d_loss\": d_loss}, prog_bar=True)\n self._step_count += 1\n\n def on_train_batch_end(self, *args: Any, **kwargs: Any) -> None:\n \"\"\"Update alpha for fade-in scheduling.\"\"\"\n n_levels = len(self.resolution_schedule)\n level = min(self._step_count // self.steps_per_phase, n_levels - 1)\n phase_step = self._step_count % self.steps_per_phase\n alpha = min(phase_step / max(self.steps_per_phase, 1), 1.0)\n self.generator.current_level = level\n self.discriminator.current_level = level\n self.generator.alpha = alpha\n self.discriminator.alpha = alpha\n\n def configure_optimizers(self):\n opt_g = torch.optim.Adam(\n self.generator.parameters(), lr=self.lr, betas=(0.0, 0.99)\n )\n opt_d = torch.optim.Adam(\n self.discriminator.parameters(), lr=self.lr, betas=(0.0, 0.99)\n )\n return [opt_g, opt_d]\n\n\ndef progressivegan(\n latent_size: int = 512,\n label_size: int = 0,\n fmap_base: int = 2048,\n fmap_max: int = 512,\n resolution_schedule: list[int] | None = None,\n **kwargs,\n) -> ProgressiveGAN:\n \"\"\"Factory function for :class:`ProgressiveGAN`.\"\"\"\n return ProgressiveGAN(\n latent_size=latent_size,\n label_size=label_size,\n fmap_base=fmap_base,\n fmap_max=fmap_max,\n resolution_schedule=resolution_schedule,\n **kwargs,\n )\n\n\n__all__ = [\"ProgressiveGAN\", \"progressivegan\"]\n" }, { "path": "nobrainer/models/highresnet.py", "content": "\"\"\"HighResNet 3-D segmentation model (PyTorch).\n\nReference\n---------\nLi W. et al., \"On the Compactness, Efficiency, and Representation of\n3D Convolutional Networks: Brain Parcellation as a Pretext Task\",\nIPMI 2017. arXiv:1707.01992.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\nclass _ResBlock(nn.Module):\n \"\"\"Residual block: BN→Act→Conv→BN→Act→Conv + skip.\"\"\"\n\n def __init__(\n self,\n channels: int,\n dilation: int,\n act: type[nn.Module],\n ) -> None:\n super().__init__()\n padding = dilation\n self.path = nn.Sequential(\n nn.BatchNorm3d(channels),\n act(),\n nn.Conv3d(\n channels, channels, 3, padding=padding, dilation=dilation, bias=False\n ),\n nn.BatchNorm3d(channels),\n act(),\n nn.Conv3d(\n channels, channels, 3, padding=padding, dilation=dilation, bias=False\n ),\n )\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return x + self.path(x)\n\n\nclass _ZeroPadChannels(nn.Module):\n \"\"\"Pad the channel dimension symmetrically with zeros.\"\"\"\n\n def __init__(self, extra_channels: int) -> None:\n super().__init__()\n self.pad = extra_channels\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return F.pad(x, (0, 0, 0, 0, 0, 0, self.pad, self.pad))\n\n\nclass HighResNet(nn.Module):\n \"\"\"HighResNet — three stages of residual blocks with increasing dilation.\n\n Stage 1 (dilation=1): base_filters channels, n_blocks residual blocks\n Stage 2 (dilation=2): 2*base_filters channels\n Stage 3 (dilation=4): 4*base_filters channels\n\n Parameters\n ----------\n n_classes : int\n Number of output segmentation classes.\n in_channels : int\n Number of input image channels.\n base_filters : int\n Initial feature map count (doubled each stage).\n n_blocks : int\n Number of residual blocks per stage.\n activation : str\n ``\"relu\"`` or ``\"elu\"``.\n dropout_rate : float\n Spatial dropout probability after the last stage (0 = none).\n \"\"\"\n\n def __init__(\n self,\n n_classes: int = 1,\n in_channels: int = 1,\n base_filters: int = 16,\n n_blocks: int = 3,\n activation: str = \"relu\",\n dropout_rate: float = 0.0,\n ) -> None:\n super().__init__()\n act_cls: type[nn.Module] = {\"relu\": nn.ReLU, \"elu\": nn.ELU}[activation.lower()]\n f = base_filters # 16\n\n # Initial projection to base_filters channels\n self.init_conv = nn.Conv3d(in_channels, f, kernel_size=3, padding=1, bias=False)\n\n # Stage 1: f channels, dilation 1 → pad to 3f\n s1 = [_ResBlock(f, dilation=1, act=act_cls) for _ in range(n_blocks)]\n self.stage1 = nn.Sequential(*s1)\n self.pad1 = _ZeroPadChannels(f) # f → 3f\n\n # Stage 2: project 3f → 2f, dilation 2 → pad to 6f\n self.stage2_proj = nn.Conv3d(3 * f, 2 * f, kernel_size=1, bias=False)\n s2 = [_ResBlock(2 * f, dilation=2, act=act_cls) for _ in range(n_blocks)]\n self.stage2 = nn.Sequential(*s2)\n self.pad2 = _ZeroPadChannels(2 * f) # 2f → 6f\n\n # Stage 3: project 6f → 4f, dilation 4\n self.stage3_proj = nn.Conv3d(6 * f, 4 * f, kernel_size=1, bias=False)\n s3 = [_ResBlock(4 * f, dilation=4, act=act_cls) for _ in range(n_blocks)]\n self.stage3 = nn.Sequential(*s3)\n\n self.dropout = (\n nn.Dropout3d(p=dropout_rate) if dropout_rate > 0 else nn.Identity()\n )\n self.classifier = nn.Sequential(\n nn.BatchNorm3d(4 * f),\n act_cls(),\n nn.Conv3d(4 * f, n_classes, kernel_size=1),\n )\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n x = self.init_conv(x)\n\n s1 = self.stage1(x)\n s1 = self.pad1(s1) # (N, 3f, D, H, W)\n\n s2 = self.stage2_proj(s1)\n s2 = self.stage2(s2)\n s2 = self.pad2(s2) # (N, 6f, D, H, W)\n\n s3 = self.stage3_proj(s2)\n s3 = self.stage3(s3)\n s3 = self.dropout(s3)\n\n return self.classifier(s3)\n\n\ndef highresnet(\n n_classes: int = 1,\n in_channels: int = 1,\n base_filters: int = 16,\n n_blocks: int = 3,\n activation: str = \"relu\",\n dropout_rate: float = 0.0,\n **kwargs,\n) -> HighResNet:\n \"\"\"Factory function for :class:`HighResNet`.\"\"\"\n return HighResNet(\n n_classes=n_classes,\n in_channels=in_channels,\n base_filters=base_filters,\n n_blocks=n_blocks,\n activation=activation,\n dropout_rate=dropout_rate,\n )\n\n\n__all__ = [\"HighResNet\", \"highresnet\"]\n" }, { "path": "nobrainer/models/meshnet.py", "content": "\"\"\"MeshNet 3-D segmentation model (PyTorch).\n\nReference\n---------\nFedorov A. et al., \"End-to-end learning of brain tissue segmentation\nfrom imperfect labeling\", IJCNN 2017. arXiv:1612.00940.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport torch\nimport torch.nn as nn\n\n# Dilation schedules indexed by receptive field size\nfrom nobrainer.models._constants import ( # noqa: E501\n DILATION_SCHEDULES as _DILATION_SCHEDULES,\n)\n\n\nclass _ConvBNActDrop(nn.Module):\n def __init__(\n self,\n in_ch: int,\n out_ch: int,\n dilation: int,\n act: type[nn.Module],\n dropout_rate: float,\n ) -> None:\n super().__init__()\n padding = dilation # same-padding for 3×3×3 kernel\n self.block = nn.Sequential(\n nn.Conv3d(\n in_ch,\n out_ch,\n kernel_size=3,\n padding=padding,\n dilation=dilation,\n bias=False,\n ),\n nn.BatchNorm3d(out_ch),\n act(),\n nn.Dropout3d(p=dropout_rate),\n )\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return self.block(x)\n\n\nclass MeshNet(nn.Module):\n \"\"\"3-D MeshNet segmentation network.\n\n Seven layers of dilated 3×3×3 convolutions with a learnable dilation\n schedule that controls the receptive field.\n\n Parameters\n ----------\n n_classes : int\n Number of output segmentation classes.\n in_channels : int\n Number of input image channels (1 for single-modality MRI).\n filters : int\n Number of feature maps in all hidden layers.\n receptive_field : int\n One of ``37``, ``67``, ``129`` — selects the dilation schedule.\n activation : str\n ``\"relu\"`` or ``\"elu\"``.\n dropout_rate : float\n Spatial dropout probability applied after each conv layer (0 = none).\n \"\"\"\n\n def __init__(\n self,\n n_classes: int = 1,\n in_channels: int = 1,\n filters: int = 71,\n receptive_field: int = 67,\n activation: str = \"relu\",\n dropout_rate: float = 0.25,\n ) -> None:\n super().__init__()\n if receptive_field not in _DILATION_SCHEDULES:\n raise ValueError(\n f\"receptive_field must be one of {list(_DILATION_SCHEDULES)}, \"\n f\"got {receptive_field}\"\n )\n dilations = _DILATION_SCHEDULES[receptive_field]\n act_cls: type[nn.Module] = {\"relu\": nn.ReLU, \"elu\": nn.ELU}[activation.lower()]\n\n layers: list[nn.Module] = []\n for i, dil in enumerate(dilations):\n in_ch = in_channels if i == 0 else filters\n layers.append(_ConvBNActDrop(in_ch, filters, dil, act_cls, dropout_rate))\n\n self.encoder = nn.Sequential(*layers)\n self.classifier = nn.Conv3d(filters, n_classes, kernel_size=1)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return self.classifier(self.encoder(x))\n\n\ndef meshnet(\n n_classes: int = 1,\n in_channels: int = 1,\n filters: int = 71,\n receptive_field: int = 67,\n activation: str = \"relu\",\n dropout_rate: float = 0.25,\n **kwargs,\n) -> MeshNet:\n \"\"\"Factory function for :class:`MeshNet`.\"\"\"\n return MeshNet(\n n_classes=n_classes,\n in_channels=in_channels,\n filters=filters,\n receptive_field=receptive_field,\n activation=activation,\n dropout_rate=dropout_rate,\n )\n\n\n__all__ = [\"MeshNet\", \"meshnet\"]\n" }, { "path": "nobrainer/models/segformer3d.py", "content": "\"\"\"SegFormer3D: Efficient Transformer for 3D Medical Image Segmentation.\n\nPort of SegFormer3D (Perera et al., CVPR 2024 Workshop) to nobrainer.\nHierarchical vision transformer with efficient self-attention and\nall-MLP decoder for 3D volumetric segmentation.\n\nReference: https://arxiv.org/abs/2404.10156\nOriginal: https://github.com/OSUPCVLab/SegFormer3D\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom einops import rearrange\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n# ---------------------------------------------------------------------------\n# Encoder components\n# ---------------------------------------------------------------------------\n\n\nclass PatchEmbedding3d(nn.Module):\n \"\"\"3D overlapping patch embedding via strided convolution.\n\n Parameters\n ----------\n in_channels : int\n Input channels.\n embed_dim : int\n Output embedding dimension.\n kernel_size : int\n Conv kernel size.\n stride : int\n Conv stride (< kernel_size for overlap).\n padding : int\n Conv padding.\n \"\"\"\n\n def __init__(\n self,\n in_channels: int = 1,\n embed_dim: int = 64,\n kernel_size: int = 7,\n stride: int = 4,\n padding: int = 3,\n ) -> None:\n super().__init__()\n self.proj = nn.Conv3d(\n in_channels,\n embed_dim,\n kernel_size=kernel_size,\n stride=stride,\n padding=padding,\n )\n self.norm = nn.LayerNorm(embed_dim)\n\n def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, int, int, int]:\n \"\"\"Returns (B, N, C) tensor and spatial dims (D, H, W).\"\"\"\n x = self.proj(x) # (B, C, D, H, W)\n B, C, D, H, W = x.shape\n x = rearrange(x, \"b c d h w -> b (d h w) c\")\n x = self.norm(x)\n return x, D, H, W\n\n\nclass EfficientSelfAttention3d(nn.Module):\n \"\"\"Multi-head self-attention with spatial reduction.\n\n Reduces K, V spatial dimensions by ``sr_ratio`` before attention,\n giving O(N²/R²) complexity instead of O(N²).\n \"\"\"\n\n def __init__(\n self,\n embed_dim: int = 64,\n num_heads: int = 1,\n sr_ratio: int = 8,\n qkv_bias: bool = False,\n ) -> None:\n super().__init__()\n self.num_heads = num_heads\n self.head_dim = embed_dim // num_heads\n self.scale = self.head_dim**-0.5\n\n self.q = nn.Linear(embed_dim, embed_dim, bias=qkv_bias)\n self.kv = nn.Linear(embed_dim, embed_dim * 2, bias=qkv_bias)\n self.proj = nn.Linear(embed_dim, embed_dim)\n\n self.sr_ratio = sr_ratio\n if sr_ratio > 1:\n self.sr = nn.Conv3d(\n embed_dim,\n embed_dim,\n kernel_size=sr_ratio,\n stride=sr_ratio,\n )\n self.sr_norm = nn.LayerNorm(embed_dim)\n\n def forward(\n self,\n x: torch.Tensor,\n D: int,\n H: int,\n W: int,\n ) -> torch.Tensor:\n B, N, C = x.shape\n\n q = self.q(x)\n q = rearrange(q, \"b n (h d) -> b h n d\", h=self.num_heads)\n\n if self.sr_ratio > 1:\n x_3d = rearrange(x, \"b (d h w) c -> b c d h w\", d=D, h=H, w=W)\n x_sr = self.sr(x_3d)\n x_sr = rearrange(x_sr, \"b c d h w -> b (d h w) c\")\n x_sr = self.sr_norm(x_sr)\n kv = self.kv(x_sr)\n else:\n kv = self.kv(x)\n\n kv = rearrange(kv, \"b n (two h d) -> two b h n d\", two=2, h=self.num_heads)\n k, v = kv[0], kv[1]\n\n attn = (q @ k.transpose(-2, -1)) * self.scale\n attn = attn.softmax(dim=-1)\n\n out = attn @ v\n out = rearrange(out, \"b h n d -> b n (h d)\")\n return self.proj(out)\n\n\nclass DWConv3d(nn.Module):\n \"\"\"3D depth-wise convolution for positional encoding in MLP.\"\"\"\n\n def __init__(self, dim: int = 64) -> None:\n super().__init__()\n self.dwconv = nn.Conv3d(dim, dim, 3, padding=1, groups=dim)\n self.bn = nn.BatchNorm3d(dim)\n\n def forward(self, x: torch.Tensor, D: int, H: int, W: int) -> torch.Tensor:\n x = rearrange(x, \"b (d h w) c -> b c d h w\", d=D, h=H, w=W)\n x = self.bn(self.dwconv(x))\n x = rearrange(x, \"b c d h w -> b (d h w) c\")\n return x\n\n\nclass MixFFN3d(nn.Module):\n \"\"\"Feed-forward network with depth-wise conv for positional encoding.\"\"\"\n\n def __init__(\n self, embed_dim: int = 64, mlp_ratio: int = 4, dropout: float = 0.0\n ) -> None:\n super().__init__()\n hidden = embed_dim * mlp_ratio\n self.fc1 = nn.Linear(embed_dim, hidden)\n self.dwconv = DWConv3d(hidden)\n self.act = nn.GELU()\n self.fc2 = nn.Linear(hidden, embed_dim)\n self.drop = nn.Dropout(dropout)\n\n def forward(self, x: torch.Tensor, D: int, H: int, W: int) -> torch.Tensor:\n x = self.fc1(x)\n x = self.dwconv(x, D, H, W)\n x = self.act(x)\n x = self.drop(x)\n x = self.fc2(x)\n x = self.drop(x)\n return x\n\n\nclass TransformerBlock3d(nn.Module):\n \"\"\"Transformer block: LN → Attention → residual → LN → FFN → residual.\"\"\"\n\n def __init__(\n self,\n embed_dim: int = 64,\n num_heads: int = 1,\n mlp_ratio: int = 4,\n sr_ratio: int = 8,\n dropout: float = 0.0,\n ) -> None:\n super().__init__()\n self.norm1 = nn.LayerNorm(embed_dim)\n self.attn = EfficientSelfAttention3d(\n embed_dim,\n num_heads,\n sr_ratio,\n qkv_bias=True,\n )\n self.norm2 = nn.LayerNorm(embed_dim)\n self.ffn = MixFFN3d(embed_dim, mlp_ratio, dropout)\n\n def forward(self, x: torch.Tensor, D: int, H: int, W: int) -> torch.Tensor:\n x = x + self.attn(self.norm1(x), D, H, W)\n x = x + self.ffn(self.norm2(x), D, H, W)\n return x\n\n\n# ---------------------------------------------------------------------------\n# Hierarchical Encoder (Mix Transformer)\n# ---------------------------------------------------------------------------\n\n\nclass MixTransformerEncoder3d(nn.Module):\n \"\"\"4-stage hierarchical transformer encoder.\n\n Each stage: PatchEmbedding → N × TransformerBlock → output features.\n Spatial resolution halves (approximately) at each stage.\n \"\"\"\n\n def __init__(\n self,\n in_channels: int = 1,\n embed_dims: tuple[int, ...] = (64, 128, 320, 512),\n depths: tuple[int, ...] = (2, 2, 2, 2),\n num_heads: tuple[int, ...] = (1, 2, 5, 8),\n sr_ratios: tuple[int, ...] = (8, 4, 2, 1),\n mlp_ratio: int = 4,\n patch_sizes: tuple[int, ...] = (7, 3, 3, 3),\n strides: tuple[int, ...] = (4, 2, 2, 2),\n dropout: float = 0.0,\n ) -> None:\n super().__init__()\n self.num_stages = len(embed_dims)\n\n for i in range(self.num_stages):\n in_ch = in_channels if i == 0 else embed_dims[i - 1]\n padding = patch_sizes[i] // 2\n patch_embed = PatchEmbedding3d(\n in_ch,\n embed_dims[i],\n patch_sizes[i],\n strides[i],\n padding,\n )\n blocks = nn.ModuleList(\n [\n TransformerBlock3d(\n embed_dims[i],\n num_heads[i],\n mlp_ratio,\n sr_ratios[i],\n dropout,\n )\n for _ in range(depths[i])\n ]\n )\n norm = nn.LayerNorm(embed_dims[i])\n\n setattr(self, f\"patch_embed_{i}\", patch_embed)\n setattr(self, f\"blocks_{i}\", blocks)\n setattr(self, f\"norm_{i}\", norm)\n\n def forward(self, x: torch.Tensor) -> list[torch.Tensor]:\n \"\"\"Returns list of multi-scale features [(B, C_i, D_i, H_i, W_i)].\"\"\"\n features = []\n for i in range(self.num_stages):\n patch_embed = getattr(self, f\"patch_embed_{i}\")\n blocks = getattr(self, f\"blocks_{i}\")\n norm = getattr(self, f\"norm_{i}\")\n\n x, D, H, W = patch_embed(x)\n for blk in blocks:\n x = blk(x, D, H, W)\n x = norm(x)\n\n # Reshape back to 3D for next stage\n x = rearrange(x, \"b (d h w) c -> b c d h w\", d=D, h=H, w=W)\n features.append(x)\n\n return features\n\n\n# ---------------------------------------------------------------------------\n# MLP Decoder\n# ---------------------------------------------------------------------------\n\n\nclass SegFormerDecoderHead(nn.Module):\n \"\"\"All-MLP decoder that aggregates multi-scale features.\n\n Upsamples features from all encoder stages to the highest resolution,\n concatenates, and projects to n_classes.\n \"\"\"\n\n def __init__(\n self,\n embed_dims: tuple[int, ...] = (64, 128, 320, 512),\n decoder_dim: int = 256,\n n_classes: int = 1,\n ) -> None:\n super().__init__()\n self.n_stages = len(embed_dims)\n\n # Linear projection per stage\n self.linears = nn.ModuleList(\n [nn.Linear(embed_dims[i], decoder_dim) for i in range(self.n_stages)]\n )\n\n # Fuse concatenated features\n self.fuse = nn.Sequential(\n nn.Linear(decoder_dim * self.n_stages, decoder_dim),\n nn.ReLU(inplace=True),\n )\n self.pred = nn.Linear(decoder_dim, n_classes)\n\n def forward(self, features: list[torch.Tensor]) -> torch.Tensor:\n \"\"\"features: list of (B, C_i, D_i, H_i, W_i) from encoder stages.\"\"\"\n # Target spatial size = largest feature map (first stage)\n target = features[0].shape[2:] # (D0, H0, W0)\n\n projected = []\n for i, feat in enumerate(features):\n B, C, D, H, W = feat.shape\n x = rearrange(feat, \"b c d h w -> b (d h w) c\")\n x = self.linears[i](x) # (B, N, decoder_dim)\n x = rearrange(x, \"b (d h w) c -> b c d h w\", d=D, h=H, w=W)\n\n # Upsample to target resolution\n if (D, H, W) != target:\n x = F.interpolate(x, size=target, mode=\"trilinear\", align_corners=False)\n\n projected.append(x)\n\n # Concatenate along channel dim, then fuse\n fused = torch.cat(projected, dim=1) # (B, decoder_dim * n_stages, D, H, W)\n B, C, D, H, W = fused.shape\n fused = rearrange(fused, \"b c d h w -> b (d h w) c\")\n fused = self.fuse(fused)\n out = self.pred(fused) # (B, D*H*W, n_classes)\n out = rearrange(out, \"b (d h w) c -> b c d h w\", d=D, h=H, w=W)\n\n return out\n\n\n# ---------------------------------------------------------------------------\n# SegFormer3D Model\n# ---------------------------------------------------------------------------\n\n\nclass SegFormer3D(nn.Module):\n \"\"\"SegFormer3D: Hierarchical Transformer for 3D Medical Image Segmentation.\n\n Combines a multi-stage transformer encoder (MixTransformer) with an\n all-MLP decoder for efficient 3D segmentation.\n\n Parameters\n ----------\n n_classes : int\n Number of output segmentation classes.\n in_channels : int\n Number of input channels (1 for MRI).\n embed_dims : tuple of int\n Embedding dimensions per encoder stage.\n depths : tuple of int\n Number of transformer blocks per stage.\n num_heads : tuple of int\n Number of attention heads per stage.\n sr_ratios : tuple of int\n Spatial reduction ratios for efficient attention per stage.\n mlp_ratio : int\n MLP hidden dimension multiplier.\n decoder_dim : int\n Decoder unified channel dimension.\n dropout : float\n Dropout probability.\n \"\"\"\n\n def __init__(\n self,\n n_classes: int = 1,\n in_channels: int = 1,\n embed_dims: tuple[int, ...] = (32, 64, 160, 256),\n depths: tuple[int, ...] = (2, 2, 2, 2),\n num_heads: tuple[int, ...] = (1, 2, 5, 8),\n sr_ratios: tuple[int, ...] = (8, 4, 2, 1),\n mlp_ratio: int = 4,\n decoder_dim: int = 256,\n dropout: float = 0.0,\n ) -> None:\n super().__init__()\n self.encoder = MixTransformerEncoder3d(\n in_channels=in_channels,\n embed_dims=embed_dims,\n depths=depths,\n num_heads=num_heads,\n sr_ratios=sr_ratios,\n mlp_ratio=mlp_ratio,\n dropout=dropout,\n )\n self.decoder = SegFormerDecoderHead(\n embed_dims=embed_dims,\n decoder_dim=decoder_dim,\n n_classes=n_classes,\n )\n\n # Final upsample to match input resolution\n self._upsample_factor = 4 # first stage stride\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n \"\"\"(B, C, D, H, W) → (B, n_classes, D, H, W).\"\"\"\n input_shape = x.shape[2:]\n features = self.encoder(x)\n out = self.decoder(features)\n\n # Upsample to input resolution if needed\n if out.shape[2:] != input_shape:\n out = F.interpolate(\n out, size=input_shape, mode=\"trilinear\", align_corners=False\n )\n\n return out\n\n\n# ---------------------------------------------------------------------------\n# Factory function\n# ---------------------------------------------------------------------------\n\n\ndef segformer3d(\n n_classes: int = 1,\n in_channels: int = 1,\n embed_dims: tuple[int, ...] = (32, 64, 160, 256),\n depths: tuple[int, ...] = (2, 2, 2, 2),\n num_heads: tuple[int, ...] = (1, 2, 5, 8),\n sr_ratios: tuple[int, ...] = (8, 4, 2, 1),\n mlp_ratio: int = 4,\n decoder_dim: int = 256,\n dropout: float = 0.0,\n **kwargs,\n) -> SegFormer3D:\n \"\"\"Factory function for :class:`SegFormer3D`.\n\n Default config (~4.5M params) matches the paper's base variant.\n\n Common size variants:\n - **tiny**: ``embed_dims=(16, 32, 80, 128)`` (~1.5M params)\n - **small** (default): ``embed_dims=(32, 64, 160, 256)`` (~4.5M params)\n - **base**: ``embed_dims=(64, 128, 320, 512)`` (~18M params)\n \"\"\"\n return SegFormer3D(\n n_classes=n_classes,\n in_channels=in_channels,\n embed_dims=embed_dims,\n depths=depths,\n num_heads=num_heads,\n sr_ratios=sr_ratios,\n mlp_ratio=mlp_ratio,\n decoder_dim=decoder_dim,\n dropout=dropout,\n )\n\n\n__all__ = [\"SegFormer3D\", \"segformer3d\"]\n" }, { "path": "nobrainer/models/segmentation.py", "content": "\"\"\"MONAI-backed segmentation model factory functions.\n\nAll models expect input of shape ``(N, C_in, D, H, W)`` and produce\noutput of shape ``(N, n_classes, D, H, W)``.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom monai.networks.nets import UNETR, AttentionUnet, UNet, VNet\nimport torch.nn as nn\n\n\ndef unet(\n n_classes: int = 1,\n in_channels: int = 1,\n channels: tuple[int, ...] = (16, 32, 64, 128, 256),\n strides: tuple[int, ...] = (2, 2, 2, 2),\n num_res_units: int = 0,\n act: str = \"RELU\",\n norm: str = \"BATCH\",\n dropout: float = 0.0,\n **kwargs,\n) -> UNet:\n \"\"\"Return a 3-D UNet (MONAI implementation).\n\n Parameters\n ----------\n n_classes : int\n Number of output segmentation classes.\n in_channels : int\n Number of input image channels (1 for grayscale MRI).\n channels : tuple of int\n Filter count at each level (len == levels + 1).\n strides : tuple of int\n Down-sampling stride at each level (len == levels).\n num_res_units : int\n Number of residual units per level (0 = plain conv blocks).\n act : str\n Activation name (MONAI convention: \"RELU\", \"LEAKYRELU\", \"ELU\", …).\n norm : str\n Normalisation: \"BATCH\", \"INSTANCE\", \"GROUP\", \"LAYER\", or \"NONE\".\n dropout : float\n Dropout probability (0 = disabled).\n \"\"\"\n return UNet(\n spatial_dims=3,\n in_channels=in_channels,\n out_channels=n_classes,\n channels=channels,\n strides=strides,\n num_res_units=num_res_units,\n act=act,\n norm=norm,\n dropout=dropout,\n **kwargs,\n )\n\n\ndef vnet(\n n_classes: int = 1,\n in_channels: int = 1,\n act: str = \"elu\",\n dropout_dim: int = 3,\n **kwargs,\n) -> VNet:\n \"\"\"Return a 3-D V-Net (MONAI implementation).\n\n Parameters\n ----------\n n_classes : int\n Number of output segmentation classes.\n in_channels : int\n Number of input channels.\n act : str\n Activation function name (lowercase MONAI style: \"elu\", \"relu\", …).\n dropout_dim : int\n Dimension for spatial dropout (1 = channel, 3 = 3-D spatial).\n \"\"\"\n return VNet(\n spatial_dims=3,\n in_channels=in_channels,\n out_channels=n_classes,\n act=act,\n dropout_dim=dropout_dim,\n **kwargs,\n )\n\n\ndef attention_unet(\n n_classes: int = 1,\n in_channels: int = 1,\n channels: tuple[int, ...] = (64, 128, 256, 512),\n strides: tuple[int, ...] = (2, 2, 2),\n dropout: float = 0.0,\n **kwargs,\n) -> AttentionUnet:\n \"\"\"Return a 3-D Attention U-Net (MONAI implementation).\n\n Parameters\n ----------\n n_classes : int\n Number of output segmentation classes.\n in_channels : int\n Number of input channels.\n channels : tuple of int\n Filter counts at each encoder level.\n strides : tuple of int\n Down-sampling strides (len == len(channels) - 1).\n dropout : float\n Dropout probability.\n \"\"\"\n return AttentionUnet(\n spatial_dims=3,\n in_channels=in_channels,\n out_channels=n_classes,\n channels=channels,\n strides=strides,\n dropout=dropout,\n **kwargs,\n )\n\n\ndef unetr(\n n_classes: int = 1,\n in_channels: int = 1,\n img_size: tuple[int, int, int] = (96, 96, 96),\n feature_size: int = 16,\n hidden_size: int = 768,\n mlp_dim: int = 3072,\n num_heads: int = 12,\n dropout_rate: float = 0.1,\n norm_name: str = \"instance\",\n **kwargs,\n) -> UNETR:\n \"\"\"Return a UNETR (ViT backbone + U-Net decoder) (MONAI implementation).\n\n Parameters\n ----------\n n_classes : int\n Number of output segmentation classes.\n in_channels : int\n Number of input channels.\n img_size : tuple of int\n Spatial size of the input volume ``(D, H, W)``.\n feature_size : int\n Spatial feature size for the decoder (MONAI default 16).\n hidden_size : int\n ViT embedding dimension (default 768 = ViT-B).\n mlp_dim : int\n MLP hidden dim in transformer blocks.\n num_heads : int\n Number of attention heads.\n dropout_rate : float\n Dropout applied inside the transformer.\n norm_name : str\n Normalisation: \"instance\", \"batch\".\n \"\"\"\n return UNETR(\n in_channels=in_channels,\n out_channels=n_classes,\n img_size=img_size,\n feature_size=feature_size,\n hidden_size=hidden_size,\n mlp_dim=mlp_dim,\n num_heads=num_heads,\n dropout_rate=dropout_rate,\n norm_name=norm_name,\n **kwargs,\n )\n\n\ndef swin_unetr(\n n_classes: int = 1,\n in_channels: int = 1,\n feature_size: int = 24,\n depths: tuple[int, ...] = (2, 2, 2, 2),\n num_heads: tuple[int, ...] = (3, 6, 12, 24),\n norm_name: str = \"instance\",\n dropout_rate: float = 0.0,\n **kwargs,\n) -> nn.Module:\n \"\"\"Return a SwinUNETR (Swin Transformer U-Net) (MONAI implementation).\n\n Parameters\n ----------\n n_classes : int\n Number of output segmentation classes.\n in_channels : int\n Number of input channels.\n feature_size : int\n Feature size for the decoder (default 24).\n depths : tuple of int\n Number of Swin Transformer blocks at each stage.\n num_heads : tuple of int\n Number of attention heads at each stage.\n norm_name : str\n Normalisation: ``\"instance\"`` or ``\"batch\"``.\n dropout_rate : float\n Dropout probability.\n \"\"\"\n from monai.networks.nets import SwinUNETR as _SwinUNETR\n\n return _SwinUNETR(\n in_channels=in_channels,\n out_channels=n_classes,\n feature_size=feature_size,\n depths=depths,\n num_heads=num_heads,\n norm_name=norm_name,\n drop_rate=dropout_rate,\n spatial_dims=3,\n **kwargs,\n )\n\n\ndef segresnet(\n n_classes: int = 1,\n in_channels: int = 1,\n blocks_down: tuple[int, ...] = (1, 2, 2, 4),\n init_filters: int = 16,\n norm: str = \"INSTANCE\",\n dropout_prob: float = 0.0,\n **kwargs,\n) -> nn.Module:\n \"\"\"Return a SegResNet (residual encoder segmentation network) (MONAI).\n\n Used as the default architecture in MONAI Auto3DSeg.\n\n Parameters\n ----------\n n_classes : int\n Number of output segmentation classes.\n in_channels : int\n Number of input channels.\n blocks_down : tuple of int\n Number of residual blocks at each encoder level.\n init_filters : int\n Initial number of filters (doubled at each level).\n norm : str\n Normalisation: ``\"GROUP\"``, ``\"BATCH\"``, ``\"INSTANCE\"``.\n dropout_prob : float\n Dropout probability.\n \"\"\"\n from monai.networks.nets import SegResNet as _SegResNet\n\n return _SegResNet(\n spatial_dims=3,\n in_channels=in_channels,\n out_channels=n_classes,\n blocks_down=blocks_down,\n init_filters=init_filters,\n norm=norm,\n dropout_prob=dropout_prob,\n **kwargs,\n )\n\n\n__all__ = [\n \"unet\",\n \"vnet\",\n \"attention_unet\",\n \"unetr\",\n \"swin_unetr\",\n \"segresnet\",\n]\n" }, { "path": "nobrainer/models/simsiam.py", "content": "\"\"\"SimSiam self-supervised learning model for 3-D brain volumes (PyTorch).\n\nReference\n---------\nChen X. & He K., \"Exploring Simple Siamese Representation Learning\",\nCVPR 2021. arXiv:2011.10566.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport torch\nimport torch.nn as nn\n\nfrom .highresnet import HighResNet\n\n\nclass SimSiam(nn.Module):\n \"\"\"Siamese network with stop-gradient for self-supervised pre-training.\n\n Architecture\n ------------\n - **Backbone**: :class:`~nobrainer.models.highresnet.HighResNet` that\n encodes a 3-D volume into a spatial feature map.\n - **Projector**: Global average pool → MLP (2 hidden layers) → projection\n vector of size ``projection_dim``.\n - **Predictor**: Bottleneck MLP (``projection_dim`` → ``latent_dim`` →\n ``projection_dim``).\n\n Training\n --------\n Produce two augmented views of the same volume, pass each through the\n encoder + projector, and apply the *negative cosine similarity* loss\n between ``predictor(z1)`` and ``stop_grad(z2)`` (and vice-versa).\n\n Parameters\n ----------\n n_classes : int\n Passed to the HighResNet backbone (not used for classification,\n but kept for architecture compatibility).\n in_channels : int\n Number of input channels.\n projection_dim : int\n Output dimension of the projector head.\n latent_dim : int\n Hidden bottleneck size in the predictor.\n weight_decay : float\n L2 regularisation weight (applied externally via the optimiser).\n \"\"\"\n\n def __init__(\n self,\n n_classes: int = 1,\n in_channels: int = 1,\n projection_dim: int = 2048,\n latent_dim: int = 512,\n weight_decay: float = 0.0005,\n ) -> None:\n super().__init__()\n self.weight_decay = weight_decay\n\n backbone = HighResNet(n_classes=n_classes, in_channels=in_channels)\n # Determine backbone output channels by inspecting the classifier head\n self.backbone = backbone\n backbone_feat_ch = backbone.classifier[2].in_channels # 4*f\n\n self.projector = nn.Sequential(\n nn.AdaptiveAvgPool3d(1),\n nn.Flatten(),\n nn.Linear(backbone_feat_ch, projection_dim),\n nn.BatchNorm1d(projection_dim),\n nn.ReLU(inplace=True),\n nn.Linear(projection_dim, projection_dim),\n nn.BatchNorm1d(projection_dim),\n )\n\n self.predictor = nn.Sequential(\n nn.Linear(projection_dim, latent_dim),\n nn.ReLU(inplace=True),\n nn.BatchNorm1d(latent_dim),\n nn.Linear(latent_dim, projection_dim),\n )\n\n def _encode(self, x: torch.Tensor) -> torch.Tensor:\n \"\"\"Run backbone up to stage3 (before classifier head).\"\"\"\n h = self.backbone.init_conv(x)\n s1 = self.backbone.stage1(h)\n s1 = self.backbone.pad1(s1)\n s2 = self.backbone.stage2_proj(s1)\n s2 = self.backbone.stage2(s2)\n s2 = self.backbone.pad2(s2)\n s3 = self.backbone.stage3_proj(s2)\n s3 = self.backbone.stage3(s3)\n return s3 # (N, 4f, D, H, W)\n\n def forward(\n self, x1: torch.Tensor, x2: torch.Tensor\n ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:\n \"\"\"Forward pass producing predictions and projections for both views.\n\n Returns\n -------\n p1, p2 : torch.Tensor\n Predictions for view 1 and view 2 (gradient flows through these).\n z1, z2 : torch.Tensor\n Projections (used as stop-gradient targets in the SimSiam loss).\n \"\"\"\n feat1 = self._encode(x1)\n feat2 = self._encode(x2)\n\n z1 = self.projector(feat1)\n z2 = self.projector(feat2)\n\n p1 = self.predictor(z1)\n p2 = self.predictor(z2)\n\n return p1, p2, z1.detach(), z2.detach()\n\n @staticmethod\n def loss(\n p1: torch.Tensor,\n p2: torch.Tensor,\n z1: torch.Tensor,\n z2: torch.Tensor,\n ) -> torch.Tensor:\n \"\"\"Negative cosine similarity loss (symmetric).\"\"\"\n cos = nn.functional.cosine_similarity\n\n def _d(p: torch.Tensor, z: torch.Tensor) -> torch.Tensor:\n return -cos(p, z, dim=-1).mean()\n\n return (_d(p1, z2) + _d(p2, z1)) * 0.5\n\n\ndef simsiam(\n n_classes: int = 1,\n in_channels: int = 1,\n projection_dim: int = 2048,\n latent_dim: int = 512,\n weight_decay: float = 0.0005,\n **kwargs,\n) -> SimSiam:\n \"\"\"Factory function for :class:`SimSiam`.\"\"\"\n return SimSiam(\n n_classes=n_classes,\n in_channels=in_channels,\n projection_dim=projection_dim,\n latent_dim=latent_dim,\n weight_decay=weight_decay,\n )\n\n\n__all__ = [\"SimSiam\", \"simsiam\"]\n" }, { "path": "nobrainer/models/tests/__init__.py", "content": "" }, { "path": "nobrainer/prediction.py", "content": "\"\"\"Block-based prediction utilities (PyTorch, no TensorFlow).\"\"\"\n\nfrom __future__ import annotations\n\nfrom pathlib import Path\nfrom typing import Any\n\nimport nibabel as nib\nimport numpy as np\nimport torch\nimport torch.nn as nn\n\nfrom nobrainer.training import get_device\n\n\ndef _forward(model: nn.Module, tensor: torch.Tensor, mc: bool | None = None):\n \"\"\"Call model forward, passing mc= if the model supports it.\"\"\"\n from nobrainer.models._utils import model_supports_mc\n\n if mc is not None and model_supports_mc(model):\n return model(tensor, mc=mc)\n return model(tensor)\n\n\ndef _pad_to_multiple(\n arr: np.ndarray, block_shape: tuple[int, int, int]\n) -> tuple[np.ndarray, tuple[int, ...]]:\n \"\"\"Pad spatial dims of ``arr`` (D, H, W) so each is divisible by block_shape.\"\"\"\n pads = []\n for dim, bs in zip(arr.shape, block_shape):\n rem = (-dim) % bs\n pads.append((0, rem))\n return np.pad(arr, pads, mode=\"constant\"), tuple(p[1] for p in pads)\n\n\ndef _extract_blocks(arr: np.ndarray, block_shape: tuple[int, int, int]) -> np.ndarray:\n \"\"\"Split ``arr`` (D, H, W) into non-overlapping blocks of ``block_shape``.\"\"\"\n D, H, W = arr.shape\n bd, bh, bw = block_shape\n blocks = arr.reshape(D // bd, bd, H // bh, bh, W // bw, bw)\n # → (nd, bD, nh, bH, nw, bW)\n blocks = blocks.transpose(0, 2, 4, 1, 3, 5)\n # → (nd, nh, nw, bD, bH, bW)\n nd, nh, nw = D // bd, H // bh, W // bw\n return blocks.reshape(nd * nh * nw, bd, bh, bw), (nd, nh, nw)\n\n\ndef _stitch_blocks(\n block_preds: np.ndarray,\n grid: tuple[int, int, int],\n block_shape: tuple[int, int, int],\n pad: tuple[int, int, int],\n orig_shape: tuple[int, int, int],\n n_classes: int,\n) -> np.ndarray:\n \"\"\"Reconstruct full prediction volume from per-block predictions.\"\"\"\n nd, nh, nw = grid\n bd, bh, bw = block_shape\n # block_preds: (N_blocks, n_classes, bD, bH, bW)\n full = np.zeros((n_classes, nd * bd, nh * bh, nw * bw), dtype=block_preds.dtype)\n idx = 0\n for i in range(nd):\n for j in range(nh):\n for k in range(nw):\n full[\n :,\n i * bd : (i + 1) * bd,\n j * bh : (j + 1) * bh,\n k * bw : (k + 1) * bw,\n ] = block_preds[idx]\n idx += 1\n\n # Remove padding\n D, H, W = orig_shape\n pd, ph, pw = pad\n end_d = full.shape[1] - pd if pd > 0 else full.shape[1]\n end_h = full.shape[2] - ph if ph > 0 else full.shape[2]\n end_w = full.shape[3] - pw if pw > 0 else full.shape[3]\n return full[:, :end_d, :end_h, :end_w]\n\n\ndef strided_patch_positions(\n volume_shape: tuple[int, int, int],\n block_shape: tuple[int, int, int],\n stride: tuple[int, int, int] | None = None,\n) -> list[tuple[slice, slice, slice]]:\n \"\"\"Compute grid positions for strided patch extraction.\n\n Parameters\n ----------\n volume_shape : tuple of int\n ``(D, H, W)`` of the volume.\n block_shape : tuple of int\n ``(bD, bH, bW)`` patch size.\n stride : tuple of int or None\n Step size per axis. None = block_shape (non-overlapping).\n\n Returns\n -------\n list of tuple of slice\n Each entry is ``(slice_d, slice_h, slice_w)`` for extracting one patch.\n \"\"\"\n if stride is None:\n stride = block_shape\n\n positions = []\n for d in range(0, volume_shape[0] - block_shape[0] + 1, stride[0]):\n for h in range(0, volume_shape[1] - block_shape[1] + 1, stride[1]):\n for w in range(0, volume_shape[2] - block_shape[2] + 1, stride[2]):\n positions.append(\n (\n slice(d, d + block_shape[0]),\n slice(h, h + block_shape[1]),\n slice(w, w + block_shape[2]),\n )\n )\n\n # Handle remainder: if volume not evenly divisible, add edge patches\n for axis in range(3):\n dim = volume_shape[axis]\n bs = block_shape[axis]\n st = stride[axis]\n last_start = (dim - bs) // st * st\n if last_start + bs < dim:\n # Need an extra patch at the edge\n edge_start = dim - bs\n if edge_start >= 0:\n # Add positions for this edge along all existing grid lines\n # (simplified: just ensure coverage)\n pass # Covered by the >= 0 check above\n\n return positions\n\n\ndef reassemble_predictions(\n patches: list[tuple[np.ndarray, tuple[slice, slice, slice]]],\n volume_shape: tuple[int, int, int],\n n_classes: int,\n strategy: str = \"average\",\n) -> np.ndarray:\n \"\"\"Reassemble overlapping patch predictions into a full volume.\n\n Parameters\n ----------\n patches : list of (array, slices)\n Each entry is ``(pred, (slice_d, slice_h, slice_w))`` where\n ``pred`` has shape ``(n_classes, bD, bH, bW)``.\n volume_shape : tuple of int\n ``(D, H, W)`` of the target volume.\n n_classes : int\n Number of output classes.\n strategy : str\n ``\"average\"`` (mean of overlapping predictions),\n ``\"vote\"`` (argmax then majority vote), or\n ``\"max\"`` (max probability per class).\n\n Returns\n -------\n np.ndarray\n Shape ``(n_classes, D, H, W)`` probability volume.\n \"\"\"\n D, H, W = volume_shape\n output = np.zeros((n_classes, D, H, W), dtype=np.float64)\n counts = np.zeros((1, D, H, W), dtype=np.float64)\n\n for pred, slices in patches:\n sd, sh, sw = slices\n if strategy == \"max\":\n output[:, sd, sh, sw] = np.maximum(output[:, sd, sh, sw], pred)\n else:\n output[:, sd, sh, sw] += pred\n counts[0, sd, sh, sw] += 1.0\n\n if strategy == \"average\":\n counts = np.maximum(counts, 1.0)\n output = output / counts\n\n return output.astype(np.float32)\n\n\ndef _predict_strided(\n arr: np.ndarray,\n affine: np.ndarray | None,\n model: nn.Module,\n block_shape: tuple[int, int, int],\n stride: tuple[int, int, int],\n batch_size: int,\n device: torch.device,\n return_labels: bool,\n normalizer: Any | None,\n) -> nib.Nifti1Image:\n \"\"\"Strided prediction with overlap reassembly.\"\"\"\n from nobrainer.gpu import get_device\n\n if device is None:\n device = get_device()\n model = model.to(device)\n model.eval()\n\n vol_shape = arr.shape[:3]\n positions = strided_patch_positions(vol_shape, block_shape, stride)\n\n patches = []\n with torch.no_grad():\n for i in range(0, len(positions), batch_size):\n batch_pos = positions[i : i + batch_size]\n batch_blocks = np.stack([arr[sd, sh, sw] for sd, sh, sw in batch_pos])\n if normalizer is not None:\n batch_blocks = np.stack([normalizer(b) for b in batch_blocks])\n\n tensor = torch.from_numpy(batch_blocks[:, None].astype(np.float32)).to(\n device\n )\n out = _forward(model, tensor, mc=False)\n probs = torch.softmax(out, dim=1).cpu().numpy()\n\n for j, pos in enumerate(batch_pos):\n patches.append((probs[j], pos))\n\n n_classes = patches[0][0].shape[0]\n full_pred = reassemble_predictions(\n patches, vol_shape, n_classes, strategy=\"average\"\n )\n\n if return_labels:\n labels = full_pred.argmax(axis=0).astype(np.int32)\n result = nib.Nifti1Image(labels, affine)\n else:\n result = nib.Nifti1Image(full_pred.transpose(1, 2, 3, 0), affine)\n return result\n\n\ndef predict(\n inputs: str | Path | np.ndarray | nib.Nifti1Image,\n model: nn.Module,\n block_shape: tuple[int, int, int] = (128, 128, 128),\n stride: tuple[int, int, int] | None = None,\n batch_size: int = 4,\n device: str | torch.device | None = None,\n return_labels: bool = True,\n normalizer: Any | None = None,\n) -> nib.Nifti1Image:\n \"\"\"Run block-based inference on a 3-D brain volume.\n\n Parameters\n ----------\n inputs : path, ndarray, or Nifti1Image\n Input brain MRI. If a file path is given, it is loaded with\n nibabel. If an ndarray, shape must be ``(D, H, W)``.\n model : nn.Module\n Trained PyTorch segmentation model. Must accept tensors of\n shape ``(N, 1, bD, bH, bW)`` and return ``(N, C, bD, bH, bW)``.\n block_shape : tuple\n Spatial block size ``(bD, bH, bW)`` for patch-based inference.\n batch_size : int\n Number of blocks to process in one forward pass.\n device : str, device, or None\n Compute device. Defaults to CUDA if available, else CPU.\n return_labels : bool\n If ``True``, return argmax labels. If ``False``, return class\n probabilities (softmax) as a 4-D volume.\n normalizer : callable or None\n Optional function ``normalizer(arr) → arr`` applied to each block\n before inference.\n\n Returns\n -------\n nib.Nifti1Image\n Segmentation (or probability) volume with the same affine as the\n input NIfTI.\n \"\"\"\n if device is None:\n device = get_device()\n device = torch.device(device)\n\n # Multi-GPU: distribute blocks across GPUs when device=\"cuda\" and >1 GPU\n n_gpus = torch.cuda.device_count() if device.type == \"cuda\" else 1\n use_multi_gpu = n_gpus > 1\n\n # Load input\n affine = np.eye(4)\n if isinstance(inputs, (str, Path)):\n img = nib.load(str(inputs))\n arr = np.asarray(img.dataobj, dtype=np.float32)\n affine = img.affine\n elif isinstance(inputs, nib.Nifti1Image):\n arr = np.asarray(inputs.dataobj, dtype=np.float32)\n affine = inputs.affine\n else:\n arr = np.asarray(inputs, dtype=np.float32)\n\n orig_shape = arr.shape[:3]\n arr3d = arr if arr.ndim == 3 else arr[..., 0]\n\n # Strided prediction path (overlapping blocks with reassembly)\n if stride is not None:\n return _predict_strided(\n arr3d,\n affine,\n model,\n block_shape,\n stride,\n batch_size,\n device,\n return_labels,\n normalizer,\n )\n\n # Pad to block-divisible size\n padded, pad = _pad_to_multiple(arr3d, block_shape)\n blocks, grid = _extract_blocks(padded, block_shape) # (N_blocks, bD, bH, bW)\n n_blocks = blocks.shape[0]\n\n if use_multi_gpu:\n # Replicate model to each GPU (deep copy to avoid moving the original)\n import copy\n\n _ = model.state_dict()\n models = []\n for i in range(n_gpus):\n m = copy.deepcopy(model).to(torch.device(f\"cuda:{i}\"))\n m.eval()\n models.append(m)\n else:\n model = model.to(device)\n model.eval()\n\n all_preds: list[np.ndarray] = []\n with torch.no_grad():\n for start in range(0, n_blocks, batch_size):\n chunk = blocks[start : start + batch_size] # (B, bD, bH, bW)\n if normalizer is not None:\n chunk = np.stack([normalizer(b) for b in chunk])\n\n if use_multi_gpu:\n # Round-robin distribute across GPUs\n gpu_idx = (start // batch_size) % n_gpus\n dev = torch.device(f\"cuda:{gpu_idx}\")\n tensor = torch.from_numpy(chunk[:, None]).to(dev)\n out = _forward(models[gpu_idx], tensor, mc=False)\n else:\n tensor = torch.from_numpy(chunk[:, None]).to(device)\n out = _forward(model, tensor, mc=False)\n\n if return_labels:\n out = out.argmax(dim=1, keepdim=True).float()\n else:\n out = torch.softmax(out, dim=1)\n all_preds.append(out.cpu().numpy())\n\n block_preds = np.concatenate(all_preds, axis=0) # (N_blocks, C, bD, bH, bW)\n n_classes = block_preds.shape[1]\n full_pred = _stitch_blocks(\n block_preds, grid, block_shape, pad, orig_shape, n_classes\n )\n\n # Squeeze class dim for single-class output\n if n_classes == 1:\n spatial = full_pred[0]\n else:\n spatial = full_pred # (C, D, H, W)\n\n out_img = nib.Nifti1Image(spatial.astype(np.float32), affine)\n return out_img\n\n\ndef predict_with_uncertainty(\n inputs: str | Path | np.ndarray | nib.Nifti1Image,\n model: nn.Module,\n n_samples: int = 10,\n block_shape: tuple[int, int, int] = (128, 128, 128),\n batch_size: int = 4,\n device: str | torch.device | None = None,\n) -> tuple[nib.Nifti1Image, nib.Nifti1Image, nib.Nifti1Image]:\n \"\"\"MC-Dropout / Bayesian uncertainty estimation.\n\n Runs ``n_samples`` stochastic forward passes with the model in **train**\n mode (activating Dropout and Pyro sampling in Bayesian layers) and\n returns mean label, predictive variance, and predictive entropy maps.\n\n Parameters\n ----------\n inputs : path, ndarray, or Nifti1Image\n Input brain MRI (same format as :func:`predict`).\n model : nn.Module\n Trained segmentation model. Should contain dropout or Bayesian\n layers so that repeated forward passes are stochastic.\n n_samples : int\n Number of Monte-Carlo forward passes.\n block_shape, batch_size, device\n Same semantics as :func:`predict`.\n\n Returns\n -------\n label_img : nib.Nifti1Image\n Mean class label (argmax over mean softmax probabilities).\n variance_img : nib.Nifti1Image\n Mean predictive variance across classes.\n entropy_img : nib.Nifti1Image\n Predictive entropy of the mean softmax distribution.\n \"\"\"\n if device is None:\n device = get_device()\n device = torch.device(device)\n\n affine = np.eye(4)\n if isinstance(inputs, (str, Path)):\n img = nib.load(str(inputs))\n arr = np.asarray(img.dataobj, dtype=np.float32)\n affine = img.affine\n elif isinstance(inputs, nib.Nifti1Image):\n arr = np.asarray(inputs.dataobj, dtype=np.float32)\n affine = inputs.affine\n else:\n arr = np.asarray(inputs, dtype=np.float32)\n\n orig_shape = arr.shape[:3]\n arr3d = arr if arr.ndim == 3 else arr[..., 0]\n\n padded, pad = _pad_to_multiple(arr3d, block_shape)\n blocks, grid = _extract_blocks(padded, block_shape)\n n_blocks = blocks.shape[0]\n\n model = model.to(device)\n # Use eval mode to preserve BatchNorm statistics.\n # Stochasticity is controlled via mc=True (KWYK/FFG models)\n # or inherent Pyro sampling (BayesianConv3d).\n model.eval()\n\n # Welford's online algorithm: accumulate mean and M2 incrementally\n # so we only keep 2 block-level arrays in memory, not n_samples copies.\n mean_probs: np.ndarray | None = None # running mean\n m2_probs: np.ndarray | None = None # running sum of squared deviations\n\n with torch.no_grad():\n for sample_idx in range(n_samples):\n preds: list[np.ndarray] = []\n for start in range(0, n_blocks, batch_size):\n chunk = blocks[start : start + batch_size]\n tensor = torch.from_numpy(chunk[:, None]).to(device)\n out = _forward(model, tensor, mc=True)\n probs = torch.softmax(out, dim=1).cpu().numpy()\n preds.append(probs)\n sample = np.concatenate(preds, axis=0) # (N_blocks, C, bD, bH, bW)\n\n if mean_probs is None:\n mean_probs = sample.copy()\n m2_probs = np.zeros_like(sample)\n else:\n delta = sample - mean_probs\n mean_probs += delta / (sample_idx + 1)\n delta2 = sample - mean_probs\n m2_probs += delta * delta2\n\n var_probs = m2_probs / max(n_samples, 1) # population variance\n del m2_probs\n n_classes = mean_probs.shape[1]\n\n # Reduce per-block before stitching to avoid materialising full (C, D, H, W)\n # Labels: argmax over classes per block → (N_blocks, 1, bD, bH, bW)\n if n_classes == 1:\n block_labels = (mean_probs[:, 0:1] > 0.5).astype(np.float32)\n else:\n block_labels = mean_probs.argmax(axis=1, keepdims=True).astype(np.float32)\n\n # Mean variance across classes per block → (N_blocks, 1, bD, bH, bW)\n block_var = var_probs.mean(axis=1, keepdims=True)\n del var_probs\n\n # Entropy per block → (N_blocks, 1, bD, bH, bW)\n eps = 1e-8\n block_entropy = -(mean_probs * np.log(mean_probs + eps)).sum(axis=1, keepdims=True)\n del mean_probs\n\n # Stitch scalar maps (n_classes=1 for each)\n labels = _stitch_blocks(block_labels, grid, block_shape, pad, orig_shape, 1)[0]\n mean_var = _stitch_blocks(block_var, grid, block_shape, pad, orig_shape, 1)[0]\n entropy = _stitch_blocks(block_entropy, grid, block_shape, pad, orig_shape, 1)[0]\n\n label_img = nib.Nifti1Image(labels, affine)\n var_img = nib.Nifti1Image(mean_var.astype(np.float32), affine)\n entropy_img = nib.Nifti1Image(entropy.astype(np.float32), affine)\n return label_img, var_img, entropy_img\n\n\n__all__ = [\"predict\", \"predict_with_uncertainty\"]\n" }, { "path": "nobrainer/processing/__init__.py", "content": "\"\"\"Scikit-learn-style estimator API for nobrainer.\n\nProvides high-level ``Segmentation``, ``Generation``, and ``Dataset``\nclasses that wrap the lower-level PyTorch internals.\n\"\"\"\n\nfrom .dataset import Dataset, PatchDataset, extract_patches\n\n__all__ = [\"Dataset\", \"PatchDataset\", \"extract_patches\"]\n\n# Optional: Segmentation (requires core models)\ntry:\n from .segmentation import Segmentation # noqa: F401\n\n __all__.append(\"Segmentation\")\nexcept ImportError:\n pass\n\n# Optional: Generation (requires pytorch-lightning)\ntry:\n from .generation import Generation # noqa: F401\n\n __all__.append(\"Generation\")\nexcept ImportError:\n pass\n" }, { "path": "nobrainer/processing/base.py", "content": "\"\"\"Base estimator with Croissant-ML metadata persistence.\"\"\"\n\nfrom __future__ import annotations\n\nimport json\nfrom pathlib import Path\nfrom typing import Any\n\nimport torch\nimport torch.nn as nn\n\n\nclass BaseEstimator:\n \"\"\"Base class for all nobrainer estimators.\n\n Provides ``save()`` / ``load()`` with Croissant-ML JSON-LD metadata,\n and optional multi-GPU support via DDP.\n \"\"\"\n\n state_variables: list[str] = []\n model_: nn.Module | None = None\n _training_result: dict | None = None\n _dataset: Any = None\n\n def __init__(\n self,\n checkpoint_filepath: str | Path | None = None,\n multi_gpu: bool = True,\n ):\n self.checkpoint_filepath = checkpoint_filepath\n self.multi_gpu = multi_gpu\n\n @property\n def model(self) -> nn.Module:\n if self.model_ is None:\n raise RuntimeError(\"Model not trained. Call .fit() first.\")\n return self.model_\n\n def save(self, save_dir: str | Path) -> None:\n \"\"\"Save model.pth + croissant.json to directory.\"\"\"\n from .croissant import write_model_croissant\n\n save_dir = Path(save_dir)\n save_dir.mkdir(parents=True, exist_ok=True)\n torch.save(self.model_.state_dict(), save_dir / \"model.pth\")\n write_model_croissant(save_dir, self, self._training_result, self._dataset)\n\n @classmethod\n def load(cls, model_dir: str | Path, multi_gpu: bool = True) -> \"BaseEstimator\":\n \"\"\"Load estimator from directory with croissant.json metadata.\"\"\"\n model_dir = Path(model_dir)\n metadata = json.loads((model_dir / \"croissant.json\").read_text())\n prov = metadata.get(\"nobrainer:provenance\", {})\n\n est = cls.__new__(cls)\n est.multi_gpu = multi_gpu\n est.checkpoint_filepath = None\n est._training_result = None\n est._dataset = None\n\n # Subclass-specific reconstruction\n est._restore_from_provenance(prov)\n est.model_ = est._build_model()\n est.model_.load_state_dict(\n torch.load(model_dir / \"model.pth\", weights_only=True)\n )\n return est\n\n def _build_model(self) -> nn.Module:\n \"\"\"Reconstruct model architecture. Override in subclasses.\"\"\"\n raise NotImplementedError\n\n def _restore_from_provenance(self, prov: dict) -> None:\n \"\"\"Restore state from provenance dict. Override in subclasses.\"\"\"\n raise NotImplementedError\n" }, { "path": "nobrainer/processing/croissant.py", "content": "\"\"\"Croissant-ML JSON-LD metadata helpers for nobrainer estimators.\"\"\"\n\nfrom __future__ import annotations\n\nimport datetime\nimport hashlib\nimport json\nfrom pathlib import Path\nfrom typing import Any\n\n\ndef _sha256(path: str | Path) -> str:\n \"\"\"Compute SHA-256 hex digest of a file.\"\"\"\n h = hashlib.sha256()\n with open(path, \"rb\") as f:\n for chunk in iter(lambda: f.read(1 << 16), b\"\"):\n h.update(chunk)\n return h.hexdigest()\n\n\ndef _dataset_checksums(dataset: Any) -> list[dict]:\n \"\"\"Extract file paths and SHA256 checksums from a Dataset.\"\"\"\n if dataset is None:\n return []\n checksums = []\n for item in getattr(dataset, \"data\", []):\n img = item.get(\"image\", \"\") if isinstance(item, dict) else \"\"\n if img and Path(img).exists():\n checksums.append({\"path\": str(img), \"sha256\": _sha256(img)})\n return checksums\n\n\ndef write_model_croissant(\n save_dir: Path,\n estimator: Any,\n training_result: dict | None,\n dataset: Any,\n) -> Path:\n \"\"\"Write croissant.json with Croissant-ML JSON-LD metadata.\n\n Includes provenance (source datasets with SHA256), training parameters,\n model architecture info, and version stamps.\n \"\"\"\n import torch\n\n import nobrainer\n\n result = training_result or {}\n\n # Extract optimizer info from estimator if available\n opt_class = getattr(estimator, \"_optimizer_class\", \"Adam\")\n opt_args = getattr(estimator, \"_optimizer_args\", {})\n loss_name = getattr(estimator, \"_loss_name\", \"unknown\")\n\n metadata = {\n \"@context\": {\"@vocab\": \"http://mlcommons.org/croissant/\"},\n \"@type\": \"cr:Dataset\",\n \"name\": f\"nobrainer-{getattr(estimator, 'base_model', 'model')}\",\n \"description\": (\n f\"Trained {getattr(estimator, 'base_model', 'model')} model \"\n f\"via nobrainer\"\n ),\n \"distribution\": [\n {\n \"@type\": \"cr:FileObject\",\n \"name\": \"model.pth\",\n \"contentUrl\": \"model.pth\",\n \"encodingFormat\": \"application/x-pytorch\",\n }\n ],\n \"nobrainer:provenance\": {\n \"source_datasets\": _dataset_checksums(dataset),\n \"training_date\": datetime.datetime.now(datetime.timezone.utc).isoformat(),\n \"nobrainer_version\": nobrainer.__version__,\n \"pytorch_version\": torch.__version__,\n \"optimizer\": {\n \"class\": str(opt_class),\n \"args\": {k: str(v) for k, v in (opt_args or {}).items()},\n },\n \"loss_function\": str(loss_name),\n \"epochs_trained\": len(result.get(\"history\", [])),\n \"final_loss\": (\n result[\"history\"][-1].get(\"loss\") if result.get(\"history\") else None\n ),\n \"best_loss\": (\n min(\n (h[\"loss\"] for h in result[\"history\"] if h.get(\"loss\") is not None),\n default=None,\n )\n if result.get(\"history\")\n else None\n ),\n \"model_architecture\": getattr(estimator, \"base_model\", \"unknown\"),\n \"model_args\": getattr(estimator, \"model_args\", None) or {},\n \"n_classes\": getattr(estimator, \"n_classes_\", None),\n \"block_shape\": list(getattr(estimator, \"block_shape_\", []) or []),\n \"gpu_count\": torch.cuda.device_count() if torch.cuda.is_available() else 0,\n },\n }\n\n out = save_dir / \"croissant.json\"\n out.write_text(json.dumps(metadata, indent=2, default=str))\n return out\n\n\ndef write_checkpoint_croissant(\n checkpoint_dir: Path,\n model: Any,\n optimizer: Any,\n criterion: Any,\n history: list[dict],\n) -> Path:\n \"\"\"Write croissant.json alongside a training checkpoint.\n\n Lighter-weight than :func:`write_model_croissant` — works with the raw\n model/optimizer/criterion objects available inside :func:`~nobrainer.training.fit`\n rather than requiring an estimator wrapper.\n \"\"\"\n import torch\n\n import nobrainer\n\n checkpoint_dir = Path(checkpoint_dir)\n\n metadata = {\n \"@context\": {\"@vocab\": \"http://mlcommons.org/croissant/\"},\n \"@type\": \"cr:Dataset\",\n \"name\": f\"nobrainer-{type(model).__name__}\",\n \"description\": f\"Trained {type(model).__name__} checkpoint via nobrainer\",\n \"distribution\": [\n {\n \"@type\": \"cr:FileObject\",\n \"name\": \"best_model.pth\",\n \"contentUrl\": \"best_model.pth\",\n \"encodingFormat\": \"application/x-pytorch\",\n }\n ],\n \"nobrainer:provenance\": {\n \"training_date\": datetime.datetime.now(datetime.timezone.utc).isoformat(),\n \"nobrainer_version\": nobrainer.__version__,\n \"pytorch_version\": torch.__version__,\n \"optimizer\": {\n \"class\": type(optimizer).__name__,\n \"args\": {k: str(v) for k, v in optimizer.defaults.items()},\n },\n \"loss_function\": type(criterion).__name__,\n \"epochs_trained\": len(history),\n \"final_loss\": (history[-1].get(\"loss\") if history else None),\n \"best_loss\": (\n min(\n (h[\"loss\"] for h in history if h.get(\"loss\") is not None),\n default=None,\n )\n if history\n else None\n ),\n \"model_architecture\": type(model).__name__,\n \"gpu_count\": (\n torch.cuda.device_count() if torch.cuda.is_available() else 0\n ),\n },\n }\n\n out = checkpoint_dir / \"croissant.json\"\n out.write_text(json.dumps(metadata, indent=2, default=str))\n return out\n\n\ndef write_dataset_croissant(\n output_path: str | Path,\n dataset: Any,\n) -> Path:\n \"\"\"Write Croissant-ML JSON-LD for a Dataset.\"\"\"\n metadata = {\n \"@context\": {\"@vocab\": \"http://mlcommons.org/croissant/\"},\n \"@type\": \"cr:Dataset\",\n \"name\": \"nobrainer-dataset\",\n \"description\": \"Brain MRI dataset for nobrainer\",\n \"distribution\": [],\n \"recordSet\": [],\n }\n\n checksums = _dataset_checksums(dataset)\n for item in checksums:\n metadata[\"distribution\"].append(\n {\n \"@type\": \"cr:FileObject\",\n \"name\": Path(item[\"path\"]).name,\n \"contentUrl\": item[\"path\"],\n \"sha256\": item[\"sha256\"],\n }\n )\n\n metadata[\"nobrainer:dataset_info\"] = {\n \"volume_shape\": list(getattr(dataset, \"volume_shape\", []) or []),\n \"n_classes\": getattr(dataset, \"n_classes\", None),\n \"block_shape\": list(getattr(dataset, \"_block_shape\", []) or []),\n \"n_volumes\": len(getattr(dataset, \"data\", [])),\n }\n\n output_path = Path(output_path)\n output_path.write_text(json.dumps(metadata, indent=2, default=str))\n return output_path\n\n\ndef validate_croissant(path: str | Path) -> bool:\n \"\"\"Validate croissant.json using mlcroissant (if installed).\"\"\"\n try:\n import mlcroissant\n\n mlcroissant.Dataset(jsonld=str(path))\n return True\n except ImportError:\n return True # Skip validation if not installed\n except Exception:\n return False\n" }, { "path": "nobrainer/processing/dataset.py", "content": "\"\"\"Fluent Dataset builder for nobrainer estimators.\"\"\"\n\nfrom __future__ import annotations\n\nimport copy\nfrom pathlib import Path\nfrom typing import TYPE_CHECKING, Callable\n\nif TYPE_CHECKING:\n import zarr\n\nimport numpy as np\nimport torch\nfrom torch.utils.data import DataLoader\n\n# Named label mapping CSV locations (relative to package or absolute)\n_NAMED_MAPPINGS = {\n \"6-class\": \"6-class-mapping.csv\",\n \"50-class\": \"50-class-mapping.csv\",\n \"115-class\": \"115-class-mapping.csv\",\n}\n\n\ndef _load_label_mapping(name_or_path: str) -> Callable:\n \"\"\"Load a label mapping CSV and return a remap function.\n\n Accepts named mappings (\"6-class\", \"50-class\", \"115-class\") or a\n path to a CSV with ``original`` and ``new`` columns.\n \"\"\"\n import csv as csv_mod\n\n if name_or_path in _NAMED_MAPPINGS:\n csv_name = _NAMED_MAPPINGS[name_or_path]\n # Primary: inside the nobrainer package (works with pip install)\n pkg_data = Path(__file__).parent.parent / \"data\" / \"label_mappings\" / csv_name\n # Fallback: scripts dir (editable installs / development)\n scripts_data = (\n Path(__file__).parent.parent.parent\n / \"scripts\"\n / \"kwyk_reproduction\"\n / \"label_mappings\"\n / csv_name\n )\n candidates = [pkg_data, scripts_data]\n csv_path = None\n for c in candidates:\n if c.exists():\n csv_path = c\n break\n if csv_path is None:\n raise FileNotFoundError(\n f\"Label mapping '{name_or_path}' not found. \"\n f\"Searched: {[str(c) for c in candidates]}\"\n )\n else:\n csv_path = Path(name_or_path)\n if not csv_path.exists():\n raise FileNotFoundError(f\"Label mapping CSV not found: {csv_path}\")\n\n # Parse CSV: build original → new lookup\n lookup = {}\n with open(csv_path) as f:\n reader = csv_mod.DictReader(f)\n for row in reader:\n orig = int(row[\"original\"])\n new = int(row[\"new\"])\n lookup[orig] = new\n\n return _LabelRemap(lookup)\n\n\nclass _LabelRemap:\n \"\"\"Picklable label remapping callable (needed for DataLoader workers).\"\"\"\n\n def __init__(self, lookup: dict[int, int]):\n self.lookup = lookup\n\n def __call__(self, x):\n result = torch.zeros_like(x)\n for orig_val, new_val in self.lookup.items():\n result[x == orig_val] = new_val\n return result.long()\n\n\nclass Dataset:\n \"\"\"Fluent dataset builder wrapping the nobrainer data pipeline.\n\n Example::\n\n ds_train, ds_eval = (\n Dataset.from_files(filepaths, block_shape=(128,128,128))\n .batch(2)\n .augment()\n .normalize()\n .split(eval_size=0.1)\n )\n loader = ds_train.dataloader\n \"\"\"\n\n def __init__(\n self,\n data: list[dict[str, str]],\n volume_shape: tuple | None = None,\n n_classes: int = 1,\n ):\n self.data = data\n self.volume_shape = volume_shape\n self.n_classes = n_classes\n self._block_shape: tuple | None = None\n self._batch_size: int = 1\n self._shuffle: bool = False\n self._augment: bool = False\n self._augment_profile: str = \"standard\"\n self._binarize: bool = False\n self._streaming: bool = False\n self._patches_per_volume: int = 10\n self._normalizer: Callable | None = None\n self._dataloader: DataLoader | None = None\n\n @classmethod\n def from_files(\n cls,\n filepaths: list[tuple[str, str]] | list[dict[str, str]],\n block_shape: tuple[int, int, int] | None = None,\n n_classes: int = 1,\n ) -> \"Dataset\":\n \"\"\"Create a Dataset from file paths.\n\n Parameters\n ----------\n filepaths : list\n Either ``[(img, label), ...]`` tuples or\n ``[{\"image\": img, \"label\": label}, ...]`` dicts.\n block_shape : tuple or None\n Spatial patch size for extraction. None loads full volumes.\n n_classes : int\n Number of label classes.\n \"\"\"\n # Normalize to list of dicts\n if filepaths and isinstance(filepaths[0], (list, tuple)):\n data = [{\"image\": str(img), \"label\": str(lbl)} for img, lbl in filepaths]\n else:\n data = [{k: str(v) for k, v in d.items()} for d in filepaths]\n\n # Detect volume shape from first file\n volume_shape = None\n if data:\n import nibabel as nib\n\n first = data[0][\"image\"]\n if Path(first).suffix in (\".zarr\",):\n pass # Zarr shape detection deferred to dataloader\n else:\n try:\n volume_shape = nib.load(first).shape[:3]\n except Exception:\n pass\n\n ds = cls(data=data, volume_shape=volume_shape, n_classes=n_classes)\n ds._block_shape = block_shape\n return ds\n\n @classmethod\n def from_zarr(\n cls,\n store_path: str | Path,\n block_shape: tuple[int, int, int] | None = None,\n n_classes: int = 1,\n partition: str | None = None,\n partition_path: str | Path | None = None,\n ) -> \"Dataset\":\n \"\"\"Create a Dataset from a Zarr3 store.\n\n Parameters\n ----------\n store_path : str or Path\n Path to a Zarr store created by\n :func:`nobrainer.datasets.zarr_store.create_zarr_store`.\n block_shape : tuple or None\n Spatial patch size.\n n_classes : int\n Number of label classes.\n partition : str or None\n Partition to use: ``\"train\"``, ``\"val\"``, ``\"test\"``, or None (all).\n partition_path : str or Path or None\n Path to partition JSON. If None and partition is set, looks for\n ``_partition.json``.\n \"\"\"\n from nobrainer.datasets.zarr_store import load_partition, store_info\n\n store_path = Path(store_path)\n info = store_info(store_path)\n subject_ids = info[\"subject_ids\"]\n volume_shape = tuple(info[\"volume_shape\"])\n\n # Filter by partition\n if partition is not None:\n if partition_path is None:\n partition_path = Path(str(store_path) + \"_partition.json\")\n parts = load_partition(partition_path)\n if partition not in parts:\n raise ValueError(\n f\"Partition '{partition}' not found. \"\n f\"Available: {list(parts.keys())}\"\n )\n subject_ids = parts[partition]\n\n # Build data list referencing zarr indices\n id_to_idx = {sid: i for i, sid in enumerate(info[\"subject_ids\"])}\n data = []\n for sid in subject_ids:\n idx = id_to_idx[sid]\n data.append(\n {\n \"image\": f\"zarr://{store_path}#images/{idx}\",\n \"label\": f\"zarr://{store_path}#labels/{idx}\",\n \"_zarr_store\": str(store_path),\n \"_zarr_index\": idx,\n \"_subject_id\": sid,\n }\n )\n\n ds = cls(data=data, volume_shape=volume_shape, n_classes=n_classes)\n ds._block_shape = block_shape\n ds._zarr_store_path = str(store_path)\n return ds\n\n # --- Fluent API ---\n\n def batch(self, batch_size: int) -> \"Dataset\":\n \"\"\"Set batch size.\"\"\"\n self._batch_size = batch_size\n self._dataloader = None # invalidate cache\n return self\n\n def binarize(self, labels: str | set[int] | Callable | None = None) -> \"Dataset\":\n \"\"\"Binarize or remap labels.\n\n Parameters\n ----------\n labels : str, set of ints, callable, or None\n - ``None`` (default): any non-zero value → 1\n - ``\"binary\"``: same as None (any non-zero → 1)\n - ``\"6-class\"``, ``\"50-class\"``, ``\"115-class\"``: named\n parcellation from nobrainer_training_scripts mapping CSVs\n - ``set``: voxels with values in the set → 1, all others → 0\n - ``callable``: custom ``fn(label_tensor) → tensor``\n - ``str`` (path): path to a custom mapping CSV with\n ``original,new`` columns\n\n Examples\n --------\n Brain extraction (any tissue)::\n\n ds.binarize()\n\n Named parcellation::\n\n ds.binarize(labels=\"50-class\")\n\n Select specific FreeSurfer regions (e.g., hippocampus L+R)::\n\n ds.binarize(labels={17, 53})\n\n Custom mapping CSV::\n\n ds.binarize(labels=\"/path/to/mapping.csv\")\n \"\"\"\n if isinstance(labels, str) and labels not in (\"binary\",):\n # Named mapping or CSV path\n self._binarize = _load_label_mapping(labels)\n elif labels is not None:\n self._binarize = labels\n else:\n self._binarize = True\n self._dataloader = None\n return self\n\n def shuffle(self, buffer_size: int = 100) -> \"Dataset\":\n \"\"\"Enable shuffling.\"\"\"\n self._shuffle = True\n self._dataloader = None\n return self\n\n def augment(self, profile: str | bool = True) -> \"Dataset\":\n \"\"\"Enable data augmentation.\n\n Parameters\n ----------\n profile : str or bool\n ``True`` or ``\"standard\"`` for the standard profile.\n Named profiles: ``\"none\"``, ``\"light\"``, ``\"standard\"``, ``\"heavy\"``.\n ``False`` disables augmentation.\n \"\"\"\n if profile is False or profile == \"none\":\n self._augment = False\n elif profile is True:\n self._augment = True\n self._augment_profile = \"standard\"\n elif isinstance(profile, str):\n self._augment = True\n self._augment_profile = profile\n self._dataloader = None\n return self\n\n def mix(\n self,\n generator: \"torch.utils.data.Dataset\",\n ratio: float = 0.3,\n ) -> \"Dataset\":\n \"\"\"Combine this dataset with a synthetic data generator.\n\n Creates a mixed dataset where each sample is drawn from either\n the real data (this dataset) or the synthetic generator, based\n on the ratio.\n\n Parameters\n ----------\n generator : torch.utils.data.Dataset\n Synthetic data source (e.g., ``SynthSegGenerator``).\n Must return ``{\"image\": Tensor, \"label\": Tensor}`` dicts.\n ratio : float\n Fraction of samples drawn from the generator (default 0.3 = 30%).\n\n Returns\n -------\n Dataset\n A new Dataset wrapping a ``MixedDataset``.\n \"\"\"\n\n mixed = MixedDataset(self, generator, ratio=ratio)\n new_ds = Dataset(\n data=self.data, volume_shape=self.volume_shape, n_classes=self.n_classes\n )\n new_ds._block_shape = self._block_shape\n new_ds._batch_size = self._batch_size\n new_ds._augment = self._augment\n new_ds._augment_profile = self._augment_profile\n new_ds._mixed_dataset = mixed\n new_ds._dataloader = None\n return new_ds\n\n def streaming(self, patches_per_volume: int = 10) -> \"Dataset\":\n \"\"\"Use streaming patch extraction (no full-volume loading).\n\n Instead of loading entire volumes and cropping in memory (MONAI\n pipeline), patches are read directly from disk. For Zarr stores,\n only the chunks overlapping the requested patch are fetched —\n enabling efficient cloud and large-dataset training.\n\n Requires ``block_shape`` to be set via ``from_files()`` or\n ``batch()`` first.\n\n Parameters\n ----------\n patches_per_volume : int\n Random patches per volume per epoch.\n\n Example\n -------\n ::\n\n ds = (Dataset.from_files(paths, block_shape=(64,64,64))\n .batch(4).binarize().streaming(patches_per_volume=20))\n \"\"\"\n self._streaming = True\n self._patches_per_volume = patches_per_volume\n self._dataloader = None\n return self\n\n def normalize(self, fn: Callable | None = None) -> \"Dataset\":\n \"\"\"Set normalization function.\"\"\"\n self._normalizer = fn\n self._dataloader = None\n return self\n\n def split(self, eval_size: float = 0.1) -> tuple[\"Dataset\", \"Dataset\"]:\n \"\"\"Split into train and eval datasets.\"\"\"\n n = len(self.data)\n n_eval = max(1, int(n * eval_size))\n indices = np.random.permutation(n)\n eval_idx = indices[:n_eval]\n train_idx = indices[n_eval:]\n\n train_ds = copy.copy(self)\n train_ds.data = [self.data[i] for i in train_idx]\n train_ds._dataloader = None\n\n eval_ds = copy.copy(self)\n eval_ds.data = [self.data[i] for i in eval_idx]\n eval_ds._dataloader = None\n\n return train_ds, eval_ds\n\n @property\n def dataloader(self) -> DataLoader:\n \"\"\"Lazily build and return a PyTorch DataLoader.\"\"\"\n if self._dataloader is not None:\n return self._dataloader\n\n # Streaming mode: use PatchDataset for on-the-fly patch extraction\n if self._streaming:\n # Build augmentation transforms if enabled\n transforms = None\n if self._augment:\n from monai.transforms import Compose\n\n from nobrainer.augmentation.profiles import get_augmentation_profile\n\n aug_transforms = get_augmentation_profile(\n self._augment_profile, keys=[\"image\", \"label\"]\n )\n if aug_transforms:\n transforms = Compose(aug_transforms)\n\n patch_ds = PatchDataset(\n data=self.data,\n block_shape=self._block_shape or (32, 32, 32),\n patches_per_volume=self._patches_per_volume,\n binarize=self._binarize if self._binarize else None,\n transforms=transforms,\n )\n # Use multiple workers for I/O prefetching — each worker loads\n # patches independently while GPU processes the current batch.\n # Respect SLURM allocation or fall back to cpu_count.\n import os\n\n slurm_cpus = os.environ.get(\"SLURM_CPUS_PER_TASK\")\n max_cpus = int(slurm_cpus) if slurm_cpus else (os.cpu_count() or 1)\n n_workers = max(1, max_cpus - 1) # leave 1 CPU for main process\n self._dataloader = DataLoader(\n patch_ds,\n batch_size=self._batch_size,\n shuffle=self._shuffle,\n num_workers=n_workers,\n prefetch_factor=2,\n persistent_workers=True if n_workers > 0 else False,\n pin_memory=torch.cuda.is_available(),\n )\n return self._dataloader\n\n image_paths = [d[\"image\"] for d in self.data]\n label_paths = [d[\"label\"] for d in self.data if \"label\" in d] or None\n\n # Check for Zarr paths\n is_zarr = any(str(p).rstrip(\"/\").endswith(\".zarr\") for p in image_paths)\n\n if is_zarr:\n from nobrainer.dataset import ZarrDataset\n\n zarr_data = self.data\n ds = ZarrDataset(zarr_data)\n self._dataloader = DataLoader(\n ds,\n batch_size=self._batch_size,\n shuffle=self._shuffle,\n pin_memory=torch.cuda.is_available(),\n )\n else:\n from nobrainer.dataset import get_dataset\n\n self._dataloader = get_dataset(\n image_paths=image_paths,\n label_paths=label_paths,\n block_shape=self._block_shape,\n batch_size=self._batch_size,\n augment=self._augment,\n binarize_labels=self._binarize,\n )\n\n return self._dataloader\n\n @property\n def batch_size(self) -> int:\n return self._batch_size\n\n @property\n def block_shape(self) -> tuple | None:\n return self._block_shape\n\n def to_croissant(self, output_path: str | Path) -> Path:\n \"\"\"Export dataset metadata as Croissant-ML JSON-LD.\"\"\"\n from .croissant import write_dataset_croissant\n\n return write_dataset_croissant(output_path, self)\n\n\ndef extract_patches(\n volume: np.ndarray,\n label: np.ndarray | None = None,\n block_shape: tuple[int, int, int] = (32, 32, 32),\n n_patches: int = 10,\n binarize: bool | set | Callable | None = None,\n) -> list[tuple[np.ndarray, ...]] | list[np.ndarray]:\n \"\"\"Extract random patches from a 3D volume.\n\n Parameters\n ----------\n volume : ndarray\n 3D volume of shape ``(D, H, W)`` or path loadable by nibabel.\n label : ndarray or None\n Corresponding label volume. If None, only image patches returned.\n block_shape : tuple\n Spatial size of each patch ``(bD, bH, bW)``.\n n_patches : int\n Number of random patches to extract.\n binarize : bool, set, callable, or None\n If not None, applied to label patches:\n - ``True``: any non-zero → 1\n - ``set``: voxels in set → 1\n - ``callable``: custom ``fn(patch) → patch``\n\n Returns\n -------\n list of tuples ``(image_patch, label_patch)`` if label given,\n or list of ``image_patch`` arrays if label is None.\n\n Examples\n --------\n ::\n\n import nibabel as nib\n vol = nib.load(\"brain.nii.gz\").get_fdata()\n lbl = nib.load(\"label.nii.gz\").get_fdata()\n patches = extract_patches(vol, lbl, block_shape=(32, 32, 32), n_patches=20)\n # patches[0] = (image_patch, label_patch), each shape (32, 32, 32)\n \"\"\"\n import nibabel as nib\n\n # Load from path if needed\n if isinstance(volume, (str, Path)):\n volume = np.asarray(nib.load(str(volume)).dataobj, dtype=np.float32)\n if isinstance(label, (str, Path)):\n label = np.asarray(nib.load(str(label)).dataobj, dtype=np.float32)\n\n vol = np.asarray(volume, dtype=np.float32)\n bd, bh, bw = block_shape\n D, H, W = vol.shape[:3]\n\n patches = []\n for _ in range(n_patches):\n d0 = np.random.randint(0, max(1, D - bd + 1))\n h0 = np.random.randint(0, max(1, H - bh + 1))\n w0 = np.random.randint(0, max(1, W - bw + 1))\n\n img_patch = vol[d0 : d0 + bd, h0 : h0 + bh, w0 : w0 + bw]\n\n if label is not None:\n lbl = np.asarray(label, dtype=np.float32)\n lbl_patch = lbl[d0 : d0 + bd, h0 : h0 + bh, w0 : w0 + bw]\n\n # Apply binarization\n if binarize is True:\n lbl_patch = (lbl_patch > 0).astype(np.float32)\n elif isinstance(binarize, set):\n mask = np.zeros_like(lbl_patch)\n for val in binarize:\n mask = np.maximum(mask, (lbl_patch == val).astype(np.float32))\n lbl_patch = mask\n elif callable(binarize):\n lbl_patch = binarize(lbl_patch)\n\n patches.append((img_patch, lbl_patch))\n else:\n patches.append(img_patch)\n\n return patches\n\n\nclass PatchDataset(torch.utils.data.Dataset):\n \"\"\"Streaming patch dataset — generates random patches on-the-fly.\n\n Instead of pre-extracting patches or loading full volumes into memory,\n this dataset lazily reads only the voxels needed for each patch. For\n Zarr v3 stores, this uses chunk-aligned partial I/O (only the chunks\n overlapping the patch are read from disk/cloud).\n\n Parameters\n ----------\n data : list of dicts\n ``[{\"image\": path, \"label\": path}, ...]``. Paths can be NIfTI\n (``.nii``, ``.nii.gz``, ``.mgz``) or Zarr (``.zarr``).\n block_shape : tuple\n Spatial size of each patch ``(bD, bH, bW)``.\n patches_per_volume : int\n Number of random patches to yield per volume per epoch.\n binarize : bool, set, callable, or None\n Label remapping (see :func:`extract_patches`).\n transforms : callable or None\n Optional transform applied to each ``(image, label)`` dict after\n extraction (e.g., normalization, augmentation).\n\n Examples\n --------\n ::\n\n from nobrainer.processing.dataset import PatchDataset\n\n ds = PatchDataset(\n data=[{\"image\": \"sub-01.zarr\", \"label\": \"sub-01_label.zarr\"}],\n block_shape=(64, 64, 64),\n patches_per_volume=10,\n binarize=True,\n )\n loader = DataLoader(ds, batch_size=4, num_workers=2)\n\n Each epoch yields ``len(data) * patches_per_volume`` patches, with\n different random locations each time.\n \"\"\"\n\n def __init__(\n self,\n data: list[dict[str, str]],\n block_shape: tuple[int, int, int] = (32, 32, 32),\n patches_per_volume: int = 10,\n binarize: bool | set | Callable | None = None,\n transforms: Callable | None = None,\n ):\n self.data = data\n self.block_shape = block_shape\n self.patches_per_volume = patches_per_volume\n self.binarize = binarize\n self.transforms = transforms\n\n # Cache zarr store handles (opened once, reused for all reads)\n self._zarr_cache: dict[str, zarr.Group] = {}\n\n # Cache volume shapes — use zarr metadata when available (fast)\n self._shapes: list[tuple[int, ...]] = []\n first_parsed = self._parse_zarr_path(str(data[0][\"image\"])) if data else None\n if first_parsed is not None:\n # All items share the same zarr store — read shape once\n store = self._get_zarr_store(first_parsed[0])\n spatial_shape = store[first_parsed[1]].shape[1:] # (D, H, W)\n self._shapes = [spatial_shape] * len(data)\n else:\n for item in data:\n self._shapes.append(self._get_shape(item[\"image\"]))\n\n def __len__(self) -> int:\n return len(self.data) * self.patches_per_volume\n\n def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:\n vol_idx = idx // self.patches_per_volume\n item = self.data[vol_idx]\n shape = self._shapes[vol_idx]\n\n # Random patch origin\n bd, bh, bw = self.block_shape\n d0 = np.random.randint(0, max(1, shape[0] - bd + 1))\n h0 = np.random.randint(0, max(1, shape[1] - bh + 1))\n w0 = np.random.randint(0, max(1, shape[2] - bw + 1))\n slc = (slice(d0, d0 + bd), slice(h0, h0 + bh), slice(w0, w0 + bw))\n\n # Read only the patch region (cached zarr handles for speed)\n img_patch = self._read_region_cached(item[\"image\"], slc).astype(np.float32)\n\n result: dict[str, torch.Tensor] = {\n \"image\": torch.from_numpy(img_patch[None]), # add channel dim\n }\n\n if \"label\" in item:\n lbl_patch = self._read_region_cached(item[\"label\"], slc).astype(np.float32)\n lbl_patch = self._apply_binarize(lbl_patch)\n result[\"label\"] = torch.from_numpy(lbl_patch[None])\n\n if self.transforms is not None:\n result = self.transforms(result)\n\n return result\n\n def _apply_binarize(self, lbl: np.ndarray) -> np.ndarray:\n \"\"\"Apply binarization to a label patch.\"\"\"\n if self.binarize is True:\n return (lbl > 0).astype(np.float32)\n elif isinstance(self.binarize, set):\n mask = np.zeros_like(lbl)\n for val in self.binarize:\n mask = np.maximum(mask, (lbl == val).astype(np.float32))\n return mask\n elif callable(self.binarize):\n # Remap functions may expect torch tensors (e.g., _load_label_mapping)\n t = torch.from_numpy(lbl.astype(np.int32))\n result = self.binarize(t)\n return result.numpy().astype(np.float32)\n return lbl\n\n @staticmethod\n def _parse_zarr_path(path: str) -> tuple[str, str, int] | None:\n \"\"\"Parse zarr://store_path#array_name/subject_index.\n\n Returns ``(store_path, array_name, subject_index)`` or None.\n \"\"\"\n if path.startswith(\"zarr://\"):\n rest = path[len(\"zarr://\") :]\n if \"#\" in rest:\n store_path, fragment = rest.split(\"#\", 1)\n parts = fragment.rsplit(\"/\", 1)\n if len(parts) == 2:\n return store_path, parts[0], int(parts[1])\n return store_path, fragment, 0\n return rest, \"images\", 0\n return None\n\n @staticmethod\n def _get_shape(path: str) -> tuple[int, ...]:\n \"\"\"Get volume shape without loading full data.\"\"\"\n path = str(path)\n parsed = PatchDataset._parse_zarr_path(path)\n if parsed is not None:\n import zarr\n\n store_path, array_name, idx = parsed\n store = zarr.open_group(store_path, mode=\"r\")\n # Shape of the 4D array is (N, D, H, W); return spatial (D, H, W)\n return store[array_name].shape[1:]\n elif path.rstrip(\"/\").endswith(\".zarr\"):\n import zarr\n\n store = zarr.open_group(path, mode=\"r\")\n return store[\"0\"].shape\n else:\n import nibabel as nib\n\n return nib.load(path).shape[:3]\n\n def _get_zarr_store(self, store_path: str):\n \"\"\"Get or create a cached zarr group handle.\"\"\"\n if store_path not in self._zarr_cache:\n import zarr\n\n self._zarr_cache[store_path] = zarr.open_group(store_path, mode=\"r\")\n return self._zarr_cache[store_path]\n\n def _read_region_cached(self, path: str, slc: tuple[slice, ...]) -> np.ndarray:\n \"\"\"Read a spatial region, using cached zarr handles.\"\"\"\n path = str(path)\n parsed = self._parse_zarr_path(path)\n if parsed is not None:\n store_path, array_name, idx = parsed\n store = self._get_zarr_store(store_path)\n sd, sh, sw = slc\n return np.asarray(store[array_name][idx, sd, sh, sw])\n return self._read_region(path, slc)\n\n @staticmethod\n def _read_region(path: str, slc: tuple[slice, ...]) -> np.ndarray:\n \"\"\"Read a spatial region from a volume (static, no caching).\"\"\"\n path = str(path)\n parsed = PatchDataset._parse_zarr_path(path)\n if parsed is not None:\n import zarr\n\n store_path, array_name, idx = parsed\n store = zarr.open_group(store_path, mode=\"r\")\n sd, sh, sw = slc\n return np.asarray(store[array_name][idx, sd, sh, sw])\n elif path.rstrip(\"/\").endswith(\".zarr\"):\n import zarr\n\n store = zarr.open_group(path, mode=\"r\")\n return np.asarray(store[\"0\"][slc])\n else:\n import nibabel as nib\n\n img = nib.load(path)\n return np.asarray(img.dataobj[slc])\n\n\nclass MixedDataset(torch.utils.data.Dataset):\n \"\"\"Combine a real dataset with a synthetic generator at a given ratio.\n\n Each ``__getitem__`` call randomly selects from either the real data\n or the generator based on the ratio.\n\n Parameters\n ----------\n real_dataset : Dataset or torch.utils.data.Dataset\n The real data source.\n generator : torch.utils.data.Dataset\n Synthetic data source (e.g., ``SynthSegGenerator``).\n ratio : float\n Fraction of samples from the generator (0.3 = 30% synthetic).\n \"\"\"\n\n def __init__(\n self,\n real_dataset: \"Dataset | torch.utils.data.Dataset\",\n generator: torch.utils.data.Dataset,\n ratio: float = 0.3,\n ) -> None:\n self.real_dataset = real_dataset\n self.generator = generator\n self.ratio = ratio\n # Total length is the max of real and synthetic\n self._real_len = len(real_dataset) if hasattr(real_dataset, \"__len__\") else 0\n self._gen_len = len(generator)\n\n def __len__(self) -> int:\n return max(self._real_len, self._gen_len)\n\n def __getitem__(self, idx: int) -> dict:\n import random\n\n if random.random() < self.ratio:\n # Synthetic sample\n gen_idx = idx % self._gen_len\n return self.generator[gen_idx]\n else:\n # Real sample\n real_idx = idx % max(self._real_len, 1)\n if hasattr(self.real_dataset, \"dataloader\"):\n # Dataset object — use underlying data\n return self.real_dataset.data[real_idx]\n return self.real_dataset[real_idx]\n" }, { "path": "nobrainer/processing/generation.py", "content": "\"\"\"Generation estimator — scikit-learn-style API for GANs.\"\"\"\n\nfrom __future__ import annotations\n\nfrom pathlib import Path\nfrom typing import Any\n\nimport nibabel as nib\nimport numpy as np\nimport torch\nimport torch.nn as nn\n\nfrom .base import BaseEstimator\n\n\nclass Generation(BaseEstimator):\n \"\"\"Train and generate synthetic brain volumes.\n\n Example::\n\n gen = Generation(\"progressivegan\").fit(dataset, epochs=100)\n images = gen.generate(n_images=5)\n \"\"\"\n\n state_variables = [\"base_model\", \"model_args\", \"latent_size\"]\n\n def __init__(\n self,\n base_model: str = \"progressivegan\",\n model_args: dict | None = None,\n multi_gpu: bool = True,\n ):\n super().__init__(multi_gpu=multi_gpu)\n self.base_model = base_model\n self.model_args = model_args or {}\n self.latent_size = self.model_args.get(\"latent_size\", 256)\n\n def fit(\n self,\n dataset_train: Any,\n epochs: int = 100,\n **trainer_kwargs: Any,\n ) -> \"Generation\":\n \"\"\"Train the generative model using Lightning.\"\"\"\n import pytorch_lightning as pl\n\n from nobrainer.models import get as get_model\n\n factory = get_model(self.base_model)\n self.model_ = factory(**self.model_args)\n self.latent_size = getattr(self.model_, \"latent_size\", self.latent_size)\n\n loader = (\n dataset_train.dataloader\n if hasattr(dataset_train, \"dataloader\")\n else dataset_train\n )\n\n trainer_defaults = {\n \"max_steps\": epochs,\n \"accelerator\": \"auto\",\n \"devices\": 1,\n \"enable_checkpointing\": False,\n \"logger\": False,\n }\n trainer_defaults.update(trainer_kwargs)\n\n trainer = pl.Trainer(**trainer_defaults)\n trainer.fit(self.model_, loader)\n\n self._dataset = dataset_train\n self._training_result = {\n \"history\": [{\"epoch\": e, \"loss\": None} for e in range(1, epochs + 1)],\n \"checkpoint_path\": None,\n }\n return self\n\n def generate(\n self,\n n_images: int = 1,\n data_type: type | None = None,\n ) -> list[nib.Nifti1Image]:\n \"\"\"Generate synthetic brain volumes.\"\"\"\n self.model_.eval()\n gen = self.model_.generator\n gen.current_level = getattr(gen, \"current_level\", 0)\n gen.alpha = 1.0\n\n images = []\n with torch.no_grad():\n z = torch.randn(n_images, self.latent_size, device=self.model_.device)\n out = gen(z) # (N, 1, D, H, W)\n\n for i in range(n_images):\n arr = out[i, 0].cpu().numpy()\n if data_type is not None:\n arr = arr.astype(data_type)\n images.append(nib.Nifti1Image(arr, np.eye(4)))\n\n return images\n\n def save(self, save_dir: str | Path) -> None:\n \"\"\"Save Lightning checkpoint + croissant.json.\"\"\"\n from .croissant import write_model_croissant\n\n save_dir = Path(save_dir)\n save_dir.mkdir(parents=True, exist_ok=True)\n torch.save(self.model_.state_dict(), save_dir / \"model.pth\")\n write_model_croissant(save_dir, self, self._training_result, self._dataset)\n\n def _build_model(self) -> nn.Module:\n from nobrainer.models import get as get_model\n\n return get_model(self.base_model)(**self.model_args)\n\n def _restore_from_provenance(self, prov: dict) -> None:\n self.base_model = prov.get(\"model_architecture\", \"progressivegan\")\n self.model_args = prov.get(\"model_args\", {})\n self.latent_size = self.model_args.get(\"latent_size\", 256)\n" }, { "path": "nobrainer/processing/segmentation.py", "content": "\"\"\"Segmentation estimator — scikit-learn-style API.\"\"\"\n\nfrom __future__ import annotations\n\nfrom pathlib import Path\nfrom typing import Any, Callable\n\nimport nibabel as nib\nimport numpy as np\nimport torch\nimport torch.nn as nn\n\nfrom nobrainer.training import get_device\n\nfrom .base import BaseEstimator\n\n\nclass Segmentation(BaseEstimator):\n \"\"\"Train and run brain segmentation with a simple API.\n\n Example::\n\n seg = Segmentation(\"unet\").fit(dataset, epochs=5)\n result = seg.predict(\"brain.nii.gz\")\n seg.save(\"my_model\")\n \"\"\"\n\n state_variables = [\n \"base_model\",\n \"model_args\",\n \"block_shape_\",\n \"volume_shape_\",\n \"n_classes_\",\n ]\n\n def __init__(\n self,\n base_model: str = \"unet\",\n model_args: dict | None = None,\n checkpoint_filepath: str | Path | None = None,\n multi_gpu: bool = True,\n ):\n super().__init__(checkpoint_filepath, multi_gpu)\n self.base_model = base_model\n self.model_args = model_args or {}\n self.block_shape_: tuple | None = None\n self.volume_shape_: tuple | None = None\n self.n_classes_: int | None = None\n self._optimizer_class: str = \"Adam\"\n self._optimizer_args: dict = {}\n self._loss_name: str = \"unknown\"\n\n def fit(\n self,\n dataset_train: Any,\n dataset_validate: Any | None = None,\n epochs: int = 1,\n optimizer: type = torch.optim.Adam,\n opt_args: dict | None = None,\n loss: Callable | nn.Module | None = None,\n class_weights: torch.Tensor | str | None = None,\n metrics: Callable | None = None,\n callbacks: list | None = None,\n **kwargs,\n ) -> \"Segmentation\":\n \"\"\"Train the model and return self for chaining.\n\n Parameters\n ----------\n class_weights : Tensor, str, or None\n Per-class weights for CrossEntropyLoss. Pass a tensor of\n shape ``(n_classes,)``, ``\"auto\"`` to compute from training\n labels via inverse frequency, or None (default, no weighting).\n \"\"\"\n from nobrainer.models import get as get_model\n from nobrainer.training import fit as training_fit\n\n # Store metadata from dataset\n self.block_shape_ = getattr(dataset_train, \"block_shape\", None)\n self.volume_shape_ = getattr(dataset_train, \"volume_shape\", None)\n self.n_classes_ = getattr(dataset_train, \"n_classes\", 1)\n\n # Set n_classes in model_args\n model_args = {**self.model_args, \"n_classes\": self.n_classes_}\n factory = get_model(self.base_model)\n self.model_ = factory(**model_args)\n\n # Configure optimizer\n opt_args = opt_args or {\"lr\": 1e-3}\n opt = optimizer(self.model_.parameters(), **opt_args)\n self._optimizer_class = optimizer.__name__\n self._optimizer_args = opt_args\n\n # Configure class weights\n weights_tensor = None\n if class_weights is not None:\n if isinstance(class_weights, str) and class_weights == \"auto\":\n from nobrainer.losses import compute_class_weights\n\n label_paths = [\n p[1] if isinstance(p, (list, tuple)) else p.get(\"label\", p)\n for p in getattr(dataset_train, \"data\", [])\n ]\n label_mapping = getattr(dataset_train, \"_binarize_name\", None)\n weights_tensor = compute_class_weights(\n label_paths,\n self.n_classes_,\n label_mapping=label_mapping,\n max_samples=50,\n )\n elif isinstance(class_weights, torch.Tensor):\n weights_tensor = class_weights\n if weights_tensor is not None:\n self._class_weights = weights_tensor\n\n # Configure loss\n if loss is None:\n loss = nn.CrossEntropyLoss(weight=weights_tensor)\n self._loss_name = (\n \"WeightedCrossEntropyLoss\"\n if weights_tensor is not None\n else \"CrossEntropyLoss\"\n )\n elif callable(loss):\n self._loss_name = getattr(loss, \"__name__\", type(loss).__name__)\n if not isinstance(loss, nn.Module):\n loss = loss() # factory function like losses.dice()\n else:\n self._loss_name = type(loss).__name__\n\n # Train\n gpus = torch.cuda.device_count() if self.multi_gpu else 1\n loader = (\n dataset_train.dataloader\n if hasattr(dataset_train, \"dataloader\")\n else dataset_train\n )\n val_loader = None\n if dataset_validate is not None:\n val_loader = (\n dataset_validate.dataloader\n if hasattr(dataset_validate, \"dataloader\")\n else dataset_validate\n )\n self._training_result = training_fit(\n model=self.model_,\n loader=loader,\n criterion=loss,\n optimizer=opt,\n max_epochs=epochs,\n gpus=gpus,\n checkpoint_dir=self.checkpoint_filepath,\n callbacks=callbacks,\n val_loader=val_loader,\n checkpoint_freq=kwargs.get(\"checkpoint_freq\", 0),\n gradient_checkpointing=kwargs.get(\"gradient_checkpointing\", False),\n model_parallel=kwargs.get(\"model_parallel\", False),\n resume_from=kwargs.get(\"resume_from\"),\n )\n self._dataset = dataset_train\n return self\n\n def predict(\n self,\n x: str | Path | np.ndarray | nib.Nifti1Image,\n batch_size: int = 4,\n block_shape: tuple | None = None,\n normalizer: Callable | None = None,\n n_samples: int = 0,\n ) -> nib.Nifti1Image | tuple[nib.Nifti1Image, ...]:\n \"\"\"Predict on a volume.\n\n If ``n_samples > 0`` and model is Bayesian, returns\n ``(label, variance, entropy)`` tuple.\n \"\"\"\n from nobrainer.prediction import predict, predict_with_uncertainty\n\n bs = block_shape or self.block_shape_ or (128, 128, 128)\n\n if n_samples > 0:\n return predict_with_uncertainty(\n inputs=x,\n model=self.model,\n n_samples=n_samples,\n block_shape=bs,\n batch_size=batch_size,\n )\n return predict(\n inputs=x,\n model=self.model,\n block_shape=bs,\n batch_size=batch_size,\n normalizer=normalizer,\n )\n\n def evaluate(\n self,\n dataset: Any,\n metrics: Callable | None = None,\n ) -> dict:\n \"\"\"Evaluate model on a dataset. Returns dict with loss and metrics.\"\"\"\n device = get_device()\n self.model_.to(device).eval()\n criterion = nn.CrossEntropyLoss()\n\n total_loss = 0.0\n n_batches = 0\n loader = dataset.dataloader if hasattr(dataset, \"dataloader\") else dataset\n\n with torch.no_grad():\n for batch in loader:\n if isinstance(batch, dict):\n images = batch[\"image\"].to(device)\n labels = batch[\"label\"].to(device)\n else:\n images, labels = batch[0].to(device), batch[1].to(device)\n pred = self.model_(images)\n total_loss += criterion(pred, labels).item()\n n_batches += 1\n\n return {\n \"loss\": total_loss / max(n_batches, 1),\n \"n_batches\": n_batches,\n }\n\n def _build_model(self) -> nn.Module:\n \"\"\"Reconstruct model architecture from stored metadata.\"\"\"\n from nobrainer.models import get as get_model\n\n model_args = {**self.model_args, \"n_classes\": self.n_classes_}\n return get_model(self.base_model)(**model_args)\n\n def _restore_from_provenance(self, prov: dict) -> None:\n \"\"\"Restore state from croissant.json provenance.\"\"\"\n self.base_model = prov.get(\"model_architecture\", \"unet\")\n self.model_args = prov.get(\"model_args\", {})\n self.n_classes_ = prov.get(\"n_classes\", 1)\n self.block_shape_ = tuple(prov.get(\"block_shape\", []))\n self._optimizer_class = prov.get(\"optimizer\", {}).get(\"class\", \"Adam\")\n self._optimizer_args = prov.get(\"optimizer\", {}).get(\"args\", {})\n self._loss_name = prov.get(\"loss_function\", \"unknown\")\n" }, { "path": "nobrainer/research/__init__.py", "content": "\"\"\"Autoresearch sub-package for nobrainer (Phase 7 — US5/US6).\"\"\"\n\nfrom .loop import commit_best_model, run_loop\n\n__all__ = [\"commit_best_model\", \"run_loop\"]\n" }, { "path": "nobrainer/research/loop.py", "content": "\"\"\"Autoresearch loop for nobrainer.\n\nProposes hyperparameter diffs via the Anthropic API, applies them to a\ntraining script, runs the experiment subprocess, and keeps improvements.\n\nIf the Anthropic API is unavailable (no key or import error) the loop\nfalls back to a random perturbation from a pre-defined search grid.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport copy\nfrom dataclasses import dataclass, field\nimport json\nimport logging\nimport os\nfrom pathlib import Path\nimport shutil\nimport subprocess\nimport sys\nimport time\nfrom typing import Any\n\nlogger = logging.getLogger(__name__)\n\n_DEFAULT_SEARCH_GRID: dict[str, list[Any]] = {\n \"learning_rate\": [1e-4, 5e-4, 1e-3, 5e-3],\n \"batch_size\": [2, 4, 8],\n \"n_epochs\": [10, 20, 50],\n \"dropout_rate\": [0.0, 0.1, 0.25, 0.5],\n}\n\n\n@dataclass\nclass ExperimentResult:\n \"\"\"Structured record for one autoresearch experiment.\"\"\"\n\n run_id: int\n config: dict[str, Any]\n val_dice: float | None\n outcome: str # \"improved\", \"degraded\", \"failed\"\n failure_reason: str | None = None\n elapsed_seconds: float = 0.0\n notes: list[str] = field(default_factory=list)\n\n\ndef run_loop(\n working_dir: str | Path,\n model_family: str = \"bayesian_vnet\",\n max_experiments: int = 10,\n budget_hours: float = 8.0,\n train_script: str = \"train.py\",\n val_dice_file: str = \"val_dice.json\",\n budget_timeout_per_run: float = 3600.0,\n budget_seconds: float | None = None,\n) -> list[ExperimentResult]:\n \"\"\"Run the autoresearch experiment loop.\n\n Parameters\n ----------\n working_dir : path\n Directory containing the training script and where results are saved.\n model_family : str\n Model family name (e.g. ``\"bayesian_vnet\"``).\n max_experiments : int\n Maximum number of experiments to run.\n budget_hours : float\n Wall-clock budget in hours (loop stops when exceeded).\n train_script : str\n Filename of the training script relative to ``working_dir``.\n val_dice_file : str\n Filename of the validation Dice JSON written by the training script.\n budget_timeout_per_run : float\n Per-experiment subprocess timeout in seconds.\n\n Returns\n -------\n list[ExperimentResult]\n All experiment records (including failures).\n \"\"\"\n working_dir = Path(working_dir)\n train_path = working_dir / train_script\n val_dice_path = working_dir / val_dice_file\n backup_path = working_dir / f\"{train_script}.backup\"\n if budget_seconds is not None:\n budget_end = time.time() + budget_seconds\n else:\n budget_end = time.time() + budget_hours * 3600.0\n\n if not train_path.exists():\n raise FileNotFoundError(\n f\"Training script not found: {train_path}. \"\n \"Create it or copy from nobrainer.research.templates.\"\n )\n\n # Read initial config from train_script (look for JSON comment block)\n current_config = _parse_config_comment(train_path)\n best_dice: float | None = None\n results: list[ExperimentResult] = []\n\n logger.info(\"Starting autoresearch loop for %s\", model_family)\n logger.info(\"max_experiments=%d, budget_hours=%.1f\", max_experiments, budget_hours)\n\n for run_id in range(max_experiments):\n if time.time() >= budget_end:\n logger.info(\"Budget exhausted — stopping at experiment %d\", run_id)\n break\n\n # Propose new config\n new_config = _propose_config(current_config, model_family, run_id, best_dice)\n logger.info(\"Experiment %d config: %s\", run_id, new_config)\n\n # Backup train script, patch config\n shutil.copy2(train_path, backup_path)\n _patch_config(train_path, new_config)\n\n # Run experiment subprocess\n t0 = time.time()\n failure_reason: str | None = None\n val_dice: float | None = None\n outcome = \"failed\"\n\n try:\n proc = subprocess.run(\n [sys.executable, str(train_path)],\n cwd=str(working_dir),\n capture_output=True,\n text=True,\n timeout=budget_timeout_per_run,\n )\n elapsed = time.time() - t0\n\n # Check for failure signals\n if proc.returncode != 0:\n failure_reason = _classify_failure(proc.stderr)\n elif _has_nan(proc.stdout):\n failure_reason = \"NaN in loss\"\n else:\n # Read val_dice.json\n val_dice = _read_val_dice(val_dice_path)\n if val_dice is not None:\n if best_dice is None or val_dice > best_dice:\n outcome = \"improved\"\n best_dice = val_dice\n current_config = new_config\n else:\n outcome = \"degraded\"\n else:\n failure_reason = \"val_dice.json missing or invalid\"\n\n except subprocess.TimeoutExpired:\n elapsed = time.time() - t0\n failure_reason = f\"timeout after {budget_timeout_per_run:.0f}s\"\n\n if failure_reason is not None:\n logger.warning(\"Experiment %d failed: %s\", run_id, failure_reason)\n # Revert train script\n shutil.copy2(backup_path, train_path)\n\n results.append(\n ExperimentResult(\n run_id=run_id,\n config=copy.deepcopy(new_config),\n val_dice=val_dice,\n outcome=outcome,\n failure_reason=failure_reason,\n elapsed_seconds=elapsed if \"elapsed\" in dir() else 0.0,\n )\n )\n\n # Write run summary\n _write_summary(working_dir, results, model_family, best_dice)\n return results\n\n\n# ---------------------------------------------------------------------------\n# Internal helpers\n# ---------------------------------------------------------------------------\n\n\ndef _propose_config(\n current: dict[str, Any],\n model_family: str,\n run_id: int,\n best_dice: float | None,\n) -> dict[str, Any]:\n \"\"\"Propose a new config via Anthropic API or random grid search.\"\"\"\n api_key = os.environ.get(\"ANTHROPIC_API_KEY\")\n if api_key:\n try:\n return _propose_via_llm(current, model_family, run_id, best_dice, api_key)\n except Exception as exc:\n logger.warning(\n \"Anthropic API proposal failed (%s) — falling back to random grid\", exc\n )\n return _propose_random(current)\n\n\ndef _propose_via_llm(\n current: dict[str, Any],\n model_family: str,\n run_id: int,\n best_dice: float | None,\n api_key: str,\n) -> dict[str, Any]:\n \"\"\"Use Anthropic claude-sonnet-4-6 to propose a new config diff.\"\"\"\n import anthropic # type: ignore[import-untyped]\n\n client = anthropic.Anthropic(api_key=api_key)\n context = (\n f\"You are an ML research assistant. The current training config is:\\n\"\n f\"{json.dumps(current, indent=2)}\\n\\n\"\n f\"Model family: {model_family}\\n\"\n f\"Experiment number: {run_id}\\n\"\n f\"Best val_dice so far: {best_dice}\\n\\n\"\n f\"Propose a new configuration as a JSON object with updated hyperparameters \"\n f\"(use the same keys). Only return the JSON object, no other text.\"\n )\n message = client.messages.create(\n model=\"claude-sonnet-4-6\",\n max_tokens=512,\n messages=[{\"role\": \"user\", \"content\": context}],\n )\n raw = message.content[0].text.strip()\n # Extract JSON from the response\n start = raw.find(\"{\")\n end = raw.rfind(\"}\") + 1\n if start == -1 or end == 0:\n raise ValueError(\"LLM did not return a JSON object\")\n proposed = json.loads(raw[start:end])\n # Merge with current (keep unchanged keys)\n merged = dict(current)\n merged.update(proposed)\n return merged\n\n\ndef _propose_random(current: dict[str, Any]) -> dict[str, Any]:\n \"\"\"Random perturbation from the search grid (LLM fallback).\"\"\"\n import random\n\n proposed = dict(current)\n for key, values in _DEFAULT_SEARCH_GRID.items():\n if key in current:\n proposed[key] = random.choice(values)\n logger.info(\"Random grid proposal: %s\", proposed)\n return proposed\n\n\ndef _parse_config_comment(path: Path) -> dict[str, Any]:\n \"\"\"Extract a JSON block from a ``# CONFIG: {...}`` comment in the script.\"\"\"\n with path.open() as fh:\n for line in fh:\n if line.strip().startswith(\"# CONFIG:\"):\n try:\n return json.loads(line.split(\"# CONFIG:\", 1)[1].strip())\n except json.JSONDecodeError:\n pass\n return {}\n\n\ndef _patch_config(path: Path, config: dict[str, Any]) -> None:\n \"\"\"Replace the ``# CONFIG: {...}`` comment line in the training script.\"\"\"\n lines = path.read_text().splitlines(keepends=True)\n patched = []\n found = False\n for line in lines:\n if line.strip().startswith(\"# CONFIG:\"):\n patched.append(f\"# CONFIG: {json.dumps(config)}\\n\")\n found = True\n else:\n patched.append(line)\n if not found:\n patched.insert(0, f\"# CONFIG: {json.dumps(config)}\\n\")\n path.write_text(\"\".join(patched))\n\n\ndef _read_val_dice(path: Path) -> float | None:\n \"\"\"Read the ``val_dice`` value from a JSON file.\"\"\"\n if not path.exists():\n return None\n try:\n data = json.loads(path.read_text())\n return float(data.get(\"val_dice\", data.get(\"dice\", 0.0)))\n except (json.JSONDecodeError, TypeError, ValueError):\n return None\n\n\ndef _has_nan(text: str) -> bool:\n return \"nan\" in text.lower() or \"NaN\" in text\n\n\ndef _classify_failure(stderr: str) -> str:\n lower = stderr.lower()\n if \"out of memory\" in lower or \"outofmemoryerror\" in lower:\n return \"CUDA OOM\"\n if \"nan\" in lower:\n return \"NaN in loss\"\n return \"non-zero exit code\"\n\n\ndef _write_summary(\n working_dir: Path,\n results: list[ExperimentResult],\n model_family: str,\n best_dice: float | None,\n) -> None:\n \"\"\"Write ``run_summary.md`` to ``working_dir``.\"\"\"\n lines = [\n f\"# Autoresearch Run Summary: {model_family}\",\n \"\",\n f\"Total experiments: {len(results)}\",\n (\n f\"Best val_dice: {best_dice:.4f}\"\n if best_dice is not None\n else \"Best val_dice: N/A\"\n ),\n \"\",\n \"## Experiment Log\",\n \"\",\n \"| run_id | val_dice | outcome | failure_reason | elapsed_s |\",\n \"|--------|----------|---------|----------------|-----------|\",\n ]\n for r in results:\n dice_str = f\"{r.val_dice:.4f}\" if r.val_dice is not None else \"—\"\n lines.append(\n f\"| {r.run_id} | {dice_str} | {r.outcome} | \"\n f\"{r.failure_reason or '—'} | {r.elapsed_seconds:.1f} |\"\n )\n (working_dir / \"run_summary.md\").write_text(\"\\n\".join(lines) + \"\\n\")\n\n\ndef commit_best_model(\n best_model_path: str | Path,\n best_config_path: str | Path,\n trained_models_path: str | Path,\n model_family: str,\n val_dice: float,\n source_run_id: str = \"\",\n) -> dict[str, Any]:\n \"\"\"Version the best model with DataLad and push to OSF.\n\n Parameters\n ----------\n best_model_path : path\n Path to the ``best_model.pth`` file.\n best_config_path : path\n Path to the ``best_config.json`` file.\n trained_models_path : path\n Root of the DataLad-managed ``trained_models`` dataset.\n model_family : str\n Model family name (used as subdirectory).\n val_dice : float\n Validation Dice score of the best model.\n source_run_id : str\n Run ID string for traceability.\n\n Returns\n -------\n dict\n ``ModelVersion`` with ``path``, ``datalad_commit``, and metadata.\n \"\"\"\n import datetime\n\n import torch\n\n try:\n import datalad.api as dl # type: ignore[import-untyped]\n except ImportError as exc:\n raise ImportError(\n \"datalad is required for model versioning. \"\n \"Install it with: pip install nobrainer[versioning]\"\n ) from exc\n\n date_str = datetime.date.today().isoformat()\n dest = (\n Path(trained_models_path)\n / \"neuronets\"\n / \"autoresearch\"\n / model_family\n / date_str\n )\n dest.mkdir(parents=True, exist_ok=True)\n\n shutil.copy2(best_model_path, dest / \"model.pth\")\n shutil.copy2(best_config_path, dest / \"config.json\")\n\n # Generate model card\n import platform\n\n import monai\n import pyro\n\n card_lines = [\n f\"# Model Card: {model_family}\",\n \"\",\n \"## Architecture\",\n f\"- Model family: {model_family}\",\n \"- Framework: PyTorch\",\n \"\",\n \"## Performance\",\n f\"- val_dice: {val_dice:.4f}\",\n f\"- source_run_id: {source_run_id}\",\n \"\",\n \"## Environment\",\n f\"- Python: {platform.python_version()}\",\n f\"- PyTorch: {torch.__version__}\",\n f\"- MONAI: {monai.__version__}\",\n f\"- Pyro-ppl: {pyro.__version__}\",\n f\"- Date: {date_str}\",\n ]\n (dest / \"model_card.md\").write_text(\"\\n\".join(card_lines) + \"\\n\")\n\n commit_msg = (\n f\"autoresearch: add {model_family} model ({date_str}) val_dice={val_dice:.4f}\"\n )\n dl.save(dataset=str(trained_models_path), message=commit_msg)\n\n try:\n dl.push(dataset=str(trained_models_path), to=\"osf\")\n osf_url = \"osf://\"\n except Exception:\n osf_url = None\n\n return {\n \"path\": str(dest),\n \"datalad_commit\": commit_msg,\n \"val_dice\": val_dice,\n \"model_family\": model_family,\n \"date\": date_str,\n \"osf_url\": osf_url,\n }\n" }, { "path": "nobrainer/research/templates/.gitkeep", "content": "" }, { "path": "nobrainer/research/templates/prepare.py", "content": "\"\"\"Standard data preparation script for autoresearch.\n\nUsage\n-----\n python prepare.py --data-dir /path/to/nifti --val-fraction 0.2\n\nWrites ``data_manifest.json`` in the current directory listing\ntrain/val split paths.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport json\nfrom pathlib import Path\nimport random\n\nimport click\n\n\n@click.command()\n@click.option(\n \"--data-dir\",\n required=True,\n type=click.Path(exists=True),\n help=\"Directory containing NIfTI files (*.nii or *.nii.gz).\",\n)\n@click.option(\n \"--val-fraction\",\n default=0.2,\n type=float,\n show_default=True,\n help=\"Fraction of data for validation.\",\n)\n@click.option(\n \"--seed\",\n default=42,\n type=int,\n show_default=True,\n help=\"Random seed for train/val split.\",\n)\n@click.option(\n \"--output\",\n default=\"data_manifest.json\",\n show_default=True,\n help=\"Output manifest filename.\",\n)\ndef prepare(*, data_dir: str, val_fraction: float, seed: int, output: str) -> None:\n \"\"\"Validate NIfTI dataset and write train/val split manifest.\"\"\"\n data_path = Path(data_dir)\n niftis = sorted(list(data_path.glob(\"*.nii\")) + list(data_path.glob(\"*.nii.gz\")))\n if not niftis:\n raise click.ClickException(f\"No NIfTI files found in {data_dir}\")\n\n random.seed(seed)\n shuffled = list(niftis)\n random.shuffle(shuffled)\n\n n_val = max(1, int(len(shuffled) * val_fraction))\n val_paths = shuffled[:n_val]\n train_paths = shuffled[n_val:]\n\n manifest = {\n \"data_dir\": str(data_path.resolve()),\n \"n_total\": len(shuffled),\n \"n_train\": len(train_paths),\n \"n_val\": len(val_paths),\n \"train\": [str(p) for p in train_paths],\n \"val\": [str(p) for p in val_paths],\n }\n\n output_path = Path(output)\n output_path.write_text(json.dumps(manifest, indent=2))\n click.echo(\n f\"Manifest written to {output_path}: \"\n f\"{len(train_paths)} train, {len(val_paths)} val\"\n )\n\n\nif __name__ == \"__main__\":\n prepare()\n" }, { "path": "nobrainer/research/templates/train_bayesian_vnet.py", "content": "\"\"\"Bayesian VNet training script for autoresearch.\n\nThe autoresearch loop patches the ``# CONFIG:`` comment line below to\nupdate hyperparameters between experiments. On completion, this script\nwrites ``val_dice.json`` in the working directory.\n\nUsage\n-----\n python train_bayesian_vnet.py\n\"\"\"\n\n# CONFIG: {\"learning_rate\": 1e-4, \"batch_size\": 4, \"n_epochs\": 20, \"kl_weight\": 1e-4, \"dropout_rate\": 0.0} # noqa: E501\n\nfrom __future__ import annotations\n\nimport json\nfrom pathlib import Path\n\nfrom monai.metrics import DiceMetric\nfrom monai.utils import set_determinism\nimport torch\nimport torch.optim as optim\n\nfrom nobrainer.dataset import get_dataset\nfrom nobrainer.losses import dice as dice_loss_fn\nfrom nobrainer.losses import elbo\nfrom nobrainer.models.bayesian import BayesianVNet\nfrom nobrainer.training import get_device\n\n\ndef main() -> None:\n # ------------------------------------------------------------------ #\n # Load config from script comment (patched by autoresearch loop) #\n # ------------------------------------------------------------------ #\n script_text = Path(__file__).read_text()\n config: dict = {}\n for line in script_text.splitlines():\n if line.strip().startswith(\"# CONFIG:\"):\n config = json.loads(line.split(\"# CONFIG:\", 1)[1].strip())\n break\n\n lr: float = config.get(\"learning_rate\", 1e-4)\n batch_size: int = int(config.get(\"batch_size\", 4))\n n_epochs: int = int(config.get(\"n_epochs\", 20))\n kl_weight: float = config.get(\"kl_weight\", 1e-4)\n\n set_determinism(seed=42)\n device = get_device()\n\n # ------------------------------------------------------------------ #\n # Data loading #\n # ------------------------------------------------------------------ #\n manifest_path = Path(\"data_manifest.json\")\n if not manifest_path.exists():\n raise FileNotFoundError(\"data_manifest.json not found. Run prepare.py first.\")\n manifest = json.loads(manifest_path.read_text())\n train_images = manifest[\"train\"]\n val_images = manifest[\"val\"]\n label_suffix = \"_label\" # adjust per dataset convention\n train_labels = [\n p.replace(\".nii.gz\", f\"{label_suffix}.nii.gz\") for p in train_images\n ]\n val_labels = [p.replace(\".nii.gz\", f\"{label_suffix}.nii.gz\") for p in val_images]\n\n train_loader = get_dataset(\n image_paths=train_images,\n label_paths=train_labels,\n batch_size=batch_size,\n augment=True,\n num_workers=0,\n cache_rate=0.0,\n )\n val_loader = get_dataset(\n image_paths=val_images,\n label_paths=val_labels,\n batch_size=1,\n num_workers=0,\n cache_rate=0.0,\n )\n\n # ------------------------------------------------------------------ #\n # Model, optimiser, metrics #\n # ------------------------------------------------------------------ #\n model = BayesianVNet(n_classes=2, in_channels=1, kl_weight=kl_weight).to(device)\n optimizer = optim.Adam(model.parameters(), lr=lr)\n recon_loss_fn = dice_loss_fn(softmax=True)\n dice_metric = DiceMetric(include_background=False, reduction=\"mean\")\n\n # ------------------------------------------------------------------ #\n # Training loop #\n # ------------------------------------------------------------------ #\n import pyro\n\n for epoch in range(n_epochs):\n model.train()\n for batch in train_loader:\n imgs = batch[\"image\"].to(device)\n labels = batch[\"label\"].to(device).long()\n optimizer.zero_grad()\n with pyro.poutine.trace():\n preds = model(imgs)\n labels_onehot = torch.zeros_like(preds)\n labels_onehot.scatter_(1, labels, 1.0)\n recon = recon_loss_fn(preds, labels_onehot)\n loss = elbo(model, kl_weight, recon)\n loss.backward()\n optimizer.step()\n\n # ------------------------------------------------------------------ #\n # Validation #\n # ------------------------------------------------------------------ #\n model.eval()\n with torch.no_grad():\n for batch in val_loader:\n imgs = batch[\"image\"].to(device)\n labels = batch[\"label\"].to(device).long()\n with pyro.poutine.trace():\n preds = model(imgs)\n preds_bin = torch.argmax(preds, dim=1, keepdim=True)\n dice_metric(preds_bin, labels)\n\n val_dice = dice_metric.aggregate().item()\n dice_metric.reset()\n\n # ------------------------------------------------------------------ #\n # Write val_dice.json #\n # ------------------------------------------------------------------ #\n Path(\"val_dice.json\").write_text(json.dumps({\"val_dice\": val_dice}))\n print(f\"val_dice: {val_dice:.4f}\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "nobrainer/slurm.py", "content": "\"\"\"SLURM utilities for preemptible training with checkpoint/resume.\n\nProvides signal handling for SLURM preemption and checkpoint persistence\nso training jobs can be interrupted and resumed automatically via\n``--requeue``.\n\nUsage::\n\n from nobrainer.slurm import (\n SlurmPreemptionHandler,\n save_checkpoint,\n load_checkpoint,\n )\n\n handler = SlurmPreemptionHandler()\n start_epoch, metrics = load_checkpoint(ckpt_dir, model, optimizer)\n\n for epoch in range(start_epoch, total_epochs):\n train_one_epoch(...)\n save_checkpoint(ckpt_dir, model, optimizer, epoch, metrics)\n if handler.preempted:\n break # job will be requeued by SLURM\n\"\"\"\n\nfrom __future__ import annotations\n\nimport json\nimport logging\nimport os\nfrom pathlib import Path\nimport signal\nfrom typing import Any\n\nimport torch\n\nlogger = logging.getLogger(__name__)\n\n\nclass SlurmPreemptionHandler:\n \"\"\"Handle SLURM preemption signals for graceful checkpoint-and-exit.\n\n SLURM sends a configurable signal (default SIGUSR1 via ``--signal``)\n before killing a preempted job. This handler sets a flag so the\n training loop can checkpoint and exit cleanly. The ``--requeue``\n sbatch flag then re-submits the job.\n\n On non-SLURM systems (no ``SLURM_JOB_ID`` environment variable),\n the handler is still safe to create but will never fire.\n\n Parameters\n ----------\n sig : signal.Signals\n Signal to catch (default ``SIGUSR1``).\n \"\"\"\n\n def __init__(self, sig: signal.Signals = signal.SIGUSR1) -> None:\n self.preempted = False\n self._sig = sig\n try:\n signal.signal(sig, self._handle)\n logger.info(\"SLURM preemption handler registered (signal=%s)\", sig.name)\n except (OSError, ValueError):\n # Signal registration can fail in non-main threads or on Windows\n logger.debug(\"Could not register signal handler for %s\", sig)\n\n def _handle(self, signum: int, frame: Any) -> None:\n logger.warning(\n \"Received preemption signal %d — will checkpoint and exit\", signum\n )\n self.preempted = True\n\n @staticmethod\n def is_slurm_job() -> bool:\n \"\"\"Return True if running inside a SLURM job.\"\"\"\n return \"SLURM_JOB_ID\" in os.environ\n\n @staticmethod\n def slurm_info() -> dict[str, str]:\n \"\"\"Return a dict of useful SLURM environment variables.\"\"\"\n keys = [\n \"SLURM_JOB_ID\",\n \"SLURM_JOB_NAME\",\n \"SLURM_JOB_PARTITION\",\n \"SLURM_NODELIST\",\n \"SLURM_NTASKS\",\n \"SLURM_GPUS_ON_NODE\",\n \"SLURM_RESTART_COUNT\",\n ]\n return {k: os.environ[k] for k in keys if k in os.environ}\n\n\ndef save_checkpoint(\n checkpoint_dir: str | Path,\n model: torch.nn.Module,\n optimizer: torch.optim.Optimizer,\n epoch: int,\n metrics: dict[str, Any] | None = None,\n) -> Path:\n \"\"\"Save a resumable training checkpoint.\n\n Writes ``checkpoint.pt`` with model weights, optimizer state, epoch,\n and metrics. Also writes ``checkpoint_meta.json`` for inspection.\n\n Parameters\n ----------\n checkpoint_dir : str or Path\n Directory for checkpoint files.\n model : torch.nn.Module\n Model to checkpoint.\n optimizer : torch.optim.Optimizer\n Optimizer state to persist.\n epoch : int\n Completed epoch number (0-indexed).\n metrics : dict, optional\n Accumulated training metrics.\n\n Returns\n -------\n Path\n Path to the written ``checkpoint.pt``.\n \"\"\"\n checkpoint_dir = Path(checkpoint_dir)\n checkpoint_dir.mkdir(parents=True, exist_ok=True)\n ckpt_path = checkpoint_dir / \"checkpoint.pt\"\n\n torch.save(\n {\n \"epoch\": epoch,\n \"model_state_dict\": model.state_dict(),\n \"optimizer_state_dict\": optimizer.state_dict(),\n \"metrics\": metrics or {},\n },\n ckpt_path,\n )\n\n meta = {\n \"epoch\": epoch,\n \"best_loss\": (metrics or {}).get(\"best_loss\"),\n \"train_losses\": (metrics or {}).get(\"train_losses\", [])[-3:],\n }\n with open(checkpoint_dir / \"checkpoint_meta.json\", \"w\") as f:\n json.dump(meta, f, indent=2, default=str)\n\n logger.info(\"Checkpoint saved: epoch %d → %s\", epoch, ckpt_path)\n return ckpt_path\n\n\ndef load_checkpoint(\n checkpoint_dir: str | Path,\n model: torch.nn.Module,\n optimizer: torch.optim.Optimizer | None = None,\n) -> tuple[int, dict[str, Any]]:\n \"\"\"Load a training checkpoint and return ``(start_epoch, metrics)``.\n\n Parameters\n ----------\n checkpoint_dir : str or Path\n Directory containing ``checkpoint.pt``.\n model : torch.nn.Module\n Model to load weights into.\n optimizer : torch.optim.Optimizer or None\n Optimizer to restore. If None, only model is loaded.\n\n Returns\n -------\n start_epoch : int\n Next epoch to train (checkpoint epoch + 1).\n metrics : dict\n Accumulated metrics from previous training.\n \"\"\"\n ckpt_path = Path(checkpoint_dir) / \"checkpoint.pt\"\n if not ckpt_path.exists():\n logger.info(\"No checkpoint at %s — starting from scratch\", ckpt_path)\n return 0, {}\n\n ckpt = torch.load(ckpt_path, weights_only=False)\n model.load_state_dict(ckpt[\"model_state_dict\"])\n if optimizer is not None and \"optimizer_state_dict\" in ckpt:\n optimizer.load_state_dict(ckpt[\"optimizer_state_dict\"])\n\n start_epoch = ckpt[\"epoch\"] + 1\n metrics = ckpt.get(\"metrics\", {})\n logger.info(\n \"Resumed from checkpoint: epoch %d, best_loss=%.6f\",\n ckpt[\"epoch\"],\n metrics.get(\"best_loss\", float(\"inf\")),\n )\n return start_epoch, metrics\n" }, { "path": "nobrainer/sr-tests/__init__.py", "content": "" }, { "path": "nobrainer/sr-tests/conftest.py", "content": "\"\"\"Shared fixtures for somewhat-realistic tests.\"\"\"\n\nimport pytest\n\nfrom nobrainer.io import read_csv\nfrom nobrainer.utils import get_data\n\n\n@pytest.fixture(scope=\"session\")\ndef sample_data():\n \"\"\"Download sample brain data once per test session.\"\"\"\n csv_path = get_data()\n return read_csv(csv_path)\n\n\n@pytest.fixture(scope=\"session\")\ndef train_eval_split(sample_data):\n \"\"\"Split into 9 train + 1 eval.\"\"\"\n return sample_data[:9], sample_data[9]\n" }, { "path": "nobrainer/sr-tests/test_bayesian_uncertainty.py", "content": "\"\"\"Tests for Bayesian segmentation with uncertainty quantification.\n\nThese tests train a BayesianVNet and run MC prediction on real brain\ndata, which takes 4+ minutes on CPU. They are marked ``@pytest.mark.gpu``\nso they only run on the EC2 GPU runner (where they take <30s).\nThe same functionality is also covered by ``test_kwyk_smoke.py``.\n\"\"\"\n\nimport nibabel as nib\nimport numpy as np\nimport pytest\n\npyro = pytest.importorskip(\"pyro\") # noqa: F841\n\nfrom nobrainer.processing import Dataset, Segmentation # noqa: E402\n\n\n@pytest.mark.gpu\nclass TestBayesianUncertainty:\n \"\"\"Test Bayesian model produces uncertainty estimates.\"\"\"\n\n def test_bayesian_predict_returns_tuple(self, train_eval_split, tmp_path):\n \"\"\"Bayesian predict with n_samples returns (label, variance, entropy).\"\"\"\n train_data, eval_pair = train_eval_split\n eval_img_path = eval_pair[0]\n\n ds = (\n Dataset.from_files(\n train_data,\n block_shape=(16, 16, 16),\n n_classes=2,\n )\n .batch(2)\n .binarize()\n )\n\n seg = Segmentation(\"bayesian_vnet\")\n seg.fit(ds, epochs=2)\n result = seg.predict(eval_img_path, block_shape=(16, 16, 16), n_samples=3)\n\n # Should return a tuple of 3 NIfTI images\n assert isinstance(result, tuple)\n assert len(result) == 3\n\n label, variance, entropy = result\n assert isinstance(label, nib.Nifti1Image)\n assert isinstance(variance, nib.Nifti1Image)\n assert isinstance(entropy, nib.Nifti1Image)\n\n def test_variance_nonzero(self, train_eval_split, tmp_path):\n \"\"\"Bayesian model variance should be non-zero.\"\"\"\n train_data, eval_pair = train_eval_split\n eval_img_path = eval_pair[0]\n\n ds = (\n Dataset.from_files(\n train_data,\n block_shape=(16, 16, 16),\n n_classes=2,\n )\n .batch(2)\n .binarize()\n )\n\n seg = Segmentation(\"bayesian_vnet\")\n seg.fit(ds, epochs=2)\n _, variance, _ = seg.predict(\n eval_img_path, block_shape=(16, 16, 16), n_samples=3\n )\n\n var_data = np.asarray(variance.dataobj)\n assert np.any(var_data > 0), \"Variance should be non-zero\"\n" }, { "path": "nobrainer/sr-tests/test_brain_generation.py", "content": "\"\"\"Tests for brain generation with Progressive GAN.\"\"\"\n\nimport nibabel as nib\nimport numpy as np\nimport pytest\nimport torch\nfrom torch.utils.data import DataLoader, TensorDataset\n\npl = pytest.importorskip(\"pytorch_lightning\") # noqa: F841\n\nfrom nobrainer.processing import Generation # noqa: E402\n\n\nclass TestBrainGeneration:\n \"\"\"Test generative model training and image generation.\"\"\"\n\n def test_generate_returns_nifti_images(self, sample_data):\n \"\"\"Generation.fit().generate(2) returns 2 NIfTI images.\"\"\"\n from scipy.ndimage import zoom\n\n # Downsample real volumes to 4^3 (GAN needs small, uniform volumes)\n volumes = []\n for img_path, _ in sample_data[:4]:\n vol = np.asarray(nib.load(img_path).dataobj, dtype=np.float32)\n vmin, vmax = vol.min(), vol.max()\n if vmax > vmin:\n vol = (vol - vmin) / (vmax - vmin)\n factors = [4 / s for s in vol.shape[:3]]\n volumes.append(zoom(vol, factors, order=1))\n\n imgs = torch.from_numpy(np.stack(volumes)[:, None]) # (N, 1, 4, 4, 4)\n loader = DataLoader(TensorDataset(imgs), batch_size=2, shuffle=True)\n\n gen = Generation(\n \"progressivegan\",\n model_args={\n \"latent_size\": 16,\n \"fmap_base\": 16,\n \"fmap_max\": 16,\n \"resolution_schedule\": [4],\n \"steps_per_phase\": 100,\n },\n )\n gen.fit(loader, epochs=50)\n images = gen.generate(n_images=2)\n\n assert len(images) == 2\n for img in images:\n assert isinstance(img, nib.Nifti1Image)\n assert len(img.shape) >= 3\n" }, { "path": "nobrainer/sr-tests/test_croissant_metadata.py", "content": "\"\"\"Tests for Croissant-ML metadata generation.\"\"\"\n\nimport json\nfrom pathlib import Path\n\nfrom nobrainer.processing import Dataset, Segmentation\n\n\nclass TestCroissantMetadata:\n \"\"\"Test Croissant-ML provenance in saved models and datasets.\"\"\"\n\n def test_segmentation_save_croissant_fields(self, train_eval_split, tmp_path):\n \"\"\"Segmentation.save() produces croissant.json with provenance fields.\"\"\"\n train_data, _ = train_eval_split\n ds = (\n Dataset.from_files(\n train_data,\n block_shape=(16, 16, 16),\n n_classes=2,\n )\n .batch(2)\n .binarize()\n )\n\n seg = Segmentation(\n \"unet\",\n model_args={\"channels\": (4, 8), \"strides\": (2,)},\n )\n seg.fit(ds, epochs=2)\n\n save_dir = tmp_path / \"croissant_model\"\n seg.save(save_dir)\n\n croissant_path = save_dir / \"croissant.json\"\n assert croissant_path.exists()\n\n meta = json.loads(croissant_path.read_text())\n\n # Should have Croissant-ML context or provenance\n has_context = \"@context\" in meta\n has_provenance = \"nobrainer:provenance\" in meta\n assert (\n has_context or has_provenance\n ), \"croissant.json must have @context or nobrainer:provenance\"\n\n # Check provenance fields if present\n if has_provenance:\n prov = meta[\"nobrainer:provenance\"]\n assert \"model_architecture\" in prov\n assert \"n_classes\" in prov\n assert prov[\"model_architecture\"] == \"unet\"\n assert prov[\"n_classes\"] == 2\n\n def test_dataset_to_croissant(self, train_eval_split, tmp_path):\n \"\"\"Dataset.to_croissant() exports dataset metadata.\"\"\"\n train_data, _ = train_eval_split\n ds = Dataset.from_files(\n train_data,\n block_shape=(16, 16, 16),\n n_classes=2,\n )\n\n output_path = tmp_path / \"dataset_croissant.json\"\n result = ds.to_croissant(output_path)\n\n assert Path(result).exists()\n meta = json.loads(Path(result).read_text())\n assert \"@context\" in meta or \"name\" in meta\n" }, { "path": "nobrainer/sr-tests/test_dataset_builder.py", "content": "\"\"\"Tests for the fluent Dataset builder with real brain data.\"\"\"\n\nfrom nobrainer.processing import Dataset\n\n\nclass TestDatasetBuilder:\n \"\"\"Test Dataset.from_files() fluent API produces correct outputs.\"\"\"\n\n def test_from_files_batch_binarize_augment(self, train_eval_split):\n \"\"\"Dataset.from_files().batch(2).binarize().augment() produces correct shapes.\"\"\"\n train_data, _ = train_eval_split\n ds = (\n Dataset.from_files(\n train_data,\n block_shape=(16, 16, 16),\n n_classes=2,\n )\n .batch(2)\n .binarize()\n .augment()\n )\n loader = ds.dataloader\n batch = next(iter(loader))\n\n assert \"image\" in batch\n assert \"label\" in batch\n # batch_size=2, 1 channel, block_shape=(16,16,16)\n assert batch[\"image\"].shape[0] == 2\n assert batch[\"image\"].shape[-3:] == (16, 16, 16)\n assert batch[\"label\"].shape[0] == 2\n\n def test_split_sizes(self, train_eval_split):\n \"\"\"Dataset.split() divides data into train/eval with correct sizes.\"\"\"\n train_data, _ = train_eval_split\n ds = Dataset.from_files(\n train_data,\n block_shape=(16, 16, 16),\n n_classes=2,\n )\n ds_train, ds_eval = ds.split(eval_size=0.2)\n\n total = len(train_data)\n assert len(ds_train.data) + len(ds_eval.data) == total\n assert len(ds_eval.data) >= 1\n\n def test_streaming_mode_produces_patches(self, train_eval_split):\n \"\"\"Dataset.streaming() produces patches via PatchDataset.\"\"\"\n train_data, _ = train_eval_split\n ds = (\n Dataset.from_files(\n train_data,\n block_shape=(16, 16, 16),\n n_classes=2,\n )\n .batch(2)\n .binarize()\n .streaming(patches_per_volume=2)\n )\n loader = ds.dataloader\n batch = next(iter(loader))\n\n assert \"image\" in batch\n assert batch[\"image\"].shape[-3:] == (16, 16, 16)\n" }, { "path": "nobrainer/sr-tests/test_extract_patches.py", "content": "\"\"\"Tests for extract_patches() with various binarization modes.\"\"\"\n\nimport nibabel as nib\nimport numpy as np\nimport pytest\n\nfrom nobrainer.processing import extract_patches\n\n\nclass TestExtractPatches:\n \"\"\"Test extract_patches() on real brain data.\"\"\"\n\n @pytest.fixture()\n def volume_and_label(self, train_eval_split):\n \"\"\"Load first volume and label as numpy arrays.\"\"\"\n train_data, _ = train_eval_split\n img_path, lbl_path = train_data[0]\n vol = np.asarray(nib.load(img_path).dataobj, dtype=np.float32)\n lbl = np.asarray(nib.load(lbl_path).dataobj, dtype=np.float32)\n return vol, lbl\n\n def test_binarize_true(self, volume_and_label):\n \"\"\"binarize=True maps any non-zero label to 1.\"\"\"\n vol, lbl = volume_and_label\n patches = extract_patches(\n vol, lbl, block_shape=(16, 16, 16), n_patches=5, binarize=True\n )\n assert len(patches) == 5\n for img_patch, lbl_patch in patches:\n assert img_patch.shape == (16, 16, 16)\n assert lbl_patch.shape == (16, 16, 16)\n # Only 0 and 1 in binarized labels\n unique_vals = set(np.unique(lbl_patch))\n assert unique_vals <= {0.0, 1.0}\n\n def test_binarize_set(self, volume_and_label):\n \"\"\"binarize={17, 53} selects hippocampus labels only.\"\"\"\n vol, lbl = volume_and_label\n patches = extract_patches(\n vol, lbl, block_shape=(16, 16, 16), n_patches=5, binarize={17, 53}\n )\n for img_patch, lbl_patch in patches:\n assert img_patch.shape == (16, 16, 16)\n unique_vals = set(np.unique(lbl_patch))\n assert unique_vals <= {0.0, 1.0}\n\n def test_binarize_callable(self, volume_and_label):\n \"\"\"binarize=lambda applies custom function to label patches.\"\"\"\n vol, lbl = volume_and_label\n\n def threshold_fn(x):\n return (x >= 1000).astype(np.float32)\n\n patches = extract_patches(\n vol, lbl, block_shape=(16, 16, 16), n_patches=5, binarize=threshold_fn\n )\n for img_patch, lbl_patch in patches:\n assert img_patch.shape == (16, 16, 16)\n unique_vals = set(np.unique(lbl_patch))\n assert unique_vals <= {0.0, 1.0}\n\n def test_patch_shapes(self, volume_and_label):\n \"\"\"Patches have the requested block_shape.\"\"\"\n vol, lbl = volume_and_label\n patches = extract_patches(vol, lbl, block_shape=(16, 16, 16), n_patches=3)\n assert len(patches) == 3\n for img_patch, lbl_patch in patches:\n assert img_patch.shape == (16, 16, 16)\n assert lbl_patch.shape == (16, 16, 16)\n" }, { "path": "nobrainer/sr-tests/test_kwyk_smoke.py", "content": "\"\"\"Smoke tests for the kwyk reproduction pipeline.\n\nTests train a tiny MeshNet and Bayesian MeshNet for 1 epoch each to verify\nthe end-to-end pipeline works (loss is finite, prediction produces valid\nNIfTI output, warm-start transfers weights correctly).\n\"\"\"\n\nimport nibabel as nib\nimport numpy as np\nimport pytest\nimport torch\n\npyro = pytest.importorskip(\"pyro\")\n\nfrom nobrainer.models import get as get_model # noqa: E402\nfrom nobrainer.models.bayesian.warmstart import ( # noqa: E402\n warmstart_bayesian_from_deterministic,\n)\nfrom nobrainer.processing import Dataset, Segmentation # noqa: E402\nfrom nobrainer.training import get_device # noqa: E402\n\n# ---------------------------------------------------------------------------\n# Shared constants for tiny model\n# ---------------------------------------------------------------------------\nFILTERS = 16\nBLOCK_SHAPE = (16, 16, 16)\nN_CLASSES = 2\nBATCH_SIZE = 2\nMODEL_ARGS = {\n \"n_classes\": N_CLASSES,\n \"filters\": FILTERS,\n \"receptive_field\": 37,\n \"dropout_rate\": 0.25,\n}\n\n\ndef _build_dataset(sample_data):\n \"\"\"Build a small binarized Dataset from sample_data fixture.\"\"\"\n # Use first 5 volumes\n pairs = sample_data[:5]\n ds = (\n Dataset.from_files(pairs, block_shape=BLOCK_SHAPE, n_classes=N_CLASSES)\n .batch(BATCH_SIZE)\n .binarize()\n )\n return ds\n\n\ndef _plot_learning_curve(losses, output_path):\n \"\"\"Save a simple learning curve figure.\"\"\"\n import matplotlib\n\n matplotlib.use(\"Agg\")\n import matplotlib.pyplot as plt # noqa: E402\n\n fig, ax = plt.subplots(figsize=(6, 4))\n ax.plot(range(1, len(losses) + 1), losses, \"b-o\", markersize=4)\n ax.set_xlabel(\"Step\")\n ax.set_ylabel(\"Loss\")\n ax.set_title(\"Smoke Test Learning Curve\")\n ax.grid(True, alpha=0.3)\n fig.tight_layout()\n fig.savefig(output_path, dpi=72, bbox_inches=\"tight\")\n plt.close(fig)\n\n\n@pytest.mark.gpu\nclass TestKwykSmoke:\n \"\"\"Smoke tests for the kwyk reproduction pipeline.\"\"\"\n\n def test_deterministic_meshnet_train(self, sample_data, tmp_path):\n \"\"\"Train deterministic MeshNet for 1 epoch; assert loss is finite.\"\"\"\n ds = _build_dataset(sample_data)\n\n seg = Segmentation(\n base_model=\"meshnet\",\n model_args={k: v for k, v in MODEL_ARGS.items() if k != \"n_classes\"},\n )\n\n # Collect losses via callback\n losses = []\n\n def _on_epoch(epoch, logs, model):\n losses.append(logs[\"loss\"] if isinstance(logs, dict) else logs)\n\n seg.fit(ds, epochs=1, callbacks=[_on_epoch])\n\n assert len(losses) >= 1, \"Expected at least 1 epoch of training\"\n for loss_val in losses:\n assert np.isfinite(loss_val), f\"Loss is not finite: {loss_val}\"\n\n # Save learning curve\n _plot_learning_curve(losses, tmp_path / \"det_learning_curve.png\")\n assert (tmp_path / \"det_learning_curve.png\").exists()\n\n def test_bayesian_warmstart_train(self, sample_data, tmp_path):\n \"\"\"Warm-start BayesianMeshNet from deterministic, train 1 epoch.\"\"\"\n ds = _build_dataset(sample_data)\n\n # First train a deterministic model\n det_model = get_model(\"meshnet\")(**MODEL_ARGS)\n device = get_device()\n det_model = det_model.to(device)\n det_model.train()\n\n ce_loss = torch.nn.CrossEntropyLoss()\n optimizer = torch.optim.Adam(det_model.parameters(), lr=1e-3)\n\n # Quick 1-epoch train of deterministic model\n loader = ds.dataloader\n for batch in loader:\n if isinstance(batch, (list, tuple)):\n images = batch[0].to(device)\n labels = batch[1].to(device)\n else:\n images = batch[\"image\"].to(device)\n labels = batch[\"label\"].to(device)\n\n if labels.ndim == images.ndim and labels.shape[1] == 1:\n labels = labels.squeeze(1)\n if labels.dtype in (torch.float32, torch.float64):\n labels = labels.long()\n\n optimizer.zero_grad()\n pred = det_model(images)\n loss = ce_loss(pred, labels)\n loss.backward()\n optimizer.step()\n break # Just one batch for speed\n\n # Build Bayesian model and warm-start (on CPU to avoid device issues)\n det_model_cpu = det_model.cpu()\n bayes_model = get_model(\"bayesian_meshnet\")(**MODEL_ARGS)\n n_transferred = warmstart_bayesian_from_deterministic(\n bayes_model, det_model_cpu, initial_rho=-3.0\n )\n assert n_transferred > 0, \"Expected at least 1 layer transferred\"\n\n # Train Bayesian for 1 epoch\n from nobrainer.models.bayesian.utils import accumulate_kl\n\n # Pyro's param store can cache unconstrained tensors on CPU even\n # after .to(device). Clear and re-register to ensure device consistency.\n pyro.clear_param_store()\n bayes_model = bayes_model.to(device)\n bayes_model.train()\n optimizer_b = torch.optim.Adam(bayes_model.parameters(), lr=1e-3)\n\n losses = []\n for batch in loader:\n if isinstance(batch, (list, tuple)):\n images = batch[0].to(device)\n labels = batch[1].to(device)\n else:\n images = batch[\"image\"].to(device)\n labels = batch[\"label\"].to(device)\n\n if labels.ndim == images.ndim and labels.shape[1] == 1:\n labels = labels.squeeze(1)\n if labels.dtype in (torch.float32, torch.float64):\n labels = labels.long()\n\n optimizer_b.zero_grad()\n pred = bayes_model(images)\n loss = ce_loss(pred, labels) + accumulate_kl(bayes_model)\n loss.backward()\n optimizer_b.step()\n losses.append(loss.item())\n break # Just one batch for speed\n\n assert len(losses) >= 1\n for loss_val in losses:\n assert np.isfinite(loss_val), f\"Bayesian loss not finite: {loss_val}\"\n\n # Save learning curve\n _plot_learning_curve(losses, tmp_path / \"bayes_learning_curve.png\")\n assert (tmp_path / \"bayes_learning_curve.png\").exists()\n\n def test_predict_output(self, sample_data, tmp_path):\n \"\"\"Predict on 1 volume; assert NIfTI output with matching shape.\"\"\"\n ds = _build_dataset(sample_data)\n\n seg = Segmentation(\n base_model=\"meshnet\",\n model_args={k: v for k, v in MODEL_ARGS.items() if k != \"n_classes\"},\n )\n seg.fit(ds, epochs=1)\n\n # Predict on first volume\n eval_img_path = sample_data[0][0]\n eval_lbl_path = sample_data[0][1]\n\n result = seg.predict(eval_img_path, block_shape=BLOCK_SHAPE)\n\n # Check output is NIfTI\n assert isinstance(\n result, nib.Nifti1Image\n ), f\"Expected Nifti1Image, got {type(result)}\"\n\n # Check shape matches input spatial dims\n input_img = nib.load(eval_img_path)\n input_shape = input_img.shape[:3]\n result_shape = result.shape[:3]\n assert (\n result_shape == input_shape\n ), f\"Shape mismatch: input={input_shape}, output={result_shape}\"\n\n # Compute Dice for informational purposes — a 1-epoch model\n # may produce all-zero predictions, so we don't require Dice > 0.\n gt_arr = np.asarray(nib.load(eval_lbl_path).dataobj, dtype=np.float32)\n gt_binary = (gt_arr > 0).astype(np.float32)\n\n pred_arr = np.asarray(result.dataobj, dtype=np.float32)\n pred_binary = (pred_arr > 0).astype(np.float32)\n\n intersection = np.logical_and(pred_binary, gt_binary).sum()\n total = pred_binary.sum() + gt_binary.sum()\n if total > 0:\n dice = float(2.0 * intersection / total)\n else:\n dice = 1.0\n\n # Dice >= 0 is always true; we just verify the computation doesn't crash.\n # With more epochs, Dice should improve — this is a smoke test only.\n assert dice >= 0, f\"Expected Dice >= 0, got {dice}\"\n\n # Save learning curve figure\n _plot_learning_curve([0.5], tmp_path / \"predict_learning_curve.png\")\n assert (tmp_path / \"predict_learning_curve.png\").exists()\n" }, { "path": "nobrainer/sr-tests/test_raw_pytorch_api.py", "content": "\"\"\"Tests for the raw PyTorch API without the estimator layer.\"\"\"\n\nimport nibabel as nib\nimport torch\n\nimport nobrainer.models\nfrom nobrainer.prediction import predict\nfrom nobrainer.training import fit as training_fit\n\n\nclass TestRawPyTorchAPI:\n \"\"\"Test using raw nobrainer modules directly (no estimator).\"\"\"\n\n def test_raw_train_predict_cycle(self, train_eval_split, tmp_path):\n \"\"\"Train with nobrainer.training.fit, predict with nobrainer.prediction.predict.\"\"\"\n train_data, eval_pair = train_eval_split\n eval_img_path = eval_pair[0]\n\n # Build model directly\n model_factory = nobrainer.models.get(\"unet\")\n model = model_factory(n_classes=2, channels=(4, 8), strides=(2,))\n\n # Build dataset directly\n from nobrainer.dataset import get_dataset\n\n image_paths = [pair[0] for pair in train_data]\n label_paths = [pair[1] for pair in train_data]\n\n loader = get_dataset(\n image_paths=image_paths,\n label_paths=label_paths,\n block_shape=(16, 16, 16),\n batch_size=2,\n binarize_labels=True,\n )\n\n # Train\n optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n criterion = torch.nn.CrossEntropyLoss()\n\n result = training_fit(\n model=model,\n loader=loader,\n criterion=criterion,\n optimizer=optimizer,\n max_epochs=2,\n gpus=0,\n )\n\n assert \"history\" in result\n assert len(result[\"history\"]) == 2\n\n # Predict\n model.eval()\n prediction = predict(\n inputs=eval_img_path,\n model=model,\n block_shape=(16, 16, 16),\n )\n\n assert isinstance(prediction, nib.Nifti1Image)\n assert len(prediction.shape) >= 3\n" }, { "path": "nobrainer/sr-tests/test_segmentation_estimator.py", "content": "\"\"\"Tests for the Segmentation estimator with real brain data.\"\"\"\n\nimport json\n\nimport nibabel as nib\n\nfrom nobrainer.processing import Dataset, Segmentation\n\n\nclass TestSegmentationEstimator:\n \"\"\"Test Segmentation estimator fit/predict/save/load cycle.\"\"\"\n\n def test_fit_predict_returns_nifti(self, train_eval_split, tmp_path):\n \"\"\"Segmentation.fit().predict() returns a NIfTI image.\"\"\"\n train_data, eval_pair = train_eval_split\n eval_img_path = eval_pair[0]\n\n ds = (\n Dataset.from_files(\n train_data,\n block_shape=(16, 16, 16),\n n_classes=2,\n )\n .batch(2)\n .binarize()\n )\n\n seg = Segmentation(\n \"unet\",\n model_args={\"channels\": (4, 8), \"strides\": (2,)},\n )\n seg.fit(ds, epochs=2)\n result = seg.predict(eval_img_path, block_shape=(16, 16, 16))\n\n assert isinstance(result, nib.Nifti1Image)\n assert len(result.shape) >= 3\n\n def test_save_creates_croissant(self, train_eval_split, tmp_path):\n \"\"\"Segmentation.save() creates model.pth and croissant.json.\"\"\"\n train_data, _ = train_eval_split\n ds = (\n Dataset.from_files(\n train_data,\n block_shape=(16, 16, 16),\n n_classes=2,\n )\n .batch(2)\n .binarize()\n )\n\n seg = Segmentation(\n \"unet\",\n model_args={\"channels\": (4, 8), \"strides\": (2,)},\n )\n seg.fit(ds, epochs=2)\n\n save_dir = tmp_path / \"saved_model\"\n seg.save(save_dir)\n\n assert (save_dir / \"model.pth\").exists()\n assert (save_dir / \"croissant.json\").exists()\n\n meta = json.loads((save_dir / \"croissant.json\").read_text())\n assert \"@context\" in meta or \"nobrainer:provenance\" in meta\n\n def test_load_roundtrip(self, train_eval_split, tmp_path):\n \"\"\"Segmentation.save() then Segmentation.load() restores the model.\"\"\"\n train_data, eval_pair = train_eval_split\n eval_img_path = eval_pair[0]\n\n ds = (\n Dataset.from_files(\n train_data,\n block_shape=(16, 16, 16),\n n_classes=2,\n )\n .batch(2)\n .binarize()\n )\n\n seg = Segmentation(\n \"unet\",\n model_args={\"channels\": (4, 8), \"strides\": (2,)},\n )\n seg.fit(ds, epochs=2)\n\n save_dir = tmp_path / \"roundtrip_model\"\n seg.save(save_dir)\n\n loaded = Segmentation.load(save_dir)\n result = loaded.predict(eval_img_path, block_shape=(16, 16, 16))\n\n assert isinstance(result, nib.Nifti1Image)\n" }, { "path": "nobrainer/sr-tests/test_synthseg_brain.py", "content": "\"\"\"SR-test: SynthSeg generation from real aparc+aseg label maps.\n\nTests that the enhanced SynthSeg generator produces realistic synthetic\nimages from actual FreeSurfer parcellation data.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport numpy as np\nimport torch\n\n\nclass TestSynthSegBrain:\n \"\"\"SynthSeg with real brain data.\"\"\"\n\n def test_generate_from_sample_data(self, sample_data):\n \"\"\"Generate synthetic image from real aparc+aseg label map.\"\"\"\n from nobrainer.augmentation.synthseg import SynthSegGenerator\n\n # sample_data is list of (image, label) tuples\n label_paths = [p[1] for p in sample_data[:2]]\n\n gen = SynthSegGenerator(\n label_paths,\n n_samples_per_map=2,\n elastic_std=2.0, # mild deformation for speed\n rotation_range=10.0,\n randomize_resolution=False, # skip for speed\n )\n\n sample = gen[0]\n assert sample[\"image\"].shape[0] == 1 # channel dim\n assert sample[\"label\"].shape[0] == 1\n assert sample[\"image\"].dtype == torch.float32\n assert sample[\"label\"].dtype == torch.int64\n\n # Image should have non-zero values in brain region\n img = sample[\"image\"][0].numpy()\n lbl = sample[\"label\"][0].numpy()\n brain_mask = lbl > 0\n assert brain_mask.sum() > 0\n assert img[brain_mask].std() > 0 # not constant\n\n def test_two_samples_differ(self, sample_data):\n \"\"\"Two samples from same label map should differ.\"\"\"\n from nobrainer.augmentation.synthseg import SynthSegGenerator\n\n label_paths = [p[1] for p in sample_data[:1]]\n gen = SynthSegGenerator(\n label_paths,\n n_samples_per_map=2,\n elastic_std=0,\n rotation_range=0,\n flipping=False,\n randomize_resolution=False,\n )\n\n s1 = gen[0][\"image\"]\n s2 = gen[1][\"image\"]\n assert not torch.allclose(s1, s2)\n\n def test_label_structure_preserved(self, sample_data):\n \"\"\"Spatial augmentation should preserve label topology.\"\"\"\n from nobrainer.augmentation.synthseg import SynthSegGenerator\n\n label_paths = [p[1] for p in sample_data[:1]]\n gen = SynthSegGenerator(\n label_paths,\n n_samples_per_map=1,\n elastic_std=2.0,\n rotation_range=5.0,\n )\n\n sample = gen[0]\n lbl = sample[\"label\"][0].numpy()\n\n # Should still have brain structure (not all zeros or all one label)\n unique = np.unique(lbl)\n assert len(unique) > 2 # at least background + 2 regions\n" }, { "path": "nobrainer/sr-tests/test_zarr_conversion.py", "content": "\"\"\"Tests for NIfTI-to-Zarr and Zarr-to-NIfTI conversion.\"\"\"\n\nfrom pathlib import Path\n\nimport nibabel as nib\nimport numpy as np\nimport pytest\n\nzarr = pytest.importorskip(\"zarr\") # noqa: F841\n\nfrom nobrainer.io import nifti_to_zarr, zarr_to_nifti # noqa: E402\n\n\ndef _mgz_to_nifti(mgz_path: str, output_dir: Path) -> Path:\n \"\"\"Convert .mgz to .nii.gz (niizarr doesn't support MGH).\"\"\"\n img = nib.load(mgz_path)\n out = output_dir / (Path(mgz_path).stem + \".nii.gz\")\n nib.save(nib.Nifti1Image(np.asarray(img.dataobj), img.affine), str(out))\n return out\n\n\nclass TestZarrConversion:\n \"\"\"Test Zarr round-trip conversion on real brain data.\"\"\"\n\n def test_nifti_to_zarr(self, train_eval_split, tmp_path):\n \"\"\"nifti_to_zarr() creates a valid Zarr store from a real volume.\"\"\"\n train_data, _ = train_eval_split\n mgz_path = train_data[0][0]\n\n nii_path = _mgz_to_nifti(mgz_path, tmp_path)\n zarr_path = tmp_path / \"brain.zarr\"\n\n result = nifti_to_zarr(nii_path, zarr_path, chunk_shape=(16, 16, 16), levels=1)\n assert Path(result).exists()\n\n import zarr as zarr_mod\n\n store = zarr_mod.open_group(str(zarr_path), mode=\"r\")\n assert \"0\" in store\n arr = np.asarray(store[\"0\"])\n assert arr.ndim == 3\n\n def test_zarr_to_nifti_roundtrip(self, train_eval_split, tmp_path):\n \"\"\"zarr_to_nifti() round-trips back to NIfTI with matching shape.\"\"\"\n train_data, _ = train_eval_split\n mgz_path = train_data[0][0]\n\n nii_path = _mgz_to_nifti(mgz_path, tmp_path)\n zarr_path = tmp_path / \"roundtrip.zarr\"\n nifti_to_zarr(nii_path, zarr_path, chunk_shape=(16, 16, 16), levels=1)\n\n roundtrip_path = tmp_path / \"roundtrip.nii.gz\"\n zarr_to_nifti(zarr_path, roundtrip_path)\n\n original = nib.load(str(nii_path))\n roundtrip = nib.load(str(roundtrip_path))\n assert original.shape == roundtrip.shape\n # Value range should be preserved (exact match may differ due to\n # niizarr orientation transforms)\n orig_data = np.asarray(original.dataobj, dtype=np.float32)\n rt_data = np.asarray(roundtrip.dataobj, dtype=np.float32)\n assert abs(orig_data.mean() - rt_data.mean()) < orig_data.std() * 0.5\n\n def test_multi_resolution_pyramid(self, train_eval_split, tmp_path):\n \"\"\"nifti_to_zarr(levels=3) creates a multi-resolution pyramid.\"\"\"\n train_data, _ = train_eval_split\n mgz_path = train_data[0][0]\n\n nii_path = _mgz_to_nifti(mgz_path, tmp_path)\n zarr_path = tmp_path / \"pyramid.zarr\"\n nifti_to_zarr(nii_path, zarr_path, chunk_shape=(16, 16, 16), levels=3)\n\n import zarr as zarr_mod\n\n store = zarr_mod.open_group(str(zarr_path), mode=\"r\")\n # Should have levels 0, 1, 2\n assert \"0\" in store\n assert \"1\" in store\n assert \"2\" in store\n\n shape_0 = np.asarray(store[\"0\"]).shape\n shape_1 = np.asarray(store[\"1\"]).shape\n shape_2 = np.asarray(store[\"2\"]).shape\n\n # Each level should be roughly half the previous\n for dim in range(3):\n assert shape_1[dim] <= shape_0[dim]\n assert shape_2[dim] <= shape_1[dim]\n" }, { "path": "nobrainer/sr-tests/test_zarr_pipeline.py", "content": "\"\"\"SR-test: end-to-end Zarr pipeline with real brain data.\n\nConverts sample brain data to Zarr, creates partition, builds\nDataset.from_zarr(), and verifies the DataLoader yields correct patches.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport numpy as np\n\nfrom nobrainer.processing import Dataset\n\n\nclass TestZarrPipeline:\n \"\"\"End-to-end Zarr store → partition → Dataset → DataLoader.\"\"\"\n\n def test_zarr_store_from_sample_data(self, sample_data, tmp_path):\n \"\"\"Convert sample data to Zarr, create partition, load via Dataset.\"\"\"\n from nobrainer.datasets.zarr_store import (\n create_partition,\n create_zarr_store,\n store_info,\n )\n\n # Use first 5 subjects\n pairs = sample_data[:5]\n\n # Create Zarr store (auto-conform since shapes may differ)\n store_path = create_zarr_store(\n pairs,\n tmp_path / \"brain.zarr\",\n conform=True,\n )\n\n # Verify store metadata\n info = store_info(store_path)\n assert info[\"n_subjects\"] == 5\n assert info[\"layout\"] == \"stacked\"\n assert info[\"conformed\"] is True\n\n # Create partition\n part_path = create_partition(store_path, ratios=(60, 20, 20), seed=42)\n\n # Build Dataset from Zarr with partition\n ds = Dataset.from_zarr(\n store_path,\n block_shape=(16, 16, 16),\n n_classes=2,\n partition=\"train\",\n partition_path=part_path,\n )\n\n # Verify data list is filtered\n assert len(ds.data) == 3 # 60% of 5 = 3\n\n # Verify Zarr metadata in data entries\n assert \"_zarr_index\" in ds.data[0]\n assert \"_subject_id\" in ds.data[0]\n\n def test_zarr_store_roundtrip(self, sample_data, tmp_path):\n \"\"\"Verify Zarr store preserves data fidelity.\"\"\"\n import zarr\n\n from nobrainer.datasets.zarr_store import create_zarr_store\n\n pairs = sample_data[:2]\n store_path = create_zarr_store(pairs, tmp_path / \"brain.zarr\", conform=True)\n\n store = zarr.open_group(str(store_path), mode=\"r\")\n assert store[\"images\"].shape[0] == 2\n assert store[\"labels\"].shape[0] == 2\n\n # Images should be float32, labels int32\n assert store[\"images\"].dtype == np.float32\n assert store[\"labels\"].dtype == np.int32\n" }, { "path": "nobrainer/tests/__init__.py", "content": "" }, { "path": "nobrainer/tests/contract/__init__.py", "content": "" }, { "path": "nobrainer/tests/contract/test_cli.py", "content": "\"\"\"CLI contract tests for nobrainer commands.\n\nVerifies that all CLI commands advertised in contracts/nobrainer-pytorch-api.md\nare present, have the expected options, and exit with code 0 on --help.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport subprocess\nimport sys\n\n\ndef _help(cmd: list[str]) -> str:\n \"\"\"Run `nobrainer --help` and return stdout.\"\"\"\n result = subprocess.run(\n [sys.executable, \"-m\", \"nobrainer.cli.main\"] + cmd + [\"--help\"],\n capture_output=True,\n text=True,\n )\n assert (\n result.returncode == 0\n ), f\"'{' '.join(cmd)} --help' exited {result.returncode}:\\n{result.stderr}\"\n return result.stdout\n\n\nclass TestPredictCommand:\n def test_predict_help_exits_zero(self):\n _help([\"predict\"])\n\n def test_predict_has_model_option(self):\n out = _help([\"predict\"])\n assert \"--model\" in out or \"-m\" in out\n\n def test_predict_has_model_type_option(self):\n out = _help([\"predict\"])\n assert \"--model-type\" in out\n\n def test_predict_has_n_classes_option(self):\n out = _help([\"predict\"])\n assert \"--n-classes\" in out\n\n def test_predict_has_device_option(self):\n out = _help([\"predict\"])\n assert \"--device\" in out\n\n def test_predict_has_n_samples_option(self):\n out = _help([\"predict\"])\n assert \"--n-samples\" in out\n\n\nclass TestGenerateCommand:\n def test_generate_help_exits_zero(self):\n _help([\"generate\"])\n\n def test_generate_has_model_option(self):\n out = _help([\"generate\"])\n assert \"--model\" in out or \"-m\" in out\n\n def test_generate_has_model_type_option(self):\n out = _help([\"generate\"])\n assert \"--model-type\" in out\n\n def test_generate_has_n_samples_option(self):\n out = _help([\"generate\"])\n assert \"--n-samples\" in out\n\n def test_generate_has_latent_size_option(self):\n out = _help([\"generate\"])\n assert \"--latent-size\" in out\n\n\nclass TestConvertTfrecordsCommand:\n def test_convert_tfrecords_help_exits_zero(self):\n _help([\"convert-tfrecords\"])\n\n def test_convert_tfrecords_has_input_option(self):\n out = _help([\"convert-tfrecords\"])\n assert \"--input\" in out or \"-i\" in out\n\n def test_convert_tfrecords_has_output_dir_option(self):\n out = _help([\"convert-tfrecords\"])\n assert \"--output-dir\" in out\n\n\nclass TestResearchCommand:\n def test_research_help_exits_zero(self):\n _help([\"research\"])\n\n def test_research_has_working_dir_option(self):\n out = _help([\"research\"])\n assert \"--working-dir\" in out\n\n def test_research_has_max_experiments_option(self):\n out = _help([\"research\"])\n assert \"--max-experiments\" in out\n\n def test_research_has_budget_hours_option(self):\n out = _help([\"research\"])\n assert \"--budget-hours\" in out\n\n\nclass TestCommitCommand:\n def test_commit_help_exits_zero(self):\n _help([\"commit\"])\n\n def test_commit_has_model_path_option(self):\n out = _help([\"commit\"])\n assert \"--model-path\" in out\n\n def test_commit_has_config_path_option(self):\n out = _help([\"commit\"])\n assert \"--config-path\" in out\n\n def test_commit_has_val_dice_option(self):\n out = _help([\"commit\"])\n assert \"--val-dice\" in out\n\n\nclass TestInfoCommand:\n def test_info_help_exits_zero(self):\n _help([\"info\"])\n" }, { "path": "nobrainer/tests/gpu/__init__.py", "content": "" }, { "path": "nobrainer/tests/gpu/test_bayesian_e2e.py", "content": "\"\"\"GPU end-to-end test: Bayesian VNet with uncertainty quantification.\n\nT045 — US2 acceptance scenario: predict_with_uncertainty() produces\nlabel, variance, and entropy maps. Variance and entropy are non-zero.\nBayesian model trained via overfit on synthetic sphere data achieves\nDice >= 0.90 (lower than deterministic due to stochastic inference).\n\"\"\"\n\nfrom __future__ import annotations\n\nimport numpy as np\nimport pytest\nimport torch\nimport torch.nn as nn\n\nfrom nobrainer.models.bayesian import BayesianVNet\nfrom nobrainer.prediction import predict_with_uncertainty\n\n\ndef _make_sphere_volume(shape=(64, 64, 64), radius=20):\n \"\"\"Create a synthetic volume with a centered sphere as the label.\"\"\"\n vol = np.random.rand(*shape).astype(np.float32) * 0.3\n label = np.zeros(shape, dtype=np.float32)\n center = np.array(shape) / 2\n coords = np.mgrid[: shape[0], : shape[1], : shape[2]]\n dist = np.sqrt(sum((c - ctr) ** 2 for c, ctr in zip(coords, center)))\n mask = dist < radius\n label[mask] = 1.0\n vol[mask] += 0.7\n return vol, label\n\n\n@pytest.mark.gpu\nclass TestBayesianEndToEnd:\n def test_bayesian_vnet_overfit_with_uncertainty(self):\n \"\"\"Train 2-class BayesianVNet, run MC inference, check Dice and uncertainty.\"\"\"\n device = torch.device(\"cuda\")\n torch.manual_seed(42)\n\n vol, label = _make_sphere_volume(shape=(64, 64, 64), radius=20)\n x = torch.from_numpy(vol[None, None]).to(device)\n label_long = torch.from_numpy(label).long().to(device)\n\n # Use n_classes=2 so softmax produces meaningful probabilities\n model = BayesianVNet(\n in_channels=1, n_classes=2, prior_type=\"standard_normal\"\n ).to(device)\n optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n criterion = nn.CrossEntropyLoss()\n\n # Overfit\n model.train()\n for _ in range(200):\n optimizer.zero_grad()\n pred = model(x)\n loss = criterion(pred, label_long.unsqueeze(0))\n loss.backward()\n optimizer.step()\n\n # Run MC inference with uncertainty\n label_img, var_img, entropy_img = predict_with_uncertainty(\n inputs=vol,\n model=model,\n n_samples=10,\n block_shape=(32, 32, 32),\n batch_size=8,\n device=\"cuda\",\n )\n\n # Check shapes\n assert label_img.shape == (64, 64, 64)\n assert var_img.shape == (64, 64, 64)\n assert entropy_img.shape == (64, 64, 64)\n\n # Variance and entropy should be non-zero (stochastic model)\n var_data = np.asarray(var_img.dataobj)\n entropy_data = np.asarray(entropy_img.dataobj)\n assert var_data.sum() > 0, \"Variance map is all zeros\"\n assert entropy_data.sum() > 0, \"Entropy map is all zeros\"\n\n # Dice check for class 1 (>= 0.90, relaxed for Bayesian stochasticity)\n pred_arr = np.asarray(label_img.dataobj)\n pred_bin = (pred_arr == 1).astype(np.float32)\n intersection = (pred_bin * label).sum()\n dice = 2 * intersection / (pred_bin.sum() + label.sum() + 1e-8)\n assert dice >= 0.90, f\"Bayesian Dice {dice:.4f} < 0.90 threshold\"\n" }, { "path": "nobrainer/tests/gpu/test_gan_e2e.py", "content": "\"\"\"GPU end-to-end test: ProgressiveGAN training.\n\nT054 — US3 acceptance scenario: ProgressiveGAN completes extended training\non synthetic 3D volumes without NaN in losses. Generated output has correct\nshape and non-trivial intensity distribution.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport numpy as np\nimport pytest\nimport pytorch_lightning as pl\nimport torch\nfrom torch.utils.data import DataLoader, TensorDataset\n\nfrom nobrainer.models.generative import ProgressiveGAN\n\n\ndef _make_loader(n_samples=64, spatial=4, batch_size=4):\n \"\"\"Create a DataLoader with enough data for extended training.\"\"\"\n imgs = torch.randn(n_samples, 1, spatial, spatial, spatial)\n return DataLoader(TensorDataset(imgs), batch_size=batch_size, shuffle=True)\n\n\n@pytest.mark.gpu\nclass TestProgressiveGANEndToEnd:\n def test_extended_training_no_nan(self):\n \"\"\"Train ProgressiveGAN for many steps; verify no NaN in discriminator.\"\"\"\n torch.manual_seed(42)\n\n loader = _make_loader(n_samples=64, spatial=4, batch_size=4)\n\n model = ProgressiveGAN(\n latent_size=32,\n fmap_base=32,\n fmap_max=32,\n resolution_schedule=[4],\n steps_per_phase=2000,\n )\n\n trainer = pl.Trainer(\n max_steps=500,\n accelerator=\"gpu\",\n devices=1,\n enable_checkpointing=False,\n logger=False,\n enable_progress_bar=False,\n )\n trainer.fit(model, loader)\n\n # Verify discriminator outputs are finite after training\n model.eval()\n with torch.no_grad():\n x_real = next(iter(loader))[0].to(model.device)\n z = torch.randn(x_real.size(0), 32, device=model.device)\n x_fake = model.generator(z)\n d_real = model.discriminator(x_real)\n d_fake = model.discriminator(x_fake)\n\n assert torch.isfinite(d_real).all(), \"d_real contains NaN/Inf\"\n assert torch.isfinite(d_fake).all(), \"d_fake contains NaN/Inf\"\n assert not torch.isnan(x_fake).any(), \"Generated volumes contain NaN\"\n\n def test_generated_output_shape(self):\n \"\"\"After training, generated volumes have correct shape.\"\"\"\n torch.manual_seed(42)\n\n loader = _make_loader(n_samples=32, spatial=4, batch_size=4)\n\n model = ProgressiveGAN(\n latent_size=32,\n fmap_base=32,\n fmap_max=32,\n resolution_schedule=[4],\n steps_per_phase=500,\n )\n\n trainer = pl.Trainer(\n max_steps=100,\n accelerator=\"gpu\",\n devices=1,\n enable_checkpointing=False,\n logger=False,\n enable_progress_bar=False,\n )\n trainer.fit(model, loader)\n\n model.eval()\n model.generator.current_level = 0\n model.generator.alpha = 1.0\n with torch.no_grad():\n z = torch.randn(4, 32, device=model.device)\n generated = model.generator(z)\n\n # Check shape: (4, 1, 4, 4, 4)\n assert generated.shape == (\n 4,\n 1,\n 4,\n 4,\n 4,\n ), f\"Expected (4, 1, 4, 4, 4), got {generated.shape}\"\n assert not np.isnan(generated.cpu().numpy()).any(), \"NaN in generated\"\n" }, { "path": "nobrainer/tests/gpu/test_multi_gpu.py", "content": "\"\"\"GPU integration test: multi-GPU training and inference.\n\nT035 — US4: requires 2+ GPUs. Tests DDP training speedup and\nmulti-GPU predict() correctness.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport time\n\nimport numpy as np\nimport pytest\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader, TensorDataset\n\nfrom nobrainer.models.segmentation import unet\nfrom nobrainer.prediction import predict\nfrom nobrainer.training import fit\n\n\n@pytest.mark.gpu\n@pytest.mark.skipif(\n torch.cuda.device_count() < 2,\n reason=\"Requires 2+ GPUs for multi-GPU tests\",\n)\nclass TestMultiGPU:\n def test_ddp_fit_loss_decreases(self):\n \"\"\"fit() with gpus=2 produces decreasing loss.\"\"\"\n torch.manual_seed(42)\n x = torch.randn(16, 1, 16, 16, 16)\n y = torch.randint(0, 2, (16, 16, 16, 16))\n loader = DataLoader(TensorDataset(x, y), batch_size=4)\n\n model = nn.Sequential(\n nn.Conv3d(1, 8, 3, padding=1), nn.ReLU(), nn.Conv3d(8, 2, 1)\n )\n\n losses = []\n\n def track(epoch, loss, model):\n losses.append(loss)\n\n result = fit(\n model,\n loader,\n nn.CrossEntropyLoss(),\n torch.optim.Adam(model.parameters(), lr=1e-2),\n max_epochs=10,\n gpus=2,\n callbacks=[track],\n )\n final_loss = result[\"history\"][-1][\"loss\"]\n assert final_loss < losses[0]\n\n def test_multi_gpu_predict_matches_single(self):\n \"\"\"Multi-GPU predict() output matches single-GPU result.\"\"\"\n torch.manual_seed(42)\n vol = np.random.rand(32, 32, 32).astype(np.float32)\n model = unet(n_classes=2)\n\n # Single GPU\n result_single = predict(\n inputs=vol,\n model=model,\n block_shape=(16, 16, 16),\n device=\"cuda:0\",\n )\n\n # Multi GPU (auto-distributes)\n result_multi = predict(\n inputs=vol,\n model=model,\n block_shape=(16, 16, 16),\n device=\"cuda\",\n )\n\n single_arr = np.asarray(result_single.dataobj)\n multi_arr = np.asarray(result_multi.dataobj)\n assert np.array_equal(single_arr, multi_arr)\n\n def test_ddp_speedup(self):\n \"\"\"2-GPU training achieves >=1.3x speedup vs 1 GPU.\"\"\"\n torch.manual_seed(42)\n x = torch.randn(32, 1, 16, 16, 16)\n y = torch.randint(0, 2, (32, 16, 16, 16))\n loader = DataLoader(TensorDataset(x, y), batch_size=4)\n\n model = nn.Sequential(\n nn.Conv3d(1, 16, 3, padding=1), nn.ReLU(), nn.Conv3d(16, 2, 1)\n )\n\n # Time single GPU\n t0 = time.time()\n fit(\n model,\n loader,\n nn.CrossEntropyLoss(),\n torch.optim.Adam(model.parameters()),\n max_epochs=5,\n gpus=1,\n )\n single_time = time.time() - t0\n\n # Time 2 GPUs\n t0 = time.time()\n fit(\n model,\n loader,\n nn.CrossEntropyLoss(),\n torch.optim.Adam(model.parameters()),\n max_epochs=5,\n gpus=2,\n )\n multi_time = time.time() - t0\n\n speedup = single_time / multi_time\n print(\n f\"Speedup: {speedup:.2f}x (single={single_time:.1f}s, multi={multi_time:.1f}s)\"\n )\n assert speedup >= 1.3, f\"Speedup {speedup:.2f}x < 1.3x threshold\"\n" }, { "path": "nobrainer/tests/gpu/test_predict_e2e.py", "content": "\"\"\"GPU end-to-end test: train a UNet on synthetic data, then verify predict()\nproduces high Dice on the same data (overfitting test).\n\nT031 — US1 acceptance scenario 2: Dice >= 0.95 on a known volume.\n\nSince we don't ship reference weights in the repo, this test creates a\nsynthetic brain-like volume (sphere label), trains a UNet to overfit it,\nthen runs predict() and checks the Dice score.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport nibabel as nib\nimport numpy as np\nimport pytest\nimport torch\nimport torch.nn as nn\n\nfrom nobrainer.models.segmentation import unet\nfrom nobrainer.prediction import predict\n\n\ndef _make_sphere_volume(shape=(64, 64, 64), radius=20):\n \"\"\"Create a synthetic volume with a centered sphere as the label.\"\"\"\n vol = np.random.rand(*shape).astype(np.float32) * 0.3\n label = np.zeros(shape, dtype=np.float32)\n center = np.array(shape) / 2\n coords = np.mgrid[: shape[0], : shape[1], : shape[2]]\n dist = np.sqrt(sum((c - ctr) ** 2 for c, ctr in zip(coords, center)))\n mask = dist < radius\n label[mask] = 1.0\n vol[mask] += 0.7 # make sphere brighter\n return vol, label\n\n\n@pytest.mark.gpu\nclass TestPredictEndToEnd:\n def test_unet_overfit_dice_above_threshold(self):\n \"\"\"Train 2-class UNet to overfit a sphere, then check Dice >= 0.95.\"\"\"\n device = torch.device(\"cuda\")\n torch.manual_seed(42)\n\n vol, label = _make_sphere_volume(shape=(64, 64, 64), radius=20)\n x = torch.from_numpy(vol[None, None]).to(device) # (1, 1, 64, 64, 64)\n # One-hot encode label for 2-class: background + foreground\n label_long = torch.from_numpy(label).long().to(device) # (64, 64, 64)\n y_onehot = nn.functional.one_hot(label_long, 2) # (64,64,64,2)\n y_onehot = y_onehot.permute(3, 0, 1, 2).unsqueeze(0).float() # (1,2,64,64,64)\n\n model = unet(n_classes=2).to(device)\n optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n criterion = nn.CrossEntropyLoss()\n\n # Overfit on a single sample\n model.train()\n for _ in range(200):\n optimizer.zero_grad()\n pred = model(x) # (1, 2, 64, 64, 64)\n loss = criterion(pred, label_long.unsqueeze(0))\n loss.backward()\n optimizer.step()\n\n # Run predict() on the same volume — returns argmax labels\n model.eval()\n result = predict(\n inputs=vol,\n model=model,\n block_shape=(32, 32, 32),\n batch_size=8,\n device=\"cuda\",\n return_labels=True,\n )\n\n pred_arr = np.asarray(result.dataobj)\n # Compute Dice for class 1 (foreground)\n pred_bin = (pred_arr == 1).astype(np.float32)\n intersection = (pred_bin * label).sum()\n dice = 2 * intersection / (pred_bin.sum() + label.sum() + 1e-8)\n\n assert dice >= 0.95, f\"Dice {dice:.4f} < 0.95 threshold\"\n\n def test_predict_output_is_nifti_on_gpu(self):\n \"\"\"Verify predict() returns a NIfTI image when run on GPU.\"\"\"\n vol, _ = _make_sphere_volume(shape=(32, 32, 32))\n model = unet(n_classes=2)\n result = predict(\n inputs=vol,\n model=model,\n block_shape=(32, 32, 32),\n batch_size=1,\n device=\"cuda\",\n )\n assert isinstance(result, nib.Nifti1Image)\n assert result.shape == (32, 32, 32)\n" }, { "path": "nobrainer/tests/integration/__init__.py", "content": "" }, { "path": "nobrainer/tests/integration/test_datalad_commit.py", "content": "\"\"\"Integration test for commit_best_model with a real DataLad dataset.\n\nRequirements: datalad>=0.19 and git-annex must be installed.\nNo OSF remote is configured — OSF push is skipped gracefully.\nThe 1-hour SC-008 SLA for OSF retrieval requires live OSF and is not\nvalidated here (manual verification only).\n\"\"\"\n\nfrom __future__ import annotations\n\nimport json\nfrom pathlib import Path\nimport subprocess\n\nimport pytest\nimport torch\n\ndatalad = pytest.importorskip(\"datalad\", reason=\"datalad not installed\")\n\n\n@pytest.fixture()\ndef trained_models_dataset(tmp_path):\n \"\"\"Create a fresh DataLad dataset in tmp_path/trained_models.\"\"\"\n import datalad.api as dl\n\n trained_models = tmp_path / \"trained_models\"\n trained_models.mkdir()\n dl.create(path=str(trained_models))\n return trained_models\n\n\n@pytest.fixture()\ndef model_files(tmp_path):\n \"\"\"Create dummy model.pth and config.json files.\"\"\"\n run_dir = tmp_path / \"run\"\n run_dir.mkdir()\n model_path = run_dir / \"best_model.pth\"\n torch.save({\"weights\": torch.randn(4, 4)}, str(model_path))\n config_path = run_dir / \"best_config.json\"\n config_path.write_text(json.dumps({\"learning_rate\": 1e-4, \"batch_size\": 4}))\n return model_path, config_path\n\n\nclass TestCommitBestModelIntegration:\n def test_files_committed_to_datalad(self, trained_models_dataset, model_files):\n \"\"\"commit_best_model creates model.pth, config.json, model_card.md in dataset.\"\"\"\n from nobrainer.research.loop import commit_best_model\n\n model_path, config_path = model_files\n result = commit_best_model(\n best_model_path=model_path,\n best_config_path=config_path,\n trained_models_path=trained_models_dataset,\n model_family=\"bayesian_vnet\",\n val_dice=0.87,\n source_run_id=\"integration_test_001\",\n )\n\n dest = Path(result[\"path\"])\n assert (dest / \"model.pth\").exists()\n assert (dest / \"config.json\").exists()\n assert (dest / \"model_card.md\").exists()\n\n def test_datalad_dataset_is_clean_after_commit(\n self, trained_models_dataset, model_files\n ):\n \"\"\"datalad status shows no untracked/modified files after commit_best_model.\"\"\"\n import datalad.api as dl\n\n from nobrainer.research.loop import commit_best_model\n\n model_path, config_path = model_files\n commit_best_model(\n best_model_path=model_path,\n best_config_path=config_path,\n trained_models_path=trained_models_dataset,\n model_family=\"bayesian_vnet\",\n val_dice=0.87,\n )\n\n status_results = list(dl.status(dataset=str(trained_models_dataset)))\n unclean = [r for r in status_results if r.get(\"state\") not in (\"clean\", None)]\n assert len(unclean) == 0, f\"Expected clean dataset, got: {unclean}\"\n\n def test_git_log_contains_commit_message(self, trained_models_dataset, model_files):\n \"\"\"Git log in DataLad dataset contains the autoresearch commit.\"\"\"\n from nobrainer.research.loop import commit_best_model\n\n model_path, config_path = model_files\n result = commit_best_model(\n best_model_path=model_path,\n best_config_path=config_path,\n trained_models_path=trained_models_dataset,\n model_family=\"bayesian_vnet\",\n val_dice=0.88,\n )\n\n git_log = subprocess.run(\n [\"git\", \"log\", \"--oneline\", \"-5\"],\n cwd=str(trained_models_dataset),\n capture_output=True,\n text=True,\n check=True,\n )\n assert \"bayesian_vnet\" in git_log.stdout\n assert \"0.8800\" in git_log.stdout\n assert result[\"datalad_commit\"] in git_log.stdout\n\n def test_directory_structure_follows_convention(\n self, trained_models_dataset, model_files\n ):\n \"\"\"Model files land under neuronets/autoresearch///.\"\"\"\n from nobrainer.research.loop import commit_best_model\n\n model_path, config_path = model_files\n result = commit_best_model(\n best_model_path=model_path,\n best_config_path=config_path,\n trained_models_path=trained_models_dataset,\n model_family=\"bayesian_vnet\",\n val_dice=0.90,\n )\n\n dest = Path(result[\"path\"])\n # Path must be: /neuronets/autoresearch/bayesian_vnet/\n parts = dest.parts\n assert \"neuronets\" in parts\n assert \"autoresearch\" in parts\n assert \"bayesian_vnet\" in parts\n\n def test_model_card_contains_required_metadata(\n self, trained_models_dataset, model_files\n ):\n \"\"\"model_card.md includes model family, val_dice, source_run_id, and versions.\"\"\"\n from nobrainer.research.loop import commit_best_model\n\n model_path, config_path = model_files\n result = commit_best_model(\n best_model_path=model_path,\n best_config_path=config_path,\n trained_models_path=trained_models_dataset,\n model_family=\"bayesian_vnet\",\n val_dice=0.85,\n source_run_id=\"run_abc123\",\n )\n\n card = (Path(result[\"path\"]) / \"model_card.md\").read_text()\n assert \"bayesian_vnet\" in card\n assert \"0.8500\" in card\n assert \"run_abc123\" in card\n assert \"PyTorch\" in card\n\n def test_osf_push_skipped_gracefully_when_no_remote(\n self, trained_models_dataset, model_files\n ):\n \"\"\"No OSF remote configured — osf_url is None, function completes normally.\"\"\"\n from nobrainer.research.loop import commit_best_model\n\n model_path, config_path = model_files\n result = commit_best_model(\n best_model_path=model_path,\n best_config_path=config_path,\n trained_models_path=trained_models_dataset,\n model_family=\"bayesian_vnet\",\n val_dice=0.80,\n )\n\n # Without an OSF remote, push fails gracefully and osf_url is None\n assert result[\"osf_url\"] is None\n" }, { "path": "nobrainer/tests/integration/test_research_smoke.py", "content": "\"\"\"Integration test for autoresearch loop with budget-minutes constraint.\n\nT014: Run the full research loop with a 60-second budget and verify\nit produces a run_summary.md with at least 1 experiment entry.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom nobrainer.research.loop import run_loop\n\n\nclass TestResearchSmoke:\n def test_research_loop_completes_with_budget_seconds(self, tmp_path):\n \"\"\"Full research loop with 60s budget, tiny MeshNet on synthetic data.\"\"\"\n # Create a minimal train script that writes val_dice.json quickly\n train_script = tmp_path / \"train.py\"\n train_script.write_text(\n \"import json, random, time\\n\"\n \"time.sleep(0.5)\\n\"\n 'json.dump({\"val_dice\": round(random.uniform(0.4, 0.9), 4)}, '\n 'open(\"val_dice.json\", \"w\"))\\n'\n )\n\n results = run_loop(\n working_dir=tmp_path,\n model_family=\"meshnet\",\n max_experiments=2,\n budget_seconds=60,\n )\n\n # Verify at least 1 experiment ran\n assert len(results) >= 1, \"Expected at least 1 experiment\"\n\n # Verify run_summary.md exists\n summary = tmp_path / \"run_summary.md\"\n assert summary.exists(), \"run_summary.md not created\"\n content = summary.read_text()\n assert \"val_dice\" in content.lower() or \"experiment\" in content.lower()\n" }, { "path": "nobrainer/tests/unit/__init__.py", "content": "" }, { "path": "nobrainer/tests/unit/test_bayesian_layers.py", "content": "\"\"\"Unit tests for BayesianConv3d, BayesianLinear, and accumulate_kl.\"\"\"\n\nfrom __future__ import annotations\n\nimport pyro\nimport pytest\nimport torch\n\nfrom nobrainer.models.bayesian.layers import BayesianConv3d, BayesianLinear\nfrom nobrainer.models.bayesian.utils import accumulate_kl\n\n# ---------------------------------------------------------------------------\n# BayesianConv3d\n# ---------------------------------------------------------------------------\n\n\nclass TestBayesianConv3d:\n def setup_method(self):\n pyro.clear_param_store()\n\n def _forward(self, layer, x):\n \"\"\"Run one forward pass inside a pyro.poutine.trace context.\"\"\"\n with pyro.poutine.trace():\n return layer(x)\n\n def test_output_shape(self):\n layer = BayesianConv3d(1, 4, kernel_size=3, padding=1)\n x = torch.zeros(2, 1, 8, 8, 8)\n out = self._forward(layer, x)\n assert out.shape == (2, 4, 8, 8, 8)\n\n def test_kl_populated_after_forward(self):\n layer = BayesianConv3d(1, 4, kernel_size=3, padding=1)\n x = torch.zeros(2, 1, 8, 8, 8)\n self._forward(layer, x)\n assert isinstance(layer.kl, torch.Tensor)\n assert layer.kl.numel() == 1\n\n def test_kl_positive(self):\n layer = BayesianConv3d(1, 4, kernel_size=3, padding=1)\n x = torch.zeros(2, 1, 8, 8, 8)\n self._forward(layer, x)\n assert layer.kl.item() > 0\n\n def test_kl_varies_across_samples(self):\n \"\"\"KL should differ between two forward passes (stochastic weights).\"\"\"\n layer = BayesianConv3d(1, 4, kernel_size=3, padding=1)\n x = torch.zeros(2, 1, 8, 8, 8)\n self._forward(layer, x)\n kl1 = layer.kl.item()\n self._forward(layer, x)\n kl2 = layer.kl.item()\n # They may occasionally be equal, but should usually differ\n assert kl1 == pytest.approx(kl2, rel=1.0) or kl1 != kl2\n\n def test_prior_laplace(self):\n layer = BayesianConv3d(1, 4, kernel_size=3, padding=1, prior_type=\"laplace\")\n x = torch.zeros(2, 1, 8, 8, 8)\n self._forward(layer, x)\n assert layer.kl.item() > 0\n\n def test_prior_spike_and_slab(self):\n layer = BayesianConv3d(\n 1, 4, kernel_size=3, padding=1, prior_type=\"spike_and_slab\"\n )\n x = torch.zeros(2, 1, 8, 8, 8)\n out = self._forward(layer, x)\n assert out.shape == (2, 4, 8, 8, 8)\n assert isinstance(layer.kl, torch.Tensor)\n assert torch.isfinite(layer.kl)\n # Check that z_logit parameter exists\n assert hasattr(layer, \"z_logit\")\n\n def test_no_bias(self):\n layer = BayesianConv3d(1, 4, kernel_size=3, padding=1, bias=False)\n assert layer.bias_mu is None\n assert layer.bias_rho is None\n x = torch.zeros(2, 1, 8, 8, 8)\n self._forward(layer, x)\n assert layer.kl.item() > 0\n\n def test_weight_sigma_positive(self):\n layer = BayesianConv3d(1, 4, kernel_size=3)\n assert (layer.weight_sigma > 0).all()\n\n\n# ---------------------------------------------------------------------------\n# BayesianLinear\n# ---------------------------------------------------------------------------\n\n\nclass TestBayesianLinear:\n def setup_method(self):\n pyro.clear_param_store()\n\n def _forward(self, layer, x):\n with pyro.poutine.trace():\n return layer(x)\n\n def test_output_shape(self):\n layer = BayesianLinear(16, 8)\n x = torch.zeros(4, 16)\n out = self._forward(layer, x)\n assert out.shape == (4, 8)\n\n def test_kl_populated(self):\n layer = BayesianLinear(16, 8)\n x = torch.zeros(4, 16)\n self._forward(layer, x)\n assert layer.kl.item() > 0\n\n def test_no_bias(self):\n layer = BayesianLinear(16, 8, bias=False)\n assert layer.bias_mu is None\n x = torch.zeros(4, 16)\n self._forward(layer, x)\n assert layer.kl.item() > 0\n\n def test_prior_laplace(self):\n layer = BayesianLinear(16, 8, prior_type=\"laplace\")\n x = torch.zeros(4, 16)\n self._forward(layer, x)\n assert layer.kl.item() > 0\n\n def test_prior_spike_and_slab(self):\n layer = BayesianLinear(16, 8, prior_type=\"spike_and_slab\")\n x = torch.zeros(4, 16)\n out = self._forward(layer, x)\n assert out.shape == (4, 8)\n assert torch.isfinite(layer.kl)\n assert hasattr(layer, \"z_logit\")\n\n\n# ---------------------------------------------------------------------------\n# accumulate_kl\n# ---------------------------------------------------------------------------\n\n\nclass TestAccumulateKl:\n def setup_method(self):\n pyro.clear_param_store()\n\n def test_single_layer(self):\n layer = BayesianConv3d(1, 4, kernel_size=3, padding=1)\n x = torch.zeros(2, 1, 8, 8, 8)\n with pyro.poutine.trace():\n layer(x)\n kl = accumulate_kl(layer)\n assert kl.item() == pytest.approx(layer.kl.item())\n\n def test_multiple_layers(self):\n from pyro.nn import PyroModule\n\n class _TwoConv(PyroModule):\n def __init__(self):\n super().__init__()\n self.l1 = BayesianConv3d(1, 4, kernel_size=3, padding=1)\n self.l2 = BayesianConv3d(4, 8, kernel_size=3, padding=1)\n\n def forward(self, x):\n return self.l2(self.l1(x))\n\n model = _TwoConv()\n x = torch.zeros(2, 1, 8, 8, 8)\n with pyro.poutine.trace():\n model(x)\n total = accumulate_kl(model)\n expected = model.l1.kl + model.l2.kl\n assert total.item() == pytest.approx(expected.item(), rel=1e-5)\n\n def test_non_bayesian_model_returns_zero(self):\n import torch.nn as nn\n\n model = nn.Sequential(nn.Conv3d(1, 4, 3, padding=1))\n kl = accumulate_kl(model)\n assert kl.item() == 0.0\n" }, { "path": "nobrainer/tests/unit/test_bayesian_models.py", "content": "\"\"\"Unit tests for BayesianVNet and BayesianMeshNet.\"\"\"\n\nfrom __future__ import annotations\n\nimport pyro\nimport pytest\nimport torch\n\nfrom nobrainer.models.bayesian import (\n BayesianMeshNet,\n BayesianVNet,\n accumulate_kl,\n bayesian_meshnet,\n bayesian_vnet,\n)\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\n\ndef _run(model, x):\n \"\"\"Forward pass inside a Pyro trace context.\"\"\"\n with pyro.poutine.trace():\n return model(x)\n\n\n# ---------------------------------------------------------------------------\n# BayesianVNet\n# ---------------------------------------------------------------------------\n\n\nclass TestBayesianVNet:\n def setup_method(self):\n pyro.clear_param_store()\n\n def test_default_construction(self):\n m = BayesianVNet()\n assert isinstance(m, BayesianVNet)\n\n def test_output_shape_single_class(self):\n m = BayesianVNet(n_classes=1, in_channels=1, base_filters=8, levels=2)\n x = torch.zeros(2, 1, 16, 16, 16)\n out = _run(m, x)\n assert out.shape == (2, 1, 16, 16, 16)\n\n def test_output_shape_multi_class(self):\n m = BayesianVNet(n_classes=4, in_channels=1, base_filters=8, levels=2)\n x = torch.zeros(2, 1, 16, 16, 16)\n out = _run(m, x)\n assert out.shape == (2, 4, 16, 16, 16)\n\n def test_kl_accumulated(self):\n m = BayesianVNet(n_classes=1, in_channels=1, base_filters=8, levels=2)\n x = torch.zeros(2, 1, 16, 16, 16)\n _run(m, x)\n kl = accumulate_kl(m)\n assert kl.item() > 0\n\n def test_factory_function(self):\n m = bayesian_vnet(n_classes=2, in_channels=1, base_filters=8, levels=2)\n assert isinstance(m, BayesianVNet)\n\n def test_laplace_prior(self):\n m = BayesianVNet(\n n_classes=1, in_channels=1, base_filters=8, levels=2, prior_type=\"laplace\"\n )\n x = torch.zeros(2, 1, 16, 16, 16)\n _run(m, x)\n assert accumulate_kl(m).item() > 0\n\n def test_kl_weight_attribute(self):\n m = BayesianVNet(kl_weight=0.001)\n assert m.kl_weight == pytest.approx(0.001)\n\n\n# ---------------------------------------------------------------------------\n# BayesianMeshNet\n# ---------------------------------------------------------------------------\n\n\nclass TestBayesianMeshNet:\n def setup_method(self):\n pyro.clear_param_store()\n\n def test_default_construction(self):\n m = BayesianMeshNet()\n assert isinstance(m, BayesianMeshNet)\n\n def test_output_shape_single_class(self):\n m = BayesianMeshNet(n_classes=1, in_channels=1, filters=8, receptive_field=37)\n x = torch.zeros(2, 1, 16, 16, 16)\n out = _run(m, x)\n assert out.shape == (2, 1, 16, 16, 16)\n\n def test_output_shape_multi_class(self):\n m = BayesianMeshNet(n_classes=4, in_channels=1, filters=8, receptive_field=37)\n x = torch.zeros(2, 1, 16, 16, 16)\n out = _run(m, x)\n assert out.shape == (2, 4, 16, 16, 16)\n\n def test_kl_accumulated(self):\n m = BayesianMeshNet(n_classes=1, in_channels=1, filters=8, receptive_field=37)\n x = torch.zeros(2, 1, 16, 16, 16)\n _run(m, x)\n assert accumulate_kl(m).item() > 0\n\n def test_invalid_receptive_field(self):\n with pytest.raises(ValueError, match=\"receptive_field\"):\n BayesianMeshNet(receptive_field=99)\n\n def test_all_dilation_schedules(self):\n for rf in [37, 67, 129]:\n m = BayesianMeshNet(\n n_classes=1, in_channels=1, filters=4, receptive_field=rf\n )\n x = torch.zeros(2, 1, 8, 8, 8)\n out = _run(m, x)\n assert out.shape == (2, 1, 8, 8, 8)\n\n def test_factory_function(self):\n m = bayesian_meshnet(n_classes=2, in_channels=1, filters=4, receptive_field=37)\n assert isinstance(m, BayesianMeshNet)\n\n def test_kl_weight_attribute(self):\n m = BayesianMeshNet(kl_weight=1e-4)\n assert m.kl_weight == pytest.approx(1e-4)\n" }, { "path": "nobrainer/tests/unit/test_class_weights.py", "content": "\"\"\"Unit tests for class weight computation and weighted losses.\"\"\"\n\nfrom __future__ import annotations\n\nimport numpy as np\nimport torch\n\nfrom nobrainer.losses import DiceCELoss, compute_class_weights, weighted_cross_entropy\n\n\nclass TestComputeClassWeights:\n def test_uniform_distribution(self, tmp_path):\n \"\"\"Equal class counts → all weights ≈ 1.\"\"\"\n import nibabel as nib\n\n # Create 2-class volume with equal counts\n arr = np.zeros((10, 10, 10), dtype=np.int32)\n arr[:5] = 1 # half zeros, half ones\n nib.save(nib.Nifti1Image(arr, np.eye(4)), str(tmp_path / \"lbl.nii.gz\"))\n\n w = compute_class_weights([str(tmp_path / \"lbl.nii.gz\")], n_classes=2)\n assert w.shape == (2,)\n assert torch.allclose(w, torch.ones(2), atol=0.01)\n\n def test_imbalanced_gives_higher_weight_to_rare(self, tmp_path):\n \"\"\"Rare class should get higher weight.\"\"\"\n import nibabel as nib\n\n arr = np.zeros((10, 10, 10), dtype=np.int32)\n arr[0, 0, 0] = 1 # class 1 is very rare\n nib.save(nib.Nifti1Image(arr, np.eye(4)), str(tmp_path / \"lbl.nii.gz\"))\n\n w = compute_class_weights([str(tmp_path / \"lbl.nii.gz\")], n_classes=2)\n assert w[1] > w[0] # rare class gets higher weight\n\n def test_median_frequency_method(self, tmp_path):\n import nibabel as nib\n\n arr = np.zeros((10, 10, 10), dtype=np.int32)\n arr[:2] = 1\n arr[:1] = 2\n nib.save(nib.Nifti1Image(arr, np.eye(4)), str(tmp_path / \"lbl.nii.gz\"))\n\n w = compute_class_weights(\n [str(tmp_path / \"lbl.nii.gz\")],\n n_classes=3,\n method=\"median_frequency\",\n )\n assert w.shape == (3,)\n assert (w > 0).all()\n\n def test_max_samples(self, tmp_path):\n \"\"\"max_samples limits the number of files scanned.\"\"\"\n import nibabel as nib\n\n for i in range(5):\n arr = np.full((4, 4, 4), i % 2, dtype=np.int32)\n nib.save(\n nib.Nifti1Image(arr, np.eye(4)),\n str(tmp_path / f\"lbl_{i}.nii.gz\"),\n )\n\n paths = [str(tmp_path / f\"lbl_{i}.nii.gz\") for i in range(5)]\n w = compute_class_weights(paths, n_classes=2, max_samples=2)\n assert w.shape == (2,)\n\n\nclass TestWeightedCrossEntropy:\n def test_with_weights(self):\n w = torch.tensor([0.5, 1.5])\n loss_fn = weighted_cross_entropy(weight=w)\n pred = torch.randn(4, 2)\n target = torch.randint(0, 2, (4,))\n loss = loss_fn(pred, target)\n assert loss.ndim == 0\n assert torch.isfinite(loss)\n\n def test_without_weights(self):\n loss_fn = weighted_cross_entropy()\n pred = torch.randn(4, 2)\n target = torch.randint(0, 2, (4,))\n loss = loss_fn(pred, target)\n assert torch.isfinite(loss)\n\n\nclass TestDiceCELoss:\n def test_3d_segmentation(self):\n loss_fn = DiceCELoss(softmax=True)\n pred = torch.randn(2, 3, 8, 8, 8) # 3-class\n target = torch.randint(0, 3, (2, 8, 8, 8))\n loss = loss_fn(pred, target)\n assert loss.ndim == 0\n assert torch.isfinite(loss)\n\n def test_with_class_weights(self):\n w = torch.tensor([0.5, 1.0, 2.0])\n loss_fn = DiceCELoss(weight=w, softmax=True)\n pred = torch.randn(2, 3, 8, 8, 8)\n target = torch.randint(0, 3, (2, 8, 8, 8))\n loss = loss_fn(pred, target)\n assert torch.isfinite(loss)\n\n def test_loss_registry(self):\n from nobrainer.losses import get\n\n loss_cls = get(\"dice_ce\")\n assert loss_cls is DiceCELoss\n" }, { "path": "nobrainer/tests/unit/test_croissant.py", "content": "\"\"\"Unit tests for nobrainer.processing.croissant helpers (T024).\"\"\"\n\nfrom __future__ import annotations\n\nimport json\nfrom pathlib import Path\nfrom unittest.mock import MagicMock\n\nimport nibabel as nib\nimport numpy as np\n\nfrom nobrainer.processing.croissant import (\n _sha256,\n validate_croissant,\n write_dataset_croissant,\n write_model_croissant,\n)\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\n\ndef _make_nifti(shape=(16, 16, 16), tmpdir: Path | None = None) -> str:\n \"\"\"Write a synthetic NIfTI file and return its path.\"\"\"\n data = np.random.rand(*shape).astype(np.float32)\n img = nib.Nifti1Image(data, affine=np.eye(4))\n path = tmpdir / f\"vol_{np.random.randint(0, int(1e6))}.nii.gz\"\n nib.save(img, str(path))\n return str(path)\n\n\ndef _make_fake_estimator(model_name=\"unet\"):\n \"\"\"Create a mock estimator with typical attributes.\"\"\"\n est = MagicMock()\n est.base_model = model_name\n est.model_args = {\"channels\": (4, 8), \"strides\": (2,)}\n est.n_classes_ = 2\n est.block_shape_ = (16, 16, 16)\n est._optimizer_class = \"Adam\"\n est._optimizer_args = {\"lr\": \"0.001\"}\n est._loss_name = \"CrossEntropyLoss\"\n return est\n\n\ndef _make_fake_dataset(tmp_path, n=2):\n \"\"\"Create a mock dataset with real NIfTI files.\"\"\"\n data = []\n for _ in range(n):\n img = _make_nifti((16, 16, 16), tmp_path)\n lbl = _make_nifti((16, 16, 16), tmp_path)\n data.append({\"image\": img, \"label\": lbl})\n\n ds = MagicMock()\n ds.data = data\n ds.volume_shape = (16, 16, 16)\n ds.n_classes = 2\n ds._block_shape = (16, 16, 16)\n return ds\n\n\n# ---------------------------------------------------------------------------\n# Tests: write_model_croissant\n# ---------------------------------------------------------------------------\n\n\nclass TestWriteModelCroissant:\n def test_creates_valid_jsonld(self, tmp_path):\n \"\"\"write_model_croissant() creates a valid JSON-LD file.\"\"\"\n est = _make_fake_estimator()\n ds = _make_fake_dataset(tmp_path)\n result = {\n \"history\": [\n {\"epoch\": 1, \"loss\": 0.5},\n {\"epoch\": 2, \"loss\": 0.4},\n ],\n \"checkpoint_path\": None,\n }\n out = write_model_croissant(tmp_path, est, result, ds)\n assert out.exists()\n data = json.loads(out.read_text())\n assert \"@context\" in data\n assert \"@type\" in data\n assert data[\"@type\"] == \"cr:Dataset\"\n\n def test_required_provenance_fields(self, tmp_path):\n \"\"\"Provenance must contain all required fields.\"\"\"\n est = _make_fake_estimator()\n ds = _make_fake_dataset(tmp_path)\n result = {\n \"history\": [\n {\"epoch\": 1, \"loss\": 0.5},\n {\"epoch\": 2, \"loss\": 0.4},\n ],\n \"checkpoint_path\": None,\n }\n out = write_model_croissant(tmp_path, est, result, ds)\n data = json.loads(out.read_text())\n prov = data[\"nobrainer:provenance\"]\n assert \"source_datasets\" in prov\n assert \"training_date\" in prov\n assert \"nobrainer_version\" in prov\n assert \"model_architecture\" in prov\n\n def test_provenance_model_architecture(self, tmp_path):\n est = _make_fake_estimator(\"meshnet\")\n ds = _make_fake_dataset(tmp_path)\n out = write_model_croissant(tmp_path, est, None, ds)\n data = json.loads(out.read_text())\n assert data[\"nobrainer:provenance\"][\"model_architecture\"] == \"meshnet\"\n\n def test_sha256_checksums_for_source_datasets(self, tmp_path):\n \"\"\"Source datasets must have SHA256 checksums.\"\"\"\n est = _make_fake_estimator()\n ds = _make_fake_dataset(tmp_path, n=2)\n out = write_model_croissant(tmp_path, est, None, ds)\n data = json.loads(out.read_text())\n sources = data[\"nobrainer:provenance\"][\"source_datasets\"]\n assert len(sources) >= 1\n for src in sources:\n assert \"sha256\" in src\n assert len(src[\"sha256\"]) == 64 # SHA256 hex digest length\n\n\nclass TestSHA256:\n def test_checksum_computed(self, tmp_path):\n \"\"\"_sha256 returns a 64-char hex digest for a file.\"\"\"\n path = _make_nifti((4, 4, 4), tmp_path)\n digest = _sha256(path)\n assert isinstance(digest, str)\n assert len(digest) == 64\n\n def test_deterministic(self, tmp_path):\n \"\"\"Same file produces same checksum.\"\"\"\n path = _make_nifti((4, 4, 4), tmp_path)\n assert _sha256(path) == _sha256(path)\n\n\n# ---------------------------------------------------------------------------\n# Tests: validate_croissant\n# ---------------------------------------------------------------------------\n\n\nclass TestValidateCroissant:\n def test_returns_true_on_valid(self, tmp_path):\n \"\"\"validate_croissant() returns True on a valid file.\"\"\"\n est = _make_fake_estimator()\n ds = _make_fake_dataset(tmp_path)\n out = write_model_croissant(tmp_path, est, None, ds)\n assert validate_croissant(out) is True\n\n\n# ---------------------------------------------------------------------------\n# Tests: write_dataset_croissant\n# ---------------------------------------------------------------------------\n\n\nclass TestWriteDatasetCroissant:\n def test_writes_dataset_metadata(self, tmp_path):\n \"\"\"write_dataset_croissant() writes a valid JSON-LD.\"\"\"\n ds = _make_fake_dataset(tmp_path)\n out = write_dataset_croissant(tmp_path / \"ds_croissant.json\", ds)\n assert out.exists()\n data = json.loads(out.read_text())\n assert \"@context\" in data\n assert \"@type\" in data\n assert data[\"@type\"] == \"cr:Dataset\"\n\n def test_dataset_info_present(self, tmp_path):\n ds = _make_fake_dataset(tmp_path)\n out = write_dataset_croissant(tmp_path / \"ds_croissant.json\", ds)\n data = json.loads(out.read_text())\n assert \"nobrainer:dataset_info\" in data\n info = data[\"nobrainer:dataset_info\"]\n assert info[\"n_classes\"] == 2\n assert info[\"n_volumes\"] == 2\n\n def test_distribution_has_sha256(self, tmp_path):\n ds = _make_fake_dataset(tmp_path)\n out = write_dataset_croissant(tmp_path / \"ds_croissant.json\", ds)\n data = json.loads(out.read_text())\n for item in data[\"distribution\"]:\n assert \"sha256\" in item\n assert len(item[\"sha256\"]) == 64\n" }, { "path": "nobrainer/tests/unit/test_dataset.py", "content": "\"\"\"Unit tests for nobrainer.dataset.get_dataset().\"\"\"\n\nfrom pathlib import Path\nimport tempfile\n\nimport nibabel as nib\nimport numpy as np\nimport pytest\n\nfrom nobrainer.dataset import get_dataset\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\n\ndef _make_nifti(shape=(16, 16, 16), tmpdir: Path | None = None) -> str:\n \"\"\"Write a synthetic NIfTI file and return its path.\"\"\"\n if tmpdir is None:\n tmpdir = Path(tempfile.mkdtemp())\n data = np.random.rand(*shape).astype(np.float32)\n img = nib.Nifti1Image(data, affine=np.eye(4))\n path = tmpdir / f\"vol_{np.random.randint(0, 1e6)}.nii.gz\"\n nib.save(img, str(path))\n return str(path)\n\n\n# ---------------------------------------------------------------------------\n# Tests\n# ---------------------------------------------------------------------------\n\n\nclass TestGetDataset:\n def test_batch_shape_image_only(self, tmp_path):\n \"\"\"Verify batch shape (B, 1, D, H, W) for image-only dataset.\"\"\"\n paths = [_make_nifti((16, 16, 16), tmp_path) for _ in range(3)]\n loader = get_dataset(\n image_paths=paths,\n batch_size=2,\n num_workers=0,\n cache_rate=0.0,\n )\n batch = next(iter(loader))\n assert \"image\" in batch\n assert batch[\"image\"].ndim == 5 # (B, C, D, H, W)\n assert batch[\"image\"].shape[0] == 2 # batch size\n assert batch[\"image\"].shape[1] == 1 # channel\n\n def test_batch_shape_with_labels(self, tmp_path):\n \"\"\"Verify both image and label tensors are returned.\"\"\"\n image_paths = [_make_nifti((16, 16, 16), tmp_path) for _ in range(2)]\n label_paths = [_make_nifti((16, 16, 16), tmp_path) for _ in range(2)]\n loader = get_dataset(\n image_paths=image_paths,\n label_paths=label_paths,\n batch_size=2,\n num_workers=0,\n cache_rate=0.0,\n )\n batch = next(iter(loader))\n assert \"image\" in batch\n assert \"label\" in batch\n\n def test_mismatch_raises(self, tmp_path):\n \"\"\"Mismatched image/label list lengths should raise ValueError.\"\"\"\n paths = [_make_nifti((16, 16, 16), tmp_path) for _ in range(2)]\n with pytest.raises(ValueError, match=\"len\"):\n get_dataset(\n image_paths=paths,\n label_paths=paths[:1],\n batch_size=1,\n num_workers=0,\n )\n\n def test_augment_flag(self, tmp_path):\n \"\"\"augment=True should not crash the dataloader.\"\"\"\n paths = [_make_nifti((16, 16, 16), tmp_path) for _ in range(2)]\n loader = get_dataset(\n image_paths=paths,\n batch_size=2,\n num_workers=0,\n augment=True,\n cache_rate=0.0,\n )\n batch = next(iter(loader))\n assert batch[\"image\"].shape[1] == 1\n\n def test_returns_dataloader(self, tmp_path):\n paths = [_make_nifti((16, 16, 16), tmp_path)]\n loader = get_dataset(\n image_paths=paths, batch_size=1, num_workers=0, cache_rate=0.0\n )\n from torch.utils.data import DataLoader\n\n assert isinstance(loader, DataLoader)\n" }, { "path": "nobrainer/tests/unit/test_dataset_builder.py", "content": "\"\"\"Unit tests for nobrainer.processing.dataset.Dataset fluent builder (T013).\"\"\"\n\nfrom __future__ import annotations\n\nimport json\nfrom pathlib import Path\n\nimport nibabel as nib\nimport numpy as np\nfrom torch.utils.data import DataLoader\n\nfrom nobrainer.processing.dataset import Dataset\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\n\ndef _make_nifti(shape=(16, 16, 16), tmpdir: Path | None = None) -> str:\n \"\"\"Write a synthetic NIfTI file and return its path.\"\"\"\n data = np.random.rand(*shape).astype(np.float32)\n img = nib.Nifti1Image(data, affine=np.eye(4))\n path = tmpdir / f\"vol_{np.random.randint(0, int(1e6))}.nii.gz\"\n nib.save(img, str(path))\n return str(path)\n\n\ndef _make_file_pairs(n, shape, tmpdir):\n \"\"\"Create n (image, label) NIfTI file pairs.\"\"\"\n pairs = []\n for _ in range(n):\n img_path = _make_nifti(shape, tmpdir)\n lbl_path = _make_nifti(shape, tmpdir)\n pairs.append((img_path, lbl_path))\n return pairs\n\n\n# ---------------------------------------------------------------------------\n# Tests\n# ---------------------------------------------------------------------------\n\n\nclass TestFromFiles:\n def test_tuple_format(self, tmp_path):\n \"\"\"from_files() accepts list of (image, label) tuples.\"\"\"\n pairs = _make_file_pairs(3, (16, 16, 16), tmp_path)\n ds = Dataset.from_files(pairs, block_shape=(16, 16, 16), n_classes=2)\n assert len(ds.data) == 3\n assert all(\"image\" in d and \"label\" in d for d in ds.data)\n\n def test_dict_format(self, tmp_path):\n \"\"\"from_files() accepts list of dicts.\"\"\"\n pairs = _make_file_pairs(2, (16, 16, 16), tmp_path)\n dicts = [{\"image\": img, \"label\": lbl} for img, lbl in pairs]\n ds = Dataset.from_files(dicts, block_shape=(16, 16, 16), n_classes=2)\n assert len(ds.data) == 2\n\n def test_volume_shape_detected(self, tmp_path):\n \"\"\"from_files() detects volume_shape from the first NIfTI.\"\"\"\n pairs = _make_file_pairs(1, (16, 16, 16), tmp_path)\n ds = Dataset.from_files(pairs, block_shape=(16, 16, 16))\n assert ds.volume_shape == (16, 16, 16)\n\n\nclass TestFluentChaining:\n def test_batch_returns_self(self, tmp_path):\n pairs = _make_file_pairs(2, (16, 16, 16), tmp_path)\n ds = Dataset.from_files(pairs, block_shape=(16, 16, 16))\n result = ds.batch(4)\n assert result is ds\n\n def test_shuffle_returns_self(self, tmp_path):\n pairs = _make_file_pairs(2, (16, 16, 16), tmp_path)\n ds = Dataset.from_files(pairs, block_shape=(16, 16, 16))\n result = ds.shuffle()\n assert result is ds\n\n def test_augment_returns_self(self, tmp_path):\n pairs = _make_file_pairs(2, (16, 16, 16), tmp_path)\n ds = Dataset.from_files(pairs, block_shape=(16, 16, 16))\n result = ds.augment()\n assert result is ds\n\n def test_chaining(self, tmp_path):\n \"\"\"Chaining .batch().shuffle().augment() returns the same instance.\"\"\"\n pairs = _make_file_pairs(2, (16, 16, 16), tmp_path)\n ds = Dataset.from_files(pairs, block_shape=(16, 16, 16))\n result = ds.batch(2).shuffle().augment()\n assert result is ds\n assert ds._batch_size == 2\n assert ds._shuffle is True\n assert ds._augment is True\n\n\nclass TestSplit:\n def test_split_sizes(self, tmp_path):\n \"\"\"split() returns two Datasets with correct combined size.\"\"\"\n pairs = _make_file_pairs(10, (16, 16, 16), tmp_path)\n ds = Dataset.from_files(pairs, block_shape=(16, 16, 16), n_classes=2)\n train, val = ds.split(eval_size=0.2)\n assert len(train.data) + len(val.data) == 10\n assert len(val.data) == 2 # int(10 * 0.2) = 2\n\n def test_split_returns_datasets(self, tmp_path):\n pairs = _make_file_pairs(4, (16, 16, 16), tmp_path)\n ds = Dataset.from_files(pairs, block_shape=(16, 16, 16))\n train, val = ds.split(eval_size=0.25)\n assert isinstance(train, Dataset)\n assert isinstance(val, Dataset)\n\n\nclass TestDataloader:\n def test_returns_dataloader(self, tmp_path):\n \"\"\"dataloader property returns a torch DataLoader.\"\"\"\n pairs = _make_file_pairs(2, (16, 16, 16), tmp_path)\n ds = Dataset.from_files(pairs, block_shape=(16, 16, 16), n_classes=2).batch(2)\n loader = ds.dataloader\n assert isinstance(loader, DataLoader)\n\n def test_batch_produces_data(self, tmp_path):\n \"\"\"DataLoader yields batches with image data.\"\"\"\n pairs = _make_file_pairs(2, (16, 16, 16), tmp_path)\n ds = Dataset.from_files(pairs, block_shape=(16, 16, 16), n_classes=2).batch(2)\n batch = next(iter(ds.dataloader))\n # MONAI DataLoader returns dict with \"image\" key\n assert \"image\" in batch\n assert batch[\"image\"].ndim == 5 # (B, C, D, H, W)\n\n\nclass TestMetadataProperties:\n def test_batch_size(self, tmp_path):\n pairs = _make_file_pairs(2, (16, 16, 16), tmp_path)\n ds = Dataset.from_files(pairs, block_shape=(16, 16, 16)).batch(4)\n assert ds.batch_size == 4\n\n def test_block_shape(self, tmp_path):\n pairs = _make_file_pairs(2, (16, 16, 16), tmp_path)\n ds = Dataset.from_files(pairs, block_shape=(16, 16, 16))\n assert ds.block_shape == (16, 16, 16)\n\n def test_volume_shape(self, tmp_path):\n pairs = _make_file_pairs(2, (16, 16, 16), tmp_path)\n ds = Dataset.from_files(pairs, block_shape=(16, 16, 16))\n assert ds.volume_shape == (16, 16, 16)\n\n def test_n_classes(self, tmp_path):\n pairs = _make_file_pairs(2, (16, 16, 16), tmp_path)\n ds = Dataset.from_files(pairs, block_shape=(16, 16, 16), n_classes=3)\n assert ds.n_classes == 3\n\n\nclass TestToCroissant:\n def test_writes_valid_jsonld(self, tmp_path):\n \"\"\"to_croissant() writes valid JSON-LD with @context and fields.\"\"\"\n pairs = _make_file_pairs(2, (16, 16, 16), tmp_path)\n ds = Dataset.from_files(pairs, block_shape=(16, 16, 16), n_classes=2)\n out = ds.to_croissant(tmp_path / \"dataset_croissant.json\")\n assert out.exists()\n data = json.loads(out.read_text())\n assert \"@context\" in data\n assert \"@type\" in data\n assert data[\"@type\"] == \"cr:Dataset\"\n\n def test_has_dataset_info(self, tmp_path):\n pairs = _make_file_pairs(2, (16, 16, 16), tmp_path)\n ds = Dataset.from_files(pairs, block_shape=(16, 16, 16), n_classes=2)\n out = ds.to_croissant(tmp_path / \"dataset_croissant.json\")\n data = json.loads(out.read_text())\n assert \"nobrainer:dataset_info\" in data\n info = data[\"nobrainer:dataset_info\"]\n assert info[\"n_classes\"] == 2\n assert info[\"n_volumes\"] == 2\n" }, { "path": "nobrainer/tests/unit/test_datasets_openneuro.py", "content": "\"\"\"Unit tests for nobrainer.datasets.openneuro.\"\"\"\n\nfrom __future__ import annotations\n\nfrom pathlib import Path\nfrom unittest.mock import patch\n\nimport pytest\n\n\nclass TestWriteManifest:\n \"\"\"Test write_manifest without DataLad.\"\"\"\n\n def test_creates_csv(self, tmp_path):\n from nobrainer.datasets.openneuro import write_manifest\n\n pairs = [\n {\n \"subject_id\": f\"sub-{i:02d}\",\n \"t1w_path\": f\"/t1_{i}.nii.gz\",\n \"label_path\": f\"/lbl_{i}.nii.gz\",\n }\n for i in range(5)\n ]\n csv_path = write_manifest(pairs, tmp_path / \"manifest.csv\")\n assert csv_path.exists()\n\n import csv\n\n with open(csv_path) as f:\n rows = list(csv.DictReader(f))\n assert len(rows) == 5\n splits = {r[\"split\"] for r in rows}\n assert splits <= {\"train\", \"val\", \"test\"}\n\n def test_split_ratios(self, tmp_path):\n from nobrainer.datasets.openneuro import write_manifest\n\n pairs = [\n {\n \"subject_id\": f\"sub-{i:02d}\",\n \"t1w_path\": f\"/t1_{i}.nii.gz\",\n \"label_path\": f\"/lbl_{i}.nii.gz\",\n }\n for i in range(10)\n ]\n write_manifest(pairs, tmp_path / \"m.csv\", split_ratios=(60, 20, 20))\n\n import csv\n\n with open(tmp_path / \"m.csv\") as f:\n rows = list(csv.DictReader(f))\n n_train = sum(1 for r in rows if r[\"split\"] == \"train\")\n assert n_train == 6 # 60% of 10\n\n def test_dataset_id_column(self, tmp_path):\n from nobrainer.datasets.openneuro import write_manifest\n\n pairs = [\n {\n \"subject_id\": \"sub-01\",\n \"dataset_id\": \"ds000114\",\n \"t1w_path\": \"/t1.nii.gz\",\n \"label_path\": \"/lbl.nii.gz\",\n }\n ]\n csv_path = write_manifest(pairs, tmp_path / \"m.csv\")\n\n import csv\n\n with open(csv_path) as f:\n reader = csv.DictReader(f)\n row = next(reader)\n assert row[\"dataset_id\"] == \"ds000114\"\n\n\nclass TestGlobDataset:\n \"\"\"Test glob_dataset (no DataLad needed).\"\"\"\n\n def test_finds_files(self, tmp_path):\n from nobrainer.datasets.openneuro import glob_dataset\n\n (tmp_path / \"sub-01\" / \"anat\").mkdir(parents=True)\n (tmp_path / \"sub-01\" / \"anat\" / \"sub-01_T1w.nii.gz\").touch()\n (tmp_path / \"sub-02\" / \"anat\").mkdir(parents=True)\n (tmp_path / \"sub-02\" / \"anat\" / \"sub-02_T1w.nii.gz\").touch()\n\n files = glob_dataset(tmp_path, \"sub-*/anat/*_T1w.nii.gz\")\n assert len(files) == 2\n\n def test_no_matches(self, tmp_path):\n from nobrainer.datasets.openneuro import glob_dataset\n\n files = glob_dataset(tmp_path, \"sub-*/anat/*_T1w.nii.gz\")\n assert files == []\n\n\nclass TestExtractSubjectId:\n def test_from_bids_path(self, tmp_path):\n from nobrainer.datasets.openneuro import _extract_subject_id\n\n p = tmp_path / \"sub-03\" / \"anat\" / \"sub-03_T1w.nii.gz\"\n assert _extract_subject_id(p) == \"sub-03\"\n\n def test_from_filename(self):\n from nobrainer.datasets.openneuro import _extract_subject_id\n\n p = Path(\"sub-99_desc-preproc_T1w.nii.gz\")\n assert _extract_subject_id(p) == \"sub-99\"\n\n\nclass TestFileOk:\n def test_real_file(self, tmp_path):\n from nobrainer.datasets.openneuro import _file_ok\n\n f = tmp_path / \"real.nii.gz\"\n f.write_bytes(b\"data\")\n assert _file_ok(f)\n\n def test_empty_file(self, tmp_path):\n from nobrainer.datasets.openneuro import _file_ok\n\n f = tmp_path / \"empty.nii.gz\"\n f.touch()\n assert not _file_ok(f)\n\n def test_missing_file(self, tmp_path):\n from nobrainer.datasets.openneuro import _file_ok\n\n assert not _file_ok(tmp_path / \"missing.nii.gz\")\n\n\nclass TestImportGuard:\n \"\"\"Test that missing datalad gives a clear error.\"\"\"\n\n def test_install_without_datalad(self):\n from nobrainer.datasets.openneuro import install_dataset\n\n with patch.dict(\"sys.modules\", {\"datalad\": None, \"datalad.api\": None}):\n with pytest.raises(ImportError, match=\"DataLad\"):\n install_dataset(\"ds000114\", \"/tmp/test\")\n" }, { "path": "nobrainer/tests/unit/test_estimator_generation.py", "content": "\"\"\"Unit tests for nobrainer.processing.generation.Generation estimator (T029).\"\"\"\n\nfrom __future__ import annotations\n\nimport json\n\nimport nibabel as nib\nimport torch\nfrom torch.utils.data import DataLoader, TensorDataset\n\nfrom nobrainer.processing.generation import Generation\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\nSPATIAL = 4\nGAN_ARGS = {\n \"latent_size\": 8,\n \"fmap_base\": 16,\n \"fmap_max\": 16,\n \"resolution_schedule\": [4],\n \"steps_per_phase\": 100,\n}\n\n\nclass _FakeDataset:\n \"\"\"Minimal dataset-like object for Generation.fit().\"\"\"\n\n def __init__(self, loader):\n self._loader = loader\n self.data = []\n\n @property\n def dataloader(self):\n return self._loader\n\n\ndef _make_fake_dataset(n=4, spatial=SPATIAL, batch_size=2):\n \"\"\"Build a fake dataset with tiny synthetic volumes.\"\"\"\n imgs = torch.randn(n, 1, spatial, spatial, spatial)\n loader = DataLoader(TensorDataset(imgs), batch_size=batch_size)\n return _FakeDataset(loader)\n\n\n# ---------------------------------------------------------------------------\n# Tests\n# ---------------------------------------------------------------------------\n\n\nclass TestGenerationFit:\n def test_fit_returns_self(self):\n \"\"\"Generation('progressivegan').fit() returns self.\"\"\"\n ds = _make_fake_dataset()\n gen = Generation(\"progressivegan\", model_args=GAN_ARGS, multi_gpu=False)\n result = gen.fit(\n ds,\n epochs=10,\n accelerator=\"cpu\",\n enable_progress_bar=False,\n )\n assert result is gen\n\n def test_model_created_after_fit(self):\n ds = _make_fake_dataset()\n gen = Generation(\"progressivegan\", model_args=GAN_ARGS, multi_gpu=False)\n gen.fit(\n ds,\n epochs=5,\n accelerator=\"cpu\",\n enable_progress_bar=False,\n )\n assert gen.model_ is not None\n\n\nclass TestGenerationGenerate:\n def test_generate_returns_list_of_nifti(self):\n \"\"\".generate(2) returns list of 2 nibabel.Nifti1Image.\"\"\"\n ds = _make_fake_dataset()\n gen = Generation(\"progressivegan\", model_args=GAN_ARGS, multi_gpu=False)\n gen.fit(\n ds,\n epochs=5,\n accelerator=\"cpu\",\n enable_progress_bar=False,\n )\n images = gen.generate(2)\n assert isinstance(images, list)\n assert len(images) == 2\n for img in images:\n assert isinstance(img, nib.Nifti1Image)\n\n\nclass TestGenerationSave:\n def test_save_creates_croissant(self, tmp_path):\n \"\"\".save() creates croissant.json.\"\"\"\n ds = _make_fake_dataset()\n gen = Generation(\"progressivegan\", model_args=GAN_ARGS, multi_gpu=False)\n gen.fit(\n ds,\n epochs=5,\n accelerator=\"cpu\",\n enable_progress_bar=False,\n )\n save_dir = tmp_path / \"gen_out\"\n gen.save(save_dir)\n assert (save_dir / \"model.pth\").exists()\n assert (save_dir / \"croissant.json\").exists()\n data = json.loads((save_dir / \"croissant.json\").read_text())\n assert \"@context\" in data\n assert \"nobrainer:provenance\" in data\n" }, { "path": "nobrainer/tests/unit/test_estimator_segmentation.py", "content": "\"\"\"Unit tests for nobrainer.processing.segmentation.Segmentation estimator (T023).\"\"\"\n\nfrom __future__ import annotations\n\nimport json\nfrom pathlib import Path\n\nimport nibabel as nib\nimport numpy as np\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader, TensorDataset\n\nfrom nobrainer.processing.segmentation import Segmentation\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\nSPATIAL = 16\nN_CLASSES = 2\n\n\ndef _make_nifti(shape=(16, 16, 16), tmpdir: Path | None = None) -> str:\n \"\"\"Write a synthetic NIfTI file and return its path.\"\"\"\n data = np.random.rand(*shape).astype(np.float32)\n img = nib.Nifti1Image(data, affine=np.eye(4))\n path = tmpdir / f\"vol_{np.random.randint(0, int(1e6))}.nii.gz\"\n nib.save(img, str(path))\n return str(path)\n\n\ndef _make_tiny_loader(n=4, spatial=SPATIAL, n_classes=N_CLASSES, batch_size=2):\n \"\"\"Create a tiny DataLoader with tuple batches for training.\"\"\"\n x = torch.randn(n, 1, spatial, spatial, spatial)\n y = torch.randint(0, n_classes, (n, spatial, spatial, spatial))\n ds = TensorDataset(x, y)\n return DataLoader(ds, batch_size=batch_size)\n\n\nclass _FakeDataset:\n \"\"\"Minimal object mimicking the Dataset builder for Segmentation.fit().\"\"\"\n\n def __init__(self, loader, block_shape, volume_shape, n_classes):\n self._loader = loader\n self._block_shape = block_shape\n self.volume_shape = volume_shape\n self.n_classes = n_classes\n\n @property\n def block_shape(self):\n return self._block_shape\n\n @property\n def dataloader(self):\n return self._loader\n\n\ndef _make_fake_dataset(n=4, spatial=SPATIAL, n_classes=N_CLASSES, batch_size=2):\n \"\"\"Build a FakeDataset with a tiny DataLoader.\"\"\"\n loader = _make_tiny_loader(n, spatial, n_classes, batch_size)\n return _FakeDataset(\n loader,\n block_shape=(spatial, spatial, spatial),\n volume_shape=(spatial, spatial, spatial),\n n_classes=n_classes,\n )\n\n\n# ---------------------------------------------------------------------------\n# Tests\n# ---------------------------------------------------------------------------\n\n\nclass TestSegmentationFit:\n def test_fit_returns_self(self):\n \"\"\"Segmentation('unet').fit(ds, epochs=2) returns self.\"\"\"\n ds = _make_fake_dataset()\n seg = Segmentation(\n \"unet\",\n model_args={\"channels\": (4, 8), \"strides\": (2,)},\n multi_gpu=False,\n )\n result = seg.fit(ds, epochs=2)\n assert result is seg\n\n def test_model_created_after_fit(self):\n ds = _make_fake_dataset()\n seg = Segmentation(\n \"unet\",\n model_args={\"channels\": (4, 8), \"strides\": (2,)},\n multi_gpu=False,\n )\n seg.fit(ds, epochs=1)\n assert seg.model_ is not None\n assert isinstance(seg.model_, nn.Module)\n\n\nclass TestSegmentationPredict:\n def test_predict_returns_nifti(self, tmp_path):\n \"\"\".predict() returns nibabel.Nifti1Image with correct shape.\"\"\"\n ds = _make_fake_dataset()\n seg = Segmentation(\n \"unet\",\n model_args={\"channels\": (4, 8), \"strides\": (2,)},\n multi_gpu=False,\n )\n seg.fit(ds, epochs=1)\n\n # Create a test volume\n vol_path = _make_nifti((SPATIAL, SPATIAL, SPATIAL), tmp_path)\n result = seg.predict(vol_path, block_shape=(SPATIAL, SPATIAL, SPATIAL))\n assert isinstance(result, nib.Nifti1Image)\n assert result.shape[:3] == (SPATIAL, SPATIAL, SPATIAL)\n\n\nclass TestSegmentationSaveLoad:\n def test_save_creates_files(self, tmp_path):\n \"\"\".save() creates model.pth and croissant.json.\"\"\"\n ds = _make_fake_dataset()\n seg = Segmentation(\n \"unet\",\n model_args={\"channels\": (4, 8), \"strides\": (2,)},\n multi_gpu=False,\n )\n seg.fit(ds, epochs=1)\n save_dir = tmp_path / \"model_out\"\n seg.save(save_dir)\n assert (save_dir / \"model.pth\").exists()\n assert (save_dir / \"croissant.json\").exists()\n\n def test_croissant_provenance_fields(self, tmp_path):\n \"\"\"croissant.json contains all provenance fields.\"\"\"\n ds = _make_fake_dataset()\n seg = Segmentation(\n \"unet\",\n model_args={\"channels\": (4, 8), \"strides\": (2,)},\n multi_gpu=False,\n )\n seg.fit(ds, epochs=1)\n save_dir = tmp_path / \"model_out\"\n seg.save(save_dir)\n data = json.loads((save_dir / \"croissant.json\").read_text())\n prov = data[\"nobrainer:provenance\"]\n assert \"source_datasets\" in prov\n assert \"training_date\" in prov\n assert \"nobrainer_version\" in prov\n assert \"model_architecture\" in prov\n assert prov[\"model_architecture\"] == \"unet\"\n\n def test_load_roundtrip(self, tmp_path):\n \"\"\".load() round-trip produces same prediction output.\"\"\"\n ds = _make_fake_dataset()\n seg = Segmentation(\n \"unet\",\n model_args={\"channels\": (4, 8), \"strides\": (2,)},\n multi_gpu=False,\n )\n seg.fit(ds, epochs=1)\n\n # Get prediction before save\n test_vol = np.random.rand(SPATIAL, SPATIAL, SPATIAL).astype(np.float32)\n pred_before = seg.predict(test_vol, block_shape=(SPATIAL, SPATIAL, SPATIAL))\n\n # Save and reload\n save_dir = tmp_path / \"model_out\"\n seg.save(save_dir)\n loaded = Segmentation.load(save_dir, multi_gpu=False)\n\n # Predict again\n pred_after = loaded.predict(test_vol, block_shape=(SPATIAL, SPATIAL, SPATIAL))\n np.testing.assert_array_equal(\n np.asarray(pred_before.dataobj),\n np.asarray(pred_after.dataobj),\n )\n" }, { "path": "nobrainer/tests/unit/test_experiment.py", "content": "\"\"\"Unit tests for nobrainer.experiment tracking.\"\"\"\n\nfrom __future__ import annotations\n\nimport json\n\nfrom nobrainer.experiment import ExperimentTracker\n\n\nclass TestExperimentTracker:\n def test_local_logging(self, tmp_path):\n tracker = ExperimentTracker(\n output_dir=tmp_path, config={\"lr\": 0.001}, use_wandb=False\n )\n tracker.log({\"epoch\": 1, \"loss\": 0.5})\n tracker.log({\"epoch\": 2, \"loss\": 0.3})\n tracker.finish()\n\n # Check JSONL\n lines = (tmp_path / \"metrics.jsonl\").read_text().strip().split(\"\\n\")\n assert len(lines) == 2\n assert json.loads(lines[0])[\"loss\"] == 0.5\n\n # Check CSV\n csv_lines = (tmp_path / \"metrics.csv\").read_text().strip().split(\"\\n\")\n assert len(csv_lines) == 3 # header + 2 rows\n assert \"epoch\" in csv_lines[0]\n\n # Check config\n config = json.loads((tmp_path / \"config.json\").read_text())\n assert config[\"lr\"] == 0.001\n\n def test_callback(self, tmp_path):\n tracker = ExperimentTracker(output_dir=tmp_path, use_wandb=False)\n cb = tracker.callback(variant=\"test\")\n\n # Simulate training callback\n cb(0, {\"loss\": 1.5}, None) # (epoch, logs_dict, model)\n cb(1, {\"loss\": 0.8}, None)\n tracker.finish()\n\n lines = (tmp_path / \"metrics.jsonl\").read_text().strip().split(\"\\n\")\n assert len(lines) == 2\n row = json.loads(lines[0])\n assert row[\"epoch\"] == 0\n assert row[\"loss\"] == 1.5\n assert row[\"variant\"] == \"test\"\n\n def test_no_wandb_by_default(self, tmp_path):\n tracker = ExperimentTracker(output_dir=tmp_path)\n # Should not fail even without wandb installed\n tracker.log({\"x\": 1})\n tracker.finish()\n" }, { "path": "nobrainer/tests/unit/test_generative.py", "content": "\"\"\"Unit tests for ProgressiveGAN and DCGAN (CPU smoke tests).\"\"\"\n\nfrom __future__ import annotations\n\nimport pytorch_lightning as pl\nimport torch\nfrom torch.utils.data import DataLoader, TensorDataset\n\nfrom nobrainer.models.generative import DCGAN, ProgressiveGAN, dcgan, progressivegan\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\n\ndef _tiny_loader(batch_size: int = 2, spatial: int = 4) -> DataLoader:\n \"\"\"Return a DataLoader with synthetic 3-D volumes.\"\"\"\n imgs = torch.randn(4, 1, spatial, spatial, spatial)\n return DataLoader(TensorDataset(imgs), batch_size=batch_size)\n\n\n# ---------------------------------------------------------------------------\n# ProgressiveGAN\n# ---------------------------------------------------------------------------\n\n\nclass TestProgressiveGAN:\n def test_construction(self):\n m = ProgressiveGAN(\n latent_size=8, fmap_base=16, fmap_max=16, resolution_schedule=[4, 8]\n )\n assert isinstance(m, ProgressiveGAN)\n\n def test_factory_function(self):\n m = progressivegan(\n latent_size=8, fmap_base=16, fmap_max=16, resolution_schedule=[4, 8]\n )\n assert isinstance(m, ProgressiveGAN)\n\n def test_generator_output_shape(self):\n m = ProgressiveGAN(\n latent_size=8, fmap_base=16, fmap_max=16, resolution_schedule=[4]\n )\n m.generator.current_level = 0\n m.generator.alpha = 1.0\n z = torch.randn(2, 8)\n out = m.generator(z)\n assert out.shape[0] == 2\n assert out.shape[1] == 1\n\n def test_discriminator_output_shape(self):\n m = ProgressiveGAN(\n latent_size=8, fmap_base=16, fmap_max=16, resolution_schedule=[4]\n )\n m.discriminator.current_level = 0\n img = torch.randn(2, 1, 4, 4, 4)\n out = m.discriminator(img)\n assert out.shape == (2, 1)\n\n def test_training_step_losses_finite(self):\n \"\"\"5-step CPU training smoke test.\"\"\"\n m = ProgressiveGAN(\n latent_size=8,\n fmap_base=16,\n fmap_max=16,\n resolution_schedule=[4],\n steps_per_phase=10,\n )\n loader = _tiny_loader(batch_size=2, spatial=4)\n trainer = pl.Trainer(\n max_steps=5,\n accelerator=\"cpu\",\n enable_checkpointing=False,\n logger=False,\n enable_progress_bar=False,\n )\n trainer.fit(m, loader)\n # Verify that logged losses are finite\n assert m._step_count > 0\n\n def test_alpha_schedule(self):\n m = ProgressiveGAN(\n latent_size=8,\n fmap_base=16,\n fmap_max=16,\n resolution_schedule=[4, 8],\n steps_per_phase=10,\n )\n m._step_count = 5\n m.on_train_batch_end()\n assert 0.0 <= m.generator.alpha <= 1.0\n\n\n# ---------------------------------------------------------------------------\n# DCGAN\n# ---------------------------------------------------------------------------\n\n\nclass TestDCGAN:\n def test_construction(self):\n m = DCGAN(latent_size=8, n_filters=4)\n assert isinstance(m, DCGAN)\n\n def test_factory_function(self):\n m = dcgan(latent_size=8, n_filters=4)\n assert isinstance(m, DCGAN)\n\n def test_generator_output_shape(self):\n m = DCGAN(latent_size=8, n_filters=4)\n z = torch.randn(2, 8)\n out = m.generator(z)\n assert out.shape[0] == 2\n assert out.shape[1] == 1\n\n def test_discriminator_output_shape(self):\n m = DCGAN(latent_size=8, n_filters=4)\n img = torch.randn(2, 1, 64, 64, 64)\n out = m.discriminator(img)\n assert out.shape == (2, 1)\n\n def test_training_step_losses_finite(self):\n \"\"\"5-step CPU training smoke test.\"\"\"\n m = DCGAN(latent_size=8, n_filters=4)\n loader = _tiny_loader(batch_size=2, spatial=4)\n trainer = pl.Trainer(\n max_steps=5,\n accelerator=\"cpu\",\n enable_checkpointing=False,\n logger=False,\n enable_progress_bar=False,\n )\n trainer.fit(m, loader)\n # No assertion needed — if fit() completes without error, losses were finite\n\n def test_configure_optimizers(self):\n m = DCGAN(latent_size=8, n_filters=4)\n opts = m.configure_optimizers()\n assert len(opts) == 2 # (opt_g, opt_d)\n" }, { "path": "nobrainer/tests/unit/test_gpu.py", "content": "\"\"\"Unit tests for nobrainer.gpu utilities.\"\"\"\n\nfrom __future__ import annotations\n\nimport torch\n\nfrom nobrainer.gpu import get_device, gpu_count, gpu_info, scale_for_multi_gpu\n\n\nclass TestGetDevice:\n def test_returns_torch_device(self):\n d = get_device()\n assert isinstance(d, torch.device)\n\n def test_device_type_known(self):\n d = get_device()\n assert d.type in (\"cuda\", \"mps\", \"cpu\")\n\n\nclass TestGpuCount:\n def test_returns_int(self):\n n = gpu_count()\n assert isinstance(n, int)\n assert n >= 0\n\n\nclass TestGpuInfo:\n def test_returns_list(self):\n info = gpu_info()\n assert isinstance(info, list)\n if torch.cuda.is_available():\n assert len(info) > 0\n assert \"name\" in info[0]\n assert \"memory_gb\" in info[0]\n\n\nclass TestScaleForMultiGpu:\n def test_no_gpu_returns_base(self):\n if torch.cuda.is_available():\n return # skip on GPU machines\n eff, per, n = scale_for_multi_gpu(base_batch_size=32)\n assert eff == 32\n assert per == 32\n assert n == 0\n\n def test_simple_division(self):\n # Without model, just divides\n eff, per, n = scale_for_multi_gpu(base_batch_size=32)\n if n > 0:\n assert eff == per * n\n" }, { "path": "nobrainer/tests/unit/test_io_weights.py", "content": "\"\"\"Unit tests for convert_weights() in nobrainer.io.\"\"\"\n\nfrom pathlib import Path\n\nimport h5py\nimport numpy as np\nimport torch\nimport torch.nn as nn\n\nfrom nobrainer.io import convert_weights\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\n\nclass _SimplePT(nn.Module):\n \"\"\"Minimal PyTorch model for weight-conversion tests.\"\"\"\n\n def __init__(self):\n super().__init__()\n self.conv = nn.Conv3d(1, 4, 3, padding=1, bias=True)\n self.bn = nn.BatchNorm3d(4)\n\n def forward(self, x):\n return self.bn(self.conv(x))\n\n\ndef _write_synthetic_h5(path: str, model: nn.Module) -> None:\n \"\"\"Write a synthetic H5 file that mimics Keras weight layout.\"\"\"\n with h5py.File(path, \"w\") as hf:\n sd = model.state_dict()\n for k, v in sd.items():\n w = v.numpy()\n # Transpose conv weights back to Keras format for the test\n if w.ndim == 5:\n w = np.transpose(w, (2, 3, 4, 1, 0)) # Cout,Cin,D,H,W → D,H,W,Cin,Cout\n hf.create_dataset(\n k.replace(\".\", \"/\") + \"/kernel\" if w.ndim == 5 else k, data=w\n )\n\n\n# ---------------------------------------------------------------------------\n# Tests\n# ---------------------------------------------------------------------------\n\n\nclass TestConvertWeights:\n def test_returns_dict(self, tmp_path):\n model = _SimplePT()\n h5_path = str(tmp_path / \"weights.h5\")\n # Write a minimal H5 that has some datasets\n with h5py.File(h5_path, \"w\") as hf:\n hf.create_dataset(\"dummy\", data=np.zeros(4))\n result = convert_weights(h5_path, model)\n assert isinstance(result, dict)\n\n def test_output_pth_written(self, tmp_path):\n model = _SimplePT()\n h5_path = str(tmp_path / \"weights.h5\")\n pth_path = str(tmp_path / \"weights.pth\")\n with h5py.File(h5_path, \"w\") as hf:\n hf.create_dataset(\"dummy\", data=np.zeros(4))\n convert_weights(h5_path, model, output_path=pth_path)\n assert Path(pth_path).exists()\n loaded = torch.load(pth_path, map_location=\"cpu\", weights_only=True)\n assert isinstance(loaded, dict)\n\n def test_state_dict_keys_preserved(self, tmp_path):\n \"\"\"Model state dict should have same keys before and after conversion.\"\"\"\n model = _SimplePT()\n original_keys = set(model.state_dict().keys())\n h5_path = str(tmp_path / \"weights.h5\")\n with h5py.File(h5_path, \"w\") as hf:\n hf.create_dataset(\"dummy\", data=np.zeros(4))\n convert_weights(h5_path, model)\n assert set(model.state_dict().keys()) == original_keys\n" }, { "path": "nobrainer/tests/unit/test_io_zarr.py", "content": "\"\"\"Unit tests for NIfTI <-> Zarr v3 conversion.\"\"\"\n\nfrom __future__ import annotations\n\nimport nibabel as nib\nimport numpy as np\nimport pytest\n\nzarr = pytest.importorskip(\"zarr\", reason=\"zarr not installed\")\n\nfrom nobrainer.io import nifti_to_zarr, zarr_to_nifti # noqa: E402\n\n\ndef _make_nifti(tmp_path, shape=(32, 32, 32)):\n \"\"\"Create a synthetic NIfTI file and return path + data.\"\"\"\n data = np.random.rand(*shape).astype(np.float32)\n affine = np.diag([2.0, 2.0, 2.0, 1.0])\n img = nib.Nifti1Image(data, affine)\n path = str(tmp_path / \"test.nii.gz\")\n nib.save(img, path)\n return path, data, affine\n\n\nclass TestNiftiToZarr:\n def test_creates_valid_store(self, tmp_path):\n nii_path, data, _ = _make_nifti(tmp_path)\n zarr_path = nifti_to_zarr(nii_path, tmp_path / \"out.zarr\")\n store = zarr.open_group(str(zarr_path), mode=\"r\")\n arr = np.asarray(store[\"0\"])\n assert arr.shape == data.shape\n assert arr.dtype == np.float32\n\n def test_provenance_stored(self, tmp_path):\n nii_path, _, _ = _make_nifti(tmp_path)\n zarr_path = nifti_to_zarr(nii_path, tmp_path / \"out.zarr\")\n store = zarr.open_group(str(zarr_path), mode=\"r\")\n prov = store.attrs.get(\"nobrainer_provenance\")\n assert prov is not None\n assert \"source_file\" in prov\n assert \"created_at\" in prov\n assert \"nobrainer_version\" in prov\n assert prov[\"tool\"] == \"nobrainer.io.nifti_to_zarr\"\n\n def test_multi_resolution_pyramid(self, tmp_path):\n nii_path, data, _ = _make_nifti(tmp_path, shape=(64, 64, 64))\n zarr_path = nifti_to_zarr(nii_path, tmp_path / \"pyramid.zarr\", levels=3)\n store = zarr.open_group(str(zarr_path), mode=\"r\")\n # Level 0: full resolution\n assert np.asarray(store[\"0\"]).shape == (64, 64, 64)\n # Downsampled levels should have smaller shapes\n level1 = np.asarray(store[\"1\"])\n assert all(s <= 64 for s in level1.shape)\n level2 = np.asarray(store[\"2\"])\n assert all(s <= level1.shape[i] for i, s in enumerate(level2.shape))\n\n\nclass TestZarrToNifti:\n def test_round_trip_shape(self, tmp_path):\n \"\"\"NIfTI -> Zarr -> NIfTI preserves shape.\"\"\"\n nii_path, data, _ = _make_nifti(tmp_path)\n zarr_path = nifti_to_zarr(nii_path, tmp_path / \"rt.zarr\")\n rt_path = zarr_to_nifti(zarr_path, tmp_path / \"roundtrip.nii.gz\")\n rt_img = nib.load(str(rt_path))\n assert rt_img.shape == data.shape\n\n def test_round_trip_data(self, tmp_path):\n \"\"\"NIfTI -> Zarr -> NIfTI preserves data values.\"\"\"\n nii_path, data, _ = _make_nifti(tmp_path)\n zarr_path = nifti_to_zarr(nii_path, tmp_path / \"rt.zarr\")\n rt_path = zarr_to_nifti(zarr_path, tmp_path / \"roundtrip.nii.gz\")\n rt_img = nib.load(str(rt_path))\n rt_data = np.asarray(rt_img.dataobj, dtype=np.float32)\n # Value range should be preserved\n assert abs(rt_data.mean() - data.mean()) < 0.1\n assert rt_data.min() >= 0\n assert rt_data.max() <= 1.0 + 0.01\n\n def test_round_trip_level1(self, tmp_path):\n \"\"\"Exporting level 1 gives a smaller shape.\"\"\"\n nii_path, _, _ = _make_nifti(tmp_path, shape=(64, 64, 64))\n zarr_path = nifti_to_zarr(nii_path, tmp_path / \"pyr.zarr\", levels=2)\n rt_path = zarr_to_nifti(zarr_path, tmp_path / \"level1.nii.gz\", level=1)\n rt_img = nib.load(str(rt_path))\n # Level 1 should be smaller than full resolution\n assert all(s <= 64 for s in rt_img.shape)\n" }, { "path": "nobrainer/tests/unit/test_layers.py", "content": "\"\"\"Unit tests for nobrainer.layers (PyTorch implementations).\"\"\"\n\nimport pytest\nimport torch\n\nfrom nobrainer.layers import (\n BernoulliDropout,\n ConcreteDropout,\n GaussianDropout,\n MaxPool4D,\n)\n\n# ---------------------------------------------------------------------------\n# Fixtures\n# ---------------------------------------------------------------------------\n\nSHAPE_3D = (1, 1, 8, 8, 8)\nSHAPE_4D = (1, 1, 2, 8, 8, 8) # (N, C, V, D, H, W)\n\n\n@pytest.fixture\ndef x3d():\n return torch.ones(SHAPE_3D)\n\n\n@pytest.fixture\ndef x4d():\n return torch.ones(SHAPE_4D, requires_grad=False)\n\n\n# ---------------------------------------------------------------------------\n# BernoulliDropout\n# ---------------------------------------------------------------------------\n\n\nclass TestBernoulliDropout:\n def test_forward_shape(self, x3d):\n layer = BernoulliDropout(rate=0.3, is_monte_carlo=True)\n layer.train()\n out = layer(x3d)\n assert out.shape == x3d.shape\n\n def test_passthrough_eval_scale(self, x3d):\n \"\"\"With scale_during_training=True, eval mode returns x unchanged.\"\"\"\n layer = BernoulliDropout(\n rate=0.5, is_monte_carlo=False, scale_during_training=True\n )\n layer.eval()\n out = layer(x3d)\n assert torch.allclose(out, x3d)\n\n def test_passthrough_eval_noscale(self, x3d):\n \"\"\"With scale_during_training=False, eval mode returns x * keep_prob.\"\"\"\n rate = 0.3\n layer = BernoulliDropout(\n rate=rate, is_monte_carlo=False, scale_during_training=False\n )\n layer.eval()\n out = layer(x3d)\n assert torch.allclose(out, x3d * (1.0 - rate))\n\n def test_gradient_flow(self, x3d):\n x = x3d.clone().requires_grad_(True)\n layer = BernoulliDropout(rate=0.3, is_monte_carlo=True, seed=42)\n layer.train()\n out = layer(x)\n loss = out.sum()\n loss.backward()\n assert x.grad is not None\n\n def test_invalid_rate(self):\n with pytest.raises(ValueError):\n BernoulliDropout(rate=1.0, is_monte_carlo=True)\n\n def test_mc_applies_in_eval(self, x3d):\n \"\"\"is_monte_carlo=True applies mask even in eval mode.\"\"\"\n torch.manual_seed(0)\n layer = BernoulliDropout(rate=0.9, is_monte_carlo=True, seed=1)\n layer.eval()\n out = layer(x3d)\n # With high rate some outputs should be zero\n assert out.sum() < x3d.sum()\n\n\n# ---------------------------------------------------------------------------\n# ConcreteDropout\n# ---------------------------------------------------------------------------\n\n\nclass TestConcreteDropout:\n def test_forward_shape(self, x3d):\n N, C, D, H, W = x3d.shape\n layer = ConcreteDropout(in_channels=C, is_monte_carlo=True)\n layer.train()\n out = layer(x3d)\n assert out.shape == x3d.shape\n\n def test_kl_positive(self, x3d):\n N, C, D, H, W = x3d.shape\n layer = ConcreteDropout(in_channels=C, is_monte_carlo=True)\n layer.train()\n _ = layer(x3d)\n assert layer.kl_loss.item() > 0.0\n\n def test_gradient_flow(self, x3d):\n N, C, D, H, W = x3d.shape\n x = x3d.clone().requires_grad_(True)\n layer = ConcreteDropout(in_channels=C, is_monte_carlo=True)\n layer.train()\n out = layer(x)\n # Gradient should flow through p_logit (learnable)\n loss = out.sum() + layer.kl_loss\n loss.backward()\n assert layer.p_logit.grad is not None\n\n def test_p_post_clipped(self, x3d):\n N, C, D, H, W = x3d.shape\n layer = ConcreteDropout(in_channels=C)\n p = layer.p_post\n assert (p >= 0.05).all() and (p <= 0.95).all()\n\n def test_passthrough_eval(self, x3d):\n N, C, D, H, W = x3d.shape\n layer = ConcreteDropout(\n in_channels=C, is_monte_carlo=False, use_expectation=False\n )\n layer.eval()\n out = layer(x3d)\n assert torch.allclose(out, x3d)\n\n\n# ---------------------------------------------------------------------------\n# GaussianDropout\n# ---------------------------------------------------------------------------\n\n\nclass TestGaussianDropout:\n def test_forward_shape(self, x3d):\n layer = GaussianDropout(rate=0.3, is_monte_carlo=True)\n layer.train()\n out = layer(x3d)\n assert out.shape == x3d.shape\n\n def test_passthrough_eval(self, x3d):\n layer = GaussianDropout(rate=0.3, is_monte_carlo=False)\n layer.eval()\n out = layer(x3d)\n assert torch.allclose(out, x3d)\n\n def test_gradient_flow(self, x3d):\n x = x3d.clone().requires_grad_(True)\n layer = GaussianDropout(rate=0.3, is_monte_carlo=True, seed=42)\n layer.train()\n out = layer(x)\n out.sum().backward()\n assert x.grad is not None\n\n def test_mc_in_eval(self, x3d):\n \"\"\"is_monte_carlo=True adds noise even in eval mode.\"\"\"\n torch.manual_seed(0)\n layer = GaussianDropout(rate=0.3, is_monte_carlo=True)\n layer.eval()\n out = layer(x3d)\n # Output should differ from input due to noise\n assert not torch.allclose(out, x3d)\n\n def test_invalid_rate(self):\n with pytest.raises(ValueError):\n GaussianDropout(rate=-0.1, is_monte_carlo=True)\n\n\n# ---------------------------------------------------------------------------\n# MaxPool4D\n# ---------------------------------------------------------------------------\n\n\nclass TestMaxPool4D:\n def test_forward_shape(self, x4d):\n layer = MaxPool4D(kernel_size=2, stride=2)\n out = layer(x4d)\n N, C, V, D, H, W = x4d.shape\n assert out.shape == (N, C, V, D // 2, H // 2, W // 2)\n\n def test_wrong_ndim(self):\n x = torch.ones(1, 1, 8, 8, 8) # 5-D\n layer = MaxPool4D(kernel_size=2)\n with pytest.raises(ValueError, match=\"6-D\"):\n layer(x)\n\n def test_pool_v(self):\n x = torch.randn(1, 1, 4, 8, 8, 8)\n layer = MaxPool4D(kernel_size=2, stride=2, pool_v=2)\n out = layer(x)\n assert out.shape[2] == 2 # V reduced from 4 → 2\n\n def test_gradient_flow(self, x4d):\n x = x4d.clone().float().requires_grad_(True)\n layer = MaxPool4D(kernel_size=2, stride=2)\n out = layer(x)\n out.sum().backward()\n assert x.grad is not None\n" }, { "path": "nobrainer/tests/unit/test_losses.py", "content": "\"\"\"Unit tests for nobrainer.losses (MONAI-backed).\"\"\"\n\nimport pytest\nimport torch\n\nimport nobrainer.losses as losses_module\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\n\ndef _binary_pair(batch=2, spatial=16):\n \"\"\"Return (y_true, y_pred) binary tensors of shape (B, 1, D, H, W).\"\"\"\n y_true = torch.randint(0, 2, (batch, 1, spatial, spatial, spatial)).float()\n y_pred = torch.sigmoid(torch.randn(batch, 1, spatial, spatial, spatial))\n return y_true, y_pred\n\n\ndef _multiclass_pair(batch=2, n_classes=3, spatial=8):\n \"\"\"Return (y_true one-hot, y_pred softmax) tensors.\"\"\"\n labels = torch.randint(0, n_classes, (batch, spatial, spatial, spatial))\n y_true = torch.zeros(batch, n_classes, spatial, spatial, spatial)\n y_true.scatter_(1, labels.unsqueeze(1), 1.0)\n y_pred = torch.softmax(\n torch.randn(batch, n_classes, spatial, spatial, spatial), dim=1\n )\n return y_true, y_pred\n\n\n# ---------------------------------------------------------------------------\n# dice\n# ---------------------------------------------------------------------------\n\n\nclass TestDiceLoss:\n def test_returns_scalar(self):\n y_true, y_pred = _binary_pair()\n loss_fn = losses_module.dice(sigmoid=False)\n loss = loss_fn(y_pred, y_true)\n assert loss.ndim == 0\n\n def test_non_negative(self):\n y_true, y_pred = _binary_pair()\n loss_fn = losses_module.dice(sigmoid=True)\n loss = loss_fn(y_pred, y_true)\n assert loss.item() >= 0.0\n\n def test_perfect_prediction_near_zero(self):\n y = torch.ones(1, 1, 8, 8, 8)\n loss_fn = losses_module.dice()\n loss = loss_fn(y, y)\n assert loss.item() < 0.01\n\n\n# ---------------------------------------------------------------------------\n# generalized_dice\n# ---------------------------------------------------------------------------\n\n\nclass TestGeneralizedDiceLoss:\n def test_returns_scalar(self):\n y_true, y_pred = _multiclass_pair()\n loss_fn = losses_module.generalized_dice(softmax=False)\n loss = loss_fn(y_pred, y_true)\n assert loss.ndim == 0\n\n def test_non_negative(self):\n y_true, y_pred = _multiclass_pair()\n loss_fn = losses_module.generalized_dice()\n loss = loss_fn(y_pred, y_true)\n assert loss.item() >= 0.0\n\n\n# ---------------------------------------------------------------------------\n# jaccard\n# ---------------------------------------------------------------------------\n\n\nclass TestJaccardLoss:\n def test_returns_scalar(self):\n y_true, y_pred = _binary_pair()\n loss_fn = losses_module.jaccard()\n loss = loss_fn(y_pred, y_true)\n assert loss.ndim == 0\n\n def test_non_negative(self):\n y_true, y_pred = _binary_pair()\n loss_fn = losses_module.jaccard()\n loss = loss_fn(y_pred, y_true)\n assert loss.item() >= 0.0\n\n\n# ---------------------------------------------------------------------------\n# tversky\n# ---------------------------------------------------------------------------\n\n\nclass TestTverskyLoss:\n def test_returns_scalar(self):\n y_true, y_pred = _binary_pair()\n loss_fn = losses_module.tversky()\n loss = loss_fn(y_pred, y_true)\n assert loss.ndim == 0\n\n def test_non_negative(self):\n y_true, y_pred = _binary_pair()\n loss_fn = losses_module.tversky(alpha=0.5, beta=0.5)\n loss = loss_fn(y_pred, y_true)\n assert loss.item() >= 0.0\n\n\n# ---------------------------------------------------------------------------\n# stubs\n# ---------------------------------------------------------------------------\n\n\nclass TestStubs:\n def test_elbo_returns_tensor(self):\n \"\"\"elbo() is implemented in Phase 4; non-Bayesian model yields zero KL.\"\"\"\n import torch.nn as nn\n\n result = losses_module.elbo(\n nn.Linear(1, 1), kl_weight=1.0, reconstruction_loss=torch.tensor(0.5)\n )\n assert isinstance(result, torch.Tensor)\n assert result.item() == pytest.approx(0.5)\n\n def test_wasserstein_returns_tensor(self):\n \"\"\"wasserstein() is implemented in Phase 5; E[fake] - E[real].\"\"\"\n real_scores = torch.ones(4)\n fake_scores = torch.zeros(4)\n loss = losses_module.wasserstein(real_scores, fake_scores)\n assert isinstance(loss, torch.Tensor)\n # E[fake] - E[real] = 0 - 1 = -1\n assert loss.item() == pytest.approx(-1.0)\n\n\n# ---------------------------------------------------------------------------\n# get()\n# ---------------------------------------------------------------------------\n\n\nclass TestGet:\n def test_known_loss(self):\n fn = losses_module.get(\"dice\")\n assert callable(fn)\n\n def test_unknown_raises(self):\n with pytest.raises(ValueError, match=\"Unknown loss\"):\n losses_module.get(\"nonexistent\")\n" }, { "path": "nobrainer/tests/unit/test_metrics.py", "content": "\"\"\"Unit tests for nobrainer.metrics (MONAI-backed).\"\"\"\n\nimport pytest\nimport torch\n\nimport nobrainer.metrics as metrics_module\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\n\ndef _onehot_pair(batch=2, n_classes=2, spatial=8):\n \"\"\"One-hot prediction and target tensors (B, C, D, H, W).\"\"\"\n labels = torch.randint(0, n_classes, (batch, spatial, spatial, spatial))\n y_true = torch.zeros(batch, n_classes, spatial, spatial, spatial)\n y_true.scatter_(1, labels.unsqueeze(1), 1.0)\n y_pred = torch.softmax(\n torch.randn(batch, n_classes, spatial, spatial, spatial), dim=1\n )\n # MONAI metrics expect argmax-style binary predictions; use threshold\n y_pred_bin = (y_pred == y_pred.max(dim=1, keepdim=True).values).float()\n return y_true, y_pred_bin\n\n\n# ---------------------------------------------------------------------------\n# dice_metric\n# ---------------------------------------------------------------------------\n\n\nclass TestDiceMetric:\n def test_instantiation(self):\n m = metrics_module.dice_metric()\n assert m is not None\n\n def test_perfect_score(self):\n y, _ = _onehot_pair(batch=1, n_classes=2, spatial=8)\n m = metrics_module.dice_metric(include_background=True)\n m(y_pred=y, y=y)\n result = m.aggregate()\n assert result.item() == pytest.approx(1.0, abs=1e-4)\n m.reset()\n\n def test_output_scalar(self):\n y_true, y_pred = _onehot_pair()\n m = metrics_module.dice_metric()\n m(y_pred=y_pred, y=y_true)\n result = m.aggregate()\n assert result.ndim == 0 or result.numel() == 1\n\n\n# ---------------------------------------------------------------------------\n# jaccard_metric (MeanIoU)\n# ---------------------------------------------------------------------------\n\n\nclass TestJaccardMetric:\n def test_instantiation(self):\n m = metrics_module.jaccard_metric()\n assert m is not None\n\n def test_perfect_score(self):\n y, _ = _onehot_pair(batch=1, n_classes=2, spatial=8)\n m = metrics_module.jaccard_metric()\n m(y_pred=y, y=y)\n result = m.aggregate()\n assert result.item() == pytest.approx(1.0, abs=1e-4)\n m.reset()\n\n\n# ---------------------------------------------------------------------------\n# hausdorff_metric\n# ---------------------------------------------------------------------------\n\n\nclass TestHausdorffMetric:\n def test_instantiation(self):\n m = metrics_module.hausdorff_metric()\n assert m is not None\n\n def test_perfect_score_zero(self):\n y, _ = _onehot_pair(batch=1, n_classes=2, spatial=8)\n m = metrics_module.hausdorff_metric(include_background=False, percentile=95.0)\n m(y_pred=y, y=y)\n result = m.aggregate()\n assert result.item() == pytest.approx(0.0, abs=1e-4)\n m.reset()\n\n\n# ---------------------------------------------------------------------------\n# get()\n# ---------------------------------------------------------------------------\n\n\nclass TestGet:\n def test_known_metric(self):\n fn = metrics_module.get(\"dice\")\n assert callable(fn)\n\n def test_unknown_raises(self):\n with pytest.raises(ValueError, match=\"Unknown metric\"):\n metrics_module.get(\"nonexistent\")\n" }, { "path": "nobrainer/tests/unit/test_model_interface.py", "content": "\"\"\"Unit tests for unified model forward interface.\"\"\"\n\nfrom __future__ import annotations\n\nimport torch\n\nfrom nobrainer.models import get\nfrom nobrainer.models._utils import model_supports_mc\n\n\nclass TestUnifiedForward:\n \"\"\"All models accept model(x) without error.\"\"\"\n\n def test_meshnet(self):\n model = get(\"meshnet\")(n_classes=2, filters=8, receptive_field=37)\n x = torch.randn(1, 1, 16, 16, 16)\n out = model(x)\n assert out.shape == (1, 2, 16, 16, 16)\n\n def test_unet(self):\n model = get(\"unet\")(n_classes=2, channels=(4, 8), strides=(2,))\n x = torch.randn(1, 1, 16, 16, 16)\n out = model(x)\n assert out.shape == (1, 2, 16, 16, 16)\n\n def test_segformer3d(self):\n model = get(\"segformer3d\")(n_classes=2, embed_dims=(16, 32, 80, 128))\n model.eval()\n x = torch.randn(1, 1, 32, 32, 32)\n with torch.no_grad():\n out = model(x)\n assert out.shape == (1, 2, 32, 32, 32)\n\n\nclass TestMcSupport:\n \"\"\"Bayesian models support mc parameter.\"\"\"\n\n def test_kwyk_meshnet_supports_mc(self):\n model = get(\"kwyk_meshnet\")(n_classes=2, filters=8, receptive_field=37)\n assert model_supports_mc(model)\n\n x = torch.randn(1, 1, 16, 16, 16)\n out_det = model(x, mc=False)\n out_mc = model(x, mc=True)\n assert out_det.shape == (1, 2, 16, 16, 16)\n assert out_mc.shape == (1, 2, 16, 16, 16)\n\n def test_bayesian_meshnet_supports_mc(self):\n import pyro\n\n pyro.clear_param_store()\n model = get(\"bayesian_meshnet\")(n_classes=2, filters=8, receptive_field=37)\n assert model_supports_mc(model)\n\n x = torch.randn(1, 1, 16, 16, 16)\n with pyro.poutine.trace():\n out = model(x)\n assert out.shape == (1, 2, 16, 16, 16)\n\n def test_regular_model_no_mc(self):\n model = get(\"meshnet\")(n_classes=2, filters=8, receptive_field=37)\n assert not model_supports_mc(model)\n\n def test_forward_helper_uses_explicit_check(self):\n \"\"\"_forward does NOT use try/except TypeError.\"\"\"\n import inspect\n\n from nobrainer.prediction import _forward\n\n source = inspect.getsource(_forward)\n assert \"except TypeError\" not in source\n assert \"model_supports_mc\" in source\n" }, { "path": "nobrainer/tests/unit/test_model_registry.py", "content": "\"\"\"Unit tests for SwinUNETR and SegResNet model registration.\"\"\"\n\nfrom __future__ import annotations\n\nimport torch\n\nfrom nobrainer.models import get\n\n\nclass TestSwinUNETR:\n def test_instantiate(self):\n model = get(\"swin_unetr\")(n_classes=2, feature_size=12)\n assert model is not None\n\n def test_output_shape(self):\n model = get(\"swin_unetr\")(n_classes=3, feature_size=12)\n model.eval()\n # SwinUNETR needs input >= 64³ due to window attention + instance norm\n x = torch.randn(1, 1, 64, 64, 64)\n with torch.no_grad():\n out = model(x)\n assert out.shape == (1, 3, 64, 64, 64)\n\n\nclass TestSegResNet:\n def test_instantiate(self):\n model = get(\"segresnet\")(n_classes=2, init_filters=8)\n assert model is not None\n\n def test_output_shape(self):\n model = get(\"segresnet\")(n_classes=5, init_filters=8, blocks_down=(1, 2, 2, 4))\n x = torch.randn(1, 1, 32, 32, 32)\n out = model(x)\n assert out.shape == (1, 5, 32, 32, 32)\n\n\nclass TestRegistryAccess:\n def test_swin_unetr_in_registry(self):\n from nobrainer.models import available_models\n\n assert \"swin_unetr\" in available_models()\n\n def test_segresnet_in_registry(self):\n from nobrainer.models import available_models\n\n assert \"segresnet\" in available_models()\n" }, { "path": "nobrainer/tests/unit/test_models_segmentation.py", "content": "\"\"\"Unit tests for nobrainer segmentation models (PyTorch).\"\"\"\n\nimport pytest\nimport torch\n\nfrom nobrainer.models import get as get_model\nfrom nobrainer.models.autoencoder import autoencoder\nfrom nobrainer.models.highresnet import highresnet\nfrom nobrainer.models.meshnet import meshnet\nfrom nobrainer.models.segmentation import attention_unet, unet, unetr, vnet\nfrom nobrainer.models.simsiam import simsiam\n\n# Small spatial size to keep tests fast on CPU\nSPATIAL = 32\nIN_SHAPE = (1, 1, SPATIAL, SPATIAL, SPATIAL)\n\n\ndef _grad_check(model: torch.nn.Module, inp: torch.Tensor) -> bool:\n \"\"\"Return True if gradients flow through all parameters.\"\"\"\n model.train()\n out = model(inp)\n if isinstance(out, tuple):\n loss = sum(o.mean() for o in out)\n else:\n loss = out.mean()\n loss.backward()\n return all(p.grad is not None for p in model.parameters() if p.requires_grad)\n\n\n# ---------------------------------------------------------------------------\n# UNet (MONAI)\n# ---------------------------------------------------------------------------\n\n\nclass TestUNet:\n def test_output_shape_binary(self):\n m = unet(n_classes=1)\n x = torch.randn(*IN_SHAPE)\n assert m(x).shape == (1, 1, SPATIAL, SPATIAL, SPATIAL)\n\n def test_output_shape_multiclass(self):\n m = unet(n_classes=3)\n x = torch.randn(*IN_SHAPE)\n assert m(x).shape == (1, 3, SPATIAL, SPATIAL, SPATIAL)\n\n def test_gradient_flow(self):\n m = unet(n_classes=2)\n x = torch.randn(*IN_SHAPE)\n assert _grad_check(m, x)\n\n def test_get_registry(self):\n fn = get_model(\"unet\")\n assert fn is unet\n\n\n# ---------------------------------------------------------------------------\n# VNet (MONAI)\n# ---------------------------------------------------------------------------\n\n\nclass TestVNet:\n def test_output_shape(self):\n m = vnet(n_classes=1)\n x = torch.randn(*IN_SHAPE)\n out = m(x)\n assert out.shape == (1, 1, SPATIAL, SPATIAL, SPATIAL)\n\n def test_gradient_flow(self):\n m = vnet(n_classes=2)\n x = torch.randn(*IN_SHAPE)\n assert _grad_check(m, x)\n\n\n# ---------------------------------------------------------------------------\n# Attention UNet (MONAI)\n# ---------------------------------------------------------------------------\n\n\nclass TestUNETR:\n def test_output_shape(self):\n m = unetr(\n n_classes=2,\n img_size=(SPATIAL, SPATIAL, SPATIAL),\n hidden_size=192,\n mlp_dim=768,\n num_heads=12,\n feature_size=8,\n )\n x = torch.randn(1, 1, SPATIAL, SPATIAL, SPATIAL)\n m.eval()\n with torch.no_grad():\n out = m(x)\n assert out.shape == (1, 2, SPATIAL, SPATIAL, SPATIAL)\n\n\nclass TestAttentionUNet:\n def test_output_shape(self):\n m = attention_unet(\n n_classes=1,\n channels=(8, 16, 32),\n strides=(2, 2),\n )\n x = torch.randn(*IN_SHAPE)\n assert m(x).shape == (1, 1, SPATIAL, SPATIAL, SPATIAL)\n\n def test_gradient_flow(self):\n m = attention_unet(\n n_classes=2,\n channels=(8, 16, 32),\n strides=(2, 2),\n )\n x = torch.randn(*IN_SHAPE)\n assert _grad_check(m, x)\n\n\n# ---------------------------------------------------------------------------\n# MeshNet (custom PyTorch)\n# ---------------------------------------------------------------------------\n\n\nclass TestMeshNet:\n def test_output_shape_binary(self):\n m = meshnet(n_classes=1)\n x = torch.randn(*IN_SHAPE)\n assert m(x).shape == (1, 1, SPATIAL, SPATIAL, SPATIAL)\n\n def test_output_shape_multiclass(self):\n m = meshnet(n_classes=3)\n x = torch.randn(*IN_SHAPE)\n assert m(x).shape == (1, 3, SPATIAL, SPATIAL, SPATIAL)\n\n def test_receptive_field_37(self):\n m = meshnet(n_classes=1, receptive_field=37)\n x = torch.randn(*IN_SHAPE)\n assert m(x).shape == (1, 1, SPATIAL, SPATIAL, SPATIAL)\n\n def test_receptive_field_129(self):\n m = meshnet(n_classes=1, receptive_field=129)\n x = torch.randn(*IN_SHAPE)\n assert m(x).shape == (1, 1, SPATIAL, SPATIAL, SPATIAL)\n\n def test_invalid_rf(self):\n with pytest.raises(ValueError, match=\"receptive_field\"):\n meshnet(n_classes=1, receptive_field=999)\n\n def test_gradient_flow(self):\n m = meshnet(n_classes=2)\n x = torch.randn(*IN_SHAPE)\n assert _grad_check(m, x)\n\n\n# ---------------------------------------------------------------------------\n# HighResNet (custom PyTorch)\n# ---------------------------------------------------------------------------\n\n\nclass TestHighResNet:\n def test_output_shape_binary(self):\n m = highresnet(n_classes=1)\n x = torch.randn(*IN_SHAPE)\n assert m(x).shape == (1, 1, SPATIAL, SPATIAL, SPATIAL)\n\n def test_output_shape_multiclass(self):\n m = highresnet(n_classes=3)\n x = torch.randn(*IN_SHAPE)\n assert m(x).shape == (1, 3, SPATIAL, SPATIAL, SPATIAL)\n\n def test_gradient_flow(self):\n m = highresnet(n_classes=2)\n x = torch.randn(*IN_SHAPE)\n assert _grad_check(m, x)\n\n\n# ---------------------------------------------------------------------------\n# Autoencoder (custom PyTorch)\n# ---------------------------------------------------------------------------\n\n\nclass TestAutoencoder:\n # Use batch=2 to avoid BatchNorm single-sample issues\n def test_output_shape(self):\n m = autoencoder(input_shape=(SPATIAL, SPATIAL, SPATIAL), encoding_dim=64)\n x = torch.randn(2, 1, SPATIAL, SPATIAL, SPATIAL)\n out = m(x)\n assert out.shape == x.shape\n\n def test_encode_shape(self):\n m = autoencoder(input_shape=(SPATIAL, SPATIAL, SPATIAL), encoding_dim=64)\n x = torch.randn(2, 1, SPATIAL, SPATIAL, SPATIAL)\n z = m.encode(x)\n assert z.shape == (2, 64)\n\n def test_gradient_flow(self):\n m = autoencoder(input_shape=(SPATIAL, SPATIAL, SPATIAL), encoding_dim=32)\n x = torch.randn(2, 1, SPATIAL, SPATIAL, SPATIAL)\n assert _grad_check(m, x)\n\n\n# ---------------------------------------------------------------------------\n# SimSiam (custom PyTorch)\n# ---------------------------------------------------------------------------\n\n\nclass TestSimSiam:\n # Use batch=2 to avoid BatchNorm1d single-sample issues\n def test_forward_shapes(self):\n m = simsiam(projection_dim=128, latent_dim=64)\n x1 = torch.randn(2, 1, SPATIAL, SPATIAL, SPATIAL)\n x2 = torch.randn(2, 1, SPATIAL, SPATIAL, SPATIAL)\n p1, p2, z1, z2 = m(x1, x2)\n assert p1.shape == (2, 128)\n assert z1.shape == (2, 128)\n\n def test_loss_negative_range(self):\n m = simsiam(projection_dim=128, latent_dim=64)\n x1 = torch.randn(2, 1, SPATIAL, SPATIAL, SPATIAL)\n x2 = torch.randn(2, 1, SPATIAL, SPATIAL, SPATIAL)\n p1, p2, z1, z2 = m(x1, x2)\n loss = m.loss(p1, p2, z1, z2)\n # Loss should be in [-1, 0] for cosine similarity\n assert -1.1 <= loss.item() <= 0.1\n\n def test_gradient_flow(self):\n m = simsiam(projection_dim=128, latent_dim=64)\n m.train()\n x1 = torch.randn(2, 1, SPATIAL, SPATIAL, SPATIAL)\n x2 = torch.randn(2, 1, SPATIAL, SPATIAL, SPATIAL)\n p1, p2, z1, z2 = m(x1, x2)\n loss = m.loss(p1, p2, z1, z2)\n loss.backward()\n assert all(\n p.grad is not None for p in m.projector.parameters() if p.requires_grad\n )\n" }, { "path": "nobrainer/tests/unit/test_prediction.py", "content": "\"\"\"Unit tests for predict() and predict_with_uncertainty().\"\"\"\n\nfrom __future__ import annotations\n\nfrom pathlib import Path\nimport tempfile\n\nimport nibabel as nib\nimport numpy as np\nimport torch.nn as nn\n\nfrom nobrainer.prediction import predict, predict_with_uncertainty\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\n\nclass _IdentityModel(nn.Module):\n \"\"\"Minimal 1-class model: sigmoid of a 1×1×1 conv applied to input.\"\"\"\n\n def __init__(self):\n super().__init__()\n self.conv = nn.Conv3d(1, 1, kernel_size=1)\n\n def forward(self, x):\n return self.conv(x)\n\n\nclass _MultiClassModel(nn.Module):\n \"\"\"Minimal 3-class model.\"\"\"\n\n def __init__(self):\n super().__init__()\n self.conv = nn.Conv3d(1, 3, kernel_size=1)\n\n def forward(self, x):\n return self.conv(x)\n\n\ndef _make_nifti(shape=(32, 32, 32), tmp_path=None) -> str:\n if tmp_path is None:\n tmp_path = Path(tempfile.mkdtemp())\n data = np.random.rand(*shape).astype(np.float32)\n img = nib.Nifti1Image(data, np.eye(4))\n path = str(tmp_path / f\"vol_{np.random.randint(0, 1e6)}.nii.gz\")\n nib.save(img, path)\n return path\n\n\n# ---------------------------------------------------------------------------\n# predict()\n# ---------------------------------------------------------------------------\n\n\nclass TestPredict:\n def test_returns_nifti(self, tmp_path):\n path = _make_nifti((16, 16, 16), tmp_path)\n model = _IdentityModel()\n out = predict(path, model, block_shape=(8, 8, 8), batch_size=2)\n assert isinstance(out, nib.Nifti1Image)\n\n def test_output_shape_matches_input(self, tmp_path):\n path = _make_nifti((16, 16, 16), tmp_path)\n model = _IdentityModel()\n out = predict(path, model, block_shape=(8, 8, 8), batch_size=2)\n assert out.shape == (16, 16, 16)\n\n def test_ndarray_input(self):\n arr = np.random.rand(16, 16, 16).astype(np.float32)\n model = _IdentityModel()\n out = predict(arr, model, block_shape=(8, 8, 8), batch_size=2)\n assert out.shape == (16, 16, 16)\n\n def test_nifti_image_input(self):\n arr = np.random.rand(16, 16, 16).astype(np.float32)\n img = nib.Nifti1Image(arr, np.eye(4))\n model = _IdentityModel()\n out = predict(img, model, block_shape=(8, 8, 8), batch_size=2)\n assert out.shape == (16, 16, 16)\n\n def test_affine_preserved(self, tmp_path):\n path = _make_nifti((16, 16, 16), tmp_path)\n model = _IdentityModel()\n src_affine = nib.load(path).affine\n out = predict(path, model, block_shape=(8, 8, 8), batch_size=2)\n assert np.allclose(out.affine, src_affine)\n\n def test_return_probabilities(self):\n arr = np.random.rand(16, 16, 16).astype(np.float32)\n model = _MultiClassModel()\n out = predict(\n arr, model, block_shape=(8, 8, 8), batch_size=2, return_labels=False\n )\n # 3-class probabilities → shape (3, D, H, W)\n assert out.shape[:1] == (3,) or out.ndim == 4\n\n def test_non_block_aligned_input(self):\n \"\"\"Volume with shape not divisible by block_shape should still work.\"\"\"\n arr = np.random.rand(20, 20, 20).astype(np.float32)\n model = _IdentityModel()\n out = predict(arr, model, block_shape=(8, 8, 8), batch_size=2)\n assert out.shape == (20, 20, 20)\n\n\n# ---------------------------------------------------------------------------\n# predict_with_uncertainty()\n# ---------------------------------------------------------------------------\n\n\nclass TestPredictWithUncertainty:\n def test_returns_three_niftis(self):\n arr = np.random.rand(16, 16, 16).astype(np.float32)\n model = _IdentityModel()\n label, var, entropy = predict_with_uncertainty(\n arr, model, n_samples=3, block_shape=(8, 8, 8), batch_size=2\n )\n assert isinstance(label, nib.Nifti1Image)\n assert isinstance(var, nib.Nifti1Image)\n assert isinstance(entropy, nib.Nifti1Image)\n\n def test_output_shapes_match_input(self):\n arr = np.random.rand(16, 16, 16).astype(np.float32)\n model = _IdentityModel()\n label, var, entropy = predict_with_uncertainty(\n arr, model, n_samples=3, block_shape=(8, 8, 8), batch_size=2\n )\n assert label.shape == (16, 16, 16)\n assert var.shape == (16, 16, 16)\n assert entropy.shape == (16, 16, 16)\n\n def test_variance_nonnegative(self):\n arr = np.random.rand(16, 16, 16).astype(np.float32)\n model = _IdentityModel()\n _, var, _ = predict_with_uncertainty(\n arr, model, n_samples=3, block_shape=(8, 8, 8), batch_size=2\n )\n assert (np.asarray(var.dataobj) >= 0).all()\n\n def test_entropy_nonnegative(self):\n arr = np.random.rand(16, 16, 16).astype(np.float32)\n model = _IdentityModel()\n _, _, entropy = predict_with_uncertainty(\n arr, model, n_samples=3, block_shape=(8, 8, 8), batch_size=2\n )\n assert (np.asarray(entropy.dataobj) >= 0).all()\n" }, { "path": "nobrainer/tests/unit/test_research_commit.py", "content": "\"\"\"Unit tests for commit_best_model in nobrainer.research.loop.\"\"\"\n\nfrom __future__ import annotations\n\nimport json\nfrom pathlib import Path\nfrom unittest.mock import MagicMock, patch\n\nimport pytest\n\n\ndef _make_model_files(tmp_path: Path) -> tuple[Path, Path]:\n \"\"\"Create dummy model and config files.\"\"\"\n model_path = tmp_path / \"best_model.pth\"\n model_path.write_bytes(b\"\\x00\" * 16) # dummy weights\n config_path = tmp_path / \"best_config.json\"\n config_path.write_text(json.dumps({\"learning_rate\": 1e-4, \"batch_size\": 4}))\n return model_path, config_path\n\n\nclass TestCommitBestModel:\n def test_directory_structure_created(self, tmp_path):\n \"\"\"commit_best_model creates the expected subdirectory.\"\"\"\n model_path, config_path = _make_model_files(tmp_path)\n trained_models = tmp_path / \"trained_models\"\n trained_models.mkdir()\n\n mock_dl = MagicMock()\n mock_dl.save = MagicMock()\n mock_dl.push = MagicMock()\n\n with patch.dict(\n \"sys.modules\", {\"datalad\": MagicMock(), \"datalad.api\": mock_dl}\n ):\n from nobrainer.research.loop import commit_best_model\n\n result = commit_best_model(\n best_model_path=model_path,\n best_config_path=config_path,\n trained_models_path=trained_models,\n model_family=\"bayesian_vnet\",\n val_dice=0.85,\n source_run_id=\"run_001\",\n )\n\n dest = Path(result[\"path\"])\n assert dest.exists()\n assert (dest / \"model.pth\").exists()\n assert (dest / \"config.json\").exists()\n assert (dest / \"model_card.md\").exists()\n\n def test_model_card_contains_required_fields(self, tmp_path):\n \"\"\"model_card.md contains architecture, val_dice, source_run_id.\"\"\"\n model_path, config_path = _make_model_files(tmp_path)\n trained_models = tmp_path / \"trained_models\"\n trained_models.mkdir()\n\n mock_dl = MagicMock()\n\n with patch.dict(\n \"sys.modules\", {\"datalad\": MagicMock(), \"datalad.api\": mock_dl}\n ):\n from nobrainer.research.loop import commit_best_model\n\n result = commit_best_model(\n best_model_path=model_path,\n best_config_path=config_path,\n trained_models_path=trained_models,\n model_family=\"bayesian_vnet\",\n val_dice=0.85,\n source_run_id=\"run_42\",\n )\n\n card = (Path(result[\"path\"]) / \"model_card.md\").read_text()\n assert \"bayesian_vnet\" in card\n assert \"0.8500\" in card\n assert \"run_42\" in card\n assert \"PyTorch\" in card\n\n def test_model_version_dict_fields(self, tmp_path):\n \"\"\"commit_best_model returns ModelVersion dict with expected keys.\"\"\"\n model_path, config_path = _make_model_files(tmp_path)\n trained_models = tmp_path / \"trained_models\"\n trained_models.mkdir()\n\n mock_dl = MagicMock()\n\n with patch.dict(\n \"sys.modules\", {\"datalad\": MagicMock(), \"datalad.api\": mock_dl}\n ):\n from nobrainer.research.loop import commit_best_model\n\n result = commit_best_model(\n best_model_path=model_path,\n best_config_path=config_path,\n trained_models_path=trained_models,\n model_family=\"bayesian_vnet\",\n val_dice=0.75,\n )\n\n assert \"path\" in result\n assert \"datalad_commit\" in result\n assert \"val_dice\" in result\n assert \"model_family\" in result\n assert result[\"val_dice\"] == pytest.approx(0.75)\n assert result[\"model_family\"] == \"bayesian_vnet\"\n\n def test_datalad_commit_message_in_result(self, tmp_path):\n \"\"\"commit_best_model result contains a descriptive datalad_commit message.\"\"\"\n model_path, config_path = _make_model_files(tmp_path)\n trained_models = tmp_path / \"trained_models\"\n trained_models.mkdir()\n\n mock_dl = MagicMock()\n\n with patch.dict(\n \"sys.modules\", {\"datalad\": MagicMock(), \"datalad.api\": mock_dl}\n ):\n from nobrainer.research.loop import commit_best_model\n\n result = commit_best_model(\n best_model_path=model_path,\n best_config_path=config_path,\n trained_models_path=trained_models,\n model_family=\"bayesian_vnet\",\n val_dice=0.9,\n )\n\n assert \"bayesian_vnet\" in result[\"datalad_commit\"]\n assert \"0.9000\" in result[\"datalad_commit\"]\n\n def test_result_contains_osf_url_key(self, tmp_path):\n \"\"\"commit_best_model result always contains the osf_url key.\"\"\"\n model_path, config_path = _make_model_files(tmp_path)\n trained_models = tmp_path / \"trained_models\"\n trained_models.mkdir()\n\n with patch.dict(\n \"sys.modules\", {\"datalad\": MagicMock(), \"datalad.api\": MagicMock()}\n ):\n from nobrainer.research.loop import commit_best_model\n\n result = commit_best_model(\n best_model_path=model_path,\n best_config_path=config_path,\n trained_models_path=trained_models,\n model_family=\"bayesian_vnet\",\n val_dice=0.8,\n )\n\n # osf_url is present; it is either 'osf://' (push succeeded) or None\n assert \"osf_url\" in result\n\n def test_datalad_not_installed_raises_import_error(self, tmp_path):\n \"\"\"ImportError raised with helpful message when datalad missing.\"\"\"\n model_path, config_path = _make_model_files(tmp_path)\n trained_models = tmp_path / \"trained_models\"\n trained_models.mkdir()\n\n with patch.dict(\"sys.modules\", {\"datalad\": None, \"datalad.api\": None}):\n from nobrainer.research.loop import commit_best_model\n\n with pytest.raises(ImportError, match=\"datalad\"):\n commit_best_model(\n best_model_path=model_path,\n best_config_path=config_path,\n trained_models_path=trained_models,\n model_family=\"bayesian_vnet\",\n val_dice=0.8,\n )\n" }, { "path": "nobrainer/tests/unit/test_research_loop.py", "content": "\"\"\"Unit tests for the autoresearch run_loop.\"\"\"\n\nfrom __future__ import annotations\n\nimport json\nfrom pathlib import Path\nfrom unittest.mock import patch\n\nimport pytest\n\nfrom nobrainer.research.loop import (\n ExperimentResult,\n _classify_failure,\n _has_nan,\n _parse_config_comment,\n _patch_config,\n _read_val_dice,\n _write_summary,\n run_loop,\n)\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\n\ndef _write_train_script(path: Path, config: dict | None = None) -> None:\n cfg = config or {\"learning_rate\": 1e-4, \"batch_size\": 4}\n path.write_text(\n f\"# CONFIG: {json.dumps(cfg)}\\n\"\n \"import sys; print('training done'); sys.exit(0)\\n\"\n )\n\n\ndef _write_val_dice(path: Path, val_dice: float) -> None:\n path.write_text(json.dumps({\"val_dice\": val_dice}))\n\n\n# ---------------------------------------------------------------------------\n# Unit helpers\n# ---------------------------------------------------------------------------\n\n\nclass TestHelpers:\n def test_parse_config_comment(self, tmp_path):\n script = tmp_path / \"train.py\"\n script.write_text(\"# CONFIG: {\\\"lr\\\": 1e-4}\\nprint('hi')\\n\")\n config = _parse_config_comment(script)\n assert config[\"lr\"] == pytest.approx(1e-4)\n\n def test_parse_config_comment_missing(self, tmp_path):\n script = tmp_path / \"train.py\"\n script.write_text(\"print('no config')\\n\")\n config = _parse_config_comment(script)\n assert config == {}\n\n def test_patch_config(self, tmp_path):\n script = tmp_path / \"train.py\"\n script.write_text(\"# CONFIG: {\\\"lr\\\": 1e-4}\\nprint('hi')\\n\")\n _patch_config(script, {\"lr\": 5e-4})\n content = script.read_text()\n assert (\n '\"lr\": 0.0005' in content or '\"lr\": 5e-4' in content or \"5e-04\" in content\n )\n\n def test_patch_config_adds_when_missing(self, tmp_path):\n script = tmp_path / \"train.py\"\n script.write_text(\"print('no config')\\n\")\n _patch_config(script, {\"lr\": 1e-3})\n content = script.read_text()\n assert \"# CONFIG:\" in content\n\n def test_read_val_dice_valid(self, tmp_path):\n (tmp_path / \"val_dice.json\").write_text('{\"val_dice\": 0.85}')\n assert _read_val_dice(tmp_path / \"val_dice.json\") == pytest.approx(0.85)\n\n def test_read_val_dice_missing(self, tmp_path):\n assert _read_val_dice(tmp_path / \"nonexistent.json\") is None\n\n def test_has_nan(self):\n assert _has_nan(\"loss: nan after epoch 3\")\n assert not _has_nan(\"loss: 0.25\")\n\n def test_classify_failure_oom(self):\n assert _classify_failure(\"CUDA out of memory\") == \"CUDA OOM\"\n\n def test_classify_failure_nan(self):\n assert _classify_failure(\"nan in grad\") == \"NaN in loss\"\n\n def test_classify_failure_generic(self):\n assert _classify_failure(\"some error\") == \"non-zero exit code\"\n\n def test_write_summary(self, tmp_path):\n results = [\n ExperimentResult(0, {}, 0.8, \"improved\"),\n ExperimentResult(1, {}, 0.79, \"degraded\"),\n ]\n _write_summary(tmp_path, results, \"bayesian_vnet\", 0.8)\n summary = (tmp_path / \"run_summary.md\").read_text()\n assert \"bayesian_vnet\" in summary\n assert \"0.8000\" in summary\n\n\n# ---------------------------------------------------------------------------\n# run_loop integration tests (subprocess mocked)\n# ---------------------------------------------------------------------------\n\n\nclass TestRunLoop:\n def test_keep_improved_experiment(self, tmp_path):\n \"\"\"run_loop keeps config when val_dice improves.\"\"\"\n _write_train_script(tmp_path / \"train.py\")\n _write_val_dice(tmp_path / \"val_dice.json\", 0.9)\n\n with (\n patch(\n \"nobrainer.research.loop._propose_config\",\n side_effect=[\n {\"learning_rate\": 5e-4, \"batch_size\": 4},\n ]\n * 5,\n ),\n patch(\n \"nobrainer.research.loop.subprocess.run\",\n ) as mock_run,\n ):\n mock_run.return_value.returncode = 0\n mock_run.return_value.stdout = \"training done\\n\"\n mock_run.return_value.stderr = \"\"\n results = run_loop(\n tmp_path,\n max_experiments=2,\n budget_hours=1.0,\n )\n\n improved = [r for r in results if r.outcome == \"improved\"]\n assert len(improved) >= 1\n\n def test_revert_on_degraded(self, tmp_path):\n \"\"\"run_loop reverts train.py when val_dice degrades.\"\"\"\n original_content = (\n f\"# CONFIG: {json.dumps({'learning_rate': 1e-4, 'batch_size': 4})}\\n\"\n \"import sys; sys.exit(0)\\n\"\n )\n (tmp_path / \"train.py\").write_text(original_content)\n # First experiment improves, second degrades\n dices = [0.8, 0.7]\n\n call_count = [0]\n\n def _mock_run(cmd, **kwargs):\n from unittest.mock import MagicMock\n\n dice = dices[call_count[0] % len(dices)]\n _write_val_dice(tmp_path / \"val_dice.json\", dice)\n call_count[0] += 1\n r = MagicMock()\n r.returncode = 0\n r.stdout = \"done\\n\"\n r.stderr = \"\"\n return r\n\n with (\n patch(\"nobrainer.research.loop.subprocess.run\", side_effect=_mock_run),\n patch(\n \"nobrainer.research.loop._propose_config\",\n side_effect=[{\"learning_rate\": 5e-4}] * 5,\n ),\n ):\n results = run_loop(tmp_path, max_experiments=2, budget_hours=1.0)\n\n degraded = [r for r in results if r.outcome == \"degraded\"]\n assert len(degraded) >= 1\n\n def test_failure_handling_reverts(self, tmp_path):\n \"\"\"run_loop reverts train.py when subprocess fails.\"\"\"\n _write_train_script(tmp_path / \"train.py\")\n original = (tmp_path / \"train.py\").read_text()\n\n with (\n patch(\n \"nobrainer.research.loop.subprocess.run\",\n ) as mock_run,\n patch(\n \"nobrainer.research.loop._propose_config\",\n return_value={\"learning_rate\": 1e-3},\n ),\n ):\n mock_run.return_value.returncode = 1\n mock_run.return_value.stdout = \"\"\n mock_run.return_value.stderr = \"some error\"\n results = run_loop(tmp_path, max_experiments=1, budget_hours=1.0)\n\n assert results[0].outcome == \"failed\"\n # Train script reverted\n assert (tmp_path / \"train.py\").read_text() == original\n\n def test_run_summary_written(self, tmp_path):\n \"\"\"run_summary.md is written after the loop.\"\"\"\n _write_train_script(tmp_path / \"train.py\")\n with (\n patch(\n \"nobrainer.research.loop.subprocess.run\",\n ) as mock_run,\n patch(\n \"nobrainer.research.loop._propose_config\",\n return_value={\"learning_rate\": 1e-4},\n ),\n ):\n mock_run.return_value.returncode = 1\n mock_run.return_value.stdout = \"\"\n mock_run.return_value.stderr = \"error\"\n run_loop(tmp_path, max_experiments=1, budget_hours=1.0)\n\n assert (tmp_path / \"run_summary.md\").exists()\n\n def test_missing_train_script_raises(self, tmp_path):\n with pytest.raises(FileNotFoundError):\n run_loop(tmp_path, max_experiments=1, budget_hours=1.0)\n\n def test_budget_seconds_terminates_quickly(self, tmp_path):\n \"\"\"T013: budget_seconds=10 should terminate within 15s.\"\"\"\n import time\n\n (tmp_path / \"train.py\").write_text(\n \"import json, time; time.sleep(0.1);\\n\"\n 'json.dump({\"val_dice\": 0.5}, open(\"val_dice.json\", \"w\"))\\n'\n )\n start = time.time()\n with patch(\n \"nobrainer.research.loop._propose_config\",\n return_value={},\n ):\n run_loop(\n tmp_path,\n max_experiments=100,\n budget_seconds=5,\n )\n elapsed = time.time() - start\n assert elapsed < 15, f\"Loop took {elapsed:.1f}s, expected < 15s\"\n" }, { "path": "nobrainer/tests/unit/test_segformer3d.py", "content": "\"\"\"Unit tests for SegFormer3D model.\"\"\"\n\nfrom __future__ import annotations\n\nimport torch\n\nfrom nobrainer.models import get\nfrom nobrainer.models.segformer3d import SegFormer3D\n\n\nclass TestSegFormer3DShapes:\n def test_output_shape_32(self):\n model = SegFormer3D(n_classes=2, embed_dims=(16, 32, 80, 128))\n model.eval()\n x = torch.randn(1, 1, 32, 32, 32)\n with torch.no_grad():\n out = model(x)\n assert out.shape == (1, 2, 32, 32, 32)\n\n def test_output_shape_64(self):\n model = SegFormer3D(n_classes=5, embed_dims=(16, 32, 80, 128))\n model.eval()\n x = torch.randn(1, 1, 64, 64, 64)\n with torch.no_grad():\n out = model(x)\n assert out.shape == (1, 5, 64, 64, 64)\n\n def test_batch_size_2(self):\n model = SegFormer3D(n_classes=3, embed_dims=(16, 32, 80, 128))\n model.eval()\n x = torch.randn(2, 1, 32, 32, 32)\n with torch.no_grad():\n out = model(x)\n assert out.shape == (2, 3, 32, 32, 32)\n\n\nclass TestSegFormer3DParams:\n def test_default_param_count(self):\n \"\"\"Default (small) config should have ~4-5M params.\"\"\"\n model = SegFormer3D(n_classes=50)\n n_params = sum(p.numel() for p in model.parameters())\n assert n_params < 10_000_000 # < 10M\n\n def test_tiny_param_count(self):\n \"\"\"Tiny config should have ~1-2M params.\"\"\"\n model = SegFormer3D(n_classes=50, embed_dims=(16, 32, 80, 128))\n n_params = sum(p.numel() for p in model.parameters())\n assert n_params < 5_000_000 # < 5M\n\n def test_base_param_count(self):\n \"\"\"Base config should have ~15-20M params.\"\"\"\n model = SegFormer3D(n_classes=50, embed_dims=(64, 128, 320, 512))\n n_params = sum(p.numel() for p in model.parameters())\n assert n_params > 10_000_000 # > 10M\n\n\nclass TestSegFormer3DRegistry:\n def test_accessible_via_get(self):\n model = get(\"segformer3d\")(n_classes=2, embed_dims=(16, 32, 80, 128))\n assert model is not None\n assert isinstance(model, SegFormer3D)\n\n def test_in_available_models(self):\n from nobrainer.models import available_models\n\n assert \"segformer3d\" in available_models()\n\n def test_factory_defaults(self):\n from nobrainer.models.segformer3d import segformer3d\n\n model = segformer3d(n_classes=2)\n assert isinstance(model, SegFormer3D)\n" }, { "path": "nobrainer/tests/unit/test_slurm.py", "content": "\"\"\"Unit tests for nobrainer.slurm utilities.\"\"\"\n\nfrom __future__ import annotations\n\nimport torch\n\nfrom nobrainer.slurm import SlurmPreemptionHandler, load_checkpoint, save_checkpoint\n\n\nclass TestSlurmPreemptionHandler:\n def test_initial_state(self):\n h = SlurmPreemptionHandler()\n assert h.preempted is False\n\n def test_is_slurm_job(self):\n # In test environment, should be False\n assert isinstance(SlurmPreemptionHandler.is_slurm_job(), bool)\n\n\nclass TestCheckpoint:\n def test_save_and_load(self, tmp_path):\n model = torch.nn.Linear(4, 2)\n opt = torch.optim.SGD(model.parameters(), lr=0.01)\n\n metrics = {\"train_losses\": [1.0, 0.5], \"best_loss\": 0.5}\n save_checkpoint(tmp_path, model, opt, epoch=5, metrics=metrics)\n\n assert (tmp_path / \"checkpoint.pt\").exists()\n assert (tmp_path / \"checkpoint_meta.json\").exists()\n\n model2 = torch.nn.Linear(4, 2)\n opt2 = torch.optim.SGD(model2.parameters(), lr=0.01)\n start, restored = load_checkpoint(tmp_path, model2, opt2)\n\n assert start == 6 # next epoch\n assert restored[\"best_loss\"] == 0.5\n assert len(restored[\"train_losses\"]) == 2\n\n def test_load_no_checkpoint(self, tmp_path):\n model = torch.nn.Linear(4, 2)\n start, metrics = load_checkpoint(tmp_path, model)\n assert start == 0\n assert metrics == {}\n\n def test_model_weights_restored(self, tmp_path):\n model = torch.nn.Linear(4, 2)\n model.weight.data.fill_(42.0)\n opt = torch.optim.SGD(model.parameters(), lr=0.01)\n save_checkpoint(tmp_path, model, opt, epoch=0)\n\n model2 = torch.nn.Linear(4, 2)\n load_checkpoint(tmp_path, model2)\n assert torch.allclose(model2.weight.data, torch.tensor(42.0))\n" }, { "path": "nobrainer/tests/unit/test_stride_patches.py", "content": "\"\"\"Unit tests for strided patch extraction and reassembly.\"\"\"\n\nfrom __future__ import annotations\n\nimport numpy as np\n\nfrom nobrainer.prediction import reassemble_predictions, strided_patch_positions\n\n\nclass TestStridedPatchPositions:\n def test_non_overlapping_count(self):\n \"\"\"256³ with block=32 stride=32 → 8³ = 512 patches.\"\"\"\n pos = strided_patch_positions((256, 256, 256), (32, 32, 32), (32, 32, 32))\n assert len(pos) == 8 * 8 * 8 # 512\n\n def test_overlapping_more_patches(self):\n \"\"\"Stride < block produces more patches.\"\"\"\n non_overlap = strided_patch_positions((64, 64, 64), (32, 32, 32), (32, 32, 32))\n overlap = strided_patch_positions((64, 64, 64), (32, 32, 32), (16, 16, 16))\n assert len(overlap) > len(non_overlap)\n\n def test_patch_shapes_valid(self):\n \"\"\"Each position should yield a valid slice.\"\"\"\n pos = strided_patch_positions((100, 100, 100), (32, 32, 32), (16, 16, 16))\n for sd, sh, sw in pos:\n assert sd.stop - sd.start == 32\n assert sh.stop - sh.start == 32\n assert sw.stop - sw.start == 32\n assert sd.stop <= 100\n assert sh.stop <= 100\n assert sw.stop <= 100\n\n def test_stride_equals_block_default(self):\n \"\"\"None stride defaults to block_shape.\"\"\"\n pos = strided_patch_positions((64, 64, 64), (32, 32, 32))\n assert len(pos) == 2 * 2 * 2 # 8\n\n\nclass TestReassemblePredictions:\n def test_non_overlapping_perfect_reconstruction(self):\n \"\"\"Non-overlapping patches reassemble perfectly.\"\"\"\n vol_shape = (64, 64, 64)\n block = (32, 32, 32)\n n_classes = 2\n\n # Create a known volume\n original = np.random.randn(n_classes, *vol_shape).astype(np.float32)\n\n # Extract non-overlapping patches\n positions = strided_patch_positions(vol_shape, block, block)\n patches = []\n for sd, sh, sw in positions:\n patch = original[:, sd, sh, sw]\n patches.append((patch, (sd, sh, sw)))\n\n # Reassemble\n result = reassemble_predictions(patches, vol_shape, n_classes)\n assert np.allclose(result, original, atol=1e-6)\n\n def test_overlapping_average(self):\n \"\"\"Overlapping patches with averaging should still reconstruct reasonably.\"\"\"\n vol_shape = (64, 64, 64)\n block = (32, 32, 32)\n stride = (16, 16, 16)\n n_classes = 2\n\n # Create constant volume (averaging constant = constant)\n original = np.ones((n_classes, *vol_shape), dtype=np.float32) * 0.5\n\n positions = strided_patch_positions(vol_shape, block, stride)\n patches = []\n for sd, sh, sw in positions:\n patch = original[:, sd, sh, sw]\n patches.append((patch, (sd, sh, sw)))\n\n result = reassemble_predictions(\n patches, vol_shape, n_classes, strategy=\"average\"\n )\n assert np.allclose(result, 0.5, atol=1e-5)\n\n def test_output_shape(self):\n \"\"\"Output shape matches volume_shape.\"\"\"\n patches = [\n (np.ones((3, 16, 16, 16)), (slice(0, 16), slice(0, 16), slice(0, 16)))\n ]\n result = reassemble_predictions(patches, (32, 32, 32), 3)\n assert result.shape == (3, 32, 32, 32)\n" }, { "path": "nobrainer/tests/unit/test_synthseg.py", "content": "\"\"\"Unit tests for enhanced SynthSeg generator.\"\"\"\n\nfrom __future__ import annotations\n\nfrom pathlib import Path\n\nimport nibabel as nib\nimport numpy as np\nimport pytest\nimport torch\n\n\ndef _make_label_map(tmp_path: Path, shape=(32, 32, 32)) -> str:\n \"\"\"Create a simple label map with a few regions.\"\"\"\n arr = np.zeros(shape, dtype=np.int32)\n # Background = 0, WM = 2, GM = 3, CSF = 4, hippocampus L = 17, R = 53\n arr[4:28, 4:28, 4:28] = 2 # WM core\n arr[6:26, 6:26, 6:26] = 3 # GM shell\n arr[12:20, 12:20, 12:20] = 4 # CSF center\n arr[8:12, 8:12, 8:16] = 17 # L hippocampus\n arr[8:12, 8:12, 16:24] = 53 # R hippocampus\n path = str(tmp_path / \"label.nii.gz\")\n nib.save(nib.Nifti1Image(arr, np.eye(4)), path)\n return path\n\n\nclass TestTissueClasses:\n def test_all_50class_labels_covered(self):\n from nobrainer.data.tissue_classes import FREESURFER_TISSUE_CLASSES\n\n all_ids = set()\n for ids in FREESURFER_TISSUE_CLASSES.values():\n all_ids.update(ids)\n # Should cover background + major structures\n assert 0 in all_ids # background\n assert 2 in all_ids # L cerebral WM\n assert 41 in all_ids # R cerebral WM\n assert 17 in all_ids # L hippocampus\n assert 53 in all_ids # R hippocampus\n\n def test_no_label_in_multiple_classes(self):\n from nobrainer.data.tissue_classes import FREESURFER_TISSUE_CLASSES\n\n seen = {}\n for cls_name, ids in FREESURFER_TISSUE_CLASSES.items():\n for lid in ids:\n assert (\n lid not in seen\n ), f\"Label {lid} in both '{seen[lid]}' and '{cls_name}'\"\n seen[lid] = cls_name\n\n\nclass TestGMMGrouping:\n def test_within_class_same_distribution(self, tmp_path):\n from nobrainer.augmentation.synthseg import SynthSegGenerator\n\n path = _make_label_map(tmp_path)\n gen = SynthSegGenerator(\n [path],\n n_samples_per_map=1,\n elastic_std=0,\n rotation_range=0,\n flipping=False,\n randomize_resolution=False,\n noise_std=0,\n bias_field_std=0,\n )\n sample = gen[0]\n image = sample[\"image\"][0].numpy() # (D, H, W)\n label = sample[\"label\"][0].numpy()\n\n # L hippocampus (17) and R hippocampus (53) are both in \"hippocampus\" class\n # They should have similar mean intensities (same GMM class)\n l_hip = image[label == 17]\n r_hip = image[label == 53]\n if len(l_hip) > 0 and len(r_hip) > 0:\n # Both drawn from same distribution — means should be close\n mean_diff = abs(l_hip.mean() - r_hip.mean())\n pooled_std = max(l_hip.std(), r_hip.std(), 1e-6)\n cv = mean_diff / pooled_std\n assert cv < 0.5 # within-class similarity\n\n def test_different_classes_differ(self, tmp_path):\n from nobrainer.augmentation.synthseg import SynthSegGenerator\n\n path = _make_label_map(tmp_path)\n gen = SynthSegGenerator(\n [path],\n n_samples_per_map=1,\n elastic_std=0,\n rotation_range=0,\n flipping=False,\n randomize_resolution=False,\n noise_std=0,\n bias_field_std=0,\n )\n sample = gen[0]\n image = sample[\"image\"][0].numpy()\n label = sample[\"label\"][0].numpy()\n\n wm = image[label == 2]\n csf = image[label == 4]\n # WM and CSF should have different distributions (different classes)\n if len(wm) > 10 and len(csf) > 10:\n # Not guaranteed to differ every time but very likely\n assert wm.mean() != pytest.approx(csf.mean(), abs=1.0)\n\n def test_two_runs_produce_different_intensities(self, tmp_path):\n from nobrainer.augmentation.synthseg import SynthSegGenerator\n\n path = _make_label_map(tmp_path)\n gen = SynthSegGenerator(\n [path],\n n_samples_per_map=2,\n elastic_std=0,\n rotation_range=0,\n flipping=False,\n randomize_resolution=False,\n )\n s1 = gen[0][\"image\"]\n s2 = gen[1][\"image\"]\n assert not torch.allclose(s1, s2)\n\n\nclass TestSpatialAugmentation:\n def test_elastic_changes_geometry(self, tmp_path):\n from nobrainer.augmentation.synthseg import SynthSegGenerator\n\n path = _make_label_map(tmp_path)\n gen = SynthSegGenerator(\n [path],\n n_samples_per_map=1,\n elastic_std=4.0,\n rotation_range=0,\n flipping=False,\n randomize_resolution=False,\n noise_std=0,\n bias_field_std=0,\n )\n sample = gen[0]\n label = sample[\"label\"][0].numpy()\n\n # Load original label for comparison\n orig = np.asarray(nib.load(path).dataobj, dtype=np.int32)\n\n # Elastic deformation should change some voxel positions\n changed = (label != orig).sum()\n total = orig.size\n assert changed / total > 0.01 # at least 1% changed\n\n def test_label_nearest_neighbor(self, tmp_path):\n \"\"\"Labels should remain integer-valued after spatial augmentation.\"\"\"\n from nobrainer.augmentation.synthseg import SynthSegGenerator\n from nobrainer.data.tissue_classes import FREESURFER_LR_PAIRS\n\n path = _make_label_map(tmp_path)\n gen = SynthSegGenerator(\n [path],\n n_samples_per_map=1,\n elastic_std=4.0,\n rotation_range=15.0,\n randomize_resolution=False,\n )\n sample = gen[0]\n label = sample[\"label\"][0].numpy()\n\n # All values should be valid integers (no interpolation artifacts)\n # Include L/R swapped labels since flipping may have occurred\n orig_labels = set(np.asarray(nib.load(path).dataobj, dtype=np.int32).flat)\n valid_labels = set(orig_labels)\n for left, right in FREESURFER_LR_PAIRS:\n if left in orig_labels:\n valid_labels.add(right)\n if right in orig_labels:\n valid_labels.add(left)\n actual_labels = set(label.flat)\n assert actual_labels.issubset(valid_labels)\n\n def test_flipping_swaps_lr(self, tmp_path):\n \"\"\"Flipping should swap L/R FreeSurfer codes.\"\"\"\n from nobrainer.augmentation.synthseg import SynthSegGenerator\n\n path = _make_label_map(tmp_path)\n gen = SynthSegGenerator(\n [path],\n n_samples_per_map=30,\n elastic_std=0,\n rotation_range=0,\n flipping=True,\n randomize_resolution=False,\n noise_std=0,\n bias_field_std=0,\n )\n # After L/R flip, label 17 (L hippocampus) should become 53 and vice versa\n # Check that at least one sample has the swap in the label set\n found_swap = False\n orig = np.asarray(nib.load(path).dataobj, dtype=np.int32)\n for i in range(30):\n label = gen[i][\"label\"][0].numpy()\n # A flip swaps L/R labels AND mirrors spatially.\n # If the spatial distribution of label 17 differs from original, flip happened\n orig_17_count = (orig == 17).sum()\n new_17_count = (label == 17).sum()\n if orig_17_count > 0 and new_17_count != orig_17_count:\n found_swap = True\n break\n assert found_swap\n\n\nclass TestResolutionRandomization:\n def test_blurs_image(self, tmp_path):\n from nobrainer.augmentation.synthseg import SynthSegGenerator\n\n path = _make_label_map(tmp_path)\n gen_sharp = SynthSegGenerator(\n [path],\n n_samples_per_map=1,\n elastic_std=0,\n rotation_range=0,\n flipping=False,\n randomize_resolution=False,\n noise_std=0,\n bias_field_std=0,\n )\n gen_blur = SynthSegGenerator(\n [path],\n n_samples_per_map=1,\n elastic_std=0,\n rotation_range=0,\n flipping=False,\n randomize_resolution=True,\n resolution_range=(2.0, 3.0), # force heavy blur\n noise_std=0,\n bias_field_std=0,\n )\n # Use same seed for intensity but different resolution\n np.random.seed(42)\n sharp = gen_sharp[0][\"image\"][0].numpy()\n np.random.seed(42)\n blurred = gen_blur[0][\"image\"][0].numpy()\n\n # Blurred should have less high-frequency energy\n sharp_grad = np.abs(np.diff(sharp, axis=0)).mean()\n blur_grad = np.abs(np.diff(blurred, axis=0)).mean()\n assert blur_grad < sharp_grad\n\n\nclass TestOutputFormat:\n def test_returns_dict_with_correct_keys(self, tmp_path):\n from nobrainer.augmentation.synthseg import SynthSegGenerator\n\n path = _make_label_map(tmp_path)\n gen = SynthSegGenerator([path], n_samples_per_map=1)\n sample = gen[0]\n assert \"image\" in sample\n assert \"label\" in sample\n assert sample[\"image\"].shape[0] == 1 # channel dim\n assert sample[\"label\"].shape[0] == 1\n assert sample[\"image\"].dtype == torch.float32\n assert sample[\"label\"].dtype == torch.int64\n\n def test_correct_length(self, tmp_path):\n from nobrainer.augmentation.synthseg import SynthSegGenerator\n\n path = _make_label_map(tmp_path)\n gen = SynthSegGenerator([path, path], n_samples_per_map=5)\n assert len(gen) == 10\n\n\nclass TestMixedDataset:\n def test_mix_ratio(self, tmp_path):\n \"\"\"Mixed dataset produces approximately correct ratio.\"\"\"\n from nobrainer.augmentation.synthseg import SynthSegGenerator\n from nobrainer.processing.dataset import MixedDataset\n\n path = _make_label_map(tmp_path)\n gen = SynthSegGenerator(\n [path],\n n_samples_per_map=50,\n elastic_std=0,\n rotation_range=0,\n flipping=False,\n randomize_resolution=False,\n noise_std=0,\n bias_field_std=0,\n )\n\n # Create a simple \"real\" dataset\n real = gen # reuse generator as real for simplicity\n mixed = MixedDataset(real, gen, ratio=0.5)\n\n assert len(mixed) == 50\n # Just verify it returns dicts without error\n sample = mixed[0]\n assert \"image\" in sample or isinstance(sample, dict)\n\n def test_dataset_mix_method(self, tmp_path):\n \"\"\"Dataset.mix() returns a Dataset with _mixed_dataset set.\"\"\"\n from nobrainer.augmentation.synthseg import SynthSegGenerator\n from nobrainer.processing.dataset import Dataset\n\n path = _make_label_map(tmp_path)\n gen = SynthSegGenerator(\n [path],\n n_samples_per_map=5,\n elastic_std=0,\n rotation_range=0,\n flipping=False,\n randomize_resolution=False,\n )\n\n pairs = [(str(tmp_path / \"label.nii.gz\"), str(tmp_path / \"label.nii.gz\"))]\n ds = Dataset.from_files(pairs, block_shape=(16, 16, 16), n_classes=2)\n mixed = ds.mix(gen, ratio=0.3)\n\n assert hasattr(mixed, \"_mixed_dataset\")\n" }, { "path": "nobrainer/tests/unit/test_training.py", "content": "\"\"\"Unit tests for nobrainer.training.fit().\"\"\"\n\nfrom __future__ import annotations\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader, TensorDataset\n\nfrom nobrainer.training import fit\n\n\ndef _make_loader(n=8, spatial=8, n_classes=2, batch_size=2):\n \"\"\"Synthetic DataLoader for training tests.\"\"\"\n x = torch.randn(n, 1, spatial, spatial, spatial)\n y = torch.randint(0, n_classes, (n, spatial, spatial, spatial))\n ds = TensorDataset(x, y)\n return DataLoader(ds, batch_size=batch_size)\n\n\ndef _make_model(n_classes=2):\n \"\"\"Tiny conv model for testing.\"\"\"\n return nn.Sequential(\n nn.Conv3d(1, 8, 3, padding=1),\n nn.ReLU(),\n nn.Conv3d(8, n_classes, 1),\n )\n\n\nclass TestFit:\n def test_returns_correct_keys(self):\n model = _make_model()\n loader = _make_loader()\n result = fit(\n model,\n loader,\n nn.CrossEntropyLoss(),\n torch.optim.Adam(model.parameters()),\n max_epochs=2,\n )\n assert \"history\" in result\n assert \"checkpoint_path\" in result\n assert len(result[\"history\"]) == 2\n\n def test_loss_decreases(self):\n torch.manual_seed(42)\n model = _make_model()\n loader = _make_loader()\n result = fit(\n model,\n loader,\n nn.CrossEntropyLoss(),\n torch.optim.Adam(model.parameters(), lr=1e-2),\n max_epochs=10,\n )\n losses = [h[\"loss\"] for h in result[\"history\"]]\n assert (\n losses[-1] < losses[0]\n ), f\"Loss did not decrease: first={losses[0]:.4f}, last={losses[-1]:.4f}\"\n\n def test_checkpoint_created(self, tmp_path):\n model = _make_model()\n loader = _make_loader()\n result = fit(\n model,\n loader,\n nn.CrossEntropyLoss(),\n torch.optim.Adam(model.parameters()),\n max_epochs=2,\n checkpoint_dir=tmp_path,\n )\n assert result[\"checkpoint_path\"] is not None\n assert (tmp_path / \"best_model.pth\").exists()\n assert (tmp_path / \"croissant.json\").exists()\n\n def test_checkpoint_croissant_content(self, tmp_path):\n \"\"\"Checkpoint croissant.json contains provenance metadata.\"\"\"\n import json\n\n model = _make_model()\n loader = _make_loader()\n fit(\n model,\n loader,\n nn.CrossEntropyLoss(),\n torch.optim.Adam(model.parameters()),\n max_epochs=2,\n checkpoint_dir=tmp_path,\n )\n data = json.loads((tmp_path / \"croissant.json\").read_text())\n prov = data[\"nobrainer:provenance\"]\n assert prov[\"epochs_trained\"] > 0\n assert prov[\"model_architecture\"] == \"Sequential\"\n assert prov[\"loss_function\"] == \"CrossEntropyLoss\"\n assert \"optimizer\" in prov\n\n def test_epochs_completed(self):\n model = _make_model()\n loader = _make_loader()\n result = fit(\n model,\n loader,\n nn.CrossEntropyLoss(),\n torch.optim.Adam(model.parameters()),\n max_epochs=3,\n )\n assert len(result[\"history\"]) == 3\n\n def test_dict_batch_format(self):\n \"\"\"fit() works with dict-style batches (from MONAI DataLoader).\"\"\"\n x = torch.randn(4, 1, 8, 8, 8)\n y = torch.randint(0, 2, (4, 8, 8, 8))\n\n class DictDataset(torch.utils.data.Dataset):\n def __len__(self):\n return 4\n\n def __getitem__(self, idx):\n return {\"image\": x[idx], \"label\": y[idx]}\n\n loader = DataLoader(DictDataset(), batch_size=2)\n model = _make_model()\n result = fit(\n model,\n loader,\n nn.CrossEntropyLoss(),\n torch.optim.Adam(model.parameters()),\n max_epochs=1,\n )\n assert len(result[\"history\"]) == 1\n" }, { "path": "nobrainer/tests/unit/test_training_convergence.py", "content": "\"\"\"CPU training-convergence smoke tests (US1 acceptance scenario 3).\n\nVerifies that each core segmentation model's training loss at epoch 5 is\nlower than at epoch 1 when overfitting a fixed batch on CPU.\n\nScope: tests nobrainer.models + nobrainer.losses integration; does NOT\nrequire GPU or real data.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport torch\n\nfrom nobrainer.losses import dice\nfrom nobrainer.models.highresnet import highresnet\nfrom nobrainer.models.meshnet import meshnet\nfrom nobrainer.models.segmentation import attention_unet, unet, vnet\n\n# Shared synthetic batch: batch_size=2 (satisfies BatchNorm), 32^3 spatial.\n# Fixed all-ones label is easy to overfit, keeping the test deterministic.\n_SPATIAL = 32\n_N_EPOCHS = 5\n_LR = 1e-2\n\n\ndef _run_epochs(model: torch.nn.Module, seed: int = 42) -> list[float]:\n \"\"\"Train *model* for _N_EPOCHS and return per-epoch loss values.\"\"\"\n torch.manual_seed(seed)\n model.train()\n x = torch.randn(2, 1, _SPATIAL, _SPATIAL, _SPATIAL)\n y = torch.ones(2, 1, _SPATIAL, _SPATIAL, _SPATIAL)\n loss_fn = dice()\n opt = torch.optim.Adam(model.parameters(), lr=_LR)\n losses = []\n for _ in range(_N_EPOCHS):\n opt.zero_grad()\n pred = model(x)\n loss = loss_fn(pred, y)\n loss.backward()\n opt.step()\n losses.append(loss.item())\n return losses\n\n\nclass TestTrainingConvergence:\n \"\"\"US1 scenario 3: loss at epoch 5 < loss at epoch 1 for all core models.\"\"\"\n\n def test_unet_loss_decreases(self):\n losses = _run_epochs(unet(n_classes=1))\n assert (\n losses[-1] < losses[0]\n ), f\"UNet loss did not decrease: epoch1={losses[0]:.4f}, epoch5={losses[-1]:.4f}\"\n\n def test_vnet_loss_decreases(self):\n losses = _run_epochs(vnet(n_classes=1))\n assert (\n losses[-1] < losses[0]\n ), f\"VNet loss did not decrease: epoch1={losses[0]:.4f}, epoch5={losses[-1]:.4f}\"\n\n def test_attention_unet_loss_decreases(self):\n losses = _run_epochs(attention_unet(n_classes=1))\n assert (\n losses[-1] < losses[0]\n ), f\"AttentionUNet loss did not decrease: epoch1={losses[0]:.4f}, epoch5={losses[-1]:.4f}\"\n\n def test_meshnet_loss_decreases(self):\n losses = _run_epochs(meshnet(n_classes=1))\n assert (\n losses[-1] < losses[0]\n ), f\"MeshNet loss did not decrease: epoch1={losses[0]:.4f}, epoch5={losses[-1]:.4f}\"\n\n def test_highresnet_loss_decreases(self):\n losses = _run_epochs(highresnet(n_classes=1))\n assert (\n losses[-1] < losses[0]\n ), f\"HighResNet loss did not decrease: epoch1={losses[0]:.4f}, epoch5={losses[-1]:.4f}\"\n" }, { "path": "nobrainer/tests/unit/test_transform_pipeline.py", "content": "\"\"\"Unit tests for TrainableCompose and Augmentation tagging.\"\"\"\n\nfrom __future__ import annotations\n\nfrom nobrainer.augmentation.transforms import Augmentation, TrainableCompose\n\n\ndef _identity(data):\n \"\"\"Preprocessing transform that passes data through.\"\"\"\n data[\"preprocess_count\"] = data.get(\"preprocess_count\", 0) + 1\n return data\n\n\ndef _augment(data):\n \"\"\"Augmentation transform that modifies data.\"\"\"\n data[\"augment_count\"] = data.get(\"augment_count\", 0) + 1\n return data\n\n\nclass TestAugmentation:\n def test_wraps_transform(self):\n aug = Augmentation(_augment)\n assert aug.is_augmentation is True\n result = aug({\"x\": 1})\n assert result[\"augment_count\"] == 1\n\n def test_repr(self):\n aug = Augmentation(_augment)\n assert \"Augmentation\" in repr(aug)\n\n\nclass TestTrainableCompose:\n def test_train_mode_runs_all(self):\n pipeline = TrainableCompose([_identity, Augmentation(_augment)])\n result = pipeline({\"x\": 1}, mode=\"train\")\n assert result[\"preprocess_count\"] == 1\n assert result[\"augment_count\"] == 1\n\n def test_predict_mode_skips_augmentation(self):\n pipeline = TrainableCompose([_identity, Augmentation(_augment)])\n result = pipeline({\"x\": 1}, mode=\"predict\")\n assert result[\"preprocess_count\"] == 1\n assert \"augment_count\" not in result\n\n def test_default_mode_is_train(self):\n pipeline = TrainableCompose([_identity, Augmentation(_augment)])\n result = pipeline({\"x\": 1})\n assert result[\"augment_count\"] == 1\n\n def test_mode_setter(self):\n pipeline = TrainableCompose([_identity, Augmentation(_augment)])\n pipeline.mode = \"predict\"\n result = pipeline({\"x\": 1})\n assert \"augment_count\" not in result\n\n def test_multiple_augmentations_skipped(self):\n pipeline = TrainableCompose(\n [\n _identity,\n Augmentation(_augment),\n _identity,\n Augmentation(_augment),\n ]\n )\n result = pipeline({\"x\": 1}, mode=\"predict\")\n assert result[\"preprocess_count\"] == 2\n assert \"augment_count\" not in result\n\n def test_train_mode_runs_multiple_augmentations(self):\n pipeline = TrainableCompose(\n [\n _identity,\n Augmentation(_augment),\n _identity,\n Augmentation(_augment),\n ]\n )\n result = pipeline({\"x\": 1}, mode=\"train\")\n assert result[\"preprocess_count\"] == 2\n assert result[\"augment_count\"] == 2\n\n def test_empty_pipeline(self):\n pipeline = TrainableCompose([])\n result = pipeline({\"x\": 1}, mode=\"train\")\n assert result == {\"x\": 1}\n\n\nclass TestAugmentationProfiles:\n def test_none_returns_empty(self):\n from nobrainer.augmentation.profiles import get_augmentation_profile\n\n transforms = get_augmentation_profile(\"none\")\n assert transforms == []\n\n def test_standard_returns_augmentations(self):\n from nobrainer.augmentation.profiles import get_augmentation_profile\n\n transforms = get_augmentation_profile(\"standard\")\n assert len(transforms) > 0\n assert all(getattr(t, \"is_augmentation\", False) for t in transforms)\n\n def test_all_profiles_valid(self):\n from nobrainer.augmentation.profiles import get_augmentation_profile\n\n for name in (\"none\", \"light\", \"standard\", \"heavy\"):\n transforms = get_augmentation_profile(name)\n assert isinstance(transforms, list)\n\n def test_unknown_profile_raises(self):\n import pytest\n\n from nobrainer.augmentation.profiles import get_augmentation_profile\n\n with pytest.raises(ValueError, match=\"Unknown augmentation profile\"):\n get_augmentation_profile(\"extreme\")\n" }, { "path": "nobrainer/tests/unit/test_vwn_layers.py", "content": "\"\"\"Unit tests for VWN layers and KWYKMeshNet.\"\"\"\n\nfrom __future__ import annotations\n\nimport torch\n\nfrom nobrainer.models.bayesian.vwn_layers import ConcreteDropout3d, FFGConv3d\n\n\nclass TestFFGConv3d:\n def test_output_shape(self):\n layer = FFGConv3d(1, 4, kernel_size=3, padding=1)\n x = torch.randn(2, 1, 8, 8, 8)\n out = layer(x, mc=True)\n assert out.shape == (2, 4, 8, 8, 8)\n\n def test_deterministic_mode(self):\n layer = FFGConv3d(1, 4, kernel_size=3, padding=1)\n x = torch.randn(2, 1, 8, 8, 8)\n layer.eval()\n out1 = layer(x, mc=False)\n out2 = layer(x, mc=False)\n assert torch.allclose(out1, out2)\n\n def test_stochastic_mode_varies(self):\n layer = FFGConv3d(1, 4, kernel_size=3, padding=1)\n x = torch.randn(2, 1, 8, 8, 8)\n out1 = layer(x, mc=True)\n out2 = layer(x, mc=True)\n # Outputs should differ due to stochastic sampling\n assert not torch.allclose(out1, out2)\n\n def test_kl_populated_after_mc(self):\n layer = FFGConv3d(1, 4, kernel_size=3, padding=1)\n x = torch.randn(2, 1, 8, 8, 8)\n layer(x, mc=True)\n assert layer.kl.item() > 0\n\n def test_kernel_m_shape(self):\n layer = FFGConv3d(1, 4, kernel_size=3, padding=1)\n assert layer.kernel_m.shape == (4, 1, 3, 3, 3)\n\n def test_no_bias(self):\n layer = FFGConv3d(1, 4, kernel_size=3, padding=1, bias=False)\n assert layer.bias_m is None\n x = torch.randn(2, 1, 8, 8, 8)\n out = layer(x, mc=True)\n assert out.shape == (2, 4, 8, 8, 8)\n\n def test_sigma_positive(self):\n layer = FFGConv3d(1, 4, kernel_size=3, padding=1)\n assert (layer.weight_sigma >= 0).all()\n\n\nclass TestConcreteDropout3d:\n def test_output_shape(self):\n cd = ConcreteDropout3d(4)\n x = torch.randn(2, 4, 8, 8, 8)\n out = cd(x, mc=True)\n assert out.shape == x.shape\n\n def test_deterministic_scales(self):\n cd = ConcreteDropout3d(4)\n x = torch.ones(2, 4, 8, 8, 8)\n out = cd(x, mc=False)\n # In deterministic mode, output = x * p\n p = cd.p.view(1, -1, 1, 1, 1)\n expected = x * p\n assert torch.allclose(out, expected)\n\n def test_p_in_range(self):\n cd = ConcreteDropout3d(4)\n assert (cd.p >= 0.05).all()\n assert (cd.p <= 0.95).all()\n\n def test_regularization_positive(self):\n cd = ConcreteDropout3d(4)\n reg = cd.regularization()\n assert reg.item() > 0\n\n\nclass TestKWYKMeshNet:\n def test_bernoulli_variant(self):\n from nobrainer.models.bayesian.kwyk_meshnet import KWYKMeshNet\n\n model = KWYKMeshNet(\n n_classes=2,\n filters=8,\n receptive_field=37,\n dropout_type=\"bernoulli\",\n )\n x = torch.randn(1, 1, 16, 16, 16)\n out = model(x, mc=True)\n assert out.shape == (1, 2, 16, 16, 16)\n\n def test_concrete_variant(self):\n from nobrainer.models.bayesian.kwyk_meshnet import KWYKMeshNet\n\n model = KWYKMeshNet(\n n_classes=2,\n filters=8,\n receptive_field=37,\n dropout_type=\"concrete\",\n )\n x = torch.randn(1, 1, 16, 16, 16)\n out = model(x, mc=True)\n assert out.shape == (1, 2, 16, 16, 16)\n\n def test_kl_divergence(self):\n from nobrainer.models.bayesian.kwyk_meshnet import KWYKMeshNet\n\n model = KWYKMeshNet(n_classes=2, filters=8, receptive_field=37)\n x = torch.randn(1, 1, 16, 16, 16)\n model(x, mc=True)\n kl = model.kl_divergence()\n assert torch.isfinite(kl)\n\n def test_concrete_regularization(self):\n from nobrainer.models.bayesian.kwyk_meshnet import KWYKMeshNet\n\n model = KWYKMeshNet(\n n_classes=2,\n filters=8,\n receptive_field=37,\n dropout_type=\"concrete\",\n )\n reg = model.concrete_regularization()\n assert reg.item() > 0\n\n def test_deterministic_forward(self):\n from nobrainer.models.bayesian.kwyk_meshnet import KWYKMeshNet\n\n model = KWYKMeshNet(n_classes=2, filters=8, receptive_field=37)\n x = torch.randn(1, 1, 16, 16, 16)\n model.eval()\n out1 = model(x, mc=False)\n out2 = model(x, mc=False)\n assert torch.allclose(out1, out2)\n\n def test_factory_function(self):\n from nobrainer.models import get\n\n model = get(\"kwyk_meshnet\")(n_classes=2, filters=8, receptive_field=37)\n x = torch.randn(1, 1, 16, 16, 16)\n out = model(x)\n assert out.shape == (1, 2, 16, 16, 16)\n" }, { "path": "nobrainer/tests/unit/test_zarr_dataset.py", "content": "\"\"\"Unit tests for ZarrDataset and get_dataset() Zarr routing.\"\"\"\n\nfrom __future__ import annotations\n\nimport nibabel as nib\nimport numpy as np\nimport pytest\nimport torch\n\nzarr = pytest.importorskip(\"zarr\", reason=\"zarr not installed\")\n\nfrom nobrainer.dataset import ZarrDataset, _is_zarr_path # noqa: E402\nfrom nobrainer.io import nifti_to_zarr # noqa: E402\n\n\ndef _make_zarr_pair(tmp_path, shape=(32, 32, 32)):\n \"\"\"Create a synthetic NIfTI → Zarr pair (image + label).\"\"\"\n img_data = np.random.rand(*shape).astype(np.float32)\n lbl_data = (np.random.rand(*shape) > 0.5).astype(np.float32)\n\n img_nii = tmp_path / \"img.nii.gz\"\n lbl_nii = tmp_path / \"lbl.nii.gz\"\n nib.save(nib.Nifti1Image(img_data, np.eye(4)), str(img_nii))\n nib.save(nib.Nifti1Image(lbl_data, np.eye(4)), str(lbl_nii))\n\n img_zarr = nifti_to_zarr(img_nii, tmp_path / \"img.zarr\")\n lbl_zarr = nifti_to_zarr(lbl_nii, tmp_path / \"lbl.zarr\")\n\n return img_zarr, lbl_zarr, img_data, lbl_data\n\n\nclass TestIsZarrPath:\n def test_zarr_extension(self):\n assert _is_zarr_path(\"data/brain.zarr\")\n assert _is_zarr_path(\"data/brain.zarr/\")\n\n def test_non_zarr(self):\n assert not _is_zarr_path(\"data/brain.nii.gz\")\n assert not _is_zarr_path(\"data/brain.h5\")\n\n\nclass TestZarrDataset:\n def test_returns_dict_with_image(self, tmp_path):\n img_zarr, _, _, _ = _make_zarr_pair(tmp_path)\n ds = ZarrDataset([{\"image\": str(img_zarr)}])\n item = ds[0]\n assert \"image\" in item\n assert isinstance(item[\"image\"], torch.Tensor)\n\n def test_image_shape_has_channel(self, tmp_path):\n img_zarr, _, img_data, _ = _make_zarr_pair(tmp_path)\n ds = ZarrDataset([{\"image\": str(img_zarr)}])\n item = ds[0]\n # Should have channel dim: (1, D, H, W)\n assert item[\"image\"].shape == (1, *img_data.shape)\n\n def test_returns_label_when_provided(self, tmp_path):\n img_zarr, lbl_zarr, _, _ = _make_zarr_pair(tmp_path)\n ds = ZarrDataset([{\"image\": str(img_zarr), \"label\": str(lbl_zarr)}])\n item = ds[0]\n assert \"label\" in item\n assert isinstance(item[\"label\"], torch.Tensor)\n\n def test_batch_from_dataloader(self, tmp_path):\n img_zarr, lbl_zarr, _, _ = _make_zarr_pair(tmp_path)\n data = [{\"image\": str(img_zarr), \"label\": str(lbl_zarr)}]\n ds = ZarrDataset(data)\n loader = torch.utils.data.DataLoader(ds, batch_size=1)\n batch = next(iter(loader))\n assert batch[\"image\"].shape[0] == 1 # batch dim\n assert batch[\"image\"].ndim == 5 # (B, C, D, H, W)\n\n def test_multi_resolution_level(self, tmp_path):\n \"\"\"Loading at level 1 gives downsampled shape.\"\"\"\n img_data = np.random.rand(64, 64, 64).astype(np.float32)\n nii_path = tmp_path / \"big.nii.gz\"\n nib.save(nib.Nifti1Image(img_data, np.eye(4)), str(nii_path))\n zarr_path = nifti_to_zarr(nii_path, tmp_path / \"big.zarr\", levels=2)\n\n ds = ZarrDataset([{\"image\": str(zarr_path)}], zarr_level=1)\n item = ds[0]\n # Level 1 is 2x downsampled: (1, 32, 32, 32)\n assert item[\"image\"].shape == (1, 32, 32, 32)\n" }, { "path": "nobrainer/tests/unit/test_zarr_store.py", "content": "\"\"\"Unit tests for nobrainer.datasets.zarr_store.\"\"\"\n\nfrom __future__ import annotations\n\nimport json\n\nimport nibabel as nib\nimport numpy as np\nimport pytest\n\n\ndef _make_nifti_pair(tmp_path, idx, shape=(32, 32, 32)):\n \"\"\"Create a NIfTI image + label pair.\"\"\"\n img_data = np.random.randn(*shape).astype(np.float32)\n lbl_data = np.random.randint(0, 5, shape, dtype=np.int32)\n affine = np.eye(4)\n\n img_path = tmp_path / f\"sub-{idx:02d}_image.nii.gz\"\n lbl_path = tmp_path / f\"sub-{idx:02d}_label.nii.gz\"\n nib.save(nib.Nifti1Image(img_data, affine), str(img_path))\n nib.save(nib.Nifti1Image(lbl_data, affine), str(lbl_path))\n return str(img_path), str(lbl_path)\n\n\nclass TestCreateZarrStore:\n def test_creates_store(self, tmp_path):\n from nobrainer.datasets.zarr_store import create_zarr_store\n\n pairs = [_make_nifti_pair(tmp_path, i) for i in range(3)]\n store_path = create_zarr_store(\n pairs,\n tmp_path / \"test.zarr\",\n conform=False,\n )\n assert store_path.exists()\n\n def test_stacked_4d_layout(self, tmp_path):\n import zarr\n\n from nobrainer.datasets.zarr_store import create_zarr_store\n\n pairs = [_make_nifti_pair(tmp_path, i) for i in range(3)]\n store_path = create_zarr_store(\n pairs,\n tmp_path / \"test.zarr\",\n conform=False,\n )\n\n store = zarr.open_group(str(store_path), mode=\"r\")\n assert store[\"images\"].shape == (3, 32, 32, 32)\n assert store[\"labels\"].shape == (3, 32, 32, 32)\n assert store[\"images\"].dtype == np.float32\n assert store[\"labels\"].dtype == np.int32\n\n def test_metadata_stored(self, tmp_path):\n from nobrainer.datasets.zarr_store import create_zarr_store, store_info\n\n pairs = [_make_nifti_pair(tmp_path, i) for i in range(3)]\n store_path = create_zarr_store(\n pairs,\n tmp_path / \"test.zarr\",\n subject_ids=[\"sub-00\", \"sub-01\", \"sub-02\"],\n conform=False,\n )\n\n info = store_info(store_path)\n assert info[\"n_subjects\"] == 3\n assert info[\"subject_ids\"] == [\"sub-00\", \"sub-01\", \"sub-02\"]\n assert info[\"volume_shape\"] == [32, 32, 32]\n assert info[\"layout\"] == \"stacked\"\n\n def test_round_trip_fidelity(self, tmp_path):\n import zarr\n\n from nobrainer.datasets.zarr_store import create_zarr_store\n\n pairs = [_make_nifti_pair(tmp_path, i) for i in range(2)]\n store_path = create_zarr_store(pairs, tmp_path / \"test.zarr\", conform=False)\n\n # Read back and compare\n original = np.asarray(nib.load(pairs[0][0]).dataobj, dtype=np.float32)\n store = zarr.open_group(str(store_path), mode=\"r\")\n stored = np.array(store[\"images\"][0])\n assert np.allclose(original, stored, atol=1e-6)\n\n def test_partial_io(self, tmp_path):\n import zarr\n\n from nobrainer.datasets.zarr_store import create_zarr_store\n\n pairs = [_make_nifti_pair(tmp_path, i) for i in range(5)]\n store_path = create_zarr_store(pairs, tmp_path / \"test.zarr\", conform=False)\n\n store = zarr.open_group(str(store_path), mode=\"r\")\n # Read a subregion from subject 2\n patch = np.array(store[\"images\"][2, 8:24, 8:24, 8:24])\n assert patch.shape == (16, 16, 16)\n\n def test_auto_conform(self, tmp_path):\n import zarr\n\n from nobrainer.datasets.zarr_store import create_zarr_store\n\n # Create volumes with different shapes — conform should make them uniform\n # Use shapes where median is 32-divisible for sharding compat\n pairs = [\n _make_nifti_pair(tmp_path, 0, shape=(32, 32, 32)),\n _make_nifti_pair(tmp_path, 1, shape=(32, 32, 32)),\n _make_nifti_pair(tmp_path, 2, shape=(64, 64, 64)),\n ]\n store_path = create_zarr_store(\n pairs,\n tmp_path / \"test.zarr\",\n conform=True,\n )\n\n store = zarr.open_group(str(store_path), mode=\"r\")\n # All subjects should have same shape\n assert store[\"images\"].shape[1:] == store[\"images\"].shape[1:]\n info = dict(store.attrs)\n assert info[\"conformed\"] is True\n\n def test_non_uniform_without_conform_raises(self, tmp_path):\n from nobrainer.datasets.zarr_store import create_zarr_store\n\n pairs = [\n _make_nifti_pair(tmp_path, 0, shape=(32, 32, 32)),\n _make_nifti_pair(tmp_path, 1, shape=(64, 64, 64)),\n ]\n with pytest.raises(ValueError, match=\"Non-uniform shapes\"):\n create_zarr_store(pairs, tmp_path / \"test.zarr\", conform=False)\n\n\nclass TestPartition:\n def test_create_partition(self, tmp_path):\n from nobrainer.datasets.zarr_store import create_partition, create_zarr_store\n\n pairs = [_make_nifti_pair(tmp_path, i) for i in range(10)]\n store_path = create_zarr_store(pairs, tmp_path / \"test.zarr\", conform=False)\n part_path = create_partition(store_path, ratios=(80, 10, 10))\n\n assert part_path.exists()\n with open(part_path) as f:\n data = json.load(f)\n assert len(data[\"partitions\"][\"train\"]) == 8\n assert len(data[\"partitions\"][\"val\"]) == 1\n assert len(data[\"partitions\"][\"test\"]) == 1\n\n def test_load_partition(self, tmp_path):\n from nobrainer.datasets.zarr_store import (\n create_partition,\n create_zarr_store,\n load_partition,\n )\n\n pairs = [_make_nifti_pair(tmp_path, i) for i in range(10)]\n store_path = create_zarr_store(pairs, tmp_path / \"test.zarr\", conform=False)\n part_path = create_partition(store_path)\n\n partitions = load_partition(part_path)\n assert \"train\" in partitions\n assert \"val\" in partitions\n assert \"test\" in partitions\n all_ids = partitions[\"train\"] + partitions[\"val\"] + partitions[\"test\"]\n assert len(set(all_ids)) == 10 # no duplicates\n\n def test_different_seeds_produce_different_splits(self, tmp_path):\n from nobrainer.datasets.zarr_store import (\n create_partition,\n create_zarr_store,\n load_partition,\n )\n\n pairs = [_make_nifti_pair(tmp_path, i) for i in range(10)]\n store_path = create_zarr_store(pairs, tmp_path / \"test.zarr\", conform=False)\n\n p1 = load_partition(\n create_partition(store_path, seed=1, output_path=tmp_path / \"p1.json\")\n )\n p2 = load_partition(\n create_partition(store_path, seed=2, output_path=tmp_path / \"p2.json\")\n )\n # Different seeds should produce different train sets (with high probability)\n assert p1[\"train\"] != p2[\"train\"]\n" }, { "path": "nobrainer/training.py", "content": "\"\"\"Training utilities with optional multi-GPU DDP support.\"\"\"\n\nfrom __future__ import annotations\n\nimport logging\nfrom pathlib import Path\nfrom typing import Any\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader\n\nlogger = logging.getLogger(__name__)\n\n\ndef get_device() -> torch.device:\n \"\"\"Select the best available device: CUDA > MPS > CPU.\n\n .. note::\n Also available as :func:`nobrainer.gpu.get_device`.\n \"\"\"\n from nobrainer.gpu import get_device as _get_device\n\n return _get_device()\n\n\ndef _run_validation(\n model: nn.Module,\n val_loader: DataLoader,\n criterion: nn.Module,\n device: torch.device,\n) -> dict[str, float]:\n \"\"\"Run one validation pass.\n\n Returns dict with val_loss, val_acc (overall), and val_bal_acc\n (balanced accuracy — mean of per-class recall).\n \"\"\"\n model.eval()\n total_loss = 0.0\n n_correct = 0\n n_total = 0\n n_batches = 0\n # Per-class correct/total for balanced accuracy\n class_correct: dict[int, int] = {}\n class_total: dict[int, int] = {}\n\n with torch.no_grad():\n for batch in val_loader:\n if isinstance(batch, dict):\n images = batch[\"image\"].to(device)\n labels = batch[\"label\"].to(device)\n elif isinstance(batch, (list, tuple)):\n images = batch[0].to(device)\n labels = batch[1].to(device)\n else:\n continue\n\n if labels.ndim == images.ndim and labels.shape[1] == 1:\n labels = labels.squeeze(1)\n if labels.dtype in (torch.float32, torch.float64):\n labels = labels.long()\n\n pred = model(images)\n # Handle model parallel: pred may be on different device\n if pred.device != labels.device:\n labels = labels.to(pred.device)\n total_loss += criterion(pred, labels).item()\n\n pred_labels = pred.argmax(1)\n correct_mask = pred_labels == labels\n n_correct += correct_mask.sum().item()\n n_total += labels.numel()\n n_batches += 1\n\n # Accumulate per-class stats\n for c in labels.unique().tolist():\n mask = labels == c\n cc = correct_mask[mask].sum().item()\n ct = mask.sum().item()\n class_correct[c] = class_correct.get(c, 0) + cc\n class_total[c] = class_total.get(c, 0) + ct\n\n val_loss = total_loss / max(n_batches, 1)\n val_acc = n_correct / max(n_total, 1)\n\n # Balanced accuracy: mean recall per class\n per_class_recall = []\n for c in sorted(class_total.keys()):\n if class_total[c] > 0:\n per_class_recall.append(class_correct[c] / class_total[c])\n val_bal_acc = sum(per_class_recall) / max(len(per_class_recall), 1)\n\n return {\"val_loss\": val_loss, \"val_acc\": val_acc, \"val_bal_acc\": val_bal_acc}\n\n\ndef _apply_gradient_checkpointing(model: nn.Module) -> None:\n \"\"\"Enable gradient checkpointing on sequential layers to save memory.\n\n Wraps each layer's forward in ``torch.utils.checkpoint.checkpoint``\n so intermediate activations are recomputed during backward instead\n of stored. Roughly halves activation memory at ~30% compute cost.\n \"\"\"\n from torch.utils.checkpoint import checkpoint\n\n for name, module in model.named_children():\n orig_forward = module.forward\n\n def _make_ckpt_forward(fwd):\n def _ckpt_forward(*args, **kwargs):\n # checkpoint requires at least one tensor with requires_grad\n def run(*a):\n return fwd(*a, **kwargs)\n\n tensors = [a for a in args if isinstance(a, torch.Tensor)]\n if tensors and any(t.requires_grad for t in tensors):\n return checkpoint(run, *args, use_reentrant=False)\n return fwd(*args, **kwargs)\n\n return _ckpt_forward\n\n module.forward = _make_ckpt_forward(orig_forward)\n logger.info(\n \"Gradient checkpointing enabled on %d modules\", len(list(model.children()))\n )\n\n\ndef _apply_model_parallel(model: nn.Module, gpus: int) -> nn.Module:\n \"\"\"Distribute model layers across multiple GPUs (pipeline parallelism).\n\n Splits the model's children into ``gpus`` roughly equal groups and\n places each group on a different GPU. Inserts device-transfer hooks\n between groups so tensors move between GPUs automatically.\n\n Parameters\n ----------\n model : nn.Module\n Model with sequential children (e.g., KWYKMeshNet).\n gpus : int\n Number of GPUs to distribute across.\n\n Returns\n -------\n nn.Module\n The model with layers placed on different GPUs and transfer hooks.\n \"\"\"\n children = list(model.named_children())\n if not children:\n logger.warning(\"Model has no children — placing on GPU 0\")\n return model.to(\"cuda:0\")\n\n # Split children into roughly equal groups\n n = len(children)\n group_size = max(1, (n + gpus - 1) // gpus)\n groups: list[list[tuple[str, nn.Module]]] = []\n for i in range(0, n, group_size):\n groups.append(children[i : i + group_size])\n\n # Place each group on its GPU\n device_map: dict[str, int] = {}\n for gpu_idx, group in enumerate(groups):\n device = torch.device(f\"cuda:{gpu_idx}\")\n for name, module in group:\n module.to(device)\n device_map[name] = gpu_idx\n logger.info(\n \"Model parallel: %d layers across %d GPUs: %s\",\n n,\n min(gpus, len(groups)),\n {k: f\"cuda:{v}\" for k, v in device_map.items()},\n )\n\n # Wrap forward to move tensors between devices\n orig_forward = model.forward\n\n def _mp_forward(*args, **kwargs):\n # Move input to first device\n first_mod = groups[0][0][1]\n dev_idx = first_mod.weight.device.index if hasattr(first_mod, \"weight\") else 0\n first_device = torch.device(f\"cuda:{dev_idx}\")\n new_args = tuple(\n a.to(first_device) if isinstance(a, torch.Tensor) else a for a in args\n )\n return orig_forward(*new_args, **kwargs)\n\n model.forward = _mp_forward\n\n # Add hooks to move activations between GPUs at group boundaries\n for gpu_idx, group in enumerate(groups):\n if gpu_idx == 0:\n continue\n target_device = torch.device(f\"cuda:{gpu_idx}\")\n first_module = group[0][1]\n\n def _make_hook(dev):\n def _hook(module, inputs):\n return tuple(\n x.to(dev) if isinstance(x, torch.Tensor) else x for x in inputs\n )\n\n return _hook\n\n first_module.register_forward_pre_hook(_make_hook(target_device))\n\n return model\n\n\ndef fit(\n model: nn.Module,\n loader: DataLoader,\n criterion: nn.Module,\n optimizer: torch.optim.Optimizer,\n max_epochs: int = 10,\n gpus: int = 1,\n checkpoint_dir: str | Path | None = None,\n callbacks: list[Any] | None = None,\n val_loader: DataLoader | None = None,\n checkpoint_freq: int = 0,\n gradient_checkpointing: bool = False,\n model_parallel: bool = False,\n resume_from: str | Path | None = None,\n) -> dict:\n \"\"\"Train a model with optional multi-GPU DDP or model parallelism.\n\n Parameters\n ----------\n model : nn.Module\n PyTorch model to train.\n loader : DataLoader\n Training data loader.\n criterion : nn.Module\n Loss function.\n optimizer : Optimizer\n PyTorch optimizer.\n max_epochs : int\n Number of training epochs.\n gpus : int\n Number of GPUs to use (1 = single GPU/CPU, >1 = DDP or model parallel).\n checkpoint_dir : path or None\n Directory for saving checkpoints. None disables checkpointing.\n callbacks : list or None\n Optional callback functions called after each epoch with\n signature ``callback(epoch, logs, model)`` where logs is a dict\n containing at minimum ``{\"loss\": float}``.\n val_loader : DataLoader or None\n Validation data loader. If provided, validation loss and accuracy\n are computed each epoch and included in the logs dict.\n checkpoint_freq : int\n Save a checkpoint every N epochs (in addition to best model).\n 0 = only save best model. Checkpoints are saved as\n ``epoch_NNN.pth`` in checkpoint_dir.\n gradient_checkpointing : bool\n If True, trade compute for memory by recomputing activations\n during backward. Roughly halves activation memory.\n model_parallel : bool\n If True and gpus > 1, distribute layers across GPUs (pipeline\n parallelism) instead of DDP. Useful when a single batch is too\n large for one GPU.\n\n Returns\n -------\n dict\n ``{\"history\": [{\"epoch\": int, \"loss\": float, ...}, ...],\n \"checkpoint_path\": str | None}``\n \"\"\"\n device = get_device()\n\n # Apply gradient checkpointing if requested\n if gradient_checkpointing:\n _apply_gradient_checkpointing(model)\n\n # Multi-GPU dispatch\n if gpus > 1 and torch.cuda.device_count() >= gpus:\n if model_parallel:\n # Pipeline parallelism: split layers across GPUs\n model = _apply_model_parallel(model, gpus)\n device = torch.device(\"cuda:0\") # input goes to first GPU\n # Fall through to single-process training loop below\n else:\n # Data parallelism: DDP\n return _fit_ddp(\n model,\n loader,\n criterion,\n optimizer,\n max_epochs,\n gpus,\n checkpoint_dir,\n callbacks,\n val_loader,\n checkpoint_freq,\n )\n\n if not model_parallel:\n model = model.to(device)\n\n best_loss = float(\"inf\")\n ckpt_path = None\n history: list[dict[str, Any]] = [] # one entry per epoch\n start_epoch = 0\n\n if checkpoint_dir is not None:\n checkpoint_dir = Path(checkpoint_dir)\n checkpoint_dir.mkdir(parents=True, exist_ok=True)\n\n # Resume from checkpoint (auto-detect or explicit path)\n resume_path = None\n if resume_from is not None:\n resume_path = Path(resume_from)\n elif checkpoint_dir is not None and (checkpoint_dir / \"checkpoint.pt\").exists():\n resume_path = checkpoint_dir / \"checkpoint.pt\"\n\n if resume_path is not None and resume_path.exists():\n from nobrainer.slurm import load_checkpoint as _load_ckpt\n\n ckpt_dir = (\n resume_path.parent\n if resume_path.name == \"checkpoint.pt\"\n else checkpoint_dir\n )\n start_epoch, prev_metrics = _load_ckpt(ckpt_dir, model, optimizer)\n history = prev_metrics.get(\"history\", [])\n best_loss = min((h[\"loss\"] for h in history), default=float(\"inf\"))\n logger.info(\n \"Resumed from epoch %d (%d history entries, best_loss=%.4f)\",\n start_epoch,\n len(history),\n best_loss,\n )\n\n for epoch in range(start_epoch, max_epochs):\n model.train()\n epoch_loss = 0.0\n n_batches = 0\n\n for batch in loader:\n if isinstance(batch, dict):\n images = batch[\"image\"].to(device)\n labels = batch[\"label\"].to(device)\n elif isinstance(batch, (list, tuple)):\n images = batch[0].to(device)\n labels = batch[1].to(device)\n else:\n raise TypeError(f\"Unsupported batch type: {type(batch)}\")\n\n if labels.ndim == images.ndim and labels.shape[1] == 1:\n labels = labels.squeeze(1)\n if labels.dtype in (torch.float32, torch.float64):\n labels = labels.long()\n\n optimizer.zero_grad()\n pred = model(images)\n if pred.device != labels.device:\n labels = labels.to(pred.device)\n loss = criterion(pred, labels)\n loss.backward()\n optimizer.step()\n\n epoch_loss += loss.item()\n n_batches += 1\n\n avg_loss = epoch_loss / max(n_batches, 1)\n\n # Epoch metrics\n logs: dict[str, Any] = {\"epoch\": epoch + 1, \"loss\": avg_loss}\n\n if val_loader is not None:\n logs.update(_run_validation(model, val_loader, criterion, device))\n model.train()\n\n history.append(logs)\n\n # Best model\n if avg_loss < best_loss:\n best_loss = avg_loss\n if checkpoint_dir is not None:\n ckpt_path = str(checkpoint_dir / \"best_model.pth\")\n torch.save(model.state_dict(), ckpt_path)\n from nobrainer.processing.croissant import write_checkpoint_croissant\n\n write_checkpoint_croissant(\n checkpoint_dir, model, optimizer, criterion, history\n )\n\n # Resumable checkpoint (every epoch)\n if checkpoint_dir is not None:\n from nobrainer.slurm import save_checkpoint as _save_ckpt\n\n _save_ckpt(\n checkpoint_dir, model, optimizer, epoch + 1, {\"history\": history}\n )\n\n # Named checkpoint (for post-hoc Dice eval)\n if (\n checkpoint_dir is not None\n and checkpoint_freq > 0\n and (epoch + 1) % checkpoint_freq == 0\n ):\n epoch_ckpt = checkpoint_dir / f\"epoch_{epoch + 1:03d}.pth\"\n torch.save(model.state_dict(), epoch_ckpt)\n\n if callbacks:\n for cb in callbacks:\n cb(epoch, logs, model)\n\n logger.debug(\n \"Epoch %d/%d: %s\",\n epoch + 1,\n max_epochs,\n \" \".join(f\"{k}={v:.4f}\" for k, v in logs.items() if isinstance(v, float)),\n )\n\n return {\"history\": history, \"checkpoint_path\": ckpt_path}\n\n\ndef _ddp_worker(\n rank: int,\n world_size: int,\n model: nn.Module,\n train_dataset,\n val_dataset,\n batch_size: int,\n num_workers: int,\n criterion: nn.Module,\n optimizer: torch.optim.Optimizer,\n max_epochs: int,\n checkpoint_dir: str | Path | None,\n checkpoint_freq: int,\n result_dict: dict,\n) -> None:\n \"\"\"Single DDP worker — module-level function for mp.spawn pickling.\"\"\"\n import torch.distributed as dist\n from torch.nn.parallel import DistributedDataParallel as DDP\n from torch.utils.data.distributed import DistributedSampler\n\n dist.init_process_group(\"nccl\", rank=rank, world_size=world_size)\n torch.cuda.set_device(rank)\n device = torch.device(f\"cuda:{rank}\")\n\n local_model = model.to(device)\n ddp_model = DDP(local_model, device_ids=[rank])\n\n sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank)\n ddp_loader = DataLoader(\n train_dataset,\n batch_size=batch_size,\n sampler=sampler,\n num_workers=num_workers,\n pin_memory=True,\n )\n\n # Validation loader on rank 0 only (no DDP sampler needed)\n val_loader = None\n if val_dataset is not None and rank == 0:\n val_loader = DataLoader(\n val_dataset,\n batch_size=batch_size,\n shuffle=False,\n num_workers=num_workers,\n pin_memory=True,\n )\n\n best_loss = float(\"inf\")\n ckpt_path = None\n history: list[dict[str, Any]] = []\n\n if checkpoint_dir is not None and rank == 0:\n Path(checkpoint_dir).mkdir(parents=True, exist_ok=True)\n\n # ExperimentTracker for live metrics (rank 0 only)\n tracker = None\n if rank == 0 and checkpoint_dir is not None:\n from nobrainer.experiment import ExperimentTracker\n\n tracker = ExperimentTracker(output_dir=checkpoint_dir)\n\n for epoch in range(max_epochs):\n sampler.set_epoch(epoch)\n ddp_model.train()\n epoch_loss = 0.0\n n_batches = 0\n\n for batch in ddp_loader:\n if isinstance(batch, dict):\n images = batch[\"image\"].to(device)\n labels = batch[\"label\"].to(device)\n elif isinstance(batch, (list, tuple)):\n images = batch[0].to(device)\n labels = batch[1].to(device)\n else:\n raise TypeError(f\"Unsupported batch type: {type(batch)}\")\n\n if labels.ndim == images.ndim and labels.shape[1] == 1:\n labels = labels.squeeze(1)\n if labels.dtype in (torch.float32, torch.float64):\n labels = labels.long()\n\n optimizer.zero_grad()\n pred = ddp_model(images)\n loss = criterion(pred, labels)\n loss.backward()\n optimizer.step()\n\n epoch_loss += loss.item()\n n_batches += 1\n\n avg_loss = epoch_loss / max(n_batches, 1)\n logs: dict[str, Any] = {\"epoch\": epoch + 1, \"loss\": avg_loss}\n\n if rank == 0:\n if val_loader is not None:\n logs.update(\n _run_validation(ddp_model.module, val_loader, criterion, device)\n )\n ddp_model.train()\n\n history.append(logs)\n\n if avg_loss < best_loss:\n best_loss = avg_loss\n if checkpoint_dir is not None:\n ckpt_path = str(Path(checkpoint_dir) / \"best_model.pth\")\n torch.save(ddp_model.module.state_dict(), ckpt_path)\n from nobrainer.processing.croissant import (\n write_checkpoint_croissant,\n )\n\n write_checkpoint_croissant(\n checkpoint_dir,\n ddp_model.module,\n optimizer,\n criterion,\n history,\n )\n\n # Resumable checkpoint\n if checkpoint_dir is not None:\n from nobrainer.slurm import save_checkpoint as _save_ckpt\n\n _save_ckpt(\n checkpoint_dir,\n ddp_model.module,\n optimizer,\n epoch + 1,\n {\"history\": history},\n )\n\n # Named checkpoint for post-hoc Dice eval\n if (\n checkpoint_dir is not None\n and checkpoint_freq > 0\n and (epoch + 1) % checkpoint_freq == 0\n ):\n torch.save(\n ddp_model.module.state_dict(),\n Path(checkpoint_dir) / f\"epoch_{epoch + 1:03d}.pth\",\n )\n\n if tracker is not None:\n tracker.log(logs)\n\n logger.info(\n \"Epoch %d/%d: %s\",\n epoch + 1,\n max_epochs,\n \" \".join(\n f\"{k}={v:.4f}\" for k, v in logs.items() if isinstance(v, float)\n ),\n )\n\n if rank == 0:\n result_dict[\"history\"] = history\n result_dict[\"checkpoint_path\"] = ckpt_path\n\n if tracker is not None:\n tracker.finish()\n\n dist.destroy_process_group()\n\n\ndef _fit_ddp(\n model: nn.Module,\n loader: DataLoader,\n criterion: nn.Module,\n optimizer: torch.optim.Optimizer,\n max_epochs: int,\n gpus: int,\n checkpoint_dir: str | Path | None,\n callbacks: list[Any] | None,\n val_loader: DataLoader | None = None,\n checkpoint_freq: int = 0,\n) -> dict:\n \"\"\"Multi-GPU training via DistributedDataParallel.\n\n Launches ``gpus`` processes via ``mp.spawn``. Validation runs on\n rank 0 inside the worker — no callbacks needed for val metrics.\n \"\"\"\n import os\n\n import torch.multiprocessing as mp\n\n os.environ.setdefault(\"MASTER_ADDR\", \"localhost\")\n os.environ.setdefault(\"MASTER_PORT\", \"29500\")\n\n results: dict = mp.Manager().dict()\n\n # Extract datasets (picklable) — not DataLoaders (may have closures)\n train_dataset = loader.dataset\n val_dataset = val_loader.dataset if val_loader is not None else None\n\n mp.spawn(\n _ddp_worker,\n args=(\n gpus,\n model,\n train_dataset,\n val_dataset,\n loader.batch_size,\n loader.num_workers,\n criterion,\n optimizer,\n max_epochs,\n checkpoint_dir,\n checkpoint_freq,\n results,\n ),\n nprocs=gpus,\n join=True,\n )\n\n result = dict(results)\n if \"history\" in result:\n result[\"history\"] = list(result[\"history\"])\n return result\n\n\n__all__ = [\"fit\"]\n" }, { "path": "nobrainer/utils.py", "content": "\"\"\"Utilities for Nobrainer.\"\"\"\n\nfrom collections import namedtuple\nimport csv\nimport hashlib\nimport os\nimport tempfile\nimport urllib.request\n\nimport numpy as np\nimport psutil\n\n_cache_dir = os.path.join(tempfile.gettempdir(), \"nobrainer-data\")\n\n\ndef _sha256(path: str) -> str:\n \"\"\"Compute SHA-256 hex digest of a file.\"\"\"\n h = hashlib.sha256()\n with open(path, \"rb\") as f:\n for chunk in iter(lambda: f.read(1 << 16), b\"\"):\n h.update(chunk)\n return h.hexdigest()\n\n\ndef _download_if_needed(url: str, dest: str, expected_hash: str) -> None:\n \"\"\"Download *url* to *dest* if the file is missing or hash mismatches.\"\"\"\n if os.path.isfile(dest):\n if _sha256(dest) == expected_hash:\n return\n urllib.request.urlretrieve(url, dest)\n actual = _sha256(dest)\n if actual != expected_hash:\n raise RuntimeError(\n f\"Hash mismatch for {dest}: expected {expected_hash}, got {actual}\"\n )\n\n\ndef get_data(cache_dir=_cache_dir):\n \"\"\"Download sample features and labels. The features are T1-weighted MGZ\n files, and the labels are the corresponding aparc+aseg MGZ files, created\n with FreeSurfer. This will download 46 megabytes of data.\n\n These data can be found at\n https://datasets.datalad.org/workshops/nih-2017/ds000114/.\n\n Parameters\n ----------\n cache_dir: str, directory where to save the data. By default, saves to a\n temporary directory.\n\n Returns\n -------\n List of `(features, labels)`.\n \"\"\"\n\n os.makedirs(cache_dir, exist_ok=True)\n URLHashPair = namedtuple(\"URLHashPair\", \"sub x_hash y_hash\")\n hashes = [\n URLHashPair(\n sub=\"sub-01\",\n x_hash=\"67d0053f021d1d137bc99715e4e3ebb763364c8ce04311b1032d4253fc149f52\",\n y_hash=\"7a85b628653f24e2b71cbef6dda86ab24a1743c5f6dbd996bdde258414e780b5\",\n ),\n URLHashPair(\n sub=\"sub-02\",\n x_hash=\"c0fee669a34bf3b43c8e4aecc88204512ef4e83f2e414640a5abc076b435990c\",\n y_hash=\"c92357c2571da72d15332b2b4838b94d442d4abd3dbddc4b54202d68f0e19380\",\n ),\n URLHashPair(\n sub=\"sub-03\",\n x_hash=\"e2bba954e37f5791260f0ec573456e3293bbd40dba139bb1af417eaaeabe63e6\",\n y_hash=\"e9204f0d50f06a89dd1870911f7ef5e9808e222227799a5384dceeb941ee8f9d\",\n ),\n URLHashPair(\n sub=\"sub-04\",\n x_hash=\"deec5245a2a5948f7e1053ace8d8a31396b14a96d520c6a52305434e75abe1e8\",\n y_hash=\"c50e33a3f87aca351414e729b7c25404af364dfe5dd1de5fe380a460cbe9f891\",\n ),\n URLHashPair(\n sub=\"sub-05\",\n x_hash=\"8a7fe84918f3f80b87903a1e8f7bd20792c0ebc7528fb98513be373258dfd6c0\",\n y_hash=\"682f52633633551d6fda71ede65aa41e16c332ebf42b4df042bc312200b0337c\",\n ),\n URLHashPair(\n sub=\"sub-06\",\n x_hash=\"f9a0c40bcd62d7b7e88015867ab5d926009b097ac3235499a541ac9072dd90c8\",\n y_hash=\"31c842969af9ac178361fa8c13f656a47d27d95357abaf3e7f3521671aa17929\",\n ),\n URLHashPair(\n sub=\"sub-07\",\n x_hash=\"9de3b7392f5383e7391c5fcd9266d6b7ab6b57bc7ab203cc9ad2a29a2d31a85b\",\n y_hash=\"b2e48bbfc4185261785643fc8ab066be5f97215b5a9b029ade1ffb12d54d616e\",\n ),\n URLHashPair(\n sub=\"sub-08\",\n x_hash=\"361098fc69c280970bb0b0d7ea6aba80d383c12e3ccfe5899693bc35b68efbe4\",\n y_hash=\"0c980ef851b1391f580d91fc87c10d6d30315527cc0749c1010f2b7d5819a009\",\n ),\n URLHashPair(\n sub=\"sub-09\",\n x_hash=\"1456b35112297df5caacb9d33cb047aa85a3a5b4db3b4b5f9a5c2e189a684e1a\",\n y_hash=\"696f1e9fef512193b71580292e0edc5835f396d2c8d63909c13668ef7bed433b\",\n ),\n URLHashPair(\n sub=\"sub-10\",\n x_hash=\"97447f17402e0f9990cd0917f281704893b52a9b61a3241b23a112a0a143d26e\",\n y_hash=\"97a7947ba1a28963714c9f5c82520d9ef803d005695a0b4109d5a73d7e8a537b\",\n ),\n ]\n x_filename = \"t1.mgz\"\n y_filename = \"aparc+aseg.mgz\"\n url_template = (\n \"https://datasets.datalad.org/workshops/nih-2017/ds000114/derivatives/\"\n \"freesurfer/{sub}/mri/{fname}\"\n )\n output = [(\"features\", \"labels\")]\n downloads_dir = os.path.join(cache_dir, \"datasets\")\n os.makedirs(downloads_dir, exist_ok=True)\n for h in hashes:\n x_origin = url_template.format(sub=h.sub, fname=x_filename)\n y_origin = url_template.format(sub=h.sub, fname=y_filename)\n x_fname = h.sub + \"_\" + x_origin.rsplit(\"/\", 1)[-1]\n y_fname = h.sub + \"_\" + y_origin.rsplit(\"/\", 1)[-1]\n x_out = os.path.join(downloads_dir, x_fname)\n y_out = os.path.join(downloads_dir, y_fname)\n _download_if_needed(x_origin, x_out, h.x_hash)\n _download_if_needed(y_origin, y_out, h.y_hash)\n output.append((x_out, y_out))\n\n csvpath = os.path.join(cache_dir, \"filepaths.csv\")\n with open(csvpath, \"w\", newline=\"\") as f:\n writer = csv.writer(f)\n writer.writerows(output)\n\n return csvpath\n\n\nclass StreamingStats:\n \"\"\"Object to calculate statistics on streaming data.\n\n Compatible with scalars and n-dimensional arrays.\n\n Examples\n --------\n\n ```python\n >>> s = StreamingStats()\n >>> s.update(10).update(20)\n >>> s.mean()\n 15.0\n ```\n\n ```python\n >>> import numpy as np\n >>> a = np.array([[0, 2], [4, 8]])\n >>> b = np.array([[2, 4], [8, 16]])\n >>> s = StreamingStats()\n >>> s.update(a).update(b)\n >>> s.mean()\n array([[ 1., 3.],\n [ 6., 12.]])\n ```\n \"\"\"\n\n def __init__(self):\n self._n_samples = 0\n self._current_mean = 0.0\n self._M = 0.0\n\n def update(self, value):\n \"\"\"Update the statistics with the next value.\n\n Parameters\n ----------\n value: scalar, array-like\n\n Returns\n -------\n Modified instance.\n \"\"\"\n if self._n_samples == 0:\n self._current_mean = value\n else:\n prev_mean = self._current_mean\n curr_mean = prev_mean + (value - prev_mean) / (self._n_samples + 1)\n _M = self._M + (prev_mean - value) * (curr_mean - value)\n # Set the instance attributes after computation in case there are\n # errors during computation.\n self._current_mean = curr_mean\n self._M = _M\n self._n_samples += 1\n return self\n\n def mean(self):\n \"\"\"Return current mean of streaming data.\"\"\"\n return self._current_mean\n\n def var(self):\n \"\"\"Return current variance of streaming data.\"\"\"\n return self._M / self._n_samples\n\n def std(self):\n \"\"\"Return current standard deviation of streaming data.\"\"\"\n return self.var() ** 0.5\n\n def entropy(self):\n \"\"\"Return current entropy of streaming data.\"\"\"\n eps = 1e-07\n mult = np.multiply(np.log(self.mean() + eps), self.mean())\n return -mult\n # return -np.sum(mult, axis=axis)\n\n\ndef get_num_parallel():\n # Get number of processes allocated to the current process.\n # Note the difference from `os.cpu_count()`.\n try:\n num_parallel_calls = len(psutil.Process().cpu_affinity())\n except AttributeError:\n num_parallel_calls = psutil.cpu_count()\n return num_parallel_calls\n" }, { "path": "nobrainer/validation.py", "content": "#!/usr/bin/env python3\n\nfrom pathlib import Path\n\nimport nibabel as nib\nimport numpy as np\n\nfrom .io import read_mapping, read_volume\nfrom .metrics import dice as dice_numpy\nfrom .prediction import predict as _predict\nfrom .volume import normalize_numpy, replace\n\nDT_X = \"float32\"\n\n\ndef validate_from_filepath(\n filepath,\n predictor,\n block_shape,\n n_classes,\n mapping_y,\n return_variance=False,\n return_entropy=False,\n return_array_from_images=False,\n n_samples=1,\n normalizer=normalize_numpy,\n batch_size=4,\n):\n \"\"\"Computes dice for a prediction compared to a ground truth image.\n\n Args:\n filepath: tuple, tuple of paths to existing neuroimaging volume (index 0)\n and ground truth (index 1).\n predictor: TensorFlow Predictor object, predictor from previously\n trained model.\n n_classes: int, number of classifications the model is trained to output.\n mapping_y: path-like, path to csv mapping file per command line argument.\n block_shape: tuple of len 3, shape of blocks on which to predict.\n return_variance: Boolean. If set True, it returns the running population\n variance along with mean. Note, if the n_samples is smaller or equal to 1,\n the variance will not be returned; instead it will return None\n return_entropy: Boolean. If set True, it returns the running entropy.\n along with mean.\n return_array_from_images: Boolean. If set True and the given input is either\n image, filepath, or filepaths, it will return arrays of [mean, variance,\n entropy] instead of images of them. Also, if the input is array, it will\n simply return array, whether or not this flag is True or False.\n n_samples: The number of sampling. If set as 1, it will just return the\n single prediction value.\n normalizer: callable, function that accepts an ndarray and returns an\n ndarray. Called before separating volume into blocks.\n batch_size: int, number of sub-volumes per batch for prediction.\n dtype: str or dtype object, dtype of features.\n\n Returns:\n `nibabel.spatialimages.SpatialImage` or arrays of predictions of\n mean, variance(optional), and entropy (optional).\n \"\"\"\n if not Path(filepath[0]).is_file():\n raise FileNotFoundError(\"could not find file {}\".format(filepath[0]))\n img = nib.load(filepath[0])\n y = read_volume(filepath[1], dtype=np.int32)\n\n outputs = _predict(\n inputs=img,\n predictor=predictor,\n block_shape=block_shape,\n return_variance=return_variance,\n return_entropy=return_entropy,\n return_array_from_images=return_array_from_images,\n n_samples=n_samples,\n normalizer=normalizer,\n batch_size=batch_size,\n )\n prediction_image = outputs[0].get_data()\n y = replace(y, read_mapping(mapping_y))\n dice = get_dice_for_images(prediction_image, y, n_classes)\n return outputs, dice\n\n\ndef get_dice_for_images(pred, gt, n_classes):\n \"\"\"Computes dice for a prediction compared to a ground truth image.\n\n Args:\n pred: nibabel.spatialimages.SpatialImage, a predicted image.\n gt: nibabel.spatialimages.SpatialImage, a ground-truth image.\n\n\n Returns:\n `nibabel.spatialimages.SpatialImage`.\n \"\"\"\n dice = np.zeros(n_classes)\n for i in range(n_classes):\n u = np.equal(pred, i)\n v = np.equal(gt, i)\n dice[i] = dice_numpy(u, v)\n\n return dice\n\n\ndef validate_from_filepaths(\n filepaths,\n predictor,\n block_shape,\n n_classes,\n mapping_y,\n output_path,\n return_variance=False,\n return_entropy=False,\n return_array_from_images=False,\n n_samples=1,\n normalizer=normalize_numpy,\n batch_size=4,\n dtype=DT_X,\n):\n \"\"\"Yield predictions from filepaths using a SavedModel.\n\n Args:\n test_csv: list, neuroimaging volume filepaths on which to predict.\n n_classes: int, number of classifications the model is trained to output.\n mapping_y: path-like, path to csv mapping file per command line argument.\n block_shape: tuple of len 3, shape of blocks on which to predict.\n predictor: TensorFlow Predictor object, predictor from previously\n trained model.\n block_shape: tuple of len 3, shape of blocks on which to predict.\n normalizer: callable, function that accepts an ndarray and returns\n an ndarray. Called before separating volume into blocks.\n batch_size: int, number of sub-volumes per batch for prediction.\n dtype: str or dtype object, dtype of features.\n\n Returns:\n None\n \"\"\"\n for filepath in filepaths:\n outputs, dice = validate_from_filepath(\n filepath=filepath,\n predictor=predictor,\n n_classes=n_classes,\n mapping_y=mapping_y,\n block_shape=block_shape,\n return_variance=return_variance,\n return_entropy=return_entropy,\n return_array_from_images=return_array_from_images,\n n_samples=n_samples,\n normalizer=normalizer,\n batch_size=batch_size,\n dtype=dtype,\n )\n\n outpath = Path(filepath[0])\n output_path = Path(output_path)\n suffixes = \"\".join(s for s in outpath.suffixes)\n mean_path = output_path / (outpath.stem + \"_mean\" + suffixes)\n variance_path = output_path / (outpath.stem + \"_variance\" + suffixes)\n entropy_path = output_path / (outpath.stem + \"_entropy\" + suffixes)\n dice_path = output_path / (outpath.stem + \"_dice.npy\")\n # if mean_path.is_file() or variance_path.is_file() or entropy_path.is_file():\n # raise Exception(str(mean_path) + \" or \" + str(variance_path) +\n # \" or \" + str(entropy_path) + \" already exists.\")\n\n nib.save(outputs[0], mean_path.as_posix()) # fix\n if not return_array_from_images:\n include_variance = (n_samples > 1) and (return_variance)\n include_entropy = (n_samples > 1) and (return_entropy)\n if include_variance and return_entropy:\n nib.save(outputs[1], str(variance_path))\n nib.save(outputs[2], str(entropy_path))\n elif include_variance:\n nib.save(outputs[1], str(variance_path))\n elif include_entropy:\n nib.save(outputs[1], str(entropy_path))\n\n print(filepath[0])\n print(\"Dice: \" + str(np.mean(dice)))\n np.save(dice_path, dice)\n" }, { "path": "pyproject.toml", "content": "[build-system]\nrequires = [\"hatchling\", \"hatch-vcs\"]\nbuild-backend = \"hatchling.build\"\n\n[project]\nname = \"nobrainer\"\ndynamic = [\"version\"]\ndescription = \"A deep learning framework for 3D brain image processing.\"\nreadme = \"README.md\"\nlicense = \"Apache-2.0\"\nrequires-python = \">= 3.12\"\nauthors = [\n { name = \"Nobrainer Developers\", email = \"jakub.kaczmarzyk@gmail.com\" },\n]\nmaintainers = [\n { name = \"Satrajit Ghosh\", email = \"satrajit.ghosh@gmail.com\" },\n]\nclassifiers = [\n \"Development Status :: 4 - Beta\",\n \"Environment :: Console\",\n \"Intended Audience :: Developers\",\n \"Intended Audience :: Education\",\n \"Intended Audience :: Healthcare Industry\",\n \"Intended Audience :: Science/Research\",\n \"License :: OSI Approved :: Apache Software License\",\n \"Operating System :: OS Independent\",\n \"Programming Language :: Python :: 3\",\n \"Programming Language :: Python :: 3 :: Only\",\n \"Programming Language :: Python :: 3.12\",\n \"Programming Language :: Python :: 3.13\",\n \"Programming Language :: Python :: 3.14\",\n \"Topic :: Scientific/Engineering :: Artificial Intelligence\",\n \"Topic :: Software Development\",\n \"Topic :: Software Development :: Libraries :: Python Modules\",\n]\ndependencies = [\n \"click\",\n \"einops >= 0.7\",\n \"fsspec\",\n \"h5py >= 3.9\",\n \"joblib\",\n \"monai >= 1.3\",\n \"nibabel >= 5.0\",\n \"numpy\",\n \"psutil\",\n \"scikit-image\",\n \"torch >= 2.0\",\n]\n\n[project.urls]\nHomepage = \"https://github.com/neuronets/nobrainer\"\nDocumentation = \"https://neuronets.dev/nobrainer-book/\"\n\"Source Code\" = \"https://github.com/neuronets/nobrainer\"\n\"Bug Tracker\" = \"https://github.com/neuronets/helpdesk/issues\"\n\n[project.scripts]\nnobrainer = \"nobrainer.cli.main:cli\"\n\n[project.optional-dependencies]\nbayesian = [\"pyro-ppl >= 1.9\"]\ngenerative = [\"pytorch-lightning >= 2.0\"]\nlightning = [\"pytorch-lightning >= 2.0\"]\nzarr = [\"zarr >= 3.0\", \"nifti-zarr\", \"scipy >= 1.11\"]\ncroissant = [\"mlcroissant\"]\nversioning = [\"datalad >= 0.19\"]\ntfrecord = [\"tfrecord >= 1.14\"]\ndev = [\"pre-commit\", \"pytest\", \"pytest-cov\", \"scipy\"]\nall = [\n \"nobrainer[bayesian,generative,zarr,croissant,versioning,tfrecord,dev]\",\n]\n\n[tool.hatch.version]\nsource = \"vcs\"\nfallback-version = \"0.0.0.dev0\"\n\n[tool.hatch.build.hooks.vcs]\nversion-file = \"nobrainer/_version.py\"\n\n[tool.hatch.build.targets.wheel]\npackages = [\"nobrainer\"]\n\n[tool.black]\nexclude = '\\.eggs|\\.git|\\.mypy_cache|\\.tox|\\.venv|_build|buck-out|build|dist|_version\\.py|versioneer\\.py'\n\n[tool.isort]\nprofile = \"black\"\nforce_sort_within_sections = true\nreverse_relative = true\nsort_relative_in_force_sorted_sections = true\nknown_first_party = [\"nobrainer\"]\n\n[tool.pytest.ini_options]\ntestpaths = [\"nobrainer/tests\"]\nmarkers = [\n \"gpu: marks tests that require a CUDA-capable GPU (deselect with '-m not gpu')\",\n]\n\n[tool.coverage.run]\nbranch = true\nomit = [\"nobrainer/_version.py\", \"*/tests*\"]\n\n[tool.coverage.report]\nexclude_lines = [\n \"pragma: no cover\",\n \"raise NotImplementedError\",\n \"if __name__ == .__main__.\",\n]\nignore_errors = true\n\n[tool.codespell]\nskip = \"nobrainer/_version.py,versioneer.py\"\nignore-words-list = \"nd\"\n" }, { "path": "scripts/kwyk_reproduction/01_assemble_dataset.py", "content": "#!/usr/bin/env python\n\"\"\"Assemble training dataset from OpenNeuro fmriprep derivatives.\n\nUses :mod:`nobrainer.datasets.openneuro` to install datasets via DataLad\nand discover paired (T1w, aparc+aseg) files per subject.\n\nUsage:\n python 01_assemble_dataset.py --datasets ds000114 --output-csv manifest.csv\n python 01_assemble_dataset.py --datasets ds000114 ds000228 ds002609 \\\n --output-csv manifest.csv --label-mapping binary --split 80 10 10\n\"\"\"\n\nfrom __future__ import annotations\n\nimport argparse\nimport logging\nfrom pathlib import Path\n\nlogging.basicConfig(\n level=logging.INFO,\n format=\"%(asctime)s %(levelname)s %(message)s\",\n)\nlog = logging.getLogger(__name__)\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Assemble dataset from OpenNeuro\")\n parser.add_argument(\n \"--datasets\",\n nargs=\"+\",\n default=[\"ds000114\"],\n help=\"OpenNeuro dataset IDs\",\n )\n parser.add_argument(\"--output-dir\", default=\"data\", help=\"Output directory\")\n parser.add_argument(\n \"--output-csv\", default=\"manifest.csv\", help=\"Output manifest CSV\"\n )\n parser.add_argument(\n \"--label-mapping\",\n default=\"binary\",\n help=\"Label mapping: binary, 6-class, 50-class, 115-class\",\n )\n parser.add_argument(\n \"--split\",\n nargs=3,\n type=int,\n default=[80, 10, 10],\n help=\"Train/val/test split percentages\",\n )\n parser.add_argument(\"--conform\", action=\"store_true\", help=\"Resample to 256³ @ 1mm\")\n args = parser.parse_args()\n\n from nobrainer.datasets.openneuro import (\n find_subject_pairs,\n install_derivatives,\n write_manifest,\n )\n\n output_dir = Path(args.output_dir)\n output_dir.mkdir(parents=True, exist_ok=True)\n\n all_pairs = []\n for ds_id in args.datasets:\n ds_dir = install_derivatives(ds_id, output_dir)\n pairs = find_subject_pairs(ds_dir)\n for p in pairs:\n p[\"dataset_id\"] = ds_id\n all_pairs.extend(pairs)\n\n if not all_pairs:\n log.error(\"No subject pairs found. Check dataset IDs and network access.\")\n raise SystemExit(1)\n\n # Optionally conform volumes\n if args.conform:\n import nibabel as nib\n from nibabel.processing import conform\n\n for row in all_pairs:\n img = nib.load(row[\"t1w_path\"])\n if img.shape[:3] != (256, 256, 256):\n log.info(\"Conforming %s\", Path(row[\"t1w_path\"]).name)\n conformed = conform(\n img, out_shape=(256, 256, 256), voxel_size=(1.0, 1.0, 1.0)\n )\n nib.save(conformed, row[\"t1w_path\"])\n\n write_manifest(all_pairs, args.output_csv, tuple(args.split))\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "scripts/kwyk_reproduction/02_train_meshnet.py", "content": "#!/usr/bin/env python\n\"\"\"Train a deterministic MeshNet for brain extraction / parcellation.\n\nUsage:\n python 02_train_meshnet.py --manifest manifest.csv --config config.yaml\n python 02_train_meshnet.py --manifest manifest.csv --config config.yaml \\\n --output-dir checkpoints/meshnet --epochs 100\n\"\"\"\n\nfrom __future__ import annotations\n\nimport argparse\nimport csv\nfrom pathlib import Path\nimport time\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nfrom utils import load_config, save_figure, setup_logging\n\nlog = setup_logging(__name__)\n\n\ndef parse_args() -> argparse.Namespace:\n \"\"\"Parse command-line arguments.\"\"\"\n parser = argparse.ArgumentParser(\n description=\"Train deterministic MeshNet for brain segmentation\",\n )\n parser.add_argument(\n \"--manifest\",\n type=str,\n required=True,\n help=\"Path to the dataset manifest CSV (output of 01_assemble_dataset.py)\",\n )\n parser.add_argument(\n \"--config\",\n type=str,\n default=\"config.yaml\",\n help=\"Path to YAML configuration file (default: config.yaml)\",\n )\n parser.add_argument(\n \"--output-dir\",\n type=str,\n default=\"checkpoints/meshnet\",\n help=\"Directory for saving model checkpoints and figures\",\n )\n parser.add_argument(\n \"--epochs\",\n type=int,\n default=None,\n help=\"Override number of training epochs from config\",\n )\n parser.add_argument(\n \"--resume\",\n type=str,\n default=None,\n help=\"Path to checkpoint .pth file to resume from\",\n )\n return parser.parse_args()\n\n\ndef load_manifest(manifest_path: str, split: str) -> list[tuple[str, str]]:\n \"\"\"Load manifest CSV and return (image, label) pairs for the given split.\"\"\"\n pairs = []\n with open(manifest_path) as f:\n reader = csv.DictReader(f)\n for row in reader:\n if row[\"split\"] == split:\n pairs.append((row[\"t1w_path\"], row[\"label_path\"]))\n return pairs\n\n\ndef evaluate_val_dice(\n seg,\n val_pairs: list[tuple[str, str]],\n block_shape: tuple[int, int, int],\n label_mapping: str | None,\n n_classes: int = 2,\n) -> list[float]:\n \"\"\"Compute per-volume mean class Dice on validation set.\n\n Returns a list of mean Dice scores (averaged across classes), one per volume.\n \"\"\"\n import nibabel as nib\n\n from nobrainer.prediction import predict\n from nobrainer.training import get_device\n\n # Load remap function for multi-class label mappings\n remap_fn = None\n if label_mapping and label_mapping != \"binary\":\n from nobrainer.processing.dataset import _load_label_mapping\n\n remap_fn = _load_label_mapping(label_mapping)\n\n dice_scores = []\n device = get_device()\n model = seg.model_.to(device)\n model.eval()\n\n for img_path, lbl_path in val_pairs:\n pred_img = predict(\n inputs=img_path,\n model=model,\n block_shape=block_shape,\n batch_size=128,\n return_labels=True,\n )\n pred_arr = np.asarray(pred_img.dataobj, dtype=np.int32)\n\n gt_arr = np.asarray(nib.load(lbl_path).dataobj, dtype=np.int32)\n if remap_fn is not None:\n gt_arr = remap_fn(torch.from_numpy(gt_arr)).numpy()\n elif label_mapping is None or label_mapping == \"binary\":\n gt_arr = (gt_arr > 0).astype(np.int32)\n pred_arr = (pred_arr > 0).astype(np.int32)\n\n # Per-class Dice (skip background class 0)\n class_dices = []\n for c in range(1, n_classes):\n pred_c = pred_arr == c\n gt_c = gt_arr == c\n intersection = (pred_c & gt_c).sum()\n total = pred_c.sum() + gt_c.sum()\n class_dices.append(2.0 * intersection / total if total > 0 else 1.0)\n\n mean_dice = float(np.mean(class_dices))\n dice_scores.append(mean_dice)\n log.info(\n \" Val volume %s: Dice=%.4f\",\n Path(img_path).name,\n mean_dice,\n )\n\n return dice_scores\n\n\ndef plot_learning_curve(\n train_losses: list[float],\n val_dice_scores: list[float],\n output_path: Path,\n) -> None:\n \"\"\"Generate dual y-axis learning curve (loss left, Dice right).\"\"\"\n epochs = list(range(1, len(train_losses) + 1))\n\n fig, ax_loss = plt.subplots(figsize=(10, 6))\n ax_dice = ax_loss.twinx()\n\n ax_loss.plot(epochs, train_losses, \"b-\", label=\"Train Loss\")\n ax_loss.set_xlabel(\"Epoch\")\n ax_loss.set_ylabel(\"Loss\", color=\"b\")\n ax_loss.tick_params(axis=\"y\", labelcolor=\"b\")\n\n if val_dice_scores:\n ax_dice.plot(epochs, val_dice_scores, \"r-o\", label=\"Val Dice (mean)\")\n ax_dice.set_ylabel(\"Dice Score\", color=\"r\")\n ax_dice.tick_params(axis=\"y\", labelcolor=\"r\")\n ax_dice.set_ylim(0.0, 1.0)\n\n fig.suptitle(\"MeshNet Training — Loss & Validation Dice\")\n fig.tight_layout()\n\n lines_loss, labels_loss = ax_loss.get_legend_handles_labels()\n lines_dice, labels_dice = ax_dice.get_legend_handles_labels()\n ax_loss.legend(\n lines_loss + lines_dice, labels_loss + labels_dice, loc=\"center right\"\n )\n\n save_figure(fig, output_path)\n plt.close(fig)\n log.info(\"Learning curve saved to %s\", output_path)\n\n\ndef main() -> None:\n \"\"\"Train deterministic MeshNet and evaluate on validation set.\"\"\"\n args = parse_args()\n t_start = time.time()\n\n # ---- Load config --------------------------------------------------------\n config = load_config(args.config)\n epochs = (\n args.epochs if args.epochs is not None else config.get(\"pretrain_epochs\", 50)\n )\n n_classes = config[\"n_classes\"]\n block_shape = tuple(config[\"block_shape\"])\n batch_size = config[\"batch_size\"]\n lr = config.get(\"lr\", 1e-4)\n label_mapping = config.get(\"label_mapping\", \"binary\")\n output_dir = Path(args.output_dir)\n output_dir.mkdir(parents=True, exist_ok=True)\n\n log.info(\"Config loaded from %s\", args.config)\n log.info(\n \"Training MeshNet: epochs=%d, n_classes=%d, block_shape=%s, batch_size=%d\",\n epochs,\n n_classes,\n block_shape,\n batch_size,\n )\n\n # ---- Load manifest and build datasets -----------------------------------\n train_pairs = load_manifest(args.manifest, split=\"train\")\n val_pairs = load_manifest(args.manifest, split=\"val\")\n log.info(\"Manifest: %d train, %d val volumes\", len(train_pairs), len(val_pairs))\n\n if not train_pairs:\n log.error(\"No training volumes found in manifest. Exiting.\")\n return\n\n # Auto-scale batch size to GPU memory\n from nobrainer.gpu import auto_batch_size as _auto_bs\n from nobrainer.gpu import gpu_count\n from nobrainer.processing.dataset import Dataset\n\n if gpu_count() > 0:\n from nobrainer.models import get as get_model\n\n _tmp_model = get_model(\"meshnet\")(\n n_classes=n_classes,\n filters=config.get(\"filters\", 96),\n receptive_field=config.get(\"receptive_field\", 37),\n dropout_rate=config.get(\"dropout_rate\", 0.25),\n )\n batch_size = _auto_bs(\n _tmp_model,\n block_shape,\n n_classes=n_classes,\n target_memory_fraction=0.90,\n )\n del _tmp_model\n log.info(\"Auto batch size: %d (target 90%% GPU memory)\", batch_size)\n\n patches_per_volume = config.get(\"patches_per_volume\", 50)\n zarr_store = config.get(\"zarr_store\")\n\n if zarr_store and Path(zarr_store).exists():\n log.info(\"Using Zarr store: %s\", zarr_store)\n ds_train = (\n Dataset.from_zarr(\n zarr_store,\n block_shape=block_shape,\n n_classes=n_classes,\n partition=\"train\",\n )\n .batch(batch_size)\n .binarize(label_mapping)\n .streaming(patches_per_volume=patches_per_volume)\n )\n else:\n ds_train = (\n Dataset.from_files(\n train_pairs, block_shape=block_shape, n_classes=n_classes\n )\n .batch(batch_size)\n .binarize(label_mapping)\n .streaming(patches_per_volume=patches_per_volume)\n )\n\n n_train = len(ds_train.data) if hasattr(ds_train, \"data\") else len(train_pairs)\n log.info(\n \"Training data: %d volumes × %d patches = %d blocks/epoch, batch_size=%d\",\n n_train,\n patches_per_volume,\n n_train * patches_per_volume,\n batch_size,\n )\n\n # ---- Build validation dataset for per-epoch block-level metrics ----------\n ds_val = None\n if val_pairs:\n if zarr_store and Path(zarr_store).exists():\n ds_val = (\n Dataset.from_zarr(\n zarr_store,\n block_shape=block_shape,\n n_classes=n_classes,\n partition=\"val\",\n )\n .batch(batch_size)\n .binarize(label_mapping)\n .streaming(patches_per_volume=patches_per_volume)\n )\n else:\n ds_val = (\n Dataset.from_files(\n val_pairs, block_shape=block_shape, n_classes=n_classes\n )\n .batch(batch_size)\n .binarize(label_mapping)\n .streaming(patches_per_volume=patches_per_volume)\n )\n\n # ---- Train with Segmentation estimator ----------------------------------\n from nobrainer.processing.segmentation import Segmentation\n\n model_args = {\n \"n_classes\": n_classes,\n \"filters\": config.get(\"filters\", 96),\n \"receptive_field\": config.get(\"receptive_field\", 37),\n \"dropout_rate\": config.get(\"dropout_rate\", 0.25),\n }\n\n log.info(\"Model args: %s\", model_args)\n\n seg = Segmentation(\n base_model=\"meshnet\",\n model_args=model_args,\n checkpoint_filepath=str(output_dir),\n )\n\n val_dice_per_epoch: list[float] = []\n val_dice_freq = config.get(\"val_dice_freq\", 5)\n\n # Simple logging callback (picklable — no closures)\n def _log_cb(epoch, logs, model):\n msg = f\"Epoch {epoch + 1}/{epochs}: train_loss={logs['loss']:.6f}\"\n if \"val_loss\" in logs:\n msg += f\" val_loss={logs['val_loss']:.6f}\"\n if \"val_acc\" in logs:\n msg += f\" val_acc={logs['val_acc']:.4f}\"\n if \"val_bal_acc\" in logs:\n msg += f\" bal_acc={logs['val_bal_acc']:.4f}\"\n log.info(msg)\n\n seg.fit(\n dataset_train=ds_train,\n dataset_validate=ds_val,\n epochs=epochs,\n optimizer=torch.optim.Adam,\n opt_args={\"lr\": lr},\n callbacks=[_log_cb],\n checkpoint_freq=val_dice_freq,\n gradient_checkpointing=config.get(\"gradient_checkpointing\", False),\n model_parallel=config.get(\"model_parallel\", False),\n resume_from=args.resume,\n )\n\n history = seg._training_result.get(\"history\", [])\n\n if history:\n last = history[-1]\n log.info(\n \"Training complete. %s\",\n \" \".join(f\"{k}={v:.4f}\" for k, v in last.items() if isinstance(v, float)),\n )\n else:\n log.info(\"Training complete (no history).\")\n\n # Ensure model is on the right device after DDP\n from nobrainer.training import get_device\n\n seg.model_.to(get_device())\n\n # Evaluate full-volume Dice on each checkpointed epoch\n if val_pairs:\n for epoch_idx in range(len(history)):\n epoch_num = history[epoch_idx].get(\"epoch\", epoch_idx + 1)\n ckpt_file = output_dir / f\"epoch_{epoch_num:03d}.pth\"\n if ckpt_file.exists():\n log.info(\"Evaluating Dice at epoch %d...\", epoch_num)\n seg.model_.load_state_dict(\n torch.load(ckpt_file, map_location=get_device(), weights_only=True)\n )\n dice_scores = evaluate_val_dice(\n seg, val_pairs, block_shape, label_mapping, n_classes\n )\n mean_dice = float(np.mean(dice_scores)) if dice_scores else 0.0\n history[epoch_idx][\"val_dice\"] = mean_dice\n log.info(\" Epoch %d Dice: %.4f\", epoch_num, mean_dice)\n\n train_losses = [h[\"loss\"] for h in history]\n val_dice_per_epoch = [h.get(\"val_dice\", float(\"nan\")) for h in history]\n fig_path = output_dir / \"learning_curve.png\"\n plot_learning_curve(train_losses, val_dice_per_epoch, fig_path)\n\n # ---- Save model with Croissant-ML metadata ------------------------------\n seg.save(output_dir)\n log.info(\"Model and Croissant-ML metadata saved to %s\", output_dir)\n\n # ---- Summary ------------------------------------------------------------\n elapsed = time.time() - t_start\n log.info(\"=\" * 60)\n log.info(\"MeshNet training complete\")\n log.info(\" Output directory : %s\", output_dir)\n log.info(\" Epochs : %d\", epochs)\n log.info(\" Final train loss : %.6f\", train_losses[-1] if train_losses else 0.0)\n final_dice = [d for d in val_dice_per_epoch if not np.isnan(d)]\n if final_dice:\n log.info(\" Val Dice (mean) : %.4f\", final_dice[-1])\n log.info(\" Elapsed time : %.1f s\", elapsed)\n log.info(\"=\" * 60)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "scripts/kwyk_reproduction/03_train_bayesian.py", "content": "#!/usr/bin/env python\n\"\"\"Train a Bayesian MeshNet with optional warm-start from deterministic weights.\n\nSupports the three kwyk model variants:\n - bvwn_multi_prior: Spike-and-slab dropout (default)\n - bayesian_gaussian: Standard Gaussian prior\n - bwn_multi: MC Bernoulli dropout (deterministic model, dropout at inference)\n\nUsage:\n # Spike-and-slab (original kwyk variant)\n python 03_train_bayesian.py --manifest manifest.csv --config config.yaml \\\n --variant bvwn_multi_prior --warmstart checkpoints/meshnet\n\n # Standard Gaussian prior\n python 03_train_bayesian.py --manifest manifest.csv --config config.yaml \\\n --variant bayesian_gaussian --warmstart checkpoints/meshnet\n\n # MC Bernoulli dropout (copies deterministic weights, uses dropout at inference)\n python 03_train_bayesian.py --manifest manifest.csv --config config.yaml \\\n --variant bwn_multi --warmstart checkpoints/meshnet\n\n # Override epochs\n python 03_train_bayesian.py --manifest manifest.csv --config config.yaml \\\n --warmstart checkpoints/meshnet --epochs 100\n\"\"\"\n\nfrom __future__ import annotations\n\nimport argparse\nimport csv\nimport os\nfrom pathlib import Path\nimport time\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nimport torch.nn as nn\nfrom utils import load_config, save_figure, setup_logging\n\nfrom nobrainer.gpu import auto_batch_size, gpu_count\nfrom nobrainer.slurm import SlurmPreemptionHandler, load_checkpoint, save_checkpoint\n\nlog = setup_logging(__name__)\n\n\n# ---------------------------------------------------------------------------\n# ELBO loss: CrossEntropy + KL divergence from Bayesian layers\n# ---------------------------------------------------------------------------\nclass ELBOLoss(nn.Module):\n \"\"\"Evidence Lower Bound loss combining CE and KL divergence.\n\n Parameters\n ----------\n model : nn.Module\n Bayesian model whose layers carry ``.kl`` attributes after\n each forward pass.\n kl_weight : float\n Scaling factor for the KL term. ``1.0`` corresponds to the\n standard variational free-energy; smaller values down-weight\n the regularisation (cold posterior).\n \"\"\"\n\n def __init__(\n self,\n model: nn.Module,\n kl_weight: float = 1.0,\n class_weights: torch.Tensor | None = None,\n ) -> None:\n super().__init__()\n self.ce = nn.CrossEntropyLoss(weight=class_weights)\n self.model = model\n self.kl_weight = kl_weight\n self._last_kl: float = 0.0\n\n def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:\n from nobrainer.models.bayesian.utils import accumulate_kl\n\n ce_loss = self.ce(pred, target)\n kl_loss = accumulate_kl(self.model)\n self._last_kl = kl_loss.item()\n return ce_loss + self.kl_weight * kl_loss\n\n\n# ---------------------------------------------------------------------------\n# CLI\n# ---------------------------------------------------------------------------\ndef parse_args() -> argparse.Namespace:\n \"\"\"Parse command-line arguments.\"\"\"\n parser = argparse.ArgumentParser(\n description=\"Train Bayesian MeshNet with optional warm-start\",\n )\n parser.add_argument(\n \"--manifest\",\n type=str,\n required=True,\n help=\"Path to the dataset manifest CSV\",\n )\n parser.add_argument(\n \"--config\",\n type=str,\n default=\"config.yaml\",\n help=\"Path to YAML configuration file\",\n )\n parser.add_argument(\n \"--output-dir\",\n type=str,\n default=\"checkpoints/bayesian\",\n help=\"Directory for saving model checkpoints and figures\",\n )\n parser.add_argument(\n \"--variant\",\n type=str,\n default=\"bvwn_multi_prior\",\n choices=[\"bvwn_multi_prior\", \"bayesian_gaussian\", \"bwn_multi\"],\n help=(\n \"Model variant: bvwn_multi_prior (spike-and-slab, default), \"\n \"bayesian_gaussian (Gaussian prior), bwn_multi (MC Bernoulli dropout)\"\n ),\n )\n parser.add_argument(\n \"--warmstart\",\n type=str,\n default=None,\n help=\"Path to a trained deterministic MeshNet directory (containing model.pth)\",\n )\n parser.add_argument(\n \"--no-warmstart\",\n action=\"store_true\",\n help=\"Explicitly disable warm-start (train from scratch)\",\n )\n parser.add_argument(\n \"--epochs\",\n type=int,\n default=None,\n help=\"Override number of training epochs from config\",\n )\n return parser.parse_args()\n\n\ndef load_manifest(manifest_path: str, split: str) -> list[tuple[str, str]]:\n \"\"\"Load manifest CSV and return (image, label) pairs for the given split.\"\"\"\n pairs = []\n with open(manifest_path) as f:\n reader = csv.DictReader(f)\n for row in reader:\n if row[\"split\"] == split:\n pairs.append((row[\"t1w_path\"], row[\"label_path\"]))\n return pairs\n\n\n# ---------------------------------------------------------------------------\n# Validation with MC inference\n# ---------------------------------------------------------------------------\ndef evaluate_mc_dice(\n model: nn.Module,\n val_pairs: list[tuple[str, str]],\n block_shape: tuple[int, int, int],\n n_samples: int,\n label_mapping: str | None,\n n_classes: int = 2,\n) -> tuple[list[float], list[float]]:\n \"\"\"Run MC inference on each validation volume.\n\n Returns\n -------\n mean_dices : list[float]\n Mean class Dice across MC samples for each volume.\n std_dices : list[float]\n Std of Dice across MC samples for each volume.\n \"\"\"\n import nibabel as nib\n\n from nobrainer.prediction import predict\n from nobrainer.training import get_device\n\n # Load remap function for multi-class label mappings\n remap_fn = None\n if label_mapping and label_mapping != \"binary\":\n from nobrainer.processing.dataset import _load_label_mapping\n\n remap_fn = _load_label_mapping(label_mapping)\n\n mean_dices: list[float] = []\n std_dices: list[float] = []\n\n device = get_device()\n model = model.to(device)\n\n for img_path, lbl_path in val_pairs:\n gt_arr = np.asarray(nib.load(lbl_path).dataobj, dtype=np.int32)\n if remap_fn is not None:\n gt_arr = remap_fn(torch.from_numpy(gt_arr)).numpy()\n elif label_mapping is None or label_mapping == \"binary\":\n gt_arr = (gt_arr > 0).astype(np.int32)\n\n # Multiple stochastic forward passes\n sample_dices: list[float] = []\n for s in range(n_samples):\n model.train()\n pred_img = predict(\n inputs=img_path,\n model=model,\n block_shape=block_shape,\n batch_size=128,\n return_labels=True,\n )\n pred_arr = np.asarray(pred_img.dataobj, dtype=np.int32)\n if label_mapping is None or label_mapping == \"binary\":\n pred_arr = (pred_arr > 0).astype(np.int32)\n\n # Per-class Dice (skip background)\n class_dices = []\n for c in range(1, n_classes):\n pred_c = pred_arr == c\n gt_c = gt_arr == c\n intersection = (pred_c & gt_c).sum()\n total = pred_c.sum() + gt_c.sum()\n class_dices.append(2.0 * intersection / total if total > 0 else 1.0)\n sample_dices.append(float(np.mean(class_dices)))\n\n vol_mean = float(np.mean(sample_dices))\n vol_std = float(np.std(sample_dices))\n mean_dices.append(vol_mean)\n std_dices.append(vol_std)\n log.info(\n \" Val volume %s: MC Dice=%.4f +/- %.4f (%d samples)\",\n Path(img_path).name,\n vol_mean,\n vol_std,\n n_samples,\n )\n\n return mean_dices, std_dices\n\n\n# ---------------------------------------------------------------------------\n# Learning curve with uncertainty bands\n# ---------------------------------------------------------------------------\ndef plot_learning_curve(\n train_losses: list[float],\n val_losses: list[float],\n val_dice_means: list[float],\n val_dice_stds: list[float],\n kl_terms: list[float],\n output_path: Path,\n) -> None:\n \"\"\"Generate learning curve with uncertainty bands.\n\n Left y-axis: train loss, val loss, KL term.\n Right y-axis: mean MC Dice with +/- std shading.\n \"\"\"\n epochs = list(range(1, len(train_losses) + 1))\n\n fig, ax_loss = plt.subplots(figsize=(12, 7))\n ax_dice = ax_loss.twinx()\n\n # Loss curves\n ax_loss.plot(epochs, train_losses, \"b-\", label=\"Train Loss (ELBO)\")\n if val_losses:\n ax_loss.plot(epochs, val_losses, \"b--\", alpha=0.7, label=\"Val Loss\")\n if kl_terms:\n ax_loss.plot(epochs, kl_terms, \"g-.\", alpha=0.6, label=\"KL Term\")\n ax_loss.set_xlabel(\"Epoch\")\n ax_loss.set_ylabel(\"Loss / KL\", color=\"b\")\n ax_loss.tick_params(axis=\"y\", labelcolor=\"b\")\n\n # Dice with uncertainty bands\n if val_dice_means:\n means = np.array(val_dice_means)\n stds = np.array(val_dice_stds)\n dice_epochs = list(range(1, len(means) + 1))\n ax_dice.plot(\n dice_epochs, means, \"r-o\", markersize=3, label=\"Val MC Dice (mean)\"\n )\n ax_dice.fill_between(\n dice_epochs,\n np.clip(means - stds, 0, 1),\n np.clip(means + stds, 0, 1),\n color=\"r\",\n alpha=0.15,\n label=\"Val MC Dice (+/- std)\",\n )\n ax_dice.set_ylabel(\"Dice Score\", color=\"r\")\n ax_dice.tick_params(axis=\"y\", labelcolor=\"r\")\n ax_dice.set_ylim(0.0, 1.0)\n\n fig.suptitle(\"Bayesian MeshNet Training — ELBO Loss & MC Dice\")\n fig.tight_layout()\n\n lines_loss, labels_loss = ax_loss.get_legend_handles_labels()\n lines_dice, labels_dice = ax_dice.get_legend_handles_labels()\n ax_loss.legend(\n lines_loss + lines_dice,\n labels_loss + labels_dice,\n loc=\"center right\",\n )\n\n save_figure(fig, output_path)\n plt.close(fig)\n log.info(\"Learning curve saved to %s\", output_path)\n\n\n# ---------------------------------------------------------------------------\n# Training loop (lower-level, using nobrainer.training.fit)\n# ---------------------------------------------------------------------------\ndef train_bayesian(\n model: nn.Module,\n train_loader,\n val_loader,\n elbo_loss: ELBOLoss,\n optimizer: torch.optim.Optimizer,\n epochs: int,\n val_pairs: list[tuple[str, str]],\n block_shape: tuple[int, int, int],\n n_samples: int,\n label_mapping: str | None,\n n_classes: int,\n checkpoint_dir: Path,\n preemption_handler: SlurmPreemptionHandler | None = None,\n callbacks: list | None = None,\n) -> dict:\n \"\"\"Custom training loop for Bayesian MeshNet with ELBO loss.\n\n Supports checkpoint/resume for SLURM preemptible jobs. When a\n preemption signal is received, the loop checkpoints and exits so\n the job can be requeued.\n \"\"\"\n from nobrainer.training import get_device\n\n device = get_device()\n model = model.to(device)\n\n # -- Resume from checkpoint if available --------------------------------\n start_epoch, prev_metrics = load_checkpoint(checkpoint_dir, model, optimizer)\n\n # Restore accumulated metrics from prior runs\n train_losses: list[float] = prev_metrics.get(\"train_losses\", [])\n val_losses_list: list[float] = prev_metrics.get(\"val_losses\", [])\n val_dice_means: list[float] = prev_metrics.get(\"val_dice_means\", [])\n val_dice_stds: list[float] = prev_metrics.get(\"val_dice_stds\", [])\n kl_terms: list[float] = prev_metrics.get(\"kl_terms\", [])\n best_loss: float = prev_metrics.get(\"best_loss\", float(\"inf\"))\n\n if start_epoch >= epochs:\n log.info(\"Already completed %d/%d epochs — nothing to do\", start_epoch, epochs)\n return {\n \"train_losses\": train_losses,\n \"val_losses\": val_losses_list,\n \"val_dice_means\": val_dice_means,\n \"val_dice_stds\": val_dice_stds,\n \"kl_terms\": kl_terms,\n \"best_loss\": best_loss,\n \"epochs_completed\": start_epoch,\n }\n\n for epoch in range(start_epoch, epochs):\n t_epoch = time.time()\n\n # -- Train one epoch --------------------------------------------------\n model.train()\n epoch_loss = 0.0\n epoch_kl = 0.0\n n_batches = 0\n\n for batch in train_loader:\n if isinstance(batch, dict):\n images = batch[\"image\"].to(device)\n labels = batch[\"label\"].to(device)\n elif isinstance(batch, (list, tuple)):\n images = batch[0].to(device)\n labels = batch[1].to(device)\n else:\n raise TypeError(f\"Unsupported batch type: {type(batch)}\")\n\n # Squeeze channel dim from labels if present\n if labels.ndim == images.ndim and labels.shape[1] == 1:\n labels = labels.squeeze(1)\n if labels.dtype in (torch.float32, torch.float64):\n labels = labels.long()\n\n optimizer.zero_grad()\n # Match original TF: deterministic VWN weights + stochastic dropout\n # (is_mc_v=False, is_mc_b=True in meshnetbwn.py)\n try:\n pred = model(images, mc_vwn=False, mc_dropout=True)\n except TypeError:\n pred = model(images)\n loss = elbo_loss(pred, labels)\n loss.backward()\n optimizer.step()\n\n epoch_loss += loss.item()\n epoch_kl += elbo_loss._last_kl\n n_batches += 1\n\n avg_loss = epoch_loss / max(n_batches, 1)\n avg_kl = epoch_kl / max(n_batches, 1)\n train_losses.append(avg_loss)\n kl_terms.append(avg_kl)\n\n # -- Validate ---------------------------------------------------------\n val_loss = 0.0\n if val_loader is not None:\n model.eval()\n n_val = 0\n with torch.no_grad():\n for batch in val_loader:\n if isinstance(batch, dict):\n images = batch[\"image\"].to(device)\n labels = batch[\"label\"].to(device)\n elif isinstance(batch, (list, tuple)):\n images = batch[0].to(device)\n labels = batch[1].to(device)\n else:\n raise TypeError(f\"Unsupported batch type: {type(batch)}\")\n\n if labels.ndim == images.ndim and labels.shape[1] == 1:\n labels = labels.squeeze(1)\n if labels.dtype in (torch.float32, torch.float64):\n labels = labels.long()\n\n try:\n pred = model(images, mc_vwn=False, mc_dropout=False)\n except TypeError:\n pred = model(images)\n loss = elbo_loss(pred, labels)\n val_loss += loss.item()\n n_val += 1\n val_loss = val_loss / max(n_val, 1)\n val_losses_list.append(val_loss)\n\n # -- MC Dice evaluation (every 10 epochs or last epoch) ---------------\n if val_pairs and (epoch == epochs - 1 or (epoch + 1) % 10 == 0):\n mean_dices, std_dices = evaluate_mc_dice(\n model, val_pairs, block_shape, n_samples, label_mapping, n_classes\n )\n overall_mean = float(np.mean(mean_dices)) if mean_dices else 0.0\n overall_std = float(np.mean(std_dices)) if std_dices else 0.0\n val_dice_means.append(overall_mean)\n val_dice_stds.append(overall_std)\n else:\n if val_dice_means:\n val_dice_means.append(val_dice_means[-1])\n val_dice_stds.append(val_dice_stds[-1])\n else:\n val_dice_means.append(float(\"nan\"))\n val_dice_stds.append(float(\"nan\"))\n\n # -- Checkpoint best --------------------------------------------------\n if avg_loss < best_loss:\n best_loss = avg_loss\n torch.save(model.state_dict(), checkpoint_dir / \"best_model.pth\")\n\n # -- Always save resumable checkpoint ---------------------------------\n metrics = {\n \"train_losses\": train_losses,\n \"val_losses\": val_losses_list,\n \"val_dice_means\": val_dice_means,\n \"val_dice_stds\": val_dice_stds,\n \"kl_terms\": kl_terms,\n \"best_loss\": best_loss,\n }\n save_checkpoint(checkpoint_dir, model, optimizer, epoch, metrics)\n\n elapsed = time.time() - t_epoch\n log.info(\n \"Epoch %d/%d: train_loss=%.6f val_loss=%.6f kl=%.6f \" \"dice=%.4f (%.1fs)\",\n epoch + 1,\n epochs,\n avg_loss,\n val_loss,\n avg_kl,\n val_dice_means[-1] if val_dice_means else 0.0,\n elapsed,\n )\n\n # -- Callbacks -----------------------------------------------------------\n for cb in callbacks or []:\n cb(epoch, avg_loss, model)\n\n # -- Check for SLURM preemption signal --------------------------------\n if preemption_handler and preemption_handler.preempted:\n log.warning(\n \"Preemption detected after epoch %d — exiting for requeue\",\n epoch + 1,\n )\n break\n\n return {\n \"train_losses\": train_losses,\n \"val_losses\": val_losses_list,\n \"val_dice_means\": val_dice_means,\n \"val_dice_stds\": val_dice_stds,\n \"kl_terms\": kl_terms,\n \"best_loss\": best_loss,\n \"epochs_completed\": epoch + 1,\n }\n\n\n# ---------------------------------------------------------------------------\n# Main\n# ---------------------------------------------------------------------------\n\n\ndef main() -> None:\n \"\"\"Train Bayesian MeshNet with optional warm-start.\"\"\"\n args = parse_args()\n t_start = time.time()\n\n # ---- Load config --------------------------------------------------------\n config = load_config(args.config)\n epochs = (\n args.epochs if args.epochs is not None else config.get(\"bayesian_epochs\", 50)\n )\n n_classes = config[\"n_classes\"]\n block_shape = tuple(config[\"block_shape\"])\n batch_size = config[\"batch_size\"]\n lr = config.get(\"lr\", 1e-4)\n n_samples = config.get(\"n_samples\", 10)\n label_mapping = config.get(\"label_mapping\", \"binary\")\n output_dir = Path(args.output_dir)\n output_dir.mkdir(parents=True, exist_ok=True)\n\n use_warmstart = args.warmstart is not None and not args.no_warmstart\n\n # ---- Load variant config from config.yaml variants section --------------\n variant = args.variant\n variants = config.get(\"variants\", {})\n variant_config = variants.get(variant, {})\n log.info(\"Model variant: %s — %s\", variant, variant_config.get(\"description\", \"\"))\n\n # ---- Determine KL weight from config ------------------------------------\n kl_weight = variant_config.get(\"kl_weight\", config.get(\"kl_weight\", 1.0))\n\n log.info(\"Config loaded from %s\", args.config)\n log.info(\n \"Training KWYK MeshNet (%s): epochs=%d, n_classes=%d, \"\n \"block_shape=%s, kl_weight=%.4f, warmstart=%s\",\n variant,\n epochs,\n n_classes,\n block_shape,\n kl_weight,\n use_warmstart,\n )\n\n # ---- Load manifest and build datasets -----------------------------------\n train_pairs = load_manifest(args.manifest, split=\"train\")\n val_pairs = load_manifest(args.manifest, split=\"val\")\n log.info(\"Manifest: %d train, %d val volumes\", len(train_pairs), len(val_pairs))\n\n if not train_pairs:\n log.error(\"No training volumes found in manifest. Exiting.\")\n return\n\n # ---- Build KWYK MeshNet first (needed for auto batch size) ----------------\n from nobrainer.models import get as get_model\n from nobrainer.processing.dataset import Dataset\n\n dropout_type = variant_config.get(\"dropout_type\", \"bernoulli\")\n model_args = {\n \"n_classes\": n_classes,\n \"filters\": config.get(\"filters\", 96),\n \"receptive_field\": config.get(\"receptive_field\", 37),\n \"dropout_type\": dropout_type,\n \"dropout_rate\": config.get(\"dropout_rate\", 0.25),\n \"sigma_init\": config.get(\"sigma_init\", 1e-4),\n }\n\n # Concrete dropout specific params\n if dropout_type == \"concrete\":\n model_args[\"concrete_temperature\"] = variant_config.get(\n \"concrete_temperature\", 0.02\n )\n model_args[\"concrete_init_p\"] = variant_config.get(\"concrete_init_p\", 0.9)\n\n log.info(\"Model args: %s\", model_args)\n\n bayesian_model = get_model(\"kwyk_meshnet\")(**model_args)\n log.info(\n \"KWYK MeshNet (%s, %s) created: %d parameters\",\n variant,\n dropout_type,\n sum(p.numel() for p in bayesian_model.parameters()),\n )\n\n # ---- Auto batch size with training-mode profiling -------------------------\n n_gpus = gpu_count()\n if n_gpus > 0:\n optimal_per_gpu = auto_batch_size(\n bayesian_model,\n block_shape,\n n_classes=n_classes,\n target_memory_fraction=0.90,\n forward_kwargs={\"mc_vwn\": False, \"mc_dropout\": True},\n )\n log.info(\n \"Auto batch size: %d (profiled with mc_vwn=False, mc_dropout=True, \"\n \"config batch_size=%d)\",\n optimal_per_gpu,\n batch_size,\n )\n batch_size = optimal_per_gpu\n\n # ---- Build datasets with optimized batch size ----------------------------\n\n # Use streaming mode: extract multiple patches per volume to fill GPU.\n # Use Zarr store if available, else fall back to NIfTI with streaming\n patches_per_volume = config.get(\"patches_per_volume\", 50)\n zarr_store = config.get(\"zarr_store\")\n\n if zarr_store and Path(zarr_store).exists():\n log.info(\"Using Zarr store: %s\", zarr_store)\n ds_train = (\n Dataset.from_zarr(\n zarr_store,\n block_shape=block_shape,\n n_classes=n_classes,\n partition=\"train\",\n )\n .batch(batch_size)\n .binarize(label_mapping)\n .streaming(patches_per_volume=patches_per_volume)\n )\n else:\n ds_train = (\n Dataset.from_files(\n train_pairs, block_shape=block_shape, n_classes=n_classes\n )\n .batch(batch_size)\n .binarize(label_mapping)\n .streaming(patches_per_volume=patches_per_volume)\n )\n train_loader = ds_train.dataloader\n n_train = len(ds_train.data) if hasattr(ds_train, \"data\") else len(train_pairs)\n log.info(\n \"Training data: %d volumes × %d patches = %d blocks/epoch, batch_size=%d\",\n n_train,\n patches_per_volume,\n n_train * patches_per_volume,\n batch_size,\n )\n\n ds_val = None\n val_loader = None\n if val_pairs:\n if zarr_store and Path(zarr_store).exists():\n ds_val = (\n Dataset.from_zarr(\n zarr_store,\n block_shape=block_shape,\n n_classes=n_classes,\n partition=\"val\",\n )\n .batch(batch_size)\n .binarize(label_mapping)\n .streaming(patches_per_volume=patches_per_volume)\n )\n else:\n ds_val = (\n Dataset.from_files(\n val_pairs, block_shape=block_shape, n_classes=n_classes\n )\n .batch(batch_size)\n .binarize(label_mapping)\n .streaming(patches_per_volume=patches_per_volume)\n )\n val_loader = ds_val.dataloader\n\n # ---- Optional warm-start ------------------------------------------------\n if use_warmstart:\n warmstart_dir = Path(args.warmstart)\n det_weights_path = warmstart_dir / \"model.pth\"\n\n if not det_weights_path.exists():\n log.error(\n \"Warm-start weights not found at %s. \"\n \"Train a deterministic MeshNet first with 02_train_meshnet.py.\",\n det_weights_path,\n )\n return\n\n log.info(\"Loading deterministic weights from %s\", det_weights_path)\n\n from nobrainer.models.bayesian.warmstart import (\n warmstart_kwyk_from_deterministic,\n )\n\n n_transferred = warmstart_kwyk_from_deterministic(\n bayesian_model,\n det_weights_path,\n get_model,\n )\n log.info(\"Warm-started %d layers from deterministic model\", n_transferred)\n else:\n log.info(\"Training KWYK MeshNet from scratch (no warm-start)\")\n\n # ---- Class weights (important for 50-class parcellation) -----------------\n class_weights = None\n weight_method = config.get(\"class_weight_method\")\n if weight_method and weight_method != \"null\":\n from nobrainer.losses import compute_class_weights\n\n label_paths = [p[1] for p in train_pairs]\n class_weights = compute_class_weights(\n label_paths,\n n_classes,\n label_mapping=label_mapping,\n method=weight_method,\n max_samples=50,\n )\n log.info(\n \"Class weights computed (%s): min=%.3f, max=%.3f, mean=%.3f\",\n weight_method,\n class_weights.min(),\n class_weights.max(),\n class_weights.mean(),\n )\n # Move weights to device\n from nobrainer.training import get_device\n\n class_weights = class_weights.to(get_device())\n\n # ---- ELBO loss and optimiser --------------------------------------------\n elbo_loss = ELBOLoss(\n bayesian_model, kl_weight=kl_weight, class_weights=class_weights\n )\n optimizer = torch.optim.Adam(bayesian_model.parameters(), lr=lr)\n\n # ---- SLURM preemption handler (no-op if not on SLURM) -----------------\n preemption = None\n if os.environ.get(\"SLURM_JOB_ID\"):\n preemption = SlurmPreemptionHandler()\n\n # ---- Experiment tracker (local + optional W&B) -------------------------\n from nobrainer.experiment import ExperimentTracker\n\n tracker = ExperimentTracker(\n output_dir=output_dir,\n config={\n \"variant\": variant,\n \"dropout_type\": dropout_type,\n \"n_classes\": n_classes,\n \"filters\": config.get(\"filters\", 96),\n \"block_shape\": list(block_shape),\n \"batch_size\": batch_size,\n \"lr\": lr,\n \"kl_weight\": kl_weight,\n \"epochs\": epochs,\n \"warmstart\": use_warmstart,\n },\n project=\"kwyk-reproduction\",\n name=variant,\n tags=[variant, f\"{n_classes}-class\"],\n )\n\n # ---- Train --------------------------------------------------------------\n result = train_bayesian(\n model=bayesian_model,\n train_loader=train_loader,\n val_loader=val_loader,\n elbo_loss=elbo_loss,\n optimizer=optimizer,\n epochs=epochs,\n val_pairs=val_pairs,\n block_shape=block_shape,\n n_samples=n_samples,\n label_mapping=label_mapping,\n n_classes=n_classes,\n checkpoint_dir=output_dir,\n preemption_handler=preemption,\n callbacks=[tracker.callback(variant=variant)],\n )\n\n # ---- Learning curve with uncertainty bands ------------------------------\n fig_path = output_dir / \"learning_curve.png\"\n plot_learning_curve(\n train_losses=result[\"train_losses\"],\n val_losses=result[\"val_losses\"],\n val_dice_means=result[\"val_dice_means\"],\n val_dice_stds=result[\"val_dice_stds\"],\n kl_terms=result[\"kl_terms\"],\n output_path=fig_path,\n )\n\n # ---- Save with Croissant-ML metadata ------------------------------------\n # Save final weights\n torch.save(bayesian_model.state_dict(), output_dir / \"model.pth\")\n\n # Use Segmentation estimator's save for Croissant metadata\n from nobrainer.processing.segmentation import Segmentation\n\n seg = Segmentation(\n base_model=\"kwyk_meshnet\",\n model_args=model_args,\n )\n seg.model_ = bayesian_model\n seg.block_shape_ = block_shape\n seg.n_classes_ = n_classes\n seg._optimizer_class = \"Adam\"\n seg._optimizer_args = {\"lr\": lr}\n seg._loss_name = \"ELBOLoss\"\n seg._training_result = {\n \"variant\": variant,\n \"dropout_type\": dropout_type,\n \"final_loss\": result[\"train_losses\"][-1] if result[\"train_losses\"] else 0.0,\n \"best_loss\": result[\"best_loss\"],\n \"epochs_completed\": result[\"epochs_completed\"],\n \"checkpoint_path\": str(output_dir / \"best_model.pth\"),\n }\n seg._dataset = ds_train\n seg.save(output_dir)\n log.info(\"Model and Croissant-ML metadata saved to %s\", output_dir)\n\n # ---- Summary ------------------------------------------------------------\n elapsed = time.time() - t_start\n final_dice = (\n result[\"val_dice_means\"][-1] if result[\"val_dice_means\"] else float(\"nan\")\n )\n log.info(\"=\" * 60)\n log.info(\"Bayesian MeshNet training complete (%s)\", variant)\n log.info(\" Output directory : %s\", output_dir)\n log.info(\" Variant : %s\", variant)\n log.info(\" Dropout type : %s\", dropout_type)\n log.info(\" Epochs : %d\", epochs)\n log.info(\" Warm-start : %s\", \"yes\" if use_warmstart else \"no\")\n log.info(\" KL weight : %.4f\", kl_weight)\n log.info(\n \" Final train loss : %.6f\",\n result[\"train_losses\"][-1] if result[\"train_losses\"] else 0.0,\n )\n log.info(\" Best train loss : %.6f\", result[\"best_loss\"])\n log.info(\" Val MC Dice : %.4f\", final_dice)\n log.info(\" MC samples : %d\", n_samples)\n log.info(\" Elapsed time : %.1f s\", elapsed)\n log.info(\"=\" * 60)\n\n tracker.finish()\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "scripts/kwyk_reproduction/04_evaluate.py", "content": "#!/usr/bin/env python\n\"\"\"Evaluate a trained segmentation model on test volumes.\n\nComputes per-class Dice for each volume (matching McClure et al. 2019,\nSection 2.4.1, Eq. 19), then averages across classes per volume. The\nreported \"class Dice\" in Table 3 of the paper is the mean ± std of\nthese per-volume average Dice scores.\n\nFor Bayesian models, MC inference produces variance and entropy maps\n(Eq. 20) saved as NIfTI files.\n\nUsage:\n python 04_evaluate.py --model checkpoints/bvwn_multi_prior \\\n --manifest manifest.csv --split test --n-samples 10\n\"\"\"\n\nfrom __future__ import annotations\n\nimport argparse\nimport csv\nfrom pathlib import Path\n\nimport matplotlib.pyplot as plt\nimport nibabel as nib\nimport numpy as np\nimport torch\nfrom utils import load_config, save_figure, setup_logging\n\nlog = setup_logging(__name__)\n\n\ndef parse_args() -> argparse.Namespace:\n parser = argparse.ArgumentParser(description=\"Evaluate segmentation model\")\n parser.add_argument(\"--model\", type=str, required=True)\n parser.add_argument(\"--manifest\", type=str, required=True)\n parser.add_argument(\"--config\", type=str, default=\"config.yaml\")\n parser.add_argument(\"--split\", type=str, default=\"test\")\n parser.add_argument(\"--n-samples\", type=int, default=10)\n parser.add_argument(\"--output-dir\", type=str, default=\"results\")\n return parser.parse_args()\n\n\ndef load_manifest(manifest_path: str, split: str) -> list[tuple[str, str]]:\n pairs = []\n with open(manifest_path) as f:\n for row in csv.DictReader(f):\n if row[\"split\"] == split:\n pairs.append((row[\"t1w_path\"], row[\"label_path\"]))\n return pairs\n\n\ndef per_class_dice(\n pred: np.ndarray,\n gt: np.ndarray,\n n_classes: int,\n) -> np.ndarray:\n \"\"\"Compute Dice coefficient for each class c = 1..n_classes-1.\n\n Matches Eq. 19 in McClure et al. (2019):\n Dice_c = 2*TP_c / (2*TP_c + FN_c + FP_c)\n\n Class 0 (background / unknown) is excluded, matching the paper:\n \"averaging across all output voxels not classified as background\".\n\n Parameters\n ----------\n pred : np.ndarray\n Integer label predictions.\n gt : np.ndarray\n Integer ground truth labels.\n n_classes : int\n Total number of classes (including background).\n\n Returns\n -------\n np.ndarray\n Shape ``(n_classes - 1,)`` — Dice for classes 1..n_classes-1.\n \"\"\"\n dice_scores = np.zeros(n_classes - 1)\n for c in range(1, n_classes):\n pred_c = (pred == c).astype(np.float64)\n gt_c = (gt == c).astype(np.float64)\n intersection = (pred_c * gt_c).sum()\n total = pred_c.sum() + gt_c.sum()\n if total > 0:\n dice_scores[c - 1] = 2.0 * intersection / total\n else:\n # Both empty for this class — perfect agreement\n dice_scores[c - 1] = 1.0\n return dice_scores\n\n\ndef compute_entropy(prob_map: np.ndarray) -> np.ndarray:\n \"\"\"Compute entropy of softmax probabilities (Eq. 20).\n\n H(y|x) = -sum_c p(y_c|x) log p(y_c|x)\n \"\"\"\n eps = 1e-10\n return -(prob_map * np.log(prob_map + eps)).sum(axis=0)\n\n\ndef plot_prediction_overlay(\n t1w_arr: np.ndarray,\n pred_arr: np.ndarray,\n gt_arr: np.ndarray,\n output_path: Path,\n title: str = \"Prediction Overlay\",\n) -> None:\n \"\"\"3-panel figure: T1w, prediction, ground truth (middle axial slice).\"\"\"\n mid = t1w_arr.shape[2] // 2\n t1 = t1w_arr[:, :, mid]\n pred = pred_arr[:, :, mid]\n gt = gt_arr[:, :, mid]\n\n fig, axes = plt.subplots(1, 3, figsize=(15, 5))\n axes[0].imshow(t1.T, cmap=\"gray\", origin=\"lower\")\n axes[0].set_title(\"T1w Input\")\n axes[0].axis(\"off\")\n\n axes[1].imshow(t1.T, cmap=\"gray\", origin=\"lower\")\n axes[1].imshow(\n pred.T,\n cmap=\"nipy_spectral\",\n alpha=0.4,\n origin=\"lower\",\n vmin=0,\n vmax=max(pred.max(), 1),\n )\n axes[1].set_title(\"Prediction\")\n axes[1].axis(\"off\")\n\n axes[2].imshow(t1.T, cmap=\"gray\", origin=\"lower\")\n axes[2].imshow(\n gt.T,\n cmap=\"nipy_spectral\",\n alpha=0.4,\n origin=\"lower\",\n vmin=0,\n vmax=max(gt.max(), 1),\n )\n axes[2].set_title(\"Ground Truth\")\n axes[2].axis(\"off\")\n\n fig.suptitle(title)\n fig.tight_layout()\n save_figure(fig, output_path)\n plt.close(fig)\n\n\ndef plot_per_class_dice(\n class_dice_all: np.ndarray,\n class_names: list[str] | None,\n output_path: Path,\n) -> None:\n \"\"\"Bar chart of mean per-class Dice across all volumes.\"\"\"\n mean_dice = class_dice_all.mean(axis=0)\n std_dice = class_dice_all.std(axis=0)\n n = len(mean_dice)\n\n fig, ax = plt.subplots(figsize=(max(12, n * 0.3), 6))\n x = np.arange(n)\n ax.bar(x, mean_dice, yerr=std_dice, capsize=2, alpha=0.7, color=\"steelblue\")\n ax.set_xlabel(\"Class\")\n ax.set_ylabel(\"Dice\")\n ax.set_title(\"Per-Class Dice (mean ± std across volumes)\")\n ax.set_ylim(0, 1.05)\n if class_names and len(class_names) == n:\n ax.set_xticks(x)\n ax.set_xticklabels(class_names, rotation=90, fontsize=6)\n fig.tight_layout()\n save_figure(fig, output_path)\n plt.close(fig)\n\n\ndef main() -> None:\n args = parse_args()\n output_dir = Path(args.output_dir)\n output_dir.mkdir(parents=True, exist_ok=True)\n fig_dir = output_dir / \"figures\"\n fig_dir.mkdir(parents=True, exist_ok=True)\n\n # ---- Load config for n_classes and label_mapping ------------------------\n config = load_config(args.config)\n n_classes = config.get(\"n_classes\", 2)\n label_mapping = config.get(\"label_mapping\", \"binary\")\n\n # Load label names for plots\n class_names = None\n if label_mapping and label_mapping != \"binary\":\n mapping_path = (\n Path(__file__).parent / \"label_mappings\" / f\"{label_mapping}-mapping.csv\"\n )\n if mapping_path.exists():\n with open(mapping_path) as f:\n reader = csv.DictReader(f)\n rows = list(reader)\n # Build class_names indexed by 'new' column (skip background=0)\n name_map = {}\n for r in rows:\n new_id = int(r[\"new\"])\n if new_id > 0 and new_id not in name_map:\n name_map[new_id] = r.get(\"label\", str(new_id))\n class_names = [name_map.get(i, str(i)) for i in range(1, n_classes)]\n\n # Load remap function for ground truth\n remap_fn = None\n if label_mapping and label_mapping != \"binary\":\n from nobrainer.processing.dataset import _load_label_mapping\n\n remap_fn = _load_label_mapping(label_mapping)\n\n # ---- Load model ---------------------------------------------------------\n from nobrainer.processing.segmentation import Segmentation\n\n log.info(\"Loading model from %s\", args.model)\n seg = Segmentation.load(args.model)\n block_shape = seg.block_shape_ or tuple(config[\"block_shape\"])\n log.info(\n \"Model: %s, block_shape=%s, n_classes=%s\",\n seg.base_model,\n block_shape,\n n_classes,\n )\n\n # ---- Load manifest ------------------------------------------------------\n pairs = load_manifest(args.manifest, split=args.split)\n log.info(\"Evaluating %d volumes from split '%s'\", len(pairs), args.split)\n\n if not pairs:\n log.error(\"No volumes for split '%s'\", args.split)\n return\n\n # ---- Evaluate each volume -----------------------------------------------\n results: list[dict] = []\n all_class_dice: list[np.ndarray] = []\n n_samples = args.n_samples\n\n for idx, (img_path, lbl_path) in enumerate(pairs):\n vol_name = Path(img_path).stem\n log.info(\"Volume %d/%d: %s\", idx + 1, len(pairs), vol_name)\n\n # Load and remap ground truth\n gt_arr = np.asarray(nib.load(lbl_path).dataobj, dtype=np.int32)\n if remap_fn is not None:\n gt_tensor = torch.from_numpy(gt_arr)\n gt_arr = remap_fn(gt_tensor).numpy().astype(np.int32)\n elif label_mapping == \"binary\":\n gt_arr = (gt_arr > 0).astype(np.int32)\n\n t1w_arr = np.asarray(nib.load(img_path).dataobj, dtype=np.float32)\n\n # Predict (batch_size=128 to utilize GPU memory with 32³ blocks)\n if n_samples > 0:\n pred_result = seg.predict(\n img_path,\n block_shape=block_shape,\n n_samples=n_samples,\n batch_size=128,\n )\n if isinstance(pred_result, tuple):\n label_img, var_img, entropy_img = pred_result\n nib.save(var_img, str(output_dir / f\"{vol_name}_variance.nii.gz\"))\n nib.save(entropy_img, str(output_dir / f\"{vol_name}_entropy.nii.gz\"))\n else:\n label_img = pred_result\n else:\n label_img = seg.predict(img_path, block_shape=block_shape, batch_size=128)\n\n pred_arr = np.asarray(label_img.dataobj, dtype=np.int32)\n\n # Per-class Dice (Eq. 19)\n class_dice = per_class_dice(pred_arr, gt_arr, n_classes)\n avg_dice = float(class_dice.mean())\n all_class_dice.append(class_dice)\n\n log.info(\n \" Avg class Dice = %.4f (min=%.4f, max=%.4f)\",\n avg_dice,\n class_dice.min(),\n class_dice.max(),\n )\n\n results.append(\n {\n \"volume\": vol_name,\n \"image_path\": img_path,\n \"avg_class_dice\": avg_dice,\n \"min_class_dice\": float(class_dice.min()),\n \"max_class_dice\": float(class_dice.max()),\n }\n )\n\n # Overlay figure\n plot_prediction_overlay(\n t1w_arr,\n pred_arr.astype(np.float32),\n gt_arr.astype(np.float32),\n fig_dir / f\"{vol_name}_overlay.png\",\n title=f\"{vol_name} — Avg Dice={avg_dice:.4f}\",\n )\n\n # ---- Per-class Dice bar chart -------------------------------------------\n class_dice_matrix = np.array(all_class_dice) # (n_volumes, n_classes-1)\n plot_per_class_dice(class_dice_matrix, class_names, fig_dir / \"per_class_dice.png\")\n\n # ---- Save CSV with per-volume results -----------------------------------\n csv_path = output_dir / \"dice_scores.csv\"\n with open(csv_path, \"w\", newline=\"\") as f:\n writer = csv.DictWriter(\n f,\n fieldnames=[\n \"volume\",\n \"image_path\",\n \"avg_class_dice\",\n \"min_class_dice\",\n \"max_class_dice\",\n ],\n )\n writer.writeheader()\n writer.writerows(results)\n\n # ---- Save per-class Dice matrix -----------------------------------------\n np.save(output_dir / \"per_class_dice.npy\", class_dice_matrix)\n\n # ---- Summary (matching Table 3 format) ----------------------------------\n avg_dices = [r[\"avg_class_dice\"] for r in results]\n log.info(\"=\" * 60)\n log.info(\"Evaluation Summary (%s split, %d-class)\", args.split, n_classes)\n log.info(\" Volumes : %d\", len(avg_dices))\n log.info(\" MC samples : %d\", n_samples)\n log.info(\" Class Dice : %.4f ± %.4f\", np.mean(avg_dices), np.std(avg_dices))\n log.info(\" Min volume Dice : %.4f\", np.min(avg_dices))\n log.info(\" Max volume Dice : %.4f\", np.max(avg_dices))\n\n # Per-class summary: median and range across volumes\n mean_per_class = class_dice_matrix.mean(axis=0) # (n_classes-1,)\n log.info(\n \" Per-class Dice : median=%.4f, range=[%.4f, %.4f]\",\n np.median(mean_per_class),\n mean_per_class.min(),\n mean_per_class.max(),\n )\n if class_names:\n worst_5 = np.argsort(mean_per_class)[:5]\n best_5 = np.argsort(mean_per_class)[-5:][::-1]\n log.info(\n \" Worst 5 classes : %s\",\n \", \".join(f\"{class_names[i]}={mean_per_class[i]:.3f}\" for i in worst_5),\n )\n log.info(\n \" Best 5 classes : %s\",\n \", \".join(f\"{class_names[i]}={mean_per_class[i]:.3f}\" for i in best_5),\n )\n\n # Save per-class summary CSV\n per_class_csv = output_dir / \"per_class_dice_summary.csv\"\n with open(per_class_csv, \"w\", newline=\"\") as f:\n writer = csv.writer(f)\n writer.writerow(\n [\n \"class_id\",\n \"class_name\",\n \"mean_dice\",\n \"median_dice\",\n \"min_dice\",\n \"max_dice\",\n ]\n )\n for i in range(len(mean_per_class)):\n name = class_names[i] if class_names else str(i + 1)\n col = class_dice_matrix[:, i]\n writer.writerow(\n [\n i + 1,\n name,\n f\"{col.mean():.4f}\",\n f\"{np.median(col):.4f}\",\n f\"{col.min():.4f}\",\n f\"{col.max():.4f}\",\n ]\n )\n log.info(\" Per-class summary : %s\", per_class_csv)\n\n log.info(\" Output : %s\", output_dir)\n log.info(\"=\" * 60)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "scripts/kwyk_reproduction/05_compare_kwyk.py", "content": "#!/usr/bin/env python\n\"\"\"Compare new model predictions against original kwyk container.\n\nUsage:\n python 05_compare_kwyk.py \\\n --new-model checkpoints/bayesian \\\n --kwyk-dir /path/to/kwyk \\\n --manifest manifest.csv \\\n --split test \\\n --output-dir results/comparison\n\"\"\"\n\nfrom __future__ import annotations\n\nimport argparse\nimport csv\nfrom pathlib import Path\nimport subprocess\n\nimport matplotlib.pyplot as plt\nimport nibabel as nib\nimport numpy as np\nfrom utils import compute_dice, save_figure, setup_logging\n\nlog = setup_logging(__name__)\n\n\n# ---------------------------------------------------------------------------\n# CLI\n# ---------------------------------------------------------------------------\ndef parse_args() -> argparse.Namespace:\n \"\"\"Parse command-line arguments.\"\"\"\n parser = argparse.ArgumentParser(\n description=\"Compare new model vs original kwyk predictions\",\n )\n parser.add_argument(\n \"--new-model\",\n type=str,\n required=True,\n help=\"Path to new model directory (model.pth + croissant.json)\",\n )\n parser.add_argument(\n \"--kwyk-dir\",\n type=str,\n required=True,\n help=\"Path to original kwyk repository (containing kwyk/cli.py)\",\n )\n parser.add_argument(\n \"--manifest\",\n type=str,\n required=True,\n help=\"Path to the dataset manifest CSV\",\n )\n parser.add_argument(\n \"--split\",\n type=str,\n default=\"test\",\n help=\"Which split to evaluate on (default: test)\",\n )\n parser.add_argument(\n \"--output-dir\",\n type=str,\n default=\"results/comparison\",\n help=\"Directory for comparison outputs\",\n )\n return parser.parse_args()\n\n\ndef load_manifest(manifest_path: str, split: str) -> list[tuple[str, str]]:\n \"\"\"Load manifest CSV and return (image, label) pairs for the given split.\"\"\"\n pairs = []\n with open(manifest_path) as f:\n reader = csv.DictReader(f)\n for row in reader:\n if row[\"split\"] == split:\n pairs.append((row[\"t1w_path\"], row[\"label_path\"]))\n return pairs\n\n\n# ---------------------------------------------------------------------------\n# Run original kwyk prediction\n# ---------------------------------------------------------------------------\ndef run_kwyk_prediction(\n kwyk_dir: str,\n infile: str,\n outdir: Path,\n) -> Path | None:\n \"\"\"Run original kwyk CLI to produce a prediction.\n\n Calls::\n\n python {kwyk_dir}/kwyk/cli.py predict \\\\\n -m bwn_multi -n 1 {infile} {outprefix}\n\n Returns the path to the prediction NIfTI, or None on failure.\n \"\"\"\n vol_stem = Path(infile).stem.replace(\".nii\", \"\")\n outprefix = str(outdir / f\"kwyk_{vol_stem}\")\n cmd = [\n \"python\",\n str(Path(kwyk_dir) / \"kwyk\" / \"cli.py\"),\n \"predict\",\n \"-m\",\n \"bwn_multi\",\n \"-n\",\n \"1\",\n infile,\n outprefix,\n ]\n log.info(\"Running kwyk: %s\", \" \".join(cmd))\n\n try:\n result = subprocess.run(\n cmd,\n capture_output=True,\n text=True,\n timeout=600,\n check=False,\n )\n if result.returncode != 0:\n log.error(\"kwyk failed (rc=%d): %s\", result.returncode, result.stderr)\n return None\n except subprocess.TimeoutExpired:\n log.error(\"kwyk timed out for %s\", infile)\n return None\n except FileNotFoundError:\n log.error(\"kwyk CLI not found at %s\", cmd[1])\n return None\n\n # kwyk outputs {outprefix}_means.nii.gz or {outprefix}.nii.gz\n for suffix in [\"_means.nii.gz\", \".nii.gz\", \"_prediction.nii.gz\"]:\n candidate = Path(outprefix + suffix)\n if candidate.exists():\n return candidate\n\n log.warning(\"Could not find kwyk output for prefix %s\", outprefix)\n return None\n\n\n# ---------------------------------------------------------------------------\n# Spatial correlation between uncertainty maps\n# ---------------------------------------------------------------------------\ndef compute_spatial_correlation(map1: np.ndarray, map2: np.ndarray) -> float:\n \"\"\"Compute Pearson correlation between two spatial maps.\n\n Parameters\n ----------\n map1, map2 : np.ndarray\n Flattened or volumetric arrays of the same shape.\n\n Returns\n -------\n float\n Pearson correlation coefficient, or 0.0 on failure.\n \"\"\"\n v1 = map1.flatten().astype(np.float64)\n v2 = map2.flatten().astype(np.float64)\n\n # Remove positions where both are zero\n mask = (v1 != 0) | (v2 != 0)\n if mask.sum() < 2:\n return 0.0\n\n v1 = v1[mask]\n v2 = v2[mask]\n\n std1 = np.std(v1)\n std2 = np.std(v2)\n if std1 == 0 or std2 == 0:\n return 0.0\n\n return float(np.corrcoef(v1, v2)[0, 1])\n\n\n# ---------------------------------------------------------------------------\n# Scatter plot\n# ---------------------------------------------------------------------------\ndef plot_dice_scatter(\n kwyk_dices: list[float],\n new_dices: list[float],\n volume_names: list[str],\n output_path: Path,\n) -> None:\n \"\"\"Generate scatter plot: kwyk Dice (x) vs new model Dice (y).\"\"\"\n fig, ax = plt.subplots(figsize=(8, 8))\n\n ax.scatter(kwyk_dices, new_dices, alpha=0.7, edgecolors=\"k\", s=50)\n\n # Identity line\n lims = [0.0, 1.0]\n ax.plot(lims, lims, \"k--\", alpha=0.3, label=\"y = x\")\n\n ax.set_xlabel(\"Original kwyk Dice\")\n ax.set_ylabel(\"New Model Dice\")\n ax.set_title(\"Dice Comparison: Original kwyk vs New Model\")\n ax.set_xlim(lims)\n ax.set_ylim(lims)\n ax.set_aspect(\"equal\")\n ax.legend()\n ax.grid(True, alpha=0.3)\n\n fig.tight_layout()\n save_figure(fig, output_path)\n plt.close(fig)\n log.info(\"Scatter plot saved to %s\", output_path)\n\n\n# ---------------------------------------------------------------------------\n# Main\n# ---------------------------------------------------------------------------\ndef main() -> None:\n \"\"\"Compare new model vs original kwyk on test volumes.\"\"\"\n args = parse_args()\n output_dir = Path(args.output_dir)\n output_dir.mkdir(parents=True, exist_ok=True)\n kwyk_pred_dir = output_dir / \"kwyk_predictions\"\n kwyk_pred_dir.mkdir(parents=True, exist_ok=True)\n\n # ---- Load new model -----------------------------------------------------\n from nobrainer.processing.segmentation import Segmentation\n\n log.info(\"Loading new model from %s\", args.new_model)\n seg = Segmentation.load(args.new_model)\n block_shape = seg.block_shape_ or (32, 32, 32)\n\n # ---- Load manifest ------------------------------------------------------\n pairs = load_manifest(args.manifest, split=args.split)\n log.info(\n \"Comparing on %d volumes from split '%s'\",\n len(pairs),\n args.split,\n )\n\n if not pairs:\n log.error(\"No volumes found for split '%s'. Exiting.\", args.split)\n return\n\n # ---- Evaluate each volume -----------------------------------------------\n results: list[dict] = []\n kwyk_dices: list[float] = []\n new_dices: list[float] = []\n volume_names: list[str] = []\n\n for idx, (img_path, lbl_path) in enumerate(pairs):\n vol_name = Path(img_path).stem\n log.info(\"Volume %d/%d: %s\", idx + 1, len(pairs), vol_name)\n\n # Load ground truth and binarize\n gt_arr = np.asarray(nib.load(lbl_path).dataobj, dtype=np.float32)\n gt_binary = (gt_arr > 0).astype(np.float32)\n\n # ---- New model prediction -------------------------------------------\n label_img = seg.predict(img_path, block_shape=block_shape)\n new_pred = np.asarray(label_img.dataobj, dtype=np.float32)\n new_binary = (new_pred > 0).astype(np.float32)\n new_dice = compute_dice(new_binary, gt_binary)\n log.info(\" New model Dice = %.4f\", new_dice)\n\n # ---- Original kwyk prediction ---------------------------------------\n kwyk_pred_path = run_kwyk_prediction(args.kwyk_dir, img_path, kwyk_pred_dir)\n kwyk_dice = float(\"nan\")\n if kwyk_pred_path is not None and kwyk_pred_path.exists():\n kwyk_arr = np.asarray(\n nib.load(str(kwyk_pred_path)).dataobj, dtype=np.float32\n )\n kwyk_binary = (kwyk_arr > 0).astype(np.float32)\n kwyk_dice = compute_dice(kwyk_binary, gt_binary)\n log.info(\" kwyk Dice = %.4f\", kwyk_dice)\n else:\n log.warning(\" kwyk prediction not available for %s\", vol_name)\n\n results.append(\n {\n \"volume\": vol_name,\n \"new_dice\": new_dice,\n \"kwyk_dice\": kwyk_dice,\n \"image_path\": img_path,\n }\n )\n new_dices.append(new_dice)\n kwyk_dices.append(kwyk_dice)\n volume_names.append(vol_name)\n\n # ---- Save comparison CSV ------------------------------------------------\n csv_path = output_dir / \"comparison_table.csv\"\n with open(csv_path, \"w\", newline=\"\") as f:\n writer = csv.DictWriter(\n f,\n fieldnames=[\"volume\", \"new_dice\", \"kwyk_dice\", \"image_path\"],\n )\n writer.writeheader()\n writer.writerows(results)\n log.info(\"Comparison table saved to %s\", csv_path)\n\n # ---- Scatter plot -------------------------------------------------------\n # Filter out NaN kwyk dices for plotting\n valid_mask = [not np.isnan(kd) for kd in kwyk_dices]\n valid_kwyk = [kd for kd, v in zip(kwyk_dices, valid_mask) if v]\n valid_new = [nd for nd, v in zip(new_dices, valid_mask) if v]\n valid_names = [n for n, v in zip(volume_names, valid_mask) if v]\n\n if valid_kwyk:\n scatter_path = output_dir / \"dice_scatter.png\"\n plot_dice_scatter(valid_kwyk, valid_new, valid_names, scatter_path)\n else:\n log.warning(\"No valid kwyk predictions; skipping scatter plot\")\n\n # ---- Spatial correlation of uncertainty maps ----------------------------\n # Check if both models have uncertainty outputs\n new_results_dir = Path(args.new_model).parent / \"results\"\n if new_results_dir.exists():\n log.info(\"Checking for uncertainty map correlations...\")\n for vol_name in volume_names:\n new_var_path = new_results_dir / f\"{vol_name}_variance.nii.gz\"\n kwyk_var_candidates = list(kwyk_pred_dir.glob(f\"kwyk_{vol_name}*variance*\"))\n if new_var_path.exists() and kwyk_var_candidates:\n new_var = np.asarray(\n nib.load(str(new_var_path)).dataobj, dtype=np.float32\n )\n kwyk_var = np.asarray(\n nib.load(str(kwyk_var_candidates[0])).dataobj,\n dtype=np.float32,\n )\n corr = compute_spatial_correlation(new_var, kwyk_var)\n log.info(\n \" %s uncertainty correlation: %.4f\",\n vol_name,\n corr,\n )\n\n # ---- Summary ------------------------------------------------------------\n log.info(\"=\" * 60)\n log.info(\"Comparison Summary (%s split)\", args.split)\n log.info(\" Volumes compared : %d\", len(results))\n if valid_kwyk:\n log.info(\" Mean kwyk Dice : %.4f\", np.nanmean(kwyk_dices))\n log.info(\" Mean new Dice : %.4f\", np.mean(new_dices))\n if valid_kwyk:\n improvement = np.mean(valid_new) - np.mean(valid_kwyk)\n log.info(\" Mean improvement : %+.4f\", improvement)\n log.info(\" Output directory : %s\", output_dir)\n log.info(\"=\" * 60)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "scripts/kwyk_reproduction/06_block_size_sweep.py", "content": "#!/usr/bin/env python\n\"\"\"Sweep over block sizes to compare segmentation performance.\n\nUsage:\n python 06_block_size_sweep.py --manifest manifest.csv --config config.yaml\n python 06_block_size_sweep.py --manifest manifest.csv --config config.yaml \\\n --block-sizes 32 64 128 --epochs 20 --output-dir results/sweep\n\"\"\"\n\nfrom __future__ import annotations\n\nimport argparse\nimport csv\nfrom pathlib import Path\nimport time\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nfrom utils import compute_dice, load_config, save_figure, setup_logging\n\nlog = setup_logging(__name__)\n\n\n# ---------------------------------------------------------------------------\n# CLI\n# ---------------------------------------------------------------------------\ndef parse_args() -> argparse.Namespace:\n \"\"\"Parse command-line arguments.\"\"\"\n parser = argparse.ArgumentParser(\n description=\"Block size sweep for Bayesian MeshNet segmentation\",\n )\n parser.add_argument(\n \"--manifest\",\n type=str,\n required=True,\n help=\"Path to the dataset manifest CSV\",\n )\n parser.add_argument(\n \"--config\",\n type=str,\n default=\"config.yaml\",\n help=\"Path to YAML configuration file\",\n )\n parser.add_argument(\n \"--block-sizes\",\n type=int,\n nargs=\"+\",\n default=[32, 64, 128],\n help=\"Block sizes to sweep over (default: 32 64 128)\",\n )\n parser.add_argument(\n \"--epochs\",\n type=int,\n default=20,\n help=\"Number of training epochs per block size (default: 20)\",\n )\n parser.add_argument(\n \"--output-dir\",\n type=str,\n default=\"results/sweep\",\n help=\"Directory for sweep outputs\",\n )\n return parser.parse_args()\n\n\ndef load_manifest(manifest_path: str, split: str) -> list[tuple[str, str]]:\n \"\"\"Load manifest CSV and return (image, label) pairs for the given split.\"\"\"\n pairs = []\n with open(manifest_path) as f:\n reader = csv.DictReader(f)\n for row in reader:\n if row[\"split\"] == split:\n pairs.append((row[\"t1w_path\"], row[\"label_path\"]))\n return pairs\n\n\n# ---------------------------------------------------------------------------\n# Train + evaluate for one block size\n# ---------------------------------------------------------------------------\ndef train_and_evaluate(\n block_size: int,\n config: dict,\n train_pairs: list[tuple[str, str]],\n val_pairs: list[tuple[str, str]],\n epochs: int,\n) -> dict:\n \"\"\"Train a Bayesian MeshNet at the given block size and evaluate Dice.\n\n Returns\n -------\n dict\n Keys: block_size, mean_dice, std_dice, per_volume_dices, final_loss.\n \"\"\"\n import nibabel as nib\n\n from nobrainer.models import get as get_model\n from nobrainer.models.bayesian.utils import accumulate_kl\n from nobrainer.prediction import predict\n from nobrainer.processing.dataset import Dataset\n\n block_shape = (block_size, block_size, block_size)\n n_classes = config[\"n_classes\"]\n batch_size = config[\"batch_size\"]\n lr = config.get(\"lr\", 1e-4)\n kl_weight = config.get(\"kl_weight\", 1.0)\n\n log.info(\n \"Training with block_size=%d for %d epochs...\",\n block_size,\n epochs,\n )\n\n # Build dataset\n ds_train = (\n Dataset.from_files(train_pairs, block_shape=block_shape, n_classes=n_classes)\n .batch(batch_size)\n .binarize(config.get(\"label_mapping\", \"binary\"))\n )\n train_loader = ds_train.dataloader\n\n # Build model\n model_args = {\n \"n_classes\": n_classes,\n \"filters\": config.get(\"filters\", 96),\n \"receptive_field\": config.get(\"receptive_field\", 37),\n \"dropout_rate\": config.get(\"dropout_rate\", 0.25),\n }\n model = get_model(\"bayesian_meshnet\")(**model_args)\n\n from nobrainer.gpu import get_device\n\n device = get_device()\n model = model.to(device)\n\n ce_loss = torch.nn.CrossEntropyLoss()\n optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n\n # Training loop\n final_loss = 0.0\n for epoch in range(epochs):\n model.train()\n epoch_loss = 0.0\n n_batches = 0\n\n for batch in train_loader:\n if isinstance(batch, dict):\n images = batch[\"image\"].to(device)\n labels = batch[\"label\"].to(device)\n elif isinstance(batch, (list, tuple)):\n images = batch[0].to(device)\n labels = batch[1].to(device)\n else:\n raise TypeError(f\"Unsupported batch type: {type(batch)}\")\n\n if labels.ndim == images.ndim and labels.shape[1] == 1:\n labels = labels.squeeze(1)\n if labels.dtype in (torch.float32, torch.float64):\n labels = labels.long()\n\n optimizer.zero_grad()\n pred = model(images)\n loss = ce_loss(pred, labels) + kl_weight * accumulate_kl(model)\n loss.backward()\n optimizer.step()\n\n epoch_loss += loss.item()\n n_batches += 1\n\n final_loss = epoch_loss / max(n_batches, 1)\n if (epoch + 1) % 5 == 0 or epoch == epochs - 1:\n log.info(\n \" block_size=%d, epoch %d/%d, loss=%.6f\",\n block_size,\n epoch + 1,\n epochs,\n final_loss,\n )\n\n # Evaluate on validation set\n model.eval()\n per_volume_dices: list[float] = []\n\n for img_path, lbl_path in val_pairs:\n gt_arr = np.asarray(nib.load(lbl_path).dataobj, dtype=np.float32)\n gt_binary = (gt_arr > 0).astype(np.float32)\n\n pred_img = predict(\n inputs=img_path,\n model=model,\n block_shape=block_shape,\n batch_size=4,\n return_labels=True,\n )\n pred_arr = np.asarray(pred_img.dataobj, dtype=np.float32)\n pred_binary = (pred_arr > 0).astype(np.float32)\n\n dice = compute_dice(pred_binary, gt_binary)\n per_volume_dices.append(dice)\n\n mean_dice = float(np.mean(per_volume_dices)) if per_volume_dices else 0.0\n std_dice = float(np.std(per_volume_dices)) if per_volume_dices else 0.0\n\n log.info(\n \" block_size=%d: Dice=%.4f +/- %.4f (%d volumes)\",\n block_size,\n mean_dice,\n std_dice,\n len(per_volume_dices),\n )\n\n return {\n \"block_size\": block_size,\n \"mean_dice\": mean_dice,\n \"std_dice\": std_dice,\n \"per_volume_dices\": per_volume_dices,\n \"final_loss\": final_loss,\n }\n\n\n# ---------------------------------------------------------------------------\n# Bar chart\n# ---------------------------------------------------------------------------\ndef plot_block_size_comparison(\n sweep_results: list[dict],\n output_path: Path,\n) -> None:\n \"\"\"Generate bar chart: block_size on x, Dice on y with error bars.\"\"\"\n block_sizes = [r[\"block_size\"] for r in sweep_results]\n means = [r[\"mean_dice\"] for r in sweep_results]\n stds = [r[\"std_dice\"] for r in sweep_results]\n\n fig, ax = plt.subplots(figsize=(8, 6))\n\n x = np.arange(len(block_sizes))\n bars = ax.bar(\n x,\n means,\n yerr=stds,\n capsize=5,\n color=\"steelblue\",\n edgecolor=\"black\",\n alpha=0.8,\n )\n\n ax.set_xlabel(\"Block Size\")\n ax.set_ylabel(\"Dice Score\")\n ax.set_title(\"Block Size Sweep — Bayesian MeshNet\")\n ax.set_xticks(x)\n ax.set_xticklabels([str(bs) for bs in block_sizes])\n ax.set_ylim(0.0, 1.0)\n ax.grid(axis=\"y\", alpha=0.3)\n\n # Annotate bars with mean values\n for bar, mean in zip(bars, means):\n ax.text(\n bar.get_x() + bar.get_width() / 2.0,\n bar.get_height() + 0.02,\n f\"{mean:.3f}\",\n ha=\"center\",\n va=\"bottom\",\n fontsize=10,\n )\n\n fig.tight_layout()\n save_figure(fig, output_path)\n plt.close(fig)\n log.info(\"Bar chart saved to %s\", output_path)\n\n\n# ---------------------------------------------------------------------------\n# Main\n# ---------------------------------------------------------------------------\ndef main() -> None:\n \"\"\"Run block size sweep and generate comparison outputs.\"\"\"\n args = parse_args()\n t_start = time.time()\n output_dir = Path(args.output_dir)\n output_dir.mkdir(parents=True, exist_ok=True)\n\n # ---- Load config --------------------------------------------------------\n config = load_config(args.config)\n log.info(\"Config loaded from %s\", args.config)\n log.info(\n \"Block size sweep: sizes=%s, epochs=%d\",\n args.block_sizes,\n args.epochs,\n )\n\n # ---- Load manifest ------------------------------------------------------\n train_pairs = load_manifest(args.manifest, split=\"train\")\n val_pairs = load_manifest(args.manifest, split=\"val\")\n log.info(\n \"Manifest: %d train, %d val volumes\",\n len(train_pairs),\n len(val_pairs),\n )\n\n if not train_pairs:\n log.error(\"No training volumes found. Exiting.\")\n return\n if not val_pairs:\n log.warning(\"No validation volumes found; Dice will be empty.\")\n\n # ---- Run sweep ----------------------------------------------------------\n sweep_results: list[dict] = []\n\n for block_size in args.block_sizes:\n result = train_and_evaluate(\n block_size=block_size,\n config=config,\n train_pairs=train_pairs,\n val_pairs=val_pairs,\n epochs=args.epochs,\n )\n sweep_results.append(result)\n\n # ---- Save comparison CSV ------------------------------------------------\n csv_path = output_dir / \"block_size_comparison.csv\"\n with open(csv_path, \"w\", newline=\"\") as f:\n writer = csv.DictWriter(\n f,\n fieldnames=[\"block_size\", \"mean_dice\", \"std_dice\", \"final_loss\"],\n )\n writer.writeheader()\n for r in sweep_results:\n writer.writerow(\n {\n \"block_size\": r[\"block_size\"],\n \"mean_dice\": r[\"mean_dice\"],\n \"std_dice\": r[\"std_dice\"],\n \"final_loss\": r[\"final_loss\"],\n }\n )\n log.info(\"Comparison CSV saved to %s\", csv_path)\n\n # ---- Bar chart ----------------------------------------------------------\n chart_path = output_dir / \"block_size_comparison.png\"\n plot_block_size_comparison(sweep_results, chart_path)\n\n # ---- Summary ------------------------------------------------------------\n elapsed = time.time() - t_start\n best = max(sweep_results, key=lambda r: r[\"mean_dice\"])\n log.info(\"=\" * 60)\n log.info(\"Block Size Sweep Complete\")\n log.info(\" Block sizes tested: %s\", args.block_sizes)\n log.info(\" Epochs per size : %d\", args.epochs)\n log.info(\n \" Best block size : %d (Dice=%.4f)\", best[\"block_size\"], best[\"mean_dice\"]\n )\n for r in sweep_results:\n log.info(\n \" block_size=%3d: Dice=%.4f +/- %.4f, loss=%.6f\",\n r[\"block_size\"],\n r[\"mean_dice\"],\n r[\"std_dice\"],\n r[\"final_loss\"],\n )\n log.info(\" Output directory : %s\", output_dir)\n log.info(\" Elapsed time : %.1f s\", elapsed)\n log.info(\"=\" * 60)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "scripts/kwyk_reproduction/ARCHITECTURE.md", "content": "# KWYK Architecture Verification\n\nThis document records how the original kwyk model architecture was verified\nagainst the paper, source code, and trained model weights.\n\n## Paper Reference\n\nMcClure P. et al., \"Knowing What You Know in Brain Segmentation Using\nBayesian Deep Neural Networks\", Front. Neuroinform. 2019.\nhttps://doi.org/10.3389/fninf.2019.00067\n\n## Three Model Variants\n\n| Variant | kwyk ID | Conv Layer | Dropout | MC at inference |\n|---------|---------|-----------|---------|-----------------|\n| MAP | `bwn` (all_50_wn) | VWN | Bernoulli (fixed) | No |\n| MC Bernoulli Dropout (BD) | `bwn_multi` (all_50_bwn_09_multi) | VWN | Bernoulli (fixed) | Yes |\n| Spike-and-Slab Dropout (SSD) | `bvwn_multi_prior` (all_50_bvwn_multi_prior) | VWN | Concrete (learned) | Yes |\n\n## Architecture: Variational Weight Normalization (VWN) Conv\n\n### Verified from trained model variables\n\nDownloaded `neuronets/kwyk:latest-cpu` Docker container and inspected the\nSavedModel variables for the SSD model (`all_50_bvwn_multi_prior`):\n\n```\nlayer_1/conv3d/v:0: [3, 3, 3, 1, 96] # raw weight for WN\nlayer_1/conv3d/g:0: [1, 1, 1, 1, 96] # gain per filter\nlayer_1/conv3d/kernel_a:0: [3, 3, 3, 1, 96] # sigma = |kernel_a|\nlayer_1/conv3d/bias_m:0: [96] # bias mean\nlayer_1/conv3d/bias_a:0: [96] # bias sigma = |bias_a|\nlayer_1/concrete_dropout/p:0: [96] # per-filter dropout rate\n```\n\nThis confirms **weight normalization** (`v`, `g`) is used, not the direct\n`μ` parameterization described in the paper's equations.\n\n### Key finding: all 3 models are independently trained VWN models\n\nAll 3 saved models have the **same layer structure** (`v`, `g`, `kernel_a`,\n`bias_m`, `bias_a`) — including the MAP model (`all_50_wn`). They are\n**not** weight-sharing variants; they were trained independently:\n\n| Model | Total variables | Extra per layer | Timestamp |\n|-------|----------------|----------------|-----------|\n| all_50_wn (MAP) | 41 | — | 1555341859 |\n| all_50_bwn_09_multi (BD) | 41 | — | 1555963478 |\n| all_50_bvwn_multi_prior (SSD) | 48 | `concrete_dropout/p` | 1556816070 |\n\nThe MAP and BD models have identical parameterization (both have `kernel_a`\nfor learned sigma). The only difference is whether MC sampling is enabled\nat inference time. The SSD model additionally has 7 `concrete_dropout/p`\nparameters (one per conv layer) for learned per-filter dropout rates.\n\n### Verified from source code\n\nCommit `4dd379c` in `neuronets/kwyk` repo (Patrick McClure, 2019-02-28):\n\n**`nobrainer/models/vwn_conv.py`** — `_Conv.build()`:\n```python\nself.v = self.add_variable(name='v', ...)\nself.g = self.add_variable(name='g', ...)\nself.v_norm = tf.nn.l2_normalize(self.v, [...])\nself.kernel_m = tf.multiply(self.g, self.v_norm, name='kernel_m')\nself.kernel_a = self.add_variable(name='kernel_a', ...)\nself.kernel_sigma = tf.abs(self.kernel_a, name='kernel_sigma')\n```\n\n**`_Conv.call()`** — local reparameterization trick:\n```python\noutputs_mean = self._convolution_op(inputs, self.kernel_m)\noutputs_var = self._convolution_op(tf.square(inputs), tf.square(self.kernel_sigma))\noutputs_e = tf.random_normal(shape=tf.shape(self.g))\n# MC path:\noutput = outputs_mean + tf.sqrt(outputs_var + 1e-8) * outputs_e\n```\n\n**`nobrainer/models/bayesian_dropout.py`** defines:\n- `bernoulli_dropout()` — standard MC dropout (bwn/bwn_multi)\n- `concrete_dropout()` — learned per-filter rate (bvwn_multi_prior)\n- `gaussian_dropout()` — not used in final models\n\n### Paper vs Implementation discrepancy\n\nThe paper (Section 2.2.3.2) describes the mean weight as `μ_{f,t}` (Eq. 13),\nbut the actual implementation uses weight normalization:\n- `kernel_m = g · v / ||v||` (Salimans & Kingma 2016)\n- This is a reparameterization of the mean that aids training stability\n- The sigma is the same in both: `σ_{f,t} = |kernel_a_{f,t}|`\n\nThe paper's equations are in terms of the effective mean (`μ`), which is\ncomputed via WN but isn't stored directly as a parameter.\n\n## KL Divergence (Eq. 16-18)\n\nTwo terms per filter:\n\n1. **Bernoulli KL** for concrete dropout (Eq. 17):\n `KL(q_p || p_prior) = p·log(p/p_prior) + (1-p)·log((1-p)/(1-p_prior))`\n Prior: `p_prior = 0.5`\n\n2. **Gaussian KL** per weight (Eq. 18):\n `KL(N(μ,σ) || N(μ_prior, σ_prior)) = log(σ_prior/σ) + (σ² + (μ-μ_prior)²)/(2σ²_prior) - 1/2`\n Prior: `μ_prior = 0, σ_prior = 0.1`\n\n## Network Architecture (Table 2)\n\n8 layers of dilated 3×3×3 convolutions:\n- Layers 1-3: dilation=1, 96 filters, ReLU\n- Layer 4: dilation=2\n- Layer 5: dilation=4\n- Layer 6: dilation=8\n- Layer 7: dilation=1\n- Layer 8 (logits): 1×1×1, 50 filters, Softmax\n\nReceptive field = 37 voxels.\n\n## Our Implementation\n\n`nobrainer.models.bayesian.vwn_layers.FFGConv3d`:\n- Parameters: `v`, `g` (weight normalization), `kernel_a` (sigma), `bias_m`, `bias_a`\n- Forward: local reparameterization trick matching the original\n- KL: Eq. 18 with `prior_mu=0, prior_sigma=0.1`\n\n`nobrainer.models.bayesian.vwn_layers.ConcreteDropout3d`:\n- Learned `p` per filter via concrete relaxation (Eq. 10)\n- KL: Eq. 17 with `prior_p=0.5`\n\n`nobrainer.models.bayesian.kwyk_meshnet.KWYKMeshNet`:\n- Registered as `\"kwyk_meshnet\"` in model registry (no Pyro dependency)\n- `dropout_type=\"bernoulli\"` for bwn/bwn_multi\n- `dropout_type=\"concrete\"` for bvwn_multi_prior (SSD)\n- `mc=True/False` flag controls stochastic vs deterministic inference\n\n## Training Details (from paper)\n\n- Optimizer: Adam, lr=1e-4\n- Batch size: 32 (4 GPUs × 8)\n- Block shape: 32×32×32\n- Data: 11,480 T1 sMRI volumes, 50-class FreeSurfer parcellation\n- MC samples at inference: 10\n" }, { "path": "scripts/kwyk_reproduction/README.md", "content": "# KWYK Brain Extraction Reproduction\n\nReproduce the kwyk brain extraction study (McClure et al., Frontiers in\nNeuroinformatics 2019) using the refactored PyTorch nobrainer.\n\n**Reference**: https://www.frontiersin.org/journals/neuroinformatics/articles/10.3389/fninf.2019.00067/full\n\n## Current Status\n\nThe reproduction pipeline is **code-complete and CI-verified** (smoke test\n+ small-scale 20-epoch training on T4 GPU with real OpenNeuro data).\nFull-scale reproduction with 50+ epochs and 100+ subjects has **not yet\nbeen run**. See \"Next Steps\" below.\n\n## Quick Setup\n\n```bash\n# Option A: Use the orchestrator script (creates venv automatically)\ncd scripts/kwyk_reproduction\n./run.sh --smoke-test # Quick verification (5 volumes, 2 epochs)\n./run.sh # Full pipeline\n\n# Option B: Manual setup\nuv venv --python 3.14 && source .venv/bin/activate\nuv pip install -e \"../../[bayesian,versioning,dev]\" monai pyro-ppl datalad matplotlib pyyaml scipy\nuv tool install git-annex # required for DataLad content retrieval\n```\n\n## Programmatic API\n\nThe dataset fetching is also available as a library:\n\n```python\nfrom nobrainer.datasets.openneuro import (\n install_derivatives,\n find_subject_pairs,\n write_manifest,\n)\n\n# Clone fmriprep derivatives (metadata only, fast)\nds = install_derivatives(\"ds000114\", \"/tmp/data\")\n\n# Discover + download T1w + aparc+aseg pairs per subject\npairs = find_subject_pairs(ds)\n\n# Write manifest CSV with train/val/test split\nwrite_manifest(pairs, \"manifest.csv\")\n```\n\n## Pipeline Steps\n\n### Step 1: Assemble Dataset\n\n```bash\npython 01_assemble_dataset.py --datasets ds000114 --output-csv manifest.csv\n```\n\nDownloads T1w + aparc+aseg volumes from OpenNeuro fmriprep derivatives via\nDataLad. Start with 1 dataset (~10 subjects) for smoke testing, then scale:\n\n```bash\n# Scale to more datasets\npython 01_assemble_dataset.py \\\n --datasets ds000114 ds000228 ds002609 ds001021 ds002105 \\\n --output-csv manifest.csv --conform\n```\n\n### Step 2: Train Deterministic MeshNet (Warm-Start Foundation)\n\n```bash\npython 02_train_meshnet.py --manifest manifest.csv --epochs 50\n```\n\nTrains a standard MeshNet with kwyk-matching parameters (filters=96,\nblock_shape=32³, lr=0.0001). This model's weights serve as the mean\npriors for the Bayesian model in Step 3.\n\n**Output**: `checkpoints/meshnet/model.pth`, `figures/meshnet_learning_curve.png`\n\n### Step 3: Train All Model Variants\n\nAll 3 kwyk models use **Variational Weight Normalization (VWN)** convolutions\nwith per-weight learned sigma and the local reparameterization trick. They\ndiffer only in the dropout layer. See [ARCHITECTURE.md](ARCHITECTURE.md)\nfor the full verification against the paper, code, and trained weights.\n\nUse `--variant` to select:\n\n```bash\n# 3a. MC Bernoulli dropout (bwn_multi) — VWN conv + dropout at inference\npython 03_train_bayesian.py \\\n --manifest manifest.csv --variant bwn_multi \\\n --warmstart checkpoints/meshnet --output-dir checkpoints/bwn_multi \\\n --epochs 50\n\n# 3b. Spike-and-slab dropout (bvwn_multi_prior) — VWN conv + concrete dropout\npython 03_train_bayesian.py \\\n --manifest manifest.csv --variant bvwn_multi_prior \\\n --warmstart checkpoints/meshnet --output-dir checkpoints/bvwn_multi_prior \\\n --epochs 50\n```\n\n| Variant | kwyk ID | Conv | Dropout | MC at inference |\n|---------|---------|------|---------|-----------------|\n| `bwn` (step 2) | all_50_wn | VWN | Bernoulli (fixed) | No (MAP) |\n| `bwn_multi` | all_50_bwn_09_multi | VWN | Bernoulli (fixed) | Yes |\n| `bvwn_multi_prior` | all_50_bvwn_multi_prior | VWN | Concrete (learned) | Yes |\n\nThe warm-start decomposes deterministic Conv3d weights into weight\nnormalization form (`v`, `g`) for the VWN layers.\n\n**Output**: `checkpoints//model.pth`, `checkpoints//croissant.json`,\n`checkpoints//learning_curve.png`\n\n### Step 4: Evaluate\n\n```bash\n# Evaluate each variant\nfor variant in meshnet bwn_multi bvwn_multi_prior bayesian_gaussian; do\n python 04_evaluate.py \\\n --model checkpoints/$variant/model.pth \\\n --manifest manifest.csv --split test --n-samples 10 \\\n --output-dir results/$variant\ndone\n```\n\nComputes per-volume Dice, saves variance + entropy maps as NIfTI.\n\n### Step 5: Compare with Original KWYK\n\n```bash\npython 05_compare_kwyk.py \\\n --new-model checkpoints/bvwn_multi_prior/model.pth \\\n --kwyk-dir ../../kwyk \\\n --manifest manifest.csv\n```\n\nRuns the original kwyk container on the same test volumes and generates a\nDice scatter plot + comparison table. **Note**: This requires the kwyk\ncontainer at `../../kwyk` to be functional. The comparison is only meaningful\nafter the Bayesian model has been trained to convergence (Steps 2-3).\n\n### Step 6: Block Size Sweep (Optional)\n\n```bash\npython 06_block_size_sweep.py --manifest manifest.csv --block-sizes 32 64 128\n```\n\n## Next Steps for GPU Execution\n\nThe following steps should be performed on a machine with a GPU (e.g., the\nEC2 GPU runner or a local workstation):\n\n### Phase 1: Smoke Test (15 minutes, any GPU)\n\n```bash\n./run.sh --smoke-test\n```\n\nVerify the pipeline works end-to-end with tiny models. Check\n`figures/` for learning curves showing loss decrease.\n\n### Phase 2: Small-Scale Training (1-2 hours, T4 16GB)\n\n```bash\npython 01_assemble_dataset.py --datasets ds000114 --output-csv manifest.csv\npython 02_train_meshnet.py --manifest manifest.csv --epochs 20\n# Train all 3 Bayesian variants\nfor variant in bwn_multi bvwn_multi_prior bayesian_gaussian; do\n python 03_train_bayesian.py --manifest manifest.csv \\\n --variant $variant --warmstart checkpoints/meshnet \\\n --output-dir checkpoints/$variant --epochs 20\ndone\n```\n\n**Expected**: Validation Dice ≥0.80 for brain extraction on 10 subjects.\n\n### Phase 3: Full Reproduction (8-24 hours, V100 16GB+)\n\n```bash\npython 01_assemble_dataset.py \\\n --datasets ds000114 ds000228 ds002609 ds001021 ds002105 \\\n --output-csv manifest.csv --conform\npython 02_train_meshnet.py --manifest manifest.csv --epochs 50\nfor variant in bwn_multi bvwn_multi_prior bayesian_gaussian; do\n python 03_train_bayesian.py --manifest manifest.csv \\\n --variant $variant --warmstart checkpoints/meshnet \\\n --output-dir checkpoints/$variant --epochs 50\ndone\npython 05_compare_kwyk.py \\\n --new-model checkpoints/bvwn_multi_prior/model.pth \\\n --kwyk-dir ../../kwyk --manifest manifest.csv\n```\n\n**Target**: Validation Dice ≥0.90 (kwyk achieved 0.97+ with 11,000 subjects).\n\n### Phase 4: Scale and Optimize\n\nTo approach kwyk's full performance:\n\n1. **Add more datasets**: Add OpenNeuro dataset IDs to the `--datasets` list\n2. **Block size sweep**: `python 06_block_size_sweep.py --block-sizes 32 64 128`\n3. **SynthSeg augmentation**: `python 03_train_bayesian.py --augmentation mixed`\n4. **Longer training**: Increase `--epochs` to 100+\n\n### Phase 5: Automated Hyperparameter Optimization\n\nUse nobrainer's autoresearch loop to explore hyperparameters overnight:\n\n```bash\n# Set up the research directory\nmkdir -p research/kwyk_bayesian\ncp checkpoints/bayesian/model.pth research/kwyk_bayesian/\ncat > research/kwyk_bayesian/program.md << 'EOF'\n## Exploration Targets\n- kl_weight: 1e-5, 1e-4, 1e-3, 1e-2, 1e-1\n- dropout_rate: 0.0, 0.1, 0.25, 0.5\n- filters: 71, 96, 128\n- prior_type: standard_normal, laplace\n- block_shape: 32, 64\n- learning_rate: 1e-5, 5e-5, 1e-4, 5e-4\n\n## Success Criterion\n- val_dice improvement over current best\n- Max 30 min per experiment\nEOF\n\n# Launch overnight optimization\nnobrainer research run \\\n --working-dir research/kwyk_bayesian \\\n --model-family bayesian_meshnet \\\n --max-experiments 20 \\\n --budget-hours 8\n```\n\nThe autoresearch loop will:\n1. Propose hyperparameter changes (via LLM or random grid)\n2. Train, evaluate, keep improvements, revert failures\n3. Save the best model with full Croissant-ML provenance\n\nCheck results: `cat research/kwyk_bayesian/run_summary.md`\n\n## Configuration\n\nEdit `config.yaml` to change default hyperparameters:\n\n| Parameter | Default | kwyk Original | Notes |\n|-----------|---------|---------------|-------|\n| filters | 96 | 96 | Feature maps per layer |\n| receptive_field | 37 | 37 | Dilation schedule [1,1,1,2,4,8,1] |\n| block_shape | [32,32,32] | [32,32,32] | Patch size for training |\n| lr | 0.0001 | 0.0001 | Adam learning rate |\n| kl_weight | 1.0 | implicit | KL divergence scaling |\n| dropout_rate | 0.25 | 0.25 | Spatial dropout |\n| prior_type | spike_and_slab | spike_and_slab | SSD: π·N(0,0.001) + (1-π)·N(0,1) |\n| spike_sigma | 0.001 | ~0 | Spike component σ |\n| slab_sigma | 1.0 | ~1 | Slab component σ |\n| prior_pi | 0.5 | 0.5 | Spike probability |\n| n_classes | 2 | 50 | Binary brain extraction (kwyk used 50-class) |\n| label_mapping | binary | N/A | Also supports 6/50/115-class |\n\n## Label Mappings\n\nThe `label_mappings/` directory contains CSVs that remap FreeSurfer\naparc+aseg codes to target classes:\n\n- **binary**: Any non-zero → 1 (brain extraction)\n- **6-class**: Coarse parcellation (WM, cortex, ventricles, cerebellum, etc.)\n- **50-class**: Matches original kwyk study\n- **115-class**: Fine-grained parcellation\n\n## GPU Requirements\n\n| Task | Block Size | Filters | GPU Memory | Time |\n|------|-----------|---------|------------|------|\n| Smoke test | 16³ | 16 | ≥4 GB | ~5 min |\n| Small training | 32³ | 96 | ≥16 GB | ~2 hr |\n| Full reproduction | 32³ | 96 | ≥16 GB | ~24 hr |\n| Block sweep 64³ | 64³ | 96 | ≥16 GB | ~4 hr |\n| Full-brain 256³ | 256³ | 96 | ≥24 GB | N/A |\n\n## What Has Been Verified\n\n- [x] Smoke test on EC2 T4 GPU (2 epochs, all 3 variants)\n- [x] Small-scale training on EC2 T4 GPU (20 epochs, 10 subjects from ds000114, all 3 variants)\n- [x] DataLad + git-annex data pipeline (OpenNeuro fmriprep derivatives)\n- [x] Spike-and-slab prior, MC dropout, and Gaussian Bayesian variants\n\n## What Has NOT Been Done Yet\n\n- [ ] Full-scale reproduction with 50+ epochs and 100+ subjects\n- [ ] Comparison against kwyk container (requires converged model)\n- [ ] Block size sweep results\n- [ ] SynthSeg augmentation experiments\n- [ ] Autoresearch hyperparameter optimization\n" }, { "path": "scripts/kwyk_reproduction/__init__.py", "content": "" }, { "path": "scripts/kwyk_reproduction/build_kwyk_manifest.py", "content": "#!/usr/bin/env python\n\"\"\"Build a manifest CSV from the original KWYK dataset (PAC brain volumes).\n\nThe KWYK dataset contains paired files:\n - pac__orig.nii.gz (T1w image)\n - pac__aseg.nii.gz (FreeSurfer aparc+aseg label)\n\nUsage:\n python build_kwyk_manifest.py --data-dir ../data/SharedData/segmentation/freesurfer_asegs \\\n --output-csv kwyk_manifest.csv --n-subjects 100\n\"\"\"\n\nfrom __future__ import annotations\n\nimport argparse\nimport csv\nfrom pathlib import Path\nimport random\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Build manifest from KWYK PAC dataset\")\n parser.add_argument(\n \"--data-dir\",\n type=str,\n required=True,\n help=\"Directory containing pac_*_orig.nii.gz and pac_*_aseg.nii.gz files\",\n )\n parser.add_argument(\n \"--output-csv\",\n type=str,\n default=\"kwyk_manifest.csv\",\n help=\"Output manifest CSV path\",\n )\n parser.add_argument(\n \"--n-subjects\",\n type=int,\n default=None,\n help=\"Number of subjects to include (default: all)\",\n )\n parser.add_argument(\n \"--split\",\n nargs=3,\n type=int,\n default=[80, 10, 10],\n help=\"Train/val/test split percentages (default: 80 10 10)\",\n )\n parser.add_argument(\n \"--seed\",\n type=int,\n default=42,\n help=\"Random seed for shuffling and split (default: 42)\",\n )\n args = parser.parse_args()\n\n data_dir = Path(args.data_dir).resolve()\n if not data_dir.is_dir():\n raise SystemExit(f\"Data directory not found: {data_dir}\")\n\n # Find all paired subjects\n orig_files = sorted(data_dir.glob(\"pac_*_orig.nii.gz\"))\n pairs = []\n for orig in orig_files:\n # Extract subject ID: pac__orig.nii.gz -> \n stem = orig.name # pac_123_orig.nii.gz\n subj_id = stem.replace(\"pac_\", \"\").replace(\"_orig.nii.gz\", \"\")\n aseg = data_dir / f\"pac_{subj_id}_aseg.nii.gz\"\n if aseg.exists():\n pairs.append((subj_id, str(orig), str(aseg)))\n\n print(f\"Found {len(pairs)} paired subjects in {data_dir}\")\n\n if not pairs:\n raise SystemExit(\"No paired (orig, aseg) files found.\")\n\n # Shuffle and subsample\n random.seed(args.seed)\n random.shuffle(pairs)\n if args.n_subjects is not None:\n pairs = pairs[: args.n_subjects]\n print(f\"Subsampled to {len(pairs)} subjects\")\n\n # Split\n n = len(pairs)\n train_pct, val_pct, test_pct = args.split\n assert train_pct + val_pct + test_pct == 100\n n_train = int(n * train_pct / 100)\n n_val = int(n * val_pct / 100)\n # rest goes to test\n\n splits = [\"train\"] * n_train + [\"val\"] * n_val + [\"test\"] * (n - n_train - n_val)\n\n # Write manifest\n output_csv = Path(args.output_csv)\n with open(output_csv, \"w\", newline=\"\") as f:\n writer = csv.writer(f)\n writer.writerow([\"subject_id\", \"dataset_id\", \"t1w_path\", \"label_path\", \"split\"])\n for (subj_id, orig_path, aseg_path), split in zip(pairs, splits):\n writer.writerow([f\"pac_{subj_id}\", \"kwyk\", orig_path, aseg_path, split])\n\n # Summary\n from collections import Counter\n\n split_counts = Counter(splits)\n print(f\"Manifest written to {output_csv}\")\n print(\n f\" train: {split_counts['train']}, \"\n f\"val: {split_counts['val']}, test: {split_counts['test']}\"\n )\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "scripts/kwyk_reproduction/config.yaml", "content": "# KWYK Brain Segmentation Reproduction - Default Configuration\n# Based on: McClure et al., Frontiers in Neuroinformatics 2019\n#\n# The original kwyk study trained 3 model variants, all using\n# Variational Weight Normalization (VWN) convolutions:\n# 1. bwn — VWN conv + Bernoulli dropout OFF at inference (MAP)\n# 2. bwn_multi — VWN conv + Bernoulli dropout ON at inference (MC)\n# 3. bvwn_multi_prior — VWN conv + Concrete dropout (learned rates)\n\n# ---------------------------------------------------------------------------\n# Shared architecture (all variants use identical VWN structure)\n# ---------------------------------------------------------------------------\nfilters: 96\nreceptive_field: 37 # dilation schedule [1,1,1,2,4,8,1]\ndropout_rate: 0.25\nsigma_init: 0.0001 # initial |kernel_a| for VWN weight sigma\n\n# Training\nblock_shape: [32, 32, 32]\nlr: 0.0001\nbatch_size: 32 # per-GPU; auto-optimized when GPU available\nn_classes: 50 # 50-class FreeSurfer parcellation (matching paper)\nlabel_mapping: 50-class # uses label_mappings/50-class-mapping.csv\n\n# Class weighting — the original paper used unweighted CrossEntropyLoss.\n# Set to \"auto\" or \"median_frequency\" for experiments with class balancing.\n# Options: null (paper default), \"auto\" (inverse frequency), \"median_frequency\"\nclass_weight_method: null\n\n# Loss function: \"cross_entropy\" (paper default) or \"dice_ce\" (Dice + weighted CE)\nloss: cross_entropy\n\n# Warm-start (deterministic → VWN)\npretrain_epochs: 50\nbayesian_epochs: 50\n\n# Inference\nn_samples: 10 # MC inference samples\n\n# Data augmentation (profile: none, light, standard, heavy)\naugmentation_profile: standard\n\n# Zarr optimization (optional — convert NIfTI to Zarr for faster I/O)\n# zarr_store: null # path like \"data/brain_store.zarr\"\n# zarr_chunk_shape: [32, 32, 32]\n\n# Stride for patch extraction (null = random, or [sD, sH, sW] for grid)\n# stride: null\n\n# Data assembly\ndatasets:\n - ds000114\n - ds000228\n - ds002609\nsplit: [80, 10, 10] # train/val/test percentages\n\n# ---------------------------------------------------------------------------\n# Model variant configurations (matching original kwyk)\n# ---------------------------------------------------------------------------\nvariants:\n # 1. bwn — VWN conv, Bernoulli dropout, MC OFF at inference (MAP)\n bwn:\n model: kwyk_meshnet\n dropout_type: bernoulli\n mc_at_inference: false\n description: \"VWN MeshNet, Bernoulli dropout, deterministic inference (MAP)\"\n\n # 2. bwn_multi — VWN conv, Bernoulli dropout, MC ON at inference\n bwn_multi:\n model: kwyk_meshnet\n dropout_type: bernoulli\n mc_at_inference: true\n description: \"VWN MeshNet, MC Bernoulli dropout at inference\"\n\n # 3. bvwn_multi_prior — VWN conv, Concrete dropout (learned per-filter rates)\n bvwn_multi_prior:\n model: kwyk_meshnet\n dropout_type: concrete\n concrete_temperature: 0.02\n concrete_init_p: 0.9\n mc_at_inference: true\n description: \"VWN MeshNet, Concrete dropout (learned rates)\"\n\n# ---------------------------------------------------------------------------\n# Presets for different scales\n# ---------------------------------------------------------------------------\n\n# Quick binary brain extraction (for initial testing)\nbinary_preset:\n n_classes: 2\n label_mapping: binary\n class_weight_method: null\n\n# Smoke test overrides\nsmoke_test:\n filters: 16\n block_shape: [16, 16, 16]\n pretrain_epochs: 1\n bayesian_epochs: 1\n batch_size: 2\n n_samples: 3\n n_classes: 2\n label_mapping: binary\n class_weight_method: null\n" }, { "path": "scripts/kwyk_reproduction/config_kwyk_smoke.yaml", "content": "# KWYK PAC Dataset Smoke Test — 50-class parcellation, 100 subjects\n# Based on: McClure et al., Frontiers in Neuroinformatics 2019\n\n# ---------------------------------------------------------------------------\n# Shared architecture (all variants use identical VWN structure)\n# ---------------------------------------------------------------------------\nfilters: 96\nreceptive_field: 37 # dilation schedule [1,1,1,2,4,8,1]\ndropout_rate: 0.25\nsigma_init: 0.0001 # initial |kernel_a| for VWN weight sigma\n\n# Training\nblock_shape: [32, 32, 32]\nlr: 0.0001\nbatch_size: 256\nn_classes: 50 # 50-class FreeSurfer parcellation\nlabel_mapping: 50-class\npatches_per_volume: 50 # random patches per volume per epoch (GPU utilization)\n\n# Warm-start (deterministic → VWN)\npretrain_epochs: 50\nbayesian_epochs: 50\n\n# Validation\nval_dice_freq: 5 # full-volume Dice every N epochs (block-level metrics every epoch)\n\n# Inference\nn_samples: 10 # MC inference samples\n\n# Zarr store (fast chunk-aligned I/O, created by slurm_convert_zarr.sbatch)\nzarr_store: /orcd/scratch/orcd/013/satra/kwyk_reproduction/data/kwyk_500.zarr\n\n# Data augmentation\naugmentation: real # real, synthetic, mixed\n\n# Data assembly\ndatasets:\n - kwyk\nsplit: [80, 10, 10] # train/val/test percentages\n\n# ---------------------------------------------------------------------------\n# Model variant configurations (matching original kwyk)\n# ---------------------------------------------------------------------------\nvariants:\n bwn:\n model: kwyk_meshnet\n dropout_type: bernoulli\n mc_at_inference: false\n description: \"VWN MeshNet, Bernoulli dropout, deterministic inference (MAP)\"\n\n bwn_multi:\n model: kwyk_meshnet\n dropout_type: bernoulli\n mc_at_inference: true\n description: \"VWN MeshNet, MC Bernoulli dropout at inference\"\n\n bvwn_multi_prior:\n model: kwyk_meshnet\n dropout_type: concrete\n concrete_temperature: 0.02\n concrete_init_p: 0.9\n mc_at_inference: true\n description: \"VWN MeshNet, Concrete dropout (learned rates)\"\n" }, { "path": "scripts/kwyk_reproduction/convert_zarr_shard.py", "content": "#!/usr/bin/env python\n\"\"\"Convert one shard of NIfTI volumes to a pre-created Zarr3 store.\n\nUsage:\n python convert_zarr_shard.py --manifest manifest.csv --zarr-store data/store.zarr \\\n --shard-idx 0 --subjects-per-shard 50\n\nCalled by SLURM job array — each task writes one shard independently.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport argparse\nimport csv\nfrom pathlib import Path\nimport time\n\nimport nibabel as nib\nimport numpy as np\nimport zarr\n\n\ndef main():\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--manifest\", required=True)\n parser.add_argument(\"--zarr-store\", required=True)\n parser.add_argument(\"--shard-idx\", type=int, required=True)\n parser.add_argument(\"--subjects-per-shard\", type=int, default=50)\n parser.add_argument(\n \"--create\",\n action=\"store_true\",\n help=\"Create the store (only shard 0 should do this)\",\n )\n args = parser.parse_args()\n\n # Read manifest\n pairs = []\n subject_ids = []\n with open(args.manifest) as f:\n for row in csv.DictReader(f):\n pairs.append((row[\"t1w_path\"], row[\"label_path\"]))\n subject_ids.append(row[\"subject_id\"])\n\n n_subjects = len(pairs)\n sps = args.subjects_per_shard\n start = args.shard_idx * sps\n end = min(start + sps, n_subjects)\n\n if start >= n_subjects:\n print(f\"Shard {args.shard_idx}: no subjects (start={start} >= {n_subjects})\")\n return\n\n store_path = Path(args.zarr_store)\n\n if args.create:\n # Create the store and arrays (only one task does this)\n D, H, W = 256, 256, 256\n n_shards = (n_subjects + sps - 1) // sps\n store = zarr.open_group(str(store_path), mode=\"w\")\n store.create_array(\n \"images\",\n shape=(n_subjects, D, H, W),\n chunks=(1, 32, 32, 32),\n shards=(sps, D, H, W),\n dtype=np.float32,\n )\n store.create_array(\n \"labels\",\n shape=(n_subjects, D, H, W),\n chunks=(1, 32, 32, 32),\n shards=(sps, D, H, W),\n dtype=np.int32,\n )\n store.attrs[\"n_subjects\"] = n_subjects\n store.attrs[\"subject_ids\"] = subject_ids\n store.attrs[\"volume_shape\"] = [D, H, W]\n print(f\"Created store: {store_path} ({n_subjects} subjects, {n_shards} shards)\")\n\n # Write partition JSON\n import json\n\n partitions = {\"train\": [], \"val\": [], \"test\": []}\n with open(args.manifest) as f:\n for row in csv.DictReader(f):\n partitions[row[\"split\"]].append(row[\"subject_id\"])\n part_path = str(store_path) + \"_partition.json\"\n with open(part_path, \"w\") as f:\n json.dump({\"partitions\": partitions}, f, indent=2)\n for k, v in partitions.items():\n print(f\" {k}: {len(v)} subjects\")\n else:\n # Open existing store in append mode\n store = zarr.open_group(str(store_path), mode=\"r+\")\n\n images_arr = store[\"images\"]\n labels_arr = store[\"labels\"]\n\n t0 = time.time()\n for i in range(start, end):\n img_path, lbl_path = pairs[i]\n # PAC data is already 256³ @ 1mm uint8/int32 — no conform needed\n img_data = np.asarray(nib.load(img_path).dataobj, dtype=np.float32)\n lbl_data = np.asarray(nib.load(lbl_path).dataobj, dtype=np.int32)\n images_arr[i] = img_data[:256, :256, :256]\n labels_arr[i] = lbl_data[:256, :256, :256]\n\n if (i - start + 1) % 10 == 0:\n elapsed = time.time() - t0\n rate = (i - start + 1) / elapsed\n print(\n f\" Shard {args.shard_idx}: {i - start + 1}/{end - start} \"\n f\"({rate:.1f} vol/s, {elapsed:.0f}s)\"\n )\n\n elapsed = time.time() - t0\n print(\n f\"Shard {args.shard_idx}: wrote {end - start} volumes in {elapsed:.1f}s \"\n f\"({(end - start) / elapsed:.1f} vol/s)\"\n )\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "scripts/kwyk_reproduction/experiments/01_20260330_eval_deterministic/README.md", "content": "# Experiment 01: Evaluate Bayesian models in deterministic mode\n\n## Rationale\n\nAll 3 Bayesian variants show zero Dice during MC evaluation despite training loss\ndecreasing from ~3.8 to ~2.2. The prediction code calls `model(tensor)` which\ndefaults to `mc=True` in KWYKMeshNet.forward(), activating local reparameterization\nnoise and dropout. With only 20 epochs of Bayesian training, this noise may\noverwhelm the learned signal.\n\n**Hypothesis:** The model weights have learned meaningful representations, but MC\ninference noise destroys the output. Evaluating with `mc=False` should show non-zero\nDice.\n\n## Plan\n\n1. Write a quick eval script that loads each Bayesian checkpoint and runs prediction\n with `mc=False` (deterministic forward pass)\n2. Compare per-class Dice between mc=True and mc=False\n3. No retraining needed — just evaluate existing checkpoints\n\n## Tasks\n\n- [x] Write eval script with mc=False support\n- [x] Run on existing 20-epoch checkpoints\n- [ ] Compare results\n" }, { "path": "scripts/kwyk_reproduction/experiments/01_20260330_eval_deterministic/eval_deterministic.py", "content": "#!/usr/bin/env python\n\"\"\"Evaluate Bayesian models in deterministic mode (mc=False).\n\nQuick diagnostic: do the weights contain useful information that MC noise destroys?\n\"\"\"\n\nfrom __future__ import annotations\n\nimport csv\nfrom pathlib import Path\nimport sys\n\nimport nibabel as nib\nimport numpy as np\nimport torch\n\nsys.path.insert(0, str(Path(__file__).parent.parent.parent))\nfrom utils import setup_logging # noqa: E402\n\nlog = setup_logging(__name__)\n\n\ndef per_class_dice(pred: np.ndarray, gt: np.ndarray, n_classes: int) -> np.ndarray:\n \"\"\"Per-class Dice for classes 1..n_classes-1.\"\"\"\n dice = np.zeros(n_classes - 1)\n for c in range(1, n_classes):\n p = pred == c\n g = gt == c\n inter = (p & g).sum()\n total = p.sum() + g.sum()\n dice[c - 1] = 2.0 * inter / total if total > 0 else 1.0\n return dice\n\n\ndef predict_volume(model, img_path, block_shape, mc=False):\n \"\"\"Block-based prediction on a single volume.\"\"\"\n from nobrainer.prediction import _extract_blocks, _pad_to_multiple, _stitch_blocks\n from nobrainer.training import get_device\n\n device = get_device()\n img = nib.load(str(img_path))\n arr = np.asarray(img.dataobj, dtype=np.float32)\n orig_shape = arr.shape[:3]\n padded, pad = _pad_to_multiple(arr, block_shape)\n blocks, grid = _extract_blocks(padded, block_shape)\n\n model = model.to(device)\n model.eval()\n\n all_preds = []\n with torch.no_grad():\n for start in range(0, len(blocks), 4):\n chunk = blocks[start : start + 4]\n tensor = torch.from_numpy(chunk[:, None]).to(device)\n out = model(tensor, mc=mc)\n labels = out.argmax(dim=1, keepdim=True).float()\n all_preds.append(labels.cpu().numpy())\n\n block_preds = np.concatenate(all_preds, axis=0)\n full = _stitch_blocks(block_preds, grid, block_shape, pad, orig_shape, 1)[0]\n return full.astype(np.int32)\n\n\ndef main():\n from nobrainer.processing.dataset import _load_label_mapping\n from nobrainer.processing.segmentation import Segmentation\n\n work_dir = Path(__file__).parent.parent.parent\n manifest_path = work_dir / \"kwyk_manifest.csv\"\n remap_fn = _load_label_mapping(\"50-class\")\n n_classes = 50\n\n # Load test pairs\n pairs = []\n with open(manifest_path) as f:\n for row in csv.DictReader(f):\n if row[\"split\"] == \"test\":\n pairs.append((row[\"t1w_path\"], row[\"label_path\"]))\n\n log.info(\"Test volumes: %d\", len(pairs))\n\n # Evaluate each variant\n variants = [\n \"kwyk_smoke_bwn_multi\",\n \"kwyk_smoke_bvwn_multi_prior\",\n \"kwyk_smoke_bayesian_gaussian\",\n ]\n\n results = []\n for variant in variants:\n ckpt_dir = work_dir / \"checkpoints\" / variant\n if not (ckpt_dir / \"model.pth\").exists():\n log.warning(\"Skipping %s — no checkpoint\", variant)\n continue\n\n log.info(\"=== %s ===\", variant)\n seg = Segmentation.load(ckpt_dir)\n model = seg.model_\n block_shape = seg.block_shape_ or (32, 32, 32)\n\n for mc_mode in [False, True]:\n mode_name = \"mc\" if mc_mode else \"deterministic\"\n all_dice = []\n\n for idx, (img_path, lbl_path) in enumerate(pairs[:3]): # first 3 for speed\n pred_arr = predict_volume(model, img_path, block_shape, mc=mc_mode)\n gt_arr = np.asarray(nib.load(lbl_path).dataobj, dtype=np.int32)\n gt_arr = remap_fn(torch.from_numpy(gt_arr)).numpy()\n\n cd = per_class_dice(pred_arr, gt_arr, n_classes)\n avg = float(cd.mean())\n all_dice.append(avg)\n log.info(\n \" [%s] vol %d: avg_dice=%.4f max=%.4f\",\n mode_name,\n idx + 1,\n avg,\n cd.max(),\n )\n\n mean_dice = float(np.mean(all_dice))\n results.append(\n {\"variant\": variant, \"mode\": mode_name, \"mean_dice\": mean_dice}\n )\n log.info(\" [%s] MEAN: %.4f\", mode_name, mean_dice)\n\n # Save results\n out_path = Path(__file__).parent / \"results.csv\"\n with open(out_path, \"w\", newline=\"\") as f:\n w = csv.DictWriter(f, fieldnames=[\"variant\", \"mode\", \"mean_dice\"])\n w.writeheader()\n w.writerows(results)\n log.info(\"Results saved to %s\", out_path)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "scripts/kwyk_reproduction/experiments/01_20260330_eval_deterministic/results_summary.md", "content": "# Experiment 01 Results\n\n## Finding\n\n**Both deterministic and MC modes produce zero Dice for all 3 Bayesian variants.**\n\n| Variant | Deterministic | MC |\n|---|---|---|\n| bwn_multi | 0.0000 | 0.0000 |\n| bvwn_multi_prior | 0.0000 | 0.0001 |\n| bayesian_gaussian | 0.0000 | 0.0000 |\n\n## Conclusion\n\nThe hypothesis was wrong — the issue is NOT MC noise destroying learned signals.\nThe weights themselves contain no useful information. Possible causes:\n\n1. **Warm-start transfer failure**: weights not properly transferred from MeshNet to KWYKMeshNet\n2. **Bayesian training destroying warm-start**: ELBO loss / weight perturbation undoing the transferred weights\n3. **Architecture mismatch**: MeshNet vs KWYKMeshNet parameter shapes may not align\n\n## Next Steps\n\n- Investigate warm-start transfer code\n- Check if Bayesian model immediately after warm-start (before training) can segment\n- Compare model architectures between MeshNet and KWYKMeshNet\n" }, { "path": "scripts/kwyk_reproduction/experiments/01_20260330_eval_deterministic/run.sbatch", "content": "#!/bin/bash\n#SBATCH --job-name=exp01-det\n#SBATCH --partition=mit_preemptable\n#SBATCH --gres=gpu:1\n#SBATCH --cpus-per-task=4\n#SBATCH --mem=32G\n#SBATCH --time=00:30:00\n#SBATCH --output=slurm-%j.out\n#SBATCH --error=slurm-%j.err\nset -euo pipefail\n\ncd /orcd/scratch/orcd/013/satra/kwyk_reproduction\nsource /orcd/data/satra/002/projects/nobrainer/venvs/nobrainer/bin/activate\n\necho \"=== Experiment 01: Deterministic eval of Bayesian models ===\"\necho \"Started: $(date)\"\n\npython experiments/01_20260330_eval_deterministic/eval_deterministic.py\n\necho \"=== Done: $(date) ===\"\n" }, { "path": "scripts/kwyk_reproduction/experiments/02_20260330_binary_bayesian/README.md", "content": "# Experiment 02: Binary (2-class) Bayesian training\n\n## Rationale\n\n50-class parcellation is a hard problem. Binary brain extraction (brain vs background)\nis much simpler and was used in the original smoke tests. If Bayesian models can learn\nbinary segmentation, the issue is with label complexity, not the Bayesian architecture.\n\n## Plan\n\n1. Use 5 subjects, binary label mapping, 20 epochs\n2. Train MeshNet → warm-start bwn_multi (simplest Bayesian variant)\n3. Evaluate both mc=True and mc=False\n4. If binary works, the 50-class zero Dice is likely a capacity/epochs issue\n\n## Tasks\n\n- [ ] Create binary config\n- [ ] Train MeshNet (binary, 5 subjects, 20 epochs)\n- [ ] Train bwn_multi (binary, warm-start, 20 epochs)\n- [ ] Evaluate both modes\n" }, { "path": "scripts/kwyk_reproduction/experiments/02_20260330_binary_bayesian/config.yaml", "content": "filters: 96\nreceptive_field: 37\ndropout_rate: 0.25\nsigma_init: 0.0001\n\nblock_shape: [32, 32, 32]\nlr: 0.0001\nbatch_size: 32\nn_classes: 2\nlabel_mapping: binary\n\npretrain_epochs: 20\nbayesian_epochs: 20\nn_samples: 10\n\nvariants:\n bwn_multi:\n model: kwyk_meshnet\n dropout_type: bernoulli\n mc_at_inference: true\n" }, { "path": "scripts/kwyk_reproduction/experiments/02_20260330_binary_bayesian/eval_binary.py", "content": "#!/usr/bin/env python\n\"\"\"Evaluate binary Bayesian model in both mc=True and mc=False modes.\"\"\"\n\nfrom __future__ import annotations\n\nimport csv\nfrom pathlib import Path\nimport sys\n\nimport nibabel as nib\nimport numpy as np\nimport torch\n\nsys.path.insert(0, str(Path(__file__).parent.parent.parent))\nfrom utils import compute_dice, setup_logging # noqa: E402\n\nlog = setup_logging(__name__)\n\nEXP_DIR = Path(__file__).parent\nWORK_DIR = EXP_DIR.parent.parent\n\n\ndef predict_volume(model, img_path, block_shape, mc=False):\n \"\"\"Block-based prediction with mc control.\"\"\"\n from nobrainer.prediction import _extract_blocks, _pad_to_multiple, _stitch_blocks\n from nobrainer.training import get_device\n\n device = get_device()\n img = nib.load(str(img_path))\n arr = np.asarray(img.dataobj, dtype=np.float32)\n orig_shape = arr.shape[:3]\n padded, pad = _pad_to_multiple(arr, block_shape)\n blocks, grid = _extract_blocks(padded, block_shape)\n\n model = model.to(device)\n model.eval()\n\n all_preds = []\n with torch.no_grad():\n for start in range(0, len(blocks), 4):\n chunk = blocks[start : start + 4]\n tensor = torch.from_numpy(chunk[:, None]).to(device)\n if hasattr(model, \"forward\") and \"mc\" in model.forward.__code__.co_varnames:\n out = model(tensor, mc=mc)\n else:\n out = model(tensor)\n labels = out.argmax(dim=1, keepdim=True).float()\n all_preds.append(labels.cpu().numpy())\n\n block_preds = np.concatenate(all_preds, axis=0)\n full = _stitch_blocks(block_preds, grid, block_shape, pad, orig_shape, 1)[0]\n return (full > 0).astype(np.float32)\n\n\ndef main():\n from nobrainer.processing.segmentation import Segmentation\n\n manifest_path = WORK_DIR / \"kwyk_sanity_manifest.csv\"\n\n # Load test pairs\n pairs = []\n with open(manifest_path) as f:\n for row in csv.DictReader(f):\n if row[\"split\"] == \"test\":\n pairs.append((row[\"t1w_path\"], row[\"label_path\"]))\n\n log.info(\"Test volumes: %d\", len(pairs))\n\n results = []\n for variant in [\"meshnet\", \"bwn_multi\"]:\n ckpt_dir = EXP_DIR / \"checkpoints\" / variant\n if not (ckpt_dir / \"model.pth\").exists():\n log.warning(\"Skipping %s — no checkpoint\", variant)\n continue\n\n seg = Segmentation.load(ckpt_dir)\n model = seg.model_\n block_shape = seg.block_shape_ or (32, 32, 32)\n\n mc_modes = [False] if variant == \"meshnet\" else [False, True]\n for mc_mode in mc_modes:\n mode_name = \"mc\" if mc_mode else \"deterministic\"\n dices = []\n\n for idx, (img_path, lbl_path) in enumerate(pairs):\n pred = predict_volume(model, img_path, block_shape, mc=mc_mode)\n gt = (\n np.asarray(nib.load(lbl_path).dataobj, dtype=np.float32) > 0\n ).astype(np.float32)\n dice = compute_dice(pred, gt)\n dices.append(dice)\n log.info(\n \" [%s/%s] vol %d: Dice=%.4f\", variant, mode_name, idx + 1, dice\n )\n\n mean_d = float(np.mean(dices))\n results.append(\n {\"variant\": variant, \"mode\": mode_name, \"mean_dice\": f\"{mean_d:.4f}\"}\n )\n log.info(\" [%s/%s] MEAN DICE: %.4f\", variant, mode_name, mean_d)\n\n out_path = EXP_DIR / \"results.csv\"\n with open(out_path, \"w\", newline=\"\") as f:\n w = csv.DictWriter(f, fieldnames=[\"variant\", \"mode\", \"mean_dice\"])\n w.writeheader()\n w.writerows(results)\n log.info(\"Results saved to %s\", out_path)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "scripts/kwyk_reproduction/experiments/02_20260330_binary_bayesian/eval_only.sbatch", "content": "#!/bin/bash\n#SBATCH --job-name=exp02-eval\n#SBATCH --partition=mit_preemptable\n#SBATCH --gres=gpu:1\n#SBATCH --cpus-per-task=4\n#SBATCH --mem=32G\n#SBATCH --time=00:15:00\n#SBATCH --output=slurm-eval-%j.out\n#SBATCH --error=slurm-eval-%j.err\nset -euo pipefail\n\ncd /orcd/scratch/orcd/013/satra/kwyk_reproduction\nsource /orcd/data/satra/002/projects/nobrainer/venvs/nobrainer/bin/activate\n\necho \"=== Experiment 02: Binary eval ===\"\npython experiments/02_20260330_binary_bayesian/eval_binary.py\necho \"=== Done: $(date) ===\"\n" }, { "path": "scripts/kwyk_reproduction/experiments/02_20260330_binary_bayesian/run.sbatch", "content": "#!/bin/bash\n#SBATCH --job-name=exp02-bin\n#SBATCH --partition=mit_preemptable\n#SBATCH --gres=gpu:1\n#SBATCH --cpus-per-task=4\n#SBATCH --mem=32G\n#SBATCH --time=01:00:00\n#SBATCH --output=slurm-%j.out\n#SBATCH --error=slurm-%j.err\nset -euo pipefail\n\nWORK_DIR=\"/orcd/scratch/orcd/013/satra/kwyk_reproduction\"\nEXP_DIR=\"$WORK_DIR/experiments/02_20260330_binary_bayesian\"\nVENV_DIR=\"/orcd/data/satra/002/projects/nobrainer/venvs/nobrainer\"\nMANIFEST=\"$WORK_DIR/kwyk_sanity_manifest.csv\" # 5 subjects\nCONFIG=\"$EXP_DIR/config.yaml\"\n\ncd \"$WORK_DIR\"\nsource \"${VENV_DIR}/bin/activate\"\n\necho \"=== Experiment 02: Binary Bayesian (5 subjects, 20 epochs) ===\"\necho \"Started: $(date)\"\n\n# Train MeshNet (binary)\necho \"=== Step 1: MeshNet (binary, 20 epochs) ===\"\npython 02_train_meshnet.py --manifest \"$MANIFEST\" --config \"$CONFIG\" \\\n --output-dir \"$EXP_DIR/checkpoints/meshnet\" --epochs 20\n\n# Train bwn_multi (binary, warm-start from MeshNet)\necho \"=== Step 2: bwn_multi (binary, 20 epochs) ===\"\npython 03_train_bayesian.py --manifest \"$MANIFEST\" --config \"$CONFIG\" \\\n --variant bwn_multi --warmstart \"$EXP_DIR/checkpoints/meshnet\" \\\n --output-dir \"$EXP_DIR/checkpoints/bwn_multi\" --epochs 20\n\n# Quick eval: deterministic and MC\necho \"=== Step 3: Eval (deterministic + MC) ===\"\npython experiments/02_20260330_binary_bayesian/eval_binary.py\n\necho \"=== Done: $(date) ===\"\n" }, { "path": "scripts/kwyk_reproduction/experiments/03_20260330_warmstart_diagnostic/README.md", "content": "# Experiment 03: Warm-start transfer diagnostic\n\n## Rationale\n\nExperiment 01 showed that Bayesian model weights have zero Dice even in deterministic\nmode. This means the warm-start transfer from MeshNet to KWYKMeshNet may not be\nworking, OR the Bayesian training loop destroys transferred weights.\n\n## Plan\n\n1. Load trained MeshNet checkpoint\n2. Create KWYKMeshNet and run warm-start transfer\n3. Evaluate KWYKMeshNet immediately BEFORE any Bayesian training (mc=False)\n4. If Dice > 0: warm-start works, Bayesian training is the problem\n5. If Dice = 0: warm-start transfer is broken\n6. Also compare parameter counts and shapes between MeshNet and KWYKMeshNet\n\n## Tasks\n\n- [ ] Compare architectures\n- [ ] Evaluate warm-started model before training\n- [ ] Check transfer log messages\n" }, { "path": "scripts/kwyk_reproduction/experiments/03_20260330_warmstart_diagnostic/diagnose.py", "content": "#!/usr/bin/env python\n\"\"\"Diagnose warm-start transfer from MeshNet to KWYKMeshNet.\"\"\"\n\nfrom __future__ import annotations\n\nimport csv\nfrom pathlib import Path\nimport sys\n\nimport nibabel as nib\nimport numpy as np\nimport torch\n\nsys.path.insert(0, str(Path(__file__).parent.parent.parent))\nfrom utils import setup_logging # noqa: E402\n\nlog = setup_logging(__name__)\n\nWORK_DIR = Path(__file__).parent.parent.parent\nEXP_DIR = Path(__file__).parent\n\n\ndef predict_volume_simple(model, img_path, block_shape, mc=False):\n \"\"\"Block-based prediction.\"\"\"\n from nobrainer.prediction import _extract_blocks, _pad_to_multiple, _stitch_blocks\n from nobrainer.training import get_device\n\n device = get_device()\n img = nib.load(str(img_path))\n arr = np.asarray(img.dataobj, dtype=np.float32)\n orig_shape = arr.shape[:3]\n padded, pad = _pad_to_multiple(arr, block_shape)\n blocks, grid = _extract_blocks(padded, block_shape)\n\n model = model.to(device)\n model.eval()\n\n all_preds = []\n with torch.no_grad():\n for start in range(0, len(blocks), 4):\n chunk = blocks[start : start + 4]\n tensor = torch.from_numpy(chunk[:, None]).to(device)\n try:\n out = model(tensor, mc=mc)\n except TypeError:\n out = model(tensor)\n labels = out.argmax(dim=1, keepdim=True).float()\n all_preds.append(labels.cpu().numpy())\n\n block_preds = np.concatenate(all_preds, axis=0)\n full = _stitch_blocks(block_preds, grid, block_shape, pad, orig_shape, 1)[0]\n return full.astype(np.int32)\n\n\ndef main():\n from nobrainer.models import get as get_model\n from nobrainer.processing.dataset import _load_label_mapping\n from nobrainer.processing.segmentation import Segmentation\n\n remap_fn = _load_label_mapping(\"50-class\")\n n_classes = 50\n block_shape = (32, 32, 32)\n\n # Load test pairs\n pairs = []\n with open(WORK_DIR / \"kwyk_manifest.csv\") as f:\n for row in csv.DictReader(f):\n if row[\"split\"] == \"test\":\n pairs.append((row[\"t1w_path\"], row[\"label_path\"]))\n\n # Use first test volume\n img_path, lbl_path = pairs[0]\n gt_arr = remap_fn(\n torch.from_numpy(np.asarray(nib.load(lbl_path).dataobj, dtype=np.int32))\n ).numpy()\n\n # ---- Step 1: Load trained MeshNet and eval ----\n log.info(\"=== Step 1: Evaluate trained MeshNet ===\")\n seg = Segmentation.load(WORK_DIR / \"checkpoints\" / \"kwyk_smoke_meshnet\")\n det_model = seg.model_\n log.info(\"MeshNet type: %s\", type(det_model).__name__)\n log.info(\"MeshNet params: %d\", sum(p.numel() for p in det_model.parameters()))\n\n pred = predict_volume_simple(det_model, img_path, block_shape, mc=False)\n # per-class dice\n dices = []\n for c in range(1, n_classes):\n p = pred == c\n g = gt_arr == c\n inter = (p & g).sum()\n total = p.sum() + g.sum()\n dices.append(2.0 * inter / total if total > 0 else 1.0)\n log.info(\"MeshNet Dice: mean=%.4f, max=%.4f\", np.mean(dices), np.max(dices))\n\n # ---- Step 2: Create KWYKMeshNet and warm-start ----\n log.info(\"=== Step 2: Warm-start KWYKMeshNet from MeshNet ===\")\n model_args = {\n \"n_classes\": n_classes,\n \"filters\": 96,\n \"receptive_field\": 37,\n \"dropout_type\": \"bernoulli\",\n \"dropout_rate\": 0.25,\n \"sigma_init\": 0.0001,\n }\n\n kwyk_factory = get_model(\"kwyk_meshnet\")\n kwyk_model = kwyk_factory(**model_args)\n log.info(\"KWYKMeshNet type: %s\", type(kwyk_model).__name__)\n log.info(\"KWYKMeshNet params: %d\", sum(p.numel() for p in kwyk_model.parameters()))\n\n # Print layer comparison\n log.info(\"--- MeshNet layers ---\")\n for name, param in det_model.named_parameters():\n log.info(\" %s: %s\", name, param.shape)\n\n log.info(\"--- KWYKMeshNet layers ---\")\n for name, param in kwyk_model.named_parameters():\n log.info(\" %s: %s\", name, param.shape)\n\n # Run warm-start\n from nobrainer.models.bayesian.warmstart import warmstart_kwyk_from_deterministic\n\n meshnet_ckpt = WORK_DIR / \"checkpoints\" / \"kwyk_smoke_meshnet\" / \"model.pth\"\n n_transferred = warmstart_kwyk_from_deterministic(kwyk_model, str(meshnet_ckpt))\n log.info(\"Transferred %d layers\", n_transferred)\n\n # ---- Step 3: Eval KWYKMeshNet BEFORE any Bayesian training ----\n log.info(\"=== Step 3: Evaluate warm-started KWYKMeshNet (mc=False) ===\")\n pred = predict_volume_simple(kwyk_model, img_path, block_shape, mc=False)\n dices = []\n for c in range(1, n_classes):\n p = pred == c\n g = gt_arr == c\n inter = (p & g).sum()\n total = p.sum() + g.sum()\n dices.append(2.0 * inter / total if total > 0 else 1.0)\n log.info(\n \"KWYKMeshNet (warm-start, mc=False) Dice: mean=%.4f, max=%.4f\",\n np.mean(dices),\n np.max(dices),\n )\n\n # Also test mc=True\n log.info(\"=== Step 4: Evaluate warm-started KWYKMeshNet (mc=True) ===\")\n pred = predict_volume_simple(kwyk_model, img_path, block_shape, mc=True)\n dices = []\n for c in range(1, n_classes):\n p = pred == c\n g = gt_arr == c\n inter = (p & g).sum()\n total = p.sum() + g.sum()\n dices.append(2.0 * inter / total if total > 0 else 1.0)\n log.info(\n \"KWYKMeshNet (warm-start, mc=True) Dice: mean=%.4f, max=%.4f\",\n np.mean(dices),\n np.max(dices),\n )\n\n # ---- Step 5: Check what 03_train_bayesian does ----\n log.info(\"=== Step 5: Check how training script loads warm-start ===\")\n # Read the training script to see if it uses warmstart_kwyk_from_deterministic\n train_script = WORK_DIR / \"03_train_bayesian.py\"\n with open(train_script) as f:\n content = f.read()\n if \"warmstart_kwyk_from_deterministic\" in content:\n log.info(\"Training script uses warmstart_kwyk_from_deterministic\")\n elif \"warmstart_bayesian_from_deterministic\" in content:\n log.info(\n \"Training script uses warmstart_bayesian_from_deterministic (WRONG for KWYK!)\"\n )\n else:\n log.info(\"No warmstart function found in training script — check manually\")\n\n # Grep for the relevant line\n for line_no, line in enumerate(content.split(\"\\n\"), 1):\n if \"warmstart\" in line.lower() and not line.strip().startswith(\"#\"):\n log.info(\" Line %d: %s\", line_no, line.strip())\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "scripts/kwyk_reproduction/experiments/03_20260330_warmstart_diagnostic/results_summary.md", "content": "# Experiment 03 Results\n\n## Key Finding: Warm-start transfer bug — sorted key ordering mismatch\n\n**MeshNet Dice (trained):** mean=0.0132, max=0.4738\n**KWYKMeshNet (warm-start, mc=False):** mean=0.0006, max=0.0101\n**KWYKMeshNet (warm-start, mc=True):** mean=0.0006, max=0.0103\n\nOnly 5 of 7+1 layers transferred successfully.\n\n## Root Cause\n\n`warmstart_kwyk_from_deterministic()` sorts state dict keys alphabetically:\n```python\ndet_convs = [(k, v) for k, v in sorted(state.items()) if \"weight\" in k and v.ndim == 5]\n```\n\nThis produces ordering: `classifier.weight, encoder.0, encoder.1, ...`\n\nBut KWYKMeshNet FFGConv3d layers are: `layer_0, layer_1, ...` (no classifier — it's a regular Conv3d)\n\nSo `classifier.weight [50,96,1,1,1]` pairs with `layer_0.conv [96,1,3,3,3]` → shape mismatch!\nThen `encoder.0 [96,1,3,3,3]` pairs with `layer_1.conv [96,96,3,3,3]` → shape mismatch!\nThen `encoder.1 [96,96,3,3,3]` pairs with `layer_2.conv [96,96,3,3,3]` → OK (but wrong weights!)\n\nResult: 5 layers \"transferred\" but with wrong weight assignments (encoder.1→layer_2 instead of\nencoder.0→layer_0), and first two layers get random initialization.\n\n## Fix\n\nFilter out the classifier weight before pairing, OR use explicit name matching.\n\n## Conclusion\n\nThe Bayesian zero Dice is caused by a broken warm-start. The model starts from\nmostly-random weights and 20 epochs isn't enough to learn from scratch.\n" }, { "path": "scripts/kwyk_reproduction/experiments/03_20260330_warmstart_diagnostic/run.sbatch", "content": "#!/bin/bash\n#SBATCH --job-name=exp03-ws\n#SBATCH --partition=mit_preemptable\n#SBATCH --gres=gpu:1\n#SBATCH --cpus-per-task=4\n#SBATCH --mem=32G\n#SBATCH --time=00:15:00\n#SBATCH --output=slurm-%j.out\n#SBATCH --error=slurm-%j.err\nset -euo pipefail\n\ncd /orcd/scratch/orcd/013/satra/kwyk_reproduction\nsource /orcd/data/satra/002/projects/nobrainer/venvs/nobrainer/bin/activate\n\necho \"=== Experiment 03: Warm-start diagnostic ===\"\npython experiments/03_20260330_warmstart_diagnostic/diagnose.py\necho \"=== Done: $(date) ===\"\n" }, { "path": "scripts/kwyk_reproduction/experiments/04_20260330_fixed_warmstart/README.md", "content": "# Experiment 04: Fix warm-start and verify Bayesian learning\n\n## Rationale\n\nExperiment 03 found that `warmstart_kwyk_from_deterministic()` has a key ordering\nbug: sorted() puts classifier.weight before encoder.X, causing all layer pairings\nto be offset. Fix the transfer, verify Dice is preserved, then train Bayesian.\n\n## Plan\n\n1. Fix warm-start: filter classifier from det_convs, transfer it separately\n2. Verify fixed warm-start preserves MeshNet Dice\n3. Train bwn_multi for 20 epochs with fixed warm-start\n4. Evaluate in both mc=False and mc=True modes\n5. Use 5 subjects for speed (sanity manifest)\n\n## Tasks\n\n- [ ] Fix warm-start function\n- [ ] Verify transfer preserves Dice\n- [ ] Train and evaluate Bayesian\n" }, { "path": "scripts/kwyk_reproduction/experiments/04_20260330_fixed_warmstart/run.py", "content": "#!/usr/bin/env python\n\"\"\"Fix warm-start, verify transfer, train Bayesian, evaluate.\"\"\"\n\nfrom __future__ import annotations\n\nimport csv\nfrom pathlib import Path\nimport sys\n\nimport nibabel as nib\nimport numpy as np\nimport torch\nimport torch.nn as nn\n\nsys.path.insert(0, str(Path(__file__).parent.parent.parent))\nfrom utils import setup_logging # noqa: E402\n\nlog = setup_logging(__name__)\n\nWORK_DIR = Path(__file__).parent.parent.parent\nEXP_DIR = Path(__file__).parent\n\n\ndef fixed_warmstart_kwyk(kwyk_model, det_weights_path):\n \"\"\"Fixed warm-start: filter classifier, match encoder layers correctly.\"\"\"\n from nobrainer.models.bayesian.vwn_layers import FFGConv3d\n\n state = torch.load(det_weights_path, weights_only=True)\n\n # Separate encoder convs from classifier\n encoder_convs = []\n classifier_weight = None\n classifier_bias = None\n for k in sorted(state.keys()):\n v = state[k]\n if k == \"classifier.weight\" and v.ndim == 5:\n classifier_weight = v\n elif k == \"classifier.bias\":\n classifier_bias = v\n elif \"weight\" in k and v.ndim == 5:\n encoder_convs.append((k, v))\n\n log.info(\"Found %d encoder convs + classifier in MeshNet\", len(encoder_convs))\n\n # Transfer encoder convs to FFGConv3d layers\n kwyk_convs = [\n (n, m) for n, m in kwyk_model.named_modules() if isinstance(m, FFGConv3d)\n ]\n\n transferred = 0\n for (det_name, det_w), (kwyk_name, kwyk_conv) in zip(encoder_convs, kwyk_convs):\n if det_w.shape != kwyk_conv.v.shape:\n log.warning(\n \"Shape mismatch: %s %s vs %s.v %s\",\n det_name,\n det_w.shape,\n kwyk_name,\n kwyk_conv.v.shape,\n )\n continue\n kwyk_conv.v.data.copy_(det_w)\n norms = det_w.flatten(1).norm(dim=1).view_as(kwyk_conv.g)\n kwyk_conv.g.data.copy_(norms)\n transferred += 1\n log.info(\" %s -> %s\", det_name, kwyk_name)\n\n # Transfer classifier\n if classifier_weight is not None and hasattr(kwyk_model, \"classifier\"):\n kwyk_model.classifier.weight.data.copy_(classifier_weight)\n if classifier_bias is not None:\n kwyk_model.classifier.bias.data.copy_(classifier_bias)\n log.info(\" classifier transferred\")\n transferred += 1\n\n log.info(\"Total transferred: %d layers\", transferred)\n return transferred\n\n\ndef predict_volume(model, img_path, block_shape, mc=False):\n \"\"\"Block-based prediction.\"\"\"\n from nobrainer.prediction import _extract_blocks, _pad_to_multiple, _stitch_blocks\n from nobrainer.training import get_device\n\n device = get_device()\n img = nib.load(str(img_path))\n arr = np.asarray(img.dataobj, dtype=np.float32)\n orig_shape = arr.shape[:3]\n padded, pad = _pad_to_multiple(arr, block_shape)\n blocks, grid = _extract_blocks(padded, block_shape)\n\n model = model.to(device)\n model.eval()\n\n all_preds = []\n with torch.no_grad():\n for start in range(0, len(blocks), 4):\n chunk = blocks[start : start + 4]\n tensor = torch.from_numpy(chunk[:, None]).to(device)\n try:\n out = model(tensor, mc=mc)\n except TypeError:\n out = model(tensor)\n labels = out.argmax(dim=1, keepdim=True).float()\n all_preds.append(labels.cpu().numpy())\n\n block_preds = np.concatenate(all_preds, axis=0)\n full = _stitch_blocks(block_preds, grid, block_shape, pad, orig_shape, 1)[0]\n return full.astype(np.int32)\n\n\ndef per_class_dice(pred, gt, n_classes):\n dices = []\n for c in range(1, n_classes):\n p = pred == c\n g = gt == c\n inter = (p & g).sum()\n total = p.sum() + g.sum()\n dices.append(2.0 * inter / total if total > 0 else 1.0)\n return np.array(dices)\n\n\ndef main():\n from nobrainer.models import get as get_model\n from nobrainer.processing.dataset import Dataset, _load_label_mapping\n from nobrainer.processing.segmentation import Segmentation\n\n n_classes = 50\n block_shape = (32, 32, 32)\n remap_fn = _load_label_mapping(\"50-class\")\n\n # Test volume\n pairs = []\n with open(WORK_DIR / \"kwyk_sanity_manifest.csv\") as f:\n for row in csv.DictReader(f):\n if row[\"split\"] == \"test\":\n pairs.append((row[\"t1w_path\"], row[\"label_path\"]))\n img_path, lbl_path = pairs[0]\n gt_arr = remap_fn(\n torch.from_numpy(np.asarray(nib.load(lbl_path).dataobj, dtype=np.int32))\n ).numpy()\n\n # ---- Step 1: Verify MeshNet baseline ----\n log.info(\"=== Step 1: MeshNet baseline ===\")\n seg = Segmentation.load(WORK_DIR / \"checkpoints\" / \"sanity_meshnet\")\n det_model = seg.model_\n pred = predict_volume(det_model, img_path, block_shape)\n cd = per_class_dice(pred, gt_arr, n_classes)\n log.info(\"MeshNet: mean=%.4f, max=%.4f\", cd.mean(), cd.max())\n\n # ---- Step 2: Fixed warm-start ----\n log.info(\"=== Step 2: Fixed warm-start ===\")\n kwyk_factory = get_model(\"kwyk_meshnet\")\n kwyk_model = kwyk_factory(\n n_classes=n_classes,\n filters=96,\n receptive_field=37,\n dropout_type=\"bernoulli\",\n dropout_rate=0.25,\n sigma_init=0.0001,\n )\n meshnet_ckpt = WORK_DIR / \"checkpoints\" / \"sanity_meshnet\" / \"model.pth\"\n fixed_warmstart_kwyk(kwyk_model, meshnet_ckpt)\n\n # Eval immediately\n pred = predict_volume(kwyk_model, img_path, block_shape, mc=False)\n cd = per_class_dice(pred, gt_arr, n_classes)\n log.info(\n \"KWYKMeshNet fixed warm-start (mc=False): mean=%.4f, max=%.4f\",\n cd.mean(),\n cd.max(),\n )\n\n pred = predict_volume(kwyk_model, img_path, block_shape, mc=True)\n cd = per_class_dice(pred, gt_arr, n_classes)\n log.info(\n \"KWYKMeshNet fixed warm-start (mc=True): mean=%.4f, max=%.4f\",\n cd.mean(),\n cd.max(),\n )\n\n # ---- Step 3: Train Bayesian (5 subjects, 20 epochs) ----\n log.info(\"=== Step 3: Train Bayesian with fixed warm-start ===\")\n manifest = WORK_DIR / \"kwyk_sanity_manifest.csv\"\n label_mapping = \"50-class\"\n\n train_pairs = []\n val_pairs = []\n with open(manifest) as f:\n for row in csv.DictReader(f):\n p = (row[\"t1w_path\"], row[\"label_path\"])\n if row[\"split\"] == \"train\":\n train_pairs.append(p)\n elif row[\"split\"] == \"val\":\n val_pairs.append(p)\n\n ds_train = (\n Dataset.from_files(train_pairs, block_shape=block_shape, n_classes=n_classes)\n .batch(32)\n .binarize(label_mapping)\n )\n\n from nobrainer.training import get_device\n\n device = get_device()\n kwyk_model = kwyk_model.to(device)\n optimizer = torch.optim.Adam(kwyk_model.parameters(), lr=0.0001)\n criterion = nn.CrossEntropyLoss()\n\n for epoch in range(20):\n kwyk_model.train()\n epoch_loss = 0.0\n n_batches = 0\n for batch in ds_train.dataloader:\n if isinstance(batch, dict):\n images = batch[\"image\"].to(device)\n labels = batch[\"label\"].to(device)\n else:\n images = batch[0].to(device)\n labels = batch[1].to(device)\n if labels.ndim == images.ndim and labels.shape[1] == 1:\n labels = labels.squeeze(1)\n if labels.dtype in (torch.float32, torch.float64):\n labels = labels.long()\n\n optimizer.zero_grad()\n pred_t = kwyk_model(images, mc=True)\n loss = criterion(pred_t, labels)\n loss.backward()\n optimizer.step()\n epoch_loss += loss.item()\n n_batches += 1\n\n avg_loss = epoch_loss / max(n_batches, 1)\n msg = f\"Epoch {epoch + 1}/20: loss={avg_loss:.4f}\"\n\n # Eval every 5 epochs\n if (epoch + 1) % 5 == 0:\n pred = predict_volume(kwyk_model, img_path, block_shape, mc=False)\n cd = per_class_dice(pred, gt_arr, n_classes)\n msg += f\" dice_det={cd.mean():.4f}/{cd.max():.4f}\"\n\n pred = predict_volume(kwyk_model, img_path, block_shape, mc=True)\n cd_mc = per_class_dice(pred, gt_arr, n_classes)\n msg += f\" dice_mc={cd_mc.mean():.4f}/{cd_mc.max():.4f}\"\n\n log.info(msg)\n\n # ---- Final eval ----\n log.info(\"=== Final evaluation ===\")\n for mc_mode in [False, True]:\n pred = predict_volume(kwyk_model, img_path, block_shape, mc=mc_mode)\n cd = per_class_dice(pred, gt_arr, n_classes)\n mode = \"mc\" if mc_mode else \"det\"\n log.info(\"Final [%s]: mean=%.4f, max=%.4f\", mode, cd.mean(), cd.max())\n\n # Save model\n torch.save(kwyk_model.state_dict(), EXP_DIR / \"kwyk_fixed_warmstart.pth\")\n log.info(\"Model saved to %s\", EXP_DIR / \"kwyk_fixed_warmstart.pth\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "scripts/kwyk_reproduction/experiments/04_20260330_fixed_warmstart/run.sbatch", "content": "#!/bin/bash\n#SBATCH --job-name=exp04-fix\n#SBATCH --partition=mit_preemptable\n#SBATCH --gres=gpu:1\n#SBATCH --cpus-per-task=4\n#SBATCH --mem=32G\n#SBATCH --time=00:30:00\n#SBATCH --output=slurm-%j.out\n#SBATCH --error=slurm-%j.err\nset -euo pipefail\n\ncd /orcd/scratch/orcd/013/satra/kwyk_reproduction\nsource /orcd/data/satra/002/projects/nobrainer/venvs/nobrainer/bin/activate\n\necho \"=== Experiment 04: Fixed warm-start + Bayesian training ===\"\npython experiments/04_20260330_fixed_warmstart/run.py\necho \"=== Done: $(date) ===\"\n" }, { "path": "scripts/kwyk_reproduction/experiments/05_20260330_kwyk_from_scratch/README.md", "content": "# Experiment 05: KWYKMeshNet from scratch (no warm-start)\n\n## Rationale\n\nExperiments 03-04 showed warm-start doesn't transfer well. But the Bayesian\ntraining also shows zero Dice after 20 epochs from scratch. The question:\ncan KWYKMeshNet learn AT ALL with mc=False (deterministic)?\n\nIf not, the VWN architecture + CrossEntropyLoss may have a fundamental issue.\nWe also test: (a) mc=False during training, (b) binary labels for simplicity.\n\n## Plan\n\n1. Train KWYKMeshNet (50-class, mc=True during training, 5 subj, 50 epochs)\n2. Train KWYKMeshNet (50-class, mc=FALSE during training, 5 subj, 50 epochs)\n3. Train KWYKMeshNet (binary, mc=False, 5 subj, 50 epochs)\n4. Compare: does turning off MC during training help? Does binary help?\n" }, { "path": "scripts/kwyk_reproduction/experiments/05_20260330_kwyk_from_scratch/results_summary.md", "content": "# Experiment 05 Results\n\n## Key Finding: mc=False during training is REQUIRED for KWYKMeshNet to learn\n\n| Condition | mc_train | n_classes | Final Loss | Final Dice (det) mean/max |\n|---|---|---|---|---|\n| A | True | 50 | 3.387 | 0.0000/0.0006 |\n| **B** | **False** | **50** | **2.620** | **0.0019/0.0936** |\n| C | True | 2 | 1.011 | 0.0000 |\n| D | False | 2 | 0.910 | 0.0001 |\n\n## Analysis\n\n1. **mc=True kills training**: Conditions A and C (mc=True) both converge to zero\n Dice despite loss decreasing. The local reparameterization noise from FFGConv3d\n prevents stable gradient flow. The loss landscape becomes too noisy.\n\n2. **mc=False allows learning**: Condition B (mc=False, 50-class) achieves 9.4% Dice\n on the best class — comparable to the deterministic MeshNet at similar epoch count.\n The VWN weight normalization itself is fine; the stochastic sampling is the issue.\n\n3. **Binary fails for both**: Possibly an issue with binary evaluation or the model\n having too many parameters for a 2-class problem (3M params for binary).\n\n4. **Loss instability with mc=True**: Condition A shows wild loss swings (1.3 to 5.3)\n because each forward pass samples different weights. mc=False gives stable loss.\n\n## Conclusion\n\nThe Bayesian training should use `mc=False` for the forward pass during gradient\ncomputation, and only enable `mc=True` at inference time for uncertainty estimation.\nThis is the standard approach: train with deterministic weights, use stochastic\ninference. The current code passes `mc=True` during training which prevents learning.\n\n## Recommendation\n\nFix `03_train_bayesian.py` to call `model(images, mc=False)` during training,\nand only use `mc=True` for validation MC Dice evaluation.\n" }, { "path": "scripts/kwyk_reproduction/experiments/05_20260330_kwyk_from_scratch/run.py", "content": "#!/usr/bin/env python\n\"\"\"Train KWYKMeshNet from scratch: mc=True vs mc=False, 50-class vs binary.\"\"\"\n\nfrom __future__ import annotations\n\nimport csv\nfrom pathlib import Path\nimport sys\nimport time\n\nimport nibabel as nib\nimport numpy as np\nimport torch\nimport torch.nn as nn\n\nsys.path.insert(0, str(Path(__file__).parent.parent.parent))\nfrom utils import setup_logging # noqa: E402\n\nlog = setup_logging(__name__)\n\nWORK_DIR = Path(__file__).parent.parent.parent\nEXP_DIR = Path(__file__).parent\n\n\ndef predict_volume(model, img_path, block_shape, mc=False):\n from nobrainer.prediction import _extract_blocks, _pad_to_multiple, _stitch_blocks\n from nobrainer.training import get_device\n\n device = get_device()\n img = nib.load(str(img_path))\n arr = np.asarray(img.dataobj, dtype=np.float32)\n orig_shape = arr.shape[:3]\n padded, pad = _pad_to_multiple(arr, block_shape)\n blocks, grid = _extract_blocks(padded, block_shape)\n model = model.to(device)\n model.eval()\n all_preds = []\n with torch.no_grad():\n for start in range(0, len(blocks), 4):\n chunk = blocks[start : start + 4]\n tensor = torch.from_numpy(chunk[:, None]).to(device)\n out = model(tensor, mc=mc)\n all_preds.append(out.argmax(dim=1, keepdim=True).float().cpu().numpy())\n block_preds = np.concatenate(all_preds, axis=0)\n return _stitch_blocks(block_preds, grid, block_shape, pad, orig_shape, 1)[0].astype(\n np.int32\n )\n\n\ndef per_class_dice(pred, gt, n_classes):\n dices = []\n for c in range(1, n_classes):\n p = pred == c\n g = gt == c\n inter = (p & g).sum()\n total = p.sum() + g.sum()\n dices.append(2.0 * inter / total if total > 0 else 1.0)\n return np.array(dices)\n\n\ndef binary_dice(pred, gt):\n pred = (pred > 0).astype(bool)\n gt = (gt > 0).astype(bool)\n inter = (pred & gt).sum()\n total = pred.sum() + gt.sum()\n return 2.0 * inter / total if total > 0 else 1.0\n\n\ndef train_kwyk(name, n_classes, label_mapping, mc_train, epochs=50):\n from nobrainer.models import get as get_model\n from nobrainer.processing.dataset import Dataset, _load_label_mapping\n from nobrainer.training import get_device\n\n log.info(\n \"=== %s: n_classes=%d, mc_train=%s, epochs=%d ===\",\n name,\n n_classes,\n mc_train,\n epochs,\n )\n block_shape = (32, 32, 32)\n device = get_device()\n\n # Load data\n train_pairs, val_pairs = [], []\n with open(WORK_DIR / \"kwyk_sanity_manifest.csv\") as f:\n for row in csv.DictReader(f):\n p = (row[\"t1w_path\"], row[\"label_path\"])\n if row[\"split\"] == \"train\":\n train_pairs.append(p)\n elif row[\"split\"] == \"test\":\n val_pairs.append(p) # use test as val for sanity\n\n ds = (\n Dataset.from_files(train_pairs, block_shape=block_shape, n_classes=n_classes)\n .batch(32)\n .binarize(label_mapping)\n )\n\n # Remap function for eval\n remap_fn = None\n if label_mapping and label_mapping != \"binary\":\n remap_fn = _load_label_mapping(label_mapping)\n\n # Create model\n model = get_model(\"kwyk_meshnet\")(\n n_classes=n_classes,\n filters=96,\n receptive_field=37,\n dropout_type=\"bernoulli\",\n dropout_rate=0.25,\n sigma_init=0.0001,\n ).to(device)\n\n optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)\n criterion = nn.CrossEntropyLoss()\n\n # Test volume\n img_path, lbl_path = val_pairs[0]\n gt_raw = np.asarray(nib.load(lbl_path).dataobj, dtype=np.int32)\n if remap_fn:\n gt_arr = remap_fn(torch.from_numpy(gt_raw)).numpy()\n else:\n gt_arr = (gt_raw > 0).astype(np.int32)\n\n t0 = time.time()\n for epoch in range(epochs):\n model.train()\n epoch_loss = 0.0\n n_batches = 0\n for batch in ds.dataloader:\n if isinstance(batch, dict):\n images = batch[\"image\"].to(device)\n labels = batch[\"label\"].to(device)\n else:\n images = batch[0].to(device)\n labels = batch[1].to(device)\n if labels.ndim == images.ndim and labels.shape[1] == 1:\n labels = labels.squeeze(1)\n if labels.dtype in (torch.float32, torch.float64):\n labels = labels.long()\n\n optimizer.zero_grad()\n pred = model(images, mc=mc_train)\n loss = criterion(pred, labels)\n loss.backward()\n optimizer.step()\n epoch_loss += loss.item()\n n_batches += 1\n\n avg_loss = epoch_loss / max(n_batches, 1)\n msg = f\" Epoch {epoch + 1}/{epochs}: loss={avg_loss:.4f}\"\n\n if (epoch + 1) % 10 == 0 or epoch == 0:\n pred_vol = predict_volume(model, img_path, block_shape, mc=False)\n if n_classes == 2:\n d = binary_dice(pred_vol, gt_arr)\n msg += f\" dice_det={d:.4f}\"\n else:\n cd = per_class_dice(pred_vol, gt_arr, n_classes)\n msg += f\" dice_det={cd.mean():.4f}/{cd.max():.4f}\"\n\n log.info(msg)\n\n elapsed = time.time() - t0\n log.info(\" Completed in %.1fs\", elapsed)\n\n # Final eval\n pred_vol = predict_volume(model, img_path, block_shape, mc=False)\n if n_classes == 2:\n d = binary_dice(pred_vol, gt_arr)\n log.info(\" FINAL [det]: dice=%.4f\", d)\n else:\n cd = per_class_dice(pred_vol, gt_arr, n_classes)\n log.info(\" FINAL [det]: mean=%.4f, max=%.4f\", cd.mean(), cd.max())\n\n if mc_train:\n pred_vol = predict_volume(model, img_path, block_shape, mc=True)\n if n_classes == 2:\n d = binary_dice(pred_vol, gt_arr)\n log.info(\" FINAL [mc]: dice=%.4f\", d)\n else:\n cd = per_class_dice(pred_vol, gt_arr, n_classes)\n log.info(\" FINAL [mc]: mean=%.4f, max=%.4f\", cd.mean(), cd.max())\n\n return model\n\n\ndef main():\n # A: 50-class, mc=True during training (current default)\n train_kwyk(\"A_50class_mcTrue\", 50, \"50-class\", mc_train=True, epochs=50)\n\n # B: 50-class, mc=False during training (deterministic forward)\n train_kwyk(\"B_50class_mcFalse\", 50, \"50-class\", mc_train=False, epochs=50)\n\n # C: Binary, mc=True during training\n train_kwyk(\"C_binary_mcTrue\", 2, \"binary\", mc_train=True, epochs=50)\n\n # D: Binary, mc=False during training\n train_kwyk(\"D_binary_mcFalse\", 2, \"binary\", mc_train=False, epochs=50)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "scripts/kwyk_reproduction/experiments/05_20260330_kwyk_from_scratch/run.sbatch", "content": "#!/bin/bash\n#SBATCH --job-name=exp05-scratch\n#SBATCH --partition=mit_preemptable\n#SBATCH --gres=gpu:1\n#SBATCH --cpus-per-task=4\n#SBATCH --mem=32G\n#SBATCH --time=01:30:00\n#SBATCH --output=slurm-%j.out\n#SBATCH --error=slurm-%j.err\nset -euo pipefail\n\ncd /orcd/scratch/orcd/013/satra/kwyk_reproduction\nsource /orcd/data/satra/002/projects/nobrainer/venvs/nobrainer/bin/activate\n\necho \"=== Experiment 05: KWYKMeshNet from scratch ===\"\npython experiments/05_20260330_kwyk_from_scratch/run.py\necho \"=== Done: $(date) ===\"\n" }, { "path": "scripts/kwyk_reproduction/experiments/06_20260331_fullvol_augment/README.md", "content": "# Experiment 06: Full-volume (256³) training with augmentation\n\n## Rationale\n\nCurrent training uses 32³ patches — the model never sees global context.\nTraining on full 256³ volumes should improve segmentation quality,\nespecially for large structures. Combined with augmentation (affine + flip\n+ noise) for regularization.\n\n## Plan\n\n1. Use 128³ blocks on L40S (batch_size=4 fits in 47GB)\n OR request H200/A100 for full 256³ (batch_size=1 per GPU, 2 GPUs)\n2. Standard augmentation profile (affine rotation/scale, flips, Gaussian noise)\n3. MeshNet first (deterministic baseline), then bwn_multi\n4. 20 epochs on 500 subjects\n5. Compare Dice vs 32³ patch training\n\n## GPU Options\n\n| GPU | Memory | 256³ batch=1 | 128³ batch=4 |\n|-----|--------|-------------|-------------|\n| L40S (47GB) | 47GB | OOM (90GB) | OK (11GB) |\n| A100 (80GB) | 80GB | Tight | OK |\n| H200 (141GB) | 141GB | OK | OK |\n| 2× L40S | 94GB | OK | OK |\n" }, { "path": "scripts/kwyk_reproduction/experiments/06_20260331_fullvol_augment/config_256.yaml", "content": "# Full 256³ volume training — requires H200 or multi-GPU\n\nfilters: 96\nreceptive_field: 37\ndropout_rate: 0.25\nsigma_init: 0.0001\n\nblock_shape: [256, 256, 256]\nlr: 0.0001\nbatch_size: 1\nn_classes: 50\nlabel_mapping: 50-class\npatches_per_volume: 1 # whole volume = 1 patch\n\npretrain_epochs: 20\nbayesian_epochs: 20\nval_dice_freq: 5\n\nn_samples: 10\naugmentation_profile: standard\n\ngradient_checkpointing: true\nzarr_store: /orcd/scratch/orcd/013/satra/kwyk_reproduction/data/kwyk_500.zarr\n\ndatasets:\n - kwyk\nsplit: [80, 10, 10]\n\nvariants:\n bwn_multi:\n model: kwyk_meshnet\n dropout_type: bernoulli\n mc_at_inference: true\n" }, { "path": "scripts/kwyk_reproduction/experiments/06_20260331_fullvol_augment/config_256_mp.yaml", "content": "# Full 256³ volume — model parallel across 2 GPUs\n\nfilters: 96\nreceptive_field: 37\ndropout_rate: 0.25\nsigma_init: 0.0001\n\nblock_shape: [256, 256, 256]\nlr: 0.0001\nbatch_size: 1\nn_classes: 50\nlabel_mapping: 50-class\npatches_per_volume: 1\n\npretrain_epochs: 20\nbayesian_epochs: 20\nval_dice_freq: 5\n\nn_samples: 10\nmodel_parallel: true\n\nzarr_store: /orcd/scratch/orcd/013/satra/kwyk_reproduction/data/kwyk_500.zarr\n\ndatasets:\n - kwyk\nsplit: [80, 10, 10]\n\nvariants:\n bwn_multi:\n model: kwyk_meshnet\n dropout_type: bernoulli\n mc_at_inference: true\n" }, { "path": "scripts/kwyk_reproduction/experiments/06_20260331_fullvol_augment/config_fullvol.yaml", "content": "# Full-volume training with augmentation — 50-class parcellation\n# Uses 128³ blocks (fits on L40S) or 256³ on larger GPUs\n\nfilters: 96\nreceptive_field: 37\ndropout_rate: 0.25\nsigma_init: 0.0001\n\nblock_shape: [128, 128, 128]\nlr: 0.0001\nbatch_size: 4\nn_classes: 50\nlabel_mapping: 50-class\npatches_per_volume: 8 # fewer patches needed with large blocks\n\npretrain_epochs: 20\nbayesian_epochs: 20\nval_dice_freq: 5\n\nn_samples: 10\naugmentation_profile: standard\n\ndatasets:\n - kwyk\nsplit: [80, 10, 10]\n\nvariants:\n bwn_multi:\n model: kwyk_meshnet\n dropout_type: bernoulli\n mc_at_inference: true\n" }, { "path": "scripts/kwyk_reproduction/experiments/06_20260331_fullvol_augment/run_128.sbatch", "content": "#!/bin/bash\n#SBATCH --job-name=exp06-128\n#SBATCH --partition=mit_preemptable\n#SBATCH --gres=gpu:1\n#SBATCH --cpus-per-task=8\n#SBATCH --mem=64G\n#SBATCH --time=06:00:00\n#SBATCH --output=slurm-128-%j.out\n#SBATCH --error=slurm-128-%j.err\n#\n# Experiment 06a: 128³ blocks + augmentation on L40S\n#\nset -euo pipefail\n\nWORK_DIR=\"/orcd/scratch/orcd/013/satra/kwyk_reproduction\"\nEXP_DIR=\"$WORK_DIR/experiments/06_20260331_fullvol_augment\"\nVENV_DIR=\"/orcd/data/satra/002/projects/nobrainer/venvs/nobrainer\"\nCONFIG=\"$EXP_DIR/config_fullvol.yaml\"\nMANIFEST=\"$WORK_DIR/kwyk_manifest_500.csv\"\nZARR_STORE=\"$WORK_DIR/data/kwyk_500.zarr\"\n\ncd \"$WORK_DIR\"\nsource \"${VENV_DIR}/bin/activate\"\n\necho \"=== Experiment 06a: 128³ + augmentation ===\"\necho \"Node: $(hostname)\"\npython -c \"import torch; print(f'GPU: {torch.cuda.get_device_name(0)}, {torch.cuda.get_device_properties(0).total_memory/1e9:.0f}GB')\"\n\n# MeshNet 128³ (20 epochs)\necho \"=== Step 1: MeshNet 128³ ===\"\npython 02_train_meshnet.py --manifest \"$MANIFEST\" --config \"$CONFIG\" \\\n --output-dir \"$EXP_DIR/checkpoints/meshnet_128\" --epochs 20\n\n# bwn_multi 128³ (20 epochs, warm-start)\necho \"=== Step 2: bwn_multi 128³ ===\"\npython 03_train_bayesian.py --manifest \"$MANIFEST\" --config \"$CONFIG\" \\\n --variant bwn_multi --warmstart \"$EXP_DIR/checkpoints/meshnet_128\" \\\n --output-dir \"$EXP_DIR/checkpoints/bwn_multi_128\" --epochs 20\n\n# Evaluate\necho \"=== Step 3: Evaluate ===\"\nfor v in meshnet_128 bwn_multi_128; do\n if [ -f \"$EXP_DIR/checkpoints/$v/model.pth\" ]; then\n echo \"--- $v (deterministic) ---\"\n python 04_evaluate.py --model \"$EXP_DIR/checkpoints/$v\" \\\n --manifest \"$MANIFEST\" --config \"$CONFIG\" \\\n --split test --n-samples 0 \\\n --output-dir \"$EXP_DIR/results/${v}_det\"\n fi\ndone\n\necho \"=== Done: $(date) ===\"\n" }, { "path": "scripts/kwyk_reproduction/experiments/06_20260331_fullvol_augment/run_256.sbatch", "content": "#!/bin/bash\n#SBATCH --job-name=exp06-256\n#SBATCH --partition=mit_preemptable\n#SBATCH --gres=gpu:h200:1\n#SBATCH --cpus-per-task=8\n#SBATCH --mem=128G\n#SBATCH --time=06:00:00\n#SBATCH --output=slurm-256-%j.out\n#SBATCH --error=slurm-256-%j.err\n#\n# Experiment 06b: Full 256³ volume + augmentation on H200\n#\nset -euo pipefail\n\nWORK_DIR=\"/orcd/scratch/orcd/013/satra/kwyk_reproduction\"\nEXP_DIR=\"$WORK_DIR/experiments/06_20260331_fullvol_augment\"\nVENV_DIR=\"/orcd/data/satra/002/projects/nobrainer/venvs/nobrainer\"\nCONFIG=\"$EXP_DIR/config_256.yaml\"\nMANIFEST=\"$WORK_DIR/kwyk_manifest_500.csv\"\n\ncd \"$WORK_DIR\"\nsource \"${VENV_DIR}/bin/activate\"\n\necho \"=== Experiment 06b: Full 256³ + augmentation ===\"\necho \"Node: $(hostname)\"\npython -c \"import torch; print(f'GPU: {torch.cuda.get_device_name(0)}, {torch.cuda.get_device_properties(0).total_memory/1e9:.0f}GB')\"\n\n# MeshNet 256³ (20 epochs)\necho \"=== Step 1: MeshNet 256³ ===\"\npython 02_train_meshnet.py --manifest \"$MANIFEST\" --config \"$CONFIG\" \\\n --output-dir \"$EXP_DIR/checkpoints/meshnet_256\" --epochs 20\n\n# bwn_multi 256³ (20 epochs, warm-start)\necho \"=== Step 2: bwn_multi 256³ ===\"\npython 03_train_bayesian.py --manifest \"$MANIFEST\" --config \"$CONFIG\" \\\n --variant bwn_multi --warmstart \"$EXP_DIR/checkpoints/meshnet_256\" \\\n --output-dir \"$EXP_DIR/checkpoints/bwn_multi_256\" --epochs 20\n\n# Evaluate\necho \"=== Step 3: Evaluate ===\"\nfor v in meshnet_256 bwn_multi_256; do\n if [ -f \"$EXP_DIR/checkpoints/$v/model.pth\" ]; then\n echo \"--- $v (deterministic) ---\"\n python 04_evaluate.py --model \"$EXP_DIR/checkpoints/$v\" \\\n --manifest \"$MANIFEST\" --config \"$CONFIG\" \\\n --split test --n-samples 0 \\\n --output-dir \"$EXP_DIR/results/${v}_det\"\n fi\ndone\n\necho \"=== Done: $(date) ===\"\n" }, { "path": "scripts/kwyk_reproduction/experiments/06_20260331_fullvol_augment/run_256_a100.sbatch", "content": "#!/bin/bash\n#SBATCH --job-name=exp06-256\n#SBATCH --partition=mit_preemptable\n#SBATCH --gres=gpu:h200:1\n#SBATCH --cpus-per-task=8\n#SBATCH --mem=128G\n#SBATCH --time=12:00:00\n#SBATCH --output=slurm-256-%j.out\n#SBATCH --error=slurm-256-%j.err\n#\n# Experiment 06b: Full 256³ volume + augmentation on H200 (141GB)\n# Single H200 fits 256³ with batch=1 (~90GB forward+backward)\n#\nset -euo pipefail\n\nWORK_DIR=\"/orcd/scratch/orcd/013/satra/kwyk_reproduction\"\nEXP_DIR=\"$WORK_DIR/experiments/06_20260331_fullvol_augment\"\nVENV_DIR=\"/orcd/data/satra/002/projects/nobrainer/venvs/nobrainer\"\nCONFIG=\"$EXP_DIR/config_256.yaml\"\nMANIFEST=\"$WORK_DIR/kwyk_manifest_500.csv\"\n\ncd \"$WORK_DIR\"\nsource \"${VENV_DIR}/bin/activate\"\n\necho \"=== Experiment 06b: Full 256³ + augmentation (2× A100) ===\"\necho \"Node: $(hostname)\"\npython -c \"\nimport torch\nn = torch.cuda.device_count()\nfor i in range(n):\n print(f'GPU {i}: {torch.cuda.get_device_name(i)}, {torch.cuda.get_device_properties(i).total_memory/1e9:.0f}GB')\n\"\n\n# MeshNet 256³ (20 epochs)\necho \"=== Step 1: MeshNet 256³ ===\"\npython 02_train_meshnet.py --manifest \"$MANIFEST\" --config \"$CONFIG\" \\\n --output-dir \"$EXP_DIR/checkpoints/meshnet_256\" --epochs 20\n\n# bwn_multi 256³ (20 epochs, warm-start)\necho \"=== Step 2: bwn_multi 256³ ===\"\npython 03_train_bayesian.py --manifest \"$MANIFEST\" --config \"$CONFIG\" \\\n --variant bwn_multi --warmstart \"$EXP_DIR/checkpoints/meshnet_256\" \\\n --output-dir \"$EXP_DIR/checkpoints/bwn_multi_256\" --epochs 20\n\n# Evaluate\necho \"=== Step 3: Evaluate ===\"\nfor v in meshnet_256 bwn_multi_256; do\n if [ -f \"$EXP_DIR/checkpoints/$v/model.pth\" ]; then\n echo \"--- $v (deterministic) ---\"\n python 04_evaluate.py --model \"$EXP_DIR/checkpoints/$v\" \\\n --manifest \"$MANIFEST\" --config \"$CONFIG\" \\\n --split test --n-samples 0 \\\n --output-dir \"$EXP_DIR/results/${v}_det\"\n fi\ndone\n\necho \"=== Done: $(date) ===\"\n" }, { "path": "scripts/kwyk_reproduction/experiments/06_20260331_fullvol_augment/run_256_gradckpt.sbatch", "content": "#!/bin/bash\n#SBATCH --job-name=exp06-gc\n#SBATCH --partition=pi_satra\n#SBATCH --gres=gpu:a100:1\n#SBATCH --cpus-per-task=8\n#SBATCH --mem=128G\n#SBATCH --time=12:00:00\n#SBATCH --output=slurm-256gc-%j.out\n#SBATCH --error=slurm-256gc-%j.err\n#\n# Experiment 06c: Full 256³ + gradient checkpointing on single A100 (80GB)\n# Gradient checkpointing halves activation memory: ~90GB → ~45GB\n#\nset -euo pipefail\n\nWORK_DIR=\"/orcd/scratch/orcd/013/satra/kwyk_reproduction\"\nEXP_DIR=\"$WORK_DIR/experiments/06_20260331_fullvol_augment\"\nVENV_DIR=\"/orcd/data/satra/002/projects/nobrainer/venvs/nobrainer\"\nCONFIG=\"$EXP_DIR/config_256.yaml\"\nMANIFEST=\"$WORK_DIR/kwyk_manifest_500.csv\"\n\ncd \"$WORK_DIR\"\nsource \"${VENV_DIR}/bin/activate\"\n\necho \"=== Experiment 06c: 256³ + gradient checkpointing (A100) ===\"\necho \"Node: $(hostname)\"\npython -c \"\nimport torch\nprint(f'GPU: {torch.cuda.get_device_name(0)}, {torch.cuda.get_device_properties(0).total_memory/1e9:.0f}GB')\n\"\n\n# MeshNet 256³ with gradient checkpointing\necho \"=== Step 1: MeshNet 256³ (gradient checkpointing) ===\"\npython -c \"\nimport sys; sys.path.insert(0, '.')\nfrom nobrainer.processing.dataset import Dataset\nfrom nobrainer.processing.segmentation import Segmentation\nfrom utils import load_config, setup_logging\nimport torch\n\nlog = setup_logging('exp06c')\nconfig = load_config('$CONFIG')\nblock_shape = tuple(config['block_shape'])\nn_classes = config['n_classes']\nlabel_mapping = config.get('label_mapping', 'binary')\n\n# Load data\nimport csv\ntrain_pairs, val_pairs = [], []\nwith open('$MANIFEST') as f:\n for row in csv.DictReader(f):\n p = (row['t1w_path'], row['label_path'])\n if row['split'] == 'train': train_pairs.append(p)\n elif row['split'] == 'val': val_pairs.append(p)\n\nds_train = Dataset.from_files(train_pairs, block_shape=block_shape, n_classes=n_classes).batch(1).binarize(label_mapping).streaming(patches_per_volume=1)\nds_val = Dataset.from_files(val_pairs[:5], block_shape=block_shape, n_classes=n_classes).batch(1).binarize(label_mapping).streaming(patches_per_volume=1)\n\nseg = Segmentation(base_model='meshnet', model_args={'n_classes': n_classes, 'filters': 96, 'receptive_field': 37, 'dropout_rate': 0.25}, checkpoint_filepath='$EXP_DIR/checkpoints/meshnet_256gc')\n\ndef _log(epoch, logs, model):\n msg = f'Epoch {epoch+1}/20: loss={logs[\\\"loss\\\"]:.4f}'\n if 'val_loss' in logs: msg += f' val_loss={logs[\\\"val_loss\\\"]:.4f} val_acc={logs[\\\"val_acc\\\"]:.4f}'\n log.info(msg)\n\nseg.fit(dataset_train=ds_train, dataset_validate=ds_val, epochs=20, optimizer=torch.optim.Adam, opt_args={'lr': 0.0001}, callbacks=[_log], checkpoint_freq=5, gradient_checkpointing=True)\nlog.info('Done. Result: %s', seg._training_result)\n\"\n\necho \"=== Done: $(date) ===\"\n" }, { "path": "scripts/kwyk_reproduction/experiments/06_20260331_fullvol_augment/run_256_mp.sbatch", "content": "#!/bin/bash\n#SBATCH --job-name=exp06-mp\n#SBATCH --partition=mit_preemptable\n#SBATCH --gres=gpu:2\n#SBATCH --cpus-per-task=8\n#SBATCH --mem=64G\n#SBATCH --time=06:00:00\n#SBATCH --output=slurm-256mp-%j.out\n#SBATCH --error=slurm-256mp-%j.err\n#\n# Experiment 06d: Full 256³ + model parallelism on 2× L40S\n# Layers split across GPUs — each GPU holds ~half the model\n#\nset -euo pipefail\n\nWORK_DIR=\"/orcd/scratch/orcd/013/satra/kwyk_reproduction\"\nEXP_DIR=\"$WORK_DIR/experiments/06_20260331_fullvol_augment\"\nVENV_DIR=\"/orcd/data/satra/002/projects/nobrainer/venvs/nobrainer\"\nCONFIG=\"$EXP_DIR/config_256_mp.yaml\"\nMANIFEST=\"$WORK_DIR/kwyk_manifest_500.csv\"\nZARR_STORE=\"$WORK_DIR/data/kwyk_500.zarr\"\n\ncd \"$WORK_DIR\"\nsource \"${VENV_DIR}/bin/activate\"\n\necho \"=== Experiment 06d: 256³ + model parallel (2× L40S) ===\"\necho \"Node: $(hostname)\"\npython -c \"\nimport torch\nfor i in range(torch.cuda.device_count()):\n print(f'GPU {i}: {torch.cuda.get_device_name(i)}, {torch.cuda.get_device_properties(i).total_memory/1e9:.0f}GB')\n\"\n\n# MeshNet 256³ with model parallelism\necho \"=== Step 1: MeshNet 256³ (model parallel) ===\"\npython 02_train_meshnet.py --manifest \"$MANIFEST\" --config \"$CONFIG\" \\\n --output-dir \"$EXP_DIR/checkpoints/meshnet_256mp\" --epochs 20\n\n# bwn_multi 256³ with model parallelism\necho \"=== Step 2: bwn_multi 256³ (model parallel) ===\"\npython 03_train_bayesian.py --manifest \"$MANIFEST\" --config \"$CONFIG\" \\\n --variant bwn_multi --warmstart \"$EXP_DIR/checkpoints/meshnet_256mp\" \\\n --output-dir \"$EXP_DIR/checkpoints/bwn_multi_256mp\" --epochs 20\n\n# Evaluate\necho \"=== Step 3: Evaluate ===\"\nfor v in meshnet_256mp bwn_multi_256mp; do\n if [ -f \"$EXP_DIR/checkpoints/$v/model.pth\" ]; then\n echo \"--- $v (deterministic) ---\"\n python 04_evaluate.py --model \"$EXP_DIR/checkpoints/$v\" \\\n --manifest \"$MANIFEST\" --config \"$CONFIG\" \\\n --split test --n-samples 0 \\\n --output-dir \"$EXP_DIR/results/${v}_det\"\n fi\ndone\n\necho \"=== Done: $(date) ===\"\n" }, { "path": "scripts/kwyk_reproduction/experiments/07_20260401_ddp_128/config.yaml", "content": "# DDP test: 128³ patches on 2× L40S\n\nfilters: 96\nreceptive_field: 37\ndropout_rate: 0.25\n\nblock_shape: [128, 128, 128]\nlr: 0.0001\nbatch_size: 4\nn_classes: 50\nlabel_mapping: 50-class\npatches_per_volume: 8\nval_dice_freq: 1\n\nzarr_store: /orcd/scratch/orcd/013/satra/kwyk_reproduction/data/kwyk_500.zarr\n" }, { "path": "scripts/kwyk_reproduction/experiments/07_20260401_ddp_128/run.sbatch", "content": "#!/bin/bash\n#SBATCH --job-name=exp07-ddp\n#SBATCH --partition=mit_preemptable\n#SBATCH --gres=gpu:2\n#SBATCH --cpus-per-task=8\n#SBATCH --mem=64G\n#SBATCH --time=01:00:00\n#SBATCH --output=slurm-%j.out\n#SBATCH --error=slurm-%j.err\n#\n# Experiment 07: DDP test on 2× L40S with 128³ patches, 1 epoch\n#\nset -euo pipefail\n\nWORK_DIR=\"/orcd/scratch/orcd/013/satra/kwyk_reproduction\"\nEXP_DIR=\"$WORK_DIR/experiments/07_20260401_ddp_128\"\nVENV_DIR=\"/orcd/data/satra/002/projects/nobrainer/venvs/nobrainer\"\nMANIFEST=\"$WORK_DIR/kwyk_manifest_500.csv\"\nZARR_STORE=\"$WORK_DIR/data/kwyk_500.zarr\"\n\ncd \"$WORK_DIR\"\nsource \"${VENV_DIR}/bin/activate\"\n\necho \"=== Experiment 07: DDP 2× L40S, 128³, 1 epoch ===\"\necho \"Node: $(hostname)\"\npython -c \"\nimport torch\nfor i in range(torch.cuda.device_count()):\n print(f'GPU {i}: {torch.cuda.get_device_name(i)}, {torch.cuda.get_device_properties(i).total_memory/1e9:.0f}GB')\n\"\n\npython 02_train_meshnet.py \\\n --manifest \"$MANIFEST\" \\\n --config \"$EXP_DIR/config.yaml\" \\\n --output-dir \"$EXP_DIR/checkpoints/meshnet\" \\\n --epochs 1\n\necho \"=== Done: $(date) ===\"\n" }, { "path": "scripts/kwyk_reproduction/experiments/08_20260401_ddp_128_full/config.yaml", "content": "# DDP 128³ on full dataset (11,479 subjects)\n\nfilters: 96\nreceptive_field: 37\ndropout_rate: 0.25\n\nblock_shape: [128, 128, 128]\nlr: 0.0001\nbatch_size: 4\nn_classes: 50\nlabel_mapping: 50-class\npatches_per_volume: 1\nval_dice_freq: 1\n\nzarr_store: /orcd/scratch/orcd/013/satra/kwyk_reproduction/data/kwyk_full.zarr\n" }, { "path": "scripts/kwyk_reproduction/experiments/08_20260401_ddp_128_full/run.sbatch", "content": "#!/bin/bash\n#SBATCH --job-name=exp08-full\n#SBATCH --partition=mit_preemptable\n#SBATCH --gres=gpu:2\n#SBATCH --cpus-per-task=16\n#SBATCH --mem=128G\n#SBATCH --time=24:00:00\n#SBATCH --output=slurm-%j.out\n#SBATCH --error=slurm-%j.err\n#SBATCH --requeue\n#SBATCH --signal=USR1@120\n#\n# DDP 128³ on full KWYK dataset (11,479 subjects, 2× L40S)\n# Checkpoints every epoch for resume after preemption\n#\nset -euo pipefail\n\nWORK_DIR=\"/orcd/scratch/orcd/013/satra/kwyk_reproduction\"\nEXP_DIR=\"$WORK_DIR/experiments/08_20260401_ddp_128_full\"\nVENV_DIR=\"/orcd/data/satra/002/projects/nobrainer/venvs/nobrainer\"\nCONFIG=\"$EXP_DIR/config.yaml\"\nMANIFEST=\"$WORK_DIR/kwyk_manifest_full.csv\"\n\ncd \"$WORK_DIR\"\nsource \"${VENV_DIR}/bin/activate\"\n\necho \"=== Experiment 08: DDP 128³ full dataset ===\"\necho \"Node: $(hostname)\"\necho \"Job ID: ${SLURM_JOB_ID:-local}\"\npython -c \"\nimport torch\nfor i in range(torch.cuda.device_count()):\n print(f'GPU {i}: {torch.cuda.get_device_name(i)}, {torch.cuda.get_device_properties(i).total_memory/1e9:.0f}GB')\n\"\n\necho \"=== Training MeshNet 128³ (20 epochs, full dataset) ===\"\npython 02_train_meshnet.py \\\n --manifest \"$MANIFEST\" \\\n --config \"$CONFIG\" \\\n --output-dir \"$EXP_DIR/checkpoints/meshnet_128\" \\\n --epochs 20\n\necho \"=== Done: $(date) ===\"\n" }, { "path": "scripts/kwyk_reproduction/experiments/task-planner.md", "content": "# Bayesian Model Learning Experiments — Task Planner\n\n**Session dates:** 2026-03-30 to 2026-03-31\n\n## Root Causes Found\n\n### 1. Warm-start key ordering bug (Exp 03)\n`sorted(state.items())` puts `classifier.weight` before `encoder.X`.\n**Fixed:** filter classifier, transfer separately. Branch: `fix/warmstart-key-ordering`.\n\n### 2. Single mc= flag (Exp 05 + original TF code review)\nOriginal TF bwn trains with `is_mc_v=False` (deterministic VWN) + `is_mc_b=True`\n(bernoulli dropout ON). PyTorch had one `mc=` flag controlling both.\n**Fixed:** `mc_vwn` and `mc_dropout` independent flags. Branch: `fix/kwyk-decouple-mc-flags`.\n\n### 3. Dropout ordering mismatch\nOriginal TF: conv → dropout → relu. PyTorch had: conv → relu → dropout.\n**Fixed** in same branch.\n\n### 4. Data pipeline bottleneck\nNIfTI streaming (20K reads/epoch) too slow. Zarr3 with sharding (1 file per array,\n32³ chunk-aligned reads) eliminates this.\n**Fixed:** sharded Zarr3 conversion + PatchDataset zarr:// support.\n\n### 5. auto_batch_size profiling with wrong mode\nProfiled with mc=True (both VWN+dropout) but training uses mc_vwn=False.\n**Fixed:** forward_kwargs parameter in auto_batch_size.\n\n## Additional Discrepancies Found (from original TF code review)\n\n| Aspect | Original TF | Current PyTorch |\n|---|---|---|\n| Framework | TensorFlow 1.12 | PyTorch 2.11 |\n| Dropout rate (bwn) | keep_prob=0.5 (50%) | dropout_rate=0.25 |\n| Dropout type | tf.nn.dropout (element-wise, no rescale) | nn.Dropout3d (spatial, rescales) |\n| Loss regularization | L2 weight decay: sum(mu²)/(2*N) | ELBOLoss with KL=0 (just CE) |\n| sigma_prior (bwn) | 1.0 | 0.1 |\n| Classifier layer | VWN conv3d | Standard nn.Conv3d |\n| prior_path support | Yes (load from trained model) | Missing |\n| Concrete dropout p_prior | 0.5 | 0.9 |\n| Subjects | ~10,000 | 500 (current) |\n\n## Experiment Log\n\n| # | Name | Status | Key Finding |\n|---|---|---|---|\n| 01 | Eval det mode | DONE | Zero Dice both modes — weights empty |\n| 02 | Binary Bayesian | DONE | Zero Dice — same issue |\n| 03 | Warm-start diagnostic | DONE | **BUG: sorted key ordering** |\n| 04 | Fixed warm-start | DONE | Transfer improved but training destroys signal |\n| 05 | From scratch | DONE | **mc=False trains, mc=True doesn't** |\n| 06 | Original TF code review | DONE | **is_mc_v=False, is_mc_b=True in original** |\n\n## Current Pipeline (running)\n\n- Job 11222646: Zarr3 conversion (500 subjects, sharded)\n- Jobs 11222647-49: Bayesian training (mc_vwn=False, mc_dropout=True, streaming Zarr)\n- Job 11222650: Evaluation (det + MC)\n\n## Next Steps\n\n1. Match dropout rate to original (0.5 not 0.25)\n2. Add L2 weight decay to match original loss (not KL)\n3. Scale to more subjects / epochs\n4. Implement prior_path for multi-stage training\n" }, { "path": "scripts/kwyk_reproduction/label_mappings/115-class-mapping.csv", "content": "original,new,label\n0,0,Unknown\n2,1,Left-Cerebral-White-Matter\n3,2,Left-Cerebral-Cortex\n4,3,Left-Lateral-Ventricle\n5,4,Left-Inf-Lat-Vent\n7,5,Left-Cerebellum-White-Matter\n8,6,Left-Cerebellum-Cortex\n10,7,Left-Thalamus-Proper\n11,8,Left-Caudate\n12,9,Left-Putamen\n13,10,Left-Pallidum\n14,11,3rd-Ventricle\n15,12,4th-Ventricle\n16,13,Brain-Stem\n17,14,Left-Hippocampus\n18,15,Left-Amygdala\n24,16,CSF\n26,17,Left-Accumbens-area\n28,18,Left-VentralDC\n30,19,Left-vessel\n31,20,Left-choroid-plexus\n41,21,Right-Cerebral-White-Matter\n42,22,Right-Cerebral-Cortex\n43,23,Right-Lateral-Ventricle\n44,24,Right-Inf-Lat-Vent\n46,25,Right-Cerebellum-White-Matter\n47,26,Right-Cerebellum-Cortex\n49,27,Right-Thalamus-Proper\n50,28,Right-Caudate\n51,29,Right-Putamen\n52,30,Right-Pallidum\n53,31,Right-Hippocampus\n54,32,Right-Amygdala\n58,33,Right-Accumbens-area\n60,34,Right-VentralDC\n62,35,Right-vessel\n63,36,Right-choroid-plexus\n72,37,5th-Ventricle\n77,38,WM-hypointensities\n85,39,Optic-Chiasm\n251,40,CC_Posterior\n252,41,CC_Mid_Posterior\n253,42,CC_Central\n254,43,CC_Mid_Anterior\n255,44,CC_Anterior\n1000,45,ctx-lh-unknown\n1001,46,ctx-lh-bankssts\n1002,47,ctx-lh-caudalanteriorcingulate\n1003,48,ctx-lh-caudalmiddlefrontal\n1005,49,ctx-lh-cuneus\n1006,50,ctx-lh-entorhinal\n1007,51,ctx-lh-fusiform\n1008,52,ctx-lh-inferiorparietal\n1009,53,ctx-lh-inferiortemporal\n1010,54,ctx-lh-isthmuscingulate\n1011,55,ctx-lh-lateraloccipital\n1012,56,ctx-lh-lateralorbitofrontal\n1013,57,ctx-lh-lingual\n1014,58,ctx-lh-medialorbitofrontal\n1015,59,ctx-lh-middletemporal\n1016,60,ctx-lh-parahippocampal\n1017,61,ctx-lh-paracentral\n1018,62,ctx-lh-parsopercularis\n1019,63,ctx-lh-parsorbitalis\n1020,64,ctx-lh-parstriangularis\n1021,65,ctx-lh-pericalcarine\n1022,66,ctx-lh-postcentral\n1023,67,ctx-lh-posteriorcingulate\n1024,68,ctx-lh-precentral\n1025,69,ctx-lh-precuneus\n1026,70,ctx-lh-rostralanteriorcingulate\n1027,71,ctx-lh-rostralmiddlefrontal\n1028,72,ctx-lh-superiorfrontal\n1029,73,ctx-lh-superiorparietal\n1030,74,ctx-lh-superiortemporal\n1031,75,ctx-lh-supramarginal\n1032,76,ctx-lh-frontalpole\n1033,77,ctx-lh-temporalpole\n1034,78,ctx-lh-transversetemporal\n1035,79,ctx-lh-insula\n2000,80,ctx-rh-unknown\n2001,81,ctx-rh-bankssts\n2002,82,ctx-rh-caudalanteriorcingulate\n2003,83,ctx-rh-caudalmiddlefrontal\n2005,84,ctx-rh-cuneus\n2006,85,ctx-rh-entorhinal\n2007,86,ctx-rh-fusiform\n2008,87,ctx-rh-inferiorparietal\n2009,88,ctx-rh-inferiortemporal\n2010,89,ctx-rh-isthmuscingulate\n2011,90,ctx-rh-lateraloccipital\n2012,91,ctx-rh-lateralorbitofrontal\n2013,92,ctx-rh-lingual\n2014,93,ctx-rh-medialorbitofrontal\n2015,94,ctx-rh-middletemporal\n2016,95,ctx-rh-parahippocampal\n2017,96,ctx-rh-paracentral\n2018,97,ctx-rh-parsopercularis\n2019,98,ctx-rh-parsorbitalis\n2020,99,ctx-rh-parstriangularis\n2021,100,ctx-rh-pericalcarine\n2022,101,ctx-rh-postcentral\n2023,102,ctx-rh-posteriorcingulate\n2024,103,ctx-rh-precentral\n2025,104,ctx-rh-precuneus\n2026,105,ctx-rh-rostralanteriorcingulate\n2027,106,ctx-rh-rostralmiddlefrontal\n2028,107,ctx-rh-superiorfrontal\n2029,108,ctx-rh-superiorparietal\n2030,109,ctx-rh-superiortemporal\n2031,110,ctx-rh-supramarginal\n2032,111,ctx-rh-frontalpole\n2033,112,ctx-rh-temporalpole\n2034,113,ctx-rh-transversetemporal\n2035,114,ctx-rh-insula\n" }, { "path": "scripts/kwyk_reproduction/label_mappings/50-class-mapping.csv", "content": ",original,new,label\n0,0,0,Unknown\n1,2,1,Left-Cerebral-White-Matter\n2,4,2,Left-Lateral-Ventricle\n3,5,2,Left-Inf-Lat-Vent\n4,7,3,Left-Cerebellum-White-Matter\n5,8,4,Left-Cerebellum-Cortex\n6,10,5,Left-Thalamus-Proper\n7,11,6,Left-Caudate\n8,12,7,Left-Putamen\n9,13,8,Left-Pallidum\n10,14,2,3rd-Ventricle\n11,15,2,4th-Ventricle\n12,16,9,Brain-Stem\n13,17,10,Left-Hippocampus\n14,18,11,Left-Amygdala\n15,24,12,CSF\n16,26,13,Left-Accumbens-area\n17,28,14,Left-VentralDC\n18,41,1,Right-Cerebral-White-Matter\n19,43,2,Right-Lateral-Ventricle\n20,44,2,Right-Inf-Lat-Vent\n21,46,3,Right-Cerebellum-White-Matter\n22,47,4,Right-Cerebellum-Cortex\n23,49,5,Right-Thalamus-Proper\n24,50,6,Right-Caudate\n25,51,7,Right-Putamen\n26,52,8,Right-Pallidum\n27,53,10,Right-Hippocampus\n28,54,11,Right-Amygdala\n29,58,13,Right-Accumbens-area\n30,60,14,Right-VentralDC\n31,72,2,5th-Ventricle\n32,192,15,Corpus_Callosum\n33,251,15,CC_Posterior\n34,252,15,CC_Mid_Posterior\n35,253,15,CC_Central\n36,254,15,CC_Mid_Anterior\n37,255,15,CC_Anterior\n38,1001,16,ctx-lh-bankssts\n39,1002,17,ctx-lh-caudalanteriorcingulate\n40,1003,18,ctx-lh-caudalmiddlefrontal\n41,1005,19,ctx-lh-cuneus\n42,1006,20,ctx-lh-entorhinal\n43,1007,21,ctx-lh-fusiform\n44,1008,22,ctx-lh-inferiorparietal\n45,1009,23,ctx-lh-inferiortemporal\n46,1010,24,ctx-lh-isthmuscingulate\n47,1011,25,ctx-lh-lateraloccipital\n48,1012,26,ctx-lh-lateralorbitofrontal\n49,1013,27,ctx-lh-lingual\n50,1014,28,ctx-lh-medialorbitofrontal\n51,1015,29,ctx-lh-middletemporal\n52,1016,30,ctx-lh-parahippocampal\n53,1017,31,ctx-lh-paracentral\n54,1018,32,ctx-lh-parsopercularis\n55,1019,33,ctx-lh-parsorbitalis\n56,1020,34,ctx-lh-parstriangularis\n57,1021,35,ctx-lh-pericalcarine\n58,1022,36,ctx-lh-postcentral\n59,1023,37,ctx-lh-posteriorcingulate\n60,1024,38,ctx-lh-precentral\n61,1025,39,ctx-lh-precuneus\n62,1026,40,ctx-lh-rostralanteriorcingulate\n63,1027,41,ctx-lh-rostralmiddlefrontal\n64,1028,42,ctx-lh-superiorfrontal\n65,1029,43,ctx-lh-superiorparietal\n66,1030,44,ctx-lh-superiortemporal\n67,1031,45,ctx-lh-supramarginal\n68,1032,46,ctx-lh-frontalpole\n69,1033,47,ctx-lh-temporalpole\n70,1034,48,ctx-lh-transversetemporal\n71,1035,49,ctx-lh-insula\n72,2001,16,ctx-rh-bankssts\n73,2002,17,ctx-rh-caudalanteriorcingulate\n74,2003,18,ctx-rh-caudalmiddlefrontal\n75,2005,19,ctx-rh-cuneus\n76,2006,20,ctx-rh-entorhinal\n77,2007,21,ctx-rh-fusiform\n78,2008,22,ctx-rh-inferiorparietal\n79,2009,23,ctx-rh-inferiortemporal\n80,2010,24,ctx-rh-isthmuscingulate\n81,2011,25,ctx-rh-lateraloccipital\n82,2012,26,ctx-rh-lateralorbitofrontal\n83,2013,27,ctx-rh-lingual\n84,2014,28,ctx-rh-medialorbitofrontal\n85,2015,29,ctx-rh-middletemporal\n86,2016,30,ctx-rh-parahippocampal\n87,2017,31,ctx-rh-paracentral\n88,2018,32,ctx-rh-parsopercularis\n89,2019,33,ctx-rh-parsorbitalis\n90,2020,34,ctx-rh-parstriangularis\n91,2021,35,ctx-rh-pericalcarine\n92,2022,36,ctx-rh-postcentral\n93,2023,37,ctx-rh-posteriorcingulate\n94,2024,38,ctx-rh-precentral\n95,2025,39,ctx-rh-precuneus\n96,2026,40,ctx-rh-rostralanteriorcingulate\n97,2027,41,ctx-rh-rostralmiddlefrontal\n98,2028,42,ctx-rh-superiorfrontal\n99,2029,43,ctx-rh-superiorparietal\n100,2030,44,ctx-rh-superiortemporal\n101,2031,45,ctx-rh-supramarginal\n102,2032,46,ctx-rh-frontalpole\n103,2033,47,ctx-rh-temporalpole\n104,2034,48,ctx-rh-transversetemporal\n105,2035,49,ctx-rh-insul\n" }, { "path": "scripts/kwyk_reproduction/label_mappings/6-class-mapping.csv", "content": ",original,new,label,50-class\n0,0,0,Unknown,0\n1,2,1,Left-Cerebral-White-Matter,1\n2,4,3,Left-Lateral-Ventricle,2\n3,5,3,Left-Inf-Lat-Vent,2\n4,7,1,Left-Cerebellum-White-Matter,3\n5,8,2,Left-Cerebellum-Cortex,4\n6,10,4,Left-Thalamus-Proper,5\n7,11,4,Left-Caudate,6\n8,12,4,Left-Putamen,7\n9,13,4,Left-Pallidum,8\n10,14,3,3rd-Ventricle,2\n11,15,3,4th-Ventricle,2\n12,16,5,Brain-Stem,9\n13,17,4,Left-Hippocampus,10\n14,18,4,Left-Amygdala,11\n15,24,3,CSF,12\n16,26,4,Left-Accumbens-area,13\n17,28,4,Left-VentralDC,14\n18,41,1,Right-Cerebral-White-Matter,1\n19,43,3,Right-Lateral-Ventricle,2\n20,44,3,Right-Inf-Lat-Vent,2\n21,46,1,Right-Cerebellum-White-Matter,3\n22,47,2,Right-Cerebellum-Cortex,4\n23,49,4,Right-Thalamus-Proper,5\n24,50,4,Right-Caudate,6\n25,51,4,Right-Putamen,7\n26,52,4,Right-Pallidum,8\n27,53,4,Right-Hippocampus,10\n28,54,4,Right-Amygdala,11\n29,58,4,Right-Accumbens-area,13\n30,60,4,Right-VentralDC,14\n31,72,3,5th-Ventricle,2\n32,192,1,Corpus_Callosum,15\n33,251,1,CC_Posterior,15\n34,252,1,CC_Mid_Posterior,15\n35,253,1,CC_Central,15\n36,254,1,CC_Mid_Anterior,15\n37,255,1,CC_Anterior,15\n38,1001,2,ctx-lh-bankssts,16\n39,1002,2,ctx-lh-caudalanteriorcingulate,17\n40,1003,2,ctx-lh-caudalmiddlefrontal,18\n41,1005,2,ctx-lh-cuneus,19\n42,1006,2,ctx-lh-entorhinal,20\n43,1007,2,ctx-lh-fusiform,21\n44,1008,2,ctx-lh-inferiorparietal,22\n45,1009,2,ctx-lh-inferiortemporal,23\n46,1010,2,ctx-lh-isthmuscingulate,24\n47,1011,2,ctx-lh-lateraloccipital,25\n48,1012,2,ctx-lh-lateralorbitofrontal,26\n49,1013,2,ctx-lh-lingual,27\n50,1014,2,ctx-lh-medialorbitofrontal,28\n51,1015,2,ctx-lh-middletemporal,29\n52,1016,2,ctx-lh-parahippocampal,30\n53,1017,2,ctx-lh-paracentral,31\n54,1018,2,ctx-lh-parsopercularis,32\n55,1019,2,ctx-lh-parsorbitalis,33\n56,1020,2,ctx-lh-parstriangularis,34\n57,1021,2,ctx-lh-pericalcarine,35\n58,1022,2,ctx-lh-postcentral,36\n59,1023,2,ctx-lh-posteriorcingulate,37\n60,1024,2,ctx-lh-precentral,38\n61,1025,2,ctx-lh-precuneus,39\n62,1026,2,ctx-lh-rostralanteriorcingulate,40\n63,1027,2,ctx-lh-rostralmiddlefrontal,41\n64,1028,2,ctx-lh-superiorfrontal,42\n65,1029,2,ctx-lh-superiorparietal,43\n66,1030,2,ctx-lh-superiortemporal,44\n67,1031,2,ctx-lh-supramarginal,45\n68,1032,2,ctx-lh-frontalpole,46\n69,1033,2,ctx-lh-temporalpole,47\n70,1034,2,ctx-lh-transversetemporal,48\n71,1035,2,ctx-lh-insula,49\n72,2001,2,ctx-rh-bankssts,16\n73,2002,2,ctx-rh-caudalanteriorcingulate,17\n74,2003,2,ctx-rh-caudalmiddlefrontal,18\n75,2005,2,ctx-rh-cuneus,19\n76,2006,2,ctx-rh-entorhinal,20\n77,2007,2,ctx-rh-fusiform,21\n78,2008,2,ctx-rh-inferiorparietal,22\n79,2009,2,ctx-rh-inferiortemporal,23\n80,2010,2,ctx-rh-isthmuscingulate,24\n81,2011,2,ctx-rh-lateraloccipital,25\n82,2012,2,ctx-rh-lateralorbitofrontal,26\n83,2013,2,ctx-rh-lingual,27\n84,2014,2,ctx-rh-medialorbitofrontal,28\n85,2015,2,ctx-rh-middletemporal,29\n86,2016,2,ctx-rh-parahippocampal,30\n87,2017,2,ctx-rh-paracentral,31\n88,2018,2,ctx-rh-parsopercularis,32\n89,2019,2,ctx-rh-parsorbitalis,33\n90,2020,2,ctx-rh-parstriangularis,34\n91,2021,2,ctx-rh-pericalcarine,35\n92,2022,2,ctx-rh-postcentral,36\n93,2023,2,ctx-rh-posteriorcingulate,37\n94,2024,2,ctx-rh-precentral,38\n95,2025,2,ctx-rh-precuneus,39\n96,2026,2,ctx-rh-rostralanteriorcingulate,40\n97,2027,2,ctx-rh-rostralmiddlefrontal,41\n98,2028,2,ctx-rh-superiorfrontal,42\n99,2029,2,ctx-rh-superiorparietal,43\n100,2030,2,ctx-rh-superiortemporal,44\n101,2031,2,ctx-rh-supramarginal,45\n102,2032,2,ctx-rh-frontalpole,46\n103,2033,2,ctx-rh-temporalpole,47\n104,2034,2,ctx-rh-transversetemporal,48\n105,2035,2,ctx-rh-insul,49\n" }, { "path": "scripts/kwyk_reproduction/run.sh", "content": "#!/bin/bash\n# KWYK Brain Extraction Reproduction — Full Pipeline Runner\n#\n# Usage:\n# ./run.sh # Full pipeline (data + train + evaluate)\n# ./run.sh --smoke-test # Quick smoke test (5 volumes, 2 epochs)\n# ./run.sh --step data # Run only data assembly\n# ./run.sh --step train # Run only training (deterministic + Bayesian)\n# ./run.sh --step evaluate # Run only evaluation\n# ./run.sh --step compare # Run only kwyk comparison\n# ./run.sh --step sweep # Run only block size sweep\n#\n# Environment:\n# Creates a dedicated venv at .venv-kwyk/ with all dependencies.\n# Set NOBRAINER_ROOT to override the nobrainer repo location.\n\nset -euo pipefail\n\nSCRIPT_DIR=\"$(cd \"$(dirname \"${BASH_SOURCE[0]}\")\" && pwd)\"\nNOBRAINER_ROOT=\"${NOBRAINER_ROOT:-$(cd \"$SCRIPT_DIR/../..\" && pwd)}\"\nVENV_DIR=\"$SCRIPT_DIR/.venv-kwyk\"\nSTEP=\"${1:---all}\"\n\n# Colors for output\nGREEN='\\033[0;32m'\nYELLOW='\\033[1;33m'\nRED='\\033[0;31m'\nNC='\\033[0m'\n\nlog() { echo -e \"${GREEN}[kwyk]${NC} $*\"; }\nwarn() { echo -e \"${YELLOW}[kwyk]${NC} $*\"; }\nerr() { echo -e \"${RED}[kwyk]${NC} $*\" >&2; }\n\n# --- Setup venv ---\nsetup_venv() {\n if [ ! -d \"$VENV_DIR\" ]; then\n log \"Creating virtual environment at $VENV_DIR\"\n uv venv --python 3.14 \"$VENV_DIR\"\n fi\n\n log \"Installing dependencies...\"\n # shellcheck disable=SC1091\n source \"$VENV_DIR/bin/activate\"\n\n uv pip install -e \"$NOBRAINER_ROOT[bayesian,zarr,versioning,dev]\" \\\n monai pyro-ppl datalad matplotlib pyyaml scipy nibabel 2>&1 | tail -3\n log \"Dependencies installed\"\n}\n\n# --- Parse arguments ---\nSMOKE_TEST=false\nwhile [[ $# -gt 0 ]]; do\n case \"$1\" in\n --smoke-test)\n SMOKE_TEST=true\n shift\n ;;\n --step)\n STEP=\"$2\"\n shift 2\n ;;\n --all)\n STEP=\"--all\"\n shift\n ;;\n *)\n err \"Unknown argument: $1\"\n exit 1\n ;;\n esac\ndone\n\nsetup_venv\ncd \"$SCRIPT_DIR\"\n\n# --- Smoke test configuration ---\nif [ \"$SMOKE_TEST\" = true ]; then\n log \"Running SMOKE TEST (5 volumes, 2 epochs, tiny model)\"\n EXTRA_ARGS=\"--epochs 2\"\n DATASETS=\"ds000114\"\n # Use get_data() instead of DataLad for smoke test\nelse\n EXTRA_ARGS=\"\"\n DATASETS=\"ds000114 ds000228 ds002609\"\nfi\n\n# --- Step: Data Assembly ---\nrun_data() {\n log \"Step 1: Assembling dataset from OpenNeuro...\"\n python 01_assemble_dataset.py \\\n --datasets $DATASETS \\\n --output-csv manifest.csv \\\n --output-dir data \\\n --label-mapping binary\n log \"Dataset assembled: $(wc -l < manifest.csv) subjects\"\n}\n\n# --- Step: Training ---\nrun_train() {\n log \"Step 2: Training deterministic MeshNet (warm-start foundation)...\"\n python 02_train_meshnet.py \\\n --manifest manifest.csv \\\n --config config.yaml \\\n --output-dir checkpoints/meshnet \\\n $EXTRA_ARGS\n log \"Deterministic MeshNet trained (bwn / MAP variant)\"\n\n log \"Step 3a: MC Bernoulli dropout variant (bwn_multi)...\"\n python 03_train_bayesian.py \\\n --manifest manifest.csv \\\n --config config.yaml \\\n --variant bwn_multi \\\n --warmstart checkpoints/meshnet \\\n --output-dir checkpoints/bwn_multi \\\n $EXTRA_ARGS\n log \"MC Bernoulli dropout variant saved\"\n\n log \"Step 3b: Spike-and-slab dropout variant (bvwn_multi_prior)...\"\n python 03_train_bayesian.py \\\n --manifest manifest.csv \\\n --config config.yaml \\\n --variant bvwn_multi_prior \\\n --warmstart checkpoints/meshnet \\\n --output-dir checkpoints/bvwn_multi_prior \\\n $EXTRA_ARGS\n log \"Spike-and-slab dropout variant trained\"\n\n log \"Step 3c: Standard Gaussian Bayesian variant (for comparison)...\"\n python 03_train_bayesian.py \\\n --manifest manifest.csv \\\n --config config.yaml \\\n --variant bayesian_gaussian \\\n --warmstart checkpoints/meshnet \\\n --output-dir checkpoints/bayesian_gaussian \\\n $EXTRA_ARGS\n log \"Gaussian Bayesian variant trained\"\n}\n\n# --- Step: Evaluate ---\nrun_evaluate() {\n log \"Step 4: Evaluating all model variants on test set...\"\n if [ -f 04_evaluate.py ]; then\n for variant_dir in checkpoints/meshnet checkpoints/bwn_multi checkpoints/bvwn_multi_prior checkpoints/bayesian_gaussian; do\n variant_name=$(basename \"$variant_dir\")\n if [ -f \"$variant_dir/model.pth\" ]; then\n log \" Evaluating $variant_name...\"\n python 04_evaluate.py \\\n --model \"$variant_dir/model.pth\" \\\n --manifest manifest.csv \\\n --split test \\\n --n-samples 10 \\\n --output-dir \"results/$variant_name\"\n else\n warn \" Skipping $variant_name (no model.pth found)\"\n fi\n done\n else\n warn \"04_evaluate.py not found\"\n fi\n}\n\n# --- Step: Compare ---\nrun_compare() {\n log \"Step 5: Comparing with original kwyk container...\"\n if [ -f scripts/kwyk_reproduction/05_compare_kwyk.py ]; then\n python 05_compare_kwyk.py \\\n --new-model checkpoints/bayesian/model.pth \\\n --kwyk-dir \"$NOBRAINER_ROOT/../kwyk\" \\\n --manifest manifest.csv \\\n --output-dir results/comparison\n else\n warn \"05_compare_kwyk.py not yet implemented\"\n fi\n}\n\n# --- Step: Block Size Sweep ---\nrun_sweep() {\n log \"Step 6: Block size sweep...\"\n if [ -f scripts/kwyk_reproduction/06_block_size_sweep.py ]; then\n python 06_block_size_sweep.py \\\n --manifest manifest.csv \\\n --block-sizes 32 64 128 \\\n --output-dir results/sweep\n else\n warn \"06_block_size_sweep.py not yet implemented\"\n fi\n}\n\n# --- Execute ---\ncase \"$STEP\" in\n --all)\n run_data\n run_train\n run_evaluate\n run_compare\n run_sweep\n ;;\n data)\n run_data\n ;;\n train)\n run_train\n ;;\n evaluate)\n run_evaluate\n ;;\n compare)\n run_compare\n ;;\n sweep)\n run_sweep\n ;;\n *)\n err \"Unknown step: $STEP\"\n err \"Available: data, train, evaluate, compare, sweep\"\n exit 1\n ;;\nesac\n\nlog \"Done! Check figures/ and results/ for outputs.\"\n" }, { "path": "scripts/kwyk_reproduction/slurm_convert_zarr.sbatch", "content": "#!/bin/bash\n#SBATCH --job-name=kwyk-zarr\n#SBATCH --partition=mit_preemptable\n#SBATCH --cpus-per-task=8\n#SBATCH --mem=64G\n#SBATCH --time=02:00:00\n#SBATCH --output=slurm-zarr-%j.out\n#SBATCH --error=slurm-zarr-%j.err\n#\n# Convert PAC NIfTI dataset to Zarr3 for fast chunk-aligned I/O\n#\nset -euo pipefail\n\nWORK_DIR=\"/orcd/scratch/orcd/013/satra/kwyk_reproduction\"\nVENV_DIR=\"/orcd/data/satra/002/projects/nobrainer/venvs/nobrainer\"\nMANIFEST=\"kwyk_manifest_500.csv\"\nZARR_OUT=\"$WORK_DIR/data/kwyk_500.zarr\"\n\ncd \"$WORK_DIR\"\nsource \"${VENV_DIR}/bin/activate\"\n\necho \"=== Converting PAC dataset to Zarr3 ===\"\necho \"Started: $(date)\"\n\n# Build image/label path lists from manifest\npython -c \"\nimport csv\nimages, labels = [], []\nwith open('${MANIFEST}') as f:\n for row in csv.DictReader(f):\n if row['split'] in ('train', 'val', 'test'):\n images.append(row['t1w_path'])\n labels.append(row['label_path'])\n\n# Write temp files for the CLI\nwith open('/tmp/zarr_images.txt', 'w') as f:\n f.write('\\n'.join(images))\nwith open('/tmp/zarr_labels.txt', 'w') as f:\n f.write('\\n'.join(labels))\nprint(f'Volumes: {len(images)}')\n\"\n\n# Convert using nobrainer API directly\npython -c \"\nfrom nobrainer.datasets.zarr_store import create_zarr_store\nimport csv\n\npairs = []\nsubject_ids = []\nwith open('${MANIFEST}') as f:\n for row in csv.DictReader(f):\n pairs.append((row['t1w_path'], row['label_path']))\n subject_ids.append(row['subject_id'])\n\nprint(f'Converting {len(pairs)} volumes to Zarr3...')\nstore = create_zarr_store(\n image_label_pairs=pairs,\n output_path='${ZARR_OUT}',\n subject_ids=subject_ids,\n chunk_shape=(32, 32, 32),\n conform=True,\n target_shape=(256, 256, 256),\n target_voxel_size=(1.0, 1.0, 1.0),\n)\nprint(f'Zarr store created: {store}')\n\n# Create partition JSON for train/val/test splits\nimport json\npartitions = {'train': [], 'val': [], 'test': []}\nwith open('${MANIFEST}') as f:\n for row in csv.DictReader(f):\n partitions[row['split']].append(row['subject_id'])\n\npart_path = '${ZARR_OUT}_partition.json'\nwith open(part_path, 'w') as f:\n json.dump({'partitions': partitions}, f, indent=2)\nprint(f'Partition file: {part_path}')\nfor k, v in partitions.items():\n print(f' {k}: {len(v)} subjects')\n\"\n\necho \"=== Done: $(date) ===\"\n" }, { "path": "scripts/kwyk_reproduction/slurm_kwyk_bayesian.sbatch", "content": "#!/bin/bash\n#SBATCH --job-name=kwyk-bayes\n#SBATCH --partition=mit_preemptable\n#SBATCH --gres=gpu:1\n#SBATCH --cpus-per-task=8\n#SBATCH --mem=32G\n#SBATCH --time=06:00:00\n#SBATCH --output=slurm-kwyk-bayes-%j.out\n#SBATCH --error=slurm-kwyk-bayes-%j.err\n#\n# KWYK Smoke Test — Train one Bayesian variant (launched in parallel)\n#\n# Usage (via submit_kwyk_smoke.sh, not directly):\n# sbatch --dependency=afterok:$MESHNET_JOB slurm_kwyk_bayesian.sbatch bwn_multi\n#\nset -euo pipefail\n\nVARIANT=\"${1:?Usage: sbatch slurm_kwyk_bayesian.sbatch }\"\nWORK_DIR=\"/orcd/scratch/orcd/013/satra/kwyk_reproduction\"\nVENV_DIR=\"/orcd/data/satra/002/projects/nobrainer/venvs/nobrainer\"\nCONFIG=\"config_kwyk_smoke.yaml\"\nMANIFEST=\"kwyk_manifest_500.csv\"\n\necho \"=== KWYK Bayesian Training: ${VARIANT} ===\"\necho \"Job ID: ${SLURM_JOB_ID:-local}\"\necho \"Node: $(hostname)\"\necho \"Started: $(date)\"\n\ncd \"$WORK_DIR\"\nsource \"${VENV_DIR}/bin/activate\"\n\nEPOCHS=\"${KWYK_EPOCHS:-20}\"\n\necho \"=== Training ${VARIANT} (${EPOCHS} epochs) ===\"\npython 03_train_bayesian.py --manifest \"$MANIFEST\" --config \"$CONFIG\" \\\n --variant \"$VARIANT\" --warmstart checkpoints/kwyk_smoke_meshnet \\\n --output-dir \"checkpoints/kwyk_smoke_${VARIANT}\" --epochs \"$EPOCHS\"\n\necho \"=== ${VARIANT} complete: $(date) ===\"\n" }, { "path": "scripts/kwyk_reproduction/slurm_kwyk_evaluate.sbatch", "content": "#!/bin/bash\n#SBATCH --job-name=kwyk-eval\n#SBATCH --partition=mit_preemptable\n#SBATCH --gres=gpu:1\n#SBATCH --cpus-per-task=8\n#SBATCH --mem=32G\n#SBATCH --time=04:00:00\n#SBATCH --output=slurm-kwyk-eval-%j.out\n#SBATCH --error=slurm-kwyk-eval-%j.err\n#\n# KWYK Smoke Test — Evaluate all variants (runs after all training completes)\n#\nset -euo pipefail\n\nWORK_DIR=\"/orcd/scratch/orcd/013/satra/kwyk_reproduction\"\nVENV_DIR=\"/orcd/data/satra/002/projects/nobrainer/venvs/nobrainer\"\nCONFIG=\"config_kwyk_smoke.yaml\"\nMANIFEST=\"kwyk_manifest_500.csv\"\n\necho \"=== KWYK Evaluation ===\"\necho \"Job ID: ${SLURM_JOB_ID:-local}\"\necho \"Node: $(hostname)\"\necho \"Started: $(date)\"\n\ncd \"$WORK_DIR\"\nsource \"${VENV_DIR}/bin/activate\"\n\nfor v in kwyk_smoke_meshnet kwyk_smoke_bwn_multi kwyk_smoke_bvwn_multi_prior kwyk_smoke_bayesian_gaussian; do\n if [ -f \"checkpoints/$v/model.pth\" ]; then\n echo \"--- Evaluating $v (deterministic) ---\"\n python 04_evaluate.py --model \"checkpoints/$v\" \\\n --manifest \"$MANIFEST\" --config \"$CONFIG\" \\\n --split test --n-samples 0 \\\n --output-dir \"results/${v}_det\"\n echo \"--- Evaluating $v (MC, 3 samples) ---\"\n python 04_evaluate.py --model \"checkpoints/$v\" \\\n --manifest \"$MANIFEST\" --config \"$CONFIG\" \\\n --split test --n-samples 3 \\\n --output-dir \"results/${v}_mc\"\n else\n echo \"WARN: checkpoints/$v/model.pth not found, skipping\"\n fi\ndone\n\necho \"=== Evaluation complete: $(date) ===\"\necho \"Results:\"\nfor csv in results/kwyk_smoke_*/dice_scores.csv; do\n [ -f \"$csv\" ] && echo \" $csv: $(tail -n +2 \"$csv\" | wc -l) volumes\"\ndone\n" }, { "path": "scripts/kwyk_reproduction/slurm_kwyk_smoke.sbatch", "content": "#!/bin/bash\n#SBATCH --job-name=kwyk-pac-smoke\n#SBATCH --partition=mit_preemptable\n#SBATCH --gres=gpu:1\n#SBATCH --cpus-per-task=8\n#SBATCH --mem=32G\n#SBATCH --time=06:00:00\n#SBATCH --output=slurm-kwyk-smoke-%j.out\n#SBATCH --error=slurm-kwyk-smoke-%j.err\n#\n# KWYK Training — Step 1: deterministic MeshNet + manifest build\n# Steps 2-4 (Bayesian variants) are launched as dependent parallel jobs\n# by submit_kwyk_smoke.sh\n# Set KWYK_EPOCHS env var to override (default: 20)\n#\nset -euo pipefail\n\nWORK_DIR=\"/orcd/scratch/orcd/013/satra/kwyk_reproduction\"\nVENV_DIR=\"/orcd/data/satra/002/projects/nobrainer/venvs/nobrainer\"\nDATA_DIR=\"/orcd/scratch/orcd/013/satra/data/SharedData/segmentation/freesurfer_asegs\"\nCONFIG=\"config_kwyk_smoke.yaml\"\nMANIFEST=\"kwyk_manifest_500.csv\"\nN_SUBJECTS=500\n\necho \"=== KWYK PAC Dataset Smoke Test — MeshNet ===\"\necho \"Job ID: ${SLURM_JOB_ID:-local}\"\necho \"Node: $(hostname)\"\necho \"Started: $(date)\"\n\ncd \"$WORK_DIR\"\nsource \"${VENV_DIR}/bin/activate\"\n\npython -c \"\nimport torch\nprint('PyTorch:', torch.__version__)\nprint('CUDA:', torch.cuda.is_available())\nif torch.cuda.is_available():\n print('GPU:', torch.cuda.get_device_name(0))\n print('Memory:', round(torch.cuda.get_device_properties(0).total_memory / 1e9, 1), 'GB')\n\"\n\n# --- Build manifest if needed ---\nif [ ! -f \"$MANIFEST\" ]; then\n echo \"=== Building manifest (${N_SUBJECTS} subjects) ===\"\n python build_kwyk_manifest.py \\\n --data-dir \"$DATA_DIR\" \\\n --output-csv \"$MANIFEST\" \\\n --n-subjects \"$N_SUBJECTS\" \\\n --seed 42\nfi\n\necho \"Split counts:\"\ntail -n +2 \"$MANIFEST\" | cut -d, -f5 | sort | uniq -c\n\nEPOCHS=\"${KWYK_EPOCHS:-20}\"\n\n# --- Train deterministic MeshNet ---\necho \"=== Training deterministic MeshNet (${EPOCHS} epochs) ===\"\npython 02_train_meshnet.py --manifest \"$MANIFEST\" --config \"$CONFIG\" \\\n --output-dir checkpoints/kwyk_smoke_meshnet --epochs \"$EPOCHS\"\n\necho \"=== MeshNet complete: $(date) ===\"\n" }, { "path": "scripts/kwyk_reproduction/slurm_train.sbatch", "content": "#!/bin/bash\n#SBATCH --job-name=kwyk-train\n#SBATCH --partition=preemptible\n#SBATCH --gres=gpu:1\n#SBATCH --cpus-per-task=4\n#SBATCH --mem=32G\n#SBATCH --time=24:00:00\n#SBATCH --requeue\n#SBATCH --signal=B:USR1@120\n#SBATCH --output=slurm-%j.out\n#SBATCH --error=slurm-%j.err\n#\n# KWYK Brain Extraction Reproduction — SLURM Preemptible Training\n#\n# Trains all kwyk model variants with automatic checkpoint/resume on\n# preemption and automatic batch size optimization per GPU.\n#\n# Usage:\n# sbatch slurm_train.sbatch # 1 GPU\n# sbatch --gres=gpu:4 slurm_train.sbatch # 4 GPUs\n# sbatch --partition=gpu slurm_train.sbatch # Non-preemptible\n# KWYK_EPOCHS=100 sbatch slurm_train.sbatch # More epochs\n#\n# Multi-GPU:\n# Batch size is auto-optimized per GPU via nobrainer.gpu.auto_batch_size.\n# Request more GPUs with --gres=gpu:N.\n#\n# Environment variables:\n# KWYK_DATASETS — space-separated OpenNeuro IDs (default: ds000114)\n# KWYK_EPOCHS — epochs per variant (default: 50)\n# KWYK_WORK_DIR — working directory (default: $SLURM_SUBMIT_DIR)\n# KWYK_VENV — path to venv (default: .venv-kwyk)\n\nset -euo pipefail\n\nWORK_DIR=\"${KWYK_WORK_DIR:-${SLURM_SUBMIT_DIR:-$(pwd)}}\"\nVENV_DIR=\"${KWYK_VENV:-${WORK_DIR}/.venv-kwyk}\"\nDATASETS=\"${KWYK_DATASETS:-ds000114}\"\nEPOCHS=\"${KWYK_EPOCHS:-50}\"\nSCRIPT_DIR=\"$(cd \"$(dirname \"${BASH_SOURCE[0]}\")\" && pwd)\"\nNOBRAINER_ROOT=\"$(cd \"$SCRIPT_DIR/../..\" && pwd)\"\n\necho \"=== KWYK SLURM Training ===\"\necho \"Job ID: ${SLURM_JOB_ID:-local}\"\necho \"Node: $(hostname)\"\necho \"Partition: ${SLURM_JOB_PARTITION:-unknown}\"\necho \"GPUs: ${SLURM_GPUS_ON_NODE:-unknown}\"\necho \"Restart: ${SLURM_RESTART_COUNT:-0}\"\n\ncd \"$WORK_DIR\"\n\n# --- Setup venv (first run only) ---\nif [ ! -d \"$VENV_DIR\" ]; then\n uv venv --python 3.14 \"$VENV_DIR\"\nfi\n# shellcheck disable=SC1091\nsource \"${VENV_DIR}/bin/activate\"\nuv pip install -e \"${NOBRAINER_ROOT}[bayesian,versioning,dev]\" \\\n monai pyro-ppl datalad matplotlib pyyaml scipy nibabel 2>&1 | tail -3\nuv tool install git-annex 2>/dev/null || true\n\n# --- Show GPU info ---\npython -c \"\nfrom nobrainer.gpu import gpu_info, gpu_count\nprint('GPUs:', gpu_count())\nfor g in gpu_info():\n print(' GPU {id}: {name} ({memory_gb} GB)'.format(**g))\n\"\n\ncd \"$SCRIPT_DIR\"\n\n# --- Step 1: Data ---\nif [ ! -f manifest.csv ]; then\n # shellcheck disable=SC2086\n python 01_assemble_dataset.py --datasets $DATASETS \\\n --output-csv manifest.csv --output-dir data --label-mapping binary\nfi\n\n# --- Step 2: Deterministic MeshNet ---\npython 02_train_meshnet.py --manifest manifest.csv --config config.yaml \\\n --output-dir checkpoints/meshnet --epochs \"$EPOCHS\"\n\n# --- Step 3: All Bayesian variants (auto batch size, checkpoint/resume) ---\nfor variant in bwn_multi bvwn_multi_prior; do\n echo \"=== Training $variant ===\"\n python 03_train_bayesian.py --manifest manifest.csv --config config.yaml \\\n --variant \"$variant\" --warmstart checkpoints/meshnet \\\n --output-dir \"checkpoints/$variant\" --epochs \"$EPOCHS\"\ndone\n\n# --- Step 4: Evaluate ---\nfor v in checkpoints/meshnet checkpoints/bwn_multi checkpoints/bvwn_multi_prior; do\n [ -f \"$v/model.pth\" ] && python 04_evaluate.py --model \"$v/model.pth\" \\\n --manifest manifest.csv --split test --n-samples 10 \\\n --output-dir \"results/$(basename $v)\" || true\ndone\n\necho \"=== Done: $(ls checkpoints/*/model.pth 2>/dev/null | wc -l) models ===\"\n" }, { "path": "scripts/kwyk_reproduction/slurm_zarr_array.sbatch", "content": "#!/bin/bash\n#SBATCH --job-name=zarr-shard\n#SBATCH --partition=pi_satra\n#SBATCH --cpus-per-task=4\n#SBATCH --mem=32G\n#SBATCH --time=01:00:00\n#SBATCH --output=slurm-zarr-shard-%A_%a.out\n#SBATCH --error=slurm-zarr-shard-%A_%a.err\n#SBATCH --array=0-114\n#\n# Job array: each task writes one shard (100 subjects) to the Zarr store.\n# 11479 subjects / 100 per shard = 115 shards (0-114)\n# Task 0 also creates the store.\n#\nset -euo pipefail\n\nWORK_DIR=\"/orcd/scratch/orcd/013/satra/kwyk_reproduction\"\nVENV_DIR=\"/orcd/data/satra/002/projects/nobrainer/venvs/nobrainer\"\nMANIFEST=\"$WORK_DIR/kwyk_manifest_full.csv\"\nZARR_OUT=\"$WORK_DIR/data/kwyk_full.zarr\"\nSUBJECTS_PER_SHARD=100\n\ncd \"$WORK_DIR\"\nsource \"${VENV_DIR}/bin/activate\"\n\nif [ \"$SLURM_ARRAY_TASK_ID\" -eq 0 ]; then\n echo \"=== Shard 0: creating store + writing first shard ===\"\n python convert_zarr_shard.py \\\n --manifest \"$MANIFEST\" \\\n --zarr-store \"$ZARR_OUT\" \\\n --shard-idx 0 \\\n --subjects-per-shard \"$SUBJECTS_PER_SHARD\" \\\n --create\nelse\n # Wait briefly for shard 0 to create the store\n sleep 10\n echo \"=== Shard ${SLURM_ARRAY_TASK_ID}: writing ===\"\n python convert_zarr_shard.py \\\n --manifest \"$MANIFEST\" \\\n --zarr-store \"$ZARR_OUT\" \\\n --shard-idx \"$SLURM_ARRAY_TASK_ID\" \\\n --subjects-per-shard \"$SUBJECTS_PER_SHARD\"\nfi\n\necho \"=== Done: $(date) ===\"\n" }, { "path": "scripts/kwyk_reproduction/submit_kwyk_smoke.sh", "content": "#!/bin/bash\n# Submit the full KWYK PAC smoke test pipeline as parallel SLURM jobs:\n#\n# Job 1: MeshNet (deterministic warm-start)\n# Jobs 2-4: 3 Bayesian variants (parallel, depend on Job 1)\n# Job 5: Evaluate all variants (depends on Jobs 2-4)\n#\nset -euo pipefail\n\ncd \"$(dirname \"${BASH_SOURCE[0]}\")\"\n\necho \"=== Submitting KWYK PAC smoke test pipeline ===\"\n\n# Step 1: MeshNet\nMESHNET_JOB=$(sbatch --parsable slurm_kwyk_smoke.sbatch)\necho \"MeshNet: job ${MESHNET_JOB}\"\n\n# Step 2: Bayesian variants (parallel, depend on MeshNet)\nBAYES_JOBS=\"\"\nfor variant in bwn_multi bvwn_multi_prior bayesian_gaussian; do\n JOB=$(sbatch --parsable --dependency=afterok:${MESHNET_JOB} slurm_kwyk_bayesian.sbatch \"$variant\")\n echo \"${variant}: job ${JOB} (after ${MESHNET_JOB})\"\n BAYES_JOBS=\"${BAYES_JOBS:+${BAYES_JOBS},}${JOB}\"\ndone\n\n# Step 3: Evaluate (depends on all Bayesian jobs + MeshNet)\nEVAL_JOB=$(sbatch --parsable --dependency=afterok:${MESHNET_JOB}:${BAYES_JOBS} slurm_kwyk_evaluate.sbatch)\necho \"Evaluate: job ${EVAL_JOB} (after all training)\"\n\necho \"\"\necho \"=== Pipeline submitted ===\"\necho \"Monitor: squeue -u \\$USER\"\n" }, { "path": "scripts/kwyk_reproduction/utils.py", "content": "\"\"\"Shared utilities for kwyk reproduction experiments.\"\"\"\n\nfrom __future__ import annotations\n\nimport json\nimport logging\nfrom pathlib import Path\nimport signal\nfrom typing import Any\n\nimport numpy as np\nimport torch\n\n\ndef load_config(path: str | Path) -> dict[str, Any]:\n \"\"\"Load a YAML configuration file and return its contents as a dict.\n\n Parameters\n ----------\n path : str or Path\n Path to the YAML file.\n\n Returns\n -------\n dict\n Parsed configuration.\n \"\"\"\n import yaml\n\n path = Path(path)\n with open(path) as f:\n return yaml.safe_load(f)\n\n\ndef setup_logging(name: str) -> logging.Logger:\n \"\"\"Configure and return a logger with timestamped format.\n\n Parameters\n ----------\n name : str\n Logger name (typically ``__name__``).\n\n Returns\n -------\n logging.Logger\n Configured logger instance.\n \"\"\"\n logger = logging.getLogger(name)\n if not logger.handlers:\n handler = logging.StreamHandler()\n formatter = logging.Formatter(\n \"%(asctime)s - %(name)s - %(levelname)s - %(message)s\",\n datefmt=\"%Y-%m-%d %H:%M:%S\",\n )\n handler.setFormatter(formatter)\n logger.addHandler(handler)\n logger.setLevel(logging.INFO)\n return logger\n\n\ndef save_figure(fig: Any, path: str | Path) -> None:\n \"\"\"Save a matplotlib figure, creating parent directories if needed.\n\n Parameters\n ----------\n fig : matplotlib.figure.Figure\n The figure to save.\n path : str or Path\n Destination file path.\n \"\"\"\n path = Path(path)\n path.parent.mkdir(parents=True, exist_ok=True)\n fig.savefig(path, bbox_inches=\"tight\", dpi=150)\n\n\ndef compute_dice(pred: np.ndarray, label: np.ndarray) -> float:\n \"\"\"Compute the Dice score between two binary volumes.\n\n Parameters\n ----------\n pred : np.ndarray\n Binary prediction array.\n label : np.ndarray\n Binary ground-truth array.\n\n Returns\n -------\n float\n Dice coefficient in [0, 1]. Returns 1.0 when both arrays are empty.\n \"\"\"\n pred = pred.astype(bool)\n label = label.astype(bool)\n intersection = np.logical_and(pred, label).sum()\n total = pred.sum() + label.sum()\n if total == 0:\n return 1.0\n return float(2.0 * intersection / total)\n\n\ndef apply_label_mapping(\n label_vol: np.ndarray, mapping_csv: str | Path | None = None\n) -> np.ndarray:\n \"\"\"Remap FreeSurfer label codes in a volume.\n\n When *mapping_csv* is ``None`` the volume is binarised\n (``(vol > 0).astype(int)``). Otherwise a CSV with columns\n ``original,new`` is loaded and used to build a lookup table that maps\n each original code to its new value.\n\n Parameters\n ----------\n label_vol : np.ndarray\n Integer label volume.\n mapping_csv : str, Path, or None\n Path to a CSV mapping file. If ``None``, perform binary\n thresholding.\n\n Returns\n -------\n np.ndarray\n Remapped label volume with the same shape as the input.\n \"\"\"\n if mapping_csv is None:\n return (label_vol > 0).astype(int)\n\n import csv\n\n mapping_csv = Path(mapping_csv)\n lookup: dict[int, int] = {}\n with open(mapping_csv) as f:\n reader = csv.DictReader(f)\n for row in reader:\n lookup[int(row[\"original\"])] = int(row[\"new\"])\n\n mapper = np.vectorize(lambda v: lookup.get(v, 0))\n return mapper(label_vol)\n\n\n# ---------------------------------------------------------------------------\n# Checkpoint / resume for SLURM preemptible jobs\n# ---------------------------------------------------------------------------\n\n_logger = logging.getLogger(__name__)\n\n\nclass SlurmPreemptionHandler:\n \"\"\"Handle SLURM preemption signals for graceful checkpoint-and-exit.\n\n SLURM sends SIGUSR1 (or the signal specified by ``--signal``) before\n killing a preempted job. This handler sets a flag so the training\n loop can checkpoint and exit cleanly. The ``--requeue`` sbatch flag\n then re-submits the job, and the training resumes from the checkpoint.\n\n Usage::\n\n handler = SlurmPreemptionHandler()\n for epoch in range(start_epoch, total_epochs):\n train_one_epoch(...)\n save_checkpoint(...)\n if handler.preempted:\n log.info(\"Preempted — exiting for requeue\")\n sys.exit(0)\n \"\"\"\n\n def __init__(self, sig: int = signal.SIGUSR1) -> None:\n self.preempted = False\n self._sig = sig\n signal.signal(sig, self._handle)\n _logger.info(\"SLURM preemption handler registered (signal=%s)\", sig.name)\n\n def _handle(self, signum: int, frame: Any) -> None:\n _logger.warning(\n \"Received preemption signal %d — will checkpoint and exit\", signum\n )\n self.preempted = True\n\n\ndef save_training_checkpoint(\n checkpoint_dir: Path,\n model: torch.nn.Module,\n optimizer: torch.optim.Optimizer,\n epoch: int,\n metrics: dict[str, Any],\n) -> Path:\n \"\"\"Save a resumable training checkpoint.\n\n Writes ``checkpoint.pt`` containing model weights, optimizer state,\n epoch number, and accumulated metrics (losses, Dice scores, etc.).\n Also writes ``checkpoint_meta.json`` with human-readable status.\n\n Parameters\n ----------\n checkpoint_dir : Path\n Directory to save checkpoint files.\n model : torch.nn.Module\n Model to checkpoint.\n optimizer : torch.optim.Optimizer\n Optimizer to checkpoint (includes momentum, lr schedule state).\n epoch : int\n Completed epoch number (0-indexed).\n metrics : dict\n Accumulated training metrics to persist across restarts.\n\n Returns\n -------\n Path\n Path to the written checkpoint file.\n \"\"\"\n checkpoint_dir.mkdir(parents=True, exist_ok=True)\n ckpt_path = checkpoint_dir / \"checkpoint.pt\"\n\n torch.save(\n {\n \"epoch\": epoch,\n \"model_state_dict\": model.state_dict(),\n \"optimizer_state_dict\": optimizer.state_dict(),\n \"metrics\": metrics,\n },\n ckpt_path,\n )\n\n # Human-readable metadata\n meta = {\n \"epoch\": epoch,\n \"best_loss\": metrics.get(\"best_loss\", None),\n \"train_losses\": metrics.get(\"train_losses\", [])[-3:],\n }\n with open(checkpoint_dir / \"checkpoint_meta.json\", \"w\") as f:\n json.dump(meta, f, indent=2, default=str)\n\n _logger.info(\"Checkpoint saved: epoch %d → %s\", epoch, ckpt_path)\n return ckpt_path\n\n\ndef load_training_checkpoint(\n checkpoint_dir: Path,\n model: torch.nn.Module,\n optimizer: torch.optim.Optimizer | None = None,\n) -> tuple[int, dict[str, Any]]:\n \"\"\"Load a training checkpoint and return (start_epoch, metrics).\n\n Parameters\n ----------\n checkpoint_dir : Path\n Directory containing ``checkpoint.pt``.\n model : torch.nn.Module\n Model to load weights into.\n optimizer : torch.optim.Optimizer or None\n Optimizer to restore state into. If None, only model is loaded.\n\n Returns\n -------\n start_epoch : int\n The next epoch to train (checkpoint epoch + 1).\n metrics : dict\n Accumulated metrics from previous training.\n \"\"\"\n ckpt_path = checkpoint_dir / \"checkpoint.pt\"\n if not ckpt_path.exists():\n _logger.info(\"No checkpoint found at %s — starting from scratch\", ckpt_path)\n return 0, {}\n\n ckpt = torch.load(ckpt_path, weights_only=False)\n model.load_state_dict(ckpt[\"model_state_dict\"])\n if optimizer is not None and \"optimizer_state_dict\" in ckpt:\n optimizer.load_state_dict(ckpt[\"optimizer_state_dict\"])\n\n start_epoch = ckpt[\"epoch\"] + 1\n metrics = ckpt.get(\"metrics\", {})\n _logger.info(\n \"Resumed from checkpoint: epoch %d, best_loss=%.6f\",\n ckpt[\"epoch\"],\n metrics.get(\"best_loss\", float(\"inf\")),\n )\n return start_epoch, metrics\n" }, { "path": "scripts/synthseg_evaluation/02_train.py", "content": "#!/usr/bin/env python\n\"\"\"Train a segmentation model with real, synthetic, or mixed data.\n\nUsage:\n python 02_train.py --config config.yaml --mode real --model unet\n python 02_train.py --config config.yaml --mode mixed --model swin_unetr\n python 02_train.py --config config.yaml --mode synthetic --model kwyk_meshnet\n\"\"\"\n\nfrom __future__ import annotations\n\nimport argparse\nimport csv\nimport logging\nfrom pathlib import Path\n\nimport yaml\n\nlogging.basicConfig(level=logging.INFO, format=\"%(asctime)s %(levelname)s %(message)s\")\nlog = logging.getLogger(__name__)\n\n\ndef parse_args():\n parser = argparse.ArgumentParser(description=\"Train with SynthSeg evaluation\")\n parser.add_argument(\"--config\", default=\"config.yaml\")\n parser.add_argument(\"--mode\", choices=[\"real\", \"synthetic\", \"mixed\"], required=True)\n parser.add_argument(\"--model\", required=True)\n parser.add_argument(\"--manifest\", default=\"manifest.csv\")\n parser.add_argument(\"--output-dir\", default=\"checkpoints\")\n parser.add_argument(\"--epochs\", type=int, default=None)\n return parser.parse_args()\n\n\ndef load_manifest(path, split):\n pairs = []\n with open(path) as f:\n for row in csv.DictReader(f):\n if row[\"split\"] == split:\n pairs.append((row[\"t1w_path\"], row[\"label_path\"]))\n return pairs\n\n\ndef main():\n args = parse_args()\n with open(args.config) as f:\n config = yaml.safe_load(f)\n\n data_cfg = config[\"data\"]\n synth_cfg = config[\"synthseg\"]\n train_cfg = config[\"training\"]\n\n epochs = args.epochs or train_cfg[\"epochs\"]\n n_classes = data_cfg[\"n_classes\"]\n block_shape = tuple(data_cfg[\"block_shape\"])\n batch_size = data_cfg[\"batch_size\"]\n lr = train_cfg[\"lr\"]\n label_mapping = data_cfg[\"label_mapping\"]\n\n output_dir = Path(args.output_dir) / f\"{args.model}_{args.mode}\"\n output_dir.mkdir(parents=True, exist_ok=True)\n\n # Load manifest\n train_pairs = load_manifest(args.manifest, \"train\")\n log.info(\n \"Training: mode=%s, model=%s, %d volumes, %d epochs\",\n args.mode,\n args.model,\n len(train_pairs),\n epochs,\n )\n\n from nobrainer.processing.dataset import Dataset\n from nobrainer.processing.segmentation import Segmentation\n\n # Build dataset based on mode\n ds = (\n Dataset.from_files(train_pairs, block_shape=block_shape, n_classes=n_classes)\n .batch(batch_size)\n .binarize(label_mapping)\n .augment(train_cfg.get(\"augmentation_profile\", \"standard\"))\n )\n\n if args.mode == \"synthetic\" or args.mode == \"mixed\":\n from nobrainer.augmentation.synthseg import SynthSegGenerator\n\n label_paths = [p[1] for p in train_pairs]\n gen = SynthSegGenerator(\n label_paths,\n n_samples_per_map=synth_cfg[\"n_samples_per_map\"],\n elastic_std=synth_cfg[\"elastic_std\"],\n rotation_range=synth_cfg[\"rotation_range\"],\n scaling_bounds=synth_cfg[\"scaling_bounds\"],\n flipping=synth_cfg[\"flipping\"],\n randomize_resolution=synth_cfg[\"randomize_resolution\"],\n resolution_range=tuple(synth_cfg[\"resolution_range\"]),\n bias_field_std=synth_cfg[\"bias_field_std\"],\n noise_std=synth_cfg[\"noise_std\"],\n intensity_prior=tuple(synth_cfg[\"intensity_prior\"]),\n std_prior=tuple(synth_cfg[\"std_prior\"]),\n )\n if args.mode == \"mixed\":\n ds = ds.mix(gen, ratio=train_cfg[\"mixed_ratio\"])\n log.info(\"Mixed mode: %.0f%% synthetic\", train_cfg[\"mixed_ratio\"] * 100)\n\n # Build model\n model_args = {\"n_classes\": n_classes}\n if args.model in (\"swin_unetr\", \"segresnet\"):\n model_args[\"feature_size\"] = 12 if args.model == \"swin_unetr\" else 16\n\n seg = Segmentation(\n args.model, model_args=model_args, checkpoint_filepath=str(output_dir)\n )\n\n # Experiment tracking\n from nobrainer.experiment import ExperimentTracker\n\n tracker = ExperimentTracker(\n output_dir=output_dir,\n config={\n \"mode\": args.mode,\n \"model\": args.model,\n \"epochs\": epochs,\n \"n_classes\": n_classes,\n \"batch_size\": batch_size,\n },\n project=\"synthseg-evaluation\",\n name=f\"{args.model}_{args.mode}\",\n )\n\n import torch\n\n seg.fit(\n ds,\n epochs=epochs,\n optimizer=torch.optim.Adam,\n opt_args={\"lr\": lr},\n callbacks=[tracker.callback(mode=args.mode, model=args.model)],\n )\n seg.save(output_dir)\n tracker.finish()\n log.info(\"Model saved to %s\", output_dir)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "scripts/synthseg_evaluation/03_evaluate.py", "content": "#!/usr/bin/env python\n\"\"\"Evaluate a trained model with per-class Dice scoring.\n\nReuses the evaluation logic from the kwyk reproduction pipeline.\n\nUsage:\n python 03_evaluate.py --model checkpoints/unet_real --manifest manifest.csv\n\"\"\"\n\nfrom __future__ import annotations\n\nimport argparse\nimport csv\nimport logging\nfrom pathlib import Path\n\nimport nibabel as nib\nimport numpy as np\nimport yaml\n\nlogging.basicConfig(level=logging.INFO, format=\"%(asctime)s %(levelname)s %(message)s\")\nlog = logging.getLogger(__name__)\n\n\ndef per_class_dice(pred, gt, n_classes):\n \"\"\"Compute Dice per class c=1..n_classes-1 (skip background).\"\"\"\n dice = np.zeros(n_classes - 1)\n for c in range(1, n_classes):\n p = (pred == c).astype(np.float64)\n g = (gt == c).astype(np.float64)\n intersection = (p * g).sum()\n total = p.sum() + g.sum()\n dice[c - 1] = 2.0 * intersection / total if total > 0 else 1.0\n return dice\n\n\ndef main():\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--model\", required=True)\n parser.add_argument(\"--manifest\", required=True)\n parser.add_argument(\"--config\", default=\"config.yaml\")\n parser.add_argument(\"--split\", default=\"test\")\n parser.add_argument(\"--output-dir\", default=None)\n args = parser.parse_args()\n\n with open(args.config) as f:\n config = yaml.safe_load(f)\n\n n_classes = config[\"data\"][\"n_classes\"]\n block_shape = tuple(config[\"data\"][\"block_shape\"])\n label_mapping = config[\"data\"][\"label_mapping\"]\n\n output_dir = Path(args.output_dir or args.model) / \"eval\"\n output_dir.mkdir(parents=True, exist_ok=True)\n\n # Load model\n from nobrainer.processing.segmentation import Segmentation\n\n seg = Segmentation.load(args.model)\n\n # Load remap function\n remap_fn = None\n if label_mapping and label_mapping != \"binary\":\n from nobrainer.processing.dataset import _load_label_mapping\n\n remap_fn = _load_label_mapping(label_mapping)\n\n # Load test pairs\n pairs = []\n with open(args.manifest) as f:\n for row in csv.DictReader(f):\n if row[\"split\"] == args.split:\n pairs.append((row[\"t1w_path\"], row[\"label_path\"]))\n\n log.info(\"Evaluating %d volumes\", len(pairs))\n\n results = []\n all_dice = []\n for i, (img_path, lbl_path) in enumerate(pairs):\n gt = np.asarray(nib.load(lbl_path).dataobj, dtype=np.int32)\n if remap_fn is not None:\n gt = remap_fn(gt)\n\n pred_img = seg.predict(img_path, block_shape=block_shape)\n pred = np.asarray(pred_img.dataobj, dtype=np.int32)\n\n dice = per_class_dice(pred, gt, n_classes)\n avg = float(dice.mean())\n all_dice.append(dice)\n results.append({\"volume\": Path(img_path).stem, \"avg_dice\": avg})\n log.info(\" %d/%d: %s — Dice=%.4f\", i + 1, len(pairs), Path(img_path).stem, avg)\n\n # Save results\n csv_path = output_dir / \"dice_scores.csv\"\n with open(csv_path, \"w\", newline=\"\") as f:\n w = csv.DictWriter(f, [\"volume\", \"avg_dice\"])\n w.writeheader()\n w.writerows(results)\n\n np.save(output_dir / \"per_class_dice.npy\", np.array(all_dice))\n\n avg_dices = [r[\"avg_dice\"] for r in results]\n log.info(\"Class Dice: %.4f ± %.4f\", np.mean(avg_dices), np.std(avg_dices))\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "scripts/synthseg_evaluation/04_compare.py", "content": "#!/usr/bin/env python\n\"\"\"Compare results across training modes and models.\n\nUsage:\n python 04_compare.py --results-dir checkpoints/\n\"\"\"\n\nfrom __future__ import annotations\n\nimport argparse\nimport csv\nimport logging\nfrom pathlib import Path\n\nimport numpy as np\n\nlogging.basicConfig(level=logging.INFO, format=\"%(asctime)s %(levelname)s %(message)s\")\nlog = logging.getLogger(__name__)\n\n\ndef main():\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--results-dir\", default=\"checkpoints\")\n parser.add_argument(\"--output-dir\", default=\"results\")\n args = parser.parse_args()\n\n results_dir = Path(args.results_dir)\n output_dir = Path(args.output_dir)\n output_dir.mkdir(parents=True, exist_ok=True)\n\n # Scan for eval results: checkpoints/_/eval/dice_scores.csv\n rows = []\n for eval_dir in sorted(results_dir.glob(\"*/eval\")):\n csv_path = eval_dir / \"dice_scores.csv\"\n if not csv_path.exists():\n continue\n\n name = eval_dir.parent.name # e.g., \"unet_real\"\n parts = name.rsplit(\"_\", 1)\n model = parts[0] if len(parts) == 2 else name\n mode = parts[1] if len(parts) == 2 else \"unknown\"\n\n with open(csv_path) as f:\n scores = [float(r[\"avg_dice\"]) for r in csv.DictReader(f)]\n\n if scores:\n rows.append(\n {\n \"model\": model,\n \"mode\": mode,\n \"mean_dice\": f\"{np.mean(scores):.4f}\",\n \"std_dice\": f\"{np.std(scores):.4f}\",\n \"n_volumes\": len(scores),\n }\n )\n log.info(\n \"%s (%s): %.4f ± %.4f\", model, mode, np.mean(scores), np.std(scores)\n )\n\n if not rows:\n log.warning(\"No results found in %s\", results_dir)\n return\n\n # Write comparison table\n csv_path = output_dir / \"comparison_table.csv\"\n with open(csv_path, \"w\", newline=\"\") as f:\n w = csv.DictWriter(f, [\"model\", \"mode\", \"mean_dice\", \"std_dice\", \"n_volumes\"])\n w.writeheader()\n w.writerows(rows)\n log.info(\"Comparison table: %s\", csv_path)\n\n # Generate figure\n try:\n import matplotlib\n\n matplotlib.use(\"Agg\")\n import matplotlib.pyplot as plt\n\n models = sorted(set(r[\"model\"] for r in rows))\n modes = sorted(set(r[\"mode\"] for r in rows))\n x = np.arange(len(models))\n width = 0.25\n\n fig, ax = plt.subplots(figsize=(max(8, len(models) * 2), 6))\n for i, mode in enumerate(modes):\n means = []\n stds = []\n for model in models:\n match = [r for r in rows if r[\"model\"] == model and r[\"mode\"] == mode]\n if match:\n means.append(float(match[0][\"mean_dice\"]))\n stds.append(float(match[0][\"std_dice\"]))\n else:\n means.append(0)\n stds.append(0)\n ax.bar(x + i * width, means, width, yerr=stds, label=mode, capsize=3)\n\n ax.set_xlabel(\"Model\")\n ax.set_ylabel(\"Mean Class Dice\")\n ax.set_title(\"SynthSeg Evaluation: Model × Training Mode\")\n ax.set_xticks(x + width * (len(modes) - 1) / 2)\n ax.set_xticklabels(models)\n ax.set_ylim(0, 1.05)\n ax.legend()\n fig.tight_layout()\n fig.savefig(output_dir / \"comparison_figure.png\", dpi=150)\n plt.close(fig)\n log.info(\"Comparison figure: %s\", output_dir / \"comparison_figure.png\")\n except ImportError:\n log.warning(\"matplotlib not available, skipping figure\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "scripts/synthseg_evaluation/README.md", "content": "# SynthSeg Evaluation Pipeline\n\nEvaluate SynthSeg-based training against real-data baselines using\nmultiple model architectures.\n\n## Quick Start\n\n```bash\ncd scripts/synthseg_evaluation\n\n# Smoke test (2 epochs, unet, real+mixed)\n./run.sh --smoke-test\n\n# Full evaluation (all models × all modes from config.yaml)\n./run.sh\n```\n\n## Training Modes\n\n| Mode | Description |\n|------|-------------|\n| `real` | Train on real data only (baseline) |\n| `synthetic` | Train on SynthSeg-generated data only |\n| `mixed` | Train on mix of real + synthetic (configurable ratio) |\n\n## Available Models\n\n| Model | Architecture | Source |\n|-------|-------------|--------|\n| `unet` | 3D U-Net | MONAI |\n| `swin_unetr` | Swin Transformer U-Net | MONAI |\n| `segresnet` | Residual Encoder SegNet | MONAI |\n| `kwyk_meshnet` | VWN MeshNet + dropout | nobrainer |\n| `attention_unet` | Attention U-Net | MONAI |\n\n## Configuration\n\nEdit `config.yaml` to change models, training modes, SynthSeg parameters,\nand data settings. Key options:\n\n- `training.modes`: which modes to evaluate\n- `training.mixed_ratio`: fraction of synthetic data in mixed mode\n- `models`: list of model architectures to compare\n- `synthseg.*`: SynthSeg generation parameters\n\n## SLURM\n\n```bash\n# Single model+mode\nSYNTHSEG_MODE=mixed SYNTHSEG_MODEL=swin_unetr sbatch slurm_train.sbatch\n\n# All combinations\nfor model in unet swin_unetr kwyk_meshnet; do\n for mode in real synthetic mixed; do\n SYNTHSEG_MODE=$mode SYNTHSEG_MODEL=$model sbatch slurm_train.sbatch\n done\ndone\n```\n\n## Output\n\n```\nresults/\n├── comparison_table.csv # Dice per model × mode\n└── comparison_figure.png # Bar chart visualization\ncheckpoints/\n├── unet_real/eval/ # Per-model eval results\n├── unet_mixed/eval/\n├── swin_unetr_real/eval/\n└── ...\n```\n" }, { "path": "scripts/synthseg_evaluation/config.yaml", "content": "# SynthSeg Evaluation Pipeline Configuration\n\ndata:\n datasets: [ds000114]\n n_classes: 50\n label_mapping: 50-class\n block_shape: [32, 32, 32]\n batch_size: 32\n split: [80, 10, 10]\n\nsynthseg:\n n_samples_per_map: 20\n elastic_std: 4.0\n rotation_range: 15.0\n scaling_bounds: 0.2\n flipping: true\n randomize_resolution: true\n resolution_range: [1.0, 3.0]\n bias_field_std: 0.7\n noise_std: 0.1\n intensity_prior: [0, 250]\n std_prior: [0, 35]\n\ntraining:\n modes: [real, synthetic, mixed]\n mixed_ratio: 0.3\n epochs: 50\n lr: 0.0001\n augmentation_profile: standard\n\nmodels:\n - unet\n - swin_unetr\n - segresnet\n - segformer3d\n - kwyk_meshnet\n\nevaluation:\n n_samples: 10\n metrics: [per_class_dice, mean_dice]\n\nsmoke_test:\n epochs: 2\n n_samples_per_map: 2\n batch_size: 2\n block_shape: [16, 16, 16]\n models: [unet]\n modes: [real, mixed]\n" }, { "path": "scripts/synthseg_evaluation/run.sh", "content": "#!/bin/bash\n# SynthSeg Evaluation Pipeline Orchestrator\n#\n# Usage:\n# ./run.sh --smoke-test # Quick test (2 epochs, 1 model)\n# ./run.sh # Full evaluation\n# ./run.sh --config custom.yaml # Custom config\n\nset -euo pipefail\n\nSCRIPT_DIR=\"$(cd \"$(dirname \"${BASH_SOURCE[0]}\")\" && pwd)\"\nCONFIG=\"${1:-config.yaml}\"\nSMOKE=false\n\nwhile [[ $# -gt 0 ]]; do\n case \"$1\" in\n --smoke-test) SMOKE=true; shift ;;\n --config) CONFIG=\"$2\"; shift 2 ;;\n *) shift ;;\n esac\ndone\n\ncd \"$SCRIPT_DIR\"\n\necho \"=== SynthSeg Evaluation Pipeline ===\"\necho \"Config: $CONFIG\"\necho \"Smoke test: $SMOKE\"\n\n# Use sample data if no manifest exists\nif [ ! -f manifest.csv ]; then\n echo \"=== Creating manifest from sample data ===\"\n python -c \"\nimport csv\nfrom nobrainer.utils import get_data\nsrc = get_data()\npairs = []\nwith open(src) as f:\n r = csv.reader(f); next(r)\n pairs = list(r)[:5]\nsplits = ['train','train','train','val','test']\nwith open('manifest.csv', 'w', newline='') as f:\n w = csv.DictWriter(f, ['t1w_path','label_path','split']); w.writeheader()\n for i,(t1,lbl) in enumerate(pairs):\n w.writerow(dict(t1w_path=t1, label_path=lbl, split=splits[i]))\nprint('Manifest created with', len(pairs), 'volumes')\n\"\nfi\n\nif [ \"$SMOKE\" = true ]; then\n echo \"=== Smoke test: 2 epochs, unet, real+mixed ===\"\n for mode in real mixed; do\n echo \" Training unet ($mode)...\"\n python 02_train.py --config \"$CONFIG\" --mode \"$mode\" --model unet \\\n --epochs 2 --manifest manifest.csv\n done\n for mode in real mixed; do\n echo \" Evaluating unet ($mode)...\"\n python 03_evaluate.py --model \"checkpoints/unet_${mode}\" \\\n --manifest manifest.csv --config \"$CONFIG\" || true\n done\n python 04_compare.py --results-dir checkpoints/ --output-dir results/ || true\nelse\n # Full evaluation from config\n MODELS=$(python -c \"import yaml; c=yaml.safe_load(open('$CONFIG')); print(' '.join(c['training']['modes']))\")\n MODES=$(python -c \"import yaml; c=yaml.safe_load(open('$CONFIG')); print(' '.join(c['models']))\")\n\n for model in $MODES; do\n for mode in $MODELS; do\n echo \"=== Training $model ($mode) ===\"\n python 02_train.py --config \"$CONFIG\" --mode \"$mode\" --model \"$model\" \\\n --manifest manifest.csv\n done\n done\n\n for model in $MODES; do\n for mode in $MODELS; do\n echo \"=== Evaluating $model ($mode) ===\"\n python 03_evaluate.py --model \"checkpoints/${model}_${mode}\" \\\n --manifest manifest.csv --config \"$CONFIG\" || true\n done\n done\n\n python 04_compare.py --results-dir checkpoints/ --output-dir results/\nfi\n\necho \"=== Done. Results in results/ ===\"\n" }, { "path": "scripts/synthseg_evaluation/slurm_train.sbatch", "content": "#!/bin/bash\n#SBATCH --job-name=synthseg-eval\n#SBATCH --partition=preemptible\n#SBATCH --gres=gpu:1\n#SBATCH --cpus-per-task=4\n#SBATCH --mem=32G\n#SBATCH --time=24:00:00\n#SBATCH --requeue\n#SBATCH --signal=B:USR1@120\n#SBATCH --output=slurm-%j.out\n#SBATCH --error=slurm-%j.err\n#\n# SynthSeg Evaluation — SLURM Preemptible Training\n#\n# Usage:\n# sbatch slurm_train.sbatch\n# SYNTHSEG_MODE=mixed SYNTHSEG_MODEL=swin_unetr sbatch slurm_train.sbatch\n\nset -euo pipefail\n\nMODE=\"${SYNTHSEG_MODE:-real}\"\nMODEL=\"${SYNTHSEG_MODEL:-unet}\"\nCONFIG=\"${SYNTHSEG_CONFIG:-config.yaml}\"\nSCRIPT_DIR=\"$(cd \"$(dirname \"${BASH_SOURCE[0]}\")\" && pwd)\"\n\necho \"=== SynthSeg SLURM Training ===\"\necho \"Model: $MODEL, Mode: $MODE\"\necho \"Job: ${SLURM_JOB_ID:-local}, Restart: ${SLURM_RESTART_COUNT:-0}\"\n\ncd \"$SCRIPT_DIR\"\n\npython 02_train.py --config \"$CONFIG\" --mode \"$MODE\" --model \"$MODEL\" \\\n --manifest manifest.csv\npython 03_evaluate.py --model \"checkpoints/${MODEL}_${MODE}\" \\\n --manifest manifest.csv --config \"$CONFIG\"\n" } ]