TDA4VP-Q1: LCAC - Clarification on TDA Variable Values

import numpy as np
from scipy.interpolate import interp1d
import os

#Base matlab code is taken from : https://e2e.ti.com/support/processors-group/processors/f/processors-forum/1544407/am62a7-the-cac-module-is-invalid

def find_block_s(imsz):
    """
    Finds the optimal block size for the down-sampled correction grid.

    The goal is to find the smallest block size `k` (in increments of 4)
    such that the total number of grid points (width * height) does not
    exceed a predefined limit of 2048. This keeps the LUT size manageable.

    Args:
        imsz (list or tuple): A list containing [height, width] of the image.

    Returns:
        int: The optimal block size, or 0 if no suitable size is found.
    """
    # Iterate through potential block sizes from 4 up to 128, in steps of 4.
    for k in range(4, 129, 4):
        # Calculate the number of grid points needed for width and height.
        # We use np.ceil and add 1 to ensure the grid covers the entire image.
        grid_w = np.ceil(imsz[1] / k) + 1
        grid_h = np.ceil(imsz[0] / k) + 1
        
        # Check if the total number of points is within the allowed limit.
        if grid_w * grid_h <= 2048:
            return k  # Return the first suitable block size found.
            
    return 0 # Return 0 if the loop completes without finding a size.

def prep_cac_dxdy(im_sz, hc, vc, lca_radial_curve, block_s):
    """
    Prepares the displacement grid (dx, dy) for a single color channel.

    This function creates a grid of points, converts their coordinates to
    polar form relative to the image center, applies the radial distortion
    by interpolating from the given curve, and then converts back to
    Cartesian displacement vectors (dx, dy).

    Args:
        im_sz (list): Image size as [height, width].
        hc (float): Horizontal center of the image.
        vc (float): Vertical center of the image.
        lca_radial_curve (np.array): A Nx2 array where column 0 is the radial
                                     distance and column 1 is the shift amount.
        block_s (int): The down-sampling block size for the grid.

    Returns:
        (np.array, np.array): A tuple containing the dx and dy displacement grids.
    """
    # 1. Create a meshgrid of coordinates for the down-sampled LUT.
    x_coords = np.arange(0, np.ceil(im_sz[1] / block_s) + 1) * block_s
    y_coords = np.arange(0, np.ceil(im_sz[0] / block_s) + 1) * block_s
    # meshgrid takes two 1D arrays (e.g., x and y) and generates two 2D arrays (X and Y) that represent a grid.
    xt, yt = np.meshgrid(x_coords, y_coords)

    # 2. Convert grid coordinates to polar coordinates (radius 'ro' and angle 'phi')
    #    relative to the optical center (hc, vc).
    phi = np.arctan2(yt - vc, xt - hc)
    # (hc, vc) is the optical center of the image (Light passes straight through this)
    # (xt, yt) is the coordinate of a point on the grid
    # 'ro' is the radial distance
    ro = np.sqrt((xt - hc)**2 + (yt - vc)**2)

    # 3. Interpolate the radial shift. For each point's radius 'ro', find the
    #    corresponding shift value from the 'lca_radial_curve'.
    #    - `interp1d` creates a function that performs linear interpolation.
    #    - `bounds_error=False` prevents errors if 'ro' is outside the curve's range.
    #    - `fill_value` specifies how to handle extrapolation: use the first value
    #      for points closer than the minimum, and the last value for points
    #      further than the maximum.
    interp_func = interp1d(
        lca_radial_curve[:, 0], 
        lca_radial_curve[:, 1], 
        kind='linear', 
        bounds_error=False, 
        fill_value=(lca_radial_curve[0, 1], lca_radial_curve[-1, 1])
    )
    # Apply the interpolation to get the new, shifted radii.
    ro1 = interp_func(ro)

    # 4. Convert the shifted polar coordinates back to Cartesian coordinates.
    #    This gives the new positions of the grid points.
    xt_shifted, yt_shifted = ro1 * np.cos(phi), ro1 * np.sin(phi)

    # 5. Calculate the displacement vectors (dx, dy). This is the negative of the
    #    calculated shift, scaled by a factor of 8 (as required by the format).
    dx = np.round(-xt_shifted * 8)
    dy = np.round(-yt_shifted * 8)
    
    return dx, dy

