Now with included sorting per cloud variation

This commit is contained in:
Si11ium
2019-08-02 13:03:56 +02:00
parent 9c100b6c43
commit 43b4c8031a
2 changed files with 22 additions and 6 deletions

View File

@ -38,6 +38,9 @@ parser.add_argument('--test_per_batches', type=int, default=1000, help='run a te
parser.add_argument('--num_workers', type=int, default=4, help='number of data loading workers')
parser.add_argument('--headers', type=strtobool, default=True, help='if raw files come with headers')
parser.add_argument('--collate_per_segment', type=strtobool, default=True, help='whether to look at pointclouds or sub')
parser.add_argument('--has_variations', type=strtobool, default=False,
help='whether a single pointcloud has variations '
'named int(id)_pc.(xyz|dat) look at pointclouds or sub')
opt = parser.parse_args()
@ -70,11 +73,13 @@ if __name__ == '__main__':
test_transform = GT.Compose([GT.NormalizeScale(), ])
dataset = ShapeNetPartSegDataset(root_dir=opt.dataset, collate_per_segment=opt.collate_per_segment,
train=True, transform=train_transform, npoints=opt.npoints, headers=opt.headers)
train=True, transform=train_transform, npoints=opt.npoints,
has_variations=opt.has_variations, headers=opt.headers)
dataLoader = DataLoader(dataset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers)
test_dataset = ShapeNetPartSegDataset(root_dir=opt.dataset, collate_per_segment=opt.collate_per_segment,
train=False, transform=test_transform, npoints=opt.npoints, headers=opt.headers)
train=False, transform=test_transform, npoints=opt.npoints,
has_variations=opt.has_variations, headers=opt.headers)
test_dataLoader = DataLoader(test_dataset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers)
num_classes = dataset.num_classes()