Xueh/import code v1

.gitattributes

@@ -0,0 +1 @@
+*.npy filter=lfs diff=lfs merge=lfs -text

.github/workflows/checks.yaml

@@ -22,6 +22,7 @@ jobs:
           ref: ${{ inputs.sha }}
           fetch-depth: 0
           fetch-tags: true
+          lfs: true
       - uses: ./.github/actions/setup-environment
         id: runner-context
       - name: check

.gitignore

@@ -33,6 +33,9 @@ bld/
 [Oo]ut/
 [Ll]og/
 [Ll]ogs/
+lightning_logs/
+run/
+.run/
 
 # Build results on 'Bin' directories
 **/[Bb]in/*

docs/.pages

@@ -0,0 +1,8 @@
+title: SNRAware
+nav:
+    - overview.md
+    - attention_layers.md
+    - block_and_cell.md
+    - backbone.md
+    - examples_and_tests.md
+    - mri_denoising.md

docs/attention_layers.md

@@ -0,0 +1,49 @@
+
+# Attention layers
+
+Given the 5D tensor, different imaging attention modules are developed to capture the signal correlation across spatial and frame/temporal dimensions. All these modules receive a 5D tensor and output a 5D tensor. Unlike the token based attention, these imaging attention modules do not convert pixel data to tokens. That is, the computation is performed in the pixel space, not in token space.
+
+![motivation](images/ifm_model_index.png)
+
+Full attention across three dimensions (F, H, W) can incur high computing costs and may not be optimal as the locality of information was not exploited. Instead, intra- and inter-frame attention can be separately computed. For spatial (along with H, W) attention, a target patch (marked in yellow) takes in information from both close neighbors (marked in green) and remote patches (marked in pink). 
+
+![attention](images/ifm_model_three_attentions.png)
+
+Let the image size be $[H, W]$, the windows size be $[w, w]$ and patch size be $[k, k]$. The number of windows will be $[\frac {H}{w}, \frac{W}{w}]$. The number of patches is $[\frac{H}{k}, \frac{W}{k}]$. For example, for an $256 \times 256$ images, $w=32$ and $k=16$, we will have $8 \times 8$ windows and each window will have $2 \times 2$ patches.
+
+**Spatial Local attention (L)** 
+
+Local attention is computed by attending to the all patches in a window for images or feature maps. The attention matrix is $\frac {w}{k} \times \frac {w}{k}$.
+
+**Spatial Global attention (G)**
+
+While the local attention only explores neighboring pixels, global attention looks at more remote pixels. This will help model learn global information over larger field-of-view. All patches with the same color are attended together. So for every patch, attention matrix size is $w^2 \times w^2$ (that is, the number of windows). 
+
+**Frame or temporal attention (F)**
+
+All frames can attend each other to compute the output values. In this case, the attention matrix size is ${F}\times{F}$.
+
+These **F, G, T** attention mechanisms are implemented as **Attention** modules. The implemented version supports multi-heads.
+
+**Spatial Local 3d attention** 
+
+Previous spatial local attention works on every 2D frame independently. For some applications where data acquisition is 3D, it may require to perform 3D attention. In this case, the [F, H, W] volume is split into 3D windows. Every window includes 3D patches. All patches within a window are attended together.
+
+**Spatial Global 3d attention**
+
+Similar to the 2D global attention, the tensor is split to 3D windows and patches. Corresponding patches from all windows attend each other.
+
+## Other attentions
+
+**Spatial ViT attention** 
+
+The vision transformer method splits the image into windows. All pixels in a window are flattened into a feature vector. Features from all windows are inputted into the attention. This layer works on every 2D frame, as the spatial local attention.
+
+**ViT 3d attention** 
+
+The image is split to 3D window to compute the attention.
+
+**Swin 3d attention**
+
+The shifted window attention is adapted to 3D. The image is split to 3D windows. Every 3d window is split to 3D patches. The attention is computed among all patches within a window or a shifted window. 
+

docs/backbone.md

@@ -0,0 +1,42 @@
+
+# Backbone
+
+Attention - Cell - Block are three levels of building blocks for a model. We can connect many blocks to assemble a backbone. Here the backbone means an "almost" full model. Different pre- and post-processing task heads can be added to a backbone, given a specific application. 
+
+In the LLMs, the stack of attentions proves to be very effective. For imaging, previous researches had explored similar architecture (e.g. [3B swin v2](https://arxiv.org/abs/2111.09883) and [20B ViT](https://arxiv.org/abs/2302.05442)). On the other hand, there are very intuitive model architectures were invented for different convolution models. Comparing to the LLM type backbone, these CNN derived architectures are more tuned for image data format, by utilizing resolution pyramid, up and downsampling, and long range skip connections. 
+
+Different backbone models are implemented by borrowing successful architectures from CNNs.
+
+### backbone U-Net
+
+![unet](images/Unet.jpg)
+
+Here every resolution stage includes one block containing multiple cells. Model can specify number of feature maps (Channel dimension C) at each resolution stage. 
+
+Return tensor of this backbone is the final output tensor after U0 layer.
+
+The [Unet with attention](https://arxiv.org/abs/1804.03999) is implemented here. Downsample and upsample are implemented with interpolation.
+
+![Attention_in_Unet](images/Attention_in_Unet.jpg)
+
+### backbone HR-Net
+
+This network is modified from the [high-resolution architecture](https://www.microsoft.com/en-us/research/blog/high-resolution-network-a-universal-neural-architecture-for-visual-recognition/).
+
+![stcnnt_hrnet](images/HRNet.jpg)
+
+The network is defined as levels and stages. Every block is numbered by its level and stage indexes (starting from 0). The downsample and upsample modules are added to link different blocks. Different up/downsample modules are implemented, with TLG attentions or 1x1 CONV. Bilinear interpolation is used to alter spatial resolution.
+
+After the fusion stage, the model outputs per-level tensors and the aggregated tensor as a list. This backbone returns a list of tensors. The outputs from every resolution level and the final aggregated tensor are returned as a list *res*. *res[0]* is the output with the highest resolution (*y_hat_0*). *res[1]* is the *y_hat_1* etc. The *res[-1]* is the aggregated tensor.
+
+### backbone Stack-of-attention (SOANet)
+
+The classical stack-of-attention model is defined by the stages. Each stage contains one block with many cells. The downsampling layers are optional between stages.
+
+We can instantiate the typical Swin or ViT model by inserting corresponding cells to SOANet. For example, a four-stage model with block strings ["S3ShS3Sh", "S3ShS3Sh", "S3ShS3ShS3ShS3ShS3ShS3Sh", "S3ShS3Sh"] is a replica of published model in the [swin paper](https://arxiv.org/abs/2103.14030).A one stage model ["V2V2V2V2V2V2V2V2V2V2V2V2"] is a version of 2D ViT base model with 12 cells. Since the backbone takes in the 5D tensor, it is often better to instantiate 3D ViT model, e.g. ["V3V3V3","V3V3V3V3V3V3", "V3V3V3V3V3V3V3V3V3V3V3V3"] for a 3 stage 3D ViT.
+
+This backbone returns a list of tensors containing outputs from all stages. *res[0]* is the first stage and *res[-1]* has the lowest resolution (if downsampling).
+
+![soanet](images/SOANet.jpg)
+
+**Please ref to [examples and tests](./examples_and_tests.md) for how to instantiate the backbone.**

docs/block_and_cell.md

@@ -0,0 +1,69 @@
+
+# Cell and Block
+
+## Cell
+A cell is a container to host an attention layer, as the basic processing unit for the 5D tensor. As the following, a cell can host either F, L or G or other attentions. Input tensor goes through the norm layers and attention. The mixer consists of another norm and two convolutions. 
+
+![Cell](./images/ifm_model_cell.png)
+
+Many attention modules are supported. These modules are represented by its name string:
+
+| Name | Description |
+| ----------- | ----------- |
+| V2 | 2D ViT attention, working on [H, W], with the mixer |
+| V3 | 3D ViT attention, working on [D/Z/T, H, W], with the mixer |
+| L1 | spatial local attention on [H, W], with mixer |
+| L0 | spatial local attention on [H, W], without mixer |
+| L3 | local 3D attention on [D/Z/T, H, W], with mixer |
+| G1 | spatial global attention on [H, W], with mixer |
+| G0 | spatial global attention on [H, W], without mixer |
+| G3 | global 3D attention on [D/Z/T, H, W], with mixer |
+| S3 | swin 3D attention over [D/Z/T, H, W], with the mixer|
+| Sh | swin 3D attention over [D/Z/T, H, W], with the mixer, shifted window|
+
+#### Norms
+
+Different $Norm$ can be configured: 
+
+| Norm | Description |
+| ----------- | ----------- |
+| LayerNorm | normalize over [C, H, W] |
+| BatchNorm 2D | normalize over [H, W], B*T are batch dimension |
+| BatchNorm 3D | normalize over [T, H, W] |
+| InstanceNorm 2D | normalize over [H, W]|
+| InstanceNorm 3D | normalize over [T, H, W]|
+
+Except the *layer* norm, all other norms support flexible image sizes, when using with the $CONV$ in attention layers.
+
+#### Mixers
+
+The mixers are added after attention layers to increase model power. Two types of mixers are implemented. *Linear* mixer are the same mixer type used in the conventional transformer, with `torch.linear` layers. The "conv" mixer replaced Linear layer with the convolution layer, to better fit for imaging tensors. The conv is performed along the $[C, H, W]$:
+
+```
+self.mlp = nn.Sequential(
+                Conv2DExt(C_out, 4*C_out, kernel_size=kernel_size, stride=stride, padding=padding, bias=True),
+                nn.GELU(),
+                Conv2DExt(4*C_out, C_out, kernel_size=kernel_size, stride=stride, padding=padding, bias=True),
+                nn.Dropout(dropout_p),
+            )
+```
+
+User can specify whether a cell has mixer or not.
+
+## Block
+
+Many cells are concatenated to enhance the model power. To standardize this configuration, the Block is introduced:
+
+![Block](./images/ifm_model_block.png)
+
+A block contains any number of cells and can be scaled up by inserting more cells. 
+
+A block is coded by the acronyms of attention layers in each cell. For example, a block is coded as the block string "L1G1T1". The letter "L", "G" or "T" means the local, global and Frame/Temporal Cell. "1" means the mixer is added on top of the attention mechanism (if "0", mixer is not added; we can have a stack of attention only layers). As an example, "L0L1G0G1T1" means a block with 5 attention layers. Mixers are not added to the first and third attentions, but added after the second and fourth attentions. The last cell is a temporal attention with its mixer added. This method of "block string specification" gives a good amount of flexibility to assemble and test different attention configurations. Depending on the image resolutions and number of feature maps, different blocks in a model can have different attention configuration. Note the convolution layers are formatted into this framework as the "C2" or "C3" cells. "C2" means the 2D conv is used in the cell. "C3" means the 3D conv is used (over F, H, W). By changing the block string from attentions to convs, we can get identical architectures with only differences being the cell structure. 
+
+The block string can also be "S2" or "S3" for 2D or 3D SWIN module. Each SWIN module includes two parts for concatenated normal and shifted window attentions. These attention modules can be concatenated in any order of mixing. For example, *L1G1T1* or *S3V2L1G1T1*, or *G1G3S3V2T1* etc. *S3ShS3ShS3Sh* is the SWIN transformer with normal and shifted window attentions interleaved. *V3V3V3V3* is the ViT transformer. By separating the backbone architecture with attention modules, it is very flexible to instantiate many types of models, tailored to the nature of datasets and amount of computing. 
+
+Two cell layouts were implemented: sequential and parallel:
+
+![Cell](./images/Cell.jpg)
+
+The reference for parallel cell is [here](https://arxiv.org/abs/2302.05442). 

docs/examples_and_tests.md

@@ -0,0 +1,89 @@
+
+# Examples and tests
+
+The modules and backbones in this package can be used to build models for different applications. The key design is to support 5D tensor [B, C, T/D/F/Z, H, W] for 2D, 2D+T, and 3D usecases.
+
+## Attention modules
+
+All attention modules are implemented as pytorch modules. They can be defined in the very straight-forward way as the following:
+
+```python
+import numpy as np
+import torch
+from snraware.components.setup import start_timer, end_timer, get_device
+from snraware.components.model.attention import ViT3DAttention
+
+# find the device, e.g. 'cuda:0'
+device = get_device()
+
+# define an input tensor
+B, C, T, H, W = 1, 2, 16, 32, 32
+C_out = 8
+test_in = torch.rand(B, C, T, H, W).to(device=device)
+print(test_in.shape)
+
+# define a vit3d module
+
+spacial_vit = ViT3DAttention(window_size=None, 
+                            num_wind=[8, 8, 4], # here we choose to set number of windows; alternatively, we can set window_size to be [2, 2, 4]
+                            attention_type="conv", # use convolution to compute Q/K/V in the attention
+                            C_in=2, # input channel, 2
+                            C_out=8, # output channel, 8
+                            H=H, W=W, D=T, # tensor sizes
+                            stride_qk=[1,1,1], # stride is 1 when computing Q and K
+                            cosine_att=True, # use the cosine attention
+                            normalize_Q_K=False, # normalize Q and K before computing attention matrix; if cosine_att is True, this option 
+                            att_with_relative_position_bias=True,
+                            att_with_output_proj=True)
+
+spacial_vit.to(device=device)
+
+test_out = spacial_vit(test_in)
+
+print(f"ViT3DAttention, input tensor {test_in.shape}, output tensor {test_out.shape}")
+
+```
+
+## Backbones
+
+The backbone models contain multiple Blocks and Cells. They are configured by supplying the block_str to define the attention structures. 
+
+Unlike the attention modules where parameters are listed out, the model has many more parameters. We use the [hydra](https://hydra.cc/) to manage the parameters. The parameter configuration files are in the `src/ifm/configs` folder. To define the model parameters for a backbone model:
+
+```
+from hydra import initialize, compose
+with initialize(version_base=None, config_path="../src/snraware/components/configs"):
+    config = compose(config_name="config")
+```
+
+The backbone parameters are in `config.backbone`. The block and cell parameters are in `config.block` and `config.cell`.
+
+```python
+
+import torch
+from snraware.components.setup import get_device
+import hydra
+from hydra import initialize, compose
+
+# get the device, e.g. 'cuda:0'
+device = get_device()
+
+# define the input tensor
+B, C, T, H, W = 1, 2, 32, 128, 128
+test_in = torch.rand(B, C, T, H, W).to(device=device)
+
+# get the mode parameter; note user can override the default configuration
+# here we order a hrnet backbone
+with initialize(version_base=None, config_path="../src/snraware/components/configs"):
+    cfg = compose(config_name="config")
+
+config = hydra.utils.instantiate(cfg.backbone)
+
+# define the model
+model = HRnet(config=config, input_feature_channels=C, H=H, W=W, D=T)
+model = model.to(device=device)
+
+# perform the forward pass
+test_out = model(test_in)
+print(f"{Fore.YELLOW}The output tensor shape is {test_out[-1].shape}.{Style.RESET_ALL}")
+```

docs/index.md

@@ -1,6 +1,79 @@
----
-title: about
----
 # SNRAware
 
-Deep Learning MR Denoising Traning and Imaging Transformer Models
+This repository contains the Pytorch code in our paper [SNRAware: Improved Deep Learning MRI Denoising with Signal-to-noise Ratio Unit Training and G-factor Map Augmentation](https://pubs.rsna.org/doi/full/10.1148/ryai.250227) published at the Radiology: Artificial Intelligence:
+
+```latex
+@article{
+    doi:10.1148/ryai.250227,
+    author = {Xue, Hui and Hooper, Sarah M. and Pierce, Iain and Davies, Rhodri H. and Stairs, John and Naegele, Joseph and Campbell-Washburn, Adrienne E. and Manisty, Charlotte and Moon, James C. and Treibel, Thomas A. and Hansen, Michael S. and Kellman, Peter},
+    title = {SNRAware: Improved Deep Learning MRI Denoising with Signal-to-noise Ratio Unit Training and G-factor Map Augmentation},
+    journal = {Radiology: Artificial Intelligence},
+    volume = {0},
+    number = {ja},
+    pages = {e250227},
+    year = {0},
+    doi = {10.1148/ryai.250227},
+    note ={PMID: 41123451},
+    URL = {https://doi.org/10.1148/ryai.250227}
+}
+```
+
+- Model type: Imaging AI, non-generative
+- License: MIT
+
+## Get started
+
+[just](https://github.com/casey/just) is used in this project. If not, please install this tool:
+
+```bash
+# install just
+wget -qO - 'https://proget.makedeb.org/debian-feeds/prebuilt-mpr.pub' | gpg --dearmor | sudo tee /usr/share/keyrings/prebuilt-mpr-archive-keyring.gpg 1> /dev/null
+echo "deb [arch=all,$(dpkg --print-architecture) signed-by=/usr/share/keyrings/prebuilt-mpr-archive-keyring.gpg] https://proget.makedeb.org prebuilt-mpr $(lsb_release -cs)" | sudo tee /etc/apt/sources.list.d/prebuilt-mpr.list
+sudo apt update
+sudo apt install just -y
+```
+
+Then, please set up the virtual environment and run tests:
+
+```bash
+# show the list
+just --list
+
+# set up virtual environment
+just setup-env
+
+# review documentation
+just serve-docs
+
+# run test
+just test
+```
+
+## Data
+Dataset for MR denoising training is not opened at this moment.
+
+## Model
+Three models are released at https://huggingface.co/microsoft/SNRAware
+
+- SNRAware-small: 27.7million parameters
+- SNRAware-medium: 55.1million parameters
+- SNRAware-large: 109million parameters
+
+## Direct intended uses
+SNRAware is shared for research and technical development purposes only, to denoise MR images.
+
+## License and Usage Notices
+The data, code, and model checkpoints described in this repository is provided for research and technical development use
+only. The data, code, and model checkpoints are not intended for use in clinical use.
+
+## Trademarks
+
+This project may contain trademarks or logos for projects, products, or services. Authorized use of Microsoft 
+trademarks or logos is subject to and must follow 
+[Microsoft's Trademark & Brand Guidelines](https://www.microsoft.com/en-us/legal/intellectualproperty/trademarks/usage/general).
+Use of Microsoft trademarks or logos in modified versions of this project must not cause confusion or imply Microsoft sponsorship.
+Any use of third-party trademarks or logos are subject to those third-party's policies.
+
+## Documentation
+
+Please find documentation in the [docs/overview](./docs/overview.md).