129 lines
4.1 KiB
Python
129 lines
4.1 KiB
Python
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 camera_name not 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().flatten()
|
||
|
||
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
|
||
self.update_undistort_maps()
|
||
return self
|
||
|
||
def undistort(self, image):
|
||
result = cv2.remap(image, *self.undistort_maps, interpolation=cv2.INTER_LINEAR,
|
||
borderMode=cv2.BORDER_CONSTANT)
|
||
return result
|
||
# def undistort(self, image):
|
||
# uimg = cv2.UMat(image)
|
||
# uresult = cv2.remap(
|
||
# uimg,
|
||
# *self.undistort_maps,
|
||
# interpolation=cv2.INTER_LINEAR,
|
||
# borderMode=cv2.BORDER_CONSTANT,
|
||
# borderValue=(0, 0, 0)
|
||
# )
|
||
# return uresult.get()
|
||
|
||
# def project(self, image):
|
||
# result = cv2.warpPerspective(image, self.project_matrix, self.project_shape)
|
||
# return result
|
||
|
||
def project(self, image):
|
||
# 转为 UMat（触发 GPU 路径）
|
||
uimg = cv2.UMat(image)
|
||
uresult = cv2.warpPerspective(
|
||
uimg,
|
||
self.project_matrix,
|
||
self.project_shape,
|
||
flags=cv2.INTER_LINEAR,
|
||
borderMode=cv2.BORDER_CONSTANT,
|
||
borderValue=(0, 0, 0)
|
||
)
|
||
return uresult.get() # 转回 numpy array 供后续使用
|
||
|
||
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 cv2.transpose(image)[:,::-1]
|
||
|
||
def save_data(self):
|
||
fs = cv2.FileStorage(self.camera_file, cv2.FILE_STORAGE_WRITE)
|
||
fs.write("camera_matrix", self.camera_matrix)
|
||
fs.write("dist_coeffs", self.dist_coeffs)
|
||
fs.write("resolution", self.resolution)
|
||
fs.write("project_matrix", self.project_matrix)
|
||
fs.write("scale_xy", np.float32(self.scale_xy))
|
||
fs.write("shift_xy", np.float32(self.shift_xy))
|
||
fs.release()
|