custom upscaling

2022-09-20 22:53:57 +00:00
parent d83e583e7f
commit fafa28cc69
3 changed files with 423 additions and 0 deletions
@@ -0,0 +1,102 @@
+# check if upscaler is better then filling
+import os
+
+from models.upscaler import UpscalerBasic as Upscaler
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+
+import tqdm
+import numpy as np
+import plotly.graph_objects as go
+
+from custom_utils.custom_utils import UV2Mesh, _maskFaces, load_model
+# from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces
+
+
+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)
+    res = 32
+    res_modifier = '_' + str(res)
+
+    upscale_enhancer = Upscaler(512, res)
+    load_model(upscale_enhancer, 'checkpoints/up_2_3')
+
+    frames = os.listdir(os.path.join(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
+        visited.add(id)
+        print(id)
+        uv_unposed_no_shape_low = np.load(os.path.join(folder, 'uv_unposed_no_shape' + res_modifier, id + '_' + frame_nr + '.npz'));  uv_unposed_no_shape_low = uv_unposed_no_shape_low['uv_unposed_no_shape']
+        uv_unposed_no_shape_filled_low = np.load(os.path.join(folder, 'uv_unposed_no_shape_filled' + res_modifier, id + '_' + frame_nr + '.npz'));  uv_unposed_no_shape_filled_low = uv_unposed_no_shape_filled_low['uv_unposed_no_shape_filled']
+        mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy'))
+        uv_mask = np.load(os.path.join(folder, 'uv_mask' + res_modifier, id + '.npz')); uv_mask = uv_mask['uv_mask']
+        mesh = uv2mesh(uv_unposed_no_shape_low[np.newaxis,:,:,:])[0].numpy()
+        mesh_filled = uv2mesh(uv_unposed_no_shape_filled_low[np.newaxis,:,:,:])[0].numpy()
+        mesh_orig = np.load(os.path.join(folder, 'mesh_unposed_no_shape', id + '_' + frame_nr + '.npz')); mesh_orig = mesh_orig['mesh_unposed_no_shape']
+
+        uv_enhanced = upscale_enhancer([uv_unposed_no_shape_low[np.newaxis,:,:,:], uv_mask[np.newaxis,:,:].astype(np.float32)], training=False)
+        mesh_enhanced = uv2mesh(uv_enhanced)[0].numpy()
+
+        metric = np.sum(np.sqrt(np.sum(mesh_mask[:,np.newaxis] * (mesh_orig - mesh) ** 2, axis=-1))) / np.sum(mesh_mask)
+        metric_filled = np.sum(np.sqrt(np.sum(mesh_mask[:,np.newaxis] * (mesh_orig - mesh_filled) ** 2, axis=-1))) / np.sum(mesh_mask)
+        metric_enhanced = np.sum(np.sqrt(np.sum(mesh_mask[:,np.newaxis] * (mesh_orig - mesh_enhanced) ** 2, axis=-1))) / np.sum(mesh_mask)
+        print(metric)
+        print(metric_filled)
+        print(metric_enhanced)
+        print()
+
+        mesh, faces = _maskFaces(mesh, F, mesh_mask)
+        mesh_filled, _ = _maskFaces(mesh_filled, F, mesh_mask)
+        mesh_enhanced, _ = _maskFaces(mesh_enhanced, F, mesh_mask)
+        x = mesh
+        y = mesh_filled
+        z = mesh_enhanced
+        y[:, 2] += 0.5
+        z[:, 2] += 1
+        meshes = [go.Mesh3d(
+                    x=x[:,0],
+                    y=x[:,2],
+                    z=x[:,1],
+                    i = faces[:,0],
+                    j = faces[:,1],
+                    k = faces[:,2]
+                    ),
+                    go.Mesh3d(
+                    x=y[:,0],
+                    y=y[:,2],
+                    z=y[:,1],
+                    i = faces[:,0],
+                    j = faces[:,1],
+                    k = faces[:,2]
+                    ),
+                    go.Mesh3d(
+                    x=z[:,0],
+                    y=z[:,2],
+                    z=z[:,1],
+                    i = faces[:,0],
+                    j = faces[:,1],
+                    k = faces[:,2]
+                    )]
+
+        fig = go.Figure(data=meshes)
+        fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=2, z=1))
+        fig.update_layout(
+            scene = dict(
+                xaxis = dict(nticks=4, range=[-1,1],),
+                            yaxis = dict(nticks=4, range=[-1,3],),
+                            zaxis = dict(nticks=4, range=[-1,1],),)
+                            )
+
+        fig.show()
+
+        print()