Wrapper.py

import cv2
import numpy as np
import glob
import os
from scipy.optimize import least_squares

folder_path = 'Calibration_Imgs'
image_files=glob.glob(os.path.join(folder_path, '*.jpg'))
image_files.sort()


#Ground truth M matrix
# The original image has 10 rows and 7 columns, but the chessboard corners are only detected in the inner corners, which are 6 rows and 9 columns.
rows=6
cols=9
inner_cols = cols - 2 
inner_rows = rows - 2
square_size=21.5 #mm

def Worldpoints(square_size, inner_cols, inner_rows):
    x_grid, y_grid = np.meshgrid(np.arange(inner_cols), np.arange(inner_rows))
    M = np.zeros((inner_rows * inner_cols, 3), np.float32)
    M[:, 0] = x_grid.flatten() * square_size
    M[:, 1] = y_grid.flatten() * square_size
    return M 

M = Worldpoints(square_size, inner_cols, inner_rows)
print(f"World points M Shape: {M.shape}")
print("M:") 
print(M)


#Detect the corners of the chessboard in the images and store the detected corners in a list. The corners should be stored in the same order as the ground truth M matrix.

image_files = glob.glob(os.path.join(folder_path, '*.jpg'))
image_files.sort()

def detect_corners(image_files, cols, rows):
    all_img_pts = []   
    all_world_pts = [] 
    good_files = []    

    #print(f"\nScanning {len(image_files)} images for {inner_cols}x{inner_rows} pattern ...")

    for fname in image_files:
        img = cv2.imread(fname)
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        ret, corners = cv2.findChessboardCorners(
            gray, (cols, rows),
            flags=cv2.CALIB_CB_ADAPTIVE_THRESH + cv2.CALIB_CB_NORMALIZE_IMAGE) 
        
        corners_grid = corners.reshape(rows, cols, 2)  # (6, 9, 2)
        corners_inner = corners_grid[1:-1, 1:-1, :]                  # (4, 7, 2)
        corners_use = corners_inner.reshape(-1, 2).astype(np.float32) # (28, 2)

        all_img_pts.append(corners_use)
        all_world_pts.append(M.copy())
        good_files.append(fname)
        #print(f"  [OK]   {os.path.basename(fname)}")

    #print(f"\n{len(good_files)}/{len(image_files)} images accepted.")
    return good_files, all_img_pts, all_world_pts

#Function to visualise the detected corners on the images, with the usable corners highlighted in red and 
# the origin corner (corner[0]) marked in green. The visualised images can be saved to a specified directory.
def verify_corners(save_dir='corner_vis'):
    good_files, all_img_pts, _ = detect_corners(image_files, cols, rows)
    for i, (fname, img_pts) in enumerate(zip(good_files, all_img_pts)):
        img = cv2.imread(fname)
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        ret, corners_full = cv2.findChessboardCorners(
            gray, (cols, rows),
            flags=cv2.CALIB_CB_ADAPTIVE_THRESH + cv2.CALIB_CB_NORMALIZE_IMAGE)
        if ret:
            cv2.drawChessboardCorners(img, (cols, rows), corners_full, ret)

        for pt in img_pts:
            cv2.circle(img, (int(pt[0]), int(pt[1])), 8, (0, 0, 255), -1)

        ox, oy = int(img_pts[0][0]), int(img_pts[0][1])
        cv2.circle(img, (ox, oy), 12, (0, 255, 0), 3)
        cv2.putText(img, 'origin', (ox + 8, oy - 8),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)

        if save_dir:
            out_path = os.path.join(save_dir, f'corners_{i:02d}.jpg')
            cv2.imwrite(out_path, img)
        



# Run verification Uncomment the following lines to save the corner visualisations to the 'corner_vis' directory.
#verify_corners(show=False, save_dir='corner_vis')
#print("Corner visualisations saved to corner_vis/")


def determineHomography(world_pts, img_pts):
    Ar = []
    for (X, Y), (u, v) in zip(world_pts, img_pts):
        Ar.append([-X, -Y, -1, 0, 0, 0, u*X, u*Y, u])
        Ar.append([0, 0, 0, -X, -Y, -1, v*X, v*Y, v])
    A = np.array(Ar)
    U, S, Vt = np.linalg.svd(A)
    H = Vt[-1].reshape(3, 3)
    return H / H[2, 2]  

def computeV(H,i,j):
    h=H.T
    return np.array([
        h[i][0]*h[j][0],                                # B11 term
        h[i][0]*h[j][1] + h[i][1]*h[j][0],              # B12 term
        h[i][1]*h[j][1],                                # B22 term
        h[i][2]*h[j][0] + h[i][0]*h[j][2],              # B13 term
        h[i][2]*h[j][1] + h[i][1]*h[j][2],              # B23 term
        h[i][2]*h[j][2]                                 # B33 term  
    ])

