autocalibrate a camera with open-cv py3

calibrate_camera_cv2.py

from glob import glob
import argparse
import cv2
import numpy as np
import os 
import json 
import datetime 

"""
This code calibrates a camera using cv2. 
- [Collection of checkboards for calibration](https://markhedleyjones.com/projects/calibration-checkerboard-collection)
"""

class CamConfig:
    """ Setting for camera calibrations evaluation weight. """
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
    corners = (9, 6) 
    calib_thresh = cv2.CALIB_CB_ADAPTIVE_THRESH + cv2.CALIB_CB_FAST_CHECK + cv2.CALIB_CB_NORMALIZE_IMAGE 
    cameras2check = 10

class KeyCodes: 
    """ Declaration of keyboard keys for serve waitKey comparation. """
    q = ord('q')
    f = ord('f')
    s = ord('s')
    b = ord('b')
    esc = 27

class Files: 
    """ Main class owner of files access. """
    def __init__(self, path: str):
        # Declare and create if needed the path. 
        self.path = path
        if self.path != "/":
            os.mkdir(self.path)

    def save_json(self, file: dict[str, any]): 
        # save json file on path. 
        with open('calibrations.json', 'w') as fp:
            json.dump(file, fp)

    def save_frame(self, frame: list[int], counter: int): 
        # Uses cv2 to save image
        cv2.imwrite(f'{self.path}/image_{(counter+1):0=4d}.png', frame)

class Camera:  
    """ Main class owner of accessing N camera. """
    buffer_frame = None
    buffer_frame_raw = None 
    captures_counter = 0
    calibration_points = []
    calibration_coefficients = None 
    calibration_matrix = None 
    
    def __init__(self, index: int, verbose: bool = False, save_captures: bool = False):
        # Declaration for arguments and initialization of camera. 
        self.index = index 
        self.verbose = verbose
        self.save_captures = save_captures
        self.file_saver = None
        if self.save_captures: 
            self.file_saver = Files(f"camera_{self.index:0=2d}")

        try:
            self.camera = cv2.VideoCapture(self.index)
            if not self.camera.isOpened():
                raise ValueError(f"Error: Camera #{index} not detected")
            _, frame = self.camera.read()
            gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            self.camera_shape = gray_frame.shape[::-1]
        except: 
            self.camera = None 
            print(f"Error: camera #{index:0=2d} not detected")
        
        corner_x, corner_y = CamConfig.corners 
        self.objp = np.zeros((corner_x * corner_y, 3), np.float32)
        self.objp[:, :2] = np.mgrid[0:corner_x, 0:corner_y].T.reshape(-1, 2) 
        print(self.calibration_points)

    def videoCalibration(self, autocalibration: bool = False, limit: int = 50): 
        # Initialize calibration through video campture. 
        while True and (limit == -1 or self.captures_counter <= limit):
            _, frame = self.camera.read() 
            frame_raw = frame 
            
            ret, corners = cv2.findChessboardCorners(frame, CamConfig.corners, None) 
            if ret: 
                gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                corners = cv2.cornerSubPix(gray_frame, corners, (11, 11), (-1, -1), CamConfig.criteria) 
                # cv2.drawChessboardCorners(frame, CamConfig.corners, corners, ret) # need improvements to speed 
                if autocalibration: 
                    self.pointsGenerator(frame_raw)

            cv2.imshow(f"Live Camera {self.index:0=2d}", frame) 
            cv2.imshow(f"Buffer Camera {self.index:0=2d}", self.buffer_frame if self.buffer_frame is not None else frame)

            key_pressed = cv2.waitKey(1) & 0xFF
            match key_pressed: 
                case KeyCodes.f: self.pointsGenerator(frame_raw) 
                case KeyCodes.s: self.buffer_frame, self.buffer_frame_raw = frame, frame_raw 
                case KeyCodes.b: self.pointsGenerator(self.buffer_frame_raw) 
                case KeyCodes.q: break 
                case KeyCodes.esc: break 

        cv2.destroyAllWindows() 

    def pointsGenerator(self, frame: list[int]): 
        # Calculation of points over original frame. 
        if self.save_captures: 
            self.file_saver.save_frame(frame, self.captures_counter)
        gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        ret, corners = cv2.findChessboardCorners(frame, CamConfig.corners, CamConfig.calib_thresh) 
        if ret: 
            corners = cv2.cornerSubPix(gray_frame, corners, (11, 11), (-1, -1), CamConfig.criteria)
            self.calibration_points.append(corners)
            self.captures_counter += 1 
            if self.verbose: 
                print(f"[+] New Corner Points Generated")
                print(f"[~] Total captures: {self.captures_counter}\n") 

    def getCalibrationCoefficients(self) -> tuple[list[int], list[int]]: 
        # Calculate matrix coefficients for calibration. 
        if self.captures_counter == 0: 
            print("Error: Not enough video calibration captures") 
            return (None, None)
        objpoints = [self.objp] * len(self.calibration_points)
        undistortion = cv2.calibrateCamera(objpoints, self.calibration_points, self.camera_shape, None, None) 
        _, self.calibration_matrix, self.calibration_coefficients, _, _ = undistortion 
        if self.verbose: 
            print(f"[+] Calibration coefficients generated")
            print(f"[>] Calibration Matrix: {self.calibration_matrix}")
            print(f"[>] Calibration Coefficients: {self.calibration_coefficients}\n") 
        return (self.calibration_matrix, self.calibration_coefficients) 

    def testVideoCalibration(self, matrix = None, coefficients = None):
        # Calibration of camera using indexes. 
        if matrix is not None and coefficients is not None: 
            self.calibration_matrix = matrix 
            self.calibration_coefficients = matrix 
            if self.verbose: 
                print(f"[+] Set new calibration indexes") 
                print(f"[>] Calibration Matrix: {self.calibration_matrix}") 
                print(f"[>] Calibration Coefficients {self.calibration_coefficients}\n")

        while True: 
            _, frame = self.camera.read() 
            frame_calibrated = cv2.undistort(frame, self.calibration_matrix, self.calibration_coefficients, None, self.calibration_matrix) 
            cv2.imshow(f"Original Live Camera {self.index:0=2d}", frame) 
            cv2.imshow(f"Calibrated Live Camera {self.index:0=2d}", frame_calibrated)
            
            key_pressed = cv2.waitKey(1) & 0xFF 
            match key_pressed: 
                case KeyCodes.q: break 
                case KeyCodes.esc: break 
        cv2.destroyAllWindows()

    def release(self): 
        # Release camera. 
        self.camera.release() 

