

In [25]:

    
import xml.etree.ElementTree as ET;
from Tools.Parsers.KdTreeXMLParser import KdTree

import re, os
import numpy as np
import math

from vispy import app
from vispy import scene
from vispy.visuals.transforms import STTransform, NullTransform, AffineTransform, BaseTransform
from vispy.color.colormap import *

from transforms3d import quaternions, axangles

def generateGridBoxLines(minP = np.array([-1.0,-1.0,-1.0]), 
              maxP = np.array([1.0,1.0,1.0]), 
              rotMatrix = None,
              subdivs=None, boundaryBox = True):
    
    centerPos = (maxP + minP)/2.0;
    extent = (maxP - minP);
    
    if boundaryBox:
        nLines = 12;
    else:
        nLinse = 0;
    # for each subdiv a quad is (4 lines) is drawn
    if subdivs:
       nLines += (subdivs[0]+1)*(subdivs[1]+1) + (subdivs[1]+1)*(subdivs[2]+1) + (subdivs[2]+1)*(subdivs[0]+1) - 4*3;

    lines = np.zeros((nLines*2,3), dtype=np.float32);
    
    # plot lines
    if boundaryBox:
        lines[0:24,0:3] = 0.5* np.array([ (-1, -1, -1),
                                        (1, -1, -1),
                                        #
                                        (1, -1, -1),
                                        (1, 1, -1),
                                        #
                                        (1, 1, -1),
                                        (-1, 1, -1), 
                                        #
                                        (-1, 1, -1),
                                        (-1, -1, -1), # end of first quad
                                        #
                                        (-1, -1, 1),
                                        (1, -1, 1),
                                        #
                                        (1, -1, 1),
                                        (1, 1, 1),
                                        #
                                        (1, 1, 1),
                                        (-1, 1, 1), 
                                        #
                                        (-1, 1, 1),
                                        (-1, -1, 1), # end of second quad
                                        #
                                        (-1, -1, -1),
                                        (-1, -1, 1),
                                        #
                                        (1, -1, -1),
                                        (1, -1, 1),
                                        #
                                        (-1, 1, -1),
                                        (-1, 1, 1), 
                                        #
                                        (1, 1, -1),
                                        (1, 1, 1)]); # end of vertical lines
    
    
    # add subdivs if necessary

    if subdivs is not None:
        counter = 12*2; # start after the cube lines
        dxyz = np.array([1.0 / subdivs[0],1.0 / subdivs[1],1.0 / subdivs[2]]);
        
        # y
        # |    +   +   +   
        # | *  *   *   *   *
        # | *  *   *   *   *
        # |    +   +   +  
        # ----------------> x
        #  * come first (, then the +
        
        # make x lines inside (*)
        for k in range(1,int(subdivs[1])):
            for l in range(0,int(subdivs[2])+1):
                lines[counter:counter+2,0:3] = np.array([ (-0.5, -0.5 + k*dxyz[1],  -0.5 + l*dxyz[2]),
                                                          (0.5, -0.5 + k*dxyz[1],  -0.5 + l*dxyz[2])  ]);
                counter+=2;
        # make the x lines (+)
        for k in (0,subdivs[1]):
            for l in range(1,int(subdivs[2])):
                lines[counter:counter+2,0:3] = np.array([ (-0.5, -0.5 + k*dxyz[1],  -0.5 + l*dxyz[2]),
                                                          (0.5, -0.5 + k*dxyz[1],  -0.5 + l*dxyz[2])  ]);
                counter+=2;
        
        
        # make y lines inside (*)     
        for k in range(1,int(subdivs[2])):
            for l in range(0,int(subdivs[0])+1):
                lines[counter:counter+2,0:3] = np.array([ (-0.5+ l*dxyz[0],  -0.5,  -0.5 + k*dxyz[2]),
                                                          (-0.5+ l*dxyz[0],  0.5,  -0.5 + k*dxyz[2])  ]);
                counter+=2;
        # make the y lines (+)
        for k in (0,subdivs[2]):
            for l in range(1,int(subdivs[0])):
                lines[counter:counter+2,0:3] = np.array([ (-0.5+ l*dxyz[0],  -0.5,  -0.5 + k*dxyz[2]),
                                                          (-0.5+ l*dxyz[0],  0.5,  -0.5 + k*dxyz[2])  ]);
                counter+=2;
                
        # make z lines inside
        for k in range(1,int(subdivs[0])):
            for l in range(0,int(subdivs[1])+1):
                lines[counter:counter+2,0:3] = np.array([ (-0.5+k*dxyz[0],  -0.5+l*dxyz[1],  -0.5),
                                                          (-0.5+k*dxyz[0],  -0.5+l*dxyz[1],  0.5)    ]); 
                counter+=2;
        # make the z lines (+)
        for k in (0,subdivs[0]):
            for l in range(1,int(subdivs[1])):
                lines[counter:counter+2,0:3] = np.array([ (-0.5+k*dxyz[0],  -0.5+l*dxyz[1],  -0.5),
                                                          (-0.5+k*dxyz[0],  -0.5+l*dxyz[1],  0.5)    ]); 
                counter+=2;
        

    lines[...] *= extent;
    if rotMatrix is not None:
        lines = np.dot(lines,rotMatrix.T); # [p1.T; p2.T] * R.transpose
    lines[...] += centerPos;    
    
    return lines;

