Add calibration from video script

modules/odm_slam/src/calibrate_video.py

@@ -0,0 +1,134 @@
+#!/usr/bin/env python
+
+import argparse
+import sys
+
+import numpy as np
+import cv2
+
+
+class Calibrator:
+    """Camera calibration using a chessboard pattern."""
+
+    def __init__(self, pattern_width, pattern_height, motion_threshold=0.05):
+        """Init the calibrator.
+
+        The parameter motion_threshold determines the minimal motion required
+        to add a new frame to the calibration data, as a ratio of image width.
+        """
+        self.pattern_size = (pattern_width, pattern_height)
+        self.motion_threshold = motion_threshold
+        self.pattern_points = np.array([
+            (i, j, 0.0)
+            for j in range(pattern_height)
+            for i in range(pattern_width)
+        ], dtype=np.float32)
+        self.object_points = []
+        self.image_points = []
+
+    def process_image(self, image, window_name):
+        """Find corners of an image and store them internally."""
+        if len(image.shape) == 3:
+            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+        else:
+            gray = image
+
+        h, w = gray.shape
+        self.image_size = (w, h)
+
+        found, corners = cv2.findChessboardCorners(gray, self.pattern_size)
+
+        if found:
+            term = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_COUNT, 30, 0.1)
+            cv2.cornerSubPix(gray, corners, (5, 5), (-1, -1), term)
+            self._add_points(corners.reshape(-1, 2))
+
+        if window_name:
+            cv2.drawChessboardCorners(image, self.pattern_size, corners, found)
+            cv2.imshow(window_name, image)
+
+        return found
+
+    def calibrate(self):
+        """Run calibration using points extracted by process_image."""
+        rms, camera_matrix, dist_coefs, rvecs, tvecs = cv2.calibrateCamera(
+            self.object_points, self.image_points, self.image_size, None, None)
+        return rms, camera_matrix, dist_coefs
+
+    def _add_points(self, image_points):
+        if self.image_points:
+            delta = np.fabs(image_points - self.image_points[-1]).max()
+            should_add = (delta > self.image_size[0] * self.motion_threshold)
+        else:
+            should_add = True
+
+        if should_add:
+            self.image_points.append(image_points)
+            self.object_points.append(self.pattern_points)
+
+
+def video_frames(filename):
+    """Yield frames in a video."""
+    cap = cv2.VideoCapture(args.video)
+    while True:
+        ret, frame = cap.read()
+        if ret:
+            yield frame
+        else:
+            break
+    cap.release()
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(
+        description="Camera calibration from video of a chessboard.")
+    parser.add_argument(
+        'video',
+        help="video of the checkerboard")
+    parser.add_argument(
+        '--output',
+        default='calibration',
+        help="base name for the output files")
+    parser.add_argument(
+        '--size',
+        default='8x6',
+        help="size of the chessboard")
+    parser.add_argument(
+        '--visual',
+        action='store_true',
+        help="display images while calibrating")
+    return parser.parse_args()
+
+
+if __name__ == '__main__':
+    args = parse_arguments()
+
+    pattern_size = [int(i) for i in args.size.split('x')]
+    calibrator = Calibrator(pattern_size[0], pattern_size[1])
+
+    window_name = None
+    if args.visual:
+        window_name = 'Chessboard detection'
+        cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
+
+    print "kept\tcurrent\tchessboard found"
+
+    for i, frame in enumerate(video_frames(args.video)):
+        found = calibrator.process_image(frame, window_name)
+
+        print "{}\t{}\t{}             \r".format(
+            len(calibrator.image_points), i, found),
+        sys.stdout.flush()
+
+        if args.visual:
+            if cv2.waitKey(1) & 0xFF == ord('q'):
+                break
+    print
+
+    cv2.destroyAllWindows()
+
+    rms, camera_matrix, dist_coefs = calibrator.calibrate()
+
+    print "RMS:", rms
+    print "camera matrix:\n", camera_matrix
+    print "distortion coefficients: ", dist_coefs.ravel()