generate curv with different radius

curvature_labels.py

@@ -153,41 +153,23 @@ def get_bins(values, nr_of_bins):
     return bins
 
 
-def generate_curvature_labels(folder, bins, curvatures):
+def generate_curvature_labels(folder, bins, curvatures, frames, neighbourhood):
 
-    ids = os.listdir(folder)#[:10]
+    os.makedirs(os.path.join(folder, f'mesh_curvature_{neighbourhood}'), exist_ok=True)
     labels = np.linspace(0, bins.shape[0] - 2, bins.shape[0] - 1)
 
-    i = 0
-    for id in tqdm.tqdm(ids):
-        garment_dir = os.path.join(folder, id)
-
-        if not os.path.isdir(garment_dir):
-            continue
-
-        garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and f != 'info.mat' and not f.endswith('.obj')]
-        # assert len (garments) == 1
-        if len (garments) != 1:
-            continue
-        garment = garments[0]
-        mask = np.squeeze(sio.loadmat(os.path.join(garment_dir, garment + '_mask.mat'))['mask'].astype(bool))
-
+    for i, frame in enumerate(tqdm.tqdm(frames)):
+        # id = frame.split('_')[0]
+        # mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy'))
+        # print(i) 
         curvature = curvatures[i]
-        i += 1
 
-        curvature_label = np.zeros((curvature.shape[0], curvature.shape[1]))
-        frame = 0
-        for frame in range(curvature.shape[0]):
-            frame_labels = np.zeros((curvature.shape[1]), dtype=np.float32)
-            for j in range(bins.shape[0] - 2):
-                frame_labels[(curvature[frame, :] >= bins[j]) * (curvature[frame, :] < bins[j+1])] = labels[j]
-            frame_labels[curvature[frame, :] >= bins[-2]] = labels[-1]
-            curvature_label[frame, :] = frame_labels
-            frame += 1
+        frame_labels = np.zeros((curvature.shape[0]), dtype=np.float32)
+        for j in range(bins.shape[0] - 2):
+            frame_labels[(curvature >= bins[j]) * (curvature < bins[j+1])] = labels[j]
+        frame_labels[curvature >= bins[-2]] = labels[-1]
 
-        unmasked_label = np.zeros((curvature.shape[0], 14475), dtype=np.float32)
-        unmasked_label[:, mask] = curvature_label
-        np.save(os.path.join(garment_dir, "curvature_label_" + str(nr_of_labels) + "_" + str(neighbourhood) + ".npy"), unmasked_label)
+        np.savez_compressed(os.path.join(folder, f'mesh_curvature_{neighbourhood}', f'{frame}.npz'), mesh_curvature=frame_labels)
         # writePC2(os.path.join(garment_dir, "curvature_label_" + str(nr_of_labels) + "_" + str(neighbourhood) + ".pc16"), unmasked_label, True)
 
 
@@ -258,24 +240,10 @@ def visualize_curvature_labels(folder):
         print()
 
 
-def read_raw_curvatures(folder):
-
-    ids = os.listdir(folder)[int(args.start):int(args.end)]
-
+def read_raw_curvatures(folder, frames, neighbourhood):
     curvatures_all = []
-
-    for id in tqdm.tqdm(ids):
-        garment_dir = os.path.join(folder, id)
-
-        if not os.path.isdir(garment_dir):
-            continue
-
-        garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and f != 'info.mat' and not f.endswith('.obj')]
-        # assert len (garments) == 1
-        if len (garments) != 1:
-            continue
-
-        curv = np.load(garment_dir + '/curvature_25_raw.npy')
+    for frame in tqdm.tqdm(frames):
+        curv = np.load(os.path.join(folder, f'mesh_curvature_raw_{neighbourhood}', f'{frame}.npz'))['mesh_curvature_raw']
         curvatures_all.append(curv)
     
     return curvatures_all
@@ -283,9 +251,9 @@ def read_raw_curvatures(folder):
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument("-f", "--folder", help = "process the given folder", default=r'/home/workspace/local-data/valid_skirt')
+    parser.add_argument("-f", "--folder", help = "process the given folder", default=r'/home/workspace/data/dataset/local-cache/train_skirt')
     parser.add_argument("-s", "--start", help = "start of range", default=0)
-    parser.add_argument("-e", "--end", help = "end of range", default=10)
+    parser.add_argument("-e", "--end", help = "end of range", default=30000)
     parser.add_argument("-n", "--name", help = "end of range", default="lhunor")
     args = parser.parse_args()
 
@@ -315,12 +283,12 @@ if __name__ == '__main__':
         lines = [line.strip() for line in lines]
     frames = lines[int(args.start):int(args.end)]
 
-    generate_curvatures(folder, neighbourhood, frames)
-    # curvatures = read_raw_curvatures(folder)
+    # generate_curvatures(folder, neighbourhood, frames)
+    curvatures = read_raw_curvatures(folder, frames, neighbourhood)
     # bins = get_bins(curvatures, nr_of_labels)
     # np.save('train_bins_skirt.npy', bins)
-    # bins = np.load('train_bins.npy')
-    # generate_curvature_labels(folder, bins, curvatures)
+    bins = np.load('custom_data/train_bins_skirt.npy')
+    generate_curvature_labels(folder, bins, curvatures, frames, neighbourhood)
     # visualize_curvature_labels(folder)
 
     pass