-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathapp.py
121 lines (86 loc) · 4.03 KB
/
app.py
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
'''
Camera Classifier v0.1 Alpha
Copyright (c) NeuralNine
Instagram: @neuralnine
YouTube: NeuralNine
Website: www.neuralnine.com
'''
import tkinter as tk
from tkinter import simpledialog
import cv2 as cv
import os
import PIL.Image, PIL.ImageTk
import model
import camera
class App:
def __init__(self, window=tk.Tk(), window_title="Camera Classifier"):
self.window = window
self.window_title = window_title
self.counters = [1, 1]
self.model = model.Model()
self.auto_predict = False
self.camera = camera.Camera()
self.init_gui()
self.delay = 15
self.update()
self.window.attributes("-topmost", True)
self.window.mainloop()
def init_gui(self):
self.canvas = tk.Canvas(self.window, width=self.camera.width, height=self.camera.height)
self.canvas.pack()
self.btn_toggleauto = tk.Button(self.window, text="Auto Prediction", width=50, command=self.auto_predict_toggle)
self.btn_toggleauto.pack(anchor=tk.CENTER, expand=True)
self.classname_one = simpledialog.askstring("Classname One", "Enter the name of the first class:", parent=self.window)
self.classname_two = simpledialog.askstring("Classname Two", "Enter the name of the second class:", parent=self.window)
self.btn_class_one = tk.Button(self.window, text=self.classname_one, width=50, command=lambda: self.save_for_class(1))
self.btn_class_one.pack(anchor=tk.CENTER, expand=True)
self.btn_class_two = tk.Button(self.window, text=self.classname_two, width=50, command=lambda: self.save_for_class(2))
self.btn_class_two.pack(anchor=tk.CENTER, expand=True)
self.btn_train = tk.Button(self.window, text="Train Model", width=50, command=lambda: self.model.train_model(self.counters))
self.btn_train.pack(anchor=tk.CENTER, expand=True)
self.btn_predict = tk.Button(self.window, text="Predcit", width=50, command=self.predict)
self.btn_predict.pack(anchor=tk.CENTER, expand=True)
self.btn_reset = tk.Button(self.window, text="Reset", width=50, command=self.reset)
self.btn_reset.pack(anchor=tk.CENTER, expand=True)
self.class_label = tk.Label(self.window, text="CLASS")
self.class_label.config(font=("Arial", 20))
self.class_label.pack(anchor=tk.CENTER, expand=True)
def auto_predict_toggle(self):
self.auto_predict = not self.auto_predict
def save_for_class(self, class_num):
ret, frame = self.camera.get_frame()
if not os.path.exists("1"):
os.mkdir("1")
if not os.path.exists("2"):
os.mkdir("2")
cv.imwrite(f'{class_num}/frame{self.counters[class_num-1]}.jpg', cv.cvtColor(frame, cv.COLOR_RGB2GRAY))
img = PIL.Image.open(f'{class_num}/frame{self.counters[class_num - 1]}.jpg')
img.thumbnail((150, 150), PIL.Image.ANTIALIAS)
img.save(f'{class_num}/frame{self.counters[class_num - 1]}.jpg')
self.counters[class_num - 1] += 1
def reset(self):
for folder in ['1', '2']:
for file in os.listdir(folder):
file_path = os.path.join(folder, file)
if os.path.isfile(file_path):
os.unlink(file_path)
self.counters = [1, 1]
self.model = model.Model()
self.class_label.config(text="CLASS")
def update(self):
if self.auto_predict:
print(self.predict())
ret, frame = self.camera.get_frame()
if ret:
self.photo = PIL.ImageTk.PhotoImage(image=PIL.Image.fromarray(frame))
self.canvas.create_image(0, 0, image=self.photo, anchor=tk.NW)
self.window.after(self.delay, self.update)
def predict(self):
frame = self.camera.get_frame()
prediction = self.model.predict(frame)
if prediction == 1:
self.class_label.config(text=self.classname_one)
return self.classname_one
if prediction == 2:
self.class_label.config(text=self.classname_two)
return self.classname_two