309 lines
13 KiB
Python
309 lines
13 KiB
Python
import os
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "6"
|
|
import sys
|
|
import time
|
|
import cv2
|
|
import random
|
|
import tqdm
|
|
import numpy as np
|
|
import trimesh as tm
|
|
import tensorflow as tf
|
|
import cmapy as cm
|
|
import plotly.graph_objects as go
|
|
import plotly.express as ex
|
|
from plotly.subplots import make_subplots
|
|
|
|
from dataset import DatasetFactory, StaticDataset
|
|
from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces, load_model
|
|
from custom_utils.smpl import SMPL
|
|
from convert_to_uv_api import convert_to_uv_api
|
|
from models.upscaler import UpscalerBasicConv as Upscaler
|
|
|
|
|
|
|
|
def color_labels(img):
|
|
colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, img.max()+1)]) / 255.0
|
|
color_img = np.zeros((img.shape[0], img.shape[1], 3), np.float32)
|
|
for i in range(img.max()+1):
|
|
color_img[img == i, :] = colors[i]
|
|
|
|
return color_img
|
|
|
|
def write_obj(file, vertices, faces, colors):
|
|
with open(file, 'w') as file:
|
|
for v, c in zip(vertices, colors):
|
|
line = 'v ' + ' '.join([str(_) for _ in v]) + ' ' + ' '.join([str(_) for _ in c]) + '\n'
|
|
file.write(line)
|
|
for f in faces:
|
|
line = 'f ' + ' '.join([str(_) for _ in f]) + '\n'
|
|
file.write(line)
|
|
|
|
def write_obj_2(file, vertices, faces):
|
|
with open(file, 'w') as file:
|
|
for v in vertices:
|
|
line = 'v ' + ' '.join([str(_) for _ in v]) + '\n'
|
|
file.write(line)
|
|
for f in faces:
|
|
line = 'f ' + ' '.join([str(_) for _ in f]) + '\n'
|
|
file.write(line)
|
|
|
|
def export_to_ply(mesh, file):
|
|
ply = tm.exchange.ply.export_ply(mesh, encoding='ascii', vertex_normal=True, include_attributes=True)
|
|
output_file = open(file, "wb+")
|
|
output_file.write(ply)
|
|
output_file.close()
|
|
|
|
|
|
def export_to_obj(mesh, file):
|
|
ply = tm.exchange.obj.export_obj(mesh)
|
|
output_file = open(file, "w+")
|
|
output_file.write(ply)
|
|
output_file.close()
|
|
|
|
|
|
def plot_fixed(mesh, faces, colors):
|
|
fig = go.Figure()
|
|
|
|
fig.add_trace(go.Mesh3d(
|
|
x=mesh[:,0],
|
|
y=mesh[:,2],
|
|
z=mesh[:,1],
|
|
i = faces[:,0],
|
|
j = faces[:,1],
|
|
k = faces[:,2],
|
|
text="",
|
|
vertexcolor=colors,
|
|
lighting={ 'ambient':0.8, 'diffuse':0.5, 'roughness':0.5 },
|
|
|
|
),)
|
|
|
|
|
|
|
|
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1))
|
|
helper = dict(xaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
|
|
yaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
|
|
zaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),)
|
|
fig.update_layout(scene = helper)
|
|
|
|
camera = dict(
|
|
up=dict(x=0, y=0, z=1),
|
|
center=dict(x=0, y=0, z=-0.05),
|
|
eye=dict(x=0.3, y=1.0, z=0.05)
|
|
)
|
|
|
|
fig.update_layout(scene_camera=camera, width=800, height=1100, autosize=False,
|
|
margin=dict(l=0, r=0, b=0, t=0, pad=0)
|
|
)
|
|
|
|
return fig
|
|
|
|
|
|
def custom_imshow(img):
|
|
fig = ex.imshow(np.flip(img, axis=0))
|
|
fig.update_layout(margin=dict(l=0, r=0, b=0, t=0, pad=0), width=img.shape[1], height=img.shape[0])
|
|
fig.update_xaxes(visible=False)
|
|
fig.update_yaxes(visible=False)
|
|
|
|
return fig
|
|
|
|
|
|
def uv_mapping_example(mesh, faces, colors, mask):
|
|
fig = plot_fixed(mesh, faces, colors)
|
|
fig.write_image('examples/' + type + '_' + id_str + '_' + frame_nr + '_body.png')
|
|
|
|
# stop inbetween
|
|
uv = convert_to_uv_api(mesh, mask)
|
|
plot = custom_imshow((uv+1)/2)
|
|
# plot = custom_imshow(result)
|
|
# plot.write_image('examples/' + type + '_' + id_str + '_' + frame_nr + '_uv_body_pre.png')
|
|
plot.write_image('examples/' + type + '_' + id_str + '_' + frame_nr + '_uv_body_inpainted.png')
|
|
|
|
print()
|
|
|
|
|
|
def curvature_example(mesh_curvature, uv_curvature, posed, faces, mask):
|
|
colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, mesh_curvature.max()+1)]) / 255.0
|
|
color_mesh = np.zeros((mesh_curvature.shape[0], 3), np.float32)
|
|
for i in range(mesh_curvature.max()+1):
|
|
color_mesh[mesh_curvature == i, :] = colors[i]
|
|
|
|
curvature_raw = np.load('/home/workspace/data/dataset/test/' + id_str + '/curvature_25_raw.npy')[int(frame_nr), :]
|
|
curvature_raw_norm = (curvature_raw - curvature_raw.min()) / (curvature_raw.max() - curvature_raw.min())
|
|
curvature_raw_colors = np.array([cm.color('jet', 1 - curvature_raw_norm[i]) for i in range(curvature_raw_norm.shape[0])]) / 255.0
|
|
|
|
fig = plot_fixed(posed, faces, curvature_raw_colors)
|
|
fig.write_image('examples/' + type + '_' + id_str + '_' + frame_nr + '_mesh_curv_raw.png')
|
|
fig = plot_fixed(posed, faces, color_mesh[mask, :])
|
|
fig.write_image('examples/' + type + '_' + id_str + '_' + frame_nr + '_mesh_curv.png')
|
|
plot = custom_imshow(uv_curvature)
|
|
plot.write_image('examples/' + type + '_' + id_str + '_' + frame_nr + '_uv_curv.png')
|
|
|
|
print()
|
|
|
|
|
|
def plot_upscaling(uv_low, uv_high, uv_mask_low, uv_mask_high, faces, color, mask):
|
|
uv2mesh = UV2Mesh(res=256)
|
|
model = Upscaler(low_res=32, kernel_size=3)
|
|
load_model(model, 'checkpoints/32_512_4')
|
|
enhanced = model([uv_low[np.newaxis,:,:,:], uv_mask_low[np.newaxis,:,:]], training=False)[0].numpy()
|
|
|
|
bad = tf.image.resize(uv_low, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5).numpy()
|
|
good = tf.image.resize(enhanced, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5).numpy()
|
|
|
|
plot = custom_imshow((bad+1)/2 * uv_mask_high[:,:,np.newaxis])
|
|
plot.write_image('examples/' + type + '_' + id_str + '_' + frame_nr + '_uv_upscale_bad.png')
|
|
plot = custom_imshow((good+1)/2 * uv_mask_high[:,:,np.newaxis])
|
|
plot.write_image('examples/' + type + '_' + id_str + '_' + frame_nr + '_uv_upscale_good.png')
|
|
plot = custom_imshow((uv_high+1)/2 * uv_mask_high[:,:,np.newaxis])
|
|
plot.write_image('examples/' + type + '_' + id_str + '_' + frame_nr + '_uv_upscale_orig.png')
|
|
plot = custom_imshow((uv_low+1)/2)
|
|
plot.write_image('examples/' + type + '_' + id_str + '_' + frame_nr + '_uv_low.png')
|
|
|
|
fig = plot_fixed(uv2mesh(bad[np.newaxis,:,:,:])[0].numpy()[mask, :], faces, color)
|
|
fig.write_image('examples/' + type + '_' + id_str + '_' + frame_nr + '_mesh_upscale_bad.png')
|
|
fig = plot_fixed(uv2mesh(good[np.newaxis,:,:,:])[0].numpy()[mask, :], faces, color)
|
|
fig.write_image('examples/' + type + '_' + id_str + '_' + frame_nr + '_mesh_upscale_good.png')
|
|
fig = plot_fixed(uv2mesh(uv_high[np.newaxis,:,:,:])[0].numpy()[mask, :], faces, color)
|
|
fig.write_image('examples/' + type + '_' + id_str + '_' + frame_nr + '_mesh_upscale_orig.png')
|
|
|
|
print()
|
|
|
|
if __name__ == '__main__':
|
|
|
|
uv2mesh = UV2Mesh()
|
|
smpl = SMPL()
|
|
|
|
F = np.load("custom_utils/smpl_faces.npy")
|
|
garment_type_table = np.load('custom_data/garment_type_table_test.npy', allow_pickle=True).item()
|
|
|
|
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_body_no_orient': (14475, 3,),
|
|
'mesh_posed': (14475, 3,),
|
|
'mesh_unposed': (14475, 3,),
|
|
'uv_unposed': (512, 256, 3),
|
|
'uv_unposed_32': (32, 16, 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_mask_32': (32, 16,),
|
|
# 'uv_shaped_body': (3, 512, 256,),
|
|
}
|
|
|
|
dataset = factory.get_dataset(folder='/home/workspace/data/dataset/test_ready',
|
|
batch_size=1,
|
|
signature=signature_dynamic,
|
|
shuffle=True,
|
|
do_repeat=False,
|
|
resolutions=[512, 32],
|
|
# frames_file='/home/workspace/data/dataset/test_ready/frames_pose_cluster_20_sample_10.txt',
|
|
frames_file='/home/workspace/data/dataset/test_ready/frames_00438_146.txt',
|
|
)
|
|
static_dataset = StaticDataset(folder='/home/workspace/data/dataset/test_ready',
|
|
resolutions=[512, 32],
|
|
signature=signature_static)
|
|
|
|
|
|
uv_res = 256
|
|
|
|
|
|
|
|
for batch in tqdm.tqdm(dataset):
|
|
batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
|
|
|
|
id_str, frame_nr = batch['id'][0].numpy().decode("utf-8"), batch['frame_nr'][0].numpy().decode("utf-8")
|
|
|
|
type = garment_type_table[id_str]
|
|
|
|
# if type != 'Jumpsuit':
|
|
# continue
|
|
|
|
shape = batch['pose_shape'][:, 72:]
|
|
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
|
|
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]).numpy()[0]
|
|
|
|
|
|
uv_posed = (np.squeeze(batch['uv_posed'].numpy()) + 1) / 2
|
|
uv_body_posed = (np.squeeze(batch['uv_body_posed'].numpy()) + 1) / 2
|
|
uv_unposed = (np.squeeze(batch['uv_unposed'].numpy()) + 1) / 2
|
|
# uv_curvature = np.tile(np.argmax(np.squeeze(batch['uv_curvature'].numpy()), axis=-1)[:,:,np.newaxis], [1,1,3]) / 6
|
|
uv_curvature = color_labels(batch['uv_curvature'][0].numpy().argmax(-1))
|
|
uv_static = (np.squeeze(batch['uv_static'].numpy()) + 1) / 2
|
|
|
|
img = np.concatenate([uv_posed, uv_body_posed, uv_unposed, uv_curvature, uv_static], axis=1)
|
|
|
|
# img_plot = custom_imshow(img)
|
|
# img_plot.write_image('examples/' + type + '_' + id_str + '_' + frame_nr + '.png')
|
|
# img_plot.show()
|
|
|
|
mask = batch['mesh_mask'][0].numpy().astype(bool)
|
|
body = batch['mesh_body_posed'][0].numpy()
|
|
posed = batch['mesh_posed'][0].numpy()[mask, :]
|
|
unposed = batch['mesh_unposed'][0].numpy()[mask, :]
|
|
static = batch['mesh_static'][0].numpy()[mask, :]
|
|
unposed_shaped = unposed + shape_offsets[mask, :]
|
|
|
|
_, faces_garment = _maskFaces(np.zeros((14475, 3), np.float32), F, mask)
|
|
faces_garment = faces_garment.astype(np.int32)
|
|
faces = np.concatenate((F, faces_garment + 14475), axis=0).astype(np.int32)
|
|
colors = np.concatenate((np.ones_like(body, dtype=np.int32)*100, np.ones_like(posed, dtype=np.int32)*230), axis=0).astype(np.float32) / 255.0
|
|
colors_garment = (np.ones_like(posed, dtype=np.int32)*230).astype(np.float32) / 255.0
|
|
|
|
x = np.concatenate((body, posed), axis=0)
|
|
|
|
# m = tm.Trimesh(x, faces, process=False, validate=False, vertex_colors=colors)
|
|
|
|
x = posed
|
|
x_1 = unposed.copy()
|
|
x_1[:, 0] = x_1[:, 0] + 0.75
|
|
x_2 = unposed_shaped.copy()
|
|
x_2[:, 0] = x_2[:, 0] + 1.75
|
|
garments = np.concatenate((x, x_1, x_2), axis=0)
|
|
all_faces = np.concatenate((faces_garment, faces_garment + x.shape[0], faces_garment + 2*x.shape[0]), axis=0).astype(np.int32)
|
|
mesh = garments
|
|
faces = all_faces
|
|
fig = go.Figure()
|
|
fig.add_trace(go.Mesh3d(
|
|
x=mesh[:,0],
|
|
y=mesh[:,2],
|
|
z=mesh[:,1],
|
|
i = faces[:,0],
|
|
j = faces[:,1],
|
|
k = faces[:,2],
|
|
text="",
|
|
lighting={ 'ambient':0.8, 'diffuse':0.5, 'roughness':0.5 },
|
|
|
|
),)
|
|
fig.show()
|
|
from stl import mesh as mesh_stl
|
|
obj = mesh_stl.Mesh(np.zeros(faces.shape[0], dtype=mesh_stl.Mesh.dtype))
|
|
for i, f in enumerate(faces):
|
|
for j in range(3):
|
|
obj.vectors[i][j] = mesh[f[j],:]
|
|
write_obj_2('temp/blender_test_2.obj', garments, all_faces)
|
|
|
|
|
|
# fig = plot_fixed(x, faces, colors)
|
|
# fig.write_image('examples/' + type + '_' + id_str + '_' + frame_nr + '_posed.svg')
|
|
# # fig.show()
|
|
|
|
# fig = plot_fixed(unposed, faces_garment, colors_garment)
|
|
# fig.write_image('examples/' + type + '_' + id_str + '_' + frame_nr + '_unposed.svg')
|
|
# # # fig.show()
|
|
|
|
# fig = plot_fixed(unposed_shaped, faces_garment, colors_garment)
|
|
# fig.write_image('examples/' + type + '_' + id_str + '_' + frame_nr + '_unposed_shaped.svg')
|
|
|
|
# uv_mapping_example(body, F, colors[:14475, :], np.ones((14475), bool))
|
|
# curvature_example(batch['mesh_curvature'][0].numpy(), uv_curvature, posed, faces_garment, mask)
|
|
# plot_upscaling(batch['uv_unposed_32'].numpy()[0], batch['uv_unposed'].numpy()[0], batch['uv_mask_32'].numpy()[0], batch['uv_mask'].numpy()[0], faces_garment, colors_garment, mask)
|
|
|
|
print() |