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

陈东豪

1. from tkinter import *
   
2. import cv2
   
3. from webbrowser import open as webopen
   
4. import numpy as np
   
5. from threading import Thread
   
6. from multiprocessing.dummy import Pool
   
7. from os import listdir, mkdir, remove, startfile, getcwd
   
8. from os.path import isdir, isfile, join, exists
   
9. from tkinter import filedialog
   
10. from windnd import hook_dropfiles
    
11. from PIL import Image, ImageTk
    
12. from tkinter.messagebox import showinfo, askyesno
    
13. from shutil import copy
    
14. from skimage import io
    
15. from pickle import dump, load
    
16. from time import time
    
17. 
    
18. class APP:
    
19. 
    
20.     def __init__(self):
    
21.         self.root = Tk()
    
22.         self.root.geometry('450x430')
    
23.         self.root.title('星猫搜图~')
    
24.         self.imgnum = 0     #用于判断文件夹内图片个数
    
25.         self.if_change = 0   #判断图片数量是否变化
    
26.         self.file = ''
    
27.         self.hashdic = {}      
    
28.         self.img11 = 0
    
29.         self.num111 = 0
    
30.         self.badimg = 0
    
31.         self.badimglist = []
    
32.         self.msg1 = StringVar(value=r'C:\Users\陈东豪\Desktop')
    
33.         self.if_top123 = 0
    
34.         self.filelist = []
    
35.         self.initimg = []
    
36.         self.pk_file = 'hashֵ.pk'
    
37.         if not isfile(self.pk_file):
    
38.             with open(self.pk_file, 'wb') as (f):
    
39.                 dump(self.hashdic, f)
    
40.         with open(self.pk_file, 'rb') as (b):
    
41.             self.hashdic = load(b)
    
42. 
    
43.     def wait(self):
    
44.         """等候时打开页面,显示加载过程"""
    
45.         self.wait1 = Label((self.root), text='请稍等...')
    
46. 
    
47.     def main(self):
    
48.         Label((self.root), text='img文件路径:').grid(row=0, column=0)
    
49.         self.hostVar = StringVar(value=r'C:\Users\Administrator\Desktop')
    
50.         self.numVar = StringVar(self.root, str(len(self.filelist)) + '张')
    
51.         self.filename = Entry((self.root), textvariable=(self.hostVar))
    
52.         self.filename.grid(row=0, column=1)
    
53.         
    
54.         self.bt = Button((self.root), text='确定', command=self.change_lujing)
    
55.         self.bt.grid(row=0, column=2)
    
56.         self.lb4 = Label((self.root), textvariable=(self.numVar))
    
57.         self.lb4.grid(row=0, column=3)
    
58.         self.bt2 = Button((self.root), text='打开目标img文件夹', command=(self.open_lujing))
    
59.         self.bt2.grid(row=1, column=0)
    
60.         self.bt4 = Button((self.root), text='此次查找详情', command=lambda:self.show_init(self.initimg))
    
61.         self.bt4.grid(row=1, column=1)
    
62.         self.bt5 = Button((self.root), text='打开星猫图片', command=(self.open_files))
    
63.         self.bt5.grid(row=1, column=3)
    
64.         self.bt6 = Button((self.root), text='查看损坏img', command=(self.see_badimg))
    
65.         self.bt6.grid(row=1, column=2)
    
66.         self.balloon = Label(self.root)
    
67. 
    
68.         self.bt.bind('<Button-3>', self.chose_dir)
    
69.         self.bt.bind('<Enter>', lambda event: self.Balloon_show(event, msg='单击鼠标右键打开文件选择框!', fg='blue'))
    
70.         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'))
    
71.         self.root.bind('<Leave>', self.Balloon_destroy)
    
72.         self.root.bind('<Double-Button-3>', self.baidu)
    
73.         hook_dropfiles(self.root, func=self.dragged_files)   #当拖入文件时触发
    
74.         self.change_lujing()
    
75.         self.display_num()
    
76.         self.root.mainloop()
    
77.         
    
78.         
    
79.         
    
80.     def show_init(self,filename):
    
81.         '''展示目标img信息'''
    
82.         file_msg = '\n'.join((item for item in filename))
    
83.         showinfo(title='提示', message=('总共有' + str(len(filename)) + '张相似图片!\n相似度为' +\
    
84.             str(100 - self.min_) + '\n' + file_msg + '\n损坏图片:' + str(self.badimg) + \
    
85.                 '张\n共耗时:' + str(self.t) + 's'))
    
86.         
    
87.     def Balloon_destroy(self, event):
    
88.         """隐藏气泡&#702;"""
    
89.         try:
    
90.             self.balloon.place_forget()
    
91.         except:
    
92.             pass
    
93. 
    
