File size: 2,902 Bytes
cc9780d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
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)
#print(train_par_img_list)
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)
|