80 lines
3.4 KiB
Python
80 lines
3.4 KiB
Python
import os
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "6"
|
|
import sys
|
|
import time
|
|
import cv2
|
|
import random
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
import cmapy as cm
|
|
import plotly.graph_objects as go
|
|
import plotly.express as ex
|
|
|
|
from dataset import DatasetFactory, StaticDataset
|
|
from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
|
|
from custom_utils.smpl import SMPL
|
|
from convert_to_uv_api import convert_to_uv_api
|
|
from tensorflow_graphics.geometry.representation.mesh.normals import vertex_normals
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
uv2mesh = UV2Mesh()
|
|
|
|
F = np.load("custom_utils/smpl_faces.npy")
|
|
|
|
factory = DatasetFactory()
|
|
signature_dynamic = { 'id': (), 'frame_nr': (),
|
|
# 'pose_shape': (72 + 11,),
|
|
'uv_posed': (512, 256, 3),
|
|
'uv_body_posed': (512, 256, 3),
|
|
'mesh_body_posed': (14475, 3,),
|
|
'mesh_posed': (14475, 3,),
|
|
'mesh_unposed': (14475, 3,),
|
|
'uv_unposed': (512, 256, 3),
|
|
'uv_curvature': (512, 256, 7),
|
|
# 'mesh_curvature': (14475)
|
|
}
|
|
signature_static = { 'uv_static': (512, 256, 3),
|
|
'mesh_static': (14475, 3,),
|
|
'mesh_mask': (14475,),
|
|
# 'uv_mask': (512, 256,),
|
|
# 'uv_shaped_body': (3, 512, 256,),
|
|
}
|
|
|
|
dataset = factory.get_dataset(folder='/home/workspace/local-data/train',
|
|
batch_size=1,
|
|
signature=signature_dynamic,
|
|
testing=True,
|
|
shuffle=True)
|
|
static_dataset = StaticDataset(folder='/home/workspace/local-data/train',
|
|
signature=signature_static)
|
|
|
|
|
|
uv_res = 256
|
|
|
|
|
|
|
|
for i, batch in enumerate(dataset):
|
|
print(i)
|
|
batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
|
|
|
|
mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
|
|
|
|
# vertex_normals(batch['mesh_posed'], tf.tile(F[None], [1, 1, 1]))
|
|
# vertex_normals(batch['mesh_body_posed'], tf.tile(F[None], [1, 1, 1]))
|
|
# vertex_normals(batch['mesh_unposed'], tf.tile(F[None], [1, 1, 1]))
|
|
# vertex_normals(batch['mesh_static'], tf.tile(F[None], [1, 1, 1]))
|
|
|
|
|
|
uv_norm_posed = (convert_to_uv_api(np.squeeze(vertex_normals(batch['mesh_posed'], tf.tile(F[None], [1, 1, 1])).numpy()) * mask[:,np.newaxis], mask) + 2) / 4
|
|
uv_norm_body_posed = (convert_to_uv_api(np.squeeze(vertex_normals(batch['mesh_body_posed'], tf.tile(F[None], [1, 1, 1])).numpy()), np.ones((14475), bool)) + 2) / 4
|
|
uv_norm_unposed = (convert_to_uv_api(np.squeeze(vertex_normals(batch['mesh_unposed'], tf.tile(F[None], [1, 1, 1])).numpy()) * mask[:,np.newaxis], mask) + 2) / 4
|
|
uv_curvature = np.tile(np.argmax(np.squeeze(batch['uv_curvature'].numpy()), axis=-1)[:,:,np.newaxis], [1,1,3]) / 6
|
|
uv_norm_static = (convert_to_uv_api(np.squeeze(vertex_normals(batch['mesh_static'], tf.tile(F[None], [1, 1, 1])).numpy()) * mask[:,np.newaxis], mask) + 2) / 4
|
|
|
|
img = np.concatenate([uv_norm_posed, uv_norm_body_posed, uv_norm_unposed, uv_curvature, uv_norm_static], axis=1)
|
|
|
|
ex.imshow(img).show()
|
|
|
|
print() |