| import os,sys |
| import numpy as np |
| import glob |
| import open3d as o3d |
| import json |
| import argparse |
|
|
| parser=argparse.ArgumentParser() |
| parser.add_argument("--cat",required=True,type=str,nargs="+") |
| parser.add_argument("--root_dir",type=str,default="../data") |
| args=parser.parse_args() |
|
|
| sdf_folder="occ_data" |
| other_folder="other_folder" |
| data_dir=args.root_dir |
| category=args.cat |
| train_path=os.path.join(data_dir,sdf_folder,category,"train.lst") |
| with open(train_path,'r') as f: |
| train_list=f.readlines() |
| train_list=[item.rstrip() for item in train_list] |
| if ".npz" in train_list[0]: |
| train_list=[item[:-4] for item in train_list] |
| val_path=os.path.join(data_dir,sdf_folder,category,"val.lst") |
| with open(val_path,'r') as f: |
| val_list=f.readlines() |
| val_list=[item.rstrip() for item in val_list] |
| if ".npz" in val_list[0]: |
| val_list=[item[:-4] for item in val_list] |
|
|
| category_list=args.cat |
| for category in category_list: |
| sdf_dir=os.path.join(data_dir,sdf_folder,category) |
| filelist=os.listdir(sdf_dir) |
| model_id_list=[item[:-4] for item in filelist if ".npz" in item] |
|
|
| train_par_img_list=[] |
| val_par_img_list=[] |
| for model_id in model_id_list: |
| image_dir=os.path.join(data_dir,other_folder,category,"6_images",model_id) |
| partial_dir=os.path.join(data_dir,other_folder,category,"5_partial_points",model_id) |
| if os.path.exists(image_dir)==False and os.path.exists(partial_dir)==False: |
| continue |
| if os.path.exists(image_dir): |
| image_list=glob.glob(image_dir+"/*.jpg")+glob.glob(image_dir+"/*.png") |
| image_list=[os.path.basename(image_path) for image_path in image_list] |
| else: |
| image_list=[] |
|
|
| if os.path.exists(partial_dir): |
| partial_list=glob.glob(partial_dir+"/partial_points_*.ply") |
| else: |
| partial_list=[] |
| partial_valid_list=[] |
| for partial_filepath in partial_list: |
| par_o3d=o3d.io.read_point_cloud(partial_filepath) |
| par_xyz=np.asarray(par_o3d.points) |
| if par_xyz.shape[0]>2048: |
| partial_valid_list.append(os.path.basename(partial_filepath)) |
| if len(image_list)==0 and len(partial_valid_list)==0: |
| continue |
| ret_dict={ |
| "model_id":model_id, |
| "image_filenames":image_list[:], |
| "partial_filenames":partial_valid_list[:] |
| } |
| if model_id in train_list: |
| train_par_img_list.append(ret_dict) |
| elif model_id in val_list: |
| val_par_img_list.append(ret_dict) |
|
|
| |
| train_save_path=os.path.join(sdf_dir,"train_par_img.json") |
| with open(train_save_path,'w') as f: |
| json.dump(train_par_img_list,f,indent=4) |
|
|
| val_save_path=os.path.join(sdf_dir,"val_par_img.json") |
| with open(val_save_path,'w') as f: |
| json.dump(val_par_img_list,f,indent=4) |
|
|
