steffen/Labeling-Primitives-With-Point2Net
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Labels can now be placed along next to the points within the datasetfile

Browse Source
This commit is contained in:
Si11ium
2019-08-06 08:36:59 +02:00
parent 54a5b48ddc
commit 97e36df1ba
2 changed files with 41 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
main.py
View File

@ -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()