class PlotKdTree:
    
    
    
    
   
    
    def plot(self,kdTree,plotPoints,plotSubDivs):
        
        topoColor=[1,0,0]
        opts={ 'topo':{'lc':topoColor,'mc':topoColor} , 
               'points':{'lc':'','ls':'','mc':topoColor,'ms':'.','mw':2}, 
               'subdivs' :{'ls':'+'} }
        
        canvas = scene.SceneCanvas(size=(800, 600), show=True, keys='interactive', )
        viewBox = canvas.central_widget.add_view(); # create view box
        viewBox.camera = scene.TurntableCamera(elevation=30, azimuth=30, up='+z')
        axis = scene.visuals.XYZAxis(parent=viewBox.scene);       
        
        # Plotting KDTREE
        levels = kdTree.leafs.keys();

        cm = self.generateColorMap(max=max(levels)-min(levels));

        # plot root box
        lines = generateGridBoxLines(minP=kdTree.aabb.min, maxP=kdTree.aabb.max);
        colors = np.ones((len(lines),4), dtype=np.float32); colors[...,0:3]*= opts['topo']['mc'];

        if plotSubDivs:

            # add all leafs
            for level,leafs in kdTree.leafs.items():
                for leaf in leafs:
                    l = generateGridBoxLines(minP=leaf.aabb.min, maxP=leaf.aabb.max);
                    color = np.random.rand(1,3);
                    c = np.ones((len(l),4), dtype=np.float32); c[...,0:3]= color;#cm.to_rgba(level);
                    colors = np.concatenate(( colors,  c));
                    lines = np.concatenate(( lines,  l));
                    #scene.visuals.Text(str(leaf.idx), parent=viewBox.scene, color=color, pos = leaf.aabb.center()[0:3]);
#                 # add all lines for the aabbs
#                 counter = 0;
#                 for level, aabbs in kdTree.aabbTree.items():
#                     if level == 0:
#                         continue;
#                     for aabb in aabbs:
#                         l = generateGridBoxLines(minP=aabb['min'], maxP=aabb['max']);
#                         c = np.ones((len(l),4), dtype=np.float32); c*= cm.to_rgba(level);
#                         colors = np.concatenate(( colors,  c));
#                         lines = np.concatenate(( lines,  l));
#                         counter+=1;

        scene.visuals.Line(pos=lines,color=colors,
                               parent=viewBox.scene,
                               name="kdTree",
                               connect='segments', method='gl', width=2)    

        app.run();
        
        
    def generateColorMap(self,max, min=0):
        hsl = colorMap = get_colormap('hsl') 
        colors = hsl.map( np.linspace(0,1,max) )
        return colors;