diff --git a/NV-Segment-CT/docs/README.md b/NV-Segment-CT/docs/README.md index 91e0f4f..f1d1290 100644 --- a/NV-Segment-CT/docs/README.md +++ b/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) diff --git a/NV-Segment-CT/requirements.txt b/NV-Segment-CT/requirements.txt index 7c31cb3..04d60a5 100644 --- a/NV-Segment-CT/requirements.txt +++ b/NV-Segment-CT/requirements.txt @@ -16,3 +16,4 @@ timm pytorch-ignite tensorboardX mlflow +huggingface_hub \ No newline at end of file diff --git a/NV-Segment-CTMR/brain_t1_preprocess/run_brain_segmentation.sh b/NV-Segment-CTMR/brain_t1_preprocess/run_brain_segmentation.sh index 22f60bd..c9b5e09 100755 --- a/NV-Segment-CTMR/brain_t1_preprocess/run_brain_segmentation.sh +++ b/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" diff --git a/NV-Segment-CTMR/docs/README.md b/NV-Segment-CTMR/docs/README.md index 9f03994..3e89a6b 100644 --- a/NV-Segment-CTMR/docs/README.md +++ b/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) diff --git a/NV-Segment-CTMR/docs/ctmr.png b/NV-Segment-CTMR/docs/ctmr.png new file mode 100644 index 0000000..c4b2885 Binary files /dev/null and b/NV-Segment-CTMR/docs/ctmr.png differ diff --git a/NV-Segment-CTMR/requirements.txt b/NV-Segment-CTMR/requirements.txt index 7c31cb3..906ac85 100644 --- a/NV-Segment-CTMR/requirements.txt +++ b/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 \ No newline at end of file diff --git a/README.md b/README.md index d493697..b615fc7 100644 --- a/README.md +++ b/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)