94.     def Balloon_show(self, event, msg, fg=None, bg=None, font=('微软雅黑', 10, 'bold')):
    
95.         try:
    
96.             self.balloon.place_forget()
    
97.         except:
    
98.             pass
    
99.         else:
    
100.             self.balloon['fg'] = fg
     
101.             self.balloon['bg'] = bg
     
102.             self.balloon['text'] = msg
     
103.             self.balloon['font'] = font
     
104.             self.balloon.place(x=35, y=190)
     
105. 
     
106.     def baidu(self, event):
     
107.         """百度搜图"""
     
108.         webopen('https://graph.baidu.com/pcpage/index?tpl_from=pc')
     
109. 
     
110.     def chose_dir(self, event):
     
111.         file = filedialog.askopenfilename()
     
112.         if file != '':
     
113.             file1 = file.rsplit('/', 1)[0]
     
114.             self.hostVar.set(file1)
     
115.             self.change_lujing()
     
116. 
     
117.     def open_lujing(self, *args):
     
118.         """打开目标img文件夹"""
     
119.         if len(self.initimg) == 0:
     
120.             showinfo(title='提示', message='无目标img!')
     
121.         else:
     
122.             files = []
     
123.             for each in self.initimg:
     
124.                 files.append(each.rsplit('\\', 1)[0])
     
125.             else:
     
126.                 files = list(set(files))
     
127.                 for each in files:
     
128.                     startfile(each)
     
129. 
     
130.     def open_files(self, *args):
     
131.         """打开文件夹"""
     
132.         fname = getcwd() + '//星猫图片'
     
133.         startfile(fname)
     
134. 
     
135. 
     
136.     def see_badimg(self, *args):
     
137.         """损坏badimg"""
     
138.         if len(self.badimglist) == 0:
     
139.             showinfo(title='提示', message='没有badimg!')
     
140.         else:
     
141.             msg = '\n'.join(self.badimglist)
     
142.             showinfo('badimg路径', msg)
     
143.             if askyesno('警告', '是否删除损坏imgs?'):
     
144.                 for each in self.badimglist:
     
145.                     try:
     
146.                         remove(each)
     
147.                         self.badimglist.remove(each)
     
148.                     except Exception as e:
     
149.                         showinfo(title='提示', message=e)
     
150. 
     
151. 
     
152.     def display_num(self):
     
153.         """提示图片数量, 判断图片个数是否不再改变"""
     
154. 
     
155.         num = len(self.filelist)
     
156.         self.numVar.set(str(num) + '张')
     
157.         if self.imgnum!=num:
     
158.             self.imgnum = num
     
159.             self.if_change = 1
     
160.         else:
     
161.             if self.if_change:
     
162.                 #如果变化停止则开始计算hash值
     
163.                 self.if_change = 0
     
164.                 self.pool_map(self.compute_hash,self.filelist)
     
165.                 # showinfo('提示','开始计算图片hash值啦!')
     
166.         self.root.after(1000, self.display_num)
     
167. 
     
168.     def change_lujing(self):
     
169.         """更改搜图路径"""
     
170. 
     
171.         self.file = self.filename.get()
     
172.         if isdir(self.file):
     
173.             pass
     
174.         else:
     
175.             lujing = getcwd()
     
176.             self.hostVar.set(lujing)
     
177.             self.file = self.filename.get()
     
178.         self.filelist = []
     
179.         self.thread_it(self.getFileList, self.file, self.filelist)
     
180. 
     
181.     def pool_map(self, func, *args):
     
182.         """多进程"""
     
183.         pool = Pool()
     
184.         pool.map_async(func, args)
     
185.         pool.close()
     
186.         pool.join()
     
187. 
     
188.     def getFileList(self, dir, imglist, ext=['jpg', 'png', 'jpeg', 'bmp', 'webp']):
     
189.         """
     
190.         获取文件夹及其子文件夹中图片列表
     
191.         输入 dir:文件夹根目录
     
192.         输入 ext: 扩展名
     
193.         返回： 文件路径列表
     
194.         """
     
195.         newDir = dir
     
196.         if isfile(dir):
     
197.             if dir.rsplit('.', 1)[(-1)] in ext:
     
198.                 imglist.append(dir)
     
199. 
     
200.         elif isdir(dir):
     
201.             try:
     
202.                 for s in listdir(dir):
     
203.                     newDir = join(dir, s)
     
204.                     self.getFileList(newDir, imglist, ext)
     
205. 
     
206.             except Exception as e:
     
207.                     print(s,'\n',e)
     
208. 
     
209.         return imglist
     
210. 
     
211. 
     
212.     #打开指定的图片，缩放到指定尺寸
     
213.     def get_img(self,filename,width=445,height=365):
     
214.         try:
     
215.             im = Image.open(filename).resize((width,height))
     
