Registro de imagem baseado em SuperPoint e SuperGlue

O registro da imagem é realizado com base em SuperPoint e SuperGlue. O endereço do projeto é https://github.com/magicleap/SuperGluePretrainedNetwork. O operador especial grid_sample é usado. Ao converter para onnx, opset_version deve ser 16 ou superior (isto é , a versão pytorch é 1.9 ou superior). O modelo SuperPoint é usado para extrair os pontos de característica e os descritores de ponto de característica da imagem (dois precisam ser executados durante o registro da imagem para realizar a extração dos pontos de característica das duas imagens), e o modelo SuperGlue é usado para extrair a característica pontos e recursos extraídos pelos descritores do modelo SuperPoint são combinados.

Use SuperPoint e SuperGlue para treinar seu próprio banco de dados, você pode visualizar as informações da biblioteca https://download.csdn.net/download/a486259/87471980, este documento registra em detalhes o uso de projetos pytorch-superpoint e pytorch-superglue para treinar seus próprios dados definem o processo de.

1. Pré-operação

Para implementar o modelo, ele pode ser implantado onnx, e alguns códigos do projeto podem ser modificados. O principal é deletar o uso do objeto dict no código, pois o onnx não suporta.

1.1 modificação de superponto

O código está em models/superpoint.py, que modifica principalmente a função forward do modelo SuperPoint (o código começa na linha 145), e não usa objetos de dicionário como parâmetros (valor de entrada e valor de saída), para evitar no operador onnx. Ao mesmo tempo, várias tentativas foram feitas para implementar a função de pontos-chave. Entre eles, o modelo SuperPoint apenas gera a confiança do ponto de coordenada (scores1) e o descritor do ponto de coordenada (descritores1) durante o treinamento.As coordenadas aqui realmente se referem aos pontos de recurso. No entanto, as informações de coordenadas são refletidas apenas nos dados da grade e as coordenadas no formato xy são necessárias para a correspondência de pontos. Para esse fim, as coordenadas dos pontos cujo valor de confiança é maior que o limite em pontuações são extraídas, portanto, keypoints1 (pontos de coordenadas) são obtidos.

    def forward(self, data):
        """ Compute keypoints, scores, descriptors for image """
        # Shared Encoder
        x = self.relu(self.conv1a(data))
        x = self.relu(self.conv1b(x))
        x = self.pool(x)
        x = self.relu(self.conv2a(x))
        x = self.relu(self.conv2b(x))
        x = self.pool(x)
        x = self.relu(self.conv3a(x))
        x = self.relu(self.conv3b(x))
        x = self.pool(x)
        x = self.relu(self.conv4a(x))
        x = self.relu(self.conv4b(x))

        # Compute the dense keypoint scores
        cPa = self.relu(self.convPa(x))
        scores = self.convPb(cPa)
        scores = torch.nn.functional.softmax(scores, 1)[:, :-1]
        b, _, h, w = scores.shape
        scores = scores.permute(0, 2, 3, 1).reshape(b, h, w, 8, 8)
        scores = scores.permute(0, 1, 3, 2, 4).reshape(b, h*8, w*8)
        scores = simple_nms(scores, self.config['nms_radius'])

        # Extract keypoints
        #keypoints = [ torch.nonzero(s > self.config['keypoint_threshold']) for s in scores]## nonzero->tensorRT not support
        #keypoints = [torch.vstack(torch.where(s > self.config['keypoint_threshold'])).T for s in scores]## vstack->onnx not support
        #keypoints = [torch.cat(torch.where(s > self.config['keypoint_threshold']),0).reshape(len(s.shape),-1).T for s in scores]# tensor.T->onnx not support
        #keypoints = [none_zero_index(s,self.config['keypoint_threshold']) for s in scores]# where->nonzero ->tensorRT not support
        keypoints = [torch.transpose(torch.cat(torch.where(s>self.config['keypoint_threshold']),0).reshape(len(s.shape),-1),1,0) for s in scores]# transpose->tensorRT not support
        scores = [s[tuple(k.t())] for s, k in zip(scores, keypoints)]

        # Discard keypoints near the image borders
        keypoints, scores = list(zip(*[
            remove_borders(k, s, self.config['remove_borders'], h*8, w*8)
            for k, s in zip(keypoints, scores)]))

        # Keep the k keypoints with highest score
        if self.config['max_keypoints'] >= 0:
            keypoints, scores = list(zip(*[
                top_k_keypoints(k, s, self.config['max_keypoints'])
                for k, s in zip(keypoints, scores)]))

        # Convert (h, w) to (x, y)
        keypoints = [torch.flip(k, [1]).float() for k in keypoints]

        # Compute the dense descriptors
        cDa = self.relu(self.convDa(x))
        descriptors = self.convDb(cDa)
        descriptors = torch.nn.functional.normalize(descriptors, p=2, dim=1)

        # Extract descriptors
        descriptors = [sample_descriptors(k[None], d[None], 8)[0]
                       for k, d in zip(keypoints, descriptors)]

        # return {
    
    
        #     'keypoints': keypoints,
        #     'scores': scores,
        #     'descriptors': descriptors,
        # }
        return  keypoints[0].unsqueeze(0),scores[0].unsqueeze(0),descriptors[0].unsqueeze(0)

