thresholding.py

import time
import cv2
import pytesseract
import numpy as np

# SmashScan Libraries
import util

# Refer to OpenCV documentation:
# https://docs.opencv.org/3.4.5/da/d97/tutorial_threshold_inRange.html
# https://docs.opencv.org/3.4.5/de/d25/imgproc_color_conversions.html
# https://docs.opencv.org/3.4/da/d6a/tutorial_trackbar.html
# https://docs.opencv.org/3.4.5/da/d97/tutorial_threshold_inRange.html


# An object that creates a parameter analyzer window for HSV ranges, and two
# separate windows that displays the results of the trackbar ranges.
class HsvParamAnalyzer:

    def __init__(self, video_location, start_fnum=0, stop_fnum=0):

        self.capture = cv2.VideoCapture(video_location)
        self.start_fnum = start_fnum
        self.stop_fnum = stop_fnum
        if stop_fnum == 0:
            self.stop_fnum = int(self.capture.get(cv2.CAP_PROP_FRAME_COUNT))

        self.window_name = 'Object Detection'
        self.low_H_name = 'Low H'
        self.low_S_name = 'Low S'
        self.low_V_name = 'Low V'
        self.high_H_name = 'High H'
        self.high_S_name = 'High S'
        self.high_V_name = 'High V'

        self.low_H, self.low_S, self.low_V = 0, 0, 0
        self.high_H, self.high_S, self.high_V = 180, 255, 255

        cv2.namedWindow(self.window_name)
        cv2.createTrackbar(self.low_H_name, self.window_name,
            self.low_H, 180, self.on_low_H_thresh_trackbar)
        cv2.createTrackbar(self.high_H_name, self.window_name,
            self.high_H, 180, self.on_high_H_thresh_trackbar)
        cv2.createTrackbar(self.low_S_name, self.window_name,
            self.low_S, 255, self.on_low_S_thresh_trackbar)
        cv2.createTrackbar(self.high_S_name, self.window_name,
            self.high_S, 255, self.on_high_S_thresh_trackbar)
        cv2.createTrackbar(self.low_V_name, self.window_name,
            self.low_V, 255, self.on_low_V_thresh_trackbar)
        cv2.createTrackbar(self.high_V_name, self.window_name,
            self.high_V, 255, self.on_high_V_thresh_trackbar)


    # The standard test iterates through the entire video with multiple track
    # bars to vary HSV thresholds. Results can be seen in a separate window.
    def standard_test(self):
        for fnum in range(self.start_fnum, self.stop_fnum):
            frame = util.get_frame(self.capture, fnum)
            frame = frame[280:, :]
            frame_HSV = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)

            mask = cv2.inRange(frame_HSV, (self.low_H, self.low_S, self.low_V),
                (self.high_H, self.high_S, self.high_V))

            res = cv2.bitwise_and(frame, frame, mask=mask)
            res_inv = cv2.bitwise_and(frame, frame, mask=cv2.bitwise_not(mask))

            cv2.imshow(self.window_name, mask)
            cv2.imshow('Video Capture AND', res)
            cv2.imshow('Video Capture INV', res_inv)

            if cv2.waitKey(30) & 0xFF == ord('q'):
                break


    # A number of methods corresponding to the various trackbars available.
    def on_low_H_thresh_trackbar(self, val):
        self.low_H = val
        self.low_H = min(self.high_H-1, self.low_H)
        cv2.setTrackbarPos(self.low_H_name, self.window_name, self.low_H)


    def on_high_H_thresh_trackbar(self, val):
        self.high_H = val
        self.high_H = max(self.high_H, self.low_H+1)
        cv2.setTrackbarPos(self.high_H_name, self.window_name, self.high_H)


    def on_low_S_thresh_trackbar(self, val):
        self.low_S = val
        self.low_S = min(self.high_S-1, self.low_S)
        cv2.setTrackbarPos(self.low_S_name, self.window_name, self.low_S)


    def on_high_S_thresh_trackbar(self, val):
        self.high_S = val
        self.high_S = max(self.high_S, self.low_S+1)
        cv2.setTrackbarPos(self.high_S_name, self.window_name, self.high_S)


    def on_low_V_thresh_trackbar(self, val):
        self.low_V = val
        self.low_V = min(self.high_V-1, self.low_V)
        cv2.setTrackbarPos(self.low_V_name, self.window_name, self.low_V)


    def on_high_V_thresh_trackbar(self, val):
        self.high_V = val
        self.high_V = max(self.high_V, self.low_V+1)
        cv2.setTrackbarPos(self.high_V_name, self.window_name, self.high_V)


