init

recontruct_mesh.py

@@ -0,0 +1,77 @@
+import os
+import sys
+import cv2
+import numpy as np
+import trimesh as tm
+import matplotlib.pyplot as plt
+
+def m2uv_array(m2uv):
+    result = np.zeros((14475, 2), dtype=np.int32)
+    for k, v in m2uv.items():
+        if len(v) == 2:
+            result[k] = list(v)
+        elif len(v) == 1:
+            x = v.pop()
+            result[k] = [x, x]
+        else:
+            assert False
+
+    return result
+
+
+def uv_to_mesh(uv):
+    uv_res = 256
+    # m2uv = np.load('/workspace/Code/featurelines/MeshUtils/m2uv.npy', allow_pickle=True).item()
+    m2uv = np.load(r'D:\Work\Code\featurelines\MeshUtils\m2uv.npy', allow_pickle=True).item()
+    m2uv = m2uv_array(m2uv)
+    # uv_map = np.load('/workspace/Code/featurelines/Meshutils/stretched_uv_map_3.npy')
+    uv_map = np.load(r'D:\Work\Code\featurelines\MeshUtils\stretched_uv_map_3.npy')
+    uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
+
+    mesh = np.zeros((14475, 3))
+    uv_pixels = uv_map_discreet[m2uv]
+    # counter = np.ones((14475))
+    # counter[m2uv[:, 0] == m2uv[:, 1]] = 2
+    mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis]
+    return mesh
+
+
+def plot(uv, mesh):
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(111, projection='3d')
+    ax.set_xlim3d(-1, 1)
+    ax.set_ylim3d(-1, 1)
+    ax.set_zlim3d(-1, 1)
+    ax.set_xlabel("X")
+    ax.set_ylabel("Y")
+    ax.set_zlabel("Z")
+    ax.scatter(mesh[:, 0], mesh[:, 1], mesh[:, 2], s=5)
+    # plt.show()
+
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(111)
+    ax.imshow(np.round( (uv - uv.min()) / (uv.max() - uv.min()) * 255).astype(np.int32))
+    plt.show()
+
+
+
+if __name__ == '__main__':
+    # folder =  '/workspace/Checkpoints/checkpoint_32100/output'
+    folder = r'D:\Work\Checkpoints\checkpoint_52000\output'
+
+    files = os.listdir(folder)
+
+    for file in files:
+        if not file.endswith('.npz'):
+            continue
+
+        uv = np.load(os.path.join(folder, file))
+        uv = np.transpose(uv['uv'][0], [1,2,0])
+        # uv[abs(uv) < 0.01] = 0
+        # print(np.sum(abs(uv) < 0.01))
+        mesh = uv_to_mesh(uv)
+
+        plot(uv, mesh)
+
+        pass
+