Finished first version of scripts necessary to finetune models

new_packages.txt

@@ -13,4 +13,5 @@ tqdm==4.66.4
 plotly==5.24.1
 pytest==8.3.3
 scikit-learn==1.5.1
-xlrd==2.0.1
+xlrd==2.0.1
+pydantic==2.10

pyment/configurations/__init__.py

@@ -0,0 +1,2 @@
+from .dataset_configuration import DatasetConfiguration
+from .finetuning_configuration import FinetuningConfiguration

pyment/configurations/data_split_configuration.py

@@ -0,0 +1,8 @@
+from typing import List
+
+from pydantic import BaseModel
+
+
+class DataSplitConfiguration(BaseModel):
+    training_fraction: float
+    stratification: List[str] = None

pyment/configurations/dataset_configuration.py

@@ -0,0 +1,287 @@
+from __future__ import annotations
+
+import logging
+import os
+import re
+import numpy as np
+import pandas as pd
+from collections import Counter
+from typing import Dict, List, Tuple, Union
+
+import nibabel as nib
+from pydantic import model_validator, BaseModel
+
+from .data_split_configuration import DataSplitConfiguration
+
+
+logging.basicConfig(
+    format='%(asctime)s - %(levelname)s - %(name)s: %(message)s',
+    level=logging.INFO
+)
+logger = logging.getLogger(__name__)
+
+def _extract_run(filename: str) -> Union[str, None]:
+    match = re.fullmatch(r'.*_run-(?P<run>[^_.]*)(?:_.*)?\.mgz', filename)
+
+    if match:
+        return match.group('run')
+
+    logger.warning('Unable to extract run from filename %s', filename)
+
+    return None
+
+def _parse_bids_folder(root: str):
+    entries = []
+
+    for subject_folder in os.listdir(root):
+        subject_match = re.fullmatch(r'sub-(?P<subject>.*)', subject_folder)
+        
+        if not subject_match:
+            logger.warning(
+                'Subject folder %s in %s does not have the expected sub-XXX '
+                'format. Skipping', subject_folder, root
+            )
+            continue
+
+        subject = subject_match.group('subject')
+
+        for session_folder in os.listdir(os.path.join(root, subject_folder)):
+            session_match = re.fullmatch(
+                r'ses-(?P<session>.*)', session_folder
+            )
+
+            if not session_match:
+                logger.warning(
+                    'Session folder %s in subject %s in folder %s does not '
+                    'match the expected ses-XXX format. Skipping', 
+                    session_folder, subject_folder, root
+                )
+                continue
+
+            session = session_match.group('session')
+
+            anat_folder = os.path.join(
+                root, subject_folder, session_folder, 'anat'
+            )
+
+            t1s = [
+                filename for filename in os.listdir(anat_folder) 
+                if 'T1' in filename
+            ]
+
+            for filename in t1s:
+                run = _extract_run(filename)
+                entries.append({
+                    'subject': subject,
+                    'session': session,
+                    'run': run,
+                    'path': os.path.join(anat_folder, filename)
+                })
+
+    return pd.DataFrame(entries, columns=['subject', 'session', 'run', 'path'])
+
+def _parse_bids_folders(folders: List[str]):
+    df = pd.concat([_parse_bids_folder(folder) for folder in folders])
+    df = df.reset_index()
+    logger.info('Parsed %d images', len(df))
+
+    return df
+
+def _parse_fastsurfer_name(name: str) -> Tuple[str, str, str]:
+    match = re.fullmatch(r'sub-(.*)_ses-(.*)_run-(.*)(?:T1w?)?', name)
+
+    if not match:
+        raise ValueError(
+            'Unable to extract subject, session, run from folder %s', name
+        )
+
+    return match.groups()
+
+def _parse_fastsurfer_folder(folder: str):
+    entries = []
+
+    for subfolder in os.listdir(folder):
+        subject, session, run = _parse_fastsurfer_name(subfolder)
+
+        mri_folder = os.path.join(folder, subfolder, 'mri')
+        brainmask = os.path.join(mri_folder, 'brainmask.mgz')
+
+        if not os.path.isfile(brainmask):
+            logger.info('Brainmask does not exist in folder %s', subfolder)
+
+            orig = os.path.join(mri_folder, 'orig.mgz')
+            mask = os.path.join(mri_folder, 'mask.mgz')
+
+            if not os.path.isfile(orig):
+                logger.error('Orig does not exist in folder %s', subfolder)
+                continue
+            elif not os.path.isfile(mask):
+                logger.error('Mask does not exist in folder %s', subfolder)
+                continue
+
+            orig_data = nib.load(orig)
+            mask_data = nib.load(mask)
+            brainmask_data = nib.Nifti1Image(
+                orig_data.get_fdata() * mask_data.get_fdata(), 
+                header=orig_data.header, 
+                affine=orig_data.affine
+            )
+
+            nib.save(brainmask_data, brainmask)
+        
+        entries.append({
+            'subject': subject,
+            'session': session,
+            'run': run,
+            'path': brainmask
+        })
+
+    return pd.DataFrame(entries, columns=['subject', 'session', 'run', 'path'])
+
+def _parse_fastsurfer_folders(folders: List[str]):
+    df = pd.concat([_parse_fastsurfer_folder(folder) for folder in folders])
+    df = df.reset_index()
+    logger.info('Parsed %d images', len(df))
+
+    return df
+
+def _summarize_values(values: np.ndarray, name: str):
+    if not np.issubdtype(values.dtype, np.number):
+        logger.info('%s: %s', name, Counter(values))
+    elif np.array_equal(
+        np.unique(values[~np.isnan(values)]), 
+        np.asarray([0, 1])
+    ):
+        nans = len(np.where(np.isnan(values))[0])
+        logger.info(
+            '%s: %s (%d NAs)', name, Counter(values[~np.isnan(values)]), nans
+        )
+    else:
+        nans = len(np.where(np.isnan(values))[0])
+        mean = np.round(np.nanmean(values), 2)
+        std = np.round(np.nanstd(values), 2)
+        logger.info('%s: %.2f+/-%.2f (%d NAs)', name, mean, std, nans)
+
+def _summarize(df: pd.DataFrame, variables: List[str], name: str):
+    logger.info('%s n=%d', name, len(df))
+
+    for variable in variables:
+        _summarize_values(df[variable].values, name=variable)
+
+
+def _split_training_validation_fold(
+    df: pd.DataFrame, 
+    labels: str, 
+    training_fraction: float,
+    target: str = None,
+    stratification: List[str] = None
+) -> Tuple[pd.DataFrame, pd.DataFrame]:
+
+    columns = set(['subject', 'session', 'run'])
+    
+    if target:
+        columns.add(target)
+    
+    if stratification:
+        columns |= set(stratification)
+    
+    labels = pd.read_csv(
+        labels,
+        usecols=list(columns),
+        dtype={'subject': object, 'session': object, 'run': object},
+    )
+
+    logger.info('Parsed %d labels', len(labels))
+
+    if not len(labels) == len(labels.drop_duplicates(['subject', 'session'])):
+        raise ValueError(
+            f'There are duplicates (subject, session)-pairs in the labels file'
+        )
+
+    df = pd.merge(
+       df, labels, 
+       how='inner', 
+       left_on=['subject', 'session'],
+       right_on=['subject', 'session']
+    )
+
+    logger.info('Merged %d data points', len(df))
+
+    if stratification is not None:
+        df = df.sort_values(stratification)
+
+    subjects = df.drop_duplicates('subject')
+    num_folds = int(1.0 / (1 - training_fraction))
+
+    if num_folds == 1:
+        raise ValueError(
+            'Training fraction %.2f yields a single fold', training_fraction
+        )
+    
+    subjects['fold'] = np.arange(len(df)) % num_folds
+    folds = {row['subject']: row['fold'] for _, row in subjects.iterrows()}
+    df['fold'] = df['subject'].map(folds)
+
+    validation_fold = num_folds // 2
+    training = df[df['fold'] != validation_fold]
+    validation = df[df['fold'] == validation_fold]
+
+    if len(
+        set(training['subject'].values) & set(validation['subject'].values)
+    ) > 0:
+        raise ValueError('Overlap between training and validation folds')
+
+    if stratification:
+        for name, df in [('Training', training), ('Validation', validation)]:
+            _summarize(df, variables=stratification, name=name)
+    
+    return training, validation
+
+class DatasetConfiguration(BaseModel):
+    input_shape: Tuple[int, int, int]
+    bids: List[str] | None = None
+    fastsurfer: List[str] | None = None
+    labels: str
+    split: DataSplitConfiguration = None
+
+    @model_validator(mode='after')
+    def check_fastsurfer_or_bids(self):
+        if self.bids is not None and self.fastsurfer is not None:
+            raise ValueError(
+                'Either \'bids\' or \'fastsurfer\'-property must be set, not '
+                'both'
+            )
+        elif self.bids is None and self.fastsurfer is None:
+            raise ValueError(
+                'Either \'bids or \'fastsurfer\'-property must be set'
+            )
+
+        return self
+
+    @staticmethod
+    def parse(
+        configuration: DatasetConfiguration, 
+        target: str = None
+    ) -> Dict[str, pd.DataFrame]:
+        if configuration.split:
+            if configuration.bids:
+                df = _parse_bids_folders(configuration.bids)
+            elif configuration.fastsurfer:
+                df = _parse_fastsurfer_folders(configuration.fastsurfer)
+            else:
+                raise ValueError(
+                    'Unable to parse DatasetConfiguration without either '
+                    '\'bids\' or \'fastsurfer\' set'
+                )
+            
+            return _split_training_validation_fold(
+                df=df,
+                labels=configuration.labels,
+                training_fraction=configuration.split.training_fraction,
+                target=target,
+                stratification=configuration.split.stratification
+            )
+
+        raise NotImplementedError(
+            f'Not sure how to parse dataset without a split configuration'
+        )

