expension of instance segmentation with a ground added to the pipeline

nibio_preprocessing/add_ground_to_inst_seg.py

@@ -23,7 +23,7 @@ class AddGroundToInstSeg():
         # get header of inst seg
         header = inst_seg.header
         if self.verbose:
-            print("Header of instance segmentation: {}".format(header))
+            print("Header of instance segmentation: {}".format(header.version))
 
         # create new header
         new_header = laspy.LasHeader(point_format=point_format.id, version=header.version)
@@ -61,14 +61,14 @@ class AddGroundToInstSeg():
         for key in tmp_dict.keys():
             las[key] = tmp_dict[key]
 
-        return las
-
-    def main(self):
-        las = self.add_ground_to_inst_seg()
         las.write(self.output_file)
 
+        # if verbose is selected write the name of the output file the results will be saved to
         if self.verbose:
-            print("The ground has been added to the instance segmentation file !")
+            print("Results are saved to: {}".format(self.output_file))
+      
+
+        return las
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -84,4 +84,4 @@ if __name__ == "__main__":
         output_file_path=args.output_file,
         verbose=args.verbose
         )
-    add_ground_to_inst_seg.main()
+    add_ground_to_inst_seg.add_ground_to_inst_seg()

nibio_preprocessing/add_ground_to_inst_seg_folders.py

+import argparse
+import glob
+import os
+from nibio_preprocessing.add_ground_to_inst_seg import AddGroundToInstSeg
+
+
+class AddGroundToInstSegFolders():
+    """Add ground to instance segmentation"""
+    def __init__(self, sem_seg_folder_path, inst_seg_folder_path, output_folder_path, verbose) -> None:
+        self.sem_seg_folder = sem_seg_folder_path
+        self.inst_seg_folder = inst_seg_folder_path
+        self.output_folder = output_folder_path
+        self.verbose = verbose
+
+    def add_ground_to_inst_seg_folders(self):
+        # get all files in the folder
+        sem_seg_files = glob.glob(self.sem_seg_folder + '/*.las')
+        inst_seg_files = glob.glob(self.inst_seg_folder + '/*.las')
+
+        # print name of files
+       
+        if self.verbose:
+            print("Sem seg files: {}".format(sem_seg_files))
+            print("Inst seg files: {}".format(inst_seg_files))
+
+        # check if same amount of files if not throw an exception and abort the programm
+        if len(sem_seg_files) != len(inst_seg_files):
+            raise Exception("The amount of files in the semantic segmentation folder and the instance segmentation folder is not the same. Please check the input folders.")
+            
+        # if output folder does not exist create it
+        if not os.path.isdir(self.output_folder):
+            os.mkdir(self.output_folder)
+
+        # sort the files
+        sem_seg_files.sort()
+        inst_seg_files.sort()
+        # match each sem seg file to a corresponding inst seg file with same core file name
+        for sem_file in sem_seg_files:
+            for inst_file in inst_seg_files:
+                if os.path.basename(sem_file).split('.')[0] == os.path.basename(inst_file).split('.')[0]:
+                    file_name = os.path.basename(inst_file)
+                    # create output file path
+                    output_file_path = os.path.join(self.output_folder, file_name)
+            
+                    # create instance of AddGroundToInstSeg
+                    add_ground = AddGroundToInstSeg(
+                        sem_file, inst_file, output_file_path , self.verbose)
+                    # add ground to instance segmentation
+                    add_ground.add_ground_to_inst_seg()
+
+if __name__ == "__main__":
+    # parse input arguments
+    parser = argparse.ArgumentParser(
+        description='Add ground to instance segmentation')
+    parser.add_argument('--sem_seg_folder', dest='sem_seg_folder',
+                        type=str, help='Semantic segmentation folder', default='./data/')  
+    parser.add_argument('--inst_seg_folder', dest='inst_seg_folder',
+                        type=str, help='Instance segmentation folder', default='./data/')
+    parser.add_argument('--output_folder', dest='output_folder',
+                        type=str, help='Output folder', default='./output_folder/00')
+    parser.add_argument('--verbose', action='store_true', help="Print information about the process")
+    args = parser.parse_args()
+
+    # create instance of AddGroundToInstSegFolders
+    add_ground_to_inst_seg_folders = AddGroundToInstSegFolders(
+        sem_seg_folder_path=args.sem_seg_folder,
+        inst_seg_folder_path=args.inst_seg_folder,
+        output_folder_path=args.output_folder,
+        verbose=args.verbose
+    )
+    # add ground to instance segmentation
+    add_ground_to_inst_seg_folders.add_ground_to_inst_seg_folders()
+
+
+

run_all_fine_grained.sh

@@ -257,6 +257,18 @@ for segmented_point_cloud_in_ply in $data_folder/results/segmented_point_clouds/
     --writers.las.extra_dims=all
 done
 
+python nibio_preprocessing/add_ground_to_inst_seg_folders.py --sem_seg_folder sample_playground/results/segmented_point_clouds/ --inst_seg_folder sample_playground/results/instance_segmented_point_clouds/ --output_folder sample_playground/instance_seg_with_ground --verbose
+
+# create the instance segmented point clouds with ground folder
+mkdir -p $data_folder/results/instance_segmented_point_clouds_with_ground
+ 
+# to add the ground to the instance segmented point clouds
+python nibio_preprocessing/add_ground_to_inst_seg_folders.py \
+--sem_seg_folder $data_folder/results/segmented_point_clouds/ \
+--inst_seg_folder $data_folder/results/instance_segmented_point_clouds/ \
+--output_folder $data_folder/results/instance_segmented_point_clouds_with_ground \
+--verbose
+
 echo " "
 echo "Done"
 # print path to the results folder and the subfolders
@@ -264,4 +276,5 @@ echo "Results can be found here: $data_folder/results"
 echo "Results containing the input point clouds can be found here:  $data_folder/results/input_data"
 echo "Results containing the segmented point clouds can be found here:  $data_folder/results/segmented_point_clouds"
 echo "Results containing the instance segmented point clouds can be found here:  $data_folder/results/instance_segmented_point_clouds"
+echo "Results containing the instance segmented point clouds with ground can be found here:  $data_folder/results/instance_segmented_point_clouds_with_ground"