Upgrade pytorch version

Dockerfile

@@ -1,4 +1,4 @@
-ARG FROM_IMAGE_NAME=nvcr.io/nvidia/pytorch:23.06-py3
+ARG FROM_IMAGE_NAME=nvcr.io/nvidia/pytorch:24.05-py3
 FROM ${FROM_IMAGE_NAME}
 
 # Set environment variables 

mist/eval_preds.py

@@ -21,6 +21,7 @@ def get_eval_args():
     p.boolean_flag("--use-native-spacing",
                    default=False,
                    help="Use native image spacing to compute Hausdorff distances")
+    p.arg("--surf-dice-tol", type=float, default=1.0, help="Tolerance for surface dice")
 
     args = p.parse_args()
     return args
@@ -37,7 +38,8 @@ def main(args):
              args.output_csv,
              args.metrics,
              args.normalize_hd,
-             args.use_native_spacing)
+             args.use_native_spacing,
+             args.surf_dice_tol)
 
 
 def mist_eval_entry():

mist/evaluate_preds/evaluate.py

@@ -28,7 +28,7 @@ def get_worst_case_haus(mask_npy, spacing):
     return np.sqrt(width ** 2 + height ** 2 + depth ** 2)
 
 
-def evaluate_single_example(pred, truth, patient_id, config, metrics, use_native_spacing):
+def evaluate_single_example(pred, truth, patient_id, config, metrics, use_native_spacing, surf_dice_tol):
     # Get dice and hausdorff distances for final prediction
     row_dict = dict()
     row_dict['id'] = patient_id
@@ -88,7 +88,7 @@ def evaluate_single_example(pred, truth, patient_id, config, metrics, use_native
             elif bool(temp_truth_sum == 0) ^ bool(temp_pred_sum == 0):
                 surf_dice_coef = 0.0
             else:
-                surf_dice_coef = compute_surface_dice_at_tolerance(distances, tolerance_mm=1.0)
+                surf_dice_coef = compute_surface_dice_at_tolerance(distances, tolerance_mm=surf_dice_tol)
             row_dict["{}_surf_dice".format(key)] = surf_dice_coef
         if "avg_surf" in metrics:
             if temp_truth_sum == 0 and temp_pred_sum == 0:
@@ -101,7 +101,7 @@ def evaluate_single_example(pred, truth, patient_id, config, metrics, use_native
     return row_dict
 
 
-def evaluate(config_json, paths, source_dir, output_csv, metrics, use_native_spacing):
+def evaluate(config_json, paths, source_dir, output_csv, metrics, use_native_spacing, surf_dice_tol):
     with open(config_json, 'r') as file:
         config = json.load(file)
 
@@ -139,7 +139,8 @@ def evaluate(config_json, paths, source_dir, output_csv, metrics, use_native_spa
                                                    patient_id,
                                                    config,
                                                    metrics,
-                                                   use_native_spacing)
+                                                   use_native_spacing,
+                                                   surf_dice_tol)
             results_df = pd.concat([results_df, pd.DataFrame(eval_results, index=[0])], ignore_index=True)
 
     results_df = compute_results_stats(results_df)

mist/main.py

@@ -64,7 +64,8 @@ def main(args):
                  os.path.join(args.results, "predictions", "train", "raw"),
                  os.path.join(args.results, "results.csv"),
                  args.metrics,
-                 args.use_native_spacing)
+                 args.use_native_spacing,
+                 args.surf_dice_tol)
 
     if args.exec_mode == "all" or args.exec_mode == "train":
         if has_test_data(args.data):

mist/post_preds.py

@@ -41,6 +41,7 @@ def get_postprocess_args():
     p.boolean_flag("--use-native-spacing",
                    default=False,
                    help="Use native image spacing to compute Hausdorff distances")
+    p.arg("--surf-dice-tol", type=float, default=1.0, help="Tolerance for surface dice")
 
     args = p.parse_args()
     return args

mist/postprocess_preds/postprocess.py

@@ -121,7 +121,8 @@ def check_transforms(self, transforms, messages, transform_kwargs):
                  self.dest_dir,
                  self.new_results_csv,
                  self.metrics,
-                 self.args.use_native_spacing)
+                 self.args.use_native_spacing,
+                 self.args.surf_dice_tol)
 
         # Compute improvement score
         new_results_df = pd.read_csv(self.new_results_csv)

mist/runtime/args.py

@@ -176,6 +176,7 @@ def get_main_args():
     p.boolean_flag("--use-native-spacing",
                    default=False,
                    help="Use native image spacing to compute Hausdorff distances")
+    p.arg("--surf-dice-tol", type=float, default=1.0, help="Tolerance for surface dice")
 
     # Uncertainty
     p.boolean_flag("--output-std", default=False, help="Output standard deviation for ensemble predictions")

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mist-medical"
-version = "0.4.7-alpha"
+version = "0.4.8-alpha"
 requires-python = ">= 3.8"
 description = "MIST is a simple, fully automated framework for 3D medical imaging segmentation."
 readme = "README.md"
@@ -37,10 +37,10 @@ keywords = [
     "unetr"
 ]
 dependencies = [
-    "torch==2.0.1",
-    "monai==1.3.0",
-    "antspyx==0.3.8",
-    "simpleitk==2.2.1",
+    "torch>=2.0.1",
+    "monai>=1.3.0",
+    "antspyx>=0.3.8",
+    "simpleitk>=2.2.1",
     "numpy",
     "pandas",
     "rich",