lot of things

2022-08-06 12:27:24 +00:00
parent cd498e79f3
commit ee8f1cb68c
28 changed files with 4595 additions and 175 deletions
@@ -0,0 +1,145 @@
+import os
+import sys
+import time
+from unicodedata import name
+import cv2
+import random
+import numpy as np
+import tensorflow as tf
+
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
+import plotly.graph_objects as go
+import plotly.express as ex
+import cmapy as cm
+
+if __name__ == '__main__':
+    os.environ["CUDA_VISIBLE_DEVICES"] = "6"
+
+    uv2mesh = UV2Mesh()
+
+    F = np.load("custom_utils/smpl_faces.npy")
+
+    factory = DatasetFactory()
+    signature_dynamic = { 'id': (),
+                        #   'pose_shape': (72 + 11,),
+                        #   'uv_posed': (3, 512, 256,),
+                        #   'uv_body_posed': (3, 512, 256,),
+                        #   'mesh_body_posed': (14475, 3,),
+                        #   'mesh_posed': (14475, 3,),
+                          'mesh_unposed': (14475, 3,),
+                          'uv_unposed': (3, 512, 256,),
+                        #   'uv_curvature': (512, 256),
+                        #   'mesh_curvature':  (14475)
+                        }
+    signature_static = { #'uv_static': (3, 512, 256,),
+                        #  'mesh_static': (14475, 3,),
+                         'mesh_mask': (14475,),
+                         'uv_mask': (512, 256,),
+                        #  'uv_shaped_body': (3, 512, 256,),
+                          }
+
+    dataset = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train',
+                                  batch_size=1,
+                                  signature=signature_dynamic)
+    static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
+                                   signature=signature_static)
+
+
+    uv_res = 256
+
+    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 unique_uv_map_mask(uv_map_discreet):
+        mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool)
+        temp = np.zeros((uv_res * 2, uv_res))
+        for p in uv_map_discreet:
+            temp[p[0], p[1]] += 1
+        # temp = temp == 1
+        for i, p in enumerate(uv_map_discreet):
+            if temp[p[0], p[1]] == 1:
+                mask[i] = True
+            elif temp[p[0], p[1]] == 2:
+                mask[i] = True
+                temp[p[0], p[1]] = 0
+        return mask
+
+    def uv_to_mesh(uv, m2uv):
+        res = uv.shape[0]
+        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
+
+
+    uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
+    m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
+    m2uv = m2uv_array(m2uv)
+    uv_mask = np.load('custom_utils/mask_256x512.npy')
+    curvature_label_type = '6_25'
+
+    uv_res = 256
+    uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
+    uv_mask = cv2.dilate(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis]
+
+    uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
+    uv_map_mask = unique_uv_map_mask(uv_map_discreet)
+    uv_counter = np.zeros((uv_map.shape[0]))
+    uv_counter[m2uv[:, 0]] = 1
+    uv_counter[m2uv[:, 1]] += 1
+    uv_counter = uv_counter[uv_map_mask]
+
+
+
+    for i, batch in enumerate(dataset):
+        print(i)
+        batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
+
+        uv_unposed = np.squeeze(batch['uv_unposed'].numpy())
+        mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
+        mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
+
+
+        values = mesh_unposed
+        uv_values = np.zeros((uv_map.shape[0], 3))
+        uv_values[m2uv[:, 0]] = values
+        uv_values[m2uv[:, 1]] += values
+        uv_counter = np.zeros((uv_map.shape[0]))
+        uv_counter[m2uv[:, 0]] = 1
+        uv_counter[m2uv[:, 1]] += 1
+        uv_values /= uv_counter[:, np.newaxis]
+        x = uv_values
+        # verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1)
+        im = ((x - x.min()) / (x.max() - x.min())) * 255
+
+
+        x = np.transpose(uv_unposed, [1,2,0])
+        x = uv2mesh(uv_unposed[tf.newaxis,:,:,:])[0].numpy()[mask, :]
+        _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
+        go.Figure(data=[
+                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]
+                )]).show()
+
+
+        print()