pyment/configurations/finetuning_configuration.py

@@ -0,0 +1,11 @@
+from pydantic import BaseModel
+
+from .dataset_configuration import DatasetConfiguration
+from .model_configuration import ModelConfiguration
+from .training_configuration import TrainingConfiguration
+
+
+class FinetuningConfiguration(BaseModel):
+    model: ModelConfiguration
+    data: DatasetConfiguration
+    training: TrainingConfiguration

pyment/configurations/learning_rate_schedule_configuration.py

@@ -0,0 +1,45 @@
+from abc import abstractmethod
+from typing import Annotated, Literal, Union
+
+from tensorflow.keras.callbacks import Callback, ReduceLROnPlateau
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class LearningRateScheduleBaseConfiguration(BaseModel):
+    model_config = ConfigDict(extra='forbid')
+
+    @abstractmethod
+    def instantiate(self) -> Callback:
+        pass
+
+class AnnealingLearningRateScheduleConfiguration(
+    LearningRateScheduleBaseConfiguration
+):
+    kind: Literal['annealing']
+    factor: float
+    patience: int
+    minimum_learning_rate: float
+
+    def instantiate(self) -> Callback:
+        return ReduceLROnPlateau(
+            factor=self.factor,
+            patience=self.patience,
+            min_lr=self.minimum_learning_rate,
+            verbose=True
+        )
+
+class StepWiseLearningRateScheduleConfiguration(
+    LearningRateScheduleBaseConfiguration
+):
+    kind: Literal['stepwise']
+
+    def instantiate(self) -> Callback:
+        return ReduceLROnPl
+
+LearningRateScheduleConfiguration = Annotated[
+    Union[
+        AnnealingLearningRateScheduleConfiguration, 
+        StepWiseLearningRateScheduleConfiguration
+    ], 
+    Field(discriminator='kind')
+]

