PythonProject/subimages.py at master · Vishalkumar12/PythonProject · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164

import os
import numpy as np
from math import sqrt
from time import time
from queue import Queue
from tkinter import *
from tkinter import ttk, filedialog, messagebox
from PIL import Image, ImageChops
from skimage.feature import match_template # requires numpy, scipy, and six

__version__ = '0.8.3'

MAX_WIDTH = 640
MAX_HEIGHT = 640
RMS_THRESHOLD = 50

class FindSubset:

    def __init__(self, master):
        self.master = master
        self.master.title('Find Subset Images')
        self.master.resizable(False, False)

        self.main_frame = ttk.Frame(self.master)
        self.main_frame.pack(padx = 5, pady = 5)

        ttk.Label(self.main_frame, text = 'Search Directory:').grid(row = 0, column = 0, sticky = 'w')
        self.path_entry = ttk.Entry(self.main_frame, width = 54)
        self.path_entry.grid(row = 1, column = 0, sticky = 'e')
        self.path_entry.insert(0, '.\\images') # default search folder
        self.browse_button = ttk.Button(self.main_frame, text = 'Browse...',
                                        command = self.browse_callback)
        self.browse_button.grid(row = 1, column = 1, sticky = 'w')

        self.search_button = ttk.Button(self.main_frame, text = 'Find Subset Images',
                                        command = self.search_callback)
        self.search_button.grid(row = 2, column = 0, columnspan = 2)

        self.results_table = ttk.Treeview(self.main_frame, column = ('subset'))
        self.results_table.heading('#0', text = 'Original Image')
        self.results_table.column('#0', width = 200)
        self.results_table.heading('subset', text = 'Subset Image')
        self.results_table.column('subset', width = 200)

        self.status_frame = ttk.Frame(self.master)
        self.status_frame.pack(fill = BOTH, expand = True)

        self.status_var = StringVar()
        self.status_label = ttk.Label(self.status_frame, textvariable = self.status_var)

        self.progress_var = DoubleVar()
        self.progressbar = ttk.Progressbar(self.status_frame, mode = 'determinate',
                                           variable = self.progress_var)

    def browse_callback(self):
        path = filedialog.askdirectory(initialdir = self.path_entry.get())
        self.path_entry.delete(0, END)
        self.path_entry.insert(0, path)

    def search_callback(self):
        self.start_time = time()

        try: # build list of all jpg image files in directory
            self.path = self.path_entry.get()
            images = list(entry for entry in os.listdir(self.path) if entry.endswith('.jpg'))
        except:
            messagebox.showerror(title = 'Invalid Directory',
                                 message = 'Invalid Search Directory:\n' + self.path)
            return

        if len(images) < 2:
            messagebox.showerror(title = 'Not Enough Images',
                                 message = 'Need at least 2 images to analyze.')
            return

        self.queue = Queue() # queue of image pairs to process
        for i in images:
            for j in images:
                if i != j:
                    self.queue.put((i, j))

        self.results_table.grid_forget() # clear previous results
        for item in self.results_table.get_children(''):
            self.results_table.delete(item)

        self.status_var.set('Beginning...')
        self.status_label.pack(side = BOTTOM, fill = BOTH, expand = True)
        self.progressbar.config(value = 0.0, maximum = self.queue.qsize())
        self.progressbar.pack(side = BOTTOM, fill = BOTH, expand = True)
        self.browse_button.state(['disabled'])
        self.search_button.state(['disabled'])

        self.master.after(10, self.process_queue) # provides time for tkinter to update GUI

    def process_queue(self):
        pair = self.queue.get()
        orig_img = Image.open(os.path.join(self.path, pair[0]))
        temp_img = Image.open(os.path.join(self.path, pair[1]))

        # verify that template image is larger than the original
        if (temp_img.size[0] < orig_img.size[0]) and (temp_img.size[1] < orig_img.size[1]):

            # determine if images need to be resized smaller
            if (orig_img.size[0] > MAX_WIDTH) or (orig_img.size[1] > MAX_HEIGHT):
                # calculate ratio for resizing
                ratio = min(MAX_WIDTH/float(orig_img.size[0]),
                            MAX_HEIGHT/float(orig_img.size[1]))

                # resize images based on ratio
                orig_img = orig_img.resize((int(ratio*orig_img.size[0]),
                                            int(ratio*orig_img.size[1])),
                                           Image.ANTIALIAS)
                temp_img = temp_img.resize((int(ratio*temp_img.size[0]),
                                            int(ratio*temp_img.size[1])),
                                           Image.ANTIALIAS)
            else:
                ratio = 1 # no resize was required

            orig_arr = np.array(orig_img.convert(mode = 'L'))
            temp_arr = np.array(temp_img.convert(mode = 'L'))

            match_arr = match_template(orig_arr, temp_arr)
            match_loc = np.unravel_index(np.argmax(match_arr), match_arr.shape)

            if ratio != 1: # if images were resized, get originals and calculate corresponding match_loc
                match_loc = (int(match_loc[0]/ratio), int(match_loc[1]/ratio))
                orig_img = Image.open(os.path.join(self.path, pair[0]))
                temp_img = Image.open(os.path.join(self.path, pair[1]))

            # get matching subsection from original image (using RGB mode)
            orig_sub_arr = np.array(orig_img)[match_loc[0]:match_loc[0] + temp_img.size[0],
                                                 match_loc[1]:match_loc[1] + temp_img.size[1]]
            orig_sub_img = Image.fromarray(orig_sub_arr, mode = 'RGB')

            # calculate the root-mean-square difference between orig_sub_img and sub_img
            h_diff = ImageChops.difference(orig_sub_img, temp_img).histogram()
            sum_of_squares = sum(value * ((idx % 256) ** 2) for idx, value in enumerate(h_diff))
            rms = sqrt(sum_of_squares/float(temp_img.size[0]*temp_img.size[1]))

            if RMS_THRESHOLD > rms: # add matches to table
                self.results_table.grid(row = 3, column = 0, columnspan = 2, padx = 5, pady = 5)
                self.results_table.insert('', 'end', str(self.progress_var.get()),text = pair[0])
                self.results_table.set(str(self.progress_var.get()), 'subset', pair[1])
                self.results_table.config(height = len(self.results_table.get_children('')))

        self.progressbar.step()
        self.status_var.set('Analyzed {} vs {} - {} pairs remaining...'.format(pair[0], pair[1], self.queue.qsize()))

        if not self.queue.empty():
            self.master.after(10, self.process_queue)
        else:
            self.progressbar.pack_forget()
            self.browse_button.state(['!disabled'])
            self.search_button.state(['!disabled'])
            elapsed_time = time() - self.start_time
            self.status_var.set('Done - Elapsed Time: {0:.2f} seconds'.format(elapsed_time))

def main():
    root = Tk()
    FindSubset(root)
    root.mainloop()

if __name__ == "__main__": main()