diff --git a/algorithms/CenterDownwards.py b/algorithms/CenterDownwards.py index 654dcd7..664b959 100644 --- a/algorithms/CenterDownwards.py +++ b/algorithms/CenterDownwards.py @@ -21,8 +21,7 @@ def __init__(self, config): # filtering parameters self.averaging_kernel_size = config.averaging_kernel_size - self.gauss_kernel_size = list( - map(int, config.gauss_kernel_size.split(','))) + self.gauss_kernel_size = list(map(int, config.gauss_kernel_size.split(','))) self.dilate_kernel_size = config.dilate_kernel_size self.sigma_x = config.sigma_x @@ -122,8 +121,7 @@ def get_contours(self, binary_mask): """ frame = np.zeros((self.HEIGHT, self.WIDTH, 3)) ret, thresh = cv.threshold(binary_mask, 0, 254, 0) - contours, hierarchy = cv.findContours( - thresh, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE) + contours, hierarchy = cv.findContours(thresh, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE) return contours, frame @@ -151,8 +149,7 @@ def ellipse_slopes(self, contours, black_frame): rect = cv.minAreaRect(con) box = cv.boxPoints(rect) box = np.int0(box) - black_frame = cv.drawContours( - black_frame, [box], 0, (255, 255, 255), 2) + black_frame = cv.drawContours(black_frame, [box], 0, (255, 255, 255), 2) ellipse = cv.fitEllipse(con) self.contours.append(ellipse) cv.ellipse(black_frame, ellipse, (255, 255, 255), 2) @@ -186,13 +183,11 @@ def ellipse_slopes(self, contours, black_frame): slopes.append(slope) lefty = int((-x * vy / vx) + y) righty = int(((cols - x) * vy / vx) + y) - black_frame = cv.line( - black_frame, (cols - 1, righty), (0, lefty), (255, 255, 0), 9) + black_frame = cv.line(black_frame, (cols - 1, righty), (0, lefty), (255, 255, 0), 9) lines.append([cols - 1, righty, 0, lefty]) else: - black_frame = cv.line(black_frame, (int(x), 0), - (int(x), rows - 1), (255, 255, 0), 9) + black_frame = cv.line(black_frame, (int(x), 0), (int(x), rows - 1), (255, 255, 0), 9) lines.append([int(x), 0, int(x), rows - 1]) return lines, slopes, black_frame diff --git a/algorithms/CenterRowAlgorithm.py b/algorithms/CenterRowAlgorithm.py index 1dcc675..fe50ec6 100644 --- a/algorithms/CenterRowAlgorithm.py +++ b/algorithms/CenterRowAlgorithm.py @@ -22,8 +22,7 @@ def __init__(self, config): # filtering parameters self.averaging_kernel_size = config.averaging_kernel_size - self.gauss_kernel_size = list( - map(int, config.gauss_kernel_size.split(','))) + self.gauss_kernel_size = list(map(int, config.gauss_kernel_size.split(','))) self.dilate_kernel_size = config.dilate_kernel_size self.sigma_x = config.sigma_x @@ -55,7 +54,7 @@ def update_lower_hsv(self, next): def update_upper_hsv(self, next): self.HIGH_GREEN = np.array(next) - + def process_frame(self, frame, show): """Uses contouring to create contours around each crop row and uses these contours to find centroid lines, row vanishing point, a center contour and the angle between the center contour and vanishing point\n @@ -77,8 +76,7 @@ def process_frame(self, frame, show): cv.drawContours(black_frame, contours, -1, self.contour_color, 3) # fillPoly fills in the polygons in the frame cv.fillPoly(black_frame, pts=contours, color=self.contour_color) - lines, slopes, ellipse_frame = self.ellipse_slopes( - contours, black_frame) + lines, slopes, ellipse_frame = self.ellipse_slopes(contours, black_frame) if show: Lines.draw_lines_on_frame(lines, black_frame) @@ -86,16 +84,14 @@ def process_frame(self, frame, show): intersections = Lines.get_intersections(lines) x_points = [point[0] for point in intersections] y_points = [point[1] for point in intersections] - vanishing_point = Lines.draw_vanishing_point( - ellipse_frame, x_points, y_points, show) + vanishing_point = Lines.draw_vanishing_point(ellipse_frame, x_points, y_points, show) if vanishing_point: center_contour, angle = self.find_center_contour(vanishing_point) if show: cv.ellipse(black_frame, center_contour, (0, 255, 0), 2) - angle = Lines.calculate_angle_from_v_point( - vanishing_point, self.WIDTH, self.HEIGHT) + angle = Lines.calculate_angle_from_v_point(vanishing_point, self.WIDTH, self.HEIGHT) return black_frame, angle @@ -125,8 +121,7 @@ def get_contours(self, binary_mask): """ frame = np.zeros((self.HEIGHT, self.WIDTH, 3)) ret, thresh = cv.threshold(binary_mask, 0, 254, 0) - contours, hierarchy = cv.findContours( - thresh, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE) + contours, hierarchy = cv.findContours(thresh, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE) cv.drawContours(frame, contours, -1, (0, 255, 0), 2) cv.fillPoly(frame, pts=contours, color=(0, 255, 0)) @@ -148,8 +143,7 @@ def ellipse_slopes(self, contours, black_frame): rect = cv.minAreaRect(cnt) box = cv.boxPoints(rect) box = np.int0(box) - black_frame = cv.drawContours( - black_frame, [box], 0, (255, 255, 255), 2) + black_frame = cv.drawContours(black_frame, [box], 0, (255, 255, 255), 2) ellipse = cv.fitEllipse(cnt) self.contours.append(ellipse) cv.ellipse(black_frame, ellipse, (255, 255, 255), 2) diff --git a/algorithms/CheckRowEnd.py b/algorithms/CheckRowEnd.py index 4572752..58b24c8 100644 --- a/algorithms/CheckRowEnd.py +++ b/algorithms/CheckRowEnd.py @@ -26,7 +26,7 @@ def update_lower_hsv(self, next): def update_upper_hsv(self, next): self.HIGH_GREEN = np.array(next) - + def process_frame(self, frame, show): """Averages values in each row in a mask of the frame. If the number of rows with an average value of zero is greater than req_rows_empty, then frame is row end\n diff --git a/algorithms/HoughAlgorithm.py b/algorithms/HoughAlgorithm.py index 868cb6c..ae2723d 100644 --- a/algorithms/HoughAlgorithm.py +++ b/algorithms/HoughAlgorithm.py @@ -41,7 +41,7 @@ def update_lower_hsv(self, next): def update_upper_hsv(self, next): self.HIGH_GREEN = np.array(next) - + # processFrame function that is called to process a frame of a video # takes in frame mat object obtained from cv2 video.read() def process_frame(self, frame, show=True): @@ -92,15 +92,12 @@ def process_frame(self, frame, show=True): y_points = [point[1] for point in intersections] if show: - vanishing_point = Lines.draw_vanishing_point( - line_img, x_points, y_points, show) + vanishing_point = Lines.draw_vanishing_point(line_img, x_points, y_points, show) else: - vanishing_point = Lines.draw_vanishing_point( - frame, x_points, y_points, show) + vanishing_point = Lines.draw_vanishing_point(frame, x_points, y_points, show) # Calculating angle from vanishing point to (self.WIDTH // 2, 0) - angle = Lines.calculate_angle_from_v_point( - vanishing_point, self.WIDTH, self.HEIGHT) + angle = Lines.calculate_angle_from_v_point(vanishing_point, self.WIDTH, self.HEIGHT) if show: return line_img, angle diff --git a/algorithms/MiniContoursAlgorithm.py b/algorithms/MiniContoursAlgorithm.py index dc3e796..ed7d5c4 100644 --- a/algorithms/MiniContoursAlgorithm.py +++ b/algorithms/MiniContoursAlgorithm.py @@ -72,8 +72,7 @@ def get_centroids(self, mask, num_strips): centroids = [] for i, strip in enumerate(strips): - contours, hierarchy = cv2.findContours( - strip, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + contours, hierarchy = cv2.findContours(strip, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) strip_centroids = [] for contour in contours: M = cv2.moments(contour) @@ -107,8 +106,7 @@ def get_center_hough_lines( # lines: list of [[votes, rho, theta]] of all lines # point_lines: list of [votes, pt1, pt2] of all lines - mask = cv2.inRange(cv2.cvtColor( - frame, cv2.COLOR_BGR2HSV), self.low_green, self.high_green) + mask = cv2.inRange(cv2.cvtColor(frame, cv2.COLOR_BGR2HSV), self.low_green, self.high_green) mask = cv2.medianBlur(mask, 9) # mask = cv2.GaussianBlur(mask, (9,9), 10) # mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, self.morphology_kernel) @@ -136,8 +134,7 @@ def get_center_hough_lines( if show: cv2.line(frame, (0, cut_off_height), (width // 2, 0), self.color_2) - cv2.line(frame, (width, cut_off_height), - (width // 2, 0), self.color_2) + cv2.line(frame, (width, cut_off_height), (width // 2, 0), self.color_2) for i, strip_centroid in enumerate(centroids): if i > int(0.3 * len(centroids)): @@ -152,16 +149,12 @@ def get_center_hough_lines( segmented_points[idx].append([int(x), int(y)]) if show: - cv2.circle(frame, (int(centroid[0]), int( - centroid[1])), 3, self.color_1, -1) - cv2.circle(mask, (int(centroid[0]), int( - centroid[1])), 3, self.color_1, -1) - points_vector.append( - [int(centroid[0]), int(centroid[1])]) + cv2.circle(frame, (int(centroid[0]), int(centroid[1])), 3, self.color_1, -1) + cv2.circle(mask, (int(centroid[0]), int(centroid[1])), 3, self.color_1, -1) + points_vector.append([int(centroid[0]), int(centroid[1])]) else: if show: - cv2.circle(frame, (int(centroid[0]), int( - centroid[1])), 3, self.color_3, -1) + cv2.circle(frame, (int(centroid[0]), int(centroid[1])), 3, self.color_3, -1) c_mask = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR) @@ -217,7 +210,7 @@ def update_lower_hsv(self, next): def update_upper_hsv(self, next): self.HIGH_GREEN = np.array(next) - + def process_frame(self, original_frame, num_strips=60, show=False): # original_frame: BGR frame @@ -243,11 +236,9 @@ def process_frame(self, original_frame, num_strips=60, show=False): x_points = [point[0] for point in intersections] y_points = [point[1] for point in intersections] - vanishing_point = Lines.draw_vanishing_point( - frame, x_points, y_points, show) + vanishing_point = Lines.draw_vanishing_point(frame, x_points, y_points, show) # Calculating angle from vanishing point to (self.WIDTH // 2, 0) - angle = Lines.calculate_angle_from_v_point( - vanishing_point, self.WIDTH, self.HEIGHT) + angle = Lines.calculate_angle_from_v_point(vanishing_point, self.WIDTH, self.HEIGHT) return frame, angle diff --git a/algorithms/MiniContoursDownwards.py b/algorithms/MiniContoursDownwards.py index 5bccaa4..c67cf59 100644 --- a/algorithms/MiniContoursDownwards.py +++ b/algorithms/MiniContoursDownwards.py @@ -24,14 +24,11 @@ def __init__(self, config): self.high_green = np.array(config.high_green) # random colors for drawing lines - # teal (for contour points that are within bounds) - self.color_1 = (255, 255, 0) + self.color_1 = (255, 255, 0) # teal (for contour points that are within bounds) # pink (for the line of best fit through all the contour points that are within bounds) self.color_2 = (200, 0, 255) - # red (for contour points that are out of bounds) - self.color_3 = (0, 0, 255) - # orange (for a reference line vertically down center of frame) - self.midline = (0, 129, 255) + self.color_3 = (0, 0, 255) # red (for contour points that are out of bounds) + self.midline = (0, 129, 255) # orange (for a reference line vertically down center of frame) # parameters for calculating centroids and drawing the best fit line among them self.num_strips = self.config.num_strips @@ -68,8 +65,7 @@ def get_centroids(self, mask, num_strips, show): centroids = [] for i, strip in enumerate(strips): - contours, hierarchy = cv2.findContours( - strip, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + contours, hierarchy = cv2.findContours(strip, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) strip_centroids = [] for contour in contours: M = cv2.moments(contour) @@ -93,8 +89,7 @@ def get_best_fit_line(self, frame, show, num_strips=60, dist_type=cv2.DIST_L2, p -line: A vector [vx, vy, x, y] representing the line of best fit through all the centroids. [vx, vy] is a vector that describes the direction of the line, where (x, y) when taken together is a point on the line """"" - mask = cv2.inRange(cv2.cvtColor( - frame, cv2.COLOR_BGR2HSV), self.low_green, self.high_green) + mask = cv2.inRange(cv2.cvtColor(frame, cv2.COLOR_BGR2HSV), self.low_green, self.high_green) mask = cv2.medianBlur(mask, 9) # mask = cv2.GaussianBlur(mask, (9,9), 10) # mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, self.morphology_kernel) @@ -116,10 +111,8 @@ def get_best_fit_line(self, frame, show, num_strips=60, dist_type=cv2.DIST_L2, p idx = int(x / width * split_factor) segmented_points[idx].append([int(x), int(y)]) if show: - cv2.circle(frame, (int(centroid[0]), int( - centroid[1])), 3, self.color_1, -1) - cv2.circle(mask, (int(centroid[0]), int( - centroid[1])), 3, self.color_1, -1) + cv2.circle(frame, (int(centroid[0]), int(centroid[1])), 3, self.color_1, -1) + cv2.circle(mask, (int(centroid[0]), int(centroid[1])), 3, self.color_1, -1) points_vector.append([int(centroid[0]), int(centroid[1])]) c_mask = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR) @@ -168,7 +161,7 @@ def update_lower_hsv(self, next): def update_upper_hsv(self, next): self.HIGH_GREEN = np.array(next) - + def process_frame(self, original_frame, num_strips=60, show=False): """"" parameters: @@ -191,8 +184,7 @@ def process_frame(self, original_frame, num_strips=60, show=False): if line is not None and line[0] is not None: angle = round(math.degrees(math.atan(-line[0] / line[1])), 2) - cv2.line(frame, (int(self.WIDTH / 2), 0), - (int(self.WIDTH / 2), int(self.HEIGHT)), self.midline, 2) + cv2.line(frame, (int(self.WIDTH / 2), 0), (int(self.WIDTH / 2), int(self.HEIGHT)), self.midline, 2) print(angle) return frame, angle else: diff --git a/algorithms/ScanningAlgorithm.py b/algorithms/ScanningAlgorithm.py index d10e567..1a02dbf 100644 --- a/algorithms/ScanningAlgorithm.py +++ b/algorithms/ScanningAlgorithm.py @@ -26,8 +26,7 @@ def __init__(self, config): self.right_x_bound = int(self.WIDTH * (1 - self.config.bounding_box_x)) self.upper_y_bound = int(self.HEIGHT * self.config.bounding_box_y) self.lower_y_bound = self.HEIGHT - 1 - self.mid_y = self.lower_y_bound - \ - (self.lower_y_bound - self.upper_y_bound) // 2 + self.mid_y = self.lower_y_bound - (self.lower_y_bound - self.upper_y_bound) // 2 self.kernel = np.ones((5, 5), np.uint8) @@ -38,16 +37,14 @@ def __init__(self, config): # creates lines from the top of the horizon to the bottom for top_x in range(0, self.WIDTH, self.pixel_gap): for bottom_x in range(0, self.WIDTH, self.pixel_gap): - line = self.create_line( - top_x, self.upper_y_bound, bottom_x, self.lower_y_bound) + line = self.create_line(top_x, self.upper_y_bound, bottom_x, self.lower_y_bound) self.lines.append(line) # creates lines from horizon to the right side for top_x in range(0, self.WIDTH, self.pixel_gap): for right_y in range(self.mid_y, self.lower_y_bound, self.pixel_gap): - line1 = self.create_line( - top_x, self.upper_y_bound, self.WIDTH - 1, right_y) + line1 = self.create_line(top_x, self.upper_y_bound, self.WIDTH - 1, right_y) self.lines.append(line1) # creates lines from horizon to the left side @@ -86,7 +83,7 @@ def update_lower_hsv(self, next): def update_upper_hsv(self, next): self.HIGH_GREEN = np.array(next) - + def process_frame(self, frame, show): hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) @@ -101,8 +98,7 @@ def process_frame(self, frame, show): return mask, None self.upper_y_bound = white_pixels[0][0] - self.mid_y = self.lower_y_bound - \ - (self.lower_y_bound - self.upper_y_bound) // 2 + self.mid_y = self.lower_y_bound - (self.lower_y_bound - self.upper_y_bound) // 2 # use this to invert the mask # mask = ~mask @@ -135,8 +131,7 @@ def process_frame(self, frame, show): most_prominent_pos_lines = pos_array[:, 1][largest_pos_indices] most_prominent_neg_lines = neg_array[:, 1][largest_neg_indices] - most_prominent_lines = numpy.concatenate( - (most_prominent_pos_lines, most_prominent_neg_lines), axis=None) + most_prominent_lines = numpy.concatenate((most_prominent_pos_lines, most_prominent_neg_lines), axis=None) # convert to lines as defined in Lines.py converted_lines = [] @@ -151,33 +146,27 @@ def process_frame(self, frame, show): x_points = [point[0] for point in intersections] y_points = [point[1] for point in intersections] - vanishing_point = Lines.draw_vanishing_point( - frame, x_points, y_points, show) + vanishing_point = Lines.draw_vanishing_point(frame, x_points, y_points, show) if show: for line in converted_lines: - frame = cv2.line( - frame, (line[0], line[1]), (line[2], line[3]), (255, 255, 255), 1) + frame = cv2.line(frame, (line[0], line[1]), (line[2], line[3]), (255, 255, 255), 1) # line for the middle of frame - frame = cv2.line(frame, (self.WIDTH // 2, self.HEIGHT), - (self.WIDTH // 2, 0), (0, 0, 255), 1) + frame = cv2.line(frame, (self.WIDTH // 2, self.HEIGHT), (self.WIDTH // 2, 0), (0, 0, 255), 1) # point with x coordinate of the vanishing point and y coordinate of the end of the crop row - frame = cv2.circle( - frame, (vanishing_point[0], self.upper_y_bound), 5, (0, 255, 0), -1) + frame = cv2.circle(frame, (vanishing_point[0], self.upper_y_bound), 5, (0, 255, 0), -1) # point in the middle of frame at midpoint between the horizon and bottom of the screen - frame = cv2.circle( - frame, (self.WIDTH // 2, self.mid_y), 5, (0, 255, 0), -1) + frame = cv2.circle(frame, (self.WIDTH // 2, self.mid_y), 5, (0, 255, 0), -1) # line between the two points above frame = cv2.line(frame, (self.WIDTH // 2, self.mid_y), (vanishing_point[0], self.upper_y_bound), (0, 255, 0), 2) # Calculating angle from vanishing point to (self.WIDTH // 2, 0) - angle = Lines.calculate_angle_from_v_point( - vanishing_point, self.WIDTH, self.HEIGHT) + angle = Lines.calculate_angle_from_v_point(vanishing_point, self.WIDTH, self.HEIGHT) return frame, angle diff --git a/algorithms/SeesawAlgorithm.py b/algorithms/SeesawAlgorithm.py index c48e23d..f729645 100644 --- a/algorithms/SeesawAlgorithm.py +++ b/algorithms/SeesawAlgorithm.py @@ -15,8 +15,7 @@ def __init__(self, config): # filtering parameters self.averaging_kernel_size = config.averaging_kernel_size - self.gauss_kernel_size = list( - map(int, config.gauss_kernel_size.split(','))) + self.gauss_kernel_size = list(map(int, config.gauss_kernel_size.split(','))) self.dilate_kernel_size = config.dilate_kernel_size self.sigma_x = config.sigma_x @@ -33,7 +32,7 @@ def update_lower_hsv(self, next): def update_upper_hsv(self, next): self.HIGH_GREEN = np.array(next) - + def process_frame(self, frame, show): black_frame, points, both_points = self.plot_points(frame) @@ -46,13 +45,11 @@ def process_frame(self, frame, show): x2 = int(x + vx * self.WIDTH) y1 = int(y - vy * self.HEIGHT) y2 = int(y + vy * self.HEIGHT) - black_frame = cv.line(black_frame, (x1, y1), - (x2, y2), (0, 255, 255), 9) + black_frame = cv.line(black_frame, (x1, y1), (x2, y2), (0, 255, 255), 9) """calculate angle""" if y1 - y2 != 0: - angle = round(math.degrees( - math.atan(int(x2 - x1) / int(y1 - y2))), 2) + angle = round(math.degrees(math.atan(int(x2 - x1) / int(y1 - y2))), 2) else: angle = None @@ -82,8 +79,7 @@ def plot_points(self, frame): while square_low < self.HEIGHT: 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] + seg_right = right[int(square_low) + 1:int(square_high), 0:half_width] left_x = int(np.sum(seg_left == 255) / bar_height) right_x = int(np.sum(seg_right == 255) / bar_height) @@ -101,12 +97,10 @@ def plot_points(self, frame): black_frame = cv.rectangle(black_frame, (half_width, square_low), ( x2, int(square_high)), (255, 255, 0), 3) - both_point = [int((x1 + x2) / 2), - int((square_high + square_low) / 2)] + both_point = [int((x1 + x2) / 2), int((square_high + square_low) / 2)] both_points.append(both_point) - black_frame = cv.circle( - black_frame, both_point, radius=0, color=(0, 0, 255), thickness=15) + black_frame = cv.circle(black_frame, both_point, radius=0, color=(0, 0, 255), thickness=15) square_high += bar_height square_low += bar_height diff --git a/algorithms/utils/Lines.py b/algorithms/utils/Lines.py index e6df41e..048d897 100644 --- a/algorithms/utils/Lines.py +++ b/algorithms/utils/Lines.py @@ -89,8 +89,7 @@ def get_intersections(lines, min_slope=1): x1R, y1R, x2R, y2R = assign_coordinate_values(lineR) # calls the getIntersection helper function - intersect = get_intersection( - ((x1L, y1L), (x2L, y2L)), ((x1R, y1R), (x2R, y2R))) + intersect = get_intersection(((x1L, y1L), (x2L, y2L)), ((x1R, y1R), (x2R, y2R))) if isinstance(intersect, bool): continue @@ -229,8 +228,7 @@ def draw_vanishing_point(frame, x_points, y_points, show, use_median=True): intersecting_x = np.mean(filtered_x) intersecting_y = np.mean(filtered_y) if show: - cv.circle(frame, (int(intersecting_x), int( - intersecting_y)), 8, (255, 0, 0), -1) + cv.circle(frame, (int(intersecting_x), int(intersecting_y)), 8, (255, 0, 0), -1) return (int(intersecting_x), int(intersecting_y)) else: return None diff --git a/algorithms/utils/delete_small_contours.py b/algorithms/utils/delete_small_contours.py index cdcd10b..ffd6f93 100644 --- a/algorithms/utils/delete_small_contours.py +++ b/algorithms/utils/delete_small_contours.py @@ -36,11 +36,9 @@ def morph_op(mask, kernel_size, op, iterations): if op == 2: mask = cv.dilate(mask, kernel, iterations=iterations) if op == 3: - mask = cv.morphologyEx(mask, cv.MORPH_CLOSE, - kernel, iterations=iterations) + mask = cv.morphologyEx(mask, cv.MORPH_CLOSE, kernel, iterations=iterations) if op == 4: - mask = cv.morphologyEx(mask, cv.MORPH_OPEN, - kernel, iterations=iterations) + mask = cv.morphologyEx(mask, cv.MORPH_OPEN, kernel, iterations=iterations) else: mask = mask diff --git a/create_dataset.py b/create_dataset.py index 5f2a4ae..0cb2b91 100644 --- a/create_dataset.py +++ b/create_dataset.py @@ -23,8 +23,7 @@ def main(): # verify that video exists in ./videos if not path.isfile(f'videos/{video_name}.mp4'): - print('--vid', video_name, - "is an invalid video name, make sure video exists in ./videos") + print('--vid', video_name, "is an invalid video name, make sure video exists in ./videos") sys.exit() # make sure any important data does not get overwritten diff --git a/gui.py b/gui.py index 9eb7417..74112bb 100644 --- a/gui.py +++ b/gui.py @@ -34,18 +34,15 @@ def __init__(self, master, img_dict, window_name): self.container.pack(side='top', anchor='nw', fill="both", expand=True) self.canvas = tk.Canvas(self.container) - self.scrollbar = tk.Scrollbar( - self.container, orient="vertical", command=self.canvas.yview) + self.scrollbar = tk.Scrollbar(self.container, orient="vertical", command=self.canvas.yview) self.scrollable_frame = tk.Frame(self.canvas) - self.scrollable_frame.bind("", lambda e: self.canvas.configure( - scrollregion=self.canvas.bbox("all"))) + self.scrollable_frame.bind("", lambda e: self.canvas.configure(scrollregion=self.canvas.bbox("all"))) - self.img_container = tk.Label(self.scrollable_frame, padx=10, pady=10) + self.img_container = tk.Label(self.scrollable_frame) self.img_container.pack(side='top', fill="both", expand=True) # self.img_container.config(width=600, height=400) - self.canvas.create_window( - (0, 0), window=self.scrollable_frame, anchor='c') + self.canvas.create_window((0, 0), window=self.scrollable_frame, anchor='c') self.canvas.configure(yscrollcommand=self.scrollbar.set) self.canvas.pack(side="left", fill="both", expand=True) @@ -73,36 +70,34 @@ def __init__(self, master, img_dict, window_name): # # self.upper_frame = tk.Frame(self.scrollable_frame) - self.upper_frame.pack(in_=self.scrollable_frame, - anchor="c", side="bottom") + self.upper_frame.pack(in_=self.scrollable_frame, anchor="c", side="bottom") self.lower_frame = tk.Frame(self.scrollable_frame) - self.lower_frame.pack(in_=self.scrollable_frame, - anchor="c", side="bottom") + self.lower_frame.pack(in_=self.scrollable_frame, anchor="c", side="bottom") self.alert_for_hsv = tk.Label( self.lower_frame, text="", fg="red") self.alert_for_hsv.pack(pady=5, side="top") # # - self.LOWER_GREEN = [35, 80, 80] + self.LOW_GREEN = [35, 80, 80] self.low_h_entry_label = tk.Label(self.lower_frame, text="LOWER-H:") self.low_h_entry_label.pack(pady=5, side="left") self.low_h_entry = tk.Entry( self.lower_frame, width=5, justify="center", font="Courier 12") - self.low_h_entry.insert(0, str(self.LOWER_GREEN[0])) + self.low_h_entry.insert(0, str(self.LOW_GREEN[0])) self.low_h_entry.pack(pady=5, side="left") # self.low_s_entry_label = tk.Label(self.lower_frame, text="S:") self.low_s_entry_label.pack(pady=5, side="left") self.low_s_entry = tk.Entry( self.lower_frame, width=5, justify="center", font="Courier 12") - self.low_s_entry.insert(0, str(self.LOWER_GREEN[1])) + self.low_s_entry.insert(0, str(self.LOW_GREEN[1])) self.low_s_entry.pack(pady=5, side="left") # self.low_v_entry_label = tk.Label(self.lower_frame, text="V:") self.low_v_entry_label.pack(pady=5, side="left") self.low_v_entry = tk.Entry( self.lower_frame, width=5, justify="center", font="Courier 12") - self.low_v_entry.insert(0, str(self.LOWER_GREEN[2])) + self.low_v_entry.insert(0, str(self.LOW_GREEN[2])) self.low_v_entry.pack(pady=5, side="left") # self.update_btn_low = tk.Button( @@ -144,70 +139,48 @@ def update_upper_hsv(self): upper_h = int(self.up_h_entry.get()) upper_s = int(self.up_s_entry.get()) upper_v = int(self.up_v_entry.get()) - if upper_h > self.LOWER_GREEN[0] and upper_s > self.LOWER_GREEN[1] and upper_v > self.LOWER_GREEN[2] and upper_h <= 255 and upper_s <= 255 and upper_v <= 255: + if upper_h > self.LOW_GREEN[0] and upper_s > self.LOW_GREEN[1] and upper_v > self.LOW_GREEN[2]: self.UPPER_GREEN = (upper_h, upper_s, upper_v) self.alert_for_hsv.config( - text="new UPPER-HSV: " + str(self.UPPER_GREEN)) + text="") else: - warning = "INVALID: " - if upper_h <= self.LOWER_GREEN[0]: - warning += "UPPER_H <= LOWER_H\t" - elif upper_s <= self.LOWER_GREEN[1]: - warning += "UPPER_S <= LOWER_S\t" - elif upper_v <= self.LOWER_GREEN[2]: - warning += "UPPER_V <= LOWER_V\t" - else: - warning += "UPPER_V > 255\t" - self.alert_for_hsv.config( - text=warning) - + text="Invalid UPPER HSV values") self.up_h_entry.delete(0, tk.END) self.up_h_entry.insert( 0, str(self.UPPER_GREEN[0])) - self.up_s_entry.delete(0, tk.END) - self.up_s_entry.insert( - 0, str(self.UPPER_GREEN[1])) self.up_v_entry.delete(0, tk.END) self.up_v_entry.insert( + 0, str(self.UPPER_GREEN[1])) + self.up_s_entry.delete(0, tk.END) + self.up_s_entry.insert( 0, str(self.UPPER_GREEN[2])) - # print(self.UPPER_GREEN) + print(self.UPPER_GREEN) def update_lower_hsv(self): lower_h = int(self.low_h_entry.get()) lower_s = int(self.low_s_entry.get()) lower_v = int(self.low_v_entry.get()) - if lower_h < self.UPPER_GREEN[0] and lower_s < self.UPPER_GREEN[1] and lower_v < self.UPPER_GREEN[2] and lower_h >= 0 and lower_s >= 0 and lower_v >= 0: + if lower_h < self.UPPER_GREEN[0] and lower_s < self.UPPER_GREEN[1] and lower_v < self.UPPER_GREEN[2]: self.LOWER_GREEN = (lower_h, lower_s, lower_v) self.alert_for_hsv.config( - text="new LOWER-HSV: " + str(self.LOWER_GREEN)) + text="") else: - warning = "INVALID: " - if lower_h >= self.UPPER_GREEN[0]: - warning += "UPPER_H <= LOWER_H\t" - elif lower_s >= self.UPPER_GREEN[1]: - warning += "UPPER_S <= LOWER_S\t" - elif lower_v >= self.UPPER_GREEN[2]: - warning += "UPPER_V <= LOWER_V\t" - else: - warning += "LOWER_V < 0\t" - self.alert_for_hsv.config( - text=warning) - + text="Invalid LOWER HSV values") self.low_h_entry.delete(0, tk.END) self.low_h_entry.insert( - 0, str(self.LOWER_GREEN[0])) - self.low_s_entry.delete(0, tk.END) - self.low_s_entry.insert( - 0, str(self.LOWER_GREEN[1])) + 0, str(self.LOW_GREEN[0])) self.low_v_entry.delete(0, tk.END) self.low_v_entry.insert( - 0, str(self.LOWER_GREEN[2])) - # print(self.LOWER_GREEN) + 0, str(self.LOW_GREEN[1])) + self.low_s_entry.delete(0, tk.END) + self.low_s_entry.insert( + 0, str(self.LOW_GREEN[2])) + print(self.LOW_GREEN) def getLowerHSV(self): - return self.LOWER_GREEN + return self.LOW_GREEN def getUpperHSV(self): return self.UPPER_GREEN @@ -258,17 +231,15 @@ def render_image(self): # Resize the image to fit within the available space, while maintaining aspect ratio new_width = int(img_width * scale_factor) new_height = int(img_height * scale_factor) - if new_height > 10 and new_width > 10: - new_width -= 10 - new_height -= 10 curr_img = curr_img.resize((new_width, new_height), Image.ANTIALIAS) curr_img = ImageTk.PhotoImage(curr_img) - self.img_container.config(image=curr_img, padx=10, pady=10) + self.img_container.config(image=curr_img) self.img_container.image = curr_img self.fps_label.config( - text="Frames Per Second: ~" + str(self.fps), padx=10, pady=10) + text="Frames Per Second: ~" + str(self.fps)) self.master.update() + def isActive(self): global isActive diff --git a/helper_scripts/DataExtractor.py b/helper_scripts/DataExtractor.py index e650d5a..5b06752 100644 --- a/helper_scripts/DataExtractor.py +++ b/helper_scripts/DataExtractor.py @@ -46,20 +46,17 @@ def extract(self): success, frame = stream.read() h, w, d = frame.shape - out = cv.VideoWriter(f'{self.video_path}/{self.name}_key.mp4', - cv.VideoWriter_fourcc(*'mp4v'), 1, (w, h)) + out = cv.VideoWriter(f'{self.video_path}/{self.name}_key.mp4', cv.VideoWriter_fourcc(*'mp4v'), 1, (w, h)) count = 0 while success: # if frame is a keyframe create a red, vertical line in the middle of the frame and save it if count % self.interval == 0: - frame = cv.line(frame, (w // 2, 0), (w // 2, h), - (0, 0, 255), thickness=2) + frame = cv.line(frame, (w // 2, 0), (w // 2, h), (0, 0, 255), thickness=2) self.frames.append(frame) self.frames_copy.append(frame.copy()) out.write(frame) - cv.imwrite( - f'{self.keyframes_path}/{self.name}_{count}.jpg', frame) + cv.imwrite(f'{self.keyframes_path}/{self.name}_{count}.jpg', frame) print('frame ', count) success, frame = stream.read() count += 1 diff --git a/test_algorithms.py b/test_algorithms.py index debc10f..1ad1564 100644 --- a/test_algorithms.py +++ b/test_algorithms.py @@ -30,8 +30,7 @@ def main(): # verify that video exists in ./videos if not path.isfile(f'videos/{args.vid}.mp4'): - print('--vid', args.vid, - "is an invalid video name, make sure it video exists in ./videos") + print('--vid', args.vid, "is an invalid video name, make sure it video exists in ./videos") sys.exit() # verify that config file for video exists @@ -41,8 +40,7 @@ def main(): # verify that config file for algorithm exists if not path.isfile(f'config/algorithm/{args.alg}.yaml'): - print( - f"--alg config for {args.alg} is not defined in ./config/algorithm/") + print(f"--alg config for {args.alg} is not defined in ./config/algorithm/") sys.exit() # set video config @@ -64,8 +62,7 @@ def main(): exists = True if not exists: - print( - f"{args.alg} is an invalid algorithm, list of valid argument values: {algo_list}") + print(f"{args.alg} is an invalid algorithm, list of valid argument values: {algo_list}") sys.exit() @@ -88,8 +85,7 @@ def run_algorithm(alg, vid_file): print(angle) if args.show: - cv.imshow( - f'{args.alg}s algorithm on {args.vid}s video', processed_image) + cv.imshow(f'{args.alg}s algorithm on {args.vid}s video', processed_image) key = cv.waitKey(25)