#!/usr/bin/env python3

import cv2

def main():
    imgSize = (640, 480)

    video = cv2.VideoCapture(0)
    video.set(cv2.CAP_PROP_FRAME_WIDTH, imgSize[0])
    video.set(cv2.CAP_PROP_FRAME_HEIGHT, imgSize[1])

    d = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_50)
    board = cv2.aruco.CharucoBoard.create(3, 3, 0.02, 0.012, d)

    allCorners = []
    allIds = []

    while True:
        (ok, frame) = video.read()
        if not ok:
            raise Error('video.read() yielded False')

        (corners, ids, rejectedImagePoints) = cv2.aruco.detectMarkers(frame, d)

        if ids is not None:
            cv2.aruco.drawDetectedMarkers(frame, corners, ids)
            (response, boardCorners, boardIds) = cv2.aruco.interpolateCornersCharuco(corners, ids, frame, board)
            if response >= 4:
                allCorners.append(boardCorners)
                allIds.append(boardIds)
                print(len(allCorners), len(allIds))

        cv2.imshow('calibrate', frame)
        if cv2.waitKey(1) != -1:
            break

    (calibration, cameraMatrix, distCoeffs, _rvecs, _tvecs) = \
        cv2.aruco.calibrateCameraCharuco(allCorners, allIds, board, imgSize, None, None)
    with open('calibration.json', 'wt') as f:
        import json
        json.dump({
            'cameraMatrix': cameraMatrix.tolist(),
            'distCoeffs': distCoeffs.tolist(),
        }, f, indent=4)

if __name__ == '__main__':
    main()