84 lines
3.2 KiB
Python
84 lines
3.2 KiB
Python
import os
|
|
import sys
|
|
import time
|
|
import cv2
|
|
import random
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
|
|
class DatasetSimple(tf.data.Dataset):
|
|
|
|
def _generator(folder):
|
|
# Opening the file
|
|
folder = folder.decode("utf-8")
|
|
uv_folder = os.path.join(folder, 'uv_posed')
|
|
pose_hape_folder = os.path.join(folder, 'poses_shapes')
|
|
frames = os.listdir(uv_folder)
|
|
random.shuffle(frames)
|
|
|
|
for frame in frames:
|
|
id, frame_nr = frame.split('.')[0].split('_')
|
|
uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']; uv_posed = np.transpose(uv_posed, [2, 0, 1])
|
|
pose_shape = np.load(os.path.join(folder, 'pose_shape', id + '_' + frame_nr + '.npy'))
|
|
uv_static = np.load(os.path.join(folder, 'uv_static', id + '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1])
|
|
mesh_static = np.load(os.path.join(folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']
|
|
mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
|
|
mask = np.load(os.path.join(folder, 'mask', id + '.npy')); mask = mask.astype(np.float32)
|
|
|
|
# uv_posed, (pose, shape, uv_static), mesh_static, mesh_body_posed, mask
|
|
yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static), mesh_static, mesh_body_posed, mask
|
|
|
|
def __new__(cls, folder, res=256):
|
|
return tf.data.Dataset.from_generator(
|
|
cls._generator,
|
|
output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), (tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475,), dtype = tf.float32) ),
|
|
args=(folder, )
|
|
)
|
|
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').batch(2)
|
|
|
|
from losses import collision_loss as collision_penalty
|
|
from losses import collision_penalty as collision_penalty_2
|
|
from custom_utils.custom_utils import UV2Mesh
|
|
import matplotlib.pyplot as plt
|
|
|
|
def plot_3d(x, y=None):
|
|
fig = plt.figure(figsize=(15, 15))
|
|
ax = fig.add_subplot(111, projection='3d')
|
|
ax.set_xlim3d(-1, 1)
|
|
ax.set_ylim3d(-1, 1)
|
|
ax.set_zlim3d(-1, 1)
|
|
ax.set_xlabel("X")
|
|
ax.set_ylabel("Y")
|
|
ax.set_zlabel("Z")
|
|
ax.grid(False)
|
|
ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5)
|
|
if y is not None:
|
|
ax.scatter(y[:, 0], y[:, 1], y[:, 2], s=5)
|
|
plt.savefig('temp.jpg')
|
|
|
|
uv2mesh = UV2Mesh()
|
|
|
|
for d in dataset:
|
|
d
|
|
|
|
v = uv2mesh(d[0])
|
|
b = d[3]
|
|
m = d[4]
|
|
|
|
b2 = tf.concat((b[1][tf.newaxis,:,:], b[0][tf.newaxis,:,:]), axis=0)
|
|
|
|
c = collision_penalty(v, b, m)
|
|
c2 = collision_penalty(v, b2, m)
|
|
|
|
c_ = collision_penalty_2(v, b, m)
|
|
c2_ = collision_penalty_2(v, b2, m)
|
|
|
|
print()
|
|
pass
|
|
|