| import numpy as np |
| import scipy |
| import os |
| import trimesh |
| from sklearn.cluster import KMeans |
| import random |
| import glob |
| import tqdm |
| import multiprocessing as mp |
| import sys |
| sys.path.append("..") |
| from datasets.taxonomy import synthetic_category_combined |
|
|
| import argparse |
| parser=argparse.ArgumentParser() |
| parser.add_argument("--category",nargs="+",type=str) |
| parser.add_argument("--root_dir",type=str,default="../data/other_data") |
| args=parser.parse_args() |
| categories=args.category |
| if categories[0]=="all": |
| categories=synthetic_category_combined["all"] |
|
|
| kmeans=KMeans( |
| init="random", |
| n_clusters=7, |
| n_init=10, |
| max_iter=300, |
| random_state=42 |
| ) |
|
|
| def process_data(src_filepath,save_path): |
| |
| src_point_tri = trimesh.load(src_filepath) |
| src_point = np.asarray(src_point_tri.vertices) |
| kmeans.fit(src_point) |
| point_cluster_index = kmeans.labels_ |
|
|
| n_cluster = random.randint(3, 6) |
| choose_cluster = np.random.choice(7, n_cluster, replace=False) |
| aug_point_list = [] |
| for cluster_index in choose_cluster: |
| cluster_point = src_point[point_cluster_index == cluster_index] |
| aug_point_list.append(cluster_point) |
| aug_point = np.concatenate(aug_point_list, axis=0) |
| aug_point_tri = trimesh.PointCloud(vertices=aug_point) |
| print("saving to %s"%(save_path)) |
| aug_point_tri.export(save_path) |
|
|
| pool=mp.Pool(10) |
| for cat in categories: |
| print("processing %s"%cat) |
| point_dir=os.path.join(args.root_dir,cat,"5_partial_points") |
| folder_list=os.listdir(point_dir) |
| for folder in folder_list[:]: |
| folder_path=os.path.join(point_dir,folder) |
| src_filelist=glob.glob(folder_path+"/partial_points_*.ply") |
| for src_filepath in src_filelist: |
| basename=os.path.basename(src_filepath) |
| save_path = os.path.join(point_dir, folder, "aug7_" + basename) |
| pool.apply_async(process_data,(src_filepath,save_path)) |
| pool.close() |
| pool.join() |
|
|
|
|
|
|
