add new files

Author: neozhaoliang

surround_view/fisheye_camera.py

@@ -0,0 +1,102 @@
+import os
+import numpy as np
+import cv2
+
+from . import param_settings as settings
+
+
+class FisheyeCameraModel(object):
+
+    """
+    Fisheye camera model, for undistorting, projecting and flipping camera frames.
+    """
+
+    def __init__(self, camera_param_file, camera_name):
+        if not os.path.isfile(camera_param_file):
+            raise ValueError("Cannot find camera param file")
+
+        if not camera_name in settings.camera_names:
+            raise ValueError("Unknown camera name: {}".format(camera_name))
+
+        self.camera_file = camera_param_file
+        self.camera_name = camera_name
+        self.scale_xy = (1.0, 1.0)
+        self.shift_xy = (0, 0)
+        self.undistort_maps = None
+        self.project_matrix = None
+        self.project_shape = settings.project_shapes[self.camera_name]
+        self.load_camera_params()
+
+    def load_camera_params(self):
+        fs = cv2.FileStorage(self.camera_file, cv2.FILE_STORAGE_READ)
+        self.camera_matrix = fs.getNode("camera_matrix").mat()
+        self.dist_coeffs = fs.getNode("dist_coeffs").mat()
+        self.resolution = fs.getNode("resolution").mat()
+
+        scale_xy = fs.getNode("scale_xy").mat()
+        if scale_xy is not None:
+            self.scale_xy = scale_xy
+
+        shift_xy = fs.getNode("shift_xy").mat()
+        if shift_xy is not None:
+            self.shift_xy = shift_xy
+
+        project_matrix = fs.getNode("project_matrix").mat()
+        if project_matrix is not None:
+            self.project_matrix = project_matrix
+
+        fs.release()
+        self.update_undistort_maps()
+
+    def update_undistort_maps(self):
+        new_matrix = self.camera_matrix.copy()
+        new_matrix[0, 0] *= self.scale_xy[0]
+        new_matrix[1, 1] *= self.scale_xy[1]
+        new_matrix[0, 2] += self.shift_xy[0]
+        new_matrix[1, 2] += self.shift_xy[1]
+        width, height = self.resolution
+
+        self.undistort_maps = cv2.fisheye.initUndistortRectifyMap(
+            self.camera_matrix,
+            self.dist_coeffs,
+            np.eye(3),
+            new_matrix,
+            (width, height),
+            cv2.CV_16SC2
+        )
+        return self
+
+    def set_scale_and_shift(self, scale_xy=(1.0, 1.0), shift_xy=(0, 0)):
+        self.scale_xy = scale_xy
+        self.shift_xy = shift_xy
+        return self
+
+    def undistort(self, image):
+        result = cv2.remap(
+            image,
+            *self.undistort_maps,
+            interpolation=cv2.INTER_LINEAR,
+            borderMode=cv2.BORDER_CONSTANT
+        )
+        return result
+
+    def project(self, image):
+        result = cv2.warpPerspective(
+            image,
+            self.project_matrix,
+            self.project_shape
+        )
+        return result
+
+    def flip(self, image):
+        if self.camera_name == "front":
+            return image.copy()
+
+        elif self.camera_name == "back":
+            return image.copy()[::-1, ::-1, :]
+
+        elif self.camera_name == "left":
+            return cv2.transpose(image)[::-1]
+
+        else:
+            return np.flip(cv2.transpose(image), 1)