registration trains

Author: humanpose1

conf/models/registration/pointnet2.yaml

@@ -8,7 +8,7 @@ models:
       radii: [[0.2], [0.4]]
       nsamples: [[64], [64]]
       down_conv_nn:
-        [[[FEAT + 3, 64, 64, 128]], [[128+3, 128, 128, 256]]]
+        [[[FEAT, 64, 64, 128]], [[128+3, 128, 128, 256]]]
     mlp_cls:
       nn: [128, 256, 32]
       dropout: 0.5

src/datasets/registration/general3dmatch.py

@@ -134,9 +134,7 @@ def get_patch(self, idx):
             if(self.transform is not None):
                 data_source = self.transform(data_source)
                 data_target = self.transform(data_target)
-            batch = Batch.from_data_list([data_source, data_target])
-            batch.pair = batch.batch
-            batch.batch = None
+            batch = make_pair(data_source, data_target)
             batch = batch.contiguous().to(torch.float)
 
             return batch
@@ -164,9 +162,7 @@ def get_fragment(self, idx):
             if(self.transform is not None):
                 data_source = self.transform(data_source)
                 data_target = self.transform(data_target)
-            batch = Batch.from_data_list([data_source, data_target])
-            batch.pair = batch.batch
-            batch.batch = None
+            batch = make_pair(data_source, data_target)
             batch.y = torch.from_numpy(match['pair'])
             return batch.contiguous().to(torch.float)
 

src/datasets/registration/pair.py

@@ -17,5 +17,5 @@ def make_pair(data_source: Data, data_target: Data):
     for key_source in data_source.keys:
         batch[key_source+"_source"] = data_source[key_source]
     for key_target in data_target.keys:
-        batch[key_target+"_source"] = data_target[key_target]
+        batch[key_target+"_target"] = data_target[key_target]
     return batch.contiguous()

src/models/registration/pointnet2.py

@@ -49,6 +49,7 @@ def set_input(self, data):
             x = torch.cat([data.x_source, data.x_target], 0).transpose(1, 2).contiguous()
         else:
             x = None
+
         pos = torch.cat([data.pos_source, data.pos_target], 0)
         rang = torch.arange(0, data.pos_source.shape[0])
         labels = torch.cat([rang, rang], 0)

test.pt

@@ -45,7 +45,6 @@ def train_epoch(
     with Ctq(train_loader) as tq_train_loader:
         for i, data in enumerate(tq_train_loader):
             data = data.to(device)  # This takes time
-
             model.set_input(data)
             t_data = time.time() - iter_data_time