pyment/configurations/model_configuration.py

@@ -0,0 +1,9 @@
+from typing import Any, Dict
+
+from pydantic import BaseModel, Field
+
+
+class ModelConfiguration(BaseModel):
+    type: str
+    hyperparameters: Dict[str, Any] = Field(default_factory=dict)
+    weights: str = None

pyment/configurations/training_configuration.py

@@ -0,0 +1,19 @@
+from typing import List
+
+from pydantic import BaseModel
+
+from .learning_rate_schedule_configuration import (
+    LearningRateScheduleConfiguration
+)
+
+
+class TrainingConfiguration(BaseModel):
+    target: str
+    loss: str
+    metrics: List[str] = None
+    optimizer: str
+    learning_rate: float
+    learning_rate_schedule: LearningRateScheduleConfiguration = None
+    batch_size: int
+    epochs: int
+    destination: str = None

pyment/factories/__init__.py

@@ -0,0 +1,3 @@
+from .loss_factory import loss_factory
+from .metric_factory import metric_factory
+from .optimizer_factory import optimizer_factory

pyment/factories/loss_factory.py

@@ -0,0 +1,10 @@
+from typing import Callable
+
+import tensorflow as tf
+
+
+def loss_factory(name: str) -> Callable[[tf.Tensor, tf.Tensor], tf.Tensor]:
+    if name.lower() == 'mse':
+        return tf.keras.losses.MeanSquaredError
+
+    raise KeyError(f'Unknown loss {name}')

