Update VORTEX demo and README (#103)

* update readme and demo

* udpate readme

README.md

@@ -39,6 +39,18 @@ os.environ["MEDDLR_DATASETS_DIR"] = "/path/to/datasets"
 
 Detailed instructions are available in [Getting Started](GETTING_STARTED.md).
 
+## Visualizations
+Use [MeddlrViz](https://github.com/ad12/meddlr-viz) to visualize your medical imaging datasets, ML models, and more!
+
+```bash
+pip install meddlr-viz
+```
+
+<div align="center">
+
+<image src="https://github.com/ad12/meddlr-viz/blob/main/static/brats.gif" height=400 alt="A gallery of images from the BRATS dataset" />
+</div>
+
 ## 🐘 Model Zoo
 Easily serve and download pretrained models from the model zoo. A (evolving) list of pre-trained models can be found [here](MODEL_ZOO.md), on [HuggingFace 🤗](https://huggingface.co/arjundd), and in [project folders](projects).
 

projects/vortex/demo.ipynb

@@ -150,7 +150,13 @@
     "## Reconstruction Models\n",
     "Meddlr offers several reconstruction models in its Model Zoo and easy-to-use [APIs](https://github.com/ad12/meddlr#-model-zoo).\n",
     "\n",
-    "Let's start by using a few pre-trained models from the VORTEX paper. These models are hosted on huggingface in the Meddlr format.\n",
+    "Let's start by using a few pre-trained models from the VORTEX paper. Let's compare these models:\n",
+    "- `Supervised`: A supervised model trained on fully-sampled data\n",
+    "- `Supervised + Aug`: A supervised model trained on fully-sampled data with physics-based augmentations\n",
+    "- [`SSDU`](https://onlinelibrary.wiley.com/doi/abs/10.1002/mrm.28378): Self-supervised via data undersampling trained with both fully-sampled and undersampled scans\n",
+    "- [`VORTEX`](https://arxiv.org/abs/2111.02549): Trained with both fully-sampled and undersampled scans\n",
+    "\n",
+    "These models are hosted on huggingface in the Meddlr format.\n",
     "Providing the urls for these models will automatically download and load them in!"
    ]
   },
@@ -195,6 +201,7 @@
     "\n",
     "class ZeroFilledModel(nn.Module):\n",
     "    def forward(self, inputs: Dict[str, torch.Tensor]):\n",
+    "        \"\"\"\"\"\"\n",
     "        A = SenseModel(inputs[\"maps\"], weights=inputs[\"mask\"])\n",
     "        return A(inputs[\"kspace\"], adjoint=True)\n",
     "\n",