update and fix

.github/workflows/codeql.yaml

@@ -23,7 +23,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        language: [ "javascript-typescript", "python" ]
+        language: [ "python" ]
     steps:
     - uses: actions/checkout@v5
       with:

.github/workflows/main.yaml

@@ -28,16 +28,6 @@ jobs:
     secrets: inherit
     with:
       sha: ${{ inputs.sha || github.sha }}
-  build-wheels:
-    uses: ./.github/workflows/build-wheels.yaml
-    secrets: inherit
-    with:
-      sha: ${{ inputs.sha || github.sha }}
-  build-containers:
-    uses: ./.github/workflows/build-containers.yaml
-    secrets: inherit
-    with:
-      sha: ${{ inputs.sha || github.sha }}
   build-docs:
     uses: ./.github/workflows/build-docs.yaml
     secrets: inherit
@@ -56,10 +46,8 @@ jobs:
     uses: ./.github/workflows/release.yaml
     secrets: inherit
     needs:
-      - build-containers
       - build-docs
-      - build-wheels
       - checks
       - codeql
     with:
-      publish-docs: ${{ github.ref_name == 'main' }}
+      publish-docs: ${{ github.ref_name == 'main' }}

docs/index.md

@@ -8,10 +8,10 @@ This repository contains the Pytorch code in our paper [SNRAware: Improved Deep
     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},
+    volume = {7},
+    number = {6},
     pages = {e250227},
