Initial commit

Author: ihnatova

dng_to_png.py

@@ -0,0 +1,43 @@
+# Copyright 2019-2020 by Andrey Ignatov. All Rights Reserved.
+
+# python dng_to_png.py path_to_my_dng_file.dng
+
+import numpy as np
+import imageio
+import rawpy
+import sys
+import os
+
+
+def extract_bayer_channels(raw):
+
+    ch_B  = raw[1::2, 1::2]
+    ch_Gb = raw[0::2, 1::2]
+    ch_R  = raw[0::2, 0::2]
+    ch_Gr = raw[1::2, 0::2]
+
+    return ch_R, ch_Gr, ch_B, ch_Gb
+
+
+if __name__ == "__main__":
+
+    raw_file = sys.argv[1]
+    print("Converting file " + raw_file)
+
+    if not os.path.isfile(raw_file):
+        print("The file doesn't exist!")
+        sys.exit()
+
+    raw = rawpy.imread(raw_file)
+    raw_image = raw.raw_image
+    del raw
+
+    # Use the following code to rotate the image (if needed)
+    # raw_image = np.rot90(raw_image, k=2)
+
+    raw_image = raw_image.astype(np.float32)
+    ch_R, ch_Gr, ch_B, ch_Gb = extract_bayer_channels(raw_image)
+
+    png_image = raw_image.astype(np.uint16)
+    new_name = raw_file.replace(".dng", ".png")
+    imageio.imwrite(new_name, png_image)

evaluate_accuracy.py

@@ -0,0 +1,73 @@
+# Copyright 2020 by Andrey Ignatov. All Rights Reserved.
+
+import tensorflow as tf
+import numpy as np
+
+import model
+import utils
+
+from load_dataset import load_test_data
+
+PATCH_WIDTH, PATCH_HEIGHT = 224, 224
+TARGET_SIZE = (PATCH_WIDTH * 2) * (PATCH_HEIGHT * 2) * 3
+config = None
+
+# Path to the dataset:
+dataset_dir = 'raw_images/'
+
+# Disable gpu (if needed):
+# config = tf.ConfigProto(device_count={'GPU': 0})
+
+with tf.Session(config=config) as sess:
+
+    # Create placeholders for input and target images
+
+    phone_ = tf.placeholder(tf.float32, [1, PATCH_HEIGHT, PATCH_HEIGHT, 4])
+    dslr_ = tf.placeholder(tf.float32, [1, int(PATCH_HEIGHT * 2), int(PATCH_WIDTH * 2), 3])
+
+    # Process raw images with your model:
+    # enhanced = your_model(phone_), e.g.:
+    enhanced, _, _, _, _, _ = model.PyNET(phone_, instance_norm=True, instance_norm_level_1=False)
+
+    # Compute losses
+
+    enhanced_flat = tf.reshape(enhanced, [-1, TARGET_SIZE])
+    dslr_flat = tf.reshape(dslr_, [-1, TARGET_SIZE])
+
+    loss_ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(enhanced, dslr_, 1.0))
+
+    loss_mse = tf.reduce_sum(tf.pow(dslr_flat - enhanced_flat, 2)) / TARGET_SIZE
+    loss_psnr = 20 * utils.log10(1.0 / tf.sqrt(loss_mse))
+
+    saver = tf.train.Saver()
+
+    # Restore your own model from a checkpoint
+    # saver.restore(sess, "path_to_your_checkpoint"), e.g.:
+    saver.restore(sess, "models/original/pynet_level_0.ckpt")
+
+    print("Loading test data...")
+    test_data, test_answ = load_test_data(dataset_dir, PATCH_WIDTH, PATCH_HEIGHT, 2.0)
+    print("Test data was loaded\n")
+
+    loss_ssim_ = 0.0
+    loss_psnr_ = 0.0
+
+    test_size = test_data.shape[0]
+    for j in range(test_size):
+
+        if j % 100 == 0:
+            print(j)
+
+        phone_images = np.reshape(test_data[j], [1, PATCH_HEIGHT, PATCH_WIDTH, 4])
+        dslr_images = np.reshape(test_answ[j], [1, int(PATCH_HEIGHT * 2), int(PATCH_WIDTH * 2), 3])
+
+        losses = sess.run([loss_psnr, loss_ms_ssim], feed_dict={phone_: phone_images, dslr_: dslr_images})
+
+        loss_psnr_ += losses[0]
+        loss_ssim_ += losses[1]
+
+    loss_psnr_ = float(loss_psnr_) / test_size
+    loss_ssim_ = float(loss_ssim_) / test_size
+
+    output_logs = "PSNR: %.4g, MS-SSIM: %.4g\n" % (loss_psnr_, loss_ssim_)
+    print(output_logs)