def prep_cac_text_lut(mesh1_dx, mesh1_dy, mesh2_dx, mesh2_dy):
    """
    Formats the two displacement grids into the final 4-column LUT text format.

    This involves clamping the values to a specific range, transposing and
    flattening the arrays to the correct order, and concatenating them.

    Args:
        mesh1_dx, mesh1_dy (np.array): Displacement grids for the first channel.
        mesh2_dx, mesh2_dy (np.array): Displacement grids for the second channel.

    Returns:
        np.array: The final formatted LUT, ready to be saved to a text file.
    """
    # --- Process the first displacement mesh ---
    # 1. Clamp the dx/dy values to their specified hardware limits.
    mesh1_dx = np.clip(mesh1_dx, -64, 64)
    mesh1_dy = np.clip(mesh1_dy, -32, 32)
    
    # 2. Transpose and flatten the arrays. The 'T' (transpose) and 'flatten'
    #    operations replicate the column-major ordering of Matlab's `(:)`.
    dx1 = mesh1_dx.T.flatten()
    dy1 = mesh1_dy.T.flatten()
    
    # 3. Stack them vertically and then transpose to get a [N, 2] array.
    mesh1_lut0 = np.vstack((dx1, dy1)).T

    # --- Process the second displacement mesh (same steps) ---
    mesh2_dx = np.clip(mesh2_dx, -64, 64)
    mesh2_dy = np.clip(mesh2_dy, -32, 32)
    dx2 = mesh2_dx.T.flatten()
    dy2 = mesh2_dy.T.flatten()
    mesh2_lut0 = np.vstack((dx2, dy2)).T

    # 4. Concatenate the two resulting arrays horizontally to create the final
    #    4-column LUT: [dx1, dy1, dx2, dy2].
    cac_text_lut = np.hstack((mesh1_lut0, mesh2_lut0))
    
    return cac_text_lut

def make_cac_lut(fout, im_sz, hc, vc, cac_radial_shift_1, cac_radial_shift_2):
    """
    Main function to generate and save the Chromatic Aberration Correction LUT.
    """
    # Step 1: Find the optimal down-sampling block size.
    block_s = find_block_s(im_sz)
    if block_s == 0:
        print("Error: Could not find a suitable block size.")
        return

    # Step 2: Prepare the displacement grids for both correction curves.
    mesh_dx1, mesh_dy1 = prep_cac_dxdy(im_sz, hc, vc, cac_radial_shift_1, block_s)
    mesh_dx2, mesh_dy2 = prep_cac_dxdy(im_sz, hc, vc, cac_radial_shift_2, block_s)
    
    # Step 3: Format the displacement grids into the final text LUT structure.
    mesh_lut = prep_cac_text_lut(mesh_dx1, mesh_dy1, mesh_dx2, mesh_dy2)
    
    # Step 4: Save the final LUT to a text file, using tabs as delimiters
    # and formatting the numbers as integers.
    np.savetxt(fout, mesh_lut, delimiter='\t', fmt='%d')
    
    # Step 5: Print summary information.
    print(f"Image size: {im_sz[1]} x {im_sz[0]}")
    print(f"CAC LUT size: {mesh_dx1.shape[1]} x {mesh_dx1.shape[0]}")
    print(f"CAC down-sampling: {block_s}")
    print(f"Successfully created '{fout}' ({os.path.getsize(fout)} bytes)")

def main():
    """
    Entry point of the script. Defines parameters and calls the LUT generator.
    """
    # --- Configuration ---
    # These arrays represent the measured chromatic aberration for the lens.
    # They are not coefficients for a formula, but rather a set of
    # [distance, shift] data points that define a correction curve.
    # Column 1: Radial distance from the optical center (in pixels).
    # Column 2: The measured radial shift for a color channel at that distance (in pixels).

cac_shift_1 = np.array([
    [      0, 0],
    [    400, 0],
    [    600, -2],
    [    800, -4],
    [    1000, -6],
    [    1100, -8],
    [    1300, -11],
    ])

cac_shift_2 = np.array([
    [      0,  0],
    [    400,  0],
    [    600,  2],
    [    800,  4],
    [    1000, 6],
    [    1100, 8],
    [   1300,  11],
    ])
   
# Image dimensions [height, width]
im_sz = [1536, 1824]
    
    # Optical center of the image
hc = im_sz[1] / 2
vc = im_sz[0] / 2

# Output filename for the generated LUT
output_filename = 'cac_lut_ganghua_py.txt'

# --- Generate LUT ---
print("Starting CAC LUT generation...")
make_cac_lut(output_filename, im_sz, hc, vc, cac_shift_1, cac_shift_2)
print("...Generation complete.")

# This standard Python construct ensures that the `main()` function is called
# only when the script is executed directly.
if __name__ == "__main__":
    main()