my_slam_python/class_Undistort.py at main · sumitsarkar1/my_slam_python · GitHub

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
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import cv2
import numpy as np
import os
import yaml
from pathlib import Path

class UndistortCamera(object):

    def __init__(self, folder_path_to_yaml: str):

        self.sensor_yaml_path = Path(folder_path_to_yaml)
        self.K        = []
        self.K_resize = []
        self.new_K    = []
        self.D        = []
        self.width    = 0
        self.height   = 0

        self.roi    = []
        self.map1   = []
        self.map2   = []


    def load_camera_params(self) -> None:

        """Load camera parameters from YAML file"""
        with open(self.sensor_yaml_path, 'r') as f:
            params = yaml.safe_load(f)

        K_array = np.array(params['intrinsics'])

        fx = K_array[0]
        fy = K_array[1]
        cx = K_array[2]
        cy = K_array[3]

        # camera matrix formed from yaml of euroc dataset
        self.K  = np.array([[fx, 0, cx],
                        [0, fy, cy],
                        [0, 0, 1]])

        # Extract distortion coefficients
        self.D = np.array(params['distortion_coefficients'])

        # Get image dimensions
        image_resolution = params['resolution']
        self.width  = image_resolution[0]
        self.height = image_resolution[1]

        # Generate undistortion maps
        print("Generating undistortion maps...")
        self.new_K, self.roi = cv2.getOptimalNewCameraMatrix(self.K, self.D, (self.width, self.height), 1, (self.width, self.height))
        self.map1, self.map2 = cv2.initUndistortRectifyMap(self.K, self.D, None, self.new_K, (self.width, self.height), cv2.CV_16SC2)

    def get_resized_camera_params(self, H, W) -> None:

        """Load camera parameters from YAML file"""
        with open(self.sensor_yaml_path, 'r') as f:
            params = yaml.safe_load(f)

        K_array = np.array(params['intrinsics'])

        # Get image dimensions
        image_resolution = params['resolution']
        #self.width  = image_resolution[0]
        #self.height = image_resolution[1]

        scale_y = H/image_resolution[1]
        scale_x  = W/image_resolution[0]

        #print(scale_y, scale_x)

        fx = K_array[0]*scale_x
        fy = K_array[1]*scale_y
        cx = K_array[2]*scale_x
        cy = K_array[3]*scale_y

        #print(fx, fy, cx,  cy)

        # camera matrix formed from yaml of euroc dataset
        self.K_resize  = np.array([[fx, 0, cx],
                        [0, fy, cy],
                        [0, 0, 1]])

        # Extract distortion coefficients
        self.D = np.array(params['distortion_coefficients'])

        return self.K_resize , self.D

        # Generate undistortion maps
        #print("Generating undistortion maps...")
        #self.new_K, self.roi = cv2.getOptimalNewCameraMatrix(self.K, self.D, (self.width, self.height), 1, (self.width, self.height))
        #self.map1, self.map2 = cv2.initUndistortRectifyMap(self.K, self.D, None, self.new_K, (self.width, self.height), cv2.CV_16SC2)

    def get_undistortImage(self, img):

        undistorted_image = cv2.remap(img, self.map1, self.map2, cv2.INTER_LINEAR)

        # Crop to ROI if needed
        x, y, w, h = self.roi
        if w > 0 and h > 0:
            undistorted_image = undistorted_image[y:y+h, x:x+w]

        return undistorted_image

    def get_new_K(self):
        return self.new_K

    def get_camera_params(self):
        return self.K , self.D