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 import scipy.io as sio from plotly.subplots import make_subplots from dataset import DatasetFactory, StaticDataset from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces, load_model, save_stl 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 faces2graph(F, v=None): if v: v = set(v) # Only this verts G = {} for f in F: for i in f: if v and i not in v: continue if not i in G: G[i] = {} for j in f: if v and j not in v: continue if i == j: continue G[i][j] = 1 return G SMOOTH_IT = 100 SMOOTH_EXHAUST_IT = 100 def smooth(M, F, segm): f_skirt = [] for f in F: f_skirt.append(f[0]) G = faces2graph(f_skirt) # Segmentation for skirt smoothing # Laplacian smoothing for _ in range(SMOOTH_EXHAUST_IT): for v in segm['exhaustive'].reshape(-1): M[v] = M[list(G[v].keys())].mean(0) for _ in range(SMOOTH_IT): # aux = M.copy() # FRONT for v in segm['front'].reshape(-1): M[v] = M[list(G[v].keys())].mean(0) for v in segm['frontBoundary'].reshape(-1): M[v,:2] = M[list(G[v].keys()),:2].mean(0) # No z-smoothing # BACK for v in segm['back'].reshape(-1): M[v] = M[list(G[v].keys())].mean(0) for v in segm['backBoundary'].reshape(-1): M[v,:2] = M[list(G[v].keys()),:2].mean(0) # No z-smoothing # M = aux.copy() return M def uv_mapping_example_skirt(mesh, faces, colors, mask): segm = sio.loadmat('custom_utils/vg_skirt_smoothing') F = sio.loadmat("custom_utils/faces_skirt.mat")['F'][0] faces = np.load('custom_utils/smpl_faces_skirt.npy') mesh = smooth(mesh, F, segm) 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, is_skirt=True) 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/iccv/iccv_32_2') 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.BILINEAR).numpy() good = tf.image.resize(enhanced, size=(512, 256), method=tf.image.ResizeMethod.BILINEAR).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') save_stl(f'temp/iccv/{id_str}_{frame_nr}_mesh_upscale_bad.stl', uv2mesh(bad[np.newaxis,:,:,:])[0].numpy()[mask, :], faces) save_stl(f'temp/iccv/{id_str}_{frame_nr}_mesh_upscale_good.stl', uv2mesh(good[np.newaxis,:,:,:])[0].numpy()[mask, :], faces) 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', # frames_file='/home/workspace/data/dataset/test_ready/frames_00228_137.txt', # frames_file='/home/workspace/data/dataset/test_ready/frames_02572_1.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() # 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)) uv_mapping_example_skirt(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()