OpenCV实现两幅图片全景拼接

课程作业的一个题目，找了代码加了注释。

import numpy as np
import cv2


class Stitcher:

    def stitch(self, images, ratio=0.75, reprojThresh=4.0,
        showMatches=False):

        # 检测出关键点，局部不变描述符
        (imageB, imageA) = images
        (kpsA, featuresA) = self.detectAndDescribe(imageA)
        (kpsB, featuresB) = self.detectAndDescribe(imageB)

        print("关键点个数",len(kpsA),len(kpsB))

        # 特征匹配
        M = self.matchKeypoints(kpsA, kpsB,
            featuresA, featuresB, ratio, reprojThresh)

        # 如果特征匹配返回None
        if M is None:
            return None

        # 将图像粘合在一起
        (matches, H, status) = M
        # 根据单应性矩阵进行矫正图片
        result = cv2.warpPerspective(imageA, H,
            (imageA.shape[1] + imageB.shape[1], imageA.shape[0]))
        # imageA.shape[1]=400,imageB.shape[1]=400,imageA.shape[0]=533
        # result.shape[0]=533,result.shape[1]=800
        result[0:imageB.shape[0], 0:imageB.shape[1]] = imageB


        # 显示匹配线
        if showMatches:
            vis = self.drawMatches(imageA, imageB, kpsA, kpsB, matches,
                status)
            return (result, vis)

        # 单独返回一个
        return result


    #接收照片，检测关键点和提取局部不变特征
    #用到了高斯差分（Difference of Gaussian (DoG)）关键点检测，和SIFT特征提取
    #detectAndCompute方法用来处理提取关键点和特征
    #返回一系列的关键点
    def detectAndDescribe(self, image):
        # 将图片转化为灰度图像
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

        # 提取特征点
        descriptor = cv2.xfeatures2d.SIFT_create()
        (kps, features) = descriptor.detectAndCompute(image, None)

        # print(kps) # 关键点
        print(features.shape[0],features.shape[1]) # 长度为128维的特征向量

        # 将关键点的座标pt存入numpy
        kps = np.float32([kp.pt for kp in kps])

        return (kps, features)


    #matchKeypoints方法需要四个参数，第一张图片的关键点和特征向量，第二张图片的关键点特征向量。
    #David Lowe’s ratio测试变量和RANSAC重投影门限也应该被提供。
    def matchKeypoints(self, kpsA, kpsB, featuresA, featuresB,
        ratio, reprojThresh):

        matcher = cv2.DescriptorMatcher_create("BruteForce")
        rawMatches = matcher.knnMatch(featuresA, featuresB, 2) # 最近邻算法设置K=2
        matches = []
        # for m in rawMatches:
        #     print(m[0].distance,m[1].distance)

        print("------------------------------")
        # 循环遍历匹配点
        for m in rawMatches:
            # Lowe’s ratio测试，用来确定高质量的特征匹配
            if len(m) == 2 and m[0].distance < m[1].distance * ratio:
                # 将第一张图像的下标值和第二张图像的下标值存入
                matches.append((m[0].trainIdx, m[0].queryIdx))

        # print(matches)
        # print(len(matches))

        # 将标注位置存入numpy
        if len(matches) > 4:
            # construct the two sets of points
            ptsA = np.float32([kpsA[i] for (_, i) in matches])
            ptsB = np.float32([kpsB[i] for (i, _) in matches])

            # 计算单应性矩阵
            # 其中H为求得的单应性矩阵矩阵
            # status则返回一个列表来表征匹配成功的特征点。
            (H, status) = cv2.findHomography(ptsA, ptsB, cv2.RANSAC,
                reprojThresh)

            return (matches, H, status)


        return None


    #连线画出两幅图的匹配
    def drawMatches(self, imageA, imageB, kpsA, kpsB, matches, status):
        # initialize the output visualization image
        (hA, wA) = imageA.shape[:2]
        (hB, wB) = imageB.shape[:2]
        # 三通道照片
        vis = np.zeros((max(hA, hB), wA + wB, 3), dtype="uint8")
        vis[0:hA, 0:wA] = imageA
        vis[0:hB, wA:] = imageB


        for ((trainIdx, queryIdx), s) in zip(matches, status):
            if s == 1:
                ptA = (int(kpsA[queryIdx][0]), int(kpsA[queryIdx][1]))
                ptB = (int(kpsB[trainIdx][0]) + wA, int(kpsB[trainIdx][1]))
                cv2.line(vis, ptA, ptB, (0, 255, 0), 1)

        return vis


if __name__ == '__main__':
# 加载图片
    imageA = cv2.imread('./hw2/building_02.jpg')
    imageB = cv2.imread('./hw2/building_03.jpg')

# 调整图片宽度
#     imageA = imutils.resize(imageA, width=400)
#     imageB = imutils.resize(imageB, width=400)

# showMatches=True 展示两幅图像特征的匹配,返回vis
    stitcher = Stitcher()
    (result, vis) = stitcher.stitch([imageA, imageB], showMatches=True)


    # vis = imutils.resize(imageA, width=800,height=800)
    # result = imutils.resize(imageB, width=800,height=800)

    cv2.imwrite('./vis1.jpg', vis)
    cv2.imwrite('./result.jpg', result)