cleanup calibrator existing code (#18)

src/Calibrator.py

@@ -23,14 +23,12 @@ 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.
@@ -40,24 +38,28 @@ class Calibrator(object):
         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 = [[c[1].pose.position.x, c[1].pose.position.y, c[1].pose.position.z]
+                for c in configurations]
 
-        pos_RE_Mx3 = np.array(pos_RE)
-        quat_RE_Mx4 = np.array(quat_RE)
+        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 (pos_RE_Mx3,quat_RE_Mx4)
+        return (np.array(pos_RE), np.array(quat_RE))
 
 
     def calibrate(self,pos_W,pos_RE,quat_RE,quat_ET,
@@ -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)