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