def rotate_and_project_batch(obj_delta, obj_beta):
obj_rot_batch = []
for i in range(minibatch_size):
obj_rot_batch.append(tf_rotate(tf.stack([obj_delta, obj_beta], axis=-1), this_theta_batch[i], interpolation='BILINEAR'))
# obj_rot = apply_rotation(obj, coord_ls[rand_proj], 'arrsize_64_64_64_ntheta_500')
obj_rot_batch = tf.stack(obj_rot_batch)
if probe_type == 'point':
exiting_batch = multislice_propagate_spherical(obj_rot_batch[:, :, :, :, 0], obj_rot_batch[:, :, :, :, 1],
probe_real, probe_imag, energy_ev,
psize_cm * ds_level, dist_to_source_cm, det_psize_cm,
theta_max, phi_max, free_prop_cm,
obj_batch_shape=[minibatch_size, *obj_size])
else:
if forward_algorithm == 'fresnel':
exiting_batch = multislice_propagate_batch(obj_rot_batch[:, :, :, :, 0], obj_rot_batch[:, :, :, :, 1],
probe_real, probe_imag, energy_ev,
psize_cm * ds_level, free_prop_cm=free_prop_cm, obj_batch_shape=[minibatch_size, *obj_size])
elif forward_algorithm == 'fd':
exiting_batch = multislice_propagate_fd(obj_rot_batch[:, :, :, :, 0], obj_rot_batch[:, :, :, :, 1],
probe_real, probe_imag, energy_ev,
psize_cm * ds_level, free_prop_cm=free_prop_cm,
obj_batch_shape=[minibatch_size, *obj_size])
loss = tf.reduce_mean(tf.squared_difference(tf.abs(exiting_batch), tf.abs(this_prj_batch)), name='loss')
return loss, exiting_batch
def rotate_crop(img,
rotation,
crop=True,
minimum_shape=[0, 0],
interpolation='NEAREST'):
with tf.name_scope('RotateCrop'):
rotated_image = tf_rotate(img, rotation, interpolation)
if crop:
rotation = tf.abs(rotation)
original_shape = tf.shape(rotated_image)[:2]
h, w = original_shape[0], original_shape[1]
# see https://stackoverflow.com/questions/16702966/rotate-image-and-crop-out-black-borders for formulae
old_l, old_s = tf.cond(h > w, lambda: [h, w], lambda: [w, h])
old_l, old_s = tf.cast(old_l,
tf.float32), tf.cast(old_s, tf.float32)
new_l = (old_l * tf.cos(rotation) -
old_s * tf.sin(rotation)) / tf.cos(2 * rotation)
new_s = (old_s - tf.sin(rotation) * new_l) / tf.cos(rotation)
new_h, new_w = tf.cond(h > w, lambda: [new_l, new_s],
lambda: [new_s, new_l])
new_h, new_w = tf.cast(new_h, tf.int32), tf.cast(new_w, tf.int32)
bb_begin = tf.cast(tf.ceil((h - new_h) / 2),
tf.int32), tf.cast(tf.ceil((w - new_w) / 2),
tf.int32)
rotated_image_crop = rotated_image[bb_begin[0]:h - bb_begin[0],
bb_begin[1]:w - bb_begin[1], :]
# If crop removes the entire image, keep the original image
rotated_image = tf.cond(tf.less_equal(
tf.reduce_min(tf.shape(rotated_image_crop)[:2]),
tf.reduce_max(minimum_shape)),
true_fn=lambda: img,
false_fn=lambda: rotated_image_crop)
return rotated_image
def rotate_crop(image: tf.Tensor, rotation: float, crop: bool=True, minimum_shape: Tuple[int, int]=[0, 0],
interpolation: str='NEAREST') -> tf.Tensor:
"""Rotates and crops the images.
:param image: image to be rotated and cropped [H, W, C]
:param rotation: angle of rotation (in radians)
:param crop: option to crop rotated image to avoid black borders due to rotation
:param minimum_shape: minimum shape of the rotated image / cropped image
:param interpolation: which interpolation to use ``NEAREST`` or ``BILINEAR``
:return:
"""
with tf.name_scope('RotateCrop'):
rotated_image = tf_rotate(image, rotation, interpolation)
if crop:
rotation = tf.abs(rotation)
original_shape = tf.shape(rotated_image)[:2]
h, w = original_shape[0], original_shape[1]
# see https://stackoverflow.com/questions/16702966/rotate-image-and-crop-out-black-borders for formulae
old_l, old_s = tf.cond(h > w, lambda: [h, w], lambda: [w, h])
old_l, old_s = tf.cast(old_l, tf.float32), tf.cast(old_s, tf.float32)
new_l = (old_l * tf.cos(rotation) - old_s * tf.sin(rotation)) / tf.cos(2 * rotation)
new_s = (old_s - tf.sin(rotation) * new_l) / tf.cos(rotation)
new_h, new_w = tf.cond(h > w, lambda: [new_l, new_s], lambda: [new_s, new_l])
new_h, new_w = tf.cast(new_h, tf.int32), tf.cast(new_w, tf.int32)
bb_begin = tf.cast(tf.ceil((h - new_h) / 2), tf.int32), tf.cast(tf.ceil((w - new_w) / 2), tf.int32)
rotated_image_crop = rotated_image[bb_begin[0]:h - bb_begin[0], bb_begin[1]:w - bb_begin[1], :]
# If crop removes the entire image, keep the original image
rotated_image = tf.cond(tf.less_equal(tf.reduce_min(tf.shape(rotated_image_crop)[:2]),
tf.reduce_max(minimum_shape)),
true_fn=lambda: image,
false_fn=lambda: rotated_image_crop)
return rotated_image
def rotate_and_project(i, obj):
# obj_rot = apply_rotation(obj, coord_ls[rand_proj], 'arrsize_64_64_64_ntheta_500')
obj_rot = tf_rotate(obj, theta_ls_tensor[i], interpolation='BILINEAR')
probe_pos_batch_ls = np.array_split(
probe_pos, int(np.ceil(float(n_pos) / n_dp_batch)))
# probe_pos_batch_ls = np.array_split(probe_pos, 6)
exiting_ls = []
# pad if needed
pad_arr = np.array([[0, 0], [0, 0]])
if probe_pos[:, 0].min() - probe_size_half[0] < 0:
pad_len = probe_size_half[0] - probe_pos[:, 0].min()
obj_rot = tf.pad(obj_rot, ((pad_len, 0), (0, 0), (0, 0), (0, 0)),
mode='CONSTANT')
pad_arr[0, 0] = pad_len
if probe_pos[:, 0].max() + probe_size_half[0] > obj_size[0]:
pad_len = probe_pos[:, 0].max() + probe_size_half[0] - obj_size[0]
obj_rot = tf.pad(obj_rot, ((0, pad_len), (0, 0), (0, 0), (0, 0)),
mode='CONSTANT')
pad_arr[0, 1] = pad_len
if probe_pos[:, 1].min() - probe_size_half[1] < 0:
pad_len = probe_size_half[1] - probe_pos[:, 1].min()
obj_rot = tf.pad(obj_rot, ((0, 0), (pad_len, 0), (0, 0), (0, 0)),
mode='CONSTANT')
pad_arr[1, 0] = pad_len
if probe_pos[:, 1].max() + probe_size_half[1] > obj_size[1]:
pad_len = probe_pos[:, 1].max() + probe_size_half[0] - obj_size[1]
obj_rot = tf.pad(obj_rot, ((0, 0), (0, pad_len), (0, 0), (0, 0)),
mode='CONSTANT')
pad_arr[1, 1] = pad_len
for k, pos_batch in enumerate(probe_pos_batch_ls):
subobj_ls = []
for j, pos in enumerate(pos_batch):
pos = [int(x) for x in pos]
pos[0] = pos[0] + pad_arr[0, 0]
pos[1] = pos[1] + pad_arr[1, 0]
subobj = obj_rot[pos[0] - probe_size_half[0]:pos[0] -
probe_size_half[0] + probe_size[0],
pos[1] - probe_size_half[1]:pos[1] -
probe_size_half[1] + probe_size[1], :, :]
subobj_ls.append(subobj)
subobj_ls = tf.stack(subobj_ls)
exiting = multislice_propagate_batch(
subobj_ls[:, :, :, :, 0],
subobj_ls[:, :, :, :, 1],
probe_real,
probe_imag,
energy_ev,
psize_cm,
h=h,
free_prop_cm='inf',
obj_batch_shape=[len(pos_batch), *probe_size, obj_size[-1]])
exiting_ls.append(exiting)
exiting_ls = tf.concat(exiting_ls, 0)
return exiting_ls
def rotate_and_project(i, obj):
obj_rot = tf_rotate(obj, theta_ls_tensor[i], interpolation='BILINEAR')
exiting = multislice_propagate(obj_rot[:, :, :, 0],
obj_rot[:, :, :, 1],
probe_real,
probe_imag,
energy_ev,
psize_cm,
free_prop_cm=free_prop_cm)
return exiting
def rotate_and_project(i, loss, obj_delta, obj_beta):
warnings.warn('Obsolete function. The output loss is scaled by minibatch_size. Proceed with caution.')
obj_rot = tf_rotate(tf.stack([obj_delta, obj_beta], axis=-1), this_theta_batch[i], interpolation='BILINEAR')
if not cpu_only:
with tf.device('/gpu:0'):
exiting = multislice_propagate(obj_rot[:, :, :, 0], obj_rot[:, :, :, 1], probe_real, probe_imag, energy_ev, psize_cm * ds_level, h=h, free_prop_cm=free_prop_cm)
else:
exiting = multislice_propagate(obj_rot[:, :, :, 0], obj_rot[:, :, :, 1], probe_real, probe_imag, energy_ev, psize_cm * ds_level, h=h, free_prop_cm=free_prop_cm)
loss += tf.reduce_mean(tf.squared_difference(tf.abs(exiting), tf.abs(this_prj_batch[i])))
# i = tf.add(i, 1)
return (i, loss, obj)
def rotate_and_project(i, obj):
# coord_old = read_origin_coords('arrsize_64_64_64_ntheta_500', i)
# obj_rot = apply_rotation(obj, coord_old, 'arrsize_64_64_64_ntheta_500')
obj_rot = tf_rotate(obj, theta_ls_tensor[i], interpolation='BILINEAR')
exiting = multislice_propagate(obj_rot[:, :, :, 0],
obj_rot[:, :, :, 1],
energy_ev,
psize_cm,
free_prop_cm=1e-4,
pad=pad)
# exiting = tf.abs(exiting)
return exiting
def rotate_and_project(i, loss, obj):
loss += tf.reduce_mean(tf.squared_difference(
tf.reduce_sum(tf_rotate(obj, theta_ls_tensor[i], interpolation='BILINEAR'), 1)[:, :, 0], prj[i]))
i = tf.add(i, 1)
return (i, loss, obj)
def rotate_and_project(i, obj_delta, obj_beta):
obj_rot = tf_rotate(tf.stack([obj_delta, obj_beta], axis=-1),
this_theta_batch[i],
interpolation='BILINEAR')
probe_pos_batch_ls = split_tasks(probe_pos, n_dp_batch)
exiting_ls = []
# loss = tf.constant(0.0)
# pad if needed
pad_arr = np.array([[0, 0], [0, 0]])
if probe_pos[:, 0].min() - probe_size_half[0] < 0:
pad_len = probe_size_half[0] - probe_pos[:, 0].min()
obj_rot = tf.pad(obj_rot, ((pad_len, 0), (0, 0), (0, 0), (0, 0)),
mode='CONSTANT')
pad_arr[0, 0] = pad_len
if probe_pos[:, 0].max() + probe_size_half[0] > obj_size[0]:
pad_len = probe_pos[:, 0].max() + probe_size_half[0] - obj_size[0]
obj_rot = tf.pad(obj_rot, ((0, pad_len), (0, 0), (0, 0), (0, 0)),
mode='CONSTANT')
pad_arr[0, 1] = pad_len
if probe_pos[:, 1].min() - probe_size_half[1] < 0:
pad_len = probe_size_half[1] - probe_pos[:, 1].min()
obj_rot = tf.pad(obj_rot, ((0, 0), (pad_len, 0), (0, 0), (0, 0)),
mode='CONSTANT')
pad_arr[1, 0] = pad_len
if probe_pos[:, 1].max() + probe_size_half[1] > obj_size[1]:
pad_len = probe_pos[:, 1].max() + probe_size_half[0] - obj_size[1]
obj_rot = tf.pad(obj_rot, ((0, 0), (0, pad_len), (0, 0), (0, 0)),
mode='CONSTANT')
pad_arr[1, 1] = pad_len
ind = 0
for k, pos_batch in enumerate(probe_pos_batch_ls):
subobj_ls = []
for j, pos in enumerate(pos_batch):
pos = [int(x) for x in pos]
# ind = np.reshape([[x, y] for x in range(int(pos[0]) - probe_size_half[0], int(pos[0]) - probe_size_half[0] + probe_size[0])
# for y in range(int(pos[1]) - probe_size_half[1], int(pos[1]) - probe_size_half[1] + probe_size[1])],
# [probe_size[0], probe_size[1], 2])
# subobj = tf.gather_nd(obj_rot, ind)
pos[0] = pos[0] + pad_arr[0, 0]
pos[1] = pos[1] + pad_arr[1, 0]
subobj = obj_rot[pos[0] - probe_size_half[0]:pos[0] -
probe_size_half[0] + probe_size[0],
pos[1] - probe_size_half[1]:pos[1] -
probe_size_half[1] + probe_size[1], :, :]
subobj_ls.append(subobj)
subobj_ls = tf.stack(subobj_ls)
if forward_algorithm == 'fresnel':
exiting = multislice_propagate_batch(subobj_ls[:, :, :, :, 0],
subobj_ls[:, :, :, :, 1],
probe_real,
probe_imag,
energy_ev,
psize_cm * ds_level,
h=h,
free_prop_cm='inf',
obj_batch_shape=[
len(pos_batch),
*probe_size,
obj_size[-1]
])
elif forward_algorithm == 'fd':
exiting = multislice_propagate_fd(subobj_ls[:, :, :, :, 0],
subobj_ls[:, :, :, :, 1],
probe_real,
probe_imag,
energy_ev,
psize_cm * ds_level,
h=h,
free_prop_cm='inf',
obj_batch_shape=[
len(pos_batch),
*probe_size, obj_size[-1]
])
# loss += tf.reduce_mean(tf.squared_difference(tf.abs(exiting), tf.abs(this_prj_batch[i][ind:ind+len(pos_batch)]))) * n_pos
# ind += len(pos_batch)
exiting_ls.append(exiting)
exiting_ls = tf.concat(exiting_ls, 0)
if probe_circ_mask is not None:
exiting_ls = exiting_ls * probe_mask
loss = tf.reduce_mean(
tf.squared_difference(tf.abs(exiting_ls), tf.abs(
this_prj_batch[i]))) * n_pos
# loss = tf.reduce_mean(tf.abs(exiting_ls) ** 2 * poisson_multiplier - tf.abs(this_prj_batch[i]) ** 2 * poisson_multiplier * tf.log(tf.abs(exiting_ls) ** 2 * poisson_multiplier), name='loss')
return loss