Калибровка камеры для ArUco

from pioneer_sdk2 import Camera, ImageViewer          # импортируем классы Camera и ImageViewer из библиотеки pioneer_sdk2
import cv2                                            # библиотека cv2 содержит функции для работы с изображениями
import numpy as np                                    # библиотека numpy нужна для работы с массивами и матрицами
import glob                                           # библиотека glob нужна для поиска файлов в папке
import time                                           # библиотека time содержит функции для работы со временем
import sys                                            # библиотека sys нужна для чтения команд из терминала
import select                                         # библиотека select нужна для неблокирующего чтения команд
from pathlib import Path                              # класс Path нужен для работы с путем к файлу калибровки


DATA_PATH = Path(__file__).with_name("data.yml")       # путь к файлу калибровки рядом со скриптом


def get_images_from_folder(folder_name, file_type="*.jpg"): # функция для получения изображений из папки
    images_list = glob.glob(folder_name + "/" + file_type)  # находит все файлы нужного типа в указанной папке
    images = []                                             # создаем пустой список для хранения изображений

    for fname in images_list:                               # перебираем все найденные файлы
        img = cv2.imread(fname)                             # считываем изображение из файла
        images.append(img)                                  # добавляем изображение в список images

    return images                                           # возвращаем список изображений


def read_terminal_command():                                # функция проверяет, введена ли команда в терминале
    if select.select([sys.stdin], [], [], 0)[0]:             # проверяем, есть ли введенная строка
        return sys.stdin.readline().strip().lower()          # считываем команду и приводим ее к нижнему регистру
    return None                                              # возвращаем None, если команды нет


def get_images_from_drone_camera(camera, viewer, save=False): # функция для получения изображений с камеры квадрокоптера
    images = []                                             # создаем пустой список для хранения снимков

    print("Сделайте 10-15 снимков шахматной доски с разных ракурсов") # выводим инструкцию для пользователя
    print("Откройте поток rtsp://10.42.0.1:8889/calibration/") # объясняем, где смотреть изображение
    print("Введите '1' и нажмите Enter, чтобы сделать снимок") # объясняем, как сделать снимок
    print("Введите 'q' и нажмите Enter для завершения")       # объясняем, как завершить съемку

    while True:                                            # запускаем бесконечный цикл
        try:                                               # пробуем получить кадр с камеры
            frame = camera.get_cv_frame(timeout=1.0)       # получаем один кадр с камеры
                                                           # timeout=1.0 - время ожидания кадра в секундах
        except TimeoutError:                               # если кадр не успел прийти
            continue                                       # пропускаем текущую итерацию цикла

        if frame is None:                                  # проверяем, что кадр не был получен
            continue                                       # пропускаем текущую итерацию цикла

        viewer.imshow("calibration", frame, fps=10)        # отправляем кадр в RTSP-трансляцию
        command = read_terminal_command()                  # считываем команду из терминала

        if command == "q":                                 # проверяем, введена ли команда выхода
            return images                                  # возвращаем список сделанных снимков

        if command == "1":                                 # проверяем, введена ли команда снимка
            images.append(frame)                           # сохраняем текущий кадр в список images
            print("Снимок №" + str(len(images)))           # выводим номер сделанного снимка

            if save:                                       # проверяем, нужно ли сохранять снимки в файлы
                cv2.imwrite("frame_" + str(len(images)) + ".jpg", frame) # сохраняем снимок в jpg-файл

            time.sleep(1)                                  # ставим паузу на 1 секунду


