diff --git a/dataset.py b/dataset.py
index 55ee42a..1478cb5 100644
--- a/dataset.py
+++ b/dataset.py
@@ -167,7 +167,7 @@ class DatasetFactory():
         return tf_signature
 
     
-    def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=False, resolutions=[512], in_memory=False, do_repeat=True, is_cape=False):
+    def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=False, resolutions=[512], in_memory=False, do_repeat=True, is_cape=False, use_one_hot=False):
         # assert (use_filled and len(resolutions) >= 1) or (not use_filled and len(resolutions) == 1)
         assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
 
@@ -179,7 +179,13 @@ class DatasetFactory():
         separator = '_'
         if is_cape:
             separator = '.'
-        
+
+        one_hot = ''
+        one_hot_2 = ''
+        if use_one_hot:
+            one_hot = '_one_hot'
+            one_hot_2 = '_one_hot_2'
+
 
         def read_item(frame):
             frame_split = tf.strings.split(frame, sep=separator)
@@ -224,7 +230,8 @@ class DatasetFactory():
             #     uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
             #     item['uv_unposed'] = uv_unposed
             if 'mesh_unposed' in signature.keys():
-                mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape', id + separator + frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape']
+                # print('one_hot')
+                mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape' + one_hot, id + separator + frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape']
                 item['mesh_unposed'] = mesh_unposed_no_shape
 
 
@@ -248,11 +255,12 @@ class DatasetFactory():
                     item['uv_curvature' + res_string] = uv_curvature
 
                 if 'uv_unposed' + res_string in signature.keys():
-                    uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + filled + res_string, id + separator + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape' + filled]
+                    uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + one_hot_2 + res_string, id + separator + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape' + filled]
                     item['uv_unposed' + res_string] = uv_unposed_no_shape
 
                 if 'uv_body_posed' + res_string in signature.keys():
-                    uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient' + filled + res_string, id + separator + frame_nr + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient' + filled]
+                    uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient_2' + filled + res_string, id + separator + frame_nr + '.npz'))
+                    uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient' + filled]
                     item['uv_body_posed' + res_string] = uv_posed_body_no_orient
 
 
@@ -318,8 +326,8 @@ if __name__ == '__main__':
     signature_dynamic = { 'id': (),
                         #   'pose_shape': (72 + 11,),
                         #   'uv_posed': (3, 512, 256,),
-                          #'uv_body_posed': (3, 512, 256,),
-                        #   'uv_body_posed_32': (3, 32, 16,),
+                        #   'uv_body_posed': (512, 256, 3),
+                        #   'uv_body_posed_32': (32, 16, 3),
                         #   'mesh_body_posed': (14475, 3,),
                         #   'mesh_posed': (14475, 3,),
                         #   'mesh_unposed': (14475, 3,),
@@ -341,13 +349,16 @@ if __name__ == '__main__':
                                   signature=signature_dynamic,
                                   use_filled=False,
                                   resolutions=[512],
-                                  do_repeat=False)
+                                  do_repeat=False,
+                                  use_one_hot=True)
     static_dataset = StaticDataset(folder='/home/workspace/local-data/train',
                                    signature=signature_static,
                                    use_filled=False,
                                    resolutions=[512])
 
 
+    # for i, batch in enumerate(dataset):
+    #     print(i)
 
 
     sum_uv = np.zeros((512, 256, 3), dtype=np.float32)
@@ -375,11 +386,11 @@ if __name__ == '__main__':
                 k = F[:,2]
             )]).show()
 
-    np.save('custom_data/train_no_shape_avg_uv.npy', avg_garment)
+    np.save('custom_data/train_no_shape_one_hot_avg_uv.npy', avg_garment)
 
-    print()
+    # print()
 
-    pass
+    # pass
 
 
     # uv_res = 256