99 lines
3.1 KiB
Python
99 lines
3.1 KiB
Python
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() |