convert_panorama_exr.py

import cv2
import numpy as np
import argparse
import math
import OpenEXR, Imath

def load_exr(in_file):
    pt = Imath.PixelType(Imath.PixelType.FLOAT)
    golden = OpenEXR.InputFile(in_file)
    dw = golden.header()['dataWindow']
    size = (dw.max.x - dw.min.x + 1, dw.max.y - dw.min.y + 1)

    redstr = golden.channel('R', pt)
    red = np.fromstring(redstr, dtype = np.float32)
    red.shape = (size[1], size[0]) # Numpy arrays are (row, col)

    greenstr = golden.channel('G', pt)
    green = np.fromstring(greenstr, dtype = np.float32)
    green.shape = (size[1], size[0]) # Numpy arrays are (row, col)

    bluestr = golden.channel('B', pt)
    blue = np.fromstring(bluestr, dtype = np.float32)
    blue.shape = (size[1], size[0]) # Numpy arrays are (row, col)

    img = np.zeros((size[1], size[0], 3), dtype=np.float32)
    img[:, :, 0] = red
    img[:, :, 1] = green
    img[:, :, 2] = blue

    return img

def write_exr(out_file, data):
    exr = OpenEXR.OutputFile(out_file, OpenEXR.Header(data.shape[1], data.shape[0]))
    red = data[:, :, 0]
    green = data[:, :, 1]
    blue = data[:, :, 2]
    exr.writePixels({'R': red.tostring(), 'G': green.tostring(), 'B': blue.tostring()})

def convert_panorama_image(img, theta=0.0, phi=0.0, move=0.5, res_x=400, res_y=800, debug=False):
    img_x = img.shape[0]
    img_y = img.shape[1]

    theta = theta / 180 * math.pi
    phi = phi / 180 * math.pi

    mid_x = res_x / 2
    mid_y = res_y / 2

    map_x = np.zeros((res_x, res_y), dtype=np.float32)
    map_y = np.zeros((res_x, res_y), dtype=np.float32)

    axis_y = math.cos(theta)
    axis_z = math.sin(theta)
    axis_x = 0

    # theta rotation matrix
    cos_theta = math.cos(theta)
    sin_theta = math.sin(theta)
    theta_rot_mat = np.array([[1, 0, 0], \
            [0, cos_theta, -sin_theta], \
            [0, sin_theta, cos_theta]], dtype=np.float32)

    # phi rotation matrix
    cos_phi = math.cos(phi)
    sin_phi = -math.sin(phi)
    phi_rot_mat = np.array([[cos_phi + axis_x**2 * (1 - cos_phi), \
            axis_x * axis_y * (1 - cos_phi) - axis_z * sin_phi, \
            axis_x * axis_z * (1 - cos_phi) + axis_y * sin_phi], \
            [axis_y * axis_x * (1 - cos_phi) + axis_z * sin_phi, \
            cos_phi + axis_y**2 * (1 - cos_phi), \
            axis_y * axis_z * (1 - cos_phi) - axis_x * sin_phi], \
            [axis_z * axis_x * (1 - cos_phi) - axis_y * sin_phi, \
            axis_z * axis_y * (1 - cos_phi) + axis_x * sin_phi, \
            cos_phi + axis_z**2 * (1 - cos_phi)]], dtype=np.float32)

    indx = np.tile(np.array(np.arange(res_x), dtype=np.float32), (res_y, 1)).T
    indy = np.tile(np.array(np.arange(res_y), dtype=np.float32), (res_x, 1))

    map_x = np.sin(indx * math.pi / res_x - math.pi / 2)
    map_y = np.sin(indy * (2 * math.pi)/ res_y) * np.cos(indx * math.pi / res_x - math.pi / 2)
    map_z = -np.cos(indy * (2 * math.pi)/ res_y) * np.cos(indx * math.pi / res_x - math.pi / 2)

    ind = np.reshape(np.concatenate((np.expand_dims(map_x, 2), np.expand_dims(map_y, 2), \
            np.expand_dims(map_z, 2)), axis=2), [-1, 3]).T

    move_dir = np.array([0, 0, -1], dtype=np.float32)
    move_dir = theta_rot_mat.dot(move_dir)
    move_dir = phi_rot_mat.dot(move_dir)

    ind = theta_rot_mat.dot(ind)
    ind = phi_rot_mat.dot(ind)

    ind += np.tile(move * move_dir, (ind.shape[1], 1)).T

    vec_len = np.sqrt(np.sum(ind**2, axis=0))
    ind /= np.tile(vec_len, (3, 1))

    cur_phi = np.arcsin(ind[0, :])
    cur_theta = np.arctan2(ind[1, :], -ind[2, :])

    map_x = (cur_phi + math.pi/2) / math.pi * img_x
    map_y = cur_theta % (2 * math.pi) / (2 * math.pi) * img_y

    map_x = np.reshape(map_x, [res_x, res_y])
    map_y = np.reshape(map_y, [res_x, res_y])

    if debug:
        for x in range(res_x):
            for y in range(res_y):
                print '(%.2f, %.2f)\t' % (map_x[x, y], map_y[x, y]),
            print

    return cv2.remap(img, map_y, map_x, cv2.INTER_LINEAR, borderMode=cv2.BORDER_WRAP)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('input_image', type=str, help='Input Panorama Image')
    parser.add_argument('output_image', type=str, help='Output Panorama Image')
    parser.add_argument('--theta', type=float, default=0.0, help='Theta angle (yaw) in range [-180, 180] degrees [default: 0.0]')
    parser.add_argument('--phi', type=float, default=0.0, help='Phi angle (pitch) in range [-90, 90) degrees [default: 0.0]')
    parser.add_argument('--resolution_x', type=int, default=800, help='Resolution of the output image width [default: 800]')
    parser.add_argument('--resolution_y', type=int, default=400, help='Resolution of the output image height [default: 400]')
    parser.add_argument('--move', type=float, default=0.5, help='Move forward a bit [default: 0.5]')
    parser.add_argument('--debug', type=bool, default=False, help='Debug mode')
    args = parser.parse_args()

    img = load_exr(args.input_image)
    out_img = convert_panorama_image(img, theta=args.theta, phi=args.phi, res_x=args.resolution_x, \
            res_y=args.resolution_y, move=args.move, debug=args.debug)
    write_exr(args.output_image, out_img)