NAFNet/scripts/make_pickle.py

56 lines
1.2 KiB
Python

import numpy as np
import os
import cv2
PATH = './datasets/SR/NTIRE22-StereoSR/Train'
LR_FOLDER = 'LR_x4'
HR_FOLDER = 'HR'
lr_lists = []
hr_lists = []
cnt = 0
for idx in range(1, 801):
L_name = f'{idx:04}_L.png'
R_name = f'{idx:04}_R.png'
LR_L = cv2.imread(os.path.join(PATH, LR_FOLDER, L_name))
LR_R = cv2.imread(os.path.join(PATH, LR_FOLDER, R_name))
HR_L = cv2.imread(os.path.join(PATH, HR_FOLDER, L_name))
HR_R = cv2.imread(os.path.join(PATH, HR_FOLDER, R_name))
LR = np.concatenate([LR_L, LR_R], axis=-1)
HR = np.concatenate([HR_L, HR_R], axis=-1)
lr_lists.append(LR)
hr_lists.append(HR)
cnt = cnt + 1
if cnt % 50 == 0:
print(f'cnt .. {cnt}, idx: {idx}')
import pickle
with open('./datasets/ntire-stereo-sr.train.lr.pickle', 'wb') as f:
pickle.dump(lr_lists, f)
with open('./datasets/ntire-stereo-sr.train.hr.pickle', 'wb') as f:
pickle.dump(hr_lists, f)
# print(f'... {lr_all_np.shape}, {lr_all_np.dtype}')
# print(f'... {hr_all_np.shape}, {hr_all_np.dtype}')
# np.save('./datasets/ntire-stereo-sr.train.lr.npy', lr_all_np)
# np.save('./datasets/ntire-stereo-sr.train.hr.npy', hr_all_np)