diff --git a/downscale_unpose_test.py b/downscale_unpose_test.py new file mode 100644 index 0000000..882166b --- /dev/null +++ b/downscale_unpose_test.py @@ -0,0 +1,145 @@ +import os +os.environ["CUDA_VISIBLE_DEVICES"] = "1" + +import tqdm +import numpy as np +import tensorflow as tf +import plotly.graph_objects as go + +from custom_utils.custom_utils import UV2Mesh, _maskFaces +from custom_utils.unpose import unpose +from custom_utils.fill_uv import fill +from custom_utils.smpl_2 import SMPL as SMPL_2 + + +if __name__ == '__main__': + folder = '/home/workspace/local-data/train_100/' + mask = np.load('custom_utils/mask_256x512.npy') + F = np.load("custom_utils/smpl_faces.npy") + uv2mesh = UV2Mesh(res=256) + smpl_2 = SMPL_2('custom_utils') + + frames = os.listdir(folder + '/pose_shape') + visited = set() + meshes = [] + count = 0 + + for frame in tqdm.tqdm(frames): + id, frame_nr = frame.split('_') + frame_nr = frame_nr.split('.')[0] + if id in visited: + continue + if id == '06798': + continue + + visited.add(id) + print(id) + + + uv_static = np.load(os.path.join(folder, 'uv_static_filled_32', id + '.npz')); uv_static = uv_static['uv_static_filled'] + uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt'] + uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask'] + mesh_mask = np.load(os.path.join('/home/workspace/local-data/train_100/mesh_mask', id + '.npy')) + mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + '_' + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed'] + mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed'] + pose_shape = np.load(os.path.join(folder, 'pose_shape', id + '_' + frame_nr + '.npy')).astype(np.float32) + pose, shape = pose_shape[:72], pose_shape[72:] + gender = int(shape[-1] > 0) + mesh_static = uv2mesh(uv_static[np.newaxis,:,:,:])[0].numpy() + mask_idx = np.array(range(14475))[mesh_mask] + + uv_posed_filled = fill(uv_posed, uv_mask) + down = tf.image.resize(uv_posed_filled, size=(32, 16), method=tf.image.ResizeMethod.LANCZOS5) + mesh_posed_low = uv2mesh(down[tf.newaxis,:,:,:]) + + unposed = unpose(mesh_static[mesh_mask,:], mesh_posed[np.newaxis,mesh_mask,:], pose[:,np.newaxis], shape, gender, mesh_mask, mask_idx, smpl_2, id) + unposed_low = unpose(mesh_static[mesh_mask,:], mesh_posed_low.numpy()[:,mesh_mask,:], pose[:,np.newaxis], shape, gender, mesh_mask, mask_idx, smpl_2, id) + + + unposed_mesh, faces = _maskFaces(unposed[0], F, mesh_mask) + x = unposed_mesh + x[:, 2] += 0.5 * count + meshes.append(go.Mesh3d( + # 8 vertices of a cube + x=x[:,0], + y=x[:,2], + z=x[:,1], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2] + )) + + + x = mesh_unposed[mesh_mask, :] + x[:, 2] += 0.5 * (count + 1) + meshes.append(go.Mesh3d( + # 8 vertices of a cube + x=x[:,0], + y=x[:,2], + z=x[:,1], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2] + )) + + x = unposed_low[0,mesh_mask, :] + x[:, 2] += 0.5 * (count + 2) + meshes.append(go.Mesh3d( + # 8 vertices of a cube + x=x[:,0], + y=x[:,2], + z=x[:,1], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2] + )) + + x = mesh_posed_low.numpy()[0,mesh_mask, :] + x[:, 1] += 0.5 * (count + 3) + meshes.append(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] + )) + + x = mesh_posed[mesh_mask, :] + x[:, 1] += 0.5 * (count + 4) + meshes.append(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] + )) + + count += 1 + + if count > 1: + break + + + + fig = go.Figure(data=meshes) + fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=count/2-1, z=1)) + fig.update_layout( + scene = dict( + xaxis = dict(nticks=4, range=[-1,1],), + yaxis = dict(nticks=4, range=[-1,count/2+1],), + zaxis = dict(nticks=4, range=[-1,1],),) + ) + + fig.show() + + print() +