def apply_distortion(self, coordinates):
if len(coordinates) // 2 == 1:
coords = np.reshape(np.asarray(coordinates), (1, 2))
else:
coords = np.asarray(coordinates)
config = get_config()
w = ii.image.width - 1
h = ii.image.height - 1
crop_w = 0
crop_h = 0
if "crop" in config:
crop_w = config["crop"]["left_top"]["y"]
crop_h = config["crop"]["left_top"]["x"]
# Compute standard radial distortion, on normalized coordinates [-1, 1],
# function of even powers of radii (distance pixel - center of distortion) and coefficients of the distortion
coords_norm = np.divide(np.multiply(coords + [crop_h, crop_w], 2), [h, w]) - [1., 1.]
radii = np.sqrt(np.power(coords_norm[:, 0] - ii.image.c_x, 2) + np.power(coords_norm[:, 1] - ii.image.c_y, 2))
L = np.multiply(np.power(radii, 2), self.k1) + \
np.multiply(np.power(radii, 4), self.k2) + \
np.multiply(np.power(radii, 6), self.k3) + 1.
transformed_coordinates_x = ii.image.c_x + np.multiply(coords_norm[:, 0] - ii.image.c_x, L)
transformed_coordinates_y = ii.image.c_y + np.multiply(coords_norm[:, 1] - ii.image.c_y, L)
transformed_coordinates = np.vstack((transformed_coordinates_x, transformed_coordinates_y))
transformed_coordinates = np.transpose(transformed_coordinates)
transformed_coordinates = np.divide(np.multiply(transformed_coordinates + [1., 1.],
[h, w]), 2) - [crop_h, crop_w]
return transformed_coordinates
def call(self, coords, **kwargs):
# x' = x_c + (1 + c_1*r^2 + c_2*r^4 + c_3*r^6)*(x - x_c)
config = get_config()
w = ii.image.width - 1
h = ii.image.height - 1
crop_w = 0
crop_h = 0
if "crop" in config:
crop_w = config["crop"]["left_top"]["y"]
crop_h = config["crop"]["left_top"]["x"]
coords_norm = tf.subtract(tf.divide(tf.multiply(tf.add(coords, tf.constant([crop_h, crop_w], dtype=tf.float32)),
2.), tf.constant([h, w], dtype=tf.float32)), [1., 1.])
radius = tf.sqrt(tf.add(tf.pow(tf.subtract(coords_norm[:, 0], ii.image.c_x), 2),
tf.pow(tf.subtract(coords_norm[:, 1], ii.image.c_y), 2)))
k1 = tf.multiply(self.k1[0], config["layer_normalization"]["radial_distortion"])
k2 = tf.multiply(self.k2[0], config["layer_normalization"]["radial_distortion_2"])
k3 = tf.multiply(self.k3[0], config["layer_normalization"]["radial_distortion_3"])
distortion = tf.add(tf.add(tf.add(tf.multiply(k1, tf.pow(radius, 2)),
tf.multiply(k2, tf.pow(radius, 4))),
tf.multiply(k3, tf.pow(radius, 6))),
1)
distortion = tf.reshape(tf.tile(distortion, [2]), [tf.shape(distortion)[0], 2])
idx = tf.add(tf.constant([ii.image.c_x, ii.image.c_y], dtype=tf.float32),
tf.multiply(distortion,
tf.subtract(coords_norm, tf.constant([ii.image.c_x, ii.image.c_y], dtype=tf.float32))))
coords_transformed = tf.subtract(tf.divide(tf.multiply(tf.add(idx, [1., 1.]),
tf.constant([h, w], dtype=tf.float32)), 2.),
tf.constant([crop_h, crop_w], dtype=tf.float32))
return coords_transformed
def call(self, coords, **kwargs):
# [x',y'] = [x + c_x, y + c_y]
config = get_config()
idx = tf.cast(coords, tf.float32)
idx = tf.add(
idx, tf.multiply(self.shift,
config["layer_normalization"]["shift"]))
return idx
def __run(inputs, outputs, moving):
with open("input/config.yaml", "rt", encoding='utf-8') as config_file:
config = yaml.load(config_file, Loader=yaml.Loader)
init_config(config)
config = get_config()
ig, extrapolation, model, bias_history = \
__information_gain(inputs,
outputs,
visible=moving,
optimizer=build_optimizer(config["train"]["optimizer"], config["train"]["batch_size"]),
refiner=build_refining_optimizer(config["train"]["refiner"]),
train_config=config["train"],
layers=get_or_default("layers", 1),
batch_size=config["train"]["batch_size"],
stages=config["train"]["stages"]
)
# print model summary to stdout
model.summary()
layer_dict = dict([(layer.name, layer) for layer in model.layers])
bias = []
bias.append(layer_dict['ShiftLayer'].get_weights())
Verbose.print(
"Shift: " +
colored(str(
(bias[-1][0]) * config["layer_normalization"]["shift"]), "green"),
Verbose.always)
bias.append(layer_dict['RotationLayer'].get_weights())
Verbose.print(
"Rotation: " + colored(
str(bias[-1][0] * config["layer_normalization"]["rotation"] * 180 /
np.pi), "green"), Verbose.always)
bias.append(layer_dict['ScaleLayer'].get_weights())
Verbose.print(
"Scale: " +
colored(str(bias[-1][0] * config["layer_normalization"]["scale"] + 1),
"green"), Verbose.always)
# bias = layer_dict['ShearLayer'].get_weights()
# Verbose.print("Shear: " + colored(str(bias[0]*0.1), "green"), Verbose.always)
plt.figure(figsize=(15, 8))
ax = plt.subplot(1, 3, 1)
ax.imshow(extrapolation[:-1], cmap="gray")
ax.set_title("Predicted")
ax = plt.subplot(1, 3, 2)
ax.imshow(outputs[:-1], cmap="gray")
ax.set_title("Expected output")
ax = plt.subplot(1, 3, 3)
ax.imshow(ig[:-1], cmap="Reds")
ax.set_title("Difference between predicted and ground truth")
plt.show()
return bias, bias_history
def call(self, coords, **kwargs):
config = get_config()
idx = tf.cast(coords, tf.float32)
idx = tf.multiply(
idx,
tf.add(
tf.multiply(self.scale,
config["layer_normalization"]["scale"]),
scale_base))
return idx
def call(self, coords, **kwargs):
config = get_config()
rotation = tf.multiply(self.rotation,
config["layer_normalization"]["rotation"])
affine = tf.reshape([[tf.math.cos(rotation), -tf.math.sin(rotation)],
[tf.math.sin(rotation),
tf.math.cos(rotation)]], (2, 2))
# cos(a) -sin(a)
# sin(a) cos(a)
# rotation matrix, rotating by angle a
idx = tf.cast(coords, tf.float32)
idx = tf.einsum("ij,kj->ik", idx, affine)
return idx
def training_batch_selection_affine(train_set_size, input_dims):
"""
Selects data with respect to the maximal expected distortion so that we have coords
effectively in range [0, image_size-max_distortion] (after transform is applied)
:param train_set_size: batch size
:param input_dims: image dimensions
:return: indices to randomly selected pixels within "safety zone"
"""
all_data_indices = np.arange(input_dims[0] * input_dims[1])
all_data_indices = all_data_indices.reshape(input_dims)
max_misplacement = int(
np.floor(get_config()["expected_max_px_misplacement"]))
selection = all_data_indices[max_misplacement:-max_misplacement,
max_misplacement:-max_misplacement]
selection = selection.reshape(-1)
selection = np.random.permutation(selection)
selection = selection[:train_set_size]
return selection
def run(inputs,
outputs,
visible):
config = get_config()
ig, extrapolation, model, bias_history = \
__information_gain(inputs,
outputs,
visible=visible,
optimizer=build_optimizer(config["train"]["optimizer"], config["train"]["batch_size"]),
refiner=build_refining_optimizer(config["train"]["refiner"]),
train_config=config["train"],
layers=get_or_default("layers", 1),
batch_size=config["train"]["batch_size"],
stages=config["train"]["stages"]
)
# print model summary to stdout
model.summary()
layer_dict = dict([(layer.name, layer) for layer in model.layers])
bias = layer_dict['ShiftLayer'].get_weights()
Verbose.print("Shift: " + colored(str((bias[0])*config["layer_normalization"]["shift"]), "green"), Verbose.always)
bias = layer_dict['RotationLayer'].get_weights()
Verbose.print("Rotation: " + colored(str(bias[0]*config["layer_normalization"]["rotation"]*180/np.pi), "green"), Verbose.always)
bias = layer_dict['ScaleLayer'].get_weights()
Verbose.print("Scale: " + colored(str(bias[0]*config["layer_normalization"]["scale"] + 1), "green"), Verbose.always)
k1 = layer_dict['RDistortionLayer'].get_weights()[0][0] * config["layer_normalization"]["radial_distortion"]
k2 = layer_dict['RDistortionLayer'].get_weights()[1][0] * config["layer_normalization"]["radial_distortion_2"]
k3 = layer_dict['RDistortionLayer'].get_weights()[2][0] * config["layer_normalization"]["radial_distortion_3"]
exp_k1 = -k1
exp_k2 = 3*k1*k1 - k2
exp_k3 = -12*k1*k1*k1 + 8*k1*k2 - k3
Verbose.print("coefs computed: " + colored(str([k1, k2, k3]), "green"), Verbose.always)
# Verbose.print("coefs inverse: " + colored(str([exp_k1, exp_k2, exp_k3]), "green"), Verbose.always)
bias = [k1, k2, k3]
# bias = layer_dict['ShearLayer'].get_weights()
# Verbose.print("Shear: " + colored(str(bias[0]*0.1), "green"), Verbose.always)
return model, bias, bias_history
def training_batch_selection(train_set_size, input_img):
"""
Selects data with respect to the maximal expected distortion so that we have coords
effectively in range [0, image_size-max_distortion] (after transform is applied)
:param train_set_size: batch size
:param input_img: image
:return: indices to randomly selected pixels within "safety zone"
"""
input_dims = input_img.shape
all_data_indices = np.arange(input_dims[0] * input_dims[1])
all_data_indices = all_data_indices.reshape(input_dims[:-1])
conf = get_config()
inside = int(np.floor(conf["inside_part"]))
outside = int(np.floor(conf["outside_part"]))
cx = ii.image.c_x
cy = ii.image.c_y
# Find the position of the crop in the image and determine part least affected by radial distortion
center_x = (conf["crop"]["left_top"]['x'] +
conf["crop"]["size"]["height"] / 2) * 2 / ii.image.height - 1.
center_y = (conf["crop"]["left_top"]['y'] +
conf["crop"]["size"]["width"] / 2) * 2 / ii.image.width - 1.
left_right_center = max(min((center_y - cy) * 2, 1), -1)
top_bot_center = max(min((center_x - cx) * 2, 1), -1)
# DEBUG: set your own center
# left_right_center = 0
# top_bot_center = 0
from_x = int(round(inside * (1. - top_bot_center)))
to_x = min(-int(round(inside * (1. + top_bot_center))), -1)
from_y = int(round(inside * (1. - left_right_center)))
to_y = min(-int(round(inside * (1. + left_right_center))), -1)
# Exclude part that is minimally affected by radial distortion
selection_exclude = all_data_indices[from_x:to_x, from_y:to_y]
selection_exclude = selection_exclude.reshape(-1)
# Exclude outer border in order to avoid index out of bounds
selection_include = all_data_indices[outside:-outside, outside:-outside]
selection_include = selection_include.reshape(-1)
selection = [x for x in selection_include if x not in selection_exclude]
selection = np.random.permutation(selection)
# DEBUG: forcing larger training set
train_set_size = int(train_set_size * 2)
selection = selection[:train_set_size]
# DEBUG: display image region for selection
image = input_img.reshape(-1)
image[:] = 0
image[selection_include] = input_img.reshape(-1)[selection_include]
image[selection_exclude] = 0
image = image.reshape(input_dims[0], input_dims[1])
Verbose.imshow(image, Verbose.debug)
return selection
def __train_networks(inputs,
outputs,
reg_layer_data,
optimizer,
refiner,
config,
layers=None,
train_set_size=50000,
batch_size=256,
stages={
"type": "polish",
"epochs": 50
}):
"""
This method builds ANN with all layers - some layers for registration other layers for information gain computation
and processes the training.
:param inputs:
input coordinates (will be randomized)
:param outputs:
output image pixel intensity values
:param reg_layer_data:
first layer data - transformation of coords to pixel intensity value
:param layers:
not used now - this will define IG layers in the future
:param train_set_size:
number of pixels used for training, default is 50k
:param batch_size:
number of pixels computed in one batch
:param epochs:
number of learning epochs
:return:
trained model and training history
"""
Verbose.print('Selecting ' + str(train_set_size) + ' samples randomly for use by algorithm.')
selection = training_batch_selection(train_set_size, reg_layer_data)
indexes = inputs[selection, :]
# define model
print('Adding input layer, width =', indexes.shape[1])
# TODO: revisit shear layer
input_layer = tf.keras.layers.Input(shape=(indexes.shape[1],),
dtype=tf.float32, name='InputLayer')
shift_layer = ShiftLayer(name='ShiftLayer')(input_layer)
scale_layer = ScaleLayer(name='ScaleLayer')(shift_layer)
rotation_layer = RotationLayer(name='RotationLayer')(scale_layer)
radial_distortion_layer = RDCompleteLayer(name='RDistortionLayer')(rotation_layer)
# radial_distortion_layer = RDistortionLayer(name='RDistortionLayer')(rotation_layer)
# radial_distortion_layer_2 = RDistortionLayer2(name='RDistortionLayer2')(radial_distortion_layer)
# radial_distortion_layer_3 = RDistortionLayer3(name='RDistortionLayer3')(radial_distortion_layer_2)
layer = Idx2PixelLayer(visible=reg_layer_data, name='Idx2PixelLayer')(radial_distortion_layer)
# TODO: Add InformationGain layers here when necessary
print('Adding ReLU output layer, width =', outputs.shape[1])
output_layer = tf.keras.layers.ReLU(max_value=1, name='Output', trainable=False)(layer)
model = tf.keras.models.Model(inputs=input_layer, outputs=output_layer)
# compile model
start_time = time()
model.compile(loss='mean_squared_error',
optimizer=optimizer,
metrics=['mean_squared_error']
)
# Set initial transformation as identity
elapsed_time = time() - start_time
print("Compiling model took {:.4f}'s.".format(elapsed_time))
# train model
start_time = time()
tf.compat.v1.keras.backend.get_session().run(tf.compat.v1.global_variables_initializer())
model.layers[1].set_weights([np.array([0, 0])]) # shift
model.layers[2].set_weights([np.array([0, 0])]) # scale
model.layers[3].set_weights([np.array([0])]) # rotation
# model.load_weights(MODEL_FILE)
# TODO: better names for stages
config = get_config()
for stage in stages:
if stage['type'] == 'blur':
output = blur_preprocessing(outputs, reg_layer_data.shape, stage['params'])
__set_train_registration(model, True)
model.compile(loss='mean_squared_error', optimizer=optimizer, metrics=['mean_squared_error'])
elif stage['type'] == 'refine':
output = outputs
__set_train_registration(model, True, target="shift")
model.compile(loss='mean_squared_error', optimizer=refiner, metrics=['mean_squared_error'])
elif stage['type'] == 'polish':
output = outputs
__set_train_registration(model, True)
model.compile(loss='mean_squared_error', optimizer=optimizer, metrics=['mean_squared_error'])
else:
stage_params = stage['type'].split('_')
if stage_params[0] == 'adam':
opt_type = 'optimizer'
opt = build_optimizer(config["train"][opt_type], config["train"]["batch_size"])
elif stage_params[0] == 'sgd':
opt_type = 'refiner'
opt = build_refining_optimizer(config["train"][opt_type])
if stage_params[1] == 'rd':
targ = "rd"
output = outputs
__set_train_registration(model, True, target=targ)
model.compile(loss='mean_squared_error', optimizer=opt, metrics=['mean_squared_error'])
reset_visible(output)
output = output[selection, :]
shift_metric = ShiftMetrics()
scale_metric = ScaleMetrics()
rotation_metric = RotationMetrics()
# distortion_metric = DistortionMetrics()
distortion_metric = RDMetrics()
mcp_save = ModelCheckpoint(MODEL_FILE,
save_best_only=True, monitor='val_loss', mode='min')
checkpoint_filepath = MODEL_FILE
model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
filepath=checkpoint_filepath,
save_weights_only=True,
monitor='val_loss',
mode='min',
save_best_only=True)
# lr_reduction = ReduceLROnPlateau(monitor='val_loss', factor=0.9, patience=100, verbose=0, mode='auto',
# min_delta=0.0001, cooldown=0, min_lr=0)
log_dir = "logs/fit/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1)
callbacks = [shift_metric, scale_metric, rotation_metric, distortion_metric, tensorboard_callback,
model_checkpoint_callback]
history = model.fit(indexes,
output,
epochs=stage['epochs'],
validation_split=0.2,
verbose=1,
callbacks=[shift_metric, scale_metric, rotation_metric, distortion_metric,
model_checkpoint_callback
],
batch_size=batch_size
)
coef_1 = [x[0][0]*config["layer_normalization"]["radial_distortion"] for x in distortion_metric.bias_history]
coef_2 = [x[1][0]*config["layer_normalization"]["radial_distortion_2"] for x in distortion_metric.bias_history]
coef_3 = [x[2][0]*config["layer_normalization"]["radial_distortion_3"] for x in distortion_metric.bias_history]
plt.plot(coef_1, label="1")
plt.plot(coef_2, label="2")
plt.plot(coef_3, label="3")
model.load_weights(checkpoint_filepath)
plt.title("Transformation %f %f %f " % (coef_1[-1], coef_2[-1], coef_3[-1]))
plt.legend()
plt.show()
elapsed_time = time() - start_time
num_epochs = len(history.history['loss'])
print("{:.4f}'s time per epoch. Epochs:".format(elapsed_time / num_epochs), num_epochs, sep="")
print("Total time {:.4f}'s".format(elapsed_time))
# calculate gain and save best model so far
model.load_weights(checkpoint_filepath)
gain = abs(outputs[selection, :] - model.predict(indexes, batch_size=batch_size)) / (outputs.shape[0] *
outputs.shape[1])
information_gain_max = gain.flatten().max()
print('Gain: {:1.4e}'.format(information_gain_max))
return model, history, shift_metric.bias_history
