Coverage for /home/travis/build/RobotLocomotion/director/build/src/director-build/lib/python2.7/dist-packages/director/terraintask : 0%

import os 

import sys 

import vtkAll as vtk 

import math 

import time 

import types 

import functools 

import numpy as np 

from collections import defaultdict 

from director import transformUtils 

from director import lcmUtils 

from director.timercallback import TimerCallback 

from director.asynctaskqueue import AsyncTaskQueue 

from director import objectmodel as om 

from director import visualization as vis 

from director import applogic as app 

from director.debugVis import DebugData 

from director import ikplanner 

from director import ioUtils 

from director.simpletimer import SimpleTimer 

from director.utime import getUtime 

from director import affordanceitems 

from director import robotstate 

from director import robotplanlistener 

from director import segmentation 

from director import planplayback 

from director import affordanceupdater 

from director import segmentationpanel 

from director import footstepsdriver 

from director import footstepsdriverpanel 

from director import pointpicker 

from director.footstepsdriver import FootstepRequestGenerator 

import director.terrain 

from director.tasks.taskuserpanel import TaskUserPanel 

from director.tasks.taskuserpanel import ImageBasedAffordanceFit 

import director.tasks.robottasks as rt 

import director.tasks.taskmanagerwidget as tmw 

import drc as lcmdrc 

import copy 

import re 

import glob 

from PythonQt import QtCore, QtGui 

blockWidth = (15 + 3/8.0) * 0.0254 

blockLength = (15 + 5/8.0) * 0.0254 

blockHeight = (5 + 5/8.0) * 0.0254 

blockDiagonal = np.linalg.norm([blockWidth, blockLength, blockHeight]) 

blockSafetyMargin = [0.03, 0.05] 

