Привет, Хабр! Сегодня с вами команда регионального научно-образовательного центра «Искусственный интеллект и анализ больших данных» при НГТУ им. Р. Е. Алексеева. Продолжаем рассказывать о нашей работе по возрождению и улучшению DPED (Deep Photo Enhancement Dataset). Это открытый проект исследователей из ETH Zurich, который включает датасет парных изображений и нейросетевую модель для повышения качества мобильных фотографий до уровня DSLR. В нашем случае мы хотим довести снимки с планшета YADRO KVADRA_T до качества полупрофессиональной камеры Sony Alpha ILCE 6600.

Отметим, что цель проекта не только исследование и обучение модели, но и последующее внедрение полученных наработок в приложение камеры планшета. Мы рассматриваем варианты локального инференса на самом устройстве, включая оптимизацию модели под мобильные вычислительные платформы с использованием TensorRT или ONNX Runtime. Так улучшать изображения можно прямо на устройстве — либо в момент съемки, либо в фоновом режиме. 

В первой статье мы рассказали, почему выбрали DPED, как его запустили на современном софте и проверили на реальной задаче. На этот раз подробно разберем, как правильно собрать и подготовить датасет, без которого невозможно добиться высоких результатов в задачах машинного обучения.

Почему важен качественный датасет

Нейросеть обучается только на тех данных, которые ей предоставили. Если у исходного датасета низкое качество, то и результат работы модели будет соответствующий. Руководствуясь принципом «чему учили, то и получили», мы внимательно изучили опыт авторов оригинального проекта DPED и решили применить их подход. 

Для обучения модели нужен датасет из парных изображений, снятых одновременно с двух устройств. Затем из снимков вырезается общая область, по которой модель будет учиться улучшать фото с планшета.

Как мы автоматизировали съемку

Авторы оригинального проекта использовали специальный стенд, на котором были закреплены четыре устройства: три смартфона и одна профессиональная камера:

Для сбора снимков мы создали похожий стенд с планшетом KVADRA_T и фотоаппаратом Sony Alpha ILCE 6600:

Основная задача — добиться одновременной съемки одной и той же сцены всеми устройствами. Способ синхронизации разработчики DPED не описали, поэтому мы нашли решение самостоятельно.

Камерой можно управлять через приложение Sony Imaging Edge Desktop с компьютера благодаря функции дистанционного управления (Remote Control). Для этого нужно подключить камеру через USB-кабель из комплекта. 

Однако, приложение не позволяет автоматически запускать режим серийной съемки — нужно нажать кнопку затвора в интерфейсе приложения. Поэтому мы решили автоматизировать этот процесс. Самым простым и быстрым решением стала Python-библиотека pyautogui, которая позволяет программно управлять курсором и кликами мыши.

Скрипт имитирует нажатие кнопки затвора в приложении Sony Imaging Edge Desktop. Да, это не идеальный способ, так как работа скрипта зависит от множества внешних факторов, включая скорость отклика GUI и задержки самой программы. Но для старта проекта его оказалось достаточно.

На планшете мы попробовали автоматизировать съемку с помощью Android-приложения IP Webcam, но его работа оказалась слишком нестабильной. Поэтому мы перешли на ADB (Android Debug Bridge) — это инструмент командной строки из состава Android SDK, который позволяет управлять Android-устройствами с компьютера: выполнять команды, снимать скриншоты, получать доступ к файловой системе и, в нашем случае, дистанционно запускать съемку. Если включить на планшете отладку по USB, то ADB может делать снимки через команду терминала. Мы также написали Python-скрипт с использованием subprocess, который запускал команду захвата кадров через ADB.

Чтобы запустить съемку дистанционно:

1. Включаем «Отладку по USB» на планшете.

2. Подключаем его к компьютеру.

3. Запускаем команду захвата кадров ADB с помощью Python-скрипта с subprocess.

Метод быстрый и стабильный, а главное для нас — он легко автоматизируется и масштабируется.

Позднее мы нашли более удобные решения автоматизации съемки на планшете, которые могут быть успешно использованы вместо метода захвата кадров через ADB: 

• Open source-библиотека Sony PTP — позволяет управлять камерой Sony по одноименному протоколу. Мы потратили немало времени на ее поиск, но оно того стоило — в библиотеке есть готовые bash-скрипты для Linux и подробная документация.

• SCRCPY — удобный инструмент для управления Android-устройствами, который лучше подходит для решения нашей задачи.

