import k3d
import numpy as np
from k3d.helpers import download
# credits: https://github.com/eladrich/pix2vertex.pytorch
pix2vertex_file = download('http://k3d-jupyter.org/DATA/pix2vertex_output.npy')
pix2vertex_image_file = download('http://k3d-jupyter.org/DATA/pix2vertex.jpg')
data = np.load(pix2vertex_file, allow_pickle=True).tolist()
image = open(pix2vertex_image_file, 'rb').read()
heights = data['Z_surface'].T
where = np.where(~np.isnan(heights))
vertices = np.dstack(where)[0].astype(np.float32)
vertices = np.append(vertices, np.zeros((vertices.shape[0], 1), dtype=np.float32), 1)
vertices[:,2] = heights[where]
pic2index = np.zeros((heights.shape[0], heights.shape[1]), dtype=np.int64)
pic2index[where] = np.arange(vertices.shape[0], dtype=np.int64)
#top-left triangle
tlt = np.where(pic2index[0:-2, 0:-2] * pic2index[1:-1,0:-2] * pic2index[0:-2, 1:-1])
#bottom-right triangle
brt = np.where(pic2index[1:-1, 1:-1] * pic2index[0:-2,1:-1] * pic2index[1:-1, 0:-2])
indices = np.dstack([pic2index[tlt],
pic2index[(tlt[0] + 1, tlt[1])],
pic2index[(tlt[0], tlt[1] + 1)]
])[0]
indices = np.append(indices,
np.dstack([pic2index[(brt[0] + 1, brt[1] + 1)],
pic2index[(brt[0], brt[1] + 1)],
pic2index[(brt[0] + 1, brt[1])]])[0])
uvs = vertices[:, 0:2].copy()
uvs[:, 0] /= heights.shape[0]
uvs[:, 1] /= heights.shape[1]
plot = k3d.plot()
obj = k3d.mesh(vertices.astype(np.float32), indices.astype(np.uint32),
flat_shading=False, color=0xb2ccff, side='double')
plot += obj
plot.display()
plot.camera = [126, 295, 413,
209, 319,-12,
0, -1, 0]
obj.uvs = uvs
obj.texture = image
obj.texture_file_format = 'jpg'
obj.color = 0xffffff
plot.lighting = 0