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):
## load image
path = args.file
im = load_image(path)
# convert PIL image to array
im = np.array(im, copy=False).astype(np.float32)
scale = args.scale # how much to downscale by
small = downsample(im, scale)
if args.verbose:
print('Downsample image:', path, file=sys.stderr)
print('From', im.shape, 'to', small.shape, file=sys.stderr)
# write the downsampled image
with open(args.output, 'wb') as f:
im = Image.fromarray(small)
if small.dtype == np.uint8:
im.save(f, 'png')
else:
im.save(f, 'tiff')
def __call__(self, args):
name, image = args
# convert PIL image to array
image = np.array(image, copy=False)
image = downsample(image, self.scale)
return name, image
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)
def main(args):
particles = pd.read_csv(args.file, sep='\t')
print('#', 'Loaded', len(particles), 'particles', file=sys.stderr)
# threshold the particles
if 'score' in particles:
particles = particles.loc[particles['score'] >= args.threshold]
print('#', 'Thresholding at', args.threshold, file=sys.stderr)
print('#', 'Extracting', len(particles), 'particles', file=sys.stderr)
N = len(particles)
size = args.size
resize = args.resize
if resize < 0:
resize = size
#
wrote_header = False
read_metadata = False
metadata = []
# write the particles iteratively
i = 0
with open(args.output, 'wb') as f:
for image_name, coords in particles.groupby('image_name'):
print('#', image_name, len(coords), 'particles', file=sys.stderr)
# load the micrograph
image_name = image_name + args.image_ext
path = os.path.join(args.image_root, image_name)
with open(path, 'rb') as fm:
content = fm.read()
micrograph, header, extended_header = mrc.parse(content)
if len(micrograph.shape) < 3:
micrograph = micrograph[
np.newaxis] # add z dim if micrograph is image
if not wrote_header: # load a/px and angles from micrograph header and write the stack header
mz = micrograph.shape[0]
dtype = micrograph.dtype
cella = (header.xlen, header.ylen, header.zlen)
cellb = (header.alpha, header.beta, header.gamma)
shape = (N * mz, resize, resize)
header = mrc.make_header(shape,
cella,
cellb,
mz=mz,
dtype=dtype)
buf = mrc.header_struct.pack(*list(header))
f.write(buf)
wrote_header = True
_, n, m = micrograph.shape
x_coord = coords['x_coord'].values
y_coord = coords['y_coord'].values
scores = None
if 'score' in coords:
scores = coords['score'].values
# crop out the particles
for j in range(len(coords)):
x = x_coord[j]
y = y_coord[j]
if scores is not None:
metadata.append((image_name, x, y, scores[j]))
else:
metadata.append((image_name, x, y))
left = x - size // 2
upper = y - size // 2
right = left + size
lower = upper + size
c = micrograph[:,
max(0, upper):min(n, lower),
max(0, left):min(m, right)]
c = (c - c.mean()) / c.std()
stack = np.zeros((mz, size, size), dtype=dtype)
#stack = np.zeros((mz, size, size), dtype=dtype) + c.mean().astype(dtype)
stack[:,
max(0, -upper):min(size + n - lower, size),
max(0, -left):min(size + m - right, size)] = c
# write particle to mrc file
if resize != size:
restack = downsample(stack, 0, shape=(resize, resize))
#print(restack.shape, restack.mean(), restack.std())
restack = (restack - restack.mean()) / restack.std()
f.write(restack.tobytes())
else:
f.write(stack.tobytes())
i += 1
#print('# wrote', i, 'out of', N, 'particles', end='\r', flush=True)
## write the particle stack mrcs
#with open(args.output, 'wb') as f:
# mrc.write(f, stack, ax=ax, ay=ay, az=az, alpha=alpha, beta=beta, gamma=gamma)
image_name = os.path.basename(args.output)
star_path = os.path.splitext(args.output)[0] + '.star'
## create the star file
columns = ['MicrographName', star.X_COLUMN_NAME, star.Y_COLUMN_NAME]
if 'score' in particles:
columns.append(star.SCORE_COLUMN_NAME)
metadata = pd.DataFrame(metadata, columns=columns)
metadata['ImageName'] = [
str(i + 1) + '@' + image_name for i in range(len(metadata))
]
if mz > 1:
metadata['NrOfFrames'] = mz
micrograph_metadata = None
if args.metadata is not None:
with open(args.metadata, 'r') as f:
micrograph_metadata = star.parse_star(f)
metadata = pd.merge(metadata,
micrograph_metadata,
on='MicrographName',
how='left')
if resize != size and 'DetectorPixelSize' in metadata:
# rescale the detector pixel size
pix = metadata['DetectorPixelSize'].values.astype(float)
metadata['DetectorPixelSize'] = pix * (size / resize)
## write the star file
with open(star_path, 'w') as f:
star.write(metadata, f)