dtrack_module.py

# Copyright (c) 2025, Advanced Realtime Tracking GmbH & Co. KG
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
# 3. Neither the name of copyright holder nor the names of its contributors
#    may be used to endorse or promote products derived from this software
#    without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.


"""Module for handling DTrack data in VRED python plugins"""



from dtrack_math import *
from dtrack_logger import *
from dtrack_base import *

importErrorDTrackSDK = False
try:
    import DTrackSDK
    dtrackLog.out( "Using DTrackSDK (pybind wrapper of the C++ DTrackSDK)")
except ImportError:
    try:
        import dtrack_sdk as DTrackSDK
        dtrackLog.out( "Using dtrack_sdk (pure python version of DTrackSDK)")
    except ImportError:
        dtrackLog.out( "No DTrack SDK module found")
        importErrorDTrackSDK = True


#from fnmatch import translate

from PySide6 import QtCore
import time


#------------------------------------------------------------------------------
# DTrackReceiveLoop
#------------------------------------------------------------------------------

class DTrackReceiveLoop( QtCore.QObject):
    """Thread receiving data from DTrack"""

    state = QtCore.Signal( VredDTrackState)


    def __init__( self, parent=None):
        super().__init__( parent)

        self.port = -1
        self.dtrackData = None
        self.localState = VredDTrackState.NotInitialized
        self.stateCount = { VredDTrackState.NoData: 0 }
        self.haveDTrackSDK  = not importErrorDTrackSDK
   

    def stateUpdate( self, state ):

        # As long as we do not have 'NoData' permanently, we assume 'Receiving'
        if state is VredDTrackState.NoData: # and self.localState is VredDTrackState.Receiving:
            self.stateCount[ state ] += 1
            if self.stateCount[ state ] % 70 == 0:      # TODO
                self.state.emit( state )
                self.localState = state
            else:
                pass

        else:
            self.stateCount[ VredDTrackState.NoData ] = 0
            if self.localState != state:
                self.state.emit( state )
                self.localState = state
        


    def runDtrack( self):

        if self.haveDTrackSDK is False:
            dtrackLog.out("No receive thread started without DTrackSDK")
            return

        thread = QtCore.QThread.currentThread()
        if dtrackLog.enabled: dtrackLog.out( "Thread {} enter".format( thread ) )
        threadExitMsg = "Thread {} exit".format( thread )

        if self.port < 0  or self.dtrackData is None:
            dtrackLog.out( "self.port < 0  or self.dtrackData is None")
            self.stateUpdate( VredDTrackState.PortInvalid )

            dtrackLog.out( threadExitMsg )
            return
                
        dtrack = DTrackSDK.dtrack( self.port )
        if dtrack is None:
            dtrackLog.out( "dtrack is NONE")

        if not dtrack.isDataInterfaceValid():
            dtrackLog.out( "NOT dtrack.isDataInterfaceValid()")
            self.stateUpdate( VredDTrackState.NoConnection )

            dtrackLog.out( threadExitMsg )
            return

        dtrack.setDataTimeoutUS( 1000 ) # usec
        logReset = False

        while not thread.isInterruptionRequested():

            if not dtrack.receive():
                self.stateUpdate( VredDTrackState.NoData )
                dtrackLog.out( "NOT dtrack.receive()", every=1000, reset=logReset )
                logReset = False
                time.sleep( 0.002) # sec
                continue

            logReset = True
            self.stateUpdate( VredDTrackState.Receiving )

            if self.dtrackData.mutex.try_lock():
                
                self.dtrackData.frame = ( 
                    self.dtrackData.frame[0] + 1,
                    dtrack.getFrameCounter() )
                                    
                nb = dtrack.getNumBody()
                self.dtrackData.body = [None] * nb
                for i in range( 0, nb ):
                    self.dtrackData.body[i] = dtrack.getBody(i)

                nf = dtrack.getNumFlyStick()
                self.dtrackData.flystick = [None] * nf
                for i in range( 0, nf ):
                    self.dtrackData.flystick[i] = dtrack.getFlyStick(i)

                nh = dtrack.getNumHand()
                self.dtrackData.hand = [None] * nh
                for i in range( 0, nh ):
                    self.dtrackData.hand[i] = dtrack.getHand(i)
               
                self.dtrackData.mutex.unlock()

        self.stateUpdate( VredDTrackState.Stopped )
        dtrackLog.out( threadExitMsg )










#------------------------------------------------------------------------------
# VredDTrack
#------------------------------------------------------------------------------


