import argparse
import cv2
import numpy as np
import os

parser = argparse.ArgumentParser(description="Camera Intrinsic Calibration")
parser.add_argument('-input', '--INPUT_TYPE', default='camera', type=str, help='Input Source: camera/video/image')
parser.add_argument('-type', '--CAMERA_TYPE', default='fisheye', type=str, help='Camera Type: fisheye/normal')
parser.add_argument('-id', '--CAMERA_ID', default=1, type=int, help='Camera ID')
parser.add_argument('-path', '--INPUT_PATH', default='./data/', type=str, help='Input Video/Image Path')
parser.add_argument('-video', '--VIDEO_FILE', default='video.mp4', type=str, help='Input Video File Name (eg.: video.mp4)')
parser.add_argument('-image', '--IMAGE_FILE', default='img_raw', type=str, help='Input Image File Name Prefix (eg.: img_raw)')
parser.add_argument('-mode', '--SELECT_MODE', default='auto', type=str, help='Image Select Mode: auto/manual')
parser.add_argument('-fw','--FRAME_WIDTH', default=1280, type=int, help='Camera Frame Width')
parser.add_argument('-fh','--FRAME_HEIGHT', default=1024, type=int, help='Camera Frame Height')
parser.add_argument('-bw','--BORAD_WIDTH', default=7, type=int, help='Chess Board Width (corners number)')
parser.add_argument('-bh','--BORAD_HEIGHT', default=6, type=int, help='Chess Board Height (corners number)')
parser.add_argument('-size','--SQUARE_SIZE', default=10, type=int, help='Chess Board Square Size (mm)')
parser.add_argument('-num','--CALIB_NUMBER', default=5, type=int, help='Least Required Calibration Frame Number')
parser.add_argument('-delay','--FRAME_DELAY', default=12, type=int, help='Capture Image Time Interval (frame number)')
parser.add_argument('-subpix','--SUBPIX_REGION', default=5, type=int, help='Corners Subpix Optimization Region')
parser.add_argument('-fps','--CAMERA_FPS', default=20, type=int, help='Camera Frame per Second(FPS)')
parser.add_argument('-fs', '--FOCAL_SCALE', default=0.5, type=float, help='Camera Undistort Focal Scale')
parser.add_argument('-ss', '--SIZE_SCALE', default=1, type=float, help='Camera Undistort Size Scale')
parser.add_argument('-store','--STORE_FLAG', default=False, type=bool, help='Store Captured Images (Ture/False)')
parser.add_argument('-store_path', '--STORE_PATH', default='./data/', type=str, help='Path to Store Captured Images')
parser.add_argument('-crop','--CROP_FLAG', default=False, type=bool, help='Crop Input Video/Image to (fw,fh) (Ture/False)')
parser.add_argument('-resize','--RESIZE_FLAG', default=False, type=bool, help='Resize Input Video/Image to (fw,fh) (Ture/False)')
args = parser.parse_args([])                 # Jupyter Notebook中直接运行时要加[], py文件则去掉

# args.INPUT_TYPE = 'image'                  # 输入形式 相机/视频/图像
# args.CAMERA_TYPE = 'normal'                # 相机类型 鱼眼/普通
# args.CAMERA_ID = 1                         # 相机编号
# args.INPUT_PATH = './data/'                # 图片、视频输入路径
# args.VIDEO_FILE = 'video.mp4'              # 输入视频文件名(含扩展名)
# args.IMAGE_FILE = 'raw'                    # 输入图像文件名前缀
# args.SELECT_MODE = 'manual'                # 选择自动/手动模式
# args.FRAME_WIDTH = 1280                    # 相机分辨率 帧宽度
# args.FRAME_HEIGHT = 720                    # 相机分辨率 帧高度
# args.BORAD_WIDTH = 7                       # 棋盘宽度 【内角点数】
# args.BORAD_HEIGHT = 6                      # 棋盘高度 【内角点数】
# args.SQUARE_SIZE = 10                      # 棋盘格边长 mm
# args.CALIB_NUMBER = 10                     # 初始化最小标定图片采样数量
# args.FRAME_DELAY = 15                      # 间隔多少帧数采样
# args.SUBPIX_REGION = 3                     # 角点坐标亚像素优化时的搜索区域大小(根据图像分辨率调整)
# args.STORE_FLAG = True                     # 是否保存抓取的图像
# args.STORE_PATH = './data/'                # 保存抓取的图像的路径
# args.CROP_FLAG = True                      # 是否将输入视频/图像尺寸裁剪至FRAME_WIDTH和FRAME_HEIGHT的设定值
# args.RESIZE_FLAG = True                    # 是否将输入视频/图像尺寸缩放至FRAME_WIDTH和FRAME_HEIGHT的设定值（图像缩放会改变相机焦距）

