Rat tracking software
Revision | 9a689d35b25cbcb1d2a1d266ebba00f18711e5c1 (tree) |
---|---|
Time | 2018-02-12 23:14:42 |
Author | Gregor <Gregor133@inte...> |
Commiter | Gregor |
First commit of the gui main file - alpha stable
@@ -0,0 +1,317 @@ | ||
1 | +#!/usr/bin/env python2 | |
2 | +# -*- coding: utf-8 -*- | |
3 | + | |
4 | + | |
5 | +""" | |
6 | +Main file of rat images analysis and tracking | |
7 | +Author: GN | |
8 | +""" | |
9 | + | |
10 | + | |
11 | +import cv2 | |
12 | +import numpy as np | |
13 | +from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg | |
14 | +from matplotlib.figure import Figure | |
15 | +import matplotlib.image as mpimg | |
16 | +import tkinter.filedialog as filedialog | |
17 | +import tkinter.messagebox as msg | |
18 | +import tkinter.scrolledtext as sctext | |
19 | +import tkinter as tk | |
20 | +import tkinter.ttk as ttk | |
21 | +import ntpath | |
22 | +import os | |
23 | +import areadetect | |
24 | +import ratdetect | |
25 | +import rattrack | |
26 | + | |
27 | + | |
28 | +class RApp(tk.Tk, object): | |
29 | + """ | |
30 | + Main class of the application | |
31 | + Create and keep the tkinter window | |
32 | + """ | |
33 | + def __init__(self): | |
34 | + super(RApp, self).__init__() | |
35 | + | |
36 | + self.title('Rat App') | |
37 | + # self.geometry("560x300") | |
38 | + | |
39 | + self.is_file_opened = False | |
40 | + self.bg = None | |
41 | + self.first_rat_frame = None | |
42 | + self.lefta = None | |
43 | + self.centera = None | |
44 | + self.righta = None | |
45 | + self.after_job = None | |
46 | + self.tracked = False | |
47 | + self.rr = np.zeros(3) | |
48 | + | |
49 | + wwidth = 770 | |
50 | + wheight = 450 | |
51 | + sw = self.winfo_screenwidth() | |
52 | + sh = self.winfo_screenheight() | |
53 | + x = (sw - wwidth)/2 | |
54 | + y = (sh - wheight)/2 | |
55 | + self.geometry('%dx%d+%d+%d' % (wwidth, wheight, x, y)) | |
56 | + | |
57 | + self.menu = tk.Menu(self, bg="lightgrey", fg="black") | |
58 | + self.file_menu = tk.Menu(self.menu, tearoff=0, bg="lightgrey", fg="black") | |
59 | + self.file_menu.add_command(label="Open", command=self.openfile, accelerator="Ctrl+o") | |
60 | + self.file_menu.add_command(label="Exit", command=self.exitapp, accelerator="Ctrl+q") | |
61 | + self.menu.add_cascade(label="File", menu=self.file_menu, underline=0) | |
62 | + self.config(menu=self.menu) | |
63 | + | |
64 | + self.bind_all('<Control-o>', self.openfile) | |
65 | + self.bind_all('<Control-q>', self.exitapp) | |
66 | + | |
67 | + self.fig = Figure() | |
68 | + self.ax = self.fig.add_subplot(111) | |
69 | + | |
70 | + self.canvas = FigureCanvasTkAgg(self.fig, master=self) | |
71 | + self.canvas.show() | |
72 | + self.canvas.get_tk_widget().grid(row=0, column=0, rowspan=5, padx=5, pady=5, sticky='nwes') | |
73 | + | |
74 | + self.button_area = tk.Button(text="Detect areas", command=self.detect_area) | |
75 | + self.button_animal = tk.Button(text="Detect animal", command=self.detect_animal) | |
76 | + self.button_track = tk.Button(text="Track", command=self.tracking) | |
77 | + self.button_summary = tk.Button(text="Summary", command=self.show_summary) | |
78 | + self.button_quit = tk.Button(text="Quit", command=self.exitappbutton) | |
79 | + self.button_area.grid(row=0, column=1, pady=5, sticky='nwe') | |
80 | + self.button_animal.grid(row=1, column=1, pady=5, sticky='nwe') | |
81 | + self.button_track.grid(row=2, column=1, pady=5, sticky='nwe') | |
82 | + self.button_summary.grid(row=3, column=1, pady=5, sticky='nwe') | |
83 | + self.button_quit.grid(row=4, column=1, pady=5, sticky='nwe') | |
84 | + | |
85 | + self.textbox = sctext.ScrolledText(self, height=7, width=60) | |
86 | + self.textbox['font'] = ('consolas', '11') | |
87 | + self.textbox.grid(row=5, column=0, columnspan=2, padx=5, pady=5, sticky='nwes') | |
88 | + self.progressbar = ttk.Progressbar(self, orient=tk.HORIZONTAL, mode="determinate") | |
89 | + self.progressbar.grid(row=6, column=0, columnspan=2, padx=5, pady=5, sticky='we') | |
90 | + | |
91 | + self.grid_rowconfigure(0, weight=0) | |
92 | + self.grid_rowconfigure(1, weight=0) | |
93 | + self.grid_rowconfigure(2, weight=0) | |
94 | + self.grid_rowconfigure(3, weight=0) | |
95 | + self.grid_rowconfigure(4, weight=1) | |
96 | + self.grid_columnconfigure(0, weight=0) | |
97 | + self.grid_columnconfigure(1, weight=1) | |
98 | + | |
99 | + self.ini_file = None | |
100 | + self.file_name_path = None | |
101 | + self.file_dir = None | |
102 | + self.file_name = None | |
103 | + self.files_no = None | |
104 | + self.files_list = None | |
105 | + self.current_file = None | |
106 | + self.tree_frame = None | |
107 | + self.tree = None | |
108 | + self.img = None | |
109 | + self.iax = None | |
110 | + | |
111 | + def openfile(self): | |
112 | + """ | |
113 | + Function for opening folder with images | |
114 | + :return: | |
115 | + """ | |
116 | + self.ini_file = filedialog.askopenfilename(initialdir='/'.join([os.getcwd(), 'images'])) | |
117 | + | |
118 | + while self.ini_file and not self.ini_file.endswith(".jpeg"): | |
119 | + msg.showerror("Wrong filetype", "Select an jpg file") | |
120 | + self.ini_file = filedialog.askopenfilename() | |
121 | + | |
122 | + if self.ini_file: | |
123 | + self.parse_ini_file() | |
124 | + # print('inifile = {}'.format(self.ini_file)) | |
125 | + | |
126 | + def parse_ini_file(self): | |
127 | + """ | |
128 | + Parse the given folder with images and show the first one | |
129 | + :param: | |
130 | + :return: | |
131 | + """ | |
132 | + self.file_name_path = ": ".join([ntpath.basename(self.ini_file), self.ini_file]) | |
133 | + self.file_name = ntpath.basename(self.ini_file) | |
134 | + self.file_dir = ntpath.dirname(self.ini_file) | |
135 | + self.files_list = sorted([f for f in os.listdir(self.file_dir) if os.path.isfile('/'.join((self.file_dir, f))) | |
136 | + and f.lower().endswith('.jpeg')]) | |
137 | + self.files_no = len(self.files_list) | |
138 | + self.textbox.insert(tk.END, "Opened file: " + self.file_name + '\n') | |
139 | + self.textbox.insert(tk.END, "Opened directory: " + self.file_dir + '\n') | |
140 | + self.textbox.insert(tk.END, "Total files in directory: " + str(self.files_no) + '\n') | |
141 | + | |
142 | + self.img = mpimg.imread(self.ini_file) | |
143 | + self.iax = self.ax.imshow(self.img) | |
144 | + self.canvas.draw() | |
145 | + self.is_file_opened = True | |
146 | + | |
147 | + def detect_area(self): | |
148 | + """ | |
149 | + Automatic detection of areas - parts of cage when the rat could be present | |
150 | + :return: | |
151 | + """ | |
152 | + if not self.is_file_opened: | |
153 | + msg.showerror("File not opened", "Open files folder first") | |
154 | + else: | |
155 | + inarea, outarea = areadetect.determine_areas(self.img) | |
156 | + self.lefta, self.centera, self.righta = areadetect.location_areas(self.img, inarea) | |
157 | + _, self.lefta = cv2.threshold(self.lefta.astype('uint8'), 128, 255, cv2.THRESH_BINARY) | |
158 | + _, self.centera = cv2.threshold(self.centera.astype('uint8'), 128, 255, cv2.THRESH_BINARY) | |
159 | + _, self.righta = cv2.threshold(self.righta.astype('uint8'), 128, 255, cv2.THRESH_BINARY) | |
160 | + rimg = areadetect.draw_areas(self.img.copy(), inarea, outarea) | |
161 | + self.iax.set_array(rimg) | |
162 | + self.canvas.draw() | |
163 | + | |
164 | + def detect_animal(self): | |
165 | + """ | |
166 | + Detection of animal at given image | |
167 | + :return: | |
168 | + """ | |
169 | + if not self.is_file_opened: | |
170 | + msg.showerror("File not opened", "Open the first file in folder") | |
171 | + ref_img = mpimg.imread(self.file_dir + '/' + self.files_list[1]) | |
172 | + ratcc = ratdetect.detect_rat(self.img, ref_img) | |
173 | + ratcc_image = self.iax.get_array() | |
174 | + cv2.drawContours(ratcc_image, [ratcc], 0, (0, 255, 0), 3) | |
175 | + cv2.drawContours(ratcc_image, [ratcc], 0, (255, 255, 0), cv2.FILLED) | |
176 | + self.first_rat_frame = ratcc_image.copy() | |
177 | + # self.ax.imshow(ratcc_image) | |
178 | + self.iax.set_array(ratcc_image) | |
179 | + self.canvas.draw() | |
180 | + | |
181 | + bs_start_img = cv2.imread(self.file_dir + '/' + self.files_list[0]) | |
182 | + bg_ref_img = cv2.imread(self.file_dir + '/' + self.files_list[3]) | |
183 | + remove_rect = (260, 585, 400, 690) | |
184 | + | |
185 | + self.bg = rattrack.remove_rat_bg(bs_start_img, bg_ref_img, remove_rect) | |
186 | + | |
187 | + def tracking(self): | |
188 | + """ | |
189 | + Function for tracking the position of rat at the images | |
190 | + :return: | |
191 | + """ | |
192 | + if not self.is_file_opened: | |
193 | + msg.showerror("Nothing for tracking", "Open file for tracking") | |
194 | + elif self.bg is None: | |
195 | + msg.showerror("Nothing for tracking", "Detect animal for tracking") | |
196 | + else: | |
197 | + # print('Tracking ...') | |
198 | + # print('Files number = {}'.format(self.files_no)) | |
199 | + self.progressbar["value"] = 0 | |
200 | + self.progressbar["maximum"] = self.files_no | |
201 | + self.current_file = 0 | |
202 | + self.read_file_and_track() | |
203 | + | |
204 | + def read_file_and_track(self): | |
205 | + """ | |
206 | + Read the current file and detect position of the rat at this file | |
207 | + :return: | |
208 | + """ | |
209 | + # print(self.current_file) | |
210 | + self.textbox.insert(tk.END, "Current file no:" + str(self.current_file) + '\n') | |
211 | + ratarea = np.zeros(self.img.shape).astype('uint8') | |
212 | + cimg = cv2.imread(self.file_dir + '/' + self.files_list[self.current_file]) | |
213 | + self.textbox.insert(tk.END, "Current file path:" + self.file_dir + '/' + self.files_list[self.current_file] + | |
214 | + '\n') | |
215 | + self.textbox.see(tk.END) | |
216 | + self.current_file += 1 | |
217 | + if self.current_file == self.files_no: | |
218 | + self.tracked = True | |
219 | + self.progressbar["value"] = self.current_file | |
220 | + rct3 = rattrack.detect_rat_cts(cimg, self.bg) | |
221 | + cv2.drawContours(cimg, rct3, -1, (0, 255, 0), 3) | |
222 | + cv2.drawContours(cimg, rct3, -1, (0, 127, 0), -1) | |
223 | + cv2.drawContours(ratarea, rct3, -1, (255, 255, 255), -1) | |
224 | + _, bwratarea = cv2.threshold(ratarea, 128, 255, cv2.THRESH_BINARY) | |
225 | + | |
226 | + leftratfield = cv2.bitwise_and(self.lefta, bwratarea) | |
227 | + centerratfield = cv2.bitwise_and(self.centera, bwratarea) | |
228 | + rightcenterfield = cv2.bitwise_and(self.righta, bwratarea) | |
229 | + | |
230 | + leftvote = np.count_nonzero(leftratfield) | |
231 | + centervote = np.count_nonzero(centerratfield) | |
232 | + rightvote = np.count_nonzero(rightcenterfield) | |
233 | + | |
234 | + self.textbox.insert(tk.END, "Left vote: " + str(leftvote) + ";" + " Center vote: " + str(centervote) + | |
235 | + ";" + " Right vote: " + str(rightvote) + '\n') | |
236 | + decision_idx = np.argmax((leftvote, centervote, rightvote)) | |
237 | + self.rr[decision_idx] += 1 | |
238 | + | |
239 | + self.textbox.insert(tk.END, "Decision: " + str(decision_idx) + '\n') | |
240 | + self.textbox.see(tk.END) | |
241 | + self.iax.set_array(cimg) | |
242 | + # self.iax.set_array(ratarea) | |
243 | + self.canvas.draw() | |
244 | + if self.current_file < self.files_no: | |
245 | + self.after_job = self.after(25, self.read_file_and_track) | |
246 | + | |
247 | + def show_summary(self): | |
248 | + """ | |
249 | + Show summary of tracking - the number of frames at which the rat was detected | |
250 | + :return: | |
251 | + """ | |
252 | + if not self.tracked: | |
253 | + msg.showerror("Nothing tracked", "Nothing tracked, nothing to show!") | |
254 | + else: | |
255 | + x = self.winfo_rootx() | |
256 | + y = self.winfo_rootx() | |
257 | + self.tree_frame = tk.Toplevel(master=self) | |
258 | + self.tree_frame.title('Tracking summary') | |
259 | + self.tree_frame.config(background="lavender") | |
260 | + self.tree = ttk.Treeview(self.tree_frame, show='headings') | |
261 | + self.tree["columns"] = ("l", "c", "r") | |
262 | + self.tree.column('l', width=101, anchor="center") | |
263 | + self.tree.column('c', width=101, anchor="center") | |
264 | + self.tree.column('r', width=101, anchor="center") | |
265 | + self.tree.heading('l', text='Left side') | |
266 | + self.tree.heading('c', text='Center size') | |
267 | + self.tree.heading('r', text='Right size') | |
268 | + self.tree.insert('', 'end', values=(self.rr[0], self.rr[1], self.rr[2])) | |
269 | + informer = tk.Button(self.tree_frame, text="More info", command=self.summary_show_info) | |
270 | + informer.pack(side="bottom", fill="both") | |
271 | + informer = tk.Button(self.tree_frame, text="Ok", command=self.tree_frame.destroy) | |
272 | + informer.pack(side="bottom", fill="both") | |
273 | + self.tree.pack() | |
274 | + self.tree_frame.geometry("+%d+%d" % (x + 240, y - 120)) | |
275 | + | |
276 | + @staticmethod | |
277 | + def summary_show_info(): | |
278 | + """ | |
279 | + Show information about the result of tracking | |
280 | + :return: | |
281 | + """ | |
282 | + tk.messagebox.showinfo("More info", "Left column represents amount of frames from left user view (left hand)\n" | |
283 | + "Center column represents amount of frames in center of user view\n" | |
284 | + "Right column represents amount of frames from right user view (right hand\n") | |
285 | + | |
286 | + def exitappbutton(self): | |
287 | + """ | |
288 | + Exit application by click of the button in the main window | |
289 | + :return: | |
290 | + """ | |
291 | + if self.after_job is not None: | |
292 | + self.after_cancel(self.after_job) | |
293 | + self.after_job = None | |
294 | + self.quit() | |
295 | + self.destroy() | |
296 | + else: | |
297 | + self.quit() | |
298 | + self.destroy() | |
299 | + | |
300 | + def exitapp(self): | |
301 | + """ | |
302 | + Exit application from the menu | |
303 | + :return: | |
304 | + """ | |
305 | + if self.after_job is not None: | |
306 | + self.after_cancel(self.after_job) | |
307 | + self.after_job = None | |
308 | + self.quit() | |
309 | + self.destroy() | |
310 | + else: | |
311 | + self.quit() | |
312 | + self.destroy() | |
313 | + | |
314 | + | |
315 | +if __name__ == "__main__": | |
316 | + | |
317 | + RApp().mainloop() |