Cura/util/util3d.py

   1 import math
   2 import numpy
   3
   4 class Vector3(object):
   5         def __init__(self, x=0.0, y=0.0, z=0.0):
   6                 self.x = x
   7                 self.y = y
   8                 self.z = z
   9
  10         def __copy__(self):
  11                 return Vector3(self.x, self.y, self.z)
  12
  13         def copy(self):
  14                 return Vector3(self.x, self.y, self.z)
  15
  16         def __repr__(self):
  17                 return 'V[%s, %s, %s]' % ( self.x, self.y, self.z )
  18
  19         def __add__(self, v):
  20                 return Vector3( self.x + v.x, self.y + v.y, self.z + v.z )
  21
  22         def __sub__(self, v):
  23                 return Vector3( self.x - v.x, self.y - v.y, self.z - v.z )
  24
  25         def __mul__(self, v):
  26                 return Vector3( self.x * v, self.y * v, self.z * v )
  27
  28         def __div__(self, v):
  29                 return Vector3( self.x / v, self.y / v, self.z / v )
  30         __truediv__ = __div__
  31
  32         def __neg__(self):
  33                 return Vector3( - self.x, - self.y, - self.z )
  34
  35         def __iadd__(self, v):
  36                 self.x += v.x
  37                 self.y += v.y
  38                 self.z += v.z
  39                 return self
  40
  41         def __isub__(self, v):
  42                 self.x += v.x
  43                 self.y += v.y
  44                 self.z += v.z
  45                 return self
  46
  47         def __imul__(self, v):
  48                 self.x *= v
  49                 self.y *= v
  50                 self.z *= v
  51                 return self
  52
  53         def __idiv__(self, v):
  54                 self.x /= v
  55                 self.y /= v
  56                 self.z /= v
  57                 return self
  58
  59         def almostEqual(self, v):
  60                 return (abs(self.x - v.x) + abs(self.y - v.y) + abs(self.z - v.z)) < 0.00001
  61
  62         def cross(self, v):
  63                 return Vector3(self.y * v.z - self.z * v.y, -self.x * v.z + self.z * v.x, self.x * v.y - self.y * v.x)
  64
  65         def vsize(self):
  66                 return math.sqrt( self.x * self.x + self.y * self.y + self.z * self.z )
  67
  68         def normalize(self):
  69                 f = self.vsize()
  70                 if f != 0.0:
  71                         self.x /= f
  72                         self.y /= f
  73                         self.z /= f
  74
  75         def min(self, v):
  76                 return Vector3(min(self.x, v.x), min(self.y, v.y), min(self.z, v.z))
  77
  78         def max(self, v):
  79                 return Vector3(max(self.x, v.x), max(self.y, v.y), max(self.z, v.z))
  80
  81 class AABB(object):
  82         def __init__(self, vMin, vMax):
  83                 self.vMin = vMin
  84                 self.vMax = vMax
  85                 self.perimeter = numpy.sum(self.vMax - self.vMin)
  86
  87         def combine(self, aabb):
  88                 return AABB(numpy.minimum(self.vMin, aabb.vMin), numpy.maximum(self.vMax, aabb.vMax))
  89
  90         def overlap(self, aabb):
  91                 if aabb.vMin[0] - self.vMax[0] > 0.0 or aabb.vMin[1] - self.vMax[1] > 0.0 or aabb.vMin[2] - self.vMax[2] > 0.0:
  92                         return False
  93                 if self.vMin[0] - aabb.vMax[0] > 0.0 or self.vMin[1] - aabb.vMax[1] > 0.0 or self.vMin[2] - aabb.vMax[2] > 0.0:
  94                         return False
  95                 return True
  96
  97         def __repr__(self):
  98                 return "AABB:%s - %s" % (str(self.vMin), str(self.vMax))
  99
 100 class _AABBNode(object):
 101         def __init__(self, aabb):
 102                 self.child1 = None
 103                 self.child2 = None
 104                 self.parent = None
 105                 self.height = 0
 106                 self.aabb = aabb
 107
 108         def isLeaf(self):
 109                 return self.child1 == None
 110
 111 class AABBTree(object):
 112         def __init__(self):
 113                 self.root = None
 114
 115         def insert(self, aabb):
 116                 newNode = _AABBNode(aabb)
 117                 if self.root == None:
 118                         self.root = newNode
 119                         return
 120
 121                 node = self.root
 122                 while not node.isLeaf():
 123                         child1 = node.child1
 124                         child2 = node.child2
 125
 126                         area = node.aabb.perimeter
 127                         combinedAABB = node.aabb.combine(aabb)
 128                         combinedArea = combinedAABB.perimeter
 129
 130                         cost = 2.0 * combinedArea
 131                         inheritanceCost = 2.0 * (combinedArea - area)
 132
 133                         if child1.isLeaf():
 134                                 cost1 = aabb.combine(child1.aabb).perimeter + inheritanceCost
 135                         else:
 136                                 oldArea = child1.aabb.perimeter
 137                                 newArea = aabb.combine(child1.aabb).perimeter
 138                                 cost1 = (newArea - oldArea) + inheritanceCost
 139
 140                         if child2.isLeaf():
 141                                 cost2 = aabb.combine(child1.aabb).perimeter + inheritanceCost
 142                         else:
 143                                 oldArea = child2.aabb.perimeter
 144                                 newArea = aabb.combine(child2.aabb).perimeter
 145                                 cost2 = (newArea - oldArea) + inheritanceCost
 146
 147                         if cost < cost1 and cost < cost2:
 148                                 break
 149
 150                         if cost1 < cost2:
 151                                 node = child1
 152                         else:
 153                                 node = child2
 154
 155                 sibling = node
 156
 157                 # Create a new parent.
 158                 oldParent = sibling.parent
 159                 newParent = _AABBNode(aabb.combine(sibling.aabb))
 160                 newParent.parent = oldParent
 161                 newParent.height = sibling.height + 1
 162
 163                 if oldParent != None:
 164                         # The sibling was not the root.
 165                         if oldParent.child1 == sibling:
 166                                 oldParent.child1 = newParent
 167                         else:
 168                                 oldParent.child2 = newParent
 169
 170                         newParent.child1 = sibling
 171                         newParent.child2 = newNode
 172                         sibling.parent = newParent
 173                         newNode.parent = newParent
 174                 else:
 175                         # The sibling was the root.
 176                         newParent.child1 = sibling
 177                         newParent.child2 = newNode
 178                         sibling.parent = newParent
 179                         newNode.parent = newParent
 180                         self.root = newParent
 181
 182                 # Walk back up the tree fixing heights and AABBs
 183                 node = newNode.parent
 184                 while node != None:
 185                         node = self._balance(node)
 186
 187                         child1 = node.child1
 188                         child2 = node.child2
 189
 190                         node.height = 1 + max(child1.height, child2.height)
 191                         node.aabb = child1.aabb.combine(child2.aabb)
 192
 193                         node = node.parent
 194
 195         def _balance(self, A):
 196                 if A.isLeaf() or A.height < 2:
 197                         return A
 198
 199                 B = A.child1
 200                 C = A.child2
 201
 202                 balance = C.height - B.height
 203
 204                 # Rotate C up
 205                 if balance > 1:
 206                         F = C.child1;
 207                         G = C.child2;
 208
 209                         # Swap A and C
 210                         C.child1 = A;
 211                         C.parent = A.parent;
 212                         A.parent = C;
 213
 214                         # A's old parent should point to C
 215                         if C.parent != None:
 216                                 if C.parent.child1 == A:
 217                                         C.parent.child1 = C
 218                                 else:
 219                                         C.parent.child2 = C
 220                         else:
 221                                 self.root = C
 222
 223                         # Rotate
 224                         if F.height > G.height:
 225                                 C.child2 = F
 226                                 A.child2 = G
 227                                 G.parent = A
 228                                 A.aabb = B.aabb.combine(G.aabb)
 229                                 C.aabb = A.aabb.combine(F.aabb)
 230
 231                                 A.height = 1 + max(B.height, G.height)
 232                                 C.height = 1 + max(A.height, F.height)
 233                         else:
 234                                 C.child2 = G
 235                                 A.child2 = F
 236                                 F.parent = A
 237                                 A.aabb = B.aabb.combine(F.aabb)
 238                                 C.aabb = A.aabb.combine(G.aabb)
 239
 240                                 A.height = 1 + max(B.height, F.height)
 241                                 C.height = 1 + max(A.height, G.height)
 242
 243                         return C;
 244
 245                 # Rotate B up
 246                 if balance < -1:
 247                         D = B.child1
 248                         E = B.child2
 249
 250                         # Swap A and B
 251                         B.child1 = A
 252                         B.parent = A.parent
 253                         A.parent = B
 254
 255                         # A's old parent should point to B
 256                         if B.parent != None:
 257                                 if B.parent.child1 == A:
 258                                         B.parent.child1 = B
 259                                 else:
 260                                         B.parent.child2 = B
 261                         else:
 262                                 self.root = B
 263
 264                         # Rotate
 265                         if D.height > E.height:
 266                                 B.child2 = D
 267                                 A.child1 = E
 268                                 E.parent = A
 269                                 A.aabb = C.aabb.combine(E.aabb)
 270                                 B.aabb = A.aabb.combine(D.aabb)
 271
 272                                 A.height = 1 + max(C.height, E.height)
 273                                 B.height = 1 + max(A.height, D.height)
 274                         else:
 275                                 B.child2 = E
 276                                 A.child1 = D
 277                                 D.parent = A
 278                                 A.aabb = C.aabb.combine(D.aabb)
 279                                 B.aabb = A.aabb.combine(E.aabb)
 280
 281                                 A.height = 1 + max(C.height, D.height)
 282                                 B.height = 1 + max(A.height, E.height)
 283
 284                         return B
 285
 286                 return A
 287
 288         def query(self, aabb):
 289                 resultList = []
 290                 if self.root != None:
 291                         self._query(self.root, aabb, resultList)
 292                 return resultList
 293
 294         def _query(self, node, aabb, resultList):
 295                 if not aabb.overlap(node.aabb):
 296                         return
 297                 if node.isLeaf():
 298                         resultList.append(node.aabb)
 299                 else:
 300                         self._query(node.child1, aabb, resultList)
 301                         self._query(node.child2, aabb, resultList)
 302
 303         def __repr__(self):
 304                 s = "AABBTree:\n"
 305                 s += str(self.root.aabb)
 306                 return s
 307
 308 if __name__ == '__main__':
 309         tree = AABBTree()
 310         tree.insert(AABB(Vector3(0,0,0), Vector3(0,0,0)))
 311         tree.insert(AABB(Vector3(1,1,1), Vector3(1,1,1)))
 312         tree.insert(AABB(Vector3(0.5,0.5,0.5), Vector3(0.5,0.5,0.5)))
 313         print(tree)
 314         print(tree.query(AABB(Vector3(0,0,0), Vector3(0,0,0))))
 315