prototype conditional pipeline

dataset.py

@@ -2,6 +2,7 @@ import os
 import sys
 import time
 import cv2
+import random
 import numpy as np
 import tensorflow as tf
 
@@ -10,32 +11,74 @@ class DatasetSimple(tf.data.Dataset):
     def _generator(folder):
         # Opening the file
         folder = folder.decode("utf-8")
-        uv_folder = os.path.join(folder, 'npz') 
+        uv_folder = os.path.join(folder, 'uv_posed') 
         pose_hape_folder = os.path.join(folder, 'poses_shapes') 
         frames = os.listdir(uv_folder)
+        random.shuffle(frames)
 
         for frame in frames:
-            uv = np.load(os.path.join(uv_folder, frame))
-            uv = uv['uv_gt']
-            # uv = cv2.resize(uv, [64, 128])
-            uv = np.transpose(uv, [2, 0, 1])
+            id, frame_nr = frame.split('.')[0].split('_')
+            uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']; uv_posed = np.transpose(uv_posed, [2, 0, 1])
+            pose_shape = np.load(os.path.join(folder, 'pose_shape',  id + '_' + frame_nr + '.npy'))
+            uv_static = np.load(os.path.join(folder, 'uv_static', id +  '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1])
+            mesh_static = np.load(os.path.join(folder, 'mesh_static', id +  '.npz')); mesh_static = mesh_static['mesh_static']
+            mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
+            mask = np.load(os.path.join(folder, 'mask', id + '.npy')); mask = mask.astype(np.float32)
 
-            pose_shape = np.load(os.path.join(pose_hape_folder, frame.split('.')[0] + '.npy'))
+            # uv_posed, (pose, shape, uv_static), mesh_static, mesh_body_posed, mask
+            yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static), mesh_static, mesh_body_posed, mask
 
-            yield uv, pose_shape
-
-    def __new__(cls, folder, res=64):
+    def __new__(cls, folder, res=256):
         return tf.data.Dataset.from_generator(
             cls._generator,
-            output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (83), dtype = tf.float32)),
+            output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), (tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475,), dtype = tf.float32) ),
             args=(folder, )
         )
 
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+
+
 if __name__ == '__main__':
     
-    dataset = DatasetSimple('/workspace/Dataset/stylegan/npz_1')
+    dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').batch(2)
+
+    from losses import collision_loss as collision_penalty
+    from losses import collision_penalty as collision_penalty_2
+    from custom_utils.custom_utils import UV2Mesh
+    import matplotlib.pyplot as plt
+
+    def plot_3d(x, y=None):
+        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.grid(False)
+        ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5)
+        if y is not None:
+            ax.scatter(y[:, 0], y[:, 1], y[:, 2], s=5)
+        plt.savefig('temp.jpg')
+
+    uv2mesh = UV2Mesh()
 
     for d in dataset:
         d
+
+        v = uv2mesh(d[0])
+        b = d[3]
+        m = d[4]
+
+        b2 = tf.concat((b[1][tf.newaxis,:,:], b[0][tf.newaxis,:,:]), axis=0)
+
+        c = collision_penalty(v, b, m)
+        c2 = collision_penalty(v, b2, m)
+
+        c_ = collision_penalty_2(v, b, m)
+        c2_ = collision_penalty_2(v, b2, m)
+
+        print()
         pass