main.py

"""
Main pipeline of the OCT Neuralizer Model, or this project in general.
Handles all the pipelines steps, i.e.:
- CLI args parsing and config file parsing
- Metadata creation (Image collection)
- Task creation
- DataLoader object creation (train / val / test)
- Fitting (with Tensorboard Logging)

Run via CUDA_VISIBLE_DEVICES=x python3 main.py, where x specifies the GPU ID to use.
CLI Args:
    --fit-neuralizer-oct: Fits Neuralizer OCT
    --fit-single-task-models: Fits Single-Task Models
    --fit-neuralizer-domain-adaption: Fits Neuralizer Domain Adaption Models (3x as we have three diff. Vendors)
    --fit-neuralizer-domain-adaption-single-task-models: Same as above, but for each task one model
    --fit-retinalizer: Fits Retinalizer OCT
    --fit-retinalizer-domain-adaption: Fits Retinalizer DA Models (3x as we have three diff. Vendors)

Default Args:
    --config: str -- Path to the config file (default: "configs/config.yaml")

Example Calls:
CUDA_VISIBLE_DEVICES=x python3 main.py --fit-neuralizer-oct
CUDA_VISIBLE_DEVICES=x python3 main.py --fit-retinalizer
CUDA_VISIBLE_DEVICES=x python3 main.py --fit-retinalizer-domain-adaption
"""

import logging
import pandas as pd
import pytorch_lightning as pl

import src.utils.utils as utils
import src.data.dataloader as dataloader
import src.tasks.taskloader as taskloader
import src.train.training as training

from typing import List, Optional
from pytorch_lightning.callbacks import EarlyStopping, LearningRateMonitor
from src.tasks.base_task import BaseTask
from src.neuralizer.lightning_model import LightningModel
from src.retinalizer.adversarial_lightning_model import Retinalizer

utils.setup_logging()
logger = logging.getLogger(__name__)


def main():
    # Parse CLI arguments and config file as dict
    CLI_ARGS: dict = utils.parse_args()
    CFG: dict = utils.load_config(CLI_ARGS.get("config"))

    # Get params which model kind to fit
    fit_neuralizer_oct: bool = CLI_ARGS.get("fit_neuralizer_oct")
    fit_single_task_models: bool = CLI_ARGS.get("fit_single_task_models")
    fit_neuralizer_domain_adaption: bool = CLI_ARGS.get("fit_neuralizer_domain_adaption")
    fit_neuralizer_domain_adaption_single_task_models: bool = CLI_ARGS.get(
        "fit_neuralizer_domain_adaption_single_task_models"
    )
    fit_retinalizer: bool = CLI_ARGS.get("fit_retinalizer")
    fit_retinalizer_domain_adaption: bool = CLI_ARGS.get("fit_retinalizer_domain_adaption")

    # Create or get metadata.csv, i.e. load all images and preprocess
    dataloader.load_data(CFG)
    metadata_df: pd.DataFrame = dataloader.get_metadata_df(CFG["dataloader"]["metadata_location"])
    enrichment_df: pd.DataFrame = dataloader.get_metadata_df(CFG["dataloader"]["enrichment_location"])

    # Get training tasks (seen)
    tasks: List[BaseTask] = taskloader.get_tasks(df=metadata_df, enrichment_df=enrichment_df, CFG=CFG["tasks"])

    # Get test tasks (tasks unseen during training)
    test_tasks: List[BaseTask] = taskloader.get_test_tasks(
        df=metadata_df, enrichment_df=enrichment_df, CFG=CFG["tasks"]
    )

    # "Random Segmentation Mask Recoloring"(= "RSMR")- Run, if it is allowed in config file
    # Has the effect of randomly recoloring segmentation masks in both label and the context labels (ctx_out)
    suffix = "RSMR" if CFG["training"]["dataset"]["allow_random_segmentation_mask_recoloring"] else None

    # Define callbacks to give to Pytorch Lightning Trainer in specific fitting settings
    callbacks = [EarlyStopping(**CFG["training"]["early_stopping"]), LearningRateMonitor(logging_interval="epoch")]

    # Fit Neuralizer OCT trained on all tasks, evaluate on them and also evaluate on test tasks (unseen)
    if fit_neuralizer_oct:
        _fit_neuralizer_oct_on_all_tasks(CFG, tasks, test_tasks, suffix, callbacks)
    # Fit Neuralizer on Single Tasks (seen+unseen) and evaluate on them
    elif fit_single_task_models:
        _fit_neuralizer_oct_on_single_tasks(CFG, tasks, test_tasks, suffix, callbacks)
    # Fit Neuralizer Domain Adaption Models (RETOUCH Vendors)
    elif fit_neuralizer_domain_adaption:
        _fit_neuralizer_retouch_domain_adaption(CFG, tasks, test_tasks, suffix, callbacks)
    # Fit Neuralizer Domain Adaption Models (RETOUCH Vendors) and for each Task on Single Model
    elif fit_neuralizer_domain_adaption_single_task_models:
        _fit_neuralizer_retouch_domain_adaption_on_single_tasks(CFG, tasks, test_tasks, suffix, callbacks)
    # Fit Retinalizer
    elif fit_retinalizer:
        _fit_retinalizer_oct(CFG, tasks, test_tasks, suffix)
    # Fit Retinalizer in Domain Generalization Scenario
    elif fit_retinalizer_domain_adaption:
        _fit_retinalizer_retouch_domain_adaption(CFG, tasks, test_tasks, suffix)
    else:
        logger.info("+++ No fitting instruction in CLI args found +++")


