Upload 4 files

eval.py

@@ -0,0 +1,62 @@
+import os
+import argparse
+import numpy as np
+from keras.models import load_model
+from PIL import Image
+import tensorflow as  tf
+from model import get_model
+
+IMAGE_SIZE = 128
+
+def autocontrast(tensor, cutoff=0):
+    tensor = tf.cast(tensor, dtype=tf.float32)
+    min_val = tf.reduce_min(tensor)
+    max_val = tf.reduce_max(tensor)
+    range_val = max_val - min_val
+    adjusted_tensor = tf.clip_by_value(tf.cast(tf.round((tensor - min_val - cutoff) * (255 / (range_val - 2 * cutoff))), tf.uint8), 0, 255)
+    return adjusted_tensor
+
+def read_image(image_path):
+    image = tf.io.read_file(image_path)
+    image = tf.image.decode_png(image, channels=3)
+    image = autocontrast(image)
+    image.set_shape([None, None, 3])
+    image = tf.cast(image, dtype=tf.float32) / 255
+    return image
+
+def load_data(low_light_image_path):
+    low_light_image = read_image(low_light_image_path)
+    return low_light_image
+
+def evaluate_model(images_path, eval_path, model_path):
+    model = get_model()
+    model.load_weights("./model.h5")
+
+    image_files = [f for f in os.listdir(images_path) if f.endswith(".jpg") or f.endswith(".png")]
+
+    for image_file in image_files:
+        image_path = os.path.join(images_path, image_file)
+        image = load_data(image_path)
+        image = np.expand_dims(image, axis=0)
+
+        generated_image = model.predict(image)
+
+        eval_file = f"eval_{image_file}"
+        eval_file_path = os.path.join(eval_path, eval_file)
+        generated_image = np.squeeze(generated_image, axis=0)
+        generated_image = np.clip(generated_image * 255, 0, 255).astype(np.uint8)  # Clip values and convert to uint8
+        generated_image = Image.fromarray(generated_image)
+        generated_image.save(eval_file_path)
+        print(f"Generated image saved at: {eval_file_path}")
+
+def main():
+    parser = argparse.ArgumentParser(description="Evaluate model on images")
+    parser.add_argument("images_path", type=str, help="Path to images directory")
+    parser.add_argument("eval_path", type=str, help="Path to evaluation directory")
+    parser.add_argument("--model_path", type=str, default="./model.h5", help="Path to model")
+    args = parser.parse_args()
+    os.makedirs(args.eval_path, exist_ok=True)
+    evaluate_model(args.images_path, args.eval_path, args.model_path)
+
+if __name__ == "__main__":
+    main()

model.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6ab872a7258a78d9b5044e8fbb98ef0439446e7a9e8388e2dd5aaebc04eb4fed
+size 6903640

model.py

@@ -0,0 +1,36 @@
+import tensorflow as tf
+
+def residual_block(inputs, filters):
+    x = tf.keras.layers.Conv2D(filters, (3, 3), padding='same', activation='relu')(inputs)
+    x = tf.keras.layers.Conv2D(filters, (3, 3), padding='same')(x)
+    x = tf.keras.layers.add([inputs, x])
+    x = tf.keras.layers.Activation('relu')(x)
+    return x
+
+def get_model():
+    inputs = tf.keras.layers.Input(shape=(None, None, 3))
+    batch_size = tf.shape(inputs)[0]
+    
+    conv1 = tf.keras.layers.Conv2D(32, (3, 3), padding='same', activation='relu')(inputs)
+    conv1 = tf.keras.layers.Conv2D(32, (3, 3), padding='same', activation='relu')(conv1)
+    
+    conv2 = tf.keras.layers.Conv2D(64, (3, 3), padding='same', activation='relu')(conv1)
+    conv2 = tf.keras.layers.Conv2D(64, (3, 3), padding='same', activation='relu')(conv2)
+    
+    conv3 = tf.keras.layers.Conv2D(128, (3, 3), padding='same', activation='relu')(conv2)
+    conv3 = tf.keras.layers.Conv2D(128, (3, 3), padding='same', activation='relu')(conv2)
+    
+    res1 = residual_block(conv3, 128)
+    res2 = residual_block(res1, 128)
+    res3 = residual_block(res2, 128)
+    res4 = residual_block(res3, 128)
+    res5 = residual_block(res4, 128)
+    
+    deconv1 = tf.keras.layers.Conv2DTranspose(64, (3, 3), padding='same', activation='relu')(res5)
+    deconv2 = tf.keras.layers.Conv2DTranspose(32, (3, 3), padding='same', activation='relu')(deconv1)
+    
+    outputs = tf.keras.layers.Conv2D(3, (3, 3), padding='same', activation='sigmoid')(deconv2)
+    outputs=tf.keras.layers.add([inputs, outputs])
+
+    model = tf.keras.models.Model(inputs=inputs, outputs=outputs)
+    return model

train.py

