app.py

import tkinter as tk
import os
import PIL.Image, PIL.ImageTk
import cv2
import camera
import model
import shutil

# TODO clean up code

class App:

    def __init__(self):
        self.window = tk.Tk()
        self.window.title = "Rep counter"

        self.counters = [1, 1]
        self.rep_counter = 0
        
        self.extended = False
        self.contracted = False
        self.last_prediction = 0

        self.model = model.Model()

        self.counting_enabled = 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.name = tk.StringVar()  

        self.btn_toggleauto = tk.Button(self.window, text="Toggle Counting", width=50, command=self.counting_toggle)
        self.btn_toggleauto.pack(anchor=tk.CENTER)

        self.choose_model = tk.Button(self.window, text="Choose Model", width=50, command=lambda: self.model.get_model(self.name.get()))
        self.choose_model.pack(anchor=tk.CENTER)

        self.input = tk.Entry(self.window, textvariable=self.name)
        self.input.pack(anchor=tk.CENTER)

        self.btn_reset = tk.Button(self.window, text="Reset Count", width=50, command=self.reset)
        self.btn_reset.pack(anchor=tk.CENTER)

        self.counter_label = tk.Label(self.window, text=f"{self.rep_counter}")
        self.counter_label.config(font=("Arial", 24))
        self.counter_label.pack(anchor=tk.CENTER)

        self.btn_class_one = tk.Button(self.window, text="Take Picture", width=50, height=5, command=lambda: self.save_for_class(1))

        self.btn_class_two = tk.Button(self.window, text="Take Picture - Count increases here", width=50, height=5, command=lambda: self.save_for_class(2))
        
        self.btn_class_train = tk.Button(self.window, text="Train Model", width=50, command=lambda: self.model.train_model(self.name.get()))

        self.counting_mode_buttons = tk.Button(self.window, text="Count Mode", width=50, command=lambda: self.count_mode())

        self.data_mode_buttons = tk.Button(self.window, text="Gather Data", width=50, command=lambda: self.data_mode())
        self.data_mode_buttons.pack(anchor=tk.CENTER)

        self.all = [
            self.btn_toggleauto,
            self.btn_class_one,
            self.btn_class_two,
            self.btn_class_train,
            self.choose_model,
            self.btn_reset,
            self.counter_label,
            self.counting_mode_buttons,
            self.data_mode_buttons,
            self.input
        ]

    def count_mode(self):
        show = [
            self.btn_toggleauto,
            self.choose_model,
            self.input,
            self.btn_reset,
            self.counter_label,
            self.data_mode_buttons
            ]
        for c in self.all:
            c.pack_forget()

        for c in show:
            c.pack(anchor=tk.CENTER)

    def data_mode(self):
        if os.path.exists("1"):
            shutil.rmtree("1")

        if os.path.exists("2"):
            shutil.rmtree("2")
        show = [
            self.btn_class_train,
            self.input,
            self.btn_class_one,
            self.btn_class_two,
            self.counting_mode_buttons
        ]
        for c in self.all:
            c.pack_forget()
        for c in show:
            c.pack(anchor=tk.CENTER)
        

    def update(self):
        if self.counting_enabled:
            self.predict()

        if self.extended and self.contracted:
            self.extended, self.contracted = False, False
            self.rep_counter += 1

        self.counter_label.config(text=f"{self.rep_counter}")

        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 != self.last_prediction:
            if prediction == 1:
                self.extended = True
                self.last_prediction = 1
            if prediction == 2:
                self.contracted = True
                self.last_prediction = 2

    def counting_toggle(self):
        self.counting_enabled = not self.counting_enabled

    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")

        cv2.imwrite(f"{class_num}/frame{self.counters[class_num-1]}.jpg", cv2.cvtColor(frame, cv2.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):
        self.rep_counter = 0