Менять работающую систему мы не стали, так как уже активно собирали снимки для датасета и любые изменения могли привести к задержкам в нашем проекте. 

Как синхронизировать кадры: многопоточный подход

Мы решили задачи автоматизации, но столкнулись с еще одной проблемой: фото на планшете и камере снимались с некоторой задержкой относительно друг друга. Использование простых пауз (time.sleep) оказалось ненадежно и неэффективно. Тогда мы реализовали многопоточное решение:

• Первый поток управляет съемкой с камеры с помощью библиотеки pyautogui.

• Второй поток управляет съемкой с планшета через ADB.

• Оба потока обмениваются информацией через очередь (queue.Queue() из стандартной библиотеки Python) — это потокобезопасная структура данных, которая позволяет одному потоку передать сигнал другому. В нашем случае очередь используется для передачи сигнала о начале съемки с камеры. Получив этот сигнал, планшет почти без задержки запускает захват изображения.

В процессе тестирования среднее время задержки составило 50 мс, но разброс данных достигал 93 мс. То есть, существуют случаи, когда мы получаем изображения с непозволительной задержкой в 100 мс и более. Мы отметили этот момент, но продолжили собирать датасет, а изображения с большой задержкой — удалять.

Скрипт автоматизации съемки кадров:

import subprocess
from threading import Thread
import pyautogui
import time
from queue import Queue

# координаты для кликов мыши

CAMERA_SHUTTER_BUTTON = (329, 748)    # кнопка затвора в приложении

FOCUS_POINT = (1189, 204)            # точка фокуса или область кадра


