try unposing in low res

2022-09-16 16:49:18 +00:00
parent d14aa43461
commit eac8aad8ab
1 changed files with 145 additions and 0 deletions
@@ -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()
+