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

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)  # находит все файлы нужного типа в указанной папке
                                                            # file_type="*.jpg" означает, что ищем только jpg-файлы

    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)        # получаем один кадр с камеры
        except TimeoutError:                                # если кадр не успел прийти
            continue                                        # пропускаем текущую итерацию цикла

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

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

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

                time.sleep(1)                               # делаем паузу 1 секунду, чтобы не сделать лишние снимки

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


def calibrate(images, viewer=None):                         # функция для калибровки камеры по снимкам шахматной доски
    CHECKERBOARD = (6, 9)                                   # количество внутренних углов шахматной доски
                                                            # 6 - количество углов по одной стороне, 9 - по другой

    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001) # критерий точности поиска углов
                                                            # 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 == True:                                     # проверяем, были ли найдены углы шахматной доски
            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()                                           # останавливаем получение кадров с камеры