145 lines
5.1 KiB
Python
145 lines
5.1 KiB
Python
import os
|
|
import sys
|
|
import time
|
|
from unicodedata import name
|
|
import cv2
|
|
import random
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
|
|
from dataset import DatasetFactory, StaticDataset
|
|
from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
|
|
import plotly.graph_objects as go
|
|
import plotly.express as ex
|
|
import cmapy as cm
|
|
|
|
if __name__ == '__main__':
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "6"
|
|
|
|
uv2mesh = UV2Mesh()
|
|
|
|
F = np.load("custom_utils/smpl_faces.npy")
|
|
|
|
factory = DatasetFactory()
|
|
signature_dynamic = { 'id': (), 'frame_nr': (),
|
|
# 'pose_shape': (72 + 11,),
|
|
# 'uv_posed': (3, 512, 256,),
|
|
# 'uv_body_posed': (3, 512, 256,),
|
|
# 'mesh_body_posed': (14475, 3,),
|
|
# 'mesh_posed': (14475, 3,),
|
|
'mesh_unposed': (14475, 3,),
|
|
'uv_unposed': (3, 512, 256,),
|
|
# 'uv_curvature': (512, 256),
|
|
# 'mesh_curvature': (14475)
|
|
}
|
|
signature_static = { #'uv_static': (3, 512, 256,),
|
|
# '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/jumpsuit_20/train',
|
|
batch_size=1,
|
|
signature=signature_dynamic)
|
|
static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
|
|
signature=signature_static)
|
|
|
|
|
|
uv_res = 256
|
|
|
|
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
|
|
|
|
def unique_uv_map_mask(uv_map_discreet):
|
|
mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool)
|
|
temp = np.zeros((uv_res * 2, uv_res))
|
|
for p in uv_map_discreet:
|
|
temp[p[0], p[1]] += 1
|
|
# temp = temp == 1
|
|
for i, p in enumerate(uv_map_discreet):
|
|
if temp[p[0], p[1]] == 1:
|
|
mask[i] = True
|
|
elif temp[p[0], p[1]] == 2:
|
|
mask[i] = True
|
|
temp[p[0], p[1]] = 0
|
|
return mask
|
|
|
|
def uv_to_mesh(uv, m2uv):
|
|
res = uv.shape[0]
|
|
mesh = np.zeros((14475, 3))
|
|
uv_pixels = uv_map_discreet[m2uv]
|
|
# counter = np.ones((14475))
|
|
# counter[m2uv[:, 0] == m2uv[:, 1]] = 2
|
|
mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis]
|
|
return mesh
|
|
|
|
|
|
uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
|
|
m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
|
|
m2uv = m2uv_array(m2uv)
|
|
uv_mask = np.load('custom_utils/mask_256x512.npy')
|
|
curvature_label_type = '6_25'
|
|
|
|
uv_res = 256
|
|
uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
|
|
uv_mask = cv2.dilate(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis]
|
|
|
|
uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
|
|
uv_map_mask = unique_uv_map_mask(uv_map_discreet)
|
|
uv_counter = np.zeros((uv_map.shape[0]))
|
|
uv_counter[m2uv[:, 0]] = 1
|
|
uv_counter[m2uv[:, 1]] += 1
|
|
uv_counter = uv_counter[uv_map_mask]
|
|
|
|
|
|
|
|
for i, batch in enumerate(dataset):
|
|
print(i)
|
|
batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
|
|
|
|
uv_unposed = np.squeeze(batch['uv_unposed'].numpy())
|
|
mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
|
|
mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
|
|
|
|
|
|
values = mesh_unposed
|
|
uv_values = np.zeros((uv_map.shape[0], 3))
|
|
uv_values[m2uv[:, 0]] = values
|
|
uv_values[m2uv[:, 1]] += values
|
|
uv_counter = np.zeros((uv_map.shape[0]))
|
|
uv_counter[m2uv[:, 0]] = 1
|
|
uv_counter[m2uv[:, 1]] += 1
|
|
uv_values /= uv_counter[:, np.newaxis]
|
|
x = uv_values
|
|
# verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1)
|
|
im = ((x - x.min()) / (x.max() - x.min())) * 255
|
|
|
|
|
|
x = np.transpose(uv_unposed, [1,2,0])
|
|
x = uv2mesh(uv_unposed[tf.newaxis,:,:,:])[0].numpy()[mask, :]
|
|
_, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
|
|
go.Figure(data=[
|
|
go.Mesh3d(
|
|
# 8 vertices of a cube
|
|
x=x[:,0],
|
|
y=x[:,1],
|
|
z=x[:,2],
|
|
# i, j and k give the vertices of triangles
|
|
i = faces[:,0],
|
|
j = faces[:,1],
|
|
k = faces[:,2]
|
|
)]).show()
|
|
|
|
|
|
print() |