Explanations of all Folders/Files

Figures

All figures generated by the code.

Data

All processed data and trained models data.

Saved_model

Model saved after training in this folder.

src

All source code for data processing, model, loss function, etc.

Jupyter Notebook file for plotting Figures

Fig_2_4.ipynb file is to plot Fig. 2 to Fig.4 in the report and similarly others.

train_mri.py

This script is used to train a U-Net model for MRI image reconstruction based image segmentation.

python train_mri.py --mri_inchannels 0 3

where the --mri_inchannels takes the index of the modality to consider for the training. Here, 0 means T1 and 3 means FLAIR.

The steps:

1. Environment: The Environmemnt file is environment.yaml.

   • conda env create -f environment.yaml

   • conda activate CS525_Project2

   • pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126

Please contact us for futher questions.

2. Load configuration: The script loads hyperparameters from a config.json file. These parameters include file paths, batch size, training and validation data details, and model parameters.

3. Prepare the dataset:

   • It collects training and validation data from the specified directory (data_dir), which contains MRI volume slices stored as .h5 files.
   • The data is loaded using a custom UNetDataset or UNetDataset_Mem class based on the config file and wrapped in PyTorch's DataLoader for batching and shuffling.

4. Model setup:

   • A U-Net model is initialized with specific input and output channels.

5. Loss function & optimizer:

   • The DiceLoss function is used to measure the accuracy of image segmentation.
   • The Adam optimizer is initialized with learning rate and other hyperparameters.

6. Training:

   • The model is trained using the train_model function, which takes the training and validation data loaders, the model, loss function, optimizer, and hyperparameters such as the number of epochs and learning rate scheduler details.

7. Model saving:

   • After training, the model is saved according to the the key "saved_model_path" of config.json. Also, the training loss/accuracy is saved according to the according to the the key "log_file of" of config.json.

The script is essential for training the MRI reconstruction model and monitoring its performance during training.

The pipeline of image reconstruction and training is shown in .