以图搜图(粗略图搜原图)(本机)

陈东豪

from tkinter import *
import cv2
from webbrowser import open as webopen
import numpy as np
from threading import Thread
from multiprocessing.dummy import Pool
from os import listdir, mkdir, remove, startfile, getcwd
from os.path import isdir, isfile, join, exists
from tkinter import filedialog
from windnd import hook_dropfiles
from PIL import Image, ImageTk
from tkinter.messagebox import showinfo, askyesno
from shutil import copy
from skimage import io
from pickle import dump, load
from time import time

class APP:

    def __init__(self):
        self.root = Tk()
        self.root.geometry('450x430')
        self.root.title('星猫搜图~')
        self.imgnum = 0     #用于判断文件夹内图片个数
        self.if_change = 0   #判断图片数量是否变化
        self.file = ''
        self.hashdic = {}       
        self.img11 = 0
        self.num111 = 0
        self.badimg = 0
        self.badimglist = []
        self.msg1 = StringVar(value=r'C:\Users\陈东豪\Desktop')
        self.if_top123 = 0
        self.filelist = []
        self.initimg = []
        self.pk_file = 'hashֵ.pk'
        if not isfile(self.pk_file):
            with open(self.pk_file, 'wb') as (f):
                dump(self.hashdic, f)
        with open(self.pk_file, 'rb') as (b):
            self.hashdic = load(b)

    def wait(self):
        """等候时打开页面,显示加载过程"""
        self.wait1 = Label((self.root), text='请稍等...')

    def main(self):
        Label((self.root), text='img文件路径:').grid(row=0, column=0)
        self.hostVar = StringVar(value=r'C:\Users\Administrator\Desktop')
        self.numVar = StringVar(self.root, str(len(self.filelist)) + '张')
        self.filename = Entry((self.root), textvariable=(self.hostVar))
        self.filename.grid(row=0, column=1)
        
        self.bt = Button((self.root), text='确定', command=self.change_lujing)
        self.bt.grid(row=0, column=2)
        self.lb4 = Label((self.root), textvariable=(self.numVar))
        self.lb4.grid(row=0, column=3)
        self.bt2 = Button((self.root), text='打开目标img文件夹', command=(self.open_lujing))
        self.bt2.grid(row=1, column=0)
        self.bt4 = Button((self.root), text='此次查找详情', command=lambda:self.show_init(self.initimg))
        self.bt4.grid(row=1, column=1)
        self.bt5 = Button((self.root), text='打开星猫图片', command=(self.open_files))
        self.bt5.grid(row=1, column=3)
        self.bt6 = Button((self.root), text='查看损坏img', command=(self.see_badimg))
        self.bt6.grid(row=1, column=2)
        self.balloon = Label(self.root)

        self.bt.bind('<Button-3>', self.chose_dir)
        self.bt.bind('<Enter>', lambda event: self.Balloon_show(event, msg='单击鼠标右键打开文件选择框!', fg='blue'))
        self.bt5.bind('<Enter>', lambda event: self.Balloon_show(event, msg='双击鼠标右键打开百度搜图\n[url]https://graph.baidu.com/pcpage/index?tpl_from=pc'[/url], fg='blue'))
        self.root.bind('<Leave>', self.Balloon_destroy)
        self.root.bind('<Double-Button-3>', self.baidu)
        hook_dropfiles(self.root, func=self.dragged_files)   #当拖入文件时触发
        self.change_lujing()
        self.display_num()
        self.root.mainloop()
        
        
        
    def show_init(self,filename):
        '''展示目标img信息'''
        file_msg = '\n'.join((item for item in filename))
        showinfo(title='提示', message=('总共有' + str(len(filename)) + '张相似图片!\n相似度为' +\
            str(100 - self.min_) + '\n' + file_msg + '\n损坏图片:' + str(self.badimg) + \
                '张\n共耗时:' + str(self.t) + 's'))
        
    def Balloon_destroy(self, event):
        """隐藏气泡ʾ"""
        try:
            self.balloon.place_forget()
        except:
            pass

    def Balloon_show(self, event, msg, fg=None, bg=None, font=('微软雅黑', 10, 'bold')):
        try:
            self.balloon.place_forget()
        except:
            pass
        else:
            self.balloon['fg'] = fg
            self.balloon['bg'] = bg
            self.balloon['text'] = msg
            self.balloon['font'] = font
            self.balloon.place(x=35, y=190)

    def baidu(self, event):
        """百度搜图"""
        webopen('https://graph.baidu.com/pcpage/index?tpl_from=pc')

    def chose_dir(self, event):
        file = filedialog.askopenfilename()
        if file != '':
            file1 = file.rsplit('/', 1)[0]
            self.hostVar.set(file1)
            self.change_lujing()

    def open_lujing(self, *args):
        """打开目标img文件夹"""
        if len(self.initimg) == 0:
            showinfo(title='提示', message='无目标img!')
        else:
            files = []
            for each in self.initimg:
                files.append(each.rsplit('\\', 1)[0])
            else:
                files = list(set(files))
                for each in files:
                    startfile(each)

    def open_files(self, *args):
        """打开文件夹"""
        fname = getcwd() + '//星猫图片'
        startfile(fname)


    def see_badimg(self, *args):
        """损坏badimg"""
        if len(self.badimglist) == 0:
            showinfo(title='提示', message='没有badimg!')
        else:
            msg = '\n'.join(self.badimglist)
            showinfo('badimg路径', msg)
            if askyesno('警告', '是否删除损坏imgs?'):
                for each in self.badimglist:
                    try:
                        remove(each)
                        self.badimglist.remove(each)
                    except Exception as e:
                        showinfo(title='提示', message=e)


    def display_num(self):
        """提示图片数量, 判断图片个数是否不再改变"""

        num = len(self.filelist)
        self.numVar.set(str(num) + '张')
        if self.imgnum!=num:
            self.imgnum = num
            self.if_change = 1
        else:
            if self.if_change:
                #如果变化停止则开始计算hash值
                self.if_change = 0
                self.pool_map(self.compute_hash,self.filelist)
                # showinfo('提示','开始计算图片hash值啦!')
        self.root.after(1000, self.display_num)

    def change_lujing(self):
        """更改搜图路径"""

        self.file = self.filename.get()
        if isdir(self.file):
            pass
        else:
            lujing = getcwd()
            self.hostVar.set(lujing)
            self.file = self.filename.get()
        self.filelist = []
        self.thread_it(self.getFileList, self.file, self.filelist)

    def pool_map(self, func, *args):
        """多进程"""
        pool = Pool()
        pool.map_async(func, args)
        pool.close()
        pool.join()

    def getFileList(self, dir, imglist, ext=['jpg', 'png', 'jpeg', 'bmp', 'webp']):
        """
        获取文件夹及其子文件夹中图片列表
        输入 dir:文件夹根目录
        输入 ext: 扩展名
        返回： 文件路径列表
        """
        newDir = dir
        if isfile(dir):
            if dir.rsplit('.', 1)[(-1)] in ext:
                imglist.append(dir)

        elif isdir(dir):
            try:
                for s in listdir(dir):
                    newDir = join(dir, s)
                    self.getFileList(newDir, imglist, ext)

            except Exception as e:
                    print(s,'\n',e)

        return imglist


    #打开指定的图片，缩放到指定尺寸
    def get_img(self,filename,width=445,height=365):
        try:
            im = Image.open(filename).resize((width,height))
            # 引用：添加一个Label，用来存储图片。使用PanedWindow也行。
            panel = Label(master=self.root)
            panel.photo = ImageTk.PhotoImage(im)  # 将原本的变量photo改为panel.photo
            return panel.photo
        except:
            return 0

    def compute_hash(self,imgname):
        '''专门用于计算大规模计算hash值'''
        if imgname not in self.hashdic:
            img2_ = self.cv_imread(imgname)
            hash2 = self.aHash(img2_, imgname)
            if hash2 != 0:
                self.hashdic[str(imgname)] = hash2

                
    def dragged_files(self, files):
        """拖拽图片"""
        sum_msg2 = [item.decode('gbk') for item in files]
        for sum_msg in sum_msg2:
            ext1 = sum_msg.rsplit('.', 1)[(-1)]
            try:
                self.img.destroy()
            except:
                pass
        
            if ext1 in ('jpg', 'png', 'jpeg', 'bmp', 'ico', 'webp'):
                if ext1 == 'ico':
                    showinfo('提示', '暂不支持ico格式图片搜索!')
                else:
                    image1 = self.get_img(sum_msg)
                    if image1 != 0:
                        self.img = Label((self.root), image=image1)
                        self.img.grid(row=2, columnspan=4)
                        self.imgdict = {}   #用来保存图片相似度
                        self.thread_it(self.Compare2, sum_msg)
                    else:
                        showinfo('提示',sum_msg+'\n该图片有问题!')
                        
            else:
                showinfo(title='提示', message=('暂不支持当前格式搜索!(' + str(ext1) + ')'))

    def Compare2(self, img):
        '''开始搜索相似图片ing'''
        t1 = time()
        img1 = self.cv_imread(img)
        self.hash1 = self.aHash(img1,img)
        if self.hash1 == 0:
            showinfo('提示', '此图片损坏,请更换其他图片!')
        else:
            # self.pool_map(self.compare, self.filelist)
            #使用多进程加速

            pool = Pool()
            pool.map_async(self.compare, self.filelist)
            pool.close()
            pool.join()

            try:

                self.min_ = min(self.imgdict.values())
                min_key = [i for i in self.imgdict.keys() if self.imgdict[i] == self.min_]
                self.initimg = min_key
                if not exists('星猫图片'):
                    mkdir('星猫图片')
                for each_img in min_key:
                    try:
                        copy(each_img, './/星猫图片')
                    except:
                        pass

                if self.min_ == 0:
                    file_img = []
                    file_img.append(min_key[0])
                    self.show_img1(file_img)
                else:
                    self.show_img1(min_key)
                with open(self.pk_file, 'wb') as (f):
                    dump(self.hashdic, f)
                t2 = time()
                self.t = t2 - t1

            except:
                showinfo('提示', '无相关图片!')

    def show_msg(self, title1='提示', msg='呱呱呱!', ms=0, fontbig=20):
        """显示消息"""
        if msg != self.msg1:
            self.msg1.set(msg)
        if not self.if_top123:
            self.top123 = Toplevel()
            self.top123.title(title1)
            self.top123.protocol('WM_DELETE_WINDOW', self.destroytop)
            msg1 = Label((self.top123), textvariable=(self.msg1), font=('微软雅黑', fontbig, 'bold'))
            msg1.pack()
        if ms != 0:
            self.top123.after(ms, self.destroytop)

    def destroytop(self):
        """关闭窗口"""
        self.top123.destroy()
        self.if_top123 = 0

    def show_img1(self, filename=[], ms=4500, num=0):
        """可切换显示多张照片"""
        try:
            if not self.img11:
                self.top11 = Toplevel()
                self.top11.title('星猫图片浏览~')
                self.top11.geometry('640x400')
                self.top11.protocol('WM_DELETE_WINDOW', lambda : self.closeimg(filename))
                self.img11 = 1
                self.img2 = Label((self.top11), image=(self.get_img(filename[num], 640, 400)))
                self.img2.pack()
            else:
                self.img2.destroy()
                self.img2 = Label((self.top11), image=(self.get_img(filename[num], 640, 400)))
                self.img2.pack()
            num += 1
            self.top11.after(3000, self.show_img1, filename, ms, num)
        except:
            self.closeimg(filename)

    def closeimg(self, filename):
        """窗口关闭时触发"""
        self.img11 = 0
        self.top11.destroy()

    def thread_it(self, func, *args):
        """多线程"""
        t = Thread(target=func, args=args)
        t.setDaemon(True)
        t.start()

    #差值感知算法
    def aHash(self,img,imgname,if_shishi=0):
        '''img:图片矩阵,imgname:图片名字,if_shishi是否'''
        try:
            #缩放8*9
            img1=cv2.resize(img,(9,8),interpolation=cv2.INTER_CUBIC)
            #转换灰度图
            gray=cv2.cvtColor(img1,cv2.COLOR_BGR2GRAY)
            hash_str=''
            #每行前一个像素大于后一个像素为1，相反为0，生成哈希
            for i in range(8):
                for j in range(8):
                    if   gray[i,j]>gray[i,j+1]:
                        hash_str=hash_str+'1'
                    else:
                        hash_str=hash_str+'0'
            return hash_str
        except:
            '''如果打开图片失败，就重新保存试一试'''
            if not if_shishi:
                '''如果是第一次失败，则重试'''
                try:
                    image = io.imread(imgname)
                    image = cv2.cvtColor(image, cv2.COLOR_RGBA2BGRA)
                    cv2.imencode(imgname.rsplit('.',1)[-1],image)[1].tofile(imgname)
                    return self.aHash(image,imgname,1)
                except:
                    print(imgname,'打开失败!')
                    self.badimg+=1
                    return 0
            else:
                '''表示获取hash失败!'''
                return 0
 
    def cmpHash(self, hash1, hash2):
        n = 0
        if len(hash1) != len(hash2):
            return -1
        for i in range(len(hash1)):
            if hash1[i] != hash2[i]:
                n = n + 1
        return n

    def cv_imread(self, filePath):
        cv_img = cv2.imdecode(np.fromfile(filePath, dtype=(np.uint8)), -1)
        return cv_img

    def compare(self, img2):
        '''比较图片相似度并保存相似度n1'''
        imgname = img2
        try:
            if str(imgname) in self.imgdict:
                pass
            elif str(imgname) not in self.hashdic:
                img2_ = self.cv_imread(img2)
                hash2 = self.aHash(img2_, imgname=img2)
                if hash2 != 0:
                    self.hashdic[str(imgname)] = hash2
                    n1 = self.cmpHash(self.hash1, self.hashdic[str(imgname)])
                    if n1 != (-1):
                        self.imgdict[str(imgname)] = n1
                    elif n1 == 0:
                        self.num111+=1
            else:
                n1 = self.cmpHash(self.hash1, self.hashdic[str(imgname)])
                if n1 != (-1):
                    self.imgdict[str(imgname)] = n1
                elif n1 == 0:
                    self.num111+=1
            # print(n1)
        except Exception as e:
                print(e)

if __name__ == '__main__':
        a = APP()
        a.main()

陈东豪

很多代码不错的