def calibrate(images, viewer=None):                        # функция для калибровки камеры по снимкам шахматной доски
    CHECKERBOARD = (6, 9)                                  # количество внутренних углов шахматной доски
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001) # критерий точности поиска углов
    calibration_flags = (                                  # флаги для поиска углов шахматной доски
        cv2.CALIB_CB_ADAPTIVE_THRESH                       # используем адаптивную обработку изображения
        + cv2.CALIB_CB_FAST_CHECK                          # ускоряем проверку наличия шахматной доски
        + cv2.CALIB_CB_NORMALIZE_IMAGE                     # нормализуем изображение для лучшего поиска углов
    )

    objpoints = []                                         # список для хранения 3D-точек шахматной доски
    imgpoints = []                                         # список для хранения 2D-точек на изображении
    image_size = None                                      # размер изображения для калибровки

    objp = np.zeros((1, CHECKERBOARD[0] * CHECKERBOARD[1], 3), np.float32)           # создаем массив 3D-точек шахматной доски
    objp[0, :, :2] = np.mgrid[0:CHECKERBOARD[0], 0:CHECKERBOARD[1]].T.reshape(-1, 2) # задаем координаты углов доски

    for img in images:                                     # перебираем все изображения
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)       # переводим изображение в черно-белый формат
        image_size = gray.shape[::-1]                      # сохраняем размер изображения

        ret, corners = cv2.findChessboardCorners(          # ищем углы шахматной доски на изображении
            gray,                                         # передаем черно-белое изображение
            CHECKERBOARD,                                 # передаем размер шахматной доски
            flags=calibration_flags                       # передаем флаги поиска углов
        )

        if ret:                                            # проверяем, были ли найдены углы шахматной доски
            objpoints.append(objp)                         # добавляем 3D-точки шахматной доски в список
            corners2 = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria) # уточняем координаты найденных углов
            imgpoints.append(corners2)                     # добавляем уточненные 2D-точки в список
            img = cv2.drawChessboardCorners(img, CHECKERBOARD, corners2, ret) # рисуем найденные углы на изображении

        if viewer is not None:                             # проверяем, нужно ли показать результат поиска углов
            viewer.imshow("calibration_result", img, fps=1) # отправляем изображение с углами в RTSP-трансляцию
            time.sleep(0.5)                                # даем время посмотреть результат

    if image_size is None or not imgpoints:                # проверяем, что есть данные для калибровки
        raise RuntimeError("Не удалось найти углы шахматной доски для калибровки")

    camera_matrix = np.zeros((3, 3), np.float64)            # матрица камеры будет рассчитана OpenCV
    dist_coeffs = np.zeros((5, 1), np.float64)              # коэффициенты искажений будут рассчитаны OpenCV

    ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(    # выполняем калибровку камеры
        objpoints,                                        # передаем 3D-точки шахматной доски
        imgpoints,                                        # передаем 2D-точки на изображениях
        image_size,                                       # передаем размер изображения
        camera_matrix,                                    # передаем матрицу камеры для заполнения
        dist_coeffs                                       # передаем коэффициенты искажений для заполнения
    )

    return mtx, dist                                       # возвращаем матрицу камеры и коэффициенты искажений


def save_coefficients(mtx, dist, path):                    # функция для сохранения коэффициентов калибровки в файл
    cv_file = cv2.FileStorage(str(path), cv2.FILE_STORAGE_WRITE) # открываем файл для записи данных калибровки
    cv_file.write("mtx", mtx)                              # записываем матрицу камеры в файл
    cv_file.write("dist", dist)                            # записываем коэффициенты искажений в файл
    cv_file.release()                                      # закрываем файл после записи


def load_coefficients(path):                               # функция для загрузки коэффициентов калибровки из файла
    cv_file = cv2.FileStorage(str(path), cv2.FILE_STORAGE_READ) # открываем файл для чтения данных калибровки
    camera_matrix = cv_file.getNode("mtx").mat()           # считываем матрицу камеры из файла
    dist_coeffs = cv_file.getNode("dist").mat()            # считываем коэффициенты искажений из файла
    cv_file.release()                                      # закрываем файл после чтения
    return camera_matrix, dist_coeffs                      # возвращаем матрицу камеры и коэффициенты искажений


camera = Camera()                                          # создаем экземпляр класса Camera для получения кадров с камеры
viewer = ImageViewer()                                     # создаем экземпляр класса ImageViewer для трансляции кадров

try:                                                       # основной код находится внутри блока try
    images = get_images_from_drone_camera(camera, viewer)  # получаем снимки шахматной доски с камеры квадрокоптера
    mtx, dist = calibrate(images, viewer)                  # выполняем калибровку камеры по полученным снимкам
    save_coefficients(mtx, dist, DATA_PATH)                # сохраняем результаты калибровки в файл data.yml рядом со скриптом

finally:                                                   # блок finally выполнится при завершении программы
    viewer.close()                                         # останавливаем RTSP-трансляцию
    camera.stop()                                          # останавливаем получение кадров с камеры