Aluode/PerceptionLabPortable

1from math import atan2, degrees2 3import numpy as np4 5from . import SRTTransform6from .Qt import QtGui7from .Transform3D import Transform3D8from .Vector import Vector9 10 11class SRTTransform3D(Transform3D):12    """4x4 Transform matrix that can always be represented as a combination of 3 matrices: scale * rotate * translate13    This transform has no shear; angles are always preserved.14    """15    def __init__(self, init=None):16        Transform3D.__init__(self)17        self.reset()18        if init is None:19            return20        if init.__class__ is QtGui.QTransform:21            init = SRTTransform.SRTTransform(init)22        23        if isinstance(init, dict):24            self.restoreState(init)25        elif isinstance(init, SRTTransform3D):26            self._state = {27                'pos': Vector(init._state['pos']),28                'scale': Vector(init._state['scale']),29                'angle': init._state['angle'],30                'axis': Vector(init._state['axis']),31            }32            self.update()33        elif isinstance(init, SRTTransform.SRTTransform):34            self._state = {35                'pos': Vector(init._state['pos']),36                'scale': Vector(init._state['scale']),37                'angle': init._state['angle'],38                'axis': Vector(0, 0, 1),39            }40            self._state['scale'][2] = 1.041            self.update()42        elif isinstance(init, QtGui.QMatrix4x4):43            self.setFromMatrix(init)44        else:45            raise Exception("Cannot build SRTTransform3D from argument type:", type(init))46 47        48    def getScale(self):49        return Vector(self._state['scale'])50        51    def getRotation(self):52        """Return (angle, axis) of rotation"""53        return self._state['angle'], Vector(self._state['axis'])54        55    def getTranslation(self):56        return Vector(self._state['pos'])57    58    def reset(self):59        self._state = {60            'pos': Vector(0,0,0),61            'scale': Vector(1,1,1),62            'angle': 0.0,  ## in degrees63            'axis': (0, 0, 1)64        }65        self.update()66        67    def translate(self, *args):68        """Adjust the translation of this transform"""69        t = Vector(*args)70        self.setTranslate(self._state['pos']+t)71        72    def setTranslate(self, *args):73        """Set the translation of this transform"""74        self._state['pos'] = Vector(*args)75        self.update()76        77    def scale(self, *args):78        """adjust the scale of this transform"""79        ## try to prevent accidentally setting 0 scale on z axis80        if len(args) == 1 and hasattr(args[0], '__len__'):81            args = args[0]82        if len(args) == 2:83            args = args + (1,)84            85        s = Vector(*args)86        self.setScale(self._state['scale'] * s)87        88    def setScale(self, *args):89        """Set the scale of this transform"""90        if len(args) == 1 and hasattr(args[0], '__len__'):91            args = args[0]92        if len(args) == 2:93            args = args + (1,)94        self._state['scale'] = Vector(*args)95        self.update()96        97    def rotate(self, angle, axis=(0,0,1)):98        """Adjust the rotation of this transform"""99        origAxis = self._state['axis']100        if axis[0] == origAxis[0] and axis[1] == origAxis[1] and axis[2] == origAxis[2]:101            self.setRotate(self._state['angle'] + angle)102        else:103            m = QtGui.QMatrix4x4()104            m.translate(*self._state['pos'])105            m.rotate(self._state['angle'], *self._state['axis'])106            m.rotate(angle, *axis)107            m.scale(*self._state['scale'])108            self.setFromMatrix(m)109        110    def setRotate(self, angle, axis=(0,0,1)):111        """Set the transformation rotation to angle (in degrees)"""112        113        self._state['angle'] = angle114        self._state['axis'] = Vector(axis)115        self.update()116    117    def setFromMatrix(self, m):118        """119        Set this transform based on the elements of *m*120        The input matrix must be affine AND have no shear,121        otherwise the conversion will most likely fail.122        """123        import numpy.linalg124        for i in range(4):125            self.setRow(i, m.row(i))126        m = self.matrix().reshape(4,4)127        ## translation is 4th column128        self._state['pos'] = m[:3,3]129        130        ## scale is vector-length of first three columns131        scale = (m[:3,:3]**2).sum(axis=0)**0.5132        ## see whether there is an inversion133        z = np.cross(m[0, :3], m[1, :3])134        if np.dot(z, m[2, :3]) < 0:135            scale[1] *= -1  ## doesn't really matter which axis we invert136        self._state['scale'] = scale137        138        ## rotation axis is the eigenvector with eigenvalue=1139        r = m[:3, :3] / scale[np.newaxis, :]140        try:141            evals, evecs = numpy.linalg.eig(r)142        except:143            print("Rotation matrix: %s" % str(r))144            print("Scale: %s" % str(scale))145            print("Original matrix: %s" % str(m))146            raise147        eigIndex = np.argwhere(np.abs(evals-1) < 1e-6)148        if len(eigIndex) < 1:149            print("eigenvalues: %s" % str(evals))150            print("eigenvectors: %s" % str(evecs))151            print("index: %s, %s" % (str(eigIndex), str(evals-1)))152            raise Exception("Could not determine rotation axis.")153        axis = evecs[:,eigIndex[0,0]].real154        axis /= ((axis**2).sum())**0.5155        self._state['axis'] = axis156        157        ## trace(r) == 2 cos(angle) + 1, so:158        cos = (r.trace()-1)*0.5  ## this only gets us abs(angle)159        160        ## The off-diagonal values can be used to correct the angle ambiguity, 161        ## but we need to figure out which element to use:162        axisInd = np.argmax(np.abs(axis))163        rInd,sign = [((1,2), -1), ((0,2), 1), ((0,1), -1)][axisInd]164        165        ## Then we have r-r.T = sin(angle) * 2 * sign * axis[axisInd];166        ## solve for sin(angle)167        sin = (r-r.T)[rInd] / (2. * sign * axis[axisInd])168        169        ## finally, we get the complete angle from arctan(sin/cos)170        self._state['angle'] = degrees(atan2(sin, cos))171        if self._state['angle'] == 0:172            self._state['axis'] = (0,0,1)173        174    def as2D(self):175        """Return a QTransform representing the x,y portion of this transform (if possible)"""176        return SRTTransform.SRTTransform(self)177 178    #def __div__(self, t):179        #"""A / B  ==  B^-1 * A"""180        #dt = t.inverted()[0] * self181        #return SRTTransform.SRTTransform(dt)182        183    #def __mul__(self, t):184        #return SRTTransform.SRTTransform(QtGui.QTransform.__mul__(self, t))185 186    def saveState(self):187        p = self._state['pos']188        s = self._state['scale']189        ax = self._state['axis']190        #if s[0] == 0:191            #raise Exception('Invalid scale: %s' % str(s))192        return {193            'pos': (p[0], p[1], p[2]), 194            'scale': (s[0], s[1], s[2]), 195            'angle': self._state['angle'], 196            'axis': (ax[0], ax[1], ax[2])197        }198 199    def restoreState(self, state):200        self._state['pos'] = Vector(state.get('pos', (0.,0.,0.)))201        scale = state.get('scale', (1.,1.,1.))202        scale = tuple(scale) + (1.,) * (3-len(scale))203        self._state['scale'] = Vector(scale)204        self._state['angle'] = state.get('angle', 0.)205        self._state['axis'] = state.get('axis', (0, 0, 1))206        self.update()207 208    def update(self):209        Transform3D.setToIdentity(self)210        ## modifications to the transform are multiplied on the right, so we need to reverse order here.211        Transform3D.translate(self, *self._state['pos'])212        Transform3D.rotate(self, self._state['angle'], *self._state['axis'])213        Transform3D.scale(self, *self._state['scale'])214 215    def __repr__(self):216        return str(self.saveState())217 218    def __reduce__(self):219        return SRTTransform3D, (self.saveState(),)220