UI updated

app.py

@@ -94,13 +94,14 @@ st.set_page_config(
 
 # Sidebar content
 with st.sidebar:
-    st.title("Segmenting vessels in the brain from a 3D Magnetic Resonance Angiograph, ideally acquired at 7T | DS6")
+    st.title("Segmenting vessels in the brain from a 3D Magnetic Resonance Angiograph | DS6")
     st.markdown("""
     This application allows you to upload a 3D NIfTI file (dims: H x W x D, where the final dim is the slice dim in the axial plane), process it through a pre-trained 3D model (from DS6 and other related works), and download the output as a `.nii.gz` file containing the vessel segmentation.
     
     **Instructions**:
-    - Upload your 3D NIfTI file (`.nii` or `.nii.gz`). It should be a single-slice cardiac long-axis dynamic CINE scan, where the first dimension represents time.
-    - Select a model from the dropdown menu.
+    - Upload your 3D NIfTI file (`.nii` or `.nii.gz`). The model was trained on `7T MRA-ToF` data, but it should work on other field strengths as well.
+    - Select a pretrained model from the dropdown menu.
+    - Select the inference mode (full volume or patch-based) from the dropdown menu.
     - Click the "Process" button to generate the latent factors.
     """)
     st.markdown("---")
@@ -108,7 +109,34 @@ with st.sidebar:
 
 # Main content
 st.header("DS6, Deformation-Aware Semi-Supervised Learning: Application to Small Vessel Segmentation with Noisy Training Data")
+st.markdown("""
+This application can be used to perform vessel segmentation by uploading a `Magnetic Resonance Angiograph (MRA-ToF)`, ideally acquired at 7T, but it should also work on other field strengths, in NIFTI format (.nii or .nii.gz).
 
+The list of weights includes those from the original [DS6 paper](https://doi.org/10.3390/jimaging8100259), from the [SMILE-UHURA challenge](https://doi.org/10.7303/syn47164761), from a follow-up research [SPOCKMIP](https://arxiv.org/abs/2407.08655), as well as a [fine-tuned version](https://www.medrxiv.org/content/10.1101/2024.10.03.24314845v1) of the DS6 model (trained on the SMILE-UHURA dataset) using the Cambridge 7T Cerebral Small Vessel Disease (CamSVD) dataset, which contains data from subjects with lacunar strokes with SVD, non-lacunar strokes without SVD, and healthy controls.
+
+Segmentation can be performed in 2 different inference modes: full volume inference and patch-based inference. All these research works used patch-based inference. However, if the volume is not large enough, and only large vessels are to be segmented, full volume inference may be performed, which is significantly faster. Nevertheless, full volume inference might result in an out-of-memory error (if the volume is very large) and may not segment the small vessels properly.  
+""")
+with st.expander("List of available pretrained models"):
+    st.markdown(
+        """
+        | Model Name | Description |
+        |------------|-------------|
+        | `DS6_UNet3D_woDeform` | UNet 3D, trained without deformation-aware learning |
+        | `DS6_UNetMSS3D_woDeform` | UNet MSS 3D, trained without deformation-aware learning |
+        | `DS6_UNetMSS3D_wDeform` | UNet MSS 3D, trained with deformation-aware learning [Proposed method, DS6] |
+        | `SMILEUHURA_DS6_UNet3D_woDeform` | UNet 3D, trained without deformation-aware learning on the SMILE-UHURA dataset |
+        | `SMILEUHURA_DS6_UNetMSS3D_woDeform` | UNet MSS 3D, trained without deformation-aware learning on the SMILE-UHURA dataset |
+        | `SMILEUHURA_DS6_UNetMSS3D_wDeform` | UNet MSS 3D, trained with deformation-aware learning on the SMILE-UHURA dataset |
+        | `SMILEUHURA_neuRoSliCCe_SPOCKMIP_UNetMSS3D_MIP` | UNet MSS 3D, trained with MIP (Maximum Intensity Projection) loss [Proposed method, SPOCKMIP] |
+        | `SMILEUHURA_neuRoSliCCe_SPOCKMIP_UNetMSS3D_mMIP` | UNet MSS 3D, trained with Multi-axis MIP loss [Proposed method, SPOCKMIP] |
+        | `SMILEUHURA_SPOCKMIP_UNet3D_MIP` | UNet 3D, trained with MIP (Maximum Intensity Projection) loss [Proposed method, SPOCKMIP] |
+        | `SMILEUHURA_SPOCKMIP_UNet3D_mMIP` | UNet 3D, trained with Multi-axis MIP loss [Proposed method, SPOCKMIP] |
+        | `SMILEUHURA_neuRoSliCCe_SPOCKMIP_UNetMSS3D_DS6MIP` | UNet MSS 3D, trained with deformation-aware learning, and then with MIP loss |
+        | `SMILEUHURA_DS6_CamSVD_UNetMSS3D_wDeform` | UNet MSS 3D, initially trained with deformation-aware learning on the SMILE-UHURA dataset, and then fine-tuned on the CamSVD dataset with deformation-aware learning |
+        """
+    )
+st.markdown("---")
+    
 # File uploader
 uploaded_file = st.file_uploader(
     "Please upload a 3D NIfTI file (.nii or .nii.gz)",
@@ -116,15 +144,28 @@ uploaded_file = st.file_uploader(
 )
 
 # Model selection
-model_options = ["SMILEUHURA_DS6_CamSVD_UNetMSS3D_wDeform"]
+model_options = [
+    "DS6_UNet3D_woDeform",
+    "DS6_UNetMSS3D_woDeform",
+    "DS6_UNetMSS3D_wDeform",
+    "SMILEUHURA_DS6_UNet3D_woDeform",
+    "SMILEUHURA_DS6_UNetMSS3D_woDeform",
+    "SMILEUHURA_DS6_UNetMSS3D_wDeform",
+    "SMILEUHURA_neuRoSliCCe_SPOCKMIP_UNetMSS3D_MIP",
+    "SMILEUHURA_neuRoSliCCe_SPOCKMIP_UNetMSS3D_mMIP",
+    "SMILEUHURA_SPOCKMIP_UNet3D_MIP",
+    "SMILEUHURA_SPOCKMIP_UNet3D_mMIP",
+    "SMILEUHURA_neuRoSliCCe_SPOCKMIP_UNetMSS3D_DS6MIP",
+    "SMILEUHURA_DS6_CamSVD_UNetMSS3D_wDeform"
+]
 selected_model = st.selectbox("Select a pretrained model:", model_options)
 
 # Mode selection
-mode_options = ["Full volume inference", "Patch-based inference [Default for DS6]"]
-selected_mode = st.selectbox("Select the running mode:", mode_options)
+mode_options = ["Full volume inference", "Patch-based inference [Default for all the published works]"]
+selected_mode = st.selectbox("Select the inference mode:", mode_options)
 
 # Parameters for patch-based inference
-if selected_mode == "Patch-based inference [Default for DS6]":
+if selected_mode == "Patch-based inference [Default for all the published works]":
     col1, col2, col3 = st.columns(3)
     with col1:
         patch_size = st.number_input("Patch size:", min_value=1, value=64)
@@ -201,7 +242,7 @@ if uploaded_file is not None and process_button:
                     st.info("Running full volume inference...")
                     output = infer_full_vol(tensor, model)
                 else:
-                    st.info("Running patch-based inference [Default for DS6]...")
+                    st.info("Running patch-based inference [Default for all the published works]...")
                     output = infer_patch_based(tensor, model, patch_size=patch_size, stride_length=stride_length, stride_width=stride_width, stride_depth=stride_depth, batch_size=batch_size, num_worker=num_worker)
             
             st.success("Processing complete.")
@@ -237,4 +278,69 @@ if uploaded_file is not None and process_button:
 elif uploaded_file is None:
     st.info("Awaiting file upload...")
 elif not process_button:
-    st.info("Click the 'Process' button to start processing.")
+    st.info("Click the 'Process' button to start processing.")
+    
+# Footer
+st.markdown(
+    """
+    ---
+    ## Credits
+    If you like this application, please click on **"Like"** on the top left!
+    
+    If you use this application and/or any of these models, please cite the following paper:
+
+    ```
+    @Article{chatterjee2022ds6,
+            AUTHOR = {Chatterjee, Soumick and Prabhu, Kartik and Pattadkal, Mahantesh and Bortsova, Gerda and Sarasaen, Chompunuch and Dubost, Florian and Mattern, Hendrik and de Bruijne, Marleen and Speck, Oliver and Nürnberger, Andreas},
+            TITLE = {DS6, Deformation-Aware Semi-Supervised Learning: Application to Small Vessel Segmentation with Noisy Training Data},
+            JOURNAL = {Journal of Imaging},
+            VOLUME = {8},
+            YEAR = {2022},
+            NUMBER = {10},
+            ARTICLE-NUMBER = {259},
+            URL = {https://www.mdpi.com/2313-433X/8/10/259},
+            ISSN = {2313-433X},
+            DOI = {10.3390/jimaging8100259}
+    }
+    ```
+    
+    If you use one of the models with the name starting with `SMILEUHURA`, please addiitonally cite the following paper:
+
+    ```
+    @article{chatterjee2023smile,
+        title={SMILE-UHURA Challenge},
+        author={Chatterjee, S and Mattern, H and Dubost, F and Schreiber, S and Nürnberger, A and Speck, O},
+        year={2023},
+            doi = {10.7303/syn47164761},
+        URL = {https://doi.org/10.7303/syn47164761}
+        }
+    ```
+    
+    If you use one of the models that contains `SPOCKMIP` in its name, please addiitonally cite the following paper:
+
+    ```
+    @article{radhakrishna2024spockmip,
+        title={SPOCKMIP: Segmentation of Vessels in MRAs with Enhanced Continuity using Maximum Intensity Projection as Loss},
+        author={Radhakrishna, Chethan and Chintalapati, Karthikesh Varma and Kumar, Sri Chandana Hudukula Ram and Sutrave, Raviteja and Mattern, Hendrik and Speck, Oliver and N{\"u}rnberger, Andreas and Chatterjee, Soumick},
+        journal={arXiv preprint arXiv:2407.08655},
+        year={2024}
+        }
+    ```
+    
+    If you use the `SMILEUHURA_DS6_CamSVD_UNetMSS3D_wDeform` model (i.e. fine-tuned on the CamSVD dataset), please addiitonally cite the following paper:
+
+    ```
+    @article{ruiDS62024,
+            author = {Li, Rui and Chatterjee, Soumick and Jiaerken, Yeerfan and Radhakrishna, Chethan and Benjamin, Philip and Nannoni, Stefania and Tozer, Daniel J. and Markus, Hugh and Rodgers, Christopher T.},
+            title = {A Deep Learning Pipeline for Analysis of the 3D Morphology of the Cerebral Small Perforating Arteries from Time-of-Flight 7 Tesla MRI},
+            year = {2024},
+            doi = {10.1101/2024.10.03.24314845},
+            publisher = {Cold Spring Harbor Laboratory Press},
+            URL = {https://www.medrxiv.org/content/early/2024/10/04/2024.10.03.24314845},
+            journal = {medRxiv}
+        }
+
+    ```
+    
+    """
+)