-
Notifications
You must be signed in to change notification settings - Fork 1
/
undistort_sequences.py
62 lines (49 loc) · 2.1 KB
/
undistort_sequences.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
"""
Undistorts all images and stores them in a new dir.
Author: Aljosa Osep (osep@vision.rwth-aachen.de)
To use:
python undistort_sequences.py /media/osep/NCLT/images /media/osep/NCLT/ladybug3_calib
"""
import numpy as np
import cv2
import matplotlib.pyplot as plt
import argparse
import re
import os
import glob
import undistort as undis
def main():
parser = argparse.ArgumentParser(description="Undistort images")
parser.add_argument('image_dir', metavar='img', type=str, help='image dir')
parser.add_argument('map_dir', metavar='map', type=str, help='undistortion maps dir')
args = parser.parse_args()
camera_list = ['Cam%d'%x for x in range(0, 6)]
# Get all dirs in image dir
img_dirs_list = os.walk(args.image_dir).next()[1]
for dir_name in img_dirs_list:
print ('Processing seq. dir: %s'%dir_name)
for cam_name in camera_list:
print ('Processing cam: %s'%cam_name)
map_path = os.path.join(args.map_dir, 'U2D_%s_1616X1232.txt'%cam_name)
undistort = undis.Undistort(map_path)
print 'Loaded camera calibration'
curr_dir = os.path.join(args.image_dir, dir_name, 'lb3', cam_name)
img_this_seq_cam_dir = os.path.join(curr_dir, '*.tiff')
print ('Processing: %s'%img_this_seq_cam_dir)
# Get all images in the image dir
all_imgs = glob.glob(img_this_seq_cam_dir)
all_imgs.sort()
for file in all_imgs:
# Call undistort on all images
im = cv2.imread(file)
im_undistorted = undistort.undistort(im)
# Save to %IMG_DIR%/undistorted/*.png
path_dir, path_name = os.path.split(file)
name, ext = os.path.splitext(path_name)
dir_out = os.path.join(path_dir, 'undistorted')
if not os.path.exists(dir_out): # Make sure output dir is there
os.makedirs(dir_out)
path_out = os.path.join(dir_out, '%s.png'%name)
cv2.imwrite(path_out, im_undistorted)
if __name__ == "__main__":
main()