label_based_low_veg_relabelling.py

from typing import Any
from nibio_postprocessing.las_to_pandas import las_to_pandas
from nibio_postprocessing.pandas_to_las import pandas_to_las


class LabelBasedLowVegRelabelling(object):
    GROUND_CLASS = 1
    VEGETATION_CLASS = 2

    def __init__(self, input_las_file, output_las_file, verbose) -> None:
        self.input_las_file = input_las_file
        self.output_las_file = output_las_file
        self.verbose = verbose

    def relabel(self):
        # Read the file
        df = las_to_pandas(self.input_las_file, self.verbose, save_csv=False)

        # get all the rows which belong to the vegetation (2) class (label) and have treeID 0 and move it to ground class (1)
        df.loc[(df['label'] == self.VEGETATION_CLASS) & (df['treeID'] == 0), 'label'] = self.GROUND_CLASS

        # save the dataframe to a new las file
        pandas_to_las(df, self.output_las_file, csv_file_provided=False, verbose=self.verbose)

        if self.verbose:
            print(f"Relabelled {self.input_las_file} and saved to {self.output_las_file}")
    
    def run(self):
        self.relabel()
        if self.verbose:
            print("Relabelling complete.")

    def __call__(self) -> Any:
        self.run()


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='Relabel low vegetation to ground class')
    parser.add_argument('-i', '--input_las_file', help='Input LAS file', required=True)
    parser.add_argument('-o', '--output_las_file', help='Output LAS file', required=True)
    parser.add_argument('-v', '--verbose', action='store_true', help="Print information about the process")

    args = vars(parser.parse_args())

    input_las_file = args['input_las_file']
    output_las_file = args['output_las_file']
    verbose = args['verbose']

    # create an instance of LabelBasedLowVegRelabelling class
    label_based_low_veg_relabelling = LabelBasedLowVegRelabelling(
        input_las_file=input_las_file,
        output_las_file=output_las_file,
        verbose=verbose
    )

    # call main function
    label_based_low_veg_relabelling()