# 外部调用修改参数
def getInCalibArgs():
    return args

def editInCalibArgs(new_args):
    global args
    args = new_args

class CalibData:                             # 标定数据类
    def __init__(self):
        self.type = None                     # 自定义数据类型
        self.camera_mat = None               # 相机内参
        self.dist_coeff = None               # 畸变参数
        self.rvecs = None                    # 旋转向量
        self.tvecs = None                    # 平移向量
        self.map1 = None                     # 映射矩阵1
        self.map2 = None                     # 映射矩阵2
        self.reproj_err = None               # 重投影误差
        self.ok = False                      # 数据采集完成标志

# cv2.fisheye.calibrate ( objectPoints,      # 角点在棋盘中的空间坐标向量
#                         imagePoints,       # 角点在图像中的坐标向量
#                         image_size,        # 图片大小
#                         K,                 # 相机内参矩阵
#                         D,                 # 畸变参数向量
#                         rvecs,             # 旋转向量
#                         tvecs,             # 平移向量
#                         flags,             # 操作标志
#                         criteria           # 迭代优化算法的停止标准
#                         )

# flags:
#     cv2.fisheye.CALIB_USE_INTRINSIC_GUESS     # 当相机内参矩阵包含有效的fx，fy，cx，cy初始值时，这些值会进一步进行优化
#                                               # 否则，（cx，cy）初始化设置为图像中心（使用imageSize），并且以最小二乘法计算焦距。
#     cv2.fisheye.CALIB_RECOMPUTE_EXTRINSIC     # 在每次内部参数优化迭代之后，将重新计算外部参数。
#     cv2.fisheye.CALIB_CHECK_COND              # 检查条件编号的有效性
#     cv2.fisheye.CALIB_FIX_SKEW                # 偏斜系数（alpha）设置为零，并保持为零
#     cv2.fisheye.CALIB_FIX_K1 (K1-K4)          # 选定的畸变系数设置为零，并保持为零  CALIB_FIX_INTRINSIC则全为零

# criteria:
#     TermCriteria (int type, int maxCount, double epsilon)      # 类型、最大计数次数、最小精度
#     Criteria type, can be one of: COUNT, EPS or COUNT + EPS    # 角点优化最大迭代次数或角点优化位置移动量小于epsilon值

# cv2.checkRange 检查元素非空及无异常值

