Skip to content

Commit

Permalink
update nnunet to 2.2.1
Browse files Browse the repository at this point in the history
  • Loading branch information
@wasserth
wasserth committed Feb 6, 2024
1 parent 350bb89 commit bf452f5
Show file tree
Hide file tree
Showing 3 changed files with 54 additions and 48 deletions.
2 changes: 1 addition & 1 deletion setup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
'p_tqdm',
'xvfbwrapper',
'fury',
'nnunetv2==2.1',
'nnunetv2>=2.2.1',
'requests==2.27.1;python_version<"3.10"',
'requests;python_version>="3.10"',
'rt_utils',
Expand Down
18 changes: 12 additions & 6 deletions tests/test_locally.py
View file
Original file line number Diff line number Diff line change
@@ -1,22 +1,26 @@
import sys
from pathlib import Path
p_dir = str(Path(__file__).absolute().parents[1])
if p_dir not in sys.path: sys.path.insert(0, p_dir)

import pytest
import os
import re
import glob
import shutil
import subprocess
from tqdm import tqdm
import pandas as pd
import numpy as np
from collections import defaultdict
import psutil
import time
import threading
import torch
import platform

import psutil
from tqdm import tqdm
import pandas as pd
import numpy as np
import torch
import nnunetv2

from totalsegmentator.python_api import totalsegmentator
from totalsegmentator.map_to_binary import class_map
from totalsegmentator.excel import set_xlsx_column_width_to_content
Expand Down Expand Up @@ -222,7 +226,8 @@ def are_logs_similar(last_log, new_log, cols, tolerance_percent=0.04):
"memory_gpu_15mm", "memory_gpu_3mm",
"cpu_utilization_15mm", "cpu_utilization_3mm",
"gpu_utilization_15mm", "gpu_utilization_3mm",
"python_version", "torch_version", "cuda_version", "cudnn_version",
"python_version", "torch_version", "nnunet_version",
"cuda_version", "cudnn_version",
"gpu_name"]
overview_file = Path(f"{base_dir}/overview.xlsx")
if overview_file.exists():
Expand All @@ -240,6 +245,7 @@ def are_logs_similar(last_log, new_log, cols, tolerance_percent=0.04):
cpu_utilization["15mm"], cpu_utilization["3mm"],
gpu_utilization["15mm"], gpu_utilization["3mm"],
platform.python_version(), torch.__version__,
nnunetv2.__version__,
float(torch.version.cuda), int(torch.backends.cudnn.version()),
torch.cuda.get_device_name(0)]

Expand Down
82 changes: 41 additions & 41 deletions totalsegmentator/nnunet.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,10 +22,10 @@
from totalsegmentator.libs import nostdout

# nnUNet 2.1
with nostdout():
from nnunetv2.inference.predict_from_raw_data import predict_from_raw_data
# with nostdout():
# from nnunetv2.inference.predict_from_raw_data import predict_from_raw_data
# nnUNet 2.2
# from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor
from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor

from nnunetv2.utilities.file_path_utilities import get_output_folder

Expand Down Expand Up @@ -178,46 +178,46 @@ def nnUNetv2_predict(dir_in, dir_out, task_id, model="3d_fullres", folds=None,
allow_tqdm = not quiet

# nnUNet 2.1
predict_from_raw_data(dir_in,
dir_out,
model_folder,
folds,
step_size,
use_gaussian=True,
use_mirroring=not disable_tta,
perform_everything_on_gpu=True,
verbose=verbose,
save_probabilities=save_probabilities,
overwrite=not continue_prediction,
checkpoint_name=chk,
num_processes_preprocessing=npp,
num_processes_segmentation_export=nps,
folder_with_segs_from_prev_stage=prev_stage_predictions,
num_parts=num_parts,
part_id=part_id,
device=device)
# predict_from_raw_data(dir_in,
# dir_out,
# model_folder,
# folds,
# step_size,
# use_gaussian=True,
# use_mirroring=not disable_tta,
# perform_everything_on_gpu=True,
# verbose=verbose,
# save_probabilities=save_probabilities,
# overwrite=not continue_prediction,
# checkpoint_name=chk,
# num_processes_preprocessing=npp,
# num_processes_segmentation_export=nps,
# folder_with_segs_from_prev_stage=prev_stage_predictions,
# num_parts=num_parts,
# part_id=part_id,
# device=device)

# nnUNet 2.2
# predictor = nnUNetPredictor(
# tile_step_size=step_size,
# use_gaussian=True,
# use_mirroring=not disable_tta,
# perform_everything_on_gpu=True,
# device=device,
# verbose=verbose,
# verbose_preprocessing=verbose,
# allow_tqdm=allow_tqdm
# )
# predictor.initialize_from_trained_model_folder(
# model_folder,
# use_folds=folds,
# checkpoint_name=chk,
# )
# predictor.predict_from_files(dir_in, dir_out,
# save_probabilities=save_probabilities, overwrite=not continue_prediction,
# num_processes_preprocessing=npp, num_processes_segmentation_export=nps,
# folder_with_segs_from_prev_stage=prev_stage_predictions,
# num_parts=num_parts, part_id=part_id)
predictor = nnUNetPredictor(
tile_step_size=step_size,
use_gaussian=True,
use_mirroring=not disable_tta,
perform_everything_on_gpu=True,
device=device,
verbose=verbose,
verbose_preprocessing=verbose,
allow_tqdm=allow_tqdm
)
predictor.initialize_from_trained_model_folder(
model_folder,
use_folds=folds,
checkpoint_name=chk,
)
predictor.predict_from_files(dir_in, dir_out,
save_probabilities=save_probabilities, overwrite=not continue_prediction,
num_processes_preprocessing=npp, num_processes_segmentation_export=nps,
folder_with_segs_from_prev_stage=prev_stage_predictions,
num_parts=num_parts, part_id=part_id)


def save_segmentation_nifti(class_map_item, tmp_dir=None, file_out=None, nora_tag=None, header=None, task_name=None, quiet=None):
Expand Down

0 comments on commit bf452f5

Please sign in to comment.