pyment/factories/metric_factory.py

@@ -0,0 +1,8 @@
+from typing import Callable
+
+import tensorflow as tf 
+
+
+def metric_factory(name: str) -> tf.keras.metrics.Metric:
+    if name.lower() == 'mae':
+        return tf.keras.metrics.MeanAbsoluteError()

pyment/factories/optimizer_factory.py

@@ -0,0 +1,8 @@
+import tensorflow as tf
+
+
+def optimizer_factory(name: str) -> tf.optimizers.Optimizer:
+    if name.lower() == 'adam':
+        return tf.optimizers.Adam
+
+    raise KeyError(f'Unknown optimizer {name}')

pyment/models/sfcn/__init__.py

@@ -1,4 +1,12 @@
 from .sfcn import SFCN
 from .sfcn_multi import MultiTaskSFCN
+from .sfcn_reg import RegressionSFCN
 
-__all__ = ['SFCN', 'MultiTaskSFCN']
+
+def sfcn_factory(model_type: str):
+    if model_type in ['sfcn-reg', 'regression']:
+        return RegressionSFCN
+
+    raise ValueError(f'Unknown SFCN type {model_type}')
+
+__all__ = ['sfcn_factory', 'SFCN', 'MultiTaskSFCN', 'RegressionSFCN']

pyment/models/sfcn/sfcn.py

@@ -82,6 +82,8 @@ class SFCN(Model):
             weights = ensure_weights(weights)
             status = self.load_weights(weights)
 
-            # Silences warnings about optimizer-status not being loaded
-            status.expect_partial()
-            status.assert_existing_objects_matched()
+            print(weights)
+            if not weights.endswith('hdf5'):
+                # Silences warnings about optimizer-status not being loaded
+                status.expect_partial()
+                status.assert_existing_objects_matched()

pyment/models/sfcn/sfcn_multi.py

@@ -7,8 +7,8 @@ from .sfcn import SFCN
 class MultiTaskSFCN(SFCN):
     @classmethod
     def construct_prediction_head(
-        cls,
-        bottleneck: Tensor,
+        cls, 
+        bottleneck: Tensor, 
         name: str
     ) -> Tensor:
         x = bottleneck

pyment/models/sfcn/sfcn_reg.py

