app.py

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] #it is used while taking the screenshots it holds the count of each one

        self.model = model.Model()

        self.auto_predict = False

        self.camera = camera.Camera()

        self.init_gui()

        self.delay = 15 #updates after some delay
        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) #canvas where camera is present
        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) #pack them and allign them in the centre
        

        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))# save_for_class isused to make screenshot and save it into right directory
        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)) #counters is used to send the screenshots to train the model
        self.btn_train.pack(anchor=tk.CENTER, expand=True)

        self.btn_predict = tk.Button(self.window, text="Predict", 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): #class_num is used to save the screenshots in the respective directory
        ret, frame = self.camera.get_frame()
        if not os.path.exists("1"):#checking if the directory exists if not then
            os.mkdir("1") #creating the new directory
        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))#imagewrite class_num is the directory name and counters[] represt and access the particular counter
        img = PIL.Image.open(f'{class_num}/frame{self.counters[class_num - 1]}.jpg') #opens the exact image
        img.thumbnail((150, 150), PIL.Image.ANTIALIAS) 
        img.save(f'{class_num}/frame{self.counters[class_num - 1]}.jpg') #saving the image

        self.counters[class_num - 1] += 1 #increasing the counters for each screenshot

    def reset(self):#reset to the beginning
        for folder in ['1', '2']:
            for file in os.listdir(folder):#for each file in that directory
                file_path = os.path.join(folder, file)# combines both directory name and the file name
                if os.path.isfile(file_path):#if file is present
                    os.unlink(file_path)#then delete the file

        self.counters = [1, 1]#counting from 1(reseting from 1)
        self.model = model.Model()
        self.class_label.config(text="CLASS")#reseting the label to the class(default value)

    def update(self):#to get the current camera data 
        if self.auto_predict:
            print(self.predict())

        ret, frame = self.camera.get_frame()#current camera frame

        if ret:#if something is present
            self.photo = PIL.ImageTk.PhotoImage(image=PIL.Image.fromarray(frame))#getting the image from the camera frame and turing it into the tk image for gui
            self.canvas.create_image(0, 0, image=self.photo, anchor=tk.NW)#image is created and allign it to north west ie,upper west

        self.window.after(self.delay, self.update)#after some delay then 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