Source code for ddd_pylib.points_manip._merge_close_points_operator

from ddd_pylib._base import PointsOperator
from scipy.spatial import KDTree
import pandas as pd




class MergeClosePointsOperator(PointsOperator):
    def __init__(self):
        super().__init__()
        self._distanceThreshold = self.defaultDistanceThreshold()



    @staticmethod
    def defaultDistanceThreshold():
        return 1.0




    def getDistanceThreshold(self):
        """
        Return the distance threshold value
        """
        return self._distanceThreshold




    def setDistanceThreshold(self, threshold):
        """
        Set the distance threshold value

        Args:
            threshold: The distance threshold value
        """
        self._distanceThreshold = threshold




    def getPrefix(self):
        return "Merged-"


    def _mergeByDistance(self, df: pd.DataFrame, threshold: float) -> pd.DataFrame:
        spatial_axes = [ax for ax in self.getValidAxes() if ax != 'T']
        used_points = set()
        merged_points = []

        if df.empty:
            return df
        
        tree = KDTree(df[spatial_axes].values)

        for i, point in df[spatial_axes].iterrows():
            if i in used_points:
                continue
            indices = tree.query_ball_point(point.values, threshold)
            merged_point = df.iloc[indices].mean()
            merged_points.append(merged_point)
            used_points.update(indices)

        return pd.DataFrame(merged_points, columns=df.columns)

    def _applyToFrame(self, frameData, t, nT):
        df = frameData[0]
        t = self.getDistanceThreshold()
        return self._mergeByDistance(df, t)



if __name__ == "__main__":
    import pandas as pd
    import numpy as np
    import napari

    points = np.random.rand(1500, 3) * 100
    df = pd.DataFrame(points, columns=['X', 'Y', 'Z'])

    op = MergeClosePointsOperator()
    op.setMainDataFrame(df)
    op.setDistanceThreshold(5.0)
    op.run()

    res = op.getResult()

    viewer = napari.Viewer()
    viewer.add_points(
        df[['Z', 'Y', 'X']].values, 
        size=1, 
        face_color='red', 
        name='Original Points'
    )
    viewer.add_points(
        res[['Z', 'Y', 'X']].values, 
        size=1, 
        face_color='transparent',
        border_color='green', 
        name='Merged Points'
    )
    napari.run()