# An object that creates a parameter analyzer window for damage OCR.
class DmgParamAnalyzer:

    def __init__(self, video_location, start_fnum=0, stop_fnum=0):
        self.capture = cv2.VideoCapture(video_location)
        self.capture.set(cv2.CAP_PROP_POS_FRAMES, start_fnum)
        self.window_name = "Damage Parameter Analyzer"

        self.start_fnum = start_fnum
        self.stop_fnum = stop_fnum
        if stop_fnum == 0:
            self.stop_fnum = int(self.capture.get(cv2.CAP_PROP_FRAME_COUNT))

        cv2.namedWindow(self.window_name)
        cv2.createTrackbar("Step Size", self.window_name,
            1, 100, self.on_step_trackbar)
        cv2.createTrackbar("Delay", self.window_name,
            10, 500, self.on_delay_trackbar)
        cv2.createTrackbar("Thresh ~ Bin, Otsu", self.window_name,
            0, 1, self.on_thresh_trackbar)
        cv2.createTrackbar("Pre Blur ~ 0, Gaus, Med", self.window_name,
            0, 2, self.on_pre_blur_trackbar)
        cv2.createTrackbar("Post Blur ~ 0, Med", self.window_name,
            0, 1, self.on_post_blur_trackbar)
        cv2.createTrackbar("Contour Filter ~ Off, On", self.window_name,
            0, 1, self.on_contour_trackbar)
        cv2.createTrackbar("Contour Display ~ Off, On", self.window_name,
            0, 1, self.on_contour_disp_trackbar)
        cv2.createTrackbar("Contour Min Area", self.window_name,
            1, 1000, self.on_contour_min_area_trackbar)
        cv2.createTrackbar("Contour Max Area", self.window_name,
            5000, 5000, self.on_contour_max_area_trackbar)
        cv2.createTrackbar("OCR ~ Off, On", self.window_name,
            0, 1, self.on_ocr_trackbar)
        cv2.createTrackbar("OCR ~ 1 Line, 1 Word", self.window_name,
            0, 1, self.on_ocr_mode_trackbar)

        self.step_size = 1
        self.step_delay = 10
        self.pre_blur_val = 0
        self.thresh_flag = False
        self.post_blur_val = 0
        self.contour_flag = False
        self.contour_disp_flag = False
        self.contour_min_area = 1
        self.contour_max_area = 5000
        self.ocr_flag = False
        self.ocr_mode_flag = False


    # The method that must be called to boot up the paramater analysis GUI.
    def standard_test(self):
        fnum = self.start_fnum
        time_queue = list()
        disp_dict = dict()
        while fnum < self.stop_fnum:
            start_time = time.time()
            fnum += self.step_size
            frame = util.get_frame(self.capture, fnum, gray_flag=True)
            frame = frame[300:340, 80:220] # 300:340, 200:320
            frame = self.param_filter(frame)
            if self.contour_flag:
                frame = self.contour_filter(frame)

            if self.ocr_flag:
                conf_text = "--psm 7"       # Single test line mode.
                if self.ocr_mode_flag:      # Single word mode.
                    conf_text = "--psm 8"

                text = pytesseract.image_to_string(cv2.bitwise_not(frame),
                    lang="eng", config=conf_text)
                disp_dict["OCR"] = text

            util.display_pa_fps(start_time, time_queue, disp_dict)
            cv2.imshow(self.window_name, frame)
            if cv2.waitKey(self.step_delay) & 0xFF == ord('q'):
                break


    # Apply filters to frame according to GUI parameters.
    def param_filter(self, frame):
        # Apply pre-blur according to trackbar value.
        if self.pre_blur_val == 1:
            frame = cv2.GaussianBlur(frame, (5, 5), 0)
        elif self.pre_blur_val == 2:
            frame = cv2.medianBlur(frame, 5)

        # Apply a thresholding method according to trackbar value.
        if self.thresh_flag:
            _, frame = cv2.threshold(frame, 127, 255, cv2.THRESH_BINARY)
        else:
            _, frame = cv2.threshold(frame, 127, 255, cv2.THRESH_OTSU)

        # Apply post-blur according to trackbar value.
        if self.post_blur_val:
            frame = cv2.medianBlur(frame, 5)

        return frame


    # Apply filterrs to frame according to contour parameters.
    def contour_filter(self, frame):
        _, contours, _ = cv2.findContours(frame,
            cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        new_frame = np.zeros(frame.shape, np.uint8)
        for i, contour in enumerate(contours):
            c_area = cv2.contourArea(contour)
            if self.contour_min_area <= c_area <= self.contour_max_area:
                mask = np.zeros(frame.shape, np.uint8)
                cv2.drawContours(mask, contours, i, 255, cv2.FILLED)
                mask = cv2.bitwise_and(frame, mask)
                new_frame = cv2.bitwise_or(new_frame, mask)
        frame = new_frame

        if self.contour_disp_flag:
            frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
            cv2.drawContours(frame, contours, -1, (255, 0, 0), 1)

        return frame


    # A number of methods corresponding to the various trackbars available.
    def on_step_trackbar(self, val):
        self.step_size = val

    def on_delay_trackbar(self, val):
        self.step_delay = val

    def on_thresh_trackbar(self, val):
        self.thresh_flag = val

    def on_pre_blur_trackbar(self, val):
        self.pre_blur_val = val

    def on_post_blur_trackbar(self, val):
        self.post_blur_val = val

    def on_contour_trackbar(self, val):
        self.contour_flag = val

    def on_contour_disp_trackbar(self, val):
        self.contour_disp_flag = val

    def on_contour_min_area_trackbar(self, val):
        self.contour_min_area = val

    def on_contour_max_area_trackbar(self, val):
        self.contour_max_area = val

    def on_ocr_trackbar(self, val):
        self.ocr_flag = val

    def on_ocr_mode_trackbar(self, val):
        self.ocr_mode_flag = val


# An object that creates a parameter analyzer for damage template matching.
class DmgTmParamAnalyzer:

    def __init__(self, video_location, start_fnum=0, stop_fnum=0):
        self.capture = cv2.VideoCapture(video_location)
        self.capture.set(cv2.CAP_PROP_POS_FRAMES, start_fnum)
        self.window_name = "Damage TM Parameter Analyzer"

        self.start_fnum = start_fnum
        self.stop_fnum = stop_fnum
        if stop_fnum == 0:
            self.stop_fnum = int(self.capture.get(cv2.CAP_PROP_FRAME_COUNT))

        cv2.namedWindow(self.window_name)
        cv2.createTrackbar("Step Size", self.window_name,
            1, 100, self.on_step_trackbar)
        cv2.createTrackbar("Delay", self.window_name,
            10, 500, self.on_delay_trackbar)
        cv2.createTrackbar("TM ~ Off, On", self.window_name,
            0, 1, self.on_tm_trackbar)

        self.step_size = 1
        self.step_delay = 10
        self.tm_flag = False

        # Read all ten of the damage integer images and extract a binary mask
        # based off of the alpha channel. Also, resize to a 360p height.
        self.orig_num_img, self.orig_num_mask = [None]*10, [None]*10
        self.num_img, self.num_mask = [None]*10, [None]*10
        for i in range(0, 10):
            self.orig_num_img[i], self.orig_num_mask[i] = util.get_image_and_mask(
                "resources/{:}.png".format(i), gray_flag=True)
            self.num_img[i] = util.resize_img(self.orig_num_img[i], 360/480)
            self.num_mask[i] = util.resize_img(self.orig_num_mask[i], 360/480)
        self.num_h, self.num_w = self.num_img[0].shape[:2]


    # The method that must be called to boot up the paramater analysis GUI.
    def standard_test(self):
        fnum = self.start_fnum
        time_queue = list()
        disp_dict = dict()
        while fnum < self.stop_fnum:
            start_time = time.time()
            fnum += self.step_size
            frame = util.get_frame(self.capture, fnum, gray_flag=True)
            frame = frame[300:340, 80:220] # 300:340, 200:320
            if self.tm_flag:
                self.match_dmg_templates(frame)

            util.display_pa_fps(start_time, time_queue, disp_dict)
            cv2.imshow(self.window_name, frame)
            if cv2.waitKey(self.step_delay) & 0xFF == ord('q'):
                break


    # Perform template matching for all ten damage integers.
    def match_dmg_templates(self, frame):
        match_mat, max_val, tl = [None]*10, [0]*10, [(0, 0)]*10
        for i in range(0, 10):
            match_mat[i] = cv2.matchTemplate(frame, self.num_img[0],
                cv2.TM_CCORR_NORMED, mask=self.num_mask[0])
            _, max_val[i], _, tl[i] = cv2.minMaxLoc(match_mat[i])
        # print(max_val[0])
        br = (tl[0][0] + self.num_w, tl[0][1] + self.num_h)
        frame = cv2.rectangle(frame, tl[0], br, (255, 255, 255), 2)

        # Multi-template result searching
        # _, max_val_1, _, tl_1 = cv2.minMaxLoc(np.array(match_mat))
        # print(tl_1)


    # A number of methods corresponding to the various trackbars available.
    def on_step_trackbar(self, val):
        self.step_size = val

    def on_delay_trackbar(self, val):
        self.step_delay = val

    def on_tm_trackbar(self, val):
        self.tm_flag = val