@@ -0,0 +1,114 @@
+import os
+import random
+import numpy as np
+from glob import glob
+from PIL import Image, ImageOps
+import matplotlib.pyplot as plt
+import tensorflow as tf
+from tensorflow import keras
+from tensorflow.keras import layers
+from model import get_model
+
+# functions to create the dataset
+random.seed(1)
+IMAGE_SIZE = 128
+BATCH_SIZE = 4
+MAX_TRAIN_IMAGES = 300
+
+def autocontrast(tensor, cutoff=0):
+    tensor = tf.cast(tensor, dtype=tf.float32)
+    min_val = tf.reduce_min(tensor)
+    max_val = tf.reduce_max(tensor)
+    range_val = max_val - min_val
+    adjusted_tensor = tf.clip_by_value(tf.cast(tf.round((tensor - min_val - cutoff) * (255 / (range_val - 2 * cutoff))), tf.uint8), 0, 255)
+    return adjusted_tensor
+
+def read_image(image_path):
+    image = tf.io.read_file(image_path)
+    image = tf.image.decode_png(image, channels=3)
+    image = autocontrast(image)
+    image.set_shape([None, None, 3])
+    image = tf.cast(image, dtype=tf.float32) / 255
+    return image
+
+
+def random_crop(low_image, enhanced_image):
+    low_image_shape = tf.shape(low_image)[:2]
+    low_w = tf.random.uniform(
+        shape=(), maxval=low_image_shape[1] - IMAGE_SIZE + 1, dtype=tf.int32
+    )
+    low_h = tf.random.uniform(
+        shape=(), maxval=low_image_shape[0] - IMAGE_SIZE + 1, dtype=tf.int32
+    )
+    enhanced_w = low_w
+    enhanced_h = low_h
+    low_image_cropped = low_image[
+        low_h : low_h + IMAGE_SIZE, low_w : low_w + IMAGE_SIZE
+    ]
+    enhanced_image_cropped = enhanced_image[
+        enhanced_h : enhanced_h + IMAGE_SIZE, enhanced_w : enhanced_w + IMAGE_SIZE
+    ]
+    return low_image_cropped, enhanced_image_cropped
+
+
+def load_data(low_light_image_path, enhanced_image_path):
+    low_light_image = read_image(low_light_image_path)
+    enhanced_image = read_image(enhanced_image_path)
+    low_light_image, enhanced_image = random_crop(low_light_image, enhanced_image)
+    return low_light_image, enhanced_image
+
+
+def get_dataset(low_light_images, enhanced_images):
+    dataset = tf.data.Dataset.from_tensor_slices((low_light_images, enhanced_images))
+    dataset = dataset.map(load_data, num_parallel_calls=tf.data.AUTOTUNE)
+    dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)
+    return dataset
+
+# Loss functions
+
+class CustomLoss:
+    def __init__(self, perceptual_loss_model):
+        self.perceptual_loss_model = perceptual_loss_model
+    def perceptual_loss(self, y_true, y_pred):
+        y_true_features = self.perceptual_loss_model(y_true)
+        y_pred_features = self.perceptual_loss_model(y_pred)
+        loss = tf.reduce_mean(tf.square(y_true_features[0] - y_pred_features[0])) + tf.reduce_mean(tf.square(y_true_features[1] - y_pred_features[1]))
+        return loss
+    def charbonnier_loss(self, y_true, y_pred):
+        return tf.reduce_mean(tf.sqrt(tf.square(y_true - y_pred) + tf.square(1e-3)))
+    def __call__(self, y_true, y_pred):
+        return  0.5*self.perceptual_loss(y_true, y_pred)  + 0.4*self.charbonnier_loss(y_true, y_pred)
+
+def peak_signal_noise_ratio(y_true, y_pred):
+    return tf.image.psnr(y_pred, y_true, max_val=255.0)
+
+def main():
+    train_low_light_images = sorted(glob("./lol_dataset/our485/low/*"))[:MAX_TRAIN_IMAGES]
+    train_enhanced_images = sorted(glob("./lol_dataset/our485/high/*"))[:MAX_TRAIN_IMAGES]
+
+    val_low_light_images = sorted(glob("./lol_dataset/our485/low/*"))[MAX_TRAIN_IMAGES:]
+    val_enhanced_images = sorted(glob("./lol_dataset/our485/high/*"))[MAX_TRAIN_IMAGES:]
+
+    train_dataset = get_dataset(train_low_light_images, train_enhanced_images)
+    val_dataset = get_dataset(val_low_light_images, val_enhanced_images)
+
+    #Model for calculating perceptual loss
+    vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet')
+    for layer in vgg.layers:
+        layer.trainable = False  #Freeze all the layers, since this model is for evaluation only
+    outputs = [vgg.get_layer('block3_conv3').output, vgg.get_layer('block4_conv3').output]
+    perceptual_loss_model = tf.keras.models.Model(inputs=vgg.input, outputs=outputs)
+
+    optimizer = tf.keras.optimizers.Adam(learning_rate=1e-4)
+    loss = CustomLoss(perceptual_loss_model)
+    model = get_model()
+
+    model.compile(
+        optimizer=optimizer, loss=loss, metrics=[peak_signal_noise_ratio]
+    )
+
+    history = model.fit(train_dataset, validation_data=val_dataset, epochs=50)
+    model.save_weights("model.h5")
+
+if __name__ == "__main__":
+    main()