test.py

import json
import os

import tensorflow as tf
import tensorflow_datasets as tfds
from PIL import Image

# Load the Div2k dataset
# div2k_data = tfds.image.Div2k(config="bicubic_x4")
# div2k_data.download_and_prepare()

# # Taking train data from div2k_data object
# train_div2k = div2k_data.as_dataset(split="train", as_supervised=True)

# Define the directories containing the high-resolution and low-resolution images
train_hr_dir = './data/train/HR'
train_lr_dir = './data/train/LR'
val_hr_dir = './data/validation/HR'
val_lr_dir = './data/validation/LR'

# Function to load custom images
# Function to load custom images


def load_custom_images(hr_dir, lr_dir):
    def image_generator():
        i = 0
        for filename in os.listdir(hr_dir):
            hr_image_path = os.path.join(hr_dir, filename)
            lr_image_path = os.path.join(lr_dir, filename)

            hr_image = Image.open(hr_image_path)
            lr_image = Image.open(lr_image_path)
            if (hr_image.mode != "RGB"):
                hr_image = hr_image.convert("RGB")

            if (lr_image.mode != "RGB"):
                lr_image = lr_image.convert("RGB")
            yield (tf.keras.preprocessing.image.img_to_array(lr_image),
                   tf.keras.preprocessing.image.img_to_array(hr_image))

    return tf.data.Dataset.from_generator(image_generator,
                                          output_types=(tf.uint8, tf.uint8),
                                          output_shapes=((None, None, 3), (None, None, 3)))


# Load custom images as a dataset
custom_train = load_custom_images(train_hr_dir, train_lr_dir)
custom_val = load_custom_images(val_hr_dir, val_lr_dir)

print("customdataset", custom_train)

y = 0
for lr, hr in custom_train:
    y = y+1
    print("y", y)
    print("hr", hr.shape[2])
    print("lr", lr.shape[2])


# Concatenate the Div2k dataset with the custom dataset
# combined_dataset = train_div2k.concatenate(custom_dataset)