Deep Learning-Based Denoising of Monte Carlo Simulations for Efficient Laser-Tissue Interaction Modeling

Overview

This repository contains the official implementation of the paper "Deep Learning-Based Denoising of Monte Carlo Simulations for Efficient Laser-Tissue Interaction Modeling".

Features

• Denoising networks for 3D Monte Carlo simulation data: U-Net, Restormer, NAFNet
• Python codes for training, evaluation, loading dataset
• Pre-trained model weights

0. Installation

Clone this repository: git clone https://github.com/CMME-Lab/LIFUSimul-DL.git
Install all prerequisites with pip install -r requirements.txt

1. Preparing dataset

• Dataset will provided upon reasonable request.
• Locate the '.h5' files in your desired root directory.
• Specify the dataset path using the '--data_path' argument when running train/test scripts.

Example usage :
python train.py --data_path ./my_root_path

Custom dataset format

• Prepare two HDF5 files following the same structure used by dataset.py.

  • frate_high.h5: high-resolution target data.

  • frate_low_{sampling}.h5: corresponding low-sampling input data. {sampling} must match the value passed to --n_sampling.

• Both HDF5 files must contain the same dataset keys. For each key, the dataloader reads the paired data as fh = frate_high[key].

• Each dataset entry should be a 3D volume. The paired high- and low-resolution volumes must have the same shape before cropping, and each key must correspond to the same sample in both HDF5 files.

• The current train/validation/test split assumes 6 cases with 1000 samples per case. If your custom dataset uses a different number of cases or samples per case, update the num_cases and length arguments in split_dataset() accordingly.

2. Training

Run the training process using train.py.

• The training script loads the paired HDF5 dataset, splits it into training and validation subsets, initializes the selected network, and saves the best model checkpoint based on the validation loss.
• The experiment outputs are saved under:

runs/{run_name}/
├── configs.pickle
├── model.pth
└── train_results.pickle

• Select the model architecture using the --net_type argument:

  • nafnet: NAFNet-based 3D denoising network
  • unet: baseline 3D U-Net
  • restormer: Restormer-based 3D denoising network

• Important training arguments:

  • --data_path: root directory containing the HDF5 dataset files
  • --run_name: name of the experiment folder under runs/
  • --net_type: network architecture to train
  • --data_shape: input/output volume size after center cropping
  • --n_sampling: sampling number used to select frate_low_{sampling}.h5
  • --gpu_num: GPU index
  • --wandb_pj: Weights & Biases project name, if logging is desired

Example usage:
python train.py --data_path ./my_root_path --run_name my_experiment --net_type nafnet --criterion mse --num_epochs 100 --batch_size 8 --gpu_num 0 --n_sampling 1000 --lr_schedule

3. Evaluation

Run the evaluation process using test.py.

• The evaluation script loads the saved training configuration and model checkpoint from runs/{run_name}/.
• The test dataset is constructed using the same dataset path, data shape, sampling number, and split setting.
• The network architecture is selected using the same --net_type argument used during training.
• The script reports PSNR, L1 error, L2 error, peak-density difference, and inference time.
• Evaluation results are saved as:

runs/{run_name}/eval_results.pickle

• To save prediction plots, add the --plot argument.
• To additionally save low-sampling input plots, add both --plot and --plot_targets.

Example usage:
python test.py --data_path ./my_root_path --run_name my_experiment --net_type nafnet --batch_size 8 --gpu_num 0 --n_sampling 1000 --plot

Authors

Minjee Seo, Seonaeng Cho, Minju Seol, Haris Ghafoor and Kyungho Yoon, School of Mathematics and Computing (Computational Science and Engineering), Yonsei University, Seoul, Republic of Korea

Acknowledgement

The work was supported by the National Research Foundation of Korea (NRF) funded by the Korean government (MSIT) under Grants RS-2024-00335185.

License

MIT License

Contact

For any queries, please reach out to Minjee Seo.