<commit_before>import os
import PythonQt
from PythonQt import QtCore, QtGui
from director import callbacks
from director import drcargs
import director.applogic as app
import director.objectmodel as om
import director.visualization as vis
import director.vtkAll as vtk
from director import jointcontrol
from director import getDRCBaseDir
from director import lcmUtils
from director import filterUtils
from director import packagepath
from director import transformUtils


import bot_core
import math
import numpy as np
import json




def getRobotGrayColor():
    return QtGui.QColor(177, 180, 190)

def getRobotOrangeColor():
    return QtGui.QColor(255, 190, 0)

def getRobotBlueColor():
    return QtGui.QColor(170, 255, 255)


class RobotModelItem(om.ObjectModelItem):

    MODEL_CHANGED_SIGNAL = 'MODEL_CHANGED_SIGNAL'

    def __init__(self, model):

        modelName = os.path.basename(model.filename())
        om.ObjectModelItem.__init__(self, modelName, om.Icons.Robot)

        self.views = []
        self.model = None
        self.callbacks.addSignal(self.MODEL_CHANGED_SIGNAL)

        self.addProperty('Filename', model.filename())
        self.addProperty('Visible', model.visible())
        self.addProperty('Alpha', model.alpha(),
                         attributes=om.PropertyAttributes(decimals=2, minimum=0, maximum=1.0, singleStep=0.1, hidden=False))
        self.addProperty('Color Mode', 0,
                         attributes=om.PropertyAttributes(enumNames=['Solid Color', 'Textures', 'URDF Colors']))
        self.addProperty('Color', model.color())



        self.setModel(model)

    def _onPropertyChanged(self, propertySet, propertyName):
        om.ObjectModelItem._onPropertyChanged(self, propertySet, propertyName)

        if propertyName == 'Alpha':
            self.model.setAlpha(self.getProperty(propertyName))
        elif propertyName == 'Visible':
            self.model.setVisible(self.getProperty(propertyName))
        elif propertyName == 'Color Mode':
            if self.getProperty(propertyName) == 1: # Textures
                self.model.setTexturesEnabled(True)
            else:
                self.model.setTexturesEnabled(False)
            self._updateModelColor()
        elif propertyName == 'Color':
            self._updateModelColor()

        self._renderAllViews()

    def hasDataSet(self, dataSet):
        return len(self.model.getLinkNameForMesh(dataSet)) != 0

    def connectModelChanged(self, func):
        return self.callbacks.connect(self.MODEL_CHANGED_SIGNAL, func)

    def disconnectModelChanged(self, callbackId):
        self.callbacks.disconnect(callbackId)

    def onModelChanged(self):
        self.callbacks.process(self.MODEL_CHANGED_SIGNAL, self)
        if self.getProperty('Visible'):
            self._renderAllViews()

    def onDisplayChanged(self):
        if self.getProperty('Visible'):
            self._renderAllViews()

    def _renderAllViews(self):
        for view in self.views:
            view.render()

    def getLinkFrame(self, linkName):
        t = vtk.vtkTransform()
        t.PostMultiply()
        if self.model.getLinkToWorld(linkName, t):
            return t
        else:
            return None

    def getHeadLink(self):
        return drcargs.getDirectorConfig().get('headLink')

    def getLinkContactPoints(self, linkName):
        pts = self.model.getBodyContactPoints(linkName)
        numberOfPoints = len(pts)/3
        return np.array(pts).reshape(numberOfPoints, 3)

    def setModel(self, model):
        assert model is not None
        if model == self.model:
            return

        model.disconnect('modelChanged()', self.onModelChanged)
        model.disconnect('displayChanged()', self.onDisplayChanged)
        views = list(self.views)
        self.removeFromAllViews()

        self.model = model
        self.model.setAlpha(self.getProperty('Alpha'))
        self.model.setVisible(self.getProperty('Visible'))
        self.model.setTexturesEnabled((self.getProperty('Color Mode') == 1))
        self._updateModelColor()

        self.setProperty('Filename', model.filename())
        model.connect('modelChanged()', self.onModelChanged)
        model.connect('displayChanged()', self.onDisplayChanged)

        for view in views:
            self.addToView(view)
        self.onModelChanged()

    def _updateModelColor(self):
        if self.getProperty('Color Mode') == 0: # Solid Color
            color = QtGui.QColor(*[c*255 for c in self.getProperty('Color')])
            self.model.setColor(color)
        elif self.getProperty('Color Mode') == 1: # Textures
            self._setupTextureColors()
        elif self.getProperty('Color Mode') == 2: # URDF Colors
            self.model.setUrdfColors()

    def _setupTextureColors(self):

        # custom colors for non-textured robotiq hand
        for name in self.model.getLinkNames():
            strs = name.split('_')
            if len(strs) >= 2 and strs[0] in ['left', 'right'] and strs[1] in ('finger', 'palm') or name.endswith('hand_force_torque'):
                self.model.setLinkColor(name, QtGui.QColor(90, 90, 90) if strs[1] == 'finger' else QtGui.QColor(20,20,20))
            else:
                self.model.setLinkColor(name, QtGui.QColor(255,255,255))

    def addToView(self, view):
        if view in self.views:
            return
        self.views.append(view)
        self.model.addToRenderer(view.renderer())
        view.render()

    def onRemoveFromObjectModel(self):
        om.ObjectModelItem.onRemoveFromObjectModel(self)
        self.removeFromAllViews()

    def removeFromAllViews(self):
        for view in list(self.views):
            self.removeFromView(view)
        assert len(self.views) == 0

    def removeFromView(self, view):
        assert view in self.views
        self.views.remove(view)
        self.model.removeFromRenderer(view.renderer())
        view.render()