1.2 Modificação do SuperGlue

Em models/superglue.py no código, corrija principalmente a imagem depois que o objeto do dicionário for excluído no superponto.

1.2.1 ajuste da função normalize_keypoints

Substitua o parâmetro image_shape da função original por altura e largura

def normalize_keypoints(kpts,  height, width):
    """ Normalize keypoints locations based on image image_shape"""
    one = kpts.new_tensor(1)
    size = torch.stack([one*width, one*height])[None]
    center = size / 2
    scaling = size.max(1, keepdim=True).values * 0.7
    return (kpts - center[:, None, :]) / scaling[:, None, :]

1.2.2 modificação da função forword

Substitua a função de encaminhamento do SuperGlue no código pelo código a seguir. O principal é modificar os parâmetros de entrada, descompactar o dicionário anterior, de modo que um parâmetro se torne 6; e modificar o valor de retorno da função e fixar o tamanho da imagem em 640*640 ao mesmo tempo .
O modelo SuperGlue gera dois conjuntos de informações de match_indices e match_mscores com base em dois conjuntos de pontos-chave de entrada, pontuações e dados de descritores. O primeiro grupo é usado para descrever a correspondência entre A->B, e o segundo grupo é usado para descrever a correspondência entre B->A.

    def forward(self, data_descriptors0, data_descriptors1, data_keypoints0, data_keypoints1, data_scores0, data_scores1):
        """Run SuperGlue on a pair of keypoints and descriptors"""
        #, height:int, width:int
        height, width=640,640
        desc0, desc1 = data_descriptors0, data_descriptors1
        kpts0, kpts1 = data_keypoints0, data_keypoints1

        if kpts0.shape[1] == 0 or kpts1.shape[1] == 0:  # no keypoints
            shape0, shape1 = kpts0.shape[:-1], kpts1.shape[:-1]
            return kpts0.new_full(shape0, -1, dtype=torch.int),kpts1.new_full(shape1, -1, dtype=torch.int),kpts0.new_zeros(shape0),kpts1.new_zeros(shape1)

        # Keypoint normalization.
        kpts0 = normalize_keypoints(kpts0, height, width)
        kpts1 = normalize_keypoints(kpts1, height, width)

        # Keypoint MLP encoder.
        desc0 = desc0 + self.kenc(kpts0, data_scores0)
        desc1 = desc1 + self.kenc(kpts1, data_scores1)

        # Multi-layer Transformer network.
        desc0, desc1 = self.gnn(desc0, desc1)

        # Final MLP projection.
        mdesc0, mdesc1 = self.final_proj(desc0), self.final_proj(desc1)

        # Compute matching descriptor distance.
        scores = torch.einsum('bdn,bdm->bnm', mdesc0, mdesc1)
        scores = scores / self.config['descriptor_dim']**.5

        # Run the optimal transport.
        scores = log_optimal_transport(
            scores, self.bin_score,
            iters=self.config['sinkhorn_iterations'])

        # Get the matches with score above "match_threshold".
        max0, max1 = scores[:, :-1, :-1].max(2), scores[:, :-1, :-1].max(1)
        indices0, indices1 = max0.indices, max1.indices
        mutual0 = arange_like(indices0, 1)[None] == indices1.gather(1, indices0)
        mutual1 = arange_like(indices1, 1)[None] == indices0.gather(1, indices1)
        zero = scores.new_tensor(0)
        mscores0 = torch.where(mutual0, max0.values.exp(), zero)
        mscores1 = torch.where(mutual1, mscores0.gather(1, indices1), zero)
        valid0 = mutual0 & (mscores0 > self.config['match_threshold'])
        valid1 = mutual1 & valid0.gather(1, indices1)
        indices0 = torch.where(valid0, indices0, indices0.new_tensor(-1))
        indices1 = torch.where(valid1, indices1, indices1.new_tensor(-1))
        # return {
    
    
        #     'matches0': indices0, # use -1 for invalid match
        #     'matches1': indices1, # use -1 for invalid match
        #     'matching_scores0': mscores0,
        #     'matching_scores1': mscores1,
        # }
        return indices0,  indices1,  mscores0,  mscores1

