def testing():
c = get_config()
c.do_load_checkpoint = True
#c.checkpoint_dir = c.base_dir + '/20190424-020641_unet_experiment' + '/checkpoint/checkpoint_current' # dice_cost train
# c.checkpoint_dir = c.base_dir + '/20190424-234657_unet_experiment' + '/checkpoint/checkpoint_last' # SDG
c.checkpoint_dir = c.base_dir + '/20190906-085449_unet_experiment' + '/checkpoint/checkpoint_current'
# c.checkpoint_file = "checkpoint_last.pth.tar"
# c.cross_vali_index = valiIndex
cross_vali_result_all_dir = os.path.join(c.base_dir, c.dataset_name
+ '_' + str(
c.batch_size) + c.cross_vali_result_all_dir + datetime.datetime.now().strftime("_%Y%m%d-%H%M%S"))
if not os.path.exists(cross_vali_result_all_dir):
os.makedirs(cross_vali_result_all_dir)
print('Created' + cross_vali_result_all_dir + '...')
c.base_dir = cross_vali_result_all_dir
c.cross_vali_result_all_dir = os.path.join(cross_vali_result_all_dir, "results")
os.makedirs(c.cross_vali_result_all_dir)
exp = UNetExperiment(config=c, name='unet_test', n_epochs=c.n_epochs,
seed=42, globs=globals())
exp.run_test(setup=True)
def training():
c = get_config()
dataset_name = 'Task04_Hippocampus'
# dataset_name = 'Task02_Heart'
# download_dataset(dest_path=c.data_root_dir, dataset=dataset_name, id=c.google_drive_id)
# c.do_load_checkpoint = True
# c.checkpoint_dir = c.base_dir + '/20190801-_unet_experiment' + '/checkpoint/checkpoint_current'
# c.checkpoint_file = "checkpoint_last.pth.tar"
if not exists(os.path.join(os.path.join(c.data_root_dir, dataset_name), 'preprocessed')):
print('Preprocessing data. [STARTED]')
preprocess_data(root_dir=os.path.join(c.data_root_dir, dataset_name))
create_splits(output_dir=c.split_dir, image_dir=c.data_dir)
print('Preprocessing data. [DONE]')
else:
print('The data has already been preprocessed. It will not be preprocessed again. Delete the folder to enforce it.')
exp = UNetExperiment(config=c, name='unet_experiment', n_epochs=c.n_epochs,
seed=42, append_rnd_to_name=c.append_rnd_string) # visdomlogger_kwargs={"auto_start": c.start_visdom}
exp.run()
exp.run_test(setup=False)
def preprocess_data(root_dir): #y_shape=64, z_shape=64):
c = get_config()
image_dir = os.path.join(root_dir, 'imagesTr')
label_dir = os.path.join(root_dir, 'labelsTr')
output_dir = os.path.join(root_dir, 'preprocessed')
classes = c.num_classes
if not os.path.exists(output_dir):
os.makedirs(output_dir)
print('Created' + output_dir + '...')
class_stats = defaultdict(int)
total = 0
nii_files = subfiles(image_dir, suffix=".nii.gz", join=False)
for i in range(0, len(nii_files)):
if nii_files[i].startswith("._"):
nii_files[i] = nii_files[i][2:]
for f in nii_files:
image, _ = load(os.path.join(image_dir, f))
label, _ = load(os.path.join(label_dir, f.replace('_0000', '')))
print(f)
for i in range(classes):
class_stats[i] += np.sum(label == i)
total += np.sum(label == i)
# normalize images
image = (image - image.min()) / (image.max() - image.min())
image = np.swapaxes(image, 0, 2)
label = np.swapaxes(label, 0, 2)
result = reshape(np.stack([image, label], axis=0),
crop_size=c.patch_size)
# image = reshape(image, append_value=0, new_shape=(image.shape[0], y_shape, z_shape))
# label = reshape(label, append_value=0, new_shape=(label.shape[0], y_shape, z_shape))
np.save(os.path.join(output_dir, f.split('.')[0] + '.npy'), result)
print(f)
print(total)
for i in range(classes):
print(class_stats[i], class_stats[i] / total)
class LabelTensorToColor(object):
def __call__(self, label):
label = label.squeeze()
colored_label = torch.zeros(3, label.size(0), label.size(1)).byte()
for i, color in enumerate(class_color):
mask = label.eq(i)
for j in range(3):
colored_label[j].masked_fill_(mask, color[j])
return colored_label
color_class_converter = LabelTensorToColor()
# Load data
c = get_config()
exp = UNetExperiment(config=c, name=c.name, n_epochs=c.n_epochs,
seed=42, append_rnd_to_name=c.append_rnd_string, globs=globals())
exp.setup()
exp.test_data_loader.do_reshuffle = True
# Load checkpoint
checkpoint = torch.load('./output_experiment/20190604-154315_fine_tune_spleen_for_heart/checkpoint/checkpoint_last.pth.tar')
exp.model.load_state_dict(checkpoint['model'])
exp.model.eval()
batch_counter = 0
with torch.no_grad():
for data_batch in exp.test_data_loader:
# Get data_batches