Update download link

NV-Segment-CT/docs/README.md

@@ -17,7 +17,9 @@ pip install -r requirements.txt;
 cd ..;
 mkdir NV-Segment-CT/models;
 # download from huggingface link
-wget -O NV-Segment-CT/models/model.pt https://huggingface.co/nvidia/NV-Segment-CT/resolve/main/vista3d_pretrained_model/model.pt
+hf download nvidia/NV-Segment-CT vista3d_pretrained_model/model.pt --local-dir NV-Segment-CT/models/ && \
+mv NV-Segment-CT/models/vista3d_pretrained_model/model.pt NV-Segment-CT/models/model.pt && \
+rmdir NV-Segment-CT/models/vista3d_pretrained_model
 ```
 
 ## 1.1 **VISTA3D-CT** [[Github]](https://github.com/NVIDIA-Medtech/NV-Segment-CTMR/tree/main/NV-Segment-CT) [[Huggingface]](https://huggingface.co/nvidia/NV-Segment-CT)

NV-Segment-CT/requirements.txt

@@ -16,3 +16,4 @@ timm
 pytorch-ignite
 tensorboardX
 mlflow
+huggingface_hub

NV-Segment-CTMR/brain_t1_preprocess/run_brain_segmentation.sh

@@ -134,9 +134,11 @@ process_single_file() {
     # Step 3: Segment the brain
     echo -e "${YELLOW}Step 3/4: Running segmentation...${NC}"
     cd "$BUNDLE_ROOT"
+    # Override output_dir in config to use our output_dir so segmentation saves to the right place
     python -m monai.bundle run \
         --config_file configs/inference.json \
         --input_dict "{'image':'$preprocess_tmp'}" \
+        --output_dir "$output_dir" \
         --modality "$MODALITY" || {
         echo -e "${RED}Error: Segmentation failed${NC}" >&2
         [[ "$KEEP_TEMP" == "false" ]] && rm -rf "$temp_dir"

NV-Segment-CTMR/docs/README.md

@@ -25,7 +25,10 @@ pip install -r requirements.txt
 # Create models directory and download pretrained model
 cd ..
 mkdir -p NV-Segment-CTMR/models
-wget -O NV-Segment-CTMR/models/model.pt https://huggingface.co/nvidia/NV-Segment-CTMR/resolve/main/vista3d_pretrained_model/model.pt
+# Option 1: Download using hf and move to expected location
+hf download nvidia/NV-Segment-CTMR vista3d_pretrained_model/model.pt --local-dir NV-Segment-CTMR/models/ && \
+mv NV-Segment-CTMR/models/vista3d_pretrained_model/model.pt NV-Segment-CTMR/models/model.pt && \
+rmdir NV-Segment-CTMR/models/vista3d_pretrained_model
 ```
 
 ## Automatic Segmentation (support multi-gpu batch processing)

NV-Segment-CTMR/requirements.txt

@@ -7,7 +7,8 @@ Pillow==10.4.0
 PyYAML==6.0.2
 scipy
 scikit-image==0.24.0
-torch==2.0.1
+torch==2.1.2
+torchvision==0.16.2
 tqdm==4.66.2
 tensorboard==2.13.0
 einops==0.6.1
@@ -16,3 +17,9 @@ timm
 pytorch-ignite
 tensorboardX
 mlflow
+pymedio
+SimpleITK
+pydicom
+statsmodels
+antspyx
+huggingface_hub

README.md

@@ -8,7 +8,7 @@ Both models follow the MONAI bundle architecture.
 
 | Feature | NV-Segment-CT | NV-Segment-CTMR |
 |---------|---------------|-----------------|
-| **Anatomical Classes** | 132 classes | 345+ classes |
+| **Anatomical Classes** | [132 classes (7 types of tumors)](NV-Segment-CT/configs/label_dict.json) | [345+ classes](NV-Segment-CTMR/configs/label_dict.json) |
 | **Modalities** | CT only | CT + MRI (body & brain) |
 | **Segmentation Type** | Automatic + Interactive (point-click) | Automatic only |
 | **Model Weights**     | [NV-Segment-CT on HuggingFace](https://huggingface.co/nvidia/NV-Segment-CT) | [NV-Segment-CTMR on HuggingFace](https://huggingface.co/nvidia/NV-Segment-CTMR) |
@@ -25,6 +25,8 @@ Both models follow the MONAI bundle architecture.
 - out-of-the-box automatic segmentation on 3D CT scans
 - share the same architecture with VISTA3D-CT model but we only trained the automatic segmentation branch with larger CT and MRI datasets.
 
+![CTMR](./NV-Segment-CTMR/docs/ctmr.png) 
+
 ## Performance on held-out test set
 
 ![Benchmark CT](./NV-Segment-CTMR/docs/benchmarkct.png) ![Benchmark MR](./NV-Segment-CTMR/docs/benchmarkmr.png)