Working on docker containers

README.md

@@ -1,9 +1,24 @@
 ## Build docker container for preprocessing
+Note that for now, building the container requires a folder called <checkpoints> that contains the FastSurfer segmentation checkpoints
+
+```
 docker build \
     -f docker/preprocess.Dockerfile \
     -t pyment/preprocessing:1.0.0 \
-    --build-arg CHECKPOINTS_FOLDER=<path_to_fastsurfer_checkpoints>
     .
+```
 
 ## Run docker container for preprocessing
-docker run --rm -v <path_to_input>:/input -v <path_to_ouput>:/output -v <path_to_licenses>:/licenses --gpus all pyment/preprocessing:1.0.0
+Running the container for preprocessing requires three volumes:
+- Inputs: A folder containing input data. All nifti-files detected in this folder or one of its subfolders will be processed
+- Outputs: A folder where the preprocessed images will be written.
+- Licenses: A folder containing the freesurfer license
+```
+docker run --rm \
+    --user $(id -u):$(id -g) \
+    --volume <path_to_input>:/input \
+    --volume <path_to_ouput>:/output \
+    --volume <path_to_licenses>:/licenses \
+    --gpus all \
+    pyment/preprocessing:1.0.0
+```

docker/preprocess.Dockerfile

@@ -31,4 +31,4 @@ COPY scripts/preprocess.sh /scripts/preprocess.sh
 CMD ["/bin/sh", "/scripts/preprocess.sh", \
      "--license", "/licenses/freesurfer.txt", \
      "--python", "/envs/fastsurfer/bin/python", \
-     "/input", "/output"]
+     "/input", "/output/fastsurfer"]

pyment/__init__.py

@@ -0,0 +1,15 @@
+import os
+import tomllib
+
+def _get_version():
+    """Get version from pyproject.toml"""
+    pyproject_path = os.path.join(
+        os.path.dirname(__file__), os.pardir, 'pyproject.toml'
+    )
+
+    with open(pyproject_path, 'rb') as f:
+        data = tomllib.load(f)
+
+    return data['project']['version']
+
+__version__ = _get_version()

pyment/models/__init__.py

@@ -0,0 +1,3 @@
+from .sfcn import SFCN, MultiTaskSFCN
+
+__all__ = ['SFCN', 'MultiTaskSFCN']

pyment/models/sfcn/__init__.py

@@ -0,0 +1,4 @@
+from .sfcn import SFCN
+from .sfcn_multi import MultiTaskSFCN
+
+__all__ = ['SFCN', 'MultiTaskSFCN']

pyment/models/sfcn/sfcn.py

@@ -0,0 +1,76 @@
+from abc import abstractmethod
+from typing import Any, Callable, Tuple
+
+import tensorflow as tf
+from tensorflow.keras.layers import (
+    Activation, BatchNormalization, Conv3D, Dropout, GlobalAveragePooling3D,
+    GlobalMaxPooling3D, Input, MaxPooling3D, Reshape
+)
+from tensorflow.keras.models import Model
+from tensorflow.keras.regularizers import Regularizer
+
+
+
+class SFCN(Model):
+    FILTERS = [32, 64, 128, 256, 256, 64]
+
+    @abstractmethod
+    def construct_prediction_head(
+        bottleneck: tf.Tensor,
+        name: str
+    ) -> tf.Tensor:
+        pass
+
+    def __init__(self, *,
+        input_shape: Tuple[int] = (224, 192, 224),
+        include_top: bool = True,
+        pooling: str = 'avg',
+        dropout: float = 0.,
+        regularizer: Regularizer = None,
+        activation: Any = 'relu',
+        name: str = 'SFCN'
+    ):
+        self.inputs = Input(input_shape, name=f'{name}_inputs')
+        x = Reshape(input_shape + (1,), name=f'{name}_expand-dims')(
+            self.inputs
+        )
+
+        for i in range(len(self.FILTERS) - 1):
+            x = Conv3D(
+                self.FILTERS[i],
+                (3, 3, 3),
+                padding='SAME',
+                activation=None,
+                kernel_regularizer=regularizer,
+                name=f'{name}_block-{i}_conv'
+            )(x)
+            x = BatchNormalization(name=f'{name}_block-{i}_norm')(x)
+            x = Activation(activation, name=f'{name}_block-{i}_activation')(x)
+            x = MaxPooling3D((2, 2, 2), name=f'{name}_block-{i}_pool')(x)
+
+        x = Conv3D(
+            self.FILTERS[-1],
+            (1, 1, 1),
+            padding='SAME',
+            activation=None,
+            name=f'{name}_top_conv'
+        )(x)
+        x = BatchNormalization(name=f'{name}_top_norm')(x)
+        x = Activation(activation, name=f'{name}_top_activation')(x)
+
+        if pooling == 'avg':
+            pooling = GlobalAveragePooling3D
+        elif pooling == 'max':
+            pooling = GlobalMaxPooling3D
+        else:
+            raise ValueError(f'Unknown pooling layer {pooling}')
+
+        x = pooling(name=f'{name}_top_pool')(x)
+
+        self.bottleneck = x
+
+        if include_top:
+            x = Dropout(dropout, name=f'{name}_dropout')(x)
+            x = self.construct_prediction_head(x, name=f'{name}_predictions')
+
+        super().__init__(self.inputs, x)