1.3 Chamada integrada

Ao realizar o registro da imagem, o fluxo de processamento de dados usando o modelo superpoint e o modelo superglue é fixo. Para simplificar o código, ele é empacotado como um modelo e o código é salvo como SPSP.py.

import torch
from superglue import SuperGlue
from superpoint import SuperPoint
import torch
import torch.nn as nn
import torch.nn.functional as F
class SPSG(nn.Module):#
    def __init__(self):
        super(SPSG, self).__init__()
        self.sp_model = SuperPoint({
    
    'max_keypoints':700})
        self.sg_model = SuperGlue({
    
    'weights': 'outdoor'})
    def forward(self,x1,x2):
        keypoints1,scores1,descriptors1=self.sp_model(x1)
        keypoints2,scores2,descriptors2=self.sp_model(x2)
        #print(scores1.shape,keypoints1.shape,descriptors1.shape)
        #example=(descriptors1.unsqueeze(0),descriptors2.unsqueeze(0),keypoints1.unsqueeze(0),keypoints2.unsqueeze(0),scores1.unsqueeze(0),scores2.unsqueeze(0))
        example=(descriptors1,descriptors2,keypoints1,keypoints2,scores1,scores2)
        indices0,  indices1,  mscores0,  mscores1=self.sg_model(*example)
        #return indices0,  indices1,  mscores0,  mscores1
        matches = indices0[0]
        
        valid = torch.nonzero(matches > -1).squeeze().detach()
        mkpts0 = keypoints1[0].index_select(0, valid);
        mkpts1 = keypoints2[0].index_select(0, matches.index_select(0, valid));
        confidence = mscores0[0].index_select(0, valid);
        return mkpts0, mkpts1, confidence

1.4 Biblioteca de processamento de imagem

O código para operações de leitura e exibição de imagens é encapsulado como a biblioteca imgutils, para obter detalhes, consulte https://hpg123.blog.csdn.net/article/details/124824892

2. Realize o registro da imagem

2.1 Obtenha pontos correspondentes

Através das etapas a seguir, os pontos característicos das duas imagens e o grau de correspondência dos pontos característicos podem ser obtidos

from imgutils import *
import torch
from SPSG import SPSG
model=SPSG()#.to('cuda')
tensor2a,img2a=read_img_as_tensor("b1.jpg",(640,640),device='cpu')
tensor2b,img2b=read_img_as_tensor("b4.jpg",(640,640),device='cpu')
mkpts0, mkpts1, confidence=model(tensor2a,tensor2b)
myimshows( [img2a,img2b],size=12)

A saída da execução do código se parece com isso:

2.2 Plotagem de pontos correspondentes

O código a seguir pode desenhar e conectar os pontos nas duas imagens cujo grau de correspondência é maior que o limite

import cv2 as cv
pt_num = mkpts0.shape[0]
im_dst,im_res=img2a,img2b
img = np.zeros((max(im_dst.shape[0], im_res.shape[0]), im_dst.shape[1]+im_res.shape[1]+10,3))
img[:,:im_res.shape[0],]=im_dst
img[:,-im_res.shape[0]:]=im_res
img=img.astype(np.uint8)
match_threshold=0.6
for i in range(0, pt_num):
    if (confidence[i] > match_threshold):
        pt0 = mkpts0[i].to('cpu').numpy().astype(np.int)
        pt1 = mkpts1[i].to('cpu').numpy().astype(np.int)
        #cv.circle(img, (pt0[0], pt0[1]), 1, (0, 0, 255), 2)
        #cv.circle(img, (pt1[0], pt1[1]+650), (0, 0, 255), 2)
        cv.line(img, pt0, (pt1[0]+im_res.shape[0], pt1[1]), (0, 255, 0), 1)
