project stuff onto a cone using tensorflow from krall

project_to_cone.py

#!/usr/bin/env python3

import numpy as np
import tensorflow as tf
import scipy.misc
import scipy.ndimage
from tensorflow import sin, cos, sqrt, atan2
from math import pi

top_circ = 265.
bottom_circ = 236.
d_top_to_bottom = 79.

r = top_circ / (2 * pi)
R = d_top_to_bottom * top_circ / (top_circ - bottom_circ)
h = sqrt(R**2 - r**2)

print(f"R = {R}")
print(f"r = {r}")
print(f"h = {h}")


def phi_d_to_x_y_z(phi, d):
    x = r * cos(phi) * (R - d) / R
    y = r * sin(phi) * (R - d) / R
    z = h * d / R

    return x, y, z


def phi_d_to_u_v(phi, d):
    theta = phi * top_circ / (2 * pi * R)
    u = (R - d) * cos(theta)
    v = (R - d) * sin(theta)
    return u, v


def u_v_to_phi_d(u, v):
    theta = atan2(v, u)
    phi = theta * R * 2 * pi / top_circ
    d = R - hypot(u, v)
    return phi, d


def hypot(a, b):
    return tf.norm(tf.stack([a, b], axis=0), axis=0)


source_image = scipy.ndimage.imread('burst_outline.png').astype(np.float32) / 65535.0
# source_image = np.zeros([12500, 12500])

dpi = 600
width_inches = 8.5
height_inches = 11

output_image = tf.ones([int(round(width_inches * dpi)), int(round(height_inches * dpi))])
print(f"output_image.shape: {output_image.shape}")

u, v = tf.meshgrid(
    tf.linspace(R - width_inches * 25.4, R, output_image.shape[0] + 1)[1:],
    tf.linspace(0.0, height_inches * 25.4, output_image.shape[1] + 1)[:-1],
    indexing='ij',
)
# print(f"uvs.shape: {uvs.shape}")

phi, d = u_v_to_phi_d(u, v)
x, y, z = phi_d_to_x_y_z(phi, d)


def x_y_z_to_source_coords(x, y, z, source_image_shape, center_x_y_z, src_x_norm, src_y_norm):
    center_x, center_y, center_z = center_x_y_z

    x = x - center_x
    y = y - center_y
    z = z - center_z

    xyz = tf.stack([x, y, z], axis=-1)
    print(f"xyz.shape: {xyz.shape}")

    src_z_norm = np.cross(src_x_norm, src_y_norm)
    src_xyz_to_xyz = np.stack([src_x_norm, src_y_norm, src_z_norm])
    print(src_xyz_to_xyz)
    xyz_to_src_xyz = np.linalg.inv(src_xyz_to_xyz)
    print(xyz_to_src_xyz)

    src_xyz = tf.tensordot(xyz, xyz_to_src_xyz.astype(np.float32), axes=1)
    print(f"retval.shape: {src_xyz.shape}")
    src_xy = tf.clip_by_value(src_xyz[:, :, :2], -1, 1.) * 0.5 + 0.5

    src_xy = src_xy * (tf.convert_to_tensor(source_image_shape, dtype=tf.float32) - 1)
    return src_xy


mask = tf.logical_and(tf.logical_and((phi <= 2*pi), (d >= 0.0)), (d <= d_top_to_bottom))

angle = 4 * pi / 3
scale1 = top_circ / pi
src_x_norm = np.array([0., 0., scale1])
src_y_norm = np.array([np.sin(angle) * scale1, np.cos(angle) * scale1, 0.])
scale2 = 1.7
src_xy = x_y_z_to_source_coords(
    x, y, z,
    source_image.shape,
    center_x_y_z=0.3 * src_y_norm + 0.2 * src_x_norm,
    src_x_norm=src_x_norm * scale2,
    src_y_norm=src_y_norm * scale2,
)
output_image *= tf.gather_nd(params=source_image, indices=tf.cast(src_xy, tf.int32))

angle = 0
scale1 = top_circ / pi
src_x_norm = np.array([0., 0., scale1])
src_y_norm = np.array([np.sin(angle) * scale1, np.cos(angle) * scale1, 0.])
scale2 = 0.7
src_xy = x_y_z_to_source_coords(
    x, y, z,
    source_image.shape,
    center_x_y_z=-0.4 * src_y_norm + 0.7 * src_x_norm,
    src_x_norm=src_x_norm * scale2,
    src_y_norm=src_y_norm * scale2,
)
output_image *= tf.gather_nd(params=source_image, indices=tf.cast(src_xy, tf.int32))

angle = pi/4
scale1 = top_circ / pi
src_x_norm = np.array([0., 0., scale1])
src_y_norm = np.array([np.sin(angle) * scale1, np.cos(angle) * scale1, 0.])
scale2 = 0.5
src_xy = x_y_z_to_source_coords(
    x, y, z,
    source_image.shape,
    center_x_y_z=0.3 * src_y_norm + 0.5 * src_x_norm,
    src_x_norm=src_x_norm * scale2,
    src_y_norm=src_y_norm * scale2,
)
output_image *= tf.gather_nd(params=source_image, indices=tf.cast(src_xy, tf.int32))


with tf.Session() as sess:
    (
        _u,
        _v,
        _phi,
        _d,
        _mask,
        _src_xy,
        _output_image,
    ) = sess.run([
        u,
        v,
        phi,
        d,
        mask,
        src_xy,
        output_image,
    ])
    print(f"_u: {_u}")
    print(f"_v: {_v}")
    print(f"_phi: {_phi}")
    print(f"_d: {_d}")
    print(f"_mask: {_mask}")
    print(f"_src_xy: {_src_xy}")
    print(f"_output_image: {_output_image}")

    imgdata = np.stack([
        # _src_xy[:, :, 0],
        _output_image * 255,
        # _src_xy[:, :, 1],
        _output_image * 255,
        # np.zeros_like(_src_xy[:, :, 0]),
        _output_image * 255,
        _mask.astype(np.float32) * 255,
    ], axis=-1)

    print(f"imgdata.shape: {imgdata.shape}")

    scipy.misc.imsave(
        'mask.png',
        _output_image * _mask.astype(np.float32),
        # imgdata.astype(np.float32),
        format='png',
    )