Upload folder using huggingface_hub

inference.py

@@ -27,9 +27,11 @@ def run_inference(model, cfg, transforms_dict, root_dir, split, scale, pc_tiling
         data_las = read_gridnet_tile(
             filepath, xyz=True, intensity=True, rgb=True, semantic=False, instance=False, remap=True
         )
+        data_las.initial_index = torch.arange(data_las.pos.shape[0]) # to keep initial order of points
 
         pos_list = []
         pred_list = []
+        indices_list = []
         pos_offset_init = None
         for x in range(2**pc_tiling): 
             data = SampleRecursiveMainXYAxisTiling(x=x, steps=pc_tiling)(data_las)
@@ -48,6 +50,7 @@ def run_inference(model, cfg, transforms_dict, root_dir, split, scale, pc_tiling
             # For full resolution level
             semantic_pred = output.full_res_semantic_pred(super_index_level0_to_level1=nag[0].super_index, sub_level0_to_raw=nag[0].sub)
             pos_list.append(data.pos.cpu())
+            indices_list.append(data.initial_index.cpu())
             
             pred_list.append(semantic_pred.cpu())
             
@@ -57,6 +60,13 @@ def run_inference(model, cfg, transforms_dict, root_dir, split, scale, pc_tiling
         merged_pos = torch.cat(pos_list, dim=0)
         merged_pred = torch.cat(pred_list, dim=0)
         merged_pos_offset = pos_offset_init + offset_initial_las
+        
+        # only for full res point cloud and keep initial order of points
+        merged_indices = torch.cat(indices_list, dim=0) 
+        sorted_indices = torch.argsort(merged_indices)
+        merged_pos = merged_pos[sorted_indices]
+        merged_pred = merged_pred[sorted_indices]
+        
 
         pos_data = (merged_pos.numpy() / scale).astype(int)
         x, y, z = pos_data[:, 0], pos_data[:, 1], pos_data[:, 2]
@@ -148,6 +158,8 @@ def main():
     parser.add_argument('--split', type=str, default='test', help="Data split to process (only used in inference mode) test or val split")
     parser.add_argument('--weights', type=str, required=True, help="Path to model checkpoint")
     parser.add_argument('--root_dir', type=str, required=True, help="Root directory of the dataset")
+    parser.add_argument('--pc_tiling', type=str, default='3', help="PC tiling for point cloud sampling")
+    
 
     args = parser.parse_args()
     cfg = init_config(overrides=["experiment=semantic/gridnet"])
@@ -158,15 +170,12 @@ def main():
     model = model.eval().cuda()
 
     SCALE = [0.001, 0.001, 0.001]
-    pc_tiling = 3
-    import time
-    start = time.time()
+    pc_tiling = int(args.pc_tiling)
 
     if args.mode == 'inference':
         run_inference(model, cfg, transforms_dict, args.root_dir, args.split, SCALE, pc_tiling)
     elif args.mode == 'export_log':
         export_logits(model, cfg, transforms_dict, args.root_dir, SCALE, pc_tiling)
-    print(time.time()-start)
 
 
 if __name__ == '__main__':