EV5_Beeldherk_Bomen/src/suite.py

#!/usr/bin/python3
import pathlib
import pygubu
import glob
from tkinter import *
from PIL import ImageTk, Image
import numpy as np
import cv2
import time
import matplotlib.pyplot as plt

PROJECT_PATH = pathlib.Path(__file__).parent
PROJECT_UI = "./src/gui/main.ui"

class MainApp:
    def __init__(self, master=None):
        self.builder = builder = pygubu.Builder()
        builder.add_resource_path(PROJECT_PATH)
        builder.add_from_file(PROJECT_UI)
        
        # Main widget
        self.mainwindow = builder.get_object("main", master)

        self.canvas = builder.get_object("output_canvas")
        self.tk_imgs = []
        
        self.canny_thr1 = None
        self.canny_thr2 = None
        
        self.img_path = None
        self.img_current = 0
        self.img_max = 0

        # Plots
        self.axs = self.createPlot(2, 2)

        # UI
        self.blur_rate = None
        self.img_size = None
        self.sobel_select = None
        self.export_id = None
        builder.import_variables(self,
                                ['canny_thr1',
                                'canny_thr2',
                                'img_path',
                                'blur_rate',
                                'img_size',
                                'sobel_select',
                                'export_id'])
        builder.connect_callbacks(self)

    def run(self):
        self.mainwindow.mainloop()
        
    def img_prev(self, event=None):
        if self.img_current == 0:
            self.img_current = self.img_max - 1
        else:
            self.img_current = self.img_current - 1
        self.update(self)

    def img_next(self, event=None):
        if self.img_current == (self.img_max - 1):
            self.img_current = 0
        else:
            self.img_current = self.img_current + 1
        self.update(self)

    def apply(self, event=None):
        img_arr = self.tk_imgs
        img_id = self.export_id.get()
        
        if (img_id >= 0 and img_id < len(img_arr)):
            print("export")
        else:
            print("Nothing to export!")

    def createPlot(self, columns, rows):
        fig, axs = plt.subplots(columns, rows)
        return axs

    def drawHist(self, image, labels, column, row):
        self.axs[column, row].clear()
        for i,lab in enumerate(labels):
            hist = cv2.calcHist(
                [image],
                [i],
                None,
                [256],
                [0, 256],
            )
            self.axs[column, row].plot(hist, label=lab)
            self.axs[column, row].grid()
            self.axs[column, row].legend()
            
    def drawCannyHM(self, img, column, row):
        self.axs[column, row].clear()
        canny_max = 500
        canny_step = 20

        results = [[] for x in range((int)(canny_max / canny_step))]

        for th1 in range(0, canny_max, canny_step):
            for th2 in range(0, canny_max, canny_step):
                # Canny Edge Detection
                edges = cv2.Canny(image=img, threshold1=th1, threshold2=th2)
                
                w_res = cv2.countNonZero(edges)
                y_ind = (int)(th1 / canny_step)
                x_ind = (int)(th2 / canny_step)
                
                results[y_ind].append(w_res)
                
                # print(f"Result at thres {th1}, {th2}; \tIndex {y_ind}, {x_ind} \t= {w_res}")
                # print(results[y_ind])
        
        self.axs[column, row].imshow(results)
        self.axs[column, row].xaxis.set_major_formatter(lambda x, pos: str(x*canny_step))
        self.axs[column, row].yaxis.set_major_formatter(lambda x, pos: str(x*canny_step))
    
    def update(self, event=None):
        path = self.img_path.get()
        print(path)
        
        drawW = self.canvas.winfo_width()
        drawX = 0
        drawY = 0
        
        if path != None and path != "":
            # Get all images at current path
            images = []
            for file in glob.glob(path + "/*.png"):
                images.append(file)
                
            self.img_max = len(images)
                
            # Get all user vars
            ct1 = self.canny_thr1.get()
            ct2 = self.canny_thr2.get()            
            br = self.blur_rate.get()
            sxy = self.sobel_select.get()
            size = self.img_size.get()
            
            # Keep array of cv2 images to output
            output = []
            
            # Clear previously printed images
            self.tk_imgs = []
            
            # Import and resize image
            img = cv2.imread(images[self.img_current])
            img = cv2.resize(img, (size, size), interpolation = cv2.INTER_AREA)
            output.append(img) # First output image is original
            
            # Set grayscale
            img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            output.append(img_gray) # Second output is grayscale
            
            # Blurred edition
            if (br == 0):
                # Disable
                img_blur = img_gray
            elif (br % 2 == 1):
                img_blur = cv2.GaussianBlur(img_gray, (br, br), 0)
            else:
                print(f"Blur rate changed to {br - 1}")
                img_blur = cv2.GaussianBlur(img_gray, (br-1, br-1), 0)
            output.append(img_blur)
            
            # Sobel edge
            if sxy in ['x', 'y', 'both']:
                if sxy == 'x':
                    dx = 1
                    dy = 0
                elif sxy == 'y':
                    dx = 0
                    dy = 1
                elif sxy == 'both':
                    dx = 1
                    dy = 1
                    
                img_sobel = cv2.Sobel(src=img_blur, ddepth=cv2.CV_8U, dx=dx, dy=dy, ksize=5)
            else:
                img_sobel = img_gray
            output.append(img_sobel)
            
            # Canny edge
            img_canny = cv2.Canny(image=img_blur,threshold1=ct1,threshold2=ct2)
            output.append(img_canny)

            # BGR
            output.append(img[:, :, 0]) # B
            output.append(img[:, :, 1]) # G
            output.append(img[:, :, 2]) # R
            if img is not None:
                self.drawHist(img, ('B', 'G', 'R'), 0, 0)
            
            # HSV
            img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
            output.append(img_hsv)
            output.append(img_hsv[:, :, 0]) # H
            output.append(img_hsv[:, :, 1]) # S
            output.append(img_hsv[:, :, 2]) # V
            if img_hsv is not None:
                self.drawHist(img_hsv, ('H', 'S', 'V'), 0, 1)
            
            self.drawCannyHM(img, 1, 1)
                
            plt.show(block=False)
            
            # Draw all output images
            for idx, data in enumerate(output):
                # Create ui image
                tk_img = cv2.cvtColor(data, cv2.COLOR_BGR2RGB)
                tk_img = ImageTk.PhotoImage(image=Image.fromarray(tk_img))
                self.tk_imgs.append(tk_img)
            
                ## Check if next item will be out of range
                if (drawX + size >= drawW):
                    drawY += size
                    drawX = 0
                    self.canvas.configure(height=(drawY+size))
                    
                self.canvas.create_image(drawX,drawY,anchor=NW,image=self.tk_imgs[idx],tags="og")
                drawX += size
            self.canvas.draw()

if __name__ == "__main__":
    app = MainApp()
    app.run()