--- a +++ b/datasets/pretreatment_rgb.py @@ -0,0 +1,71 @@ +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))