216.             # 引用：添加一个Label，用来存储图片。使用PanedWindow也行。
     
217.             panel = Label(master=self.root)
     
218.             panel.photo = ImageTk.PhotoImage(im)  # 将原本的变量photo改为panel.photo
     
219.             return panel.photo
     
220.         except:
     
221.             return 0
     
222. 
     
223.     def compute_hash(self,imgname):
     
224.         '''专门用于计算大规模计算hash值'''
     
225.         if imgname not in self.hashdic:
     
226.             img2_ = self.cv_imread(imgname)
     
227.             hash2 = self.aHash(img2_, imgname)
     
228.             if hash2 != 0:
     
229.                 self.hashdic[str(imgname)] = hash2
     
230. 
     
231.                
     
232.     def dragged_files(self, files):
     
233.         """拖拽图片"""
     
234.         sum_msg2 = [item.decode('gbk') for item in files]
     
235.         for sum_msg in sum_msg2:
     
236.             ext1 = sum_msg.rsplit('.', 1)[(-1)]
     
237.             try:
     
238.                 self.img.destroy()
     
239.             except:
     
240.                 pass
     
241.         
     
242.             if ext1 in ('jpg', 'png', 'jpeg', 'bmp', 'ico', 'webp'):
     
243.                 if ext1 == 'ico':
     
244.                     showinfo('提示', '暂不支持ico格式图片搜索!')
     
245.                 else:
     
246.                     image1 = self.get_img(sum_msg)
     
247.                     if image1 != 0:
     
248.                         self.img = Label((self.root), image=image1)
     
249.                         self.img.grid(row=2, columnspan=4)
     
250.                         self.imgdict = {}   #用来保存图片相似度
     
251.                         self.thread_it(self.Compare2, sum_msg)
     
252.                     else:
     
253.                         showinfo('提示',sum_msg+'\n该图片有问题!')
     
254.                         
     
255.             else:
     
256.                 showinfo(title='提示', message=('暂不支持当前格式搜索!(' + str(ext1) + ')'))
     
257. 
     
258.     def Compare2(self, img):
     
259.         '''开始搜索相似图片ing'''
     
260.         t1 = time()
     
261.         img1 = self.cv_imread(img)
     
262.         self.hash1 = self.aHash(img1,img)
     
263.         if self.hash1 == 0:
     
264.             showinfo('提示', '此图片损坏,请更换其他图片!')
     
265.         else:
     
266.             # self.pool_map(self.compare, self.filelist)
     
267.             #使用多进程加速
     
268. 
     
269.             pool = Pool()
     
270.             pool.map_async(self.compare, self.filelist)
     
271.             pool.close()
     
272.             pool.join()
     
273. 
     
274.             try:
     
275. 
     
276.                 self.min_ = min(self.imgdict.values())
     
277.                 min_key = [i for i in self.imgdict.keys() if self.imgdict[i] == self.min_]
     
278.                 self.initimg = min_key
     
279.                 if not exists('星猫图片'):
     
280.                     mkdir('星猫图片')
     
281.                 for each_img in min_key:
     
282.                     try:
     
283.                         copy(each_img, './/星猫图片')
     
284.                     except:
     
285.                         pass
     
286. 
     
287.                 if self.min_ == 0:
     
288.                     file_img = []
     
289.                     file_img.append(min_key[0])
     
290.                     self.show_img1(file_img)
     
291.                 else:
     
292.                     self.show_img1(min_key)
     
293.                 with open(self.pk_file, 'wb') as (f):
     
294.                     dump(self.hashdic, f)
     
295.                 t2 = time()
     
296.                 self.t = t2 - t1
     
297. 
     
298.             except:
     
299.                 showinfo('提示', '无相关图片!')
     
300. 
     
301.     def show_msg(self, title1='提示', msg='呱呱呱!', ms=0, fontbig=20):
     
302.         """显示消息"""
     
303.         if msg != self.msg1:
     
304.             self.msg1.set(msg)
     
305.         if not self.if_top123:
     
306.             self.top123 = Toplevel()
     
307.             self.top123.title(title1)
     
308.             self.top123.protocol('WM_DELETE_WINDOW', self.destroytop)
     
309.             msg1 = Label((self.top123), textvariable=(self.msg1), font=('微软雅黑', fontbig, 'bold'))
     
310.             msg1.pack()
     
311.         if ms != 0:
     
312.             self.top123.after(ms, self.destroytop)
     
313. 
     
314.     def destroytop(self):
     
315.         """关闭窗口"""
     
316.         self.top123.destroy()
     
317.         self.if_top123 = 0
     
318. 
     
319.     def show_img1(self, filename=[], ms=4500, num=0):
     
320.         """可切换显示多张照片"""
     
321.         try:
     
