fcd/data/data_loader.py at main · i4Ds/fcd · 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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
import os
import sys
import numpy as np

from .data_utils import norm_data
from ..utils.print_utils import *


"""
Options:
key = "tr" or "test"
"""


def get_data(
    path="/path/to/your/data/",
    key="tr",
    count=100000,
    x_key="vis",
    y_key="bp",
):
    ds_path = os.path.join(path, key)

    if os.path.exists(ds_path):
        p("Loading data (size={})...".format(count))
        data_x = []
        data_y = []

        sr_count = count // 4
        # below line is hard-coded !!
        # It should be changed if the number of samples in the dataset is changed.
        residual = 200000 // 4

        for k in range(1, 5):
            sr_path = "data_" + key + "_sr-" + str(k) + "_"
            for i in range(sr_count):
                idx = i + residual * (k - 1)
                data_path = os.path.join(ds_path, sr_path + str(idx) + ".npz")
                data = np.load(data_path)
                data_x.append(data[x_key])
                data_y.append(data[y_key])
            d(f"Number of {k} sources -> Done!")

        data_x = np.array(data_x)
        data_y = np.array(data_y)
        p("Done!")
        return data_x, data_y
    else:
        f("Data not found in the specified path:", ds_path)


def get_npz_data(
    path="/path/to/your/data/", x_key="tr_vis", y_key="tr_bp", zip_it=False
):
    if os.path.exists(path):
        p("Loading data...")
        data = np.load(path)
        data_x = data[x_key]
        data_y = data[y_key]
        p("Done!")

        if zip_it:
            return zip(data_x, data_y)
        else:
            return data_x, data_y
    else:
        f("Data not found in the specified path:", path)


def get_norm_npz_data(
    path, x_key="vis", y_key="gt", z_key=None
):
    if os.path.exists(path):
        p("Loading data...")
        data = np.load(path)
        data_x = data[x_key]
        data_y = data[y_key]
        data_z = data[z_key] if z_key in data else None
        p("Done!")

        p("Normalising data...")
        data_x, data_y, alpha_vec = norm_data(data_x, data_y)
        p("Done!")

        return data_x, data_y, alpha_vec, data_z
    else:
        f("Data not found in the specified path:", path)


if __name__ == "__main__":
    data_x, data_y, alpha_vec, data_z = get_norm_npz_data(sys.argv[1])
    if len(sys.argv) > 3:
        np.savez(sys.argv[2], vis=data_x, gt=data_y, alpha=alpha_vec, sigamp=data_z)
        p("Data saved with sigamp!")
    else:
        np.savez(sys.argv[2], vis=data_x, gt=data_y, alpha=alpha_vec)
        p("Data saved!")