'''
@Description: The script is to transfer the normal map from world coordinate to camera coordinate
'''
import os, os.path as osp, numpy as np, cv2, argparse, tqdm, pdb

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--root', type=str)
    parser.add_argument('--out_dir', type=str, default='normal_test')
    return parser.parse_args()

def convert_normal_map():
    args = parse_args()
    normal_dir = osp.join(args.root, 'normals')
    normal_paths = sorted([osp.join(normal_dir, name) for name in os.listdir(normal_dir)])
    out_dir = osp.join(args.root, args.out_dir)
    os.makedirs(out_dir, exist_ok=True) 
    
    for idx, norm_f in enumerate(normal_paths):
        normals = cv2.imread(norm_f)[:, :, ::-1]
        normals = 2.* normals.astype(np.float32)/255. - 1. # [0, 255] -> [-1, 1]
        normals = 
        cv2.imwrite(osp.join(out_dir, f'{idx:05d}.png'), normals)
        pdb.set_trace()
        




def main():
    convert_normal_map()
    
if __name__ == '__main__':
    main()