<commit_before>import os
import copy
import math
import functools
import numpy as np

from ddapp import transformUtils

from ddapp.asynctaskqueue import AsyncTaskQueue
from ddapp import objectmodel as om
from ddapp import visualization as vis
from ddapp import robotstate
from ddapp import segmentation
from ddapp import planplayback
from ddapp.pointpicker import PointPicker
from ddapp import vtkAll as vtk
from ddapp.simpletimer import SimpleTimer
from ddapp import affordanceupdater

from ddapp.debugVis import DebugData
from ddapp import affordanceitems
from ddapp import ikplanner
from ddapp import vtkNumpy
from numpy import array
from ddapp.uuidutil import newUUID
import affordancepanel
import ioUtils

from ddapp.tasks.taskuserpanel import TaskUserPanel
from ddapp.tasks.taskuserpanel import ImageBasedAffordanceFit

import ddapp.tasks.robottasks as rt


class TableDemo(object):

    def __init__(self, robotStateModel, playbackRobotModel, ikPlanner, manipPlanner, footstepPlanner,
                 atlasDriver, lhandDriver, rhandDriver, multisenseDriver, view, sensorJointController,
                 planPlaybackFunction, teleopPanel):
        self.planPlaybackFunction = planPlaybackFunction
        self.robotStateModel = robotStateModel
        self.playbackRobotModel = playbackRobotModel
        self.ikPlanner = ikPlanner
        self.manipPlanner = manipPlanner
        self.footstepPlanner = footstepPlanner
        self.atlasDriver = atlasDriver
        self.lhandDriver = lhandDriver
        self.rhandDriver = rhandDriver
        self.multisenseDriver = multisenseDriver
        self.sensorJointController = sensorJointController
        self.view = view
        self.teleopPanel = teleopPanel

        # live operation flags:
        self.useFootstepPlanner = True
        self.visOnly = False
        self.planFromCurrentRobotState = True
        self.useDevelopment = False
        if (self.useDevelopment):
            self.visOnly = True
            self.planFromCurrentRobotState = False
            extraModels = [self.robotStateModel]
            self.affordanceUpdater  = affordanceupdater.AffordanceGraspUpdater(self.playbackRobotModel, self.ikPlanner, extraModels)
        else:
            extraModels = [self.playbackRobotModel]
            self.affordanceUpdater  = affordanceupdater.AffordanceGraspUpdater(self.robotStateModel, self.ikPlanner, extraModels)

        segmentation.affordanceManager.setAffordanceUpdater(self.affordanceUpdater)
        self.optionalUserPromptEnabled = True
        self.requiredUserPromptEnabled = True

        self.plans = []
        self.frameSyncs = {}

        self.graspingHand = 'left' # left, right, both

        self.tableData = None
        self.binFrame = None

        # top level switch between BDI or IHMC (locked base) and MIT (moving base and back)
        self.lockBack = True
        self.lockBase = True

        self.constraintSet = []

        self.reachDist = 0.12

        # Switch indicating whether to use affordances as a collision environment
        self.useCollisionEnvironment = True

        self.sceneID = None

    # Switch between simulation/visualisation and real robot operation
    def setMode(self, mode='visualization'):
        '''
        Switches between visualization and real robot operation.
        mode='visualization'
        mode='robot'
        '''

        if (mode == 'visualization'):
            print "Setting mode to VISUALIZATION"
            self.useDevelopment = True

            self.visOnly = True
            self.planFromCurrentRobotState = False
            extraModels = [self.robotStateModel]
            self.affordanceUpdater  = affordanceupdater.AffordanceGraspUpdater(self.playbackRobotModel, self.ikPlanner, extraModels)
        else:
            print "Setting mode to ROBOT OPERATION"
            self.useDevelopment = False

            extraModels = [self.playbackRobotModel]
            self.affordanceUpdater  = affordanceupdater.AffordanceGraspUpdater(self.robotStateModel, self.ikPlanner, extraModels)


    def addPlan(self, plan):
        self.plans.append(plan)


    ### Table and Bin Focused Functions
    def userFitTable(self):
        self.tableData = None
        self.picker = PointPicker(self.view, numberOfPoints=2, drawLines=True, callback=self.onSegmentTable)
        self.picker.start()

    def userFitBin(self):
        self.binFrame = None
        self.picker = PointPicker(self.view, numberOfPoints=2, drawLines=True, callback=self.onSegmentBin)
        self.picker.start()

    def waitForTableFit(self):
        while not self.tableData:
            yield

    def waitForBinFit(self):
        while not self.binFrame:
            yield

    def getInputPointCloud(self):
        polyData = segmentation.getCurrentRevolutionData()
        if polyData is None:
            obj = om.findObjectByName('scene')
            if obj:
                polyData = obj.polyData
            else: # fall back to map in case we used mapping rather than loading of a scene
                obj = om.findObjectByName('map')
                if obj:
                    polyData = obj.polyData
                else: # fall back to kinect source and get a frame copy
                    obj = om.findObjectByName('kinect source')
                    if obj:
                        polyData = obj.polyData

        return polyData

    def onSegmentTable(self, p1, p2):
        print p1
        print p2
        self.picker.stop()
        om.removeFromObjectModel(self.picker.annotationObj)
        self.picker = None

        tableData = segmentation.segmentTableEdge(self.getInputPointCloud(), p1, p2)

        pose = transformUtils.poseFromTransform(tableData.frame)
        desc = dict(classname='MeshAffordanceItem', Name='table', Color=[0,1,0], pose=pose)
        aff = segmentation.affordanceManager.newAffordanceFromDescription(desc)
        aff.setPolyData(tableData.mesh)

        self.tableData = tableData

        self.tableBox = vis.showPolyData(tableData.box, 'table box', parent=aff, color=[0,1,0], visible=False)
        #self.tableObj.actor.SetUserTransform(self.tableFrame.transform)
        self.tableBox.actor.SetUserTransform(tableData.frame)


        if self.useCollisionEnvironment:
            self.addCollisionObject(aff)

        return

        self.tableObj = vis.showPolyData(tableData.mesh, 'table', parent='affordances', color=[0,1,0])
        self.tableFrame = vis.showFrame(tableData.frame, 'table frame', parent=self.tableObj, scale=0.2)
        self.tableBox = vis.showPolyData(tableData.box, 'table box', parent=self.tableObj, color=[0,1,0], visible=False)
        self.tableObj.actor.SetUserTransform(self.tableFrame.transform)
        self.tableBox.actor.SetUserTransform(self.tableFrame.transform)


    def onSegmentBin(self, p1, p2):
        print p1
        print p2
        self.picker.stop()
        om.removeFromObjectModel(self.picker.annotationObj)
        self.picker = None

        om.removeFromObjectModel(om.findObjectByName('bin frame'))

        binEdge = p2 - p1
        zaxis = [0.0, 0.0, 1.0]
        xaxis = np.cross(binEdge, zaxis)
        xaxis /= np.linalg.norm(xaxis)
        yaxis = np.cross(zaxis, xaxis)

        t = transformUtils.getTransformFromAxes(xaxis, yaxis, zaxis)
        t.PostMultiply()
        t.Translate(p1)

        pose = transformUtils.poseFromTransform(t)
        desc = dict(classname='BoxAffordanceItem', Name='bin', uuid=newUUID(), pose=pose, Color=[1, 0, 0], Dimensions=[0.02,0.02,0.02])
        obj = affordancepanel.panel.affordanceFromDescription(desc)


        #self.binFrame = vis.showFrame(t, 'bin frame', parent=None, scale=0.2)

    def sortClustersOnTable(self, clusters):
        '''
        returns list copy of clusters, sorted left to right using the
        table coordinate system.  (Table y axis points right to left)
        '''
        tableTransform = om.findObjectByName('table').getChildFrame().transform

        tableYAxis = transformUtils.getAxesFromTransform(tableTransform)[1]
        tableOrigin = np.array(tableTransform.GetPosition())

        origins = [np.array(c.frame.GetPosition()) for c in clusters]
        dists = [np.dot(origin-tableOrigin, -tableYAxis) for origin in origins]

        return [clusters[i] for i in np.argsort(dists)]


    def cleanupSegmentedObjects(self):
        om.removeFromObjectModel(om.findObjectByName('segmentation'))
        self.clusterObjects = None
        self.segmentationData = None

    def segmentTableObjects(self):
        tableFrame = om.findObjectByName('table').getChildFrame()

        #tableCentroid = segmentation.computeCentroid(self.tableData.box)
        #self.tableData.frame.TransformPoint(tableCentroid, tableFrame)

        data = segmentation.segmentTableScene(self.getInputPointCloud(), tableFrame.transform.GetPosition() )
        data.clusters = self.sortClustersOnTable(data.clusters)

        objects = vis.showClusterObjects(data.clusters, parent='affordances')
        self.segmentationData = data

        self.clusterObjects = []
        for i, cluster in enumerate(objects):
            affObj = affordanceitems.MeshAffordanceItem.promotePolyDataItem(cluster)
            segmentation.affordanceManager.registerAffordance(affObj)
            self.clusterObjects.append(affObj)



    def graspTableObject(self, side):

        obj, objFrame = self.getNextTableObject(side)

        self.affordanceUpdater.graspAffordance( obj.getProperty('Name') , side)

        if self.ikPlanner.fixedBaseArm: # if we're dealing with the real world, close hand
            self.closeHand(side)
            return self.delay(5) # wait for three seconds to allow for hand to close

    def dropTableObject(self, side='left'):

        obj, _ = self.getNextTableObject(side)
        obj.setProperty('Visible', False)<commit_msg>tabledemo: graspTableObject also grasps affordance<commit_after>import os
