Add scripts for deep learning reconstruction of prostate scans using Varnet

Author: tarund1996

.gitignore

@@ -127,3 +127,10 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# log directory
+logs/
+varnet/
+
+# bash files
+*.sh

DL_reconstruction/README.md

@@ -0,0 +1,42 @@
+# Prostate Diffusion MRI Reconstruction
+
+This repository contains code for reconstruction accelerated prostate diffusion MRI, derived from the [fastMRI](https://github.com/facebookresearch/fastMRI) repository and modified specifically for the fastmri prostate diffusion data. Our model supports reconstruction for b50 and b1000 diffusion images, outputting b50 or b1000 trace images accordingly.
+
+## Features
+
+- **Data Compatibility**: Designed for reconstruction of prostate diffusion MRI data, specifically compatible with the fastMRI prostate dataset.
+- **Model Output**: Generates b50 or b1000 trace images
+- **Based on fastMRI**: Leverages the framework and model from the fastMRI repository.
+
+## Usage
+
+To run the model, you can use the following commands for training and testing:
+
+### Training
+
+```bash
+python train_varnet_demo.py --mode train --data_path [path_to_data] --bvalue b50 --test_path [path_to_test_data] --state_dict_file [path_to_checkpoint_for_testing] --batch_size 1 --num_workers 4
+```
+
+Replace `[script_name]` with the name of your script, `[path_to_data]` with the path to your training data, `[path_to_test_data]` with the path to your testing data, and `[path_to_checkpoint_for_testing]` with the path to a model checkpoint file if you have one for testing.
+
+### Testing
+
+```bash
+python train_varnet_demo.py --mode test --data_path [path_to_data] --bvalue b1000 --test_path [path_to_test_data] --state_dict_file [path_to_checkpoint] --batch_size 1 --num_workers 4
+```
+
+Ensure you specify the `--state_dict_file` argument during testing to provide the path to your model checkpoint.
+
+## Configuration
+
+The script accepts various command-line arguments to customize the data paths, model parameters, and training settings. Refer to the source code for a complete list of available options.
+
+## Contributing
+
+Contributions to improve this project are welcome. Please consider submitting a pull request or opening an issue for any bugs or feature requests.
+
+## License
+
+This project is licensed under the MIT License - see the LICENSE file for details.
+

DL_reconstruction/__init__.py

@@ -0,0 +1,21 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+from .losses import SSIMLoss
+from .utils import save_reconstructions
+from .coil_combine import rss, rss_complex
+from .fftc import fft2c_new as fft2c
+from .fftc import fftshift
+from .fftc import ifft2c_new as ifft2c
+from .fftc import ifftshift, roll
+from .math_fn import (
+    complex_abs,
+    complex_abs_sq,
+    complex_conj,
+    complex_mul,
+    tensor_to_complex_np,
+)

DL_reconstruction/coil_combine.py

@@ -0,0 +1,41 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import torch
+
+from .math_fn import complex_abs_sq
+
+def rss(data: torch.Tensor, dim: int = 0) -> torch.Tensor:
+    """
+    Compute the Root Sum of Squares (RSS).
+
+    RSS is computed assuming that dim is the coil dimension.
+
+    Args:
+        data: The input tensor
+        dim: The dimensions along which to apply the RSS transform
+
+    Returns:
+        The RSS value.
+    """
+    return torch.sqrt((data ** 2).sum(dim))
+
+
+def rss_complex(data: torch.Tensor, dim: int = 0) -> torch.Tensor:
+    """
+    Compute the Root Sum of Squares (RSS) for complex inputs.
+
+    RSS is computed assuming that dim is the coil dimension.
+
+    Args:
+        data: The input tensor
+        dim: The dimensions along which to apply the RSS transform
+
+    Returns:
+        The RSS value.
+    """
+    return torch.sqrt(complex_abs_sq(data).sum(dim))

DL_reconstruction/data/__init__.py

@@ -0,0 +1,9 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+from .mri_data import SliceDataset
+from .volume_sampler import VolumeSampler