File size: 1,267 Bytes
0e868b4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 | import os
import cv2
import numpy as np
from skimage import data, img_as_ubyte
from skimage.transform import resize
def prepare_data():
os.makedirs("test_data", exist_ok=True)
# Load a standard color image (Astronaut)
print("Loading Ground Truth image...")
gt = img_as_ubyte(data.astronaut())
# Save GT
cv2.imwrite("test_data/ground_truth.jpg", cv2.cvtColor(gt, cv2.COLOR_RGB2BGR))
resolutions = {
"128": (128, 128),
"512": (512, 512),
"1080p": (1920, 1080)
}
for name, size in resolutions.items():
print(f"Generating {name}...")
# Resize GT
# Note: resize expects float, so we convert back to ubyte
resized_gt = resize(gt, (size[1], size[0]), anti_aliasing=True) # size is (h, w)
resized_gt = img_as_ubyte(resized_gt)
# Save Resized GT
gt_path = f"test_data/{name}_gt.jpg"
cv2.imwrite(gt_path, cv2.cvtColor(resized_gt, cv2.COLOR_RGB2BGR))
# Convert to Grayscale
gray = cv2.cvtColor(resized_gt, cv2.COLOR_RGB2GRAY)
# Save Grayscale Input
gray_path = f"test_data/{name}_gray.jpg"
cv2.imwrite(gray_path, gray)
print("Data preparation complete.")
if __name__ == "__main__":
prepare_data()
|