@@ -0,0 +1,16 @@
+from tensorflow import Tensor
+from tensorflow.keras.layers import Dense
+
+from .sfcn import SFCN
+
+
+class RegressionSFCN(SFCN):
+    @classmethod
+    def construct_prediction_head(
+        cls, 
+        bottleneck: Tensor, 
+        name: str
+    ) -> Tensor:
+        layer = Dense(1, activation=None, name=f'{name}/predictions')
+
+        return layer(bottleneck)

pyment/models/utils/ensure_weights.py

@@ -21,12 +21,17 @@ def ensure_weights(identifier: str) -> str:
     ------
     KeyError
         If the identifier is not a valid identifier and there does not
-        exist files <identifier>.index and
-        <identifier>.data-00000-of-00001 on the local file system.
+        exist either a single file <identifier> or files 
+        <identifier>.index and <identifier>.data-00000-of-00001 on the
+        local file system.
     """
     if not (
-        os.path.isfile(f'{identifier}.index') and 
-        os.path.isfile(f'{identifier}.data-00000-of-00001')
+        (
+            os.path.isfile(f'{identifier}.index') and 
+            os.path.isfile(f'{identifier}.data-00000-of-00001')
+        ) or (
+            os.path.isfile(identifier)
+        )
     ):
         raise NotImplementedError(
             f'Identifier-based lookups are not supported'

pyment/models/utils/load_select_pretrained_weights.py

@@ -0,0 +1,46 @@
+import logging
+
+import tensorflow as tf
+
+from ..sfcn import MultiTaskSFCN
+
+
+logging.basicConfig(
+    format='%(asctime)s - %(levelname)s - %(name)s: %(message)s',
+    level=logging.DEBUG
+)
+logger = logging.getLogger(__name__)
+
+def load_select_pretrained_weights(
+    model: tf.keras.Model, 
+    weights: str, 
+    target: str = None
+) -> tf.keras.Model:
+    logger.info('Loading pretrained weights from %s', weights)
+
+    backbone = MultiTaskSFCN(input_shape=(224, 192, 224), pooling='max')
+    checkpoint = tf.train.Checkpoint(backbone)
+
+    checkpoint.restore(weights).expect_partial()
+
+    conv_layers = [2, 6, 10, 14, 18, 22]
+    norm_layers = [3, 7, 11, 15, 19, 23]
+
+    for idx in conv_layers + norm_layers:
+        model.layers[idx].set_weights(backbone.layers[idx].get_weights())
+
+    # Loading weights from the specific dense-layer corresponding to the 
+    # given prediction-task in the multi-task model
+    if target == 'age':
+        logger.info('Loaded age weights for the prediction head')
+        model.layers[27].set_weights(backbone.layers[27].get_weights())
+    elif target == 'sex':
+        logger.info('Loaded sex weights for the prediction head')
+        model.layers[27].set_weights(backbone.layers[28].get_weights())
+    else:
+        logger.warning(
+            'Unknown target %s. Not loading weights for prediction layer',
+            target
+        )
+
+    return model

pyment/utils/json_serialize.py

@@ -0,0 +1,16 @@
+import numpy as np
+from typing import Any
+
+def json_serialize(obj: Any) -> Any:
+    if isinstance(obj, dict):
+        return {json_serialize(k): json_serialize(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [json_serialize(v) for v in obj]
+    elif isinstance(obj, (np.integer,)):
+        return int(obj)
+    elif isinstance(obj, (np.floating,)):
+        return float(obj)
+    elif isinstance(obj, (np.ndarray,)):
+        return obj.tolist()
+    else:
+        return obj

scripts/finetune_from_bids_folder.py

@@ -0,0 +1,216 @@
+import argparse
+import json
+import logging
+import os
+import pandas as pd
+from typing import Any, Callable, Dict, List, Tuple
+
+import tensorflow as tf
+from tensorflow_neuroimaging.preprocessing import center_crop_or_pad
+from tensorflow_neuroimaging.loaders.mgh import load_mgh
+
+from pyment.configurations import DatasetConfiguration, FinetuningConfiguration
+from pyment.factories import loss_factory, metric_factory, optimizer_factory
+from pyment.models.sfcn import sfcn_factory
+from pyment.models.utils.load_select_pretrained_weights import (
+    load_select_pretrained_weights
+)
+from pyment.utils.json_serialize import json_serialize
+
+
+logging.basicConfig(
+    format='%(asctime)s - %(levelname)s - %(name)s: %(message)s',
+    level=logging.DEBUG
+)
+logger = logging.getLogger(__name__)
+
+def _create_tensorflow_dataset(
+    df: pd.DataFrame, *,
+    target: str,
+    input_shape: Tuple[int, int, int],
+    batch_size: str,
+    shuffle: bool = False
+) -> tf.data.Dataset:
+    input_shape = tf.constant(input_shape)
+
+    df = df.copy()
+    df = df.sample(frac=1.)
+
+    dataset = tf.data.Dataset.from_tensor_slices((df['path'], df[target]))
+
+    if shuffle:
+        dataset = dataset.shuffle(buffer_size=5*batch_size)
+    
+    dataset = dataset.map(
+        lambda path, label: (load_mgh(path), label),
+        num_parallel_calls=tf.data.AUTOTUNE
+    )
+    dataset = dataset.map(
+        lambda image, label: (center_crop_or_pad(image, input_shape), label),
+        num_parallel_calls=tf.data.AUTOTUNE
+    )
+
+    dataset = dataset.batch(batch_size)
+    dataset = dataset.prefetch(tf.data.AUTOTUNE)
+
+    return dataset
+
+def _create_checkpointing_callback(
+    destination: str, 
+    metrics: List[tf.keras.metrics.Metric] = None
+):
+    os.mkdir(destination)
+
+    train_metrics = []
+    val_metrics = []
+    
+    if metrics is not None:
+        for metric in metrics:
+            name = metric.name.replace('_', '-')
+            train_metrics.append(f'{name}={{{metric.name}:.2f}}')
+            val_metrics.append(f'val-{name}={{val_{metric.name}:.2f}}')
+
+    terms = [
+        'epoch={epoch:03d}',
+        'loss={loss:.2f}'
+    ] + train_metrics + [
+        'val-loss={val_loss:.2f}'
+    ] + val_metrics
+    filename = '_'.join(terms) + '.hdf5'
+    filepath = os.path.join(destination, filename)
+
+    return tf.keras.callbacks.ModelCheckpoint(
+        filepath,
+        monitor='val_loss',
+        save_best_only=True,
+        save_weights_only=True
+    )
+
+def finetune(
+    model_type: str,
+    model_constructor_arguments: Dict[str, Any],
+    weights: str, 
+    input_shape: Tuple[int, int, int],
+    target: str,
+    loss: Callable[[tf.Tensor, tf.Tensor], tf.Tensor],
+    metrics: List[tf.keras.metrics.Metric],
+    optimizer: tf.optimizers.Optimizer,
+    learning_rate_scheduler: tf.keras.callbacks.Callback,
+    training: pd.DataFrame, 
+    validation: pd.DataFrame,
+    batch_size: int,
+    epochs: int,
+    destination: str
+):
+    if destination is not None:
+        if os.path.isdir(destination):
+            raise ValueError(f'Destination {destination} already exists')
+
+        logger.info('Creating destination folder %s', destination)
+        os.mkdir(destination)
+
+    model_class = sfcn_factory(model_type)
+    model = model_class(
+        input_shape=input_shape,
+        **model_constructor_arguments
+    )
+    load_select_pretrained_weights(model, weights, target=target)
+
+    model.compile(loss=loss, optimizer=optimizer, metrics=metrics)
+
+    training_dataset = _create_tensorflow_dataset(
+        training, 
+        input_shape=input_shape,
+        target=target, 
+        batch_size=batch_size,
+        shuffle=True
+    )
+    validation_dataset = _create_tensorflow_dataset(
+        validation,
+        input_shape=input_shape,
+        target=target,
+        batch_size=batch_size,
+        shuffle=False
+    )
+
+    callbacks = [
+        _create_checkpointing_callback(
+            os.path.join(destination, 'checkpoints'),
+            metrics=metrics
+        ),
+        learning_rate_scheduler
+    ]
+
+    history = model.fit(
+        training_dataset, 
+        validation_data=validation_dataset,
+        epochs=epochs,
+        callbacks=callbacks
+    )
+
+    with open(os.path.join(destination, 'history.json'), 'w') as f:
+        json.dump(json_serialize(history.history), f)
+
+def finetune_from_configuration(configuration: str):
+    with open(configuration, 'r') as f:
+        configuration = json.load(f)
+    
+    configuration = FinetuningConfiguration.model_validate(configuration)
+
+    training, validation = DatasetConfiguration.parse(
+        configuration.data,
+        target=configuration.training.target
+    )
+
+    # strategy = tf.distribute.MirroredStrategy()
+
+    # with strategy.scope():
+
+    loss_cls = loss_factory(configuration.training.loss)
+    loss = loss_cls()
+
+    optimizer_cls = optimizer_factory(configuration.training.optimizer)
+    optimizer = optimizer_cls(configuration.training.learning_rate)
+
+    metrics = None
+    
+    if configuration.training.metrics is not None:
+        metrics = [
+            metric_factory(metric)
+            for metric in configuration.training.metrics
+        ]
+
+    learning_rate_scheduler = None
+
+    if configuration.training.learning_rate_schedule:
+        learning_rate_scheduler = (
+            configuration.training.learning_rate_schedule.instantiate()
+        )
+
+    finetune(
+        model_type=configuration.model.type,
+        model_constructor_arguments=configuration.model.hyperparameters,
+        weights=configuration.model.weights,
+        input_shape=configuration.data.input_shape,
+        target=configuration.training.target,
+        loss=loss,
+        metrics=metrics,
+        optimizer=optimizer,
+        learning_rate_scheduler=learning_rate_scheduler,
+        training=training,
+        validation=validation,
+        batch_size=configuration.training.batch_size,
+        epochs=configuration.training.epochs,
+        destination=configuration.training.destination
+    )
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        'Finetunes a multi-task SFCN according to the given configuration'
+    )
+
+    parser.add_argument('configuration', help='Path to configuration JSON')
+
+    args = parser.parse_args()
+
+    finetune_from_configuration(args.configuration)

scripts/predict_from_bids_folder.py

@@ -4,11 +4,12 @@ import os
 import re
 import numpy as np
 import pandas as pd
+from typing import List
 from tqdm import tqdm
 
 import nibabel as nib
 
-from pyment.models import MultiTaskSFCN
+from pyment.models.sfcn import sfcn_factory
 from pyment.preprocessing.conform import conform
 
 
@@ -29,7 +30,11 @@ def _extract_run(filename: str) -> str:
 
 def predict_from_bids_folder(
     source: str, 
-    weights: str, 
+    weights: str,
+    model_name: str = 'sfcn-multi',
+    targets: List[str] = [
+        'age', 'sex', 'handedness', 'bmi', 'fluid_intelligence', 'neuroticism'
+    ],
     destination: str = None,
     per_image_normalization: bool = False
 ) -> pd.DataFrame:
@@ -37,7 +42,8 @@ def predict_from_bids_folder(
         raise ValueError(f'Destination {destination} already exists')
     
     logger.info('Loading multi-task model with weights %s', weights)
-    model = MultiTaskSFCN(weights=weights)
+    model_class = sfcn_factory(model_name)
+    model = model_class(weights=weights)
 
     results = []
 
@@ -73,16 +79,13 @@ def predict_from_bids_folder(
                 )
 
                 results.append({
-                    'source': os.path.join(anat_folder, filename),
-                    'subject': subject,
-                    'session': session,
-                    'run': run,
-                    'age': predictions[0],
-                    'sex': predictions[1],
-                    'handedness': predictions[2],
-                    'bmi': predictions[3],
-                    'fluid_intelligence': predictions[4],
-                    'neuroticism': predictions[5]
+                    **{
+                        'source': path,
+                        'subject': subject,
+                        'session': session,
+                        'run': run
+                    },
+                    **{targets[i]: predictions[i] for i in range(len(targets))}
                 })
 
     results = pd.DataFrame(results)
@@ -108,6 +111,24 @@ if __name__ == '__main__':
             'exist files named <path>.index and <path>.data-00000-of-00001'
         )
     )
+    parser.add_argument(
+        '-m', '--model', 
+        required=False,
+        default='sfcn-multi',
+        help=(
+            'Name of the model to use'
+        )
+    )
+    parser.add_argument(
+        '-t', '--targets',
+        required=False,
+        nargs='+',
+        default=[
+            'age', 'sex', 'handedness', 'bmi', 'fluid_intelligence', 
+            'neuroticism'
+        ],
+        help='Name to use for each of the prediction heads in the output CSV'
+    )
     parser.add_argument(
         '-d', '--destination',
         required=False,
@@ -128,6 +149,8 @@ if __name__ == '__main__':
     predict_from_bids_folder(
         source=args.bids,
         weights=args.weights,
+        model_name=args.model,
+        targets=args.targets,
         destination=args.destination,
         per_image_normalization=args.per_image_normalization
     )

scripts/predict_from_fastsurfer_folder.py

@@ -5,11 +5,11 @@ import re
 import numpy as np
 import pandas as pd
 from tqdm import tqdm
-from typing import Tuple
+from typing import List, Tuple
 
 import nibabel as nib
 
-from pyment.models import MultiTaskSFCN
+from pyment.models.sfcn import sfcn_factory
 from pyment.preprocessing.conform import conform
 
 
@@ -29,14 +29,20 @@ def _parse_folder_name(name: str) -> Tuple[str, str, str]:
 
 def predict_from_fastsurfer_folder(
     source: str, 
-    weights: str, 
+    weights: str,
+    model_name: str = 'sfcn-multi',
+    targets: List[str] = [
+        'age', 'sex', 'handedness', 'bmi', 'fluid_intelligence', 'neuroticism'
+    ],
     destination: str = None
 ) -> 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 = MultiTaskSFCN(weights=weights)
+
+    model_class = sfcn_factory(model_name)
+    model = model_class(weights=weights)
 
     results = []
 
@@ -66,19 +72,18 @@ 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),
-            'subject': subject,
-            'session': session,
-            'run': run,
-            'age': predictions[0],
-            'sex': predictions[1],
-            'handedness': predictions[2],
-            'bmi': predictions[3],
-            'fluid_intelligence': predictions[4],
-            'neuroticism': predictions[5]
+            **{
+                'source': os.path.join(source, folder),
+                'subject': subject,
+                'session': session,
+                'run': run
+            },
+            **{targets[i]: predictions[i] for i in range(len(targets))}
         })
 
     results = pd.DataFrame(results)
@@ -110,6 +115,24 @@ if __name__ == '__main__':
             'exist files named <path>.index and <path>.data-00000-of-00001'
         )
     )
+    parser.add_argument(
+        '-m', '--model', 
+        required=False,
+        default='sfcn-multi',
+        help=(
+            'Name of the model to use'
+        )
+    )
+    parser.add_argument(
+        '-t', '--targets',
+        required=False,
+        nargs='+',
+        default=[
+            'age', 'sex', 'handedness', 'bmi', 'fluid_intelligence', 
+            'neuroticism'
+        ],
+        help='Name to use for each of the prediction heads in the output CSV'
+    )
     parser.add_argument(
         '-d', '--destination',
         required=False,
@@ -121,7 +144,9 @@ if __name__ == '__main__':
 
     predict_from_fastsurfer_folder(
         source=args.root,
+        model_name=args.model,
         weights=args.weights,
+        targets=args.targets,
         destination=args.destination
     )