octree.py

#!/usr/bin/env python

import numpy as np
from open3d import *


class Node:
    def __init__(self, center = [0,0,0], size=0.0, children=[], data=None):
        self.center =  center
        self.size = size 
        self.children = children
        self.data = None

    def __repr__(self):
        return "Node()"

    def __str__(self):
        return "Node with Center: {} Size: {} IsLeaf: {}".format(self.center, self.size, self.isleaf())


    def isleaf(self):
        return self.data is not None

    def split(self):
        halfsize = self.size/2
        self.children = [
            Node([self.center[0]-halfsize, self.center[1]-halfsize, self.center[2]-halfsize], halfsize),
            Node([self.center[0]+halfsize, self.center[1]-halfsize, self.center[2]-halfsize], halfsize),
            Node([self.center[0]+halfsize, self.center[1]+halfsize, self.center[2]-halfsize], halfsize),
            Node([self.center[0]+halfsize, self.center[1]+halfsize, self.center[2]+halfsize], halfsize),
            Node([self.center[0]-halfsize, self.center[1]+halfsize, self.center[2]+halfsize], halfsize),
            Node([self.center[0]-halfsize, self.center[1]-halfsize, self.center[2]+halfsize], halfsize),
            Node([self.center[0]+halfsize, self.center[1]-halfsize, self.center[2]+halfsize], halfsize),
            Node([self.center[0]-halfsize, self.center[1]+halfsize, self.center[2]-halfsize], halfsize)
        ]

        self.add_to_child(self.data)
        self.data = None

    def add(self, point=[]):
        if self.isleaf():
            self.split()

        if len(self.children) == 0:
            #implicitly makes this a leaf node
            self.data = point
            return

        self.add_to_child(point)

    def add_to_child(self, point):
        pts = [point, self.center]
        for child in self.children:
            pts.append(child.center)
            if child.contains(point):
                child.add(point)
                return

    def contains(self, point=[]):
        return (
            self.between(point[0], self.center[0]+self.size, self.center[0]-self.size) and
            self.between(point[1], self.center[1]+self.size, self.center[1]-self.size) and
            self.between(point[2], self.center[2]+self.size, self.center[2]-self.size)
        )


    def between(self, p1, p2, p3):
        return (p1<=p2 and p1>=p3) or (p1>=p2 and p1<=p3)

    def get_contents(self):
        contents = []
        if self.isleaf():
            contents.append(self.data)
        else:
            for child in self.children:
                contents += child.get_contents()

        return contents

    def get_geoms(self, level, maxlevel=3):
        if level>=maxlevel:
            return []

        geoms = [self.draw()]
        for child in self.children:
            geoms += child.get_geoms(level + 1, maxlevel)

        return geoms

    def draw(self):
        points = [
            [self.center[0]-self.size, self.center[1]-self.size, self.center[2]-self.size],
            [self.center[0]-self.size, self.center[1]+self.size, self.center[2]-self.size],
            [self.center[0]-self.size, self.center[1]-self.size, self.center[2]+self.size],
            [self.center[0]-self.size, self.center[1]+self.size, self.center[2]+self.size],
            [self.center[0]+self.size, self.center[1]-self.size, self.center[2]-self.size],
            [self.center[0]+self.size, self.center[1]+self.size, self.center[2]-self.size],
            [self.center[0]+self.size, self.center[1]-self.size, self.center[2]+self.size],
            [self.center[0]+self.size, self.center[1]+self.size, self.center[2]+self.size],
        ]
        lines = [[0,1],[0,2],[0,4],[1,3],
                 [1,5],[2,3],[2,6],[3,7],
                 [4,5],[4,6],[5,7],[6,7]]
        colors = [[0, 0, 1] for i in range(len(lines))]
        line_set = LineSet()
        line_set.points = Vector3dVector(points)
        line_set.lines = Vector2iVector(lines)
        line_set.colors = Vector3dVector(colors)

        return line_set

class Octree:

    def __init__(self, data=[]):
        #Get max/min points, call it bbox
        bbox = [np.amax(data, axis=0), np.amin(data, axis=0)]

        #Compute center as halfway between max/min
        center = [
            ( (bbox[0][0] - bbox[1][0]) /2 ) + bbox[1][0],
            ( (bbox[0][1] - bbox[1][1]) /2 ) + bbox[1][1],
            ( (bbox[0][2] - bbox[1][2]) /2 ) + bbox[1][2]
        ]

        #Make Cubic based on max delta between bounds 
        maxsize = 0.0
        for i in range(3):
            size = bbox[0][i] - bbox[1][i]
            if size>maxsize:
                maxsize = size

        self.root = Node(center, maxsize/2 + 0.001)

        for i in range(len(data)):
            self.root.add(data[i])
            if i % 100 == 0:
                print("indexed {} points".format(i))
        
    def nearest_neighbor(self, point, count):
        if not self.root.contains(point):
            return []


        #Build path to leaf node containing this datapoint
        curnode = self.root
        stack = [curnode]
        found_leaf = False
        while not found_leaf:

            if curnode.isleaf():
                found_leaf = True 
                continue

            for child in curnode.children:
                if child.contains(point):
                    curnode = child
                    stack.append(curnode)
                    break

        #Look at siblings first and backtrack 
        contenders = []
        for i in range(len(stack)-1):
            for sibling in stack[(-i)+1].children:
                contenders += sibling.get_contents()

            if len(contenders)>count:
                break

        contenders.sort(key=lambda c: np.linalg.norm(c-point))

        return contenders[:count]

if __name__ == "__main__":
    pcd = read_point_cloud("/home/ben/Open3D/examples/TestData/fragment.ply")
    pcd = voxel_down_sample(pcd, voxel_size = 0.05)

    points = np.asarray(pcd.points)

    ot = Octree(points)

    draw_geometries([pcd]+ot.root.get_geoms(0, 4))

    nn = ot.nearest_neighbor(points[10], 3)
    print(nn)