def main(args=None):
args, no_config_file = get_args(args)
setup_log(args.verbosity)
logger = logging.getLogger(__name__)
try:
# set random seeds for reproducibility
torch.manual_seed(args.seed)
np.random.seed(args.seed)
# since prediction only uses one gpu (at most), make the batch size small enough to fit
if args.n_gpus > 1: args.batch_size = args.batch_size // args.n_gpus
learner = Learner.predict_setup(args)
# determine how many samples we will use in prediction
nsyn = args.monte_carlo or 1
# get relevant prediction directories and determine extension
predict_dir = args.predict_dir or args.valid_source_dir
output_dir = args.predict_out or os.getcwd() + '/syn_'
ext = determine_ext(predict_dir[0])
# setup and start prediction loop
axis = args.sample_axis or 0
if axis < 0 or axis > 2 and not isinstance(axis, int):
raise ValueError(
'sample_axis must be an integer between 0 and 2 inclusive')
n_imgs = len(glob_imgs(predict_dir[0], ))
if n_imgs == 0:
raise SynthNNError(
'Prediction directory does not contain valid images.')
if any([len(glob_imgs(pd, ext)) != n_imgs for pd in predict_dir]):
raise SynthNNError(
'Number of images in prediction directories must have an equal number of images in each '
'directory (e.g., so that img_t1_1 aligns with img_t2_1 etc. for multimodal synth)'
)
predict_fns = zip(*[glob_imgs(pd, ext) for pd in predict_dir])
if args.net3d and args.patch_size > 0 and args.calc_var:
raise SynthNNError(
'Patch-based 3D variance calculation not currently supported.')
for k, fn in enumerate(predict_fns):
_, base, ext = split_filename(fn[0])
logger.info(
f'Starting synthesis of image: {base} ({k+1}/{n_imgs})')
out_imgs = learner.predict(fn, nsyn, args.temperature_map,
args.calc_var)
for i, oin in enumerate(out_imgs):
out_fn = output_dir + f'{k}_{i}{ext}'
if hasattr(oin, 'to_filename'):
oin.to_filename(out_fn)
else:
oin.save(out_fn)
logger.info(f'Finished synthesis. Saved as: {out_fn}')
return 0
except Exception as e:
logger.exception(e)
return 1
def setUp(self):
wd = os.path.dirname(os.path.abspath(__file__))
self.nii_dir = os.path.join(wd, 'test_data', 'nii')
self.mask_dir = os.path.join(wd, 'test_data', 'masks')
self.tif_dir = os.path.join(wd, 'test_data', 'tif')
self.png_dir = os.path.join(wd, 'test_data', 'png')
self.out_dir = tempfile.mkdtemp()
os.mkdir(os.path.join(self.out_dir, 'models'))
self.train_dir = os.path.join(self.out_dir, 'imgs')
os.mkdir(self.train_dir)
os.mkdir(os.path.join(self.train_dir, 'mask'))
os.mkdir(os.path.join(self.train_dir, 'tif'))
os.mkdir(os.path.join(self.train_dir, 'png'))
nii = glob_imgs(self.nii_dir)[0]
msk = glob_imgs(self.mask_dir)[0]
tif = os.path.join(self.tif_dir, 'test.tif')
png = os.path.join(self.png_dir, 'test.png')
path, base, ext = split_filename(nii)
for i in range(8):
shutil.copy(nii, os.path.join(self.train_dir, base + str(i) + ext))
shutil.copy(
msk, os.path.join(self.train_dir, 'mask', base + str(i) + ext))
shutil.copy(
tif, os.path.join(self.train_dir, 'tif',
base + str(i) + '.tif'))
shutil.copy(
png, os.path.join(self.train_dir, 'png',
base + str(i) + '.png'))
self.train_args = f'-s {self.train_dir} -t {self.train_dir}'.split()
self.predict_args = f'-s {self.train_dir} -o {self.out_dir}/test'.split(
)
self.jsonfn = f'{self.out_dir}/test.json'
def determine_ext(d):
""" given a directory determine if it contains supported images """
exts = ('*.nii*', '*.tif*', '*.png')
contains = [len(glob_imgs(d, ext)) > 0 for ext in exts]
if sum(contains) == 0:
raise SynthtorchError(f'Directory {d} contains no supported images.')
if sum(contains) > 1:
raise SynthtorchError(f'Directory {d} contains more than two types of supported images, '
f'remove unwanted images from directory')
ext = [e for c, e in zip(contains, exts) if c][0]
return ext