<s> import os <EOL> import vtkAll as vtk <EOL> import math <EOL> import time <EOL> import re <EOL> import numpy as np <EOL> from director . timercallback import TimerCallback <EOL> from director import objectmodel as om <EOL> from director . simpletimer import SimpleTimer <EOL> from director . utime import getUtime <EOL> from director import robotstate <EOL> import pickle <EOL> import scipy . interpolate <EOL> class PlanPlayback ( object ) : <EOL> def __init__ ( self ) : <EOL> self . animationCallback = None <EOL> self . animationTimer = None <EOL> self . interpolationMethod = '<STR_LIT>' <EOL> self . playbackSpeed = <NUM_LIT:1.0> <EOL> self . jointNameRegex = '<STR_LIT>' <EOL> @ staticmethod <EOL> def getPlanPoses ( msgOrList ) : <EOL> if isinstance ( msgOrList , list ) : <EOL> messages = msgOrList <EOL> allPoseTimes , allPoses = PlanPlayback . getPlanPoses ( messages [ <NUM_LIT:0> ] ) <EOL> for msg in messages [ <NUM_LIT:1> : ] : <EOL> poseTimes , poses = PlanPlayback . getPlanPoses ( msg ) <EOL> poseTimes += allPoseTimes [ - <NUM_LIT:1> ] <EOL> allPoseTimes = np . hstack ( ( allPoseTimes , poseTimes [ <NUM_LIT:1> : ] ) ) <EOL> allPoses += poses [ <NUM_LIT:1> : ] <EOL> return allPoseTimes , allPoses <EOL> else : <EOL> msg = robotstate . asRobotPlan ( msgOrList ) <EOL> poses = [ ] <EOL> poseTimes = [ ] <EOL> for plan in msg . plan : <EOL> pose = robotstate . convertStateMessageToDrakePose ( plan ) <EOL> poseTimes . append ( plan . utime / <NUM_LIT> ) <EOL> poses . append ( pose ) <EOL> return np . array ( poseTimes ) , poses <EOL> @ staticmethod <EOL> def getPlanElapsedTime ( msg ) : <EOL> msg = robotstate . asRobotPlan ( msg ) <EOL> startTime = msg . plan [ <NUM_LIT:0> ] . utime <EOL> endTime = msg . plan [ - <NUM_LIT:1> ] . utime <EOL> return ( endTime - startTime ) / <NUM_LIT> <EOL> def stopAnimation ( self ) : <EOL> if self . animationTimer : <EOL> self . animationTimer . stop ( ) <EOL> def setInterpolationMethod ( method ) : <EOL> self . interpolationMethod = method <EOL> def playPlan ( self , msg , jointController ) : <EOL> self . playPlans ( poseTimes , [ msg ] , jointController ) <EOL> def playPlans ( self , messages , jointController ) : <EOL> assert len ( messages ) <EOL> poseTimes , poses = self . getPlanPoses ( messages ) <EOL> self . playPoses ( poseTimes , poses , jointController ) <EOL> def getPoseInterpolatorFromPlan ( self , message ) : <EOL> poseTimes , poses = self . getPlanPoses ( message ) <EOL> return self . getPoseInterpolator ( poseTimes , poses ) <EOL> def getPoseInterpolator ( self , poseTimes , poses , unwrap_rpy = True ) : <EOL> if unwrap_rpy : <EOL> poses = np . array ( poses , copy = True ) <EOL> poses [ : , <NUM_LIT:3> : <NUM_LIT:6> ] = np . unwrap ( poses [ : , <NUM_LIT:3> : <NUM_LIT:6> ] , axis = <NUM_LIT:0> ) <EOL> if self . interpolationMethod in [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] : <EOL> f = scipy . interpolate . interp1d ( poseTimes , poses , axis = <NUM_LIT:0> , kind = self . interpolationMethod ) <EOL> elif self . interpolationMethod == '<STR_LIT>' : <EOL> f = scipy . interpolate . pchip ( poseTimes , poses , axis = <NUM_LIT:0> ) <EOL> return f <EOL> def getPlanPoseMeshes ( self , messages , jointController , robotModel , numberOfSamples ) : <EOL> poseTimes , poses = self . getPlanPoses ( messages ) <EOL> f = self . getPoseInterpolator ( poseTimes , poses ) <EOL> sampleTimes = np . linspace ( poseTimes [ <NUM_LIT:0> ] , poseTimes [ - <NUM_LIT:1> ] , numberOfSamples ) <EOL> meshes = [ ] <EOL> for sampleTime in sampleTimes : <EOL> pose = f ( sampleTime ) <EOL> jointController . setPose ( '<STR_LIT>' , pose ) <EOL> polyData = vtk . vtkPolyData ( ) <EOL> robotModel . model . getModelMesh ( polyData ) <EOL> meshes . append ( polyData ) <EOL> return meshes <EOL> def showPoseAtTime ( self , time , jointController , poseInterpolator ) : <EOL> pose = poseInterpolator ( time ) <EOL> jointController . setPose ( '<STR_LIT>' , pose ) <EOL> def playPoses ( self , poseTimes , poses , jointController ) : <EOL> f = self . getPoseInterpolator ( poseTimes , poses ) <EOL> timer = SimpleTimer ( ) <EOL> def updateAnimation ( ) : <EOL> tNow = timer . elapsed ( ) * self . playbackSpeed <EOL> if tNow > poseTimes [ - <NUM_LIT:1> ] : <EOL> pose = poses [ - <NUM_LIT:1> ] <EOL> jointController . setPose ( '<STR_LIT>' , pose ) <EOL> if self . animationCallback : <EOL> self . animationCallback ( ) <EOL> return False <EOL> pose = f ( tNow ) <EOL> jointController . setPose ( '<STR_LIT>' , pose ) <EOL> if self . animationCallback : <EOL> self . animationCallback ( ) <EOL> self . animationTimer = TimerCallback ( ) <EOL> self . animationTimer . targetFps = <NUM_LIT> <EOL> self . animationTimer . callback = updateAnimation <EOL> self . animationTimer . start ( ) <EOL> updateAnimation ( ) <EOL> def picklePlan ( self , filename , msg ) : <EOL> poseTimes , poses = self . getPlanPoses ( msg ) <EOL> pickle . dump ( ( poseTimes , poses ) , open ( filename , '<STR_LIT:w>' ) ) <EOL> def getMovingJointNames ( self , msg ) : <EOL> poseTimes , poses = self . getPlanPoses ( msg ) <EOL> diffs = np . diff ( poses , axis = <NUM_LIT:0> ) <EOL> jointIds = np . unique ( np . where ( diffs != <NUM_LIT:0.0> ) [ <NUM_LIT:1> ] ) <EOL> jointNames = [ robotstate . getDrakePoseJointNames ( ) [ jointId ] for jointId in jointIds ] <EOL> return jointNames <EOL> def plotPlan ( self , msg ) : <EOL> poseTimes , poses = self . getPlanPoses ( msg ) <EOL> self . plotPoses ( poseTimes , poses ) <EOL> def plotPoses ( self , poseTimes , poses ) : <EOL> import matplotlib . pyplot as plt <EOL> poses = np . array ( poses ) <EOL> if self . jointNameRegex : <EOL> jointIds = range ( poses . shape [ <NUM_LIT:1> ] ) <EOL> else : <EOL> diffs = np . diff ( poses , axis = <NUM_LIT:0> ) <EOL> jointIds = np . unique ( np . where ( diffs != <NUM_LIT:0.0> ) [ <NUM_LIT:1> ] ) <EOL> jointNames = [ robotstate . getDrakePoseJointNames ( ) [ jointId ] for jointId in jointIds ] <EOL> jointTrajectories = [ poses [ : , jointId ] for jointId in jointIds ] <EOL> seriesNames = [ ] <EOL> sampleResolutionInSeconds = <NUM_LIT> <EOL> numberOfSamples = ( poseTimes [ - <NUM_LIT:1> ] - poseTimes [ <NUM_LIT:0> ] ) / sampleResolutionInSeconds <EOL> xnew = np . linspace ( poseTimes [ <NUM_LIT:0> ] , poseTimes [ - <NUM_LIT:1> ] , numberOfSamples ) <EOL> fig = plt . figure ( ) <EOL> ax = fig . add_subplot ( <NUM_LIT> ) <EOL> for jointId , jointName , jointTrajectory in zip ( jointIds , jointNames , jointTrajectories ) : <EOL> if self . jointNameRegex and not re . match ( self . jointNameRegex , jointName ) : <EOL> continue <EOL> x = poseTimes <EOL> y = jointTrajectory <EOL> y = np . rad2deg ( y ) <EOL> if self . interpolationMethod in [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] : <EOL> f = scipy . interpolate . interp1d ( x , y , kind = self . interpolationMethod ) <EOL> elif self . interpolationMethod == '<STR_LIT>' : <EOL> f = scipy . interpolate . pchip ( x , y ) <EOL> ax . plot ( x , y , '<STR_LIT>' ) <EOL> seriesNames . append ( jointName + '<STR_LIT>' ) <EOL> ax . plot ( xnew , f ( xnew ) , '<STR_LIT:->' ) <EOL> seriesNames . append ( jointName + '<STR_LIT:U+0020>' + self . interpolationMethod ) <EOL> ax . legend ( seriesNames , loc = '<STR_LIT>' ) . draggable ( ) <EOL> ax . set_xlabel ( '<STR_LIT>' ) <EOL> ax . set_ylabel ( '<STR_LIT>' ) <EOL> ax . set_title ( '<STR_LIT>' ) <EOL> plt . show ( ) </s>
