def __call__(self, path): # load the image x = np.array(load_image(path), copy=False).astype(np.float32) if self.scale > 1: x = downsample(x, self.scale) # normalize it method = 'gmm' if self.affine: method = 'affine' x,metadata = normalize(x, alpha=self.alpha, beta=self.beta, num_iters=self.num_iters , method=method, sample=self.sample, use_cuda=self.use_cuda) # save the image and the metadata name,_ = os.path.splitext(os.path.basename(path)) base = os.path.join(self.dest, name) for f in self.formats: save_image(x, base, f=f) if self.metadata: # save the metadata in json format mdpath = base + '.metadata.json' if not self.affine: metadata['mus'] = metadata['mus'].tolist() metadata['stds'] = metadata['stds'].tolist() metadata['pis'] = metadata['pis'].tolist() metadata['logps'] = metadata['logps'].tolist() with open(mdpath, 'w') as f: json.dump(metadata, f, indent=4) return name
def main(args): verbose = args.verbose N = len(args.files) image_iterator = load_images(args.files) ## downsample the images names = [] images = [] for name, image in downsample_images(image_iterator, args.scale, num_workers=args.num_workers): names.append(name) images.append(image) if verbose: print('# [{} of {}] downsampled: {}'.format(len(images), N, name)) niters = args.niters samples = args.pixel_sampling seed = args.seed if verbose: print('# fit scaled GMM, niters={}, sample={}, seed={}'.format( niters, samples, seed)) images, metadata = sgmm_scaling(images, niters, samples, seed, verbose) if verbose: print('# weights:', metadata['weights']) print('# means:', metadata['means']) print('# variances:', metadata['variances']) scales = {} for name, scale in zip(names, metadata['scales']): scales[name] = scale metadata['scales'] = scales destdir = args.destdir if not os.path.exists(destdir): os.makedirs(destdir) ## what image formats are we writing formats = args.format_.split(',') for name, x in zip(names, images): base = os.path.join(destdir, name) for f in formats: save_image(x, base, f=f, verbose=verbose) ## save the metadata in json format path = os.path.join(destdir, 'metadata.json') with open(path, 'w') as f: if verbose: print('# saving metadata:', path) json.dump(metadata, f, indent=4)
def main(args): images = load_images(args.files) names = list(images.keys()) images = list(images.values()) if args.verbose: print('# fit scaled GMM, niters={}, sample={}, seed={}'.format( args.niters, args.sample, args.seed)) images, metadata = sgmm_scaling(images, args.niters, args.sample, args.seed, args.verbose) if args.verbose: print('# weights:', metadata['weights']) print('# means:', metadata['means']) print('# variances:', metadata['variances']) scales = {} for name, scale in zip(names, metadata['scales']): scales[name] = scale metadata['scales'] = scales destdir = args.destdir if not os.path.exists(destdir): os.makedirs(destdir) ## what image formats are we writing verbose = args.verbose formats = args.format_.split(',') for name, x in zip(names, images): base = os.path.join(destdir, name) for f in formats: save_image(x, base, f=f, verbose=verbose) ## save the metadata in json format path = os.path.join(destdir, 'metadata.json') with open(path, 'w') as f: if args.verbose: print('# saving metadata:', path) json.dump(metadata, f, indent=4)
def main(args): # set the number of threads num_threads = args.num_threads from topaz.torch import set_num_threads set_num_threads(num_threads) ## set the device use_cuda = topaz.cuda.set_device(args.device) print('# using device={} with cuda={}'.format(args.device, use_cuda), file=sys.stderr) cutoff = args.pixel_cutoff # pixel truncation limit do_train = (args.dir_a is not None and args.dir_b is not None) or (args.hdf is not None) if do_train: method = args.method paired = (method == 'noise2noise') preload = args.preload holdout = args.holdout # fraction of image pairs to holdout for validation if args.hdf is None: #use dirA/dirB crop = args.crop dir_as = args.dir_a dir_bs = args.dir_b dset_train = [] dset_val = [] for dir_a, dir_b in zip(dir_as, dir_bs): random = np.random.RandomState(44444) if paired: dataset_train, dataset_val = make_paired_images_datasets( dir_a, dir_b, crop, random=random, holdout=holdout, preload=preload, cutoff=cutoff) else: dataset_train, dataset_val = make_images_datasets( dir_a, dir_b, crop, cutoff=cutoff, random=random, holdout=holdout) dset_train.append(dataset_train) dset_val.append(dataset_val) dataset_train = dset_train[0] for i in range(1, len(dset_train)): dataset_train.x += dset_train[i].x if paired: dataset_train.y += dset_train[i].y dataset_val = dset_val[0] for i in range(1, len(dset_val)): dataset_val.x += dset_val[i].x if paired: dataset_val.y += dset_val[i].y shuffle = True else: # make HDF datasets dataset_train, dataset_val = make_hdf5_datasets(args.hdf, paired=paired, cutoff=cutoff, holdout=holdout, preload=preload) shuffle = preload # initialize the model arch = args.arch if arch == 'unet': model = dn.UDenoiseNet() elif arch == 'unet-small': model = dn.UDenoiseNetSmall() elif arch == 'unet2': model = dn.UDenoiseNet2() elif arch == 'unet3': model = dn.UDenoiseNet3() elif arch == 'fcnet': model = dn.DenoiseNet(32) elif arch == 'fcnet2': model = dn.DenoiseNet2(64) elif arch == 'affine': model = dn.AffineDenoise() else: raise Exception('Unknown architecture: ' + arch) if use_cuda: model = model.cuda() # train optim = args.optim lr = args.lr batch_size = args.batch_size num_epochs = args.num_epochs digits = int(np.ceil(np.log10(num_epochs))) num_workers = args.num_workers print('epoch', 'loss_train', 'loss_val') #criteria = nn.L1Loss() criteria = args.criteria if method == 'noise2noise': iterator = dn.train_noise2noise(model, dataset_train, lr=lr, optim=optim, batch_size=batch_size, criteria=criteria, num_epochs=num_epochs, dataset_val=dataset_val, use_cuda=use_cuda, num_workers=num_workers, shuffle=shuffle) elif method == 'masked': iterator = dn.train_mask_denoise(model, dataset_train, lr=lr, optim=optim, batch_size=batch_size, criteria=criteria, num_epochs=num_epochs, dataset_val=dataset_val, use_cuda=use_cuda, num_workers=num_workers, shuffle=shuffle) for epoch, loss_train, loss_val in iterator: print(epoch, loss_train, loss_val) sys.stdout.flush() # save the model if args.save_prefix is not None: path = args.save_prefix + ('_epoch{:0' + str(digits) + '}.sav').format(epoch) #path = args.save_prefix + '_epoch{}.sav'.format(epoch) model.cpu() model.eval() torch.save(model, path) if use_cuda: model.cuda() models = [model] else: # load the saved model(s) models = [] for arg in args.model: if arg == 'none': print('# Warning: no denoising model will be used', file=sys.stderr) else: print('# Loading model:', arg, file=sys.stderr) model = dn.load_model(arg) model.eval() if use_cuda: model.cuda() models.append(model) # using trained model # denoise the images normalize = args.normalize if args.format_ == 'png' or args.format_ == 'jpg': # always normalize png and jpg format normalize = True format_ = args.format_ suffix = args.suffix lowpass = args.lowpass gaus = args.gaussian if gaus > 0: gaus = dn.GaussianDenoise(gaus) if use_cuda: gaus.cuda() else: gaus = None inv_gaus = args.inv_gaussian if inv_gaus > 0: inv_gaus = dn.InvGaussianFilter(inv_gaus) if use_cuda: inv_gaus.cuda() else: inv_gaus = None deconvolve = args.deconvolve deconv_patch = args.deconv_patch ps = args.patch_size padding = args.patch_padding count = 0 # we are denoising a single MRC stack if args.stack: with open(args.micrographs[0], 'rb') as f: content = f.read() stack, _, _ = mrc.parse(content) print('# denoising stack with shape:', stack.shape, file=sys.stderr) total = len(stack) denoised = np.zeros_like(stack) for i in range(len(stack)): mic = stack[i] # process and denoise the micrograph mic = denoise_image(mic, models, lowpass=lowpass, cutoff=cutoff, gaus=gaus, inv_gaus=inv_gaus, deconvolve=deconvolve, deconv_patch=deconv_patch, patch_size=ps, padding=padding, normalize=normalize, use_cuda=use_cuda) denoised[i] = mic count += 1 print('# {} of {} completed.'.format(count, total), file=sys.stderr, end='\r') print('', file=sys.stderr) # write the denoised stack path = args.output print('# writing', path, file=sys.stderr) with open(path, 'wb') as f: mrc.write(f, denoised) else: # stream the micrographs and denoise them total = len(args.micrographs) # make the output directory if it doesn't exist if not os.path.exists(args.output): os.makedirs(args.output) for path in args.micrographs: name, _ = os.path.splitext(os.path.basename(path)) mic = np.array(load_image(path), copy=False).astype(np.float32) # process and denoise the micrograph mic = denoise_image(mic, models, lowpass=lowpass, cutoff=cutoff, gaus=gaus, inv_gaus=inv_gaus, deconvolve=deconvolve, deconv_patch=deconv_patch, patch_size=ps, padding=padding, normalize=normalize, use_cuda=use_cuda) # write the micrograph if not args.output: if suffix == '' or suffix is None: suffix = '.denoised' # write the file to the same location as input no_ext, ext = os.path.splitext(path) outpath = no_ext + suffix + '.' + format_ else: outpath = args.output + os.sep + name + suffix + '.' + format_ save_image(mic, outpath) #, mi=None, ma=None) count += 1 print('# {} of {} completed.'.format(count, total), file=sys.stderr, end='\r') print('', file=sys.stderr)
def main(args): ## set the device use_cuda = False if args.device >= 0: use_cuda = torch.cuda.is_available() if use_cuda: torch.cuda.set_device(args.device) print('# using device={} with cuda={}'.format(args.device, use_cuda), file=sys.stderr) do_train = (args.dir_a is not None and args.dir_b is not None) or (args.hdf is not None) if do_train: if args.hdf is None: #use dirA/dirB crop = args.crop dir_a = args.dir_a dir_b = args.dir_b random = np.random.RandomState(44444) dataset_train, dataset_val = make_paired_images_datasets( dir_a, dir_b, crop, random=random) shuffle = True else: # make HDF datasets dataset_train, dataset_val = make_hdf5_datasets(args.hdf) shuffle = False # initialize the model #model = dn.DenoiseNet(32) model = dn.UDenoiseNet() if use_cuda: model = model.cuda() # train lr = args.lr batch_size = args.batch_size num_epochs = args.num_epochs num_workers = args.num_workers print('epoch', 'loss_train', 'loss_val') #criteria = nn.L1Loss() criteria = args.criteria for epoch, loss_train, loss_val in dn.train_noise2noise( model, dataset_train, lr=lr, batch_size=batch_size, criteria=criteria, num_epochs=num_epochs, dataset_val=dataset_val, use_cuda=use_cuda, num_workers=num_workers, shuffle=shuffle): print(epoch, loss_train, loss_val) sys.stdout.flush() # save the model if args.save_prefix is not None: path = args.save_prefix + '_epoch{}.sav'.format(epoch) model.cpu() model.eval() torch.save(model, path) if use_cuda: model.cuda() else: # load the saved model if args.model in ['L0', 'L1', 'L2']: if args.model in ['L0', 'L1']: print( 'ERROR: L0 and L1 models are not implemented in the current version', file=sys.stderr) sys.exit(1) model = dn.load_model(args.model) else: model = torch.load(args.model) print('# using model:', args.model, file=sys.stderr) model.eval() if use_cuda: model.cuda() if args.stack: # we are denoising a single MRC stack with open(args.micrographs[0], 'rb') as f: content = f.read() stack, _, _ = mrc.parse(content) print('# denoising stack with shape:', stack.shape, file=sys.stderr) denoised = dn.denoise_stack(model, stack, use_cuda=use_cuda) # write the denoised stack path = args.output print('# writing', path, file=sys.stderr) with open(path, 'wb') as f: mrc.write(f, denoised) else: # using trained model # stream the micrographs and denoise as we go normalize = args.normalize if args.format_ == 'png' or args.format_ == 'jpg': # always normalize png and jpg format normalize = True format_ = args.format_ count = 0 total = len(args.micrographs) bin_ = args.bin ps = args.patch_size padding = args.patch_padding # now, stream the micrographs and denoise them for path in args.micrographs: name, _ = os.path.splitext(os.path.basename(path)) mic = np.array(load_image(path), copy=False) if bin_ > 1: mic = downsample(mic, bin_) mu = mic.mean() std = mic.std() # denoise mic = (mic - mu) / std mic = dn.denoise(model, mic, patch_size=ps, padding=padding, use_cuda=use_cuda) if normalize: mic = (mic - mic.mean()) / mic.std() else: # add back std. dev. and mean mic = std * mic + mu # write the micrograph outpath = args.output + os.sep + name + '.' + format_ save_image(mic, outpath) count += 1 print('# {} of {} completed.'.format(count, total), file=sys.stderr, end='\r') print('', file=sys.stderr)