Refactoring of the codebase

The codebase has gotten too big and some refactoring is needed. Some ideas to address:

- [x] Data loading: The data loader has been extracted out into [caked](https://github.com/alan-turing-institute/caked), in order to be able to load many more types of datasets, and from different pipelines. This simplifies what happens in `data.py`. Refactoring happening on #293

- [x] Training and evaluation: Too many things are happening in these scripts that can be moved in to its own files, so these only show the pipeline for training (#303):
- Load data
- Define model
- Define loss function
- Run loop for training
- Visualise and save results.

- [ ] Functions with too many parameters: train, eval, pass_batch, etc have too many parameters and it can get confusing and introduce bugs, we already have a a pedantic data model that can be passed around and in some cases, a class seems better than a function. Some functions can be broken into smaller parts as well.

- [ ] Many minor bug fixes that have been opened for a while.

Please add any more ideas for refactoring, I'll keep this as a meta issue for discussion.

Before this refactoring begins I think we should merge the current version of the code to `main` as it that has been used to produce results for the paper.