def _fit_neuralizer_oct_on_all_tasks(
    CFG: dict,
    tasks: List[BaseTask],
    test_tasks: List[BaseTask],
    suffix: Optional[str],
    callbacks: List,
) -> None:
    """Model is trained on all tasks and tested on tasks seen, as well as tested on tasks unseen.

    :param CFG: dict -- Parsed config file containing all hyperparameters
    :param tasks: List[BaseTask] -- List of initialized tasks to be seen
    :param test_tasks: List[BaseTask] -- List of initialized tasks to be unseen
    :param suffix: Optional[str] -- Suffix to append to "OCTNeuralizer" for logging purposes
    :param callbacks: List -- List of Callbacks, e.g. EarlyStopping(.)
    :return: None
    """
    assert CFG["training"]["tensorboard_logger"]["name"] == "OCTNeuralizer", "Use 'OCTNeuralizer' as name in config."
    # Get model
    model: LightningModel = training.get_model(CFG=CFG)

    # Get train, validation and test set as DataLoader objects
    train_dataloader, val_dataloader, test_dataloader = training.get_dataloaders(CFG=CFG, tasks=tasks)

    # Get test-tasks as Dataloader
    _, _, test_unseen_dataloader = training.get_dataloaders(CFG=CFG, tasks=test_tasks)

    # Fit model and return fitted Trainer object
    trainer: pl.Trainer = training.fit_model(
        CFG=CFG,
        model=model,
        train_dataloader=train_dataloader,
        val_dataloader=val_dataloader,
        test_dataloader=test_dataloader,
        callbacks=callbacks,
        suffix=suffix,
    )

    # Test generalization ability to unseen tasks
    model.test_loss_suffix = "unseen"
    trainer.test(model=model, dataloaders=test_unseen_dataloader, verbose=True)


def _fit_neuralizer_oct_on_single_tasks(
    CFG: dict,
    tasks: List[BaseTask],
    test_tasks: List[BaseTask],
    suffix: Optional[str],
    callbacks: List,
) -> None:
    """Model is trained on one task and tested on that task again with 60:20:20 split.
    Repeat for all tasks, therefore len(tasks) + len(test_tasks) models are trained and evaluated

    :param CFG: dict -- Parsed config file containing all hyperparameters
    :param tasks: List[BaseTask]  -- List of initialized tasks to be seen
    :param test_tasks: List[BaseTask]  -- List of initialized tasks to be seen (single-tasks models are trained on all)
    :param suffix: Optional[str] -- Suffix to append to "OCTNeuralizer" for logging purposes
    :param callbacks: List -- List of Callbacks, e.g. EarlyStopping(.)
    :return: None
    """
    assert CFG["training"]["tensorboard_logger"]["name"] == "OCTNeuralizer", "Use 'OCTNeuralizer' as name in config."
    # Combine task, as single models are trained on all available tasks
    tasks_combined = tasks + test_tasks

    for task in tasks_combined:
        task_name: str = task.get_task_name()
        logger.info(f"+++ Fitting Single Model for Task '{task_name}' +++")

        # Instantiate and get model
        model: LightningModel = training.get_model(CFG=CFG)

        # Get train, validation and test set as DataLoader objects
        train_dataloader, val_dataloader, test_dataloader = training.get_dataloaders(CFG=CFG, tasks=[task])

        # Fit model and return fitted Trainer object
        trainer: pl.Trainer = training.fit_model(
            CFG=CFG,
            model=model,
            train_dataloader=train_dataloader,
            val_dataloader=val_dataloader,
            test_dataloader=test_dataloader,
            callbacks=callbacks,
            task_name=task_name,
            suffix=suffix,
        )


