65 lines
3.4 KiB
Python
65 lines
3.4 KiB
Python
import numpy as np
|
|
import tensorflow as tf
|
|
|
|
|
|
def m2uv_array(m2uv):
|
|
result = np.zeros((14475, 2), dtype=np.int32)
|
|
for k, v in m2uv.items():
|
|
if len(v) == 2:
|
|
result[k] = list(v)
|
|
elif len(v) == 1:
|
|
x = v.pop()
|
|
result[k] = [x, x]
|
|
else:
|
|
assert False
|
|
|
|
return result
|
|
|
|
|
|
class UV2Mesh():
|
|
def __init__(self, res=256):
|
|
super(UV2Mesh, self).__init__()
|
|
self.res = res
|
|
uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
|
|
self.uv_map_discreet = np.round(uv_map * self.res).astype(np.int32)
|
|
m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
|
|
m2uv = m2uv_array(m2uv)
|
|
self.uv_pixels = self.uv_map_discreet[m2uv]
|
|
# self.batch_size = batch_size
|
|
|
|
# self.x1 = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.zeros((self.batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
|
|
# self.y1 = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
|
|
# self.z1 = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32) * 2), axis=-1)
|
|
# self.x2 = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.zeros((self.batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
|
|
# self.y2 = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
|
|
# self.z2 = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32) * 2), axis=-1)
|
|
|
|
# input = channels_first
|
|
def uv_to_mesh(self, uv):
|
|
uv = tf.transpose(uv, [0,2,3,1])
|
|
batch_size = uv.shape[0]
|
|
|
|
x1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
|
|
y1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
|
|
z1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, 14475, 1), dtype=tf.int32) * 2), axis=-1)
|
|
x2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
|
|
y2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
|
|
z2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, 14475, 1), dtype=tf.int32) * 2), axis=-1)
|
|
|
|
|
|
x1 = tf.gather_nd(uv, x1_ind, batch_dims=1)
|
|
y1 = tf.gather_nd(uv, y1_ind, batch_dims=1)
|
|
z1 = tf.gather_nd(uv, z1_ind, batch_dims=1)
|
|
x2 = tf.gather_nd(uv, x2_ind, batch_dims=1)
|
|
y2 = tf.gather_nd(uv, y2_ind, batch_dims=1)
|
|
z2 = tf.gather_nd(uv, z2_ind, batch_dims=1)
|
|
|
|
p1 = tf.stack((x1, y1, z1), axis=-1)
|
|
p2 = tf.stack((x2, y2, z2), axis=-1)
|
|
|
|
mesh = (p1 + p2) / 2
|
|
return mesh
|
|
|
|
def __call__(self, uv):
|
|
return self.uv_to_mesh(uv)
|
|
|