Merge branch 'CapnBry-gcodeinterpreter-upstream' into devel

master
Gina Häußge 2013-08-11 17:39:53 +02:00
commit ec9c7adf70
4 changed files with 168 additions and 488 deletions

View File

@ -63,8 +63,8 @@ class GcodeManager:
dirty = True dirty = True
if gcode.extrusionAmount: if gcode.extrusionAmount:
analysisResult["filament"] = "%.2fm" % (gcode.extrusionAmount / 1000) analysisResult["filament"] = "%.2fm" % (gcode.extrusionAmount / 1000)
if gcode.extrusionVolume: if gcode.calculateVolumeCm3():
analysisResult["filament"] += " / %.2fcm³" % gcode.extrusionVolume analysisResult["filament"] += " / %.2fcm³" % gcode.calculateVolumeCm3()
dirty = True dirty = True
if dirty: if dirty:

View File

@ -2,6 +2,5 @@ The code in this sub package mostly originates from the Cura project (https://gi
slightly reorganized and adapted. The mapping to the original Cura source is the following: slightly reorganized and adapted. The mapping to the original Cura source is the following:
* avr_isp.* => Cura.avr_isp.* * avr_isp.* => Cura.avr_isp.*
* comm => Cura.util.machineCom * comm => Cura.util.machineCom (highly modified now)
* gcodeInterpreter => Cura.util.gcodeInterpreter * gcodeInterpreter => Cura.util.gcodeInterpreter
* util3d => Cura.util.util3d

View File

@ -1,15 +1,13 @@
from __future__ import absolute_import from __future__ import absolute_import
__copyright__ = "Copyright (C) 2013 David Braam - Released under terms of the AGPLv3 License"
import sys import sys
import math import math
import re
import os import os
import base64 import base64
import zlib import zlib
import logging import logging
from octoprint.util import util3d
preferences = { preferences = {
"extruder_offset_x1": -22.0, "extruder_offset_x1": -22.0,
"extruder_offset_y1": 0.0, "extruder_offset_y1": 0.0,
@ -28,232 +26,224 @@ def getPreference(key, default=None):
class AnalysisAborted(Exception): class AnalysisAborted(Exception):
pass pass
class gcodePath(object): def gcodePath(newType, pathType, layerThickness, startPoint):
def __init__(self, newType, pathType, layerThickness, startPoint): return {'type': newType,
self.type = newType 'pathType': pathType,
self.pathType = pathType 'layerThickness': layerThickness,
self.layerThickness = layerThickness 'points': [startPoint],
self.list = [startPoint] 'extrusion': [0.0]}
class gcode(object): class gcode(object):
def __init__(self): def __init__(self):
self._logger = logging.getLogger(__name__) self._logger = logging.getLogger(__name__)
self.regMatch = {} self.layerList = None
self.layerList = []
self.extrusionAmount = 0 self.extrusionAmount = 0
self.extrusionVolume = None
self.totalMoveTimeMinute = 0 self.totalMoveTimeMinute = 0
self.filename = None
self.progressCallback = None self.progressCallback = None
self._abort = False self._abort = False
self._filamentDiameter = 0
def load(self, filename): def load(self, filename):
if os.path.isfile(filename): if os.path.isfile(filename):
self.filename = filename
self._fileSize = os.stat(filename).st_size self._fileSize = os.stat(filename).st_size
gcodeFile = open(filename, 'r') gcodeFile = open(filename, 'r')
self._load(gcodeFile) self._load(gcodeFile)
gcodeFile.close() gcodeFile.close()
def loadList(self, l): def loadList(self, l):
self.filename = None
self._load(l) self._load(l)
def abort(self): def abort(self):
self._abort = True self._abort = True
def calculateVolumeCm3(self):
radius = self._filamentDiameter / 2
return (self.extrusionAmount * (math.pi * radius * radius)) / 1000
def calculateWeight(self):
#Calculates the weight of the filament in kg
volumeM3 = self.calculateVolumeCm3() /(1000*1000)
return volumeM3 * getPreference('filament_physical_density')
def calculateCost(self):
cost_kg = getPreference('filament_cost_kg')
cost_meter = getPreference('filament_cost_meter')
if cost_kg > 0.0 and cost_meter > 0.0:
return "%.2f / %.2f" % (self.calculateWeight() * cost_kg, self.extrusionAmount / 1000 * cost_meter)
elif cost_kg > 0.0:
return "%.2f" % (self.calculateWeight() * cost_kg)
elif cost_meter > 0.0:
return "%.2f" % (self.extrusionAmount / 1000 * cost_meter)
return None
def _load(self, gcodeFile): def _load(self, gcodeFile):
filePos = 0 filePos = 0
pos = util3d.Vector3() self.layerList = []
posOffset = util3d.Vector3() pos = [0.0, 0.0, 0.0]
posOffset = [0.0, 0.0, 0.0]
currentE = 0.0 currentE = 0.0
totalExtrusion = 0.0 totalExtrusion = 0.0
maxExtrusion = 0.0 maxExtrusion = 0.0
currentExtruder = 0 currentExtruder = 0
extrudeAmountMultiply = 1.0
totalMoveTimeMinute = 0.0 totalMoveTimeMinute = 0.0
filamentDiameter = 0.0 absoluteE = True
scale = 1.0 scale = 1.0
posAbs = True posAbs = True
posAbsExtruder = True; feedRate = 3600.0
feedRate = 3600 unknownGcodes = {}
layerThickness = 0.1 unknownMcodes = {}
pathType = 'CUSTOM';
currentLayer = []
unknownGcodes={}
unknownMcodes={}
currentPath = gcodePath('move', pathType, layerThickness, pos.copy())
currentPath.list[0].e = totalExtrusion
currentPath.list[0].extrudeAmountMultiply = extrudeAmountMultiply
currentLayer.append(currentPath)
for line in gcodeFile: for line in gcodeFile:
if self._abort: if self._abort:
raise AnalysisAborted() raise AnalysisAborted()
if type(line) is tuple: if type(line) is tuple:
line = line[0] line = line[0]
if self.progressCallback != None: filePos += 1
if self.progressCallback is not None and (filePos % 1000 == 0):
if isinstance(gcodeFile, (file)): if isinstance(gcodeFile, (file)):
self.progressCallback(float(filePos) / float(self._fileSize)) self.progressCallback(float(gcodeFile.tell()) / float(self._fileSize))
elif isinstance(gcodeFile, (list)): elif isinstance(gcodeFile, (list)):
self.progressCallback(float(filePos) / float(len(gcodeFile))) self.progressCallback(float(filePos) / float(len(gcodeFile)))
filePos += 1
#Parse Cura_SF comments
if line.startswith(';TYPE:'):
pathType = line[6:].strip()
if pathType != "CUSTOM":
startCodeDone = True
if ';' in line: if ';' in line:
# Slic3r GCode comment parser
comment = line[line.find(';')+1:].strip() comment = line[line.find(';')+1:].strip()
if comment == 'fill': if comment.startswith("filament_diameter"):
pathType = 'FILL' self._filamentDiameter = float(comment.split("=", 1)[1].strip())
elif comment == 'perimeter':
pathType = 'WALL-INNER'
elif comment == 'skirt':
pathType = 'SKIRT'
elif comment.startswith("filament_diameter"):
filamentDiameter = float(line.split("=", 1)[1].strip())
# Cura Gcode comment parser
if comment.startswith('LAYER:'):
self.layerList.append(currentLayer)
currentLayer = []
elif comment.startswith("CURA_PROFILE_STRING"): elif comment.startswith("CURA_PROFILE_STRING"):
curaOptions = self._parseCuraProfileString(comment) curaOptions = self._parseCuraProfileString(comment)
if "filament_diameter" in curaOptions:
if "filament_diameter" in curaOptions.keys():
try: try:
filamentDiameter = float(curaOptions["filament_diameter"]) self._filamentDiameter = float(curaOptions["filament_diameter"])
except: except:
filamentDiameter = 0.0 self._filamentDiameter = 0.0
line = line[0:line.find(';')] line = line[0:line.find(';')]
T = self.getCodeInt(line, 'T')
T = getCodeInt(line, 'T')
if T is not None: if T is not None:
if currentExtruder > 0: if currentExtruder > 0:
posOffset.x -= getPreference('extruder_offset_x%d' % (currentExtruder), 0.0) posOffset[0] -= getPreference('extruder_offset_x%d' % (currentExtruder), 0.0)
posOffset.y -= getPreference('extruder_offset_y%d' % (currentExtruder), 0.0) posOffset[1] -= getPreference('extruder_offset_y%d' % (currentExtruder), 0.0)
currentExtruder = T currentExtruder = T
if currentExtruder > 0: if currentExtruder > 0:
posOffset.x += getPreference('extruder_offset_x%d' % (currentExtruder), 0.0) posOffset[0] += getPreference('extruder_offset_x%d' % (currentExtruder), 0.0)
posOffset.y += getPreference('extruder_offset_y%d' % (currentExtruder), 0.0) posOffset[1] += getPreference('extruder_offset_y%d' % (currentExtruder), 0.0)
G = self.getCodeInt(line, 'G') G = getCodeInt(line, 'G')
if G is not None: if G is not None:
if G == 0 or G == 1: #Move if G == 0 or G == 1: #Move
x = self.getCodeFloat(line, 'X') x = getCodeFloat(line, 'X')
y = self.getCodeFloat(line, 'Y') y = getCodeFloat(line, 'Y')
z = self.getCodeFloat(line, 'Z') z = getCodeFloat(line, 'Z')
e = self.getCodeFloat(line, 'E') e = getCodeFloat(line, 'E')
f = self.getCodeFloat(line, 'F') f = getCodeFloat(line, 'F')
oldPos = pos.copy() oldPos = pos
if x is not None: pos = pos[:]
if posAbs: if posAbs:
pos.x = x * scale + posOffset.x if x is not None:
else: pos[0] = x * scale + posOffset[0]
pos.x += x * scale if y is not None:
if y is not None: pos[1] = y * scale + posOffset[1]
if posAbs: if z is not None:
pos.y = y * scale + posOffset.y pos[2] = z * scale + posOffset[2]
else: else:
pos.y += y * scale if x is not None:
if z is not None: pos[0] += x * scale
if posAbs: if y is not None:
pos.z = z * scale + posOffset.z pos[1] += y * scale
else: if z is not None:
pos.z += z * scale pos[2] += z * scale
if f is not None: if f is not None:
feedRate = f feedRate = f
if x is not None or y is not None or z is not None: if x is not None or y is not None or z is not None:
totalMoveTimeMinute += (oldPos - pos).vsize() / feedRate diffX = oldPos[0] - pos[0]
diffY = oldPos[1] - pos[1]
totalMoveTimeMinute += math.sqrt(diffX * diffX + diffY * diffY) / feedRate
moveType = 'move' moveType = 'move'
if e is not None: if e is not None:
if posAbsExtruder: if absoluteE:
if e > currentE: e -= currentE
moveType = 'extrude' if e > 0.0:
if e < currentE: moveType = 'extrude'
moveType = 'retract' if e < 0.0:
totalExtrusion += e - currentE moveType = 'retract'
currentE = e totalExtrusion += e
else: currentE += e
if e > 0:
moveType = 'extrude'
if e < 0:
moveType = 'retract'
totalExtrusion += e
currentE += e
if totalExtrusion > maxExtrusion: if totalExtrusion > maxExtrusion:
maxExtrusion = totalExtrusion maxExtrusion = totalExtrusion
if moveType == 'move' and oldPos.z != pos.z: else:
if oldPos.z > pos.z and abs(oldPos.z - pos.z) > 5.0 and pos.z < 1.0: e = 0.0
oldPos.z = 0.0 if moveType == 'move' and oldPos[2] != pos[2]:
layerThickness = abs(oldPos.z - pos.z) if oldPos[2] > pos[2] and abs(oldPos[2] - pos[2]) > 5.0 and pos[2] < 1.0:
if currentPath.type != moveType or currentPath.pathType != pathType: oldPos[2] = 0.0
currentPath = gcodePath(moveType, pathType, layerThickness, currentPath.list[-1])
currentLayer.append(currentPath)
newPos = pos.copy()
newPos.e = totalExtrusion
newPos.extrudeAmountMultiply = extrudeAmountMultiply
currentPath.list.append(newPos)
elif G == 4: #Delay elif G == 4: #Delay
S = self.getCodeFloat(line, 'S') S = getCodeFloat(line, 'S')
if S is not None: if S is not None:
totalMoveTimeMinute += S / 60 totalMoveTimeMinute += S / 60.0
P = self.getCodeFloat(line, 'P') P = getCodeFloat(line, 'P')
if P is not None: if P is not None:
totalMoveTimeMinute += P / 60 / 1000 totalMoveTimeMinute += P / 60.0 / 1000.0
elif G == 20: #Units are inches elif G == 20: #Units are inches
scale = 25.4 scale = 25.4
elif G == 21: #Units are mm elif G == 21: #Units are mm
scale = 1.0 scale = 1.0
elif G == 28: #Home elif G == 28: #Home
x = self.getCodeFloat(line, 'X') x = getCodeFloat(line, 'X')
y = self.getCodeFloat(line, 'Y') y = getCodeFloat(line, 'Y')
z = self.getCodeFloat(line, 'Z') z = getCodeFloat(line, 'Z')
if x is None and y is None and z is None: if getPreference('machine_center_is_zero') == 'True':
pos = util3d.Vector3() center = [getPreference('machine_width') / 2, getPreference('machine_depth') / 2,0.0]
else: else:
center = [0.0,0.0,0.0]
if x is None and y is None and z is None:
pos = center
else:
pos = pos[:]
if x is not None: if x is not None:
pos.x = 0.0 pos[0] = center[0]
if y is not None: if y is not None:
pos.y = 0.0 pos[1] = center[1]
if z is not None: if z is not None:
pos.z = 0.0 pos[2] = center[2]
elif G == 90: #Absolute position elif G == 90: #Absolute position
posAbs = True posAbs = True
posAbsExtruder = True
elif G == 91: #Relative position elif G == 91: #Relative position
posAbs = False posAbs = False
posAbsExtruder = False
elif G == 92: elif G == 92:
x = self.getCodeFloat(line, 'X') x = getCodeFloat(line, 'X')
y = self.getCodeFloat(line, 'Y') y = getCodeFloat(line, 'Y')
z = self.getCodeFloat(line, 'Z') z = getCodeFloat(line, 'Z')
e = self.getCodeFloat(line, 'E') e = getCodeFloat(line, 'E')
if e is not None: if e is not None:
currentE = e currentE = e
if x is not None: if x is not None:
posOffset.x = pos.x - x posOffset[0] = pos[0] - x
if y is not None: if y is not None:
posOffset.y = pos.y - y posOffset[1] = pos[1] - y
if z is not None: if z is not None:
posOffset.z = pos.z - z posOffset[2] = pos[2] - z
else: else:
if G not in unknownGcodes: if G not in unknownGcodes:
self._logger.info("Unknown G code: %r" % G) self._logger.info("Unknown G code: %r" % G)
unknownGcodes[G] = True unknownGcodes[G] = True
else: else:
M = self.getCodeInt(line, 'M') M = getCodeInt(line, 'M')
if M is not None: if M is not None:
if M == 1: #Message with possible wait (ignored) if M == 0: #Message with possible wait (ignored)
pass
elif M == 1: #Message with possible wait (ignored)
pass pass
elif M == 80: #Enable power supply elif M == 80: #Enable power supply
pass pass
elif M == 81: #Suicide/disable power supply elif M == 81: #Suicide/disable power supply
pass pass
elif M == 82: # Use absolute extruder positions elif M == 82: #Absolute E
posAbsExtruder = True absoluteE = True
elif M == 83: # Use relative extruder positions elif M == 83: #Relative E
posAbsExtruder = False absoluteE = False
elif M == 84: #Disable step drivers elif M == 84: #Disable step drivers
pass pass
elif M == 92: #Set steps per unit elif M == 92: #Set steps per unit
@ -278,50 +268,58 @@ class gcode(object):
pass pass
elif M == 113: #Extruder PWM (these should not be in the final GCode, but they are) elif M == 113: #Extruder PWM (these should not be in the final GCode, but they are)
pass pass
elif M == 117: #LCD message
pass
elif M == 140: #Set bed temperature elif M == 140: #Set bed temperature
pass pass
elif M == 190: #Set bed temperature & wait elif M == 190: #Set bed temperature & wait
pass pass
elif M == 221: #Extrude amount multiplier elif M == 221: #Extrude amount multiplier
s = self.getCodeFloat(line, 'S') s = getCodeFloat(line, 'S')
if s != None:
extrudeAmountMultiply = s / 100.0
else: else:
if M not in unknownMcodes: if M not in unknownMcodes:
self._logger.info("Unknown M code: %r" % M) self._logger.info("Unknown M code: %r" % M)
unknownMcodes[M] = True unknownMcodes[M] = True
self.layerList.append(currentLayer) if self.progressCallback is not None:
self.progressCallback(100.0)
self.extrusionAmount = maxExtrusion self.extrusionAmount = maxExtrusion
if filamentDiameter is not None and filamentDiameter > 0:
self.extrusionVolume = math.pi * math.pow(filamentDiameter / 2.0, 2) * maxExtrusion / 1000.0
self.totalMoveTimeMinute = totalMoveTimeMinute self.totalMoveTimeMinute = totalMoveTimeMinute
def getCodeInt(self, line, code):
if code not in self.regMatch:
self.regMatch[code] = re.compile(code + '([^\s]+)')
m = self.regMatch[code].search(line)
if m == None:
return None
try:
return int(m.group(1))
except:
return None
def getCodeFloat(self, line, code):
if code not in self.regMatch:
self.regMatch[code] = re.compile(code + '([^\s]+)')
m = self.regMatch[code].search(line)
if m == None:
return None
try:
return float(m.group(1))
except:
return None
def _parseCuraProfileString(self, comment): def _parseCuraProfileString(self, comment):
return {key: value for (key, value) in map(lambda x: x.split("=", 1), zlib.decompress(base64.b64decode(comment[len("CURA_PROFILE_STRING:"):])).split("\b"))} return {key: value for (key, value) in map(lambda x: x.split("=", 1), zlib.decompress(base64.b64decode(comment[len("CURA_PROFILE_STRING:"):])).split("\b"))}
if __name__ == '__main__': def getCodeInt(line, code):
for filename in sys.argv[1:]: n = line.find(code) + 1
gcode().load(filename) if n < 1:
return None
m = line.find(' ', n)
try:
if m < 0:
return int(line[n:])
return int(line[n:m])
except:
return None
def getCodeFloat(line, code):
n = line.find(code) + 1
if n < 1:
return None
m = line.find(' ', n)
try:
if m < 0:
return float(line[n:])
return float(line[n:m])
except:
return None
if __name__ == '__main__':
from time import time
t = time()
for filename in sys.argv[1:]:
g = gcode()
g.load(filename)
print g.totalMoveTimeMinute
print g.extrusionAmount
print g.calculateVolumeCm3()
print time() - t

View File

@ -1,317 +0,0 @@
from __future__ import absolute_import
import math
import numpy
class Vector3(object):
def __init__(self, x=0.0, y=0.0, z=0.0):
self.x = x
self.y = y
self.z = z
def __copy__(self):
return Vector3(self.x, self.y, self.z)
def copy(self):
return Vector3(self.x, self.y, self.z)
def __repr__(self):
return 'V[%s, %s, %s]' % ( self.x, self.y, self.z )
def __add__(self, v):
return Vector3( self.x + v.x, self.y + v.y, self.z + v.z )
def __sub__(self, v):
return Vector3( self.x - v.x, self.y - v.y, self.z - v.z )
def __mul__(self, v):
return Vector3( self.x * v, self.y * v, self.z * v )
def __div__(self, v):
return Vector3( self.x / v, self.y / v, self.z / v )
__truediv__ = __div__
def __neg__(self):
return Vector3( - self.x, - self.y, - self.z )
def __iadd__(self, v):
self.x += v.x
self.y += v.y
self.z += v.z
return self
def __isub__(self, v):
self.x += v.x
self.y += v.y
self.z += v.z
return self
def __imul__(self, v):
self.x *= v
self.y *= v
self.z *= v
return self
def __idiv__(self, v):
self.x /= v
self.y /= v
self.z /= v
return self
def almostEqual(self, v):
return (abs(self.x - v.x) + abs(self.y - v.y) + abs(self.z - v.z)) < 0.00001
def cross(self, v):
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)
def vsize(self):
return math.sqrt( self.x * self.x + self.y * self.y + self.z * self.z )
def normalize(self):
f = self.vsize()
if f != 0.0:
self.x /= f
self.y /= f
self.z /= f
def min(self, v):
return Vector3(min(self.x, v.x), min(self.y, v.y), min(self.z, v.z))
def max(self, v):
return Vector3(max(self.x, v.x), max(self.y, v.y), max(self.z, v.z))
class AABB(object):
def __init__(self, vMin, vMax):
self.vMin = vMin
self.vMax = vMax
self.perimeter = numpy.sum(self.vMax - self.vMin)
def combine(self, aabb):
return AABB(numpy.minimum(self.vMin, aabb.vMin), numpy.maximum(self.vMax, aabb.vMax))
def overlap(self, aabb):
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:
return False
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:
return False
return True
def __repr__(self):
return "AABB:%s - %s" % (str(self.vMin), str(self.vMax))
class _AABBNode(object):
def __init__(self, aabb):
self.child1 = None
self.child2 = None
self.parent = None
self.height = 0
self.aabb = aabb
def isLeaf(self):
return self.child1 == None
class AABBTree(object):
def __init__(self):
self.root = None
def insert(self, aabb):
newNode = _AABBNode(aabb)
if self.root == None:
self.root = newNode
return
node = self.root
while not node.isLeaf():
child1 = node.child1
child2 = node.child2
area = node.aabb.perimeter
combinedAABB = node.aabb.combine(aabb)
combinedArea = combinedAABB.perimeter
cost = 2.0 * combinedArea
inheritanceCost = 2.0 * (combinedArea - area)
if child1.isLeaf():
cost1 = aabb.combine(child1.aabb).perimeter + inheritanceCost
else:
oldArea = child1.aabb.perimeter
newArea = aabb.combine(child1.aabb).perimeter
cost1 = (newArea - oldArea) + inheritanceCost
if child2.isLeaf():
cost2 = aabb.combine(child1.aabb).perimeter + inheritanceCost
else:
oldArea = child2.aabb.perimeter
newArea = aabb.combine(child2.aabb).perimeter
cost2 = (newArea - oldArea) + inheritanceCost
if cost < cost1 and cost < cost2:
break
if cost1 < cost2:
node = child1
else:
node = child2
sibling = node
# Create a new parent.
oldParent = sibling.parent
newParent = _AABBNode(aabb.combine(sibling.aabb))
newParent.parent = oldParent
newParent.height = sibling.height + 1
if oldParent != None:
# The sibling was not the root.
if oldParent.child1 == sibling:
oldParent.child1 = newParent
else:
oldParent.child2 = newParent
newParent.child1 = sibling
newParent.child2 = newNode
sibling.parent = newParent
newNode.parent = newParent
else:
# The sibling was the root.
newParent.child1 = sibling
newParent.child2 = newNode
sibling.parent = newParent
newNode.parent = newParent
self.root = newParent
# Walk back up the tree fixing heights and AABBs
node = newNode.parent
while node != None:
node = self._balance(node)
child1 = node.child1
child2 = node.child2
node.height = 1 + max(child1.height, child2.height)
node.aabb = child1.aabb.combine(child2.aabb)
node = node.parent
def _balance(self, A):
if A.isLeaf() or A.height < 2:
return A
B = A.child1
C = A.child2
balance = C.height - B.height
# Rotate C up
if balance > 1:
F = C.child1;
G = C.child2;
# Swap A and C
C.child1 = A;
C.parent = A.parent;
A.parent = C;
# A's old parent should point to C
if C.parent != None:
if C.parent.child1 == A:
C.parent.child1 = C
else:
C.parent.child2 = C
else:
self.root = C
# Rotate
if F.height > G.height:
C.child2 = F
A.child2 = G
G.parent = A
A.aabb = B.aabb.combine(G.aabb)
C.aabb = A.aabb.combine(F.aabb)
A.height = 1 + max(B.height, G.height)
C.height = 1 + max(A.height, F.height)
else:
C.child2 = G
A.child2 = F
F.parent = A
A.aabb = B.aabb.combine(F.aabb)
C.aabb = A.aabb.combine(G.aabb)
A.height = 1 + max(B.height, F.height)
C.height = 1 + max(A.height, G.height)
return C;
# Rotate B up
if balance < -1:
D = B.child1
E = B.child2
# Swap A and B
B.child1 = A
B.parent = A.parent
A.parent = B
# A's old parent should point to B
if B.parent != None:
if B.parent.child1 == A:
B.parent.child1 = B
else:
B.parent.child2 = B
else:
self.root = B
# Rotate
if D.height > E.height:
B.child2 = D
A.child1 = E
E.parent = A
A.aabb = C.aabb.combine(E.aabb)
B.aabb = A.aabb.combine(D.aabb)
A.height = 1 + max(C.height, E.height)
B.height = 1 + max(A.height, D.height)
else:
B.child2 = E
A.child1 = D
D.parent = A
A.aabb = C.aabb.combine(D.aabb)
B.aabb = A.aabb.combine(E.aabb)
A.height = 1 + max(C.height, D.height)
B.height = 1 + max(A.height, E.height)
return B
return A
def query(self, aabb):
resultList = []
if self.root != None:
self._query(self.root, aabb, resultList)
return resultList
def _query(self, node, aabb, resultList):
if not aabb.overlap(node.aabb):
return
if node.isLeaf():
resultList.append(node.aabb)
else:
self._query(node.child1, aabb, resultList)
self._query(node.child2, aabb, resultList)
def __repr__(self):
s = "AABBTree:\n"
s += str(self.root.aabb)
return s
if __name__ == '__main__':
tree = AABBTree()
tree.insert(AABB(Vector3(0,0,0), Vector3(0,0,0)))
tree.insert(AABB(Vector3(1,1,1), Vector3(1,1,1)))
tree.insert(AABB(Vector3(0.5,0.5,0.5), Vector3(0.5,0.5,0.5)))
print(tree)
print(tree.query(AABB(Vector3(0,0,0), Vector3(0,0,0))))