def getArgs() -> argparse.Namespace: 
    # Cast arguments on single array. 
    parser = argparse.ArgumentParser() 
    parser.add_argument('-c', '--captures')
    parser.add_argument('-s', '--save', action='store_true')
    parser.add_argument('-a', '--autocalibration', action='store_true')
    parser.add_argument('-v', '--verbose', action='store_true')
    args = parser.parse_args() 
    return args

def calibrateCamera(args: argparse.Namespace): 
    # Running function, enumeration & calibrate each camera. 
    configurations = {'Execution': "01", 'Cameras': []}
    for cam_index, cam_name in enumerate(glob("/dev/video*") or range(CamConfig.cameras2check)): 
        camera = Camera(cam_index, verbose = args.verbose, save_captures = args.save)

        if not hasattr(camera, "camera"):
            continue 

        if camera.camera is None: 
            continue 

        camera.videoCalibration(autocalibration = args.autocalibration, limit = int(args.captures or 50))
        mat, coff = camera.getCalibrationCoefficients()
        camera.testVideoCalibration()
        camera.release()

        configurations['Cameras'].append({'Camera index': cam_index, 
                               'Source': cam_name, 
                               'Coefficients': coff.tolist(), 
                               'Matrix': mat.tolist()})
        del camera  

    if args.verbose: 
        print("[+] Configurations dictionary created. ")
        print(configurations) 

    if args.save: 
        saver = Files("/")
        saver.save_json(configurations)

if __name__ == "__main__": 
    args = getArgs() 
    calibrateCamera(args)