pose clustering

cluster_poses.py

@@ -0,0 +1,99 @@
+import os
+import tqdm
+import random
+import numpy as np
+import pandas as pd
+import plotly.graph_objects as go
+
+from sklearn.cluster import KMeans, MiniBatchKMeans
+
+
+def generate_clusters(folder, n_cluster):
+    frames = os.listdir(os.path.join(folder, 'pose_shape'))
+    frames = [frame.split('.')[0] for frame in frames]
+    # random.shuffle(frames)
+    # frames = frames[:5000]
+
+    all_poses = np.zeros((len(frames), 69), np.float32)
+
+    i = 0
+    for frame in tqdm.tqdm(frames):
+        # id, frame_nr = frame.split('_')
+        # frame_nr = frame_nr.split('.')[0]
+
+        pose_shape = np.load(os.path.join(folder, 'pose_shape', frame + '.npy'))
+        pose = pose_shape[3:72]
+
+        all_poses[i, :] = pose
+
+        i+= 1
+
+
+    # kmeans = KMeans(n_clusters=20, random_state=0).fit(all_poses)
+    kmeans = MiniBatchKMeans(n_clusters=n_cluster, random_state=0, verbose=1).fit(all_poses)
+
+    data = {'frame': frames, 'label': kmeans.labels_}
+
+    df = pd.DataFrame(data)
+    df.to_csv('custom_data/' + folder.split('/')[-1] + '_pose_cluster_' + str(n_cluster) + '.csv', index=False)
+
+
+def sample_from_cluster(file, n_sample, output_file):
+    df = pd.read_csv(file)
+    samples = df.groupby('label', group_keys=False).apply(lambda x: x.sample(n_sample))
+
+    sampled_frames = np.asarray(samples['frame'].array)
+
+    with open(os.path.join(folder, output_file), 'w') as f:
+        for i, frame in enumerate(sampled_frames):
+            f.write(frame + '\n')
+
+
+if __name__ == '__main__':
+    folder = '/home/workspace/data/dataset/test_ready'
+
+    is_skirt = False
+    F = np.load("custom_utils/smpl_faces" + ("_skirt" if is_skirt == 1 else "") + ".npy")
+
+    # generate_clusters(folder, 20)
+    sample_from_cluster('custom_data/test_ready_pose_cluster_20.csv', 10, '/home/workspace/data/dataset/test_ready/frames_pose_cluster_20_sample_10.txt')
+
+
+    # samples = df.groupby('label', group_keys=False).apply(lambda x: x.sample())
+
+    # sampled_frames = np.asarray(samples['frame'].array)
+    # sampled_labels = np.asarray(samples['label'].array)
+
+    # for i in range(sampled_frames.shape[0]):
+    #     body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', sampled_frames[i] + '.npz')); body_no_orient = body_no_orient['mesh_posed_body_no_orient']
+
+    #     x = body_no_orient
+    #     fig = 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 = F[:,0],
+    #                 j = F[:,1],
+    #                 k = F[:,2],
+    #                 text='cluster: ' + str(sampled_labels[i])
+    #             ),
+    #         ])
+
+    #     fig.update_layout(scene_aspectmode='manual',
+    #                     scene_aspectratio=dict(x=1, y=1, z=1))
+    #     fig.update_layout(
+    #         scene = dict(
+    #             xaxis = dict(nticks=4, range=[-1.5,1.5],),
+    #             yaxis = dict(nticks=4, range=[-1.5,1.5],),
+    #             zaxis = dict(nticks=4, range=[-1.5,1.5],),)
+    #             )
+
+    #     fig.show()
+
+    #     print()
+
+
+    print()