class Fisheye:           # 鱼眼相机
    def __init__(self):
        self.data = CalibData()
        self.inited = False
        self.BOARD = np.array([ [(j * args.SQUARE_SIZE, i * args.SQUARE_SIZE, 0.)]
                               for i in range(args.BORAD_HEIGHT) 
                               for j in range(args.BORAD_WIDTH) ],dtype=np.float32)     # 棋盘角点二维坐标（乘上尺寸）
        
    # 更新标定数据，分为初始化和精调
    def update(self, corners, frame_size):
        board = [self.BOARD] * len(corners)
        if not self.inited:
            self._update_init(board, corners, frame_size)
            self.inited = True
        else:
            self._update_refine(board, corners, frame_size)
        self._calc_reproj_err(corners)
        self._get_undistort_maps()
    
    # 得到一定数量标定样本时进行初始标定
    def _update_init(self, board, corners, frame_size):
        data = self.data
        data.type = "FISHEYE"
        data.camera_mat = np.eye(3, 3)
        data.dist_coeff = np.zeros((4, 1))
        data.ok, data.camera_mat, data.dist_coeff, data.rvecs, data.tvecs = cv2.fisheye.calibrate(
            board, corners, frame_size, data.camera_mat, data.dist_coeff,
            flags=cv2.fisheye.CALIB_FIX_SKEW|cv2.fisheye.CALIB_RECOMPUTE_EXTRINSIC,
            criteria=(cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_COUNT, 30, 1e-6)) 
        data.ok = data.ok and cv2.checkRange(data.camera_mat) and cv2.checkRange(data.dist_coeff)

    # 精调时启用CALIB_USE_INTRINSIC_GUESS
    def _update_refine(self, board, corners, frame_size):
        data = self.data
        data.ok, data.camera_mat, data.dist_coeff, data.rvecs, data.tvecs = cv2.fisheye.calibrate(
            board, corners, frame_size, data.camera_mat, data.dist_coeff,
            flags=cv2.fisheye.CALIB_FIX_SKEW|cv2.fisheye.CALIB_RECOMPUTE_EXTRINSIC|cv2.CALIB_USE_INTRINSIC_GUESS,
            criteria=(cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_COUNT, 10, 1e-6))
        data.ok = data.ok and cv2.checkRange(data.camera_mat) and cv2.checkRange(data.dist_coeff)

    # 计算重投影误差，单位为像素
    def _calc_reproj_err(self, corners):
        if not self.inited: return
        data = self.data
        data.reproj_err = []
        for i in range(len(corners)):
            corners_reproj, _ = cv2.fisheye.projectPoints(self.BOARD, data.rvecs[i], data.tvecs[i], data.camera_mat, data.dist_coeff)
            err = cv2.norm(corners_reproj, corners[i], cv2.NORM_L2) / len(corners_reproj)
            data.reproj_err.append(err)
            
    # 计算去畸变的新的相机内参，可以改变焦距和画幅
    def _get_camera_mat_dst(self, camera_mat):
        camera_mat_dst = camera_mat.copy()
        camera_mat_dst[0][0] *= args.FOCAL_SCALE
        camera_mat_dst[1][1] *= args.FOCAL_SCALE
        camera_mat_dst[0][2] = args.FRAME_WIDTH / 2 * args.SIZE_SCALE
        camera_mat_dst[1][2] = args.FRAME_HEIGHT / 2 * args.SIZE_SCALE
        return camera_mat_dst
    
    # 计算去畸变的映射矩阵
    def _get_undistort_maps(self):
        data = self.data
        camera_mat_dst = self._get_camera_mat_dst(data.camera_mat)
        data.map1, data.map2 = cv2.fisheye.initUndistortRectifyMap(
                                 data.camera_mat, data.dist_coeff, np.eye(3, 3), camera_mat_dst, 
                                 (int(args.FRAME_WIDTH * args.SIZE_SCALE), int(args.FRAME_HEIGHT * args.SIZE_SCALE)), cv2.CV_16SC2)

# cv2.calibrateCamera (objectPoints,         # 角点在棋盘中的空间坐标向量        
#                      imagePoints,          # 角点在图像中的坐标向量
#                      image_size,           # 图片大小
#                      K,                    # 相机内参矩阵
#                      D,                    # 畸变参数向量
#                      rvecs,                # 旋转向量
#                      tvecs,                # 平移向量
#                      flags,                # 操作标志
#                      criteria              # 迭代优化算法的停止标准
#                     )

# flags:
#     cv2.CALIB_USE_INTRINSIC_GUESS          # 当相机内参矩阵包含有效的fx，fy，cx，cy初始值时，这些值会进一步进行优化
#                                            # 否则，（cx，cy）初始化设置为图像中心（使用imageSize），并且以最小二乘法计算焦距
#     cv2.CALIB_FIX_PRINCIPAL_POINT          # 固定光轴点（当设置CALIB_USE_INTRINSIC_GUESS时可以使用）
#     cv2.CALIB_FIX_ASPECT_RATIO             # 固定fx/fy的值，函数仅将fy视为自由参数
#     cv2.CALIB_ZERO_TANGENT_DIST            # 切向畸变系数（p1，p2） 设置为零并保持为零
#     cv2.CALIB_FIX_FOCAL_LENGTH             # 如果设置了CALIB_USE_INTRINSIC_GUESS，则在全局优化过程中不会更改焦距
#     cv2.CALIB_FIX_K1 (K1-K6)               # 固定相应的径向畸变系数为0或给定的初始值
#     cv2.CALIB_RATIONAL_MODEL               # 理想模型：启用系数k4，k5和k6。此时返回8个或更多的系数
#     cv2.CALIB_THIN_PRISM_MODEL             # 薄棱镜模型：启用系数s1，s2，s3和s4。此时返回12个或更多的系数
#     cv2.CALIB_FIX_S1_S2_S3_S4              # 固定薄棱镜畸变系数为0或给定的初始值
#     cv2.CALIB_TILTED_MODEL                 # 倾斜模型：启用系数tauX和tauY。此时返回14个系数
#     cv2.CALIB_FIX_TAUX_TAUY                # 固定倾斜传感器模型的系数为0或给定的初始值