def loadRobotModel(name, view=None, parent='planning', urdfFile=None, color=None, visible=True, colorMode='URDF Colors'):

    if not urdfFile:
        urdfFile = drcargs.getDirectorConfig()['urdfConfig']['default']

    if isinstance(parent, str):
        parent = om.getOrCreateContainer(parent)

    model = loadRobotModelFromFile(urdfFile)
    if not model:
        raise Exception('Error loading robot model from file: %s' % urdfFile)

    obj = RobotModelItem(model)
    om.addToObjectModel(obj, parent)

    obj.setProperty('Visible', visible)
    obj.setProperty('Name', name)
    obj.setProperty('Color', color or getRobotGrayColor())
    if colorMode == 'Textures':
        obj.setProperty('Color Mode', 1)
    elif colorMode == 'Solid Color':
        obj.setProperty('Color Mode', 0)
    elif colorMode == 'URDF Colors':
        obj.setProperty('Color Mode', 2)

    if view is not None:
        obj.addToView(view)

    jointController = jointcontrol.JointController([obj])

    fixedPointFile = drcargs.getDirectorConfig()['fixedPointFile']
    jointController.setPose('q_nom', jointController.loadPoseFromFile(fixedPointFile))<commit_msg>modify optional args for roboturdf.loadRobotModel()<commit_after>import os
import PythonQt
from PythonQt import QtCore, QtGui
from director import callbacks
from director import drcargs
import director.applogic as app
import director.objectmodel as om
import director.visualization as vis
import director.vtkAll as vtk
from director import jointcontrol
from director import getDRCBaseDir
from director import lcmUtils
from director import filterUtils
from director import packagepath
from director import transformUtils


import bot_core
import math
import numpy as np
import json




def getRobotGrayColor():
    return QtGui.QColor(177, 180, 190)

def getRobotOrangeColor():
    return QtGui.QColor(255, 190, 0)

def getRobotBlueColor():
    return QtGui.QColor(170, 255, 255)


