def get_files_list(root_dir, categories):
files = find_files(root_dir, 'obj')
_, dirnames = get_cat_names(root_dir, return_dirnames=True)
newfiles = []
for file in files:
id = get_file_id(file)
cat = categories[id]
if file.split(os.sep)[-4] == dirnames[cat]:
newfiles.append(file)
return newfiles
def parse_shapenet_split(root_dir, categories):
files = find_files(root_dir, 'obj')
split = {}
coding = {
'train': 0,
'test': 1,
}
counter = [0] * 55
for file in files:
id = get_file_id(file)
cat = categories[id]
if counter[cat] % 10 in [0, 5]:
split[id] = 1
else:
split[id] = 0
counter[cat] += 1
return split
def get_metadata(root_dir):
files = find_files(root_dir, 'off')
categories = {}
split = {}
coding = {
'train':0,
'test':1
}
cats = get_cat_names(root_dir)
cats_mapping = {}
for i in range(len(cats)):
cats_mapping[cats[i]] = i
for file in files:
splited = file.split(os.sep)
name = os.path.join(Path(file).stem)
categories[name] = cats_mapping[splited[-3]]
split[name] = coding[splited[-2]]
cat_names = get_cat_names(root_dir)
return categories, split, cat_names
def get_files_list(root_dir, categories):
return find_files(root_dir, 'off')