Update scripts

mathpluscode · mathpluscode · commit 57904c83e4f8 · 2025-05-19T02:47:41.000+01:00
diff --git a/.gitignore b/.gitignore
@@ -149,5 +149,5 @@ wandb/
 **/*.safetensors
 
 # docs
-examples/*.png
 out/
+**/*.png
diff --git a/README.md b/README.md
@@ -4,7 +4,7 @@
   <img alt="CineMA logo" src="logo_light.svg" height="256">
 </picture>
 
-# CineMA: A Foundation Model for Cine Cardiac MRI 🎥🫀
+# CineMA: A Vision Foundation Model for Cine Cardiac MRI 🎥🫀
 
 ![python](https://img.shields.io/badge/Python-3.11-3776AB.svg?style=flat&logo=python&logoColor=white)
 ![pytorch](https://img.shields.io/badge/PyTorch-EE4C2C?style=flat&logo=pytorch&logoColor=white)
@@ -15,7 +15,7 @@
 
 ## Overview
 
-**CineMA** is a foundation model for **Cine** cardiac magnetic resonance (CMR) imaging based on
+**CineMA** is a vision foundation model for **Cine** cardiac magnetic resonance (CMR) imaging based on
 **M**asked-**A**utoencoder. CineMA has been pre-trained on UK Biobank data and fine-tuned on multiple clinically
 relevant tasks such as ventricle and myocaridum segmentation, ejection fraction (EF) regression, cardiovascular disease
 (CVD) detection and classification, and mid-valve plane and apical landmark localization. The model has been evaluated
@@ -73,7 +73,7 @@ python examples/inference/landmark_coordinate.py
 | Landmark localization by heatmap regression     | LAX 2C or LAX 4C | 1                | [landmark_heatmap.py](examples/inference/landmark_heatmap.py)           |
 | Landmark localization by coordinates regression | LAX 2C or LAX 4C | 1                | [landmark_coordinate.py](examples/inference/landmark_coordinate.py)     |
 
-### Use pre-trained models for fine-tuning
+### Use pre-trained models
 
 The pre-trained CineMA model backbone is available at https://huggingface.co/mathpluscode/CineMA. Following scripts
 demonstrated how to fine-tune this backbone using
@@ -91,10 +91,17 @@ python examples/train/regression.py
 | Cardiovascular disease classification | [classification.py](examples/train/classification.py) |
 | Ejection fraction regression          | [regression.py](examples/train/regression.py)         |
 
-For other datasets, pre-process can be performed using the provided scripts following the documentations. Note that it
-is recommended to download the data under `~/.cache/cinema_datasets` as the integration tests uses this path. For
-instance, the mnms preprocessed data would be `~/.cache/cinema_datasets/mnms/processed`. Otherwise define the path using
-environment variable `CINEMA_DATA_DIR`.
+Another two scripts demonstrated the masking and prediction process of MAE and the feature extraction from MAE.
+
+```bash
+python examples/inference/mae.py
+python examples/inference/mae_feature_extraction.py
+```
+
+For fine-tuning CineMA on other datasets, pre-process can be performed using the provided scripts following the
+documentations. Note that it is recommended to download the data under `~/.cache/cinema_datasets` as the integration
+tests uses this path. For instance, the mnms preprocessed data would be `~/.cache/cinema_datasets/mnms/processed`.
+Otherwise define the path using environment variable `CINEMA_DATA_DIR`.
 
 | Training Data | Documentations                               |
 | ------------- | -------------------------------------------- |
diff --git a/cinema/mae/mae.py b/cinema/mae/mae.py
@@ -629,7 +629,7 @@ def from_pretrained(cls, **kwargs) -> CineMA:  # type: ignore[no-untyped-def]
         # download config
         config_path = hf_hub_download(
             repo_id="mathpluscode/CineMA",
-            filename="pretrained/cinema.yaml",
+            filename="pretrained/config.yaml",
             **kwargs,
         )
         logger.info(f"Cached model config to {config_path}.")
diff --git a/examples/inference/classification_cvd.py b/examples/inference/classification_cvd.py
@@ -17,16 +17,16 @@ def run(trained_dataset: str, view: str, seed: int) -> None:
     # load config to get class names
     config_path = hf_hub_download(
         repo_id="mathpluscode/CineMA",
-        filename=f"finetuned/classification_cvd/{trained_dataset}_{view}.yaml",
+        filename=f"finetuned/classification_cvd/{trained_dataset}_{view}/config.yaml",
     )
     config = OmegaConf.load(config_path)
     classes = list(config.data[config.data.class_column])
 
     # load model
     model = ConvViT.from_finetuned(
         repo_id="mathpluscode/CineMA",
-        model_filename=f"finetuned/classification_cvd/{trained_dataset}_{view}_{seed}.safetensors",
-        config_filename=f"finetuned/classification_cvd/{trained_dataset}_{view}.yaml",
+        model_filename=f"finetuned/classification_cvd/{trained_dataset}_{view}/{trained_dataset}_{view}_{seed}.safetensors",
+        config_filename=f"finetuned/classification_cvd/{trained_dataset}_{view}/config.yaml",
     )
 
     # load sample data from mnms2 of class HCM and form a batch of size 1
diff --git a/examples/inference/classification_sex.py b/examples/inference/classification_sex.py
@@ -18,16 +18,16 @@ def run(seed: int) -> None:
     # load config to get class names
     config_path = hf_hub_download(
         repo_id="mathpluscode/CineMA",
-        filename=f"finetuned/classification_sex/{trained_dataset}_{view}.yaml",
+        filename=f"finetuned/classification_sex/{trained_dataset}_{view}/config.yaml",
     )
     config = OmegaConf.load(config_path)
     classes = list(config.data[config.data.class_column])
 
     # load model
     model = ConvViT.from_finetuned(
         repo_id="mathpluscode/CineMA",
-        model_filename=f"finetuned/classification_sex/{trained_dataset}_{view}_{seed}.safetensors",
-        config_filename=f"finetuned/classification_sex/{trained_dataset}_{view}.yaml",
+        model_filename=f"finetuned/classification_sex/{trained_dataset}_{view}/{trained_dataset}_{view}_{seed}.safetensors",
+        config_filename=f"finetuned/classification_sex/{trained_dataset}_{view}/config.yaml",
     )
 
     # load sample data from mnms2 of class HCM and form a batch of size 1
diff --git a/examples/inference/classification_vendor.py b/examples/inference/classification_vendor.py
@@ -18,16 +18,16 @@ def run(view: str, seed: int) -> None:
     # load config to get class names
     config_path = hf_hub_download(
         repo_id="mathpluscode/CineMA",
-        filename=f"finetuned/classification_vendor/{trained_dataset}_{view}.yaml",
+        filename=f"finetuned/classification_vendor/{trained_dataset}_{view}/config.yaml",
     )
     config = OmegaConf.load(config_path)
     classes = list(config.data[config.data.class_column])
 
     # load model
     model = ConvViT.from_finetuned(
         repo_id="mathpluscode/CineMA",
-        model_filename=f"finetuned/classification_vendor/{trained_dataset}_{view}_{seed}.safetensors",
-        config_filename=f"finetuned/classification_vendor/{trained_dataset}_{view}.yaml",
+        model_filename=f"finetuned/classification_vendor/{trained_dataset}_{view}/{trained_dataset}_{view}_{seed}.safetensors",
+        config_filename=f"finetuned/classification_vendor/{trained_dataset}_{view}/config.yaml",
     )
 
     # load sample data from mnms2 of class HCM and form a batch of size 1
diff --git a/examples/inference/landmark_coordinate.py b/examples/inference/landmark_coordinate.py
@@ -17,8 +17,8 @@ def run(view: str, seed: int) -> None:
     # load model
     model = ConvViT.from_finetuned(
         repo_id="mathpluscode/CineMA",
-        model_filename=f"finetuned/landmark_coordinate/{view}_{seed}.safetensors",
-        config_filename=f"finetuned/landmark_coordinate/{view}.yaml",
+        model_filename=f"finetuned/landmark_coordinate/{view}/{view}_{seed}.safetensors",
+        config_filename=f"finetuned/landmark_coordinate/{view}/config.yaml",
     )
 
     # load sample data and form a batch of size 1
diff --git a/examples/inference/landmark_heatmap.py b/examples/inference/landmark_heatmap.py
@@ -17,8 +17,8 @@ def run(view: str, seed: int) -> None:
     # load model
     model = ConvUNetR.from_finetuned(
         repo_id="mathpluscode/CineMA",
-        model_filename=f"finetuned/landmark_heatmap/{view}_{seed}.safetensors",
-        config_filename=f"finetuned/landmark_heatmap/{view}.yaml",
+        model_filename=f"finetuned/landmark_heatmap/{view}/{view}_{seed}.safetensors",
+        config_filename=f"finetuned/landmark_heatmap/{view}/config.yaml",
     )
 
     # load sample data and form a batch of size 1
diff --git a/examples/inference/regression_age.py b/examples/inference/regression_age.py
@@ -18,7 +18,7 @@ def run(seed: int) -> None:
     # load config to get class names
     config_path = hf_hub_download(
         repo_id="mathpluscode/CineMA",
-        filename=f"finetuned/regression_age/{trained_dataset}_{view}.yaml",
+        filename=f"finetuned/regression_age/{trained_dataset}_{view}/config.yaml",
     )
     config = OmegaConf.load(config_path)
     mean = config.data[config.data.regression_column].mean
@@ -27,8 +27,8 @@ def run(seed: int) -> None:
     # load model
     model = ConvViT.from_finetuned(
         repo_id="mathpluscode/CineMA",
-        model_filename=f"finetuned/regression_age/{trained_dataset}_{view}_{seed}.safetensors",
-        config_filename=f"finetuned/regression_age/{trained_dataset}_{view}.yaml",
+        model_filename=f"finetuned/regression_age/{trained_dataset}_{view}/{trained_dataset}_{view}_{seed}.safetensors",
+        config_filename=f"finetuned/regression_age/{trained_dataset}_{view}/config.yaml",
     )
 
     # load sample data
diff --git a/examples/inference/regression_bmi.py b/examples/inference/regression_bmi.py
@@ -18,7 +18,7 @@ def run(seed: int) -> None:
     # load config to get class names
     config_path = hf_hub_download(
         repo_id="mathpluscode/CineMA",
-        filename=f"finetuned/regression_bmi/{trained_dataset}_{view}.yaml",
+        filename=f"finetuned/regression_bmi/{trained_dataset}_{view}/config.yaml",
     )
     config = OmegaConf.load(config_path)
     mean = config.data[config.data.regression_column].mean
@@ -27,8 +27,8 @@ def run(seed: int) -> None:
     # load model
     model = ConvViT.from_finetuned(
         repo_id="mathpluscode/CineMA",
-        model_filename=f"finetuned/regression_bmi/{trained_dataset}_{view}_{seed}.safetensors",
-        config_filename=f"finetuned/regression_bmi/{trained_dataset}_{view}.yaml",
+        model_filename=f"finetuned/regression_bmi/{trained_dataset}_{view}/{trained_dataset}_{view}_{seed}.safetensors",
+        config_filename=f"finetuned/regression_bmi/{trained_dataset}_{view}/config.yaml",
     )
 
     # load sample data
diff --git a/examples/inference/regression_ef.py b/examples/inference/regression_ef.py
@@ -17,7 +17,7 @@ def run(trained_dataset: str, view: str, seed: int) -> None:
     # load config to get class names
     config_path = hf_hub_download(
         repo_id="mathpluscode/CineMA",
-        filename=f"finetuned/regression_ef/{trained_dataset}_{view}.yaml",
+        filename=f"finetuned/regression_ef/{trained_dataset}_{view}/config.yaml",
     )
     config = OmegaConf.load(config_path)
     mean = config.data[config.data.regression_column].mean
@@ -26,8 +26,8 @@ def run(trained_dataset: str, view: str, seed: int) -> None:
     # load model
     model = ConvViT.from_finetuned(
         repo_id="mathpluscode/CineMA",
-        model_filename=f"finetuned/regression_ef/{trained_dataset}_{view}_{seed}.safetensors",
-        config_filename=f"finetuned/regression_ef/{trained_dataset}_{view}.yaml",
+        model_filename=f"finetuned/regression_ef/{trained_dataset}_{view}/{trained_dataset}_{view}_{seed}.safetensors",
+        config_filename=f"finetuned/regression_ef/{trained_dataset}_{view}/config.yaml",
     )
 
     # load sample data
diff --git a/examples/inference/segmentation_lax_4c.py b/examples/inference/segmentation_lax_4c.py
@@ -46,25 +46,26 @@ def run(seed: int) -> None:
     trained_dataset = "mnms2"
 
     # load model
+    view = "lax_4c"
     model = ConvUNetR.from_finetuned(
         repo_id="mathpluscode/CineMA",
-        model_filename=f"finetuned/segmentation/{trained_dataset}_lax_4c_{seed}.safetensors",
-        config_filename="finetuned/segmentation/lax_4c.yaml",
+        model_filename=f"finetuned/segmentation/{trained_dataset}_{view}/{trained_dataset}_{view}_{seed}.safetensors",
+        config_filename=f"finetuned/segmentation/{trained_dataset}_{view}/config.yaml",
     )
 
     # load sample data and form a batch of size 1
-    transform = ScaleIntensityd(keys="lax_4c")
+    transform = ScaleIntensityd(keys=view)
 
     # (x, y, 1, t)
     exp_dir = Path(__file__).parent.parent.resolve()
     images = np.transpose(sitk.GetArrayFromImage(sitk.ReadImage(exp_dir / "data/ukb/1/1_lax_4c.nii.gz")))
     n_frames = images.shape[-1]
     labels_list = []
     for t in tqdm(range(n_frames), total=n_frames):
-        batch = transform({"lax_4c": torch.from_numpy(images[None, ..., 0, t]).to(dtype=torch.float32)})
+        batch = transform({view: torch.from_numpy(images[None, ..., 0, t]).to(dtype=torch.float32)})
         batch = {k: v[None, ...] for k, v in batch.items()}  # batch size 1
         with torch.no_grad(), torch.autocast("cuda", enabled=torch.cuda.is_available()):
-            logits = model(batch)["lax_4c"]  # (1, 4, x, y)
+            logits = model(batch)[view]  # (1, 4, x, y)
         labels = torch.argmax(logits, dim=1)[0].detach().numpy()  # (x, y)
 
         # the model seems to hallucinate an additional right ventricle and myocardium sometimes
@@ -88,7 +89,7 @@ def run(seed: int) -> None:
             if j == 0:
                 axs[i, j].set_ylabel(f"t = {t}")
     plt.subplots_adjust(wspace=0.02, hspace=0.02)
-    plt.savefig(f"segmentation_lax_4c_mask_{trained_dataset}_{seed}.png", dpi=300, bbox_inches="tight")
+    plt.savefig(f"segmentation_{view}_mask_{trained_dataset}_{seed}.png", dpi=300, bbox_inches="tight")
     plt.show(block=False)
 
     # visualise area changes
@@ -105,7 +106,7 @@ def run(seed: int) -> None:
     plt.ylabel("Area (mm2)")
     plt.title(f"LVEF = {lvef:.2f}%, RVEF = {rvef:.2f}%")
     plt.legend(loc="lower right")
-    plt.savefig(f"segmentation_lax_4c_mask_area_{trained_dataset}_{seed}.png", dpi=300, bbox_inches="tight")
+    plt.savefig(f"segmentation_{view}_mask_area_{trained_dataset}_{seed}.png", dpi=300, bbox_inches="tight")
     plt.show(block=False)
 
 
diff --git a/examples/inference/segmentation_sax.py b/examples/inference/segmentation_sax.py
@@ -15,17 +15,18 @@
 def run(trained_dataset: str, seed: int) -> None:
     """Run segmentation on SAX images using fine-tuned checkpoint."""
     # load model
+    view = "sax"
     model = ConvUNetR.from_finetuned(
         repo_id="mathpluscode/CineMA",
-        model_filename=f"finetuned/segmentation/{trained_dataset}_sax_{seed}.safetensors",
-        config_filename="finetuned/segmentation/sax.yaml",  # same config for all models
+        model_filename=f"finetuned/segmentation/{trained_dataset}_{view}/{trained_dataset}_{view}_{seed}.safetensors",
+        config_filename=f"finetuned/segmentation/{trained_dataset}_{view}/config.yaml",  # same config for all models
     )
 
     # load sample data and form a batch of size 1
     transform = Compose(
         [
-            ScaleIntensityd(keys="sax"),
-            SpatialPadd(keys="sax", spatial_size=(192, 192, 16), method="end", lazy=True, allow_missing_keys=True),
+            ScaleIntensityd(keys=view),
+            SpatialPadd(keys=view, spatial_size=(192, 192, 16), method="end", lazy=True, allow_missing_keys=True),
         ]
     )
 
@@ -35,10 +36,10 @@ def run(trained_dataset: str, seed: int) -> None:
     n_slices, n_frames = images.shape[-2:]
     labels_list = []
     for t in tqdm(range(n_frames), total=n_frames):
-        batch = transform({"sax": torch.from_numpy(images[None, ..., t]).to(dtype=torch.float32)})
+        batch = transform({view: torch.from_numpy(images[None, ..., t]).to(dtype=torch.float32)})
         batch = {k: v[None, ...] for k, v in batch.items()}  # batch size 1
         with torch.no_grad(), torch.autocast("cuda", enabled=torch.cuda.is_available()):
-            logits = model(batch)["sax"]  # (1, 4, x, y, z)
+            logits = model(batch)[view]  # (1, 4, x, y, z)
         labels_list.append(torch.argmax(logits, dim=1)[0, ..., :n_slices])
     labels = torch.stack(labels_list, dim=-1).detach().numpy()  # (x, y, z, t)
 
@@ -57,7 +58,7 @@ def run(trained_dataset: str, seed: int) -> None:
                 axs[i, z].set_ylabel(f"t = {t}")
     axs[0, 4].set_title("SAX Slices")
     plt.subplots_adjust(wspace=0.02, hspace=0.02)
-    plt.savefig(f"segmentation_sax_mask_{trained_dataset}_{seed}.png", dpi=300, bbox_inches="tight")
+    plt.savefig(f"segmentation_{view}_mask_{trained_dataset}_{seed}.png", dpi=300, bbox_inches="tight")
     plt.show(block=False)
 
     # visualise volume changes
@@ -74,7 +75,7 @@ def run(trained_dataset: str, seed: int) -> None:
     plt.ylabel("Volume (ml)")
     plt.title(f"LVEF = {lvef:.2f}%, RVEF = {rvef:.2f}%")
     plt.legend(loc="lower right")
-    plt.savefig(f"segmentation_sax_mask_volume_{trained_dataset}_{seed}.png", dpi=300, bbox_inches="tight")
+    plt.savefig(f"segmentation_{view}_mask_volume_{trained_dataset}_{seed}.png", dpi=300, bbox_inches="tight")
     plt.show(block=False)
 
 

Original file line number	Diff line number	Diff line change
`@@ -629,7 +629,7 @@ def from_pretrained(cls, **kwargs) -> CineMA: # type: ignore[no-untyped-def]`
`629`	`629`	`# download config`
`630`	`630`	`config_path = hf_hub_download(`
`631`	`631`	`repo_id="mathpluscode/CineMA",`
`632`		`- filename="pretrained/cinema.yaml",`
	`632`	`+ filename="pretrained/config.yaml",`
`633`	`633`	`**kwargs,`
`634`	`634`	`)`
`635`	`635`	`logger.info(f"Cached model config to {config_path}.")`