load_dataset.py

@@ -0,0 +1,81 @@
+# Copyright 2020 by Andrey Ignatov. All Rights Reserved.
+
+from __future__ import print_function
+from scipy import misc
+from PIL import Image
+import imageio
+import os
+import numpy as np
+
+
+def extract_bayer_channels(raw):
+
+    # Reshape the input bayer image
+
+    ch_B  = raw[1::2, 1::2]
+    ch_Gb = raw[0::2, 1::2]
+    ch_R  = raw[0::2, 0::2]
+    ch_Gr = raw[1::2, 0::2]
+
+    RAW_combined = np.dstack((ch_B, ch_Gb, ch_R, ch_Gr))
+    RAW_norm = RAW_combined.astype(np.float32) / (4 * 255)
+
+    return RAW_norm
+
+
+def load_test_data(dataset_dir, PATCH_WIDTH, PATCH_HEIGHT, DSLR_SCALE):
+
+    test_directory_dslr = dataset_dir + 'test/canon/'
+    test_directory_phone = dataset_dir + 'test/huawei_raw/'
+
+    # NUM_TEST_IMAGES = 1204
+    NUM_TEST_IMAGES = len([name for name in os.listdir(test_directory_phone)
+                           if os.path.isfile(os.path.join(test_directory_phone, name))])
+
+    test_data = np.zeros((NUM_TEST_IMAGES, PATCH_WIDTH, PATCH_HEIGHT, 4))
+    test_answ = np.zeros((NUM_TEST_IMAGES, int(PATCH_WIDTH * DSLR_SCALE), int(PATCH_HEIGHT * DSLR_SCALE), 3))
+
+    for i in range(0, NUM_TEST_IMAGES):
+
+        I = np.asarray(imageio.imread((test_directory_phone + str(i) + '.png')))
+        I = extract_bayer_channels(I)
+        test_data[i, :] = I
+        
+        I = np.asarray(Image.open(test_directory_dslr + str(i) + '.jpg'))
+        I = misc.imresize(I, DSLR_SCALE / 2, interp='bicubic')
+        I = np.float16(np.reshape(I, [1, int(PATCH_WIDTH * DSLR_SCALE), int(PATCH_HEIGHT * DSLR_SCALE), 3])) / 255
+        test_answ[i, :] = I
+
+    return test_data, test_answ
+
+
+def load_training_batch(dataset_dir, TRAIN_SIZE, PATCH_WIDTH, PATCH_HEIGHT, DSLR_SCALE):
+
+    train_directory_dslr = dataset_dir + 'train/canon/'
+    train_directory_phone = dataset_dir + 'train/huawei_raw/'
+
+    # NUM_TRAINING_IMAGES = 46839
+    NUM_TRAINING_IMAGES = len([name for name in os.listdir(train_directory_phone)
+                               if os.path.isfile(os.path.join(train_directory_phone, name))])
+
+    TRAIN_IMAGES = np.random.choice(np.arange(0, NUM_TRAINING_IMAGES), TRAIN_SIZE, replace=False)
+
+    train_data = np.zeros((TRAIN_SIZE, PATCH_WIDTH, PATCH_HEIGHT, 4))
+    train_answ = np.zeros((TRAIN_SIZE, int(PATCH_WIDTH * DSLR_SCALE), int(PATCH_HEIGHT * DSLR_SCALE), 3))
+
+    i = 0
+    for img in TRAIN_IMAGES:
+
+        I = np.asarray(imageio.imread((train_directory_phone + str(img) + '.png')))
+        I = extract_bayer_channels(I)
+        train_data[i, :] = I
+
+        I = np.asarray(Image.open(train_directory_dslr + str(img) + '.jpg'))
+        I = misc.imresize(I, DSLR_SCALE / 2, interp='bicubic')
+        I = np.float16(np.reshape(I, [1, int(PATCH_WIDTH * DSLR_SCALE), int(PATCH_HEIGHT * DSLR_SCALE), 3])) / 255
+        train_answ[i, :] = I
+
+        i += 1
+
+    return train_data, train_answ
+