duplicate_image_detection_system/codes/feature.py

from __future__ import print_function
import cv2
import numpy as np
import os
from flask import render_template
import zipfile

MAX_MATCHES = 500
GOOD_MATCH_PERCENT = 0.15



class feat:
    def call_feature_extraction_1(self, folder_path, list, index, tmp):
        # 按相似度排序
        list = sorted(list, key=(lambda x: [x[2], x[5]]),reverse=True)
        print(list)
        print('---------------')
        print(index)

        for i in range(int(index/10)):
            # 取第一组比较并返回
            refFilename = list[i][0]
            # imgname1 = '/home/Jupyterlab/wrl/pic/xiagao/pic/11.jpeg'
            print("Reading reference image : ", refFilename)
            imReference = cv2.imread(refFilename, cv2.IMREAD_COLOR)

            imFilename = list[i][1]
            # imgname2 = '/home/Jupyterlab/wrl/pic/xiagao/pic/12.jpeg'
            print("Reading image to align : ", imFilename);
            im = cv2.imread(imFilename, cv2.IMREAD_COLOR)

            print(refFilename)
            print(imFilename)
            print(folder_path[7:])
            p1 = refFilename.rfind('/')
            name1 = refFilename[p1:-4]
            print(name1)
            p2 = imFilename.rfind('/')
            name2 = imFilename[p2 + 1:]
            print(name2)

            # Write aligned image to disk.
            outFilename = "output/" + folder_path[7:]
            pre = os.getcwd()
            print("Saving aligned image : ", outFilename)

            print("Aligning images ...")
            # Registered image will be resotred in imReg.
            # The estimated homography will be stored in h.
            imReg, h, img5 = feat().alignImages(im, imReference)
            if (str(img5) == 'white'):
                print('white')
                # continue
            else:
                print(outFilename)
                # imgwrite需要建好路径！！！！！！
                if not os.path.exists(outFilename):
                    os.makedirs(outFilename)
                outFilename1 = outFilename + name1 + name2
                cv2.imwrite(outFilename1, img5)
                # continue

        outFullName = feat.zipDir(outFilename)
        pre1 = "/home/wwwroot/default/" + outFilename + ".zip"
        # return str(pre) + '/' + outFilename
        # return pre1
        return outFullName




    def alignImages(self, im1, im2):
        # Convert images to grayscale
        im1Gray = cv2.cvtColor(im1, cv2.COLOR_BGR2GRAY)
        im2Gray = cv2.cvtColor(im2, cv2.COLOR_BGR2GRAY)

        # Detect ORB features and compute descriptors.
        orb = cv2.ORB_create(MAX_MATCHES)
        keypoints1, descriptors1 = orb.detectAndCompute(im1Gray, None)
        keypoints2, descriptors2 = orb.detectAndCompute(im2Gray, None)

        if keypoints1  and keypoints2:
            # Match features.
            matcher = cv2.DescriptorMatcher_create(cv2.DESCRIPTOR_MATCHER_BRUTEFORCE_HAMMING)
            matches = matcher.match(descriptors1, descriptors2, None)

            # Sort matches by score
            matches.sort(key=lambda x: x.distance, reverse=False)

            # Remove not so good matches
            numGoodMatches = int(len(matches) * GOOD_MATCH_PERCENT)
            matches = matches[:numGoodMatches]

            # Draw top matches
            imMatches = cv2.drawMatches(im1, keypoints1, im2, keypoints2, matches, None)
            cv2.imwrite("matches.jpg", imMatches)

            # Extract location of good matches
            points1 = np.zeros((len(matches), 2), dtype=np.float32)
            points2 = np.zeros((len(matches), 2), dtype=np.float32)

            for i, match in enumerate(matches):
                points1[i, :] = keypoints1[match.queryIdx].pt
                points2[i, :] = keypoints2[match.trainIdx].pt

            if (points1.size == 0) or (points2.size == 0):
                return 'white', 'white', 'white'

            # Find homography
            h, mask = cv2.findHomography(points1, points2, cv2.RANSAC)

            # Use homography
            height, width, channels = im2.shape
            im1Reg = cv2.warpPerspective(im1, h, (width, height))
            img1 = cv2.drawKeypoints(im1, keypoints1, im1, color=(255, 0, 255))
            img2 = cv2.drawKeypoints(im2, keypoints2, im2, color=(255, 0, 255))
            img5 = cv2.drawMatches(img1, keypoints1, img2, keypoints2, matches, None, flags=2)

            return im1Reg, h, img5

        else:
            return 'white', 'white', 'white'

    def return_img_stream(self, img_local_path):
        # """
        # 工具函数:
        # 获取本地图片流
        # :param img_local_path:文件单张图片的本地绝对路径
        # :return: 图片流
        # """
        import base64
        img_stream = ''
        with open(img_local_path, 'r') as img_f:
            img_stream = img_f.read()
            img_stream = base64.b64encode(img_stream)
        return img_local_path



    def zipDir(dirpath):
        """
        压缩指定文件夹
        :param dirpath: 目标文件夹路径
        :param outFullName: 压缩文件保存路径+xxxx.zip
        :return: 无
        """
        outFullName = "/home/wwwroot/default/" + dirpath[21:] + ".zip"
        outFullName1 = "http://106.75.226.23/"  + dirpath[21:] + ".zip"
        # if not os.path.exists(outFullName):
        #     os.makedirs(outFullName)
        zip = zipfile.ZipFile(outFullName, "w", zipfile.ZIP_DEFLATED)
        for path, dirnames, filenames in os.walk(dirpath):
            # 去掉目标跟路径，只对目标文件夹下边的文件及文件夹进行压缩
            fpath = path.replace(dirpath, '')

            for filename in filenames:
                zip.write(os.path.join(path, filename), os.path.join(fpath, filename))
        zip.close()
        return outFullName1



# im1 = '/Users/wrl/Desktop/test0506/pic/2303.png'
# im2 = '/Users/wrl/Desktop/test0506/pic/2304.png'
# img1 = cv2.imread(im1, cv2.IMREAD_COLOR)
# img2 = cv2.imread(im2, cv2.IMREAD_COLOR)
# imReg, h, img5 =feat().alignImages(img1,img2)