updated weights download for contrast phase algorithm

Author: malteekj

comp2comp/contrast_phase/contrast_inf.py

@@ -454,8 +454,9 @@ def predict_phase(TS_path, scan_path, outputPath=None, save_sample=False):
             text_file.write('{},{:.3f}\n'.format(phase_dict[i], y_pred_proba[i]))
 
     print('Predicted phase: ' + pred_phase)
+    print('\nProbabilities:')
     for i in range(len(y_pred_proba)):
-        print('{},{:.3f}'.format(phase_dict[i], y_pred_proba[i]))
+        print('{:<20}{:.3f}'.format(phase_dict[i], y_pred_proba[i]))
 
     output_path_images = os.path.join(outputPath, "images")
     if not os.path.exists(output_path_images):

comp2comp/contrast_phase/contrast_phase.py

@@ -2,9 +2,11 @@
 from pathlib import Path
 from time import time
 from typing import Union
+import subprocess
+import zipfile
 
 from totalsegmentator.libs import (
-    download_pretrained_weights,
+    # download_pretrained_weights,
     nostdout,
     setup_nnunet,
 )
@@ -67,9 +69,12 @@ def run_segmentation(
         task_id = [251]
 
         setup_nnunet()
-        for task_id in [251]:
-            download_pretrained_weights(task_id)
+        # for task_id in [251]:
+        #     download_pretrained_weights(task_id)
 
+        # download with weight for id 251
+        self.download_pretrained_weights_updated(task_id[0])
+        
         from totalsegmentator.nnunet import nnUNet_predict_image
 
         with nostdout():
@@ -130,7 +135,40 @@ def run_segmentation(
 
         # return seg, img
         return seg, img
-
+    
+    def download_pretrained_weights_updated(self, task_id):
+        '''
+        Download the weights with curl to resolve problems
+        with downloading from Zenodo
+        '''
+        home_path = Path(os.environ["SCRATCH"])
+        config_dir = home_path / ".totalsegmentator/nnunet/results/nnUNet"
+        (config_dir / "3d_fullres").mkdir(exist_ok=True, parents=True)
+        (config_dir / "2d").mkdir(exist_ok=True, parents=True)
+        
+        url = "https://zenodo.org/records/6802342/files/Task251_TotalSegmentator_part1_organs_1139subj.zip?download=1"
+        config_dir = config_dir / "3d_fullres"
+        weights_path = config_dir / "Task251_TotalSegmentator_part1_organs_1139subj"
+        tempfile = config_dir / "tmp_download_file.zip"
+        
+        if not weights_path.exists():
+            print('Downloading weights..')
+            subprocess.run(
+                ["curl", "-L", url, "-o", tempfile],
+                check=True
+            )
+            
+            print('Unzipping..')
+            with zipfile.ZipFile(config_dir / "tmp_download_file.zip", 'r') as zip_f:
+                zip_f.extractall(config_dir)
+            # print(f"  downloaded in {time.time()-st:.2f}s")
+            if tempfile.exists():
+                os.remove(tempfile)
+            print('Done.')
+        else:
+            print('Weights are already downloaded')
+            
+                
     def convertNibToNumpy(self, TSNib, ImageNib):
         """Convert nifti to numpy array.
 

setup.py

@@ -65,6 +65,7 @@ def get_version():
         "dicom2nifti<2.6",
         "torch==2.5.1",
         "torchvision==0.20.1",
+        "xgboost",
     ],
     extras_require={
         "all": ["shapely", "psutil"],