2.5D image rotation

project3d.py

"""
Code for 2.5D rotation
"""

import cv2
from PIL import Image

import os
import numpy
import numpy as np
import argparse


def rotate_2d_image(image, depth_map, angle):
    orig_depth_map = depth_map
    # Convert depth map is in range[0, 255], scale it and shift points by "focal" distance
    depth_map = (1 - depth_map.astype(np.float32) / 255.0) * 150 + 50

    # Calculate the scene center
    scene_center = (image.shape[1] // 2, image.shape[0] // 2)

    # Convert angle to radians
    th = np.radians(angle)

    width, height = image.shape[1], image.shape[0]
    ar = np.zeros_like(image, dtype=np.float32)

    # Create 3D points with center at the image center
    ar[:, :, 0] = np.tile(np.arange(width) - width / 2, (height, 1))
    ar[:, :, 1] = np.flip(np.tile(np.arange(height) - height / 2, (width, 1)).T, axis=0)
    ar[:, :, 2] = depth_map

    # Rotation matrix around y-axis
    R = np.array([[np.cos(th), 0, np.sin(th)],
                  [0, 1, 0],
                  [-np.sin(th), 0, np.cos(th)]])

    # Rotate the 3D points
    rotated_ar = np.matmul(ar.reshape(-1, 3), R.T).reshape(ar.shape)

    # Perform projection and get the new image
    new_img = numpy.zeros_like(image)
    max_d = 1000000
    new_d = np.ones_like(depth_map) * max_d
    result_d = np.zeros_like(orig_depth_map)
    mask = np.zeros_like(depth_map, dtype=numpy.uint8)
    for x in range(width):
        for y in range(height):
            x0, y0, d0 = ar[y, x] 
            x1, y1, d1 = rotated_ar[y, x]
            pix_y0 = round(height // 2 - y0)
            pix_x0 = round(x0 + width // 2)
            pix_y1 = round(height // 2 - y1)
            pix_x1 = round(x1 + width // 2)
            if 0 <= pix_x1 and pix_x1 < width and 0 <= pix_y1 and pix_y1 < height:
                if d1 < new_d[pix_y1, pix_x1]:
                    new_img[pix_y1, pix_x1] = image[pix_y0, pix_x0]
                    new_d[pix_y1, pix_x1] = d1
                    mask[pix_y1, pix_x1] = 255
                    result_d[pix_y1, pix_x1] = orig_depth_map[pix_y0, pix_x0]
    return new_img, mask, result_d


def parse_args():
    parser = argparse.ArgumentParser(description='Pseudo 3D Image Rotation')
    parser.add_argument('image', type=str, help='path to the source image')
    parser.add_argument('depth', type=str, help='path to the depth map')
    parser.add_argument('angle', type=float, help='rotation angle in degrees')
    parser.add_argument('--postfix', type=str, help='postfix for saving the results', default='_res', required=False)
    args = parser.parse_args()
    return args

def main():
    args = parse_args()
    
    # Load the input image and depth map
    image = cv2.imread(args.image, cv2.IMREAD_COLOR)  # Load the image
    depth_map = cv2.imread(args.depth, cv2.IMREAD_GRAYSCALE)  # Load 
    
    # Define the rotation angle
    angle = args.angle
    # Perform the pseudo 3D rotation
    output_image, mask, new_d = rotate_2d_image(np.asarray(image), depth_map, angle)
    
    
    # Save the output images with the specified postfix
    base_name = os.path.splitext(os.path.basename(args.image))[0]
    output_image_path = base_name + '_' + args.postfix + '.jpg'
    output_depth_map_path = base_name + '_' + args.postfix + '_depth.png'
    output_mask_path = base_name + '_' + args.postfix + '_mask.png'

    # Save the output image
    cv2.imwrite(output_image_path, output_image)
    cv2.imwrite(output_mask_path, mask)
    cv2.imwrite(output_depth_map_path, new_d)


if __name__ == '__main__':
    main()