Rectify Image using Pitch and Roll angles

rectify_from_inclination.py

'''
https://stackoverflow.com/questions/20445147/transform-image-using-roll-pitch-yaw-angles-image-rectification/20469709#20469709
'''

import cv2
import numpy as np


# CORNERS FROM 3D ROTATION AND PROJECTION TO SCREENSPACE
def screenspace_from_3D(image_width, image_height, diagonal_fov, camera_position, pitch, roll, points_3d, pixel_units=True):
    
    def get_camera_intrinsics(image_width, image_height, diagonal_fov):
        # Define principal point
        principal_point = (image_width / 2, image_height / 2)

        # Convert FOV to radians
        diagonal_fov = np.deg2rad(diagonal_fov) # convert to radians

        # Calculate focal length
        f = (image_width / 2) / np.tan(diagonal_fov / 2)

        return principal_point, f

    def get_camera_extrinsics(camera_position, principal_point, f):
        # Create camera matrix
        K = np.array([[f, 0, principal_point[0]],
                [0, f, principal_point[1]],
                [0, 0, 1]])

        # Define camera orientation
        R = np.eye(3)

        # Define camera position
        t = -R @ camera_position

        return K, R, t

    def get_object_rotation_matrix(pitch, roll):
        pitch = np.deg2rad(pitch) # convert to radians
        roll = np.deg2rad(roll) # convert to radians

        Rx = np.array([[1, 0, 0],
                    [0, np.cos(pitch), -np.sin(pitch)],
                    [0, np.sin(pitch), np.cos(pitch)]])
        
        Rz = np.array([[np.cos(roll), -np.sin(roll), 0],
                    [np.sin(roll), np.cos(roll), 0],
                    [0, 0, 1]])
        
        R = Rz @ Rx

        return R


    # Calculate principal point and focal length
    principal_point, f = get_camera_intrinsics(image_width, image_height, diagonal_fov)

    # Calculate camera matrix, orientation and position
    K, R_cam, t = get_camera_extrinsics(camera_position, principal_point, f)

    # Calculate rotation matrix for the points
    R_points = get_object_rotation_matrix(pitch, roll)

    # Rotate points around their local coordinates
    points_3d = (R_points @ points_3d.T).T

    # Transform points from world coordinates to camera coordinates
    points_cam = R_cam @ points_3d.T + t.reshape(3, 1)

    # Project points onto image plane
    points_2d = K @ points_cam
    points_2d /= points_2d[2]

    # Convert pixel coordinates to screenspace coordinates
    screenspace_positions = points_2d[:2].T

    if pixel_units is False:
        screenspace_positions[:, 0] /= image_width
        screenspace_positions[:, 1] /= image_height

    return screenspace_positions


# REORDER CORNERS COUNTER-CLOCKWISE
def reorder(points):
    newPoints = np.zeros((4, 2), dtype=np.int32)
    add = points.sum(1)
    newPoints[0] = points[np.argmin(add)]
    newPoints[2] = points[np.argmax(add)]
    diff = np.diff(points, axis=1)
    newPoints[1] = points[np.argmin(diff)]
    newPoints[3] = points[np.argmax(diff)]
    return newPoints

# WARP FROM HOMOGRAPHY
def rectify_image(image, coords, aspect_ratio=1, crop=False):
    # Convert coordinates to numpy array
    pts1 = np.float32(coords)

    # Calculate the width and height of the rectangle
    width = int(np.sqrt(((pts1[1][0] - pts1[0][0]) ** 2) + ((pts1[1][1] - pts1[0][1]) ** 2)))
    height = int(width / aspect_ratio)

    if crop is False:
        # Calculate the size and position of the output image
        x_min = min(point[0] for point in coords)
        x_max = max(point[0] for point in coords)
        y_min = min(point[1] for point in coords)
        y_max = max(point[1] for point in coords)
        output_x_min = min(0, x_min)
        output_x_max = max(image.shape[1], x_max + width - (x_max - x_min))
        output_y_min = min(0, y_min)
        output_y_max = max(image.shape[0], y_max + height - (y_max - y_min))
        output_width = output_x_max - output_x_min
        output_height = output_y_max - output_y_min

        # Define the coordinates of the destination image
        pts2 = np.float32(
            [[x_min - output_x_min, y_min - output_y_min], [x_min + width - output_x_min, y_min - output_y_min],
            [x_min + width - output_x_min, y_min + height - output_y_min],
            [x_min - output_x_min, y_min + height - output_y_min]])
        
        width = output_width
        height = output_height

    else:
        # Define the coordinates of the destination image
        pts2 = np.float32([[0, 0], [width, 0], [width, height], [0, height]])

    # Calculate the perspective transform matrix
    M = cv2.getPerspectiveTransform(pts1, pts2)

    # Warp the image using the perspective transform matrix
    dst = cv2.warpPerspective(image, M, (width, height))

    return dst

# DRAW EDGES AND CORNERS
def draw_rectangle(image, pts_float, thickness=2, color=(0,0,255), draw_corners=True):
    pts = pts_float.astype(int)
    img = image.copy()
    lineType=cv2.LINE_AA

    for i in range(4):
        cv2.line(img, pts[i], pts[(i + 1) % 4],  color, thickness, lineType)

    if draw_corners:
        unsqueezed_array = pts.reshape(-1, 1, 2)
        img = cv2.drawContours(img, unsqueezed_array, -1, (0, 255, 0), 6)

    return img


crop = False

image_path = "images/initial_image.jpg"
camera_position = (0,0,1000)
diagonal_fov = 22
    
# Define inclination in degrees
pitch = 28.5 - 90
roll =  10 

# Define 3D points in world coordinates
plane = np.array([100,100]) / 2
aspect_ratio = plane[0] / plane[1]
points_3d = np.array([[-plane[0], plane[1], 0],[plane[0],plane[1],0],[plane[0],-plane[1],0],[-plane[0],-plane[1],0]])

# Load image
image = cv2.imread(image_path)
image_width = image.shape[1]
image_height = image.shape[0]

# Calculate 2D vertex positions of hypothetical plane in screenspace
corners = screenspace_from_3D(image_width, image_height, diagonal_fov, camera_position, pitch, roll, points_3d, pixel_units=True)
corners = reorder(corners)
#print("SCREENSPACE POSITIONS:\n", corners)

# Draw Rectangle on the input image
img_with_contours = draw_rectangle(image, corners)
cv2.imshow('Input', img_with_contours)

# Rectify the image
rectified_image = rectify_image(image, corners, aspect_ratio=aspect_ratio, crop=crop)
cv2.imshow('Rectified', rectified_image)

cv2.waitKey(0)
cv2.destroyAllWindows