Время прочтения: 4 мин.

В продолжении статьи «Kornia — python библиотека для обработки изображений в задачах CV», хотел бы познакомить читателей с еще одной интересной функцией данной библиотеки.

Речь идет о функции LoFTR из состава модуля kornia.feature, которая предназначена для поиска одинаковых локальных объектов на разных изображениях. В основе инструмента лежат нейросетевые технологии, предобученные модели есть в открытом доступе. Ниже приведен пример работы данного алгоритма.

На вход поданы два следующих изображения:

            Результат сравнения:

Код для сравнения двух изображений и отрисовки результата:

import matplotlib.pyplot as plt
import cv2
import kornia as K
import kornia.feature as KF
import numpy as np
import torch
from kornia_moons.feature import *
import torchvision


def load_torch_image(fname):
    img = K.image_to_tensor(cv2.imread(fname), False).float() / 255
    img = K.color.bgr_to_rgb(img)
    return img
    

fname1 = 'test1.png'
fname2 = 'test2.png'

img1 = load_torch_image(fname1)
img2 = load_torch_image(fname2)

matcher = KF.LoFTR(pretrained='outdoor')

input_dict = {'image0': K.color.rgb_to_grayscale(img1),
              'image1': K.color.rgb_to_grayscale(img2)}

with torch.no_grad():
    correspondences = matcher(input_dict)

for k, v in correspondences.items():
    print(k)

mask = correspondences['confidence'] > 0.9
indices = torch.nonzero(mask)
correspondences['confidence'] = correspondences['confidence'][indices]
correspondences['keypoints0'] = correspondences['keypoints0'][indices]
correspondences['keypoints1'] = correspondences['keypoints1'][indices]
correspondences['batch_indexes'] = correspondences['batch_indexes'][indices]

mkpts0 = correspondences['keypoints0'].cpu().numpy()
mkpts1 = correspondences['keypoints1'].cpu().numpy()
H, inliers = cv2.findFundamentalMat(mkpts0, mkpts1, cv2.USAC_MAGSAC, 0.5, 0.999, 100000)
inliers = inliers > 0


#Функция отрисовки совпадений
draw_LAF_matches(
    KF.laf_from_center_scale_ori(torch.from_numpy(mkpts0).view(1,-1, 2),
                                torch.ones(mkpts0.shape[0]).view(1,-1, 1, 1),
                                torch.ones(mkpts0.shape[0]).view(1,-1, 1)),

    KF.laf_from_center_scale_ori(torch.from_numpy(mkpts1).view(1,-1, 2),
                                torch.ones(mkpts1.shape[0]).view(1,-1, 1, 1),
                                torch.ones(mkpts1.shape[0]).view(1,-1, 1)),
    torch.arange(mkpts0.shape[0]).view(-1,1).repeat(1,2),
    K.tensor_to_image(img1),
    K.tensor_to_image(img2),
    inliers,
    draw_dict={'inlier_color': (0.2, 1, 0.2),
               'tentative_color': None,
               'feature_color': (0.2, 0.2, 1), 'vertical': False})

plt.axis('off')
plt.savefig('matches1.png')
plt.show()

Разберем код. В самом начале производится чтение изображений с диска, преобразование их в torch-массивы и смена цветовой модели (BGR на RGB). Затем инициализируется модель, создается словарь, значениям которого присваиваются изображения в цветовой модели grayscale (модель не может работать с цветными изображениями). Этот словарь подается на вход модели, она возвращает массив correspondences, в котором перечисляются координаты объектов первого и второго изображения, признанных моделью похожими, значения «уверенности» модели в похожести областей находятся в столбце confidence.

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

Еще пример работы функции:

Обратите внимание, перед началом отрисовки совпадений, выполняется функция из состава OpenCV – findFundamentalMat(), которая служит дополнительным фильтром совпадающих точек. На этом этапе вычисляются варианты искажения перспективы в связи со сменой ракурса съемки, отбирается наилучший вариант и на основании лучшего формируется массив numpy с булевыми значениями, указывающими обосновано ли перемещение координат точки на разных изображениях сменой ракурса. На втором результирующем изображении, точки, отсеянные данной функцией, отмечены синим цветом. Эту дополнительную проверку можно не выполнять, для этого в аргументы функции draw_LAF_matches() нужно передать массив со значениями True, длиной равной длине массива correspondences.

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