pyment/models/sfcn/sfcn_multi.py

@@ -0,0 +1,42 @@
+from tensorflow import concat, Tensor
+from tensorflow.keras.layers import Activation, BatchNormalization, Dense
+
+from .sfcn import SFCN
+
+
+class MultiTaskSFCN(SFCN):
+    @classmethod
+    def construct_prediction_head(
+        cls,
+        bottleneck: Tensor,
+        name: str
+    ) -> Tensor:
+        x = bottleneck
+        depths = []
+
+        for i in range(len(depths)):
+            x = Dense(depths[i], activation=None,
+                      name=f'{name}_dense-{i+1}_dense')(x)
+            x = BatchNormalization(name=f'{name}_dense-{i+1}_norm')(x)
+            x = Activation('relu', name=f'{name}_dense-{i+1}_activation')(x)
+
+        heads = [
+            Dense(1, activation=None, name=f'{name}_predictions_age'),
+            Dense(1, activation='sigmoid', name=f'{name}_predictions_sex'),
+            Dense(1, activation='sigmoid',
+                  name=f'{name}_predictions_handedness'),
+            Dense(1, activation=None, name=f'{name}_predictions_bmi'),
+            Dense(1, activation=None,
+                  name=f'{name}_predictions_fluid_intelligence'),
+            Dense(1, activation=None, name=f'{name}_predictions_neuroticism')
+        ]
+        heads = [head(x) for head in heads]
+
+        return concat(heads, axis=-1)
+
+    def __init__(self, *args,
+        pooling: str = 'max',
+        name: str = 'SFCNMulti',
+        **kwargs
+    ):
+        super().__init__(*args, pooling=pooling, name=name, **kwargs)

pyproject.toml

@@ -12,6 +12,15 @@ license-files = [
 readme = "README.md"
 requires-python = ">=3.10.2"
 
+[tool.poetry]
+packages = [{include = "pyment"}]
+
+[tool.poetry.dependencies]
+python = "^3.10.2"
+tensorflow = "^2.15.0"
+numpy = "^1.24.0"
+protobuf = "^4.25.0"
+
 [build-system]
 requires = ["poetry-core>=2.0.0,<3.0.0"]
 build-backend = "poetry.core.masonry.api"

scripts/predict_from_fastsurfer_folder.py

@@ -0,0 +1,27 @@
+import argparse
+
+from pyment.models import MultiTaskSFCN
+
+def predict_from_fastsurfer_folder(folder: str, output: str, weights: str):
+    model = MultiTaskSFCN()
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        'Generates predictions from SFCN multi and a given weights file for '
+        'all images in a fastsurfer folder'
+    )
+
+    parser.add_argument('folder', help='Path to fastsurfer folder')
+    parser.add_argument(
+        'output', help='Path where CSV with predictions are written'
+    )
+    parser.add_argument('-w', '--weights', help='Path to weights file')
+
+    args = parser.parse_args()
+
+    predict_from_fastsurfer_folder(
+        folder=args.folder,
+        output=args.output,
+        weights=args.weights
+    )
+