|
import os |
|
import numpy as np |
|
import glob |
|
import open3d as o3d |
|
import json |
|
import argparse |
|
import glob |
|
|
|
parser=argparse.ArgumentParser() |
|
parser.add_argument("--cat",required=True,type=str,nargs="+") |
|
parser.add_argument("--keyword",default="lowres",type=str) |
|
parser.add_argument("--root_dir",type=str,default="../data") |
|
args=parser.parse_args() |
|
|
|
keyword=args.keyword |
|
sdf_folder="occ_data" |
|
other_folder="other_data" |
|
data_dir=args.root_dir |
|
|
|
align_dir=os.path.join(args.root_dir,"align_mat_all") |
|
|
|
align_filelist=glob.glob(align_dir+"/*/*.txt") |
|
valid_model_list=[] |
|
for align_filepath in align_filelist: |
|
if "-v" in align_filepath: |
|
align_mat=np.loadtxt(align_filepath) |
|
if align_mat.shape[0]!=4: |
|
continue |
|
model_id=os.path.basename(align_filepath).split("-")[0] |
|
valid_model_list.append(model_id) |
|
|
|
print("there are %d valid lowres models"%(len(valid_model_list))) |
|
|
|
category_list=args.cat |
|
for category in category_list: |
|
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] |
|
|
|
|
|
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: |
|
if model_id not in valid_model_list: |
|
continue |
|
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+"/%s_partial_points_*.ply"%(keyword)) |
|
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 model_id in val_list: |
|
if "%s_partial_points_0.ply"%(keyword) in partial_valid_list: |
|
partial_valid_list=["%s_partial_points_0.ply"%(keyword)] |
|
else: |
|
partial_valid_list=[] |
|
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,"%s_train_par_img.json"%(keyword)) |
|
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,"%s_val_par_img.json"%(keyword)) |
|
with open(val_save_path,'w') as f: |
|
json.dump(val_par_img_list,f,indent=4) |
|
|
