def save_nii_files(merged, image_pair, nii_res_dir, save_input_files):
# Extract data if nii files
try:
merged = merged.get_data()
except AttributeError:
merged = nib.Nifti1Image(merged, affine=image_pair.affine)
volumes = [merged, image_pair.image_obj, image_pair.labels_obj]
labels = ["%s_PRED.nii.gz" % image_pair.id,
"%s_IMAGE.nii.gz" % image_pair.id,
"%s_LABELS.nii.gz" % image_pair.id]
if not save_input_files:
volumes = volumes[:1]
labels = labels[:1]
p = os.path.abspath(nii_res_dir) # Save file directly in nii_res_dir
else:
# Create sub-folder under nii_res_dir
p = os.path.join(nii_res_dir, image_pair.id)
create_folders(p)
for nii, fname in zip(volumes, labels):
try:
nib.save(nii, "%s/%s" % (p, fname))
except AttributeError:
# No labels file?
pass
def entry_func(args=None):
args = get_argparser().parse_args(args)
out_dir = os.path.abspath(os.path.abspath(args.out_dir))
if os.path.exists(out_dir) and os.listdir(out_dir):
from sys import exit
print("Path '{}' already exists and is not empty!".format(
args.out_dir))
exit(1)
create_folders([out_dir], create_deep=False)
n_train, n_val, n_test = [args.N] * 3
n_train = args.N_train or n_train
n_val = args.N_val or n_val
n_test = args.N_test or n_test
np.random.seed(args.seed)
for dataset, N in zip(["train", "val", "test"], [n_train, n_val, n_test]):
if not N:
continue
else:
print("[*] Creating dataset '{}' of {} samples".format(dataset, N))
create_toy_dataset(N=N,
image_size=args.image_size,
out_dir=os.path.join(out_dir, dataset),
image_subdir=args.image_subdir,
label_subdir=args.label_subdir)
def create_toy_dataset(N, image_size, out_dir, image_subdir, label_subdir):
image_out_dir = os.path.join(out_dir, image_subdir)
label_out_dir = os.path.join(out_dir, label_subdir)
create_folders([out_dir, image_out_dir, label_out_dir])
for img_id in range(N):
print("--", img_id)
file_name = "toy_data_{}.nii.gz".format(img_id)
im_path = os.path.join(image_out_dir, file_name)
lab_path = os.path.join(label_out_dir, file_name)
image, label = create_toy_data_point(image_size)
affine = np.eye(4)
nib.save(nib.Nifti1Image(image, affine=affine), im_path)
nib.save(nib.Nifti1Image(label, affine=affine), lab_path)
def entry_func(args=None):
# Get command line arguments
args = get_argparser().parse_args(args)
assert_args(args)
# Get most important paths
project_dir = os.path.abspath(args.project_dir)
out_dir = os.path.abspath(args.out_dir)
# Check if valid dir structures
validate_folders(project_dir, out_dir,
overwrite=args.overwrite,
_continue=vars(args)["continue"])
nii_res_dir = os.path.join(out_dir, "nii_files")
create_folders(nii_res_dir, create_deep=True)
# Get settings from YAML file
hparams = load_hparams(project_dir)
# Get dataset
image_pair_loader, image_pair_dict = get_image_pair_loader(args, hparams,
out_dir)
# Wait for PID to terminate before continuing, if specified
if args.wait_for:
await_PIDs(args.wait_for, check_every=120)
# Set GPU device
set_gpu_vis(args)
# Get views
views = np.load("%s/views.npz" % project_dir)["arr_0"]
del hparams['fit']['views']
# Prepare result dicts
results, per_view_results = None, None
if not args.no_eval:
results, per_view_results = get_results_dicts(out_dir, views,
image_pair_dict,
hparams["build"]["n_classes"],
vars(args)["continue"])
# Get model and load weights, assign to one or more GPUs
model, weights_name = get_model(project_dir, hparams['build'])
fusion_model = None
if not args.sum_fusion:
fusion_model = get_fusion_model(n_views=len(views),
n_classes=hparams["build"]["n_classes"],
project_dir=project_dir,
weights_name=weights_name)
run_predictions_and_eval(
image_pair_loader=image_pair_loader,
image_pair_dict=image_pair_dict,
model=model,
fusion_model=fusion_model,
views=views,
hparams=hparams,
args=args,
results=results,
per_view_results=per_view_results,
out_dir=out_dir,
nii_res_dir=nii_res_dir
)
if not args.no_eval:
# Write final results
save_all(results, per_view_results, out_dir)