def computeIntrinsicParams(H_list):
    V = []
    for H in H_list:
        V.append(computeV(H, 0, 1))  # v12
        V.append(computeV(H, 0, 0) - computeV(H, 1, 1))  # v11 - v22
    V = np.array(V)
    U, S, Vt = np.linalg.svd(V)
    b = Vt[-1]
    B11, B12, B22, B13, B23, B33 = b

    # Compute intrinsic parameters from B
    v0 = (B12*B13 - B11*B23) / (B11*B22 - B12**2)
    lambda_ = B33 - (B13**2 + v0*(B12*B13 - B11*B23)) / B11
    alpha = np.sqrt(lambda_ / B11)
    beta = np.sqrt(lambda_ * B11 / (B11*B22 - B12**2))
    gamma = -B12 * alpha**2 * beta / lambda_
    u0 = gamma * v0 / beta - B13 * alpha**2 / lambda_

    print(f"Computed Intrinsic Parameters:\nalpha: {alpha}\nbeta: {beta}\ngamma: {gamma}\nu0: {u0}\nv0: {v0}")

    K = np.array([[alpha, gamma, u0],
                  [0, beta, v0],
                  [0, 0, 1]])
    
    print("Computed Intrinsic Matrix K:\n", K)
    return K

def calculateExtrinsics(H, K):
    K_inv = np.linalg.inv(K)
    h1 = H[:, 0]
    h2 = H[:, 1]
    h3 = H[:, 2]

    lambda_ = 1 / np.linalg.norm(np.dot(K_inv, h1))
    r1 = lambda_ * np.dot(K_inv, h1)
    r2 = lambda_ * np.dot(K_inv, h2)
    r3 = np.cross(r1, r2)
    t = lambda_ * np.dot(K_inv, h3)

    R = np.column_stack((r1, r2, r3))
    #Implementing Appendix C to ensure R is a perfectly valid rotation matrix
    U,S,Vt=np.linalg.svd(R)
    R=np.dot(U,Vt)

    return R, t

def matrixtoRodrigues(R):
    val=(np.trace(R) - 1) / 2
    theta = np.arccos(np.clip(val, -1.0, 1.0))  # Clip to handle numerical issues
    if theta < 1e-6:
        return np.zeros(3)  # No rotation
    else:
        rx = (R[2, 1] - R[1, 2]) / (2 * np.sin(theta))
        ry = (R[0, 2] - R[2, 0]) / (2 * np.sin(theta))
        rz = (R[1, 0] - R[0, 1]) / (2 * np.sin(theta))
        return theta * np.array([rx, ry, rz])
    
def rodriquesToMatrix(r):
    theta = np.linalg.norm(r)
    if theta < 1e-6:
        return np.eye(3)  # No rotation
    else:
        k = r / theta
        K = np.array([[0, -k[2], k[1]],
                      [k[2], 0, -k[0]],
                      [-k[1], k[0], 0]])
        R = np.eye(3) + np.sin(theta) * K + (1 - np.cos(theta)) * np.dot(K, K)
        return R

def params3Dto1Darray(K,extrinsics_list,k1=0.0, k2=0.0):
    alpha, gamma, u0 = K[0, 0], K[0, 1], K[0, 2]
    beta, v0 = K[1, 1], K[1, 2]
    
    # Start the flat list
    params = [alpha, beta, gamma, u0, v0, k1, k2]
    for R, t in extrinsics_list:
        rod = matrixtoRodrigues(R)
        params.extend(rod.flatten())
        params.extend(t.flatten())
        
    return np.array(params, dtype=np.float64)

def params1Dto3Darray(params,num_images):
    alpha, beta, gamma, u0, v0, k1, k2 = params[:7]
    
    K = np.array([[alpha, gamma, u0],
                  [0.0,   beta,   v0],
                  [0.0,   0.0,    1.0]])
    
    extrinsics_list = []
    idx = 7 
    
    for _ in range(num_images):
        rod = params[idx : idx+3]
        t = params[idx+3 : idx+6]
        R = rodriquesToMatrix(rod)
        t_vec = t.reshape(3, 1)
        
        extrinsics_list.append((R, t_vec))
        idx += 6 
        
    return K, extrinsics_list, k1, k2


def distortioncalibration(ideal_img_pts, observed_img_pts, K):
    for i in range(len(ideal_img_pts)):
        for j in range(len(ideal_img_pts[i])):
            u_ideal = ideal_img_pts[i][j]
            u_observed = observed_img_pts[i][j]

