def __init__(self, opt):
BaseDataset.__init(self, opt)
opt.downsample = 1
opt.pack_n = 10000
opt.excl_subclass = None
classes = opt.class_3d
assert classes is not None, 'dataset argument [classes] has to be set'
self.res = opt.voxel_res
dataset_paths = paths['canon']
# parse classes
classlist = _parse_class(classes)
self._class_str = '+'.join(classlist)
assert (opt.excl_subclass is None) or (classlist == [
'03001627'
]), 'subclass exclusion only supported for chair'
class_size = dict()
class_to_ind = dict()
filelist = list()
labellist = list()
self._matname = dataset_paths['matname']
dataset_merged_csv = dataset_paths['filelist']
dataset_path_format = dataset_paths['filepath']
for i, c in enumerate(classlist):
merged_csvfile = dataset_merged_csv.format(classid=c, res=self.res)
if os.path.isfile(merged_csvfile):
with open(merged_csvfile, 'r') as fin:
class_filelist = list(map(str.strip, fin.readlines()))
else:
class_filelist = sorted(
glob.glob(dataset_path_format.format(classid=c,
res=self.res),
recursive=False))
if len(class_filelist) == 0:
raise ValueError('No .mat files found for class: ' + c)
class_size[c] = len(class_filelist)
class_to_ind[c] = i
filelist += sorted(class_filelist)
labellist += [class_to_label[c]] * len(class_filelist)
# filter out specific subclasses
subclasses_excl = opt.excl_subclass
if subclasses_excl is not None:
assert (classlist == ['03001627']), \
"Excluding subclasses only supported for chairs"
if subclasses_excl in chair_subclass_aliases:
subclasses_excl = chair_subclass_aliases[subclasses_excl]
subclasses_excl = (set(subclasses_excl.split('+')))
with open(chair_subclass_file, 'r') as f:
lines = f.readlines()
list_obj_ids = [l.split(' ')[0] for l in lines]
list_subclasses = [
set(l.split(' ')[1].replace('\n', '').split(','))
for l in lines
]
subclass_dict = dict(zip(list_obj_ids, list_subclasses))
idlist = [
basename(dirname(dirname(filename)))
for filename in filelist
]
intersection_list = [
(len(subclasses_excl.intersection(subclass_dict[id_]))
if id_ in subclass_dict else -1) for id_ in idlist
]
print(str_verbose, 'subclass exclusion: ')
print(
str_verbose, '\tvoxels kept: %d' %
(np.array(intersection_list) == 0).sum())
print(
str_verbose, '\tvoxels excluded: %d' %
(np.array(intersection_list) > 0).sum())
print(
str_verbose, '\tvoxels missing subclass: %d' %
(np.array(intersection_list) < 0).sum())
# update filelist and labellist
filelist = [
filelist[i] for i in range(len(filelist))
if intersection_list[i] == 0
]
labellist = [
labellist[i] for i in range(len(labellist))
if intersection_list[i] == 0
]
self._classes = classlist
self._class_to_ind = class_to_ind
self._filelist = filelist
self._labellist = labellist
self._transform = self.get_transform(opt, opt.downsample)
self.size = len(self._filelist)