class Normal:           # 平面相机
    def __init__(self):
        self.data = CalibData()
        self.inited = False
        self.BOARD = np.array([ [(j * args.SQUARE_SIZE, i * args.SQUARE_SIZE, 0.)]
                               for i in range(args.BORAD_HEIGHT) 
                               for j in range(args.BORAD_WIDTH) ],dtype=np.float32)
        
    def update(self, corners, frame_size):
        board = [self.BOARD] * len(corners)
        if not self.inited:
            self._update_init(board, corners, frame_size)
            self.inited = True
        else:
            self._update_refine(board, corners, frame_size)
        self._calc_reproj_err(corners)
        self._get_undistort_maps()
        
    def _update_init(self, board, corners, frame_size):
        data = self.data
        data.type = "NORMAL"
        data.camera_mat = np.eye(3, 3)
        data.dist_coeff = np.zeros((5, 1))     # 畸变向量的尺寸根据使用模型修改
        data.ok, data.camera_mat, data.dist_coeff, data.rvecs, data.tvecs = cv2.calibrateCamera(
            board, corners, frame_size, data.camera_mat, data.dist_coeff, 
            criteria=(cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_COUNT, 30, 1e-6))
        data.ok = data.ok and cv2.checkRange(data.camera_mat) and cv2.checkRange(data.dist_coeff)
        
    def _update_refine(self, board, corners, frame_size):
        data = self.data
        data.ok, data.camera_mat, data.dist_coeff, data.rvecs, data.tvecs = cv2.calibrateCamera(
            board, corners, frame_size, data.camera_mat, data.dist_coeff,  
            flags = cv2.CALIB_USE_INTRINSIC_GUESS,
            criteria=(cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_COUNT, 10, 1e-6))
        data.ok = data.ok and cv2.checkRange(data.camera_mat) and cv2.checkRange(data.dist_coeff)
        
    def _calc_reproj_err(self, corners):
        if not self.inited: return
        data = self.data
        data.reproj_err = []
        for i in range(len(corners)):
            corners_reproj, _ = cv2.projectPoints(self.BOARD, data.rvecs[i], data.tvecs[i], data.camera_mat, data.dist_coeff)
            err = cv2.norm(corners_reproj, corners[i], cv2.NORM_L2) / len(corners_reproj)
            data.reproj_err.append(err)
            
    def _get_camera_mat_dst(self, camera_mat):
        camera_mat_dst = camera_mat.copy()
        camera_mat_dst[0][0] *= args.FOCAL_SCALE
        camera_mat_dst[1][1] *= args.FOCAL_SCALE
        camera_mat_dst[0][2] = args.FRAME_WIDTH / 2 * args.SIZE_SCALE
        camera_mat_dst[1][2] = args.FRAME_HEIGHT / 2 * args.SIZE_SCALE
        return camera_mat_dst
    
    def _get_undistort_maps(self):
        data = self.data
        camera_mat_dst = self._get_camera_mat_dst(data.camera_mat)
        data.map1, data.map2 = cv2.initUndistortRectifyMap(
                                 data.camera_mat, data.dist_coeff, np.eye(3, 3), camera_mat_dst, 
                                 (int(args.FRAME_WIDTH * args.SIZE_SCALE), int(args.FRAME_HEIGHT * args.SIZE_SCALE)), cv2.CV_16SC2)

# cv2.findChessboardCorners ( image,         # 棋盘图像
#                             patternSize,   # 棋盘格行和列的【内角点】数量
#                             corners,       # 输出数组
#                             flags          # 操作标志
#                             )
# flags:
#     CV_CALIB_CB_ADAPTIVE_THRESH            # 使用自适应阈值处理将图像转换为黑白图像
#     CV_CALIB_CB_NORMALIZE_IMAGE            # 对图像进行归一化。
#     CV_CALIB_CB_FILTER_QUADS               # 过滤在轮廓检索阶段提取的假四边形。
#     CALIB_CB_FAST_CHECK                    # 对查找棋盘角的图像进行快速检查

