Code changes

README.md

@@ -48,9 +48,11 @@ nohup python train_brain2vec_PCA.py \
 
 # model inference
 python inference_brain2vec_PCA.py \
-    --pca_model /path/to/pca_model.joblib \
+    --pca_model pca_model.joblib \
     --input_images /path/to/img1.nii.gz /path/to/img2.nii.gz \
-    --output_dir /path/to/out
+    --output_dir pca_output \
+    --embeddings_filename pca_output/pca_embeddings_2 \
+    --save_recons
     
 # or if you have a CSV with image paths:
 python inference_brain2vec_PCA.py \

inference_brain2vec_PCA.py

@@ -163,6 +163,18 @@ def main() -> None:
         "--csv_input", type=str, default=None,
         help="Path to a CSV containing column 'image_path'."
     )
+    parser.add_argument(
+        "--embeddings_filename", 
+        type=str, 
+        required=True,
+        help="Filename (without path) to save the stacked embeddings (e.g., 'pca_embeddings.npy')."
+    )
+    parser.add_argument(
+        "--save_recons",
+        action="store_true",
+        help="If set, save each reconstruction as .npy. Default is not to save."
+    )
+
     args = parser.parse_args()
 
     os.makedirs(args.output_dir, exist_ok=True)
@@ -205,15 +217,18 @@ def main() -> None:
         all_embeddings.append(embedding_np)
 
         # Optionally save or visualize reconstructions
-        out_recon_path = os.path.join(
-            args.output_dir, f"reconstruction_{i}.npy"
-        )
-        np.save(out_recon_path, recon_np)
-        print(f"[INFO] Saved reconstruction to: {out_recon_path}")
+        if args.save_recons:
+            out_recon_path = os.path.join(args.output_dir, f"reconstruction_{i}.npy")
+            np.save(out_recon_path, recon_np)
+            print(f"[INFO] Saved reconstruction to: {out_recon_path}")
 
     # Save all embeddings stacked
     stacked_embeddings = np.vstack(all_embeddings)  # (N, n_components)
-    out_embed_path = os.path.join(args.output_dir, "all_pca_embeddings.npy")
+    filename = args.embeddings_filename
+    if not filename.lower().endswith(".npy"):
+        filename += ".npy"
+
+    out_embed_path = os.path.join(args.output_dir, filename)
     np.save(out_embed_path, stacked_embeddings)
     print(f"[INFO] Saved embeddings of shape {stacked_embeddings.shape} to: {out_embed_path}")
 

train_brain2vec_PCA.py

@@ -276,7 +276,7 @@ def main() -> None:
     )
 
     # Fit the PCA model
-    print(f"Fitting {args.pca_type.capitalize()}PCA in batches...")
+    print(f"Fitting {args.pca_type.capitalize()}PCA")
     model.fit(X)
     print("Done fitting PCA. Transforming data to embeddings...")