diff --git a/Dockerfile b/Dockerfile index 21f8e64..629e34f 100755 --- a/Dockerfile +++ b/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 diff --git a/mist/eval_preds.py b/mist/eval_preds.py index d7834da..4cedbb6 100755 --- a/mist/eval_preds.py +++ b/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(): diff --git a/mist/evaluate_preds/evaluate.py b/mist/evaluate_preds/evaluate.py index b1a9342..df11cec 100755 --- a/mist/evaluate_preds/evaluate.py +++ b/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) diff --git a/mist/main.py b/mist/main.py index bb05a47..6ab9fbe 100755 --- a/mist/main.py +++ b/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): diff --git a/mist/post_preds.py b/mist/post_preds.py index 38c5cbc..85dd06a 100755 --- a/mist/post_preds.py +++ b/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 diff --git a/mist/postprocess_preds/postprocess.py b/mist/postprocess_preds/postprocess.py index fbb6e72..b8e9cf9 100755 --- a/mist/postprocess_preds/postprocess.py +++ b/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) diff --git a/mist/runtime/args.py b/mist/runtime/args.py index f9bda84..1be9379 100755 --- a/mist/runtime/args.py +++ b/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") diff --git a/pyproject.toml b/pyproject.toml index a11d747..2e431e0 100755 --- a/pyproject.toml +++ b/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",