def _fit_neuralizer_retouch_domain_adaption(
    CFG: dict, tasks: List[BaseTask], test_tasks: List[BaseTask], suffix: Optional[str], callbacks: List
) -> None:
    """In this scenario, multiple models are trained: For each vendor, one model.

    :param CFG: dict -- Parsed config file containing all hyperparameters
    :param tasks: List[BaseTask] -- List of initialized tasks to be seen
    :param test_tasks: List[BaseTask] -- List of initialized tasks to be unseen
    :param suffix: Optional[str] -- Suffix to append to "OCTNeuralizer" for logging purposes
    :param callbacks: List -- List of Callbacks, e.g. EarlyStopping(.)
    :return: None
    """
    assert CFG["training"]["tensorboard_logger"]["name"] == "OCTNeuralizer", "Use 'OCTNeuralizer' as name in config."
    vendors: List[str] = CFG["misc"]["retouch_vendors"]

    # We fit three models in total: For each vendor left out, one Model
    for vendor in vendors:
        logger.info(f"+++ Fitting Domain Adaption Model for vendor '{vendor}' +++")
        model: LightningModel = training.get_model(CFG=CFG)

        # Get train, validation and test set as DataLoader objects with one vendor left out
        train_dataloader, val_dataloader, test_dataloader = training.get_dataloaders(
            CFG=CFG, tasks=tasks, vendor=vendor
        )

        # Get test-tasks as DataLoader, with one vendor left out
        _, _, test_unseen_dataloader = training.get_dataloaders(CFG=CFG, tasks=test_tasks, vendor=vendor)

        # Fit model and return fitted Trainer object
        trainer: pl.Trainer = training.fit_model(
            CFG=CFG,
            model=model,
            train_dataloader=train_dataloader,
            val_dataloader=val_dataloader,
            test_dataloader=test_dataloader,
            callbacks=callbacks,
            vendor=vendor,
            suffix=suffix,
        )

        # Test generalization ability to unseen tasks
        model.test_loss_suffix = "unseen"
        trainer.test(model=model, dataloaders=test_unseen_dataloader, verbose=True)


def _fit_retinalizer_retouch_domain_adaption(
    CFG: dict, tasks: List[BaseTask], test_tasks: List[BaseTask], suffix: Optional[str]
) -> None:
    """In this scenario, multiple models are trained: For each vendor, one model.

    :param CFG: dict -- Parsed config file containing all hyperparameters
    :param tasks: List[BaseTask] -- List of initialized tasks to be seen
    :param test_tasks: List[BaseTask] -- List of initialized tasks to be unseen
    :param suffix: Optional[str] -- Suffix to append to "OCTNeuralizer" for logging purposes
    :return: None
    """
    assert CFG["training"]["tensorboard_logger"]["name"] == "Retinalizer", "Use 'Retinalizer' as name in config."
    vendors: List[str] = CFG["misc"]["retouch_vendors"]

    # We fit three models in total: For each vendor left out, one Model
    for vendor in vendors:
        logger.info(f"+++ Fitting Domain Adaption Model for vendor '{vendor}' +++")
        model: Retinalizer = training.get_retinalizer_model(CFG=CFG)

        # Get train, validation and test set as DataLoader objects with one vendor left out
        train_dataloader, val_dataloader, test_dataloader = training.get_dataloaders(
            CFG=CFG, tasks=tasks, vendor=vendor
        )

        # Get test-tasks as DataLoader, with one vendor left out
        _, _, test_unseen_dataloader = training.get_dataloaders(CFG=CFG, tasks=test_tasks, vendor=vendor)

        # Fit model and return fitted Trainer object
        trainer: pl.Trainer = training.fit_model(
            CFG=CFG,
            model=model,
            train_dataloader=train_dataloader,
            val_dataloader=val_dataloader,
            test_dataloader=test_dataloader,
            vendor=vendor,
            suffix=suffix,
        )

        # Test generalization ability to unseen tasks
        model.test_loss_suffix = "unseen"
        trainer.test(model=model, dataloaders=test_unseen_dataloader, verbose=True)


