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