License change to apache-2.0

README.md

@@ -1,5 +1,5 @@
 ---
-license: mit
+license: apache-2.0
 language:
   - en
 task_categories:
@@ -52,11 +52,13 @@ nohup python brain2vec.py train \
   --n_epochs    10 \
 > train_log.txt 2>&1 &
 
-# run model inference to create *_embeddings.npz files
-python brain2vec.py infererence \
-  --dataset_csv home/ubuntu/brain2vec/inputs.csv \
-  --aekl_ckpt /home/ubuntu/brain2vec/autoencoder_final.pth \
-  --output_dir /home/ubuntu/brain2vec
+# model inference
+python inference_brain2vec.py \
+  --checkpoint_path /path/to/model.pth \
+  --input_images /path/to/img1.nii.gz /path/to/img2.nii.gz \
+  --output_dir ./vae_inference_outputs \
+  --embeddings_filename pca_output/pca_embeddings_2.npy \
+  --save_recons
 ```
 
 # Methods

inference_brain2vec.py

@@ -156,6 +156,18 @@ def main() -> None:
         "--csv_input", type=str,
         help="Path to a CSV file with an 'image_path' column."
     )
+    parser.add_argument(
+        "--embeddings_filename",
+        type=str,
+        required=True,
+        help="Filename (in output_dir) to save the stacked z_mu embeddings (e.g. 'all_z_mu.npy')."
+    )
+    parser.add_argument(
+        "--save_recons",
+        action="store_true",
+        help="If set, saves each reconstruction as .npy. Default is not to save."
+    )
+
     args = parser.parse_args()
 
     os.makedirs(args.output_dir, exist_ok=True)
@@ -198,6 +210,7 @@ def main() -> None:
         z_sigma_np = z_sigma.detach().cpu().numpy()
 
         # Save each reconstruction (per image) as .npy
+    if args.save_recons:
         recon_path = os.path.join(args.output_dir, f"reconstruction_{i}.npy")
         np.save(recon_path, recon_np)
         print(f"[INFO] Saved reconstruction to {recon_path}")
@@ -210,8 +223,13 @@ def main() -> None:
     stacked_mu = np.concatenate(all_z_mu, axis=0)       # e.g., shape (N, latent_channels, ...)
     stacked_sigma = np.concatenate(all_z_sigma, axis=0) # e.g., shape (N, latent_channels, ...)
 
-    mu_path = os.path.join(args.output_dir, "all_z_mu.npy")
+    mu_filename = args.embeddings_filename
+    if not mu_filename.lower().endswith(".npy"):
+        mu_filename += ".npy"
+
+    mu_path = os.path.join(args.output_dir, mu_filename)
     sigma_path = os.path.join(args.output_dir, "all_z_sigma.npy")
+
     np.save(mu_path, stacked_mu)
     np.save(sigma_path, stacked_sigma)