colab

.gitignore

@@ -12,7 +12,6 @@ toy.py
 tmp
 core
 *.pkl
-*.pth
 *.pt
 *.npy
 *.npz

brainnet/plmodel.py

@@ -16,16 +16,19 @@
 
 
 class PLModel(pl.LightningModule):
-    def __init__(self, config, backbone, cached=True, draw=True):
+    def __init__(self, config, backbone, cached=True, draw=True, skip_data=False):
         super().__init__()
         self.config = config
         self.draw = draw
 
         # dataset
-        self.dataset = BrainDataset(config.DATASET.DATA_DIR, config.DATASET.RESOLUTION)
-        self.train_dataset, self.val_dataset = torch.utils.data.random_split(
-            self.dataset, [9000, 841]
-        )
+        if not skip_data:
+            self.dataset = BrainDataset(
+                config.DATASET.DATA_DIR, config.DATASET.RESOLUTION
+            )
+            self.train_dataset, self.val_dataset = torch.utils.data.random_split(
+                self.dataset, [9000, 841]
+            )
 
         # coordinates
         self.coords = load_coords()[nsdgeneral_indices]
@@ -249,7 +252,7 @@ def get_top_channels(self, x):
         for roi in roi_names:
             # roi_indices
             roi_indices = rh_roi_dict[roi]
-            fsaverage = np.zeros(327684)   # convert fsaverage and nsdgeneral space
+            fsaverage = np.zeros(327684)  # convert fsaverage and nsdgeneral space
             fsaverage[nsdgeneral_indices] = np.arange(nsdgeneral_indices.shape[0])
             roi_indices = fsaverage[roi_indices]
             # pca of weights

conda_env.yml

@@ -33,6 +33,7 @@ dependencies:
     - scipy
     - fast-pytorch-kmeans
     - einops
+    - colorstamps
     - open_clip_torch==2.20.0
     - git+https://github.com/facebookresearch/segment-anything.git@6fdee8f2727f4506cfbbe553e23b895e27956588
     - git+https://github.com/facebookresearch/dinov2.git@ebc1cba1096de0a5388527817cf1e5d4d3896699