Note
Go to the end to download the full example code.
Deep-learning habitat embeddings
Background. Deep-learning encoders turn an image into a stack of feature maps (channels). Averaging those channels inside each habitat gives every habitat a learned descriptor that can sit next to hand-crafted radiomics in a downstream model.
Purpose. You get a table with one row per habitat and one column per channel (64 here), and a heatmap of the first 16 dimensions. The feature volume on this page is random noise, so the numbers only show the mechanics, not meaningful biology.
Key terms.
embedding – a fixed-length vector of learned numbers describing a region; here one 64-D vector per habitat.
masked spatial average pooling – the mean of each channel over the voxels where the habitat mask is non-zero.
(C, z, y, x) – channels first, then the same spatial axes as the habitat label array; the two must match voxel for voxel.
Bridge HABIT habitat maps to deep-learning pipelines: extract a binary
mask per habitat and apply masked spatial average pooling on a 3-D
feature tensor (simulated here with NumPy; swap in a MONAI / PyTorch
encoder output with the same (C, z, y, x) layout).
Habitat map from HABIT, synthetic 64-channel feature volume from a DL model. sphinx_gallery_thumbnail_number = 1
from pathlib import Path
from typing import Dict, Tuple
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from habit.contracts import cohort_from_directory
from habit.datasets import fetch_demo
from habit.recipes import one_step_habitat
def masked_spatial_average_pooling(
feature_map: np.ndarray,
mask: np.ndarray,
) -> np.ndarray:
"""
Average-pool a 4-D feature tensor over voxels where ``mask > 0``.
Parameters
----------
feature_map:
Deep feature volume, shape ``(C, z, y, x)``.
mask:
Binary or integer habitat mask, shape ``(z, y, x)``.
Returns
-------
np.ndarray
Per-channel mean, shape ``(C,)``.
"""
if feature_map.ndim != 4:
raise ValueError(
f"feature_map must be 4-D (C, z, y, x); got shape {feature_map.shape}"
)
if mask.shape != feature_map.shape[1:]:
raise ValueError(
f"mask shape {mask.shape} must match spatial dims {feature_map.shape[1:]}"
)
region = mask > 0
if not np.any(region):
return np.full(feature_map.shape[0], np.nan, dtype=np.float64)
pooled = feature_map[:, region].mean(axis=1)
return np.asarray(pooled, dtype=np.float64)
def habitat_embedding_table(
feature_map: np.ndarray,
label_array: np.ndarray,
habitat_ids: Tuple[int, ...],
) -> pd.DataFrame:
"""
Build a table of masked-pooled embeddings, one row per habitat.
Parameters
----------
feature_map:
Deep feature volume, shape ``(C, z, y, x)``.
label_array:
Integer habitat labels, shape ``(z, y, x)``.
habitat_ids:
Habitat ids to embed (model order).
Returns
-------
pd.DataFrame
Rows indexed by habitat id; columns ``emb_0`` … ``emb_{C-1}``.
"""
rows: Dict[int, np.ndarray] = {}
for hid in habitat_ids:
# Binary mask of this one habitat; absent habitats pool to NaN.
mask = label_array == hid
rows[hid] = masked_spatial_average_pooling(feature_map, mask)
n_channels = int(feature_map.shape[0])
columns = [f"emb_{i}" for i in range(n_channels)]
frame = pd.DataFrame.from_dict(rows, orient="index", columns=columns)
frame.index.name = "habitat_id"
return frame
DATA = fetch_demo()
MODALITIES = ("LAP",)
ROI = "LAP"
cohort = cohort_from_directory(DATA, modalities=MODALITIES, roi=ROI)[:1]
result = one_step_habitat(
modalities=MODALITIES, n_habitats=3, random_seed=0, roi=ROI
).fit_predict(cohort)
habitat_map = result.habitat_maps[0]
labels = np.asarray(habitat_map.label_array)
spatial = labels.shape
rng = np.random.default_rng(0)
# Simulate a 3-D encoder output: 64 channels, same (z, y, x) as the ROI.
n_channels = 64
feature_map = rng.standard_normal((n_channels,) + spatial, dtype=np.float32)
# Optional: torch users can replace the array with ``tensor.detach().cpu().numpy()``.
embeddings = habitat_embedding_table(
feature_map,
labels,
habitat_map.habitat_ids,
)
print(f"Embedding shape: {embeddings.shape} (habitats x {n_channels}-D)")
print(embeddings.iloc[:, :6].round(3))
embeddings.iloc[:, :6]
HABIT demo data (cached)
DATA (preprocessed root): C:\Users\dongm\.habit_data\demo-data-v1\preprocessed
On-disk inventory of this folder:
subjects (5): subj001, subj002, subj003, subj004, subj005
image series: LAP, PVP, delay_3min, pre_contrast
mask keys: LAP, PVP, delay_3min, pre_contrast
example image: images/subj001/delay_3min/WATER__BH_Ax_LAVA_Flex_3min_Series0012.nrrd
example mask: masks/subj001/delay_3min/WATER__BH_Ax_LAVA_Flex_10min_Series0017_mask.nrrd
Your own data must use the same folder tree (change IDs / series names):
DATA/
images/<subject_id>/<modality>/<one image file>
masks/<subject_id>/<roi>/<one mask file>
Then load it with the same call the demos use:
cohort = cohort_from_directory(DATA, modalities=("LAP",), roi="LAP")
Swap DATA / modalities / roi to match your tree. Mask key is often the
same as one image series (here LAP).
Cohort.map[_DefineAndLabelWithinSubject]: 0%| | 0/1 [00:00<?, ?it/s]
Cohort.map[_DefineAndLabelWithinSubject]: 100%|██████████| 1/1 [00:01<00:00, 1.27s/it]
Cohort.map[_DefineAndLabelWithinSubject]: 100%|██████████| 1/1 [00:01<00:00, 1.27s/it]
Embedding shape: (3, 64) (habitats x 64-D)
emb_0 emb_1 emb_2 emb_3 emb_4 emb_5
habitat_id
1 0.004 -0.004 -0.008 -0.009 -0.003 -0.001
2 -0.001 -0.005 0.002 0.012 -0.010 -0.010
3 -0.003 -0.005 -0.006 0.006 -0.004 -0.009
Heatmap of the first 16 embedding dimensions per habitat (English labels only).
Path("out").mkdir(exist_ok=True)
fig, ax = plt.subplots(figsize=(6, 3))
subset = embeddings.iloc[:, :16].to_numpy()
im = ax.imshow(subset, aspect="auto", cmap="viridis")
ax.set_xlabel("Embedding dimension")
ax.set_ylabel("Habitat id")
ax.set_yticks(range(len(embeddings)))
ax.set_yticklabels([f"H{idx}" for idx in embeddings.index])
ax.set_title("Masked spatial average pooling (64-D, first 16 shown)")
fig.colorbar(im, ax=ax, label="Pooled activation")
fig.savefig("out/dl_habitat_embeddings.png", dpi=150, bbox_inches="tight")
plt.show()
Total running time of the script: (0 minutes 20.251 seconds)