def _build(self, batch):
(img, ) = batch
img = gaussian_filter2d(img, filter_shape=[6, 6])
img_before_autoencoder = (img - tf.reduce_min(img)) / (
tf.reduce_max(img) - tf.reduce_min(img))
tf.summary.image(f'img_before_autoencoder',
img_before_autoencoder,
step=self.step)
encoded_img = self.encoder(img)
print(encoded_img.shape)
mn = self.mn(encoded_img)
std = self.std(encoded_img)
epsilon = tf.random.normal(
tf.stack([tf.shape(encoded_img)[0], self.n_latent]))
z = mn + tf.multiply(epsilon, tf.exp(tf.multiply(0.5, std)))
decoded_img = self.decoder(z)
print(decoded_img.shape)
img_after_autoencoder = (decoded_img - tf.reduce_min(decoded_img)) / (
tf.reduce_max(decoded_img) - tf.reduce_min(decoded_img))
tf.summary.image(f'img_after_autoencoder',
img_after_autoencoder,
step=self.step)
return mn, std, z, decoded_img
def build_dataset(tfrecords_dirs, batch_size, train_test_dir='train'):
"""
Build data set from a directory of tfrecords. With graph batching
Args:
tfrecords_dirs: str, path to *.tfrecords
batch_size: size of the batch
train_test_dir: whether to load 'train' or 'test' data
Returns: Dataset obj.
"""
tfrecords = []
for tfrecords_dir in tfrecords_dirs:
tfrecords += glob.glob(
os.path.join(tfrecords_dir, train_test_dir, '*.tfrecords'))
random.shuffle(tfrecords)
print(f'Number of {train_test_dir} tfrecord files : {len(tfrecords)}')
dataset = tf.data.TFRecordDataset(tfrecords).map(
partial(decode_examples,
voxels_shape=(64, 64, 64, 10),
image_shape=(256, 256, 1))) # (voxels, image, idx)
dataset = dataset.map(lambda voxels, img, cluster_idx, projection_idx,
vprime: (voxels, gaussian_filter2d(img)))
dataset = dataset.shuffle(buffer_size=52).batch(batch_size=batch_size)
return dataset
def build_dataset(tfrecords, batch_size):
"""
Build data set from a directory of tfrecords. With graph batching
Args:
data_dir: str, path to *.tfrecords
Returns: Dataset obj.
"""
dataset = tf.data.TFRecordDataset(tfrecords).map(
partial(decode_examples,
node_shape=(10, ),
edge_shape=(2, ),
image_shape=(1024, 1024, 1))) # (graph, image, idx)
dataset = dataset.map(
lambda graph_data_dict, img, cluster_idx, projection_idx, vprime: tf.
concat([
double_downsample(img),
gaussian_filter2d(double_downsample(img), filter_shape=[6, 6])
],
axis=-1))
dataset = dataset.shuffle(buffer_size=50).batch(batch_size=batch_size)
# dataset = batch_dataset_set_graph_tuples(all_graphs_same_size=True, dataset=dataset, batch_size=batch_size)
return dataset
def _build(self, batch):
(img, ) = batch
img = gaussian_filter2d(img, filter_shape=[6, 6])
img_before_autoencoder = (img - tf.reduce_min(img)) / (
tf.reduce_max(img) - tf.reduce_min(img))
tf.summary.image(f'img_before_autoencoder',
img_before_autoencoder,
step=self.step)
encoded_img = self.encoder(img)
vq_dict = self.VQVAE(encoded_img, is_training=True)
quantized_img = (vq_dict['quantize'] - tf.reduce_min(
vq_dict['quantize'])) / (tf.reduce_max(vq_dict['quantize']) -
tf.reduce_min(vq_dict['quantize']))
# quantized_img has shape [32, 32, 32, 64] (batch, x, y, channels),
# for the tf.summary we only show one of the 64 channels.
tf.summary.image(
f'quantized_img',
quantized_img[:, :, :, np.random.randint(low=0, high=64)][:, :, :,
None],
step=self.step)
decoded_img = self.decoder(vq_dict['quantize'])
img_after_autoencoder = (decoded_img - tf.reduce_min(decoded_img)) / (
tf.reduce_max(decoded_img) - tf.reduce_min(decoded_img))
tf.summary.image(f'img_after_autoencoder',
img_after_autoencoder,
step=self.step)
return vq_dict['loss'], decoded_img
def test_gaussian_filter2d_different_sigma():
image = np.arange(40 * 40).reshape(40, 40).astype(np.float32)
sigma = [1.0, 2.0]
test_utils.assert_allclose_according_to_type(
gaussian_filter2d(image, [9, 17], sigma).numpy(),
gaussian_filter(image, sigma, mode="mirror"),
)
def test_gaussian_filter2d_batch(image_shape):
test_image_tf = tf.random.uniform([1, 40, 40, 1],
minval=0,
maxval=255,
dtype=tf.float32)
gb = gaussian_filter2d(test_image_tf, 5, 1, padding="SYMMETRIC")
gb = gb.numpy()
gb1 = np.resize(gb, (40, 40))
test_image_np = test_image_tf.numpy()
test_image_np = np.resize(test_image_np, [40, 40])
gb2 = gaussian(test_image_np, 1, mode="reflect")
np.testing.assert_allclose(gb2, gb1, 0.06)
def test_gaussian_filter2d_constant():
test_image_tf = tf.random.uniform([1, 40, 40, 1],
minval=0,
maxval=255,
dtype=tf.float64)
gb = gaussian_filter2d(test_image_tf, 5, 1, padding="CONSTANT")
gb = gb.numpy()
gb1 = np.resize(gb, (40, 40))
test_image_np = test_image_tf.numpy()
test_image_np = np.resize(test_image_np, [40, 40])
gb2 = gaussian(test_image_np, 1, mode="constant")
np.testing.assert_allclose(gb2, gb1, 0.06)
def test_gaussian_filter2d_channels():
test_image_tf = tf.constant(
[[
[
[0.0, 0.0, 0.0],
[2.0, 2.0, 0.0],
[4.0, 4.0, 0.0],
[6.0, 6.0, 0.0],
[8.0, 8.0, 0.0],
],
[
[10.0, 10.0, 0.0],
[12.0, 12.0, 0.0],
[14.0, 14.0, 0.0],
[16.0, 16.0, 0.0],
[18.0, 18.0, 0.0],
],
[
[20.0, 20.0, 0.0],
[22.0, 22.0, 0.0],
[24.0, 24.0, 0.0],
[26.0, 26.0, 0.0],
[28.0, 28.0, 0.0],
],
[
[30.0, 30.0, 0.0],
[32.0, 32.0, 0.0],
[34.0, 34.0, 0.0],
[36.0, 36.0, 0.0],
[38.0, 38.0, 0.0],
],
[
[40.0, 40.0, 0.0],
[42.0, 42.0, 0.0],
[44.0, 44.0, 0.0],
[46.0, 46.0, 0.0],
[48.0, 48.0, 0.0],
],
]],
dtype=tf.float32,
)
gb = gaussian_filter2d(test_image_tf,
5,
1,
padding="SYMMETRIC",
name="gaussian")
gb = gb.numpy()
gb1 = np.resize(gb, (5, 5, 3))
test_image_np = test_image_tf.numpy()
test_image_np = np.resize(test_image_np, [5, 5, 3])
gb2 = gaussian(test_image_np, sigma=1, mode="reflect", multichannel=True)
np.testing.assert_allclose(gb2, gb1, 0.06)
def loss(model_outputs, batch):
(img, ) = batch
vq_loss, decoded_img = model_outputs
print('im shape', img.shape)
print('dec im shape', decoded_img.shape)
# reconstruction_loss = tf.reduce_mean(tf.reduce_sum(
# keras.losses.binary_crossentropy(img, decoded_img), axis=(1, 2)
# ))
reconstruction_loss = tf.reduce_mean(
(gaussian_filter2d(img, filter_shape=[6, 6]) -
decoded_img[:, 12:-12, 12:-12, :])**2)
total_loss = reconstruction_loss + vq_loss
return total_loss
def loss(model_outputs, batch):
(img, ) = batch
mn, std, z, decoded_img = model_outputs
# reconstruction_loss = tf.reduce_mean(tf.reduce_sum(
# keras.losses.binary_crossentropy(img, decoded_img), axis=(1, 2)
# ))
reconstruction_loss = tf.reduce_mean(
(gaussian_filter2d(img, filter_shape=[6, 6]) -
decoded_img[:, 12:-12, 12:-12, :])**2)
kl_loss = -0.5 * (1 + std - tf.square(mn) - tf.exp(std))
kl_loss = tf.reduce_mean(tf.reduce_sum(kl_loss, axis=1))
total_loss = reconstruction_loss + kl_loss
print(f"recon_loss = {reconstruction_loss}")
print(f"kl_loss = {kl_loss}")
return total_loss
def _build(self, batch):
(img, ) = batch
img = gaussian_filter2d(img, filter_shape=[6, 6])
img_before_autoencoder = (img - tf.reduce_min(img)) / (
tf.reduce_max(img) - tf.reduce_min(img))
tf.summary.image(f'img_before_autoencoder',
img_before_autoencoder,
step=self.step)
encoded_img = self.encoder(img)
print(encoded_img.shape)
decoded_img = self.decoder(encoded_img)
img_after_autoencoder = (decoded_img - tf.reduce_min(decoded_img)) / (
tf.reduce_max(decoded_img) - tf.reduce_min(decoded_img))
tf.summary.image(f'img_after_autoencoder',
img_after_autoencoder,
step=self.step)
return decoded_img
def test_gaussian_filter2d(shape, padding, dtype):
modes = {
"SYMMETRIC": "reflect",
"CONSTANT": "constant",
"REFLECT": "mirror",
}
image = np.arange(np.prod(shape)).reshape(*shape).astype(dtype)
ndims = len(shape)
sigma = [1.0, 1.0]
if ndims == 3:
sigma = [1.0, 1.0, 0.0]
elif ndims == 4:
sigma = [0.0, 1.0, 1.0, 0.0]
test_utils.assert_allclose_according_to_type(
gaussian_filter2d(image, 9, 1, padding=padding).numpy(),
gaussian_filter(image, sigma, mode=modes[padding]),
)
def build_dataset(tfrecords_dirs, batch_size, type='train'):
"""
Build data set from a directory of tfrecords. With graph batching
Args:
data_dir: str, path to *.tfrecords
Returns: Dataset obj.
"""
tfrecords = []
for tfrecords_dir in tfrecords_dirs:
tfrecords += glob.glob(os.path.join(tfrecords_dir, type,
'*.tfrecords'))
random.shuffle(tfrecords)
print(f'Number of {type} tfrecord files : {len(tfrecords)}')
dataset = tf.data.TFRecordDataset(tfrecords).map(
partial(decode_examples,
node_shape=(10, ),
edge_shape=(2, ),
image_shape=(256, 256, 1))) # (graph, image, idx)
dataset = dataset.map(
lambda graph_data_dict, img, cluster_idx, projection_idx, vprime:
(GraphsTuple(**graph_data_dict).replace(nodes=tf.concat([
GraphsTuple(**graph_data_dict).nodes[:, :3],
GraphsTuple(**graph_data_dict).nodes[:, 6:8]
],
axis=-1)),
gaussian_filter2d(img))).shuffle(buffer_size=52).batch(
batch_size=batch_size)
return dataset
def loss(model_outputs, batch):
(img, ) = batch
decoded_img = model_outputs
return tf.reduce_mean((gaussian_filter2d(img, filter_shape=[6, 6]) -
decoded_img[:, 12:-12, 12:-12, :])**2)
def _build(self, batch, *args, **kwargs):
(graph, img, c) = batch
del c
print(graph.nodes.shape)
# The encoded cluster graph has globals which can be compared against the encoded image graph
encoded_graph = self.epd_graph(graph, self._core_steps)
# # Add an extra dimension to the image (tf.summary expects a Tensor of rank 4)
# img = img[None, ...]
print("IMG SHAPE:", img.shape)
print("IMG MIN MAX:", tf.math.reduce_min(img), tf.math.reduce_max(img))
img_before_cnn = (img - tf.reduce_min(img)) / \
(tf.reduce_max(img) - tf.reduce_min(img))
tf.summary.image(f'img_before_cnn', img_before_cnn, step=self.step)
# Smooth the image and use the encoder from the autoencoder to reduce the dimensionality of the image
# The autoencoder was trained on images that were smoothed in the same way
img = gaussian_filter2d(img, filter_shape=[6, 6])
img = self.auto_encoder.encoder(img)
# Prevent the autoencoder from learning
try:
for variable in self.auto_encoder.encoder.trainable_variables:
variable._trainable = False
for variable in self.auto_encoder.decoder.trainable_variables:
variable._trainable = False
except:
pass
print("IMG SHAPE AFTER CNN:", img.shape)
print("IMG MIN MAX AFTER CNN:", tf.math.reduce_min(img),
tf.math.reduce_max(img))
img_after_cnn = (img - tf.reduce_min(img)) / \
(tf.reduce_max(img) - tf.reduce_min(img))
tf.summary.image(
f'quantized_img',
img_after_cnn[:, :, :, np.random.randint(low=0, high=64)][:, :, :,
None],
step=self.step)
decoded_img = self.auto_encoder.decoder(img)
decoded_img = (decoded_img - tf.reduce_min(decoded_img)) / \
(tf.reduce_max(decoded_img) - tf.reduce_min(decoded_img))
tf.summary.image(f'decoded_img', decoded_img, step=self.step)
# Reshape the encoded image so it can be used for the nodes
img_nodes = tf.reshape(img, (-1, self.image_feature_size))
print(img_nodes.shape)
# Create a graph that has a node for every encoded pixel. The features of each node
# are the channels of the corresponding pixel. Then connect each node with every other
# node.
img_graph = GraphsTuple(nodes=img_nodes,
edges=None,
globals=None,
receivers=None,
senders=None,
n_node=tf.shape(img_nodes)[0:1],
n_edge=tf.constant([0]))
connected_graph = fully_connect_graph_dynamic(img_graph)
# The encoded image graph has globals which can be compared against the encoded cluster graph
encoded_img = self.epd_image(connected_graph, 1)
# Compare the globals from the encoded cluster graph and encoded image graph
# to estimate the similarity between the input graph and input image
print(encoded_img.globals.shape)
print(encoded_graph.globals.shape)
distance = self.compare(
tf.concat([encoded_graph.globals, encoded_img.globals], axis=1)
) + self.compare(
tf.concat([encoded_img.globals, encoded_graph.globals], axis=1))
return distance