# cv2.cornerSubPix (image,                        # 棋盘图像
#                   corners,                      # 棋盘角点
#                   winSize,                      # 搜索窗口边长的一半
#                   zeroZone,                     # 搜索区域死区大小的一半, (-1,-1)代表无
#                   criteria                      # 迭代停止标准
#                  )

# cv2.fisheye.initUndistortRectifyMap (K,         # 相机内参矩阵
#                                      D,         # 畸变向量
#                                      R,         # 旋转矩阵
#                                      P,         # 新的相机矩阵
#                                      size,      # 输出图像大小
#                                      m1type,    # 映射矩阵类型
#                                      map1,      # 输出映射矩阵1
#                                      map2       # 输出映射矩阵2
#                                     )

class InCalibrator:                  # 内参标定器
    def __init__(self, camera):
        if camera == 'fisheye':
            self.camera = Fisheye()  # 鱼眼相机类
        elif camera == 'normal':
            self.camera = Normal()   # 普通相机类
        else:
            raise Exception("camera should be fisheye/normal")
        self.corners = []
    
    # 获取args参数，供外部调用修改参数
    @staticmethod
    def get_args():
        return args
    
    # 获取棋盘格角点坐标
    def get_corners(self, img):
        ok, corners = cv2.findChessboardCorners(img, (args.BORAD_WIDTH, args.BORAD_HEIGHT),
                      flags = cv2.CALIB_CB_ADAPTIVE_THRESH|cv2.CALIB_CB_NORMALIZE_IMAGE|cv2.CALIB_CB_FAST_CHECK)
        if ok: 
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            # 角点坐标亚像素优化
            corners = cv2.cornerSubPix(gray, corners, (args.SUBPIX_REGION, args.SUBPIX_REGION), (-1, -1),
                                       (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.01))
        return ok, corners
    
    # 在图上绘制棋盘格角点
    def draw_corners(self, img):
        ok, corners = self.get_corners(img)
        cv2.drawChessboardCorners(img, (args.BORAD_WIDTH, args.BORAD_HEIGHT), corners, ok)
        return img
    
    # 图像去畸变
    def undistort(self, img):
        data = self.camera.data
        return cv2.remap(img, data.map1, data.map2, cv2.INTER_LINEAR)
    
    # 使用现有角点坐标标定
    def calibrate(self, img):
        if len(self.corners) >= args.CALIB_NUMBER:
            self.camera.update(self.corners, img.shape[1::-1])  # 更新标定数据
        return self.camera.data
    
    def __call__(self, raw_frame):
        ok, corners = self.get_corners(raw_frame)
        result = self.camera.data
        if ok:
            self.corners.append(corners)          # 加入新的角点坐标
            result = self.calibrate(raw_frame)    # 得到标定结果
        return result

# 居中裁剪
def centerCrop(img,width,height):
    if img.shape[1] < width or img.shape[0] < height:
        raise Exception("CROP size should be smaller than original size")
    img = img[round((img.shape[0]-height)/2) : round((img.shape[0]-height)/2)+height,
              round((img.shape[1]-width)/2) : round((img.shape[1]-width)/2)+width ]  
    return img

# 筛选图片文件
def get_images(PATH, NAME):
    filePath = [os.path.join(PATH, x) for x in os.listdir(PATH) 
                if any(x.endswith(extension) for extension in ['.png', '.jpg', '.jpeg', '.PNG', '.JPG', '.JPEG'])
               ]                                                            # 得到给定路径下所有图片文件
    filenames = [filename for filename in filePath if NAME in filename]     # 再筛选出包含给定名字的图片
    if len(filenames) == 0:
        raise Exception("from {} read images failed".format(PATH))
    return filenames