-    year = {0},
+    year = {2025},
     doi = {10.1148/ryai.250227},
     note ={PMID: 41123451},
     URL = {https://doi.org/10.1148/ryai.250227}
@@ -23,34 +23,47 @@ This repository contains the Pytorch code in our paper [SNRAware: Improved Deep
 
 ## Get started
 
-[just](https://github.com/casey/just) is used in this project. If not, please install this tool:
+`uv` is used in this project. Please install it as:
 
 ```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
+# install uv
+curl -LsSf https://astral.sh/uv/install.sh | sh
+
+# install git-lfs
 sudo apt update
-sudo apt install just -y
+sudo apt install git-lfs direnv
 ```
 
-Then, please set up the virtual environment and run tests:
+Make sure commands `uv` are on your path.
+
+Also, this project requires NVIDIA GPU. To check whether your GPU is available and is working:
 
 ```bash
-# show the list
-just --list
+nvidia-smi
+```
+If the GPU is working correctly, this command will display detailed information, including driver version, GPU usage, memory usage, and temperature.
 
-# set up virtual environment
-just setup-env
+Make sure the command `uv` are on your path. Then please clone the repo, set up the virtual environment and run tests:
 
-# review documentation
-just serve-docs
+```bash
+# clone the repo
+git clone git@github.com:microsoft/SNRAware.git
 
-# run test
-just test
+# set up env
+direnv allow
+cd ./SNRAware
+uv sync
+
+# pull down test data
+git lfs pull
+
+# run the test
+uv run pytest -m gpu ./test
 ```
 
-## Data
-Dataset for MR denoising training is not opened at this moment.
+## Training data
+
+Dataset for MR denoising training is not opened at this moment. More information will be provided once training data is released.
 
 ## Model
 Three models are released at https://huggingface.co/microsoft/SNRAware
@@ -59,8 +72,34 @@ Three models are released at https://huggingface.co/microsoft/SNRAware
 - SNRAware-medium: 55.1million parameters
 - SNRAware-large: 109million parameters
 
+To test the model, 
+```bash
+# download the model from the huggingface
+# small model
+wget https://huggingface.co/microsoft/SNRAware/resolve/main/small/snraware_small_model.pts
+wget https://huggingface.co/microsoft/SNRAware/resolve/main/small/snraware_small_model.yaml
+
+# a test data is provided at ./test/data/inference
+# input data are [H, W, Frame] 3D complex tensor, input_real.npy and input_imag.npy store the 
+# real and imaginary part
+# gmap.npy is the g-factor map for all frames or for every frame, [H, W, 1 or Frame]
+
+# let's use the small model to run a inference
+export model_file=snraware_small_model.pts
+export config_file=snraware_small_model.yaml
+
+# run the inference
+uv run python3 ./src/snraware/projects/mri/denoising/run_inference.py --input_dir ./test/data/phantom --output_dir /tmp/phantom_res_inference --saved_model_path $model_file --saved_config_path $config_file --batch_size 1 --input_fname input --gmap_fname gmap
+```
+
+After the run, the result is saved in the `/tmp/phantom_res_inference` as numpy files. 
+
+![alt text](./docs/images/image.png)
+
+raw, model output, difference
+
 ## Direct intended uses
-SNRAware is shared for research and technical development purposes only, to denoise MR images.
+SNRAware is shared for research and technical development purposes only, to denoisegit  MR images.
 
 ## License and Usage Notices
 The data, code, and model checkpoints described in this repository is provided for research and technical development use

docs/mri_denoising.md

@@ -45,14 +45,4 @@ User needs to log into the wandb as `wandb login`. Training and validation sampl
 
 # Run inference
 
-After training the model, user can run inference with `run_inference.py` in the `src/snraware/projects/mri/denoising` folder:
-
-```bash
-export model_file="SNRAware-small.pts"
-export config_file="SNRAware-small_config.yaml"
-export res_dir="res"
-export data_dir="/data"
-
-python3 ./src/snraware/projects/mri/denoising/run_inference.py --input_dir ${data_dir} --output_dir ${data_dir}/${res_dir} --saved_model_path $model_file --saved_config_path $config_file --batch_size 1 --input_fname input --gmap_fname gmap
-
-```
+After training the model, user can run inference with `run_inference.py` in the `src/snraware/projects/mri/denoising` folder. Examples to run the model inference is given in the [README](../README.md).

pyproject.toml

@@ -3,15 +3,16 @@ name = "snraware"
 description = "A deep learning imaging AI framework"
 version = "0.1.0"
 readme = "README.md"
-requires-python = ">=3.12,<4.0"
+requires-python = ">=3.12,<=3.14"
 authors = [
     {name = "Hui Xue", email = "xueh@microsoft.com"},
 ]
 keywords = []
 license = "MIT"
 dependencies = [
-    "mkl-fft>=2.0.0",
-    "mkl-service>=2.4.2",
+    "mkl",
+    "mkl-fft",
+    "mkl-service",
     "typing_extensions",
     "click",
     "h5py",
@@ -78,15 +79,6 @@ build-backend = "uv_build"
 [tool.uv.build-backend]
 module-name = ["snraware"]
 
-[[tool.uv.index]]
-name = "mkl"
-url = "https://urob.github.io/numpy-mkl"
-
-[tool.uv.sources]
-numpy = { index = "mkl" }
-scipy = { index = "mkl" }
-mkl-service = { index = "mkl" }
-
 [project.scripts]
 
 [tool.pyright]

test/components/test_backbone_hrnet.py

@@ -57,6 +57,9 @@ def test(self, backbone):
         with_timer = True
         device = get_device()
 
+        if device != "cuda":
+            pytest.skip("GPU only test")
+
         _B, C, T, H, W = 1, 2, 16, 16, 16
         test_in = torch.from_numpy(self.test_in).to(dtype=torch.float32, device=device)
         assert np.linalg.norm(self.test_in - test_in.cpu().numpy()) < 1e-3

test/components/test_backbone_soanet.py

@@ -55,6 +55,8 @@ def teardown_class(self):
     def test(self, downsample, backbone):
         with_timer = True
         device = get_device()
+        if device != "cuda":
+            pytest.skip("GPU only test")
 
         _B, C, T, H, W = 1, 4, 8, 32, 32
         test_in = torch.from_numpy(self.test_in).to(dtype=torch.float32, device=device)

test/components/test_backbone_unet.py

@@ -55,6 +55,8 @@ def teardown_class(self):
     def test(self, backbone):
         with_timer = True
         device = get_device()
+        if device != "cuda":
+            pytest.skip("GPU only test")
 
         _B, C, T, H, W = 1, 2, 16, 32, 32
         test_in = torch.from_numpy(self.test_in).to(dtype=torch.float32, device=device)

test/components/test_block.py

@@ -72,6 +72,9 @@ def test(self, block_str):
         test_in = torch.rand(B, T, C, H, W)
 
         device = get_device()
+        if device != "cuda":
+            pytest.skip("GPU only test")
+
         test_in = test_in.to(device=device, dtype=torch.float32)
         test_in = torch.permute(test_in, [0, 2, 1, 3, 4])
 

test/components/test_cell.py

@@ -31,6 +31,9 @@ def test(self):
         test_in = torch.rand(B, T, C, H, W).to(torch.float32)
 
         device = get_device()
+        if device != "cuda":
+            pytest.skip("GPU only test")
+
         test_in = test_in.to(device=device)
 
         att_types = [

test/components/test_convolution_module.py

@@ -39,6 +39,8 @@ def test(self):
 
         with_timer = True
         device = get_device()
+        if device != "cuda":
+            pytest.skip("GPU only test")
 
         test_in = torch.rand(B, T, C, H, W, device=device)
         assert np.linalg.norm(self.test_in - test_in.cpu().numpy()) < 1e-3

test/components/test_global_3d_attention.py

@@ -29,6 +29,10 @@ def teardown_class(self):
 
     @pytest.mark.gpu
     def test(self):
+        device = get_device()
+        if device != "cuda":
+            pytest.skip("GPU only test")
+
         t = np.arange(256)
         t = np.reshape(t, (16, 16))
 
@@ -72,8 +76,6 @@ def test(self):
 
         with_timer = True
 
-        device = get_device()
-
         B, T, C, H1, W1 = 1, 16, 2, 64, 64
         C_out = 8
         test_in = torch.rand(B, T, C, H1, W1).to(device=device)