Labels can now be placed along next to the points within the datasetfile
This commit is contained in:
Si11ium
parent
54a5b48ddc
commit
97e36df1ba
@ -23,16 +23,15 @@ class CustomShapeNet(InMemoryDataset):
|
||||
categories = {key: val for val, key in enumerate(['Box', 'Cone', 'Cylinder', 'Sphere'])}
|
||||
modes = {key: val for val, key in enumerate(['train', 'test', 'predict'])}
|
||||
|
||||
def __init__(self, root, collate_per_segment=True, mode='train', transform=None, pre_filter=None, pre_transform=None,
|
||||
headers=True, has_variations=False, refresh=False):
|
||||
def __init__(self, root_dir, collate_per_segment=True, mode='train', transform=None, pre_filter=None,
|
||||
pre_transform=None, headers=True, has_variations=False, refresh=False, labels_within=False):
|
||||
assert mode in self.modes.keys(), f'"mode" must be one of {self.modes.keys()}'
|
||||
assert not (collate_per_segment and has_variations), 'Either use each element or pointclouds - with variations'
|
||||
self.has_headers = headers
|
||||
self.has_variations = has_variations
|
||||
self.collate_per_element = collate_per_segment
|
||||
self.mode = mode
|
||||
self.refresh = refresh
|
||||
super(CustomShapeNet, self).__init__(root, transform, pre_transform, pre_filter)
|
||||
|
||||
#Set the Dataset Parameters
|
||||
self.has_headers, self.has_variations, self.labels_within = headers, has_variations, labels_within
|
||||
self.collate_per_element, self.mode, self.refresh = collate_per_segment, mode, refresh
|
||||
super(CustomShapeNet, self).__init__(root_dir, transform, pre_transform, pre_filter)
|
||||
self.data, self.slices = self._load_dataset()
|
||||
print("Initialized")
|
||||
|
||||
@ -116,19 +115,6 @@ class CustomShapeNet(InMemoryDataset):
|
||||
if pattern.match(os.path.split(element)[-1]):
|
||||
continue
|
||||
else:
|
||||
# The following two lines were intendations
|
||||
# check if it is "just" a pc.dat or pc.xyz
|
||||
# if all([x not in os.path.split(element)[-1] for x in ['pc.dat', 'pc.xyz']]):
|
||||
|
||||
# Assign training data to the data container
|
||||
# Following the original logic;
|
||||
# y should be the label;
|
||||
# pos should be the six dimensional vector describing: !its pos not points!!
|
||||
# x,y,z,x_rot,y_rot,z_rot
|
||||
|
||||
# Get the y - Label
|
||||
y_raw = next(i for i, v in enumerate(self.categories.keys()) if v.lower() in element.lower())
|
||||
# y_raw = os.path.splitext(element)[0].split('_')[-2]
|
||||
with open(element,'r') as f:
|
||||
if self.has_headers:
|
||||
headers = f.__next__()
|
||||
@ -142,7 +128,15 @@ class CustomShapeNet(InMemoryDataset):
|
||||
points = torch.tensor(src, dtype=None).squeeze()
|
||||
if not len(points.shape) > 1:
|
||||
continue
|
||||
y_all = ([y_raw] if self.mode != 'predict' else [-1]) * points.shape[0]
|
||||
# Place Fake Labels to hold the given structure
|
||||
if self.labels_within:
|
||||
y_all = points[:, -1]
|
||||
points = points[:, :-1]
|
||||
else:
|
||||
# Get the y - Label
|
||||
y_raw = next(i for i, v in enumerate(self.categories.keys()) if v.lower() in element.lower())
|
||||
y_all = ([y_raw] if self.mode != 'predict' else [-1]) * points.shape[0]
|
||||
|
||||
y = torch.as_tensor(y_all, dtype=torch.int)
|
||||
if self.collate_per_element:
|
||||
data = Data(y=y, pos=points[:, :3]) # , points=points, norm=points[:, 3:])
|
||||
@ -178,12 +172,11 @@ class ShapeNetPartSegDataset(Dataset):
|
||||
Resample raw point cloud to fixed number of points.
|
||||
Map raw label from range [1, N] to [0, N-1].
|
||||
"""
|
||||
def __init__(self, root_dir, collate_per_segment=True, mode='train', transform=None, refresh=False,
|
||||
has_variations=False, npoints=1024, headers=True):
|
||||
def __init__(self, root_dir, npoints=1024, **kwargs):
|
||||
super(ShapeNetPartSegDataset, self).__init__()
|
||||
kwargs.update(dict(root_dir=root_dir))
|
||||
self.npoints = npoints
|
||||
self.dataset = CustomShapeNet(root=root_dir, collate_per_segment=collate_per_segment, refresh=refresh,
|
||||
mode=mode, transform=transform, headers=headers, has_variations=has_variations)
|
||||
self.dataset = CustomShapeNet(**kwargs)
|
||||
|
||||
def __getitem__(self, index):
|
||||
data = self.dataset[index]
|
||||
|
||||
39
main.py
39
main.py
@ -33,6 +33,7 @@ parser.add_argument('--npoints', type=int, default=1024, help='resample points n
|
||||
parser.add_argument('--model', type=str, default='', help='model path')
|
||||
parser.add_argument('--nepoch', type=int, default=250, help='number of epochs to train for')
|
||||
parser.add_argument('--outf', type=str, default='checkpoint', help='output folder')
|
||||
parser.add_argument('--labels_within', type=strtobool, default=False, help='defines the label location')
|
||||
parser.add_argument('--batch_size', type=int, default=8, help='input batch size')
|
||||
parser.add_argument('--test_per_batches', type=int, default=1000, help='run a test batch per training batches number')
|
||||
parser.add_argument('--num_workers', type=int, default=4, help='number of data loading workers')
|
||||
@ -56,30 +57,34 @@ torch.manual_seed(opt.manual_seed)
|
||||
torch.cuda.manual_seed(opt.manual_seed)
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
# Dataset and transform
|
||||
print('Construct dataset ..')
|
||||
if True:
|
||||
rot_max_angle = 15
|
||||
trans_max_distance = 0.01
|
||||
|
||||
RotTransform = GT.Compose([GT.RandomRotate(rot_max_angle, 0),
|
||||
GT.RandomRotate(rot_max_angle, 1),
|
||||
GT.RandomRotate(rot_max_angle, 2)]
|
||||
)
|
||||
TransTransform = GT.RandomTranslate(trans_max_distance)
|
||||
rot_max_angle = 15
|
||||
trans_max_distance = 0.01
|
||||
|
||||
train_transform = GT.Compose([GT.NormalizeScale(), RotTransform, TransTransform])
|
||||
test_transform = GT.Compose([GT.NormalizeScale(), ])
|
||||
RotTransform = GT.Compose([GT.RandomRotate(rot_max_angle, 0),
|
||||
GT.RandomRotate(rot_max_angle, 1),
|
||||
GT.RandomRotate(rot_max_angle, 2)]
|
||||
)
|
||||
|
||||
dataset = ShapeNetPartSegDataset(root_dir=opt.dataset, collate_per_segment=opt.collate_per_segment,
|
||||
mode='train', transform=train_transform, npoints=opt.npoints,
|
||||
has_variations=opt.has_variations, headers=opt.headers)
|
||||
TransTransform = GT.RandomTranslate(trans_max_distance)
|
||||
train_transform = GT.Compose([GT.NormalizeScale(), RotTransform, TransTransform])
|
||||
test_transform = GT.Compose([GT.NormalizeScale(), ])
|
||||
|
||||
params = dict(root_dir=opt.dataset,
|
||||
collate_per_segment=opt.collate_per_segment,
|
||||
transform=train_transform,
|
||||
npoints=opt.npoints,
|
||||
labels_within=opt.labels_within,
|
||||
has_variations=opt.has_variations,
|
||||
headers=opt.headers
|
||||
)
|
||||
|
||||
dataset = ShapeNetPartSegDataset(mode='train', **params)
|
||||
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,
|
||||
mode='test', transform=test_transform, npoints=opt.npoints,
|
||||
has_variations=opt.has_variations, headers=opt.headers)
|
||||
test_dataset = ShapeNetPartSegDataset(mode='test', **params)
|
||||
test_dataLoader = DataLoader(test_dataset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers)
|
||||
|
||||
num_classes = dataset.num_classes()
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user