Working on setting up automatic weight downloading

README.md

@@ -1,5 +1,34 @@
 # Installation
 
+## Configure system
+<details>
+<summary>Ubuntu</summary>
+
+First we need to download and install CUDA 11.2:
+```
+wget https://developer.download.nvidia.com/compute/cuda/11.2.2/local_installers/cuda_11.2.2_460.32.03_linux.run
+sudo sh cuda_11.2.2_460.32.03_linux.run --silent --toolkit --installpath=/usr/local/cuda-11.2
+```
+
+Next, cudnn must be installed. Download a suitable deb-file from
+https://developer.nvidia.com/rdp/cudnn-archive. Then install the file:
+```
+sudo dpkg -i ~/Downloads/cudnn-local-repo-ubuntu2204-8.9.7.29_1.0-1_amd64.deb
+sudo cp /var/cudnn-local-repo-ubuntu2204-8.9.7.29/cudnn-local-*-keyring.gpg /usr/share/keyrings/
+sudo apt update
+sudo apt install libcudnn8 libcudnn8-dev
+sudo cp /usr/include/cudnn*.h /usr/local/cuda-11.2/include/
+sudo cp -P /usr/lib/x86_64-linux-gnu/libcudnn*.so* /usr/local/cuda-11.2/lib64/
+sudo ldconfig
+```
+Finally, we must configure the system paths (or add these lines to ~/.bashrc:
+```
+echo 'export CUDA_HOME=/usr/local/cuda-11.2' >> ~/.bashrc
+echo 'export PATH=$CUDA_HOME/bin:$PATH' >> ~/.bashrc
+echo 'export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$CUDA_HOME/extras/CUPTI/lib64' >> ~/.bashrc
+```
+
+</details>
 
 ## Install pyenv and Python
 
@@ -100,4 +129,14 @@ Finally, we can run the preprocessing script, pointing towards the python from t
 ```
 sh scripts/preprocess.sh --license <path-to-license> --python ~/venvs/fastsurfer/bin/python ~/data/ixi/images ~/data/ixi/preprocessed
 ```
+
+### Generate predictions
+After preprocessing, we can generate predictions for the IXI dataset using the scripts in the repository. First, ensure the virtual environment is loaded:
+```
+eval $(poetry env activate)
+```
+Next, run the prediction-script:
+```
+python scripts/predict_from_fastsurfer_folder.py
+```
 </details>

poetry.lock

@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 2.2.0 and should not be changed by hand.
+# This file is automatically @generated by Poetry 2.1.3 and should not be changed by hand.
 
 [[package]]
 name = "absl-py"
@@ -2807,6 +2807,13 @@ optional = false
 python-versions = ">=3.8"
 groups = ["main"]
 files = [
+    {file = "PyYAML-6.0.3-cp38-cp38-macosx_10_13_x86_64.whl", hash = "sha256:c2514fceb77bc5e7a2f7adfaa1feb2fb311607c9cb518dbc378688ec73d8292f"},
+    {file = "PyYAML-6.0.3-cp38-cp38-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9c57bb8c96f6d1808c030b1687b9b5fb476abaa47f0db9c0101f5e9f394e97f4"},
+    {file = "PyYAML-6.0.3-cp38-cp38-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:efd7b85f94a6f21e4932043973a7ba2613b059c4a000551892ac9f1d11f5baf3"},
+    {file = "PyYAML-6.0.3-cp38-cp38-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:22ba7cfcad58ef3ecddc7ed1db3409af68d023b7f940da23c6c2a1890976eda6"},
+    {file = "PyYAML-6.0.3-cp38-cp38-musllinux_1_2_x86_64.whl", hash = "sha256:6344df0d5755a2c9a276d4473ae6b90647e216ab4757f8426893b5dd2ac3f369"},
+    {file = "PyYAML-6.0.3-cp38-cp38-win32.whl", hash = "sha256:3ff07ec89bae51176c0549bc4c63aa6202991da2d9a6129d7aef7f1407d3f295"},
+    {file = "PyYAML-6.0.3-cp38-cp38-win_amd64.whl", hash = "sha256:5cf4e27da7e3fbed4d6c3d8e797387aaad68102272f8f9752883bc32d61cb87b"},
     {file = "pyyaml-6.0.3-cp310-cp310-macosx_10_13_x86_64.whl", hash = "sha256:214ed4befebe12df36bcc8bc2b64b396ca31be9304b8f59e25c11cf94a4c033b"},
     {file = "pyyaml-6.0.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:02ea2dfa234451bbb8772601d7b8e426c2bfa197136796224e50e35a78777956"},
     {file = "pyyaml-6.0.3-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b30236e45cf30d2b8e7b3e85881719e98507abed1011bf463a8fa23e9c3e98a8"},

pyment/models/sfcn/sfcn.py

@@ -79,7 +79,10 @@ class SFCN(Model):
         super().__init__(self.inputs, x)
 
         if weights:
+            print(weights)
             weights = ensure_weights(weights)
+            print(weights)
+            weights = 'checkpoints/pyment/sfcn-multi'
             status = self.load_weights(weights)
 
             print(weights)

pyment/models/utils/ensure_weights.py

@@ -1,15 +1,60 @@
 import os
 
+from pyment.utils.download_file import download_file
 
-def ensure_weights(identifier: str) -> str:
-    """Takes either a path or an identifier for a valid weight 
+
+IDENTIFIERS = {
+    'multi-2025': {
+        'data': 'f4054d701fa59971fb7000d38cf9e63a202bd66a',
+        'index': '9c208ca0bcc3969ceb281ba63a8cee4944a63157'
+    }
+}
+BASE_URL = 'https://api.github.com/repos/estenhl/pyment-public/git/blobs'
+
+def _lookup_identifier(identifier: str, local_cache: str) -> str:
+    if not (
+        os.path.isfile(
+            os.path.join(local_cache, f'{identifier}.index')
+        ) and os.path.isfile(
+            os.path.join(local_cache, f'{identifier}.data-00000-of-00001')
+        )
+    ):
+        if not os.path.isdir(local_cache):
+            os.makedirs(local_cache, exist_ok=True)
+
+        download_file(
+            url=BASE_URL + '/' + IDENTIFIERS[identifier]['data'],
+            destination=os.path.join(
+                local_cache,
+                f'{identifier}.data-00000-of-00001'
+            ),
+            description=f'Downloading {identifier} data'
+        )
+        download_file(
+            url=BASE_URL + '/' + IDENTIFIERS[identifier]['index'],
+            destination=os.path.join(
+                local_cache,
+                f'{identifier}.index'
+            ),
+            description=f'Downloading {identifier} index'
+        )
+
+    return os.path.join(local_cache, identifier)
+
+def ensure_weights(
+    identifier: str,
+    local_cache: str = os.path.join(
+        os.path.expanduser('~'), '.pyment', 'weights'
+    )
+) -> str:
+    """Takes either a path or an identifier for a valid weight
     configuration as an argument, and returns a path-prefix to files
     containing the weights. If necessary, the weights are downloaded.
 
     Parameters
     ----------
     identifier : str
-        Points to either a filename or a valid keyword identifiying a 
+        Points to either a filename or a valid keyword identifiying a
         weight file.
 
     Returns
@@ -21,20 +66,24 @@ def ensure_weights(identifier: str) -> str:
     ------
     KeyError
         If the identifier is not a valid identifier and there does not
-        exist either a single file <identifier> or files 
+        exist either a single file <identifier> or files
         <identifier>.index and <identifier>.data-00000-of-00001 on the
         local file system.
     """
-    if not (
+    if (
         (
-            os.path.isfile(f'{identifier}.index') and 
+            os.path.isfile(f'{identifier}.index') and
             os.path.isfile(f'{identifier}.data-00000-of-00001')
         ) or (
             os.path.isfile(identifier)
         )
     ):
+        return identifier
+    elif identifier in IDENTIFIERS:
+        return _lookup_identifier(identifier, local_cache)
+    else:
         raise NotImplementedError(
             f'Identifier-based lookups are not supported'
         )
 
-    return identifier
+    return identifier

pyment/utils/download_file.py

@@ -0,0 +1,30 @@
+import math
+import requests
+from tqdm import tqdm
+
+
+def download_file(
+    url: str,
+    destination: str,
+    description: str = None
+) -> str:
+    with requests.get(url, stream=True) as response:
+        response.raise_for_status()
+        total_size = int(response.headers.get('content-length', 0))
+
+        # 1 MB chunks
+        chunk_size = 1<<20
+
+        progress_bar = tqdm(
+            response.iter_content(chunk_size=chunk_size),
+            total=int(math.ceil(total_size / chunk_size)),
+            unit='mb',
+            unit_scale=True,
+            unit_divisor=1024,
+            desc=description
+        )
+        progress_bar.format_dict['rate'] = f'mb/s'
+
+        with open(destination, 'wb') as f:
+            for chunk in progress_bar:
+                f.write(chunk)

scripts/predict_from_fastsurfer_folder.py

@@ -28,7 +28,7 @@ def _parse_folder_name(name: str) -> Tuple[str, str, str]:
     return match.groups()
 
 def predict_from_fastsurfer_folder(
-    source: str, 
+    source: str,
     weights: str,
     model_name: str = 'sfcn-multi',
     targets: List[str] = [
@@ -38,7 +38,7 @@ def predict_from_fastsurfer_folder(
 ) -> pd.DataFrame:
     if destination is not None and os.path.isfile(destination):
         raise ValueError(f'Destination {destination} already exists')
-    
+
     logger.info('Loading multi-task model with weights %s', weights)
 
     model_class = sfcn_factory(model_name)
@@ -62,7 +62,7 @@ def predict_from_fastsurfer_folder(
         if not os.path.isfile(brainmask):
             logger.warning('No mask.mgz file for folder %s', folder)
             continue
-        
+
         brainmask = nib.load(brainmask)
         brainmask = brainmask.get_fdata()
 
@@ -72,10 +72,8 @@ def predict_from_fastsurfer_folder(
         image = conform(image)
 
         predictions = model.predict(np.expand_dims(image, axis=0))[0]
-        print(predictions.shape)
-        print(predictions)
         logger.debug('Predictions for %s: %s', folder, str(predictions))
-        
+
         results.append({
             **{
                 'source': os.path.join(source, folder),
@@ -100,7 +98,7 @@ if __name__ == '__main__':
     )
 
     parser.add_argument(
-        'root', 
+        'root',
         help=(
             'Path to FastSurfer folder. Should contain subfolders that have '
             'an \'mri\' subfolder that contains files orig.mgz and mask.mgz'
@@ -108,15 +106,16 @@ if __name__ == '__main__':
     )
     parser.add_argument(
         '-w', '--weights',
-        required=True,
+        required=False,
+        default='multi-2025',
         help=(
             'Weights to use. Should either point to a local file path, or a '
-            'known keyword. If a local file path <path> is used, there should '
+            'known identifier. If a local file path <path> is used, there should '
             'exist files named <path>.index and <path>.data-00000-of-00001'
         )
     )
     parser.add_argument(
-        '-m', '--model', 
+        '-m', '--model',
         required=False,
         default='sfcn-multi',
         help=(
@@ -128,7 +127,7 @@ if __name__ == '__main__':
         required=False,
         nargs='+',
         default=[
-            'age', 'sex', 'handedness', 'bmi', 'fluid_intelligence', 
+            'age', 'sex', 'handedness', 'bmi', 'fluid_intelligence',
             'neuroticism'
         ],
         help='Name to use for each of the prediction heads in the output CSV'

scripts/utils/upload_weights_to_github.py

@@ -0,0 +1,68 @@
+import argparse
+import base64
+import os
+import requests
+
+
+def upload_weights_to_github(filename: str, token: str, user: str, repo: str):
+    with open(filename, 'rb') as f:
+        bytes = f.read()
+
+    bytes = base64.b64encode(bytes).decode()
+
+    if os.path.isfile(token):
+        with open(token, 'r') as f:
+            token = f.read().strip()
+
+    headers = {
+        'Accept': 'application/vnd.github+json',
+        'Authorization': f'Bearer {token}',
+        'X-GitHub-Api-Version': '2022-11-28'
+    }
+    content = {
+        'content': bytes,
+        'encoding': 'base64'
+    }
+
+    url = f'https://api.github.com/repos/{user}/{repo}/git/blobs'
+
+    response = requests.post(url, json=content, headers=headers)
+
+    return response
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser('Uploads a weights-file to github')
+
+    parser.add_argument(
+        '-f', '--filename',
+        required=True,
+        help='Path to file containing weights'
+    )
+    parser.add_argument(
+        '-t', '--token',
+        required=True,
+        help='Token for the GitHub API'
+    )
+    parser.add_argument(
+        '-u', '--user',
+        required=False,
+        default='estenhl',
+        help='Owner of the github repo'
+    )
+    parser.add_argument(
+        '-r', '--repo',
+        required=False,
+        default='pyment-public',
+        help='Name of the github repo'
+    )
+
+    args = parser.parse_args()
+
+    response = upload_weights_to_github(
+        filename=args.filename,
+        token=args.token,
+        user=args.user,
+        repo=args.repo
+    )
+
+    print(response.text)

tutorials/download_ixi.py

@@ -5,34 +5,10 @@ import requests
 import tarfile
 from tqdm import tqdm
 
+from pyment.utils.download_file import download_file
 
-DEFAULT_DESTINATION = os.path.join(os.path.expanduser('~'), 'data', 'ixi')
-
-def download_file(
-    url: str,
-    destination: str,
-    description: str = None
-) -> str:
-    with requests.get(url, stream=True) as response:
-        response.raise_for_status()
-        total_size = int(response.headers.get('content-length', 0))
-
-        # 1 MB chunks
-        chunk_size = 1<<20
 
-        progress_bar = tqdm(
-            response.iter_content(chunk_size=chunk_size),
-            total=int(math.ceil(total_size / chunk_size)),
-            unit='mb',
-            unit_scale=True,
-            unit_divisor=1024,
-            desc=description
-        )
-        progress_bar.format_dict['rate'] = f'mb/s'
-
-        with open(destination, 'wb') as f:
-            for chunk in progress_bar:
-                f.write(chunk)
+DEFAULT_DESTINATION = os.path.join(os.path.expanduser('~'), 'data', 'ixi')
 
 def download_tar(tar_path: str) -> str:
     url = (
@@ -82,4 +58,4 @@ if __name__ == '__main__':
 
     args = parser.parse_args()
 
-    download_ixi(args.destination)
+    download_ixi(args.destination)