Finished

Author: Aftaab99

.gitignore

@@ -1,4 +1,3 @@
 .idea/
-*.png
 Data/
-Checkpoints/
+__pycache__/

Checkpoints/weights-epoch-6loss-0.005/checkpoint

@@ -0,0 +1,6 @@
+model_checkpoint_path: "weights-epoch-6loss-0.005.ckpt"
+all_model_checkpoint_paths: "../weights-epoch-2loss-0.009/weights-epoch-2loss-0.009.ckpt"
+all_model_checkpoint_paths: "../weights-epoch-3loss-0.007/weights-epoch-3loss-0.007.ckpt"
+all_model_checkpoint_paths: "../weights-epoch-4loss-0.006/weights-epoch-4loss-0.006.ckpt"
+all_model_checkpoint_paths: "../weights-epoch-5loss-0.005/weights-epoch-5loss-0.005.ckpt"
+all_model_checkpoint_paths: "weights-epoch-6loss-0.005.ckpt"

Checkpoints/weights-epoch-6loss-0.005/weights-epoch-6loss-0.005.ckpt.data-00000-of-00001

@@ -1,2 +1,45 @@
-### Denoising Autoencoder
-Implementation of a denoising autoencoder trained on the RENOIR dataset. 
+## Denoising Autoencoder
+Implementation of a denoising autoencoder trained on the RENOIR dataset(MI 3 images). 
+
+## Setting up locally
+
+	pip install -r requirements.txt
+
+
+## Dataset
+50x50px patches were taken from the reference and noisy images in the dataset. I've serialised these into TFRecords, which can be downloaded using,
+
+	python download_data.py
+	
+This will download the train and validation records required for training.
+
+## Training and inference
+1. For training you can run,
+		
+		python train.py -e <num_of_epochs> -c <checkpoint_after> -v <validation_enabled, 1 or 0>
+	Example:
+
+		python train.py -e 50 -c 5 -v 1
+Default values are training for 10 epochs, checkpointing every 1 epoch with validation enabled
+
+2. For inference,
+
+		python predict.py -i <input_file> -o <output_file>
+
+
+## Results
+I've trained the model for only 6 epochs(which is a very very small fraction of what a lot of papers recommend), so the results aren't particularly good. 
+
+1. Reference:
+![Reference Image](https://github.com/Aftaab99/DenoisingAutoencoder/blob/master/images/reference.bmp  "Reference Image")
+
+2. Noisy
+![Noisy Image](https://github.com/Aftaab99/DenoisingAutoencoder/blob/master/images/noisy.png  "Noisy Image")
+
+3. Denoised
+
+![Denoised Image](https://github.com/Aftaab99/DenoisingAutoencoder/blob/master/images/denoised.png  "Denoised Image")
+
+### References
+1. J. Anaya, A. Barbu. RENOIR - A Dataset for Real Low-Light Image Noise Reduction.([arxiv](https://arxiv.org/abs/1409.8230))
+

Checkpoints/weights-epoch-6loss-0.005/weights-epoch-6loss-0.005.ckpt.index

@@ -3,27 +3,52 @@
 from PIL import Image
 from sklearn.model_selection import train_test_split
 import tensorflow as tf
+import numpy as np
+from random import shuffle
 
 
 def _bytes_feature(value):
 	return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
 
 
-def load_image(addr):
-	img = Image.open(addr).resize([1240, 1240])
-	return img
+def extract_patches(ref, noisy):
+	patch_list = []
+	for x in range(50, 3000, 50):
+		for y in range(50, 3000, 50):
+			patch_ref = ref[x - 50:x, y - 50:y, :]
+			patch_noisy = noisy[x - 50:x, y - 50:y, :]
+			if patch_ref.shape[0] != 50 or patch_ref.shape[1] != 50:
+				continue
+			patch_list.append({'ref': patch_ref, 'noisy': patch_noisy})
+	return patch_list
+
+
+def get_patches(addr):
+	ref_img = Image.open(addr['reference']).convert('RGB')
+	noisy_image = Image.open(addr['noisy']).convert('RGB')
+	ref_img_t = np.array(ref_img)
+	noisy_image_t = np.array(noisy_image)
+	patch_list = extract_patches(ref_img_t, noisy_image_t)
+
+	return patch_list
 
 
 def create_data_record(out_filename, addrs):
 	# open the TFRecords file
 
 	writer = tf.python_io.TFRecordWriter(out_filename)
-	for addr, i in zip(addrs, range(1, len(addrs)+1)):
 
-		ref = load_image(addr['reference'])
-		noisy = load_image(addr['noisy'])
+	patch_list = []
 
+	for addr, i in zip(addrs, range(1, len(addrs) + 1)):
+		patch_list = patch_list + get_patches(addr)
+
+	shuffle(patch_list)
+
+	for item, i in zip(patch_list, range(1, len(patch_list) + 1)):
 		# Create a feature
+		ref = Image.fromarray(item['ref'])
+		noisy = Image.fromarray(item['noisy'])
 		feature = {
 			'reference': _bytes_feature(ref.tobytes()),
 			'noisy': _bytes_feature(noisy.tobytes())
@@ -33,14 +58,13 @@ def create_data_record(out_filename, addrs):
 
 		# Serialize to string and write on the file
 		writer.write(example.SerializeToString())
-		print('Image {} wrote to record'.format(i))
+		print('Patch {} wrote to record'.format(i))
 	writer.close()
 	sys.stdout.flush()
 
 
 base_path = '/home/aftaab/Datasets/Mi3_Aligned'
 
-
 addrs = []
 
 for directory in os.listdir(base_path):
@@ -55,4 +79,4 @@ def create_data_record(out_filename, addrs):
 train_addrs, test_addrs = train_test_split(addrs, test_size=0.2)
 
 create_data_record('Data/train.tfrecords', train_addrs)
-create_data_record('Data/test.tfrecords', test_addrs)
+create_data_record('Data/val.tfrecords', test_addrs)

Checkpoints/weights-epoch-6loss-0.005/weights-epoch-6loss-0.005.ckpt.meta

@@ -0,0 +1,50 @@
+import requests
+import os
+
+
+def download_file_from_google_drive(id, destination):
+	URL = "https://drive.google.com/uc?export=download"
+
+	session = requests.Session()
+
+	response = session.get(URL, params={'id': id}, stream=True)
+	token = get_confirm_token(response)
+
+	if token:
+		params = {'id': id, 'confirm': token}
+		response = session.get(URL, params=params, stream=True)
+
+	save_response_content(response, destination)
+
+
+def get_confirm_token(response):
+	for key, value in response.cookies.items():
+		if key.startswith('download_warning'):
+			return value
+
+	return None
+
+
+def save_response_content(response, destination):
+	CHUNK_SIZE = 32768
+
+	with open(destination, "wb") as f:
+		for chunk in response.iter_content(CHUNK_SIZE):
+			if chunk:  # filter out keep-alive new chunks
+				f.write(chunk)
+
+
+if __name__ == "__main__":
+
+	train_file_id = '12ctvUrf-Jivr0P9kHThW2GxIZzQf-R3I'
+	train_file_dest = 'Data/train.tfrecords'
+	val_file_id = '	1YovsQgVVNeUyDGyhD0XpAf83HxThkKAZ'
+	val_file_dest = 'Data/val.tfrecords'
+
+	if not os.path.exists('./Data'):
+		os.mkdir('./Data')
+		print('Downloading train records...')
+		download_file_from_google_drive(train_file_id, train_file_dest)
+
+		print('Downloading validation records...')
+		download_file_from_google_drive(val_file_id, val_file_dest)

README.md

@@ -1,59 +1,76 @@
 import tensorflow as tf
-from tensorflow.layers import conv2d, dropout, max_pooling2d, conv2d_transpose
+from tensorflow.layers import conv2d, max_pooling2d, conv2d_transpose
 from PIL import Image
 import numpy as np
 
 
 class DenoisingAutoEncoder:
 
-	def __init__(self, input_shape: tuple, optimizer, is_training: bool):
+	def __init__(self, input_shape: tuple, batch_input_shape: tuple, optimizer, is_training: bool):
 		self.sess = tf.Session()
 		self.input_shape = input_shape
-		self.input_image = tf.placeholder(tf.float32, shape=(None, 1240, 1240, 3), name="input_image")
-		self.target_image = tf.placeholder(tf.float32, shape=(None, 1240, 1240, 3), name="target_image")
+		self.input_image = tf.placeholder(tf.float32, shape=batch_input_shape, name="input_image")
+		self.target_image = tf.placeholder(tf.float32, shape=batch_input_shape, name="target_image")
+		self.training = is_training
 
 		with tf.name_scope('Encoder'):
-			self.conv1 = tf.nn.leaky_relu(conv2d(self.input_image, 16, (7, 7), padding='same', use_bias=False))
-			self.pool1 = max_pooling2d(self.conv1, (4, 4), (4, 4))
-			self.dropout1 = dropout(self.pool1, 0.2, training=is_training)
-			self.conv2 = tf.nn.leaky_relu(conv2d(self.dropout1, 20, (5, 5), padding='same', use_bias=False))
-			self.pool2 = max_pooling2d(self.conv2, (2, 2), (2, 2))
-			self.dropout2 = dropout(self.pool2, 0.3, training=is_training)
-			self.conv3 = tf.nn.leaky_relu(conv2d(self.dropout2, 32, (5, 5), padding='same', use_bias=False))
+			self.conv1 = tf.nn.leaky_relu(conv2d(self.input_image, 16, (5, 5), padding='same'))
+			self.pool1 = max_pooling2d(self.conv1, (2, 2), (2, 2))
+			self.conv2 = tf.nn.leaky_relu(conv2d(self.pool1, 32, (3, 3), padding='same'))
+			self.pool2 = max_pooling2d(self.conv2, (5, 5), (5, 5))
+			self.conv3 = tf.nn.leaky_relu(conv2d(self.pool2, 64, (3, 3), padding='same'))
 			self.pool3 = max_pooling2d(self.conv3, (5, 5), (5, 5))
-			self.dropout3 = dropout(self.pool3, 0.3, training=is_training)
-			self.latent_repr = tf.nn.leaky_relu(conv2d(self.dropout3, 128, (3, 3), padding='same', use_bias=False))
-			# self.latent_repr = max_pooling2d(self.conv4, (5, 5), (5, 5))
+			self.latent_repr = tf.nn.leaky_relu(conv2d(self.pool3, 256, (3, 3), padding='same'))
 
 		with tf.name_scope('Decoder'):
-			self.upsampling1 = tf.image.resize_images(self.latent_repr, (31, 31),
+			self.upsampling1 = tf.image.resize_images(self.latent_repr, (5, 5),
 													  tf.image.ResizeMethod.BICUBIC)
 			self.conv5 = tf.nn.leaky_relu(
-				conv2d_transpose(self.upsampling1, 32, (3, 3), padding='same', use_bias=False))
-			self.dropout4 = dropout(self.conv5, 0.3, training=is_training)
-			self.upsampling2 = tf.image.resize_images(self.dropout4, (155, 155), tf.image.ResizeMethod.BICUBIC)
+				conv2d_transpose(self.upsampling1, 64, (3, 3), padding='same'))
+			self.upsampling2 = tf.image.resize_images(self.conv5, (25, 25), tf.image.ResizeMethod.BICUBIC)
 			self.conv6 = tf.nn.leaky_relu(
-				conv2d_transpose(self.upsampling2, 16, (5, 5), padding='same', use_bias=False))
-			self.upsampling3 = tf.image.resize_images(self.conv6, (310, 310), tf.image.ResizeMethod.BICUBIC)
-			self.conv7 = tf.nn.leaky_relu(conv2d_transpose(self.upsampling3, 3, (5, 5), padding='same', use_bias=False))
-			self.upsampling4 = tf.image.resize_images(self.conv7, (1240, 1240), tf.image.ResizeMethod.BICUBIC)
-			self.conv8 = tf.nn.leaky_relu(conv2d_transpose(self.upsampling4, 3, (1, 1), padding='same', use_bias=True))
+				conv2d_transpose(self.upsampling2, 32, (5, 5), padding='same'))
+			self.upsampling3 = tf.image.resize_images(self.conv6, (50, 50), tf.image.ResizeMethod.BICUBIC)
+			self.conv7 = tf.nn.leaky_relu(conv2d_transpose(self.upsampling3, 3, (5, 5), padding='same'))
 
-		self.output_image = tf.nn.sigmoid(self.conv8)
+		self.output_image = tf.nn.sigmoid(self.conv7)
 		self.loss = tf.losses.mean_squared_error(self.target_image, self.output_image)
 		self.batch_loss = tf.reduce_mean(self.loss)
 
 		self.train_step = optimizer.minimize(self.batch_loss)
 		self.sess.run(tf.global_variables_initializer())
 		self.saver = tf.train.Saver()
+		self.__load_weights()
 
-	def train(self, epochs: int, ckpt_every: int):
+	def validate(self):
+		noisy_batch, target_batch = self.input_fn('Data/val.tfrecords', False, 1024)
+		val_loss = 0
+		n_batch = 0
+		while True:
+			try:
+				noisies, targets = self.sess.run([noisy_batch, target_batch])
+				n_batch += 1
+				noisies /= 255
+				targets /= 255
+
+				l = self.sess.run([self.batch_loss], feed_dict={self.input_image: noisies,
+																self.target_image: targets})
+				print(l)
+				val_loss += l[0]
+			except tf.errors.OutOfRangeError:
+				val_loss = val_loss / n_batch
+				return val_loss
+
+	def train(self, epochs: int, ckpt_every: int, validate: bool):
 		for e in range(1, epochs + 1):
-			noisy_batch, target_batch = self.input_fn('Data/train.tfrecords', True, 2)
+			noisy_batch, target_batch = self.input_fn('Data/train.tfrecords', True, 1024)
 			epoch_loss = self.train_epoch(noisy_batch, target_batch)
 			if e % ckpt_every == 0:
 				self.checkpoint(e, epoch_loss)
-			print('Epoch Loss = {}, epoch={}'.format(epoch_loss, e))
+			if validate:
+				print('Epoch {}, train_loss ={}, val_loss={}'.format(e, epoch_loss, self.validate()))
+			else:
+				print('Epoch Loss = {}, epoch={}'.format(epoch_loss, e))
 
 	def train_epoch(self, noisy_batch, target_batch):
 		epoch_loss = 0
@@ -71,23 +88,43 @@ def train_epoch(self, noisy_batch, target_batch):
 			except tf.errors.OutOfRangeError:
 				return epoch_loss / n_batch
 
+	def __load_weights(self):
+		weights_file = "Checkpoints/weights-epoch-6loss-0.005/weights-epoch-6loss-0.005.ckpt"
+		if not self.training:
+			print('Loaded weights')
+			self.saver.restore(self.sess, weights_file)
+
 	def checkpoint(self, epoch, loss):
 		epoch = str(epoch)
 		loss = "{:.3f}".format(loss)
 		file_name = 'weights-epoch-' + epoch + 'loss-' + loss
 		save_path = self.saver.save(self.sess, 'Checkpoints/' + file_name + "/" + file_name + '.ckpt')
+
 		print('Checkpoint for epoch {}, loss {} saved in {}'.format(epoch, loss, save_path))
 
-	def load(self, ckpt_path):
-		self.saver.restore(self.sess, ckpt_path)
+	def load(self, saved_path):
+		self.saver.restore(self.sess, saved_path)
 
-	def denoise(self, noisy_image):
-		latent, output_t = self.sess.run([self.conv8, self.output_image], feed_dict={self.input_image: noisy_image})
-		print(latent)
+	def denoise_patch(self, image_patch):
+		image_patch = image_patch.reshape(1, 50, 50, 3)
+		latent, output_t = self.sess.run([self.conv7, self.output_image], feed_dict={self.input_image: image_patch})
 		output_t = np.array(output_t) * 255.0
 		output_t = output_t.reshape(self.input_shape)
-		# print(output_t)
-		return Image.fromarray(output_t.astype('uint8')).convert('RGB')
+		return output_t
+
+	def denoise(self, image_array):
+		d_image = np.zeros(shape=image_array.shape)
+		for x in range(50, 3000, 50):
+			for y in range(50, 3000, 50):
+				patch = image_array[x - 50:x, y - 50:y, :]
+
+				if patch.shape[0] != 50 or patch.shape[1] != 50:
+					continue
+				patch = self.denoise_patch(patch)
+				d_image[x - 50:x, y - 50:y, :] = patch
+
+		# print(d_image)
+		return Image.fromarray(d_image.astype('uint8')).convert('RGB')
 
 	def close_session(self):
 		self.sess.close()
@@ -101,10 +138,10 @@ def parser(record):
 		parsed = tf.parse_single_example(record, keys_to_feature)
 		target_image = tf.decode_raw(parsed['reference'], tf.uint8)
 		target_image = tf.cast(target_image, tf.float32)
-		target_image = tf.reshape(target_image, [1240, 1240, 3])
+		target_image = tf.reshape(target_image, [50, 50, 3])
 		noisy_image = tf.decode_raw(parsed['noisy'], tf.uint8)
 		noisy_image = tf.cast(noisy_image, tf.float32)
-		noisy_image = tf.reshape(noisy_image, [1240, 1240, 3])
+		noisy_image = tf.reshape(noisy_image, [50, 50, 3])
 		return noisy_image, target_image
 
 	def input_fn(self, filename, train, batch_size=4, buffer_size=2048):
@@ -116,15 +153,3 @@ def input_fn(self, filename, train, batch_size=4, buffer_size=2048):
 		iterator = dataset.make_one_shot_iterator()
 		noisy_batch, target_batch = iterator.get_next()
 		return noisy_batch, target_batch
-
-
-d = DenoisingAutoEncoder((1240, 1240, 3), tf.train.AdamOptimizer(), True)
-# d.train(100, 10)
-d.load('Checkpoints/weights-epoch-100loss-0.033/weights-epoch-100loss-0.033.ckpt')
-sample_img = Image.open('/home/aftaab/Datasets/Mi3_Aligned/Batch_017//IMG_20151116_151714Noisy.bmp').convert(
-	'RGB').resize([1240, 1240])
-sample_img_t = np.array(sample_img).reshape((1, 1240, 1240, 3)) / 255.0
-d_img = d.denoise(sample_img_t)
-d_img.save('denoised.png', 'PNG')
-sample_img.save('noisy.png', 'PNG')
-d.close_session()