def test_gru_mixture_fprop():
hidden_sizes = 50
with Timer("Creating dataset", newline=True):
volume_manager = neurotools.VolumeManager()
trainset = make_dummy_dataset(volume_manager)
print("Dataset sizes:", len(trainset))
batch_scheduler = batch_schedulers.TractographyBatchScheduler(
trainset, batch_size=16, noisy_streamlines_sigma=None, seed=1234)
print("An epoch will be composed of {} updates.".format(
batch_scheduler.nb_updates_per_epoch))
print(volume_manager.data_dimension, hidden_sizes,
batch_scheduler.target_size)
with Timer("Creating model"):
hyperparams = {
'model': 'gru_mixture',
'n_gaussians': 2,
'SGD': "1e-2",
'hidden_sizes': hidden_sizes,
'learn_to_stop': False,
'normalize': False,
'feed_previous_direction': False
}
model = factories.model_factory(
hyperparams,
input_size=volume_manager.data_dimension,
output_size=batch_scheduler.target_size,
volume_manager=volume_manager)
model.initialize(
factories.weigths_initializer_factory("orthogonal", seed=1234))
# Test fprop with missing streamlines from one subject in a batch
output = model.get_output(trainset.symb_inputs)
fct = theano.function([trainset.symb_inputs],
output,
updates=model.graph_updates)
batch_inputs, batch_targets, batch_mask = batch_scheduler._next_batch(2)
out = fct(batch_inputs)
with Timer("Building optimizer"):
loss = factories.loss_factory(hyperparams, model, trainset)
optimizer = factories.optimizer_factory(hyperparams, loss)
fct_loss = theano.function(
[trainset.symb_inputs, trainset.symb_targets, trainset.symb_mask],
loss.loss,
updates=model.graph_updates)
loss_value = fct_loss(batch_inputs, batch_targets, batch_mask)
print("Loss:", loss_value)
fct_optim = theano.function(
[trainset.symb_inputs, trainset.symb_targets, trainset.symb_mask],
list(optimizer.directions.values()),
updates=model.graph_updates)
dirs = fct_optim(batch_inputs, batch_targets, batch_mask)
def test_gru_multistep_fprop_k3():
hidden_sizes = 50
hyperparams = {
'model': 'gru_multistep',
'k': 3,
'm': 3,
'batch_size': 16,
'SGD': "1e-2",
'hidden_sizes': hidden_sizes,
'learn_to_stop': False,
'normalize': False,
'noisy_streamlines_sigma': None,
'shuffle_streamlines': True,
'seed': 1234
}
with Timer("Creating dataset", newline=True):
volume_manager = neurotools.VolumeManager()
trainset = make_dummy_dataset(volume_manager)
print("Dataset sizes:", len(trainset))
batch_scheduler = factories.batch_scheduler_factory(hyperparams,
trainset,
train_mode=True)
print("An epoch will be composed of {} updates.".format(
batch_scheduler.nb_updates_per_epoch))
print(volume_manager.data_dimension, hidden_sizes,
batch_scheduler.target_size)
with Timer("Creating model"):
model = factories.model_factory(
hyperparams,
input_size=volume_manager.data_dimension,
output_size=batch_scheduler.target_size,
volume_manager=volume_manager)
model.initialize(
factories.weigths_initializer_factory("orthogonal", seed=1234))
# Test fprop
output = model.get_output(trainset.symb_inputs)
fct = theano.function([trainset.symb_inputs],
output,
updates=model.graph_updates)
batch_inputs, batch_targets, batch_mask = batch_scheduler._next_batch(2)
out = fct(batch_inputs)
with Timer("Building optimizer"):
loss = factories.loss_factory(hyperparams, model, trainset)
optimizer = factories.optimizer_factory(hyperparams, loss)
fct_loss = theano.function(
[trainset.symb_inputs, trainset.symb_targets, trainset.symb_mask],
loss.loss,
updates=model.graph_updates)
loss_value = fct_loss(batch_inputs, batch_targets, batch_mask)
print("Loss:", loss_value)
fct_optim = theano.function(
[trainset.symb_inputs, trainset.symb_targets, trainset.symb_mask],
list(optimizer.directions.values()),
updates=model.graph_updates)
dirs = fct_optim(batch_inputs, batch_targets, batch_mask)
def main():
parser = build_argparser()
args = parser.parse_args()
# Get experiment folder
experiment_path = args.name
if not os.path.isdir(experiment_path):
# If not a directory, it must be the name of the experiment.
experiment_path = pjoin(".", "experiments", args.name)
if not os.path.isdir(experiment_path):
parser.error('Cannot find experiment: {0}!'.format(args.name))
# Load experiments hyperparameters
try:
hyperparams = smartutils.load_dict_from_json_file(
pjoin(experiment_path, "hyperparams.json"))
except FileNotFoundError:
hyperparams = smartutils.load_dict_from_json_file(
pjoin(experiment_path, "..", "hyperparams.json"))
with Timer("Loading DWIs"):
# Load gradients table
dwi_name = args.dwi
if dwi_name.endswith(".gz"):
dwi_name = dwi_name[:-3]
if dwi_name.endswith(".nii"):
dwi_name = dwi_name[:-4]
try:
bvals_filename = dwi_name + ".bvals"
bvecs_filename = dwi_name + ".bvecs"
bvals, bvecs = dipy.io.gradients.read_bvals_bvecs(
bvals_filename, bvecs_filename)
except FileNotFoundError:
try:
bvals_filename = dwi_name + ".bval"
bvecs_filename = dwi_name + ".bvec"
bvals, bvecs = dipy.io.gradients.read_bvals_bvecs(
bvals_filename, bvecs_filename)
except FileNotFoundError as e:
print("Could not find .bvals/.bvecs or .bval/.bvec files...")
raise e
dwi = nib.load(args.dwi)
if hyperparams["use_sh_coeffs"]:
# Use 45 spherical harmonic coefficients to represent the diffusion signal.
weights = neurotools.get_spherical_harmonics_coefficients(
dwi, bvals, bvecs).astype(np.float32)
else:
# Resample the diffusion signal to have 100 directions.
weights = neurotools.resample_dwi(dwi, bvals,
bvecs).astype(np.float32)
affine_rasmm2dwivox = np.linalg.inv(dwi.affine)
with Timer("Loading model"):
if hyperparams["model"] == "gru_regression":
from learn2track.models import GRU_Regression
model_class = GRU_Regression
elif hyperparams['model'] == 'gru_gaussian':
from learn2track.models import GRU_Gaussian
model_class = GRU_Gaussian
elif hyperparams['model'] == 'gru_mixture':
from learn2track.models import GRU_Mixture
model_class = GRU_Mixture
elif hyperparams['model'] == 'gru_multistep':
from learn2track.models import GRU_Multistep_Gaussian
model_class = GRU_Multistep_Gaussian
elif hyperparams['model'] == 'ffnn_regression':
from learn2track.models import FFNN_Regression
model_class = FFNN_Regression
else:
raise ValueError("Unknown model!")
kwargs = {}
volume_manager = neurotools.VolumeManager()
volume_manager.register(weights)
kwargs['volume_manager'] = volume_manager
# Load the actual model.
model = model_class.create(
pjoin(experiment_path),
**kwargs) # Create new instance and restore model.
model.drop_prob = 0.
print(str(model))
mask = None
if args.mask is not None:
with Timer("Loading mask"):
mask_nii = nib.load(args.mask)
mask = mask_nii.get_data()
# Compute the affine allowing to evaluate the mask at some coordinates correctly.
# affine_maskvox2dwivox = mask_vox => rasmm space => dwi_vox
affine_maskvox2dwivox = np.dot(affine_rasmm2dwivox,
mask_nii.affine)
if args.dilate_mask:
import scipy
mask = scipy.ndimage.morphology.binary_dilation(mask).astype(
mask.dtype)
with Timer("Generating seeds"):
seeds = []
for filename in args.seeds:
if filename.endswith('.trk') or filename.endswith('.tck'):
tfile = nib.streamlines.load(filename)
# Send the streamlines to voxel since that's where we'll track.
tfile.tractogram.apply_affine(affine_rasmm2dwivox)
# Use extremities of the streamlines as seeding points.
seeds += [s[0] for s in tfile.streamlines]
seeds += [s[-1] for s in tfile.streamlines]
else:
# Assume it is a binary mask.
rng = np.random.RandomState(args.seeding_rng_seed)
nii_seeds = nib.load(filename)
# affine_seedsvox2dwivox = mask_vox => rasmm space => dwi_vox
affine_seedsvox2dwivox = np.dot(affine_rasmm2dwivox,
nii_seeds.affine)
nii_seeds_data = nii_seeds.get_data()
if args.dilate_seeding_mask:
import scipy
nii_seeds_data = scipy.ndimage.morphology.binary_dilation(
nii_seeds_data).astype(nii_seeds_data.dtype)
indices = np.array(np.where(nii_seeds_data)).T
for idx in indices:
seeds_in_voxel = idx + rng.uniform(
-0.5, 0.5, size=(args.nb_seeds_per_voxel, 3))
seeds_in_voxel = nib.affines.apply_affine(
affine_seedsvox2dwivox, seeds_in_voxel)
seeds.extend(seeds_in_voxel)
seeds = np.array(seeds, dtype=theano.config.floatX)
with Timer("Tracking in the diffusion voxel space"):
voxel_sizes = np.asarray(dwi.header.get_zooms()[:3])
if not np.all(voxel_sizes == dwi.header.get_zooms()[0]):
print("* Careful voxel are anisotropic {}!".format(
tuple(voxel_sizes)))
# Since we are tracking in diffusion voxel space, convert step_size (in mm) to voxel.
if args.step_size is not None:
step_size = np.float32(args.step_size / voxel_sizes.max())
# Also convert max length (in mm) to voxel.
max_nb_points = int(np.ceil(args.max_length / args.step_size))
else:
step_size = None
max_nb_points = args.max_length
if args.theta is not None:
theta = np.deg2rad(args.theta)
elif args.curvature is not None and args.curvature > 0:
theta = get_max_angle_from_curvature(args.curvature, step_size)
else:
theta = np.deg2rad(45)
print("Angle: {}".format(np.rad2deg(theta)))
print("Step size (vox): {}".format(step_size))
print("Max nb. points: {}".format(max_nb_points))
is_outside_mask = make_is_outside_mask(mask,
affine_maskvox2dwivox,
threshold=args.mask_threshold)
is_too_long = make_is_too_long(max_nb_points)
is_too_curvy = make_is_too_curvy(np.rad2deg(theta))
is_unlikely = make_is_unlikely(0.5)
is_stopping = make_is_stopping({
STOPPING_MASK: is_outside_mask,
STOPPING_LENGTH: is_too_long,
STOPPING_CURVATURE: is_too_curvy,
STOPPING_LIKELIHOOD: is_unlikely
})
is_stopping.max_nb_points = max_nb_points # Small hack
tractogram = batch_track(model,
weights,
seeds,
step_size=step_size,
is_stopping=is_stopping,
batch_size=args.batch_size,
args=args)
# Streamlines have been generated in voxel space.
# Transform them them back to RAS+mm space using the dwi's affine.
tractogram.affine_to_rasmm = dwi.affine
tractogram.to_world() # Performed in-place.
nb_streamlines = len(tractogram)
if args.save_rejected:
rejected_tractogram = Tractogram()
rejected_tractogram.affine_to_rasmm = tractogram._affine_to_rasmm
print("Generated {:,} (compressed) streamlines".format(nb_streamlines))
with Timer("Cleaning streamlines", newline=True):
# Flush streamlines that have no points.
if args.save_rejected:
rejected_tractogram += tractogram[
np.array(list(map(len, tractogram))) <= 0]
tractogram = tractogram[np.array(list(map(len, tractogram))) > 0]
print("Removed {:,} empty streamlines".format(nb_streamlines -
len(tractogram)))
# Remove small streamlines
nb_streamlines = len(tractogram)
lengths = dipy.tracking.streamline.length(tractogram.streamlines)
if args.save_rejected:
rejected_tractogram += tractogram[lengths < args.min_length]
tractogram = tractogram[lengths >= args.min_length]
lengths = lengths[lengths >= args.min_length]
if len(lengths) > 0:
print("Average length: {:.2f} mm.".format(lengths.mean()))
print("Minimum length: {:.2f} mm. Maximum length: {:.2f}".format(
lengths.min(), lengths.max()))
print("Removed {:,} streamlines smaller than {:.2f} mm".format(
nb_streamlines - len(tractogram), args.min_length))
if args.discard_stopped_by_curvature:
nb_streamlines = len(tractogram)
stopping_curvature_flag_is_set = is_flag_set(
tractogram.data_per_streamline['stopping_flags'][:, 0],
STOPPING_CURVATURE)
if args.save_rejected:
rejected_tractogram += tractogram[
stopping_curvature_flag_is_set]
tractogram = tractogram[np.logical_not(
stopping_curvature_flag_is_set)]
print(
"Removed {:,} streamlines stopped for having a curvature higher than {:.2f} degree"
.format(nb_streamlines - len(tractogram), np.rad2deg(theta)))
if args.filter_threshold is not None:
# Remove streamlines that produces a reconstruction error higher than a certain threshold.
nb_streamlines = len(tractogram)
losses = compute_loss_errors(tractogram.streamlines, model,
hyperparams)
print("Mean loss: {:.4f} ± {:.4f}".format(
np.mean(losses),
np.std(losses, ddof=1) / np.sqrt(len(losses))))
if args.save_rejected:
rejected_tractogram += tractogram[
losses > args.filter_threshold]
tractogram = tractogram[losses <= args.filter_threshold]
print(
"Removed {:,} streamlines producing a loss lower than {:.2f} mm"
.format(nb_streamlines - len(tractogram),
args.filter_threshold))
with Timer("Saving {:,} (compressed) streamlines".format(len(tractogram))):
filename = args.out
if args.out is None:
prefix = args.prefix
if prefix is None:
dwi_name = os.path.basename(args.dwi)
if dwi_name.endswith(".nii.gz"):
dwi_name = dwi_name[:-7]
else: # .nii
dwi_name = dwi_name[:-4]
prefix = os.path.basename(os.path.dirname(args.dwi)) + dwi_name
prefix = prefix.replace(".", "_")
seed_mask_type = args.seeds[0].replace(".", "_").replace(
"_", "").replace("/", "-")
if "int" in args.seeds[0]:
seed_mask_type = "int"
elif "wm" in args.seeds[0]:
seed_mask_type = "wm"
elif "rois" in args.seeds[0]:
seed_mask_type = "rois"
elif "bundles" in args.seeds[0]:
seed_mask_type = "bundles"
mask_type = ""
if "fa" in args.mask:
mask_type = "fa"
elif "wm" in args.mask:
mask_type = "wm"
if args.dilate_seeding_mask:
seed_mask_type += "D"
if args.dilate_mask:
mask_type += "D"
filename_items = [
"{}",
"useMaxComponent-{}",
# "seed-{}",
# "mask-{}",
"step-{:.2f}mm",
"nbSeeds-{}",
"maxAngleDeg-{:.1f}"
# "keepCurv-{}",
# "filtered-{}",
# "minLen-{}",
# "pftRetry-{}",
# "pftHist-{}",
# "trackLikePeter-{}",
]
filename = ('_'.join(filename_items) + ".tck").format(
prefix,
args.use_max_component,
# seed_mask_type,
# mask_type,
args.step_size,
args.nb_seeds_per_voxel,
np.rad2deg(theta)
# not args.discard_stopped_by_curvature,
# args.filter_threshold,
# args.min_length,
# args.pft_nb_retry,
# args.pft_nb_backtrack_steps,
# args.track_like_peter
)
save_path = pjoin(experiment_path, filename)
try: # Create dirs, if needed.
os.makedirs(os.path.dirname(save_path))
except:
pass
print("Saving to {}".format(save_path))
nib.streamlines.save(tractogram, save_path)
if args.save_rejected:
with Timer("Saving {:,} (compressed) rejected streamlines".format(
len(rejected_tractogram))):
rejected_filename_items = filename_items.copy()
rejected_filename_items.insert(1, "rejected")
rejected_filename = (
'_'.join(rejected_filename_items) + ".tck"
).format(
prefix,
args.use_max_component,
# seed_mask_type,
# mask_type,
args.step_size,
args.nb_seeds_per_voxel,
np.rad2deg(theta)
# not args.discard_stopped_by_curvature,
# args.filter_threshold,
# args.min_length,
# args.pft_nb_retry,
# args.pft_nb_backtrack_steps,
# args.track_like_peter
)
rejected_save_path = pjoin(experiment_path, rejected_filename)
try: # Create dirs, if needed.
os.makedirs(os.path.dirname(rejected_save_path))
except:
pass
print("Saving rejected streamlines to {}".format(rejected_save_path))
nib.streamlines.save(rejected_tractogram, rejected_save_path)
def test_gru_mixture_track():
hidden_sizes = 50
with Timer("Creating dummy volume", newline=True):
volume_manager = neurotools.VolumeManager()
dwi, gradients = make_dummy_dwi(nb_gradients=30,
volume_shape=(10, 10, 10),
seed=1234)
volume = neurotools.resample_dwi(dwi, gradients.bvals,
gradients.bvecs).astype(np.float32)
volume_manager.register(volume)
with Timer("Creating model"):
hyperparams = {
'model': 'gru_mixture',
'SGD': "1e-2",
'hidden_sizes': hidden_sizes,
'learn_to_stop': False,
'normalize': False,
'activation': 'tanh',
'feed_previous_direction': False,
'predict_offset': False,
'use_layer_normalization': False,
'drop_prob': 0.,
'use_zoneout': False,
'skip_connections': False,
'neighborhood_radius': None,
'nb_seeds_per_voxel': 2,
'step_size': 0.5,
'batch_size': 200,
'n_gaussians': 2,
'seed': 1234
}
model = factories.model_factory(
hyperparams,
input_size=volume_manager.data_dimension,
output_size=3,
volume_manager=volume_manager)
model.initialize(
factories.weigths_initializer_factory("orthogonal", seed=1234))
rng = np.random.RandomState(1234)
mask = np.ones(volume.shape[:3])
seeding_mask = np.random.randint(2, size=mask.shape)
seeds = []
indices = np.array(np.where(seeding_mask)).T
for idx in indices:
seeds_in_voxel = idx + rng.uniform(
-0.5, 0.5, size=(hyperparams['nb_seeds_per_voxel'], 3))
seeds.extend(seeds_in_voxel)
seeds = np.array(seeds, dtype=theano.config.floatX)
is_outside_mask = make_is_outside_mask(mask, np.eye(4), threshold=0.5)
is_too_long = make_is_too_long(150)
is_too_curvy = make_is_too_curvy(np.rad2deg(30))
is_unlikely = make_is_unlikely(0.5)
is_stopping = make_is_stopping({
STOPPING_MASK: is_outside_mask,
STOPPING_LENGTH: is_too_long,
STOPPING_CURVATURE: is_too_curvy,
STOPPING_LIKELIHOOD: is_unlikely
})
is_stopping.max_nb_points = 150
args = SimpleNamespace()
args.track_like_peter = False
args.pft_nb_retry = 0
args.pft_nb_backtrack_steps = 0
args.use_max_component = False
args.flip_x = False
args.flip_y = False
args.flip_z = False
args.verbose = True
tractogram = batch_track(model,
volume,
seeds,
step_size=hyperparams['step_size'],
is_stopping=is_stopping,
batch_size=hyperparams['batch_size'],
args=args)
return True
def test_gru_mixture_fprop_neighborhood():
hyperparams = {
'model': 'gru_mixture',
'SGD': "1e-2",
'hidden_sizes': 50,
'batch_size': 16,
'learn_to_stop': False,
'normalize': True,
'activation': 'tanh',
'feed_previous_direction': False,
'predict_offset': False,
'use_layer_normalization': False,
'drop_prob': 0.,
'use_zoneout': False,
'skip_connections': False,
'seed': 1234,
'noisy_streamlines_sigma': None,
'keep_step_size': True,
'sort_streamlines': False,
'n_gaussians': 2,
'neighborhood_radius': 0.5
}
with Timer("Creating dataset", newline=True):
volume_manager = neurotools.VolumeManager()
trainset = make_dummy_dataset(volume_manager)
print("Dataset sizes:", len(trainset))
batch_scheduler = factories.batch_scheduler_factory(hyperparams,
dataset=trainset)
print("An epoch will be composed of {} updates.".format(
batch_scheduler.nb_updates_per_epoch))
print(volume_manager.data_dimension, hyperparams['hidden_sizes'],
batch_scheduler.target_size)
with Timer("Creating model"):
model = factories.model_factory(
hyperparams,
input_size=volume_manager.data_dimension,
output_size=batch_scheduler.target_size,
volume_manager=volume_manager)
model.initialize(
factories.weigths_initializer_factory("orthogonal", seed=1234))
print("Input size: {}".format(model.model_input_size))
# Test fprop with missing streamlines from one subject in a batch
output = model.get_output(trainset.symb_inputs)
fct = theano.function([trainset.symb_inputs],
output,
updates=model.graph_updates)
batch_inputs, batch_targets, batch_mask = batch_scheduler._next_batch(2)
out = fct(batch_inputs)
with Timer("Building optimizer"):
loss = factories.loss_factory(hyperparams, model, trainset)
optimizer = factories.optimizer_factory(hyperparams, loss)
fct_loss = theano.function(
[trainset.symb_inputs, trainset.symb_targets, trainset.symb_mask],
loss.loss,
updates=model.graph_updates)
loss_value = fct_loss(batch_inputs, batch_targets, batch_mask)
print("Loss:", loss_value)
fct_optim = theano.function(
[trainset.symb_inputs, trainset.symb_targets, trainset.symb_mask],
list(optimizer.directions.values()),
updates=model.graph_updates)
dirs = fct_optim(batch_inputs, batch_targets, batch_mask)
return True
def main():
parser = build_argparser()
args = parser.parse_args()
# Get experiment folder
experiment_path = args.name
if not os.path.isdir(experiment_path):
# If not a directory, it must be the name of the experiment.
experiment_path = pjoin(".", "experiments", args.name)
if not os.path.isdir(experiment_path):
parser.error('Cannot find experiment: {0}!'.format(args.name))
# Load experiments hyperparameters
try:
hyperparams = smartutils.load_dict_from_json_file(
pjoin(experiment_path, "hyperparams.json"))
except FileNotFoundError:
hyperparams = smartutils.load_dict_from_json_file(
pjoin(experiment_path, "..", "hyperparams.json"))
with Timer("Loading DWIs"):
# Load gradients table
dwi_name = args.dwi
if dwi_name.endswith(".gz"):
dwi_name = dwi_name[:-3]
if dwi_name.endswith(".nii"):
dwi_name = dwi_name[:-4]
bvals_filename = dwi_name + ".bvals"
bvecs_filename = dwi_name + ".bvecs"
bvals, bvecs = dipy.io.gradients.read_bvals_bvecs(
bvals_filename, bvecs_filename)
dwi = nib.load(args.dwi)
if hyperparams["use_sh_coeffs"]:
# Use 45 spherical harmonic coefficients to represent the diffusion signal.
weights = neurotools.get_spherical_harmonics_coefficients(
dwi, bvals, bvecs).astype(np.float32)
else:
# Resample the diffusion signal to have 100 directions.
weights = neurotools.resample_dwi(dwi, bvals,
bvecs).astype(np.float32)
with Timer("Loading model"):
if hyperparams["model"] == "ffnn_classification":
from learn2track.models import FFNN_Classification
model_class = FFNN_Classification
else:
raise ValueError("Unknown model!")
kwargs = {}
volume_manager = neurotools.VolumeManager()
volume_manager.register(weights)
kwargs['volume_manager'] = volume_manager
# Load the actual model.
model = model_class.create(
pjoin(experiment_path),
**kwargs) # Create new instance and restore model.
print(str(model))
with Timer("Generating mask"):
symb_input = T.matrix(name="input")
model_symb_pred = model.get_output(symb_input)
f = theano.function(inputs=[symb_input], outputs=[model_symb_pred])
generated_mask = np.zeros(dwi.shape[:3]).astype(np.float32)
# all_coords.shape = (n_coords, 3)
all_coords = np.argwhere(generated_mask == 0)
volume_ids = np.zeros((all_coords.shape[0], 1))
all_coords_and_volume_ids = np.concatenate((all_coords, volume_ids),
axis=1).astype(np.float32)
batch_size = args.batch_size if args.batch_size else len(
all_coords_and_volume_ids)
probs = []
while batch_size > 1:
print("Trying to to process batches of size {} out of {}".format(
batch_size, len(all_coords_and_volume_ids)))
nb_batches = int(
np.ceil(len(all_coords_and_volume_ids) / batch_size))
try:
for batch_count in range(nb_batches):
start = batch_count * batch_size
end = (batch_count + 1) * batch_size
probs.extend(f(all_coords_and_volume_ids[start:end])[-1])
print("Generated batch {} out of {}".format(
batch_count + 1, nb_batches))
break
except MemoryError:
print("{} coordinates at the same time is too much!".format(
batch_size))
batch_size //= 2
except RuntimeError:
print("{} coordinates at the same time is too much!".format(
batch_size))
batch_size //= 2
if not probs:
raise RuntimeError("Could not generate predictions...")
generated_mask[np.where(generated_mask == 0)] = np.array(probs) > 0.5
with Timer("Saving generated mask"):
filename = args.out
if args.out is None:
prefix = args.prefix
if prefix is None:
dwi_name = os.path.basename(args.dwi)
if dwi_name.endswith(".nii.gz"):
dwi_name = dwi_name[:-7]
else: # .nii
dwi_name = dwi_name[:-4]
prefix = os.path.basename(os.path.dirname(args.dwi)) + dwi_name
prefix = prefix.replace(".", "_")
filename = "{}.nii.gz".format(prefix)
save_path = pjoin(experiment_path, filename)
try: # Create dirs, if needed.
os.makedirs(os.path.dirname(save_path))
except:
pass
print("Saving to {}".format(save_path))
mask = nib.Nifti1Image(generated_mask, dwi.affine)
nib.save(mask, save_path)