myimshow( img,size=12)

2.3 Interceptar a área de sobreposição

Primeiro, chame a função getGoodMatchPoint para filtrar os pontos de recurso com alto grau de correspondência de acordo com o limite, depois calcule e perspetive a matriz de alteração H e, finalmente, extraia a área de sobreposição

import cv2
def getGoodMatchPoint(mkpts0, mkpts1, confidence,  match_threshold:float=0.5):
    n = min(mkpts0.size(0), mkpts1.size(0))
    srcImage1_matchedKPs, srcImage2_matchedKPs=[],[]

    if (match_threshold > 1 or match_threshold < 0):
        print("match_threshold error!")

    for i in range(n):
        kp0 = mkpts0[i]
        kp1 = mkpts1[i]
    
        pt0=(kp0[0].item(),kp0[1].item());
        pt1=(kp1[0].item(),kp1[1].item());
        c = confidence[i].item();
        if (c > match_threshold):
            srcImage1_matchedKPs.append(pt0);
            srcImage2_matchedKPs.append(pt1);
    
    return np.array(srcImage1_matchedKPs),np.array(srcImage2_matchedKPs)
pts_src, pts_dst=getGoodMatchPoint(mkpts0, mkpts1, confidence)

h1, status = cv2.findHomography(pts_src, pts_dst, cv.RANSAC, 8)
im_out1 = cv2.warpPerspective(im_dst, h1, (im_dst.shape[1],im_dst.shape[0]))
im_out2 = cv2.warpPerspective(im_res, h1, (im_dst.shape[1],im_dst.shape[0]),16)
#这里 im_dst和im_out1是严格配准的状态
myimshowsCL([im_dst,im_out1,im_res,im_out2],rows=2,cols=2, size=6)

2.4 Exportação do modelo

Use o código a seguir para obter a exportação do modelo

input_names = ["input1","input2"]
output_names = ['mkpts0', 'mkpts1', 'confidence']
dummy_input=(tensor2a,tensor2b)
example_outputs=model(tensor2a,tensor2b)
ONNX_name="model.onnx"
torch.onnx.export(model.eval(), dummy_input, ONNX_name, verbose=True, input_names=input_names,opset_version=16,
                  dynamic_axes={
    
    
                        'confidence': {
    
    0: 'point_num',},
                        'mkpts0': {
    
    0: 'batch_size',},
                        'mkpts1': {
    
    0: 'batch_size',}
                       },
                   output_names=output_names)#,example_outputs=example_outputs

3. Use o superponto sozinho

Os pontos característicos da imagem podem ser extraídos usando apenas o modelo SuperPoint

from imgutils import *
import torch
from superpoint import SuperPoint
import cv2

config={
    
    'max_keypoints': 400,'keypoint_threshold':0.1}
sp_model=SuperPoint(config).to('cuda')
sp_model=sp_model.eval()

tensor2a,img2a=read_img_as_tensor(r"D:\SuperGluePretrainedNetwork-master\assets\freiburg_sequence\1341847986.762616.png",(640,640),device='cuda')
tensor2b,img2b=read_img_as_tensor(r"D:\SuperGluePretrainedNetwork-master\assets\freiburg_sequence\1341847987.758741.png",(640,640),device='cuda')

keypoints1,scores1,descriptors1=sp_model(tensor2a)
keypoints2,scores2,descriptors2=sp_model(tensor2b)

yanse=(0,0,255)
points=keypoints1[0].int().cpu().numpy()
for i in range(len(points)):
    X,Y=points[i]
    cv2.circle(img2a,(X,Y),3,yanse,2)

points=keypoints2[0].int().cpu().numpy()
for i in range(len(points)):
    X,Y=points[i]
    cv2.circle(img2b,(X,Y),3,yanse,2)
myimshows([img2a,img2b])