cleanup calibrator existing code (#18)

src/Calibrator.py

@@ -23,41 +23,43 @@ from Controller import getR_ET
 
 class Calibrator(object):
 
-    def __init__(self,
-        timeScale=0.25):
-
+    def __init__(self, timeScale=0.25):
         self.timeScale = timeScale
-
         self.controller = None
 
-    def measureConfigurations(self,pos_W):
+
+    def measureConfigurations(self, pos_W, prior_calibs):
         """
         Calibrate the arm for the robot's pose relative to the world
         and the position of the tool tip relative to the end effector.
-        
+
         pos_W: Nx3 sequence of positions to visit for calibration
         """
         if self.controller is None:
             if DEBUG:
                 print("Creating KeyboardController")
-            self.controller = KeyboardController(waitForExtruder=False,timeScale=self.timeScale,disableExtruder=True)
-
-        nPos = pos_W.shape[0]
-
-        ## Form prompts to the user to satisfy
-        prompts = ['Move robot tip to {}'.format(np.array2string(pos_W[iPos,])) for iPos in range(nPos)]
-
-        ## Gather configurations
-        # List of (list of joint angle, ROS pose message)
+            self.controller = KeyboardController(
+                    waitForExtruder=False,
+                    timeScale=self.timeScale,
+                    disableExtruder=True
+                    )
+
+        # form prompts to the user to satisfy
+        prompts = ['Move robot tip to {}'.format(np.array2string(pos_W[iPos,]))
+                for iPos in range(pos_W.shape[0])]
+
+        # gather configurations
+        # [ ( [joint angles] , ROS pose message ) , ... ]
         configurations = self.controller.prompt_configurations(prompts)
-    
-        pos_RE = [[c[1].pose.position.x,c[1].pose.position.y,c[1].pose.position.z] for c in configurations]
-        quat_RE = [[c[1].pose.orientation.x,c[1].pose.orientation.y,c[1].pose.orientation.z,c[1].pose.orientation.w] for c in configurations]
 
-        pos_RE_Mx3 = np.array(pos_RE)
-        quat_RE_Mx4 = np.array(quat_RE)
+        pos_RE = [[c[1].pose.position.x, c[1].pose.position.y, c[1].pose.position.z]
+                for c in configurations]
 
-        return (pos_RE_Mx3,quat_RE_Mx4)
+        quat_RE = [[c[1].pose.orientation.x, c[1].pose.orientation.y,
+                    c[1].pose.orientation.z, c[1].pose.orientation.w]
+                for c in configurations]
+
+        return (np.array(pos_RE), np.array(quat_RE))
 
 
     def calibrate(self,pos_W,pos_RE,quat_RE,quat_ET,
@@ -67,7 +69,7 @@ class Calibrator(object):
         """
         Calibrate the arm for the robot's pose relative to the world
         and the position of the tool tip relative to the end effector.
-        
+
         pos_W: Mx3 set of positions visited for calibration
         pos_RE: Mx3 set of positions of the end effector relative to the robot's base
         quat_RE: Mx4 set of orientations of the end effector relative to the robot's base
@@ -94,7 +96,7 @@ def calibrate_robotPose_toolPos(pos_W,pos_RE,quat_RE,quat_ET,
     Calibration for:
     - the robot's pose from the world
     - the tool's position from the end effector
-    
+
     pos_W: Nx3 [x,y,z] positions
     pos_RE: Nx3 [x,y,z] position of end effector relative to robot base
     quat_RE: Nx4 [x,y,z,w] orientation of end effector relative to robot base
@@ -102,7 +104,7 @@ def calibrate_robotPose_toolPos(pos_W,pos_RE,quat_RE,quat_ET,
     initPos_WR: 3x (default: [0,0,0])
     initQuat_WR: 4x (default: [0,0,0,1])
     initPos_ET: 3x (default: [0,0,0])
-    
+
     Returns: (pos_WR,quat_WR,pos_ET)
     pos_WR: 3x
     quat_WR: 4x
@@ -149,7 +151,7 @@ def err_robotPose_toolTrans(x,pos_W,pos_RE,quat_RE,quat_ET,pos_T):
     quat_RE: Nx4 [x,y,z,w]
     pos_ET: Nx3 [x,y,z]
     pos_T: Nx3 [x,y,z]
-    
+
     Returns N3x dimension array
     """
     nPose = pos_RE.shape[0]
@@ -182,7 +184,7 @@ def getPos_WT(pos_WR,quat_WR,pos_RE,quat_RE,pos_ET,quat_ET,pos_T):
     pos_ET: Nx3 [x,y,z]
     quat_ET: Nx4 [x,y,z,w]
     pos_T: Nx3 [x,y,z]
-    
+
     Returns: Nx3
     """
     nPose = pos_WR.shape[0]
@@ -206,7 +208,7 @@ def pose2transform(pos,quat):
     """
     pos: Nx3 [x,y,z]
     quat: Nx4 [x,y,z,w]
-    
+
     Returns: Nx4x4
     """
     nPose = pos.shape[0]
@@ -245,30 +247,10 @@ def printTransforms(T):
     print(','.join(strings))
 
 
-if __name__=="__main__":
-    parser = argparse.ArgumentParser(description='Calibration for UR5e arm')
-    parser.add_argument('calibFile_in',
-        metavar='calibFile_read',
-        help='file storing calibration targets and calibration parameters')
-    parser.add_argument('-o','--out',
-        metavar='calibFile_write',
-        dest='calibFile_out',
-        default='',
-        help='file storing calibration targets and calibration parameters')
-
-    args = parser.parse_args()
-
-    calibFile_in = args.calibFile_in
-    if len(args.calibFile_out) > 0:
-        calibFile_out = args.calibFile_out
-    else:
-        # Overwrite the file to read from
-        calibFile_out = args.calibFile_in
-
-    # Read calibration file
-    with open(calibFile_in,'r') as file:
-        calibData = yaml.load(file,
-            Loader=yaml.FullLoader)
+def main(calibFile_in, calibFile_out, calibData):
+    """
+    main Calibration process
+    """
 
     if DEBUG:
         print("Creating Calibrator")
@@ -308,10 +290,19 @@ if __name__=="__main__":
         measurements['pos_RE'] = pos_RE
         measurements['quat_RE'] = quat_RE
 
+    ## Solve nonlinear equations to obtain calibration parameters
+    # If there is any calibration parameters from before
+    # also used for faster calibration
+    if 'calibration' in calibData:
+        calibration = calibData['calibration']
+    else:
+        calibration = {}
+        calibData['calibration'] = calibration
+
     # If the number of measurements does not match the number of targets
     if len(pos_RE) != nTarget or len(quat_RE) != nTarget:
         # Take measurements of the targets
-        (pos_RE_Mx3,quat_RE_Mx4) = cal.measureConfigurations(pos_W_Mx3)
+        (pos_RE_Mx3,quat_RE_Mx4) = cal.measureConfigurations(pos_W_Mx3, calibration)
         pos_RE.extend(pos_RE_Mx3.tolist())
         quat_RE.extend(quat_RE_Mx4.tolist())
 
@@ -323,13 +314,6 @@ if __name__=="__main__":
         pos_RE_Mx3 = np.array(pos_RE)
         quat_RE_Mx4 = np.array(quat_RE)
 
-    ## Solve nonlinear equations to obtain calibration parameters
-    # If there is any calibration parameters from before
-    if 'calibration' in calibData:
-        calibration = calibData['calibration']
-    else:
-        calibration = {}
-        calibData['calibration'] = calibration
 
     if 'pos_WR' in calibration and 'quat_WR' in calibration and 'pos_ET' in calibration:
         # Use as initial estimate for solver
@@ -391,3 +375,31 @@ if __name__=="__main__":
     with open(calibFile_out,'w') as file:
         yaml.dump(calibData,file,
             default_flow_style=None)
+
+
+if __name__=="__main__":
+    parser = argparse.ArgumentParser(description='Calibration for UR5e arm')
+    parser.add_argument('calibFile_in',
+        metavar='calibFile_read',
+        help='file storing calibration targets and calibration parameters')
+    parser.add_argument('-o','--out',
+        metavar='calibFile_write',
+        dest='calibFile_out',
+        default='',
+        help='file storing calibration targets and calibration parameters')
+
+    args = parser.parse_args()
+
+    calibFile_in = args.calibFile_in
+    if len(args.calibFile_out) > 0:
+        calibFile_out = args.calibFile_out
+    else:
+        # Overwrite the file to read from
+        calibFile_out = args.calibFile_in
+
+    # Read calibration file
+    with open(calibFile_in,'r') as file:
+        calibData = yaml.load(file,
+            Loader=yaml.FullLoader)
+
+    main(calibFile_in, calibFile_out, calibData)