Source code for xrprimer.transform.camera.distortion

from typing import Tuple

import cv2
import numpy as np

from xrprimer.data_structure.camera import (
    FisheyeCameraParameter,
    PinholeCameraParameter,
)
from xrprimer.transform.convention.camera import convert_camera_parameter


[docs]def undistort_camera( distorted_cam: FisheyeCameraParameter) -> PinholeCameraParameter: """Undistort a FisheyeCameraParameter to PinholeCameraParameter. Args: distorted_cam (FisheyeCameraParameter): An instance of FisheyeCameraParameter. Convention will be checked, resolution, intrinsic mat and distortion coefficients will be used. Raises: NotImplementedError: Camera convention not supported. Returns: PinholeCameraParameter: Undistorted camera parameter. """ assert isinstance(distorted_cam, FisheyeCameraParameter), \ 'distorted_cam must be an instance of class FisheyeCameraParameter.' # prepare input of cv2.undistort if distorted_cam.convention != 'opencv': distorted_cam = convert_camera_parameter( cam_param=distorted_cam, dst='opencv') dist_coeff_list = distorted_cam.get_dist_coeff() distorted_intrinsic33 = np.array(distorted_cam.get_intrinsic(k_dim=3)) resolution_wh = np.array([distorted_cam.width, distorted_cam.height]) # prepare output of cv2.undistort corrected_intrinsic33 = np.zeros_like(distorted_intrinsic33) corrected_intrinsic33, _ = cv2.getOptimalNewCameraMatrix( distorted_intrinsic33, np.array(dist_coeff_list), resolution_wh, 0, resolution_wh) corrected_cam_param = PinholeCameraParameter( K=corrected_intrinsic33, R=distorted_cam.get_extrinsic_r(), T=distorted_cam.get_extrinsic_t(), name=f'undistort_{distorted_cam.name}', height=distorted_cam.height, width=distorted_cam.width, world2cam=distorted_cam.world2cam, convention='opencv') return corrected_cam_param
[docs]def undistort_images( distorted_cam: FisheyeCameraParameter, image_array: np.ndarray) -> Tuple[PinholeCameraParameter, np.ndarray]: """Undistort a FisheyeCameraParameter to PinholeCameraParameter, and undistort an array of images shot on a fisheye camera. Args: distorted_cam (FisheyeCameraParameter): An instance of FisheyeCameraParameter. Convention will be checked, resolution, intrinsic mat and distortion coefficients will be used. image_array (np.ndarray): An array of images, in shape [n_frame, height, width, n_channel]. Raises: NotImplementedError: Camera convention not supported. Returns: Tuple[PinholeCameraParameter, np.ndarray]: PinholeCameraParameter: Undistorted camera parameter. np.ndarray: Corrected images in the same shape as input. """ # prepare input of cv2.undistort if distorted_cam.convention != 'opencv': distorted_cam = convert_camera_parameter( cam_param=distorted_cam, dst='opencv') distorted_intrinsic33 = np.array(distorted_cam.get_intrinsic(k_dim=3)) dist_coeff_list = distorted_cam.get_dist_coeff() dist_coeff_np = np.array(dist_coeff_list) corrected_cam_param = undistort_camera(distorted_cam=distorted_cam) corrected_intrinsic33 = np.array( corrected_cam_param.get_intrinsic(k_dim=3)) corrected_image_array = np.ones_like(image_array) for image_index, image_np in enumerate(image_array): corrected_image_array[image_index] = cv2.undistort( image_np, distorted_intrinsic33, dist_coeff_np, newCameraMatrix=corrected_intrinsic33) return corrected_cam_param, corrected_image_array
[docs]def undistort_points( distorted_cam: FisheyeCameraParameter, points: np.ndarray) -> Tuple[PinholeCameraParameter, np.ndarray]: """Undistort a FisheyeCameraParameter to PinholeCameraParameter, and undistort an array of points in fisheye camera screen. Parameters and points will be casted to np.float64 before operation. Args: distorted_cam (FisheyeCameraParameter): An instance of FisheyeCameraParameter. Convention will be checked, resolution, intrinsic mat and distortion coefficients will be used. points (np.ndarray): An array of points, in shape [..., 2], int or float. ... could be [n_point, ], [n_frame, n_point, ] [n_frame, n_object, n_point, ], etc. Raises: NotImplementedError: Camera convention not supported. Returns: Tuple[PinholeCameraParameter, np.ndarray]: PinholeCameraParameter: Undistorted camera parameter. np.ndarray: Corrected points location in the same shape as input, dtype is np.float64. """ # prepare input of cv2.undistortPoints if distorted_cam.convention != 'opencv': distorted_cam = convert_camera_parameter( cam_param=distorted_cam, dst='opencv') distorted_intrinsic33 = np.array(distorted_cam.get_intrinsic(k_dim=3)) dist_coeff_list = distorted_cam.get_dist_coeff() dist_coeff_np = np.array(dist_coeff_list) corrected_cam_param = undistort_camera(distorted_cam=distorted_cam) corrected_intrinsic33 = np.array( corrected_cam_param.get_intrinsic(k_dim=3)) shape_backup = points.shape # opencv expects (n, 1, 2) points = points.reshape(-1, 1, 2) corrected_points = cv2.undistortPoints( points.astype(np.float64), cameraMatrix=distorted_intrinsic33, distCoeffs=dist_coeff_np, P=corrected_intrinsic33) corrected_points = corrected_points.reshape(*shape_backup) return corrected_cam_param, corrected_points