czzhangheng
/
TrafficWheel
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update data_selector

This commit is contained in:
czzhangheng 2025-11-20 10:51:32 +08:00
parent 3b6dc18742
commit a9e7cd5d3b
1 changed files with 6 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
dataloader/data_selector.py
View File

@ -20,40 +20,22 @@ def load_st_dataset(config):
data = np.load(data_path)["data"][:, :, 0] data = np.load(data_path)["data"][:, :, 0]
case "PEMSD4": case "PEMSD4":
data_path = os.path.join("./data/PEMS04/PEMS04.npz") data_path = os.path.join("./data/PEMS04/PEMS04.npz")
data = np.load(data_path)["data"][ data = np.load(data_path)["data"][:, :, 0]
:, :, 0
]
case "PEMSD7": case "PEMSD7":
data_path = os.path.join("./data/PEMS07/PEMS07.npz") data_path = os.path.join("./data/PEMS07/PEMS07.npz")
data = np.load(data_path)["data"][ data = np.load(data_path)["data"][:, :, 0]
:, :, 0
]
case "PEMSD8": case "PEMSD8":
data_path = os.path.join("./data/PEMS08/PEMS08.npz") data_path = os.path.join("./data/PEMS08/PEMS08.npz")
data = np.load(data_path)["data"][ data = np.load(data_path)["data"][:, :, 0]
:, :, 0
]
case "PEMSD7(L)": case "PEMSD7(L)":
data_path = os.path.join("./data/PEMS07(L)/PEMS07L.npz") data_path = os.path.join("./data/PEMS07(L)/PEMS07L.npz")
data = np.load(data_path)["data"][ data = np.load(data_path)["data"][:, :, 0]
:, :, 0
]
case "PEMSD7(M)": case "PEMSD7(M)":
data_path = os.path.join("./data/PEMS07(M)/V_228.csv") data_path = os.path.join("./data/PEMS07(M)/V_228.csv")
data = np.genfromtxt( data = np.genfromtxt(data_path, delimiter=",")
data_path, delimiter=","
)
case "METR-LA":
data_path = os.path.join("./data/METR-LA/METR.h5")
with h5py.File(
data_path, "r"
) as f:
data = np.array(f["data"])
case "BJ": case "BJ":
data_path = os.path.join("./data/BJ/BJ500.csv") data_path = os.path.join("./data/BJ/BJ500.csv")
data = np.genfromtxt( data = np.genfromtxt(data_path, delimiter=",", skip_header=1)
data_path, delimiter=",", skip_header=1
)
case "Hainan": case "Hainan":
data_path = os.path.join("./data/Hainan/Hainan.npz") data_path = os.path.join("./data/Hainan/Hainan.npz")
data = np.load(data_path)["data"][:, :, 0] data = np.load(data_path)["data"][:, :, 0]