|
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) |
|
|