class TerrainTask(object): 

    def __init__(self, robotSystem): 

        self.robotSystem = robotSystem 

        self.cinderblockPrefix = 'cinderblock' 

        self.stairPrefix = 'stair' 

        self.currentRow = {'left':-1, 'right':-1} 

        self.stanceBlocks = {'left':None, 'right':None} 

        self.useTextures = False 

        self.constrainBlockSize = True 

        self.blockFitAlgo = 1 

        self.timer = TimerCallback(targetFps=10) 

        self.timer.callback = self.updateBlockState 

        self.defaultLeadingFoot = 'right' 

        self.removeDuringOrganize = False 

        self.numberOfPrefabSteps = -1 

        self.prefabStepStartOffset = 0 

        self.relativeStartPos = np.array([0,0]) 

        self.relativeStartYaw = 0 

        # terrain config file stuff 

        self.terrainConfigDir = os.path.join(director.getDRCBaseDir(), 'software','config','terrain') 

        files = glob.glob(os.path.join(self.terrainConfigDir,'*.py')) 

        self.terrainConfigList = [] 

        for f in files: 

            self.terrainConfigList.append(os.path.basename(f).split('.')[0]) 

        if len(self.terrainConfigList)>0: 

            self.terrainConfigList.sort() 

            self.loadTerrainConfig(self.terrainConfigList[0]) 

    def startBlockUpdater(self): 

        self.timer.start() 

    def stopBlockUpdater(self): 

        self.timer.stop() 

    def requestRaycastTerrain(self): 

        affs = self.robotSystem.affordanceManager.getCollisionAffordances() 

        xy = self.robotSystem.robotStateJointController.q[:2] 

        self.robotSystem.raycastDriver.requestRaycast(affs, xy-4.5, xy+4.5) 

    def walkToTiltedCinderblocks(self): 

        frame = om.findObjectByName('cinderblock stance frame') 

        assert frame 

        frameCopy = transformUtils.copyFrame(frame.transform) 

        footstepsdriverpanel.panel.onNewWalkingGoal(frameCopy) 

    def requestBlockFit(self): 

        for block in self.getIndexedBlockAffordances(): 

            om.removeFromObjectModel(block) 

        blockSize = self.terrainConfig['blockSize'].tolist() 

        msg = lcmdrc.block_fit_request_t() 

        msg.utime = getUtime() 

        if self.constrainBlockSize: 

            msg.dimensions = blockSize 

        else: 

            msg.dimensions = [0, 0, blockSize[2]] 

        msg.name_prefix = self.terrainConfig['blockName'] 

        msg.algorithm = self.blockFitAlgo 

        lcmUtils.publish('BLOCK_FIT_TRIGGER', msg) 

    def spawnGroundAffordance(self): 

        polyData = segmentation.getCurrentRevolutionData() 

        groundOrigin, normal, groundPoints, _ = segmentation.segmentGround(polyData) 

        stanceFrame = FootstepRequestGenerator.getRobotStanceFrame(self.robotSystem.robotStateModel) 

        origin = np.array(stanceFrame.GetPosition()) 

        origin = segmentation.projectPointToPlane(origin, groundOrigin, normal) 

        zaxis = normal 

        xaxis = transformUtils.getAxesFromTransform(stanceFrame)[0] 

        yaxis = np.cross(zaxis, xaxis) 

        yaxis /= np.linalg.norm(yaxis) 

        xaxis = np.cross(yaxis, zaxis) 

        xaxis /= np.linalg.norm(xaxis) 

        boxThickness = 0.01 

        t = transformUtils.getTransformFromAxes(xaxis, yaxis, zaxis) 

        t.PreMultiply() 

        t.Translate(0.0, 0.0, -boxThickness/2.0) 

        t.PostMultiply() 

        t.Translate(origin) 

        om.removeFromObjectModel(om.findObjectByName('ground affordance')) 

        pose = transformUtils.poseFromTransform(t) 

        desc = dict(classname='BoxAffordanceItem', Name='ground affordance', Dimensions=[100, 100, boxThickness], pose=pose) 

        aff = segmentation.affordanceManager.newAffordanceFromDescription(desc) 

        aff.setProperty('Visible', False) 

        aff.setProperty('Alpha', 0.2) 

    def spawnFootplaneGroundAffordance(self, side): 

        boxThickness = 0.01 

        footFrame = self.getFootFrameAtSole(self.sideToFootLinkName(side)) 

        t = transformUtils.copyFrame(footFrame) 

        t.PreMultiply() 

        t.Translate(0.0, 0.0, -boxThickness/2.0) 

        om.removeFromObjectModel(om.findObjectByName('ground affordance')) 

        pose = transformUtils.poseFromTransform(t) 

        desc = dict(classname='BoxAffordanceItem', Name='ground affordance', Dimensions=[100, 100, boxThickness], pose=pose) 

        aff = segmentation.affordanceManager.newAffordanceFromDescription(desc) 

        aff.setProperty('Visible', False) 

        aff.setProperty('Alpha', 0.2) 

    def getPlanningStartPose(self): 

        return self.robotSystem.robotStateJointController.q.copy() 

    # 

    # TERRAIN SPAWNING STUFF STARTS HERE 

    # 

    def loadTerrainConfig(self, terrainType): 

        # first clear out existing affordances 

        if hasattr(self, 'terrainConfig') and self.terrainConfig is not None: 

            blocks = self.getAllBlockAffordances() 

            if blocks is not None: 

                for obj in blocks: 

                    om.removeFromObjectModel(obj) 

        # load file 

        configFile = os.path.join(self.terrainConfigDir, terrainType+'.py') 

        self.terrainConfig = {} 

        execfile(configFile, self.terrainConfig) 

    def getAllBlockAffordances(self): 

        return [obj for obj in om.getObjects() if obj.getProperty('Name').startswith(self.terrainConfig['blockName'])] 

    def getTabularBlockAffordances(self): 

        blocks = [] 

        for obj in om.getObjects(): 

            name = obj.getProperty('Name') 

            if re.match('^%s \(\d+,\d+\)$' % self.terrainConfig['blockName'], name): 

                blocks.append(obj) 

        return blocks 

    def getIndexedBlockAffordances(self): 

        blocks = [] 

        for obj in om.getObjects(): 

            name = obj.getProperty('Name') 

            if re.match(r'^%s \d+$' % self.terrainConfig['blockName'], name): 

                blocks.append(obj) 

        return blocks 

    def createStartingGoal(self): 

        blockTable = self.createBlockObjectTable() 

        if len(blockTable)>1 and len(blockTable[0])>0: 

            firstBlockFrame = transformUtils.copyFrame(blockTable[0][0].getChildFrame().transform) 

            secondBlockFrame = transformUtils.copyFrame(blockTable[1][0].getChildFrame().transform) 

            secondBlockDir = np.array(secondBlockFrame.GetPosition()) - np.array(firstBlockFrame.GetPosition()) 

            blockXDir = transformUtils.getAxesFromTransform(firstBlockFrame)[0] 

            yawOffset = np.arctan2(secondBlockDir[1],secondBlockDir[0]) - np.arctan2(blockXDir[1],blockXDir[0]) 

            yawDiff = self.relativeYawOffset - yawOffset 

            yawAngle = np.radians(90) * np.round(yawDiff / np.radians(90)) 

            startingFrame = transformUtils.copyFrame(self.relativeStartFrame) 

            startingFrame.PostMultiply() 

            startingFrame.RotateZ(np.degrees(-yawAngle)) 

            startingFrame.Concatenate(firstBlockFrame) 

            goalFrame = startingFrame 

            footstepsdriverpanel.panel.onNewWalkingGoal(goalFrame) 

        else: 

            print 'error: no blocks defined; use Spawn Terrain button' 

            return 

    def correctFrameYaw(self, t1, t2): 

        R1 = np.array(transformUtils.getAxesFromTransform(t1)).T 

        R2 = np.array(transformUtils.getAxesFromTransform(t2)).T 

        R = R1.dot(np.round(R1.T.dot(R2))) 

        t = transformUtils.getTransformFromAxesAndOrigin(R[:,0],R[:,1],R[:,2],t1.GetPosition()) 

        return t 

    def createFootstepsForTerrain(self): 

        footstepData = self.terrainConfig['footstepData'] 

        # check that we have required data 

        blockObjectTable = self.createBlockObjectTable() 

        if len(blockObjectTable) == 0: 

            print 'error: no blocks available for footsteps' 

            return 

        stanceFrame = FootstepRequestGenerator.getRobotStanceFrame(self.robotSystem.robotStateModel) 

        stanceAxes = transformUtils.getAxesFromTransform(stanceFrame) 

        # get coord frame of first block 

        firstBlockFrame = transformUtils.copyFrame(blockObjectTable[0][0].getChildFrame().transform) 

        firstBlockFrame = self.correctFrameYaw(firstBlockFrame, stanceFrame) 

        toFirstBlock = firstBlockFrame.GetLinearInverse() 

        # get current frames of feet 

        leftFootFrame = self.getFootFrameAtSole(self.sideToFootLinkName('left')) 

        rightFootFrame = self.getFootFrameAtSole(self.sideToFootLinkName('right')) 

        leftPos = np.array(leftFootFrame.GetPosition()[0:2]) 

        leftLine = np.array(transformUtils.getAxesFromTransform(leftFootFrame)[0][0:2]) 

        leftLine = np.append(leftLine, -np.dot(leftLine,leftPos)) 

        rightPos = np.array(rightFootFrame.GetPosition()[0:2]) 

        rightLine = np.array(transformUtils.getAxesFromTransform(rightFootFrame)[0][0:2]) 

        rightLine = np.append(rightLine, -np.dot(rightLine,rightPos)) 

        # generate new footstep frames 

        steps = [] 

        for foot, blockIndex, offset, supportType in footstepData: 

            block = blockObjectTable[blockIndex[0]][blockIndex[1]] 

            if block is None: 

                print 'error: no block for footstep (%d,%d)' % blockIndex 

                return 

            z = np.array(block.getProperty('Dimensions'))[2]/2.0 

            pt = np.array([offset[0], offset[1], z]) 

            t = transformUtils.copyFrame(block.getChildFrame().transform) 

            # TODO: should feet be aligned with forward direction or block? 

            blockAxes = transformUtils.getAxesFromTransform(t) 

            x = np.round(np.dot(stanceAxes[0], blockAxes[0])) 

            y = np.round(np.dot(stanceAxes[0], blockAxes[1])) 

            theta = np.degrees(np.arctan2(y,x)) 

            # compute and apply lateral offset if desired 

            if self.terrainConfig['forceZeroLateralFootstepOffset']: 

                curBlockFrame = self.correctFrameYaw(t, stanceFrame) 

                curBlockToFirst = transformUtils.copyFrame(curBlockFrame) 

                R = np.array(transformUtils.getAxesFromTransform(curBlockToFirst)).T 

                curBlockToFirst.PostMultiply() 

                curBlockToFirst.Concatenate(toFirstBlock) 

                p = np.array(curBlockToFirst.GetPosition()) 

                yDir = np.array(transformUtils.getAxesFromTransform(curBlockToFirst)[1]) 

                dist = p[1]/yDir[1] 

                pt[1] -= dist 

            # compute new foot frame 

            t.PreMultiply() 

            t.RotateZ(theta); 

            t.Translate(pt) 

            # only append footstep if it is forward of current foot 

            pos = np.array(t.GetPosition()[0:2]) 

            if foot=='left': 

                dist = leftLine[0:2].dot(pos)+leftLine[2] 

            else: 

                dist = rightLine[0:2].dot(pos)+rightLine[2] 

            steps.append({'transform':t, 'foot':foot, 'support':supportType, 'dist':dist}) 

            #obj = vis.showFrame(t, '%s step frame' % block.getProperty('Name'), parent='step frames', scale=0.2) 

        # find first valid start step 

        startIndex = 0 

        for step in steps: 

            if (step['dist'] < 0.1): 

                startIndex += 1 

            else: 

                break 

        # adjust start index by manually specified offset 

        startIndex += self.prefabStepStartOffset 

        startIndex = max(0, startIndex) 

        startIndex = min(len(steps), startIndex) 

        steps = steps[startIndex:] 

        # remove footsteps beyond specified limit 

        if self.numberOfPrefabSteps >= 0: 

            endIndex = min(len(steps), self.numberOfPrefabSteps) 

            steps = steps[0:endIndex] 

        # send footstep request 

        if len(steps) > 0: 

            leadingFoot = steps[0]['foot'] 

        else: 

            print 'error: no footsteps' 

            return 

        startPose = self.getPlanningStartPose() 

        helper = FootstepRequestGenerator(self.robotSystem.footstepsDriver) 

        stepFrames = [step['transform'] for step in steps] 

        request = helper.makeFootstepRequest(startPose, stepFrames, leadingFoot,  snapToTerrain=False) 

        for i in range(len(stepFrames)): 

            request.goal_steps[i].params.support_contact_groups = steps[i]['support'] 

        self.robotSystem.footstepsDriver.sendFootstepPlanRequest(request, waitForResponse=True) 

    def createBlockObjectTable(self): 

        types = self.terrainConfig['blockTypes'] 

        table = [[None for i in range(len(types[0]))] for j in range(len(types))] 

        for obj in self.getTabularBlockAffordances(): 

            name = obj.getProperty('Name') 

            m = re.match('%s \((\d+)\,(\d+)\)' % self.terrainConfig['blockName'], name) 

            if m is not None: 

                row = int(m.group(1)) 

                col = int(m.group(2)) 

                table[row][col] = obj 

        return table 

    def spawnCinderblockTerrain(self, cols=[]): 

        # clear old affordances 

        for obj in self.getTabularBlockAffordances(): 

            om.removeFromObjectModel(obj) 

        blockTypes = self.terrainConfig['blockTypes'] 

        blockLevels = self.terrainConfig['blockLevels'] 

        blockAngleMap = self.terrainConfig['blockAngleMap'] 

        blockSize = self.terrainConfig['blockSize'] 

        if len(cols) == 0: 

            cols = range(len(blockTypes[0])) 

        # create frame for starting position 

        yaw = self.terrainConfig['startingYaw'] 

        pos = np.copy(self.terrainConfig['startingPosition']) 

        pos[1] -= blockSize[0]*np.mean(np.array(cols)) 

        offsetFrame = transformUtils.frameFromPositionAndRPY(pos, np.array([0,0,yaw])) 

        stanceFrame = FootstepRequestGenerator.getRobotStanceFrame(self.robotSystem.robotStateModel) 

        blockFrame = transformUtils.copyFrame(stanceFrame) 

        blockFrame.PreMultiply() 

        blockFrame.Concatenate(offsetFrame.GetLinearInverse()) 

        firstBlockFrame = None 

        secondBlockFrame = None 

        for row in range(len(blockTypes)): 

            for col in cols: 

                blockType = blockTypes[row][col] 

                blockPitch = np.radians(blockAngleMap[blockType][0]) 

                blockYaw = np.radians(blockAngleMap[blockType][1]) 

                blockLevel = blockLevels[row][col] 

                verticalOffset = blockSize[2]/2.0 + np.sin(blockPitch)*blockSize[1]/2.0 

                rpy = np.degrees([0.0, blockPitch, blockYaw]) 

                pos = [blockSize[1]*row, -blockSize[0]*col, blockSize[2]*blockLevel + verticalOffset] 

                offsetFrame = transformUtils.frameFromPositionAndRPY(pos, rpy) 

                offsetFrame.PostMultiply() 

                offsetFrame.Concatenate(blockFrame) 

                if row==0 and col==0: 

                    firstBlockFrame = offsetFrame 

                if row==1 and col==0: 

                    secondBlockFrame = offsetFrame 

                pose = transformUtils.poseFromTransform(offsetFrame) 

                desc = dict(classname='BoxAffordanceItem', Name='%s (%d,%d)' % (self.terrainConfig['blockName'], row, col), Dimensions=blockSize.tolist(), pose=pose, Color=self.terrainConfig['blockColor']) 

                block = self.robotSystem.affordanceManager.newAffordanceFromDescription(desc) 

        if firstBlockFrame is not None and secondBlockFrame is not None: 

            relativeTransform = transformUtils.copyFrame(stanceFrame) 

            relativeTransform.PostMultiply() 

            relativeTransform.Concatenate(firstBlockFrame.GetLinearInverse()) 

            self.relativeStartFrame = relativeTransform 

            firstBlockPos = np.array(firstBlockFrame.GetPosition()) 

            secondBlockPos = np.array(secondBlockFrame.GetPosition()) 

            secondBlockDir = (secondBlockPos - firstBlockPos)[0:2] 

            blockXDir = transformUtils.getAxesFromTransform(firstBlockFrame)[0] 

            yawOffset = np.arctan2(secondBlockDir[1],secondBlockDir[0]) - np.arctan2(blockXDir[1],blockXDir[0]) 

            self.relativeYawOffset = yawOffset 

        #for block in self.getCinderblockAffordances(): 

        #    frameSync.addFrame(block.getChildFrame(), ignoreIncoming=True) 

    def showBlocks(self, blocks, visible): 

        for b in blocks: 

            b.setProperty('Visible', visible) 

            b.setProperty('Collision Enabled', visible) 

    def assignBlocks(self): 

        # get list of ideal blocks and detected blocks 

        idealBlocks = self.getTabularBlockAffordances() 

        detectedBlocks = self.getIndexedBlockAffordances() 

        # check that we have both spawned and detected blocks 

        if len(detectedBlocks) == 0: 

            print 'error: no blocks detected; use Fit Blocks button' 

            return 

        elif len(idealBlocks) == 0: 

            print 'error: no blocks spawned; use Spawn Terrain button' 

            return 

        # hide ideal blocks, show detected blocks 

        self.showBlocks(idealBlocks, False) 

        self.showBlocks(detectedBlocks, True) 

        # pick point in detected blocks 

        picker = pointpicker.AffordancePicker(app.getDRCView(), self.robotSystem.affordanceManager, filterFunc=lambda x: x in detectedBlocks) 

        picker.callbackFunc = functools.partial(self.assignBlocks2) 

        picker.start() 

    def assignBlocks2(self, pickedDetectedBlock): 

        proceed = len(pickedDetectedBlock)==1 

        # get list of ideal blocks and detected blocks 

        idealBlocks = self.getTabularBlockAffordances() 

        detectedBlocks = self.getIndexedBlockAffordances() 

        # hide detected blocks, show ideal blocks 

        self.showBlocks(detectedBlocks, not proceed) 

        self.showBlocks(idealBlocks, True) 

        # TODO: make an abort function 

        if not proceed: 

            print 'cannot proceed: no spawned block was clicked' 

            return 

        # pick point in ideal blocks 

        pickedDetectedBlock = pickedDetectedBlock[0] 

        picker = pointpicker.AffordancePicker(app.getDRCView(), self.robotSystem.affordanceManager, filterFunc=lambda x: x in idealBlocks) 

        picker.callbackFunc = functools.partial(self.assignBlocks3, pickedDetectedBlock) 

        picker.start() 

    def assignBlocks3(self, pickedDetectedBlock, pickedIdealBlock): 

        proceed = len(pickedIdealBlock)==1 

        if not proceed: 

            return 

        pickedIdealBlock = pickedIdealBlock[0] 

        # get list of ideal blocks and detected blocks 

        idealBlocks = self.getTabularBlockAffordances() 

        detectedBlocks = self.getIndexedBlockAffordances() 

        # grab poses of selected blocks 

        t1 = transformUtils.copyFrame(pickedIdealBlock.getChildFrame().transform) 

        t2 = transformUtils.copyFrame(pickedDetectedBlock.getChildFrame().transform) 

        # due to symmetry, find minimal of the 4 possible rotations 

        axes1 = transformUtils.getAxesFromTransform(t1) 

        axes2 = transformUtils.getAxesFromTransform(t2) 

        R1 = np.array(axes1).T 

        R2 = np.array(axes2).T 

        R = R2.T.dot(R1) 

        permIndices = np.argmax(np.fabs(R),axis=1) 

        R2 = R2[:,permIndices] 

        values = [R[permIndices[0]][0], R[permIndices[1]][1], R[permIndices[2]][2]] 

        R2 = R2.dot(np.sign(np.diag(np.array(values)))) 

        rot = R2.T.dot(R1) 

        # find translation so that tops of blocks align 

        d1 = pickedIdealBlock.getProperty('Dimensions') 

        d2 = pickedDetectedBlock.getProperty('Dimensions') 

        pos1 = np.array(t1.GetPosition()) + np.array(axes1[2])*d1[2]/2 

        pos2 = np.array(t2.GetPosition()) + np.array(axes2[2])*d2[2]/2 

        # compute initial (minimal) transform between selected blocks 

        correction = vtk.vtkTransform() 

        correction.PostMultiply() 

        correction.Translate(-np.array(pos1)) 

        rotTransform = transformUtils.getTransformFromAxes(rot[0,:],rot[1,:],rot[2,:]) 

        correction.Concatenate(rotTransform) 

        correction.Translate(pos2) 

        # move all ideal blocks and reset color 

        for b in idealBlocks: 

            t = b.getChildFrame().transform 

            t.PostMultiply() 

            t.Concatenate(correction) 

            t.Modified() 

            b.setProperty('Color', self.terrainConfig['blockColor']) 

        # associate blocks (only use xy distances) 

        matches = [] 

        blockSize = self.terrainConfig['blockSize'] 

        for b in detectedBlocks: 

            t1 = b.getChildFrame().transform 

            p1 = np.array(t1.GetPosition()) 

            normal1 = np.array(transformUtils.getAxesFromTransform(t1)[2]) 

            bestMatch = None 

            minDist = 1e10 

            for ib in idealBlocks: 

                t2 = ib.getChildFrame().transform 

                p2 = np.array(t2.GetPosition()) 

                dist = np.linalg.norm(p2[0:2]-p1[0:2])