Camera class with methods to transform coordinates, transform images, and change parameters intuitively (e.g. zoom, orbit around or turn towards point, center principal point)

camera.py

import cv2
import numpy as np


class Camera:
    def __init__(self, eye_point, rot_matrix_world_to_cam, intrinsic_matrix, distortion_coeffs):
        """Initializes camera.

        Args:
            eye_point: position of the camera in world coordinates
            rot_matrix_world_to_cam: 3x3 rotation matrix for transforming column vectors
                from world to camera reference frame as follows:
                column_point_camera = rot_matrix_world_to_cam @ (column_point_world - eye_point)
            intrinsic_matrix: 3x3 matrix that maps 3D points in camera space to homogeneous
                coordinates in image (pixel) space. Last row must be (0,0,1).
            distortion_coeffs: parameters describing radial and tangential lens distortions,
                following OpenCV's model and order: k1, k2, p1, p2, k3
                or None, if there is no distortion
        """
        self.R = np.asarray(rot_matrix_world_to_cam, np.float32)
        self.t = np.asarray(eye_point, np.float32)
        self.intrinsic_matrix = np.asarray(intrinsic_matrix, np.float32)
        self.distortion_coeffs = np.asarray(distortion_coeffs, np.float32)

        if not np.allclose(intrinsic_matrix[2, :], [0, 0, 1]):
            raise Exception(f'Bottom row of camera\'s intrinsic matrix must be (0,0,1), '
                            f'got {intrinsic_matrix[2, :]}.')

    def camera_to_image(self, points):
        return cv2.projectPoints(
            np.expand_dims(points, 0), np.float32([0, 0, 0]), np.float32([0, 0, 0]),
            self.intrinsic_matrix, self.distortion_coeffs)[0][:, 0, :]

        # Equivalently:
        #
        # projected = points[:, 0:2] / points[:, 2:3]
        #
        # if self.distortion_coeffs is not None:
        #     r2 = np.sum(projected[:, 0:2] ** 2, axis=1, keepdims=True)
        #
        #     k = self.distortion_coeffs[(0, 1, 4)]
        #     radial = 1 + np.hstack([r2, r2 ** 2, r2 ** 3]) @ k
        #
        #     p_flipped = self.distortion_coeffs[3:1:-1]
        #     tagential = projected @ (p_flipped * 2)
        #     distorted = projected * np.expand_dims(radial + tagential, -1) + p_flipped * r2
        # else:
        #     distorted = projected
        #
        # return distorted @ self.intrinsic_matrix[:2, :2].T + self.intrinsic_matrix[:2, 2]

    def world_to_camera(self, points):
        points = np.asarray(points, np.float32)
        return (points - self.t.T) @ self.R.T

    def camera_to_world(self, points):
        points = np.asarray(points, np.float32)
        return points @ self.R + self.t.T

    def world_to_image(self, points):
        return self.camera_to_image(self.world_to_camera(points))

    def image_to_camera(self, points):
        points = np.expand_dims(np.asarray(points, np.float32), 0)
        new_image_points = cv2.undistortPoints(
            points, self.intrinsic_matrix, self.distortion_coeffs, None, None, None)[0]
        return cv2.convertPointsToHomogeneous(new_image_points)[:, 0, :]

    def image_to_world(self, points):
        return self.camera_to_world(self.image_to_camera(points))

    def zoom(self, factor):
        """Zooms the camera (>1 makes objects look larger), keeping the principal point fixed."""
        self.intrinsic_matrix[0:2, 0:2] *= factor

    def scale_output(self, factor):
        """Adjusts the camera such that the images become scaled by `factor`.
        The difference with `self.zoom` is that this method also moves the principal point,
        multiplying its coordinates by `factor`."""
        self.intrinsic_matrix[0:2] *= factor

    def undistort(self):
        self.distortion_coeffs = None

    def horizontal_flip(self):
        self.R[0] *= -1

    def center_principal_point(self, imshape):
        """Adjusts the intrinsic matrix so that the principal point becomes located at the center
        of an image sized imshape (height, width)"""

        self.intrinsic_matrix[:2, 2] = [imshape[1] / 2, imshape[0] / 2]

    def shift_to_center(self, desired_center_image_point, imshape):
        """"""
        self.intrinsic_matrix[:2, 2] += (
                np.float32([imshape[1], imshape[0]]) / 2 - desired_center_image_point)

    def turn_towards(self, target_image_point=None, target_world_point=None):
        """Turns the camera so that its optical axis goes through a desired target point."""

        assert (target_image_point is None) != (target_world_point is None)
        if target_image_point is not None:
            target_world_point = self.image_to_world([target_image_point])[0]

        def unit_vec(v):
            return v / np.linalg.norm(v)

        world_up_vector = [0, 0, -1]
        new_z = unit_vec(target_world_point - self.t)
        new_x = unit_vec(np.cross(world_up_vector, new_z))
        new_y = np.cross(new_z, new_x)

        # row_stack because we need the inverse transform (we make a matrix that transforms
        # points from one coord system to another), which is the same as the transpose
        # for rotation matrices.
        self.R = np.row_stack([new_x, new_y, new_z]).astype(np.float32)

    def horizontal_orbit_around(self, world_point, angle_radians):
        """Moves the camera on a circular orbit around `world point` parallel to the ground by
        `angle_radians`, keeping its optical axis unchanged relative to `world_point` (i.e.
        if it pointed towards `world_point` it will still point towards it afterwards)."""

        world_up_vector = np.array([0, 0, -1])
        rot_matrix = cv2.Rodrigues(world_up_vector * angle_radians)[0]
        # The eye position rotates simply as any point
        self.t = (rot_matrix @ (self.t - world_point)) + world_point

        # R is rotated by a transform expressed in world coords, so it (its inverse since its a
        # coord transform matrix, not a point transform matrix) is applied on the right.
        # (inverse = transpose for rotation matrices, they are orthogonal)
        self.R = self.R @ rot_matrix.T

    @staticmethod
    def reproject_image(image, old_camera, new_camera, output_imshape):
        """Transforms an image captured with `old_camera` so it looks like it was captured by
        `new_camera`. The world position (eye point) of the cameras must be the same, otherwise
        we'd have parallax effects and no way to construct the other image."""

        assert np.allclose(old_camera.t, new_camera.t)
        output_size = (output_imshape[1], output_imshape[0])

        # If only the intrinsics have changed we can use a simpler and quicker affine warp
        if (np.allclose(new_camera.R, old_camera.R) and
                allclose_or_nones(new_camera.distortion_coeffs, old_camera.distortion_coeffs)):
            relative_intrinsics = (
                    old_camera.intrinsic_matrix @ np.linalg.inv(new_camera.intrinsic_matrix))
            return cv2.warpAffine(
                image, relative_intrinsics[:2], output_size, flags=cv2.WARP_INVERSE_MAP)

        # If there are no new distortions we can use cv2.initUndistortRectifyMap
        if new_camera.distortion_coeffs is None:
            relative_rotation = new_camera.R @ old_camera.R.T
            map1, map2 = cv2.initUndistortRectifyMap(
                old_camera.intrinsic_matrix, old_camera.distortion_coeffs, relative_rotation,
                new_camera.intrinsic_matrix, output_size, cv2.CV_32FC1)
            return cv2.remap(image, map1, map2, cv2.INTER_LINEAR)

        # We must do the general case (i.e. new distortions) by hand
        y, x = np.mgrid[0:output_imshape[0], 0:output_imshape[1]].astype(np.float32)
        new_maps = np.stack([x, y], axis=-1)
        newim_coords = new_maps.reshape([-1, 2])
        world_coords = new_camera.image_to_world(newim_coords)
        oldim_coords = old_camera.world_to_image(world_coords)
        old_maps = oldim_coords.reshape(new_maps.shape)
        # For cv2.remap, we need to provide a grid of lookup pixel coordinates for
        # each output pixel.
        cv2.remap(image, old_maps, None, cv2.INTER_LINEAR)


def allclose_or_nones(a, b):
    if a is None and b is None:
        return True

    if a is None or b is None:
        return False

    return np.allclose(a, b)