def _fit_neuralizer_retouch_domain_adaption_on_single_tasks(
    CFG: dict,
    tasks: List[BaseTask],
    test_tasks: List[BaseTask],
    suffix: Optional[str],
    callbacks: List,
) -> None:
    """Here we fit the Neuralizer Domain Adaption Models. For one task and one vendor: One Model.

    :param CFG: dict -- Parsed config file containing all hyperparameters
    :param tasks: List[BaseTask] -- List of initialized tasks to be seen
    :param test_tasks: List[BaseTask] -- List of initialized tasks to be unseen
    :param suffix: Optional[str] -- Suffix to append to "OCTNeuralizer" for logging purposes
    :param callbacks: List -- List of Callbacks, e.g. EarlyStopping(.)
    :return: None
    """
    assert CFG["training"]["tensorboard_logger"]["name"] == "OCTNeuralizer", "Use 'OCTNeuralizer' as name in config."
    logger.info("Fitting Neuralizer Domain Adaption (DA) Single Task Models ...")

    # Get valid list of OCT vendors
    vendors: List[str] = CFG["misc"]["retouch_vendors"]
    # Combine task, as single models are trained on all available tasks and filter for only RETOUCH tasks
    tasks_combined = tasks + test_tasks
    tasks_combined = [task for task in tasks_combined if "RETOUCH" in task.get_task_name()]

    # We fit multiple single tasks models for each vendor
    for vendor in vendors:
        logger.info(f"+++ Fitting Domain Adaption Single Task Models for vendor '{vendor}' +++")

        # We fit one model for each task and vendor
        for task in tasks_combined:
            task_name: str = task.get_task_name()
            logger.info(f"+++ Fitting Single DA-Model for Task '{task_name}' +++")

            # Load non-fitted model
            model: LightningModel = training.get_model(CFG=CFG)

            # Get train, validation and test set as DataLoader objects with one vendor left out
            train_dataloader, val_dataloader, test_dataloader = training.get_dataloaders(
                CFG=CFG, tasks=[task], vendor=vendor
            )

            # Fit model and return fitted Trainer object
            trainer: pl.Trainer = training.fit_model(
                CFG=CFG,
                model=model,
                train_dataloader=train_dataloader,
                val_dataloader=val_dataloader,
                test_dataloader=test_dataloader,
                task_name=task_name,
                callbacks=callbacks,
                vendor=vendor,
                suffix=suffix,
            )


def _fit_retinalizer_oct(CFG: dict, tasks: List[BaseTask], test_tasks: List[BaseTask], suffix: Optional[str]) -> None:
    """Fits Retinalizer, i.e. our Neuralizer architecture improvement.
    Will fit Retinalizer on all (seen) tasks and tests them on seen and unseen tasks.

    :param CFG: dict -- Parsed config file containing all hyperparameters
    :param tasks: List[BaseTask] -- List of initialized tasks to be seen
    :param test_tasks: List[BaseTask] -- List of initialized tasks to be unseen
    :param suffix: Optional[str] -- Suffix to append to "Retinalizer" for logging purposes, e.g. "Retinalizer-RSMR"
    :return: None
    """
    # Get model
    model = training.get_retinalizer_model(CFG=CFG)
    assert CFG["training"]["tensorboard_logger"]["name"] == "Retinalizer", "Please use 'Retinalizer' as name in config!"

    # Get train, validation and test set as DataLoader objects
    train_dataloader, val_dataloader, test_dataloader = training.get_dataloaders(CFG=CFG, tasks=tasks)

    # Get test-tasks as Dataloader
    _, _, test_unseen_dataloader = training.get_dataloaders(CFG=CFG, tasks=test_tasks)

    logger.info("Fitting Retinalizer ...")
    # Fit model and return fitted Trainer object
    trainer: pl.Trainer = training.fit_model(
        CFG=CFG,
        model=model,
        train_dataloader=train_dataloader,
        val_dataloader=val_dataloader,
        test_dataloader=test_dataloader,
        suffix=suffix,
    )

    # Test generalization ability to unseen tasks
    model.test_loss_suffix = "unseen"
    trainer.test(model=model, dataloaders=test_unseen_dataloader, verbose=True)


if __name__ == "__main__":
    main()