class VredDTrack( QtCore.QObject):
    """DTrack data and VRED render loop management"""

    state = QtCore.Signal( VredDTrackState)
    
    def __init__(self, parent=None):
        super().__init__( parent)

        dtrackLog.enable( DTRACK_LOG_ENABLED)     # TMP

        self.dtrackData = DTrackData()
        self.dtrackData.mutex = QtCore.QMutex()
        
        self.renderHooks = []
        
        self.receiveLoop = DTrackReceiveLoop()
        self.receiveLoop.dtrackData = self.dtrackData

        self.receiveLoop.state.connect( self.receiveLoopState)
        
        self.receiveThread = QtCore.QThread()
        self.receiveLoop.moveToThread( self.receiveThread)
        self.receiveThread.started.connect( self.receiveLoop.runDtrack)

        self.rotation = Rotation()

    
    def addToLoop( self, hook):
        """Add a hook to the render loop"""

        if self.receiveLoop.haveDTrackSDK:
            self.renderHooks.append( hook)
            self.renderHooks[-1].module = self


    def clearHooks( self):
        """Remove all hooks from the render loop"""

        for hook in self.renderHooks:
            hook.subLoop()
        self.renderHooks.clear()

 

    def start( self):
        """Start receiving DTrack data, and application in the VRED render loop"""

        if not self.receiveLoop.haveDTrackSDK:
            return

        if dtrackLog.enabled: dtrackLog.out( "Main thread {}".format( QtCore.QThread.currentThread() ))

        self.receiveThread.start()

        for hook in self.renderHooks:
            hook.addLoop()


    def stop( self):
        """Stop receiving and rendering"""

        if not self.receiveLoop.haveDTrackSDK:
            return

        self.receiveThread.requestInterruption()
        self.receiveThread.quit()
        self.receiveThread.wait( 4000)

        for hook in self.renderHooks:
            hook.subLoop()



    @QtCore.Slot( VredDTrackState)
    def receiveLoopState( self, VredDTrackState):

        if not self.receiveLoop.haveDTrackSDK:
            return

        # In case 'VredDTrackState.Stopped' self.stop() was already called
        if VredDTrackState in { VredDTrackState.PortInvalid, VredDTrackState.NoConnection }:
            self.receiveThread.quit()
            self.receiveThread.wait( 4000)
            for hook in self.renderHooks:
                hook.subLoop()

        # Forward signal
        self.state.emit( VredDTrackState )
        
        
    def setPort( self, port):
        """Set the local UDP port for receiving DTrack data"""
        self.receiveLoop.port = port



    def getTranslationNP( self, targetType, targetID, scale = 1.0):
        """Get the translation of the specified A.R.T. target as a numpy array"""

        if self.receiveLoop.haveDTrackSDK:
            if targetType is TargetType.Body and targetID < len(self.dtrackData.body):
                b = self.dtrackData.body[ targetID ]
                return b.loc * scale
            elif targetType is TargetType.Flystick and targetID < len(self.dtrackData.flystick):
                b = self.dtrackData.flystick[ targetID ]
                return b.loc * scale

        return None


    def getTranslation( self, targetType, targetID, scale = 1.0):
        """Get the translation of the specified A.R.T. target as a QVector3D"""

        loc = self.getTranslationNP( targetType, targetID, scale )
        if loc is not None:
            return QVector3D( loc[0,0], loc[1,0], loc[2,0])
        return None


    def getRotationAsQuaternionNP( self, targetType, targetID):
        """Get the rotation of the specified A.R.T. target as quaternion [w,x,y,z] in a numpy array"""

        if self.receiveLoop.haveDTrackSDK:
            if targetType is TargetType.Body and targetID < len(self.dtrackData.body):
                b = self.dtrackData.body[ targetID ]
                return b.getQuaternion()
            elif targetType is TargetType.Flystick and targetID < len(self.dtrackData.flystick):
                b = self.dtrackData.flystick[ targetID ]
                return b.getQuaternion()

        return None


    def getRotationAsQuaternion( self, targetType, targetID):
        """Get the rotation of the specified A.R.T. target as a QQuaternion"""

        rot = self.getRotationAsQuaternionNP( targetType, targetID)
        if rot is not None:
            return QQuaternion( rot[0], rot[1], rot[2], rot[3])
        return None


    def getRotationAsMatrixNP( self, targetType, targetID):
        """Get the rotation matrix of the specified A.R.T. target as a numpy array"""

        if self.receiveLoop.haveDTrackSDK:
            if targetType is TargetType.Body and targetID < len(self.dtrackData.body):
                b = self.dtrackData.body[ targetID ]
                return b.rot
            elif targetType is TargetType.Flystick and targetID < len(self.dtrackData.flystick):
                b = self.dtrackData.flystick[ targetID ]
                return b.rot

        return None



    def isTracked( self, targetType, targetID):
        """Check if the specified A.R.T. target is tracked"""

        if self.receiveLoop.haveDTrackSDK:
            if targetType is TargetType.Body and targetID < len(self.dtrackData.body):
                return self.dtrackData.body[targetID].isTracked()
            elif targetType is TargetType.Flystick and targetID < len(self.dtrackData.flystick):
                return self.dtrackData.flystick[targetID].isTracked()

            elif targetType is TargetType.Hand and targetID < len(self.dtrackData.hand):
                return self.dtrackData.hand[targetID].isTracked()

        return False


    def getFlyButtons( self, targetID ):
        """Get the button information of the specified A.R.T. Flystick"""

        if self.receiveLoop.haveDTrackSDK:
            if targetID < len(self.dtrackData.flystick):
                return self.dtrackData.flystick[targetID].button
        return None


    def getFlyJoystick( self, targetID ):
        """Get the joystick information of the specified A.R.T. Flystick"""

        if self.receiveLoop.haveDTrackSDK:
            if targetID < len(self.dtrackData.flystick):
                return self.dtrackData.flystick[targetID].joystick
        return None



    def getHandedness( self, targetID ):
        """Check if left or right hand"""

        if self.receiveLoop.haveDTrackSDK:
            if targetID < len(self.dtrackData.hand):
                if self.dtrackData.hand[ targetID].lr == 0:
                    return Handedness.Left
                return Handedness.Right
        return Handedness.Undefined



    def updateHandJointPoses( self, vrHandJointPoses ):
        """Calculate all hand joint poses in world space"""

        if not self.receiveLoop.haveDTrackSDK:
            return False

        vrHandJointPoses.isTracked = self.isTracked( TargetType.Hand, vrHandJointPoses.targetID )
        if not vrHandJointPoses.isTracked:
            return False
        
        dtHand  = self.dtrackData.hand[ vrHandJointPoses.targetID ]
        artHand = ArtHandJointPoses( vrHandJointPoses.handedness, vrHandJointPoses.targetID )
        vrHand  = vrHandJointPoses


        artHand.palm["base"].loc = dtHand.loc
        artHand.palm["base"].rot = dtHand.rot

        # Unpack ART joint poses
        for F in range( 0, HandJointPoses.nFingers ):
            artFinger = artHand.finger[F]

            # Finger tip (in Hand CS)
            artFinger.joint[ JointT.Tip].loc  =  dtHand.finger[F].loc
            artFinger.joint[ JointT.Tip].rot  =  dtHand.finger[F].rot

            # Outer phalanx (in Hand CS)
            artFinger.joint[ JointT.Outer].rot  =  dtHand.finger[F].rot
            phalanxOuter = np.array( [ [ -dtHand.finger[F].lengthphalanx[0] ], [0], [0] ])
            artFinger.joint[ JointT.Outer].loc  =  artFinger.joint[ JointT.Tip].rot  @ phalanxOuter  + artFinger.joint[ JointT.Tip].loc

            # Middle phalanx (in Hand CS)
            jointOuter  =  Rotation.Y( dtHand.finger[F].anglephalanx[0] )
            artFinger.joint[ JointT.Middle].rot  =  artFinger.joint[ JointT.Outer].rot  @ jointOuter
            phalanxMiddle = np.array( [ [ -dtHand.finger[F].lengthphalanx[1] ], [0], [0] ])
            artFinger.joint[ JointT.Middle].loc  =  artFinger.joint[ JointT.Middle].rot   @ phalanxMiddle  + artFinger.joint[ JointT.Outer].loc 

            # Inner phalanx (in Hand CS)
            jointInner  =  Rotation.Y( dtHand.finger[F].anglephalanx[1] )
            artFinger.joint[ JointT.Root].rot  =  artFinger.joint[ JointT.Middle].rot  @ jointInner
            phalanxInner  =  np.array([ [ -dtHand.finger[F].lengthphalanx[2] ], [0], [0] ])
            artFinger.joint[ JointT.Root].loc  =  artFinger.joint[ JointT.Root].rot  @ phalanxInner  + artFinger.joint[ JointT.Middle].loc

            # Trafo to World CS
            for joint in artFinger.joint:
                joint.rot  =  artHand.palm["base"].rot  @ joint.rot
                joint.loc  =  artHand.palm["base"].rot  @ joint.loc  + artHand.palm["base"].loc 


        # Update VRED hand model: VRED joint poses expressed through  
        # ART joint poses. May still need some finetuning (TODO)

        if vrHand.handedness == Handedness.Left:
            flipThumbA  = np.eye(3,3)
            flipThumbB  = self.rotation.X90_
            flipOtherA  = self.rotation.X90
            flipOtherB  = self.rotation.Z90_
            flipMiddle  = self.rotation.X180
            flipRing    = self.rotation.X180

        else:
            flipEye  =  -1 * np.eye(3,3)

            flipThumbA  = self.rotation.Z180 @ self.rotation.X180   @ flipEye
            flipThumbB  = self.rotation.Z180 @ self.rotation.X90    @ flipEye
            flipOtherA  = self.rotation.Z180 @ self.rotation.X90_   @ flipEye
            flipOtherB  = self.rotation.Z90  @ self.rotation.Y180   @ flipEye
            flipMiddle  = self.rotation.X180
            flipRing    = self.rotation.X180

        # Thumb
        for J in range( 0, FingerJointPoses.nJoints-1):
            vrHand.finger[ FingerT.Thumb].joint[J].loc = artHand.finger[ FingerT.Thumb].joint[J].loc
            vrHand.finger[ FingerT.Thumb].joint[J].rot = artHand.finger[ FingerT.Thumb].joint[J].rot                        @ flipThumbA
        vrHand.finger[ FingerT.Thumb].joint[ JointT.Tip].loc = artHand.finger[ FingerT.Thumb].joint[ JointT.Tip].loc
        vrHand.finger[ FingerT.Thumb].joint[ JointT.Tip].rot = artHand.finger[ FingerT.Thumb].joint[ JointT.Tip].rot        @ flipThumbB

        # Other fingers
        for F in range( 1, HandJointPoses.nFingers):
            for J in range( 0, FingerJointPoses.nJoints-1):
                vrHand.finger[F].joint[J].loc = artHand.finger[F].joint[J].loc
                vrHand.finger[F].joint[J].rot = artHand.finger[F].joint[J].rot                                              @ flipOtherA
            vrHand.finger[F].joint[ JointT.Tip].loc = artHand.finger[F].joint[ JointT.Tip].loc
            vrHand.finger[F].joint[ JointT.Tip].rot = artHand.finger[F].joint[ JointT.Tip].rot                              @ flipOtherB

        # Ring and middle finger root joint
        vrHand.finger[ FingerT.Middle].joint[ JointT.Root].rot  =  vrHand.finger[ FingerT.Middle].joint[ JointT.Root].rot   @ flipMiddle
        vrHand.finger[ FingerT.Ring  ].joint[ JointT.Root].rot  =  vrHand.finger[ FingerT.Ring  ].joint[ JointT.Root].rot   @ flipRing 

        # Wrist
        vrHand.palm["wrist"].rot      =  artHand.palm["base"].rot  @ vrHand.relative["wrist_artBase"].rot
        vrHand.palm["wrist"].loc      =  artHand.palm["base"].rot  @ vrHand.relative["wrist_artBase"].loc   + artHand.palm["base"].loc

        # Finger sockets (vrHand.socket[])
        vrHand.palm["palmEnd"].rot    =  vrHand.palm["wrist"].rot  @ vrHand.relative["palmEnd_wrist"].rot
        vrHand.palm["palmEnd"].loc    =  vrHand.palm["wrist"].rot  @ vrHand.relative["palmEnd_wrist"].loc   + vrHand.palm["wrist"].loc
        vrHand.palm["thumbBall"].rot  =  vrHand.palm["wrist"].rot  @ vrHand.relative["thumbBall_wrist"].rot
        vrHand.palm["thumbBall"].loc  =  vrHand.palm["wrist"].rot  @ vrHand.relative["thumbBall_wrist"].loc + vrHand.palm["wrist"].loc
        vrHand.palm["pinkyBase"].rot  =  vrHand.palm["wrist"].rot  @ vrHand.relative["pinkyBase_wrist"].rot
        vrHand.palm["pinkyBase"].loc  =  vrHand.palm["wrist"].rot  @ vrHand.relative["pinkyBase_wrist"].loc + vrHand.palm["wrist"].loc

        # Scale pivot
        vrHand.scalePivot  = artHand.palm["base"].loc

        return vrHandJointPoses.isTracked