SoftGroup/data/scannetv2.py
2022-04-08 03:19:49 +00:00

15 lines
670 B
Python

from .custom import CustomDataset
class ScanNetDataset(CustomDataset):
CLASSES = ('cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window', 'bookshelf', 'picture',
'counter', 'desk', 'curtain', 'refrigerator', 'shower curtain', 'toilet', 'sink',
'bathtub', 'otherfurniture')
def getInstanceInfo(self, xyz, instance_label, label):
ret = super().getInstanceInfo(xyz, instnace_label, label)
instance_num, instance_pointnum, instance_cls, pt_offset_label = ret
instance_cls = [x - 2 if x != -100 else x for x in instance_cls]
return instance_num, instance_pointnum, instance_cls, pt_offset_label