Image-Colorizer/create_dataset.py at main · MasonAndrewHarrison/Image-Colorizer · GitHub

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
import kaggle
import numpy as np
import os
import shutil
import yaml

"""
This uses the CIELAB (Lab) format
and is from the MIRFLICKR-25K dataset.
"""
kaggle.api.authenticate()
kaggle.api.dataset_download_files("shravankumar9892/image-colorization", path="temp/all/", unzip=True)


with open("config.yaml", "r") as f:
    config = yaml.safe_load(f)["create_dataset"]

test_size = config["test_size"]

L = (np.load("temp/all/l/gray_scale.npy") / 256) * 100

ab1 = np.load("temp/all/ab/ab/ab1.npy")
ab2 = np.load("temp/all/ab/ab/ab2.npy")
ab3 = np.load("temp/all/ab/ab/ab3.npy")

ab = np.concat([ab1, ab2, ab3], axis=0)

os.makedirs("CIELAB-Dataset/train", exist_ok=True)
os.makedirs("CIELAB-Dataset/test", exist_ok=True)

L_test = L[:test_size, :, :]
L_train = L[test_size:, :, :]
ab_test = ab[:test_size, :, : ,:]
ab_train = ab[test_size:, :, : ,:]

np.save("CIELAB-Dataset/test/L.npy", L_test)
np.save("CIELAB-Dataset/test/ab.npy", ab_test)
np.save("CIELAB-Dataset/train/L.npy", L_train)
np.save("CIELAB-Dataset/train/ab.npy", ab_train)
shutil.rmtree("temp/")