Dongyang Jin
72 lines
2.5 KiB
Python
72 lines
2.5 KiB
Python
import os
|
|
from time import time
|
|
from multiprocessing import Pool
|
|
from tqdm import tqdm
|
|
import numpy as np
|
|
import os
|
|
import pickle
|
|
import numpy as np
|
|
import cv2
|
|
from tqdm import tqdm
|
|
|
|
SRC_0 = 'rgb_input_path'
|
|
DST_0 = 'rgb_pkl_output_path'
|
|
|
|
SRC = SRC_0 # Path_of_RGB_rearranged
|
|
DST = DST_0 # Path_of_RGB_256128pkl_PadResized
|
|
|
|
def resize_with_padding(img, target_size):
|
|
h, w, _ = img.shape
|
|
target_h, target_w = target_size
|
|
resized_img = cv2.resize(img, (int(w * target_h / h), target_h))
|
|
padded_img = np.zeros((target_h, target_w, 3), dtype=np.uint8)
|
|
x_offset = (target_w - resized_img.shape[1]) // 2
|
|
if x_offset < 0 :
|
|
x_offset = abs(x_offset)
|
|
padded_img = resized_img[:, x_offset:x_offset+target_w,:]
|
|
else:
|
|
padded_img[:, x_offset:x_offset + resized_img.shape[1]] = resized_img
|
|
return padded_img
|
|
|
|
def job(src, id):
|
|
for ty in sorted(os.listdir(os.path.join(src, id))):
|
|
for vi in sorted(os.listdir(os.path.join(src, id, ty))):
|
|
exist_file = os.path.join(DST, id, ty, vi, vi+"-aligned-rgbs.pkl")
|
|
if os.path.exists(exist_file):
|
|
print('Have Passed: ' + DST + '/' + id + '/' + ty)
|
|
continue
|
|
ratios = []
|
|
aligned_imgs = []
|
|
for img_file in sorted(os.listdir(os.path.join(src, id, ty, vi))):
|
|
img_path = os.path.join(src, id, ty, vi, img_file)
|
|
img = cv2.imread(img_path)
|
|
ratio = img.shape[1]/img.shape[0]
|
|
ratios.append(ratio)
|
|
aligned_img = np.transpose(cv2.cvtColor(resize_with_padding(img, (256, 128)), cv2.COLOR_BGR2RGB), (2, 0, 1))
|
|
aligned_imgs.append(aligned_img)
|
|
if len(aligned_imgs) > 0:
|
|
output_path = os.path.join(DST, id, ty, vi)
|
|
os.makedirs(output_path, exist_ok=True)
|
|
pickle.dump(np.asarray(aligned_imgs), open(os.path.join(output_path, vi+"-aligned-rgbs.pkl"), "wb"))
|
|
pickle.dump(np.asarray(ratios), open(os.path.join(output_path, vi+"-ratios.pkl"), "wb"))
|
|
print('Successfully saved: ' + DST + '/' + id + '/' + ty + '/' + vi)
|
|
|
|
if __name__ == '__main__':
|
|
a = time()
|
|
po = Pool(8)
|
|
src_path = SRC
|
|
|
|
cnt = 0
|
|
need_data = sorted(os.listdir(src_path))
|
|
for id in tqdm(need_data[:]):
|
|
po.apply_async(job,(src_path, id,))
|
|
cnt = cnt + 1
|
|
|
|
print('---START---')
|
|
po.close()
|
|
po.join()
|
|
print(cnt)
|
|
|
|
t = time() - a
|
|
print('---END---{}'.format(t))
|