def tablet(q):
    time.sleep(0.1)
    if q.get() == 1:
        p = subprocess.Popen(r'.\adb.exe shell', stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        p.stdin.write(b'input keyevent 27')
        p.stdin.close()

def camera(q):
    pyautogui.click(*CAMERA_SHUTTER_BUTTON)
    pyautogui.moveTo(*FOCUS_POINT)
    q.put(1)
    pyautogui.mouseDown()
    time.sleep(0.02)
    pyautogui.mouseUp()

q = Queue()
thread1 = Thread(target=camera, args=(q,))
thread2 = Thread(target=tablet, args=(q,))
thread1.start()
thread2.start()

В оригинальной работе DPED точные значения задержки не указывались: авторы фиксировали устройства на механическом стенде и выполняли съемку вручную, без программной синхронизации или последующего анализа временного лага между кадрами. Насколько нам удалось выяснить, синхронизация производилась «на глаз», что не позволяет оценить точность в миллисекундах. Таким образом, можно утверждать, что наша реализация обеспечивает более детерминированный и измеримый результат по синхронизации.

Как мы изменяли конструкцию стенда для съемки

Мы несколько раз дорабатывали конструкцию стенда, чтобы добиться его максимальной устойчивости.

Первая итерация оказалась неустойчивой, и часть снимков получалась с критичными для нас искажениями. 

Вторую версию стенда мы спроектировали сами и напечатали на 3D-принтере: это была легкая стабилизирующая платформа, которая улучшила фиксацию, но не обеспечивала достаточной жесткости конструкции.

На фото ниже — финальная версия стенда в сборе: планшет и камера закреплены на массивной платформе, установленной на штативе с уровнем. Вы можете скачать STL-файлы этих моделей.

Собираем датасет

Теперь наш стенд с оборудованием готов к сбору материалов для датасета. На этом этапе команда разделилась на две группы. Первая занялась съемкой парных фотографий, а вторая начала работу над предварительной обработкой полученных снимков: выделением общей области и извлечением из них патчей, то есть небольших фрагментов, которые нам нужны для обучения нейросети.

Напомним, что для обучения модели нам нужны снимки одной и той же сцены, сделанные одновременно на планшет KVADRA_T и камеру Sony. Поскольку у объективов устройств разные фокусные расстояния, углы обзора, светосила и так далее, то и фотографии всегда отличаются. Поэтому мы выделяем общую для парных снимков область и дальше работаем уже с ней.

Чтобы выделить такую область, мы находим на снимках ключевые точки: четко выраженные и устойчивые элементы изображения, такие как угол дома, край окна или характерная текстура. Их можно легко определить на разных снимках одной и той же сцены. 

Каждая ключевая точка описывается специальным дескриптором, который характеризует ее окружение, что позволяет однозначно сопоставить точки на двух разных снимках. По найденным соответствиям ключевых точек затем вычисляется матрица гомографии — математическое преобразование, которое связывает два изображения между собой. Подробнее о ней мы расскажем ниже.

Есть несколько популярных алгоритмов поиска ключевых точек:

• SIFT (Scale-Invariant Feature Transform) — инвариантный к масштабу и повороту метод, дающий хорошие результаты практически на любых сценах. Начиная с OpenCV 4.4, включен в состав opencv_contrib и доступен через cv2.SIFT_create(), но требует установки opencv-contrib-python.

• SURF (Speeded-Up Robust Features) — быстрый и стабильно работающий метод, но он защищен патентами, поэтому отсутствует даже в opencv-contrib и требует самостоятельной пересборки OpenCV с включенным модулем xfeatures2d.

• ORB (Oriented FAST and Rotated BRIEF) — бесплатный, открытый и быстрый метод. Он доступен в стандартной OpenCV, включая opencv-python, без дополнительных зависимостей.

• AKAZE — современный и эффективный метод, также доступен в стандартной сборке OpenCV. Поддерживает нелинейные шкалированные пространства, подходит для множества сценариев.

• SuperPoint — не входит в OpenCV. Это внешняя реализация на основе нейросетей (TensorFlow/PyTorch), которая требует отдельной загрузки весов и подключения через сторонний код.

Авторы оригинального проекта DPED использовали метод SIFT, и мы, протестировав несколько подходов, тоже выбрали его как самый надежный. Еще мы проверили работу алгоритма ORB, но его результаты на сценах с выраженными перспективными искажениями нас не удовлетворили, и мы от него отказались.

На изображениях ниже показаны результаты работы алгоритмов SIFT и ORB. Зеленые точки обозначают найденные ключевые точки, а линии между изображениями — найденные соответствия между точками на паре изображений. Линии могут быть пунктирными или цветными в зависимости от выбранной визуализации (например, при использовании cv2.drawMatches() в OpenCV). Рекомендуем увеличить изображение, чтобы разглядеть точки и линии.

Обратите внимание, что SIFT находит больше стабильных точек, особенно в зонах с неоднородной текстурой и при изменениях перспективы.

Время выполнения (в секундах):

• real — 31.47

• user — 189.89

• sys — 3.68

Время выполнения (в секундах):

• real — 1.27

• user — 2.70

• sys — 0.30

Хотя алгоритм ORB работает значительно быстрее SIFT, но в большинстве случаев проигрывает ему в точности и стабильности обнаружения ключевых точек, особенно на изображениях с выраженными перспективными искажениями и сложными текстурами.

Мы нашли ключевые точки в парах изображений, и теперь нужно найти их соответствия друг другу. Для этого применяются специальные методы сопоставления, такие как RANSAC (Random Sample Consensus) и FLANN (Fast Library for Approximate Nearest Neighbors). Эти алгоритмы позволяют определить наиболее устойчивые пары точек, исключая «выбросы» — случайно совпавшие точки, которые могут ухудшить качество дальнейших вычислений.

Коротко поясним принцип работы этих алгоритмов:

• FLANN основан на методе приближенного поиска ближайших соседей (nearest neighbors). Он быстро и эффективно сопоставляет дескрипторы ключевых точек, что позволяет оперативно отбирать хорошие соответствия и исключать ошибочные совпадения.

• RANSAC работает по итеративному принципу: многократно случайным образом выбирает небольшие подмножества точек и на их основе строит модель преобразования (например, матрицу гомографии). После этого сам алгоритм выбирает ту модель преобразования, которая наилучшим образом согласуется с наибольшим числом точек (inliers), автоматически исключая выбросы (outliers), что не вписываются в построенную гипотезу.

После фильтрации и отбора устойчивых пар ключевых точек мы переходим к построению матрицы гомографии размером 3×3, которая описывает проективное преобразование между двумя плоскостями изображений. Другими словами, она позволяет проецировать одно изображение на другое, компенсируя перспективные искажения и различия в ракурсах съемки. Именно матрица гомографии поможет определить и вырезать общую область из пары изображений, полученных с разных устройств. Говоря упрощенно, матрица гомографии описывает, каким образом одно изображение можно проецировать на другое, учитывая перспективные искажения перспективы.

С математической точки зрения, изображение можно представить в виде матрицы пикселей. Умножив координаты точек изображения на матрицу гомографии размером 3×3, мы получим координаты соответствующих точек на втором изображении.

Даже если снимки сделаны на устройства с разными фокусными расстояниями, углами обзора и положением в пространстве, с помощью гомографии можно перевести одно изображение в проекцию другого. Так мы корректно совместим сцены и выделим общую область, присутствующую на обоих кадрах в одной и той же перспективе, — что и нужно для дальнейшей нарезки патчей и обучения модели.

Давайте разберемся, как именно мы выделяем и сохраняем общую область между изображениями с обоих устройств.

Сначала мы используем ранее вычисленную матрицу гомографии для преобразования изображения с планшета в координатную систему камеры. Это позволяет привести обе сцены к одной проекции и обеспечить их пространственное согласование.

Затем создаем логические маски для каждого изображения. Такая маска представляет собой двумерную матрицу, где каждый пиксель принимает значение 1, если он содержит информацию, и 0, если это фоновая или неиспользуемая область (например, черные поля, появившиеся в результате трансформации). 

Далее выполняется побитовая операция AND (&) между этими двумя масками — преобразованного изображения с планшета и исходного изображения с камеры.  В результате получается маска, обозначающая общую непустую область — ту часть сцены, которая одновременно присутствует на обоих изображениях после преобразования. Этот шаг также позволяет автоматически избавиться от черных полей, образовавшихся после применения гомографии к изображению с планшета, — они просто не попадут в итоговую общую маску. 

На следующем шаге определяем границы результирующей области по логической маске. После вырезаем эту общую область из обоих изображений для последующей работы (например, для извлечения патчей). 

Важно отметить, что из-за разных фокусных расстояний и углов обзора у камеры и планшета даже идеально выровненные кадры будут содержать локальные искажения. Например, объекты в центре сцены могут располагаться с небольшим сдвигом или масштабным искажением относительно друг друга. Такие эффекты — следствие физики оптики и проекции, и полностью устранить их средствами постобработки невозможно. Посмотрите на фото ниже, где хорошо видны такие искажения.

На переднем плане располагается столб, а сразу за ним — здание нашего университета. Хотя камера и планшет были установлены рядом, на одной поверхности и ориентированы в одном направлении, но геометрия сцены на полученных изображениях все равно различается. Часть здания слева от столба на снимке с камеры выглядит шире, чем на изображении с планшета. Это связано с различиями в фокусных расстояниях и углах обзора объективов, а также с характером проекционного искажения. 

Такие локальные расхождения невозможно полностью компенсировать ни с помощью гомографии, ни другими способами выравнивания. Поэтому на этапе подготовки датасета мы отбирали патчи с высокой степенью соответствия. Для этого мы установили порог по нормализованной кросс-корреляции не ниже 0.9, чтобы отобрать только те патчи, где изображения максимально совпадают по содержанию. Это позволяло исключать области с геометрическими искажениями, локальными сдвигами или недостаточной информативностью, которые могли бы ухудшить обучение модели.

На фотографиях выше показаны вырезанные общие области с изображений, которые получены с камеры и планшета. Эти участки сцены полностью совпадают по содержанию и могут быть использованы для дальнейшей нарезки патчей.

Извлекаем патчи

После выделения общих областей пора начинать подготовку обучающих пар патчей. Для этого мы реализовали базовый алгоритм, основанный на двух вложенных циклах, который проходится по изображению с заданным шагом — в нашем случае 100 пикселей по ширине и высоте.

В каждой итерации алгоритм извлекает квадратные фрагменты (патчи) фиксированного размера из обоих изображений. Затем для каждой пары патчей вычисляется кросс-корреляция: метрика схожести двух изображений. Если значение метрики превышает заданный порог (в нашем случае 0.9), то патч считается качественным и сохраняется в датасет. Так мы сразу отбрасываем фрагменты, которые содержат смещения, искажения или малоинформативные области. 

Для подсчета метрики мы использовали нормализованную кросс-корреляцию (NCC), которая измеряет степень линейной зависимости между двумя патчами:

NCC(A, B) = Σ[(Aᵢ - Ā)(Bᵢ - B̄)] / (sqrt(Σ(Aᵢ - Ā)²) * sqrt(Σ(Bᵢ - B̄)²))

Здесь Aᵢ и Bᵢ — значения пикселей в патчах A и B, Ā и B̄ — средние значения интенсивности в каждом патче. Метрика принимает значения от –1 до 1, где 1 означает полное совпадение, 0 — отсутствие линейной зависимости, а –1 — полное обратное соответствие. Мы считали патчи приемлемыми, если NCC ≥ 0.9.

Однако в процессе работы мы увидели, что даже при гомографическом выравнивании некоторые пары изображений сохраняют остаточные локальные искажения. Это связано как с различиями в оптике, так и с микроскопическими сдвигами в момент съемки. Поэтому мы предложили два подхода к улучшению качества отбора патчей.

Учет глобального смещения (поиск коэффициента «сдвига»). Мы предположили, что на большинстве пар изображений существует постоянное смещение одной сцены относительно другой — например, влево или вверх. Это может быть результатом небольшого поворота одного из устройств, разной высоты установки или других аппаратных факторов.

Мы решили оценить это смещение (в пикселях) в центральной части кадра, а затем сдвигать область извлечения патчей на полученную величину. Такой подход позволил бы сосредоточиться на тех участках, где значение кросс-корреляции было наибольшим, и сократить число неудачных совпадений.

Поиск по локальной области (расширенное сопоставление) — более ресурсоемкий, но, как выяснилось, более надежный метод. Вместо того чтобы сравнивать патчи строго по одинаковым координатам, мы брали один патч с камеры и искали для него наилучшее соответствие среди множества патчей из окрестной области на снимке с планшета — с площадью от 100×100 до 400×400 пикселей.

Отметим, что локальные смещения внутри патча могут быть неравномерными и вызваны перспективными искажениями. Поэтому задачу можно рассматривать как локальный поиск матрицы преобразования (например, аффинной), и для такой локальной подгонки можно применять метод cv2.findTransformECC(), который вычисляет преобразование на основе максимизации корреляции между патчами.

В итоге для каждого патча мы вычисляли несколько значений кросс-корреляции и сохраняли наиболее похожую пару. Это позволило нивелировать небольшие смещения объектов и компенсировать локальные искажения проекции.

Выбор между двумя предложенными методами отбора патчей оказался непростым. Второй вариант, с перебором патчей по области, уже был реализован, однако у него был существенный недостаток — высокая вычислительная нагрузка. 

Первый метод, основанный на оценке смещения, казался более быстрым, но имел существенные ограничения: на тестовых изображениях смещение наблюдалось далеко не всегда, а в некоторых случаях оно отличалось не только по направлению, но и по величине. Такая непредсказуемость объясняется реальными условиями съемки: даже при стабильной установке оборудования невозможно гарантировать абсолютно идентичное позиционирование и проекцию сцены. 

Более того, смещение может быть неравномерным — например, минимальным в центре изображения и значительно большим по краям. В таких условиях попытка сдвигать все патчи на фиксированный коэффициент может привести к искажению данных для обучения модели. Именно поэтому мы приняли решение использовать более надежный подход — поиск соответствий патчей по локальной области. Хотя метод и потребовал больше времени, он оказался более устойчивым к реальным искажениям и позволил сформировать более качественный датасет.

Оптимизируем процесс генерации патчей

В оригинальном проекте авторы сохраняли патчи в отдельных папках и затем вручную разделяли их на тренировочные и тестовые. В итоге получалось в среднем около 130 тысяч патчей. Мы решили автоматизировать процесс с помощью скрипта, который который работает в многопоточном режиме. Скрипт решает сразу три задачи:

• выделяет общую область на паре изображений,

• извлекает патчи с учетом кросс-корреляции,

• сохраняет результат в бинарный файл формата .h5, что позволяет легко масштабировать и быстро загружать данные при обучении модели.

Сам скрипт:

from multiprocessing import Queue, Process
from concurrent.futures import ProcessPoolExecutor, as_completed
from h5 import Dataset
from pathlib import Path
from process_image import process_image
import argparse


parser = argparse.ArgumentParser()
parser.add_argument('path')
filename = 'dataset.h5'
path_dataset = parser.parse_args().path

def writer_process(queue):
    with Dataset(filename, 'a') as dataset:
        while True:
            data = queue.get()
            if data == 'stop':
                break
            if data is not None:
                [dataset.insert(pair[0], pair[1]) for pair in data]

queue = Queue()
writer = Process(target=writer_process, args=(queue,))
writer.start()

image_paths = [(str(path.name), i) for i, path in enumerate(Path(path_dataset).glob('*'))]

with ProcessPoolExecutor(max_workers=10) as executor:

    futures = [
        executor.submit(process_image, filename)
        for filename in image_paths
    ]

    for future in as_completed(futures):
        chunk = future.result()
        queue.put(chunk)

queue.put('stop')
writer.join()



def process_image(arg: tuple[str, int]) -> list | None:
    filename, index = arg
    path = 'dataset/'
    img1 = cv2.imread(f'{path}/tablet/{filename}')
    img2 = cv2.imread(f'{path}/camera/{filename}')

    gray1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    gray2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)

    sift = cv2.SIFT_create()

    kp1, des1 = sift.detectAndCompute(gray1, None)
    kp2, des2 = sift.detectAndCompute(gray2, None)
    bf = cv2.BFMatcher(cv2.NORM_L2, crossCheck=True)
    matches = bf.match(des1, des2)
    matches = sorted(matches, key=lambda x: x.distance)

    good_matches = matches[:50]

    src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1,1,2)
    dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).reshape(-1,1,2)

    H, mask = cv2.findHomography(dst_pts, src_pts, cv2.RANSAC, 5.0)

    height, width = img1.shape[:2]
    warped_img2 = cv2.warpPerspective(img2, H, (width, height))

    overlap_mask = (cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY) > 0) & (cv2.cvtColor(warped_img2, cv2.COLOR_BGR2GRAY) > 0)

    if not (np.count_nonzero(overlap_mask) == 0):
        ys, xs = np.where(overlap_mask)
        x_min, x_max = xs.min(), xs.max()
        y_min, y_max = ys.min(), ys.max()

        crop1 = img1[y_min:y_max, x_min:x_max]
        crop2 = warped_img2[y_min:y_max, x_min:x_max]

        img1, img2 = crop1, crop2
        if img1.shape != img2.shape:
            return None

        height, width, _ = img1.shape
        c = 0
        counter = 0
        cross_correlation = 0.9

        patch_arr = list()
        cor_arr = list()
        patches_list = list()

        x1, y1 = 0, 0
        for y in trange(1, height, 100):
            print(f"************ {index + 1} step ****************")
            for x in range(1, width, 100):
                if x + 109 < width and y + 109 < height:
                    start_x = 1
                    finish_x = width
                    start_y = 1
                    finish_y = height

                    kx = 0
                    ky = 0

                    if x - start_x <= 10:
                        kx = x - start_x
                    elif finish_x - x <= 10:
                        kx = finish_x - x
                    else:
                        kx = 10
                        x1 = x - 10
                    
                    if y - start_y <= 10:
                        ky = y - start_y
                    elif finish_y - y <= 10:
                        ky = finish_y - y
                    else:
                        ky = 10
                        y1 = y - 10
                    
                    patch1 = img1[y:y+100, x:x+100]
                    if np.any(patch1 == 0):
                         continue
                    patch1_copy = patch1.copy()
                    for j in range(1, 11 + ky):
                        for i in range(1, 11 + kx):
                            patch2 = img2[y1+j:y1+100+j, x1+i:x1+100+i]
                            if np.any(patch1 == 0) or np.any(patch2 == 0):
                                continue
                            patch2_copy = patch2.copy()
                            flat1 = cv2.cvtColor(patch1_copy, cv2.COLOR_BGR2GRAY).flatten()
                            flat2 = cv2.cvtColor(patch2_copy, cv2.COLOR_BGR2GRAY).flatten()
                            correlation, _ = pearsonr(flat1, flat2)
                            patch_arr.append(patch2)
                            cor_arr.append(correlation)
                    max_cor = max(cor_arr)
                    if max_cor >= cross_correlation:
                        i = cor_arr.index(max_cor)
                            patches_list.append(
                        [np.transpose(patch1, (2, 0, 1)), np.transpose(patch_arr[i], (2, 0, 1))]
                        )
                        c += 1
                    cor_arr = []
                    patch_arr = []
        return patches_list

В однопоточном режиме процесс сборки датасета занимал около 9 часов. Использование 10 параллельных процессов позволило сократить это время в несколько раз. 

В итоге мы собрали около 500 качественных пар снимков, из которых извлекли примерно 185 тысяч качественных патчей. 15 пар исключили — в них не удалось добиться приемлемой синхронизации или геометрического совпадения.

Заключение

В первом материале мы рассказали, зачем через восемь лет оживлять проект DPED, как именно мы его реанимировали и какого результата добились «в боевых условиях». В этой статье описали сбор и подготовку датасета, который планируем дополнить тысячей пар изображений. В следующий раз приступим к обучению модели, а заодно разберемся, как мы переписывали код и адаптировали его под современные версии библиотек. 

А еще мы обновили наш проект, пока готовился этот материал: переписали алгоритм сбора изображений, добавили автоматическую сортировку и сохранение снимков, а также перешли на SCRCPY. Подробнее расскажем в третьей статье.