322.             if not self.img11:
     
323.                 self.top11 = Toplevel()
     
324.                 self.top11.title('星猫图片浏览~')
     
325.                 self.top11.geometry('640x400')
     
326.                 self.top11.protocol('WM_DELETE_WINDOW', lambda : self.closeimg(filename))
     
327.                 self.img11 = 1
     
328.                 self.img2 = Label((self.top11), image=(self.get_img(filename[num], 640, 400)))
     
329.                 self.img2.pack()
     
330.             else:
     
331.                 self.img2.destroy()
     
332.                 self.img2 = Label((self.top11), image=(self.get_img(filename[num], 640, 400)))
     
333.                 self.img2.pack()
     
334.             num += 1
     
335.             self.top11.after(3000, self.show_img1, filename, ms, num)
     
336.         except:
     
337.             self.closeimg(filename)
     
338. 
     
339.     def closeimg(self, filename):
     
340.         """窗口关闭时触发"""
     
341.         self.img11 = 0
     
342.         self.top11.destroy()
     
343. 
     
344.     def thread_it(self, func, *args):
     
345.         """多线程"""
     
346.         t = Thread(target=func, args=args)
     
347.         t.setDaemon(True)
     
348.         t.start()
     
349. 
     
350.     #差值感知算法
     
351.     def aHash(self,img,imgname,if_shishi=0):
     
352.         '''img:图片矩阵,imgname:图片名字,if_shishi是否'''
     
353.         try:
     
354.             #缩放8*9
     
355.             img1=cv2.resize(img,(9,8),interpolation=cv2.INTER_CUBIC)
     
356.             #转换灰度图
     
357.             gray=cv2.cvtColor(img1,cv2.COLOR_BGR2GRAY)
     
358.             hash_str=''
     
359.             #每行前一个像素大于后一个像素为1，相反为0，生成哈希
     
360.             for i in range(8):
     
361.                 for j in range(8):
     
362.                     if   gray[i,j]>gray[i,j+1]:
     
363.                         hash_str=hash_str+'1'
     
364.                     else:
     
365.                         hash_str=hash_str+'0'
     
366.             return hash_str
     
367.         except:
     
368.             '''如果打开图片失败，就重新保存试一试'''
     
369.             if not if_shishi:
     
370.                 '''如果是第一次失败，则重试'''
     
371.                 try:
     
372.                     image = io.imread(imgname)
     
373.                     image = cv2.cvtColor(image, cv2.COLOR_RGBA2BGRA)
     
374.                     cv2.imencode(imgname.rsplit('.',1)[-1],image)[1].tofile(imgname)
     
375.                     return self.aHash(image,imgname,1)
     
376.                 except:
     
377.                     print(imgname,'打开失败!')
     
378.                     self.badimg+=1
     
379.                     return 0
     
380.             else:
     
381.                 '''表示获取hash失败!'''
     
382.                 return 0
     
383. 
     
384.     def cmpHash(self, hash1, hash2):
     
385.         n = 0
     
386.         if len(hash1) != len(hash2):
     
387.             return -1
     
388.         for i in range(len(hash1)):
     
389.             if hash1[i] != hash2[i]:
     
390.                 n = n + 1
     
391.         return n
     
392. 
     
393.     def cv_imread(self, filePath):
     
394.         cv_img = cv2.imdecode(np.fromfile(filePath, dtype=(np.uint8)), -1)
     
395.         return cv_img
     
396. 
     
397.     def compare(self, img2):
     
398.         '''比较图片相似度并保存相似度n1'''
     
399.         imgname = img2
     
400.         try:
     
401.             if str(imgname) in self.imgdict:
     
402.                 pass
     
403.             elif str(imgname) not in self.hashdic:
     
404.                 img2_ = self.cv_imread(img2)
     
405.                 hash2 = self.aHash(img2_, imgname=img2)
     
406.                 if hash2 != 0:
     
407.                     self.hashdic[str(imgname)] = hash2
     
408.                     n1 = self.cmpHash(self.hash1, self.hashdic[str(imgname)])
     
409.                     if n1 != (-1):
     
410.                         self.imgdict[str(imgname)] = n1
     
411.                     elif n1 == 0:
     
412.                         self.num111+=1
     
413.             else:
     
414.                 n1 = self.cmpHash(self.hash1, self.hashdic[str(imgname)])
     
415.                 if n1 != (-1):
     
416.                     self.imgdict[str(imgname)] = n1
     
417.                 elif n1 == 0:
     
418.                     self.num111+=1
     
419.             # print(n1)
     
420.         except Exception as e:
     
421.                 print(e)
     
422. 
     
423. if __name__ == '__main__':
     
424.         a = APP()
     
425.         a.main()
     

复制代码

陈东豪

很多代码不错的