class RobotModelItem(om.ObjectModelItem):

    MODEL_CHANGED_SIGNAL = 'MODEL_CHANGED_SIGNAL'

    def __init__(self, model):

        modelName = os.path.basename(model.filename())
        om.ObjectModelItem.__init__(self, modelName, om.Icons.Robot)

        self.views = []
        self.model = None
        self.callbacks.addSignal(self.MODEL_CHANGED_SIGNAL)

        self.addProperty('Filename', model.filename())
        self.addProperty('Visible', model.visible())
        self.addProperty('Alpha', model.alpha(),
                         attributes=om.PropertyAttributes(decimals=2, minimum=0, maximum=1.0, singleStep=0.1, hidden=False))
        self.addProperty('Color Mode', 0,
                         attributes=om.PropertyAttributes(enumNames=['Solid Color', 'Textures', 'URDF Colors']))
        self.addProperty('Color', model.color())



        self.setModel(model)

    def _onPropertyChanged(self, propertySet, propertyName):
        om.ObjectModelItem._onPropertyChanged(self, propertySet, propertyName)

        if propertyName == 'Alpha':
            self.model.setAlpha(self.getProperty(propertyName))
        elif propertyName == 'Visible':
            self.model.setVisible(self.getProperty(propertyName))
        elif propertyName == 'Color Mode':
            if self.getProperty(propertyName) == 1: # Textures
                self.model.setTexturesEnabled(True)
            else:
                self.model.setTexturesEnabled(False)
            self._updateModelColor()
        elif propertyName == 'Color':
            self._updateModelColor()

        self._renderAllViews()

    def hasDataSet(self, dataSet):
        return len(self.model.getLinkNameForMesh(dataSet)) != 0

    def connectModelChanged(self, func):
        return self.callbacks.connect(self.MODEL_CHANGED_SIGNAL, func)

    def disconnectModelChanged(self, callbackId):
        self.callbacks.disconnect(callbackId)

    def onModelChanged(self):
        self.callbacks.process(self.MODEL_CHANGED_SIGNAL, self)
        if self.getProperty('Visible'):
            self._renderAllViews()

    def onDisplayChanged(self):
        if self.getProperty('Visible'):
            self._renderAllViews()

    def _renderAllViews(self):
        for view in self.views:
            view.render()

    def getLinkFrame(self, linkName):
        t = vtk.vtkTransform()
        t.PostMultiply()
        if self.model.getLinkToWorld(linkName, t):
            return t
        else:
            return None

    def getHeadLink(self):
        return drcargs.getDirectorConfig().get('headLink')

    def getLinkContactPoints(self, linkName):
        pts = self.model.getBodyContactPoints(linkName)
        numberOfPoints = len(pts)/3
        return np.array(pts).reshape(numberOfPoints, 3)

    def setModel(self, model):
        assert model is not None
        if model == self.model:
            return

        model.disconnect('modelChanged()', self.onModelChanged)
        model.disconnect('displayChanged()', self.onDisplayChanged)
        views = list(self.views)
        self.removeFromAllViews()

        self.model = model
        self.model.setAlpha(self.getProperty('Alpha'))
        self.model.setVisible(self.getProperty('Visible'))
        self.model.setTexturesEnabled((self.getProperty('Color Mode') == 1))
        self._updateModelColor()

        self.setProperty('Filename', model.filename())
        model.connect('modelChanged()', self.onModelChanged)
        model.connect('displayChanged()', self.onDisplayChanged)

        for view in views:
            self.addToView(view)
        self.onModelChanged()

    def _updateModelColor(self):
        if self.getProperty('Color Mode') == 0: # Solid Color
            color = QtGui.QColor(*[c*255 for c in self.getProperty('Color')])
            self.model.setColor(color)
        elif self.getProperty('Color Mode') == 1: # Textures
            self._setupTextureColors()
        elif self.getProperty('Color Mode') == 2: # URDF Colors
            self.model.setUrdfColors()

    def _setupTextureColors(self):

        # custom colors for non-textured robotiq hand
        for name in self.model.getLinkNames():
            strs = name.split('_')
            if len(strs) >= 2 and strs[0] in ['left', 'right'] and strs[1] in ('finger', 'palm') or name.endswith('hand_force_torque'):
                self.model.setLinkColor(name, QtGui.QColor(90, 90, 90) if strs[1] == 'finger' else QtGui.QColor(20,20,20))
            else:
                self.model.setLinkColor(name, QtGui.QColor(255,255,255))

    def addToView(self, view):
        if view in self.views:
            return
        self.views.append(view)
        self.model.addToRenderer(view.renderer())
        view.render()

    def onRemoveFromObjectModel(self):
        om.ObjectModelItem.onRemoveFromObjectModel(self)
        self.removeFromAllViews()

    def removeFromAllViews(self):
        for view in list(self.views):
            self.removeFromView(view)
        assert len(self.views) == 0

    def removeFromView(self, view):
        assert view in self.views
        self.views.remove(view)
        self.model.removeFromRenderer(view.renderer())
        view.render()


def loadRobotModel(name=None, view=None, parent='scene', urdfFile=None, color=None, visible=True, colorMode='URDF Colors', useConfigFile=True):

    if not urdfFile:
        urdfFile = drcargs.getDirectorConfig()['urdfConfig']['default']

    if isinstance(parent, str):
        parent = om.getOrCreateContainer(parent)

    model = loadRobotModelFromFile(urdfFile)
    if not model:
        raise Exception('Error loading robot model from file: %s' % urdfFile)

    obj = RobotModelItem(model)
    om.addToObjectModel(obj, parent)

    if name:
        obj.setProperty('Name', name)

    obj.setProperty('Visible', visible)
    obj.setProperty('Color', color or getRobotGrayColor())
    if colorMode == 'Textures':
        obj.setProperty('Color Mode', 1)
    elif colorMode == 'Solid Color':
        obj.setProperty('Color Mode', 0)
    elif colorMode == 'URDF Colors':
        obj.setProperty('Color Mode', 2)

    if view is not None:
        obj.addToView(view)

    if useConfigFile:
        jointNames = None
    else:
        jointNames = model.getJointNames()

    jointController = jointcontrol.JointController([obj], jointNames=jointNames)

    if useConfigFile:
        fixedPointFile = drcargs.getDirectorConfig()['fixedPointFile']
        jointController.setPose('q_nom', jointController.loadPoseFromFile(fixedPointFile))
{'fraction_non_alphanumeric': 0.10811374539386776, 'fraction_numerical': 0.010011819509142738, 'mean_word_length': 2.9322033898305087, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2625755', 'n_tokens_mistral': 4296, 'n_tokens_neox': 4074, 'n_words': 892}