寻求图像去畸变（鱼眼校正）的TensorFlow开源代码或内置函数

Great question! Integrating fisheye undistortion directly into your TensorFlow pipeline is a smart move for end-to-end computer vision workflows—no more juggling preprocessing outside your model graph. Since there’s no built-in TF function for this, you’ll need to implement the fisheye distortion model using TensorFlow’s vectorized tensor operations (which is totally doable, since the math behind fisheye correction is well-defined).

方法1：手动实现OpenCV兼容的鱼眼校正（推荐） #

OpenCV’s fisheye correction uses the equidistant distortion model (the most common for wide-angle lenses), which we can replicate in TensorFlow. Here’s a complete, TF2.x-compatible implementation that mirrors OpenCV’s cv2.fisheye.undistortImage behavior:

import tensorflow as tf

def tf_fisheye_undistort(image, K, D):
    """
    Args:
        image: Input tensor of shape (H, W, 3) or (B, H, W, 3) (batch support)
        K: Camera intrinsic matrix tensor of shape (3,3) or (B,3,3)
        D: Fisheye distortion coefficients tensor of shape (4,) or (B,4) [k1, k2, k3, k4]
    Returns:
        Undistorted image tensor of same shape as input
    """
    # Handle batch vs single image input
    is_batch = len(image.shape) == 4
    if not is_batch:
        image = tf.expand_dims(image, 0)
        K = tf.expand_dims(K, 0)
        D = tf.expand_dims(D, 0)
    
    # Extract intrinsic parameters from camera matrix
    fx = K[:, 0, 0]
    fy = K[:, 1, 1]
    cx = K[:, 0, 2]
    cy = K[:, 1, 2]
    k1, k2, k3, k4 = tf.split(D, 4, axis=-1)
    
    # Generate pixel coordinate grid
    H, W = image.shape[1], image.shape[2]
    u = tf.range(W, dtype=tf.float32)
    v = tf.range(H, dtype=tf.float32)
    u_grid, v_grid = tf.meshgrid(u, v)
    u_grid = tf.expand_dims(u_grid, 0)
    v_grid = tf.expand_dims(v_grid, 0)
    
    # Convert pixel coordinates to normalized camera space
    x = (u_grid - cx[:, None, None]) / fx[:, None, None]
    y = (v_grid - cy[:, None, None]) / fy[:, None, None]
    
    # Compute distortion factor using equidistant model
    r_sq = x**2 + y**2
    distortion = 1 + k1[:, None, None] * r_sq + k2[:, None, None] * r_sq**2 + \
                k3[:, None, None] * r_sq**3 + k4[:, None, None] * r_sq**4
    
    # Undistort coordinates and convert back to pixel space
    x_undist = x * distortion
    y_undist = y * distortion
    u_undist = x_undist * fx[:, None, None] + cx[:, None, None]
    v_undist = y_undist * fy[:, None, None] + cy[:, None, None]
    
    # Normalize coordinates for bilinear sampling
    u_norm = (u_undist / (W - 1)) * 2 - 1
    v_norm = (v_undist / (H - 1)) * 2 - 1
    grid = tf.stack([u_norm, v_norm], axis=-1)
    
    # Sample from original image to get undistorted result
    undistorted_image = tf.image.sample_bilinear(image, grid, align_corners=True)
    
    if not is_batch:
        undistorted_image = tf.squeeze(undistorted_image, 0)
    
    return undistorted_image

代码说明 #

• Supports both single images and batches, which is critical for training pipelines
• Reuses your existing OpenCV-calibrated intrinsic parameters (K and D) directly
• Uses tf.image.sample_bilinear for smooth, high-quality interpolation matching OpenCV’s INTER_LINEAR flag

方法2：导入OpenCV预计算的映射（静态场景） #

If you’re working with a fixed camera (so K and D never change), you can precompute distortion maps with OpenCV, then convert them to TensorFlow tensors for faster inference:

import cv2
import numpy as np
import tensorflow as tf

# Precompute maps with OpenCV
K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float32)
D = np.array([k1, k2, k3, k4], dtype=np.float32)
H, W = 480, 640
map1, map2 = cv2.fisheye.initUndistortRectifyMap(K, D, np.eye(3), K, (W, H), cv2.CV_32FC1)

# Convert maps to TF tensors
map1_tf = tf.convert_to_tensor(map1, dtype=tf.float32)
map2_tf = tf.convert_to_tensor(map2, dtype=tf.float32)

# TF function to apply undistortion
@tf.function
def tf_undistort_with_opencv_maps(image):
    u_norm = (map1_tf / (W - 1)) * 2 - 1
    v_norm = (map2_tf / (H - 1)) * 2 - 1
    grid = tf.stack([u_norm, v_norm], axis=-1)
    return tf.image.sample_bilinear(tf.expand_dims(image, 0), grid, align_corners=True)[0]

关键注意事项 #

• Vectorization: Avoid loops entirely—TensorFlow’s strength lies in vectorized operations that run efficiently on GPUs/TPUs. The first method uses meshgrids to handle all pixels in parallel.
• TF Version Compatibility: For TF1.x, replace tf.image.sample_bilinear with tf.contrib.image.transform and adjust grid normalization accordingly.
• Trainable Distortion: If you want to learn distortion parameters end-to-end (instead of using calibrated values), make K and D trainable variables and add regularization to keep them physically plausible.

内容的提问来源于stack exchange，提问作者bremen_matt