def residuals(params, all_world_pts, all_img_pts):
    K, extrinsics_list, k1, k2 = params1Dto3Darray(params, len(all_world_pts))
    alpha, beta, gamma, u0, v0 = K[0, 0], K[1, 1], K[0, 1], K[0, 2], K[1, 2]
    residuals = []
    for i in range(len(all_world_pts)):
        world_pts = all_world_pts[i]
        img_pts = all_img_pts[i]
        R, t = extrinsics_list[i]
        t_flat = t.flatten()
        for j in range(len(world_pts)):
            X = world_pts[j]
            P_cam = np.dot(R, X) + t_flat
            X,Y,Z = P_cam[0], P_cam[1], P_cam[2]
            x=X/Z
            y=Y/Z
            r2 = x**2 + y**2
            factor = 1.0 + k1 * r2 + k2 * r2**2
            x_distorted = x * factor
            y_distorted = y * factor
            u_proj = alpha * x_distorted + gamma * y_distorted + u0
            v_proj = beta * y_distorted + v0
            u_obs, v_obs = img_pts[j]
            residuals.extend([u_obs - u_proj, v_obs - v_proj])
    return np.array(residuals)

def optimize_parameters(K,extrinsics_list, all_world_pts, all_img_pts):
    initial_params = params3Dto1Darray(K, extrinsics_list)
    result = least_squares(residuals, initial_params,x_scale='jac' ,method='lm', args=(all_world_pts, all_img_pts),verbose=2)
    K_opt, extrinsics_list_opt, k1_opt, k2_opt = params1Dto3Darray(result.x, len(all_world_pts))
    final_residuals = residuals(result.x, all_world_pts, all_img_pts)
    mean_error = np.mean(np.abs(final_residuals))
    #Final output of the optimized parameters and the mean reprojection error after optimization
    print("\n--- Initial K ---")
    print(K)
    print("\n--- Optimized K ---")
    print(K_opt)
    
    print(f"\nFinal Distortion Coefficients:")
    print(f"k1: {k1_opt:.6f}")
    print(f"k2: {k2_opt:.6f}")
    print(f"\nFinal Mean Reprojection Error: {mean_error:.4f} pixels")
    
    return K_opt, extrinsics_list_opt, k1_opt, k2_opt, mean_error

def generate_rectified_images(good_files, all_world_pts, K_opt, extrinsics_opt, k1, k2, save_dir='Rectified_Images'):
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
        
    dist_coeffs = np.array([k1, k2, 0.0, 0.0, 0.0], dtype=np.float64)
    
    alpha, gamma, u0 = K_opt[0, 0], K_opt[0, 1], K_opt[0, 2]
    beta, v0 = K_opt[1, 1], K_opt[1, 2]
    
    print(f"\nGenerating {len(good_files)} rectified images to '{save_dir}'...")
    
    for i, fname in enumerate(good_files):
        img = cv2.imread(fname)
        rectified_img = cv2.undistort(img, K_opt, dist_coeffs)
        
        R, t = extrinsics_opt[i]
        t_flat = t.flatten()
        world_pts = all_world_pts[i]
        
        for j in range(len(world_pts)):
            X_world = world_pts[j]
            
            P_cam = np.dot(R, X_world) + t_flat
            X, Y, Z = P_cam[0], P_cam[1], P_cam[2]
            
            # Normalize
            x = X / Z
            y = Y / Z
            
            u_proj = int(alpha * x + gamma * y + u0)
            v_proj = int(beta * y + v0)
            
            cv2.circle(rectified_img, (u_proj, v_proj), 8, (0, 255, 0), -1)
            
        base_name = os.path.basename(fname)
        save_path = os.path.join(save_dir, f"rectified_{base_name}")
        cv2.imwrite(save_path, rectified_img)
        
    print("Done! Check the 'Rectified_Images' folder.")

if __name__ == "__main__":

    good_files, all_img_pts, all_world_pts = detect_corners(image_files, cols, rows)
    H_list = []
    Extrincsics_list = []
    for img_pts, world_pts in zip(all_img_pts, all_world_pts):
        H = determineHomography(world_pts[:, :2], img_pts)
        H_list.append(H)
        print("Homography H:\n", H)
    K = computeIntrinsicParams(H_list)
    for i, H in enumerate(H_list):
        R, t = calculateExtrinsics(H, K)
        Extrincsics_list.append((R,t))
        print(f"\nImage: {os.path.basename(good_files[i])}")
        print("Rotation Matrix R:\n", R)
        print("Translation Vector t:\n", t)
    K_opt, extrinsics_list_opt, k1_opt, k2_opt, mean_error = optimize_parameters(K, Extrincsics_list, all_world_pts, all_img_pts)
    generate_rectified_images(good_files, all_world_pts, K_opt, extrinsics_list_opt, k1_opt, k2_opt)