import copy
import math
import functools
import numpy as np

from ddapp import transformUtils

from ddapp.asynctaskqueue import AsyncTaskQueue
from ddapp import objectmodel as om
from ddapp import visualization as vis
from ddapp import robotstate
from ddapp import segmentation
from ddapp import planplayback
from ddapp.pointpicker import PointPicker
from ddapp import vtkAll as vtk
from ddapp.simpletimer import SimpleTimer
from ddapp import affordanceupdater

from ddapp.debugVis import DebugData
from ddapp import affordanceitems
from ddapp import ikplanner
from ddapp import vtkNumpy
from numpy import array
from ddapp.uuidutil import newUUID
import affordancepanel
import ioUtils

from ddapp.tasks.taskuserpanel import TaskUserPanel
from ddapp.tasks.taskuserpanel import ImageBasedAffordanceFit

import ddapp.tasks.robottasks as rt


class TableDemo(object):

    def __init__(self, robotStateModel, playbackRobotModel, ikPlanner, manipPlanner, footstepPlanner,
                 atlasDriver, lhandDriver, rhandDriver, multisenseDriver, view, sensorJointController,
                 planPlaybackFunction, teleopPanel):
        self.planPlaybackFunction = planPlaybackFunction
        self.robotStateModel = robotStateModel
        self.playbackRobotModel = playbackRobotModel
        self.ikPlanner = ikPlanner
        self.manipPlanner = manipPlanner
        self.footstepPlanner = footstepPlanner
        self.atlasDriver = atlasDriver
        self.lhandDriver = lhandDriver
        self.rhandDriver = rhandDriver
        self.multisenseDriver = multisenseDriver
        self.sensorJointController = sensorJointController
        self.view = view
        self.teleopPanel = teleopPanel

        # live operation flags:
        self.useFootstepPlanner = True
        self.visOnly = False
        self.planFromCurrentRobotState = True
        self.useDevelopment = False
        if (self.useDevelopment):
            self.visOnly = True
            self.planFromCurrentRobotState = False
            extraModels = [self.robotStateModel]
            self.affordanceUpdater  = affordanceupdater.AffordanceGraspUpdater(self.playbackRobotModel, self.ikPlanner, extraModels)
        else:
            extraModels = [self.playbackRobotModel]
            self.affordanceUpdater  = affordanceupdater.AffordanceGraspUpdater(self.robotStateModel, self.ikPlanner, extraModels)

        segmentation.affordanceManager.setAffordanceUpdater(self.affordanceUpdater)
        self.optionalUserPromptEnabled = True
        self.requiredUserPromptEnabled = True

        self.plans = []
        self.frameSyncs = {}

        self.graspingHand = 'left' # left, right, both

        self.tableData = None
        self.binFrame = None

        # top level switch between BDI or IHMC (locked base) and MIT (moving base and back)
        self.lockBack = True
        self.lockBase = True

        self.constraintSet = []

        self.reachDist = 0.12

        # Switch indicating whether to use affordances as a collision environment
        self.useCollisionEnvironment = True

        self.sceneID = None

    # Switch between simulation/visualisation and real robot operation
    def setMode(self, mode='visualization'):
        '''
        Switches between visualization and real robot operation.
        mode='visualization'
        mode='robot'
        '''

        if (mode == 'visualization'):
            print "Setting mode to VISUALIZATION"
            self.useDevelopment = True

            self.visOnly = True
            self.planFromCurrentRobotState = False
            extraModels = [self.robotStateModel]
            self.affordanceUpdater  = affordanceupdater.AffordanceGraspUpdater(self.playbackRobotModel, self.ikPlanner, extraModels)
        else:
            print "Setting mode to ROBOT OPERATION"
            self.useDevelopment = False

            extraModels = [self.playbackRobotModel]
            self.affordanceUpdater  = affordanceupdater.AffordanceGraspUpdater(self.robotStateModel, self.ikPlanner, extraModels)


    def addPlan(self, plan):
        self.plans.append(plan)


    ### Table and Bin Focused Functions
    def userFitTable(self):
        self.tableData = None
        self.picker = PointPicker(self.view, numberOfPoints=2, drawLines=True, callback=self.onSegmentTable)
        self.picker.start()

    def userFitBin(self):
        self.binFrame = None
        self.picker = PointPicker(self.view, numberOfPoints=2, drawLines=True, callback=self.onSegmentBin)
        self.picker.start()

    def waitForTableFit(self):
        while not self.tableData:
            yield

    def waitForBinFit(self):
        while not self.binFrame:
            yield

    def getInputPointCloud(self):
        polyData = segmentation.getCurrentRevolutionData()
        if polyData is None:
            obj = om.findObjectByName('scene')
            if obj:
                polyData = obj.polyData
            else: # fall back to map in case we used mapping rather than loading of a scene
                obj = om.findObjectByName('map')
                if obj:
                    polyData = obj.polyData
                else: # fall back to kinect source and get a frame copy
                    obj = om.findObjectByName('kinect source')
                    if obj:
                        polyData = obj.polyData

        return polyData

    def onSegmentTable(self, p1, p2):
        print p1
        print p2
        self.picker.stop()
        om.removeFromObjectModel(self.picker.annotationObj)
        self.picker = None

        tableData = segmentation.segmentTableEdge(self.getInputPointCloud(), p1, p2)

        pose = transformUtils.poseFromTransform(tableData.frame)
        desc = dict(classname='MeshAffordanceItem', Name='table', Color=[0,1,0], pose=pose)
        aff = segmentation.affordanceManager.newAffordanceFromDescription(desc)
        aff.setPolyData(tableData.mesh)

        self.tableData = tableData

        self.tableBox = vis.showPolyData(tableData.box, 'table box', parent=aff, color=[0,1,0], visible=False)
        #self.tableObj.actor.SetUserTransform(self.tableFrame.transform)
        self.tableBox.actor.SetUserTransform(tableData.frame)


        if self.useCollisionEnvironment:
            self.addCollisionObject(aff)

        return

        self.tableObj = vis.showPolyData(tableData.mesh, 'table', parent='affordances', color=[0,1,0])
        self.tableFrame = vis.showFrame(tableData.frame, 'table frame', parent=self.tableObj, scale=0.2)
        self.tableBox = vis.showPolyData(tableData.box, 'table box', parent=self.tableObj, color=[0,1,0], visible=False)
        self.tableObj.actor.SetUserTransform(self.tableFrame.transform)
        self.tableBox.actor.SetUserTransform(self.tableFrame.transform)


    def onSegmentBin(self, p1, p2):
        print p1
        print p2
        self.picker.stop()
        om.removeFromObjectModel(self.picker.annotationObj)
        self.picker = None

        om.removeFromObjectModel(om.findObjectByName('bin frame'))

        binEdge = p2 - p1
        zaxis = [0.0, 0.0, 1.0]
        xaxis = np.cross(binEdge, zaxis)
        xaxis /= np.linalg.norm(xaxis)
        yaxis = np.cross(zaxis, xaxis)

        t = transformUtils.getTransformFromAxes(xaxis, yaxis, zaxis)
        t.PostMultiply()
        t.Translate(p1)

        pose = transformUtils.poseFromTransform(t)
        desc = dict(classname='BoxAffordanceItem', Name='bin', uuid=newUUID(), pose=pose, Color=[1, 0, 0], Dimensions=[0.02,0.02,0.02])
        obj = affordancepanel.panel.affordanceFromDescription(desc)


        #self.binFrame = vis.showFrame(t, 'bin frame', parent=None, scale=0.2)

    def sortClustersOnTable(self, clusters):
        '''
        returns list copy of clusters, sorted left to right using the
        table coordinate system.  (Table y axis points right to left)
        '''
        tableTransform = om.findObjectByName('table').getChildFrame().transform

        tableYAxis = transformUtils.getAxesFromTransform(tableTransform)[1]
        tableOrigin = np.array(tableTransform.GetPosition())

        origins = [np.array(c.frame.GetPosition()) for c in clusters]
        dists = [np.dot(origin-tableOrigin, -tableYAxis) for origin in origins]

        return [clusters[i] for i in np.argsort(dists)]


    def cleanupSegmentedObjects(self):
        om.removeFromObjectModel(om.findObjectByName('segmentation'))
        self.clusterObjects = None
        self.segmentationData = None

    def segmentTableObjects(self):
        tableFrame = om.findObjectByName('table').getChildFrame()

        #tableCentroid = segmentation.computeCentroid(self.tableData.box)
        #self.tableData.frame.TransformPoint(tableCentroid, tableFrame)

        data = segmentation.segmentTableScene(self.getInputPointCloud(), tableFrame.transform.GetPosition() )
        data.clusters = self.sortClustersOnTable(data.clusters)

        objects = vis.showClusterObjects(data.clusters, parent='affordances')
        self.segmentationData = data

        self.clusterObjects = []
        for i, cluster in enumerate(objects):
            affObj = affordanceitems.MeshAffordanceItem.promotePolyDataItem(cluster)
            segmentation.affordanceManager.registerAffordance(affObj)
            self.clusterObjects.append(affObj)



    def graspTableObject(self, side='left'):

        obj, objFrame = self.getNextTableObject(side)
        if self.useCollisionEnvironment:
            objAffordance = om.findObjectByName(obj.getProperty('Name') + ' affordance')

        self.affordanceUpdater.graspAffordance(obj.getProperty('Name'), side)
        if self.useCollisionEnvironment:
            self.affordanceUpdater.graspAffordance(objAffordance.getProperty('Name'), side)

        if self.ikPlanner.fixedBaseArm: # if we're dealing with the real world, close hand
            self.closeHand(side)
            return self.delay(5) # wait for three seconds to allow for hand to close

    def dropTableObject(self, side='left'):

        obj, _ = self.getNextTableObject(side)
        obj.setProperty('Visible', False)
{'fraction_non_alphanumeric': 0.08225475988206486, 'fraction_numerical': 0.005397031632602069, 'mean_word_length': 2.8381281166091292, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2620479', 'n_tokens_mistral': 5789, 'n_tokens_neox': 5557, 'n_words': 1245}