class CalibMode():
    def __init__(self, calibrator, input_type, mode):
        self.calibrator = calibrator
        self.input_type = input_type
        self.mode = mode
    
    # 图片预处理
    def imgPreprocess(self, img):
        if args.CROP_FLAG:          # 裁剪图片尺寸
            img = centerCrop(img, args.FRAME_WIDTH, args.FRAME_HEIGHT)
        elif args.RESIZE_FLAG:      # 缩放图片尺寸
            img = cv2.resize(img, (args.FRAME_WIDTH, args.FRAME_HEIGHT))
        return img
    
    # 设置相机
    def setCamera(self, cap):
        cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter.fourcc('M','J','P','G'))    # 设置编码格式为MJPG
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, args.FRAME_WIDTH)                      # 设置相机分辨率
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, args.FRAME_HEIGHT)
        cap.set(cv2.CAP_PROP_FPS, args.CAMERA_FPS)                               # 设置相机帧率
        return cap
    
    # 运行标定程序
    def runCalib(self, raw_frame, display_raw=True, display_undist=True):
        calibrator = self.calibrator
        raw_frame = self.imgPreprocess(raw_frame)
        result = calibrator(raw_frame)
        raw_frame = calibrator.draw_corners(raw_frame)
        if display_raw:
            cv2.namedWindow("raw_frame", flags = cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)
            cv2.imshow("raw_frame", raw_frame)
        if len(calibrator.corners) > args.CALIB_NUMBER and display_undist: 
            undist_frame = calibrator.undistort(raw_frame)
            cv2.namedWindow("undist_frame", flags = cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)
            cv2.imshow("undist_frame", undist_frame)   
        cv2.waitKey(1)
        return result
    
    # 图片输入自动标定
    def imageAutoMode(self):
        calibrator = self.calibrator
        filenames = get_images(args.INPUT_PATH, args.IMAGE_FILE)
        for filename in filenames:
            print(filename)
            raw_frame = cv2.imread(filename)
            result = self.runCalib(raw_frame)
            key = cv2.waitKey(1)
            if key == 27: break
        cv2.destroyAllWindows() 
        return result
    
    # 图片输入手动挑选 按空格键确认 其他键丢弃该图片
    def imageManualMode(self):
        filenames = get_images(args.INPUT_PATH, args.IMAGE_FILE)
        for filename in filenames:
            print(filename)
            raw_frame = cv2.imread(filename)
            raw_frame = self.imgPreprocess(raw_frame)
            img = raw_frame.copy()
            img = self.calibrator.draw_corners(img)
            display = "raw_frame: press SPACE to SELECT, other key to SKIP, press ESC to QUIT"
            cv2.namedWindow(display, flags = cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)
            cv2.imshow(display, img)
            key = cv2.waitKey(0)
            if key == 32:
                result = self.runCalib(raw_frame, display_raw = False)
            if key == 27: break
        cv2.destroyAllWindows() 
        return result
    
    # 视频输入自动标定
    def videoAutoMode(self):
        cap = cv2.VideoCapture(args.INPUT_PATH + args.VIDEO_FILE)
        if not cap.isOpened(): 
            raise Exception("from {} read video failed".format(args.INPUT_PATH + args.VIDEO_FILE))
        frame_id = 0
        while True:
            key = cv2.waitKey(1)
            ok, raw_frame = cap.read()
            raw_frame = self.imgPreprocess(raw_frame)
            if frame_id % args.FRAME_DELAY == 0:
                if args.STORE_FLAG:  # 存储该帧图像
                    cv2.imwrite(args.STORE_PATH + 'img_raw{}.jpg'.format(len(self.calibrator.corners)), raw_frame)
                result = self.runCalib(raw_frame) 
                print(len(self.calibrator.corners))
            frame_id += 1 
            key = cv2.waitKey(1)
            if key == 27: break
        cap.release()
        cv2.destroyAllWindows() 
        return result
    
    # 视频输入手动挑选 按空格键采集图片
    def videoManualMode(self):
        cap = cv2.VideoCapture(args.INPUT_PATH + args.VIDEO_FILE)
        if not cap.isOpened(): 
            raise Exception("from {} read video failed".format(args.INPUT_PATH + args.VIDEO_FILE))
        while True:
            key = cv2.waitKey(1)
            ok, raw_frame = cap.read()
            raw_frame = self.imgPreprocess(raw_frame)
            display = "raw_frame: press SPACE to capture image"
            cv2.namedWindow(display, flags = cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)
            cv2.imshow(display, raw_frame)
            if key == 32:
                if args.STORE_FLAG:  # 存储该帧图像
                    cv2.imwrite(args.STORE_PATH + 'img_raw{}.jpg'.format(len(self.calibrator.corners)), raw_frame)
                result = self.runCalib(raw_frame) 
                print(len(self.calibrator.corners))
            if key == 27: break
        cap.release()
        cv2.destroyAllWindows() 
        return result
    
    # 相机输入在线标定
    def cameraAutoMode(self):
        cap = cv2.VideoCapture(args.CAMERA_ID)
        if not cap.isOpened(): 
            raise Exception("from {} read video failed".format(args.CAMERA_ID))
        cap = self.setCamera(cap)
        frame_id = 0
        start_flag = False
        while True:
            key = cv2.waitKey(1)
            ok, raw_frame = cap.read()
            raw_frame = self.imgPreprocess(raw_frame)
            if key == 32: start_flag = True
            if key == 27: break
            if not start_flag:
                cv2.putText(raw_frame, 'press SPACE to start!', (args.FRAME_WIDTH//4,args.FRAME_HEIGHT//2), 
                             cv2.FONT_HERSHEY_COMPLEX, 1.5, (0,0,255), 2)
                cv2.imshow("raw_frame", raw_frame)
                continue
            if frame_id % args.FRAME_DELAY == 0:
                if args.STORE_FLAG:  # 存储该帧图像
                    cv2.imwrite(args.STORE_PATH + 'img_raw{}.jpg'.format(len(self.calibrator.corners)), raw_frame)
                result = self.runCalib(raw_frame) 
                print(len(self.calibrator.corners))
            frame_id += 1 
        cap.release()
        cv2.destroyAllWindows() 
        return result
    
    # 相机输入手动挑选 按空格键采集图片
    def cameraManualMode(self):
        cap = cv2.VideoCapture(args.CAMERA_ID)
        if not cap.isOpened(): 
            raise Exception("from {} read video failed".format(args.CAMERA_ID))
        cap = self.setCamera(cap)
        frame_id = 0
        while True:
            key = cv2.waitKey(1)
            ok, raw_frame = cap.read()
            raw_frame = self.imgPreprocess(raw_frame)
            display = "raw_frame: press SPACE to capture image"
            cv2.namedWindow(display, flags = cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)
            cv2.imshow(display, raw_frame)
            if key == 32:
                if args.STORE_FLAG:  # 存储该帧图像
                    cv2.imwrite(args.STORE_PATH + 'img_raw{}.jpg'.format(len(self.calibrator.corners)), raw_frame)
                result = self.runCalib(raw_frame) 
                print(len(self.calibrator.corners))
            if key == 27: break
        cap.release()
        cv2.destroyAllWindows() 
        return result

    def __call__(self):
        input_type = self.input_type
        mode = self.mode
        if input_type == 'image' and mode == 'auto':
            result = self.imageAutoMode()
        if input_type == 'image' and mode == 'manual':
            result = self.imageManualMode()
        if input_type == 'video' and mode == 'auto':
            result = self.videoAutoMode()
        if input_type == 'video' and mode == 'manual':
            result = self.videoManualMode()
        if input_type == 'camera' and mode == 'auto':
            result = self.cameraAutoMode()
        if input_type == 'camera' and mode == 'manual':
            result = self.cameraManualMode()
        return result

def main():
    calibrator = InCalibrator(args.CAMERA_TYPE)                            # 初始化内参标定器
    calib = CalibMode(calibrator, args.INPUT_TYPE, args.SELECT_MODE)       # 选择标定模式
    result = calib()                                                       # 开始标定
                  
    if len(calibrator.corners) == 0:                      # 标定失败 未找到棋盘或参数设置错误
        raise Exception("Calibration failed. Chessboard not found, check the parameters")  
    if len(calibrator.corners) < args.CALIB_NUMBER:       # 标定样本小于初始化标定所需的图片数
        raise Exception("Warning: Calibration images are not enough. At least {} valid images are needed.".format(args.CALIB_NUMBER))            

    print("Calibration Complete")
    print("Camera Matrix is : {}".format(result.camera_mat.tolist()))                 # 相机内参
    print("Distortion Coefficient is : {}".format(result.dist_coeff.tolist()))        # 畸变向量
    print("Reprojection Error is : {}".format(np.mean(result.reproj_err)))            # 平均重投影误差
    np.save('camera_{}_K.npy'.format(args.CAMERA_ID),result.camera_mat.tolist())
    np.save('camera_{}_D.npy'.format(args.CAMERA_ID),result.dist_coeff.tolist())      # 输出并存储数据
        
if __name__ == '__main__':
    main()