diff --git a/algorithms/SeesawAlgorithm.py b/algorithms/SeesawAlgorithm.py index c5f38e3..3da0a3c 100644 --- a/algorithms/SeesawAlgorithm.py +++ b/algorithms/SeesawAlgorithm.py @@ -8,7 +8,10 @@ class SeesawAlgorithm(Algorithm): def __init__(self, config): - + """ + Sets seesaw algorithm configurations + :param config: config params + """ # masking range for green self.LOW_GREEN = np.array(config.lower_hsv_threshold) self.HIGH_GREEN = np.array(config.upper_hsv_threshold) @@ -23,13 +26,25 @@ def __init__(self, config): self.HEIGHT = int(config.frame_length) self.WIDTH = int(config.frame_width) + # visual parameters + self.BAR_HEIGHT = config.bar_height + self.NORM_FACTOR = config.normalization_factor + def process_frame(self, frame, show): + """ + Divides screen into horizontal strips and draw bars according to the amount + of green present on each strip. Draw the best fit line based on the centre of each bar, + then calculate the angle between the best fit line and a horizontal line. + :param frame: current frame (mat) + :param show: show/hide frames on screen for debugging + :type show: bool + :return: processed frame (mat), angle [-90, 90] + """ - black_frame, points, both_points = self.plot_points(frame) - points = np.array(points) + black_frame, both_points = self.plot_points(frame) both_points = np.array(both_points) - """get best fit line for middle points""" + """get best fit line for centre points""" [vx, vy, x, y] = cv.fitLine(both_points, cv.DIST_L2, 0, 0.01, 0.01) x1 = int(x - vx * self.WIDTH) x2 = int(x + vx * self.WIDTH) @@ -37,18 +52,36 @@ def process_frame(self, frame, show): y2 = int(y + vy * self.HEIGHT) black_frame = cv.line(black_frame, (x1, y1), (x2, y2), (0, 255, 255), 9) - """calculate angle""" + # calculate angle if y1 - y2 != 0: angle = round(math.degrees(math.atan(int(x2 - x1) / int(y1 - y2))), 2) else: angle = None + # alternative way of calculating angle + ''' + angle = 0 + for point in both_points: + if point[0] > self.WIDTH / 2: + angle += 1 + elif point[0] < self.WIDTH / 2: + angle -= 1 + angle = angle/(len(both_points))*90 + ''' + return black_frame, angle def plot_points(self, frame): - """This value needs to be changed to change the height of the bars""" - bar_height = int(90) + """ + Divides screen into equally sized rectangles on the left and right side + and draw bars according to the amount of green present on each strip. + Calculates the centre point of each horizontal strip. + :param frame: current frame (mat) + :return: processed frame (mat), list of centre points + """ + # initializing parameters + bar_height = int(self.BAR_HEIGHT) mask = self.create_binary_mask(frame) half_width = int(self.WIDTH / 2) @@ -58,16 +91,18 @@ def plot_points(self, frame): square_low = 0 square_high = bar_height both_points = [] - xs = [] normalized = False + isnull = True black_frame = frame - """draw rectangle and point for every square, add points to array""" while square_low < self.HEIGHT: + + # for each area, calculates the amount of green present normalized = False + seg_left = left[int(square_low) + 1:int(square_high), 0:half_width] seg_right = right[int(square_low) + 1:int(square_high), 0:half_width] @@ -77,16 +112,21 @@ def plot_points(self, frame): if left_x > half_width / 2 or right_x > half_width / 2: normalized = True + if left_x != 0 or right_x != 0: + isnull = False + xs.append([left_x, right_x]) + # draw bars based on the amount of green present x1 = half_width - left_x x2 = half_width + right_x - black_frame = cv.rectangle(black_frame, (half_width, square_low), ( - x1, int(square_high)), (255, 255, 0), 3) - black_frame = cv.rectangle(black_frame, (half_width, square_low), ( - x2, int(square_high)), (255, 255, 0), 3) + black_frame = cv.rectangle(black_frame, (half_width, square_low), + (x1, int(square_high)), (255, 255, 0), 3) + black_frame = cv.rectangle(black_frame, (half_width, square_low), + (x2, int(square_high)), (255, 255, 0), 3) + # draw centre points of the bars, add points to a list both_point = [int((x1 + x2) / 2), int((square_high + square_low) / 2)] both_points.append(both_point) @@ -95,17 +135,25 @@ def plot_points(self, frame): square_high += bar_height square_low += bar_height - if normalized is False: + # normalize the lengths of the bars + if isnull: + normalized = True + + while normalized is False: for points in xs: - points[0] *= 2 - points[1] *= 2 + points[0] = float(self.NORM_FACTOR * points[0]) + points[1] = float(self.NORM_FACTOR * points[1]) + + if points[0] > half_width / 2 or points[1] > half_width / 2: + normalized = True + # draw normalized bars square_low = 0 square_high = bar_height for points in xs: - x1 = half_width - points[0] - x2 = half_width + points[1] + x1 = int(half_width - points[0]) + x2 = int(half_width + points[1]) black_frame = cv.rectangle(black_frame, (half_width, square_low), ( x1, int(square_high)), (255, 255, 255), 3) @@ -115,7 +163,7 @@ def plot_points(self, frame): square_high += bar_height square_low += bar_height - return black_frame, points, both_points + return black_frame, both_points def create_binary_mask(self, frame): """ diff --git a/config/algorithm/seesaw.yaml b/config/algorithm/seesaw.yaml index 40027f4..9eebc45 100644 --- a/config/algorithm/seesaw.yaml +++ b/config/algorithm/seesaw.yaml @@ -4,8 +4,8 @@ gauss_kernel_size: 3,3 dilate_kernel_size: (5,5) sigma_x: 0 -contour_color: 0,129,255 -bar_height: 90 +bar_height: 30 +normalization_factor: 1.5 #### size of frame frame_width: 1280 @@ -13,5 +13,4 @@ frame_length: 720 #### hsv thresholds for crop video lower_hsv_threshold: [35, 80, 80] -upper_hsv_threshold: [80, 255, 255] - +upper_hsv_threshold: [80, 255, 255] \ No newline at end of file