Added inference notebooks

microsoft · Dec 4, 2024 · 3634dc4 · 3634dc4
1 parent 6fa21ec
commit 3634dc4
Show file tree

Hide file tree

Showing 6 changed files with 955 additions and 54 deletions.
diff --git a/README.md b/README.md
@@ -9,6 +9,7 @@ This repository hosts the code and resources for the paper **"A Foundation Model
 ![Example Predictions](assets/readmes/biomedparse_prediction_examples.png)
 
 ## News
+- Dec. 3, 2024: We added inference notebook examples in inference_example_RGB.ipynb and inference_example_NIFTI.ipynb
 - Nov. 22, 2024: We added negative prediction p-value example in inference_example_DICOM.ipynb
 - Nov. 18, 2024: BiomedParse is officially online in [*Nature Methods*](https://aka.ms/biomedparse-paper)!
 

diff --git a/examples/amos_0328.nii.gz b/examples/amos_0328.nii.gz
diff --git a/inference_examples_DICOM.ipynb b/inference_examples_DICOM.ipynb
diff --git a/inference_examples_NIFTI.ipynb b/inference_examples_NIFTI.ipynb
diff --git a/inference_examples_RGB.ipynb b/inference_examples_RGB.ipynb
diff --git a/modeling/BaseModel.py b/modeling/BaseModel.py
@@ -27,15 +27,15 @@ def forward(self, *inputs, **kwargs):
     def save_pretrained(self, save_dir):
         torch.save(self.model.state_dict(), os.path.join(save_dir, "model_state_dict.pt"))
 
-    def from_pretrained(self, pretrained, 
+    def from_pretrained(self, pretrained, filename: str = "biomedparse_v1.pt",
                         local_dir: str = "./pretrained", config_dir: str = "./configs"):
         if pretrained.startswith("hf_hub:"):
             hub_name = pretrained.split(":")[1]
-            huggingface_hub.hf_hub_download(hub_name, filename="biomedparse_v1.pt", 
+            huggingface_hub.hf_hub_download(hub_name, filename=filename, 
                                             local_dir=local_dir)
             huggingface_hub.hf_hub_download(hub_name, filename="config.yaml", 
                                             local_dir=config_dir)
-            load_dir = os.path.join(local_dir, "biomedparse_v1.pt")
+            load_dir = os.path.join(local_dir, filename)
         else:
             load_dir = pretrained