# Data Organization for Virtual Staining

Here we document our conventions for storing data, metadata, configs, and models.

## Data flow in the pipeline

TODO: Following diagram captures the planned flow of data and metadata through the new version of the pipeline.

### Preprocessing

Parameters are provided via the CLI, and stored in the attributes of
the OME-Zarr datasets using [iohub](https://github.com/czbiohub/iohub).
The metadata is in json format that is practical to edit by hand.

### Training

We use the PyTorch Lightning framework for training,
which provides [good defaults](https://pytorch-lightning.readthedocs.io/en/1.6.5/common/lightning_cli.html)
for CLI and training configs,
and organized TensorBoard [logs](https://lightning.ai/docs/pytorch/stable/extensions/logging.html).

### Inference

The inference module does not depend on lightning, but just on PyTorch.
Parameters are provided with CLI and stored with the OME-Zarr datasets,
similar to preprocessing.

> This can be further incorporated into the lightning pipeline

### Evaluation

Evaluating the models requires human proof-reading of ground truth.
Currently computing the evaluation metrics does not depend on PyTorch.

> This can be further incorporated into the lightning pipeline

### Deployment

For run-time deployment, we export the model to the ONNX format.

## Data hierarchy (Lightning Framework)

Data generated by the pipeline is stored on the file system following this schema:

```yaml
# project root directory
virtual_staining:
    # registerd, deconvolved, preprocessed OME-Zarr stores
    datasets:
        train:
            yyyymmdd_dataname0.zarr
            yyyymmdd_dataname1.zarr
            ...
        test:
            yyyymmdd_dataname0.zarr
            yyyymmdd_dataname1.zarr
            ...
        # FOVs with human-proof-read ground truth segmentation masks
        ground_truth:
            yyyymmdd_dataname0_gt:
                fovs.zarr
                <frame0>.tif
                <frame0>_mask.png
                <frame1>.tif
                <frame1>_mask.png
                ...
            yyyymmdd_dataname1_gt:
                ...
            ...
    # computational experiments
    models:
        experiment0:
            # working training configs (may not specify default fields)
            training_config0.yaml
            training_config1.yaml
            ...
            lightning_logs:
                # output of one run
                yyyymmdd-hhmmss:
                    # model weights
                    checkpoints:
                        epoch0.ckpt
                        epoch1.ckpt
                        ...
                    # autosaved full config
                    config.yaml
                yyyymmdd-hhmmss:
                    ...
            # Inference and/or Evaluation of selected models.
            test: 
                # config for prediction with test dataset.
                test_<suffix>.yml # config used for inference, optionally copies ground truth and input for evaluation. This config will follow the lightning CLI/config format.

                # inference output on test dataset, may include copies of input and ground truth to facilitate visualization of model performance. 
                test_<suffix>.zarr # Not all test datasets need to have human curated ground truth.
                ...

                # config for evaluation: checkpoint path, test data path that have ground turth included, and choice of metrics.
                evaluation_<suffix>.yaml
                ...

                # evaluation metrics
                evaluation_metrics_<suffix>.csv
                ...
                # (optional) tensorboard logs generated to visualize distribution of metrics or specific samples of input, prediction, ground truth.
                evaluation_logs: 
            # (optional) exported models for deployment
            deployment:
                <task>_<input_shape>.onnx
                README
        experiment1:
            ...
        ...
```

## Data hierarchy (Gunpowder Framework)

This data hierarchy is deprecated, and is documented for archiving purposes.
The hierarchy organizes subdirectories of config files, models/training logs and data first according to the data related to the computational experiment, then by the specific experiment. Each set of config files should have a corresponding sibling-level training log, and parent-level dataset in their respective directories

```yaml
# project root directory 
torch_microDL:

    #training and test data
    data:
        <yyyy_mm_dd>_<data_name>: # data-level sibling to dataset dir in config files
            # Due to evolving data format, no single standard for each dataset's format
            # Generally these directories are populated by one of the following:
            #   tile directories (from old preprocessing)
            #   single page tiff directories (from old raw data)
            #   zarr stores (from new dataloading)
            <data_directory1>:
            <data_directory2>:
            ...
        ...

    # configuration files (preprocessing, training, inference, etc)
    config_files:
        <yyyy_mm_dd>_<data_name>: # data_name is often an abbreviated tag to the microscopy experiment sourcing this data
            <mm_dd_yyyy>_<experiment_name>: # config files often stored under additional subdirectories. No standard format
                config0_<config_type>.yml
                config1_<config_type>.yml
                ...
            ... 
        ... 

    # training logs and saved models
    models:
        <yyyy_mm_dd>_<data_name>: # data-level sibling to dataset dir in config files
            <mm_dd_yyyy>_<experiment_name>:
                model_<model_type>0:
                    # Sometimes there is an additional subdirectory here deliniating different runs.
                    # There are often *many* training models. These should be cleaned or sorted by size.
                    training_model_<yyyy_mm_dd_mm_ss>:
                        <tensorboard_logs>
                        data_splits.yml
                        saved_model_ep_<ep0>_testloss_<loss>.pt
                        saved_model_ep_<ep1>_testloss_<loss>.pt
                        ...
                        prediction_ep_<ep0>.png
                        prediction_ep_<ep1>.png
                        ...
                     inference_results_<yyyy_mm_dd_mm_ss>:
                        <tensorboard_logs>
                        <inference_result_1>.tiff
                        <inference_result_2>.tiff
                        ...
                ...
            ...
        ...
          
```
